feat: add worker control plane grpc server

2026-07-31 11:46:40 +08:00 · 2026-07-31 11:46:40 +08:00 · 2f8a62cad7
commit 2f8a62cad7
parent 5a678dc66f
4 changed files with 411 additions and 0 deletions
--- a/internal/controller/worker/identity.go
+++ b/internal/controller/worker/identity.go
@ -0,0 +1,72 @@
+package worker
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/peer"
+)
+
+var ErrUnauthorizedIdentity = errors.New("worker identity is not authorized")
+
+// AllowLoopbackIdentity is used only by the validated loopback plaintext mode.
+// It deliberately does not inspect transport credentials because that mode has no TLS peer.
+type AllowLoopbackIdentity struct{}
+
+func (AllowLoopbackIdentity) Authorize(context.Context, string) error {
+	return nil
+}
+
+type SPIFFEIdentityAuthorizer struct {
+	trustDomain string
+	environment string
+}
+
+func NewSPIFFEIdentityAuthorizer(trustDomain, environment string) (*SPIFFEIdentityAuthorizer, error) {
+	if strings.TrimSpace(trustDomain) == "" || strings.TrimSpace(environment) == "" {
+		return nil, fmt.Errorf("spiffe identity authorizer: trust domain and environment are required")
+	}
+	return &SPIFFEIdentityAuthorizer{trustDomain: trustDomain, environment: environment}, nil
+}
+
+func (authorizer *SPIFFEIdentityAuthorizer) Authorize(ctx context.Context, workerID string) error {
+	if authorizer == nil || workerID == "" {
+		return ErrUnauthorizedIdentity
+	}
+	peerInfo, ok := peer.FromContext(ctx)
+	if !ok || peerInfo.AuthInfo == nil {
+		return ErrUnauthorizedIdentity
+	}
+	tlsInfo, ok := peerInfo.AuthInfo.(credentials.TLSInfo)
+	if !ok || len(tlsInfo.State.VerifiedChains) == 0 {
+		return ErrUnauthorizedIdentity
+	}
+
+	expectedPath := "/" + authorizer.environment + "/worker/" + workerID
+	workerIdentityCount := 0
+	for _, chain := range tlsInfo.State.VerifiedChains {
+		if len(chain) == 0 || chain[0] == nil {
+			continue
+		}
+		for _, uri := range chain[0].URIs {
+			if uri == nil || uri.Scheme != "spiffe" || uri.Host != authorizer.trustDomain {
+				continue
+			}
+			segments := strings.Split(strings.Trim(uri.Path, "/"), "/")
+			if len(segments) != 3 || segments[1] != "worker" {
+				continue
+			}
+			workerIdentityCount++
+			if workerIdentityCount > 1 || uri.RawQuery != "" || uri.Fragment != "" || uri.Path != expectedPath {
+				return ErrUnauthorizedIdentity
+			}
+		}
+	}
+	if workerIdentityCount != 1 {
+		return ErrUnauthorizedIdentity
+	}
+	return nil
+}
--- a/internal/controller/worker/identity_test.go
+++ b/internal/controller/worker/identity_test.go
@ -0,0 +1,84 @@
+package worker
+
+import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"net/url"
+	"testing"
+
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/peer"
+)
+
+func TestSPIFFEIdentityAuthorizer(t *testing.T) {
+	authorizer, err := NewSPIFFEIdentityAuthorizer("proxy.example", "prod")
+	if err != nil {
+		t.Fatalf("NewSPIFFEIdentityAuthorizer(): %v", err)
+	}
+
+	tests := []struct {
+		name      string
+		workerID  string
+		uri       string
+		withTLS   bool
+		wantError bool
+	}{
+		{name: "matching worker uri", workerID: "worker-a", uri: "spiffe://proxy.example/prod/worker/worker-a", withTLS: true},
+		{name: "different worker", workerID: "worker-a", uri: "spiffe://proxy.example/prod/worker/worker-b", withTLS: true, wantError: true},
+		{name: "different environment", workerID: "worker-a", uri: "spiffe://proxy.example/staging/worker/worker-a", withTLS: true, wantError: true},
+		{name: "different trust domain", workerID: "worker-a", uri: "spiffe://other.example/prod/worker/worker-a", withTLS: true, wantError: true},
+		{name: "missing peer tls", workerID: "worker-a", wantError: true},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			ctx := context.Background()
+			if test.withTLS {
+				ctx = tlsPeerContext(t, test.uri)
+			}
+			err := authorizer.Authorize(ctx, test.workerID)
+			if (err != nil) != test.wantError {
+				t.Fatalf("Authorize() error = %v, wantError %v", err, test.wantError)
+			}
+		})
+	}
+}
+
+func TestSPIFFEIdentityAuthorizerRejectsMultipleWorkerURIs(t *testing.T) {
+	authorizer, err := NewSPIFFEIdentityAuthorizer("proxy.example", "prod")
+	if err != nil {
+		t.Fatalf("NewSPIFFEIdentityAuthorizer(): %v", err)
+	}
+	first, err := url.Parse("spiffe://proxy.example/prod/worker/worker-a")
+	if err != nil {
+		t.Fatal(err)
+	}
+	second, err := url.Parse("spiffe://proxy.example/prod/worker/worker-b")
+	if err != nil {
+		t.Fatal(err)
+	}
+	ctx := peer.NewContext(context.Background(), &peer.Peer{AuthInfo: credentials.TLSInfo{
+		State: tls.ConnectionState{VerifiedChains: [][]*x509.Certificate{{{URIs: []*url.URL{first, second}}}}},
+	}})
+	if err := authorizer.Authorize(ctx, "worker-a"); err == nil {
+		t.Fatal("Authorize() error = nil, want rejection for multiple worker identities")
+	}
+}
+
+func TestAllowLoopbackIdentity(t *testing.T) {
+	if err := (AllowLoopbackIdentity{}).Authorize(context.Background(), "worker-a"); err != nil {
+		t.Fatalf("Authorize(): %v", err)
+	}
+}
+
+func tlsPeerContext(t *testing.T, identityURI string) context.Context {
+	t.Helper()
+	uri, err := url.Parse(identityURI)
+	if err != nil {
+		t.Fatalf("url.Parse(): %v", err)
+	}
+	return peer.NewContext(context.Background(), &peer.Peer{AuthInfo: credentials.TLSInfo{
+		State: tls.ConnectionState{VerifiedChains: [][]*x509.Certificate{{{URIs: []*url.URL{uri}}}}},
+	}})
+}
--- a/internal/controller/worker/server.go
+++ b/internal/controller/worker/server.go
@ -0,0 +1,169 @@
+package worker
+
+import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"errors"
+	"fmt"
+	"net"
+	"os"
+	"strings"
+	"time"
+
+	controlplanev1 "proxy-pool/gen/controlplane/v1"
+	"proxy-pool/internal/config"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/keepalive"
+)
+
+var ErrInvalidServer = errors.New("invalid worker control server configuration")
+
+type ServerOptions struct {
+	ShutdownTimeout time.Duration
+}
+
+func DefaultServerOptions() ServerOptions {
+	return ServerOptions{ShutdownTimeout: 15 * time.Second}
+}
+
+type Server struct {
+	listen          string
+	grpcServer      *grpc.Server
+	shutdownTimeout time.Duration
+}
+
+func NewServer(controlPlane config.ControlPlane, service Service, options ServerOptions) (*Server, error) {
+	if service == nil || !controlPlane.Enabled || !validServerConfig(controlPlane) {
+		return nil, ErrInvalidServer
+	}
+	if options.ShutdownTimeout < 0 {
+		return nil, fmt.Errorf("%w: shutdown timeout must not be negative", ErrInvalidServer)
+	}
+	if options.ShutdownTimeout == 0 {
+		options = DefaultServerOptions()
+	}
+
+	identity, serverOptions, err := serverTransportOptions(controlPlane)
+	if err != nil {
+		return nil, err
+	}
+	serverOptions = append(serverOptions,
+		grpc.MaxRecvMsgSize(controlPlane.MaxMessageBytes),
+		grpc.MaxSendMsgSize(controlPlane.MaxMessageBytes),
+		grpc.MaxConcurrentStreams(controlPlane.MaxConcurrentStreams),
+		grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{
+			MinTime:             10 * time.Second,
+			PermitWithoutStream: false,
+		}),
+	)
+	grpcServer := grpc.NewServer(serverOptions...)
+	controlplanev1.RegisterWorkerControlPlaneServer(grpcServer, NewGRPCHandler(service, identity))
+	return &Server{listen: controlPlane.Listen, grpcServer: grpcServer, shutdownTimeout: options.ShutdownTimeout}, nil
+}
+
+func (server *Server) Run(ctx context.Context) error {
+	if server == nil || server.grpcServer == nil || server.listen == "" {
+		return ErrInvalidServer
+	}
+	listener, err := net.Listen("tcp", server.listen)
+	if err != nil {
+		return fmt.Errorf("listen worker control plane: %w", err)
+	}
+	return server.Serve(ctx, listener)
+}
+
+func (server *Server) Serve(ctx context.Context, listener net.Listener) error {
+	if server == nil || server.grpcServer == nil || listener == nil || ctx == nil {
+		return ErrInvalidServer
+	}
+	completed := make(chan struct{})
+	go func() {
+		select {
+		case <-ctx.Done():
+			server.gracefulStop()
+		case <-completed:
+		}
+	}()
+
+	err := server.grpcServer.Serve(listener)
+	close(completed)
+	if ctx.Err() != nil || errors.Is(err, grpc.ErrServerStopped) {
+		return nil
+	}
+	return err
+}
+
+func (server *Server) gracefulStop() {
+	stopped := make(chan struct{})
+	go func() {
+		server.grpcServer.GracefulStop()
+		close(stopped)
+	}()
+	timer := time.NewTimer(server.shutdownTimeout)
+	defer timer.Stop()
+	select {
+	case <-stopped:
+	case <-timer.C:
+		server.grpcServer.Stop()
+		<-stopped
+	}
+}
+
+func serverTransportOptions(controlPlane config.ControlPlane) (IdentityAuthorizer, []grpc.ServerOption, error) {
+	switch controlPlane.TLS.Mode {
+	case "disabled":
+		if !loopbackListen(controlPlane.Listen) {
+			return nil, nil, fmt.Errorf("%w: plaintext listener must be loopback", ErrInvalidServer)
+		}
+		return AllowLoopbackIdentity{}, nil, nil
+	case "mtls":
+		identity, err := NewSPIFFEIdentityAuthorizer(controlPlane.TLS.TrustDomain, controlPlane.TLS.Environment)
+		if err != nil {
+			return nil, nil, fmt.Errorf("%w: %v", ErrInvalidServer, err)
+		}
+		certificate, err := tls.LoadX509KeyPair(controlPlane.TLS.CertFile, controlPlane.TLS.KeyFile)
+		if err != nil {
+			return nil, nil, fmt.Errorf("%w: load server certificate: %v", ErrInvalidServer, err)
+		}
+		caPEM, err := os.ReadFile(controlPlane.TLS.ClientCAFile)
+		if err != nil {
+			return nil, nil, fmt.Errorf("%w: read client ca: %v", ErrInvalidServer, err)
+		}
+		clientCAs := x509.NewCertPool()
+		if !clientCAs.AppendCertsFromPEM(caPEM) {
+			return nil, nil, fmt.Errorf("%w: parse client ca", ErrInvalidServer)
+		}
+		transport := credentials.NewTLS(&tls.Config{
+			MinVersion:   tls.VersionTLS13,
+			Certificates: []tls.Certificate{certificate},
+			ClientAuth:   tls.RequireAndVerifyClientCert,
+			ClientCAs:    clientCAs,
+		})
+		return identity, []grpc.ServerOption{grpc.Creds(transport)}, nil
+	default:
+		return nil, nil, fmt.Errorf("%w: unsupported tls mode", ErrInvalidServer)
+	}
+}
+
+func validServerConfig(controlPlane config.ControlPlane) bool {
+	return controlPlane.Listen != "" && controlPlane.ProtocolVersion == 1 &&
+		controlPlane.HeartbeatInterval.Value() > 0 && controlPlane.SessionTTL.Value() > 0 &&
+		controlPlane.MaxStaleAge.Value() > 0 && controlPlane.MaxMessageBytes > 0 &&
+		controlPlane.MaxRuntimeCounters > 0 && controlPlane.MaxConcurrentStreams > 0
+}
+
+func loopbackListen(listen string) bool {
+	host, _, err := net.SplitHostPort(listen)
+	if err != nil {
+		return false
+	}
+	host = strings.Trim(host, "[]")
+	if strings.EqualFold(host, "localhost") {
+		return true
+	}
+	ip := net.ParseIP(host)
+	return ip != nil && ip.IsLoopback()
+}
--- a/internal/controller/worker/server_test.go
+++ b/internal/controller/worker/server_test.go
@ -0,0 +1,86 @@
+package worker
+
+import (
+	"context"
+	"errors"
+	"net"
+	"testing"
+	"time"
+
+	controlplanev1 "proxy-pool/gen/controlplane/v1"
+	"proxy-pool/internal/config"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+)
+
+func TestNewServerRejectsInvalidOptions(t *testing.T) {
+	controlPlane := validServerControlPlane()
+	service := &grpcServiceStub{}
+	if _, err := NewServer(controlPlane, nil, DefaultServerOptions()); !errors.Is(err, ErrInvalidServer) {
+		t.Fatalf("NewServer(nil service) error = %v, want ErrInvalidServer", err)
+	}
+	if _, err := NewServer(controlPlane, service, ServerOptions{ShutdownTimeout: -time.Second}); !errors.Is(err, ErrInvalidServer) {
+		t.Fatalf("NewServer(negative shutdown timeout) error = %v, want ErrInvalidServer", err)
+	}
+	controlPlane.Enabled = false
+	if _, err := NewServer(controlPlane, service, DefaultServerOptions()); !errors.Is(err, ErrInvalidServer) {
+		t.Fatalf("NewServer(disabled) error = %v, want ErrInvalidServer", err)
+	}
+}
+
+func TestServerServesAndStopsOnContextCancellation(t *testing.T) {
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("net.Listen(): %v", err)
+	}
+	controlPlane := validServerControlPlane()
+	controlPlane.Listen = listener.Addr().String()
+	service := &grpcServiceStub{registration: Registration{WorkerID: "worker-a", SessionID: "session-a", OwnershipEpoch: 5, HeartbeatInterval: time.Second, MaxStaleAge: 2 * time.Second}}
+	server, err := NewServer(controlPlane, service, ServerOptions{ShutdownTimeout: time.Second})
+	if err != nil {
+		t.Fatalf("NewServer(): %v", err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	result := make(chan error, 1)
+	go func() { result <- server.Serve(ctx, listener) }()
+
+	connection, err := grpc.NewClient(listener.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		t.Fatalf("grpc.NewClient(): %v", err)
+	}
+	client := controlplanev1.NewWorkerControlPlaneClient(connection)
+	requestCtx, requestCancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer requestCancel()
+	response, err := client.RegisterWorker(requestCtx, &controlplanev1.RegisterWorkerRequest{WorkerId: "worker-a", SupportedProtocolVersion: 1})
+	if err != nil || response.GetSessionId() != "session-a" {
+		t.Fatalf("RegisterWorker() = %+v, %v", response, err)
+	}
+	if err := connection.Close(); err != nil {
+		t.Fatalf("connection.Close(): %v", err)
+	}
+	cancel()
+	select {
+	case err := <-result:
+		if err != nil {
+			t.Fatalf("Serve() error = %v", err)
+		}
+	case <-time.After(3 * time.Second):
+		t.Fatal("Serve() did not stop after context cancellation")
+	}
+}
+
+func validServerControlPlane() config.ControlPlane {
+	return config.ControlPlane{
+		Enabled:              true,
+		Listen:               "127.0.0.1:8443",
+		ProtocolVersion:      1,
+		HeartbeatInterval:    config.Duration(10 * time.Second),
+		SessionTTL:           config.Duration(30 * time.Second),
+		MaxStaleAge:          config.Duration(10 * time.Second),
+		MaxMessageBytes:      1 << 20,
+		MaxRuntimeCounters:   100,
+		MaxConcurrentStreams: 10,
+		TLS:                  config.ControlPlaneTLS{Mode: "disabled"},
+	}
+}