Fix the handling on local mode in AwsProxyEncoderHelper

bplessis-swi · bplessis-swi · commit d582239339ae · 2025-11-27T19:31:49.000+01:00
diff --git a/proxy-socket-core/src/main/java/net/airvantage/proxysocket/core/v2/ProxyProtocolV2Decoder.java b/proxy-socket-core/src/main/java/net/airvantage/proxysocket/core/v2/ProxyProtocolV2Decoder.java
@@ -65,17 +65,11 @@ public static ProxyHeader parse(byte[] data, int offset, int length, boolean par
             throw new ProxyProtocolParseException("Invalid version");
         }
 
-        Command command;
-        switch (cmd) {
-            case 0x00:
-                // Early return for LOCAL command
-                return new ProxyHeader(Command.LOCAL, AddressFamily.AF_UNSPEC, TransportProtocol.UNSPEC, null, null, null, PROTOCOL_SIGNATURE_FIXED_LENGTH);
-            case 0x01:
-                command = Command.PROXY;
-                break;
-            default:
-                throw new ProxyProtocolParseException("Invalid command");
-        }
+        Command command = switch (cmd) {
+            case 0x00 -> Command.LOCAL;
+            case 0x01 -> Command.PROXY;
+            default -> throw new ProxyProtocolParseException("Invalid command");
+        };
 
         // Byte 14: address family and protocol
         int famProto = data[pos++] & 0xFF;
