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[string]map[string]spec.Spec

type State struct {
	agentID datastore.AgentID
	specs   Specs
}

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

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

			for name, versions := range specs {
				if _, exists := rawSpecs[name]; !exists {
					rawSpecs[name] = make(map[string]*spec.RawSpec)
				}

				for version, sp := range versions {
					rawSpecs[name][version] = &spec.RawSpec{
						DefinitionName:    sp.SpecDefinitionName(),
						DefinitionVersion: sp.SpecDefinitionVersion(),
						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[string]map[string]*spec.RawSpec `json:"specs"`
	}{}

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

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

		for name, versions := range rawSpecs {
			if _, exists := specs[name]; !exists {
				specs[name] = make(map[string]spec.Spec)
			}
			for version, raw := range versions {
				specs[name][version] = 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[string]map[string]spec.Spec)

	return s
}

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

	name := sp.SpecDefinitionName()

	if _, exists := s.specs[name]; !exists {
		s.specs[name] = make(map[string]spec.Spec)
	}

	version := sp.SpecDefinitionVersion()

	s.specs[name][version] = sp

	return s
}

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

	spec, exists := versions[version]
	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
}