From 24a1b8d642c6eaa6453c15debb7a5e58f8b68d73 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Wed, 15 May 2019 13:59:05 +0200
Subject: [PATCH] loopdb: initialize new database with current version

---
 loopdb/meta.go       | 12 ++++----
 loopdb/store.go      | 25 ++++++++++++----
 loopdb/store_test.go | 71 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 97 insertions(+), 11 deletions(-)

diff --git a/loopdb/meta.go b/loopdb/meta.go
index a1d9a18..1a3c8c4 100644
--- a/loopdb/meta.go
+++ b/loopdb/meta.go
@@ -10,7 +10,7 @@ import (
 var (
 	// metaBucket stores all the meta information concerning the state of
 	// the database.
-	metaBucket = []byte("metadata")
+	metaBucketKey = []byte("metadata")
 
 	// dbVersionKey is a boltdb key and it's used for storing/retrieving
 	// current database version.
@@ -41,7 +41,7 @@ func getDBVersion(db *bbolt.DB) (uint32, error) {
 	var version uint32
 
 	err := db.View(func(tx *bbolt.Tx) error {
-		metaBucket := tx.Bucket(metaBucket)
+		metaBucket := tx.Bucket(metaBucketKey)
 		if metaBucket == nil {
 			return errors.New("bucket does not exist")
 		}
@@ -61,11 +61,11 @@ func getDBVersion(db *bbolt.DB) (uint32, error) {
 	return version, nil
 }
 
-// getDBVersion updates the current db version.
+// setDBVersion updates the current db version.
 func setDBVersion(tx *bbolt.Tx, version uint32) error {
-	metaBucket := tx.Bucket(metaBucket)
-	if metaBucket == nil {
-		return errors.New("bucket does not exist")
+	metaBucket, err := tx.CreateBucketIfNotExists(metaBucketKey)
+	if err != nil {
+		return fmt.Errorf("set db version: %v", err)
 	}
 
 	scratch := make([]byte, 4)
diff --git a/loopdb/store.go b/loopdb/store.go
index 2d731b5..2a90ce8 100644
--- a/loopdb/store.go
+++ b/loopdb/store.go
@@ -100,18 +100,33 @@ func NewBoltSwapStore(dbPath string, chainParams *chaincfg.Params) (
 	// We'll create all the buckets we need if this is the first time we're
 	// starting up. If they already exist, then these calls will be noops.
 	err = bdb.Update(func(tx *bbolt.Tx) error {
-		_, err := tx.CreateBucketIfNotExists(loopOutBucketKey)
-		if err != nil {
-			return err
+		// Check if the meta bucket exists. If it exists, we consider
+		// the database as initialized and assume the meta bucket
+		// contains the db version.
+		metaBucket := tx.Bucket(metaBucketKey)
+		if metaBucket == nil {
+			log.Infof("Initializing new database with version %v",
+				latestDBVersion)
+
+			// Set db version to the current version.
+			err := setDBVersion(tx, latestDBVersion)
+			if err != nil {
+				return err
+			}
 		}
-		_, err = tx.CreateBucketIfNotExists(loopInBucketKey)
+
+		// Try creating these buckets, because loop in was added without
+		// bumping the db version number.
+		_, err = tx.CreateBucketIfNotExists(loopOutBucketKey)
 		if err != nil {
 			return err
 		}
-		_, err = tx.CreateBucketIfNotExists(metaBucket)
+
+		_, err = tx.CreateBucketIfNotExists(loopInBucketKey)
 		if err != nil {
 			return err
 		}
+
 		return nil
 	})
 	if err != nil {
diff --git a/loopdb/store_test.go b/loopdb/store_test.go
index 1e1a61f..64d9451 100644
--- a/loopdb/store_test.go
+++ b/loopdb/store_test.go
@@ -4,11 +4,13 @@ import (
 	"crypto/sha256"
 	"io/ioutil"
 	"os"
+	"path/filepath"
 	"reflect"
 	"testing"
 	"time"
 
 	"github.com/btcsuite/btcd/chaincfg"
+	"github.com/coreos/bbolt"
 	"github.com/lightninglabs/loop/test"
 	"github.com/lightningnetwork/lnd/lntypes"
 )
@@ -293,3 +295,72 @@ func TestLoopInStore(t *testing.T) {
 	}
 	checkSwap(StateFailInsufficientValue)
 }
+
+// TestVersionNew tests that a new database is initialized with the current
+// version.
+func TestVersionNew(t *testing.T) {
+	tempDirName, err := ioutil.TempDir("", "clientstore")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(tempDirName)
+
+	store, err := NewBoltSwapStore(tempDirName, &chaincfg.MainNetParams)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ver, err := getDBVersion(store.db)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if ver != latestDBVersion {
+		t.Fatal("db not at latest version")
+	}
+}
+
+// TestVersionNew tests that an existing version zero database is migrated to
+// the latest version.
+func TestVersionMigrated(t *testing.T) {
+	tempDirName, err := ioutil.TempDir("", "clientstore")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(tempDirName)
+
+	createVersionZeroDb(t, tempDirName)
+
+	store, err := NewBoltSwapStore(tempDirName, &chaincfg.MainNetParams)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ver, err := getDBVersion(store.db)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if ver != latestDBVersion {
+		t.Fatal("db not at latest version")
+	}
+}
+
+// createVersionZeroDb creates a database with an empty meta bucket. In version
+// zero, there was no version key specified yet.
+func createVersionZeroDb(t *testing.T, dbPath string) {
+	path := filepath.Join(dbPath, dbFileName)
+	bdb, err := bbolt.Open(path, 0600, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer bdb.Close()
+
+	err = bdb.Update(func(tx *bbolt.Tx) error {
+		_, err := tx.CreateBucket(metaBucketKey)
+		return err
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+}