forked from scylladb/scylla-rust-driver
-
Notifications
You must be signed in to change notification settings - Fork 0
/
host_filter.rs
78 lines (66 loc) · 2 KB
/
host_filter.rs
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
//! Host filters.
//!
//! Host filters are essentially just a predicate over
//! [`Peer`](crate::transport::topology::Peer)s. Currently, they are used
//! by the [`Session`](crate::transport::session::Session) to determine whether
//! connections should be opened to a given node or not.
use std::collections::HashSet;
use std::io::Error;
use std::net::{SocketAddr, ToSocketAddrs};
use crate::transport::topology::Peer;
/// The `HostFilter` trait.
pub trait HostFilter: Send + Sync {
/// Returns whether a peer should be accepted or not.
fn accept(&self, peer: &Peer) -> bool;
}
/// Unconditionally accepts all nodes.
pub struct AcceptAllHostFilter;
impl HostFilter for AcceptAllHostFilter {
fn accept(&self, _peer: &Peer) -> bool {
true
}
}
/// Accepts nodes whose addresses are present in the allow list provided
/// during filter's construction.
pub struct AllowListHostFilter {
allowed: HashSet<SocketAddr>,
}
impl AllowListHostFilter {
/// Creates a new `AllowListHostFilter` which only accepts nodes from the
/// list.
pub fn new<I, A>(allowed_iter: I) -> Result<Self, Error>
where
I: IntoIterator<Item = A>,
A: ToSocketAddrs,
{
// I couldn't get the iterator combinators to work
let mut allowed = HashSet::new();
for item in allowed_iter {
for addr in item.to_socket_addrs()? {
allowed.insert(addr);
}
}
Ok(Self { allowed })
}
}
impl HostFilter for AllowListHostFilter {
fn accept(&self, peer: &Peer) -> bool {
self.allowed.contains(&peer.address)
}
}
/// Accepts nodes from given DC.
pub struct DcHostFilter {
local_dc: String,
}
impl DcHostFilter {
/// Creates a new `DcHostFilter` that accepts nodes only from the
/// `local_dc`.
pub fn new(local_dc: String) -> Self {
Self { local_dc }
}
}
impl HostFilter for DcHostFilter {
fn accept(&self, peer: &Peer) -> bool {
peer.datacenter.as_ref() == Some(&self.local_dc)
}
}