From 6e4c285ed1d187e15c8b63d71c66feaff208c47e Mon Sep 17 00:00:00 2001
From: Austin Roberts <austin.roberts@openrelay.xyz>
Date: Mon, 12 Jul 2021 14:45:42 -0500
Subject: [PATCH] Add state update hooks

---
 cmd/geth/main.go           |  3 ++-
 cmd/geth/plugin_hooks.go   | 20 ++++++++++++++++++++
 core/state/plugin_hooks.go | 27 +++++++++++++++++++++++++++
 core/state/statedb.go      |  1 +
 plugins/plugin_loader.go   | 18 ++++++++++++++++--
 5 files changed, 66 insertions(+), 3 deletions(-)
 create mode 100644 core/state/plugin_hooks.go

diff --git a/cmd/geth/main.go b/cmd/geth/main.go
index e2ef7c7ed..bef31b517 100644
--- a/cmd/geth/main.go
+++ b/cmd/geth/main.go
@@ -313,7 +313,7 @@ func prepare(ctx *cli.Context) {
 // It creates a default node based on the command line arguments and runs it in
 // blocking mode, waiting for it to be shut down.
 func geth(ctx *cli.Context) error {
-	if err := plugins.Initialize(path.Join(ctx.GlobalString(utils.DataDirFlag.Name), "plugins")); err != nil { return err }
+	if err := plugins.Initialize(path.Join(ctx.GlobalString(utils.DataDirFlag.Name), "plugins"), ctx); err != nil { return err }
 	if ok, err := plugins.RunSubcommand(ctx); ok { return err }
 	if !plugins.ParseFlags(ctx.Args()) {
 		if args := ctx.Args(); len(args) > 0 {
@@ -323,6 +323,7 @@ func geth(ctx *cli.Context) error {
 
 	prepare(ctx)
 	stack, backend := makeFullNode(ctx)
+	pluginsInitializeNode(stack, backend)
 	defer stack.Close()
 	stack.RegisterAPIs(pluginGetAPIs(stack, backend))
 
diff --git a/cmd/geth/plugin_hooks.go b/cmd/geth/plugin_hooks.go
index 46bb3f902..a05ca94b7 100644
--- a/cmd/geth/plugin_hooks.go
+++ b/cmd/geth/plugin_hooks.go
@@ -31,3 +31,23 @@ func pluginGetAPIs(stack *node.Node, backend interfaces.Backend) []rpc.API {
 	 }
 	return GetAPIsFromLoader(plugins.DefaultPluginLoader, stack, backend)
 }
+
+func InitializeNode(pl *plugins.PluginLoader, stack *node.Node, backend interfaces.Backend) {
+  fnList := pl.Lookup("InitializeNode", func(item interface{}) bool {
+    _, ok := item.(func(*node.Node, interfaces.Backend))
+      return ok
+  })
+  for _, fni := range fnList {
+    if fn, ok := fni.(func(*node.Node, interfaces.Backend)); ok {
+      fn(stack, backend)
+    }
+  }
+}
+
+func pluginsInitializeNode(stack *node.Node, backend interfaces.Backend) {
+	if plugins.DefaultPluginLoader == nil {
+		log.Warn("Attempting InitializeNode, but default PluginLoader has not been initialized")
+		return
+	}
+	InitializeNode(plugins.DefaultPluginLoader, stack, backend)
+}
diff --git a/core/state/plugin_hooks.go b/core/state/plugin_hooks.go
new file mode 100644
index 000000000..5d11c4602
--- /dev/null
+++ b/core/state/plugin_hooks.go
@@ -0,0 +1,27 @@
+package state
+
+import (
+  "github.com/ethereum/go-ethereum/common"
+  "github.com/ethereum/go-ethereum/plugins"
+  "github.com/ethereum/go-ethereum/log"
+)
+
+func PluginStateUpdate(pl *plugins.PluginLoader, blockRoot, parentRoot common.Hash, destructs map[common.Hash]struct{}, accounts map[common.Hash][]byte, storage map[common.Hash]map[common.Hash][]byte) {
+  fnList := pl.Lookup("StateUpdate", func(item interface{}) bool  {
+    _, ok := item.(func(common.Hash, common.Hash, map[common.Hash]struct{}, map[common.Hash][]byte, map[common.Hash]map[common.Hash][]byte))
+      return ok
+  })
+  for _, fni := range fnList {
+    if fn, ok := fni.(func(common.Hash, common.Hash, map[common.Hash]struct{}, map[common.Hash][]byte, map[common.Hash]map[common.Hash][]byte)); ok {
+      fn(blockRoot, parentRoot, destructs, accounts, storage)
+    }
+  }
+}
+
+func pluginStateUpdate(blockRoot, parentRoot common.Hash, destructs map[common.Hash]struct{}, accounts map[common.Hash][]byte, storage map[common.Hash]map[common.Hash][]byte) {
+	if plugins.DefaultPluginLoader == nil {
+		log.Warn("Attempting StateUpdate, but default PluginLoader has not been initialized")
+    return
+	 }
+	PluginStateUpdate(plugins.DefaultPluginLoader, blockRoot, parentRoot, destructs, accounts, storage)
+}
diff --git a/core/state/statedb.go b/core/state/statedb.go
index 203556c6b..b6811b5a5 100644
--- a/core/state/statedb.go
+++ b/core/state/statedb.go
@@ -967,6 +967,7 @@ func (s *StateDB) Commit(deleteEmptyObjects bool) (common.Hash, error) {
 		}
 		// Only update if there's a state transition (skip empty Clique blocks)
 		if parent := s.snap.Root(); parent != root {
+			pluginStateUpdate(root, parent, s.snapDestructs, s.snapAccounts, s.snapStorage)
 			if err := s.snaps.Update(root, parent, s.snapDestructs, s.snapAccounts, s.snapStorage); err != nil {
 				log.Warn("Failed to update snapshot tree", "from", parent, "to", root, "err", err)
 			}
diff --git a/plugins/plugin_loader.go b/plugins/plugin_loader.go
index 96eec36ea..dae58955e 100644
--- a/plugins/plugin_loader.go
+++ b/plugins/plugin_loader.go
@@ -104,9 +104,23 @@ func NewPluginLoader(target string) (*PluginLoader, error) {
 	return pl, nil
 }
 
-func Initialize(target string) (err error) {
+func Initialize(target string, ctx *cli.Context) (err error) {
 	DefaultPluginLoader, err = NewPluginLoader(target)
-	return err
+	if err != nil { return err }
+	DefaultPluginLoader.Initialize(ctx)
+	return nil
+}
+
+func (pl *PluginLoader) Initialize(ctx *cli.Context) {
+	fns := pl.Lookup("Initialize", func(i interface{}) bool {
+		_, ok := i.(func(*cli.Context) error)
+		return ok
+	})
+	for _, fni := range fns {
+		if fn, ok := fni.(func(*cli.Context)); ok {
+			fn(ctx)
+		}
+	}
 }
 
 func (pl *PluginLoader) RunSubcommand(ctx *cli.Context) (bool, error) {