netbird/management/internals/server/grpc_extension.go

package server

import (
	"context"

	"google.golang.org/grpc"
)

// GRPCExtension bundles an external module's contribution to the management
// gRPC server: the registration of one or more services onto the shared
// grpc.Server, any server-wide interceptors those services require, and an
// optional shutdown hook. It is a generic extension point with no knowledge of
// any specific service.
type GRPCExtension struct {
	// Register is invoked with the shared grpc.Server (as a ServiceRegistrar)
	// after the built-in services are registered. It may register any number of
	// services. May be nil.
	Register func(grpc.ServiceRegistrar)
	// UnaryInterceptors are appended to the server's unary interceptor chain,
	// running after the built-in interceptors. May be empty.
	UnaryInterceptors []grpc.UnaryServerInterceptor
	// StreamInterceptors are appended to the server's stream interceptor chain,
	// running after the built-in interceptors. May be empty.
	StreamInterceptors []grpc.StreamServerInterceptor
	// Shutdown, if non-nil, is called once during Stop() with the context
	// governing server shutdown, which carries a deadline. The hook MUST
	// return promptly and MUST abandon its work once that context is
	// cancelled or expires: it runs before the rest of Stop()'s cleanup
	// (store, event store, embedded IdP) and before Stop() itself checks the
	// context's deadline, so a hook that ignores the context will delay all
	// of that cleanup and prevent Stop() from returning on time. May be nil.
	Shutdown func(ctx context.Context)
}

// RegisterGRPCExtension registers a gRPC extension. Call before the gRPC server
// is first built (i.e. before Start); registrations after that have no effect.
func (s *BaseServer) RegisterGRPCExtension(ext GRPCExtension) {
	s.grpcExtensions = append(s.grpcExtensions, ext)
}

// appendExtensionInterceptors appends each extension's interceptors to the gRPC
// server options as additional chained interceptors. grpc.ChainUnaryInterceptor
// and grpc.ChainStreamInterceptor are additive, so the returned options run the
// extension interceptors after any interceptors already present in opts.
func appendExtensionInterceptors(opts []grpc.ServerOption, exts []GRPCExtension) []grpc.ServerOption {
	for _, ext := range exts {
		if len(ext.UnaryInterceptors) > 0 {
			opts = append(opts, grpc.ChainUnaryInterceptor(ext.UnaryInterceptors...))
		}
		if len(ext.StreamInterceptors) > 0 {
			opts = append(opts, grpc.ChainStreamInterceptor(ext.StreamInterceptors...))
		}
	}
	return opts
}

// registerExtensions registers each extension's services onto reg.
func registerExtensions(reg grpc.ServiceRegistrar, exts []GRPCExtension) {
	for _, ext := range exts {
		if ext.Register != nil {
			ext.Register(reg)
		}
	}
}

// runExtensionShutdownHooks calls each extension's shutdown hook, if set,
// passing ctx through so hooks can honor its deadline/cancellation.
func runExtensionShutdownHooks(ctx context.Context, exts []GRPCExtension) {
	for _, ext := range exts {
		if ext.Shutdown != nil {
			ext.Shutdown(ctx)
		}
	}
}