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