[NRBF] Address issues discovered by Threat Model

src/libraries/System.Formats.Nrbf/ref/System.Formats.Nrbf.cs

@@ -11,6 +11,7 @@ public abstract partial class ArrayRecord : System.Formats.Nrbf.SerializationRec
         internal ArrayRecord() { }
         public override System.Formats.Nrbf.SerializationRecordId Id { get { throw null; } }
         public abstract System.ReadOnlySpan<int> Lengths { get; }
+        public virtual long TotalElementsCount { get; }
         public int Rank { get { throw null; } }
         [System.Diagnostics.CodeAnalysis.RequiresDynamicCode("The code for an array of the specified type might not be available.")]
         public System.Array GetArray(System.Type expectedArrayType, bool allowNulls = true) { throw null; }

src/libraries/System.Formats.Nrbf/src/Resources/Strings.resx

@@ -136,13 +136,13 @@
     <value>Member reference was pointing to a record of unexpected type.</value>
   </data>
   <data name="Serialization_InvalidTypeName" xml:space="preserve">
-    <value>Invalid type name: `{0}`.</value>
+    <value>Invalid type name.</value>
   </data>
   <data name="Serialization_TypeMismatch" xml:space="preserve">
     <value>Expected the array to be of type {0}, but its element type was {1}.</value>
   </data>
   <data name="Serialization_InvalidTypeOrAssemblyName" xml:space="preserve">
-    <value>Invalid type or assembly name: `{0},{1}`.</value>
+    <value>Invalid type or assembly name.</value>
   </data>
   <data name="Serialization_DuplicateMemberName" xml:space="preserve">
     <value>Duplicate member name: `{0}`.</value>
@@ -160,6 +160,6 @@
     <value>Only arrays with zero offsets are supported.</value>
   </data>
   <data name="Serialization_InvalidAssemblyName" xml:space="preserve">
-    <value>Invalid assembly name: `{0}`.</value>
+    <value>Invalid assembly name.</value>
   </data>
 </root>

src/libraries/System.Formats.Nrbf/src/System/Formats/Nrbf/ArrayRecord.cs

@@ -27,6 +27,12 @@ private protected ArrayRecord(ArrayInfo arrayInfo)
     /// <value>A buffer of integers that represent the number of elements in every dimension.</value>
     public abstract ReadOnlySpan<int> Lengths { get; }
 
+    /// <summary>
+    /// When overridden in a derived class, gets the total number of all elements in every dimension.
+    /// </summary>
+    /// <value>A number that represent the total number of all elements in every dimension.</value>
+    public virtual long TotalElementsCount => ArrayInfo.TotalElementsCount;
+
     /// <summary>
     /// Gets the rank of the array.
     /// </summary>
@@ -44,7 +50,7 @@ private protected ArrayRecord(ArrayInfo arrayInfo)
 
     internal long ValuesToRead { get; private protected set; }
 
-    private protected ArrayInfo ArrayInfo { get; }
+    internal ArrayInfo ArrayInfo { get; }
 
     /// <summary>
     /// Allocates an array and fills it with the data provided in the serialized records (in case of primitive types like <see cref="string"/> or <see cref="int"/>) or the serialized records themselves.

src/libraries/System.Formats.Nrbf/src/System/Formats/Nrbf/BinaryArrayRecord.cs

@@ -6,6 +6,7 @@
 using System.IO;
 using System.Reflection.Metadata;
 using System.Formats.Nrbf.Utils;
+using System.Diagnostics;
 
 namespace System.Formats.Nrbf;
 
@@ -27,19 +28,96 @@ internal sealed class BinaryArrayRecord : ArrayRecord
     ];
 
     private TypeName? _typeName;
+    private long _totalElementsCount;
 
     private BinaryArrayRecord(ArrayInfo arrayInfo, MemberTypeInfo memberTypeInfo)
         : base(arrayInfo)
     {
         MemberTypeInfo = memberTypeInfo;
         Values = [];
+
+        // We need to parse all elements of the jagged array to obtain total elements count;
+        _totalElementsCount = arrayInfo.ArrayType != BinaryArrayType.Jagged ? arrayInfo.TotalElementsCount : -1;
     }
 
     public override SerializationRecordType RecordType => SerializationRecordType.BinaryArray;
 
     /// <inheritdoc/>
     public override ReadOnlySpan<int> Lengths => new int[1] { Length };
 
