From 1f910c50dd44e53924ab4f94a1d681bd07a4e871 Mon Sep 17 00:00:00 2001
From: zcq98 <zhaocongqi_yewu@cmss.chinamobile.com>
Date: Fri, 31 May 2024 17:23:56 +0800
Subject: [PATCH] fix: IP Add/Sub overflow or underflow

Signed-off-by: zcq98 <zhaocongqi_yewu@cmss.chinamobile.com>
---
 pkg/ipam/ip.go      | 16 ++++++++++++++--
 pkg/ipam/ip_test.go | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/pkg/ipam/ip.go b/pkg/ipam/ip.go
index 7de50ff8464..e70011e75e4 100644
--- 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
index 81b886719b0..134fbd4d9ef 100644
--- 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 {