Scrub extension properties in OpenAPI document during serialization (#56688)

Author: captainsafia

src/OpenApi/src/Comparers/OpenApiAnyComparer.cs

@@ -42,6 +42,7 @@ public bool Equals(IOpenApiAny? x, IOpenApiAny? y)
                 OpenApiByte byteX => y is OpenApiByte byteY && byteX.Value.SequenceEqual(byteY.Value),
                 OpenApiDate dateX => y is OpenApiDate dateY && dateX.Value == dateY.Value,
                 OpenApiDateTime dateTimeX => y is OpenApiDateTime dateTimeY && dateTimeX.Value == dateTimeY.Value,
+                ScrubbedOpenApiAny scrubbedX => y is ScrubbedOpenApiAny scrubbedY && scrubbedX.Value == scrubbedY.Value,
                 _ => x.Equals(y)
             });
     }
@@ -73,6 +74,7 @@ public int GetHashCode(IOpenApiAny obj)
             OpenApiPassword password => password.Value,
             OpenApiDate date => date.Value,
             OpenApiDateTime dateTime => dateTime.Value,
+            ScrubbedOpenApiAny scrubbed => scrubbed.Value,
             _ => null
         });
 

src/OpenApi/src/Extensions/OpenApiEndpointRouteBuilderExtensions.cs

@@ -9,7 +9,6 @@
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Options;
 using Microsoft.OpenApi.Extensions;
-using Microsoft.OpenApi.Writers;
 
 namespace Microsoft.AspNetCore.Builder;
 
