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, 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
349/// Stores the SQLite connection, which is lazily constructed and can be closed/shutdown.
350enum ConnectionCell {
351    Uninitialized,
352    Initialized(Connection),
353    Closed,
354}
355
356#[cfg(test)]
357mod tests {
358    use super::Storage;
359    use crate::{
360        client::CollectionMetadata, client::CollectionSignature, Attachment, RemoteSettingsRecord,
361        Result, RsJsonObject,
362    };
363    use sha2::{Digest, Sha256};
364
365    #[test]
366    fn test_storage_set_and_get_records() -> Result<()> {
367        let mut storage = Storage::new(":memory:".into());
368
369        let collection_url = "https://example.com/api";
370        let records = vec![
371            RemoteSettingsRecord {
372                id: "1".to_string(),
373                last_modified: 100,
374                deleted: false,
375                attachment: None,
376                fields: serde_json::json!({"key": "value1"})
377                    .as_object()
378                    .unwrap()
379                    .clone(),
380            },
381            RemoteSettingsRecord {
382                id: "2".to_string(),
383                last_modified: 200,
384                deleted: false,
385                attachment: None,
386                fields: serde_json::json!({"key": "value2"})
387                    .as_object()
388                    .unwrap()
389                    .clone(),
390            },
391        ];
392
393        // Set records
394        storage.insert_collection_content(
395            collection_url,
396            &records,
397            300,
398            CollectionMetadata::default(),
399        )?;
400
401        // Get records
402        let fetched_records = storage.get_records(collection_url)?;
403        assert!(fetched_records.is_some());
404        let fetched_records = fetched_records.unwrap();
405        assert_eq!(fetched_records.len(), 2);
406        assert_eq!(fetched_records, records);
407
408        assert_eq!(fetched_records[0].fields["key"], "value1");
409
410        // Get last modified timestamp
411        let last_modified = storage.get_last_modified_timestamp(collection_url)?;
412        assert_eq!(last_modified, Some(300));
413
414        Ok(())
415    }
416
417    #[test]
418    fn test_storage_get_records_none() -> Result<()> {
419        let mut storage = Storage::new(":memory:".into());
420
421        let collection_url = "https://example.com/api";
422
423        // Get records when none are set
424        let fetched_records = storage.get_records(collection_url)?;
425        assert!(fetched_records.is_none());
426
427        // Get last modified timestamp when no records
428        let last_modified = storage.get_last_modified_timestamp(collection_url)?;
429        assert!(last_modified.is_none());
430
431        Ok(())
432    }
433
434    #[test]
435    fn test_storage_get_records_empty() -> Result<()> {
436        let mut storage = Storage::new(":memory:".into());
437
438        let collection_url = "https://example.com/api";
439
440        // Set empty records
441        storage.insert_collection_content(
442            collection_url,
443            &Vec::<RemoteSettingsRecord>::default(),
444            42,
445            CollectionMetadata::default(),
446        )?;
447
448        // Get records
449        let fetched_records = storage.get_records(collection_url)?;
450        assert_eq!(fetched_records, Some(Vec::new()));
451
452        // Get last modified timestamp when no records
453        let last_modified = storage.get_last_modified_timestamp(collection_url)?;
454        assert_eq!(last_modified, Some(42));
455
456        Ok(())
457    }
458
459    #[test]
460    fn test_storage_set_and_get_attachment() -> Result<()> {
461        let mut storage = Storage::new(":memory:".into());
462
463        let attachment = &[0x18, 0x64];
464        let collection_url = "https://example.com/api";
465        let attachment_metadata = Attachment {
466            filename: "abc".to_string(),
467            mimetype: "application/json".to_string(),
468            location: "tmp".to_string(),
469            hash: format!("{:x}", Sha256::digest(attachment)),
470            size: attachment.len() as u64,
471        };
472
473        // Store attachment
474        storage.set_attachment(collection_url, &attachment_metadata.location, attachment)?;
475
476        // Get attachment
477        let fetched_attachment = storage.get_attachment(collection_url, attachment_metadata)?;
478        assert!(fetched_attachment.is_some());
479        let fetched_attachment = fetched_attachment.unwrap();
480        assert_eq!(fetched_attachment, attachment);
481
482        Ok(())
483    }
484
485    #[test]
486    fn test_storage_set_and_replace_attachment() -> Result<()> {
487        let mut storage = Storage::new(":memory:".into());
488
489        let collection_url = "https://example.com/api";
490
491        let attachment_1 = &[0x18, 0x64];
492        let attachment_2 = &[0x12, 0x48];
493
494        let attachment_metadata_1 = Attachment {
495            filename: "abc".to_string(),
496            mimetype: "application/json".to_string(),
497            location: "tmp".to_string(),
498            hash: format!("{:x}", Sha256::digest(attachment_1)),
499            size: attachment_1.len() as u64,
500        };
501
502        let attachment_metadata_2 = Attachment {
503            filename: "def".to_string(),
504            mimetype: "application/json".to_string(),
505            location: "tmp".to_string(),
506            hash: format!("{:x}", Sha256::digest(attachment_2)),
507            size: attachment_2.len() as u64,
508        };
509
510        // Store first attachment
511        storage.set_attachment(
512            collection_url,
513            &attachment_metadata_1.location,
514            attachment_1,
515        )?;
516
517        // Replace attachment with new data
518        storage.set_attachment(
519            collection_url,
520            &attachment_metadata_2.location,
521            attachment_2,
522        )?;
523
524        // Get attachment
525        let fetched_attachment = storage.get_attachment(collection_url, attachment_metadata_2)?;
526        assert!(fetched_attachment.is_some());
527        let fetched_attachment = fetched_attachment.unwrap();
528        assert_eq!(fetched_attachment, attachment_2);
529
530        Ok(())
531    }
532
533    #[test]
534    fn test_storage_set_attachment_delete_others() -> Result<()> {
535        let mut storage = Storage::new(":memory:".into());
536
537        let collection_url_1 = "https://example.com/api1";
538        let collection_url_2 = "https://example.com/api2";
539
540        let attachment_1 = &[0x18, 0x64];
541        let attachment_2 = &[0x12, 0x48];
542
543        let attachment_metadata_1 = Attachment {
544            filename: "abc".to_string(),
545            mimetype: "application/json".to_string(),
546            location: "first_tmp".to_string(),
547            hash: format!("{:x}", Sha256::digest(attachment_1)),
548            size: attachment_1.len() as u64,
549        };
550
551        let attachment_metadata_2 = Attachment {
552            filename: "def".to_string(),
553            mimetype: "application/json".to_string(),
554            location: "second_tmp".to_string(),
555            hash: format!("{:x}", Sha256::digest(attachment_2)),
556            size: attachment_2.len() as u64,
557        };
558
559        // Set attachments for two different collections
560        storage.set_attachment(
561            collection_url_1,
562            &attachment_metadata_1.location,
563            attachment_1,
564        )?;
565        storage.set_attachment(
566            collection_url_2,
567            &attachment_metadata_2.location,
568            attachment_2,
569        )?;
570
571        // Verify that only the attachment for the second collection remains
572        let fetched_attachment_1 =
573            storage.get_attachment(collection_url_1, attachment_metadata_1)?;
574        assert!(fetched_attachment_1.is_none());
575
576        let fetched_attachment_2 =
577            storage.get_attachment(collection_url_2, attachment_metadata_2)?;
578        assert!(fetched_attachment_2.is_some());
579        let fetched_attachment_2 = fetched_attachment_2.unwrap();
580        assert_eq!(fetched_attachment_2, attachment_2);
581
582        Ok(())
583    }
584
585    /// Test that orphaned attachments are cleaned up when a record's attachment location changes.
586    /// This reproduces the 1.1GB bloat observed in production. The `quicksuggest-amp` collection
587    /// has records like:
588    /// {
589    ///   "type": "amp",
590    ///   "country": "US",
591    ///   "form_factor": "phone",
592    ///   "filter_expression": "env.country == 'US' && env.formFactor == 'phone'",
593    ///   "id": "sponsored-suggestions-us-phone",
594    ///   "attachment": { "hash": "992ed42aa...", "size": 9795975, "filename": "sponsored-suggestions-us-phone.json", ... }
595    /// }
596    /// When the data refreshes (schema timestamp changes), the record keeps its ID but gets
597    /// a new attachment with a different hash/location. The old cached blob was never cleaned up.
598    #[test]
599    fn test_storage_orphaned_attachments_cleaned_up_on_update() -> Result<()> {
600        let mut storage = Storage::new(":memory:".into());
601        let collection_url = "https://example.com/api";
602
603        let attachment_v1 = b"version 1 data";
604        let attachment_v2 = b"version 2 data";
605
606        let attachment_meta_v1 = Attachment {
607            filename: "sponsored-suggestions-us-phone.json".to_string(),
608            mimetype: "application/json".to_string(),
609            location: "main-workspace/quicksuggest-amp/attachment-v1.json".to_string(),
610            hash: format!("{:x}", Sha256::digest(attachment_v1)),
611            size: attachment_v1.len() as u64,
612        };
613
614        let attachment_meta_v2 = Attachment {
615            filename: "sponsored-suggestions-us-phone.json".to_string(),
616            mimetype: "application/json".to_string(),
617            location: "main-workspace/quicksuggest-amp/attachment-v2.json".to_string(),
618            hash: format!("{:x}", Sha256::digest(attachment_v2)),
619            size: attachment_v2.len() as u64,
620        };
621
622        // Initially record points to attachment_v1
623        let records_v1 = vec![RemoteSettingsRecord {
624            id: "sponsored-suggestions-us-phone".to_string(),
625            last_modified: 100,
626            deleted: false,
627            attachment: Some(attachment_meta_v1.clone()),
628            fields: serde_json::json!({"type": "amp"})
629                .as_object()
630                .unwrap()
631                .clone(),
632        }];
633
634        storage.insert_collection_content(
635            collection_url,
636            &records_v1,
637            100,
638            CollectionMetadata::default(),
639        )?;
640        storage.set_attachment(collection_url, &attachment_meta_v1.location, attachment_v1)?;
641
642        let fetched = storage.get_attachment(collection_url, attachment_meta_v1.clone())?;
643        assert!(fetched.is_some(), "v1 attachment should be stored");
644
645        // Simulate a server data refresh while keeping the same record ID
646        // but a new attachment location
647        let records_v2 = vec![RemoteSettingsRecord {
648            id: "sponsored-suggestions-us-phone".to_string(),
649            last_modified: 200,
650            deleted: false,
651            attachment: Some(attachment_meta_v2.clone()),
652            fields: serde_json::json!({"type": "amp"})
653                .as_object()
654                .unwrap()
655                .clone(),
656        }];
657
658        storage.insert_collection_content(
659            collection_url,
660            &records_v2,
661            200,
662            CollectionMetadata::default(),
663        )?;
664        storage.set_attachment(collection_url, &attachment_meta_v2.location, attachment_v2)?;
665
666        let fetched_v2 = storage.get_attachment(collection_url, attachment_meta_v2)?;
667        assert!(fetched_v2.is_some(), "v2 attachment should be stored");
668
669        let fetched_v1 = storage.get_attachment(collection_url, attachment_meta_v1)?;
670        assert!(
671            fetched_v1.is_none(),
672            "v1 attachment should be cleaned up after record points to v2"
673        );
674
675        Ok(())
676    }
677
678    /// Test that orphaned attachments are cleaned up when a record is deleted via tombstone.
679    /// The `quicksuggest-amp` changeset uses filter_expression to target specific countries
680    /// and form factors. If the server removes a record (e.g. drops a country), a tombstone
681    /// is sent:
682    /// {
683    ///   "id": "sponsored-suggestions-gb-phone",
684    ///   "last_modified": 1774549156905,
685    ///   "deleted": true
686    /// }
687    /// The record row gets deleted, but the cached attachment blob was never cleaned up.
688    #[test]
689    fn test_storage_orphaned_attachments_cleaned_up_on_delete() -> Result<()> {
690        let mut storage = Storage::new(":memory:".into());
691        let collection_url = "https://example.com/api";
692
693        let attachment_data = b"sponsored suggestions for GB phone";
694
695        let attachment_meta = Attachment {
696            filename: "sponsored-suggestions-gb-phone.json".to_string(),
697            mimetype: "application/json".to_string(),
698            location: "main-workspace/quicksuggest-amp/attachment-gb.json".to_string(),
699            hash: format!("{:x}", Sha256::digest(attachment_data)),
700            size: attachment_data.len() as u64,
701        };
702
703        // Initially we have two records, one with an attachment
704        let initial_records = vec![
705            RemoteSettingsRecord {
706                id: "sponsored-suggestions-gb-phone".to_string(),
707                last_modified: 100,
708                deleted: false,
709                attachment: Some(attachment_meta.clone()),
710                fields: serde_json::json!({"type": "amp"})
711                    .as_object()
712                    .unwrap()
713                    .clone(),
714            },
715            RemoteSettingsRecord {
716                id: "sponsored-suggestions-us-phone".to_string(),
717                last_modified: 100,
718                deleted: false,
719                attachment: None,
720                fields: serde_json::json!({"type": "amp"})
721                    .as_object()
722                    .unwrap()
723                    .clone(),
724            },
725        ];
726
727        storage.insert_collection_content(
728            collection_url,
729            &initial_records,
730            100,
731            CollectionMetadata::default(),
732        )?;
733        storage.set_attachment(collection_url, &attachment_meta.location, attachment_data)?;
734
735        let fetched = storage.get_attachment(collection_url, attachment_meta.clone())?;
736        assert!(fetched.is_some(), "GB attachment should be stored");
737
738        // Simulate server sending a tombstone, GB record is deleted
739        let updated_records = vec![RemoteSettingsRecord {
740            id: "sponsored-suggestions-gb-phone".to_string(),
741            last_modified: 200,
742            deleted: true,
743            attachment: None,
744            fields: RsJsonObject::new(),
745        }];
746
747        storage.insert_collection_content(
748            collection_url,
749            &updated_records,
750            200,
751            CollectionMetadata::default(),
752        )?;
753
754        let fetched = storage.get_attachment(collection_url, attachment_meta)?;
755        assert!(
756            fetched.is_none(),
757            "GB attachment should be cleaned up after record is deleted via tombstone"
758        );
759
760        Ok(())
761    }
762
763    #[test]
764    fn test_storage_get_attachment_not_found() -> Result<()> {
765        let mut storage = Storage::new(":memory:".into());
766
767        let collection_url = "https://example.com/api";
768        let metadata = Attachment::default();
769
770        // Get attachment that doesn't exist
771        let fetched_attachment = storage.get_attachment(collection_url, metadata)?;
772        assert!(fetched_attachment.is_none());
773
774        Ok(())
775    }
776
777    #[test]
778    fn test_storage_empty() -> Result<()> {
779        let mut storage = Storage::new(":memory:".into());
780
781        let collection_url = "https://example.com/api";
782        let attachment = &[0x18, 0x64];
783
784        let records = vec![
785            RemoteSettingsRecord {
786                id: "1".to_string(),
787                last_modified: 100,
788                deleted: false,
789                attachment: None,
790                fields: serde_json::json!({"key": "value1"})
791                    .as_object()
792                    .unwrap()
793                    .clone(),
794            },
795            RemoteSettingsRecord {
796                id: "2".to_string(),
797                last_modified: 200,
798                deleted: false,
799                attachment: Some(Attachment {
800                    filename: "abc".to_string(),
801                    mimetype: "application/json".to_string(),
802                    location: "tmp".to_string(),
803                    hash: format!("{:x}", Sha256::digest(attachment)),
804                    size: attachment.len() as u64,
805                }),
806                fields: serde_json::json!({"key": "value2"})
807                    .as_object()
808                    .unwrap()
809                    .clone(),
810            },
811        ];
812
813        let metadata = records[1]
814            .clone()
815            .attachment
816            .expect("No attachment metadata for record");
817
818        // Set records and attachment
819        storage.insert_collection_content(
820            collection_url,
821            &records,
822            42,
823            CollectionMetadata::default(),
824        )?;
825        storage.set_attachment(collection_url, &metadata.location, attachment)?;
826
827        // Verify they are stored
828        let fetched_records = storage.get_records(collection_url)?;
829        assert!(fetched_records.is_some());
830        let fetched_attachment = storage.get_attachment(collection_url, metadata.clone())?;
831        assert!(fetched_attachment.is_some());
832
833        // Empty the storage
834        storage.empty()?;
835
836        // Verify they are deleted
837        let fetched_records = storage.get_records(collection_url)?;
838        assert!(fetched_records.is_none());
839        let fetched_attachment = storage.get_attachment(collection_url, metadata)?;
840        assert!(fetched_attachment.is_none());
841
842        Ok(())
843    }
844
845    #[test]
846    fn test_storage_collection_url_isolation() -> Result<()> {
847        let mut storage = Storage::new(":memory:".into());
848
849        let collection_url1 = "https://example.com/api1";
850        let collection_url2 = "https://example.com/api2";
851        let records_collection_url1 = vec![RemoteSettingsRecord {
852            id: "1".to_string(),
853            last_modified: 100,
854            deleted: false,
855            attachment: None,
856            fields: serde_json::json!({"key": "value1"})
857                .as_object()
858                .unwrap()
859                .clone(),
860        }];
861        let records_collection_url2 = vec![RemoteSettingsRecord {
862            id: "2".to_string(),
863            last_modified: 200,
864            deleted: false,
865            attachment: None,
866            fields: serde_json::json!({"key": "value2"})
867                .as_object()
868                .unwrap()
869                .clone(),
870        }];
871
872        // Set records for collection_url1
873        storage.insert_collection_content(
874            collection_url1,
875            &records_collection_url1,
876            42,
877            CollectionMetadata::default(),
878        )?;
879        // Verify records for collection_url1
880        let fetched_records = storage.get_records(collection_url1)?;
881        assert!(fetched_records.is_some());
882        let fetched_records = fetched_records.unwrap();
883        assert_eq!(fetched_records.len(), 1);
884        assert_eq!(fetched_records, records_collection_url1);
885
886        // Set records for collection_url2, which will clear records for all collections
887        storage.insert_collection_content(
888            collection_url2,
889            &records_collection_url2,
890            300,
891            CollectionMetadata::default(),
892        )?;
893
894        // Verify that records for collection_url1 have been cleared
895        let fetched_records = storage.get_records(collection_url1)?;
896        assert!(fetched_records.is_none());
897
898        // Verify records for collection_url2 are correctly stored
899        let fetched_records = storage.get_records(collection_url2)?;
900        assert!(fetched_records.is_some());
901        let fetched_records = fetched_records.unwrap();
902        assert_eq!(fetched_records.len(), 1);
903        assert_eq!(fetched_records, records_collection_url2);
904
905        // Verify last modified timestamps only for collection_url2
906        let last_modified1 = storage.get_last_modified_timestamp(collection_url1)?;
907        assert_eq!(last_modified1, None);
908        let last_modified2 = storage.get_last_modified_timestamp(collection_url2)?;
909        assert_eq!(last_modified2, Some(300));
910
911        Ok(())
912    }
913
914    #[test]
915    fn test_storage_insert_collection_content() -> Result<()> {
916        let mut storage = Storage::new(":memory:".into());
917
918        let collection_url = "https://example.com/api";
919        let initial_records = vec![RemoteSettingsRecord {
920            id: "2".to_string(),
921            last_modified: 200,
922            deleted: false,
923            attachment: None,
924            fields: serde_json::json!({"key": "value2"})
925                .as_object()
926                .unwrap()
927                .clone(),
928        }];
929
930        // Set initial records
931        storage.insert_collection_content(
932            collection_url,
933            &initial_records,
934            42,
935            CollectionMetadata::default(),
936        )?;
937
938        // Verify initial records
939        let fetched_records = storage.get_records(collection_url)?;
940        assert!(fetched_records.is_some());
941        assert_eq!(fetched_records.unwrap(), initial_records);
942
943        // Update records
944        let updated_records = vec![RemoteSettingsRecord {
945            id: "2".to_string(),
946            last_modified: 200,
947            deleted: false,
948            attachment: None,
949            fields: serde_json::json!({"key": "value2_updated"})
950                .as_object()
951                .unwrap()
952                .clone(),
953        }];
954        storage.insert_collection_content(
955            collection_url,
956            &updated_records,
957            300,
958            CollectionMetadata::default(),
959        )?;
960
961        // Verify updated records
962        let fetched_records = storage.get_records(collection_url)?;
963        assert!(fetched_records.is_some());
964        assert_eq!(fetched_records.unwrap(), updated_records);
965
966        // Verify last modified timestamp
967        let last_modified = storage.get_last_modified_timestamp(collection_url)?;
968        assert_eq!(last_modified, Some(300));
969
970        Ok(())
971    }
972
973    // Quick way to generate the fields data for our mock records
974    fn test_fields(data: &str) -> RsJsonObject {
975        let mut map = serde_json::Map::new();
976        map.insert("data".into(), data.into());
977        map
978    }
979
980    #[test]
981    fn test_storage_merge_records() -> Result<()> {
982        let mut storage = Storage::new(":memory:".into());
983
984        let collection_url = "https://example.com/api";
985
986        let initial_records = vec![
987            RemoteSettingsRecord {
988                id: "a".into(),
989                last_modified: 100,
990                deleted: false,
991                attachment: None,
992                fields: test_fields("a"),
993            },
994            RemoteSettingsRecord {
995                id: "b".into(),
996                last_modified: 200,
997                deleted: false,
998                attachment: None,
999                fields: test_fields("b"),
1000            },
1001            RemoteSettingsRecord {
1002                id: "c".into(),
1003                last_modified: 300,
1004                deleted: false,
1005                attachment: None,
1006                fields: test_fields("c"),
1007            },
1008        ];
1009        let updated_records = vec![
1010            // d is new
1011            RemoteSettingsRecord {
1012                id: "d".into(),
1013                last_modified: 1300,
1014                deleted: false,
1015                attachment: None,
1016                fields: test_fields("d"),
1017            },
1018            // b was deleted
1019            RemoteSettingsRecord {
1020                id: "b".into(),
1021                last_modified: 1200,
1022                deleted: true,
1023                attachment: None,
1024                fields: RsJsonObject::new(),
1025            },
1026            // a was updated
1027            RemoteSettingsRecord {
1028                id: "a".into(),
1029                last_modified: 1100,
1030                deleted: false,
1031                attachment: None,
1032                fields: test_fields("a-with-new-data"),
1033            },
1034            // c was not modified, so it's not present in the new response
1035        ];
1036        let expected_records = vec![
1037            // a was updated
1038            RemoteSettingsRecord {
1039                id: "a".into(),
1040                last_modified: 1100,
1041                deleted: false,
1042                attachment: None,
1043                fields: test_fields("a-with-new-data"),
1044            },
1045            RemoteSettingsRecord {
1046                id: "c".into(),
1047                last_modified: 300,
1048                deleted: false,
1049                attachment: None,
1050                fields: test_fields("c"),
1051            },
1052            RemoteSettingsRecord {
1053                id: "d".into(),
1054                last_modified: 1300,
1055                deleted: false,
1056                attachment: None,
1057                fields: test_fields("d"),
1058            },
1059        ];
1060
1061        // Set initial records
1062        storage.insert_collection_content(
1063            collection_url,
1064            &initial_records,
1065            1000,
1066            CollectionMetadata::default(),
1067        )?;
1068
1069        // Verify initial records
1070        let fetched_records = storage.get_records(collection_url)?.unwrap();
1071        assert_eq!(fetched_records, initial_records);
1072
1073        // Update records
1074        storage.insert_collection_content(
1075            collection_url,
1076            &updated_records,
1077            1300,
1078            CollectionMetadata::default(),
1079        )?;
1080
1081        // Verify updated records
1082        let mut fetched_records = storage.get_records(collection_url)?.unwrap();
1083        fetched_records.sort_by_cached_key(|r| r.id.clone());
1084        assert_eq!(fetched_records, expected_records);
1085
1086        // Verify last modified timestamp
1087        let last_modified = storage.get_last_modified_timestamp(collection_url)?;
1088        assert_eq!(last_modified, Some(1300));
1089        Ok(())
1090    }
1091    #[test]
1092    fn test_storage_get_collection_metadata() -> Result<()> {
1093        let mut storage = Storage::new(":memory:".into());
1094
1095        let collection_url = "https://example.com/api";
1096        let initial_records = vec![RemoteSettingsRecord {
1097            id: "2".to_string(),
1098            last_modified: 200,
1099            deleted: false,
1100            attachment: None,
1101            fields: serde_json::json!({"key": "value2"})
1102                .as_object()
1103                .unwrap()
1104                .clone(),
1105        }];
1106
1107        // Set initial records
1108        storage.insert_collection_content(
1109            collection_url,
1110            &initial_records,
1111            1337,
1112            CollectionMetadata {
1113                bucket: "main".into(),
1114                signatures: vec![
1115                    CollectionSignature {
1116                        signature: "b64encodedsig".into(),
1117                        x5u: "http://15u/".into(),
1118                    },
1119                    CollectionSignature {
1120                        signature: "b64encodedsig2".into(),
1121                        x5u: "http://15u2/".into(),
1122                    },
1123                ],
1124            },
1125        )?;
1126
1127        let metadata = storage.get_collection_metadata(collection_url)?.unwrap();
1128
1129        assert_eq!(metadata.signatures[0].signature, "b64encodedsig");
1130        assert_eq!(metadata.signatures[0].x5u, "http://15u/");
1131        assert_eq!(metadata.signatures[1].signature, "b64encodedsig2");
1132        assert_eq!(metadata.signatures[1].x5u, "http://15u2/");
1133
1134        Ok(())
1135    }
1136}
remote_settings/storage.rs

remote_settings/
storage.rs
