Fix for AsyncHttpClient#91 [websocket] Add support a new Close Listener with code & reason

Author: jfarcand

src/main/java/com/ning/http/client/providers/netty/NettyAsyncHttpProvider.java

@@ -91,6 +91,7 @@
 import org.jboss.netty.handler.codec.http.HttpResponse;
 import org.jboss.netty.handler.codec.http.HttpResponseDecoder;
 import org.jboss.netty.handler.codec.http.HttpVersion;
+import org.jboss.netty.handler.codec.http.websocketx.CloseWebSocketFrame;
 import org.jboss.netty.handler.codec.http.websocketx.WebSocket08FrameDecoder;
 import org.jboss.netty.handler.codec.http.websocketx.WebSocket08FrameEncoder;
 import org.jboss.netty.handler.codec.http.websocketx.WebSocketFrame;
@@ -2370,6 +2371,15 @@ public void setContent(ChannelBuffer content) {
                     NettyWebSocket webSocket = NettyWebSocket.class.cast(h.onCompleted());
                     webSocket.onMessage(rp.getBodyPartBytes());
                     webSocket.onTextMessage(frame.getBinaryData().toString("UTF-8"));
+
+                    if (CloseWebSocketFrame.class.isAssignableFrom(frame.getClass())) {
+                        try {
+                            webSocket.onClose(CloseWebSocketFrame.class.cast(frame).getStatusCode(), CloseWebSocketFrame.class.cast(frame).getReasonText());
+                        } catch (Throwable t) {
+                            // Swallow any exception that may comes from a Netty version released before 3.4.0
+                            log.trace("", t);
+                        }
+                    }
                 }
             } else {
                 log.error("Invalid attachment {}", ctx.getAttachment());

src/main/java/com/ning/http/client/providers/netty/NettyWebSocket.java

@@ -14,19 +14,23 @@
 
 import com.ning.http.client.websocket.WebSocket;
 import com.ning.http.client.websocket.WebSocketByteListener;
+import com.ning.http.client.websocket.WebSocketCloseCodeReasonListener;
 import com.ning.http.client.websocket.WebSocketListener;
 import com.ning.http.client.websocket.WebSocketTextListener;
 import org.jboss.netty.channel.Channel;
 import org.jboss.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
 import org.jboss.netty.handler.codec.http.websocketx.PingWebSocketFrame;
 import org.jboss.netty.handler.codec.http.websocketx.PongWebSocketFrame;
 import org.jboss.netty.handler.codec.http.websocketx.TextWebSocketFrame;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.concurrent.ConcurrentLinkedQueue;
 
 import static org.jboss.netty.buffer.ChannelBuffers.wrappedBuffer;
 
 public class NettyWebSocket implements WebSocket {
+    private final static Logger logger = LoggerFactory.getLogger(NettyWebSocket.class);
 
     private final Channel channel;
     private final ConcurrentLinkedQueue<WebSocketListener> listeners = new ConcurrentLinkedQueue<WebSocketListener>();
@@ -67,7 +71,7 @@ public WebSocket sendPing(byte[] payload) {
         channel.write(new PingWebSocketFrame(wrappedBuffer(payload)));
         return this;
     }
-    
+
     @Override
     public WebSocket sendPong(byte[] payload) {
         channel.write(new PongWebSocketFrame(wrappedBuffer(payload)));
@@ -101,28 +105,49 @@ public void close() {
     protected void onMessage(byte[] message) {
         for (WebSocketListener l : listeners) {
             if (WebSocketByteListener.class.isAssignableFrom(l.getClass())) {
-                WebSocketByteListener.class.cast(l).onMessage(message);
+                try {
+                    WebSocketByteListener.class.cast(l).onMessage(message);
+                } catch (Exception ex) {
+                    l.onError(ex);
+                }
             }
         }
     }
 
     protected void onTextMessage(String message) {
         for (WebSocketListener l : listeners) {
             if (WebSocketTextListener.class.isAssignableFrom(l.getClass())) {
-                WebSocketTextListener.class.cast(l).onMessage(message);
+                try {
+                    WebSocketTextListener.class.cast(l).onMessage(message);
+                } catch (Exception ex) {
+                    l.onError(ex);
+                }
             }
         }
     }
 
     protected void onError(Throwable t) {
         for (WebSocketListener l : listeners) {
-            l.onError(t);
+            try {
+                l.onError(t);
+            } catch (Throwable t2) {
+                logger.error("", t2);
+            }
+
         }
     }
 
     protected void onClose() {
+        onClose(1000, "Normal closure; the connection successfully completed whatever purpose for which it was created.");
+    }
+
+    protected void onClose(int code, String reason) {
         for (WebSocketListener l : listeners) {
-            l.onClose(this);
+            if (WebSocketCloseCodeReasonListener.class.isAssignableFrom(l.getClass())) {
+                WebSocketCloseCodeReasonListener.class.cast(l).onClose(this, code, reason);
+            } else {
+                l.onClose(this);
+            }
         }
     }
 

src/main/java/com/ning/http/client/websocket/WebSocketListener.java

@@ -19,18 +19,21 @@ public interface WebSocketListener {
 
     /**
      * Invoked when the {@link WebSocket} is open.
+     *
      * @param websocket
      */
     void onOpen(WebSocket websocket);
 
     /**
      * Invoked when the {@link WebSocket} is close.
+     *
      * @param websocket
      */
     void onClose(WebSocket websocket);
 
     /**
      * Invoked when the {@link WebSocket} is open.
+     *
      * @param t a {@link Throwable}
      */
     void onError(Throwable t);

src/test/java/com/ning/http/client/websocket/netty/NettyCloseCodeReasonMessageTest.java

@@ -0,0 +1,91 @@
+/*
+ * Copyright (c) 2010-2012 Sonatype, Inc. All rights reserved.
+ *
+ * This program is licensed to you under the Apache License Version 2.0,
+ * and you may not use this file except in compliance with the Apache License Version 2.0.
+ * You may obtain a copy of the Apache License Version 2.0 at http://www.apache.org/licenses/LICENSE-2.0.
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the Apache License Version 2.0 is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the Apache License Version 2.0 for the specific language governing permissions and limitations there under.
+ */
+package com.ning.http.client.websocket.netty;
+
+import com.ning.http.client.AsyncHttpClient;
+import com.ning.http.client.AsyncHttpClientConfig;
+import com.ning.http.client.async.ProviderUtil;
+import com.ning.http.client.websocket.TextMessageTest;
+import com.ning.http.client.websocket.WebSocket;
+import com.ning.http.client.websocket.WebSocketCloseCodeReasonListener;
+import com.ning.http.client.websocket.WebSocketListener;
+import com.ning.http.client.websocket.WebSocketUpgradeHandler;
+import org.eclipse.jetty.server.nio.SelectChannelConnector;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.testng.Assert.assertEquals;
+
+public class NettyCloseCodeReasonMessageTest extends NettyTextMessageTest {
+
+    @Test(timeOut = 60000)
+    public void onCloseWithCode() throws Throwable {
+        AsyncHttpClient c = getAsyncHttpClient(new AsyncHttpClientConfig.Builder().build());
+        final CountDownLatch latch = new CountDownLatch(1);
+        final AtomicReference<String> text = new AtomicReference<String>("");
+
+        WebSocket websocket = c.prepareGet(getTargetUrl())
+                .execute(new WebSocketUpgradeHandler.Builder().addWebSocketListener(new Listener(latch, text)).build()).get();
+
+        websocket.close();
+
+        latch.await();
+        assertEquals(text.get(), "1000-Normal closure; the connection successfully completed whatever purpose for which it was created.");
+    }
+
+    @Test(timeOut = 60000)
+    public void onCloseWithCodeServerClose() throws Throwable {
+        AsyncHttpClient c = getAsyncHttpClient(new AsyncHttpClientConfig.Builder().build());
+        final CountDownLatch latch = new CountDownLatch(1);
+        final AtomicReference<String> text = new AtomicReference<String>("");
+
+        c.prepareGet(getTargetUrl())
+                .execute(new WebSocketUpgradeHandler.Builder().addWebSocketListener(new Listener(latch, text)).build()).get();
+
+        latch.await();
+        assertEquals(text.get(), "1000-Idle for 10022ms > 10000ms");
+    }
+
+    public final static class Listener implements WebSocketListener, WebSocketCloseCodeReasonListener {
+
+        final CountDownLatch latch;
+        final AtomicReference<String> text;
+
+        public Listener(CountDownLatch latch, AtomicReference<String> text) {
+            this.latch = latch;
+            this.text = text;
+        }
+
+        //@Override
+        public void onOpen(com.ning.http.client.websocket.WebSocket websocket) {
+        }
+
+        //@Override
+        public void onClose(com.ning.http.client.websocket.WebSocket websocket) {
+        }
+
+        public void onClose(WebSocket websocket, int code, String reason) {
+            text.set(code + "-" + reason);
+            latch.countDown();
+        }
+
+        //@Override
+        public void onError(Throwable t) {
+            t.printStackTrace();
+            latch.countDown();
+        }
+    }
+}