DDI/server/internal/agentapi/rpcs.go

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

package agentapi

import (
	"context"
	"errors"
	"io"
	"log/slog"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/durationpb"

	"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/service"
)

// maxHealthChecks bounds the checks stored per report.
const maxHealthChecks = 100

func roleName(r agentv1.Role) string {
	switch r {
	case agentv1.Role_ROLE_DNS:
		return service.RoleDNS
	case agentv1.Role_ROLE_DHCP:
		return service.RoleDHCP
	}
	return ""
}

func modeProto(m string) agentv1.OperatingMode {
	if m == agenthub.ModeManual {
		return agentv1.OperatingMode_OPERATING_MODE_MANUAL
	}
	return agentv1.OperatingMode_OPERATING_MODE_AUTO_PULL
}

func healthStateName(h agentv1.HealthState) string {
	switch h {
	case agentv1.HealthState_HEALTH_STATE_OK:
		return service.HealthOK
	case agentv1.HealthState_HEALTH_STATE_DEGRADED:
		return service.HealthDegraded
	case agentv1.HealthState_HEALTH_STATE_FAILED:
		return service.HealthFailed
	}
	return service.HealthUnknown
}

func healthFromProto(h *agentv1.HealthStatus) *service.AgentHealth {
	if h == nil {
		return nil
	}
	out := &service.AgentHealth{State: healthStateName(h.GetState()), Message: h.GetMessage(), Checks: []service.HealthCheck{}}
	if h.GetCheckedAt() != nil {
		t := h.GetCheckedAt().AsTime()
		out.CheckedAt = &t
	}
	for i, c := range h.GetChecks() {
		if i == maxHealthChecks {
			break
		}
		out.Checks = append(out.Checks, service.HealthCheck{Name: c.GetName(), State: healthStateName(c.GetState()), Message: c.GetMessage()})
	}
	return out
}

// Enroll implements agentv1.AgentServiceServer.
func (s *Server) Enroll(ctx context.Context, req *agentv1.EnrollRequest) (*agentv1.EnrollResponse, error) {
	c := callerFrom(ctx)
	role := roleName(req.GetRole())
	if role == "" {
		return nil, status.Error(codes.InvalidArgument, "unknown role")
	}
	res, err := s.svc.EnrollAgent(ctx, service.EnrollInput{
		Token: req.GetToken(), CSR: req.GetCsrDer(), Role: role, Hostname: req.GetHostname(),
		AgentVersion: req.GetAgentVersion(), RemoteIP: c.RemoteIP, RequestID: c.RequestID,
	})
	if err != nil {
		return nil, err
	}
	c.AgentID = res.AgentID
	s.log.Info("agent enrolled", slog.String("request_id", c.RequestID), slog.String("agent_id", res.AgentID),
		slog.String("role", role), slog.Any("remote", c.RemoteIP))
	return &agentv1.EnrollResponse{
		AgentId:                res.AgentID,
		CertificateDer:         res.Certificate.Raw,
		CaCertificatesDer:      [][]byte{s.pki.CACertificate().Raw},
		BundleSigningPublicKey: s.pki.BundlePublicKey(),
	}, nil
}

// RenewCertificate implements agentv1.AgentServiceServer.
func (s *Server) RenewCertificate(ctx context.Context, req *agentv1.RenewCertificateRequest) (*agentv1.RenewCertificateResponse, error) {
	c := callerFrom(ctx)
	cert, err := s.svc.RenewAgentCertificate(ctx, c.AgentID, req.GetCsrDer())
	if err != nil {
		return nil, err
	}
	s.log.Info("agent certificate renewed", slog.String("request_id", c.RequestID), slog.String("agent_id", c.AgentID),
		slog.Time("not_after", cert.NotAfter))
	return &agentv1.RenewCertificateResponse{
		CertificateDer: cert.Raw, CaCertificatesDer: [][]byte{s.pki.CACertificate().Raw},
	}, nil
}

