1 files changed, 110 insertions, 2 deletions

diff --git a/roseh.moe.go b/roseh.moe.go
index 95f303c..a512191 100644
--- a/roseh.moe.go
+++ b/roseh.moe.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"archive/tar"
+	"bytes"
 	"cmp"
 	"compress/gzip"
 	"crypto/hkdf"
@@ -31,6 +32,7 @@ import (
 	"strings"
 	"sync"
 	"time"
+	"unicode"
 
 	"github.com/skip2/go-qrcode"
 	"roseh.moe/pkg/oae2"
@@ -47,7 +49,7 @@ var (
 	https       = flag.String("https", "", "directory containing cert.pem and key.pem, or empty string to use unencrypted http")
 	holeTempDir = flag.String("hole-temp-dir", "hole", "directory to store temporary files for the wormhole")
 	codeDir     = flag.String("code-dir", "code", "directory to serve git repos")
-	nezuko      = flag.String("nezuko", "nezuko", "nezuko partition")
+	nezuko      = flag.String("nezuko", "/mnt/nezuko", "nezuko partition")
 )
 
 var notepadPassword, secretKey, authenticatorKey []byte
@@ -297,6 +299,7 @@ func qrHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	w.Header().Set("Content-Type", "image/png")
+	w.Header().Set("Content-Length", strconv.Itoa(len(qr)))
 	w.Write(qr)
 }
 
@@ -951,6 +954,111 @@ func (p *jellyfinReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request)
 	p.proxy.ServeHTTP(w, r)
 }
 
+type nopCloser struct {
+	io.Writer
+}
+
+func (nopCloser) Close() error { return nil }
+
+func (w nopCloser) ReadFrom(r io.Reader) (int64, error) {
+	return io.Copy(w.Writer, r)
+}
+
+type gzipResponseWriter struct {
+	http.ResponseWriter
+
+	w    io.WriteCloser
+	buf  []byte
+	code int
+}
+
+func (w *gzipResponseWriter) ReadFrom(r io.Reader) (int64, error) {
+	if w.w != nil {
+		return io.Copy(w.w, r)
+	}
+	var nn int64
+	if len(w.Header().Values("Content-Type")) == 0 {
+		if w.buf == nil {
+			w.buf = make([]byte, 0, 512)
+		}
+		n, err := io.ReadFull(r, w.buf[len(w.buf):cap(w.buf)])
+		w.buf = w.buf[:len(w.buf)+n]
+		nn += int64(n)
+		if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
+			return nn, err
+		}
+		if len(w.buf) < 512 {
+			return nn, nil
+		}
+		contentType := http.DetectContentType(w.buf)
+		w.Header().Set("Content-Type", contentType)
+	}
+	if disposition, _, _ := strings.Cut(w.Header().Get("Content-Disposition"), ";"); disposition == "attachment" && len(w.Header().Values("Accept-Ranges")) > 0 || !strings.HasPrefix(w.Header().Get("Content-Type"), "text/") {
+		w.w = nopCloser{w.ResponseWriter}
+	} else {
+		var err error
+		w.w, err = gzip.NewWriterLevel(w.ResponseWriter, gzip.BestSpeed)
+		if err != nil {
+			return nn, err
+		}
+		w.Header().Set("Content-Encoding", "gzip")
+		w.Header().Del("Content-Length")
+		w.Header().Del("Accept-Ranges")
+	}
+	if w.code != 0 {
+		w.ResponseWriter.WriteHeader(w.code)
+	}
+	if len(w.buf) > 0 {
+		if n, err := w.w.Write(w.buf); err != nil {
+			return max(int64(n)-(512-nn), 0), err
+		}
+	}
+	n, err := io.Copy(w.w, r)
+	nn += n
+	return nn, err
+}
+
+func (w *gzipResponseWriter) Write(p []byte) (int, error) {
+	n, err := w.ReadFrom(bytes.NewReader(p))
+	return int(n), err
+}
+
+func (w *gzipResponseWriter) WriteHeader(statusCode int) {
+	w.code = statusCode
+}
+
+func (w *gzipResponseWriter) close() {
+	if w.w == nil {
+		if w.code != 0 {
+			w.ResponseWriter.WriteHeader(w.code)
+		}
+		w.ResponseWriter.Write(w.buf)
+	} else {
+		w.w.Close()
+	}
+}
+
+type gzipHandler struct {
+	h http.Handler
+}
+
+func (h *gzipHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if len(r.Header.Values("Range")) == 0 {
+		for _, hdr := range r.Header.Values("Accept-Encoding") {
+			for encoding := range strings.FieldsFuncSeq(hdr, func(r rune) bool { return r == ',' || unicode.IsSpace(r) }) {
+				encoding, _, _ := strings.Cut(encoding, ";")
+				if encoding == "gzip" {
+					gzipWriter := &gzipResponseWriter{ResponseWriter: w}
+					h.h.ServeHTTP(gzipWriter, r)
+					gzipWriter.close()
+					return
+				}
+			}
+		}
+	}
+	h.h.ServeHTTP(w, r)
+}
+
 func main() {
 	flag.Parse()
 
@@ -991,7 +1099,7 @@ func main() {
 	apiMux.HandleFunc("PUT /api/wormhole/{name}", wormholeAPIUpload)
 	mux.Handle("/api/", apiMux)
 	mux.HandleFunc("/", notFound)
-	http.Handle("/", http.NewCrossOriginProtection().Handler(mux))
+	http.Handle("/", http.NewCrossOriginProtection().Handler(&gzipHandler{mux}))
 	jellyfinURL, err := url.Parse("http://127.0.0.1:8096")
 	if err != nil {
 		log.Fatal(err)