Support token, host and port flags

1 files changed, 20 insertions, 4 deletions

diff --git a/main.go b/main.go
index 6625dfc..b416971 100644
--- a/main.go
+++ b/main.go
@@ -10,21 +10,24 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"flag"
 )
 
 //go:embed templates
 var tmplFS embed.FS
 
+
 type BoxHandler struct {
 	dataPath string
+	token string
 }
 
-func serve(w http.ResponseWriter, views ...string) {
+func serve(w http.ResponseWriter, token string, views ...string) {
 	t, err := template.New("index.html").ParseFS(tmplFS, views...)
 	if err != nil {
 		log.Fatal(err)
 	}
-	if err := t.Execute(w, nil); err != nil {
+	if err := t.Execute(w, token); err != nil {
 		log.Fatal(err)
 	}
 }
@@ -32,13 +35,20 @@ func serve(w http.ResponseWriter, views ...string) {
 func (handler BoxHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	switch r.Method {
 	case http.MethodGet:
-		serve(w, "templates/index.html")
+		serve(w, handler.token, "templates/index.html")
 		return
 	case http.MethodPost:
+		token := r.Header.Get("X-Upload-Token")
+		if token != handler.token {
+			log.Println("unauthorized")
+			w.WriteHeader(http.StatusUnauthorized)
+			return
+		}
 		filename := filepath.Join(handler.dataPath, path.Base(r.URL.Path))
 		log.Printf("boxing %s\n", filename)
 		f, err := os.Create(filename)
 		if err != nil {
+			log.Println(err)
 			fmt.Fprint(w, err.Error())
 			w.WriteHeader(http.StatusInternalServerError)
 			return
@@ -52,12 +62,18 @@ func (handler BoxHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 }
 
 func main() {
+	host := flag.String("n", "", "The hostname to listen on")
+	port := flag.Int("p", 8080, "The port to listen on")
+	token := flag.String("t", "", "The token to use to protect uploads")
+	flag.Parse()
 	boxHandler := BoxHandler {
 		"data",
+		*token,
 	}
 	err := os.MkdirAll(boxHandler.dataPath, 0750)
 	if err != nil {
 		log.Fatal(err)
 	}
-	log.Fatal(http.ListenAndServe(":8080", boxHandler))
+	log.Printf("Listening on %s:%d", *host, *port)
+	log.Fatal(http.ListenAndServe(fmt.Sprintf("%s:%d", *host, *port), boxHandler))
 }