trz-gateway-server 0.2.2

Secure Proxy / Agents implementation in Rust
Documentation 
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

use axum::http;
use tonic::body::Body;
use tonic::transport::channel::ResponseFuture;
use tower::Service;
use tower::load::Load as _;
use tower::util::rng::HasherRng;
use tower::util::rng::Rng;
use tracing::info;
use trz_gateway_common::is_global::IsGlobal;

use super::connection_id::ConnectionId;
use super::pending_requests::PendingRequests;

/// A struct that maintains a list of channels for a given client.
///
/// A client can open multiple tunnels with the Terrazzo Gateway.
pub struct IncomingClients<S: Service<http::Request<Body>>> {
    channels: Vec<ChannelWithId<S>>,
    rng: HasherRng,
}

struct ChannelWithId<S> {
    connection_id: ConnectionId,
    channel: PendingRequests<S>,
}

impl<S> IncomingClients<S>
where
    S: Service<
            http::Request<Body>,
            Response = http::Response<Body>,
            Future = ResponseFuture,
            Error = tonic::transport::Error,
        > + Clone
        + IsGlobal,
{
    pub fn new() -> Self {
        Self {
            channels: vec![],
            rng: Default::default(),
        }
    }

    /// Adds a channel
    pub fn add_channel(&mut self, connection_id: ConnectionId, channel: S) {
        info!("Adding channel");
        let channel = PendingRequests::new(channel);
        self.channels.push(ChannelWithId {
            connection_id,
            channel,
        });
    }

    /// Removes a channel
    pub fn remove_channel(&mut self, connection_id: ConnectionId) {
        info!("Removing channel");
        self.channels = std::mem::take(&mut self.channels)
            .into_iter()
            .filter(|c| c.connection_id != connection_id)
            .collect();
    }

    pub fn is_empty(&self) -> bool {
        self.channels.is_empty()
    }

    /// Returns a channel.
    ///
    /// If a client has ≥ 2 tunnels, the load-balancing algorithm choses
    /// channels that have less load based on the number of running requests.
    pub fn get_channel(&mut self) -> Option<PendingRequests<S>> {
        let count = self.channels.len();
        if count < 2 {
            return if count == 0 {
                None
            } else {
                Some(self.channels[0].channel.clone())
            };
        }
        let [a, b] =
            sample_floyd2(&mut self.rng, count as u64).map(|i| &self.channels[i as usize].channel);
        Some(if a.load() < b.load() { a } else { b }.clone())
    }
}

fn sample_floyd2<R: Rng>(rng: &mut R, length: u64) -> [u64; 2] {
    debug_assert!(2 <= length);
    let aidx = rng.next_range(0..length - 1);
    let bidx = rng.next_range(0..length);
    let aidx = if aidx == bidx { length - 1 } else { aidx };
    [aidx, bidx]
}