xiaomihome/miot/miot_mdns.go

// Package miot provides MIoT core client for Xiaomi Home devices.
// miot_mdns.go — mDNS service discovery, ported from py-miot/miot_mdns.py.
package miot

import (
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/mdns"

	"xiaomihome/logger"
)

// ============================================================================
// Constants (aligned with Python)
// ============================================================================

const (
	mipsMDNSType             = "_miot-central._tcp.local."
	mipsMDNSRequestTimeoutMs = 5000
	mipsMDNSUpdateIntervalS  = 600
)

// ============================================================================
// MipsServiceState — service state enum
// ============================================================================

// MipsServiceState represents a change in an mDNS-discovered service.
type MipsServiceState int

const (
	MipsServiceAdded   MipsServiceState = 1
	MipsServiceRemoved MipsServiceState = 2
	MipsServiceUpdated MipsServiceState = 3
)

// String returns a human-readable name.
func (s MipsServiceState) String() string {
	switch s {
	case MipsServiceAdded:
		return "added"
	case MipsServiceRemoved:
		return "removed"
	case MipsServiceUpdated:
		return "updated"
	default:
		return "unknown"
	}
}

// ============================================================================
// MipsServiceData — parsed mDNS service entry
// ============================================================================

// MipsServiceData holds the parsed information from a _miot-central._tcp.local.
// mDNS TXT record. The profile binary is base64-decoded and parsed per the
// profile_bin byte layout (see parseProfile).
type MipsServiceData struct {
	Name      string   `json:"name"`
	Addresses []string `json:"addresses"`
	Port      int      `json:"port"`
	Type      string   `json:"type"`
	Server    string   `json:"server"`
	DID       string   `json:"did"`
	GroupID   string   `json:"group_id"`
	Role      int      `json:"role"`
	SuiteMQTT bool     `json:"suite_mqtt"`
}

// validService returns true when role==1 and suiteMQTT is enabled.
// Only primary-role services that support MQTT connections are accepted.
func (d *MipsServiceData) validService() bool {
	return d.Role == 1 && d.SuiteMQTT
}

// ToDict returns the service data as a map, matching the Python to_dict().
func (d *MipsServiceData) ToDict() map[string]interface{} {
	return map[string]interface{}{
		"name":       d.Name,
		"addresses":  d.Addresses,
		"port":       d.Port,
		"type":       d.Type,
		"server":     d.Server,
		"did":        d.DID,
		"group_id":   d.GroupID,
		"role":       d.Role,
		"suite_mqtt": d.SuiteMQTT,
	}
}

// String implements fmt.Stringer.
func (d *MipsServiceData) String() string {
	return fmt.Sprintf("%v", d.ToDict())
}

// ============================================================================
// Profile parsing
// ============================================================================

// parseProfile decodes a base64 profile string and extracts did, groupID,
// role, and suiteMQTT per the documented byte layout:
//
//	did:       bytes[1:9]  big-endian uint64
//	groupID:   bytes[9:17] reversed, then hex-encoded
//	role:      byte 20 >> 4
//	suiteMQTT: (byte 22 >> 1) & 0x01 == 0x01
func parseProfile(profileB64 string) (did, groupID string, role int, suiteMQTT bool, err error) {
	if profileB64 == "" {
		return "", "", 0, false, NewMipsServiceError("empty profile")
	}

	profileBin, err := base64.StdEncoding.DecodeString(profileB64)
	if err != nil {
		return "", "", 0, false, NewMipsServiceError("invalid profile base64: " + err.Error())
	}
	if len(profileBin) < 23 {
		return "", "", 0, false, NewMipsServiceError("profile too short")
	}

	// did: bytes[1:9] big-endian
	did = fmt.Sprintf("%d", binary.BigEndian.Uint64(profileBin[1:9]))

	// groupID: bytes[9:17] reversed + hex
	groupBytes := make([]byte, 8)
	copy(groupBytes, profileBin[9:17])
	reverseBytes(groupBytes)
	groupID = hex.EncodeToString(groupBytes)

	// role: byte 20 >> 4
	role = int(profileBin[20] >> 4)

	// suiteMQTT: (byte 22 >> 1) & 0x01 == 0x01
	suiteMQTT = ((profileBin[22] >> 1) & 0x01) == 0x01

	return did, groupID, role, suiteMQTT, nil
}

