webterm_agent/messaging/
process_f2a.rs

1use crate::config::Config;
2use crate::models::agent_error::AgentError;
3use crate::models::frontend_registry::FrontendRegistry;
4use crate::models::send_payload::SendPayload;
5use crate::models::session_registry::SessionRegistry;
6use webterm_core::flatbuffers_helpers::read_message;
7use webterm_core::generated::flatbuffers_schema::talk_v1::{
8    A2fErrorType, F2aEncryptedRoot, F2aMessage, F2aMessageFormat, F2aPlainMessage, F2aRoot,
9};
10use webterm_core::models::webterm_error::WebtermError;
11use webterm_core::serialisers::talk_v1::a2f_builder::A2fBuilder;
12use webterm_core::serialisers::talk_v1::terminal_output_builder::ActivityInputBlob;
13use webterm_core::types::{ActivityId, Bits96, FrontendId, SessionId};
14
15pub async fn process_f2a(
16    frontend_id: FrontendId,
17    message: &[u8],
18    send: SendPayload,
19    config: &Config,
20) -> Result<SendPayload, AgentError> {
21    let root = read_message::<F2aRoot>(message)?;
22
23    match root.format() {
24        F2aMessageFormat::Plain => process_plain(root, frontend_id, send, config).await,
25        _ => process_encrypted(root, frontend_id, send).await,
26    }
27}
28
29async fn process_plain(
30    message: F2aRoot<'_>,
31    frontend_id: FrontendId,
32    mut send: SendPayload,
33    config: &Config,
34) -> Result<SendPayload, AgentError> {
35    let a2f = A2fBuilder::new();
36
37    match message.plain_message_type() {
38        F2aPlainMessage::AuthRequestPreamble => {
39            let frontend = FrontendRegistry::build_frontend(frontend_id).await?;
40            let frontend = frontend.lock().await;
41            let version_str = env!("CARGO_PKG_VERSION");
42            let version = semver::Version::parse(version_str).unwrap();
43            let payload = a2f
44                .build_preamble(
45                    version,
46                    frontend.salt(),
47                    frontend.pbkdf2_iterations(),
48                    frontend.challenge_nonce()?,
49                )
50                .to_flatbuffers_plain();
51
52            send.prepare_for_frontend(frontend.frontend_id(), payload)
53        }
54
55        F2aPlainMessage::AuthPresentVerification => {
56            let message = message
57                .plain_message_as_auth_present_verification()
58                .unwrap();
59            let frontend_arc = FrontendRegistry::find(frontend_id).await?;
60            let mut frontend = frontend_arc.lock().await;
61
62            let mut success = false;
63            frontend.init_cryptographer(config.secret_key());
64
65            let decrypted = frontend.cryptographer()?.decrypt(
66                message
67                    .challenge_aes256gcm_solution()
68                    .ok_or(AgentError::FBParseError(
69                    "Expected challenge aes256gcm solution for auth present verification, got None"
70                        .to_string(),
71                ))?.bytes(),
72                &Bits96::from(message.challenge_iv().ok_or(AgentError::FBParseError(
73                    "Expected challenge iv for auth present verification, got None".to_string(),
74                ))?),
75                false,
76            );
77
78            if let Ok(decrypted) = decrypted {
79                if decrypted == frontend.challenge_nonce()?.0.to_vec() {
80                    success = true;
81                }
82            } else {
83                success = false;
84            }
85
86            let session_arc = if SessionId(message.resume_session_id()) == SessionId(0) {
87                SessionRegistry::build_session().await?
88            } else {
89                SessionRegistry::find(SessionId(message.resume_session_id())).await?
90            };
91
92            frontend.register_session(session_arc.clone());
93
94            let mut session = session_arc.lock().await;
95            session.set_current_frontend(frontend_arc.clone());
96            frontend.register_session(session_arc.clone());
97            let payload = a2f
98                .build_auth_result(success, session.session_id())
99                .to_flatbuffers_plain();
100
101            send.prepare_for_frontend(frontend.frontend_id(), payload)
102        }
103
104        _ => {
105            return Err(AgentError::FBParseError(format!(
106                "Unknown plain message type: {:?}",
107                message.plain_message_type()
108            )))
109        }
110    }
111
112    Ok(send)
113}
114
115async fn process_encrypted(
116    root: F2aRoot<'_>,
117    frontend_id: FrontendId,
118    mut send: SendPayload,
119) -> Result<SendPayload, AgentError> {
120    let compressed = root.format() == F2aMessageFormat::Aes256GcmDeflateRaw;
121
122    let frontend = FrontendRegistry::find(frontend_id).await?;
123    let frontend = frontend.lock().await;
124
125    let encrypted_payload = root.encrypted_payload();
126
127    let encrypted_payload = encrypted_payload.ok_or(AgentError::FBParseError(format!(
128        "Expected a2f encrypted payload, got None for frontend: {:?}",
129        frontend.frontend_id(),
130    )))?;
131
132    let iv = Bits96::from(root.iv().ok_or(AgentError::FBParseError(
133        "Expected iv for encrypted message, got None".to_string(),
134    ))?);
135
136    let decrypted =
137        match frontend
138            .cryptographer()?
139            .decrypt(encrypted_payload.bytes(), &iv, compressed)
140        {
141            Ok(decrypted) => decrypted,
142            Err(e) => {
143                return match e {
144                    WebtermError::DecryptionError(_) => {
145                        let a2f = A2fBuilder::new();
146                        let error_payload = a2f
147                            .build_error(A2fErrorType::ErrorDecryptionFailed)
148                            .to_flatbuffers_encrypted(frontend.cryptographer()?)?;
149                        send.prepare_for_frontend(frontend.frontend_id(), error_payload);
150                        Ok(send)
151                    }
152                    _ => Err(e.into()),
153                }
154            }
155        };
156
157    let message = read_message::<F2aEncryptedRoot>(&decrypted)?;
158
159    let a2f = A2fBuilder::new();
160    let session = frontend.session().await?;
161    let session = session.lock().await;
162
163    match message.message_type() {
164        F2aMessage::ActivityInput => {
165            let message = message
166                .message_as_activity_input()
167                .ok_or(AgentError::FBParseError(format!(
168                    "Expected activity input for frontend: {:?}, got None",
169                    frontend.frontend_id()
170                )))?;
171            let activity_id = ActivityId(message.activity_id());
172            let input = message.input().ok_or(AgentError::FBParseError(format!(
173                "Expected input for activity: {:?}, got None",
174                activity_id
175            )))?;
176
177            let input = ActivityInputBlob(input.bytes().to_vec());
178            let activity = session.get_activity(&activity_id).await?;
179            send.prepare_for_activity(activity, input)
180        }
181
182        F2aMessage::ActivityCreateTerminal => {
183            let activity = session.create_terminal_activity().await?;
184
185            let payload = a2f
186                .build_activity_create_terminal(activity.activity_id())
187                .to_flatbuffers_encrypted(frontend.cryptographer()?)?;
188            send.prepare_for_frontend(frontend.frontend_id(), payload);
189        }
190
191        _ => {
192            return Err(AgentError::FBParseError(format!(
193                "Unknown encrypted message type: {:#?}",
194                message.message_type()
195            )))
196        }
197    }
198
199    Ok(send)
200}