fix: forward proxy response path with custom backend ports (#25)

Author: stillya

ebpf/include/backend.h

@@ -31,6 +31,15 @@ struct {
 	__type(value, __u32);
 } backend_count SEC(".maps");
 
+// Set of ports that backends listen on, used to detect FROM_WG return traffic.
+// Key: backend port (host byte order), Value: unused.
+struct {
+	__uint(type, BPF_MAP_TYPE_HASH);
+	__uint(max_entries, MAX_BACKENDS);
+	__type(key, __u16);
+	__type(value, __u8);
+} backend_port_set SEC(".maps");
+
 // jhash - Jenkins hash for consistent backend selection
 static __always_inline __u32 jhash_2words(__u32 a, __u32 b, __u32 initval) {
 	__u32 c = initval;

ebpf/wg_forward_proxy.c

@@ -240,50 +240,10 @@ int wg_forward_proxy(struct xdp_md *xdp_ctx) {
 	__u16 src_port = ctx.src_port;
 	__u16 dst_port = ctx.dst_port;
 	__u16 wg_port = CONFIG(wg_port);
-	if (dst_port != wg_port && src_port != wg_port)
-		return XDP_PASS;
 
 	__u8 is_to_wg = (dst_port == wg_port) ? 1 : 0;
-	__u8 is_from_wg = (src_port == wg_port) ? 1 : 0;
-
 	__u32 pkt_len = (void *)(long)xdp_ctx->data_end - (void *)(long)xdp_ctx->data;
 
-	if (likely(is_from_wg)) {
-		struct connection_key original_conn = { 0 };
-		if (restore_nat_connection(&ctx, &original_conn) < 0) {
-			DEBUG_PRINTK("Failed to restore NAT connection for FROM WG packet, passing "
-				     "through");
-			return XDP_PASS;
-		}
-
-		struct connection_value *conn_value = bpf_map_lookup_elem(&connection_map, &original_conn);
-		if (!conn_value) {
-			DEBUG_PRINTK("No connection value found for FROM WG packet");
-			return XDP_PASS;
-		}
-
-		__u8 backend_index = conn_value->backend_index;
-
-		// FROM_WG path: backend->proxy (upstream rx), proxy->client (downstream tx)
-		update_metrics(backend_index, METRIC_UPSTREAM, pkt_len, 1, METRIC_REASON_FORWARDED);
-
-		if (instr_deobfuscate_xdp(&ctx) < 0) {
-			DEBUG_PRINTK("Deobfuscation failed, dropping packet");
-			update_metrics(backend_index, METRIC_UPSTREAM, pkt_len, 1, METRIC_REASON_DROPPED);
-			return XDP_DROP;
-		}
-
-		__u32 tx_pkt_len = (void *)(long)xdp_ctx->data_end - (void *)(long)xdp_ctx->data;
-
-		__u32 dst_addr = bpf_ntohl(ctx.ip->daddr);
-		__u32 client_ip = bpf_ntohl(original_conn.client_ip);
-		__u16 server_port = original_conn.server_port;
-		__u16 client_port = original_conn.client_port;
-
-		update_metrics(backend_index, METRIC_DOWNSTREAM, tx_pkt_len, 0, METRIC_REASON_FORWARDED);
-		return forward_packet(&ctx, dst_addr, server_port, client_ip, client_port);
-	}
-
 	if (unlikely(is_to_wg)) {
 		struct backend_entry backend = { 0 };
 		if (create_nat_connection(&ctx, &backend) < 0) {
@@ -326,6 +286,44 @@ int wg_forward_proxy(struct xdp_md *xdp_ctx) {
 		return forward_packet(&ctx, proxy_ip, conn_value->nat_port, server_ip, target_port);
 	}
 
+	__u8 is_from_wg = bpf_map_lookup_elem(&backend_port_set, &src_port) != NULL ? 1 : 0;
+
+	if (likely(is_from_wg)) {
+		struct connection_key original_conn = { 0 };
+		if (restore_nat_connection(&ctx, &original_conn) < 0) {
+			DEBUG_PRINTK("Failed to restore NAT connection for FROM WG packet, passing "
+				     "through");
+			return XDP_PASS;
+		}
+
+		struct connection_value *conn_value = bpf_map_lookup_elem(&connection_map, &original_conn);
+		if (!conn_value) {
+			DEBUG_PRINTK("No connection value found for FROM WG packet");
+			return XDP_PASS;
+		}
+
+		__u8 backend_index = conn_value->backend_index;
+
+		// FROM_WG path: backend->proxy (upstream rx), proxy->client (downstream tx)
+		update_metrics(backend_index, METRIC_UPSTREAM, pkt_len, 1, METRIC_REASON_FORWARDED);
+
+		if (instr_deobfuscate_xdp(&ctx) < 0) {
+			DEBUG_PRINTK("Deobfuscation failed, dropping packet");
+			update_metrics(backend_index, METRIC_UPSTREAM, pkt_len, 1, METRIC_REASON_DROPPED);
+			return XDP_DROP;
+		}
+
+		__u32 tx_pkt_len = (void *)(long)xdp_ctx->data_end - (void *)(long)xdp_ctx->data;
+
+		__u32 dst_addr = bpf_ntohl(ctx.ip->daddr);
+		__u32 client_ip = bpf_ntohl(original_conn.client_ip);
+		__u16 server_port = original_conn.server_port;
+		__u16 client_port = original_conn.client_port;
+
+		update_metrics(backend_index, METRIC_DOWNSTREAM, tx_pkt_len, 0, METRIC_REASON_FORWARDED);
+		return forward_packet(&ctx, dst_addr, server_port, client_ip, client_port);
+	}
+
 	DEBUG_PRINTK("No matching handler for WG packet, passing through");
 	return XDP_PASS;
 }

ebpf/wg_forward_proxy_test.go

@@ -332,6 +332,40 @@ func TestMultipleBackends(t *testing.T) {
 	}
 }
 
+func TestBackendCustomPort(t *testing.T) {
+	spec, err := LoadWgForwardProxy()
+	if err != nil {
+		t.Fatalf("Failed to load spec: %v", err)
+	}
+
+	setVar(t, spec, "__cfg_xor_enabled", false)
+	setVar(t, spec, "__cfg_wg_port", uint16(wgPort))
+
+	objs := &WgForwardProxyObjects{}
+	if err := spec.LoadAndAssign(objs, nil); err != nil {
+		t.Fatalf("Failed to load objects: %v", err)
+	}
+	defer objs.Close()
+
+	if err := configureBackends(objs, []config.BackendServer{
+		{IP: "10.0.0.1", Port: 11580},
+	}); err != nil {
+		t.Fatalf("Failed to configure backends: %v", err)
+	}
+
+	packet := createWGPacket("192.168.1.1", "192.168.1.2", 12345, wgPort)
+	result, outputPacket, err := objs.WgForwardProxy.Test(packet)
+	if err != nil {
+		t.Fatalf("Failed to run program: %v", err)
+	}
+
+	if int(result) != xdpRedirect {
+		t.Errorf("Expected XDP_REDIRECT, got %d", result)
+	}
+
+	verifyPacket(t, outputPacket, "10.0.0.1", 11580)
+}
+
 func TestWgPortConfig(t *testing.T) {
 	tests := []struct {
 		name          string
@@ -568,5 +602,20 @@ func configureBackends(objs *WgForwardProxyObjects, backends []config.BackendSer
 
 	countKey := uint32(0)
 	count := uint32(len(backends)) //nolint:gosec // G304: it's fine
-	return objs.BackendCount.Put(&countKey, &count)
+	if err := objs.BackendCount.Put(&countKey, &count); err != nil {
+		return err
+	}
+
+	dummy := uint8(1)
+	for _, backend := range backends {
+		port := backend.Port
+		if port == 0 {
+			port = wgPort
+		}
+		if err := objs.BackendPortSet.Put(&port, &dummy); err != nil {
+			return err
+		}
+	}
+
+	return nil
 }

ebpf/wg_reverse_proxy_test.go

@@ -451,6 +451,79 @@ func createObfuscatedAndPaddedWGPacket(srcIP, dstIP string, srcPort, dstPort uin
 	return packet
 }
 
+func TestReverseProxyWgPortConfig(t *testing.T) {
+	xorKey := "test-key-1234567890abcdef12345678"
+	keyBytes := []byte(xorKey)
+	var keyArray [32]byte
+	copy(keyArray[:], keyBytes)
+
+	tests := []struct {
+		name          string
+		wgPort        uint16
+		packetDstPort uint16
+		shouldProcess bool
+	}{
+		{
+			name:          "default_port",
+			wgPort:        51820,
+			packetDstPort: 51820,
+			shouldProcess: true,
+		},
+		{
+			name:          "custom_port",
+			wgPort:        11580,
+			packetDstPort: 11580,
+			shouldProcess: true,
+		},
+		{
+			name:          "wrong_port",
+			wgPort:        51820,
+			packetDstPort: 9999,
+			shouldProcess: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			spec, err := LoadWgReverseProxy()
+			if err != nil {
+				t.Fatalf("Failed to load spec: %v", err)
+			}
+
+			if err := spec.Variables["__cfg_xor_enabled"].Set(true); err != nil {
+				t.Fatalf("Failed to set xor_enabled: %v", err)
+			}
+			if err := spec.Variables["__cfg_xor_key"].Set(keyArray); err != nil {
+				t.Fatalf("Failed to set xor_key: %v", err)
+			}
+			if err := spec.Variables["__cfg_padding_enabled"].Set(false); err != nil {
+				t.Fatalf("Failed to set padding_enabled: %v", err)
+			}
+			if err := spec.Variables["__cfg_wg_port"].Set(tt.wgPort); err != nil {
+				t.Fatalf("Failed to set wg_port: %v", err)
+			}
+
+			objs := &WgReverseProxyObjects{}
+			if err := spec.LoadAndAssign(objs, nil); err != nil {
+				t.Fatalf("Failed to load objects: %v", err)
+			}
+			defer objs.Close()
+
+			inputPacket := createObfuscatedWGPacket("192.168.1.1", "192.168.1.2", 12345, tt.packetDstPort, keyBytes)
+			_, outputPacket, err := objs.WgReverseProxy.Test(inputPacket)
+			if err != nil {
+				t.Fatalf("Failed to run program: %v", err)
+			}
+
+			if tt.shouldProcess {
+				verifyXORDeobfuscation(t, inputPacket, outputPacket, keyBytes)
+			} else {
+				verifyPayloadUnchanged(t, inputPacket, outputPacket)
+			}
+		})
+	}
+}
+
 // TestReverseProxyPaddingObfuscateMTUExceededDrop verifies that a FROM_WG packet is dropped
 // when adding padding would cause it to exceed the configured link MTU.
 func TestReverseProxyPaddingObfuscateMTUExceededDrop(t *testing.T) {