diff --git a/src/Client/features/authentication.ts b/src/Client/features/authentication.ts new file mode 100644 index 0000000..655e525 --- /dev/null +++ b/src/Client/features/authentication.ts @@ -0,0 +1,195 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +/* + * Bridges AAD token requests from the language server to VS Code's built-in + * Microsoft authentication provider. The server uses this fallback path only + * when its native Kusto.Data authentication flow cannot complete in-process + * (e.g. running under Remote-SSH / WSL / Codespaces, or when the local token + * cache is wedged). Routing sign-in through VS Code lets the user manage + * accounts via the standard Accounts gear and avoids hosting any sign-in UI + * in the language-server child process. + */ + +import { authentication, AuthenticationSession, Disposable, window } from 'vscode'; + +/** Parameters for the server's getAuthenticationToken request. */ +export interface GetAuthenticationTokenParams { + clusterUri: string; +} + +/** Result of the server's getAuthenticationToken request. */ +export interface GetAuthenticationTokenResult { + accessToken: string | null; +} + +/** + * All Kusto clusters within the same Azure cloud share a single AAD resource + * (e.g. `https://kusto.kusto.windows.net` for the public cloud). Issuing the + * token against that shared resource means one sign-in covers every cluster + * the user has configured, including all the clusters whose schemas are + * fetched on extension startup. Per-cluster scopes would prompt the user + * once per cluster. + * + * In some environments (notably the Extension Development Host and some + * remote setups) `getSession({ silent: true })` returns nothing even + * immediately after a successful `createIfNone`, which would cause every + * subsequent token request to re-prompt. To stay robust we keep the + * resolved session in an in-process cache and only re-acquire it when + * VS Code tells us via `onDidChangeSessions` that it has changed (e.g. + * the user signed out via the Accounts gear). + */ +const inFlightSessions = new Map>(); +const cachedSessions = new Map(); +let sessionChangeListener: Disposable | undefined; + +function ensureSessionChangeListener(): void { + if (sessionChangeListener) { + return; + } + sessionChangeListener = authentication.onDidChangeSessions(e => { + if (e.provider.id === 'microsoft') { + // Any change to Microsoft sessions (sign-in, sign-out, switch account) + // invalidates our cache; the next token request will re-acquire. + cachedSessions.clear(); + } + }); +} + +/** + * Returns true if the JWT access token expires within the next 5 minutes, + * or if its expiry cannot be determined (in which case we err on the side + * of refreshing). Used to decide when the cached session is no longer safe + * to reuse and we need to ask VS Code for a fresh token. + */ +function isTokenExpiringSoon(accessToken: string): boolean { + const refreshSkewSeconds = 5 * 60; + try { + const parts = accessToken.split('.'); + if (parts.length < 2 || !parts[1]) { + return true; + } + // base64url -> base64 + const b64 = parts[1].replace(/-/g, '+').replace(/_/g, '/'); + const padded = b64 + '='.repeat((4 - (b64.length % 4)) % 4); + const json = Buffer.from(padded, 'base64').toString('utf8'); + const payload = JSON.parse(json) as { exp?: number }; + if (typeof payload.exp !== 'number') { + return true; + } + const nowSeconds = Math.floor(Date.now() / 1000); + return payload.exp - nowSeconds <= refreshSkewSeconds; + } catch { + return true; + } +} + +/** + * Returns the AAD resource scope for the cloud that hosts the given Kusto + * cluster, or null if the URI cannot be parsed. The convention for Kusto + * first-party AAD resources is `kusto.`: + * + * *.kusto.windows.net → https://kusto.kusto.windows.net/.default + * *.kusto.usgovcloudapi.net → https://kusto.kusto.usgovcloudapi.net/.default + * *.kusto.chinacloudapi.cn → https://kusto.kusto.chinacloudapi.cn/.default + * *.kusto.cloudapi.de → https://kusto.kusto.cloudapi.de/.default + * + * For any other host (custom DNS / non-public cluster) we fall back to the + * cluster's own URI as the scope - that still works for that cluster, the + * user just may be prompted more than once. + */ +export function getKustoScope(clusterUri: string): string | null { + let hostname: string; + try { + // hostname (not host) deliberately excludes any :port suffix — AAD + // resource URIs for Kusto never include a port. + hostname = new URL(clusterUri).hostname.toLowerCase(); + } catch { + return null; + } + // Match the public Kusto host shape "*.kusto." (apex may have one + // or more labels, e.g. "windows.net" or "usgovcloudapi.net"). + const match = hostname.match(/\.kusto\.([^/]+)$/); + if (match) { + return `https://kusto.kusto.${match[1]}/.default`; + } + // Unknown host shape: fall back to a per-cluster scope (port stripped). + // Restrict the fallback to https — public Kusto and any sane on-prem / + // sovereign / private-link deployment is HTTPS, and refusing to derive + // a scope from an http:// URI prevents a malicious or misconfigured + // connection string from coaxing AAD into issuing a token bound to a + // plaintext resource. AAD itself only issues tokens for resources + // registered in the user's tenant, but blocking http here is a cheap + // additional guard with no legitimate-user impact. + try { + const u = new URL(clusterUri); + if (u.protocol !== 'https:') { + return null; + } + return `${u.protocol}//${u.hostname}/.default`; + } catch { + return null; + } +} + +/** + * Acquires an AAD access token for a Kusto cluster using VS Code's built-in + * Microsoft authentication provider. Returns a result whose accessToken is + * null if the user cancelled or no session could be obtained. + */ +export async function acquireMicrosoftAuthenticationToken( + clusterUri: string +): Promise { + const scope = getKustoScope(clusterUri); + if (!scope) { + return { accessToken: null }; + } + + ensureSessionChangeListener(); + + // Fast path: a cached session whose access token is not yet near expiry. + // We always prefer reusing the cached token over calling getSession again, + // because in some environments (notably the Extension Development Host + // and certain remote configurations) `getSession({ silent: true })` can + // return null even when a valid session exists, and `createIfNone` would + // then re-prompt the user. The cache is cleared whenever VS Code reports + // a session change (sign-out, account switch). + const cached = cachedSessions.get(scope); + if (cached && !isTokenExpiringSoon(cached.accessToken)) { + return { accessToken: cached.accessToken }; + } + + let pending = inFlightSessions.get(scope); + if (!pending) { + pending = (async () => { + // Silent first: returns a fresh access token from VS Code's persistent + // session cache without showing UI. Survives VS Code restarts. + let session = await authentication.getSession('microsoft', [scope], { silent: true }); + if (!session && !cached) { + // First-ever request for this scope and no persisted session: + // prompt the user via VS Code's Accounts UI. + session = await authentication.getSession('microsoft', [scope], { createIfNone: true }); + } + // If silent unexpectedly fails but we still have a cached session + // whose token has not yet expired, fall back to the cached session. + if (!session && cached && !isTokenExpiringSoon(cached.accessToken)) { + session = cached; + } + return session; + })(); + inFlightSessions.set(scope, pending); + pending.finally(() => inFlightSessions.delete(scope)); + } + + try { + const session = await pending; + if (session) { + cachedSessions.set(scope, session); + } + return { accessToken: session?.accessToken ?? null }; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + window.showErrorMessage(`Kusto: failed to acquire Microsoft account token for ${clusterUri}. ${message}`); + return { accessToken: null }; + } +} diff --git a/src/Client/features/connectionsPanel.ts b/src/Client/features/connectionsPanel.ts index 7ee7158..1523db1 100644 --- a/src/Client/features/connectionsPanel.ts +++ b/src/Client/features/connectionsPanel.ts @@ -210,7 +210,12 @@ export class ConnectionsPanel { }; context.subscriptions.push( this.treeView.onDidChangeVisibility((e) => { - if (e.visible) { void tryPromptImport(); } + if (e.visible) { + void tryPromptImport(); + // Re-sync tree selection with the active document now that the view + // is visible again (selection updates were skipped while hidden). + void this.updateTreeSelectionForActiveDocument(); + } }) ); // Also check if the panel is already visible at activation time @@ -461,6 +466,14 @@ export class ConnectionsPanel { } private async programmaticSelectTreeItem(item: KustoTreeItem, options?: { select?: boolean; focus?: boolean; expand?: boolean }): Promise { + // Skip reveal when the tree view is not visible. TreeView.reveal() forces the + // containing view to become visible, which would steal focus away from whatever + // activity bar panel (e.g. Source Control) the user is currently using. + // Selection will be re-applied via the onDidChangeVisibility handler when the + // user navigates back to the Kusto Explorer panel. + if (!this.treeView.visible) { + return; + } this.programmaticSelectionCount++; try { await this.treeView.reveal(item, options); diff --git a/src/Client/features/server.ts b/src/Client/features/server.ts index 6c46791..1810b10 100644 --- a/src/Client/features/server.ts +++ b/src/Client/features/server.ts @@ -3,6 +3,11 @@ import { LanguageClient, InitializeResult } from 'vscode-languageclient/node'; import { CancellationToken, Disposable, ExtensionContext } from 'vscode'; +import { + acquireMicrosoftAuthenticationToken, + GetAuthenticationTokenParams, + GetAuthenticationTokenResult, +} from './authentication'; /** * Interface for the Kusto Language Server, used by all components. @@ -63,6 +68,15 @@ export class Server implements IServer { this.client.onRequest('kusto/setData', async (params: SetDataParams) => { await context.globalState.update(params.key, params.data); }); + // Bridge server AAD token requests to VS Code's built-in Microsoft + // authentication provider. This lets sign-in UI live in the host + // (the VS Code window) rather than in the language server process, + // which avoids "non-interactive environment" failures on remote-SSH / + // WSL / Codespaces and lets the user manage sign-in through the + // standard Accounts gear. + this.client.onRequest('kusto/getAuthenticationToken', async (params: GetAuthenticationTokenParams) => { + return await acquireMicrosoftAuthenticationToken(params.clusterUri); + }); } // ─── LSP Requests ────────────────────────────────────────────────── diff --git a/src/Client/tests/unit/authentication.test.ts b/src/Client/tests/unit/authentication.test.ts new file mode 100644 index 0000000..e62befa --- /dev/null +++ b/src/Client/tests/unit/authentication.test.ts @@ -0,0 +1,130 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +import { describe, it, expect } from 'vitest'; +import { getKustoScope } from '../../features/authentication'; + +// ─── getKustoScope ─────────────────────────────────────────────────────────── +// +// Validates the host-to-scope mapping used to bridge AAD token requests from +// the language server through VS Code's Microsoft authentication provider. +// Regressions here would cause unnecessary re-prompts (per-cluster scopes +// instead of the shared cloud-wide resource) or, worse, silent failures. + +describe('getKustoScope', () => { + describe('public clouds', () => { + it('maps a public Azure cluster to the kusto.windows.net resource', () => { + expect(getKustoScope('https://help.kusto.windows.net')) + .toBe('https://kusto.kusto.windows.net/.default'); + }); + + it('maps a US Government cluster to the usgovcloudapi.net resource', () => { + expect(getKustoScope('https://mycluster.kusto.usgovcloudapi.net')) + .toBe('https://kusto.kusto.usgovcloudapi.net/.default'); + }); + + it('maps a China cluster to the chinacloudapi.cn resource', () => { + expect(getKustoScope('https://mycluster.kusto.chinacloudapi.cn')) + .toBe('https://kusto.kusto.chinacloudapi.cn/.default'); + }); + + it('maps a Germany cluster to the cloudapi.de resource', () => { + expect(getKustoScope('https://mycluster.kusto.cloudapi.de')) + .toBe('https://kusto.kusto.cloudapi.de/.default'); + }); + + it('returns the same shared scope for every cluster in the same cloud', () => { + // The whole point of the shared cloud apex resource is that one + // sign-in covers every cluster a user has configured. + const a = getKustoScope('https://help.kusto.windows.net'); + const b = getKustoScope('https://other.kusto.windows.net'); + const c = getKustoScope('https://yet-another.kusto.windows.net'); + expect(a).toBe(b); + expect(b).toBe(c); + }); + + it('lowercases the hostname', () => { + expect(getKustoScope('https://HELP.KUSTO.WINDOWS.NET')) + .toBe('https://kusto.kusto.windows.net/.default'); + }); + + it('ignores any path on the cluster URI', () => { + expect(getKustoScope('https://help.kusto.windows.net/v1/rest/query')) + .toBe('https://kusto.kusto.windows.net/.default'); + }); + + it('ignores a port suffix when matching the public-cloud apex', () => { + // Public-cloud Kusto resources never include a port; the scope + // must collapse to the shared apex regardless of the cluster URI. + expect(getKustoScope('https://help.kusto.windows.net:443')) + .toBe('https://kusto.kusto.windows.net/.default'); + }); + + it('handles an ingest- prefixed cluster (still inside the apex)', () => { + expect(getKustoScope('https://ingest-help.kusto.windows.net')) + .toBe('https://kusto.kusto.windows.net/.default'); + }); + + it('handles a follower / regional cluster suffix inside the apex', () => { + expect(getKustoScope('https://mycluster.eastus.kusto.windows.net')) + .toBe('https://kusto.kusto.windows.net/.default'); + }); + }); + + describe('custom DNS / non-public hosts', () => { + it('falls back to a per-cluster scope when the host is not under .kusto.', () => { + expect(getKustoScope('https://mycluster.example.com')) + .toBe('https://mycluster.example.com/.default'); + }); + + it('strips the port from the per-cluster fallback scope', () => { + // AAD resource URIs never include a port; the fallback path must + // also strip it so the scope is a valid resource identifier. + expect(getKustoScope('https://mycluster.example.com:8080')) + .toBe('https://mycluster.example.com/.default'); + }); + + it('returns null for an http:// (non-TLS) custom host', () => { + // Refusing to derive a scope from a plaintext URI prevents a + // malicious or misconfigured connection string from coaxing AAD + // into issuing a token bound to an http resource. Public Kusto + // and supported on-prem / sovereign deployments are all HTTPS. + expect(getKustoScope('http://mycluster.example.com')).toBeNull(); + }); + + it('lowercases the hostname in the per-cluster fallback', () => { + expect(getKustoScope('https://MyCluster.Example.COM')) + .toBe('https://mycluster.example.com/.default'); + }); + + it('does not match a host that merely contains "kusto" as a label without the .kusto. apex pattern', () => { + // "kusto-test.example.com" must not be misinterpreted as a public + // Kusto cloud host. + expect(getKustoScope('https://kusto-test.example.com')) + .toBe('https://kusto-test.example.com/.default'); + }); + + it('does not match a host whose name starts with "kusto." (no left-side cluster label)', () => { + // The regex requires ".kusto." (with a leading dot) so a bare + // "kusto.windows.net" host falls through to the fallback path. + expect(getKustoScope('https://kusto.windows.net')) + .toBe('https://kusto.windows.net/.default'); + }); + }); + + describe('invalid input', () => { + it('returns null for a non-URI string', () => { + expect(getKustoScope('not a uri')).toBeNull(); + }); + + it('returns null for an empty string', () => { + expect(getKustoScope('')).toBeNull(); + }); + + it('returns null for a bare hostname without scheme', () => { + // URL parsing requires a scheme; a bare hostname is rejected + // rather than silently producing a malformed scope. + expect(getKustoScope('help.kusto.windows.net')).toBeNull(); + }); + }); +}); diff --git a/src/Server/Connections/ConnectionManager.cs b/src/Server/Connections/ConnectionManager.cs index 39277d5..99ba3a2 100644 --- a/src/Server/Connections/ConnectionManager.cs +++ b/src/Server/Connections/ConnectionManager.cs @@ -7,9 +7,11 @@ using Kusto.Data.Common; using Kusto.Data.Common.Impl; using Kusto.Data.Data; +using Kusto.Data.Exceptions; using Kusto.Data.Net.Client; using Kusto.Language; using Kusto.Language.Editor; +using System.Collections.Concurrent; using System.Collections.Immutable; using System.Data; using System.Diagnostics.CodeAnalysis; @@ -22,6 +24,8 @@ namespace Kusto.Vscode; /// public class ConnectionManager : IConnectionManager { + private readonly IAuthenticationProvider? _authProvider; + private ImmutableDictionary _connectionStringsToInfoMap = ImmutableDictionary.Empty; @@ -30,6 +34,25 @@ public class ConnectionManager : IConnectionManager private string _defaultDomain = KustoFacts.KustoWindowsNet; + // Tracks which clusters have been observed to fail native authentication. + // Once a cluster is in this set, we route its connections through the + // host-bridged auth provider (e.g. VS Code) without first re-attempting + // native auth. The set is per-process and reset on restart. + private readonly ConcurrentDictionary _fallbackRequiredByCluster = + new(StringComparer.OrdinalIgnoreCase); + + // Carries the caller's CancellationToken into the + // WithAadTokenProviderAuthentication callback that Kusto.Data invokes + // synchronously while servicing a query. The callback signature has no + // CancellationToken parameter, so we use AsyncLocal to flow it through + // the same async context as ExecuteCoreAsync. + private static readonly AsyncLocal _ambientAuthToken = new(); + + public ConnectionManager(IAuthenticationProvider? authProvider = null) + { + _authProvider = authProvider; + } + private class ConnectionInfo { public required KustoConnection Connection { get; set; } @@ -82,12 +105,81 @@ ConnectionInfo CreateConnectionInfo(string connectionString) builder.FederatedSecurity = true; } + // The "primary" builder always uses Kusto.Data's native auth path, + // which on a normal Windows desktop transparently does WAM/SSO with + // the user's signed-in work account - no prompt at all. When that + // path can't proceed (remote-SSH / WSL / Codespaces / dev tunnels, + // or a stuck local cache) it throws KustoClientAuthenticationException + // and we fall back to host-bridged auth via the IAuthenticationProvider. + // See KustoConnection.ExecuteAsync for the retry logic. var connection = new KustoConnection(this, builder); - return new ConnectionInfo { Connection = connection }; } } + /// + /// Returns true if a previous query against the given cluster failed + /// native authentication and the cluster has been promoted to use the + /// host-bridged auth path. + /// + internal bool ShouldUseFallback(string clusterHostName) + => _fallbackRequiredByCluster.ContainsKey(clusterHostName); + + /// + /// Marks a cluster as requiring host-bridged authentication for the + /// remainder of this server process. Called by + /// after observing a from + /// the native auth path. + /// + internal void MarkFallbackRequired(string clusterHostName) + => _fallbackRequiredByCluster[clusterHostName] = true; + + /// + /// Builds a host-bridged variant of the supplied builder that routes AAD + /// token acquisition through the registered . + /// Returns null if no provider is registered or the supplied builder + /// already specifies an explicit authentication method. + /// + internal KustoConnectionStringBuilder? CreateFallbackBuilder(KustoConnectionStringBuilder primary) + { + if (_authProvider == null || HasExplicitAuthentication(primary)) + { + return null; + } + + var fallback = new KustoConnectionStringBuilder(primary); + var clusterUri = fallback.DataSource; + var provider = _authProvider; + return fallback.WithAadTokenProviderAuthentication(async () => + { + // Pick up the caller's CancellationToken if one was published into + // the ambient async context by ExecuteCoreAsync. Falls back to + // None if Kusto.Data invokes the callback outside our query path + // (e.g. background metadata refresh). + var token = await provider + .GetAccessTokenAsync(clusterUri, _ambientAuthToken.Value) + .ConfigureAwait(false); + return token ?? throw new InvalidOperationException( + $"Authentication failed: no access token was returned for '{clusterUri}'."); + }); + } + + private static bool HasExplicitAuthentication(KustoConnectionStringBuilder builder) + { + // Treat the connection as having explicit auth if any credential-bearing + // field has been set, or if an interactive auth flow was already requested. + return builder.AzCliInteractiveLogin + || !string.IsNullOrEmpty(builder.ApplicationClientId) + || !string.IsNullOrEmpty(builder.ApplicationKey) + || builder.ApplicationCertificateBlob != null + || !string.IsNullOrEmpty(builder.ApplicationCertificateThumbprint) + || !string.IsNullOrEmpty(builder.ApplicationToken) + || !string.IsNullOrEmpty(builder.UserToken) + || !string.IsNullOrEmpty(builder.UserID) + || builder.TokenProviderCallback != null + || builder.KustoTokenCredentialsProvider != null; + } + public bool TryGetConnection(string cluster, [NotNullWhen(true)] out IConnection? connection) { if (_clusterToInfoMap.TryGetValue(cluster, out var info)) @@ -103,33 +195,71 @@ public bool TryGetConnection(string cluster, [NotNullWhen(true)] out IConnection private class KustoConnection : IConnection { private readonly ConnectionManager _manager; - private readonly KustoConnectionStringBuilder _builder; + private readonly KustoConnectionStringBuilder _primaryBuilder; + private KustoConnectionStringBuilder? _fallbackBuilder; // lazily built on first need public KustoConnection(ConnectionManager manager, KustoConnectionStringBuilder builder) { _manager = manager; - _builder = builder; + _primaryBuilder = builder; } - public string Cluster => _builder.Hostname; - public string? Database => _builder.InitialCatalog; + /// + /// Exposed for tests. Returns the currently active connection-string builder, + /// which is the fallback (host-bridged) builder if this cluster has been + /// promoted to fallback mode, otherwise the primary (native auth) builder. + /// + internal KustoConnectionStringBuilder Builder => ActiveBuilder; + + /// + /// Exposed for tests. Returns the primary (native auth) builder. + /// + internal KustoConnectionStringBuilder PrimaryBuilder => _primaryBuilder; + + /// + /// Exposed for tests. Returns the fallback (host-bridged auth) builder, + /// constructing it on demand. May be null if no + /// is registered or the connection string already specifies explicit auth. + /// + internal KustoConnectionStringBuilder? FallbackBuilder => GetFallbackBuilder(); + + private KustoConnectionStringBuilder ActiveBuilder + => UseFallback ? (GetFallbackBuilder() ?? _primaryBuilder) : _primaryBuilder; + + private bool UseFallback => _manager.ShouldUseFallback(_primaryBuilder.Hostname); + + private KustoConnectionStringBuilder? GetFallbackBuilder() + { + if (_fallbackBuilder == null) + { + var fb = _manager.CreateFallbackBuilder(_primaryBuilder); + if (fb != null) + { + Interlocked.CompareExchange(ref _fallbackBuilder, fb, null); + } + } + return _fallbackBuilder; + } + + public string Cluster => _primaryBuilder.Hostname; + public string? Database => _primaryBuilder.InitialCatalog; public IConnection WithCluster(string clusterName) { var clusterUri = KustoFacts.GetFullHostName(clusterName, _manager._defaultDomain); if (!clusterUri.Contains("://") - && !string.IsNullOrEmpty(_builder.ConnectionScheme)) + && !string.IsNullOrEmpty(_primaryBuilder.ConnectionScheme)) { - clusterUri = _builder.ConnectionScheme + "://" + clusterUri; + clusterUri = _primaryBuilder.ConnectionScheme + "://" + clusterUri; } // borrow most security settings from default cluster connection - var builder = new KustoConnectionStringBuilder(_builder); + var builder = new KustoConnectionStringBuilder(_primaryBuilder); builder.DataSource = clusterUri; - builder.ApplicationCertificateBlob = _builder.ApplicationCertificateBlob; - builder.ApplicationKey = _builder.ApplicationKey; + builder.ApplicationCertificateBlob = _primaryBuilder.ApplicationCertificateBlob; + builder.ApplicationKey = _primaryBuilder.ApplicationKey; builder.InitialCatalog = "NetDefaultDB"; return new KustoConnection(_manager, builder); @@ -137,35 +267,63 @@ public IConnection WithCluster(string clusterName) public IConnection WithDatabase(string database) { - var newBuilder = new KustoConnectionStringBuilder(_builder) { InitialCatalog = database }; + var newBuilder = new KustoConnectionStringBuilder(_primaryBuilder) { InitialCatalog = database }; return new KustoConnection(_manager, newBuilder); } - // These providers are IDisposable but are intentionally not disposed. - // They are cached for the lifetime of the server process and cleaned up on exit. - private ICslQueryProvider? _queryProvider; + // Cached providers. Kept per auth-mode so that switching from primary + // to fallback after an auth failure produces a fresh provider built + // from the fallback builder. These providers are IDisposable but are + // intentionally not disposed - they are cached for the lifetime of + // the server process and cleaned up on exit. + private ICslQueryProvider? _primaryQueryProvider; + private ICslQueryProvider? _fallbackQueryProvider; + private ICslAdminProvider? _primaryAdminProvider; + private ICslAdminProvider? _fallbackAdminProvider; + public ICslQueryProvider QueryProvider { get { - if (_queryProvider == null) + if (UseFallback && GetFallbackBuilder() is { } fb) + { + if (_fallbackQueryProvider == null) + { + Interlocked.CompareExchange(ref _fallbackQueryProvider, + KustoClientFactory.CreateCslQueryProvider(fb), null); + } + return _fallbackQueryProvider; + } + + if (_primaryQueryProvider == null) { - Interlocked.CompareExchange(ref _queryProvider, KustoClientFactory.CreateCslQueryProvider(_builder), null); + Interlocked.CompareExchange(ref _primaryQueryProvider, + KustoClientFactory.CreateCslQueryProvider(_primaryBuilder), null); } - return _queryProvider; + return _primaryQueryProvider; } } - private ICslAdminProvider? _adminProvider; public ICslAdminProvider AdminProvider { - get - { - if (_adminProvider == null) + get + { + if (UseFallback && GetFallbackBuilder() is { } fb) { - Interlocked.CompareExchange(ref _adminProvider, KustoClientFactory.CreateCslAdminProvider(_builder), null); + if (_fallbackAdminProvider == null) + { + Interlocked.CompareExchange(ref _fallbackAdminProvider, + KustoClientFactory.CreateCslAdminProvider(fb), null); + } + return _fallbackAdminProvider; + } + + if (_primaryAdminProvider == null) + { + Interlocked.CompareExchange(ref _primaryAdminProvider, + KustoClientFactory.CreateCslAdminProvider(_primaryBuilder), null); } - return _adminProvider; + return _primaryAdminProvider; } } @@ -179,7 +337,63 @@ public async Task ExecuteAsync( { try { - var properties = CreateClientRequestProperties(options ?? ImmutableDictionary.Empty, parameters ?? ImmutableDictionary.Empty, clientRequestId); + return await ExecuteCoreAsync(query, options, parameters, clientRequestId, cancellationToken) + .ConfigureAwait(false); + } + catch (KustoClientAuthenticationException) when (CanRetryWithFallback()) + { + // Native authentication failed (e.g. Kusto.Data could not run + // its WAM/MSAL prompt because the server process is hosted in + // a non-interactive environment, or its cached token is stuck). + // Promote this cluster to host-bridged auth and retry once. + _manager.MarkFallbackRequired(_primaryBuilder.Hostname); + try + { + return await ExecuteCoreAsync(query, options, parameters, clientRequestId, cancellationToken) + .ConfigureAwait(false); + } + catch (Exception retryEx) + { + return new ExecuteResult + { + Diagnostics = [ErrorDecoder.GetDiagnostic(retryEx, query)] + }; + } + } + catch (Exception ex) + { + return new ExecuteResult + { + Diagnostics = [ErrorDecoder.GetDiagnostic(ex, query)] + }; + } + } + + private bool CanRetryWithFallback() + => !UseFallback && GetFallbackBuilder() != null; + + private async Task ExecuteCoreAsync( + EditString query, + ImmutableDictionary? options, + ImmutableDictionary? parameters, + string? clientRequestId, + CancellationToken cancellationToken) + { + // Publish the caller's CancellationToken into the ambient async + // context so the WithAadTokenProviderAuthentication callback (if + // Kusto.Data triggers fallback auth during this call) can observe + // and respect query cancellation. Save and restore the previous + // value so a later token request (e.g. a background metadata + // refresh) doesn't observe a stale/cancelled token from this + // query's context. + var previousAuthToken = _ambientAuthToken.Value; + _ambientAuthToken.Value = cancellationToken; + try + { + var properties = CreateClientRequestProperties( + options ?? ImmutableDictionary.Empty, + parameters ?? ImmutableDictionary.Empty, + clientRequestId); var resultReader = (Kusto.Language.KustoCode.GetKind(query) == CodeKinds.Command) ? await this.AdminProvider.ExecuteControlCommandAsync(this.Database, query, properties).ConfigureAwait(false) : await this.QueryProvider.ExecuteQueryAsync(this.Database, query, properties, cancellationToken).ConfigureAwait(false); @@ -203,12 +417,9 @@ public async Task ExecuteAsync( }; } } - catch (Exception ex) + finally { - return new ExecuteResult - { - Diagnostics = [ErrorDecoder.GetDiagnostic(ex, query)] - }; + _ambientAuthToken.Value = previousAuthToken; } } diff --git a/src/Server/Connections/IAuthenticationProvider.cs b/src/Server/Connections/IAuthenticationProvider.cs new file mode 100644 index 0000000..e1556ca --- /dev/null +++ b/src/Server/Connections/IAuthenticationProvider.cs @@ -0,0 +1,22 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +namespace Kusto.Vscode; + +/// +/// Provides AAD bearer tokens for Kusto cluster authentication. +/// Implementations typically delegate to the host (e.g. VS Code's +/// vscode.authentication API) so that sign-in UI can be shown +/// in the host process rather than in the language server process. +/// +public interface IAuthenticationProvider +{ + /// + /// Acquires an access token for the given cluster URI. + /// + /// The cluster URI (e.g. https://help.kusto.windows.net). + /// The implementation will derive the appropriate AAD scope from this. + /// Cancellation token. + /// A bearer access token, or null if no token could be acquired. + Task GetAccessTokenAsync(string clusterUri, CancellationToken cancellationToken); +} diff --git a/src/Server/Server.cs b/src/Server/Server.cs index d13d1d2..ca8a629 100644 --- a/src/Server/Server.cs +++ b/src/Server/Server.cs @@ -16,7 +16,7 @@ namespace Kusto.Vscode; -public class Server : LspServer, ILogger, ISettingSource, IStorage +public class Server : LspServer, ILogger, ISettingSource, IStorage, IAuthenticationProvider { private readonly ILogger _logger; private readonly ISettingSource _settingSource; @@ -71,7 +71,14 @@ public Server( _settingSource = this; _storage = this; _optionsManager = new OptionsManager(_settingSource); - _connectionManager = new ConnectionManager(); + // When running as a VS Code extension server ("vscode" arg passed by the + // client), route AAD authentication through the host so that VS Code's + // built-in Microsoft account UI can be used instead of Kusto.Data's + // in-process MSAL/WAM flow. This avoids "non-interactive environment" + // failures on remote-SSH/WSL/Codespaces and gives users a single place + // (the Accounts gear) to manage their sign-in. + var tokenProvider = _args.Contains("vscode") ? (IAuthenticationProvider)this : null; + _connectionManager = new ConnectionManager(tokenProvider); var schemaSource = new ServerSchemaSource(_connectionManager, _logger); _schemaManager = new SchemaManager(schemaSource, _storage, _logger); _symbolManager = new SymbolManager(_schemaManager, _optionsManager, _logger); @@ -2625,6 +2632,47 @@ public class SetDataParams } #endregion + + #region IAuthenticationProvider + + /// + /// Asks the client (e.g. VS Code) to acquire an AAD access token for the + /// given cluster. The client uses its own authentication UI / cache so the + /// server process never has to host any sign-in interaction. + /// + /// + /// A safety timeout caps how long we wait for the client. The user may + /// legitimately take a while to interact with a sign-in prompt, but if + /// the client process is hung or has dropped the request entirely we + /// don't want to keep a server-side query stuck forever. Five minutes + /// is comfortably longer than any normal sign-in flow. + /// + async Task IAuthenticationProvider.GetAccessTokenAsync(string clusterUri, CancellationToken cancellationToken) + { + using var timeoutCts = new CancellationTokenSource(TimeSpan.FromMinutes(5)); + using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token); + var result = await SendRequestAsync( + "kusto/getAuthenticationToken", + new GetAuthenticationTokenParams { ClusterUri = clusterUri }, + linkedCts.Token).ConfigureAwait(false); + return result?.AccessToken; + } + + [DataContract] + public class GetAuthenticationTokenParams + { + [DataMember(Name = "clusterUri")] + public required string ClusterUri { get; set; } + } + + [DataContract] + public class GetAuthenticationTokenResult + { + [DataMember(Name = "accessToken")] + public string? AccessToken { get; set; } + } + + #endregion } diff --git a/src/ServerTests/Features/ConnectionManagerTests.cs b/src/ServerTests/Features/ConnectionManagerTests.cs index a0a0d40..47e0a19 100644 --- a/src/ServerTests/Features/ConnectionManagerTests.cs +++ b/src/ServerTests/Features/ConnectionManagerTests.cs @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT license. +using System.Reflection; +using Kusto.Data; using Kusto.Vscode; namespace Tests.Features; @@ -271,4 +273,165 @@ public void ConnectionChaining_MultipleWithCalls_ProducesCorrectResult() } #endregion + + #region IAuthenticationProvider Tests + + private sealed class CountingAuthProvider : IAuthenticationProvider + { + public int CallCount; + public string? LastClusterUri; + public string? Token = "test-access-token"; + + public Task GetAccessTokenAsync(string clusterUri, CancellationToken cancellationToken) + { + CallCount++; + LastClusterUri = clusterUri; + return Task.FromResult(Token); + } + } + + /// + /// Reads the named non-public property from the internal KustoConnection + /// wrapper. Used to verify how the connection has wired up authentication + /// without actually executing a query against a real cluster. Reading the + /// property (rather than a private field) keeps these tests decoupled from + /// internal field naming. + /// + private static KustoConnectionStringBuilder? GetBuilderProperty(IConnection connection, string propertyName) + { + var prop = connection.GetType().GetProperty(propertyName, BindingFlags.Instance | BindingFlags.NonPublic); + Assert.IsNotNull(prop, $"Expected KustoConnection to expose an internal {propertyName} property."); + return prop!.GetValue(connection) as KustoConnectionStringBuilder; + } + + private static KustoConnectionStringBuilder? GetFallbackBuilder(IConnection connection) + => GetBuilderProperty(connection, "FallbackBuilder"); + + private static void MarkFallbackRequired(ConnectionManager manager, string clusterHostName) + { + var method = typeof(ConnectionManager).GetMethod("MarkFallbackRequired", BindingFlags.Instance | BindingFlags.NonPublic); + Assert.IsNotNull(method, "Expected ConnectionManager to expose an internal MarkFallbackRequired method."); + method!.Invoke(manager, new object[] { clusterHostName }); + } + + [TestMethod] + public void PrimaryBuilder_NeverHasTokenProviderCallback() + { + // The primary builder always uses Kusto.Data's native authentication + // path (WAM/MSAL/SSO) - the host-bridged auth provider is only used + // as a fallback after native auth fails. So even when an + // IAuthenticationProvider is registered, the primary builder must + // not have a TokenProviderCallback installed. + var auth = new CountingAuthProvider(); + var manager = new ConnectionManager(auth); + + var connection = manager.GetOrAddConnection("https://mycluster.kusto.windows.net"); + + var primary = GetBuilderProperty(connection, "PrimaryBuilder"); + Assert.IsNotNull(primary); + Assert.IsNull(primary!.TokenProviderCallback, + "Primary builder must use native authentication, not host-bridged."); + Assert.AreEqual(0, auth.CallCount, + "Provider must not be invoked while the cluster is on the native auth path."); + } + + [TestMethod] + public void FallbackBuilder_AvailableWhenAuthProviderIsSupplied() + { + var auth = new CountingAuthProvider(); + var manager = new ConnectionManager(auth); + + var connection = manager.GetOrAddConnection("https://mycluster.kusto.windows.net"); + + var fallback = GetFallbackBuilder(connection); + Assert.IsNotNull(fallback, "Fallback builder should be available when an IAuthenticationProvider is supplied."); + Assert.IsNotNull(fallback!.TokenProviderCallback, + "Fallback builder must route authentication through the host-supplied provider."); + } + + [TestMethod] + public async Task FallbackBuilder_CallbackInvokesProvider() + { + var auth = new CountingAuthProvider(); + var manager = new ConnectionManager(auth); + + var connection = manager.GetOrAddConnection("https://mycluster.kusto.windows.net"); + var fallback = GetFallbackBuilder(connection); + Assert.IsNotNull(fallback); + + // Simulate Kusto.Data asking the fallback builder for a token. + var token = await fallback!.TokenProviderCallback!(); + + Assert.AreEqual(1, auth.CallCount, "Provider should be called exactly once per token request."); + Assert.AreEqual("https://mycluster.kusto.windows.net", auth.LastClusterUri); + Assert.AreEqual("test-access-token", token); + } + + [TestMethod] + public async Task FallbackBuilder_NullToken_CallbackThrows() + { + var auth = new CountingAuthProvider { Token = null }; + var manager = new ConnectionManager(auth); + + var connection = manager.GetOrAddConnection("https://mycluster.kusto.windows.net"); + var fallback = GetFallbackBuilder(connection); + Assert.IsNotNull(fallback); + + await Assert.ThrowsExactlyAsync( + async () => await fallback!.TokenProviderCallback!()); + Assert.AreEqual(1, auth.CallCount); + } + + [TestMethod] + public void FallbackBuilder_NotAvailableWithoutAuthProvider() + { + var manager = new ConnectionManager(); + + var connection = manager.GetOrAddConnection("https://mycluster.kusto.windows.net"); + + var fallback = GetFallbackBuilder(connection); + Assert.IsNull(fallback, + "Without an IAuthenticationProvider there is no fallback builder."); + } + + [TestMethod] + public void FallbackBuilder_NotAvailableWhenConnectionStringHasExplicitAuth() + { + var auth = new CountingAuthProvider(); + var manager = new ConnectionManager(auth); + + var connection = manager.GetOrAddConnection( + "Data Source=https://mycluster.kusto.windows.net;AppClientId=client;AppKey=secret;Authority Id=tenant"); + + var fallback = GetFallbackBuilder(connection); + Assert.IsNull(fallback, + "Explicit credentials in the connection string suppress the fallback path."); + Assert.AreEqual(0, auth.CallCount); + } + + [TestMethod] + public void ActiveBuilder_SwitchesToFallback_WhenClusterIsMarked() + { + var auth = new CountingAuthProvider(); + var manager = new ConnectionManager(auth); + + var connection = manager.GetOrAddConnection("https://mycluster.kusto.windows.net"); + + // Before marking: the active builder is the primary (native auth). + var beforeBuilder = GetBuilderProperty(connection, "PrimaryBuilder"); + Assert.IsNotNull(beforeBuilder); + var activeBefore = GetBuilderProperty(connection, "Builder"); + Assert.AreSame(beforeBuilder, activeBefore, + "Before fallback is required, the active builder is the primary builder."); + + // Simulate the auth-failure retry path marking the cluster. + MarkFallbackRequired(manager, "mycluster.kusto.windows.net"); + + var fallback = GetFallbackBuilder(connection); + var activeAfter = GetBuilderProperty(connection, "Builder"); + Assert.AreSame(fallback, activeAfter, + "After the cluster is marked, the active builder is the fallback builder."); + } + + #endregion }