From d55ba2ca7d95731ce49166fc4115ce52229cbe64 Mon Sep 17 00:00:00 2001
From: Christopher Goes <cwgoes@protonmail.ch>
Date: Tue, 15 May 2018 16:00:17 +0200
Subject: [PATCH] Add p2p filter functions & tests

---
 baseapp/baseapp.go      | 45 ++++++++++++++++++++++++++++++++++++++---
 baseapp/baseapp_test.go | 33 ++++++++++++++++++++++++++++++
 types/abci.go           |  3 +++
 3 files changed, 78 insertions(+), 3 deletions(-)

diff --git a/baseapp/baseapp.go b/baseapp/baseapp.go
index cbe69d1b7c..8de8492005 100644
--- a/baseapp/baseapp.go
+++ b/baseapp/baseapp.go
@@ -40,9 +40,11 @@ type BaseApp struct {
 	txGasLimit  sdk.Gas         // per-transaction gas limit
 
 	// may be nil
-	initChainer  sdk.InitChainer  // initialize state with validators and state blob
-	beginBlocker sdk.BeginBlocker // logic to run before any txs
-	endBlocker   sdk.EndBlocker   // logic to run after all txs, and to determine valset changes
+	initChainer      sdk.InitChainer  // initialize state with validators and state blob
+	beginBlocker     sdk.BeginBlocker // logic to run before any txs
+	endBlocker       sdk.EndBlocker   // logic to run after all txs, and to determine valset changes
+	addrPeerFilter   sdk.PeerFilter   // filter peers by address and port
+	pubkeyPeerFilter sdk.PeerFilter   // filter peers by public key
 
 	//--------------------
 	// Volatile
@@ -142,6 +144,12 @@ func (app *BaseApp) SetEndBlocker(endBlocker sdk.EndBlocker) {
 func (app *BaseApp) SetAnteHandler(ah sdk.AnteHandler) {
 	app.anteHandler = ah
 }
+func (app *BaseApp) SetAddrPeerFilter(pf sdk.PeerFilter) {
+	app.addrPeerFilter = pf
+}
+func (app *BaseApp) SetPubKeyPeerFilter(pf sdk.PeerFilter) {
+	app.pubkeyPeerFilter = pf
+}
 func (app *BaseApp) Router() Router { return app.router }
 
 // load latest application version
@@ -282,6 +290,22 @@ func (app *BaseApp) InitChain(req abci.RequestInitChain) (res abci.ResponseInitC
 	return
 }
 
+// Filter peers by address / port
+func (app *BaseApp) FilterPeerByAddrPort(info string) abci.ResponseQuery {
+	if app.addrPeerFilter != nil {
+		return app.addrPeerFilter(info)
+	}
+	return abci.ResponseQuery{}
+}
+
+// Filter peers by public key
+func (app *BaseApp) FilterPeerByPubKey(info string) abci.ResponseQuery {
+	if app.pubkeyPeerFilter != nil {
+		return app.pubkeyPeerFilter(info)
+	}
+	return abci.ResponseQuery{}
+}
+
 // Implements ABCI.
 // Delegates to CommitMultiStore if it implements Queryable
 func (app *BaseApp) Query(req abci.RequestQuery) (res abci.ResponseQuery) {
@@ -318,6 +342,21 @@ func (app *BaseApp) Query(req abci.RequestQuery) (res abci.ResponseQuery) {
 		req.Path = req.Path[6:] // slice off "/store"
 		return queryable.Query(req)
 	}
+	// "/p2p" prefix for p2p queries
+	if strings.HasPrefix(path, "/p2p") {
+		path = path[4:]
+		if strings.HasPrefix(path, "/filter") {
+			path = path[7:]
+			if strings.HasPrefix(path, "/addr") {
+				path = path[6:]
+				return app.FilterPeerByAddrPort(path)
+			}
+			if strings.HasPrefix(path, "/pubkey") {
+				path = path[8:]
+				return app.FilterPeerByPubKey(path)
+			}
+		}
+	}
 	msg := "unknown query path"
 	return sdk.ErrUnknownRequest(msg).QueryResult()
 }
diff --git a/baseapp/baseapp_test.go b/baseapp/baseapp_test.go
index a605935c79..e54993648f 100644
--- a/baseapp/baseapp_test.go
+++ b/baseapp/baseapp_test.go
@@ -399,6 +399,39 @@ func TestQuery(t *testing.T) {
 	assert.Equal(t, value, res.Value)
 }
 
+// Test p2p filter queries
+func TestP2PQuery(t *testing.T) {
+	app := newBaseApp(t.Name())
+
+	// make a cap key and mount the store
+	capKey := sdk.NewKVStoreKey("main")
+	app.MountStoresIAVL(capKey)
+	err := app.LoadLatestVersion(capKey) // needed to make stores non-nil
+	assert.Nil(t, err)
+
+	app.SetAddrPeerFilter(func(addrport string) abci.ResponseQuery {
+		require.Equal(t, "1.1.1.1:8000", addrport)
+		return abci.ResponseQuery{Code: uint32(3)}
+	})
+
+	app.SetPubKeyPeerFilter(func(pubkey string) abci.ResponseQuery {
+		require.Equal(t, "testpubkey", pubkey)
+		return abci.ResponseQuery{Code: uint32(4)}
+	})
+
+	addrQuery := abci.RequestQuery{
+		Path: "/p2p/filter/addr/1.1.1.1:8000",
+	}
+	res := app.Query(addrQuery)
+	require.Equal(t, uint32(3), res.Code)
+
+	pubkeyQuery := abci.RequestQuery{
+		Path: "/p2p/filter/pubkey/testpubkey",
+	}
+	res = app.Query(pubkeyQuery)
+	require.Equal(t, uint32(4), res.Code)
+}
+
 //----------------------
 // TODO: clean this up
 
diff --git a/types/abci.go b/types/abci.go
index 40651163c4..a46e797ebe 100644
--- a/types/abci.go
+++ b/types/abci.go
@@ -10,3 +10,6 @@ type BeginBlocker func(ctx Context, req abci.RequestBeginBlock) abci.ResponseBeg
 
 // run code after the transactions in a block and return updates to the validator set
 type EndBlocker func(ctx Context, req abci.RequestEndBlock) abci.ResponseEndBlock
+
+// respond to p2p filtering queries from Tendermint
+type PeerFilter func(info string) abci.ResponseQuery