use std::fmt;
use std::io::{Read, Write};
use std::net::{SocketAddr, TcpStream as StdTcpStream};
use std::result;
use std::str;
use base64;
use futures::StreamExt;
use rand;
use tokio::{
io::{AsyncRead, AsyncWrite, AsyncWriteExt},
net::TcpStream as TokioTcpStream,
};
use tokio_util::codec::{Decoder, Encoder, Framed};
use url::{self, Url};
use websocket_codec::UpgradeCodec;
use crate::ssl;
use crate::sync;
use crate::{AsyncClient, AsyncNetworkStream, Client, MessageCodec, NetworkStream, Result};
fn replace_codec<T, C1, C2>(framed: Framed<T, C1>, codec: C2) -> Framed<T, C2>
where
T: AsyncRead + AsyncWrite,
C1: Encoder + Decoder,
C2: Encoder + Decoder,
{
let parts1 = framed.into_parts();
let mut parts2 = Framed::new(parts1.io, codec).into_parts();
parts2.read_buf = parts1.read_buf;
parts2.write_buf = parts1.write_buf;
Framed::from_parts(parts2)
}
macro_rules! writeok {
($dst:expr, $($arg:tt)*) => {
let _ = fmt::Write::write_fmt(&mut $dst, format_args!($($arg)*));
}
}
fn resolve(url: &Url) -> Result<SocketAddr> {
url.socket_addrs(|| None)?
.into_iter()
.next()
.ok_or_else(|| "can't resolve host".to_owned().into())
}
fn make_key(key: Option<[u8; 16]>, key_base64: &mut [u8; 24]) -> &str {
let key_bytes = key.unwrap_or_else(rand::random);
assert_eq!(
24,
base64::encode_config_slice(&key_bytes, base64::STANDARD, key_base64)
);
str::from_utf8(key_base64).unwrap()
}
fn build_request(url: &Url, key: &str, headers: &Vec<(String, String)>) -> String {
let mut s = String::new();
writeok!(s, "GET {path}", path = url.path());
if let Some(query) = url.query() {
writeok!(s, "?{query}", query = query);
}
s += " HTTP/1.1\r\n";
if let Some(host) = url.host() {
writeok!(s, "Host: {host}", host = host);
if let Some(port) = url.port_or_known_default() {
writeok!(s, ":{port}", port = port);
}
s += "\r\n";
}
writeok!(
s,
"Upgrade: websocket\r\n\
Connection: Upgrade\r\n\
Sec-WebSocket-Key: {key}\r\n\
Sec-WebSocket-Version: 13\r\n",
key = key
);
for (name, value) in headers.iter() {
writeok!(s, "{name}: {value}\r\n", name = name, value = value);
}
writeok!(s, "\r\n");
s
}
pub struct ClientBuilder {
url: Url,
key: Option<[u8; 16]>,
headers: Vec<(String, String)>,
}
impl ClientBuilder {
pub fn new(url: &str) -> result::Result<Self, url::ParseError> {
Ok(Self::from_url(Url::parse(url)?))
}
pub fn from_url(url: Url) -> Self {
ClientBuilder {
url,
key: None,
headers: Vec::new(),
}
}
pub fn add_header(&mut self, name: String, value: String) {
self.headers.push((name, value));
}
pub async fn async_connect_insecure(self) -> Result<AsyncClient<TokioTcpStream>> {
let addr = resolve(&self.url)?;
let stream = TokioTcpStream::connect(&addr).await?;
Ok(self.async_connect_on(stream).await?)
}
pub fn connect_insecure(self) -> Result<Client<StdTcpStream>> {
let addr = resolve(&self.url)?;
let stream = StdTcpStream::connect(&addr)?;
self.connect_on(stream)
}
pub async fn async_connect(
self,
) -> Result<AsyncClient<Box<dyn AsyncNetworkStream + Sync + Send + Unpin + 'static>>> {
let addr = resolve(&self.url)?;
let stream = TokioTcpStream::connect(&addr).await?;
let stream: Box<dyn AsyncNetworkStream + Sync + Send + Unpin + 'static> = if self.url.scheme() == "wss" {
let domain = self.url.domain().unwrap_or("").to_owned();
let stream = ssl::async_wrap(domain, stream).await?;
Box::new(stream)
} else {
Box::new(stream)
};
self.async_connect_on(stream).await
}
pub fn connect(self) -> Result<Client<Box<dyn NetworkStream + Sync + Send + 'static>>> {
let addr = resolve(&self.url)?;
let stream = StdTcpStream::connect(&addr)?;
let stream: Box<dyn NetworkStream + Sync + Send + 'static> = if self.url.scheme() == "wss" {
let domain = self.url.domain().unwrap_or("");
let stream = ssl::wrap(domain, stream)?;
Box::new(stream)
} else {
Box::new(stream)
};
self.connect_on(stream)
}
pub async fn async_connect_on<S: AsyncRead + AsyncWrite + Unpin>(self, mut stream: S) -> Result<AsyncClient<S>> {
let mut key_base64 = [0; 24];
let key = make_key(self.key, &mut key_base64);
let upgrade_codec = UpgradeCodec::new(key);
let request = build_request(&self.url, key, &self.headers);
AsyncWriteExt::write_all(&mut stream, request.as_bytes()).await?;
let (opt, framed) = upgrade_codec.framed(stream).into_future().await;
opt.ok_or_else(|| "no HTTP Upgrade response".to_owned())??;
Ok(replace_codec(framed, MessageCodec::new()))
}
pub fn connect_on<S: Read + Write>(self, mut stream: S) -> Result<Client<S>> {
let mut key_base64 = [0; 24];
let key = make_key(self.key, &mut key_base64);
let upgrade_codec = UpgradeCodec::new(key);
let request = build_request(&self.url, key, &self.headers);
Write::write_all(&mut stream, request.as_bytes())?;
let mut framed = sync::Framed::new(stream, upgrade_codec);
framed.receive()?.ok_or_else(|| "no HTTP Upgrade response".to_owned())?;
Ok(framed.replace_codec(MessageCodec::new()))
}
#[cfg(test)]
fn key(mut self, key: &[u8]) -> Self {
let mut a = [0; 16];
a.copy_from_slice(key);
self.key = Some(a);
self
}
}
#[cfg(test)]
mod tests {
use std::fmt;
use std::io::{self, Cursor, Read, Write};
use std::pin::Pin;
use std::result;
use std::str;
use std::task::{Context, Poll};
use base64;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::runtime::Runtime;
use crate::ClientBuilder;
type Result<T> = result::Result<T, crate::Error>;
pub struct ReadWritePair<R, W>(pub R, pub W);
impl<R: Read, W> Read for ReadWritePair<R, W> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
self.0.read(buf)
}
fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
self.0.read_to_end(buf)
}
fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
self.0.read_to_string(buf)
}
fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
self.0.read_exact(buf)
}
}
impl<R, W: Write> Write for ReadWritePair<R, W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.1.write(buf)
}
fn flush(&mut self) -> io::Result<()> {
self.1.flush()
}
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
self.1.write_all(buf)
}
fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> io::Result<()> {
self.1.write_fmt(fmt)
}
}
impl<R: AsyncRead + Unpin, W: Unpin> AsyncRead for ReadWritePair<R, W> {
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
Pin::new(&mut self.get_mut().0).poll_read(cx, buf)
}
}
impl<R: Unpin, W: AsyncWrite + Unpin> AsyncWrite for ReadWritePair<R, W> {
fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
Pin::new(&mut self.get_mut().1).poll_write(cx, buf)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.get_mut().1).poll_flush(cx)
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.get_mut().1).poll_shutdown(cx)
}
}
static REQUEST: &'static str = "GET /stream?query HTTP/1.1\r\n\
Host: localhost:8000\r\n\
Upgrade: websocket\r\n\
Connection: Upgrade\r\n\
Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
Sec-WebSocket-Version: 13\r\n\
\r\n";
static RESPONSE: &'static str = "HTTP/1.1 101 Switching Protocols\r\n\
Upgrade: websocket\r\n\
Connection: Upgrade\r\n\
sec-websocket-accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n\
\r\n";
#[test]
fn can_async_connect_on() -> Result<()> {
let mut input = Cursor::new(&RESPONSE[..]);
let mut output = Vec::new();
Runtime::new()?.block_on(async {
ClientBuilder::new("ws://localhost:8000/stream?query")?
.key(&base64::decode(b"dGhlIHNhbXBsZSBub25jZQ==")?)
.async_connect_on(ReadWritePair(&mut input, &mut output))
.await
})?;
assert_eq!(REQUEST, str::from_utf8(&output)?);
Ok(())
}
#[test]
fn can_connect_on() -> Result<()> {
let mut input = Cursor::new(&RESPONSE[..]);
let mut output = Vec::new();
ClientBuilder::new("ws://localhost:8000/stream?query")?
.key(&base64::decode(b"dGhlIHNhbXBsZSBub25jZQ==")?)
.connect_on(ReadWritePair(&mut input, &mut output))?;
assert_eq!(REQUEST, str::from_utf8(&output)?);
Ok(())
}
}