From be67bd0fb6fc3f64dcafe8d78bc0450520307518 Mon Sep 17 00:00:00 2001
From: Atridad Lahiji <me@atri.dad>
Date: Sun, 28 Dec 2025 02:03:52 -0700
Subject: [PATCH] fix persisting keys

---
 internal/database.go | 14 ++++++++++++++
 internal/model.go    | 24 ++++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/internal/database.go b/internal/database.go
index e764152..7cf0e8b 100644
--- a/internal/database.go
+++ b/internal/database.go
@@ -30,6 +30,9 @@ func InitDB() {
         identity_key BLOB,
         prekey BLOB,
         prekey_signature BLOB,
+        enc_priv_key BLOB,
+        priv_key_salt BLOB,
+        priv_key_nonce BLOB,
         created_at DATETIME
     );
     CREATE TABLE IF NOT EXISTS rooms (
@@ -97,6 +100,17 @@ func authorizeUser(pubkey, username string, identityKey, prekey, prekeySignature
 	return err
 }
 
+func storeUserEncryptedKey(username string, encKey, salt, nonce []byte) error {
+	_, err := db.Exec("UPDATE users SET enc_priv_key=?, priv_key_salt=?, priv_key_nonce=? WHERE username=?",
+		encKey, salt, nonce, username)
+	return err
+}
+
+func getUserEncryptedKey(username string) (encKey, salt, nonce []byte, err error) {
+	err = db.QueryRow("SELECT enc_priv_key, priv_key_salt, priv_key_nonce FROM users WHERE username = ?", username).Scan(&encKey, &salt, &nonce)
+	return
+}
+
 func getUsername(pubkey string) (string, error) {
 	var username string
 	err := db.QueryRow("SELECT username FROM users WHERE pubkey = ?", pubkey).Scan(&username)
diff --git a/internal/model.go b/internal/model.go
index 18dae1e..2a9d769 100644
--- a/internal/model.go
+++ b/internal/model.go
@@ -190,15 +190,39 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 			if m.state == 0 {
 				pass := m.input.Value()
 				if len(pass) > 0 {
+					// REGISTERING USER FLOW
 					if m.identityKey != nil {
 						ciphertext, nonce, salt, err := encryptUserKey(m.identityKey.PrivateKey, pass)
 						if err != nil {
 							m.err = err
 							return m, nil
 						}
+						// Store encrypted private key
+						if err := storeUserEncryptedKey(m.username, ciphertext, salt, nonce); err != nil {
+							m.err = err
+							return m, nil
+						}
 						m.encryptedPrivKey = ciphertext
 						m.privKeyNonce = nonce
 						m.privKeySalt = salt
+					} else {
+						// LOGGING IN USER FLOW
+						encKey, salt, nonce, err := getUserEncryptedKey(m.username)
+						if err != nil {
+							m.err = fmt.Errorf("failed to load keys: %v", err)
+							return m, nil
+						}
+						privKey, err := decryptUserKey(encKey, nonce, salt, pass)
+						if err != nil {
+							m.err = fmt.Errorf("incorrect passphrase")
+							return m, nil
+						}
+						// Reconstruct IdentityKeyPair
+						pubKey, _ := getMemberIdentityKey(m.username)
+						m.identityKey = &IdentityKeyPair{
+							PrivateKey: privKey,
+							PublicKey:  pubKey,
+						}
 					}
 
 					m.passphrase = pass