Enable cookie store at request level

Author: tranchitam

client/src/main/java/org/asynchttpclient/DefaultRequest.java

@@ -17,6 +17,7 @@
 import io.netty.handler.codec.http.cookie.Cookie;
 import io.netty.resolver.NameResolver;
 import org.asynchttpclient.channel.ChannelPoolPartitioning;
+import org.asynchttpclient.cookie.CookieStore;
 import org.asynchttpclient.proxy.ProxyServer;
 import org.asynchttpclient.request.body.generator.BodyGenerator;
 import org.asynchttpclient.request.body.multipart.Part;
@@ -43,6 +44,7 @@ public class DefaultRequest implements Request {
   private final InetAddress localAddress;
   private final HttpHeaders headers;
   private final List<Cookie> cookies;
+  private final CookieStore cookieStore;
   private final byte[] byteData;
   private final List<byte[]> compositeByteData;
   private final String stringData;
@@ -70,6 +72,7 @@ public DefaultRequest(String method,
                         InetAddress localAddress,
                         HttpHeaders headers,
                         List<Cookie> cookies,
+                        CookieStore cookieStore,
                         byte[] byteData,
                         List<byte[]> compositeByteData,
                         String stringData,
@@ -95,6 +98,7 @@ public DefaultRequest(String method,
     this.localAddress = localAddress;
     this.headers = headers;
     this.cookies = cookies;
+    this.cookieStore = cookieStore;
     this.byteData = byteData;
     this.compositeByteData = compositeByteData;
     this.stringData = stringData;
@@ -150,6 +154,11 @@ public HttpHeaders getHeaders() {
   public List<Cookie> getCookies() {
     return cookies;
   }
+  
+  @Override
+  public CookieStore getCookieStore() {
+    return cookieStore;
+  }
 
   @Override
   public byte[] getByteData() {

client/src/main/java/org/asynchttpclient/Request.java

@@ -20,6 +20,7 @@
 import io.netty.handler.codec.http.cookie.Cookie;
 import io.netty.resolver.NameResolver;
 import org.asynchttpclient.channel.ChannelPoolPartitioning;
+import org.asynchttpclient.cookie.CookieStore;
 import org.asynchttpclient.proxy.ProxyServer;
 import org.asynchttpclient.request.body.generator.BodyGenerator;
 import org.asynchttpclient.request.body.multipart.Part;
@@ -80,6 +81,11 @@ public interface Request {
    * @return the HTTP cookies
    */
   List<Cookie> getCookies();
+  
+  /**
+   * @return the cookie store
+   */
+  CookieStore getCookieStore();
 
   /**
    * @return the request's body byte array (only non null if it was set this way)

client/src/main/java/org/asynchttpclient/RequestBuilderBase.java

@@ -23,6 +23,7 @@
 import io.netty.resolver.NameResolver;
 import io.netty.util.concurrent.ImmediateEventExecutor;
 import org.asynchttpclient.channel.ChannelPoolPartitioning;
+import org.asynchttpclient.cookie.CookieStore;
 import org.asynchttpclient.proxy.ProxyServer;
 import org.asynchttpclient.request.body.generator.BodyGenerator;
 import org.asynchttpclient.request.body.generator.ReactiveStreamsBodyGenerator;
@@ -68,6 +69,7 @@ public abstract class RequestBuilderBase<T extends RequestBuilderBase<T>> {
   protected InetAddress localAddress;
   protected HttpHeaders headers;
   protected ArrayList<Cookie> cookies;
+  protected CookieStore cookieStore;
   protected byte[] byteData;
   protected List<byte[]> compositeByteData;
   protected String stringData;
@@ -113,6 +115,7 @@ protected RequestBuilderBase(Request prototype, boolean disableUrlEncoding, bool
     if (isNonEmpty(prototype.getCookies())) {
       this.cookies = new ArrayList<>(prototype.getCookies());
     }
+    this.cookieStore = prototype.getCookieStore();
     this.byteData = prototype.getByteData();
     this.compositeByteData = prototype.getCompositeByteData();
     this.stringData = prototype.getStringData();
@@ -335,6 +338,11 @@ public void resetCookies() {
     if (this.cookies != null)
       this.cookies.clear();
   }
+  
+  public T setCookieStore(CookieStore cookieStore) {
+    this.cookieStore = cookieStore;
+    return asDerivedType();
+  }
 
   public void resetQuery() {
     queryParams = null;
@@ -580,6 +588,7 @@ private RequestBuilderBase<?> executeSignatureCalculator() {
     rb.charset = this.charset;
     rb.channelPoolPartitioning = this.channelPoolPartitioning;
     rb.nameResolver = this.nameResolver;
+    rb.cookieStore = this.cookieStore;
     Request unsignedRequest = rb.build();
     signatureCalculator.calculateAndAddSignature(unsignedRequest, rb);
     return rb;
@@ -624,6 +633,7 @@ public Request build() {
             rb.localAddress,
             rb.headers,
             cookiesCopy,
+            rb.cookieStore,
             rb.byteData,
             rb.compositeByteData,
             rb.stringData,

client/src/main/java/org/asynchttpclient/netty/handler/intercept/Interceptors.java

@@ -71,12 +71,19 @@ public boolean exitAfterIntercept(Channel channel,
 
     // This MUST BE called before Redirect30xInterceptor because latter assumes cookie store is already updated
     CookieStore cookieStore = config.getCookieStore();
-    if (cookieStore != null) {
+    CookieStore requestCookieStore = request.getCookieStore();
+    if (cookieStore != null || requestCookieStore != null) {
       for (String cookieStr : responseHeaders.getAll(SET_COOKIE)) {
         Cookie c = cookieDecoder.decode(cookieStr);
         if (c != null) {
           // Set-Cookie header could be invalid/malformed
-          cookieStore.add(future.getCurrentRequest().getUri(), c);
+          if (cookieStore != null) {
+            cookieStore.add(future.getCurrentRequest().getUri(), c);
+          }
+
+          if (requestCookieStore != null) {
+            cookieStore.add(request.getUri(), c);
+          }
         }
       }
     }

client/src/main/java/org/asynchttpclient/netty/handler/intercept/Redirect30xInterceptor.java

@@ -91,6 +91,7 @@ public boolean exitAfterHandlingRedirect(Channel channel,
         boolean keepBody = statusCode == TEMPORARY_REDIRECT_307 || statusCode == PERMANENT_REDIRECT_308 || (statusCode == FOUND_302 && config.isStrict302Handling());
 
         final RequestBuilder requestBuilder = new RequestBuilder(switchToGet ? GET : originalMethod)
+                .setCookieStore(request.getCookieStore())
                 .setChannelPoolPartitioning(request.getChannelPoolPartitioning())
                 .setFollowRedirect(true)
                 .setLocalAddress(request.getLocalAddress())
@@ -126,7 +127,8 @@ else if (request.getBodyGenerator() != null)
 
         LOGGER.debug("Redirecting to {}", newUri);
 
-        CookieStore cookieStore = config.getCookieStore();
+        CookieStore cookieStore =
+            request.getCookieStore() != null ? request.getCookieStore() : config.getCookieStore();
         if (cookieStore != null) {
           // Update request's cookies assuming that cookie store is already updated by Interceptors
           List<Cookie> cookies = cookieStore.get(newUri);