easier wsupgrade api for protocols/exts, ability for client to verify protocols, etc.

Author: illegalprime

examples/hyper.rs

@@ -38,8 +38,7 @@ fn main() {
 				return;
 			}
 
-			// TODO: same check like in server.rs
-			let mut client = connection.accept().unwrap();
+			let mut client = connection.use_protocol("rust-websocket").accept().unwrap();
 
 			let ip = client.peer_addr().unwrap();
 

examples/server.rs

@@ -4,14 +4,6 @@ use std::thread;
 use websocket::{Server, Message};
 use websocket::message::Type;
 
-// TODO: I think the .reject() call is only for malformed packets
-// there should be an easy way to accept the socket with the given protocols
-// this would mean there should be a way to accept or reject on the client
-// Do you send the protocol you want to talk when you are not given it as an
-// option? What is a rejection response? Does the client check for it?
-// Client should expose what the decided protocols/extensions/etc are.
-// can you accept only one protocol??
-
 fn main() {
 	let server = Server::bind("127.0.0.1:2794").unwrap();
 
@@ -23,7 +15,7 @@ fn main() {
 				return;
 			}
 
-			let mut client = request.accept().unwrap();
+			let mut client = request.use_protocol("rust-websocket").accept().unwrap();
 
 			let ip = client.peer_addr().unwrap();
 

src/client/builder.rs

@@ -20,25 +20,6 @@ use result::{WSUrlErrorKind, WebSocketResult, WebSocketError};
 use stream::Stream;
 use super::Client;
 
-macro_rules! upsert_header {
-    ($headers:expr; $header:ty; {
-        Some($pat:pat) => $some_match:expr,
-        None => $default:expr
-    }) => {{
-        match $headers.has::<$header>() {
-            true => {
-                match $headers.get_mut::<$header>() {
-                    Some($pat) => { $some_match; },
-                    None => (),
-                };
-            }
-            false => {
-                $headers.set($default);
-            },
-        };
-    }}
-}
-
 /// Build clients with a builder-style API
 #[derive(Clone, Debug)]
 pub struct ClientBuilder<'u> {
@@ -305,6 +286,6 @@ impl<'u> ClientBuilder<'u> {
 			return Err(WebSocketError::ResponseError("Connection field must be 'Upgrade'"));
 		}
 
-		Ok(Client::unchecked(stream))
+		Ok(Client::unchecked(stream, response.headers))
 	}
 }

src/client/mod.rs

@@ -4,6 +4,7 @@ extern crate url;
 use std::net::TcpStream;
 use std::net::SocketAddr;
 use std::io::Result as IoResult;
+use hyper::header::Headers;
 
 use ws;
 use ws::sender::Sender as SenderTrait;
@@ -12,6 +13,8 @@ use ws::receiver::Receiver as ReceiverTrait;
 use result::WebSocketResult;
 use stream::{AsTcpStream, Stream, Splittable, Shutdown};
 use dataframe::DataFrame;
+use header::{WebSocketProtocol, WebSocketExtensions, Origin};
+use header::extensions::Extension;
 
 use ws::dataframe::DataFrame as DataFrameable;
 use sender::Sender;
@@ -54,6 +57,7 @@ pub struct Client<S>
 	where S: Stream
 {
 	pub stream: S,
+	headers: Headers,
 	sender: Sender,
 	receiver: Receiver,
 }
@@ -109,8 +113,9 @@ impl<S> Client<S>
 	/// **without sending any handshake** this is meant to only be used with
 	/// a stream that has a websocket connection already set up.
 	/// If in doubt, don't use this!
