Fix for #313 [websocket] Improve WebSocketHandshakeFilter logic

Atmosphere · Apr 27, 2012 · 5cee0a1 · 5cee0a1
1 parent d83f8e9
commit 5cee0a1
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 13 deletions.
diff --git a/modules/cpr/src/main/java/org/atmosphere/util/Utils.java b/modules/cpr/src/main/java/org/atmosphere/util/Utils.java
@@ -15,8 +15,7 @@
  */
 package org.atmosphere.util;
 
-import org.atmosphere.cpr.AtmosphereRequest;
-
+import javax.servlet.http.HttpServletRequest;
 import java.util.Enumeration;
 
 import static org.atmosphere.cpr.HeaderConfig.WEBSOCKET_UPGRADE;
@@ -28,10 +27,10 @@
  */
 public final class Utils {
 
-    public static boolean webSocketEnabled(AtmosphereRequest request) {
+    public static boolean webSocketEnabled(HttpServletRequest request) {
         boolean webSocketEnabled = false;
         Enumeration<String> connection = request.getHeaders("Connection");
-        if (connection == null) {
+        if (connection == null || !connection.hasMoreElements()) {
             connection = request.getHeaders("connection");
         }
 

diff --git a/modules/cpr/src/main/java/org/atmosphere/websocket/WebSocketHandshakeFilter.java b/modules/cpr/src/main/java/org/atmosphere/websocket/WebSocketHandshakeFilter.java
@@ -16,6 +16,7 @@
 package org.atmosphere.websocket;
 
 import org.atmosphere.cpr.ApplicationConfig;
+import org.atmosphere.util.Utils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -28,6 +29,7 @@
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import java.io.IOException;
+import java.util.Enumeration;
 
 import static org.atmosphere.cpr.HeaderConfig.X_ATMOSPHERE_ERROR;
 
@@ -53,26 +55,24 @@ public void init(FilterConfig filterConfig) throws ServletException {
     @Override
     public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
 
-        if (HttpServletRequest.class.cast(request).getHeader("Connection") != null && HttpServletRequest.class.cast(request).getHeader("Connection").equalsIgnoreCase("upgrade")) {
-            int draft = HttpServletRequest.class.cast(request).getIntHeader("Sec-WebSocket-Version");
+        HttpServletRequest r = HttpServletRequest.class.cast(request);
+        if (Utils.webSocketEnabled(r)) {
+            int draft =r.getIntHeader("Sec-WebSocket-Version");
             if (draft < 0) {
-                draft = HttpServletRequest.class.cast(request).getIntHeader("Sec-WebSocket-Draft");
+                draft = r.getIntHeader("Sec-WebSocket-Draft");
             }
 
             if (bannedVersion != null) {
                 for (String s : bannedVersion) {
-                    if (Integer.getInteger(s) == draft) {
+                    if (Integer.parseInt(s) == draft) {
+                        logger.error("Invalid WebSocket Specification {} with {} ",
+                                r.getHeader("Connection"), r.getIntHeader("Sec-WebSocket-Version"));
                         HttpServletResponse.class.cast(response).addHeader(X_ATMOSPHERE_ERROR, "Websocket protocol not supported");
                         HttpServletResponse.class.cast(response).sendError(202, "Websocket protocol not supported");
                         return;
                     }
                 }
             }
-        } else if (HttpServletRequest.class.cast(request).getIntHeader("Sec-WebSocket-Version") > 0) {
-            logger.error("Invalid WebSocket Specification {} with {} ", HttpServletRequest.class.cast(request).getHeader("Connection"), HttpServletRequest.class.cast(request).getIntHeader("Sec-WebSocket-Version"));
-            HttpServletResponse.class.cast(response).addHeader(X_ATMOSPHERE_ERROR, "Websocket protocol not supported");
-            HttpServletResponse.class.cast(response).sendError(202, "Websocket protocol not supported");
-            return;
         }
         chain.doFilter(request, response);
     }

diff --git a/samples/chat/src/main/webapp/WEB-INF/web.xml b/samples/chat/src/main/webapp/WEB-INF/web.xml
@@ -6,6 +6,14 @@
 
     <description>Atmosphere Chat</description>
     <display-name>Atmosphere Chat</display-name>
+    <filter>
+        <filter-name>WebSocketHandshakeFilter</filter-name>
+        <filter-class>org.atmosphere.websocket.WebSocketHandshakeFilter</filter-class>
+        <init-param>
+            <param-name>org.atmosphere.websocket.bannedVersion</param-name>
+            <param-value>-1</param-value>
+        </init-param>
+    </filter>
     <servlet>
         <description>AtmosphereServlet</description>
         <servlet-name>AtmosphereServlet</servlet-name>
@@ -17,5 +25,10 @@
         <servlet-name>AtmosphereServlet</servlet-name>
         <url-pattern>/chat/*</url-pattern>
     </servlet-mapping>
+    <filter-mapping>
+        <filter-name>WebSocketHandshakeFilter</filter-name>
+        <url-pattern>/*</url-pattern>
+    </filter-mapping>
+
 </web-app>