Socks5AcceptorHandler extends IoHandlerAdapter { public void messageReceived(final IoSession session, Object message) throws Exception { IoBuffer buffer = (IoBuffer) message; ProxyState state = getState(session); try { if (state == ProxyState.OPENED) { OpeningPacket packet = new OpeningPacket(buffer); if (packet.isAuthSupported()) { setState(session, ProxyState.INITIATED); session.write(packet.createResponse()); log.info("SOCKS5 session initiated"); } else { session.write(packet.createResponse()); session.close(false); log.info("SOCKS5 session closed"); } } else if (state == ProxyState.INITIATED) { InitiatingPacket packet = new InitiatingPacket(buffer); session.write(packet.createResponse()); setState(session, ProxyState.CONNECTED); String hash = packet.getAddress(); Socks5Pair pair = connections.register(hash, session); session.setAttribute(PAIR_KEY, pair); } else { Socks5Pair pair = getPair(session); if (pair != null && pair.isActivated()) { pair.getOther(session).write(message); } else { pair.close(); } } } catch (IllegalArgumentException e) { session.close(false); log.info("SOCKS5 session closed"); } } Socks5AcceptorHandler(Socks5ConnectionsRegistry connections); void sessionOpened(IoSession session); void sessionIdle(IoSession session, IdleStatus status); void messageReceived(final IoSession session, Object message); void sessionClosed(IoSession session); void exceptionCaught(IoSession session, Throwable cause); static final AttributeKey STATE_KEY; static final AttributeKey PAIR_KEY; }
@Test public void messageReceivedOpened() throws Exception { when(session.getAttribute(Socks5AcceptorHandler.STATE_KEY)).thenReturn(ProxyState.OPENED); IoBuffer buffer = IoBuffer.wrap(new byte[] {5, 1, 0}); handler.messageReceived(session, buffer); verify(session).setAttribute(Socks5AcceptorHandler.STATE_KEY, ProxyState.INITIATED); verify(session).write(IoBuffer.wrap(new byte[]{5, 0})); }
@Test public void messageReceivedOpenedUsernamePasswordAuth() throws Exception { when(session.getAttribute(Socks5AcceptorHandler.STATE_KEY)).thenReturn(ProxyState.OPENED); IoBuffer buffer = IoBuffer.wrap(new byte[] {5, 1, 2}); handler.messageReceived(session, buffer); verify(session).write(IoBuffer.wrap(new byte[]{5, 0x55})); verify(session).close(false); }
@Test public void messageReceivedOpenedInvalid() throws Exception { when(session.getAttribute(Socks5AcceptorHandler.STATE_KEY)).thenReturn(ProxyState.OPENED); IoBuffer buffer = IoBuffer.wrap(new byte[] {1,2,3,4}); handler.messageReceived(session, buffer); verify(session).close(false); }
@Test public void messageReceivedInitiated() throws Exception { when(session.getAttribute(Socks5AcceptorHandler.STATE_KEY)).thenReturn(ProxyState.INITIATED); IoBuffer buffer = IoBuffer.wrap(new byte[]{5, 1, 0, 3, 9, 0x6C, 0x6F, 0x63, 0x61, 0x6C, 0x68, 0x6F, 0x73, 0x74, 5, 6}); handler.messageReceived(session, buffer); verify(session).setAttribute(Socks5AcceptorHandler.STATE_KEY, ProxyState.CONNECTED); IoBuffer expected = IoBuffer.wrap(new byte[]{5, 0, 0, 3, 9, 0x6C, 0x6F, 0x63, 0x61, 0x6C, 0x68, 0x6F, 0x73, 0x74, 5, 6}); verify(session).write(expected); }
@Test public void messageReceivedInitiatedInvalid() throws Exception { when(session.getAttribute(Socks5AcceptorHandler.STATE_KEY)).thenReturn(ProxyState.INITIATED); IoBuffer buffer = IoBuffer.wrap(new byte[] {1,2,3,4}); handler.messageReceived(session, buffer); verify(session).close(false); }
@Test public void messageReceivedActivated() throws Exception { IoSession other = mock(IoSession.class); when(session.getAttribute(Socks5AcceptorHandler.STATE_KEY)).thenReturn(ProxyState.CONNECTED); Socks5Pair pair = new Socks5Pair(session, connectionsRegistry, "foo"); pair.setRequester(other); pair.activate(); when(session.getAttribute(Socks5AcceptorHandler.PAIR_KEY)).thenReturn(pair); IoBuffer buffer = IoBuffer.wrap(new byte[]{0x6C, 0x6F, 0x63}); handler.messageReceived(session, buffer); verify(other).write(buffer); }
@Test public void messageReceivedNotActivated() throws Exception { IoSession other = mock(IoSession.class); when(session.getAttribute(Socks5AcceptorHandler.STATE_KEY)).thenReturn(ProxyState.CONNECTED); Socks5Pair pair = new Socks5Pair(session, connectionsRegistry, "foo"); pair.setRequester(other); when(session.getAttribute(Socks5AcceptorHandler.PAIR_KEY)).thenReturn(pair); IoBuffer buffer = IoBuffer.wrap(new byte[]{0x6C, 0x6F, 0x63}); handler.messageReceived(session, buffer); verify(session).close(false); verify(other).close(false); }