-	pub fn unchecked(stream: S) -> Self {
+	pub fn unchecked(stream: S, headers: Headers) -> Self {
 		Client {
+			headers: headers,
 			stream: stream,
 			// NOTE: these are always true & false, see
 			// https://tools.ietf.org/html/rfc6455#section-5
@@ -152,6 +157,28 @@ impl<S> Client<S>
 		self.receiver.recv_message(self.stream.reader())
 	}
 
+	pub fn headers(&self) -> &Headers {
+		&self.headers
+	}
+
+	pub fn protocols(&self) -> &[String] {
+		self.headers
+		    .get::<WebSocketProtocol>()
+		    .map(|p| p.0.as_slice())
+		    .unwrap_or(&[])
+	}
+
+	pub fn extensions(&self) -> &[Extension] {
+		self.headers
+		    .get::<WebSocketExtensions>()
+		    .map(|e| e.0.as_slice())
+		    .unwrap_or(&[])
+	}
+
+	pub fn origin(&self) -> Option<&str> {
+		self.headers.get::<Origin>().map(|o| &o.0 as &str)
+	}
+
 	pub fn stream_ref(&self) -> &S {
 		&self.stream
 	}
@@ -160,6 +187,10 @@ impl<S> Client<S>
 		&mut self.stream
 	}
 
+	pub fn into_stream(self) -> S {
+		self.stream
+	}
+
 	/// Returns an iterator over incoming messages.
 	///
 	///```no_run

src/lib.rs

@@ -60,6 +60,26 @@ pub use self::stream::Stream;
 pub use self::ws::Sender;
 pub use self::ws::Receiver;
 
+macro_rules! upsert_header {
+    ($headers:expr; $header:ty; {
+        Some($pat:pat) => $some_match:expr,
+        None => $default:expr
+    }) => {{
+        match $headers.has::<$header>() {
+            true => {
+                match $headers.get_mut::<$header>() {
+                    Some($pat) => { $some_match; },
+                    None => (),
+                };
+            }
+            false => {
+                $headers.set($default);
+            },
+        };
+    }}
+}
+
+
 pub mod ws;
 pub mod client;
 pub mod server;

src/server/upgrade/hyper.rs

@@ -30,6 +30,7 @@ impl<'a, 'b> IntoWs for HyperRequest<'a, 'b> {
 		let stream = reader.into_inner().get_mut();
 
 		Ok(WsUpgrade {
+		       headers: Headers::new(),
 		       stream: stream,
 		       request: Incoming {
 		           version: version,

src/server/upgrade/mod.rs

@@ -37,41 +37,74 @@ pub mod hyper;
 pub struct WsUpgrade<S>
 	where S: Stream
 {
-	stream: S,
-	request: Request,
+	pub headers: Headers,
+	pub stream: S,
+	pub request: Request,
 }
 
 impl<S> WsUpgrade<S>
     where S: Stream
 {
-	pub fn accept(self) -> IoResult<Client<S>> {
+	pub fn use_protocol<P>(mut self, protocol: P) -> Self
+		where P: Into<String>
+	{
+		upsert_header!(self.headers; WebSocketProtocol; {
+            Some(protos) => protos.0.push(protocol.into()),
+            None => WebSocketProtocol(vec![protocol.into()])
+        });
+		self
+	}
+
+	pub fn use_extension(mut self, extension: Extension) -> Self {
+		upsert_header!(self.headers; WebSocketExtensions; {
+            Some(protos) => protos.0.push(extension),
+            None => WebSocketExtensions(vec![extension])
+        });
+		self
+	}
+
+	pub fn use_extensions<I>(mut self, extensions: I) -> Self
+		where I: IntoIterator<Item = Extension>
+	{
+		let mut extensions: Vec<Extension> =
+			extensions.into_iter().collect();
+		upsert_header!(self.headers; WebSocketExtensions; {
+            Some(protos) => protos.0.append(&mut extensions),
+            None => WebSocketExtensions(extensions)
+        });
+		self
+	}
+
+	pub fn accept(self) -> Result<Client<S>, (S, IoError)> {
 		self.accept_with(&Headers::new())
 	}
 
-	pub fn accept_with(mut self, custom_headers: &Headers) -> IoResult<Client<S>> {
-		let mut headers = Headers::new();
-		headers.extend(custom_headers.iter());
-		headers.set(WebSocketAccept::new(
-            // NOTE: we know there is a key because this is a valid request
-            // i.e. to construct this you must go through the validate function
-            self.request.headers.get::<WebSocketKey>().unwrap()
-        ));
-		headers.set(Connection(vec![
+	pub fn accept_with(mut self, custom_headers: &Headers) -> Result<Client<S>, (S, IoError)> {
+		self.headers.extend(custom_headers.iter());
+		self.headers
+		    .set(WebSocketAccept::new(// NOTE: we know there is a key because this is a valid request
+		                              // i.e. to construct this you must go through the validate function
+		                              self.request.headers.get::<WebSocketKey>().unwrap()));
+		self.headers
+		    .set(Connection(vec![
 			      ConnectionOption::ConnectionHeader(UniCase("Upgrade".to_string()))
 		    ]));
-		headers.set(Upgrade(vec![Protocol::new(ProtocolName::WebSocket, None)]));
+		self.headers.set(Upgrade(vec![Protocol::new(ProtocolName::WebSocket, None)]));
 
-		try!(self.send(StatusCode::SwitchingProtocols, &headers));
+		if let Err(e) = self.send(StatusCode::SwitchingProtocols) {
+			return Err((self.stream, e));
+		}
 
-		Ok(Client::unchecked(self.stream))
+		Ok(Client::unchecked(self.stream, self.headers))
 	}
 
 	pub fn reject(self) -> Result<S, (S, IoError)> {
 		self.reject_with(&Headers::new())
 	}
 
 	pub fn reject_with(mut self, headers: &Headers) -> Result<S, (S, IoError)> {
-		match self.send(StatusCode::BadRequest, headers) {
+		self.headers.extend(headers.iter());
+		match self.send(StatusCode::BadRequest) {
 			Ok(()) => Ok(self.stream),
 			Err(e) => Err((self.stream, e)),
 		}
@@ -113,12 +146,12 @@ impl<S> WsUpgrade<S>
 		self.stream
 	}
 
-	fn send(&mut self, status: StatusCode, headers: &Headers) -> IoResult<()> {
+	fn send(&mut self, status: StatusCode) -> IoResult<()> {
 		try!(write!(self.stream.writer(),
 		            "{} {}\r\n",
 		            self.request.version,
 		            status));
-		try!(write!(self.stream.writer(), "{}\r\n", headers));
+		try!(write!(self.stream.writer(), "{}\r\n", self.headers));
 		Ok(())
 	}
 }
@@ -173,6 +206,7 @@ impl<S> IntoWs for S
 		match validate(&request.subject.0, &request.version, &request.headers) {
 			Ok(_) => {
 				Ok(WsUpgrade {
+				       headers: Headers::new(),
 				       stream: self,
 				       request: request,
 				   })
@@ -192,6 +226,7 @@ impl<S> IntoWs for RequestStreamPair<S>
 		match validate(&self.1.subject.0, &self.1.version, &self.1.headers) {
 			Ok(_) => {
 				Ok(WsUpgrade {
+				       headers: Headers::new(),
 				       stream: self.0,
 				       request: self.1,
 				   })