process review step 2

Author: bplessis-swi

proxy-socket-core/src/main/java/net/airvantage/proxysocket/tools/SubnetPredicate.java

@@ -2,24 +2,24 @@
  * BSD-3-Clause License.
  * Copyright (c) 2025 Semtech
  */
-package net.airvantage.proxysocket.udp;
+package net.airvantage.proxysocket.tools;
 
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.UnknownHostException;
 import java.util.function.Predicate;
 
 /**
- * Predicate that tests whether an InetSocketAddress belongs to a given subnet (CIDR).
+ * Predicate compatible class that tests whether an InetSocketAddress belongs to a given subnet (CIDR).
  * Supports both IPv4 and IPv6 CIDR notation.
  *
  * <p>Example usage:
  * <pre>
  * // Single subnet
- * socket.setTrustedProxy(new SubnetPredicate("10.0.0.0/8"));
+ * Predicate<InetSocketAddress> predicate = new SubnetPredicate("10.0.0.0/8");
  *
  * // Multiple subnets
- * socket.setTrustedProxy(
+ * Predicate<InetSocketAddress> predicate =
  *     new SubnetPredicate("10.0.0.0/8")
  *         .or(new SubnetPredicate("192.168.0.0/16"))
  *         .or(new SubnetPredicate("2001:db8::/32"))

proxy-socket-core/src/test/java/net/airvantage/proxysocket/tools/SubnetPredicateTest.java

@@ -8,6 +8,7 @@
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.UnknownHostException;
+import java.util.function.Predicate;
 
 import static org.junit.jupiter.api.Assertions.*;
 
@@ -189,7 +190,7 @@ void testIPv4vsIPv6_NoMatch() throws UnknownHostException {
     void testPredicateComposition_Or() throws UnknownHostException {
         SubnetPredicate predicate1 = new SubnetPredicate("192.168.1.0/24");
         SubnetPredicate predicate2 = new SubnetPredicate("10.0.0.0/8");
-        SubnetPredicate combined = predicate1.or(predicate2);
+        Predicate<InetSocketAddress> combined = predicate1.or(predicate2);
 
         assertTrue(combined.test(addr("192.168.1.100", 8080)));
         assertTrue(combined.test(addr("10.20.30.40", 8080)));
@@ -200,7 +201,7 @@ void testPredicateComposition_Or() throws UnknownHostException {
     void testPredicateComposition_And() throws UnknownHostException {
         SubnetPredicate predicate1 = new SubnetPredicate("192.168.0.0/16");
         SubnetPredicate predicate2 = new SubnetPredicate("192.168.1.0/24");
-        SubnetPredicate combined = predicate1.and(predicate2);
+        Predicate<InetSocketAddress> combined = predicate1.and(predicate2);
 
         assertTrue(combined.test(addr("192.168.1.100", 8080)));
         assertFalse(combined.test(addr("192.168.2.100", 8080)));

proxy-socket-udp/pom.xml

@@ -18,12 +18,10 @@
             <artifactId>proxy-socket-core</artifactId>
             <version>${project.version}</version>
         </dependency>
-        <!-- slf4j as optional -->
         <dependency>
             <groupId>org.slf4j</groupId>
             <artifactId>slf4j-api</artifactId>
-            <version>2.0.16</version>
-            <optional>true</optional>
+            <version>2.0.17</version>
         </dependency>
         <!-- tests -->
         <dependency>
@@ -36,7 +34,7 @@
         <dependency>
             <groupId>org.slf4j</groupId>
             <artifactId>slf4j-simple</artifactId>
-            <version>2.0.16</version>
+            <version>2.0.17</version>
             <scope>test</scope>
         </dependency>
     </dependencies>

proxy-socket-udp/src/main/java/net/airvantage/proxysocket/udp/ProxyDatagramSocket.java

@@ -6,15 +6,19 @@
 
 import net.airvantage.proxysocket.core.ProxyAddressCache;
 import net.airvantage.proxysocket.core.ProxyProtocolMetricsListener;
+import net.airvantage.proxysocket.core.ProxyProtocolParseException;
 import net.airvantage.proxysocket.core.v2.ProxyHeader;
 import net.airvantage.proxysocket.core.v2.ProxyProtocolV2Decoder;
 
 import java.io.IOException;
 import java.net.DatagramPacket;
 import java.net.DatagramSocket;
 import java.net.InetSocketAddress;
+import java.net.PortUnreachableException;
 import java.net.SocketAddress;
 import java.net.SocketException;
+import java.net.SocketTimeoutException;
+import java.nio.channels.IllegalBlockingModeException;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import java.util.function.Predicate;
@@ -30,30 +34,38 @@
 public class ProxyDatagramSocket extends DatagramSocket {
     private static final Logger LOG = Logger.getLogger(ProxyDatagramSocket.class.getName());
 
-    private ProxyAddressCache addressCache;
-    private ProxyProtocolMetricsListener metrics;
-    private Predicate<InetSocketAddress> trustedProxyPredicate;
+    private final ProxyAddressCache addressCache;
+    private final ProxyProtocolMetricsListener metrics;
+    private final Predicate<InetSocketAddress> trustedProxyPredicate;
 
-    public ProxyDatagramSocket() throws SocketException {
+    public ProxyDatagramSocket(ProxyAddressCache cache, ProxyProtocolMetricsListener metrics, Predicate<InetSocketAddress> predicate) throws SocketException {
         super();
+        this.addressCache = cache;
+        this.metrics = metrics;
+        this.trustedProxyPredicate = predicate;
     }
 
-    public ProxyDatagramSocket(SocketAddress bindaddr) throws SocketException {
+    public ProxyDatagramSocket(SocketAddress bindaddr, ProxyAddressCache cache, ProxyProtocolMetricsListener metrics, Predicate<InetSocketAddress> predicate) throws SocketException {
         super(bindaddr);
+        this.addressCache = cache;
+        this.metrics = metrics;
+        this.trustedProxyPredicate = predicate;
     }
 
-    public ProxyDatagramSocket(int port) throws SocketException {
+    public ProxyDatagramSocket(int port, ProxyAddressCache cache, ProxyProtocolMetricsListener metrics, Predicate<InetSocketAddress> predicate) throws SocketException {
         super(port);
+        this.addressCache = cache;
+        this.metrics = metrics;
+        this.trustedProxyPredicate = predicate;
     }
 
-    public ProxyDatagramSocket(int port, java.net.InetAddress laddr) throws SocketException {
+    public ProxyDatagramSocket(int port, java.net.InetAddress laddr, ProxyAddressCache cache, ProxyProtocolMetricsListener metrics, Predicate<InetSocketAddress> predicate) throws SocketException {
         super(port, laddr);
+        this.addressCache = cache;
+        this.metrics = metrics;
+        this.trustedProxyPredicate = predicate;
     }
 
-    public ProxyDatagramSocket setCache(ProxyAddressCache cache) { this.addressCache = cache; return this; }
-    public ProxyDatagramSocket setMetrics(ProxyProtocolMetricsListener metrics) { this.metrics = metrics; return this; }
-    public ProxyDatagramSocket setTrustedProxy(Predicate<InetSocketAddress> predicate) { this.trustedProxyPredicate = predicate; return this; }
-
     @Override
     public void receive(DatagramPacket packet)
         throws IOException, SocketTimeoutException, PortUnreachableException, IllegalBlockingModeException {
@@ -102,12 +114,14 @@ public void send(DatagramPacket packet) throws IOException {
         if (lb != null) {
             packet.setSocketAddress(lb);
             if (metrics != null) metrics.onCacheHit(client);
-        } else {
+        } else if (addressCache != null) {
+            // Cache miss: unable to map client to load balancer address,
+            LOG.log(Level.DEBUG, "Cache miss for client {0}; unable to map to load balancer address, dropping packet.", client);
             if (metrics != null) metrics.onCacheMiss(client);
+            return;
+        // } else {
+            // No cache: deliver original packet
         }
-
         super.send(packet);
     }
 }
-
-

…dp/ProxyDatagramSocketIPMappingTest.java …xysocket/udp/ProxyDatagramSockeTest.javaproxy-socket-udp/src/test/java/net/airvantage/proxysocket/udp/ProxyDatagramSocketIPMappingTest.java renamed to proxy-socket-udp/src/test/java/net/airvantage/proxysocket/udp/ProxyDatagramSockeTest.java proxy-socket-udp/src/test/java/net/airvantage/proxysocket/udp/ProxyDatagramSocketIPMappingTest.java renamed to proxy-socket-udp/src/test/java/net/airvantage/proxysocket/udp/ProxyDatagramSockeTest.java

@@ -36,11 +36,7 @@ void setUp() throws Exception {
         mockCache = mock(ProxyAddressCache.class);
         mockMetrics = mock(ProxyProtocolMetricsListener.class);
 
-        socket = new ProxyDatagramSocket(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0))
-                .setCache(mockCache)
-                .setMetrics(mockMetrics)
-                .setTrustedProxy(addr -> true); // Trust all for these tests
-
+        socket = new ProxyDatagramSocket(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), mockCache, mockMetrics, null);
         localPort = socket.getLocalPort();
     }
 
@@ -58,8 +54,8 @@ void receive_withValidProxyHeader_populatesCache() throws Exception {
         InetSocketAddress lbAddress = new InetSocketAddress("127.0.0.1", 54321);
         byte[] payload = "test-data".getBytes(StandardCharsets.UTF_8);
 
-        byte[] proxyHeader = new ProxyProtocolV2Encoder()
-                .family(ProxyHeader.AddressFamily.INET4)
+        var proxyHeader = new AwsProxyEncoderHelper()
+                .family(ProxyHeader.AddressFamily.AF_INET)
                 .socket(ProxyHeader.TransportProtocol.DGRAM)
                 .source(realClient)
                 .destination(new InetSocketAddress("127.0.0.1", localPort))
@@ -177,17 +173,13 @@ void send_withCacheMiss_usesOriginalAddress() throws Exception {
         }
     }
 
-    @Test
-    void receive_withUntrustedProxy_skipsProcessing() throws Exception {
-        // Arrange - configure to reject all sources
-        socket.setTrustedProxy(addr -> false);
 
-        byte[] payload = "test".getBytes(StandardCharsets.UTF_8);
-        byte[] proxyHeader = new ProxyProtocolV2Encoder()
-                .family(ProxyHeader.AddressFamily.INET4)
-                .socket(ProxyHeader.TransportProtocol.DGRAM)
-                .source(new InetSocketAddress("10.1.2.3", 12345))
-                .destination(new InetSocketAddress("127.0.0.1", localPort))
+    @Test
+    void receive_withLocalCommand_doesNotPopulateCache() throws Exception {
+        // Arrange - create LOCAL command (not proxied)
+        byte[] payload = "local".getBytes(StandardCharsets.UTF_8);
+        byte[] proxyHeader = new AwsProxyEncoderHelper()
+                .command(ProxyHeader.Command.LOCAL)
                 .build();
 
         byte[] packet = new byte[proxyHeader.length + payload.length];
@@ -204,20 +196,25 @@ void receive_withUntrustedProxy_skipsProcessing() throws Exception {
         DatagramPacket receivePacket = new DatagramPacket(receiveBuf, receiveBuf.length);
         socket.receive(receivePacket);
 
-        // Assert - packet should be delivered unchanged, no parsing
-        verify(mockMetrics, never()).onHeaderParsed(any());
+        // Assert - cache should NOT be populated for LOCAL commands
         verify(mockCache, never()).put(any(), any());
 
-        // Packet length should include proxy header (not stripped)
-        assertEquals(packet.length, receivePacket.getLength());
+        // But metrics should still be called
+        verify(mockMetrics).onHeaderParsed(any());
+
+        // Payload should be stripped of header
+        assertEquals(payload.length, receivePacket.getLength());
     }
 
     @Test
-    void receive_withLocalCommand_doesNotPopulateCache() throws Exception {
-        // Arrange - create LOCAL command (not proxied)
-        byte[] payload = "local".getBytes(StandardCharsets.UTF_8);
+    void receive_withTcpProtocol_doesNotPopulateCache() throws Exception {
+        // Arrange - create header with TCP (not DGRAM) protocol
+        byte[] payload = "tcp".getBytes(StandardCharsets.UTF_8);
         byte[] proxyHeader = new ProxyProtocolV2Encoder()
-                .command(ProxyHeader.Command.LOCAL)
+                .family(ProxyHeader.AddressFamily.INET4)
+                .socket(ProxyHeader.TransportProtocol.STREAM) // TCP, not UDP
+                .source(new InetSocketAddress("10.1.2.3", 12345))
+                .destination(new InetSocketAddress("127.0.0.1", localPort))
                 .build();
 
         byte[] packet = new byte[proxyHeader.length + payload.length];
@@ -234,31 +231,32 @@ void receive_withLocalCommand_doesNotPopulateCache() throws Exception {
         DatagramPacket receivePacket = new DatagramPacket(receiveBuf, receiveBuf.length);
         socket.receive(receivePacket);
 
-        // Assert - cache should NOT be populated for LOCAL commands
+        // Assert - cache should NOT be populated for non-DGRAM protocols
         verify(mockCache, never()).put(any(), any());
 
-        // But metrics should still be called
+        // Metrics should still be called
         verify(mockMetrics).onHeaderParsed(any());
-
-        // Payload should be stripped of header
-        assertEquals(payload.length, receivePacket.getLength());
     }
 
+
     @Test
-    void receive_withTcpProtocol_doesNotPopulateCache() throws Exception {
-        // Arrange - create header with TCP (not DGRAM) protocol
-        byte[] payload = "tcp".getBytes(StandardCharsets.UTF_8);
-        byte[] proxyHeader = new ProxyProtocolV2Encoder()
+    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.INET4)
-                .socket(ProxyHeader.TransportProtocol.STREAM) // TCP, not UDP
-                .source(new InetSocketAddress("10.1.2.3", 12345))
+                .socket(ProxyHeader.TransportProtocol.DGRAM)
+                .source(realClient)
                 .destination(new InetSocketAddress("127.0.0.1", localPort))
                 .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);
 
+        // 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)));
@@ -269,11 +267,99 @@ void receive_withTcpProtocol_doesNotPopulateCache() throws Exception {
         DatagramPacket receivePacket = new DatagramPacket(receiveBuf, receiveBuf.length);
         socket.receive(receivePacket);
 
-        // Assert - cache should NOT be populated for non-DGRAM protocols
-        verify(mockCache, never()).put(any(), any());
+        // Assert - onHeaderParsed should be called
+        ArgumentCaptor<ProxyHeader> headerCaptor = ArgumentCaptor.forClass(ProxyHeader.class);
+        verify(mockMetrics).onHeaderParsed(headerCaptor.capture());
 
-        // Metrics should still be called
-        verify(mockMetrics).onHeaderParsed(any());
+        ProxyHeader capturedHeader = headerCaptor.getValue();
+        assertNotNull(capturedHeader);
+        assertEquals(ProxyHeader.TransportProtocol.DGRAM, capturedHeader.getProtocol());
+        assertEquals(realClient, capturedHeader.getSourceAddress());
+    }
+
+    @Test
+    void receive_withInvalidData_callsMetricsOnParseError() throws Exception {
+        // Arrange - send garbage data
+        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)));
+        }
+
+        // Act
+        byte[] receiveBuf = new byte[2048];
+        DatagramPacket receivePacket = new DatagramPacket(receiveBuf, receiveBuf.length);
+        socket.receive(receivePacket);
+
+        // Assert - onParseError should be called
+        verify(mockMetrics).onParseError(any(Exception.class));
+
+        // Original packet should be delivered unchanged
+        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();
+        }
     }
 }