crypto: save source of megolm sessions

crypto/keybackup.go

@@ -200,13 +200,14 @@ func (mach *OlmMachine) ImportRoomKeyFromBackupWithoutSaving(
 		SigningKey:       keyBackupData.SenderClaimedKeys.Ed25519,
 		SenderKey:        keyBackupData.SenderKey,
 		RoomID:           roomID,
-		ForwardingChains: append(keyBackupData.ForwardingKeyChain, keyBackupData.SenderKey.String()),
+		ForwardingChains: keyBackupData.ForwardingKeyChain,
 		id:               sessionID,
 
 		ReceivedAt:       time.Now().UTC(),
 		MaxAge:           maxAge.Milliseconds(),
 		MaxMessages:      maxMessages,
 		KeyBackupVersion: version,
+		KeySource:        id.KeySourceBackup,
 	}, nil
 }
 

crypto/keyimport.go

@@ -108,14 +108,13 @@ func (mach *OlmMachine) importExportedRoomKey(ctx context.Context, session Expor
 		return false, ErrMismatchingExportedSessionID
 	}
 	igs := &InboundGroupSession{
-		Internal:   igsInternal,
-		SigningKey: session.SenderClaimedKeys.Ed25519,
-		SenderKey:  session.SenderKey,
-		RoomID:     session.RoomID,
-		// TODO should we add something here to mark the signing key as unverified like key requests do?
+		Internal:         igsInternal,
+		SigningKey:       session.SenderClaimedKeys.Ed25519,
+		SenderKey:        session.SenderKey,
+		RoomID:           session.RoomID,
 		ForwardingChains: session.ForwardingChains,
-
-		ReceivedAt: time.Now().UTC(),
+		KeySource:        id.KeySourceImport,
+		ReceivedAt:       time.Now().UTC(),
 	}
 	existingIGS, _ := mach.CryptoStore.GetGroupSession(ctx, igs.RoomID, igs.ID())
 	firstKnownIndex := igs.Internal.FirstKnownIndex()

crypto/keysharing.go

@@ -189,6 +189,7 @@ func (mach *OlmMachine) importForwardedRoomKey(ctx context.Context, evt *Decrypt
 		MaxAge:      maxAge.Milliseconds(),
 		MaxMessages: maxMessages,
 		IsScheduled: content.IsScheduled,
+		KeySource:   id.KeySourceForward,
 	}
 	existingIGS, _ := mach.CryptoStore.GetGroupSession(ctx, igs.RoomID, igs.ID())
 	if existingIGS != nil && existingIGS.Internal.FirstKnownIndex() <= igs.Internal.FirstKnownIndex() {

crypto/sessions.go

@@ -117,6 +117,7 @@ type InboundGroupSession struct {
 	MaxMessages      int
 	IsScheduled      bool
 	KeyBackupVersion id.KeyBackupVersion
+	KeySource        id.KeySource
 
 	id id.SessionID
 }
@@ -136,6 +137,7 @@ func NewInboundGroupSession(senderKey id.SenderKey, signingKey id.Ed25519, roomI
 		MaxAge:           maxAge.Milliseconds(),
 		MaxMessages:      maxMessages,
 		IsScheduled:      isScheduled,
+		KeySource:        id.KeySourceDirect,
 	}, nil
 }
 

crypto/sql_store.go

