/*
 * Copyright (c) 2002-2011 JSCAPE
 * 1147 S. 53rd Pl., Mesa, Arizona, 85206, U.S.A.
 * All rights reserved.
 *
 * This software is the confidential and proprietary information of
 * JSCAPE. ("Confidential Information").  You shall not disclose such
 * Confidential Information and shall use it only in accordance with
 *  the terms of the license agreement you entered into with JSCAPE.
 */

package com.jscape.util.net.ssl;

import javax.net.ssl.SSLEngine;
import javax.net.ssl.SSLEngineResult;
import javax.net.ssl.SSLException;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.nio.channels.ReadableByteChannel;
import java.nio.channels.WritableByteChannel;

import static com.jscape.util.Parameters.assertNotNull;

/**
 * @author Alex Usun
 */
public class SslLayer {

    private final SSLEngine engine;
    private final PacketReader reader;
    private final PacketWriter writer;
    private final Handshaker handshaker;
    private final InputStream in;
    private final OutputStream out;

    public SslLayer(SSLEngine engine, InputStream sslIn, OutputStream sslOut) {
        assertNotNull(engine);
        this.engine = engine;

        this.reader = new PacketReader(engine, sslIn);
        this.writer = new PacketWriter(engine, sslOut);
        this.handshaker = new Handshaker(engine, this.reader, this.writer);

        this.in = new SslInputStream();
        this.out = new SslOutputStream();
    }

    public InputStream getInputStream() {
        return this.in;
    }

    public OutputStream getOutputStream() {
        return this.out;
    }

    private void write(ByteBuffer data)
            throws IOException {
        while (!this.writer.isClosed() && data.hasRemaining()) {
            this.writer.write(data);
            handleClose(this.writer.isClosed());
            handleHandshake();
        }
    }

    private ByteBuffer read(ByteBuffer buffer)
            throws IOException {
        while (!this.reader.isClosed() && buffer.position() == 0) {
            buffer = this.reader.read(buffer);
            handleClose(this.reader.isClosed());
            handleHandshake();
        }

        buffer.flip();

        return buffer;
    }

    private void handleClose(boolean closed)
            throws IOException {
        if (closed) {
            this.handshaker.shutdown();
        }
    }

    private void handleHandshake()
            throws IOException {
        this.handshaker.handleHandshake();
    }

    private void closeInbound()
            throws SSLException {
        this.engine.closeInbound();
    }

    private void closeOutbound() {
        this.engine.closeOutbound();
    }

    private enum Allocator {

        PACKET() {
            @Override
            public int size(SSLEngine engine) {
                return engine.getSession().getPacketBufferSize();
            }
        },

        APPLICATION() {
            @Override
            public int size(SSLEngine engine) {
                return engine.getSession().getApplicationBufferSize();
            }
        };

        public abstract int size(SSLEngine engine);

        public ByteBuffer allocate(SSLEngine engine) {
            return ByteBuffer.allocate(size(engine));
        }

        public ByteBuffer reallocate(SSLEngine engine, ByteBuffer source) {
            ByteBuffer ret = ByteBuffer.allocate(Math.max(size(engine), source.capacity() * 2));
            ret.put(source);

            return ret;
        }
    }

    private static class PacketWriter {

        private final SSLEngine engine;
        private final WritableByteChannel out;
        private ByteBuffer packet;
        private SSLEngineResult wrapResult;
        private boolean closed;

        private PacketWriter(SSLEngine engine, OutputStream out) {
            this.engine = engine;
            this.out = Channels.newChannel(out);
            this.packet = Allocator.PACKET.allocate(this.engine);
        }

        public synchronized void write(ByteBuffer data)
                throws IOException {
            makePacket(data);
            writePacket();
        }

        public boolean isClosed() {
            return this.closed;
        }

        private void makePacket(ByteBuffer data)
                throws SSLException {
            this.packet.clear();

            do {
                wrap(data);

                if (packetTooSmall()) {
                    adjustPacketSize();
                }
            } while (packetTooSmall());

            handleEngineClosed();
        }

