viceroy_lib/config/
geolocation.rs

1use {
2    crate::error::GeolocationConfigError,
3    serde_json::{
4        Map, Number, Value as SerdeValue, Value::Number as SerdeNumber,
5        Value::String as SerdeString,
6    },
7    std::{collections::HashMap, fs, iter::FromIterator, net::IpAddr, path::Path, path::PathBuf},
8};
9
10#[derive(Clone, Debug)]
11pub struct Geolocation {
12    mapping: GeolocationMapping,
13    use_default_loopback: bool,
14}
15
16#[derive(Clone, Debug)]
17pub enum GeolocationMapping {
18    Empty,
19    InlineToml {
20        addresses: HashMap<IpAddr, GeolocationData>,
21    },
22    Json {
23        file: PathBuf,
24    },
25}
26
27#[derive(Clone, Debug)]
28pub struct GeolocationData {
29    data: Map<String, SerdeValue>,
30}
31
32impl Default for Geolocation {
33    fn default() -> Self {
34        Self {
35            mapping: GeolocationMapping::default(),
36            use_default_loopback: true,
37        }
38    }
39}
40
41impl Geolocation {
42    pub fn new() -> Self {
43        Self::default()
44    }
45
46    pub fn lookup(&self, addr: &IpAddr) -> Option<GeolocationData> {
47        self.mapping.get(addr).or_else(|| {
48            if self.use_default_loopback && addr.is_loopback() {
49                Some(GeolocationData::default())
50            } else {
51                None
52            }
53        })
54    }
55}
56mod deserialization {
57    use std::{net::IpAddr, str::FromStr};
58
59    use serde_json::Number;
60
61    use {
62        super::{Geolocation, GeolocationData, GeolocationMapping},
63        crate::error::{FastlyConfigError, GeolocationConfigError},
64        serde_json::Value as SerdeValue,
65        std::path::PathBuf,
66        std::{collections::HashMap, convert::TryFrom},
67        toml::value::{Table, Value},
68    };
69
70    impl TryFrom<Table> for Geolocation {
71        type Error = FastlyConfigError;
72
73        fn try_from(toml: Table) -> Result<Self, Self::Error> {
74            fn process_config(mut toml: Table) -> Result<Geolocation, GeolocationConfigError> {
75                let use_default_loopback = toml.remove("use_default_loopback").map_or(
76                    Ok(true),
77                    |use_default_loopback| match use_default_loopback {
78                        Value::Boolean(use_default_loopback) => Ok(use_default_loopback),
79                        _ => Err(GeolocationConfigError::InvalidEntryType),
80                    },
81                )?;
82
83                let mapping = match toml.remove("format") {
84                    Some(Value::String(value)) => match value.as_str() {
85                        "inline-toml" => process_inline_toml_dictionary(&mut toml)?,
86                        "json" => process_json_entries(&mut toml)?,
87                        "" => return Err(GeolocationConfigError::EmptyFormatEntry),
88                        format => {
89                            return Err(GeolocationConfigError::InvalidGeolocationMappingFormat(
90                                format.to_string(),
91                            ))
92                        }
93                    },
94                    Some(_) => return Err(GeolocationConfigError::InvalidFormatEntry),
95                    None => GeolocationMapping::Empty,
96                };
97
98                Ok(Geolocation {
99                    mapping,
100                    use_default_loopback,
101                })
102            }
103
104            process_config(toml).map_err(|err| FastlyConfigError::InvalidGeolocationDefinition {
105                name: "geolocation_mapping".to_string(),
106                err,
107            })
108        }
109    }
110
111    pub fn parse_ip_address(address: &str) -> Result<IpAddr, GeolocationConfigError> {
112        IpAddr::from_str(address)
113            .map_err(|err| GeolocationConfigError::InvalidAddressEntry(err.to_string()))
114    }
115
116    fn process_inline_toml_dictionary(
117        toml: &mut Table,
118    ) -> Result<GeolocationMapping, GeolocationConfigError> {
119        fn convert_value_to_json(value: Value) -> Option<SerdeValue> {
120            match value {
121                Value::String(value) => Some(SerdeValue::String(value)),
122                Value::Integer(value) => Number::try_from(value).ok().map(SerdeValue::Number),
123                Value::Float(value) => Number::from_f64(value).map(SerdeValue::Number),
124                Value::Boolean(value) => Some(SerdeValue::Bool(value)),
125                _ => None,
126            }
127        }
128
129        // Take the `addresses` field from the provided TOML table.
130        let toml = match toml
131            .remove("addresses")
132            .ok_or(GeolocationConfigError::MissingAddresses)?
133        {
134            Value::Table(table) => table,
135            _ => return Err(GeolocationConfigError::InvalidAddressesType),
136        };
137
138        let mut addresses = HashMap::<IpAddr, GeolocationData>::with_capacity(toml.len());
139
140        for (address, value) in toml {
141            let address = parse_ip_address(address.as_str())?;
142            let table = value
143                .as_table()
144                .ok_or(GeolocationConfigError::InvalidInlineEntryType)?
145                .to_owned();
146
147            let mut geolocation_data = GeolocationData::new();
148
149            for (field, value) in table {
150                let value = convert_value_to_json(value)
151                    .ok_or(GeolocationConfigError::InvalidInlineEntryType)?;
152                geolocation_data.insert(field, value);
153            }
154
155            addresses.insert(address, geolocation_data);
156        }
157
158        Ok(GeolocationMapping::InlineToml { addresses })
159    }
160
161    fn process_json_entries(
162        toml: &mut Table,
163    ) -> Result<GeolocationMapping, GeolocationConfigError> {
164        let file: PathBuf = match toml
165            .remove("file")
166            .ok_or(GeolocationConfigError::MissingFile)?
167        {
168            Value::String(file) => {
169                if file.is_empty() {
170                    return Err(GeolocationConfigError::EmptyFileEntry);
171                } else {
172                    file.into()
173                }
174            }
175            _ => return Err(GeolocationConfigError::InvalidFileEntry),
176        };
177
178        GeolocationMapping::read_json_contents(&file)?;
179
180        Ok(GeolocationMapping::Json { file })
181    }
182}
183
184impl Default for GeolocationMapping {
185    fn default() -> Self {
186        Self::Empty
187    }
188}
189
190impl GeolocationMapping {
191    pub fn get(&self, address: &IpAddr) -> Option<GeolocationData> {
192        match self {
193            Self::Empty => None,
194            Self::InlineToml { addresses } => addresses
195                .get(address)
196                .map(|geolocation_data| geolocation_data.to_owned()),
197            Self::Json { file } => Self::read_json_contents(file)
198                .ok()
199                .map(|addresses| {
200                    addresses
201                        .get(address)
202                        .map(|geolocation_data| geolocation_data.to_owned())
203                })
204                .unwrap(),
205        }
206    }
207
208    pub fn read_json_contents(
209        file: &Path,
210    ) -> Result<HashMap<IpAddr, GeolocationData>, GeolocationConfigError> {
211        let data = fs::read_to_string(file).map_err(GeolocationConfigError::IoError)?;
212
213        // Deserialize the contents of the given JSON file.
214        let json = match serde_json::from_str(&data)
215            .map_err(|_| GeolocationConfigError::GeolocationFileWrongFormat)?
216        {
217            // Check that we were given an object.
218            serde_json::Value::Object(obj) => obj,
219            _ => {
220                return Err(GeolocationConfigError::GeolocationFileWrongFormat);
221            }
222        };
223
224        let mut addresses = HashMap::<IpAddr, GeolocationData>::with_capacity(json.len());
225
226        for (address, value) in json {
227            let address = deserialization::parse_ip_address(address.as_str())?;
228            let table = value
229                .as_object()
230                .ok_or(GeolocationConfigError::InvalidInlineEntryType)?
231                .to_owned();
232
233            let geolocation_data = GeolocationData::from(&table);
234
235            addresses.insert(address, geolocation_data);
236        }
237
238        Ok(addresses)
239    }
240}
241
242impl Default for GeolocationData {
243    fn default() -> Self {
244        let default_entries = HashMap::<&str, SerdeValue>::from([
245            ("as_name", SerdeString(String::from("Fastly, Inc"))),
246            ("as_number", SerdeNumber(Number::from(54113))),
247            ("area_code", SerdeNumber(Number::from(415))),
248            ("city", SerdeString(String::from("San Francisco"))),
249            ("conn_speed", SerdeString(String::from("broadband"))),
250            ("conn_type", SerdeString(String::from("wired"))),
251            ("continent", SerdeString(String::from("NA"))),
252            ("country_code", SerdeString(String::from("US"))),
253            ("country_code3", SerdeString(String::from("USA"))),
254            (
255                "country_name",
256                SerdeString(String::from("United States of America")),
257            ),
258            ("latitude", SerdeNumber(Number::from_f64(37.77869).unwrap())),
259            (
260                "longitude",
261                SerdeNumber(Number::from_f64(-122.39557).unwrap()),
262            ),
263            ("metro_code", SerdeNumber(Number::from(0))),
264            ("postal_code", SerdeString(String::from("94107"))),
265            ("proxy_description", SerdeString(String::from("?"))),
266            ("proxy_type", SerdeString(String::from("?"))),
267            ("region", SerdeString(String::from("CA"))),
268            ("utc_offset", SerdeNumber(Number::from(-700))),
269        ]);
270
271        Self::from(default_entries)
272    }
273}
274
275impl From<HashMap<&str, SerdeValue>> for GeolocationData {
276    fn from(value: HashMap<&str, SerdeValue>) -> Self {
277        let entries = value
278            .iter()
279            .map(|(&field, value)| (field.to_string(), value.to_owned()));
280
281        Self {
282            data: Map::from_iter(entries),
283        }
284    }
285}
286
287impl From<&Map<String, SerdeValue>> for GeolocationData {
288    fn from(data: &Map<String, SerdeValue>) -> Self {
289        Self {
290            data: data.to_owned(),
291        }
292    }
293}
294
295impl GeolocationData {
296    pub fn new() -> Self {
297        Self { data: Map::new() }
298    }
299
300    pub fn insert(&mut self, field: String, value: SerdeValue) {
301        self.data.insert(field, value);
302    }
303}
304
305impl ToString for GeolocationData {
306    fn to_string(&self) -> String {
307        serde_json::to_string(&self.data).unwrap_or_else(|_| "".to_string())
308    }
309}
viceroy_lib/config/geolocation.rs

viceroy_lib/config/
geolocation.rs
