Skip to content

Commit 5b0eef9

Browse files
RAV-2958 - Add UDP Socket support (#2)
Provide a DatagramSocket compatible sub-class that automatically process the ProxyProtocol header. Support: * An optional IP-cache data structure to dynamically replace the proxy IP with the client IP in the DatagramPacket * An optional metric interface for storing performance statistics * A predicate to conditionally enable ProxyProtocol processing only for some sources --------- Co-authored-by: cthirouin <113358856+cthirouin-swi@users.noreply.github.com>
1 parent 2a42037 commit 5b0eef9

18 files changed

Lines changed: 1085 additions & 37 deletions

File tree

.gitignore

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,3 +22,6 @@
2222
# virtual machine crash logs, see http://www.java.com/en/download/help/error_hotspot.xml
2323
hs_err_pid*
2424
replay_pid*
25+
26+
# maven build directories
27+
target/*

pom.xml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
<modules>
2424
<module>proxy-socket-core</module>
25+
<module>proxy-socket-udp</module>
2526
<module>proxy-socket-guava</module>
2627
</modules>
2728

proxy-socket-core/pom.xml

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,24 @@
3333
</dependency>
3434
</dependencies>
3535

36+
<build>
37+
<plugins>
38+
<!-- Create test-jar to share test utilities with other modules -->
39+
<plugin>
40+
<groupId>org.apache.maven.plugins</groupId>
41+
<artifactId>maven-jar-plugin</artifactId>
42+
<version>3.3.0</version>
43+
<executions>
44+
<execution>
45+
<goals>
46+
<goal>test-jar</goal>
47+
</goals>
48+
</execution>
49+
</executions>
50+
</plugin>
51+
</plugins>
52+
</build>
53+
3654
</project>
3755

3856

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
package net.airvantage.proxysocket.core;
66

77
import net.airvantage.proxysocket.core.v2.ProxyHeader;
8+
9+
import java.net.InetAddress;
810
import java.net.InetSocketAddress;
911

1012
/**
@@ -16,4 +18,7 @@ default void onHeaderParsed(ProxyHeader header) {}
1618
default void onParseError(Exception e) {}
1719
default void onCacheHit(InetSocketAddress client) {}
1820
default void onCacheMiss(InetSocketAddress client) {}
21+
default void onUntrustedProxy(InetAddress proxy) {}
22+
default void onTrustedProxy(InetAddress proxy) {}
23+
default void onLocal(InetAddress proxy) {}
1924
}

proxy-socket-core/src/main/java/net/airvantage/proxysocket/core/v2/ProxyProtocolV2Decoder.java

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -65,17 +65,11 @@ public static ProxyHeader parse(byte[] data, int offset, int length, boolean par
6565
throw new ProxyProtocolParseException("Invalid version");
6666
}
6767

68-
Command command;
69-
switch (cmd) {
70-
case 0x00:
71-
// Early return for LOCAL command
72-
return new ProxyHeader(Command.LOCAL, AddressFamily.AF_UNSPEC, TransportProtocol.UNSPEC, null, null, null, PROTOCOL_SIGNATURE_FIXED_LENGTH);
73-
case 0x01:
74-
command = Command.PROXY;
75-
break;
76-
default:
77-
throw new ProxyProtocolParseException("Invalid command");
78-
}
68+
Command command = switch (cmd) {
69+
case 0x00 -> Command.LOCAL;
70+
case 0x01 -> Command.PROXY;
71+
default -> throw new ProxyProtocolParseException("Invalid command");
72+
};
7973

8074
// Byte 14: address family and protocol
8175
int famProto = data[pos++] & 0xFF;
Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
/**
2+
* BSD-3-Clause License.
3+
* Copyright (c) 2025 Semtech
4+
*/
5+
package net.airvantage.proxysocket.tools;
6+
7+
import java.net.InetAddress;
8+
import java.net.InetSocketAddress;
9+
import java.net.UnknownHostException;
10+
import java.util.function.Predicate;
11+
12+
/**
13+
* Predicate compatible class that tests whether an InetSocketAddress belongs to a given subnet (CIDR).
14+
* Supports both IPv4 and IPv6 CIDR notation.
15+
*
16+
* <p>Example usage:
17+
* <pre>
18+
* // Single subnet
19+
* Predicate<InetSocketAddress> predicate = new SubnetPredicate("10.0.0.0/8");
20+
*
21+
* // Multiple subnets
22+
* Predicate<InetSocketAddress> predicate =
23+
* new SubnetPredicate("10.0.0.0/8")
24+
* .or(new SubnetPredicate("192.168.0.0/16"))
25+
* .or(new SubnetPredicate("2001:db8::/32"))
26+
* );
27+
* </pre>
28+
*
29+
* Note: it's possible to create a SubnetPredicate with a hostname instead of an IP address,
30+
* the address will be resolved to an IP address using InetAddress.getByName(hostname).
31+
* If the hostname is not resolvable, an IllegalArgumentException will be thrown.
32+
* But if the hostname resolves to a mix of IPv4/IPv6 addresses or multiple addresses,
33+
* the predicate will only match the first address found.
34+
*
35+
* Thread-safety: This class is immutable and thread-safe.
36+
*/
37+
public class SubnetPredicate implements Predicate<InetSocketAddress> {
38+
private final byte[] networkAddress;
39+
private final int prefixLength;
40+
private final int addressLength; // 4 for IPv4, 16 for IPv6
41+
42+
/**
43+
* Creates a predicate for the given CIDR subnet.
44+
*
45+
* @param cidr CIDR notation string (e.g., "10.0.0.0/8" or "2001:db8::/32")
46+
* @throws IllegalArgumentException if the CIDR notation is invalid
47+
*/
48+
public SubnetPredicate(String cidr) {
49+
if (cidr == null || cidr.isEmpty()) {
50+
throw new IllegalArgumentException("CIDR notation cannot be null or empty");
51+
}
52+
53+
int slashIndex = cidr.indexOf('/');
54+
if (slashIndex == -1) {
55+
throw new IllegalArgumentException("Invalid CIDR notation: missing '/' separator");
56+
}
57+
58+
String addressPart = cidr.substring(0, slashIndex);
59+
String prefixLengthPart = cidr.substring(slashIndex + 1);
60+
61+
try {
62+
this.prefixLength = Integer.parseInt(prefixLengthPart);
63+
} catch (NumberFormatException e) {
64+
throw new IllegalArgumentException("Invalid prefix length: " + prefixLengthPart, e);
65+
}
66+
67+
byte[] rawAddress;
68+
try {
69+
InetAddress addr = InetAddress.getByName(addressPart);
70+
rawAddress = addr.getAddress();
71+
} catch (UnknownHostException e) {
72+
throw new IllegalArgumentException("Invalid IP address: " + addressPart, e);
73+
}
74+
75+
this.addressLength = rawAddress.length;
76+
77+
// Validate prefix length
78+
int maxPrefixLength = addressLength * 8; // converts address length to bits
79+
if (prefixLength < 0 || prefixLength > maxPrefixLength) {
80+
throw new IllegalArgumentException(
81+
"Invalid prefix length " + prefixLength + " for address type (must be 0-" + maxPrefixLength + ")"
82+
);
83+
}
84+
85+
// Apply the mask to the network address to normalize it
86+
this.networkAddress = applyMask(rawAddress);
87+
}
88+
89+
/**
90+
* Tests whether the given socket address belongs to this subnet.
91+
*
92+
* @param socketAddress the socket address to test
93+
* @return true if the address is in this subnet, false otherwise
94+
*/
95+
@Override
96+
public boolean test(InetSocketAddress socketAddress) {
97+
if (socketAddress == null) {
98+
return false;
99+
}
100+
101+
InetAddress address = socketAddress.getAddress();
102+
if (address == null) {
103+
return false;
104+
}
105+
106+
byte[] testAddress = address.getAddress();
107+
108+
// Different address families don't match
109+
if (testAddress.length != addressLength) {
110+
return false;
111+
}
112+
113+
byte[] maskedTestAddress = applyMask(testAddress);
114+
115+
// Compare network portions
116+
for (int i = 0; i < networkAddress.length; i++) {
117+
if (networkAddress[i] != maskedTestAddress[i]) {
118+
return false;
119+
}
120+
}
121+
122+
return true;
123+
}
124+
125+
/**
126+
* Applies a subnet mask to an IP address.
127+
*
128+
* @param address the raw IP address bytes
129+
* @return the masked address bytes
130+
*/
131+
private byte[] applyMask(byte[] address) {
132+
byte[] result = new byte[address.length];
133+
134+
int fullBytes = prefixLength / 8;
135+
int remainingBits = prefixLength % 8;
136+
137+
// Copy the full bytes
138+
System.arraycopy(address, 0, result, 0, fullBytes);
139+
140+
// Apply mask to the partial byte if any
141+
if (remainingBits > 0) {
142+
int mask = 0xFF << (8 - remainingBits);
143+
result[fullBytes] = (byte) (address[fullBytes] & mask);
144+
}
145+
146+
// Remaining bytes are already 0
147+
return result;
148+
}
149+
150+
@Override
151+
public String toString() {
152+
try {
153+
InetAddress addr = InetAddress.getByAddress(networkAddress);
154+
return addr.getHostAddress() + "/" + prefixLength;
155+
} catch (UnknownHostException e) {
156+
return "SubnetPredicate[invalid]";
157+
}
158+
}
159+
}
160+

proxy-socket-core/src/main/java/net/airvantage/proxysocket/core/cache/ConcurrentMapProxyAddressCache.java renamed to proxy-socket-core/src/main/java/net/airvantage/proxysocket/tools/cache/ConcurrentMapProxyAddressCache.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
* BSD-3-Clause License.
33
* Copyright (c) 2025 Semtech
44
*/
5-
package net.airvantage.proxysocket.core.cache;
5+
package net.airvantage.proxysocket.tools.cache;
66

77
import net.airvantage.proxysocket.core.ProxyAddressCache;
88
import java.net.InetSocketAddress;

proxy-socket-core/src/test/java/net/airvantage/proxysocket/core/v2/AwsProxyEncoderHelper.java

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
/*
2-
* MIT License
1+
/**
2+
* BSD-3-Clause License.
33
* Copyright (c) 2025 Semtech
44
*
55
* Helper class to encode PROXY protocol v2 headers using AWS ProProt library.
@@ -28,9 +28,10 @@ public final class AwsProxyEncoderHelper {
2828
private final Header header = new Header();
2929

3030
public AwsProxyEncoderHelper command(ProxyHeader.Command cmd) {
31-
this.command = cmd == ProxyHeader.Command.LOCAL
32-
? ProxyProtocolSpec.Command.LOCAL
33-
: ProxyProtocolSpec.Command.PROXY;
31+
this.command = switch (cmd) {
32+
case LOCAL -> ProxyProtocolSpec.Command.LOCAL;
33+
case PROXY -> ProxyProtocolSpec.Command.PROXY;
34+
};
3435
return this;
3536
}
3637

@@ -73,16 +74,10 @@ public AwsProxyEncoderHelper addTlv(int type, byte[] value) {
7374

7475
public byte[] build() throws IOException {
7576
header.setCommand(command);
76-
header.setAddressFamily(family);
77-
header.setTransportProtocol(protocol);
7877

79-
// AWS ProProt validates addresses even for LOCAL command, set dummy values
80-
if (command == ProxyProtocolSpec.Command.LOCAL && source == null) {
81-
header.setSrcAddress(new byte[]{0, 0, 0, 0});
82-
header.setDstAddress(new byte[]{0, 0, 0, 0});
83-
header.setSrcPort(0);
84-
header.setDstPort(0);
85-
} else {
78+
if (command != ProxyProtocolSpec.Command.LOCAL) {
79+
header.setAddressFamily(family);
80+
header.setTransportProtocol(protocol);
8681
if (source != null) {
8782
header.setSrcAddress(source.getAddress().getAddress());
8883
header.setSrcPort(source.getPort());
@@ -92,6 +87,12 @@ public byte[] build() throws IOException {
9287
header.setDstAddress(destination.getAddress().getAddress());
9388
header.setDstPort(destination.getPort());
9489
}
90+
} else {
91+
// Spec clearly state that for LOCAL command, we
92+
// 1. must discard the protocol block including the family and
93+
// 2. \x00 is expected to be used for the protocol field.
94+
header.setAddressFamily(ProxyProtocolSpec.AddressFamily.AF_UNSPEC);
95+
header.setTransportProtocol(ProxyProtocolSpec.TransportProtocol.UNSPEC);
9596
}
9697

9798
ByteArrayOutputStream out = new ByteArrayOutputStream();

proxy-socket-core/src/test/java/net/airvantage/proxysocket/core/v2/ProxyProtocolV2DecoderTest.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
/*
2-
* MIT License
1+
/**
2+
* BSD-3-Clause License.
33
* Copyright (c) 2025 Semtech
4-
4+
*
55
* Validation of ProxyProtocolV2Decoder against hardcoded headers for known cases
66
*/
77
package net.airvantage.proxysocket.core.v2;

proxy-socket-core/src/test/java/net/airvantage/proxysocket/core/v2/ProxyProtocolV2Test.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
/*
2-
* MIT License
1+
/**
2+
* BSD-3-Clause License.
33
* Copyright (c) 2025 Semtech
4-
4+
*
55
* Validation of ProxyProtocolV2Decoder using AWS ProProt library
66
*/
77
package net.airvantage.proxysocket.core.v2;

0 commit comments

Comments
 (0)