emissary/internal/agent/controller/spec/controller.go

package spec

import (
	"context"

	"forge.cadoles.com/Cadoles/emissary/internal/agent"
	"forge.cadoles.com/Cadoles/emissary/internal/agent/machineid"
	"forge.cadoles.com/Cadoles/emissary/internal/client"
	"forge.cadoles.com/Cadoles/emissary/internal/datastore"
	"forge.cadoles.com/Cadoles/emissary/internal/server"
	"github.com/pkg/errors"
	"gitlab.com/wpetit/goweb/api"
	"gitlab.com/wpetit/goweb/logger"
)

type Controller struct {
	client *client.Client
}

// Name implements node.Controller.
func (c *Controller) Name() string {
	return "spec-controller"
}

// Reconcile implements node.Controller.
func (c *Controller) Reconcile(ctx context.Context, state *agent.State) error {
	machineID, err := machineid.Get()
	if err != nil {
		return errors.WithStack(err)
	}

	ctx = logger.With(ctx, logger.F("machineID", machineID))

	agent, err := c.client.RegisterAgent(ctx, machineID)
	isAlreadyRegisteredErr, _ := isAPIError(err, server.ErrCodeAlreadyRegistered)

	switch {
	case isAlreadyRegisteredErr:
		agents, _, err := c.client.QueryAgents(
			ctx,
			client.WithQueryAgentsLimit(1),
			client.WithQueryAgentsRemoteID(machineID),
		)
		if err != nil {
			return errors.WithStack(err)
		}

		if len(agents) == 0 {
			logger.Error(ctx, "could not find remote matching agent")

			return nil
		}

		if err := c.reconcileAgent(ctx, state, agents[0]); err != nil {
			return errors.WithStack(err)
		}

		return nil

	case agent != nil:
		if err := c.reconcileAgent(ctx, state, agent); err != nil {
			return errors.WithStack(err)
		}

		return nil

	case err != nil:
		logger.Error(ctx, "could not contact server", logger.E(errors.WithStack(err)))

		return nil
	}

	return nil
}

func (c *Controller) reconcileAgent(ctx context.Context, state *agent.State, agent *datastore.Agent) error {
	ctx = logger.With(ctx, logger.F("agentID", agent.ID))

	if agent.Status != datastore.AgentStatusAccepted {
		logger.Error(ctx, "unexpected agent status", logger.F("status", agent.Status))

		return nil
	}

	specs, err := c.client.GetAgentSpecs(ctx, agent.ID)
	if err != nil {
		logger.Error(ctx, "could not retrieve agent specs", logger.E(errors.WithStack(err)))

		return nil
	}

	state.ClearSpecs()

	for _, spec := range specs {
		state.SetSpec(spec)
	}

	return nil
}

func NewController(serverURL string) *Controller {
	client := client.New(serverURL)

	return &Controller{client}
}

func isAPIError(err error, code api.ErrorCode) (bool, any) {
	apiError := &api.Error{}
	if errors.As(err, &apiError) && apiError.Code == code {
		return true, apiError.Data
	}

	return false, nil
}

var _ agent.Controller = &Controller{}