remote_settings/
storage.rs

1/* This Source Code Form is subject to the terms of the Mozilla Public
2 * License, v. 2.0. If a copy of the MPL was not distributed with this
3 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
4
5use crate::{
6    client::CollectionMetadata, client::CollectionSignature,
7    schema::RemoteSettingsConnectionInitializer, Attachment, Error, RemoteSettingsRecord, Result,
8};
9use camino::Utf8PathBuf;
10use rusqlite::{params, Connection, OpenFlags, OptionalExtension, Transaction};
11use serde_json;
12use sha2::{Digest, Sha256};
13use std::io;
14
15use sql_support::{open_database::open_database_with_flags, run_maintenance, ConnExt};
16
17/// Internal storage type
18///
19/// This will store downloaded records/attachments in a SQLite database.  Nothing is implemented
20/// yet other than the initial API.
21///
22/// Most methods input a `collection_url` parameter, is a URL that includes the remote settings
23/// server, bucket, and collection. If the `collection_url` for a get method does not match the one
24/// for a set method, then this means the application has switched their remote settings config and
25/// [Storage] should pretend like nothing is stored in the database.
26///
27/// The reason for this is the [crate::RemoteSettingsService::update_config] method.  If a consumer
28/// passes a new server or bucket to `update_config`, we don't want to be using cached data from
29/// the previous config.
30pub struct Storage {
31    path: Utf8PathBuf,
32    conn: ConnectionCell,
33}
34
35impl Storage {
36    pub fn new(path: Utf8PathBuf) -> Self {
37        Self {
38            path,
39            conn: ConnectionCell::Uninitialized,
40        }
41    }
42
43    fn transaction(&mut self) -> Result<Transaction<'_>> {
44        match &self.conn {
45            ConnectionCell::Uninitialized => {
46                self.ensure_dir()?;
47                self.conn = ConnectionCell::Initialized(open_database_with_flags(
48                    &self.path,
49                    OpenFlags::default(),
50                    &RemoteSettingsConnectionInitializer,
51                )?);
52            }
53            ConnectionCell::Initialized(_) => (),
54            ConnectionCell::Closed => return Err(Error::DatabaseClosed),
55        }
56        match &mut self.conn {
57            ConnectionCell::Initialized(conn) => Ok(conn.transaction()?),
58            _ => unreachable!(),
59        }
60    }
61
62    pub fn ensure_dir(&self) -> Result<()> {
63        if self.path == ":memory:" {
64            return Ok(());
65        }
66        let Some(dir) = self.path.parent() else {
67            return Ok(());
68        };
69        if !std::fs::exists(dir).map_err(Error::CreateDirError)? {
70            match std::fs::create_dir(dir) {
71                Ok(()) => (),
72                // Another thread created the directory, just ignore the error.
73                Err(e) if e.kind() == io::ErrorKind::AlreadyExists => (),
74                Err(e) => return Err(Error::CreateDirError(e)),
75            }
76        }
77        Ok(())
78    }
79
80    pub fn close(&mut self) {
81        self.conn = ConnectionCell::Closed;
82    }
83
84    /// Get the last modified timestamp for the stored records
85    ///
86    /// Returns None if no records are stored or if `collection_url` does not match the
87    /// last `collection_url` passed to `insert_collection_content`
88    pub fn get_last_modified_timestamp(&mut self, collection_url: &str) -> Result<Option<u64>> {
89        let tx = self.transaction()?;
90        let mut stmt =
91            tx.prepare("SELECT last_modified FROM collection_metadata WHERE collection_url = ?")?;
92        let result: Option<u64> = stmt
93            .query_row((collection_url,), |row| row.get(0))
94            .optional()?;
95        Ok(result)
96    }
97
98    /// Get cached records for this collection
99    ///
100    /// Returns None if no records are stored or if `collection_url` does not match the `collection_url` passed
101    /// to `insert_collection_content`.
102    pub fn get_records(
103        &mut self,
104        collection_url: &str,
105    ) -> Result<Option<Vec<RemoteSettingsRecord>>> {
106        let tx = self.transaction()?;
107
108        let fetched = tx.exists(
109            "SELECT 1 FROM collection_metadata WHERE collection_url = ?",
110            (collection_url,),
111        )?;
112        let result = if fetched {
113            // If fetched before, get the records from the records table
114            let records: Vec<RemoteSettingsRecord> = tx
115                .prepare("SELECT data FROM records WHERE collection_url = ?")?
116                .query_map(params![collection_url], |row| row.get::<_, Vec<u8>>(0))?
117                .map(|data| serde_json::from_slice(&data.unwrap()).unwrap())
118                .collect();
119
120            Ok(Some(records))
121        } else {
122            Ok(None)
123        };
124
125        tx.commit()?;
126        result
127    }
128
129    /// Get cached metadata for this collection
130    ///
131    /// Returns None if no data is stored or if `collection_url` does not match the `collection_url` passed
132    /// to `insert_collection_content`.
133    pub fn get_collection_metadata(
134        &mut self,
135        collection_url: &str,
136    ) -> Result<Option<CollectionMetadata>> {
137        let tx = self.transaction()?;
138        // signatures is a JSON array of objects with "signature" and "x5u" fields,
139        // so we need to iterate through the rows and construct the list of signatures.
140        // we use LEFT JOIN to return a row even if list of signatures is empty.
141        let mut stmt_metadata = tx.prepare(
142            "
143            SELECT
144                cm.bucket,
145                json_extract(sig.value, '$.x5u') AS x5u,
146                json_extract(sig.value, '$.signature') AS signature,
147                json_extract(sig.value, '$.mode') AS mode
148            FROM collection_metadata AS cm
149            LEFT JOIN json_each(cm.signatures) AS sig ON true
150            WHERE cm.collection_url = ?
151            ",
152        )?;
153
154        let mut rows = stmt_metadata.query(params![collection_url])?;
155        let mut bucket: Option<String> = None;
156        let mut signatures = Vec::new();
157
158        while let Some(row) = rows.next()? {
159            // bucket should be the same for every row, so just set it once
160            if bucket.is_none() {
161                bucket = Some(row.get(0)?);
162            }
163            let x5u: Option<String> = row.get(1)?;
164            let signature: Option<String> = row.get(2)?;
165            let mode: Option<String> = row.get(3)?;
166            if let (Some(x5u), Some(signature), Some(mode)) = (x5u, signature, mode) {
167                signatures.push(CollectionSignature {
168                    signature,
169                    x5u,
170                    mode,
171                });
172            }
173        }
174        match bucket {
175            Some(bucket) => Ok(Some(CollectionMetadata { bucket, signatures })),
176            None => Ok(None),
177        }
178    }
179
180    /// Get cached attachment data
181    ///
182    /// This returns the last attachment data sent to [Self::set_attachment].
183    ///
184    /// Returns None if no attachment data is stored or if `collection_url` does not match the `collection_url`
185    /// passed to `set_attachment`.
186    pub fn get_attachment(
187        &mut self,
188        collection_url: &str,
189        metadata: Attachment,
190    ) -> Result<Option<Vec<u8>>> {
191        let tx = self.transaction()?;
192        let mut stmt =
193            tx.prepare("SELECT data FROM attachments WHERE id = ? AND collection_url = ?")?;
194
195        if let Some(data) = stmt
196            .query_row((metadata.location, collection_url), |row| {
197                row.get::<_, Vec<u8>>(0)
198            })
199            .optional()?
200        {
201            // Return None if data doesn't match expected metadata
202            if data.len() as u64 != metadata.size {
203                return Ok(None);
204            }
205            let hash = format!("{:x}", Sha256::digest(&data));
206            if hash != metadata.hash {
207                return Ok(None);
208            }
209            Ok(Some(data))
210        } else {
211            Ok(None)
212        }
213    }
214
215    /// Set cached content for this collection.
216    pub fn insert_collection_content(
217        &mut self,
218        collection_url: &str,
219        records: &[RemoteSettingsRecord],
220        last_modified: u64,
221        metadata: CollectionMetadata,
222    ) -> Result<()> {
223        let tx = self.transaction()?;
224
225        // Delete ALL existing records and metadata for with different collection_urls.
226        //
227        // This way, if a user (probably QA) switches the remote settings server in the middle of a
228        // browser sessions, we'll delete the stale data from the previous server.
229        tx.execute(
230            "DELETE FROM records where collection_url <> ?",
231            [collection_url],
232        )?;
233        tx.execute(
234            "DELETE FROM collection_metadata where collection_url <> ?",
235            [collection_url],
236        )?;
237
238        Self::update_record_rows(&tx, collection_url, records)?;
239        Self::update_collection_metadata(&tx, collection_url, last_modified, metadata)?;
240        Self::cleanup_orphaned_attachments(&tx, collection_url)?;
241        tx.commit()?;
242        Ok(())
243    }
244
245    /// Insert/remove/update rows in the records table based on a records list
246    ///
247    /// Returns the max last modified record from the list
248    fn update_record_rows(
249        tx: &Transaction<'_>,
250        collection_url: &str,
251        records: &[RemoteSettingsRecord],
252    ) -> Result<u64> {
253        // Find the max last_modified time while inserting records
254        let mut max_last_modified = 0;
255        {
256            let mut insert_stmt = tx.prepare(
257                "INSERT OR REPLACE INTO records (id, collection_url, data) VALUES (?, ?, ?)",
258            )?;
259            let mut delete_stmt = tx.prepare("DELETE FROM records WHERE id=?")?;
260            for record in records {
261                if record.deleted {
262                    delete_stmt.execute(params![&record.id])?;
263                } else {
264                    max_last_modified = max_last_modified.max(record.last_modified);
265                    let data = serde_json::to_vec(&record)?;
266                    insert_stmt.execute(params![record.id, collection_url, data])?;
267                }
268            }
269        }
270        Ok(max_last_modified)
271    }
272
273    /// Update the collection metadata after setting/merging records
274    fn update_collection_metadata(
275        tx: &Transaction<'_>,
276        collection_url: &str,
277        last_modified: u64,
278        metadata: CollectionMetadata,
279    ) -> Result<()> {
280        let signatures_json = serde_json::to_string(&metadata.signatures)
281            .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?;
282
283        let mut stmt = tx.prepare(
284            "INSERT OR REPLACE INTO collection_metadata
285            (collection_url, last_modified, bucket, signatures)
286            VALUES (?, ?, ?, ?)",
287        )?;
288
289        stmt.execute((
290            collection_url,
291            last_modified,
292            &metadata.bucket,
293            &signatures_json,
294        ))?;
295        Ok(())
296    }
297
298    /// Set the attachment data stored in the database, clearing out any previously stored data
299    pub fn set_attachment(
300        &mut self,
301        collection_url: &str,
302        location: &str,
303        attachment: &[u8],
304    ) -> Result<()> {
305        let tx = self.transaction()?;
306
307        // Delete ALL existing attachments for every collection_url
308        tx.execute(
309            "DELETE FROM attachments WHERE collection_url != ?",
310            params![collection_url],
311        )?;
312
313        tx.execute(
314            "INSERT OR REPLACE INTO ATTACHMENTS \
315            (id, collection_url, data) \
316            VALUES (?, ?, ?)",
317            params![location, collection_url, attachment,],
318        )?;
319
320        tx.commit()?;
321
322        Ok(())
323    }
324
325    /// Empty out all cached values and start from scratch.  This is called when
326    /// RemoteSettingsService::update_config() is called, since that could change the remote
327    /// settings server which would invalidate all cached data.
328    pub fn empty(&mut self) -> Result<()> {
329        let tx = self.transaction()?;
330        tx.execute("DELETE FROM records", [])?;
331        tx.execute("DELETE FROM attachments", [])?;
332        tx.execute("DELETE FROM collection_metadata", [])?;
333        tx.commit()?;
334        Ok(())
335    }
336
337    /// Remove attachments that are no longer referenced by any current record.
338    /// When the server updates an attachment, the record gets a new UUID/hash in its location
339    /// field. Without this cleanup, the old attachment blob persists forever, causing unbounded
340    /// database growth (1+ GB observed in production for `quicksuggest-amp.sql`).
341    fn cleanup_orphaned_attachments(tx: &Transaction<'_>, collection_url: &str) -> Result<()> {
342        tx.execute(
343            "DELETE FROM attachments
344             WHERE collection_url = ?1
345             AND NOT EXISTS (
346                 SELECT 1 FROM records WHERE collection_url = ?1
347                 AND json_extract(data, '$.attachment.location') = attachments.id
348             )",
349            params![collection_url],
350        )?;
351        Ok(())
352    }
353
354    pub fn run_maintenance(&mut self) -> Result<()> {
355        if let ConnectionCell::Initialized(conn) = &self.conn {
356            run_maintenance(conn)?;
357        }
358
359        Ok(())
360    }
361}
362
363/// Stores the SQLite connection, which is lazily constructed and can be closed/shutdown.
364enum ConnectionCell {
365    Uninitialized,
366    Initialized(Connection),
367    Closed,
368}
369
370#[cfg(test)]
371mod tests {
372    use super::Storage;
373    use crate::{
374        client::CollectionMetadata, client::CollectionSignature, Attachment, RemoteSettingsRecord,
375        Result, RsJsonObject,
376    };
377    use sha2::{Digest, Sha256};
378
379    #[test]
380    fn test_storage_set_and_get_records() -> Result<()> {
381        let mut storage = Storage::new(":memory:".into());
382
383        let collection_url = "https://example.com/api";
384        let records = vec![
385            RemoteSettingsRecord {
386                id: "1".to_string(),
387                last_modified: 100,
388                deleted: false,
389                attachment: None,
390                fields: serde_json::json!({"key": "value1"})
391                    .as_object()
392                    .unwrap()
393                    .clone(),
394            },
395            RemoteSettingsRecord {
396                id: "2".to_string(),
397                last_modified: 200,
398                deleted: false,
399                attachment: None,
400                fields: serde_json::json!({"key": "value2"})
401                    .as_object()
402                    .unwrap()
403                    .clone(),
404            },
405        ];
406
407        // Set records
408        storage.insert_collection_content(
409            collection_url,
410            &records,
411            300,
412            CollectionMetadata::default(),
413        )?;
414
415        // Get records
416        let fetched_records = storage.get_records(collection_url)?;
417        assert!(fetched_records.is_some());
418        let fetched_records = fetched_records.unwrap();
419        assert_eq!(fetched_records.len(), 2);
420        assert_eq!(fetched_records, records);
421
422        assert_eq!(fetched_records[0].fields["key"], "value1");
423
424        // Get last modified timestamp
425        let last_modified = storage.get_last_modified_timestamp(collection_url)?;
426        assert_eq!(last_modified, Some(300));
427
428        Ok(())
429    }
430
431    #[test]
432    fn test_storage_get_records_none() -> Result<()> {
433        let mut storage = Storage::new(":memory:".into());
434
435        let collection_url = "https://example.com/api";
436
437        // Get records when none are set
438        let fetched_records = storage.get_records(collection_url)?;
439        assert!(fetched_records.is_none());
440
441        // Get last modified timestamp when no records
442        let last_modified = storage.get_last_modified_timestamp(collection_url)?;
443        assert!(last_modified.is_none());
444
445        Ok(())
446    }
447
448    #[test]
449    fn test_storage_get_records_empty() -> Result<()> {
450        let mut storage = Storage::new(":memory:".into());
451
452        let collection_url = "https://example.com/api";
453
454        // Set empty records
455        storage.insert_collection_content(
456            collection_url,
457            &Vec::<RemoteSettingsRecord>::default(),
458            42,
459            CollectionMetadata::default(),
460        )?;
461
462        // Get records
463        let fetched_records = storage.get_records(collection_url)?;
464        assert_eq!(fetched_records, Some(Vec::new()));
465
466        // Get last modified timestamp when no records
467        let last_modified = storage.get_last_modified_timestamp(collection_url)?;
468        assert_eq!(last_modified, Some(42));
469
470        Ok(())
471    }
472
473    #[test]
474    fn test_storage_set_and_get_attachment() -> Result<()> {
475        let mut storage = Storage::new(":memory:".into());
476
477        let attachment = &[0x18, 0x64];
478        let collection_url = "https://example.com/api";
479        let attachment_metadata = Attachment {
480            filename: "abc".to_string(),
481            mimetype: "application/json".to_string(),
482            location: "tmp".to_string(),
483            hash: format!("{:x}", Sha256::digest(attachment)),
484            size: attachment.len() as u64,
485        };
486
487        // Store attachment
488        storage.set_attachment(collection_url, &attachment_metadata.location, attachment)?;
489
490        // Get attachment
491        let fetched_attachment = storage.get_attachment(collection_url, attachment_metadata)?;
492        assert!(fetched_attachment.is_some());
493        let fetched_attachment = fetched_attachment.unwrap();
494        assert_eq!(fetched_attachment, attachment);
495
496        Ok(())
497    }
498
499    #[test]
500    fn test_storage_set_and_replace_attachment() -> Result<()> {
501        let mut storage = Storage::new(":memory:".into());
502
503        let collection_url = "https://example.com/api";
504
505        let attachment_1 = &[0x18, 0x64];
506        let attachment_2 = &[0x12, 0x48];
507
508        let attachment_metadata_1 = Attachment {
509            filename: "abc".to_string(),
510            mimetype: "application/json".to_string(),
511            location: "tmp".to_string(),
512            hash: format!("{:x}", Sha256::digest(attachment_1)),
513            size: attachment_1.len() as u64,
514        };
515
516        let attachment_metadata_2 = Attachment {
517            filename: "def".to_string(),
518            mimetype: "application/json".to_string(),
519            location: "tmp".to_string(),
520            hash: format!("{:x}", Sha256::digest(attachment_2)),
521            size: attachment_2.len() as u64,
522        };
523
524        // Store first attachment
525        storage.set_attachment(
526            collection_url,
527            &attachment_metadata_1.location,
528            attachment_1,
529        )?;
530
531        // Replace attachment with new data
532        storage.set_attachment(
533            collection_url,
534            &attachment_metadata_2.location,
535            attachment_2,
536        )?;
537
538        // Get attachment
539        let fetched_attachment = storage.get_attachment(collection_url, attachment_metadata_2)?;
540        assert!(fetched_attachment.is_some());
541        let fetched_attachment = fetched_attachment.unwrap();
542        assert_eq!(fetched_attachment, attachment_2);
543
544        Ok(())
545    }
546
547    #[test]
548    fn test_storage_set_attachment_delete_others() -> Result<()> {
549        let mut storage = Storage::new(":memory:".into());
550
551        let collection_url_1 = "https://example.com/api1";
552        let collection_url_2 = "https://example.com/api2";
553
554        let attachment_1 = &[0x18, 0x64];
555        let attachment_2 = &[0x12, 0x48];
556
557        let attachment_metadata_1 = Attachment {
558            filename: "abc".to_string(),
559            mimetype: "application/json".to_string(),
560            location: "first_tmp".to_string(),
561            hash: format!("{:x}", Sha256::digest(attachment_1)),
562            size: attachment_1.len() as u64,
563        };
564
565        let attachment_metadata_2 = Attachment {
566            filename: "def".to_string(),
567            mimetype: "application/json".to_string(),
568            location: "second_tmp".to_string(),
569            hash: format!("{:x}", Sha256::digest(attachment_2)),
570            size: attachment_2.len() as u64,
571        };
572
573        // Set attachments for two different collections
574        storage.set_attachment(
575            collection_url_1,
576            &attachment_metadata_1.location,
577            attachment_1,
578        )?;
579        storage.set_attachment(
580            collection_url_2,
581            &attachment_metadata_2.location,
582            attachment_2,
583        )?;
584
585        // Verify that only the attachment for the second collection remains
586        let fetched_attachment_1 =
587            storage.get_attachment(collection_url_1, attachment_metadata_1)?;
588        assert!(fetched_attachment_1.is_none());
589
590        let fetched_attachment_2 =
591            storage.get_attachment(collection_url_2, attachment_metadata_2)?;
592        assert!(fetched_attachment_2.is_some());
593        let fetched_attachment_2 = fetched_attachment_2.unwrap();
594        assert_eq!(fetched_attachment_2, attachment_2);
595
596        Ok(())
597    }
598
599    /// Test that orphaned attachments are cleaned up when a record's attachment location changes.
600    /// This reproduces the 1.1GB bloat observed in production. The `quicksuggest-amp` collection
601    /// has records like:
602    /// {
603    ///   "type": "amp",
604    ///   "country": "US",
605    ///   "form_factor": "phone",
606    ///   "filter_expression": "env.country == 'US' && env.formFactor == 'phone'",
607    ///   "id": "sponsored-suggestions-us-phone",
608    ///   "attachment": { "hash": "992ed42aa...", "size": 9795975, "filename": "sponsored-suggestions-us-phone.json", ... }
609    /// }
610    /// When the data refreshes (schema timestamp changes), the record keeps its ID but gets
611    /// a new attachment with a different hash/location. The old cached blob was never cleaned up.
612    #[test]
613    fn test_storage_orphaned_attachments_cleaned_up_on_update() -> Result<()> {
614        let mut storage = Storage::new(":memory:".into());
615        let collection_url = "https://example.com/api";
616
617        let attachment_v1 = b"version 1 data";
618        let attachment_v2 = b"version 2 data";
619
620        let attachment_meta_v1 = Attachment {
621            filename: "sponsored-suggestions-us-phone.json".to_string(),
622            mimetype: "application/json".to_string(),
623            location: "main-workspace/quicksuggest-amp/attachment-v1.json".to_string(),
624            hash: format!("{:x}", Sha256::digest(attachment_v1)),
625            size: attachment_v1.len() as u64,
626        };
627
628        let attachment_meta_v2 = Attachment {
629            filename: "sponsored-suggestions-us-phone.json".to_string(),
630            mimetype: "application/json".to_string(),
631            location: "main-workspace/quicksuggest-amp/attachment-v2.json".to_string(),
632            hash: format!("{:x}", Sha256::digest(attachment_v2)),
633            size: attachment_v2.len() as u64,
634        };
635
636        // Initially record points to attachment_v1
637        let records_v1 = vec![RemoteSettingsRecord {
638            id: "sponsored-suggestions-us-phone".to_string(),
639            last_modified: 100,
640            deleted: false,
641            attachment: Some(attachment_meta_v1.clone()),
642            fields: serde_json::json!({"type": "amp"})
643                .as_object()
644                .unwrap()
645                .clone(),
646        }];
647
648        storage.insert_collection_content(
649            collection_url,
650            &records_v1,
651            100,
652            CollectionMetadata::default(),
653        )?;
654        storage.set_attachment(collection_url, &attachment_meta_v1.location, attachment_v1)?;
655
656        let fetched = storage.get_attachment(collection_url, attachment_meta_v1.clone())?;
657        assert!(fetched.is_some(), "v1 attachment should be stored");
658
659        // Simulate a server data refresh while keeping the same record ID
660        // but a new attachment location
661        let records_v2 = vec![RemoteSettingsRecord {
662            id: "sponsored-suggestions-us-phone".to_string(),
663            last_modified: 200,
664            deleted: false,
665            attachment: Some(attachment_meta_v2.clone()),
666            fields: serde_json::json!({"type": "amp"})
667                .as_object()
668                .unwrap()
669                .clone(),
670        }];
671
672        storage.insert_collection_content(
673            collection_url,
674            &records_v2,
675            200,
676            CollectionMetadata::default(),
677        )?;
678        storage.set_attachment(collection_url, &attachment_meta_v2.location, attachment_v2)?;
679
680        let fetched_v2 = storage.get_attachment(collection_url, attachment_meta_v2)?;
681        assert!(fetched_v2.is_some(), "v2 attachment should be stored");
682
683        let fetched_v1 = storage.get_attachment(collection_url, attachment_meta_v1)?;
684        assert!(
685            fetched_v1.is_none(),
686            "v1 attachment should be cleaned up after record points to v2"
687        );
688
689        Ok(())
690    }
691
692    /// Test that orphaned attachments are cleaned up when a record is deleted via tombstone.
693    /// The `quicksuggest-amp` changeset uses filter_expression to target specific countries
694    /// and form factors. If the server removes a record (e.g. drops a country), a tombstone
695    /// is sent:
696    /// {
697    ///   "id": "sponsored-suggestions-gb-phone",
698    ///   "last_modified": 1774549156905,
699    ///   "deleted": true
700    /// }
701    /// The record row gets deleted, but the cached attachment blob was never cleaned up.
702    #[test]
703    fn test_storage_orphaned_attachments_cleaned_up_on_delete() -> Result<()> {
704        let mut storage = Storage::new(":memory:".into());
705        let collection_url = "https://example.com/api";
706
707        let attachment_data = b"sponsored suggestions for GB phone";
708
709        let attachment_meta = Attachment {
710            filename: "sponsored-suggestions-gb-phone.json".to_string(),
711            mimetype: "application/json".to_string(),
712            location: "main-workspace/quicksuggest-amp/attachment-gb.json".to_string(),
713            hash: format!("{:x}", Sha256::digest(attachment_data)),
714            size: attachment_data.len() as u64,
715        };
716
717        // Initially we have two records, one with an attachment
718        let initial_records = vec![
719            RemoteSettingsRecord {
720                id: "sponsored-suggestions-gb-phone".to_string(),
721                last_modified: 100,
722                deleted: false,
723                attachment: Some(attachment_meta.clone()),
724                fields: serde_json::json!({"type": "amp"})
725                    .as_object()
726                    .unwrap()
727                    .clone(),
728            },
729            RemoteSettingsRecord {
730                id: "sponsored-suggestions-us-phone".to_string(),
731                last_modified: 100,
732                deleted: false,
733                attachment: None,
734                fields: serde_json::json!({"type": "amp"})
735                    .as_object()
736                    .unwrap()
737                    .clone(),
738            },
739        ];
740
741        storage.insert_collection_content(
742            collection_url,
743            &initial_records,
744            100,
745            CollectionMetadata::default(),
746        )?;
747        storage.set_attachment(collection_url, &attachment_meta.location, attachment_data)?;
748
749        let fetched = storage.get_attachment(collection_url, attachment_meta.clone())?;
750        assert!(fetched.is_some(), "GB attachment should be stored");
751
752        // Simulate server sending a tombstone, GB record is deleted
753        let updated_records = vec![RemoteSettingsRecord {
754            id: "sponsored-suggestions-gb-phone".to_string(),
755            last_modified: 200,
756            deleted: true,
757            attachment: None,
758            fields: RsJsonObject::new(),
759        }];
760
761        storage.insert_collection_content(
762            collection_url,
763            &updated_records,
764            200,
765            CollectionMetadata::default(),
766        )?;
767
768        let fetched = storage.get_attachment(collection_url, attachment_meta)?;
769        assert!(
770            fetched.is_none(),
771            "GB attachment should be cleaned up after record is deleted via tombstone"
772        );
773
774        Ok(())
775    }
776
777    #[test]
778    fn test_storage_get_attachment_not_found() -> Result<()> {
779        let mut storage = Storage::new(":memory:".into());
780
781        let collection_url = "https://example.com/api";
782        let metadata = Attachment::default();
783
784        // Get attachment that doesn't exist
785        let fetched_attachment = storage.get_attachment(collection_url, metadata)?;
786        assert!(fetched_attachment.is_none());
787
788        Ok(())
789    }
790
791    #[test]
792    fn test_storage_empty() -> Result<()> {
793        let mut storage = Storage::new(":memory:".into());
794
795        let collection_url = "https://example.com/api";
796        let attachment = &[0x18, 0x64];
797
798        let records = vec![
799            RemoteSettingsRecord {
800                id: "1".to_string(),
801                last_modified: 100,
802                deleted: false,
803                attachment: None,
804                fields: serde_json::json!({"key": "value1"})
805                    .as_object()
806                    .unwrap()
807                    .clone(),
808            },
809            RemoteSettingsRecord {
810                id: "2".to_string(),
811                last_modified: 200,
812                deleted: false,
813                attachment: Some(Attachment {
814                    filename: "abc".to_string(),
815                    mimetype: "application/json".to_string(),
816                    location: "tmp".to_string(),
817                    hash: format!("{:x}", Sha256::digest(attachment)),
818                    size: attachment.len() as u64,
819                }),
820                fields: serde_json::json!({"key": "value2"})
821                    .as_object()
822                    .unwrap()
823                    .clone(),
824            },
825        ];
826
827        let metadata = records[1]
828            .clone()
829            .attachment
830            .expect("No attachment metadata for record");
831
832        // Set records and attachment
833        storage.insert_collection_content(
834            collection_url,
835            &records,
836            42,
837            CollectionMetadata::default(),
838        )?;
839        storage.set_attachment(collection_url, &metadata.location, attachment)?;
840
841        // Verify they are stored
842        let fetched_records = storage.get_records(collection_url)?;
843        assert!(fetched_records.is_some());
844        let fetched_attachment = storage.get_attachment(collection_url, metadata.clone())?;
845        assert!(fetched_attachment.is_some());
846
847        // Empty the storage
848        storage.empty()?;
849
850        // Verify they are deleted
851        let fetched_records = storage.get_records(collection_url)?;
852        assert!(fetched_records.is_none());
853        let fetched_attachment = storage.get_attachment(collection_url, metadata)?;
854        assert!(fetched_attachment.is_none());
855
856        Ok(())
857    }
858
859    #[test]
860    fn test_storage_collection_url_isolation() -> Result<()> {
861        let mut storage = Storage::new(":memory:".into());
862
863        let collection_url1 = "https://example.com/api1";
864        let collection_url2 = "https://example.com/api2";
865        let records_collection_url1 = vec![RemoteSettingsRecord {
866            id: "1".to_string(),
867            last_modified: 100,
868            deleted: false,
869            attachment: None,
870            fields: serde_json::json!({"key": "value1"})
871                .as_object()
872                .unwrap()
873                .clone(),
874        }];
875        let records_collection_url2 = vec![RemoteSettingsRecord {
876            id: "2".to_string(),
877            last_modified: 200,
878            deleted: false,
879            attachment: None,
880            fields: serde_json::json!({"key": "value2"})
881                .as_object()
882                .unwrap()
883                .clone(),
884        }];
885
886        // Set records for collection_url1
887        storage.insert_collection_content(
888            collection_url1,
889            &records_collection_url1,
890            42,
891            CollectionMetadata::default(),
892        )?;
893        // Verify records for collection_url1
894        let fetched_records = storage.get_records(collection_url1)?;
895        assert!(fetched_records.is_some());
896        let fetched_records = fetched_records.unwrap();
897        assert_eq!(fetched_records.len(), 1);
898        assert_eq!(fetched_records, records_collection_url1);
899
900        // Set records for collection_url2, which will clear records for all collections
901        storage.insert_collection_content(
902            collection_url2,
903            &records_collection_url2,
904            300,
905            CollectionMetadata::default(),
906        )?;
907
908        // Verify that records for collection_url1 have been cleared
909        let fetched_records = storage.get_records(collection_url1)?;
910        assert!(fetched_records.is_none());
911
912        // Verify records for collection_url2 are correctly stored
913        let fetched_records = storage.get_records(collection_url2)?;
914        assert!(fetched_records.is_some());
915        let fetched_records = fetched_records.unwrap();
916        assert_eq!(fetched_records.len(), 1);
917        assert_eq!(fetched_records, records_collection_url2);
918
919        // Verify last modified timestamps only for collection_url2
920        let last_modified1 = storage.get_last_modified_timestamp(collection_url1)?;
921        assert_eq!(last_modified1, None);
922        let last_modified2 = storage.get_last_modified_timestamp(collection_url2)?;
923        assert_eq!(last_modified2, Some(300));
924
925        Ok(())
926    }
927
928    #[test]
929    fn test_storage_insert_collection_content() -> Result<()> {
930        let mut storage = Storage::new(":memory:".into());
931
932        let collection_url = "https://example.com/api";
933        let initial_records = vec![RemoteSettingsRecord {
934            id: "2".to_string(),
935            last_modified: 200,
936            deleted: false,
937            attachment: None,
938            fields: serde_json::json!({"key": "value2"})
939                .as_object()
940                .unwrap()
941                .clone(),
942        }];
943
944        // Set initial records
945        storage.insert_collection_content(
946            collection_url,
947            &initial_records,
948            42,
949            CollectionMetadata::default(),
950        )?;
951
952        // Verify initial records
953        let fetched_records = storage.get_records(collection_url)?;
954        assert!(fetched_records.is_some());
955        assert_eq!(fetched_records.unwrap(), initial_records);
956
957        // Update records
958        let updated_records = vec![RemoteSettingsRecord {
959            id: "2".to_string(),
960            last_modified: 200,
961            deleted: false,
962            attachment: None,
963            fields: serde_json::json!({"key": "value2_updated"})
964                .as_object()
965                .unwrap()
966                .clone(),
967        }];
968        storage.insert_collection_content(
969            collection_url,
970            &updated_records,
971            300,
972            CollectionMetadata::default(),
973        )?;
974
975        // Verify updated records
976        let fetched_records = storage.get_records(collection_url)?;
977        assert!(fetched_records.is_some());
978        assert_eq!(fetched_records.unwrap(), updated_records);
979
980        // Verify last modified timestamp
981        let last_modified = storage.get_last_modified_timestamp(collection_url)?;
982        assert_eq!(last_modified, Some(300));
983
984        Ok(())
985    }
986
987    // Quick way to generate the fields data for our mock records
988    fn test_fields(data: &str) -> RsJsonObject {
989        let mut map = serde_json::Map::new();
990        map.insert("data".into(), data.into());
991        map
992    }
993
994    #[test]
995    fn test_storage_merge_records() -> Result<()> {
996        let mut storage = Storage::new(":memory:".into());
997
998        let collection_url = "https://example.com/api";
999
1000        let initial_records = vec![
1001            RemoteSettingsRecord {
1002                id: "a".into(),
1003                last_modified: 100,
1004                deleted: false,
1005                attachment: None,
1006                fields: test_fields("a"),
1007            },
1008            RemoteSettingsRecord {
1009                id: "b".into(),
1010                last_modified: 200,
1011                deleted: false,
1012                attachment: None,
1013                fields: test_fields("b"),
1014            },
1015            RemoteSettingsRecord {
1016                id: "c".into(),
1017                last_modified: 300,
1018                deleted: false,
1019                attachment: None,
1020                fields: test_fields("c"),
1021            },
1022        ];
1023        let updated_records = vec![
1024            // d is new
1025            RemoteSettingsRecord {
1026                id: "d".into(),
1027                last_modified: 1300,
1028                deleted: false,
1029                attachment: None,
1030                fields: test_fields("d"),
1031            },
1032            // b was deleted
1033            RemoteSettingsRecord {
1034                id: "b".into(),
1035                last_modified: 1200,
1036                deleted: true,
1037                attachment: None,
1038                fields: RsJsonObject::new(),
1039            },
1040            // a was updated
1041            RemoteSettingsRecord {
1042                id: "a".into(),
1043                last_modified: 1100,
1044                deleted: false,
1045                attachment: None,
1046                fields: test_fields("a-with-new-data"),
1047            },
1048            // c was not modified, so it's not present in the new response
1049        ];
1050        let expected_records = vec![
1051            // a was updated
1052            RemoteSettingsRecord {
1053                id: "a".into(),
1054                last_modified: 1100,
1055                deleted: false,
1056                attachment: None,
1057                fields: test_fields("a-with-new-data"),
1058            },
1059            RemoteSettingsRecord {
1060                id: "c".into(),
1061                last_modified: 300,
1062                deleted: false,
1063                attachment: None,
1064                fields: test_fields("c"),
1065            },
1066            RemoteSettingsRecord {
1067                id: "d".into(),
1068                last_modified: 1300,
1069                deleted: false,
1070                attachment: None,
1071                fields: test_fields("d"),
1072            },
1073        ];
1074
1075        // Set initial records
1076        storage.insert_collection_content(
1077            collection_url,
1078            &initial_records,
1079            1000,
1080            CollectionMetadata::default(),
1081        )?;
1082
1083        // Verify initial records
1084        let fetched_records = storage.get_records(collection_url)?.unwrap();
1085        assert_eq!(fetched_records, initial_records);
1086
1087        // Update records
1088        storage.insert_collection_content(
1089            collection_url,
1090            &updated_records,
1091            1300,
1092            CollectionMetadata::default(),
1093        )?;
1094
1095        // Verify updated records
1096        let mut fetched_records = storage.get_records(collection_url)?.unwrap();
1097        fetched_records.sort_by_cached_key(|r| r.id.clone());
1098        assert_eq!(fetched_records, expected_records);
1099
1100        // Verify last modified timestamp
1101        let last_modified = storage.get_last_modified_timestamp(collection_url)?;
1102        assert_eq!(last_modified, Some(1300));
1103        Ok(())
1104    }
1105    #[test]
1106    fn test_storage_get_collection_metadata() -> Result<()> {
1107        let mut storage = Storage::new(":memory:".into());
1108
1109        let collection_url = "https://example.com/api";
1110        let initial_records = vec![RemoteSettingsRecord {
1111            id: "2".to_string(),
1112            last_modified: 200,
1113            deleted: false,
1114            attachment: None,
1115            fields: serde_json::json!({"key": "value2"})
1116                .as_object()
1117                .unwrap()
1118                .clone(),
1119        }];
1120
1121        // Set initial records
1122        storage.insert_collection_content(
1123            collection_url,
1124            &initial_records,
1125            1337,
1126            CollectionMetadata {
1127                bucket: "main".into(),
1128                signatures: vec![
1129                    CollectionSignature {
1130                        signature: "b64encodedsig".into(),
1131                        x5u: "http://15u/".into(),
1132                        mode: "mldsa".into(),
1133                    },
1134                    CollectionSignature {
1135                        signature: "b64encodedsig2".into(),
1136                        x5u: "http://15u2/".into(),
1137                        mode: "p384ecdsa".into(),
1138                    },
1139                ],
1140            },
1141        )?;
1142
1143        let metadata = storage.get_collection_metadata(collection_url)?.unwrap();
1144
1145        assert_eq!(metadata.signatures[0].signature, "b64encodedsig");
1146        assert_eq!(metadata.signatures[0].x5u, "http://15u/");
1147        assert_eq!(metadata.signatures[0].mode, "mldsa");
1148        assert_eq!(metadata.signatures[1].signature, "b64encodedsig2");
1149        assert_eq!(metadata.signatures[1].x5u, "http://15u2/");
1150        assert_eq!(metadata.signatures[1].mode, "p384ecdsa");
1151
1152        Ok(())
1153    }
1154}
remote_settings/storage.rs

remote_settings/
storage.rs
