Fix interface srcIP

v-byte-cpu · Mar 10, 2021 · 92ff7f4 · 92ff7f4
1 parent a02b6ea
commit 92ff7f4
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 6 deletions.
diff --git a/command/arp.go b/command/arp.go
@@ -55,6 +55,9 @@ var arpCmd = &cobra.Command{
 			if iface, err = net.InterfaceByName(interfaceFlag); err != nil {
 				return err
 			}
+			if srcIP, err = ip.GetSubnetInterfaceIP(iface, dstSubnet); err != nil {
+				return err
+			}
 		} else {
 			if iface, srcIP, err = ip.GetSubnetInterface(dstSubnet); err != nil {
 				return err

diff --git a/pkg/ip/ip.go b/pkg/ip/ip.go
@@ -39,21 +39,33 @@ func ParseIPNet(subnet string) (*net.IPNet, error) {
 }
 
 func GetSubnetInterface(dstSubnet *net.IPNet) (*net.Interface, net.IP, error) {
-	dstSubnetIP := dstSubnet.IP.Mask(dstSubnet.Mask)
 	ifaces, err := net.Interfaces()
 	if err != nil {
 		return nil, nil, err
 	}
 	for _, iface := range ifaces {
-		addrs, err := iface.Addrs()
+		iface := iface
+		ifaceIP, err := GetSubnetInterfaceIP(&iface, dstSubnet)
 		if err != nil {
 			return nil, nil, err
 		}
-		for _, addr := range addrs {
-			if ipnet, ok := addr.(*net.IPNet); ok && ipnet.Contains(dstSubnetIP) {
-				return &iface, ipnet.IP.To4(), nil
-			}
+		if ifaceIP != nil {
+			return &iface, ifaceIP.To4(), nil
 		}
 	}
 	return nil, nil, ErrSubnetInterface
 }
+
+func GetSubnetInterfaceIP(iface *net.Interface, dstSubnet *net.IPNet) (net.IP, error) {
+	dstSubnetIP := dstSubnet.IP.Mask(dstSubnet.Mask)
+	addrs, err := iface.Addrs()
+	if err != nil {
+		return nil, err
+	}
+	for _, addr := range addrs {
+		if ipnet, ok := addr.(*net.IPNet); ok && ipnet.Contains(dstSubnetIP) {
+			return ipnet.IP.To4(), nil
+		}
+	}
+	return nil, nil
+}