A simple Java SOCKS5 server

Socks5Server.java

import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;

public class Socks5Server {
   private static final int BUFFER_SIZE = 8192;
   private Selector selector;
   private final Dialer dialer;

   private Map<SocketChannel, SocketChannel> associatedChannels = new HashMap<>();

   public Socks5Server(Dialer dialer) {
	  this.dialer = dialer;
   }

   public void start(int port) throws IOException {
	  selector = Selector.open();
	  ServerSocketChannel serverSocket = ServerSocketChannel.open();
	  serverSocket.bind(new InetSocketAddress(port));
	  serverSocket.configureBlocking(false);
	  serverSocket.register(selector, SelectionKey.OP_ACCEPT);

	  while (true) {
		 selector.select();
		 for (SelectionKey key : selector.selectedKeys()) {
			if (key.isAcceptable()) {
			   handleAccept(key);
			} else if (key.isReadable()) {
			   handleRead(key);
			}
		 }
		 selector.selectedKeys().clear();
	  }
   }

   private PacketHandler getPacketHandler(Dialer dialer, SocketChannel clientChannel, ByteBuffer buffer) throws IOException {
	  // Check SOCKS5 version
	  byte version = buffer.get();
	  if (version != 5) {
		 throw new IOException("Unsupported SOCKS version");
	  }


	  // Check the command
	  byte command = buffer.get();

	  byte reserved = buffer.get();
	  if (reserved != 0) {
		 throw new IOException("Invalid SOCKS5 reserved byte");
	  }

	  switch (command) {
		 case 1:
			return new ConnectCommandHandler(dialer, clientChannel, buffer);
		 case 2:
			return new BindCommandHandler(dialer, clientChannel, buffer);
		 default:
			throw new IOException("Unsupported SOCKS command");
	  }
   }

   private void handleAccept(SelectionKey key) throws IOException {
	  ServerSocketChannel serverSocket = (ServerSocketChannel) key.channel();
	  SocketChannel clientChannel = serverSocket.accept();
	  clientChannel.configureBlocking(false);
	  clientChannel.register(selector, SelectionKey.OP_READ, new byte[BUFFER_SIZE]);
   }

   private void handleGreeting(SocketChannel channel, ByteBuffer buffer) throws IOException {
	  if (buffer.get() != 0x05) {
		 channel.close();
		 return;
	  }

	  int numMethods = Byte.toUnsignedInt(buffer.get());
	  boolean noAuthRequired = false;
	  for (int i = 0; i < numMethods; i++) {
		 byte method = buffer.get();
		 if (method == 0x00) {
			noAuthRequired = true;
		 }
	  }

	  ByteBuffer response = ByteBuffer.allocate(2);
	  response.put((byte) 0x05); // SOCKS version
	  if (noAuthRequired) {
		 response.put((byte) 0x00); // No authentication required
	  } else {
		 response.put((byte) 0xFF); // No acceptable methods
	  }
	  response.flip();
	  channel.write(response);

	  if (!noAuthRequired) {
		 channel.close();
	  }
   }

   private void handleRead(SelectionKey key) throws IOException {
	  SocketChannel channel = (SocketChannel) key.channel();
	  ByteBuffer buffer = ByteBuffer.allocate(8192);
	  int bytesRead = channel.read(buffer);

	  if (bytesRead > 0) {
		 SocketChannel otherChannel = associatedChannels.get(channel);
		 if (otherChannel != null) {
			buffer.flip();
			while (buffer.hasRemaining()) {
			   otherChannel.write(buffer);
			}
		 } else {
			buffer.flip();

			if (key.attachment() != null) {
			   handleGreeting(channel, buffer);
			   key.attach(null);

			   if (!buffer.hasRemaining())
				  return;
			}

			PacketHandler handler = getPacketHandler(dialer, channel, buffer);
			handler.handle();
		 }
	  } else if (bytesRead < 0) {
		 SocketChannel otherChannel = associatedChannels.remove(channel);
		 if (otherChannel != null) {
			otherChannel.close();
			associatedChannels.remove(otherChannel);
		 }
		 channel.close();
	  }
   }


   abstract static class PacketHandler {
	  protected Dialer dialer;
	  protected SocketChannel clientChannel;
	  protected ByteBuffer buffer;

	  public PacketHandler(Dialer dialer, SocketChannel clientChannel, ByteBuffer buffer) {
		 this.dialer = dialer;
		 this.clientChannel = clientChannel;
		 this.buffer = buffer;
	  }

