TLSStreamReader.java 7.74 KB
Newer Older
1 2 3 4 5
/**
 * $RCSfile$
 * $Revision: $
 * $Date: $
 *
6
 * Copyright (C) 2005-2008 Jive Software. All rights reserved.
7 8
 *
 * This software is published under the terms of the GNU Public License (GPL),
9 10
 * a copy of which is included in this distribution, or a commercial license
 * agreement with Jive.
11 12
 */

13
package org.jivesoftware.openfire.net;
14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50

import java.io.IOException;
import java.io.InputStream;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.nio.channels.ReadableByteChannel;

/**
 * A <code>TLSStreamReader</code> that returns a special InputStream that hides the ByteBuffers
 * used by the underlying Channels.
 * 
 * @author Hao Chen
 */
public class TLSStreamReader {

	/**
	 * <code>TLSWrapper</code> is a TLS wrapper for connections requiring TLS protocol.
	 */
	private TLSWrapper wrapper;

	private ReadableByteChannel rbc;

	/**
	 * <code>inNetBB</code> buffer keeps data read from socket.
	 */
	private ByteBuffer inNetBB;

	/**
	 * <code>inAppBB</code> buffer keeps decypted data.
	 */
	private ByteBuffer inAppBB;

    private TLSStatus lastStatus;

    public TLSStreamReader(TLSWrapper tlsWrapper, Socket socket) throws IOException {
		wrapper = tlsWrapper;
51 52
        // DANIELE: Add code to use directly the socket channel
        if (socket.getChannel() != null) {
53
            rbc = ServerTrafficCounter.wrapReadableChannel(socket.getChannel());
54 55
        }
        else {
56 57
            rbc = Channels.newChannel(
                    ServerTrafficCounter.wrapInputStream(socket.getInputStream()));
58 59
        }
        inNetBB = ByteBuffer.allocate(wrapper.getNetBuffSize());
60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195
		inAppBB = ByteBuffer.allocate(wrapper.getAppBuffSize());
	}

	/*
	 * Read TLS encrpyted data from SocketChannel, and use <code>decrypt</code> method to decypt.
	 */
	private void doRead() throws IOException {
        //System.out.println("doRead inNet position: " + inNetBB.position() + " capacity: " + inNetBB.capacity() + " (before read)");

        // Read from the channel and fill inNetBB with the encrypted data
        final int cnt = rbc.read(inNetBB);
		if (cnt > 0) {
            //System.out.println("doRead inNet position: " + inNetBB.position() + " capacity: " + inNetBB.capacity() + " (after read)");
            //System.out.println("doRead inAppBB (before decrypt) position: " + inAppBB.position() + " limit: " + inAppBB.limit() + " capacity: " + inAppBB.capacity());

            // Decode encrypted data
            inAppBB = decrypt(inNetBB, inAppBB);

            ///System.out.println("doRead inAppBB (after decrypt) position: " + inAppBB.position() + " limit: " + inAppBB.limit() + " capacity: " + inAppBB.capacity() + " lastStatus: " + lastStatus);

            if (lastStatus == TLSStatus.OK) {
                // All the data contained in inNetBB was read and decrypted so we can safely
                // set the position of inAppBB to 0 to process it.
                inAppBB.flip();
            }
            else {
                // Some data in inNetBB was not decrypted since it is not complete. A
                // bufferunderflow was detected since the TLS packet is not complete to be
                // decrypted. We need to read more data from the channel to decrypt the whole
                // TLS packet. The inNetBB byte buffer has been compacted so the read and
                // decrypted is discarded and only the unread and encrypted data is left in the
                // buffer. The inAppBB has been completed with the decrypted data and we must
                // leave the position at the end of the written so that in the next doRead the
                // decrypted data is appended to the end of the buffer.
                //System.out.println("Reading more data from the channel (UNDERFLOW state)");
                doRead();
            }
        } else {
			if (cnt == -1) {
                inAppBB.flip();
                rbc.close();
			}
		}
	}

	/*
	 * This method uses <code>TLSWrapper</code> to decrypt TLS encrypted data.
	 */
	private ByteBuffer decrypt(ByteBuffer input, ByteBuffer output) throws IOException {
		ByteBuffer out = output;
        input.flip();
        do {
            // Decode SSL/TLS network data and place it in the app buffer
			out = wrapper.unwrap(input, out);

			lastStatus = wrapper.getStatus();
        }
        while ((lastStatus == TLSStatus.NEED_READ || lastStatus == TLSStatus.OK) &&
                input.hasRemaining());

        if (input.hasRemaining()) {
            // Complete TLS packets have been read, decrypted and written to the output buffer.
            // However, the input buffer contains incomplete TLS packets that cannot be decrpted.
            // Discard the read data and keep the unread data in the input buffer. The channel will
            // be read again to obtain the missing data to complete the TLS packet. So in the next
            // round the TLS packet will be decrypted and written to the output buffer
            input.compact();
		} else {
            // All the encrypted data in the inpu buffer was decrypted so we can clear
            // the input buffer.
            input.clear();
		}

		return out;
	}

	public InputStream getInputStream() {
		return createInputStream();
	}

	/*
	 * Returns an input stream for a ByteBuffer. The read() methods use the relative ByteBuffer
	 * get() methods.
	 */
	private InputStream createInputStream() {
		return new InputStream() {
			public synchronized int read() throws IOException {
				doRead();
				if (!inAppBB.hasRemaining()) {
					return -1;
				}
				return inAppBB.get();
			}

			public synchronized int read(byte[] bytes, int off, int len) throws IOException {
                // Check if in the previous read the inAppBB ByteBuffer remained with unread data.
                // If all the data was consumed then read from the socket channel. Otherwise,
                // consume the data contained in the buffer.
                if (inAppBB.position() == 0) {
                    // Read from the channel the encrypted data, decrypt it and load it
                    // into inAppBB
                    doRead();
                }
                else {
                    //System.out.println("#createInputStream. Detected previously unread data. position: " + inAppBB.position());

                    // The inAppBB contains data from a previous read so set the position to 0
                    // to consume it
                    inAppBB.flip();
                }
				len = Math.min(len, inAppBB.remaining());
                if (len == 0) {
                    // Nothing was read so the end of stream should have been reached.
                    return -1;
                }
                inAppBB.get(bytes, off, len);
                // If the requested length is less than the limit of inAppBB then all the data
                // inside inAppBB was not read. In that case we need to discard the read data and
                // keep only the unread data to be consume the next time this method is called
                if (inAppBB.hasRemaining()) {
                    // Discard read data and move unread data to the begining of the buffer. Leave
                    // the position at the end of the buffer as a way to indicate that there is
                    // unread data
                    inAppBB.compact();

                    //System.out.println("#createInputStream. Data left unread. inAppBB compacted. position: " + inAppBB.position() + " limit: " + inAppBB.limit());
                }
                else {
                    // Everything was read so reset the buffer
                    inAppBB.clear();
                }
				return len;
			}
		};
	}
}