cipherstash-client 0.12.5

use std::{
    iter::{once, Once},
    slice::Iter,
};

use blake3::Hasher;
use serde_json::{Number, Value};

use crate::encryption::{EncryptionError, IndexTerm};

/// Indexer that is designed to support JSONB containment queries that use one or more of the following operators: `@>`, `<@`
///
/// # Does left contain right?
///
/// `cs_ste_vec_v1(jsonb) @> cs_ste_vec_v1(jsonb) → boolean`
///
/// # Does right contain left?
///
/// `cs_ste_vec_v1(jsonb) <@ cs_ste_vec_v1(jsonb) → boolean`
///
pub struct JsonContainmentIndexer {
    key: [u8; 32],
    prefix: Vec<u8>,
}

impl JsonContainmentIndexer {
    pub fn new(key: [u8; 32], prefix: String) -> Self {
        Self {
            key,
            prefix: prefix.as_bytes().into(),
        }
    }

    /// Produces a `Vec` of keyed Blake3 hashes: one per non-null leaf node of the JSON value.
    ///
    /// The input to the keyed hash is `self.prefix`, path of leaf node, value of leaf node.  All parts are converted to
    /// `[u8]` before being joined with a `0u8` (the NULL byte) and hashed.
    pub fn encrypt(&self, json: Value) -> Result<IndexTerm, EncryptionError> {
        let mut hashes: Vec<blake3::Hash> = Vec::new();

        for PathValue(path, leaf) in PathHashTargets::new(&json) {
            let mut hasher = Hasher::new_keyed(&self.key);

            hasher.update(&self.prefix);

            // `path[1..]`` skips the first entry which is the same for all `PathValue`s `PathSegment::Root`.
            for segment in path[1..].iter() {
                segment.update_hash(&mut hasher);
                hashes.push(hasher.finalize());
            }

            leaf.update_hash(&mut hasher);
            hashes.push(hasher.finalize());
        }

        Ok(IndexTerm::BinaryVec(
            hashes
                .into_iter()
                .map(|hash| Vec::from(hash.as_bytes()))
                .collect(),
        ))
    }
}

/// A node that will be hashed.
///
/// `Array` and `Object` are empty variants. They are more like a marker in order to support prefix queries where the
/// deepest node in the query is an array or object. If only *leaf* nodes were indexed, then querying for the presence
/// of an object or array would not be possible.
///
/// This replicates the behaviour of Postgres containment queries.
#[derive(Debug, PartialEq, Eq, Hash)]
enum HashTarget<'a> {
    Bool(bool),
    Number(&'a Number),
    String(&'a str),
    Array,
    Object,
}

impl<'a> HashTarget<'a> {
    fn update_hash(&self, hasher: &mut Hasher) {
        match self {
            HashTarget::Bool(x) => hasher.update(format!("Bool-{x}").as_bytes()),
            HashTarget::Number(x) => hasher.update(format!("Number-{x}").as_bytes()),
            HashTarget::String(x) => hasher.update(format!("String-{x}").as_bytes()),
            HashTarget::Array => hasher.update("Array".as_bytes()),
            HashTarget::Object => hasher.update("Object".as_bytes()),
        };
    }
}

/// A path and the `Leaf` value at that path.
#[derive(Debug, PartialEq, Eq, Hash)]
struct PathValue<'a>(Vec<PathSegment<'a>>, HashTarget<'a>);

/// An iterator over the content of a single JSON node. Not recursive.
enum JsonIter<'a> {
    Bool(PathSegment<'a>, Once<&'a Value>),
    Number(PathSegment<'a>, Once<&'a Value>),
    String(PathSegment<'a>, Once<&'a Value>),
    Array(PathSegment<'a>, Iter<'a, Value>),
    Object(PathSegment<'a>, serde_json::map::Iter<'a>),
}

impl<'a> JsonIter<'a> {
    fn idx(&self) -> PathSegment<'a> {
        match self {
            JsonIter::Bool(idx, _) => idx.clone(),
            JsonIter::Number(idx, _) => idx.clone(),
            JsonIter::String(idx, _) => idx.clone(),
            JsonIter::Array(idx, _) => idx.clone(),
            JsonIter::Object(idx, _) => idx.clone(),
        }
    }
}

/// An [`Iterator`] implementation that yields [`PathValue`] values from a JSON AST.
///
/// It works by maintaining a stack of [`JsonIter`] iterators. The top of the stack represents the current node.
struct PathHashTargets<'a> {
    stack: Vec<JsonIter<'a>>,
}

