load configuration and database only once

This commit is contained in:
Miraty 2024-04-23 01:33:39 +02:00
parent 973d1c5b41
commit b2a6d8d02a
5 changed files with 158 additions and 86 deletions

47
Cargo.lock generated
View File

@ -397,6 +397,7 @@ checksum = "03fc05c17098f21b89bc7d98fe1dd3cce2c11c2ad8e145f2a44fe08ed28eb559"
dependencies = [ dependencies = [
"diesel_derives", "diesel_derives",
"libsqlite3-sys", "libsqlite3-sys",
"r2d2",
"time", "time",
] ]
@ -1530,6 +1531,17 @@ dependencies = [
"proc-macro2", "proc-macro2",
] ]
[[package]]
name = "r2d2"
version = "0.8.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93"
dependencies = [
"log",
"parking_lot",
"scheduled-thread-pool",
]
[[package]] [[package]]
name = "rand" name = "rand"
version = "0.8.5" version = "0.8.5"
@ -1782,6 +1794,31 @@ dependencies = [
"uncased", "uncased",
] ]
[[package]]
name = "rocket_sync_db_pools"
version = "0.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d83f32721ed79509adac4328e97f817a8f55a47c4b64799f6fd6cc3adb6e42ff"
dependencies = [
"diesel",
"r2d2",
"rocket",
"rocket_sync_db_pools_codegen",
"serde",
"tokio",
"version_check",
]
[[package]]
name = "rocket_sync_db_pools_codegen"
version = "0.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5cc890925dc79370c28eb15c9957677093fdb7e8c44966d189f38cedb995ee68"
dependencies = [
"devise",
"quote",
]
[[package]] [[package]]
name = "rustc-demangle" name = "rustc-demangle"
version = "0.1.23" version = "0.1.23"
@ -1869,6 +1906,15 @@ dependencies = [
"winapi-util", "winapi-util",
] ]
[[package]]
name = "scheduled-thread-pool"
version = "0.2.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19"
dependencies = [
"parking_lot",
]
[[package]] [[package]]
name = "scoped-tls" name = "scoped-tls"
version = "1.0.1" version = "1.0.1"
@ -1960,6 +2006,7 @@ dependencies = [
"reqwest", "reqwest",
"rocket", "rocket",
"rocket-accept-language", "rocket-accept-language",
"rocket_sync_db_pools",
"serde", "serde",
"serde_json", "serde_json",
"tera", "tera",

View File

@ -34,6 +34,11 @@ version = "0.12"
default-features = false default-features = false
features = ["rustls-tls", "socks"] features = ["rustls-tls", "socks"]
[dependencies.rocket_sync_db_pools]
version = "0.1"
default-features = false
features = ["diesel_sqlite_pool"]
[dependencies.tera] [dependencies.tera]
version = "1.19" version = "1.19"
default-features = false default-features = false

View File

@ -4,8 +4,6 @@ pub mod schema;
use core::fmt::{Debug, Display, Formatter, Result}; use core::fmt::{Debug, Display, Formatter, Result};
use std::fs::read_to_string; use std::fs::read_to_string;
use diesel::prelude::*;
use diesel::sqlite::SqliteConnection;
use kdl::KdlDocument; use kdl::KdlDocument;
pub struct Config { pub struct Config {
@ -40,11 +38,6 @@ pub fn get_config() -> Config {
} }
} }
pub fn establish_connection() -> SqliteConnection {
SqliteConnection::establish(&get_config().database)
.unwrap_or_else(|_| panic!("Error connecting to database"))
}
#[derive(rocket::FromFormField, Debug, PartialEq)] #[derive(rocket::FromFormField, Debug, PartialEq)]
pub enum Software { pub enum Software {
PeerTube, PeerTube,

View File

@ -3,7 +3,7 @@ extern crate rocket;
mod check; mod check;
use crate::check::check; use crate::check::check;
use services::{establish_connection, get_config, models::*, Software}; use services::{get_config, models::*, Config, Software};
use std::fs::File; use std::fs::File;
@ -12,6 +12,7 @@ use diesel::prelude::*;
use fluent_templates::FluentLoader; use fluent_templates::FluentLoader;
use rocket::{ use rocket::{
fairing::{Fairing, Info, Kind}, fairing::{Fairing, Info, Kind},
figment::{util::map, value::Value},
form::{self, Error, Form, Strict}, form::{self, Error, Form, Strict},
http::Header, http::Header,
response::content::RawHtml, response::content::RawHtml,
@ -19,10 +20,11 @@ use rocket::{
Request, Response, State, Request, Response, State,
}; };
use rocket_accept_language::{language, AcceptLanguage, LanguageIdentifier}; use rocket_accept_language::{language, AcceptLanguage, LanguageIdentifier};
use rocket_sync_db_pools::{database, diesel};
use serde::Serialize; use serde::Serialize;
use serde_json::json; use serde_json::json;
use std::io::BufReader; use std::io::BufReader;
use tera::{Context, Tera, Value}; use tera::{Context, Tera};
use unic_langid::{langid, subtags::Language}; use unic_langid::{langid, subtags::Language};
const LANGUAGE_DEFAULT: Language = language!("en"); const LANGUAGE_DEFAULT: Language = language!("en");
@ -37,7 +39,7 @@ fluent_templates::static_loader! {
}; };
} }
fn gen_context(accept_language: &AcceptLanguage, values: Value) -> Context { fn gen_context(accept_language: &AcceptLanguage, values: tera::Value) -> Context {
let mut cont: Context = Context::from_value(json!({ let mut cont: Context = Context::from_value(json!({
"ln": &accept_language "ln": &accept_language
.get_appropriate_language_region(&SL) .get_appropriate_language_region(&SL)
@ -66,37 +68,49 @@ impl Fairing for HttpHeaders {
} }
} }
#[database("main_db")]
struct DbConn(diesel::SqliteConnection);
#[launch] #[launch]
fn rocket() -> _ { fn rocket() -> _ {
let config = get_config();
let mut tera = Tera::new("templates/*.html.tera").unwrap(); let mut tera = Tera::new("templates/*.html.tera").unwrap();
tera.register_function("fluent", FluentLoader::new(&*LOCALES)); tera.register_function("fluent", FluentLoader::new(&*LOCALES));
let db: Database = Database::from_reader(BufReader::new( rocket::custom(rocket::Config::figment().merge((
File::open(get_config().ip_to_asn).expect("unable to open ip2asn TSV file"), "databases",
)) map!["main_db" => map!{
.unwrap(); "url" => Into::<Value>::into(config.database.clone()),
}],
rocket::build() )))
.manage(tera) .manage(
.manage(db) Database::from_reader(BufReader::new(
.attach(Shield::new().enable(Referrer::NoReferrer)) File::open(&config.ip_to_asn).expect("unable to open ip2asn TSV file"),
.attach(HttpHeaders) ))
.mount( .unwrap(),
"/", )
routes![ .manage(config)
list_services, .manage(tera)
list_scans, .attach(DbConn::fairing())
add_service_get, .attach(Shield::new().enable(Referrer::NoReferrer))
add_service_post, .attach(HttpHeaders)
dl, .mount(
about, "/",
], routes![
) list_services,
list_scans,
add_service_get,
add_service_post,
dl,
about,
],
)
} }
#[get("/services.db")] #[get("/services.db")]
fn dl() -> File { fn dl(config: &State<Config>) -> File {
File::open(get_config().database).unwrap() File::open(&config.database).unwrap()
} }
#[derive(Serialize, Debug)] #[derive(Serialize, Debug)]
@ -119,22 +133,27 @@ struct TemplateServices {
ip_info: Vec<IpInfo>, ip_info: Vec<IpInfo>,
} }
#[get("/?<software>")] #[get("/?<software>")]
fn list_services( async fn list_services(
conn: DbConn,
tera: &State<Tera>, tera: &State<Tera>,
db: &State<Database>, ipdb: &State<Database>,
al: &AcceptLanguage, al: &AcceptLanguage,
software: Option<Strict<Software>>, software: Option<Strict<Software>>,
) -> RawHtml<String> { ) -> RawHtml<String> {
let mut request = services::schema::services::dsl::services.into_boxed(); let services = conn
if let Some(s) = software { .run(|c| {
request = let mut request = services::schema::services::dsl::services.into_boxed();
request.filter(services::schema::services::software.eq(s.to_string().to_lowercase())); if let Some(s) = software {
} request = request
let services = request .filter(services::schema::services::software.eq(s.to_string().to_lowercase()));
.limit(300) }
.select(Services::as_select()) request
.load(&mut establish_connection()) .limit(300)
.unwrap(); .select(Services::as_select())
.load(c)
.unwrap()
})
.await;
let mut templates: Vec<TemplateServices> = vec![]; let mut templates: Vec<TemplateServices> = vec![];
for service in &services { for service in &services {
@ -145,7 +164,7 @@ fn list_services(
.iter() .iter()
.filter(|ip| !ip.is_empty()) .filter(|ip| !ip.is_empty())
.for_each(|ip| { .for_each(|ip| {
match db.lookup(ip.parse().unwrap()).unwrap() { match ipdb.lookup(ip.parse().unwrap()).unwrap() {
asn_db2::IpEntry::V6(info) => ip_info.push(IpInfo { asn_db2::IpEntry::V6(info) => ip_info.push(IpInfo {
ip: ip.to_string(), ip: ip.to_string(),
subnet: info.subnet.to_string(), subnet: info.subnet.to_string(),
@ -190,18 +209,18 @@ fn list_services(
} }
#[get("/list-scans")] #[get("/list-scans")]
fn list_scans(tera: &State<Tera>, al: &AcceptLanguage) -> RawHtml<String> { async fn list_scans(conn: DbConn, tera: &State<Tera>, al: &AcceptLanguage) -> RawHtml<String> {
RawHtml( RawHtml(
tera.render( tera.render(
"list-scans.html.tera", "list-scans.html.tera",
&gen_context( &gen_context(
al, al,
json!({ json!({
"scans": services::schema::scans::dsl::scans "scans": conn.run(|c| services::schema::scans::dsl::scans
.limit(1000) .limit(1000)
.select(Scans::as_select()) .select(Scans::as_select())
.load(&mut establish_connection()) .load(c)
.unwrap() .unwrap()).await
}), }),
), ),
) )
@ -242,6 +261,7 @@ struct Submission<'r> {
#[post("/add-service", data = "<submission>")] #[post("/add-service", data = "<submission>")]
async fn add_service_post( async fn add_service_post(
conn: DbConn,
submission: Form<Strict<Submission<'_>>>, submission: Form<Strict<Submission<'_>>>,
tera: &State<Tera>, tera: &State<Tera>,
al: &AcceptLanguage, al: &AcceptLanguage,
@ -262,20 +282,23 @@ async fn add_service_post(
} }
}; };
diesel::insert_into(services::table()) conn.run(|c| {
.values(Services { diesel::insert_into(services::table())
url: service.url, .values(Services {
software: service.software, url: service.url,
server: service.server, software: service.software,
ipv6: "".to_string(), server: service.server,
ipv4: "".to_string(), ipv6: "".to_string(),
availability_ipv6: "".to_string(), ipv4: "".to_string(),
availability_ipv4: "".to_string(), availability_ipv6: "".to_string(),
address_ipv6: "".to_string(), availability_ipv4: "".to_string(),
address_ipv4: "".to_string(), address_ipv6: "".to_string(),
}) address_ipv4: "".to_string(),
.execute(&mut establish_connection()) })
.unwrap(); .execute(c)
.unwrap()
})
.await;
RawHtml( RawHtml(
tera.render("add-service.html.tera", &gen_context(al, json!({}))) tera.render("add-service.html.tera", &gen_context(al, json!({})))
@ -284,14 +307,14 @@ async fn add_service_post(
} }
#[get("/about")] #[get("/about")]
fn about(tera: &State<Tera>, al: &AcceptLanguage) -> RawHtml<String> { fn about(config: &State<Config>, tera: &State<Tera>, al: &AcceptLanguage) -> RawHtml<String> {
RawHtml( RawHtml(
tera.render( tera.render(
"about.html.tera", "about.html.tera",
&gen_context( &gen_context(
al, al,
json!({ json!({
"source_code": get_config().source_code, "source_code": config.source_code,
}), }),
), ),
) )

View File

@ -1,28 +1,32 @@
mod check; mod check;
use crate::check::check; use crate::check::check;
use ::services::{
use ::services::establish_connection; get_config,
use ::services::models::Scans; models::{Scans, Services},
use ::services::models::Services; schema::{
use ::services::schema::scans::dsl::scans; scans::{dsl::scans, installation},
use ::services::schema::scans::installation; services::dsl::*,
use ::services::schema::services::dsl::*; },
use diesel::associations::HasTable; };
use diesel::ExpressionMethods;
use diesel::QueryDsl;
use diesel::RunQueryDsl;
use diesel::SelectableHelper;
use dns_lookup::lookup_host;
use std::net::IpAddr; use std::net::IpAddr;
use diesel::{
associations::HasTable, Connection, ExpressionMethods, QueryDsl, RunQueryDsl, SelectableHelper,
SqliteConnection,
};
use dns_lookup::lookup_host;
#[tokio::main] #[tokio::main]
async fn main() { async fn main() {
let connection = &mut establish_connection(); let config = get_config();
let mut conn = SqliteConnection::establish(&config.database)
.unwrap_or_else(|_| panic!("Error connecting to database"));
for service in services for service in services
.limit(300) .limit(300)
.select(Services::as_select()) .select(Services::as_select())
.load(connection) .load(&mut conn)
.expect("Error loading services") .expect("Error loading services")
{ {
let result_ipv6 = match check(&service.url, Some(6)).await { let result_ipv6 = match check(&service.url, Some(6)).await {
@ -46,14 +50,14 @@ async fn main() {
result_ipv6: result_ipv6.clone(), result_ipv6: result_ipv6.clone(),
result_ipv4: result_ipv4.clone(), result_ipv4: result_ipv4.clone(),
}) })
.execute(connection) .execute(&mut conn)
.unwrap(); .unwrap();
let installation_scans = scans let installation_scans = scans
.filter(installation.eq(service.url.to_string())) .filter(installation.eq(service.url.to_string()))
.limit(100) .limit(100)
.select(Scans::as_select()) .select(Scans::as_select())
.load(connection) .load(&mut conn)
.unwrap(); .unwrap();
let ipv6_successes = installation_scans let ipv6_successes = installation_scans
@ -99,7 +103,7 @@ async fn main() {
address_ipv6.eq(&addr_ipv6.join(",")), address_ipv6.eq(&addr_ipv6.join(",")),
address_ipv4.eq(&addr_ipv4.join(",")), address_ipv4.eq(&addr_ipv4.join(",")),
)) ))
.execute(connection) .execute(&mut conn)
.unwrap(); .unwrap();
} }
} }