Allow setting headers from a websocket handshake [IMMUTANT-580]

Author: tobias

project.clj

@@ -83,7 +83,7 @@
                         core.memoize               "0.5.6"
                         io.pedestal                "0.3.1"
                         http.async.client          "0.5.2"
-                        gniazdo                    "0.4.1a"
+                        gniazdo                    "0.4.1b"
                         compojure                  "1.3.4"
                         org.clojure/java.jdbc      "0.3.6"
                         h2                         "1.3.176"
@@ -94,7 +94,7 @@
                         environ                    "1.0.0"
 
                         ;; org.projectodd.wunderboss  "0.9.0"
-                        org.projectodd.wunderboss  "1.x.incremental.296"
+                        org.projectodd.wunderboss  "1.x.incremental.297"
                         ;; org.projectodd.wunderboss  "0.9.1-SNAPSHOT"
 
                         org.immutant               :version

web/dev-resources/testing/app.clj

@@ -85,6 +85,7 @@
        :on-error (fn [ch err]
                    (println "Error on websocket")
                    (.printStackTrace err))})
+    :headers {"ham" "biscuit"}
     :session (assoc (:session request) :ham :sandwich)))
 
 (def client-defined-handler (atom (fn [_] (throw (Exception. "no handler given")))))

web/src/immutant/web/async.clj

@@ -284,8 +284,8 @@
   milliseconds.
 
   When the ring handler is called during a WebSocket upgrade request,
-  any headers returned in the response map are ignored, but any changes to
-  the session are applied.
+  any changes to the session in the response map are applied, and any
+  headers from the response map are included in the upgrade response.
 
   Returns a ring response map, at least the :body of which *must* be
   returned in the response map from the calling ring handler."

web/src/immutant/web/internal/servlet.clj

@@ -26,7 +26,8 @@
               WebsocketChannel WebsocketChannel$OnMessage]
              [javax.servlet.http HttpServlet HttpServletRequest HttpServletResponse HttpSession]
              [javax.servlet Servlet ServletConfig ServletContext]
-             [javax.websocket Session Endpoint EndpointConfig MessageHandler$Whole CloseReason]
+             [javax.websocket Session
+              Endpoint EndpointConfig HandshakeResponse MessageHandler$Whole CloseReason]
              [javax.websocket.server ServerContainer
               ServerEndpointConfig ServerEndpointConfig$Builder ServerEndpointConfig$Configurator]))
 
@@ -104,6 +105,17 @@
 (defn ^ServerContainer server-container [^ServletContext context]
   (.getAttribute context "javax.websocket.server.ServerContainer"))
 
