Skip to content

Improve allocations in NegotiateStreamPal #71280

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.

Sign up for GitHub

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

Merged
wfurt merged 5 commits into from
Jun 29, 2022
Merged

Improve allocations in NegotiateStreamPal #71280

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
1dd747d
Reduce buffer allocations during NTLM/Negotiate authentication
filipnavara Jun 25, 2022
49c1ed0
Update ReadWriteAdapter.WriteAsync prototype to use Memory<byte> inst…
filipnavara Jun 25, 2022
2fe4491
Spanify NTAuthentication.Decrypt and avoid couple of offset/count checks
filipnavara Jun 25, 2022
3b4aff6
Spanify NegotiateStreamPal.VerifySignature/MakeSignature.
filipnavara Jun 25, 2022
46e05a4
Update src/libraries/Common/src/System/Net/NTAuthentication.Common.cs
filipnavara Jun 28, 2022
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
19 changes: 8 additions & 11 deletions ...libraries/Common/src/Interop/Unix/System.Net.Security.Native/Interop.NetSecurityNative.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -210,10 +210,10 @@ private static unsafe partial Status Wrap(
ref GssBuffer outBuffer);

[LibraryImport(Interop.Libraries.NetSecurityNative, EntryPoint="NetSecurityNative_Unwrap")]
private static partial Status Unwrap(
private static unsafe partial Status Unwrap(
out Status minorStatus,
SafeGssContextHandle? contextHandle,
byte[] inputBytes,
byte* inputBytes,
Copy link
Member

@wfurt wfurt Jun 27, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need offset any more? It feels like if we pass pointer we can just do count.

Copy link
Member Author

@filipnavara filipnavara Jun 27, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't. Unfortunately, dotnet/sqlclient uses the native APIs, so I didn't feel confident in changing it. I am quite sure they don't use this particular API though.

Copy link
Contributor

@MichalPetryka MichalPetryka Jun 27, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wouldn't the pointer change by itself be breaking for them?

Copy link
Member Author

@filipnavara filipnavara Jun 27, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It doesn't change the native code side, the prototype remains unchanged there. They basically copied the managed side of the interop. (We cannot change the signature, add parameters, remove parameters, or change their types. We can change how they are marshalled on the C# side though.)

Copy link
Member

@wfurt wfurt Jun 27, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc: @JRahnama just in case.

Copy link
Member Author

@filipnavara filipnavara Jun 29, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I cross-checked with the SqlClient source and this particular native method is not referenced so we can remove the offset parameter. I will do that in a follow-up PR (#71373) since I need to update the native interop there anyway.

int offset,
int count,
ref GssBuffer outBuffer);
Expand All @@ -231,19 +231,16 @@ internal static unsafe Status WrapBuffer(
}
}

internal static Status UnwrapBuffer(
internal static unsafe Status UnwrapBuffer(
out Status minorStatus,
SafeGssContextHandle? contextHandle,
byte[] inputBytes,
int offset,
int count,
ReadOnlySpan<byte> inputBytes,
ref GssBuffer outBuffer)
{
Debug.Assert(inputBytes != null, "inputBytes must be valid value");
Debug.Assert(offset >= 0 && offset <= inputBytes.Length, "offset must be valid");
Debug.Assert(count >= 0 && count <= inputBytes.Length, "count must be valid");

return Unwrap(out minorStatus, contextHandle, inputBytes, offset, count, ref outBuffer);
fixed (byte* inputBytesPtr = inputBytes)
{
return Unwrap(out minorStatus, contextHandle, inputBytesPtr, 0, inputBytes.Length, ref outBuffer);
}
}
}
}
2 changes: 1 addition & 1 deletion src/libraries/Common/src/Interop/Windows/SspiCli/ISSPIInterface.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ internal interface ISSPIInterface
int QueryContextChannelBinding(SafeDeleteContext phContext, Interop.SspiCli.ContextAttribute attribute, out SafeFreeContextBufferChannelBinding refHandle);
int QueryContextAttributes(SafeDeleteContext phContext, Interop.SspiCli.ContextAttribute attribute, Span<byte> buffer, Type? handleType, out SafeHandle? refHandle);
int QuerySecurityContextToken(SafeDeleteContext phContext, out SecurityContextTokenHandle phToken);
int CompleteAuthToken(ref SafeDeleteSslContext? refContext, in SecurityBuffer inputBuffer);
int CompleteAuthToken(ref SafeDeleteSslContext? refContext, in InputSecurityBuffer inputBuffer);
int ApplyControlToken(ref SafeDeleteContext? refContext, in SecurityBuffer inputBuffer);
}
}
2 changes: 1 addition & 1 deletion src/libraries/Common/src/Interop/Windows/SspiCli/SSPIAuthType.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -168,7 +168,7 @@ public int QuerySecurityContextToken(SafeDeleteContext phContext, out SecurityCo
return GetSecurityContextToken(phContext, out phToken);
}

