DDI/server/internal/agentapi/server.go

// SPDX-License-Identifier: LicenseRef-DDI-Proprietary

// Package agentapi is the gRPC AgentService of the management server
// (proto/ddi/agent/v1, ADR 0007, 0008, 0021). It listens on its own TLS 1.3
// port, requests client certificates and checks them itself: Enroll needs
// none, RenewCertificate accepts a recently expired one, every other call
// needs the agent's current, unrevoked certificate.
package agentapi

import (
	"context"
	"crypto/rand"
	"crypto/tls"
	"encoding/hex"
	"errors"
	"log/slog"
	"net"
	"net/netip"
	"regexp"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/peer"
	"google.golang.org/grpc/status"

	"git.beugsiud.de/beug.k/DDI/internal/bundle"
	agentv1 "git.beugsiud.de/beug.k/DDI/proto/gen/ddi/agent/v1"
	"git.beugsiud.de/beug.k/DDI/server/internal/agenthub"
	"git.beugsiud.de/beug.k/DDI/server/internal/pki"
	"git.beugsiud.de/beug.k/DDI/server/internal/service"
)

// MaxRecvMsgSize bounds requests (large full lease syncs).
const MaxRecvMsgSize = 32 << 20

// MaxSendMsgSize allows the largest bundle (bundle.MaxTotalSize) plus the
// manifest and framing.
const MaxSendMsgSize = bundle.MaxTotalSize + 16<<20

// Options configures the server.
type Options struct {
	Service *service.Service
	PKI     *pki.Authority
	Hub     *agenthub.Hub
	// ServerNames become the SANs of the server certificate.
	ServerNames []string
	Logger      *slog.Logger
}

// Server implements agentv1.AgentServiceServer.
type Server struct {
	agentv1.UnimplementedAgentServiceServer

	svc  *service.Service
	pki  *pki.Authority
	hub  *agenthub.Hub
	log  *slog.Logger
	grpc *grpc.Server
}

// New creates the server and issues its certificate.
func New(opts Options) (*Server, error) {
	if opts.Service == nil || opts.PKI == nil || opts.Hub == nil {
		return nil, errors.New("agentapi: service, PKI and hub are required")
	}
	if opts.Logger == nil {
		opts.Logger = slog.New(slog.DiscardHandler)
	}
	cert, err := opts.PKI.ServerCertificate(opts.ServerNames)
	if err != nil {
		return nil, err
	}
	s := &Server{svc: opts.Service, pki: opts.PKI, hub: opts.Hub, log: opts.Logger}
	tlsCfg := &tls.Config{
		MinVersion:   tls.VersionTLS13,
		Certificates: []tls.Certificate{cert},
		// Checked by the interceptors (Enroll has no certificate,
		// RenewCertificate may present an expired one).
		ClientAuth: tls.RequestClientCert,
	}
	s.grpc = grpc.NewServer(
		grpc.Creds(credentials.NewTLS(tlsCfg)),
		grpc.ChainUnaryInterceptor(s.unaryInterceptor),
		grpc.ChainStreamInterceptor(s.streamInterceptor),
		// Agents ping every 30 s, also without an open stream.
		grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{MinTime: 20 * time.Second, PermitWithoutStream: true}),
		grpc.KeepaliveParams(keepalive.ServerParameters{Time: time.Minute, Timeout: 20 * time.Second}),
		grpc.MaxRecvMsgSize(MaxRecvMsgSize),
		grpc.MaxSendMsgSize(MaxSendMsgSize),
	)
	agentv1.RegisterAgentServiceServer(s.grpc, s)
	return s, nil
}

// Serve accepts connections on lis until Stop or Shutdown.
func (s *Server) Serve(lis net.Listener) error {
	err := s.grpc.Serve(lis)
	if errors.Is(err, grpc.ErrServerStopped) {
		return nil
	}
	return err
}

// Shutdown closes the agents' streams, waits for running calls and stops
// hard when ctx ends.
func (s *Server) Shutdown(ctx context.Context) {
	s.hub.CloseAll(agenthub.ReasonShutdown)
	done := make(chan struct{})
	go func() {
		s.grpc.GracefulStop()
		close(done)
	}()
	select {
	case <-done:
	case <-ctx.Done():
		s.grpc.Stop()
		<-done
	}
}

// caller is the authenticated agent of a call.
type caller struct {
	AgentID   string
	Role      string
	Mode      string
	RequestID string
	RemoteIP  *netip.Addr
}

type callerKey struct{}

func callerFrom(ctx context.Context) *caller {
	c, _ := ctx.Value(callerKey{}).(*caller)
	if c == nil {
		return &caller{}
	}
	return c
}

var requestIDPattern = regexp.MustCompile(`^[A-Za-z0-9._-]{1,64}$`)

// requestID takes the agent's x-request-id if it is well-formed, so that
// agent and server logs correlate; otherwise it makes one up.
func requestID(ctx context.Context) string {
	if md, ok := metadata.FromIncomingContext(ctx); ok {
		if v := md.Get("x-request-id"); len(v) == 1 && requestIDPattern.MatchString(v[0]) {
			return v[0]
		}
	}
	var b [8]byte
	_, _ = rand.Read(b[:])
	return hex.EncodeToString(b[:])
}

