sslhandler.java

apache 的一个socket框架

/*
 *  Licensed to the Apache Software Foundation (ASF) under one
 *  or more contributor license agreements.  See the NOTICE file
 *  distributed with this work for additional information
 *  regarding copyright ownership.  The ASF licenses this file
 *  to you under the Apache License, Version 2.0 (the
 *  "License"); you may not use this file except in compliance
 *  with the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing,
 *  software distributed under the License is distributed on an
 *  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 *  KIND, either express or implied.  See the License for the
 *  specific language governing permissions and limitations
 *  under the License.
 *
 */
package org.apache.mina.filter.support;

import java.nio.ByteBuffer;
import java.util.LinkedList;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;

import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLEngine;
import javax.net.ssl.SSLEngineResult;
import javax.net.ssl.SSLException;
import javax.net.ssl.SSLHandshakeException;
import javax.net.ssl.SSLSession;

import org.apache.mina.common.IoSession;
import org.apache.mina.common.WriteFuture;
import org.apache.mina.common.IoFilter.NextFilter;
import org.apache.mina.common.IoFilter.WriteRequest;
import org.apache.mina.common.support.DefaultWriteFuture;
import org.apache.mina.filter.SSLFilter;
import org.apache.mina.util.SessionLog;

/**
 * A helper class using the SSLEngine API to decrypt/encrypt data.
 * <p>
 * Each connection has a SSLEngine that is used through the lifetime of the connection.
 * We allocate byte buffers for use as the outbound and inbound network buffers.
 * These buffers handle all of the intermediary data for the SSL connection. To make things easy,
 * we'll require outNetBuffer be completely flushed before trying to wrap any more data.
 *
 * @author The Apache Directory Project ([e-mail])
 * @version $Rev: 557169 $, $Date: 2007-07-18 15:26:04 +0900 (수, 18  7월 2007) $
 */
public class SSLHandler {
    private final SSLFilter parent;

    private final SSLContext ctx;

    private final IoSession session;

    private final Queue<Event> preHandshakeEventQueue = new LinkedList<Event>();

    private final Queue<Event> filterWriteEventQueue = new ConcurrentLinkedQueue<Event>();

    private final Queue<Event> messageReceivedEventQueue = new ConcurrentLinkedQueue<Event>();

    private SSLEngine sslEngine;

    /**
     * Encrypted data from the net
     */
    private ByteBuffer inNetBuffer;

    /**
     * Encrypted data to be written to the net
     */
    private ByteBuffer outNetBuffer;

    /**
     * Applicaton cleartext data to be read by application
     */
    private ByteBuffer appBuffer;

    /**
     * Empty buffer used during initial handshake and close operations
     */
    private final ByteBuffer hsBB = ByteBuffer.allocate(0);

    /**
     * Handshake status
     */
    private SSLEngineResult.HandshakeStatus initialHandshakeStatus;

    /**
     * Initial handshake complete?
     */
    private boolean initialHandshakeComplete;

    private boolean writingEncryptedData;

    /**
     * Constuctor.
     *
     * @param sslc
     * @throws SSLException
     */
    public SSLHandler(SSLFilter parent, SSLContext sslc, IoSession session)
            throws SSLException {
        this.parent = parent;
        this.session = session;
        this.ctx = sslc;
        init();
    }

