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
|
package modserver
import (
"context"
grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware/v2"
"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/module/modshared"
"google.golang.org/grpc"
)
type rpcApiKeyType int
const (
rpcApiKey rpcApiKeyType = iota
)
// RpcApi provides the API for the module's gRPC handlers to use.
type RpcApi interface {
modshared.RpcApi
}
type RpcApiFactory func(ctx context.Context, fullMethodName string) RpcApi
func InjectRpcApi(ctx context.Context, rpcApi RpcApi) context.Context {
return context.WithValue(ctx, rpcApiKey, rpcApi)
}
func RpcApiFromContext(ctx context.Context) RpcApi {
rpcApi, ok := ctx.Value(rpcApiKey).(RpcApi)
if !ok {
// This is a programmer error, so panic.
panic("modserver.RpcApi not attached to context. Make sure you are using interceptors")
}
return rpcApi
}
// UnaryRpcApiInterceptor returns a new unary server interceptor that augments connection context with a RpcApi.
func UnaryRpcApiInterceptor(factory RpcApiFactory) grpc.UnaryServerInterceptor {
return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
return handler(augmentContextWithRpcApi(ctx, info.FullMethod, factory), req)
}
}
// StreamRpcApiInterceptor returns a new stream server interceptor that augments connection context with a RpcApi.
func StreamRpcApiInterceptor(factory RpcApiFactory) grpc.StreamServerInterceptor {
return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
wrapper := grpc_middleware.WrapServerStream(ss)
wrapper.WrappedContext = augmentContextWithRpcApi(wrapper.WrappedContext, info.FullMethod, factory)
return handler(srv, wrapper)
}
}
func augmentContextWithRpcApi(ctx context.Context, fullMethodName string, factory RpcApiFactory) context.Context {
return InjectRpcApi(ctx, factory(ctx, fullMethodName))
}
|
