move chunk reconstruction to btrfsinspect, add --mappings

2 files changed, 101 insertions, 1 deletions

diff --git a/cmd/btrfs-rec/inspect_recoverchunks.go b/cmd/btrfs-rec/inspect_recoverchunks.go
new file mode 100644
index 0000000..bc32248
--- /dev/null
+++ b/cmd/btrfs-rec/inspect_recoverchunks.go
@@ -0,0 +1,75 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/datawire/dlib/dlog"
+	"github.com/datawire/ocibuild/pkg/cliutil"
+	"github.com/spf13/cobra"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect"
+)
+
+func init() {
+	inspectors = append(inspectors, subcommand{
+		Command: cobra.Command{
+			Use:   "recover-chunks",
+			Short: "Rebuild broken chunk/dev-extent/blockgroup trees",
+			Long: "" +
+				"The rebuilt information is printed as JSON on stdout, and can\n" +
+				"be loaded by the --mappings flag.\n" +
+				"\n" +
+				"This is very similar to `btrfs rescue chunk-recover`, but (1)\n" +
+				"does a better job, (2) is less buggy, and (3) doesn't actually\n" +
+				"write the info back to the filesystem; instead writing it\n" +
+				"out-of-band to stdout.",
+			Args: cliutil.WrapPositionalArgs(cobra.NoArgs),
+		},
+		RunE: func(fs *btrfs.FS, cmd *cobra.Command, _ []string) error {
+			ctx := cmd.Context()
+
+			dlog.Info(ctx, "Reading superblock...")
+			superblock, err := fs.Superblock()
+			if err != nil {
+				return err
+			}
+
+			for _, dev := range fs.LV.PhysicalVolumes() {
+				dlog.Infof(ctx, "dev[%q] Scanning for unreachable nodes...", dev.Name())
+				devResult, err := btrfsinspect.ScanOneDev(ctx, dev, superblock.Data)
+				if err != nil {
+					return err
+				}
+
+				dlog.Infof(ctx, "dev[%q] Re-inserting lost+found mappings...", dev.Name())
+				devResult.AddToLV(ctx, fs, dev)
+			}
+
+			dlog.Infof(ctx, "Writing reconstructed mappings to stdout...")
+
+			mappings := fs.LV.Mappings()
+			_, _ = io.WriteString(os.Stdout, "{\n  \"Mappings\": [\n")
+			for i, mapping := range mappings {
+				suffix := ","
+				if i == len(mappings)-1 {
+					suffix = ""
+				}
+				bs, err := json.Marshal(mapping)
+				if err != nil {
+					return err
+				}
+				fmt.Printf("    %s%s\n", bs, suffix)
+			}
+			_, _ = io.WriteString(os.Stdout, "  ]\n}\n")
+			return nil
+		},
+	})
+}
diff --git a/cmd/btrfs-rec/main.go b/cmd/btrfs-rec/main.go
index b68bdd8..b17c8ec 100644
--- a/cmd/btrfs-rec/main.go
+++ b/cmd/btrfs-rec/main.go
@@ -6,6 +6,7 @@ package main
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"os"
 
@@ -17,6 +18,7 @@ import (
 	"github.com/spf13/pflag"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsutil"
 )
 
@@ -45,6 +47,7 @@ func main() {
 		Level: logrus.InfoLevel,
 	}
 	var pvsFlag []string
+	var mappingsFlag string
 
 	argparser := &cobra.Command{
 		Use:   "btrfs-rec {[flags]|SUBCOMMAND}",
@@ -70,6 +73,10 @@ func main() {
 	if err := argparser.MarkPersistentFlagRequired("pv"); err != nil {
 		panic(err)
 	}
+	argparser.PersistentFlags().StringVar(&mappingsFlag, "mappings", "", "load chunk/dev-extent/blockgroup data from external JSON file `mappings.json`")
+	if err := argparser.MarkPersistentFlagFilename("mappings"); err != nil {
+		panic(err)
+	}
 
 	var openFlag int = os.O_RDONLY
 
@@ -121,7 +128,7 @@ func main() {
 							err = _err
 						}
 					}
-					fs, err := btrfsutil.Open(openFlag, pvsFlag...)
+					fs, err := btrfsutil.Open(ctx, openFlag, pvsFlag...)
 					if err != nil {
 						return err
 					}
@@ -129,6 +136,24 @@ func main() {
 						maybeSetErr(fs.Close())
 					}()
 
+					if mappingsFlag != "" {
+						bs, err := os.ReadFile(mappingsFlag)
+						if err != nil {
+							return err
+						}
+						var mappingsJSON struct {
+							Mappings []btrfsvol.Mapping
+						}
+						if err := json.Unmarshal(bs, &mappingsJSON); err != nil {
+							return err
+						}
+						for _, mapping := range mappingsJSON.Mappings {
+							if err := fs.LV.AddMapping(mapping); err != nil {
+								return err
+							}
+						}
+					}
+
 					cmd.SetContext(ctx)
 					return runE(fs, cmd, args)
 				})