blueprint_engine_core/
context.rs

1use std::collections::HashSet;
2use std::io::{self, Write};
3use std::sync::Arc;
4use tokio::sync::RwLock;
5use tokio::task_local;
6
7use crate::{BlueprintError, PermissionCheck, Permissions, Result};
8
9task_local! {
10    static PERMISSIONS: Arc<Permissions>;
11    static PROMPT_STATE: Arc<PromptState>;
12}
13
14pub struct PromptState {
15    session_allowed: RwLock<HashSet<String>>,
16    session_denied: RwLock<HashSet<String>>,
17    interactive: bool,
18}
19
20impl PromptState {
21    pub fn new(interactive: bool) -> Self {
22        Self {
23            session_allowed: RwLock::new(HashSet::new()),
24            session_denied: RwLock::new(HashSet::new()),
25            interactive,
26        }
27    }
28}
29
30impl Default for PromptState {
31    fn default() -> Self {
32        Self::new(true)
33    }
34}
35
36pub fn with_permissions<F, R>(permissions: Arc<Permissions>, f: F) -> R
37where
38    F: FnOnce() -> R,
39{
40    let prompt_state = Arc::new(PromptState::default());
41    PERMISSIONS.sync_scope(permissions, || PROMPT_STATE.sync_scope(prompt_state, f))
42}
43
44pub async fn with_permissions_async<F, Fut, R>(permissions: Arc<Permissions>, f: F) -> R
45where
46    F: FnOnce() -> Fut,
47    Fut: std::future::Future<Output = R>,
48{
49    let prompt_state = Arc::new(PromptState::default());
50    PERMISSIONS
51        .scope(permissions, async {
52            PROMPT_STATE.scope(prompt_state, f()).await
53        })
54        .await
55}
56
57pub async fn with_permissions_and_prompt<F, Fut, R>(
58    permissions: Arc<Permissions>,
59    prompt_state: Arc<PromptState>,
60    f: F,
61) -> R
62where
63    F: FnOnce() -> Fut,
64    Fut: std::future::Future<Output = R>,
65{
66    PERMISSIONS
67        .scope(permissions, async {
68            PROMPT_STATE.scope(prompt_state, f()).await
69        })
70        .await
71}
72
73pub fn get_permissions() -> Option<Arc<Permissions>> {
74    PERMISSIONS.try_with(|p| p.clone()).ok()
75}
76
77fn get_prompt_state() -> Option<Arc<PromptState>> {
78    PROMPT_STATE.try_with(|p| p.clone()).ok()
79}
80
81async fn handle_permission_check(
82    check: PermissionCheck,
83    operation: &str,
84    resource: Option<&str>,
85) -> Result<()> {
86    match check {
87        PermissionCheck::Allow => Ok(()),
88        PermissionCheck::Deny => {
89            let resource_str = resource.unwrap_or("");
90            Err(BlueprintError::PermissionDenied {
91                operation: operation.into(),
92                resource: resource_str.into(),
93                hint: format!(
94                    "Add '{}:{}' to permissions.allow in BP.toml",
95                    operation,
96                    if resource_str.is_empty() {
97                        "*"
98                    } else {
99                        resource_str
100                    }
101                ),
102            })
103        }
104        PermissionCheck::Ask => {
105            let key = match resource {
106                Some(r) => format!("{}:{}", operation, r),
107                None => operation.to_string(),
108            };
109
110            if let Some(state) = get_prompt_state() {
111                if state.session_allowed.read().await.contains(&key) {
112                    return Ok(());
113                }
114                if state.session_denied.read().await.contains(&key) {
115                    return Err(BlueprintError::PermissionDenied {
116                        operation: operation.into(),
117                        resource: resource.unwrap_or("").into(),
118                        hint: "Permission was denied earlier in this session".into(),
119                    });
120                }
121
122                if state.interactive {
123                    let allowed = prompt_user(operation, resource).await?;
124                    if allowed {
125                        state.session_allowed.write().await.insert(key);
126                        return Ok(());
127                    } else {
128                        state.session_denied.write().await.insert(key);
129                        return Err(BlueprintError::PermissionDenied {
130                            operation: operation.into(),
131                            resource: resource.unwrap_or("").into(),
132                            hint: "Permission denied by user".into(),
133                        });
134                    }
135                }
136            }
137
138            Err(BlueprintError::PermissionDenied {
139                operation: operation.into(),
140                resource: resource.unwrap_or("").into(),
141                hint: format!(
142                    "Add '{}:{}' to permissions.allow in BP.toml (or run interactively to be prompted)",
143                    operation,
144                    resource.unwrap_or("*")
145                ),
146            })
147        }
148    }
149}
150
151async fn prompt_user(operation: &str, resource: Option<&str>) -> Result<bool> {
152    let resource_display = resource.unwrap_or("");
153
154    eprintln!();
155    eprintln!("┌─────────────────────────────────────────────────────────────────┐");
156    eprintln!("│ Permission Request                                              │");
157    eprintln!("├─────────────────────────────────────────────────────────────────┤");
158    eprintln!("│ Operation: {:<52} │", operation);
159    if !resource_display.is_empty() {
160        let truncated = if resource_display.len() > 52 {
161            format!("...{}", &resource_display[resource_display.len() - 49..])
162        } else {
163            resource_display.to_string()
164        };
165        eprintln!("│ Resource:  {:<52} │", truncated);
166    }
167    eprintln!("├─────────────────────────────────────────────────────────────────┤");
168    eprintln!("│ [y] Allow   [n] Deny   [Y] Allow all similar   [N] Deny all    │");
169    eprintln!("└─────────────────────────────────────────────────────────────────┘");
170    eprint!("Choice: ");
171    io::stderr().flush().ok();
172
173    let mut input = String::new();
174
175    tokio::task::spawn_blocking(move || {
176        io::stdin().read_line(&mut input).ok();
177        input.trim().to_lowercase()
178    })
179    .await
180    .map(|response| matches!(response.as_str(), "y" | "yes" | ""))
181    .map_err(|e| BlueprintError::IoError {
182        path: "stdin".into(),
183        message: e.to_string(),
184    })
185}
186
187pub async fn check_fs_read(path: &str) -> Result<()> {
188    match get_permissions() {
189        None => Ok(()),
190        Some(p) => {
191            let check = p.check_fs_read(path);
192            handle_permission_check(check, "fs.read", Some(path)).await
193        }
194    }
195}
196
197pub async fn check_fs_write(path: &str) -> Result<()> {
198    match get_permissions() {
199        None => Ok(()),
200        Some(p) => {
201            let check = p.check_fs_write(path);
202            handle_permission_check(check, "fs.write", Some(path)).await
203        }
204    }
205}
206
207pub async fn check_fs_delete(path: &str) -> Result<()> {
208    match get_permissions() {
209        None => Ok(()),
210        Some(p) => {
211            let check = p.check_fs_delete(path);
212            handle_permission_check(check, "fs.delete", Some(path)).await
213        }
214    }
215}
216
217pub async fn check_http(url: &str) -> Result<()> {
218    match get_permissions() {
219        None => Ok(()),
220        Some(p) => {
221            let check = p.check_http(url);
222            handle_permission_check(check, "net.http", Some(url)).await
223        }
224    }
225}
226
227pub async fn check_ws(url: &str) -> Result<()> {
228    match get_permissions() {
229        None => Ok(()),
230        Some(p) => {
231            let check = p.check_ws(url);
232            handle_permission_check(check, "net.ws", Some(url)).await
233        }
234    }
235}
236
237pub async fn check_process_run(binary: &str) -> Result<()> {
238    match get_permissions() {
239        None => Ok(()),
240        Some(p) => {
241            let check = p.check_process_run(binary);
242            handle_permission_check(check, "process.run", Some(binary)).await
243        }
244    }
245}
246
247pub async fn check_process_shell() -> Result<()> {
248    match get_permissions() {
249        None => Ok(()),
250        Some(p) => {
251            let check = p.check_process_shell();
252            handle_permission_check(check, "process.shell", None).await
253        }
254    }
255}
256
257pub async fn check_env_read(var: &str) -> Result<()> {
258    match get_permissions() {
259        None => Ok(()),
260        Some(p) => {
261            let check = p.check_env_read(var);
262            handle_permission_check(check, "env.read", Some(var)).await
263        }
264    }
265}
266
267pub async fn check_env_write() -> Result<()> {
268    match get_permissions() {
269        None => Ok(()),
270        Some(p) => {
271            let check = p.check_env_write();
272            handle_permission_check(check, "env.write", None).await
273        }
274    }
275}