Reuse channel on 401, close AsyncHttpClient#785

Author: Stephane Landelle

src/main/java/com/ning/http/client/providers/netty/channel/ChannelManager.java

@@ -451,7 +451,7 @@ public void call() {
         };
     }
 
-    public void drainChannel(final Channel channel, final NettyResponseFuture<?> future) {
+    public void drainChannelAndOffer(final Channel channel, final NettyResponseFuture<?> future) {
         Channels.setAttribute(channel, newDrainCallback(future, channel, future.isKeepAlive(), getPartitionId(future)));
     }
 

src/main/java/com/ning/http/client/providers/netty/handler/HttpProtocol.java

@@ -76,8 +76,7 @@ private Realm kerberosChallenge(Channel channel,//
             FluentCaseInsensitiveStringsMap headers,//
             Realm realm,//
             NettyResponseFuture<?> future,//
-            boolean proxyInd)
-            throws NTLMEngineException {
+            boolean proxyInd) throws NTLMEngineException {
 
         Uri uri = request.getUri();
         String host = request.getVirtualHost() == null ? uri.getHost() : request.getVirtualHost();
@@ -146,9 +145,8 @@ private Realm ntlmChallenge(String wwwAuth, Request request, ProxyServer proxySe
         }
     }
 
-    private Realm ntlmProxyChallenge(String wwwAuth, Request request, ProxyServer proxyServer,
-            FluentCaseInsensitiveStringsMap headers, Realm realm, NettyResponseFuture<?> future, boolean proxyInd)
-            throws NTLMEngineException {
+    private Realm ntlmProxyChallenge(String wwwAuth, Request request, ProxyServer proxyServer, FluentCaseInsensitiveStringsMap headers,
+            Realm realm, NettyResponseFuture<?> future, boolean proxyInd) throws NTLMEngineException {
         future.getAndSetAuth(false);
         headers.remove(HttpHeaders.Names.PROXY_AUTHORIZATION);
 
@@ -161,8 +159,8 @@ private Realm ntlmProxyChallenge(String wwwAuth, Request request, ProxyServer pr
                 .setMethodName(request.getMethod()).build();
     }
 
-    private void addType3NTLMAuthorizationHeader(String auth, FluentCaseInsensitiveStringsMap headers, String username,
-            String password, String domain, String workstation, boolean proxyInd) throws NTLMEngineException {
+    private void addType3NTLMAuthorizationHeader(String auth, FluentCaseInsensitiveStringsMap headers, String username, String password,
+            String domain, String workstation, boolean proxyInd) throws NTLMEngineException {
         headers.remove(authorizationHeaderName(proxyInd));
 
         if (isNonEmpty(auth) && auth.startsWith("NTLM ")) {
@@ -176,7 +174,7 @@ private void finishUpdate(final NettyResponseFuture<?> future, Channel channel,
 
         boolean keepAlive = future.isKeepAlive();
         if (expectOtherChunks && keepAlive)
-            channelManager.drainChannel(channel, future);
+            channelManager.drainChannelAndOffer(channel, future);
         else
             channelManager.tryToOfferChannelToPool(channel, keepAlive, channelManager.getPartitionId(future));
         markAsDone(future, channel);
@@ -228,17 +226,11 @@ private boolean exitAfterHandling401(//
                     // NTLM
                     newRealm = ntlmChallenge(ntlmAuthenticate, request, proxyServer, request.getHeaders(), realm, future, false);
 
-                    // don't forget to reuse channel: NTLM authenticates a connection
-                    future.setReuseChannel(true);
-
                 } else if (negociate) {
                     // SPNEGO KERBEROS
                     newRealm = kerberosChallenge(channel, wwwAuthHeaders, request, proxyServer, request.getHeaders(), realm, future, false);
                     if (newRealm == null)
                         return true;
-                    else
-                        // don't forget to reuse channel: KERBEROS authenticates a connection
-                        future.setReuseChannel(true);
 
                 } else {
                     newRealm = new Realm.RealmBuilder()//
@@ -255,20 +247,22 @@ private boolean exitAfterHandling401(//
                 final Request nextRequest = new RequestBuilder(future.getRequest()).setHeaders(request.getHeaders()).setRealm(nr).build();
 
                 logger.debug("Sending authentication to {}", request.getUri());
-                Callback callback = new Callback(future) {
-                    public void call() throws IOException {
-                        channelManager.drainChannel(channel, future);
+                if (future.isKeepAlive()) {
+                    if (HttpHeaders.isTransferEncodingChunked(response)) {
+                        // we must first drain the channel, let's use a fresh one for performing auth
+                        Channels.setAttribute(channel, new Callback(future) {
+                            public void call() throws IOException {
+                                requestSender.drainChannelAndExecuteNextRequest(channel, future, nextRequest);
+                            }
+                        });
+                    } else {
+                        future.setReuseChannel(true);
                         requestSender.sendNextRequest(nextRequest, future);
                     }
-                };
-
-                if (future.isKeepAlive() && HttpHeaders.isTransferEncodingChunked(response))
-                    // We must make sure there is no bytes left
-                    // before executing the next request.
-                    Channels.setAttribute(channel, callback);
-                else
-                    // call might crash with an IOException
-                    callback.call();
+                } else {
+                    channelManager.closeChannel(channel);
+                    requestSender.sendNextRequest(nextRequest, future);
+                }
 
                 return true;
             }
@@ -374,8 +368,8 @@ private boolean exitAfterHandlingConnect(//
         return false;
     }
 
-    private boolean exitAfterHandlingStatus(Channel channel, NettyResponseFuture<?> future, HttpResponse response,
-            AsyncHandler<?> handler, NettyResponseStatus status) throws IOException, Exception {
+    private boolean exitAfterHandlingStatus(Channel channel, NettyResponseFuture<?> future, HttpResponse response, AsyncHandler<?> handler,
+            NettyResponseStatus status) throws IOException, Exception {
         if (!future.getAndSetStatusReceived(true) && handler.onStatusReceived(status) != STATE.CONTINUE) {
             finishUpdate(future, channel, HttpHeaders.isTransferEncodingChunked(response));
             return true;
@@ -393,8 +387,8 @@ private boolean exitAfterHandlingHeaders(Channel channel, NettyResponseFuture<?>
     }
 
     // Netty 3: if the response is not chunked, the full body comes with the response
-    private boolean exitAfterHandlingBody(Channel channel, NettyResponseFuture<?> future, HttpResponse response,
-            AsyncHandler<?> handler) throws Exception {
+    private boolean exitAfterHandlingBody(Channel channel, NettyResponseFuture<?> future, HttpResponse response, AsyncHandler<?> handler)
+            throws Exception {
         if (!response.isChunked()) {
             // no chunks expected, exiting
             if (response.getContent().readableBytes() > 0)

src/main/java/com/ning/http/client/providers/netty/request/NettyRequestSender.java

@@ -46,6 +46,7 @@
 import com.ning.http.client.filter.FilterException;
 import com.ning.http.client.filter.IOExceptionFilter;
 import com.ning.http.client.listener.TransferCompletionHandler;
+import com.ning.http.client.providers.netty.Callback;
 import com.ning.http.client.providers.netty.NettyAsyncHttpProviderConfig;
 import com.ning.http.client.providers.netty.channel.ChannelManager;
 import com.ning.http.client.providers.netty.channel.Channels;
@@ -529,12 +530,26 @@ public void replayRequest(final NettyResponseFuture<?> future, FilterContext fc,
         if (future.getAsyncHandler() instanceof AsyncHandlerExtensions)
             AsyncHandlerExtensions.class.cast(future.getAsyncHandler()).onRetry();
 
-        channelManager.drainChannel(channel, future);
+        channelManager.drainChannelAndOffer(channel, future);
         sendNextRequest(newRequest, future);
         return;
     }
 
     public boolean isClosed() {
         return closed.get();
     }
+   
+    public final Callback newExecuteNextRequestCallback(final NettyResponseFuture<?> future, final Request nextRequest) {
+
+        return new Callback(future) {
+            @Override
+            public void call() throws IOException {
+                sendNextRequest(nextRequest, future);
+            }
+        };
+    }
+    
+    public void drainChannelAndExecuteNextRequest(final Channel channel, final NettyResponseFuture<?> future, Request nextRequest) {
+        Channels.setAttribute(channel, newExecuteNextRequestCallback(future, nextRequest));
+    }
 }

src/test/java/com/ning/http/client/async/BasicAuthTest.java

@@ -24,7 +24,6 @@
 import com.ning.http.client.Realm;
 import com.ning.http.client.Realm.AuthScheme;
 import com.ning.http.client.Response;
-import com.ning.http.client.generators.InputStreamBodyGenerator;
 
 import org.apache.log4j.ConsoleAppender;
 import org.apache.log4j.Level;