Add document discriminator sanitizer for JSON protocols

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")

aws/client/aws-client-awsjson/src/it/java/software/amazon/smithy/java/client/aws/jsonprotocols/AwsJson11ProtocolTests.java

@@ -0,0 +1,41 @@
+/*
+ * 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
+    public void responseTest(Runnable test) {
+        test.run();
+    }
+}

aws/client/aws-client-awsjson/src/it/java/software/amazon/smithy/java/client/aws/jsonprotocols/AwsJson1ProtocolTests.java

@@ -51,12 +51,6 @@ public void requestTest(DataStream expected, DataStream actual) {
             skipTests = {
                     "AwsJson10ClientPopulatesDefaultsValuesWhenMissingInResponse",
                     "AwsJson10ClientIgnoresDefaultValuesIfMemberValuesArePresentInResponse",
-                    //The below fail because we haven't implemented code based exception handling
-                    "AwsJson10FooErrorUsingCode",
-                    "AwsJson10FooErrorUsingCodeAndNamespace",
-                    "AwsJson10FooErrorUsingCodeUriAndNamespace",
-                    "AwsJson10FooErrorWithDunderTypeUriAndNamespace"
-
             },
             skipOperations = "aws.protocoltests.json10#OperationWithRequiredMembersWithDefaults")
     public void responseTest(Runnable test) {

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.client.http.HttpErrorDeserializer;
 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, HttpErrorDeserializer.DEFAULT_ERROR_PAYLOAD_PARSER);
     }
 
     @Override

aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJson1Protocol.java

@@ -10,6 +10,9 @@
 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.client.http.ErrorTypeUtils;
+import software.amazon.smithy.java.core.serde.document.Document;
+import software.amazon.smithy.java.core.serde.document.DocumentDeserializer;
 import software.amazon.smithy.model.shapes.ShapeId;
 
 /**
@@ -24,7 +27,13 @@ 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, AwsJson1Protocol::extractErrorType);
+    }
+
+    private static ShapeId extractErrorType(Document document, String namespace) {
+        return DocumentDeserializer.parseDiscriminator(
+                ErrorTypeUtils.removeUri(ErrorTypeUtils.readTypeAndCode(document)),
+                namespace);
     }
 
     @Override

aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJsonProtocol.java

@@ -39,14 +39,22 @@ 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,
+            HttpErrorDeserializer.ErrorPayloadParser errorPayloadParser
+    ) {
         super(trait);
         this.service = service;
-        this.codec = JsonCodec.builder().defaultNamespace(service.getNamespace()).build();
+        this.codec = JsonCodec.builder()
+                .defaultNamespace(service.getNamespace())
+                .useTimestampFormat(true)
+                .build();
 
         this.errorDeserializer = HttpErrorDeserializer.builder()
                 .codec(codec)
                 .serviceId(service)
+                .errorPayloadParser(errorPayloadParser)
                 .headerErrorExtractor(new AmznErrorHeaderExtractor())
                 .build();
     }

aws/client/aws-client-restjson/src/it/java/software/amazon/smithy/java/client/aws/restjson/RestJson1ProtocolTests.java

@@ -56,13 +56,6 @@ public void requestTest(DataStream expected, DataStream actual) {
     }
 
     @HttpClientResponseTests
-    @ProtocolTestFilter(
-            skipTests = {
-                    "RestJsonFooErrorUsingCode",
-                    "RestJsonFooErrorUsingCodeAndNamespace",
-                    "RestJsonFooErrorUsingCodeUriAndNamespace",
-                    "RestJsonFooErrorWithDunderTypeUriAndNamespace"
-            })
     public void responseTest(Runnable test) {
         test.run();
     }

aws/client/aws-client-restxml/src/main/java/software/amazon/smithy/java/aws/client/restxml/RestXmlClientProtocol.java

@@ -18,11 +18,13 @@
 import software.amazon.smithy.java.client.http.binding.HttpBindingClientProtocol;
 import software.amazon.smithy.java.client.http.binding.HttpBindingErrorFactory;
 import software.amazon.smithy.java.context.Context;
+import software.amazon.smithy.java.core.error.CallException;
 import software.amazon.smithy.java.core.error.ModeledException;
 import software.amazon.smithy.java.core.schema.InputEventStreamingApiOperation;
 import software.amazon.smithy.java.core.schema.OutputEventStreamingApiOperation;
 import software.amazon.smithy.java.core.serde.Codec;
 import software.amazon.smithy.java.core.serde.TypeRegistry;
+import software.amazon.smithy.java.core.serde.document.Document;
 import software.amazon.smithy.java.core.serde.event.EventDecoderFactory;
 import software.amazon.smithy.java.core.serde.event.EventEncoderFactory;
 import software.amazon.smithy.java.core.serde.event.EventStreamingException;
@@ -95,24 +97,34 @@ protected EventDecoderFactory<AwsEventFrame> getEventDecoderFactory(
         return AwsEventDecoderFactory.forOutputStream(outputOperation, payloadCodec(), f -> f);
     }
 
-    private static final HttpErrorDeserializer.ErrorPayloadParser XML_ERROR_PAYLOAD_PARSER = (
-            Context context,
-            Codec codec,
-            HttpErrorDeserializer.KnownErrorFactory knownErrorFactory,
-            ShapeId serviceId,
-            TypeRegistry typeRegistry,
-            HttpResponse response,
-            ByteBuffer buffer) -> {
-        var deserializer = codec.createDeserializer(buffer);
-        String code = XmlUtil.parseErrorCodeName(deserializer);
-        var nameSpace = serviceId.getNamespace();
-        var id = ShapeId.fromOptionalNamespace(nameSpace, code);
-        var builder = typeRegistry.createBuilder(id, ModeledException.class);
-        if (builder != null) {
-            return knownErrorFactory.createError(context, codec, response, builder);
-        }
-        return null;
-    };
+    private static final HttpErrorDeserializer.ErrorPayloadParser XML_ERROR_PAYLOAD_PARSER =
+            new HttpErrorDeserializer.ErrorPayloadParser() {
+                @Override
+                public CallException parsePayload(
+                        Context context,
+                        Codec codec,
+                        HttpErrorDeserializer.KnownErrorFactory knownErrorFactory,
+                        ShapeId serviceId,
+                        TypeRegistry typeRegistry,
+                        HttpResponse response,
+                        ByteBuffer buffer
+                ) {
+                    var deserializer = codec.createDeserializer(buffer);
+                    String code = XmlUtil.parseErrorCodeName(deserializer);
+                    var nameSpace = serviceId.getNamespace();
+                    var id = ShapeId.fromOptionalNamespace(nameSpace, code);
+                    var builder = typeRegistry.createBuilder(id, ModeledException.class);
+                    if (builder != null) {
+                        return knownErrorFactory.createError(context, codec, response, builder);
+                    }
+                    return null;
+                }
+
+                @Override
+                public ShapeId extractErrorType(Document document, String namespace) {
+                    return null;
+                }
+            };
 
     public static final class Factory implements ClientProtocolFactory<RestXmlTrait> {
         @Override

client/client-http/src/main/java/software/amazon/smithy/java/client/http/ErrorTypeUtils.java

@@ -0,0 +1,86 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.java.client.http;
+
+import software.amazon.smithy.java.core.serde.document.Document;
+import software.amazon.smithy.model.shapes.ShapeType;
+
+public final class ErrorTypeUtils {
+
+    private ErrorTypeUtils() {}
+
+    /**
+     * Read the error type from __type field of the document.
+     *
+     * @param document the document of the payload to read.
+     * @return the extracted error type from the document.
+     */
+    public static String readType(Document document) {
+        String errorType = null;
+        var member = document.getMember("__type");
+        if (member != null && member.type() == ShapeType.STRING) {
+            errorType = member.asString();
+        }
+        return errorType;
+    }
+
+    /**
+     * Read the error type from __type or code field of the document.
+     *
+     * @param document the document of the payload to read.
+     * @return the extracted error type from the document.
+     */
+    public static String readTypeAndCode(Document document) {
+        String errorType = readType(document);
+        if (errorType == null) {
+            var member = document.getMember("code");
+            if (member != null && member.type() == ShapeType.STRING) {
+                errorType = member.asString();
+            }
+        }
+        return errorType;
+    }
+
+    /**
+     * Removes the trailing URI in {@code __type} field or {@code code} field of the document.
+     *
+     * <p>For example, given {@code __type = "aws.protocoltests.restjson#FooError:http://abc.com"},
+     * protocols like restJSON should ignore the trailing URI and keep the namespace of the error type.
+     *
+     * @param text The error type string.
+     * @return The error type string without the trailing URI.
+     */
+    public static String removeUri(String text) {
+        if (text == null) {
+            return null;
+        }
+        var colon = text.indexOf(':');
+        if (colon > 0) {
+            text = text.substring(0, colon);
+        }
+        return text;
+    }
+
+    /**
+     * Removes the namespace and trailing URI in {@code __type} field or {@code code} field of the document.
+     *
+     * <p>For example, given {@code __type = "aws.protocoltests.restjson#FooError:http://abc.com"},
+     * protocols like awsJSON 1.1 should ignore the namespace and the trailing URI of the error type.
+     *
+     * @param text The error type string.
+     * @return The error type string without the trailing URI.
+     */
+    public static String removeNamespaceAndUri(String text) {
+        if (text == null) {
+            return null;
+        }
+        var hash = text.indexOf('#');
+        if (hash > 0) {
+            text = text.substring(hash + 1);
+        }
+        return removeUri(text);
+    }
+}