@@ -346,22 +346,23 @@ func (store *SQLCryptoStore) PutGroupSession(ctx context.Context, session *Inbou
 		Int("max_messages", session.MaxMessages).
 		Bool("is_scheduled", session.IsScheduled).
 		Stringer("key_backup_version", session.KeyBackupVersion).
+		Stringer("key_source", session.KeySource).
 		Msg("Upserting megolm inbound group session")
 	_, err = store.DB.Exec(ctx, `
 		INSERT INTO crypto_megolm_inbound_session (
 			session_id, sender_key, signing_key, room_id, session, forwarding_chains,
-			ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version, account_id
-		) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)
+			ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version, key_source, account_id
+		) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)
 		ON CONFLICT (session_id, account_id) DO UPDATE
 		    SET withheld_code=NULL, withheld_reason=NULL, sender_key=excluded.sender_key, signing_key=excluded.signing_key,
 		        room_id=excluded.room_id, session=excluded.session, forwarding_chains=excluded.forwarding_chains,
 		        ratchet_safety=excluded.ratchet_safety, received_at=excluded.received_at,
 		        max_age=excluded.max_age, max_messages=excluded.max_messages, is_scheduled=excluded.is_scheduled,
-		        key_backup_version=excluded.key_backup_version
+		        key_backup_version=excluded.key_backup_version, key_source=excluded.key_source
 	`,
 		session.ID(), session.SenderKey, session.SigningKey, session.RoomID, sessionBytes, forwardingChains,
 		ratchetSafety, datePtr(session.ReceivedAt), dbutil.NumPtr(session.MaxAge), dbutil.NumPtr(session.MaxMessages),
-		session.IsScheduled, session.KeyBackupVersion, store.AccountID,
+		session.IsScheduled, session.KeyBackupVersion, session.KeySource, store.AccountID,
 	)
 	return err
 }
@@ -374,12 +375,13 @@ func (store *SQLCryptoStore) GetGroupSession(ctx context.Context, roomID id.Room
 	var maxAge, maxMessages sql.NullInt64
 	var isScheduled bool
 	var version id.KeyBackupVersion
+	var keySource id.KeySource
 	err := store.DB.QueryRow(ctx, `
-		SELECT sender_key, signing_key, session, forwarding_chains, withheld_code, withheld_reason, ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version
+		SELECT sender_key, signing_key, session, forwarding_chains, withheld_code, withheld_reason, ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version, key_source
 		FROM crypto_megolm_inbound_session
 		WHERE room_id=$1 AND session_id=$2 AND account_id=$3`,
 		roomID, sessionID, store.AccountID,
-	).Scan(&senderKey, &signingKey, &sessionBytes, &forwardingChains, &withheldCode, &withheldReason, &ratchetSafetyBytes, &receivedAt, &maxAge, &maxMessages, &isScheduled, &version)
+	).Scan(&senderKey, &signingKey, &sessionBytes, &forwardingChains, &withheldCode, &withheldReason, &ratchetSafetyBytes, &receivedAt, &maxAge, &maxMessages, &isScheduled, &version, &keySource)
 	if errors.Is(err, sql.ErrNoRows) {
 		return nil, nil
 	} else if err != nil {
@@ -410,6 +412,7 @@ func (store *SQLCryptoStore) GetGroupSession(ctx context.Context, roomID id.Room
 		MaxMessages:      int(maxMessages.Int64),
 		IsScheduled:      isScheduled,
 		KeyBackupVersion: version,
+		KeySource:        keySource,
 	}, nil
 }
 
