Merge pull request #273 from coinbase/patrick/protect-nil-handler

[storage] Protect against nil Helper and/or Handler in BalanceStorage

storage/errors/errors.go

@@ -334,6 +334,8 @@ var (
 	// to save cannot be parsed.
 	ErrInvalidValue = errors.New("invalid value")
 
+	ErrHelperHandlerMissing = errors.New("balance storage helper or handler is missing")
+
 	BalanceStorageErrs = []error{
 		ErrNegativeBalance,
 		ErrInvalidLiveBalance,
@@ -342,6 +344,7 @@ var (
 		ErrAccountMissing,
 		ErrInvalidChangeValue,
 		ErrInvalidValue,
+		ErrHelperHandlerMissing,
 	}
 )
 

storage/modules/balance_storage.go

@@ -188,6 +188,10 @@ func (b *BalanceStorage) AddingBlock(
 	block *types.Block,
 	transaction database.Transaction,
 ) (database.CommitWorker, error) {
+	if b.handler == nil {
+		return nil, storageErrs.ErrHelperHandlerMissing
+	}
+
 	changes, err := b.parser.BalanceChanges(ctx, block, false)
 	if err != nil {
 		return nil, fmt.Errorf("%w: unable to calculate balance changes", err)
@@ -244,6 +248,10 @@ func (b *BalanceStorage) RemovingBlock(
 	block *types.Block,
 	transaction database.Transaction,
 ) (database.CommitWorker, error) {
+	if b.handler == nil {
+		return nil, storageErrs.ErrHelperHandlerMissing
+	}
+
 	changes, err := b.parser.BalanceChanges(ctx, block, true)
 	if err != nil {
 		return nil, fmt.Errorf("%w: unable to calculate balance changes", err)
@@ -316,6 +324,10 @@ func (b *BalanceStorage) SetBalance(
 	amount *types.Amount,
 	block *types.BlockIdentifier,
 ) error {
+	if b.handler == nil {
+		return storageErrs.ErrHelperHandlerMissing
+	}
+
 	// Remove all account-related items
 	if err := b.deleteAccountRecords(
 		ctx,
@@ -424,6 +436,10 @@ func (b *BalanceStorage) Reconciled(
 // get an idea of the reconciliation coverage without
 // doing an expensive DB scan across all accounts.
 func (b *BalanceStorage) EstimatedReconciliationCoverage(ctx context.Context) (float64, error) {
+	if b.helper == nil {
+		return -1, storageErrs.ErrHelperHandlerMissing
+	}
+
 	dbTx := b.db.ReadTransaction(ctx)
 	defer dbTx.Discard(ctx)
 
@@ -498,6 +514,10 @@ func (b *BalanceStorage) existingValue(
 	parentBlock *types.BlockIdentifier,
 	existingValue string,
 ) (string, error) {
+	if b.helper == nil {
+		return "", storageErrs.ErrHelperHandlerMissing
+	}
+
 	if exists {
 		return existingValue, nil
 	}
@@ -547,6 +567,10 @@ func (b *BalanceStorage) applyExemptions(
 	change *parser.BalanceChange,
 	newVal string,
 ) (string, error) {
+	if b.helper == nil {
+		return "", storageErrs.ErrHelperHandlerMissing
+	}
+
 	// Find exemptions that are applicable to the *parser.BalanceChange
 	exemptions := b.parser.FindExemptions(change.Account, change.Currency)
 	if len(exemptions) == 0 {
@@ -610,6 +634,10 @@ func (b *BalanceStorage) deleteAccountRecords(
 	account *types.AccountIdentifier,
 	currency *types.Currency,
 ) error {
+	if b.handler == nil {
+		return storageErrs.ErrHelperHandlerMissing
+	}
+
 	// Remove historical balance records
 	if err := b.removeHistoricalBalances(
 		ctx,

storage/modules/balance_storage_test.go

@@ -1035,7 +1035,6 @@ func TestBootstrapBalances(t *testing.T) {
 	mockHelper.On("Asserter").Return(baseAsserter())
 	mockHelper.On("ExemptFunc").Return(exemptFunc())
 	mockHelper.On("BalanceExemptions").Return([]*types.BalanceExemption{})
-	storage.Initialize(mockHelper, mockHandler)
 	bootstrapBalancesFile := path.Join(newDir, "balances.csv")
 
 	t.Run("File doesn't exist", func(t *testing.T) {
@@ -1047,29 +1046,40 @@ func TestBootstrapBalances(t *testing.T) {
 		assert.Contains(t, err.Error(), "no such file or directory")
 	})
 
-	t.Run("Set balance successfully", func(t *testing.T) {
-		amount := &types.Amount{
-			Value: "10",
-			Currency: &types.Currency{
-				Symbol:   "BTC",
-				Decimals: 8,
-			},
-		}
+	// Initialize file
+	amount := &types.Amount{
+		Value: "10",
+		Currency: &types.Currency{
+			Symbol:   "BTC",
+			Decimals: 8,
+		},
+	}
 
-		file, err := json.MarshalIndent([]*BootstrapBalance{
-			{
-				Account:  account,
-				Value:    amount.Value,
-				Currency: amount.Currency,
-			},
-		}, "", " ")
-		assert.NoError(t, err)
+	file, err := json.MarshalIndent([]*BootstrapBalance{
+		{
+			Account:  account,
+			Value:    amount.Value,
+			Currency: amount.Currency,
+		},
+	}, "", " ")
+	assert.NoError(t, err)
 
-		assert.NoError(
-			t,
-			ioutil.WriteFile(bootstrapBalancesFile, file, utils.DefaultFilePermissions),
+	assert.NoError(
+		t,
+		ioutil.WriteFile(bootstrapBalancesFile, file, utils.DefaultFilePermissions),
+	)
+
+	t.Run("run before initializing helper/handler", func(t *testing.T) {
+		err = storage.BootstrapBalances(
+			ctx,
+			bootstrapBalancesFile,
+			genesisBlockIdentifier,
 		)
+		assert.True(t, errors.Is(err, storageErrs.ErrHelperHandlerMissing))
+	})
 
+	storage.Initialize(mockHelper, mockHandler)
+	t.Run("Set balance successfully", func(t *testing.T) {
 		mockHandler.On("AccountsSeen", ctx, mock.Anything, 1).Return(nil).Once()
 		err = storage.BootstrapBalances(
 			ctx,
@@ -1248,8 +1258,14 @@ func TestBalanceReconciliation(t *testing.T) {
 	mockHelper.On("Asserter").Return(baseAsserter())
 	mockHelper.On("ExemptFunc").Return(exemptFunc())
 	mockHelper.On("BalanceExemptions").Return([]*types.BalanceExemption{})
-	storage.Initialize(mockHelper, mockHandler)
 
+	t.Run("test estimated before helper/handler", func(t *testing.T) {
+		coverage, err := storage.EstimatedReconciliationCoverage(ctx)
+		assert.Equal(t, float64(-1), coverage)
+		assert.True(t, errors.Is(err, storageErrs.ErrHelperHandlerMissing))
+	})
+
+	storage.Initialize(mockHelper, mockHandler)
 	t.Run("attempt to store reconciliation for non-existent account", func(t *testing.T) {
 		err := storage.Reconciled(ctx, account, currency, genesisBlock)
 		assert.NoError(t, err)