From 970ba1a907f5da4fbbbc2cf5c06ce275570305da Mon Sep 17 00:00:00 2001
From: Toni Spets <toni.spets@beeper.com>
Date: Mon, 15 Jan 2024 09:41:55 +0200
Subject: [PATCH] Store own device keys on init

---
 crypto/decryptolm.go |  2 +-
 crypto/machine.go    | 15 ++++++++++++---
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/crypto/decryptolm.go b/crypto/decryptolm.go
index f99c7dbe..68eaa875 100644
--- a/crypto/decryptolm.go
+++ b/crypto/decryptolm.go
@@ -216,7 +216,7 @@ func (mach *OlmMachine) createInboundSession(ctx context.Context, senderKey id.S
 	if err != nil {
 		return nil, err
 	}
-	mach.saveAccount()
+	mach.saveAccount(ctx)
 	err = mach.CryptoStore.AddSession(ctx, senderKey, session)
 	if err != nil {
 		zerolog.Ctx(ctx).Error().Err(err).Msg("Failed to store created inbound session")
diff --git a/crypto/machine.go b/crypto/machine.go
index 9892536a..b7c41ab0 100644
--- a/crypto/machine.go
+++ b/crypto/machine.go
@@ -145,8 +145,8 @@ func (mach *OlmMachine) Load(ctx context.Context) (err error) {
 	return nil
 }
 
-func (mach *OlmMachine) saveAccount() {
-	err := mach.CryptoStore.PutAccount(context.TODO(), mach.account)
+func (mach *OlmMachine) saveAccount(ctx context.Context) {
+	err := mach.CryptoStore.PutAccount(ctx, mach.account)
 	if err != nil {
 		mach.Log.Error().Err(err).Msg("Failed to save account")
 	}
@@ -655,6 +655,15 @@ func (mach *OlmMachine) ShareKeys(ctx context.Context, currentOTKCount int) erro
 	var deviceKeys *mautrix.DeviceKeys
 	if !mach.account.Shared {
 		deviceKeys = mach.account.getInitialKeys(mach.Client.UserID, mach.Client.DeviceID)
+		err := mach.CryptoStore.PutDevice(ctx, mach.Client.UserID, &id.Device{
+			UserID:      mach.Client.UserID,
+			DeviceID:    mach.Client.DeviceID,
+			IdentityKey: deviceKeys.Keys.GetCurve25519(mach.Client.DeviceID),
+			SigningKey:  deviceKeys.Keys.GetEd25519(mach.Client.DeviceID),
+		})
+		if err != nil {
+			return fmt.Errorf("failed to save initial keys: %w", err)
+		}
 		log.Debug().Msg("Going to upload initial account keys")
 	}
 	oneTimeKeys := mach.account.getOneTimeKeys(mach.Client.UserID, mach.Client.DeviceID, currentOTKCount)
@@ -673,7 +682,7 @@ func (mach *OlmMachine) ShareKeys(ctx context.Context, currentOTKCount int) erro
 	}
 	mach.lastOTKUpload = time.Now()
 	mach.account.Shared = true
-	mach.saveAccount()
+	mach.saveAccount(ctx)
 	return nil
 }