//
//  ========================================================================
//  Copyright (c) 1995-2017 Mort Bay Consulting Pty. Ltd.
//  ------------------------------------------------------------------------
//  All rights reserved. This program and the accompanying materials
//  are made available under the terms of the Eclipse Public License v1.0
//  and Apache License v2.0 which accompanies this distribution.
//
//      The Eclipse Public License is available at
//      http://www.eclipse.org/legal/epl-v10.html
//
//      The Apache License v2.0 is available at
//      http://www.opensource.org/licenses/apache2.0.php
//
//  You may elect to redistribute this code under either of these licenses.
//  ========================================================================
//

package org.eclipse.jetty.websocket.common;

import java.nio.ByteBuffer;
import java.util.Arrays;

import org.eclipse.jetty.util.BufferUtil;
import org.eclipse.jetty.util.StringUtil;
import org.eclipse.jetty.util.log.Log;
import org.eclipse.jetty.util.log.Logger;
import org.eclipse.jetty.websocket.api.StatusCode;
import org.eclipse.jetty.websocket.api.WebSocketPolicy;
import org.eclipse.jetty.websocket.api.extensions.Frame;
import org.eclipse.jetty.websocket.common.frames.BinaryFrame;
import org.eclipse.jetty.websocket.common.frames.CloseFrame;
import org.eclipse.jetty.websocket.common.frames.PingFrame;
import org.eclipse.jetty.websocket.common.frames.TextFrame;
import org.eclipse.jetty.websocket.common.test.IncomingFramesCapture;
import org.eclipse.jetty.websocket.common.test.UnitGenerator;
import org.eclipse.jetty.websocket.common.test.UnitParser;
import org.eclipse.jetty.websocket.common.util.Hex;
import org.junit.Assert;
import org.junit.Test;

import static org.hamcrest.Matchers.is;

public class GeneratorTest
{
    private static final Logger LOG = Log.getLogger(GeneratorTest.WindowHelper.class);

    public static class WindowHelper
    {
        final int windowSize;
        int totalParts;
        int totalBytes;

        public WindowHelper(int windowSize)
        {
            this.windowSize = windowSize;
            this.totalParts = 0;
            this.totalBytes = 0;
        }

        public ByteBuffer generateWindowed(Frame... frames)
        {
            // Create Buffer to hold all generated frames in a single buffer
            int completeBufSize = 0;
            for (Frame f : frames)
            {
                completeBufSize += Generator.MAX_HEADER_LENGTH + f.getPayloadLength();
            }

            ByteBuffer completeBuf = ByteBuffer.allocate(completeBufSize);
            BufferUtil.clearToFill(completeBuf);

            // Generate from all frames
            Generator generator = new UnitGenerator();

            for (Frame f : frames)
            {
                ByteBuffer header = generator.generateHeaderBytes(f);
                totalBytes += BufferUtil.put(header,completeBuf);

                if (f.hasPayload())
                {
                    ByteBuffer payload=f.getPayload();
                    totalBytes += payload.remaining();
                    totalParts++;
                    completeBuf.put(payload.slice());
                }
            }

            // Return results
            BufferUtil.flipToFlush(completeBuf,0);
            return completeBuf;
        }

        public void assertTotalParts(int expectedParts)
        {
            Assert.assertThat("Generated Parts",totalParts,is(expectedParts));
        }

        public void assertTotalBytes(int expectedBytes)
        {
            Assert.assertThat("Generated Bytes",totalBytes,is(expectedBytes));
        }
    }

    private void assertGeneratedBytes(CharSequence expectedBytes, Frame... frames)
    {
        // collect up all frames as single ByteBuffer
        ByteBuffer allframes = UnitGenerator.generate(frames);
        // Get hex String form of all frames bytebuffer.
        String actual = Hex.asHex(allframes);
        // Validate
        Assert.assertThat("Buffer",actual,is(expectedBytes.toString()));
    }

    private String asMaskedHex(String str, byte[] maskingKey)
    {
        byte utf[] = StringUtil.getUtf8Bytes(str);
        mask(utf,maskingKey);
        return Hex.asHex(utf);
    }

    private void mask(byte[] buf, byte[] maskingKey)
    {
        int size = buf.length;
        for (int i = 0; i < size; i++)
        {
            buf[i] ^= maskingKey[i % 4];
        }
    }

    @Test
    public void testClose_Empty()
    {
        // 0 byte payload (no status code)
        assertGeneratedBytes("8800",new CloseFrame());
    }

    @Test
    public void testClose_CodeNoReason()
    {
        CloseInfo close = new CloseInfo(StatusCode.NORMAL);
        // 2 byte payload (2 bytes for status code)
        assertGeneratedBytes("880203E8",close.asFrame());
    }

    @Test
    public void testClose_CodeOkReason()
    {
        CloseInfo close = new CloseInfo(StatusCode.NORMAL,"OK");
        // 4 byte payload (2 bytes for status code, 2 more for "OK")
        assertGeneratedBytes("880403E84F4B",close.asFrame());
    }

    @Test
    public void testText_Hello()
    {
        WebSocketFrame frame = new TextFrame().setPayload("Hello");
        byte utf[] = StringUtil.getUtf8Bytes("Hello");
        assertGeneratedBytes("8105" + Hex.asHex(utf),frame);
    }

