Carefully release ByteBuf when processing WebSocketFrame (#430)

* Safer bytebuf copying procedure

* Carefully release bytebuf when working with websocket frames

* Additional comment for a closing frame

* Simplify branching when checking WebSocketFrames

Author: kachayev

src/aleph/http/server.clj

@@ -543,35 +543,45 @@
 
          :channel-read
          ([_ ctx msg]
-          (try
-            (let [ch (.channel ctx)]
-              (if-not (instance? WebSocketFrame msg)
-                (.fireChannelRead ctx msg)
-                (let [^WebSocketFrame msg msg]
-                  (condp instance? msg
-
-                    TextWebSocketFrame
-                    (netty/put! ch in (.text ^TextWebSocketFrame msg))
-
-                    BinaryWebSocketFrame
-                    (let [body (.content ^BinaryWebSocketFrame msg)]
-                      (netty/put! ch in
-                                  (if raw-stream?
-                                    (netty/acquire body)
-                                    (netty/buf->array body))))
-
-                    PingWebSocketFrame
-                    (netty/write-and-flush ch (PongWebSocketFrame. (netty/acquire (.content msg))))
-
-                    PongWebSocketFrame
-                    (http/resolve-pings! pending-pings true)
-
-                    CloseWebSocketFrame
-                    (.close handshaker ch (netty/acquire msg))
-
-                    (.fireChannelRead ctx msg)))))
-            (finally
-              (netty/release msg)))))]))))
+         (let [ch (.channel ctx)]
+           (cond
+             (instance? TextWebSocketFrame msg)
+             (let [text (.text ^TextWebSocketFrame msg)]
+               ;; working with text now, so we do not need
+               ;; ByteBuf inside TextWebSocketFrame
+               ;; note, that all *WebSocketFrame classes are
+               ;; subclasses of DefaultByteBufHolder, meaning
+               ;; there's no difference between releasing
+               ;; frame & frame's content
+               (netty/release msg)
+               (netty/put! ch in text))
+
+             (instance? BinaryWebSocketFrame msg)
+             (let [body (.content ^BinaryWebSocketFrame msg)]
+               (netty/put! ch in
+                 (if raw-stream?
+                   body
+                   ;; copied data into byte array, deallocating ByteBuf
+                   (netty/release-buf->array body))))
+
+             (instance? PingWebSocketFrame msg)
+             (let [body (.content ^PingWebSocketFrame msg)]
+               ;; reusing the same buffer
+               ;; will be deallocated by Netty
+               (netty/write-and-flush ch (PongWebSocketFrame. body)))
+
+             (instance? PongWebSocketFrame msg)
+             (http/resolve-pings! pending-pings true)
+
+             (instance? CloseWebSocketFrame msg)
+             ;; reusing the same buffer
+             ;; will be deallocated by Netty
+             (.close handshaker ch msg)
+
+             :else
+             ;; no need to release buffer when passing to a next handler
+             (.fireChannelRead ctx msg)))))]))))
+
 
 ;; note, as we set `keep-alive?` to `false`, `send-message` will close the connection
 ;; after writes are done, which is exactly what we expect to happen

src/aleph/netty.clj

@@ -119,9 +119,10 @@
     (.array dst)))
 
 (defn release-buf->array [^ByteBuf buf]
-  (let [ary (buf->array buf)]
-    (release buf)
-    ary))
+  (try
+    (buf->array buf)
+    (finally
+      (release buf))))
 
 (defn bufs->array [bufs]
   (let [bufs' (mapcat #(.nioBuffers ^ByteBuf %) bufs)