From 0e6a677b05e875e89a0fbe11dc4ea55f5ef1370b Mon Sep 17 00:00:00 2001
From: Jordan Doyle <jordan@doyle.la>
Date: Sun, 25 Jun 2023 12:55:19 +0100
Subject: [PATCH] Flush to audits to disk after not receiving a write for 5 seconds

---
 flake.nix    | 10 ----------
 src/audit.rs | 39 ++++++++++++++++++++++++++++++---------
 2 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/flake.nix b/flake.nix
index 84bc4e7..e267263 100644
--- a/flake.nix
+++ b/flake.nix
@@ -40,16 +40,6 @@
             };
 
             config = mkIf cfg.enable {
-              systemd.sockets.pisshoff = {
-                wantedBy = [ "sockets.target" ];
-
-                socketConfig = {
-                  ListenStream = 22;
-                  BindIPv6Only = "both";
-                  Accept = "no";
-                };
-              };
-
               systemd.services.pisshoff = {
                 enable = true;
                 wantedBy = [ "multi-user.target" ];
diff --git a/src/audit.rs b/src/audit.rs
index d7c842d..4071cd7 100644
--- a/src/audit.rs
+++ b/src/audit.rs
@@ -1,14 +1,18 @@
 use crate::config::Config;
 use serde::Serialize;
-use std::io::ErrorKind;
-use std::sync::Arc;
 use std::{
     fmt::{Debug, Formatter},
+    io::ErrorKind,
     net::SocketAddr,
+    sync::Arc,
     time::{Duration, Instant},
 };
-use tokio::io::{AsyncWriteExt, BufWriter};
-use tokio::{fs::OpenOptions, task::JoinHandle};
+use tokio::{
+    fs::OpenOptions,
+    io::{AsyncWriteExt, BufWriter},
+    task::JoinHandle,
+};
+use tracing::debug;
 use uuid::Uuid;
 
 pub fn start_audit_writer(
@@ -26,11 +30,28 @@ pub fn start_audit_writer(
             .open(&config.audit_output_file)
             .await?;
         let mut writer = BufWriter::new(file);
-
-        while let Some(log) = recv.recv().await {
-            let log =
-                serde_json::to_vec(&log).map_err(|e| std::io::Error::new(ErrorKind::Other, e))?;
-            writer.write_all(&log).await?;
+        let mut shutdown = false;
+
+        loop {
+            tokio::select! {
+                log = recv.recv(), if !shutdown => {
+                    match log {
+                        Some(log) => {
+                            let log = serde_json::to_vec(&log)
+                                .map_err(|e| std::io::Error::new(ErrorKind::Other, e))?;
+                            writer.write_all(&log).await?;
+                        }
+                        None => {
+                            shutdown = true;
+                        }
+                    }
+                }
+                _ = tokio::time::sleep(Duration::from_secs(5)), if !writer.buffer().is_empty() && !shutdown => {
+                    debug!("Flushing audits to disk");
+                    writer.flush().await?;
+                }
+                else => break,
+            }
         }
 
         writer.flush().await?;
--
libgit2 1.7.2