    public void init() throws SSLException {
        if (sslEngine != null) {
            return;
        }

        sslEngine = ctx.createSSLEngine();
        sslEngine.setUseClientMode(parent.isUseClientMode());

        if (parent.isWantClientAuth()) {
            sslEngine.setWantClientAuth(true);
        }

        if (parent.isNeedClientAuth()) {
            sslEngine.setNeedClientAuth(true);
        }

        if (parent.getEnabledCipherSuites() != null) {
            sslEngine.setEnabledCipherSuites(parent.getEnabledCipherSuites());
        }

        if (parent.getEnabledProtocols() != null) {
            sslEngine.setEnabledProtocols(parent.getEnabledProtocols());
        }

        sslEngine.beginHandshake();
        initialHandshakeStatus = sslEngine.getHandshakeStatus();//SSLEngineResult.HandshakeStatus.NEED_UNWRAP;
        initialHandshakeComplete = false;

        SSLByteBufferPool.initiate(sslEngine);

        appBuffer = SSLByteBufferPool.getApplicationBuffer();

        inNetBuffer = SSLByteBufferPool.getPacketBuffer();
        outNetBuffer = SSLByteBufferPool.getPacketBuffer();
        outNetBuffer.position(0);
        outNetBuffer.limit(0);

        writingEncryptedData = false;
    }

    /**
     * Release allocated ByteBuffers.
     */
    public void destroy() {
        if (sslEngine == null) {
            return;
        }

        // Close inbound and flush all remaining data if available.
        try {
            sslEngine.closeInbound();
        } catch (SSLException e) {
            SessionLog.debug(session,
                    "Unexpected exception from SSLEngine.closeInbound().", e);
        }

        try {
            do {
                outNetBuffer.clear();
            } while (sslEngine.wrap(hsBB, outNetBuffer).bytesProduced() > 0);
        } catch (SSLException e) {
            SessionLog.debug(session,
                    "Unexpected exception from SSLEngine.wrap().", e);
        }
        sslEngine.closeOutbound();
        sslEngine = null;

        SSLByteBufferPool.release(appBuffer);
        SSLByteBufferPool.release(inNetBuffer);
        SSLByteBufferPool.release(outNetBuffer);
        preHandshakeEventQueue.clear();
        //postHandshakeEventQueue.clear();
    }

    public SSLFilter getParent() {
        return parent;
    }

    public IoSession getSession() {
        return session;
    }

    /**
     * Check we are writing encrypted data.
     */
    public boolean isWritingEncryptedData() {
        return writingEncryptedData;
    }

    /**
     * Check if initial handshake is completed.
     */
    public boolean isInitialHandshakeComplete() {
        return initialHandshakeComplete;
    }

    public boolean isInboundDone() {
        return sslEngine == null || sslEngine.isInboundDone();
    }

    public boolean isOutboundDone() {
        return sslEngine == null || sslEngine.isOutboundDone();
    }

    /**
     * Check if there is any need to complete initial handshake.
     */
    public boolean needToCompleteInitialHandshake() {
        return (initialHandshakeStatus == SSLEngineResult.HandshakeStatus.NEED_WRAP && !isInboundDone());
    }

    public void schedulePreHandshakeWriteRequest(NextFilter nextFilter,
            WriteRequest writeRequest) {
        preHandshakeEventQueue.offer(new Event(EventType.FILTER_WRITE,
                nextFilter, writeRequest));
    }

    public void flushPreHandshakeEvents() throws SSLException {
        Event scheduledWrite;

        while ((scheduledWrite = preHandshakeEventQueue.poll()) != null) {
            if (SessionLog.isDebugEnabled(session)) {
                SessionLog.debug(session, " Flushing buffered write request: "
                        + scheduledWrite.data);
            }
            parent.filterWrite(scheduledWrite.nextFilter, session,
                    (WriteRequest) scheduledWrite.data);
        }
    }

    public void scheduleFilterWrite(NextFilter nextFilter,
            WriteRequest writeRequest) {
        filterWriteEventQueue.offer(new Event(EventType.FILTER_WRITE,
                nextFilter, writeRequest));
    }

    public void scheduleMessageReceived(NextFilter nextFilter,
            Object message) {
        messageReceivedEventQueue.offer(new Event(EventType.RECEIVED, nextFilter,
                message));
    }
    
