From a27d8e6f125cb97fd308f66f894d0f4007d2fc2a Mon Sep 17 00:00:00 2001 From: Xander Bil Date: Tue, 5 Mar 2024 22:13:32 +0100 Subject: [PATCH] Requests uses sqlite now --- src/cache.rs => migration | 0 .../2024-03-03-220459_create_records/down.sql | 2 + .../2024-03-03-220459_create_records/up.sql | 11 +++ src/db/lib.rs | 13 +++ src/db/mod.rs | 2 + src/db/models.rs | 47 ++++++++++ src/errors.rs | 11 +++ src/main.rs | 85 +++++++++++++++---- src/parser.rs | 12 +-- src/schema.rs | 13 +++ src/structs.rs | 1 + 11 files changed, 176 insertions(+), 21 deletions(-) rename src/cache.rs => migration (100%) create mode 100644 migrations/2024-03-03-220459_create_records/down.sql create mode 100644 migrations/2024-03-03-220459_create_records/up.sql create mode 100644 src/db/lib.rs create mode 100644 src/db/mod.rs create mode 100644 src/db/models.rs create mode 100644 src/schema.rs diff --git a/src/cache.rs b/migration similarity index 100% rename from src/cache.rs rename to migration diff --git a/migrations/2024-03-03-220459_create_records/down.sql b/migrations/2024-03-03-220459_create_records/down.sql new file mode 100644 index 0000000..80cac6e --- /dev/null +++ b/migrations/2024-03-03-220459_create_records/down.sql @@ -0,0 +1,2 @@ +-- This file should undo anything in `up.sql` +DROP TABLE records diff --git a/migrations/2024-03-03-220459_create_records/up.sql b/migrations/2024-03-03-220459_create_records/up.sql new file mode 100644 index 0000000..be11487 --- /dev/null +++ b/migrations/2024-03-03-220459_create_records/up.sql @@ -0,0 +1,11 @@ +-- Your SQL goes here +CREATE TABLE records ( + name TEXT NOT NULL, + type INT NOT NULL, + class INT NOT NULL, + ttl INT NOT NULL, + rdlength INT NOT NULL, + rdata BLOB NOT NULL, + + PRIMARY KEY (name,type,class) +) diff --git a/src/db/lib.rs b/src/db/lib.rs new file mode 100644 index 0000000..9f2d667 --- /dev/null +++ b/src/db/lib.rs @@ -0,0 +1,13 @@ +use diesel::prelude::*; +use diesel::sqlite::SqliteConnection; +use dotenvy::dotenv; +use std::env; + + +pub fn establish_connection() -> SqliteConnection { + dotenv().ok(); + + let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set"); + SqliteConnection::establish(&database_url) + .unwrap_or_else(|_| panic!("Error connecting to {}", database_url)) +} diff --git a/src/db/mod.rs b/src/db/mod.rs new file mode 100644 index 0000000..5253185 --- /dev/null +++ b/src/db/mod.rs @@ -0,0 +1,2 @@ +pub mod models; +pub mod lib; diff --git a/src/db/models.rs b/src/db/models.rs new file mode 100644 index 0000000..80832a2 --- /dev/null +++ b/src/db/models.rs @@ -0,0 +1,47 @@ +use self::schema::records; +use diesel::prelude::*; + +mod schema { + diesel::table! { + records (name, _type, class) { + name -> Text, + #[sql_name = "type"] + _type -> Integer, + class -> Integer, + ttl -> Integer, + rdlength -> Integer, + rdata -> Binary, + } + } +} + +#[derive(Insertable, Queryable, Selectable)] +#[diesel(table_name = records)] +pub struct Record { + pub name: String, + pub _type: i32, + pub class: i32, + pub ttl: i32, + pub rdlength: i32, + pub rdata: Vec, +} + +impl Record { + pub fn get( + db: &mut SqliteConnection, + name: String, + _type: i32, + class: i32, + ) -> Result { + records::table.find((name, _type, class)).get_result(db) + } + + pub fn create( + db: &mut SqliteConnection, + new_record: Record, + ) -> Result { + diesel::insert_into(records::table) + .values(&new_record) + .execute(db) + } +} diff --git a/src/errors.rs b/src/errors.rs index 743ae6e..cb147c0 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -22,3 +22,14 @@ impl fmt::Display for ParseError { write!(f, "Parse Error for {}: {}", self.object, self.message) } } + +#[derive(Debug)] +pub struct DatabaseError { + pub message: String, +} + +impl fmt::Display for DatabaseError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "Database Error: {}", self.message) + } +} diff --git a/src/main.rs b/src/main.rs index 1242bb0..2720046 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,14 @@ use std::{error::Error, net::SocketAddr, sync::Arc}; +use db::{lib::establish_connection, models::Record}; +use errors::DatabaseError; use parser::FromBytes; -use structs::Message; +use structs::{Message, Question}; use tokio::net::UdpSocket; use crate::structs::{Class, Type, RR}; +mod db; mod errors; mod parser; mod structs; @@ -13,28 +16,80 @@ mod worker; const MAX_DATAGRAM_SIZE: usize = 40_96; -async fn create_query(message: Message) -> Message { - let mut response = message.clone(); - let ip = String::from("93.184.216.34"); - let rr = RR { - name: vec![String::from("example"), String::from("org")], - _type: Type::A, - class: Class::IN, - ttl: 4096, - rdlength: ip.len() as u16, - rdata: vec![1, 2, 3, 4], +async fn get_from_database(question: Question) -> Result { + let db_connection = &mut establish_connection(); + let record = Record::get( + db_connection, + question.qname.join("."), + question.qtype as i32, + question.qclass as i32, + ) + .map_err(|e| DatabaseError { + message: e.to_string(), + })?; + + Ok(RR { + name: record.name.split(".").map(str::to_string).collect(), + _type: Type::try_from(record._type as u16).map_err(|e| DatabaseError { message: e })?, + class: Class::try_from(record.class as u16).map_err(|e| DatabaseError { message: e })?, + ttl: record.ttl, + rdlength: record.rdlength as u16, + rdata: record.rdata, + }) +} + +async fn insert_into_database(rr: RR) -> Result<(), DatabaseError> { + let db_connection = &mut establish_connection(); + let record = Record { + name: rr.name.join("."), + _type: rr._type as i32, + class: rr.class as i32, + ttl: rr.ttl, + rdlength: rr.rdlength as i32, + rdata: rr.rdata, }; - response.header.flags |= 0b1000010110000000; - response.header.ancount = 1; - response.header.arcount = 0; - response.answer = Some(rr); + Record::create(db_connection, record).map_err(|e| DatabaseError { + message: e.to_string(), + })?; + + Ok(()) +} + +async fn create_query(message: Message) -> Message { + let mut response = message.clone(); + + let answer = get_from_database(message.question).await; + + match answer { + Ok(rr) => { + response.header.flags |= 0b1000010110000000; + response.header.ancount = 1; + response.header.arcount = 0; + response.answer = Some(rr) + } + Err(e) => { + response.header.flags |= 0b1000010110000011; + eprintln!("{}", e); + } + } response } #[tokio::main] async fn main() -> Result<(), Box> { + // insert_into_database(RR{ + // name: vec![String::from("example"),String::from("org")], + // _type: Type::A, + // class: Class::IN, + // ttl: 100, + // rdlength: 4, + // rdata: vec![1,2,3,4] + // }).await; + // + // return Ok(()); + let local_addr: SocketAddr = "127.0.0.1:8080".parse()?; let socket_shared = Arc::new(UdpSocket::bind(local_addr).await?); diff --git a/src/parser.rs b/src/parser.rs index 8fb6316..9180519 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -8,23 +8,23 @@ use crate::{ type Result = std::result::Result; impl TryFrom for Type { - type Error = (); //TODO: user better error + type Error = String; - fn try_from(value: u16) -> std::result::Result { + fn try_from(value: u16) -> std::result::Result { match value { x if x == Type::A as u16 => Ok(Type::A), - _ => Err(()), + _ => Err(format!("Invalid Type value: {}",value)), } } } impl TryFrom for Class { - type Error = (); //TODO: user better error + type Error = String; - fn try_from(value: u16) -> std::result::Result { + fn try_from(value: u16) -> std::result::Result { match value { x if x == Class::IN as u16 => Ok(Class::IN), - _ => Err(()), + _ => Err(format!("Invalid Class value: {}",value)), } } } diff --git a/src/schema.rs b/src/schema.rs new file mode 100644 index 0000000..4d63d37 --- /dev/null +++ b/src/schema.rs @@ -0,0 +1,13 @@ +// @generated automatically by Diesel CLI. + +diesel::table! { + records (name, type_, class) { + name -> Text, + #[sql_name = "type"] + type_ -> Integer, + class -> Integer, + ttl -> Integer, + rdlength -> Integer, + rdata -> Binary, + } +} diff --git a/src/structs.rs b/src/structs.rs index 06350f2..8751164 100644 --- a/src/structs.rs +++ b/src/structs.rs @@ -2,6 +2,7 @@ #[derive(Debug, Clone)] pub enum Type { A = 1, + AAAA = 28 } #[repr(u16)]