fix async dns client

migrations/2021-03-26-164945_create_users/down.sql

@@ -1,3 +1,4 @@
 -- This file should undo anything in `up.sql`
 DROP TABLE localuser;
 DROP TABLE user;
+DROP TABLE user_zone;

migrations/2021-03-26-164945_create_users/up.sql

@@ -10,3 +10,10 @@ CREATE TABLE user (
   id VARCHAR NOT NULL PRIMARY KEY,
   role TEXT CHECK(role IN ('admin', 'zoneadmin')) NOT NULL -- note: migrate to postgres so enum are actually a thing
 );
+
+CREATE TABLE user_zone (
+  user_id VARCHAR NOT NULL,
+  zone VARCHAR NOT NULL,
+  PRIMARY KEY(user_id, zone),
+  FOREIGN KEY(user_id) REFERENCES user(id)
+)

src/main.rs

@@ -4,16 +4,6 @@
 #[macro_use] extern crate rocket_contrib;
 #[macro_use] extern crate diesel;
 
-use trust_dns_client::client::AsyncClient;
-use trust_dns_client::tcp::TcpClientStream;
-use trust_dns_proto::xfer::dns_multiplexer::DnsMultiplexer;
-use trust_dns_proto::iocompat::AsyncIoTokioAsStd;
-use trust_dns_client::rr::dnssec::Signer;
-use tokio::net::TcpStream as TokioTcpStream;
-use tokio::task;
-
-use std::sync::{Arc, Mutex};
-
 mod models;
 mod config;
 mod schema;
@@ -26,22 +16,12 @@ use routes::zones::*;
 #[database("db")]
 pub struct DbConn(diesel::SqliteConnection);
 
-type DnsClient = Arc<Mutex<AsyncClient>>;
-
-
 #[launch]
 async fn rocket() -> rocket::Rocket {
     let app_config = config::load("config.toml".into());
     println!("{:#?}", app_config);
 
-    let (stream, handle) = TcpClientStream::<AsyncIoTokioAsStd<TokioTcpStream>>::new(app_config.dns.server);
-    let multiplexer = DnsMultiplexer::<_, Signer>::new(stream, handle, None);
-    let client = AsyncClient::connect(multiplexer);
-    let (client, bg) = client.await.expect("connection failed");
-    task::spawn(bg);
-
     rocket::ignite()
-        .manage(Arc::new(Mutex::new(client)))
         .manage(app_config)
         .attach(DbConn::fairing())
         .mount("/api/v1", routes![get_zone_records, create_auth_token, create_user])

src/models/dns.rs

@@ -1,10 +1,22 @@
 use std::net::{Ipv6Addr, Ipv4Addr};
 use std::fmt;
+use std::ops::{Deref, DerefMut};
+
+
+use rocket::{Request, State, http::Status, request::{FromParam, FromRequest, Outcome}};
 
 use serde::{Serialize, Deserialize};
-use trust_dns_client::serialize::binary::BinEncoder;
 
-use super::trust_dns_types;
+use tokio::{net::TcpStream as TokioTcpStream, task};
+
+use trust_dns_client::{client::AsyncClient, serialize::binary::BinEncoder, tcp::TcpClientStream};
+use trust_dns_proto::error::{ProtoError};
+use trust_dns_proto::iocompat::AsyncIoTokioAsStd;
+
+
+use super::trust_dns_types::{self, Name};
+use crate::config::Config;
+use crate::models::errors::make_500;
 
 
 #[derive(Deserialize, Serialize)]
@@ -238,3 +250,64 @@ impl From<trust_dns_types::Record> for Record {
         }
     }
 }
+
+pub struct AbsoluteName(Name);
+
+impl<'r> FromParam<'r> for AbsoluteName {
+    type Error = ProtoError;
+
+    fn from_param(param: &'r str) -> Result<Self, Self::Error> {
+        let mut name = Name::from_utf8(&param).unwrap();
+        if !name.is_fqdn() {
+            name.set_fqdn(true);
+        }
+        Ok(AbsoluteName(name))
+    }
+}
+
+impl Deref for AbsoluteName {
+    type Target = Name;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+pub struct DnsClient(AsyncClient);
+
+impl Deref for DnsClient {
+    type Target = AsyncClient;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for DnsClient {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+
+#[rocket::async_trait]
+impl<'r> FromRequest<'r> for DnsClient {
+    type Error = ();
+    async fn from_request(request: &'r Request<'_>) -> Outcome<Self, Self::Error> {
+        let config = try_outcome!(request.guard::<State<Config>>().await);
+        let (stream, handle) = TcpClientStream::<AsyncIoTokioAsStd<TokioTcpStream>>::new(config.dns.server);
+        let client = AsyncClient::with_timeout(
+            stream,
+            handle,
+            std::time::Duration::from_secs(5),
+            None);
+        let (client, bg) = match client.await {
+            Err(e) => {
+                println!("Failed to connect to DNS server {:#?}", e);
+                return Outcome::Failure((Status::InternalServerError, ()))
+            },
+            Ok(c) => c
+        };
+        task::spawn(bg);
+        Outcome::Success(DnsClient(client))
+    }
+}

src/routes/zones.rs

@@ -1,38 +1,33 @@
-use rocket::State;
 use rocket::http::Status;
 
 use rocket_contrib::json::Json;
 
 use trust_dns_client::client::ClientHandle;
 use trust_dns_client::op::{DnsResponse, ResponseCode};
-use trust_dns_client::rr::{DNSClass, Name, RecordType};
+use trust_dns_client::rr::{DNSClass, RecordType};
 
 use crate::models::dns;
 use crate::models::errors::{ErrorResponse, make_500};
 use crate::models::users::UserInfo;
-use crate::DnsClient;
 
 
 #[get("/zones/<zone>/records")]
 pub async fn get_zone_records(
-    client: State<'_, DnsClient>,
+    mut client: dns::DnsClient,
     user_info: Result<UserInfo, ErrorResponse>,
-    zone: String
+    zone: dns::AbsoluteName
 ) -> Result<Json<Vec<dns::Record>>, ErrorResponse> {
     println!("{:#?}", user_info?);
 
-    // TODO: Implement FromParam for Name
-    let name = Name::from_utf8(&zone).unwrap();
-
     let response: DnsResponse = {
-        let query = client.lock().unwrap().query(name.clone(), DNSClass::IN, RecordType::AXFR);
+        let query = client.query((*zone).clone(), DNSClass::IN, RecordType::AXFR);
         query.await.map_err(make_500)?
     };
 
     if response.response_code() != ResponseCode::NoError {
         return ErrorResponse::new(
             Status::NotFound,
-            format!("zone {} could not be found", name.to_utf8())
+            format!("zone {} could not be found", *zone)
         ).err()
     }
 

src/schema.rs

@@ -18,9 +18,20 @@ table! {
     }
 }
 
+table! {
+    use diesel::sql_types::*;
+
+    user_zone (user_id, zone) {
+        user_id -> Text,
+        zone -> Text,
+    }
+}
+
 joinable!(localuser -> user (user_id));
+joinable!(user_zone -> user (user_id));
 
 allow_tables_to_appear_in_same_query!(
     localuser,
     user,
+    user_zone,
 );