Hooking into and providing custom behavior for the host discovery process

src/com/esotericsoftware/kryonet/Client.java

@@ -51,6 +51,7 @@ public class Client extends Connection implements EndPoint {
 	private int connectTcpPort;
 	private int connectUdpPort;
 	private boolean isClosed;
+	private ClientDiscoveryHandler discoveryHandler;
 
 	/** Creates a Client with a write buffer size of 8192 and an object buffer size of 2048. */
 	public Client () {
@@ -80,6 +81,8 @@ public Client (int writeBufferSize, int objectBufferSize, Serialization serializ
 		endPoint = this;
 
 		this.serialization = serialization;
+
+		this.discoveryHandler = ClientDiscoveryHandler.DEFAULT;
 
 		initialize(serialization, writeBufferSize, objectBufferSize);
 
@@ -89,6 +92,10 @@ public Client (int writeBufferSize, int objectBufferSize, Serialization serializ
 			throw new RuntimeException("Error opening selector.", ex);
 		}
 	}
+
+	public void setDiscoveryHandler(ClientDiscoveryHandler newDiscoveryHandler) {
+		discoveryHandler = newDiscoveryHandler;
+	}
 
 	public Serialization getSerialization () {
 		return serialization;
@@ -450,20 +457,22 @@ public InetAddress discoverHost (int udpPort, int timeoutMillis) {
 			socket = new DatagramSocket();
 			broadcast(udpPort, socket);
 			socket.setSoTimeout(timeoutMillis);
-			DatagramPacket packet = new DatagramPacket(new byte[0], 0);
+			DatagramPacket packet = discoveryHandler.onRequestNewDatagramPacket();
 			try {
 				socket.receive(packet);
 			} catch (SocketTimeoutException ex) {
 				if (INFO) info("kryonet", "Host discovery timed out.");
 				return null;
 			}
 			if (INFO) info("kryonet", "Discovered server: " + packet.getAddress());
+			discoveryHandler.onDiscoveredHost(packet, getKryo());
 			return packet.getAddress();
 		} catch (IOException ex) {
 			if (ERROR) error("kryonet", "Host discovery failed.", ex);
 			return null;
 		} finally {
 			if (socket != null) socket.close();
+			discoveryHandler.onFinally();
 		}
 	}
 
@@ -478,21 +487,23 @@ public List<InetAddress> discoverHosts (int udpPort, int timeoutMillis) {
 			broadcast(udpPort, socket);
 			socket.setSoTimeout(timeoutMillis);
 			while (true) {
-				DatagramPacket packet = new DatagramPacket(new byte[0], 0);
+				DatagramPacket packet = discoveryHandler.onRequestNewDatagramPacket();
 				try {
 					socket.receive(packet);
 				} catch (SocketTimeoutException ex) {
 					if (INFO) info("kryonet", "Host discovery timed out.");
 					return hosts;
 				}
 				if (INFO) info("kryonet", "Discovered server: " + packet.getAddress());
+				discoveryHandler.onDiscoveredHost(packet, getKryo());
 				hosts.add(packet.getAddress());
 			}
 		} catch (IOException ex) {
 			if (ERROR) error("kryonet", "Host discovery failed.", ex);
 			return hosts;
 		} finally {
 			if (socket != null) socket.close();
+			discoveryHandler.onFinally();
 		}
 	}
 }

src/com/esotericsoftware/kryonet/ClientDiscoveryHandler.java

@@ -0,0 +1,61 @@
+package com.esotericsoftware.kryonet;
+
+import java.net.DatagramPacket;
+
+import com.esotericsoftware.kryo.Kryo;
+import com.esotericsoftware.kryo.io.Input;
+
+public interface ClientDiscoveryHandler {
+
+	/**
+	 * This implementation of the {@link ClientDiscoveryHandler} is responsible
+	 * for providing the {@link Client} with it's default behavior.
+	 */
+	public static final ClientDiscoveryHandler DEFAULT = new ClientDiscoveryHandler() {
+
+		@Override
+		public DatagramPacket onRequestNewDatagramPacket() {
+			return new DatagramPacket(new byte[0], 0);
+		}
+
+		@Override
+		public void onDiscoveredHost(DatagramPacket datagramPacket, Kryo kryo) {
+			//
+		}
+
+		@Override
+		public void onFinally() {
+			//
+		}
+
+	};
+
+	/**
+	 * Implementations of this method should return a new {@link DatagramPacket}
+	 * that the {@link Client} will use to fill with the incoming packet data
+	 * sent by the {@link ServerDiscoveryHandler}.
+	 * 
+	 * @return a new {@link DatagramPacket}
+	 */
+	public DatagramPacket onRequestNewDatagramPacket();
+
+	/**
+	 * Called when the {@link Client} discovers a host.
+	 * 
+	 * @param datagramPacket
+	 *            the same {@link DatagramPacket} from
+	 *            {@link #onRequestNewDatagramPacket()}, after being filled with
+	 *            the incoming packet data.
+	 * @param kryo
+	 *            the {@link Kryo} instance
+	 */
+	public void onDiscoveredHost(DatagramPacket datagramPacket, Kryo kryo);
+
+	/**
+	 * Called right before the {@link Client#discoverHost(int, int)} or
+	 * {@link Client#discoverHosts(int, int)} method exits. This allows the
+	 * implementation to clean up any resources used, i.e. an {@link Input}.
+	 */
+	public void onFinally();
+
+}

