Java: Socks 4 Proxy work with non-blocking sockets

Starting with version 1.4, j2se introduced package java.nio, which allows you to work with sockets in non-blocking mode, which often improves performance, simplifies the code and provides additional features and functionality. And starting with j2se 1.6 on servers under OS Linux management (kernel 2.6), the implementation of the Selector class is performed using epoll, which ensures the highest possible performance.

In the example described below, I will try to demonstrate the basic principle of working with non-blocking sockets, using the example of a very real problem - the implementation of Socks 4 proxy server.

During life, anything can happen to a non-blocking socket, namely

ServerSocketChannel

• OP_ACCEPT - incoming connection

SocketChannel

• OP_READ - on the nipple data or disconnect
• OP_WRITE - the nipple is ready for recording or disconnected
• OP_CONNECT - connected or established or not

')

Sockets are selected on which something happened using one of the methods

• select () - blocking method, wakes up on an event or on wakeUp ()
• select (long) - same theme only with timeout
• selectNow () - well, non-blocking option

In our case, the proxy is passive, so the basic block select () is more suitable for us.
After that, you need to ask the selector for the keys that were active in the last sample and using the methods isAcceptable () , isReadable () , isWriteable () , isConnectable () to find out what happened to them.

The basic algorithm of our proxy server is:

1. Accept connection
2. Parsim header (to simplify this step, we assume that the header size is always smaller than the buffer size)
3. We establish connection with the purpose
4. We respond to the client that everything is OK
5. Proxy
6. Close connections

To avoid problems with full socket buffers, we will proxy as follows:
Let us have two ends A and B while A.in = B.out and vice versa, therefore A.interestOps () | OP_READ! = B.interestOps () | OP_WRITE (so that one buffer is not simultaneously used by two channels).
After one of the parties closes the connection, it is necessary to add data from the buffer to the second side and close the connection.

Well, actually the code itself, the functions tried to arrange in the order of actions to simplify the understanding of the algorithm, the comments are attached.

package ru.habrahabr ;

import java.io.IOException ;
import java.net.InetAddress ;
import java.net.InetSocketAddress ;
import java.net.UnknownHostException ;
import java.nio.ByteBuffer ;
import java.nio.channels.ClosedChannelException ;
import java.nio.channels.SelectionKey ;
import java.nio.channels.Selector ;
import java.nio.channels.ServerSocketChannel ;
import java.nio.channels.SocketChannel ;
import java.nio.channels.spi.SelectorProvider ;
import java.util.Iterator ;

