From a865e922d485b6dfc58955620fcbd565a4709c04 Mon Sep 17 00:00:00 2001
From: Qiu Jian <qiujian@yunionyun.com>
Date: Tue, 17 Aug 2021 22:43:50 +0800
Subject: [PATCH] fix: dhcp classless route encoding error

---
 pkg/util/dhcp/helpers.go      | 14 +++++---
 pkg/util/dhcp/helpers_test.go | 68 +++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 4 deletions(-)
 create mode 100644 pkg/util/dhcp/helpers_test.go

diff --git a/pkg/util/dhcp/helpers.go b/pkg/util/dhcp/helpers.go
index 5744298e5a..20e91c5256 100644
--- a/pkg/util/dhcp/helpers.go
+++ b/pkg/util/dhcp/helpers.go
@@ -70,10 +70,13 @@ func GetOptTime(d time.Duration) []byte {
 	return timeBytes
 }
 
-func GetClasslessRoutePack(route []string) []byte {
+func getClasslessRoutePack(route []string) []byte {
 	var snet, gw = route[0], route[1]
 	tmp := strings.Split(snet, "/")
 	netaddr := net.ParseIP(tmp[0])
+	if netaddr != nil {
+		netaddr = netaddr.To4()
+	}
 	masklen, _ := strconv.Atoi(tmp[1])
 	netlen := masklen / 8
 	if masklen%8 > 0 {
@@ -83,10 +86,13 @@ func GetClasslessRoutePack(route []string) []byte {
 		netaddr = netaddr[0:netlen]
 	}
 	gwaddr := net.ParseIP(gw)
+	if gwaddr != nil {
+		gwaddr = gwaddr.To4()
+	}
 
 	res := []byte{byte(masklen)}
-	res = append(res, []byte(netaddr.To4())...)
-	return append(res, []byte(gwaddr.To4())...)
+	res = append(res, []byte(netaddr)...)
+	return append(res, []byte(gwaddr)...)
 }
 
 func MakeReplyPacket(pkt Packet, conf *ResponseConfig) (Packet, error) {
@@ -159,7 +165,7 @@ func makeDHCPReplyPacket(req Packet, conf *ResponseConfig, msgType MessageType)
 			optCode = OptClasslessRouteWin
 		}
 		for _, route := range conf.Routes {
-			routeBytes := GetClasslessRoutePack(route)
+			routeBytes := getClasslessRoutePack(route)
 			resp.AddOption(optCode, routeBytes)
 		}
 	}
diff --git a/pkg/util/dhcp/helpers_test.go b/pkg/util/dhcp/helpers_test.go
new file mode 100644
index 0000000000..555486cb13
--- /dev/null
+++ b/pkg/util/dhcp/helpers_test.go
@@ -0,0 +1,68 @@
+// Copyright 2019 Yunion
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dhcp
+
+import (
+	"fmt"
+	"testing"
+)
+
+func compareBytes(b1, b2 []byte) error {
+	if len(b1) != len(b2) {
+		return fmt.Errorf("length not match")
+	}
+	for i := range b1 {
+		if b1[i] != b2[i] {
+			return fmt.Errorf("%dth not equal", i)
+		}
+	}
+	return nil
+}
+
+func TestGetClasslessRoutePack(t *testing.T) {
+	cases := []struct {
+		net  string
+		gw   string
+		want []byte
+	}{
+		{
+			net: "10.0.0.0/8",
+			gw:  "10.168.120.1",
+			want: []byte{
+				8, 10, 10, 168, 120, 1,
+			},
+		},
+		{
+			net: "192.168.0.0/16",
+			gw:  "10.168.120.1",
+			want: []byte{
+				16, 192, 168, 10, 168, 120, 1,
+			},
+		},
+		{
+			net: "172.16.0.0/12",
+			gw:  "10.168.222.1",
+			want: []byte{
+				12, 172, 16, 10, 168, 222, 1,
+			},
+		},
+	}
+	for _, c := range cases {
+		got := getClasslessRoutePack([]string{c.net, c.gw})
+		if err := compareBytes(c.want, got); err != nil {
+			t.Errorf("net: %s gw: %s want: %#v got: %#v", c.net, c.gw, c.want, got)
+		}
+	}
+}