Add support for @eventHeader and @eventPayload for RPC protocols (#864)

* Add support for @eventHeader and @eventPayload for RPC protocols

Adds support to encode and decode members with the @eventHeader and
@eventPayload traits. For the decoding part, it's necessary to decode
the union member independently of the union itself for which we need
its builder, but we currently have no means to get one. For this, the
schema was enhanced wiht a method that to return a supplier to
generate new instances of builder for that schema which is only set
for structures and unions.

* Return the builder instead of a supplier

* Fix the javadoc for the shapeBuilder method

Author: sugmanue

aws/aws-event-streams/build.gradle.kts

@@ -7,8 +7,13 @@ description = "This module provides AWS event streaming support"
 extra["displayName"] = "Smithy :: Java :: AWS :: Event Streams"
 extra["moduleName"] = "software.amazon.smithy.java.aws.events"
 
+tasks.test {
+    dependsOn(":codecs:json-codec:shadowJar")
+}
+
 dependencies {
     api(project(":core"))
     implementation(project(":logging"))
     api("software.amazon.eventstream:eventstream:1.0.1")
+    testImplementation(project(":codecs:json-codec"))
 }

aws/aws-event-streams/src/main/java/software/amazon/smithy/java/aws/events/AwsEventShapeDecoder.java

@@ -5,9 +5,13 @@
 
 package software.amazon.smithy.java.aws.events;
 
+import java.nio.ByteBuffer;
+import java.time.Instant;
+import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.Flow;
 import java.util.function.Supplier;
+import software.amazon.eventstream.HeaderValue;
 import software.amazon.eventstream.Message;
 import software.amazon.smithy.java.core.schema.Schema;
 import software.amazon.smithy.java.core.schema.SerializableStruct;
@@ -52,7 +56,7 @@ public final class AwsEventShapeDecoder<E extends SerializableStruct, IR extends
     public SerializableStruct decode(AwsEventFrame frame) {
         var message = frame.unwrap();
         var eventType = getEventType(message);
-        if (initialEventType.getName().equals(eventType)) {
+        if (initialEventType.value().equals(eventType)) {
             return decodeInitialResponse(frame);
         }
         return decodeEvent(frame);
@@ -71,9 +75,14 @@ private E decodeEvent(AwsEventFrame frame) {
             throw new IllegalArgumentException("Unsupported event type: " + eventType);
         }
         var codecDeserializer = codec.createDeserializer(message.getPayload());
-        var eventDeserializer = new AwsEventDeserializer(memberSchema, codecDeserializer);
+        var headers = message.getHeaders();
+        var deserializer = new EventStreamDeserializer(codecDeserializer, new HeadersDeserializer(headers));
+        var memberTarget = memberSchema.memberTarget();
+        var shapeBuilder = memberTarget.shapeBuilder();
+        shapeBuilder.deserialize(deserializer);
         var builder = eventBuilder.get();
-        return builder.deserialize(eventDeserializer).build();
+        builder.setMemberValue(memberSchema, shapeBuilder.build());
+        return builder.build();
     }
 
     private IR decodeInitialResponse(AwsEventFrame frame) {
@@ -82,8 +91,13 @@ private IR decodeInitialResponse(AwsEventFrame frame) {
         var builder = initialEventBuilder.get();
         builder.deserialize(codecDeserializer);
         var publisherMember = getPublisherMember(builder.schema());
-        var responseDeserializer = new EventStreamDeserializer(publisherMember, publisher);
+        // Set the publisher member
+        var responseDeserializer = new InitialResponseDeserializer(publisherMember, publisher);
         builder.deserialize(responseDeserializer);
+        // Deserialize the rest of the members if any
+        var headers = message.getHeaders();
+        var deserializer = new EventStreamDeserializer(codecDeserializer, new HeadersDeserializer(headers));
+        builder.deserialize(deserializer);
         return builder.build();
     }
 
@@ -100,11 +114,11 @@ private String getEventType(Message message) {
         return message.getHeaders().get(":event-type").getString();
     }
 
-    static class EventStreamDeserializer extends SpecificShapeDeserializer {
+    static class InitialResponseDeserializer extends SpecificShapeDeserializer {
         private final Schema publisherMember;
         private final Flow.Publisher<? extends SerializableStruct> publisher;
 
-        EventStreamDeserializer(Schema publisherMember, Flow.Publisher<? extends SerializableStruct> publisher) {
+        InitialResponseDeserializer(Schema publisherMember, Flow.Publisher<? extends SerializableStruct> publisher) {
             this.publisherMember = publisherMember;
             this.publisher = publisher;
         }
@@ -119,4 +133,99 @@ public <T> void readStruct(Schema schema, T state, ShapeDeserializer.StructMembe
             consumer.accept(state, publisherMember, this);
         }
     }
+
+    static class EventStreamDeserializer extends SpecificShapeDeserializer {
+        private final ShapeDeserializer codecDeserializer;
+        private final HeadersDeserializer headersDeserializer;
+
+        EventStreamDeserializer(ShapeDeserializer codecDeserializer, HeadersDeserializer headersDeserializer) {
+            this.codecDeserializer = codecDeserializer;
+            this.headersDeserializer = headersDeserializer;
+        }
+
+        @Override
+        public <T> void readStruct(Schema schema, T builder, ShapeDeserializer.StructMemberConsumer<T> consumer) {
+            var payloadWritten = false;
+            for (Schema member : schema.members()) {
+                if (member.hasTrait(TraitKey.EVENT_HEADER_TRAIT)) {
+                    consumer.accept(builder, member, headersDeserializer);
+                } else if (member.hasTrait(TraitKey.EVENT_PAYLOAD_TRAIT)) {
+                    consumer.accept(builder, member, codecDeserializer);
+                    payloadWritten = true;
+                }
+            }
+            // Deserialize from the payload if still needed.
+            if (!payloadWritten) {
+                codecDeserializer.readStruct(schema, builder, consumer);
+            }
+        }
+    }
+
+    static class HeadersDeserializer extends SpecificShapeDeserializer {
+        private final Map<String, HeaderValue> headers;
+
+        HeadersDeserializer(Map<String, HeaderValue> headers) {
+            this.headers = headers;
+        }
+
+        @Override
+        public ByteBuffer readBlob(Schema schema) {
+            return getValueForShapeType(schema);
+        }
+
+        @Override
+
+        public byte readByte(Schema schema) {
+            return getValueForShapeType(schema);
+        }
+
+        @Override
+        public short readShort(Schema schema) {
+            return getValueForShapeType(schema);
+        }
+
+        @Override
+        public int readInteger(Schema schema) {
+            return getValueForShapeType(schema);
+        }
+
+        @Override
+        public long readLong(Schema schema) {
+            return getValueForShapeType(schema);
+        }
+
+        @Override
+        public String readString(Schema schema) {
+            return getValueForShapeType(schema);
+        }
+
+        @Override
+        public boolean readBoolean(Schema schema) {
+            return getValueForShapeType(schema);
+        }
+
+        @Override
+        public Instant readTimestamp(Schema schema) {
+            return getValueForShapeType(schema);
+        }
+
+        @SuppressWarnings("unchecked")
+        private <T> T getValueForShapeType(Schema member) {
+            HeaderValue value = headers.get(member.memberName());
+            if (value == null) {
+                return null;
+            }
+            return (T) switch (member.type()) {
+                case BLOB -> value.getByteBuffer();
+                case BOOLEAN -> value.getBoolean();
+                case BYTE -> value.getByte();
+                case SHORT -> value.getShort();
+                case INTEGER, INT_ENUM -> value.getInteger();
+                case LONG -> value.getLong();
+                case TIMESTAMP -> value.getTimestamp();
+                case STRING -> value.getString();
+                default -> throw new IllegalArgumentException("Unsupported shape type: " + member.type());
+            };
+        }
+    }
 }