GH-3410: Add UDP SocketCustomizer

Resolves https://github.com/spring-projects/spring-integration/issues/3410

**cherry-pick to 5.3.x**

* Doc polishing.

spring-integration-ip/src/main/java/org/springframework/integration/ip/config/IpAdapterParserUtils.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -133,6 +133,8 @@ public abstract class IpAdapterParserUtils {
 
 	public static final String CONNECT_TIMEOUT = "connect-timeout";
 
+	public static final String UDP_SOCKET_CUSTOMIZER = "socket-customizer";
+
 	private IpAdapterParserUtils() {
 	}
 

spring-integration-ip/src/main/java/org/springframework/integration/ip/config/UdpInboundChannelAdapterParser.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -51,6 +51,8 @@ public class UdpInboundChannelAdapterParser extends AbstractChannelAdapterParser
 				IpAdapterParserUtils.TASK_EXECUTOR);
 		IntegrationNamespaceUtils.setValueIfAttributeDefined(builder, element,
 				IpAdapterParserUtils.LOOKUP_HOST);
+		IntegrationNamespaceUtils.setReferenceIfAttributeDefined(builder, element,
+				IpAdapterParserUtils.UDP_SOCKET_CUSTOMIZER);
 		return builder.getBeanDefinition();
 	}
 

spring-integration-ip/src/main/java/org/springframework/integration/ip/config/UdpOutboundChannelAdapterParser.java

@@ -36,6 +36,7 @@ import org.springframework.util.StringUtils;
  */
 public class UdpOutboundChannelAdapterParser extends AbstractOutboundChannelAdapterParser {
 
+	@Override
 	protected AbstractBeanDefinition parseConsumer(Element element, ParserContext parserContext) {
 		BeanDefinitionBuilder builder = parseUdp(element, parserContext);
 		IpAdapterParserUtils.addCommonSocketOptions(builder, element);
@@ -86,6 +87,8 @@ public class UdpOutboundChannelAdapterParser extends AbstractOutboundChannelAdap
 				IpAdapterParserUtils.TASK_EXECUTOR);
 		IntegrationNamespaceUtils.setValueIfAttributeDefined(builder, element,
 				"socket-expression", "socketExpressionString");
+		IntegrationNamespaceUtils.setReferenceIfAttributeDefined(builder, element,
+				IpAdapterParserUtils.UDP_SOCKET_CUSTOMIZER);
 		return builder;
 	}
 

spring-integration-ip/src/main/java/org/springframework/integration/ip/dsl/AbstractUdpOutboundChannelAdapterSpec.java

@@ -21,6 +21,7 @@ import java.util.function.Function;
 
 import org.springframework.integration.dsl.MessageHandlerSpec;
 import org.springframework.integration.expression.FunctionExpression;
+import org.springframework.integration.ip.udp.SocketCustomizer;
 import org.springframework.integration.ip.udp.UnicastSendingMessageHandler;
 import org.springframework.messaging.Message;
 
@@ -133,4 +134,15 @@ public abstract class AbstractUdpOutboundChannelAdapterSpec<S extends AbstractUd
 		return _this();
 	}
 
+	/**
+	 * Configure the socket.
+	 * @param customizer the customizer.
+	 * @return the spec.
+	 * @since 5.3.3
+	 */
+	public S configureSocket(SocketCustomizer customizer) {
+		this.target.setSocketCustomizer(customizer);
+		return _this();
+	}
+
 }

spring-integration-ip/src/main/java/org/springframework/integration/ip/dsl/UdpInboundChannelAdapterSpec.java

@@ -21,6 +21,7 @@ import java.util.concurrent.Executor;
 
 import org.springframework.integration.dsl.MessageProducerSpec;
 import org.springframework.integration.ip.udp.MulticastReceivingChannelAdapter;
+import org.springframework.integration.ip.udp.SocketCustomizer;
 import org.springframework.integration.ip.udp.UnicastReceivingChannelAdapter;
 import org.springframework.scheduling.TaskScheduler;
 
