1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151

//! Contains the structs and serialization implementations for encrypted vectors returned by the
//! indexer.

use ore_rs::{scheme::bit2::OREAES128, CipherText};
use serde::{ser::SerializeSeq, Serialize};
use serde_bytes::ByteBuf;

/// A single encrypted term - which can be one or more ORE cipher texts.
#[derive(Debug, PartialEq, PartialOrd)]
pub struct EncryptedTerm(pub Vec<CipherText<OREAES128, 8>>);

impl From<CipherText<OREAES128, 8>> for EncryptedTerm {
    fn from(term: CipherText<OREAES128, 8>) -> Self {
        EncryptedTerm(vec![term])
    }
}

impl EncryptedTerm {
    pub fn to_vec(self) -> Vec<Vec<u8>> {
        let mut output = vec![];

        for term in self.0.into_iter() {
            output.push(term.to_bytes());
        }

        output
    }
}

impl Serialize for EncryptedTerm {
    /// The ORE CipherText has a special `to_bytes` method for serialization.
    /// This custom serde implementation allows the EncryptedTerm type to be stored
    /// as an array of these `to_bytes` calls.
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        let mut seq = serializer.serialize_seq(Some(self.0.len()))?;

        for term in self.0.iter() {
            seq.serialize_element(&ByteBuf::from(term.to_bytes()))?;
        }

        seq.end()
    }
}

/// A single term stored inside a [`TermVector`].
///
/// It links back to the record it was generated from.
#[derive(Debug, PartialEq, PartialOrd, Serialize)]
pub struct VectorTerm {
    pub term: EncryptedTerm,
    #[serde(with = "super::fixed_byte_buf")]
    pub link: [u8; 16],
}

/// A collection of ORE encrypted terms for a specific document and index.
///
/// This struct matches the "Document::Vector" GRPC type.
#[derive(Serialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct TermVector {
    pub terms: Vec<VectorTerm>,
    #[serde(with = "super::fixed_byte_buf")]
    pub index_id: [u8; 16],
}

#[derive(Debug)]
pub enum ConstraintCondition {
    Exact {
        term: EncryptedTerm,
    },
    Range {
        upper: EncryptedTerm,
        lower: EncryptedTerm,
    },
}

#[derive(Debug)]
pub struct Constraint {
    pub index_id: [u8; 16],
    pub condition: ConstraintCondition,
}

impl Constraint {
    #[cfg(test)]
    pub fn unwrap_exact_term(self) -> EncryptedTerm {
        match self.condition {
            ConstraintCondition::Exact { term } => term,
            _ => panic!("Expected constraint to have exact condition"),
        }
    }

    #[cfg(test)]
    pub fn unwrap_range_terms(self) -> (EncryptedTerm, EncryptedTerm) {
        match self.condition {
            ConstraintCondition::Range { lower, upper, .. } => (lower, upper),
            _ => panic!("Expected constraint to have exact condition"),
        }
    }
}

#[cfg(test)]
mod tests {
    use ore_rs::{scheme::bit2::OREAES128, CipherText, ORECipher};

    use super::{EncryptedTerm, TermVector, VectorTerm};

    #[test]
    fn test_serialize_as_cbor() {
        let ore = OREAES128::init([0; 16], [1; 16], &[2; 8]).expect("Failed to init ORE");

        let ct_bytes = ore
            .encrypt(&[1; 8])
            .expect("Failed to encrypt buffer")
            .to_bytes();
        let ct: CipherText<OREAES128, 8> = CipherText::from_bytes(&ct_bytes).unwrap();

        let index_id = [3; 16];
        let record_id = [4; 16];

        let vector = TermVector {
            terms: vec![VectorTerm {
                term: EncryptedTerm(vec![ct]),
                link: record_id,
            }],
            index_id,
        };

        let mut cbor_buffer = vec![];
        ciborium::ser::into_writer(&vector, &mut cbor_buffer).expect("Failed to encode as cbor");

        let mut decoder = minicbor::Decoder::new(&cbor_buffer[..]);

        // Enter TermVector
        assert_eq!(decoder.map().unwrap(), Some(2));
        assert_eq!(decoder.str().unwrap(), "terms");
        assert_eq!(decoder.array().unwrap(), Some(1));
        // Enter DocumentTerm
        assert_eq!(decoder.map().unwrap(), Some(2));
        assert_eq!(decoder.str().unwrap(), "term");
        // Enter EncryptedTerm
        assert_eq!(decoder.array().unwrap(), Some(1));
        assert_eq!(decoder.bytes().unwrap(), &ct_bytes);
        // Leave EncryptedTerm
        assert_eq!(decoder.str().unwrap(), "link");
        assert_eq!(decoder.bytes().unwrap(), &record_id);
        // Leave DocumentTerm
        assert_eq!(decoder.str().unwrap(), "indexId");
        assert_eq!(decoder.bytes().unwrap(), &index_id);
        // Finished!
    }
}