diff --git a/internal/controller/worker/identity.go b/internal/controller/worker/identity.go new file mode 100644 index 0000000..4a213de --- /dev/null +++ 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 +} diff --git a/internal/controller/worker/identity_test.go b/internal/controller/worker/identity_test.go new file mode 100644 index 0000000..8a01ae1 --- /dev/null +++ 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}}}}}, + }}) +} diff --git a/internal/controller/worker/server.go b/internal/controller/worker/server.go new file mode 100644 index 0000000..08de9dd --- /dev/null +++ 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() +} diff --git a/internal/controller/worker/server_test.go b/internal/controller/worker/server_test.go new file mode 100644 index 0000000..fe55df1 --- /dev/null +++ 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"}, + } +}