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

remote_settings/
storage.rs
