add 'bbolt surgery revert-meta-page' command

Signed-off-by: Benjamin Wang <wachao@vmware.com>

cmd/bbolt/main.go

@@ -130,6 +130,8 @@ func (m *Main) Run(args ...string) error {
 		return newPagesCommand(m).Run(args[1:]...)
 	case "stats":
 		return newStatsCommand(m).Run(args[1:]...)
+	case "surgery":
+		return newSurgeryCommand(m).Run(args[1:]...)
 	default:
 		return ErrUnknownCommand
 	}
@@ -159,6 +161,7 @@ The commands are:
     pages       print list of pages with their types
     page-item   print the key and value of a page item.
     stats       iterate over all pages and generate usage stats
+    surgery     perform surgery on bbolt database
 
 Use "bbolt [command] -h" for more information about a command.
 `, "\n")

cmd/bbolt/surgery_commands.go

@@ -0,0 +1,156 @@
+package main
+
+import (
+	"errors"
+	"flag"
+	"fmt"
+	"io"
+	"os"
+	"strings"
+
+	"go.etcd.io/bbolt/internal/surgeon"
+)
+
+// SurgeryCommand represents the "surgery" command execution.
+type SurgeryCommand struct {
+	baseCommand
+}
+
+// newSurgeryCommand returns a SurgeryCommand.
+func newSurgeryCommand(m *Main) *SurgeryCommand {
+	c := &SurgeryCommand{}
+	c.baseCommand = m.baseCommand
+	return c
+}
+
+// Run executes the `surgery` program.
+func (cmd *SurgeryCommand) Run(args ...string) error {
+	// Require a command at the beginning.
+	if len(args) == 0 || strings.HasPrefix(args[0], "-") {
+		fmt.Fprintln(cmd.Stderr, cmd.Usage())
+		return ErrUsage
+	}
+
+	// Execute command.
+	switch args[0] {
+	case "help":
+		fmt.Fprintln(cmd.Stderr, cmd.Usage())
+		return ErrUsage
+	case "revert-meta-page":
+		return newRevertMetaPageCommand(cmd).Run(args[1:]...)
+	default:
+		return ErrUnknownCommand
+	}
+}
+
+// Usage returns the help message.
+func (cmd *SurgeryCommand) Usage() string {
+	return strings.TrimLeft(`
+Surgery is a command for performing low level update on bbolt databases.
+
+Usage:
+
+	bbolt surgery command [arguments]
+
+The commands are:
+
+    help                   print this screen
+    revert-meta-page       revert the meta page change made by the last transaction
+
+Use "bbolt surgery [command] -h" for more information about a command.
+`, "\n")
+}
+
+// RevertMetaPageCommand represents the "surgery revert-meta-page" command execution.
+type RevertMetaPageCommand struct {
+	baseCommand
+
+	SrcPath string
+	DstPath string
+}
+
+// newRevertMetaPageCommand returns a RevertMetaPageCommand.
+func newRevertMetaPageCommand(m *SurgeryCommand) *RevertMetaPageCommand {
+	c := &RevertMetaPageCommand{}
+	c.baseCommand = m.baseCommand
+	return c
+}
+
+// Run executes the command.
+func (cmd *RevertMetaPageCommand) Run(args ...string) error {
+	// Parse flags.
+	fs := flag.NewFlagSet("", flag.ContinueOnError)
+	fs.SetOutput(io.Discard)
+	fs.StringVar(&cmd.DstPath, "o", "", "")
+	if err := fs.Parse(args); err == flag.ErrHelp {
+		fmt.Fprintln(cmd.Stderr, cmd.Usage())
+		return ErrUsage
+	} else if err != nil {
+		return err
+	} else if cmd.DstPath == "" {
+		return errors.New("output file required")
+	}
+
+	// Require database paths.
+	cmd.SrcPath = fs.Arg(0)
+	if cmd.SrcPath == "" {
+		return ErrPathRequired
+	}
+
+	// Ensure source file exists.
+	fi, err := os.Stat(cmd.SrcPath)
+	if os.IsNotExist(err) {
+		return ErrFileNotFound
+	} else if err != nil {
+		return err
+	}
+	initialSize := fi.Size()
+
+	// Ensure output file not exist.
+	_, err = os.Stat(cmd.DstPath)
+	if err == nil {
+		return fmt.Errorf("output file %q already exists", cmd.DstPath)
+	} else if !os.IsNotExist(err) {
+		return err
+	}
+
+	// Copy database from SrcPath to DstPath
+	srcDB, err := os.Open(cmd.SrcPath)
+	if err != nil {
+		return fmt.Errorf("failed to open source file %q: %w", cmd.SrcPath, err)
+	}
+	defer srcDB.Close()
+	dstDB, err := os.Create(cmd.DstPath)
+	if err != nil {
+		return fmt.Errorf("failed to create output file %q: %w", cmd.DstPath, err)
+	}
+	defer dstDB.Close()
+	written, err := io.Copy(dstDB, srcDB)
+	if err != nil {
+		return fmt.Errorf("failed to copy database file from %q to %q: %w", cmd.SrcPath, cmd.DstPath, err)
+	}
+	if initialSize != written {
+		return fmt.Errorf("the byte copied (%q: %d) isn't equal to the initial db size (%q: %d)", cmd.DstPath, written, cmd.SrcPath, initialSize)
+	}
+
+	// revert the meta page
+	if err = surgeon.RevertMetaPage(cmd.DstPath); err != nil {
+		return err
+	}
+
+	fmt.Fprintln(cmd.Stdout, "The meta page is reverted.")
+	return nil
+}
+
+// Usage returns the help message.
+func (cmd *RevertMetaPageCommand) Usage() string {
+	return strings.TrimLeft(`
+usage: bolt surgery revert-meta-page -o DST SRC
+
+RevertMetaPage copies the database file at SRC to a newly created database
+file at DST. Afterwards, it reverts the meta page on the newly created
+database at DST.
+
+The original database is left untouched.
+`, "\n")
+}

cmd/bbolt/surgery_commands_test.go

@@ -0,0 +1,69 @@
+package main_test
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"go.etcd.io/bbolt/internal/btesting"
+	"go.etcd.io/bbolt/internal/guts_cli"
+)
+
+func TestSurgery_RevertMetaPage(t *testing.T) {
+	pageSize := os.Getpagesize()
+	db := btesting.MustCreateDB(t)
+	srcPath := db.Path()
+
+	srcFile, err := os.Open(srcPath)
+	require.NoError(t, err)
+	defer srcFile.Close()
+
+	// Read both meta0 and meta1 from srcFile
+	srcBuf0, srcBuf1 := readBothMetaPages(t, srcPath, pageSize)
+	meta0Page := guts_cli.LoadPageMeta(srcBuf0)
+	meta1Page := guts_cli.LoadPageMeta(srcBuf1)
+
+	// Get the non-active meta page
+	nonActiveSrcBuf := srcBuf0
+	nonActiveMetaPageId := 0
+	if meta0Page.Txid() > meta1Page.Txid() {
+		nonActiveSrcBuf = srcBuf1
+		nonActiveMetaPageId = 1
+	}
+	t.Logf("non active meta page id: %d", nonActiveMetaPageId)
+
+	// revert the meta page
+	dstPath := filepath.Join(t.TempDir(), "dstdb")
+	m := NewMain()
+	err = m.Run("surgery", "revert-meta-page", "-o", dstPath, srcPath)
+	require.NoError(t, err)
+
+	// read both meta0 and meta1 from dst file
+	dstBuf0, dstBuf1 := readBothMetaPages(t, dstPath, pageSize)
+
+	// check result. Note we should skip the page ID
+	assert.Equal(t, nonActiveSrcBuf[8:], dstBuf0[8:])
+	assert.Equal(t, nonActiveSrcBuf[8:], dstBuf1[8:])
+}
+
+func readBothMetaPages(t *testing.T, filePath string, pageSize int) ([]byte, []byte) {
+	dbFile, err := os.Open(filePath)
+	require.NoError(t, err)
+	defer dbFile.Close()
+
+	buf0 := make([]byte, pageSize)
+	buf1 := make([]byte, pageSize)
+
+	meta0Len, err := dbFile.ReadAt(buf0, 0)
+	require.NoError(t, err)
+	require.Equal(t, pageSize, meta0Len)
+
+	meta1Len, err := dbFile.ReadAt(buf1, int64(pageSize))
+	require.NoError(t, err)
+	require.Equal(t, pageSize, meta1Len)
+
+	return buf0, buf1
+}

internal/guts_cli/guts_cli.go

@@ -64,6 +64,10 @@ func (m *Meta) RootBucket() *Bucket {
 	return &m.root
 }
 
+func (m *Meta) Txid() uint64 {
+	return uint64(m.txid)
+}
+
 func (m *Meta) Print(w io.Writer) {
 	fmt.Fprintf(w, "Version:    %d\n", m.version)
 	fmt.Fprintf(w, "Page Size:  %d bytes\n", m.pageSize)