feat: kill server after n downloads

TODO

@@ -5,5 +5,4 @@ feat: improve flags
 feat: replace uuid dep with custom id generator
 feat: multiple files
 fix: remove `_` prefix
-fix: kill server after `n` downloads
 feat: use zeroconf for mDNS

cmd/lash/main.go

@@ -14,6 +14,7 @@ import (
 var versionFlag = flag.Bool("version", false, "Print out version and exit.")
 
 var port = flag.Int("p", 1337, "Set the port for LASH exchanges.")
+var numberOfShares = flag.Int("n", -1, "Number of shares to allow before killing the server.")
 
 func main() {
 	flag.Parse()
@@ -23,14 +24,22 @@ func main() {
 		return
 	}
 
-	srv := app.New()
+	ctx := lash.LashContext{
+		N: *numberOfShares,
+	}
+
+	srv := app.New(&ctx)
 
 	// TODO(vxc): Make this more portable
 	errx.Log("Your share link is %s", share.GenerateShareLink(*port))
 	errx.Log("Your token is \033[1;92m%s\033[0m", lash.Token)
 
 	errx.Log("starting server at http://0.0.0.0:%d", *port)
-	if err := http.ListenAndServe(fmt.Sprintf(":%d", *port), srv); err != nil {
+	server := http.Server{
+		Addr:    fmt.Sprintf(":%d", *port),
+		Handler: srv,
+	}
+	if err := server.ListenAndServe(); err != nil {
 		errx.FatalPerror(err)
 	}
 }

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

lash.go

@@ -2,10 +2,16 @@ package lash
 
 import (
 	"embed"
+	"net/http"
 
 	"github.com/google/uuid"
 )
 
+type LashContext struct {
+	N      int
+	Server *http.Server
+}
+
 //go:embed templates/*.html
 var Templates embed.FS
 

version

@@ -1 +1 @@
-1.0.0
+3