From d249a67be47c1f50312d49f615ee9463daf43737 Mon Sep 17 00:00:00 2001
From: Vivian Roest <vivian@0x76.dev>
Date: Tue, 19 Dec 2023 10:21:13 +0100
Subject: [PATCH] add tries, and timeout

---
 flake.nix      |  4 ++--
 src/main.rs    | 33 +++++++++++++++++++++++++--------
 src/tpm/mod.rs |  5 ++---
 3 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/flake.nix b/flake.nix
index fdb20f9..7f6ee89 100644
--- a/flake.nix
+++ b/flake.nix
@@ -19,6 +19,7 @@
 
           doCheck = false;
 
+          buildInputs = with pkgs; [ openssl tpm2-tss ];
           nativeBuildInputs = with pkgs; [
             llvmPackages.libclang
             llvmPackages.libcxxClang
@@ -26,7 +27,6 @@
             pkg-config
           ];
           LIBCLANG_PATH = "${pkgs.llvmPackages.libclang.lib}/lib";
-          buildInputs = with pkgs; [ openssl tpm2-tss ];
 
           preBuild = ''
             export BINDGEN_EXTRA_CLANG_ARGS="$(< ${stdenv.cc}/nix-support/libc-crt1-cflags) \
@@ -53,7 +53,7 @@
           '';
         };
 
-        devShell = pkgs.mkShell {
+        devShells.default = pkgs.mkShell {
           shellHook = "${packages.default.preBuild}";
           inherit (packages.default) nativeBuildInputs buildInputs LIBCLANG_PATH;
         };
diff --git a/src/main.rs b/src/main.rs
index 54bcc66..9b0a758 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -8,6 +8,8 @@ use std::io::{Read, Write};
 use std::os::unix::net::UnixStream;
 use std::path::PathBuf;
 use std::process::exit;
+use std::thread::sleep;
+use std::time::Duration;
 use tpm::tpm_objects::TPM2Config;
 
 fn get_control_socket() -> Option<PathBuf> {
@@ -25,6 +27,7 @@ fn get_control_socket() -> Option<PathBuf> {
 }
 
 #[derive(Debug, Clone, Copy)]
+#[allow(dead_code)]
 enum ControlOp {
     Initialize = 0,
     Unlock = 1,
@@ -52,6 +55,7 @@ impl ControlResult {
         match num {
             0 => Some(Self::Ok),
             1 => Some(Self::Denied),
+            2 => Some(Self::Failed),
             3 => Some(Self::NoDaemon),
             _ => None,
         }
@@ -122,7 +126,14 @@ struct Cli {
 #[derive(Subcommand)]
 enum Commands {
     /// Unlock gnome keyring using encrypted password stored in tpm
-    Unlock,
+    Unlock {
+        #[arg(default_value_t = 5)]
+        tries: usize,
+
+        // timeout in seconds
+        #[arg(default_value_t = 1)]
+        timeout: u64,
+    },
     /// Enroll a password into the tpm to use when unlocking
     Enroll,
 }
@@ -137,15 +148,21 @@ fn main() -> color_eyre::Result<()> {
         .ok_or_else(|| eyre!("Token path not found"))?;
 
     match cli.command {
-        Commands::Unlock => {
+        Commands::Unlock { tries, timeout } => {
             if token_path.exists() {
                 let token = read_to_string(token_path)?;
-                let password =
-                    tpm::perform_decrypt(token.as_bytes()).map_err(|err| eyre!("{err:?}"))?;
-                let res = unlock_keyring(password.as_slice())?;
-                if res != ControlResult::Ok {
-                    eprintln!("Failed to unlock keyring: {res:?}");
-                    exit(2);
+
+                for _ in 0..tries {
+                    let password =
+                        tpm::perform_decrypt(token.as_bytes()).map_err(|err| eyre!("{err:?}"))?;
+                    let res = unlock_keyring(password.as_slice())?;
+                    if res == ControlResult::Ok {
+                        break;
+                    } else {
+                        eprintln!("Failed to unlock keyring: {res:?}");
+                    }
+
+                    sleep(Duration::from_secs(timeout));
                 }
             } else {
                 bail!("password token file not found")
diff --git a/src/tpm/mod.rs b/src/tpm/mod.rs
index 6210eee..db01361 100644
--- a/src/tpm/mod.rs
+++ b/src/tpm/mod.rs
@@ -1,6 +1,5 @@
-// Copyright 2020 Patrick Uiterwijk
-//
-// Licensed under the MIT license
+// Adapted from https://github.com/fedora-iot/clevis-pin-tpm2/
+// by Patrick Uiterwijk under the MIT License
 
 use std::convert::{TryFrom, TryInto};