Add KDL configuration + use combined IP2ASN DB

Cargo.lock

@@ -1067,6 +1067,17 @@ dependencies = [
  "wasm-bindgen",
 ]
 
+[[package]]
+name = "kdl"
+version = "4.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "062c875482ccb676fd40c804a40e3824d4464c18c364547456d1c8e8e951ae47"
+dependencies = [
+ "miette",
+ "nom",
+ "thiserror",
+]
+
 [[package]]
 name = "lazy_static"
 version = "1.4.0"
@@ -1147,6 +1158,29 @@ version = "2.7.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6c8640c5d730cb13ebd907d8d04b52f55ac9a2eec55b440c8892f40d56c76c1d"
 
+[[package]]
+name = "miette"
+version = "5.10.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "59bb584eaeeab6bd0226ccf3509a69d7936d148cf3d036ad350abe35e8c6856e"
+dependencies = [
+ "miette-derive",
+ "once_cell",
+ "thiserror",
+ "unicode-width",
+]
+
+[[package]]
+name = "miette-derive"
+version = "5.10.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "49e7bc1560b95a3c4a25d03de42fe76ca718ab92d1a22a55b9b4cf67b3ae635c"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "mime"
 version = "0.3.17"
@@ -1922,6 +1956,7 @@ dependencies = [
  "diesel",
  "dns-lookup",
  "fluent-templates",
+ "kdl",
  "reqwest",
  "rocket",
  "rocket-accept-language",
@@ -2549,6 +2584,12 @@ dependencies = [
  "tinyvec",
 ]
 
+[[package]]
+name = "unicode-width"
+version = "0.1.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e51733f11c9c4f72aa0c160008246859e340b00807569a0da0e7a1079b27ba85"
+
 [[package]]
 name = "unicode-xid"
 version = "0.2.4"

Cargo.toml

@@ -10,6 +10,7 @@ path = "src/updater.rs"
 [dependencies]
 asn-db2 = { version = "0.2", default-features = false }
 dns-lookup = { version = "2.0", default-features = false }
+kdl = { version = "4.6", default-features = false }
 rocket = { version = "0.5", default-features = false }
 rocket-accept-language = { version = "0.8", default-features = false }
 serde = { version = "1.0", default-features = false }

README.md

@@ -27,3 +27,6 @@ contains 2 binaries:
 
 [asn-db2](https://docs.rs/asn-db2/)
 : parse [IPtoASN](https://iptoasn.com/) tsv format
+
+[kdl-rs](https://docs.rs/kdl-rs/)
+: parse [KDL](https://kdl.dev/) configuration

config.kdl

@@ -0,0 +1,3 @@
+database "services.db"
+
+ip-to-asn "ip2asn-combined.tsv"

src/lib.rs

@@ -2,15 +2,40 @@ pub mod models;
 pub mod schema;
 
 use core::fmt::{Debug, Display, Formatter, Result};
+use std::fs::read_to_string;
 
 use diesel::prelude::*;
 use diesel::sqlite::SqliteConnection;
+use kdl::KdlDocument;
 
-pub const DATABASE_URL: &str = "services.db";
+pub struct Config {
+	pub database: String,
+	pub ip_to_asn: String,
+}
+pub fn get_config() -> Config {
+	let kdl: KdlDocument = read_to_string("config.kdl")
+		.expect("unable to open config.kdl")
+		.parse()
+		.expect("failed to parse KDL");
+	Config {
+		database: kdl
+			.get_arg("database")
+			.unwrap()
+			.as_string()
+			.unwrap()
+			.to_string(),
+		ip_to_asn: kdl
+			.get_arg("ip-to-asn")
+			.unwrap()
+			.as_string()
+			.unwrap()
+			.to_string(),
+	}
+}
 
 pub fn establish_connection() -> SqliteConnection {
-	SqliteConnection::establish(DATABASE_URL)
-		.unwrap_or_else(|_| panic!("Error connecting to {}", &DATABASE_URL))
+	SqliteConnection::establish(&get_config().database)
+		.unwrap_or_else(|_| panic!("Error connecting to database"))
 }
 
 #[derive(rocket::FromFormField, Debug, PartialEq)]

src/main.rs

@@ -3,11 +3,11 @@ extern crate rocket;
 
 mod check;
 use crate::check::check;
-use services::{establish_connection, models::*, Software, DATABASE_URL};
+use services::{establish_connection, get_config, models::*, Software};
 
 use std::fs::File;
 
-use asn_db2::{Ipv4Database, Ipv6Database};
+use asn_db2::Database;
 use diesel::prelude::*;
 use fluent_templates::FluentLoader;
 use rocket::{
@@ -70,19 +70,14 @@ fn rocket() -> _ {
 	let mut tera = Tera::new("templates/*.html.tera").unwrap();
 	tera.register_function("fluent", FluentLoader::new(&*LOCALES));
 
-	let db6 = Ipv6Database::from_reader(BufReader::new(
-		File::open("ip2asn-v6.tsv").expect("unable to open ip2asn-v6.tsv"),
-	))
-	.unwrap();
-	let db4 = Ipv4Database::from_reader(BufReader::new(
-		File::open("ip2asn-v4.tsv").expect("unable to open ip2asn-v4.tsv"),
+	let db: Database = Database::from_reader(BufReader::new(
+		File::open(get_config().ip_to_asn).expect("unable to open ip2asn TSV file"),
 	))
 	.unwrap();
 
 	rocket::build()
 		.manage(tera)
-		.manage(db6)
-		.manage(db4)
+		.manage(db)
 		.attach(Shield::new().enable(Referrer::NoReferrer))
 		.attach(HttpHeaders)
 		.mount(
@@ -99,7 +94,7 @@ fn rocket() -> _ {
 
 #[get("/services.db")]
 fn dl() -> File {
-	File::open(DATABASE_URL).unwrap()
+	File::open(get_config().database).unwrap()
 }
 
 #[derive(Serialize, Debug)]
@@ -124,8 +119,7 @@ struct TemplateServices {
 #[get("/?<software>")]
 fn list_services(
 	tera: &State<Tera>,
-	db6: &State<Ipv6Database>,
-	db4: &State<Ipv4Database>,
+	db: &State<Database>,
 	al: &AcceptLanguage,
 	software: Option<Strict<Software>>,
 ) -> RawHtml<String> {
@@ -143,33 +137,28 @@ fn list_services(
 	let mut templates: Vec<TemplateServices> = vec![];
 	for service in &services {
 		let mut ip_info = vec![];
-		service
-			.address_ipv6
-			.split(',')
+		let mut ip_combined: Vec<&str> = service.address_ipv6.split(',').collect();
+		ip_combined.append(&mut service.address_ipv4.split(',').collect());
+		ip_combined
+			.iter()
 			.filter(|ip| !ip.is_empty())
 			.for_each(|ip| {
-				let info = db6.lookup(ip.parse().unwrap()).unwrap();
-				ip_info.push(IpInfo {
-					ip: ip.to_string(),
-					subnet: info.subnet.to_string(),
-					asn: info.as_number.to_string(),
-					country: info.country.to_string(),
-					as_owner: info.owner.to_string(),
-				})
-			});
-		service
-			.address_ipv4
-			.split(',')
-			.filter(|ip| !ip.is_empty())
-			.for_each(|ip| {
-				let info = db4.lookup(ip.parse().unwrap()).unwrap();
-				ip_info.push(IpInfo {
-					ip: ip.to_string(),
-					subnet: info.subnet.to_string(),
-					asn: info.as_number.to_string(),
-					country: info.country.to_string(),
-					as_owner: info.owner.to_string(),
-				})
+				match db.lookup(ip.parse().unwrap()).unwrap() {
+					asn_db2::IpEntry::V6(info) => ip_info.push(IpInfo {
+						ip: ip.to_string(),
+						subnet: info.subnet.to_string(),
+						asn: info.as_number.to_string(),
+						country: info.country.to_string(),
+						as_owner: info.owner.to_string(),
+					}),
+					asn_db2::IpEntry::V4(info) => ip_info.push(IpInfo {
+						ip: ip.to_string(),
+						subnet: info.subnet.to_string(),
+						asn: info.as_number.to_string(),
+						country: info.country.to_string(),
+						as_owner: info.owner.to_string(),
+					}),
+				};
 			});
 
 		templates.push(TemplateServices {