summaryrefslogtreecommitdiffstats
path: root/src/network/connection.rs
blob: 34194c182f99b29ce95b2e6f7752b1a2982e3a20 (plain)
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
use ::std::fmt;
use ::std::net::Ipv4Addr;

use ::std::mem::swap;
use ::std::net::SocketAddr;

use std::cmp::Ordering;
use std::hash::{Hash, Hasher};

#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord)]
pub enum Protocol {
    Tcp,
    Udp,
}

impl fmt::Display for Protocol {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Protocol::Tcp => write!(f, "tcp"),
            Protocol::Udp => write!(f, "udp"),
        }
    }
}

#[derive(Clone)]
pub struct Socket {
    pub ip: Ipv4Addr,
    pub port: u16,
    host_addr: Option<String>
}

impl Socket {
    pub fn clone_host_or_ip(&self) -> String {
        match &self.host_addr {
            Some(host_addr) => host_addr.clone(),
            None => self.ip.to_string()
        }
    }
}

impl Ord for Socket {
    fn cmp(&self, other: &Self) -> Ordering {
        let ip_eq = self.ip.cmp(&other.ip); // TODO: also port
        match ip_eq {
            Ordering::Equal => self.port.cmp(&other.port),
            _ => ip_eq
        }
    }
}

impl PartialOrd for Socket {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl PartialEq for Socket {
    fn eq(&self, other: &Self) -> bool {
        self.ip == other.ip && self.port == other.port
    }
}

impl Eq for Socket {}

impl Hash for Socket {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.ip.hash(state);
        self.port.hash(state);
    }
}

impl fmt::Display for Socket {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match &self.host_addr {
            Some(host_addr) => {
                write!(f, "{}:{}", host_addr, self.port)
            },
            None => {
                write!(f, "{}:{}", self.ip, self.port)
            }
        }
    }
}

#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord)]
pub struct Connection {
    pub local_socket: Socket,
    pub remote_socket: Socket,
    pub protocol: Protocol,
}

impl fmt::Display for Connection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{} => {} ({})",
            self.local_socket, self.remote_socket, self.protocol
        )
    }
}

impl Connection {
    pub fn new(
        local_socket: SocketAddr,
        remote_socket: SocketAddr,
        protocol: Protocol,
    ) -> Option<Self> {
        match (local_socket, remote_socket) {
            (SocketAddr::V4(local_socket), SocketAddr::V4(remote_socket)) => {
                Some(Connection {
                    local_socket: Socket {
                        ip: *local_socket.ip(),
                        port: local_socket.port(),
                        host_addr: None,
                    },
                    remote_socket: Socket {
                        ip: *remote_socket.ip(),
                        port: remote_socket.port(),
                        host_addr: None,
                    },
                    protocol,
                })
            },
            (_, _) => None,
        }
    }
    pub fn swap_direction(&mut self) {
        swap(&mut self.local_socket, &mut self.remote_socket);
    }
    pub fn set_local_host_addr(&mut self, addr: &str) {
        self.local_socket.host_addr = Some(String::from(addr));
    }
    pub fn set_remote_host_addr(&mut self, addr: &str) {
        self.remote_socket.host_addr = Some(String::from(addr));
    }
}