mirror of
https://github.com/ZeusWPI/ZNS.git
synced 2024-11-22 05:41:11 +01:00
Requests uses sqlite now
This commit is contained in:
parent
ec0b93e7c0
commit
a27d8e6f12
11 changed files with 176 additions and 21 deletions
2
migrations/2024-03-03-220459_create_records/down.sql
Normal file
2
migrations/2024-03-03-220459_create_records/down.sql
Normal file
|
@ -0,0 +1,2 @@
|
|||
-- This file should undo anything in `up.sql`
|
||||
DROP TABLE records
|
11
migrations/2024-03-03-220459_create_records/up.sql
Normal file
11
migrations/2024-03-03-220459_create_records/up.sql
Normal 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
13
src/db/lib.rs
Normal 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
2
src/db/mod.rs
Normal file
|
@ -0,0 +1,2 @@
|
|||
pub mod models;
|
||||
pub mod lib;
|
47
src/db/models.rs
Normal file
47
src/db/models.rs
Normal 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)
|
||||
}
|
||||
}
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
85
src/main.rs
85
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<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?);
|
||||
|
|
|
@ -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
13
src/schema.rs
Normal 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,
|
||||
}
|
||||
}
|
|
@ -2,6 +2,7 @@
|
|||
#[derive(Debug, Clone)]
|
||||
pub enum Type {
|
||||
A = 1,
|
||||
AAAA = 28
|
||||
}
|
||||
|
||||
#[repr(u16)]
|
||||
|
|
Loading…
Reference in a new issue