Updates for Hyper 1.0

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "MaxMind-DB"]
+	path = MaxMind-DB
+	url = https://github.com/maxmind/MaxMind-DB/

Cargo.toml

@@ -2,19 +2,22 @@
 name = "geoip-server"
 version = "0.1.0"
 authors = ["Donovan Schönknecht <[email protected]>"]
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-tokio-io = "0.1.2"
-tokio-core = "0.1.7"
-tokio-proto = "0.1.1"
-tokio-service = "0.1.0"
-futures = "0.1.14"
-futures-cpupool = "0.1.5"
-serde = "1.0.8"
-serde_derive = "1.0.8"
-serde_json = "1.0.2"
-maxminddb = "0.7.2"
-num_cpus = "1.5.0"
+http-body-util = "0.1.0"
+hyper = { version = "1.0.1", features = ["full"] }
+hyper-util = { version = "0.1.1", features = ["full"] }
+maxminddb = { version = "0.23.0", features = ["mmap"] }
+serde = "1.0.192"
+serde_json = "1.0.108"
+tokio = { version = "1.34.0", features = ["full"] }
+
+[dev-dependencies]
+tokio-test = "0.4.3"
+
+[[bin]]
+name = "geoip-server"
 
-[dependencies.tokio-minihttp]
-git = "https://github.com/tokio-rs/tokio-minihttp"

README.md

@@ -2,13 +2,13 @@
 Fast, lightweight MaxMind GeoIP lookup server written in Rust (experimental)
 
 ### Starting the server
-Usage: `geoip-server 0.0.0.0:3000 /path/to/GeoIP2-City.mmdb`
+Usage: `./geoip-server localhost:3000 MaxMind-DB/test-data/GeoIP2-City-Test.mmdb`
 
 ### Querying
-Make a GET request to `http://localhost:3000/<ip>`
+Make a GET request to `http://localhost:3000/89.160.20.128`
 
 Example response:
 ```json
-{"ip":"169.0.183.91","latitude":-33.9185,"longitude":18.4131,"time_zone":"Africa/Johannesburg","iso_code":"ZA","city":{"de":"Kapstadt","en":"Cape Town","es":"Ciudad del Cabo","fr":"Le Cap","ja":"ケープタウン","pt-BR":"Cidade do Cabo","ru":"Кейптаун"},"subdivisions":[{"en":"Western Cape","pt-BR":"Cabo Ocidental"}],"country":{"de":"Südafrika","en":"South Africa","es":"Sudáfrica","fr":"Afrique du Sud","ja":"南アフリカ","pt-BR":"África do Sul","ru":"ЮАР","zh-CN":"南非"},"registered_country":{"de":"Südafrika","en":"South Africa","es":"Sudáfrica","fr":"Afrique du Sud","ja":"南アフリカ","pt-BR":"África do Sul","ru":"ЮАР","zh-CN":"南非"}}
+{"city":{"geoname_id":2694762,"names":{"de":"Linköping","en":"Linköping","fr":"Linköping","ja":"リンシェーピング","zh-CN":"林雪平"}},"continent":{"code":"EU","geoname_id":6255148,"names":{"de":"Europa","en":"Europe","es":"Europa","fr":"Europe","ja":"ヨーロッパ","pt-BR":"Europa","ru":"Европа","zh-CN":"欧洲"}},"country":{"geoname_id":2661886,"is_in_european_union":true,"iso_code":"SE","names":{"de":"Schweden","en":"Sweden","es":"Suecia","fr":"Suède","ja":"スウェーデン王国","pt-BR":"Suécia","ru":"Швеция","zh-CN":"瑞典"}},"location":{"accuracy_radius":76,"latitude":58.4167,"longitude":15.6167,"metro_code":null,"time_zone":"Europe/Stockholm"},"postal":null,"registered_country":{"geoname_id":2921044,"is_in_european_union":true,"iso_code":"DE","names":{"de":"Deutschland","en":"Germany","es":"Alemania","fr":"Allemagne","ja":"ドイツ連邦共和国","pt-BR":"Alemanha","ru":"Германия","zh-CN":"德国"}},"represented_country":null,"subdivisions":[{"geoname_id":2685867,"iso_code":"E","names":{"en":"Östergötland County","fr":"Comté d'Östergötland"}}],"traits":null}
 ```
 

src/main.rs

@@ -1,179 +1,130 @@
-#[macro_use]
-extern crate serde_derive;
-extern crate maxminddb;
-extern crate futures;
-extern crate futures_cpupool;
-extern crate serde;
-extern crate serde_json;
-extern crate tokio_minihttp;
-extern crate tokio_proto;
-extern crate tokio_service;
-extern crate num_cpus;
-
-use std::sync::Arc;
-use std::io;
-use std::net::{IpAddr, SocketAddr, AddrParseError};
-use std::str::FromStr;
-use std::collections::BTreeMap;
-use std::env::args;
-use futures::{BoxFuture, Future};
-use futures_cpupool::CpuPool;
-use tokio_minihttp::{Request, Response};
-use tokio_proto::TcpServer;
-use tokio_service::Service;
+use http_body_util::Full;
+use hyper::body::Bytes;
+use hyper::server::conn::http1;
+use hyper::Error;
+use hyper::{service::service_fn, Response};
+use hyper::header::CONTENT_TYPE;
+use hyper::StatusCode;
+use hyper_util::rt::TokioIo;
 use maxminddb::geoip2;