@@ -50,7 +51,7 @@ public class UdpInboundChannelAdapterSpec
 	 */
 	public UdpInboundChannelAdapterSpec soTimeout(int soTimeout) {
 		this.target.setSoTimeout(soTimeout);
-		return _this();
+		return this;
 	}
 
 	/**
@@ -60,7 +61,7 @@ public class UdpInboundChannelAdapterSpec
 	 */
 	public UdpInboundChannelAdapterSpec taskScheduler(TaskScheduler taskScheduler) {
 		this.target.setTaskScheduler(taskScheduler);
-		return _this();
+		return this;
 	}
 
 	/**
@@ -70,7 +71,7 @@ public class UdpInboundChannelAdapterSpec
 	 */
 	public UdpInboundChannelAdapterSpec soReceiveBufferSize(int soReceiveBufferSize) {
 		this.target.setSoReceiveBufferSize(soReceiveBufferSize);
-		return _this();
+		return this;
 	}
 
 	/**
@@ -80,7 +81,7 @@ public class UdpInboundChannelAdapterSpec
 	 */
 	public UdpInboundChannelAdapterSpec receiveBufferSize(int receiveBufferSize) {
 		this.target.setReceiveBufferSize(receiveBufferSize);
-		return _this();
+		return this;
 	}
 
 	/**
@@ -90,7 +91,7 @@ public class UdpInboundChannelAdapterSpec
 	 */
 	public UdpInboundChannelAdapterSpec lengthCheck(boolean lengthCheck) {
 		this.target.setLengthCheck(lengthCheck);
-		return _this();
+		return this;
 	}
 
 	/**
@@ -100,7 +101,7 @@ public class UdpInboundChannelAdapterSpec
 	 */
 	public UdpInboundChannelAdapterSpec localAddress(String localAddress) {
 		this.target.setLocalAddress(localAddress);
-		return _this();
+		return this;
 	}
 
 	/**
@@ -110,7 +111,7 @@ public class UdpInboundChannelAdapterSpec
 	 */
 	public UdpInboundChannelAdapterSpec poolSize(int poolSize) {
 		this.target.setPoolSize(poolSize);
-		return _this();
+		return this;
 	}
 
 	/**
@@ -120,7 +121,7 @@ public class UdpInboundChannelAdapterSpec
 	 */
 	public UdpInboundChannelAdapterSpec taskExecutor(Executor taskExecutor) {
 		this.target.setTaskExecutor(taskExecutor);
-		return _this();
+		return this;
 	}
 
 	/**
@@ -130,7 +131,7 @@ public class UdpInboundChannelAdapterSpec
 	 */
 	public UdpInboundChannelAdapterSpec socket(DatagramSocket socket) {
 		this.target.setSocket(socket);
-		return _this();
+		return this;
 	}
 
 	/**
@@ -140,7 +141,7 @@ public class UdpInboundChannelAdapterSpec
 	 */
 	public UdpInboundChannelAdapterSpec soSendBufferSize(int soSendBufferSize) {
 		this.target.setSoSendBufferSize(soSendBufferSize);
-		return _this();
+		return this;
 	}
 
 	/**
@@ -150,7 +151,18 @@ public class UdpInboundChannelAdapterSpec
 	 */
 	public UdpInboundChannelAdapterSpec lookupHost(boolean lookupHost) {
 		this.target.setLookupHost(lookupHost);
-		return _this();
+		return this;
+	}
+
+	/**
+	 * Configure the socket.
+	 * @param customizer the customizer.
+	 * @return the spec.
+	 * @since 5.3.3
+	 */
+	public UdpInboundChannelAdapterSpec configureSocket(SocketCustomizer customizer) {
+		this.target.setSocketCustomizer(customizer);
+		return this;
 	}
 
 }

spring-integration-ip/src/main/java/org/springframework/integration/ip/dsl/package-info.java

@@ -1,4 +1,5 @@
 /**
  * Provides TCP/UDP Component support for the Java DSL.
  */
