Fixes AsyncHttpClient#110

Author: jfarcand

api/src/main/java/org/asynchttpclient/websocket/WebSocketUpgradeHandler.java

@@ -36,6 +36,7 @@ public class WebSocketUpgradeHandler implements UpgradeHandler<WebSocket>, Async
  private final long maxTextSize;
  private final AtomicBoolean ok = new AtomicBoolean(false);
  private final AtomicBoolean onSuccessCalled = new AtomicBoolean(false);
+ private int status;
 
  private WebSocketUpgradeHandler(Builder b) {
  l = b.l;
@@ -52,7 +53,7 @@ public final void onThrowable(Throwable t) {
  onFailure(t);
  }
 
- public boolean touchSuccess(){
+ public boolean touchSuccess() {
  return onSuccessCalled.getAndSet(true);
  }
 
@@ -69,6 +70,7 @@ public final STATE onBodyPartReceived(HttpResponseBodyPart bodyPart) throws Exce
  */
  @Override
  public final STATE onStatusReceived(HttpResponseStatus responseStatus) throws Exception {
+ status = responseStatus.getStatusCode();
  if (responseStatus.getStatusCode() == 101) {
  return STATE.UPGRADE;
  } else {
@@ -89,6 +91,14 @@ public final STATE onHeadersReceived(HttpResponseHeaders headers) throws Excepti
  */
  @Override
  public final WebSocket onCompleted() throws Exception {
+
+ if (status != 101) {
+ for (WebSocketListener w : l) {
+ w.onError(new IllegalStateException(String.format("Invalid Status Code %d", status)));
+ }
+ return null;
+ }
+
  if (webSocket == null) {
  throw new IllegalStateException("WebSocket is null");
  }
@@ -202,6 +212,7 @@ public Builder setMaxTextSize(long maxTextSize) {
 
  /**
  * Build a {@link WebSocketUpgradeHandler}
+ *
  * @return a {@link WebSocketUpgradeHandler}
  */
  public WebSocketUpgradeHandler build() {

api/src/test/java/org/asynchttpclient/websocket/CloseCodeReasonMessageTest.java

@@ -12,16 +12,18 @@
  */
 package org.asynchttpclient.websocket;
 
-import static org.testng.Assert.*;
-
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.atomic.AtomicReference;
-
 import org.asynchttpclient.AsyncHttpClient;
 import org.eclipse.jetty.websocket.server.WebSocketHandler;
 import org.eclipse.jetty.websocket.servlet.WebSocketServletFactory;
 import org.testng.annotations.Test;
 
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+
 public abstract class CloseCodeReasonMessageTest extends AbstractBasicTest {
 
  @Override
@@ -99,4 +101,45 @@ public void onError(Throwable t) {
  latch.countDown();
  }
  }
+
+
+ @Test(timeOut = 60000)
+ public void wrongStatusCode() throws Throwable {
+ AsyncHttpClient c = getAsyncHttpClient(null);
+ try {
+ final CountDownLatch latch = new CountDownLatch(1);
+ final AtomicReference<Throwable> throwable = new AtomicReference<Throwable>();
+
+ WebSocket websocket = c.prepareGet("http://apache.org").execute(new WebSocketUpgradeHandler.Builder().addWebSocketListener(new WebSocketTextListener() {
+
+ @Override
+ public void onMessage(String message) {
+ }
+
+ @Override
+ public void onFragment(String fragment, boolean last) {
+ }
+
+ @Override
+ public void onOpen(WebSocket websocket) {
+ }
+
+ @Override
+ public void onClose(WebSocket websocket) {
+ }
+
+ @Override
+ public void onError(Throwable t) {
+ throwable.set(t);
+ latch.countDown();
+ }
+ }).build()).get();
+
+ latch.await();
+ assertNotNull(throwable.get());
+ assertEquals(throwable.get().getClass(), IllegalStateException.class);
+ } finally {
+ c.close();
+ }
+ }
 }