Add upload handler.

1 files changed, 52 insertions, 1 deletions

diff --git a/main.go b/main.go
index cf428ed..36781f2 100644
--- a/main.go
+++ b/main.go
@@ -3,9 +3,11 @@ package main
 import (
 	"flag"
 	"fmt"
+	"io"
 	"log"
 	"net/http"
 	"os"
+	"strings"
 )
 
 func accessLog(h http.Handler) http.Handler {
@@ -15,10 +17,52 @@ func accessLog(h http.Handler) http.Handler {
 	})
 }
 
+func uploadHandler(dir string) http.Handler {
+	fn := func(w http.ResponseWriter, r *http.Request) {
+		file, header, err := r.FormFile("file")
+
+		if err != nil {
+			w.WriteHeader(500)
+			w.Write([]byte(err.Error()))
+			log.Println("ERROR", err.Error())
+			return
+		}
+
+		defer file.Close()
+
+		if !strings.HasSuffix(dir, "/") {
+			dir = dir + "/"
+		}
+
+		out, err := os.Create(dir + header.Filename)
+		if err != nil {
+			w.WriteHeader(500)
+			w.Write([]byte(err.Error()))
+			log.Println("ERROR", err.Error())
+			return
+		}
+
+		defer out.Close()
+		_, err = io.Copy(out, file)
+		if err != nil {
+			w.WriteHeader(500)
+			w.Write([]byte(err.Error()))
+			log.Println("ERROR", err.Error())
+			return
+		}
+
+		log.Println(r.Method, r.URL.Path, header.Filename, r.RemoteAddr)
+		w.Write([]byte("Uploaded " + header.Filename))
+	}
+	return http.HandlerFunc(fn)
+}
+
 func main() {
 	ip_f := flag.String("ip", "0.0.0.0", "IP adress to listen on.")
 	port_f := flag.String("port", "3000", "Port to listen on.")
 	dir_f := flag.String("dir", ".", "Directory to serve.")
+	disallow_upl_f := flag.Bool("disallow-upload", false, "Disallow uploads to /upload.")
+	upl_dir_f := flag.String("upload-dir", ".", "Directory to store uploaded files.")
 
 	flag.Parse()
 
@@ -28,5 +72,12 @@ func main() {
 	}
 
 	fmt.Println("Starting uhttpd serving \"" + *dir_f + "\" on " + *ip_f + ":" + port + ".")
-	log.Fatal(http.ListenAndServe(*ip_f+":"+port, accessLog(http.FileServer(http.Dir(*dir_f)))))
+
+	mux := http.NewServeMux()
+	if !*disallow_upl_f {
+		os.MkdirAll(*upl_dir_f, 0755)
+		mux.Handle("/upload", uploadHandler(*upl_dir_f))
+	}
+	mux.Handle("/", accessLog(http.FileServer(http.Dir(*dir_f))))
+	log.Fatal(http.ListenAndServe(*ip_f+":"+port, mux))
 }