github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/storage/badger/dkg_state_test.go (about)

     1  package badger_test
     2  
     3  import (
     4  	"errors"
     5  	"math/rand"
     6  	"testing"
     7  
     8  	"github.com/dgraph-io/badger/v2"
     9  	"github.com/stretchr/testify/assert"
    10  	"github.com/stretchr/testify/require"
    11  
    12  	"github.com/onflow/flow-go/model/flow"
    13  	"github.com/onflow/flow-go/module/metrics"
    14  	"github.com/onflow/flow-go/storage"
    15  	bstorage "github.com/onflow/flow-go/storage/badger"
    16  	"github.com/onflow/flow-go/utils/unittest"
    17  )
    18  
    19  func TestDKGState_DKGStarted(t *testing.T) {
    20  	unittest.RunWithTypedBadgerDB(t, bstorage.InitSecret, func(db *badger.DB) {
    21  		metrics := metrics.NewNoopCollector()
    22  		store, err := bstorage.NewDKGState(metrics, db)
    23  		require.NoError(t, err)
    24  
    25  		epochCounter := rand.Uint64()
    26  
    27  		// check dkg-started flag for non-existent epoch
    28  		t.Run("DKGStarted should default to false", func(t *testing.T) {
    29  			started, err := store.GetDKGStarted(rand.Uint64())
    30  			assert.NoError(t, err)
    31  			assert.False(t, started)
    32  		})
    33  
    34  		// store dkg-started flag for epoch
    35  		t.Run("should be able to set DKGStarted", func(t *testing.T) {
    36  			err = store.SetDKGStarted(epochCounter)
    37  			assert.NoError(t, err)
    38  		})
    39  
    40  		// retrieve flag for epoch
    41  		t.Run("should be able to read DKGStarted", func(t *testing.T) {
    42  			started, err := store.GetDKGStarted(epochCounter)
    43  			assert.NoError(t, err)
    44  			assert.True(t, started)
    45  		})
    46  	})
    47  }
    48  
    49  func TestDKGState_BeaconKeys(t *testing.T) {
    50  	unittest.RunWithTypedBadgerDB(t, bstorage.InitSecret, func(db *badger.DB) {
    51  		metrics := metrics.NewNoopCollector()
    52  		store, err := bstorage.NewDKGState(metrics, db)
    53  		require.NoError(t, err)
    54  
    55  		epochCounter := rand.Uint64()
    56  
    57  		// attempt to get a non-existent key
    58  		t.Run("should error if retrieving non-existent key", func(t *testing.T) {
    59  			_, err = store.RetrieveMyBeaconPrivateKey(epochCounter)
    60  			assert.True(t, errors.Is(err, storage.ErrNotFound))
    61  		})
    62  
    63  		// attempt to store a nil key should fail  - use DKGState.SetEndState(flow.DKGEndStateNoKey)
    64  		t.Run("should fail to store a nil key instead)", func(t *testing.T) {
    65  			err = store.InsertMyBeaconPrivateKey(epochCounter, nil)
    66  			assert.Error(t, err)
    67  		})
    68  
    69  		// store a key in db
    70  		expected := unittest.RandomBeaconPriv()
    71  		t.Run("should be able to store and read a key", func(t *testing.T) {
    72  			err = store.InsertMyBeaconPrivateKey(epochCounter, expected)
    73  			require.NoError(t, err)
    74  		})
    75  
    76  		// retrieve the key by epoch counter
    77  		t.Run("should be able to retrieve stored key", func(t *testing.T) {
    78  			actual, err := store.RetrieveMyBeaconPrivateKey(epochCounter)
    79  			require.NoError(t, err)
    80  			assert.Equal(t, expected, actual)
    81  		})
    82  
    83  		// test storing same key
    84  		t.Run("should fail to store a key twice", func(t *testing.T) {
    85  			err = store.InsertMyBeaconPrivateKey(epochCounter, expected)
    86  			require.True(t, errors.Is(err, storage.ErrAlreadyExists))
    87  		})
    88  	})
    89  }
    90  
    91  func TestDKGState_EndState(t *testing.T) {
    92  	unittest.RunWithTypedBadgerDB(t, bstorage.InitSecret, func(db *badger.DB) {
    93  		metrics := metrics.NewNoopCollector()
    94  		store, err := bstorage.NewDKGState(metrics, db)
    95  		require.NoError(t, err)
    96  
    97  		epochCounter := rand.Uint64()
    98  		endState := flow.DKGEndStateNoKey
    99  
   100  		t.Run("should be able to store an end state", func(t *testing.T) {
   101  			err = store.SetDKGEndState(epochCounter, endState)
   102  			require.NoError(t, err)
   103  		})
   104  
   105  		t.Run("should be able to read an end state", func(t *testing.T) {
   106  			readEndState, err := store.GetDKGEndState(epochCounter)
   107  			require.NoError(t, err)
   108  			assert.Equal(t, endState, readEndState)
   109  		})
   110  	})
   111  }
   112  
   113  func TestSafeBeaconPrivateKeys(t *testing.T) {
   114  	unittest.RunWithTypedBadgerDB(t, bstorage.InitSecret, func(db *badger.DB) {
   115  		metrics := metrics.NewNoopCollector()
   116  		dkgState, err := bstorage.NewDKGState(metrics, db)
   117  		require.NoError(t, err)
   118  		safeKeys := bstorage.NewSafeBeaconPrivateKeys(dkgState)
   119  
   120  		t.Run("non-existent key -> should return ErrNotFound", func(t *testing.T) {
   121  			epochCounter := rand.Uint64()
   122  			key, safe, err := safeKeys.RetrieveMyBeaconPrivateKey(epochCounter)
   123  			assert.Nil(t, key)
   124  			assert.False(t, safe)
   125  			assert.ErrorIs(t, err, storage.ErrNotFound)
   126  		})
   127  
   128  		t.Run("existent key, non-existent end state -> should return ErrNotFound", func(t *testing.T) {
   129  			epochCounter := rand.Uint64()
   130  
   131  			// store a key
   132  			expected := unittest.RandomBeaconPriv().PrivateKey
   133  			err := dkgState.InsertMyBeaconPrivateKey(epochCounter, expected)
   134  			assert.NoError(t, err)
   135  
   136  			key, safe, err := safeKeys.RetrieveMyBeaconPrivateKey(epochCounter)
   137  			assert.Nil(t, key)
   138  			assert.False(t, safe)
   139  			assert.ErrorIs(t, err, storage.ErrNotFound)
   140  		})
   141  
   142  		t.Run("existent key, unsuccessful end state -> not safe", func(t *testing.T) {
   143  			epochCounter := rand.Uint64()
   144  
   145  			// store a key
   146  			expected := unittest.RandomBeaconPriv().PrivateKey
   147  			err := dkgState.InsertMyBeaconPrivateKey(epochCounter, expected)
   148  			assert.NoError(t, err)
   149  			// mark dkg unsuccessful
   150  			err = dkgState.SetDKGEndState(epochCounter, flow.DKGEndStateInconsistentKey)
   151  			assert.NoError(t, err)
   152  
   153  			key, safe, err := safeKeys.RetrieveMyBeaconPrivateKey(epochCounter)
   154  			assert.Nil(t, key)
   155  			assert.False(t, safe)
   156  			assert.NoError(t, err)
   157  		})
   158  
   159  		t.Run("existent key, inconsistent key end state -> not safe", func(t *testing.T) {
   160  			epochCounter := rand.Uint64()
   161  
   162  			// store a key
   163  			expected := unittest.RandomBeaconPriv().PrivateKey
   164  			err := dkgState.InsertMyBeaconPrivateKey(epochCounter, expected)
   165  			assert.NoError(t, err)
   166  			// mark dkg result as inconsistent
   167  			err = dkgState.SetDKGEndState(epochCounter, flow.DKGEndStateInconsistentKey)
   168  			assert.NoError(t, err)
   169  
   170  			key, safe, err := safeKeys.RetrieveMyBeaconPrivateKey(epochCounter)
   171  			assert.Nil(t, key)
   172  			assert.False(t, safe)
   173  			assert.NoError(t, err)
   174  		})
   175  
   176  		t.Run("non-existent key, no key end state -> not safe", func(t *testing.T) {
   177  			epochCounter := rand.Uint64()
   178  
   179  			// mark dkg result as no key
   180  			err = dkgState.SetDKGEndState(epochCounter, flow.DKGEndStateNoKey)
   181  			assert.NoError(t, err)
   182  
   183  			key, safe, err := safeKeys.RetrieveMyBeaconPrivateKey(epochCounter)
   184  			assert.Nil(t, key)
   185  			assert.False(t, safe)
   186  			assert.NoError(t, err)
   187  		})
   188  
   189  		t.Run("existent key, successful end state -> safe", func(t *testing.T) {
   190  			epochCounter := rand.Uint64()
   191  
   192  			// store a key
   193  			expected := unittest.RandomBeaconPriv().PrivateKey
   194  			err := dkgState.InsertMyBeaconPrivateKey(epochCounter, expected)
   195  			assert.NoError(t, err)
   196  			// mark dkg successful
   197  			err = dkgState.SetDKGEndState(epochCounter, flow.DKGEndStateSuccess)
   198  			assert.NoError(t, err)
   199  
   200  			key, safe, err := safeKeys.RetrieveMyBeaconPrivateKey(epochCounter)
   201  			assert.NotNil(t, key)
   202  			assert.True(t, expected.Equals(key))
   203  			assert.True(t, safe)
   204  			assert.NoError(t, err)
   205  		})
   206  
   207  		t.Run("non-existent key, successful end state -> exception!", func(t *testing.T) {
   208  			epochCounter := rand.Uint64()
   209  
   210  			// mark dkg successful
   211  			err = dkgState.SetDKGEndState(epochCounter, flow.DKGEndStateSuccess)
   212  			assert.NoError(t, err)
   213  
   214  			key, safe, err := safeKeys.RetrieveMyBeaconPrivateKey(epochCounter)
   215  			assert.Nil(t, key)
   216  			assert.False(t, safe)
   217  			assert.Error(t, err)
   218  			assert.NotErrorIs(t, err, storage.ErrNotFound)
   219  		})
   220  
   221  	})
   222  }
   223  
   224  // TestSecretDBRequirement tests that the DKGState constructor will return an
   225  // error if instantiated using a database not marked with the correct type.
   226  func TestSecretDBRequirement(t *testing.T) {
   227  	unittest.RunWithBadgerDB(t, func(db *badger.DB) {
   228  		metrics := metrics.NewNoopCollector()
   229  		_, err := bstorage.NewDKGState(metrics, db)
   230  		require.Error(t, err)
   231  	})
   232  }