How to check the connection type (plain/websocket)?

[mhanusek]

Hello, rustaceans!
I want to create a server with multipurpose sockets (web and classic sockets).
How do this? How do you check the type of connection in runtime?

// classic socket handler
async fn handle_connection(peer_addr: SocketAddr, 
                            mut tls_stream: TlsStream<TcpStream>, 
                            clients_map: Arc<Mutex<HashMap<SocketAddr, TlsStream<TcpStream>>>>) 
{
}

// websocket handler
async fn handle_ws_connection(peer_addr: SocketAddr, 
                              tls_stream: TlsStream<TcpStream>,
                              clients_map: Arc<Mutex<HashMap<SocketAddr, WebSocketStream<TlsStream<TcpStream>>>>>)
{
}

    let acceptor = TlsAcceptor::from(Arc::new(config));
    let listener = TcpListener::bind("0.0.0.0:9900").await?;

    let clients_map  = Arc::new(Mutex::new(HashMap::<SocketAddr, _>::new()));
    
    loop {

        let (stream, peer_addr) = listener.accept().await?;
        let acceptor = acceptor.clone();
        let tls_stream = acceptor.accept(stream).await?;
        
        let clients_map_cloned = clients_map.clone();

        tokio::spawn(async move {
            // let websocket = ... 
            // TODO: check connection
            if websocket {
                 handle_ws_connection(peer_addr, tokio_rustls::TlsStream::Server(tls_stream), clients_map_cloned).await;
            }
            else {
               handle_connection(peer_addr, tokio_rustls::TlsStream::Server(tls_stream), clients_map_cloned).await;
            }
        });
    }

Full example:

1. Server
2. Client
3. Websocket Client

Used libraries:

• tokio,
• tokio-rustls,
• tokio-tungstenite,
• async-tungstenite,
• rustls-pemfile

[semicoleon]

Why do you want to accept both WebSocket and plain TCP connections from the same port? Generally just having every client use WebSockets should be fine.

You'd have to buffer and then pre-parse the incoming message to see if it matches a WebSocket connection or your plain TCP protocol. It's certainly possible to do but I'm not sure that it would actually be useful. You'll have to maintain separate code paths for every message you receive and reply you send since your plain TCP protocol and the WebSocket protocol would be different^[1]

---

1. and if they weren't different then you would have just reinvented WebSockets ↩︎

[mhanusek]

I would like to create a server for communication with Desktop Apps (C++, Qt, Windows, Linux) and Web Apps (Yew.rs).

[semicoleon]

You can use WebSockets in the desktop apps.

[mhanusek]

the desktop app is already. I don't want to modify the desktop app. I create a new server in Rust for it.
I would like to do a possibility for the creation of a future web app.

[semicoleon]

Ahh, then you should probably just use different ports for the TCP connection and the WebSocket connection.

[yuriy0]

