Statediff API to add an address to be watched

statediff/api.go

@@ -149,3 +149,8 @@ func (api *PublicStateDiffAPI) WriteStateDiffAt(ctx context.Context, blockNumber
 func (api *PublicStateDiffAPI) WriteStateDiffFor(ctx context.Context, blockHash common.Hash, params Params) error {
 	return api.sds.WriteStateDiffFor(blockHash, params)
 }
+
+// WatchAddress adds the given address to a list of watched addresses to which the direct statediff process is restricted
+func (api *PublicStateDiffAPI) WatchAddress(address common.Address) error {
+	return api.sds.WatchAddress(address)
+}

statediff/api_test.go

@@ -0,0 +1,101 @@
+// Copyright 2019 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package statediff_test
+
+import (
+	"fmt"
+	"os"
+	"reflect"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/statediff"
+)
+
+func init() {
+	if os.Getenv("MODE") != "statediff" {
+		fmt.Println("Skipping statediff test")
+		os.Exit(0)
+	}
+}
+
+var (
+	address1Hex = "0x1ca7c995f8eF0A2989BbcE08D5B7Efe50A584aa1"
+	address2Hex = "0xe799eE0191652c864E49F3A3344CE62535B15afe"
+	address1    = common.HexToAddress(address1Hex)
+	address2    = common.HexToAddress(address2Hex)
+
+	watchedAddresses0 []common.Address
+	watchedAddresses1 = []common.Address{address1}
+	watchedAddresses2 = []common.Address{address1, address2}
+
+	expectedError = fmt.Errorf("Address %s already watched", address1Hex)
+
+	service = statediff.Service{}
+)
+
+func TestWatchAddress(t *testing.T) {
+	watchedAddresses := service.GetWathchedAddresses()
+	if !reflect.DeepEqual(watchedAddresses, watchedAddresses0) {
+		t.Error("Test failure:", t.Name())
+		t.Logf("Actual watched addresses not equal expected watched addresses.\nactual: %+v\nexpected: %+v", watchedAddresses, watchedAddresses0)
+	}
+
+	testWatchUnwatchedAddress(t)
+	testWatchWatchedAddress(t)
+}
+
+func testWatchUnwatchedAddress(t *testing.T) {
+	err := service.WatchAddress(address1)
+	if err != nil {
+		t.Error("Test failure:", t.Name())
+		t.Logf("Unexpected error %s thrown on an attempt to watch an unwatched address.", err.Error())
+	}
+	watchedAddresses := service.GetWathchedAddresses()
+	if !reflect.DeepEqual(watchedAddresses, watchedAddresses1) {
+		t.Error("Test failure:", t.Name())
+		t.Logf("Actual watched addresses not equal expected watched addresses.\nactual: %+v\nexpected: %+v", watchedAddresses, watchedAddresses1)
+	}
+
+	err = service.WatchAddress(address2)
+	if err != nil {
+		t.Error("Test failure:", t.Name())
+		t.Logf("Unexpected error %s thrown on an attempt to watch an unwatched address.", err.Error())
+	}
+	watchedAddresses = service.GetWathchedAddresses()
+	if !reflect.DeepEqual(watchedAddresses, watchedAddresses2) {
+		t.Error("Test failure:", t.Name())
+		t.Logf("Actual watched addresses not equal expected watched addresses.\nactual: %+v\nexpected: %+v", watchedAddresses, watchedAddresses2)
+	}
+}
+
+func testWatchWatchedAddress(t *testing.T) {
+	err := service.WatchAddress(address1)
+	if err == nil {
+		t.Error("Test failure:", t.Name())
+		t.Logf("Expected error %s not thrown on an attempt to watch an already watched address.", expectedError.Error())
+	}
+	if err.Error() != expectedError.Error() {
+		t.Error("Test failure:", t.Name())
+		t.Logf("Actual thrown error not equal expected error.\nactual: %+v\nexpected: %+v", err.Error(), expectedError.Error())
+	}
+	watchedAddresses := service.GetWathchedAddresses()
+	if !reflect.DeepEqual(watchedAddresses, watchedAddresses2) {
+		t.Error("Test failure:", t.Name())
+		t.Logf("Actual watched addresses not equal expected watched addresses.\nactual: %+v\nexpected: %+v", watchedAddresses, watchedAddresses2)
+	}
+}

statediff/service.go

@@ -18,7 +18,11 @@ package statediff
 
 import (
 	"bytes"
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
 	"math/big"
+	"os"
 	"strconv"
 	"strings"
 	"sync"
@@ -54,6 +58,9 @@ const (
 	deadlockDetected   = "deadlock detected" // 40P01 https://www.postgresql.org/docs/current/errcodes-appendix.html
 )
 
+// TODO: Take the watched addresses file path as a CLI arg.
+const watchedAddressesFile = "./watched-addresses.json"
+
 var writeLoopParams = Params{
 	IntermediateStateNodes:   true,
 	IntermediateStorageNodes: true,
@@ -101,6 +108,10 @@ type IService interface {
 	WriteStateDiffFor(blockHash common.Hash, params Params) error
 	// Event loop for progressively processing and writing diffs directly to DB
 	WriteLoop(chainEventCh chan core.ChainEvent)
+	// Method to add an address to be watched to write loop params
+	WatchAddress(address common.Address) error
+	// Method to get currently watched addresses from write loop params
+	GetWathchedAddresses() []common.Address
 }
 
 // Wraps consructor parameters
@@ -200,6 +211,12 @@ func New(stack *node.Node, ethServ *eth.Ethereum, cfg *ethconfig.Config, params
 	}
 	stack.RegisterLifecycle(sds)
 	stack.RegisterAPIs(sds.APIs())
+
+	err := loadWatchedAddresses()
+	if err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -713,3 +730,91 @@ func (sds *Service) writeStateDiffWithRetry(block *types.Block, parentRoot commo
 	}
 	return err
 }
+
+// Adds the provided address to the list of watched addresses in write loop params and to the watched addresses file
+func (sds *Service) WatchAddress(address common.Address) error {
+	// Check if address is already being watched
+	if containsAddress(writeLoopParams.WatchedAddresses, address) {
+		return fmt.Errorf("Address %s already watched", address)
+	}
+
+	// Check if the watched addresses file exists
+	fileExists, err := doesFileExist(watchedAddressesFile)
+	if err != nil {
+		return err
+	}
+
+	// Create the watched addresses file if doesn't exist
+	if !fileExists {
+		_, err := os.Create(watchedAddressesFile)
+		if err != nil {
+			return err
+		}
+	}
+
+	watchedAddresses := append(writeLoopParams.WatchedAddresses, address)
+
+	// Write the updated list of watched address to a json file
+	content, err := json.Marshal(watchedAddresses)
+	err = ioutil.WriteFile(watchedAddressesFile, content, 0644)
+	if err != nil {
+		return err
+	}
+
+	// Update the in-memory params as well
+	writeLoopParams.WatchedAddresses = watchedAddresses
+
+	return nil
+}
+
+// Gets currently watched addresses from the in-memory write loop params
+func (sds *Service) GetWathchedAddresses() []common.Address {
+	return writeLoopParams.WatchedAddresses
+}
+
+// loadWatchedAddresses is used to load watched addresses to the in-memory write loop params from a json file if it exists
+func loadWatchedAddresses() error {
+	// Check if the watched addresses file exists
+	fileExists, err := doesFileExist(watchedAddressesFile)
+	if err != nil {
+		return err
+	}
+
+	if fileExists {
+		content, err := ioutil.ReadFile(watchedAddressesFile)
+		if err != nil {
+			return err
+		}
+
+		var watchedAddresses []common.Address
+		err = json.Unmarshal(content, &watchedAddresses)
+		if err != nil {
+			return err
+		}
+
+		writeLoopParams.WatchedAddresses = watchedAddresses
+	}
+
+	return nil
+}
+
+// containsAddress is used to check if an address is present in the provided list of watched addresses
+func containsAddress(watchedAddresses []common.Address, address common.Address) bool {
+	for _, addr := range watchedAddresses {
+		if addr == address {
+			return true
+		}
+	}
+	return false
+}
+
+// doesFileExist is used to check if file at a given path exists
+func doesFileExist(path string) (bool, error) {
+	_, err := os.Stat(path)
+	if err == nil {
+		return true, nil
+	} else if os.IsNotExist(err) {
+		return false, nil
+	}
+	return false, err
+}

statediff/testhelpers/mocks/service.go

@@ -332,3 +332,11 @@ func sendNonBlockingQuit(id rpc.ID, sub statediff.Subscription) {
 		log.Info("unable to close subscription %s; channel has no receiver", id)
 	}
 }
+
+func (sds *MockStateDiffService) WatchAddress(address common.Address) error {
+	return nil
+}
+
+func (sds *MockStateDiffService) GetWathchedAddresses() []common.Address {
+	return []common.Address{}
+}