Support for cookies in okhttp implementation, fixes #632

apache · Oct 9, 2018 · 4c6a274 · 4c6a274
1 parent 214c7bc
commit 4c6a274
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 2 deletions.
diff --git a/core/src/main/java/com/digitalpebble/stormcrawler/protocol/AbstractHttpProtocol.java b/core/src/main/java/com/digitalpebble/stormcrawler/protocol/AbstractHttpProtocol.java
@@ -47,6 +47,8 @@ public abstract class AbstractHttpProtocol implements Protocol {
 
     protected boolean useCookies = false;
 
+    protected static final String RESPONSE_COOKIES_HEADER = "set-cookie";
+
     @Override
     public void configure(Config conf) {
         this.skipRobots = ConfUtils.getBoolean(conf, "http.skip.robots", false);

diff --git a/core/src/main/java/com/digitalpebble/stormcrawler/protocol/httpclient/HttpProtocol.java b/core/src/main/java/com/digitalpebble/stormcrawler/protocol/httpclient/HttpProtocol.java
@@ -82,8 +82,6 @@ public class HttpProtocol extends AbstractHttpProtocol implements
 
     private RequestConfig requestConfig;
 
-    public static final String RESPONSE_COOKIES_HEADER = "set-cookie";
-
     @Override
     public void configure(final Config conf) {
 

diff --git a/core/src/main/java/com/digitalpebble/stormcrawler/protocol/okhttp/HttpProtocol.java b/core/src/main/java/com/digitalpebble/stormcrawler/protocol/okhttp/HttpProtocol.java
@@ -20,7 +20,9 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.InetSocketAddress;
+import java.net.MalformedURLException;
 import java.net.Proxy;
+import java.net.URL;
 import java.security.cert.CertificateException;
 import java.util.Base64;
 import java.util.LinkedList;
@@ -37,6 +39,7 @@
 
 import org.apache.commons.lang.StringUtils;
 import org.apache.commons.lang.mutable.MutableBoolean;
+import org.apache.http.cookie.Cookie;
 import org.apache.http.util.ByteArrayBuffer;
 import org.apache.storm.Config;
 import org.slf4j.LoggerFactory;
@@ -45,6 +48,7 @@
 import com.digitalpebble.stormcrawler.protocol.AbstractHttpProtocol;
 import com.digitalpebble.stormcrawler.protocol.ProtocolResponse;
 import com.digitalpebble.stormcrawler.util.ConfUtils;
+import com.digitalpebble.stormcrawler.util.CookieConverter;
 
 import okhttp3.Call;
 import okhttp3.Headers;
@@ -168,6 +172,21 @@ public boolean verify(String hostname, SSLSession session) {
         client = builder.build();
     }
 
+    private void addCookiesToRequest(Builder rb, String url, Metadata md) {
+        String[] cookieStrings = md.getValues(RESPONSE_COOKIES_HEADER);
+        if (cookieStrings == null || cookieStrings.length == 0) {
+            return;
+        }
+        try {
+            List<Cookie> cookies = CookieConverter.getCookies(cookieStrings,
+                    new URL(url));
+            for (Cookie c : cookies) {
+                rb.addHeader("Cookie", c.getName() + "=" + c.getValue());
+            }
+        } catch (MalformedURLException e) { // Bad url , nothing to do
+        }
+    }
+
     @Override
     public ProtocolResponse getProtocolOutput(String url, final Metadata metadata) throws Exception {
         Builder rb = new Request.Builder().url(url);
@@ -186,6 +205,10 @@ public ProtocolResponse getProtocolOutput(String url, final Metadata metadata) t
             if (StringUtils.isNotBlank(ifNoneMatch)) {
                 rb.header("If-None-Match", ifNoneMatch);
             }
+
+            if (useCookies) {
+                addCookiesToRequest(rb, url, metadata);
+            }
         }
 
         Request request = rb.build();