1 files changed, 47 insertions, 68 deletions

diff --git a/roseh.moe.go b/roseh.moe.go
index 696b64e..3a53208 100644
--- a/roseh.moe.go
+++ b/roseh.moe.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"crypto/hmac"
 	"crypto/rand"
 	"crypto/sha256"
@@ -13,6 +14,7 @@ import (
 	"html/template"
 	"io"
 	"log"
+	"mime/multipart"
 	"net/http"
 	"os"
 	"strings"
@@ -198,24 +200,59 @@ func makeHole() string {
 	return strings.Join(words, "-")
 }
 
+func newWormhole(w http.ResponseWriter, r *http.Request) {
+	http.Redirect(w, r, "/wormhole/"+makeHole()+"/upload", http.StatusSeeOther)
+}
+
 func wormhole(w http.ResponseWriter, r *http.Request) {
-	if err := wormholeTemplate.Execute(w, wormholeTemplateArgs{Self: *selfURL, Hole: makeHole()}); err != nil {
+	if err := wormholeTemplate.Execute(w, wormholeTemplateArgs{Self: *selfURL, Hole: r.PathValue("hole")}); err != nil {
 		log.Printf("Warning: wormhole: %s", err)
 	}
 }
 
 type wormholeConn struct {
 	done chan struct{}
+	r    *multipart.Part
 	w    http.ResponseWriter
 }
 
 var wormholeConnsMu sync.Mutex
-var wormholeConns = make(map[string]*wormholeConn)
+var wormholeConns = make(map[string]wormholeConn)
 
-var wormholeNotifyMu sync.Mutex
-var wormholeNotify = make(map[string]chan struct{})
+func (c wormholeConn) wormholeCopy(ctx context.Context, hole string) error {
+	wormholeConnsMu.Lock()
+	prevConn, ok := wormholeConns[hole]
+	if !ok {
+		wormholeConns[hole] = c
+		wormholeConnsMu.Unlock()
+		defer func() {
+			wormholeConnsMu.Lock()
+			delete(wormholeConns, hole)
+			wormholeConnsMu.Unlock()
+		}()
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-c.done:
+			return nil
+		}
+	}
+	wormholeConnsMu.Unlock()
+	defer close(prevConn.done)
+	if c.w == nil {
+		c.w = prevConn.w
+	} else {
+		c.r = prevConn.r
+	}
+	c.w.Header().Set("Content-Disposition", "attachment; filename*=UTF-8''"+c.r.FileName())
+	if _, err := io.Copy(c.w, c.r); err != nil {
+		return fmt.Errorf("copy: %s", err)
+	}
+	return nil
+}
 
 func wormholeSend(w http.ResponseWriter, r *http.Request) {
+	ctx := r.Context()
 	reader, err := r.MultipartReader()
 	if err != nil {
 		http.Error(w, fmt.Sprintf("Not a multipart/form-data request: %s", err), http.StatusBadRequest)
@@ -230,19 +267,8 @@ func wormholeSend(w http.ResponseWriter, r *http.Request) {
 			continue
 		}
 		hole := r.PathValue("hole")
-		wormholeConnsMu.Lock()
-		conn := wormholeConns[hole]
-		delete(wormholeConns, hole)
-		wormholeConnsMu.Unlock()
-		if conn == nil {
-			http.Error(w, "no such connection", http.StatusBadRequest)
-			return
-		}
-		defer close(conn.done)
-		conn.w.Header().Set("Content-Disposition", "attachment; filename*=UTF-8''"+part.FileName())
-		if _, err := io.Copy(conn.w, part); err != nil {
+		if err := (wormholeConn{done: make(chan struct{}), r: part}).wormholeCopy(ctx, hole); err != nil {
 			http.Error(w, fmt.Sprintf("Error during copy: %s", err), http.StatusServiceUnavailable)
-			return
 		}
 		fmt.Fprintf(w, "uploaded!")
 		return
@@ -253,56 +279,9 @@ func wormholeSend(w http.ResponseWriter, r *http.Request) {
 func wormholeRecv(w http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 	hole := r.PathValue("hole")
-	conn := &wormholeConn{
-		done: make(chan struct{}),
-		w:    w,
-	}
-	wormholeConnsMu.Lock()
-	wormholeConns[hole] = conn
-	wormholeConnsMu.Unlock()
-	defer func() {
-		wormholeConnsMu.Lock()
-		delete(wormholeConns, hole)
-		wormholeConnsMu.Unlock()
-	}()
-	wormholeNotifyMu.Lock()
-	notify := wormholeNotify[hole]
-	wormholeNotifyMu.Unlock()
-	if notify == nil {
-		http.Error(w, "no such connection", http.StatusBadRequest)
-		return
-	}
-	select {
-	case notify <- struct{}{}:
-	default:
-		http.Error(w, "connection not ready", http.StatusBadRequest)
-		return
-	}
-	select {
-	case <-ctx.Done():
-	case <-conn.done:
-	}
-}
-
-func wormholeReady(w http.ResponseWriter, r *http.Request) {
-	ctx := r.Context()
-	hole := r.PathValue("hole")
-	notify := make(chan struct{})
-	wormholeNotifyMu.Lock()
-	wormholeNotify[hole] = notify
-	wormholeNotifyMu.Unlock()
-	defer func() {
-		wormholeNotifyMu.Lock()
-		delete(wormholeNotify, hole)
-		wormholeNotifyMu.Unlock()
-	}()
-	w.Header().Set("Content-Type", "text/event-stream")
-	select {
-	case <-ctx.Done():
-		return
-	case <-notify:
+	if err := (wormholeConn{done: make(chan struct{}), w: w}).wormholeCopy(ctx, hole); err != nil {
+		http.Error(w, fmt.Sprintf("Error during copy: %s", err), http.StatusServiceUnavailable)
 	}
-	fmt.Fprintf(w, "event: ready\ndata:\n\n")
 }
 
 func mac(msg []byte) []byte {
@@ -551,10 +530,10 @@ func main() {
 	}
 
 	http.HandleFunc("GET /pong", pong)
-	http.HandleFunc("GET /wormhole", wormhole)
-	http.HandleFunc("POST /wormhole/{hole}", wormholeSend)
+	http.HandleFunc("GET /wormhole", newWormhole)
+	http.HandleFunc("GET /wormhole/{hole}/upload", wormhole)
+	http.HandleFunc("POST /wormhole/{hole}/upload", wormholeSend)
 	http.HandleFunc("GET /wormhole/{hole}", wormholeRecv)
-	http.HandleFunc("GET /wormhole/{hole}/ready", wormholeReady)
 	http.HandleFunc("POST /login", login)
 	http.HandleFunc("GET /notepad", notepad)
 	http.HandleFunc("POST /notepad", autosave)