emissary/internal/agent/controller/registration/controller.go

package registration

import (
	"context"
	"net/http"
	"sync/atomic"

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

type Status struct {
	Agent      *datastore.Agent
	Connected  bool
	Claimed    bool
	Thumbprint string
	ServerURL  string
}

type Controller struct {
	status *atomic.Value
	server *atomic.Value
	addr   string
}

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

// Reconcile implements node.Controller.
func (c *Controller) Reconcile(ctx context.Context, state *agent.State) error {
	cl := agent.Client(ctx)
	thumbprint := agent.Thumbprint(ctx)

	connected := true
	agent, err := cl.GetAgent(
		ctx,
		state.AgentID(),
	)
	if err != nil {
		logger.Error(ctx, "could not get agent", logger.E(errors.WithStack(err)))
		connected = false
	}

	claimed := agent != nil && agent.TenantID != nil

	c.status.Store(Status{
		Agent:      agent,
		Connected:  connected,
		Claimed:    claimed,
		Thumbprint: thumbprint,
		ServerURL:  cl.ServerURL(),
	})

	if err := c.reconcileAgent(ctx, connected, claimed); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

func (c *Controller) reconcileAgent(ctx context.Context, connected bool, claimed bool) error {
	shouldStart := !connected || !claimed

	if shouldStart {
		if err := c.startServer(ctx); err != nil {
			return errors.WithStack(err)
		}
	} else {
		if err := c.stopServer(ctx); err != nil {
			return errors.WithStack(err)
		}
	}

	return nil
}

func (c *Controller) startServer(ctx context.Context) error {
	server := c.getServer()
	if server != nil {
		return nil
	}

	server = &http.Server{
		Addr: c.addr,
		Handler: &Handler{
			status: c.status,
		},
	}

	go func() {
		defer c.setServer(nil)

		if err := server.ListenAndServe(); err != nil {
			logger.Error(ctx, "could not start server", logger.E(errors.WithStack(err)))
		}
	}()

	c.setServer(server)

	return nil
}

func (c *Controller) stopServer(ctx context.Context) error {
	server := c.getServer()
	if server == nil {
		return nil
	}

	defer c.setServer(nil)

	if err := server.Close(); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

func (c *Controller) setServer(s *http.Server) {
	c.server.Store(s)
}

func (c *Controller) getServer() *http.Server {
	server, ok := c.server.Load().(*http.Server)
	if !ok {
		return nil
	}

	return server
}

func NewController(addr string) *Controller {
	return &Controller{
		addr:   addr,
		status: &atomic.Value{},
		server: &atomic.Value{},
	}
}

var _ agent.Controller = &Controller{}