/*

 * Copyright (c) 2019, Oracle and/or its affiliates. All rights reserved.

 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.

 *

 * This code is free software; you can redistribute it and/or modify it

 * under the terms of the GNU General Public License version 2 only, as

 * published by the Free Software Foundation.

 *

 * This code is distributed in the hope that it will be useful, but WITHOUT

 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or

 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License

 * version 2 for more details (a copy is included in the LICENSE file that

 * accompanied this code).

 *

 * You should have received a copy of the GNU General Public License version

 * 2 along with this work; if not, write to the Free Software Foundation,

 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.

 *

 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA

 * or visit www.oracle.com if you need additional information or have any

 * questions.

 */

/**

 * @test

 * @modules java.base/java.net:+open java.base/sun.nio.ch:+open

 * @run testng/othervm CustomSocketImpls

 * @run testng/othervm -Djdk.net.usePlainSocketImpl CustomSocketImpls

 */

import java.io.FileDescriptor;

import java.io.IOException;

import java.io.InputStream;

import java.io.OutputStream;

import java.lang.reflect.Field;

import java.net.InetAddress;

import java.net.InetSocketAddress;

import java.net.Proxy;

import java.net.ServerSocket;

import java.net.Socket;

import java.net.SocketAddress;

import java.net.SocketImpl;

import java.net.SocketImplFactory;

import java.nio.channels.ServerSocketChannel;

import java.nio.channels.SocketChannel;

import java.util.function.Consumer;

import org.testng.annotations.Test;

import static org.testng.Assert.*;

@Test

public class CustomSocketImpls {

    /**

     * Test ServerSocket is created with the expected SocketImpl.

     */

    public void testServerSocketSocketImpl() throws Exception {

        try (ServerSocket ss = new ServerSocket()) {

            SocketImpl si = getSocketImpl(ss);

            assertTrue(isPlatformSocketImpl(si));

        }

        SocketImpl impl = new CustomSocketImpl(true);

        try (ServerSocket ss = new ServerSocket(impl){}) {

            SocketImpl si = getSocketImpl(ss);

            assertTrue(si instanceof CustomSocketImpl);

        }

        setSocketFactory(() -> new CustomSocketImpl(true));

        try (ServerSocket ss = new ServerSocket()) {

            SocketImpl si = getSocketImpl(ss);

            assertTrue(si instanceof CustomSocketImpl);

        } finally {

            setSocketFactory(null);

        }

    }

    /**

     * Test Socket is created with the expected SocketImpl.

     */

    public void tesSocketSocketImpl() throws Exception {

        try (Socket s = new Socket()) {

            SocketImpl si = getSocketImpl(s);

            assertTrue(isSocksSocketImpl(si));

            SocketImpl delegate = getDelegate(si);

            assertTrue(isPlatformSocketImpl(delegate));

        }

        try (Socket s = new Socket((SocketImpl) null){}) {

            assertTrue(getSocketImpl(s) == null);

        }

        try (Socket s = new Socket(Proxy.NO_PROXY)) {

            SocketImpl si = getSocketImpl(s);

            assertTrue(isPlatformSocketImpl(si));

        }

        var address = new InetSocketAddress("127.0.0.1", 1000);

        var socksProxy = new Proxy(Proxy.Type.SOCKS, address);

        try (Socket s = new Socket(socksProxy)) {

            SocketImpl si = getSocketImpl(s);

            assertTrue(isSocksSocketImpl(si));

            SocketImpl delegate = getDelegate(si);

            assertTrue(isPlatformSocketImpl(delegate));

        }

        var httpProxy = new Proxy(Proxy.Type.HTTP, address);

        try (Socket s = new Socket(httpProxy)) {

            SocketImpl si = getSocketImpl(s);

            assertTrue(isHttpConnectSocketImpl(si));

            SocketImpl delegate = getDelegate(si);

            assertTrue(isPlatformSocketImpl(delegate));

        }

        try (Socket s = new Socket(new CustomSocketImpl(false)){}) {

            SocketImpl si = getSocketImpl(s);

            assertTrue(si instanceof CustomSocketImpl);

        }

        setSocketImplFactory(() -> new CustomSocketImpl(false));

        try (Socket s = new Socket()) {

            SocketImpl si = getSocketImpl(s);

            assertTrue(si instanceof CustomSocketImpl);

        } finally {

            setSocketImplFactory(null);

        }

    }