diff --git a/proxy-socket-core/src/test/java/net/airvantage/proxysocket/core/v2/AwsProxyEncoderHelper.java b/proxy-socket-core/src/test/java/net/airvantage/proxysocket/core/v2/AwsProxyEncoderHelper.java
@@ -28,9 +28,18 @@ public final class AwsProxyEncoderHelper {
     private final Header header = new Header();
 
     public AwsProxyEncoderHelper command(ProxyHeader.Command cmd) {
-        this.command = cmd == ProxyHeader.Command.LOCAL
-                ? ProxyProtocolSpec.Command.LOCAL
-                : ProxyProtocolSpec.Command.PROXY;
+        if (cmd == ProxyHeader.Command.LOCAL) {
+            this.command = ProxyProtocolSpec.Command.LOCAL;
+
+            // Spec clearly state that for LOCAL command, we
+            // 1. must discard the protocol block including the family and
+            // 2. \x00 is expected to be used for the protocol field.
+            this.family = ProxyProtocolSpec.AddressFamily.AF_UNSPEC;
+            this.protocol = ProxyProtocolSpec.TransportProtocol.UNSPEC;
+        } else {
+            this.command = ProxyProtocolSpec.Command.PROXY;
+        }
+
         return this;
     }
 
@@ -76,13 +85,7 @@ public byte[] build() throws IOException {
         header.setAddressFamily(family);
         header.setTransportProtocol(protocol);
 
-        // AWS ProProt validates addresses even for LOCAL command, set dummy values
-        if (command == ProxyProtocolSpec.Command.LOCAL && source == null) {
-            header.setSrcAddress(new byte[]{0, 0, 0, 0});
-            header.setDstAddress(new byte[]{0, 0, 0, 0});
-            header.setSrcPort(0);
-            header.setDstPort(0);
-        } else {
+        if (command != ProxyProtocolSpec.Command.LOCAL) {
             if (source != null) {
                 header.setSrcAddress(source.getAddress().getAddress());
                 header.setSrcPort(source.getPort());
diff --git a/proxy-socket-udp/src/main/java/net/airvantage/proxysocket/udp/ProxyDatagramSocket.java b/proxy-socket-udp/src/main/java/net/airvantage/proxysocket/udp/ProxyDatagramSocket.java
@@ -101,7 +101,7 @@ public void receive(DatagramPacket packet)
             }
 
             int headerLen = header.getHeaderLength();
-            LOG.trace("Stripping header: {} bytes, remaining length: {}", headerLen, packet.getLength() - headerLen);
+            LOG.trace("Stripping header: {} bytes, original length: {}", headerLen, packet.getLength());
             packet.setData(packet.getData(), packet.getOffset() + headerLen, packet.getLength() - headerLen);
         } catch (ProxyProtocolParseException e) {
             LOG.warn("Proxy socket parse error; delivering original packet.", e);
diff --git a/proxy-socket-udp/src/test/java/net/airvantage/proxysocket/udp/ProxyDatagramSocketTest.java b/proxy-socket-udp/src/test/java/net/airvantage/proxysocket/udp/ProxyDatagramSocketTest.java
@@ -21,23 +21,33 @@
 import static org.junit.jupiter.api.Assertions.*;
 import static org.mockito.Mockito.*;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 /**
  * Unit tests for ProxyDatagramSocket IP address mapping and cache behavior.
  */
-class ProxyDatagramSocketIPMappingTest {
+class ProxyDatagramSocketTest {
+    private static final Logger LOG = LoggerFactory.getLogger(ProxyDatagramSocket.class);
 
     private ProxyDatagramSocket socket;
     private ProxyAddressCache mockCache;
     private ProxyProtocolMetricsListener mockMetrics;
+
+    private InetSocketAddress realClient;
+    private InetSocketAddress serviceAddress;
+    private InetSocketAddress backendAddress;
     private int localPort;
 
     @BeforeEach
     void setUp() throws Exception {
         mockCache = mock(ProxyAddressCache.class);
         mockMetrics = mock(ProxyProtocolMetricsListener.class);
 
+        realClient = new InetSocketAddress(InetAddress.getLoopbackAddress(), 12345);
+        serviceAddress = new InetSocketAddress(InetAddress.getLoopbackAddress(), 54321);
         socket = new ProxyDatagramSocket(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), mockCache, mockMetrics, null);
         localPort = socket.getLocalPort();
+        backendAddress = new InetSocketAddress(InetAddress.getLoopbackAddress(), localPort);
     }
 
     @AfterEach
@@ -48,27 +58,24 @@ void tearDown() {
     }
 
     @Test
-    void receive_withValidProxyHeader_populatesCache() throws Exception {
+    void receive_withValidProxyHeader() throws Exception {
         // Arrange
-        InetSocketAddress realClient = new InetSocketAddress("10.1.2.3", 12345);
-        InetSocketAddress lbAddress = new InetSocketAddress("127.0.0.1", 54321);
         byte[] payload = "test-data".getBytes(StandardCharsets.UTF_8);
 
         var proxyHeader = new AwsProxyEncoderHelper()
                 .family(ProxyHeader.AddressFamily.AF_INET)
                 .socket(ProxyHeader.TransportProtocol.DGRAM)
                 .source(realClient)
-                .destination(new InetSocketAddress("127.0.0.1", localPort))
+                .destination(serviceAddress)
                 .build();
 
         byte[] packet = new byte[proxyHeader.length + payload.length];
         System.arraycopy(proxyHeader, 0, packet, 0, proxyHeader.length);
         System.arraycopy(payload, 0, packet, proxyHeader.length, payload.length);
 
         // Create a loopback socket to send from
-        try (java.net.DatagramSocket sender = new java.net.DatagramSocket(lbAddress)) {
-            sender.send(new DatagramPacket(packet, packet.length,
-                    new InetSocketAddress("127.0.0.1", localPort)));
+        try (java.net.DatagramSocket sender = new java.net.DatagramSocket(serviceAddress)) {
+            sender.send(new DatagramPacket(packet, packet.length, backendAddress));
         }
 
         // Act
@@ -82,8 +89,8 @@ void receive_withValidProxyHeader_populatesCache() throws Exception {
         verify(mockCache).put(clientCaptor.capture(), lbCaptor.capture());
 
         assertEquals(realClient, clientCaptor.getValue());
-        assertEquals(lbAddress.getAddress(), lbCaptor.getValue().getAddress());
-        assertEquals(lbAddress.getPort(), lbCaptor.getValue().getPort());
+        assertEquals(serviceAddress.getAddress(), lbCaptor.getValue().getAddress());
+        assertEquals(serviceAddress.getPort(), lbCaptor.getValue().getPort());
 
         // Verify packet was modified to show real client address
         assertEquals(realClient, receivePacket.getSocketAddress());
@@ -96,24 +103,21 @@ void receive_withValidProxyHeader_populatesCache() throws Exception {
                         receivePacket.getOffset() + receivePacket.getLength()));
 
         verify(mockMetrics).onHeaderParsed(any(ProxyHeader.class));
-        verify(mockMetrics).onTrustedProxy(lbAddress.getAddress());
+        verify(mockMetrics).onTrustedProxy(serviceAddress.getAddress());
         verify(mockMetrics, never()).onUntrustedProxy(any());
         verify(mockMetrics, never()).onParseError(any());
         verify(mockMetrics, never()).onLocal(any());
     }
 
     @Test
     void send_withCacheHit_usesLoadBalancerAddress() throws Exception {
-        // Arrange
-        InetSocketAddress realClient = new InetSocketAddress("10.1.2.3", 12345);
-        InetSocketAddress lbAddress = new InetSocketAddress("127.0.0.1", 54321);
         byte[] payload = "response".getBytes(StandardCharsets.UTF_8);
 
         // Mock cache to return lb address
-        when(mockCache.get(realClient)).thenReturn(lbAddress);
+        when(mockCache.get(realClient)).thenReturn(serviceAddress);
 
         // Create a receiver to verify the packet destination
-        java.net.DatagramSocket receiver = new java.net.DatagramSocket(lbAddress);
+        java.net.DatagramSocket receiver = new java.net.DatagramSocket(serviceAddress);
         receiver.setSoTimeout(1000);
 
         try {
@@ -142,37 +146,35 @@ void send_withCacheHit_usesLoadBalancerAddress() throws Exception {
     }
 
     @Test
-    void send_withCacheMiss_usesOriginalAddress() throws Exception {
+    void send_withCacheMiss_dropsPacket() throws Exception {
         // Arrange
-        InetSocketAddress clientAddress = new InetSocketAddress("127.0.0.1", 55555);
         byte[] payload = "response".getBytes(StandardCharsets.UTF_8);
 
         // Mock cache to return null (cache miss)
-        when(mockCache.get(clientAddress)).thenReturn(null);
+        when(mockCache.get(realClient)).thenReturn(null);
 
-        // Create a receiver at the client address
-        java.net.DatagramSocket receiver = new java.net.DatagramSocket(clientAddress);
-        receiver.setSoTimeout(1000);
+        // Create a receiver at the client address to verify packet is NOT sent
+        java.net.DatagramSocket receiver = new java.net.DatagramSocket(realClient);
+        receiver.setSoTimeout(500); // Short timeout since we expect no packet
 
         try {
             // Act - send to client address
-            DatagramPacket sendPacket = new DatagramPacket(payload, payload.length, clientAddress);
+            DatagramPacket sendPacket = new DatagramPacket(payload, payload.length, realClient);
             socket.send(sendPacket);
 
-            // Verify packet was sent to original address
+            // Try to receive - should timeout since packet was dropped
             byte[] receiveBuf = new byte[2048];
             DatagramPacket receivePacket = new DatagramPacket(receiveBuf, receiveBuf.length);
-            receiver.receive(receivePacket);
 
-            // Assert
-            assertArrayEquals(payload,
-                    java.util.Arrays.copyOfRange(receivePacket.getData(), 0, receivePacket.getLength()));
+            assertThrows(java.net.SocketTimeoutException.class, () -> {
+                receiver.receive(receivePacket);
+            }, "Expected packet to be dropped on cache miss");
 
             // Verify cache was queried
-            verify(mockCache).get(clientAddress);
+            verify(mockCache).get(realClient);
 
             // Verify metrics - cache miss
-            verify(mockMetrics).onCacheMiss(clientAddress);
+            verify(mockMetrics).onCacheMiss(realClient);
             verify(mockMetrics, never()).onCacheHit(any());
         } finally {
             receiver.close();
@@ -187,20 +189,24 @@ void receive_withLocalCommand_doesNotPopulateCache() throws Exception {
                 .command(ProxyHeader.Command.LOCAL)
                 .build();
 
+        LOG.trace("Payload:\n{}", Utility.hexdump(payload, 0, payload.length));
+        LOG.trace("Proxy header:\n{}", Utility.hexdump(proxyHeader, 0, proxyHeader.length));
+
         byte[] packet = new byte[proxyHeader.length + payload.length];
         System.arraycopy(proxyHeader, 0, packet, 0, proxyHeader.length);
         System.arraycopy(payload, 0, packet, proxyHeader.length, payload.length);
 
         try (java.net.DatagramSocket sender = new java.net.DatagramSocket()) {
-            sender.send(new DatagramPacket(packet, packet.length,
-                    new InetSocketAddress("127.0.0.1", localPort)));
+            sender.send(new DatagramPacket(packet, packet.length, backendAddress));
         }
 
         // Act
         byte[] receiveBuf = new byte[2048];
         DatagramPacket receivePacket = new DatagramPacket(receiveBuf, receiveBuf.length);
         socket.receive(receivePacket);
 
+        LOG.trace("Received packet length={}, content:\n{}", receivePacket.getLength(), Utility.hexdump(receivePacket.getData(), receivePacket.getOffset(), receivePacket.getLength()));
+
         // Assert - cache should NOT be populated for LOCAL commands
         verify(mockCache, never()).put(any(), any());
 
@@ -218,17 +224,16 @@ void receive_withTcpProtocol_doesNotPopulateCache() throws Exception {
         byte[] proxyHeader = new AwsProxyEncoderHelper()
                 .family(ProxyHeader.AddressFamily.AF_INET)
                 .socket(ProxyHeader.TransportProtocol.STREAM) // TCP, not UDP
-                .source(new InetSocketAddress("10.1.2.3", 12345))
-                .destination(new InetSocketAddress("127.0.0.1", localPort))
+                .source(realClient)
+                .destination(serviceAddress)
                 .build();
 
         byte[] packet = new byte[proxyHeader.length + payload.length];
         System.arraycopy(proxyHeader, 0, packet, 0, proxyHeader.length);
         System.arraycopy(payload, 0, packet, proxyHeader.length, payload.length);
 
         try (java.net.DatagramSocket sender = new java.net.DatagramSocket()) {
-            sender.send(new DatagramPacket(packet, packet.length,
-                    new InetSocketAddress("127.0.0.1", localPort)));
+            sender.send(new DatagramPacket(packet, packet.length, backendAddress));
         }
 
         // Act
@@ -247,14 +252,13 @@ void receive_withTcpProtocol_doesNotPopulateCache() throws Exception {
     @Test
     void receive_withValidProxyHeader_callsMetricsOnHeaderParsed() throws Exception {
         // Arrange
-        InetSocketAddress realClient = new InetSocketAddress("10.1.2.3", 12345);
         byte[] payload = "test".getBytes(StandardCharsets.UTF_8);
 
         byte[] proxyHeader = new AwsProxyEncoderHelper()
                 .family(ProxyHeader.AddressFamily.AF_INET)
                 .socket(ProxyHeader.TransportProtocol.DGRAM)
                 .source(realClient)
-                .destination(new InetSocketAddress("127.0.0.1", localPort))
+                .destination(serviceAddress)
                 .build();
 
         byte[] packet = new byte[proxyHeader.length + payload.length];
@@ -263,8 +267,7 @@ void receive_withValidProxyHeader_callsMetricsOnHeaderParsed() throws Exception
 
         // Send packet
         try (java.net.DatagramSocket sender = new java.net.DatagramSocket()) {
-            sender.send(new DatagramPacket(packet, packet.length,
-                    new InetSocketAddress("127.0.0.1", localPort)));
+            sender.send(new DatagramPacket(packet, packet.length, backendAddress));
         }
 
         // Act
@@ -288,8 +291,7 @@ void receive_withInvalidData_callsMetricsOnParseError() throws Exception {
         byte[] garbage = "not-a-proxy-header".getBytes(StandardCharsets.UTF_8);
 
         try (java.net.DatagramSocket sender = new java.net.DatagramSocket()) {
-            sender.send(new DatagramPacket(garbage, garbage.length,
-                    new InetSocketAddress("127.0.0.1", localPort)));
+            sender.send(new DatagramPacket(garbage, garbage.length, backendAddress));
         }
 
         // Act
@@ -304,67 +306,5 @@ void receive_withInvalidData_callsMetricsOnParseError() throws Exception {
         assertEquals(garbage.length, receivePacket.getLength());
     }
 
-    @Test
-    void send_withCacheHit_callsMetricsOnCacheHit() throws Exception {
-        // Arrange
-        InetSocketAddress realClient = new InetSocketAddress("10.1.2.3", 12345);
-        InetSocketAddress lbAddress = new InetSocketAddress("127.0.0.1", 54321);
-        byte[] payload = "response".getBytes(StandardCharsets.UTF_8);
-
-        // Mock cache to return lb address
-        when(mockCache.get(realClient)).thenReturn(lbAddress);
-
-        // Create a receiver to verify the packet destination
-        java.net.DatagramSocket receiver = new java.net.DatagramSocket(lbAddress);
-        receiver.setSoTimeout(1000);
-
-        try {
-            // Act - send to real client, should be redirected to LB
-            DatagramPacket sendPacket = new DatagramPacket(payload, payload.length, realClient);
-            socket.send(sendPacket);
-
-            // Receive the packet (to avoid timeout)
-            byte[] receiveBuf = new byte[2048];
-            DatagramPacket receivePacket = new DatagramPacket(receiveBuf, receiveBuf.length);
-            receiver.receive(receivePacket);
-
-            // Assert - onCacheHit should be called
-            verify(mockMetrics).onCacheHit(realClient);
-            verify(mockMetrics, never()).onCacheMiss(any());
-        } finally {
-            receiver.close();
-        }
-    }
-
-    @Test
-    void send_withCacheMiss_callsMetricsOnCacheMiss() throws Exception {
-        // Arrange
-        InetSocketAddress clientAddress = new InetSocketAddress("127.0.0.1", 55555);
-        byte[] payload = "response".getBytes(StandardCharsets.UTF_8);
-
-        // Mock cache to return null (cache miss)
-        when(mockCache.get(clientAddress)).thenReturn(null);
-
-        // Create a receiver at the client address
-        java.net.DatagramSocket receiver = new java.net.DatagramSocket(clientAddress);
-        receiver.setSoTimeout(1000);
-
-        try {
-            // Act - send to client address
-            DatagramPacket sendPacket = new DatagramPacket(payload, payload.length, clientAddress);
-            socket.send(sendPacket);
-
-            // Receive the packet (to avoid timeout)
-            byte[] receiveBuf = new byte[2048];
-            DatagramPacket receivePacket = new DatagramPacket(receiveBuf, receiveBuf.length);
-            receiver.receive(receivePacket);
-
-            // Assert - onCacheMiss should be called
-            verify(mockMetrics).onCacheMiss(clientAddress);
-            verify(mockMetrics, never()).onCacheHit(any());
-        } finally {
-            receiver.close();
-        }
-    }
 }
 
diff --git a/proxy-socket-udp/src/test/java/net/airvantage/proxysocket/udp/ProxyDatagramSocketUnTrustedProxyTest.java b/proxy-socket-udp/src/test/java/net/airvantage/proxysocket/udp/ProxyDatagramSocketUnTrustedProxyTest.java
@@ -25,7 +25,7 @@
 /**
  * Unit tests for ProxyDatagramSocket with untrusted proxy source.
  */
-class ProxyDatagramSocketMetricsTest {
+class ProxyDatagramSocketUnTrustedProxyTest {
 
     private ProxyDatagramSocket socket;
     private ProxyAddressCache mockCache;
diff --git a/proxy-socket-udp/src/test/java/net/airvantage/proxysocket/udp/Utility.java b/proxy-socket-udp/src/test/java/net/airvantage/proxysocket/udp/Utility.java

Original file line number	Diff line number	Diff line change
`@@ -101,7 +101,7 @@ public void receive(DatagramPacket packet)`
`101`	`101`	`}`
`102`	`102`
`103`	`103`	`int headerLen = header.getHeaderLength();`
`104`		`- LOG.trace("Stripping header: {} bytes, remaining length: {}", headerLen, packet.getLength() - headerLen);`
	`104`	`+ LOG.trace("Stripping header: {} bytes, original length: {}", headerLen, packet.getLength());`
`105`	`105`	`packet.setData(packet.getData(), packet.getOffset() + headerLen, packet.getLength() - headerLen);`
`106`	`106`	`} catch (ProxyProtocolParseException e) {`
`107`	`107`	`LOG.warn("Proxy socket parse error; delivering original packet.", e);`