From a53442d66bbe059a89eed3b5769a7d4c11575134 Mon Sep 17 00:00:00 2001
From: Eric Larson <elarson@digitalocean.com>
Date: Sun, 6 Aug 2017 20:15:26 -0500
Subject: [PATCH] Allow defining a cache as "shared"

In the spec (https://tools.ietf.org/html/rfc7234#section-5.2.2.6) it talks
about the `private` directive and that a shared cache MUST NOT store
the response.

This starts by defining the cache as being "shared" and testing for
the "shared" attribute on the cache.

Fixes: #141
---
 cachecontrol/cache.py              |  6 ++++-
 cachecontrol/caches/file_cache.py  |  5 +++-
 cachecontrol/caches/redis_cache.py |  3 ++-
 cachecontrol/controller.py         | 13 ++++++++++
 tests/test_cache_control.py        | 40 +++++++++++++++++++++++++++++-
 5 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/cachecontrol/cache.py b/cachecontrol/cache.py
index 7389a73f..925f2187 100644
--- a/cachecontrol/cache.py
+++ b/cachecontrol/cache.py
@@ -7,6 +7,9 @@
 
 class BaseCache(object):
 
+    def __init__(self, shared=False):
+        self.shared = shared
+
     def get(self, key):
         raise NotImplemented()
 
@@ -22,7 +25,8 @@ def close(self):
 
 class DictCache(BaseCache):
 
-    def __init__(self, init_dict=None):
+    def __init__(self, init_dict=None, shared=False):
+        super(DictCache, self).__init__(shared)
         self.lock = Lock()
         self.data = init_dict or {}
 
diff --git a/cachecontrol/caches/file_cache.py b/cachecontrol/caches/file_cache.py
index d4d9e21f..3b43cf98 100644
--- a/cachecontrol/caches/file_cache.py
+++ b/cachecontrol/caches/file_cache.py
@@ -54,7 +54,10 @@ def _secure_open_write(filename, fmode):
 
 class FileCache(BaseCache):
     def __init__(self, directory, forever=False, filemode=0o0600,
-                 dirmode=0o0700, use_dir_lock=None, lock_class=None):
+                 dirmode=0o0700, use_dir_lock=None, lock_class=None,
+                 shared=False):
+
+        super(FileCache, self).__init__(shared)
 
         if use_dir_lock is not None and lock_class is not None:
             raise ValueError("Cannot use use_dir_lock and lock_class together")
diff --git a/cachecontrol/caches/redis_cache.py b/cachecontrol/caches/redis_cache.py
index a5a15aa8..7abb4029 100644
--- a/cachecontrol/caches/redis_cache.py
+++ b/cachecontrol/caches/redis_cache.py
@@ -16,7 +16,8 @@ def total_seconds(td):
 
 class RedisCache(BaseCache):
 
-    def __init__(self, conn):
+    def __init__(self, conn, shared=False):
+        super(RedisCache, self).__init__(shared)
         self.conn = conn
 
     def get(self, key):
diff --git a/cachecontrol/controller.py b/cachecontrol/controller.py
index 8436b955..cecdfd77 100644
--- a/cachecontrol/controller.py
+++ b/cachecontrol/controller.py
@@ -284,10 +284,23 @@ def cache_response(self, request, response, body=None,
         if 'no-store' in cc_req:
             no_store = True
             logger.debug('Request header has "no-store"')
+
         if no_store and self.cache.get(cache_url):
             logger.debug('Purging existing cache entry to honor "no-store"')
             self.cache.delete(cache_url)
 
+
+        # Check to see if the cache shared or not
+        is_shared_cache = False
+        try:
+            is_shared_cache = self.cache.shared
+        except AttributeError:
+            pass
+
+        if 'private' in cc and is_shared_cache is True:
+            logger.debug('Request header has "private" and the cache is shared')
+            return
+
         # If we've been given an etag, then keep the response
         if self.cache_etags and 'etag' in response_headers:
             logger.debug('Caching due to etag')
diff --git a/tests/test_cache_control.py b/tests/test_cache_control.py
index 77b8f90d..754a9b5a 100644
--- a/tests/test_cache_control.py
+++ b/tests/test_cache_control.py
@@ -2,7 +2,7 @@
 Unit tests that verify our caching methods work correctly.
 """
 import pytest
-from mock import ANY, Mock
+from mock import ANY, Mock, patch
 import time
 
 from cachecontrol import CacheController
@@ -115,6 +115,44 @@ def test_cache_response_no_store(self):
         cc.cache_response(self.req(), resp)
         assert not cc.cache.get(cache_url)
 
+    def test_cache_response_private_with_shared_cache(self):
+        '''
+        When a cache store is shared, a private directive should turn off caching.
+
+        In this example, the etag is set, which should trigger a
+        cache, but since the private directive is set and the cache is
+        considered shared, we should not cache.
+        '''
+        resp = Mock()
+        cache = DictCache(shared=True)
+        cc = CacheController(cache)
+
+        cache_url = cc.cache_url(self.url)
+
+        resp = self.resp({'cache-control': 'max-age=3600, private'})
+
+        cc.cache_response(self.req(), resp)
+        assert not cc.cache.get(cache_url)
+
+    def test_cache_response_private_with_legacy_cache(self):
+        # Not all cache objects will have the "shared" attribute.
+        resp = Mock()
+        cache = Mock()
+        cache.shared.side_effect = AttributeError
+        cc = CacheController(cache)
+        cc.serializer = Mock()
+
+        cache_url = cc.cache_url(self.url)
+
+        now = time.strftime(TIME_FMT, time.gmtime())
+        resp = self.resp({
+            'cache-control': 'max-age=3600, private',
+            'date': now,
+        })
+
+        cc.cache_response(self.req(), resp)
+        assert cc.cache.set.called
+
     def test_update_cached_response_with_valid_headers(self):
         cached_resp = Mock(headers={'ETag': 'jfd9094r808', 'Content-Length': 100})