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
86#[derive(Clone)]
87pub struct DecodedSpecsResponse {
88    pub specs: SpecsResponse,
89    pub decompression_dict: Option<DictionaryDecoder>,
90}
91
92impl DecodedSpecsResponse {
93    pub fn from_slice(
94        slice: &[u8],
95        decompression_dict: Option<&DictionaryDecoder>,
96    ) -> Result<DecodedSpecsResponse, StatsigErr> {
97        serde_json::from_slice::<CompressedSpecsResponse>(slice)
98            .map_err(|e| {
99                StatsigErr::JsonParseError("CompressedSpecsResponse".to_string(), e.to_string())
100            })
101            .and_then(|compressed| Self::from_compressed(compressed, decompression_dict))
102    }
103
104    fn from_compressed(
105        compressed: CompressedSpecsResponse,
106        decompression_dict: Option<&DictionaryDecoder>,
107    ) -> Result<DecodedSpecsResponse, StatsigErr> {
108        match compressed {
109            CompressedSpecsResponse::DictCompressed(compressed_response) => {
110                eprintln!("Parsing dict compressed specs");
111                compressed_response.decompress(decompression_dict)
112            }
113            CompressedSpecsResponse::Uncompressed(specs) => {
114                eprintln!("Parsing uncompressed specs");
115                Ok(DecodedSpecsResponse {
116                    specs,
117                    decompression_dict: decompression_dict.cloned(),
118                })
119            }
120        }
121    }
122}