feat: kill server after n downloads

internal/app/routes.go

@@ -8,7 +8,7 @@ import (
 	"os"
 )
 
-func New() http.Handler {
+func New(ctx *lash.LashContext) http.Handler {
 	mux := http.NewServeMux()
 
 	fp, err := GetFilePath()
@@ -25,9 +25,12 @@ func New() http.Handler {
 		Version: lash.Version,
 	}
 
-	file := handlers.FileData{
-		Contents: contents,
-		FileName: fp,
+	file := handlers.FileHandler{
+		Ctx: ctx,
+		FileData: handlers.FileData{
+			Contents: contents,
+			FileName: fp,
+		},
 	}
 
 	mux.HandleFunc("/", share.Handler)

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++
 }

internal/shareLink/shareLink.go

@@ -17,13 +17,13 @@ func GenerateShareLink(port int) string {
 }
 
 func getLocalIP() string {
-    conn, err := net.Dial("udp", "8.8.8.8:80")
-    if err != nil {
-	    errx.FatalPerror(err)
-    }
-    defer conn.Close()
+	conn, err := net.Dial("udp", "8.8.8.8:80")
+	if err != nil {
+		errx.FatalPerror(err)
+	}
+	defer conn.Close()
 
-    localAddress := conn.LocalAddr().(*net.UDPAddr)
+	localAddress := conn.LocalAddr().(*net.UDPAddr)
 
-    return localAddress.IP.String()
+	return localAddress.IP.String()
 }