UCSPI-TCP support

[uploader.git] / main.go

diff --git a/main.go b/main.go
index f97f4df68647fcd4ce7b5a18a17d997577b2a86a..a6109acfb6345cffcf9dae244549398394ec5237 100644 (file)
--- a/main.go
+++ b/main.go
@@ -35,6 +35,7 @@ import (
        "os/exec"
        "strconv"
        "strings"
+       "sync"
        "time"
 
        "go.cypherpunks.ru/recfile"
@@ -57,7 +58,7 @@ var (
 Example command line usage:
 <pre>
 $ curl -F file=@somedata.tar.gpg [-F comment="optional comment"] http://.../upload/
-$ b2sum -a blake2b somedata.tar.gpg # to verify checksum
+$ b2sum somedata.tar.gpg # to verify BLAKE2b-512 checksum
 </pre>
 <form enctype="multipart/form-data" action="/upload/" method="post">
 <label for="file">File to upload:</label><br/>
@@ -68,9 +69,11 @@ $ b2sum -a blake2b somedata.tar.gpg # to verify checksum
 </form></body></html>`))
        NotifyFromAddr *string
        NotifyToAddr   *string
+       Jobs           sync.WaitGroup
 )
 
 func notify(tai, filename string, size int64, comment string) {
+       defer Jobs.Done()
        if *NotifyToAddr == "" {
                return
        }
@@ -181,41 +184,70 @@ func upload(w http.ResponseWriter, r *http.Request) {
        }
        io.Copy(w, &rec)
        log.Println(r.RemoteAddr, tai, fnOrig, n, sum)
+
+       rec.Reset()
+       wr = recfile.NewWriter(&rec)
+       if _, err = wr.WriteFields(
+               recfile.Field{Name: "Filename", Value: fnOrig},
+       ); err != nil {
+               log.Println(r.RemoteAddr, tai, fnOrig, n, sum, err)
+               return
+       }
+
+       var comment []byte
        p, err = mr.NextPart()
        if err != nil || p.FormName() != CommentFieldName {
-               go notify(fnOrig, tai, n, "")
-               return
+               goto Notify
        }
-       comment, err := ioutil.ReadAll(p)
+       comment, err = ioutil.ReadAll(p)
        if err != nil || len(comment) == 0 {
-               go notify(tai, fnOrig, n, "")
-               return
+               goto Notify
        }
-       ioutil.WriteFile(tai+".txt", comment, os.FileMode(0666))
+       wr.WriteFieldMultiline("Comment", strings.Split(string(comment), "\n"))
+
+Notify:
+       Jobs.Add(1)
        go notify(tai, fnOrig, n, string(comment))
+       ioutil.WriteFile(tai+".rec", rec.Bytes(), os.FileMode(0666))
 }
 
 func main() {
+       doUCSPI := flag.Bool("ucspi", false, "Work as UCSPI-TCP service")
        bind := flag.String("bind", "[::]:8086", "Address to bind to")
        conns := flag.Int("conns", 2, "Maximal number of connections")
        NotifyFromAddr = flag.String("notify-from", "uploader@example.com", "Address notifications are sent to")
        NotifyToAddr = flag.String("notify-to", "", "Address notifications are sent from")
        flag.Parse()
        log.SetFlags(log.Lshortfile)
-       log.SetOutput(os.Stdout)
+       if !*doUCSPI {
+               log.SetOutput(os.Stdout)
+       }
        if len(*NotifyFromAddr) == 0 && len(*NotifyToAddr) > 0 {
                log.Fatalln("notify-from address can not be empty, if notify-to is set")
        }
-       ln, err := net.Listen("tcp", *bind)
-       if err != nil {
-               log.Fatalln(err)
-       }
-       log.Println("listening", *bind)
-       ln = netutil.LimitListener(ln, *conns)
        s := &http.Server{
                ReadHeaderTimeout: 10 * time.Second,
                MaxHeaderBytes:    10 * (1 << 10),
        }
+       s.SetKeepAlivesEnabled(false)
        http.HandleFunc("/upload/", upload)
-       s.Serve(ln)
+       if *doUCSPI {
+               ln := &UCSPI{}
+               s.ConnState = connStater
+               err := s.Serve(ln)
+               if _, ok := err.(UCSPIAlreadyAccepted); !ok {
+                       log.Fatalln(err)
+               }
+               Jobs.Wait()
+       } else {
+               ln, err := net.Listen("tcp", *bind)
+               if err != nil {
+                       log.Fatalln(err)
+               }
+               log.Println("listening", *bind)
+               ln = netutil.LimitListener(ln, *conns)
+               if err = s.Serve(ln); err != nil {
+                       log.Fatalln(err)
+               }
+       }
 }