From 40e2cf850e7cff21a510aa2e4d5230c718e2972e Mon Sep 17 00:00:00 2001
From: JustAnotherArchivist <JustAnotherArchivist@users.noreply.github.com>
Date: Mon, 8 Feb 2021 01:41:56 +0000
Subject: [PATCH] On-the-fly Zstd decompression

---
 server/handlers.go | 65 +++++++++++++++++++++++++++++++++++++---------
 1 file changed, 53 insertions(+), 12 deletions(-)

diff --git a/server/handlers.go b/server/handlers.go
index 5c61a8a..ef9740f 100644
--- a/server/handlers.go
+++ b/server/handlers.go
@@ -61,6 +61,8 @@ import (
 	"github.com/gorilla/mux"
 	"github.com/microcosm-cc/bluemonday"
 	"github.com/skip2/go-qrcode"
+
+	"github.com/klauspost/compress/zstd"
 )
 
 const getPathPart = "get"
@@ -935,16 +937,27 @@ func (s *Server) headHandler(w http.ResponseWriter, r *http.Request) {
 
 	token := vars["token"]
 	filename := vars["filename"]
+	realFilename := filename
+	useZstd := false
 
 	metadata, err := s.CheckMetadata(token, filename, false)
 
+	var contentType string
 	if err != nil {
-		log.Printf("Error metadata: %s", err.Error())
-		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
-		return
+		log.Printf("Error metadata: %s; trying .zst", err.Error())
+		useZstd = true
+		filename += ".zst"
+		metadata, err = s.CheckMetadata(token, filename, false)
+		if err != nil {
+			log.Printf("Error metadata .zst: %s", err.Error())
+			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
+			return
+		}
+		contentType = mime.TypeByExtension(filepath.Ext(realFilename))
+	} else {
+		contentType = metadata.ContentType
 	}
 
-	contentType := metadata.ContentType
 	contentLength, err := s.storage.Head(token, filename)
 	if s.storage.IsNotExist(err) {
 		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
@@ -958,7 +971,9 @@ func (s *Server) headHandler(w http.ResponseWriter, r *http.Request) {
 	remainingDownloads, remainingDays := metadata.remainingLimitHeaderValues()
 
 	w.Header().Set("Content-Type", contentType)
-	w.Header().Set("Content-Length", strconv.FormatUint(contentLength, 10))
+	if !useZstd {
+		w.Header().Set("Content-Length", strconv.FormatUint(contentLength, 10))
+	}
 	w.Header().Set("Connection", "close")
 	w.Header().Set("X-Remaining-Downloads", remainingDownloads)
 	w.Header().Set("X-Remaining-Days", remainingDays)
@@ -970,13 +985,21 @@ func (s *Server) getHandler(w http.ResponseWriter, r *http.Request) {
 	action := vars["action"]
 	token := vars["token"]
 	filename := vars["filename"]
+	realFilename := filename
+	useZstd := false
 
 	metadata, err := s.CheckMetadata(token, filename, true)
 
 	if err != nil {
-		log.Printf("Error metadata: %s", err.Error())
-		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
-		return
+		log.Printf("Error metadata: %s; trying .zst", err.Error())
+		useZstd = true
+		filename += ".zst"
+		metadata, err = s.CheckMetadata(token, filename, true)
+		if err != nil {
+			log.Printf("Error metadata .zst: %s", err.Error())
+			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
+			return
+		}
 	}
 
 	contentType := metadata.ContentType
@@ -992,8 +1015,22 @@ func (s *Server) getHandler(w http.ResponseWriter, r *http.Request) {
 
 	defer reader.Close()
 
-	var disposition string
+	if useZstd {
+		contentType = mime.TypeByExtension(filepath.Ext(realFilename))
+
+		d, err := zstd.NewReader(reader)
+		if err != nil {
+			log.Printf("Failed to create zstd reader: %s", err.Error())
+			http.Error(w, "Could not retrieve file.", 500)
+			return
+		}
 
+		defer d.Close()
+
+		reader = d.IOReadCloser()
+	}
+
+	var disposition string
 	if action == "inline" {
 		disposition = "inline"
 	} else {
@@ -1003,8 +1040,12 @@ func (s *Server) getHandler(w http.ResponseWriter, r *http.Request) {
 	remainingDownloads, remainingDays := metadata.remainingLimitHeaderValues()
 
 	w.Header().Set("Content-Type", contentType)
-	w.Header().Set("Content-Length", strconv.FormatUint(contentLength, 10))
-	w.Header().Set("Content-Disposition", fmt.Sprintf("%s; filename=\"%s\"", disposition, filename))
+	if !useZstd {
+		w.Header().Set("Content-Length", strconv.FormatUint(contentLength, 10))
+	} else {
+		w.Header().Set("Transfer-Encoding", "chunked")
+	}
+	w.Header().Set("Content-Disposition", fmt.Sprintf("%s; filename=\"%s\"", disposition, realFilename))
 	w.Header().Set("Connection", "keep-alive")
 	w.Header().Set("X-Remaining-Downloads", remainingDownloads)
 	w.Header().Set("X-Remaining-Days", remainingDays)
@@ -1043,7 +1084,7 @@ func (s *Server) getHandler(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
-	http.ServeContent(w, r, filename, time.Now(), file)
+	http.ServeContent(w, r, filename, time.Unix(0, 0), file)
 }
 
 func (s *Server) RedirectHandler(h http.Handler) http.HandlerFunc {