From b552c7324ff14da59bd335cfdffd8a489f72e833 Mon Sep 17 00:00:00 2001
From: Mahdi <mkfard@outlook.com>
Date: Thu, 24 Mar 2022 15:01:45 +0430
Subject: [PATCH] Fix swagger schema not being updated when `response_body` is
 set in the .proto file.

---
 .../GrpcHttpApiDescriptionProvider.cs         | 32 ++++++++++++++++++-
 .../src/Shared/ServiceDescriptorHelpers.cs    | 12 +++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/src/GrpcHttpApi/src/Microsoft.AspNetCore.Grpc.Swagger/Internal/GrpcHttpApiDescriptionProvider.cs b/src/GrpcHttpApi/src/Microsoft.AspNetCore.Grpc.Swagger/Internal/GrpcHttpApiDescriptionProvider.cs
index 112236813..9d5faa69e 100644
--- a/src/GrpcHttpApi/src/Microsoft.AspNetCore.Grpc.Swagger/Internal/GrpcHttpApiDescriptionProvider.cs
+++ b/src/GrpcHttpApi/src/Microsoft.AspNetCore.Grpc.Swagger/Internal/GrpcHttpApiDescriptionProvider.cs
@@ -72,10 +72,40 @@ private static ApiDescription CreateApiDescription(RouteEndpoint routeEndpoint,
             };
             apiDescription.RelativePath = pattern.TrimStart('/');
             apiDescription.SupportedRequestFormats.Add(new ApiRequestFormat { MediaType = "application/json" });
+
+            var modelMetadata = new GrpcModelMetadata(ModelMetadataIdentity.ForType(methodDescriptor.OutputType.ClrType));
+            var responseBody = httpRule.ResponseBody;
+
+            if (!string.IsNullOrEmpty(responseBody))
+            {
+                var responseBodyDescriptor = ServiceDescriptorHelpers.ResolveResponseBodyDescriptor(responseBody, methodDescriptor);
+                if (responseBodyDescriptor.IsRepeated)
+                {
+                    if (responseBodyDescriptor.IsMap)
+                    {
+                        var key = responseBodyDescriptor.MessageType.FindFieldByName("key");
+                        var value = responseBodyDescriptor.MessageType.FindFieldByName("value");
+
+                        var keyType = MessageDescriptorHelpers.ResolveFieldType(key);
+                        var valueType = MessageDescriptorHelpers.ResolveFieldType(value);
+
+                        modelMetadata = new GrpcModelMetadata(ModelMetadataIdentity.ForType(typeof(Dictionary<,>).MakeGenericType(keyType, valueType)));
+                    }
+                    else
+                    {
+                        modelMetadata = new GrpcModelMetadata(ModelMetadataIdentity.ForType(typeof(List<>).MakeGenericType(MessageDescriptorHelpers.ResolveFieldType(responseBodyDescriptor))));
+                    }
+                }
+                else
+                {
+                    modelMetadata = new GrpcModelMetadata(ModelMetadataIdentity.ForType(MessageDescriptorHelpers.ResolveFieldType(responseBodyDescriptor)));
+                }
+            }
+
             apiDescription.SupportedResponseTypes.Add(new ApiResponseType
             {
                 ApiResponseFormats = { new ApiResponseFormat { MediaType = "application/json" } },
-                ModelMetadata = new GrpcModelMetadata(ModelMetadataIdentity.ForType(methodDescriptor.OutputType.ClrType)),
+                ModelMetadata = modelMetadata,
                 StatusCode = 200
             });
             var explorerSettings = routeEndpoint.Metadata.GetMetadata<ApiExplorerSettingsAttribute>();
diff --git a/src/GrpcHttpApi/src/Shared/ServiceDescriptorHelpers.cs b/src/GrpcHttpApi/src/Shared/ServiceDescriptorHelpers.cs
index 640603ffe..9e346da27 100644
--- a/src/GrpcHttpApi/src/Shared/ServiceDescriptorHelpers.cs
+++ b/src/GrpcHttpApi/src/Shared/ServiceDescriptorHelpers.cs
@@ -318,6 +318,18 @@ public static Dictionary<string, List<FieldDescriptor>> ResolveRouteParameterDes
             return null;
         }
 
+        public static FieldDescriptor ResolveResponseBodyDescriptor(string responseBody, MethodDescriptor methodDescriptor)
+        {
+            if (!TryResolveDescriptors(methodDescriptor.OutputType, responseBody, out var responseBodyFieldDescriptors))
+            {
+                throw new InvalidOperationException($"Couldn't find matching field for response-body '{responseBody}' on {methodDescriptor.OutputType.Name}.");
+            }
+
+            var fieldDescriptor = responseBodyFieldDescriptors.Last();
+            return fieldDescriptor;
+
+        }
+
         public record BodyDescriptorInfo(
             MessageDescriptor Descriptor,
             List<FieldDescriptor>? FieldDescriptors,