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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
//! The `SocketManager` is responsible for managing the socket connection and
//! handling the messages sent and received over the socket.
//!
//! It is responsible for binding the socket, sending and receiving messages,
//! and shutting down the socket when it is no longer needed.
use crate::{
Error, TCP_TIMEOUT_GENERAL_INACTIVITY,
client::ClientOptions,
connection,
message_codec::MessageCodec,
messages::{MessageError, OwnedMessage},
};
use futures::{SinkExt, StreamExt};
use std::{net::SocketAddr, string::ToString};
use tokio::{
net::tcp::{OwnedReadHalf, OwnedWriteHalf},
select,
sync::mpsc,
};
use tokio_util::codec::{FramedRead, FramedWrite};
use tracing::{debug, error, info, trace};
/// 1-to-1 mapping of the socket manager to the client (currently)
/// There is only one socket manager per client.
#[derive(Debug)]
pub(crate) struct SocketManager<Conn> {
/// Receiver used to receive messages from the socket
/// This is the channel that the socket manager uses to send messages back up to the client
receiver: mpsc::Receiver<Result<OwnedMessage, MessageError>>,
/// Sender used to send messages to the socket
sender: mpsc::Sender<OwnedMessage>,
local_port: u16,
session_id: u16,
_phantom: std::marker::PhantomData<Conn>,
}
impl<Conn> SocketManager<Conn>
where
Conn: connection::Connector + 'static + Send + Sync,
{
/// Creates a new `SocketManager` instance
///
/// TCP socket is bound to the specified address
///
/// # Errors
/// Returns an [`Error`] if the client logical address is invalid or the
/// TCP connection cannot be established
pub async fn bind(
client_options: ClientOptions,
gateway_address: SocketAddr,
) -> Result<Self, Error> {
if !client_options
.client_logical_address
.is_valid_client_address()
{
return Err(Error::InvalidClientLogicalAddress(
client_options.client_logical_address,
));
}
// Call the connection - this might be overridden by the user
let (rx, tx) = match Conn::establish_connection(gateway_address).await {
Ok((rx, tx)) => (rx, tx),
Err(e) => {
error!("Failed to establish connection: {e} on {gateway_address}");
return Err(e);
}
};
let socket_read_stream = FramedRead::new(rx, MessageCodec::new());
let socket_write_sink = FramedWrite::new(tx, MessageCodec::new());
let (rx_tx, rx_rx) = mpsc::channel(16);
let (tx_tx, tx_rx) = mpsc::channel(16);
Self::spawn_socket_loop(rx_tx, tx_rx, socket_read_stream, socket_write_sink);
Ok(Self {
receiver: rx_rx,
sender: tx_tx,
local_port: gateway_address.port(),
session_id: 0,
_phantom: std::marker::PhantomData,
})
}
/// Send a message to the target address
///
/// # Errors
/// Returns an [`Error::ConnectionClosed`] if the message cannot be sent
pub async fn send(&mut self, message: OwnedMessage) -> Result<(), Error> {
self.sender.send(message).await.map_err(|e| {
error!("Failed to send message: {}", e);
Error::ConnectionClosed
})?;
self.session_id += 1;
Ok(())
}
/// Receive a message from the receiver/Request channel
pub async fn receive(&mut self) -> Option<Result<OwnedMessage, MessageError>> {
self.receiver.recv().await
}
/// Return the local TCP port this socket is bound to
#[must_use]
pub fn port(&self) -> u16 {
self.local_port
}
/// Shutdown the socket manager
/// This will close the socket and stop the event loop
/// It will also drop the sender and receiver channels
pub async fn shut_down(self) {
let Self {
sender,
mut receiver,
..
} = self;
trace!("Shutting down socket manager - Sender");
// First stop accepting messages before we drop
receiver.close();
drop(sender);
trace!("receive any remaining messages");
_ = receiver.recv().await;
}
/// Spawn the socket loop to get messages from the socket
fn spawn_socket_loop(
rx_tx: mpsc::Sender<Result<OwnedMessage, MessageError>>,
mut tx_rx: mpsc::Receiver<OwnedMessage>,
mut socket_read_stream: FramedRead<OwnedReadHalf, MessageCodec>,
mut socket_write_sink: FramedWrite<OwnedWriteHalf, MessageCodec>,
) {
tokio::spawn(async move {
// General TCP activity timeout
// this is used to close the socket if there is no activity for a while
let mut last_activity = tokio::time::Instant::now();
loop {
select! {
() = tokio::time::sleep_until(last_activity + TCP_TIMEOUT_GENERAL_INACTIVITY) => {
info!("General inactivity timeout reached, closing socket");
// Breaking out of this loop drops the socket, which closes the
// connection; callers observe that through the stream ending.
break;
}
// Once there is information in the Response/Read stream we'll do work on it
// and send it along to the receiver on the other end
//
result = socket_read_stream.next() => {
match result {
// Decoding the message can fail, so we handle that here
Some(Err(e)) => {
last_activity = tokio::time::Instant::now();
// Socket-level errors from the tokio layer arrive as
// `MessageError::Std`, preserving OS error detail.
// `MessageError::Io` is encode-side only (embedded-io
// short writes) and is not expected on this RX path, but
// is classified identically so the match stays honest.
let reset = match &e {
MessageError::Std(io_err) => {
Some(io_err.kind() == std::io::ErrorKind::ConnectionReset)
}
MessageError::Io(kind) => {
Some(*kind == embedded_io::ErrorKind::ConnectionReset)
}
_ => None,
};
if let Some(was_reset) = reset {
if was_reset {
info!("Connection reset by peer, closing socket: {e}");
} else {
error!(concat!("{}\n",
"Check that you are not sending too many requests to the server.",
"The server may be closing the connection due to overload."
), e);
}
// Either way the socket is unusable; exit the read loop.
break;
}
error!("Error decoding message: {:?}", e.to_string());
// send a MessageError to the receiver
let _ = rx_tx.send(Err(e)).await;
}
Some(message) => {
// Update the last activity time
last_activity = tokio::time::Instant::now();
trace!("A: STREAM INCOMING: {:?}", message);
if rx_tx.send(message).await.is_err() {
info!("Socket Dropping");
// The receiver has been dropped, so we should exit
break;
}
}
None => {
info!("Socket Dropping");
// The sender has been dropped, so we should exit
break;
}
}
},
// maps to self.receiver
message = tx_rx.recv() => {
let Some(message) = message else {
debug!("Socket Dropping");
// The sender has been dropped, so we should exit
break;
};
// Update the last activity time
last_activity = tokio::time::Instant::now();
trace!("OUTGOING: {:?}", message);
if let Err(e) = socket_write_sink.send(&message).await {
error!("Error sending message to socket: {:?}", e);
break;
}
}
}
}
});
}
}