Skip to main content

dope_session/pool/client/
write.rs

1use crate::Codec;
2use crate::protocol::client::ClientProtocol;
3use dope_core::driver::{DriverOps, PoolDriver};
4use dope_core::profile::Profile;
5use dope_transport::Transport;
6
7use super::super::Pool;
8use super::super::core::PushError;
9use super::super::slot::{ConnFlags, CoreSlot};
10use super::Client;
11use super::connect_source::ConnectSource;
12use super::state::CLIENT_BYTES_CAP;
13use crate::SlotId;
14
15impl<P, T, S, C, F, B> Pool<Client<P, T, S>, C, F, B>
16where
17    P: ClientProtocol,
18    T: Transport,
19    T::Addr: Clone,
20    S: ConnectSource<T>,
21    C: Codec,
22    F: Profile,
23    B: PoolDriver,
24{
25    #[inline]
26    pub(super) fn try_submit_outbound(&mut self, driver: &mut B::Driver, slot_idx: usize) {
27        if !self.core.slots[slot_idx]
28            .conn
29            .flags
30            .contains(ConnFlags::LIVE)
31        {
32            return;
33        }
34        if self.core.slots[slot_idx]
35            .conn
36            .flags
37            .contains(ConnFlags::CLOSE_AFTER)
38        {
39            return;
40        }
41        if C::IS_PLAIN {
42            self.try_submit_outbound_plain(driver, slot_idx);
43        } else {
44            self.try_submit_outbound_tls(driver, slot_idx);
45        }
46    }
47
48    #[inline]
49    fn try_submit_outbound_plain(&mut self, driver: &mut B::Driver, slot_idx: usize) {
50        if self.core.slots[slot_idx].io.write.has_in_flight {
51            return;
52        }
53        let conn_fd;
54        let token_epoch: Option<(B::Token, u32)>;
55        let msg_ptr: *const libc::msghdr;
56        let close;
57        {
58            let external = self.core.slot_id_base + slot_idx;
59            let entry = &mut self.core.slots[slot_idx];
60            conn_fd = entry.conn.fd;
61            let user = entry.user.get_mut();
62            match user.session.fill_msghdr(CLIENT_BYTES_CAP) {
63                Some(m) => {
64                    msg_ptr = m as *const _;
65                    token_epoch = Some(entry.io.write.next_token::<B::Token>(external));
66                    close = false;
67                }
68                None => {
69                    msg_ptr = std::ptr::null();
70                    token_epoch = None;
71                    close = self.direction.protocol.wants_close(user.session.state());
72                }
73            }
74        }
75        if let Some((token, epoch)) = token_epoch {
76            if unsafe { driver.submit_send_msg_tagged(token, conn_fd, msg_ptr) } {
77                self.core.slots[slot_idx].io.write.mark_inflight(epoch);
78            }
79        } else if close {
80            self.disconnect_and_free(driver, slot_idx);
81        }
82    }
83
84    #[cold]
85    fn try_submit_outbound_tls(&mut self, driver: &mut B::Driver, slot_idx: usize) {
86        let cap = self.core.slots[slot_idx].buf.write_buf.len();
87        let head = self.core.slots[slot_idx].conn.write_head as usize;
88        let room = cap.saturating_sub(head);
89
90        if room > 0 {
91            let CoreSlot {
92                conn, buf, user, ..
93            } = &mut self.core.slots[slot_idx];
94            let user = user.get_mut();
95            let req_cap = room.min(CLIENT_BYTES_CAP);
96            if let Some(msg) = user.session.fill_msghdr(req_cap) {
97                let iov_count = msg.msg_iovlen as usize;
98
99                let iov_slice: &[libc::iovec] =
100                    unsafe { std::slice::from_raw_parts(msg.msg_iov, iov_count) };
101                let mut copied = 0u32;
102                for iov in iov_slice {
103                    let len = iov.iov_len;
104                    if (head + copied as usize) + len > cap {
105                        break;
106                    }
107
108                    let src = unsafe { std::slice::from_raw_parts(iov.iov_base.cast::<u8>(), len) };
109                    let off = head + copied as usize;
110                    buf.write_buf[off..off + len].copy_from_slice(src);
111                    copied += len as u32;
112                }
113                conn.write_head = (head + copied as usize).min(u16::MAX as usize) as u16;
114                user.session.ack_send(copied);
115            }
116        }
117
118        let close = {
119            let user = self.core.slots[slot_idx].user.get();
120            self.direction.protocol.wants_close(user.session.state())
121        };
122        if close {
123            self.core.slots[slot_idx]
124                .conn
125                .flags
126                .set(ConnFlags::CLOSE_AFTER, true);
127        }
128
129        self.core.encrypt_write_buf(slot_idx);
130        self.core.try_submit_codec(driver, slot_idx);
131    }
132
133    pub fn poke_outbound(
134        &mut self,
135        driver: &mut B::Driver,
136        conn_id: SlotId,
137    ) -> Result<(), PushError> {
138        let slot_idx = self.core.resolve_live_slot(conn_id)?;
139        self.try_submit_outbound(driver, slot_idx);
140        Ok(())
141    }
142
143    pub fn pump_outbound(&mut self, driver: &mut B::Driver) -> usize {
144        let mut count = 0usize;
145        for slot_idx in 0..self.core.slots.len() {
146            if !self.core.slots[slot_idx]
147                .conn
148                .flags
149                .contains(ConnFlags::LIVE)
150            {
151                continue;
152            }
153            self.try_submit_outbound(driver, slot_idx);
154            count += 1;
155        }
156        count
157    }
158
159    #[cold]
160    fn handle_write_event_tls(&mut self, driver: &mut B::Driver, slot_idx: usize, result: i32) {
161        match self.core.handle_tls_inflight(slot_idx, result) {
162            super::super::core::TlsAck::NoOp => {}
163            super::super::core::TlsAck::Err { slot_idx } => {
164                self.disconnect_and_free(driver, slot_idx)
165            }
166            super::super::core::TlsAck::Progressed { slot_idx } => {
167                self.core.try_submit_codec(driver, slot_idx);
168                self.try_submit_outbound(driver, slot_idx);
169                let close = self.core.slots[slot_idx]
170                    .conn
171                    .flags
172                    .contains(ConnFlags::CLOSE_AFTER)
173                    || self.core.codec_close_pending(slot_idx);
174                if self.core.codec_idle(slot_idx) && close {
175                    self.disconnect_and_free(driver, slot_idx);
176                }
177            }
178        }
179    }
180}
181
182impl<P, T, S, C, F, B> super::super::PoolWrite<B> for Pool<Client<P, T, S>, C, F, B>
183where
184    P: ClientProtocol,
185    T: Transport,
186    T::Addr: Clone,
187    S: ConnectSource<T>,
188    C: Codec,
189    F: Profile,
190    B: PoolDriver,
191{
192    fn on_write_event(
193        &mut self,
194        driver: &mut B::Driver,
195        token: B::Token,
196        result: i32,
197        notif: bool,
198    ) {
199        match self.core.classify_write(token, result, notif) {
200            super::super::core::WriteDispatch::Skipped => {}
201            super::super::core::WriteDispatch::Tls { slot_idx, result } => {
202                self.handle_write_event_tls(driver, slot_idx, result);
203            }
204            super::super::core::WriteDispatch::Plain { slot_idx, result } => {
205                if !self.core.slots[slot_idx]
206                    .conn
207                    .flags
208                    .contains(ConnFlags::LIVE)
209                {
210                    return;
211                }
212                if result < 0 {
213                    self.disconnect_and_free(driver, slot_idx);
214                    return;
215                }
216                let n = result as u32;
217                if n > 0 {
218                    let user = self.core.slots[slot_idx].user.get_mut();
219                    user.session.ack_send(n);
220                }
221                if self.core.slots[slot_idx]
222                    .conn
223                    .flags
224                    .contains(ConnFlags::CLOSE_AFTER)
225                {
226                    self.disconnect_and_free(driver, slot_idx);
227                    return;
228                }
229                self.try_submit_outbound(driver, slot_idx);
230            }
231        }
232    }
233}