diff --git a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/NitroProxyProvider.java b/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/NitroProxyProvider.java index 3cb47a1..b8c0810 100644 --- a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/NitroProxyProvider.java +++ b/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; + } } diff --git a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketClient.java b/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketClient.java index aabb946..59bf027 100644 --- a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketClient.java +++ b/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, ""); diff --git a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketServer.java b/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketServer.java index 04b238c..b54b96c 100644 --- a/G-Earth/src/main/java/gearth/protocol/connection/proxy/nitro/websocket/NitroWebsocketServer.java +++ b/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 implements NitroSession { +public class NitroWebsocketServer extends Endpoint 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> 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 - public void onOpen(Session Session) { - this.activeSession = Session; + @Override + public void onOpen(Session session, EndpointConfig config) { + this.activeSession = session; this.activeSession.setMaxBinaryMessageBufferSize(NitroConstants.WEBSOCKET_BUFFER_SIZE); + this.activeSession.addMessageHandler(new MessageHandler.Whole() { + @Override + public void onMessage(byte[] message) { + try { + packetHandler.act(message); + } catch (IOException e) { + e.printStackTrace(); + } + } + }); } - @OnMessage - public void onMessage(byte[] b, Session session) throws IOException { - packetHandler.act(b); - } - - @OnClose - public void onClose(Session userSession, CloseReason reason) { + @Override + public void onClose(Session session, CloseReason closeReason) { // Hotel closed connection. client.shutdownProxy(); } - @OnError + @Override public void onError(Session session, Throwable throwable) { throwable.printStackTrace();