/*
 * Copyright 2012 The Netty Project
 *
 * The Netty Project 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.
 */
// (BSD License: http://www.opensource.org/licenses/bsd-license)
//
// Copyright (c) 2011, Joe Walnes and contributors
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or
// without modification, are permitted provided that the
// following conditions are met:
//
// * Redistributions of source code must retain the above
// copyright notice, this list of conditions and the
// following disclaimer.
//
// * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the
// following disclaimer in the documentation and/or other
// materials provided with the distribution.
//
// * Neither the name of the Webbit nor the names of
// its contributors may be used to endorse or promote products
// derived from this software without specific prior written
// permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
// CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE
// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
// BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
// OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

package org.jboss.netty.handler.codec.http.websocketx;

import org.jboss.netty.buffer.ChannelBuffer;
import org.jboss.netty.buffer.ChannelBuffers;
import org.jboss.netty.channel.Channel;
import org.jboss.netty.channel.ChannelFutureListener;
import org.jboss.netty.channel.ChannelHandlerContext;
import org.jboss.netty.handler.codec.frame.CorruptedFrameException;
import org.jboss.netty.handler.codec.frame.TooLongFrameException;
import org.jboss.netty.handler.codec.replay.ReplayingDecoder;
import org.jboss.netty.logging.InternalLogger;
import org.jboss.netty.logging.InternalLoggerFactory;

/**
 * Decodes a web socket frame from wire protocol version 8 format. This code was forked from <a
 * href="https://github.com/joewalnes/webbit">webbit</a> and modified.
 */
public class WebSocket08FrameDecoder extends ReplayingDecoder<WebSocket08FrameDecoder.State> {

    private static final InternalLogger logger = InternalLoggerFactory.getInstance(WebSocket08FrameDecoder.class);

    private static final byte OPCODE_CONT = 0x0;
    private static final byte OPCODE_TEXT = 0x1;
    private static final byte OPCODE_BINARY = 0x2;
    private static final byte OPCODE_CLOSE = 0x8;
    private static final byte OPCODE_PING = 0x9;
    private static final byte OPCODE_PONG = 0xA;

    private UTF8Output fragmentedFramesText;
    private int fragmentedFramesCount;

    private final long maxFramePayloadLength;
    private boolean frameFinalFlag;
    private int frameRsv;
    private int frameOpcode;
    private long framePayloadLength;
    private ChannelBuffer framePayload;
    private int framePayloadBytesRead;
    private ChannelBuffer maskingKey;

    private final boolean allowExtensions;
    private final boolean maskedPayload;
    private boolean receivedClosingHandshake;

    public enum State {
        FRAME_START, MASKING_KEY, PAYLOAD, CORRUPT
    }

    /**
     * Constructor with default values
     *
     * @param maskedPayload
     *            Web socket servers must set this to true processed incoming masked payload. Client implementations
     *            must set this to false.
     * @param allowExtensions
     *            Flag to allow reserved extension bits to be used or not
     */
    public WebSocket08FrameDecoder(boolean maskedPayload, boolean allowExtensions) {
        this(maskedPayload, allowExtensions, Long.MAX_VALUE);
    }

    /**
     * Constructor
     *
     * @param maskedPayload
     *            Web socket servers must set this to true processed incoming masked payload. Client implementations
     *            must set this to false.
     * @param allowExtensions
     *            Flag to allow reserved extension bits to be used or not
     * @param maxFramePayloadLength
     *            Maximum length of a frame's payload. Setting this to an appropriate value for you application
     *            helps check for denial of services attacks.
     */
    public WebSocket08FrameDecoder(boolean maskedPayload, boolean allowExtensions, long maxFramePayloadLength) {
        super(State.FRAME_START);
        this.maskedPayload = maskedPayload;
        this.allowExtensions = allowExtensions;
        this.maxFramePayloadLength = maxFramePayloadLength;
    }

