feat(ls): configurable workspace files caching (VirusTotal#568)

Now configuration is stored in the state of the language server. The
configuration is resolved with `workspace.didChangeConfiguration` LSP
method. VS Code always sends all `YARA.*` configuration attributes after
LSP communication is initialized and when the configuration changes
using this LSP method. This is most cheaper than calling `workspace.configuration`
for specific language feature configuration every time.

Additionally, the language server now caches the entire workspace by storing a 
CST for each YARA file. This still can be disabled in YARA LS configuration in the
editor settings. For maintaining consistency between cache and FS state
it is leveraging LSP method `workspace/didChangeWatchedFiles` to watch
YARA files that were not opened in the editor but were changed (e.g. git
commands or any other manipulation from the terminal...).

Author: prosperritty

Cargo.lock

@@ -57,6 +57,7 @@ walkdir = "2.5.0"
 [target.'cfg(any(target_arch = "wasm32", target_arch = "wasm64"))'.dependencies]
 async-lsp = {version = "0.2.2", default-features=false, features=["omni-trait"]}
 tower = { version = "0.5.2" }
+wasm-bindgen-futures = {version = "0.4.61"}
 
 [dev-dependencies]
 goldenfile = { workspace = true }

ls/Cargo.toml

@@ -104,6 +104,11 @@
           "type": "string",
           "default": "",
           "description": "A regular expression that all rule names must match (example: APT_.*)."
+        },
+        "YARA.cacheWorkspace": {
+          "type": "boolean",
+          "default": true,
+          "description": "Specifies if the language server should cache all files from the workspace."
         }
       }
     },

ls/editors/code/package.json

