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 const 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 slash_commands: BTreeMap::default(),
111 snippets: None,
112 capabilities: vec![],
113 debug_adapters: Default::default(),
114 debug_locators: Default::default(),
115 }
116 }
117
118 #[test]
119 fn test_grant_exec() {
120 let manifest = Arc::new(ExtensionManifest {
121 capabilities: vec![ExtensionCapability::ProcessExec(ProcessExecCapability {
122 command: "ls".to_string(),
123 args: vec!["-la".to_string()],
124 })],
125 ..extension_manifest()
126 });
127
128 // It returns an error when the extension host has no granted capabilities.
129 let granter = CapabilityGranter::new(Vec::new(), manifest.clone());
130 assert!(granter.grant_exec("ls", &["-la"]).is_err());
131
132 // It succeeds when the extension host has the exact capability.
133 let granter = CapabilityGranter::new(
134 vec![ExtensionCapability::ProcessExec(ProcessExecCapability {
135 command: "ls".to_string(),
136 args: vec!["-la".to_string()],
137 })],
138 manifest.clone(),
139 );
140 assert!(granter.grant_exec("ls", &["-la"]).is_ok());
141
142 // It succeeds when the extension host has a wildcard capability.
143 let granter = CapabilityGranter::new(
144 vec![ExtensionCapability::ProcessExec(ProcessExecCapability {
145 command: "*".to_string(),
146 args: vec!["**".to_string()],
147 })],
148 manifest,
149 );
150 assert!(granter.grant_exec("ls", &["-la"]).is_ok());
151 }
152}