From 54f109f48ad7d9330197100eee907f44cf74241d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Christian=20Gr=C3=BCnhage?=
 <jan.christian@gruenhage.xyz>
Date: Tue, 28 Mar 2023 10:23:04 +0200
Subject: [PATCH] hacky hacks

---
 Cargo.toml          |  1 +
 src/vici_structs.rs | 28 ++++++++++++++++++++--------
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 54ad01d..1ec5017 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -14,3 +14,4 @@ actix-web-httpauth = "0.8.0"
 actix-web = "4.2.1"
 tokio = "1.21.2"
 rsvici = "0.1"
+anyhow = "1.0.70"
diff --git a/src/vici_structs.rs b/src/vici_structs.rs
index fde8266..1b4dec4 100644
--- a/src/vici_structs.rs
+++ b/src/vici_structs.rs
@@ -3,9 +3,11 @@
 use serde::Deserialize;
 use std::collections::HashMap;
 
-use prometheus_client::{
-    encoding::text::Encode,
-};
+use futures_util::    stream::StreamExt;
+
+use prometheus_client::encoding::text::Encode;
+
+use anyhow::Result;
 
 #[derive(Debug, Deserialize)]
 pub struct VICIState {
@@ -19,21 +21,29 @@ pub struct VICIState {
     pub pools:                  Pools,
 }
 
-impl<E> VICIState {
-    async fn update(client: rsvici::Client) -> Result<VICIState, E> {
-        VICIState {
+impl VICIState {
+    async fn update(client: &mut rsvici::Client) -> Result<VICIState> {
+        Ok(VICIState {
                 version:                client.request("version", ()).await?,
                 statistics:             client.request("statistics", ()).await?,
-                policies:               client.stream_request::<(), Policies>("list-policies", "list-policy", ()).await?,
+                policies:               collected_stream::<NamedPolicy, Policies>(client, "list-policies", "list-policy").await,
                 connections:            client.stream_request::<(), Connections>("list-connections", "list-conn", ()).await?,
                 security_associations:  client.stream_request::<(), SecurityAssociations>("list-sas", "list-sa", ()).await?,
                 certificates:           client.stream_request::<(), Certificates>("list-certs", "list-cert", ()).await?,
                 authorities:            client.stream_request::<(), Authorities>("list-authoroties", "list-authoroty", ()).await?,
                 pools:                  client.stream_request::<(), Pools>("list-pools", "list-pool", ()).await?,
-        }
+        })
     }
 }
 
+async fn collected_stream<N,C>(client: &mut rsvici::Client, command: &str, event: &str) -> C
+where
+    N: for<'de> serde::Deserialize<'de>,
+    C: std::iter::Extend<N> + Default,
+{
+    client.stream_request::<(), N>(command, event, ()).filter_map(|event| async move {event.ok()}).collect::<C>().await
+}
+
 // Structs for parsing the control interface
 
 #[derive(Debug, Deserialize)]
@@ -99,6 +109,8 @@ pub struct StatisticsMallinfo {
 
 pub type Policies = HashMap<String, Policy>;
 
+pub type NamedPolicy = (String, Policy);
+
 #[derive(Debug, Deserialize)]
 pub struct Policy {
     pub child:     String,