    /**

     * ServerSocket using default SocketImpl. Test accept returning a Socket

     * that initially doesn't have a SocketImpl.

     */

    public void test1() throws Exception {

        var socket = new Socket((SocketImpl) null) { };

        assertTrue(getSocketImpl(socket) == null);

        testAccept(socket, s -> {

            assertTrue(s == socket);

            SocketImpl si = getSocketImpl(s);

            assertTrue(isPlatformSocketImpl(si));

            checkFields(si);

        });

    }

    /**

     * ServerSocket using default SocketImpl. Test accept returning a Socket

     * that has an existing default SocketImpl.

     */

    public void test2() throws Exception {

        var socket = new Socket();

        SocketImpl si = getSocketImpl(socket);

        assertTrue(isSocksSocketImpl(si));

        SocketImpl delegate = getDelegate(si);

        assertTrue(isPlatformSocketImpl(delegate));

        testAccept(socket, s -> {

            assertTrue(s == socket);

            assertTrue(getSocketImpl(s) == si);

            assertTrue(getDelegate(si) == delegate);

            checkFields(delegate);

        });

    }

    /**

     * ServerSocket using default SocketImpl. A SocketImplFactory is set to

     * return a custom SocketImpl. Test accept returning a Socket that initially

     * doesn't have a SocketImpl.

     */

    public void test3() throws Exception {

        var socket = new Socket((SocketImpl) null) { };

        assertTrue(getSocketImpl(socket) == null);

        testAccept(socket, () -> new CustomSocketImpl(false), s -> {

            assertTrue(s == socket);

            SocketImpl si = getSocketImpl(s);

            assertTrue(si instanceof CustomSocketImpl);

            checkFields(si);

        });

    }

    /**

     * ServerSocket using default SocketImpl. Test accept returning a Socket

     * that has an existing custom SocketImpl.

     */

    public void test4() throws Exception {

        SocketImpl si = new CustomSocketImpl(false);

        Socket socket = new Socket(si) { };

        assertTrue(getSocketImpl(socket) == si);

        testAccept(socket, s -> {

            assertTrue(s == socket);

            assertTrue(getSocketImpl(s) == si);

            checkFields(si);

        });

    }

    /**

     * ServerSocket using a custom SocketImpl. Test accept returning a Socket

     * that initially doesn't have a SocketImpl.

     */

    public void test5() throws Exception {

        SocketImpl impl = new CustomSocketImpl(true);

        Socket socket = new Socket((SocketImpl) null) { };

        assertTrue(getSocketImpl(socket) == null);

        testAccept(impl, socket, s -> {

            assertTrue(s == socket);

            SocketImpl si = getSocketImpl(s);

            assertTrue(isPlatformSocketImpl(si));

            checkFields(si);

        });

    }

    /**

     * ServerSocket using a custom SocketImpl. Test accept returning a Socket

     * that has an existing default SocketImpl.

     */

    public void test6() throws Exception {

        SocketImpl impl = new CustomSocketImpl(true);

        var socket = new Socket();

        SocketImpl si = getSocketImpl(socket);

        assertTrue(isSocksSocketImpl(si));

        SocketImpl delegate = getDelegate(si);

        assertTrue(isPlatformSocketImpl(delegate));

        testAccept(impl, socket, s -> {

            assertTrue(s == socket);

            assertTrue(getSocketImpl(s) == si);

            assertTrue(getDelegate(si) == delegate);

            checkFields(delegate);

        });

    }

    /**

     * ServerSocket using a custom SocketImpl. A SocketImplFactory is set to

     * return a custom SocketImpl. Test accept returning a Socket that initially

     * doesn't have a SocketImpl.

     */

    public void test7() throws Exception {

        var socket = new Socket((SocketImpl) null) { };

        assertTrue(getSocketImpl(socket) == null);

        Socket s1 = null;

        Socket s2 = null;

        try (ServerSocket ss = prepareToAccept(new CustomSocketImpl(true), socket)) {

            s1 = new Socket(ss.getInetAddress(), ss.getLocalPort());

            setSocketImplFactory(() -> new CustomSocketImpl(false));

            try {

                s2 = ss.accept();

                SocketImpl si = getSocketImpl(s2);

                assertTrue(si instanceof CustomSocketImpl);

                checkFields(si);

            } finally {

                setSocketImplFactory(null);

            }

        } finally {

            if (s1 != null) s1.close();

            if (s2 != null) s1.close();

        }

    }

