cdp_core/storage/
manager.rs

1use crate::error::Result;
2use crate::page::Page;
3use async_trait::async_trait;
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::sync::Arc;
7
8/// Storage item representing a key-value pair
9#[derive(Debug, Serialize, Deserialize, Clone)]
10pub struct StorageItem {
11    pub key: String,
12    pub value: String,
13}
14
15/// Storage type enumeration
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum StorageType {
18    Local,
19    Session,
20}
21
22impl StorageType {
23    fn as_str(&self) -> &'static str {
24        match self {
25            StorageType::Local => "localStorage",
26            StorageType::Session => "sessionStorage",
27        }
28    }
29}
30
31/// Storage manager trait for localStorage and sessionStorage operations
32#[async_trait]
33pub trait StorageManager {
34    /// Gets all items from the specified storage type
35    async fn get_storage_items(
36        self: &Arc<Self>,
37        storage_type: StorageType,
38    ) -> Result<Vec<StorageItem>>;
39
40    /// Gets a specific item from storage by key
41    async fn get_storage_item(
42        self: &Arc<Self>,
43        storage_type: StorageType,
44        key: &str,
45    ) -> Result<Option<String>>;
46
47    /// Sets an item in storage
48    async fn set_storage_item(
49        self: &Arc<Self>,
50        storage_type: StorageType,
51        key: &str,
52        value: &str,
53    ) -> Result<()>;
54
55    /// Removes an item from storage by key
56    async fn remove_storage_item(
57        self: &Arc<Self>,
58        storage_type: StorageType,
59        key: &str,
60    ) -> Result<()>;
61
62    /// Clears all items from the specified storage type
63    async fn clear_storage(self: &Arc<Self>, storage_type: StorageType) -> Result<()>;
64
65    /// Gets the number of items in storage
66    async fn get_storage_length(self: &Arc<Self>, storage_type: StorageType) -> Result<usize>;
67}
68
69#[async_trait]
70impl StorageManager for Page {
71    async fn get_storage_items(
72        self: &Arc<Self>,
73        storage_type: StorageType,
74    ) -> Result<Vec<StorageItem>> {
75        let storage_name = storage_type.as_str();
76        let script = format!(
77            r#"
78            (() => {{
79                const storage = {};
80                const items = [];
81                for (let i = 0; i < storage.length; i++) {{
82                    const key = storage.key(i);
83                    if (key !== null) {{
84                        items.push({{ key: key, value: storage.getItem(key) }});
85                    }}
86                }}
87                return items;
88            }})()
89            "#,
90            storage_name
91        );
92
93        let main_frame = self.main_frame().await?;
94        let result = main_frame.evaluate(&script).await?;
95
96        let items: Vec<StorageItem> = serde_json::from_value(result)?;
97        Ok(items)
98    }
99
100    async fn get_storage_item(
101        self: &Arc<Self>,
102        storage_type: StorageType,
103        key: &str,
104    ) -> Result<Option<String>> {
105        let storage_name = storage_type.as_str();
106        let script = format!(
107            r#"
108            (() => {{
109                const storage = {};
110                return storage.getItem({});
111            }})()
112            "#,
113            storage_name,
114            serde_json::to_string(key)?
115        );
116
117        let main_frame = self.main_frame().await?;
118        let result = main_frame.evaluate(&script).await?;
119
120        if result.is_null() {
121            Ok(None)
122        } else {
123            Ok(Some(result.as_str().unwrap_or("").to_string()))
124        }
125    }
126
127    async fn set_storage_item(
128        self: &Arc<Self>,
129        storage_type: StorageType,
130        key: &str,
131        value: &str,
132    ) -> Result<()> {
133        let storage_name = storage_type.as_str();
134        let script = format!(
135            r#"
136            (() => {{
137                const storage = {};
138                storage.setItem({}, {});
139            }})()
140            "#,
141            storage_name,
142            serde_json::to_string(key)?,
143            serde_json::to_string(value)?
144        );
145
146        let main_frame = self.main_frame().await?;
147        main_frame.evaluate(&script).await?;
148        Ok(())
149    }
150
151    async fn remove_storage_item(
152        self: &Arc<Self>,
153        storage_type: StorageType,
154        key: &str,
155    ) -> Result<()> {
156        let storage_name = storage_type.as_str();
157        let script = format!(
158            r#"
159            (() => {{
160                const storage = {};
161                storage.removeItem({});
162            }})()
163            "#,
164            storage_name,
165            serde_json::to_string(key)?
166        );
167
168        let main_frame = self.main_frame().await?;
169        main_frame.evaluate(&script).await?;
170        Ok(())
171    }
172
173    async fn clear_storage(self: &Arc<Self>, storage_type: StorageType) -> Result<()> {
174        let storage_name = storage_type.as_str();
175        let script = format!(
176            r#"
177            (() => {{
178                const storage = {};
179                storage.clear();
180            }})()
181            "#,
182            storage_name
183        );
184
185        let main_frame = self.main_frame().await?;
186        main_frame.evaluate(&script).await?;
187        Ok(())
188    }
189
190    async fn get_storage_length(self: &Arc<Self>, storage_type: StorageType) -> Result<usize> {
191        let storage_name = storage_type.as_str();
192        let script = format!(
193            r#"
194            (() => {{
195                const storage = {};
196                return storage.length;
197            }})()
198            "#,
199            storage_name
200        );
201
202        let main_frame = self.main_frame().await?;
203        let result = main_frame.evaluate(&script).await?;
204
205        let length = result.as_u64().unwrap_or(0) as usize;
206        Ok(length)
207    }
208}
209
210/// Helper methods for convenient access to localStorage
211#[async_trait]
212pub trait LocalStorageExt {
213    async fn get_local_storage(self: &Arc<Self>) -> Result<HashMap<String, String>>;
214    async fn get_local_storage_item(self: &Arc<Self>, key: &str) -> Result<Option<String>>;
215    async fn set_local_storage_item(self: &Arc<Self>, key: &str, value: &str) -> Result<()>;
216    async fn remove_local_storage_item(self: &Arc<Self>, key: &str) -> Result<()>;
217    async fn clear_local_storage(self: &Arc<Self>) -> Result<()>;
218}
219
220#[async_trait]
221impl LocalStorageExt for Page {
222    async fn get_local_storage(self: &Arc<Self>) -> Result<HashMap<String, String>> {
223        let items = self.get_storage_items(StorageType::Local).await?;
224        Ok(items
225            .into_iter()
226            .map(|item| (item.key, item.value))
227            .collect())
228    }
229
230    async fn get_local_storage_item(self: &Arc<Self>, key: &str) -> Result<Option<String>> {
231        self.get_storage_item(StorageType::Local, key).await
232    }
233
234    async fn set_local_storage_item(self: &Arc<Self>, key: &str, value: &str) -> Result<()> {
235        self.set_storage_item(StorageType::Local, key, value).await
236    }
237
238    async fn remove_local_storage_item(self: &Arc<Self>, key: &str) -> Result<()> {
239        self.remove_storage_item(StorageType::Local, key).await
240    }
241
242    async fn clear_local_storage(self: &Arc<Self>) -> Result<()> {
243        self.clear_storage(StorageType::Local).await
244    }
245}
246
247/// Helper methods for convenient access to sessionStorage
248#[async_trait]
249pub trait SessionStorageExt {
250    async fn get_session_storage(self: &Arc<Self>) -> Result<HashMap<String, String>>;
251    async fn get_session_storage_item(self: &Arc<Self>, key: &str) -> Result<Option<String>>;
252    async fn set_session_storage_item(self: &Arc<Self>, key: &str, value: &str) -> Result<()>;
253    async fn remove_session_storage_item(self: &Arc<Self>, key: &str) -> Result<()>;
254    async fn clear_session_storage(self: &Arc<Self>) -> Result<()>;
255}
256
257#[async_trait]
258impl SessionStorageExt for Page {
259    async fn get_session_storage(self: &Arc<Self>) -> Result<HashMap<String, String>> {
260        let items = self.get_storage_items(StorageType::Session).await?;
261        Ok(items
262            .into_iter()
263            .map(|item| (item.key, item.value))
264            .collect())
265    }
266
267    async fn get_session_storage_item(self: &Arc<Self>, key: &str) -> Result<Option<String>> {
268        self.get_storage_item(StorageType::Session, key).await
269    }
270
271    async fn set_session_storage_item(self: &Arc<Self>, key: &str, value: &str) -> Result<()> {
272        self.set_storage_item(StorageType::Session, key, value)
273            .await
274    }
275
276    async fn remove_session_storage_item(self: &Arc<Self>, key: &str) -> Result<()> {
277        self.remove_storage_item(StorageType::Session, key).await
278    }
279
280    async fn clear_session_storage(self: &Arc<Self>) -> Result<()> {
281        self.clear_storage(StorageType::Session).await
282    }
283}
cdp_core/storage/manager.rs

cdp_core/storage/
manager.rs
