libp2prs_core/transport/
memory.rs

1// Copyright 2018 Parity Technologies (UK) Ltd.
2// Copyright 2020 Netwarps Ltd.
3//
4// Permission is hereby granted, free of charge, to any person obtaining a
5// copy of this software and associated documentation files (the "Software"),
6// to deal in the Software without restriction, including without limitation
7// the rights to use, copy, modify, merge, publish, distribute, sublicense,
8// and/or sell copies of the Software, and to permit persons to whom the
9// Software is furnished to do so, subject to the following conditions:
10//
11// The above copyright notice and this permission notice shall be included in
12// all copies or substantial portions of the Software.
13//
14// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
15// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
16// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
17// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
18// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
19// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
20// DEALINGS IN THE SOFTWARE.
21
22use async_trait::async_trait;
23use fnv::FnvHashMap;
24use futures::{channel::mpsc, prelude::*, task::Context, task::Poll};
25use futures::{SinkExt, StreamExt};
26use pin_project::pin_project;
27use std::{collections::hash_map::Entry, fmt, io, num::NonZeroU64, pin::Pin};
28
29use lazy_static::lazy_static;
30use libp2prs_multiaddr::{protocol, Multiaddr, Protocol};
31use parking_lot::Mutex;
32use rw_stream_sink::RwStreamSink;
33
34use crate::muxing::{IReadWrite, ReadWriteEx, StreamInfo};
35use crate::transport::{ConnectionInfo, IListener, ITransport, ListenerEvent, TransportListener};
36use crate::{transport::TransportError, Transport};
37// use libp2prs_traits::SplitEx;
38// use futures::io::{ReadHalf, WriteHalf};
39
40lazy_static! {
41    static ref HUB: Mutex<FnvHashMap<NonZeroU64, mpsc::Sender<Channel>>> = Mutex::new(FnvHashMap::default());
42}
43
44/// Transport that supports `/memory/N` multiaddresses.
45///
46/// MemoryTransport is mainly for test purpose, used to write test code for basic transport functionality.
47#[derive(Debug, Clone, Default)]
48pub struct MemoryTransport;
49
50#[async_trait]
51impl Transport for MemoryTransport {
52    type Output = Channel;
53
54    fn listen_on(&mut self, addr: Multiaddr) -> Result<IListener<Self::Output>, TransportError> {
55        let port = if let Ok(port) = parse_memory_addr(&addr) {
56            port
57        } else {
58            return Err(TransportError::MultiaddrNotSupported(addr));
59        };
60
61        let mut hub = (&*HUB).lock();
62
63        let port = if let Some(port) = NonZeroU64::new(port) {
64            port
65        } else {
66            loop {
67                let port = match NonZeroU64::new(rand::random()) {
68                    Some(p) => p,
69                    None => continue,
70                };
71                if !hub.contains_key(&port) {
72                    break port;
73                }
74            }
75        };
76
77        let (tx, rx) = mpsc::channel(2);
78        match hub.entry(port) {
79            Entry::Occupied(_) => return Err(TransportError::Unreachable),
80            Entry::Vacant(e) => e.insert(tx),
81        };
82
83        let listener = Box::new(Listener {
84            port,
85            addr: Protocol::Memory(port.get()).into(),
86            receiver: rx,
87        });
88
89        Ok(listener)
90    }
91
92    async fn dial(&mut self, addr: Multiaddr) -> Result<Self::Output, TransportError> {
93        let port = if let Ok(port) = parse_memory_addr(&addr) {
94            if let Some(port) = NonZeroU64::new(port) {
95                port
96            } else {
97                return Err(TransportError::Unreachable);
98            }
99        } else {
100            return Err(TransportError::MultiaddrNotSupported(addr));
101        };
102
103        // get a cloned sender, unlock the HUB asap
104        let mut sender = {
105            let hub = HUB.lock();
106            if let Some(sender) = hub.get(&port) {
107                sender.clone()
108            } else {
109                return Err(TransportError::Unreachable);
110            }
111        };
112
113        let (a_tx, a_rx) = mpsc::channel(4096);
114        let (b_tx, b_rx) = mpsc::channel(4096);
115
116        let la = Multiaddr::empty();
117        let ra = addr;
118
119        let channel_to_send = Channel {
120            io: RwStreamSink::new(Chan {
121                incoming: a_rx,
122                outgoing: b_tx,
123            }),
124            la: la.clone(),
125            ra: ra.clone(),
126        };
127        let channel_to_return = Channel {
128            io: RwStreamSink::new(Chan {
129                incoming: b_rx,
130                outgoing: a_tx,
131            }),
132            la: la.clone(),
133            ra: ra.clone(),
134        };
135        sender.send(channel_to_send).await.map_err(|_| TransportError::Unreachable)?;
136        Ok(channel_to_return)
137    }
138
139    fn box_clone(&self) -> ITransport<Self::Output> {
140        Box::new(self.clone())
141    }
142
143    fn protocols(&self) -> Vec<u32> {
144        vec![protocol::MEMORY]
145    }
146}
147
148/// Listener for memory connections.
149pub struct Listener {
150    /// Port we're listening on.
151    port: NonZeroU64,
152    /// The address we are listening on.
153    addr: Multiaddr,
154    /// Receives incoming connections.
155    receiver: mpsc::Receiver<Channel>,
156}
157
158#[async_trait]
159impl TransportListener for Listener {
160    type Output = Channel;
161
162    async fn accept(&mut self) -> Result<ListenerEvent<Self::Output>, TransportError> {
163        self.receiver
164            .next()
165            .await
166            .map(ListenerEvent::Accepted)
167            .ok_or(TransportError::Unreachable)
168    }
169
170    fn multi_addr(&self) -> Option<&Multiaddr> {
171        Some(&self.addr)
172    }
173}
174
175impl Drop for Listener {
176    fn drop(&mut self) {
177        let val_in = HUB.lock().remove(&self.port);
178        debug_assert!(val_in.is_some());
179    }
180}
181
182/// If the address is `/memory/n`, returns the value of `n`.
183fn parse_memory_addr(a: &Multiaddr) -> Result<u64, ()> {
184    let mut iter = a.iter();
185
186    let port = if let Some(Protocol::Memory(port)) = iter.next() {
187        port
188    } else {
189        return Err(());
190    };
191
192    if iter.next().is_some() {
193        return Err(());
194    }
195
196    Ok(port)
197}
198
199/// A channel represents an established, in-memory, logical connection between two endpoints.
200///
201/// Implements `ReadEx` and `WriteEx`.
202#[pin_project]
203pub struct Channel {
204    #[pin]
205    io: RwStreamSink<Chan<Vec<u8>>>,
206    la: Multiaddr,
207    ra: Multiaddr,
208}
209
210impl fmt::Debug for Channel {
211    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
212        f.debug_struct("Channel").field("la", &self.la).field("ra", &self.ra).finish()
213    }
214}
215
216// Implements AsyncRead & AsyncWrite
217impl AsyncRead for Channel {
218    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
219        let this = self.project();
220        this.io.poll_read(cx, buf)
221    }
222}
223impl AsyncWrite for Channel {
224    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
225        let this = self.project();
226        this.io.poll_write(cx, buf)
227    }
228
229    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
230        let this = self.project();
231        this.io.poll_flush(cx)
232    }
233
234    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
235        let this = self.project();
236        this.io.poll_close(cx)
237    }
238}
239
240/// A channel represents an established, in-memory, logical connection between two endpoints.
241///
242/// Implements `Sink` and `Stream`.
243pub struct Chan<T = Vec<u8>> {
244    incoming: mpsc::Receiver<T>,
245    outgoing: mpsc::Sender<T>,
246}
247
248impl<T> Unpin for Chan<T> {}
249
250impl<T> Stream for Chan<T> {
251    type Item = Result<T, io::Error>;
252
253    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
254        match Stream::poll_next(Pin::new(&mut self.incoming), cx) {
255            Poll::Pending => Poll::Pending,
256            Poll::Ready(None) => Poll::Ready(Some(Err(io::ErrorKind::BrokenPipe.into()))),
257            Poll::Ready(Some(v)) => Poll::Ready(Some(Ok(v))),
258        }
259    }
260}
261
262impl<T> Sink<T> for Chan<T> {
263    type Error = io::Error;
264
265    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
266        self.outgoing
267            .poll_ready(cx)
268            .map(|v| v.map_err(|_| io::ErrorKind::BrokenPipe.into()))
269    }
270
271    fn start_send(mut self: Pin<&mut Self>, item: T) -> Result<(), Self::Error> {
272        self.outgoing.start_send(item).map_err(|_| io::ErrorKind::BrokenPipe.into())
273    }
274
275    fn poll_flush(self: Pin<&mut Self>, _: &mut Context) -> Poll<Result<(), Self::Error>> {
276        Poll::Ready(Ok(()))
277    }
278
279    fn poll_close(self: Pin<&mut Self>, _: &mut Context) -> Poll<Result<(), Self::Error>> {
280        Poll::Ready(Ok(()))
281    }
282}
283
284impl ConnectionInfo for Channel {
285    fn local_multiaddr(&self) -> Multiaddr {
286        self.la.clone()
287    }
288
289    fn remote_multiaddr(&self) -> Multiaddr {
290        self.ra.clone()
291    }
292}
293
294impl StreamInfo for Channel {
295    fn id(&self) -> usize {
296        0
297    }
298}
299
300impl ReadWriteEx for Channel {
301    fn box_clone(&self) -> IReadWrite {
302        unimplemented!()
303    }
304}
305
306#[cfg(test)]
307mod tests {
308    use super::*;
309    use libp2prs_runtime::task;
310
311    #[test]
312    fn parse_memory_addr_works() {
313        assert_eq!(parse_memory_addr(&"/memory/5".parse().unwrap()), Ok(5));
314        assert_eq!(parse_memory_addr(&"/tcp/150".parse().unwrap()), Err(()));
315        assert_eq!(parse_memory_addr(&"/memory/0".parse().unwrap()), Ok(0));
316        assert_eq!(parse_memory_addr(&"/memory/5/tcp/150".parse().unwrap()), Err(()));
317        assert_eq!(parse_memory_addr(&"/tcp/150/memory/5".parse().unwrap()), Err(()));
318        assert_eq!(parse_memory_addr(&"/memory/1234567890".parse().unwrap()), Ok(1_234_567_890));
319    }
320
321    #[test]
322    fn listening_twice() {
323        let mut transport = MemoryTransport::default();
324        assert!(transport.listen_on("/memory/1639174018481".parse().unwrap()).is_ok());
325        assert!(transport.listen_on("/memory/1639174018481".parse().unwrap()).is_ok());
326        let _listener = transport.listen_on("/memory/1639174018481".parse().unwrap()).unwrap();
327        assert!(transport.listen_on("/memory/1639174018481".parse().unwrap()).is_err());
328        assert!(transport.listen_on("/memory/1639174018481".parse().unwrap()).is_err());
329        drop(_listener);
330        assert!(transport.listen_on("/memory/1639174018481".parse().unwrap()).is_ok());
331        assert!(transport.listen_on("/memory/1639174018481".parse().unwrap()).is_ok());
332    }
333
334    #[test]
335    fn port_not_in_use() {
336        task::block_on(async move {
337            let mut transport = MemoryTransport::default();
338            assert!(transport.dial("/memory/810172461024613".parse().unwrap()).await.is_err());
339            let _listener = transport.listen_on("/memory/810172461024613".parse().unwrap()).unwrap();
340            assert!(transport.dial("/memory/810172461024613".parse().unwrap()).await.is_ok());
341        });
342    }
343
344    #[test]
345    fn communicating_between_dialer_and_listener() {
346        let msg = [1, 2, 3];
347
348        // Setup listener.
349        let rand_port = rand::random::<u64>().saturating_add(1);
350        let t1_addr: Multiaddr = format!("/memory/{}", rand_port).parse().unwrap();
351        let cloned_t1_addr = t1_addr.clone();
352
353        let mut t1 = MemoryTransport::default();
354        let listener = async move {
355            let mut listener = t1.listen_on(t1_addr.clone()).unwrap();
356            let mut socket = match listener.accept().await.unwrap() {
357                ListenerEvent::Accepted(socket) => socket,
358                _ => panic!("unreachable"),
359            };
360
361            let mut buf = [0; 3];
362            socket.read_exact(&mut buf).await.unwrap();
363            assert_eq!(buf, msg);
364        };
365
366        // Setup dialer.
367        let mut t2 = MemoryTransport::default();
368        let dialer = async move {
369            let mut socket = t2.dial(cloned_t1_addr).await.unwrap();
370            socket.write_all(&msg).await.unwrap();
371        };
372
373        // Wait for both to finish.
374        task::block_on(futures::future::join(listener, dialer));
375    }
376}
libp2prs_core/transport/memory.rs

libp2prs_core/transport/
memory.rs