+@org.springframework.lang.NonNullApi
 package org.springframework.integration.ip.dsl;

spring-integration-ip/src/main/java/org/springframework/integration/ip/udp/SocketCustomizer.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2020 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.integration.ip.udp;
+
+import java.net.DatagramSocket;
+import java.net.SocketException;
+
+/**
+ * Configures a socket.
+ *
+ * @author Gary Russell
+ * @since 5.3.3
+ */
+@FunctionalInterface
+public interface SocketCustomizer {
+
+	/**
+	 * Configure the socket ({code setTrafficClass()}, etc).
+	 * @param socket the socket.
+	 * @throws SocketException a socket exception.
+	 */
+	void configure(DatagramSocket socket) throws SocketException;
+
+}

spring-integration-ip/src/main/java/org/springframework/integration/ip/udp/UnicastReceivingChannelAdapter.java

@@ -35,6 +35,7 @@ import org.springframework.lang.Nullable;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageHeaders;
 import org.springframework.messaging.MessagingException;
+import org.springframework.util.Assert;
 
 /**
  * A channel adapter to receive incoming UDP packets. Packets can optionally be preceded by a
@@ -58,6 +59,7 @@ public class UnicastReceivingChannelAdapter extends AbstractInternetProtocolRece
 
 	private int soSendBufferSize = -1;
 
+	private SocketCustomizer socketCustomizer = socket -> { };
 
 	/**
 	 * Constructs a UnicastReceivingChannelAdapter that listens on the specified port.
@@ -89,6 +91,16 @@ public class UnicastReceivingChannelAdapter extends AbstractInternetProtocolRece
 		this.mapper.setLengthCheck(lengthCheck);
 	}
 
+	/**
+	 * Set a customizer to further configure the socket after creation.
+	 * @param socketCustomizer the customizer.
+	 * @since 5.3.3
+	 */
+	public void setSocketCustomizer(SocketCustomizer socketCustomizer) {
+		Assert.notNull(socketCustomizer, "'socketCustomizer' cannot be null");
+		this.socketCustomizer = socketCustomizer;
+	}
+
 	@Override
 	public boolean isLongLived() {
 		return true;
@@ -169,6 +181,7 @@ public class UnicastReceivingChannelAdapter extends AbstractInternetProtocolRece
 			if (this.soSendBufferSize > 0) {
 				out.setSendBufferSize(this.soSendBufferSize);
 			}
+			this.socketCustomizer.configure(out);
 			out.send(ackPack);
 			out.close();
 		}
@@ -258,7 +271,7 @@ public class UnicastReceivingChannelAdapter extends AbstractInternetProtocolRece
 	}
 
 	/**
-	 * Sets timeout and receive buffer size
+	 * Sets timeout and receive buffer size; calls the socket customizer.
 	 *
 	 * @param socket The socket.
 	 * @throws SocketException Any socket exception.
@@ -269,6 +282,7 @@ public class UnicastReceivingChannelAdapter extends AbstractInternetProtocolRece
 		if (soReceiveBufferSize > 0) {
 			socket.setReceiveBufferSize(soReceiveBufferSize);
 		}
+		this.socketCustomizer.configure(socket);
 	}
 
 	@Override

spring-integration-ip/src/main/java/org/springframework/integration/ip/udp/UnicastSendingMessageHandler.java

@@ -101,6 +101,8 @@ public class UnicastSendingMessageHandler extends
 
 	private EvaluationContext evaluationContext;
 
+	private SocketCustomizer socketCustomizer = socket -> { };
+
 	private volatile CountDownLatch ackLatch;
 
 	private volatile boolean ackThreadRunning;
@@ -206,6 +208,25 @@ public class UnicastSendingMessageHandler extends
 				ackTimeout);
 	}
 
+	/**
+	 * @param lengthCheck if true, a four byte binary length header is added to the
+	 * packet, allowing the receiver to check for data truncation.
+	 * @since 5.0
+	 */
+	public void setLengthCheck(boolean lengthCheck) {
+		this.mapper.setLengthCheck(lengthCheck);
+	}
+
+	/**
+	 * Set a customizer to further configure the socket after creation.
+	 * @param socketCustomizer the customizer.
+	 * @since 5.3.3
+	 */
+	public void setSocketCustomizer(SocketCustomizer socketCustomizer) {
+		Assert.notNull(socketCustomizer, "'socketCustomizer' cannot be null");
+		this.socketCustomizer = socketCustomizer;
+	}
+
 	protected final void setReliabilityAttributes(boolean lengthCheck,
 			boolean acknowledge, String ackHost, int ackPort, int ackTimeout) {
 
@@ -224,15 +245,6 @@ public class UnicastSendingMessageHandler extends
 		}
 	}
 