    /**

     * ServerSocket using a custom SocketImpl. A SocketImplFactory is set to

     * return a custom SocketImpl. Test accept returning a Socket has an existing

     * custom SocketImpl.

     */

    public void test8() throws Exception {

        SocketImpl si = new CustomSocketImpl(false);

        Socket socket = new Socket(si) { };

        assertTrue(getSocketImpl(socket) == si);

        Socket s1 = null;

        Socket s2 = null;

        try (ServerSocket ss = prepareToAccept(new CustomSocketImpl(true), socket)) {

            s1 = new Socket(ss.getInetAddress(), ss.getLocalPort());

            setSocketImplFactory(() -> new CustomSocketImpl(false));

            try {

                s2 = ss.accept();

                assertTrue(getSocketImpl(s2) == si);

                checkFields(si);

            } finally {

                setSocketImplFactory(null);

            }

        } finally {

            if (s1 != null) s1.close();

            if (s2 != null) s1.close();

        }

    }

    /**

     * Creates a ServerSocket that returns the given Socket from accept.

     * The consumer is invoked with the accepted socket.

     */

    static void testAccept(Socket socket, Consumer<Socket> consumer)

        throws IOException

    {

        Socket s1 = null;

        Socket s2 = null;

        try (ServerSocket ss = prepareToAccept(socket)) {

            s1 = new Socket(ss.getInetAddress(), ss.getLocalPort());

            s2 = ss.accept();

            consumer.accept(s2);

        } finally {

            if (s1 != null) s1.close();

            if (s2 != null) s1.close();

        }

    }

    /**

     * Creates a ServerSocket that returns the given Socket from accept. The

     * given SocketImplFactory is set during the accept and the consumer is

     * invoked when the accepted socket.

     */

    static void testAccept(Socket socket, SocketImplFactory factory, Consumer<Socket> consumer)

        throws IOException

    {

        Socket s1 = null;

        Socket s2 = null;

        try (ServerSocket ss = prepareToAccept(socket)) {

            s1 = new Socket(ss.getInetAddress(), ss.getLocalPort());

            setSocketImplFactory(factory);

            try {

                s2 = ss.accept();

            } finally {

                setSocketImplFactory(null);

            }

            consumer.accept(s2);

        } finally {

            if (s1 != null) s1.close();

            if (s2 != null) s1.close();

        }

    }

    /**

     * Creates a ServerSocket with a SocketImpl that returns the given Socket

     * from accept. The consumer is invoked with the accepted socket.

     */

    static void testAccept(SocketImpl impl, Socket socket, Consumer<Socket> consumer)

        throws IOException

    {

        Socket s1 = null;

        Socket s2 = null;

        try (ServerSocket ss = prepareToAccept(impl, socket)) {

            s1 = new Socket(ss.getInetAddress(), ss.getLocalPort());

            s2 = ss.accept();

            consumer.accept(s2);

        } finally {

            if (s1 != null) s1.close();

            if (s2 != null) s1.close();

        }

    }

    /**

     * Creates a ServerSocket that returns the given Socket from accept.

     */

    static ServerSocket prepareToAccept(Socket s) throws IOException {

        return new ServerSocket(0) {

            @Override

            public Socket accept() throws IOException {

                implAccept(s);

                return s;

            }

        };

    }

    /**

     * Creates a ServerSocket with a SocketImpl that returns the given Socket

     * from accept.

     */

    static ServerSocket prepareToAccept(SocketImpl impl, Socket s) throws IOException {

        ServerSocket ss = new ServerSocket(impl) {

            @Override

            public Socket accept() throws IOException {

                implAccept(s);

                return s;

            }

        };

        ss.bind(new InetSocketAddress(0));

        return ss;

    }

    /**

     * Returns the socket's SocketImpl

     */

