mlock sensitive memory

1 files changed, 76 insertions, 34 deletions

diff --git a/src/pinentry.rs b/src/pinentry.rs
index ff778e7..4a2196e 100644
--- a/src/pinentry.rs
+++ b/src/pinentry.rs
@@ -6,7 +6,7 @@ pub async fn getpin(
     prompt: &str,
     desc: &str,
     tty: Option<&str>,
-) -> Result<String> {
+) -> Result<crate::locked::Password> {
     let mut opts = tokio::process::Command::new("pinentry");
     let opts = opts
         .stdin(std::process::Stdio::piped())
@@ -17,42 +17,84 @@ pub async fn getpin(
         opts
     };
     let mut child = opts.spawn().context(crate::error::Spawn)?;
-    {
-        let stdin = child.stdin.as_mut().unwrap();
-
-        stdin
-            .write_all(b"SETTITLE rbw\n")
-            .await
-            .context(crate::error::WriteStdin)?;
-        stdin
-            .write_all(format!("SETPROMPT {}\n", prompt).as_bytes())
-            .await
-            .context(crate::error::WriteStdin)?;
-        stdin
-            .write_all(format!("SETDESC {}\n", desc).as_bytes())
-            .await
-            .context(crate::error::WriteStdin)?;
-        stdin
-            .write_all(b"GETPIN\n")
-            .await
-            .context(crate::error::WriteStdin)?;
-    }
+    let mut stdin = child.stdin.take().unwrap();
 
-    let out = child
-        .wait_with_output()
+    stdin
+        .write_all(b"SETTITLE rbw\n")
+        .await
+        .context(crate::error::WriteStdin)?;
+    stdin
+        .write_all(format!("SETPROMPT {}\n", prompt).as_bytes())
+        .await
+        .context(crate::error::WriteStdin)?;
+    stdin
+        .write_all(format!("SETDESC {}\n", desc).as_bytes())
         .await
-        .context(crate::error::ProcessWaitOutput)?
-        .stdout;
-    let out_str = String::from_utf8(out.clone()).context(
-        crate::error::FailedToParsePinentryUtf8 { out: out.clone() },
-    )?;
-    for line in out_str.lines() {
-        if line.starts_with("D ") {
-            return Ok(line[2..line.len()].to_string());
-        } else if !line.starts_with("OK") {
-            break;
+        .context(crate::error::WriteStdin)?;
+    stdin
+        .write_all(b"GETPIN\n")
+        .await
+        .context(crate::error::WriteStdin)?;
+    drop(stdin);
+
+    let mut buf = crate::locked::Vec::new();
+    buf.extend(std::iter::repeat(0));
+    let len =
+        read_password(buf.data_mut(), child.stdout.as_mut().unwrap()).await?;
+    buf.truncate(len);
+
+    child.await.context(crate::error::PinentryWait)?;
+
+    Ok(crate::locked::Password::new(buf))
+}
+
+async fn read_password<
+    R: tokio::io::AsyncRead + tokio::io::AsyncReadExt + Unpin,
+>(
+    data: &mut [u8],
+    mut r: R,
+) -> Result<usize> {
+    let mut len = 0;
+    loop {
+        let nl = data.iter().take(len).position(|c| *c == b'\n');
+        if let Some(nl) = nl {
+            if data.starts_with(b"OK") {
+                data.copy_within((nl + 1).., 0);
+                len -= nl + 1;
+            } else if data.starts_with(b"D ") {
+                data.copy_within(2..nl, 0);
+                len = nl - 2;
+                break;
+            } else {
+                return Err(Error::FailedToParsePinentry {
+                    out: data.to_vec(),
+                });
+            }
+        } else {
+            let bytes = r
+                .read(&mut data[len..])
+                .await
+                .context(crate::error::PinentryReadOutput)?;
+            len += bytes;
         }
     }
 
-    Err(Error::FailedToParsePinentry { out })
+    Ok(len)
+}
+
+#[test]
+fn test_read_password() {
+    let good_inputs = &[
+        &b"D super secret password\n"[..],
+        &b"OK\nOK\nOK\nD super secret password\nOK\n"[..],
+        &b"OK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nD super secret password\nOK\n"[..],
+        &b"OK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nD super secret password\nOK\n"[..],
+    ];
+    for input in good_inputs {
+        let mut buf = [0; 64];
+        tokio::runtime::Runtime::new().unwrap().block_on(async {
+            let len = read_password(&mut buf, &input[..]).await.unwrap();
+            assert_eq!(&buf[0..len], b"super secret password");
+        });
+    }
 }