story/internal/devicefleet/service_test.go

package devicefleet

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"net"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/mochi-mqtt/server/v2/packets"
)

type fakeStore struct{ objects map[string][]byte }

type fakeDownload struct{ *bytes.Reader }

func (fakeDownload) Close() error { return nil }

func (f *fakeStore) PresignUpload(_ context.Context, key string, _ time.Duration) (string, error) {
	return "https://example.test/" + key, nil
}
func (f *fakeStore) PresignMediaDownload(_ context.Context, key string, _ time.Duration, _, _ string, _ bool) (string, error) {
	return "https://example.test/" + key, nil
}
func (f *fakeStore) OpenDownload(_ context.Context, key string) (io.ReadSeekCloser, error) {
	content, ok := f.objects[key]
	if !ok {
		return nil, errors.New("missing object")
	}
	return fakeDownload{bytes.NewReader(content)}, nil
}
func (f *fakeStore) VerifyDigest(_ context.Context, key, hash string, size int64) error {
	b := f.objects[key]
	sum := sha256.Sum256(b)
	if int64(len(b)) != size || hex.EncodeToString(sum[:]) != hash {
		return errors.New("bad object")
	}
	return nil
}
func (f *fakeStore) FreezeFirmware(ctx context.Context, source, target, hash string, size int64) error {
	f.objects[target] = bytes.Clone(f.objects[source])
	return f.VerifyDigest(ctx, target, hash, size)
}
func TestHeartbeatValidation(t *testing.T) {
	for _, b := range []Heartbeat{{true, "", 0, true}, {true, "x", -1, true}, {true, strings.Repeat("x", 81), 1, true}} {
		if ValidateHeartbeat(b) {
			t.Fatal(b)
		}
	}
	if !ValidateHeartbeat(Heartbeat{true, "1.0.0", 0, true}) || !ValidateHeartbeat(Heartbeat{}) {
		t.Fatal("valid heartbeat rejected")
	}
}
func TestPasswordDeterministicDeviceScoped(t *testing.T) {
	key := bytes.Repeat([]byte{7}, 32)
	a := Password(key, "device-a")
	if len(a) != 43 || a != Password(key, "device-a") || a == Password(key, "device-b") {
		t.Fatal("password derivation")
	}
}

func TestBrokerRejectsTLSURL(t *testing.T) {
	s := &Service{BrokerURL: "mqtts://127.0.0.1:8883"}
	if _, err := s.StartBroker(context.Background(), "127.0.0.1:0", slog.New(slog.NewTextHandler(io.Discard, nil))); err == nil {
		t.Fatal("TLS broker URL accepted for a plaintext listener")
	}
}

