github.com/aavshr/aws-sdk-go@v1.41.3/private/protocol/eventstream/eventstreamapi/signer.go (about)

     1  package eventstreamapi
     2  
     3  import (
     4  	"bytes"
     5  	"strings"
     6  	"time"
     7  
     8  	"github.com/aavshr/aws-sdk-go/private/protocol/eventstream"
     9  )
    10  
    11  var timeNow = time.Now
    12  
    13  // StreamSigner defines an interface for the implementation of signing of event stream payloads
    14  type StreamSigner interface {
    15  	GetSignature(headers, payload []byte, date time.Time) ([]byte, error)
    16  }
    17  
    18  // SignEncoder envelopes event stream messages
    19  // into an event stream message payload with included
    20  // signature headers using the provided signer and encoder.
    21  type SignEncoder struct {
    22  	signer     StreamSigner
    23  	encoder    Encoder
    24  	bufEncoder *BufferEncoder
    25  
    26  	closeErr error
    27  	closed   bool
    28  }
    29  
    30  // NewSignEncoder returns a new SignEncoder using the provided stream signer and
    31  // event stream encoder.
    32  func NewSignEncoder(signer StreamSigner, encoder Encoder) *SignEncoder {
    33  	// TODO: Need to pass down logging
    34  
    35  	return &SignEncoder{
    36  		signer:     signer,
    37  		encoder:    encoder,
    38  		bufEncoder: NewBufferEncoder(),
    39  	}
    40  }
    41  
    42  // Close encodes a final event stream signing envelope with an empty event stream
    43  // payload. This final end-frame is used to mark the conclusion of the stream.
    44  func (s *SignEncoder) Close() error {
    45  	if s.closed {
    46  		return s.closeErr
    47  	}
    48  
    49  	if err := s.encode([]byte{}); err != nil {
    50  		if strings.Contains(err.Error(), "on closed pipe") {
    51  			return nil
    52  		}
    53  
    54  		s.closeErr = err
    55  		s.closed = true
    56  		return s.closeErr
    57  	}
    58  
    59  	return nil
    60  }
    61  
    62  // Encode takes the provided message and add envelopes the message
    63  // with the required signature.
    64  func (s *SignEncoder) Encode(msg eventstream.Message) error {
    65  	payload, err := s.bufEncoder.Encode(msg)
    66  	if err != nil {
    67  		return err
    68  	}
    69  
    70  	return s.encode(payload)
    71  }
    72  
    73  func (s SignEncoder) encode(payload []byte) error {
    74  	date := timeNow()
    75  
    76  	var msg eventstream.Message
    77  	msg.Headers.Set(DateHeader, eventstream.TimestampValue(date))
    78  	msg.Payload = payload
    79  
    80  	var headers bytes.Buffer
    81  	if err := eventstream.EncodeHeaders(&headers, msg.Headers); err != nil {
    82  		return err
    83  	}
    84  
    85  	sig, err := s.signer.GetSignature(headers.Bytes(), msg.Payload, date)
    86  	if err != nil {
    87  		return err
    88  	}
    89  
    90  	msg.Headers.Set(ChunkSignatureHeader, eventstream.BytesValue(sig))
    91  
    92  	return s.encoder.Encode(msg)
    93  }
    94  
    95  // BufferEncoder is a utility that provides a buffered
    96  // event stream encoder
    97  type BufferEncoder struct {
    98  	encoder Encoder
    99  	buffer  *bytes.Buffer
   100  }
   101  
   102  // NewBufferEncoder returns a new BufferEncoder initialized
   103  // with a 1024 byte buffer.
   104  func NewBufferEncoder() *BufferEncoder {
   105  	buf := bytes.NewBuffer(make([]byte, 1024))
   106  	return &BufferEncoder{
   107  		encoder: eventstream.NewEncoder(buf),
   108  		buffer:  buf,
   109  	}
   110  }
   111  
   112  // Encode returns the encoded message as a byte slice.
   113  // The returned byte slice will be modified on the next encode call
   114  // and should not be held onto.
   115  func (e *BufferEncoder) Encode(msg eventstream.Message) ([]byte, error) {
   116  	e.buffer.Reset()
   117  
   118  	if err := e.encoder.Encode(msg); err != nil {
   119  		return nil, err
   120  	}
   121  
   122  	return e.buffer.Bytes(), nil
   123  }