Unregister push token on signout (#8661)

* unregister push - WIP

* create agent and submit push token revokation

* mock unregisterpush

* fix import

* Add proxy headers

---------

Co-authored-by: Eric Bailey <git@esb.lol>

src/lib/constants.ts

@@ -241,6 +241,10 @@ export const BLUESKY_MOD_SERVICE_HEADERS = {
   'atproto-proxy': `${BSKY_LABELER_DID}#atproto_labeler`,
 }
 
+export const BLUESKY_NOTIF_SERVICE_HEADERS = {
+  'atproto-proxy': `${BLUESKY_PROXY_DID}#bsky_notif`,
+}
+
 export const webLinks = {
   tos: `https://bsky.social/about/support/tos`,
   privacy: `https://bsky.social/about/support/privacy-policy`,

src/lib/notifications/notifications.ts

@@ -2,10 +2,15 @@ import {useCallback, useEffect} from 'react'
 import {Platform} from 'react-native'
 import * as Notifications from 'expo-notifications'
 import {getBadgeCountAsync, setBadgeCountAsync} from 'expo-notifications'
-import {type AppBskyNotificationRegisterPush, type AtpAgent} from '@atproto/api'
+import {type AtpAgent} from '@atproto/api'
+import {type AppBskyNotificationRegisterPush} from '@atproto/api'
 import debounce from 'lodash.debounce'
 
-import {PUBLIC_APPVIEW_DID, PUBLIC_STAGING_APPVIEW_DID} from '#/lib/constants'
+import {
+  BLUESKY_NOTIF_SERVICE_HEADERS,
+  PUBLIC_APPVIEW_DID,
+  PUBLIC_STAGING_APPVIEW_DID,
+} from '#/lib/constants'
 import {logger as notyLogger} from '#/lib/notifications/util'
 import {isNetworkError} from '#/lib/strings/errors'
 import {isNative} from '#/platform/detection'
@@ -44,7 +49,9 @@ async function _registerPushToken({
 
     notyLogger.debug(`registerPushToken: registering`, {...payload})
 
-    await agent.app.bsky.notification.registerPush(payload)
+    await agent.app.bsky.notification.registerPush(payload, {
+      headers: BLUESKY_NOTIF_SERVICE_HEADERS,
+    })
 
     notyLogger.debug(`registerPushToken: success`)
   } catch (error) {
@@ -286,3 +293,33 @@ export async function resetBadgeCount() {
   await BackgroundNotificationHandler.setBadgeCountAsync(0)
   await setBadgeCountAsync(0)
 }
+
+export async function unregisterPushToken(agents: AtpAgent[]) {
+  if (!isNative) return
+
+  try {
+    const token = await getPushToken()
+    if (token) {
+      for (const agent of agents) {
+        await agent.app.bsky.notification.unregisterPush(
+          {
+            serviceDid: agent.serviceUrl.hostname.includes('staging')
+              ? PUBLIC_STAGING_APPVIEW_DID
+              : PUBLIC_APPVIEW_DID,
+            platform: Platform.OS,
+            token: token.data,
+            appId: 'xyz.blueskyweb.app',
+          },
+          {
+            headers: BLUESKY_NOTIF_SERVICE_HEADERS,
+          },
+        )
+        notyLogger.debug(`Push token unregistered for ${agent.session?.handle}`)
+      }
+    } else {
+      notyLogger.debug('Tried to unregister push token, but could not find one')
+    }
+  } catch (error) {
+    notyLogger.debug('Failed to unregister push token', {message: error})
+  }
+}

src/state/session/__tests__/session-test.ts

@@ -12,6 +12,11 @@ jest.mock('jwt-decode', () => ({
 
 jest.mock('../../birthdate')
 jest.mock('../../../ageAssurance/data')
+jest.mock('#/lib/notifications/notifications', () => ({
+  unregisterPushToken(_agents: BskyAgent[]) {
+    return Promise.resolve()
+  },
+}))
 
 describe('session', () => {
   it('can log in and out', () => {

src/state/session/reducer.ts

@@ -1,8 +1,11 @@
-import {type AtpSessionEvent, type BskyAgent} from '@atproto/api'
+import {type AtpAgent, type AtpSessionEvent} from '@atproto/api'
 
+import {unregisterPushToken} from '#/lib/notifications/notifications'
+import {logger} from '#/lib/notifications/util'
 import {createPublicAgent} from './agent'
 import {wrapSessionReducerForLogging} from './logging'
 import {type SessionAccount} from './types'
+import {createTemporaryAgentsAndResume} from './util'
 
 // A hack so that the reducer can't read anything from the agent.
 // From the reducer's point of view, it should be a completely opaque object.
@@ -137,6 +140,23 @@ let reducer = (state: State, action: Action): State => {
     }
     case 'removed-account': {
       const {accountDid} = action
+
+      // side effect
+      const account = state.accounts.find(a => a.did === accountDid)
+      if (account) {
+        createTemporaryAgentsAndResume([account])
+          .then(agents => unregisterPushToken(agents))
+          .then(() =>
+            logger.debug('Push token unregistered', {did: accountDid}),
+          )
+          .catch(err => {
+            logger.error('Failed to unregister push token', {
+              did: accountDid,
+              error: err,
+            })
+          })
+      }
+
       return {
         accounts: state.accounts.filter(a => a.did !== accountDid),
         currentAgentState:
@@ -148,9 +168,26 @@ let reducer = (state: State, action: Action): State => {
     }
     case 'logged-out-current-account': {
       const {currentAgentState} = state
+      const accountDid = currentAgentState.did
+      // side effect
+      const account = state.accounts.find(a => a.did === accountDid)
+      if (account && accountDid) {
+        createTemporaryAgentsAndResume([account])
+          .then(agents => unregisterPushToken(agents))
+          .then(() =>
+            logger.debug('Push token unregistered', {did: accountDid}),
+          )
+          .catch(err => {
+            logger.error('Failed to unregister push token', {
+              did: accountDid,
+              error: err,
+            })
+          })
+      }
+
       return {
         accounts: state.accounts.map(a =>
-          a.did === currentAgentState.did
+          a.did === accountDid
             ? {
                 ...a,
                 refreshJwt: undefined,
@@ -163,6 +200,15 @@ let reducer = (state: State, action: Action): State => {
       }
     }
     case 'logged-out-every-account': {
+      createTemporaryAgentsAndResume(state.accounts)
+        .then(agents => unregisterPushToken(agents))
+        .then(() => logger.debug('Push token unregistered'))
+        .catch(err => {
+          logger.error('Failed to unregister push token', {
+            error: err,
+          })
+        })
+
       return {
         accounts: state.accounts.map(a => ({
           ...a,
@@ -187,7 +233,7 @@ let reducer = (state: State, action: Action): State => {
     }
     case 'partial-refresh-session': {
       const {accountDid, patch} = action
-      const agent = state.currentAgentState.agent as BskyAgent
+      const agent = state.currentAgentState.agent as AtpAgent
 
       /*
        * Only mutating values that are safe. Be very careful with this.

src/state/session/util.ts

@@ -1,8 +1,10 @@
+import AtpAgent from '@atproto/api'
 import {jwtDecode} from 'jwt-decode'
 
 import {isJwtExpired} from '#/lib/jwt'
 import {hasProp} from '#/lib/type-guards'
 import * as persisted from '#/state/persisted'
+import {sessionAccountToSession} from './agent'
 import {type SessionAccount} from './types'
 
 export function readLastActiveAccount() {
@@ -28,3 +30,32 @@ export function isSessionExpired(account: SessionAccount) {
     return true
   }
 }
+
+/**
+ * Creates and attempted to resumeSession for every stored session.
+ * Intended to be used to send push token revokations just before logout.
+ */
+export async function createTemporaryAgentsAndResume(
+  accounts: SessionAccount[],
+) {
+  const agents = await Promise.allSettled(
+    accounts.map(async account => {
+      const agent: AtpAgent = new AtpAgent({service: account.service})
+      if (account.pdsUrl) {
+        agent.sessionManager.pdsUrl = new URL(account.pdsUrl)
+      }
+
+      const session = sessionAccountToSession(account)
+      const res = await agent.resumeSession(session)
+      if (!res.success) throw new Error('Failed to resume session')
+
+      agent.assertAuthenticated() // confirm auth success
+
+      return agent
+    }),
+  )
+
+  return agents
+    .filter(x => x.status === 'fulfilled')
+    .map(promise => promise.value)
+}