Java NIO 实现简单代理

通过简单代理的实现，学习Java 非阻塞IO的使用，供大家参考。本文不涉及异步IO处理。

1、首先是SocketChannel的创建，创建方式与ServerSocket相似。在有新的请求来后，启动新的线程对该Socket连接进行处理。

package jim.proxy;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;

/***
 * Start up the proxy.
 * @author Jim-Zhang
 *
 */
public class ProxyBootstrap {

	/**
	 * Entry.
	 * @param args arguments
	 */
	public static void main(String[] args) {
		
		// Check arguments
		
		// If the first argument is '-help',print usage.
		if (args == null || args.length == 0){
			
			// Start the server
			startServer(Consts.LISTEN_PORT);
			return;
		}
		
		// Check '-p'.
		for (int i = 0; i <args.length; i ++){
			String arg = args[i];
			if (!arg.equals(Consts.CMD_ARG_PORT)){
				printHelp();
				return;
			}
			
			if (i == args.length -1){
				printHelp();
				return;
			}
			
			// Check port
			String sport = args[i+1].trim();
			if (sport.length() == 0){ // invalid
				printHelp();
				return;
			}
			
			try{
				int usePort = Integer.valueOf(sport);
				
				// start the server
				startServer(usePort);
			}catch(NumberFormatException e){
				printHelp();
				e.printStackTrace();
				return;
			}
		}
		
	}
	
	/**
	 * Start up the server on specific port.
	 * @param port Listen port
	 */
	private static void startServer(int port){
		
		try {
			ServerSocketChannel serverSocketChannel = ServerSocketChannel.open();
			SocketAddress socketAddress = new InetSocketAddress(port);
			
			// Bind serverSocketChannel
			serverSocketChannel.bind(socketAddress);
			
			System.out.println("Listening on port:"+port+".");
			while(true){
				
				// block mode
				SocketChannel socketChannel = serverSocketChannel.accept();
				
				ConnectionHandler.getInstance().add(socketChannel);
				
			}
		} catch (IOException e) {
			e.printStackTrace();
		}
	}
	
	/**
	 * Print help info.
	 */
	private static void printHelp(){
		System.out.println("Usage:");
		System.out.println("-p port");
		System.out.println("For example: -p 8200 will listen on port 8200.");
	}
}

2、SocketChannel处理类，主要通过线程池对连接进行处理，涉及的知识较多，主要是HTTP协议、非阻塞IO的使用。

package jim.proxy;

import java.io.IOException;
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.SocketChannel;
import java.util.Arrays;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;

/**
 * Connection handler.
 * Need to finish.
 * @author Jim-Zhang
 */
public class ConnectionHandler{
	
	/** Singleton */
	private static final ConnectionHandler instance = new ConnectionHandler();
	
	/** Scheduled worker */
	private ScheduledExecutorService executorService = Executors.newScheduledThreadPool(Consts.THREAD_COUNT);
	
	/** 
	 * Add channel.
	 * @param socketChannel SocketChannel to handle 
	 * */
	public void add(SocketChannel socketChannel){
		
		// One thread per request's SocketChannel
		executorService.execute(new HandlerThread(socketChannel));
	}
	
	/** 
	 * Get {@link ConnectionHandler} object.
	 * @return {@link ConnectionHandler} Instance
	 * */
	public static ConnectionHandler getInstance(){
		return instance;
	}
	
	/** Thread for handling socket channel. */
	private class HandlerThread extends Thread{
		
		/** Reqeust SocketChannel */
		private SocketChannel socketChannel;
		
		private final BlockingQueue<String> sourceRequestData = new ArrayBlockingQueue<String>(1000);
		
		/** Target */
		private SocketChannel targetSocketChannel;
		
		/** Constructor */
		public HandlerThread(SocketChannel socketChannel){
			HandlerThread.this.socketChannel = socketChannel;
			setDaemon(true);
		}
		
