feat: add --proto tcp|udp flag

README.md

@@ -39,10 +39,10 @@ sudo ./gotproxy [flags]
 | **--p-pid** | The process ID of the proxy. If not provided, the program will automatically start a forwarding proxy. |
 | **--p-port** | The proxy port. |
 | **--socks5**	| The SOCKS5 proxy Server network address. If configured, SOCKS5 proxying will be used. |
+| **--proto** | Proxy protocol selection: `both` (default) / `tcp` / `udp`. When set to `tcp`, only TCP traffic will be redirected; when set to `udp`, only UDP traffic will be redirected. |
 
 Features Under Development：
 IPv6 support
-UDP support
 
 ***Examples***
 1. Proxy a specific command:
@@ -57,6 +57,16 @@ sudo ./gotproxy --socks5 192.168.1.2:1080
  ```
 Where '192.168.1.2:1080' is the IP and port of the SOCKS5 proxy server.
 
+3. TCP-only proxy:
+```bash
+sudo ./gotproxy --proto tcp
+```
+
+4. UDP-only proxy:
+```bash
+sudo ./gotproxy --proto udp
+```
+
 
 ## Known Limitations ##
 * Theoretically, a connection should be determined by a 5-tuple, but for most cases, connection mapping is currently based only on protocol type and source port.

README_CN.md

@@ -40,10 +40,11 @@ sudo ./gotproxy [flags]
 | **--p-pid** | 代理程序的进程id. 会自动过滤不代理该进程的网络通信，以免网络循环。如果没有配置, 本程序会自动启动一个转发代理服务. |
 | **--p-port** | 代理服务监听的端口。 |
 | **--socks5** | socks5代理的服务端地址，如果配置，会进行socks5代理. |
+| **--proto** | 代理协议选择：`both`（默认）/ `tcp` / `udp`。当设置为 `tcp` 时只重定向 TCP 流量；设置为 `udp` 时只重定向 UDP 流量。 |
 
 
 正在开发中的功能：
-支持ipv6,支持udp
+支持ipv6
 
 
 
@@ -61,6 +62,16 @@ sudo ./gotproxy --socks5 192.168.1.2:1080
  ```
 其中‘192.168.1.2:1080’是socks5代理服务器的ip和端口
 
+3. 仅代理 TCP:
+```bash
+sudo ./gotproxy --proto tcp
+```
+
+4. 仅代理 UDP:
+```bash
+sudo ./gotproxy --proto udp
+```
+
 ## 已知限制：
 * 理论上应该根据5元组确定一个连接，但是考虑大多数情况目前只根据协议类型和源端口进行连接映射。
 * 在根据进程名称进行代理的场景中，如果进程启动了子进程并使用了execve执行一个新命令，会无法进行代理。

cmd/cmd.go