src/com/esotericsoftware/kryonet/Server.java

@@ -39,7 +39,7 @@ public class Server implements EndPoint {
 	private volatile boolean shutdown;
 	private Object updateLock = new Object();
 	private Thread updateThread;
-	private ByteBuffer emptyBuffer = ByteBuffer.allocate(0);
+	private ServerDiscoveryHandler discoveryHandler;
 
 	private Listener dispatchListener = new Listener() {
 		public void connected (Connection connection) {
@@ -96,13 +96,19 @@ public Server (int writeBufferSize, int objectBufferSize, Serialization serializ
 		this.objectBufferSize = objectBufferSize;
 
 		this.serialization = serialization;
+
+		this.discoveryHandler = ServerDiscoveryHandler.DEFAULT;
 
 		try {
 			selector = Selector.open();
 		} catch (IOException ex) {
 			throw new RuntimeException("Error opening selector.", ex);
 		}
 	}
+
+	public void setDiscoveryHandler(ServerDiscoveryHandler newDiscoveryHandler) {
+		discoveryHandler = newDiscoveryHandler;
+	}
 
 	public Serialization getSerialization () {
 		return serialization;
@@ -305,8 +311,8 @@ public void update (int timeout) throws IOException {
 							}
 							if (object instanceof DiscoverHost) {
 								try {
-									udp.datagramChannel.send(emptyBuffer, fromAddress);
-									if (DEBUG) debug("kryonet", "Responded to host discovery from: " + fromAddress);
+									boolean responseSent = discoveryHandler.onDiscoverHost(udp, fromAddress, serialization);
+									if (DEBUG && responseSent) debug("kryonet", "Responded to host discovery from: " + fromAddress);
 								} catch (IOException ex) {
 									if (WARN) warn("kryonet", "Error replying to host discovery from: " + fromAddress, ex);
 								}

src/com/esotericsoftware/kryonet/ServerDiscoveryHandler.java

@@ -0,0 +1,47 @@
+package com.esotericsoftware.kryonet;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.nio.ByteBuffer;
+import java.nio.channels.DatagramChannel;
+
+import com.esotericsoftware.kryonet.FrameworkMessage.DiscoverHost;
+
+public interface ServerDiscoveryHandler {
+
+	/**
+	 * This implementation of {@link ServerDiscoveryHandler} is responsible for
+	 * providing the {@link Server} with it's default behavior.
+	 */
+	public static final ServerDiscoveryHandler DEFAULT = new ServerDiscoveryHandler() {
+		private ByteBuffer emptyBuffer = ByteBuffer.allocate(0);
+
+		@Override
+		public boolean onDiscoverHost(UdpConnection udp,
+				InetSocketAddress fromAddress, Serialization serialization)
+				throws IOException {
+			udp.datagramChannel.send(emptyBuffer, fromAddress);
+			return true;
+		}
+	};
+
+	/**
+	 * Called when the {@link Server} receives a {@link DiscoverHost} packet.
+	 * 
+	 * @param udp
+	 *            the {@link UdpConnection}
+	 * @param fromAddress
+	 *            {@link InetSocketAddress} the {@link DiscoverHost} came from
+	 * @param serialization
+	 *            the {@link Server}'s {@link Serialization} instance
+	 * @return true if a response was sent to {@code fromAddress}, false
+	 *         otherwise
+	 * @throws IOException
+	 *             from the use of
+	 *             {@link DatagramChannel#send(ByteBuffer, java.net.SocketAddress)}
+	 */
+	public boolean onDiscoverHost(UdpConnection udp,
+			InetSocketAddress fromAddress, Serialization serialization)
+			throws IOException;
+
+}

test/com/esotericsoftware/kryonet/DiscoverHostTest.java

@@ -1,22 +1,130 @@
-
 package com.esotericsoftware.kryonet;
 
 import java.io.IOException;
+import java.net.DatagramPacket;
 import java.net.InetAddress;
+import java.net.InetSocketAddress;
+import java.nio.ByteBuffer;
+
+import com.esotericsoftware.kryo.Kryo;
+import com.esotericsoftware.kryo.io.Input;
+import com.sun.xml.internal.ws.api.message.Packet;
+
+import static com.esotericsoftware.minlog.Log.*;
 
 public class DiscoverHostTest extends KryoNetTestCase {
-	public void testBroadcast () throws IOException {
+
+	public void testBroadcast() throws IOException {
+		// This server exists solely to reply to Client#discoverHost.
+		// It wouldn't be needed if the real server was using UDP.
+		final Server broadcastServer = new Server();
+		startEndPoint(broadcastServer);
+		broadcastServer.bind(0, udpPort);
+
+		final Server server = new Server();
+		startEndPoint(server);
+		server.bind(54555);
+		server.addListener(new Listener() {
+			public void disconnected(Connection connection) {
+				broadcastServer.stop();
+				server.stop();
+			}
+		});
+
+		// ----
+
+		Client client = new Client();
+		InetAddress host = client.discoverHost(udpPort, 2000);
+		if (host == null) {
+			stopEndPoints();
+			fail("No servers found.");
+			return;
+		}
+
+		startEndPoint(client);
+		client.connect(2000, host, tcpPort);
+		client.stop();
+
+		waitForThreads();
+	}
+
+	public void testCustomBroadcast() throws IOException {
+
+		ServerDiscoveryHandler serverDiscoveryHandler = new ServerDiscoveryHandler() {
+
+			@Override
+			public boolean onDiscoverHost(UdpConnection udp,
+					InetSocketAddress fromAddress, Serialization serialization)
+					throws IOException {
+
+				DiscoveryResponsePacket packet = new DiscoveryResponsePacket();
+				packet.id = 42;
+				packet.gameName = "gameName";
+				packet.playerName = "playerName";
+
+				ByteBuffer buffer = ByteBuffer.allocate(256);
+				serialization.write(null, buffer, packet);
+				buffer.flip();
+
+				udp.datagramChannel.send(buffer, fromAddress);
+
+				return true;
+			}
+
+		};
+
+		ClientDiscoveryHandler clientDiscoveryHandler = new ClientDiscoveryHandler() {
+
+			private Input input = null;
+
+			@Override
+			public DatagramPacket onRequestNewDatagramPacket() {
+				byte[] buffer = new byte[1024];
+				input = new Input(buffer);
+				return new DatagramPacket(buffer, buffer.length);
+			}
+
+			@Override
+			public void onDiscoveredHost(DatagramPacket datagramPacket, Kryo kryo) {
+				if (input != null) {
+					DiscoveryResponsePacket packet;
+					packet = (DiscoveryResponsePacket) kryo.readClassAndObject(input);
+					info("test", "packet.id = " + packet.id);
+					info("test", "packet.gameName = " + packet.gameName);
+					info("test", "packet.playerName = " + packet.playerName);
+					info("test", "datagramPacket.getAddress() = " + datagramPacket.getAddress());
+					info("test", "datagramPacket.getPort() = " + datagramPacket.getPort());
+					assertEquals(42, packet.id);
+					assertEquals("gameName", packet.gameName);
+					assertEquals("playerName", packet.playerName);
+					assertEquals(udpPort, datagramPacket.getPort());
+				}
+			}
+
+			@Override
+			public void onFinally() {
+				if (input != null) {
+					input.close();
+				}
+			}
+
+		};
+
 		// This server exists solely to reply to Client#discoverHost.
 		// It wouldn't be needed if the real server was using UDP.
 		final Server broadcastServer = new Server();
+
+		broadcastServer.getKryo().register(DiscoveryResponsePacket.class);
+		broadcastServer.setDiscoveryHandler(serverDiscoveryHandler);
+
 		startEndPoint(broadcastServer);
 		broadcastServer.bind(0, udpPort);
 
 		final Server server = new Server();
 		startEndPoint(server);
 		server.bind(54555);
 		server.addListener(new Listener() {
-			public void disconnected (Connection connection) {
+			public void disconnected(Connection connection) {
 				broadcastServer.stop();
 				server.stop();
 			}
@@ -25,6 +133,10 @@ public void disconnected (Connection connection) {
 		// ----
 
 		Client client = new Client();
+
+		client.getKryo().register(DiscoveryResponsePacket.class);
+		client.setDiscoveryHandler(clientDiscoveryHandler);
+
 		InetAddress host = client.discoverHost(udpPort, 2000);
 		if (host == null) {
 			stopEndPoints();
@@ -38,4 +150,16 @@ public void disconnected (Connection connection) {
 
 		waitForThreads();
 	}
+
+	public static class DiscoveryResponsePacket {
+
+		public DiscoveryResponsePacket() {
+			//
+		}
+
+		public int id;
+		public String gameName;
+		public String playerName;
+	}
+
 }