1package httpcache
  2
  3import (
  4	"encoding/binary"
  5	"encoding/json"
  6	"errors"
  7	"net/http"
  8	"strconv"
  9	"strings"
 10	"time"
 11)
 12
 13// responseWriter buffers an upstream response so it can be stored. A body
 14// larger than maxBody is not stored: what was buffered is sent and the rest
 15// streams straight through.
 16type responseWriter struct {
 17	http.ResponseWriter
 18	statusCode int
 19	body       []byte
 20	maxBody    int64
 21	streaming  bool
 22	// streamStatus is the cache-status sent with a streamed response.
 23	streamStatus string
 24}
 25
 26func (rw *responseWriter) WriteHeader(code int) {
 27	rw.statusCode = code
 28	size, err := strconv.ParseInt(rw.Header().Get("content-length"), 10, 64)
 29	if err == nil && rw.tooLarge(size) {
 30		_ = rw.stream()
 31	}
 32}
 33
 34func (rw *responseWriter) Write(data []byte) (int, error) {
 35	if !rw.streaming && rw.tooLarge(int64(len(rw.body)+len(data))) {
 36		if err := rw.stream(); err != nil {
 37			return 0, err
 38		}
 39	}
 40	if rw.streaming {
 41		return rw.ResponseWriter.Write(data)
 42	}
 43	rw.body = append(rw.body, data...)
 44	return len(data), nil
 45}
 46
 47// Flush passes through once the response is streaming. Flushing a buffered
 48// response would send its headers before it is known whether it is stored.
 49func (rw *responseWriter) Flush() {
 50	if rw.streaming {
 51		_ = http.NewResponseController(rw.ResponseWriter).Flush()
 52	}
 53}
 54
 55func (rw *responseWriter) tooLarge(size int64) bool {
 56	return rw.maxBody > 0 && size > rw.maxBody
 57}
 58
 59func (rw *responseWriter) stream() error {
 60	if rw.streaming {
 61		return nil
 62	}
 63	rw.streaming = true
 64	rw.Header().Set("cache-status", rw.streamStatus)
 65	rw.ResponseWriter.WriteHeader(rw.StatusCode())
 66	body := rw.body
 67	rw.body = nil
 68	if len(body) == 0 {
 69		return nil
 70	}
 71	_, err := rw.ResponseWriter.Write(body)
 72	return err
 73}
 74
 75// Body returns the captured response body.
 76func (rw *responseWriter) Body() []byte {
 77	return rw.body
 78}
 79
 80// StatusCode returns the captured status code.
 81func (rw *responseWriter) StatusCode() int {
 82	if rw.statusCode == 0 {
 83		return http.StatusOK
 84	}
 85	return rw.statusCode
 86}
 87
 88func (rw *responseWriter) Send() {
 89	rw.ResponseWriter.WriteHeader(rw.StatusCode())
 90	// RFC 9110 15.4.5: 304 responses MUST NOT contain a body.
 91	if rw.StatusCode() == http.StatusNotModified {
 92		return
 93	}
 94	_, _ = rw.ResponseWriter.Write(rw.body)
 95}
 96
 97func (rw *responseWriter) ToCacheValue(r *http.Request) *CacheValue {
 98	// Normalize header keys to lowercase to avoid case-sensitivity issues
 99	// in the cached map (e.g., "ETag" vs "Etag" as separate keys).
100	headers := make(map[string][]string)
101	for k, v := range rw.Header() {
102		headers[strings.ToLower(k)] = v
103	}
104
105	// Snapshot the request header values named by the response Vary header so
106	// matchVary can compare them on future cache lookups (Vary lists request
107	// header names, not response header names).
108	varyReqHdrs := make(map[string]string)
109	if vary := headers["vary"]; len(vary) > 0 {
110		for _, field := range strings.FieldsFunc(vary[0], func(c rune) bool { return c == ',' }) {
111			field = strings.TrimSpace(strings.ToLower(field))
112			if field != "" && field != "*" {
113				varyReqHdrs[field] = r.Header.Get(field)
114			}
115		}
116	}
117
118	cv := &CacheValue{
119		Header:             headers,
120		Body:               rw.body,
121		CreatedAt:          time.Now(),
122		StatusCode:         rw.StatusCode(),
123		VaryRequestHeaders: varyReqHdrs,
124	}
125
126	return cv
127}
128
129type CacheValue struct {
130	Header             map[string][]string `json:"headers"`
131	Body               []byte              `json:"-"`
132	CreatedAt          time.Time           `json:"created_at"`
133	StatusCode         int                 `json:"status_code"`
134	VaryRequestHeaders map[string]string   `json:"vary_request_headers,omitempty"`
135}
136
137// encode stores the body after the JSON metadata so a cache hit decodes only
138// the metadata.
139func (cv *CacheValue) encode() []byte {
140	meta, _ := json.Marshal(cv)
141	out := make([]byte, 0, binary.MaxVarintLen64+len(meta)+len(cv.Body))
142	out = binary.AppendUvarint(out, uint64(len(meta)))
143	out = append(out, meta...)
144	return append(out, cv.Body...)
145}
146
147// decodeCacheValue reads an entry written by encode. The body shares memory
148// with data, which the cache never modifies.
149func decodeCacheValue(data []byte) (*CacheValue, error) {
150	n, k := binary.Uvarint(data)
151	if k <= 0 || n > uint64(len(data)-k) {
152		return nil, errors.New("corrupt cache entry")
153	}
154	cv := &CacheValue{}
155	if err := json.Unmarshal(data[k:k+int(n)], cv); err != nil {
156		return nil, err
157	}
158	cv.Body = data[k+int(n):]
159	return cv, nil
160}