1use anyhow::Result;
2use collections::HashMap;
3use gpui::{AsyncApp, SharedString, Task};
4use project::agent_server_store::*;
5use std::{any::Any, collections::HashSet, fmt::Write as _, path::PathBuf};
6// A simple fake that implements ExternalAgentServer without needing async plumbing.
7struct NoopExternalAgent;
8
9impl ExternalAgentServer for NoopExternalAgent {
10 fn get_command(
11 &mut self,
12 _extra_env: HashMap<String, String>,
13 _new_version_available_tx: Option<watch::Sender<Option<String>>>,
14 _cx: &mut AsyncApp,
15 ) -> Task<Result<AgentServerCommand>> {
16 Task::ready(Ok(AgentServerCommand {
17 path: PathBuf::from("noop"),
18 args: Vec::new(),
19 env: None,
20 }))
21 }
22
23 fn as_any_mut(&mut self) -> &mut dyn Any {
24 self
25 }
26}
27
28#[test]
29fn external_agent_server_name_display() {
30 let name = AgentId(SharedString::from("Ext: Tool"));
31 let mut s = String::new();
32 write!(&mut s, "{name}").unwrap();
33 assert_eq!(s, "Ext: Tool");
34}
35
36#[test]
37fn sync_extension_agents_removes_previous_extension_entries() {
38 let mut store = AgentServerStore::collab();
39
40 // Seed with a couple of agents that will be replaced by extensions
41 store.external_agents.insert(
42 AgentId(SharedString::from("foo-agent")),
43 ExternalAgentEntry::new(
44 Box::new(NoopExternalAgent) as Box<dyn ExternalAgentServer>,
45 ExternalAgentSource::Custom,
46 None,
47 None,
48 ),
49 );
50 store.external_agents.insert(
51 AgentId(SharedString::from("bar-agent")),
52 ExternalAgentEntry::new(
53 Box::new(NoopExternalAgent) as Box<dyn ExternalAgentServer>,
54 ExternalAgentSource::Custom,
55 None,
56 None,
57 ),
58 );
59 store.external_agents.insert(
60 AgentId(SharedString::from("custom")),
61 ExternalAgentEntry::new(
62 Box::new(NoopExternalAgent) as Box<dyn ExternalAgentServer>,
63 ExternalAgentSource::Custom,
64 None,
65 None,
66 ),
67 );
68
69 // Simulate the removal phase: if we're syncing extensions that provide
70 // "foo-agent" and "bar-agent", those should be removed first
71 let extension_agent_names: HashSet<String> = ["foo-agent".to_string(), "bar-agent".to_string()]
72 .into_iter()
73 .collect();
74
75 let keys_to_remove: Vec<_> = store
76 .external_agents
77 .keys()
78 .filter(|name| extension_agent_names.contains(name.0.as_ref()))
79 .cloned()
80 .collect();
81
82 for key in keys_to_remove {
83 store.external_agents.remove(&key);
84 }
85
86 // Only the custom entry should remain.
87 let remaining: Vec<_> = store
88 .external_agents
89 .keys()
90 .map(|k| k.0.to_string())
91 .collect();
92 assert_eq!(remaining, vec!["custom".to_string()]);
93}
94
95#[test]
96fn resolve_extension_icon_path_allows_valid_paths() {
97 // Create a temporary directory structure for testing
98 let temp_dir = tempfile::tempdir().unwrap();
99 let extensions_dir = temp_dir.path();
100 let ext_dir = extensions_dir.join("my-extension");
101 std::fs::create_dir_all(&ext_dir).unwrap();
102
103 // Create a valid icon file
104 let icon_path = ext_dir.join("icon.svg");
105 std::fs::write(&icon_path, "<svg></svg>").unwrap();
106
107 // Test that a valid relative path works
108 let result = project::agent_server_store::resolve_extension_icon_path(
109 extensions_dir,
110 "my-extension",
111 "icon.svg",
112 );
113 assert!(result.is_some());
114 assert!(result.unwrap().ends_with("icon.svg"));
115}
116
117#[test]
118fn resolve_extension_icon_path_allows_nested_paths() {
119 let temp_dir = tempfile::tempdir().unwrap();
120 let extensions_dir = temp_dir.path();
121 let ext_dir = extensions_dir.join("my-extension");
122 let icons_dir = ext_dir.join("assets").join("icons");
123 std::fs::create_dir_all(&icons_dir).unwrap();
124
125 let icon_path = icons_dir.join("logo.svg");
126 std::fs::write(&icon_path, "<svg></svg>").unwrap();
127
128 let result = project::agent_server_store::resolve_extension_icon_path(
129 extensions_dir,
130 "my-extension",
131 "assets/icons/logo.svg",
132 );
133 assert!(result.is_some());
134 assert!(result.unwrap().ends_with("logo.svg"));
135}
136
137#[test]
138fn resolve_extension_icon_path_blocks_path_traversal() {
139 let temp_dir = tempfile::tempdir().unwrap();
140 let extensions_dir = temp_dir.path();
141
142 // Create two extension directories
143 let ext1_dir = extensions_dir.join("extension1");
144 let ext2_dir = extensions_dir.join("extension2");
145 std::fs::create_dir_all(&ext1_dir).unwrap();
146 std::fs::create_dir_all(&ext2_dir).unwrap();
147
148 // Create a file in extension2
149 let secret_file = ext2_dir.join("secret.svg");
150 std::fs::write(&secret_file, "<svg>secret</svg>").unwrap();
151
152 // Try to access extension2's file from extension1 using path traversal
153 let result = project::agent_server_store::resolve_extension_icon_path(
154 extensions_dir,
155 "extension1",
156 "../extension2/secret.svg",
157 );
158 assert!(
159 result.is_none(),
160 "Path traversal to sibling extension should be blocked"
161 );
162}
163
164#[test]
165fn resolve_extension_icon_path_blocks_absolute_escape() {
166 let temp_dir = tempfile::tempdir().unwrap();
167 let extensions_dir = temp_dir.path();
168 let ext_dir = extensions_dir.join("my-extension");
169 std::fs::create_dir_all(&ext_dir).unwrap();
170
171 // Create a file outside the extensions directory
172 let outside_file = temp_dir.path().join("outside.svg");
173 std::fs::write(&outside_file, "<svg>outside</svg>").unwrap();
174
175 // Try to escape to parent directory
176 let result = project::agent_server_store::resolve_extension_icon_path(
177 extensions_dir,
178 "my-extension",
179 "../outside.svg",
180 );
181 assert!(
182 result.is_none(),
183 "Path traversal to parent directory should be blocked"
184 );
185}
186
187#[test]
188fn resolve_extension_icon_path_blocks_deep_traversal() {
189 let temp_dir = tempfile::tempdir().unwrap();
190 let extensions_dir = temp_dir.path();
191 let ext_dir = extensions_dir.join("my-extension");
192 std::fs::create_dir_all(&ext_dir).unwrap();
193
194 // Try deep path traversal
195 let result = project::agent_server_store::resolve_extension_icon_path(
196 extensions_dir,
197 "my-extension",
198 "../../../../../../etc/passwd",
199 );
200 assert!(
201 result.is_none(),
202 "Deep path traversal should be blocked (file doesn't exist)"
203 );
204}
205
206#[test]
207fn resolve_extension_icon_path_returns_none_for_nonexistent() {
208 let temp_dir = tempfile::tempdir().unwrap();
209 let extensions_dir = temp_dir.path();
210 let ext_dir = extensions_dir.join("my-extension");
211 std::fs::create_dir_all(&ext_dir).unwrap();
212
213 // Try to access a file that doesn't exist
214 let result = project::agent_server_store::resolve_extension_icon_path(
215 extensions_dir,
216 "my-extension",
217 "nonexistent.svg",
218 );
219 assert!(result.is_none(), "Nonexistent file should return None");
220}