harborshield 0.1.0

A Rust port of Whalewall, to automate management of firewall rules for Docker containers
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

use super::CleanupTracker;
use crate::Result;
use bon::bon;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
use tracing::{debug, error};

/// RAII guard that ensures cleanup happens when dropped
pub struct CleanupGuard {
    tracker: Arc<CleanupTracker>,
    token: CancellationToken,
    resource_id: String,
    cleanup_on_drop: bool,
}

#[bon]
impl CleanupGuard {
    #[builder]
    /// Create a new cleanup guard
    pub fn new(tracker: Arc<CleanupTracker>, resource_id: String) -> Self {
        let token = tracker.child_token();
        Self {
            tracker,
            token,
            resource_id,
            cleanup_on_drop: true,
        }
    }

    /// Get the cancellation token for this guard
    pub fn token(&self) -> &CancellationToken {
        &self.token
    }

    /// Check if the operation has been cancelled
    pub fn is_cancelled(&self) -> bool {
        self.token.is_cancelled()
    }

    /// Disarm the guard - cleanup won't happen on drop
    pub fn disarm(mut self) {
        self.cleanup_on_drop = false;
        debug!("Cleanup guard disarmed for resource: {}", self.resource_id);
    }

    /// Complete the operation successfully
    pub async fn complete(mut self) -> Result<()> {
        self.cleanup_on_drop = false;
        debug!(
            "Operation completed successfully for resource: {}",
            self.resource_id
        );
        Ok(())
    }
}

impl Drop for CleanupGuard {
    fn drop(&mut self) {
        if self.cleanup_on_drop {
            debug!(
                "CleanupGuard dropped with cleanup enabled for resource: {}",
                self.resource_id
            );

            // We can't await in drop, so we spawn a task
            let tracker = self.tracker.clone();
            let resource_id = self.resource_id.clone();

            tokio::spawn(async move {
                if let Err(e) = tracker.cleanup_all().await {
                    error!("Failed to cleanup resources for {}: {}", resource_id, e);
                }
            });
        }
    }
}

#[async_trait::async_trait]
/// Extension trait for operations that need cleanup
pub trait WithCleanup {
    /// Execute with automatic cleanup on cancellation or error
    async fn with_cleanup<F, R>(
        &self,
        tracker: Arc<CleanupTracker>,
        resource_id: String,
        f: F,
    ) -> Result<R>
    where
        F: std::future::Future<Output = Result<R>>;
}

/// Macro to simplify cleanup guard usage
#[macro_export]
macro_rules! with_cleanup {
    ($tracker:expr, $resource_id:expr, $body:expr) => {{
        let guard = $crate::cleanup::guard::CleanupGuard::builder()
            .tracker($tracker)
            .resource_id($resource_id)
            .build();
        let token = guard.token().clone();

        tokio::select! {
            result = $body => {
                match result {
                    Ok(value) => {
                        guard.complete().await?;
                        Ok(value)
                    }
                    Err(e) => {
                        // Guard will cleanup on drop
                        Err(e)
                    }
                }
            }
            _ = token.cancelled() => {
                Err($crate::Error::invalid_state("Operation cancelled", "active", "cancelled"))
            }
        }
    }};
}