From 21004eb3452d1c34420a29dba0bdb5390820a590 Mon Sep 17 00:00:00 2001
From: JustAnotherArchivist <JustAnotherArchivist@users.noreply.github.com>
Date: Sun, 7 Feb 2021 04:49:39 +0000
Subject: [PATCH] On-the-fly .zst decompression, take 2

---
 server/handlers.go | 64 ++++++++++++++++++++++++++++++++++++----------
 1 file changed, 51 insertions(+), 13 deletions(-)

diff --git a/server/handlers.go b/server/handlers.go
index b4f1a7c..7296b64 100644
--- a/server/handlers.go
+++ b/server/handlers.go
@@ -61,6 +61,8 @@ import (
 
 	"encoding/base64"
 	qrcode "github.com/skip2/go-qrcode"
+
+	"github.com/klauspost/compress/zstd"
 )
 
 var (
@@ -838,23 +840,34 @@ func (s *Server) headHandler(w http.ResponseWriter, r *http.Request) {
 	filename := vars["filename"]
 
 	if err := s.CheckMetadata(token, filename); err != nil {
-		log.Printf("Error metadata: %s", err.Error())
-		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
-		return
+		if err2 := s.CheckMetadata(token, filename + ".zst"); err2 != nil {
+			log.Printf("Error metadata: %s", err.Error())
+			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
+			return
+		}
 	}
 
 	contentType, contentLength, err := s.storage.Head(token, filename)
 	if s.storage.IsNotExist(err) {
-		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
-		return
+		_, _, err := s.storage.Head(token, filename + ".zst")
+		if s.storage.IsNotExist(err) {
+			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
+			return
+		} else if err != nil {
+			log.Printf("%s", err.Error())
+			http.Error(w, "Could not retrieve file.", 500)
+			return
+		}
+		contentType = mime.TypeByExtension(filepath.Ext(filename))
 	} else if err != nil {
 		log.Printf("%s", err.Error())
 		http.Error(w, "Could not retrieve file.", 500)
 		return
+	} else {
+		w.Header().Set("Content-Length", strconv.FormatUint(contentLength, 10))
 	}
 
 	w.Header().Set("Content-Type", contentType)
-	w.Header().Set("Content-Length", strconv.FormatUint(contentLength, 10))
 	w.Header().Set("Connection", "close")
 }
 
@@ -866,19 +879,45 @@ func (s *Server) getHandler(w http.ResponseWriter, r *http.Request) {
 	filename := vars["filename"]
 
 	if err := s.CheckMetadata(token, filename); err != nil {
-		log.Printf("Error metadata: %s", err.Error())
-		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
-		return
+		if err2 := s.CheckMetadata(token, filename + ".zst"); err2 != nil {
+			log.Printf("Error metadata: %s and %s", err.Error(), err2.Error())
+			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
+			return
+		}
 	}
 
 	reader, contentType, contentLength, err := s.storage.Get(token, filename)
 	if s.storage.IsNotExist(err) {
-		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
-		return
+		reader, _, _, err = s.storage.Get(token, filename + ".zst")
+		if s.storage.IsNotExist(err) {
+			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
+			return
+		} else if err != nil {
+			log.Printf("Failed to get .zst from storage: %s", err.Error())
+			http.Error(w, "Could not retrieve file.", 500)
+			return
+		}
+		defer reader.Close()
+		if w.Header().Get("Range") != "" {
+			log.Printf("Range request with decompression")
+			http.Error(w, "Range requests with decompression are not supported", 400)
+			return
+		}
+		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
+		}
+		reader = d.IOReadCloser()
+		contentType = mime.TypeByExtension(filepath.Ext(filename))
+		w.Header().Set("Transfer-Encoding", "chunked")
 	} else if err != nil {
-		log.Printf("%s", err.Error())
+		log.Printf("Failed to get from storage: %s", err.Error())
 		http.Error(w, "Could not retrieve file.", 500)
 		return
+	} else {
+		w.Header().Set("Content-Length", strconv.FormatUint(contentLength, 10))
 	}
 
 	defer reader.Close()
@@ -892,7 +931,6 @@ func (s *Server) getHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	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))
 	w.Header().Set("Connection", "keep-alive")