src

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

middleware_headers.go (5259B)

      1 package http
      2 
      3 import (
      4 	"context"
      5 	"fmt"
      6 
      7 	"github.com/aws/smithy-go/middleware"
      8 )
      9 
     10 type isContentTypeAutoSet struct{}
     11 
     12 // SetIsContentTypeDefaultValue returns a Context specifying if the request's
     13 // content-type header was set to a default value.
     14 func SetIsContentTypeDefaultValue(ctx context.Context, isDefault bool) context.Context {
     15 	return context.WithValue(ctx, isContentTypeAutoSet{}, isDefault)
     16 }
     17 
     18 // GetIsContentTypeDefaultValue returns if the content-type HTTP header on the
     19 // request is a default value that was auto assigned by an operation
     20 // serializer. Allows middleware post serialization to know if the content-type
     21 // was auto set to a default value or not.
     22 //
     23 // Also returns false if the Context value was never updated to include if
     24 // content-type was set to a default value.
     25 func GetIsContentTypeDefaultValue(ctx context.Context) bool {
     26 	v, _ := ctx.Value(isContentTypeAutoSet{}).(bool)
     27 	return v
     28 }
     29 
     30 // AddNoPayloadDefaultContentTypeRemover Adds the DefaultContentTypeRemover
     31 // middleware to the stack after the operation serializer. This middleware will
     32 // remove the content-type header from the request if it was set as a default
     33 // value, and no request payload is present.
     34 //
     35 // Returns error if unable to add the middleware.
     36 func AddNoPayloadDefaultContentTypeRemover(stack *middleware.Stack) (err error) {
     37 	err = stack.Serialize.Insert(removeDefaultContentType{},
     38 		"OperationSerializer", middleware.After)
     39 	if err != nil {
     40 		return fmt.Errorf("failed to add %s serialize middleware, %w",
     41 			removeDefaultContentType{}.ID(), err)
     42 	}
     43 
     44 	return nil
     45 }
     46 
     47 // RemoveNoPayloadDefaultContentTypeRemover removes the
     48 // DefaultContentTypeRemover middleware from the stack. Returns an error if
     49 // unable to remove the middleware.
     50 func RemoveNoPayloadDefaultContentTypeRemover(stack *middleware.Stack) (err error) {
     51 	_, err = stack.Serialize.Remove(removeDefaultContentType{}.ID())
     52 	if err != nil {
     53 		return fmt.Errorf("failed to remove %s serialize middleware, %w",
     54 			removeDefaultContentType{}.ID(), err)
     55 
     56 	}
     57 	return nil
     58 }
     59 
     60 // removeDefaultContentType provides after serialization middleware that will
     61 // remove the content-type header from an HTTP request if the header was set as
     62 // a default value by the operation serializer, and there is no request payload.
     63 type removeDefaultContentType struct{}
     64 
     65 // ID returns the middleware ID
     66 func (removeDefaultContentType) ID() string { return "RemoveDefaultContentType" }
     67 
     68 // HandleSerialize implements the serialization middleware.
     69 func (removeDefaultContentType) HandleSerialize(
     70 	ctx context.Context, input middleware.SerializeInput, next middleware.SerializeHandler,
     71 ) (
     72 	out middleware.SerializeOutput, meta middleware.Metadata, err error,
     73 ) {
     74 	req, ok := input.Request.(*Request)
     75 	if !ok {
     76 		return out, meta, fmt.Errorf(
     77 			"unexpected request type %T for removeDefaultContentType middleware",
     78 			input.Request)
     79 	}
     80 
     81 	if GetIsContentTypeDefaultValue(ctx) && req.GetStream() == nil {
     82 		req.Header.Del("Content-Type")
     83 		input.Request = req
     84 	}
     85 
     86 	return next.HandleSerialize(ctx, input)
     87 }
     88 
     89 type headerValue struct {
     90 	header string
     91 	value  string
     92 	append bool
     93 }
     94 
     95 type headerValueHelper struct {
     96 	headerValues []headerValue
     97 }
     98 
     99 func (h *headerValueHelper) addHeaderValue(value headerValue) {
    100 	h.headerValues = append(h.headerValues, value)
    101 }
    102 
    103 func (h *headerValueHelper) ID() string {
    104 	return "HTTPHeaderHelper"
    105 }
    106 
    107 func (h *headerValueHelper) HandleBuild(ctx context.Context, in middleware.BuildInput, next middleware.BuildHandler) (out middleware.BuildOutput, metadata middleware.Metadata, err error) {
    108 	req, ok := in.Request.(*Request)
    109 	if !ok {
    110 		return out, metadata, fmt.Errorf("unknown transport type %T", in.Request)
    111 	}
    112 
    113 	for _, value := range h.headerValues {
    114 		if value.append {
    115 			req.Header.Add(value.header, value.value)
    116 		} else {
    117 			req.Header.Set(value.header, value.value)
    118 		}
    119 	}
    120 
    121 	return next.HandleBuild(ctx, in)
    122 }
    123 
    124 func getOrAddHeaderValueHelper(stack *middleware.Stack) (*headerValueHelper, error) {
    125 	id := (*headerValueHelper)(nil).ID()
    126 	m, ok := stack.Build.Get(id)
    127 	if !ok {
    128 		m = &headerValueHelper{}
    129 		err := stack.Build.Add(m, middleware.After)
    130 		if err != nil {
    131 			return nil, err
    132 		}
    133 	}
    134 
    135 	requestUserAgent, ok := m.(*headerValueHelper)
    136 	if !ok {
    137 		return nil, fmt.Errorf("%T for %s middleware did not match expected type", m, id)
    138 	}
    139 
    140 	return requestUserAgent, nil
    141 }
    142 
    143 // AddHeaderValue returns a stack mutator that adds the header value pair to header.
    144 // Appends to any existing values if present.
    145 func AddHeaderValue(header string, value string) func(stack *middleware.Stack) error {
    146 	return func(stack *middleware.Stack) error {
    147 		helper, err := getOrAddHeaderValueHelper(stack)
    148 		if err != nil {
    149 			return err
    150 		}
    151 		helper.addHeaderValue(headerValue{header: header, value: value, append: true})
    152 		return nil
    153 	}
    154 }
    155 
    156 // SetHeaderValue returns a stack mutator that adds the header value pair to header.
    157 // Replaces any existing values if present.
    158 func SetHeaderValue(header string, value string) func(stack *middleware.Stack) error {
    159 	return func(stack *middleware.Stack) error {
    160 		helper, err := getOrAddHeaderValueHelper(stack)
    161 		if err != nil {
    162 			return err
    163 		}
    164 		helper.addHeaderValue(headerValue{header: header, value: value, append: false})
    165 		return nil
    166 	}
    167 }