How to retrieve information from the tokio-proto connection handshake? - rust

I'm figuring out how to use the tokio-proto crate, particularly on the handshake made when a connection is established. I've got the example from the official documentation working:
impl<T: AsyncRead + AsyncWrite + 'static> ClientProto<T> for ClientLineProto {
type Request = String;
type Response = String;
/// `Framed<T, LineCodec>` is the return value of `io.framed(LineCodec)`
type Transport = Framed<T, line::LineCodec>;
type BindTransport = Box<Future<Item = Self::Transport, Error = io::Error>>;
fn bind_transport(&self, io: T) -> Self::BindTransport {
// Construct the line-based transport
let transport = io.framed(line::LineCodec);
// Send the handshake frame to the server.
let handshake = transport.send("You ready?".to_string())
// Wait for a response from the server, if the transport errors out,
// we don't care about the transport handle anymore, just the error
.and_then(|transport| transport.into_future().map_err(|(e, _)| e))
.and_then(|(line, transport)| {
// The server sent back a line, check to see if it is the
// expected handshake line.
match line {
Some(ref msg) if msg == "Bring it!" => {
println!("CLIENT: received server handshake");
Ok(transport)
}
Some(ref msg) if msg == "No! Go away!" => {
// At this point, the server is at capacity. There are a
// few things that we could do. Set a backoff timer and
// try again in a bit. Or we could try a different
// remote server. However, we're just going to error out
// the connection.
println!("CLIENT: server is at capacity");
let err = io::Error::new(io::ErrorKind::Other, "server at capacity");
Err(err)
}
_ => {
println!("CLIENT: server handshake INVALID");
let err = io::Error::new(io::ErrorKind::Other, "invalid handshake");
Err(err)
}
}
});
Box::new(handshake)
}
}
But the official docs only mention a handshake without stateful information. Is there a common way to retrieve and store useful data from the handshake?
For example, if during the handshake (in the first message after the connection is established) the server sends some key that should be used later by the client, how should the ClientProto implementation look into that key? And where should it be stored?

You can add fields to ClientLineProto, so this should work:
pub struct ClientLineProto {
handshakes: Arc<Mutex<HashMap<String, String>>>
}
And then you can reference it and store data as needed:
let mut handshakes = self.handshakes.lock();
handshakes.insert(handshake_key, "Blah blah handshake data")
This sort of access would work in bind_transport() for storing things. Then when you create the Arc::Mutex::HashMap in your main() function and you will have access to the whole thing in the serve() method as well, which means you can pass it in to the Service object instantiation and then the handshakes will be available during call().

Related

Rust, libzmq, client only send message, when followed by receive

I am testing Rust libzmq client
https://crates.io/crates/libzmq
https://docs.rs/libzmq/0.2.5/libzmq/struct.ClientBuilder.html
and stumbled upon this weird behavior.
This Client would not send the message:
use libzmq::{prelude::*, *, ServerBuilder, ClientBuilder, TcpAddr};
fn main() -> Result<(), String> {
let saddr = format!("{}:{}", "192.168.1.206","5540");
println!("Strating client for {}", saddr);
let addr2: TcpAddr = saddr.try_into().unwrap();
let client = ClientBuilder::new()
.connect(&addr2)
.build().unwrap();
client.send("test").unwrap();
println!("Finished");
Ok(())
}
this would send:
use libzmq::{prelude::*, *, ServerBuilder, ClientBuilder, TcpAddr};
fn main() -> Result<(), String> {
let saddr = format!("{}:{}", "192.168.1.206","5540");
println!("Strating client for {}", saddr);
let addr2: TcpAddr = saddr.try_into().unwrap();
let client = ClientBuilder::new()
.connect(&addr2)
.build().unwrap();
client.send("test").unwrap();
let mut msg = Msg::new();
client.recv(&mut msg).unwrap();
println!("Finished");
Ok(())
}
For receiving I am using python server (but I tested on Rust as well):
import zmq
if __name__ == '__main__':
context = zmq.Context()
socket = context.socket(zmq.SERVER)
socket.bind("tcp://0.0.0.0:5540")
print("waiting for hand shake")
m = socket.recv(copy=False)
print("got it...")
socket.send(b'READY', routing_id=m.routing_id)
print("sent reply")
socket.close()
context.term()
Output on server side from code 1:
waiting for hand shake
[blocked]
Output on server side from code 2:
waiting for hand shake
got it...
sent reply
using this version for Rust:
[dependencies]
libzmq = "0.2.5"
The libzmq crate is just a wrapper for the libzmq C library. In the documentation for the C library, you will find this note:
NOTE: A successful invocation of zmq_msg_send() does not indicate that the
message has been transmitted to the network, only that it has been queued on
the 'socket' and 0MQ has assumed responsibility for the message.
In your first example, since the program exits immediately after calling send, the message is still in the libzmq queue and has not been transmitted yet. In your second example, the message gets transmitted while your program is waiting in the recv call.

