1"""
  2This module contains helper functions for controlling caching. It does so by
  3managing the "Vary" header of responses. It includes functions to patch the
  4header of response objects directly and decorators that change functions to do
  5that header-patching themselves.
  6
  7For information on the Vary header, see RFC 9110 Section 12.5.5.
  8
  9Essentially, the "Vary" HTTP header defines which headers a cache should take
 10into account when building its cache key. Requests with the same path but
 11different header content for headers named in "Vary" need to get different
 12cache keys to prevent delivery of wrong content.
 13
 14An example: i18n middleware would need to distinguish caches by the
 15"Accept-language" header.
 16"""
 17
 18import time
 19from collections import defaultdict
 20from typing import TYPE_CHECKING, Any
 21
 22from .http import http_date
 23from .regex_helper import _lazy_re_compile
 24
 25if TYPE_CHECKING:
 26    from plain.http import Response
 27
 28_cc_delim_re = _lazy_re_compile(r"\s*,\s*")
 29
 30
 31def patch_response_headers(response: Response, cache_timeout: float) -> None:
 32    """
 33    Add HTTP caching headers to the given HttpResponse: Expires and
 34    Cache-Control.
 35
 36    Each header is only added if it isn't already set.
 37    """
 38    cache_timeout = max(cache_timeout, 0)  # Can't have max-age negative
 39    if "Expires" not in response.headers:
 40        response.headers["Expires"] = http_date(time.time() + cache_timeout)
 41    patch_cache_control(response, max_age=int(cache_timeout))
 42
 43
 44def add_never_cache_headers(response: Response) -> None:
 45    """
 46    Add headers to a response to indicate that a page should never be cached.
 47    """
 48    patch_response_headers(response, cache_timeout=-1)
 49    patch_cache_control(
 50        response, no_cache=True, no_store=True, must_revalidate=True, private=True
 51    )
 52
 53
 54def patch_cache_control(
 55    response: Response,
 56    *,
 57    max_age: int | None = None,
 58    s_maxage: int | None = None,
 59    stale_while_revalidate: int | None = None,
 60    stale_if_error: int | None = None,
 61    no_cache: bool = False,
 62    no_store: bool = False,
 63    no_transform: bool = False,
 64    must_revalidate: bool = False,
 65    proxy_revalidate: bool = False,
 66    must_understand: bool = False,
 67    public: bool = False,
 68    private: bool = False,
 69    immutable: bool = False,
 70) -> None:
 71    """
 72    Patch the Cache-Control header on the response, merging the given
 73    directives with any already present.
 74
 75    Parameter names map to their hyphenated Cache-Control directives
 76    (``max_age`` -> ``max-age``). Integer directives are emitted as
 77    ``name=value`` when set; boolean directives are emitted as a bare
 78    directive name when True. If ``max-age`` is already present, the smaller of
 79    the existing and new value wins (this happens when a decorator and a piece
 80    of middleware both operate on a given view). Setting ``public`` clears an
 81    existing ``private`` directive and vice versa.
 82    """
 83
 84    def dictitem(s: str) -> tuple[str, str | bool]:
 85        t = s.split("=", 1)
 86        if len(t) > 1:
 87            return (t[0].lower(), t[1])
 88        else:
 89            return (t[0].lower(), True)
 90
 91    def dictvalue(directive: str, value: str | bool) -> str:
 92        if value is True:
 93            return directive
 94        else:
 95            return f"{directive}={value}"
 96
 97    # Collect the requested directives under their hyphenated names.
 98    new: dict[str, int | bool] = {}
 99    for directive, number in (
100        ("max-age", max_age),
101        ("s-maxage", s_maxage),
102        ("stale-while-revalidate", stale_while_revalidate),
103        ("stale-if-error", stale_if_error),
104    ):
105        if number is not None:
106            new[directive] = number
107    for directive, flag in (
108        ("no-cache", no_cache),
109        ("no-store", no_store),
110        ("no-transform", no_transform),
111        ("must-revalidate", must_revalidate),
112        ("proxy-revalidate", proxy_revalidate),
113        ("must-understand", must_understand),
114        ("public", public),
115        ("private", private),
116        ("immutable", immutable),
117    ):
118        if flag:
119            new[directive] = True
120
121    cc: defaultdict[str, Any] = defaultdict(set)
122    if response.headers.get("Cache-Control"):
123        for field in _cc_delim_re.split(response.headers["Cache-Control"]):
124            directive, value = dictitem(field)
125            if directive == "no-cache":
126                # no-cache supports multiple field names.
127                cc[directive].add(value)
128            else:
129                cc[directive] = value
130
131    # If there's already a max-age header but we're being asked to set a new
132    # max-age, use the minimum of the two ages.
133    if "max-age" in cc and "max-age" in new:
134        new["max-age"] = min(int(cc["max-age"]), int(new["max-age"]))
135
136    # Allow overriding private caching and vice versa.
137    if "private" in cc and "public" in new:
138        del cc["private"]
139    elif "public" in cc and "private" in new:
140        del cc["public"]
141
142    for directive, value in new.items():
143        if directive == "no-cache":
144            # no-cache supports multiple field names.
145            cc[directive].add(value)
146        else:
147            cc[directive] = value
148
149    directives = []
150    for directive, values in cc.items():
151        if isinstance(values, set):
152            if True in values:
153                # True takes precedence.
154                values = {True}
155            directives.extend([dictvalue(directive, value) for value in values])
156        else:
157            directives.append(dictvalue(directive, values))
158    response.headers["Cache-Control"] = ", ".join(directives)
159
160
161def patch_vary_headers(response: Response, newheaders: list[str]) -> None:
162    """
163    Add (or update) the "Vary" header in the given Response object.
164    newheaders is a list of header names that should be in "Vary". If headers
165    contains an asterisk, then "Vary" header will consist of a single asterisk
166    '*'. Otherwise, existing headers in "Vary" aren't removed.
167    """
168    # Note that we need to keep the original order intact, because cache
169    # implementations may rely on the order of the Vary contents in, say,
170    # computing an MD5 hash.
171    if "Vary" in response.headers:
172        vary_headers = _cc_delim_re.split(response.headers["Vary"])
173    else:
174        vary_headers = []
175    # Use .lower() here so we treat headers as case-insensitive.
176    existing_headers = {header.lower() for header in vary_headers}
177    additional_headers = [
178        newheader
179        for newheader in newheaders
180        if newheader.lower() not in existing_headers
181    ]
182    vary_headers += additional_headers
183    if "*" in vary_headers:
184        response.headers["Vary"] = "*"
185    else:
186        response.headers["Vary"] = ", ".join(vary_headers)