From 55a5ae413f03eddc2f06bce96d1692074b6c5c24 Mon Sep 17 00:00:00 2001
From: JustAnotherArchivist <JustAnotherArchivist@users.noreply.github.com>
Date: Thu, 4 Feb 2021 00:25:27 +0000
Subject: [PATCH] Attempt to implement on-the-fly .zst decompression on GET;
 does not work

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

diff --git a/server/handlers.go b/server/handlers.go
index b4f1a7c..42aa518 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 (
@@ -866,23 +868,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)
+	isZstd := false
+	var d zstd.Decoder
+	_ = d // Only used when isZstd is true; silence compiler
 	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()
+		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()
+		isZstd = true
+		contentType = mime.TypeByExtension(filepath.Ext(filename))
 	} 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 {
+		defer reader.Close()
 	}
 
-	defer reader.Close()
-
 	var disposition string
 
 	if action == "inline" {
@@ -892,17 +916,47 @@ 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))
+	if !isZstd {
+		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, filename))
 	w.Header().Set("Connection", "keep-alive")
 
 	if w.Header().Get("Range") == "" {
-		if _, err = io.Copy(w, reader); err != nil {
-			log.Printf("%s", err.Error())
-			http.Error(w, "Error occurred copying to output stream", 500)
-			return
+		if !isZstd {
+			if _, err = io.Copy(w, reader); err != nil {
+				log.Printf("%s", err.Error())
+				http.Error(w, "Error occurred copying to output stream", 500)
+				return
+			}
+		} else {
+			buffer := make([]byte, 1024)
+			for {
+				log.Printf("Reading from decoded stream")
+				n, err := d.Read(buffer)
+				log.Printf("Read from decoded stream")
+				if err != nil && err != io.EOF {
+					log.Printf("Failed to read from file: %s", err.Error())
+					panic("Error reading data")
+				}
+				log.Printf("Writing to HTTP")
+				w.Write(buffer[0:n])
+				log.Printf("Trying to flush")
+				if f, ok := w.(http.Flusher); ok {
+					f.Flush()
+				}
+				if err == io.EOF {
+					break
+				}
+			}
 		}
 
+		return
+	} else if isZstd {
+		log.Printf("Range request with decompression")
+		http.Error(w, "Range requests with decompression are not supported", 400)
 		return
 	}