// reverseBytes reverses a byte slice in place.
func reverseBytes(b []byte) {
	for i, j := 0, len(b)-1; i < j; i, j = i+1, j-1 {
		b[i], b[j] = b[j], b[i]
	}
}

// ============================================================================
// MipsServiceChangeHandler — callback type
// ============================================================================

// MipsServiceChangeHandler is called when a mDNS service changes state.
// Parameters match the Python callback: groupID, state, data.
type MipsServiceChangeHandler func(groupID string, state MipsServiceState, data map[string]interface{})

// ============================================================================
// MipsService — mDNS browser
// ============================================================================

// MipsService discovers MIoT central hub gateways via mDNS (Bonjour/Zeroconf).
// It continuously polls for _miot-central._tcp.local. services and notifies
// subscribers of additions and changes.
type MipsService struct {
	mu       sync.Mutex
	services map[string]*MipsServiceData // keyed by groupID
	subs     map[subKey]MipsServiceChangeHandler
	lgr      logger.Logger

	stopCh chan struct{}
}

// subKey is the composite key for subscriber lookup: (key, groupID).
type subKey struct {
	key     string
	groupID string
}

// NewMipsService creates a new MipsService.
func NewMipsService() *MipsService {
	s := &MipsService{
		services: make(map[string]*MipsServiceData),
		subs:     make(map[subKey]MipsServiceChangeHandler),
		lgr:      logger.Default(),
	}
	return s
}

// SetLogger sets a custom logger for MipsService.
func (s *MipsService) SetLogger(l logger.Logger) { s.lgr = l }

// Init starts the mDNS browser in a background goroutine.
// Safe to call multiple times; subsequent calls after the first are no-ops
// unless preceded by Deinit().
func (s *MipsService) Init() {
	s.mu.Lock()
	if s.stopCh != nil {
		s.mu.Unlock()
		return
	}
	s.stopCh = make(chan struct{})
	s.mu.Unlock()

	go s.runBrowser()
	s.lgr.Infof("[mdns] mips service browser started")
}

// Deinit stops the browser goroutine and clears all cached services and
// subscriptions.
func (s *MipsService) Deinit() {
	s.mu.Lock()
	ch := s.stopCh
	if ch != nil {
		s.stopCh = nil
		s.mu.Unlock()
		close(ch) // signal goroutine to stop
	} else {
		s.mu.Unlock()
		return
	}

	s.mu.Lock()
	s.services = make(map[string]*MipsServiceData)
	s.subs = make(map[subKey]MipsServiceChangeHandler)
	s.mu.Unlock()

	s.lgr.Infof("[mdns] mips service deinitialized")
}

// GetServices returns cached services. When groupID is non-empty, only that
// groupID's entry is returned.
func (s *MipsService) GetServices(groupID string) map[string]*MipsServiceData {
	s.mu.Lock()
	defer s.mu.Unlock()

	result := make(map[string]*MipsServiceData)
	if groupID != "" {
		if svc, ok := s.services[groupID]; ok {
			result[groupID] = svc
		}
		return result
	}
	for k, v := range s.services {
		result[k] = v
	}
	return result
}

