feat: add backbone of wasm

Author: jjeffcaii

Cargo.toml

@@ -54,8 +54,8 @@ moka = { version = "0.12", features = ["future", "sync"] }
 url = "2.5"
 regex = "1.10"
 hex = "0.4"
-strum = { version = "0.26", default-features = false, features = ["strum_macros", "derive"] }
-strum_macros = "0.26"
+strum = { version = "0.27", default-features = false, features = ["strum_macros", "derive"] }
+strum_macros = "0.27"
 deadpool = "0.12"
 socket2 = "0.5"
 mlua = { version = "0.10", features = ["luajit", "vendored", "serialize", "async", "macros", "send", "anyhow"] }
@@ -66,11 +66,16 @@ tokio-rustls = "0.26"
 httparse = "1.9"
 http = "1.2"
 string_cache = "0.8"
-rand = "0.8"
+rand = "0.9"
 resolv-conf = "0.7"
 bitflags = "2.6"
 sha2 = "0.10"
 
+wasmedge-sdk = "0.13.2"
+wasmedge-sys = "0.17.5"
+wasmedge-types = "0.4.4"
+wasmedge-macro = "0.6.1"
+
 [dev-dependencies]
 hex = "0.4"
 pretty_env_logger = "0.5"

README.md

@@ -2,14 +2,14 @@
 
 # ZeroDNS
 
