emissary/internal/agent/state.go

package agent

import (
	"encoding/json"

	"forge.cadoles.com/Cadoles/emissary/internal/datastore"
	"forge.cadoles.com/Cadoles/emissary/internal/spec"
	"github.com/mitchellh/mapstructure"
	"github.com/pkg/errors"
)

var ErrSpecNotFound = errors.New("spec not found")

type Specs map[spec.Name]spec.Spec

type State struct {
	agentID datastore.AgentID
	specs   Specs
}

func NewState() *State {
	return &State{
		specs: make(map[spec.Name]spec.Spec),
	}
}

func (s *State) MarshalJSON() ([]byte, error) {
	state := struct {
		ID    datastore.AgentID           `json:"agentId"`
		Specs map[spec.Name]*spec.RawSpec `json:"specs"`
	}{
		ID: s.agentID,
		Specs: func(specs map[spec.Name]spec.Spec) map[spec.Name]*spec.RawSpec {
			rawSpecs := make(map[spec.Name]*spec.RawSpec)

			for name, sp := range specs {
				rawSpecs[name] = &spec.RawSpec{
					Name:     sp.SpecName(),
					Revision: sp.SpecRevision(),
					Data:     sp.SpecData(),
				}
			}

			return rawSpecs
		}(s.specs),
	}

	data, err := json.Marshal(state)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return data, nil
}

func (s *State) UnmarshalJSON(data []byte) error {
	state := struct {
		AgentID datastore.AgentID           `json:"agentId"`
		Specs   map[spec.Name]*spec.RawSpec `json:"specs"`
	}{}

	if err := json.Unmarshal(data, &state); err != nil {
		return errors.WithStack(err)
	}

	s.specs = func(rawSpecs map[spec.Name]*spec.RawSpec) Specs {
		specs := make(Specs)

		for name, raw := range rawSpecs {
			specs[name] = spec.Spec(raw)
		}

		return specs
	}(state.Specs)

	return nil
}

func (s *State) AgentID() datastore.AgentID {
	return s.agentID
}

func (s *State) Specs() Specs {
	return s.specs
}

func (s *State) ClearSpecs() *State {
	s.specs = make(map[spec.Name]spec.Spec)

	return s
}

func (s *State) SetSpec(sp spec.Spec) *State {
	if s.specs == nil {
		s.specs = make(map[spec.Name]spec.Spec)
	}

	s.specs[sp.SpecName()] = sp

	return s
}

func (s *State) GetSpec(name spec.Name, dest any) error {
	spec, exists := s.specs[name]
	if !exists {
		return errors.WithStack(ErrSpecNotFound)
	}

	if err := mapstructure.Decode(spec, dest); err != nil {
		return errors.WithStack(err)
	}

	if err := mapstructure.Decode(spec.SpecData(), dest); err != nil {
		return errors.WithStack(err)
	}

	return nil
}