/ **
* A class that implements a simple non-blocking Socks 4 Proxy Server Implementing
* connect command only
*
* @author dgreen
* @date 09/19/2009
*
* /
public class Socks4Proxy implements Runnable {
int bufferSize = 8192 ;
/ **
* Port
* /
int port ;
/ **
* Host
* /
String host ;

/ **
* Additional information clinging to each key {@link SelectionKey}
*
* @author dgreen
* @date 09/19/2009
*
* /
static class Attachment {
/ **
* Buffer for reading, at the time of proxying becomes a buffer for
* entries for key stored in peer
*
* IMPORTANT: When parsing Socks4 header, we assume that the size
* Buffer, larger than normal header size, Mozilla browser
* Firefox, header size is 12 bytes 1 version + 1 command + 2 port +
* 4 ip + 3 id (MOZ) + 1 \ 0
* /

ByteBuffer in ;
/ **
* Buffer for writing, at the time of proxying, is equal to read buffer for
* key stored in peer
* /
ByteBuffer out ;
/ **
* Where are we proxying
* /
SelectionKey peer ;

}

/ **
* the answer is OK or Service is provided
* /
static final byte [ ] OK = new byte [ ] { 0x00, 0x5a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 } ;

/ **
* The heart of a non-blocking server, practically does not change from application to
* application, except when using a non-blocking server in
* multithreaded application, and working with keys from other threads, it is necessary
* will add some KeyChangeRequest, but we are in this application without
* needs
* /
@ Override
public void run ( ) {
try {
// Create Selector
Selector selector = SelectorProvider . provider ( ) . openSelector ( ) ;
// Open the server channel
ServerSocketChannel serverChannel = ServerSocketChannel . open ( ) ;
// Remove the lock
serverChannel. configureBlocking ( false ) ;
// Hang on the port
serverChannel. socket ( ) . bind ( new InetSocketAddress ( host, port ) ) ;
// Register in the selector
serverChannel. register ( selector, serverChannel. validOps ( ) ) ;
// The main cycle of the non-blocking server
// This cycle will be the same for almost any non-blocking
// server
while ( selector. select ( ) > - 1 ) {
// Get the keys on which the events occurred at the moment
// last sample
Iterator < SelectionKey > iterator = selector. selectedKeys ( ) . iterator ( ) ;
while ( iterator. hasNext ( ) ) {
SelectionKey key = iterator. next ( ) ;
iterator. remove ( ) ;
if ( key. isValid ( ) ) {
// Handle all possible key events
try {
if ( key. isAcceptable ( ) ) {
// accept connection
accept ( key ) ;
} else if ( key. isConnectable ( ) ) {
// Establish connection
connect ( key ) ;
} else if ( key. isReadable ( ) ) {
// Read the data
read ( key ) ;
} else if ( key. isWritable ( ) ) {
// Write data
write ( key ) ;
}
} catch ( Exception e ) {
e. printStackTrace ( ) ;
close ( key ) ;
}
}
}
}

} catch ( Exception e ) {
e. printStackTrace ( ) ;
throw new IllegalStateException ( e ) ;
}
}

/ **
* The function accepts the connection, registers the key with the action of interest.
* read data (OP_READ)
*
* @param key
* key on which the event occurred
* @throws IOException
* @throws ClosedChannelException
* /
private void accept ( Selection key Key ) throws IOException , ClosedChannelException {
// Accepted
SocketChannel newChannel = ( ( ServerSocketChannel ) key. Channel ( ) ) . accept ( ) ;
// Non-blocking
newChannel. configureBlocking ( false ) ;
// Register in the selector
newChannel. register ( key. selector ( ) , SelectionKey . OP_READ ) ;
}

/ **
* We read data available at the moment. The function is in two states -
* read request header and direct proxying
*
* @param key
* key on which the event occurred
* @throws IOException
* @throws UnknownHostException
* @throws ClosedChannelException
* /
private void read ( SelectionKey key ) throws IOException , UnknownHostException , ClosedChannelException {
SocketChannel channel = ( ( SocketChannel ) key. Channel ( ) ) ;
Attachment attachment = ( ( Attachment ) key. Attachment ( ) ) ;
if ( attachment == null ) {
// Lazily initialize the buffers
key. attach ( attachment = new Attachment ( ) ) ;
attachment. in = ByteBuffer . allocate ( bufferSize ) ;
}
if ( channel. read ( attachment. in ) < 1 ) {
// -1 - break 0 - there is no space in the buffer, this can only be if
// header exceeded buffer size
close ( key ) ;
} else if ( attachment. peer == null ) {
// if there is no second end :) therefore we read the title
readHeader ( key, attachment ) ;
} else {
// well, if we proxify, then we add interest to the second end
// write
attachment. peer . interestOps ( attachment. peer . interestOps ( ) | SelectionKey . OP_WRITE ) ;
// and remove the interest of the first to read, because it has not yet been recorded
// current data, we will not read anything
key. interestOps ( key. interestOps ( ) ^ SelectionKey . OP_READ ) ;
// prepare buffer for writing
attachment. in . flip ( ) ;
}
}

private void readHeader ( SelectionKey key, Attachment attachment ) throws IllegalStateException , IOException ,
UnknownHostException , ClosedChannelException {
byte [ ] ar = attachment. in . array ( ) ;
if ( ar [ attachment. in . position ( ) - 1 ] == 0 ) {
// If the last byte \ 0 is the end of the user ID.
if ( ar [ 0 ] ! = 4 && ar [ 1 ] ! = 1 || attachment. in . position ( ) < 8 ) {
// A simple check on the protocol version and on the validity
// commands
// We support only conect
throw new IllegalStateException ( "Bad Request" ) ;
} else {
// Create a connection
SocketChannel peer = SocketChannel . open ( ) ;
peer. configureBlocking ( false ) ;
// Get the address and port from the packet
byte [ ] addr = new byte [ ] { ar [ 4 ] , ar [ 5 ] , ar [ 6 ] , ar [ 7 ] } ;
int p = ( ( ( 0xFF & ar [ 2 ] ) << 8 ) + ( 0xFF & ar [ 3 ] ) ) ;
// Start to connect
peer. connect ( new InetSocketAddress ( InetAddress . getByAddress ( addr ) , p ) ) ;
// Register in the selector
SelectionKey peerKey = peer. register ( key. selector ( ) , SelectionKey . OP_CONNECT ) ;
// Hear the requesting connection
key. interestOps ( 0 ) ;
// Key exchange :)
attachment. peer = peerKey ;
Attachment peerAttachemtn = new Attachment ( ) ;
peerAttachemtn. peer = key ;
peerKey attach ( peerAttachemtn ) ;
// Clear the buffer with headers
attachment. in . clear ( ) ;
}
}
}

/ **
* Write data from the buffer
*
* @param key
* @throws IOException
* /
private void write ( Selection key key ) throws IOException {
// Close the socket only by writing all the data
SocketChannel channel = ( ( SocketChannel ) key. Channel ( ) ) ;
Attachment attachment = ( ( Attachment ) key. Attachment ( ) ) ;
if ( channel. write ( attachment. out ) == - 1 ) {
close ( key ) ;
} else if ( attachment. out . remaining ( ) == 0 ) {
if ( attachment. peer == null ) {
// Write what was in the buffer and close
close ( key ) ;
} else {
// if everything is written, clear the buffer
attachment. out . clear ( ) ;
// Add a read interest to the second end
attachment. peer . interestOps ( attachment. peer . interestOps ( ) | SelectionKey . OP_READ ) ;
// And we remove the interest on the record
key. interestOps ( key. interestOps ( ) ^ SelectionKey . OP_WRITE ) ;
}
}
}

/ **
* Complete the connection
*
* @param key
* key on which the event occurred
* @throws IOException
* /
private void connect ( SelectionKey key ) throws IOException {
SocketChannel channel = ( ( SocketChannel ) key. Channel ( ) ) ;
Attachment attachment = ( ( Attachment ) key. Attachment ( ) ) ;
// Finish the connection
channel. finishConnect ( ) ;
// Create a buffer and respond OK
attachment. in = ByteBuffer . allocate ( bufferSize ) ;
attachment. in . put ( ok ) . flip ( ) ;
attachment. out = ( ( Attachment ) attachment. peer . attachment ( ) ) . in ;
( ( Attachment ) attachment. Peer . Attachment ( ) ) . out = attachment. in ;
// Put the second end of the flags on the write and read
// as soon as it writes OK, it will switch the second end to read and everything
// will be happy
attachment. peer . interestOps ( SelectionKey . OP_WRITE | SelectionKey . OP_READ ) ;
key. interestOps ( 0 ) ;
}

/ **
* No Comments
*
* @param key
* @throws IOException
* /
private void close ( SelectionKey key ) throws IOException {
key. cancel ( ) ;
key. channel ( ) . close ( ) ;
SelectionKey peerKey = ( ( Attachment ) key. Attachment ( ) ) . peer ;
if ( peerKey ! = null ) {
( ( Attachment ) peerKey. Attachment ( ) ) . peer = null ;
if ( ( peerKey. interestOps ( ) & SelectionKey. OP_WRITE ) == 0 ) {
( ( Attachment ) peerKey. Attachment ( ) ) . out . flip ( ) ;
}
peerKey interestOps ( SelectionKey . OP_WRITE ) ;
}
}

public static void main ( String [ ] args ) {
Socks4Proxy server = new Socks4Proxy ( ) ;
server. host = "127.0.0.1" ;
server. port = 1080 ;
server. run ( ) ;
}
}

Next, open your favorite browser, select socks 4 proxy, enter 127.0.0.1:1080 and check the performance.