fixes customizer logic in the GsonGrpcJsonMarshallerBuilder

grpc/src/main/java/com/linecorp/armeria/common/grpc/GsonGrpcJsonMarshallerBuilder.java

@@ -19,6 +19,7 @@
 import static java.util.Objects.requireNonNull;
 
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 import com.google.protobuf.Message;
 import com.google.protobuf.util.JsonFormat;
@@ -32,40 +33,98 @@
 public final class GsonGrpcJsonMarshallerBuilder {
 
     @Nullable
-    private Consumer<JsonFormat.Parser> jsonParserCustomizer;
+    private Function<JsonFormat.Parser, JsonFormat.Parser> jsonParserCustomizer;
 
     @Nullable
-    private Consumer<JsonFormat.Printer> jsonPrinterCustomizer;
+    private Function<JsonFormat.Printer, JsonFormat.Printer> jsonPrinterCustomizer;
 
     GsonGrpcJsonMarshallerBuilder() {}
 
+    /**
+     * Adds a {@link Function} that returns customized the {@link JsonFormat.Parser}
+     * used when deserializing a JSON payload into a {@link Message}.
+     */
+    public GsonGrpcJsonMarshallerBuilder jsonParserCustomizer(
+            Function<? super JsonFormat.Parser, JsonFormat.Parser> jsonParserCustomizer) {
+        requireNonNull(jsonParserCustomizer, "jsonParserCustomizer");
+        if (this.jsonParserCustomizer == null) {
+            @SuppressWarnings("unchecked")
+            final Function<JsonFormat.Parser, JsonFormat.Parser> cast =
+                    (Function<JsonFormat.Parser, JsonFormat.Parser>) jsonParserCustomizer;
+            this.jsonParserCustomizer = cast;
+        } else {
+            this.jsonParserCustomizer = this.jsonParserCustomizer.andThen(jsonParserCustomizer);
+        }
+        return this;
+    }
+
     /**
      * Adds a {@link Consumer} that can customize the {@link JsonFormat.Parser}
      * used when deserializing a JSON payload into a {@link Message}.
+     *
+     * @deprecated {@link JsonFormat.Parser} is immutable so all changes applied in the {@link Consumer}
+     *     will be lost. Please use the {@link #jsonParserCustomizer(Function) jsonParserCustomizer}
+     *     which accepts {@link Function} parameter instead.
      */
+    @Deprecated
     public GsonGrpcJsonMarshallerBuilder jsonParserCustomizer(
             Consumer<? super JsonFormat.Parser> jsonParserCustomizer) {
         requireNonNull(jsonParserCustomizer, "jsonParserCustomizer");
         if (this.jsonParserCustomizer == null) {
             @SuppressWarnings("unchecked")
             final Consumer<JsonFormat.Parser> cast = (Consumer<JsonFormat.Parser>) jsonParserCustomizer;
-            this.jsonParserCustomizer = cast;
+            this.jsonParserCustomizer = parser -> {
+                cast.accept(parser);
+                return parser;
+            };
         } else {
-            this.jsonParserCustomizer = this.jsonParserCustomizer.andThen(jsonParserCustomizer);
+            this.jsonParserCustomizer = this.jsonParserCustomizer.andThen(parser -> {
+                jsonParserCustomizer.accept(parser);
+                return parser;
+            });
         }
         return this;
     }
 
     /**
      * Adds a {@link Consumer} that can customize the {@link JsonFormat.Printer}
      * used when serializing a {@link Message} into a JSON payload.
+     *
+     * @deprecated {@link JsonFormat.Printer} is immutable so all changes applied in the {@link Consumer}
+     *     will be lost. Please use the {@link #jsonPrinterCustomizer(Function) jsonParserCustomizer}
+     *     which accepts {@link Function} parameter instead.
      */
+    @Deprecated
     public GsonGrpcJsonMarshallerBuilder jsonPrinterCustomizer(
             Consumer<? super JsonFormat.Printer> jsonPrinterCustomizer) {
         requireNonNull(jsonPrinterCustomizer, "jsonPrinterCustomizer");
         if (this.jsonPrinterCustomizer == null) {
             @SuppressWarnings("unchecked")
             final Consumer<JsonFormat.Printer> cast = (Consumer<JsonFormat.Printer>) jsonPrinterCustomizer;
+            this.jsonPrinterCustomizer = printer -> {
+                cast.accept(printer);
+                return printer;
+            };
+        } else {
+            this.jsonPrinterCustomizer = this.jsonPrinterCustomizer.andThen(printer -> {
+                jsonPrinterCustomizer.accept(printer);
+                return printer;
+            });
+        }
+        return this;
+    }
+
+    /**
+     * Adds a {@link Function} that returns customized the {@link JsonFormat.Printer}
+     * used when serializing a {@link Message} into a JSON payload.
+     */
+    public GsonGrpcJsonMarshallerBuilder jsonPrinterCustomizer(
+            Function<? super JsonFormat.Printer, JsonFormat.Printer> jsonPrinterCustomizer) {
+        requireNonNull(jsonPrinterCustomizer, "jsonPrinterCustomizer");
+        if (this.jsonPrinterCustomizer == null) {
+            @SuppressWarnings("unchecked")
+            final Function<JsonFormat.Printer, JsonFormat.Printer> cast =
+                    (Function<JsonFormat.Printer, JsonFormat.Printer>) jsonPrinterCustomizer;
             this.jsonPrinterCustomizer = cast;
         } else {
             this.jsonPrinterCustomizer = this.jsonPrinterCustomizer.andThen(jsonPrinterCustomizer);
@@ -77,14 +136,14 @@ public GsonGrpcJsonMarshallerBuilder jsonPrinterCustomizer(
      * Returns a newly-created {@link GrpcJsonMarshaller}.
      */
     public GrpcJsonMarshaller build() {
-        final JsonFormat.Printer printer = JsonFormat.printer().omittingInsignificantWhitespace();
+        JsonFormat.Printer printer = JsonFormat.printer().omittingInsignificantWhitespace();
         if (jsonPrinterCustomizer != null) {
-            jsonPrinterCustomizer.accept(printer);
+            printer = jsonPrinterCustomizer.apply(printer);
         }
 
-        final JsonFormat.Parser parser = JsonFormat.parser().ignoringUnknownFields();
+        JsonFormat.Parser parser = JsonFormat.parser().ignoringUnknownFields();
         if (jsonParserCustomizer != null) {
-            jsonParserCustomizer.accept(parser);
+            parser = jsonParserCustomizer.apply(parser);
         }
         return new GsonGrpcJsonMarshaller(printer, parser);
     }

grpc/src/test/java/com/linecorp/armeria/common/grpc/GsonGrpcJsonMarshallerBuilderTest.java

@@ -0,0 +1,125 @@
+/*
+ *  Copyright 2025 LINE Corporation
+ *
+ *  LINE Corporation licenses this file to you 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:
+ *
+ *    https://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 com.linecorp.armeria.common.grpc;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import java.io.IOException;
+import java.io.InputStream;
+
+import org.apache.tools.ant.filters.StringInputStream;
+import org.jetbrains.annotations.Nullable;
+import org.junit.jupiter.api.Test;
+
+import com.google.api.client.testing.util.TestableByteArrayOutputStream;
+import com.google.protobuf.util.JsonFormat;
+
+import io.grpc.MethodDescriptor;
+import testing.grpc.Messages;
+import testing.grpc.TestServiceGrpc;
+
+class GsonGrpcJsonMarshallerBuilderTest {
+    private static final Messages.SimpleRequest testData = Messages.SimpleRequest.newBuilder()
+            .setFillUsername(true)
+            .setPayload(
+                    Messages.Payload.newBuilder()
+                            .setType(Messages.PayloadType.RANDOM)
+                            .build()
+            )
+            .build();
+
+    private static final MethodDescriptor.Marshaller<Messages.SimpleRequest> customRequestMarshaller =
+            new MethodDescriptor.PrototypeMarshaller<Messages.SimpleRequest>() {
+                @Override
+                public Class<Messages.SimpleRequest> getMessageClass() {
+                    return Messages.SimpleRequest.class;
+                }
+
+                @Nullable
+                @Override
+                public Messages.SimpleRequest getMessagePrototype() {
+                    return Messages.SimpleRequest.getDefaultInstance();
+                }
+
+                @Override
+                public InputStream stream(Messages.SimpleRequest value) {
+                    return TestServiceGrpc.getUnaryCallMethod().getRequestMarshaller().stream(value);
+                }
+
+                @Override
+                public Messages.SimpleRequest parse(InputStream stream) {
+                    return TestServiceGrpc.getUnaryCallMethod().getRequestMarshaller().parse(stream);
+                }
+            };
+
+    @Test
+    void createJsonPrinterWithDefaultSettingsIfNoCustomizerRegistered() throws IOException {
+        final GrpcJsonMarshaller jsonMarshaller = GrpcJsonMarshaller.builderForGson().build();
+        final String json = serializeToJson(jsonMarshaller);
+        assertThat(json)
+                .isEqualTo("{\"payload\":{\"type\":\"RANDOM\"},\"fillUsername\":true}");
+    }
+
+    @Test
+    void createJsonPrinterWithCustomizer() throws IOException {
+        final GrpcJsonMarshaller jsonMarshaller = GrpcJsonMarshaller.builderForGson()
+                .jsonPrinterCustomizer(JsonFormat.Printer::preservingProtoFieldNames)
+                .jsonPrinterCustomizer(JsonFormat.Printer::printingEnumsAsInts)
+                .build();
+        final String json = serializeToJson(jsonMarshaller);
+        assertThat(json)
+                .isEqualTo("{\"payload\":{\"type\":2},\"fill_username\":true}");
+    }
+
+    @Test
+    void createJsonParserWithDefaultSettingsIfNoCustomizerRegistered() throws IOException {
+        final GrpcJsonMarshaller jsonMarshaller = GrpcJsonMarshaller.builderForGson().build();
+        assertThat(parseJson(jsonMarshaller, "{\"test\": true,\"fill_username\":true}").getFillUsername())
+                .isEqualTo(true);
+    }
+
+    @Test
+    void createJsonParserWithCustomizerNotIgnoringUnknownFields() {
+        final GrpcJsonMarshaller jsonMarshaller = GrpcJsonMarshaller.builderForGson()
+                .jsonParserCustomizer(parser -> {
+                    return parser.usingTypeRegistry(
+                            JsonFormat.TypeRegistry.newBuilder()
+                                    .add(Messages.SimpleRequest.getDescriptor())
+                                    .build()
+                    );
+                })
+                .jsonParserCustomizer(parser -> {
+                    return JsonFormat.parser();
+                })
+                .build();
+        assertThatThrownBy(() -> parseJson(jsonMarshaller, "{\"test\": true}"))
+                .hasMessageStartingWith("Cannot find field");
+    }
+
+    private static String serializeToJson(GrpcJsonMarshaller jsonMarshaller) throws IOException {
+        final TestableByteArrayOutputStream outputStream = new TestableByteArrayOutputStream();
+        jsonMarshaller.serializeMessage(customRequestMarshaller, testData, outputStream);
+        return outputStream.toString();
+    }
+
+    private static Messages.SimpleRequest parseJson(
+            GrpcJsonMarshaller jsonMarshaller, String input
+    ) throws IOException {
+        return jsonMarshaller.deserializeMessage(customRequestMarshaller, new StringInputStream(input));
+    }
+}