diff --git a/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mip00KeyPackages/KeyPackageRotationManager.kt b/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mip00KeyPackages/KeyPackageRotationManager.kt new file mode 100644 index 000000000..196c8eeb3 --- /dev/null +++ b/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mip00KeyPackages/KeyPackageRotationManager.kt @@ -0,0 +1,218 @@ +/* + * Copyright (c) 2025 Vitor Pamplona + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * this software and associated documentation files (the "Software"), to deal in + * the Software without restriction, including without limitation the rights to use, + * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the + * Software, and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS + * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR + * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN + * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION + * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ +package com.vitorpamplona.quartz.marmot.mip00KeyPackages + +import com.vitorpamplona.quartz.marmot.mls.crypto.Ed25519 +import com.vitorpamplona.quartz.marmot.mls.crypto.MlsCryptoProvider +import com.vitorpamplona.quartz.marmot.mls.crypto.X25519 +import com.vitorpamplona.quartz.marmot.mls.messages.KeyPackageBundle +import com.vitorpamplona.quartz.marmot.mls.messages.MlsKeyPackage +import com.vitorpamplona.quartz.marmot.mls.tree.Capabilities +import com.vitorpamplona.quartz.marmot.mls.tree.Credential +import com.vitorpamplona.quartz.marmot.mls.tree.LeafNode +import com.vitorpamplona.quartz.marmot.mls.tree.LeafNodeSource +import com.vitorpamplona.quartz.marmot.mls.tree.Lifetime + +/** + * Manages KeyPackage creation and rotation lifecycle (MIP-00). + * + * Key responsibilities: + * - Generate fresh KeyPackages for publishing + * - Track consumed KeyPackages that need rotation after Welcome processing + * - Provide bundles for group joins + * - Handle periodic rotation for long-lived KeyPackages + * + * After a KeyPackage is consumed by a Welcome message: + * 1. The init_key is effectively spent — cannot be reused + * 2. A new KeyPackage MUST be published to the same d-tag slot + * 3. The old KeyPackageBundle MUST be discarded + * + * Per MIP-00 spec, each user should maintain up to [KeyPackageUtils.MAX_SLOTS] + * KeyPackage slots, rotating consumed ones promptly. + */ +class KeyPackageRotationManager { + private val activeBundles = mutableMapOf() + private val pendingRotations = mutableSetOf() + + /** + * Generate a new KeyPackage and its associated private bundle. + * + * @param identity the user's identity bytes (typically 32-byte Nostr pubkey) + * @param dTagSlot the d-tag slot for addressable replacement + * @return a [KeyPackageBundle] containing the KeyPackage and all private keys + */ + fun generateKeyPackage( + identity: ByteArray, + dTagSlot: String = KeyPackageUtils.PRIMARY_SLOT, + ): KeyPackageBundle { + val initKp = X25519.generateKeyPair() + val encKp = X25519.generateKeyPair() + val sigKp = Ed25519.generateKeyPair() + + val leafNode = + buildKeyPackageLeafNode( + encryptionKey = encKp.publicKey, + signatureKey = sigKp.publicKey, + identity = identity, + signingKey = sigKp.privateKey, + ) + + val unsigned = + MlsKeyPackage( + initKey = initKp.publicKey, + leafNode = leafNode, + signature = ByteArray(0), + ) + val keyPackage = + unsigned.copy( + signature = + MlsCryptoProvider.signWithLabel( + sigKp.privateKey, + "KeyPackageTBS", + unsigned.encodeTbs(), + ), + ) + + val bundle = KeyPackageBundle(keyPackage, initKp.privateKey, encKp.privateKey, sigKp.privateKey) + activeBundles[dTagSlot] = bundle + return bundle + } + + /** + * Get the active bundle for a d-tag slot. + * Used when processing a Welcome that references one of our KeyPackages. + */ + fun getBundle(dTagSlot: String): KeyPackageBundle? = activeBundles[dTagSlot] + + /** + * Find the bundle whose KeyPackage reference matches the given ref. + * Used when we receive a Welcome and need to find the matching bundle. + */ + fun findBundleByRef(keyPackageRef: ByteArray): KeyPackageBundle? = + activeBundles.values.find { bundle -> + bundle.keyPackage.reference().contentEquals(keyPackageRef) + } + + /** + * Mark a KeyPackage slot as consumed (used in a Welcome). + * The slot will be included in [pendingRotationSlots] and should be + * rotated by the caller. + */ + fun markConsumed(dTagSlot: String) { + activeBundles.remove(dTagSlot) + pendingRotations.add(dTagSlot) + } + + /** + * Mark a slot as consumed by looking up the KeyPackage reference. + */ + fun markConsumedByRef(keyPackageRef: ByteArray) { + val entry = + activeBundles.entries.find { (_, bundle) -> + bundle.keyPackage.reference().contentEquals(keyPackageRef) + } + if (entry != null) { + activeBundles.remove(entry.key) + pendingRotations.add(entry.key) + } + } + + /** + * Get the d-tag slots that need rotation (KeyPackage was consumed). + */ + fun pendingRotationSlots(): Set = pendingRotations.toSet() + + /** + * Clear a slot from the pending rotation set after a new KeyPackage + * has been published. + */ + fun clearPendingRotation(dTagSlot: String) { + pendingRotations.remove(dTagSlot) + } + + /** + * Check if any slots need rotation. + */ + fun needsRotation(): Boolean = pendingRotations.isNotEmpty() + + /** + * Rotate a consumed slot: generate a new KeyPackage for the same d-tag. + * + * @param identity the user's identity bytes + * @param dTagSlot the slot to rotate + * @return the new [KeyPackageBundle] ready for publishing + */ + fun rotateSlot( + identity: ByteArray, + dTagSlot: String, + ): KeyPackageBundle { + val bundle = generateKeyPackage(identity, dTagSlot) + pendingRotations.remove(dTagSlot) + return bundle + } + + /** + * Check if a KeyPackage should be proactively rotated based on age. + * + * MIP-00 recommends rotating KeyPackages periodically even if they + * haven't been consumed, to limit the exposure window. + * + * @param createdAtSeconds the KeyPackage event's created_at timestamp + * @param nowSeconds current time in seconds + * @return true if the KeyPackage should be rotated + */ + fun shouldRotateByAge( + createdAtSeconds: Long, + nowSeconds: Long, + ): Boolean = (nowSeconds - createdAtSeconds) > MAX_KEY_PACKAGE_AGE_SECONDS + + private fun buildKeyPackageLeafNode( + encryptionKey: ByteArray, + signatureKey: ByteArray, + identity: ByteArray, + signingKey: ByteArray, + ): LeafNode { + val now = System.currentTimeMillis() / 1000 + val unsigned = + LeafNode( + encryptionKey = encryptionKey, + signatureKey = signatureKey, + credential = Credential.Basic(identity), + capabilities = Capabilities(), + leafNodeSource = LeafNodeSource.KEY_PACKAGE, + lifetime = Lifetime(notBefore = now, notAfter = now + KEY_PACKAGE_LIFETIME_SECONDS), + extensions = emptyList(), + signature = ByteArray(0), + ) + + val tbs = unsigned.encodeTbs(groupId = null, leafIndex = null) + val signature = MlsCryptoProvider.signWithLabel(signingKey, "LeafNodeTBS", tbs) + return unsigned.copy(signature = signature) + } + + companion object { + /** KeyPackage lifetime: 30 days */ + const val KEY_PACKAGE_LIFETIME_SECONDS = 30L * 24 * 60 * 60 + + /** Proactive rotation after 7 days even if not consumed */ + const val MAX_KEY_PACKAGE_AGE_SECONDS = 7L * 24 * 60 * 60 + } +} diff --git a/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/group/MlsGroup.kt b/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/group/MlsGroup.kt index ae0494055..d69864f12 100644 --- a/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/group/MlsGroup.kt +++ b/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/group/MlsGroup.kt @@ -108,6 +108,47 @@ class MlsGroup private constructor( val epoch: Long get() = groupContext.epoch val leafIndex: Int get() = myLeafIndex + // --- State Persistence --- + + /** + * Capture the current group state as a serializable snapshot. + * + * The returned [MlsGroupState] contains all secret key material and + * MUST be stored in encrypted local storage. Call this after every + * epoch transition (commit, processCommit, processWelcome) to ensure + * the group can be restored after an app restart. + */ + fun saveState(): MlsGroupState { + val treeWriter = TlsWriter() + tree.encodeTls(treeWriter) + + return MlsGroupState( + groupContext = groupContext, + treeBytes = treeWriter.toByteArray(), + myLeafIndex = myLeafIndex, + epochSecrets = epochSecrets, + initSecret = initSecret, + signingPrivateKey = signingPrivateKey, + encryptionPrivateKey = encryptionPrivateKey, + interimTranscriptHash = interimTranscriptHash, + encryptionSecret = epochSecrets.encryptionSecret, + ) + } + + /** + * Extract retained epoch secrets for late-message decryption. + * + * Call this BEFORE an epoch transition to capture the outgoing epoch's + * decryption secrets. These are kept in a bounded window by [MlsGroupManager]. + */ + fun retainedSecrets(): RetainedEpochSecrets = + RetainedEpochSecrets( + epoch = epoch, + senderDataSecret = epochSecrets.senderDataSecret, + encryptionSecret = epochSecrets.encryptionSecret, + leafCount = tree.leafCount, + ) + val memberCount: Int get() { var count = 0 @@ -216,6 +257,48 @@ class MlsGroup private constructor( return proposal } + /** + * Create an Update proposal to rotate the signing key within the group. + * + * Per MIP-00, members SHOULD perform a self-update within 24 hours + * of joining a group. This replaces the leaf node with fresh key material, + * improving forward secrecy. + * + * The new signing key pair is generated internally and takes effect + * after the next Commit. + * + * @return the Update proposal (already queued for the next commit) + */ + fun proposeSigningKeyRotation(): Proposal.Update { + val newSigKp = Ed25519.generateKeyPair() + val newEncKp = X25519.generateKeyPair() + + val currentLeaf = tree.getLeaf(myLeafIndex) + val identity = + (currentLeaf?.credential as? Credential.Basic)?.identity + ?: ByteArray(0) + + val newLeafNode = + buildLeafNode( + encryptionKey = newEncKp.publicKey, + signatureKey = newSigKp.publicKey, + identity = identity, + source = LeafNodeSource.UPDATE, + signingKey = newSigKp.privateKey, + groupId = groupId, + leafIndex = myLeafIndex, + ) + + val proposal = Proposal.Update(newLeafNode) + pendingProposals.add(PendingProposal(proposal, myLeafIndex)) + + // Stage the new keys — they take effect when commit() is called + signingPrivateKey = newSigKp.privateKey + encryptionPrivateKey = newEncKp.privateKey + + return proposal + } + /** * Create a PSK proposal to include a pre-shared key in the next epoch. * The PSK must be registered via registerPsk() before committing. @@ -1408,6 +1491,32 @@ class MlsGroup private constructor( return Pair(group, commitBytes) } + /** + * Restore a group from a previously saved [MlsGroupState]. + * + * The SecretTree is reconstructed from the stored encryption_secret. + * Note: SecretTree ratchet state (per-sender generation counters) is + * NOT preserved — messages sent/received before the save point cannot + * be re-decrypted, which is acceptable because they would already + * have been processed. + */ + fun restore(state: MlsGroupState): MlsGroup { + val tree = RatchetTree.decodeTls(TlsReader(state.treeBytes)) + val secretTree = SecretTree(state.encryptionSecret, tree.leafCount) + + return MlsGroup( + groupContext = state.groupContext, + tree = tree, + myLeafIndex = state.myLeafIndex, + epochSecrets = state.epochSecrets, + secretTree = secretTree, + initSecret = state.initSecret, + signingPrivateKey = state.signingPrivateKey, + encryptionPrivateKey = state.encryptionPrivateKey, + interimTranscriptHash = state.interimTranscriptHash, + ) + } + /** * Build a LeafNode with signature. */ diff --git a/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/group/MlsGroupManager.kt b/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/group/MlsGroupManager.kt new file mode 100644 index 000000000..8ad481ebf --- /dev/null +++ b/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/group/MlsGroupManager.kt @@ -0,0 +1,442 @@ +/* + * Copyright (c) 2025 Vitor Pamplona + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * this software and associated documentation files (the "Software"), to deal in + * the Software without restriction, including without limitation the rights to use, + * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the + * Software, and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS + * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR + * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN + * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION + * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ +package com.vitorpamplona.quartz.marmot.mls.group + +import com.vitorpamplona.quartz.marmot.mls.codec.TlsReader +import com.vitorpamplona.quartz.marmot.mls.codec.TlsWriter +import com.vitorpamplona.quartz.marmot.mls.crypto.MlsCryptoProvider +import com.vitorpamplona.quartz.marmot.mls.messages.CommitResult +import com.vitorpamplona.quartz.marmot.mls.messages.KeyPackageBundle +import com.vitorpamplona.quartz.marmot.mls.schedule.SecretTree +import com.vitorpamplona.quartz.nip01Core.core.HexKey + +/** + * High-level coordinator for MLS group lifecycle and state management. + * + * Manages: + * - In-memory group instances (keyed by Nostr group ID) + * - Epoch secret retention window (N-1 epochs for late messages) + * - Secure deletion of consumed init_keys after Welcome processing + * - Persistence of group state through [MlsGroupStateStore] + * - KeyPackage rotation scheduling after group joins + * + * This class is the bridge between the low-level MLS engine ([MlsGroup]) + * and the application layer. It ensures that MLS state survives app restarts + * and that cryptographic hygiene (key deletion, rotation) is maintained. + * + * Thread safety: All public methods are suspending and should be called + * from a single coroutine context (e.g., the Account's scope). + */ +class MlsGroupManager( + private val store: MlsGroupStateStore, +) { + private val groups = mutableMapOf() + private val retainedEpochs = mutableMapOf>() + + /** + * Restore all groups from persistent storage on startup. + * Call this once during Account initialization. + */ + suspend fun restoreAll() { + val groupIds = store.listGroups() + for (nostrGroupId in groupIds) { + try { + val stateBytes = store.load(nostrGroupId) ?: continue + val state = MlsGroupState.decodeTls(stateBytes) + groups[nostrGroupId] = MlsGroup.restore(state) + + // Restore retained epochs + val retained = store.loadRetainedEpochs(nostrGroupId) + if (retained.isNotEmpty()) { + retainedEpochs[nostrGroupId] = + retained + .map { RetainedEpochSecrets.decodeTls(TlsReader(it)) } + .toMutableList() + } + } catch (_: Exception) { + // Corrupted state — remove it so it doesn't block future joins + store.delete(nostrGroupId) + } + } + } + + /** + * Get an active group by its Nostr group ID. + */ + fun getGroup(nostrGroupId: HexKey): MlsGroup? = groups[nostrGroupId] + + /** + * List all active Nostr group IDs. + */ + fun activeGroupIds(): Set = groups.keys.toSet() + + /** + * Check if we are a member of the given group. + */ + fun isMember(nostrGroupId: HexKey): Boolean = groups.containsKey(nostrGroupId) + + // --- Group Creation --- + + /** + * Create a new MLS group and persist it. + * + * @param nostrGroupId hex-encoded Nostr routing ID for the group + * @param identity the creator's identity bytes (typically Nostr pubkey) + * @param signingKey optional Ed25519 signing key (generated if null) + * @return the new [MlsGroup] + */ + suspend fun createGroup( + nostrGroupId: HexKey, + identity: ByteArray, + signingKey: ByteArray? = null, + ): MlsGroup { + val group = MlsGroup.create(identity, signingKey) + groups[nostrGroupId] = group + persistGroup(nostrGroupId) + return group + } + + // --- Joining --- + + /** + * Join a group by processing a Welcome message. + * + * After processing: + * 1. The init_key from the KeyPackageBundle is securely deleted + * (the bundle should be discarded by the caller after this returns) + * 2. Group state is persisted + * 3. A KeyPackage rotation should be triggered (see [needsKeyPackageRotation]) + * + * @param nostrGroupId hex-encoded Nostr group ID (from Welcome event tags) + * @param welcomeBytes TLS-serialized Welcome message + * @param bundle the KeyPackageBundle that was used for the invitation + * @return the joined [MlsGroup] + */ + suspend fun processWelcome( + nostrGroupId: HexKey, + welcomeBytes: ByteArray, + bundle: KeyPackageBundle, + ): MlsGroup { + val group = MlsGroup.processWelcome(welcomeBytes, bundle) + groups[nostrGroupId] = group + + // init_key is consumed — the bundle's initPrivateKey should not be + // reused. Caller must discard the bundle and rotate KeyPackages. + + persistGroup(nostrGroupId) + return group + } + + /** + * Join a group via external commit. + * + * @param nostrGroupId hex-encoded Nostr group ID + * @param groupInfoBytes TLS-serialized GroupInfo + * @param identity the joiner's identity bytes + * @param signingKey optional signing key + * @return pair of (joined group, commit bytes to publish) + */ + suspend fun externalJoin( + nostrGroupId: HexKey, + groupInfoBytes: ByteArray, + identity: ByteArray, + signingKey: ByteArray? = null, + ): Pair { + val (group, commitBytes) = MlsGroup.externalJoin(groupInfoBytes, identity, signingKey) + groups[nostrGroupId] = group + persistGroup(nostrGroupId) + return Pair(group, commitBytes) + } + + // --- Epoch Transitions --- + + /** + * Create and apply a Commit, advancing the epoch. + * + * Retains the outgoing epoch's decryption secrets and persists + * the new state before returning. + * + * @param nostrGroupId hex-encoded Nostr group ID + * @return the [CommitResult] to publish + */ + suspend fun commit(nostrGroupId: HexKey): CommitResult { + val group = requireGroup(nostrGroupId) + + // Retain current epoch secrets before transition + retainEpochSecrets(nostrGroupId, group) + + val result = group.commit() + persistGroup(nostrGroupId) + return result + } + + /** + * Process a received Commit, advancing the epoch. + * + * @param nostrGroupId hex-encoded Nostr group ID + * @param commitBytes TLS-serialized Commit + * @param senderLeafIndex sender's leaf index + * @param confirmationTag optional confirmation tag for verification + */ + suspend fun processCommit( + nostrGroupId: HexKey, + commitBytes: ByteArray, + senderLeafIndex: Int, + confirmationTag: ByteArray? = null, + ) { + val group = requireGroup(nostrGroupId) + + // Retain current epoch secrets before transition + retainEpochSecrets(nostrGroupId, group) + + group.processCommit(commitBytes, senderLeafIndex, confirmationTag) + persistGroup(nostrGroupId) + } + + // --- Message Encryption/Decryption --- + + /** + * Encrypt an application message. + */ + fun encrypt( + nostrGroupId: HexKey, + plaintext: ByteArray, + ): ByteArray = requireGroup(nostrGroupId).encrypt(plaintext) + + /** + * Decrypt an application message. + * + * Tries the current epoch first, then falls back to retained epoch + * secrets for late-arriving messages from previous epochs. + */ + fun decrypt( + nostrGroupId: HexKey, + messageBytes: ByteArray, + ): DecryptedMessage { + val group = requireGroup(nostrGroupId) + + // Try current epoch + val current = group.decryptOrNull(messageBytes) + if (current != null) return current + + // Try retained epochs + val retained = retainedEpochs[nostrGroupId] ?: emptyList() + for (epochSecrets in retained) { + val result = tryDecryptWithRetainedEpoch(messageBytes, epochSecrets) + if (result != null) return result + } + + // No epoch could decrypt — rethrow from current epoch for diagnostics + return group.decrypt(messageBytes) + } + + /** + * Decrypt with null return on failure. + */ + fun decryptOrNull( + nostrGroupId: HexKey, + messageBytes: ByteArray, + ): DecryptedMessage? = + try { + decrypt(nostrGroupId, messageBytes) + } catch (_: Exception) { + null + } + + // --- Member Management --- + + /** + * Add a member and create a Commit. + */ + suspend fun addMember( + nostrGroupId: HexKey, + keyPackageBytes: ByteArray, + ): CommitResult { + val group = requireGroup(nostrGroupId) + retainEpochSecrets(nostrGroupId, group) + val result = group.addMember(keyPackageBytes) + persistGroup(nostrGroupId) + return result + } + + /** + * Remove a member and create a Commit. + */ + suspend fun removeMember( + nostrGroupId: HexKey, + targetLeafIndex: Int, + ): CommitResult { + val group = requireGroup(nostrGroupId) + retainEpochSecrets(nostrGroupId, group) + val result = group.removeMember(targetLeafIndex) + persistGroup(nostrGroupId) + return result + } + + /** + * Rotate the signing key within a group and commit. + * + * Per MIP-00, members SHOULD self-update within 24 hours of joining. + * This creates an Update proposal with a fresh signing key and commits it. + * + * @param nostrGroupId hex-encoded Nostr group ID + * @return the [CommitResult] to publish + */ + suspend fun rotateSigningKey(nostrGroupId: HexKey): CommitResult { + val group = requireGroup(nostrGroupId) + retainEpochSecrets(nostrGroupId, group) + group.proposeSigningKeyRotation() + val result = group.commit() + persistGroup(nostrGroupId) + return result + } + + /** + * Leave a group (self-remove). + * Returns the SelfRemove proposal bytes to publish, then removes + * local state. + */ + suspend fun leaveGroup(nostrGroupId: HexKey): ByteArray { + val group = requireGroup(nostrGroupId) + val proposalBytes = group.selfRemove() + groups.remove(nostrGroupId) + retainedEpochs.remove(nostrGroupId) + store.delete(nostrGroupId) + return proposalBytes + } + + // --- Key Export --- + + /** + * Export the Marmot outer encryption key for GroupEvent wrapping. + * MLS-Exporter("marmot", "group-event", 32) + */ + fun exporterSecret(nostrGroupId: HexKey): ByteArray = + requireGroup(nostrGroupId).exporterSecret( + "marmot", + "group-event".encodeToByteArray(), + 32, + ) + + // --- Private Helpers --- + + private fun requireGroup(nostrGroupId: HexKey): MlsGroup = + groups[nostrGroupId] + ?: throw IllegalStateException("Not a member of group $nostrGroupId") + + private suspend fun persistGroup(nostrGroupId: HexKey) { + val group = groups[nostrGroupId] ?: return + val state = group.saveState() + store.save(nostrGroupId, state.encodeTls()) + + // Also persist retained epochs + val retained = retainedEpochs[nostrGroupId] + if (retained != null) { + val retainedBytes = + retained.map { epoch -> + val writer = TlsWriter() + epoch.encodeTls(writer) + writer.toByteArray() + } + store.saveRetainedEpochs(nostrGroupId, retainedBytes) + } + } + + private fun retainEpochSecrets( + nostrGroupId: HexKey, + group: MlsGroup, + ) { + val retained = retainedEpochs.getOrPut(nostrGroupId) { mutableListOf() } + retained.add(group.retainedSecrets()) + + // Trim to retention window (keep only the most recent N-1 epochs) + while (retained.size > EPOCH_RETENTION_WINDOW) { + retained.removeAt(0) + } + } + + private fun tryDecryptWithRetainedEpoch( + messageBytes: ByteArray, + retained: RetainedEpochSecrets, + ): DecryptedMessage? = + try { + val secretTree = SecretTree(retained.encryptionSecret, retained.leafCount) + val mlsMsg = + com.vitorpamplona.quartz.marmot.mls.framing.MlsMessage + .decodeTls(TlsReader(messageBytes)) + + if (mlsMsg.wireFormat != com.vitorpamplona.quartz.marmot.mls.framing.WireFormat.PRIVATE_MESSAGE) { + return null + } + + val privMsg = + com.vitorpamplona.quartz.marmot.mls.framing.PrivateMessage + .decodeTls(TlsReader(mlsMsg.payload)) + if (privMsg.epoch != retained.epoch) return null + + // Decrypt sender data + val senderDataKey = + MlsCryptoProvider.expandWithLabel( + retained.senderDataSecret, + "key", + ByteArray(0), + MlsCryptoProvider.AEAD_KEY_LENGTH, + ) + val senderDataNonce = + MlsCryptoProvider.expandWithLabel( + retained.senderDataSecret, + "nonce", + ByteArray(0), + MlsCryptoProvider.AEAD_NONCE_LENGTH, + ) + val senderDataPlain = + MlsCryptoProvider.aeadDecrypt( + senderDataKey, + senderDataNonce, + ByteArray(0), + privMsg.encryptedSenderData, + ) + val senderReader = TlsReader(senderDataPlain) + val senderLeafIndex = senderReader.readUint32().toInt() + val generation = senderReader.readUint32().toInt() + + val kng = secretTree.applicationKeyNonceForGeneration(senderLeafIndex, generation) + val plaintext = + MlsCryptoProvider.aeadDecrypt(kng.key, kng.nonce, ByteArray(0), privMsg.ciphertext) + + DecryptedMessage( + senderLeafIndex = senderLeafIndex, + contentType = privMsg.contentType, + content = plaintext, + epoch = privMsg.epoch, + ) + } catch (_: Exception) { + null + } + + companion object { + /** + * Number of past epochs to retain for late-arriving message decryption. + * MLS forward secrecy guarantees mean we want to limit this window. + */ + const val EPOCH_RETENTION_WINDOW = 2 + } +} diff --git a/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/group/MlsGroupState.kt b/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/group/MlsGroupState.kt new file mode 100644 index 000000000..f46fdb7f9 --- /dev/null +++ b/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/group/MlsGroupState.kt @@ -0,0 +1,199 @@ +/* + * Copyright (c) 2025 Vitor Pamplona + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * this software and associated documentation files (the "Software"), to deal in + * the Software without restriction, including without limitation the rights to use, + * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the + * Software, and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS + * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR + * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN + * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION + * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ +package com.vitorpamplona.quartz.marmot.mls.group + +import com.vitorpamplona.quartz.marmot.mls.codec.TlsReader +import com.vitorpamplona.quartz.marmot.mls.codec.TlsWriter +import com.vitorpamplona.quartz.marmot.mls.messages.GroupContext +import com.vitorpamplona.quartz.marmot.mls.schedule.EpochSecrets + +/** + * Serializable snapshot of an MLS group's complete state. + * + * This captures everything needed to restore an [MlsGroup] after an app restart: + * - Group context (group ID, epoch, tree hash, transcript hash, extensions) + * - Ratchet tree (all member leaf nodes and parent nodes) + * - Key schedule secrets (epoch secrets for decryption) + * - Local member state (leaf index, signing key, encryption key) + * - Transcript hash for next epoch + * + * The state is serialized using TLS encoding (compact binary, same codec + * used throughout the MLS stack) for efficient storage. + * + * Security: This blob contains secret key material (signing key, encryption key, + * epoch secrets). It MUST be stored in encrypted local storage. + */ +data class MlsGroupState( + val groupContext: GroupContext, + val treeBytes: ByteArray, + val myLeafIndex: Int, + val epochSecrets: EpochSecrets, + val initSecret: ByteArray, + val signingPrivateKey: ByteArray, + val encryptionPrivateKey: ByteArray, + val interimTranscriptHash: ByteArray, + val encryptionSecret: ByteArray, +) { + fun encodeTls(): ByteArray { + val writer = TlsWriter() + + // Version tag for future-proofing + writer.putUint16(STATE_VERSION) + + // GroupContext + groupContext.encodeTls(writer) + + // Ratchet tree (already TLS-encoded) + writer.putOpaqueVarInt(treeBytes) + + // Local member state + writer.putUint32(myLeafIndex.toLong()) + + // Epoch secrets (all 12 fields) + writer.putOpaqueVarInt(epochSecrets.joinerSecret) + writer.putOpaqueVarInt(epochSecrets.welcomeSecret) + writer.putOpaqueVarInt(epochSecrets.epochSecret) + writer.putOpaqueVarInt(epochSecrets.senderDataSecret) + writer.putOpaqueVarInt(epochSecrets.encryptionSecret) + writer.putOpaqueVarInt(epochSecrets.exporterSecret) + writer.putOpaqueVarInt(epochSecrets.epochAuthenticator) + writer.putOpaqueVarInt(epochSecrets.externalSecret) + writer.putOpaqueVarInt(epochSecrets.confirmationKey) + writer.putOpaqueVarInt(epochSecrets.membershipKey) + writer.putOpaqueVarInt(epochSecrets.resumptionPsk) + writer.putOpaqueVarInt(epochSecrets.initSecret) + + // Init secret for next epoch + writer.putOpaqueVarInt(initSecret) + + // Private keys + writer.putOpaqueVarInt(signingPrivateKey) + writer.putOpaqueVarInt(encryptionPrivateKey) + + // Transcript state + writer.putOpaqueVarInt(interimTranscriptHash) + + // Encryption secret for SecretTree reconstruction + writer.putOpaqueVarInt(encryptionSecret) + + return writer.toByteArray() + } + + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (other !is MlsGroupState) return false + return groupContext == other.groupContext && myLeafIndex == other.myLeafIndex + } + + override fun hashCode(): Int { + var result = groupContext.hashCode() + result = 31 * result + myLeafIndex + return result + } + + companion object { + private const val STATE_VERSION = 1 + + fun decodeTls(data: ByteArray): MlsGroupState { + val reader = TlsReader(data) + + val version = reader.readUint16() + require(version == STATE_VERSION) { "Unsupported state version: $version" } + + val groupContext = GroupContext.decodeTls(reader) + val treeBytes = reader.readOpaqueVarInt() + val myLeafIndex = reader.readUint32().toInt() + + val epochSecrets = + EpochSecrets( + joinerSecret = reader.readOpaqueVarInt(), + welcomeSecret = reader.readOpaqueVarInt(), + epochSecret = reader.readOpaqueVarInt(), + senderDataSecret = reader.readOpaqueVarInt(), + encryptionSecret = reader.readOpaqueVarInt(), + exporterSecret = reader.readOpaqueVarInt(), + epochAuthenticator = reader.readOpaqueVarInt(), + externalSecret = reader.readOpaqueVarInt(), + confirmationKey = reader.readOpaqueVarInt(), + membershipKey = reader.readOpaqueVarInt(), + resumptionPsk = reader.readOpaqueVarInt(), + initSecret = reader.readOpaqueVarInt(), + ) + + val initSecret = reader.readOpaqueVarInt() + val signingPrivateKey = reader.readOpaqueVarInt() + val encryptionPrivateKey = reader.readOpaqueVarInt() + val interimTranscriptHash = reader.readOpaqueVarInt() + val encryptionSecret = reader.readOpaqueVarInt() + + return MlsGroupState( + groupContext = groupContext, + treeBytes = treeBytes, + myLeafIndex = myLeafIndex, + epochSecrets = epochSecrets, + initSecret = initSecret, + signingPrivateKey = signingPrivateKey, + encryptionPrivateKey = encryptionPrivateKey, + interimTranscriptHash = interimTranscriptHash, + encryptionSecret = encryptionSecret, + ) + } + } +} + +/** + * A retained epoch's decryption secrets for processing late-arriving messages. + * + * When an epoch transition occurs, the previous epoch's secrets are kept + * in a bounded window so messages that arrive out of order can still be + * decrypted. Only the minimum secrets needed for decryption are retained. + */ +data class RetainedEpochSecrets( + val epoch: Long, + val senderDataSecret: ByteArray, + val encryptionSecret: ByteArray, + val leafCount: Int, +) { + override fun equals(other: Any?): Boolean { + if (this === other) return true + if (other !is RetainedEpochSecrets) return false + return epoch == other.epoch + } + + override fun hashCode(): Int = epoch.hashCode() + + fun encodeTls(writer: TlsWriter) { + writer.putUint64(epoch) + writer.putOpaqueVarInt(senderDataSecret) + writer.putOpaqueVarInt(encryptionSecret) + writer.putUint32(leafCount.toLong()) + } + + companion object { + fun decodeTls(reader: TlsReader): RetainedEpochSecrets = + RetainedEpochSecrets( + epoch = reader.readUint64(), + senderDataSecret = reader.readOpaqueVarInt(), + encryptionSecret = reader.readOpaqueVarInt(), + leafCount = reader.readUint32().toInt(), + ) + } +} diff --git a/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/group/MlsGroupStateStore.kt b/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/group/MlsGroupStateStore.kt new file mode 100644 index 000000000..d674931a9 --- /dev/null +++ b/quartz/src/commonMain/kotlin/com/vitorpamplona/quartz/marmot/mls/group/MlsGroupStateStore.kt @@ -0,0 +1,92 @@ +/* + * Copyright (c) 2025 Vitor Pamplona + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * this software and associated documentation files (the "Software"), to deal in + * the Software without restriction, including without limitation the rights to use, + * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the + * Software, and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS + * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR + * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN + * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION + * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ +package com.vitorpamplona.quartz.marmot.mls.group + +/** + * Interface for encrypted local storage of MLS group state. + * + * Implementations MUST encrypt all data at rest — the stored blobs contain + * private keys and epoch secrets. Platform-specific implementations: + * - Android: EncryptedSharedPreferences or encrypted DataStore + * - Desktop: Encrypted file with AES-GCM keyed from system keychain + * + * Each group's state is stored independently, keyed by the hex-encoded + * Nostr group ID (the `h` tag value from MIP-01, NOT the internal MLS + * group ID). This allows the storage layer to list/prune groups without + * parsing MLS state. + * + * Retained epoch secrets are stored alongside the current state so + * late-arriving messages from the previous epoch can still be decrypted. + */ +interface MlsGroupStateStore { + /** + * Persist the current group state. + * Overwrites any previously saved state for this group. + * + * @param nostrGroupId hex-encoded Nostr group ID (`h` tag) + * @param state TLS-encoded group state bytes from [MlsGroupState.encodeTls] + */ + suspend fun save( + nostrGroupId: String, + state: ByteArray, + ) + + /** + * Load a previously saved group state. + * + * @param nostrGroupId hex-encoded Nostr group ID + * @return TLS-encoded state bytes, or null if no state exists + */ + suspend fun load(nostrGroupId: String): ByteArray? + + /** + * Delete all stored state for a group (after leaving). + * This should securely delete the encrypted blob. + * + * @param nostrGroupId hex-encoded Nostr group ID + */ + suspend fun delete(nostrGroupId: String) + + /** + * List all Nostr group IDs that have saved state. + * Used during startup to restore active group memberships. + */ + suspend fun listGroups(): List + + /** + * Save retained epoch secrets for late-message decryption. + * + * @param nostrGroupId hex-encoded Nostr group ID + * @param retainedSecrets list of TLS-encoded [RetainedEpochSecrets] + */ + suspend fun saveRetainedEpochs( + nostrGroupId: String, + retainedSecrets: List, + ) + + /** + * Load retained epoch secrets. + * + * @param nostrGroupId hex-encoded Nostr group ID + * @return list of TLS-encoded [RetainedEpochSecrets], empty if none + */ + suspend fun loadRetainedEpochs(nostrGroupId: String): List +} diff --git a/quartz/src/jvmAndroidTest/kotlin/com/vitorpamplona/quartz/marmot/mls/MlsGroupManagerTest.kt b/quartz/src/jvmAndroidTest/kotlin/com/vitorpamplona/quartz/marmot/mls/MlsGroupManagerTest.kt new file mode 100644 index 000000000..47c76842b --- /dev/null +++ b/quartz/src/jvmAndroidTest/kotlin/com/vitorpamplona/quartz/marmot/mls/MlsGroupManagerTest.kt @@ -0,0 +1,223 @@ +/* + * Copyright (c) 2025 Vitor Pamplona + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * this software and associated documentation files (the "Software"), to deal in + * the Software without restriction, including without limitation the rights to use, + * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the + * Software, and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS + * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR + * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN + * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION + * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ +package com.vitorpamplona.quartz.marmot.mls + +import com.vitorpamplona.quartz.marmot.mls.group.MlsGroupManager +import com.vitorpamplona.quartz.marmot.mls.group.MlsGroupStateStore +import kotlinx.coroutines.runBlocking +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertNotNull +import kotlin.test.assertNull +import kotlin.test.assertTrue + +/** + * In-memory implementation of [MlsGroupStateStore] for testing. + */ +class InMemoryGroupStateStore : MlsGroupStateStore { + private val states = mutableMapOf() + private val retainedEpochs = mutableMapOf>() + + override suspend fun save( + nostrGroupId: String, + state: ByteArray, + ) { + states[nostrGroupId] = state + } + + override suspend fun load(nostrGroupId: String): ByteArray? = states[nostrGroupId] + + override suspend fun delete(nostrGroupId: String) { + states.remove(nostrGroupId) + retainedEpochs.remove(nostrGroupId) + } + + override suspend fun listGroups(): List = states.keys.toList() + + override suspend fun saveRetainedEpochs( + nostrGroupId: String, + retainedSecrets: List, + ) { + retainedEpochs[nostrGroupId] = retainedSecrets + } + + override suspend fun loadRetainedEpochs(nostrGroupId: String): List = retainedEpochs[nostrGroupId] ?: emptyList() +} + +class MlsGroupManagerTest { + private val groupId = "a".repeat(64) // 32-byte hex Nostr group ID + + @Test + fun testCreateGroupAndPersist() { + runBlocking { + val store = InMemoryGroupStateStore() + val manager = MlsGroupManager(store) + + val group = manager.createGroup(groupId, "alice".encodeToByteArray()) + + assertNotNull(group) + assertEquals(0L, group.epoch) + assertTrue(manager.isMember(groupId)) + assertTrue(manager.activeGroupIds().contains(groupId)) + + // State should be persisted + assertNotNull(store.load(groupId)) + } + } + + @Test + fun testRestoreAfterRestart() { + runBlocking { + val store = InMemoryGroupStateStore() + + // Session 1: create group + val manager1 = MlsGroupManager(store) + manager1.createGroup(groupId, "alice".encodeToByteArray()) + + // Session 2: restore + val manager2 = MlsGroupManager(store) + manager2.restoreAll() + + assertTrue(manager2.isMember(groupId)) + val group = manager2.getGroup(groupId) + assertNotNull(group) + assertEquals(0L, group.epoch) + } + } + + @Test + fun testEncryptDecryptThroughManager() { + runBlocking { + val store = InMemoryGroupStateStore() + val manager = MlsGroupManager(store) + manager.createGroup(groupId, "alice".encodeToByteArray()) + + val plaintext = "Hello via manager".encodeToByteArray() + val encrypted = manager.encrypt(groupId, plaintext) + val decrypted = manager.decrypt(groupId, encrypted) + + assertContentEquals(plaintext, decrypted.content) + } + } + + @Test + fun testAddMemberPersistsState() { + runBlocking { + val store = InMemoryGroupStateStore() + val manager = MlsGroupManager(store) + val group = manager.createGroup(groupId, "alice".encodeToByteArray()) + + val bobBundle = group.createKeyPackage("bob".encodeToByteArray(), ByteArray(0)) + val result = manager.addMember(groupId, bobBundle.keyPackage.toTlsBytes()) + + assertNotNull(result.welcomeBytes) + + // Restore and verify epoch advanced + val manager2 = MlsGroupManager(store) + manager2.restoreAll() + val restored = manager2.getGroup(groupId) + assertNotNull(restored) + assertEquals(1L, restored.epoch) + assertEquals(2, restored.memberCount) + } + } + + @Test + fun testLeaveGroupDeletesState() { + runBlocking { + val store = InMemoryGroupStateStore() + val manager = MlsGroupManager(store) + manager.createGroup(groupId, "alice".encodeToByteArray()) + + assertTrue(manager.isMember(groupId)) + + manager.leaveGroup(groupId) + + assertTrue(!manager.isMember(groupId)) + assertNull(store.load(groupId)) + } + } + + @Test + fun testExporterSecret() { + runBlocking { + val store = InMemoryGroupStateStore() + val manager = MlsGroupManager(store) + manager.createGroup(groupId, "alice".encodeToByteArray()) + + val key = manager.exporterSecret(groupId) + assertEquals(32, key.size) + } + } + + @Test + fun testSigningKeyRotation() { + runBlocking { + val store = InMemoryGroupStateStore() + val manager = MlsGroupManager(store) + manager.createGroup(groupId, "alice".encodeToByteArray()) + + val keyBefore = manager.exporterSecret(groupId) + + val result = manager.rotateSigningKey(groupId) + assertTrue(result.commitBytes.isNotEmpty()) + + val keyAfter = manager.exporterSecret(groupId) + assertTrue(!keyBefore.contentEquals(keyAfter)) + + // Verify state persisted after rotation + val manager2 = MlsGroupManager(store) + manager2.restoreAll() + val restoredKey = manager2.exporterSecret(groupId) + assertContentEquals(keyAfter, restoredKey) + } + } + + @Test + fun testMultipleGroupsIndependent() { + runBlocking { + val store = InMemoryGroupStateStore() + val manager = MlsGroupManager(store) + + val groupId1 = "1".repeat(64) + val groupId2 = "2".repeat(64) + + manager.createGroup(groupId1, "alice".encodeToByteArray()) + manager.createGroup(groupId2, "alice".encodeToByteArray()) + + assertTrue(manager.isMember(groupId1)) + assertTrue(manager.isMember(groupId2)) + assertEquals(2, manager.activeGroupIds().size) + + // Leave one group + manager.leaveGroup(groupId1) + assertTrue(!manager.isMember(groupId1)) + assertTrue(manager.isMember(groupId2)) + + // Restore + val manager2 = MlsGroupManager(store) + manager2.restoreAll() + assertTrue(!manager2.isMember(groupId1)) + assertTrue(manager2.isMember(groupId2)) + } + } +} diff --git a/quartz/src/jvmAndroidTest/kotlin/com/vitorpamplona/quartz/marmot/mls/MlsGroupStateTest.kt b/quartz/src/jvmAndroidTest/kotlin/com/vitorpamplona/quartz/marmot/mls/MlsGroupStateTest.kt new file mode 100644 index 000000000..68cbfbe99 --- /dev/null +++ b/quartz/src/jvmAndroidTest/kotlin/com/vitorpamplona/quartz/marmot/mls/MlsGroupStateTest.kt @@ -0,0 +1,222 @@ +/* + * Copyright (c) 2025 Vitor Pamplona + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * this software and associated documentation files (the "Software"), to deal in + * the Software without restriction, including without limitation the rights to use, + * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the + * Software, and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS + * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR + * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN + * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION + * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ +package com.vitorpamplona.quartz.marmot.mls + +import com.vitorpamplona.quartz.marmot.mls.group.MlsGroup +import com.vitorpamplona.quartz.marmot.mls.group.MlsGroupState +import com.vitorpamplona.quartz.marmot.mls.group.RetainedEpochSecrets +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * Tests for MLS group state serialization and restoration. + * + * Verifies that group state survives a save/restore cycle and that + * the restored group can continue encrypting/decrypting messages. + */ +class MlsGroupStateTest { + @Test + fun testSaveAndRestoreState() { + val group = MlsGroup.create("alice".encodeToByteArray()) + val plaintext = "Hello from epoch 0".encodeToByteArray() + val encrypted = group.encrypt(plaintext) + + // Save state + val state = group.saveState() + val stateBytes = state.encodeTls() + assertTrue(stateBytes.isNotEmpty()) + + // Restore + val restoredState = MlsGroupState.decodeTls(stateBytes) + val restoredGroup = MlsGroup.restore(restoredState) + + assertEquals(group.epoch, restoredGroup.epoch) + assertEquals(group.leafIndex, restoredGroup.leafIndex) + assertEquals(group.memberCount, restoredGroup.memberCount) + assertContentEquals(group.groupId, restoredGroup.groupId) + } + + @Test + fun testRestoredGroupCanEncrypt() { + val group = MlsGroup.create("alice".encodeToByteArray()) + + // Save and restore + val state = group.saveState() + val restoredGroup = MlsGroup.restore(MlsGroupState.decodeTls(state.encodeTls())) + + // Encrypt with restored group + val plaintext = "Message after restore".encodeToByteArray() + val encrypted = restoredGroup.encrypt(plaintext) + assertTrue(encrypted.isNotEmpty()) + + // Decrypt with restored group + val decrypted = restoredGroup.decrypt(encrypted) + assertContentEquals(plaintext, decrypted.content) + } + + @Test + fun testSaveRestoreAfterAddMember() { + val group = MlsGroup.create("alice".encodeToByteArray()) + val bobBundle = group.createKeyPackage("bob".encodeToByteArray(), ByteArray(0)) + group.addMember(bobBundle.keyPackage.toTlsBytes()) + + assertEquals(1L, group.epoch) + assertEquals(2, group.memberCount) + + // Save and restore + val state = group.saveState() + val restoredGroup = MlsGroup.restore(MlsGroupState.decodeTls(state.encodeTls())) + + assertEquals(1L, restoredGroup.epoch) + assertEquals(2, restoredGroup.memberCount) + + // Can still encrypt/decrypt + val plaintext = "Post-add message".encodeToByteArray() + val encrypted = restoredGroup.encrypt(plaintext) + val decrypted = restoredGroup.decrypt(encrypted) + assertContentEquals(plaintext, decrypted.content) + } + + @Test + fun testSaveRestorePreservesExporterSecret() { + val group = MlsGroup.create("alice".encodeToByteArray()) + + val keyBefore = group.exporterSecret("marmot", "group-event".encodeToByteArray(), 32) + + val state = group.saveState() + val restoredGroup = MlsGroup.restore(MlsGroupState.decodeTls(state.encodeTls())) + + val keyAfter = restoredGroup.exporterSecret("marmot", "group-event".encodeToByteArray(), 32) + + assertContentEquals(keyBefore, keyAfter) + } + + @Test + fun testRetainedEpochSecretsSerialization() { + val group = MlsGroup.create("alice".encodeToByteArray()) + val retained = group.retainedSecrets() + + assertEquals(0L, retained.epoch) + assertTrue(retained.senderDataSecret.isNotEmpty()) + assertTrue(retained.encryptionSecret.isNotEmpty()) + assertEquals(1, retained.leafCount) + + // Serialize and deserialize + val writer = + com.vitorpamplona.quartz.marmot.mls.codec + .TlsWriter() + retained.encodeTls(writer) + val bytes = writer.toByteArray() + + val restored = + RetainedEpochSecrets.decodeTls( + com.vitorpamplona.quartz.marmot.mls.codec + .TlsReader(bytes), + ) + + assertEquals(retained.epoch, restored.epoch) + assertContentEquals(retained.senderDataSecret, restored.senderDataSecret) + assertContentEquals(retained.encryptionSecret, restored.encryptionSecret) + assertEquals(retained.leafCount, restored.leafCount) + } + + @Test + fun testMultipleEpochSaveRestore() { + val group = MlsGroup.create("alice".encodeToByteArray()) + + // Advance through multiple epochs + for (i in 0 until 3) { + val bundle = group.createKeyPackage("member$i".encodeToByteArray(), ByteArray(0)) + group.addMember(bundle.keyPackage.toTlsBytes()) + } + + assertEquals(3L, group.epoch) + assertEquals(4, group.memberCount) + + // Save and restore + val state = group.saveState() + val restoredGroup = MlsGroup.restore(MlsGroupState.decodeTls(state.encodeTls())) + + assertEquals(3L, restoredGroup.epoch) + assertEquals(4, restoredGroup.memberCount) + + // Verify encrypt/decrypt still works + val plaintext = "After 3 epochs".encodeToByteArray() + val encrypted = restoredGroup.encrypt(plaintext) + val decrypted = restoredGroup.decrypt(encrypted) + assertContentEquals(plaintext, decrypted.content) + } + + @Test + fun testStateVersionInSerialization() { + val group = MlsGroup.create("alice".encodeToByteArray()) + val state = group.saveState() + val bytes = state.encodeTls() + + // First two bytes should be the version (uint16 = 1) + val reader = + com.vitorpamplona.quartz.marmot.mls.codec + .TlsReader(bytes) + val version = reader.readUint16() + assertEquals(1, version) + } + + @Test + fun testSigningKeyRotation() { + val group = MlsGroup.create("alice".encodeToByteArray()) + val exporterBefore = group.exporterSecret("marmot", "test".encodeToByteArray(), 32) + + // Rotate signing key + group.proposeSigningKeyRotation() + group.commit() + + assertEquals(1L, group.epoch) + + // Exporter secret changes after epoch transition + val exporterAfter = group.exporterSecret("marmot", "test".encodeToByteArray(), 32) + assertTrue(!exporterBefore.contentEquals(exporterAfter)) + + // Can still encrypt/decrypt after rotation + val plaintext = "After key rotation".encodeToByteArray() + val encrypted = group.encrypt(plaintext) + val decrypted = group.decrypt(encrypted) + assertContentEquals(plaintext, decrypted.content) + } + + @Test + fun testSaveRestoreAfterSigningKeyRotation() { + val group = MlsGroup.create("alice".encodeToByteArray()) + group.proposeSigningKeyRotation() + group.commit() + + val state = group.saveState() + val restoredGroup = MlsGroup.restore(MlsGroupState.decodeTls(state.encodeTls())) + + assertEquals(1L, restoredGroup.epoch) + + val plaintext = "After restore + rotation".encodeToByteArray() + val encrypted = restoredGroup.encrypt(plaintext) + val decrypted = restoredGroup.decrypt(encrypted) + assertContentEquals(plaintext, decrypted.content) + } +}