diff --git a/deployment/environment/devenv/jd.go b/deployment/environment/devenv/jd.go index 48150340cae..818f9b09400 100644 --- a/deployment/environment/devenv/jd.go +++ b/deployment/environment/devenv/jd.go @@ -20,6 +20,7 @@ type JDConfig struct { WSRPC string Creds credentials.TransportCredentials Auth oauth2.TokenSource + GAP string NodeInfo []NodeInfo } @@ -44,14 +45,40 @@ func authTokenInterceptor(source oauth2.TokenSource) grpc.UnaryClientInterceptor } } +func gapTokenInterceptor(token string) grpc.UnaryClientInterceptor { + return func( + ctx context.Context, + method string, + req, reply any, + cc *grpc.ClientConn, + invoker grpc.UnaryInvoker, + opts ...grpc.CallOption, + ) error { + return invoker( + metadata.AppendToOutgoingContext(ctx, "x-authorization-github-jwt", "Bearer "+token), + method, req, reply, cc, opts..., + ) + } +} + func NewJDConnection(cfg JDConfig) (*grpc.ClientConn, error) { opts := []grpc.DialOption{} + interceptors := []grpc.UnaryClientInterceptor{} + if cfg.Creds != nil { opts = append(opts, grpc.WithTransportCredentials(cfg.Creds)) } if cfg.Auth != nil { - opts = append(opts, grpc.WithUnaryInterceptor(authTokenInterceptor(cfg.Auth))) + interceptors = append(interceptors, authTokenInterceptor(cfg.Auth)) } + if cfg.GAP != "" { + interceptors = append(interceptors, gapTokenInterceptor(cfg.GAP)) + } + + if len(interceptors) > 0 { + opts = append(opts, grpc.WithChainUnaryInterceptor(interceptors...)) + } + conn, err := grpc.NewClient(cfg.GRPC, opts...) if err != nil { return nil, fmt.Errorf("failed to connect Job Distributor service. Err: %w", err)