Java: Socks 4 Proxy working with non-blocking sockets

Starting with version 1.4, package java.nio has appeared in j2se, which allows you to work with sockets in non-blocking mode, which often improves performance, simplifies code and provides additional features and functionality. And since version j2se 1.6 on servers running Linux OS (kernel 2.6), the Selector class was implemented using epoll, which provides the highest possible performance.

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

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

ServerSocketChannel

• OP_ACCEPT - incoming connection

Socketchannel

• OP_READ - data or disconnect on the nipple
• OP_WRITE - nipple ready for recording or disconnect
• OP_CONNECT - connection either established or not

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

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

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

The basic algorithm of our proxy server is as follows:

1. Accept Connection
2. Parsing the header (to simplify this step, we assume that the header size is always smaller than the buffer size)
3. We establish a connection with the goal
4. We answer the client that everything is OK
5. Proxied
6. Close connections

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

Well, actually the code itself, I tried to arrange the functions in the order of actions to simplify the understanding of the algorithm, 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;
 
/ **
 * Class implementing a simple non-blocking Socks 4 Proxy Server Implementing
 * only connect
 * 
 * @author dgreen
 * @date September 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 it becomes a buffer for
         * writing for key stored in peer
         * 
         * IMPORTANT: When parsing a Socks4 header, we assume that the
         * buffer size is larger than the normal header size for Mozilla
         * Firefox browser , the 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 it is equal to the buffer for reading for
         * the key stored in peer
         * /
        ByteBuffer out;
        / **
         * Where to proxy
         * /
        SelectionKey peer;
 
    }
 
    / **
     * the answer looks OK or the 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, unless when using a non-blocking server in a
     * multi-threaded application, and working with keys from other threads, you
     * will need to add some KeyChangeRequest, but this application doesn’t need us
     * 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));
            // Registration in the
            serverChannel.register selector (selector, serverChannel.validOps ());
            // The main loop the operation of a non-blocking server
            // This cycle will be the same for almost any non-blocking
            // server
            while (selector.select ()> -1) {
                // Get the keys on which events occurred at the moment
                // last
                iterator fetch 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 the connection
                                accept (key);
                            } else if (key.isConnectable ()) {
                                // Make a connection
                                connect (key);
                            } else if (key.isReadable ()) {
                                // Read 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
     * the key on which the event occurred
     * @throws IOException
     * @throws ClosedChannelException
     * /
    private void accept (SelectionKey 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);
    }
 
    / **
     * Read the data currently available. The function is in two states -
     * reading the request header and directly proxing
     * 
     * @param key
     * the 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 - gap 0 - there is no space in the buffer, this can only happen if the
            // header exceeds the buffer size
            close (key);
        } else if (attachment.peer == null) {
            // if there is no second end :) so we read the header
            readHeader (key, attachment);
        } else {
            // well, if we proxy, then we add interest to the second end
            // write
            attachment.peer.interestOps (attachment.peer.interestOps () | SelectionKey.OP_WRITE);
            // and we remove the interest from the first one to read, because we have not written
            // the current data, we will not read anything
            key.interestOps (key.interestOps () ^ SelectionKey.OP_READ);
            // prepare a 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 of
                // commands,
                // We only support 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]));
                // We begin to establish a connection
                peer.connect (new InetSocketAddress (InetAddress.getByAddress (addr), p));
                // Registration in the
                SelectionKey selector peerKey = peer.register (key.selector (), SelectionKey.OP_CONNECT);
                // Mute 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 (SelectionKey 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) {
                // Add what was in the buffer and close
                close (key);
            } else {
                // if everything is recorded, clear the
                attachment.out.clear () buffer ;
                // Add to the second end the interest in reading
                attachment.peer.interestOps (attachment.peer.interestOps () | SelectionKey.OP_READ);
                // And we remove our interest in writing
                key.interestOps (key.interestOps () ^ SelectionKey.OP_WRITE);
            }
        }
    }
 
    / **
     * End 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 ());
        // End 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;
        // Set the second end of the flags for writing and reading
        // as soon as she writes OK, switches the second end to reading and all
        // 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.