emissary/internal/agent/state.go

145 lines
3.0 KiB
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
}