Add Infra

2025-03-03 18:59:25 +07:00
parent c02d8cda69
commit c9f1e33c24
1782 changed files with 692894 additions and 21 deletions
@@ -0,0 +1,81 @@
+package eventstreamapi
+
+import (
+	"fmt"
+	"sync"
+)
+
+// InputWriterCloseErrorCode is used to denote an error occurred
+// while closing the event stream input writer.
+const InputWriterCloseErrorCode = "EventStreamInputWriterCloseError"
+
+type messageError struct {
+	code string
+	msg  string
+}
+
+func (e messageError) Code() string {
+	return e.code
+}
+
+func (e messageError) Message() string {
+	return e.msg
+}
+
+func (e messageError) Error() string {
+	return fmt.Sprintf("%s: %s", e.code, e.msg)
+}
+
+func (e messageError) OrigErr() error {
+	return nil
+}
+
+// OnceError wraps the behavior of recording an error
+// once and signal on a channel when this has occurred.
+// Signaling is done by closing of the channel.
+//
+// Type is safe for concurrent usage.
+type OnceError struct {
+	mu  sync.RWMutex
+	err error
+	ch  chan struct{}
+}
+
+// NewOnceError return a new OnceError
+func NewOnceError() *OnceError {
+	return &OnceError{
+		ch: make(chan struct{}, 1),
+	}
+}
+
+// Err acquires a read-lock and returns an
+// error if one has been set.
+func (e *OnceError) Err() error {
+	e.mu.RLock()
+	err := e.err
+	e.mu.RUnlock()
+
+	return err
+}
+
+// SetError acquires a write-lock and will set
+// the underlying error value if one has not been set.
+func (e *OnceError) SetError(err error) {
+	if err == nil {
+		return
+	}
+
+	e.mu.Lock()
+	if e.err == nil {
+		e.err = err
+		close(e.ch)
+	}
+	e.mu.Unlock()
+}
+
+// ErrorSet returns a channel that will be used to signal
+// that an error has been set. This channel will be closed
+// when the error value has been set for OnceError.
+func (e *OnceError) ErrorSet() <-chan struct{} {
+	return e.ch
+}
@@ -0,0 +1,173 @@
+package eventstreamapi
+
+import (
+	"fmt"
+
+	"github.com/aws/aws-sdk-go/private/protocol"
+	"github.com/aws/aws-sdk-go/private/protocol/eventstream"
+)
+
+// Unmarshaler provides the interface for unmarshaling a EventStream
+// message into a SDK type.
+type Unmarshaler interface {
+	UnmarshalEvent(protocol.PayloadUnmarshaler, eventstream.Message) error
+}
+
+// EventReader provides reading from the EventStream of an reader.
+type EventReader struct {
+	decoder *eventstream.Decoder
+
+	unmarshalerForEventType func(string) (Unmarshaler, error)
+	payloadUnmarshaler      protocol.PayloadUnmarshaler
+
+	payloadBuf []byte
+}
+
+// NewEventReader returns a EventReader built from the reader and unmarshaler
+// provided.  Use ReadStream method to start reading from the EventStream.
+func NewEventReader(
+	decoder *eventstream.Decoder,
+	payloadUnmarshaler protocol.PayloadUnmarshaler,
+	unmarshalerForEventType func(string) (Unmarshaler, error),
+) *EventReader {
+	return &EventReader{
+		decoder:                 decoder,
+		payloadUnmarshaler:      payloadUnmarshaler,
+		unmarshalerForEventType: unmarshalerForEventType,
+		payloadBuf:              make([]byte, 10*1024),
+	}
+}
+
+// ReadEvent attempts to read a message from the EventStream and return the
+// unmarshaled event value that the message is for.
+//
+// For EventStream API errors check if the returned error satisfies the
+// awserr.Error interface to get the error's Code and Message components.
+//
+// EventUnmarshalers called with EventStream messages must take copies of the
+// message's Payload. The payload will is reused between events read.
+func (r *EventReader) ReadEvent() (event interface{}, err error) {
+	msg, err := r.decoder.Decode(r.payloadBuf)
+	if err != nil {
+		return nil, err
+	}
+	defer func() {
+		// Reclaim payload buffer for next message read.
+		r.payloadBuf = msg.Payload[0:0]
+	}()
+
+	typ, err := GetHeaderString(msg, MessageTypeHeader)
+	if err != nil {
+		return nil, err
+	}
+
+	switch typ {
+	case EventMessageType:
+		return r.unmarshalEventMessage(msg)
+	case ExceptionMessageType:
+		return nil, r.unmarshalEventException(msg)
+	case ErrorMessageType:
+		return nil, r.unmarshalErrorMessage(msg)
+	default:
+		return nil, &UnknownMessageTypeError{
+			Type: typ, Message: msg.Clone(),
+		}
+	}
+}
+
+// UnknownMessageTypeError provides an error when a message is received from
+// the stream, but the reader is unable to determine what kind of message it is.
+type UnknownMessageTypeError struct {
+	Type    string
+	Message eventstream.Message
+}
+
+func (e *UnknownMessageTypeError) Error() string {
+	return "unknown eventstream message type, " + e.Type
+}
+
+func (r *EventReader) unmarshalEventMessage(
+	msg eventstream.Message,
+) (event interface{}, err error) {
+	eventType, err := GetHeaderString(msg, EventTypeHeader)
+	if err != nil {
+		return nil, err
+	}
+
+	ev, err := r.unmarshalerForEventType(eventType)
+	if err != nil {
+		return nil, err
+	}
+
+	err = ev.UnmarshalEvent(r.payloadUnmarshaler, msg)
+	if err != nil {
+		return nil, err
+	}
+
+	return ev, nil
+}
+
+func (r *EventReader) unmarshalEventException(
+	msg eventstream.Message,
+) (err error) {
+	eventType, err := GetHeaderString(msg, ExceptionTypeHeader)
+	if err != nil {
+		return err
+	}
+
+	ev, err := r.unmarshalerForEventType(eventType)
+	if err != nil {
+		return err
+	}
+
+	err = ev.UnmarshalEvent(r.payloadUnmarshaler, msg)
+	if err != nil {
+		return err
+	}
+
+	var ok bool
+	err, ok = ev.(error)
+	if !ok {
+		err = messageError{
+			code: "SerializationError",
+			msg: fmt.Sprintf(
+				"event stream exception %s mapped to non-error %T, %v",
+				eventType, ev, ev,
+			),
+		}
+	}
+
+	return err
+}
+
+func (r *EventReader) unmarshalErrorMessage(msg eventstream.Message) (err error) {
+	var msgErr messageError
+
+	msgErr.code, err = GetHeaderString(msg, ErrorCodeHeader)
+	if err != nil {
+		return err
+	}
+
+	msgErr.msg, err = GetHeaderString(msg, ErrorMessageHeader)
+	if err != nil {
+		return err
+	}
+
+	return msgErr
+}
+
+// GetHeaderString returns the value of the header as a string. If the header
+// is not set or the value is not a string an error will be returned.
+func GetHeaderString(msg eventstream.Message, headerName string) (string, error) {
+	headerVal := msg.Headers.Get(headerName)
+	if headerVal == nil {
+		return "", fmt.Errorf("error header %s not present", headerName)
+	}
+
+	v, ok := headerVal.Get().(string)
+	if !ok {
+		return "", fmt.Errorf("error header value is not a string, %T", headerVal)
+	}
+
+	return v, nil
+}
@@ -0,0 +1,23 @@
+package eventstreamapi
+
+// EventStream headers with specific meaning to async API functionality.
+const (
+	ChunkSignatureHeader = `:chunk-signature` // chunk signature for message
+	DateHeader           = `:date`            // Date header for signature
+
+	// Message header and values
+	MessageTypeHeader    = `:message-type` // Identifies type of message.
+	EventMessageType     = `event`
+	ErrorMessageType     = `error`
+	ExceptionMessageType = `exception`
+
+	// Message Events
+	EventTypeHeader = `:event-type` // Identifies message event type e.g. "Stats".
+
+	// Message Error
+	ErrorCodeHeader    = `:error-code`
+	ErrorMessageHeader = `:error-message`
+
+	// Message Exception
+	ExceptionTypeHeader = `:exception-type`
+)
@@ -0,0 +1,123 @@
+package eventstreamapi
+
+import (
+	"bytes"
+	"strings"
+	"time"
+
+	"github.com/aws/aws-sdk-go/private/protocol/eventstream"
+)
+
+var timeNow = time.Now
+
+// StreamSigner defines an interface for the implementation of signing of event stream payloads
+type StreamSigner interface {
+	GetSignature(headers, payload []byte, date time.Time) ([]byte, error)
+}
+
+// SignEncoder envelopes event stream messages
+// into an event stream message payload with included
+// signature headers using the provided signer and encoder.
+type SignEncoder struct {
+	signer     StreamSigner
+	encoder    Encoder
+	bufEncoder *BufferEncoder
+
+	closeErr error
+	closed   bool
+}
+
+// NewSignEncoder returns a new SignEncoder using the provided stream signer and
+// event stream encoder.
+func NewSignEncoder(signer StreamSigner, encoder Encoder) *SignEncoder {
+	// TODO: Need to pass down logging
+
+	return &SignEncoder{
+		signer:     signer,
+		encoder:    encoder,
+		bufEncoder: NewBufferEncoder(),
+	}
+}
+
+// Close encodes a final event stream signing envelope with an empty event stream
+// payload. This final end-frame is used to mark the conclusion of the stream.
+func (s *SignEncoder) Close() error {
+	if s.closed {
+		return s.closeErr
+	}
+
+	if err := s.encode([]byte{}); err != nil {
+		if strings.Contains(err.Error(), "on closed pipe") {
+			return nil
+		}
+
+		s.closeErr = err
+		s.closed = true
+		return s.closeErr
+	}
+
+	return nil
+}
+
+// Encode takes the provided message and add envelopes the message
+// with the required signature.
+func (s *SignEncoder) Encode(msg eventstream.Message) error {
+	payload, err := s.bufEncoder.Encode(msg)
+	if err != nil {
+		return err
+	}
+
+	return s.encode(payload)
+}
+
+func (s SignEncoder) encode(payload []byte) error {
+	date := timeNow()
+
+	var msg eventstream.Message
+	msg.Headers.Set(DateHeader, eventstream.TimestampValue(date))
+	msg.Payload = payload
+
+	var headers bytes.Buffer
+	if err := eventstream.EncodeHeaders(&headers, msg.Headers); err != nil {
+		return err
+	}
+
+	sig, err := s.signer.GetSignature(headers.Bytes(), msg.Payload, date)
+	if err != nil {
+		return err
+	}
+
+	msg.Headers.Set(ChunkSignatureHeader, eventstream.BytesValue(sig))
+
+	return s.encoder.Encode(msg)
+}
+
+// BufferEncoder is a utility that provides a buffered
+// event stream encoder
+type BufferEncoder struct {
+	encoder Encoder
+	buffer  *bytes.Buffer
+}
+
+// NewBufferEncoder returns a new BufferEncoder initialized
+// with a 1024 byte buffer.
+func NewBufferEncoder() *BufferEncoder {
+	buf := bytes.NewBuffer(make([]byte, 1024))
+	return &BufferEncoder{
+		encoder: eventstream.NewEncoder(buf),
+		buffer:  buf,
+	}
+}
+
+// Encode returns the encoded message as a byte slice.
+// The returned byte slice will be modified on the next encode call
+// and should not be held onto.
+func (e *BufferEncoder) Encode(msg eventstream.Message) ([]byte, error) {
+	e.buffer.Reset()
+
+	if err := e.encoder.Encode(msg); err != nil {
+		return nil, err
+	}
+
+	return e.buffer.Bytes(), nil
+}
@@ -0,0 +1,129 @@
+package eventstreamapi
+
+import (
+	"fmt"
+	"io"
+	"sync"
+
+	"github.com/aws/aws-sdk-go/aws"
+)
+
+// StreamWriter provides concurrent safe writing to an event stream.
+type StreamWriter struct {
+	eventWriter *EventWriter
+	stream      chan eventWriteAsyncReport
+
+	done      chan struct{}
+	closeOnce sync.Once
+	err       *OnceError
+
+	streamCloser io.Closer
+}
+
+// NewStreamWriter returns a StreamWriter for the event writer, and stream
+// closer provided.
+func NewStreamWriter(eventWriter *EventWriter, streamCloser io.Closer) *StreamWriter {
+	w := &StreamWriter{
+		eventWriter:  eventWriter,
+		streamCloser: streamCloser,
+		stream:       make(chan eventWriteAsyncReport),
+		done:         make(chan struct{}),
+		err:          NewOnceError(),
+	}
+	go w.writeStream()
+
+	return w
+}
+
+// Close terminates the writers ability to write new events to the stream. Any
+// future call to Send will fail with an error.
+func (w *StreamWriter) Close() error {
+	w.closeOnce.Do(w.safeClose)
+	return w.Err()
+}
+
+func (w *StreamWriter) safeClose() {
+	close(w.done)
+}
+
+// ErrorSet returns a channel which will be closed
+// if an error occurs.
+func (w *StreamWriter) ErrorSet() <-chan struct{} {
+	return w.err.ErrorSet()
+}
+
+// Err returns any error that occurred while attempting to write an event to the
+// stream.
+func (w *StreamWriter) Err() error {
+	return w.err.Err()
+}
+
+// Send writes a single event to the stream returning an error if the write
+// failed.
+//
+// Send may be called concurrently. Events will be written to the stream
+// safely.
+func (w *StreamWriter) Send(ctx aws.Context, event Marshaler) error {
+	if err := w.Err(); err != nil {
+		return err
+	}
+
+	resultCh := make(chan error)
+	wrapped := eventWriteAsyncReport{
+		Event:  event,
+		Result: resultCh,
+	}
+
+	select {
+	case w.stream <- wrapped:
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-w.done:
+		return fmt.Errorf("stream closed, unable to send event")
+	}
+
+	select {
+	case err := <-resultCh:
+		return err
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-w.done:
+		return fmt.Errorf("stream closed, unable to send event")
+	}
+}
+
+func (w *StreamWriter) writeStream() {
+	defer w.Close()
+
+	for {
+		select {
+		case wrapper := <-w.stream:
+			err := w.eventWriter.WriteEvent(wrapper.Event)
+			wrapper.ReportResult(w.done, err)
+			if err != nil {
+				w.err.SetError(err)
+				return
+			}
+
+		case <-w.done:
+			if err := w.streamCloser.Close(); err != nil {
+				w.err.SetError(err)
+			}
+			return
+		}
+	}
+}
+
+type eventWriteAsyncReport struct {
+	Event  Marshaler
+	Result chan<- error
+}
+
+func (e eventWriteAsyncReport) ReportResult(cancel <-chan struct{}, err error) bool {
+	select {
+	case e.Result <- err:
+		return true
+	case <-cancel:
+		return false
+	}
+}
@@ -0,0 +1,10 @@
+//go:build go1.18
+// +build go1.18
+
+package eventstreamapi
+
+import "github.com/aws/aws-sdk-go/aws/request"
+
+// ApplyHTTPTransportFixes is a no-op for Go 1.18 and above.
+func ApplyHTTPTransportFixes(r *request.Request) {
+}
@@ -0,0 +1,19 @@
+//go:build !go1.18
+// +build !go1.18
+
+package eventstreamapi
+
+import "github.com/aws/aws-sdk-go/aws/request"
+
+// ApplyHTTPTransportFixes applies fixes to the HTTP request for proper event
+// stream functionality. Go 1.15 through 1.17 HTTP client could hang forever
+// when an HTTP/2 connection failed with an non-200 status code and err. Using
+// Expect 100-Continue, allows the HTTP client to gracefully handle the non-200
+// status code, and close the connection.
+//
+// This is a no-op for Go 1.18 and above.
+func ApplyHTTPTransportFixes(r *request.Request) {
+	r.Handlers.Sign.PushBack(func(r *request.Request) {
+		r.HTTPRequest.Header.Set("Expect", "100-Continue")
+	})
+}
@@ -0,0 +1,63 @@
+package eventstreamapi
+
+import (
+	"github.com/aws/aws-sdk-go/private/protocol"
+	"github.com/aws/aws-sdk-go/private/protocol/eventstream"
+)
+
+// Marshaler provides a marshaling interface for event types to event stream
+// messages.
+type Marshaler interface {
+	MarshalEvent(protocol.PayloadMarshaler) (eventstream.Message, error)
+}
+
+// Encoder is an stream encoder that will encode an event stream message for
+// the transport.
+type Encoder interface {
+	Encode(eventstream.Message) error
+}
+
+// EventWriter provides a wrapper around the underlying event stream encoder
+// for an io.WriteCloser.
+type EventWriter struct {
+	encoder          Encoder
+	payloadMarshaler protocol.PayloadMarshaler
+	eventTypeFor     func(Marshaler) (string, error)
+}
+
+// NewEventWriter returns a new event stream writer, that will write to the
+// writer provided. Use the WriteEvent method to write an event to the stream.
+func NewEventWriter(encoder Encoder, pm protocol.PayloadMarshaler, eventTypeFor func(Marshaler) (string, error),
+) *EventWriter {
+	return &EventWriter{
+		encoder:          encoder,
+		payloadMarshaler: pm,
+		eventTypeFor:     eventTypeFor,
+	}
+}
+
+// WriteEvent writes an event to the stream. Returns an error if the event
+// fails to marshal into a message, or writing to the underlying writer fails.
+func (w *EventWriter) WriteEvent(event Marshaler) error {
+	msg, err := w.marshal(event)
+	if err != nil {
+		return err
+	}
+
+	return w.encoder.Encode(msg)
+}
+
+func (w *EventWriter) marshal(event Marshaler) (eventstream.Message, error) {
+	eventType, err := w.eventTypeFor(event)
+	if err != nil {
+		return eventstream.Message{}, err
+	}
+
+	msg, err := event.MarshalEvent(w.payloadMarshaler)
+	if err != nil {
+		return eventstream.Message{}, err
+	}
+
+	msg.Headers.Set(EventTypeHeader, eventstream.StringValue(eventType))
+	return msg, nil
+}