DoGaMa-serv/internal/agent/agent_registry.go

package agent

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"sync"
)

const registryVersion = 1

// RegisteredInstance binds a DoGaMa instance to its expected Docker identity.
type RegisteredInstance struct {
	InstanceID  string `json:"instance_id"`
	ContainerID string `json:"container_id"`
	PlanDigest  string `json:"plan_digest"`
}

type registryPayload struct {
	Version   int                  `json:"version"`
	Instances []RegisteredInstance `json:"instances"`
}

type registryEnvelope struct {
	Payload registryPayload `json:"payload"`
	MAC     string          `json:"mac"`
}

// Registry is the authenticated durable agent-local registration store.
type Registry struct {
	mu        sync.RWMutex
	path      string
	macKey    []byte
	instances []RegisteredInstance
}

// OpenRegistry loads and verifies the registry or creates an empty one.
func OpenRegistry(path string, secret []byte) (*Registry, error) {
	if !filepath.IsAbs(path) {
		return nil, errors.New("registry path must be absolute")
	}
	keyMAC := hmac.New(sha256.New, secret)
	_, _ = keyMAC.Write([]byte("dogama-agent-registry-v1"))
	registry := &Registry{path: path, macKey: keyMAC.Sum(nil)}
	if err := registry.load(); err != nil {
		if !errors.Is(err, os.ErrNotExist) {
			return nil, err
		}
		if err := registry.save(nil); err != nil {
			return nil, err
		}
	}
	return registry, nil
}

// List returns a stable copy of every authenticated registration.
func (r *Registry) List() []RegisteredInstance {
	r.mu.RLock()
	defer r.mu.RUnlock()
	result := append([]RegisteredInstance(nil), r.instances...)
	sort.Slice(result, func(i, j int) bool { return result[i].InstanceID < result[j].InstanceID })
	return result
}

func (r *Registry) Get(instanceID string) (RegisteredInstance, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	for _, entry := range r.instances {
		if entry.InstanceID == instanceID {
			return entry, true
		}
	}
	return RegisteredInstance{}, false
}

func (r *Registry) Register(entry RegisteredInstance) error {
	if entry.InstanceID == "" || entry.ContainerID == "" || entry.PlanDigest == "" {
		return errors.New("incomplete agent registration")
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	for _, existing := range r.instances {
		if existing.InstanceID == entry.InstanceID || existing.ContainerID == entry.ContainerID {
			if existing == entry {
				return nil
			}
			return errors.New("instance registration conflict")
		}
	}
	instances := append(append([]RegisteredInstance(nil), r.instances...), entry)
	return r.saveLocked(instances)
}

func (r *Registry) Remove(instanceID string) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	instances := make([]RegisteredInstance, 0, len(r.instances))
	found := false
	for _, entry := range r.instances {
		if entry.InstanceID == instanceID {
			found = true
			continue
		}
		instances = append(instances, entry)
	}
	if !found {
		return errors.New("instance is not registered")
	}
	return r.saveLocked(instances)
}

// Replace atomically persists a new binding for an already registered instance.
func (r *Registry) Replace(entry RegisteredInstance) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	instances := append([]RegisteredInstance(nil), r.instances...)
	for index := range instances {
		if instances[index].InstanceID == entry.InstanceID {
			instances[index] = entry
			return r.saveLocked(instances)
		}
	}
	return errors.New("instance registration not found")
}

func (r *Registry) load() error {
	info, err := os.Lstat(r.path)
	if err != nil {
		return err
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return errors.New("registry must be a regular file")
	}
	body, err := os.ReadFile(r.path)
	if err != nil {
		return fmt.Errorf("read agent registry: %w", err)
	}
	var envelope registryEnvelope
	decoderErr := json.Unmarshal(body, &envelope)
	if decoderErr != nil || envelope.Payload.Version != registryVersion || !r.validMAC(envelope.Payload, envelope.MAC) {
		return errors.New("agent registry integrity check failed")
	}
	r.instances = append([]RegisteredInstance(nil), envelope.Payload.Instances...)
	return nil
}

func (r *Registry) save(instances []RegisteredInstance) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.saveLocked(instances)
}

func (r *Registry) saveLocked(instances []RegisteredInstance) error {
	payload := registryPayload{Version: registryVersion, Instances: append([]RegisteredInstance(nil), instances...)}
	envelope := registryEnvelope{Payload: payload, MAC: r.mac(payload)}
	body, err := json.Marshal(envelope)
	if err != nil {
		return fmt.Errorf("encode agent registry: %w", err)
	}
	if err := os.MkdirAll(filepath.Dir(r.path), 0o700); err != nil {
		return fmt.Errorf("create agent registry directory: %w", err)
	}
	temporary, err := os.CreateTemp(filepath.Dir(r.path), ".registry-*")
	if err != nil {
		return fmt.Errorf("create agent registry temporary file: %w", err)
	}
	temporaryPath := temporary.Name()
	defer os.Remove(temporaryPath)
	if err = temporary.Chmod(0o600); err == nil {
		_, err = temporary.Write(body)
	}
	if err == nil {
		err = temporary.Sync()
	}
	if closeErr := temporary.Close(); err == nil {
		err = closeErr
	}
	if err != nil {
		return fmt.Errorf("write agent registry: %w", err)
	}
	if err := os.Rename(temporaryPath, r.path); err != nil {
		return fmt.Errorf("replace agent registry: %w", err)
	}
	r.instances = append([]RegisteredInstance(nil), instances...)
	return nil
}

func (r *Registry) mac(payload registryPayload) string {
	body, _ := json.Marshal(payload)
	mac := hmac.New(sha256.New, r.macKey)
	_, _ = mac.Write(body)
	return base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
}

func (r *Registry) validMAC(payload registryPayload, provided string) bool {
	providedBytes, err := base64.RawURLEncoding.DecodeString(provided)
	if err != nil {
		return false
	}
	expectedBytes, _ := base64.RawURLEncoding.DecodeString(r.mac(payload))
	return hmac.Equal(providedBytes, expectedBytes)
}