// SubServiceChange subscribes to state changes for a specific groupID.
// key is a unique subscriber identifier. groupID can be "*" to match all.
func (s *MipsService) SubServiceChange(key, groupID string, handler MipsServiceChangeHandler) {
	if key == "" || groupID == "" || handler == nil {
		return
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	s.subs[subKey{key: key, groupID: groupID}] = handler
}

// UnsubServiceChange removes all subscriptions matching the given key.
func (s *MipsService) UnsubServiceChange(key string) {
	if key == "" {
		return
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	for sk := range s.subs {
		if sk.key == key {
			delete(s.subs, sk)
		}
	}
}

// ============================================================================
// Browser goroutine
// ============================================================================

// runBrowser polls for _miot-central._tcp.local. services in a loop.
func (s *MipsService) runBrowser() {
	for {
		entriesCh := make(chan *mdns.ServiceEntry, 32)
		params := &mdns.QueryParam{
			Service: mipsMDNSType,
			Timeout: time.Duration(mipsMDNSRequestTimeoutMs) * time.Millisecond,
			Entries: entriesCh,
		}

		// Run the blocking Query in a goroutine so we can detect stop.
		go func() {
			if err := mdns.Query(params); err != nil {
				s.lgr.Infof("[mdns] query error: %v", err)
			}
			close(entriesCh)
		}()

		for entry := range entriesCh {
			s.handleEntry(entry)
		}

		select {
		case <-s.stopCh:
			return
		case <-time.After(time.Duration(mipsMDNSUpdateIntervalS) * time.Second):
			// poll again
		}
	}
}

// handleEntry parses and processes a single mDNS service entry.
func (s *MipsService) handleEntry(entry *mdns.ServiceEntry) {
	svc, err := entryToServiceData(entry)
	if err != nil {
		s.lgr.Infof("[mdns] skip entry: %v", err)
		return
	}
	if !svc.validService() {
		return
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	if existing, ok := s.services[svc.GroupID]; ok {
		// Compare relevant fields to determine if an update is needed.
		if existing.DID != svc.DID ||
			!strSlicesEqual(existing.Addresses, svc.Addresses) ||
			existing.Port != svc.Port {
			s.services[svc.GroupID] = svc
			s.callSubsLocked(svc.GroupID, MipsServiceUpdated, svc.ToDict())
		}
	} else {
		// New service.
		s.services[svc.GroupID] = svc
		s.callSubsLocked(svc.GroupID, MipsServiceAdded, svc.ToDict())
	}
}

// callSubsLocked notifies all subscribers whose groupID matches.
// Caller MUST hold s.mu.
func (s *MipsService) callSubsLocked(groupID string, state MipsServiceState, data map[string]interface{}) {
	s.lgr.Infof("[mdns] service change: groupID=%s, state=%s", groupID, state)
	for sk, handler := range s.subs {
		if sk.groupID == groupID || sk.groupID == "*" {
			go handler(groupID, state, data)
		}
	}
}

// ============================================================================
// Service entry conversion
// ============================================================================

// entryToServiceData converts a hashicorp/mdns ServiceEntry into a
// MipsServiceData. It extracts the "profile" field from TXT records,
// base64-decodes it, and parses the binary layout.
func entryToServiceData(entry *mdns.ServiceEntry) (*MipsServiceData, error) {
	if entry == nil {
		return nil, NewMipsServiceError("nil service entry")
	}

	// Extract profile from TXT records (InfoFields contains "key=value" pairs).
	profile := ""
	for _, field := range entry.InfoFields {
		if strings.HasPrefix(field, "profile=") {
			profile = strings.TrimPrefix(field, "profile=")
			break
		}
	}
	if profile == "" {
		return nil, NewMipsServiceError("missing profile in TXT records")
	}

	did, groupID, role, suiteMQTT, err := parseProfile(profile)
	if err != nil {
		return nil, err
	}

	if entry.Port == 0 {
		return nil, NewMipsServiceError("invalid port (0)")
	}

	addresses := make([]string, 0, 1)
	if entry.AddrV4 != nil {
		addresses = append(addresses, entry.AddrV4.String())
	}
	if len(addresses) == 0 {
		return nil, NewMipsServiceError("no IPv4 address")
	}

	server := entry.Host
	if server == "" {
		server = entry.Name
	}

	return &MipsServiceData{
		Name:      entry.Name,
		Addresses: addresses,
		Port:      entry.Port,
		Type:      mipsMDNSType,
		Server:    server,
		DID:       did,
		GroupID:   groupID,
		Role:      role,
		SuiteMQTT: suiteMQTT,
	}, nil
}

// ============================================================================
// Helpers
// ============================================================================

// strSlicesEqual compares two string slices for equality.
func strSlicesEqual(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}