@@ -17,16 +17,31 @@ var (
 	pids            []string
 	ipStr           string
 	socks5ProxyAddr string
+	proto           string
 )
 
 var rootCmd = &cobra.Command{
 	Use:   "gotproxy",
 	Short: "A simple tcp transparent proxy tool for Linux",
 	Run: func(cmd *cobra.Command, args []string) {
+		var enableTCP, enableUDP bool
+		switch proto {
+		case "both":
+			enableTCP, enableUDP = true, true
+		case "tcp":
+			enableTCP, enableUDP = true, false
+		case "udp":
+			enableTCP, enableUDP = false, true
+		default:
+			log.Fatalf("Invalid --proto value %q, expected one of: both|tcp|udp", proto)
+		}
+
 		Options := &Options{
 			Command:   command,
 			ProxyPid:  proxyPid,
 			ProxyPort: proxyPort,
+			EnableTCP: enableTCP,
+			EnableUDP: enableUDP,
 		}
 
 		if ok, err := common.HasPermission(); err != nil {
@@ -74,4 +89,5 @@ func init() {
 	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")
 	rootCmd.PersistentFlags().StringVar(&socks5ProxyAddr, "socks5", "", "The socks5 proxyAddr.")
+	rootCmd.PersistentFlags().StringVar(&proto, "proto", "both", "Proxy protocol: both|tcp|udp")
 }

cmd/loadBpf.go

@@ -31,6 +31,8 @@ type Options struct {
 	Pids      []uint64
 	Ip4       uint32
 	Ip4Mask   uint8
+	EnableTCP bool
+	EnableUDP bool
 }
 
 func LoadBpf(options *Options) {
@@ -46,9 +48,20 @@ func LoadBpf(options *Options) {
 	}
 	defer objs.Close()
 
+	// Log the effective proxy protocol before starting user-space forwarding.
+	mode := "both"
+	if options.EnableTCP && !options.EnableUDP {
+		mode = "tcp"
+	} else if !options.EnableTCP && options.EnableUDP {
+		mode = "udp"
+	} else if !options.EnableTCP && !options.EnableUDP {
+		mode = "none"
+	}
+	log.Printf("Proxy protocol enabled: %s (tcp=%v udp=%v)", mode, options.EnableTCP, options.EnableUDP)
+
 	// Start TCP (and UDP) proxy so it can use objs.MapUdpDest for UDP original-dest lookup
 	if options.ProxyPid == 0 {
-		StartProxy(objs.MapUdpDest)
+		StartProxy(objs.MapUdpDest, options.EnableTCP, options.EnableUDP)
 	}
 
 	// Attach eBPF programs to the root cgroup
@@ -101,6 +114,8 @@ func LoadBpf(options *Options) {
 		FilterByPid:  len(options.Pids) > 0,
 		FilterIp:     options.Ip4,
 		FilterIpMask: options.Ip4Mask,
+		EnableTcp:    options.EnableTCP,
+		EnableUdp:    options.EnableUDP,
 	}
 	stringToInt8Array(config.Command[:], options.Command)
 	err = objs.proxyMaps.MapConfig.Update(&key, &config, ebpf.UpdateAny)

cmd/proxy.c

@@ -24,6 +24,8 @@ struct Config {
   __u32  filter_ip;
   __u8 filter_ip_mask;
   bool filter_by_pid;
+  bool enable_tcp;
+  bool enable_udp;
   char command[TASK_COMM_LEN];
 };
 
@@ -142,6 +144,7 @@ int cg_connect4(struct bpf_sock_addr *ctx) {
   __u16 dst_port = bpf_ntohl(ctx->user_port) >> 16;
 
   if (ctx->protocol == IPPROTO_TCP) {
+    if (!conf->enable_tcp) return 1;
     __u64 cookie = bpf_get_socket_cookie(ctx);
     struct Socket sock;
     __builtin_memset(&sock, 0, sizeof(sock));
@@ -163,6 +166,7 @@ int cg_connect4(struct bpf_sock_addr *ctx) {
    * dereferences through it (e.g. ctx->sk).  Scalar reads/writes
    * (ctx->user_ip4, ctx->user_port) remain fine.
    */
+  if (!conf->enable_udp) return 1;
   struct bpf_sock *sk = ctx->sk;
   if (!sk) return 1;
   __u16 src_port = sk->src_port;

cmd/tcpProxy.go

@@ -15,18 +15,27 @@ import (
 	"golang.org/x/net/proxy"
 )
 
-// StartProxy starts TCP (and optionally UDP when udpMap is not nil) proxy on proxyPort.
-func StartProxy(udpMap *ebpf.Map) {
+// StartProxy starts TCP/UDP proxy on proxyPort based on enableTCP/enableUDP.
+func StartProxy(udpMap *ebpf.Map, enableTCP bool, enableUDP bool) {
 	proxyAddr := fmt.Sprintf("127.0.0.1:%d", proxyPort)
-	listener, err := net.Listen("tcp", proxyAddr)
-	if err != nil {
-		log.Fatalf("Failed to start proxy server: %v", err)
+
+	if !enableTCP && !enableUDP {
+		log.Printf("Proxy: enableTCP and enableUDP are both false, nothing to start")
+		return
+	}
+
+	if enableTCP {
+		listener, err := net.Listen("tcp", proxyAddr)
+		if err != nil {
+			log.Fatalf("Failed to start TCP proxy server: %v", err)
+		}
+		log.Printf("TCP proxy server with PID %d listening on %s", os.Getpid(), proxyAddr)
+		go acceptLoop(listener)
 	}
 
-	log.Printf("Proxy server with PID %d listening on %s", os.Getpid(), proxyAddr)
-	go acceptLoop(listener)
-	if udpMap != nil {
+	if enableUDP && udpMap != nil {
 		go StartUDPProxy(proxyAddr, udpMap)
+		log.Printf("UDP proxy server with PID %d listening on %s", os.Getpid(), proxyAddr)
 	}
 }