Issue #5868 - allow request attributes to be set in websocket upgrade

jetty-websocket/websocket-core-client/src/main/java/org/eclipse/jetty/websocket/core/client/CoreClientUpgradeRequest.java

@@ -37,6 +37,7 @@
 import org.eclipse.jetty.http.HttpScheme;
 import org.eclipse.jetty.http.HttpStatus;
 import org.eclipse.jetty.http.HttpVersion;
+import org.eclipse.jetty.io.ByteBufferPool;
 import org.eclipse.jetty.io.EndPoint;
 import org.eclipse.jetty.util.Callback;
 import org.eclipse.jetty.util.QuotedStringTokenizer;
@@ -432,7 +433,8 @@ else if (values.length == 1)
         customizer.customize(coreSession);
 
         HttpClient httpClient = wsClient.getHttpClient();
-        WebSocketConnection wsConnection = new WebSocketConnection(endPoint, httpClient.getExecutor(), httpClient.getScheduler(), httpClient.getByteBufferPool(), coreSession);
+        ByteBufferPool bufferPool = wsClient.getWebSocketComponents().getBufferPool();
+        WebSocketConnection wsConnection = new WebSocketConnection(endPoint, httpClient.getExecutor(), httpClient.getScheduler(), bufferPool, coreSession);
         wsClient.getEventListeners().forEach(wsConnection::addEventListener);
         coreSession.setWebSocketConnection(wsConnection);
         notifyUpgradeListeners((listener) -> listener.onHandshakeResponse(this, response));

jetty-websocket/websocket-core-server/src/main/java/org/eclipse/jetty/websocket/core/server/ServerUpgradeRequest.java

@@ -19,15 +19,16 @@
 import java.net.URI;
 import java.security.Principal;
 import java.security.cert.X509Certificate;
-import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.Enumeration;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.TreeMap;
+import java.util.stream.Collectors;
 import javax.servlet.ServletRequest;
-import javax.servlet.http.Cookie;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpSession;
 
@@ -38,11 +39,15 @@
 import org.eclipse.jetty.websocket.core.WebSocketConstants;
 import org.eclipse.jetty.websocket.core.server.internal.UpgradeHttpServletRequest;
 
