rustenium_core/
network.rs

1use crate::error::{PostDataError, SessionSendError};
2use crate::transport::ConnectionTransport;
3use crate::{CommandResponseState, Session};
4use form_urlencoded;
5use tokio::sync::oneshot;
6
7/// How a network request was handled
8#[derive(Debug, Clone)]
9pub enum NetworkRequestHandledState {
10    /// Request was continued (with or without modifications)
11    Continued,
12    /// Request was aborted/failed
13    Aborted,
14    /// Request was responded with a custom response
15    Responded,
16}
17use rustenium_bidi_commands::network::commands::{
18    ContinueRequest, ContinueRequestParameters, ContinueWithAuth, ContinueWithAuthParameters,
19    FailRequest, NetworkFailRequestMethod, FailRequestParameters, NetworkContinueRequestMethod,
20    NetworkContinueWithAuthMethod, ProvideResponse, NetworkProvideResponseMethod,
21    ProvideResponseParameters,
22};
23use rustenium_bidi_commands::network::types::{
24    AuthRequiredParameters, BaseParameters, BeforeRequestSentParameters,
25    ContinueWithAuthCredentials, ContinueWithAuthNoCredentials,
26    ContinueWithAuthCredentialsContinueWithAuthNoCredentialsUnion,
27    ContinueWithAuthNoCredentialsactionUnion, ProvideCredentialsEnum,
28};
29use rustenium_bidi_commands::network::types::{BytesValue, CookieHeader, SetCookieHeader, Header};
30use rustenium_bidi_commands::{CommandData, NetworkCommand};
31use serde_json;
32use std::collections::HashMap;
33use std::sync::Arc;
34use tokio::sync::Mutex;
35
36/// Represents a network request that can be intercepted
37pub struct  NetworkRequest<T: ConnectionTransport> {
38    pub base: BaseParameters,
39    session: Arc<Mutex<Session<T>>>,
40}
41
42impl<T: ConnectionTransport> std::fmt::Debug for NetworkRequest<T> {
43    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44        f.debug_struct("NetworkRequest")
45            .field("base", &self.base)
46            .field("session", &"<Arc<Mutex<Session>>>")
47            .finish()
48    }
49}
50
51impl<T: ConnectionTransport> NetworkRequest<T> {
52    pub fn new(params: BeforeRequestSentParameters, session: Arc<Mutex<Session<T>>>) -> Self {
53        NetworkRequest {
54            base: params.base_parameters,
55            session
56        }
57    }
58
59    pub fn from_auth_required(params: AuthRequiredParameters, session: Arc<Mutex<Session<T>>>) -> Self {
60        NetworkRequest {
61            base: params.base_parameters,
62            session
63        }
64    }
65
66    /// Get the request ID
67    pub fn request_id(&self) -> &str {
68        &self.base.request.request
69    }
70
71    /// Get the request URL
72    pub fn url(&self) -> &str {
73        &self.base.request.url
74    }
75
76    /// Get the request headers
77    pub fn headers(&self) -> &Vec<Header> {
78        &self.base.request.headers
79    }
80
81    /// Get the request method
82    pub fn method(&self) -> &str {
83        &self.base.request.method
84    }
85
86    /// Check if the request has POST data (Chrome-specific)
87    pub fn has_post_data(&self) -> bool {
88        self.base
89            .request
90            .extensible
91            .get("goog:hasPostData")
92            .and_then(|v| v.as_bool())
93            .unwrap_or(false)
94    }
95
96    /// Get the POST data as a raw string if available (Chrome-specific)
97    pub fn post_data(&self) -> Option<&str> {
98        self.base
99            .request
100            .extensible
101            .get("goog:postData")
102            .and_then(|v| v.as_str())
103    }
104
105    /// Parse the POST data as JSON (Chrome-specific)
106    /// Returns an error if POST data is missing, not valid JSON, or not a JSON object
107    pub fn post_data_json(&self) -> Result<serde_json::Map<String, serde_json::Value>, PostDataError> {
108        let data = self.post_data()
109            .ok_or(PostDataError::NoPostData)?;
110
111        let value: serde_json::Value = serde_json::from_str(data)?;
112
113        value.as_object()
114            .cloned()
115            .ok_or(PostDataError::NotJsonObject)
116    }
117
118    /// Parse the POST data as URL-encoded form data (Chrome-specific)
119    /// Returns an error if POST data is missing
120    pub fn post_data_form(&self) -> Result<HashMap<String, String>, PostDataError> {
121        let data = self.post_data()
122            .ok_or(PostDataError::NoPostData)?;
123
124        let parsed: HashMap<String, String> = form_urlencoded::parse(data.as_bytes())
125            .into_owned()
126            .collect();
127
128        Ok(parsed)
129    }
130
131    /// Check if this request has already been handled
132    pub async fn is_handled(&self) -> bool {
133        self.session.lock().await.handled_network_requests
134            .lock()
135            .unwrap()
136            .contains_key(&self.base.request.request)
137    }
138
139    /// Get the handled state if the request was already handled
140    pub async fn get_handled_state(&self) -> Option<NetworkRequestHandledState> {
141        self.session.lock().await
142            .handled_network_requests
143            .lock()
144            .unwrap()
145            .get(&self.base.request.request)
146            .cloned()
147    }
148
149    /// Mark this request as handled with the given state
150    async fn mark_handled(&self, state: NetworkRequestHandledState) {
151        let session = self.session.lock().await;
152        session
153            .handled_network_requests
154            .lock()
155            .unwrap()
156            .insert(
157                self.base.request.request.clone(),
158                state,
159            );
160    }
161
162    /// Continue the request without modifications
163    pub async fn continue_(&self) -> oneshot::Receiver<CommandResponseState> {
164        let command =
165            CommandData::NetworkCommand(NetworkCommand::ContinueRequest(ContinueRequest {
166                method: NetworkContinueRequestMethod::NetworkContinueRequest,
167                params: ContinueRequestParameters {
168                    request: self.base.request.request.clone(),
169                    body: None,
170                    cookies: None,
171                    headers: None,
172                    method: None,
173                    url: None,
174                },
175            }));
176
177        let rx = self.session.lock().await.send_and_get_receiver(command).await;
178        self.mark_handled(NetworkRequestHandledState::Continued).await;
179        rx
180    }
181
182    /// Continue the request with modifications
183    pub async fn continue_with(
184        &self,
185        headers: Option<Vec<Header>>,
186        cookies: Option<Vec<CookieHeader>>,
187        url: Option<String>,
188        method: Option<String>,
189        body: Option<BytesValue>,
190    ) -> oneshot::Receiver<CommandResponseState> {
191        let command =
192            CommandData::NetworkCommand(NetworkCommand::ContinueRequest(ContinueRequest {
193                method: NetworkContinueRequestMethod::NetworkContinueRequest,
194                params: ContinueRequestParameters {
195                    request: self.base.request.request.clone(),
196                    body,
197                    cookies,
198                    headers,
199                    method,
200                    url,
201                },
202            }));
203
204        let rx = self.session.lock().await.send_and_get_receiver(command).await;
205        self.mark_handled(NetworkRequestHandledState::Continued).await;
206        rx
207    }
208
209    /// Abort/fail the request
210    pub async fn abort(&self) -> oneshot::Receiver<CommandResponseState> {
211        let command = CommandData::NetworkCommand(NetworkCommand::FailRequest(FailRequest {
212            method: NetworkFailRequestMethod::NetworkFailRequest,
213            params: FailRequestParameters {
214                request: self.base.request.request.clone(),
215            },
216        }));
217
218        let rx = self.session.lock().await.send_and_get_receiver(command).await;
219        self.mark_handled(NetworkRequestHandledState::Aborted).await;
220        rx
221    }
222
223    /// Provide a custom response
224    pub async fn respond(
225        &self,
226        status_code: Option<u64>,
227        reason_phrase: Option<String>,
228        headers: Option<Vec<Header>>,
229        cookies: Option<Vec<SetCookieHeader>>,
230        body: Option<BytesValue>,
231    ) -> oneshot::Receiver<CommandResponseState> {
232        let command =
233            CommandData::NetworkCommand(NetworkCommand::ProvideResponse(ProvideResponse {
234                method: NetworkProvideResponseMethod::NetworkProvideResponse,
235                params: ProvideResponseParameters {
236                    request: self.base.request.request.clone(),
237                    body,
238                    cookies,
239                    headers,
240                    reason_phrase,
241                    status_code,
242                },
243            }));
244
245        let rx = self.session.lock().await.send_and_get_receiver(command).await;
246        self.mark_handled(NetworkRequestHandledState::Responded).await;
247        rx
248    }
249
250    /// Continue with HTTP authentication
251    pub async fn continue_with_auth(&self, credentials: Option<rustenium_bidi_commands::network::types::AuthCredentials>) -> Result<(), SessionSendError> {
252        let auth_union = match credentials {
253            Some(credentials) => {
254                ContinueWithAuthCredentialsContinueWithAuthNoCredentialsUnion::ContinueWithAuthCredentials(
255                    ContinueWithAuthCredentials {
256                        action: ProvideCredentialsEnum::ProvideCredentials,
257                        credentials,
258                    }
259                )
260            },
261            None => {
262                ContinueWithAuthCredentialsContinueWithAuthNoCredentialsUnion::ContinueWithAuthNoCredentials(
263                    ContinueWithAuthNoCredentials {
264                        action: ContinueWithAuthNoCredentialsactionUnion::Default,
265                    }
266                )
267            }
268        };
269
270        let command = CommandData::NetworkCommand(
271            NetworkCommand::ContinueWithAuth(ContinueWithAuth {
272                method: NetworkContinueWithAuthMethod::NetworkContinueWithAuth,
273                params: ContinueWithAuthParameters {
274                    request: self.base.request.request.clone(),
275                    continue_with_auth_credentials_continue_with_auth_no_credentials_union: auth_union,
276                },
277            })
278        );
279
280        self.session.lock().await.send(command).await.map(|_| ())
281    }
282}
rustenium_core/network.rs

rustenium_core/
network.rs