-use maxminddb::geoip2::model::{City, Country, Subdivision, Location};
-
-#[derive(Serialize)]
-struct Result {
-	ip: String,
-	latitude: f64,
-	longitude: f64,
-	time_zone: String,
-	iso_code: String,
-	city: BTreeMap<String, String>,
-	subdivisions: Vec<BTreeMap<String, String>>,
-	country: BTreeMap<String, String>,
-	registered_country: BTreeMap<String, String>
-}
-
-struct Server {
-	thread_pool: CpuPool,
-	reader: Arc<maxminddb::Reader>
-}
-
-fn get_city_names(city: &Option<City>) -> BTreeMap<String, String> {
-	match *city {
-		Some(ref xcity) => xcity.names.clone().unwrap(),
-		None => std::collections::BTreeMap::new()
-	}
+use std::net::{IpAddr, SocketAddr, ToSocketAddrs};
+use std::sync::Arc;
+use std::env;
+use tokio::net::TcpListener;
+
+fn parse_args() -> Result<(SocketAddr, String), Box<dyn std::error::Error>> {
+    let args: Vec<String> = env::args().collect();
+    if args.len() < 3 {
+
+        eprintln!("Usage: {} <bind ip:port> <mmdb file>", args[0]);
+        std::process::exit(1);
+    }
+
+    let bind_host = &args[1];
+    let addr = bind_host.to_socket_addrs()?.next().ok_or("Invalid bind host")?;
+    let mmdb_file = args[2].clone();
+
+    Ok((addr, mmdb_file))
 }
 
-fn get_country_names(country: &Option<Country>) -> BTreeMap<String, String> {
-	match *country {
-		Some(ref xcountry) => xcountry.names.clone().unwrap(),
-		None => std::collections::BTreeMap::new()
-	}
+async fn run_server(listener: TcpListener, db: Arc<maxminddb::Reader<Vec<u8>>>) -> Result<(), Box<dyn std::error::Error>> {
+    println!("Listening on http://{}", listener.local_addr()?);
+    loop {
+        let (stream, _) = listener.accept().await?;
+        let io = TokioIo::new(stream);
+        let db = db.clone();
+
+        let service = service_fn(move |req| {
+            let db = db.clone();
+            let ip = req.uri().path().trim_start_matches('/').to_string();
+
+            async move {
+                handle_request(&ip, db).await
+            }
+        });
+
+        if let Err(err) = http1::Builder::new().serve_connection(io, service).await {
+            println!("Error serving connection: {:?}", err);
+        }
+    }
 }
 
-fn get_country_iso_code(country: &Option<Country>) -> String {
-	match *country {
-		Some(ref xcountry) => xcountry.iso_code.clone().unwrap(),
-		None => String::from("")
-	}
+async fn handle_request(ip: &str, db: Arc<maxminddb::Reader<Vec<u8>>>) -> Result<Response<Full<Bytes>>, Error> {
+    let body: Full<Bytes>;
+    let status: StatusCode;
+
+    if let Ok(ipaddr) = ip.parse::<IpAddr>() {
+        let entry: Option<geoip2::City> = db.lookup(ipaddr).ok();
+        status = match entry {
+            Some(_) => StatusCode::OK,
+            None => StatusCode::NOT_FOUND
+        };
+        body = match entry {
+            Some(result) => Full::new(Bytes::from(
+                serde_json::to_string(&result).unwrap().to_string(),
+            )),
+            None => Full::new(Bytes::from("{\"error\": \"not_found\"}"))
+        };
+    } else {
+        status = StatusCode::BAD_REQUEST;
+        body = Full::new(Bytes::from("{\"error\": \"invalid_ip\"}"))
+    }
+
+    Ok::<_, Error>(Response::builder()
+    .header(CONTENT_TYPE, "application/json")
+    .status(status)
+    .body(body).unwrap())
 }
 
-fn get_location_latitude(location: &Option<Location>) -> f64 {
-	match *location {
-		Some(ref xlocation) => xlocation.clone().latitude.unwrap(),
-		None => 0.0
-	}
-}
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let (addr, mmdb_file) = parse_args()?;
 
-fn get_location_longitude(location: &Option<Location>) -> f64 {
-	match *location {
-		Some(ref xlocation) => xlocation.clone().longitude.unwrap(),
-		None => 0.0
-	}
-}
+    let listener = TcpListener::bind(addr).await?;
 
-fn get_location_time_zone(location: &Option<Location>) -> String {
-	match *location {
-		Some(ref xlocation) => xlocation.clone().time_zone.unwrap(),
-		None => String::from("")
-	}
-}
+    let db = Arc::new(maxminddb::Reader::open_readfile(mmdb_file)?);
 
-fn get_subdivision_names(subdivisions: &Option<Vec<Subdivision>>) -> Vec<BTreeMap<String, String>> {
-	match *subdivisions {
-		Some(ref xsubdivisions) => {
-			let mut subdivisions = vec![];
-			for subdivision in xsubdivisions {
-				subdivisions.push(subdivision.names.clone().unwrap());
-			}
-			subdivisions
-		}
-		None => Vec::new()
-	}
-}
+    run_server(listener, db).await?;
 
-impl Service for Server {
-	type Request = Request;
-	type Response = Response;
-	type Error = io::Error;
-	type Future = BoxFuture<Response, io::Error>;
-
-	fn call(&self, req: Request) -> Self::Future {
-		let db = self.reader.clone();
-		let msg = self.thread_pool.spawn_fn(move || {
-			let path = req.path();
-			let foo = &path[1..];
-			//println!("Looking up {}", foo);
-
-			match IpAddr::from_str(foo) {
-				Ok(ip) => {
-					match db.lookup(ip) {
-						Ok(xlookup) => {
-							let lookup: geoip2::City = xlookup;
-							let result = Result {
-								ip: String::from_str(foo).unwrap(),
-								longitude: get_location_longitude(&lookup.location),
-								latitude: get_location_latitude(&lookup.location),
-								time_zone: get_location_time_zone(&lookup.location),
-								iso_code: get_country_iso_code(&lookup.country),
-								city: get_city_names(&lookup.city),
-								subdivisions: get_subdivision_names(&lookup.subdivisions),
-								country: get_country_names(&lookup.country),
-								registered_country: get_country_names(&lookup.registered_country),
-							};
-							Ok(serde_json::to_string(&result).unwrap())
-						},
-						Err(_) => Ok(String::from("{\"error\": \"not_found\"}"))
-					}
-				},
-				Err(AddrParseError(_)) => Ok(String::from("{\"error\": \"invalid_address\"}")),
-			}
-		});
-
-		msg.map(|body| {
-			let mut response = Response::new();
-			response.header("Content-Type", "application/json; charset=UTF-8");
-			response.body(&body);
-			response
-		}).boxed()
-	}
+    Ok(())
 }
 
-fn main() {
-	let args = args().collect::<Vec<_>>();
-
-	if args.len() < 3 {
-		println!("Usage: {} <bind ip:port> <mmdb file>", args[0]);
-		return;
-	}
-
-	let addr = match SocketAddr::from_str(&args[1]) {
-		Ok(addr) => addr,
-		Err(AddrParseError(_)) => {
-			println!("Error: Invalid bind address: \"{}\".", &args[1]);
-			return
-		}
-	};
-
-	let reader = match maxminddb::Reader::open(&args[2]) {
-		Ok(reader) => reader,
-		Err(_) => {
-			println!("Error: Unable to open database file: \"{}\"", &args[2]);
-			return
-		}
-	};
-
-	let db = Arc::new(reader);
-	let thread_pool = CpuPool::new(num_cpus::get());
-	TcpServer::new(tokio_minihttp::Http, addr).serve(move || {
-		Ok(Server {
-			thread_pool: thread_pool.clone(),
-			reader: db.clone(),
-		})
-	});
-}
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tokio_test::assert_ok;
+
+    async fn mock_handle_request(ip: &str) -> Result<Response<Full<Bytes>>, Error> {
+        let db = Arc::new(maxminddb::Reader::open_readfile("MaxMind-DB/test-data/GeoIP2-City-Test.mmdb").unwrap());
+        handle_request(&ip, db).await
+    }
+
+    #[tokio::test]
+    async fn test_handle_request_with_valid_ip() {
+        let ip = "89.160.20.128";
+        let response = mock_handle_request(ip).await;
+        assert_ok!(&response);
+        let data = response.unwrap();
+        assert_eq!(data.headers().get(CONTENT_TYPE).unwrap(), "application/json");
+        assert_eq!(data.status(), StatusCode::OK);
+    }
+
+    #[tokio::test]
+    async fn test_handle_request_with_unknown_ip() {
+        let ip = "192.168.0.1";
+        let response = mock_handle_request(ip).await;
+        assert_ok!(&response);
+        let data = response.unwrap();
+        assert_eq!(data.headers().get(CONTENT_TYPE).unwrap(), "application/json");
+        assert_eq!(data.status(), StatusCode::NOT_FOUND);
+    }
+
+    #[tokio::test]
+    async fn test_handle_request_with_invalid_ip() {
+        let ip = "invalid-ip"; // Example of an invalid IP
+        let response = mock_handle_request(ip).await;
+        assert_ok!(&response);
+        let data = response.unwrap();
+        assert_eq!(data.status(), StatusCode::BAD_REQUEST);
+    }
+}