	  protected String readAddress() throws IOException {
		 // read address type
		 byte type = buffer.get();
		 byte[] bytes;
		 switch (type) {
			case 1: // IPv4
			   bytes = new byte[4];
			   buffer.get(bytes);
			   return InetAddress.getByAddress(bytes).getHostAddress();
			case 3: // Domain name
			   int length = buffer.get() & 0xff;
			   bytes = new byte[length];
			   buffer.get(bytes);
			   return new String(bytes, StandardCharsets.UTF_8);
			case 4: // IPv6
			   bytes = new byte[16];
			   buffer.get(bytes);
			   return InetAddress.getByAddress(bytes).getHostAddress();
			default:
			   throw new IOException("Unsupported address type: " + type);
		 }
	  }

	  protected int readPort() {
		 return buffer.getShort() & 0xffff;
	  }

	  protected void sendResponse(byte rep, SocketAddress bndAddr) throws IOException {
		 buffer.clear();
		 buffer.put((byte) 5); // version
		 buffer.put(rep); // rep
		 buffer.put((byte) 0); // rsv

		 if (bndAddr instanceof InetSocketAddress) {
			InetSocketAddress inetBndAddr = (InetSocketAddress) bndAddr;
			byte[] addrBytes = inetBndAddr.getAddress().getAddress();
			if (addrBytes.length == 4) { // IPv4
			   buffer.put((byte) 1); // atyp
			   buffer.put(addrBytes); // bnd.addr
			} else { // IPv6
			   buffer.put((byte) 4); // atyp
			   buffer.put(addrBytes); // bnd.addr
			}
			buffer.putShort((short) inetBndAddr.getPort()); // bnd.port
		 } else {
			buffer.put((byte) 1); // atyp (IPv4)
			buffer.putInt(0); // bnd.addr (0.0.0.0)
			buffer.putShort((short) 0); // bnd.port (0)
		 }

		 buffer.flip();
		 clientChannel.write(buffer);
	  }

	  public abstract void handle() throws IOException;
   }

   class ConnectCommandHandler extends PacketHandler {
	  public ConnectCommandHandler(Dialer dialer, SocketChannel clientChannel, ByteBuffer buffer) {
		 super(dialer, clientChannel, buffer);
	  }

	  @Override
	  public void handle() throws IOException {
		 String host = readAddress();
		 int port = readPort();

		 try {
			SocketChannel upstreamChannel = dialer.connect(host, port);
			if (upstreamChannel != null) {
			   upstreamChannel.configureBlocking(false);
			   upstreamChannel.register(selector, SelectionKey.OP_READ, clientChannel);
			   associatedChannels.put(upstreamChannel, clientChannel);
			   clientChannel.register(selector, SelectionKey.OP_READ, upstreamChannel);
			   associatedChannels.put(clientChannel, upstreamChannel);
			   sendResponse((byte) 0, upstreamChannel.getLocalAddress()); // succeeded
			} else {
			   sendResponse((byte) 1, null); // general SOCKS server failure
			}
		 } catch (IOException e) {
			sendResponse((byte) 5, null); // connection refused
		 }
	  }
   }

   class BindCommandHandler extends PacketHandler {
	  public BindCommandHandler(Dialer dialer, SocketChannel clientChannel, ByteBuffer buffer) {
		 super(dialer, clientChannel, buffer);
	  }

	  @Override
	  public void handle() throws IOException {
		 readAddress();
		 readPort();

		 try {
			ServerSocketChannel serverChannel = dialer.bind(null, 0);
			if (serverChannel != null) {
			   serverChannel.configureBlocking(false);
			   sendResponse((byte) 0, serverChannel.getLocalAddress()); // succeeded
			   SocketChannel upstreamChannel = serverChannel.accept();
			   if (upstreamChannel != null) {
				  upstreamChannel.configureBlocking(false);
				  upstreamChannel.register(selector, SelectionKey.OP_READ, clientChannel);
				  associatedChannels.put(upstreamChannel, clientChannel);
				  clientChannel.register(selector, SelectionKey.OP_READ, upstreamChannel);
				  associatedChannels.put(clientChannel, upstreamChannel);
				  sendResponse((byte) 0, upstreamChannel.getLocalAddress()); // succeeded
			   } else {
				  sendResponse((byte) 1, null); // general SOCKS server failure
			   }
			} else {
			   sendResponse((byte) 1, null); // general SOCKS server failure
			}
		 } catch (IOException e) {
			sendResponse((byte) 5, null); // connection refused
		 }
	  }
   }
}