vector_core/tls/
maybe_tls.rs

1use std::{
2    fmt,
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7use pin_project::pin_project;
8use tokio::io::{self, AsyncRead, AsyncWrite, ReadBuf};
9
10/// A type wrapper for objects that can exist in either a raw state or
11/// wrapped by TLS handling.
12#[pin_project(project = MaybeTlsProj)]
13pub enum MaybeTls<R, T> {
14    Raw(#[pin] R),
15    Tls(#[pin] T),
16}
17
18impl<R, T> MaybeTls<R, T> {
19    pub const fn is_raw(&self) -> bool {
20        matches!(self, Self::Raw(_))
21    }
22
23    pub const fn is_tls(&self) -> bool {
24        matches!(self, Self::Tls(_))
25    }
26
27    pub const fn raw(&self) -> Option<&R> {
28        match self {
29            Self::Raw(raw) => Some(raw),
30            Self::Tls(_) => None,
31        }
32    }
33
34    pub const fn tls(&self) -> Option<&T> {
35        match self {
36            Self::Raw(_) => None,
37            Self::Tls(tls) => Some(tls),
38        }
39    }
40}
41
42impl<O> From<Option<O>> for MaybeTls<(), O> {
43    fn from(tls: Option<O>) -> Self {
44        match tls {
45            Some(tls) => Self::Tls(tls),
46            None => Self::Raw(()),
47        }
48    }
49}
50
51// Conditionally implement Clone for Cloneable types
52impl<R: Clone, T: Clone> Clone for MaybeTls<R, T> {
53    fn clone(&self) -> Self {
54        match self {
55            Self::Raw(raw) => Self::Raw(raw.clone()),
56            Self::Tls(tls) => Self::Tls(tls.clone()),
57        }
58    }
59}
60
61// Conditionally implement Debug for Debuggable types
62impl<R: fmt::Debug, T: fmt::Debug> fmt::Debug for MaybeTls<R, T> {
63    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
64        match self {
65            Self::Raw(raw) => write!(fmt, "MaybeTls::Raw({raw:?})"),
66            Self::Tls(tls) => write!(fmt, "MaybeTls::Tls({tls:?})"),
67        }
68    }
69}
70
71impl<R: AsyncRead, T: AsyncRead> AsyncRead for MaybeTls<R, T> {
72    fn poll_read(
73        self: Pin<&mut Self>,
74        cx: &mut Context,
75        buf: &mut ReadBuf<'_>,
76    ) -> Poll<io::Result<()>> {
77        match self.project() {
78            MaybeTlsProj::Tls(s) => s.poll_read(cx, buf),
79            MaybeTlsProj::Raw(s) => s.poll_read(cx, buf),
80        }
81    }
82}
83
84impl<R: AsyncWrite, T: AsyncWrite> AsyncWrite for MaybeTls<R, T> {
85    fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<io::Result<usize>> {
86        match self.project() {
87            MaybeTlsProj::Tls(s) => s.poll_write(cx, buf),
88            MaybeTlsProj::Raw(s) => s.poll_write(cx, buf),
89        }
90    }
91
92    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
93        match self.project() {
94            MaybeTlsProj::Tls(s) => s.poll_flush(cx),
95            MaybeTlsProj::Raw(s) => s.poll_flush(cx),
96        }
97    }
98
99    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
100        match self.project() {
101            MaybeTlsProj::Tls(s) => s.poll_shutdown(cx),
102            MaybeTlsProj::Raw(s) => s.poll_shutdown(cx),
103        }
104    }
105}
vector_core/tls/maybe_tls.rs

vector_core/tls/
maybe_tls.rs
