1use std::sync::Arc;
2
3use anyhow::{Result, bail};
4use extension::{ExtensionCapability, ExtensionManifest};
5use url::Url;
6
7pub struct CapabilityGranter {
8 granted_capabilities: Vec<ExtensionCapability>,
9 manifest: Arc<ExtensionManifest>,
10}
11
12impl CapabilityGranter {
13 pub fn new(
14 granted_capabilities: Vec<ExtensionCapability>,
15 manifest: Arc<ExtensionManifest>,
16 ) -> Self {
17 Self {
18 granted_capabilities,
19 manifest,
20 }
21 }
22
23 pub fn grant_exec(
24 &self,
25 desired_command: &str,
26 desired_args: &[impl AsRef<str> + std::fmt::Debug],
27 ) -> Result<()> {
28 self.manifest.allow_exec(desired_command, desired_args)?;
29
30 let is_allowed = self
31 .granted_capabilities
32 .iter()
33 .any(|capability| match capability {
34 ExtensionCapability::ProcessExec(capability) => {
35 capability.allows(desired_command, desired_args)
36 }
37 _ => false,
38 });
39
40 if !is_allowed {
41 bail!(
42 "capability for process:exec {desired_command} {desired_args:?} is not granted by the extension host",
43 );
44 }
45
46 Ok(())
47 }
48
49 pub fn grant_download_file(&self, desired_url: &Url) -> Result<()> {
50 let is_allowed = self
51 .granted_capabilities
52 .iter()
53 .any(|capability| match capability {
54 ExtensionCapability::DownloadFile(capability) => capability.allows(desired_url),
55 _ => false,
56 });
57
58 if !is_allowed {
59 bail!(
60 "capability for download_file {desired_url} is not granted by the extension host",
61 );
62 }
63
64 Ok(())
65 }
66
67 pub fn grant_npm_install_package(&self, package_name: &str) -> Result<()> {
68 let is_allowed = self
69 .granted_capabilities
70 .iter()
71 .any(|capability| match capability {
72 ExtensionCapability::NpmInstallPackage(capability) => {
73 capability.allows(package_name)
74 }
75 _ => false,
76 });
77
78 if !is_allowed {
79 bail!("capability for npm:install {package_name} is not granted by the extension host",);
80 }
81
82 Ok(())
83 }
84}
85
86#[cfg(test)]
87mod tests {
88 use std::collections::BTreeMap;
89
90 use extension::{ProcessExecCapability, SchemaVersion};
91
92 use super::*;
93
94 fn extension_manifest() -> ExtensionManifest {
95 ExtensionManifest {
96 id: "test".into(),
97 name: "Test".to_string(),
98 version: "1.0.0".into(),
99 schema_version: SchemaVersion::ZERO,
100 description: None,
101 repository: None,
102 authors: vec![],
103 lib: Default::default(),
104 themes: vec![],
105 icon_themes: vec![],
106 languages: vec![],
107 grammars: BTreeMap::default(),
108 language_servers: BTreeMap::default(),
109 context_servers: BTreeMap::default(),
110 agent_servers: BTreeMap::default(),
111 slash_commands: BTreeMap::default(),
112 snippets: None,
113 capabilities: vec![],
114 debug_adapters: Default::default(),
115 debug_locators: Default::default(),
116 }
117 }
118
119 #[test]
120 fn test_grant_exec() {
121 let manifest = Arc::new(ExtensionManifest {
122 capabilities: vec![ExtensionCapability::ProcessExec(ProcessExecCapability {
123 command: "ls".to_string(),
124 args: vec!["-la".to_string()],
125 })],
126 ..extension_manifest()
127 });
128
129 // It returns an error when the extension host has no granted capabilities.
130 let granter = CapabilityGranter::new(Vec::new(), manifest.clone());
131 assert!(granter.grant_exec("ls", &["-la"]).is_err());
132
133 // It succeeds when the extension host has the exact capability.
134 let granter = CapabilityGranter::new(
135 vec![ExtensionCapability::ProcessExec(ProcessExecCapability {
136 command: "ls".to_string(),
137 args: vec!["-la".to_string()],
138 })],
139 manifest.clone(),
140 );
141 assert!(granter.grant_exec("ls", &["-la"]).is_ok());
142
143 // It succeeds when the extension host has a wildcard capability.
144 let granter = CapabilityGranter::new(
145 vec![ExtensionCapability::ProcessExec(ProcessExecCapability {
146 command: "*".to_string(),
147 args: vec!["**".to_string()],
148 })],
149 manifest,
150 );
151 assert!(granter.grant_exec("ls", &["-la"]).is_ok());
152 }
153}