verity/build_verity_tree.cpp

   1 #include <openssl/evp.h>
   2 #include <sparse/sparse.h>
   3
   4 #undef NDEBUG
   5
   6 #include <assert.h>
   7 #include <errno.h>
   8 #include <getopt.h>
   9 #include <fcntl.h>
  10 #include <inttypes.h>
  11 #include <stdbool.h>
  12 #include <stdio.h>
  13 #include <stdlib.h>
  14 #include <string.h>
  15 #include <unistd.h>
  16
  17 struct sparse_hash_ctx {
  18     unsigned char *hashes;
  19     const unsigned char *salt;
  20     uint64_t salt_size;
  21     uint64_t hash_size;
  22     uint64_t block_size;
  23     const unsigned char *zero_block_hash;
  24     const EVP_MD *md;
  25 };
  26
  27 #define div_round_up(x,y) (((x) + (y) - 1)/(y))
  28
  29 #define round_up(x,y) (div_round_up(x,y)*(y))
  30
  31 #define FATAL(x...) { \
  32     fprintf(stderr, x); \
  33     exit(1); \
  34 }
  35
  36 size_t verity_tree_blocks(uint64_t data_size, size_t block_size, size_t hash_size,
  37                           int level)
  38 {
  39     size_t level_blocks = div_round_up(data_size, block_size);
  40     int hashes_per_block = div_round_up(block_size, hash_size);
  41
  42     do {
  43         level_blocks = div_round_up(level_blocks, hashes_per_block);
  44     } while (level--);
  45
  46     return level_blocks;
  47 }
  48
  49 int hash_block(const EVP_MD *md,
  50                const unsigned char *block, size_t len,
  51                const unsigned char *salt, size_t salt_len,
  52                unsigned char *out, size_t *out_size)
  53 {
  54     EVP_MD_CTX *mdctx;
  55     unsigned int s;
  56     int ret = 1;
  57
  58     mdctx = EVP_MD_CTX_create();
  59     assert(mdctx);
  60     ret &= EVP_DigestInit_ex(mdctx, md, NULL);
  61     ret &= EVP_DigestUpdate(mdctx, salt, salt_len);
  62     ret &= EVP_DigestUpdate(mdctx, block, len);
  63     ret &= EVP_DigestFinal_ex(mdctx, out, &s);
  64     EVP_MD_CTX_destroy(mdctx);
  65     assert(ret == 1);
  66     if (out_size) {
  67         *out_size = s;
  68     }
  69     return 0;
  70 }
  71
  72 int hash_blocks(const EVP_MD *md,
  73                 const unsigned char *in, size_t in_size,
  74                 unsigned char *out, size_t *out_size,
  75                 const unsigned char *salt, size_t salt_size,
  76                 size_t block_size)
  77 {
  78     size_t s;
  79     *out_size = 0;
  80     for (size_t i = 0; i < in_size; i += block_size) {
  81         hash_block(md, in + i, block_size, salt, salt_size, out, &s);
  82         out += s;
  83         *out_size += s;
  84     }
  85
  86     return 0;
  87 }
  88
  89 int hash_chunk(void *priv, const void *data, int len)
  90 {
  91     struct sparse_hash_ctx *ctx = (struct sparse_hash_ctx *)priv;
  92     assert(len % ctx->block_size == 0);
  93     if (data) {
  94         size_t s;
  95         hash_blocks(ctx->md, (const unsigned char *)data, len,
  96                     ctx->hashes, &s,
  97                     ctx->salt, ctx->salt_size, ctx->block_size);
  98         ctx->hashes += s;
  99     } else {
 100         for (size_t i = 0; i < (size_t)len; i += ctx->block_size) {
 101             memcpy(ctx->hashes, ctx->zero_block_hash, ctx->hash_size);
 102             ctx->hashes += ctx->hash_size;
 103         }
 104     }
 105     return 0;
 106 }
 107
 108 void usage(void)
 109 {
 110     printf("usage: build_verity_tree [ <options> ] -s <size> | <data> <verity>\n"
 111            "options:\n"
 112            "  -a,--salt-str=<string>       set salt to <string>\n"
 113            "  -A,--salt-hex=<hex digits>   set salt to <hex digits>\n"
 114            "  -h                           show this help\n"
 115            "  -s,--verity-size=<data size> print the size of the verity tree\n"
 116            "  -S                           treat <data image> as a sparse file\n"
 117         );
 118 }
 119
 120 int main(int argc, char **argv)
 121 {
 122     char *data_filename;
 123     char *verity_filename;
 124     unsigned char *salt = NULL;
 125     size_t salt_size = 0;
 126     bool sparse = false;
 127     size_t block_size = 4096;
 128     size_t calculate_size = 0;
 129
 130     while (1) {
 131         const static struct option long_options[] = {
 132             {"salt-str", required_argument, 0, 'a'},
 133             {"salt-hex", required_argument, 0, 'A'},
 134             {"help", no_argument, 0, 'h'},
 135             {"sparse", no_argument, 0, 'S'},
 136             {"verity-size", required_argument, 0, 's'},
 137         };
 138         int c = getopt_long(argc, argv, "a:A:hSs:", long_options, NULL);
 139         if (c < 0) {
 140             break;
 141         }
 142
 143         switch (c) {
 144         case 'a':
 145             salt_size = strlen(optarg);
 146             salt = new unsigned char[salt_size]();
 147             if (salt == NULL) {
 148                 FATAL("failed to allocate memory for salt\n");
 149             }
 150             memcpy(salt, optarg, salt_size);
 151             break;
 152         case 'A': {
 153                 BIGNUM *bn = NULL;
 154                 if(!BN_hex2bn(&bn, optarg)) {
 155                     FATAL("failed to convert salt from hex\n");
 156                 }
 157                 salt_size = BN_num_bytes(bn);
 158                 salt = new unsigned char[salt_size]();
 159                 if (salt == NULL) {
 160                     FATAL("failed to allocate memory for salt\n");
 161                 }
 162                 if((size_t)BN_bn2bin(bn, salt) != salt_size) {
 163                     FATAL("failed to convert salt to bytes\n");
 164                 }
 165             }
 166             break;
 167         case 'h':
 168             usage();
 169             return 1;
 170         case 'S':
 171             sparse = true;
 172             break;
 173         case 's':
 174             calculate_size = strtoul(optarg, NULL, 0);
 175             break;
 176         case '?':
 177             usage();
 178             return 1;
 179         default:
 180             abort();
 181         }
 182     }
 183
 184     argc -= optind;
 185     argv += optind;
 186
 187     const EVP_MD *md = EVP_sha256();
 188     if (!md) {
 189         FATAL("failed to get digest\n");
 190     }
 191
 192     size_t hash_size = EVP_MD_size(md);
 193     assert(hash_size * 2 < block_size);
 194
 195     if (!salt || !salt_size) {
 196         salt_size = hash_size;
 197         salt = new unsigned char[salt_size];
 198         if (salt == NULL) {
 199             FATAL("failed to allocate memory for salt\n");
 200         }
 201
 202         int random_fd = open("/dev/urandom", O_RDONLY);
 203         if (random_fd < 0) {
 204             FATAL("failed to open /dev/urandom\n");
 205         }
 206
 207         ssize_t ret = read(random_fd, salt, salt_size);
 208         if (ret != (ssize_t)salt_size) {
 209             FATAL("failed to read %zu bytes from /dev/urandom: %zd %d\n", salt_size, ret, errno);
 210         }
 211         close(random_fd);
 212     }
 213
 214     if (calculate_size) {
 215         if (argc != 0) {
 216             usage();
 217             return 1;
 218         }
 219         size_t verity_blocks = 0;
 220         size_t level_blocks;
 221         int levels = 0;
 222         do {
 223             level_blocks = verity_tree_blocks(calculate_size, block_size, hash_size, levels);
 224             levels++;
 225             verity_blocks += level_blocks;
 226         } while (level_blocks > 1);
 227
 228         printf("%zu\n", verity_blocks * block_size);
 229         return 0;
 230     }
 231
 232     if (argc != 2) {
 233         usage();
 234         return 1;
 235     }
 236
 237     data_filename = argv[0];
 238     verity_filename = argv[1];
 239
 240     int fd = open(data_filename, O_RDONLY);
 241     if (fd < 0) {
 242         FATAL("failed to open %s\n", data_filename);
 243     }
 244
 245     struct sparse_file *file;
 246     if (sparse) {
 247         file = sparse_file_import(fd, false, false);
 248     } else {
 249         file = sparse_file_import_auto(fd, false);
 250     }
 251
 252     if (!file) {
 253         FATAL("failed to read file %s\n", data_filename);
 254     }
 255
 256     int64_t len = sparse_file_len(file, false, false);
 257     if (len % block_size != 0) {
 258         FATAL("file size %" PRIu64 " is not a multiple of %zu bytes\n",
 259                 len, block_size);
 260     }
 261
 262     int levels = 0;
 263     size_t verity_blocks = 0;
 264     size_t level_blocks;
 265
 266     do {
 267         level_blocks = verity_tree_blocks(len, block_size, hash_size, levels);
 268         levels++;
 269         verity_blocks += level_blocks;
 270     } while (level_blocks > 1);
 271
 272     unsigned char *verity_tree = new unsigned char[verity_blocks * block_size]();
 273     unsigned char **verity_tree_levels = new unsigned char *[levels + 1]();
 274     size_t *verity_tree_level_blocks = new size_t[levels]();
 275     if (verity_tree == NULL || verity_tree_levels == NULL || verity_tree_level_blocks == NULL) {
 276         FATAL("failed to allocate memory for verity tree\n");
 277     }
 278
 279     unsigned char *ptr = verity_tree;
 280     for (int i = levels - 1; i >= 0; i--) {
 281         verity_tree_levels[i] = ptr;
 282         verity_tree_level_blocks[i] = verity_tree_blocks(len, block_size, hash_size, i);
 283         ptr += verity_tree_level_blocks[i] * block_size;
 284     }
 285     assert(ptr == verity_tree + verity_blocks * block_size);
 286     assert(verity_tree_level_blocks[levels - 1] == 1);
 287
 288     unsigned char zero_block_hash[hash_size];
 289     unsigned char zero_block[block_size];
 290     memset(zero_block, 0, block_size);
 291     hash_block(md, zero_block, block_size, salt, salt_size, zero_block_hash, NULL);
 292
 293     unsigned char root_hash[hash_size];
 294     verity_tree_levels[levels] = root_hash;
 295
 296     struct sparse_hash_ctx ctx;
 297     ctx.hashes = verity_tree_levels[0];
 298     ctx.salt = salt;
 299     ctx.salt_size = salt_size;
 300     ctx.hash_size = hash_size;
 301     ctx.block_size = block_size;
 302     ctx.zero_block_hash = zero_block_hash;
 303     ctx.md = md;
 304
 305     sparse_file_callback(file, false, false, hash_chunk, &ctx);
 306
 307     sparse_file_destroy(file);
 308     close(fd);
 309
 310     for (int i = 0; i < levels; i++) {
 311         size_t out_size;
 312         hash_blocks(md,
 313                 verity_tree_levels[i], verity_tree_level_blocks[i] * block_size,
 314                 verity_tree_levels[i + 1], &out_size,
 315                 salt, salt_size, block_size);
 316           if (i < levels - 1) {
 317               assert(div_round_up(out_size, block_size) == verity_tree_level_blocks[i + 1]);
 318           } else {
 319               assert(out_size == hash_size);
 320           }
 321     }
 322
 323     for (size_t i = 0; i < hash_size; i++) {
 324         printf("%02x", root_hash[i]);
 325     }
 326     printf(" ");
 327     for (size_t i = 0; i < salt_size; i++) {
 328         printf("%02x", salt[i]);
 329     }
 330     printf("\n");
 331
 332     fd = open(verity_filename, O_WRONLY|O_CREAT, 0666);
 333     if (fd < 0) {
 334         FATAL("failed to open output file '%s'\n", verity_filename);
 335     }
 336     write(fd, verity_tree, verity_blocks * block_size);
 337     close(fd);
 338
 339     delete[] verity_tree_levels;
 340     delete[] verity_tree_level_blocks;
 341     delete[] verity_tree;
 342     delete[] salt;
 343 }