reftable/block.c

   1 /*
   2 Copyright 2020 Google LLC
   3
   4 Use of this source code is governed by a BSD-style
   5 license that can be found in the LICENSE file or at
   6 https://developers.google.com/open-source/licenses/bsd
   7 */
   8
   9 #include "block.h"
  10
  11 #include "blocksource.h"
  12 #include "constants.h"
  13 #include "record.h"
  14 #include "reftable-error.h"
  15 #include "system.h"
  16 #include <zlib.h>
  17
  18 int header_size(int version)
  19 {
  20         switch (version) {
  21         case 1:
  22                 return 24;
  23         case 2:
  24                 return 28;
  25         }
  26         abort();
  27 }
  28
  29 int footer_size(int version)
  30 {
  31         switch (version) {
  32         case 1:
  33                 return 68;
  34         case 2:
  35                 return 72;
  36         }
  37         abort();
  38 }
  39
  40 static int block_writer_register_restart(struct block_writer *w, int n,
  41                                          int is_restart, struct strbuf *key)
  42 {
  43         int rlen = w->restart_len;
  44         if (rlen >= MAX_RESTARTS) {
  45                 is_restart = 0;
  46         }
  47
  48         if (is_restart) {
  49                 rlen++;
  50         }
  51         if (2 + 3 * rlen + n > w->block_size - w->next)
  52                 return -1;
  53         if (is_restart) {
  54                 if (w->restart_len == w->restart_cap) {
  55                         w->restart_cap = w->restart_cap * 2 + 1;
  56                         w->restarts = reftable_realloc(
  57                                 w->restarts, sizeof(uint32_t) * w->restart_cap);
  58                 }
  59
  60                 w->restarts[w->restart_len++] = w->next;
  61         }
  62
  63         w->next += n;
  64
  65         strbuf_reset(&w->last_key);
  66         strbuf_addbuf(&w->last_key, key);
  67         w->entries++;
  68         return 0;
  69 }
  70
  71 void block_writer_init(struct block_writer *bw, uint8_t typ, uint8_t *buf,
  72                        uint32_t block_size, uint32_t header_off, int hash_size)
  73 {
  74         bw->buf = buf;
  75         bw->hash_size = hash_size;
  76         bw->block_size = block_size;
  77         bw->header_off = header_off;
  78         bw->buf[header_off] = typ;
  79         bw->next = header_off + 4;
  80         bw->restart_interval = 16;
  81         bw->entries = 0;
  82         bw->restart_len = 0;
  83         bw->last_key.len = 0;
  84 }
  85
  86 uint8_t block_writer_type(struct block_writer *bw)
  87 {
  88         return bw->buf[bw->header_off];
  89 }
  90
  91 /* adds the reftable_record to the block. Returns -1 if it does not fit, 0 on
  92    success */
  93 int block_writer_add(struct block_writer *w, struct reftable_record *rec)
  94 {
  95         struct strbuf empty = STRBUF_INIT;
  96         struct strbuf last =
  97                 w->entries % w->restart_interval == 0 ? empty : w->last_key;
  98         struct string_view out = {
  99                 .buf = w->buf + w->next,
 100                 .len = w->block_size - w->next,
 101         };
 102
 103         struct string_view start = out;
 104
 105         int is_restart = 0;
 106         struct strbuf key = STRBUF_INIT;
 107         int n = 0;
 108
 109         reftable_record_key(rec, &key);
 110         n = reftable_encode_key(&is_restart, out, last, key,
 111                                 reftable_record_val_type(rec));
 112         if (n < 0)
 113                 goto done;
 114         string_view_consume(&out, n);
 115
 116         n = reftable_record_encode(rec, out, w->hash_size);
 117         if (n < 0)
 118                 goto done;
 119         string_view_consume(&out, n);
 120
 121         if (block_writer_register_restart(w, start.len - out.len, is_restart,
 122                                           &key) < 0)
 123                 goto done;
 124
 125         strbuf_release(&key);
 126         return 0;
 127
 128 done:
 129         strbuf_release(&key);
 130         return -1;
 131 }
 132
 133 int block_writer_finish(struct block_writer *w)
 134 {
 135         int i;
 136         for (i = 0; i < w->restart_len; i++) {
 137                 put_be24(w->buf + w->next, w->restarts[i]);
 138                 w->next += 3;
 139         }
 140
 141         put_be16(w->buf + w->next, w->restart_len);
 142         w->next += 2;
 143         put_be24(w->buf + 1 + w->header_off, w->next);
 144
 145         if (block_writer_type(w) == BLOCK_TYPE_LOG) {
 146                 int block_header_skip = 4 + w->header_off;
 147                 uLongf src_len = w->next - block_header_skip;
 148                 uLongf dest_cap = src_len * 1.001 + 12;
 149
 150                 uint8_t *compressed = reftable_malloc(dest_cap);
 151                 while (1) {
 152                         uLongf out_dest_len = dest_cap;
 153                         int zresult = compress2(compressed, &out_dest_len,
 154                                                 w->buf + block_header_skip,
 155                                                 src_len, 9);
 156                         if (zresult == Z_BUF_ERROR && dest_cap < LONG_MAX) {
 157                                 dest_cap *= 2;
 158                                 compressed =
 159                                         reftable_realloc(compressed, dest_cap);
 160                                 if (compressed)
 161                                         continue;
 162                         }
 163
 164                         if (Z_OK != zresult) {
 165                                 reftable_free(compressed);
 166                                 return REFTABLE_ZLIB_ERROR;
 167                         }
 168
 169                         memcpy(w->buf + block_header_skip, compressed,
 170                                out_dest_len);
 171                         w->next = out_dest_len + block_header_skip;
 172                         reftable_free(compressed);
 173                         break;
 174                 }
 175         }
 176         return w->next;
 177 }
 178
 179 uint8_t block_reader_type(struct block_reader *r)
 180 {
 181         return r->block.data[r->header_off];
 182 }
 183
 184 int block_reader_init(struct block_reader *br, struct reftable_block *block,
 185                       uint32_t header_off, uint32_t table_block_size,
 186                       int hash_size)
 187 {
 188         uint32_t full_block_size = table_block_size;
 189         uint8_t typ = block->data[header_off];
 190         uint32_t sz = get_be24(block->data + header_off + 1);
 191
 192         uint16_t restart_count = 0;
 193         uint32_t restart_start = 0;
 194         uint8_t *restart_bytes = NULL;
 195
 196         if (!reftable_is_block_type(typ))
 197                 return REFTABLE_FORMAT_ERROR;
 198
 199         if (typ == BLOCK_TYPE_LOG) {
 200                 int block_header_skip = 4 + header_off;
 201                 uLongf dst_len = sz - block_header_skip; /* total size of dest
 202                                                             buffer. */
 203                 uLongf src_len = block->len - block_header_skip;
 204                 /* Log blocks specify the *uncompressed* size in their header.
 205                  */
 206                 uint8_t *uncompressed = reftable_malloc(sz);
 207
 208                 /* Copy over the block header verbatim. It's not compressed. */
 209                 memcpy(uncompressed, block->data, block_header_skip);
 210
 211                 /* Uncompress */
 212                 if (Z_OK !=
 213                     uncompress2(uncompressed + block_header_skip, &dst_len,
 214                                 block->data + block_header_skip, &src_len)) {
 215                         reftable_free(uncompressed);
 216                         return REFTABLE_ZLIB_ERROR;
 217                 }
 218
 219                 if (dst_len + block_header_skip != sz)
 220                         return REFTABLE_FORMAT_ERROR;
 221
 222                 /* We're done with the input data. */
 223                 reftable_block_done(block);
 224                 block->data = uncompressed;
 225                 block->len = sz;
 226                 block->source = malloc_block_source();
 227                 full_block_size = src_len + block_header_skip;
 228         } else if (full_block_size == 0) {
 229                 full_block_size = sz;
 230         } else if (sz < full_block_size && sz < block->len &&
 231                    block->data[sz] != 0) {
 232                 /* If the block is smaller than the full block size, it is
 233                    padded (data followed by '\0') or the next block is
 234                    unaligned. */
 235                 full_block_size = sz;
 236         }
 237
 238         restart_count = get_be16(block->data + sz - 2);
 239         restart_start = sz - 2 - 3 * restart_count;
 240         restart_bytes = block->data + restart_start;
 241
 242         /* transfer ownership. */
 243         br->block = *block;
 244         block->data = NULL;
 245         block->len = 0;
 246
 247         br->hash_size = hash_size;
 248         br->block_len = restart_start;
 249         br->full_block_size = full_block_size;
 250         br->header_off = header_off;
 251         br->restart_count = restart_count;
 252         br->restart_bytes = restart_bytes;
 253
 254         return 0;
 255 }
 256
 257 static uint32_t block_reader_restart_offset(struct block_reader *br, int i)
 258 {
 259         return get_be24(br->restart_bytes + 3 * i);
 260 }
 261
 262 void block_reader_start(struct block_reader *br, struct block_iter *it)
 263 {
 264         it->br = br;
 265         strbuf_reset(&it->last_key);
 266         it->next_off = br->header_off + 4;
 267 }
 268
 269 struct restart_find_args {
 270         int error;
 271         struct strbuf key;
 272         struct block_reader *r;
 273 };
 274
 275 static int restart_key_less(size_t idx, void *args)
 276 {
 277         struct restart_find_args *a = args;
 278         uint32_t off = block_reader_restart_offset(a->r, idx);
 279         struct string_view in = {
 280                 .buf = a->r->block.data + off,
 281                 .len = a->r->block_len - off,
 282         };
 283
 284         /* the restart key is verbatim in the block, so this could avoid the
 285            alloc for decoding the key */
 286         struct strbuf rkey = STRBUF_INIT;
 287         struct strbuf last_key = STRBUF_INIT;
 288         uint8_t unused_extra;
 289         int n = reftable_decode_key(&rkey, &unused_extra, last_key, in);
 290         int result;
 291         if (n < 0) {
 292                 a->error = 1;
 293                 return -1;
 294         }
 295
 296         result = strbuf_cmp(&a->key, &rkey);
 297         strbuf_release(&rkey);
 298         return result;
 299 }
 300
 301 void block_iter_copy_from(struct block_iter *dest, struct block_iter *src)
 302 {
 303         dest->br = src->br;
 304         dest->next_off = src->next_off;
 305         strbuf_reset(&dest->last_key);
 306         strbuf_addbuf(&dest->last_key, &src->last_key);
 307 }
 308
 309 int block_iter_next(struct block_iter *it, struct reftable_record *rec)
 310 {
 311         struct string_view in = {
 312                 .buf = it->br->block.data + it->next_off,
 313                 .len = it->br->block_len - it->next_off,
 314         };
 315         struct string_view start = in;
 316         struct strbuf key = STRBUF_INIT;
 317         uint8_t extra = 0;
 318         int n = 0;
 319
 320         if (it->next_off >= it->br->block_len)
 321                 return 1;
 322
 323         n = reftable_decode_key(&key, &extra, it->last_key, in);
 324         if (n < 0)
 325                 return -1;
 326
 327         string_view_consume(&in, n);
 328         n = reftable_record_decode(rec, key, extra, in, it->br->hash_size);
 329         if (n < 0)
 330                 return -1;
 331         string_view_consume(&in, n);
 332
 333         strbuf_reset(&it->last_key);
 334         strbuf_addbuf(&it->last_key, &key);
 335         it->next_off += start.len - in.len;
 336         strbuf_release(&key);
 337         return 0;
 338 }
 339
 340 int block_reader_first_key(struct block_reader *br, struct strbuf *key)
 341 {
 342         struct strbuf empty = STRBUF_INIT;
 343         int off = br->header_off + 4;
 344         struct string_view in = {
 345                 .buf = br->block.data + off,
 346                 .len = br->block_len - off,
 347         };
 348
 349         uint8_t extra = 0;
 350         int n = reftable_decode_key(key, &extra, empty, in);
 351         if (n < 0)
 352                 return n;
 353
 354         return 0;
 355 }
 356
 357 int block_iter_seek(struct block_iter *it, struct strbuf *want)
 358 {
 359         return block_reader_seek(it->br, it, want);
 360 }
 361
 362 void block_iter_close(struct block_iter *it)
 363 {
 364         strbuf_release(&it->last_key);
 365 }
 366
 367 int block_reader_seek(struct block_reader *br, struct block_iter *it,
 368                       struct strbuf *want)
 369 {
 370         struct restart_find_args args = {
 371                 .key = *want,
 372                 .r = br,
 373         };
 374         struct reftable_record rec = reftable_new_record(block_reader_type(br));
 375         struct strbuf key = STRBUF_INIT;
 376         int err = 0;
 377         struct block_iter next = {
 378                 .last_key = STRBUF_INIT,
 379         };
 380
 381         int i = binsearch(br->restart_count, &restart_key_less, &args);
 382         if (args.error) {
 383                 err = REFTABLE_FORMAT_ERROR;
 384                 goto done;
 385         }
 386
 387         it->br = br;
 388         if (i > 0) {
 389                 i--;
 390                 it->next_off = block_reader_restart_offset(br, i);
 391         } else {
 392                 it->next_off = br->header_off + 4;
 393         }
 394
 395         /* We're looking for the last entry less/equal than the wanted key, so
 396            we have to go one entry too far and then back up.
 397         */
 398         while (1) {
 399                 block_iter_copy_from(&next, it);
 400                 err = block_iter_next(&next, &rec);
 401                 if (err < 0)
 402                         goto done;
 403
 404                 reftable_record_key(&rec, &key);
 405                 if (err > 0 || strbuf_cmp(&key, want) >= 0) {
 406                         err = 0;
 407                         goto done;
 408                 }
 409
 410                 block_iter_copy_from(it, &next);
 411         }
 412
 413 done:
 414         strbuf_release(&key);
 415         strbuf_release(&next.last_key);
 416         reftable_record_destroy(&rec);
 417
 418         return err;
 419 }
 420
 421 void block_writer_release(struct block_writer *bw)
 422 {
 423         FREE_AND_NULL(bw->restarts);
 424         strbuf_release(&bw->last_key);
 425         /* the block is not owned. */
 426 }
 427
 428 void reftable_block_done(struct reftable_block *blockp)
 429 {
 430         struct reftable_block_source source = blockp->source;
 431         if (blockp && source.ops)
 432                 source.ops->return_block(source.arg, blockp);
 433         blockp->data = NULL;
 434         blockp->len = 0;
 435         blockp->source.ops = NULL;
 436         blockp->source.arg = NULL;
 437 }