maint: changes

internal/handlers/file.go

@@ -1,9 +1,11 @@
 package handlers
 
 import (
+	"archive/zip"
 	"encoding/json"
 	"fmt"
 	"lash"
+	"lash/internal/errx"
 	"net/http"
 	"os"
 	"strconv"
@@ -11,7 +13,7 @@ import (
 
 type FileHandler struct {
 	Ctx      *lash.LashContext
-	FileData FileData
+	FileData []FileData
 }
 
 type FileData struct {
@@ -23,10 +25,17 @@ type ValidateRequest struct {
 	Token string
 }
 
+type headers map[string]string
+
+var zipHeaders headers = headers{
+	"Content-Disposition": "attachment; filename=lash.zip",
+	"Content-Type":        "application/octet-stream",
+}
+
 var sent int = 0
 
 func (h FileHandler) APIHandler(w http.ResponseWriter, r *http.Request) {
-	if sent >= h.Ctx.N && h.Ctx.N != -1 {
+	if h.hasHitMax() {
 		w.WriteHeader(http.StatusTooManyRequests)
 		os.Exit(0)
 		return
@@ -45,27 +54,71 @@ func (h FileHandler) APIHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", h.FileData.FileName))
-	w.Header().Set("Content-Type", "application/octet-stream")
-	w.Header().Set("Content-Length", strconv.Itoa(len(h.FileData.Contents)))
-
 	w.WriteHeader(http.StatusOK)
-	w.Write(h.FileData.Contents)
+	if len(h.FileData) == 1 {
+		headers := headers{
+			"Content-Disposition": fmt.Sprintf("attachment; filename=%s", h.FileData[0].FileName),
+			"Content-Type":        "application/octet-stream",
+			"Content-Length":      strconv.Itoa(len(h.FileData[0].FileName)),
+		}
+		headers.set(w)
+
+		w.Write(h.FileData[0].Contents)
+	} else {
+		zipHeaders.set(w)
+		h.writeZip(w)
+	}
 	sent++
 }
 
 func (h FileHandler) FileHandler(w http.ResponseWriter, r *http.Request) {
-	if sent >= h.Ctx.N && h.Ctx.N != -1 {
+	if h.hasHitMax() {
 		w.WriteHeader(http.StatusTooManyRequests)
 		os.Exit(0)
 		return
 	}
 
-	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", h.FileData.FileName))
-	w.Header().Set("Content-Type", "application/octet-stream")
-	w.Header().Set("Content-Length", strconv.Itoa(len(h.FileData.Contents)))
-
 	w.WriteHeader(http.StatusOK)
-	w.Write(h.FileData.Contents)
+	if len(h.FileData) == 1 {
+		headers := headers{
+			"Content-Disposition": fmt.Sprintf("attachment; filename=%s", h.FileData[0].FileName),
+			"Content-Type":        "application/octet-stream",
+			"Content-Length":      strconv.Itoa(len(h.FileData[0].FileName)),
+		}
+		headers.set(w)
+
+		w.Write(h.FileData[0].Contents)
+
+	} else {
+		zipHeaders.set(w)
+		h.writeZip(w)
+	}
 	sent++
 }
+
+func (h FileHandler) writeZip(hw http.ResponseWriter) {
+	zw := zip.NewWriter(hw)
+	defer zw.Close()
+
+	for _, f := range h.FileData {
+		w, err := zw.Create(f.FileName)
+		if err != nil {
+			errx.FatalPerror(err)
+		}
+
+		_, err = w.Write(f.Contents)
+		if err != nil {
+			errx.FatalPerror(err)
+		}
+	}
+}
+
+func (h headers) set(w http.ResponseWriter) {
+	for k, v := range h {
+		w.Header().Set(k, v)
+	}
+}
+
+func (h FileHandler) hasHitMax() bool {
+	return sent >= h.Ctx.N && h.Ctx.N != -1
+}