//
//  ========================================================================
//  Copyright (c) 1995-2020 Mort Bay Consulting Pty Ltd and others.
//  ------------------------------------------------------------------------
//  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.client;

import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocket;

import org.eclipse.jetty.http.HttpScheme;
import org.eclipse.jetty.util.ssl.SslContextFactory;
import org.eclipse.jetty.util.thread.QueuedThreadPool;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

public class Socks4ProxyTest
{
    private ServerSocketChannel server;
    private HttpClient client;

    @BeforeEach
    public void prepare() throws Exception
    {
        server = ServerSocketChannel.open();
        server.bind(new InetSocketAddress("localhost", 0));

        QueuedThreadPool clientThreads = new QueuedThreadPool();
        clientThreads.setName("client");
        client = new HttpClient(new SslContextFactory.Client());
        client.setExecutor(clientThreads);
        client.start();
    }

    @AfterEach
    public void dispose() throws Exception
    {
        client.stop();
        server.close();
    }

    @Test
    public void testSocks4Proxy() throws Exception
    {
        int proxyPort = server.socket().getLocalPort();
        client.getProxyConfiguration().getProxies().add(new Socks4Proxy("localhost", proxyPort));

        CountDownLatch latch = new CountDownLatch(1);

        byte ip1 = 127;
        byte ip2 = 0;
        byte ip3 = 0;
        byte ip4 = 13;
        String serverHost = ip1 + "." + ip2 + "." + ip3 + "." + ip4;
        int serverPort = proxyPort + 1; // Any port will do
        String method = "GET";
        String path = "/path";
        client.newRequest(serverHost, serverPort)
            .method(method)
            .path(path)
            .timeout(5, TimeUnit.SECONDS)
            .send(result ->
            {
                if (result.isSucceeded())
                    latch.countDown();
            });

        try (SocketChannel channel = server.accept())
        {
            int socks4MessageLength = 9;
            ByteBuffer buffer = ByteBuffer.allocate(socks4MessageLength);
            int read = channel.read(buffer);
            assertEquals(socks4MessageLength, read);
            assertEquals(4, buffer.get(0) & 0xFF);
            assertEquals(1, buffer.get(1) & 0xFF);
            assertEquals(serverPort, buffer.getShort(2) & 0xFFFF);
            assertEquals(ip1, buffer.get(4) & 0xFF);
            assertEquals(ip2, buffer.get(5) & 0xFF);
            assertEquals(ip3, buffer.get(6) & 0xFF);
            assertEquals(ip4, buffer.get(7) & 0xFF);
            assertEquals(0, buffer.get(8) & 0xFF);

            // Socks4 response.
            channel.write(ByteBuffer.wrap(new byte[]{0, 0x5A, 0, 0, 0, 0, 0, 0}));

            buffer = ByteBuffer.allocate(method.length() + 1 + path.length());
            read = channel.read(buffer);
            assertEquals(buffer.capacity(), read);
            buffer.flip();
            assertEquals(method + " " + path, StandardCharsets.UTF_8.decode(buffer).toString());

            // Response
            String response =
                "HTTP/1.1 200 OK\r\n" +
                    "Content-Length: 0\r\n" +
                    "Connection: close\r\n" +
                    "\r\n";
            channel.write(ByteBuffer.wrap(response.getBytes(StandardCharsets.UTF_8)));

            assertTrue(latch.await(5, TimeUnit.SECONDS));
        }
    }

