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

use std::{
    collections::HashMap,
    io,
};

use async_trait::async_trait;
use tokio::task::JoinHandle;
use url::Url;

use crate::{
    prelude::{
        EdgeInfo,
        EndpointInfo,
        TunnelCloser,
        TunnelInfo,
    },
    session::RpcError,
    Tunnel,
};

/// An ngrok forwarder.
///
/// Represents a tunnel that is being forwarded to a URL.
pub struct Forwarder<T> {
    pub(crate) join: JoinHandle<Result<(), io::Error>>,
    pub(crate) inner: T,
}

impl<T> Forwarder<T> {
    /// Wait for the forwarding task to exit.
    pub fn join(&mut self) -> &mut JoinHandle<Result<(), io::Error>> {
        &mut self.join
    }
}

#[async_trait]
impl<T> TunnelCloser for Forwarder<T>
where
    T: TunnelCloser + Send,
{
    async fn close(&mut self) -> Result<(), RpcError> {
        self.inner.close().await
    }
}

impl<T> TunnelInfo for Forwarder<T>
where
    T: TunnelInfo,
{
    fn id(&self) -> &str {
        self.inner.id()
    }

    fn forwards_to(&self) -> &str {
        self.inner.forwards_to()
    }

    fn metadata(&self) -> &str {
        self.inner.metadata()
    }
}

impl<T> EndpointInfo for Forwarder<T>
where
    T: EndpointInfo,
{
    fn proto(&self) -> &str {
        self.inner.proto()
    }

    fn url(&self) -> &str {
        self.inner.url()
    }
}

impl<T> EdgeInfo for Forwarder<T>
where
    T: EdgeInfo,
{
    fn labels(&self) -> &HashMap<String, String> {
        self.inner.labels()
    }
}

pub(crate) fn forward<T>(mut listener: T, info: T, to_url: Url) -> Result<Forwarder<T>, RpcError>
where
    T: Tunnel + Send + 'static,
    <T as Tunnel>::Conn: crate::tunnel_ext::ConnExt,
{
    let handle =
        tokio::spawn(async move { crate::tunnel_ext::forward_tunnel(&mut listener, to_url).await });

    Ok(Forwarder {
        join: handle,
        inner: info,
    })
}