Extract origin url from config path

G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/NitroProxyProvider.java

@@ -11,6 +11,8 @@ import gearth.protocol.connection.proxy.nitro.http.NitroHttpProxyServerCallback;
 import gearth.protocol.connection.proxy.nitro.websocket.NitroWebsocketProxy;
 
 import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
 
 public class NitroProxyProvider implements ProxyProvider, NitroHttpProxyServerCallback, StateChangeListener {
 
@@ -21,6 +23,7 @@ public class NitroProxyProvider implements ProxyProvider, NitroHttpProxyServerCa
     private final NitroWebsocketProxy nitroWebsocketProxy;
 
     private String originalWebsocketUrl;
+    private String originalOriginUrl;
 
     public NitroProxyProvider(HProxySetter proxySetter, HStateSetter stateSetter, HConnection connection) {
         this.proxySetter = proxySetter;
@@ -34,6 +37,10 @@ public class NitroProxyProvider implements ProxyProvider, NitroHttpProxyServerCa
         return originalWebsocketUrl;
     }
 
+    public String getOriginalOriginUrl() {
+        return originalOriginUrl;
+    }
+
     @Override
     public void start() throws IOException {
         connection.getStateObservable().addListener(this);
@@ -79,6 +86,8 @@ public class NitroProxyProvider implements ProxyProvider, NitroHttpProxyServerCa
     @Override
     public String replaceWebsocketServer(String configUrl, String websocketUrl) {
         originalWebsocketUrl = websocketUrl;
+        originalOriginUrl = extractOriginUrl(configUrl);
+
         return String.format("ws://127.0.0.1:%d", NitroConstants.WEBSOCKET_PORT);
     }
 
@@ -90,4 +99,15 @@ public class NitroProxyProvider implements ProxyProvider, NitroHttpProxyServerCa
             nitroHttpProxy.pause();
         }
     }
+
+    private static String extractOriginUrl(String url) {
+        try {
+            final URI uri = new URI(url);
+            return String.format("%s://%s/", uri.getScheme(), uri.getHost());
+        } catch (URISyntaxException e) {
+            e.printStackTrace();
+        }
+
+        return null;
+    }
 }

G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketClient.java

@@ -39,7 +39,7 @@ public class NitroWebsocketClient implements NitroSession {
         activeSession = session;
         activeSession.setMaxBinaryMessageBufferSize(NitroConstants.WEBSOCKET_BUFFER_SIZE);
 
-        server.connect(proxyProvider.getOriginalWebsocketUrl());
+        server.connect(proxyProvider.getOriginalWebsocketUrl(), proxyProvider.getOriginalOriginUrl());
 
         final HProxy proxy = new HProxy(HClient.NITRO, "", "", -1, -1, "");
 

G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketServer.java

@@ -8,9 +8,11 @@ import gearth.protocol.packethandler.PacketHandler;
 import javax.websocket.*;
 import java.io.IOException;
 import java.net.URI;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
 
-@ClientEndpoint
+public class NitroWebsocketServer extends Endpoint implements NitroSession {
-public class NitroWebsocketServer implements NitroSession {
 
     private final PacketHandler packetHandler;
     private final NitroWebsocketClient client;
@@ -21,32 +23,50 @@ public class NitroWebsocketServer implements NitroSession {
         this.packetHandler = new NitroPacketHandler(HMessage.Direction.TOCLIENT, client, connection.getExtensionHandler(), connection.getTrafficObservables());
     }
 
-    public void connect(String websocketUrl) throws IOException {
+    public void connect(String websocketUrl, String originUrl) throws IOException {
         try {
-            ContainerProvider.getWebSocketContainer().connectToServer(this, URI.create(websocketUrl));
+            ClientEndpointConfig.Builder builder = ClientEndpointConfig.Builder.create();
+
+            if (originUrl != null) {
+                builder.configurator(new ClientEndpointConfig.Configurator() {
+                    @Override
+                    public void beforeRequest(Map<String, List<String>> headers) {
+                        headers.put("Origin", Collections.singletonList(originUrl));
+                    }
+                });
+            }
+
+            ClientEndpointConfig config = builder.build();
+
+            ContainerProvider.getWebSocketContainer().connectToServer(this, config, URI.create(websocketUrl));
         } catch (DeploymentException e) {
             throw new IOException("Failed to deploy websocket client", e);
         }
     }
 
-    @OnOpen
+    @Override
-    public void onOpen(Session Session) {
+    public void onOpen(Session session, EndpointConfig config) {
-        this.activeSession = Session;
+        this.activeSession = session;
         this.activeSession.setMaxBinaryMessageBufferSize(NitroConstants.WEBSOCKET_BUFFER_SIZE);
+        this.activeSession.addMessageHandler(new MessageHandler.Whole<byte[]>() {
+            @Override
+            public void onMessage(byte[] message) {
+                try {
+                    packetHandler.act(message);
+                } catch (IOException e) {
+                    e.printStackTrace();
+                }
+            }
+        });
     }
 
-    @OnMessage
+    @Override
-    public void onMessage(byte[] b, Session session) throws IOException {
+    public void onClose(Session session, CloseReason closeReason) {
-        packetHandler.act(b);
-    }
-
-    @OnClose
-    public void onClose(Session userSession, CloseReason reason) {
         // Hotel closed connection.
         client.shutdownProxy();
     }
 
-    @OnError
+    @Override
     public void onError(Session session, Throwable throwable) {
         throwable.printStackTrace();