Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: Defend against non-KCL files #798

Merged
merged 2 commits into from
Oct 25, 2023
Merged

fix: Defend against non-KCL files #798

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3c473ae
fix: Defend against panics caused by requests from non-KCL files. Whe…
He1pa Oct 23, 2023
3fe8feb
feat: verify request path valid at begin of handle request
He1pa Oct 25, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 36 additions & 5 deletions kclvm/tools/src/LSP/src/request.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,19 @@ impl LanguageServerState {
}

impl LanguageServerSnapshot {
// defend against non-kcl files
pub(crate) fn verify_request_path(&self, path: &VfsPath, sender: &Sender<Task>) -> bool {
let res = self.vfs.read().file_id(path).is_some()
&& self
.db
.get(&self.vfs.read().file_id(path).unwrap())
.is_some();
if !res {
let _ = log_message("Not a valid kcl path, request failed".to_string(), &sender);
}
res
}

pub(crate) fn get_db(&self, path: &VfsPath) -> anyhow::Result<&AnalysisDatabase> {
match self.vfs.read().file_id(path) {
Some(id) => match self.db.get(&id) {
Expand Down Expand Up @@ -128,7 +141,10 @@ pub(crate) fn handle_goto_definition(
) -> anyhow::Result<Option<lsp_types::GotoDefinitionResponse>> {
let file = file_path_from_url(&params.text_document_position_params.text_document.uri)?;
let path = from_lsp::abs_path(&params.text_document_position_params.text_document.uri)?;
let db = snapshot.get_db(&path.into())?;
if !snapshot.verify_request_path(&path.clone().into(), &sender) {
return Ok(None);
}
let db = snapshot.get_db(&path.clone().into())?;
let kcl_pos = kcl_pos(&file, params.text_document_position_params.position);
let res = goto_definition(&db.prog, &kcl_pos, &db.scope);
if res.is_none() {
Expand All @@ -145,6 +161,9 @@ pub(crate) fn handle_reference(
) -> anyhow::Result<Option<Vec<Location>>> {
let file = file_path_from_url(&params.text_document_position.text_document.uri)?;
let path = from_lsp::abs_path(&params.text_document_position.text_document.uri)?;
if !snapshot.verify_request_path(&path.clone().into(), &sender) {
return Ok(None);
}
let db = snapshot.get_db(&path.clone().into())?;
let pos = kcl_pos(&file, params.text_document_position.position);
let log = |msg: String| log_message(msg, &sender);
Expand Down Expand Up @@ -172,7 +191,10 @@ pub(crate) fn handle_completion(
) -> anyhow::Result<Option<lsp_types::CompletionResponse>> {
let file = file_path_from_url(&params.text_document_position.text_document.uri)?;
let path = from_lsp::abs_path(&params.text_document_position.text_document.uri)?;
let db = snapshot.get_db(&path.into())?;
if !snapshot.verify_request_path(&path.clone().into(), &sender) {
return Ok(None);
}
let db = snapshot.get_db(&path.clone().into())?;
let kcl_pos = kcl_pos(&file, params.text_document_position.position);
let completion_trigger_character = params
.context
Expand All @@ -193,7 +215,10 @@ pub(crate) fn handle_hover(
) -> anyhow::Result<Option<lsp_types::Hover>> {
let file = file_path_from_url(&params.text_document_position_params.text_document.uri)?;
let path = from_lsp::abs_path(&params.text_document_position_params.text_document.uri)?;
let db = snapshot.get_db(&path.into())?;
if !snapshot.verify_request_path(&path.clone().into(), &sender) {
return Ok(None);
}
let db = snapshot.get_db(&path.clone().into())?;
let kcl_pos = kcl_pos(&file, params.text_document_position_params.position);
let res = hover::hover(&db.prog, &kcl_pos, &db.scope);
if res.is_none() {
Expand All @@ -210,7 +235,10 @@ pub(crate) fn handle_document_symbol(
) -> anyhow::Result<Option<lsp_types::DocumentSymbolResponse>> {
let file = file_path_from_url(&params.text_document.uri)?;
let path = from_lsp::abs_path(&params.text_document.uri)?;
let db = snapshot.get_db(&path.into())?;
if !snapshot.verify_request_path(&path.clone().into(), &sender) {
return Ok(None);
}
let db = snapshot.get_db(&path.clone().into())?;
let res = document_symbol(&file, &db.prog, &db.scope);
if res.is_none() {
log_message(format!("File {file} Document symbol not found"), &sender)?;
Expand All @@ -233,7 +261,10 @@ pub(crate) fn handle_rename(
// 2. find all the references of the symbol
let file = file_path_from_url(&params.text_document_position.text_document.uri)?;
let path = from_lsp::abs_path(&params.text_document_position.text_document.uri)?;
let db = snapshot.get_db(&path.into())?;
if !snapshot.verify_request_path(&path.clone().into(), &sender) {
return Ok(None);
}
let db = snapshot.get_db(&path.clone().into())?;
let kcl_pos = kcl_pos(&file, params.text_document_position.position);
let log = |msg: String| log_message(msg, &sender);
let references = find_refs(
Expand Down
41 changes: 41 additions & 0 deletions kclvm/tools/src/LSP/src/tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ use lsp_types::CompletionParams;
use lsp_types::CompletionResponse;
use lsp_types::CompletionTriggerKind;
use lsp_types::DocumentFormattingParams;
use lsp_types::DocumentSymbolParams;
use lsp_types::GotoDefinitionParams;
use lsp_types::GotoDefinitionResponse;
use lsp_types::Hover;
Expand Down Expand Up @@ -630,6 +631,46 @@ fn close_file_test() {
);
}

#[test]
fn non_kcl_file_test() {
let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));

let server = Project {}.server(InitializeParams::default());
let mut path = root.clone();
path.push("src/test_data/diagnostics.kcl");

// Mock open a Non-KCL file
server.notification::<lsp_types::notification::DidOpenTextDocument>(
lsp_types::DidOpenTextDocumentParams {
text_document: TextDocumentItem {
uri: Url::from_file_path(path.clone()).unwrap(),
language_id: "KCL".to_string(),
version: 0,
text: "".to_string(),
},
},
);

let id = server.next_request_id.get();
server.next_request_id.set(id.wrapping_add(1));

let r: Request = Request::new(
id.into(),
"textDocument/documentSymbol".to_string(),
DocumentSymbolParams {
text_document: TextDocumentIdentifier {
uri: Url::from_file_path(path).unwrap(),
},
work_done_progress_params: Default::default(),
partial_result_params: Default::default(),
},
);

// Send request and wait for it's response
let res = server.send_and_receive(r);
assert!(res.result.is_some());
}

#[test]
fn goto_def_test() {
let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
Expand Down
Loading