Merge pull request #603 from Flow-Launcher/FixJsonRPCBufferExceed

Add a MemoryStream buffer to ReadStream first

Author: jjw24

Flow.Launcher.Core/Flow.Launcher.Core.csproj

@@ -55,6 +55,7 @@
   <ItemGroup>
     <PackageReference Include="Droplex" Version="1.3.1" />
     <PackageReference Include="FSharp.Core" Version="4.7.1" />
+    <PackageReference Include="Microsoft.IO.RecyclableMemoryStream" Version="2.1.3" />
     <PackageReference Include="squirrel.windows" Version="1.5.2" />
   </ItemGroup>
 

Flow.Launcher.Core/Plugin/JsonRPCPlugin.cs

@@ -11,7 +11,9 @@
 using System.Windows.Forms;
 using Flow.Launcher.Infrastructure.Logger;
 using Flow.Launcher.Plugin;
+using ICSharpCode.SharpZipLib.Zip;
 using JetBrains.Annotations;
+using Microsoft.IO;
 
 namespace Flow.Launcher.Core.Plugin
 {
@@ -33,9 +35,11 @@ internal abstract class JsonRPCPlugin : IAsyncPlugin, IContextMenu
         protected abstract string ExecuteCallback(JsonRPCRequestModel rpcRequest);
         protected abstract string ExecuteContextMenu(Result selectedResult);
 
+        private static readonly RecyclableMemoryStreamManager BufferManager = new();
+
         public List<Result> LoadContextMenus(Result selectedResult)
         {
-            string output = ExecuteContextMenu(selectedResult);
+            var output = ExecuteContextMenu(selectedResult);
             try
             {
                 return DeserializedResult(output);
@@ -61,12 +65,23 @@ private async Task<List<Result>> DeserializedResultAsync(Stream output)
         {
             if (output == Stream.Null) return null;
 
-            var queryResponseModel = await
-                JsonSerializer.DeserializeAsync<JsonRPCQueryResponseModel>(output, options);
+            try
+            {
+                var queryResponseModel =
+                    await JsonSerializer.DeserializeAsync<JsonRPCQueryResponseModel>(output, options);
 
-            await output.DisposeAsync();
-            
-            return ParseResults(queryResponseModel);
+                return ParseResults(queryResponseModel);
+            }
+            catch (JsonException e)
+            {
+                Log.Exception(GetType().FullName, "Unexpected Json Input", e);
+            }
+            finally
+            {
+                await output.DisposeAsync();
+            }
+
+            return null;
         }
 
         private List<Result> DeserializedResult(string output)
@@ -81,15 +96,14 @@ private List<Result> DeserializedResult(string output)
 
         private List<Result> ParseResults(JsonRPCQueryResponseModel queryResponseModel)
         {
-            var results = new List<Result>();
             if (queryResponseModel.Result == null) return null;
 
             if (!string.IsNullOrEmpty(queryResponseModel.DebugMessage))
             {
                 context.API.ShowMsg(queryResponseModel.DebugMessage);
             }
 
-            foreach (JsonRPCResult result in queryResponseModel.Result)
+            foreach (var result in queryResponseModel.Result)
             {
                 result.Action = c =>
                 {
@@ -114,7 +128,8 @@ private List<Result> ParseResults(JsonRPCQueryResponseModel queryResponseModel)
                             return !result.JsonRPCAction.DontHideAfterAction;
                         }
 
-                        var jsonRpcRequestModel = JsonSerializer.Deserialize<JsonRPCRequestModel>(actionResponse, options);
+                        var jsonRpcRequestModel =
+                            JsonSerializer.Deserialize<JsonRPCRequestModel>(actionResponse, options);
 
                         if (jsonRpcRequestModel?.Method?.StartsWith("Flow.Launcher.") ?? false)
                         {
@@ -125,9 +140,12 @@ private List<Result> ParseResults(JsonRPCQueryResponseModel queryResponseModel)
 
                     return !result.JsonRPCAction.DontHideAfterAction;
                 };
-                results.Add(result);
             }
 
+            var results = new List<Result>();
+
+            results.AddRange(queryResponseModel.Result);
+
             return results;
         }
 
@@ -217,16 +235,42 @@ protected string Execute(ProcessStartInfo startInfo)
 
         protected async Task<Stream> ExecuteAsync(ProcessStartInfo startInfo, CancellationToken token = default)
         {
+            Process process = null;
+            bool disposed = false;
             try
             {
-                using var process = Process.Start(startInfo);
+                process = Process.Start(startInfo);
                 if (process == null)
                 {
                     Log.Error("|JsonRPCPlugin.ExecuteAsync|Can't start new process");
                     return Stream.Null;
                 }
 
-                var result = process.StandardOutput.BaseStream;
+                await using var source = process.StandardOutput.BaseStream;
+
+                var buffer = BufferManager.GetStream();
+                
+                token.Register(() =>
+                {
+                    // ReSharper disable once AccessToModifiedClosure
+                    // Manually Check whether disposed
+                    if (!disposed && !process.HasExited)
+                        process.Kill();
+                });
+
+                try
+                {
+                    // token expire won't instantly trigger the exception, 
+                    // manually kill process at before
+                    await source.CopyToAsync(buffer, token);
+                }
+                catch (OperationCanceledException)
+                {
+                    await buffer.DisposeAsync();
+                    return Stream.Null;
+                }
+
+                buffer.Seek(0, SeekOrigin.Begin);
 
                 token.ThrowIfCancellationRequested();
 
@@ -245,7 +289,7 @@ protected async Task<Stream> ExecuteAsync(ProcessStartInfo startInfo, Cancellati
                     return Stream.Null;
                 }
 
-                return result;
+                return buffer;
             }
             catch (Exception e)
             {
@@ -254,15 +298,24 @@ protected async Task<Stream> ExecuteAsync(ProcessStartInfo startInfo, Cancellati
                     e);
                 return Stream.Null;
             }
+            finally
+            {
+                process?.Dispose();
+                disposed = true;
+            }
         }
 
         public async Task<List<Result>> QueryAsync(Query query, CancellationToken token)
         {
-            var output = await ExecuteQueryAsync(query, token);
             try
             {
+                var output = await ExecuteQueryAsync(query, token);
                 return await DeserializedResultAsync(output);
             }
+            catch (OperationCanceledException)
+            {
+                return null;
+            }
             catch (Exception e)
             {
                 Log.Exception($"|JsonRPCPlugin.Query|Exception when query <{query}>", e);