How to know when the server has received the whole request?

I am implementing the HTTP/1.1 protocol from scratch for academic purpose. I have implemented the RequestBuilder which builds the request object successively from the buffer passed. This is the code to handle the opened socket.
async fn process_socket(stream: TcpStream) -> Result<Request> {
let mut request_builder = RequestBuilder::new();
let mut buffer: [u8; 1024] = unsafe { MaybeUninit::uninit().assume_init() };
loop {
stream.readable().await?;
match stream.try_read(&mut buffer) {
Ok(0) => {
break;
}
Ok(n) => (),
Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
continue;
}
Err(e) => {
return Err(e.into());
}
}
request_builder.parse(&buffer);
}
let request = request_builder.build()?;
Ok(request)
}
request_builder.parse(&buffer); will take the next part of the buffer and parses the request further. My question is, how to break the loop when the client has sent the whole request. When I make a request to the server using curl localhost:8080, the whole request is parsed.
Expected behaviour
The loop would have been broken after reading the whole request stream.
Actual behaviour
The loop is stuck at stream.readable().await?; after reading the whole request into buffer. Currently, when I kill curl command using Ctrl+C, the loop is broken using Ok(0), but I want it to break after reading the who
You need to interpret the HTTP request, as the TCP connection will not get half-closed by a client. A FIN by the client which would violate the protocol) is the only way readable() returns (with an Err) unless the client sends more data (which breaks the HTTP specification).

How to inject an immutable object into service code in actix-web?

