refactor: pass company ID in WebSocket code analysis requests

Backup/Conviso.Platform.VisualStudio.sln

@@ -0,0 +1,25 @@
+
+Microsoft Visual Studio Solution File, Format Version 12.00
+# Visual Studio Version 17
+VisualStudioVersion = 17.0.31903.59
+MinimumVisualStudioVersion = 10.0.40219.1
+Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Conviso.Platform.VisualStudio", "src\Conviso.Platform.VisualStudio\Conviso.Platform.VisualStudio.csproj", "{60E0A468-A01B-4C56-9AC9-9B2FCD90B20E}"
+EndProject
+Global
+	GlobalSection(SolutionConfigurationPlatforms) = preSolution
+		Debug|Any CPU = Debug|Any CPU
+		Release|Any CPU = Release|Any CPU
+	EndGlobalSection
+	GlobalSection(ProjectConfigurationPlatforms) = postSolution
+		{60E0A468-A01B-4C56-9AC9-9B2FCD90B20E}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{60E0A468-A01B-4C56-9AC9-9B2FCD90B20E}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{60E0A468-A01B-4C56-9AC9-9B2FCD90B20E}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{60E0A468-A01B-4C56-9AC9-9B2FCD90B20E}.Release|Any CPU.Build.0 = Release|Any CPU
+	EndGlobalSection
+	GlobalSection(SolutionProperties) = preSolution
+		HideSolutionNode = FALSE
+	EndGlobalSection
+	GlobalSection(ExtensibilityGlobals) = postSolution
+		SolutionGuid = {4F23F0B7-B5F8-42D2-A711-F725B2C5B7E2}
+	EndGlobalSection
+EndGlobal

Conviso.Platform.VisualStudio.sln

@@ -1,4 +1,3 @@
-
 Microsoft Visual Studio Solution File, Format Version 12.00
 # Visual Studio Version 17
 VisualStudioVersion = 17.0.31903.59

src/Conviso.Platform.VisualStudio/AssemblyInfo.cs

@@ -5,6 +5,6 @@
 [assembly: AssemblyDescription("Conviso Platform Visual Studio extension")]
 [assembly: AssemblyCompany("Conviso")]
 [assembly: AssemblyProduct("Conviso Platform Visual Studio")]
-[assembly: AssemblyVersion("0.2.1.0")]
-[assembly: AssemblyFileVersion("0.2.1.0")]
+[assembly: AssemblyVersion("0.2.3.0")]
+[assembly: AssemblyFileVersion("0.2.3.0")]
 [assembly: ComVisible(false)]

src/Conviso.Platform.VisualStudio/ConvisoPlatformPackage.cs

