okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

nekomimi.petapi: zstd responses, streams and request bodies; images served as raw bytes; drop redundant re-encodes7fd8cf5

main
branchesmainall refs →
2.1 KiB71 linesraw
1package daemon
2
3import (
4	"io"
5	"net/http"
6	"strings"
7
8	"github.com/klauspost/compress/zstd"
9)
10
11// Request bodies at least this large are compressed when the daemon
12// advertises zstd_requests; smaller ones are not worth a frame.
13const compressedRequestBytes = 16 * 1024
14
15// The daemon compresses with a small window; anything larger is refused.
16const zstdWindowBytes = 8 << 20
17
18var zstdEncoder, _ = zstd.NewWriter(nil, zstd.WithEncoderLevel(zstd.SpeedFastest), zstd.WithEncoderConcurrency(1))
19
20// requestBody returns the payload to send and its Content-Encoding, empty
21// for identity.
22func requestBody(conn *Connection, payload []byte) ([]byte, string) {
23	if len(payload) < compressedRequestBytes || conn.capability("zstd_requests") < 1 {
24		return payload, ""
25	}
26	return zstdEncoder.EncodeAll(payload, nil), "zstd"
27}
28
29// codingTransport asks for zstd responses and decodes them, the way
30// net/http does for gzip when no Accept-Encoding is set.
31type codingTransport struct{ base http.RoundTripper }
32
33func (transport codingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
34	if req.Header.Get("Accept-Encoding") == "" {
35		req = req.Clone(req.Context())
36		req.Header.Set("Accept-Encoding", "zstd")
37	}
38	res, err := transport.base.RoundTrip(req)
39	if err != nil || !strings.EqualFold(res.Header.Get("Content-Encoding"), "zstd") {
40		return res, err
41	}
42	decoder, err := zstd.NewReader(res.Body, zstd.WithDecoderConcurrency(1), zstd.WithDecoderMaxWindow(zstdWindowBytes))
43	if err != nil {
44		_ = res.Body.Close()
45		return nil, err
46	}
47	res.Body = zstdBody{decoder: decoder, body: res.Body}
48	res.Header.Del("Content-Encoding")
49	res.Header.Del("Content-Length")
50	res.ContentLength = -1
51	res.Uncompressed = true
52	return res, nil
53}
54
55func (transport codingTransport) CloseIdleConnections() {
56	if closer, ok := transport.base.(interface{ CloseIdleConnections() }); ok {
57		closer.CloseIdleConnections()
58	}
59}
60
61type zstdBody struct {
62	decoder *zstd.Decoder
63	body    io.ReadCloser
64}
65
66func (body zstdBody) Read(target []byte) (int, error) { return body.decoder.Read(target) }
67
68func (body zstdBody) Close() error {
69	body.decoder.Close()
70	return body.body.Close()
71}