-t![GitHub Actions Workflow Status](https://img.shields.io/github/actions/workflow/status/jjeffcaii/zerodns/rust.yml)
+![GitHub Actions Workflow Status](https://img.shields.io/github/actions/workflow/status/jjeffcaii/zerodns/rust.yml)
 [![Codecov](https://img.shields.io/codecov/c/github/jjeffcaii/zerodns)](https://app.codecov.io/gh/jjeffcaii/zerodns)
 [![Crates.io Version](https://img.shields.io/crates/v/zerodns)](https://crates.io/crates/zerodns)
 [![Crates.io Total Downloads](https://img.shields.io/crates/d/zerodns)](https://crates.io/crates/zerodns)
 ![GitHub Tag](https://img.shields.io/github/v/tag/jjeffcaii/zerodns)
 ![GitHub License](https://img.shields.io/github/license/jjeffcaii/zerodns)
 
-a DNS server in Rust, which is inspired from chinadns/dnsmasq.
+a flexible DNS server in Rust, which is inspired from chinadns/dnsmasq.
 
 > :construction_worker: WARNING: still in an active development!!!
 

justfile

@@ -6,6 +6,7 @@ use crate::{Error, Result};
 use hashbrown::HashMap;
 use moka::future::Cache;
 use once_cell::sync::Lazy;
+use rand::Rng;
 use smallvec::SmallVec;
 use std::net::Ipv4Addr;
 use std::time::Duration;
@@ -70,12 +71,7 @@ impl LookupCache {
             .opcode(OpCode::StandardQuery)
             .build();
 
-        let id = {
-            use rand::prelude::*;
-
-            let mut rng = thread_rng();
-            rng.gen_range(1..u16::MAX)
-        };
+        let id = rand::rng().random_range(1..u16::MAX);
 
         let req0 = Message::builder()
             .id(id)

src/client/lookup.rs

@@ -1,7 +1,6 @@
 use futures::StreamExt;
 use hashbrown::HashMap;
 use once_cell::sync::Lazy;
-use rand::{thread_rng, Rng};
 use socket2::{Domain, Protocol, Type};
 use std::fmt::{Display, Formatter};
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
@@ -202,8 +201,8 @@ impl MultiplexUdpClient {
         });
 
         let seq = {
-            let mut rng = thread_rng();
-            rng.gen_range(1..u16::MAX)
+            use rand::prelude::*;
+            rand::rng().random_range(1..u16::MAX)
         };
 
         Self {

src/client/udp.rs

@@ -72,9 +72,8 @@ pub(crate) async fn execute(sm: &ArgMatches) -> Result<()> {
     let req = {
         let mut bu = Message::builder()
             .id({
-                use rand::Rng;
-                let mut rng = rand::thread_rng();
-                rng.gen_range(1024..u16::MAX)
+                use rand::prelude::*;
+                rand::rng().random_range(1024..u16::MAX)
             })
             .flags(flags);
 

src/cmds/resolve.rs

@@ -323,10 +323,15 @@ struct LuaContext(*mut Context, *mut Message, *mut Option<Message>);
 
 impl UserData for LuaContext {
     fn add_fields<F: LuaUserDataFields<Self>>(fields: &mut F) {
-        fields.add_field_method_get("request", |lua, this| {
+        fields.add_field_method_get("request", |_lua, this| {
             let msg = LuaMessage(Clone::clone(unsafe { this.1.as_ref().unwrap() }));
             Ok(msg)
         });
+
+        fields.add_field_method_get("peer", |_lua, this| {
+            let ctx = unsafe { this.0.as_ref().unwrap() };
+            Ok(ctx.client_addr().to_string())
+        });
     }
 
     fn add_methods<M: LuaUserDataMethods<Self>>(methods: &mut M) {

src/filter/lua.rs

@@ -20,3 +20,4 @@ mod noop;
 mod proto;
 mod proxyby;
 mod registry;
+mod wasm;

src/filter/mod.rs

@@ -1,5 +1,6 @@
 use crate::protocol::Message;
 use crate::Result;
+use std::net::SocketAddr;
 
 #[derive(Debug, Copy, Clone, Default, Hash, PartialEq, Eq)]
 pub struct ContextFlags(u64);
@@ -13,6 +14,13 @@ bitflags! {
 #[derive(Debug, Default)]
 pub struct Context {
     pub flags: ContextFlags,
+    pub(crate) peer: Option<SocketAddr>,
+}
+
+impl Context {
+    pub fn client_addr(&self) -> SocketAddr {
+        self.peer.unwrap()
+    }
 }
 
 #[async_trait::async_trait]

src/filter/proto.rs

@@ -0,0 +1,66 @@
+use super::{Context, Filter, FilterFactory, Options};
+use crate::protocol::Message;
+use crate::Result;
+use async_trait::async_trait;
+use std::result::Result as StdResult;
+use wasmedge_sdk::{params, Module, Store, VmBuilder, WasmVal};
+
+struct WasmFilter {
+    next: Option<Box<dyn Filter>>,
+    module: Module,
+}
+
+#[async_trait]
+impl Filter for WasmFilter {
+    async fn handle(
+        &self,
+        ctx: &mut Context,
+        req: &mut Message,
+        res: &mut Option<Message>,
+    ) -> Result<()> {
+        let vm = {
+            let store = Store::new()?;
+            let vm = VmBuilder::new().with_store(store).build()?;
+            vm.register_module(Some("extern"), Clone::clone(&self.module))?
+        };
+
+        let res = vm.run_func(Some("extern"), "add", params!(1_i32, 2_i32))?;
+
+        todo!()
+    }
+
+    fn set_next(&mut self, next: Box<dyn Filter>) {
+        self.next.replace(next);
+    }
+}
+
+struct WasmFilterFactory {
+    module: Module,
+}
+
+impl FilterFactory for WasmFilterFactory {
+    type Item = WasmFilter;
+
+    fn get(&self) -> Result<Self::Item> {
+        Ok(WasmFilter {
+            next: None,
+            module: Clone::clone(&self.module),
+        })
+    }
+}
+
+impl TryFrom<&Options> for WasmFilterFactory {
+    type Error = anyhow::Error;
+
+    fn try_from(value: &Options) -> StdResult<Self, Self::Error> {
+        let path = value
+            .get("path")
+            .ok_or_else(|| anyhow!("the property of 'path' is required"))?
+            .as_str()
+            .ok_or_else(|| anyhow!("invalid 'path' format"))?;
+
+        let module = Module::from_file(None, path)?;
+
+        Ok(Self { module })
+    }
+}

src/filter/wasm.rs

@@ -123,6 +123,7 @@ mod tests {
             .unwrap();
 
         let mut ctx = Context::default();
+
         let res = h.handle(&mut ctx, &mut req).await;
         assert!(res.is_ok_and(|it| it.is_none()));
     }

src/handler/filtered.rs

@@ -4,6 +4,7 @@ use crate::filter::Context;
 use crate::handler::Handler;
 use crate::protocol::{Flags, Message, RCode};
 use crate::{Error as ZError, Result};
+use std::net::SocketAddr;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
@@ -85,18 +86,25 @@ fn convert_error_to_message(
 }
 
 #[inline]
-async fn handle_<H>(req: &Message, h: Arc<H>) -> Result<Message>
+async fn handle_<H>(peer: SocketAddr, req: &Message, h: Arc<H>) -> Result<Message>
 where
     H: Handler,
 {
     let mut req = Clone::clone(req);
     let mut ctx = Context::default();
+    ctx.peer.replace(peer);
+
     h.handle(&mut ctx, &mut req)
         .await?
         .ok_or_else(|| anyhow!(ZError::ResolveNothing))
 }
 
-pub(super) async fn handle<H, C>(req: Message, h: Arc<H>, cache: Option<Arc<C>>) -> (Message, bool)
+pub(super) async fn handle<H, C>(
+    peer: SocketAddr,
+    req: Message,
+    h: Arc<H>,
+    cache: Option<Arc<C>>,
+) -> (Message, bool)
 where
     H: Handler,
     C: LoadingCache,
@@ -106,7 +114,7 @@ where
     }
 
     let (res, cached) = match cache.as_deref() {
-        None => (handle_(&req, h).await, false),
+        None => (handle_(peer, &req, h).await, false),
         Some(lc) => {
             let cached = Arc::new(AtomicBool::new(true));
 
@@ -115,7 +123,7 @@ where
                 let cached = Clone::clone(&cached);
                 lc.try_get_with_fixed(req, move |req| {
                     cached.store(false, Ordering::SeqCst);
-                    async move { handle_(&req, h).await }
+                    async move { handle_(peer, &req, h).await }
                 })
                 .await
             };

src/server/helper.rs

@@ -90,7 +90,7 @@ where
             let req = next?;
             let handler = Clone::clone(&handler);
             let cache = Clone::clone(&cache);
-            let (res, cached) = super::helper::handle(req, handler, cache).await;
+            let (res, cached) = super::helper::handle(addr, req, handler, cache).await;
 
             if res.answer_count() > 0 {
                 for next in res.answers() {

src/server/tcp.rs

@@ -41,7 +41,7 @@ where
         h: Arc<H>,
         cache: Option<Arc<C>>,
     ) {
-        let (res, cached) = helper::handle(req, h, cache).await;
+        let (res, cached) = helper::handle(peer, req, h, cache).await;
 
         if res.answer_count() > 0 {
             for next in res.answers() {