Custom private blocks (#1705)
Co-authored-by: Asim Aslam <asim@aslam.me>
This commit is contained in:
		| @@ -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 { | func isPrivateIP(ipAddr string) bool { | ||||||
| 	ip := net.ParseIP(ipAddr) | 	ip := net.ParseIP(ipAddr) | ||||||
| 	for _, priv := range privateBlocks { | 	for _, priv := range privateBlocks { | ||||||
|   | |||||||
| @@ -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) | ||||||
|  | 			} | ||||||
|  | 		}) | ||||||
|  | 	} | ||||||
|  | } | ||||||
|   | |||||||
		Reference in New Issue
	
	Block a user