impl<'a> PathHashTargets<'a> {
    fn new(json: &'a Value) -> Self {
        Self {
            stack: match json {
                Value::Null => vec![],
                Value::Bool(_) => vec![JsonIter::Bool(PathSegment::Root, once(json))],
                Value::Number(_) => vec![JsonIter::Number(PathSegment::Root, once(json))],
                Value::String(_) => vec![JsonIter::String(PathSegment::Root, once(json))],
                Value::Array(array) => {
                    vec![JsonIter::Array(PathSegment::Root, array.iter())]
                }
                Value::Object(object) => {
                    vec![JsonIter::Object(PathSegment::Root, object.into_iter())]
                }
            },
        }
    }

    fn path(&self, segment: PathSegment<'a>) -> Vec<PathSegment<'a>> {
        Vec::from_iter(
            self.stack
                .iter()
                .map(|node| node.idx())
                .chain(once(segment)),
        )
    }
}

/// An index into a [`Map`] or `Vec` of [`Value`].
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
enum PathSegment<'a> {
    /// The root node a JSON is not contained within a map or an array
    Root,

    /// An item in an array. The indexes are not used in contains queries and are discarded.
    ArrayItem,

    /// A property of a JSON object. The property is used to generate hashes used in containment queries so it is
    /// retained.
    ObjectItem(&'a str),
}

impl<'a> PathSegment<'a> {
    fn update_hash(&self, hasher: &mut Hasher) {
        match self {
            PathSegment::Root => hasher.update("Root".as_bytes()),
            PathSegment::ArrayItem => hasher.update("ArrayItem".as_bytes()),
            PathSegment::ObjectItem(key) => hasher.update(format!("ObjectItem-{}", key).as_bytes()),
        };
    }
}

impl<'a> Iterator for JsonIter<'a> {
    type Item = (PathSegment<'a>, &'a Value);

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            JsonIter::Number(path_segment, iter) => iter.next().map(|v| (path_segment.clone(), v)),
            JsonIter::String(path_segment, iter) => iter.next().map(|v| (path_segment.clone(), v)),
            JsonIter::Bool(path_segment, iter) => iter.next().map(|v| (path_segment.clone(), v)),
            JsonIter::Array(_, iter) => iter.next().map(|v| (PathSegment::ArrayItem, v)),
            JsonIter::Object(_, iter) => iter
                .next()
                .map(|(key, v)| (PathSegment::ObjectItem(key), v)),
        }
    }
}

