feat: 实现 WebSocket、mochi broker 与下行发布
This commit is contained in:
Nixevol
@@ -0,0 +1,379 @@
|
||||
package broker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
|
||||
"git.asio.asia/nixevol/NixMsg/internal/app/port"
|
||||
mqtt "github.com/mochi-mqtt/server/v2"
|
||||
"github.com/mochi-mqtt/server/v2/packets"
|
||||
)
|
||||
|
||||
const (
|
||||
maxClients = 2000
|
||||
maxPacketSize = 786432
|
||||
uplinkQueueSize = 256
|
||||
largeFrameBytes = 64 * 1024
|
||||
largeFrameSlots = 64
|
||||
packetOverheadBudget = 128 // 主题与 MQTT 包头预留
|
||||
keepaliveMin = 10
|
||||
keepaliveMax = 600
|
||||
)
|
||||
|
||||
// ErrPayloadTooLarge 下行超过客户端 Maximum Packet Size（减包头预留）或 max_receive_bytes。
|
||||
var ErrPayloadTooLarge = errors.New("broker: payload exceeds client limit")
|
||||
|
||||
// ErrNoConnection 目标端没有当前连接。
|
||||
var ErrNoConnection = errors.New("broker: no active connection")
|
||||
|
||||
// AuthResult 是登录校验结论（N3 实现真实逻辑；N2 默认拒绝）。
|
||||
type AuthResult struct {
|
||||
OK bool
|
||||
SessionToken string // 密码登录成功时由 N3 填写
|
||||
}
|
||||
|
||||
// Authenticator 由 N3 实现；内部故障必须返回 error，不得当成密码错误。
|
||||
type Authenticator interface {
|
||||
Authenticate(ctx context.Context, endpointID string, password []byte, remoteIP string) (AuthResult, error)
|
||||
}
|
||||
|
||||
// RejectAuthenticator 默认拒绝所有客户端（CONNACK 用户名密码错误）。
|
||||
type RejectAuthenticator struct{}
|
||||
|
||||
func (RejectAuthenticator) Authenticate(context.Context, string, []byte, string) (AuthResult, error) {
|
||||
return AuthResult{OK: false}, nil
|
||||
}
|
||||
|
||||
// AllowAuthenticator 测试用：允许任意编号。
|
||||
type AllowAuthenticator struct{}
|
||||
|
||||
func (AllowAuthenticator) Authenticate(context.Context, string, []byte, string) (AuthResult, error) {
|
||||
return AuthResult{OK: true}, nil
|
||||
}
|
||||
|
||||
// Options 装配 broker。
|
||||
type Options struct {
|
||||
Authenticator Authenticator
|
||||
Uplink port.UplinkHandler
|
||||
Logger *slog.Logger
|
||||
}
|
||||
|
||||
// Broker 内置 mochi，不自带监听端口。
|
||||
type Broker struct {
|
||||
server *mqtt.Server
|
||||
auth Authenticator
|
||||
uplink port.UplinkHandler
|
||||
log *slog.Logger
|
||||
|
||||
hook *nixHook
|
||||
|
||||
connsMu sync.RWMutex
|
||||
current map[string]*connState
|
||||
byClient map[*mqtt.Client]*connState
|
||||
|
||||
queuesMu sync.Mutex
|
||||
queues map[string]*uplinkQueue
|
||||
|
||||
largeSem chan struct{}
|
||||
closed atomic.Bool
|
||||
}
|
||||
|
||||
type connState struct {
|
||||
connID port.ConnID
|
||||
endpointID string
|
||||
transport port.Transport
|
||||
remoteIP string
|
||||
client *mqtt.Client
|
||||
maxPacketSize uint32
|
||||
maxRecvBytes int
|
||||
authOK bool
|
||||
authErr error
|
||||
sessionToken string
|
||||
largeHeld int
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// New 创建并 Serve mochi（无监听器）。
|
||||
func New(opts Options) (*Broker, error) {
|
||||
auth := opts.Authenticator
|
||||
if auth == nil {
|
||||
auth = RejectAuthenticator{}
|
||||
}
|
||||
uplink := opts.Uplink
|
||||
if uplink == nil {
|
||||
uplink = port.StubUplinkHandler{}
|
||||
}
|
||||
log := opts.Logger
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
|
||||
caps := mqtt.NewDefaultServerCapabilities()
|
||||
caps.MaximumClients = maxClients
|
||||
caps.MaximumQos = 1
|
||||
caps.MaximumPacketSize = maxPacketSize
|
||||
caps.MaximumSessionExpiryInterval = 0
|
||||
caps.ReceiveMaximum = 1024
|
||||
caps.MaximumInflight = 1024
|
||||
caps.MaximumClientWritesPending = 1024
|
||||
caps.RetainAvailable = 0
|
||||
caps.WildcardSubAvailable = 0
|
||||
caps.SharedSubAvailable = 0
|
||||
caps.TopicAliasMaximum = 0
|
||||
caps.Compatibilities.ObscureNotAuthorized = true
|
||||
|
||||
srv := mqtt.New(&mqtt.Options{
|
||||
InlineClient: true,
|
||||
Capabilities: caps,
|
||||
Logger: log,
|
||||
})
|
||||
|
||||
b := &Broker{
|
||||
server: srv,
|
||||
auth: auth,
|
||||
uplink: uplink,
|
||||
log: log,
|
||||
current: make(map[string]*connState),
|
||||
byClient: make(map[*mqtt.Client]*connState),
|
||||
queues: make(map[string]*uplinkQueue),
|
||||
largeSem: make(chan struct{}, largeFrameSlots),
|
||||
}
|
||||
b.hook = &nixHook{b: b}
|
||||
if err := srv.AddHook(b.hook, nil); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := srv.Serve(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
// Server 返回底层 mochi（测试用）。
|
||||
func (b *Broker) Server() *mqtt.Server { return b.server }
|
||||
|
||||
// Close 关闭 broker。
|
||||
func (b *Broker) Close() error {
|
||||
if b.closed.Swap(true) {
|
||||
return nil
|
||||
}
|
||||
b.queuesMu.Lock()
|
||||
for _, q := range b.queues {
|
||||
q.close()
|
||||
}
|
||||
b.queuesMu.Unlock()
|
||||
return b.server.Close()
|
||||
}
|
||||
|
||||
// AttachTCP 把裸 TCP/TLS 连接交给 mochi；阻塞到连接结束。
|
||||
func (b *Broker) AttachTCP(conn net.Conn) error {
|
||||
return b.server.EstablishConnection("tcp", conn)
|
||||
}
|
||||
|
||||
// AttachWS 把 WebSocket NetConn 交给 mochi；阻塞到连接结束。
|
||||
func (b *Broker) AttachWS(conn net.Conn) error {
|
||||
return b.server.EstablishConnection("ws", conn)
|
||||
}
|
||||
|
||||
// PublishDown 实现 port.Downlink。
|
||||
func (b *Broker) PublishDown(ctx context.Context, endpointID string, connID port.ConnID, payload []byte, opts port.PublishOpts) error {
|
||||
if b.closed.Load() {
|
||||
return errors.New("broker: closed")
|
||||
}
|
||||
st := b.lookupConn(endpointID, connID)
|
||||
if st == nil {
|
||||
return ErrNoConnection
|
||||
}
|
||||
|
||||
limit := effectivePayloadLimit(st.maxPacketSize, st.maxRecvBytes)
|
||||
if limit > 0 && len(payload) > limit {
|
||||
return ErrPayloadTooLarge
|
||||
}
|
||||
|
||||
qos := opts.QoS
|
||||
if qos > 1 {
|
||||
qos = 1
|
||||
}
|
||||
topic := downTopic(endpointID)
|
||||
large := len(payload) > largeFrameBytes
|
||||
|
||||
if large {
|
||||
select {
|
||||
case b.largeSem <- struct{}{}:
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
st.mu.Lock()
|
||||
st.largeHeld++
|
||||
st.mu.Unlock()
|
||||
}
|
||||
|
||||
if err := b.server.Publish(topic, payload, false, qos); err != nil {
|
||||
if large {
|
||||
b.releaseOneLarge(st)
|
||||
}
|
||||
return err
|
||||
}
|
||||
if large && qos == 0 {
|
||||
b.releaseOneLarge(st)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *Broker) releaseOneLarge(st *connState) {
|
||||
st.mu.Lock()
|
||||
if st.largeHeld > 0 {
|
||||
st.largeHeld--
|
||||
st.mu.Unlock()
|
||||
select {
|
||||
case <-b.largeSem:
|
||||
default:
|
||||
}
|
||||
return
|
||||
}
|
||||
st.mu.Unlock()
|
||||
}
|
||||
|
||||
func (b *Broker) releaseAllLarge(st *connState) {
|
||||
st.mu.Lock()
|
||||
n := st.largeHeld
|
||||
st.largeHeld = 0
|
||||
st.mu.Unlock()
|
||||
for i := 0; i < n; i++ {
|
||||
select {
|
||||
case <-b.largeSem:
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Disconnect 实现 port.ConnControl。
|
||||
func (b *Broker) Disconnect(_ context.Context, endpointID string, connID port.ConnID, reason port.DisconnectReason) error {
|
||||
st := b.lookupConn(endpointID, connID)
|
||||
if st == nil {
|
||||
return ErrNoConnection
|
||||
}
|
||||
code := packets.CodeDisconnect
|
||||
switch reason {
|
||||
case port.DisconnectTakenOver:
|
||||
code = packets.ErrSessionTakenOver
|
||||
case port.DisconnectKicked, port.DisconnectFatal:
|
||||
code = packets.ErrAdministrativeAction
|
||||
}
|
||||
return b.server.DisconnectClient(st.client, code)
|
||||
}
|
||||
|
||||
func (b *Broker) lookupConn(endpointID string, connID port.ConnID) *connState {
|
||||
b.connsMu.RLock()
|
||||
defer b.connsMu.RUnlock()
|
||||
if connID != "" {
|
||||
for _, st := range b.byClient {
|
||||
if st.endpointID == endpointID && st.connID == connID {
|
||||
return st
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return b.current[endpointID]
|
||||
}
|
||||
|
||||
func downTopic(endpointID string) string {
|
||||
return "nix/c/" + endpointID + "/down"
|
||||
}
|
||||
|
||||
func upTopic(endpointID string) string {
|
||||
return "nix/c/" + endpointID + "/up"
|
||||
}
|
||||
|
||||
func effectivePayloadLimit(maxPacketSize uint32, maxRecvBytes int) int {
|
||||
limit := 0
|
||||
if maxPacketSize > 0 {
|
||||
if maxPacketSize > packetOverheadBudget {
|
||||
limit = int(maxPacketSize) - packetOverheadBudget
|
||||
}
|
||||
}
|
||||
if maxRecvBytes > 0 {
|
||||
if limit == 0 || maxRecvBytes < limit {
|
||||
limit = maxRecvBytes
|
||||
}
|
||||
}
|
||||
return limit
|
||||
}
|
||||
|
||||
func randomConnID() port.ConnID {
|
||||
var b [16]byte
|
||||
_, _ = rand.Read(b[:])
|
||||
return port.ConnID(hex.EncodeToString(b[:]))
|
||||
}
|
||||
|
||||
func transportOf(cl *mqtt.Client) port.Transport {
|
||||
if cl != nil && cl.Net.Listener == "ws" {
|
||||
return port.TransportWS
|
||||
}
|
||||
return port.TransportTCP
|
||||
}
|
||||
|
||||
func remoteIPOf(cl *mqtt.Client) string {
|
||||
if cl == nil {
|
||||
return ""
|
||||
}
|
||||
addr := cl.Net.Remote
|
||||
if addr == "" && cl.Net.Conn != nil && cl.Net.Conn.RemoteAddr() != nil {
|
||||
addr = cl.Net.Conn.RemoteAddr().String()
|
||||
}
|
||||
host, _, err := net.SplitHostPort(addr)
|
||||
if err != nil {
|
||||
return addr
|
||||
}
|
||||
return host
|
||||
}
|
||||
|
||||
// SetMaxReceiveBytes 供 N3 握手后设置；0 表示不限。
|
||||
func (b *Broker) SetMaxReceiveBytes(endpointID string, connID port.ConnID, n int) {
|
||||
st := b.lookupConn(endpointID, connID)
|
||||
if st == nil {
|
||||
return
|
||||
}
|
||||
st.mu.Lock()
|
||||
st.maxRecvBytes = n
|
||||
st.mu.Unlock()
|
||||
}
|
||||
|
||||
// ConnInfoOf 返回连接信息（测试/N3）。
|
||||
func (b *Broker) ConnInfoOf(endpointID string) (port.ConnInfo, bool) {
|
||||
b.connsMu.RLock()
|
||||
st := b.current[endpointID]
|
||||
b.connsMu.RUnlock()
|
||||
if st == nil {
|
||||
return port.ConnInfo{}, false
|
||||
}
|
||||
return port.ConnInfo{
|
||||
ConnID: st.connID,
|
||||
EndpointID: st.endpointID,
|
||||
Transport: st.transport,
|
||||
RemoteIP: st.remoteIP,
|
||||
SessionToken: st.sessionToken,
|
||||
MaxPacketSize: st.maxPacketSize,
|
||||
}, true
|
||||
}
|
||||
|
||||
func (b *Broker) enqueueUplink(endpointID string, conn port.ConnInfo, payload []byte) {
|
||||
b.queuesMu.Lock()
|
||||
q, ok := b.queues[endpointID]
|
||||
if !ok {
|
||||
q = newUplinkQueue(b, endpointID)
|
||||
b.queues[endpointID] = q
|
||||
}
|
||||
b.queuesMu.Unlock()
|
||||
q.push(uplinkItem{conn: conn, payload: payload})
|
||||
}
|
||||
|
||||
var (
|
||||
_ port.Downlink = (*Broker)(nil)
|
||||
_ port.ConnControl = (*Broker)(nil)
|
||||
)
|
||||
Reference in New Issue
Block a user