    public void flushScheduledEvents() {
        // Fire events only when no lock is hold for this handler.
        if (Thread.holdsLock(this)) {
            return;
        }

        Event e;
        
        // We need synchronization here inevitably because filterWrite can be
        // called simultaneously and cause 'bad record MAC' integrity error.
        synchronized (this) {
            while ((e = filterWriteEventQueue.poll()) != null) {
                e.nextFilter.filterWrite(session, (WriteRequest) e.data);
            }
        }

        while ((e = messageReceivedEventQueue.poll()) != null) {
            e.nextFilter.messageReceived(session, e.data);
        }
    }

    /**
     * Call when data read from net. Will perform inial hanshake or decrypt provided
     * Buffer.
     * Decrytpted data reurned by getAppBuffer(), if any.
     *
     * @param buf buffer to decrypt
     * @throws SSLException on errors
     */
    public void messageReceived(NextFilter nextFilter, ByteBuffer buf)
            throws SSLException {
        if (buf.limit() > inNetBuffer.remaining()) {
            // We have to expand inNetBuffer
            inNetBuffer = SSLByteBufferPool.expandBuffer(inNetBuffer,
                    inNetBuffer.capacity() + (buf.limit() * 2));
            // We also expand app. buffer (twice the size of in net. buffer)
            appBuffer = SSLByteBufferPool.expandBuffer(appBuffer, inNetBuffer
                    .capacity() * 2);
            appBuffer.position(0);
            appBuffer.limit(0);
            if (SessionLog.isDebugEnabled(session)) {
                SessionLog.debug(session, " expanded inNetBuffer:"
                        + inNetBuffer);
                SessionLog.debug(session, " expanded appBuffer:" + appBuffer);
            }
        }

        // append buf to inNetBuffer
        inNetBuffer.put(buf);
        if (!initialHandshakeComplete) {
            handshake(nextFilter);
        } else {
            decrypt();
        }

        if (isInboundDone()) {
            // Rewind the MINA buffer if not all data is processed and inbound is finished.
            buf.position(buf.position() - inNetBuffer.position());
            inNetBuffer.clear();
        }
    }

    /**
     * Get decrypted application data.
     *
     * @return buffer with data
     */
    public ByteBuffer getAppBuffer() {
        return appBuffer;
    }

    /**
     * Get encrypted data to be sent.
     *
     * @return buffer with data
     */
    public ByteBuffer getOutNetBuffer() {
        return outNetBuffer;
    }

    /**
     * Encrypt provided buffer. Encytpted data reurned by getOutNetBuffer().
     *
     * @param src data to encrypt
     * @throws SSLException on errors
     */
    public void encrypt(ByteBuffer src) throws SSLException {
        if (!initialHandshakeComplete) {
            throw new IllegalStateException();
        }

        // The data buffer is (must be) empty, we can reuse the entire
        // buffer.
        outNetBuffer.clear();

        // Loop until there is no more data in src
        while (src.hasRemaining()) {

            if (src.remaining() > ((outNetBuffer.capacity() - outNetBuffer
                    .position()) / 2)) {
                // We have to expand outNetBuffer
                // Note: there is no way to know the exact size required, but enrypted data
                // shouln't need to be larger than twice the source data size?
                outNetBuffer = SSLByteBufferPool.expandBuffer(outNetBuffer, src
                        .capacity() * 2);
                if (SessionLog.isDebugEnabled(session)) {
                    SessionLog.debug(session, " expanded outNetBuffer:"
                            + outNetBuffer);
                }
            }

            SSLEngineResult result = sslEngine.wrap(src, outNetBuffer);
            if (SessionLog.isDebugEnabled(session)) {
                SessionLog.debug(session, " Wrap res:" + result);
            }

            if (result.getStatus() == SSLEngineResult.Status.OK) {
                if (result.getHandshakeStatus() == SSLEngineResult.HandshakeStatus.NEED_TASK) {
                    doTasks();
                }
            } else {
                throw new SSLException("SSLEngine error during encrypt: "
                        + result.getStatus() + " src: " + src
                        + "outNetBuffer: " + outNetBuffer);
            }
        }

        outNetBuffer.flip();
    }