@@ -1,4 +1,4 @@
-import { ExtensionContext, window } from "vscode";
+import { ExtensionContext, window, workspace } from "vscode";
 import {
   Executable,
   LanguageClient,
@@ -11,6 +11,7 @@ import * as path from "path";
 let client: LanguageClient | null = null;
 
 export async function activate(context: ExtensionContext) {
+  const config = workspace.getConfiguration("YARA");
   const platform = os.platform();
   const arch = os.arch();
 
@@ -39,6 +40,7 @@ export async function activate(context: ExtensionContext) {
     documentSelector: [{ scheme: "file", language: "yara" }],
     outputChannel: outputChannel,
     traceOutputChannel: outputChannel,
+    initializationOptions: config
   };
 
   client = new LanguageClient(

ls/editors/code/src/extension.ts

@@ -0,0 +1,53 @@
+use serde::Deserialize;
+
+/// This structure contains all client-side configuration settings,
+/// which user can specify in the code editor.
+#[derive(Deserialize, Default, Debug)]
+#[serde(rename_all = "camelCase")]
+pub struct Config {
+    pub code_formatting: FormattingConfiguration,
+    pub metadata_validation: Vec<MetadataValidationRule>,
+    pub rule_name_validation: Option<String>,
+    pub cache_workspace: bool,
+}
+
+/// This structure represents settings for the YARA-X formatter.
+#[derive(Clone, Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub(crate) struct FormattingConfiguration {
+    pub align_metadata: bool,
+    pub align_patterns: bool,
+    pub indent_section_headers: bool,
+    pub indent_section_contents: bool,
+    pub newline_before_curly_brace: bool,
+    pub empty_line_before_section_header: bool,
+    pub empty_line_after_section_header: bool,
+}
+
+impl Default for FormattingConfiguration {
+    fn default() -> Self {
+        Self {
+            align_metadata: true,
+            align_patterns: true,
+            indent_section_headers: true,
+            indent_section_contents: true,
+            newline_before_curly_brace: false,
+            empty_line_before_section_header: false,
+            empty_line_after_section_header: false,
+        }
+    }
+}
+
+/// Rule that describes a how to validate a metadata entry in a rule.
+#[derive(Deserialize, Debug, Clone, Default)]
+#[serde(rename_all = "camelCase")]
+pub struct MetadataValidationRule {
+    /// Metadata identifier
+    pub identifier: String,
+    /// Whether the metadata entry is required or not.
+    #[serde(default)]
+    pub required: bool,
+    /// Type of the metadata entry.
+    #[serde(rename = "type")]
+    pub ty: Option<String>,
+}

ls/src/configuration.rs

@@ -24,6 +24,12 @@ impl Document {
         Self { uri, text, cst, line_index }
     }
 
+    /// Creates a new document with precached CST.
+    pub fn new_with_cst(uri: Url, text: String, cst: CST) -> Self {
+        let line_index = LineIndex::new(text.as_str());
+        Self { uri, text, cst, line_index }
+    }
+
     /// Updates all stored structures.
     pub fn update(&mut self, text: String) {
         self.cst = CST::from(text.as_str());

ls/src/documents/document.rs

@@ -3,7 +3,7 @@ use std::collections::{HashMap, HashSet};
 #[cfg(not(any(target_arch = "wasm32", target_arch = "wasm64")))]
 use std::fs;
 
-use async_lsp::lsp_types::Url;
+use async_lsp::lsp_types::{FileChangeType, FileEvent, Url};
 use dashmap::{mapref::one::Ref, DashMap};
 use yara_x_parser::cst::{Immutable, Node, SyntaxKind, Token, CST};
 
@@ -23,6 +23,7 @@ pub struct OccurrencesResult {
 #[derive(Default)]
 pub struct DocumentStorage {
     opened: DashMap<Url, Document>,
+    cached: DashMap<Url, CST>,
     workspace: Option<Url>,
 }
 
@@ -38,7 +39,12 @@ impl DocumentStorage {
 
     /// Inserts a new document with the specified content.
     pub fn insert(&self, uri: Url, text: String) {
-        self.opened.insert(uri.clone(), Document::new(uri, text));
+        // Checks if the opened document is already in cache and remove it.
+        let document = match self.cached.remove(&uri) {
+            Some((_, cst)) => Document::new_with_cst(uri.clone(), text, cst),
+            None => Document::new(uri.clone(), text),
+        };
+        self.opened.insert(uri.clone(), document);
     }
 
     /// Updates document content and its internal structures.
@@ -49,8 +55,15 @@ impl DocumentStorage {
     }
 
     /// Removes the document and returns (key,value) before removing.
-    pub fn remove(&self, uri: &Url) -> Option<(Url, Document)> {
-        self.opened.remove(uri)
+    pub fn remove(&self, uri: &Url, cache_enabled: bool) {
+        if let Some((uri, document)) = self.opened.remove(uri) {
+            // If the workspace caching is enabled we want to
+            // move CST of the closed document to the cache.
+            if cache_enabled {
+                let cst = document.cst;
+                self.cached.insert(uri, cst);
+            }
+        }
     }
 
     /// Sets workspace folder to the specified URI.
@@ -64,6 +77,8 @@ impl DocumentStorage {
     fn get_document_cst_root(&self, uri: &Url) -> Option<Node<Immutable>> {
         if let Some(doc) = self.get(uri) {
             Some(doc.cst.root())
+        } else if let Some(cst) = self.cached.get(uri) {
+            Some(cst.root())
         } else {
             uri.to_file_path().ok().and_then(|path| {
                 fs::read_to_string(path)
@@ -321,4 +336,70 @@ impl DocumentStorage {
 
         rules
     }
+
+    /// Reads all files from the workspace and initializes cache iwth CST
+    /// for each YARA file.
+    #[cfg(not(any(target_arch = "wasm32", target_arch = "wasm64")))]
+    pub fn cache_workspace(&self) {
+        if let Some(workspace_files) = self
+            .walk_workspace()
+            .map(|workspace| workspace.collect::<Vec<Url>>())
+        {
+            for entry in workspace_files {
+                if self.opened.contains_key(&entry) {
+                    continue;
+                }
+
+                if let Some(cst) = entry
+                    .to_file_path()
+                    .ok()
+                    .and_then(|path| fs::read_to_string(path).ok())
+                    .map(|content| CST::from(content.as_str()))
+                {
+                    self.cached.insert(entry, cst);
+                }
+            }
+        }
+    }
+
+    #[cfg(any(target_arch = "wasm32", target_arch = "wasm64"))]
+    pub fn cache_workspace(&self) {}
+
+    /// Clears the cache.
+    pub fn clear_cache(&self) {
+        self.cached.clear();
+    }
+
+    /// Reacts to the file system changes within the workspace by making
+    /// changes in the cache.
+    #[cfg(not(any(target_arch = "wasm32", target_arch = "wasm64")))]
+    pub fn react_watched_files_changes(&self, changes: Vec<FileEvent>) {
+        for change in changes {
+            // Opened files are synchronized in `textDocument/did*` methods.
+            if self.opened.contains_key(&change.uri) {
+                continue;
+            }
+
+            match change.typ {
+                FileChangeType::CHANGED | FileChangeType::CREATED => {
+                    if let Some(cst) =
+                        change.uri.to_file_path().ok().and_then(|path| {
+                            fs::read_to_string(path)
+                                .ok()
+                                .map(|content| CST::from(content.as_str()))
+                        })
+                    {
+                        self.cached.insert(change.uri, cst);
+                    }
+                }
+                FileChangeType::DELETED => {
+                    self.cached.remove(&change.uri);
+                }
+                _ => {}
+            }
+        }
+    }
+
+    #[cfg(any(target_arch = "wasm32", target_arch = "wasm64"))]
+    pub fn react_watched_files_changes(&self, changes: Vec<FileEvent>) {}
 }