btrfs-image.c

   1 /*
   2  * Copyright (C) 2008 Oracle.  All rights reserved.
   3  *
   4  * This program is free software; you can redistribute it and/or
   5  * modify it under the terms of the GNU General Public
   6  * License v2 as published by the Free Software Foundation.
   7  *
   8  * This program is distributed in the hope that it will be useful,
   9  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  10  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  11  * General Public License for more details.
  12  *
  13  * You should have received a copy of the GNU General Public
  14  * License along with this program; if not, write to the
  15  * Free Software Foundation, Inc., 59 Temple Place - Suite 330,
  16  * Boston, MA 021110-1307, USA.
  17  */
  18
  19 #define _XOPEN_SOURCE 500
  20 #define _GNU_SOURCE 1
  21 #include <pthread.h>
  22 #include <stdio.h>
  23 #include <stdlib.h>
  24 #include <sys/types.h>
  25 #include <sys/stat.h>
  26 #include <fcntl.h>
  27 #include <unistd.h>
  28 #include <dirent.h>
  29 #include <zlib.h>
  30 #include "kerncompat.h"
  31 #include "crc32c.h"
  32 #include "ctree.h"
  33 #include "disk-io.h"
  34 #include "transaction.h"
  35 #include "utils.h"
  36 #include "version.h"
  37 #include "volumes.h"
  38 #include "extent_io.h"
  39
  40 #define HEADER_MAGIC            0xbd5c25e27295668bULL
  41 #define MAX_PENDING_SIZE        (256 * 1024)
  42 #define BLOCK_SIZE              1024
  43 #define BLOCK_MASK              (BLOCK_SIZE - 1)
  44
  45 #define COMPRESS_NONE           0
  46 #define COMPRESS_ZLIB           1
  47
  48 struct meta_cluster_item {
  49         __le64 bytenr;
  50         __le32 size;
  51 } __attribute__ ((__packed__));
  52
  53 struct meta_cluster_header {
  54         __le64 magic;
  55         __le64 bytenr;
  56         __le32 nritems;
  57         u8 compress;
  58 } __attribute__ ((__packed__));
  59
  60 /* cluster header + index items + buffers */
  61 struct meta_cluster {
  62         struct meta_cluster_header header;
  63         struct meta_cluster_item items[];
  64 } __attribute__ ((__packed__));
  65
  66 #define ITEMS_PER_CLUSTER ((BLOCK_SIZE - sizeof(struct meta_cluster)) / \
  67                            sizeof(struct meta_cluster_item))
  68
  69 struct fs_chunk {
  70         u64 logical;
  71         u64 physical;
  72         u64 bytes;
  73         struct rb_node n;
  74 };
  75
  76 struct async_work {
  77         struct list_head list;
  78         struct list_head ordered;
  79         u64 start;
  80         u64 size;
  81         u8 *buffer;
  82         size_t bufsize;
  83         int error;
  84 };
  85
  86 struct metadump_struct {
  87         struct btrfs_root *root;
  88         FILE *out;
  89
  90         struct meta_cluster *cluster;
  91
  92         pthread_t *threads;
  93         size_t num_threads;
  94         pthread_mutex_t mutex;
  95         pthread_cond_t cond;
  96         struct rb_root name_tree;
  97
  98         struct list_head list;
  99         struct list_head ordered;
 100         size_t num_items;
 101         size_t num_ready;
 102
 103         u64 pending_start;
 104         u64 pending_size;
 105
 106         int compress_level;
 107         int done;
 108         int data;
 109         int sanitize_names;
 110 };
 111
 112 struct name {
 113         struct rb_node n;
 114         char *val;
 115         char *sub;
 116         u32 len;
 117 };
 118
 119 struct mdrestore_struct {
 120         FILE *in;
 121         FILE *out;
 122
 123         pthread_t *threads;
 124         size_t num_threads;
 125         pthread_mutex_t mutex;
 126         pthread_cond_t cond;
 127
 128         struct rb_root chunk_tree;
 129         struct list_head list;
 130         size_t num_items;
 131         u64 leafsize;
 132         u64 devid;
 133         u8 uuid[BTRFS_UUID_SIZE];
 134         u8 fsid[BTRFS_FSID_SIZE];
 135
 136         int compress_method;
 137         int done;
 138         int error;
 139         int old_restore;
 140         int fixup_offset;
 141         int multi_devices;
 142         struct btrfs_fs_info *info;
 143 };
 144
 145 static void print_usage(void) __attribute__((noreturn));
 146 static int search_for_chunk_blocks(struct mdrestore_struct *mdres,
 147                                    u64 search, u64 cluster_bytenr);
 148 static struct extent_buffer *alloc_dummy_eb(u64 bytenr, u32 size);
 149
 150 static void csum_block(u8 *buf, size_t len)
 151 {
 152         char result[BTRFS_CRC32_SIZE];
 153         u32 crc = ~(u32)0;
 154         crc = crc32c(crc, buf + BTRFS_CSUM_SIZE, len - BTRFS_CSUM_SIZE);
 155         btrfs_csum_final(crc, result);
 156         memcpy(buf, result, BTRFS_CRC32_SIZE);
 157 }
 158
 159 static int has_name(struct btrfs_key *key)
 160 {
 161         switch (key->type) {
 162         case BTRFS_DIR_ITEM_KEY:
 163         case BTRFS_DIR_INDEX_KEY:
 164         case BTRFS_INODE_REF_KEY:
 165         case BTRFS_INODE_EXTREF_KEY:
 166         case BTRFS_XATTR_ITEM_KEY:
 167                 return 1;
 168         default:
 169                 break;
 170         }
 171
 172         return 0;
 173 }
 174
 175 static char *generate_garbage(u32 name_len)
 176 {
 177         char *buf = malloc(name_len);
 178         int i;
 179
 180         if (!buf)
 181                 return NULL;
 182
 183         for (i = 0; i < name_len; i++) {
 184                 char c = rand() % 94 + 33;
 185
 186                 if (c == '/')
 187                         c++;
 188                 buf[i] = c;
 189         }
 190
 191         return buf;
 192 }
 193
 194 static int name_cmp(struct rb_node *a, struct rb_node *b, int fuzz)
 195 {
 196         struct name *entry = rb_entry(a, struct name, n);
 197         struct name *ins = rb_entry(b, struct name, n);
 198         u32 len;
 199
 200         len = min(ins->len, entry->len);
 201         return memcmp(ins->val, entry->val, len);
 202 }
 203
 204 static int chunk_cmp(struct rb_node *a, struct rb_node *b, int fuzz)
 205 {
 206         struct fs_chunk *entry = rb_entry(a, struct fs_chunk, n);
 207         struct fs_chunk *ins = rb_entry(b, struct fs_chunk, n);
 208
 209         if (fuzz && ins->logical >= entry->logical &&
 210             ins->logical < entry->logical + entry->bytes)
 211                 return 0;
 212
 213         if (ins->logical < entry->logical)
 214                 return -1;
 215         else if (ins->logical > entry->logical)
 216                 return 1;
 217         return 0;
 218 }
 219
 220 static void tree_insert(struct rb_root *root, struct rb_node *ins,
 221                         int (*cmp)(struct rb_node *a, struct rb_node *b,
 222                                    int fuzz))
 223 {
 224         struct rb_node ** p = &root->rb_node;
 225         struct rb_node * parent = NULL;
 226         int dir;
 227
 228         while(*p) {
 229                 parent = *p;
 230
 231                 dir = cmp(*p, ins, 0);
 232                 if (dir < 0)
 233                         p = &(*p)->rb_left;
 234                 else if (dir > 0)
 235                         p = &(*p)->rb_right;
 236                 else
 237                         BUG();
 238         }
 239
 240         rb_link_node(ins, parent, p);
 241         rb_insert_color(ins, root);
 242 }
 243
 244 static struct rb_node *tree_search(struct rb_root *root,
 245                                    struct rb_node *search,
 246                                    int (*cmp)(struct rb_node *a,
 247                                               struct rb_node *b, int fuzz),
 248                                    int fuzz)
 249 {
 250         struct rb_node *n = root->rb_node;
 251         int dir;
 252
 253         while (n) {
 254                 dir = cmp(n, search, fuzz);
 255                 if (dir < 0)
 256                         n = n->rb_left;
 257                 else if (dir > 0)
 258                         n = n->rb_right;
 259                 else
 260                         return n;
 261         }
 262
 263         return NULL;
 264 }
 265
 266 static char *find_collision(struct metadump_struct *md, char *name,
 267                             u32 name_len)
 268 {
 269         struct name *val;
 270         struct rb_node *entry;
 271         struct name tmp;
 272         unsigned long checksum;
 273         int found = 0;
 274         int i;
 275
 276         tmp.val = name;
 277         tmp.len = name_len;
 278         entry = tree_search(&md->name_tree, &tmp.n, name_cmp, 0);
 279         if (entry) {
 280                 val = rb_entry(entry, struct name, n);
 281                 free(name);
 282                 return val->sub;
 283         }
 284
 285         val = malloc(sizeof(struct name));
 286         if (!val) {
 287                 fprintf(stderr, "Couldn't sanitize name, enomem\n");
 288                 free(name);
 289                 return NULL;
 290         }
 291
 292         memset(val, 0, sizeof(*val));
 293
 294         val->val = name;
 295         val->len = name_len;
 296         val->sub = malloc(name_len);
 297         if (!val->sub) {
 298                 fprintf(stderr, "Couldn't sanitize name, enomem\n");
 299                 free(val);
 300                 free(name);
 301                 return NULL;
 302         }
 303
 304         checksum = crc32c(~1, val->val, name_len);
 305         memset(val->sub, ' ', name_len);
 306         i = 0;
 307         while (1) {
 308                 if (crc32c(~1, val->sub, name_len) == checksum &&
 309                     memcmp(val->sub, val->val, val->len)) {
 310                         found = 1;
 311                         break;
 312                 }
 313
 314                 if (val->sub[i] == 127) {
 315                         do {
 316                                 i++;
 317                                 if (i >= name_len)
 318                                         break;
 319                         } while (val->sub[i] == 127);
 320
 321                         if (i >= name_len)
 322                                 break;
 323                         val->sub[i]++;
 324                         if (val->sub[i] == '/')
 325                                 val->sub[i]++;
 326                         memset(val->sub, ' ', i);
 327                         i = 0;
 328                         continue;
 329                 } else {
 330                         val->sub[i]++;
 331                         if (val->sub[i] == '/')
 332                                 val->sub[i]++;
 333                 }
 334         }
 335
 336         if (!found) {
 337                 fprintf(stderr, "Couldn't find a collision for '%.*s', "
 338                         "generating normal garbage, it won't match indexes\n",
 339                         val->len, val->val);
 340                 for (i = 0; i < name_len; i++) {
 341                         char c = rand() % 94 + 33;
 342
 343                         if (c == '/')
 344                                 c++;
 345                         val->sub[i] = c;
 346                 }
 347         }
 348
 349         tree_insert(&md->name_tree, &val->n, name_cmp);
 350         return val->sub;
 351 }
 352
 353 static void sanitize_dir_item(struct metadump_struct *md, struct extent_buffer *eb,
 354                               int slot)
 355 {
 356         struct btrfs_dir_item *dir_item;
 357         char *buf;
 358         char *garbage;
 359         unsigned long name_ptr;
 360         u32 total_len;
 361         u32 cur = 0;
 362         u32 this_len;
 363         u32 name_len;
 364         int free_garbage = (md->sanitize_names == 1);
 365
 366         dir_item = btrfs_item_ptr(eb, slot, struct btrfs_dir_item);
 367         total_len = btrfs_item_size_nr(eb, slot);
 368         while (cur < total_len) {
 369                 this_len = sizeof(*dir_item) +
 370                         btrfs_dir_name_len(eb, dir_item) +
 371                         btrfs_dir_data_len(eb, dir_item);
 372                 name_ptr = (unsigned long)(dir_item + 1);
 373                 name_len = btrfs_dir_name_len(eb, dir_item);
 374
 375                 if (md->sanitize_names > 1) {
 376                         buf = malloc(name_len);
 377                         if (!buf) {
 378                                 fprintf(stderr, "Couldn't sanitize name, "
 379                                         "enomem\n");
 380                                 return;
 381                         }
 382                         read_extent_buffer(eb, buf, name_ptr, name_len);
 383                         garbage = find_collision(md, buf, name_len);
 384                 } else {
 385                         garbage = generate_garbage(name_len);
 386                 }
 387                 if (!garbage) {
 388                         fprintf(stderr, "Couldn't sanitize name, enomem\n");
 389                         return;
 390                 }
 391                 write_extent_buffer(eb, garbage, name_ptr, name_len);
 392                 cur += this_len;
 393                 dir_item = (struct btrfs_dir_item *)((char *)dir_item +
 394                                                      this_len);
 395                 if (free_garbage)
 396                         free(garbage);
 397         }
 398 }
 399
 400 static void sanitize_inode_ref(struct metadump_struct *md,
 401                                struct extent_buffer *eb, int slot, int ext)
 402 {
 403         struct btrfs_inode_extref *extref;
 404         struct btrfs_inode_ref *ref;
 405         char *garbage, *buf;
 406         unsigned long ptr;
 407         unsigned long name_ptr;
 408         u32 item_size;
 409         u32 cur_offset = 0;
 410         int len;
 411         int free_garbage = (md->sanitize_names == 1);
 412
 413         item_size = btrfs_item_size_nr(eb, slot);
 414         ptr = btrfs_item_ptr_offset(eb, slot);
 415         while (cur_offset < item_size) {
 416                 if (ext) {
 417                         extref = (struct btrfs_inode_extref *)(ptr +
 418                                                                cur_offset);
 419                         name_ptr = (unsigned long)(&extref->name);
 420                         len = btrfs_inode_extref_name_len(eb, extref);
 421                         cur_offset += sizeof(*extref);
 422                 } else {
 423                         ref = (struct btrfs_inode_ref *)(ptr + cur_offset);
 424                         len = btrfs_inode_ref_name_len(eb, ref);
 425                         name_ptr = (unsigned long)(ref + 1);
 426                         cur_offset += sizeof(*ref);
 427                 }
 428                 cur_offset += len;
 429
 430                 if (md->sanitize_names > 1) {
 431                         buf = malloc(len);
 432                         if (!buf) {
 433                                 fprintf(stderr, "Couldn't sanitize name, "
 434                                         "enomem\n");
 435                                 return;
 436                         }
 437                         read_extent_buffer(eb, buf, name_ptr, len);
 438                         garbage = find_collision(md, buf, len);
 439                 } else {
 440                         garbage = generate_garbage(len);
 441                 }
 442
 443                 if (!garbage) {
 444                         fprintf(stderr, "Couldn't sanitize name, enomem\n");
 445                         return;
 446                 }
 447                 write_extent_buffer(eb, garbage, name_ptr, len);
 448                 if (free_garbage)
 449                         free(garbage);
 450         }
 451 }
 452
 453 static void sanitize_xattr(struct metadump_struct *md,
 454                            struct extent_buffer *eb, int slot)
 455 {
 456         struct btrfs_dir_item *dir_item;
 457         unsigned long data_ptr;
 458         u32 data_len;
 459
 460         dir_item = btrfs_item_ptr(eb, slot, struct btrfs_dir_item);
 461         data_len = btrfs_dir_data_len(eb, dir_item);
 462
 463         data_ptr = (unsigned long)((char *)(dir_item + 1) +
 464                                    btrfs_dir_name_len(eb, dir_item));
 465         memset_extent_buffer(eb, 0, data_ptr, data_len);
 466 }
 467
 468 static void sanitize_name(struct metadump_struct *md, u8 *dst,
 469                           struct extent_buffer *src, struct btrfs_key *key,
 470                           int slot)
 471 {
 472         struct extent_buffer *eb;
 473
 474         eb = alloc_dummy_eb(src->start, src->len);
 475         if (!eb) {
 476                 fprintf(stderr, "Couldn't sanitize name, no memory\n");
 477                 return;
 478         }
 479
 480         memcpy(eb->data, dst, eb->len);
 481
 482         switch (key->type) {
 483         case BTRFS_DIR_ITEM_KEY:
 484         case BTRFS_DIR_INDEX_KEY:
 485                 sanitize_dir_item(md, eb, slot);
 486                 break;
 487         case BTRFS_INODE_REF_KEY:
 488                 sanitize_inode_ref(md, eb, slot, 0);
 489                 break;
 490         case BTRFS_INODE_EXTREF_KEY:
 491                 sanitize_inode_ref(md, eb, slot, 1);
 492                 break;
 493         case BTRFS_XATTR_ITEM_KEY:
 494                 sanitize_xattr(md, eb, slot);
 495                 break;
 496         default:
 497                 break;
 498         }
 499
 500         memcpy(dst, eb->data, eb->len);
 501         free(eb);
 502 }
 503
 504 /*
 505  * zero inline extents and csum items
 506  */
 507 static void zero_items(struct metadump_struct *md, u8 *dst,
 508                        struct extent_buffer *src)
 509 {
 510         struct btrfs_file_extent_item *fi;
 511         struct btrfs_item *item;
 512         struct btrfs_key key;
 513         u32 nritems = btrfs_header_nritems(src);
 514         size_t size;
 515         unsigned long ptr;
 516         int i, extent_type;
 517
 518         for (i = 0; i < nritems; i++) {
 519                 item = btrfs_item_nr(i);
 520                 btrfs_item_key_to_cpu(src, &key, i);
 521                 if (key.type == BTRFS_CSUM_ITEM_KEY) {
 522                         size = btrfs_item_size_nr(src, i);
 523                         memset(dst + btrfs_leaf_data(src) +
 524                                btrfs_item_offset_nr(src, i), 0, size);
 525                         continue;
 526                 }
 527
 528                 if (md->sanitize_names && has_name(&key)) {
 529                         sanitize_name(md, dst, src, &key, i);
 530                         continue;
 531                 }
 532
 533                 if (key.type != BTRFS_EXTENT_DATA_KEY)
 534                         continue;
 535
 536                 fi = btrfs_item_ptr(src, i, struct btrfs_file_extent_item);
 537                 extent_type = btrfs_file_extent_type(src, fi);
 538                 if (extent_type != BTRFS_FILE_EXTENT_INLINE)
 539                         continue;
 540
 541                 ptr = btrfs_file_extent_inline_start(fi);
 542                 size = btrfs_file_extent_inline_item_len(src, item);
 543                 memset(dst + ptr, 0, size);
 544         }
 545 }
 546
 547 /*
 548  * copy buffer and zero useless data in the buffer
 549  */
 550 static void copy_buffer(struct metadump_struct *md, u8 *dst,
 551                         struct extent_buffer *src)
 552 {
 553         int level;
 554         size_t size;
 555         u32 nritems;
 556
 557         memcpy(dst, src->data, src->len);
 558         if (src->start == BTRFS_SUPER_INFO_OFFSET)
 559                 return;
 560
 561         level = btrfs_header_level(src);
 562         nritems = btrfs_header_nritems(src);
 563
 564         if (nritems == 0) {
 565                 size = sizeof(struct btrfs_header);
 566                 memset(dst + size, 0, src->len - size);
 567         } else if (level == 0) {
 568                 size = btrfs_leaf_data(src) +
 569                         btrfs_item_offset_nr(src, nritems - 1) -
 570                         btrfs_item_nr_offset(nritems);
 571                 memset(dst + btrfs_item_nr_offset(nritems), 0, size);
 572                 zero_items(md, dst, src);
 573         } else {
 574                 size = offsetof(struct btrfs_node, ptrs) +
 575                         sizeof(struct btrfs_key_ptr) * nritems;
 576                 memset(dst + size, 0, src->len - size);
 577         }
 578         csum_block(dst, src->len);
 579 }
 580
 581 static void *dump_worker(void *data)
 582 {
 583         struct metadump_struct *md = (struct metadump_struct *)data;
 584         struct async_work *async;
 585         int ret;
 586
 587         while (1) {
 588                 pthread_mutex_lock(&md->mutex);
 589                 while (list_empty(&md->list)) {
 590                         if (md->done) {
 591                                 pthread_mutex_unlock(&md->mutex);
 592                                 goto out;
 593                         }
 594                         pthread_cond_wait(&md->cond, &md->mutex);
 595                 }
 596                 async = list_entry(md->list.next, struct async_work, list);
 597                 list_del_init(&async->list);
 598                 pthread_mutex_unlock(&md->mutex);
 599
 600                 if (md->compress_level > 0) {
 601                         u8 *orig = async->buffer;
 602
 603                         async->bufsize = compressBound(async->size);
 604                         async->buffer = malloc(async->bufsize);
 605
 606                         ret = compress2(async->buffer,
 607                                          (unsigned long *)&async->bufsize,
 608                                          orig, async->size, md->compress_level);
 609
 610                         if (ret != Z_OK)
 611                                 async->error = 1;
 612
 613                         free(orig);
 614                 }
 615
 616                 pthread_mutex_lock(&md->mutex);
 617                 md->num_ready++;
 618                 pthread_mutex_unlock(&md->mutex);
 619         }
 620 out:
 621         pthread_exit(NULL);
 622 }
 623
 624 static void meta_cluster_init(struct metadump_struct *md, u64 start)
 625 {
 626         struct meta_cluster_header *header;
 627
 628         md->num_items = 0;
 629         md->num_ready = 0;
 630         header = &md->cluster->header;
 631         header->magic = cpu_to_le64(HEADER_MAGIC);
 632         header->bytenr = cpu_to_le64(start);
 633         header->nritems = cpu_to_le32(0);
 634         header->compress = md->compress_level > 0 ?
 635                            COMPRESS_ZLIB : COMPRESS_NONE;
 636 }
 637
 638 static int metadump_init(struct metadump_struct *md, struct btrfs_root *root,
 639                          FILE *out, int num_threads, int compress_level,
 640                          int sanitize_names)
 641 {
 642         int i, ret = 0;
 643
 644         memset(md, 0, sizeof(*md));
 645         pthread_cond_init(&md->cond, NULL);
 646         pthread_mutex_init(&md->mutex, NULL);
 647         INIT_LIST_HEAD(&md->list);
 648         INIT_LIST_HEAD(&md->ordered);
 649         md->root = root;
 650         md->out = out;
 651         md->pending_start = (u64)-1;
 652         md->compress_level = compress_level;
 653         md->cluster = calloc(1, BLOCK_SIZE);
 654         md->sanitize_names = sanitize_names;
 655         if (sanitize_names > 1)
 656                 crc32c_optimization_init();
 657
 658         if (!md->cluster) {
 659                 pthread_cond_destroy(&md->cond);
 660                 pthread_mutex_destroy(&md->mutex);
 661                 return -ENOMEM;
 662         }
 663
 664         meta_cluster_init(md, 0);
 665         if (!num_threads)
 666                 return 0;
 667
 668         md->name_tree.rb_node = NULL;
 669         md->num_threads = num_threads;
 670         md->threads = calloc(num_threads, sizeof(pthread_t));
 671         if (!md->threads) {
 672                 free(md->cluster);
 673                 pthread_cond_destroy(&md->cond);
 674                 pthread_mutex_destroy(&md->mutex);
 675                 return -ENOMEM;
 676         }
 677
 678         for (i = 0; i < num_threads; i++) {
 679                 ret = pthread_create(md->threads + i, NULL, dump_worker, md);
 680                 if (ret)
 681                         break;
 682         }
 683
 684         if (ret) {
 685                 pthread_mutex_lock(&md->mutex);
 686                 md->done = 1;
 687                 pthread_cond_broadcast(&md->cond);
 688                 pthread_mutex_unlock(&md->mutex);
 689
 690                 for (i--; i >= 0; i--)
 691                         pthread_join(md->threads[i], NULL);
 692
 693                 pthread_cond_destroy(&md->cond);
 694                 pthread_mutex_destroy(&md->mutex);
 695                 free(md->cluster);
 696                 free(md->threads);
 697         }
 698
 699         return ret;
 700 }
 701
 702 static void metadump_destroy(struct metadump_struct *md)
 703 {
 704         int i;
 705         struct rb_node *n;
 706
 707         pthread_mutex_lock(&md->mutex);
 708         md->done = 1;
 709         pthread_cond_broadcast(&md->cond);
 710         pthread_mutex_unlock(&md->mutex);
 711
 712         for (i = 0; i < md->num_threads; i++)
 713                 pthread_join(md->threads[i], NULL);
 714
 715         pthread_cond_destroy(&md->cond);
 716         pthread_mutex_destroy(&md->mutex);
 717
 718         while ((n = rb_first(&md->name_tree))) {
 719                 struct name *name;
 720
 721                 name = rb_entry(n, struct name, n);
 722                 rb_erase(n, &md->name_tree);
 723                 free(name->val);
 724                 free(name->sub);
 725                 free(name);
 726         }
 727         free(md->threads);
 728         free(md->cluster);
 729 }
 730
 731 static int write_zero(FILE *out, size_t size)
 732 {
 733         static char zero[BLOCK_SIZE];
 734         return fwrite(zero, size, 1, out);
 735 }
 736
 737 static int write_buffers(struct metadump_struct *md, u64 *next)
 738 {
 739         struct meta_cluster_header *header = &md->cluster->header;
 740         struct meta_cluster_item *item;
 741         struct async_work *async;
 742         u64 bytenr = 0;
 743         u32 nritems = 0;
 744         int ret;
 745         int err = 0;
 746
 747         if (list_empty(&md->ordered))
 748                 goto out;
 749
 750         /* wait until all buffers are compressed */
 751         while (md->num_items > md->num_ready) {
 752                 struct timespec ts = {
 753                         .tv_sec = 0,
 754                         .tv_nsec = 10000000,
 755                 };
 756                 pthread_mutex_unlock(&md->mutex);
 757                 nanosleep(&ts, NULL);
 758                 pthread_mutex_lock(&md->mutex);
 759         }
 760
 761         /* setup and write index block */
 762         list_for_each_entry(async, &md->ordered, ordered) {
 763                 item = md->cluster->items + nritems;
 764                 item->bytenr = cpu_to_le64(async->start);
 765                 item->size = cpu_to_le32(async->bufsize);
 766                 nritems++;
 767         }
 768         header->nritems = cpu_to_le32(nritems);
 769
 770         ret = fwrite(md->cluster, BLOCK_SIZE, 1, md->out);
 771         if (ret != 1) {
 772                 fprintf(stderr, "Error writing out cluster: %d\n", errno);
 773                 return -EIO;
 774         }
 775
 776         /* write buffers */
 777         bytenr += le64_to_cpu(header->bytenr) + BLOCK_SIZE;
 778         while (!list_empty(&md->ordered)) {
 779                 async = list_entry(md->ordered.next, struct async_work,
 780                                    ordered);
 781                 list_del_init(&async->ordered);
 782
 783                 bytenr += async->bufsize;
 784                 if (!err)
 785                         ret = fwrite(async->buffer, async->bufsize, 1,
 786                                      md->out);
 787                 if (ret != 1) {
 788                         err = -EIO;
 789                         ret = 0;
 790                         fprintf(stderr, "Error writing out cluster: %d\n",
 791                                 errno);
 792                 }
 793
 794                 free(async->buffer);
 795                 free(async);
 796         }
 797
 798         /* zero unused space in the last block */
 799         if (!err && bytenr & BLOCK_MASK) {
 800                 size_t size = BLOCK_SIZE - (bytenr & BLOCK_MASK);
 801
 802                 bytenr += size;
 803                 ret = write_zero(md->out, size);
 804                 if (ret != 1) {
 805                         fprintf(stderr, "Error zeroing out buffer: %d\n",
 806                                 errno);
 807                         err = -EIO;
 808                 }
 809         }
 810 out:
 811         *next = bytenr;
 812         return err;
 813 }
 814
 815 static int read_data_extent(struct metadump_struct *md,
 816                             struct async_work *async)
 817 {
 818         struct btrfs_multi_bio *multi = NULL;
 819         struct btrfs_device *device;
 820         u64 bytes_left = async->size;
 821         u64 logical = async->start;
 822         u64 offset = 0;
 823         u64 bytenr;
 824         u64 read_len;
 825         ssize_t done;
 826         int fd;
 827         int ret;
 828
 829         while (bytes_left) {
 830                 read_len = bytes_left;
 831                 ret = btrfs_map_block(&md->root->fs_info->mapping_tree, READ,
 832                                       logical, &read_len, &multi, 0, NULL);
 833                 if (ret) {
 834                         fprintf(stderr, "Couldn't map data block %d\n", ret);
 835                         return ret;
 836                 }
 837
 838                 device = multi->stripes[0].dev;
 839
 840                 if (device->fd == 0) {
 841                         fprintf(stderr,
 842                                 "Device we need to read from is not open\n");
 843                         free(multi);
 844                         return -EIO;
 845                 }
 846                 fd = device->fd;
 847                 bytenr = multi->stripes[0].physical;
 848                 free(multi);
 849
 850                 read_len = min(read_len, bytes_left);
 851                 done = pread64(fd, async->buffer+offset, read_len, bytenr);
 852                 if (done < read_len) {
 853                         if (done < 0)
 854                                 fprintf(stderr, "Error reading extent %d\n",
 855                                         errno);
 856                         else
 857                                 fprintf(stderr, "Short read\n");
 858                         return -EIO;
 859                 }
 860
 861                 bytes_left -= done;
 862                 offset += done;
 863                 logical += done;
 864         }
 865
 866         return 0;
 867 }
 868
 869 static int flush_pending(struct metadump_struct *md, int done)
 870 {
 871         struct async_work *async = NULL;
 872         struct extent_buffer *eb;
 873         u64 blocksize = md->root->nodesize;
 874         u64 start;
 875         u64 size;
 876         size_t offset;
 877         int ret = 0;
 878
 879         if (md->pending_size) {
 880                 async = calloc(1, sizeof(*async));
 881                 if (!async)
 882                         return -ENOMEM;
 883
 884                 async->start = md->pending_start;
 885                 async->size = md->pending_size;
 886                 async->bufsize = async->size;
 887                 async->buffer = malloc(async->bufsize);
 888                 if (!async->buffer) {
 889                         free(async);
 890                         return -ENOMEM;
 891                 }
 892                 offset = 0;
 893                 start = async->start;
 894                 size = async->size;
 895
 896                 if (md->data) {
 897                         ret = read_data_extent(md, async);
 898                         if (ret) {
 899                                 free(async->buffer);
 900                                 free(async);
 901                                 return ret;
 902                         }
 903                 }
 904
 905                 while (!md->data && size > 0) {
 906                         u64 this_read = min(blocksize, size);
 907                         eb = read_tree_block(md->root, start, this_read, 0);
 908                         if (!eb) {
 909                                 free(async->buffer);
 910                                 free(async);
 911                                 fprintf(stderr,
 912                                         "Error reading metadata block\n");
 913                                 return -EIO;
 914                         }
 915                         copy_buffer(md, async->buffer + offset, eb);
 916                         free_extent_buffer(eb);
 917                         start += this_read;
 918                         offset += this_read;
 919                         size -= this_read;
 920                 }
 921
 922                 md->pending_start = (u64)-1;
 923                 md->pending_size = 0;
 924         } else if (!done) {
 925                 return 0;
 926         }
 927
 928         pthread_mutex_lock(&md->mutex);
 929         if (async) {
 930                 list_add_tail(&async->ordered, &md->ordered);
 931                 md->num_items++;
 932                 if (md->compress_level > 0) {
 933                         list_add_tail(&async->list, &md->list);
 934                         pthread_cond_signal(&md->cond);
 935                 } else {
 936                         md->num_ready++;
 937                 }
 938         }
 939         if (md->num_items >= ITEMS_PER_CLUSTER || done) {
 940                 ret = write_buffers(md, &start);
 941                 if (ret)
 942                         fprintf(stderr, "Error writing buffers %d\n",
 943                                 errno);
 944                 else
 945                         meta_cluster_init(md, start);
 946         }
 947         pthread_mutex_unlock(&md->mutex);
 948         return ret;
 949 }
 950
 951 static int add_extent(u64 start, u64 size, struct metadump_struct *md,
 952                       int data)
 953 {
 954         int ret;
 955         if (md->data != data ||
 956             md->pending_size + size > MAX_PENDING_SIZE ||
 957             md->pending_start + md->pending_size != start) {
 958                 ret = flush_pending(md, 0);
 959                 if (ret)
 960                         return ret;
 961                 md->pending_start = start;
 962         }
 963         readahead_tree_block(md->root, start, size, 0);
 964         md->pending_size += size;
 965         md->data = data;
 966         return 0;
 967 }
 968
 969 #ifdef BTRFS_COMPAT_EXTENT_TREE_V0
 970 static int is_tree_block(struct btrfs_root *extent_root,
 971                          struct btrfs_path *path, u64 bytenr)
 972 {
 973         struct extent_buffer *leaf;
 974         struct btrfs_key key;
 975         u64 ref_objectid;
 976         int ret;
 977
 978         leaf = path->nodes[0];
 979         while (1) {
 980                 struct btrfs_extent_ref_v0 *ref_item;
 981                 path->slots[0]++;
 982                 if (path->slots[0] >= btrfs_header_nritems(leaf)) {
 983                         ret = btrfs_next_leaf(extent_root, path);
 984                         if (ret < 0)
 985                                 return ret;
 986                         if (ret > 0)
 987                                 break;
 988                         leaf = path->nodes[0];
 989                 }
 990                 btrfs_item_key_to_cpu(leaf, &key, path->slots[0]);
 991                 if (key.objectid != bytenr)
 992                         break;
 993                 if (key.type != BTRFS_EXTENT_REF_V0_KEY)
 994                         continue;
 995                 ref_item = btrfs_item_ptr(leaf, path->slots[0],
 996                                           struct btrfs_extent_ref_v0);
 997                 ref_objectid = btrfs_ref_objectid_v0(leaf, ref_item);
 998                 if (ref_objectid < BTRFS_FIRST_FREE_OBJECTID)
 999                         return 1;
1000                 break;
1001         }
1002         return 0;
1003 }
1004 #endif
1005
1006 static int copy_tree_blocks(struct btrfs_root *root, struct extent_buffer *eb,
1007                             struct metadump_struct *metadump, int root_tree)
1008 {
1009         struct extent_buffer *tmp;
1010         struct btrfs_root_item *ri;
1011         struct btrfs_key key;
1012         u64 bytenr;
1013         int level;
1014         int nritems = 0;
1015         int i = 0;
1016         int ret;
1017
1018         ret = add_extent(btrfs_header_bytenr(eb), root->leafsize, metadump, 0);
1019         if (ret) {
1020                 fprintf(stderr, "Error adding metadata block\n");
1021                 return ret;
1022         }
1023
1024         if (btrfs_header_level(eb) == 0 && !root_tree)
1025                 return 0;
1026
1027         level = btrfs_header_level(eb);
1028         nritems = btrfs_header_nritems(eb);
1029         for (i = 0; i < nritems; i++) {
1030                 if (level == 0) {
1031                         btrfs_item_key_to_cpu(eb, &key, i);
1032                         if (key.type != BTRFS_ROOT_ITEM_KEY)
1033                                 continue;
1034                         ri = btrfs_item_ptr(eb, i, struct btrfs_root_item);
1035                         bytenr = btrfs_disk_root_bytenr(eb, ri);
1036                         tmp = read_tree_block(root, bytenr, root->leafsize, 0);
1037                         if (!tmp) {
1038                                 fprintf(stderr,
1039                                         "Error reading log root block\n");
1040                                 return -EIO;
1041                         }
1042                         ret = copy_tree_blocks(root, tmp, metadump, 0);
1043                         free_extent_buffer(tmp);
1044                         if (ret)
1045                                 return ret;
1046                 } else {
1047                         bytenr = btrfs_node_blockptr(eb, i);
1048                         tmp = read_tree_block(root, bytenr, root->leafsize, 0);
1049                         if (!tmp) {
1050                                 fprintf(stderr, "Error reading log block\n");
1051                                 return -EIO;
1052                         }
1053                         ret = copy_tree_blocks(root, tmp, metadump, root_tree);
1054                         free_extent_buffer(tmp);
1055                         if (ret)
1056                                 return ret;
1057                 }
1058         }
1059
1060         return 0;
1061 }
1062
1063 static int copy_log_trees(struct btrfs_root *root,
1064                           struct metadump_struct *metadump,
1065                           struct btrfs_path *path)
1066 {
1067         u64 blocknr = btrfs_super_log_root(root->fs_info->super_copy);
1068
1069         if (blocknr == 0)
1070                 return 0;
1071
1072         if (!root->fs_info->log_root_tree ||
1073             !root->fs_info->log_root_tree->node) {
1074                 fprintf(stderr, "Error copying tree log, it wasn't setup\n");
1075                 return -EIO;
1076         }
1077
1078         return copy_tree_blocks(root, root->fs_info->log_root_tree->node,
1079                                 metadump, 1);
1080 }
1081
1082 static int copy_space_cache(struct btrfs_root *root,
1083                             struct metadump_struct *metadump,
1084                             struct btrfs_path *path)
1085 {
1086         struct extent_buffer *leaf;
1087         struct btrfs_file_extent_item *fi;
1088         struct btrfs_key key;
1089         u64 bytenr, num_bytes;
1090         int ret;
1091
1092         root = root->fs_info->tree_root;
1093
1094         key.objectid = 0;
1095         key.type = BTRFS_EXTENT_DATA_KEY;
1096         key.offset = 0;
1097
1098         ret = btrfs_search_slot(NULL, root, &key, path, 0, 0);
1099         if (ret < 0) {
1100                 fprintf(stderr, "Error searching for free space inode %d\n",
1101                         ret);
1102                 return ret;
1103         }
1104
1105         while (1) {
1106                 leaf = path->nodes[0];
1107                 if (path->slots[0] >= btrfs_header_nritems(leaf)) {
1108                         ret = btrfs_next_leaf(root, path);
1109                         if (ret < 0) {
1110                                 fprintf(stderr, "Error going to next leaf "
1111                                         "%d\n", ret);
1112                                 return ret;
1113                         }
1114                         if (ret > 0)
1115                                 break;
1116                         leaf = path->nodes[0];
1117                 }
1118
1119                 btrfs_item_key_to_cpu(leaf, &key, path->slots[0]);
1120                 if (key.type != BTRFS_EXTENT_DATA_KEY) {
1121                         path->slots[0]++;
1122                         continue;
1123                 }
1124
1125                 fi = btrfs_item_ptr(leaf, path->slots[0],
1126                                     struct btrfs_file_extent_item);
1127                 if (btrfs_file_extent_type(leaf, fi) !=
1128                     BTRFS_FILE_EXTENT_REG) {
1129                         path->slots[0]++;
1130                         continue;
1131                 }
1132
1133                 bytenr = btrfs_file_extent_disk_bytenr(leaf, fi);
1134                 num_bytes = btrfs_file_extent_disk_num_bytes(leaf, fi);
1135                 ret = add_extent(bytenr, num_bytes, metadump, 1);
1136                 if (ret) {
1137                         fprintf(stderr, "Error adding space cache blocks %d\n",
1138                                 ret);
1139                         btrfs_release_path(path);
1140                         return ret;
1141                 }
1142                 path->slots[0]++;
1143         }
1144
1145         return 0;
1146 }
1147
1148 static int copy_from_extent_tree(struct metadump_struct *metadump,
1149                                  struct btrfs_path *path)
1150 {
1151         struct btrfs_root *extent_root;
1152         struct extent_buffer *leaf;
1153         struct btrfs_extent_item *ei;
1154         struct btrfs_key key;
1155         u64 bytenr;
1156         u64 num_bytes;
1157         int ret;
1158
1159         extent_root = metadump->root->fs_info->extent_root;
1160         bytenr = BTRFS_SUPER_INFO_OFFSET + 4096;
1161         key.objectid = bytenr;
1162         key.type = BTRFS_EXTENT_ITEM_KEY;
1163         key.offset = 0;
1164
1165         ret = btrfs_search_slot(NULL, extent_root, &key, path, 0, 0);
1166         if (ret < 0) {
1167                 fprintf(stderr, "Error searching extent root %d\n", ret);
1168                 return ret;
1169         }
1170         ret = 0;
1171
1172         while (1) {
1173                 leaf = path->nodes[0];
1174                 if (path->slots[0] >= btrfs_header_nritems(leaf)) {
1175                         ret = btrfs_next_leaf(extent_root, path);
1176                         if (ret < 0) {
1177                                 fprintf(stderr, "Error going to next leaf %d"
1178                                         "\n", ret);
1179                                 break;
1180                         }
1181                         if (ret > 0) {
1182                                 ret = 0;
1183                                 break;
1184                         }
1185                         leaf = path->nodes[0];
1186                 }
1187
1188                 btrfs_item_key_to_cpu(leaf, &key, path->slots[0]);
1189                 if (key.objectid < bytenr ||
1190                     (key.type != BTRFS_EXTENT_ITEM_KEY &&
1191                      key.type != BTRFS_METADATA_ITEM_KEY)) {
1192                         path->slots[0]++;
1193                         continue;
1194                 }
1195
1196                 bytenr = key.objectid;
1197                 if (key.type == BTRFS_METADATA_ITEM_KEY)
1198                         num_bytes = extent_root->leafsize;
1199                 else
1200                         num_bytes = key.offset;
1201
1202                 if (btrfs_item_size_nr(leaf, path->slots[0]) > sizeof(*ei)) {
1203                         ei = btrfs_item_ptr(leaf, path->slots[0],
1204                                             struct btrfs_extent_item);
1205                         if (btrfs_extent_flags(leaf, ei) &
1206                             BTRFS_EXTENT_FLAG_TREE_BLOCK) {
1207                                 ret = add_extent(bytenr, num_bytes, metadump,
1208                                                  0);
1209                                 if (ret) {
1210                                         fprintf(stderr, "Error adding block "
1211                                                 "%d\n", ret);
1212                                         break;
1213                                 }
1214                         }
1215                 } else {
1216 #ifdef BTRFS_COMPAT_EXTENT_TREE_V0
1217                         ret = is_tree_block(extent_root, path, bytenr);
1218                         if (ret < 0) {
1219                                 fprintf(stderr, "Error checking tree block "
1220                                         "%d\n", ret);
1221                                 break;
1222                         }
1223
1224                         if (ret) {
1225                                 ret = add_extent(bytenr, num_bytes, metadump,
1226                                                  0);
1227                                 if (ret) {
1228                                         fprintf(stderr, "Error adding block "
1229                                                 "%d\n", ret);
1230                                         break;
1231                                 }
1232                         }
1233                         ret = 0;
1234 #else
1235                         fprintf(stderr, "Either extent tree corruption or "
1236                                 "you haven't built with V0 support\n");
1237                         ret = -EIO;
1238                         break;
1239 #endif
1240                 }
1241                 bytenr += num_bytes;
1242         }
1243
1244         btrfs_release_path(path);
1245
1246         return ret;
1247 }
1248
1249 static int create_metadump(const char *input, FILE *out, int num_threads,
1250                            int compress_level, int sanitize, int walk_trees)
1251 {
1252         struct btrfs_root *root;
1253         struct btrfs_path *path = NULL;
1254         struct metadump_struct metadump;
1255         int ret;
1256         int err = 0;
1257
1258         root = open_ctree(input, 0, 0);
1259         if (!root) {
1260                 fprintf(stderr, "Open ctree failed\n");
1261                 return -EIO;
1262         }
1263
1264         BUG_ON(root->nodesize != root->leafsize);
1265
1266         ret = metadump_init(&metadump, root, out, num_threads,
1267                             compress_level, sanitize);
1268         if (ret) {
1269                 fprintf(stderr, "Error initing metadump %d\n", ret);
1270                 close_ctree(root);
1271                 return ret;
1272         }
1273
1274         ret = add_extent(BTRFS_SUPER_INFO_OFFSET, 4096, &metadump, 0);
1275         if (ret) {
1276                 fprintf(stderr, "Error adding metadata %d\n", ret);
1277                 err = ret;
1278                 goto out;
1279         }
1280
1281         path = btrfs_alloc_path();
1282         if (!path) {
1283                 fprintf(stderr, "Out of memory allocing path\n");
1284                 err = -ENOMEM;
1285                 goto out;
1286         }
1287
1288         if (walk_trees) {
1289                 ret = copy_tree_blocks(root, root->fs_info->chunk_root->node,
1290                                        &metadump, 1);
1291                 if (ret) {
1292                         err = ret;
1293                         goto out;
1294                 }
1295
1296                 ret = copy_tree_blocks(root, root->fs_info->tree_root->node,
1297                                        &metadump, 1);
1298                 if (ret) {
1299                         err = ret;
1300                         goto out;
1301                 }
1302         } else {
1303                 ret = copy_from_extent_tree(&metadump, path);
1304                 if (ret) {
1305                         err = ret;
1306                         goto out;
1307                 }
1308         }
1309
1310         ret = copy_log_trees(root, &metadump, path);
1311         if (ret) {
1312                 err = ret;
1313                 goto out;
1314         }
1315
1316         ret = copy_space_cache(root, &metadump, path);
1317 out:
1318         ret = flush_pending(&metadump, 1);
1319         if (ret) {
1320                 if (!err)
1321                         err = ret;
1322                 fprintf(stderr, "Error flushing pending %d\n", ret);
1323         }
1324
1325         metadump_destroy(&metadump);
1326
1327         btrfs_free_path(path);
1328         ret = close_ctree(root);
1329         return err ? err : ret;
1330 }
1331
1332 static void update_super_old(u8 *buffer)
1333 {
1334         struct btrfs_super_block *super = (struct btrfs_super_block *)buffer;
1335         struct btrfs_chunk *chunk;
1336         struct btrfs_disk_key *key;
1337         u32 sectorsize = btrfs_super_sectorsize(super);
1338         u64 flags = btrfs_super_flags(super);
1339
1340         flags |= BTRFS_SUPER_FLAG_METADUMP;
1341         btrfs_set_super_flags(super, flags);
1342
1343         key = (struct btrfs_disk_key *)(super->sys_chunk_array);
1344         chunk = (struct btrfs_chunk *)(super->sys_chunk_array +
1345                                        sizeof(struct btrfs_disk_key));
1346
1347         btrfs_set_disk_key_objectid(key, BTRFS_FIRST_CHUNK_TREE_OBJECTID);
1348         btrfs_set_disk_key_type(key, BTRFS_CHUNK_ITEM_KEY);
1349         btrfs_set_disk_key_offset(key, 0);
1350
1351         btrfs_set_stack_chunk_length(chunk, (u64)-1);
1352         btrfs_set_stack_chunk_owner(chunk, BTRFS_EXTENT_TREE_OBJECTID);
1353         btrfs_set_stack_chunk_stripe_len(chunk, 64 * 1024);
1354         btrfs_set_stack_chunk_type(chunk, BTRFS_BLOCK_GROUP_SYSTEM);
1355         btrfs_set_stack_chunk_io_align(chunk, sectorsize);
1356         btrfs_set_stack_chunk_io_width(chunk, sectorsize);
1357         btrfs_set_stack_chunk_sector_size(chunk, sectorsize);
1358         btrfs_set_stack_chunk_num_stripes(chunk, 1);
1359         btrfs_set_stack_chunk_sub_stripes(chunk, 0);
1360         chunk->stripe.devid = super->dev_item.devid;
1361         btrfs_set_stack_stripe_offset(&chunk->stripe, 0);
1362         memcpy(chunk->stripe.dev_uuid, super->dev_item.uuid, BTRFS_UUID_SIZE);
1363         btrfs_set_super_sys_array_size(super, sizeof(*key) + sizeof(*chunk));
1364         csum_block(buffer, 4096);
1365 }
1366
1367 static int update_super(u8 *buffer)
1368 {
1369         struct btrfs_super_block *super = (struct btrfs_super_block *)buffer;
1370         struct btrfs_chunk *chunk;
1371         struct btrfs_disk_key *disk_key;
1372         struct btrfs_key key;
1373         u32 new_array_size = 0;
1374         u32 array_size;
1375         u32 cur = 0;
1376         u32 new_cur = 0;
1377         u8 *ptr, *write_ptr;
1378         int old_num_stripes;
1379
1380         write_ptr = ptr = super->sys_chunk_array;
1381         array_size = btrfs_super_sys_array_size(super);
1382
1383         while (cur < array_size) {
1384                 disk_key = (struct btrfs_disk_key *)ptr;
1385                 btrfs_disk_key_to_cpu(&key, disk_key);
1386
1387                 new_array_size += sizeof(*disk_key);
1388                 memmove(write_ptr, ptr, sizeof(*disk_key));
1389
1390                 write_ptr += sizeof(*disk_key);
1391                 ptr += sizeof(*disk_key);
1392                 cur += sizeof(*disk_key);
1393                 new_cur += sizeof(*disk_key);
1394
1395                 if (key.type == BTRFS_CHUNK_ITEM_KEY) {
1396                         chunk = (struct btrfs_chunk *)ptr;
1397                         old_num_stripes = btrfs_stack_chunk_num_stripes(chunk);
1398                         chunk = (struct btrfs_chunk *)write_ptr;
1399
1400                         memmove(write_ptr, ptr, sizeof(*chunk));
1401                         btrfs_set_stack_chunk_num_stripes(chunk, 1);
1402                         btrfs_set_stack_chunk_sub_stripes(chunk, 0);
1403                         btrfs_set_stack_chunk_type(chunk,
1404                                                    BTRFS_BLOCK_GROUP_SYSTEM);
1405                         chunk->stripe.devid = super->dev_item.devid;
1406                         memcpy(chunk->stripe.dev_uuid, super->dev_item.uuid,
1407                                BTRFS_UUID_SIZE);
1408                         new_array_size += sizeof(*chunk);
1409                         new_cur += sizeof(*chunk);
1410                 } else {
1411                         fprintf(stderr, "Bogus key in the sys chunk array "
1412                                 "%d\n", key.type);
1413                         return -EIO;
1414                 }
1415                 write_ptr += sizeof(*chunk);
1416                 ptr += btrfs_chunk_item_size(old_num_stripes);
1417                 cur += btrfs_chunk_item_size(old_num_stripes);
1418         }
1419
1420         btrfs_set_super_sys_array_size(super, new_array_size);
1421         csum_block(buffer, 4096);
1422
1423         return 0;
1424 }
1425
1426 static struct extent_buffer *alloc_dummy_eb(u64 bytenr, u32 size)
1427 {
1428         struct extent_buffer *eb;
1429
1430         eb = malloc(sizeof(struct extent_buffer) + size);
1431         if (!eb)
1432                 return NULL;
1433         memset(eb, 0, sizeof(struct extent_buffer) + size);
1434
1435         eb->start = bytenr;
1436         eb->len = size;
1437         return eb;
1438 }
1439
1440 static void truncate_item(struct extent_buffer *eb, int slot, u32 new_size)
1441 {
1442         struct btrfs_item *item;
1443         u32 nritems;
1444         u32 old_size;
1445         u32 old_data_start;
1446         u32 size_diff;
1447         u32 data_end;
1448         int i;
1449
1450         old_size = btrfs_item_size_nr(eb, slot);
1451         if (old_size == new_size)
1452                 return;
1453
1454         nritems = btrfs_header_nritems(eb);
1455         data_end = btrfs_item_offset_nr(eb, nritems - 1);
1456
1457         old_data_start = btrfs_item_offset_nr(eb, slot);
1458         size_diff = old_size - new_size;
1459
1460         for (i = slot; i < nritems; i++) {
1461                 u32 ioff;
1462                 item = btrfs_item_nr(i);
1463                 ioff = btrfs_item_offset(eb, item);
1464                 btrfs_set_item_offset(eb, item, ioff + size_diff);
1465         }
1466
1467         memmove_extent_buffer(eb, btrfs_leaf_data(eb) + data_end + size_diff,
1468                               btrfs_leaf_data(eb) + data_end,
1469                               old_data_start + new_size - data_end);
1470         item = btrfs_item_nr(slot);
1471         btrfs_set_item_size(eb, item, new_size);
1472 }
1473
1474 static int fixup_chunk_tree_block(struct mdrestore_struct *mdres,
1475                                   struct async_work *async, u8 *buffer,
1476                                   size_t size)
1477 {
1478         struct extent_buffer *eb;
1479         size_t size_left = size;
1480         u64 bytenr = async->start;
1481         int i;
1482
1483         if (size_left % mdres->leafsize)
1484                 return 0;
1485
1486         eb = alloc_dummy_eb(bytenr, mdres->leafsize);
1487         if (!eb)
1488                 return -ENOMEM;
1489
1490         while (size_left) {
1491                 eb->start = bytenr;
1492                 memcpy(eb->data, buffer, mdres->leafsize);
1493
1494                 if (btrfs_header_bytenr(eb) != bytenr)
1495                         break;
1496                 if (memcmp(mdres->fsid,
1497                            eb->data + offsetof(struct btrfs_header, fsid),
1498                            BTRFS_FSID_SIZE))
1499                         break;
1500
1501                 if (btrfs_header_owner(eb) != BTRFS_CHUNK_TREE_OBJECTID)
1502                         goto next;
1503
1504                 if (btrfs_header_level(eb) != 0)
1505                         goto next;
1506
1507                 for (i = 0; i < btrfs_header_nritems(eb); i++) {
1508                         struct btrfs_chunk chunk;
1509                         struct btrfs_key key;
1510                         u64 type;
1511
1512                         btrfs_item_key_to_cpu(eb, &key, i);
1513                         if (key.type != BTRFS_CHUNK_ITEM_KEY)
1514                                 continue;
1515                         truncate_item(eb, i, sizeof(chunk));
1516                         read_extent_buffer(eb, &chunk,
1517                                            btrfs_item_ptr_offset(eb, i),
1518                                            sizeof(chunk));
1519
1520                         /* Zero out the RAID profile */
1521                         type = btrfs_stack_chunk_type(&chunk);
1522                         type &= (BTRFS_BLOCK_GROUP_DATA |
1523                                  BTRFS_BLOCK_GROUP_SYSTEM |
1524                                  BTRFS_BLOCK_GROUP_METADATA);
1525                         btrfs_set_stack_chunk_type(&chunk, type);
1526
1527                         btrfs_set_stack_chunk_num_stripes(&chunk, 1);
1528                         btrfs_set_stack_chunk_sub_stripes(&chunk, 0);
1529                         btrfs_set_stack_stripe_devid(&chunk.stripe, mdres->devid);
1530                         memcpy(chunk.stripe.dev_uuid, mdres->uuid,
1531                                BTRFS_UUID_SIZE);
1532                         write_extent_buffer(eb, &chunk,
1533                                             btrfs_item_ptr_offset(eb, i),
1534                                             sizeof(chunk));
1535                 }
1536                 memcpy(buffer, eb->data, eb->len);
1537                 csum_block(buffer, eb->len);
1538 next:
1539                 size_left -= mdres->leafsize;
1540                 buffer += mdres->leafsize;
1541                 bytenr += mdres->leafsize;
1542         }
1543
1544         free(eb);
1545         return 0;
1546 }
1547
1548 static void write_backup_supers(int fd, u8 *buf)
1549 {
1550         struct btrfs_super_block *super = (struct btrfs_super_block *)buf;
1551         struct stat st;
1552         u64 size;
1553         u64 bytenr;
1554         int i;
1555         int ret;
1556
1557         if (fstat(fd, &st)) {
1558                 fprintf(stderr, "Couldn't stat restore point, won't be able "
1559                         "to write backup supers: %d\n", errno);
1560                 return;
1561         }
1562
1563         size = btrfs_device_size(fd, &st);
1564
1565         for (i = 1; i < BTRFS_SUPER_MIRROR_MAX; i++) {
1566                 bytenr = btrfs_sb_offset(i);
1567                 if (bytenr + 4096 > size)
1568                         break;
1569                 btrfs_set_super_bytenr(super, bytenr);
1570                 csum_block(buf, 4096);
1571                 ret = pwrite64(fd, buf, 4096, bytenr);
1572                 if (ret < 4096) {
1573                         if (ret < 0)
1574                                 fprintf(stderr, "Problem writing out backup "
1575                                         "super block %d, err %d\n", i, errno);
1576                         else
1577                                 fprintf(stderr, "Short write writing out "
1578                                         "backup super block\n");
1579                         break;
1580                 }
1581         }
1582 }
1583
1584 static u64 logical_to_physical(struct mdrestore_struct *mdres, u64 logical, u64 *size)
1585 {
1586         struct fs_chunk *fs_chunk;
1587         struct rb_node *entry;
1588         struct fs_chunk search;
1589         u64 offset;
1590
1591         if (logical == BTRFS_SUPER_INFO_OFFSET)
1592                 return logical;
1593
1594         search.logical = logical;
1595         entry = tree_search(&mdres->chunk_tree, &search.n, chunk_cmp, 1);
1596         if (!entry) {
1597                 if (mdres->in != stdin)
1598                         printf("Couldn't find a chunk, using logical\n");
1599                 return logical;
1600         }
1601         fs_chunk = rb_entry(entry, struct fs_chunk, n);
1602         if (fs_chunk->logical > logical || fs_chunk->logical + fs_chunk->bytes < logical)
1603                 BUG();
1604         offset = search.logical - fs_chunk->logical;
1605
1606         *size = min(*size, fs_chunk->bytes + fs_chunk->logical - logical);
1607         return fs_chunk->physical + offset;
1608 }
1609
1610 static void *restore_worker(void *data)
1611 {
1612         struct mdrestore_struct *mdres = (struct mdrestore_struct *)data;
1613         struct async_work *async;
1614         size_t size;
1615         u8 *buffer;
1616         u8 *outbuf;
1617         int outfd;
1618         int ret;
1619         int compress_size = MAX_PENDING_SIZE * 4;
1620
1621         outfd = fileno(mdres->out);
1622         buffer = malloc(compress_size);
1623         if (!buffer) {
1624                 fprintf(stderr, "Error allocing buffer\n");
1625                 pthread_mutex_lock(&mdres->mutex);
1626                 if (!mdres->error)
1627                         mdres->error = -ENOMEM;
1628                 pthread_mutex_unlock(&mdres->mutex);
1629                 goto out;
1630         }
1631
1632         while (1) {
1633                 u64 bytenr;
1634                 off_t offset = 0;
1635                 int err = 0;
1636
1637                 pthread_mutex_lock(&mdres->mutex);
1638                 while (!mdres->leafsize || list_empty(&mdres->list)) {
1639                         if (mdres->done) {
1640                                 pthread_mutex_unlock(&mdres->mutex);
1641                                 goto out;
1642                         }
1643                         pthread_cond_wait(&mdres->cond, &mdres->mutex);
1644                 }
1645                 async = list_entry(mdres->list.next, struct async_work, list);
1646                 list_del_init(&async->list);
1647                 pthread_mutex_unlock(&mdres->mutex);
1648
1649                 if (mdres->compress_method == COMPRESS_ZLIB) {
1650                         size = compress_size;
1651                         ret = uncompress(buffer, (unsigned long *)&size,
1652                                          async->buffer, async->bufsize);
1653                         if (ret != Z_OK) {
1654                                 fprintf(stderr, "Error decompressing %d\n",
1655                                         ret);
1656                                 err = -EIO;
1657                         }
1658                         outbuf = buffer;
1659                 } else {
1660                         outbuf = async->buffer;
1661                         size = async->bufsize;
1662                 }
1663
1664                 if (!mdres->multi_devices) {
1665                         if (async->start == BTRFS_SUPER_INFO_OFFSET) {
1666                                 if (mdres->old_restore) {
1667                                         update_super_old(outbuf);
1668                                 } else {
1669                                         ret = update_super(outbuf);
1670                                         if (ret)
1671                                                 err = ret;
1672                                 }
1673                         } else if (!mdres->old_restore) {
1674                                 ret = fixup_chunk_tree_block(mdres, async, outbuf, size);
1675                                 if (ret)
1676                                         err = ret;
1677                         }
1678                 }
1679
1680                 if (!mdres->fixup_offset) {
1681                         while (size) {
1682                                 u64 chunk_size = size;
1683                                 if (!mdres->multi_devices)
1684                                         bytenr = logical_to_physical(mdres,
1685                                                                      async->start + offset,
1686                                                                      &chunk_size);
1687                                 else
1688                                         bytenr = async->start + offset;
1689
1690                                 ret = pwrite64(outfd, outbuf+offset, chunk_size,
1691                                                bytenr);
1692                                 if (ret != chunk_size) {
1693                                         if (ret < 0) {
1694                                                 fprintf(stderr, "Error writing to "
1695                                                         "device %d\n", errno);
1696                                                 err = errno;
1697                                                 break;
1698                                         } else {
1699                                                 fprintf(stderr, "Short write\n");
1700                                                 err = -EIO;
1701                                                 break;
1702                                         }
1703                                 }
1704                                 size -= chunk_size;
1705                                 offset += chunk_size;
1706                         }
1707                 } else if (async->start != BTRFS_SUPER_INFO_OFFSET) {
1708                         ret = write_data_to_disk(mdres->info, outbuf, async->start, size, 0);
1709                         if (ret) {
1710                                 printk("Error write data\n");
1711                                 exit(1);
1712                         }
1713                 }
1714
1715
1716                 /* backup super blocks are already there at fixup_offset stage */
1717                 if (!mdres->multi_devices && async->start == BTRFS_SUPER_INFO_OFFSET)
1718                         write_backup_supers(outfd, outbuf);
1719
1720                 pthread_mutex_lock(&mdres->mutex);
1721                 if (err && !mdres->error)
1722                         mdres->error = err;
1723                 mdres->num_items--;
1724                 pthread_mutex_unlock(&mdres->mutex);
1725
1726                 free(async->buffer);
1727                 free(async);
1728         }
1729 out:
1730         free(buffer);
1731         pthread_exit(NULL);
1732 }
1733
1734 static void mdrestore_destroy(struct mdrestore_struct *mdres)
1735 {
1736         struct rb_node *n;
1737         int i;
1738
1739         while ((n = rb_first(&mdres->chunk_tree))) {
1740                 struct fs_chunk *entry;
1741
1742                 entry = rb_entry(n, struct fs_chunk, n);
1743                 rb_erase(n, &mdres->chunk_tree);
1744                 free(entry);
1745         }
1746         pthread_mutex_lock(&mdres->mutex);
1747         mdres->done = 1;
1748         pthread_cond_broadcast(&mdres->cond);
1749         pthread_mutex_unlock(&mdres->mutex);
1750
1751         for (i = 0; i < mdres->num_threads; i++)
1752                 pthread_join(mdres->threads[i], NULL);
1753
1754         pthread_cond_destroy(&mdres->cond);
1755         pthread_mutex_destroy(&mdres->mutex);
1756         free(mdres->threads);
1757 }
1758
1759 static int mdrestore_init(struct mdrestore_struct *mdres,
1760                           FILE *in, FILE *out, int old_restore,
1761                           int num_threads, int fixup_offset,
1762                           struct btrfs_fs_info *info, int multi_devices)
1763 {
1764         int i, ret = 0;
1765
1766         memset(mdres, 0, sizeof(*mdres));
1767         pthread_cond_init(&mdres->cond, NULL);
1768         pthread_mutex_init(&mdres->mutex, NULL);
1769         INIT_LIST_HEAD(&mdres->list);
1770         mdres->in = in;
1771         mdres->out = out;
1772         mdres->old_restore = old_restore;
1773         mdres->chunk_tree.rb_node = NULL;
1774         mdres->fixup_offset = fixup_offset;
1775         mdres->info = info;
1776         mdres->multi_devices = multi_devices;
1777
1778         if (!num_threads)
1779                 return 0;
1780
1781         mdres->num_threads = num_threads;
1782         mdres->threads = calloc(num_threads, sizeof(pthread_t));
1783         if (!mdres->threads)
1784                 return -ENOMEM;
1785         for (i = 0; i < num_threads; i++) {
1786                 ret = pthread_create(mdres->threads + i, NULL, restore_worker,
1787                                      mdres);
1788                 if (ret)
1789                         break;
1790         }
1791         if (ret)
1792                 mdrestore_destroy(mdres);
1793         return ret;
1794 }
1795
1796 static int fill_mdres_info(struct mdrestore_struct *mdres,
1797                            struct async_work *async)
1798 {
1799         struct btrfs_super_block *super;
1800         u8 *buffer = NULL;
1801         u8 *outbuf;
1802         int ret;
1803
1804         /* We've already been initialized */
1805         if (mdres->leafsize)
1806                 return 0;
1807
1808         if (mdres->compress_method == COMPRESS_ZLIB) {
1809                 size_t size = MAX_PENDING_SIZE * 2;
1810
1811                 buffer = malloc(MAX_PENDING_SIZE * 2);
1812                 if (!buffer)
1813                         return -ENOMEM;
1814                 ret = uncompress(buffer, (unsigned long *)&size,
1815                                  async->buffer, async->bufsize);
1816                 if (ret != Z_OK) {
1817                         fprintf(stderr, "Error decompressing %d\n", ret);
1818                         free(buffer);
1819                         return -EIO;
1820                 }
1821                 outbuf = buffer;
1822         } else {
1823                 outbuf = async->buffer;
1824         }
1825
1826         super = (struct btrfs_super_block *)outbuf;
1827         mdres->leafsize = btrfs_super_leafsize(super);
1828         memcpy(mdres->fsid, super->fsid, BTRFS_FSID_SIZE);
1829         memcpy(mdres->uuid, super->dev_item.uuid,
1830                        BTRFS_UUID_SIZE);
1831         mdres->devid = le64_to_cpu(super->dev_item.devid);
1832         free(buffer);
1833         return 0;
1834 }
1835
1836 static int add_cluster(struct meta_cluster *cluster,
1837                        struct mdrestore_struct *mdres, u64 *next)
1838 {
1839         struct meta_cluster_item *item;
1840         struct meta_cluster_header *header = &cluster->header;
1841         struct async_work *async;
1842         u64 bytenr;
1843         u32 i, nritems;
1844         int ret;
1845
1846         BUG_ON(mdres->num_items);
1847         mdres->compress_method = header->compress;
1848
1849         bytenr = le64_to_cpu(header->bytenr) + BLOCK_SIZE;
1850         nritems = le32_to_cpu(header->nritems);
1851         for (i = 0; i < nritems; i++) {
1852                 item = &cluster->items[i];
1853                 async = calloc(1, sizeof(*async));
1854                 if (!async) {
1855                         fprintf(stderr, "Error allocating async\n");
1856                         return -ENOMEM;
1857                 }
1858                 async->start = le64_to_cpu(item->bytenr);
1859                 async->bufsize = le32_to_cpu(item->size);
1860                 async->buffer = malloc(async->bufsize);
1861                 if (!async->buffer) {
1862                         fprintf(stderr, "Error allocing async buffer\n");
1863                         free(async);
1864                         return -ENOMEM;
1865                 }
1866                 ret = fread(async->buffer, async->bufsize, 1, mdres->in);
1867                 if (ret != 1) {
1868                         fprintf(stderr, "Error reading buffer %d\n", errno);
1869                         free(async->buffer);
1870                         free(async);
1871                         return -EIO;
1872                 }
1873                 bytenr += async->bufsize;
1874
1875                 pthread_mutex_lock(&mdres->mutex);
1876                 if (async->start == BTRFS_SUPER_INFO_OFFSET) {
1877                         ret = fill_mdres_info(mdres, async);
1878                         if (ret) {
1879                                 fprintf(stderr, "Error setting up restore\n");
1880                                 pthread_mutex_unlock(&mdres->mutex);
1881                                 free(async->buffer);
1882                                 free(async);
1883                                 return ret;
1884                         }
1885                 }
1886                 list_add_tail(&async->list, &mdres->list);
1887                 mdres->num_items++;
1888                 pthread_cond_signal(&mdres->cond);
1889                 pthread_mutex_unlock(&mdres->mutex);
1890         }
1891         if (bytenr & BLOCK_MASK) {
1892                 char buffer[BLOCK_MASK];
1893                 size_t size = BLOCK_SIZE - (bytenr & BLOCK_MASK);
1894
1895                 bytenr += size;
1896                 ret = fread(buffer, size, 1, mdres->in);
1897                 if (ret != 1) {
1898                         fprintf(stderr, "Error reading in buffer %d\n", errno);
1899                         return -EIO;
1900                 }
1901         }
1902         *next = bytenr;
1903         return 0;
1904 }
1905
1906 static int wait_for_worker(struct mdrestore_struct *mdres)
1907 {
1908         int ret = 0;
1909
1910         pthread_mutex_lock(&mdres->mutex);
1911         ret = mdres->error;
1912         while (!ret && mdres->num_items > 0) {
1913                 struct timespec ts = {
1914                         .tv_sec = 0,
1915                         .tv_nsec = 10000000,
1916                 };
1917                 pthread_mutex_unlock(&mdres->mutex);
1918                 nanosleep(&ts, NULL);
1919                 pthread_mutex_lock(&mdres->mutex);
1920                 ret = mdres->error;
1921         }
1922         pthread_mutex_unlock(&mdres->mutex);
1923         return ret;
1924 }
1925
1926 static int read_chunk_block(struct mdrestore_struct *mdres, u8 *buffer,
1927                             u64 bytenr, u64 item_bytenr, u32 bufsize,
1928                             u64 cluster_bytenr)
1929 {
1930         struct extent_buffer *eb;
1931         int ret = 0;
1932         int i;
1933
1934         eb = alloc_dummy_eb(bytenr, mdres->leafsize);
1935         if (!eb) {
1936                 ret = -ENOMEM;
1937                 goto out;
1938         }
1939
1940         while (item_bytenr != bytenr) {
1941                 buffer += mdres->leafsize;
1942                 item_bytenr += mdres->leafsize;
1943         }
1944
1945         memcpy(eb->data, buffer, mdres->leafsize);
1946         if (btrfs_header_bytenr(eb) != bytenr) {
1947                 fprintf(stderr, "Eb bytenr doesn't match found bytenr\n");
1948                 ret = -EIO;
1949                 goto out;
1950         }
1951
1952         if (memcmp(mdres->fsid, eb->data + offsetof(struct btrfs_header, fsid),
1953                    BTRFS_FSID_SIZE)) {
1954                 fprintf(stderr, "Fsid doesn't match\n");
1955                 ret = -EIO;
1956                 goto out;
1957         }
1958
1959         if (btrfs_header_owner(eb) != BTRFS_CHUNK_TREE_OBJECTID) {
1960                 fprintf(stderr, "Does not belong to the chunk tree\n");
1961                 ret = -EIO;
1962                 goto out;
1963         }
1964
1965         for (i = 0; i < btrfs_header_nritems(eb); i++) {
1966                 struct btrfs_chunk chunk;
1967                 struct fs_chunk *fs_chunk;
1968                 struct btrfs_key key;
1969
1970                 if (btrfs_header_level(eb)) {
1971                         u64 blockptr = btrfs_node_blockptr(eb, i);
1972
1973                         ret = search_for_chunk_blocks(mdres, blockptr,
1974                                                       cluster_bytenr);
1975                         if (ret)
1976                                 break;
1977                         continue;
1978                 }
1979
1980                 /* Yay a leaf!  We loves leafs! */
1981                 btrfs_item_key_to_cpu(eb, &key, i);
1982                 if (key.type != BTRFS_CHUNK_ITEM_KEY)
1983                         continue;
1984
1985                 fs_chunk = malloc(sizeof(struct fs_chunk));
1986                 if (!fs_chunk) {
1987                         fprintf(stderr, "Erorr allocating chunk\n");
1988                         ret = -ENOMEM;
1989                         break;
1990                 }
1991                 memset(fs_chunk, 0, sizeof(*fs_chunk));
1992                 read_extent_buffer(eb, &chunk, btrfs_item_ptr_offset(eb, i),
1993                                    sizeof(chunk));
1994
1995                 fs_chunk->logical = key.offset;
1996                 fs_chunk->physical = btrfs_stack_stripe_offset(&chunk.stripe);
1997                 fs_chunk->bytes = btrfs_stack_chunk_length(&chunk);
1998                 tree_insert(&mdres->chunk_tree, &fs_chunk->n, chunk_cmp);
1999         }
2000 out:
2001         free(eb);
2002         return ret;
2003 }
2004
2005 /* If you have to ask you aren't worthy */
2006 static int search_for_chunk_blocks(struct mdrestore_struct *mdres,
2007                                    u64 search, u64 cluster_bytenr)
2008 {
2009         struct meta_cluster *cluster;
2010         struct meta_cluster_header *header;
2011         struct meta_cluster_item *item;
2012         u64 current_cluster = cluster_bytenr, bytenr;
2013         u64 item_bytenr;
2014         u32 bufsize, nritems, i;
2015         u32 max_size = MAX_PENDING_SIZE * 2;
2016         u8 *buffer, *tmp = NULL;
2017         int ret = 0;
2018
2019         cluster = malloc(BLOCK_SIZE);
2020         if (!cluster) {
2021                 fprintf(stderr, "Error allocating cluster\n");
2022                 return -ENOMEM;
2023         }
2024
2025         buffer = malloc(max_size);
2026         if (!buffer) {
2027                 fprintf(stderr, "Error allocing buffer\n");
2028                 free(cluster);
2029                 return -ENOMEM;
2030         }
2031
2032         if (mdres->compress_method == COMPRESS_ZLIB) {
2033                 tmp = malloc(max_size);
2034                 if (!tmp) {
2035                         fprintf(stderr, "Error allocing tmp buffer\n");
2036                         free(cluster);
2037                         free(buffer);
2038                         return -ENOMEM;
2039                 }
2040         }
2041
2042         bytenr = current_cluster;
2043         while (1) {
2044                 if (fseek(mdres->in, current_cluster, SEEK_SET)) {
2045                         fprintf(stderr, "Error seeking: %d\n", errno);
2046                         ret = -EIO;
2047                         break;
2048                 }
2049
2050                 ret = fread(cluster, BLOCK_SIZE, 1, mdres->in);
2051                 if (ret == 0) {
2052                         if (cluster_bytenr != 0) {
2053                                 cluster_bytenr = 0;
2054                                 current_cluster = 0;
2055                                 bytenr = 0;
2056                                 continue;
2057                         }
2058                         printf("ok this is where we screwed up?\n");
2059                         ret = -EIO;
2060                         break;
2061                 } else if (ret < 0) {
2062                         fprintf(stderr, "Error reading image\n");
2063                         break;
2064                 }
2065                 ret = 0;
2066
2067                 header = &cluster->header;
2068                 if (le64_to_cpu(header->magic) != HEADER_MAGIC ||
2069                     le64_to_cpu(header->bytenr) != current_cluster) {
2070                         fprintf(stderr, "bad header in metadump image\n");
2071                         ret = -EIO;
2072                         break;
2073                 }
2074
2075                 bytenr += BLOCK_SIZE;
2076                 nritems = le32_to_cpu(header->nritems);
2077                 for (i = 0; i < nritems; i++) {
2078                         size_t size;
2079
2080                         item = &cluster->items[i];
2081                         bufsize = le32_to_cpu(item->size);
2082                         item_bytenr = le64_to_cpu(item->bytenr);
2083
2084                         if (bufsize > max_size) {
2085                                 fprintf(stderr, "item %u size %u too big\n",
2086                                         i, bufsize);
2087                                 ret = -EIO;
2088                                 break;
2089                         }
2090
2091                         if (mdres->compress_method == COMPRESS_ZLIB) {
2092                                 ret = fread(tmp, bufsize, 1, mdres->in);
2093                                 if (ret != 1) {
2094                                         fprintf(stderr, "Error reading: %d\n",
2095                                                 errno);
2096                                         ret = -EIO;
2097                                         break;
2098                                 }
2099
2100                                 size = max_size;
2101                                 ret = uncompress(buffer,
2102                                                  (unsigned long *)&size, tmp,
2103                                                  bufsize);
2104                                 if (ret != Z_OK) {
2105                                         fprintf(stderr, "Error decompressing "
2106                                                 "%d\n", ret);
2107                                         ret = -EIO;
2108                                         break;
2109                                 }
2110                         } else {
2111                                 ret = fread(buffer, bufsize, 1, mdres->in);
2112                                 if (ret != 1) {
2113                                         fprintf(stderr, "Error reading: %d\n",
2114                                                 errno);
2115                                         ret = -EIO;
2116                                         break;
2117                                 }
2118                                 size = bufsize;
2119                         }
2120                         ret = 0;
2121
2122                         if (item_bytenr <= search &&
2123                             item_bytenr + size > search) {
2124                                 ret = read_chunk_block(mdres, buffer, search,
2125                                                        item_bytenr, size,
2126                                                        current_cluster);
2127                                 if (!ret)
2128                                         ret = 1;
2129                                 break;
2130                         }
2131                         bytenr += bufsize;
2132                 }
2133                 if (ret) {
2134                         if (ret > 0)
2135                                 ret = 0;
2136                         break;
2137                 }
2138                 if (bytenr & BLOCK_MASK)
2139                         bytenr += BLOCK_SIZE - (bytenr & BLOCK_MASK);
2140                 current_cluster = bytenr;
2141         }
2142
2143         free(tmp);
2144         free(buffer);
2145         free(cluster);
2146         return ret;
2147 }
2148
2149 static int build_chunk_tree(struct mdrestore_struct *mdres,
2150                             struct meta_cluster *cluster)
2151 {
2152         struct btrfs_super_block *super;
2153         struct meta_cluster_header *header;
2154         struct meta_cluster_item *item = NULL;
2155         u64 chunk_root_bytenr = 0;
2156         u32 i, nritems;
2157         u64 bytenr = 0;
2158         u8 *buffer;
2159         int ret;
2160
2161         /* We can't seek with stdin so don't bother doing this */
2162         if (mdres->in == stdin)
2163                 return 0;
2164
2165         ret = fread(cluster, BLOCK_SIZE, 1, mdres->in);
2166         if (ret <= 0) {
2167                 fprintf(stderr, "Error reading in cluster: %d\n", errno);
2168                 return -EIO;
2169         }
2170         ret = 0;
2171
2172         header = &cluster->header;
2173         if (le64_to_cpu(header->magic) != HEADER_MAGIC ||
2174             le64_to_cpu(header->bytenr) != 0) {
2175                 fprintf(stderr, "bad header in metadump image\n");
2176                 return -EIO;
2177         }
2178
2179         bytenr += BLOCK_SIZE;
2180         mdres->compress_method = header->compress;
2181         nritems = le32_to_cpu(header->nritems);
2182         for (i = 0; i < nritems; i++) {
2183                 item = &cluster->items[i];
2184
2185                 if (le64_to_cpu(item->bytenr) == BTRFS_SUPER_INFO_OFFSET)
2186                         break;
2187                 bytenr += le32_to_cpu(item->size);
2188                 if (fseek(mdres->in, le32_to_cpu(item->size), SEEK_CUR)) {
2189                         fprintf(stderr, "Error seeking: %d\n", errno);
2190                         return -EIO;
2191                 }
2192         }
2193
2194         if (!item || le64_to_cpu(item->bytenr) != BTRFS_SUPER_INFO_OFFSET) {
2195                 fprintf(stderr, "Huh, didn't find the super?\n");
2196                 return -EINVAL;
2197         }
2198
2199         buffer = malloc(le32_to_cpu(item->size));
2200         if (!buffer) {
2201                 fprintf(stderr, "Error allocing buffer\n");
2202                 return -ENOMEM;
2203         }
2204
2205         ret = fread(buffer, le32_to_cpu(item->size), 1, mdres->in);
2206         if (ret != 1) {
2207                 fprintf(stderr, "Error reading buffer: %d\n", errno);
2208                 free(buffer);
2209                 return -EIO;
2210         }
2211
2212         if (mdres->compress_method == COMPRESS_ZLIB) {
2213                 size_t size = MAX_PENDING_SIZE * 2;
2214                 u8 *tmp;
2215
2216                 tmp = malloc(MAX_PENDING_SIZE * 2);
2217                 if (!tmp) {
2218                         free(buffer);
2219                         return -ENOMEM;
2220                 }
2221                 ret = uncompress(tmp, (unsigned long *)&size,
2222                                  buffer, le32_to_cpu(item->size));
2223                 if (ret != Z_OK) {
2224                         fprintf(stderr, "Error decompressing %d\n", ret);
2225                         free(buffer);
2226                         free(tmp);
2227                         return -EIO;
2228                 }
2229                 free(buffer);
2230                 buffer = tmp;
2231         }
2232
2233         super = (struct btrfs_super_block *)buffer;
2234         chunk_root_bytenr = btrfs_super_chunk_root(super);
2235         mdres->leafsize = btrfs_super_leafsize(super);
2236         memcpy(mdres->fsid, super->fsid, BTRFS_FSID_SIZE);
2237         memcpy(mdres->uuid, super->dev_item.uuid,
2238                        BTRFS_UUID_SIZE);
2239         mdres->devid = le64_to_cpu(super->dev_item.devid);
2240         free(buffer);
2241
2242         return search_for_chunk_blocks(mdres, chunk_root_bytenr, 0);
2243 }
2244
2245 static int __restore_metadump(const char *input, FILE *out, int old_restore,
2246                               int num_threads, int fixup_offset,
2247                               const char *target, int multi_devices)
2248 {
2249         struct meta_cluster *cluster = NULL;
2250         struct meta_cluster_header *header;
2251         struct mdrestore_struct mdrestore;
2252         struct btrfs_fs_info *info = NULL;
2253         u64 bytenr = 0;
2254         FILE *in = NULL;
2255         int ret = 0;
2256
2257         if (!strcmp(input, "-")) {
2258                 in = stdin;
2259         } else {
2260                 in = fopen(input, "r");
2261                 if (!in) {
2262                         perror("unable to open metadump image");
2263                         return 1;
2264                 }
2265         }
2266
2267         /* NOTE: open with write mode */
2268         if (fixup_offset) {
2269                 BUG_ON(!target);
2270                 info = open_ctree_fs_info(target, 0, 0,
2271                                           OPEN_CTREE_WRITES |
2272                                           OPEN_CTREE_RESTORE |
2273                                           OPEN_CTREE_PARTIAL);
2274                 if (!info) {
2275                         fprintf(stderr, "%s: open ctree failed\n", __func__);
2276                         ret = -EIO;
2277                         goto failed_open;
2278                 }
2279         }
2280
2281         cluster = malloc(BLOCK_SIZE);
2282         if (!cluster) {
2283                 fprintf(stderr, "Error allocating cluster\n");
2284                 ret = -ENOMEM;
2285                 goto failed_info;
2286         }
2287
2288         ret = mdrestore_init(&mdrestore, in, out, old_restore, num_threads,
2289                              fixup_offset, info, multi_devices);
2290         if (ret) {
2291                 fprintf(stderr, "Error initing mdrestore %d\n", ret);
2292                 goto failed_cluster;
2293         }
2294
2295         if (!multi_devices) {
2296                 ret = build_chunk_tree(&mdrestore, cluster);
2297                 if (ret)
2298                         goto out;
2299         }
2300
2301         if (in != stdin && fseek(in, 0, SEEK_SET)) {
2302                 fprintf(stderr, "Error seeking %d\n", errno);
2303                 goto out;
2304         }
2305
2306         while (1) {
2307                 ret = fread(cluster, BLOCK_SIZE, 1, in);
2308                 if (!ret)
2309                         break;
2310
2311                 header = &cluster->header;
2312                 if (le64_to_cpu(header->magic) != HEADER_MAGIC ||
2313                     le64_to_cpu(header->bytenr) != bytenr) {
2314                         fprintf(stderr, "bad header in metadump image\n");
2315                         ret = -EIO;
2316                         break;
2317                 }
2318                 ret = add_cluster(cluster, &mdrestore, &bytenr);
2319                 if (ret) {
2320                         fprintf(stderr, "Error adding cluster\n");
2321                         break;
2322                 }
2323
2324                 ret = wait_for_worker(&mdrestore);
2325                 if (ret) {
2326                         fprintf(stderr, "One of the threads errored out %d\n",
2327                                 ret);
2328                         break;
2329                 }
2330         }
2331 out:
2332         mdrestore_destroy(&mdrestore);
2333 failed_cluster:
2334         free(cluster);
2335 failed_info:
2336         if (fixup_offset && info)
2337                 close_ctree(info->chunk_root);
2338 failed_open:
2339         if (in != stdin)
2340                 fclose(in);
2341         return ret;
2342 }
2343
2344 static int restore_metadump(const char *input, FILE *out, int old_restore,
2345                             int num_threads, int multi_devices)
2346 {
2347         return __restore_metadump(input, out, old_restore, num_threads, 0, NULL,
2348                                   multi_devices);
2349 }
2350
2351 static int fixup_metadump(const char *input, FILE *out, int num_threads,
2352                           const char *target)
2353 {
2354         return __restore_metadump(input, out, 0, num_threads, 1, target, 1);
2355 }
2356
2357 static int update_disk_super_on_device(struct btrfs_fs_info *info,
2358                                        const char *other_dev, u64 cur_devid)
2359 {
2360         struct btrfs_key key;
2361         struct extent_buffer *leaf;
2362         struct btrfs_path path;
2363         struct btrfs_dev_item *dev_item;
2364         struct btrfs_super_block *disk_super;
2365         char dev_uuid[BTRFS_UUID_SIZE];
2366         char fs_uuid[BTRFS_UUID_SIZE];
2367         u64 devid, type, io_align, io_width;
2368         u64 sector_size, total_bytes, bytes_used;
2369         char *buf;
2370         int fp;
2371         int ret;
2372
2373         key.objectid = BTRFS_DEV_ITEMS_OBJECTID;
2374         key.type = BTRFS_DEV_ITEM_KEY;
2375         key.offset = cur_devid;
2376
2377         btrfs_init_path(&path);
2378         ret = btrfs_search_slot(NULL, info->chunk_root, &key, &path, 0, 0);
2379         if (ret) {
2380                 fprintf(stderr, "search key fails\n");
2381                 exit(1);
2382         }
2383
2384         leaf = path.nodes[0];
2385         dev_item = btrfs_item_ptr(leaf, path.slots[0],
2386                                   struct btrfs_dev_item);
2387
2388         devid = btrfs_device_id(leaf, dev_item);
2389         if (devid != cur_devid) {
2390                 printk("devid %llu mismatch with %llu\n", devid, cur_devid);
2391                 exit(1);
2392         }
2393
2394         type = btrfs_device_type(leaf, dev_item);
2395         io_align = btrfs_device_io_align(leaf, dev_item);
2396         io_width = btrfs_device_io_width(leaf, dev_item);
2397         sector_size = btrfs_device_sector_size(leaf, dev_item);
2398         total_bytes = btrfs_device_total_bytes(leaf, dev_item);
2399         bytes_used = btrfs_device_bytes_used(leaf, dev_item);
2400         read_extent_buffer(leaf, dev_uuid, (unsigned long)btrfs_device_uuid(dev_item), BTRFS_UUID_SIZE);
2401         read_extent_buffer(leaf, fs_uuid, (unsigned long)btrfs_device_fsid(dev_item), BTRFS_UUID_SIZE);
2402
2403         btrfs_release_path(&path);
2404
2405         printk("update disk super on %s devid=%llu\n", other_dev, devid);
2406
2407         /* update other devices' super block */
2408         fp = open(other_dev, O_CREAT | O_RDWR, 0600);
2409         if (fp < 0) {
2410                 fprintf(stderr, "could not open %s\n", other_dev);
2411                 exit(1);
2412         }
2413
2414         buf = malloc(BTRFS_SUPER_INFO_SIZE);
2415         if (!buf) {
2416                 ret = -ENOMEM;
2417                 exit(1);
2418         }
2419
2420         memcpy(buf, info->super_copy, BTRFS_SUPER_INFO_SIZE);
2421
2422         disk_super = (struct btrfs_super_block *)buf;
2423         dev_item = &disk_super->dev_item;
2424
2425         btrfs_set_stack_device_type(dev_item, type);
2426         btrfs_set_stack_device_id(dev_item, devid);
2427         btrfs_set_stack_device_total_bytes(dev_item, total_bytes);
2428         btrfs_set_stack_device_bytes_used(dev_item, bytes_used);
2429         btrfs_set_stack_device_io_align(dev_item, io_align);
2430         btrfs_set_stack_device_io_width(dev_item, io_width);
2431         btrfs_set_stack_device_sector_size(dev_item, sector_size);
2432         memcpy(dev_item->uuid, dev_uuid, BTRFS_UUID_SIZE);
2433         memcpy(dev_item->fsid, fs_uuid, BTRFS_UUID_SIZE);
2434         csum_block((u8 *)buf, BTRFS_SUPER_INFO_SIZE);
2435
2436         ret = pwrite64(fp, buf, BTRFS_SUPER_INFO_SIZE, BTRFS_SUPER_INFO_OFFSET);
2437         if (ret != BTRFS_SUPER_INFO_SIZE) {
2438                 ret = -EIO;
2439                 goto out;
2440         }
2441
2442         write_backup_supers(fp, (u8 *)buf);
2443
2444 out:
2445         free(buf);
2446         close(fp);
2447         return 0;
2448 }
2449
2450 static void print_usage(void)
2451 {
2452         fprintf(stderr, "usage: btrfs-image [options] source target\n");
2453         fprintf(stderr, "\t-r      \trestore metadump image\n");
2454         fprintf(stderr, "\t-c value\tcompression level (0 ~ 9)\n");
2455         fprintf(stderr, "\t-t value\tnumber of threads (1 ~ 32)\n");
2456         fprintf(stderr, "\t-o      \tdon't mess with the chunk tree when restoring\n");
2457         fprintf(stderr, "\t-s      \tsanitize file names, use once to just use garbage, use twice if you want crc collisions\n");
2458         fprintf(stderr, "\t-w      \twalk all trees instead of using extent tree, do this if your extent tree is broken\n");
2459         exit(1);
2460 }
2461
2462 int main(int argc, char *argv[])
2463 {
2464         char *source;
2465         char *target;
2466         int num_threads = 0;
2467         int compress_level = 0;
2468         int create = 1;
2469         int old_restore = 0;
2470         int walk_trees = 0;
2471         int multi_devices = 0;
2472         int ret;
2473         int sanitize = 0;
2474         int dev_cnt = 0;
2475         FILE *out;
2476
2477         while (1) {
2478                 int c = getopt(argc, argv, "rc:t:oswm");
2479                 if (c < 0)
2480                         break;
2481                 switch (c) {
2482                 case 'r':
2483                         create = 0;
2484                         break;
2485                 case 't':
2486                         num_threads = atoi(optarg);
2487                         if (num_threads <= 0 || num_threads > 32)
2488                                 print_usage();
2489                         break;
2490                 case 'c':
2491                         compress_level = atoi(optarg);
2492                         if (compress_level < 0 || compress_level > 9)
2493                                 print_usage();
2494                         break;
2495                 case 'o':
2496                         old_restore = 1;
2497                         break;
2498                 case 's':
2499                         sanitize++;
2500                         break;
2501                 case 'w':
2502                         walk_trees = 1;
2503                         break;
2504                 case 'm':
2505                         create = 0;
2506                         multi_devices = 1;
2507                         break;
2508                 default:
2509                         print_usage();
2510                 }
2511         }
2512
2513         if ((old_restore) && create)
2514                 print_usage();
2515
2516         argc = argc - optind;
2517         dev_cnt = argc - 1;
2518
2519         if (multi_devices && dev_cnt < 2)
2520                 print_usage();
2521         if (!multi_devices && dev_cnt != 1)
2522                 print_usage();
2523
2524         source = argv[optind];
2525         target = argv[optind + 1];
2526
2527         if (create && !strcmp(target, "-")) {
2528                 out = stdout;
2529         } else {
2530                 out = fopen(target, "w+");
2531                 if (!out) {
2532                         perror("unable to create target file");
2533                         exit(1);
2534                 }
2535         }
2536
2537         if (num_threads == 0 && compress_level > 0) {
2538                 num_threads = sysconf(_SC_NPROCESSORS_ONLN);
2539                 if (num_threads <= 0)
2540                         num_threads = 1;
2541         }
2542
2543         if (create)
2544                 ret = create_metadump(source, out, num_threads,
2545                                       compress_level, sanitize, walk_trees);
2546         else
2547                 ret = restore_metadump(source, out, old_restore, 1,
2548                                        multi_devices);
2549         if (ret) {
2550                 printk("%s failed (%s)\n", (create) ? "create" : "restore",
2551                        strerror(errno));
2552                 goto out;
2553         }
2554
2555          /* extended support for multiple devices */
2556         if (!create && multi_devices) {
2557                 struct btrfs_fs_info *info;
2558                 u64 total_devs;
2559                 int i;
2560
2561                 info = open_ctree_fs_info(target, 0, 0,
2562                                           OPEN_CTREE_PARTIAL |
2563                                           OPEN_CTREE_RESTORE);
2564                 if (!info) {
2565                         int e = errno;
2566                         fprintf(stderr, "unable to open %s error = %s\n",
2567                                 target, strerror(e));
2568                         return 1;
2569                 }
2570
2571                 total_devs = btrfs_super_num_devices(info->super_copy);
2572                 if (total_devs != dev_cnt) {
2573                         printk("it needs %llu devices but has only %d\n",
2574                                 total_devs, dev_cnt);
2575                         close_ctree(info->chunk_root);
2576                         goto out;
2577                 }
2578
2579                 /* update super block on other disks */
2580                 for (i = 2; i <= dev_cnt; i++) {
2581                         ret = update_disk_super_on_device(info,
2582                                         argv[optind + i], (u64)i);
2583                         if (ret) {
2584                                 printk("update disk super failed devid=%d (error=%d)\n",
2585                                         i, ret);
2586                                 close_ctree(info->chunk_root);
2587                                 exit(1);
2588                         }
2589                 }
2590
2591                 close_ctree(info->chunk_root);
2592
2593                 /* fix metadata block to map correct chunk */
2594                 ret = fixup_metadump(source, out, 1, target);
2595                 if (ret) {
2596                         fprintf(stderr, "fix metadump failed (error=%d)\n",
2597                                 ret);
2598                         exit(1);
2599                 }
2600         }
2601
2602 out:
2603         if (out == stdout)
2604                 fflush(out);
2605         else
2606                 fclose(out);
2607
2608         return !!ret;
2609 }