1 files changed, 36 insertions, 12 deletions
diff --git a/roseh.moe.go b/roseh.moe.go
index cd56353..b8e6e21 100644
--- a/roseh.moe.go
+++ b/roseh.moe.go
@@ -3,6 +3,7 @@ package main
 import (
 	"bytes"
 	"crypto/ed25519"
+	"crypto/rand"
 	"crypto/sha512"
 	"crypto/subtle"
 	"embed"
@@ -146,31 +147,49 @@ func attachCookie(w http.ResponseWriter) error {
 	return nil
 }
 
-func cookieAuth(w http.ResponseWriter, r *http.Request) bool {
+func cookieAuth(w http.ResponseWriter, r *http.Request) (string, bool) {
+	const csrfTokenLen = 32
+
 	cookie, err := r.Cookie("auth")
 	if err != nil {
-		return false
+		return "", false
 	}
 	bytes, err := base64.RawStdEncoding.DecodeString(cookie.Value)
 	if err != nil {
-		return false
+		return "", false
 	}
 	if len(bytes) < ed25519.SignatureSize {
-		return false
+		return "", false
 	}
 	msg, sig := bytes[:len(bytes)-ed25519.SignatureSize], bytes[len(bytes)-ed25519.SignatureSize:]
 	if !ed25519.Verify(publicKey, msg, sig) {
-		return false
+		return "", false
 	}
 	var t time.Time
 	if err := t.UnmarshalBinary(msg); err != nil {
-		return false
+		return "", false
 	}
 	if time.Since(t) > cookieExpiration {
-		return false
+		return "", false
 	}
 	attachCookie(w)
-	return true
+
+	if csrfToken, err := r.Cookie("csrf-token"); err == nil {
+		return csrfToken.Value, true
+	}
+	buf := make([]byte, csrfTokenLen)
+	rand.Read(buf)
+	csrfToken := base64.RawStdEncoding.EncodeToString(buf)
+	http.SetCookie(w, &http.Cookie{
+		Name:        "csrf-token",
+		Value:       csrfToken,
+		Path:        "/notepad",
+		Secure:      true,
+		HttpOnly:    true,
+		SameSite:    http.SameSiteStrictMode,
+		Partitioned: true,
+	})
+	return csrfToken, true
 }
 
 var (
@@ -203,11 +222,13 @@ var (
 )
 
 type notepadTemplateArgs struct {
-	Content string
+	Content   string
+	CSRFToken string
 }
 
 func notepad(w http.ResponseWriter, r *http.Request) {
-	if !cookieAuth(w, r) {
+	csrfToken, ok := cookieAuth(w, r)
+	if !ok {
 		if err := loginTemplate.Execute(w, loginTemplateArgs{Redirect: "/notepad"}); err != nil {
 			log.Printf("Warning: login: %s", err)
 		}
@@ -219,13 +240,16 @@ func notepad(w http.ResponseWriter, r *http.Request) {
 	} else {
 		currentContent = string(currentContentBytes)
 	}
-	if err := notepadTemplate.Execute(w, notepadTemplateArgs{Content: currentContent}); err != nil {
+	if err := notepadTemplate.Execute(w, notepadTemplateArgs{Content: currentContent, CSRFToken: csrfToken}); err != nil {
 		log.Printf("Warning: notepad: %s", err)
 	}
 }
 
 func saveNote(w http.ResponseWriter, r *http.Request) error {
-	if !cookieAuth(w, r) {
+	if csrfToken, err := r.Cookie("csrf-token"); err != nil || subtle.ConstantTimeCompare([]byte(csrfToken.Value), []byte(r.FormValue("csrf-token"))) == 0 {
+		return fmt.Errorf("bad CSRF token")
+	}
+	if _, ok := cookieAuth(w, r); !ok {
 		return fmt.Errorf("not logged in")
 	}
 	f, err := os.CreateTemp(filepath.Dir(*notepadFile), "notepad")