359 lines
8.4 KiB
Go
359 lines
8.4 KiB
Go
package accept
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"io"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.asio.asia/nixevol/NixMsg/internal/protocol"
|
|
"git.asio.asia/nixevol/NixMsg/test/harness"
|
|
"github.com/mochi-mqtt/server/v2/packets"
|
|
)
|
|
|
|
// MQTTSession 是验收/弱网用的端侧 MQTT 会话（WebSocket + hello + 应用帧）。
|
|
type MQTTSession struct {
|
|
t *testing.T
|
|
mc harness.MQTTClient
|
|
EndpointID string
|
|
pktID uint16
|
|
mu sync.Mutex
|
|
inbox []map[string]any
|
|
closed bool
|
|
done chan struct{}
|
|
}
|
|
|
|
// AppResp 是 type=resp 的解析结果。
|
|
type AppResp struct {
|
|
OK bool
|
|
Error map[string]any
|
|
Data any
|
|
Raw map[string]any
|
|
}
|
|
|
|
// MQTTLoginOpts 控制握手参数。
|
|
type MQTTLoginOpts struct {
|
|
// MaxReceiveBytes 非 nil 时写入 hello.max_receive_bytes。
|
|
MaxReceiveBytes *int
|
|
}
|
|
|
|
// MQTTLogin 用密码连上 /mqtt、订阅 down、完成 hello。
|
|
func MQTTLogin(t *testing.T, httpBase, endpointID, password string) *MQTTSession {
|
|
t.Helper()
|
|
return MQTTLoginWith(t, httpBase, endpointID, password, MQTTLoginOpts{})
|
|
}
|
|
|
|
// MQTTLoginWith 同 MQTTLogin，可声明接收上限等。
|
|
func MQTTLoginWith(t *testing.T, httpBase, endpointID, password string, opts MQTTLoginOpts) *MQTTSession {
|
|
t.Helper()
|
|
mc, err := harness.DialMQTTWebSocket(httpBase, 15*time.Second)
|
|
if err != nil {
|
|
t.Fatalf("dial mqtt: %v", err)
|
|
}
|
|
s := &MQTTSession{t: t, mc: mc, EndpointID: endpointID, pktID: 10, done: make(chan struct{})}
|
|
s.connectSubscribeHello(password, opts)
|
|
go s.readLoop()
|
|
return s
|
|
}
|
|
|
|
// Close 关闭底层连接。
|
|
func (s *MQTTSession) Close() {
|
|
s.mu.Lock()
|
|
if s.closed {
|
|
s.mu.Unlock()
|
|
return
|
|
}
|
|
s.closed = true
|
|
s.mu.Unlock()
|
|
_ = s.mc.Close()
|
|
select {
|
|
case <-s.done:
|
|
case <-time.After(3 * time.Second):
|
|
}
|
|
}
|
|
|
|
func (s *MQTTSession) nextPkt() uint16 {
|
|
s.pktID++
|
|
if s.pktID == 0 {
|
|
s.pktID = 1
|
|
}
|
|
return s.pktID
|
|
}
|
|
|
|
func (s *MQTTSession) connectSubscribeHello(password string, opts MQTTLoginOpts) {
|
|
t := s.t
|
|
pk := packets.Packet{
|
|
FixedHeader: packets.FixedHeader{Type: packets.Connect},
|
|
ProtocolVersion: 5,
|
|
Connect: packets.ConnectParams{
|
|
ProtocolName: []byte("MQTT"),
|
|
Clean: true,
|
|
ClientIdentifier: s.EndpointID,
|
|
Keepalive: 30,
|
|
UsernameFlag: true,
|
|
Username: []byte(s.EndpointID),
|
|
PasswordFlag: true,
|
|
Password: []byte(password),
|
|
},
|
|
}
|
|
var buf bytes.Buffer
|
|
if err := pk.ConnectEncode(&buf); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := s.mc.Send(buf.Bytes()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
ack, err := s.mc.Recv()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(ack) < 4 || ack[0]>>4 != packets.Connack || ack[3] != 0 {
|
|
t.Fatalf("connack %x", ack)
|
|
}
|
|
|
|
sub := packets.Packet{
|
|
FixedHeader: packets.FixedHeader{Type: packets.Subscribe, Qos: 1},
|
|
ProtocolVersion: 5,
|
|
PacketID: s.nextPkt(),
|
|
Filters: packets.Subscriptions{
|
|
{Filter: "nix/c/" + s.EndpointID + "/down", Qos: 1},
|
|
},
|
|
}
|
|
buf.Reset()
|
|
if err := sub.SubscribeEncode(&buf); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := s.mc.Send(buf.Bytes()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := s.mc.Recv(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
helloFrame := protocol.Hello{V: protocol.Version, Type: protocol.TypeHello, RID: "h0"}
|
|
if opts.MaxReceiveBytes != nil {
|
|
helloFrame.MaxReceiveBytes = opts.MaxReceiveBytes
|
|
}
|
|
hello, _ := protocol.Marshal(helloFrame)
|
|
s.publishRaw(hello)
|
|
deadline := time.Now().Add(10 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
raw, err := s.mc.Recv()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
m := s.handlePacket(raw)
|
|
if m == nil {
|
|
continue
|
|
}
|
|
if m["type"] == "resp" && m["ok"] == true {
|
|
return
|
|
}
|
|
if m["type"] == "resp" {
|
|
t.Fatalf("hello failed: %v", m)
|
|
}
|
|
s.push(m)
|
|
}
|
|
t.Fatal("hello timeout")
|
|
}
|
|
|
|
func (s *MQTTSession) publishRaw(payload []byte) {
|
|
t := s.t
|
|
pub := packets.Packet{
|
|
FixedHeader: packets.FixedHeader{Type: packets.Publish, Qos: 1},
|
|
ProtocolVersion: 5,
|
|
TopicName: "nix/c/" + s.EndpointID + "/up",
|
|
PacketID: s.nextPkt(),
|
|
Payload: payload,
|
|
}
|
|
var buf bytes.Buffer
|
|
if err := pub.PublishEncode(&buf); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := s.mc.Send(buf.Bytes()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func (s *MQTTSession) readLoop() {
|
|
defer close(s.done)
|
|
for {
|
|
raw, err := s.mc.Recv()
|
|
if err != nil {
|
|
return
|
|
}
|
|
m := s.handlePacket(raw)
|
|
if m != nil {
|
|
s.push(m)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *MQTTSession) handlePacket(raw []byte) map[string]any {
|
|
if len(raw) < 2 {
|
|
return nil
|
|
}
|
|
typ := raw[0] >> 4
|
|
qos := (raw[0] >> 1) & 0x3
|
|
switch typ {
|
|
case packets.Puback, packets.Pingresp, packets.Suback:
|
|
return nil
|
|
case packets.Publish:
|
|
payload, err := decodePublishPayload(raw)
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
if qos == 1 {
|
|
rem, n, _ := decodeRemainingLength(raw[1:])
|
|
body := raw[1+n:]
|
|
pk := packets.Packet{ProtocolVersion: 5, FixedHeader: packets.FixedHeader{Type: packets.Publish, Remaining: rem, Qos: qos}}
|
|
if decErr := pk.PublishDecode(body); decErr == nil && pk.PacketID != 0 {
|
|
ack := packets.Packet{
|
|
FixedHeader: packets.FixedHeader{Type: packets.Puback},
|
|
ProtocolVersion: 5,
|
|
PacketID: pk.PacketID,
|
|
}
|
|
var buf bytes.Buffer
|
|
if encErr := ack.PubackEncode(&buf); encErr == nil {
|
|
_ = s.mc.Send(buf.Bytes())
|
|
}
|
|
}
|
|
}
|
|
var m map[string]any
|
|
if json.Unmarshal(payload, &m) != nil {
|
|
return nil
|
|
}
|
|
return m
|
|
default:
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func (s *MQTTSession) push(m map[string]any) {
|
|
s.mu.Lock()
|
|
s.inbox = append(s.inbox, m)
|
|
s.mu.Unlock()
|
|
}
|
|
|
|
// Request 发上行帧并等同 rid 的 resp。
|
|
func (s *MQTTSession) Request(t *testing.T, frame map[string]any) AppResp {
|
|
t.Helper()
|
|
rid, _ := frame["rid"].(string)
|
|
payload, err := protocol.Marshal(frame)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
s.publishRaw(payload)
|
|
deadline := time.Now().Add(15 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
m := s.takeMatching(func(x map[string]any) bool {
|
|
return x["type"] == "resp" && x["rid"] == rid
|
|
})
|
|
if m != nil {
|
|
r := AppResp{OK: m["ok"] == true, Raw: m, Data: m["data"]}
|
|
if e, ok := m["error"].(map[string]any); ok {
|
|
r.Error = e
|
|
}
|
|
return r
|
|
}
|
|
time.Sleep(5 * time.Millisecond)
|
|
}
|
|
t.Fatalf("timeout waiting resp rid=%s", rid)
|
|
return AppResp{}
|
|
}
|
|
|
|
// WaitType 等到指定 type 的下行帧。
|
|
func (s *MQTTSession) WaitType(t *testing.T, typ string, timeout time.Duration) map[string]any {
|
|
t.Helper()
|
|
deadline := time.Now().Add(timeout)
|
|
for time.Now().Before(deadline) {
|
|
m := s.takeMatching(func(x map[string]any) bool { return x["type"] == typ })
|
|
if m != nil {
|
|
return m
|
|
}
|
|
time.Sleep(5 * time.Millisecond)
|
|
}
|
|
t.Fatalf("timeout waiting type=%s", typ)
|
|
return nil
|
|
}
|
|
|
|
// TryType 在超时内尝试取指定 type；超时返回 nil。
|
|
func (s *MQTTSession) TryType(typ string, timeout time.Duration) map[string]any {
|
|
deadline := time.Now().Add(timeout)
|
|
for time.Now().Before(deadline) {
|
|
m := s.takeMatching(func(x map[string]any) bool { return x["type"] == typ })
|
|
if m != nil {
|
|
return m
|
|
}
|
|
time.Sleep(10 * time.Millisecond)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *MQTTSession) takeMatching(pred func(map[string]any) bool) map[string]any {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
for i, m := range s.inbox {
|
|
if pred(m) {
|
|
s.inbox = append(s.inbox[:i], s.inbox[i+1:]...)
|
|
return m
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// DrainEvents 排空 group_event / presence，避免干扰断言。
|
|
func DrainEvents(t *testing.T, s *MQTTSession, d time.Duration) {
|
|
t.Helper()
|
|
deadline := time.Now().Add(d)
|
|
for time.Now().Before(deadline) {
|
|
_ = s.takeMatching(func(x map[string]any) bool {
|
|
typ, _ := x["type"].(string)
|
|
return typ == "group_event" || typ == "presence"
|
|
})
|
|
time.Sleep(20 * time.Millisecond)
|
|
}
|
|
}
|
|
|
|
func decodePublishPayload(raw []byte) ([]byte, error) {
|
|
if len(raw) < 2 {
|
|
return nil, io.ErrUnexpectedEOF
|
|
}
|
|
rem, n, err := decodeRemainingLength(raw[1:])
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
body := raw[1+n:]
|
|
if len(body) != rem {
|
|
return nil, io.ErrUnexpectedEOF
|
|
}
|
|
pk := packets.Packet{
|
|
ProtocolVersion: 5,
|
|
FixedHeader: packets.FixedHeader{
|
|
Type: packets.Publish,
|
|
Remaining: rem,
|
|
Qos: (raw[0] >> 1) & 0x3,
|
|
},
|
|
}
|
|
if err := pk.PublishDecode(body); err != nil {
|
|
return nil, err
|
|
}
|
|
return pk.Payload, nil
|
|
}
|
|
|
|
func decodeRemainingLength(b []byte) (value int, n int, err error) {
|
|
var mul uint32 = 1
|
|
var v uint32
|
|
for i := 0; i < len(b) && i < 4; i++ {
|
|
v += uint32(b[i]&127) * mul
|
|
n++
|
|
if b[i]&128 == 0 {
|
|
return int(v), n, nil
|
|
}
|
|
mul *= 128
|
|
}
|
|
return 0, 0, io.ErrUnexpectedEOF
|
|
}
|