From d9684e913b98d13d99ad60f74a4efd7276ac87b2 Mon Sep 17 00:00:00 2001
From: Hunter Stern <hunter@archive.org>
Date: Fri, 16 Oct 2015 14:05:56 -0700
Subject: [PATCH] Restrict number of cookies per domain to 50

---
 .../modules/fetcher/BdbCookieStore.java       | 23 +++++++++++++++++++
 .../modules/fetcher/CookieStoreTest.java      | 23 +++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/modules/src/main/java/org/archive/modules/fetcher/BdbCookieStore.java b/modules/src/main/java/org/archive/modules/fetcher/BdbCookieStore.java
index 0cf5c34e..13efc983 100644
--- a/modules/src/main/java/org/archive/modules/fetcher/BdbCookieStore.java
+++ b/modules/src/main/java/org/archive/modules/fetcher/BdbCookieStore.java
@@ -26,6 +26,8 @@ import java.util.Iterator;
 import java.util.List;
 import java.util.ListIterator;
 import java.util.SortedMap;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 
 import org.apache.commons.collections.collection.CompositeCollection;
 import org.apache.http.client.CookieStore;
@@ -59,6 +61,10 @@ import com.sleepycat.je.DatabaseException;
 public class BdbCookieStore extends AbstractCookieStore implements
         FetchHTTPCookieStore, CookieStore {
 
+    public static final int MAX_COOKIES_FOR_DOMAIN = 50;
+    
+    private static Logger logger = Logger.getLogger(BdbCookieStore.class.getName());
+    
     /**
      * A {@link List} implementation that wraps a {@link Collection}. Needed
      * because httpclient requires {@code List<Cookie>}.
@@ -130,6 +136,17 @@ public class BdbCookieStore extends AbstractCookieStore implements
     }
 
     public void addCookie(Cookie cookie) {
+        
+        if (isCookieCountMaxedForDomain(cookie.getDomain())) {
+            logger.log(
+                    Level.FINEST,
+                    "Maximum number of cookies reached for domain "
+                            + cookie.getDomain() + ". Will not add new cookie "
+                            + cookie.getName() + " with value "
+                            + cookie.getValue());
+            return;
+        }
+        
         byte[] key;
         try {
             key = sortableKey(cookie).getBytes("UTF-8");
@@ -143,6 +160,12 @@ public class BdbCookieStore extends AbstractCookieStore implements
             cookies.remove(key);
         }
     }
+    
+    protected boolean isCookieCountMaxedForDomain(String domain) {
+        Collection<Cookie> subset = hostSubset(normalizeHost(domain));
+        
+        return (subset != null && subset.size() >= MAX_COOKIES_FOR_DOMAIN);
+    }
 
     protected Collection<Cookie> hostSubset(String host) { 
         try {
diff --git a/modules/src/test/java/org/archive/modules/fetcher/CookieStoreTest.java b/modules/src/test/java/org/archive/modules/fetcher/CookieStoreTest.java
index e0b311cc..3b8f5226 100644
--- a/modules/src/test/java/org/archive/modules/fetcher/CookieStoreTest.java
+++ b/modules/src/test/java/org/archive/modules/fetcher/CookieStoreTest.java
@@ -163,6 +163,25 @@ public class CookieStoreTest extends TmpDirTestCase {
         basicCookieStore().addCookie(cookie);
         assertCookieStoresEquivalent(basicCookieStore(), bdbCookieStore());
     }
+    
+    public void testMaxCookieDomain() throws IOException {
+        bdbCookieStore().clear();
+
+        for (int i = 1; i <= BdbCookieStore.MAX_COOKIES_FOR_DOMAIN; ++i) {
+            BasicClientCookie cookie = new BasicClientCookie("name" + i, "value" + i);
+            bdbCookieStore().addCookie(cookie);
+            
+            assertCookieStoreCountEquals(bdbCookieStore, i);
+        }
+        
+        BasicClientCookie cookie = new BasicClientCookie("nametoomany1", "valuetoomany1");
+        bdbCookieStore().addCookie(cookie);
+        assertCookieStoreCountEquals(bdbCookieStore, BdbCookieStore.MAX_COOKIES_FOR_DOMAIN);
+        
+        cookie = new BasicClientCookie("nametoomany2", "valuetoomany2");
+        bdbCookieStore().addCookie(cookie);
+        assertCookieStoreCountEquals(bdbCookieStore, BdbCookieStore.MAX_COOKIES_FOR_DOMAIN);
+    }
 
     public void testPaths() throws IOException {
         bdbCookieStore().clear();
@@ -318,6 +337,10 @@ public class CookieStoreTest extends TmpDirTestCase {
         assertCookieListsEquivalent(bdbCookieList, basicCookieStore().getCookies());
     }
 
+    protected void assertCookieStoreCountEquals(BdbCookieStore bdb, int count) {
+        assertEquals(bdb.getCookies().size(), count);
+    }
+    
     protected void assertCookieListsEquivalent(List<Cookie> list1,
             List<Cookie> list2) {
         Comparator<Cookie> comparator = new Comparator<Cookie>() {