+(extend-type HandshakeResponse
+  hdr/Headers
+  (set-header [response ^String key value] (-> response .getHeaders (.put key [value])))
+  (add-header [response ^String key value]
+    (let [headers (.getHeaders response)
+          curr (.get headers key)]
+      (.put headers key
+        (if curr
+          (conj (vec curr) value)
+          [value])))))
+
 (defn add-endpoint-with-handler
   [^Servlet servlet ^ServletConfig config handler]
     (let [servlet-context (.getServletContext config)
@@ -118,15 +130,16 @@
           (create DelegatingJavaxEndpoint path)
           (configurator (proxy [ServerEndpointConfig$Configurator] []
                           (getEndpointInstance [_] (DelegatingJavaxEndpoint.))
-                          (modifyHandshake [_ _ _]
+                          (modifyHandshake [_ _ ^HandshakeResponse response]
                             (let [request (.get WebSocketHelpyHelpertonFilter/requestTL)
-                                  body (:body (handler (ring/ring-request-map request
-                                                         [:handler-type :servlet]
-                                                         [:servlet-request request]
-                                                         [:websocket? true])))]
+                                  {:keys [body headers]} (handler (ring/ring-request-map request
+                                                                    [:handler-type :servlet]
+                                                                    [:servlet-request request]
+                                                                    [:websocket? true]))]
+                              (hdr/set-headers response headers)
                               (when (instance? WebsocketChannel body)
-                          (DelegatingJavaxEndpoint/setCurrentDelegate
-                            (.endpoint ^WebsocketChannel body)))))))
+                                (DelegatingJavaxEndpoint/setCurrentDelegate
+                                  (.endpoint ^WebsocketChannel body)))))))
           build))))
 
 (defn add-endpoint

web/src/immutant/web/internal/undertow.clj

@@ -14,15 +14,14 @@
 
 (ns ^{:no-doc true} immutant.web.internal.undertow
     (:require [immutant.web.async            :as async]
-              [immutant.web.internal.ring    :as ring]
               [immutant.web.internal.headers :as hdr]
+              [immutant.web.internal.ring    :as ring]
               [ring.middleware.session       :as ring-session])
-    (:import java.net.URI
-             [io.undertow.server HttpHandler HttpServerExchange]
+    (:import [io.undertow.server HttpHandler HttpServerExchange]
              [io.undertow.server.session Session SessionConfig SessionCookieConfig]
              [io.undertow.util HeaderMap Headers HttpString Sessions]
              [io.undertow.websockets.core CloseMessage WebSocketChannel]
-             [io.undertow.websockets.spi WebSocketHttpExchange]
+             java.net.URI
              [org.projectodd.wunderboss.web.async Channel
               Channel$OnOpen Channel$OnClose Channel$OnError]
              [org.projectodd.wunderboss.web.async.websocket WebsocketChannel
@@ -180,20 +179,18 @@
           (on-message ch message))))))
 
 (defn ^:internal create-websocket-init-handler [handler-fn request-map-fn]
-  (let [http-exchange-tl (ThreadLocal.)
-        downstream-handler (create-http-handler handler-fn)]
+  (let [downstream-handler (create-http-handler handler-fn)]
     (UndertowWebsocket/createHandler
-      http-exchange-tl
       (reify WebsocketInitHandler
         (shouldConnect [_ exchange endpoint-wrapper]
-          (let [http-exchange (.get http-exchange-tl)]
-            (boolean
-              (let [body (:body (handler-fn (request-map-fn http-exchange
-                                              [:websocket? true]
-                                              [:server-exchange http-exchange]
-                                              [:handler-type :undertow])))]
-                (when (instance? WebsocketChannel body)
-                  (.setEndpoint endpoint-wrapper
-                    (.endpoint ^WebsocketChannel body))
-                  true))))))
+          (boolean
+            (let [{:keys [body headers] :as r} (handler-fn (request-map-fn exchange
+                                                             [:websocket? true]
+                                                             [:server-exchange exchange]
+                                                             [:handler-type :undertow]))]
+              (hdr/set-headers (.getResponseHeaders exchange) headers)
+              (when (instance? WebsocketChannel body)
+                (.setEndpoint endpoint-wrapper
+                  (.endpoint ^WebsocketChannel body))
+                true)))))
       downstream-handler)))

web/test-integration/immutant/web/integ_test.clj

@@ -207,6 +207,17 @@
       (is (= "HELLO" (deref result 5000 :failure)))
       (is (= {:count 1 :ham :sandwich} (decode (get-body (str (url) "session"))))))))
 
+(marktest websocket-upgrade-request-can-set-headers
+  (let [result (promise)]
+    (with-open [socket (ws/connect (str (url "ws") "ws")
+                         :on-receive (fn [m] (deliver result m))
+                         :cookies @cookies)
+                session (ws/session socket)]
+      (ws/send-msg socket "hello")
+      (is (= "HELLO" (deref result 5000 :failure)))
+      (let [headers (into {} (-> session .getUpgradeResponse .getHeaders))]
+        (is (= "biscuit" (first (headers "ham"))))))))
+
 (marktest nested-ws-routes
   (doseq [path ["" "foo" "foo/bar"]]
     (let [result (promise)]