feat: add ipmask

cmd/cmd.go

@@ -1,11 +1,9 @@
 package main
 
 import (
-	"encoding/binary"
 	"fmt"
 	"gotproxy/common"
 	"log"
-	"net"
 	"os"
 	"strconv"
 
@@ -39,11 +37,12 @@ var rootCmd = &cobra.Command{
 			return
 		}
 
-		ip, err := ipStrToUnit32()
+		ip, mask, err := common.ParseIPWithMask(ipStr)
 		if err != nil {
 			log.Fatal(err)
 		}
 		Options.Ip4 = ip
+		Options.Ip4Mask = mask
 
 		if proxyPid == 0 {
 			StartProxy()
@@ -60,22 +59,6 @@ var rootCmd = &cobra.Command{
 	},
 }
 
-func ipStrToUnit32() (uint32, error) {
-	if ipStr == "" {
-		return 0, nil
-	}
-	ip := net.ParseIP(ipStr)
-	if ip == nil {
-		return 0, fmt.Errorf("invalid ip: %s", ipStr)
-	}
-
-	ip = ip.To4()
-	if ip == nil {
-		return 0, fmt.Errorf("is not ipV4: %s", ipStr)
-	}
-	return binary.LittleEndian.Uint32(ip), nil
-}
-
 func Execute() {
 	if err := rootCmd.Execute(); err != nil {
 		fmt.Println(err)

cmd/loadBpf.go

@@ -30,6 +30,7 @@ type Options struct {
 	Command   string
 	Pids      []uint64
 	Ip4       uint32
+	Ip4Mask   uint8
 }
 
 func LoadBpf(options *Options) {
@@ -91,10 +92,11 @@ func LoadBpf(options *Options) {
 		pid = options.ProxyPid
 	}
 	config := proxyConfig{
-		ProxyPort:   options.ProxyPort,
-		ProxyPid:    pid,
-		FilterByPid: len(options.Pids) > 0,
-		FilterIp:    options.Ip4,
+		ProxyPort:    options.ProxyPort,
+		ProxyPid:     pid,
+		FilterByPid:  len(options.Pids) > 0,
+		FilterIp:     options.Ip4,
+		FilterIpMask: options.Ip4Mask,
 	}
 	stringToInt8Array(config.Command[:], options.Command)
 	err = objs.proxyMaps.MapConfig.Update(&key, &config, ebpf.UpdateAny)

cmd/proxy.c

@@ -20,6 +20,7 @@ struct Config {
   __u16 proxy_port;
   __u64 proxy_pid;
   __u32  filter_ip;
+  __u8 filter_ip_mask;
   bool filter_by_pid;
   char command[TASK_COMM_LEN];
 };
@@ -97,8 +98,11 @@ int cg_connect4(struct bpf_sock_addr *ctx) {
 
   if (!match_process(conf)) return 1;
 
-  if (conf->filter_ip){ 
-    if (ctx->user_ip4 != conf->filter_ip) {
+  if (conf->filter_ip)
+  {
+    __u32 mask = 0xFFFFFFFF >> (32 - conf->filter_ip_mask);
+    if ((ctx->user_ip4 & mask) != (conf->filter_ip & mask))
+    {
       BPF_LOG_DEBUG("not match ip\n");
       return 1;
     }

common/ip.go

@@ -0,0 +1,44 @@
+package common
+
+import (
+	"encoding/binary"
+	"fmt"
+	"net"
+	"strings"
+)
+
+func ParseIPWithMask(ipStr string) (uint32, uint8, error) {
+	if ipStr == "" {
+		return 0, 0, nil
+	}
+	if strings.Contains(ipStr, "/") {
+		_, ipNet, err := net.ParseCIDR(ipStr)
+		if err != nil {
+			return 0, 0, fmt.Errorf("invalid CIDR: %s", ipStr)
+		}
+
+		ip4 := ipNet.IP.To4()
+		if ip4 == nil {
+			return 0, 0, fmt.Errorf("not an IPv4 CIDR: %s", ipStr)
+		}
+
+		maskSize, _ := ipNet.Mask.Size()
+		ipVal := binary.LittleEndian.Uint32(ip4)
+		return ipVal, uint8(maskSize), nil
+	} else {
+		ip := net.ParseIP(ipStr)
+		if ip == nil {
+			return 0, 0, fmt.Errorf("invalid IP: %s", ipStr)
+		}
+
+		ip4 := ip.To4()
+		if ip4 == nil {
+			return 0, 0, fmt.Errorf("not an IPv4 address: %s", ipStr)
+		}
+
+		ipVal := binary.LittleEndian.Uint32(ip4)
+
+		fmt.Printf("Parsed IP: %s, Mask=32\n", ipStr)
+		return ipVal, 32, nil
+	}
+}