@@ -534,7 +537,8 @@ func (store *SQLCryptoStore) scanInboundGroupSession(rows dbutil.Scannable) (*In
 	var maxAge, maxMessages sql.NullInt64
 	var isScheduled bool
 	var version id.KeyBackupVersion
-	err := rows.Scan(&roomID, &senderKey, &signingKey, &sessionBytes, &forwardingChains, &ratchetSafetyBytes, &receivedAt, &maxAge, &maxMessages, &isScheduled, &version)
+	var keySource id.KeySource
+	err := rows.Scan(&roomID, &senderKey, &signingKey, &sessionBytes, &forwardingChains, &ratchetSafetyBytes, &receivedAt, &maxAge, &maxMessages, &isScheduled, &version, &keySource)
 	if err != nil {
 		return nil, err
 	}
@@ -554,12 +558,13 @@ func (store *SQLCryptoStore) scanInboundGroupSession(rows dbutil.Scannable) (*In
 		MaxMessages:      int(maxMessages.Int64),
 		IsScheduled:      isScheduled,
 		KeyBackupVersion: version,
+		KeySource:        keySource,
 	}, nil
 }
 
 func (store *SQLCryptoStore) GetGroupSessionsForRoom(ctx context.Context, roomID id.RoomID) dbutil.RowIter[*InboundGroupSession] {
 	rows, err := store.DB.Query(ctx, `
-		SELECT room_id, sender_key, signing_key, session, forwarding_chains, ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version
+		SELECT room_id, sender_key, signing_key, session, forwarding_chains, ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version, key_source
 		FROM crypto_megolm_inbound_session WHERE room_id=$1 AND account_id=$2 AND session IS NOT NULL`,
 		roomID, store.AccountID,
 	)
@@ -568,7 +573,7 @@ func (store *SQLCryptoStore) GetGroupSessionsForRoom(ctx context.Context, roomID
 
 func (store *SQLCryptoStore) GetAllGroupSessions(ctx context.Context) dbutil.RowIter[*InboundGroupSession] {
 	rows, err := store.DB.Query(ctx, `
-		SELECT room_id, sender_key, signing_key, session, forwarding_chains, ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version
+		SELECT room_id, sender_key, signing_key, session, forwarding_chains, ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version, key_source
 		FROM crypto_megolm_inbound_session WHERE account_id=$1 AND session IS NOT NULL`,
 		store.AccountID,
 	)
@@ -577,7 +582,7 @@ func (store *SQLCryptoStore) GetAllGroupSessions(ctx context.Context) dbutil.Row
 
 func (store *SQLCryptoStore) GetGroupSessionsWithoutKeyBackupVersion(ctx context.Context, version id.KeyBackupVersion) dbutil.RowIter[*InboundGroupSession] {
 	rows, err := store.DB.Query(ctx, `
-		SELECT room_id, sender_key, signing_key, session, forwarding_chains, ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version
+		SELECT room_id, sender_key, signing_key, session, forwarding_chains, ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version, key_source
 		FROM crypto_megolm_inbound_session WHERE account_id=$1 AND session IS NOT NULL AND key_backup_version != $2`,
 		store.AccountID, version,
 	)

crypto/sql_store_upgrade/00-latest-revision.sql

@@ -1,4 +1,4 @@
--- v0 -> v18 (compatible with v15+): Latest revision
+-- v0 -> v19 (compatible with v15+): Latest revision
 CREATE TABLE IF NOT EXISTS crypto_account (
 	account_id         TEXT    PRIMARY KEY,
 	device_id          TEXT    NOT NULL,
@@ -71,6 +71,7 @@ CREATE TABLE IF NOT EXISTS crypto_megolm_inbound_session (
 	max_messages       INTEGER,
 	is_scheduled       BOOLEAN NOT NULL DEFAULT false,
 	key_backup_version TEXT NOT NULL DEFAULT '',
+	key_source         TEXT NOT NULL DEFAULT '',
 	PRIMARY KEY (account_id, session_id)
 );
 -- Useful index to find keys that need backing up

crypto/sql_store_upgrade/19-megolm-session-source.sql

@@ -0,0 +1,2 @@
+-- v19 (compatible with v15+): Store megolm session source
+ALTER TABLE crypto_megolm_inbound_session ADD COLUMN key_source TEXT NOT NULL DEFAULT '';

id/crypto.go

@@ -53,6 +53,34 @@ const (
 	KeyBackupAlgorithmMegolmBackupV1 KeyBackupAlgorithm = "m.megolm_backup.v1.curve25519-aes-sha2"
 )
 
+type KeySource string
+
+func (source KeySource) String() string {
+	return string(source)
+}
+
+func (source KeySource) Int() int {
+	switch source {
+	case KeySourceDirect:
+		return 100
+	case KeySourceBackup:
+		return 90
+	case KeySourceImport:
+		return 80
+	case KeySourceForward:
+		return 50
+	default:
+		return 0
+	}
+}
+
+const (
+	KeySourceDirect  KeySource = "direct"
+	KeySourceBackup  KeySource = "backup"
+	KeySourceImport  KeySource = "import"
+	KeySourceForward KeySource = "forward"
+)
+
 // BackupVersion is an arbitrary string that identifies a server side key backup.
 type KeyBackupVersion string