From 08d4ca82f5ca18406674b6fe86f3914671e98c54 Mon Sep 17 00:00:00 2001 From: Xander Bil Date: Tue, 26 Mar 2024 23:22:12 +0100 Subject: [PATCH] Add delete support --- migration | 0 .../2024-03-03-220459_create_records/up.sql | 2 +- src/db/models.rs | 78 +++++++++++++++---- src/parser.rs | 11 +-- src/resolver.rs | 36 ++++++--- src/schema.rs | 13 ++++ src/structs.rs | 5 +- 7 files changed, 110 insertions(+), 35 deletions(-) delete mode 100644 migration create mode 100644 src/schema.rs diff --git a/migration b/migration deleted file mode 100644 index e69de29..0000000 diff --git a/migrations/2024-03-03-220459_create_records/up.sql b/migrations/2024-03-03-220459_create_records/up.sql index be11487..67e9c48 100644 --- a/migrations/2024-03-03-220459_create_records/up.sql +++ b/migrations/2024-03-03-220459_create_records/up.sql @@ -7,5 +7,5 @@ CREATE TABLE records ( rdlength INT NOT NULL, rdata BLOB NOT NULL, - PRIMARY KEY (name,type,class) + PRIMARY KEY (name,type,class,rdlength,rdata) ) diff --git a/src/db/models.rs b/src/db/models.rs index f9cf2c6..34497e9 100644 --- a/src/db/models.rs +++ b/src/db/models.rs @@ -2,15 +2,15 @@ use crate::{ errors::DatabaseError, structs::{Class, Question, Type, RR}, }; -use diesel::prelude::*; +use diesel::{dsl, prelude::*}; -use self::schema::records; +use self::schema::records::{self}; use super::lib::establish_connection; mod schema { diesel::table! { - records (name, _type, class) { + records (name, _type, class, rdlength, rdata) { name -> Text, #[sql_name = "type"] _type -> Integer, @@ -39,8 +39,14 @@ impl Record { name: String, _type: i32, class: i32, - ) -> Result { - records::table.find((name, _type, class)).get_result(db) + ) -> Result, diesel::result::Error> { + records::table + .filter( + records::name + .eq(name) + .and(records::_type.eq(_type).and(records::class.eq(class))), + ) + .get_results(db) } pub fn create( @@ -51,6 +57,28 @@ impl Record { .values(&new_record) .execute(db) } + + pub fn delete( + db: &mut SqliteConnection, + name: String, + _type: Option, + class: i32, + rdata: Option>, + ) -> Result { + let mut query = diesel::delete(records::table).into_boxed(); + + query = query.filter(records::name.eq(name).and(records::class.eq(class))); + + if let Some(_type) = _type { + query = query.filter(records::_type.eq(_type)); + } + + if let Some(rdata) = rdata { + query = query.filter(records::rdata.eq(rdata)); + } + + query.execute(db) + } } pub async fn insert_into_database(rr: RR) -> Result<(), DatabaseError> { @@ -71,9 +99,9 @@ pub async fn insert_into_database(rr: RR) -> Result<(), DatabaseError> { Ok(()) } -pub async fn get_from_database(question: &Question) -> Result { +pub async fn get_from_database(question: &Question) -> Result, DatabaseError> { let db_connection = &mut establish_connection(); - let record = Record::get( + let records = Record::get( db_connection, question.qname.join("."), question.qtype.clone() as i32, @@ -83,12 +111,32 @@ pub async fn get_from_database(question: &Question) -> Result 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, - }) + Ok(records + .into_iter() + .filter_map(|record| { + Some(RR { + name: record.name.split(".").map(str::to_string).collect(), + _type: Type::try_from(record._type as u16) + .map_err(|e| DatabaseError { message: e }) + .ok()?, + class: Class::try_from(record.class as u16) + .map_err(|e| DatabaseError { message: e }) + .ok()?, + ttl: record.ttl, + rdlength: record.rdlength as u16, + rdata: record.rdata, + }) + }) + .collect()) +} + +//TODO: cleanup models +pub async fn delete_from_database( + name: Vec, + _type: Option, + class: Class, + rdata: Option>, +) { + let db_connection = &mut establish_connection(); + let _ = Record::delete(db_connection, name.join("."), _type.map(|f| f as i32), class as i32, rdata); } diff --git a/src/parser.rs b/src/parser.rs index ef73d46..1c7f7cd 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -12,9 +12,11 @@ impl TryFrom for Type { fn try_from(value: u16) -> std::result::Result { match value { + //TODO: clean this up x if x == Type::A as u16 => Ok(Type::A), x if x == Type::OPT as u16 => Ok(Type::OPT), x if x == Type::SOA as u16 => Ok(Type::SOA), + x if x == Type::ANY as u16 => Ok(Type::ANY), _ => Err(format!("Invalid Type value: {}", value)), } } @@ -26,6 +28,8 @@ impl TryFrom for Class { fn try_from(value: u16) -> std::result::Result { match value { x if x == Class::IN as u16 => Ok(Class::IN), + x if x == Class::ANY as u16 => Ok(Class::ANY), + x if x == Class::NONE as u16 => Ok(Class::NONE), _ => Err(format!("Invalid Class value: {}", value)), } } @@ -90,8 +94,7 @@ impl Type { }) } } - Type::SOA => todo!(), - Type::OPT => todo!(), + _ => todo!() } } pub fn from_data(&self, bytes: &[u8]) -> Result { @@ -107,8 +110,7 @@ impl Type { }) } } - Type::SOA => todo!(), - Type::OPT => todo!(), + _ => todo!() } } } @@ -246,7 +248,6 @@ impl FromBytes for Question { impl FromBytes for RR { fn from_bytes(bytes: &[u8], i: &mut usize) -> Result { let name = LabelString::from_bytes(bytes, i)?; - println!("{:#?}", name); if bytes.len() - *i < size_of::() + size_of::() + 6 { Err(ParseError { object: String::from("RR"), diff --git a/src/resolver.rs b/src/resolver.rs index 823535d..89137a2 100644 --- a/src/resolver.rs +++ b/src/resolver.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use tokio::net::UdpSocket; -use crate::db::models::{get_from_database, insert_into_database}; +use crate::db::models::{delete_from_database, get_from_database, insert_into_database}; use crate::errors::ParseError; use crate::parser::FromBytes; use crate::structs::{Class, Header, Message, Opcode, Type, RCODE}; @@ -25,13 +25,13 @@ async fn handle_query(message: Message) -> Message { response.header.arcount = 0; //TODO: fix this, handle unknown class values for question in message.question { - let answer = get_from_database(&question).await; + let answers = get_from_database(&question).await; - match answer { - Ok(rr) => { + match answers { + Ok(rrs) => { response.header.flags = set_response_flags(response.header.flags, RCODE::NOERROR); response.header.ancount = 1; - response.answer = vec![rr] + response.answer = rrs } Err(e) => { response.header.flags = set_response_flags(response.header.flags, RCODE::NXDOMAIN); @@ -81,22 +81,34 @@ async fn handle_update(message: Message) -> Message { if (rr.class == Class::ANY && (rr.ttl != 0 || rr.rdlength != 0)) || (rr.class == Class::NONE && rr.ttl != 0) - || rr.class != zone.qclass + || ![Class::NONE, Class::ANY, zone.qclass.clone()].contains(&rr.class) { response.header.flags = set_response_flags(response.header.flags, RCODE::FORMERR); return response; } } + //FIX: with nsupdate delete, I get `dns_request_getresponse: unexpected end of input` for rr in message.authority { if rr.class == zone.qclass { - insert_into_database(rr).await; + let _ = insert_into_database(rr).await; } else if rr.class == Class::ANY { - response.header.flags = set_response_flags(response.header.flags, RCODE::NOTIMP); - return response; - } else if rr.class == Class::ANY { - response.header.flags = set_response_flags(response.header.flags, RCODE::NOTIMP); - return response; + if rr._type == Type::ANY { + if rr.name == zone.qname { + response.header.flags = + set_response_flags(response.header.flags, RCODE::NOTIMP); + return response; + } else { + delete_from_database(rr.name, None, Class::IN, None).await; + } + } else { + delete_from_database(rr.name, Some(rr._type), Class::IN, None).await; + } + } else if rr.class == Class::NONE { + if rr._type == Type::SOA { + continue; + } + delete_from_database(rr.name, Some(rr._type),Class::IN, Some(rr.rdata)).await; } } diff --git a/src/schema.rs b/src/schema.rs new file mode 100644 index 0000000..d03d571 --- /dev/null +++ b/src/schema.rs @@ -0,0 +1,13 @@ +// @generated automatically by Diesel CLI. + +diesel::table! { + records (name, type_, class, rdlength, rdata) { + 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 3e85e15..569e183 100644 --- a/src/structs.rs +++ b/src/structs.rs @@ -1,11 +1,12 @@ use serde::Deserialize; #[repr(u16)] -#[derive(Debug, Clone, Deserialize)] +#[derive(Debug, Clone, Deserialize, PartialEq)] pub enum Type { A = 1, SOA = 6, - OPT = 41 + OPT = 41, + ANY = 255 } #[repr(u16)]