1 files changed, 67 insertions, 47 deletions

diff --git a/roseh.moe.go b/roseh.moe.go
index 0c7003f..b0614bb 100644
--- a/roseh.moe.go
+++ b/roseh.moe.go
@@ -6,11 +6,12 @@ import (
 	"crypto/cipher"
 	"crypto/hmac"
 	"crypto/rand"
-	"crypto/sha512"
+	"crypto/sha256"
 	"crypto/subtle"
 	"embed"
 	"encoding/base64"
 	"encoding/hex"
+	"errors"
 	"flag"
 	"fmt"
 	"html/template"
@@ -59,8 +60,8 @@ func loadSecrets() error {
 			}
 			notepadPasswordSalt, notepadPassword = buf[:pwhash.SaltLen], buf[pwhash.SaltLen:]
 		} else if key, ok := bytes.CutPrefix(line, []byte("secret-key=")); ok {
-			if hex.DecodedLen(len(key)) != sha512.Size256 {
-				return fmt.Errorf("invalid HMAC-SHA512/256 key")
+			if hex.DecodedLen(len(key)) != sha256.Size {
+				return fmt.Errorf("invalid HMAC-SHA256 key")
 			}
 			privateKey = make([]byte, hex.DecodedLen(len(key)))
 			if _, err := hex.Decode(privateKey, key); err != nil {
@@ -71,6 +72,54 @@ func loadSecrets() error {
 	return nil
 }
 
+func sign(msg []byte) []byte {
+	mac := hmac.New(sha256.New, privateKey)
+	mac.Write(msg)
+	return mac.Sum(msg)
+}
+
+func verify(msg []byte) ([]byte, bool) {
+	if len(msg) < sha256.Size {
+		return nil, false
+	}
+	msg, sig := msg[:len(msg)-sha256.Size], msg[len(msg)-sha256.Size:]
+	mac := hmac.New(sha256.New, privateKey)
+	mac.Write(msg)
+	return msg, hmac.Equal(sig, mac.Sum(nil))
+}
+
+var errNoKey = errors.New("not logged in")
+
+func loadEncryptionKey() (cipher.AEAD, error) {
+	encryptionKeyMu.Lock()
+	key := encryptionKey
+	encryptionKeyMu.Unlock()
+	if key == nil {
+		return nil, errNoKey
+	}
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+	return cipher.NewGCMWithRandomNonce(block)
+}
+
+func encrypt(msg []byte) ([]byte, error) {
+	aead, err := loadEncryptionKey()
+	if err != nil {
+		return nil, err
+	}
+	return aead.Seal(nil, nil, msg, nil), nil
+}
+
+func decrypt(msg []byte) ([]byte, error) {
+	aead, err := loadEncryptionKey()
+	if err != nil {
+		return nil, err
+	}
+	return aead.Open(nil, nil, msg, nil)
+}
+
 var (
 	//go:embed templates/404.html.template
 	notFoundString   string
@@ -143,11 +192,9 @@ func attachCookie(w http.ResponseWriter) error {
 	if err != nil {
 		return err
 	}
-	mac := hmac.New(sha512.New512_256, privateKey)
-	mac.Write(nowBytes)
 	http.SetCookie(w, &http.Cookie{
 		Name:        "auth",
-		Value:       base64.RawStdEncoding.EncodeToString(append(nowBytes, mac.Sum(nil)...)),
+		Value:       base64.RawStdEncoding.EncodeToString(sign(nowBytes)),
 		Path:        "/notepad",
 		Expires:     time.Now().Add(cookieExpiration),
 		Secure:      true,
@@ -165,18 +212,12 @@ func cookieAuth(w http.ResponseWriter, r *http.Request) (string, bool) {
 	if err != nil {
 		return "", false
 	}
-	bytes, err := base64.RawStdEncoding.DecodeString(cookie.Value)
+	authCookie, err := base64.RawStdEncoding.DecodeString(cookie.Value)
 	if err != nil {
 		return "", false
 	}
-	mac := hmac.New(sha512.New512_256, privateKey)
-	macSize := mac.Size()
-	if len(bytes) < macSize {
-		return "", false
-	}
-	msg, sig := bytes[:len(bytes)-macSize], bytes[len(bytes)-macSize:]
-	mac.Write(msg)
-	if !hmac.Equal(sig, mac.Sum(nil)) {
+	msg, ok := verify(authCookie)
+	if !ok {
 		return "", false
 	}
 	var t time.Time
@@ -238,20 +279,12 @@ func login(w http.ResponseWriter, r *http.Request) {
 	http.Redirect(w, r, "/notepad", http.StatusSeeOther)
 }
 
-func readNotepad(key []byte) (string, error) {
+func readNotepad() (string, error) {
 	encrypted, err := os.ReadFile(*notepadFile)
 	if err != nil {
 		return "", err
 	}
-	block, err := aes.NewCipher(key)
-	if err != nil {
-		return "", err
-	}
-	aead, err := cipher.NewGCMWithRandomNonce(block)
-	if err != nil {
-		return "", err
-	}
-	decrypted, err := aead.Open(nil, nil, encrypted, nil)
+	decrypted, err := decrypt(encrypted)
 	if err != nil {
 		return "", err
 	}
@@ -277,17 +310,14 @@ func notepad(w http.ResponseWriter, r *http.Request) {
 		}
 		return
 	}
-	encryptionKeyMu.Lock()
-	key := encryptionKey
-	encryptionKeyMu.Unlock()
-	if key == nil {
-		if err := loginTemplate.Execute(w, loginTemplateArgs{}); err != nil {
-			log.Printf("Warning: login: %s", err)
-		}
-		return
-	}
-	currentContent, err := readNotepad(key)
+	currentContent, err := readNotepad()
 	if err != nil {
+		if err == errNoKey {
+			if err := loginTemplate.Execute(w, loginTemplateArgs{}); err != nil {
+				log.Printf("Warning: login: %s", err)
+			}
+			return
+		}
 		currentContent = fmt.Sprintf("Error reading notepad file: %s", err)
 	}
 	if err := notepadTemplate.Execute(w, notepadTemplateArgs{Content: currentContent, CSRFToken: csrfToken}); err != nil {
@@ -302,17 +332,7 @@ func saveNote(w http.ResponseWriter, r *http.Request) error {
 	if _, ok := cookieAuth(w, r); !ok {
 		return fmt.Errorf("not logged in")
 	}
-	encryptionKeyMu.Lock()
-	key := encryptionKey
-	encryptionKeyMu.Unlock()
-	if key == nil {
-		return fmt.Errorf("not logged in")
-	}
-	block, err := aes.NewCipher(key)
-	if err != nil {
-		return err
-	}
-	aead, err := cipher.NewGCMWithRandomNonce(block)
+	encrypted, err := encrypt([]byte(r.FormValue("content")))
 	if err != nil {
 		return err
 	}
@@ -321,7 +341,7 @@ func saveNote(w http.ResponseWriter, r *http.Request) error {
 		return err
 	}
 	defer f.Close()
-	if _, err = f.Write(aead.Seal(nil, nil, []byte(r.FormValue("content")), nil)); err != nil {
+	if _, err = f.Write(encrypted); err != nil {
 		return err
 	}
 	if err := f.Close(); err != nil {