secureniochannel.java

业界著名的tomcat服务器的最新6.0的源代码。

package org.apache.tomcat.util.net;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.SocketChannel;
import javax.net.ssl.SSLEngine;
import javax.net.ssl.SSLEngineResult;
import javax.net.ssl.SSLEngineResult.HandshakeStatus;
import javax.net.ssl.SSLEngineResult.Status;
import java.nio.channels.Selector;

/**
 * 
 * Implementation of a secure socket channel
 * @author Filip Hanik
 * @version 1.0
 */

public class SecureNioChannel extends NioChannel  {
    
    protected ByteBuffer netInBuffer;
    protected ByteBuffer netOutBuffer;
    
    protected SSLEngine sslEngine;
    
    protected boolean initHandshakeComplete = false;
    protected HandshakeStatus initHandshakeStatus; //gets set by begin handshake
    
    protected boolean closed = false;
    protected boolean closing = false;
    
    protected NioSelectorPool pool;
    
    public SecureNioChannel(SocketChannel channel, SSLEngine engine, 
                            ApplicationBufferHandler bufHandler, NioSelectorPool pool) throws IOException {
        super(channel,bufHandler);
        this.sslEngine = engine;
        int appBufSize = sslEngine.getSession().getApplicationBufferSize();
        int netBufSize = sslEngine.getSession().getPacketBufferSize();
        //allocate network buffers - TODO, add in optional direct non-direct buffers
        if ( netInBuffer == null ) netInBuffer = ByteBuffer.allocateDirect(netBufSize);
        if ( netOutBuffer == null ) netOutBuffer = ByteBuffer.allocateDirect(netBufSize);
        
        //selector pool for blocking operations
        this.pool = pool;
        
        //ensure that the application has a large enough read/write buffers
        //by doing this, we should not encounter any buffer overflow errors
        bufHandler.expand(bufHandler.getReadBuffer(), appBufSize);
        bufHandler.expand(bufHandler.getWriteBuffer(), appBufSize);
        reset();
    }
    
    public void reset(SSLEngine engine) throws IOException {
        this.sslEngine = engine;
        reset();
    }
    public void reset() throws IOException {
        super.reset();
        netOutBuffer.position(0);
        netOutBuffer.limit(0);
        netInBuffer.position(0);
        netInBuffer.limit(0);
        initHandshakeComplete = false;
        closed = false;
        closing = false;
        //initiate handshake
        sslEngine.beginHandshake();
        initHandshakeStatus = sslEngine.getHandshakeStatus();
    }
    
    public int getBufferSize() {
        int size = super.getBufferSize();
        size += netInBuffer!=null?netInBuffer.capacity():0;
        size += netOutBuffer!=null?netOutBuffer.capacity():0;
        return size;
    }

    
//===========================================================================================    
//                  NIO SSL METHODS
//===========================================================================================
    /**
     * returns true if the network buffer has 
     * been flushed out and is empty
     * @return boolean
     */
    public boolean flush(Selector s, long timeout) throws IOException {
        pool.write(netOutBuffer,this,s,timeout);
        return !netOutBuffer.hasRemaining();
    }
    
    /**
     * Flushes the buffer to the network, non blocking
     * @param buf ByteBuffer
     * @return boolean true if the buffer has been emptied out, false otherwise
     * @throws IOException
     */
    protected boolean flush(ByteBuffer buf) throws IOException {
        int remaining = buf.remaining();
        if ( remaining > 0 ) {
            int written = sc.write(buf);
            return written >= remaining;
        }else {
            return true;
        }
    }
    
