From 0041f5823b8434eb31807892b3ee0b72f3852d01 Mon Sep 17 00:00:00 2001
From: i-norden <iansnorden@gmail.com>
Date: Tue, 14 Jun 2022 23:10:43 -0500
Subject: [PATCH] config and env updates

---
 pkg/snapshot/config.go | 32 ++++++++++++++++++++++++++++----
 pkg/snapshot/env.go    |  3 +++
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/pkg/snapshot/config.go b/pkg/snapshot/config.go
index 4d2b76b..26ecf56 100644
--- a/pkg/snapshot/config.go
+++ b/pkg/snapshot/config.go
@@ -19,6 +19,8 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/ethereum/go-ethereum/common"
+
 	"github.com/sirupsen/logrus"
 
 	"github.com/ethereum/go-ethereum/statediff/indexer/database/sql/postgres"
@@ -38,9 +40,10 @@ const (
 
 // Config contains params for both databases the service uses
 type Config struct {
-	Eth  *EthConfig
-	DB   *DBConfig
-	File *FileConfig
+	Eth     *EthConfig
+	DB      *DBConfig
+	File    *FileConfig
+	Service *ServiceConfig
 }
 
 // EthConfig is config parameters for the chain.
@@ -60,11 +63,17 @@ type FileConfig struct {
 	OutputDir string
 }
 
+type ServiceConfig struct {
+	AllowedPaths    [][]byte
+	AllowedAccounts []common.Address
+}
+
 func NewConfig(mode SnapshotMode) (*Config, error) {
 	ret := &Config{
 		&EthConfig{},
 		&DBConfig{},
 		&FileConfig{},
+		&ServiceConfig{},
 	}
 	return ret, ret.Init(mode)
 }
@@ -110,7 +119,7 @@ func (c *Config) Init(mode SnapshotMode) error {
 	default:
 		return fmt.Errorf("no output mode specified")
 	}
-	return nil
+	return c.Service.Init()
 }
 
 func (c *DBConfig) Init() {
@@ -148,3 +157,18 @@ func (c *FileConfig) Init() error {
 	}
 	return nil
 }
+
+func (c *ServiceConfig) Init() error {
+	viper.BindEnv(SNAPSHOT_ACCOUNTS_TOML, SNAPSHOT_ACCOUNTS)
+	allowedAccounts := viper.GetStringSlice(SNAPSHOT_ACCOUNTS)
+	accountsLen := len(allowedAccounts)
+	if accountsLen != 0 {
+		c.AllowedAccounts = make([]common.Address, accountsLen)
+		for i, allowedAccount := range allowedAccounts {
+			c.AllowedAccounts[i] = common.HexToAddress(allowedAccount)
+		}
+	} else {
+		logrus.Infof("no snapshot addresses specified, will perform snapshot of entire trie(s)")
+	}
+	return nil
+}
diff --git a/pkg/snapshot/env.go b/pkg/snapshot/env.go
index 5c9927a..a592726 100644
--- a/pkg/snapshot/env.go
+++ b/pkg/snapshot/env.go
@@ -23,6 +23,7 @@ const (
 	SNAPSHOT_MODE          = "SNAPSHOT_MODE"
 	SNAPSHOT_START_HEIGHT  = "SNAPSHOT_START_HEIGHT"
 	SNAPSHOT_END_HEIGHT    = "SNAPSHOT_END_HEIGHT"
+	SNAPSHOT_ACCOUNTS      = "SNAPSHOT_ACCOUNTS"
 
 	LOGRUS_LEVEL = "LOGRUS_LEVEL"
 	LOGRUS_FILE  = "LOGRUS_FILE"
@@ -62,6 +63,7 @@ const (
 	SNAPSHOT_MODE_TOML          = "snapshot.mode"
 	SNAPSHOT_START_HEIGHT_TOML  = "snapshot.startHeight"
 	SNAPSHOT_END_HEIGHT_TOML    = "snapshot.endHeight"
+	SNAPSHOT_ACCOUNTS_TOML      = "snapshot.accounts"
 
 	LOGRUS_LEVEL_TOML = "log.level"
 	LOGRUS_FILE_TOML  = "log.file"
@@ -101,6 +103,7 @@ const (
 	SNAPSHOT_MODE_CLI          = "snapshot-mode"
 	SNAPSHOT_START_HEIGHT_CLI  = "start-height"
 	SNAPSHOT_END_HEIGHT_CLI    = "end-height"
+	SNAPSHOT_ACCOUNTS_CLI      = "snapshot-accounts"
 
 	LOGRUS_LEVEL_CLI = "log-level"
 	LOGRUS_FILE_CLI  = "log-file"