story/internal/devicefleet/mqtt.go

package devicefleet

import (
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/url"
	"sync"
	"time"

	"github.com/google/uuid"
	mqtt "github.com/mochi-mqtt/server/v2"
	"github.com/mochi-mqtt/server/v2/listeners"
	"github.com/mochi-mqtt/server/v2/packets"
)

type connection struct{ device, token string }
type fleetHook struct {
	mqtt.HookBase
	service *Service
	clients sync.Map
}

func (h *fleetHook) ID() string { return "voicelink" }
func (h *fleetHook) Provides(b byte) bool {
	switch b {
	case mqtt.OnConnectAuthenticate, mqtt.OnACLCheck, mqtt.OnSessionEstablished, mqtt.OnDisconnect, mqtt.OnPublish:
		return true
	}
	return false
}
func bounded() (context.Context, context.CancelFunc) {
	return context.WithTimeout(context.Background(), 5*time.Second)
}
func (h *fleetHook) OnConnectAuthenticate(cl *mqtt.Client, pk packets.Packet) bool {
	if cl.ID == "" || cl.ID != string(pk.Connect.Username) || !pk.Connect.Clean {
		return false
	}
	// Will and application data are never retained, and cannot target another device.
	if pk.Connect.WillFlag && (pk.Connect.WillTopic != StatusTopic(cl.ID) || pk.Connect.WillRetain) {
		return false
	}
	ctx, cancel := bounded()
	defer cancel()
	var device string
	var hash []byte
	err := h.service.Pool.QueryRow(ctx, `SELECT d.id::text,p.password_hash FROM device d JOIN fleet_presence p ON p.device_id=d.id JOIN tenant t ON t.id=d.tenant_id WHERE d.factory_device_id=$1 AND d.status='registered' AND d.activation_state NOT IN ('LOST','REVOKED') AND t.status='active'`, cl.ID).Scan(&device, &hash)
	actual := sha256.Sum256(pk.Connect.Password)
	if err != nil || subtle.ConstantTimeCompare(hash, actual[:]) != 1 {
		return false
	}
	h.clients.Store(cl, connection{device, uuid.NewString()})
	return true
}
func (h *fleetHook) OnACLCheck(cl *mqtt.Client, topic string, write bool) bool {
	if cl.Net.Inline {
		return true
	}
	if _, ok := h.clients.Load(cl); !ok {
		return false
	}
	if write {
		return topic == StatusTopic(cl.ID)
	}
	return topic == OTATopic(cl.ID)
}
func (h *fleetHook) OnSessionEstablished(cl *mqtt.Client, _ packets.Packet) {
	value, ok := h.clients.Load(cl)
	if !ok {
		return
	}
	c := value.(connection)
	ctx, cancel := bounded()
	defer cancel()
	_, err := h.service.Pool.Exec(ctx, `UPDATE fleet_presence SET connection_id=$2,connected=true,last_seen_at=NULL WHERE device_id=$1`, c.device, c.token)
	if err != nil {
		h.Log.Error("MQTT connection persistence failed")
	}
}
func (h *fleetHook) OnDisconnect(cl *mqtt.Client, _ error, _ bool) {
	value, ok := h.clients.LoadAndDelete(cl)
	if !ok {
		return
	}
	c := value.(connection)
	ctx, cancel := bounded()
	defer cancel()
	_, err := h.service.Pool.Exec(ctx, `UPDATE fleet_presence SET connected=false WHERE device_id=$1 AND connection_id=$2`, c.device, c.token)
	if err != nil {
		h.Log.Error("MQTT disconnect persistence failed")
	}
}

type Heartbeat struct {
	Online       bool   `json:"online"`
	Version      string `json:"firmware_version"`
	Code         int64  `json:"version_code"`
	OTASupported bool   `json:"ota_supported"`
}

func ValidateHeartbeat(b Heartbeat) bool {
	return b.Code >= 0 && len(b.Version) <= 80 && (!b.Online || b.Version != "")
}
func (h *fleetHook) OnPublish(cl *mqtt.Client, pk packets.Packet) (packets.Packet, error) {
	if cl.Net.Inline {
		return pk, nil
	}
	value, ok := h.clients.Load(cl)
	if !ok || pk.FixedHeader.Retain || pk.TopicName != StatusTopic(cl.ID) || len(pk.Payload) > 2048 {
		return pk, packets.ErrRejectPacket
	}
	var beat Heartbeat
	if json.Unmarshal(pk.Payload, &beat) != nil || !ValidateHeartbeat(beat) {
		return pk, packets.ErrRejectPacket
	}
	c := value.(connection)
	ctx, cancel := bounded()
	defer cancel()
	tag, err := h.service.Pool.Exec(ctx, `UPDATE fleet_presence p SET connected=$3,last_seen_at=now(),firmware_version=CASE WHEN $3 THEN $4 ELSE firmware_version END,version_code=CASE WHEN $3 THEN $5 ELSE version_code END,ota_supported=CASE WHEN $3 THEN $6 ELSE ota_supported END FROM device d WHERE p.device_id=$1 AND p.connection_id=$2 AND d.id=p.device_id AND d.status='registered' AND d.activation_state NOT IN ('LOST','REVOKED')`, c.device, c.token, beat.Online, beat.Version, beat.Code, beat.OTASupported)
	if err != nil || tag.RowsAffected() != 1 {
		return pk, packets.ErrRejectPacket
	}
	return pk, nil
}

// StartBroker serves one API instance over plain MQTT/TCP.
func (s *Service) StartBroker(ctx context.Context, address string, logger *slog.Logger) (func(), error) {
	if s.BrokerURL == "" {
		return func() {}, nil
	}
	publicURL, err := url.Parse(s.BrokerURL)
	if err != nil || publicURL.Scheme != "mqtt" || publicURL.Host == "" {
		return nil, fmt.Errorf("DEVICE_MQTT_PUBLIC_URL must use mqtt:// with a host")
	}
	broker := mqtt.New(&mqtt.Options{InlineClient: true, Logger: logger})
	broker.Options.Capabilities.MaximumPacketSize = 8192
	if err := broker.AddHook(&fleetHook{service: s}, nil); err != nil {
		return nil, err
	}
	err = broker.AddListener(listeners.NewTCP(listeners.Config{ID: "devices", Address: address}))
	if err != nil {
		return nil, err
	}
	if _, err = s.Pool.Exec(ctx, `UPDATE fleet_presence SET connected=false`); err != nil {
		_ = broker.Close()
		return nil, err
	}
	if err = broker.Serve(); err != nil {
		_ = broker.Close()
		return nil, err
	}
	s.Publish = func(topic string, body []byte) error {
		if ctx.Err() != nil {
			return errors.New("broker stopping")
		}
		return broker.Publish(topic, body, false, 1)
	}
	go func() {
		timer := time.NewTicker(5 * time.Second)
		defer timer.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-timer.C:
				if e := s.Expire(ctx); e != nil {
					logger.Error("OTA expiry check failed")
				}
			}
		}
	}()
	return func() { _ = broker.Close() }, nil
}