main.rs

use std::{env, io::Error as IoError, net::SocketAddr};

use futures_util::StreamExt;

use tokio::net::{TcpListener, TcpStream};
use tokio_tungstenite::accept_async;

use tungstenite::Error;

async fn handle_connection(raw_stream: TcpStream, addr: SocketAddr) -> tungstenite::Result<()> {
  println!("{}", addr);

  let ws_stream = accept_async(raw_stream)
    .await
    .expect("handshake error");

  println!("connected: {}", addr);

  let (write, read) = ws_stream.split();
  read.forward(write).await.unwrap_or_else(|e: Error| {
    eprintln!("{:?}", e);
  });

  Ok(())
}

async fn accept_connection(stream: TcpStream, peer: SocketAddr) {
  if let Err(e) = handle_connection(stream, peer).await {
    match e {
      Error::ConnectionClosed | Error::Protocol(_) | Error::Utf8 => (),
      err => eprintln!("error: {}", err),
    }
  }
}

#[tokio::main]
async fn main() -> Result<(), IoError> {
  let host = env::var("WS_HOST").unwrap_or(String::from("0.0.0.0"));
  let port = env::var("WS_PORT")
    .unwrap_or(String::from("8080"))
    .parse::<u16>()
    .unwrap_or(8080);

  let address = format!("{}:{}", host, port);
  let mut listener = TcpListener::bind(&address).await.expect("cant bind");

  println!("listening: {}", address);

  while let Ok((stream, addr)) = listener.accept().await {
    tokio::spawn(accept_connection(stream, addr));
  }

  Ok(())
}