cgbystrom · ctpahhik · Jul 29, 2013 · Jul 29, 2013 · Jul 29, 2013 · Jul 30, 2013
diff --git a/pom.xml b/pom.xml
@@ -8,7 +8,7 @@
     <modelVersion>4.0.0</modelVersion>
     <groupId>com.cgbystrom</groupId>
     <artifactId>sockjs-netty</artifactId>
-    <version>0.1.0-SNAPSHOT</version>
+    <version>0.1.2-SNAPSHOT</version>
     <name>SockJS for JBoss Netty</name>
     <packaging>jar</packaging>
     <description>
@@ -96,7 +96,18 @@
                     </execution>
                 </executions>
             </plugin>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-resources-plugin</artifactId>
+                <version>2.2</version>
+            </plugin>
         </plugins>
+        <resources>
+            <resource>
+                <directory>resources</directory>
+                <targetPath>resources</targetPath>
+            </resource>
+        </resources>
     </build>
 
     <distributionManagement>

diff --git a/resources/sockjs-0.3.4.min.js b/resources/sockjs-0.3.4.min.js
diff --git a/src/main/java/com/cgbystrom/sockjs/IframePage.java b/src/main/java/com/cgbystrom/sockjs/IframePage.java
@@ -5,7 +5,7 @@
 import org.jboss.netty.handler.codec.http.*;
 import org.jboss.netty.util.CharsetUtil;
 
