add dns server to embassy-os (#1572)

* add dns server to embassy-os

* fix initialization

* multiple ip addresses

backend/src/net/dns.rs

@@ -0,0 +1,171 @@
+use std::borrow::Borrow;
+use std::collections::{BTreeMap, BTreeSet};
+use std::net::{Ipv4Addr, SocketAddr};
+use std::sync::Arc;
+use std::time::Duration;
+
+use futures::TryFutureExt;
+use helpers::NonDetachingJoinHandle;
+use models::PackageId;
+use tokio::net::{TcpListener, UdpSocket};
+use tokio::sync::RwLock;
+use trust_dns_server::authority::MessageResponseBuilder;
+use trust_dns_server::client::op::{Header, ResponseCode};
+use trust_dns_server::client::rr::{Name, Record, RecordType};
+use trust_dns_server::server::{Request, RequestHandler, ResponseHandler, ResponseInfo};
+use trust_dns_server::ServerFuture;
+
+use crate::net::mdns::resolve_mdns;
+use crate::{Error, ErrorKind, ResultExt};
+
+pub struct DnsController {
+    services: Arc<RwLock<BTreeMap<PackageId, BTreeSet<Ipv4Addr>>>>,
+    #[allow(dead_code)]
+    dns_server: NonDetachingJoinHandle<Result<(), Error>>,
+}
+
+struct Resolver {
+    services: Arc<RwLock<BTreeMap<PackageId, BTreeSet<Ipv4Addr>>>>,
+}
+impl Resolver {
+    async fn resolve(&self, name: &Name) -> Option<Vec<Ipv4Addr>> {
+        match name.iter().next_back() {
+            Some(b"local") => match resolve_mdns(&format!(
+                "{}.local",
+                name.iter()
+                    .rev()
+                    .skip(1)
+                    .next()
+                    .and_then(|v| std::str::from_utf8(v).ok())
+                    .unwrap_or_default()
+            ))
+            .await
+            {
+                Ok(ip) => Some(vec![ip]),
+                Err(e) => {
+                    tracing::error!("{}", e);
+                    tracing::debug!("{:?}", e);
+                    None
+                }
+            },
+            Some(b"embassy") => {
+                if let Some(pkg) = name.iter().rev().skip(1).next() {
+                    if let Some(ip) = self
+                        .services
+                        .read()
+                        .await
+                        .get(std::str::from_utf8(pkg).unwrap_or_default())
+                    {
+                        Some(ip.iter().copied().collect())
+                    } else {
+                        None
+                    }
+                } else {
+                    None
+                }
+            }
+            _ => None,
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl RequestHandler for Resolver {
+    async fn handle_request<R: ResponseHandler>(
+        &self,
+        request: &Request,
+        mut response_handle: R,
+    ) -> ResponseInfo {
+        let query = request.request_info().query;
+        if let Some(ip) = self.resolve(query.name().borrow()).await {
+            if query.query_type() != RecordType::A {
+                tracing::warn!("Non A-Record requested for {}", query.name());
+            }
+            response_handle
+                .send_response(
+                    MessageResponseBuilder::from_message_request(&*request).build(
+                        Header::response_from_request(request.header()),
+                        &ip.into_iter()
+                            .map(|ip| {
+                                Record::from_rdata(
+                                    request.request_info().query.name().to_owned().into(),
+                                    0,
+                                    trust_dns_server::client::rr::RData::A(ip),
+                                )
+                            })
+                            .collect::<Vec<_>>(),
+                        [],
+                        [],
+                        [],
+                    ),
+                )
+                .await
+        } else {
+            let mut res = Header::response_from_request(request.header());
+            res.set_response_code(ResponseCode::NXDomain);
+            response_handle
+                .send_response(
+                    MessageResponseBuilder::from_message_request(&*request).build(
+                        res.into(),
+                        [],
+                        [],
+                        [],
+                        [],
+                    ),
+                )
+                .await
+        }
+        .unwrap_or_else(|e| {
+            tracing::error!("{}", e);
+            tracing::debug!("{:?}", e);
+            let mut res = Header::response_from_request(request.header());
+            res.set_response_code(ResponseCode::ServFail);
+            res.into()
+        })
+    }
+}
+
+impl DnsController {
+    pub async fn init(bind: &[SocketAddr]) -> Result<Self, Error> {
+        let services = Arc::new(RwLock::new(BTreeMap::new()));
+
+        let mut server = ServerFuture::new(Resolver {
+            services: services.clone(),
+        });
+        server.register_listener(
+            TcpListener::bind(bind)
+                .await
+                .with_kind(ErrorKind::Network)?,
+            Duration::from_secs(30),
+        );
+        server.register_socket(UdpSocket::bind(bind).await.with_kind(ErrorKind::Network)?);
+
+        let dns_server = tokio::spawn(
+            server
+                .block_until_done()
+                .map_err(|e| Error::new(e, ErrorKind::Network)),
+        )
+        .into();
+
+        Ok(Self {
+            services,
+            dns_server,
+        })
+    }
+
+    pub async fn add(&self, pkg_id: &PackageId, ip: Ipv4Addr) {
+        let mut writable = self.services.write().await;
+        let mut ips = writable.remove(pkg_id).unwrap_or_default();
+        ips.insert(ip);
+        writable.insert(pkg_id.clone(), ips);
+    }
+
+    pub async fn remove(&self, pkg_id: &PackageId, ip: Ipv4Addr) {
+        let mut writable = self.services.write().await;
+        let mut ips = writable.remove(pkg_id).unwrap_or_default();
+        ips.remove(&ip);
+        if !ips.is_empty() {
+            writable.insert(pkg_id.clone(), ips);
+        }
+    }
+}

backend/src/net/mdns.rs

@@ -1,5 +1,5 @@
 use std::collections::BTreeMap;
-use std::net::IpAddr;
+use std::net::Ipv4Addr;
 
 use avahi_sys::{
     self, avahi_client_errno, avahi_entry_group_add_service, avahi_entry_group_commit,
@@ -17,7 +17,7 @@ use crate::s9pk::manifest::PackageId;
 use crate::util::Invoke;
 use crate::Error;
 
-pub async fn resolve_mdns(hostname: &str) -> Result<IpAddr, Error> {
+pub async fn resolve_mdns(hostname: &str) -> Result<Ipv4Addr, Error> {
     Ok(String::from_utf8(
         Command::new("avahi-resolve-host-name")
             .arg("-4")

backend/src/net/mod.rs

@@ -14,11 +14,13 @@ use self::mdns::MdnsController;
 use self::nginx::NginxController;
 use self::ssl::SslManager;
 use self::tor::TorController;
+use crate::net::dns::DnsController;
 use crate::net::interface::TorConfig;
 use crate::net::nginx::InterfaceMetadata;
 use crate::s9pk::manifest::PackageId;
 use crate::Error;
 
+pub mod dns;
 pub mod interface;
 #[cfg(feature = "avahi")]
 pub mod mdns;
@@ -45,6 +47,7 @@ pub struct NetController {
     pub mdns: MdnsController,
     pub nginx: NginxController,
     pub ssl: SslManager,
+    pub dns: DnsController,
 }
 impl NetController {
     #[instrument(skip(db))]
@@ -52,6 +55,7 @@ impl NetController {
         embassyd_addr: SocketAddr,
         embassyd_tor_key: TorSecretKeyV3,
         tor_control: SocketAddr,
+        dns_bind: &[SocketAddr],
         db: SqlitePool,
         import_root_ca: Option<(PKey<Private>, X509)>,
     ) -> Result<Self, Error> {
@@ -65,6 +69,7 @@ impl NetController {
             mdns: MdnsController::init(),
             nginx: NginxController::init(PathBuf::from("/etc/nginx"), &ssl).await?,
             ssl,
+            dns: DnsController::init(dns_bind).await?,
         })
     }
 
@@ -92,7 +97,7 @@ impl NetController {
                 Some(cfg) => Some((i.0, cfg, i.2)),
             })
             .collect::<Vec<(InterfaceId, TorConfig, TorSecretKeyV3)>>();
-        let (tor_res, _, nginx_res) = tokio::join!(
+        let (tor_res, _, nginx_res, _) = tokio::join!(
             self.tor.add(pkg_id, ip, interfaces_tor),
             {
                 #[cfg(feature = "avahi")]
@@ -123,7 +128,8 @@ impl NetController {
                         )),
                     });
                 self.nginx.add(&self.ssl, pkg_id.clone(), ip, interfaces)
-            }
+            },
+            self.dns.add(pkg_id, ip),
         );
         tor_res?;
         nginx_res?;
@@ -135,9 +141,10 @@ impl NetController {
     pub async fn remove<I: IntoIterator<Item = InterfaceId> + Clone>(
         &self,
         pkg_id: &PackageId,
+        ip: Ipv4Addr,
         interfaces: I,
     ) -> Result<(), Error> {
-        let (tor_res, _, nginx_res) = tokio::join!(
+        let (tor_res, _, nginx_res, _) = tokio::join!(
             self.tor.remove(pkg_id, interfaces.clone()),
             {
                 #[cfg(feature = "avahi")]
@@ -146,7 +153,8 @@ impl NetController {
                 let mdns_fut = futures::future::ready(());
                 mdns_fut
             },
-            self.nginx.remove(pkg_id)
+            self.nginx.remove(pkg_id),
+            self.dns.remove(pkg_id, ip),
         );
         tor_res?;
         nginx_res?;