        private void writePacket()
                throws IOException {
            if (this.packet.hasRemaining()) {
                this.out.write(this.packet);
            }
        }

        private void wrap(ByteBuffer data)
                throws SSLException {
            this.wrapResult = this.engine.wrap(data, this.packet);
            this.packet.flip();
        }

        private boolean packetTooSmall() {
            return this.wrapResult.getStatus() == SSLEngineResult.Status.BUFFER_OVERFLOW;
        }

        private void adjustPacketSize() {
            this.packet = Allocator.PACKET.reallocate(this.engine, this.packet);
        }

        private void handleEngineClosed() {
            if (this.wrapResult.getStatus() == SSLEngineResult.Status.CLOSED) {
                this.closed = true;
            }
        }
    }

    private static class PacketReader {

        private final SSLEngine engine;
        private final ReadableByteChannel in;
        private ByteBuffer packet;
        private SSLEngineResult unwrapResult;
        private boolean closed;

        private PacketReader(SSLEngine engine, InputStream in) {
            this.engine = engine;
            this.in = Channels.newChannel(in);
            this.packet = Allocator.PACKET.allocate(this.engine);
            this.packet.flip();
        }

        public synchronized ByteBuffer read(ByteBuffer data)
                throws IOException {
            do {
                unwrap(data);

                if (dataBufferTooSmall()) {
                    data = adjustDataBuffer(data);
                } else if (insufficientPacketData()) {
                    adjustPacketSize();
                    readPacket();
                }
            } while (hasBufferProblems());

            handleEngineClosed();

            return data;
        }

        public boolean isClosed() {
            return this.closed;
        }

        private void unwrap(ByteBuffer data)
                throws SSLException {
            this.unwrapResult = this.engine.unwrap(this.packet, data);
        }

        private boolean dataBufferTooSmall() {
            return this.unwrapResult.getStatus() == SSLEngineResult.Status.BUFFER_OVERFLOW;
        }

        private ByteBuffer adjustDataBuffer(ByteBuffer data) {
            data.flip();
            return Allocator.APPLICATION.reallocate(this.engine, data);
        }

        private boolean insufficientPacketData() {
            return this.unwrapResult.getStatus() == SSLEngineResult.Status.BUFFER_UNDERFLOW;
        }

        private void adjustPacketSize() {
            if (this.packet.remaining() == this.packet.capacity()) {
                this.packet = Allocator.PACKET.reallocate(this.engine, this.packet);
            } else if (this.packet.hasRemaining()) {
                this.packet.compact();
            } else {
                this.packet.clear();
            }
        }

        private boolean hasBufferProblems() {
            return dataBufferTooSmall() || insufficientPacketData();
        }

        private void readPacket()
                throws IOException {
            this.in.read(this.packet);
            this.packet.flip();
        }

        private void handleEngineClosed() {
            if (this.unwrapResult.getStatus() == SSLEngineResult.Status.CLOSED) {
                this.closed = true;
            }
        }
    }

    private static class Handshaker {

        private final SSLEngine engine;
        private final PacketReader reader;
        private final PacketWriter writer;
        private ByteBuffer data;

        private Handshaker(SSLEngine engine, PacketReader reader, PacketWriter writer) {
            this.engine = engine;
            this.reader = reader;
            this.writer = writer;
        }

        public synchronized void handleHandshake()
                throws IOException {
            if (handshakeRequired()) {
                handshake();
            }
        }

        public synchronized void shutdown()
                throws IOException {
            this.data = Allocator.APPLICATION.allocate(this.engine);

            do {
                writeData();
            } while (!this.engine.isOutboundDone());

            this.data = null;
        }

        private boolean handshakeRequired() {
            SSLEngineResult.HandshakeStatus status = this.engine.getHandshakeStatus();
            return status != SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING
                    && status != SSLEngineResult.HandshakeStatus.FINISHED;
        }

        private void handshake()
                throws IOException {
            this.data = Allocator.APPLICATION.allocate(this.engine);

            do {
                if (runDelegatedTaskRequired()) {
                    runDelegatedTasks();
                    writeData();
                } else if (writeDataRequired()) {
                    writeData();
                } else if (readDataRequired()) {
                    readData();
                }
            } while (handshakeRequired());

            this.data = null;
        }

