From ef49a98ee62ebdec93d1c04352b6da4de0d5c79d Mon Sep 17 00:00:00 2001
From: Fionera <fionera@fionera.de>
Date: Sun, 7 Feb 2021 05:40:29 +0100
Subject: [PATCH] refactor: Remove duplicate code of zstHandler

---
 server/handlers.go | 95 ++++++++++++++++------------------------------
 1 file changed, 32 insertions(+), 63 deletions(-)

diff --git a/server/handlers.go b/server/handlers.go
index 0971e33..227d45f 100644
--- a/server/handlers.go
+++ b/server/handlers.go
@@ -859,73 +859,25 @@ func (s *Server) headHandler(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Connection", "close")
 }
 
-func (s *Server) getHandlerZst(w http.ResponseWriter, r *http.Request) {
-	vars := mux.Vars(r)
-
-	action := vars["action"]
-	token := vars["token"]
-	filename := vars["filename"]
-
-	if err := s.CheckMetadata(token, filename+".zst"); err != nil {
-		log.Printf("Error metadata: %v", err)
-		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
-		return
-	}
-
-	reader, contentType, _, err := s.storage.Get(token, filename+".zst")
-	if s.storage.IsNotExist(err) {
-		log.Println(err)
-		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
-		return
-	}
-	if err != nil {
-		log.Printf("Failed to get from storage: %s", err.Error())
-		http.Error(w, "Could not retrieve file.", 500)
-		return
-	}
-	defer reader.Close()
-	contentType = mime.TypeByExtension(filepath.Ext(filename))
-
-	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()
-
-	var disposition string
-	if action == "inline" {
-		disposition = "inline"
-	} else {
-		disposition = "attachment"
-	}
-
-	w.Header().Set("Content-Type", contentType)
-	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") != "" {
-		log.Printf("Range request with decompression")
-		http.Error(w, "Range requests with decompression are not supported", 400)
-		return
-	}
-
-	_, _ = io.Copy(w, d)
-}
-
 func (s *Server) getHandler(w http.ResponseWriter, r *http.Request) {
 	vars := mux.Vars(r)
 
 	action := vars["action"]
 	token := vars["token"]
 	filename := vars["filename"]
+	realFilename := filename
+	useZST := false
 
 	if err := s.CheckMetadata(token, filename); err != nil {
 		log.Printf("Error metadata: %v; trying with .zst", err)
-		s.getHandlerZst(w, r)
-		return
+
+		useZST = true
+		filename += ".zst"
+		if err := s.CheckMetadata(token, filename); err != nil {
+			log.Printf("Error metadata: %v", err)
+			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
+			return
+		}
 	}
 
 	reader, contentType, contentLength, err := s.storage.Get(token, filename)
@@ -937,11 +889,24 @@ func (s *Server) getHandler(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, "Could not retrieve file.", 500)
 		return
 	}
-
 	defer reader.Close()
 
-	var disposition string
+	if useZST {
+		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 {
@@ -949,8 +914,12 @@ 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))
+	if useZST {
+		w.Header().Set("Transfer-Encoding", "chunked")
+	} else {
+		w.Header().Set("Content-Length", strconv.FormatUint(contentLength, 10))
+	}
+	w.Header().Set("Content-Disposition", fmt.Sprintf("%s; filename=\"%s\"", disposition, realFilename))
 	w.Header().Set("Connection", "keep-alive")
 
 	if w.Header().Get("Range") == "" {
@@ -987,7 +956,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 {