// GetDesiredState implements agentv1.AgentServiceServer.
func (s *Server) GetDesiredState(ctx context.Context, _ *agentv1.GetDesiredStateRequest) (*agentv1.GetDesiredStateResponse, error) {
	c := callerFrom(ctx)
	st, err := s.svc.AgentDesiredState(ctx, c.AgentID, c.RemoteIP)
	if err != nil {
		return nil, err
	}
	return &agentv1.GetDesiredStateResponse{
		DesiredRevision: st.DesiredRevision, Mode: modeProto(st.Mode), PullInterval: durationpb.New(service.PullInterval),
		DeploymentId: st.PushDeploymentID,
	}, nil
}

// FetchBundle implements agentv1.AgentServiceServer. The bundle is signed
// for the calling agent and its role; equal input gives equal bytes.
func (s *Server) FetchBundle(ctx context.Context, req *agentv1.FetchBundleRequest) (*agentv1.FetchBundleResponse, error) {
	c := callerFrom(ctx)
	createdAt, arts, err := s.svc.RevisionArtifacts(ctx, req.GetRevision(), c.Role)
	if errors.Is(err, service.ErrNotFound) {
		return nil, status.Errorf(codes.NotFound, "revision %d does not exist", req.GetRevision())
	}
	if err != nil {
		return nil, err
	}
	manifest, sig, err := bundle.Sign(s.pki.BundleSigningKey(), c.AgentID, c.Role, req.GetRevision(), createdAt, arts)
	if err != nil {
		return nil, err
	}
	out := &agentv1.FetchBundleResponse{Manifest: manifest, Signature: sig, Artifacts: make([]*agentv1.Artifact, len(arts))}
	for i, a := range arts {
		out.Artifacts[i] = &agentv1.Artifact{Path: a.Path, Content: a.Content}
	}
	return out, nil
}

var (
	triggerNames = map[agentv1.DeploymentTrigger]string{
		agentv1.DeploymentTrigger_DEPLOYMENT_TRIGGER_PULL: "pull",
		agentv1.DeploymentTrigger_DEPLOYMENT_TRIGGER_PUSH: "push",
	}
	outcomeNames = map[agentv1.DeploymentOutcome]string{
		agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_SUCCEEDED:   "succeeded",
		agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_REJECTED:    "rejected",
		agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_ROLLED_BACK: "rolled_back",
		agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_FAILED:      "failed",
	}
	leaseStateNames = map[agentv1.LeaseState]string{
		agentv1.LeaseState_LEASE_STATE_ACTIVE:   "active",
		agentv1.LeaseState_LEASE_STATE_EXPIRED:  "expired",
		agentv1.LeaseState_LEASE_STATE_RELEASED: "released",
		agentv1.LeaseState_LEASE_STATE_DECLINED: "declined",
	}
)

// ReportDeployment implements agentv1.AgentServiceServer.
func (s *Server) ReportDeployment(ctx context.Context, req *agentv1.ReportDeploymentRequest) (*agentv1.ReportDeploymentResponse, error) {
	c := callerFrom(ctx)
	trigger, ok := triggerNames[req.GetTrigger()]
	if !ok {
		return nil, status.Error(codes.InvalidArgument, "unknown trigger")
	}
	outcome, ok := outcomeNames[req.GetOutcome()]
	if !ok {
		return nil, status.Error(codes.InvalidArgument, "unknown outcome")
	}
	rep := service.DeploymentReport{
		CorrelationID: req.GetDeploymentId(), Revision: req.GetRevision(), Trigger: trigger, Outcome: outcome,
		Message: req.GetMessage(), ActiveRevision: req.GetActiveRevision(),
	}
	if req.GetStartedAt() != nil {
		rep.StartedAt = req.GetStartedAt().AsTime()
	}
	if req.GetFinishedAt() != nil {
		rep.FinishedAt = req.GetFinishedAt().AsTime()
	}
	if err := s.svc.RecordDeployment(ctx, c.AgentID, rep); err != nil {
		return nil, err
	}
	s.log.Info("deployment reported", slog.String("request_id", c.RequestID), slog.String("agent_id", c.AgentID),
		slog.Uint64("revision", req.GetRevision()), slog.String("outcome", outcome),
		slog.Uint64("active_revision", req.GetActiveRevision()), slog.String("deployment_id", req.GetDeploymentId()))
	return &agentv1.ReportDeploymentResponse{}, nil
}