I'm using actix framework to build a server that should support an opportunity to show both age/balance to a user given a user_id:
fn show_balance(req: &HttpRequest) -> HttpResponse {
let client = create_client();
let user_id = req.match_info().get("user_id").unwrap();
let balance = client.load_grade(user_id); // Returns a balance as a String
HttpResponse::Ok()
.content_type("text/plain")
.body(format!("Hello! Your balance is {}", balance))
}
fn show_age(req: &HttpRequest) -> HttpResponse {
let client = create_client();
let user_id = req.match_info().get("user_id").unwrap();
let age = client.load_grade(user_id); // Returns an age as a String
HttpResponse::Ok()
.content_type("text/plain")
.body(format!("Hello! Your balance is {}", age))
}
fn main() {
env::set_var("RUST_LOG", "actix_web=debug");
env::set_var("RUST_BACKTRACE", "1");
env_logger::init();
let sys = actix::System::new("basic-example");
let addr = server::new(
|| App::new()
// enable logger
.middleware(middleware::Logger::default())
.resource("/balance/{user_id}", |r| r.method(Method::GET).f(show_balance))
.resource("/age/{user_id}", |r| r.method(Method::GET).f(show_age))
.bind("127.0.0.1:8080").expect("Can not bind to 127.0.0.1:8080")
.start();
println!("Starting http server: 127.0.0.1:8080");
let _ = sys.run();
}
fn create_client() -> UserDataClient {
let enviroment = grpcio::EnvBuilder::new().build();
let channel = grpcio::ChannelBuilder::new(enviroment)
.connect(API_URL);
UserDataClient::new(channel)
}
This code works, but my concern is that I have to create a client (and open a channel) for every incoming request which is inefficient and readable, I think it's a good idea to make sort of a singleton instead (since I can reuse it). I looked through the example folder and found that todo example is kinda similar to what I'm doing. So I found the following two options to inject my client object (after I create a single instance of it in main():
Put it in an app state
Inject it as a middleware (1, 2)?
What's the best/correct one to implement?
I thought about just passing a client object to every handler as an argument but I didn't manage to make it work (and doesn't look good anyway).

Asynchronously reconnecting a client to a server in an infinite loop

I'm not able to create a client that tries to connect to a server and:
if the server is down it has to try again in an infinite loop
if the server is up and connection is successful, when the connection is lost (i.e. server disconnects the client) the client has to restart the infinite loop to try to connect to the server
Here's the code to connect to a server; currently when the connection is lost the program exits. I'm not sure what the best way to implement it is; maybe I have to create a Future with an infinite loop?
extern crate tokio_line;
use tokio_line::LineCodec;
fn get_connection(handle: &Handle) -> Box<Future<Item = (), Error = io::Error>> {
let remote_addr = "127.0.0.1:9876".parse().unwrap();
let tcp = TcpStream::connect(&remote_addr, handle);
let client = tcp.and_then(|stream| {
let (sink, from_server) = stream.framed(LineCodec).split();
let reader = from_server.for_each(|message| {
println!("{}", message);
Ok(())
});
reader.map(|_| {
println!("CLIENT DISCONNECTED");
()
}).map_err(|err| err)
});
let client = client.map_err(|_| { panic!()});
Box::new(client)
}
fn main() {
let mut core = Core::new().unwrap();
let handle = core.handle();
let client = get_connection(&handle);
let client = client.and_then(|c| {
println!("Try to reconnect");
get_connection(&handle);
Ok(())
});
core.run(client).unwrap();
}
Add the tokio-line crate with:
tokio-line = { git = "https://github.com/tokio-rs/tokio-line" }
The key question seems to be: how do I implement an infinite loop using Tokio? By answering this question, we can tackle the problem of reconnecting infinitely upon disconnection. From my experience writing asynchronous code, recursion seems to be a straightforward solution to this problem.
UPDATE: as pointed out by Shepmaster (and the folks of the Tokio Gitter), my original answer leaks memory since we build a chain of futures that grows on each iteration. Here follows a new one:
Updated answer: use loop_fn
There is a function in the futures crate that does exactly what you need. It is called loop_fn. You can use it by changing your main function to the following:
fn main() {
let mut core = Core::new().unwrap();
let handle = core.handle();
let client = future::loop_fn((), |_| {
// Run the get_connection function and loop again regardless of its result
get_connection(&handle).map(|_| -> Loop<(), ()> {
Loop::Continue(())
})
});
core.run(client).unwrap();
}
The function resembles a for loop, which can continue or break depending on the result of get_connection (see the documentation for the Loop enum). In this case, we choose to always continue, so it will infinitely keep reconnecting.
Note that your version of get_connection will panic if there is an error (e.g. if the client cannot connect to the server). If you also want to retry after an error, you should remove the call to panic!.
Old answer: use recursion
Here follows my old answer, in case anyone finds it interesting.
WARNING: using the code below results in unbounded memory growth.
Making get_connection loop infinitely
We want to call the get_connection function each time the client is disconnected, so that is exactly what we are going to do (look at the comment after reader.and_then):
fn get_connection(handle: &Handle) -> Box<Future<Item = (), Error = io::Error>> {
let remote_addr = "127.0.0.1:9876".parse().unwrap();
let tcp = TcpStream::connect(&remote_addr, handle);
let handle_clone = handle.clone();
let client = tcp.and_then(|stream| {
let (sink, from_server) = stream.framed(LineCodec).split();
let reader = from_server.for_each(|message| {
println!("{}", message);
Ok(())
});
reader.and_then(move |_| {
println!("CLIENT DISCONNECTED");
// Attempt to reconnect in the future
get_connection(&handle_clone)
})
});
let client = client.map_err(|_| { panic!()});
Box::new(client)
}
Remember that get_connection is non-blocking. It just constructs a Box<Future>. This means that when calling it recursively, we still don't block. Instead, we get a new future, which we can link to the previous one by using and_then. As you can see, this is different to normal recursion since the stack doesn't grow on each iteration.
Note that we need to clone the handle (see handle_clone), and move it into the closure passed to reader.and_then. This is necessary because the closure is going to live longer than the function (it will be contained in the future we are returning).
Handling errors
The code you provided doesn't handle the case in which the client is unable to connect to the server (nor any other errors). Following the same principle shown above, we can handle errors by changing the end of get_connection to the following:
let handle_clone = handle.clone();
let client = client.or_else(move |err| {
// Note: this code will infinitely retry, but you could pattern match on the error
// to retry only on certain kinds of error
println!("Error connecting to server: {}", err);
get_connection(&handle_clone)
});
Box::new(client)
Note that or_else is like and_then, but it operates on the error produced by the future.
Removing unnecessary code from main
Finally, it is not necessary to use and_then in the main function. You can replace your main by the following code:
fn main() {
let mut core = Core::new().unwrap();
let handle = core.handle();
let client = get_connection(&handle);
core.run(client).unwrap();
}

Rust Type Inference Error

I'm writing a chat server over TCP as a learning project. I've been tinkering with the ws crate today, but I've come across an issue. This is the code I wrote, modifying their server example.
extern crate ws;
extern crate env_logger;
use ws::listen;
fn main() {
// Setup logging
env_logger::init().unwrap();
// Listen on an address and call the closure for each connection
if let Err(error) = listen("127.0.0.1:3012", |out| {
let mut message: String;
// The handler needs to take ownership of out, so we use move
move |message| {
message = message.trim();
// Handle messages received on this connection
println!("Server got message '{}'. ", message);
// Use the out channel to send messages back
out.send(message)
}
}) {
// Inform the user of failure
println!("Failed to create WebSocket due to {:?}", error);
}
}
When I try compiling it I get an error:
error: the type of this value must be known in this context
--> src/main.rs:15:23
|
15 | message = message.trim();
| ^^^^^^^^^^^^^^
Why is this happening? How may I fix this?
move |message| shadows the message variable you've declared outside the closure. So within the closure.. message is said to be a ws::Message ... except you've done this:
message = message.trim();
The compiler goes "oh no! trim()? That doesn't exist for ws::Message".. and so now it doesn't quite know what to do.
Option 1
The first fix involves delegating the trim() call to the client who sends the message.
The fix is to not make any assumptions about what the message is inside this closure. If you keep this:
move |message|
..but remove the trim() call, the compiler happily infers its type as ws::Message and will build:
if let Err(error) = listen("127.0.0.1:3012", |out| {
// The handler needs to take ownership of out, so we use move
move |message| {
// --- REMOVED trim() call ---
// Handle messages received on this connection
println!("Server got message '{}'. ", message);
// Use the out channel to send messages back
out.send(message)
}
}
This gives you the option of delegating the trim() call to the client instead.
Option 2
Option 2 involves inspecting the type of message you've received, and making sure you trim it only if it is text:
// The handler needs to take ownership of out, so we use move
move |mut message: ws::Message| {
// Only do it if the Message is text
if message.is_text() {
message = ws::Message::Text(message.as_text().unwrap().trim().into());
}
// Handle messages received on this connection
println!("Server got message '{}'. ", message);
// Use the out channel to send messages back
out.send(message)
}
This is perhaps a little more verbose than it needs to be.. but hopefully it shows you what the actual issue is with your original snippet of code.

Resources