src

Go monorepo.
git clone git://code.dwrz.net/src
Log | Files | Refs

eventstream_middleware.go (1922B)

      1 package http
      2 
      3 import (
      4 	"context"
      5 	"fmt"
      6 	"io"
      7 
      8 	"github.com/aws/smithy-go/middleware"
      9 )
     10 
     11 type eventStreamWriterKey struct{}
     12 
     13 // GetInputStreamWriter returns the io.WriteCloser pipe used for the
     14 // operation's input event stream.
     15 func GetInputStreamWriter(ctx context.Context) io.WriteCloser {
     16 	writeCloser, _ := middleware.GetStackValue(ctx, eventStreamWriterKey{}).(io.WriteCloser)
     17 	return writeCloser
     18 }
     19 
     20 func setInputStreamWriter(ctx context.Context, writeCloser io.WriteCloser) context.Context {
     21 	return middleware.WithStackValue(ctx, eventStreamWriterKey{}, writeCloser)
     22 }
     23 
     24 // InitializeStreamWriter is a Finalize middleware that creates an in-memory
     25 // pipe and sets it as the HTTP request body so event stream messages can be
     26 // written after the request is sent.
     27 type InitializeStreamWriter struct{}
     28 
     29 // AddInitializeStreamWriter adds the InitializeStreamWriter middleware to the
     30 // provided stack.
     31 func AddInitializeStreamWriter(stack *middleware.Stack) error {
     32 	return stack.Finalize.Add(&InitializeStreamWriter{}, middleware.After)
     33 }
     34 
     35 // ID returns the identifier for the middleware.
     36 func (i *InitializeStreamWriter) ID() string {
     37 	return "InitializeStreamWriter"
     38 }
     39 
     40 // HandleFinalize is the middleware implementation.
     41 func (i *InitializeStreamWriter) HandleFinalize(
     42 	ctx context.Context, in middleware.FinalizeInput, next middleware.FinalizeHandler,
     43 ) (
     44 	out middleware.FinalizeOutput, metadata middleware.Metadata, err error,
     45 ) {
     46 	request, ok := in.Request.(*Request)
     47 	if !ok {
     48 		return out, metadata, fmt.Errorf("unknown transport type: %T", in.Request)
     49 	}
     50 
     51 	inputReader, inputWriter := io.Pipe()
     52 	defer func() {
     53 		if err == nil {
     54 			return
     55 		}
     56 		_ = inputReader.Close()
     57 		_ = inputWriter.Close()
     58 	}()
     59 
     60 	request, err = request.SetStream(inputReader)
     61 	if err != nil {
     62 		return out, metadata, err
     63 	}
     64 	in.Request = request
     65 
     66 	ctx = setInputStreamWriter(ctx, inputWriter)
     67 
     68 	return next.HandleFinalize(ctx, in)
     69 }