Properly 2 TcpStream - rust

I am working on a TCP proxy server where a client connects to it and it starts a new connection to a backend and forward all packets (bidirectional). The function that handles this part looks like this.
pub fn route(source: TcpStream, worker: Worker) -> Result<()> {
// `source` and `destination` are both `TcpStream`
let mut source = source;
let mut destination = TcpStream::connect(worker.address)?;
let mut source_copy = source.try_clone()?;
let mut destination_copy = destination.try_clone()?;
let src2dst = std::thread::spawn(move || {
std::io::copy(&mut source_copy, &mut destination_copy).unwrap();
});
let dst2src = std::thread::spawn(move || {
std::io::copy(&mut destination, &mut source).unwrap();
// source.shutdown(Shutdown::Both).unwrap(); <--- if this line is commented, it will stuck
});
src2dst.join().unwrap();
dst2src.join().unwrap();
Ok(())
}
However, in its current form, this function will stuck. In particular, if I shutdown the source when destination stops writing, it won't block. But I am still not sure why this works (or why it does not). I am currently using it only to proxy HTTP traffic, and it does not appear to have problems. But I am not sure if it works for generic TCP. What is the proper way to do this?

std::io::copy reads from the reader until it gets an end-of-file condition. The way to signal an end-of-file condition on a TCP connection is to send a FIN packet, which is realized in software by calling shutdown to shut down the write half of the connection.
HTTP allows reusing a single connection to send multiple requests, one after the other. There is enough information in a request for the recipient to determine where the request ends, and likewise for responses. In your case, it appears the backend shuts down its write half of the connection, which causes your io::copy to return. If the HTTP response has Connection: close in its headers, then you must shut down the write half, or else the client will hang waiting for the response body because it never receives a FIN packet.
TL;DR: When the copy from A to B returns normally, you should shut down the write half of B to forward the end-of-file condition (FIN packet). This goes for the copies in both directions, and is valid for TCP in general. Shutting down the read half of sockets doesn't appear to be necessary (it doesn't generate any TCP packets).

Related

Rust on ESP32 - How to send (and receive) data using the MQTT protocol to AWS IoT (Core)?

