ip2location_ip2location/
bin_format.rs

1//
2#[cfg(feature = "tokio_fs")]
3pub type TokioFile = async_compat::Compat<tokio::fs::File>;
4
5#[cfg(feature = "async_fs")]
6pub type AsyncFsFile = async_fs::File;
7
8use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
9
10use futures_util::{AsyncRead, AsyncSeek};
11use ip2location_bin_format::querier::{
12    LookupError as QuerierLookupError, NewError as QuerierNewError, Querier,
13};
14
15use crate::record::{OptionRecord, Record, RecordField};
16
17//
18pub struct Database<S> {
19    pub inner: Querier<S>,
20}
21
22impl<S> core::fmt::Debug for Database<S>
23where
24    Querier<S>: core::fmt::Debug,
25{
26    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
27        f.debug_struct("Database")
28            .field("inner", &self.inner)
29            .finish()
30    }
31}
32
33#[cfg(feature = "tokio_fs")]
34impl Database<async_compat::Compat<tokio::fs::File>> {
35    pub async fn new(
36        path: impl AsRef<std::path::Path>,
37        pool_max_size: usize,
38    ) -> Result<Self, DatabaseNewError> {
39        use futures_util::TryFutureExt as _;
40
41        let path = path.as_ref().to_owned();
42
43        let inner = Querier::new(
44            || Box::pin(tokio::fs::File::open(path.clone()).map_ok(async_compat::Compat::new)),
45            pool_max_size,
46        )
47        .await
48        .map_err(DatabaseNewError::QuerierNewError)?;
49
50        if !inner.header.r#type.is_ip2location() {
51            return Err(DatabaseNewError::TypeMismatch);
52        }
53
54        Ok(Self { inner })
55    }
56}
57
58#[cfg(feature = "async_fs")]
59impl Database<async_fs::File> {
60    pub async fn new(
61        path: impl AsRef<std::path::Path>,
62        pool_max_size: usize,
63    ) -> Result<Self, DatabaseNewError> {
64        let path = path.as_ref().to_owned();
65
66        let inner = Querier::new(
67            || Box::pin(async_fs::File::open(path.clone())),
68            pool_max_size,
69        )
70        .await
71        .map_err(DatabaseNewError::QuerierNewError)?;
72
73        if !inner.header.r#type.is_ip2location() {
74            return Err(DatabaseNewError::TypeMismatch);
75        }
76
77        Ok(Self { inner })
78    }
79}
80
81//
82#[derive(Debug)]
83pub enum DatabaseNewError {
84    QuerierNewError(QuerierNewError),
85    TypeMismatch,
86}
87
88impl core::fmt::Display for DatabaseNewError {
89    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
90        write!(f, "{self:?}")
91    }
92}
93
94impl std::error::Error for DatabaseNewError {}
95
96//
97//
98//
99impl<S> Database<S>
100where
101    S: AsyncSeek + AsyncRead + Unpin,
102{
103    pub async fn lookup(
104        &self,
105        ip: IpAddr,
106        selected_fields: impl Into<Option<&[RecordField]>>,
107    ) -> Result<Option<Record>, DatabaseLookupError> {
108        match ip {
109            IpAddr::V4(ip) => self.lookup_ipv4(ip, selected_fields).await,
110            IpAddr::V6(ip) => self.lookup_ipv6(ip, selected_fields).await,
111        }
112    }
113
114    pub async fn lookup_ipv4(
115        &self,
116        ip: Ipv4Addr,
117        selected_fields: impl Into<Option<&[RecordField]>>,
118    ) -> Result<Option<Record>, DatabaseLookupError> {
119        let selected_fields: Option<Vec<ip2location_bin_format::record_field::RecordField>> =
120            selected_fields
121                .into()
122                .map(|x| x.iter().map(Into::into).collect::<Vec<_>>());
123        let selected_fields = selected_fields.as_deref();
124
125        //
126        match self
127            .inner
128            .lookup_ipv4(ip, selected_fields)
129            .await
130            .map_err(DatabaseLookupError::QuerierLookupError)?
131        {
132            Some(x) => Ok(OptionRecord::try_from(x)
133                .map_err(DatabaseLookupError::ToRecordFailed)?
134                .0),
135            None => Ok(None),
136        }
137    }
138
139    pub async fn lookup_ipv6(
140        &self,
141        ip: Ipv6Addr,
142        selected_fields: impl Into<Option<&[RecordField]>>,
143    ) -> Result<Option<Record>, DatabaseLookupError> {
144        let selected_fields: Option<Vec<ip2location_bin_format::record_field::RecordField>> =
145            selected_fields
146                .into()
147                .map(|x| x.iter().map(Into::into).collect::<Vec<_>>());
148        let selected_fields = selected_fields.as_deref();
149
150        //
151        match self
152            .inner
153            .lookup_ipv6(ip, selected_fields)
154            .await
155            .map_err(DatabaseLookupError::QuerierLookupError)?
156        {
157            Some(x) => Ok(OptionRecord::try_from(x)
158                .map_err(DatabaseLookupError::ToRecordFailed)?
159                .0),
160            None => Ok(None),
161        }
162    }
163}
164
165//
166#[derive(Debug)]
167pub enum DatabaseLookupError {
168    QuerierLookupError(QuerierLookupError),
169    ToRecordFailed(Box<str>),
170}
171
172impl core::fmt::Display for DatabaseLookupError {
173    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
174        write!(f, "{self:?}")
175    }
176}
177
178impl std::error::Error for DatabaseLookupError {}
179
180#[cfg(feature = "tokio_fs")]
181#[cfg(test)]
182mod tests {
183    use super::*;
184
185    use std::io::ErrorKind as IoErrorKind;
186
187    #[tokio::test]
188    async fn test_new_and_lookup_latest() -> Result<(), Box<dyn std::error::Error>> {
189        let path_bin_v4 = "data/ip2location-lite/latest/IP2LOCATION-LITE-DB11.BIN";
190        let path_bin_v6 = "data/ip2location-lite/latest/IP2LOCATION-LITE-DB11.IPV6.BIN";
191
192        let db_v4 = match Database::<TokioFile>::new(path_bin_v4, 1).await {
193            Ok(x) => Some(x),
194            Err(DatabaseNewError::QuerierNewError(QuerierNewError::OpenFailed(err)))
195                if err.kind() == IoErrorKind::NotFound =>
196            {
197                None
198            }
199            Err(err) => panic!("{err:?}"),
200        };
201        let db_v6 = match Database::<TokioFile>::new(path_bin_v6, 1).await {
202            Ok(x) => Some(x),
203            Err(DatabaseNewError::QuerierNewError(QuerierNewError::OpenFailed(err)))
204                if err.kind() == IoErrorKind::NotFound =>
205            {
206                None
207            }
208            Err(err) => panic!("{err:?}"),
209        };
210
211        if let Some(db_v4) = db_v4 {
212            let record_1 = db_v4
213                .lookup(Ipv4Addr::from(16777216).into(), None)
214                .await?
215                .unwrap();
216            assert_eq!(record_1.country_code.to_string(), "US");
217            assert_eq!(record_1.latitude.unwrap(), 34.052_86);
218
219            let selected_fields = &[
220                RecordField::CountryCodeAndName,
221                RecordField::RegionName,
222                RecordField::Latitude,
223            ];
224            let record_2 = db_v4
225                .lookup(Ipv4Addr::from(16777472).into(), selected_fields.as_ref())
226                .await?
227                .unwrap();
228            assert_eq!(record_2.country_code.to_string(), "CN");
229            println!("{record_2:?}");
230        }
231
232        if let Some(db_v6) = db_v6 {
233            let record_1 = db_v6
234                .lookup(Ipv6Addr::from(281470698520576).into(), None)
235                .await?
236                .unwrap();
237            assert_eq!(record_1.country_code.to_string(), "US");
238
239            let selected_fields = &[
240                RecordField::CountryCodeAndName,
241                RecordField::RegionName,
242                RecordField::Latitude,
243            ];
244            let record_2 = db_v6
245                .lookup(
246                    Ipv6Addr::from(281470698520832).into(),
247                    selected_fields.as_ref(),
248                )
249                .await?
250                .unwrap();
251            assert_eq!(record_2.country_code.to_string(), "CN");
252            println!("{record_2:?}");
253        }
254
255        Ok(())
256    }
257}