    /**
     * Performs SSL handshake, non blocking, but performs NEED_TASK on the same thread.<br>
     * Hence, you should never call this method using your Acceptor thread, as you would slow down
     * your system significantly.<br>
     * The return for this operation is 0 if the handshake is complete and a positive value if it is not complete.
     * In the event of a positive value coming back, reregister the selection key for the return values interestOps.
     * @param read boolean - true if the underlying channel is readable
     * @param write boolean - true if the underlying channel is writable
     * @return int - 0 if hand shake is complete, otherwise it returns a SelectionKey interestOps value
     * @throws IOException
     */
    public int handshake(boolean read, boolean write) throws IOException {
        if ( initHandshakeComplete ) return 0; //we have done our initial handshake
        
        if (!flush(netOutBuffer)) return SelectionKey.OP_WRITE; //we still have data to write
        
        SSLEngineResult handshake = null;
        
        while (!initHandshakeComplete) {
            switch ( initHandshakeStatus ) {
                case NOT_HANDSHAKING: {
                    //should never happen
                    throw new IOException("NOT_HANDSHAKING during handshake");
                }
                case FINISHED: {
                    //we are complete if we have delivered the last package
                    initHandshakeComplete = !netOutBuffer.hasRemaining();
                    //return 0 if we are complete, otherwise we still have data to write
                    return initHandshakeComplete?0:SelectionKey.OP_WRITE; 
                }
                case NEED_WRAP: {
                    //perform the wrap function
                    handshake = handshakeWrap(write);
                    if ( handshake.getStatus() == Status.OK ){
                        if (initHandshakeStatus == HandshakeStatus.NEED_TASK) 
                            initHandshakeStatus = tasks();
                    } else {
                        //wrap should always work with our buffers
                        throw new IOException("Unexpected status:" + handshake.getStatus() + " during handshake WRAP.");
                    }
                    if ( initHandshakeStatus != HandshakeStatus.NEED_UNWRAP || (!flush(netOutBuffer)) ) {
                        //should actually return OP_READ if we have NEED_UNWRAP
                        return SelectionKey.OP_WRITE;
                    }
                    //fall down to NEED_UNWRAP on the same call, will result in a 
                    //BUFFER_UNDERFLOW if it needs data
                }
                case NEED_UNWRAP: {
                    //perform the unwrap function
                    handshake = handshakeUnwrap(read);
                    if ( handshake.getStatus() == Status.OK ) {
                        if (initHandshakeStatus == HandshakeStatus.NEED_TASK) 
                            initHandshakeStatus = tasks();
                    } else if ( handshake.getStatus() == Status.BUFFER_UNDERFLOW ){
                        //read more data, reregister for OP_READ
                        return SelectionKey.OP_READ;
                    } else {
                        throw new IOException("Invalid handshake status:"+initHandshakeStatus+" during handshake UNWRAP.");
                    }//switch
                    break;
                }
                case NEED_TASK: {
                    initHandshakeStatus = tasks();
                    break;
                }
                default: throw new IllegalStateException("Invalid handshake status:"+initHandshakeStatus);
            }//switch
        }//while      
        //return 0 if we are complete, otherwise reregister for any activity that 
        //would cause this method to be called again.
        return initHandshakeComplete?0:(SelectionKey.OP_WRITE|SelectionKey.OP_READ);
    }
    
    /**
     * Executes all the tasks needed on the same thread.
     * @return HandshakeStatus
     */
    protected SSLEngineResult.HandshakeStatus tasks() {
        Runnable r = null;
        while ( (r = sslEngine.getDelegatedTask()) != null) {
            r.run();
        }
        return sslEngine.getHandshakeStatus();
    }

    /**
     * Performs the WRAP function
     * @param doWrite boolean
     * @return SSLEngineResult
     * @throws IOException
     */
    protected SSLEngineResult handshakeWrap(boolean doWrite) throws IOException {
        //this should never be called with a network buffer that contains data
        //so we can clear it here.
        netOutBuffer.clear();
        //perform the wrap
        SSLEngineResult result = sslEngine.wrap(bufHandler.getWriteBuffer(), netOutBuffer);
        //prepare the results to be written
        netOutBuffer.flip();
        //set the status
        initHandshakeStatus = result.getHandshakeStatus();
        //optimization, if we do have a writable channel, write it now
        if ( doWrite ) flush(netOutBuffer);
        return result;
    }
    
    /**
     * Perform handshake unwrap
     * @param doread boolean
     * @return SSLEngineResult
     * @throws IOException
     */
    protected SSLEngineResult handshakeUnwrap(boolean doread) throws IOException {