Optimize `AssemblyLoadContext.LoadFromStream` and fix partial reads. (#72783)

src/coreclr/nativeaot/System.Private.CoreLib/src/Internal/Reflection/Augments/ReflectionAugments.cs

@@ -121,7 +121,7 @@ internal static ReflectionCoreCallbacks ReflectionCoreCallbacks
     public abstract class ReflectionCoreCallbacks
     {
         public abstract Assembly Load(AssemblyName refName, bool throwOnFileNotFound);
-        public abstract Assembly Load(byte[] rawAssembly, byte[] pdbSymbolStore);
+        public abstract Assembly Load(ReadOnlySpan<byte> rawAssembly, ReadOnlySpan<byte> pdbSymbolStore);
         public abstract Assembly Load(string assemblyPath);
 
         public abstract MethodBase GetMethodFromHandle(RuntimeMethodHandle runtimeMethodHandle);

src/coreclr/nativeaot/System.Private.CoreLib/src/Internal/Reflection/Core/AssemblyBinder.cs

@@ -36,7 +36,7 @@ public abstract class AssemblyBinder
 
         public abstract bool Bind(RuntimeAssemblyName refName, bool cacheMissedLookups, out AssemblyBindResult result, out Exception exception);
 
-        public abstract bool Bind(byte[] rawAssembly, byte[] rawSymbolStore, out AssemblyBindResult result, out Exception exception);
+        public abstract bool Bind(ReadOnlySpan<byte> rawAssembly, ReadOnlySpan<byte> rawSymbolStore, out AssemblyBindResult result, out Exception exception);
 
         public abstract bool Bind(string assemblyPath, out AssemblyBindResult bindResult, out Exception exception);
 

src/coreclr/nativeaot/System.Private.CoreLib/src/System/Reflection/Runtime/General/Dispensers.cs

@@ -42,7 +42,7 @@ internal static RuntimeAssembly GetRuntimeAssembly(RuntimeAssemblyName assemblyR
         /// <summary>
         /// Returns non-null or throws.
         /// </summary>
-        internal static RuntimeAssembly GetRuntimeAssemblyFromByteArray(byte[] rawAssembly, byte[] pdbSymbolStore)
+        internal static RuntimeAssembly GetRuntimeAssemblyFromByteArray(ReadOnlySpan<byte> rawAssembly, ReadOnlySpan<byte> pdbSymbolStore)
         {
             AssemblyBinder binder = ReflectionCoreExecution.ExecutionDomain.ReflectionDomainSetup.AssemblyBinder;
             if (!binder.Bind(rawAssembly, pdbSymbolStore, out AssemblyBindResult bindResult, out Exception exception))

src/coreclr/nativeaot/System.Private.CoreLib/src/System/Reflection/Runtime/General/ReflectionCoreCallbacksImplementation.cs

@@ -40,9 +40,9 @@ public sealed override Assembly Load(AssemblyName assemblyRef, bool throwOnFileN
                 return RuntimeAssemblyInfo.GetRuntimeAssemblyIfExists(assemblyRef.ToRuntimeAssemblyName());
         }
 
-        public sealed override Assembly Load(byte[] rawAssembly, byte[] pdbSymbolStore)
+        public sealed override Assembly Load(ReadOnlySpan<byte> rawAssembly, ReadOnlySpan<byte> pdbSymbolStore)
         {
-            if (rawAssembly == null)
+            if (rawAssembly.IsEmpty)
                 throw new ArgumentNullException(nameof(rawAssembly));
 
             return RuntimeAssemblyInfo.GetRuntimeAssemblyFromByteArray(rawAssembly, pdbSymbolStore);

src/coreclr/nativeaot/System.Private.CoreLib/src/System/Runtime/Loader/AssemblyLoadContext.NativeAot.cs

@@ -48,7 +48,7 @@ private static Assembly InternalLoadFromPath(string? assemblyPath, string? nativ
         }
 
 #pragma warning disable CA1822
-        internal Assembly InternalLoad(byte[] arrAssembly, byte[] arrSymbols)
+        internal Assembly InternalLoad(ReadOnlySpan<byte> arrAssembly, ReadOnlySpan<byte> arrSymbols)
         {
             return ReflectionAugments.ReflectionCoreCallbacks.Load(arrAssembly, arrSymbols);
         }

src/coreclr/nativeaot/System.Private.DisabledReflection/src/Internal/Reflection/ReflectionCoreCallbacksImplementation.cs

@@ -47,7 +47,7 @@ public override DynamicInvokeInfo GetDelegateDynamicInvokeInfo(Type type)
         public override Type GetTypeFromCLSID(Guid clsid, string server, bool throwOnError) => throw new NotSupportedException(SR.Reflection_Disabled);
 #endif
         public override Assembly Load(AssemblyName refName, bool throwOnFileNotFound) => throw new NotSupportedException(SR.Reflection_Disabled);
-        public override Assembly Load(byte[] rawAssembly, byte[] pdbSymbolStore) => throw new NotSupportedException(SR.Reflection_Disabled);
+        public override Assembly Load(ReadOnlySpan<byte> rawAssembly, ReadOnlySpan<byte> pdbSymbolStore) => throw new NotSupportedException(SR.Reflection_Disabled);
         public override Assembly Load(string assemblyPath) => throw new NotSupportedException(SR.Reflection_Disabled);
         public override void MakeTypedReference(object target, FieldInfo[] flds, out Type type, out int offset) => throw new NotSupportedException(SR.Reflection_Disabled);
         public override void RunModuleConstructor(Module module) => throw new NotSupportedException(SR.Reflection_Disabled);

src/coreclr/nativeaot/System.Private.TypeLoader/src/Internal/Reflection/Execution/AssemblyBinderImplementation.cs

@@ -34,7 +34,7 @@ private AssemblyBinderImplementation()
         public static AssemblyBinderImplementation Instance { get; } = new AssemblyBinderImplementation();
 
         partial void BindEcmaFilePath(string assemblyPath, ref AssemblyBindResult bindResult, ref Exception exception, ref bool? result);
-        partial void BindEcmaByteArray(byte[] rawAssembly, byte[] rawSymbolStore, ref AssemblyBindResult bindResult, ref Exception exception, ref bool? result);
+        partial void BindEcmaBytes(ReadOnlySpan<byte> rawAssembly, ReadOnlySpan<byte> rawSymbolStore, ref AssemblyBindResult bindResult, ref Exception exception, ref bool? result);
         partial void BindEcmaAssemblyName(RuntimeAssemblyName refName, bool cacheMissedLookups, ref AssemblyBindResult result, ref Exception exception, ref Exception preferredException, ref bool resultBoolean);
         partial void InsertEcmaLoadedAssemblies(List<AssemblyBindResult> loadedAssemblies);
 
@@ -53,13 +53,13 @@ public sealed override bool Bind(string assemblyPath, out AssemblyBindResult bin
                 return result.Value;
         }
 
-        public sealed override bool Bind(byte[] rawAssembly, byte[] rawSymbolStore, out AssemblyBindResult bindResult, out Exception exception)
+        public sealed override bool Bind(ReadOnlySpan<byte> rawAssembly, ReadOnlySpan<byte> rawSymbolStore, out AssemblyBindResult bindResult, out Exception exception)
         {
             bool? result = null;
             exception = null;
             bindResult = default(AssemblyBindResult);
 
-            BindEcmaByteArray(rawAssembly, rawSymbolStore, ref bindResult, ref exception, ref result);
+            BindEcmaBytes(rawAssembly, rawSymbolStore, ref bindResult, ref exception, ref result);
 
             // If the Ecma assembly binder isn't linked in, simply throw PlatformNotSupportedException
             if (!result.HasValue)

src/libraries/System.Private.CoreLib/src/System/Runtime/Loader/AssemblyLoadContext.cs

@@ -374,34 +374,54 @@ public Assembly LoadFromStream(Stream assembly, Stream? assemblySymbols)
         {
             ArgumentNullException.ThrowIfNull(assembly);
 
-            int iAssemblyStreamLength = (int)assembly.Length;
-
-            if (iAssemblyStreamLength <= 0)
+            ReadOnlySpan<byte> spanAssembly = ReadAllBytes(assembly);
+            if (spanAssembly.IsEmpty)
             {
                 throw new BadImageFormatException(SR.BadImageFormat_BadILFormat);
             }
 
-            // Allocate the byte[] to hold the assembly
-            byte[] arrAssembly = new byte[iAssemblyStreamLength];
-
-            // Copy the assembly to the byte array
-            assembly.Read(arrAssembly, 0, iAssemblyStreamLength);
-
-            // Get the symbol stream in byte[] if provided
-            byte[]? arrSymbols = null;
+            // Read the symbol stream if provided
+            ReadOnlySpan<byte> spanSymbols = default;
             if (assemblySymbols != null)
             {
-                int iSymbolLength = (int)assemblySymbols.Length;
-                arrSymbols = new byte[iSymbolLength];
-
-                assemblySymbols.Read(arrSymbols, 0, iSymbolLength);
+                spanSymbols = ReadAllBytes(assemblySymbols);
             }
 
             lock (_unloadLock)
             {
                 VerifyIsAlive();
 
-                return InternalLoad(arrAssembly, arrSymbols);
+                return InternalLoad(spanAssembly, spanSymbols);
+            }
+
+            static ReadOnlySpan<byte> ReadAllBytes(Stream stream)
+            {
+                if (stream.GetType() == typeof(MemoryStream) && ((MemoryStream)stream).TryGetBuffer(out ArraySegment<byte> memoryStreamBuffer))
+                {
+                    int position = (int)stream.Position;
+                    // Simulate that we read the stream to its end.
+                    stream.Seek(0, SeekOrigin.End);
+                    return memoryStreamBuffer.AsSpan(position);
+                }
+
+                long length = stream.Length - stream.Position;
+
+                if (length == 0)
+                {
+                    return ReadOnlySpan<byte>.Empty;
+                }
+
+                if (((ulong)length) > (ulong)Array.MaxLength)
+                {
+                    throw new BadImageFormatException(SR.BadImageFormat_BadILFormat);
+                }
+
+                byte[] bytes = GC.AllocateUninitializedArray<byte>((int)length);
+
+                // Copy the stream to the byte array
+                stream.ReadExactly(bytes);
+
+                return bytes;
             }
         }
 

src/mono/System.Private.CoreLib/src/System/Runtime/Loader/AssemblyLoadContext.Mono.cs

@@ -47,15 +47,14 @@ private Assembly InternalLoadFromPath(string? assemblyPath, string? nativeImageP
 #pragma warning restore IDE0060
 
         [RequiresUnreferencedCode("Types and members the loaded assembly depends on might be removed")]
-        internal Assembly InternalLoad(byte[] arrAssembly, byte[]? arrSymbols)
+        internal Assembly InternalLoad(ReadOnlySpan<byte> arrAssembly, ReadOnlySpan<byte> arrSymbols)
         {
             unsafe
             {
-                int symbolsLength = arrSymbols?.Length ?? 0;
                 fixed (byte* ptrAssembly = arrAssembly, ptrSymbols = arrSymbols)
                 {
                     return InternalLoadFromStream(NativeALC, new IntPtr(ptrAssembly), arrAssembly.Length,
-                                       new IntPtr(ptrSymbols), symbolsLength);
+                                       new IntPtr(ptrSymbols), arrSymbols.Length);
                 }
             }
         }