fix(push): read badge counts only for recipients with a device (#3521)

fluxer_api/src/api/admin/tests/LimitConfigRemovedKeys.test.ts

@@ -0,0 +1,142 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+import type {TestAccount} from '@app/api/auth/tests/AuthTestUtils';
+import {createTestAccount, setUserACLs} from '@app/api/auth/tests/AuthTestUtils';
+import {getConfig} from '@app/api/Config';
+import {getInstanceConfigRepository, getLimitConfigService} from '@app/api/middleware/ServiceSingletons';
+import {type ApiTestHarness, createApiTestHarness} from '@app/api/test/ApiTestHarness';
+import {HTTP_STATUS} from '@app/api/test/TestConstants';
+import {createBuilder} from '@app/api/test/TestRequestBuilder';
+import {AdminACLs} from '@fluxer/constants/src/AdminACLs';
+import {APIErrorCodes} from '@fluxer/constants/src/ApiErrorCodes';
+import {ValidationErrorCodes} from '@fluxer/constants/src/ValidationErrorCodes';
+import type {LimitRule} from '@fluxer/limits/src/LimitTypes';
+import type {LimitConfigGetResponse} from '@fluxer/schema/src/domains/admin/AdminSchemas';
+import {afterAll, beforeAll, beforeEach, describe, expect, it} from 'vitest';
+import type {z} from 'zod';
+
+type LimitConfigResponse = z.infer<typeof LimitConfigGetResponse>;
+
+const REMOVED_KEY = 'max_created_packs';
+
+describe.each([true, false])('limit config with a key from an older release (self-hosted: %s)', (selfHosted) => {
+	let harness: ApiTestHarness;
+	let originalSelfHosted: boolean;
+
+	beforeAll(async () => {
+		originalSelfHosted = getConfig().instance.selfHosted;
+		getConfig().instance.selfHosted = selfHosted;
+		harness = await createApiTestHarness();
+	});
+
+	beforeEach(async () => {
+		await harness.reset();
+		await getLimitConfigService().refreshCache();
+	});
+
+	afterAll(async () => {
+		await harness.shutdown();
+		getConfig().instance.selfHosted = originalSelfHosted;
+	});
+
+	async function createAdmin(): Promise<TestAccount> {
+		return await setUserACLs(harness, await createTestAccount(harness), [
+			AdminACLs.AUTHENTICATE,
+			AdminACLs.INSTANCE_LIMIT_CONFIG_VIEW,
+			AdminACLs.INSTANCE_LIMIT_CONFIG_UPDATE,
+		]);
+	}
+
+	async function storeRuleFromOlderRelease(): Promise<void> {
+		const config = getLimitConfigService().getConfigSnapshot();
+		const legacyRule = {
+			id: 'vip',
+			filters: {traits: ['vip']},
+			limits: {max_bio_length: 200, [REMOVED_KEY]: 5},
+		} as LimitRule;
+		await getInstanceConfigRepository().setLimitConfig({
+			traitDefinitions: [...config.traitDefinitions, 'vip'],
+			rules: [...config.rules, legacyRule],
+		});
+		await getLimitConfigService().refreshCache();
+	}
+
+	it('leaves the removed key out when reading the stored config', async () => {
+		const admin = await createAdmin();
+		await storeRuleFromOlderRelease();
+		const response = await createBuilder<LimitConfigResponse>(harness, admin.token)
+			.get('/admin/limit-config')
+			.execute();
+		const vip = response.limit_config.rules.find((rule) => rule.id === 'vip');
+		expect(vip?.limits).toEqual({max_bio_length: 200});
+	});
+
+	it('saves the config it read back and drops the removed key from storage', async () => {
+		const admin = await createAdmin();
+		await storeRuleFromOlderRelease();
+		const {limit_config: config} = await createBuilder<LimitConfigResponse>(harness, admin.token)
+			.get('/admin/limit-config')
+			.execute();
+		await createBuilder<LimitConfigResponse>(harness, admin.token)
+			.put('/admin/limit-config')
+			.body({
+				limit_config: {
+					traitDefinitions: config.traitDefinitions,
+					rules: config.rules.map(({modifiedFields: _modifiedFields, ...rule}) => rule),
+				},
+			})
+			.execute();
+		const {config: stored} = await getInstanceConfigRepository().readLimitConfigInputs();
+		const vip = stored?.rules.find((rule) => rule.id === 'vip');
+		expect(vip?.limits).toEqual({max_bio_length: 200});
+	});
+
+	it('drops a rule whose only limits were removed keys', async () => {
+		const admin = await createAdmin();
+		const config = getLimitConfigService().getConfigSnapshot();
+		await getInstanceConfigRepository().setLimitConfig({
+			traitDefinitions: [...config.traitDefinitions, 'packs'],
+			rules: [...config.rules, {id: 'packs', filters: {traits: ['packs']}, limits: {[REMOVED_KEY]: 5}} as LimitRule],
+		});
+		await getLimitConfigService().refreshCache();
+		const {limit_config: read} = await createBuilder<LimitConfigResponse>(harness, admin.token)
+			.get('/admin/limit-config')
+			.execute();
+		expect(read.rules.some((rule) => rule.id === 'packs')).toBe(false);
+		await createBuilder<LimitConfigResponse>(harness, admin.token)
+			.put('/admin/limit-config')
+			.body({
+				limit_config: {
+					traitDefinitions: read.traitDefinitions,
+					rules: read.rules.map(({modifiedFields: _modifiedFields, ...rule}) => rule),
+				},
+			})
+			.execute();
+	});
+
+	it('names the unknown key in the validation error', async () => {
+		const admin = await createAdmin();
+		const {limit_config: config} = await createBuilder<LimitConfigResponse>(harness, admin.token)
+			.get('/admin/limit-config')
+			.execute();
+		const response = await createBuilder<{errors: Array<{path: string; code: string}>}>(harness, admin.token)
+			.put('/admin/limit-config')
+			.body({
+				limit_config: {
+					traitDefinitions: [...config.traitDefinitions, 'vip'],
+					rules: [
+						...config.rules.map(({modifiedFields: _modifiedFields, ...rule}) => rule),
+						{id: 'vip', filters: {traits: ['vip']}, limits: {max_bio_length: 200, [REMOVED_KEY]: 5}},
+					],
+				},
+			})
+			.expect(HTTP_STATUS.BAD_REQUEST, APIErrorCodes.INVALID_FORM_BODY)
+			.execute();
+		expect(response.errors).toEqual([
+			expect.objectContaining({
+				path: `limit_config.rules.${config.rules.length}.limits.${REMOVED_KEY}`,
+				code: ValidationErrorCodes.INVALID_FORMAT,
+			}),
+		]);
+	});
+});

fluxer_api/src/api/constants/LimitConfig.ts

@@ -103,7 +103,7 @@ export function mergeWithCurrentDefaults(
 	const mergedRules: Array<LimitRule> = [];
 	const existingRulesMap = new Map<string, LimitRule>();
 	for (const rule of stored.rules) {
-		existingRulesMap.set(rule.id, {...rule, limits: clampToAllowedRange(rule.limits)});
+		existingRulesMap.set(rule.id, {...rule, limits: readKnownLimits(rule.limits)});
 	}
 	for (const defaultRule of newDefaults.rules) {
 		const existingRule = existingRulesMap.get(defaultRule.id);
@@ -125,6 +125,9 @@ export function mergeWithCurrentDefaults(
 		existingRulesMap.delete(defaultRule.id);
 	}
 	for (const [, rule] of existingRulesMap) {
+		if (Object.keys(rule.limits).length === 0) {
+			continue;
+		}
 		const modifiedFields = findModifiedLimits(rule.limits, defaultBaseRule.limits);
 		mergedRules.push({
 			id: rule.id,
@@ -231,16 +234,17 @@ function wasSavedUnderOtherDefaults(
 	);
 }
 
-function clampToAllowedRange(limits: Partial<Record<LimitKey, number>>): Partial<Record<LimitKey, number>> {
+function readKnownLimits(limits: Partial<Record<LimitKey, number>>): Partial<Record<LimitKey, number>> {
-	const clamped = {...limits};
+	const known: Partial<Record<LimitKey, number>> = {};
 	for (const key of LIMIT_KEYS) {
-		const value = clamped[key];
+		const value = limits[key];
+		if (value === undefined) {
+			continue;
+		}
 		const {min = 0, max} = LIMIT_KEY_METADATA[key];
-		if (value !== undefined && max !== undefined) {
+		known[key] = max === undefined ? value : Math.min(Math.max(value, min), max);
-			clamped[key] = Math.min(Math.max(value, min), max);
 	}
-	}
+	return known;
-	return clamped;
 }
 
 function pickModifiedLimits(

fluxer_push/src/delivery.rs

@@ -9,7 +9,7 @@ use crate::metrics::{JobKind, JobRejection, Provider, elapsed_ms};
 use crate::payload;
 use crate::providers::{self, SendOutcome};
 use crate::retry::{self, RETRY_DEADLINE};
-use crate::rpc::RpcError;
+use crate::rpc::{RpcClient, RpcError};
 use crate::server::AppState;
 use crate::subscription::{Platform, Subscription};
 use crate::unix_millis;
@@ -353,13 +353,7 @@ async fn run_message_job(
 ) -> anyhow::Result<()> {
     let started_ms = now_ms();
     let deadline = Instant::now() + RETRY_DEADLINE;
-    let (badges, subscriptions) = lookup(deadline, || async {
+    let (badges, subscriptions) = lookup_recipients(&state.rpc, &job.user_ids, deadline).await?;
-        tokio::try_join!(
-            state.rpc.badge_counts(&job.user_ids),
-            state.rpc.push_subscriptions(&job.user_ids)
-        )
-    })
-    .await?;
     state.metrics.record_recipients(job.user_ids.len() as u64);
 
     let envelopes = job
@@ -410,13 +404,7 @@ async fn run_clear_job(state: &AppState, sends: &Semaphore, job: ClearJob) -> an
     let started_ms = now_ms();
     let deadline = Instant::now() + RETRY_DEADLINE;
     let user_ids = std::slice::from_ref(&job.user_id);
-    let (badges, subscriptions) = lookup(deadline, || async {
+    let (badges, subscriptions) = lookup_recipients(&state.rpc, user_ids, deadline).await?;
-        tokio::try_join!(
-            state.rpc.badge_counts(user_ids),
-            state.rpc.push_subscriptions(user_ids)
-        )
-    })
-    .await?;
     state.metrics.record_recipients(1);
 
     let envelope = payload::web_push_clear(&job, badge_of(&badges, &job.user_id));
@@ -503,6 +491,24 @@ async fn run_ring_job(state: &AppState, sends: &Semaphore, job: RingJob) -> anyh
     Ok(())
 }
 
+async fn lookup_recipients(
+    rpc: &RpcClient,
+    user_ids: &[String],
+    deadline: Instant,
+) -> Result<(HashMap<String, u32>, HashMap<String, Vec<Subscription>>), RpcError> {
+    let subscriptions = lookup(deadline, || rpc.push_subscriptions(user_ids)).await?;
+    let subscribed_ids: Vec<String> = user_ids
+        .iter()
+        .filter(|user_id| subscribed(&subscriptions, user_id))
+        .cloned()
+        .collect();
+    if subscribed_ids.is_empty() {
+        return Ok((HashMap::new(), subscriptions));
+    }
+    let badges = lookup(deadline, || rpc.badge_counts(&subscribed_ids)).await?;
+    Ok((badges, subscriptions))
+}
+
 async fn lookup<T, F, Fut>(deadline: Instant, mut call: F) -> Result<T, RpcError>
 where
     F: FnMut() -> Fut,
@@ -708,4 +714,95 @@ mod tests {
         assert!(Audience::Standard.admits(&subscription("ios_apns", false)));
         assert!(Audience::Standard.admits(&subscription("ios_apns", true)));
     }
+
+    async fn fake_api(
+        subscribed: &'static [&'static str],
+    ) -> (RpcClient, Arc<std::sync::Mutex<Vec<Value>>>) {
+        let calls = Arc::new(std::sync::Mutex::new(Vec::new()));
+        let recorded = Arc::clone(&calls);
+        let app = axum::Router::new().route(
+            "/",
+            axum::routing::post(move |body: axum::body::Bytes| {
+                let recorded = Arc::clone(&recorded);
+                async move {
+                    let request: Value = serde_json::from_slice(&body).unwrap();
+                    recorded.lock().unwrap().push(request.clone());
+                    let user_ids = request["user_ids"].as_array().cloned().unwrap_or_default();
+                    let data = match request["type"].as_str().unwrap() {
+                        "get_push_subscriptions" => serde_json::json!(
+                            subscribed
+                                .iter()
+                                .map(|user_id| {
+                                    (
+                                        (*user_id).to_owned(),
+                                        serde_json::json!([{
+                                            "subscription_id": "sub",
+                                            "endpoint": "token",
+                                            "p256dh_key": null,
+                                            "auth_key": null,
+                                            "platform": "android_fcm",
+                                            "app_id": null,
+                                            "provider_environment": null
+                                        }]),
+                                    )
+                                })
+                                .collect::<serde_json::Map<_, _>>()
+                        ),
+                        "get_badge_counts" => serde_json::json!({
+                            "badge_counts": user_ids
+                                .iter()
+                                .map(|user_id| (user_id.as_str().unwrap().to_owned(), serde_json::json!(3)))
+                                .collect::<serde_json::Map<_, _>>()
+                        }),
+                        other => panic!("unexpected rpc {other}"),
+                    };
+                    serde_json::json!({ "data": data }).to_string()
+                }
+            }),
+        );
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let url = format!("http://{}/", listener.local_addr().unwrap());
+        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
+        let rpc = RpcClient::new(
+            &crate::config::RpcConfig {
+                url,
+                auth_token: fluxer_common::secret::SecretString::new("token".to_owned()),
+            },
+            reqwest::Client::new(),
+            Arc::new(crate::metrics::Metrics::new()),
+        );
+        (rpc, calls)
+    }
+
+    fn badge_requests(calls: &std::sync::Mutex<Vec<Value>>) -> Vec<Value> {
+        calls
+            .lock()
+            .unwrap()
+            .iter()
+            .filter(|call| call["type"] == "get_badge_counts")
+            .map(|call| call["user_ids"].clone())
+            .collect()
+    }
+
+    #[tokio::test]
+    async fn badge_counts_are_read_only_for_recipients_with_a_device() {
+        let (rpc, calls) = fake_api(&["2"]).await;
+        let user_ids: Vec<String> = (1..=500).map(|id| id.to_string()).collect();
+        let deadline = Instant::now() + Duration::from_secs(5);
+        let (badges, subscriptions) = lookup_recipients(&rpc, &user_ids, deadline).await.unwrap();
+        assert_eq!(badge_requests(&calls), vec![serde_json::json!(["2"])]);
+        assert_eq!(badges, HashMap::from([("2".to_owned(), 3)]));
+        assert!(subscribed(&subscriptions, "2"));
+    }
+
+    #[tokio::test]
+    async fn no_badge_counts_are_read_when_nobody_has_a_device() {
+        let (rpc, calls) = fake_api(&[]).await;
+        let user_ids: Vec<String> = (1..=500).map(|id| id.to_string()).collect();
+        let deadline = Instant::now() + Duration::from_secs(5);
+        let (badges, _subscriptions) = lookup_recipients(&rpc, &user_ids, deadline).await.unwrap();
+        assert!(badge_requests(&calls).is_empty());
+        assert!(badges.is_empty());
+        assert_eq!(calls.lock().unwrap().len(), 1);
+    }
 }

packages/schema/src/domains/admin/AdminSchemas.ts

@@ -951,14 +951,12 @@ const LimitRuleSchema = z.object({
 	filters: LimitFilterSchema.optional().describe('Optional filters that scope the rule'),
 	limits: z
 		.record(z.string(), NonNegativeSafeIntegerType)
-		.refine(
-			(limits) => {
-				const limitKeys = Object.keys(limits);
-				return limitKeys.every((key) => (LIMIT_KEYS as ReadonlyArray<string>).includes(key));
-			},
-			{error: 'Invalid limit key detected'},
-		)
 		.superRefine((limits, ctx) => {
+			for (const key of Object.keys(limits)) {
+				if (!(LIMIT_KEYS as ReadonlyArray<string>).includes(key)) {
+					ctx.addIssue({code: 'custom', message: ValidationErrorCodes.INVALID_FORMAT, path: [key]});
+				}
+			}
 			for (const key of LIMIT_KEYS) {
 				const value = limits[key];
 				const {min = 0, max} = LIMIT_KEY_METADATA[key];