From 3b77f5a35ad1b332b675bc44818f815cd0d08ace Mon Sep 17 00:00:00 2001
From: Tobias Klauser <tobias@cilium.io>
Date: Thu, 24 Aug 2023 15:22:24 +0200
Subject: [PATCH] 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 <tobias@cilium.io>
---
 netipx.go      | 15 ++++++++++++++
 netipx_test.go | 54 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+)

diff --git a/netipx.go b/netipx.go
index 57e8146..d241fb3 100644
--- 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
index f33805c..b3e5333 100644
--- 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