rustenium_core/
network.rs

1use crate::error::{PostDataError, SessionSendError};
2use crate::transport::ConnectionTransport;
3use crate::{BidiSession, CommandResponseState};
4use form_urlencoded;
5use rustenium_bidi_definitions::Command;
6use rustenium_bidi_definitions::network::command_builders::{
7    ContinueRequestBuilder, ContinueWithAuthBuilder, FailRequestBuilder, ProvideResponseBuilder
8};
9use rustenium_bidi_definitions::network::events::{AuthRequiredParams, BeforeRequestSentParams
10};
11use rustenium_bidi_definitions::network::type_builders::{ContinueWithAuthCredentialsBuilder, ContinueWithAuthNoCredentialsBuilder};
12use rustenium_bidi_definitions::network::types::{AuthCredentials, BaseParameters, ContinueWithAuthCredentialsAction, ContinueWithAuthCredentialsContinueWithAuthNoCredentialsUnion, ContinueWithAuthNoCredentialsAction, Header, Request};
13use tokio::sync::oneshot;
14
15/// How a network request was handled
16#[derive(Debug, Clone)]
17pub enum NetworkRequestHandledState {
18    /// Request was continued (with or without modifications)
19    Continued,
20    /// Request was aborted/failed
21    Aborted,
22    /// Request was responded with a custom response
23    Responded,
24}
25use serde_json;
26use std::collections::HashMap;
27use std::sync::Arc;
28use tokio::sync::Mutex;
29
30/// Represents a network request that can be intercepted
31pub struct NetworkRequest<T: ConnectionTransport> {
32    pub base: BaseParameters,
33    session: Arc<Mutex<BidiSession<T>>>,
34}
35
36impl<T: ConnectionTransport> std::fmt::Debug for NetworkRequest<T> {
37    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38        f.debug_struct("NetworkRequest")
39            .field("base", &self.base)
40            .field("session", &"<Arc<Mutex<Session>>>")
41            .finish()
42    }
43}
44
45impl<T: ConnectionTransport> NetworkRequest<T> {
46    pub fn new(params: BeforeRequestSentParams, session: Arc<Mutex<BidiSession<T>>>) -> Self {
47        NetworkRequest {
48            base: params.base_parameters,
49            session,
50        }
51    }
52
53    pub fn from_auth_required(params: AuthRequiredParams, session: Arc<Mutex<BidiSession<T>>>) -> Self {
54        NetworkRequest {
55            base: params.base_parameters,
56            session,
57        }
58    }
59
60    /// Get the request ID
61    pub fn request_id(&self) -> &Request {
62        &self.base.request.request
63    }
64
65    /// Get the request URL
66    pub fn url(&self) -> &str {
67        &self.base.request.url
68    }
69
70    /// Get the request headers
71    pub fn headers(&self) -> &Vec<Header> {
72        &self.base.request.headers
73    }
74
75    /// Get the request method
76    pub fn method(&self) -> &str {
77        &self.base.request.method
78    }
79
80    /// Check if the request has POST data (Chrome-specific)
81    pub fn has_post_data(&self) -> bool {
82        self.base
83            .request
84            .extensible
85            .get("goog:hasPostData")
86            .and_then(|v| v.as_bool())
87            .unwrap_or(false)
88    }
89
90    /// Get the POST data as a raw string if available (Chrome-specific)
91    pub fn post_data(&self) -> Option<&str> {
92        self.base
93            .request
94            .extensible
95            .get("goog:postData")
96            .and_then(|v| v.as_str())
97    }
98
99    /// Parse the POST data as JSON (Chrome-specific)
100    /// Returns an error if POST data is missing, not valid JSON, or not a JSON object
101    pub fn post_data_json(
102        &self,
103    ) -> Result<serde_json::Map<String, serde_json::Value>, PostDataError> {
104        let data = self.post_data().ok_or(PostDataError::NoPostData)?;
105
106        let value: serde_json::Value = serde_json::from_str(data)?;
107
108        value
109            .as_object()
110            .cloned()
111            .ok_or(PostDataError::NotJsonObject)
112    }
113
114    /// Parse the POST data as URL-encoded form data (Chrome-specific)
115    /// Returns an error if POST data is missing
116    pub fn post_data_form(&self) -> Result<HashMap<String, String>, PostDataError> {
117        let data = self.post_data().ok_or(PostDataError::NoPostData)?;
118
119        let parsed: HashMap<String, String> = form_urlencoded::parse(data.as_bytes())
120            .into_owned()
121            .collect();
122
123        Ok(parsed)
124    }
125
126    /// Check if this request has already been handled
127    pub async fn is_handled(&self) -> bool {
128        self.session
129            .lock()
130            .await
131            .handled_network_requests
132            .lock()
133            .unwrap()
134            .contains_key(self.request_id().as_ref())
135    }
136
137    /// Get the handled state if the request was already handled
138    pub async fn get_handled_state(&self) -> Option<NetworkRequestHandledState> {
139        self.session
140            .lock()
141            .await
142            .handled_network_requests
143            .lock()
144            .unwrap()
145            .get(self.request_id().as_ref())
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(self.base.request.request.clone().into(), state);
157    }
158
159    /// Continue the request without modifications
160    pub async fn continue_(&self) -> oneshot::Receiver<CommandResponseState> {
161        let command: Command = ContinueRequestBuilder::default()
162            .request(self.base.request.request.clone())
163            .build()
164            .unwrap()
165            .into(); // This should never result in an error
166
167        let rx = self
168            .session
169            .lock()
170            .await
171            .send_and_get_receiver(command)
172            .await;
173        self.mark_handled(NetworkRequestHandledState::Continued)
174            .await;
175        rx
176    }
177
178    /// Continue the request with modifications
179    pub async fn continue_with(
180        &self,
181        continue_request: ContinueRequestBuilder,
182    ) -> oneshot::Receiver<CommandResponseState> {
183        let continue_request = continue_request.request(self.base.request.request.clone()).build().unwrap();
184        let rx = self
185            .session
186            .lock()
187            .await
188            .send_and_get_receiver(continue_request)
189            .await;
190        self.mark_handled(NetworkRequestHandledState::Continued)
191            .await;
192        rx
193    }
194
195    /// Abort/fail the request
196    pub async fn abort(&self) -> oneshot::Receiver<CommandResponseState> {
197        let command = FailRequestBuilder::default()
198            .request(self.base.request.request.clone())
199            .build().unwrap();
200
201        let rx = self
202            .session
203            .lock()
204            .await
205            .send_and_get_receiver(command)
206            .await;
207        self.mark_handled(NetworkRequestHandledState::Aborted).await;
208        rx
209    }
210
211    /// Provide a custom response
212    pub async fn respond(
213        &self,
214        provide_response_builder: ProvideResponseBuilder
215    ) -> oneshot::Receiver<CommandResponseState> {
216        let command = provide_response_builder.request(self.request_id().clone()).build().unwrap();
217
218        let rx = self
219            .session
220            .lock()
221            .await
222            .send_and_get_receiver(command)
223            .await;
224        self.mark_handled(NetworkRequestHandledState::Responded)
225            .await;
226        rx
227    }
228
229    /// Continue with HTTP authentication
230    pub async fn continue_with_auth(
231        &self,
232        credentials: AuthCredentials,
233    ) -> Result<(), SessionSendError> {
234        let command = 
235                ContinueWithAuthBuilder::default().continue_with_auth_credentials_continue_with_auth_no_credentials_union(ContinueWithAuthCredentialsContinueWithAuthNoCredentialsUnion::ContinueWithAuthCredentials(
236                    ContinueWithAuthCredentialsBuilder::default().action(ContinueWithAuthCredentialsAction::ProvideCredentials).credentials(
237                        credentials).build().unwrap())).request(self.request_id().clone()).build().unwrap();
238
239        self.session.lock().await.send(command).await.map(|_| ())
240    }
241    pub async fn continue_with_no_auth(&self, action: ContinueWithAuthNoCredentialsAction) -> Result<(), SessionSendError> {
242        let command = ContinueWithAuthBuilder::default().continue_with_auth_credentials_continue_with_auth_no_credentials_union(ContinueWithAuthCredentialsContinueWithAuthNoCredentialsUnion::ContinueWithAuthNoCredentials(
243                    ContinueWithAuthNoCredentialsBuilder::default().action(action).build().unwrap()
244                )).request(self.request_id().clone()).build().unwrap();
245
246                
247        self.session.lock().await.send(command).await.map(|_| ())
248    }
249}
rustenium_core/network.rs

rustenium_core/
network.rs
