kingfisher-scanner 1.2.0

High-level scanning API for Kingfisher secret scanner
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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

//! Network resource policy scoped to one future, never to the process.
//!
//! A scope follows the future across executor threads. Spawned tasks must be
//! scoped explicitly; unrelated concurrent callers retain their own defaults.
use std::{
    future::{Future, IntoFuture},
    time::Duration,
};

tokio::task_local! {
    static POLICY: NetworkLimits;
}

#[derive(Clone, Copy, Debug, Default)]
pub struct NetworkLimits {
    pub no_timeouts: bool,
    pub unlimited_response: bool,
    pub unlimited_results: bool,
}

impl NetworkLimits {
    pub fn current() -> Self {
        POLICY.try_with(|value| *value).unwrap_or_default()
    }

    pub async fn scope<F: Future>(self, future: F) -> F::Output {
        POLICY.scope(self, future).await
    }

    pub fn sync_scope<T>(self, work: impl FnOnce() -> T) -> T {
        POLICY.sync_scope(self, work)
    }
}

pub async fn timeout<F: IntoFuture>(
    duration: Duration,
    work: F,
) -> Result<F::Output, tokio::time::error::Elapsed> {
    if duration.is_zero() || NetworkLimits::current().no_timeouts {
        Ok(work.await)
    } else {
        tokio::time::timeout(duration, work.into_future()).await
    }
}

/// Apply a timeout only when the enclosing operation has a finite budget.
pub trait ResourceTimeout: Sized {
    fn resource_timeout(self, duration: Duration) -> Self;
}
impl ResourceTimeout for reqwest::ClientBuilder {
    fn resource_timeout(self, duration: Duration) -> Self {
        if duration.is_zero() || NetworkLimits::current().no_timeouts {
            self
        } else {
            self.timeout(duration)
        }
    }
}
impl ResourceTimeout for reqwest::RequestBuilder {
    fn resource_timeout(self, duration: Duration) -> Self {
        if duration.is_zero() || NetworkLimits::current().no_timeouts {
            self
        } else {
            self.timeout(duration)
        }
    }
}

/// Apply a connection timeout only when the enclosing operation has a finite budget.
pub trait ResourceConnectTimeout: Sized {
    fn resource_connect_timeout(self, duration: Duration) -> Self;
}
impl ResourceConnectTimeout for reqwest::ClientBuilder {
    fn resource_connect_timeout(self, duration: Duration) -> Self {
        if duration.is_zero() || NetworkLimits::current().no_timeouts {
            self
        } else {
            self.connect_timeout(duration)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn unlimited_client_waits_for_response_without_affecting_bounded_clients() {
        let app = axum::Router::new().route(
            "/",
            axum::routing::get(|| async {
                tokio::time::sleep(Duration::from_millis(50)).await;
                "complete"
            }),
        );
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let url = format!("http://{}/", listener.local_addr().unwrap());
        let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
        let work = || async {
            reqwest::Client::builder()
                .no_proxy()
                .resource_connect_timeout(Duration::from_secs(1))
                .resource_timeout(Duration::from_millis(5))
                .build()
                .unwrap()
                .get(&url)
                .send()
                .await
        };
        let (unlimited, bounded) = tokio::join!(
            NetworkLimits { no_timeouts: true, ..Default::default() }.scope(work()),
            work(),
        );
        assert_eq!(unlimited.unwrap().text().await.unwrap(), "complete");
        assert!(bounded.unwrap_err().is_timeout());
        assert!(!NetworkLimits::current().no_timeouts);
        server.abort();
    }

    #[tokio::test]
    async fn unlimited_timeout_is_scoped_and_preserves_other_tasks() {
        let work = || async {
            timeout(Duration::from_millis(1), async {
                tokio::time::sleep(Duration::from_millis(25)).await;
                42
            })
            .await
        };
        let (unlimited, bounded) = tokio::join!(
            NetworkLimits { no_timeouts: true, ..Default::default() }.scope(work()),
            work(),
        );
        assert_eq!(unlimited.unwrap(), 42);
        assert!(bounded.is_err());
        assert!(!NetworkLimits::current().no_timeouts);
    }
}