|
| 1 | +//! MCP resource URI codec and subscription tracking for LSP diagnostics. |
| 2 | +//! |
| 3 | +//! Resources in mcpls use the `lsp-diagnostics:///` scheme (RFC 3986 compliant, |
| 4 | +//! empty authority, percent-encoded path). Each resource corresponds to a single |
| 5 | +//! file whose diagnostics are cached from LSP `textDocument/publishDiagnostics` |
| 6 | +//! notifications. |
| 7 | +
|
| 8 | +use std::collections::HashSet; |
| 9 | +use std::path::{Path, PathBuf}; |
| 10 | + |
| 11 | +use thiserror::Error; |
| 12 | +use tokio::sync::RwLock; |
| 13 | +use url::Url; |
| 14 | + |
| 15 | +/// URI scheme used for diagnostic resources. |
| 16 | +const SCHEME: &str = "lsp-diagnostics"; |
| 17 | + |
| 18 | +/// Full scheme + authority prefix (`scheme://`). |
| 19 | +/// |
| 20 | +/// Three-slash form (`lsp-diagnostics:///`) is produced by appending an empty |
| 21 | +/// authority and the absolute path: `{PREFIX}{path}`. |
| 22 | +const PREFIX: &str = "lsp-diagnostics://"; |
| 23 | + |
| 24 | +/// Maximum number of resource URIs a single client session may subscribe to. |
| 25 | +/// |
| 26 | +/// Guards against memory exhaustion from a misbehaving or adversarial client. |
| 27 | +pub const MAX_SUBSCRIPTIONS: usize = 1_000; |
| 28 | + |
| 29 | +/// Errors produced by the resource URI codec. |
| 30 | +#[derive(Debug, Error)] |
| 31 | +pub enum ResourceUriError { |
| 32 | + /// The path is relative or contains non-UTF-8 components. |
| 33 | + #[error("path must be absolute and valid UTF-8: {0}")] |
| 34 | + InvalidPath(String), |
| 35 | + |
| 36 | + /// The URI has the wrong scheme or malformed structure. |
| 37 | + #[error("expected '{SCHEME}:///' prefix in URI: {0}")] |
| 38 | + InvalidScheme(String), |
| 39 | + |
| 40 | + /// The URI path could not be decoded to a filesystem path. |
| 41 | + #[error("failed to decode URI to filesystem path: {0}")] |
| 42 | + DecodeFailed(String), |
| 43 | +} |
| 44 | + |
| 45 | +/// Encode an absolute filesystem path into a `lsp-diagnostics:///…` resource URI. |
| 46 | +/// |
| 47 | +/// Percent-encoding is delegated to [`url::Url::from_file_path`], which |
| 48 | +/// handles spaces, unicode, `%`, `?`, `#`, and platform separators correctly. |
| 49 | +/// |
| 50 | +/// # Errors |
| 51 | +/// |
| 52 | +/// Returns [`ResourceUriError::InvalidPath`] if the path is relative or |
| 53 | +/// cannot be expressed as a valid file URI. |
| 54 | +/// |
| 55 | +/// # Examples |
| 56 | +/// |
| 57 | +/// ``` |
| 58 | +/// use std::path::Path; |
| 59 | +/// use mcpls_core::bridge::resources::make_uri; |
| 60 | +/// |
| 61 | +/// let uri = make_uri(Path::new("/home/user/main.rs")).unwrap(); |
| 62 | +/// assert!(uri.starts_with("lsp-diagnostics:///")); |
| 63 | +/// ``` |
| 64 | +pub fn make_uri(path: &Path) -> Result<String, ResourceUriError> { |
| 65 | + let file_url = Url::from_file_path(path) |
| 66 | + .map_err(|()| ResourceUriError::InvalidPath(path.display().to_string()))?; |
| 67 | + |
| 68 | + // Replace the "file" scheme with our custom scheme while keeping the |
| 69 | + // already-percent-encoded path and authority (empty) components. |
| 70 | + let uri = format!("{SCHEME}://{}", &file_url[url::Position::BeforeHost..]); |
| 71 | + Ok(uri) |
| 72 | +} |
| 73 | + |
| 74 | +/// Decode a `lsp-diagnostics:///…` resource URI back to an absolute filesystem path. |
| 75 | +/// |
| 76 | +/// # Errors |
| 77 | +/// |
| 78 | +/// Returns an error if the URI does not start with the expected scheme, |
| 79 | +/// or if the percent-encoded path cannot be mapped to a filesystem path. |
| 80 | +/// |
| 81 | +/// # Examples |
| 82 | +/// |
| 83 | +/// ``` |
| 84 | +/// use std::path::Path; |
| 85 | +/// use mcpls_core::bridge::resources::{make_uri, parse_uri}; |
| 86 | +/// |
| 87 | +/// let path = Path::new("/home/user/main.rs"); |
| 88 | +/// let uri = make_uri(path).unwrap(); |
| 89 | +/// let recovered = parse_uri(&uri).unwrap(); |
| 90 | +/// assert_eq!(recovered, path); |
| 91 | +/// ``` |
| 92 | +pub fn parse_uri(uri: &str) -> Result<PathBuf, ResourceUriError> { |
| 93 | + if !uri.starts_with(PREFIX) { |
| 94 | + return Err(ResourceUriError::InvalidScheme(uri.to_string())); |
| 95 | + } |
| 96 | + |
| 97 | + // Require empty authority: the character immediately after `://` must be `/`. |
| 98 | + // This blocks `lsp-diagnostics://evil-host/path` → UNC path on Windows. |
| 99 | + let after_prefix = &uri[PREFIX.len()..]; |
| 100 | + if !after_prefix.starts_with('/') { |
| 101 | + return Err(ResourceUriError::InvalidScheme(format!( |
| 102 | + "non-empty authority in URI: {uri}" |
| 103 | + ))); |
| 104 | + } |
| 105 | + |
| 106 | + let file_uri = format!("file://{after_prefix}"); |
| 107 | + let url = Url::parse(&file_uri).map_err(|e| ResourceUriError::DecodeFailed(e.to_string()))?; |
| 108 | + |
| 109 | + url.to_file_path() |
| 110 | + .map_err(|()| ResourceUriError::DecodeFailed(file_uri)) |
| 111 | +} |
| 112 | + |
| 113 | +/// Tracks which MCP resource URIs the client has subscribed to. |
| 114 | +/// |
| 115 | +/// The hot read path (pump tasks checking before sending notifications) uses |
| 116 | +/// a `RwLock` so concurrent readers do not block each other. |
| 117 | +#[derive(Debug)] |
| 118 | +pub struct ResourceSubscriptions(RwLock<HashSet<String>>); |
| 119 | + |
| 120 | +impl Default for ResourceSubscriptions { |
| 121 | + fn default() -> Self { |
| 122 | + Self::new() |
| 123 | + } |
| 124 | +} |
| 125 | + |
| 126 | +impl ResourceSubscriptions { |
| 127 | + /// Create an empty subscription set. |
| 128 | + #[must_use] |
| 129 | + pub fn new() -> Self { |
| 130 | + Self(RwLock::new(HashSet::new())) |
| 131 | + } |
| 132 | + |
| 133 | + /// Add a URI to the subscription set. |
| 134 | + /// |
| 135 | + /// Returns `Ok(true)` if newly inserted, `Ok(false)` if already present. |
| 136 | + /// Returns `Err` if the subscription set has reached [`MAX_SUBSCRIPTIONS`]. |
| 137 | + /// |
| 138 | + /// # Errors |
| 139 | + /// |
| 140 | + /// Returns an error string when the cap is exceeded. |
| 141 | + pub async fn subscribe(&self, uri: String) -> Result<bool, String> { |
| 142 | + let mut set = self.0.write().await; |
| 143 | + if !set.contains(&uri) && set.len() >= MAX_SUBSCRIPTIONS { |
| 144 | + return Err(format!("subscription limit of {MAX_SUBSCRIPTIONS} reached")); |
| 145 | + } |
| 146 | + Ok(set.insert(uri)) |
| 147 | + } |
| 148 | + |
| 149 | + /// Check whether the subscription set is empty. |
| 150 | + /// |
| 151 | + /// Used as a fast path in the diagnostics pump to skip URI construction |
| 152 | + /// when no client has subscribed yet. |
| 153 | + pub async fn is_empty(&self) -> bool { |
| 154 | + self.0.read().await.is_empty() |
| 155 | + } |
| 156 | + |
| 157 | + /// Remove a URI from the subscription set. |
| 158 | + /// |
| 159 | + /// Returns `true` if the URI was present and removed. |
| 160 | + pub async fn unsubscribe(&self, uri: &str) -> bool { |
| 161 | + self.0.write().await.remove(uri) |
| 162 | + } |
| 163 | + |
| 164 | + /// Check if a URI is currently subscribed. |
| 165 | + pub async fn contains(&self, uri: &str) -> bool { |
| 166 | + self.0.read().await.contains(uri) |
| 167 | + } |
| 168 | + |
| 169 | + /// Return a snapshot of all subscribed URIs (primarily for tests). |
| 170 | + pub async fn snapshot(&self) -> Vec<String> { |
| 171 | + self.0.read().await.iter().cloned().collect() |
| 172 | + } |
| 173 | +} |
| 174 | + |
| 175 | +#[cfg(test)] |
| 176 | +#[allow(clippy::unwrap_used, clippy::expect_used)] |
| 177 | +mod tests { |
| 178 | + use super::*; |
| 179 | + |
| 180 | + // ------------------------------------------------------------------ |
| 181 | + // URI codec |
| 182 | + // ------------------------------------------------------------------ |
| 183 | + |
| 184 | + #[test] |
| 185 | + fn test_make_uri_rejects_relative_path() { |
| 186 | + let result = make_uri(Path::new("relative/path.rs")); |
| 187 | + assert!(result.is_err()); |
| 188 | + } |
| 189 | + |
| 190 | + #[test] |
| 191 | + fn test_parse_uri_rejects_wrong_scheme() { |
| 192 | + let result = parse_uri("file:///home/user/main.rs"); |
| 193 | + assert!(result.is_err()); |
| 194 | + } |
| 195 | + |
| 196 | + #[test] |
| 197 | + fn test_parse_uri_rejects_http_scheme() { |
| 198 | + let result = parse_uri("https://example.com/file.rs"); |
| 199 | + assert!(result.is_err()); |
| 200 | + } |
| 201 | + |
| 202 | + #[cfg(unix)] |
| 203 | + #[test] |
| 204 | + fn test_make_uri_simple_path() { |
| 205 | + let uri = make_uri(Path::new("/home/user/main.rs")).unwrap(); |
| 206 | + assert_eq!(uri, "lsp-diagnostics:///home/user/main.rs"); |
| 207 | + } |
| 208 | + |
| 209 | + #[cfg(unix)] |
| 210 | + #[test] |
| 211 | + fn test_make_uri_scheme_prefix() { |
| 212 | + let uri = make_uri(Path::new("/tmp/file.rs")).unwrap(); |
| 213 | + assert!(uri.starts_with("lsp-diagnostics:///")); |
| 214 | + } |
| 215 | + |
| 216 | + #[cfg(unix)] |
| 217 | + #[test] |
| 218 | + fn test_parse_uri_simple() { |
| 219 | + let path = PathBuf::from("/home/user/main.rs"); |
| 220 | + let uri = make_uri(&path).unwrap(); |
| 221 | + let recovered = parse_uri(&uri).unwrap(); |
| 222 | + assert_eq!(recovered, path); |
| 223 | + } |
| 224 | + |
| 225 | + /// Round-trip: paths with spaces, unicode, `%`, `?`, `#`. |
| 226 | + #[cfg(unix)] |
| 227 | + #[test] |
| 228 | + fn test_round_trip_special_chars() { |
| 229 | + let paths = [ |
| 230 | + "/home/user/my file.rs", |
| 231 | + "/tmp/café/main.rs", |
| 232 | + "/data/100%/test.rs", |
| 233 | + "/workspace/query?param/file.rs", |
| 234 | + "/repo/branch#fragment/src.rs", |
| 235 | + "/путь/к/файлу.rs", |
| 236 | + ]; |
| 237 | + |
| 238 | + for raw in &paths { |
| 239 | + let path = PathBuf::from(raw); |
| 240 | + let uri = make_uri(&path).expect(raw); |
| 241 | + assert!( |
| 242 | + uri.starts_with("lsp-diagnostics:///"), |
| 243 | + "URI should start with correct scheme: {uri}" |
| 244 | + ); |
| 245 | + let recovered = parse_uri(&uri).expect(&uri); |
| 246 | + assert_eq!(recovered, path, "Round-trip failed for: {raw}"); |
| 247 | + } |
| 248 | + } |
| 249 | + |
| 250 | + /// Snapshot test: verify the on-wire form uses three slashes and percent-encoding. |
| 251 | + #[cfg(unix)] |
| 252 | + #[test] |
| 253 | + fn test_wire_format_percent_encoded() { |
| 254 | + let path = Path::new("/home/user/my file.rs"); |
| 255 | + let uri = make_uri(path).unwrap(); |
| 256 | + // Space must be percent-encoded as %20 |
| 257 | + assert!(uri.contains("%20"), "Expected %20 in: {uri}"); |
| 258 | + assert!(uri.starts_with("lsp-diagnostics:///")); |
| 259 | + } |
| 260 | + |
| 261 | + // ------------------------------------------------------------------ |
| 262 | + // ResourceSubscriptions |
| 263 | + // ------------------------------------------------------------------ |
| 264 | + |
| 265 | + #[tokio::test] |
| 266 | + async fn test_subscribe_and_contains() { |
| 267 | + let subs = ResourceSubscriptions::new(); |
| 268 | + let uri = "lsp-diagnostics:///home/user/main.rs".to_string(); |
| 269 | + |
| 270 | + assert!(!subs.contains(&uri).await); |
| 271 | + assert!(subs.subscribe(uri.clone()).await.unwrap()); |
| 272 | + assert!(subs.contains(&uri).await); |
| 273 | + } |
| 274 | + |
| 275 | + #[tokio::test] |
| 276 | + async fn test_subscribe_duplicate_returns_false() { |
| 277 | + let subs = ResourceSubscriptions::new(); |
| 278 | + let uri = "lsp-diagnostics:///tmp/file.rs".to_string(); |
| 279 | + assert!(subs.subscribe(uri.clone()).await.unwrap()); |
| 280 | + assert!(!subs.subscribe(uri).await.unwrap()); |
| 281 | + } |
| 282 | + |
| 283 | + #[tokio::test] |
| 284 | + async fn test_unsubscribe() { |
| 285 | + let subs = ResourceSubscriptions::new(); |
| 286 | + let uri = "lsp-diagnostics:///tmp/file.rs".to_string(); |
| 287 | + subs.subscribe(uri.clone()).await.unwrap(); |
| 288 | + assert!(subs.unsubscribe(&uri).await); |
| 289 | + assert!(!subs.contains(&uri).await); |
| 290 | + } |
| 291 | + |
| 292 | + #[tokio::test] |
| 293 | + async fn test_unsubscribe_nonexistent_returns_false() { |
| 294 | + let subs = ResourceSubscriptions::new(); |
| 295 | + assert!(!subs.unsubscribe("lsp-diagnostics:///nonexistent.rs").await); |
| 296 | + } |
| 297 | + |
| 298 | + #[tokio::test] |
| 299 | + async fn test_subscribe_cap_exceeded() { |
| 300 | + let subs = ResourceSubscriptions::new(); |
| 301 | + for i in 0..MAX_SUBSCRIPTIONS { |
| 302 | + subs.subscribe(format!("lsp-diagnostics:///file{i}.rs")) |
| 303 | + .await |
| 304 | + .unwrap(); |
| 305 | + } |
| 306 | + let result = subs |
| 307 | + .subscribe("lsp-diagnostics:///overflow.rs".to_string()) |
| 308 | + .await; |
| 309 | + assert!(result.is_err()); |
| 310 | + } |
| 311 | + |
| 312 | + #[tokio::test] |
| 313 | + async fn test_snapshot() { |
| 314 | + let subs = ResourceSubscriptions::new(); |
| 315 | + subs.subscribe("lsp-diagnostics:///a.rs".to_string()) |
| 316 | + .await |
| 317 | + .unwrap(); |
| 318 | + subs.subscribe("lsp-diagnostics:///b.rs".to_string()) |
| 319 | + .await |
| 320 | + .unwrap(); |
| 321 | + let mut snap = subs.snapshot().await; |
| 322 | + snap.sort(); |
| 323 | + assert_eq!(snap, ["lsp-diagnostics:///a.rs", "lsp-diagnostics:///b.rs"]); |
| 324 | + } |
| 325 | +} |
0 commit comments