feat: kill server after n downloads

internal/handlers/file.go

@@ -5,9 +5,15 @@ import (
 	"fmt"
 	"lash"
 	"net/http"
+	"os"
 	"strconv"
 )
 
+type FileHandler struct {
+	Ctx      *lash.LashContext
+	FileData FileData
+}
+
 type FileData struct {
 	Contents []byte
 	FileName string
@@ -17,7 +23,14 @@ type ValidateRequest struct {
 	Token string
 }
 
-func (h FileData) APIHandler(w http.ResponseWriter, r *http.Request) {
+var sent int = 0
+
+func (h FileHandler) APIHandler(w http.ResponseWriter, r *http.Request) {
+	if sent >= h.Ctx.N && h.Ctx.N != -1 {
+		w.WriteHeader(http.StatusTooManyRequests)
+		os.Exit(0)
+		return
+	}
 	decoder := json.NewDecoder(r.Body)
 
 	var t ValidateRequest
@@ -32,21 +45,27 @@ func (h FileData) APIHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// Send the file over as a stream of bytes
-	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", h.FileName))
+	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.Contents)))
+	w.Header().Set("Content-Length", strconv.Itoa(len(h.FileData.Contents)))
 
 	w.WriteHeader(http.StatusOK)
-	w.Write(h.Contents)
+	w.Write(h.FileData.Contents)
+	sent++
 }
 
-func (h FileData) FileHandler(w http.ResponseWriter, r *http.Request) {
-	// Send the file over as a stream of bytes
-	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", h.FileName))
+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)
+		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.Contents)))
+	w.Header().Set("Content-Length", strconv.Itoa(len(h.FileData.Contents)))
 
 	w.WriteHeader(http.StatusOK)
-	w.Write(h.Contents)
+	w.Write(h.FileData.Contents)
+	sent++
 }