		@Override
		public void run() {
			
			// Use SocketChannel to read http data
			// final Socket socket = socketChannel.socket();
			String resourceUrl = null;
			try {
				// set non-block
				socketChannel.configureBlocking(false);
				Selector selector = Selector.open();
				SelectionKey sourceKey = socketChannel.register(selector, SelectionKey.OP_READ | SelectionKey.OP_WRITE);
				
				// for reading line 
				StringBuilder lineBuilder = new StringBuilder();
				ByteBuffer buffer = ByteBuffer.allocate(1024);
				String line = null; // null for no data
				
				// for target
				Selector targetSelector = null;
				SelectionKey targetSelectionKey = null;
				
				// Read request data
				while (true){ 
					
					Thread.sleep(20);
					
					buffer.clear();
					
					// Read request data
					selector.selectNow();
					
					if (sourceKey.isValid() && sourceKey.isReadable()){
						int readLen = socketChannel.read(buffer);
						if (readLen > 0){
							byte[] datas = new byte[readLen];
							buffer.rewind();
							buffer.get(datas);
							lineBuilder.append(new String(datas.clone()));
							byte[] dataEnd = new byte[]{datas[datas.length -2],datas[datas.length -1]};
							if (endof(dataEnd)){
								line = lineBuilder.toString();
								break;
							}
						}
						
						if (readLen == -1){
							line = lineBuilder.toString();
							break;
						}
					}
					
				}
				
				
				if (line == null || line.isEmpty()){
					return;
				}
				
				if (line != null){
					
					// Push request data to Queue to be handled later
					String requestStr = line.replace(HttpHeader.PROXY_CONNECTION, "Connection");
					sourceRequestData.add(requestStr);
				}
				
				// Handle target channel
				while(true && socketChannel.isOpen()){
					
					String lineToHandle = sourceRequestData.poll();
					// GET http://xx/ss/dd HTTP/1.1 
					// Immediately open the SocketChannel
					if (targetSocketChannel == null 
							&& lineToHandle != null 
							&& lineToHandle.indexOf("HTTP/1.1") != -1){
						String targetHostStr = lineToHandle.split(" ")[1];
						targetHostStr = targetHostStr.substring(targetHostStr.indexOf("//")+2);
						targetHostStr = targetHostStr.substring(0,targetHostStr.indexOf("/")); 
								
						// use SocketChannel to request data
						String[] hostParams = targetHostStr.split(HttpHeader.SPLIT_CHAR);
						String host = hostParams[0].trim();
						int port = hostParams.length == 1 ?  80 : Integer.valueOf(hostParams[1].trim());
						
						SocketAddress targetAddress = new InetSocketAddress(host, port);
							
						targetSocketChannel =  SocketChannel.open(targetAddress);
						targetSocketChannel.configureBlocking(false);
						
						targetSelector = Selector.open();
						targetSelectionKey = targetSocketChannel.register(targetSelector,
									SelectionKey.OP_READ | SelectionKey.OP_WRITE);
						
						resourceUrl = lineToHandle;
					}
					
					if(targetSocketChannel != null){
					
						targetSelector.select();
					    if (targetSelectionKey.isWritable()){
							String toTargetData = lineToHandle;
							if (toTargetData != null){
								byte[] datas = toTargetData.getBytes();
								ByteBuffer buffers = ByteBuffer.wrap(datas);
								// buffers.flip();
								targetSocketChannel.write(buffers);
								
								targetSocketChannel.write(ByteBuffer.wrap(Consts.LINE_END));
								targetSocketChannel.write(ByteBuffer.wrap(Consts.LINE_END));
								// targetSocketChannel.write(ByteBuffer.wrap(Consts.LINE_END));
								
								// line end
								// targetSocketChannel.write(ByteBuffer.wrap(Consts.LINE_END));
							}else{
								// targetSocketChannel.write(ByteBuffer.wrap(Consts.STREAM_END));
								// targetSocketChannel.write(ByteBuffer.wrap(Consts.LINE_END));
							}
						}
						
						if (targetSelectionKey.isReadable()){
							
							ByteBuffer buffers = ByteBuffer.allocate(1024);
							
							int len = targetSocketChannel.read(buffers);
							if (len == 0){
								continue;
							}
							
							if (len == -1){
								break;
							}
							
							buffers.flip();
							socketChannel.write(buffers);
							buffers.compact();
						}
					}
					
					Thread.sleep(20);
				}
				
				
			} catch (Throwable e) {
				
				try {
					System.out.println(resourceUrl);
					socketChannel.close();
					targetSocketChannel.close();
				} catch (IOException e1) {
					e1.printStackTrace();
				}
				e.printStackTrace();
			}
		}
		
		/** 
		 * End of stream or line.
		 * @param data Data
		 * @return true-end of line or stream
		 * */
		private boolean endof(byte[] datas){
			return Arrays.equals(Consts.LINE_END,datas) 
						|| Arrays.equals(Consts.STREAM_END,datas);
		}
	}
}

最后是常量定义类。

package jim.proxy;

/**
 * Constants for proxy.Should not instatiate.
 * @author Jim-Zhang
 *
 */
public class Consts {
	
	// private constructor
	private Consts(){throw new UnsupportedOperationException("Can not instatiate.");}
	
	/** The number of threads handle requests. */
	public static final int THREAD_COUNT = 200;
	
	/** Default listening port. */
	public static final int LISTEN_PORT = 4186;
	
	/** Listening port cmd argument */
	public static String CMD_ARG_PORT = "-p";
	
	/** Line end */
	public static String LINE_END_S = "\r\n";
	
	/** Line end */
	public static byte[] LINE_END = "\r\n".getBytes();
	
	/** Stream end */
	public static byte[] STREAM_END = new byte[]{0,0};
}

以上是代理的简单实现，仅是一个HTTP代理的简单思路，目的是为了使用NIO。HTTP协议的头结束、数据压缩等都有涉及，本文只对头做了处理，而对数据压缩未做处理。

目前我另外一个实现为采用Apache HttpComponents，采用其中NHttpServer\NHttpClient使用比较完整的代理。有兴趣可以参考下HttpComponents对Http的封装。