-import java.io.UnsupportedEncodingException;
+import java.io.*;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.util.Formatter;
@@ -15,7 +15,11 @@ public class IframePage {
     private String etag;
 
     public IframePage(String url) {
-        content = createContent(url);
+        content = createContent(getJsClient(url));
+    }
+
+    public IframePage(InputStream clientResource) throws IOException {
+        content = createContent(getJsClient(clientResource));
     }
 
     public void handle(HttpRequest request, HttpResponse response) {
@@ -44,7 +48,7 @@ public void handle(HttpRequest request, HttpResponse response) {
         response.setHeader(HttpHeaders.Names.ETAG, etag);
     }
 
-    private ChannelBuffer createContent(String url) {
+    private ChannelBuffer createContent(String jsClient) {
         String content = "<!DOCTYPE html>\n" +
         "<html>\n" +
         "<head>\n" +
@@ -54,7 +58,7 @@ private ChannelBuffer createContent(String url) {
         "    document.domain = document.domain;\n" +
         "    _sockjs_onload = function(){SockJS.bootstrap_iframe();};\n" +
         "  </script>\n" +
-        "  <script src=\"" + url + "\"></script>\n" +
+        jsClient +
         "</head>\n" +
         "<body>\n" +
         "  <h2>Don't panic!</h2>\n" +
@@ -68,6 +72,22 @@ private ChannelBuffer createContent(String url) {
         return ChannelBuffers.copiedBuffer(content, CharsetUtil.UTF_8);
     }
 
+    private String getJsClient(String url) {
+        return "<script src=\"" + url + "\"></script>\n";
+    }
+
+    private String getJsClient(InputStream clientResource) throws IOException {
+        BufferedReader reader = new BufferedReader(new InputStreamReader(clientResource));
+        StringBuilder builder = new StringBuilder();
+        builder.append("<script>");
+        String line;
+        while ((line = reader.readLine()) != null) {
+            builder.append(line).append("\n");
+        }
+        builder.append("</script>\n");
+        return builder.toString();
+    }
+
     private static String generateMd5(String value) {
         String encryptedString = null;
         byte[] bytesToBeEncrypted;

diff --git a/src/main/java/com/cgbystrom/sockjs/ServiceRouter.java b/src/main/java/com/cgbystrom/sockjs/ServiceRouter.java
@@ -12,6 +12,7 @@
 import static org.jboss.netty.handler.codec.http.HttpHeaders.Values.KEEP_ALIVE;
 
 import java.io.IOException;
+import java.io.InputStream;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Random;
@@ -27,6 +28,7 @@ private enum SessionCreation { CREATE_OR_REUSE, FORCE_REUSE, FORCE_CREATE }
 
     private final Map<String, ServiceMetadata> services = new LinkedHashMap<String, ServiceMetadata>();
     private IframePage iframe;
+    private volatile boolean noNewConnections;
 
     /**
      *
@@ -37,6 +39,22 @@ public ServiceRouter(String clientUrl) {
         this.iframe = new IframePage(clientUrl);
     }
 
+    public ServiceRouter() throws IOException {
+        this(ServiceRouter.class.getResourceAsStream("/resources/sockjs-0.3.4.min.js"));
+    }
+
+    public ServiceRouter(InputStream clientResource) throws IOException {
+        try {
+            this.iframe = new IframePage(clientResource);
+        } finally {
+            if (clientResource != null) {
+                try {
+                    clientResource.close();
+                } catch (Exception ignore) {}
+            }
+        }
+    }
+
     public synchronized ServiceMetadata registerService(String baseUrl, final SessionCallback service,
                 boolean isWebSocketEnabled, int maxResponseSize) {
         return registerService(baseUrl, new SessionCallbackFactory() {
@@ -103,9 +121,15 @@ private void handleService(ChannelHandlerContext ctx, MessageEvent e, ServiceMet
             iframe.handle(request, response);
             writeResponse(e.getChannel(), request, response);
         } else if (path.startsWith("/info")) {
+            if (noNewConnections) {
+                response.setStatus(HttpResponseStatus.SERVICE_UNAVAILABLE);
+                response.setHeader(CACHE_CONTROL, "no-store, no-cache, must-revalidate, max-age=0");
+                response.setContent(ChannelBuffers.copiedBuffer("Temporary unavailable", CharsetUtil.UTF_8));
+            } else {
             response.setHeader(CONTENT_TYPE, "application/json; charset=UTF-8");
             response.setHeader(CACHE_CONTROL, "no-store, no-cache, must-revalidate, max-age=0");
             response.setContent(getInfo(serviceMetadata));
+            }
             writeResponse(e.getChannel(), request, response);
         } else if (path.startsWith("/websocket")) {
             // Raw web socket
@@ -236,6 +260,30 @@ private ChannelBuffer getInfo(ServiceMetadata metadata) {
         return ChannelBuffers.copiedBuffer(sb.toString(), CharsetUtil.UTF_8);
     }
 
+    /**
+     * Check if server is accepting new connections
+     *
+     * @return true if new connections are not allowed
+     */
+    public boolean isNoNewConnections() {
+        return noNewConnections;
+    }
+
+    /**
+     * Stop accepting new connections.
+     * Existed sessions still to be handled
+     */
+    public void stopAcceptingNewConnections() {
+        noNewConnections = true;
+    }
+
+    /**
+     * Continue accepting new connections.
+     */
+    public void startAcceptingNewConnections() {
+        noNewConnections = false;
+    }
+
     public static class ServiceMetadata {
         private ServiceMetadata(String url, SessionCallbackFactory factory, ConcurrentHashMap<String,
                 SessionHandler> sessions, boolean isWebSocketEnabled, int maxResponseSize) {

diff --git a/src/main/java/com/cgbystrom/sockjs/Session.java b/src/main/java/com/cgbystrom/sockjs/Session.java
@@ -1,6 +1,15 @@
 package com.cgbystrom.sockjs;
 
+import java.net.SocketAddress;
+
 public interface Session {
     public void send(String message);
     public void close();
+
+    /**
+     * Allow to get client's Address
+     *
+     * @return SocketAddress (ip, port etc.) of the client.
+     */
+    public SocketAddress getRemoteAddress();
 }
diff --git a/src/main/java/com/cgbystrom/sockjs/SessionHandler.java b/src/main/java/com/cgbystrom/sockjs/SessionHandler.java
@@ -5,6 +5,7 @@
 import org.jboss.netty.logging.InternalLoggerFactory;
 import org.jboss.netty.util.CharsetUtil;
 
+import java.net.SocketAddress;
 import java.util.ArrayList;
 import java.util.LinkedList;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -100,6 +101,7 @@ public synchronized void channelClosed(ChannelHandlerContext ctx, ChannelStateEv
         if (state == State.OPEN && !serverHasInitiatedClose.get()) {
             logger.debug("Session " + id + " underlying channel closed unexpectedly. Flagging session as interrupted." + e.getChannel());
             setState(State.INTERRUPTED);
+            sessionCallback.onClose();
         } else {
             logger.debug("Session " + id + " underlying channel closed " + e.getChannel());
         }
@@ -139,7 +141,11 @@ public synchronized void send(String message) {
 
     @Override
     public void close() {
-        close(3000, "Go away!");
+        try {
+            flush();
+        } finally {
+            close(3000, "Go away!");
+        }
     }
 
     public synchronized void close(int code, String message) {
@@ -152,7 +158,7 @@ public synchronized void close(int code, String message) {
                 channel.write(closeReason);
             }
             // FIXME: Should we really call onClose here? Potentially calling it twice for same session close?
-            sessionCallback.onClose();
+            //sessionCallback.onClose();
         }
     }
 
@@ -197,4 +203,12 @@ public LockException(Channel channel) {
             super("Session is locked by channel " + channel + ". Please disconnect other channel first before trying to register it with a session.");
         }
     }
+
+    @Override
+    public SocketAddress getRemoteAddress() {
+        if (channel != null) {
+            return channel.getRemoteAddress();
+        }
+        return null;
+    }
 }
diff --git a/src/test/java/com/cgbystrom/sockjs/TestServer.java b/src/test/java/com/cgbystrom/sockjs/TestServer.java
@@ -14,6 +14,7 @@
 import org.jboss.netty.handler.codec.http.HttpRequestDecoder;
 import org.jboss.netty.handler.codec.http.HttpResponseEncoder;
 
+import org.jboss.netty.handler.execution.OrderedMemoryAwareThreadPoolExecutor;
 import org.jboss.netty.logging.InternalLoggerFactory;
 import org.jboss.netty.logging.Slf4JLoggerFactory;
 import org.slf4j.LoggerFactory;
@@ -24,7 +25,7 @@
 import static org.jboss.netty.channel.Channels.pipeline;
 
 public class TestServer {
-    public static void main(String[] args) {
+    public static void main(String[] args) throws Exception {
         Logger rootLogger = (Logger) LoggerFactory.getLogger(Logger.ROOT_LOGGER_NAME);
         LoggerContext loggerContext = rootLogger.getLoggerContext();
         loggerContext.reset();
@@ -43,10 +44,10 @@ public static void main(String[] args) {
 
         ServerBootstrap bootstrap = new ServerBootstrap(
                 new NioServerSocketChannelFactory(
-                        Executors.newCachedThreadPool(),
-                        Executors.newCachedThreadPool()));
+                        Executors.newCachedThreadPool(), new OrderedMemoryAwareThreadPoolExecutor(20, 0, 0)));
 
-        final ServiceRouter router = new ServiceRouter("http://cdn.sockjs.org/sockjs-0.3.4.min.js");
+        //final ServiceRouter router = new ServiceRouter("http://cdn.sockjs.org/sockjs-0.3.4.min.js");
+        final ServiceRouter router = new ServiceRouter(TestServer.class.getClassLoader().getResourceAsStream("resources/sockjs-0.3.4.min.js"));
         router.registerService("/echo", new SessionCallbackFactory() {
             @Override
             public EchoSession getSession(String id) throws Exception {