s2n_quic_dc/stream/socket/
handle.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use super::{Protocol, TransportFeatures};
5use crate::msg::{self, addr::Addr, cmsg};
6use core::task::{Context, Poll};
7use s2n_quic_core::inet::ExplicitCongestionNotification;
8use std::{
9    io::{self, IoSlice, IoSliceMut},
10    net::SocketAddr,
11    sync::Arc,
12};
13
14pub type Flags = libc::c_int;
15
16pub trait Socket: 'static + Send + Sync {
17    /// Returns the local address for the socket
18    fn local_addr(&self) -> io::Result<SocketAddr>;
19
20    /// Returns the local port for the socket
21    #[inline]
22    fn local_port(&self) -> io::Result<u16> {
23        Ok(self.local_addr()?.port())
24    }
25
26    fn protocol(&self) -> Protocol;
27
28    /// Returns the [`TransportFeatures`] that the socket supports
29    fn features(&self) -> TransportFeatures;
30
31    /// Returns the amount of buffered data on the socket
32    fn poll_peek_len(&self, cx: &mut Context) -> Poll<io::Result<usize>>;
33
34    #[inline]
35    fn poll_recv_buffer(
36        &self,
37        cx: &mut Context,
38        msg: &mut msg::recv::Message,
39    ) -> Poll<io::Result<usize>> {
40        #[cfg(debug_assertions)]
41        if !self.features().is_stream() {
42            assert!(
43                msg.is_empty(),
44                "receive buffer should be empty for datagram protocols"
45            );
46        }
47
48        msg.poll_recv_with(|addr, cmsg, buffer| self.poll_recv(cx, addr, cmsg, buffer))
49    }
50
51    /// Receives data on the socket
52    fn poll_recv(
53        &self,
54        cx: &mut Context,
55        addr: &mut Addr,
56        cmsg: &mut cmsg::Receiver,
57        buffer: &mut [IoSliceMut],
58    ) -> Poll<io::Result<usize>>;
59
60    #[inline]
61    fn try_send_buffer(&self, msg: &mut msg::send::Message) -> io::Result<usize> {
62        msg.send_with(|addr, ecn, iov| self.try_send(addr, ecn, iov))
63    }
64
65    /// Tries to send data on the socket, returning `Err(WouldBlock)` if none could be sent.
66    fn try_send(
67        &self,
68        addr: &Addr,
69        ecn: ExplicitCongestionNotification,
70        buffer: &[IoSlice],
71    ) -> io::Result<usize>;
72
73    #[inline]
74    fn poll_send_buffer(
75        &self,
76        cx: &mut Context,
77        msg: &mut msg::send::Message,
78    ) -> Poll<io::Result<usize>> {
79        msg.poll_send_with(|addr, ecn, iov| self.poll_send(cx, addr, ecn, iov))
80    }
81
82    /// Sends data on the socket
83    fn poll_send(
84        &self,
85        cx: &mut Context,
86        addr: &Addr,
87        ecn: ExplicitCongestionNotification,
88        buffer: &[IoSlice],
89    ) -> Poll<io::Result<usize>>;
90
91    /// Shuts down the sender half of the socket, if a concept exists
92    fn send_finish(&self) -> io::Result<()>;
93}
94
95pub trait Ext: Socket {
96    #[inline]
97    fn recv_buffer<'a>(&'a self, msg: &'a mut msg::recv::Message) -> ExtRecvBuffer<'a, Self> {
98        ExtRecvBuffer { socket: self, msg }
99    }
100}
101
102pub struct ExtRecvBuffer<'a, T: Socket + ?Sized> {
103    socket: &'a T,
104    msg: &'a mut msg::recv::Message,
105}
106
107impl<T: Socket> core::future::Future for ExtRecvBuffer<'_, T> {
108    type Output = io::Result<usize>;
109
110    fn poll(mut self: core::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
111        self.socket.poll_recv_buffer(cx, self.msg)
112    }
113}
114
115impl<T: Socket> Ext for T {}
116
117macro_rules! impl_box {
118    ($b:ident) => {
119        impl<T: Socket> Socket for $b<T> {
120            #[inline(always)]
121            fn local_addr(&self) -> io::Result<SocketAddr> {
122                (**self).local_addr()
123            }
124
125            #[inline(always)]
126            fn protocol(&self) -> Protocol {
127                (**self).protocol()
128            }
129
130            #[inline(always)]
131            fn features(&self) -> TransportFeatures {
132                (**self).features()
133            }
134
135            #[inline(always)]
136            fn poll_peek_len(&self, cx: &mut Context) -> Poll<io::Result<usize>> {
137                (**self).poll_peek_len(cx)
138            }
139
140            #[inline(always)]
141            fn poll_recv(
142                &self,
143                cx: &mut Context,
144                addr: &mut Addr,
145                cmsg: &mut cmsg::Receiver,
146                buffer: &mut [IoSliceMut],
147            ) -> Poll<io::Result<usize>> {
148                (**self).poll_recv(cx, addr, cmsg, buffer)
149            }
150
151            #[inline(always)]
152            fn try_send(
153                &self,
154                addr: &Addr,
155                ecn: ExplicitCongestionNotification,
156                buffer: &[IoSlice],
157            ) -> io::Result<usize> {
158                (**self).try_send(addr, ecn, buffer)
159            }
160
161            #[inline(always)]
162            fn poll_send(
163                &self,
164                cx: &mut Context,
165                addr: &Addr,
166                ecn: ExplicitCongestionNotification,
167                buffer: &[IoSlice],
168            ) -> Poll<io::Result<usize>> {
169                (**self).poll_send(cx, addr, ecn, buffer)
170            }
171
172            #[inline(always)]
173            fn send_finish(&self) -> io::Result<()> {
174                (**self).send_finish()
175            }
176        }
177    };
178}
179
180impl_box!(Box);
181impl_box!(Arc);
s2n_quic_dc/stream/socket/handle.rs

s2n_quic_dc/stream/socket/
handle.rs
