Notify Pings and Pongs, close AsyncHttpClient#517

Author: Stephane Landelle

providers/netty/src/main/java/org/asynchttpclient/providers/netty/handler/WebSocketProtocol.java

@@ -22,6 +22,9 @@
 import io.netty.handler.codec.http.LastHttpContent;
 import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.CloseWebSocketFrame;
+import io.netty.handler.codec.http.websocketx.PingWebSocketFrame;
+import io.netty.handler.codec.http.websocketx.PongWebSocketFrame;
+import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.WebSocketFrame;
 
 import java.io.IOException;
@@ -74,7 +77,7 @@ public void handle(Channel channel, NettyResponseFuture<?> future, Object e) thr
             HttpResponse response = (HttpResponse) e;
             // we buffer the response until we get the LastHttpContent
             future.setPendingResponse(response);
-            
+
         } else if (e instanceof LastHttpContent) {
             HttpResponse response = future.getPendingResponse();
             future.setPendingResponse(null);
@@ -115,8 +118,7 @@ public void handle(Channel channel, NettyResponseFuture<?> future, Object e) thr
             }
 
             String accept = response.headers().get(HttpHeaders.Names.SEC_WEBSOCKET_ACCEPT);
-            String key = getAcceptKey(future.getNettyRequest().getHttpRequest().headers()
-                    .get(HttpHeaders.Names.SEC_WEBSOCKET_KEY));
+            String key = getAcceptKey(future.getNettyRequest().getHttpRequest().headers().get(HttpHeaders.Names.SEC_WEBSOCKET_KEY));
             if (accept == null || !accept.equals(key)) {
                 requestSender.abort(future, new IOException(String.format("Invalid challenge. Actual: %s. Expected: %s", accept, key)));
             }
@@ -141,13 +143,17 @@ public void handle(Channel channel, NettyResponseFuture<?> future, Object e) thr
                     ByteBuf buf = frame.content();
                     if (buf != null && buf.readableBytes() > 0) {
                         try {
-                            NettyResponseBodyPart rp = nettyConfig.getBodyPartFactory().newResponseBodyPart(buf, frame.isFinalFragment());
-                            handler.onBodyPartReceived(rp);
+                            NettyResponseBodyPart part = nettyConfig.getBodyPartFactory().newResponseBodyPart(buf, frame.isFinalFragment());
+                            handler.onBodyPartReceived(part);
 
                             if (frame instanceof BinaryWebSocketFrame) {
-                                webSocket.onBinaryFragment(rp);
-                            } else {
-                                webSocket.onTextFragment(rp);
+                                webSocket.onBinaryFragment(part);
+                            } else if (frame instanceof TextWebSocketFrame) {
+                                webSocket.onTextFragment(part);
+                            } else if (frame instanceof PingWebSocketFrame) {
+                                webSocket.onPing(part);
+                            } else if (frame instanceof PongWebSocketFrame) {
+                                webSocket.onPong(part);
                             }
                         } finally {
                             buf.release();
@@ -196,7 +202,8 @@ public void onClose(Channel channel) {
             WebSocketUpgradeHandler h = WebSocketUpgradeHandler.class.cast(nettyResponse.getAsyncHandler());
             NettyWebSocket webSocket = NettyWebSocket.class.cast(h.onCompleted());
 
-            // FIXME How could this test not succeed, we just checked above that attribute is a NettyResponseFuture????
+            // FIXME How could this test not succeed, we just checked above that
+            // attribute is a NettyResponseFuture????
             logger.trace("Connection was closed abnormally (that is, with no close frame being sent).");
             if (attribute != DiscardEvent.INSTANCE && webSocket != null)
                 webSocket.close(1006, "Connection was closed abnormally (that is, with no close frame being sent).");

providers/netty/src/main/java/org/asynchttpclient/providers/netty/ws/NettyWebSocket.java

@@ -37,6 +37,8 @@
 import org.asynchttpclient.websocket.WebSocketByteListener;
 import org.asynchttpclient.websocket.WebSocketCloseCodeReasonListener;
 import org.asynchttpclient.websocket.WebSocketListener;
+import org.asynchttpclient.websocket.WebSocketPingListener;
+import org.asynchttpclient.websocket.WebSocketPongListener;
 import org.asynchttpclient.websocket.WebSocketTextFragmentListener;
 import org.asynchttpclient.websocket.WebSocketTextListener;
 import org.slf4j.Logger;
@@ -300,4 +302,20 @@ public void onTextFragment(HttpResponseBodyPart part) {
                 bufferFragment(fragment);
         }
     }
+
+    public void onPing(HttpResponseBodyPart part) {
+        for (WebSocketListener listener : listeners) {
+            if (listener instanceof WebSocketPingListener)
+                // bytes are cached in the part
+                WebSocketPingListener.class.cast(listener).onPing(part.getBodyPartBytes());
+        }
+    }
+
+    public void onPong(HttpResponseBodyPart part) {
+        for (WebSocketListener listener : listeners) {
+            if (listener instanceof WebSocketPongListener)
+                // bytes are cached in the part
+                WebSocketPongListener.class.cast(listener).onPong(part.getBodyPartBytes());
+        }
+    }
 }