crypto: store session creation time to prevent out of order redaction

crypto/devicelist.go

@@ -12,6 +12,7 @@ import (
 	"fmt"
 	"slices"
 	"strings"
+	"time"
 
 	"github.com/rs/zerolog"
 	"go.mau.fi/util/exzerolog"
@@ -280,7 +281,7 @@ func (mach *OlmMachine) FetchKeys(ctx context.Context, users []id.UserID, includ
 						Str("identity_key", device.IdentityKey.String()).
 						Str("signing_key", device.SigningKey.String()).
 						Logger()
-					sessionIDs, err := mach.CryptoStore.RedactGroupSessions(ctx, "", device.IdentityKey, "device removed")
+					sessionIDs, err := mach.CryptoStore.RedactGroupSessions(ctx, "", device.IdentityKey, "device removed", time.Time{})
 					if err != nil {
 						log.Err(err).Msg("Failed to redact megolm sessions from deleted device")
 					} else {

crypto/encryptmegolm.go

@@ -211,7 +211,7 @@ func (mach *OlmMachine) newOutboundGroupSession(ctx context.Context, roomID id.R
 		signingKey, idKey := mach.account.Keys()
 		err = mach.createGroupSession(
 			ctx, mach.Client.UserID, idKey, signingKey, roomID, session.ID(), session.Internal.Key(),
-			session.MaxAge, session.MaxMessages, session.SharedHistory, false,
+			session.MaxAge, session.MaxMessages, session.SharedHistory, false, session.CreationTime,
 		)
 		if err != nil {
 			return nil, err

crypto/machine.go

@@ -653,7 +653,7 @@ func (mach *OlmMachine) SendEncryptedToDevice(ctx context.Context, device *id.De
 
 func (mach *OlmMachine) createGroupSession(
 	ctx context.Context, sender id.UserID, senderKey id.SenderKey, signingKey id.Ed25519, roomID id.RoomID, sessionID id.SessionID,
-	sessionKey string, maxAge time.Duration, maxMessages int, sharedHistory *bool, isScheduled bool,
+	sessionKey string, maxAge time.Duration, maxMessages int, sharedHistory *bool, isScheduled bool, creationTS time.Time,
 ) error {
 	log := zerolog.Ctx(ctx)
 	igs, err := NewInboundGroupSession(senderKey, signingKey, roomID, sessionKey, maxAge, maxMessages, sharedHistory, isScheduled)
@@ -667,6 +667,7 @@ func (mach *OlmMachine) createGroupSession(
 		return fmt.Errorf("mismatched session ID while creating inbound group session")
 	}
 	igs.SourceUser = sender
+	igs.CreationTS = creationTS
 	err = mach.StoreGroupSession(ctx, igs)
 	if err != nil {
 		log.Err(err).Stringer("session_id", sessionID).Msg("Failed to store new inbound group session")
@@ -803,8 +804,10 @@ func (mach *OlmMachine) receiveRoomKey(ctx context.Context, evt *DecryptedOlmEve
 	}
 	// TODO(history sharing): fill shared history with current state if it's unset?
 	if mach.DeletePreviousKeysOnReceive && !content.IsScheduled {
-		log.Debug().Msg("Redacting previous megolm sessions from sender in room")
-		sessionIDs, err := mach.CryptoStore.RedactGroupSessions(ctx, content.RoomID, evt.SenderKey, "received new key from device")
+		// Only redact sessions the sender created before this one, in case they arrive out of order
+		before := content.CreationTS.Time
+		log.Debug().Time("before", before).Msg("Redacting previous megolm sessions from sender in room")
+		sessionIDs, err := mach.CryptoStore.RedactGroupSessions(ctx, content.RoomID, evt.SenderKey, "received new key from device", before)
 		if err != nil {
 			log.Err(err).Msg("Failed to redact previous megolm sessions")
 		} else {
@@ -815,7 +818,7 @@ func (mach *OlmMachine) receiveRoomKey(ctx context.Context, evt *DecryptedOlmEve
 	}
 	err = mach.createGroupSession(
 		ctx, evt.Sender, evt.SenderKey, evt.Keys.Ed25519, content.RoomID, content.SessionID, content.SessionKey,
-		maxAge, maxMessages, content.SharedHistory, content.IsScheduled,
+		maxAge, maxMessages, content.SharedHistory, content.IsScheduled, content.CreationTS.Time,
 	)
 	if err != nil {
 		log.Err(err).Msg("Failed to create inbound group session")

crypto/sessions.go

@@ -114,6 +114,7 @@ type InboundGroupSession struct {
 	RatchetSafety    RatchetSafety
 
 	ReceivedAt       time.Time
+	CreationTS       time.Time
 	MaxAge           int64
 	MaxMessages      int
 	SharedHistory    *bool

crypto/sql_store.go

@@ -351,6 +351,7 @@ func (store *SQLCryptoStore) PutGroupSession(ctx context.Context, session *Inbou
 		Stringer("signing_key", session.SigningKey).
 		Stringer("room_id", session.RoomID).
 		Time("received_at", session.ReceivedAt).
+		Time("session_creation_ts", session.CreationTS).
 		Int64("max_age", session.MaxAge).
 		Int("max_messages", session.MaxMessages).
 		Bool("is_scheduled", session.IsScheduled).
@@ -364,18 +365,19 @@ func (store *SQLCryptoStore) PutGroupSession(ctx context.Context, session *Inbou
 		INSERT INTO crypto_megolm_inbound_session (
 			session_id, sender_key, signing_key, room_id, session, forwarding_chains, shared_history,
 			ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version, key_source,
-			source_user, account_id
-		) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)
+			source_user, session_creation_ts, account_id
+		) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17)
 		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,
 		        shared_history=excluded.shared_history, 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_source=excluded.key_source, source_user=excluded.source_user
+		        key_backup_version=excluded.key_backup_version, key_source=excluded.key_source, source_user=excluded.source_user,
+		        session_creation_ts=excluded.session_creation_ts
 	`,
 		session.ID(), session.SenderKey, session.SigningKey, session.RoomID, sessionBytes, forwardingChains, session.SharedHistory,
 		ratchetSafety, datePtr(session.ReceivedAt), dbutil.NumPtr(session.MaxAge), dbutil.NumPtr(session.MaxMessages),
-		session.IsScheduled, session.KeyBackupVersion, session.KeySource, session.SourceUser, store.AccountID,
+		session.IsScheduled, session.KeyBackupVersion, session.KeySource, session.SourceUser, datePtr(session.CreationTS), store.AccountID,
 	)
 	return err
 }
@@ -384,7 +386,7 @@ func (store *SQLCryptoStore) PutGroupSession(ctx context.Context, session *Inbou
 func (store *SQLCryptoStore) GetGroupSession(ctx context.Context, roomID id.RoomID, sessionID id.SessionID) (*InboundGroupSession, error) {
 	var senderKey, signingKey, forwardingChains, withheldCode, withheldReason sql.NullString
 	var sessionBytes, ratchetSafetyBytes []byte
-	var receivedAt sql.NullTime
+	var receivedAt, creationTS sql.NullTime
 	var maxAge, maxMessages sql.NullInt64
 	var isScheduled bool
 	var sharedHistory *bool
@@ -393,13 +395,13 @@ func (store *SQLCryptoStore) GetGroupSession(ctx context.Context, roomID id.Room
 	var sourceUser id.UserID
 	err := store.DB.QueryRow(ctx, `
 		SELECT sender_key, signing_key, session, forwarding_chains, withheld_code, withheld_reason, shared_history,
-		       ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version, key_source, source_user
+		       ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version, key_source, source_user, session_creation_ts
 		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, &sharedHistory,
-		&ratchetSafetyBytes, &receivedAt, &maxAge, &maxMessages, &isScheduled, &version, &keySource, &sourceUser,
+		&ratchetSafetyBytes, &receivedAt, &maxAge, &maxMessages, &isScheduled, &version, &keySource, &sourceUser, &creationTS,
 	)
 	if errors.Is(err, sql.ErrNoRows) {
 		return nil, nil
@@ -427,6 +429,7 @@ func (store *SQLCryptoStore) GetGroupSession(ctx context.Context, roomID id.Room
 		ForwardingChains: chains,
 		RatchetSafety:    rs,
 		ReceivedAt:       receivedAt.Time,
+		CreationTS:       creationTS.Time,
 		MaxAge:           maxAge.Int64,
 		MaxMessages:      int(maxMessages.Int64),
 		SharedHistory:    sharedHistory,
@@ -446,17 +449,23 @@ func (store *SQLCryptoStore) RedactGroupSession(ctx context.Context, _ id.RoomID
 	return err
 }
 
-func (store *SQLCryptoStore) RedactGroupSessions(ctx context.Context, roomID id.RoomID, senderKey id.SenderKey, reason string) ([]id.SessionID, error) {
+func (store *SQLCryptoStore) RedactGroupSessions(ctx context.Context, roomID id.RoomID, senderKey id.SenderKey, reason string, before time.Time) ([]id.SessionID, error) {
 	if roomID == "" && senderKey == "" {
 		return nil, fmt.Errorf("room ID or sender key must be provided for redacting sessions")
 	}
+	args := []any{event.RoomKeyWithheldBeeperRedacted, "Session redacted: " + reason, roomID, senderKey, store.AccountID}
+	creationTSFilter := ""
+	if !before.IsZero() {
+		creationTSFilter = " AND session_creation_ts IS NOT NULL AND session_creation_ts < $6"
+		args = append(args, before)
+	}
 	res, err := store.DB.Query(ctx, `
 		UPDATE crypto_megolm_inbound_session
 		SET withheld_code=$1, withheld_reason=$2, session=NULL, forwarding_chains=NULL
 		WHERE (room_id=$3 OR $3='') AND (sender_key=$4 OR $4='') AND account_id=$5
-		  AND session IS NOT NULL AND is_scheduled=false AND received_at IS NOT NULL
+		  AND session IS NOT NULL AND is_scheduled=false AND received_at IS NOT NULL`+creationTSFilter+`
 		RETURNING session_id
-	`, event.RoomKeyWithheldBeeperRedacted, "Session redacted: "+reason, roomID, senderKey, store.AccountID)
+	`, args...)
 	return dbutil.NewRowIterWithError(res, dbutil.ScanSingleColumn[id.SessionID], err).AsList()
 }
 
@@ -554,7 +563,7 @@ func (store *SQLCryptoStore) scanInboundGroupSession(rows dbutil.Scannable) (*In
 	var roomID id.RoomID
 	var signingKey, senderKey, forwardingChains sql.NullString
 	var sessionBytes, ratchetSafetyBytes []byte
-	var receivedAt sql.NullTime
+	var receivedAt, creationTS sql.NullTime
 	var maxAge, maxMessages sql.NullInt64
 	var isScheduled bool
 	var sharedHistory *bool
@@ -563,7 +572,7 @@ func (store *SQLCryptoStore) scanInboundGroupSession(rows dbutil.Scannable) (*In
 	var sourceUser id.UserID
 	err := rows.Scan(
 		&roomID, &senderKey, &signingKey, &sessionBytes, &forwardingChains, &sharedHistory, &ratchetSafetyBytes,
-		&receivedAt, &maxAge, &maxMessages, &isScheduled, &version, &keySource, &sourceUser,
+		&receivedAt, &maxAge, &maxMessages, &isScheduled, &version, &keySource, &sourceUser, &creationTS,
 	)
 	if err != nil {
 		return nil, err
@@ -580,6 +589,7 @@ func (store *SQLCryptoStore) scanInboundGroupSession(rows dbutil.Scannable) (*In
 		ForwardingChains: chains,
 		RatchetSafety:    rs,
 		ReceivedAt:       receivedAt.Time,
+		CreationTS:       creationTS.Time,
 		MaxAge:           maxAge.Int64,
 		MaxMessages:      int(maxMessages.Int64),
 		SharedHistory:    sharedHistory,
@@ -593,7 +603,7 @@ func (store *SQLCryptoStore) scanInboundGroupSession(rows dbutil.Scannable) (*In
 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, shared_history,
-		       ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version, key_source, source_user
+		       ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version, key_source, source_user, session_creation_ts
 		FROM crypto_megolm_inbound_session WHERE room_id=$1 AND account_id=$2 AND session IS NOT NULL`,
 		roomID, store.AccountID,
 	)
@@ -603,7 +613,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, shared_history,
-		       ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version, key_source, source_user
+		       ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version, key_source, source_user, session_creation_ts
 		FROM crypto_megolm_inbound_session WHERE account_id=$1 AND session IS NOT NULL`,
 		store.AccountID,
 	)
@@ -613,7 +623,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, shared_history,
-		       ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version, key_source, source_user
+		       ratchet_safety, received_at, max_age, max_messages, is_scheduled, key_backup_version, key_source, source_user, session_creation_ts
 		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 -> v21 (compatible with v20+): Latest revision
+-- v0 -> v22 (compatible with v20+): Latest revision
 CREATE TABLE crypto_account (
 	account_id         TEXT    PRIMARY KEY,
 	device_id          TEXT    NOT NULL,
@@ -73,6 +73,7 @@ CREATE TABLE crypto_megolm_inbound_session (
 	key_backup_version TEXT NOT NULL DEFAULT '',
 	key_source         TEXT NOT NULL DEFAULT '',
 	source_user        TEXT NOT NULL DEFAULT '',
+	session_creation_ts timestamp,
 	PRIMARY KEY (account_id, session_id)
 );
 -- Useful index to find keys that need backing up

crypto/sql_store_upgrade/22-megolm-session-creation-ts.sql

@@ -0,0 +1,2 @@
+-- v22 (compatible with v20+): Store the sender-provided creation timestamp for megolm sessions
+ALTER TABLE crypto_megolm_inbound_session ADD COLUMN session_creation_ts timestamp;

crypto/store.go

@@ -72,8 +72,9 @@ type Store interface {
 	GetGroupSession(context.Context, id.RoomID, id.SessionID) (*InboundGroupSession, error)
 	// RedactGroupSession removes the session data for the given inbound Megolm session from the store.
 	RedactGroupSession(context.Context, id.RoomID, id.SessionID, string) error
-	// RedactGroupSessions removes the session data for all inbound Megolm sessions from a specific device and/or in a specific room.
-	RedactGroupSessions(context.Context, id.RoomID, id.SenderKey, string) ([]id.SessionID, error)
+	// RedactGroupSessions removes session data for inbound Megolm sessions from a specific device and/or in a
+	// specific room. If before is set only sessions created before will be redacted.
+	RedactGroupSessions(ctx context.Context, roomID id.RoomID, senderKey id.SenderKey, reason string, before time.Time) ([]id.SessionID, error)
 	// RedactExpiredGroupSessions removes the session data for all inbound Megolm sessions that have expired.
 	RedactExpiredGroupSessions(context.Context) ([]id.SessionID, error)
 	// RedactOutdatedGroupSessions removes the session data for all inbound Megolm sessions that are lacking the expiration metadata.
@@ -350,30 +351,37 @@ func (gs *MemoryStore) RedactGroupSession(_ context.Context, roomID id.RoomID, s
 	return gs.save()
 }
 
-func (gs *MemoryStore) RedactGroupSessions(_ context.Context, roomID id.RoomID, senderKey id.SenderKey, reason string) ([]id.SessionID, error) {
+func (gs *MemoryStore) RedactGroupSessions(_ context.Context, roomID id.RoomID, senderKey id.SenderKey, reason string, before time.Time) ([]id.SessionID, error) {
 	gs.lock.Lock()
 	defer gs.lock.Unlock()
 	var sessionIDs []id.SessionID
+	olderThanCutoff := func(session *InboundGroupSession) bool {
+		return before.IsZero() || (!session.CreationTS.IsZero() && session.CreationTS.Before(before))
+	}
 	if roomID != "" && senderKey != "" {
 		sessions := gs.getGroupSessions(roomID)
 		for sessionID, session := range sessions {
-			if session.SenderKey == senderKey {
+			if session.SenderKey == senderKey && olderThanCutoff(session) {
 				sessionIDs = append(sessionIDs, sessionID)
 				delete(sessions, sessionID)
 			}
 		}
 	} else if senderKey != "" {
 		for _, room := range gs.GroupSessions {
 			for sessionID, session := range room {
-				if session.SenderKey == senderKey {
+				if session.SenderKey == senderKey && olderThanCutoff(session) {
 					sessionIDs = append(sessionIDs, sessionID)
 					delete(room, sessionID)
 				}
 			}
 		}
 	} else if roomID != "" {
-		sessionIDs = maps.Keys(gs.GroupSessions[roomID])
-		delete(gs.GroupSessions, roomID)
+		for sessionID, session := range gs.GroupSessions[roomID] {
+			if olderThanCutoff(session) {
+				sessionIDs = append(sessionIDs, sessionID)
+				delete(gs.GroupSessions[roomID], sessionID)
+			}
+		}
 	} else {
 		return nil, fmt.Errorf("room ID or sender key must be provided for redacting sessions")
 	}

crypto/store_test.go

@@ -11,6 +11,7 @@ import (
 	"database/sql"
 	"strconv"
 	"testing"
+	"time"
 
 	_ "github.com/mattn/go-sqlite3"
 	"github.com/stretchr/testify/assert"
@@ -163,6 +164,53 @@ func TestStoreMegolmSession(t *testing.T) {
 	}
 }
 
+func TestRedactGroupSessionsGatedByCreationTS(t *testing.T) {
+	stores := getCryptoStores(t)
+	for storeName, store := range stores {
+		t.Run(storeName, func(t *testing.T) {
+			acc := NewOlmAccount()
+			senderKey := acc.IdentityKey()
+
+			mkSession := func(creationTS time.Time) *InboundGroupSession {
+				ogs, err := olm.NewOutboundGroupSession()
+				require.NoError(t, err)
+				internal, err := olm.NewInboundGroupSession([]byte(ogs.Key()))
+				require.NoError(t, err)
+				igs := &InboundGroupSession{
+					Internal:   internal,
+					SigningKey: acc.SigningKey(),
+					SenderKey:  senderKey,
+					RoomID:     "room1",
+					ReceivedAt: time.Now().UTC(),
+					CreationTS: creationTS,
+				}
+				require.NoError(t, store.PutGroupSession(context.TODO(), igs))
+				return igs
+			}
+
+			oldTS := time.UnixMilli(1000).UTC()
+			newTS := time.UnixMilli(2000).UTC()
+			oldSession := mkSession(oldTS)
+			newSession := mkSession(newTS)
+			untimestampedSession := mkSession(time.Time{})
+
+			// A share created at newTS arriving late must only redact sessions the sender created before it.
+			redacted, err := store.RedactGroupSessions(context.TODO(), "room1", senderKey, "received new key from device", newTS)
+			require.NoError(t, err)
+			assert.ElementsMatch(t, []id.SessionID{oldSession.ID()}, redacted)
+
+			survivor, err := store.GetGroupSession(context.TODO(), "room1", newSession.ID())
+			require.NoError(t, err)
+			require.NotNil(t, survivor, "newer session must survive a delayed older share")
+
+			// A zero cutoff (e.g. device removed) redacts everything that remains.
+			redacted, err = store.RedactGroupSessions(context.TODO(), "room1", senderKey, "device removed", time.Time{})
+			require.NoError(t, err)
+			assert.ElementsMatch(t, []id.SessionID{newSession.ID(), untimestampedSession.ID()}, redacted)
+		})
+	}
+}
+
 func TestStoreOutboundMegolmSession(t *testing.T) {
 	stores := getCryptoStores(t)
 	for storeName, store := range stores {

event/encryption.go

@@ -11,6 +11,7 @@ import (
 	"fmt"
 
 	"go.mau.fi/util/jsonbytes"
+	"go.mau.fi/util/jsontime"
 
 	"maunium.net/go/mautrix/id"
 )
@@ -108,9 +109,10 @@ type RoomKeyEventContent struct {
 	SessionKey    string       `json:"session_key"`
 	SharedHistory *bool        `json:"shared_history,omitempty"`
 
-	MaxAge      int64 `json:"com.beeper.max_age_ms,omitempty"`
-	MaxMessages int   `json:"com.beeper.max_messages,omitempty"`
-	IsScheduled bool  `json:"com.beeper.is_scheduled,omitempty"`
+	MaxAge      int64              `json:"com.beeper.max_age_ms,omitempty"`
+	MaxMessages int                `json:"com.beeper.max_messages,omitempty"`
+	IsScheduled bool               `json:"com.beeper.is_scheduled,omitempty"`
+	CreationTS  jsontime.UnixMilli `json:"com.beeper.session_creation_ts,omitzero"`
 }
 
 // ForwardedRoomKeyEventContent represents the content of a m.forwarded_room_key to_device event.