Much progress

src/handler.rs

@@ -0,0 +1,215 @@
+use std::net::{IpAddr, Ipv4Addr, SocketAddr};
+use std::future::Future;
+use std::path::{Path, PathBuf};
+use std::sync::Arc;
+use async_lock::Mutex;
+use async_tftp::async_trait;
+use async_tftp::packet::Error;
+use async_tftp::server::handlers::{DirHandler, DirHandlerMode};
+use futures::{AsyncRead, AsyncWrite, TryStreamExt};
+use mlua::{FromLua, Lua, UserData, UserDataFields, Value};
+
+#[derive(Clone)]
+pub struct Handler {
+    lua: Arc<async_lock::Mutex<mlua::Lua>>,
+    call_key: Arc<mlua::RegistryKey>,
+    dir_handler: Arc<DirHandler>,
+    http: reqwest::Client,
+}
+
+#[derive(Debug)]
+pub enum Resource {
+    Http(String), // Parameter is URL
+    File(String), // Parameter is content
+    Data(Vec<u8>),
+    Error(Error)
+}
+
+impl Resource {
+}
+
+impl UserData for Resource {
+
+}
+
+impl<'lua> FromLua<'lua> for Resource {
+    fn from_lua(value: Value<'lua>, _lua: &'lua Lua) -> mlua::Result<Self> {
+        value.as_userdata().ok_or(mlua::Error::UserDataTypeMismatch).and_then(|value| value.take())
+    }
+}
+
+#[derive(Clone, Debug)]
+struct Client {
+    address: SocketAddr,
+    for_write: bool,
+}
+
+impl UserData for Client {
+    fn add_fields<'lua, F: UserDataFields<'lua, Self>>(fields: &mut F) {
+        fields.add_field_method_get("mode", |_lua, client| Ok(if client.for_write { "w" } else { "r" }));
+        fields.add_field_method_get("address", |lua, client| lua.create_any_userdata(client.address.ip()));
+    }
+}
+
+impl Handler {
+    pub fn new(srv_path: impl AsRef<Path>) -> Result<Self, anyhow::Error> {
+        let lua = mlua::Lua::new();
+
+        let http = reqwest::Client::builder()
+            .user_agent(concat!("sftpd/", env!("CARGO_PKG_VERSION")))
+            .build()?;
+
+        lua.register_userdata_type::<IpAddr>(|registry| {
+            registry.add_field_method_get("version", |_, ip| Ok(if ip.is_ipv4() { 4} else {6}));
+        })?;
+
+        {
+            // prepare resource types...
+            let resources = lua.create_table()?;
+            resources.set("HTTP", lua.create_function(|_lua, url: String| Ok(Resource::Http(url)))?)?;
+            resources.set("FILE", lua.create_function(|_lua, path: String| Ok(Resource::File(path)))?)?;
+            resources.set("DATA", lua.create_function(|_lua, url: mlua::String| Ok(Resource::Data(url.as_bytes().to_vec())))?)?;
+
+            let err_tbl = lua.create_table()?;
+            let err_mtbl = lua.create_table()?;
+            err_tbl.set_metatable(Some(err_mtbl.clone()));
+
+            let err_fn = lua.create_function(|_, msg: String| Ok(Resource::Error(Error::Msg(msg))))?;
+            err_mtbl.set("__call", err_fn.clone())?;
+            err_tbl.set("FileNotFound", Resource::Error(Error::FileNotFound))?;
+            err_tbl.set("Unknown", Resource::Error(Error::UnknownError))?;
+            err_tbl.set("PermissionDenied", Resource::Error(Error::PermissionDenied))?;
+            err_tbl.set("DiskFull", Resource::Error(Error::DiskFull))?;
+            err_tbl.set("IllegalOperation", Resource::Error(Error::IllegalOperation))?;
+            err_tbl.set("FileAlreadyExists", Resource::Error(Error::FileAlreadyExists))?;
+            err_tbl.set("NoSuchUser", Resource::Error(Error::NoSuchUser))?;
+            err_tbl.set("Message", err_fn)?;
+
+            resources.set("ERROR", err_tbl)?;
+            lua.globals().set("resource", resources)?;
+            lua.globals().set("state", lua.create_table()?)?;
+
+            //
+
+        }
+        let handler_fn = lua.create_registry_value(0)?;
+
+        Ok(Self {
+            lua: Arc::new(Mutex::new(lua)),
+            call_key: Arc::new(handler_fn),
+            http,
+            dir_handler: Arc::new(DirHandler::new(srv_path, DirHandlerMode::ReadWrite)?),
+        })
+    }
+
+    pub fn load_script(&mut self, data: PathBuf) -> impl Future<Output=anyhow::Result<()>> {
+        let lua = self.lua.clone();
+        let key = self.call_key.clone();
+        async move {
+            let lua = lua.lock_arc().await;
+            let chunk = lua.load(data);
+            let script_fn = chunk.into_function()?;
+
+            // Prepare a fake client to determine whether the script should just be run in full for each request
+            // or if it returns a function
+            let client = Client {
+                address: SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0),
+                for_write: false,
+            };
+
+            let result: mlua::Value = script_fn.call_async(("", client, None as Option<u64>)).await?;
+            if result.is_function() {
+                lua.replace_registry_value(&*key, result)?
+            } else if Resource::from_lua(result, &*lua).is_ok() {
+                // We must run the script, it seems
+                lua.replace_registry_value(&*key, script_fn)?
+            }
+
+            Ok(())
+        }
+    }
+
+    async fn call_handler_int(lua: Arc<Mutex<mlua::Lua>>, key: Arc<mlua::RegistryKey>, path: String, client: Client, size: Option<u64>) -> Result<Resource, Error> {
+        lua.lock().await
+            .registry_value::<mlua::Function>(&*key).map_err(|_| Error::UnknownError)?
+            .call_async((path, client, size)).await
+            .map_err(|_| Error::UnknownError)
+
+    }
+
+    async fn call_handler(&mut self, path: String, client: Client, size: Option<u64>) -> Result<Resource, Error> {
+
+        let handle = tokio::task::spawn_local(Self::call_handler_int(self.lua.clone(), self.call_key.clone(), path, client, size));
+        handle.await.map_err(|err| Error::Msg(err.to_string()))
+            .and_then(|x| x)
+    }
+}
+
+#[async_trait]
+impl async_tftp::server::Handler for Handler {
+    type Reader = Box<dyn AsyncRead + Send + Unpin + 'static>;
+    type Writer = Box<dyn AsyncWrite + Send + Unpin + 'static>;
+
+    async fn read_req_open(&mut self, client: &SocketAddr, path: &Path) -> Result<(Self::Reader, Option<u64>), Error> {
+        let lua_client = Client {
+            address: client.clone(),
+            for_write: false,
+        };
+        let resource: Resource = self.call_handler(path.to_str().ok_or(Error::FileNotFound)?.to_owned(), lua_client, None).await?;
+
+        match resource {
+            Resource::Http(url) => {
+                // TODO: Add headers describing client
+                let req = self.http.get(url).send().await.map_err(|err| Error::Msg(err.to_string()))?;
+                let size = req.content_length();
+                let stream = req.bytes_stream()
+                    .map_err(|e| futures::io::Error::new(futures::io::ErrorKind::Other, e))
+                    .into_async_read();
+
+                Ok((Box::new(stream), size))
+            }
+            Resource::File(path) => {
+                let (rdr, size) = (*self.dir_handler).clone().read_req_open(client, Path::new(path.as_str())).await?;
+                Ok((Box::new(rdr), size))
+            }
+            Resource::Data(data) => {
+                let len = data.len() as u64;
+                Ok((
+                    Box::new(futures::io::Cursor::new(data)),
+                    Some(len),
+                ))
+            }
+            Resource::Error(err) => { return Err(err) }
+        }
+    }
+
+    async fn write_req_open(&mut self, _client: &SocketAddr, _path: &Path, _size: Option<u64>) -> Result<Self::Writer, Error> {
+        todo!();
+
+        #[cfg(ignore)]
+        {
+            let lua_client = Client {
+                address: client.clone(),
+                for_write: true,
+            };
+
+            let resource: Resource = {
+                let mut lua = self.lua.lock_arc().await;
+                let lc = lua_client.clone().into_lua(&*lua).map_err(|_| Error::UnknownError)?;
+                let handle_fn: mlua::Function = lua.registry_value(&self.call_key).map_err(|_| Error::UnknownError)?;
+
+                let result = handle_fn.call_async((path.to_str().ok_or(Error::FileNotFound)?, lc, size)).await.map_err(|_| Error::UnknownError)?;
+
+                // let lc = lc.as_userdata().and_then(|ud| ud.take().ok()).unwrap_or(lua_client);
+                result
+            };
+
+            match resource {
+                Resource::Http(_) => { todo!() }
+                Resource::File(_) => { todo!() }
+                Resource::Data(_) => { todo!() }
+                Resource::Error(_) => { todo!() }
+            }
+        }
+    }
+}