fix: IP Add/Sub overflow or underflow

Signed-off-by: zcq98 <[email protected]>
kubeovn · May 31, 2024 · f940fd2 · f940fd2
1 parent 45c1e47
commit f940fd2
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 2 deletions.
diff --git a/pkg/ipam/ip.go b/pkg/ipam/ip.go
@@ -60,11 +60,23 @@ func bytes2IP(buff []byte, length int) IP {
 }
 
 func (a IP) Add(n int64) IP {
-	return bytes2IP(big.NewInt(0).Add(big.NewInt(0).SetBytes([]byte(a)), big.NewInt(n)).Bytes(), len(a))
+	if a.To4() != nil {
+		return bytes2IP(big.NewInt(0).Add(big.NewInt(0).SetBytes([]byte(a.To4())), big.NewInt(n)).Bytes(), len(a.To4()))
+	}
+	return bytes2IP(big.NewInt(0).Add(big.NewInt(0).SetBytes([]byte(a.To16())), big.NewInt(n)).Bytes(), len(a.To16()))
 }
 
 func (a IP) Sub(n int64) IP {
-	return bytes2IP(big.NewInt(0).Sub(big.NewInt(0).SetBytes([]byte(a)), big.NewInt(n)).Bytes(), len(a))
+	ipInt := big.NewInt(0).SetBytes(a.To16())
+	ipInt.Sub(ipInt, big.NewInt(n))
+	if ipInt.Sign() < 0 {
+		ipInt.Add(ipInt, big.NewInt(0).Exp(big.NewInt(2), big.NewInt(128), nil))
+	}
+
+	if a.To4() != nil {
+		return bytes2IP(ipInt.Bytes(), len(a.To4()))
+	}
+	return bytes2IP(ipInt.Bytes(), len(a.To16()))
 }
 
 func (a IP) String() string {

diff --git a/pkg/ipam/ip_test.go b/pkg/ipam/ip_test.go
@@ -145,12 +145,24 @@ func TestIPAdd(t *testing.T) {
 			n:    1,
 			want: IP(net.ParseIP("192.168.1.2")),
 		},
+		{
+			name: "IPv4 add 2",
+			a:    IP(net.ParseIP("192.168.1.1")),
+			n:    10,
+			want: IP(net.ParseIP("192.168.1.11")),
+		},
 		{
 			name: "IPv6 add",
 			a:    IP(net.ParseIP("2001:db8::1")),
 			n:    1,
 			want: IP(net.ParseIP("2001:db8::2")),
 		},
+		{
+			name: "IPv6 add 2",
+			a:    IP(net.ParseIP("1:db8::1")),
+			n:    1,
+			want: IP(net.ParseIP("1:db8::2")),
+		},
 		{
 			name: "IPv4 add overflow",
 			a:    IP(net.ParseIP("255.255.255.255")),
@@ -163,6 +175,12 @@ func TestIPAdd(t *testing.T) {
 			n:    2,
 			want: IP(net.ParseIP("0.0.0.0")),
 		},
+		{
+			name: "IPv4 add overflow 3",
+			a:    IP(net.ParseIP("255.255.255.254")),
+			n:    3,
+			want: IP(net.ParseIP("0.0.0.1")),
+		},
 		{
 			name: "IPv6 add overflow",
 			a:    IP(net.ParseIP("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff")),
@@ -175,6 +193,12 @@ func TestIPAdd(t *testing.T) {
 			n:    2,
 			want: IP(net.ParseIP("::")),
 		},
+		{
+			name: "IPv6 add overflow 3",
+			a:    IP(net.ParseIP("ffff:ffff:ffff:ffff:ffff:ffff:ffff:fffe")),
+			n:    3,
+			want: IP(net.ParseIP("::1")),
+		},
 	}
 
 	for _, tt := range tests {
@@ -217,6 +241,12 @@ func TestIPSub(t *testing.T) {
 			n:    2,
 			want: IP(net.ParseIP("255.255.255.255")),
 		},
+		{
+			name: "IPv4 sub underflow 3",
+			a:    IP(net.ParseIP("0.0.0.1")),
+			n:    3,
+			want: IP(net.ParseIP("255.255.255.254")),
+		},
 		{
 			name: "IPv6 sub underflow",
 			a:    IP(net.ParseIP("::")),
@@ -229,6 +259,12 @@ func TestIPSub(t *testing.T) {
 			n:    2,
 			want: IP(net.ParseIP("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff")),
 		},
+		{
+			name: "IPv6 sub underflow 3",
+			a:    IP(net.ParseIP("::1")),
+			n:    3,
+			want: IP(net.ParseIP("ffff:ffff:ffff:ffff:ffff:ffff:ffff:fffe")),
+		},
 	}
 
 	for _, tt := range tests {