func peerIP(ctx context.Context) *netip.Addr {
	p, ok := peer.FromContext(ctx)
	if !ok || p.Addr == nil {
		return nil
	}
	ap, err := netip.ParseAddrPort(p.Addr.String())
	if err != nil {
		return nil
	}
	ip := ap.Addr().Unmap()
	return &ip
}

// authenticate checks the client certificate for method.
func (s *Server) authenticate(ctx context.Context, method string) (*caller, error) {
	c := &caller{RequestID: requestID(ctx), RemoteIP: peerIP(ctx)}
	if method == agentv1.AgentService_Enroll_FullMethodName {
		return c, nil
	}
	p, ok := peer.FromContext(ctx)
	if !ok {
		return c, status.Error(codes.Unauthenticated, "client certificate required")
	}
	info, ok := p.AuthInfo.(credentials.TLSInfo)
	if !ok || len(info.State.PeerCertificates) == 0 {
		return c, status.Error(codes.Unauthenticated, "client certificate required")
	}
	cert := info.State.PeerCertificates[0]
	mode := pki.VerifyCurrent
	if method == agentv1.AgentService_RenewCertificate_FullMethodName {
		mode = pki.VerifyRenewal
	}
	id, err := s.pki.VerifyAgentCertificate(cert, mode)
	if err != nil {
		s.log.Warn("agent certificate rejected", slog.String("request_id", c.RequestID), slog.String("method", method),
			slog.Any("remote", c.RemoteIP), slog.Any("error", err))
		return c, status.Error(codes.Unauthenticated, "invalid client certificate")
	}
	c.AgentID = id
	ident, err := s.svc.AgentIdentity(ctx, id)
	if errors.Is(err, service.ErrNotFound) {
		return c, status.Error(codes.Unauthenticated, "unknown agent")
	}
	if err != nil {
		return c, s.toStatus(c, err)
	}
	if ident.Revoked {
		return c, status.Error(codes.PermissionDenied, "agent is revoked")
	}
	if pki.SerialHex(cert.SerialNumber) != ident.CertSerial {
		return c, status.Error(codes.Unauthenticated, "certificate has been replaced")
	}
	c.Role, c.Mode = ident.Role, ident.Mode
	return c, nil
}

func (s *Server) logCall(c *caller, method string, start time.Time, err error) {
	code := status.Code(err)
	attrs := []any{slog.String("request_id", c.RequestID), slog.String("method", method),
		slog.String("code", code.String()), slog.Duration("duration", time.Since(start))}
	if c.AgentID != "" {
		attrs = append(attrs, slog.String("agent_id", c.AgentID))
	}
	switch code {
	case codes.OK, codes.Canceled:
		s.log.Debug("agent call", attrs...)
	case codes.Internal, codes.Unknown:
		s.log.Error("agent call failed", append(attrs, slog.Any("error", err))...)
	default:
		s.log.Info("agent call rejected", append(attrs, slog.String("reason", status.Convert(err).Message()))...)
	}
}

func (s *Server) unaryInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
	start := time.Now()
	c, err := s.authenticate(ctx, info.FullMethod)
	if err == nil {
		var resp any
		resp, err = handler(context.WithValue(ctx, callerKey{}, c), req)
		if err != nil {
			err = s.toStatus(c, err)
		}
		s.logCall(c, info.FullMethod, start, err)
		return resp, err
	}
	s.logCall(c, info.FullMethod, start, err)
	return nil, err
}

type wrappedStream struct {
	grpc.ServerStream
	ctx context.Context
}

func (w *wrappedStream) Context() context.Context { return w.ctx }

func (s *Server) streamInterceptor(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
	start := time.Now()
	c, err := s.authenticate(ss.Context(), info.FullMethod)
	if err == nil {
		err = handler(srv, &wrappedStream{ServerStream: ss, ctx: context.WithValue(ss.Context(), callerKey{}, c)})
		if err != nil {
			err = s.toStatus(c, err)
		}
	}
	s.logCall(c, info.FullMethod, start, err)
	return err
}

var kindCodes = map[service.Kind]codes.Code{
	service.KindInvalid:      codes.InvalidArgument,
	service.KindNotFound:     codes.NotFound,
	service.KindConflict:     codes.FailedPrecondition,
	service.KindForbidden:    codes.PermissionDenied,
	service.KindUnauthorized: codes.Unauthenticated,
	service.KindRateLimited:  codes.ResourceExhausted,
	service.KindUnavailable:  codes.Unavailable,
}

// toStatus converts service errors; internal details are logged, not sent.
func (s *Server) toStatus(c *caller, err error) error {
	if err == nil {
		return nil
	}
	if _, ok := status.FromError(err); ok {
		return err
	}
	var se *service.Error
	if errors.As(err, &se) && se.Kind != 0 {
		return status.Error(kindCodes[se.Kind], se.Message)
	}
	switch {
	case errors.Is(err, context.Canceled):
		return status.Error(codes.Canceled, "canceled")
	case errors.Is(err, context.DeadlineExceeded):
		return status.Error(codes.DeadlineExceeded, "deadline exceeded")
	}
	s.log.Error("agent call failed", slog.String("request_id", c.RequestID), slog.String("agent_id", c.AgentID), slog.Any("error", err))
	return status.Error(codes.Internal, "internal error")
}