// ReportLeases implements agentv1.AgentServiceServer.
func (s *Server) ReportLeases(ctx context.Context, req *agentv1.ReportLeasesRequest) (*agentv1.ReportLeasesResponse, error) {
	c := callerFrom(ctx)
	in := make([]service.LeaseInput, len(req.GetLeases()))
	for i, l := range req.GetLeases() {
		in[i] = service.LeaseInput{IP: l.GetIp(), MAC: l.GetMac(), Hostname: l.GetHostname(),
			State: leaseStateNames[l.GetState()], SubnetID: l.GetSubnetId(), ClientID: l.GetClientId()}
		if l.GetExpiresAt() != nil {
			t := l.GetExpiresAt().AsTime()
			in[i].ExpiresAt = &t
		}
		if l.GetUpdatedAt() != nil {
			t := l.GetUpdatedAt().AsTime()
			in[i].UpdatedAt = &t
		}
	}
	res, err := s.svc.StoreLeaseBatch(ctx, c.AgentID, req.GetSequence(), req.GetFullSync(), in)
	if err != nil {
		return nil, err
	}
	attrs := []any{slog.String("request_id", c.RequestID), slog.String("agent_id", c.AgentID),
		slog.Uint64("sequence", req.GetSequence()), slog.Bool("full_sync", req.GetFullSync()), slog.Int("leases", len(in))}
	switch {
	case res.Duplicate:
		s.log.Info("dropped duplicate lease batch", append(attrs, slog.Uint64("acknowledged", res.Acknowledged))...)
	case res.Skipped > 0:
		s.log.Warn("skipped invalid leases", append(attrs, slog.Int("skipped", res.Skipped))...)
	default:
		s.log.Debug("stored lease batch", attrs...)
	}
	return &agentv1.ReportLeasesResponse{AcknowledgedSequence: res.Acknowledged}, nil
}

