Move complex crypto logic into a helper package

6 files changed, 302 insertions, 152 deletions

diff --git a/internal/cryptoutil/cryptoutil.go b/internal/cryptoutil/cryptoutil.go
new file mode 100644
index 0000000..c5d2a86
--- /dev/null
+++ b/internal/cryptoutil/cryptoutil.go
@@ -0,0 +1,167 @@
+// Package cryptoutil contains friendly wrappers around the algorithms from the
+// standard library's crypto package.
+package cryptoutil
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/hkdf"
+	"crypto/hmac"
+	"crypto/pbkdf2"
+	"crypto/rand"
+	"crypto/sha256"
+	"crypto/subtle"
+	"encoding/hex"
+	"errors"
+)
+
+const (
+	defaultIterations = 3749890 // from cmd/finditers
+
+	saltSize   = 2 * sha256.Size
+	hashSize   = 32
+	certSize   = sha256.Size
+	aesKeySize = 32
+	nonceSize  = aes.BlockSize
+)
+
+// A PasswordHash must be PasswordHashSize bytes.
+const PasswordHashSize = saltSize + hashSize
+
+// A PasswordHash is derived from the user's password and can be passed to
+// CheckPassword to verify if two passwords match. A PasswordHash must be
+// PasswordHashSize bytes.
+type PasswordHash []byte
+
+func (h PasswordHash) String() string { return hex.EncodeToString(h) }
+
+// A RawKey is generated from a password hash and can be used to derive further
+// key material.
+type RawKey struct {
+	key []byte
+}
+
+// An HMACKey must be HMACKeySize bytes.
+const HMACKeySize = sha256.Size
+
+// An HMAC key can be used to symmetrically sign and verify messages
+// using HMAC-SHA256. An HMACKey must be HMACKeyLen bytes.
+type HMACKey []byte
+
+// A SignedMessage is a message that has been cryptographically signed
+// with HMAC-SHA256.
+type SignedMessage []byte
+
+// An EncryptionKey must be EncryptionKeySize bytes.
+const EncryptionKeySize = aesKeySize + HMACKeySize
+
+// An EncryptionKey is used for encrypting and decrypting data. An EncryptionKey
+// must be EncryptionKeySize bytes.
+type EncryptionKey []byte
+
+// An EncryptedMessage is encrypted with AES-256-CTR-HMAC-SHA256.
+type EncryptedMessage []byte
+
+// HashIter runs pbkdf2-sha256 for iter iterations. Useful for benchmarking.
+func HashIter(password string, salt []byte, iter int) ([]byte, error) {
+	return pbkdf2.Key(sha256.New, password, salt, iter, sha256.Size)
+}
+
+func hashWithSalt(password string, salt []byte) (RawKey, []byte, error) {
+	hash, err := HashIter(password, salt[:sha256.Size], defaultIterations)
+	if err != nil {
+		return RawKey{}, nil, err
+	}
+	key, err := hkdf.Extract(sha256.New, hash, salt[sha256.Size:])
+	if err != nil {
+		return RawKey{}, nil, err
+	}
+	pwHash, err := hkdf.Expand(sha256.New, key, "pwhash", hashSize)
+	if err != nil {
+		return RawKey{}, nil, err
+	}
+	return RawKey{key}, pwHash, nil
+}
+
+// Hash hashes a user password using PBKDF2-SHA256.
+func Hash(password string) (PasswordHash, error) {
+	salt := make([]byte, saltSize, saltSize+hashSize)
+	rand.Read(salt)
+	_, pwHash, err := hashWithSalt(password, salt)
+	if err != nil {
+		return nil, err
+	}
+	return append(salt, pwHash...), nil
+}
+
+// CheckPassword verifies password against expectedHash and returns an
+// EncryptionKey derived from the password if successful.
+func CheckPassword(password string, expectedHash PasswordHash) (RawKey, error) {
+	salt, expectedHash := expectedHash[:saltSize], expectedHash[saltSize:]
+	key, pwHash, err := hashWithSalt(password, salt)
+	if err != nil {
+		return RawKey{}, err
+	}
+	if subtle.ConstantTimeCompare(pwHash, expectedHash) == 0 {
+		return RawKey{}, errors.New("incorrect password")
+	}
+	return key, nil
+}
+
+// Sign generates an HMAC-SHA256 signature and appends it to msg.
+func (k HMACKey) Sign(msg []byte) SignedMessage {
+	mac := hmac.New(sha256.New, k)
+	mac.Write(msg)
+	return mac.Sum(msg)
+}
+
+// Verify checks whether the given message has a valid signature, and returns
+// the raw message if it does.
+func (k HMACKey) Verify(msg SignedMessage) ([]byte, bool) {
+	if len(msg) < certSize {
+		return nil, false
+	}
+	msg, sig := msg[:len(msg)-certSize], msg[len(msg)-certSize:]
+	mac := hmac.New(sha256.New, k)
+	mac.Write(msg)
+	if !hmac.Equal(sig, mac.Sum(nil)) {
+		return nil, false
+	}
+	return msg, true
+}
+
+// EncryptionKey derives an EncryptionKey.
+func (k RawKey) EncryptionKey() (EncryptionKey, error) {
+	return hkdf.Expand(sha256.New, k.key, "encrypt", EncryptionKeySize)
+}
+
+// Encrypt encrypts a message with AES-256-CTR-HMAC-SHA256.
+func (k EncryptionKey) Encrypt(msg []byte) (EncryptedMessage, error) {
+	aesKey, hmacKey := k[:aesKeySize], HMACKey(k[aesKeySize:])
+	block, err := aes.NewCipher(aesKey)
+	if err != nil {
+		return nil, err
+	}
+	cipherText := make([]byte, len(msg)+nonceSize, len(msg)+nonceSize+certSize)
+	nonce := cipherText[len(msg) : len(msg)+nonceSize]
+	rand.Read(nonce)
+	cipher.NewCTR(block, nonce).XORKeyStream(cipherText, msg)
+	return EncryptedMessage(hmacKey.Sign(cipherText)), nil
+}
+
+// Decrypt verifies and decrypts an encrypted message. It overwrites msg with
+// the resulting plain text.
+func (k EncryptionKey) Decrypt(msg EncryptedMessage) ([]byte, error) {
+	aesKey, hmacKey := k[:aesKeySize], HMACKey(k[aesKeySize:])
+	msg, ok := hmacKey.Verify(SignedMessage(msg))
+	if !ok {
+		return nil, errors.New("bad signature")
+	}
+	block, err := aes.NewCipher(aesKey)
+	if err != nil {
+		return nil, err
+	}
+	msg, nonce := msg[:len(msg)-nonceSize], msg[len(msg)-nonceSize:]
+	cipher.NewCTR(block, nonce).XORKeyStream(msg, msg)
+	return msg, nil
+}
diff --git a/internal/cryptoutil/cryptoutil_test.go b/internal/cryptoutil/cryptoutil_test.go
new file mode 100644
index 0000000..6235878
--- /dev/null
+++ b/internal/cryptoutil/cryptoutil_test.go
@@ -0,0 +1,85 @@
+package cryptoutil
+
+import (
+	"bytes"
+	"encoding/hex"
+	"testing"
+)
+
+func mustHex(t *testing.T, s string) []byte {
+	t.Helper()
+	bytes, err := hex.DecodeString(s)
+	if err != nil {
+		t.Fatalf("Bad hex %q: %s", s, err)
+	}
+	return bytes
+}
+
+func TestPassword(t *testing.T) {
+	const password = "eclair scroll gliding sled shining removed ascend android cheer confetti"
+	hash, err := Hash(password)
+	if err != nil {
+		t.Fatalf("Hash(%q) failed: %s", password, err)
+	}
+	_, err = CheckPassword(password, hash)
+	if err != nil {
+		t.Errorf("CheckPassword(%q, %q) failed: %s", password, hash, err)
+	}
+}
+
+func TestSignature(t *testing.T) {
+	key := HMACKey(mustHex(t, "669e06ec457778b9a8133edb0a87ea82c6b141ffbbc63c038da96258175eb35c"))
+	msg := []byte("test message")
+	signedMsg := key.Sign(msg)
+	got, ok := key.Verify(signedMsg)
+	if !ok {
+		t.Fatalf("Verify(%x) rejected the message", signedMsg)
+	}
+	if !bytes.Equal(got, msg) {
+		t.Errorf("Verify(%x) = %x, want %x", signedMsg, got, msg)
+	}
+}
+
+func TestEncrypt(t *testing.T) {
+	key := EncryptionKey(mustHex(t, "b6de26860e0a39aa134732e58055c06ba028675453f73a6912bc932e58d5743d0e423217f06487d3e96a59980301fcc97dfcc4c6b19765f2947de3c33ab7ef9e"))
+	msg := []byte("test message")
+	encryptedMsg, err := key.Encrypt(msg)
+	if err != nil {
+		t.Fatalf("Encrypt(%q) failed: %s", msg, err)
+	}
+	got, err := key.Decrypt(encryptedMsg)
+	if err != nil {
+		t.Fatalf("Decrypt(%x) failed: %s", encryptedMsg, err)
+	}
+	if !bytes.Equal(got, msg) {
+		t.Errorf("Decrypt(%x) = %x, want %x", encryptedMsg, got, msg)
+	}
+}
+
+func TestDeriveKey(t *testing.T) {
+	const password = "automaker moisture botch tubular kelp rinse rule giggly donated stock"
+	hash, err := Hash(password)
+	if err != nil {
+		t.Fatalf("Hash(%q) failed: %s", password, err)
+	}
+	rawKey, err := CheckPassword(password, hash)
+	if err != nil {
+		t.Fatalf("CheckPassword(%q, %x) failed: %s", password, hash, err)
+	}
+	key, err := rawKey.EncryptionKey()
+	if err != nil {
+		t.Fatalf("EncryptionKey() failed: %s", err)
+	}
+	msg := []byte("test message")
+	encryptedMsg, err := key.Encrypt(msg)
+	if err != nil {
+		t.Fatalf("Encrypt(%q) failed: %s", msg, err)
+	}
+	got, err := key.Decrypt(encryptedMsg)
+	if err != nil {
+		t.Fatalf("Decrypt(%x) failed: %s", encryptedMsg, err)
+	}
+	if !bytes.Equal(got, msg) {
+		t.Errorf("Decrypt(%x) = %x, want %x", encryptedMsg, got, msg)
+	}
+}
diff --git a/internal/pwhash/pwhash.go b/internal/pwhash/pwhash.go
deleted file mode 100644
index 9ef5d2a..0000000
--- a/internal/pwhash/pwhash.go
+++ /dev/null
@@ -1,33 +0,0 @@
-package pwhash
-
-import (
-	"crypto/hkdf"
-	"crypto/pbkdf2"
-	"crypto/sha256"
-)
-
-const SaltLen = 2 * sha256.Size
-
-const defaultIterations = 3749890 // from cmd/finditers
-
-func HashIter(password string, salt []byte, iter int) ([]byte, error) {
-	return pbkdf2.Key(sha256.New, password, salt, iter, sha256.Size)
-}
-
-func Hash(password string, salt []byte) (rawKey, pwHash []byte, err error) {
-	const hashLen = 32
-
-	key, err := HashIter(password, salt[:sha256.Size], defaultIterations)
-	if err != nil {
-		return nil, nil, err
-	}
-	key, err = hkdf.Extract(sha256.New, key, salt[sha256.Size:])
-	if err != nil {
-		return nil, nil, err
-	}
-	pwHash, err = hkdf.Expand(sha256.New, key, "pwhash", hashLen)
-	if err != nil {
-		return nil, nil, err
-	}
-	return key, pwHash, nil
-}
diff --git a/roseh.moe.go b/roseh.moe.go
index 3c5e7a9..6624706 100644
--- a/roseh.moe.go
+++ b/roseh.moe.go
@@ -2,12 +2,7 @@ package main
 
 import (
 	"bytes"
-	"crypto/aes"
-	"crypto/cipher"
-	"crypto/hkdf"
-	"crypto/hmac"
 	"crypto/rand"
-	"crypto/sha256"
 	"crypto/subtle"
 	"embed"
 	"encoding/base64"
@@ -25,7 +20,7 @@ import (
 	"sync"
 	"time"
 
-	"gitlab.com/rhogenson/roseh.moe/internal/pwhash"
+	"gitlab.com/rhogenson/roseh.moe/internal/cryptoutil"
 )
 
 var (
@@ -37,12 +32,11 @@ var (
 )
 
 var (
-	notepadPassword     []byte
-	notepadPasswordSalt []byte
-	privateKey          []byte
+	notepadPassword cryptoutil.PasswordHash
+	secretKey       cryptoutil.HMACKey
 
 	encryptionKeyMu sync.Mutex
-	encryptionKey   []byte
+	encryptionKey   cryptoutil.EncryptionKey
 )
 
 func loadSecrets() error {
@@ -52,20 +46,19 @@ func loadSecrets() error {
 	}
 	for _, line := range bytes.Split(bytes.TrimSuffix(secrets, []byte("\n")), []byte("\n")) {
 		if pw, ok := bytes.CutPrefix(line, []byte("notepad-password=")); ok {
-			buf := make([]byte, hex.DecodedLen(len(pw)))
-			if _, err := hex.Decode(buf, pw); err != nil {
-				return err
+			if hex.DecodedLen(len(pw)) != cryptoutil.PasswordHashSize {
+				return fmt.Errorf("invalid PBKDF2-SHA256 hash")
 			}
-			if len(buf) < pwhash.SaltLen {
-				return fmt.Errorf("bad password hash")
+			notepadPassword = make(cryptoutil.PasswordHash, hex.DecodedLen(len(pw)))
+			if _, err := hex.Decode(notepadPassword, pw); err != nil {
+				return err
 			}
-			notepadPasswordSalt, notepadPassword = buf[:pwhash.SaltLen], buf[pwhash.SaltLen:]
 		} else if key, ok := bytes.CutPrefix(line, []byte("secret-key=")); ok {
-			if hex.DecodedLen(len(key)) != sha256.Size {
+			if hex.DecodedLen(len(key)) != cryptoutil.HMACKeySize {
 				return fmt.Errorf("invalid HMAC-SHA256 key")
 			}
-			privateKey = make([]byte, hex.DecodedLen(len(key)))
-			if _, err := hex.Decode(privateKey, key); err != nil {
+			secretKey = make(cryptoutil.HMACKey, hex.DecodedLen(len(key)))
+			if _, err := hex.Decode(secretKey, key); err != nil {
 				return err
 			}
 		}
@@ -73,78 +66,6 @@ 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")
-
-const (
-	aesKeyLen = 32
-	nonceLen  = aes.BlockSize
-	certLen   = sha256.Size
-)
-
-func encrypt(msg string) ([]byte, error) {
-	encryptionKeyMu.Lock()
-	key := encryptionKey
-	encryptionKeyMu.Unlock()
-	if key == nil {
-		return nil, errNoKey
-	}
-	block, err := aes.NewCipher(key[:aesKeyLen])
-	if err != nil {
-		return nil, err
-	}
-	buf := make([]byte, nonceLen+len(msg)+certLen)
-	nonce := buf[:nonceLen]
-	rand.Read(nonce)
-	cipher.NewCTR(block, nonce).XORKeyStream(buf[nonceLen:], []byte(msg))
-	cipherText := buf[:nonceLen+len(msg)]
-	mac := hmac.New(sha256.New, key[aesKeyLen:])
-	mac.Write(cipherText)
-	return mac.Sum(cipherText), nil
-}
-
-func decrypt(msg []byte) (string, error) {
-	encryptionKeyMu.Lock()
-	key := encryptionKey
-	encryptionKeyMu.Unlock()
-	if key == nil {
-		return "", errNoKey
-	}
-	if len(msg) < certLen+nonceLen {
-		return "", fmt.Errorf("message too short")
-	}
-	msg, messageMAC := msg[:len(msg)-certLen], msg[len(msg)-certLen:]
-	mac := hmac.New(sha256.New, key[aesKeyLen:])
-	mac.Write(msg)
-	expectedMAC := mac.Sum(nil)
-	if !hmac.Equal(messageMAC, expectedMAC) {
-		return "", fmt.Errorf("bad signature")
-	}
-	block, err := aes.NewCipher(key[:aesKeyLen])
-	if err != nil {
-		return "", err
-	}
-	nonce, msg := msg[:nonceLen], msg[nonceLen:]
-	buf := make([]byte, len(msg))
-	cipher.NewCTR(block, nonce).XORKeyStream(buf, msg)
-	return string(buf), nil
-}
-
 var (
 	//go:embed templates/404.html.template
 	notFoundString   string
@@ -219,7 +140,7 @@ func attachCookie(w http.ResponseWriter) error {
 	}
 	http.SetCookie(w, &http.Cookie{
 		Name:        "auth",
-		Value:       base64.RawStdEncoding.EncodeToString(sign(nowBytes)),
+		Value:       base64.RawStdEncoding.EncodeToString(secretKey.Sign(nowBytes)),
 		Path:        "/notepad",
 		Expires:     time.Now().Add(cookieExpiration),
 		Secure:      true,
@@ -241,7 +162,7 @@ func cookieAuth(w http.ResponseWriter, r *http.Request) (string, bool) {
 	if err != nil {
 		return "", false
 	}
-	msg, ok := verify(authCookie)
+	msg, ok := secretKey.Verify(authCookie)
 	if !ok {
 		return "", false
 	}
@@ -284,12 +205,8 @@ type loginTemplateArgs struct {
 }
 
 func login(w http.ResponseWriter, r *http.Request) {
-	rawKey, pwHash, err := pwhash.Hash(r.FormValue("password"), notepadPasswordSalt)
+	rawKey, err := cryptoutil.CheckPassword(r.FormValue("password"), notepadPassword)
 	if err != nil {
-		http.Error(w, fmt.Sprintf("Unable to hash password: %s", err), http.StatusInternalServerError)
-		return
-	}
-	if subtle.ConstantTimeCompare(pwHash[:], notepadPassword) == 0 {
 		if err := loginTemplate.Execute(w, loginTemplateArgs{Error: true}); err != nil {
 			log.Printf("Warning: login: %s", err)
 		}
@@ -299,7 +216,7 @@ func login(w http.ResponseWriter, r *http.Request) {
 	currentKey := encryptionKey
 	encryptionKeyMu.Unlock()
 	if currentKey == nil {
-		key, err := hkdf.Expand(sha256.New, rawKey, "encrypt", aesKeyLen+certLen)
+		key, err := rawKey.EncryptionKey()
 		if err != nil {
 			http.Error(w, fmt.Sprintf("failed to derive encryption key: %s", err), http.StatusInternalServerError)
 			return
@@ -314,16 +231,16 @@ func login(w http.ResponseWriter, r *http.Request) {
 	http.Redirect(w, r, "/notepad", http.StatusSeeOther)
 }
 
-func readNotepad() (string, error) {
+func readNotepad(key cryptoutil.EncryptionKey) (string, error) {
 	encrypted, err := os.ReadFile(*notepadFile)
 	if err != nil {
 		return "", err
 	}
-	decrypted, err := decrypt(encrypted)
+	decrypted, err := key.Decrypt(encrypted)
 	if err != nil {
 		return "", err
 	}
-	return decrypted, nil
+	return string(decrypted), nil
 }
 
 var (
@@ -345,14 +262,17 @@ func notepad(w http.ResponseWriter, r *http.Request) {
 		}
 		return
 	}
-	currentContent, err := readNotepad()
-	if err != nil {
-		if err == errNoKey {
-			if err := loginTemplate.Execute(w, loginTemplateArgs{}); err != nil {
-				log.Printf("Warning: login: %s", err)
-			}
-			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)
+	if err != nil {
 		currentContent = fmt.Sprintf("Error reading notepad file: %s", err)
 	}
 	if err := notepadTemplate.Execute(w, notepadTemplateArgs{Content: currentContent, CSRFToken: csrfToken}); err != nil {
@@ -367,7 +287,13 @@ func saveNote(w http.ResponseWriter, r *http.Request) error {
 	if _, ok := cookieAuth(w, r); !ok {
 		return fmt.Errorf("not logged in")
 	}
-	encrypted, err := encrypt(r.FormValue("content"))
+	encryptionKeyMu.Lock()
+	key := encryptionKey
+	encryptionKeyMu.Unlock()
+	if key == nil {
+		return errors.New("not logged in")
+	}
+	encrypted, err := key.Encrypt([]byte(r.FormValue("content")))
 	if err != nil {
 		return err
 	}
diff --git a/tools/finditers/finditers.go b/tools/finditers/finditers.go
index 2a256c8..cd144fa 100644
--- a/tools/finditers/finditers.go
+++ b/tools/finditers/finditers.go
@@ -8,19 +8,27 @@ import (
 	"testing"
 	"time"
 
-	"gitlab.com/rhogenson/roseh.moe/internal/pwhash"
+	"gitlab.com/rhogenson/roseh.moe/internal/cryptoutil"
 )
 
+func mustHex(s string) []byte {
+	bytes, err := hex.DecodeString(s)
+	if err != nil {
+		panic(fmt.Sprintf("mustHex: bad hex %q: %s", s, err))
+	}
+	return bytes
+}
+
 var (
 	iterations int
 
-	salt, _ = hex.DecodeString("3fb84513fc3afcd6d3b230bf9ece91aaae2d2a99da17efbf7de83b21fafe3f08")
+	salt = mustHex("3fb84513fc3afcd6d3b230bf9ece91aaae2d2a99da17efbf7de83b21fafe3f08")
 )
 
 func BenchmarkHashIter(b *testing.B) {
 	const password = "oboe shortness ether ideology undesired fresh freezable catching mashing glimpse"
 	for b.Loop() {
-		pwhash.HashIter(password, salt, iterations)
+		cryptoutil.HashIter(password, salt, iterations)
 	}
 }
 
diff --git a/tools/hashpw/hashpw.go b/tools/hashpw/hashpw.go
index 5bc740d..7705b31 100644
--- a/tools/hashpw/hashpw.go
+++ b/tools/hashpw/hashpw.go
@@ -1,11 +1,10 @@
 package main
 
 import (
-	"crypto/rand"
 	"fmt"
 	"os"
 
-	"gitlab.com/rhogenson/roseh.moe/internal/pwhash"
+	"gitlab.com/rhogenson/roseh.moe/internal/cryptoutil"
 	"golang.org/x/term"
 )
 
@@ -16,13 +15,11 @@ func run() error {
 	if err != nil {
 		return err
 	}
-	salt := make([]byte, pwhash.SaltLen)
-	rand.Read(salt)
-	_, hashedPassword, err := pwhash.Hash(string(password), salt)
+	hashedPassword, err := cryptoutil.Hash(string(password))
 	if err != nil {
 		return err
 	}
-	fmt.Printf("notepad-password=%x%x\n", salt, hashedPassword)
+	fmt.Printf("notepad-password=%s\n", hashedPassword)
 	return nil
 }