From e5edd3b983189790391dca5b2ae4a0e460cb7f42 Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Tue, 22 Nov 2016 20:52:31 +0100
Subject: [PATCH] cmd/bootnode, cmd/geth, cmd/bzzd: add --netrestrict

---
 cmd/bootnode/main.go | 14 ++++++++++++--
 cmd/bzzd/main.go     |  1 +
 cmd/geth/main.go     |  1 +
 cmd/utils/flags.go   | 15 +++++++++++++++
 node/config.go       |  5 +++++
 node/node.go         |  1 +
 6 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/cmd/bootnode/main.go b/cmd/bootnode/main.go
index abecac3d8c..9b5ba19368 100644
--- a/cmd/bootnode/main.go
+++ b/cmd/bootnode/main.go
@@ -29,6 +29,7 @@ import (
 	"github.com/ethereum/go-ethereum/p2p/discover"
 	"github.com/ethereum/go-ethereum/p2p/discv5"
 	"github.com/ethereum/go-ethereum/p2p/nat"
+	"github.com/ethereum/go-ethereum/p2p/netutil"
 )
 
 func main() {
@@ -39,6 +40,7 @@ func main() {
 		nodeKeyFile = flag.String("nodekey", "", "private key filename")
 		nodeKeyHex  = flag.String("nodekeyhex", "", "private key as hex (for testing)")
 		natdesc     = flag.String("nat", "none", "port mapping mechanism (any|none|upnp|pmp|extip:<IP>)")
+		netrestrict = flag.String("netrestrict", "", "restrict network communication to the given IP networks (CIDR masks)")
 		runv5       = flag.Bool("v5", false, "run a v5 topic discovery bootnode")
 
 		nodeKey *ecdsa.PrivateKey
@@ -81,12 +83,20 @@ func main() {
 		os.Exit(0)
 	}
 
+	var restrictList *netutil.Netlist
+	if *netrestrict != "" {
+		restrictList, err = netutil.ParseNetlist(*netrestrict)
+		if err != nil {
+			utils.Fatalf("-netrestrict: %v", err)
+		}
+	}
+
 	if *runv5 {
-		if _, err := discv5.ListenUDP(nodeKey, *listenAddr, natm, ""); err != nil {
+		if _, err := discv5.ListenUDP(nodeKey, *listenAddr, natm, "", restrictList); err != nil {
 			utils.Fatalf("%v", err)
 		}
 	} else {
-		if _, err := discover.ListenUDP(nodeKey, *listenAddr, natm, ""); err != nil {
+		if _, err := discover.ListenUDP(nodeKey, *listenAddr, natm, "", restrictList); err != nil {
 			utils.Fatalf("%v", err)
 		}
 	}
diff --git a/cmd/bzzd/main.go b/cmd/bzzd/main.go
index b2f14a4a95..a3e87dc8a1 100644
--- a/cmd/bzzd/main.go
+++ b/cmd/bzzd/main.go
@@ -96,6 +96,7 @@ func init() {
 		utils.BootnodesFlag,
 		utils.KeyStoreDirFlag,
 		utils.ListenPortFlag,
+		utils.NetrestrictFlag,
 		utils.MaxPeersFlag,
 		utils.NATFlag,
 		utils.NodeKeyFileFlag,
diff --git a/cmd/geth/main.go b/cmd/geth/main.go
index 0eab77f7eb..e1d0b84c8c 100644
--- a/cmd/geth/main.go
+++ b/cmd/geth/main.go
@@ -149,6 +149,7 @@ participating.
 		utils.NatspecEnabledFlag,
 		utils.NoDiscoverFlag,
 		utils.DiscoveryV5Flag,
+		utils.NetrestrictFlag,
 		utils.NodeKeyFileFlag,
 		utils.NodeKeyHexFlag,
 		utils.RPCEnabledFlag,
diff --git a/cmd/utils/flags.go b/cmd/utils/flags.go
index dd18fd78ca..e97fcc0258 100644
--- a/cmd/utils/flags.go
+++ b/cmd/utils/flags.go
@@ -46,6 +46,7 @@ import (
 	"github.com/ethereum/go-ethereum/p2p/discover"
 	"github.com/ethereum/go-ethereum/p2p/discv5"
 	"github.com/ethereum/go-ethereum/p2p/nat"
+	"github.com/ethereum/go-ethereum/p2p/netutil"
 	"github.com/ethereum/go-ethereum/params"
 	"github.com/ethereum/go-ethereum/pow"
 	"github.com/ethereum/go-ethereum/rpc"
@@ -367,10 +368,16 @@ var (
 		Name:  "v5disc",
 		Usage: "Enables the experimental RLPx V5 (Topic Discovery) mechanism",
 	}
+	NetrestrictFlag = cli.StringFlag{
+		Name:  "netrestrict",
+		Usage: "Restricts network communication to the given IP networks (CIDR masks)",
+	}
+
 	WhisperEnabledFlag = cli.BoolFlag{
 		Name:  "shh",
 		Usage: "Enable Whisper",
 	}
+
 	// ATM the url is left to the user and deployment to
 	JSpathFlag = cli.StringFlag{
 		Name:  "jspath",
@@ -694,6 +701,14 @@ func MakeNode(ctx *cli.Context, name, gitCommit string) *node.Node {
 		config.MaxPeers = 0
 		config.ListenAddr = ":0"
 	}
+	if netrestrict := ctx.GlobalString(NetrestrictFlag.Name); netrestrict != "" {
+		list, err := netutil.ParseNetlist(netrestrict)
+		if err != nil {
+			Fatalf("Option %q: %v", NetrestrictFlag.Name, err)
+		}
+		config.NetRestrict = list
+	}
+
 	stack, err := node.New(config)
 	if err != nil {
 		Fatalf("Failed to create the protocol stack: %v", err)
diff --git a/node/config.go b/node/config.go
index 8d85b7ff82..8d75e441b7 100644
--- a/node/config.go
+++ b/node/config.go
@@ -34,6 +34,7 @@ import (
 	"github.com/ethereum/go-ethereum/p2p/discover"
 	"github.com/ethereum/go-ethereum/p2p/discv5"
 	"github.com/ethereum/go-ethereum/p2p/nat"
+	"github.com/ethereum/go-ethereum/p2p/netutil"
 )
 
 var (
@@ -103,6 +104,10 @@ type Config struct {
 	// Listener address for the V5 discovery protocol UDP traffic.
 	DiscoveryV5Addr string
 
+	// Restrict communication to white listed IP networks.
+	// The whitelist only applies when non-nil.
+	NetRestrict *netutil.Netlist
+
 	// BootstrapNodes used to establish connectivity with the rest of the network.
 	BootstrapNodes []*discover.Node
 
diff --git a/node/node.go b/node/node.go
index d49ae3a457..4b56fba4c5 100644
--- a/node/node.go
+++ b/node/node.go
@@ -165,6 +165,7 @@ func (n *Node) Start() error {
 		TrustedNodes:     n.config.TrusterNodes(),
 		NodeDatabase:     n.config.NodeDB(),
 		ListenAddr:       n.config.ListenAddr,
+		NetRestrict:      n.config.NetRestrict,
 		NAT:              n.config.NAT,
 		Dialer:           n.config.Dialer,
 		NoDial:           n.config.NoDial,