// Connect implements agentv1.AgentServiceServer: Hello → Welcome, then
// heartbeats and acknowledgements from the agent and commands from the
// server until either side ends the stream.
func (s *Server) Connect(stream agentv1.AgentService_ConnectServer) error {
	ctx := stream.Context()
	c := callerFrom(ctx)
	first, err := stream.Recv()
	if err != nil {
		return err
	}
	hello := first.GetHello()
	if hello == nil {
		return status.Error(codes.InvalidArgument, "the first message must be Hello")
	}
	// Register before reading the desired revision, so that a revision
	// created in between is pushed rather than missed.
	conn := s.hub.Attach(c.AgentID, c.Role, c.Mode)
	defer func() {
		s.hub.Detach(conn)
		s.svc.AgentDisconnected(c.AgentID)
	}()
	settings, err := s.svc.AgentHello(ctx, c.AgentID, service.HelloInput{
		Hostname: hello.GetHostname(), AgentVersion: hello.GetAgentVersion(), ServiceVersion: hello.GetServiceVersion(),
		AppliedRevision: hello.GetAppliedRevision(), LastGoodRevision: hello.GetLastGoodRevision(),
		Health: healthFromProto(hello.GetHealth()), RemoteIP: c.RemoteIP,
	})
	if err != nil {
		return err
	}
	// The mode read with Hello is the one Welcome announces; it may be
	// newer than the one from the authentication (a concurrent mode change).
	s.hub.SetMode(c.AgentID, settings.Mode)
	if err := stream.Send(&agentv1.ConnectResponse{Message: &agentv1.ConnectResponse_Welcome{Welcome: &agentv1.Welcome{
		Mode: modeProto(settings.Mode), PullInterval: durationpb.New(service.PullInterval),
		HeartbeatInterval: durationpb.New(service.HeartbeatInterval), DesiredRevision: settings.DesiredRevision,
	}}}); err != nil {
		return err
	}
	s.log.Info("agent connected", slog.String("request_id", c.RequestID), slog.String("agent_id", c.AgentID),
		slog.String("role", c.Role), slog.Uint64("applied_revision", hello.GetAppliedRevision()),
		slog.Uint64("desired_revision", settings.DesiredRevision))

	// An open push is sent again: the agent may have missed it while it was
	// offline (ADR 0022). Only the latest one is still pending.
	if settings.PushDeploymentID != "" {
		cmd := agenthub.ApplyCommand(settings.DesiredRevision, settings.PushDeploymentID)
		if err := stream.Send(&agentv1.ConnectResponse{Message: &agentv1.ConnectResponse_Command{Command: cmd}}); err != nil {
			return err
		}
		s.log.Info("pending push sent", slog.String("request_id", c.RequestID), slog.String("agent_id", c.AgentID),
			slog.String("command_id", cmd.GetCommandId()), slog.Uint64("revision", settings.DesiredRevision),
			slog.String("deployment_id", settings.PushDeploymentID))
	}

	recvErr := make(chan error, 1)
	go func() {
		for {
			msg, err := stream.Recv()
			if err != nil {
				recvErr <- err
				return
			}
			s.handleMessage(ctx, c, msg)
		}
	}()
	for {
		select {
		case cmd := <-conn.Commands():
			if err := stream.Send(&agentv1.ConnectResponse{Message: &agentv1.ConnectResponse_Command{Command: cmd}}); err != nil {
				return err
			}
			s.log.Info("command sent", slog.String("request_id", c.RequestID), slog.String("agent_id", c.AgentID),
				slog.String("command_id", cmd.GetCommandId()), slog.Uint64("revision", cmd.GetApplyRevision().GetRevision()),
				slog.String("deployment_id", cmd.GetApplyRevision().GetDeploymentId()))
		case <-conn.Done():
			reason := conn.Reason()
			s.log.Info("agent disconnected by server", slog.String("request_id", c.RequestID),
				slog.String("agent_id", c.AgentID), slog.String("reason", reason))
			switch reason {
			case agenthub.ReasonRevoked, agenthub.ReasonDeleted:
				return status.Error(codes.PermissionDenied, reason)
			case agenthub.ReasonReplaced:
				return status.Error(codes.Aborted, reason)
			default:
				return status.Error(codes.Unavailable, reason)
			}
		case err := <-recvErr:
			if errors.Is(err, io.EOF) {
				s.log.Info("agent closed its stream", slog.String("request_id", c.RequestID), slog.String("agent_id", c.AgentID))
				return nil
			}
			return err
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (s *Server) handleMessage(ctx context.Context, c *caller, msg *agentv1.ConnectRequest) {
	switch m := msg.GetMessage().(type) {
	case *agentv1.ConnectRequest_Heartbeat:
		hb := m.Heartbeat
		err := s.svc.AgentHeartbeat(ctx, c.AgentID, service.HeartbeatInput{
			AppliedRevision: hb.GetAppliedRevision(), LastGoodRevision: hb.GetLastGoodRevision(),
			Health: healthFromProto(hb.GetHealth()),
		})
		if errors.Is(err, service.ErrAgentRevoked) {
			s.hub.Disconnect(c.AgentID, agenthub.ReasonRevoked)
		} else if err != nil && ctx.Err() == nil {
			s.log.Warn("heartbeat not stored", slog.String("request_id", c.RequestID), slog.String("agent_id", c.AgentID),
				slog.Any("error", err))
		}
	case *agentv1.ConnectRequest_CommandAck:
		ack := m.CommandAck
		s.log.Info("command acknowledged", slog.String("request_id", c.RequestID), slog.String("agent_id", c.AgentID),
			slog.String("command_id", ack.GetCommandId()), slog.Bool("accepted", ack.GetAccepted()),
			slog.String("message", ack.GetMessage()))
	case *agentv1.ConnectRequest_Hello:
		s.log.Warn("ignoring repeated Hello", slog.String("request_id", c.RequestID), slog.String("agent_id", c.AgentID))
	default:
		s.log.Warn("ignoring unknown stream message", slog.String("request_id", c.RequestID), slog.String("agent_id", c.AgentID))
	}
}