|
| 1 | +import * as fs from 'fs-extra'; |
| 2 | +import * as yaml from 'js-yaml'; |
| 3 | +import * as tmp from 'tmp'; |
| 4 | +import * as vscode from "vscode"; |
| 5 | +import { decodeSourceArchiveUri, zipArchiveScheme } from "./archive-filesystem-provider"; |
| 6 | +import { ColumnKindCode, EntityValue, getResultSetSchema, LineColumnLocation, UrlValue } from "./bqrs-cli-types"; |
| 7 | +import { CodeQLCliServer } from "./cli"; |
| 8 | +import { DatabaseItem, DatabaseManager } from "./databases"; |
| 9 | +import * as helpers from './helpers'; |
| 10 | +import { CachedOperation } from './helpers'; |
| 11 | +import * as messages from "./messages"; |
| 12 | +import { QueryServerClient } from "./queryserver-client"; |
| 13 | +import { compileAndRunQueryAgainstDatabase, QueryWithResults } from "./run-queries"; |
| 14 | + |
| 15 | +/** |
| 16 | + * Run templated CodeQL queries to find definitions and references in |
| 17 | + * source-language files. We may eventually want to find a way to |
| 18 | + * generalize this to other custom queries, e.g. showing dataflow to |
| 19 | + * or from a selected identifier. |
| 20 | + */ |
| 21 | + |
| 22 | +const TEMPLATE_NAME = "selectedSourceFile"; |
| 23 | +const SELECT_QUERY_NAME = "#select"; |
| 24 | + |
| 25 | +enum KeyType { |
| 26 | + DefinitionQuery = 'DefinitionQuery', |
| 27 | + ReferenceQuery = 'ReferenceQuery', |
| 28 | +} |
| 29 | + |
| 30 | +function tagOfKeyType(keyType: KeyType): string { |
| 31 | + switch (keyType) { |
| 32 | + case KeyType.DefinitionQuery: return "ide-contextual-queries/local-definitions"; |
| 33 | + case KeyType.ReferenceQuery: return "ide-contextual-queries/local-references"; |
| 34 | + } |
| 35 | +} |
| 36 | + |
| 37 | +async function resolveQueries(cli: CodeQLCliServer, qlpack: string, keyType: KeyType): Promise<string[]> { |
| 38 | + const suiteFile = tmp.fileSync({ postfix: '.qls' }).name; |
| 39 | + const suiteYaml = { qlpack, include: { kind: 'definitions', 'tags contain': tagOfKeyType(keyType) } }; |
| 40 | + await fs.writeFile(suiteFile, yaml.safeDump(suiteYaml), 'utf8'); |
| 41 | + |
| 42 | + const queries = await cli.resolveQueriesInSuite(suiteFile, helpers.getOnDiskWorkspaceFolders()); |
| 43 | + if (queries.length === 0) { |
| 44 | + throw new Error("Couldn't find any queries for qlpack"); |
| 45 | + } |
| 46 | + return queries; |
| 47 | +} |
| 48 | + |
| 49 | +async function qlpackOfDatabase(cli: CodeQLCliServer, db: DatabaseItem): Promise<string | undefined> { |
| 50 | + if (db.contents === undefined) |
| 51 | + return undefined; |
| 52 | + const datasetPath = db.contents.datasetUri.fsPath; |
| 53 | + const { qlpack } = await helpers.resolveDatasetFolder(cli, datasetPath); |
| 54 | + return qlpack; |
| 55 | +} |
| 56 | + |
| 57 | +interface FullLocationLink extends vscode.LocationLink { |
| 58 | + originUri: vscode.Uri; |
| 59 | +} |
| 60 | + |
| 61 | +export class TemplateQueryDefinitionProvider implements vscode.DefinitionProvider { |
| 62 | + private cache: CachedOperation<vscode.LocationLink[]>; |
| 63 | + |
| 64 | + constructor( |
| 65 | + private cli: CodeQLCliServer, |
| 66 | + private qs: QueryServerClient, |
| 67 | + private dbm: DatabaseManager, |
| 68 | + ) { |
| 69 | + this.cache = new CachedOperation<vscode.LocationLink[]>(this.getDefinitions.bind(this)); |
| 70 | + } |
| 71 | + |
| 72 | + async getDefinitions(uriString: string): Promise<vscode.LocationLink[]> { |
| 73 | + return getLinksForUriString(this.cli, this.qs, this.dbm, uriString, (src, _dest) => src === uriString); |
| 74 | + } |
| 75 | + |
| 76 | + async provideDefinition(document: vscode.TextDocument, position: vscode.Position, _token: vscode.CancellationToken): Promise<vscode.LocationLink[]> { |
| 77 | + const fileLinks = await this.cache.get(document.uri.toString()); |
| 78 | + const locLinks: vscode.LocationLink[] = []; |
| 79 | + for (const link of fileLinks) { |
| 80 | + if (link.originSelectionRange!.contains(position)) { |
| 81 | + locLinks.push(link); |
| 82 | + } |
| 83 | + } |
| 84 | + return locLinks; |
| 85 | + } |
| 86 | +} |
| 87 | + |
| 88 | +export class TemplateQueryReferenceProvider implements vscode.ReferenceProvider { |
| 89 | + private cache: CachedOperation<FullLocationLink[]>; |
| 90 | + |
| 91 | + constructor( |
| 92 | + private cli: CodeQLCliServer, |
| 93 | + private qs: QueryServerClient, |
| 94 | + private dbm: DatabaseManager, |
| 95 | + ) { |
| 96 | + this.cache = new CachedOperation<FullLocationLink[]>(this.getReferences.bind(this)); |
| 97 | + } |
| 98 | + |
| 99 | + async getReferences(uriString: string): Promise<FullLocationLink[]> { |
| 100 | + return getLinksForUriString(this.cli, this.qs, this.dbm, uriString, (_src, dest) => dest === uriString); |
| 101 | + } |
| 102 | + |
| 103 | + async provideReferences(document: vscode.TextDocument, position: vscode.Position, _context: vscode.ReferenceContext, _token: vscode.CancellationToken): Promise<vscode.Location[]> { |
| 104 | + const fileLinks = await this.cache.get(document.uri.toString()); |
| 105 | + const locLinks: vscode.Location[] = []; |
| 106 | + for (const link of fileLinks) { |
| 107 | + if (link.targetRange!.contains(position)) { |
| 108 | + locLinks.push({ range: link.originSelectionRange!, uri: link.originUri }); |
| 109 | + } |
| 110 | + } |
| 111 | + return locLinks; |
| 112 | + } |
| 113 | +} |
| 114 | + |
| 115 | +interface FileRange { |
| 116 | + file: vscode.Uri; |
| 117 | + range: vscode.Range; |
| 118 | +} |
| 119 | + |
| 120 | +async function getLinksFromResults(results: QueryWithResults, cli: CodeQLCliServer, db: DatabaseItem, filter: (srcFile: string, destFile: string) => boolean): Promise<FullLocationLink[]> { |
| 121 | + const localLinks: FullLocationLink[] = []; |
| 122 | + const bqrsPath = results.query.resultsPaths.resultsPath; |
| 123 | + const info = await cli.bqrsInfo(bqrsPath); |
| 124 | + const selectInfo = getResultSetSchema(SELECT_QUERY_NAME, info); |
| 125 | + if (selectInfo && selectInfo.columns.length == 3 |
| 126 | + && selectInfo.columns[0].kind == ColumnKindCode.ENTITY |
| 127 | + && selectInfo.columns[1].kind == ColumnKindCode.ENTITY |
| 128 | + && selectInfo.columns[2].kind == ColumnKindCode.STRING) { |
| 129 | + // TODO: Page this |
| 130 | + const allTuples = await cli.bqrsDecode(bqrsPath, SELECT_QUERY_NAME); |
| 131 | + for (const tuple of allTuples.tuples) { |
| 132 | + const src = tuple[0] as EntityValue; |
| 133 | + const dest = tuple[1] as EntityValue; |
| 134 | + const srcFile = src.url && fileRangeFromURI(src.url, db); |
| 135 | + const destFile = dest.url && fileRangeFromURI(dest.url, db); |
| 136 | + if (srcFile && destFile && filter(srcFile.file.toString(), destFile.file.toString())) { |
| 137 | + localLinks.push({ targetRange: destFile.range, targetUri: destFile.file, originSelectionRange: srcFile.range, originUri: srcFile.file }); |
| 138 | + } |
| 139 | + } |
| 140 | + } |
| 141 | + return localLinks; |
| 142 | +} |
| 143 | + |
| 144 | +async function getLinksForUriString( |
| 145 | + cli: CodeQLCliServer, |
| 146 | + qs: QueryServerClient, |
| 147 | + dbm: DatabaseManager, |
| 148 | + uriString: string, |
| 149 | + filter: (src: string, dest: string) => boolean |
| 150 | +) { |
| 151 | + const uri = decodeSourceArchiveUri(vscode.Uri.parse(uriString)); |
| 152 | + const sourceArchiveUri = vscode.Uri.file(uri.sourceArchiveZipPath).with({ scheme: zipArchiveScheme }); |
| 153 | + |
| 154 | + const db = dbm.findDatabaseItemBySourceArchive(sourceArchiveUri); |
| 155 | + if (db) { |
| 156 | + const qlpack = await qlpackOfDatabase(cli, db); |
| 157 | + if (qlpack === undefined) { |
| 158 | + throw new Error("Can't infer qlpack from database source archive"); |
| 159 | + } |
| 160 | + const links: FullLocationLink[] = [] |
| 161 | + for (const query of await resolveQueries(cli, qlpack, KeyType.ReferenceQuery)) { |
| 162 | + const templates: messages.TemplateDefinitions = { |
| 163 | + [TEMPLATE_NAME]: { |
| 164 | + values: { |
| 165 | + tuples: [[{ |
| 166 | + stringValue: uri.pathWithinSourceArchive |
| 167 | + }]] |
| 168 | + } |
| 169 | + } |
| 170 | + }; |
| 171 | + const results = await compileAndRunQueryAgainstDatabase(cli, qs, db, false, vscode.Uri.file(query), templates); |
| 172 | + if (results.result.resultType == messages.QueryResultType.SUCCESS) { |
| 173 | + links.push(...await getLinksFromResults(results, cli, db, filter)); |
| 174 | + } |
| 175 | + } |
| 176 | + return links; |
| 177 | + } else { |
| 178 | + return []; |
| 179 | + } |
| 180 | +} |
| 181 | + |
| 182 | +function fileRangeFromURI(uri: UrlValue, db: DatabaseItem): FileRange | undefined { |
| 183 | + if (typeof uri === "string") { |
| 184 | + return undefined; |
| 185 | + } else if ('startOffset' in uri) { |
| 186 | + return undefined; |
| 187 | + } else { |
| 188 | + const loc = uri as LineColumnLocation; |
| 189 | + const range = new vscode.Range(Math.max(0, loc.startLine - 1), |
| 190 | + Math.max(0, loc.startColumn - 1), |
| 191 | + Math.max(0, loc.endLine - 1), |
| 192 | + Math.max(0, loc.endColumn)); |
| 193 | + try { |
| 194 | + const parsed = vscode.Uri.parse(uri.uri, true); |
| 195 | + if (parsed.scheme === "file") { |
| 196 | + return { file: db.resolveSourceFile(parsed.fsPath), range }; |
| 197 | + } |
| 198 | + return undefined; |
| 199 | + } catch (e) { |
| 200 | + return undefined; |
| 201 | + } |
| 202 | + } |
| 203 | +} |
0 commit comments