webterm-agent 0.2.5

Lightweight agent component for Webterm that enables secure browser-based terminal access
Documentation 
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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

use crate::config::{Config, RELAY_RECONNECT_INTERVAL};
use crate::models::agent_error::AgentError;
use crate::models::handshaker::Handshaker;
use crate::models::relay::Relay;
use crate::models::relay_connection::RelayConnection;
use crate::models::socket_reader::SocketSubscriber;
use crate::models::socket_writer::SocketPublisher;
use std::sync::Arc;
use tokio::sync::Notify;
use tokio::sync::RwLock;
use tokio::time::sleep;
use tracing::{error, info};

pub struct ConnectionManager {
    config: Arc<Config>,
    state: RwLock<Option<State>>,
    connect_signal: Arc<Notify>,
    disconnect_signal: Arc<Notify>,
}

struct State {
    rc: RelayConnection,
}

impl ConnectionManager {
    pub async fn new(config: Arc<Config>) -> Arc<Self> {
        let record = Arc::new(Self {
            config,
            state: RwLock::new(None),
            connect_signal: Arc::new(Notify::new()),
            disconnect_signal: Arc::new(Notify::new()),
        });

        let record_clone = record.clone();
        tokio::spawn(async move {
            Self::run(record_clone).await;
        });

        record
    }

    async fn run(self: Arc<Self>) {
        loop {
            let result = Handshaker::run(self.config.clone()).await;
            match result {
                Err(e) => {
                    error!("Failed to initiate handshake: {}", e)
                }
                Ok(handshake) => {
                    let relay = handshake.result_relay;
                    let nonce = handshake.result_nonce;
                    match self.try_connect(relay.clone(), self.clone(), nonce).await {
                        Ok(_) => {
                            info!("Connected to relay: {}", relay.websocket_url(None));
                            self.connect_signal.notify_waiters();
                            info!("Waiting for disconnect signal");
                            self.wait_for_disconnect().await;
                            info!("Received disconnect signal");
                        }
                        Err(e) => {
                            error!(
                                "Failed to connect to relay {}: {}",
                                relay.websocket_url(None),
                                e
                            );
                        }
                    }
                }
            }
            sleep(RELAY_RECONNECT_INTERVAL).await;
        }
    }

    async fn connect(
        &self,
        relay: Arc<Relay>,
        self_arc: Arc<Self>,
        nonce: String,
    ) -> Result<(), AgentError> {
        let rc = RelayConnection::new(relay.clone(), nonce, self_arc.clone()).await?;
        let mut state = self.state.write().await;
        *state = Some(State { rc });

        Ok(())
    }

    async fn try_connect(
        &self,
        relay: Arc<Relay>,
        self_arc: Arc<Self>,
        nonce: String,
    ) -> Result<(), AgentError> {
        if self.connected().await {
            return Ok(());
        }
        self.connect(relay, self_arc, nonce).await
    }

    pub async fn connected(&self) -> bool {
        self.state.read().await.is_some()
    }

    pub async fn pub_sub(&self) -> Option<(SocketPublisher, SocketSubscriber)> {
        if !self.connected().await {
            return None;
        }

        let state = self.state.read().await;

        if let Some(ref s) = *state {
            let pub_ = s.rc.publisher().await;
            let sub = s.rc.subscriber().await;

            return Some((pub_, sub));
        }

        None
    }

    // TODO: Is this race condition safe?
    // If the connect_signal.notify_waiters() is called by another task right after the connected()
    // check, and right before the notified() call, notified() will wait forever.
    //
    // One easy fix would be to have a separate task which calls connect_signal.notify_waiters()
    // every few seconds if connected() is true.
    pub async fn wait_for_connect(&self) {
        if self.connected().await {
            return;
        }
        self.connect_signal.notified().await;
    }

    pub async fn wait_for_disconnect(&self) {
        if !self.connected().await {
            return;
        }
        self.disconnect_signal.notified().await;
    }

    pub async fn disconnect(&self) {
        let mut state = self.state.write().await;
        *state = None;
        self.disconnect_signal.notify_waiters();
    }

    pub async fn disconnect_with_error(&self, e: AgentError) {
        error!("Disconnecting on error: {}", e);
        self.disconnect().await;
    }
}