    @Override
    protected Object decode(ChannelHandlerContext ctx, Channel channel, ChannelBuffer buffer, State state)
            throws Exception {

        // Discard all data received if closing handshake was received before.
        if (receivedClosingHandshake) {
            buffer.skipBytes(actualReadableBytes());
            return null;
        }

        switch (state) {
        case FRAME_START:
            framePayloadBytesRead = 0;
            framePayloadLength = -1;
            framePayload = null;

            // FIN, RSV, OPCODE
            byte b = buffer.readByte();
            frameFinalFlag = (b & 0x80) != 0;
            frameRsv = (b & 0x70) >> 4;
            frameOpcode = b & 0x0F;

            if (logger.isDebugEnabled()) {
                logger.debug("Decoding WebSocket Frame opCode=" + frameOpcode);
            }

            // MASK, PAYLOAD LEN 1
            b = buffer.readByte();
            boolean frameMasked = (b & 0x80) != 0;
            int framePayloadLen1 = b & 0x7F;

            if (frameRsv != 0 && !allowExtensions) {
                protocolViolation(channel, "RSV != 0 and no extension negotiated, RSV:" + frameRsv);
                return null;
            }

            if (maskedPayload && !frameMasked) {
                protocolViolation(channel, "unmasked client to server frame");
                return null;
            }
            if (frameOpcode > 7) { // control frame (have MSB in opcode set)

                // control frames MUST NOT be fragmented
                if (!frameFinalFlag) {
                    protocolViolation(channel, "fragmented control frame");
                    return null;
                }

                // control frames MUST have payload 125 octets or less
                if (framePayloadLen1 > 125) {
                    protocolViolation(channel, "control frame with payload length > 125 octets");
                    return null;
                }

                // check for reserved control frame opcodes
                if (!(frameOpcode == OPCODE_CLOSE || frameOpcode == OPCODE_PING || frameOpcode == OPCODE_PONG)) {
                    protocolViolation(channel, "control frame using reserved opcode " + frameOpcode);
                    return null;
                }

                // close frame : if there is a body, the first two bytes of the
                // body MUST be a 2-byte unsigned integer (in network byte
                // order) representing a status code
                if (frameOpcode == 8 && framePayloadLen1 == 1) {
                    protocolViolation(channel, "received close control frame with payload len 1");
                    return null;
                }
            } else { // data frame
                // check for reserved data frame opcodes
                if (!(frameOpcode == OPCODE_CONT || frameOpcode == OPCODE_TEXT || frameOpcode == OPCODE_BINARY)) {
                    protocolViolation(channel, "data frame using reserved opcode " + frameOpcode);
                    return null;
                }

                // check opcode vs message fragmentation state 1/2
                if (fragmentedFramesCount == 0 && frameOpcode == OPCODE_CONT) {
                    protocolViolation(channel, "received continuation data frame outside fragmented message");
                    return null;
                }

                // check opcode vs message fragmentation state 2/2
                if (fragmentedFramesCount != 0 && frameOpcode != OPCODE_CONT && frameOpcode != OPCODE_PING) {
                    protocolViolation(channel, "received non-continuation data frame while inside fragmented message");
                    return null;
                }
            }

            // Read frame payload length
            if (framePayloadLen1 == 126) {
                framePayloadLength = buffer.readUnsignedShort();
                if (framePayloadLength < 126) {
                    protocolViolation(channel, "invalid data frame length (not using minimal length encoding)");
                    return null;
                }
            } else if (framePayloadLen1 == 127) {
                framePayloadLength = buffer.readLong();
                // TODO: check if it's bigger than 0x7FFFFFFFFFFFFFFF, Maybe
                // just check if it's negative?

                if (framePayloadLength < 65536) {
                    protocolViolation(channel, "invalid data frame length (not using minimal length encoding)");
                    return null;
                }
            } else {
                framePayloadLength = framePayloadLen1;
            }

            if (framePayloadLength > maxFramePayloadLength) {
                protocolViolation(channel, "Max frame length of " + maxFramePayloadLength + " has been exceeded.");
                return null;
            }
            if (logger.isDebugEnabled()) {
                logger.debug("Decoding WebSocket Frame length=" + framePayloadLength);
            }

            checkpoint(State.MASKING_KEY);
        case MASKING_KEY:
            if (maskedPayload) {
                maskingKey = buffer.readBytes(4);
            }
            checkpoint(State.PAYLOAD);
        case PAYLOAD:
            // Sometimes, the payload may not be delivered in 1 nice packet
            // We need to accumulate the data until we have it all
            int rbytes = actualReadableBytes();
            ChannelBuffer payloadBuffer = null;

            long  willHaveReadByteCount = framePayloadBytesRead + rbytes;
            // logger.debug("Frame rbytes=" + rbytes + " willHaveReadByteCount="
            // + willHaveReadByteCount + " framePayloadLength=" +
            // framePayloadLength);
            if (willHaveReadByteCount == framePayloadLength) {
                // We have all our content so proceed to process
                payloadBuffer = buffer.readBytes(rbytes);
            } else if (willHaveReadByteCount < framePayloadLength) {
                // We don't have all our content so accumulate payload.
                // Returning null means we will get called back
                payloadBuffer = buffer.readBytes(rbytes);
                if (framePayload == null) {
                    framePayload = channel.getConfig().getBufferFactory().getBuffer(toFrameLength(framePayloadLength));
                }
                framePayload.writeBytes(payloadBuffer);
                framePayloadBytesRead += rbytes;

                // Return null to wait for more bytes to arrive
                return null;
            } else if (willHaveReadByteCount > framePayloadLength) {
                // We have more than what we need so read up to the end of frame
                // Leave the remainder in the buffer for next frame
                payloadBuffer = buffer.readBytes(toFrameLength(framePayloadLength - framePayloadBytesRead));
            }

            // Now we have all the data, the next checkpoint must be the next
            // frame
            checkpoint(State.FRAME_START);

            // Take the data that we have in this packet
            if (framePayload == null) {
                framePayload = payloadBuffer;
            } else {
                framePayload.writeBytes(payloadBuffer);
            }

            // Unmask data if needed
            if (maskedPayload) {
                unmask(framePayload);
            }

            // Processing ping/pong/close frames because they cannot be
            // fragmented
            if (frameOpcode == OPCODE_PING) {
                return new PingWebSocketFrame(frameFinalFlag, frameRsv, framePayload);
            } else if (frameOpcode == OPCODE_PONG) {
                return new PongWebSocketFrame(frameFinalFlag, frameRsv, framePayload);
            } else if (frameOpcode == OPCODE_CLOSE) {
                checkCloseFrameBody(channel, framePayload);
                receivedClosingHandshake = true;
                return new CloseWebSocketFrame(frameFinalFlag, frameRsv, framePayload);
            }

            // Processing for possible fragmented messages for text and binary
            // frames
            String aggregatedText = null;
            if (frameFinalFlag) {
                // Final frame of the sequence. Apparently ping frames are
                // allowed in the middle of a fragmented message
                if (frameOpcode != OPCODE_PING) {
                    fragmentedFramesCount = 0;

                    // Check text for UTF8 correctness
                    if (frameOpcode == OPCODE_TEXT || fragmentedFramesText != null) {
                        // Check UTF-8 correctness for this payload
                        checkUTF8String(channel, framePayload.array());

                        // This does a second check to make sure UTF-8
                        // correctness for entire text message
                        aggregatedText = fragmentedFramesText.toString();

                        fragmentedFramesText = null;
                    }
                }
            } else {
                // Not final frame so we can expect more frames in the
                // fragmented sequence
                if (fragmentedFramesCount == 0) {
                    // First text or binary frame for a fragmented set
                    fragmentedFramesText = null;
                    if (frameOpcode == OPCODE_TEXT) {
                        checkUTF8String(channel, framePayload.array());
                    }
                } else {
                    // Subsequent frames - only check if init frame is text
                    if (fragmentedFramesText != null) {
                        checkUTF8String(channel, framePayload.array());
                    }
                }

                // Increment counter
                fragmentedFramesCount++;
            }

            // Return the frame
            if (frameOpcode == OPCODE_TEXT) {
                return new TextWebSocketFrame(frameFinalFlag, frameRsv, framePayload);
            } else if (frameOpcode == OPCODE_BINARY) {
                return new BinaryWebSocketFrame(frameFinalFlag, frameRsv, framePayload);
            } else if (frameOpcode == OPCODE_CONT) {
                return new ContinuationWebSocketFrame(frameFinalFlag, frameRsv, framePayload, aggregatedText);
            } else {
                throw new UnsupportedOperationException("Cannot decode web socket frame with opcode: " + frameOpcode);
            }
        case CORRUPT:
            // If we don't keep reading Netty will throw an exception saying
            // we can't return null if no bytes read and state not changed.
            buffer.readByte();
            return null;
        default:
            throw new Error("Shouldn't reach here.");
        }
    }

