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
|
package modserver
import (
"context"
grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware/v2"
"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/api"
"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/module/modshared"
"go.uber.org/zap"
"google.golang.org/grpc"
)
type agentRpcApiKeyType int
const (
agentRpcApiKey agentRpcApiKeyType = iota
)
// AgentRpcApi provides the API for the module's gRPC handlers to use.
// It should be used only by modules, that handle requests from agents.
type AgentRpcApi interface {
modshared.RpcApi
// AgentToken returns the token of an agent making the RPC.
AgentToken() api.AgentToken
// AgentInfo returns information about the agent making the RPC.
// Returns a gRPC-compatible error.
// Returns an error with the Unavailable code if there was a retriable error.
// If there was an error, it takes care of tracking it via HandleProcessingError().
AgentInfo(ctx context.Context, log *zap.Logger) (*api.AgentInfo, error)
}
type AgentRpcApiFactory func(ctx context.Context, fullMethodName string) (AgentRpcApi, error)
func InjectAgentRpcApi(ctx context.Context, rpcApi AgentRpcApi) context.Context {
return context.WithValue(ctx, agentRpcApiKey, rpcApi)
}
func AgentRpcApiFromContext(ctx context.Context) AgentRpcApi {
rpcApi, ok := ctx.Value(agentRpcApiKey).(AgentRpcApi)
if !ok {
// This is a programmer error, so panic.
panic("modserver.AgentRpcApi not attached to context. Make sure you are using interceptors")
}
return rpcApi
}
// UnaryAgentRpcApiInterceptor returns a new unary server interceptor that augments connection context with a AgentRpcApi.
func UnaryAgentRpcApiInterceptor(factory AgentRpcApiFactory) grpc.UnaryServerInterceptor {
return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
rpcApi, err := factory(ctx, info.FullMethod)
if err != nil {
return nil, err
}
return handler(InjectAgentRpcApi(ctx, rpcApi), req)
}
}
// StreamAgentRpcApiInterceptor returns a new stream server interceptor that augments connection context with a AgentRpcApi.
func StreamAgentRpcApiInterceptor(factory AgentRpcApiFactory) grpc.StreamServerInterceptor {
return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
wrapper := grpc_middleware.WrapServerStream(ss)
rpcApi, err := factory(wrapper.WrappedContext, info.FullMethod)
if err != nil {
return err
}
wrapper.WrappedContext = InjectAgentRpcApi(wrapper.WrappedContext, rpcApi)
return handler(srv, wrapper)
}
}
|
