save

internal/generator/generator.go

@@ -34,6 +34,10 @@ var keys []string = []string{
 	"wash",
 	"vote",
 	"onyx",
+	"snow",
+	"john",
+	"json",
+	"fang",
 }
 
 func Generate(n int) (code string) {

internal/handlers/file.go

@@ -7,7 +7,6 @@ import (
 	"lash"
 	"lash/internal/errx"
 	"net/http"
-	"os"
 	"strconv"
 )
 
@@ -47,36 +46,38 @@ func (h FileHandler) APIHandler(w http.ResponseWriter, r *http.Request) {
 
 func (h FileHandler) FileHandler(w http.ResponseWriter, r *http.Request) {
 	if sent >= h.Ctx.N && h.Ctx.N != -1 {
-		w.WriteHeader(http.StatusTooManyRequests)
-		os.Exit(0)
+		http.Error(w, "Too many requests", http.StatusTooManyRequests)
 		return
 	}
 
-	w.WriteHeader(http.StatusOK)
+	filename := h.FileData[0].FileName
+	if len(h.FileData) > 1 {
+		filename = "lash.zip"
+	}
+
+	w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, filename))
+	w.Header().Set("Content-Type", "application/octet-stream")
+
 	if len(h.FileData) == 1 {
-		w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", h.FileData[0].FileName))
-		w.Header().Set("Content-Type", "application/octet-stream")
-		w.Header().Set("Content-Length", strconv.Itoa(len(h.FileData[0].FileName)))
-
-		w.Write(h.FileData[0].Contents)
-
+		w.Header().Set("Content-Length", strconv.Itoa(len(h.FileData[0].Contents)))
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write(h.FileData[0].Contents)
 	} else {
-		w.Header().Set("Content-Disposition", "attachment; filename=lash.zip")
-		w.Header().Set("Content-Type", "application/octet-stream")
+		w.WriteHeader(http.StatusOK)
 		zw := zip.NewWriter(w)
 		defer zw.Close()
 
 		for _, f := range h.FileData {
-			w, err := zw.Create(f.FileName)
+			fw, err := zw.Create(f.FileName)
 			if err != nil {
 				errx.FatalPerror(err)
 			}
-
-			_, err = w.Write(f.Contents)
+			_, err = fw.Write(f.Contents)
 			if err != nil {
 				errx.FatalPerror(err)
 			}
 		}
 	}
+
 	sent++
 }