Skip to content

Use NativeMemory in ManagedPSEntry #25817

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: master
Choose a base branch
from
Draft
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
137 changes: 61 additions & 76 deletions src/powershell/Program.cs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
#nullable enable

using System;
using System.Diagnostics;
using System.IO;
using System.Management.Automation;
using System.Reflection;
Expand Down Expand Up @@ -81,7 +82,7 @@ public static int Main(string[] args)
/// In the event of success, we use an exec() call, so this method never returns.
/// </summary>
/// <param name="args">The startup arguments to pwsh.</param>
private static void AttemptExecPwshLogin(string[] args)
private static unsafe void AttemptExecPwshLogin(string[] args)
{
// If the login environment variable is set, we have already done the login logic and have been exec'd
if (Environment.GetEnvironmentVariable(LOGIN_ENV_VAR_NAME) != null)
Expand All @@ -90,16 +91,14 @@ private static void AttemptExecPwshLogin(string[] args)
return;
}

bool isLinux = Platform.IsLinux;

// The first byte (ASCII char) of the name of this process, used to detect '-' for login
byte procNameFirstByte;

// The path to the executable this process was started from
string? pwshPath;

// On Linux, we can simply use the /proc filesystem
if (isLinux)
if (Platform.IsLinux)
{
// Read the process name byte
using (FileStream fs = File.OpenRead("/proc/self/cmdline"))
Expand All @@ -114,104 +113,82 @@ private static void AttemptExecPwshLogin(string[] args)
}

// Read the symlink to the startup executable
IntPtr linkPathPtr = Marshal.AllocHGlobal(LINUX_PATH_MAX);
IntPtr bufSize = ReadLink("/proc/self/exe", linkPathPtr, (UIntPtr)LINUX_PATH_MAX);
pwshPath = Marshal.PtrToStringAnsi(linkPathPtr, (int)bufSize);
Marshal.FreeHGlobal(linkPathPtr);

ArgumentNullException.ThrowIfNull(pwshPath);
byte* linkPathPtr = (byte*)NativeMemory.Alloc(LINUX_PATH_MAX);
nint len = ReadLink("/proc/self/exe", linkPathPtr, LINUX_PATH_MAX);
pwshPath = new string((sbyte*)linkPathPtr, 0, checked((int)len));
NativeMemory.Free(linkPathPtr);

// exec pwsh
ThrowOnFailure("exec", ExecPwshLogin(args, pwshPath, isMacOS: false));
return;
}

// At this point, we are on macOS
Debug.Assert(Platform.IsMacOS);

// Set up the mib array and the query for process maximum args size
Span<int> mib = stackalloc int[3];
int mibLength = 2;
mib[0] = MACOS_CTL_KERN;
mib[1] = MACOS_KERN_ARGMAX;
int size = IntPtr.Size / 2;
int argmax = 0;
Span<int> mib = [MACOS_CTL_KERN, MACOS_KERN_ARGMAX];
nuint argmax = 0;
nuint size = (uint)sizeof(nuint);

// Get the process args size
unsafe
{
fixed (int *mibptr = mib)
{
ThrowOnFailure(nameof(argmax), SysCtl(mibptr, mibLength, &argmax, &size, IntPtr.Zero, 0));
}
}
ThrowOnFailure(nameof(argmax), SysCtl(mib, &argmax, &size, null, 0));

// Get the PID so we can query this process' args
int pid = GetPid();

// The following logic is based on https://gist.github.com/nonowarn/770696

