add compression

main.go

@@ -1,12 +1,16 @@
 package main
 
 import (
+	"compress/gzip"
 	"flag"
 	"fmt"
+	"io"
 	"log"
 	"net/http"
+	"strings"
 
 	"git.teamworkapps.com/shortcut/enhancedfileserver"
+	"github.com/andybalholm/brotli"
 	"github.com/gorilla/mux"
 )
 
@@ -51,6 +55,9 @@ func main() {
 		r.Use(noCache)
 	}
 
+	// Always enable compression (both brotli and gzip)
+	r.Use(compression)
+
 	if c.noindex {
 		r.PathPrefix("/").Handler(http.StripPrefix("/", enhancedfileserver.FileServerNoIndex(http.Dir("./"))))
 	} else {
@@ -92,3 +99,58 @@ func queryLog(next http.Handler) http.Handler {
 		next.ServeHTTP(w, r)
 	})
 }
+
+// compressionWriter wraps http.ResponseWriter to provide compression
+type compressionWriter struct {
+	http.ResponseWriter
+	writer io.Writer
+}
+
+func (cw *compressionWriter) Write(b []byte) (int, error) {
+	return cw.writer.Write(b)
+}
+
+// compression middleware that supports both brotli and gzip
+func compression(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Get the Accept-Encoding header
+		acceptEncoding := r.Header.Get("Accept-Encoding")
+
+		// Check if client supports brotli (preferred)
+		if strings.Contains(acceptEncoding, "br") {
+			w.Header().Set("Content-Encoding", "br")
+			w.Header().Set("Vary", "Accept-Encoding")
+
+			bw := brotli.NewWriter(w)
+			defer bw.Close()
+
+			cw := &compressionWriter{
+				ResponseWriter: w,
+				writer:         bw,
+			}
+
+			next.ServeHTTP(cw, r)
+			return
+		}
+
+		// Check if client supports gzip
+		if strings.Contains(acceptEncoding, "gzip") {
+			w.Header().Set("Content-Encoding", "gzip")
+			w.Header().Set("Vary", "Accept-Encoding")
+
+			gw := gzip.NewWriter(w)
+			defer gw.Close()
+
+			cw := &compressionWriter{
+				ResponseWriter: w,
+				writer:         gw,
+			}
+
+			next.ServeHTTP(cw, r)
+			return
+		}
+
+		// No compression supported, serve normally
+		next.ServeHTTP(w, r)
+	})
+}