Files
stftpd/src/main.rs
2023-11-18 17:14:33 +01:00

115 lines
4.2 KiB
Rust

// SPDX-License-Identifier: MIT
use std::ffi::c_void;
use std::net;
use std::net::{Ipv4Addr, Ipv6Addr};
use std::path::{Path, PathBuf};
use anyhow::anyhow;
use structopt::StructOpt;
mod handler;
#[derive(StructOpt, Debug)]
struct Options {
#[structopt(short="s", long="script", env="STFTPD_SCRIPT")]
/// The lua script to determine how to handle requests
script: PathBuf,
/// Systemd socket activated mode. Can also be used for inetd activation
#[structopt(long)]
systemd: bool,
/// The address and port to listen on
#[structopt(short="l", env="STFTPD_LISTEN", default_value=":69")]
listen: String,
#[structopt(short="u")]
/// User to drop privileges to
user: Option<String>,
#[structopt(short="g")]
/// User to drop privileges to
group: Option<String>,
#[structopt(short="d")]
/// Directory to serve files from
serve: Option<PathBuf>,
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let opts = Options::from_args();
let mut handler = handler::Handler::new(
opts.serve.as_ref()
.map(PathBuf::as_path)
.unwrap_or_else(|| Path::new(""))
)?;
let load_fut = tokio::task::spawn_local(handler.load_script(opts.script.clone()));
load_fut.await??;
let sock = if opts.systemd {
let mut lfds = listenfd::ListenFd::from_env();
if lfds.len() > 0 {
lfds.take_udp_socket(0)?.ok_or(anyhow!("Failed to receive socket from systemd"))?
} else {
// inetd activation
let sock_fd = 0;
let sock = unsafe {
// validate the socket
let mut sockaddr : libc::sockaddr = std::mem::zeroed();
let mut sockaddr_sz = std::mem::size_of_val(&sockaddr) as libc::socklen_t;
let mut ty : libc::c_int = 0;
let mut ty_sz = std::mem::size_of_val(&ty) as libc::socklen_t;
let ret = libc::getsockname(sock_fd, &mut sockaddr, &mut sockaddr_sz);
if ret != 0 {
return Err(std::io::Error::last_os_error().into())
}
let ret = libc::getsockopt(
sock_fd,
libc::SOL_SOCKET,
libc::SO_TYPE,
&mut ty as *mut libc::c_int as *mut c_void,
&mut ty_sz
);
if ret != 0 {
return Err(std::io::Error::last_os_error().into())
}
if sockaddr.sa_family as libc::c_int != libc::AF_INET && sockaddr.sa_family as libc::c_int != libc::AF_INET6 || ty != libc::SOCK_DGRAM {
return Err(anyhow!("Can only listen on inet or inet6 UDP sockets"))
}
let owned = std::os::fd::BorrowedFd::borrow_raw(sock_fd).try_clone_to_owned()?;
// Putz around with stdin and stdout so that we don't accidentally write to the socket.
let mut pipe_fds = [0 as libc::c_int; 2];
if libc::pipe(&mut pipe_fds[0]) < 0 {
return Err(std::io::Error::last_os_error().into())
}
if libc::close(pipe_fds[1]) < 0 {
return Err(std::io::Error::last_os_error().into())
}
if libc::dup2(pipe_fds[0], 0) < 0 || libc::dup2(2, 1) < 0 {
return Err(std::io::Error::last_os_error().into())
}
net::UdpSocket::from(owned)
};
sock
}
} else {
let (host, port) = opts.listen.split_once(':').ok_or(anyhow!("Invalid listen address"))?;
let port = u16::from_str_radix(port, 10).map_err(|_| anyhow!("Invalid listen address"))?;
if host.is_empty() {
net::UdpSocket::bind((Ipv6Addr::UNSPECIFIED, port)).or_else(
|_| net::UdpSocket::bind((Ipv4Addr::UNSPECIFIED, port))
)?
} else {
net::UdpSocket::bind((host, port))?
}
};
let server = async_tftp::server::TftpServerBuilder::with_handler(handler.clone())
.std_socket(sock)?
.build().await?;
server.serve().await?;
println!("{opts:#?}");
Ok(())
}