internal/protocol/protocol.go

@@ -0,0 +1,102 @@
+// Package protocol defines the messages agents and the master exchange over
+// their WebSocket connection. Every message is an Envelope whose Payload is
+// one of the payload types below, selected by Type.
+package protocol
+
+import (
+	"encoding/json"
+	"fmt"
+)
+
+// Message types sent from agent to master.
+const (
+	TypeHello    = "hello"
+	TypeSnapshot = "snapshot"
+	TypeUpsert   = "upsert"
+	TypeRemove   = "remove"
+)
+
+// Envelope wraps every message on the wire.
+type Envelope struct {
+	Type    string          `json:"type"`
+	Payload json.RawMessage `json:"payload,omitempty"`
+}
+
+// Hello is the first message an agent sends on every connection.
+type Hello struct {
+	AgentVersion   string `json:"agent_version"`
+	Hostname       string `json:"hostname"`
+	SwarmManager   bool   `json:"swarm_manager"`
+	SwarmClusterID string `json:"swarm_cluster_id,omitempty"`
+}
+
+// Kind says how a service is deployed.
+type Kind string
+
+// Service kinds.
+const (
+	KindCompose   Kind = "compose"
+	KindSwarm     Kind = "swarm"
+	KindContainer Kind = "container"
+)
+
+// Service is one deployable unit as an agent sees it. Which location fields
+// are set depends on Kind.
+type Service struct {
+	Kind           Kind   `json:"kind"`
+	ComposeProject string `json:"compose_project,omitempty"`
+	ComposeService string `json:"compose_service,omitempty"`
+	WorkingDir     string `json:"working_dir,omitempty"`
+	SwarmService   string `json:"swarm_service,omitempty"`
+	ContainerName  string `json:"container_name,omitempty"`
+	Image          string `json:"image"`
+	Digest         string `json:"digest,omitempty"`
+	Status         string `json:"status"`
+	Self           bool   `json:"self,omitempty"`
+}
+
+// Key identifies the service within one agent's report. It is stable across
+// restarts and is what Remove refers to.
+func (s Service) Key() string {
+	switch s.Kind {
+	case KindCompose:
+		return "compose|" + s.ComposeProject + "|" + s.ComposeService
+	case KindSwarm:
+		return "swarm|" + s.SwarmService
+	default:
+		return "container|" + s.ContainerName
+	}
+}
+
+// Snapshot is the agent's complete inventory. Sent after every Hello.
+type Snapshot struct {
+	Services []Service `json:"services"`
+}
+
+// Upsert reports a service that appeared or changed.
+type Upsert struct {
+	Service Service `json:"service"`
+}
+
+// Remove reports that the service with Key is gone.
+type Remove struct {
+	Key string `json:"key"`
+}
+
+// Encode wraps payload in an Envelope of the given type.
+func Encode(typ string, payload any) (Envelope, error) {
+	raw, err := json.Marshal(payload)
+	if err != nil {
+		return Envelope{}, fmt.Errorf("encode %s: %w", typ, err)
+	}
+	return Envelope{Type: typ, Payload: raw}, nil
+}
+
+// Decode unpacks env's payload into T.
+func Decode[T any](env Envelope) (T, error) {
+	var v T
+	if err := json.Unmarshal(env.Payload, &v); err != nil {
+		return v, fmt.Errorf("decode %s: %w", env.Type, err)
+	}
+	return v, nil
+}

internal/protocol/protocol_test.go

@@ -0,0 +1,49 @@
+package protocol
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestServiceKey(t *testing.T) {
+	cases := []struct {
+		name string
+		svc  Service
+		want string
+	}{
+		{"compose", Service{Kind: KindCompose, ComposeProject: "slusk", ComposeService: "web"}, "compose|slusk|web"},
+		{"swarm", Service{Kind: KindSwarm, SwarmService: "gitea_gitea"}, "swarm|gitea_gitea"},
+		{"container", Service{Kind: KindContainer, ContainerName: "adhoc"}, "container|adhoc"},
+	}
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			assert.Equal(t, c.want, c.svc.Key())
+		})
+	}
+}
+
+func TestEncodeDecodeRoundTrip(t *testing.T) {
+	in := Snapshot{Services: []Service{
+		{Kind: KindCompose, ComposeProject: "slusk", ComposeService: "web", Image: "slusk:1.4.2", Status: "running"},
+	}}
+	env, err := Encode(TypeSnapshot, in)
+	require.NoError(t, err)
+
+	raw, err := json.Marshal(env)
+	require.NoError(t, err)
+	var back Envelope
+	require.NoError(t, json.Unmarshal(raw, &back))
+
+	assert.Equal(t, TypeSnapshot, back.Type)
+	out, err := Decode[Snapshot](back)
+	require.NoError(t, err)
+	assert.Equal(t, in, out)
+}
+
+func TestDecodeRejectsMalformedPayload(t *testing.T) {
+	_, err := Decode[Hello](Envelope{Type: TypeHello, Payload: json.RawMessage(`{"hostname": 5}`)})
+	assert.ErrorContains(t, err, "decode hello")
+}