func TestFleetMQTTAndOTALifecycle(t *testing.T) {
	dsn := os.Getenv("TEST_DATABASE_URL")
	if dsn == "" {
		t.Skip("isolated integration database required")
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	pool, err := pgxpool.New(ctx, dsn)
	must(t, err)
	defer pool.Close()
	const device = "10000000-0000-0000-0000-000000000502"
	const other = "10000000-0000-0000-0000-000000000501"
	const tenant = "10000000-0000-0000-0000-000000000001"
	const factory = "KP01-020000001234"
	_, err = pool.Exec(ctx, `UPDATE device SET factory_device_id=$2 WHERE id=$1`, device, factory)
	must(t, err)
	store := &fakeStore{objects: map[string][]byte{}}
	s := &Service{Pool: pool, Store: store, BrokerURL: "mqtt://127.0.0.1:1883", DownloadBaseURL: "https://82.157.201.114"}
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	must(t, err)
	address := listener.Addr().String()
	listener.Close()
	stop, err := s.StartBroker(ctx, address, slog.New(slog.NewTextHandler(io.Discard, nil)))
	must(t, err)
	defer stop()
	config, err := s.Config(ctx, device, factory, bytes.Repeat([]byte{7}, 32))
	must(t, err)
	connect := func(password string) net.Conn {
		t.Helper()
		c, e := net.Dial("tcp", address)
		must(t, e)
		pk := packets.Packet{FixedHeader: packets.FixedHeader{Type: packets.Connect}, ProtocolVersion: 4, Connect: packets.ConnectParams{ProtocolName: []byte("MQTT"), ClientIdentifier: factory, Username: []byte(factory), Password: []byte(password), UsernameFlag: true, PasswordFlag: true, Clean: true, Keepalive: 60}}
		var b bytes.Buffer
		must(t, pk.ConnectEncode(&b))
		_, e = c.Write(b.Bytes())
		must(t, e)
		return c
	}
	bad := connect("wrong")
	_, ack := readPacket(t, bad)
	if ack[len(ack)-1] == 0 {
		t.Fatal("bad credentials accepted")
	}
	bad.Close()
	c := connect(config["password"].(string))
	defer c.Close()
	kind, ack := readPacket(t, c)
	if kind>>4 != packets.Connack || ack[len(ack)-1] != 0 {
		t.Fatal("connect", kind, ack)
	}
	state, err := s.State(ctx, device)
	must(t, err)
	if state.Online {
		t.Fatal("connection without heartbeat is not online")
	}
	subscribe := func(topic string, id uint16) byte {
		pk := packets.Packet{FixedHeader: packets.FixedHeader{Type: packets.Subscribe, Qos: 1}, ProtocolVersion: 4, PacketID: id, Filters: packets.Subscriptions{{Filter: topic, Qos: 1}}}
		var b bytes.Buffer
		must(t, pk.SubscribeEncode(&b))
		_, e := c.Write(b.Bytes())
		must(t, e)
		_, data := readPacket(t, c)
		return data[len(data)-1]
	}
	if subscribe("vl/devices/another/ota", 1) != 0x80 {
		t.Fatal("cross-device subscription allowed")
	}
	if subscribe(OTATopic(factory), 2) != 1 {
		t.Fatal("own OTA subscription failed")
	}
	heartbeat := func(version string, code int64) {
		payload, _ := json.Marshal(Heartbeat{true, version, code, true})
		pk := packets.Packet{FixedHeader: packets.FixedHeader{Type: packets.Publish, Qos: 1}, ProtocolVersion: 4, PacketID: 3, TopicName: StatusTopic(factory), Payload: payload}
		var b bytes.Buffer
		must(t, pk.PublishEncode(&b))
		_, e := c.Write(b.Bytes())
		must(t, e)
		k, _ := readPacket(t, c)
		if k>>4 != packets.Puback {
			t.Fatal("heartbeat not acknowledged")
		}
	}
	heartbeat("factory", 0)
	state, err = s.State(ctx, device)
	must(t, err)
	if !state.Online || !state.OTASupported {
		t.Fatal("heartbeat not visible", state)
	}
	authorize := func(tx pgx.Tx) (string, error) {
		var id string
		e := tx.QueryRow(ctx, `SELECT device_id::text FROM device_twin WHERE device_id=$1 FOR UPDATE`, device).Scan(&id)
		return id, e
	}
	sum := sha256.Sum256([]byte("firmware-v1"))
	input := ReleaseInput{Model: "simulated-terminal", Version: "1.0.0", Notes: "test", Size: 11, SHA256: hex.EncodeToString(sum[:])}
	r, err := s.Prepare(ctx, tenant, input)
	must(t, err)
	if r.Code <= 0 {
		t.Fatal("missing generated sequence")
	}
	if _, e := s.Finalize(ctx, tenant, r.ID); e == nil {
		t.Fatal("missing file published")
	}
	store.objects[r.Object] = []byte("firmware-v1")
	input.Size = int64(len(store.objects[r.Object])) // correct the expected size by retrying the staged upload
	r, err = s.Prepare(ctx, tenant, input)
	must(t, err)
	store.objects[r.Object] = []byte("firmware-v1")
	r, err = s.Finalize(ctx, tenant, r.ID)
	must(t, err)
	download, err := s.Download(ctx, tenant, r.ID)
	must(t, err)
	if len(download.DownloadURL) > 64 || !strings.HasPrefix(download.DownloadURL, "https://82.157.201.114/f/") {
		t.Fatal("firmware API did not return a compact download URL", download.DownloadURL)
	}
	adminToken := strings.TrimPrefix(download.DownloadURL, "https://82.157.201.114/f/")
	adminGrant, err := s.ResolveDownload(ctx, adminToken)
	must(t, err)
	if adminGrant.SHA256 != r.SHA256 || adminGrant.Size != r.Size {
		t.Fatal("firmware API ticket resolved to the wrong firmware", adminGrant)
	}
	if _, e := s.Prepare(ctx, tenant, input); e == nil {
		t.Fatal("published version overwritten")
	}
	state, err = s.State(ctx, device)
	must(t, err)
	if !state.UpdateAvailable || state.Latest.Code != r.Code {
		t.Fatal("latest version missing")
	}
	// The older rollout path can have an ASSIGNED device after its command has
	// expired. It must still exclude a second OTA through the MQTT fleet path.
	keyID, legacyReleaseID, legacyRolloutID := uuid.NewString(), uuid.NewString(), uuid.NewString()
	const creator = "10000000-0000-0000-0000-000000000101"
	const household = "10000000-0000-0000-0000-000000000201"
	_, err = pool.Exec(ctx, `INSERT INTO firmware_signing_key(tenant_id,key_id,public_key_base64,created_by) VALUES($1,$2,'test',$3)`, tenant, keyID, creator)
	must(t, err)
	_, err = pool.Exec(ctx, `INSERT INTO firmware_release(id,tenant_id,household_id,model,version,object_ref,expected_size,expected_digest,signature,signing_key_id,created_by) VALUES($1,$2,$3,'simulated-terminal',$4,$4,11,$5,'test',$6,$7)`, legacyReleaseID, tenant, household, legacyReleaseID, r.SHA256, keyID, creator)
	must(t, err)
	_, err = pool.Exec(ctx, `INSERT INTO firmware_rollout(id,tenant_id,household_id,release_id,name,status,target_percentage,created_by) VALUES($1,$2,$3,$4,'mutual exclusion test','RUNNING',100,$5)`, legacyRolloutID, tenant, household, legacyReleaseID, creator)
	must(t, err)
	_, err = pool.Exec(ctx, `INSERT INTO firmware_assignment(rollout_id,tenant_id,device_id,status) VALUES($1,$2,$3,'ASSIGNED')`, legacyRolloutID, tenant, device)
	must(t, err)
	_, err = s.Start(ctx, uuid.NewString(), r.Code, authorize)
	var busyError *Error
	if !errors.As(err, &busyError) || busyError.Code != "DEVICE_BUSY" {
		t.Fatalf("ASSIGNED rollout allowed a parallel fleet OTA: %v", err)
	}
	_, err = pool.Exec(ctx, `DELETE FROM firmware_rollout WHERE id=$1`, legacyRolloutID)
	must(t, err)
	_, err = pool.Exec(ctx, `DELETE FROM firmware_release WHERE id=$1`, legacyReleaseID)
	must(t, err)
	_, err = pool.Exec(ctx, `DELETE FROM firmware_signing_key WHERE tenant_id=$1 AND key_id=$2`, tenant, keyID)
	must(t, err)
	request := uuid.NewString()
	task, err := s.Start(ctx, request, r.Code, authorize)
	must(t, err)
	typ, payload := readPacket(t, c)
	if typ>>4 != packets.Publish || typ&1 != 0 {
		t.Fatal("expected non-retained OTA publish")
	}
	topicLength := int(payload[0])<<8 | int(payload[1])
	var command map[string]any
	must(t, json.Unmarshal(payload[2+topicLength+2:], &command))
	if command["task_id"] != task.ID || command["sha256"] != r.SHA256 {
		t.Fatal("bad OTA command")
	}
	downloadURL, ok := command["download_url"].(string)
	if !ok || len(downloadURL) > 64 || !strings.HasPrefix(downloadURL, "https://82.157.201.114/f/") {
		t.Fatal("OTA command did not contain a compact download URL", downloadURL)
	}
	token := strings.TrimPrefix(downloadURL, "https://82.157.201.114/f/")
	grant, err := s.ResolveDownload(ctx, token)
	must(t, err)
	if grant.SHA256 != r.SHA256 || grant.Size != r.Size {
		t.Fatal("short download ticket resolved to the wrong firmware", grant)
	}
	reader, err := s.OpenDownload(ctx, grant)
	must(t, err)
	downloaded, err := io.ReadAll(reader)
	must(t, err)
	must(t, reader.Close())
	if !bytes.Equal(downloaded, []byte("firmware-v1")) {
		t.Fatal("short download returned wrong content")
	}
	again, err := s.Start(ctx, request, r.Code, authorize)
	must(t, err)
	if again.ID != task.ID {
		t.Fatal("retry duplicated task")
	}
	if _, e := s.Start(ctx, uuid.NewString(), r.Code, authorize); e == nil {
		t.Fatal("parallel OTA allowed")
	}
	if e := s.Report(ctx, other, Report{TaskID: task.ID, Status: "success", FirmwareVersion: r.Version, Code: r.Code}); e == nil {
		t.Fatal("another device completed task")
	}
	if e := s.Report(ctx, device, Report{TaskID: task.ID, Status: "success", FirmwareVersion: "wrong", Code: r.Code}); e == nil {
		t.Fatal("wrong firmware accepted")
	}
	manifest, err := s.Manifest(ctx, device, task.ID)
	must(t, err)
	if manifest["sha256"] != r.SHA256 || manifest["version_code"] != r.Code || manifest["model"] != r.Model {
		t.Fatal("manifest differs from immutable release", manifest)
	}
	manifestToken := strings.TrimPrefix(manifest["download_url"].(string), "https://82.157.201.114/f/")
	manifestGrant, err := s.ResolveDownload(ctx, manifestToken)
	must(t, err)
	if manifestGrant.SHA256 != r.SHA256 {
		t.Fatal("HTTPS manifest ticket returned different bytes")
	}
	if _, err = s.Manifest(ctx, other, task.ID); err == nil {
		t.Fatal("foreign device obtained OTA manifest")
	}
	must(t, s.Report(ctx, device, Report{TaskID: task.ID, Status: "accepted"}))
	must(t, s.Report(ctx, device, Report{TaskID: task.ID, Status: "success", FirmwareVersion: r.Version, Code: r.Code}))
	if _, err = s.Manifest(ctx, device, task.ID); err == nil {
		t.Fatal("completed task gained a new download grant")
	}
	must(t, s.Report(ctx, device, Report{TaskID: task.ID, Status: "success", FirmwareVersion: r.Version, Code: r.Code}))
	if _, err = s.ResolveDownload(ctx, token); !errors.Is(err, ErrDownloadTicketUnavailable) {
		t.Fatal("completed OTA download ticket remained usable", err)
	}
	state, err = s.State(ctx, device)
	must(t, err)
	if state.UpdateAvailable || state.CurrentCode != r.Code || state.Task.Status != "success" {
		t.Fatal("success not reflected", state)
	}
	_, err = pool.Exec(ctx, `UPDATE fleet_presence SET last_seen_at=now()-interval '91 seconds' WHERE device_id=$1`, device)
	must(t, err)
	state, err = s.State(ctx, device)
	must(t, err)
	if state.Online {
		t.Fatal("stale heartbeat is online")
	}
	if _, e := s.Start(ctx, uuid.NewString(), r.Code, authorize); e == nil {
		t.Fatal("offline upgrade allowed")
	}
	_, err = pool.Exec(ctx, `INSERT INTO fleet_ota(device_id,release_id,request_id,created_at) VALUES($1,$2,$3,now()-interval '31 seconds')`, device, r.ID, uuid.NewString())
	must(t, err)
	must(t, s.Expire(ctx))
	state, err = s.State(ctx, device)
	must(t, err)
	var pending int
	must(t, pool.QueryRow(ctx, `SELECT count(*) FROM fleet_ota WHERE device_id=$1 AND status='pending'`, device).Scan(&pending))
	if pending != 0 {
		t.Fatal("pending task not expired")
	}
}
func must(t *testing.T, err error) {
	t.Helper()
	if err != nil {
		t.Fatal(err)
	}
}
func readPacket(t *testing.T, c net.Conn) (byte, []byte) {
	t.Helper()
	must(t, c.SetReadDeadline(time.Now().Add(5*time.Second)))
	var first [1]byte
	_, e := io.ReadFull(c, first[:])
	must(t, e)
	size, mul := 0, 1
	for i := 0; i < 4; i++ {
		var b [1]byte
		_, e = io.ReadFull(c, b[:])
		must(t, e)
		size += int(b[0]&127) * mul
		if b[0]&128 == 0 {
			break
		}
		mul *= 128
	}
	body := make([]byte, size)
	_, e = io.ReadFull(c, body)
	must(t, e)
	return first[0], body
}