+    /// <inheritdoc/>
+    public override long TotalElementsCount
+    {
+        get
+        {
+            if (_totalElementsCount < 0)
+            {
+                _totalElementsCount = GetJaggedArrayTotalElementsCount(this);
+            }
+
+            return _totalElementsCount;
+        }
+    }
+
+    private static long GetJaggedArrayTotalElementsCount(BinaryArrayRecord jaggedArrayRecord)
+    {
+        long result = 0;
+        Queue<BinaryArrayRecord>? jaggedArrayRecords = null;
+
+        do
+        {
+            if (jaggedArrayRecords is not null)
+            {
+                jaggedArrayRecord = jaggedArrayRecords.Dequeue();
+            }
+
+            Debug.Assert(jaggedArrayRecord.ArrayInfo.ArrayType == BinaryArrayType.Jagged);
+
+            foreach (object value in jaggedArrayRecord.Values)
+            {
+                object item = value is MemberReferenceRecord referenceRecord
+                    ? referenceRecord.GetReferencedRecord()
+                    : value;
+
+                if (item is not SerializationRecord record)
+                {
+                    result++;
+                    continue;
+                }
+
+                switch (record.RecordType)
+                {
+                    case SerializationRecordType.BinaryArray:
+                    case SerializationRecordType.ArraySinglePrimitive:
+                    case SerializationRecordType.ArraySingleObject:
+                    case SerializationRecordType.ArraySingleString:
+                        ArrayRecord nestedArrayRecord = (ArrayRecord)record;
+                        if (nestedArrayRecord.ArrayInfo.ArrayType == BinaryArrayType.Jagged)
+                        {
+                            (jaggedArrayRecords ??= new()).Enqueue((BinaryArrayRecord)nestedArrayRecord);
+                        }
+                        else
+                        {
+                            result += nestedArrayRecord.TotalElementsCount;
+                        }
+                        break;
+                    case SerializationRecordType.ObjectNull:
+                    case SerializationRecordType.ObjectNullMultiple256:
+                    case SerializationRecordType.ObjectNullMultiple:
+                        // Null Records nested inside jagged array do not increase total elements count.
+                        // Example: "int[][] input = [[1, 2, 3], null]" is just 3 elements in total.
+                        break;
+                    default:
+                        result++;
+                        break;
+                }
+            }
+        }
+        while (jaggedArrayRecords is not null && jaggedArrayRecords.Count > 0);
+
+        return result;
+    }
+
     public override TypeName TypeName
         => _typeName ??= MemberTypeInfo.GetArrayTypeName(ArrayInfo);
 

src/libraries/System.Formats.Nrbf/src/System/Formats/Nrbf/BinaryLibraryRecord.cs

@@ -57,7 +57,7 @@ internal static BinaryLibraryRecord Decode(BinaryReader reader, PayloadOptions o
         }
         else if (!options.UndoTruncatedTypeNames)
         {
-            ThrowHelper.ThrowInvalidAssemblyName(rawName);
+            ThrowHelper.ThrowInvalidAssemblyName();
         }
 
         return new BinaryLibraryRecord(id, rawName);

src/libraries/System.Formats.Nrbf/src/System/Formats/Nrbf/Utils/ThrowHelper.cs

@@ -6,6 +6,8 @@
 
 namespace System.Formats.Nrbf.Utils;
 
+// The exception messages do not contain type or assembly names on purpose,
+// as it's most likely corrupted/tampered/malicious data.
 internal static class ThrowHelper
 {
     internal static void ThrowInvalidValue(object value)
@@ -14,8 +16,8 @@ internal static void ThrowInvalidValue(object value)
     internal static void ThrowInvalidReference()
         => throw new SerializationException(SR.Serialization_InvalidReference);
 
-    internal static void ThrowInvalidTypeName(string name)
-        => throw new SerializationException(SR.Format(SR.Serialization_InvalidTypeName, name));
+    internal static void ThrowInvalidTypeName()
+        => throw new SerializationException(SR.Serialization_InvalidTypeName);
 
     internal static void ThrowUnexpectedNullRecordCount()
         => throw new SerializationException(SR.Serialization_UnexpectedNullRecordCount);
@@ -26,8 +28,8 @@ internal static void ThrowMaxArrayLength(long limit, long actual)
     internal static void ThrowArrayContainedNulls()
         => throw new SerializationException(SR.Serialization_ArrayContainedNulls);
 
-    internal static void ThrowInvalidAssemblyName(string rawName)
-        => throw new SerializationException(SR.Format(SR.Serialization_InvalidAssemblyName, rawName));
+    internal static void ThrowInvalidAssemblyName()
+        => throw new SerializationException(SR.Serialization_InvalidAssemblyName);
 
     internal static void ThrowEndOfStreamException()
         => throw new EndOfStreamException();

src/libraries/System.Formats.Nrbf/src/System/Formats/Nrbf/Utils/TypeNameHelpers.cs

@@ -125,7 +125,7 @@ internal static TypeName ParseNonSystemClassRecordTypeName(this string rawName,
 
         if (typeName is null)
         {
-            throw new SerializationException(SR.Format(SR.Serialization_InvalidTypeOrAssemblyName, rawName, libraryRecord.RawLibraryName));
+            throw new SerializationException(SR.Serialization_InvalidTypeOrAssemblyName);
         }
 
         if (typeName.AssemblyName is null)
@@ -169,7 +169,7 @@ private static TypeName With(this TypeName typeName, AssemblyNameInfo assemblyNa
             else
             {
                 // BinaryFormatter can not serialize pointers or references.
-                ThrowHelper.ThrowInvalidTypeName(typeName.FullName);
+                ThrowHelper.ThrowInvalidTypeName();
             }
         }
 

src/libraries/System.Formats.Nrbf/tests/ArraySinglePrimitiveRecordTests.cs

@@ -94,6 +94,7 @@ private void Test<T>(int size, bool canSeek)
         SZArrayRecord<T> arrayRecord = (SZArrayRecord<T>)NrbfDecoder.Decode(stream);
 
         Assert.Equal(size, arrayRecord.Length);
+        Assert.Equal(size, arrayRecord.TotalElementsCount);
         T?[] output = arrayRecord.GetArray();
         Assert.Equal(input, output);
         Assert.Same(output, arrayRecord.GetArray());

src/libraries/System.Formats.Nrbf/tests/JaggedArraysTests.cs

@@ -19,6 +19,19 @@ public void CanReadJaggedArraysOfPrimitiveTypes_2D()
 
         Verify(input, arrayRecord);
         Assert.Equal(input, arrayRecord.GetArray(input.GetType()));
+        Assert.Equal(input.Length * 3, arrayRecord.TotalElementsCount);
+    }
+
+    [Fact]
+    public void TotalElementsCountDoesNotIncludeNullArrays()
+    {
+        int[][] input = [[1, 2, 3], null];
+
+        var arrayRecord = (ArrayRecord)NrbfDecoder.Decode(Serialize(input));
+
+        Verify(input, arrayRecord);
+        Assert.Equal(input, arrayRecord.GetArray(input.GetType()));
+        Assert.Equal(3, arrayRecord.TotalElementsCount);
     }
 
     [Fact]
@@ -36,6 +49,7 @@ public void CanReadJaggedArraysOfPrimitiveTypes_3D()
         Verify(input, arrayRecord);
         Assert.Equal(input, arrayRecord.GetArray(input.GetType()));
         Assert.Equal(1, arrayRecord.Rank);
+        Assert.Equal(input.Length * 1 * 3, arrayRecord.TotalElementsCount);
     }
 
     [Fact]
