mirror of
https://github.com/ZeusWPI/ZNS.git
synced 2024-10-30 05:24:26 +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)
|
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 std::{error::Error, net::SocketAddr, sync::Arc};
|
||||||
|
|
||||||
|
use db::{lib::establish_connection, models::Record};
|
||||||
|
use errors::DatabaseError;
|
||||||
use parser::FromBytes;
|
use parser::FromBytes;
|
||||||
use structs::Message;
|
use structs::{Message, Question};
|
||||||
use tokio::net::UdpSocket;
|
use tokio::net::UdpSocket;
|
||||||
|
|
||||||
use crate::structs::{Class, Type, RR};
|
use crate::structs::{Class, Type, RR};
|
||||||
|
|
||||||
|
mod db;
|
||||||
mod errors;
|
mod errors;
|
||||||
mod parser;
|
mod parser;
|
||||||
mod structs;
|
mod structs;
|
||||||
|
@ -13,28 +16,80 @@ mod worker;
|
||||||
|
|
||||||
const MAX_DATAGRAM_SIZE: usize = 40_96;
|
const MAX_DATAGRAM_SIZE: usize = 40_96;
|
||||||
|
|
||||||
async fn create_query(message: Message) -> Message {
|
async fn get_from_database(question: Question) -> Result<RR, DatabaseError> {
|
||||||
let mut response = message.clone();
|
let db_connection = &mut establish_connection();
|
||||||
let ip = String::from("93.184.216.34");
|
let record = Record::get(
|
||||||
let rr = RR {
|
db_connection,
|
||||||
name: vec![String::from("example"), String::from("org")],
|
question.qname.join("."),
|
||||||
_type: Type::A,
|
question.qtype as i32,
|
||||||
class: Class::IN,
|
question.qclass as i32,
|
||||||
ttl: 4096,
|
)
|
||||||
rdlength: ip.len() as u16,
|
.map_err(|e| DatabaseError {
|
||||||
rdata: vec![1, 2, 3, 4],
|
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;
|
Record::create(db_connection, record).map_err(|e| DatabaseError {
|
||||||
response.header.ancount = 1;
|
message: e.to_string(),
|
||||||
response.header.arcount = 0;
|
})?;
|
||||||
response.answer = Some(rr);
|
|
||||||
|
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
|
response
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::main]
|
#[tokio::main]
|
||||||
async fn main() -> Result<(), Box<dyn Error>> {
|
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 local_addr: SocketAddr = "127.0.0.1:8080".parse()?;
|
||||||
|
|
||||||
let socket_shared = Arc::new(UdpSocket::bind(local_addr).await?);
|
let socket_shared = Arc::new(UdpSocket::bind(local_addr).await?);
|
||||||
|
|
|
@ -8,23 +8,23 @@ use crate::{
|
||||||
type Result<T> = std::result::Result<T, ParseError>;
|
type Result<T> = std::result::Result<T, ParseError>;
|
||||||
|
|
||||||
impl TryFrom<u16> for Type {
|
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 {
|
match value {
|
||||||
x if x == Type::A as u16 => Ok(Type::A),
|
x if x == Type::A as u16 => Ok(Type::A),
|
||||||
_ => Err(()),
|
_ => Err(format!("Invalid Type value: {}",value)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl TryFrom<u16> for Class {
|
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 {
|
match value {
|
||||||
x if x == Class::IN as u16 => Ok(Class::IN),
|
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)]
|
#[derive(Debug, Clone)]
|
||||||
pub enum Type {
|
pub enum Type {
|
||||||
A = 1,
|
A = 1,
|
||||||
|
AAAA = 28
|
||||||
}
|
}
|
||||||
|
|
||||||
#[repr(u16)]
|
#[repr(u16)]
|
||||||
|
|
Loading…
Reference in a new issue