-	/**
-	 * @param lengthCheck if true, a four byte binary length header is added to the
-	 * packet, allowing the receiver to check for data truncation.
-	 * @since 5.0
-	 */
-	public void setLengthCheck(boolean lengthCheck) {
-		this.mapper.setLengthCheck(lengthCheck);
-	}
-
 	@Override
 	public void doStart() {
 		if (this.acknowledge) {
@@ -495,6 +507,7 @@ public class UnicastSendingMessageHandler extends
 		if (soSendBufferSize > 0) {
 			socket.setSendBufferSize(soSendBufferSize);
 		}
+		this.socketCustomizer.configure(socket);
 	}
 
 	/**

spring-integration-ip/src/main/resources/org/springframework/integration/ip/config/spring-integration-ip.xsd

@@ -873,6 +873,18 @@
 			<xsd:extension base="common-attributes">
 				<xsd:attribute name="check-length" type="xsd:string"/>
 				<xsd:attribute name="multicast" type="xsd:string"/>
+				<xsd:attribute name="socket-customizer" type="xsd:string">
+					<xsd:annotation>
+						<xsd:documentation>
+							A SocketCustomizer bean.
+						</xsd:documentation>
+						<xsd:appinfo>
+							<tool:annotation kind="ref">
+								<tool:expected-type type="org.springframework.integration.ip.udp.SocketCustomizer"/>
+							</tool:annotation>
+						</xsd:appinfo>
+					</xsd:annotation>
+				</xsd:attribute>
 			</xsd:extension>
 		</xsd:complexContent>
 	</xsd:complexType>

spring-integration-ip/src/test/java/org/springframework/integration/ip/config/ParserTests-context.xml

@@ -19,6 +19,8 @@
 
 	<task:executor id="externalTE" pool-size="10"/>
 
+	<bean id="udpCust" class = "org.springframework.integration.ip.config.ParserUnitTests$SocketCust"/>
+
 	<ip:udp-inbound-channel-adapter id="testInUdp"
 		channel="udpChannel"
 		check-length="true"
@@ -35,6 +37,7 @@
 		lookup-host="false"
 		auto-startup="false"
 		phase="1234"
+		socket-customizer="udpCust"
 	/>
 
 	<ip:udp-inbound-channel-adapter id="testInUdpMulticast"
@@ -141,6 +144,7 @@
 		local-address="127.0.0.1"
 		task-executor="externalTE"
 		order="23"
+		socket-customizer="udpCust"
 	/>
 
 	<ip:udp-outbound-channel-adapter id="testOutUdpiMulticast"

spring-integration-ip/src/test/java/org/springframework/integration/ip/config/ParserUnitTests.java

@@ -20,6 +20,8 @@ import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.BDDMockito.given;
 import static org.mockito.Mockito.mock;
 
+import java.net.DatagramSocket;
+import java.net.SocketException;
 import java.time.Duration;
 import java.util.Iterator;
 import java.util.Set;
@@ -66,6 +68,7 @@ import org.springframework.integration.ip.tcp.connection.TcpSocketSupport;
 import org.springframework.integration.ip.udp.DatagramPacketMessageMapper;
 import org.springframework.integration.ip.udp.MulticastReceivingChannelAdapter;
 import org.springframework.integration.ip.udp.MulticastSendingMessageHandler;
+import org.springframework.integration.ip.udp.SocketCustomizer;
 import org.springframework.integration.ip.udp.UnicastReceivingChannelAdapter;
 import org.springframework.integration.ip.udp.UnicastSendingMessageHandler;
 import org.springframework.integration.test.util.TestUtils;
@@ -273,6 +276,9 @@ public class ParserUnitTests {
 	@Autowired
 	TcpMessageMapper mapper;
 
+	@Autowired
+	SocketCust socketCustomizer;
+
 	private static CountDownLatch adviceCalled = new CountDownLatch(1);
 
 	@Test
@@ -293,6 +299,7 @@ public class ParserUnitTests {
 		assertThat((Boolean) mapperAccessor.getPropertyValue("lookupHost")).isFalse();
 		assertThat(TestUtils.getPropertyValue(udpIn, "autoStartup", Boolean.class)).isFalse();
 		assertThat(dfa.getPropertyValue("phase")).isEqualTo(1234);
+		assertThat(dfa.getPropertyValue("socketCustomizer")).isSameAs(this.socketCustomizer);
 	}
 
 	@Test
@@ -365,6 +372,7 @@ public class ParserUnitTests {
 		assertThat(dfa.getPropertyValue("order")).isEqualTo(23);
 		assertThat(udpOut.getComponentName()).isEqualTo("testOutUdp");
 		assertThat(udpOut.getComponentType()).isEqualTo("ip:udp-outbound-channel-adapter");
+		assertThat(dfa.getPropertyValue("socketCustomizer")).isSameAs(this.socketCustomizer);
 	}
 
 	@Test
@@ -697,4 +705,12 @@ public class ParserUnitTests {
 
 	}
 
+	public static class SocketCust implements SocketCustomizer {
+
+		@Override
+		public void configure(DatagramSocket socket) throws SocketException {
+		}
+
+	}
+
 }

spring-integration-ip/src/test/java/org/springframework/integration/ip/dsl/IpIntegrationTests.java

@@ -18,6 +18,7 @@ package org.springframework.integration.ip.dsl;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
+import java.net.DatagramSocket;
 import java.util.Collections;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
@@ -56,6 +57,7 @@ import org.springframework.integration.ip.tcp.serializer.TcpCodecs;
 import org.springframework.integration.ip.udp.MulticastSendingMessageHandler;
 import org.springframework.integration.ip.udp.UdpServerListeningEvent;
 import org.springframework.integration.ip.udp.UnicastReceivingChannelAdapter;
+import org.springframework.integration.ip.udp.UnicastSendingMessageHandler;
 import org.springframework.integration.ip.util.TestingUtilities;
 import org.springframework.integration.support.MessageBuilder;
 import org.springframework.integration.test.util.TestUtils;
@@ -100,6 +102,9 @@ public class IpIntegrationTests {
 	@Autowired
 	private UnicastReceivingChannelAdapter udpInbound;
 
+	@Autowired
+	private UnicastSendingMessageHandler udpOutbound;
+
 	@Autowired
 	private QueueChannel udpIn;
 
@@ -176,6 +181,8 @@ public class IpIntegrationTests {
 		Message<?> received = this.udpIn.receive(10000);
 		assertThat(received).isNotNull();
 		assertThat(Transformers.objectToString().transform(received).getPayload()).isEqualTo("foo");
+		assertThat(TestUtils.getPropertyValue(this.udpOutbound, "socket", DatagramSocket.class).getTrafficClass())
+				.isEqualTo(0x10);
 	}
 
 	@Test
@@ -303,7 +310,8 @@ public class IpIntegrationTests {
 
 		@Bean
 		public IntegrationFlow outUdpAdapter() {
-			return f -> f.handle(Udp.outboundAdapter(m -> m.getHeaders().get("udp_dest")));
+			return f -> f.handle(Udp.outboundAdapter(m -> m.getHeaders().get("udp_dest"))
+					.configureSocket(socket -> socket.setTrafficClass(0x10)));
 		}
 
 		@Bean