// Now read the process args into the allocated space
IntPtr procargs = Marshal.AllocHGlobal(argmax);
IntPtr executablePathPtr = IntPtr.Zero;
byte* procargs = (byte*)NativeMemory.Alloc(argmax);
byte* executablePathPtr = null;
try
{
mib[0] = MACOS_CTL_KERN;
mib[1] = MACOS_KERN_PROCARGS2;
mib[2] = pid;
mibLength = 3;

unsafe
{
fixed (int *mibptr = mib)
{
ThrowOnFailure(nameof(procargs), SysCtl(mibptr, mibLength, procargs.ToPointer(), &argmax, IntPtr.Zero, 0));
}

// The memory block we're reading is a series of null-terminated strings
// that looks something like this:
//
// | argc | <int>
// | exec_path | ... \0
// | argv[0] | ... \0
// | argv[1] | ... \0
// ...
//
// We care about argv[0], since that's the name the process was started with.
// If argv[0][0] == '-', we have been invoked as login.
// Doing this, the buffer we populated also recorded `exec_path`,
// which is the path to our executable `pwsh`.
// We can reuse this value later to prevent needing to call a .NET API
// to generate our exec invocation.

// We don't care about argc's value, since argv[0] must always exist.
// Skip over argc, but remember where exec_path is for later
executablePathPtr = IntPtr.Add(procargs, sizeof(int));

// Skip over exec_path
byte *argvPtr = (byte *)executablePathPtr;
while (*argvPtr != 0) { argvPtr++; }
while (*argvPtr == 0) { argvPtr++; }

// First char in argv[0]
procNameFirstByte = *argvPtr;
}
mib = [MACOS_CTL_KERN, MACOS_KERN_PROCARGS2, pid];

ThrowOnFailure(nameof(procargs), SysCtl(mib, procargs, &argmax, null, 0));

// The memory block we're reading is a series of null-terminated strings
// that looks something like this:
//
// | argc | <int>
// | exec_path | ... \0
// | argv[0] | ... \0
// | argv[1] | ... \0
// ...
//
// We care about argv[0], since that's the name the process was started with.
// If argv[0][0] == '-', we have been invoked as login.
// Doing this, the buffer we populated also recorded `exec_path`,
// which is the path to our executable `pwsh`.
// We can reuse this value later to prevent needing to call a .NET API
// to generate our exec invocation.

// We don't care about argc's value, since argv[0] must always exist.
// Skip over argc, but remember where exec_path is for later
executablePathPtr = procargs + sizeof(int);

// Skip over exec_path
byte* argvPtr = executablePathPtr;
while (*argvPtr != 0) { argvPtr++; }
while (*argvPtr == 0) { argvPtr++; }

// First char in argv[0]
procNameFirstByte = *argvPtr;

if (!IsLogin(procNameFirstByte, args))
{
return;
}

// Get the pwshPath from exec_path
pwshPath = Marshal.PtrToStringAnsi(executablePathPtr);

ArgumentNullException.ThrowIfNull(pwshPath);
pwshPath = new string((sbyte*)executablePathPtr);

// exec pwsh
ThrowOnFailure("exec", ExecPwshLogin(args, pwshPath, isMacOS: true));
}
finally
{
Marshal.FreeHGlobal(procargs);
NativeMemory.Free(procargs);
}
}

Expand Down Expand Up @@ -451,14 +428,14 @@ private static void ThrowOnFailure(string call, int code)
/// </summary>
/// <param name="pathname">The path to the symlink to read.</param>
/// <param name="buf">Pointer to a buffer to fill with the result.</param>
/// <param name="size">The size of the buffer we have supplied.</param>
/// <param name="bufsiz">The size of the buffer we have supplied.</param>
/// <returns>The number of bytes placed in the buffer.</returns>
[DllImport("libc",
EntryPoint = "readlink",
CallingConvention = CallingConvention.Cdecl,
CharSet = CharSet.Ansi,
SetLastError = true)]
private static extern IntPtr ReadLink(string pathname, IntPtr buf, UIntPtr size);
private static extern unsafe nint ReadLink(string pathname, byte* buf, nuint bufsiz);

/// <summary>
/// The `getpid` POSIX syscall we use to quickly get the current process PID on macOS.
Expand Down Expand Up @@ -488,19 +465,27 @@ private static void ThrowOnFailure(string call, int code)
/// <summary>
/// The `sysctl` BSD sycall used to get system information on macOS.
/// </summary>
/// <param name="mib">The Management Information Base name, used to query information.</param>
/// <param name="mibLength">The length of the MIB name.</param>
/// <param name="name">The Management Information Base name, used to query information.</param>
/// <param name="namelen">The length of the MIB name.</param>
/// <param name="oldp">The object passed out of sysctl (may be null)</param>
/// <param name="oldlenp">The size of the object passed out of sysctl.</param>
/// <param name="newp">The object passed in to sysctl.</param>
/// <param name="newlenp">The length of the object passed in to sysctl.</param>
/// <param name="newlen">The length of the object passed in to sysctl.</param>
/// <returns></returns>
[DllImport("libc",
EntryPoint = "sysctl",
CallingConvention = CallingConvention.Cdecl,
CharSet = CharSet.Ansi,
SetLastError = true)]
private static extern unsafe int SysCtl(int *mib, int mibLength, void *oldp, int *oldlenp, IntPtr newp, int newlenp);
private static extern unsafe int SysCtl(int* name, uint namelen, void* oldp, nuint* oldlenp, void* newp, nuint newlen);

private static unsafe int SysCtl(Span<int> mib, void* oldp, nuint* oldlenp, void* newp, nuint newlen)
{
fixed (int* name = &MemoryMarshal.GetReference(mib))
{
return SysCtl(name, (uint)mib.Length, oldp, oldlenp, newp, newlen);
}
}
#endif
}
}