Pro: turn the badge on when an account gets its first ever proof

app/src/main/java/org/thoughtcrime/securesms/pro/ProProofGenerationWorker.kt

@@ -15,7 +15,9 @@ import dagger.assisted.Assisted
 import dagger.assisted.AssistedInject
 import kotlinx.coroutines.CancellationException
 import network.loki.messenger.libsession_util.ED25519
+import network.loki.messenger.libsession_util.MutableUserProfile
 import network.loki.messenger.libsession_util.pro.ProConfig
+import network.loki.messenger.libsession_util.pro.ProProof
 import org.session.libsession.network.SnodeClock
 import org.session.libsession.utilities.ConfigFactoryProtocol
 import org.session.libsession.utilities.withMutableUserConfigs
@@ -140,15 +142,10 @@ class ProProofGenerationWorker @AssistedInject constructor(
                     val proof = requireNotNull(response.proof) { "generate-proof returned ok without a proof" }
 
                     configFactory.withMutableUserConfigs { configs ->
-                        // Upgrade guard: only replace the proof if it extends coverage (monotonic merge;
-                        // same-period races round to the same expiry -> byte-identical -> no-op). Avoids
-                        // churning a proof another device just landed.
-                        val current = configs.userProfile.getProConfig()?.proProof
-                        if (current == null || proof.expirySeconds > current.expirySeconds) {
-                            configs.userProfile.setProConfig(ProConfig(
-                                proProof = proof,
-                                rotatingPrivateKey = rotatingPrivateKey))
-                        }
+                        // Before the access-expiry write below: that is one of the three values
+                        // `storeProof` reads to decide whether this account has ever held Pro.
+                        storeProof(configs.userProfile, proof, rotatingPrivateKey)
+
                         // Refresh the cached access-expiry from the advisory account_expiry that rides the
                         // proof response, so the renewal path keeps E fresh without a separate get_pro_status.
                         response.accountExpiry?.let { configs.userProfile.setProAccessExpiry(it.epochSecond) }
@@ -307,6 +304,43 @@ class ProProofGenerationWorker @AssistedInject constructor(
     companion object {
         private const val WORK_NAME = "ProProofGenerationWorker"
 
+        /**
+         * Store a freshly minted [proof] in config, and on an account's first-ever proof turn its pro
+         * badge on.
+         *
+         * Must run before anything writes the access expiry from the same response — that is one of the
+         * three values read here.
+         */
+        internal fun storeProof(
+            userProfile: MutableUserProfile,
+            proof: ProProof,
+            rotatingPrivateKey: ByteArray,
+        ) {
+            // Upgrade guard: only replace the proof if it extends coverage (monotonic merge;
+            // same-period races round to the same expiry -> byte-identical -> no-op). Avoids
+            // churning a proof another device just landed.
+            val current = userProfile.getProConfig()?.proProof
+            if (current == null || proof.expirySeconds > current.expirySeconds) {
+                // First-ever proof: enable the pro badge feature. The badge is off by default because
+                // being visible as a subscriber is the user's choice, and these three absent values are
+                // the only evidence the account has never had one to express. Asking "is Pro active"
+                // instead would re-enable the badge at every renewal for a subscriber who had turned it
+                // off, leaving them no way to make it stick.
+                //
+                // Mirrors iOS `SessionProManager.applyProofSuccess` and Desktop `ducks/proBackendData.ts`.
+                if (current == null &&
+                    userProfile.getProAccessExpiry() == null &&
+                    userProfile.getProFeatures().isEmpty
+                ) {
+                    userProfile.setProBadge(true)
+                }
+
+                userProfile.setProConfig(ProConfig(
+                    proProof = proof,
+                    rotatingPrivateKey = rotatingPrivateKey))
+            }
+        }
+
         /**
          * Minimum spacing between proof requests. **Shared cross-client contract** — iOS
          * (`SessionProManager.reconcileProofRenewal`) and Desktop use exactly these values; keep them

app/src/test/java/org/thoughtcrime/securesms/pro/ProBadgeFirstProofTest.kt

@@ -0,0 +1,106 @@
+package org.thoughtcrime.securesms.pro
+
+import io.mockk.mockk
+import io.mockk.every
+import io.mockk.verify
+import network.loki.messenger.libsession_util.MutableUserProfile
+import network.loki.messenger.libsession_util.pro.ProConfig
+import network.loki.messenger.libsession_util.pro.ProProof
+import network.loki.messenger.libsession_util.protocol.ProProfileFeature
+import network.loki.messenger.libsession_util.protocol.ProProfileFeatures
+import network.loki.messenger.libsession_util.util.toBitSet
+import org.junit.Test
+import org.thoughtcrime.securesms.pro.ProProofGenerationWorker.Companion.storeProof
+
+/**
+ * Covers the pro badge being enabled on an account's first-ever proof.
+ *
+ * The badge is off by default and the settings toggle is the user's own choice, so the guard asks
+ * whether this account has EVER held Pro rather than whether it holds Pro now. The second question
+ * is the one that ships easily and is annoying to diagnose: a subscriber who turns the badge off has
+ * it turned back on at the next renewal, forever, with no way to make it stick. The cases below that
+ * assert the badge is NOT touched are the ones carrying that regression.
+ */
+class ProBadgeFirstProofTest {
+
+    private val rotatingPrivateKey = ByteArray(64) { 1 }
+
+    @Test
+    fun `a first-ever proof turns the pro badge on`() {
+        val userProfile = neverHadPro()
+
+        storeProof(userProfile, proof(expirySeconds = 2_000), rotatingPrivateKey)
+
+        verify(exactly = 1) { userProfile.setProBadge(true) }
+        verify(exactly = 1) { userProfile.setProConfig(any()) }
+    }
+
+    @Test
+    fun `a subscriber who turned the badge off keeps it off when the proof renews`() {
+        val userProfile = userProfile(
+            proof = proof(expirySeconds = 1_000),
+            accessExpirySeconds = 1_500L
+        )
+
+        storeProof(userProfile, proof(expirySeconds = 2_000), rotatingPrivateKey)
+
+        verify(exactly = 0) { userProfile.setProBadge(any()) }
+        verify(exactly = 1) { userProfile.setProConfig(any()) }
+    }
+
+    @Test
+    fun `a lapsed subscriber who turned the badge off keeps it off when the plan is renewed`() {
+        // The proof is cleared when a plan lapses, so the access expiry is the only evidence left that
+        // this account has been Pro before — and the only thing standing between the user's choice and
+        // a fresh proof reinstating the badge.
+        val userProfile = userProfile(proof = null, accessExpirySeconds = 1_500L)
+
+        storeProof(userProfile, proof(expirySeconds = 2_000), rotatingPrivateKey)
+
+        verify(exactly = 0) { userProfile.setProBadge(any()) }
+        verify(exactly = 1) { userProfile.setProConfig(any()) }
+    }
+
+    @Test
+    fun `a profile feature already in config means the account has been Pro before`() {
+        val userProfile = userProfile(
+            features = listOf(ProProfileFeature.ANIMATED_AVATAR).toBitSet()
+        )
+
+        storeProof(userProfile, proof(expirySeconds = 2_000), rotatingPrivateKey)
+
+        verify(exactly = 0) { userProfile.setProBadge(any()) }
+    }
+
+    @Test
+    fun `a proof that does not extend coverage is not stored, and enables nothing`() {
+        val userProfile = userProfile(
+            proof = proof(expirySeconds = 3_000),
+            accessExpirySeconds = 3_500L
+        )
+
+        storeProof(userProfile, proof(expirySeconds = 2_000), rotatingPrivateKey)
+
+        verify(exactly = 0) { userProfile.setProConfig(any()) }
+        verify(exactly = 0) { userProfile.setProBadge(any()) }
+    }
+
+    private fun neverHadPro(): MutableUserProfile = userProfile()
+
+    private fun userProfile(
+        proof: ProProof? = null,
+        accessExpirySeconds: Long? = null,
+        features: ProProfileFeatures = emptyList<ProProfileFeature>().toBitSet(),
+    ): MutableUserProfile = mockk(relaxed = true) {
+        every { getProConfig() } returns proof?.let { ProConfig(it, rotatingPrivateKey) }
+        every { getProAccessExpiry() } returns accessExpirySeconds
+        every { getProFeatures() } returns features
+    }
+
+    private fun proof(expirySeconds: Long): ProProof = ProProof(
+        revocationTagHex = "aa".repeat(32),
+        rotatingPubKeyHex = "bb".repeat(32),
+        expirySeconds = expirySeconds,
+        signatureHex = "cc".repeat(64),
+    )
+}