emissary/internal/agent/controller/mdns/controller.go

package mdns

import (
	"context"
	"net"
	"sync"

	"forge.cadoles.com/Cadoles/emissary/internal/agent"
	mdns "forge.cadoles.com/Cadoles/emissary/internal/agent/controller/mdns/spec"
	"github.com/brutella/dnssd"
	"github.com/mitchellh/hashstructure/v2"
	"github.com/pkg/errors"
	"gitlab.com/wpetit/goweb/logger"
)

const (
	DefaultDomain = "local"
)

type Controller struct {
	serviceDefHash uint64
	cancel         context.CancelFunc
	responder      dnssd.Responder
	mutex          sync.RWMutex
}

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

// Reconcile implements node.Controller.
func (c *Controller) Reconcile(ctx context.Context, state *agent.State) error {
	mdnsSpec := mdns.NewSpec()

	if err := state.GetSpec(mdns.Name, mdnsSpec); err != nil {
		if errors.Is(err, agent.ErrSpecNotFound) {
			logger.Info(ctx, "could not find mdns spec")

			c.stopResponder(ctx)

			return nil
		}

		return errors.WithStack(err)
	}

	logger.Info(ctx, "retrieved spec", logger.F("spec", mdnsSpec.SpecName()), logger.F("revision", mdnsSpec.SpecRevision()))

	if err := c.updateResponder(ctx, mdnsSpec); err != nil {
		return errors.Wrap(err, "could not update responder")
	}

	return nil
}

func (c *Controller) stopResponder(ctx context.Context) {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	if c.responder == nil {
		return
	}

	c.cancel()
	c.responder = nil
	c.cancel = nil
}

func (c *Controller) updateResponder(ctx context.Context, spec *mdns.Spec) error {
	serviceDef := struct {
		Services map[string]mdns.Service
	}{
		Services: spec.Services,
	}

	newServerDefHash, err := hashstructure.Hash(serviceDef, hashstructure.FormatV2, nil)
	if err != nil {
		return errors.WithStack(err)
	}

	c.mutex.RLock()
	if newServerDefHash == c.serviceDefHash && c.responder != nil {
		c.mutex.RUnlock()
		return nil
	}
	c.mutex.RUnlock()

	c.stopResponder(ctx)

	defaultIfaces, err := c.getDefaultIfaces()
	if err != nil {
		return errors.WithStack(err)
	}

	services := make([]dnssd.Service, 0, len(spec.Services))

	for name, service := range spec.Services {
		domain := service.Domain
		if domain == "" {
			domain = DefaultDomain
		}

		ifaces := service.Ifaces
		if len(ifaces) == 0 {
			ifaces = defaultIfaces
		}

		config := dnssd.Config{
			Name:   name,
			Type:   service.Type,
			Domain: domain,
			Host:   service.Host,
			Ifaces: ifaces,
			Port:   service.Port,
		}

		service, err := dnssd.NewService(config)
		if err != nil {
			logger.Error(ctx, "could not create mdns service", logger.E(errors.WithStack(err)))

			continue
		}

		services = append(services, service)
	}

	responder, err := dnssd.NewResponder()
	if err != nil {
		return errors.WithStack(err)
	}

	for _, service := range services {
		if _, err := responder.Add(service); err != nil {
			logger.Error(ctx, "could not add mdns service", logger.E(errors.WithStack(err)))

			continue
		}
	}

	ctx, cancel := context.WithCancel(context.Background())

	c.responder = responder
	c.cancel = cancel
	c.serviceDefHash = newServerDefHash

	go func() {
		defer c.stopResponder(ctx)

		if err := responder.Respond(ctx); err != nil && !errors.Is(err, context.Canceled) {
			logger.Error(ctx, "could not respond to mdns queries", logger.E(errors.WithStack(err)))
		}
	}()

	return nil
}

func (c *Controller) getDefaultIfaces() ([]string, error) {
	ifaces, err := net.Interfaces()
	if err != nil {
		return nil, errors.WithStack(err)
	}

	ifaceNames := make([]string, len(ifaces))

	for idx, ifa := range ifaces {
		ifaceNames[idx] = ifa.Name
	}

	return ifaceNames, nil
}

func NewController() *Controller {
	return &Controller{
		cancel:         nil,
		responder:      nil,
		serviceDefHash: 0,
	}
}

var _ agent.Controller = &Controller{}