import java.net.Socket;
import java.net.ServerSocket;
import java.io.*;
import javax.naming.*;
import javax.naming.directory.*;
import java.security.Permission;
import java.util.Hashtable;
import java.util.concurrent.*;

/**
 * @test
 * @compile TestDnsProvider.java
 * @run main/othervm LdapDnsProviderTest
 * @modules java.naming/com.sun.jndi.ldap
 * @bug 8160768
 * @summary ctx provider tests for ldap
 */

class DNSSecurityManager extends SecurityManager {

    private boolean perm = false;

    public void allowSetFactory() {
        perm = true;
    }

    @Override
    public void checkPermission(Permission p) {
        if (p.getName().equals("setFactory") && !perm) {
            throw new SecurityException(p.getName());
        }
    }
}

class ProviderTest implements Callable {

    private final String expected;
    private final ScheduledExecutorService killSwitchPool;
    private final LdapTestServer server;
    private final Hashtable env = new Hashtable(11);
    private final int HANGING_TEST_TIMEOUT = 1000;

    public ProviderTest(LdapTestServer server,
                        ScheduledExecutorService killSwitchPool,
                        String expected)
            throws IOException
    {
        this.server = server;
        this.killSwitchPool = killSwitchPool;
        this.expected = expected;
        env.put(Context.INITIAL_CONTEXT_FACTORY, "com.sun.jndi.ldap.LdapCtxFactory");
    }

    public ProviderTest(LdapTestServer server,
                        ScheduledExecutorService killSwitchPool,
                        String expected,
                        String dnsProvider)
            throws IOException
    {
        this(server, killSwitchPool, expected);
        env.put("com.sun.jndi.ldap.dnsProvider", dnsProvider);
    }

    boolean shutItDown(InitialContext ctx) {
        try {
            if (ctx != null) ctx.close();
            return true;
        } catch (NamingException ex) {
            return false;
        }
    }

    public Boolean call() {
        boolean passed;
        InitialContext ctx = null;
        ScheduledFuture killer = null;
        // String providerUrl = "ldap://localhost:" + server.getLocalPort();
        String providerUrl = "ldap:///dc=example,dc=com";

        try {
            while(!server.accepting()) {
                Thread.sleep(200); // allow the server to start up
            }
            Thread.sleep(200); // to be sure

            // if this is a hanging test, scheduled a thread to
            // interrupt after a certain time
            if (killSwitchPool != null) {
                final Thread current = Thread.currentThread();
                killer = killSwitchPool.schedule(
                        (Callable<Void>) () -> {
                            current.interrupt();
                            return null;
                        }, HANGING_TEST_TIMEOUT, TimeUnit.MILLISECONDS);
            }

            env.put(Context.PROVIDER_URL, providerUrl);

            try {
                ctx = new InitialDirContext(env);
                SearchControls scl = new SearchControls();
                scl.setSearchScope(SearchControls.SUBTREE_SCOPE);
                ((InitialDirContext)ctx).search(
                        "ou=People,o=Test", "(objectClass=*)", scl);
                throw new RuntimeException("Search should not complete");
            } catch (NamingException e) {
                System.err.println(e);
                passed = e.toString().indexOf(expected) > -1;
            } finally {
                if (killer != null && !killer.isDone()) {
                    killer.cancel(true);
                }
                shutItDown(ctx);
                server.close();
            }
            return passed;
        } catch (IOException|InterruptedException e) {
            throw new RuntimeException(e);
        }
    }
}


class LdapTestServer extends Thread {
    ServerSocket serverSock;
    boolean accepting = false;
    private byte[] bindResponse = {
        0x30, 0x0C, 0x02, 0x01, 0x01, 0x61, 0x07, 0x0A,
        0x01, 0x00, 0x04, 0x00, 0x04, 0x00
    };

    public LdapTestServer() throws IOException {
        this.serverSock = new ServerSocket(0);
        start();
    }

    public int getLocalPort() {
        return serverSock.getLocalPort();
    }

    public boolean accepting() {
        return accepting;
    }

    public void close() throws IOException {
        serverSock.close();
    }

    public void run() {
        try {
            accepting = true;
            Socket socket = serverSock.accept();
            InputStream in = socket.getInputStream();
            OutputStream out = socket.getOutputStream();

            // Read the LDAP BindRequest
            while (in.read() != -1) {
                in.skip(in.available());
                break;
            }

            // Write an LDAP BindResponse
            out.write(bindResponse);
            out.flush();
        } catch (IOException e) {
            // ignore
        }
    }
}

public class LdapDnsProviderTest {

    private static final ExecutorService testPool =
        Executors.newFixedThreadPool(1);
    private static final ScheduledExecutorService killSwitchPool =
        Executors.newScheduledThreadPool(1);
    private static DNSSecurityManager sm = new DNSSecurityManager();

    public static void main(String[] args) throws Exception {
        LdapTestServer server = new LdapTestServer();
        try {
            // TestDnsProvider, no SecurityManager
            runTest(server, "yupyupyup:389", TestDnsProvider.class.getName());
            // DefaultDnsProvider, no SecurityManager
            runTest(server, "localhost:389", "com.sun.jndi.ldap.DefaultLdapDnsProvider");
            // no (default) dnsProvider, no SecurityManager
            runTest(server, "localhost:389", "");
            // no dnsProvider context property, no SecurityManager
            runTest(server, "localhost:389");

            System.setSecurityManager(sm);

            // TestDnsProvider with SecurityManager and no permissions
            runTest(server,
                    "setFactory",
                    TestDnsProvider.class.getName());

            sm.allowSetFactory();

            // TestDnsProvider with SecurityManager and correct permissions
            runTest(server,
                    "yupyupyup:389",
                    TestDnsProvider.class.getName());
        } finally {
            LdapDnsProviderTest.killSwitchPool.shutdown();
            LdapDnsProviderTest.testPool.shutdown();
        }
    }

    private static boolean runTest(LdapTestServer server,
                                   String expected)
            throws IOException, InterruptedException, ExecutionException
    {
        Future test = testPool.submit(
                new ProviderTest(
                        server, killSwitchPool, expected));

        while (!test.isDone()) {
            if ((Boolean) test.get()) {
                return true;
            }
        }
        throw new AssertionError("FAILED: " + expected);
    }

    private static boolean runTest(LdapTestServer server,
                                   String expected,
                                   String dnsProvider)
            throws IOException, InterruptedException, ExecutionException
    {
        Future test = testPool.submit(
                new ProviderTest(
                        server, killSwitchPool, expected, dnsProvider));

        while (!test.isDone()) {
            if ((Boolean) test.get()) {
                return true;
            }
        }
        throw new AssertionError("FAILED: " + expected + ", " + dnsProvider);
    }

}