Allow specifying top-level file and format for KVStore (#365)

* 0;9uAllow specifying a top-level file and format fields for KVStore. This change allows keeping KVStore key/values in a separate file for easier sharing between environments, and allowing users to only gitignore the actual KVStore content without needing to ignore the whole fastly.toml file.0;9u * Apply suggestions from code review Rename object_stores to kv_stores Co-authored-by: Kevin P. Fleming <[email protected]> * changed author block, broadened error assertion on file not found errors, and renamed Object store to KVStore one place * CHANGELOG entry --------- Co-authored-by: Kevin P. Fleming <[email protected]> Co-authored-by: Kevin P. Fleming <[email protected]>
fastly · Jul 26, 2024 · 0596f51 · 0596f51
1 parent 7aedece
commit 0596f51
Show file tree

Hide file tree

Showing 5 changed files with 216 additions and 5 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,7 @@
 ## Unreleased
 
+- Add support for JSON files in local_server.kv_stores ([#365](https://github.com/fastly/Viceroy/pull/365))
+
 ## 0.10.2 (2024-07-22)
 
 - Add support for supplying client certificates in fastly.toml, through the use of the

diff --git a/cli/tests/integration/kv_store.rs b/cli/tests/integration/kv_store.rs
@@ -61,6 +61,36 @@ viceroy_test!(object_stores_backward_compat, |is_component| {
     Ok(())
 });
 
+viceroy_test!(kv_store_allows_fetching_of_key_from_file, |is_component| {
+    // This test checks that we can provide a "format" and a "file"
+    // with a JSON dictionary inside it for a KV Store
+    // and have the KV Store populated with it.
+    const FASTLY_TOML: &str = r#"
+        name = "kv-store-test"
+        description = "kv store test"
+        authors = ["Gustav Wengel <[email protected]>"]
+        language = "rust"
+        [local_server]
+        kv_stores.empty_store = []
+        kv_stores.store_one = { file = "../test-fixtures/data/json-kv_store.json", format = "json" }
+    "#;
+
+    let resp = Test::using_fixture("kv_store.wasm")
+        .adapt_component(is_component)
+        .using_fastly_toml(FASTLY_TOML)?
+        .against_empty()
+        .await?;
+
+    assert_eq!(resp.status(), StatusCode::OK);
+    assert!(to_bytes(resp.into_body())
+        .await
+        .expect("can read body")
+        .to_vec()
+        .is_empty());
+
+    Ok(())
+});
+
 viceroy_test!(object_store_backward_compat, |is_component| {
     // Previously the "object_stores" key was named "object_store" and
     // the "file" key was named "path".  This test ensures that both
@@ -232,6 +262,83 @@ viceroy_test!(kv_store_bad_configs, |is_component| {
       _ => panic!(),
     }
 
+    const BAD_10_FASTLY_TOML: &str = r#"
+        name = "kv-store-test"
+        description = "kv store test"
+        authors = ["Gustav Wengel <[email protected]>"]
+        language = "rust"
+        [local_server]
+        kv_stores.empty_store = []
+        kv_stores.store_one = { file = "../test-fixtures/data/json-kv_store.json" }
+    "#;
+    match Test::using_fixture("kv_store.wasm").using_fastly_toml(BAD_10_FASTLY_TOML) {
+        Err(e) => assert_eq!("invalid configuration for 'store_one': When using a top-level 'file' to load data, both 'file' and 'format' must be set.", &e.to_string()),
+        _ => panic!(),
+    }
+
+    const BAD_11_FASTLY_TOML: &str = r#"
+        name = "kv-store-test"
+        description = "kv store test"
+        authors = ["Gustav Wengel <[email protected]>"]
+        language = "rust"
+        [local_server]
+        kv_stores.empty_store = []
+        kv_stores.store_one = { format = "json" }
+    "#;
+    match Test::using_fixture("kv_store.wasm").using_fastly_toml(BAD_11_FASTLY_TOML) {
+        Err(e) => assert_eq!("invalid configuration for 'store_one': When using a top-level 'file' to load data, both 'file' and 'format' must be set.", &e.to_string()),
+        _ => panic!(),
+    }
+
+    const BAD_12_FASTLY_TOML: &str = r#"
+        name = "kv-store-test"
+        description = "kv store test"
+        authors = ["Gustav Wengel <[email protected]>"]
+        language = "rust"
+        [local_server]
+        kv_stores.empty_store = []
+        kv_stores.store_one = { file = "../test-fixtures/data/json-kv_store.json", format = "INVALID" }
+    "#;
+    match Test::using_fixture("kv_store.wasm").using_fastly_toml(BAD_12_FASTLY_TOML) {
+        Err(e) => assert_eq!("invalid configuration for 'store_one': 'INVALID' is not a valid format for the config store. Supported format(s) are: 'json'.", &e.to_string()),
+        _ => panic!(),
+    }
+
+    const BAD_13_FASTLY_TOML: &str = r#"
+        name = "kv-store-test"
+        description = "kv store test"
+        authors = ["Gustav Wengel <[email protected]>"]
+        language = "rust"
+        [local_server]
+        kv_stores.empty_store = []
+        kv_stores.store_one = { file = "../test-fixtures/data/ABSOLUTELY_NOT_A_REAL_PATH", format = "json" }
+    "#;
+    match Test::using_fixture("kv_store.wasm").using_fastly_toml(BAD_13_FASTLY_TOML) {
+        Err(e) => {
+            // We can't assert on the whole error message here as the next part of the string is platform-specific
+            // where it says that it cannot find the file.
+            assert!(e
+                .to_string()
+                .contains("invalid configuration for 'store_one'"));
+        }
+        _ => panic!(),
+    }
+
+    // Not a valid JSON file
+    const BAD_14_FASTLY_TOML: &str = r#"
+        name = "kv-store-test"
+        description = "kv store test"
+        authors = ["Gustav Wengel <[email protected]>"]
+        language = "rust"
+        [local_server]
+        kv_stores.empty_store = []
+        kv_stores.store_one = { file = "../test-fixtures/data/kv-store.txt", format = "json" }
+    "#;
+    match Test::using_fixture("kv_store.wasm").using_fastly_toml(BAD_14_FASTLY_TOML) {
+        Err(e) => assert_eq!("invalid configuration for 'store_one': The file is of the wrong format. The file is expected to contain a single JSON Object.", &e.to_string()),
+        _ => panic!(),
+    }
+
     Ok(())
 });
 

diff --git a/lib/src/config/object_store.rs b/lib/src/config/object_store.rs
@@ -1,3 +1,6 @@
+use std::collections::HashMap;
+use std::path::{Path, PathBuf};
+use toml::{toml, Value};
 use {
     crate::{
         error::{FastlyConfigError, ObjectStoreConfigError},
@@ -15,12 +18,69 @@ impl TryFrom<Table> for ObjectStoreConfig {
     fn try_from(toml: Table) -> Result<Self, Self::Error> {
         let obj_store = ObjectStores::new();
         for (store, items) in toml.iter() {
-            let items = items.as_array().ok_or_else(|| {
-                FastlyConfigError::InvalidObjectStoreDefinition {
-                    name: store.to_string(),
-                    err: ObjectStoreConfigError::NotAnArray,
+            // Either the items here is from a top-level file with "file" and "format" keys
+            // or it's an inline array.
+            // We try to parse either one of them to the same Vec<toml::Value>
+            // to allow them to run through the same validation path further down
+            let file_path = items
+                .as_table()
+                .and_then(|table| table.get("file"))
+                .and_then(|file| file.as_str());
+            let file_format = items
+                .as_table()
+                .and_then(|table| table.get("format"))
+                .and_then(|format| format.as_str());
+
+            let items: Vec<toml::Value> = match (file_path, file_format) {
+                (Some(file_path), Some(file_type)) => {
+                    if file_type != "json" {
+                        return Err(FastlyConfigError::InvalidObjectStoreDefinition {
+                            name: store.to_string(),
+                            err: ObjectStoreConfigError::InvalidFileFormat(file_type.to_string()),
+                        });
+                    }
+
+                    let path = PathBuf::from(&file_path);
+
+                    let json = read_json_contents(&path).map_err(|err| {
+                        FastlyConfigError::InvalidObjectStoreDefinition {
+                            name: store.to_string(),
+                            err,
+                        }
+                    })?;
+
+                    let toml: Vec<Value> = json
+                        .into_iter()
+                        .map(|(key, value)| {
+                            toml! {
+                                key = key
+                                data = value
+                            }
+                        })
+                        .collect();
+
+                    toml
                 }
-            })?;
+                (None, None) => {
+                    // No file or format specified, parse the TOML as an array
+                    items
+                        .as_array()
+                        .ok_or_else(|| FastlyConfigError::InvalidObjectStoreDefinition {
+                            name: store.to_string(),
+                            err: ObjectStoreConfigError::NotAnArray,
+                        })?
+                        .clone()
+                }
+                // This means that *either* `format` or `file` is set, which isn't allowed
+                // we need both or neither.
+                (_, _) => {
+                    return Err(FastlyConfigError::InvalidObjectStoreDefinition {
+                        name: store.to_string(),
+                        err: ObjectStoreConfigError::OnlyOneFormatOrFileSet,
+                    });
+                }
+            };
+
             // Handle the case where there are no items to insert, but the store
             // exists and needs to be in the ObjectStore
             if items.is_empty() {
@@ -111,6 +171,34 @@ impl TryFrom<Table> for ObjectStoreConfig {
                     .expect("Lock was not poisoned");
             }
         }
+
         Ok(ObjectStoreConfig(obj_store))
     }
 }
+
+fn read_json_contents(file: &Path) -> Result<HashMap<String, String>, ObjectStoreConfigError> {
+    // Read the contents of the given file.
+    let data = fs::read_to_string(file).map_err(ObjectStoreConfigError::IoError)?;
+
+    // Deserialize the contents of the given JSON file.
+    let json =
+        match serde_json::from_str(&data).map_err(|_| ObjectStoreConfigError::FileWrongFormat)? {
+            // Check that we were given an object.
+            serde_json::Value::Object(obj) => obj,
+            _ => {
+                return Err(ObjectStoreConfigError::FileWrongFormat);
+            }
+        };
+
+    // Check that each dictionary entry has a string value.
+    let mut contents = HashMap::with_capacity(json.len());
+    for (key, value) in json {
+        let value = value
+            .as_str()
+            .ok_or_else(|| ObjectStoreConfigError::FileValueWrongFormat { key: key.clone() })?
+            .to_owned();
+        contents.insert(key, value);
+    }
+
+    Ok(contents)
+}
diff --git a/lib/src/error.rs b/lib/src/error.rs
@@ -647,6 +647,16 @@ pub enum ObjectStoreConfigError {
     ObjectStoreError(#[from] crate::object_store::ObjectStoreError),
     #[error("Invalid `key` value used: {0}.")]
     KeyValidationError(#[from] crate::object_store::KeyValidationError),
+    #[error("'{0}' is not a valid format for the config store. Supported format(s) are: 'json'.")]
+    InvalidFileFormat(String),
+    #[error("When using a top-level 'file' to load data, both 'file' and 'format' must be set.")]
+    OnlyOneFormatOrFileSet,
+    #[error(
+        "The file is of the wrong format. The file is expected to contain a single JSON Object."
+    )]
+    FileWrongFormat,
+    #[error("Item value under key named '{key}' is of the wrong format. The value is expected to be a JSON String.")]
+    FileValueWrongFormat { key: String },
 }
 
 /// Errors that may occur while validating secret store configurations.

diff --git a/test-fixtures/data/json-kv_store.json b/test-fixtures/data/json-kv_store.json
@@ -0,0 +1,4 @@
+{
+    "first": "This is some data",
+    "second": "More data"
+}