package com.jwrp.javawebsocketreverseproxy;

import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.client.standard.StandardWebSocketClient;

import java.io.IOException;
import java.util.concurrent.TimeUnit;

/**
 * Represents a 'hop' in the proxying chain, establishes a 'client' to
 * communicate with the next server, with a {@link WebSocketProxyClientHandler}
 * to copy data from the 'client' to the supplied 'server' session.
 */
public class NextHop {

    private final WebSocketSession webSocketClientSession;

    public NextHop(WebSocketSession webSocketServerSession) {
        webSocketClientSession = createWebSocketClientSession(webSocketServerSession);
    }

    private WebSocketSession createWebSocketClientSession(WebSocketSession webSocketServerSession) {
        try {
            return new StandardWebSocketClient()
                    .doHandshake(new WebSocketProxyClientHandler(webSocketServerSession), "ws://localhost:9999")
                    .get(1000, TimeUnit.MILLISECONDS);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public void sendMessageToNextHop(WebSocketMessage<?> webSocketMessage) throws IOException {
        webSocketClientSession.sendMessage(webSocketMessage);
    }
}

@@ -38,16 +38,348 @@
 proxy;%0A%0A
+import org.slf4j.Logger;%0Aimport org.slf4j.LoggerFactory;%0Aimport org.springframework.http.HttpHeaders;%0Aimport org.springframework.security.oauth2.provider.OAuth2Authentication;%0Aimport org.springframework.security.oauth2.provider.authentication.OAuth2AuthenticationDetails;%0Aimport org.springframework.web.socket.WebSocketHttpHeaders;%0A
 import o
@@ -586,24 +586,86 @@
 OException;%0A
+import java.security.Principal;%0Aimport java.util.Collections;%0A
 import java.
@@ -1000,16 +1000,92 @@
 Session;
+%0A    private final Logger logger = LoggerFactory.getLogger(this.getClass());
 %0A%0A    pu
@@ -1230,24 +1230,1184 @@
 on);%0A    %7D%0A%0A
+    private WebSocketHttpHeaders getWebSocketHttpHeaders(final WebSocketSession userAgentSession) %7B%0A        WebSocketHttpHeaders headers = new WebSocketHttpHeaders();%0A        Principal principal = userAgentSession.getPrincipal();%0A        if (principal != null && OAuth2Authentication.class.isAssignableFrom(principal.getClass())) %7B%0A            OAuth2Authentication oAuth2Authentication = (OAuth2Authentication) principal;%0A            OAuth2AuthenticationDetails details = (OAuth2AuthenticationDetails) oAuth2Authentication.getDetails();%0A            String accessToken = details.getTokenValue();%0A            headers.put(HttpHeaders.AUTHORIZATION, Collections.singletonList(%22Bearer %22 + accessToken));%0A            if(logger.isDebugEnabled()) %7B%0A                logger.debug(%22Added Oauth2 bearer token authentication header for user %22 +%0A                        principal.getName() + %22 to web sockets http headers%22);%0A            %7D%0A        %7D%0A        else %7B%0A            if(logger.isDebugEnabled()) %7B%0A                logger.debug(%22Skipped adding basic authentication header since user session principal is null%22);%0A            %7D%0A        %7D%0A        return headers;%0A    %7D%0A%0A
     private 
@@ -2509,16 +2509,108 @@
   try %7B%0A
+            WebSocketHttpHeaders headers = getWebSocketHttpHeaders(webSocketServerSession);%0A
         
@@ -2740,16 +2740,33 @@
 ssion), 
+headers, new URI(
 %22ws://lo
@@ -2779,16 +2779,17 @@
 t:9999%22)
+)
 %0A       

send authorization to NextNop if the proxy layer provided