client/client-http/src/main/java/software/amazon/smithy/java/client/http/HttpErrorDeserializer.java

@@ -17,6 +17,7 @@
 import software.amazon.smithy.java.core.serde.TypeRegistry;
 import software.amazon.smithy.java.core.serde.document.DiscriminatorException;
 import software.amazon.smithy.java.core.serde.document.Document;
+import software.amazon.smithy.java.core.serde.document.DocumentDeserializer;
 import software.amazon.smithy.java.http.api.HttpResponse;
 import software.amazon.smithy.java.io.datastream.DataStream;
 import software.amazon.smithy.model.shapes.ShapeId;
@@ -132,15 +133,41 @@ public interface ErrorPayloadParser {
          *
          * @return the created error.
          */
-        CallException parsePayload(
+        default CallException parsePayload(
                 Context context,
                 Codec codec,
                 KnownErrorFactory knownErrorFactory,
                 ShapeId serviceId,
                 TypeRegistry typeRegistry,
                 HttpResponse response,
                 ByteBuffer buffer
-        ) throws SerializationException, DiscriminatorException;
+        ) {
+            var document = codec.createDeserializer(buffer).readDocument();
+            var id = extractErrorType(document, serviceId.getNamespace());
+            var builder = typeRegistry.createBuilder(id, ModeledException.class);
+            if (builder != null) {
+                return knownErrorFactory.createErrorFromDocument(
+                        context,
+                        codec,
+                        response,
+                        buffer,
+                        document,
+                        builder);
+            }
+            return null;
+        }
+
+        /**
+         * This method should extract the error type from converted document based on the
+         * protocol requirement from either __type or code and apply necessary sanitizing to
+         * the error type.
+         *
+         * @param document The converted document of the payload.
+         * @param namespace The default namespace used for error type's ShapeId.
+         *
+         * @return the created error.
+         */
+        ShapeId extractErrorType(Document document, String namespace);
     }
 
     // Does not check for any error headers by default.
