cmd: add meta page related surgery commands

Two command:
- surgery meta validate
- surgery meta update

Signed-off-by: Benjamin Wang <[email protected]>

cmd/bbolt/command_surgery.go

@@ -28,6 +28,7 @@ func newSurgeryCobraCommand() *cobra.Command {
 	surgeryCmd.AddCommand(newSurgeryClearPageCommand())
 	surgeryCmd.AddCommand(newSurgeryClearPageElementsCommand())
 	surgeryCmd.AddCommand(newSurgeryFreelistCommand())
+	surgeryCmd.AddCommand(newSurgeryMetaCommand())
 
 	return surgeryCmd
 }
@@ -311,15 +312,23 @@ func surgeryClearPageElementFunc(srcDBPath string, cfg surgeryClearPageElementsO
 }
 
 func readMetaPage(path string) (*common.Meta, error) {
-	_, activeMetaPageId, err := guts_cli.GetRootPage(path)
+	pageSize, _, err := guts_cli.ReadPageAndHWMSize(path)
 	if err != nil {
-		return nil, fmt.Errorf("read root page failed: %w", err)
+		return nil, fmt.Errorf("read Page size failed: %w", err)
 	}
-	_, buf, err := guts_cli.ReadPage(path, uint64(activeMetaPageId))
-	if err != nil {
-		return nil, fmt.Errorf("read active mage page failed: %w", err)
+
+	m := make([]*common.Meta, 2)
+	for i := 0; i < 2; i++ {
+		m[i], _, err = ReadMetaPageAt(path, uint32(i), uint32(pageSize))
+		if err != nil {
+			return nil, fmt.Errorf("read meta page %d failed: %w", i, err)
+		}
+	}
+
+	if m[0].Txid() > m[1].Txid() {
+		return m[0], nil
 	}
-	return common.LoadPageMeta(buf), nil
+	return m[1], nil
 }
 
 func checkSourceDBPath(srcPath string) (os.FileInfo, error) {

cmd/bbolt/command_surgery_meta.go

@@ -0,0 +1,280 @@
+package main
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"os"
+	"strconv"
+	"strings"
+
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+
+	"go.etcd.io/bbolt/internal/common"
+)
+
+func newSurgeryMetaCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "meta <subcommand>",
+		Short: "meta page related surgery commands",
+	}
+
+	cmd.AddCommand(newSurgeryMetaValidateCommand())
+	cmd.AddCommand(newSurgeryMetaUpdateCommand())
+
+	return cmd
+}
+
+func newSurgeryMetaValidateCommand() *cobra.Command {
+	metaValidateCmd := &cobra.Command{
+		Use:   "validate <bbolt-file> [options]",
+		Short: "Validate both meta pages",
+		Args: func(cmd *cobra.Command, args []string) error {
+			if len(args) == 0 {
+				return errors.New("db file path not provided")
+			}
+			if len(args) > 1 {
+				return errors.New("too many arguments")
+			}
+			return nil
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return surgeryMetaValidateFunc(args[0])
+		},
+	}
+	return metaValidateCmd
+}
+
+func surgeryMetaValidateFunc(srcDBPath string) error {
+	if _, err := checkSourceDBPath(srcDBPath); err != nil {
+		return err
+	}
+
+	var pageSize uint32
+
+	for i := 0; i <= 1; i++ {
+		m, _, err := ReadMetaPageAt(srcDBPath, uint32(i), pageSize)
+		if err != nil {
+			return fmt.Errorf("read meta page %d failed: %w", i, err)
+		}
+		if mValidateErr := m.Validate(); mValidateErr != nil {
+			fmt.Fprintf(os.Stdout, "WARNING: The meta page %d isn't valid: %v!\n", i, mValidateErr)
+		} else {
+			fmt.Fprintf(os.Stdout, "The meta page %d is valid!\n", i)
+		}
+
+		pageSize = m.PageSize()
+	}
+
+	return nil
+}
+
+type surgeryMetaUpdateOptions struct {
+	surgeryBaseOptions
+	fields     []string
+	metaPageId uint32
+}
+
+var allowedMetaUpdateFields = map[string]struct{}{
+	"pageSize": {},
+	"root":     {},
+	"freelist": {},
+	"pgid":     {},
+}
+
+// AddFlags sets the flags for `meta update` command.
+// Example: --fields root:16,freelist:8 --fields pgid:128 --fields txid:1234
+// Result: []string{"root:16", "freelist:8", "pgid:128", "txid:1234"}
+func (o *surgeryMetaUpdateOptions) AddFlags(fs *pflag.FlagSet) {
+	o.surgeryBaseOptions.AddFlags(fs)
+	fs.StringSliceVarP(&o.fields, "fields", "", o.fields, "comma separated list of fields (supported fields: pageSize, root, freelist, pgid and txid) to be updated, and each item is a colon-separated key-value pair")
+	fs.Uint32VarP(&o.metaPageId, "meta-page", "", o.metaPageId, "the meta page ID to operate on, valid values are 0 and 1")
+}
+
+func (o *surgeryMetaUpdateOptions) Validate() error {
+	if err := o.surgeryBaseOptions.Validate(); err != nil {
+		return err
+	}
+
+	if o.metaPageId > 1 {
+		return fmt.Errorf("invalid meta page id: %d", o.metaPageId)
+	}
+
+	for _, field := range o.fields {
+		kv := strings.SplitN(field, ":", 2)
+		if len(kv) != 2 {
+			return fmt.Errorf("invalid key-value pair: %s", field)
+		}
+
+		if _, ok := allowedMetaUpdateFields[kv[0]]; !ok {
+			return fmt.Errorf("field %q isn't allowed to be updated", kv[0])
+		}
+
+		if _, err := strconv.ParseUint(kv[1], 10, 64); err != nil {
+			return fmt.Errorf("invalid value %q for field %q", kv[1], kv[0])
+		}
+	}
+
+	return nil
+}
+
+func newSurgeryMetaUpdateCommand() *cobra.Command {
+	var o surgeryMetaUpdateOptions
+	metaUpdateCmd := &cobra.Command{
+		Use:   "update <bbolt-file> [options]",
+		Short: "Update fields in meta pages",
+		Args: func(cmd *cobra.Command, args []string) error {
+			if len(args) == 0 {
+				return errors.New("db file path not provided")
+			}
+			if len(args) > 1 {
+				return errors.New("too many arguments")
+			}
+			return nil
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			if err := o.Validate(); err != nil {
+				return err
+			}
+			return surgeryMetaUpdateFunc(args[0], o)
+		},
+	}
+	o.AddFlags(metaUpdateCmd.Flags())
+	return metaUpdateCmd
+}
+
+func surgeryMetaUpdateFunc(srcDBPath string, cfg surgeryMetaUpdateOptions) error {
+	if _, err := checkSourceDBPath(srcDBPath); err != nil {
+		return err
+	}
+
+	if err := common.CopyFile(srcDBPath, cfg.outputDBFilePath); err != nil {
+		return fmt.Errorf("[meta update] copy file failed: %w", err)
+	}
+
+	// read the page size from the first meta page if we want to edit the second meta page.
+	var pageSize uint32
+	if cfg.metaPageId == 1 {
+		m0, _, err := ReadMetaPageAt(cfg.outputDBFilePath, 0, pageSize)
+		if err != nil {
+			return fmt.Errorf("read the first meta page failed: %w", err)
+		}
+		pageSize = m0.PageSize()
+	}
+
+	// update the specified meta page
+	m, buf, err := ReadMetaPageAt(cfg.outputDBFilePath, cfg.metaPageId, pageSize)
+	if err != nil {
+		return fmt.Errorf("read meta page %d failed: %w", cfg.metaPageId, err)
+	}
+	mChanged := updateMetaField(m, parseFields(cfg.fields))
+	if mChanged {
+		if err := writeMetaPageAt(cfg.outputDBFilePath, buf, cfg.metaPageId, pageSize); err != nil {
+			return fmt.Errorf("[meta update] write meta page %d failed: %w", cfg.metaPageId, err)
+		}
+	}
+
+	if cfg.metaPageId == 1 && pageSize != m.PageSize() {
+		fmt.Fprintf(os.Stdout, "WARNING: The page size (%d) in the first meta page doesn't match the second meta page (%d)\n", pageSize, m.PageSize())
+	}
+
+	// Display results
+	if !mChanged {
+		fmt.Fprintln(os.Stdout, "Nothing changed!")
+	}
+
+	if mChanged {
+		fmt.Fprintf(os.Stdout, "The meta page %d has been updated!\n", cfg.metaPageId)
+	}
+
+	return nil
+}
+
+func parseFields(fields []string) map[string]uint64 {
+	fieldsMap := make(map[string]uint64)
+	for _, field := range fields {
+		kv := strings.SplitN(field, ":", 2)
+		val, _ := strconv.ParseUint(kv[1], 10, 64)
+		fieldsMap[kv[0]] = val
+	}
+	return fieldsMap
+}
+
+func updateMetaField(m *common.Meta, fields map[string]uint64) bool {
+	changed := false
+	for key, val := range fields {
+		switch key {
+		case "pageSize":
+			m.SetPageSize(uint32(val))
+		case "root":
+			m.SetRootBucket(common.NewInBucket(common.Pgid(val), 0))
+		case "freelist":
+			m.SetFreelist(common.Pgid(val))
+		case "pgid":
+			m.SetPgid(common.Pgid(val))
+		}
+
+		changed = true
+	}
+
+	if m.Magic() != common.Magic {
+		m.SetMagic(common.Magic)
+		changed = true
+	}
+	if m.Version() != common.Version {
+		m.SetVersion(common.Version)
+		changed = true
+	}
+	if m.Flags() != common.MetaPageFlag {
+		m.SetFlags(common.MetaPageFlag)
+		changed = true
+	}
+
+	newChecksum := m.Sum64()
+	if m.Checksum() != newChecksum {
+		m.SetChecksum(newChecksum)
+		changed = true
+	}
+
+	return changed
+}
+
+func ReadMetaPageAt(dbPath string, metaPageId uint32, pageSize uint32) (*common.Meta, []byte, error) {
+	if metaPageId > 1 {
+		return nil, nil, fmt.Errorf("invalid metaPageId: %d", metaPageId)
+	}
+
+	f, err := os.OpenFile(dbPath, os.O_RDONLY, 0444)
+	if err != nil {
+		return nil, nil, err
+	}
+	defer f.Close()
+
+	buf := make([]byte, 1024)
+	n, err := f.ReadAt(buf, int64(metaPageId*pageSize))
+	if n == len(buf) && (err == nil || err == io.EOF) {
+		return common.LoadPageMeta(buf), buf, nil
+	}
+
+	return nil, nil, err
+}
+
+func writeMetaPageAt(dbPath string, buf []byte, metaPageId uint32, pageSize uint32) error {
+	if metaPageId > 1 {
+		return fmt.Errorf("invalid metaPageId: %d", metaPageId)
+	}
+
+	f, err := os.OpenFile(dbPath, os.O_RDWR, 0666)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	n, err := f.WriteAt(buf, int64(metaPageId*pageSize))
+	if n == len(buf) && (err == nil || err == io.EOF) {
+		return nil
+	}
+
+	return err
+}