        private boolean runDelegatedTaskRequired() {
            return this.engine.getHandshakeStatus() == SSLEngineResult.HandshakeStatus.NEED_TASK;
        }

        private void runDelegatedTasks() {
            Runnable task;
            while ((task = this.engine.getDelegatedTask()) != null) {
                task.run();
            }
        }

        private boolean writeDataRequired() {
            return this.engine.getHandshakeStatus() == SSLEngineResult.HandshakeStatus.NEED_WRAP;
        }

        private void writeData()
                throws IOException {
            this.data.clear();
            this.data.flip();
            this.writer.write(this.data);
        }

        private boolean readDataRequired() {
            return this.engine.getHandshakeStatus() == SSLEngineResult.HandshakeStatus.NEED_UNWRAP;
        }

        private void readData()
                throws IOException {
            this.data.clear();
            this.data = this.reader.read(this.data);
        }
    }

    private class SslInputStream
            extends InputStream {

        private ByteBuffer buffer;
        private final byte[] oneByteBuffer;

        private SslInputStream() {
            this.buffer = Allocator.APPLICATION.allocate(engine);
            this.buffer.flip();
            this.oneByteBuffer = new byte[1];
        }

        @Override
        public int read()
                throws IOException {
            int read;

            do {
                read = read(this.oneByteBuffer, 0, 1);
            } while (read == 0);

            return read != -1 ? this.oneByteBuffer[0] & 0xFF : -1;
        }

        @Override
        public int read(byte[] data, int offset, int length)
                throws IOException {
            if (reader.isClosed()) {
                return -1;
            }

            fillBuffer();

            return copyData(data, offset, length);
        }

        @Override
        public int available()
                throws IOException {
            return this.buffer.remaining();
        }

        @Override
        public void close()
                throws IOException {
            closeInbound();
        }

        private boolean eof(int read) {
            return read == -1;
        }

        private void fillBuffer()
                throws IOException {
            if (!this.buffer.hasRemaining()) {
                this.buffer.clear();
                this.buffer = SslLayer.this.read(this.buffer);
            }
        }

        private int copyData(byte[] data, int offset, int length) {
            if (!this.buffer.hasRemaining()) {
                return -1;
            }

            int bytesToRead = Math.min(length, this.buffer.remaining());
            this.buffer.get(data, offset, bytesToRead);

            return bytesToRead;
        }

    }

    private class SslOutputStream
            extends OutputStream {

        private final ByteBuffer buffer;
        private final byte[] oneByteBuffer;

        private SslOutputStream() {
            this.buffer = Allocator.APPLICATION.allocate(engine);
            this.oneByteBuffer = new byte[1];
        }

        @Override
        public void write(int data)
                throws IOException {
            this.oneByteBuffer[0] = (byte) data;
            write(this.oneByteBuffer);
        }

        @Override
        public void write(byte[] data, int offset, int length)
                throws IOException {
            assertNotClosed();

            while (length > 0) {
                fillBuffer(data, offset, length);
                length -= this.buffer.remaining();
                offset += this.buffer.remaining();
                writeBuffer();
                assertNotClosed(length);
            }
        }

        private void writeBuffer()
                throws IOException {
            SslLayer.this.write(this.buffer);
        }

        @Override
        public void close()
                throws IOException {
            closeOutbound();
        }

        private void assertNotClosed()
                throws IOException {
            if (writer.isClosed()) {
                throw new IOException("Stream is closed.");
            }
        }

        private void fillBuffer(byte[] data, int offset, int length) {
            this.buffer.clear();
            int bytesToWrite = Math.min(length, this.buffer.capacity());
            this.buffer.put(data, offset, bytesToWrite);
            this.buffer.flip();
        }

        private void assertNotClosed(int remainedDataLength)
                throws IOException {
            if (writer.isClosed() && remainedDataLength > 0) {
                throw new IOException("Stream is closed.");
            }
        }
    }
}
