From 6add74b4f65888d38627ec1cf9fc90e7e624b4b6 Mon Sep 17 00:00:00 2001 From: Di Wu Date: Wed, 17 Jun 2020 00:05:42 +0800 Subject: [PATCH] Custom private blocks (#1705) Co-authored-by: Asim Aslam --- util/addr/addr.go | 9 +++++++++ util/addr/addr_test.go | 21 +++++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/util/addr/addr.go b/util/addr/addr.go index 469e3673..d814a95c 100644 --- a/util/addr/addr.go +++ b/util/addr/addr.go @@ -17,6 +17,15 @@ func init() { } } +// AppendPrivateBlocks append private network blocks +func AppendPrivateBlocks(bs ...string) { + for _, b := range bs { + if _, block, err := net.ParseCIDR(b); err == nil { + privateBlocks = append(privateBlocks, block) + } + } +} + func isPrivateIP(ipAddr string) bool { ip := net.ParseIP(ipAddr) for _, priv := range privateBlocks { diff --git a/util/addr/addr_test.go b/util/addr/addr_test.go index 06aae9a6..e728552b 100644 --- a/util/addr/addr_test.go +++ b/util/addr/addr_test.go @@ -56,3 +56,24 @@ func TestExtractor(t *testing.T) { } } + +func TestAppendPrivateBlocks(t *testing.T) { + tests := []struct { + addr string + expect bool + }{ + {addr: "9.134.71.34", expect: true}, + {addr: "8.10.110.34", expect: false}, // not in private blocks + } + + AppendPrivateBlocks("9.134.0.0/16") + + for _, test := range tests { + t.Run(test.addr, func(t *testing.T) { + res := isPrivateIP(test.addr) + if res != test.expect { + t.Fatalf("expected %t got %t", test.expect, res) + } + }) + } +}