nftables_async/
driver.rs

1use std::{ffi::OsStr, future::Future, process::Output};
2
3#[cfg(feature = "async-process-driver")]
4use futures_lite::AsyncWriteExt as FuturesAsyncWriteExt;
5#[cfg(any(feature = "async-process-driver", feature = "tokio-driver"))]
6use std::process::Stdio;
7#[cfg(feature = "tokio-driver")]
8use tokio::io::AsyncWriteExt as TokioAsyncWriteExt;
9
10/// A process driver internally used by the helper to run the "nft" process asynchronously, write to its stdin and
11/// retrieve its output.
12pub trait Driver {
13    /// Run the provided program with the provided arguments and retrieve its output.
14    /// If stdin is [None], stdin should be nulled by the driver and no interaction should occur with the process.
15    /// If stdin is [Some], stdin should be piped and the byte payload as well as a subsequent EOF (closure of the pipe)
16    /// should be transmitted before the process is waited on.
17    fn run_process(
18        program: &OsStr,
19        args: &[&OsStr],
20        stdin: Option<&[u8]>,
21    ) -> impl Future<Output = Result<Output, std::io::Error>> + Send;
22}
23
24/// A [Driver] implementation using the tokio crate for I/O.
25#[cfg(feature = "tokio-driver")]
26#[cfg_attr(docsrs, doc(cfg(feature = "tokio-process")))]
27pub struct TokioDriver;
28
29#[cfg(feature = "tokio-driver")]
30#[cfg_attr(docsrs, doc(cfg(feature = "tokio-process")))]
31impl Driver for TokioDriver {
32    async fn run_process(
33        program: &OsStr,
34        args: &[&OsStr],
35        stdin: Option<&[u8]>,
36    ) -> Result<Output, std::io::Error> {
37        let mut command = tokio::process::Command::new(program);
38        command.args(args);
39
40        match stdin {
41            Some(stdin) => {
42                let mut child = command.stdin(Stdio::piped()).spawn()?;
43                child
44                    .stdin
45                    .take()
46                    .ok_or_else(|| {
47                        std::io::Error::new(
48                            std::io::ErrorKind::InvalidInput,
49                            "Stdin not redirected successfully by tokio",
50                        )
51                    })?
52                    .write_all(stdin)
53                    .await?;
54
55                child.wait_with_output().await
56            }
57            None => command.stdin(Stdio::null()).output().await,
58        }
59    }
60}
61
62/// A [Driver] implementation using the async-process crate for I/O.
63#[cfg(feature = "async-process-driver")]
64#[cfg_attr(docsrs, doc(cfg(feature = "async-process")))]
65pub struct AsyncProcessDriver;
66
67#[cfg(feature = "async-process-driver")]
68#[cfg_attr(docsrs, doc(cfg(feature = "async-process")))]
69impl Driver for AsyncProcessDriver {
70    async fn run_process(
71        program: &OsStr,
72        args: &[&OsStr],
73        stdin: Option<&[u8]>,
74    ) -> Result<Output, std::io::Error> {
75        let mut command = async_process::Command::new(program);
76        command.args(args);
77
78        match stdin {
79            Some(stdin) => {
80                let mut child = command.stdin(Stdio::piped()).spawn()?;
81                child
82                    .stdin
83                    .take()
84                    .ok_or_else(|| {
85                        std::io::Error::new(
86                            std::io::ErrorKind::InvalidInput,
87                            "Stdin not redirected successfully by async-process",
88                        )
89                    })?
90                    .write_all(stdin)
91                    .await?;
92
93                child.output().await
94            }
95            None => command.stdin(Stdio::null()).output().await,
96        }
97    }
98}