As I understand, a websocket connection starts with a handshake - an HTTP request/response pair in which the client and server agree to communicate via the websocket protocol. If the client fails to produce a handshake of the appropriate form, you can go ahead with trying to communicate via a "plain" TCP connection (whatever that entails for your application, I'm not sure).

I'm not really familiar with all the libs involved, but this should be fairly straightforward with tungstenite which will accept an abstract stream type.

The only issue is that to fallback to the plain stream correctly, you would need to rewind the stream to the beginning (since the handshake protocol would try to consume some data from the start in order to read the HTTP request). If your stream is seekable you can try something like (untested)

enum MyStream<S>
{
  Websocket(WebSocketStream<S>),
  Plain(S)
}

async fn get_connection_type(tls_stream: S) -> MyStream<S>
{
    accept_async(tls_stream).await
      .map(|s| MyStream::Websocket(s))
      .unwrap_or_else(|_| {
         s.rewind(); // pretend we never read from the stream
         MyStream::Plain(s)
      })
}

If not you would probably have to wrap the stream type with your own stream which saves the first X MB of the stream and provides a method like rewind.

[mhanusek]

@yuriy0 @semicoleon thanks for your replies.
I have another problem Broadcast server - what is the best way?

[mhanusek]

I have a problem with binding on two ports.
The first binding blocks the second.
How to solve it?

  let plain_listener = TcpListener::bind("0.0.0.0:9900").await?;
  let websocket_listener = TcpListener::bind("0.0.0.0:9999").await?;

Could I use tokio::select! ?
or it?

[ndusart]

I think you may keep it as is.

bind syscall is always blocking anyway. TcpListener::bind is async because it may have to resolve DNS hostname. But this is not your case, so just make the two binds one after the other is not less performant.

But, if you want to call these concurrently, you could use either tokio::select! or spawn two tasks. Depending on what you will do, one or the other may be more adapted.

The simplest is still to do it sequentially and move them to different tasks if the need for separate tasks occurs.

[mhanusek]

I would like to create - Framed<WebSocketStream<TlsStream>, BytesCodec>.

How to use WebSocketStream<TlsStream<TcpStream>> and Framed?

#[derive(Debug)]
enum MyStream
{
  WebsocketTls(WebSocketStream<TlsStream<TcpStream>>),
  Tls(tokio_rustls::server::TlsStream<tokio::net::TcpStream>),
}

async fn process_stream(
    state: Arc<Mutex<Shared>>,
    stream: MyStream,
    addr: SocketAddr,
) -> Result<(), Box<dyn Error>> 
{   
    match stream
    {
        MyStream::Tls(s) => { 
            let framed = Framed::new(s, BytesCodec::new());
        },
        MyStream::WebsocketTls(s) => {
          //TODO: how to do it?  
          //Framed::new(s, BytesCodec::new())
        }
    };

    Ok(())
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> 
{
    tracing_subscriber::fmt().init();

    let certs = load_certs(&Path::new("../keys/eddsa-ca.cert"))?;
    let keys = load_keys(&Path::new("../keys/Ed25519_private_key.pem"))?;

    let private_key = keys.get(0).unwrap().clone();
    let _cert = certs.get(0).unwrap().clone();

    let config = rustls::ServerConfig::builder()
        .with_safe_defaults()
        .with_no_client_auth()
        .with_single_cert(certs, private_key)
        .map_err(|err| IoError::new(ErrorKind::InvalidInput, err))?;

    let cfg = Arc::new(config);

    let plain_listener = TcpListener::bind("0.0.0.0:9900").await?;
    let websocket_listener = TcpListener::bind("0.0.0.0:9999").await?;

    let state = Arc::new(Mutex::new(Shared::new()));

    loop {

        let (stream, addr) = tokio::select! {
            res1 = plain_listener.accept() =>  {
                let (stream, addr) = res1.unwrap();
                let acceptor = TlsAcceptor::from(cfg.clone());
                let tls_stream = acceptor.accept(stream).await.unwrap();
                (MyStream::Tls(tls_stream), addr)
            }
            res2 = websocket_listener.accept() =>  { 
                let (stream, addr) = res2.unwrap();
                let acceptor = TlsAcceptor::from(cfg.clone());
                let tls_stream = acceptor.accept(stream).await.unwrap();
                let ws_stream = tokio_tungstenite::accept_async(tls_stream).await.unwrap();        
                (MyStream::WebsocketTls(ws_stream), addr)
            },
        };

        let state = Arc::clone(&state);

        tokio::spawn(async move {
            tracing::debug!("accepted connection");
            if let Err(e) = process_stream(state, stream, addr).await {
                tracing::info!("an error occurred; error = {:?}", e);
            }
        });
    }
}

[ndusart]

Note that WebSocketStream is already Stream + Sink, you wouldn't need a codec for it.

The difference is that it is a stream over websocket messages and not bytes. You have to choose how to convert these into bytes (you could only extra bytes from data messages and ignore other messages for example, but it really depends on what you want). You can use map or filter_map on WebSocketStream for this.

This will be a different type than Framed. You'll have to store it in a Box<dyn Stream + Sink> or use a custom enum type as you did with MyStream.

[mhanusek]

Thanks for your reply.

My sandboxes:

1. tokio-tls-example/main.rs at websockets · hanusek/tokio-tls-example · GitHub
2. tokio-tls-example/main.rs at websocket2 · hanusek/tokio-tls-example · GitHub