Make wormhole a little bit simpler

1 files changed, 144 insertions, 83 deletions

diff --git a/roseh.moe.go b/roseh.moe.go
index 993282a..b1f659e 100644
--- a/roseh.moe.go
+++ b/roseh.moe.go
@@ -1,7 +1,7 @@
 package main
 
 import (
-	"context"
+	"crypto/hkdf"
 	"crypto/hmac"
 	"crypto/rand"
 	"crypto/sha1"
@@ -10,6 +10,7 @@ import (
 	"embed"
 	"encoding/base64"
 	"encoding/binary"
+	"errors"
 	"flag"
 	"fmt"
 	"html/template"
@@ -27,6 +28,7 @@ import (
 
 	"github.com/skip2/go-qrcode"
 	"roseh.moe/pkg/ccl"
+	"roseh.moe/pkg/oae2"
 	"roseh.moe/pkg/roseh.moe/internal/config"
 	"roseh.moe/pkg/roseh.moe/internal/pwhash"
 	"roseh.moe/pkg/wordlist"
@@ -39,6 +41,7 @@ var (
 	secretsFile = flag.String("secrets", "secrets.ccl", "secrets file")
 	configFile  = flag.String("config", "config.ccl", "configuration file (ccl format https://pkg.go.dev/roseh.moe/pkg/ccl)")
 	https       = flag.String("https", "", "directory containing cert.pem and key.pem, or empty string to use unencrypted http")
+	holeTempDir = flag.String("hole-temp-dir", "hole", "directory to store temporary files for the wormhole")
 
 	serverStartTime = time.Now()
 )
@@ -187,18 +190,32 @@ func favicon(w http.ResponseWriter, r *http.Request) {
 	serveStaticFile(w, r, "static/favicon.ico")
 }
 
+func qrHandler(w http.ResponseWriter, r *http.Request) {
+	sig, err := base64.RawURLEncoding.DecodeString(r.FormValue("s"))
+	if err != nil {
+		http.Error(w, fmt.Sprintf("Bad signature: %s", err), http.StatusBadRequest)
+		return
+	}
+	text := r.FormValue("t")
+	if !verify([]byte(text), sig, "qr") {
+		http.Error(w, "Bad signature", http.StatusBadRequest)
+		return
+	}
+	qr, err := qrcode.Encode(text, qrcode.Medium, 200)
+	if err != nil {
+		http.Error(w, fmt.Sprintf("failed to encode QR code: %s", err), http.StatusInternalServerError)
+		return
+	}
+	w.Header().Set("Content-Type", "image/png")
+	w.Write(qr)
+}
+
 var (
 	//go:embed templates/wormhole.html.template
 	wormholeTemplateString string
 	wormholeTemplate       = template.Must(template.Must(outlineTemplate.Clone()).New("body").Parse(wormholeTemplateString)).Lookup("outline")
 )
 
-type wormholeTemplateArgs struct {
-	Self string
-	Hole string
-	Sig  string
-}
-
 func makeHole() string {
 	const nWords = 10
 	buf := make([]byte, 2*nWords)
@@ -210,59 +227,72 @@ func makeHole() string {
 	return strings.Join(words, "-")
 }
 
-func newWormhole(w http.ResponseWriter, r *http.Request) {
-	hole := makeHole()
-	http.Redirect(w, r, "/wormhole/"+hole+"/upload?s="+base64.RawURLEncoding.EncodeToString(mac([]byte(hole), "hole")), http.StatusSeeOther)
-}
+const blockSize = 4 * 1024 * 1024
 
-func wormhole(w http.ResponseWriter, r *http.Request) {
-	if err := wormholeTemplate.Execute(w, wormholeTemplateArgs{Self: *selfURL, Hole: r.PathValue("hole"), Sig: r.FormValue("s")}); err != nil {
-		log.Printf("Warning: wormhole: %s", err)
+func writeHoleFile(part *multipart.Part) (string, error) {
+	hole := makeHole()
+	key, err := hkdf.Key(sha256.New, []byte(hole), nil, "", 48)
+	if err != nil {
+		return "", err
 	}
+	primaryKey, encryptionKey := key[:32], key[32:]
+	f, err := os.Create(fmt.Sprintf("%s/%x", *holeTempDir, primaryKey))
+	if err != nil {
+		return "", err
+	}
+	defer f.Close()
+	w := oae2.NewWriter(f, encryptionKey, blockSize)
+	fileName := part.FileName()
+	buf := make([]byte, 8)
+	binary.BigEndian.PutUint64(buf, uint64(len(fileName)))
+	if _, err := w.Write(buf); err != nil {
+		return "", err
+	}
+	if _, err := io.WriteString(w, fileName); err != nil {
+		return "", err
+	}
+	if _, err := io.Copy(w, part); err != nil {
+		return "", err
+	}
+	if err := w.Close(); err != nil {
+		return "", err
+	}
+	return hole, f.Close()
 }
 
-type wormholeConn struct {
-	done chan struct{}
-	r    *multipart.Part
-	w    http.ResponseWriter
-}
-
-var wormholeConnsMu sync.Mutex
-var wormholeConns = make(map[string]wormholeConn)
+var errNoHole = errors.New("no such hole")
 
-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
-		}
+func readHoleFile(w http.ResponseWriter, hole string) error {
+	key, err := hkdf.Key(sha256.New, []byte(hole), nil, "", 48)
+	if err != nil {
+		return err
 	}
-	wormholeConnsMu.Unlock()
-	if c.w == nil {
-		c.w = prevConn.w
-	} else {
-		c.r = prevConn.r
+	primaryKey, encryptionKey := key[:32], key[32:]
+	f, err := os.Open(fmt.Sprintf("%s/%x", *holeTempDir, primaryKey))
+	if err != nil {
+		return errNoHole
 	}
-	if c.w == nil || c.r == nil {
-		return fmt.Errorf("connection already registered")
+	defer f.Close()
+	r := oae2.NewReader(f, encryptionKey, blockSize)
+	buf := make([]byte, 8)
+	if _, err := io.ReadFull(r, buf); err != nil {
+		return err
 	}
-	defer close(prevConn.done)
-	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)
+	fileNameLen := binary.BigEndian.Uint64(buf)
+	buf = make([]byte, fileNameLen)
+	if _, err := io.ReadFull(r, buf); err != nil {
+		return err
+	}
+	fileName := string(buf)
+	w.Header().Set("Content-Disposition", "attachment; filename*=UTF-8''"+fileName)
+	_, err = io.Copy(w, r)
+	return err
+}
+
+func wormhole(w http.ResponseWriter, r *http.Request) {
+	if err := wormholeTemplate.Execute(w, nil); err != nil {
+		log.Printf("Warning: wormhole: %s", err)
 	}
-	return nil
 }
 
 var (
@@ -271,8 +301,12 @@ var (
 	wormholeSuccessTemplate       = template.Must(template.Must(outlineTemplate.Clone()).New("body").Parse(wormholeSuccessTemplateString)).Lookup("outline")
 )
 
-func wormholeSend(w http.ResponseWriter, r *http.Request) {
-	ctx := r.Context()
+type wormholeSuccessTemplateArgs struct {
+	URL string
+	QR  string
+}
+
+func wormholeUpload(w http.ResponseWriter, r *http.Request) {
 	reader, err := r.MultipartReader()
 	if err != nil {
 		http.Error(w, fmt.Sprintf("Not a multipart/form-data request: %s", err), http.StatusBadRequest)
@@ -286,11 +320,17 @@ func wormholeSend(w http.ResponseWriter, r *http.Request) {
 		if part.FormName() != "file" {
 			continue
 		}
-		hole := r.PathValue("hole")
-		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)
+		hole, err := writeHoleFile(part)
+		if err != nil {
+			log.Printf("Warning: failed to open hole file: %s", err)
+			http.Error(w, "Internal error", http.StatusInternalServerError)
+			return
 		}
-		if err := wormholeSuccessTemplate.Execute(w, nil); err != nil {
+		downloadURL := *selfURL + "/wormhole/" + hole
+		if err := wormholeSuccessTemplate.Execute(w, wormholeSuccessTemplateArgs{
+			URL: downloadURL,
+			QR:  "/qr.png?t=" + url.QueryEscape(downloadURL) + "&s=" + base64.RawURLEncoding.EncodeToString(mac([]byte(downloadURL), "qr")),
+		}); err != nil {
 			log.Printf("Warning: wormholeSend: %s", err)
 		}
 		return
@@ -298,31 +338,51 @@ func wormholeSend(w http.ResponseWriter, r *http.Request) {
 	http.Error(w, "file not found", http.StatusBadRequest)
 }
 
-func wormholeRecv(w http.ResponseWriter, r *http.Request) {
-	ctx := r.Context()
-	hole := r.PathValue("hole")
-	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)
+func wormholeDownload(w http.ResponseWriter, r *http.Request) {
+	if err := readHoleFile(w, r.PathValue("hole")); err != nil {
+		if err == errNoHole {
+			notFound(w, r)
+			return
+		}
+		log.Printf("Warning: read hole file: %s", err)
+		return
 	}
 }
 
-func wormholeQR(w http.ResponseWriter, r *http.Request) {
-	sig, err := base64.RawURLEncoding.DecodeString(r.FormValue("s"))
-	if err != nil {
-		http.Error(w, fmt.Sprintf("Bad signature: %s", err), http.StatusBadRequest)
-		return
-	}
-	if !verify([]byte(r.PathValue("hole")), sig, "hole") {
-		http.Error(w, "Bad signature", http.StatusBadRequest)
-		return
-	}
-	qr, err := qrcode.Encode(*selfURL+"/wormhole/"+r.PathValue("hole"), qrcode.Medium, 200)
-	if err != nil {
-		http.Error(w, fmt.Sprintf("failed to encode QR code: %s", err), http.StatusInternalServerError)
-		return
+func cleanHole() {
+	for range time.NewTicker(time.Hour).C {
+		log.Printf("Cleaning hole...")
+		holeDir, err := os.Open(*holeTempDir)
+		if err != nil {
+			log.Printf("Warning: clean hole: %s", err)
+			continue
+		}
+		defer holeDir.Close()
+		deadline := time.Now().Add(-time.Hour)
+		for {
+			files, err := holeDir.ReadDir(1024)
+			for _, file := range files {
+				stat, err := file.Info()
+				if err != nil {
+					log.Printf("Warning: clean hole: %s", err)
+					continue
+				}
+				if !stat.ModTime().Before(deadline) {
+					continue
+				}
+				if err := os.Remove(*holeTempDir + "/" + file.Name()); err != nil {
+					log.Printf("Warning: clean hole: %s", err)
+					continue
+				}
+			}
+			if err != nil {
+				if err != io.EOF {
+					log.Printf("Warning: clean hole: %s", err)
+				}
+				break
+			}
+		}
 	}
-	w.Header().Set("Content-Type", "image/png")
-	w.Write(qr)
 }
 
 const macSize = sha256.Size
@@ -605,13 +665,14 @@ func main() {
 		log.Fatal(err)
 	}
 
+	go cleanHole()
+
 	mux := http.NewServeMux()
 	mux.HandleFunc("GET /pong", pong)
-	mux.HandleFunc("GET /wormhole", newWormhole)
-	mux.HandleFunc("GET /wormhole/{hole}/upload", wormhole)
-	mux.HandleFunc("POST /wormhole/{hole}/upload", wormholeSend)
-	mux.HandleFunc("GET /wormhole/{hole}", wormholeRecv)
-	mux.HandleFunc("GET /wormhole/{hole}/qr.png", wormholeQR)
+	mux.HandleFunc("GET /qr.png", qrHandler)
+	mux.HandleFunc("GET /wormhole", wormhole)
+	mux.HandleFunc("POST /wormhole", wormholeUpload)
+	mux.HandleFunc("GET /wormhole/{hole}", wormholeDownload)
 	mux.HandleFunc("GET /login", serveLogin)
 	mux.HandleFunc("POST /login", login)
 	mux.HandleFunc("GET /notepad", notepad)