    static SocketImpl getSocketImpl(Socket s) {

        try {

            Field f = Socket.class.getDeclaredField("impl");

            f.setAccessible(true);

            return (SocketImpl) f.get(s);

        } catch (Exception e) {

            throw new RuntimeException(e);

        }

    }

    /**

     * Returns the server socket's SocketImpl

     */

    static SocketImpl getSocketImpl(ServerSocket ss) {

        try {

            Field f = ServerSocket.class.getDeclaredField("impl");

            f.setAccessible(true);

            return (SocketImpl) f.get(ss);

        } catch (Exception e) {

            throw new RuntimeException(e);

        }

    }

    /**

     * Returns the SocketImpl that the given SocketImpl delegates to

     */

    static SocketImpl getDelegate(SocketImpl si) {

        try {

            Class<?> clazz = Class.forName("java.net.DelegatingSocketImpl");

            Field f = clazz.getDeclaredField("delegate");

            f.setAccessible(true);

            return (SocketImpl) f.get(si);

        } catch (Exception e) {

            throw new RuntimeException(e);

        }

    }

    /**

     * Returns the value of a SocketImpl field

     */

    static <T> T get(SocketImpl si, String name) {

        try {

            Field f = SocketImpl.class.getDeclaredField(name);

            f.setAccessible(true);

            return (T) f.get(si);

        } catch (Exception e) {

            throw new RuntimeException(e);

        }

    }

    /**

     * Sets the value of SocketImpl field

     */

    static void set(SocketImpl si, String name, Object value) {

        try {

            Field f = SocketImpl.class.getDeclaredField(name);

            f.setAccessible(true);

            f.set(si, value);

        } catch (Exception e) {

            throw new RuntimeException(e);

        }

    }

    /**

     * Returns true if the SocketImpl is a PlatformSocketImpl

     */

    static boolean isPlatformSocketImpl(SocketImpl si) {

        try {

            Class<?> clazz = Class.forName("sun.net.PlatformSocketImpl");

            return clazz.isInstance(si);

        } catch (Exception e) {

            throw new RuntimeException(e);

        }

    }

    /**

     * Returns true if the SocketImpl is a SocksSocketImpl

     */

    static boolean isSocksSocketImpl(SocketImpl si) {

        try {

            Class<?> clazz = Class.forName("java.net.SocksSocketImpl");

            return clazz.isInstance(si);

        } catch (Exception e) {

            throw new RuntimeException(e);

        }

    }

    /**

     * Returns true if the SocketImpl is a HttpConnectSocketImpl

     */

    static boolean isHttpConnectSocketImpl(SocketImpl si) {

        try {

            Class<?> clazz = Class.forName("java.net.HttpConnectSocketImpl");

            return clazz.isInstance(si);

        } catch (Exception e) {

            throw new RuntimeException(e);

        }

    }

    /**

     * Socket.setSocketImplFactory

     */

    static void setSocketImplFactory(SocketImplFactory factory) {

        try {

            Field f = Socket.class.getDeclaredField("factory");

            f.setAccessible(true);

            f.set(null, factory);

        } catch (Exception e) {

            throw new RuntimeException(e);

        }

    }

    /**

     * ServerSocket.setSocketFactory

     */

    static void setSocketFactory(SocketImplFactory factory) {

        try {

            Field f = ServerSocket.class.getDeclaredField("factory");

            f.setAccessible(true);

            f.set(null, factory);

        } catch (Exception e) {

            throw new RuntimeException(e);

        }

    }

    /**

     * Checks the 4 protected fields of a SocketImpl to make sure that they

     * have been initialized.

     */

    static void checkFields(SocketImpl si) {

        FileDescriptor fd = get(si, "fd");

        InetAddress address = get(si, "address");

        int port = get(si, "port");

        int localport = get(si, "localport");

        assertTrue(fd.valid() && address != null && port != 0 && localport != 0);

    }

    /**

     * Custom SocketImpl that is layed on a SocketChannel or ServerSocketChannel

     */

    static class CustomSocketImpl extends SocketImpl {

        private final boolean server;

        private ServerSocketChannel ssc;

        private SocketChannel sc;

        CustomSocketImpl(boolean server) {

            this.server = server;

        }

        @Override

        protected void create(boolean stream) throws IOException {

            if (server) {

                ssc = ServerSocketChannel.open();