-
Notifications
You must be signed in to change notification settings - Fork 246
/
requestid.go
120 lines (99 loc) · 3.33 KB
/
requestid.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
package requestid
import (
"context"
"math/rand"
log "github.com/authzed/spicedb/internal/logging"
"github.com/authzed/authzed-go/pkg/responsemeta"
"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
)
// RequestIDMetadataKey is the key in which request IDs are passed to metadata.
const RequestIDMetadataKey = "x-request-id"
// Option instances control how the middleware is initialized.
type Option func(*handleRequestID)
// GenerateIfMissing will instruct the middleware to create a request ID if one
// isn't already on the incoming request.
//
// default: false
func GenerateIfMissing(enable bool) Option {
return func(reporter *handleRequestID) {
reporter.generateIfMissing = enable
}
}
// IDGenerator functions are used to generate request IDs if a new one is needed.
type IDGenerator func() string
// WithIDGenerator gives the middleware a function to use for generating requestIDs.
//
// default: 32 character hex string
func WithIDGenerator(genFunc IDGenerator) Option {
return func(reporter *handleRequestID) {
reporter.requestIDGenerator = genFunc
}
}
type handleRequestID struct {
generateIfMissing bool
requestIDGenerator IDGenerator
}
func (r *handleRequestID) ServerReporter(ctx context.Context, _ interceptors.CallMeta) (interceptors.Reporter, context.Context) {
var requestID string
var haveRequestID bool
md, ok := metadata.FromIncomingContext(ctx)
if ok {
var requestIDs []string
requestIDs, haveRequestID = md[RequestIDMetadataKey]
if haveRequestID {
requestID = requestIDs[0]
}
}
if !haveRequestID && r.generateIfMissing {
requestID, haveRequestID = r.requestIDGenerator(), true
// Inject the newly generated request ID into the metadata
md.Set(RequestIDMetadataKey, requestID)
ctx = metadata.NewIncomingContext(ctx, md)
}
if haveRequestID {
ctx = metadata.AppendToOutgoingContext(ctx, RequestIDMetadataKey, requestID)
err := responsemeta.SetResponseHeaderMetadata(ctx, map[responsemeta.ResponseMetadataHeaderKey]string{
responsemeta.RequestID: requestID,
})
// if context is cancelled, the stream will be closed, and gRPC will return ErrIllegalHeaderWrite
// this prevents logging unnecessary error messages
if ctx.Err() != nil {
return interceptors.NoopReporter{}, ctx
}
if err != nil {
log.Ctx(ctx).Warn().Err(err).Msg("requestid: could not report metadata")
}
}
return interceptors.NoopReporter{}, ctx
}
// UnaryServerInterceptor returns a new interceptor which handles request IDs according
// to the provided options.
func UnaryServerInterceptor(opts ...Option) grpc.UnaryServerInterceptor {
return interceptors.UnaryServerInterceptor(createReporter(opts))
}
// StreamServerInterceptor returns a new interceptor which handles request IDs according
// to the provided options.
func StreamServerInterceptor(opts ...Option) grpc.StreamServerInterceptor {
return interceptors.StreamServerInterceptor(createReporter(opts))
}
func createReporter(opts []Option) *handleRequestID {
reporter := &handleRequestID{
requestIDGenerator: func() string {
return randSeq(32)
},
}
for _, opt := range opts {
opt(reporter)
}
return reporter
}
var letters = []rune("0123456789abcdef")
func randSeq(n int) string {
b := make([]rune, n)
for i := range b {
b[i] = letters[rand.Intn(len(letters))]
}
return string(b)
}