addressed Issue #17

added logic to look at the previous char to see if symbol is a variable did some extra regex to filter out variables that don't match
infosec-intern · Feb 14, 2018 · 5d29c34 · 5d29c34
1 parent f2d0e58
commit 5d29c34
Showing 1 changed file with 32 additions and 13 deletions.
diff --git a/yara/src/extension.ts b/yara/src/extension.ts
@@ -39,22 +39,41 @@ export class YaraDefinitionProvider implements vscode.DefinitionProvider {
 export class YaraReferenceProvider implements vscode.ReferenceProvider {
     public provideReferences(doc: vscode.TextDocument, pos: vscode.Position, options: { includeDeclaration: boolean }, token: vscode.CancellationToken): Thenable<vscode.Location[]> {
         return new Promise((resolve, reject) => {
-            let references: vscode.Location[] = new Array<vscode.Location>();
+            const varFirstChar = new Set(["$", "#", "@", "!"]);
             const fileUri: vscode.Uri = vscode.Uri.file(doc.fileName);
             const range: vscode.Range = doc.getWordRangeAtPosition(pos);
-            const symbol: string = doc.getText(range);
-            // console.log(`Providing references for symbol '${symbol}'`);
             let lines: string[] = doc.getText().split("\n");
-            let lineNo = 0;
-            lines.forEach(line => {
-                let character: number = line.indexOf(symbol);
-                if (character != -1) {
-                    // console.log(`Found ${symbol} on line ${lineNo} at character ${character}`);
-                    let refPosition: vscode.Position = new vscode.Position(lineNo, character);
-                    references.push(new vscode.Location(fileUri, refPosition));
-                }
-                lineNo++;
-            });
+            let references: vscode.Location[] = new Array<vscode.Location>();
+            let symbol: string = doc.getText(range);
+            // console.log(`Providing references for symbol '${symbol}'`);
+            let possibleVarStart: vscode.Position = new vscode.Position(range.start.line, range.start.character-1);
+            let possibleVarRange: vscode.Range = new vscode.Range(possibleVarStart, range.end);
+            let possibleVar: string = doc.getText(possibleVarRange);
+            if (varFirstChar.has(possibleVar.charAt(0))) {
+                // console.log(`Identified symbol as a variable: ${symbol}`);
+                let lineNo = 0;
+                lines.forEach(line => {
+                    let character: number = line.search(`[\$#@!]${symbol}[^a-zA-Z0-9_]`);
+                    if (character != -1) {
+                        // console.log(`Found ${symbol} on line ${lineNo} at character ${character}`);
+                        let refPosition: vscode.Position = new vscode.Position(lineNo, character);
+                        references.push(new vscode.Location(fileUri, refPosition));
+                    }
+                    lineNo++;
+                });
+            }
+            else {
+                let lineNo = 0;
+                lines.forEach(line => {
+                    let character: number = line.indexOf(symbol);
+                    if (character != -1) {
+                        // console.log(`Found ${symbol} on line ${lineNo} at character ${character}`);
+                        let refPosition: vscode.Position = new vscode.Position(lineNo, character);
+                        references.push(new vscode.Location(fileUri, refPosition));
+                    }
+                    lineNo++;
+                });
+            }
             if (references != null) {
                 resolve(references);
             }