mtproto: add compat transport quick ack support
(cherry picked from commit d051bc37bd14076fdd0a83ad41cd507929b20ece)
This commit is contained in:
A
parent
47cab2c9b0
commit
091d8f084b
11 changed files with 915 additions and 21 deletions
|
|
@ -57,6 +57,9 @@ type Conn struct {
|
|||
|
||||
// sentContentMessages 只由 outbound actor 访问，用于生成 MTProto seq_no。
|
||||
sentContentMessages int32
|
||||
// outboundPlain/outboundWire 只由 outbound actor 访问，用于复用出站加密缓冲。
|
||||
outboundPlain bin.Buffer
|
||||
outboundWire bin.Buffer
|
||||
|
||||
identityMu sync.RWMutex
|
||||
businessAuthKeyID [8]byte
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@ package mtprotoedge
|
|||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/binary"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
|
@ -125,6 +127,9 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
|
|||
)
|
||||
return current, s.sendBadMsg(ctx, current, data.MessageID, data.SeqNo, code)
|
||||
}
|
||||
if err := sendQuickAckIfRequested(ctx, tc, key, data); err != nil {
|
||||
return current, err
|
||||
}
|
||||
|
||||
content := clientMessageNeedsAck(typeID)
|
||||
if record, ok := cs.seenRecord(data.MessageID); ok {
|
||||
|
|
@ -170,6 +175,30 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
|
|||
return current, nil
|
||||
}
|
||||
|
||||
func sendQuickAckIfRequested(ctx context.Context, tc transport.Conn, key crypto.AuthKey, data *crypto.EncryptedMessageData) error {
|
||||
q, ok := tc.(quickAckTransport)
|
||||
if !ok || !q.ConsumeQuickAckRequested() {
|
||||
return nil
|
||||
}
|
||||
token, err := clientQuickAckToken(key, data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return q.SendQuickAck(ctx, token)
|
||||
}
|
||||
|
||||
func clientQuickAckToken(key crypto.AuthKey, data *crypto.EncryptedMessageData) (uint32, error) {
|
||||
var plain bin.Buffer
|
||||
if err := data.Encode(&plain); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
h := sha256.New()
|
||||
_, _ = h.Write(key.Value[88:120])
|
||||
_, _ = h.Write(plain.Raw())
|
||||
sum := h.Sum(nil)
|
||||
return binary.LittleEndian.Uint32(sum[:4]) &^ quickAckResponseFlag, nil
|
||||
}
|
||||
|
||||
// dispatch 处理一条明文消息：解包 container/gzip，处理服务消息，其余转 RPC 路由。
|
||||
// content-related 消息（ping、RPC）的 msg_id 会收集到 acks 以便统一确认。
|
||||
func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int64, seqNo int32, b *bin.Buffer, acks *[]int64) error {
|
||||
|
|
@ -560,6 +589,9 @@ func validateClientEnvelope(now time.Time, msgID int64, seqNo int32, typeID uint
|
|||
if msgTime.After(now.Add(30 * time.Second)) {
|
||||
return badMsgIDTooHigh
|
||||
}
|
||||
if clientMessageAllowsEitherSeqParity(typeID) {
|
||||
return 0
|
||||
}
|
||||
if clientMessageNeedsAck(typeID) {
|
||||
if seqNo%2 == 0 {
|
||||
return badMsgSeqNotOdd
|
||||
|
|
@ -593,6 +625,9 @@ func validateClientContainerEnvelope(msgID int64, seqNo int32, typeID uint32) in
|
|||
if msgID == 0 || proto.MessageID(msgID).Type() != proto.MessageFromClient {
|
||||
return badMsgIDInvalidBits
|
||||
}
|
||||
if clientMessageAllowsEitherSeqParity(typeID) {
|
||||
return 0
|
||||
}
|
||||
if clientMessageNeedsAck(typeID) {
|
||||
if seqNo%2 == 0 {
|
||||
return badMsgSeqNotOdd
|
||||
|
|
@ -603,6 +638,15 @@ func validateClientContainerEnvelope(msgID int64, seqNo int32, typeID uint32) in
|
|||
return 0
|
||||
}
|
||||
|
||||
func clientMessageAllowsEitherSeqParity(typeID uint32) bool {
|
||||
switch typeID {
|
||||
case mt.PingDelayDisconnectRequestTypeID:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func clientMessageNeedsAck(typeID uint32) bool {
|
||||
switch typeID {
|
||||
case proto.MessageContainerTypeID,
|
||||
|
|
|
|||
|
|
@ -301,6 +301,32 @@ func TestPingDelayDisconnectEvenSeqAccepted(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestPingDelayDisconnectOddSeqAccepted(t *testing.T) {
|
||||
const dc = 2
|
||||
addr, pub, _ := startTestServer(t, Options{DC: dc})
|
||||
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
|
||||
|
||||
clientMsgID := proto.NewMessageIDGen(time.Now)
|
||||
reqMsgID := clientMsgID.New(proto.MessageFromClient)
|
||||
sendEncryptedWithSeq(t, conn, cipher, auth, reqMsgID, 1, &mt.PingDelayDisconnectRequest{
|
||||
PingID: 10,
|
||||
DisconnectDelay: 60,
|
||||
})
|
||||
|
||||
replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.PongTypeID)
|
||||
if _, ok := replies[mt.BadMsgNotificationTypeID]; ok {
|
||||
t.Fatalf("odd ping_delay_disconnect seq_no produced bad_msg_notification")
|
||||
}
|
||||
buf := mustHave(t, replies, mt.PongTypeID, "pong")
|
||||
var pong mt.Pong
|
||||
if err := pong.Decode(buf); err != nil {
|
||||
t.Fatalf("decode pong: %v", err)
|
||||
}
|
||||
if pong.MsgID != reqMsgID || pong.PingID != 10 {
|
||||
t.Fatalf("pong = %+v, want msg_id=%d ping_id=10", pong, reqMsgID)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDestroyAuthKey 验证 MTProto service message destroy_auth_key 由连接层直接响应，
|
||||
// 避免 TDesktop 清理旧 key 时落到业务 RPC fallback。
|
||||
func TestDestroyAuthKey(t *testing.T) {
|
||||
|
|
|
|||
|
|
@ -2,10 +2,14 @@ package mtprotoedge
|
|||
|
||||
import (
|
||||
"context"
|
||||
"crypto/aes"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"time"
|
||||
|
||||
"github.com/gotd/ige"
|
||||
|
||||
"github.com/gotd/td/bin"
|
||||
"github.com/gotd/td/crypto"
|
||||
"github.com/gotd/td/mt"
|
||||
|
|
@ -45,12 +49,19 @@ type outboundOp struct {
|
|||
ctx context.Context
|
||||
msgType proto.MessageType
|
||||
msg bin.Encoder
|
||||
encoded *encodedOutboundMessage
|
||||
ids []int64
|
||||
reqMsgID int64
|
||||
enqueuedAt time.Time
|
||||
done chan outboundResult
|
||||
}
|
||||
|
||||
type encodedOutboundMessage struct {
|
||||
body []byte
|
||||
typeID uint32
|
||||
reqMsgID int64
|
||||
}
|
||||
|
||||
type outboundResult struct {
|
||||
info []byte
|
||||
resent bool
|
||||
|
|
@ -119,6 +130,14 @@ func (c *Conn) SendPriority(ctx context.Context, t proto.MessageType, msg bin.En
|
|||
// SendBestEffort 只等待消息进入普通 outbound 队列，不等待网络写完成。
|
||||
// 用于 updates fanout：队列拥塞时返回 ErrOutboundQueueFull，durable outbox/getDifference 负责兜底。
|
||||
func (c *Conn) SendBestEffort(ctx context.Context, t proto.MessageType, msg bin.Encoder, timeout time.Duration) error {
|
||||
return c.sendBestEffort(ctx, t, msg, nil, timeout)
|
||||
}
|
||||
|
||||
func (c *Conn) SendBestEffortEncoded(ctx context.Context, t proto.MessageType, encoded *encodedOutboundMessage, timeout time.Duration) error {
|
||||
return c.sendBestEffort(ctx, t, nil, encoded, timeout)
|
||||
}
|
||||
|
||||
func (c *Conn) sendBestEffort(ctx context.Context, t proto.MessageType, msg bin.Encoder, encoded *encodedOutboundMessage, timeout time.Duration) error {
|
||||
if c.outbound == nil || c.outboundControl == nil {
|
||||
return ErrConnClosed
|
||||
}
|
||||
|
|
@ -131,6 +150,7 @@ func (c *Conn) SendBestEffort(ctx context.Context, t proto.MessageType, msg bin.
|
|||
ctx: writeCtx,
|
||||
msgType: t,
|
||||
msg: msg,
|
||||
encoded: encoded,
|
||||
enqueuedAt: time.Now(),
|
||||
}
|
||||
if timeout == 0 {
|
||||
|
|
@ -164,6 +184,14 @@ func (c *Conn) SendBestEffort(ctx context.Context, t proto.MessageType, msg bin.
|
|||
}
|
||||
|
||||
func (c *Conn) send(ctx context.Context, t proto.MessageType, msg bin.Encoder, control bool) error {
|
||||
return c.sendOutbound(ctx, t, msg, nil, control)
|
||||
}
|
||||
|
||||
func (c *Conn) SendEncoded(ctx context.Context, t proto.MessageType, encoded *encodedOutboundMessage) error {
|
||||
return c.sendOutbound(ctx, t, nil, encoded, false)
|
||||
}
|
||||
|
||||
func (c *Conn) sendOutbound(ctx context.Context, t proto.MessageType, msg bin.Encoder, encoded *encodedOutboundMessage, control bool) error {
|
||||
if c.outbound == nil || c.outboundControl == nil {
|
||||
return ErrConnClosed
|
||||
}
|
||||
|
|
@ -173,6 +201,7 @@ func (c *Conn) send(ctx context.Context, t proto.MessageType, msg bin.Encoder, c
|
|||
ctx: ctx,
|
||||
msgType: t,
|
||||
msg: msg,
|
||||
encoded: encoded,
|
||||
enqueuedAt: time.Now(),
|
||||
done: make(chan outboundResult, 1),
|
||||
}
|
||||
|
|
@ -391,7 +420,7 @@ func (c *Conn) handleOutboundOp(state *outboundState, op outboundOp) {
|
|||
}
|
||||
|
||||
func (c *Conn) handleOutboundSend(state *outboundState, op outboundOp) {
|
||||
frame, err := c.buildFrame(op.msgType, op.msg)
|
||||
frame, err := c.buildFrame(op.msgType, op.msg, op.encoded)
|
||||
if err == nil {
|
||||
err = c.writeFrame(op.ctx, frame)
|
||||
}
|
||||
|
|
@ -465,7 +494,26 @@ func (op outboundOp) finish(res outboundResult) {
|
|||
}
|
||||
}
|
||||
|
||||
func (c *Conn) buildFrame(t proto.MessageType, msg bin.Encoder) (*outboundFrame, error) {
|
||||
func (c *Conn) buildFrame(t proto.MessageType, msg bin.Encoder, encoded *encodedOutboundMessage) (*outboundFrame, error) {
|
||||
if encoded == nil {
|
||||
var err error
|
||||
encoded, err = encodeOutboundMessage(msg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
content := frameNeedsAck(encoded.typeID)
|
||||
msgID := c.msgID.New(t)
|
||||
return &outboundFrame{
|
||||
msgID: msgID,
|
||||
seqNo: c.nextSeqNo(content),
|
||||
typeID: encoded.typeID,
|
||||
body: encoded.body,
|
||||
reqMsgID: encoded.reqMsgID,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func encodeOutboundMessage(msg bin.Encoder) (*encodedOutboundMessage, error) {
|
||||
if msg == nil {
|
||||
return nil, errors.New("nil outbound message")
|
||||
}
|
||||
|
|
@ -477,13 +525,9 @@ func (c *Conn) buildFrame(t proto.MessageType, msg bin.Encoder) (*outboundFrame,
|
|||
if err != nil {
|
||||
return nil, fmt.Errorf("peek outbound type id: %w", err)
|
||||
}
|
||||
content := frameNeedsAck(typeID)
|
||||
msgID := c.msgID.New(t)
|
||||
return &outboundFrame{
|
||||
msgID: msgID,
|
||||
seqNo: c.nextSeqNo(content),
|
||||
return &encodedOutboundMessage{
|
||||
typeID: typeID,
|
||||
body: body.Copy(),
|
||||
body: body.Raw(),
|
||||
reqMsgID: outboundRequestMsgID(msg),
|
||||
}, nil
|
||||
}
|
||||
|
|
@ -501,15 +545,8 @@ func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error {
|
|||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
var out bin.Buffer
|
||||
if err := c.cipher.Encrypt(c.key, crypto.EncryptedMessageData{
|
||||
Salt: c.salt,
|
||||
SessionID: c.sessionID,
|
||||
MessageID: frame.msgID,
|
||||
SeqNo: frame.seqNo,
|
||||
MessageDataLen: int32(len(frame.body)),
|
||||
MessageDataWithPadding: frame.body,
|
||||
}, &out); err != nil {
|
||||
out, err := c.encryptOutboundFrame(frame)
|
||||
if err != nil {
|
||||
return fmt.Errorf("encrypt: %w", err)
|
||||
}
|
||||
|
||||
|
|
@ -523,7 +560,7 @@ func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error {
|
|||
if writer == nil {
|
||||
writer = c.transport
|
||||
}
|
||||
if err := writer.Send(sendCtx, &out); err != nil {
|
||||
if err := writer.Send(sendCtx, out); err != nil {
|
||||
return fmt.Errorf("send: %w", err)
|
||||
}
|
||||
if frame.sentAt.IsZero() {
|
||||
|
|
@ -533,6 +570,61 @@ func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error {
|
|||
return nil
|
||||
}
|
||||
|
||||
func (c *Conn) encryptOutboundFrame(frame *outboundFrame) (*bin.Buffer, error) {
|
||||
plain := &c.outboundPlain
|
||||
plain.Reset()
|
||||
plain.PutLong(c.salt)
|
||||
plain.PutLong(c.sessionID)
|
||||
plain.PutLong(frame.msgID)
|
||||
plain.PutInt32(frame.seqNo)
|
||||
plain.PutInt32(int32(len(frame.body)))
|
||||
plain.Put(frame.body)
|
||||
|
||||
paddingOffset := plain.Len()
|
||||
paddingLen := encryptedPaddingLen(paddingOffset)
|
||||
growBinBufferLen(plain, paddingOffset+paddingLen)
|
||||
if _, err := io.ReadFull(c.cipher.Rand(), plain.Buf[paddingOffset:]); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
msgKey := crypto.MessageKey(c.key.Value, plain.Raw(), crypto.Server)
|
||||
key, iv := crypto.Keys(c.key.Value, msgKey, crypto.Server)
|
||||
aesBlock, err := aes.NewCipher(key[:])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
wireLen := len(c.key.ID) + len(msgKey) + plain.Len()
|
||||
wire := &c.outboundWire
|
||||
ensureBinBufferLen(wire, wireLen)
|
||||
copy(wire.Buf[:len(c.key.ID)], c.key.ID[:])
|
||||
copy(wire.Buf[len(c.key.ID):len(c.key.ID)+len(msgKey)], msgKey[:])
|
||||
ige.EncryptBlocks(aesBlock, iv[:], wire.Buf[len(c.key.ID)+len(msgKey):], plain.Raw())
|
||||
return wire, nil
|
||||
}
|
||||
|
||||
func encryptedPaddingLen(l int) int {
|
||||
return 16 + (16 - (l % 16))
|
||||
}
|
||||
|
||||
func ensureBinBufferLen(b *bin.Buffer, n int) {
|
||||
if cap(b.Buf) < n {
|
||||
b.Buf = make([]byte, n)
|
||||
return
|
||||
}
|
||||
b.Buf = b.Buf[:n]
|
||||
}
|
||||
|
||||
func growBinBufferLen(b *bin.Buffer, n int) {
|
||||
if cap(b.Buf) < n {
|
||||
next := make([]byte, n)
|
||||
copy(next, b.Buf)
|
||||
b.Buf = next
|
||||
return
|
||||
}
|
||||
b.Buf = b.Buf[:n]
|
||||
}
|
||||
|
||||
func frameNeedsAck(typeID uint32) bool {
|
||||
switch typeID {
|
||||
case mt.MsgsAckTypeID,
|
||||
|
|
|
|||
|
|
@ -1,17 +1,58 @@
|
|||
package mtprotoedge
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gotd/td/bin"
|
||||
"github.com/gotd/td/crypto"
|
||||
"github.com/gotd/td/mt"
|
||||
"github.com/gotd/td/proto"
|
||||
"github.com/gotd/td/tg"
|
||||
)
|
||||
|
||||
func TestEncryptOutboundFrameDecryptsWithGotdCipher(t *testing.T) {
|
||||
var key crypto.Key
|
||||
if _, err := rand.Read(key[:]); err != nil {
|
||||
t.Fatalf("rand key: %v", err)
|
||||
}
|
||||
authKey := key.WithID()
|
||||
body := mustEncodeTL(t, &mt.NewSessionCreated{
|
||||
FirstMsgID: 111,
|
||||
UniqueID: 222,
|
||||
ServerSalt: 333,
|
||||
})
|
||||
c := &Conn{
|
||||
cipher: crypto.NewServerCipher(rand.Reader),
|
||||
key: authKey,
|
||||
salt: 12345,
|
||||
sessionID: 67890,
|
||||
}
|
||||
out, err := c.encryptOutboundFrame(&outboundFrame{
|
||||
msgID: 7649066000000000001,
|
||||
seqNo: 1,
|
||||
typeID: mt.NewSessionCreatedTypeID,
|
||||
body: body,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("encrypt: %v", err)
|
||||
}
|
||||
data, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(authKey, &bin.Buffer{Buf: append([]byte(nil), out.Raw()...)})
|
||||
if err != nil {
|
||||
t.Fatalf("decrypt: %v", err)
|
||||
}
|
||||
if data.Salt != c.salt || data.SessionID != c.sessionID {
|
||||
t.Fatalf("salt/session = %d/%d, want %d/%d", data.Salt, data.SessionID, c.salt, c.sessionID)
|
||||
}
|
||||
if got := data.Data(); !bytes.Equal(got, body) {
|
||||
t.Fatalf("body = %x, want %x", got, body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOutboundActorSerializesConcurrentSends(t *testing.T) {
|
||||
const dc = 2
|
||||
addr, pub, srv := startTestServer(t, Options{DC: dc})
|
||||
|
|
|
|||
|
|
@ -221,7 +221,7 @@ func (s *Server) Serve(ctx context.Context, ln net.Listener) error {
|
|||
if s.obfuscated {
|
||||
transportListener = transport.ObfuscatedListener(ln)
|
||||
}
|
||||
l := transport.ListenCodec(s.codec, transportListener)
|
||||
l := newCompatTransportListener(s.codec, transportListener)
|
||||
s.log.Info("Serving", zap.String("addr", ln.Addr().String()), zap.Int("dc", s.dc), zap.Bool("obfuscated_tcp", s.obfuscated))
|
||||
defer s.log.Info("Stopped")
|
||||
|
||||
|
|
|
|||
|
|
@ -163,3 +163,63 @@ func TestServerAcceptObfuscatedAbridged(t *testing.T) {
|
|||
t.Fatal("server did not stop after ctx cancel")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerAcceptObfuscatedAbridgedQuickAckFrame(t *testing.T) {
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatalf("listen: %v", err)
|
||||
}
|
||||
|
||||
frames := make(chan int, 1)
|
||||
srv := New(Options{Logger: zaptest.NewLogger(t), ObfuscatedTCP: true})
|
||||
srv.onFrame = func(n int) {
|
||||
select {
|
||||
case frames <- n:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
serveErr := make(chan error, 1)
|
||||
go func() { serveErr <- srv.Serve(ctx, ln) }()
|
||||
|
||||
raw, err := net.Dial("tcp", ln.Addr().String())
|
||||
if err != nil {
|
||||
t.Fatalf("dial: %v", err)
|
||||
}
|
||||
obfs := obfuscator.Obfuscated2(rand.Reader, raw)
|
||||
if err := obfs.Handshake((codec.Abridged{}).ObfuscatedTag(), 2, mtproxy.Secret{}); err != nil {
|
||||
t.Fatalf("obfuscated handshake: %v", err)
|
||||
}
|
||||
|
||||
var b bin.Buffer
|
||||
b.PutInt32(0x12345678)
|
||||
b.PutInt32(0x0badf00d)
|
||||
packet := append([]byte{0x80 | byte(b.Len()/4)}, b.Raw()...)
|
||||
if _, err := obfs.Write(packet); err != nil {
|
||||
t.Fatalf("write quick ack frame: %v", err)
|
||||
}
|
||||
|
||||
select {
|
||||
case n := <-frames:
|
||||
if n != b.Len() {
|
||||
t.Fatalf("received frame len = %d, want %d", n, b.Len())
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("server did not receive quick-ack abridged frame in time")
|
||||
}
|
||||
|
||||
_ = raw.Close()
|
||||
|
||||
cancel()
|
||||
select {
|
||||
case err := <-serveErr:
|
||||
if err != nil {
|
||||
t.Fatalf("serve returned error: %v", err)
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("server did not stop after ctx cancel")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -488,8 +488,16 @@ func (m *SessionManager) PushToUserExceptAuthKeySession(ctx context.Context, use
|
|||
}
|
||||
|
||||
func (m *SessionManager) pushToUser(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
|
||||
getEncoded := onceEncodedOutbound(msg)
|
||||
return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, func(c *Conn) error {
|
||||
return c.Send(ctx, t, msg)
|
||||
if c.outbound == nil || c.outboundControl == nil {
|
||||
return ErrConnClosed
|
||||
}
|
||||
encoded, err := getEncoded()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return c.SendEncoded(ctx, t, encoded)
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -502,11 +510,32 @@ func (m *SessionManager) PushToUserExceptAuthKeySessionBestEffort(ctx context.Co
|
|||
}
|
||||
|
||||
func (m *SessionManager) pushToUserBestEffort(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
|
||||
getEncoded := onceEncodedOutbound(msg)
|
||||
return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, func(c *Conn) error {
|
||||
return c.SendBestEffort(ctx, t, msg, timeout)
|
||||
if c.outbound == nil || c.outboundControl == nil {
|
||||
return ErrConnClosed
|
||||
}
|
||||
encoded, err := getEncoded()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return c.SendBestEffortEncoded(ctx, t, encoded, timeout)
|
||||
})
|
||||
}
|
||||
|
||||
func onceEncodedOutbound(msg bin.Encoder) func() (*encodedOutboundMessage, error) {
|
||||
var (
|
||||
encoded *encodedOutboundMessage
|
||||
err error
|
||||
)
|
||||
return func() (*encodedOutboundMessage, error) {
|
||||
if encoded == nil && err == nil {
|
||||
encoded, err = encodeOutboundMessage(msg)
|
||||
}
|
||||
return encoded, err
|
||||
}
|
||||
}
|
||||
|
||||
func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, send func(*Conn) error) (int, error) {
|
||||
m.mu.Lock()
|
||||
conns := make([]*Conn, 0, len(m.byUser[userID]))
|
||||
|
|
|
|||
|
|
@ -8,11 +8,21 @@ import (
|
|||
|
||||
"go.uber.org/zap/zaptest"
|
||||
|
||||
"github.com/gotd/td/bin"
|
||||
"github.com/gotd/td/mt"
|
||||
"github.com/gotd/td/proto"
|
||||
"github.com/gotd/td/tg"
|
||||
)
|
||||
|
||||
type countingOutboundEncoder struct {
|
||||
count *int
|
||||
}
|
||||
|
||||
func (e *countingOutboundEncoder) Encode(b *bin.Buffer) error {
|
||||
*e.count++
|
||||
return (&tg.UpdatesTooLong{}).Encode(b)
|
||||
}
|
||||
|
||||
// TestSessionManagerRegistry 验证注册表的注册/注销/查找语义（不涉及网络发送）。
|
||||
func TestSessionManagerRegistry(t *testing.T) {
|
||||
sm := NewSessionManager(zaptest.NewLogger(t))
|
||||
|
|
@ -51,6 +61,43 @@ func TestSessionManagerRegistry(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestSessionManagerBestEffortFanoutPreencodesOnce(t *testing.T) {
|
||||
sm := NewSessionManager(zaptest.NewLogger(t))
|
||||
const userID = int64(100)
|
||||
for i := 0; i < 2; i++ {
|
||||
c := &Conn{
|
||||
sessionID: int64(i + 1),
|
||||
authKeyID: [8]byte{byte(i + 1)},
|
||||
outbound: make(chan outboundOp, 1),
|
||||
outboundControl: make(chan outboundOp, 1),
|
||||
outboundStop: make(chan struct{}),
|
||||
}
|
||||
c.userID.Store(userID)
|
||||
c.userIDResolved.Store(true)
|
||||
c.receivesUpdates.Store(true)
|
||||
sm.Register(c)
|
||||
}
|
||||
|
||||
encodes := 0
|
||||
sent, err := sm.PushToUserExceptSessionBestEffort(
|
||||
context.Background(),
|
||||
userID,
|
||||
0,
|
||||
proto.MessageFromServer,
|
||||
&countingOutboundEncoder{count: &encodes},
|
||||
0,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("push: %v", err)
|
||||
}
|
||||
if sent != 2 {
|
||||
t.Fatalf("sent = %d, want 2", sent)
|
||||
}
|
||||
if encodes != 1 {
|
||||
t.Fatalf("encoded %d times, want 1", encodes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionManagerScopesSameSessionIDByAuthKey(t *testing.T) {
|
||||
sm := NewSessionManager(zaptest.NewLogger(t))
|
||||
raw1 := [8]byte{1}
|
||||
|
|
|
|||
456
internal/mtprotoedge/transport_compat.go
Normal file
456
internal/mtprotoedge/transport_compat.go
Normal file
|
|
@ -0,0 +1,456 @@
|
|||
package mtprotoedge
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/go-faster/errors"
|
||||
"go.uber.org/multierr"
|
||||
|
||||
"github.com/gotd/td/bin"
|
||||
tdcrypto "github.com/gotd/td/crypto"
|
||||
"github.com/gotd/td/proto/codec"
|
||||
"github.com/gotd/td/transport"
|
||||
)
|
||||
|
||||
const maxTransportMessageSize = 1 << 24
|
||||
const quickAckResponseFlag = uint32(1 << 31)
|
||||
|
||||
type transportListener interface {
|
||||
Accept() (transport.Conn, error)
|
||||
Close() error
|
||||
Addr() net.Addr
|
||||
}
|
||||
|
||||
type quickAckTransport interface {
|
||||
ConsumeQuickAckRequested() bool
|
||||
SendQuickAck(ctx context.Context, token uint32) error
|
||||
}
|
||||
|
||||
type compatTransportListener struct {
|
||||
codec func() transport.Codec
|
||||
listener net.Listener
|
||||
}
|
||||
|
||||
func newCompatTransportListener(codec func() transport.Codec, listener net.Listener) transportListener {
|
||||
if codec != nil {
|
||||
return transport.ListenCodec(codec, listener)
|
||||
}
|
||||
return &compatTransportListener{listener: listener}
|
||||
}
|
||||
|
||||
func (l *compatTransportListener) Accept() (_ transport.Conn, rErr error) {
|
||||
conn, err := l.listener.Accept()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() {
|
||||
if rErr != nil {
|
||||
multierr.AppendInto(&rErr, conn.Close())
|
||||
}
|
||||
}()
|
||||
|
||||
connCodec, reader, err := detectCompatCodec(conn)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "detect codec")
|
||||
}
|
||||
|
||||
return &compatTransportConn{
|
||||
conn: wrappedCompatConn{
|
||||
reader: reader,
|
||||
Conn: conn,
|
||||
},
|
||||
codec: connCodec,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (l *compatTransportListener) Close() error {
|
||||
return l.listener.Close()
|
||||
}
|
||||
|
||||
func (l *compatTransportListener) Addr() net.Addr {
|
||||
return l.listener.Addr()
|
||||
}
|
||||
|
||||
type wrappedCompatConn struct {
|
||||
reader io.Reader
|
||||
net.Conn
|
||||
}
|
||||
|
||||
func (w wrappedCompatConn) Read(p []byte) (int, error) {
|
||||
return w.reader.Read(p)
|
||||
}
|
||||
|
||||
type compatTransportConn struct {
|
||||
conn net.Conn
|
||||
codec transport.Codec
|
||||
|
||||
readMux sync.Mutex
|
||||
writeMux sync.Mutex
|
||||
}
|
||||
|
||||
func (c *compatTransportConn) Send(ctx context.Context, b *bin.Buffer) error {
|
||||
c.writeMux.Lock()
|
||||
defer c.writeMux.Unlock()
|
||||
|
||||
if err := c.conn.SetWriteDeadline(time.Time{}); err != nil {
|
||||
return errors.Wrap(err, "reset write deadline")
|
||||
}
|
||||
if deadline, ok := ctx.Deadline(); ok {
|
||||
if err := c.conn.SetWriteDeadline(deadline); err != nil {
|
||||
return errors.Wrap(err, "set write deadline")
|
||||
}
|
||||
}
|
||||
|
||||
if err := c.codec.Write(c.conn, b); err != nil {
|
||||
return errors.Wrap(err, "write")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *compatTransportConn) ConsumeQuickAckRequested() bool {
|
||||
q, ok := c.codec.(quickAckCodec)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return q.consumeQuickAckRequested()
|
||||
}
|
||||
|
||||
func (c *compatTransportConn) SendQuickAck(ctx context.Context, token uint32) error {
|
||||
q, ok := c.codec.(quickAckCodec)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
c.writeMux.Lock()
|
||||
defer c.writeMux.Unlock()
|
||||
|
||||
if err := c.conn.SetWriteDeadline(time.Time{}); err != nil {
|
||||
return errors.Wrap(err, "reset write deadline")
|
||||
}
|
||||
if deadline, ok := ctx.Deadline(); ok {
|
||||
if err := c.conn.SetWriteDeadline(deadline); err != nil {
|
||||
return errors.Wrap(err, "set write deadline")
|
||||
}
|
||||
}
|
||||
|
||||
raw := q.quickAckResponse(token)
|
||||
if err := writeAll(c.conn, raw[:]); err != nil {
|
||||
return errors.Wrap(err, "write quick ack")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *compatTransportConn) Recv(ctx context.Context, b *bin.Buffer) error {
|
||||
c.readMux.Lock()
|
||||
defer c.readMux.Unlock()
|
||||
|
||||
if err := c.conn.SetReadDeadline(time.Time{}); err != nil {
|
||||
return errors.Wrap(err, "reset read deadline")
|
||||
}
|
||||
if deadline, ok := ctx.Deadline(); ok {
|
||||
if err := c.conn.SetReadDeadline(deadline); err != nil {
|
||||
return errors.Wrap(err, "set read deadline")
|
||||
}
|
||||
}
|
||||
|
||||
if err := c.codec.Read(c.conn, b); err != nil {
|
||||
return errors.Wrap(err, "read")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *compatTransportConn) Close() error {
|
||||
return c.conn.Close()
|
||||
}
|
||||
|
||||
func detectCompatCodec(c io.Reader) (transport.Codec, io.Reader, error) {
|
||||
var buf [4]byte
|
||||
if _, err := io.ReadFull(c, buf[:1]); err != nil {
|
||||
return nil, nil, errors.Wrap(err, "read first byte")
|
||||
}
|
||||
|
||||
if buf[0] == codec.AbridgedClientStart[0] {
|
||||
return &quickAckAbridgedCodec{}, c, nil
|
||||
}
|
||||
|
||||
if _, err := io.ReadFull(c, buf[1:4]); err != nil {
|
||||
return nil, nil, errors.Wrap(err, "read header")
|
||||
}
|
||||
switch buf {
|
||||
case codec.IntermediateClientStart:
|
||||
return &quickAckIntermediateCodec{}, c, nil
|
||||
case codec.PaddedIntermediateClientStart:
|
||||
return &quickAckPaddedIntermediateCodec{}, c, nil
|
||||
default:
|
||||
return transport.Full.Codec(), io.MultiReader(bytes.NewReader(buf[:]), c), nil
|
||||
}
|
||||
}
|
||||
|
||||
type quickAckCodec interface {
|
||||
transport.Codec
|
||||
consumeQuickAckRequested() bool
|
||||
quickAckResponse(token uint32) [4]byte
|
||||
}
|
||||
|
||||
type quickAckAbridgedCodec struct {
|
||||
quickAckRequested bool
|
||||
}
|
||||
|
||||
func (*quickAckAbridgedCodec) WriteHeader(w io.Writer) error {
|
||||
return (codec.Abridged{}).WriteHeader(w)
|
||||
}
|
||||
|
||||
func (*quickAckAbridgedCodec) ReadHeader(r io.Reader) error {
|
||||
return (codec.Abridged{}).ReadHeader(r)
|
||||
}
|
||||
|
||||
func (*quickAckAbridgedCodec) Write(w io.Writer, b *bin.Buffer) error {
|
||||
if err := validateOutgoingCompatMessage(b); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
words := b.Len() >> 2
|
||||
var header [4]byte
|
||||
headerLen := 1
|
||||
if words < 0x7f {
|
||||
header[0] = byte(words)
|
||||
} else {
|
||||
header[0] = 0x7f
|
||||
header[1] = byte(words)
|
||||
header[2] = byte(words >> 8)
|
||||
header[3] = byte(words >> 16)
|
||||
headerLen = 4
|
||||
}
|
||||
return writeCompatPacket(w, header[:headerLen], b.Raw())
|
||||
}
|
||||
|
||||
func (q *quickAckAbridgedCodec) Read(r io.Reader, b *bin.Buffer) error {
|
||||
requested, err := readQuickAckAbridged(r, b)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "read abridged")
|
||||
}
|
||||
q.quickAckRequested = requested
|
||||
return checkCompatProtocolError(b)
|
||||
}
|
||||
|
||||
func (q *quickAckAbridgedCodec) consumeQuickAckRequested() bool {
|
||||
v := q.quickAckRequested
|
||||
q.quickAckRequested = false
|
||||
return v
|
||||
}
|
||||
|
||||
func (*quickAckAbridgedCodec) quickAckResponse(token uint32) [4]byte {
|
||||
var raw [4]byte
|
||||
binary.BigEndian.PutUint32(raw[:], (token&^quickAckResponseFlag)|quickAckResponseFlag)
|
||||
return raw
|
||||
}
|
||||
|
||||
type quickAckIntermediateCodec struct {
|
||||
quickAckRequested bool
|
||||
}
|
||||
|
||||
func (*quickAckIntermediateCodec) WriteHeader(w io.Writer) error {
|
||||
return (codec.Intermediate{}).WriteHeader(w)
|
||||
}
|
||||
|
||||
func (*quickAckIntermediateCodec) ReadHeader(r io.Reader) error {
|
||||
return (codec.Intermediate{}).ReadHeader(r)
|
||||
}
|
||||
|
||||
func (*quickAckIntermediateCodec) Write(w io.Writer, b *bin.Buffer) error {
|
||||
if err := validateOutgoingCompatMessage(b); err != nil {
|
||||
return err
|
||||
}
|
||||
var header [4]byte
|
||||
binary.LittleEndian.PutUint32(header[:], uint32(b.Len()))
|
||||
return writeCompatPacket(w, header[:], b.Raw())
|
||||
}
|
||||
|
||||
func (q *quickAckIntermediateCodec) Read(r io.Reader, b *bin.Buffer) error {
|
||||
requested, err := readQuickAckIntermediate(r, b, false)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "read intermediate")
|
||||
}
|
||||
q.quickAckRequested = requested
|
||||
return checkCompatProtocolError(b)
|
||||
}
|
||||
|
||||
func (q *quickAckIntermediateCodec) consumeQuickAckRequested() bool {
|
||||
v := q.quickAckRequested
|
||||
q.quickAckRequested = false
|
||||
return v
|
||||
}
|
||||
|
||||
func (*quickAckIntermediateCodec) quickAckResponse(token uint32) [4]byte {
|
||||
var raw [4]byte
|
||||
binary.LittleEndian.PutUint32(raw[:], (token&^quickAckResponseFlag)|quickAckResponseFlag)
|
||||
return raw
|
||||
}
|
||||
|
||||
type quickAckPaddedIntermediateCodec struct {
|
||||
quickAckRequested bool
|
||||
}
|
||||
|
||||
func (*quickAckPaddedIntermediateCodec) WriteHeader(w io.Writer) error {
|
||||
return (codec.PaddedIntermediate{}).WriteHeader(w)
|
||||
}
|
||||
|
||||
func (*quickAckPaddedIntermediateCodec) ReadHeader(r io.Reader) error {
|
||||
return (codec.PaddedIntermediate{}).ReadHeader(r)
|
||||
}
|
||||
|
||||
func (*quickAckPaddedIntermediateCodec) Write(w io.Writer, b *bin.Buffer) error {
|
||||
if err := validateOutgoingCompatMessage(b); err != nil {
|
||||
return err
|
||||
}
|
||||
var padding [4]byte
|
||||
if _, err := io.ReadFull(tdcrypto.DefaultRand(), padding[:]); err != nil {
|
||||
return err
|
||||
}
|
||||
n := int(padding[0] % 4)
|
||||
payload := b.Raw()
|
||||
if n > 0 {
|
||||
payload = append(append([]byte(nil), payload...), padding[:n]...)
|
||||
}
|
||||
var header [4]byte
|
||||
binary.LittleEndian.PutUint32(header[:], uint32(len(payload)))
|
||||
return writeCompatPacket(w, header[:], payload)
|
||||
}
|
||||
|
||||
func (q *quickAckPaddedIntermediateCodec) Read(r io.Reader, b *bin.Buffer) error {
|
||||
requested, err := readQuickAckIntermediate(r, b, true)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "read padded intermediate")
|
||||
}
|
||||
q.quickAckRequested = requested
|
||||
return checkCompatProtocolError(b)
|
||||
}
|
||||
|
||||
func (q *quickAckPaddedIntermediateCodec) consumeQuickAckRequested() bool {
|
||||
v := q.quickAckRequested
|
||||
q.quickAckRequested = false
|
||||
return v
|
||||
}
|
||||
|
||||
func (*quickAckPaddedIntermediateCodec) quickAckResponse(token uint32) [4]byte {
|
||||
var raw [4]byte
|
||||
binary.LittleEndian.PutUint32(raw[:], (token&^quickAckResponseFlag)|quickAckResponseFlag)
|
||||
return raw
|
||||
}
|
||||
|
||||
func readQuickAckAbridged(r io.Reader, b *bin.Buffer) (bool, error) {
|
||||
var first [1]byte
|
||||
if _, err := io.ReadFull(r, first[:]); err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
requested := first[0]&0x80 != 0
|
||||
lengthByte := first[0] & 0x7f
|
||||
var n int
|
||||
if lengthByte == 0x7f {
|
||||
var tail [3]byte
|
||||
if _, err := io.ReadFull(r, tail[:]); err != nil {
|
||||
return false, err
|
||||
}
|
||||
words := uint32(tail[0]) | uint32(tail[1])<<8 | uint32(tail[2])<<16
|
||||
n = int(words << 2)
|
||||
} else {
|
||||
n = int(lengthByte) << 2
|
||||
}
|
||||
|
||||
if err := validateCompatTransportLength(n); err != nil {
|
||||
return false, err
|
||||
}
|
||||
resetCompatBufferN(b, n)
|
||||
if _, err := io.ReadFull(r, b.Buf); err != nil {
|
||||
return false, errors.Wrap(err, "read payload")
|
||||
}
|
||||
return requested, nil
|
||||
}
|
||||
|
||||
func readQuickAckIntermediate(r io.Reader, b *bin.Buffer, padding bool) (bool, error) {
|
||||
var lengthBuf [4]byte
|
||||
if _, err := io.ReadFull(r, lengthBuf[:]); err != nil {
|
||||
return false, errors.Wrap(err, "read length")
|
||||
}
|
||||
rawLength := binary.LittleEndian.Uint32(lengthBuf[:])
|
||||
requested := rawLength&quickAckResponseFlag != 0
|
||||
n := int(rawLength &^ quickAckResponseFlag)
|
||||
if err := validateCompatTransportLength(n); err != nil {
|
||||
return false, err
|
||||
}
|
||||
resetCompatBufferN(b, n)
|
||||
if _, err := io.ReadFull(r, b.Buf); err != nil {
|
||||
return false, errors.Wrap(err, "read payload")
|
||||
}
|
||||
if padding {
|
||||
paddingLength := n % 4
|
||||
b.Buf = b.Buf[:n-paddingLength]
|
||||
}
|
||||
return requested, nil
|
||||
}
|
||||
|
||||
func validateOutgoingCompatMessage(b *bin.Buffer) error {
|
||||
n := b.Len()
|
||||
if err := validateCompatTransportLength(n); err != nil {
|
||||
return err
|
||||
}
|
||||
if n%bin.Word != 0 {
|
||||
return fmt.Errorf("invalid message length %d: not aligned to %d", n, bin.Word)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func writeCompatPacket(w io.Writer, header, payload []byte) error {
|
||||
packet := make([]byte, len(header)+len(payload))
|
||||
copy(packet, header)
|
||||
copy(packet[len(header):], payload)
|
||||
return writeAll(w, packet)
|
||||
}
|
||||
|
||||
func writeAll(w io.Writer, p []byte) error {
|
||||
for len(p) > 0 {
|
||||
n, err := w.Write(p)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n == 0 {
|
||||
return io.ErrShortWrite
|
||||
}
|
||||
p = p[n:]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateCompatTransportLength(n int) error {
|
||||
if n <= 0 || n > maxTransportMessageSize {
|
||||
return fmt.Errorf("invalid message length %d", n)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func resetCompatBufferN(b *bin.Buffer, n int) {
|
||||
if cap(b.Buf) < n {
|
||||
b.Buf = make([]byte, n)
|
||||
return
|
||||
}
|
||||
b.Buf = b.Buf[:n]
|
||||
}
|
||||
|
||||
func checkCompatProtocolError(b *bin.Buffer) error {
|
||||
if b.Len() != bin.Word {
|
||||
return nil
|
||||
}
|
||||
code, err := b.Int32()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return &codec.ProtocolErr{Code: -code}
|
||||
}
|
||||
96
internal/mtprotoedge/transport_compat_test.go
Normal file
96
internal/mtprotoedge/transport_compat_test.go
Normal file
|
|
@ -0,0 +1,96 @@
|
|||
package mtprotoedge
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"testing"
|
||||
|
||||
"github.com/gotd/td/bin"
|
||||
)
|
||||
|
||||
type countWriteBuffer struct {
|
||||
bytes.Buffer
|
||||
writes int
|
||||
}
|
||||
|
||||
func (w *countWriteBuffer) Write(p []byte) (int, error) {
|
||||
w.writes++
|
||||
return w.Buffer.Write(p)
|
||||
}
|
||||
|
||||
func TestQuickAckResponseEncoding(t *testing.T) {
|
||||
const token = 0x01020304
|
||||
|
||||
abridged := (&quickAckAbridgedCodec{}).quickAckResponse(token)
|
||||
if want := []byte{0x81, 0x02, 0x03, 0x04}; !bytes.Equal(abridged[:], want) {
|
||||
t.Fatalf("abridged quick ack = %x, want %x", abridged, want)
|
||||
}
|
||||
|
||||
intermediate := (&quickAckIntermediateCodec{}).quickAckResponse(token)
|
||||
if want := []byte{0x04, 0x03, 0x02, 0x81}; !bytes.Equal(intermediate[:], want) {
|
||||
t.Fatalf("intermediate quick ack = %x, want %x", intermediate, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQuickAckReadFlags(t *testing.T) {
|
||||
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
|
||||
abridgedPacket := append([]byte{0x80 | byte(len(payload)/bin.Word)}, payload...)
|
||||
|
||||
var got bin.Buffer
|
||||
requested, err := readQuickAckAbridged(bytes.NewReader(abridgedPacket), &got)
|
||||
if err != nil {
|
||||
t.Fatalf("read abridged: %v", err)
|
||||
}
|
||||
if !requested {
|
||||
t.Fatal("abridged quick ack flag was not detected")
|
||||
}
|
||||
if !bytes.Equal(got.Raw(), payload) {
|
||||
t.Fatalf("abridged payload = %x, want %x", got.Raw(), payload)
|
||||
}
|
||||
|
||||
var header [4]byte
|
||||
binary.LittleEndian.PutUint32(header[:], uint32(len(payload))|quickAckResponseFlag)
|
||||
intermediatePacket := append(header[:], payload...)
|
||||
requested, err = readQuickAckIntermediate(bytes.NewReader(intermediatePacket), &got, false)
|
||||
if err != nil {
|
||||
t.Fatalf("read intermediate: %v", err)
|
||||
}
|
||||
if !requested {
|
||||
t.Fatal("intermediate quick ack flag was not detected")
|
||||
}
|
||||
if !bytes.Equal(got.Raw(), payload) {
|
||||
t.Fatalf("intermediate payload = %x, want %x", got.Raw(), payload)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompatTransportCodecsWriteSinglePacket(t *testing.T) {
|
||||
var payload bin.Buffer
|
||||
payload.PutInt32(0x01020304)
|
||||
payload.PutInt32(0x05060708)
|
||||
|
||||
t.Run("abridged", func(t *testing.T) {
|
||||
var out countWriteBuffer
|
||||
if err := (&quickAckAbridgedCodec{}).Write(&out, &payload); err != nil {
|
||||
t.Fatalf("write: %v", err)
|
||||
}
|
||||
if out.writes != 1 {
|
||||
t.Fatalf("writes = %d, want 1", out.writes)
|
||||
}
|
||||
if got, want := out.Bytes()[0], byte(payload.Len()/bin.Word); got != want {
|
||||
t.Fatalf("abridged header = %#x, want %#x", got, want)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("intermediate", func(t *testing.T) {
|
||||
var out countWriteBuffer
|
||||
if err := (&quickAckIntermediateCodec{}).Write(&out, &payload); err != nil {
|
||||
t.Fatalf("write: %v", err)
|
||||
}
|
||||
if out.writes != 1 {
|
||||
t.Fatalf("writes = %d, want 1", out.writes)
|
||||
}
|
||||
if got, want := binary.LittleEndian.Uint32(out.Bytes()[:4]), uint32(payload.Len()); got != want {
|
||||
t.Fatalf("intermediate length = %d, want %d", got, want)
|
||||
}
|
||||
})
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue