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

//! Host filters.
//!
//! Host filters are essentially just a predicate over
//! [`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 {
        match peer.address {
            super::NodeAddr::Translatable(addr) => self.allowed.contains(&addr),
            // If the address is Untranslatable, then the node either was originally
            // an ContactPoint, or a Translatable node for which the host filter
            // returned true.
            super::NodeAddr::Untranslatable(_) => true,
        }
    }
}

/// 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)
    }
}