feat: add filter by ip

cmd/cmd.go

@@ -1,7 +1,10 @@
 package main
 
 import (
+	"encoding/binary"
 	"fmt"
+	"log"
+	"net"
 	"os"
 	"strconv"
 
@@ -13,20 +16,28 @@ var (
 	proxyPort uint16
 	proxyPid  uint64
 	pids      []string
+	ipStr     string
 )
 
 var rootCmd = &cobra.Command{
 	Use:   "gotproxy",
 	Short: "A simple tcp transparent proxy tool for Linux",
 	Run: func(cmd *cobra.Command, args []string) {
-		if proxyPid == 0 {
-			StartProxy(proxyPort)
-		}
 		Options := &Options{
 			Command:   command,
 			ProxyPid:  proxyPid,
 			ProxyPort: proxyPort,
 		}
+
+		ip, err := ipStrToUnit32()
+		if err != nil {
+			log.Fatal(err)
+		}
+		Options.Ip4 = ip
+
+		if proxyPid == 0 {
+			StartProxy(proxyPort)
+		}
 		for _, pid := range pids {
 			pidInt, err := strconv.ParseUint(pid, 10, 64)
 			if err != nil {
@@ -39,6 +50,22 @@ 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)
@@ -55,4 +82,5 @@ func init() {
 	rootCmd.PersistentFlags().Uint16Var(&proxyPort, "p-port", 18000, "The proxy port")
 	rootCmd.PersistentFlags().Uint64Var(&proxyPid, "p-pid", 0, "The process ID of the proxy. If not provided, the program will automatically start a forwarding proxy.")
 	rootCmd.PersistentFlags().StringSliceVar(&pids, "pids", []string{}, "The pid to be proxied, seperate by ','")
+	rootCmd.PersistentFlags().StringVar(&ipStr, "ip", "", "The ip to be proxied,only support ipv4")
 }

cmd/loadBpf.go

@@ -29,6 +29,7 @@ type Options struct {
 	ProxyPid  uint64 // PID of the proxy server
 	Command   string
 	Pids      []uint64
+	Ip4       uint32
 }
 
 func LoadBpf(options *Options) {
@@ -81,8 +82,7 @@ func LoadBpf(options *Options) {
 		log.Print("Attaching TcpSetState program as kprobe:", err)
 	}
 	defer kprobeLink.Close()
-	// Update the proxyMaps map with the proxy server configuration, because we need to know the proxy server PID in order
-	// to filter out eBPF events generated by the proxy server itself so it would not proxy its own packets in a loop.
+
 	var key uint32 = 0
 	var pid uint64
 	if options.ProxyPid == 0 {
@@ -94,6 +94,7 @@ func LoadBpf(options *Options) {
 		ProxyPort:   options.ProxyPort,
 		ProxyPid:    pid,
 		FilterByPid: len(options.Pids) > 0,
+		FilterIp:    options.Ip4,
 	}
 	stringToInt8Array(config.Command[:], options.Command)
 	err = objs.proxyMaps.MapConfig.Update(&key, &config, ebpf.UpdateAny)

cmd/proxy.c

@@ -9,9 +9,17 @@
 #define MAX_CONNECTIONS 20000
 #define MAX_PIDS 64
 
+#define DEBUG
+
+#ifdef DEBUG
+  #define BPF_LOG_DEBUG(fmt, ...) bpf_printk(fmt, ##__VA_ARGS__)
+#endif
+
+
 struct Config {
   __u16 proxy_port;
   __u64 proxy_pid;
+  __u32  filter_ip;
   bool filter_by_pid;
   char command[TASK_COMM_LEN];
 };
@@ -54,12 +62,10 @@ struct {
 #define SO_ORIGINAL_DST 80
 
 #define AF_INET 2
-
-
 static __always_inline bool
-match(struct Config *conf)
+match_process(struct Config *conf)
 {
-  if (conf->command[0] == '\0' || !conf->filter_by_pid){
+  if (conf->command[0] == '\0' && !conf->filter_by_pid){
     return true;
   }
 
@@ -89,14 +95,17 @@ int cg_connect4(struct bpf_sock_addr *ctx) {
   if (!conf) return 1;
   if ((bpf_get_current_pid_tgid() >> 32) == conf->proxy_pid) return 1;
 
-  if (!match(conf)) return 1;
+  if (!match_process(conf)) return 1;
+
+  if (conf->filter_ip){ 
+    if (ctx->user_ip4 != conf->filter_ip) {
+      BPF_LOG_DEBUG("not match ip\n");
+      return 1;
+    }
+  }
 
-  // This field contains the IPv4 address passed to the connect() syscall
-  // a.k.a. connect to this socket destination address and port
   __u32 dst_addr = bpf_ntohl(ctx->user_ip4);
-  // This field contains the port number passed to the connect() syscall
   __u16 dst_port = bpf_ntohl(ctx->user_port) >> 16;
-  // Unique identifier for the destination socket
   __u64 cookie = bpf_get_socket_cookie(ctx);
 
   // Store destination socket under cookie key
@@ -110,7 +119,7 @@ int cg_connect4(struct bpf_sock_addr *ctx) {
   ctx->user_ip4 = bpf_htonl(0x7f000001); // 127.0.0.1 == proxy IP
   ctx->user_port = bpf_htonl(conf->proxy_port << 16); // Proxy port
 
-  bpf_printk("Redirecting client connection to proxy\n");
+  BPF_LOG_DEBUG("Redirecting client connection to proxy\n");
 
   return 1;
 }
@@ -119,22 +128,20 @@ int cg_connect4(struct bpf_sock_addr *ctx) {
 // This is just to record client source address and port after succesful connection establishment to the proxy
 SEC("sockops")
 int cg_sock_ops(struct bpf_sock_ops *ctx) {
-  //  bpf_printk("sockops");
   // Only forward on IPv4 connections
   if (ctx->family != AF_INET) return 0;
 
   // Active socket with an established connection
   if (ctx->op == BPF_SOCK_OPS_ACTIVE_ESTABLISHED_CB) {
     __u64 cookie = bpf_get_socket_cookie(ctx);
-    // Lookup the socket in the map for the corresponding cookie
-    // In case the socket is present, store the source port and socket mapping
+
     struct Socket *sock = bpf_map_lookup_elem(&map_socks, &cookie);
     if (sock) {
       __u16 src_port = ctx->local_port;
       bpf_map_update_elem(&map_ports, &src_port, &cookie, 0);
     }
   }
-  // bpf_printk("sockops hook successful\n");
+  BPF_LOG_DEBUG("sockops hook successful\n");
   return 0;
 }
 
@@ -144,19 +151,11 @@ int cg_sock_ops(struct bpf_sock_ops *ctx) {
 // then establishes a connection with the original target and forwards the client's request.
 SEC("cgroup/getsockopt")
 int cg_sock_opt(struct bpf_sockopt *ctx) {
-  // The SO_ORIGINAL_DST socket option is a specialized option used primarily in the context of network address translation (NAT) and transparent proxying.
-  // In a typical NAT or transparent proxy setup, incoming packets are redirected from their original destination to a proxy server. 
-  // The proxy server, upon receiving the packets, often needs to know the original destination address in order to handle the traffic appropriately. 
-  // This is where SO_ORIGINAL_DST comes into play.
-  // bpf_printk("cg_sock_opt");
   if (ctx->optname != SO_ORIGINAL_DST) return 1;
   // Only forward IPv4 TCP connections
   if (ctx->sk->family != AF_INET) return 1;
   if (ctx->sk->protocol != IPPROTO_TCP) return 1;
 
-  // Get the clients source port
-  // It's actually sk->dst_port because getsockopt() syscall with SO_ORIGINAL_DST socket option
-  // is retrieving the original dst port of the client so it's "querying" the destination port of the client
   __u16 src_port = bpf_ntohs(ctx->sk->dst_port);
 
   // Retrieve the socket cookie using the clients' src_port 
@@ -191,7 +190,7 @@ int tcp_set_state(struct pt_regs *ctx)
     __u64 *cookie = bpf_map_lookup_elem(&map_ports, &src_port);
     if (cookie)
     {
-      // bpf_printk("tcp close\n");
+      BPF_LOG_DEBUG("tcp close\n");
       bpf_map_delete_elem(&map_ports, &src_port);
       bpf_map_delete_elem(&map_socks, &cookie);
     }