Implement support to get the peer address

pull/1/head
Jeffrey Esquivel S 6 years ago
parent eb7b3df7c7
commit eda57a012c
No known key found for this signature in database
GPG Key ID: 9020CB6FE67C1864
  1. 4
      examples/client.rs
  2. 1
      examples/server.rs
  3. 18
      src/connect.rs
  4. 16
      src/lib.rs
  5. 16
      src/stream.rs

@ -25,6 +25,7 @@ use futures::{Future, Sink, Stream};
use tungstenite::protocol::Message;
use tokio_tungstenite::connect_async;
use tokio_tungstenite::stream::PeerAddr;
fn main() {
// Specify the server address to which the client will be connecting.
@ -61,6 +62,9 @@ fn main() {
let client = connect_async(url).and_then(move |(ws_stream, _)| {
println!("WebSocket handshake has been successfully completed");
let addr = ws_stream.peer_addr().expect("connected streams should have a peer address");
println!("Peer address: {}", addr);
// `sink` is the stream of messages going out.
// `stream` is the stream of incoming messages.
let (sink, stream) = ws_stream.split();

@ -49,6 +49,7 @@ fn main() {
let srv = socket.incoming().for_each(move |stream| {
let addr = stream.peer_addr().expect("connected streams should have a peer address");
println!("Peer address: {}", addr);
// We have to clone both of these values, because the `and_then`
// function below constructs a new future, `and_then` requires

@ -3,6 +3,7 @@
extern crate tokio_dns;
extern crate tokio_tcp;
use std::net::SocketAddr;
use std::io::Result as IoResult;
use self::tokio_tcp::TcpStream;
@ -14,7 +15,7 @@ use tungstenite::Error;
use tungstenite::client::url_mode;
use tungstenite::handshake::client::Response;
use stream::NoDelay;
use stream::{NoDelay, PeerAddr};
use super::{WebSocketStream, Request, client_async};
impl NoDelay for TcpStream {
@ -23,6 +24,12 @@ impl NoDelay for TcpStream {
}
}
impl PeerAddr for TcpStream {
fn peer_addr(&self) -> IoResult<SocketAddr> {
self.peer_addr()
}
}
#[cfg(feature="tls")]
mod encryption {
extern crate native_tls;
@ -31,6 +38,7 @@ mod encryption {
use self::native_tls::TlsConnector;
use self::tokio_tls::{TlsConnector as TokioTlsConnector, TlsStream};
use std::net::SocketAddr;
use std::io::{Read, Write, Result as IoResult};
use futures::{future, Future};
@ -39,7 +47,7 @@ mod encryption {
use tungstenite::Error;
use tungstenite::stream::Mode;
use stream::{NoDelay, Stream as StreamSwitcher};
use stream::{NoDelay, PeerAddr, Stream as StreamSwitcher};
/// A stream that might be protected with TLS.
pub type MaybeTlsStream<S> = StreamSwitcher<S, TlsStream<S>>;
@ -52,6 +60,12 @@ mod encryption {
}
}
impl<S: Read + Write + PeerAddr> PeerAddr for TlsStream<S> {
fn peer_addr(&self) -> IoResult<SocketAddr> {
self.get_ref().get_ref().peer_addr()
}
}
pub fn wrap_stream<S>(socket: S, domain: String, mode: Mode)
-> Box<Future<Item=AutoStream<S>, Error=Error> + Send>
where

@ -28,6 +28,12 @@ pub mod stream;
use std::io::ErrorKind;
#[cfg(feature="stream")]
use std::{
net::SocketAddr,
io::Result as IoResult,
};
use futures::{Poll, Future, Async, AsyncSink, Stream, Sink, StartSend};
use tokio_io::{AsyncRead, AsyncWrite};
@ -45,6 +51,9 @@ use tungstenite::{
#[cfg(feature="connect")]
pub use connect::{connect_async, client_async_tls};
#[cfg(feature="stream")]
pub use stream::PeerAddr;
#[cfg(all(feature="connect", feature="tls"))]
pub use connect::MaybeTlsStream;
@ -194,6 +203,13 @@ impl<S> WebSocketStream<S> {
}
}
#[cfg(feature="stream")]
impl<S: PeerAddr> PeerAddr for WebSocketStream<S> {
fn peer_addr(&self) -> IoResult<SocketAddr> {
self.inner.get_ref().peer_addr()
}
}
impl<T> Stream for WebSocketStream<T> where T: AsyncRead + AsyncWrite {
type Item = Message;
type Error = WsError;

@ -6,6 +6,7 @@
extern crate bytes;
use std::net::SocketAddr;
use std::io::{Read, Write, Result as IoResult, Error as IoError};
use self::bytes::{Buf, BufMut};
@ -18,6 +19,12 @@ pub trait NoDelay {
fn set_nodelay(&mut self, nodelay: bool) -> IoResult<()>;
}
/// Trait to get the remote address from the underlying stream.
pub trait PeerAddr {
/// Returns the remote address that this stream is connected to.
fn peer_addr(&self) -> IoResult<SocketAddr>;
}
/// Stream, either plain TCP or TLS.
pub enum Stream<S, T> {
/// Unencrypted socket stream.
@ -59,6 +66,15 @@ impl<S: NoDelay, T: NoDelay> NoDelay for Stream<S, T> {
}
}
impl<S: PeerAddr, T: PeerAddr> PeerAddr for Stream<S, T> {
fn peer_addr(&self) -> IoResult<SocketAddr> {
match *self {
Stream::Plain(ref s) => s.peer_addr(),
Stream::Tls(ref s) => s.peer_addr(),
}
}
}
impl<S: AsyncRead, T: AsyncRead> AsyncRead for Stream<S, T> {
unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
match *self {

Loading…
Cancel
Save