    private void unmask(ChannelBuffer frame) {
        byte[] bytes = frame.array();
        for (int i = 0; i < bytes.length; i++) {
            frame.setByte(i, frame.getByte(i) ^ maskingKey.getByte(i % 4));
        }
    }

    private void protocolViolation(Channel channel, String reason) throws CorruptedFrameException {
        checkpoint(State.CORRUPT);
        if (channel.isConnected()) {
            channel.write(ChannelBuffers.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
            channel.close().awaitUninterruptibly();
        }
        throw new CorruptedFrameException(reason);
    }

    private static int toFrameLength(long l) throws TooLongFrameException {
        if (l > Integer.MAX_VALUE) {
            throw new TooLongFrameException("Length:" + l);
        } else {
            return (int) l;
        }
    }

    private void checkUTF8String(Channel channel, byte[] bytes) throws CorruptedFrameException {
        try {
            // StringBuilder sb = new StringBuilder("UTF8 " + bytes.length +
            // " bytes: ");
            // for (byte b : bytes) {
            // sb.append(Integer.toHexString(b)).append(" ");
            // }
            // logger.debug(sb.toString());

            if (fragmentedFramesText == null) {
                fragmentedFramesText = new UTF8Output(bytes);
            } else {
                fragmentedFramesText.write(bytes);
            }
        } catch (UTF8Exception ex) {
            protocolViolation(channel, "invalid UTF-8 bytes");
        }
    }

    protected void checkCloseFrameBody(Channel channel, ChannelBuffer buffer) throws CorruptedFrameException {
        if (buffer == null || buffer.capacity() == 0) {
            return;
        }
        if (buffer.capacity() == 1) {
            protocolViolation(channel, "Invalid close frame body");
        }

        // Save reader index
        int idx = buffer.readerIndex();
        buffer.readerIndex(0);

        // Must have 2 byte integer within the valid range
        int statusCode = buffer.readShort();
        if (statusCode >= 0 && statusCode <= 999 || statusCode >= 1004 && statusCode <= 1006
                || statusCode >= 1012 && statusCode <= 2999) {
            protocolViolation(channel, "Invalid close frame status code: " + statusCode);
        }

        // May have UTF-8 message
        if (buffer.readableBytes() > 0) {
            byte[] b = new byte[buffer.readableBytes()];
            buffer.readBytes(b);
            try {
                new UTF8Output(b);
            } catch (UTF8Exception ex) {
                protocolViolation(channel, "Invalid close frame reason text. Invalid UTF-8 bytes");
            }
        }

        // Restore reader index
        buffer.readerIndex(idx);
    }
}