    @Test
    public void testText_Masked()
    {
        WebSocketFrame frame = new TextFrame().setPayload("Hello");
        byte maskingKey[] = Hex.asByteArray("11223344");
        frame.setMask(maskingKey);

        // what is expected
        StringBuilder expected = new StringBuilder();
        expected.append("8185").append("11223344");
        expected.append(asMaskedHex("Hello",maskingKey));

        // validate
        assertGeneratedBytes(expected,frame);
    }

    @Test
    public void testText_Masked_OffsetSourceByteBuffer()
    {
        ByteBuffer payload = ByteBuffer.allocate(100);
        payload.position(5);
        payload.put(StringUtil.getUtf8Bytes("Hello"));
        payload.flip();
        payload.position(5);
        // at this point, we have a ByteBuffer of 100 bytes.
        // but only a few bytes in the middle are made available for the payload.
        // we are testing that masking works as intended, even if the provided
        // payload does not start at position 0.
        LOG.debug("Payload = {}",BufferUtil.toDetailString(payload));
        WebSocketFrame frame = new TextFrame().setPayload(payload);
        byte maskingKey[] = Hex.asByteArray("11223344");
        frame.setMask(maskingKey);

        // what is expected
        StringBuilder expected = new StringBuilder();
        expected.append("8185").append("11223344");
        expected.append(asMaskedHex("Hello",maskingKey));

        // validate
        assertGeneratedBytes(expected,frame);
    }

    /**
     * Prevent regression of masking of many packets.
     */
    @Test
    public void testManyMasked()
    {
        int pingCount = 2;

        // Prepare frames
        WebSocketFrame[] frames = new WebSocketFrame[pingCount + 1];
        for (int i = 0; i < pingCount; i++)
        {
            frames[i] = new PingFrame().setPayload(String.format("ping-%d",i));
        }
        frames[pingCount] = new CloseInfo(StatusCode.NORMAL).asFrame();

        // Mask All Frames
        byte maskingKey[] = Hex.asByteArray("11223344");
        for (WebSocketFrame f : frames)
        {
            f.setMask(maskingKey);
        }

        // Validate result of generation
        StringBuilder expected = new StringBuilder();
        expected.append("8986").append("11223344");
        expected.append(asMaskedHex("ping-0",maskingKey)); // ping 0
        expected.append("8986").append("11223344");
        expected.append(asMaskedHex("ping-1",maskingKey)); // ping 1
        expected.append("8882").append("11223344");
        byte closure[] = Hex.asByteArray("03E8");
        mask(closure,maskingKey);
        expected.append(Hex.asHex(closure)); // normal closure

        assertGeneratedBytes(expected,frames);
    }

    /**
     * Test the windowed generate of a frame that has no masking.
     */
    @Test
    public void testWindowedGenerate()
    {
        // A decent sized frame, no masking
        byte payload[] = new byte[10240];
        Arrays.fill(payload,(byte)0x44);

        WebSocketFrame frame = new BinaryFrame().setPayload(payload);

        // Generate
        int windowSize = 1024;
        WindowHelper helper = new WindowHelper(windowSize);
        ByteBuffer completeBuffer = helper.generateWindowed(frame);

        // Validate
        int expectedHeaderSize = 4;
        int expectedSize = payload.length + expectedHeaderSize;
        int expectedParts = 1;

        helper.assertTotalParts(expectedParts);
        helper.assertTotalBytes(payload.length + expectedHeaderSize);

        Assert.assertThat("Generated Buffer",completeBuffer.remaining(),is(expectedSize));
    }

    @Test
    public void testWindowedGenerateWithMasking()
    {
        // A decent sized frame, with masking
        byte payload[] = new byte[10240];
        Arrays.fill(payload,(byte)0x55);

        byte mask[] = new byte[]
        { 0x2A, (byte)0xF0, 0x0F, 0x00 };

        WebSocketFrame frame = new BinaryFrame().setPayload(payload);
        frame.setMask(mask); // masking!

        // Generate
        int windowSize = 2929;
        WindowHelper helper = new WindowHelper(windowSize);
        ByteBuffer completeBuffer = helper.generateWindowed(frame);

        // Validate
        int expectedHeaderSize = 8;
        int expectedSize = payload.length + expectedHeaderSize;
        int expectedParts = 1;

        helper.assertTotalParts(expectedParts);
        helper.assertTotalBytes(payload.length + expectedHeaderSize);

        Assert.assertThat("Generated Buffer",completeBuffer.remaining(),is(expectedSize));

        // Parse complete buffer.
        WebSocketPolicy policy = WebSocketPolicy.newServerPolicy();
        Parser parser = new UnitParser(policy);
        IncomingFramesCapture capture = new IncomingFramesCapture();
        parser.setIncomingFramesHandler(capture);

        parser.parse(completeBuffer);

        // Assert validity of frame
        WebSocketFrame actual = capture.getFrames().poll();
        Assert.assertThat("Frame.opcode",actual.getOpCode(),is(OpCode.BINARY));
        Assert.assertThat("Frame.payloadLength",actual.getPayloadLength(),is(payload.length));

        // Validate payload contents for proper masking
        ByteBuffer actualData = actual.getPayload().slice();
        Assert.assertThat("Frame.payload.remaining",actualData.remaining(),is(payload.length));
        while (actualData.remaining() > 0)
        {
            Assert.assertThat("Actual.payload[" + actualData.position() + "]",actualData.get(),is((byte)0x55));
        }
    }
}