@@ -191,30 +218,12 @@ public ModeledException createErrorFromDocument(
         }
     };
 
-    // This default parser should work for most protocols, but other  protocols
+    // This default parser should work for most protocols, but other protocols
     // that do not support document types will need a custom parser to extract error ShapeId.
-    private static final ErrorPayloadParser DEFAULT_ERROR_PAYLOAD_PARSER = (
-            Context context,
-            Codec codec,
-            KnownErrorFactory knownErrorFactory,
-            ShapeId serviceId,
-            TypeRegistry typeRegistry,
-            HttpResponse response,
-            ByteBuffer buffer) -> {
-        var document = codec.createDeserializer(buffer).readDocument();
-        var id = document.discriminator();
-        var builder = typeRegistry.createBuilder(id, ModeledException.class);
-        if (builder != null) {
-            return knownErrorFactory.createErrorFromDocument(
-                    context,
-                    codec,
-                    response,
-                    buffer,
-                    document,
-                    builder);
-        }
-        return null;
-    };
+    public static final ErrorPayloadParser DEFAULT_ERROR_PAYLOAD_PARSER =
+            (document, namespace) -> DocumentDeserializer.parseDiscriminator(
+                    ErrorTypeUtils.removeNamespaceAndUri(ErrorTypeUtils.readTypeAndCode(document)),
+                    namespace);
 
     private final Codec codec;
     private final HeaderErrorExtractor headerErrorExtractor;