data_loader/
lib.rs

1#[cfg(feature = "impl")]
2use anyhow::{anyhow, Context};
3#[cfg(feature = "impl")]
4use config::{any_err, from_lua_value, get_or_create_sub_module};
5#[cfg(feature = "impl")]
6use mlua::Lua;
7use serde::{Deserialize, Serialize};
8#[cfg(feature = "impl")]
9use vaultrs::client::{VaultClient, VaultClientSettingsBuilder};
10
11#[derive(Deserialize, Serialize, Clone, Hash, PartialEq, Eq, Debug)]
12#[serde(untagged)]
13pub enum KeySource {
14    File(String),
15    Data {
16        #[serde(with = "serde_bytes")]
17        key_data: Vec<u8>,
18    },
19    Vault {
20        vault_address: Option<String>,
21        vault_token: Option<String>,
22        vault_mount: String,
23        vault_path: String,
24        #[serde(default = "default_vault_key")]
25        vault_key: String,
26    },
27    Event {
28        event_name: String,
29        event_args: Vec<serde_json::Value>,
30    },
31}
32
33fn default_vault_key() -> String {
34    "key".to_string()
35}
36
37#[cfg(feature = "impl")]
38impl KeySource {
39    pub async fn get(&self) -> anyhow::Result<Vec<u8>> {
40        match self {
41            Self::File(path) => Ok(tokio::fs::read(path)
42                .await
43                .with_context(|| format!("KeySource failed to load data from file `{}`", path))?),
44            Self::Data { key_data } => Ok(key_data.to_vec()),
45            Self::Vault {
46                vault_address,
47                vault_token,
48                vault_mount,
49                vault_path,
50                vault_key,
51            } => {
52                let address = match vault_address {
53                    Some(a) => a.to_string(),
54                    None => std::env::var("VAULT_ADDR").map_err(|err| {
55                        anyhow!(
56                            "vault_address was not specified and $VAULT_ADDR is not set/usable: {self:?} {err:#}"
57                        )
58                    })?,
59                };
60                let token = match vault_token {
61                    Some(a) => a.to_string(),
62                    None => std::env::var("VAULT_TOKEN").map_err(|err| {
63                        anyhow!(
64                            "vault_token was not specified and $VAULT_TOKEN is not set/usable: {self:?} {err:#}"
65                        )
66                    })?,
67                };
68
69                let client = VaultClient::new(
70                    VaultClientSettingsBuilder::default()
71                        .address(address)
72                        .token(token)
73                        .build()?,
74                )?;
75
76                let entry: serde_json::Value = vaultrs::kv2::read(&client, vault_mount, vault_path)
77                    .await
78                    .with_context(|| {
79                        format!(
80                            "kv2::read vault_mount={vault_mount}, vault_path={vault_path} {self:?}"
81                        )
82                    })?;
83
84                let value = entry
85                    .get(vault_key)
86                    .and_then(|v| v.as_str())
87                    .ok_or_else(|| {
88                        anyhow!("vault secret at {vault_path} does not contain key '{vault_key}'")
89                    })?;
90
91                Ok(value.as_bytes().to_vec())
92            }
93            Self::Event {
94                event_name,
95                event_args,
96            } => {
97                let mut config = config::load_config().await?;
98                let sig = config::CallbackSignature::<mlua::MultiValue, mlua::String>::new(
99                    event_name.clone(),
100                );
101
102                let args = config.convert_args_to_multi(event_args)?;
103                let result = config.async_call_callback_non_default(&sig, args).await?;
104                config.put();
105
106                Ok(result.as_bytes().to_vec())
107            }
108        }
109    }
110}
111
112#[cfg(feature = "impl")]
113pub fn register(lua: &Lua) -> anyhow::Result<()> {
114    let secrets_mod = get_or_create_sub_module(lua, "secrets")?;
115
116    secrets_mod.set(
117        "load",
118        lua.create_async_function(|lua, source: mlua::Value| async move {
119            let source: KeySource = from_lua_value(&lua, source)?;
120            lua.create_string(&source.get().await.map_err(any_err)?)
121        })?,
122    )?;
123
124    Ok(())
125}
126
127#[cfg(test)]
128#[cfg(feature = "impl")]
129mod test {
130    use super::*;
131    use anyhow::Context;
132    use std::process::Stdio;
133    use std::time::Duration;
134    use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt};
135    use tokio::process::{Child, Command};
136    use tokio::time::timeout;
137    use vaultrs::client::Client;
138
139    /// Ask the kernel to assign a free port.
140    /// We pass this to sshd and tell it to listen on that port.
141    /// This is racy, as releasing the socket technically makes
142    /// that port available to others using the same technique.
143    fn allocate_port() -> u16 {
144        let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind 127.0.0.1:0 failed");
145        listener.local_addr().unwrap().port()
146    }
147
148    const KEY: &str = "woot";
149
150    struct VaultServer {
151        port: u16,
152        _daemon: Child,
153    }
154
155    impl VaultServer {
156        pub async fn spawn() -> anyhow::Result<Self> {
157            let mut errors = vec![];
158
159            for _ in 0..2 {
160                let port = allocate_port();
161                match timeout(Duration::from_secs(5), Self::spawn_with_port(port)).await? {
162                    Ok(me) => return Ok(me),
163                    Err(err) => {
164                        errors.push(format!("{err:#}"));
165                    }
166                }
167            }
168            anyhow::bail!("failed to spawn vault: {}", errors.join(". "));
169        }
170
171        pub fn address(&self) -> String {
172            format!("http://127.0.0.1:{}", self.port)
173        }
174
175        async fn spawn_with_port(port: u16) -> anyhow::Result<Self> {
176            eprintln!("Trying to start vault on port {port}");
177
178            let mut daemon = Command::new("vault")
179                .args([
180                    "server",
181                    "-dev",
182                    &format!("-dev-listen-address=127.0.0.1:{port}"),
183                    &format!("-dev-root-token-id={KEY}"),
184                ])
185                .stdin(Stdio::piped())
186                .stderr(Stdio::piped())
187                .stdout(Stdio::piped())
188                .kill_on_drop(true)
189                .spawn()
190                .context("spawning vault")?;
191
192            let mut stderr = daemon.stderr.take().unwrap();
193            tokio::spawn(async move {
194                copy_stream_with_line_prefix("vault stderr", &mut stderr, &mut tokio::io::stderr())
195                    .await
196            });
197            let mut stdout = daemon.stdout.take().unwrap();
198            tokio::spawn(async move {
199                copy_stream_with_line_prefix("vault stdout", &mut stdout, &mut tokio::io::stderr())
200                    .await
201            });
202
203            let mut ok = false;
204            for _ in 0..25 {
205                let client = VaultClient::new(
206                    VaultClientSettingsBuilder::default()
207                        .address(format!("http://127.0.0.1:{port}"))
208                        .token(KEY)
209                        .build()?,
210                )?;
211                let status = client.status().await;
212                eprintln!("checking status: {status:?}");
213                if let Ok(vaultrs::sys::ServerStatus::OK) = status {
214                    ok = true;
215                    break;
216                }
217                tokio::time::sleep(Duration::from_millis(100)).await;
218            }
219
220            anyhow::ensure!(ok, "server didn't startup successfully");
221
222            if let Ok(Some(status)) = daemon.try_wait() {
223                anyhow::bail!("daemon exited already: {status:?}");
224            }
225
226            Ok(Self {
227                _daemon: daemon,
228                port,
229            })
230        }
231
232        pub async fn put_from_file(&self, vault_path: &str, path: &str) -> anyhow::Result<()> {
233            let output = Command::new("vault")
234                .args([
235                    "kv",
236                    "put",
237                    &format!("-address=http://127.0.0.1:{}", self.port),
238                    "-mount=secret",
239                    vault_path,
240                    &format!("key=@{path}"),
241                ])
242                .output()
243                .await?;
244            let stdout = String::from_utf8_lossy(&output.stdout);
245            if !stdout.is_empty() {
246                eprintln!("put_from_file: {stdout}");
247            }
248            let stderr = String::from_utf8_lossy(&output.stderr);
249            if !stderr.is_empty() {
250                eprintln!("put_from_file: {stderr}");
251            }
252            anyhow::ensure!(output.status.success(), "{:?}", output.status);
253            Ok(())
254        }
255
256        pub async fn put(&self, vault_path: &str, value: &str) -> anyhow::Result<()> {
257            let output = Command::new("vault")
258                .args([
259                    "kv",
260                    "put",
261                    &format!("-address=http://127.0.0.1:{}", self.port),
262                    "-mount=secret",
263                    vault_path,
264                    &format!("key={value}"),
265                ])
266                .output()
267                .await?;
268            let stdout = String::from_utf8_lossy(&output.stdout);
269            if !stdout.is_empty() {
270                eprintln!("put: {stdout}");
271            }
272            let stderr = String::from_utf8_lossy(&output.stderr);
273            if !stderr.is_empty() {
274                eprintln!("put: {stderr}");
275            }
276            anyhow::ensure!(output.status.success(), "{:?}", output.status);
277            Ok(())
278        }
279
280        pub fn make_source(&self, path: &str) -> KeySource {
281            KeySource::Vault {
282                vault_address: Some(format!("http://127.0.0.1:{}", self.port)),
283                vault_token: Some(KEY.to_string()),
284                vault_mount: "secret".to_string(),
285                vault_path: path.to_string(),
286                vault_key: "key".to_string(),
287            }
288        }
289    }
290
291    async fn copy_stream_with_line_prefix<SRC, DEST>(
292        prefix: &str,
293        src: SRC,
294        mut dest: DEST,
295    ) -> std::io::Result<()>
296    where
297        SRC: AsyncRead + Unpin,
298        DEST: AsyncWrite + Unpin,
299    {
300        let mut src = tokio::io::BufReader::new(src);
301        loop {
302            let mut line = String::new();
303            src.read_line(&mut line).await?;
304            if !line.is_empty() {
305                dest.write_all(format!("{prefix}: {line}").as_bytes())
306                    .await?;
307            }
308        }
309    }
310
311    #[tokio::test]
312    async fn test_vault() -> anyhow::Result<()> {
313        if which::which("vault").is_err() {
314            return Ok(());
315        }
316        let vault = VaultServer::spawn().await?;
317
318        vault
319            .put_from_file("example.com", "../../example-private-dkim-key.pem")
320            .await?;
321
322        let source = vault.make_source("example.com");
323        let data = source.get().await?;
324
325        assert_eq!(
326            data,
327            std::fs::read("../../example-private-dkim-key.pem").unwrap()
328        );
329
330        vault.put("foo", "bar").await?;
331
332        let source = vault.make_source("foo");
333        let data = source.get().await?;
334
335        assert_eq!(data, b"bar");
336
337        let lua = Lua::new();
338        register(&lua).unwrap();
339        lua.globals().set("ADDR", vault.address())?;
340        lua.globals().set("KEY", KEY)?;
341        let pw: String = lua
342            .load(
343                r#"
344            local kumo = require 'kumo';
345            print("ADDR", ADDR)
346            print("KEY", KEY)
347            return kumo.secrets.load {
348                vault_path = "foo",
349                vault_address = ADDR,
350                vault_mount = "secret",
351                vault_token = KEY,
352            }
353            "#,
354            )
355            .eval_async()
356            .await
357            .unwrap();
358
359        assert_eq!(pw, "bar");
360
361        // Test with a different key name
362        vault.put("custom_key", "custom_value").await?;
363
364        let source = KeySource::Vault {
365            vault_address: Some(vault.address()),
366            vault_token: Some(KEY.to_string()),
367            vault_mount: "secret".to_string(),
368            vault_path: "custom_key".to_string(),
369            vault_key: "custom_field".to_string(),
370        };
371
372        // This should fail because the vault secret has "key" but we're looking for "custom_field"
373        let result = source.get().await;
374        let err = result.unwrap_err();
375        assert!(err
376            .to_string()
377            .contains("vault secret at custom_key does not contain key 'custom_field'"));
378
379        // Test with the correct key name
380        let source = KeySource::Vault {
381            vault_address: Some(vault.address()),
382            vault_token: Some(KEY.to_string()),
383            vault_mount: "secret".to_string(),
384            vault_path: "custom_key".to_string(),
385            vault_key: "key".to_string(),
386        };
387
388        let data = source.get().await?;
389        assert_eq!(data, b"custom_value");
390
391        Ok(())
392    }
393}
data_loader/lib.rs

data_loader/
lib.rs