@@ -38,7 +38,7 @@ protected override async Task InitializeAsync(CancellationToken cancellationToke
             var settingsService = new SettingsService(this);
             var apiClient = new PlatformApiClient(settingsService);
             var platformFacade = new PlatformFacade(apiClient, settingsService);
-            var brokerClient = new BrokerClient();
+            var brokerClient = new BrokerClient(settingsService);
             var editorContextService = new EditorContextService(this);
             var patchService = new DocumentPatchService(this);
             ToolWindowContext = new ToolWindowContext(settingsService, platformFacade, brokerClient, editorContextService, patchService);

src/Conviso.Platform.VisualStudio/Services/Broker/BrokerClient.cs

@@ -5,26 +5,35 @@
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
+using Conviso.Platform.VisualStudio.Configuration;
 using Conviso.Platform.VisualStudio.Models;
 
 namespace Conviso.Platform.VisualStudio.Services.Broker
 {
     internal sealed class BrokerClient : IBrokerClient
     {
         private const int ConnectTimeoutMilliseconds = 15000;
+        private readonly ISettingsService settingsService;
         private ClientWebSocket? socket;
         private CancellationTokenSource? receiveLoopCancellation;
         private TaskCompletionSource<bool>? authenticationCompletionSource;
+        private volatile bool isAuthenticated;
         private readonly object exclusiveRequestsLock = new object();
         private readonly HashSet<string> exclusiveRequestIds = new HashSet<string>(StringComparer.Ordinal);
         private event Action<BrokerEvent>? InternalEventReceived;
 
         public event Action<BrokerEvent>? EventReceived;
 
-        public bool IsConnected => socket != null && socket.State == WebSocketState.Open;
+        public BrokerClient(ISettingsService settingsService)
+        {
+            this.settingsService = settingsService;
+        }
+
+        public bool IsConnected => isAuthenticated && socket != null && socket.State == WebSocketState.Open;
 
         public async Task ConnectAsync(BrokerConnectionOptions options, CancellationToken cancellationToken)
         {
+            isAuthenticated = false;
             await DisconnectAsync(cancellationToken);
 
             socket = new ClientWebSocket();
@@ -35,7 +44,8 @@ public async Task ConnectAsync(BrokerConnectionOptions options, CancellationToke
                 throw new InvalidOperationException("Missing chat API key.");
             }
 
-            authenticationCompletionSource = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+            var authenticationCompletion = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+            authenticationCompletionSource = authenticationCompletion;
 
             await socket.ConnectAsync(new Uri(endpoint), cancellationToken);
             receiveLoopCancellation = new CancellationTokenSource();
@@ -47,7 +57,7 @@ public async Task ConnectAsync(BrokerConnectionOptions options, CancellationToke
 
             string authRequestId = CreateRequestId("auth");
             await SendMessageAsync(
-                socket,
+                activeSocket,
                 new
                 {
                     type = "auth",
@@ -64,11 +74,11 @@ await SendMessageAsync(
                 cancellationToken,
                 timeoutCancellation.Token);
 
-            using (linkedCancellation.Token.Register(() => authenticationCompletionSource.TrySetCanceled(), useSynchronizationContext: false))
+            using (linkedCancellation.Token.Register(() => authenticationCompletion.TrySetCanceled(), useSynchronizationContext: false))
             {
                 try
                 {
-                    await authenticationCompletionSource.Task;
+                    await authenticationCompletion.Task;
                 }
                 catch (TaskCanceledException) when (timeoutCancellation.IsCancellationRequested)
                 {
@@ -83,14 +93,11 @@ await SendMessageAsync(
 
         public async Task<string> SendChatMessageAsync(ChatMessage message, CancellationToken cancellationToken)
         {
-            if (socket == null || socket.State != WebSocketState.Open)
-            {
-                throw new InvalidOperationException("Broker is not connected.");
-            }
+            ClientWebSocket activeSocket = GetAuthenticatedSocket();
 
             string requestId = CreateRequestId("req");
             await SendMessageAsync(
-                socket,
+                activeSocket,
                 new
                 {
                     type = "analyze_code",
@@ -99,6 +106,7 @@ await SendMessageAsync(
                     {
                         code = message.Content,
                         language = string.IsNullOrWhiteSpace(message.Language) ? "text" : message.Language,
+                        company_id = GetCompanyId(),
                     },
                 },
                 cancellationToken);
@@ -107,10 +115,7 @@ await SendMessageAsync(
 
         public async Task<AutoFixResult> RequestAutoFixAsync(string findingId, CancellationToken cancellationToken)
         {
-            if (socket == null || socket.State != WebSocketState.Open)
-            {
-                throw new InvalidOperationException("Broker is not connected.");
-            }
+            ClientWebSocket activeSocket = GetAuthenticatedSocket();
 
             if (string.IsNullOrWhiteSpace(findingId))
             {
@@ -153,7 +158,7 @@ void HandleEvent(BrokerEvent brokerEvent)
             try
             {
                 await SendMessageAsync(
-                    socket,
+                    activeSocket,
                     new
                     {
                         type = "analyze_code",
@@ -166,6 +171,7 @@ await SendMessageAsync(
                                 "Explain the risk and provide the corrected code in a fenced code block when possible.",
                                 "Vulnerability ID: " + findingId),
                             language = "text",
+                            company_id = GetCompanyId(),
                         },
                     },
                     cancellationToken);
@@ -223,10 +229,7 @@ await SendMessageAsync(
 
         public async Task UpdateExtractorAcceptedAsync(int extractorId, CancellationToken cancellationToken)
         {
-            if (socket == null || socket.State != WebSocketState.Open)
-            {
-                throw new InvalidOperationException("Broker is not connected.");
-            }
+            ClientWebSocket activeSocket = GetAuthenticatedSocket();
 
             if (extractorId <= 0)
             {
@@ -263,7 +266,7 @@ void HandleEvent(BrokerEvent brokerEvent)
             try
             {
                 await SendMessageAsync(
-                    socket,
+                    activeSocket,
                     new
                     {
                         type = "update_extractor",
@@ -302,6 +305,7 @@ await SendMessageAsync(
 
         public async Task DisconnectAsync(CancellationToken cancellationToken)
         {
+            isAuthenticated = false;
             authenticationCompletionSource?.TrySetCanceled();
             authenticationCompletionSource = null;
             receiveLoopCancellation?.Cancel();
@@ -352,6 +356,7 @@ private async Task ReceiveLoopAsync(ClientWebSocket activeSocket, CancellationTo
                     result = await activeSocket.ReceiveAsync(new ArraySegment<byte>(buffer), cancellationToken);
                     if (result.MessageType == WebSocketMessageType.Close)
                     {
+                        isAuthenticated = false;
                         authenticationCompletionSource?.TrySetException(
                             new InvalidOperationException("Chat connection closed before authentication completed."));
                         return;
@@ -377,6 +382,7 @@ private async Task RunReceiveLoopSafelyAsync(ClientWebSocket activeSocket, Cance
             }
             catch (Exception error)
             {
+                isAuthenticated = false;
                 authenticationCompletionSource?.TrySetException(error);
                 Infrastructure.DiagnosticsLogger.LogError("Chat receive loop stopped: " + error);
             }
@@ -390,11 +396,13 @@ private void ProcessIncomingMessage(string raw)
             {
                 if (brokerEvent.Status == "success")
                 {
+                    isAuthenticated = true;
                     authenticationCompletionSource?.TrySetResult(true);
                     authenticationCompletionSource = null;
                 }
                 else
                 {
+                    isAuthenticated = false;
                     authenticationCompletionSource?.TrySetException(
                         new InvalidOperationException(string.IsNullOrWhiteSpace(brokerEvent.Content)
                             ? "Chat authentication failed."
@@ -407,6 +415,7 @@ private void ProcessIncomingMessage(string raw)
 
             if (brokerEvent.Type == "auth_error")
             {
+                isAuthenticated = false;
                 authenticationCompletionSource?.TrySetException(
                     new InvalidOperationException(string.IsNullOrWhiteSpace(brokerEvent.Content)
                         ? "Chat authentication failed."
@@ -415,6 +424,15 @@ private void ProcessIncomingMessage(string raw)
                 return;
             }
 
+            if ((brokerEvent.Type == "error" || brokerEvent.Type == "analysis_error") &&
+                IsAuthenticationFailure(brokerEvent.Content))
+            {
+                // The server can keep the socket open after the authenticated session
+                // expires. Mark it unusable so the next operation performs a full
+                // reconnect and authentication handshake.
+                isAuthenticated = false;
+            }
+
             InternalEventReceived?.Invoke(brokerEvent);
 
             bool isExclusiveRequest;
@@ -495,6 +513,48 @@ private static BrokerEvent ParseEvent(string raw)
             return new BrokerEvent(type, requestId, content, raw, status);
         }
 
+        private static bool IsAuthenticationFailure(string message)
+        {
+            if (string.IsNullOrWhiteSpace(message))
+            {
+                return false;
+            }
+
+            bool mentionsPortugueseAuthentication =
+                (message.IndexOf("conex", StringComparison.OrdinalIgnoreCase) >= 0 ||
+                 message.IndexOf("connection", StringComparison.OrdinalIgnoreCase) >= 0) &&
+                message.IndexOf("autenticad", StringComparison.OrdinalIgnoreCase) >= 0;
+
+            return message.IndexOf("not authenticated", StringComparison.OrdinalIgnoreCase) >= 0 ||
+                   message.IndexOf("unauthenticated", StringComparison.OrdinalIgnoreCase) >= 0 ||
+                   message.IndexOf("authentication required", StringComparison.OrdinalIgnoreCase) >= 0 ||
+                   message.IndexOf("unauthorized", StringComparison.OrdinalIgnoreCase) >= 0 ||
+                   message.IndexOf("nao autentic", StringComparison.OrdinalIgnoreCase) >= 0 ||
+                   mentionsPortugueseAuthentication;
+        }
+
+        private ClientWebSocket GetAuthenticatedSocket()
+        {
+            ClientWebSocket? activeSocket = socket;
+            if (!isAuthenticated || activeSocket == null || activeSocket.State != WebSocketState.Open)
+            {
+                throw new InvalidOperationException("Broker is not authenticated.");
+            }
+
+            return activeSocket;
+        }
+
+        private int GetCompanyId()
+        {
+            string companyId = settingsService.GetString(ConvisoOptions.CompanyIdKey, string.Empty);
+            if (!int.TryParse(companyId, out int numericCompanyId))
+            {
+                throw new InvalidOperationException("Configure a valid numeric Company ID before analyzing code.");
+            }
+
+            return numericCompanyId;
+        }
+
         private static async Task SendMessageAsync(
             ClientWebSocket socket,
             object payload,

src/Conviso.Platform.VisualStudio/ViewModels/SettingsToolWindowViewModel.cs

@@ -129,7 +129,7 @@ private async Task TestApiAsync()
 
         private async Task TestBrokerAsync()
         {
-            var brokerClient = new BrokerClient();
+            var brokerClient = new BrokerClient(settingsService);
             try
             {
                 Status = "Testing broker...";

src/Conviso.Platform.VisualStudio/source.extension.vsixmanifest

@@ -1,7 +1,7 @@
 <?xml version="1.0" encoding="utf-8"?>
 <PackageManifest Version="2.0.0" xmlns="http://schemas.microsoft.com/developer/vsx-schema/2011" xmlns:d="http://schemas.microsoft.com/developer/vsx-schema-design/2011">
   <Metadata>
-    <Identity Id="Conviso.Platform.VisualStudio" Version="0.2.1" Language="en-US" Publisher="Conviso Application Security" />
+    <Identity Id="Conviso.Platform.VisualStudio" Version="0.2.3" Language="en-US" Publisher="Conviso Application Security" />
     <DisplayName>Conviso Platform</DisplayName>
     <Description xml:space="preserve">Conviso Platform integration for Visual Studio.</Description>
     <Icon>Resources\pluginIcon.png</Icon>