Update CachedResponse headers with 304 response headers (RFC7234)

3 files changed, 18 insertions, 11 deletions

diff --git a/requests_cache/backends/base.py b/requests_cache/backends/base.py
index 0096d0f..ddd47e7 100644
--- a/requests_cache/backends/base.py
+++ b/requests_cache/backends/base.py
@@ -92,12 +92,8 @@ class BaseCache:
             expire_after: Time in seconds until this cache item should expire
         """
         cache_key = cache_key or self.create_key(response.request)
-        if isinstance(response, CachedResponse):
-            cached_response = response
-            cached_response.expires = expires
-        else:
-            cached_response = CachedResponse.from_response(response, expires=expires)
-            cached_response = redact_response(cached_response, self.ignored_parameters)
+        cached_response = CachedResponse.from_response(response, expires=expires)
+        cached_response = redact_response(cached_response, self.ignored_parameters)
         self.responses[cache_key] = cached_response
         for r in response.history:
             self.redirects[self.create_key(r.request)] = cache_key
diff --git a/requests_cache/models/response.py b/requests_cache/models/response.py
index 974d878..ec65034 100755
--- a/requests_cache/models/response.py
+++ b/requests_cache/models/response.py
@@ -2,6 +2,7 @@ from datetime import datetime, timedelta, timezone
 from logging import getLogger
 from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
+import attr
 from attr import define, field
 from requests import PreparedRequest, Response
 from requests.cookies import RequestsCookieJar
@@ -46,9 +47,13 @@ class CachedResponse(Response):
             self.raw.headers = HTTPHeaderDict(self.headers)
 
     @classmethod
-    def from_response(cls, original_response: Response, **kwargs):
-        """Create a CachedResponse based on an original response object"""
-        obj = cls(**kwargs)
+    def from_response(
+        cls, original_response: Union[Response, 'CachedResponse'], expires: datetime = None, **kwargs
+    ):
+        """Create a CachedResponse based on an original Response or another CachedResponse object"""
+        if isinstance(original_response, CachedResponse):
+            return attr.evolve(original_response, expires=expires)
+        obj = cls(expires=expires, **kwargs)
 
         # Copy basic attributes
         for k in Response.__attrs__:
diff --git a/requests_cache/session.py b/requests_cache/session.py
index 2468bdb..f45b7c5 100644
--- a/requests_cache/session.py
+++ b/requests_cache/session.py
@@ -183,9 +183,15 @@ class CacheMixin(MIXIN_BASE):
             logger.debug(
                 f'Response for URL {request.url} has not been modified; using cached response and updating expiration date.'
             )
-            # Update the cache expiration date. Since we performed validation,
+            # Update the cache expiration date and the headers (see RFC7234 4.3.4, p.18). Since we performed validation,
             # the cache entry may be marked as fresh again.
-            self.cache.save_response(cached_response, actions.cache_key, actions.expires)
+            cached_response.headers.update(response.headers)
+            # Since it is a 304 response we have to update cache control once again with combination of
+            # cached_response's and 304 response's Cache-Control directives.
+            response.headers = cached_response.headers
+            actions.update_from_response(response)
+            cached_response.expires = actions.expires
+            self.cache.save_response(cached_response, actions.cache_key, cached_response.expires)
             return cached_response
         else:
             logger.debug(f'Skipping cache write for URL: {request.url}')