public int CompleteAuthToken(ref SafeDeleteSslContext? refContext, in SecurityBuffer inputBuffer)
public int CompleteAuthToken(ref SafeDeleteSslContext? refContext, in InputSecurityBuffer inputBuffer)
{
return SafeDeleteContext.CompleteAuthToken(ref refContext, in inputBuffer);
}
Expand Down
2 changes: 1 addition & 1 deletion src/libraries/Common/src/Interop/Windows/SspiCli/SSPISecureChannelType.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,7 @@ public int QuerySecurityContextToken(SafeDeleteContext phContext, out SecurityCo
throw new NotSupportedException();
}

public int CompleteAuthToken(ref SafeDeleteSslContext? refContext, in SecurityBuffer inputBuffer)
public int CompleteAuthToken(ref SafeDeleteSslContext? refContext, in InputSecurityBuffer inputBuffer)
{
throw new NotSupportedException();
}
Expand Down
143 changes: 1 addition & 142 deletions src/libraries/Common/src/Interop/Windows/SspiCli/SSPIWrapper.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,7 @@ internal static int AcceptSecurityContext(ISSPIInterface secModule, SafeFreeCred
return errorCode;
}

internal static int CompleteAuthToken(ISSPIInterface secModule, ref SafeDeleteSslContext? context, in SecurityBuffer inputBuffer)
internal static int CompleteAuthToken(ISSPIInterface secModule, ref SafeDeleteSslContext? context, in InputSecurityBuffer inputBuffer)
{
int errorCode = secModule.CompleteAuthToken(ref context, in inputBuffer);

Expand All @@ -186,147 +186,6 @@ public static int QuerySecurityContextToken(ISSPIInterface secModule, SafeDelete
return secModule.QuerySecurityContextToken(context, out token);
}

public static int EncryptMessage(ISSPIInterface secModule, SafeDeleteContext context, Span<SecurityBuffer> input, uint sequenceNumber)
{
return EncryptDecryptHelper(OP.Encrypt, secModule, context, input, sequenceNumber);
}

public static int DecryptMessage(ISSPIInterface secModule, SafeDeleteContext context, Span<SecurityBuffer> input, uint sequenceNumber)
{
return EncryptDecryptHelper(OP.Decrypt, secModule, context, input, sequenceNumber);
}

internal static int MakeSignature(ISSPIInterface secModule, SafeDeleteContext context, Span<SecurityBuffer> input, uint sequenceNumber)
{
return EncryptDecryptHelper(OP.MakeSignature, secModule, context, input, sequenceNumber);
}

public static int VerifySignature(ISSPIInterface secModule, SafeDeleteContext context, Span<SecurityBuffer> input, uint sequenceNumber)
{
return EncryptDecryptHelper(OP.VerifySignature, secModule, context, input, sequenceNumber);
}

private enum OP
{
Encrypt = 1,
Decrypt,
MakeSignature,
VerifySignature
}

[StructLayout(LayoutKind.Sequential)]
private ref struct ThreeByteArrays
{
public const int NumItems = 3;
internal byte[] _item0;
private byte[] _item1;
private byte[] _item2;
}

private static unsafe int EncryptDecryptHelper(OP op, ISSPIInterface secModule, SafeDeleteContext context, Span<SecurityBuffer> input, uint sequenceNumber)
{
Debug.Assert(Enum.IsDefined<OP>(op), $"Unknown op: {op}");
Debug.Assert(input.Length <= 3, "The below logic only works for 3 or fewer buffers.");

Interop.SspiCli.SecBufferDesc sdcInOut = new Interop.SspiCli.SecBufferDesc(input.Length);
Span<Interop.SspiCli.SecBuffer> unmanagedBuffer = stackalloc Interop.SspiCli.SecBuffer[input.Length];
unmanagedBuffer.Clear();

fixed (Interop.SspiCli.SecBuffer* unmanagedBufferPtr = unmanagedBuffer)
fixed (byte* pinnedBuffer0 = input.Length > 0 ? input[0].token : null)
fixed (byte* pinnedBuffer1 = input.Length > 1 ? input[1].token : null)
fixed (byte* pinnedBuffer2 = input.Length > 2 ? input[2].token : null)
{
sdcInOut.pBuffers = unmanagedBufferPtr;

ThreeByteArrays byteArrayStruct = default;
Span<byte[]> buffers = MemoryMarshal.CreateSpan(ref byteArrayStruct._item0!, ThreeByteArrays.NumItems).Slice(0, input.Length);

for (int i = 0; i < input.Length; i++)
{
ref readonly SecurityBuffer iBuffer = ref input[i];
unmanagedBuffer[i].cbBuffer = iBuffer.size;
unmanagedBuffer[i].BufferType = iBuffer.type;
if (iBuffer.token == null || iBuffer.token.Length == 0)
{
unmanagedBuffer[i].pvBuffer = IntPtr.Zero;
}
else
{
unmanagedBuffer[i].pvBuffer = Marshal.UnsafeAddrOfPinnedArrayElement(iBuffer.token, iBuffer.offset);
buffers[i] = iBuffer.token;
}
}

// The result is written in the input Buffer passed as type=BufferType.Data.
int errorCode = op switch
{
OP.Encrypt => secModule.EncryptMessage(context, ref sdcInOut, sequenceNumber),
OP.Decrypt => secModule.DecryptMessage(context, ref sdcInOut, sequenceNumber),
OP.MakeSignature => secModule.MakeSignature(context, ref sdcInOut, sequenceNumber),
_ /* OP.VerifySignature */ => secModule.VerifySignature(context, ref sdcInOut, sequenceNumber),
};

// Marshalling back returned sizes / data.
for (int i = 0; i < input.Length; i++)
{
ref SecurityBuffer iBuffer = ref input[i];
iBuffer.size = unmanagedBuffer[i].cbBuffer;
iBuffer.type = unmanagedBuffer[i].BufferType;

if (iBuffer.size == 0)
{
iBuffer.offset = 0;
iBuffer.token = null;
}
else
{

// Find the buffer this is inside of. Usually they all point inside buffer 0.
int j;
for (j = 0; j < input.Length; j++)
{
if (buffers[j] != null)
{
checked
{
byte* bufferAddress = (byte*)Marshal.UnsafeAddrOfPinnedArrayElement(buffers[j], 0);
if ((byte*)unmanagedBuffer[i].pvBuffer >= bufferAddress &&
(byte*)unmanagedBuffer[i].pvBuffer + iBuffer.size <= bufferAddress + buffers[j].Length)
{
iBuffer.offset = (int)((byte*)unmanagedBuffer[i].pvBuffer - bufferAddress);
iBuffer.token = buffers[j];
break;
}
}
}
}

if (j >= input.Length)
{
Debug.Fail("Output buffer out of range.");
iBuffer.size = 0;
iBuffer.offset = 0;
iBuffer.token = null;
}
}

// Backup validate the new sizes.
Debug.Assert(iBuffer.offset >= 0 && iBuffer.offset <= (iBuffer.token == null ? 0 : iBuffer.token.Length), $"'offset' out of range. [{iBuffer.offset}]");
Debug.Assert(iBuffer.size >= 0 && iBuffer.size <= (iBuffer.token == null ? 0 : iBuffer.token.Length - iBuffer.offset), $"'size' out of range. [{iBuffer.size}]");
}

if (NetEventSource.Log.IsEnabled() && errorCode != 0)
{
NetEventSource.Error(null, errorCode == Interop.SspiCli.SEC_I_RENEGOTIATE ?
SR.Format(SR.event_OperationReturnedSomething, op, "SEC_I_RENEGOTIATE") :
SR.Format(SR.net_log_operation_failed_with_error, op, $"0x{0:X}"));
}

return errorCode;
}
}

public static SafeFreeContextBufferChannelBinding? QueryContextChannelBinding(ISSPIInterface secModule, SafeDeleteContext securityContext, Interop.SspiCli.ContextAttribute contextAttribute)
{
SafeFreeContextBufferChannelBinding result;
Expand Down
27 changes: 14 additions & 13 deletions src/libraries/Common/src/Interop/Windows/SspiCli/SecuritySafeHandles.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -472,9 +472,13 @@ internal static unsafe int InitializeSecurityContext(
// Get unmanaged buffer with index 0 as the only one passed into PInvoke.
outSecBuffer.size = outUnmanagedBuffer.cbBuffer;
outSecBuffer.type = outUnmanagedBuffer.BufferType;
outSecBuffer.token = outSecBuffer.size > 0 ?
new Span<byte>((byte*)outUnmanagedBuffer.pvBuffer, outUnmanagedBuffer.cbBuffer).ToArray() :
null;

if (isSspiAllocated)
{
outSecBuffer.token = outSecBuffer.size > 0 ?
new Span<byte>((byte*)outUnmanagedBuffer.pvBuffer, outUnmanagedBuffer.cbBuffer).ToArray() :
null;
}

if (inSecBuffers.Count > 1 && inUnmanagedBuffer[1].BufferType == SecurityBufferType.SECBUFFER_EXTRA && inSecBuffers._item1.Type == SecurityBufferType.SECBUFFER_EMPTY)
{
Expand Down Expand Up @@ -952,25 +956,22 @@ private static unsafe int MustRunAcceptSecurityContext_SECURITY(

internal static unsafe int CompleteAuthToken(
ref SafeDeleteSslContext? refContext,
in SecurityBuffer inSecBuffer)
in InputSecurityBuffer inSecBuffer)
{
if (NetEventSource.Log.IsEnabled()) NetEventSource.Info(null, $"refContext = {refContext}, inSecBuffer = {inSecBuffer}");
if (NetEventSource.Log.IsEnabled()) NetEventSource.Info(null, $"refContext = {refContext}");

var inSecurityBufferDescriptor = new Interop.SspiCli.SecBufferDesc(1);
int errorCode = (int)Interop.SECURITY_STATUS.InvalidHandle;

Interop.SspiCli.SecBuffer inUnmanagedBuffer = default;
inSecurityBufferDescriptor.pBuffers = &inUnmanagedBuffer;
fixed (byte* pinnedToken = inSecBuffer.token)
fixed (byte* pinnedToken = inSecBuffer.Token)
{
inUnmanagedBuffer.cbBuffer = inSecBuffer.size;
inUnmanagedBuffer.BufferType = inSecBuffer.type;

// Use the unmanaged token if it's not null; otherwise use the managed buffer.
Debug.Assert(inSecBuffer.UnmanagedToken != null);
inUnmanagedBuffer.cbBuffer = inSecBuffer.Token.Length;
inUnmanagedBuffer.BufferType = inSecBuffer.Type;
inUnmanagedBuffer.pvBuffer =
inSecBuffer.unmanagedToken != null ? inSecBuffer.unmanagedToken.DangerousGetHandle() :
inSecBuffer.token == null || inSecBuffer.token.Length == 0 ? IntPtr.Zero :
(IntPtr)(pinnedToken + inSecBuffer.offset);
inSecBuffer.Token.IsEmpty ? IntPtr.Zero : (IntPtr)pinnedToken;

Interop.SspiCli.CredHandle contextHandle = refContext != null ? refContext._handle : default;
if (refContext == null || refContext.IsInvalid)
Expand Down
33 changes: 21 additions & 12 deletions src/libraries/Common/src/System/Net/NTAuthentication.Common.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ internal sealed partial class NTAuthentication
private string? _spn;

private int _tokenSize;
private byte[]? _tokenBuffer;
private ContextFlagsPal _requestedContextFlags;
private ContextFlagsPal _contextFlags;

Expand Down Expand Up @@ -158,14 +159,14 @@ internal void CloseContext()
_isCompleted = false;
}

internal int VerifySignature(byte[] buffer, int offset, int count)
internal int VerifySignature(ReadOnlySpan<byte> buffer)
{
return NegotiateStreamPal.VerifySignature(_securityContext!, buffer, offset, count);
return NegotiateStreamPal.VerifySignature(_securityContext!, buffer);
}

internal int MakeSignature(byte[] buffer, int offset, int count, [AllowNull] ref byte[] output)
internal int MakeSignature(ReadOnlySpan<byte> buffer, [AllowNull] ref byte[] output)
{
return NegotiateStreamPal.MakeSignature(_securityContext!, buffer, offset, count, ref output);
return NegotiateStreamPal.MakeSignature(_securityContext!, buffer, ref output);
}

internal string? GetOutgoingBlob(string? incomingBlob)
Expand Down Expand Up @@ -221,9 +222,10 @@ internal int MakeSignature(byte[] buffer, int offset, int count, [AllowNull] ref

internal byte[]? GetOutgoingBlob(ReadOnlySpan<byte> incomingBlob, bool throwOnError, out SecurityStatusPal statusCode)
{
byte[]? result = new byte[_tokenSize];
_tokenBuffer ??= _tokenSize == 0 ? Array.Empty<byte>() : new byte[_tokenSize];

bool firstTime = _securityContext == null;
int resultBlobLength;
try
{
if (!_isServer)
Expand All @@ -236,18 +238,19 @@ internal int MakeSignature(byte[] buffer, int offset, int count, [AllowNull] ref
_requestedContextFlags,
incomingBlob,
_channelBinding,
ref result,
ref _tokenBuffer,
out resultBlobLength,
ref _contextFlags);

if (NetEventSource.Log.IsEnabled()) NetEventSource.Info(this, $"SSPIWrapper.InitializeSecurityContext() returns statusCode:0x{((int)statusCode.ErrorCode):x8} ({statusCode})");

if (statusCode.ErrorCode == SecurityStatusPalErrorCode.CompleteNeeded)
{
statusCode = NegotiateStreamPal.CompleteAuthToken(ref _securityContext, result);
statusCode = NegotiateStreamPal.CompleteAuthToken(ref _securityContext, _tokenBuffer.AsSpan(0, resultBlobLength));

if (NetEventSource.Log.IsEnabled()) NetEventSource.Info(this, $"SSPIWrapper.CompleteAuthToken() returns statusCode:0x{((int)statusCode.ErrorCode):x8} ({statusCode})");

result = null;
resultBlobLength = 0;
}
}
else
Expand All @@ -259,7 +262,8 @@ internal int MakeSignature(byte[] buffer, int offset, int count, [AllowNull] ref
_requestedContextFlags,
incomingBlob,
_channelBinding,
ref result,
ref _tokenBuffer,
out resultBlobLength,
ref _contextFlags);

if (NetEventSource.Log.IsEnabled()) NetEventSource.Info(this, $"SSPIWrapper.AcceptSecurityContext() returns statusCode:0x{((int)statusCode.ErrorCode):x8} ({statusCode})");
Expand All @@ -284,6 +288,7 @@ internal int MakeSignature(byte[] buffer, int offset, int count, [AllowNull] ref
{
CloseContext();
_isCompleted = true;
_tokenBuffer = null;
if (throwOnError)
{
throw NegotiateStreamPal.CreateExceptionFromError(statusCode);
Expand All @@ -297,12 +302,18 @@ internal int MakeSignature(byte[] buffer, int offset, int count, [AllowNull] ref
SSPIHandleCache.CacheCredential(_credentialsHandle);
}

byte[]? result =
resultBlobLength == 0 || _tokenBuffer == null ? null :
_tokenBuffer.Length == resultBlobLength ? _tokenBuffer :
_tokenBuffer[0..resultBlobLength];

// The return value will tell us correctly if the handshake is over or not
if (statusCode.ErrorCode == SecurityStatusPalErrorCode.OK
|| (_isServer && statusCode.ErrorCode == SecurityStatusPalErrorCode.CompleteNeeded))
{
// Success.
_isCompleted = true;
_tokenBuffer = null;
}
else
{
Expand Down Expand Up @@ -335,13 +346,11 @@ internal int Encrypt(ReadOnlySpan<byte> buffer, [NotNull] ref byte[]? output, ui
sequenceNumber);
}

internal int Decrypt(byte[] payload, int offset, int count, out int newOffset, uint expectedSeqNumber)
internal int Decrypt(Span<byte> payload, out int newOffset, uint expectedSeqNumber)
{
return NegotiateStreamPal.Decrypt(
_securityContext!,
payload,
offset,
count,
(_contextFlags & ContextFlagsPal.Confidentiality) != 0,
IsNTLM,
out newOffset,
Expand Down
Loading