10
0
Fork 0
mirror of https://github.com/ZeusWPI/ZNS.git synced 2024-11-21 21:41:10 +01:00

Requests uses sqlite now

This commit is contained in:
Xander Bil 2024-03-05 22:13:32 +01:00
parent ec0b93e7c0
commit a27d8e6f12
No known key found for this signature in database
GPG key ID: EC9706B54A278598
11 changed files with 176 additions and 21 deletions

View file

@ -0,0 +1,2 @@
-- This file should undo anything in `up.sql`
DROP TABLE records

View file

@ -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)
)

13
src/db/lib.rs Normal file
View file

@ -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))
}

2
src/db/mod.rs Normal file
View file

@ -0,0 +1,2 @@
pub mod models;
pub mod lib;

47
src/db/models.rs Normal file
View file

@ -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<u8>,
}
impl Record {
pub fn get(
db: &mut SqliteConnection,
name: String,
_type: i32,
class: i32,
) -> Result<Record, diesel::result::Error> {
records::table.find((name, _type, class)).get_result(db)
}
pub fn create(
db: &mut SqliteConnection,
new_record: Record,
) -> Result<usize, diesel::result::Error> {
diesel::insert_into(records::table)
.values(&new_record)
.execute(db)
}
}

View file

@ -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)
}
}

View file

@ -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<RR, DatabaseError> {
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<dyn Error>> {
// 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?);

View file

@ -8,23 +8,23 @@ use crate::{
type Result<T> = std::result::Result<T, ParseError>;
impl TryFrom<u16> for Type {
type Error = (); //TODO: user better error
type Error = String;
fn try_from(value: u16) -> std::result::Result<Self, ()> {
fn try_from(value: u16) -> std::result::Result<Self, String> {
match value {
x if x == Type::A as u16 => Ok(Type::A),
_ => Err(()),
_ => Err(format!("Invalid Type value: {}",value)),
}
}
}
impl TryFrom<u16> for Class {
type Error = (); //TODO: user better error
type Error = String;
fn try_from(value: u16) -> std::result::Result<Self, ()> {
fn try_from(value: u16) -> std::result::Result<Self, String> {
match value {
x if x == Class::IN as u16 => Ok(Class::IN),
_ => Err(()),
_ => Err(format!("Invalid Class value: {}",value)),
}
}
}

13
src/schema.rs Normal file
View file

@ -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,
}
}

View file

@ -2,6 +2,7 @@
#[derive(Debug, Clone)]
pub enum Type {
A = 1,
AAAA = 28
}
#[repr(u16)]