fix persisting keys

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)

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