Merge pull request #603 from Flow-Launcher/FixJsonRPCBufferExceed

Add a MemoryStream buffer to ReadStream first
This commit is contained in:
Jeremy Wu 2021-07-27 19:10:26 +10:00 committed by GitHub
commit 27777956cc
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 68 additions and 14 deletions

View file

@ -55,6 +55,7 @@
<ItemGroup> <ItemGroup>
<PackageReference Include="Droplex" Version="1.3.1" /> <PackageReference Include="Droplex" Version="1.3.1" />
<PackageReference Include="FSharp.Core" Version="4.7.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" /> <PackageReference Include="squirrel.windows" Version="1.5.2" />
</ItemGroup> </ItemGroup>

View file

@ -11,7 +11,9 @@ using System.Threading.Tasks;
using System.Windows.Forms; using System.Windows.Forms;
using Flow.Launcher.Infrastructure.Logger; using Flow.Launcher.Infrastructure.Logger;
using Flow.Launcher.Plugin; using Flow.Launcher.Plugin;
using ICSharpCode.SharpZipLib.Zip;
using JetBrains.Annotations; using JetBrains.Annotations;
using Microsoft.IO;
namespace Flow.Launcher.Core.Plugin namespace Flow.Launcher.Core.Plugin
{ {
@ -33,9 +35,11 @@ namespace Flow.Launcher.Core.Plugin
protected abstract string ExecuteCallback(JsonRPCRequestModel rpcRequest); protected abstract string ExecuteCallback(JsonRPCRequestModel rpcRequest);
protected abstract string ExecuteContextMenu(Result selectedResult); protected abstract string ExecuteContextMenu(Result selectedResult);
private static readonly RecyclableMemoryStreamManager BufferManager = new();
public List<Result> LoadContextMenus(Result selectedResult) public List<Result> LoadContextMenus(Result selectedResult)
{ {
string output = ExecuteContextMenu(selectedResult); var output = ExecuteContextMenu(selectedResult);
try try
{ {
return DeserializedResult(output); return DeserializedResult(output);
@ -61,12 +65,23 @@ namespace Flow.Launcher.Core.Plugin
{ {
if (output == Stream.Null) return null; if (output == Stream.Null) return null;
var queryResponseModel = await try
JsonSerializer.DeserializeAsync<JsonRPCQueryResponseModel>(output, options); {
var queryResponseModel =
await JsonSerializer.DeserializeAsync<JsonRPCQueryResponseModel>(output, options);
await output.DisposeAsync(); return ParseResults(queryResponseModel);
}
catch (JsonException e)
{
Log.Exception(GetType().FullName, "Unexpected Json Input", e);
}
finally
{
await output.DisposeAsync();
}
return ParseResults(queryResponseModel); return null;
} }
private List<Result> DeserializedResult(string output) private List<Result> DeserializedResult(string output)
@ -81,7 +96,6 @@ namespace Flow.Launcher.Core.Plugin
private List<Result> ParseResults(JsonRPCQueryResponseModel queryResponseModel) private List<Result> ParseResults(JsonRPCQueryResponseModel queryResponseModel)
{ {
var results = new List<Result>();
if (queryResponseModel.Result == null) return null; if (queryResponseModel.Result == null) return null;
if (!string.IsNullOrEmpty(queryResponseModel.DebugMessage)) if (!string.IsNullOrEmpty(queryResponseModel.DebugMessage))
@ -89,7 +103,7 @@ namespace Flow.Launcher.Core.Plugin
context.API.ShowMsg(queryResponseModel.DebugMessage); context.API.ShowMsg(queryResponseModel.DebugMessage);
} }
foreach (JsonRPCResult result in queryResponseModel.Result) foreach (var result in queryResponseModel.Result)
{ {
result.Action = c => result.Action = c =>
{ {
@ -114,7 +128,8 @@ namespace Flow.Launcher.Core.Plugin
return !result.JsonRPCAction.DontHideAfterAction; 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) if (jsonRpcRequestModel?.Method?.StartsWith("Flow.Launcher.") ?? false)
{ {
@ -125,9 +140,12 @@ namespace Flow.Launcher.Core.Plugin
return !result.JsonRPCAction.DontHideAfterAction; return !result.JsonRPCAction.DontHideAfterAction;
}; };
results.Add(result);
} }
var results = new List<Result>();
results.AddRange(queryResponseModel.Result);
return results; return results;
} }
@ -217,16 +235,42 @@ namespace Flow.Launcher.Core.Plugin
protected async Task<Stream> ExecuteAsync(ProcessStartInfo startInfo, CancellationToken token = default) protected async Task<Stream> ExecuteAsync(ProcessStartInfo startInfo, CancellationToken token = default)
{ {
Process process = null;
bool disposed = false;
try try
{ {
using var process = Process.Start(startInfo); process = Process.Start(startInfo);
if (process == null) if (process == null)
{ {
Log.Error("|JsonRPCPlugin.ExecuteAsync|Can't start new process"); Log.Error("|JsonRPCPlugin.ExecuteAsync|Can't start new process");
return Stream.Null; 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(); token.ThrowIfCancellationRequested();
@ -245,7 +289,7 @@ namespace Flow.Launcher.Core.Plugin
return Stream.Null; return Stream.Null;
} }
return result; return buffer;
} }
catch (Exception e) catch (Exception e)
{ {
@ -254,15 +298,24 @@ namespace Flow.Launcher.Core.Plugin
e); e);
return Stream.Null; return Stream.Null;
} }
finally
{
process?.Dispose();
disposed = true;
}
} }
public async Task<List<Result>> QueryAsync(Query query, CancellationToken token) public async Task<List<Result>> QueryAsync(Query query, CancellationToken token)
{ {
var output = await ExecuteQueryAsync(query, token);
try try
{ {
var output = await ExecuteQueryAsync(query, token);
return await DeserializedResultAsync(output); return await DeserializedResultAsync(output);
} }
catch (OperationCanceledException)
{
return null;
}
catch (Exception e) catch (Exception e)
{ {
Log.Exception($"|JsonRPCPlugin.Query|Exception when query <{query}>", e); Log.Exception($"|JsonRPCPlugin.Query|Exception when query <{query}>", e);