I was able to hook into java.net.Socket and java.net.ServerSocket and spy all new instances of those classes. The complete code can be seen in the source repository. Here is an overview of the approach:
When a Socket or ServerSocket is instantiated, the first thing in its constructor is a call to setImpl() which instantiates the object which really implements the socket functionality. The default implementation is an instance of java.net.SocksSocketImpl, but it's possible to override that by setting a custom java.net.SocketImplFactory through java.net.Socket#setSocketImplFactory and java.net.ServerSocket#setSocketFactory.
This is complicated a bit by all implementations of java.net.SocketImpl being package-private, but with a little bit of reflection that's not too hard:
private static SocketImpl newSocketImpl() {
try {
Class<?> defaultSocketImpl = Class.forName("java.net.SocksSocketImpl");
Constructor<?> constructor = defaultSocketImpl.getDeclaredConstructor();
constructor.setAccessible(true);
return (SocketImpl) constructor.newInstance();
} catch (Exception e) {
throw new RuntimeException(e);
}
}
The SocketImplFactory implementation for spying on all sockets as they are created looks something like this:
final List<SocketImpl> allSockets = Collections.synchronizedList(new ArrayList<SocketImpl>());
ServerSocket.setSocketFactory(new SocketImplFactory() {
public SocketImpl createSocketImpl() {
SocketImpl socket = newSocketImpl();
allSockets.add(socket);
return socket;
}
});
Note that setSocketFactory/setSocketImplFactory can be called only once, so you either need to have only one test which does that (like I have it), or you must create a static singleton (yuck!) for holding that spy.
Then the question is that that how to find out whether the socket is closed? Both Socket and ServerSocket have a method isClosed(), but that uses a boolean internal to those classes for keeping track of whether it was closed - the SocketImpl instance does not have an easy way of checking whether it was closed. (BTW, both Socket and ServerSocket are backed by a SocketImpl - there is no "ServerSocketImpl".)
Thankfully the SocketImpl has a reference to the Socket or ServerSocket which it is backing. The aforementioned setImpl() method calls impl.setSocket(this) or impl.setServerSocket(this), and it's possible to get that reference back by calling java.net.SocketImpl#getSocket or java.net.SocketImpl#getServerSocket.
Once again those methods are package-private, so a little bit of reflection is needed:
private static Socket getSocket(SocketImpl impl) {
try {
Method getSocket = SocketImpl.class.getDeclaredMethod("getSocket");
getSocket.setAccessible(true);
return (Socket) getSocket.invoke(impl);
} catch (Exception e) {
throw new RuntimeException(e);
}
}
private static ServerSocket getServerSocket(SocketImpl impl) {
try {
Method getServerSocket = SocketImpl.class.getDeclaredMethod("getServerSocket");
getServerSocket.setAccessible(true);
return (ServerSocket) getServerSocket.invoke(impl);
} catch (Exception e) {
throw new RuntimeException(e);
}
}
Note that getSocket/getServerSocket may not be called inside the SocketImplFactory, because Socket/ServerSocket sets them only after the SocketImpl is returned from there.
Now there is all the infrastructure necessary for checking in our tests whatever we want about the Socket/ServerSocket:
for (SocketImpl impl : allSockets) {
assertIsClosed(getSocket(impl));
}
The full source code is here.