First off: I know running Rust on an ESP32 isn't a very common practice yet, and some (quite a bit of) trouble is to be expected. But I seem to have hit a roadblock.
What works:
flashing and running the code on an ESP32
passing along the certificates in the src/certificates directory
WiFi connection (simple WPA Personal, nothing fancy like WPA Enterprise)
publishing and suscribing to topics using MQTT
What doesn't work:
publising and subscribing to AWS IoT (Core) using MQTT. This needs certificates, and as far as I'm aware I'm handling this properly (see code below).
Some additional info (see code below):
server.cert.crt is renamed from the AWS provided root-CA.crt
client.cert.pem is renamed from the AWS provided my-thing-rev1.cert.pem
client.private.key is renamed from the AWS provided my-thing-rev1.private.key
I also received my-thing-rev1.public.key and my-thing-rev1-Policy, but I don't think I need these...?
I know this is not the proper way of implementing this (I should not provide the certificates directly, instead use a service to get them, but this is a very basic POC)
the code works fine if I don't want to connect to AWS, but instead use my own broker or broker.emqx.io for testing (even with the certificates included)
This is the code I'm currently using (heavily based on Rust on ESP32 STD demo app):
use embedded_svc::httpd::Result;
use embedded_svc::mqtt::client::{Connection, MessageImpl, QoS};
use esp_idf_svc::mqtt::client::{EspMqttClient, MqttClientConfiguration};
use esp_idf_svc::tls::X509;
use esp_idf_sys::EspError;
// other needed imports (not relevant here)
extern crate dotenv_codegen;
extern crate core;
const AWS_IOT_ENDPOINT: &str = dotenv!("AWS_IOT_ENDPOINT");
const AWS_IOT_CLIENT_ID: &str = dotenv!("AWS_IOT_CLIENT_ID");
const AWS_IOT_TOPIC: &str = dotenv!("AWS_IOT_TOPIC");
fn main() -> Result<()> {
esp_idf_sys::link_patches();
// other code
let mqtt_client: EspMqttClient<ConnState<MessageImpl, EspError>> = test_mqtt_client()?;
// more code
Ok(())
}
fn convert_certificate(mut certificate_bytes: Vec<u8>) -> X509<'static> {
// append NUL
certificate_bytes.push(0);
// convert the certificate
let certificate_slice: &[u8] = unsafe {
let ptr: *const u8 = certificate_bytes.as_ptr();
let len: usize = certificate_bytes.len();
mem::forget(certificate_bytes);
slice::from_raw_parts(ptr, len)
};
// return the certificate file in the correct format
X509::pem_until_nul(certificate_slice)
}
fn test_mqtt_client() -> Result<EspMqttClient<ConnState<MessageImpl, EspError>>> {
info!("About to start MQTT client");
let server_cert_bytes: Vec<u8> = include_bytes!("certificates/server.cert.crt").to_vec();
let client_cert_bytes: Vec<u8> = include_bytes!("certificates/client.cert.pem").to_vec();
let private_key_bytes: Vec<u8> = include_bytes!("certificates/client.private.key").to_vec();
let server_cert: X509 = convert_certificate(server_cert_bytes);
let client_cert: X509 = convert_certificate(client_cert_bytes);
let private_key: X509 = convert_certificate(private_key_bytes);
// TODO: fix the following error: `E (16903) esp-tls-mbedtls: mbedtls_ssl_handshake returned -0x7280`
let conf = MqttClientConfiguration {
client_id: Some(AWS_IOT_CLIENT_ID),
crt_bundle_attach: Some(esp_idf_sys::esp_crt_bundle_attach),
server_certificate: Some(server_cert),
client_certificate: Some(client_cert),
private_key: Some(private_key),
..Default::default()
};
let (mut client, mut connection) =
EspMqttClient::new_with_conn(AWS_IOT_ENDPOINT, &conf)?;
info!("MQTT client started");
// Need to immediately start pumping the connection for messages, or else subscribe() and publish() below will not work
// Note that when using the alternative constructor - `EspMqttClient::new` - you don't need to
// spawn a new thread, as the messages will be pumped with a backpressure into the callback you provide.
// Yet, you still need to efficiently process each message in the callback without blocking for too long.
//
// Note also that if you go to http://tools.emqx.io/ and then connect and send a message to the specified topic,
// the client configured here should receive it.
thread::spawn(move || {
info!("MQTT Listening for messages");
while let Some(msg) = connection.next() {
match msg {
Err(e) => info!("MQTT Message ERROR: {}", e),
Ok(msg) => info!("MQTT Message: {:?}", msg),
}
}
info!("MQTT connection loop exit");
});
client.subscribe(AWS_IOT_TOPIC, QoS::AtMostOnce)?;
info!("Subscribed to all topics ({})", AWS_IOT_TOPIC);
client.publish(
AWS_IOT_TOPIC,
QoS::AtMostOnce,
false,
format!("Hello from {}!", AWS_IOT_TOPIC).as_bytes(),
)?;
info!("Published a hello message to topic \"{}\".", AWS_IOT_TOPIC);
Ok(client)
}
Here are the final lines of output when I try to run this on the device (it's setup to compile and flash to the device and monitor (debug mode) when running cargo run):
I (16913) esp32_aws_iot_with_std: About to start MQTT client
I (16923) esp32_aws_iot_with_std: MQTT client started
I (16923) esp32_aws_iot_with_std: MQTT Listening for messages
I (16933) esp32_aws_iot_with_std: MQTT Message: BeforeConnect
I (17473) esp-x509-crt-bundle: Certificate validated
E (19403) MQTT_CLIENT: mqtt_message_receive: transport_read() error: errno=119 # <- This is the actual error
E (19403) MQTT_CLIENT: esp_mqtt_connect: mqtt_message_receive() returned -1
E (19413) MQTT_CLIENT: MQTT connect failed
I (19413) esp32_aws_iot_with_std: MQTT Message ERROR: ESP_FAIL
I (19423) esp32_aws_iot_with_std: MQTT Message: Disconnected
E (19433) MQTT_CLIENT: Client has not connected
I (19433) esp32_aws_iot_with_std: MQTT connection loop exit
I (24423) esp_idf_svc::eventloop: Dropped
I (24423) esp_idf_svc::wifi: Stop requested
I (24423) wifi:state: run -> init (0)
I (24423) wifi:pm stop, total sleep time: 10737262 us / 14862601 us
W (24423) wifi:<ba-del>idx
I (24433) wifi:new:<1,0>, old:<1,1>, ap:<1,1>, sta:<1,0>, prof:1
W (24443) wifi:hmac tx: ifx0 stop, discard
I (24473) wifi:flush txq
I (24473) wifi:stop sw txq
I (24473) wifi:lmac stop hw txq
I (24473) esp_idf_svc::wifi: Stopping
I (24473) esp_idf_svc::wifi: Disconnect requested
I (24473) esp_idf_svc::wifi: Stop requested
I (24483) esp_idf_svc::wifi: Stopping
I (24483) wifi:Deinit lldesc rx mblock:10
I (24503) esp_idf_svc::wifi: Driver deinitialized
I (24503) esp_idf_svc::wifi: Dropped
I (24503) esp_idf_svc::eventloop: Dropped
Error: ESP_FAIL
This error seems to indicate the buffer holding the incoming data is full and can't hold any more data, but I'm not sure. And I definately don't know how to fix it.
(I assume the actual certificate handling is done properly)
When I run the following command, I do get the message in AWS IoT (MQTT test client):
mosquitto_pub -h my.amazonawsIoT.com --cafile server.cert.crt --cert client.cert.pem --key client.private.key -i basicPubSub -t my/topic -m 'test'
Does anyone have some more experience with this who can point me in the right direction?
Is this actually a buffer error, and if so: how do I mitigate this error? Do I need to increase the buffer size somehow (it is running on a basic ESP32 revision 1, ESP32_Devkitc_v4, if that helps). As far as I can tell this version has a 4MB flash size, so that might explain the buffer overlow, although I think this should be enough. The total memory used is under 35% of the total storage (App/part. size: 1347344/4128768 bytes, 32.63%)
UPDATE 1: I have been made aware that this data is stored in RAM, not in flash memory (didn't cross my mind at the time), but I'm not entirely sure on how large the RAM on my specific device is (ESP32 revision 1, ESP32_Devkitc_v4). My best guess is 320KB, but I'm not sure.
UPDATE 2: I've tried changing the buffer size like so:
let conf = MqttClientConfiguration {
client_id: Some(AWS_IOT_CLIENT_ID),
crt_bundle_attach: Some(esp_idf_sys::esp_crt_bundle_attach),
server_certificate: Some(server_cert),
client_certificate: Some(client_cert),
private_key: Some(private_key),
buffer_size: 50, // added this (tried various sizes)
out_buffer_size: 50, // added this (tried various sizes)
..Default::default()
};
I've tried various combinations, but this doesn't seem to change much: either I get the exact same error, or this one (when choosing smaller numbers, for example 10):
E (18303) MQTT_CLIENT: Connect message cannot be created
E (18303) MQTT_CLIENT: MQTT connect failed
E (18313) MQTT_CLIENT: Client has not connected
I'm not sure how big this buffer size should be (when sending simple timestamps to AWS IoT), and can't find any documentation on what this number represents: is it in Bit, KiloBit, ... No idea.

how to gracefully shutdown websocket clients using tokio

I am implementing a websocket chat application where I want to gracefully shutdown all clients when the server stop because of the ctrl+c signal.
I am listening to incoming events using mio poll and tokens. Any new socket connection is registered with mio poll and any event received on the socket is successfully captured on the polling.
My initial idea was to use tokio::select with listen events and shutdown as 2 branches. But, I guess this design needs some modification to enable graceful shutdown.
// get_events function
fn get_events(poll, conn) -> Option<mio::Event>{
let shutdown = tokio::signal::ctrl_c();
let res = tokio::select!{
res = get_poll_events(poll) => { // returns events asynchronously, its an async method
// ... handler when events are received
res
},
_ = shutdown => {
// ... handler when ctrl+c signal is received
// send close connection message to TcpStream
return None;
}
};
Some(res)
}
// main function
let poll = Poll::new();
let shared_poll = Arc::new(Mutex::new(poll));
let conn: Arc<Mutex<HashMap<Token, WebSocketClient>>> = Arc::new(Mutex::new(HashMap::new()));
let (shutdown_notifier, shutdown_receiver) = mpsc::channel(10);
loop {
let res = WebSocketServer::get_events(shared_poll.clone(), conn.clone()).await;
if let None = res {
drop(shutdown_notifier); // drop the original mpsc::Sender, cloned in all clients.
let _ = shutdown_receiver.recv().await; // mpsc::Receiver only returns error when all clients are dropped, thus dropping the senders inside them.
break; // stop the program
}
// use event to register a new client with mpsc::sender
// if readable; accept incoming message
// broadcast message to other subscribers
// ... other tasks
...
}
On Ctrl+c signal received the shutdown handler gets activated. Post that, close connection message is sent on the TcpStream to all active clients. At this point, everything is as expected. All the clients receive close messages. But the real issue starts here, since the shutdown handler executes, the get_poll_events branch of tokio::select gets dropped. Meaning, no more polling for the incoming events on the registered sources (TcpStreams).
Ideally, the server should only close when it has received back the close message response from all the clients. The clients communicate back on the TcpStream, but since there is no active mechanism to listen to these events, I am unable to capture them and hence not able to drop the clients. However, instead of sending close message to clients if I dropped all the clients manually, then there was no need to listen to close message response and things worked albeit an incorrect implementation.
Tokio::select isn't an ideal choice as far as I can guess, but I am unable to come up with a solution on how to implement this case, where the server is still active to listen to all the close message responses from clients. It can then close when all the clients have gone out of scope in the received close messages.
What would be a way to achieve this functionality? TIA.

How to exchange data in a block_on section?

I'm learning Rust and Tokio and I suspect I may be going in the wrong direction.
I'm trying to open a connection to a remote server and perform a handshake. I want to use non-blocking IO so I'm using Tokio's thread pool. The handshake needs to be performed quickly or the remote will close the socket so I'm trying to chain the message exchange in a single block_on section:
let result: Result<(), Box<dyn std::error::Error>> = session
.runtime()
.borrow_mut()
.block_on(async {
let startup = startup(session.configuration());
stream.write_all(startup.as_ref()).await?;
let mut buffer:Vec<u8> = Vec::new();
let mut tmp = [0u8; 1];
loop {
let total = stream.read(&mut tmp).await;
/*
if total == 0 {
break;
}
*/
if total.is_err() {
break;
}
buffer.extend(&tmp);
}
Ok(())
});
My problem is what to do when there are no more bytes in the socket to read. My current implementation reads the response and after the last byte hangs, I believe because the socket is not closed. I thought checking for 0 bytes read would be enough but the call to read() never returns.
What's the best way to handle this?
From your comment:
Nope, the connection is meant to remain open.
If you read from an open connection, the read will block until there are enough bytes to satisfy it or the other end closes the connection, similar to how blocking reads work in C. Tokio is working as-intended.
If closing the stream does not signal the end of a message, then you will have to do your own work to figure out when to stop reading and start processing. A simple way would to just prefix the request with a length, and only read that many bytes.
Note that you'd have to do the above no matter what API you'd use. The fact that you use tokio or not doesn't really answer the fundamental question of "when is the message over".

How to cancel future/close stream in multithreaded tokio?

Based on tokio's example at https://github.com/tokio-rs/tokio/blob/master/examples/proxy.rs
let (mut ri, mut wi) = inbound.split();
let (mut ro, mut wo) = outbound.split();
let client_to_server = io::copy(&mut ri, &mut wo);
let server_to_client = io::copy(&mut ro, &mut wi);
try_join(client_to_server, server_to_client).await?;
Ok(())
I have a modified version so that I can handle the termination of each connection as in:
// Server will disconnect their side normally 8s later, from what I've observed
let server_to_client = io::copy(&mut ro, &mut wi).map(|f| {
server_session_time = server_start_time.elapsed().unwrap();
f
});
// Normally, this will stop first, as the client disconnects as soon as he has the results...
let client_to_server = io::copy(&mut ri, &mut wo).map(|f| {
client_session_time = client_start_time.elapsed().unwrap();
f
});
// Join on both
match try_join(client_to_server, server_to_client).await {...}
This has allowed me to time correctly the connected time for the client side, since the clients immediately close connection upon receiving the answer, while the proxied server seems to take (in my case 8s) to close.
Given this structure of code, is there any possibility to terminate the downstream connection from server_to_client, once I exit the future of the client_to_server (i.e. not wait the 8s that I observe that it takes to be shutdown)?
Ok with a few more examples, was able to understand what I had to do.
For any people coming back to this question in the future, what is needed is that you implement the bidirectional copy yourself based on the 4 futures of each of the reads and writes with tokio::select!.
That will allow to access to all the streams and when one of them terminates, it is your option if you want to complete processing the others or just stop.
As it is above there is no way to "cancel" the "other" copy...
You can look both at the implementation of io::copy https://github.com/tokio-rs/tokio-io/blob/master/src/copy.rs and tokio::select https://docs.rs/tokio/0.2.20/tokio/macro.select.html, to build your 4-way select.

How to use tokio's UdpSocket to handle messages in a 1 server: N clients setup?

What I want to do:
... write a (1) server/ (N) clients (network-game-)architecture that uses UDP sockets as underlying base for communication.
Messages are sent as Vec<u8>, encoded via bincode (crate)
I also want to be able to occasionally send datagrams that can exceed the typical max MTU of ~1500 bytes and be correctly assembled on receiver end, including sending of ack-messages etc. (I assume I'll have to implement that myself, right?)
For the UdpSocket I thought about using tokio's implementation and maybe framed. I am not sure whether this is a good choice though, as it seems that this would introduce an unnecessary step of mapping Vec<u8> (serialized by bincode) to Vec<u8> (needed by UdpCodec of tokio) (?)
Consider this minimal code-example:
Cargo.toml (server)
bincode = "1.0"
futures = "0.1"
tokio-core = "^0.1"
(Serde and serde-derive are used in shared crate where the protocol is defined!)
(I want to replace tokio-core with tokio asap)
fn main() -> () {
let addr = format!("127.0.0.1:{port}", port = 8080);
let addr = addr.parse::<SocketAddr>().expect(&format!("Couldn't create valid SocketAddress out of {}", addr));
let mut core = Core::new().unwrap();
let handle = core.handle();
let socket = UdpSocket::bind(&addr, &handle).expect(&format!("Couldn't bind socket to address {}", addr));
let udp_future = socket.framed(MyCodec {}).for_each(|(addr, data)| {
socket.send_to(&data, &addr); // Just echo back the data
Ok(())
});
core.run(udp_future).unwrap();
}
struct MyCodec;
impl UdpCodec for MyCodec {
type In = (SocketAddr, Vec<u8>);
type Out = (SocketAddr, Vec<u8>);
fn decode(&mut self, src: &SocketAddr, buf: &[u8]) -> io::Result<Self::In> {
Ok((*src, buf.to_vec()))
}
fn encode(&mut self, msg: Self::Out, buf: &mut Vec<u8>) -> SocketAddr {
let (addr, mut data) = msg;
buf.append(&mut data);
addr
}
}
The problem here is:
let udp_future = socket.framed(MyCodec {}).for_each(|(addr, data)| {
| ------ value moved here ^^^^^^^^^^^^^^ value captured here after move
|
= note: move occurs because socket has type tokio_core::net::UdpSocket, which does not implement the Copy trait
The error makes total sense, yet I am not sure how I would create such a simple echo-service. In reality, the handling of a message involves a bit more logic ofc, but for the sake of a minimal example, this should be enough to give a rough idea.
My workaround is an ugly hack: creating a second socket.
Here's the signature of UdpSocket::framed from Tokio's documentation:
pub fn framed<C: UdpCodec>(self, codec: C) -> UdpFramed<C>
Note that it takes self, not &self; that is, calling this function consumes the socket. The UdpFramed wrapper owns the underlying socket when you call this. Your compilation error is telling you that you're moving socket when you call this method, but you're also trying to borrow socket inside your closure (to call send_to).
This probably isn't what you want for real code. The whole point of using framed() is to turn your socket into something higher-level, so you can send your codec's items directly instead of having to assemble datagrams. Using send or send_to directly on the socket will probably break the framing of your message protocol. In this code, where you're trying to implement a simple echo server, you don't need to use framed at all. But if you do want to have your cake and eat it and use both framed and send_to, luckily UdpFramed still allows you to borrow the underlying UdpSocket, using get_ref. You can fix your problem this way:
let framed = {
let socket = UdpSocket::bind(&addr, &handle).expect(&format!("Couldn't bind socket to address {}", addr));
socket.framed(MyCodec {})
}
let udp_future = framed.for_each(|(addr, data)| {
info!(self.logger, "Udp packet received from {}: length: {}", addr, data.len());
framed.get_ref().send_to(&data, &addr); // Just echo back the data
Ok(())
});
I haven't checked this code, since (as Shepmaster rightly pointed out) your code snippet has other problems, but it should give you the idea anyway. I'll repeat my warning from earlier: if you do this in real code, it will break the network protocol you're using. get_ref's documentation puts it like this:
Note that care should be taken to not tamper with the underlying stream of data coming in as it may corrupt the stream of frames otherwise being worked with.
To answer the new part of your question: yes, you need to handle reassembly yourself, which means your codec does actually need to do some framing on the bytes you're sending. Typically this might involve a start sequence which cannot occur in the Vec<u8>. The start sequence lets you recognise the start of the next message after a packet was lost (which happens a lot with UDP). If there's no byte sequence that can't occur in the Vec<u8>, you need to escape it when it does occur. You might then send the length of the message, followed by the data itself; or just the data, followed by an end sequence and a checksum so you know none was lost. There are pros and cons to these designs, and it's a big topic in itself.
You also need your UdpCodec to contain data: a map from SocketAddr to the partially-reassembled message that's currently in progress. In decode, if you are given the start of a message, copy it into the map and return Ok. If you are given the middle of a message, and you already have the start of a message in the map (for that SocketAddr), append the buffer to the existing buffer and return Ok. When you get to the end of the message, return the whole thing and empty the buffer. The methods on UdpCodec take &mut self in order to enable this use case. (NB In theory, you should also deal with packets arriving out of order, but that's actually quite rare in the real world.)
encode is a lot simpler: you just need to add the same framing and copy the message into the buffer.
Let me reiterate here that you don't need to and shouldn't use the underlying socket after calling framed() on it. UdpFramed is both a source and a sink, so you use that one object to send the replies as well. You can even use split() to get separate Stream and Sink implementations out of it, if that makes the ownership easier in your application.
Overall, now I've seen how much of the problem you're struggling with, I'd recommend just using several TCP sockets instead of UDP. If you want a connection-oriented, reliable protocol, TCP already exists and does that for you. It's very easy to spend a lot of time making a "reliable" layer on top of UDP that is both slower and less reliable than TCP.

Resources