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

//! # warp-websockify
//! websockify implementation for warp
//!
//! ```
//! use warp::Filter;
//! use warp_websockify::{Destination, websockify};
//!
//! let dest = Destination::tcp("localhost:5901").unwrap();
//! let serve = websockify(dest);
//! ```

mod error;

pub use error::{WebsockifyError, WebsockifyErrorKind};
use futures::prelude::*;
use log::{debug, error, info};
use std::io;
use std::net::{SocketAddr, ToSocketAddrs};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::{TcpStream, UnixStream};
use warp::ws::{Message, WebSocket, Ws};
use warp::{reject::Rejection, reply::Reply, Filter};

/// WebSockify upstream
pub enum Destination {
    /// Connect to TCP
    Tcp(Vec<SocketAddr>),

    /// Connect to unix domain socket
    Unix(PathBuf),
}

impl Destination {
    /// Create destination to unix domain socket
    pub fn unix<P: AsRef<Path>>(path: P) -> Destination {
        Destination::Unix(path.as_ref().to_path_buf())
    }

    /// Create destination to TCP
    pub fn tcp(addr: impl ToSocketAddrs) -> io::Result<Destination> {
        Ok(Destination::Tcp(addr.to_socket_addrs()?.collect()))
    }

    async fn connect(&self) -> io::Result<NetStream> {
        match self {
            Destination::Tcp(addrs) => {
                let mut last_error = None;
                for one in addrs {
                    match TcpStream::connect(one).await {
                        Ok(stream) => return Ok(NetStream::Tcp(stream)),
                        Err(e) => last_error = Some(e),
                    }
                }
                Err(last_error.unwrap())
            }
            Destination::Unix(path) => Ok(NetStream::Unix(UnixStream::connect(path).await?)),
        }
    }
}

enum NetStream {
    Unix(UnixStream),
    Tcp(TcpStream),
}

/// Creates a `Filter` that connet to TCP or unix domain socket.
pub fn websockify(
    dest: Destination,
) -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone {
    let dest: Arc<Destination> = Arc::new(dest);
    let dest = warp::any().map(move || dest.clone());
    warp::ws().and(dest).and_then(websockify_connect)
}

/// Connect to TCP or unix domain socket and redirect to websocket
pub fn websockify_connect(
    ws: Ws,
    dest: impl AsRef<Destination>,
) -> impl TryFuture<Ok = impl Reply, Error = Rejection> {
    async move {
        match dest.as_ref().connect().await {
            Ok(stream) => Ok(ws.on_upgrade(|ws| connect(ws, stream))),
            Err(_) => Err(warp::reject::reject()),
        }
    }
}

async fn connect(ws: WebSocket, stream: NetStream) {
    match stream {
        NetStream::Unix(x) => unix_connect(ws, x).await,
        NetStream::Tcp(x) => tcp_connect(ws, x).await,
    }
}

async fn tcp_connect(ws: WebSocket, mut stream: TcpStream) {
    let (tcp_rx, tcp_tx) = stream.split();
    let (tx, rx) = ws.split();

    let x = handle_ws_rx(rx, tcp_tx);
    let y = handle_ws_tx(tx, tcp_rx);

    if let Err(e) = tokio::try_join!(x, y) {
        error!("tcp connection error: {}", e);
    }
}

async fn unix_connect(ws: WebSocket, mut stream: UnixStream) {
    let (tcp_rx, tcp_tx) = stream.split();
    let (tx, rx) = ws.split();

    let x = handle_ws_rx(rx, tcp_tx);
    let y = handle_ws_tx(tx, tcp_rx);
    if let Err(e) = tokio::try_join!(x, y) {
        info!("unix connection error: {}", e);
    }
}

async fn handle_ws_rx<
    WSR: Stream<Item = Result<Message, warp::Error>> + std::marker::Unpin + std::marker::Send,
    SW: AsyncWrite + std::marker::Unpin + std::marker::Send,
>(
    mut rx: WSR,
    mut tcp_tx: SW,
) -> Result<(), WebsockifyError> {
    while let Some(message) = rx.next().await {
        let message = message?;
        if message.is_binary() {
            tcp_tx.write_all(message.as_bytes()).await?;
        }
    }
    info!("finish 1");
    Ok(())
}

async fn handle_ws_tx<
    WST: Sink<Message, Error = warp::Error> + std::marker::Unpin + std::marker::Send,
    SR: AsyncRead + std::marker::Unpin + std::marker::Send,
>(
    mut tx: WST,
    mut tcp_rx: SR,
) -> Result<(), WebsockifyError> {
    let mut buffer: Vec<u8> = Vec::with_capacity(10000);
    for _ in 0..10000 {
        buffer.push(0);
    }

    loop {
        let read_bytes = tcp_rx.read(&mut buffer).await?;
        if read_bytes > 0 {
            //println!("send bytes: {}", read_bytes);
            tx.send(Message::binary(&buffer[0..read_bytes])).await?;
        } else {
            debug!("read zero bytes");
            tx.close().await?;
            break;
        }
    }

    info!("finish 2");
    Ok(())
}