netipx: add ParsePrefixOrAddr

ParsePrefixOrAddr parses a string as an IP address prefix or IP address. If the string parses as a prefix, its address part (netip.Prefix.Addr) is returned. Signed-off-by: Tobias Klauser <[email protected]>
go4org · Aug 24, 2023 · 73b0654 · 73b0654
1 parent ec4c8b8
commit 73b0654
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 0 deletions.
diff --git a/netipx.go b/netipx.go
@@ -51,6 +51,21 @@ func FromStdIPRaw(std net.IP) (ip netip.Addr, ok bool) {
 	return netip.AddrFromSlice(std)
 }
 
+// ParsePrefixOrAddr parses s as an IP address prefix or IP address. If s parses
+// as an IP address prefix, its [net/netip.Prefix.Addr] is returned. The string
+// s can be an IPv4 address ("192.0.2.1"), IPv6 address ("2001:db8::68"), IPv4
+// prefix ("192.0.2.1/32"), or IPv6 prefix ("2001:db:68/96").
+func ParsePrefixOrAddr(s string) (netip.Addr, error) {
+	// factored out of netip.ParsePrefix to avoid allocating an empty netip.Prefix in case it's
+	// an address and not a prefix.
+	i := strings.LastIndexByte(s, '/')
+	if i < 0 {
+		return netip.ParseAddr(s)
+	}
+	prefix, err := netip.ParsePrefix(s)
+	return prefix.Addr(), err
+}
+
 // AddrNext returns the IP following ip.
 // If there is none, it returns the IP zero value.
 //

diff --git a/netipx_test.go b/netipx_test.go
@@ -210,6 +210,60 @@ func TestFromStdIPNet(t *testing.T) {
 	}
 }
 
+func TestParsePrefixOrAddr(t *testing.T) {
+	tests := []struct {
+		name    string
+		s       string
+		want    netip.Addr
+		wantErr bool
+	}{
+		{
+			name:    "empty IP",
+			s:       "",
+			wantErr: true,
+		},
+		{
+			name:    "invalid IP",
+			s:       "192.168.",
+			wantErr: true,
+		},
+		{
+			name: "IPv4 address",
+			s:    "127.0.0.1",
+			want: netip.MustParseAddr("127.0.0.1"),
+		},
+		{
+			name: "IPv6 address",
+			s:    "2001:db8::",
+			want: netip.MustParseAddr("2001:db8::"),
+		},
+		{
+			name: "IPv4 prefix",
+			s:    "192.0.2.1/32",
+			want: netip.MustParseAddr("192.0.2.1"),
+		},
+		{
+			name: "IPv6 prefix",
+			s:    "2001:db8::68/96",
+			want: netip.MustParseAddr("2001:db8::68"),
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := ParsePrefixOrAddr(tt.s)
+			if err != nil {
+				if tt.wantErr {
+					return
+				}
+				t.Fatal(err)
+			}
+			if got != tt.want {
+				t.Errorf("ParsePrefixOrAddr(%q) = %v; want %v", tt.s, got, tt.want)
+			}
+		})
+	}
+}
+
 func TestAddrIPNet(t *testing.T) {
 	tests := []struct {
 		name string