1package http
 2
 3import (
 4	"context"
 5	"fmt"
 6	"net/http"
 7
 8	"github.com/aws/smithy-go/middleware"
 9)
10
11// WithHeaderComment instruments a middleware stack to append an HTTP field
12// comment to the given header as specified in RFC 9110
13// (https://www.rfc-editor.org/rfc/rfc9110#name-comments).
14//
15// The header is case-insensitive. If the provided header exists when the
16// middleware runs, the content will be inserted as-is enclosed in parentheses.
17//
18// Note that per the HTTP specification, comments are only allowed in fields
19// containing "comment" as part of their field value definition, but this API
20// will NOT verify whether the provided header is one of them.
21//
22// WithHeaderComment MAY be applied more than once to a middleware stack and/or
23// more than once per header.
24func WithHeaderComment(header, content string) func(*middleware.Stack) error {
25	return func(s *middleware.Stack) error {
26		m, err := getOrAddHeaderComment(s)
27		if err != nil {
28			return fmt.Errorf("get or add header comment: %v", err)
29		}
30
31		m.values.Add(header, content)
32		return nil
33	}
34}
35
36type headerCommentMiddleware struct {
37	values http.Header // hijack case-insensitive access APIs
38}
39
40func (*headerCommentMiddleware) ID() string {
41	return "headerComment"
42}
43
44func (m *headerCommentMiddleware) HandleBuild(ctx context.Context, in middleware.BuildInput, next middleware.BuildHandler) (
45	out middleware.BuildOutput, metadata middleware.Metadata, err error,
46) {
47	r, ok := in.Request.(*Request)
48	if !ok {
49		return out, metadata, fmt.Errorf("unknown transport type %T", in.Request)
50	}
51
52	for h, contents := range m.values {
53		for _, c := range contents {
54			if existing := r.Header.Get(h); existing != "" {
55				r.Header.Set(h, fmt.Sprintf("%s (%s)", existing, c))
56			}
57		}
58	}
59
60	return next.HandleBuild(ctx, in)
61}
62
63func getOrAddHeaderComment(s *middleware.Stack) (*headerCommentMiddleware, error) {
64	id := (*headerCommentMiddleware)(nil).ID()
65	m, ok := s.Build.Get(id)
66	if !ok {
67		m := &headerCommentMiddleware{values: http.Header{}}
68		if err := s.Build.Add(m, middleware.After); err != nil {
69			return nil, fmt.Errorf("add build: %v", err)
70		}
71
72		return m, nil
73	}
74
75	hc, ok := m.(*headerCommentMiddleware)
76	if !ok {
77		return nil, fmt.Errorf("existing middleware w/ id %s is not *headerCommentMiddleware", id)
78	}
79
80	return hc, nil
81}