Merge pull request #32 from wiremock/any-type-issue

Added support for Any type in a slightly hacky way

src/main/java/org/wiremock/grpc/dsl/WireMockGrpc.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 Thomas Akehurst
+ * Copyright (C) 2023-2024 Thomas Akehurst
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,7 +15,9 @@
  */
 package org.wiremock.grpc.dsl;
 
+import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.github.tomakehurst.wiremock.client.WireMock;
+import com.github.tomakehurst.wiremock.common.Json;
 import com.github.tomakehurst.wiremock.matching.StringValuePattern;
 import com.google.protobuf.MessageOrBuilder;
 import org.wiremock.annotations.Beta;
@@ -48,6 +50,15 @@ public static GrpcResponseDefinitionBuilder message(MessageOrBuilder messageOrBu
     return new GrpcResponseDefinitionBuilder(Status.OK).fromJson(json);
   }
 
+  public static GrpcResponseDefinitionBuilder messageAsAny(MessageOrBuilder messageOrBuilder) {
+    final String initialJson = JsonMessageUtils.toJson(messageOrBuilder);
+    final ObjectNode jsonObject = Json.read(initialJson, ObjectNode.class);
+    jsonObject.put(
+        "@type", "type.googleapis.com/" + messageOrBuilder.getDescriptorForType().getFullName());
+    String finalJson = Json.write(jsonObject);
+    return new GrpcResponseDefinitionBuilder(Status.OK).fromJson(finalJson);
+  }
+
   public enum Status {
     OK(0),
     CANCELLED(1),

src/main/java/org/wiremock/grpc/internal/BaseCallHandler.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2023-2024 Thomas Akehurst
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.wiremock.grpc.internal;
+
+import com.github.tomakehurst.wiremock.http.StubRequestHandler;
+import com.google.protobuf.Descriptors;
+
+public abstract class BaseCallHandler {
+
+  protected final StubRequestHandler stubRequestHandler;
+  protected final Descriptors.ServiceDescriptor serviceDescriptor;
+  protected final Descriptors.MethodDescriptor methodDescriptor;
+
+  protected final JsonMessageConverter jsonMessageConverter;
+
+  protected BaseCallHandler(
+      StubRequestHandler stubRequestHandler,
+      Descriptors.ServiceDescriptor serviceDescriptor,
+      Descriptors.MethodDescriptor methodDescriptor,
+      JsonMessageConverter jsonMessageConverter) {
+    this.stubRequestHandler = stubRequestHandler;
+    this.serviceDescriptor = serviceDescriptor;
+    this.methodDescriptor = methodDescriptor;
+    this.jsonMessageConverter = jsonMessageConverter;
+  }
+}

src/main/java/org/wiremock/grpc/internal/ClientStreamingServerCallHandler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 Thomas Akehurst
+ * Copyright (C) 2023-2024 Thomas Akehurst
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -29,20 +29,15 @@
 import java.util.concurrent.atomic.AtomicReference;
 import org.wiremock.grpc.dsl.WireMockGrpc;
 
-public class ClientStreamingServerCallHandler
+public class ClientStreamingServerCallHandler extends BaseCallHandler
     implements ServerCalls.ClientStreamingMethod<DynamicMessage, DynamicMessage> {
 
-  private final StubRequestHandler stubRequestHandler;
-  private final Descriptors.ServiceDescriptor serviceDescriptor;
-  private final Descriptors.MethodDescriptor methodDescriptor;
-
   public ClientStreamingServerCallHandler(
       StubRequestHandler stubRequestHandler,
       Descriptors.ServiceDescriptor serviceDescriptor,
-      Descriptors.MethodDescriptor methodDescriptor) {
-    this.stubRequestHandler = stubRequestHandler;
-    this.serviceDescriptor = serviceDescriptor;
-    this.methodDescriptor = methodDescriptor;
+      Descriptors.MethodDescriptor methodDescriptor,
+      JsonMessageConverter jsonMessageConverter) {
+    super(stubRequestHandler, serviceDescriptor, methodDescriptor, jsonMessageConverter);
   }
 
   @Override
@@ -67,7 +62,7 @@ public void onNext(DynamicMessage request) {
                 serverAddress.port,
                 serviceDescriptor.getFullName(),
                 methodDescriptor.getName(),
-                request);
+                jsonMessageConverter.toJson(request));
 
         stubRequestHandler.handle(
             wireMockRequest,
@@ -97,7 +92,7 @@ public void onNext(DynamicMessage request) {
                   DynamicMessage.newBuilder(methodDescriptor.getOutputType());
 
               final DynamicMessage response =
-                  JsonMessageUtils.toMessage(resp.getBodyAsString(), messageBuilder);
+                  jsonMessageConverter.toMessage(resp.getBodyAsString(), messageBuilder);
 
               responseStatus.set(WireMockGrpc.Status.OK);
               firstResponse.set(response);

src/main/java/org/wiremock/grpc/internal/GrpcFilter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 Thomas Akehurst
+ * Copyright (C) 2023-2024 Thomas Akehurst
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@
 import com.github.tomakehurst.wiremock.http.StubRequestHandler;
 import com.google.protobuf.Descriptors;
 import com.google.protobuf.DynamicMessage;
+import com.google.protobuf.TypeRegistry;
 import io.grpc.*;
 import io.grpc.protobuf.ProtoUtils;
 import io.grpc.servlet.jakarta.GrpcServlet;
@@ -42,10 +43,22 @@ public class GrpcFilter extends HttpFilter {
   private final StubRequestHandler stubRequestHandler;
   private final List<Descriptors.FileDescriptor> fileDescriptors;
 
+  private final JsonMessageConverter jsonMessageConverter;
+
   public GrpcFilter(
       StubRequestHandler stubRequestHandler, List<Descriptors.FileDescriptor> fileDescriptors) {
     this.stubRequestHandler = stubRequestHandler;
     this.fileDescriptors = fileDescriptors;
+
+    final TypeRegistry.Builder typeRegistryBuilder = TypeRegistry.newBuilder();
+    fileDescriptors.forEach(
+        fileDescriptor -> {
+          fileDescriptor.getMessageTypes().forEach(typeRegistryBuilder::add);
+        });
+
+    final TypeRegistry typeRegistry = typeRegistryBuilder.build();
+    jsonMessageConverter = new JsonMessageConverter(typeRegistry);
+
     grpcServlet = new GrpcServlet(buildServices());
   }
 
@@ -77,9 +90,10 @@ private ServerCallHandler<DynamicMessage, DynamicMessage> buildHandler(
     return methodDescriptor.isClientStreaming()
         ? ServerCalls.asyncClientStreamingCall(
             new ClientStreamingServerCallHandler(
-                stubRequestHandler, serviceDescriptor, methodDescriptor))
+                stubRequestHandler, serviceDescriptor, methodDescriptor, jsonMessageConverter))
         : ServerCalls.asyncUnaryCall(
-            new UnaryServerCallHandler(stubRequestHandler, serviceDescriptor, methodDescriptor));
+            new UnaryServerCallHandler(
+                stubRequestHandler, serviceDescriptor, methodDescriptor, jsonMessageConverter));
   }
 
   private static MethodDescriptor<DynamicMessage, DynamicMessage> buildMessageDescriptorInstance(

src/main/java/org/wiremock/grpc/internal/GrpcHttpServerFactory.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 Thomas Akehurst
+ * Copyright (C) 2023-2024 Thomas Akehurst
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -35,10 +35,33 @@
 
 public class GrpcHttpServerFactory implements HttpServerFactory {
 
-  private final BlobStore protoDescriptorStore;
+  private final List<Descriptors.FileDescriptor> fileDescriptors = new ArrayList<>();
 
   public GrpcHttpServerFactory(BlobStore protoDescriptorStore) {
-    this.protoDescriptorStore = protoDescriptorStore;
+    protoDescriptorStore
+        .getAllKeys()
+        .filter(key -> key.endsWith(".dsc") || key.endsWith(".desc"))
+        .map(
+            key ->
+                protoDescriptorStore
+                    .get(key)
+                    .map(
+                        data ->
+                            Exceptions.uncheck(
+                                () -> DescriptorProtos.FileDescriptorSet.parseFrom(data),
+                                DescriptorProtos.FileDescriptorSet.class)))
+        .filter(Optional::isPresent)
+        .map(Optional::get)
+        .flatMap(fileDescriptorSet -> fileDescriptorSet.getFileList().stream())
+        .forEach(
+            fileDescriptorProto ->
+                Exceptions.uncheck(
+                    () ->
+                        fileDescriptors.add(
+                            Descriptors.FileDescriptor.buildFrom(
+                                fileDescriptorProto,
+                                fileDescriptors.toArray(Descriptors.FileDescriptor[]::new),
+                                true))));
   }
 
   @Override
@@ -56,33 +79,6 @@ public HttpServer buildHttpServer(
       protected void decorateMockServiceContextBeforeConfig(
           ServletContextHandler mockServiceContext) {
 
-        List<Descriptors.FileDescriptor> fileDescriptors = new ArrayList<>();
-
-        protoDescriptorStore
-            .getAllKeys()
-            .filter(key -> key.endsWith(".dsc") || key.endsWith(".desc"))
-            .map(
-                key ->
-                    protoDescriptorStore
-                        .get(key)
-                        .map(
-                            data ->
-                                Exceptions.uncheck(
-                                    () -> DescriptorProtos.FileDescriptorSet.parseFrom(data),
-                                    DescriptorProtos.FileDescriptorSet.class)))
-            .filter(Optional::isPresent)
-            .map(Optional::get)
-            .flatMap(fileDescriptorSet -> fileDescriptorSet.getFileList().stream())
-            .forEach(
-                fileDescriptorProto ->
-                    Exceptions.uncheck(
-                        () ->
-                            fileDescriptors.add(
-                                Descriptors.FileDescriptor.buildFrom(
-                                    fileDescriptorProto,
-                                    fileDescriptors.toArray(Descriptors.FileDescriptor[]::new),
-                                    true))));
-
         final GrpcFilter grpcFilter = new GrpcFilter(stubRequestHandler, fileDescriptors);
         final FilterHolder filterHolder = new FilterHolder(grpcFilter);
         mockServiceContext.addFilter(filterHolder, "/*", EnumSet.of(DispatcherType.REQUEST));

src/main/java/org/wiremock/grpc/internal/GrpcRequest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 Thomas Akehurst
+ * Copyright (C) 2023-2024 Thomas Akehurst
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,7 +20,6 @@
 import com.github.tomakehurst.wiremock.common.Encoding;
 import com.github.tomakehurst.wiremock.common.Strings;
 import com.github.tomakehurst.wiremock.http.*;
-import com.google.protobuf.DynamicMessage;
 import java.util.Collection;
 import java.util.Map;
 import java.util.Optional;
@@ -37,19 +36,13 @@ public class GrpcRequest implements Request {
   private final String body;
 
   public GrpcRequest(
-      String scheme,
-      String host,
-      int port,
-      String serviceName,
-      String methodName,
-      DynamicMessage message) {
+      String scheme, String host, int port, String serviceName, String methodName, String body) {
     this.scheme = scheme;
     this.host = host;
     this.port = port;
     this.serviceName = serviceName;
     this.methodName = methodName;
-
-    body = JsonMessageUtils.toJson(message);
+    this.body = body;
   }
 
   @Override

src/main/java/org/wiremock/grpc/internal/JsonMessageConverter.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2023-2024 Thomas Akehurst
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.wiremock.grpc.internal;
+
+import com.github.tomakehurst.wiremock.common.Exceptions;
+import com.google.protobuf.Message;
+import com.google.protobuf.MessageOrBuilder;
+import com.google.protobuf.TypeRegistry;
+import com.google.protobuf.util.JsonFormat;
+
+public class JsonMessageConverter {
+
+  private final JsonFormat.Printer jsonPrinter;
+  private final JsonFormat.Parser jsonParser;
+
+  public JsonMessageConverter(TypeRegistry typeRegistry) {
+    jsonPrinter = JsonFormat.printer().usingTypeRegistry(typeRegistry);
+    jsonParser = JsonFormat.parser().usingTypeRegistry(typeRegistry);
+  }
+
+  public String toJson(MessageOrBuilder message) {
+    return Exceptions.uncheck(() -> jsonPrinter.print(message), String.class);
+  }
+
+  @SuppressWarnings("unchecked")
+  public <T extends Message, B extends Message.Builder> T toMessage(String json, B builder) {
+    Exceptions.uncheck(() -> jsonParser.merge(json, builder));
+    return (T) builder.build();
+  }
+}

src/main/java/org/wiremock/grpc/internal/JsonMessageUtils.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 Thomas Akehurst
+ * Copyright (C) 2023-2024 Thomas Akehurst
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,25 +15,22 @@
  */
 package org.wiremock.grpc.internal;
 
-import com.github.tomakehurst.wiremock.common.Exceptions;
 import com.google.protobuf.Message;
 import com.google.protobuf.MessageOrBuilder;
-import com.google.protobuf.util.JsonFormat;
+import com.google.protobuf.TypeRegistry;
 
 public class JsonMessageUtils {
 
-  private static final JsonFormat.Printer jsonPrinter = JsonFormat.printer();
-  private static final JsonFormat.Parser jsonParser = JsonFormat.parser();
+  private static final JsonMessageConverter converter =
+      new JsonMessageConverter(TypeRegistry.getEmptyTypeRegistry());
 
   private JsonMessageUtils() {}
 
   public static String toJson(MessageOrBuilder message) {
-    return Exceptions.uncheck(() -> jsonPrinter.print(message), String.class);
+    return converter.toJson(message);
   }
 
-  @SuppressWarnings("unchecked")
   public static <T extends Message, B extends Message.Builder> T toMessage(String json, B builder) {
-    Exceptions.uncheck(() -> jsonParser.merge(json, builder));
-    return (T) builder.build();
+    return converter.toMessage(json, builder);
   }
 }