hacks

(tidier examples of) random scripts from throughout the years
Log | Files | Refs | README

sha.cu (6941B)


      1 /* Editor's note: a partially parallelized SHA256 implementation.
      2  * Written around December 2021 for my high school computer security class.
      3  * If I remember correctly the parallelization was done very poorly and there's room
      4  * for a lot more optimization. 
      5  *
      6  * The message schedules for the chunks processed by SHA256 are not sequentially
      7  * dependent and so this program generates them in parallel and then processes them
      8  * sequentially when updating the hash.
      9  */
     10 
     11 #include <cstdint>
     12 #include <cstring>
     13 #include <unistd.h>
     14 #include <stdio.h>
     15 #include <fcntl.h>
     16 #include <sys/types.h>
     17 #include <sys/stat.h>
     18 #include <cuda_runtime_api.h>
     19 #include <cuda.h>
     20 
     21 // 32-bit bit right rotation
     22 #define rotr(a,b) (((a) >> (b)) | ((a) << (32-(b))))
     23 
     24 // 32-bit byte swap
     25 #define bswap32(x) ((x>>24)&0xff) |             \
     26     ((x<<8)&0xff0000) |                         \
     27     ((x>>8)&0xff00) |                           \
     28     ((x<<24)&0xff000000)                        \
     29 
     30 uint64_t k[64] = {
     31     0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
     32     0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
     33     0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
     34     0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
     35     0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
     36     0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
     37     0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
     38     0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2,
     39 };
     40 
     41 struct sha_ctx {
     42     uint32_t hash[8] = {
     43         0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a,
     44         0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19,
     45     };
     46     uint64_t len;
     47     
     48     sha_ctx(uint64_t len);
     49     void compress(uint32_t* w);
     50     void dump_hash();
     51 };
     52 
     53 sha_ctx::sha_ctx(uint64_t length) :len(length) {}
     54 
     55 /*
     56  * GPU Kernel responsible for generating message schedules.
     57  * Takes:
     58  * - `bytes`, a pointer to a buffer containing the input from the file.
     59  * - `w`, a pointer to the large buffer to write all the message schedules to.
     60  * - `iters`, the number of message schedules to generate per thread.
     61  * Does not return anything but instead modifies `w`
     62  */
     63 __global__ void process(const uint8_t* bytes, uint32_t* w, uint32_t iters) {
     64     const size_t start_chunk = (blockIdx.y*gridDim.x*blockDim.x)+(blockIdx.x*blockDim.x)+threadIdx.x;
     65     for (int off = 0; off < iters; off++) {
     66         uint32_t* w_adj = w+(((start_chunk*iters)+off)*64);
     67         memcpy(w_adj, bytes+(((start_chunk*iters)+off)*64), 64);
     68         for (int i = 0; i < 16; i++) w_adj[i] = bswap32(w_adj[i]);
     69         for (int i = 16; i < 64; i++) {
     70             uint32_t s0 = (rotr(w_adj[i-15], 7) ^ rotr(w_adj[i-15], 18) ^ (w_adj[i-15] >> 3));
     71             uint32_t s1 = (rotr(w_adj[i-2], 17) ^ rotr(w_adj[i-2], 19)  ^ (w_adj[i-2] >> 10));
     72             w_adj[i] = w_adj[i-16] + s0 + w_adj[i-7] + s1;
     73         }
     74     }
     75 }
     76 
     77 /*
     78  * CPU-side function that updates the hash for a given message schedule.
     79  * Takes:
     80  * - `w`, a pointer to a single message schedule.
     81  * Does not return anything but modifies `sha_ctx::hash`
     82  */
     83 void sha_ctx::compress(uint32_t* w) {
     84     uint32_t a[8] = {0};
     85     memcpy(a, hash, 8*sizeof(uint32_t));
     86     for (int i = 0; i < 64; i++) {
     87         uint32_t s1 = (rotr(a[4], 6) ^ rotr(a[4], 11) ^ rotr(a[4], 25));
     88         uint32_t ch = (a[4] & a[5]) ^ ((~a[4]) & a[6]);
     89         uint32_t temp1 = a[7] + s1 + ch + k[i] + w[i];
     90         uint32_t s0 = (rotr(a[0], 2) ^ rotr(a[0], 13) ^ rotr(a[0], 22));
     91         uint32_t maj = (a[0] & a[1]) ^ (a[0] & a[2]) ^ (a[1] & a[2]);
     92         uint32_t temp2 = s0 + maj;
     93         for (int i = 7; i > 0; i--) a[i] = a[i-1];
     94         a[4] += temp1;
     95         a[0] = temp1 + temp2;
     96     }
     97     for (int i = 0; i < 8; i++) hash[i] += a[i];
     98 }
     99 
    100 /*
    101  * CPU-side function that prints out the current hash.
    102  * Returns nothing but writes to stdout.
    103  */
    104 void sha_ctx::dump_hash() {
    105     uint32_t temp_hash[8];
    106     memcpy(temp_hash, hash, 8*sizeof(uint32_t));
    107     for (int i = 0; i < 8; i++) temp_hash[i] = bswap32(temp_hash[i]);
    108     auto u8_ptr = reinterpret_cast<uint8_t*>(temp_hash);
    109     for (int i = 0; i < 32; i++) printf("%02x", u8_ptr[i]);
    110 }
    111 
    112 int main(int argc, char** argv) {
    113     // Open file and initialize sha_ctx
    114     int fd = open(argv[1], O_RDONLY | O_NONBLOCK);
    115     // Advise kernel we're reading sequentially
    116     posix_fadvise(fd, 0, 0, POSIX_FADV_SEQUENTIAL);
    117     struct stat stat; 
    118     fstat(fd, &stat); // Get size of file
    119     sha_ctx sha(stat.st_size);
    120 
    121     // Allocate input and message schedule buffers on the FPU
    122     const size_t BUFFER_SIZE = 268435456; // 2 GiB
    123     uint8_t* buf; 
    124     cudaMallocManaged(&buf, BUFFER_SIZE);
    125     uint32_t* w;
    126     cudaMallocManaged(&w, BUFFER_SIZE*4);
    127 
    128     bool padded = false; // Store if we've padded the input yet (for handling edge cases)
    129     size_t bytes_read = read(fd, buf, BUFFER_SIZE);    
    130 
    131     do {
    132         // Maximum size_t (represented by underflow here) is read's return code for error
    133         if (bytes_read == (size_t)-1) { 
    134             printf("Error reading file.");
    135             exit(1);
    136         }
    137         if (!bytes_read) break; // If we're done reading, stop.
    138         else if (bytes_read < BUFFER_SIZE) {
    139             // Pad buffer in accordance with SHA256's standards
    140             size_t buffer_len = 64 * (((bytes_read + 9) / 64) + 1);
    141             for (size_t i = bytes_read; i < buffer_len; i++) buf[i] = 0;
    142             buf[bytes_read] = 0b10000000;
    143             for (int i = 1; i <= 8; i++) buf[buffer_len-i] = sha.len*8 >> (i-1)*8;
    144             process<<<dim3{8,8,1}, 64>>>(buf, w, 1024);
    145             cudaDeviceSynchronize();
    146             for (int i = 0; i < buffer_len/64; i++) sha.compress(w+(i*64));
    147             padded=true;
    148         } else {
    149             process<<<dim3{8,8,1}, 64>>>(buf, w, 1024);
    150             cudaDeviceSynchronize();
    151             for (int i = 0; i < BUFFER_SIZE/64; i++) sha.compress(w+(i*64));
    152         }
    153     } while ((bytes_read = read(fd, buf, BUFFER_SIZE)));
    154     if (padded == false) {
    155         // Pad buffer in accordance with SHA256
    156         for (size_t i = 0; i < 64; i++) buf[i] = 0;
    157         buf[0] = 0b10000000;
    158         for (int i = 1; i <= 8; i++) buf[64-i] = sha.len*8 >> (i-1)*8;
    159 
    160         // Run last chunk on CPU
    161         memcpy(w, buf, 64);
    162         for (int i = 0; i < 16; i++) w[i] = bswap32(w[i]);
    163         for (int i = 16; i < 64; i++) {
    164             uint32_t s0 = (rotr(w[i-15], 7) ^ rotr(w[i-15], 18) ^ (w[i-15] >> 3));
    165             uint32_t s1 = (rotr(w[i-2], 17) ^ rotr(w[i-2], 19)  ^ (w[i-2] >> 10));
    166             w[i] = w[i-16] + s0 + w[i-7] + s1;
    167         }
    168         sha.compress(w);
    169     }
    170     sha.dump_hash();
    171     printf("  %s\n", argv[1]);
    172 }