fork of https://github.com/sourcegraph/zoekt
1package propagator
2
3import (
4 "context"
5
6 "google.golang.org/grpc"
7 "google.golang.org/grpc/metadata"
8)
9
10// Propagator is a type that can extract some information from a context.Context,
11// returning it in the form of metadata.MD and can also inject that same metadata
12// back into a context on the server side of an RPC call.
13type Propagator interface {
14 // FromContext extracts the information to be propagated from a context,
15 // converting it to a metadata.MD. This will be called on the client side
16 // of an RPC.
17 FromContext(context.Context) metadata.MD
18
19 // InjectContext takes a context and some metadata and creates a new context
20 // with the information from the metadata injected into the context.
21 // This will be called on the server side of an RPC.
22 InjectContext(context.Context, metadata.MD) (context.Context, error)
23}
24
25// StreamServerPropagator returns an interceptor that will use the given propagator
26// to translate some metadata back into the context for the RPC handler. The client
27// should be configured with an interceptor that uses the same propagator.
28func StreamServerPropagator(prop Propagator) grpc.StreamServerInterceptor {
29 return func(
30 srv any,
31 ss grpc.ServerStream,
32 info *grpc.StreamServerInfo,
33 handler grpc.StreamHandler,
34 ) error {
35 ctx := ss.Context()
36 md, ok := metadata.FromIncomingContext(ctx)
37 if ok {
38 var err error
39 ctx, err = prop.InjectContext(ss.Context(), md)
40 if err != nil {
41 return err
42 }
43 ss = contextedServerStream{ss, ctx}
44 }
45 return handler(srv, ss)
46 }
47}
48
49// UnaryServerPropagator returns an interceptor that will use the given propagator
50// to translate some metadata back into the context for the RPC handler. The client
51// should be configured with an interceptor that uses the same propagator.
52func UnaryServerPropagator(prop Propagator) grpc.UnaryServerInterceptor {
53 return func(
54 ctx context.Context,
55 req any,
56 info *grpc.UnaryServerInfo,
57 handler grpc.UnaryHandler,
58 ) (resp any, err error) {
59 md, ok := metadata.FromIncomingContext(ctx)
60 if ok {
61 ctx, err = prop.InjectContext(ctx, md)
62 if err != nil {
63 return nil, err
64 }
65 }
66 return handler(ctx, req)
67 }
68}
69
70type contextedServerStream struct {
71 grpc.ServerStream
72 ctx context.Context
73}
74
75func (css contextedServerStream) Context() context.Context {
76 return css.ctx
77}