@@ -49,7 +48,7 @@ public static IEndpointConventionBuilder MapOpenApi(this IEndpointRouteBuilder e
                     using var writer = Utf8BufferTextWriter.Get(output);
                     try
                     {
-                        document.Serialize(new OpenApiJsonWriter(writer), documentOptions.OpenApiVersion);
+                        document.Serialize(new ScrubbingOpenApiJsonWriter(writer), documentOptions.OpenApiVersion);
                         context.Response.ContentType = "application/json;charset=utf-8";
                         await context.Response.BodyWriter.WriteAsync(output.ToArray(), context.RequestAborted);
                         await context.Response.BodyWriter.FlushAsync(context.RequestAborted);

src/OpenApi/src/Schemas/OpenApiJsonSchema.Helpers.cs

@@ -304,7 +304,7 @@ public static void ReadProperty(ref Utf8JsonReader reader, string propertyName,
                 break;
             case OpenApiConstants.SchemaId:
                 reader.Read();
-                schema.Extensions.Add(OpenApiConstants.SchemaId, new OpenApiString(reader.GetString()));
+                schema.Extensions.Add(OpenApiConstants.SchemaId, new ScrubbedOpenApiAny(reader.GetString()));
                 break;
             // OpenAPI does not support the `const` keyword in its schema implementation, so
             // we map it to its closest approximation, an enum with a single value, here.

src/OpenApi/src/Services/OpenApiDocumentProvider.cs

@@ -6,7 +6,6 @@
 using Microsoft.Extensions.Options;
 using Microsoft.OpenApi;
 using Microsoft.OpenApi.Extensions;
-using Microsoft.OpenApi.Writers;
 using System.Linq;
 
 namespace Microsoft.Extensions.ApiDescriptions;
@@ -41,7 +40,7 @@ public async Task GenerateAsync(string documentName, TextWriter writer, OpenApiS
         // more info.
         var targetDocumentService = serviceProvider.GetRequiredKeyedService<OpenApiDocumentService>(documentName);
         var document = await targetDocumentService.GetOpenApiDocumentAsync();
-        var jsonWriter = new OpenApiJsonWriter(writer);
+        var jsonWriter = new ScrubbingOpenApiJsonWriter(writer);
         document.Serialize(jsonWriter, openApiSpecVersion);
     }
 

src/OpenApi/src/Services/OpenApiDocumentService.cs

@@ -20,7 +20,6 @@
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Options;
-using Microsoft.OpenApi.Any;
 using Microsoft.OpenApi.Models;
 
 namespace Microsoft.AspNetCore.OpenApi;
@@ -34,7 +33,6 @@ internal sealed class OpenApiDocumentService(
 {
     private readonly OpenApiOptions _options = optionsMonitor.Get(documentName);
     private readonly OpenApiSchemaService _componentService = serviceProvider.GetRequiredKeyedService<OpenApiSchemaService>(documentName);
-    private readonly IOpenApiDocumentTransformer _scrubExtensionsTransformer = new ScrubExtensionsTransformer();
     private readonly IOpenApiDocumentTransformer _schemaReferenceTransformer = new OpenApiSchemaReferenceTransformer();
 
     private static readonly OpenApiEncoding _defaultFormEncoding = new OpenApiEncoding { Style = ParameterStyle.Form, Explode = true };
@@ -82,8 +80,6 @@ private async Task ApplyTransformersAsync(OpenApiDocument document, Cancellation
         }
         // Move duplicated JSON schemas to the global components.schemas object and map references after all transformers have run.
         await _schemaReferenceTransformer.TransformAsync(document, documentTransformerContext, cancellationToken);
-        // Remove `x-aspnetcore-id` and `x-schema-id` extensions from operations after all transformers have run.
-        await _scrubExtensionsTransformer.TransformAsync(document, documentTransformerContext, cancellationToken);
     }
 
     // Note: Internal for testing.
@@ -126,7 +122,7 @@ private async Task<Dictionary<OperationType, OpenApiOperation>> GetOperationsAsy
         foreach (var description in descriptions)
         {
             var operation = await GetOperationAsync(description, capturedTags, cancellationToken);
-            operation.Extensions.Add(OpenApiConstants.DescriptionId, new OpenApiString(description.ActionDescriptor.Id));
+            operation.Extensions.Add(OpenApiConstants.DescriptionId, new ScrubbedOpenApiAny(description.ActionDescriptor.Id));
             _operationTransformerContextCache.TryAdd(description.ActionDescriptor.Id, new OpenApiOperationTransformerContext
             {
                 DocumentName = documentName,

src/OpenApi/src/Services/Schemas/OpenApiSchemaStore.cs

@@ -4,7 +4,6 @@
 using System.IO.Pipelines;
 using System.Text.Json.Nodes;
 using Microsoft.AspNetCore.Http;
-using Microsoft.OpenApi.Any;
 using Microsoft.OpenApi.Models;
 
 namespace Microsoft.AspNetCore.OpenApi;
@@ -106,7 +105,7 @@ public void PopulateSchemaIntoReferenceCache(OpenApiSchema schema, bool captureS
             // ID to support disambiguating between a derived type on its own and a derived type
             // as part of a polymorphic schema.
             var baseTypeSchemaId = schema.Extensions.TryGetValue(OpenApiConstants.SchemaId, out var schemaId)
-                ? ((OpenApiString)schemaId).Value
+                ? ((ScrubbedOpenApiAny)schemaId).Value
                 : null;
             foreach (var anyOfSchema in schema.AnyOf)
             {
@@ -178,7 +177,7 @@ private void AddOrUpdateSchemaByReference(OpenApiSchema schema, string? baseType
     private static string? GetSchemaReferenceId(OpenApiSchema schema)
     {
         if (schema.Extensions.TryGetValue(OpenApiConstants.SchemaId, out var referenceIdAny)
-            && referenceIdAny is OpenApiString { Value: string referenceId })
+            && referenceIdAny is ScrubbedOpenApiAny { Value: string referenceId })
         {
             return referenceId;
         }

src/OpenApi/src/Transformers/DelegateOpenApiDocumentTransformer.cs

@@ -2,7 +2,6 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using Microsoft.Extensions.DependencyInjection;
-using Microsoft.OpenApi.Any;
 using Microsoft.OpenApi.Models;
 
 namespace Microsoft.AspNetCore.OpenApi;
@@ -43,7 +42,7 @@ public async Task TransformAsync(OpenApiDocument document, OpenApiDocumentTransf
                     }
 
                     if (operation.Extensions.TryGetValue(OpenApiConstants.DescriptionId, out var descriptionIdExtension) &&
-                        descriptionIdExtension is OpenApiString { Value: var descriptionId } &&
+                        descriptionIdExtension is ScrubbedOpenApiAny { Value: string descriptionId } &&
                         documentService.TryGetCachedOperationTransformerContext(descriptionId, out var operationContext))
                     {
                         await _operationTransformer(operation, operationContext, cancellationToken);

src/OpenApi/src/Transformers/Implementations/ScrubExtensionsTransformer.cs

@@ -0,0 +1,30 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Microsoft.OpenApi;
+using Microsoft.OpenApi.Any;
+using Microsoft.OpenApi.Writers;
+
+namespace Microsoft.AspNetCore.OpenApi;
+
+/// <summary>
+/// Represents an <see cref="IOpenApiAny"/> instance that does not serialize itself to
+/// the outgoing document.
+///
+/// The no-op implementation of the <see cref="Write(IOpenApiWriter, OpenApiSpecVersion)"/> method
+/// prevents the value of these properties from being written to disk. When used in conjunction with
+/// the logic to exempt these properties from serialization in <see cref="ScrubbingOpenApiJsonWriter"/>,
+/// we achieve the desired result of not serializing these properties to the output document but retaining
+/// them in the in-memory document.
+/// </summary>
+internal sealed class ScrubbedOpenApiAny(string? value) : IOpenApiAny
+{
+    public AnyType AnyType { get; } = AnyType.Primitive;
+
+    public string? Value { get; } = value;
+
+    public void Write(IOpenApiWriter writer, OpenApiSpecVersion specVersion)
+    {
+        return;
+    }
+}

src/OpenApi/src/Writers/ScrubbedOpenApiAny.cs

@@ -0,0 +1,30 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Microsoft.OpenApi.Writers;
+
+namespace Microsoft.AspNetCore.OpenApi;
+
+/// <summary>
+/// Represents a JSON writer that scrubs certain properties from the output,
+/// specifically the schema ID and description ID that are used for schema resolution
+/// and action descriptor resolution in the in-memory OpenAPI document.
+///
+/// In conjunction with <see cref="ScrubbedOpenApiAny" /> this allows us to work around
+/// the lack of an in-memory property bag on the OpenAPI object model and allows us to
+/// avoid having to scrub the properties in the OpenAPI document prior to serialization.
+///
+/// For more information, see  https://github.com/microsoft/OpenAPI.NET/issues/1719.
+/// </summary>
+internal sealed class ScrubbingOpenApiJsonWriter(TextWriter textWriter) : OpenApiJsonWriter(textWriter)
+{
+    public override void WritePropertyName(string name)
+    {
+        if (name == OpenApiConstants.SchemaId || name == OpenApiConstants.DescriptionId)
+        {
+            return;
+        }
+
+        base.WritePropertyName(name);
+    }
+}

src/OpenApi/src/Writers/ScrubbingOpenApiJsonWriter.cs

@@ -39,7 +39,10 @@ public class OpenApiAnyComparerTests
         [new OpenApiArray { new OpenApiString("value") }, new OpenApiArray { new OpenApiString("value") }, true],
         [new OpenApiArray { new OpenApiString("value") }, new OpenApiArray { new OpenApiString("value2") }, false],
         [new OpenApiArray { new OpenApiString("value2"), new OpenApiString("value") }, new OpenApiArray { new OpenApiString("value"), new OpenApiString("value2") }, false],
-        [new OpenApiArray { new OpenApiString("value"), new OpenApiString("value") }, new OpenApiArray { new OpenApiString("value"), new OpenApiString("value") }, true]
+        [new OpenApiArray { new OpenApiString("value"), new OpenApiString("value") }, new OpenApiArray { new OpenApiString("value"), new OpenApiString("value") }, true],
+        [new ScrubbedOpenApiAny("value"), new ScrubbedOpenApiAny("value"), true],
+        [new ScrubbedOpenApiAny("value"), new ScrubbedOpenApiAny("value2"), false],
+        [new ScrubbedOpenApiAny(null), new ScrubbedOpenApiAny(null), true]
     ];
 
     [Theory]

src/OpenApi/test/Comparers/OpenApiAnyComparerTests.cs

@@ -32,7 +32,7 @@ await Verifier.Verify(GetOpenApiJson(document))
     private static string GetOpenApiJson(OpenApiDocument document)
     {
         using var textWriter = new StringWriter(CultureInfo.InvariantCulture);
-        var jsonWriter = new OpenApiJsonWriter(textWriter);
+        var jsonWriter = new ScrubbingOpenApiJsonWriter(textWriter);
         document.SerializeAsV3(jsonWriter);
         return textWriter.ToString();
     }