jsonrpsee_client_transport/ws/
stream.rs

1// Copyright 2019-2021 Parity Technologies (UK) Ltd.
2//
3// Permission is hereby granted, free of charge, to any
4// person obtaining a copy of this software and associated
5// documentation files (the "Software"), to deal in the
6// Software without restriction, including without
7// limitation the rights to use, copy, modify, merge,
8// publish, distribute, sublicense, and/or sell copies of
9// the Software, and to permit persons to whom the Software
10// is furnished to do so, subject to the following
11// conditions:
12//
13// The above copyright notice and this permission notice
14// shall be included in all copies or substantial portions
15// of the Software.
16//
17// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
18// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
19// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
20// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
21// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
22// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
23// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
24// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
25// DEALINGS IN THE SOFTWARE.
26
27//! Convenience wrapper for a stream (AsyncRead + AsyncWrite) which can either be plain TCP or TLS.
28
29use std::io::Error as IoError;
30use std::pin::Pin;
31use std::task::Context;
32use std::task::Poll;
33
34use pin_project::pin_project;
35use tokio::io::{AsyncRead, AsyncWrite};
36use tokio::net::TcpStream;
37
38/// Stream to represent either a unencrypted or encrypted socket stream.
39#[pin_project(project = EitherStreamProj)]
40#[derive(Debug)]
41#[allow(clippy::large_enum_variant)]
42pub enum EitherStream {
43	/// Unencrypted socket stream.
44	Plain(#[pin] TcpStream),
45	/// Encrypted socket stream.
46	#[cfg(feature = "tls")]
47	Tls(#[pin] tokio_rustls::client::TlsStream<TcpStream>),
48}
49
50impl AsyncRead for EitherStream {
51	fn poll_read(
52		self: Pin<&mut Self>,
53		cx: &mut Context,
54		buf: &mut tokio::io::ReadBuf<'_>,
55	) -> Poll<Result<(), IoError>> {
56		match self.project() {
57			EitherStreamProj::Plain(stream) => AsyncRead::poll_read(stream, cx, buf),
58			#[cfg(feature = "tls")]
59			EitherStreamProj::Tls(stream) => AsyncRead::poll_read(stream, cx, buf),
60		}
61	}
62}
63
64impl AsyncWrite for EitherStream {
65	fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<Result<usize, IoError>> {
66		match self.project() {
67			EitherStreamProj::Plain(stream) => AsyncWrite::poll_write(stream, cx, buf),
68			#[cfg(feature = "tls")]
69			EitherStreamProj::Tls(stream) => AsyncWrite::poll_write(stream, cx, buf),
70		}
71	}
72
73	fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), IoError>> {
74		match self.project() {
75			EitherStreamProj::Plain(stream) => AsyncWrite::poll_flush(stream, cx),
76			#[cfg(feature = "tls")]
77			EitherStreamProj::Tls(stream) => AsyncWrite::poll_flush(stream, cx),
78		}
79	}
80
81	fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), IoError>> {
82		match self.project() {
83			EitherStreamProj::Plain(stream) => AsyncWrite::poll_shutdown(stream, cx),
84			#[cfg(feature = "tls")]
85			EitherStreamProj::Tls(stream) => AsyncWrite::poll_shutdown(stream, cx),
86		}
87	}
88}