//
//  ========================================================================
//  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.client;

import static org.hamcrest.Matchers.*;

import java.io.IOException;
import java.lang.reflect.Field;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.SocketChannel;
import java.util.Arrays;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

import org.eclipse.jetty.io.EndPoint;
import org.eclipse.jetty.io.EofException;
import org.eclipse.jetty.io.SelectChannelEndPoint;
import org.eclipse.jetty.io.SelectorManager.ManagedSelector;
import org.eclipse.jetty.toolchain.test.EventQueue;
import org.eclipse.jetty.toolchain.test.TestTracker;
import org.eclipse.jetty.util.BufferUtil;
import org.eclipse.jetty.util.log.Log;
import org.eclipse.jetty.util.log.Logger;
import org.eclipse.jetty.util.log.StacklessLogging;
import org.eclipse.jetty.util.thread.Scheduler;
import org.eclipse.jetty.websocket.api.ProtocolException;
import org.eclipse.jetty.websocket.api.Session;
import org.eclipse.jetty.websocket.api.StatusCode;
import org.eclipse.jetty.websocket.api.WebSocketAdapter;
import org.eclipse.jetty.websocket.client.io.ConnectionManager;
import org.eclipse.jetty.websocket.client.io.WebSocketClientSelectorManager;
import org.eclipse.jetty.websocket.common.CloseInfo;
import org.eclipse.jetty.websocket.common.OpCode;
import org.eclipse.jetty.websocket.common.Parser;
import org.eclipse.jetty.websocket.common.WebSocketFrame;
import org.eclipse.jetty.websocket.common.WebSocketSession;
import org.eclipse.jetty.websocket.common.frames.TextFrame;
import org.eclipse.jetty.websocket.common.io.AbstractWebSocketConnection;
import org.eclipse.jetty.websocket.common.test.BlockheadServer;
import org.eclipse.jetty.websocket.common.test.BlockheadServer.ServerConnection;
import org.eclipse.jetty.websocket.common.test.IncomingFramesCapture;
import org.eclipse.jetty.websocket.common.test.RawFrameBuilder;
import org.hamcrest.Matcher;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;

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

    private static class CloseTrackingSocket extends WebSocketAdapter
    {
        private static final Logger LOG = Log.getLogger(ClientCloseTest.CloseTrackingSocket.class);

        public int closeCode = -1;
        public String closeReason = null;
        public CountDownLatch closeLatch = new CountDownLatch(1);
        public AtomicInteger closeCount = new AtomicInteger(0);
        public CountDownLatch openLatch = new CountDownLatch(1);

        public EventQueue<String> messageQueue = new EventQueue<>();
        public EventQueue<Throwable> errorQueue = new EventQueue<>();

        public void assertNoCloseEvent()
        {
            Assert.assertThat("Client Close Event",closeLatch.getCount(),is(1L));
            Assert.assertThat("Client Close Event Status Code ",closeCode,is(-1));
        }

        public void assertReceivedCloseEvent(int clientTimeoutMs, Matcher<Integer> statusCodeMatcher, Matcher<String> reasonMatcher)
                throws InterruptedException
        {
            long maxTimeout = clientTimeoutMs * 2;

            Assert.assertThat("Client Close Event Occurred",closeLatch.await(maxTimeout,TimeUnit.MILLISECONDS),is(true));
            Assert.assertThat("Client Close Event Count",closeCount.get(),is(1));
            Assert.assertThat("Client Close Event Status Code",closeCode,statusCodeMatcher);
            if (reasonMatcher == null)
            {
                Assert.assertThat("Client Close Event Reason",closeReason,nullValue());
            }
            else
            {
                Assert.assertThat("Client Close Event Reason",closeReason,reasonMatcher);
            }
        }

        public void assertReceivedError(Class<? extends Throwable> expectedThrownClass, Matcher<String> messageMatcher) throws TimeoutException,
                InterruptedException
        {
            errorQueue.awaitEventCount(1,500,TimeUnit.MILLISECONDS);
            Throwable actual = errorQueue.poll();
            Assert.assertThat("Client Error Event",actual,instanceOf(expectedThrownClass));
            if (messageMatcher == null)
            {
                Assert.assertThat("Client Error Event Message",actual.getMessage(),nullValue());
            }
            else
            {
                Assert.assertThat("Client Error Event Message",actual.getMessage(),messageMatcher);
            }
        }

        public void clearQueues()
        {
            messageQueue.clear();
            errorQueue.clear();
        }

        @Override
        public void onWebSocketClose(int statusCode, String reason)
        {
            LOG.debug("onWebSocketClose({},{})",statusCode,reason);
            super.onWebSocketClose(statusCode,reason);
            closeCount.incrementAndGet();
            closeCode = statusCode;
            closeReason = reason;
            closeLatch.countDown();
        }

        @Override
        public void onWebSocketConnect(Session session)
        {
            super.onWebSocketConnect(session);
            openLatch.countDown();
        }

        @Override
        public void onWebSocketError(Throwable cause)
        {
            LOG.debug("onWebSocketError",cause);
            Assert.assertThat("Error capture",errorQueue.offer(cause),is(true));
        }

        @Override
        public void onWebSocketText(String message)
        {
            LOG.debug("onWebSocketText({})",message);
            messageQueue.offer(message);
        }

        public EndPoint getEndPoint() throws Exception
        {
            Session session = getSession();
            Assert.assertThat("Session type",session,instanceOf(WebSocketSession.class));

            WebSocketSession wssession = (WebSocketSession)session;
            Field fld = wssession.getClass().getDeclaredField("connection");
            fld.setAccessible(true);
            Assert.assertThat("Field: connection",fld,notNullValue());

            Object val = fld.get(wssession);
            Assert.assertThat("Connection type",val,instanceOf(AbstractWebSocketConnection.class));
            @SuppressWarnings("resource")
            AbstractWebSocketConnection wsconn = (AbstractWebSocketConnection)val;
            return wsconn.getEndPoint();
        }
    }

    @Rule
    public TestTracker tt = new TestTracker();

    private BlockheadServer server;
    private WebSocketClient client;

    private void confirmConnection(CloseTrackingSocket clientSocket, Future<Session> clientFuture, ServerConnection serverConn) throws Exception
    {
        // Wait for client connect on via future
        clientFuture.get(500,TimeUnit.MILLISECONDS);

        // Wait for client connect via client websocket
        Assert.assertThat("Client WebSocket is Open",clientSocket.openLatch.await(500,TimeUnit.MILLISECONDS),is(true));

        try
        {
            // Send message from client to server
            final String echoMsg = "echo-test";
            Future<Void> testFut = clientSocket.getRemote().sendStringByFuture(echoMsg);

            // Wait for send future
            testFut.get(500,TimeUnit.MILLISECONDS);

            // Read Frame on server side
            IncomingFramesCapture serverCapture = serverConn.readFrames(1,500,TimeUnit.MILLISECONDS);
            serverCapture.assertNoErrors();
            serverCapture.assertFrameCount(1);
            WebSocketFrame frame = serverCapture.getFrames().poll();
            Assert.assertThat("Server received frame",frame.getOpCode(),is(OpCode.TEXT));
            Assert.assertThat("Server received frame payload",frame.getPayloadAsUTF8(),is(echoMsg));

            // Server send echo reply
            serverConn.write(new TextFrame().setPayload(echoMsg));

            // Wait for received echo
            clientSocket.messageQueue.awaitEventCount(1,1,TimeUnit.SECONDS);

            // Verify received message
            String recvMsg = clientSocket.messageQueue.poll();
            Assert.assertThat("Received message",recvMsg,is(echoMsg));

            // Verify that there are no errors
            Assert.assertThat("Error events",clientSocket.errorQueue,empty());
        }
        finally
        {
            clientSocket.clearQueues();
        }
    }

    private void confirmServerReceivedCloseFrame(ServerConnection serverConn, int expectedCloseCode, Matcher<String> closeReasonMatcher) throws IOException,
            TimeoutException
    {
        IncomingFramesCapture serverCapture = serverConn.readFrames(1,500,TimeUnit.MILLISECONDS);
        serverCapture.assertNoErrors();
        serverCapture.assertFrameCount(1);
        serverCapture.assertHasFrame(OpCode.CLOSE,1);
        WebSocketFrame frame = serverCapture.getFrames().poll();
        Assert.assertThat("Server received close frame",frame.getOpCode(),is(OpCode.CLOSE));
        CloseInfo closeInfo = new CloseInfo(frame);
        Assert.assertThat("Server received close code",closeInfo.getStatusCode(),is(expectedCloseCode));
        if (closeReasonMatcher == null)
        {
            Assert.assertThat("Server received close reason",closeInfo.getReason(),nullValue());
        }
        else
        {
            Assert.assertThat("Server received close reason",closeInfo.getReason(),closeReasonMatcher);
        }
    }

    public static class TestWebSocketClient extends WebSocketClient
    {
        @Override
        protected ConnectionManager newConnectionManager()
        {
            return new TestConnectionManager(this);
        }
    }

    public static class TestConnectionManager extends ConnectionManager
    {
        public TestConnectionManager(WebSocketClient client)
        {
            super(client);
        }

        @Override
        protected WebSocketClientSelectorManager newWebSocketClientSelectorManager(WebSocketClient client)
        {
            return new TestSelectorManager(client);
        }
    }

    public static class TestSelectorManager extends WebSocketClientSelectorManager
    {
        public TestSelectorManager(WebSocketClient client)
        {
            super(client);
        }

        @Override
        protected EndPoint newEndPoint(SocketChannel channel, ManagedSelector selectSet, SelectionKey selectionKey) throws IOException
        {
            return new TestEndPoint(channel,selectSet,selectionKey,getScheduler(),getPolicy().getIdleTimeout());
        }
    }

    public static class TestEndPoint extends SelectChannelEndPoint
    {
        public AtomicBoolean congestedFlush = new AtomicBoolean(false);

        public TestEndPoint(SocketChannel channel, ManagedSelector selector, SelectionKey key, Scheduler scheduler, long idleTimeout)
        {
            super(channel,selector,key,scheduler,idleTimeout);
        }

        @Override
        public boolean flush(ByteBuffer... buffers) throws IOException
        {
            boolean flushed = super.flush(buffers);
            congestedFlush.set(!flushed);
            return flushed;
        }
    }

    @Before
    public void startClient() throws Exception
    {
        client = new TestWebSocketClient();
        client.start();
    }

    @Before
    public void startServer() throws Exception
    {
        server = new BlockheadServer();
        server.start();
    }

    @After
    public void stopClient() throws Exception
    {
        if (client.isRunning())
        {
            client.stop();
        }
    }

    @After
    public void stopServer() throws Exception
    {
        server.stop();
    }

    @Test
    public void testHalfClose() throws Exception
    {
        // Set client timeout
        final int timeout = 1000;
        client.setMaxIdleTimeout(timeout);

        // Client connects
        CloseTrackingSocket clientSocket = new CloseTrackingSocket();
        Future<Session> clientConnectFuture = client.connect(clientSocket,server.getWsUri());

        // Server accepts connect
        ServerConnection serverConn = server.accept();
        serverConn.upgrade();

        // client confirms connection via echo
        confirmConnection(clientSocket,clientConnectFuture,serverConn);

        // client sends close frame (code 1000, normal)
        final String origCloseReason = "Normal Close";
        clientSocket.getSession().close(StatusCode.NORMAL,origCloseReason);

        // server receives close frame
        confirmServerReceivedCloseFrame(serverConn,StatusCode.NORMAL,is(origCloseReason));

        // server sends 2 messages
        serverConn.write(new TextFrame().setPayload("Hello"));
        serverConn.write(new TextFrame().setPayload("World"));

        // server sends close frame (code 1000, no reason)
        CloseInfo sclose = new CloseInfo(StatusCode.NORMAL,"From Server");
        serverConn.write(sclose.asFrame());

        // client receives 2 messages
        clientSocket.messageQueue.awaitEventCount(2,1,TimeUnit.SECONDS);

        // Verify received messages
        String recvMsg = clientSocket.messageQueue.poll();
        Assert.assertThat("Received message 1",recvMsg,is("Hello"));
        recvMsg = clientSocket.messageQueue.poll();
        Assert.assertThat("Received message 2",recvMsg,is("World"));

        // Verify that there are no errors
        Assert.assertThat("Error events",clientSocket.errorQueue,empty());

        // client close event on ws-endpoint
        clientSocket.assertReceivedCloseEvent(timeout,is(StatusCode.NORMAL),containsString("From Server"));
    }

    @Test
    public void testNetworkCongestion() throws Exception
    {
        // Set client timeout
        final int timeout = 1000;
        client.setMaxIdleTimeout(timeout);

        // Client connects
        CloseTrackingSocket clientSocket = new CloseTrackingSocket();
        Future<Session> clientConnectFuture = client.connect(clientSocket,server.getWsUri());

        // Server accepts connect
        ServerConnection serverConn = server.accept();
        serverConn.upgrade();

        // client confirms connection via echo
        confirmConnection(clientSocket,clientConnectFuture,serverConn);

        // client sends BIG frames (until it cannot write anymore)
        // server must not read (for test purpose, in order to congest connection)
        // when write is congested, client enqueue close frame
        // client initiate write, but write never completes
        EndPoint endp = clientSocket.getEndPoint();
        Assert.assertThat("EndPoint is testable",endp,instanceOf(TestEndPoint.class));
        TestEndPoint testendp = (TestEndPoint)endp;

        char msg[] = new char[10240];
        int writeCount = 0;
        long writeSize = 0;
        int i = 0;
        while (!testendp.congestedFlush.get())
        {
            int z = i - ((i / 26) * 26);
            char c = (char)('a' + z);
            Arrays.fill(msg,c);
            clientSocket.getRemote().sendStringByFuture(String.valueOf(msg));
            writeCount++;
            writeSize += msg.length;
        }
        LOG.debug("Wrote {} frames totalling {} bytes of payload before congestion kicked in",writeCount,writeSize);

        // Verify that there are no errors
        Assert.assertThat("Error events",clientSocket.errorQueue,empty());

        // client idle timeout triggers close event on client ws-endpoint
        // client close event on ws-endpoint
        clientSocket.assertReceivedCloseEvent(timeout,
                anyOf(is(StatusCode.SHUTDOWN),is(StatusCode.ABNORMAL)),
                anyOf(containsString("Timeout"),containsString("timeout"),containsString("Write")));
    }

    @Test
    public void testProtocolException() throws Exception
    {
        // Set client timeout
        final int timeout = 1000;
        client.setMaxIdleTimeout(timeout);

        // Client connects
        CloseTrackingSocket clientSocket = new CloseTrackingSocket();
        Future<Session> clientConnectFuture = client.connect(clientSocket,server.getWsUri());

        // Server accepts connect
        ServerConnection serverConn = server.accept();
        serverConn.upgrade();

        // client confirms connection via echo
        confirmConnection(clientSocket,clientConnectFuture,serverConn);

        // client should not have received close message (yet)
        clientSocket.assertNoCloseEvent();

        // server sends bad close frame (too big of a reason message)
        byte msg[] = new byte[400];
        Arrays.fill(msg,(byte)'x');
        ByteBuffer bad = ByteBuffer.allocate(500);
        RawFrameBuilder.putOpFin(bad,OpCode.CLOSE,true);
        RawFrameBuilder.putLength(bad,msg.length + 2,false);
        bad.putShort((short)StatusCode.NORMAL);
        bad.put(msg);
        BufferUtil.flipToFlush(bad,0);
        try (StacklessLogging quiet = new StacklessLogging(Parser.class))
        {
            serverConn.write(bad);

            // client should have noticed the error
            clientSocket.assertReceivedError(ProtocolException.class,containsString("Invalid control frame"));

            // client parse invalid frame, notifies server of close (protocol error)
            confirmServerReceivedCloseFrame(serverConn,StatusCode.PROTOCOL,allOf(containsString("Invalid control frame"),containsString("length")));
        }

        // server disconnects
        serverConn.disconnect();

        // client triggers close event on client ws-endpoint
        clientSocket.assertReceivedCloseEvent(timeout,is(StatusCode.PROTOCOL),allOf(containsString("Invalid control frame"),containsString("length")));
    }

    @Test
    public void testReadEOF() throws Exception
    {
        // Set client timeout
        final int timeout = 1000;
        client.setMaxIdleTimeout(timeout);

        // Client connects
        CloseTrackingSocket clientSocket = new CloseTrackingSocket();
        Future<Session> clientConnectFuture = client.connect(clientSocket,server.getWsUri());

        // Server accepts connect
        ServerConnection serverConn = server.accept();
        serverConn.upgrade();

        // client confirms connection via echo
        confirmConnection(clientSocket,clientConnectFuture,serverConn);

        // client sends close frame
        final String origCloseReason = "Normal Close";
        clientSocket.getSession().close(StatusCode.NORMAL,origCloseReason);

        // server receives close frame
        confirmServerReceivedCloseFrame(serverConn,StatusCode.NORMAL,is(origCloseReason));

        // client should not have received close message (yet)
        clientSocket.assertNoCloseEvent();

        // server shuts down connection (no frame reply)
        serverConn.disconnect();

        // client reads -1 (EOF)
        // client triggers close event on client ws-endpoint
        clientSocket.assertReceivedCloseEvent(timeout,is(StatusCode.ABNORMAL),containsString("EOF"));
    }

    @Test
    public void testServerNoCloseHandshake() throws Exception
    {
        // Set client timeout
        final int timeout = 1000;
        client.setMaxIdleTimeout(timeout);

        // Client connects
        CloseTrackingSocket clientSocket = new CloseTrackingSocket();
        Future<Session> clientConnectFuture = client.connect(clientSocket,server.getWsUri());

        // Server accepts connect
        ServerConnection serverConn = server.accept();
        serverConn.upgrade();

        // client confirms connection via echo
        confirmConnection(clientSocket,clientConnectFuture,serverConn);

        // client sends close frame
        final String origCloseReason = "Normal Close";
        clientSocket.getSession().close(StatusCode.NORMAL,origCloseReason);

        // server receives close frame
        confirmServerReceivedCloseFrame(serverConn,StatusCode.NORMAL,is(origCloseReason));

        // client should not have received close message (yet)
        clientSocket.assertNoCloseEvent();

        // server never sends close frame handshake
        // server sits idle

        // client idle timeout triggers close event on client ws-endpoint
        clientSocket.assertReceivedCloseEvent(timeout,is(StatusCode.SHUTDOWN),containsString("Timeout"));
    }

    @Test
    public void testStopLifecycle() throws Exception
    {
        // Set client timeout
        final int timeout = 1000;
        client.setMaxIdleTimeout(timeout);

        int clientCount = 3;
        CloseTrackingSocket clientSockets[] = new CloseTrackingSocket[clientCount];
        ServerConnection serverConns[] = new ServerConnection[clientCount];

        // Connect Multiple Clients
        for (int i = 0; i < clientCount; i++)
        {
            // Client Request Upgrade
            clientSockets[i] = new CloseTrackingSocket();
            Future<Session> clientConnectFuture = client.connect(clientSockets[i],server.getWsUri());

            // Server accepts connection
            serverConns[i] = server.accept();
            serverConns[i].upgrade();

            // client confirms connection via echo
            confirmConnection(clientSockets[i],clientConnectFuture,serverConns[i]);
        }

        // client lifecycle stop
        client.stop();

        // clients send close frames (code 1001, shutdown)
        for (int i = 0; i < clientCount; i++)
        {
            // server receives close frame
            confirmServerReceivedCloseFrame(serverConns[i],StatusCode.SHUTDOWN,containsString("Shutdown"));
        }

        // clients disconnect
        for (int i = 0; i < clientCount; i++)
        {
            clientSockets[i].assertReceivedCloseEvent(timeout,is(StatusCode.SHUTDOWN),containsString("Shutdown"));
        }
    }

    @Test
    public void testWriteException() throws Exception
    {
        // Set client timeout
        final int timeout = 1000;
        client.setMaxIdleTimeout(timeout);

        // Client connects
        CloseTrackingSocket clientSocket = new CloseTrackingSocket();
        Future<Session> clientConnectFuture = client.connect(clientSocket,server.getWsUri());

        // Server accepts connect
        ServerConnection serverConn = server.accept();
        serverConn.upgrade();

        // client confirms connection via echo
        confirmConnection(clientSocket,clientConnectFuture,serverConn);

        // setup client endpoint for write failure (test only)
        EndPoint endp = clientSocket.getEndPoint();
        endp.shutdownOutput();

        // client enqueue close frame
        // client write failure
        final String origCloseReason = "Normal Close";
        clientSocket.getSession().close(StatusCode.NORMAL,origCloseReason);

        clientSocket.assertReceivedError(EofException.class,null);

        // client triggers close event on client ws-endpoint
        // assert - close code==1006 (abnormal)
        // assert - close reason message contains (write failure)
        clientSocket.assertReceivedCloseEvent(timeout,is(StatusCode.ABNORMAL),containsString("EOF"));
    }
}