    @Test
    public void testSocks4ProxyWithSplitResponse() throws Exception
    {
        int proxyPort = server.socket().getLocalPort();
        client.getProxyConfiguration().getProxies().add(new Socks4Proxy("localhost", proxyPort));

        CountDownLatch latch = new CountDownLatch(1);

        String serverHost = "127.0.0.13"; // Test expects an IP address.
        int serverPort = proxyPort + 1; // Any port will do
        String method = "GET";
        client.newRequest(serverHost, serverPort)
            .method(method)
            .path("/path")
            .timeout(5, TimeUnit.SECONDS)
            .send(result ->
            {
                if (result.isSucceeded())
                    latch.countDown();
                else
                    result.getFailure().printStackTrace();
            });

        try (SocketChannel channel = server.accept())
        {
            int socks4MessageLength = 9;
            ByteBuffer buffer = ByteBuffer.allocate(socks4MessageLength);
            int read = channel.read(buffer);
            assertEquals(socks4MessageLength, read);

            // Socks4 response, with split bytes.
            byte[] chunk1 = new byte[]{0, 0x5A, 0};
            byte[] chunk2 = new byte[]{0, 0, 0, 0, 0};
            channel.write(ByteBuffer.wrap(chunk1));

            // Wait before sending the second chunk.
            Thread.sleep(1000);

            channel.write(ByteBuffer.wrap(chunk2));

            buffer = ByteBuffer.allocate(method.length());
            read = channel.read(buffer);
            assertEquals(buffer.capacity(), read);
            buffer.flip();
            assertEquals(method, StandardCharsets.UTF_8.decode(buffer).toString());

            // Response
            String response =
                "HTTP/1.1 200 OK\r\n" +
                    "Content-Length: 0\r\n" +
                    "Connection: close\r\n" +
                    "\r\n";
            channel.write(ByteBuffer.wrap(response.getBytes(StandardCharsets.UTF_8)));

            assertTrue(latch.await(5, TimeUnit.SECONDS));
        }
    }

    @Test
    public void testSocks4ProxyWithTLSServer() throws Exception
    {
        String proxyHost = "localhost";
        int proxyPort = server.socket().getLocalPort();

        String serverHost = "127.0.0.13"; // Server host different from proxy host.
        int serverPort = proxyPort + 1; // Any port will do.

        SslContextFactory clientTLS = client.getSslContextFactory();
        clientTLS.reload(ssl ->
        {
            // The client keystore contains the trustedCertEntry for the
            // self-signed server certificate, so it acts as a truststore.
            ssl.setTrustStorePath("src/test/resources/client_keystore.p12");
            ssl.setTrustStorePassword("storepwd");
            // Disable TLS hostname verification, but
            // enable application hostname verification.
            ssl.setEndpointIdentificationAlgorithm(null);
            // The hostname must be that of the server, not of the proxy.
            ssl.setHostnameVerifier((hostname, session) -> serverHost.equals(hostname));
        });
        client.getProxyConfiguration().getProxies().add(new Socks4Proxy(proxyHost, proxyPort));

        CountDownLatch latch = new CountDownLatch(1);
        client.newRequest(serverHost, serverPort)
            .scheme(HttpScheme.HTTPS.asString())
            .path("/path")
            .send(result ->
            {
                if (result.isSucceeded())
                    latch.countDown();
                else
                    result.getFailure().printStackTrace();
            });

        try (SocketChannel channel = server.accept())
        {
            int socks4MessageLength = 9;
            ByteBuffer buffer = ByteBuffer.allocate(socks4MessageLength);
            int read = channel.read(buffer);
            assertEquals(socks4MessageLength, read);

            // Socks4 response.
            channel.write(ByteBuffer.wrap(new byte[]{0, 0x5A, 0, 0, 0, 0, 0, 0}));

            // Wrap the socket with TLS.
            SslContextFactory.Server serverTLS = new SslContextFactory.Server();
            serverTLS.setKeyStorePath("src/test/resources/keystore.p12");
            serverTLS.setKeyStorePassword("storepwd");
            serverTLS.start();
            SSLContext sslContext = serverTLS.getSslContext();
            SSLSocket sslSocket = (SSLSocket)sslContext.getSocketFactory().createSocket(channel.socket(), serverHost, serverPort, false);
            sslSocket.setUseClientMode(false);

            // Read the request.
            int crlfs = 0;
            InputStream input = sslSocket.getInputStream();
            while (true)
            {
                read = input.read();
                if (read < 0)
                    break;
                if (read == '\r' || read == '\n')
                    ++crlfs;
                else
                    crlfs = 0;
                if (crlfs == 4)
                    break;
            }

            // Send the response.
            String response =
                "HTTP/1.1 200 OK\r\n" +
                    "Content-Length: 0\r\n" +
                    "Connection: close\r\n" +
                    "\r\n";
            OutputStream output = sslSocket.getOutputStream();
            output.write(response.getBytes(StandardCharsets.UTF_8));
            output.flush();

            assertTrue(latch.await(5, TimeUnit.SECONDS));
        }
    }
}
