Add document discrinator sanitizer for JSON protocols

joewyz · joewyz · commit cb64f4f7a7b1 · 2025-12-08T13:36:46.000-08:00
diff --git a/aws/client/aws-client-awsjson/build.gradle.kts b/aws/client/aws-client-awsjson/build.gradle.kts
@@ -20,3 +20,4 @@ dependencies {
 
 val generator = "software.amazon.smithy.java.protocoltests.generators.ProtocolTestGenerator"
 addGenerateSrcsTask(generator, "awsJson1_0", "aws.protocoltests.json10#JsonRpc10")
+addGenerateSrcsTask(generator, "awsJson1_1", "aws.protocoltests.json#JsonProtocol")
diff --git a/aws/client/aws-client-awsjson/src/it/java/software/amazon/smithy/java/client/aws/jsonprotocols/AwsJson11ProtocolTests.java b/aws/client/aws-client-awsjson/src/it/java/software/amazon/smithy/java/client/aws/jsonprotocols/AwsJson11ProtocolTests.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.java.client.aws.jsonprotocols;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import java.nio.charset.StandardCharsets;
+import software.amazon.smithy.java.io.ByteBufferUtils;
+import software.amazon.smithy.java.io.datastream.DataStream;
+import software.amazon.smithy.java.protocoltests.harness.*;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.node.ObjectNode;
+
+@ProtocolTest(
+        service = "aws.protocoltests.json#JsonProtocol",
+        testType = TestType.CLIENT)
+public class AwsJson11ProtocolTests {
+    @HttpClientRequestTests
+    @ProtocolTestFilter(
+            skipTests = {
+                    "SDKAppliedContentEncoding_awsJson1_1",
+                    "SDKAppendsGzipAndIgnoresHttpProvidedEncoding_awsJson1_1",
+            })
+    public void requestTest(DataStream expected, DataStream actual) {
+        Node expectedNode = ObjectNode.objectNode();
+        if (expected.contentLength() != 0) {
+            expectedNode = Node.parse(new String(ByteBufferUtils.getBytes(expected.asByteBuffer()),
+                    StandardCharsets.UTF_8));
+        }
+        Node actualNode = Node.parse(new StringBuildingSubscriber(actual).getResult());
+        assertEquals(expectedNode, actualNode);
+    }
+
+    @HttpClientResponseTests
+    @ProtocolTestFilter(
+            skipTests = {
+                    "AwsJson11FooErrorUsingCode",
+                    "AwsJson11FooErrorUsingCodeAndNamespace",
+                    "AwsJson11FooErrorUsingCodeUriAndNamespace",
+            })
+    public void responseTest(Runnable test) {
+        test.run();
+    }
+}
diff --git a/aws/client/aws-client-awsjson/src/it/java/software/amazon/smithy/java/client/aws/jsonprotocols/AwsJson1ProtocolTests.java b/aws/client/aws-client-awsjson/src/it/java/software/amazon/smithy/java/client/aws/jsonprotocols/AwsJson1ProtocolTests.java
@@ -55,8 +55,6 @@ public void requestTest(DataStream expected, DataStream actual) {
                     "AwsJson10FooErrorUsingCode",
                     "AwsJson10FooErrorUsingCodeAndNamespace",
                     "AwsJson10FooErrorUsingCodeUriAndNamespace",
-                    "AwsJson10FooErrorWithDunderTypeUriAndNamespace"
-
             },
             skipOperations = "aws.protocoltests.json10#OperationWithRequiredMembersWithDefaults")
     public void responseTest(Runnable test) {
diff --git a/aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJson11Protocol.java b/aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJson11Protocol.java
@@ -10,6 +10,7 @@
 import software.amazon.smithy.java.client.core.ClientProtocol;
 import software.amazon.smithy.java.client.core.ClientProtocolFactory;
 import software.amazon.smithy.java.client.core.ProtocolSettings;
+import software.amazon.smithy.java.json.ErrorTypeSanitizer;
 import software.amazon.smithy.model.shapes.ShapeId;
 
 /**
@@ -24,7 +25,7 @@ public final class AwsJson11Protocol extends AwsJsonProtocol {
      *                discriminator of documents that use relative shape IDs.
      */
     public AwsJson11Protocol(ShapeId service) {
-        super(TRAIT_ID, service);
+        super(TRAIT_ID, service, ErrorTypeSanitizer::REMOVE_NAMESPACE_AND_URI);
     }
 
     @Override
diff --git a/aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJson1Protocol.java b/aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJson1Protocol.java
@@ -10,6 +10,7 @@
 import software.amazon.smithy.java.client.core.ClientProtocol;
 import software.amazon.smithy.java.client.core.ClientProtocolFactory;
 import software.amazon.smithy.java.client.core.ProtocolSettings;
+import software.amazon.smithy.java.json.ErrorTypeSanitizer;
 import software.amazon.smithy.model.shapes.ShapeId;
 
 /**
@@ -24,7 +25,7 @@ public final class AwsJson1Protocol extends AwsJsonProtocol {
      *                discriminator of documents that use relative shape IDs.
      */
     public AwsJson1Protocol(ShapeId service) {
-        super(TRAIT_ID, service);
+        super(TRAIT_ID, service, ErrorTypeSanitizer::REMOVE_URI);
     }
 
     @Override
diff --git a/aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJsonProtocol.java b/aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJsonProtocol.java
@@ -9,6 +9,7 @@
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 import software.amazon.smithy.java.client.http.AmznErrorHeaderExtractor;
 import software.amazon.smithy.java.client.http.HttpClientProtocol;
 import software.amazon.smithy.java.client.http.HttpErrorDeserializer;
@@ -39,10 +40,14 @@ abstract sealed class AwsJsonProtocol extends HttpClientProtocol permits AwsJson
      * @param service The service ID used to make X-Amz-Target, and the namespace is used when finding the
      *                discriminator of documents that use relative shape IDs.
      */
-    public AwsJsonProtocol(ShapeId trait, ShapeId service) {
+    public AwsJsonProtocol(ShapeId trait, ShapeId service, Function<String, String> errorTypeSanitizer) {
         super(trait);
         this.service = service;
-        this.codec = JsonCodec.builder().defaultNamespace(service.getNamespace()).build();
+        this.codec = JsonCodec.builder()
+                .defaultNamespace(service.getNamespace())
+                .errorTypeSanitizer(errorTypeSanitizer)
+                .useTimestampFormat(true)
+                .build();
 
         this.errorDeserializer = HttpErrorDeserializer.builder()
                 .codec(codec)
diff --git a/aws/client/aws-client-restjson/src/it/java/software/amazon/smithy/java/client/aws/restjson/RestJson1ProtocolTests.java b/aws/client/aws-client-restjson/src/it/java/software/amazon/smithy/java/client/aws/restjson/RestJson1ProtocolTests.java
@@ -61,7 +61,6 @@ public void requestTest(DataStream expected, DataStream actual) {
                     "RestJsonFooErrorUsingCode",
                     "RestJsonFooErrorUsingCodeAndNamespace",
                     "RestJsonFooErrorUsingCodeUriAndNamespace",
-                    "RestJsonFooErrorWithDunderTypeUriAndNamespace"
             })
     public void responseTest(Runnable test) {
         test.run();
diff --git a/aws/client/aws-client-restjson/src/main/java/software/amazon/smithy/java/aws/client/restjson/RestJsonClientProtocol.java b/aws/client/aws-client-restjson/src/main/java/software/amazon/smithy/java/aws/client/restjson/RestJsonClientProtocol.java
@@ -29,6 +29,7 @@
 import software.amazon.smithy.java.core.serde.event.EventStreamingException;
 import software.amazon.smithy.java.http.api.HttpRequest;
 import software.amazon.smithy.java.http.binding.RequestSerializer;
+import software.amazon.smithy.java.json.ErrorTypeSanitizer;
 import software.amazon.smithy.java.json.JsonCodec;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.shapes.ShapeType;
@@ -52,6 +53,7 @@ public RestJsonClientProtocol(ShapeId service) {
                 .useJsonName(true)
                 .useTimestampFormat(true)
                 .defaultNamespace(service.getNamespace())
+                .errorTypeSanitizer(ErrorTypeSanitizer::REMOVE_NAMESPACE_AND_URI)
                 .build();
 
         this.errorDeserializer = HttpErrorDeserializer.builder()
diff --git a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/ErrorTypeSanitizer.java b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/ErrorTypeSanitizer.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.java.json;
+
+/**
+ * A util class contains the sanitizer functions to process {@code __type} field from the document
+ *
+ * <p>For example, given {@code __type = "aws.protocoltests.restjson#FooError:http://abc.com"},
+ * different protocols have different requirements for the {@code __type} field. This class provides
+ * sanitizers to remove trailing URIs or leading namespaces from {@code __type}.
+ */
+public class ErrorTypeSanitizer {
+    public static String REMOVE_URI(String text) {
+        if (text == null) {
+            return null;
+        }
+        var colon = text.indexOf(':');
+        if (colon > 0) {
+            text = text.substring(0, colon);
+        }
+        return text;
+    }
+
+    public static String REMOVE_NAMESPACE_AND_URI(String text) {
+        if (text == null) {
+            return null;
+        }
+        var hash = text.indexOf('#');
+        if (hash > 0) {
+            text = text.substring(hash + 1);
+        }
+        return REMOVE_URI(text);
+    }
+}
diff --git a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonCodec.java b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonCodec.java
@@ -7,6 +7,7 @@
 
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
+import java.util.function.Function;
 import software.amazon.smithy.java.core.serde.Codec;
 import software.amazon.smithy.java.core.serde.ShapeDeserializer;
 import software.amazon.smithy.java.core.serde.ShapeSerializer;
@@ -174,5 +175,16 @@ Builder overrideSerdeProvider(JsonSerdeProvider provider) {
             settingsBuilder.overrideSerdeProvider(provider);
             return this;
         }
+
+        /**
+         * Uses a custom error type sanitizer to process the {@code __type} field.
+         *
+         * @param errorTypeSanitizer the type sanitizer to use.
+         * @return the builder.
+         */
+        public Builder errorTypeSanitizer(Function<String, String> errorTypeSanitizer) {
+            settingsBuilder.errorTypeSanitizer(errorTypeSanitizer);
+            return this;
+        }
     }
 }
diff --git a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonDocuments.java b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonDocuments.java
@@ -250,6 +250,9 @@ public ShapeId discriminator() {
             if (member != null && member.type() == ShapeType.STRING) {
                 discriminator = member.asString();
             }
+            if (settings.errorTypeSanitizer() != null) {
+                discriminator = settings.errorTypeSanitizer().apply(discriminator);
+            }
             return DocumentDeserializer.parseDiscriminator(discriminator, settings.defaultNamespace());
         }
 
diff --git a/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonSettings.java b/codecs/json-codec/src/main/java/software/amazon/smithy/java/json/JsonSettings.java
@@ -7,6 +7,7 @@
 
 import java.util.Objects;
 import java.util.ServiceLoader;
+import java.util.function.Function;
 import software.amazon.smithy.java.core.serde.TimestampFormatter;
 
 /**
@@ -45,6 +46,7 @@ public final class JsonSettings {
     private final JsonSerdeProvider provider;
     private final boolean serializeTypeInDocuments;
     private final boolean prettyPrint;
+    private final Function<String, String> errorTypeSanitizer;
 
     private JsonSettings(Builder builder) {
         this.timestampResolver = builder.useTimestampFormat
@@ -58,6 +60,7 @@ private JsonSettings(Builder builder) {
         this.provider = builder.provider;
         this.serializeTypeInDocuments = builder.serializeTypeInDocuments;
         this.prettyPrint = builder.prettyPrint;
+        this.errorTypeSanitizer = builder.errorTypeSanitizer;
     }
 
     /**
@@ -115,6 +118,15 @@ public boolean serializeTypeInDocuments() {
         return serializeTypeInDocuments;
     }
 
+    /**
+     * The error type sanitizer to use for {@code __type}. Default is null
+     *
+     * @return the sanitizer used or null
+     */
+    public Function<String, String> errorTypeSanitizer() {
+        return errorTypeSanitizer;
+    }
+
     /**
      * Whether to format the JSON output with pretty printing (indentation and line breaks).
      *
@@ -165,6 +177,7 @@ public static final class Builder {
         private JsonSerdeProvider provider = PROVIDER;
         private boolean serializeTypeInDocuments = true;
         private boolean prettyPrint = false;
+        private Function<String, String> errorTypeSanitizer;
 
         private Builder() {}
 
@@ -275,5 +288,16 @@ Builder overrideSerdeProvider(JsonSerdeProvider provider) {
             this.provider = Objects.requireNonNull(provider);
             return this;
         }
+
+        /**
+         * Uses a custom error type sanitizer for error type
+         *
+         * @param errorTypeSanitizer the sanitizer to use for error type.
+         * @return the builder.
+         */
+        Builder errorTypeSanitizer(Function<String, String> errorTypeSanitizer) {
+            this.errorTypeSanitizer = errorTypeSanitizer;
+            return this;
+        }
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -20,3 +20,4 @@ dependencies {`
`20`	`20`
`21`	`21`	`val generator = "software.amazon.smithy.java.protocoltests.generators.ProtocolTestGenerator"`
`22`	`22`	`addGenerateSrcsTask(generator, "awsJson1_0", "aws.protocoltests.json10#JsonRpc10")`
	`23`	`+addGenerateSrcsTask(generator, "awsJson1_1", "aws.protocoltests.json#JsonProtocol")`