1 files changed, 55 insertions, 20 deletions

diff --git a/roseh.moe.go b/roseh.moe.go
index 2cdf68f..3c5e7a9 100644
--- a/roseh.moe.go
+++ b/roseh.moe.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"crypto/aes"
 	"crypto/cipher"
+	"crypto/hkdf"
 	"crypto/hmac"
 	"crypto/rand"
 	"crypto/sha256"
@@ -90,34 +91,58 @@ func verify(msg []byte) ([]byte, bool) {
 
 var errNoKey = errors.New("not logged in")
 
-func loadEncryptionKey() (cipher.AEAD, error) {
+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)
+	block, err := aes.NewCipher(key[:aesKeyLen])
 	if err != nil {
 		return nil, err
 	}
-	return cipher.NewGCMWithRandomNonce(block)
+	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 encrypt(msg []byte) ([]byte, error) {
-	aead, err := loadEncryptionKey()
-	if err != nil {
-		return nil, err
+func decrypt(msg []byte) (string, error) {
+	encryptionKeyMu.Lock()
+	key := encryptionKey
+	encryptionKeyMu.Unlock()
+	if key == nil {
+		return "", errNoKey
 	}
-	return aead.Seal(nil, nil, msg, nil), nil
-}
-
-func decrypt(msg []byte) ([]byte, error) {
-	aead, err := loadEncryptionKey()
+	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 nil, err
+		return "", err
 	}
-	return aead.Open(nil, nil, msg, nil)
+	nonce, msg := msg[:nonceLen], msg[nonceLen:]
+	buf := make([]byte, len(msg))
+	cipher.NewCTR(block, nonce).XORKeyStream(buf, msg)
+	return string(buf), nil
 }
 
 var (
@@ -259,7 +284,7 @@ type loginTemplateArgs struct {
 }
 
 func login(w http.ResponseWriter, r *http.Request) {
-	key, pwHash, err := pwhash.Hash(r.FormValue("password"), notepadPasswordSalt)
+	rawKey, pwHash, err := pwhash.Hash(r.FormValue("password"), notepadPasswordSalt)
 	if err != nil {
 		http.Error(w, fmt.Sprintf("Unable to hash password: %s", err), http.StatusInternalServerError)
 		return
@@ -271,10 +296,20 @@ func login(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	encryptionKeyMu.Lock()
-	if encryptionKey == nil {
-		encryptionKey = key
-	}
+	currentKey := encryptionKey
 	encryptionKeyMu.Unlock()
+	if currentKey == nil {
+		key, err := hkdf.Expand(sha256.New, rawKey, "encrypt", aesKeyLen+certLen)
+		if err != nil {
+			http.Error(w, fmt.Sprintf("failed to derive encryption key: %s", err), http.StatusInternalServerError)
+			return
+		}
+		encryptionKeyMu.Lock()
+		if encryptionKey == nil {
+			encryptionKey = key
+		}
+		encryptionKeyMu.Unlock()
+	}
 	attachCookie(w)
 	http.Redirect(w, r, "/notepad", http.StatusSeeOther)
 }
@@ -288,7 +323,7 @@ func readNotepad() (string, error) {
 	if err != nil {
 		return "", err
 	}
-	return string(decrypted), nil
+	return decrypted, nil
 }
 
 var (
@@ -332,7 +367,7 @@ func saveNote(w http.ResponseWriter, r *http.Request) error {
 	if _, ok := cookieAuth(w, r); !ok {
 		return fmt.Errorf("not logged in")
 	}
-	encrypted, err := encrypt([]byte(r.FormValue("content")))
+	encrypted, err := encrypt(r.FormValue("content"))
 	if err != nil {
 		return err
 	}