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

//! The default implementation of a WebSocket Receiver.

use std::io::Read;
use std::io::Result as IoResult;
use hyper::buffer::BufReader;

use dataframe::{DataFrame, Opcode};
use result::{WebSocketResult, WebSocketError};
use stream::WebSocketStream;
use stream::Shutdown;
use ws;

/// A Receiver that wraps a Reader and provides a default implementation using
/// DataFrames and Messages.
pub struct Receiver<R> {
	inner: BufReader<R>,
	buffer: Vec<DataFrame>,
	mask: bool,
}

impl<R> Receiver<R>
where R: Read {
	/// Create a new Receiver using the specified Reader.
	pub fn new(reader: R, mask: bool) -> Receiver<R> {
		Receiver {
			inner: BufReader::new(reader),
			buffer: Vec::new(),
			mask: mask,
		}
	}
	/// Returns a reference to the underlying Reader.
	pub fn get_ref(&self) -> &BufReader<R> {
		&self.inner
	}
	/// Returns a mutable reference to the underlying Reader.
	pub fn get_mut(&mut self) -> &mut BufReader<R> {
		&mut self.inner
	}
}

impl Receiver<WebSocketStream> {
    /// Closes the receiver side of the connection, will cause all pending and future IO to
    /// return immediately with an appropriate value.
    pub fn shutdown(&mut self) -> IoResult<()> {
        self.inner.get_mut().shutdown(Shutdown::Read)
    }

    /// Shuts down both Sender and Receiver, will cause all pending and future IO to
    /// return immediately with an appropriate value.
    pub fn shutdown_all(&mut self) -> IoResult<()> {
        self.inner.get_mut().shutdown(Shutdown::Both)
    }
}

impl<R: Read> ws::Receiver<DataFrame> for Receiver<R> {
	/// Reads a single data frame from the remote endpoint.
	fn recv_dataframe(&mut self) -> WebSocketResult<DataFrame> {
		DataFrame::read_dataframe(&mut self.inner, self.mask)
	}
	/// Returns the data frames that constitute one message.
	fn recv_message_dataframes(&mut self) -> WebSocketResult<Vec<DataFrame>> {
		let mut finished = if self.buffer.is_empty() {
			let first = try!(self.recv_dataframe());

			if first.opcode == Opcode::Continuation {
				return Err(WebSocketError::ProtocolError(
					"Unexpected continuation data frame opcode"
				));
			}

			let finished = first.finished;
			self.buffer.push(first);
			finished
		}
		else {
			false
		};

		while !finished {
			let next = try!(self.recv_dataframe());
			finished = next.finished;

			match next.opcode as u8 {
				// Continuation opcode
				0 => self.buffer.push(next),
				// Control frame
				8...15 => {
					return Ok(vec![next]);
				}
				// Others
				_ => return Err(WebSocketError::ProtocolError(
					"Unexpected data frame opcode"
				)),
			}
		}

		let buffer = self.buffer.clone();
		self.buffer.clear();

		Ok(buffer)
	}
}