From 1bab313fae2562b25e8cc12d42fabd792c21f3d1 Mon Sep 17 00:00:00 2001
From: Vitaliy Vasylenko <me@vitalvas.com>
Date: Mon, 9 Sep 2024 22:52:15 -0400
Subject: [PATCH] feat: add xnet strip address

---
 xnet/strip.go      |  48 ++++++++++++++++++++
 xnet/strip_test.go | 109 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 157 insertions(+)
 create mode 100644 xnet/strip.go
 create mode 100644 xnet/strip_test.go

diff --git a/xnet/strip.go b/xnet/strip.go
new file mode 100644
index 0000000..66317e7
--- /dev/null
+++ b/xnet/strip.go
@@ -0,0 +1,48 @@
+package xnet
+
+import (
+	"errors"
+	"net"
+)
+
+var (
+	ErrInvalidIPAddress    = errors.New("invalid IP address")
+	ErrInvalidIPv4MaskSize = errors.New("invalid IPv4 mask size")
+	ErrInvalidIPv6MaskSize = errors.New("invalid IPv6 mask size")
+	ErrInvalidIPMaskSize   = errors.New("both IPv4 and IPv6 mask sizes cannot be zero")
+	ErrNonStandardIP       = errors.New("this is a non-standard IP or IPv6 address")
+)
+
+func GetStripedAddress(addr net.IP, ipv4, ipv6 int) (net.IP, error) {
+	if addr == nil {
+		return nil, ErrInvalidIPAddress
+	}
+
+	if ipv4 < 0 || ipv4 > 32 {
+		return nil, ErrInvalidIPv4MaskSize
+	}
+
+	if ipv6 < 0 || ipv6 > 128 {
+		return nil, ErrInvalidIPv6MaskSize
+	}
+
+	if ipv4 == 0 && ipv6 == 0 {
+		return nil, ErrInvalidIPMaskSize
+	}
+
+	if addr.To4() != nil && ipv4 > 0 {
+		mask := net.CIDRMask(ipv4, 32)
+		network := addr.Mask(mask)
+
+		return network, nil
+	}
+
+	if addr.To16() != nil && ipv6 > 0 {
+		mask := net.CIDRMask(ipv6, 128)
+		network := addr.Mask(mask)
+
+		return network, nil
+	}
+
+	return nil, ErrNonStandardIP
+}
diff --git a/xnet/strip_test.go b/xnet/strip_test.go
new file mode 100644
index 0000000..ec15b8e
--- /dev/null
+++ b/xnet/strip_test.go
@@ -0,0 +1,109 @@
+package xnet
+
+import (
+	"net"
+	"testing"
+)
+
+func TestGetStripedAddress(t *testing.T) {
+	tests := []struct {
+		name    string
+		addr    net.IP
+		ipv4    int
+		ipv6    int
+		want    net.IP
+		wantErr error
+	}{
+		{
+			name:    "Valid IPv4 Address with valid mask",
+			addr:    net.ParseIP("192.168.1.10"),
+			ipv4:    24,
+			want:    net.ParseIP("192.168.1.0"),
+			wantErr: nil,
+		},
+		{
+			name:    "Valid IPv6 Address with valid mask",
+			addr:    net.ParseIP("2001:db8::1"),
+			ipv6:    64,
+			want:    net.ParseIP("2001:db8::"),
+			wantErr: nil,
+		},
+		{
+			name:    "Valid IPv4 Address with /18 mask",
+			addr:    net.ParseIP("192.168.99.22"),
+			ipv4:    18,
+			want:    net.ParseIP("192.168.64.0"),
+			wantErr: nil,
+		},
+		{
+			name:    "Valid IPv6 Address with /56 mask",
+			addr:    net.ParseIP("2001:db8:1234:5678:90ab:cdef:0000:0001"),
+			ipv6:    56,
+			want:    net.ParseIP("2001:db8:1234:5600::"),
+			wantErr: nil,
+		},
+		{
+			name:    "Nil IP Address",
+			addr:    nil,
+			want:    nil,
+			wantErr: ErrInvalidIPAddress,
+		},
+		{
+			name:    "Invalid IPv4 Mask Size (negative)",
+			addr:    net.ParseIP("192.168.1.10"),
+			ipv4:    -1,
+			want:    nil,
+			wantErr: ErrInvalidIPv4MaskSize,
+		},
+		{
+			name:    "Invalid IPv4 Mask Size (too large)",
+			addr:    net.ParseIP("192.168.1.10"),
+			ipv4:    33,
+			want:    nil,
+			wantErr: ErrInvalidIPv4MaskSize,
+		},
+		{
+			name:    "Invalid IPv6 Mask Size (negative)",
+			addr:    net.ParseIP("2001:db8::1"),
+			ipv6:    -1,
+			want:    nil,
+			wantErr: ErrInvalidIPv6MaskSize,
+		},
+		{
+			name:    "Invalid IPv6 Mask Size (too large)",
+			addr:    net.ParseIP("2001:db8::1"),
+			ipv6:    129,
+			want:    nil,
+			wantErr: ErrInvalidIPv6MaskSize,
+		},
+		{
+			name:    "Both IPv4 and IPv6 Mask Sizes Are Zero",
+			addr:    net.ParseIP("192.168.1.10"),
+			ipv4:    0,
+			ipv6:    0,
+			want:    nil,
+			wantErr: ErrInvalidIPMaskSize,
+		},
+		{
+			name:    "Invalid IP Address",
+			addr:    net.IP{0xff, 0x00, 0x00, 0x00, 0xff},
+			ipv4:    24,
+			want:    nil,
+			wantErr: ErrNonStandardIP,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := GetStripedAddress(tt.addr, tt.ipv4, tt.ipv6)
+
+			if err != nil && tt.wantErr == nil || err == nil && tt.wantErr != nil || err != nil && err.Error() != tt.wantErr.Error() {
+				t.Errorf("GetStripedAddress() error = %v, wantErr %v", err, tt.wantErr)
+			}
+
+			if !got.Equal(tt.want) && !(got == nil && tt.want == nil) {
+				t.Errorf("GetStripedAddress() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}