import java.io.IOException; 
import java.net.Authenticator; 
import java.net.PasswordAuthentication; 
import java.net.URI; 
import java.net.http.HttpClient; 
import java.net.http.HttpConnectTimeoutException; 
import java.net.http.WebSocket; 
import java.nio.ByteBuffer; 
import java.util.Base64; 
import java.util.Collection; 
import java.util.List; 
import java.util.concurrent.CompletionStage; 
import java.util.concurrent.Executor; 
import java.util.concurrent.Future; 

import javax.servlet.ServletException; 
import javax.servlet.http.HttpServlet; 
import javax.servlet.http.HttpServletRequest; 
import javax.servlet.http.HttpServletResponse; 

import org.bouncycastle.util.encoders.Base64Encoder; 
import org.eclipse.jetty.io.ByteBufferPool; 
import org.eclipse.jetty.io.EndPoint; 
import org.eclipse.jetty.security.ConstraintMapping; 
import org.eclipse.jetty.security.ConstraintSecurityHandler; 
import org.eclipse.jetty.security.HashLoginService; 
import org.eclipse.jetty.security.RoleInfo; 
import org.eclipse.jetty.security.SecurityHandler; 
import org.eclipse.jetty.security.UserStore; 
import org.eclipse.jetty.security.authentication.BasicAuthenticator; 
import org.eclipse.jetty.server.ConnectionFactory; 
import org.eclipse.jetty.server.Connector; 
import org.eclipse.jetty.server.HttpConfiguration; 
import org.eclipse.jetty.server.Request; 
import org.eclipse.jetty.server.Response; 
import org.eclipse.jetty.server.Server; 
import org.eclipse.jetty.server.ServerConnector; 
import org.eclipse.jetty.server.UserIdentity; 
import org.eclipse.jetty.servlet.ServletContextHandler; 
import org.eclipse.jetty.servlet.ServletHolder; 
import org.eclipse.jetty.util.component.Container; 
import org.eclipse.jetty.util.component.LifeCycle; 
import org.eclipse.jetty.util.log.Log; 
import org.eclipse.jetty.util.log.StdErrLog; 
import org.eclipse.jetty.util.security.Constraint; 
import org.eclipse.jetty.util.security.Credential; 
import org.eclipse.jetty.util.thread.Scheduler; 
import org.eclipse.jetty.websocket.api.Session; 
import org.eclipse.jetty.websocket.api.WebSocketAdapter; 
import org.eclipse.jetty.websocket.servlet.WebSocketServlet; 
import org.eclipse.jetty.websocket.servlet.WebSocketServletFactory; 
import org.junit.After; 
import org.junit.Before; 
import org.junit.Test; 

public class WebSocketAuthenticationTest extends WebSocketAdapter { 
    public class TestServlet extends WebSocketServlet { 
        @Override 
        public void configure(WebSocketServletFactory factory) { 
            factory.register(WebSocketAuthenticationTest.class); 
        } 
    } 

    @Override 
    public void onWebSocketConnect(Session sess) { 
        super.onWebSocketConnect(sess); 
    } 

    private String username = "MyUserName"; 
    private String password = "MyPassword"; 

    private Server server; 

    private WebSocket.Listener websocketListener = new WebSocket.Listener() { 
        @Override 
        public void onOpen(WebSocket webSocket) { 
        } 

        @Override 
        public CompletionStage<?> onText(WebSocket webSocket, CharSequence data, boolean last) { 
            return null; 
        } 

        @Override 
        public CompletionStage<?> onBinary(WebSocket webSocket, ByteBuffer data, boolean last) { 
            return null; 
        } 

        @Override 
        public CompletionStage<?> onPing(WebSocket webSocket, ByteBuffer message) { 
            return null; 
        } 

        @Override 
        public CompletionStage<?> onPong(WebSocket webSocket, ByteBuffer message) { 
            return null; 
        } 

        @Override 
        public CompletionStage<?> onClose(WebSocket webSocket, int statusCode, String reason) { 
            return null; 
        } 

        @Override 
        public void onError(WebSocket webSocket, Throwable error) { 
        } 
    }; 

    @Before 
    public void setUp() throws Exception { 
        StdErrLog logger = new StdErrLog(); 
        logger.setDebugEnabled(true); 
        Log.setLog(logger); 

        server = new Server(); 
        ServerConnector connector = new ServerConnector(server); 
        connector.setPort(8080); 
        server.addConnector(connector); 
        ServletContextHandler ctx = new ServletContextHandler(ServletContextHandler.SESSIONS); 
        ctx.setContextPath("/"); 
        ctx.addServlet(new ServletHolder("ws-events", new TestServlet()), "/endpoint/*"); 
        ConstraintSecurityHandler securityHandler = new ConstraintSecurityHandler(); 
        HashLoginService loginService = new HashLoginService(); 
        UserStore us = new UserStore(); 
        us.addUser(username, Credential.getCredential(password), new String[]{"user"}); 
        Constraint constraint = new Constraint(); 
        constraint.setName(Constraint.__BASIC_AUTH); 
        constraint.setRoles(new String[]{"user"}); 
        constraint.setAuthenticate(true); 
        loginService.setUserStore(us); 
        securityHandler.setLoginService(loginService); 
        ConstraintMapping constraintMapping = new ConstraintMapping(); 
        constraintMapping.setConstraint(constraint); 
        constraintMapping.setPathSpec("/endpoint/*"); 
        securityHandler.addConstraintMapping(constraintMapping); 
        ctx.setSecurityHandler(securityHandler); 
        server.setHandler(ctx); 
        server.start(); 
    } 

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

    @Test 
    public void connectToWithAuthenticatorServer() throws Exception { 
        final HttpClient httpClient = HttpClient.newBuilder().authenticator(new Authenticator() { 
            @Override 
            protected PasswordAuthentication getPasswordAuthentication() { 
                return new PasswordAuthentication(username, password.toCharArray()); 
            } 
        }).build(); 
        httpClient.newWebSocketBuilder().buildAsync(new URI("ws://localhost:8080/endpoint/"), websocketListener).get(); 
    } 

    @Test 
    public void connectToWithHeaderServer() throws Exception { 
        final HttpClient httpClient = HttpClient.newBuilder().build(); 
        httpClient.newWebSocketBuilder().header("Authorization", "Basic " + Base64.getEncoder().encodeToString((username+ ":" + password).getBytes())) 
                .buildAsync(new URI("ws://localhost:8080/endpoint/"), websocketListener).get(); 
    } 

}
