simple-http-upload-server.go

//
// Start server:
//   $ ./simple-upload-server -port 10000 -root /path/to/upload/dir
//
// Upload file:
//   $ curl -i -T a-file-name host:10000/a-file-name
//

package main

import (
	"crypto/md5"
	"encoding/hex"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"syscall"
)

var Root string

const (
	minFreeSpace = 1 << 24 // 16 MiB
)

var nameRegexp = regexp.MustCompile(`^(?:[\w-]+/)*[\w.-]+$`)

func upload(w http.ResponseWriter, req *http.Request) {
	log.Println("request", req)

	if req.Method != "PUT" {
		http.Error(w, "only PUT is supported", http.StatusMethodNotAllowed)
		return
	}

	size := req.ContentLength
	if size < 0 {
		http.Error(w, "bad size", http.StatusLengthRequired)
		return
	}

	name := req.URL.Path[1:]
	if !nameRegexp.MatchString(name) {
		http.Error(w, "bad name", http.StatusBadRequest)
		return
	}

	freeSpace := freeDiskSpace(Root)
	if freeSpace < uint64(size)+minFreeSpace {
		log.Println("not enough space", freeSpace, size)
		http.Error(w, "disk full", http.StatusInsufficientStorage)
		return
	}

	path := filepath.Join(Root, name)
	dir, filename := filepath.Split(path)
	log.Println("mkdir", dir)
	err := os.MkdirAll(dir, 0755)
	if err != nil {
		log.Println(err)
		http.Error(w, "can't create directory", http.StatusInternalServerError)
		return
	}

	f, err := ioutil.TempFile(dir, filename)
	if err != nil {
		log.Println(err)
		http.Error(w, "can't create temp file", http.StatusInternalServerError)
		return
	}
	tmppath := f.Name()
	defer os.Remove(tmppath)
	log.Println("tempfile", tmppath)

	h := md5.New()
	r := io.TeeReader(req.Body, h)
	_, err = io.CopyN(f, r, size)
	if err != nil {
		log.Println(err)
		http.Error(w, "can't copy file", http.StatusInternalServerError)
		return
	}

	err = f.Close()
	if err != nil {
		log.Println(err)
		http.Error(w, "can't close file", http.StatusInternalServerError)
		return
	}

	log.Printf("renaming %s -> %s\n", tmppath, path)
	err = os.Rename(tmppath, path)
	if err != nil {
		log.Println(err)
		http.Error(w, "couldn't rename file", http.StatusInternalServerError)
		os.Remove(path)
		return
	}

	md5 := hex.EncodeToString(h.Sum(nil))
	log.Println("md5", md5)
	w.Header().Set("Etag", fmt.Sprintf(`"%s"`, md5))
}

func freeDiskSpace(path string) uint64 {
	var stat syscall.Statfs_t
	syscall.Statfs(path, &stat)
	return stat.Bavail * uint64(stat.Bsize)
}

func main() {
	var port int
	flag.IntVar(&port, "port", 12345, "port to listen on")
	flag.StringVar(&Root, "root", os.TempDir(), "upload directory")
	flag.Parse()

	log.Println("port", port)
	log.Println("root", Root)

	err := os.Chdir(Root)
	if err != nil {
		log.Fatal(err)
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/", upload)

	srv := &http.Server{
		Addr:    fmt.Sprintf(":%d", port),
		Handler: mux,
	}
	log.Fatal(srv.ListenAndServe())
}