pingora_core/protocols/
digest.rs1use std::sync::Arc;
18use std::time::{Duration, SystemTime};
19
20use once_cell::sync::OnceCell;
21
22use super::l4::ext::{get_original_dest, get_recv_buf, get_snd_buf, get_tcp_info, TCP_INFO};
23use super::l4::socket::SocketAddr;
24use super::raw_connect::ProxyDigest;
25use super::tls::digest::SslDigest;
26
27#[derive(Clone, Debug, Default)]
29pub struct Digest {
30 pub ssl_digest: Option<Arc<SslDigest>>,
32 pub timing_digest: Vec<Option<TimingDigest>>,
34 pub proxy_digest: Option<Arc<ProxyDigest>>,
36 pub socket_digest: Option<Arc<SocketDigest>>,
38}
39
40pub trait ProtoDigest {
42 fn get_digest(&self) -> Option<&Digest> {
43 None
44 }
45}
46
47#[derive(Clone, Debug)]
49pub struct TimingDigest {
50 pub established_ts: SystemTime,
52 pub establishment_duration: Option<Duration>,
59 pub offload_wait_duration: Option<Duration>,
66}
67
68impl Default for TimingDigest {
69 fn default() -> Self {
70 TimingDigest {
71 established_ts: SystemTime::UNIX_EPOCH,
72 establishment_duration: None,
73 offload_wait_duration: None,
74 }
75 }
76}
77
78#[derive(Debug)]
79pub struct SocketDigest {
81 #[cfg(unix)]
82 raw_fd: std::os::unix::io::RawFd,
83 #[cfg(windows)]
84 raw_sock: std::os::windows::io::RawSocket,
85 pub peer_addr: OnceCell<Option<SocketAddr>>,
87 pub local_addr: OnceCell<Option<SocketAddr>>,
89 pub original_dst: OnceCell<Option<SocketAddr>>,
91}
92
93impl SocketDigest {
94 #[cfg(unix)]
95 pub fn from_raw_fd(raw_fd: std::os::unix::io::RawFd) -> SocketDigest {
96 SocketDigest {
97 raw_fd,
98 peer_addr: OnceCell::new(),
99 local_addr: OnceCell::new(),
100 original_dst: OnceCell::new(),
101 }
102 }
103
104 #[cfg(windows)]
105 pub fn from_raw_socket(raw_sock: std::os::windows::io::RawSocket) -> SocketDigest {
106 SocketDigest {
107 raw_sock,
108 peer_addr: OnceCell::new(),
109 local_addr: OnceCell::new(),
110 original_dst: OnceCell::new(),
111 }
112 }
113
114 #[cfg(unix)]
119 pub fn socket_cookie(&self) -> std::io::Result<u64> {
120 super::l4::ext::get_socket_cookie(self.raw_fd)
121 }
122
123 #[cfg(unix)]
124 pub fn peer_addr(&self) -> Option<&SocketAddr> {
125 self.peer_addr
126 .get_or_init(|| SocketAddr::from_raw_fd(self.raw_fd, true))
127 .as_ref()
128 }
129
130 #[cfg(windows)]
131 pub fn peer_addr(&self) -> Option<&SocketAddr> {
132 self.peer_addr
133 .get_or_init(|| SocketAddr::from_raw_socket(self.raw_sock, true))
134 .as_ref()
135 }
136
137 #[cfg(unix)]
138 pub fn local_addr(&self) -> Option<&SocketAddr> {
139 self.local_addr
140 .get_or_init(|| SocketAddr::from_raw_fd(self.raw_fd, false))
141 .as_ref()
142 }
143
144 #[cfg(windows)]
145 pub fn local_addr(&self) -> Option<&SocketAddr> {
146 self.local_addr
147 .get_or_init(|| SocketAddr::from_raw_socket(self.raw_sock, false))
148 .as_ref()
149 }
150
151 fn is_inet(&self) -> bool {
152 self.local_addr().and_then(|p| p.as_inet()).is_some()
153 }
154
155 #[cfg(unix)]
156 pub fn tcp_info(&self) -> Option<TCP_INFO> {
157 if self.is_inet() {
158 get_tcp_info(self.raw_fd).ok()
159 } else {
160 None
161 }
162 }
163
164 #[cfg(windows)]
165 pub fn tcp_info(&self) -> Option<TCP_INFO> {
166 if self.is_inet() {
167 get_tcp_info(self.raw_sock).ok()
168 } else {
169 None
170 }
171 }
172
173 #[cfg(unix)]
174 pub fn get_recv_buf(&self) -> Option<usize> {
175 if self.is_inet() {
176 get_recv_buf(self.raw_fd).ok()
177 } else {
178 None
179 }
180 }
181
182 #[cfg(windows)]
183 pub fn get_recv_buf(&self) -> Option<usize> {
184 if self.is_inet() {
185 get_recv_buf(self.raw_sock).ok()
186 } else {
187 None
188 }
189 }
190
191 #[cfg(unix)]
192 pub fn get_snd_buf(&self) -> Option<usize> {
193 if self.is_inet() {
194 get_snd_buf(self.raw_fd).ok()
195 } else {
196 None
197 }
198 }
199
200 #[cfg(windows)]
201 pub fn get_snd_buf(&self) -> Option<usize> {
202 if self.is_inet() {
203 get_snd_buf(self.raw_sock).ok()
204 } else {
205 None
206 }
207 }
208
209 #[cfg(unix)]
210 pub fn original_dst(&self) -> Option<&SocketAddr> {
211 self.original_dst
212 .get_or_init(|| {
213 get_original_dest(self.raw_fd)
214 .ok()
215 .flatten()
216 .map(SocketAddr::Inet)
217 })
218 .as_ref()
219 }
220
221 #[cfg(windows)]
222 pub fn original_dst(&self) -> Option<&SocketAddr> {
223 self.original_dst
224 .get_or_init(|| {
225 get_original_dest(self.raw_sock)
226 .ok()
227 .flatten()
228 .map(SocketAddr::Inet)
229 })
230 .as_ref()
231 }
232}
233
234pub trait GetTimingDigest {
236 fn get_timing_digest(&self) -> Vec<Option<TimingDigest>>;
238 fn get_read_pending_time(&self) -> Duration {
239 Duration::ZERO
240 }
241 fn get_write_pending_time(&self) -> Duration {
242 Duration::ZERO
243 }
244}
245
246pub trait GetProxyDigest {
248 fn get_proxy_digest(&self) -> Option<Arc<ProxyDigest>>;
249 fn set_proxy_digest(&mut self, _digest: ProxyDigest) {}
250}
251
252pub trait GetSocketDigest {
254 fn get_socket_digest(&self) -> Option<Arc<SocketDigest>>;
255 fn set_socket_digest(&mut self, _socket_digest: SocketDigest) {}
256}
257
258#[cfg(all(test, target_os = "linux"))]
259mod tests {
260 use super::SocketDigest;
261 use std::os::unix::io::AsRawFd;
262
263 #[test]
264 fn socket_cookie_returns_cookie_for_tcp_socket() {
265 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
266 let client = std::net::TcpStream::connect(listener.local_addr().unwrap()).unwrap();
267 let (server, _) = listener.accept().unwrap();
268
269 let client_digest = SocketDigest::from_raw_fd(client.as_raw_fd());
270 let server_digest = SocketDigest::from_raw_fd(server.as_raw_fd());
271
272 assert_ne!(client_digest.socket_cookie().unwrap(), 0);
273 assert_ne!(server_digest.socket_cookie().unwrap(), 0);
274 }
275
276 #[test]
277 fn socket_cookie_returns_cookie_for_unix_socket() {
278 let (client, server) = std::os::unix::net::UnixStream::pair().unwrap();
279
280 let client_digest = SocketDigest::from_raw_fd(client.as_raw_fd());
281 let server_digest = SocketDigest::from_raw_fd(server.as_raw_fd());
282
283 assert_ne!(client_digest.socket_cookie().unwrap(), 0);
284 assert_ne!(server_digest.socket_cookie().unwrap(), 0);
285 }
286}