1package ext
  2
  3import (
  4	"compress/gzip"
  5	"compress/lzw"
  6	"errors"
  7	"io"
  8	"log/slog"
  9	"net/http"
 10	"strconv"
 11	"strings"
 12
 13	"git.gabrielgio.me/cerrado/pkg/u"
 14	"github.com/andybalholm/brotli"
 15	"github.com/klauspost/compress/zstd"
 16)
 17
 18var errInvalidParam = errors.New("Invalid weighted param")
 19
 20type CompressionResponseWriter struct {
 21	innerWriter    http.ResponseWriter
 22	compressWriter io.Writer
 23}
 24
 25func Compress(next HandlerFunc) HandlerFunc {
 26	return func(w http.ResponseWriter, r *Request) {
 27		// TODO: hand this better
 28		if strings.HasSuffix(r.URL.Path, ".tar.gz") {
 29			next(w, r)
 30			return
 31		}
 32
 33		if accept, ok := r.Header["Accept-Encoding"]; ok {
 34			if compress, algo := GetCompressionWriter(u.FirstOrZero(accept), w); algo != "" {
 35				defer compress.Close()
 36				w.Header().Add("Content-Encoding", algo)
 37				w = &CompressionResponseWriter{
 38					innerWriter:    w,
 39					compressWriter: compress,
 40				}
 41			}
 42		}
 43		next(w, r)
 44	}
 45}
 46
 47func Decompress(next http.HandlerFunc) http.HandlerFunc {
 48	return func(w http.ResponseWriter, r *http.Request) {
 49		// TODO: hand this better
 50		if strings.HasSuffix(r.URL.Path, ".tar.gz") {
 51			next(w, r)
 52			return
 53		}
 54
 55		if accept, ok := r.Header["Accept-Encoding"]; ok {
 56			if compress, algo := GetCompressionWriter(u.FirstOrZero(accept), w); algo != "" {
 57				defer compress.Close()
 58				w.Header().Add("Content-Encoding", algo)
 59				w = &CompressionResponseWriter{
 60					innerWriter:    w,
 61					compressWriter: compress,
 62				}
 63			}
 64		}
 65		next(w, r)
 66	}
 67}
 68
 69func GetCompressionWriter(header string, inner io.Writer) (io.WriteCloser, string) {
 70	c := GetCompression(header)
 71	switch c {
 72	case "br":
 73		return GetBrotliWriter(inner), c
 74	case "gzip":
 75		return GetGZIPWriter(inner), c
 76	case "compress":
 77		return GetLZWWriter(inner), c
 78	case "zstd":
 79		return GetZSTDWriter(inner), c
 80	default:
 81		return nil, ""
 82	}
 83}
 84
 85func (c *CompressionResponseWriter) Header() http.Header {
 86	return c.innerWriter.Header()
 87}
 88
 89func (c *CompressionResponseWriter) Write(b []byte) (int, error) {
 90	return c.compressWriter.Write(b)
 91}
 92
 93func (c *CompressionResponseWriter) WriteHeader(statusCode int) {
 94	c.innerWriter.WriteHeader(statusCode)
 95}
 96
 97func GetCompression(header string) string {
 98	c := "*"
 99	q := 0.0
100
101	if header == "" {
102		return c
103	}
104
105	// https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Accept-Encoding
106	for _, e := range strings.Split(header, ",") {
107		ps := strings.Split(e, ";")
108		if len(ps) == 2 {
109			w, err := getWeighedValue(ps[1])
110			if err != nil {
111				slog.Error(
112					"Error parsing weighting from Accept-Encoding",
113					"error", err,
114				)
115				continue
116			}
117			// gettting weighting value
118			if w > q {
119				q = w
120				c = strings.Trim(ps[0], " ")
121			}
122		} else {
123			if 1 > q {
124				q = 1
125				c = strings.Trim(ps[0], " ")
126			}
127		}
128	}
129
130	return c
131}
132
133func GetGZIPWriter(w io.Writer) io.WriteCloser {
134	// error can be ignored here since it will only err when compression level
135	// is not valid
136	r, _ := gzip.NewWriterLevel(w, gzip.BestCompression)
137	return r
138}
139
140func GetBrotliWriter(w io.Writer) io.WriteCloser {
141	return brotli.NewWriterLevel(w, brotli.BestCompression)
142}
143
144func GetZSTDWriter(w io.Writer) io.WriteCloser {
145	// error can be ignored here since it will only opts are given
146	r, _ := zstd.NewWriter(w)
147	return r
148}
149
150func GetLZWWriter(w io.Writer) io.WriteCloser {
151	return lzw.NewWriter(w, lzw.LSB, 8)
152}
153
154func getWeighedValue(part string) (float64, error) {
155	ps := strings.SplitN(part, "=", 2)
156	if len(ps) != 2 {
157		return 0, errInvalidParam
158	}
159	if name := strings.TrimSpace(ps[0]); name == "q" {
160		w, err := strconv.ParseFloat(ps[1], 64)
161		if err != nil {
162			return 0, err
163		}
164		return w, nil
165	}
166
167	return 0, errInvalidParam
168}