263 lines
11 KiB
C#
263 lines
11 KiB
C#
// Copyright Epic Games, Inc. All Rights Reserved.
|
|
|
|
using System;
|
|
using System.Collections;
|
|
using System.Collections.Concurrent;
|
|
using System.Collections.Generic;
|
|
using System.IO;
|
|
using System.Runtime.InteropServices;
|
|
|
|
namespace EpicGames.UBA.Impl
|
|
{
|
|
internal class SessionServerCreateInfoImpl : ISessionServerCreateInfo
|
|
{
|
|
nint _handle = IntPtr.Zero;
|
|
readonly IStorageServer _storage;
|
|
readonly IServer _client;
|
|
readonly ILogger _logger;
|
|
|
|
#region DllImport
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern nint SessionServerCreateInfo_Create(nint storage, nint client, nint logger, string rootDir, string traceOutputFile,
|
|
byte disableCustomAllocator, byte launchVisualizer, byte resetCas, byte writeToDisk, byte detailedTrace, byte allowWaitOnMem, byte allowKillOnMem, byte storeObjFilesCompressed);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern void SessionServerCreateInfo_Destroy(nint server);
|
|
#endregion
|
|
|
|
public SessionServerCreateInfoImpl(IStorageServer storage, IServer client, ILogger logger, SessionServerCreateInfo info)
|
|
{
|
|
_storage = storage;
|
|
_client = client;
|
|
_logger = logger;
|
|
_handle = SessionServerCreateInfo_Create(_storage.GetHandle(), _client.GetHandle(), _logger.GetHandle(), info.RootDirectory, info.TraceOutputFile, (byte)(info.DisableCustomAllocator?1:0), (byte)(info.LaunchVisualizer?1:0), (byte)(info.ResetCas?1:0), (byte)(info.WriteToDisk?1:0), (byte)(info.DetailedTrace?1:0), (byte)(info.AllowWaitOnMem?1:0), (byte)(info.AllowKillOnMem?1:0), (byte)(info.StoreObjFilesCompressed?1:0));
|
|
}
|
|
|
|
#region IDisposable
|
|
~SessionServerCreateInfoImpl() => Dispose(false);
|
|
|
|
public void Dispose()
|
|
{
|
|
Dispose(true);
|
|
GC.SuppressFinalize(this);
|
|
}
|
|
|
|
protected virtual void Dispose(bool disposing)
|
|
{
|
|
if (disposing)
|
|
{
|
|
}
|
|
|
|
if (_handle != IntPtr.Zero)
|
|
{
|
|
SessionServerCreateInfo_Destroy(_handle);
|
|
_handle = IntPtr.Zero;
|
|
}
|
|
}
|
|
#endregion
|
|
|
|
#region ISessionServerCreateInfo
|
|
public nint GetHandle() => _handle;
|
|
#endregion
|
|
}
|
|
|
|
internal class SessionServerImpl : ISessionServer
|
|
{
|
|
delegate void RemoteProcessSlotAvailableCallback(nint userData, byte isCrossArchitecture);
|
|
delegate void RemoteProcessReturnedCallback(nint handle);
|
|
|
|
nint _handle = IntPtr.Zero;
|
|
readonly ISessionServerCreateInfo _info;
|
|
readonly RemoteProcessSlotAvailableCallback _remoteProcessSlotAvailableCallbackDelegate;
|
|
readonly RemoteProcessReturnedCallback _remoteProcessReturnedCallbackDelegate;
|
|
readonly ConcurrentDictionary<ulong, IProcess> _remoteProcesses = new();
|
|
|
|
#region DllImport
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern nint SessionServer_Create(nint info, byte[] environment, uint environmentSize);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern void SessionServer_SetRemoteProcessAvailable(nint server, RemoteProcessSlotAvailableCallback func, nint userData);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern void SessionServer_SetRemoteProcessReturned(nint server, RemoteProcessReturnedCallback func);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern void SessionServer_RefreshDirectory(nint server, string directory);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern void SessionServer_RegisterNewFile(nint server, string filename);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern byte SessionServer_RegisterVirtualFile(nint server, string filename, string sourceFile, ulong sourceOffset, ulong sourceSize);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern uint SessionServer_BeginExternalProcess(nint server, string description);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern void SessionServer_EndExternalProcess(nint server, uint id, uint exitCode);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern void SessionServer_UpdateProgress(nint server, uint processesTotal, uint processesDone, uint errorCount);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern void SessionServer_UpdateStatus(nint server, uint statusRow, uint statusColumn, string statusText, byte statusType, string? statusLink);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern nint SessionServer_RunProcess(nint server, nint info, bool async, bool enableDetour);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern nint SessionServer_RunProcessRemote(nint server, nint info, float weight, byte[]? knownInputs, uint knownInputsCount, bool allowCrossArchitecture);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern ulong SessionServer_RegisterRoots(nint server, byte[] rootsData, uint rootsDataSize);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern void SessionServer_SetMaxRemoteProcessCount(nint server, uint count);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern void SessionServer_DisableRemoteExecution(nint server);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern void SessionServer_SetCustomCasKeyFromTrackedInputs(nint server, nint process, string filename, string workingdir);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern void SessionServer_RegisterCrossArchitectureMapping(nint server, string from, string to);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern void SessionServer_PrintSummary(nint server);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern void SessionServer_CancelAll(nint server);
|
|
|
|
[DllImport("UbaHost", CharSet = CharSet.Auto)]
|
|
static extern void SessionServer_Destroy(nint server);
|
|
#endregion
|
|
|
|
public SessionServerImpl(ISessionServerCreateInfo info)
|
|
{
|
|
_info = info;
|
|
_remoteProcessSlotAvailableCallbackDelegate = RaiseRemoteProcessSlotAvailable;
|
|
_remoteProcessReturnedCallbackDelegate = RaiseRemoteProcessReturned;
|
|
|
|
// We need to manually transfer environment variables on non-windows platforms since they are not automatically propagated from c# to native.
|
|
using MemoryStream environmentMemory = new();
|
|
{
|
|
if (!RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
|
|
{
|
|
using (BinaryWriter writer = new(environmentMemory, System.Text.Encoding.UTF8, true))
|
|
{
|
|
foreach (DictionaryEntry de in Environment.GetEnvironmentVariables())
|
|
{
|
|
writer.Write($"{de.Key}={de.Value}");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
_handle = SessionServer_Create(_info.GetHandle(), environmentMemory.GetBuffer(), (uint)environmentMemory.Position);
|
|
SessionServer_SetRemoteProcessAvailable(_handle, _remoteProcessSlotAvailableCallbackDelegate, 0);
|
|
SessionServer_SetRemoteProcessReturned(_handle, _remoteProcessReturnedCallbackDelegate);
|
|
|
|
foreach (KeyValuePair<string, string> kv in Utils.CrossArchitecturePaths)
|
|
{
|
|
SessionServer_RegisterCrossArchitectureMapping(_handle, kv.Key, kv.Value);
|
|
}
|
|
}
|
|
|
|
#region IDisposable
|
|
~SessionServerImpl() => Dispose(false);
|
|
|
|
public void Dispose()
|
|
{
|
|
Dispose(true);
|
|
GC.SuppressFinalize(this);
|
|
}
|
|
|
|
protected virtual void Dispose(bool disposing)
|
|
{
|
|
if (disposing)
|
|
{
|
|
}
|
|
|
|
if (_handle != IntPtr.Zero)
|
|
{
|
|
SessionServer_Destroy(_handle);
|
|
_handle = IntPtr.Zero;
|
|
}
|
|
}
|
|
#endregion
|
|
|
|
#region ISessionServer
|
|
public nint GetHandle() => _handle;
|
|
|
|
public event ISessionServer.RemoteProcessSlotAvailableEventHandler? RemoteProcessSlotAvailable;
|
|
|
|
public event ISessionServer.RemoteProcessReturnedEventHandler? RemoteProcessReturned;
|
|
|
|
public IProcess RunProcess(ProcessStartInfo info, bool async, IProcess.ExitedEventHandler? exitedEventHandler, bool enableDetour)
|
|
{
|
|
IProcessStartInfo startInfo = IProcessStartInfo.CreateProcessStartInfo(info, exitedEventHandler != null);
|
|
nint processPtr = SessionServer_RunProcess(_handle, startInfo.GetHandle(), async, enableDetour);
|
|
IProcess process = IProcess.CreateProcess(processPtr, startInfo, exitedEventHandler, info.UserData);
|
|
return process;
|
|
}
|
|
|
|
public IProcess RunProcessRemote(ProcessStartInfo info, IProcess.ExitedEventHandler? exitedEventHandler, double weight, byte[]? knownInputs, uint knownInputsCount, bool canExecuteCrossArchitecture)
|
|
{
|
|
IProcessStartInfo startInfo = IProcessStartInfo.CreateProcessStartInfo(info, exitedEventHandler != null);
|
|
nint processPtr = SessionServer_RunProcessRemote(_handle, startInfo.GetHandle(), (float)weight, knownInputs, knownInputsCount, canExecuteCrossArchitecture);
|
|
IProcess process = IProcess.CreateProcess(processPtr, startInfo, exitedEventHandler, info.UserData);
|
|
_remoteProcesses.AddOrUpdate(process.Hash, process, (k, v) => process);
|
|
return process;
|
|
}
|
|
|
|
public ulong RegisterRoots(byte[] rootsData, uint rootsDataSize) => SessionServer_RegisterRoots(_handle, rootsData, rootsDataSize);
|
|
|
|
public void DisableRemoteExecution() => SessionServer_DisableRemoteExecution(_handle);
|
|
|
|
public void SetMaxRemoteProcessCount(uint count) => SessionServer_SetMaxRemoteProcessCount(_handle, count);
|
|
|
|
public void RefreshDirectories(params string[] directories) => Array.ForEach(directories, (directory) => SessionServer_RefreshDirectory(_handle, directory));
|
|
|
|
public void RegisterNewFiles(params string[] files) => Array.ForEach(files, (file) => SessionServer_RegisterNewFile(_handle, file));
|
|
public bool RegisterVirtualFile(string name, string sourceFile, ulong sourceOffset, ulong sourceSize) => SessionServer_RegisterVirtualFile(_handle, name, sourceFile, sourceOffset, sourceSize) != 0;
|
|
|
|
public uint BeginExternalProcess(string description) => SessionServer_BeginExternalProcess(_handle, description);
|
|
|
|
public void EndExternalProcess(uint id, uint exitCode) => SessionServer_EndExternalProcess(_handle, id, exitCode);
|
|
|
|
public void UpdateProgress(uint processesTotal, uint processesDone, uint errorCount) => SessionServer_UpdateProgress(_handle, processesTotal, processesDone, errorCount);
|
|
|
|
public void UpdateStatus(uint statusRow, uint statusColumn, string statusText, LogEntryType statusType, string? statusLink) => SessionServer_UpdateStatus(_handle, statusRow, statusColumn, statusText, (byte)statusType, statusLink);
|
|
|
|
public void SetCustomCasKeyFromTrackedInputs(string file, string workingDirectory, IProcess process) => SessionServer_SetCustomCasKeyFromTrackedInputs(_handle, process.GetHandle(), file, workingDirectory);
|
|
|
|
public void PrintSummary() => SessionServer_PrintSummary(_handle);
|
|
|
|
public void CancelAll() => SessionServer_CancelAll(_handle);
|
|
#endregion
|
|
|
|
void RaiseRemoteProcessSlotAvailable(nint userData, byte isCrossArchitecture)
|
|
{
|
|
RemoteProcessSlotAvailable?.Invoke(this, new RemoteProcessSlotAvailableEventArgs(isCrossArchitecture != 0));
|
|
}
|
|
|
|
void RaiseRemoteProcessReturned(nint handle)
|
|
{
|
|
IProcess? process = null;
|
|
try
|
|
{
|
|
if (_remoteProcesses.Remove((ulong)handle.ToInt64(), out process))
|
|
{
|
|
RemoteProcessReturned?.Invoke(this, new RemoteProcessReturnedEventArgs(process));
|
|
}
|
|
}
|
|
finally
|
|
{
|
|
process?.Dispose();
|
|
}
|
|
}
|
|
}
|
|
}
|