lnd.xprv/channeldb/migtest/raw_db.go

package migtest

import (
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"strings"

	"github.com/coreos/bbolt"
)

// DumpDB dumps go code describing the contents of the database to stdout. This
// function is only intended for use during development.
//
// Example output:
//
// map[string]interface{}{
// 	hex("1234"): map[string]interface{}{
// 		"human-readable": hex("102030"),
// 		hex("1111"): hex("5783492373"),
// 	},
// }
func DumpDB(tx *bbolt.Tx, rootKey []byte) error {
	bucket := tx.Bucket(rootKey)
	if bucket == nil {
		return fmt.Errorf("bucket %v not found", string(rootKey))
	}

	return dumpBucket(bucket)
}

func dumpBucket(bucket *bbolt.Bucket) error {
	fmt.Printf("map[string]interface{} {\n")
	err := bucket.ForEach(func(k, v []byte) error {
		key := toString(k)
		fmt.Printf("%v: ", key)

		subBucket := bucket.Bucket(k)
		if subBucket != nil {
			err := dumpBucket(subBucket)
			if err != nil {
				return err
			}
		} else {
			fmt.Print(toHex(v))
		}
		fmt.Printf(",\n")

		return nil
	})
	if err != nil {
		return err
	}
	fmt.Printf("}")

	return nil
}

// RestoreDB primes the database with the given data set.
func RestoreDB(tx *bbolt.Tx, rootKey []byte, data map[string]interface{}) error {
	bucket, err := tx.CreateBucket(rootKey)
	if err != nil {
		return err
	}

	return restoreDB(bucket, data)
}

func restoreDB(bucket *bbolt.Bucket, data map[string]interface{}) error {
	for k, v := range data {
		key := []byte(k)

		switch value := v.(type) {

		// Key contains value.
		case string:
			err := bucket.Put(key, []byte(value))
			if err != nil {
				return err
			}

		// Key contains a sub-bucket.
		case map[string]interface{}:
			subBucket, err := bucket.CreateBucket(key)
			if err != nil {
				return err
			}

			if err := restoreDB(subBucket, value); err != nil {
				return err
			}

		default:
			return errors.New("invalid type")
		}
	}

	return nil
}

// VerifyDB verifies the database against the given data set.
func VerifyDB(tx *bbolt.Tx, rootKey []byte, data map[string]interface{}) error {
	bucket := tx.Bucket(rootKey)
	if bucket == nil {
		return fmt.Errorf("bucket %v not found", string(rootKey))
	}

	return verifyDB(bucket, data)
}

func verifyDB(bucket *bbolt.Bucket, data map[string]interface{}) error {
	for k, v := range data {
		key := []byte(k)

		switch value := v.(type) {

		// Key contains value.
		case string:
			expectedValue := []byte(value)
			dbValue := bucket.Get(key)

			if !bytes.Equal(dbValue, expectedValue) {
				return errors.New("value mismatch")
			}

		// Key contains a sub-bucket.
		case map[string]interface{}:
			subBucket := bucket.Bucket(key)
			if subBucket == nil {
				return fmt.Errorf("bucket %v not found", k)
			}

			err := verifyDB(subBucket, value)
			if err != nil {
				return err
			}

		default:
			return errors.New("invalid type")
		}
	}

	keyCount := 0
	err := bucket.ForEach(func(k, v []byte) error {
		keyCount++
		return nil
	})
	if err != nil {
		return err
	}
	if keyCount != len(data) {
		return errors.New("unexpected keys in database")
	}

	return nil
}

func toHex(v []byte) string {
	if len(v) == 0 {
		return "nil"
	}

	return "hex(\"" + hex.EncodeToString(v) + "\")"
}

func toString(v []byte) string {
	readableChars := "abcdefghijklmnopqrstuvwxyz0123456789-"

	for _, c := range v {
		if !strings.Contains(readableChars, string(c)) {
			return toHex(v)
		}
	}

	return "\"" + string(v) + "\""
}

// Hex is a test helper function to convert readable hex arrays to raw byte
// strings.
func Hex(value string) string {
	b, err := hex.DecodeString(value)
	if err != nil {
		panic(err)
	}
	return string(b)
}