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