@@ -60,6 +74,7 @@ public void CanReadJaggedArrayOfRectangularArrays()
         Verify(input, arrayRecord);
         Assert.Equal(input, arrayRecord.GetArray(input.GetType()));
         Assert.Equal(1, arrayRecord.Rank);
+        Assert.Equal(input.Length * 3 * 3, arrayRecord.TotalElementsCount);
     }
 
     [Fact]
@@ -75,6 +90,7 @@ public void CanReadJaggedArraysOfStrings()
 
         Verify(input, arrayRecord);
         Assert.Equal(input, arrayRecord.GetArray(input.GetType()));
+        Assert.Equal(input.Length * 3, arrayRecord.TotalElementsCount);
     }
 
     [Fact]
@@ -90,6 +106,7 @@ public void CanReadJaggedArraysOfObjects()
 
         Verify(input, arrayRecord);
         Assert.Equal(input, arrayRecord.GetArray(input.GetType()));
+        Assert.Equal(input.Length * 3, arrayRecord.TotalElementsCount);
     }
 
     [Serializable]
@@ -102,14 +119,17 @@ public class ComplexType
     public void CanReadJaggedArraysOfComplexTypes()
     {
         ComplexType[][] input = new ComplexType[3][];
+        long totalElementsCount = 0;
         for (int i = 0; i < input.Length; i++)
         {
             input[i] = Enumerable.Range(0, i + 1).Select(j => new ComplexType { SomeField = j }).ToArray();
+            totalElementsCount += input[i].Length;
         }
 
         var arrayRecord = (ArrayRecord)NrbfDecoder.Decode(Serialize(input));
 
         Verify(input, arrayRecord);
+        Assert.Equal(totalElementsCount, arrayRecord.TotalElementsCount);
         var output = (ClassRecord?[][])arrayRecord.GetArray(input.GetType());
         for (int i = 0; i < input.Length; i++)
         {

src/libraries/System.Formats.Nrbf/tests/RectangularArraysTests.cs

@@ -223,10 +223,13 @@ public void CanReadRectangularArraysOfComplexTypes_3D()
     internal static void Verify(Array input, ArrayRecord arrayRecord)
     {
         Assert.Equal(input.Rank, arrayRecord.Lengths.Length);
+        long totalElementsCount = 1;
         for (int i = 0; i < input.Rank; i++)
         {
             Assert.Equal(input.GetLength(i), arrayRecord.Lengths[i]);
+            totalElementsCount *= input.GetLength(i);
         }
+        Assert.Equal(totalElementsCount, arrayRecord.TotalElementsCount);
         Assert.Equal(input.GetType().FullName, arrayRecord.TypeName.FullName);
         Assert.Equal(input.GetType().GetAssemblyNameIncludingTypeForwards(), arrayRecord.TypeName.AssemblyName!.FullName);
     }