+/**
+ * Upgrade request used for websocket negotiation.
+ * Provides getters for things like the requested extensions and subprotocols so that the headers don't have to be parsed manually.
+ */
 public class ServerUpgradeRequest
 {
     private final URI requestURI;
     private final String queryString;
-    private final UpgradeHttpServletRequest request;
+    private final HttpServletRequest request;
     private final boolean secure;
     private final WebSocketNegotiation negotiation;
     private List<HttpCookie> cookies;
@@ -51,20 +56,17 @@ public class ServerUpgradeRequest
     public ServerUpgradeRequest(WebSocketNegotiation negotiation) throws BadMessageException
     {
         this.negotiation = negotiation;
-        HttpServletRequest httpRequest = negotiation.getRequest();
-        this.queryString = httpRequest.getQueryString();
-        this.secure = httpRequest.isSecure();
+        this.request = negotiation.getRequest();
+        this.queryString = request.getQueryString();
+        this.secure = request.isSecure();
 
         try
         {
-            // TODO why is this URL and not URI?
-            StringBuffer uri = httpRequest.getRequestURL();
-            // WHY?
+            StringBuffer uri = request.getRequestURL();
             if (this.queryString != null)
                 uri.append("?").append(this.queryString);
             uri.replace(0, uri.indexOf(":"), secure ? "wss" : "ws");
             this.requestURI = new URI(uri.toString());
-            this.request = new UpgradeHttpServletRequest(httpRequest);
         }
         catch (Throwable t)
         {
@@ -88,17 +90,9 @@ public List<HttpCookie> getCookies()
     {
         if (cookies == null)
         {
-            Cookie[] requestCookies = request.getCookies();
-            if (requestCookies != null)
-            {
-                cookies = new ArrayList<>();
-                for (Cookie requestCookie : requestCookies)
-                {
-                    HttpCookie cookie = new HttpCookie(requestCookie.getName(), requestCookie.getValue());
-                    // No point handling domain/path/expires/secure/httponly on client request cookies
-                    cookies.add(cookie);
-                }
-            }
+            cookies = Arrays.stream(request.getCookies())
+                .map(c -> new HttpCookie(c.getName(), c.getValue()))
+                .collect(Collectors.toList());
         }
 
         return cookies;
@@ -130,12 +124,7 @@ public String getHeader(String name)
      */
     public int getHeaderInt(String name)
     {
-        String val = request.getHeader(name);
-        if (val == null)
-        {
-            return -1;
-        }
-        return Integer.parseInt(val);
+        return request.getIntHeader(name);
     }
 
     /**
@@ -144,7 +133,14 @@ public int getHeaderInt(String name)
      */
     public Map<String, List<String>> getHeadersMap()
     {
-        return request.getHeaders();
+        Map<String, List<String>> headers = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
+        Enumeration<String> headerNames = request.getHeaderNames();
+        while (headerNames.hasMoreElements())
+        {
+            String name = headerNames.nextElement();
+            headers.put(name, Collections.list(request.getHeaders(name)));
+        }
+        return headers;
     }
 
     /**
@@ -154,7 +150,10 @@ public Map<String, List<String>> getHeadersMap()
      */
     public List<String> getHeaders(String name)
     {
-        return request.getHeaders().get(name);
+        Enumeration<String> headers = request.getHeaders(name);
+        if (headers == null || !headers.hasMoreElements())
+            return null;
+        return Collections.list(headers);
     }
 
     /**
@@ -163,7 +162,6 @@ public List<String> getHeaders(String name)
      */
     public String getHost()
     {
-        // TODO why is this not HttpServletRequest#getHost ?
         return requestURI.getHost();
     }
 
@@ -209,7 +207,7 @@ public Enumeration<Locale> getLocales()
      */
     public SocketAddress getLocalSocketAddress()
     {
-        // TODO: fix when HttpServletRequest can use Unix Socket stuff
+        // TODO: fix when HttpServletRequest can use Unix Socket stuff.
         return new InetSocketAddress(request.getLocalAddr(), request.getLocalPort());
     }
 
@@ -326,11 +324,17 @@ public Object getServletAttribute(String name)
 
     /**
      * @return Request attribute map
-     * @see UpgradeHttpServletRequest#getAttributes()
      */
     public Map<String, Object> getServletAttributes()
     {
-        return request.getAttributes();
+        Map<String, Object> attributes = new HashMap<>(2);
+        Enumeration<String> attributeNames = request.getAttributeNames();
+        while (attributeNames.hasMoreElements())
+        {
+            String name = attributeNames.nextElement();
+            attributes.put(name, request.getAttribute(name));
+        }
+        return attributes;
     }
 
     /**
@@ -377,9 +381,7 @@ public boolean hasSubProtocol(String subprotocol)
         for (String protocol : getSubProtocols())
         {
             if (protocol.equalsIgnoreCase(subprotocol))
-            {
                 return true;
-            }
         }
         return false;
     }

jetty-websocket/websocket-core-server/src/main/java/org/eclipse/jetty/websocket/core/server/ServerUpgradeResponse.java

@@ -26,6 +26,10 @@
 import org.eclipse.jetty.http.HttpHeader;
 import org.eclipse.jetty.websocket.core.ExtensionConfig;
 
+/**
+ * Upgrade response used for websocket negotiation.
+ * Allows setting of extensions and subprotocol without using headers directly.
+ */
 public class ServerUpgradeResponse
 {
     private final HttpServletResponse response;

jetty-websocket/websocket-core-server/src/main/java/org/eclipse/jetty/websocket/core/server/internal/AbstractHandshaker.java

@@ -59,7 +59,11 @@ public boolean upgradeRequest(WebSocketNegotiator negotiator, HttpServletRequest
         if (!validateRequest(request))
             return false;
 
-        WebSocketNegotiation negotiation = newNegotiation(request, response, components);
+        // After negotiation these can be set to copy data from request and disable unavailable methods.
+        UpgradeHttpServletRequest upgradeRequest = new UpgradeHttpServletRequest(request);
+        UpgradeHttpServletResponse upgradeResponse = new UpgradeHttpServletResponse(response);
+
+        WebSocketNegotiation negotiation = newNegotiation(upgradeRequest, upgradeResponse, components);
         if (LOG.isDebugEnabled())
             LOG.debug("negotiation {}", negotiation);
         negotiation.negotiate();
@@ -127,7 +131,7 @@ public boolean upgradeRequest(WebSocketNegotiator negotiator, HttpServletRequest
         Negotiated negotiated = new Negotiated(baseRequest.getHttpURI().toURI(), protocol, baseRequest.isSecure(), extensionStack, WebSocketConstants.SPEC_VERSION_STRING);
 
         // Create the Session
-        WebSocketCoreSession coreSession = newWebSocketCoreSession(request, handler, negotiated, components);
+        WebSocketCoreSession coreSession = newWebSocketCoreSession(upgradeRequest, handler, negotiated, components);
         if (defaultCustomizer != null)
             defaultCustomizer.customize(coreSession);
         negotiator.customize(coreSession);
@@ -159,6 +163,10 @@ public boolean upgradeRequest(WebSocketNegotiator negotiator, HttpServletRequest
 
         baseRequest.setAttribute(HttpTransport.UPGRADE_CONNECTION_ATTRIBUTE, connection);
 
+        // Save state from request/response and remove reference to the base request/response.
+        upgradeRequest.upgrade();
+        upgradeResponse.upgrade();
+
         if (LOG.isDebugEnabled())
             LOG.debug("upgrade connection={} session={} framehandler={}", connection, coreSession, handler);