/*
 * Copyright (c) 2016, 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
 * @bug 8129566
 * @modules java.httpclient
 * @summary Check permissions for WebSocket APIs
 * @run main/othervm WSPermissionTest
 */

import java.net.InetSocketAddress;
import java.net.NetPermission;
import java.net.SocketPermission;
import java.net.URI;
import java.net.URLPermission;
import java.net.http.WebSocket;
import java.net.http.WebSocket.Listener;
import java.nio.channels.SocketChannel;
import java.security.AccessControlContext;
import java.security.AccessControlException;
import java.security.AccessController;
import java.security.CodeSource;
import java.security.Permission;
import java.security.PermissionCollection;
import java.security.Permissions;
import java.security.Policy;
import java.security.PrivilegedExceptionAction;
import java.security.ProtectionDomain;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.function.Consumer;

public class WSPermissionTest {

    public static void setupSecurityManager() throws Exception {
        // All permissions, a specific ACC will be used to when testing
        // with a reduced permission set.
        Policy.setPolicy(new Policy() {
            final PermissionCollection perms = new Permissions();

            {
                perms.add(new java.security.AllPermission());
            }

            public PermissionCollection getPermissions(
                    ProtectionDomain domain) {
                return perms;
            }

            public PermissionCollection getPermissions(CodeSource codesource) {
                return perms;
            }

            public boolean implies(ProtectionDomain domain, Permission perm) {
                return perms.implies(perm);
            }
        });
        System.setSecurityManager(new SecurityManager());
    }

    static final AccessControlContext RESTRICTED_ACC = getAccessControlContext();

    private static AccessControlContext getAccessControlContext(
            Permission... ps) {
        Permissions perms = new Permissions();
        for (Permission p : ps) {
            perms.add(p);
        }
        /*
         * Create an AccessControlContext that consist a single protection
         * domain with only the permissions calculated above
         */
        ProtectionDomain pd = new ProtectionDomain(null, perms);
        return new AccessControlContext(new ProtectionDomain[] { pd });
    }

    public static void main(String[] args) throws Exception {
        setupSecurityManager();
        defaultHttpClientTest();
        urlPermissionTest();
    }

    static void defaultHttpClientTest() throws Exception {
        AccessControlContext acc = getAccessControlContext(
                new NetPermission("getDefaultHttpClient"));
        // negative
        checkPermission((uri) -> {
            try {
                AccessController
                        .doPrivileged((PrivilegedExceptionAction<Void>) () -> {
                            WebSocket.newBuilder(uri, new Listener() {
                            });
                            return null;
                        }, RESTRICTED_ACC);
                throw new RuntimeException(
                        "It didn't throw ExecutionException!");
            } catch (SecurityException e) {
                if (!(e instanceof AccessControlException)) {
                    throw new RuntimeException(
                            "Expected cause is SecurityException, but actually is "
                                    + e.getClass().getName());
                }
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        });

        // Positive
        checkPermission((uri) -> {
            try {
                AccessController
                        .doPrivileged((PrivilegedExceptionAction<Void>) () -> {
                            WebSocket.newBuilder(uri, new Listener() {
                            });
                            return null;
                        }, acc);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        });

    }

    static void urlPermissionTest() throws Exception {

        AccessControlContext restrictedAcc = getAccessControlContext(
                new RuntimePermission("enableContextClassLoaderOverride"),
                new NetPermission("getDefaultHttpClient"));

        // Negative
        checkPermission((uri) -> {
            try {
                AccessController
                        .doPrivileged((PrivilegedExceptionAction<Void>) () -> {
                            try {
                                WebSocket.newBuilder(uri, new Listener() {
                                }).buildAsync().get();
                                throw new RuntimeException(
                                        "It didn't throw ExecutionException!");
                            } catch (ExecutionException e) {
                                if (!(e.getCause() instanceof AccessControlException)) {
                                    throw new RuntimeException(
                                            "Expected cause is AccessControlException, but actually is "
                                                    + e.getCause().getClass()
                                                            .getName());
                                }
                            }
                            return null;
                        }, restrictedAcc);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        });

        // Positive
        checkPermission((uri) -> {
            try {
                System.out.println("uri:" + uri.toString());
                AccessControlContext acc = getAccessControlContext(
                        //not documented
                        new RuntimePermission("enableContextClassLoaderOverride"),
                        new NetPermission("getDefaultHttpClient"),
                        new URLPermission(
                                "http://" + uri.getHost() + ":" + uri.getPort() + uri.getPath(),
                                "GET:Sec-WebSocket-Key,Sec-WebSocket-Version")
                        //SocketPermission is not documented in javadoc
                        // Even granted, it still fail with ACE
                        , new SocketPermission(
                                uri.getHost() + ":" + uri.getPort(),
                                "connect,resolve")
                        );
                AccessController
                        .doPrivileged((PrivilegedExceptionAction<Void>) () -> {
                            WebSocket.newBuilder(uri, new Listener() {
                            }).buildAsync().get();
                            return null;
                        }, acc);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        });

    }

    /**
     * Check actions which need specific permissions. Automatically closes
     * everything after the check has been performed
     * 
     * @param c
     *            Consumer which performs actions to be tested.
     * @param clz
     *            expected exception class. Null if no exception is expected.
     */
    private static void checkPermission(Consumer<URI> c) throws Exception {
        HandshakePhase HandshakePhase = new HandshakePhase(
                new InetSocketAddress("127.0.0.1", 0));
        URI serverURI = HandshakePhase.getURI();
        CompletableFuture<SocketChannel> cfc = HandshakePhase.afterHandshake();

        try {
            c.accept(serverURI);
        } finally {
            try {
                SocketChannel now = cfc.getNow(null);
                if (now != null) {
                    now.close();
                }
            } catch (Throwable ignored) {
            }
        }
    }
}