statsig_rust/
spec_types_encoded.rs

1use serde::Deserialize;
2
3use crate::{
4    compression::zstd_decompression_dict::DictionaryDecoder, spec_types::SpecsResponse, StatsigErr,
5};
6
7#[derive(Deserialize)]
8struct DictCompressedSpecsResponse {
9    #[serde(rename = "s")]
10    pub specs: Vec<u8>,
11    #[serde(rename = "d_id")]
12    pub dict_id: Option<String>,
13    #[serde(rename = "d")]
14    pub dict_buff: Option<Vec<u8>>,
15}
16
17impl DictCompressedSpecsResponse {
18    fn decompress(
19        self,
20        cached_dict: Option<&DictionaryDecoder>,
21    ) -> Result<DecodedSpecsResponse, StatsigErr> {
22        let decompression_dict_to_use =
23            select_decompression_dict_for_response(self.dict_id, self.dict_buff, cached_dict)?;
24
25        match decompression_dict_to_use {
26            None => {
27                // Response was not compressed
28                let parsed = serde_json::from_slice::<SpecsResponse>(&self.specs).map_err(|e| {
29                    StatsigErr::JsonParseError("SpecsResponse".to_string(), e.to_string())
30                })?;
31                Ok(DecodedSpecsResponse {
32                    specs: parsed,
33                    decompression_dict: None,
34                })
35            }
36            Some(dict) => {
37                // Response was compressed, so we need to decompress first then parse
38                let uncompressed = dict.decompress(&self.specs)?;
39                let parsed =
40                    serde_json::from_slice::<SpecsResponse>(&uncompressed).map_err(|e| {
41                        StatsigErr::JsonParseError("SpecsResponse".to_string(), e.to_string())
42                    })?;
43                Ok(DecodedSpecsResponse {
44                    specs: parsed,
45                    decompression_dict: Some(dict),
46                })
47            }
48        }
49    }
50}
51
52fn select_decompression_dict_for_response(
53    response_dict_id: Option<String>,
54    response_dict_buff: Option<Vec<u8>>,
55    cached_dict: Option<&DictionaryDecoder>,
56) -> Result<Option<DictionaryDecoder>, StatsigErr> {
57    match response_dict_id {
58        None => Ok(None),
59        Some(response_dict_id) => {
60            if let Some(cached_dict) = cached_dict.filter(|d| d.get_dict_id() == response_dict_id) {
61                return Ok(Some(cached_dict.clone()));
62            }
63
64            response_dict_buff
65                .map(|dict_buff| DictionaryDecoder::new(None, response_dict_id.clone(), &dict_buff))
66                .map(|dict| Ok(Some(dict)))
67                .unwrap_or_else(|| {
68                    Err(StatsigErr::ZstdDictCompressionError(format!(
69                        "Cannot decompress response compressed with dict_id: {}, \
70                                because the appropriate dictionary is not cached \
71                                and the response does not contain one.",
72                        response_dict_id
73                    )))
74                })
75        }
76    }
77}
78
79#[derive(Deserialize)]
80#[serde(untagged)]
81enum CompressedSpecsResponse {
82    DictCompressed(DictCompressedSpecsResponse),
83    Uncompressed(SpecsResponse),
84}
85
86pub struct DecodedSpecsResponse {
87    pub specs: SpecsResponse,
88    pub decompression_dict: Option<DictionaryDecoder>,
89}
90
91impl DecodedSpecsResponse {
92    pub fn from_str(
93        response_str: &str,
94        decompression_dict: Option<&DictionaryDecoder>,
95    ) -> Result<DecodedSpecsResponse, StatsigErr> {
96        serde_json::from_str::<CompressedSpecsResponse>(response_str)
97            .map_err(|e| {
98                StatsigErr::JsonParseError("CompressedSpecsResponse".to_string(), e.to_string())
99            })
100            .and_then(|compressed| Self::from_compressed(compressed, decompression_dict))
101    }
102
103    fn from_compressed(
104        compressed: CompressedSpecsResponse,
105        decompression_dict: Option<&DictionaryDecoder>,
106    ) -> Result<DecodedSpecsResponse, StatsigErr> {
107        match compressed {
108            CompressedSpecsResponse::DictCompressed(compressed_response) => {
109                eprintln!("Parsing dict compressed specs");
110                compressed_response.decompress(decompression_dict)
111            }
112            CompressedSpecsResponse::Uncompressed(specs) => {
113                eprintln!("Parsing uncompressed specs");
114                Ok(DecodedSpecsResponse {
115                    specs,
116                    decompression_dict: decompression_dict.cloned(),
117                })
118            }
119        }
120    }
121}