DKG migration tool from v1.* to v2.0.0 (#1215)

* DKG migration tool from v1.* to v2.0.0
* implemented the migration
* added a CLI command for running it automagically for each beaconID
* reload beacon upon migration
* added integration test for restoring node state from migration

chain/beacon/node_test.go

@@ -58,6 +58,10 @@ func (t *testBeaconServer) SyncChain(req *drand.SyncRequest, p drand.Protocol_Sy
 	return SyncChain(t.h.l, t.h.chain, req, p)
 }
 
+func (t *testBeaconServer) Migrate(context.Context, *drand.Empty) (*drand.Empty, error) {
+	return &drand.Empty{}, nil
+}
+
 func dkgShares(_ *testing.T, n, t int, sch *crypto.Scheme) ([]*key.Share, []kyber.Point) {
 	var priPoly *share.PriPoly
 	var pubPoly *share.PubPoly

cmd/drand-cli/dkg_cli.go

@@ -116,9 +116,34 @@ var dkgCommand = &cli.Command{
 				return generateProposalCmd(c, l)
 			},
 		},
+		{
+			Name: "migrate",
+			Flags: toArray(
+				controlFlag,
+			),
+			Action: migrateDKG,
+		},
 	},
 }
 
+func migrateDKG(c *cli.Context) error {
+	if !c.IsSet(controlFlag.Name) {
+		return errors.New("you must set the control port")
+	}
+
+	port := c.String(controlFlag.Name)
+	ctrl, err := net.NewControlClientWithLogger(log.DefaultLogger(), port)
+	if err != nil {
+		return err
+	}
+
+	err = ctrl.Migrate()
+	if err == nil {
+		fmt.Println("Migration completed successfully!")
+	}
+	return err
+}
+
 var joinerFlag = &cli.StringSliceFlag{
 	Name: "joiner",
 	Usage: "the address of a joiner you wish to add to a DKG proposal. You can pass it multiple times. " +

core/drand_beacon_test.go

@@ -3,9 +3,13 @@ package core
 import (
 	"context"
 	"os"
+	"path"
 	"testing"
 	"time"
 
+	"github.com/drand/drand/dkg"
+	"github.com/drand/drand/protobuf/drand"
+
 	"github.com/jonboulle/clockwork"
 	"github.com/stretchr/testify/require"
 
@@ -200,3 +204,48 @@ func TestMemDBBeaconJoinsNetworkAfterDKG(t *testing.T) {
 	err = ts.WaitUntilRound(t, memDBNode, expectedRound-1)
 	require.NoError(t, err)
 }
+
+func TestMigrateMissingDKGDatabase(t *testing.T) {
+	const nodeCount = 3
+	const thr = 2
+	const period = 1 * time.Second
+	sleepDuration := 100 * time.Millisecond
+
+	// set up a few nodes and run a DKG
+	beaconName := t.Name()
+	ts := NewDrandTestScenario(t, nodeCount, thr, period, beaconName, clockwork.NewFakeClockAt(time.Now()))
+
+	group, err := ts.RunDKG()
+	require.NoError(t, err)
+
+	ts.SetMockClock(t, group.GenesisTime)
+	ts.AdvanceMockClock(t, period)
+	time.Sleep(sleepDuration)
+
+	err = ts.WaitUntilRound(t, ts.nodes[0], 2)
+	require.NoError(t, err)
+
+	// nuke the DKG state for a node and reload
+	// the DKG process to clear any open handles
+	node := ts.nodes[0]
+	err = os.Remove(path.Join(node.daemon.opts.configFolder, dkg.BoltFileName))
+	require.NoError(t, err)
+	dkgStore, err := dkg.NewDKGStore(node.daemon.opts.configFolder, node.daemon.opts.boltOpts)
+	require.NoError(t, err)
+	node.daemon.dkg = dkg.NewDKGProcess(
+		dkgStore, node.daemon, node.daemon.completedDKGs, node.daemon.privGateway, dkg.Config{}, node.daemon.log,
+	)
+	require.NoError(t, err)
+
+	// there should be no completed DKGs now for that node
+	status, err := node.daemon.DKGStatus(context.Background(), &drand.DKGStatusRequest{BeaconID: ts.beaconID})
+	require.NoError(t, err)
+	require.Nil(t, status.Complete)
+
+	// run the migration and check that there now is a completed DKG
+	_, err = node.daemon.Migrate(context.Background(), &drand.Empty{})
+	require.NoError(t, err)
+	status2, err := node.daemon.DKGStatus(context.Background(), &drand.DKGStatusRequest{BeaconID: ts.beaconID})
+	require.NoError(t, err)
+	require.NotNil(t, status2.Complete)
+}

core/drand_daemon_control.go

@@ -235,3 +235,35 @@ func (dd *DrandDaemon) Stop(ctx context.Context) {
 func (dd *DrandDaemon) WaitExit() chan bool {
 	return dd.exitCh
 }
+
+func (dd *DrandDaemon) Migrate(context.Context, *drand.Empty) (*drand.Empty, error) {
+	for beaconID, bp := range dd.beaconProcesses {
+		dd.log.Debugw("Migrating DKG from group file...", "beaconID", beaconID)
+
+		// first we fetch the old group and share from disk
+		group, err := bp.store.LoadGroup()
+		if err != nil {
+			return nil, err
+		}
+		share, err := bp.store.LoadShare(group.Scheme)
+		if err != nil {
+			return nil, err
+		}
+
+		// then we run the migration using them
+		if err := dd.dkg.Migrate(beaconID, group, share); err != nil {
+			return nil, err
+		}
+
+		// then stop and start the beacon process to load the new DKG
+		// state and start listening for messages correctly
+		bp.StopBeacon()
+		_, err = dd.LoadBeaconFromStore(beaconID, bp.store)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	dd.log.Debugw("Completed migration from group file")
+	return &drand.Empty{}, nil
+}

core/drand_test.go

@@ -305,7 +305,7 @@ func TestRunDKGReshareAbsentNodeForExecutionStart(t *testing.T) {
 
 	dt.SetMockClock(t, group1.GenesisTime)
 	// Note: Removing this sleep will cause the test to randomly break.
-	time.Sleep(1*time.Second)
+	time.Sleep(1 * time.Second)
 	err = dt.WaitUntilChainIsServing(t, dt.nodes[0])
 	require.NoError(t, err)
 
@@ -355,7 +355,7 @@ func TestRunDKGReshareTimeout(t *testing.T) {
 	beaconPeriod := 2 * time.Second
 	offline := 1
 	beaconID := test.GetBeaconIDFromEnv()
-	sleepDuration := 100*time.Millisecond
+	sleepDuration := 100 * time.Millisecond
 
 	dt := NewDrandTestScenario(t, oldNodes, oldThreshold, beaconPeriod, beaconID, clockwork.NewFakeClockAt(time.Now()))
 
@@ -538,6 +538,8 @@ func TestDrandPublicChainInfo(t *testing.T) {
 }
 
 // Test if we can correctly fetch the rounds after a DKG using the PublicRand RPC call
+//
+//nolint:funlen // this is a test
 func TestDrandPublicRand(t *testing.T) {
 	if os.Getenv("CI") == "true" {
 		t.Skip("test is flacky in CI")
@@ -636,7 +638,7 @@ func TestDrandPublicStream(t *testing.T) {
 	thr := key.DefaultThreshold(n)
 	p := 1 * time.Second
 	beaconID := test.GetBeaconIDFromEnv()
-	sleepDuration := 100*time.Millisecond
+	sleepDuration := 100 * time.Millisecond
 
 	dt := NewDrandTestScenario(t, n, thr, p, beaconID, clockwork.NewFakeClockAt(time.Now()))
 

dkg/actions_active_test.go

@@ -547,3 +547,8 @@ func (m *MockStore) Close() error {
 	args := m.Called()
 	return args.Error(0)
 }
+
+func (m *MockStore) MigrateFromGroupfile(beaconID string, group *key.Group, share *key.Share) error {
+	args := m.Called(beaconID, group, share)
+	return args.Error(0)
+}

dkg/dkg_migration_test.go

@@ -0,0 +1,183 @@
+package dkg
+
+import (
+	"testing"
+	"time"
+
+	"github.com/drand/drand/crypto"
+	"github.com/drand/drand/key"
+	"github.com/drand/kyber"
+	"github.com/drand/kyber/share"
+	"github.com/drand/kyber/share/dkg"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNilGroupFails(t *testing.T) {
+	store, err := NewDKGStore(t.TempDir(), nil)
+	require.NoError(t, err)
+
+	err = store.MigrateFromGroupfile("some-beacon", nil, fakeShare())
+	require.Error(t, err)
+}
+
+func TestNilKeyShareFails(t *testing.T) {
+	store, err := NewDKGStore(t.TempDir(), nil)
+	require.NoError(t, err)
+
+	group := fakeGroup()
+
+	err = store.MigrateFromGroupfile("some-beacon", group, nil)
+	require.Error(t, err)
+}
+
+func TestEmptyBeaconIDFails(t *testing.T) {
+	store, err := NewDKGStore(t.TempDir(), nil)
+	require.NoError(t, err)
+
+	keyShare := fakeShare()
+	group := fakeGroup()
+
+	err = store.MigrateFromGroupfile("", group, keyShare)
+	require.Error(t, err)
+}
+
+func TestStateAlreadyInDBForBeaconIDFails(t *testing.T) {
+	// create a new store
+	beaconID := "banana"
+	store, err := NewDKGStore(t.TempDir(), nil)
+	require.NoError(t, err)
+
+	// save an existing state in it
+	now := time.Now()
+	err = store.SaveFinished(beaconID, &DBState{
+		BeaconID:       beaconID,
+		Epoch:          1,
+		State:          Complete,
+		Threshold:      1,
+		Timeout:        now,
+		SchemeID:       crypto.DefaultSchemeID,
+		GenesisTime:    now,
+		GenesisSeed:    []byte("deadbeef"),
+		TransitionTime: now,
+		CatchupPeriod:  1,
+		BeaconPeriod:   3,
+		Leader:         nil,
+		Remaining:      nil,
+		Joining:        nil,
+		Leaving:        nil,
+		Acceptors:      nil,
+		Rejectors:      nil,
+		FinalGroup:     nil,
+		KeyShare:       nil,
+	})
+	require.NoError(t, err)
+
+	err = store.MigrateFromGroupfile(beaconID, fakeGroup(), fakeShare())
+	require.Error(t, err)
+}
+
+func TestStateInDBForDifferentBeaconIDDoesntFail(t *testing.T) {
+	// create a new store
+	beaconID := "banana"
+	aDifferentBeaconID := "different-beacon-id"
+	store, err := NewDKGStore(t.TempDir(), nil)
+	require.NoError(t, err)
+
+	// save an existing state but for a differen beacon ID
+	now := time.Now()
+	err = store.SaveFinished(aDifferentBeaconID, &DBState{
+		BeaconID:       aDifferentBeaconID,
+		Epoch:          1,
+		State:          Complete,
+		Threshold:      1,
+		Timeout:        now,
+		SchemeID:       crypto.DefaultSchemeID,
+		GenesisTime:    now,
+		GenesisSeed:    []byte("deadbeef"),
+		TransitionTime: now,
+		CatchupPeriod:  1,
+		BeaconPeriod:   3,
+		Leader:         nil,
+		Remaining:      nil,
+		Joining:        nil,
+		Leaving:        nil,
+		Acceptors:      nil,
+		Rejectors:      nil,
+		FinalGroup:     nil,
+		KeyShare:       nil,
+	})
+	require.NoError(t, err)
+
+	err = store.MigrateFromGroupfile(beaconID, fakeGroup(), fakeShare())
+	require.NoError(t, err)
+}
+
+func TestValidMigrationIsRetrievable(t *testing.T) {
+	// create a new store
+	beaconID := "banana"
+	store, err := NewDKGStore(t.TempDir(), nil)
+	require.NoError(t, err)
+
+	// perform the migration
+	err = store.MigrateFromGroupfile(beaconID, fakeGroup(), fakeShare())
+	require.NoError(t, err)
+
+	// get the finished migrated state and check some of its fields
+	state, err := store.GetFinished(beaconID)
+	require.NoError(t, err)
+	require.Equal(t, state.BeaconID, beaconID)
+	require.Equal(t, state.State, Complete)
+}
+
+func TestInvalidMigrationIsNotRetrievable(t *testing.T) {
+	// create a new store
+	beaconID := "banana"
+	store, err := NewDKGStore(t.TempDir(), nil)
+	require.NoError(t, err)
+
+	// perform an invalid migration
+	err = store.MigrateFromGroupfile(beaconID, nil, nil)
+	require.Error(t, err)
+
+	// get the finished migrated state and check some of its fields
+	state, err := store.GetFinished(beaconID)
+	require.NoError(t, err)
+	require.Nil(t, state)
+}
+
+func fakeShare() *key.Share {
+	sch := crypto.NewPedersenBLSChained()
+	scalarOne := sch.KeyGroup.Scalar().One()
+	s := &share.PriShare{I: 2, V: scalarOne}
+	return &key.Share{DistKeyShare: dkg.DistKeyShare{Share: s}, Scheme: sch}
+}
+
+func fakeGroup() *key.Group {
+	sch := crypto.NewPedersenBLSChained()
+	return &key.Group{
+		Threshold:     1,
+		Period:        3,
+		Scheme:        sch,
+		ID:            "default",
+		CatchupPeriod: 2,
+		Nodes: []*key.Node{{
+			Index: 0,
+			Identity: &key.Identity{
+				Key:       sch.KeyGroup.Point(),
+				Addr:      "localhost:1234",
+				TLS:       false,
+				Signature: []byte("abcd1234"),
+				Scheme:    sch,
+			},
+		}},
+		GenesisTime:    time.Now().Unix(),
+		GenesisSeed:    []byte("deadbeef"),
+		TransitionTime: time.Now().Unix(),
+		PublicKey:      fakePublic(),
+	}
+}
+
+func fakePublic() *key.DistPublic {
+	sch := crypto.NewPedersenBLSChained()
+	return &key.DistPublic{Coefficients: []kyber.Point{sch.KeyGroup.Point()}}
+}

dkg/dkg_process.go

@@ -67,6 +67,11 @@ type Store interface {
 
 	// Close closes and cleans up any database handles
 	Close() error
+
+	// MigrateFromGroupfile takes an existing groupfile and keyshare, and creates a first epoch DKG state for them.
+	// It will fail if DKG state already exists for the given beaconID
+	// Deprecated: will only exist in 2.0.0 for migration from v1.5.* to 2.0.0
+	MigrateFromGroupfile(beaconID string, groupFile *key.Group, share *key.Share) error
 }
 
 type Network interface {
@@ -122,3 +127,10 @@ func (d *DKGProcess) Close() {
 		d.log.Errorw("error closing the database", "err", err)
 	}
 }
+
+// Migrate takes an existing groupfile and keyshare, and creates a first epoch DKG state for them.
+// It will fail if DKG state already exists for the given beaconID
+// Deprecated: will only exist in 2.0.0 for migration from v1.5.* to 2.0.0
+func (d *DKGProcess) Migrate(beaconID string, groupfile *key.Group, share *key.Share) error {
+	return d.store.MigrateFromGroupfile(beaconID, groupfile, share)
+}