impl<'a> Iterator for PathHashTargets<'a> {
    type Item = PathValue<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.stack.last_mut() {
                Some(value_iter) => match value_iter.next() {
                    Some((segment, value)) => match value {
                        Value::Null => {}
                        Value::Bool(x) => {
                            return Some(PathValue(self.path(segment), HashTarget::Bool(*x)))
                        }
                        Value::Number(x) => {
                            return Some(PathValue(self.path(segment), HashTarget::Number(x)))
                        }
                        Value::String(x) => {
                            return Some(PathValue(self.path(segment), HashTarget::String(x)))
                        }
                        Value::Array(array) => {
                            let arr_value =
                                PathValue(self.path(segment.clone()), HashTarget::Array);
                            self.stack.push(JsonIter::Array(segment, array.iter()));

                            return Some(arr_value);
                        }
                        Value::Object(map) => {
                            let object_value =
                                PathValue(self.path(segment.clone()), HashTarget::Object);
                            self.stack.push(JsonIter::Object(segment, map.into_iter()));

                            return Some(object_value);
                        }
                    },

                    None => {
                        self.stack.pop();
                    }
                },
                None => return None,
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use super::*;
    use serde_json::Number;

    #[test]
    fn test_flattening() {
        let json = serde_json::json!({
            "name": "John",
            "age": 30,
            "cars": [
                "Ford",
                "BMW",
                "Fiat",
                { "Tesla": "Model S" }
            ]
        });

        let flattened: HashSet<PathValue> = HashSet::from_iter(PathHashTargets::new(&json));

        assert!(flattened.contains(&PathValue(
            vec![PathSegment::Root, PathSegment::ObjectItem("name")],
            HashTarget::String("John")
        )));
        assert!(flattened.contains(&PathValue(
            vec![PathSegment::Root, PathSegment::ObjectItem("age")],
            HashTarget::Number(&Number::from(30))
        )));
        assert!(flattened.contains(&PathValue(
            vec![
                PathSegment::Root,
                PathSegment::ObjectItem("cars"),
                PathSegment::ArrayItem,
            ],
            HashTarget::String("Ford")
        )));
        assert!(flattened.contains(&PathValue(
            vec![
                PathSegment::Root,
                PathSegment::ObjectItem("cars"),
                PathSegment::ArrayItem,
            ],
            HashTarget::String("BMW")
        )));
        assert!(flattened.contains(&PathValue(
            vec![
                PathSegment::Root,
                PathSegment::ObjectItem("cars"),
                PathSegment::ArrayItem,
            ],
            HashTarget::String("Fiat")
        )));
        assert!(flattened.contains(&PathValue(
            vec![
                PathSegment::Root,
                PathSegment::ObjectItem("cars"),
                PathSegment::ArrayItem,
                PathSegment::ObjectItem("Tesla"),
            ],
            HashTarget::String("Model S")
        )));
    }

    #[test]
    fn test_encrypt() {
        let indexer = JsonContainmentIndexer::new([0; 32], "cars".into());
        let json = serde_json::json!({
            "name": "John",
            "null": null,
            "age": 30,
            "cars": [
                "Ford",
                "BMW",
                null,
                "Fiat"
            ]
        });

        if let IndexTerm::BinaryVec(terms) = indexer.encrypt(json).unwrap() {
            assert_eq!(terms.len(), 15);
        } else {
            panic!();
        }
    }

    #[test]
    fn find_by_leaf_node() {
        let json = serde_json::json!({
            "name": "John",
            "age": 30,
            "cars": [
                "Ford",
                "BMW",
                "Fiat",
                { "Tesla": "Model S" }
            ]
        });

        let indexer = JsonContainmentIndexer::new([0; 32], "cars".into());

        if let IndexTerm::BinaryVec(terms) = indexer.encrypt(json).unwrap() {
            let Ok(IndexTerm::BinaryVec(search_terms)) = indexer.encrypt(serde_json::json!({
                "name": "John"
            })) else {
                panic!();
            };

            assert!(search_terms
                .iter()
                .all(|term| terms.iter().any(|candidate| term == candidate)));

            let Ok(IndexTerm::BinaryVec(search_terms)) = indexer.encrypt(serde_json::json!({
                "cars": [
                    { "Tesla": "Model S" }
                ]
            })) else {
                panic!();
            };

            assert!(search_terms
                .iter()
                .all(|term| terms.iter().any(|candidate| term == candidate)));
        } else {
            panic!();
        }
    }

    #[test]
    fn find_by_non_existing_leaf_node() {
        let json = serde_json::json!({
            "name": "John",
            "age": 30,
            "cars": [
                "Ford",
                "BMW",
                "Fiat",
                { "Tesla": "Model S" }
            ]
        });

        let indexer = JsonContainmentIndexer::new([0; 32], "cars".into());

        if let IndexTerm::BinaryVec(terms) = indexer.encrypt(json).unwrap() {
            // This one should fail to be found because the search terms does not include the
            // the path through the JSON.
            let Ok(IndexTerm::BinaryVec(search_terms)) =
                indexer.encrypt(serde_json::json!(["Ford"]))
            else {
                panic!();
            };

            assert!(!search_terms
                .iter()
                .all(|term| terms.iter().any(|candidate| term == candidate)));
        } else {
            panic!();
        }
    }

    #[test]
    fn find_by_non_leaf_node() {
        let json = serde_json::json!({
            "name": "John",
            "age": 30,
            "cars": [
                "Ford",
                "BMW",
                "Fiat",
                { "Tesla": "Model S" }
            ]
        });

        let indexer = JsonContainmentIndexer::new([0; 32], "cars".into());

        if let IndexTerm::BinaryVec(terms) = indexer.encrypt(json).unwrap() {
            let Ok(IndexTerm::BinaryVec(search_terms)) = indexer.encrypt(serde_json::json!({
                "cars": []
            })) else {
                panic!();
            };

            assert_eq!(search_terms.len(), 2);

            assert!(search_terms
                .iter()
                .all(|term| terms.iter().any(|candidate| term == candidate)));
        } else {
            panic!();
        }
    }

    #[test]
    fn find_by_non_existing_non_leaf_node() {
        let json = serde_json::json!({
            "name": "John",
            "age": 30,
            "cars": [
                "Ford",
                "BMW",
                "Fiat",
                { "Tesla": "Model S" }
            ]
        });

        let indexer = JsonContainmentIndexer::new([0; 32], "cars".into());

        if let IndexTerm::BinaryVec(terms) = indexer.encrypt(json).unwrap() {
            let Ok(IndexTerm::BinaryVec(search_terms)) = indexer.encrypt(serde_json::json!({
                // Cars is a JSON array (not an object) so we shouldn't be able to find a result
                // with this query.
                "cars": {}
            })) else {
                panic!();
            };

            assert_eq!(search_terms.len(), 2);

            assert!(!search_terms
                .iter()
                .all(|term| terms.iter().any(|candidate| term == candidate)));
        } else {
            panic!();
        }
    }
}