internal/mtprotoedge/exchange.go
Rama main · réplica de lectura del repositorio original.
package mtprotoedge
import (
"context"
"encoding/binary"
"errors"
"fmt"
"go.uber.org/zap"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/exchange"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/transport"
"telesrv/internal/store"
)
// emptyAuthKeyID is the auth_key_id (all zero) of unencrypted messages (key exchange).
var emptyAuthKeyID [8]byte
// peekAuthKeyID reads the auth_key_id from the first 8 bytes of the message
// without consuming the buffer.
func peekAuthKeyID(b *bin.Buffer) (id [8]byte, err error) {
err = b.PeekN(id[:], len(id))
return id, err
}
// handleExchange performs the server-side MTProto key exchange after receiving
// the first frame with auth_key_id==0.
//
// first is the already-read first frame (req_pq*), handed back to the exchange
// flow via bufferedConn so it can read the handshake messages from the start.
// The auth key + server salt are persisted to AuthKeyStore before DhGenOk is
// sent; on persistence failure the handshake is not confirmed to the client.
func (s *Server) handleExchange(ctx context.Context, conn transport.Conn, first *bin.Buffer) (*bin.Buffer, error) {
if s.key.Zero() {
s.log.Error("Key exchange requested but server RSA key is not configured")
return nil, s.sendProtoError(ctx, conn, codec.CodeAuthKeyNotFound)
}
buffered := newBufferedConn(conn)
buffered.push(first)
// Put an overall duration bound on the whole key exchange. HandshakeIdleTimeout
// only constrains a single read idle and is useless against a client that keeps
// sending—if a client falls into a "ResPQ→nonce desync→resent req_pq" handshake
// restart loop, an unbounded serverExchange would blindly reply ResPQ to every
// req_pq, spinning forever and flushing logs/burning CPU. On timeout this
// handshake is abandoned and the connection closed; the client reconnects with
// a fresh handshake (no residual phase offset) and recovers.
runCtx := ctx
if s.handshakeMaxDur > 0 {
var cancel context.CancelFunc
runCtx, cancel = context.WithTimeout(ctx, s.handshakeMaxDur)
defer cancel()
}
start := s.clock.Now()
res, err := s.runServerExchange(runCtx, buffered)
if err != nil {
// gotd v0.158: reading a non-zero auth_key_id frame mid-handshake (client uses
// an existing auth key instead of re-exchanging) is surfaced via a typed
// UnexpectedEncryptedError with the original frame attached (older versions
// only matched on error text, which broke after the upgrade). Replay that
// frame as the first frame of an existing session; never reply -404—TDesktop
// would decide the temp key was destroyed, discard it and rerun the key
// exchange, triggering a reconnect/re-exchange storm.
var encErr *exchange.UnexpectedEncryptedError
if errors.As(err, &encErr) {
replay := encErr.Frame
if len(replay) > 0 {
s.log.Debug("Key exchange interrupted by encrypted frame; replaying as existing session")
return &bin.Buffer{Buf: replay}, nil
}
}
// req_pq frame count over the bound (client handshake restart loop): drop the
// connection immediately, prompting the client to reconnect with a fresh handshake.
if errors.Is(err, errTooManyHandshakeReqPQ) {
s.log.Info("Key exchange aborted: too many req_pq retries (client handshake restart loop)",
zap.Int("max", maxHandshakeReqPQ))
return nil, err
}
// Only this handshake's overall duration bound expired (ctx itself is not
// canceled): abandon and close, prompting the client to reconnect.
if errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil {
s.log.Info("Key exchange aborted: exceeded max duration (possible client req_pq restart loop)",
zap.Duration("max", s.handshakeMaxDur))
return nil, err
}
var exErr *exchange.ServerExchangeError
if errors.As(err, &exErr) {
s.log.Info("Key exchange rejected", zap.Int32("code", exErr.Code), zap.Error(err))
return nil, s.sendProtoError(ctx, conn, exErr.Code)
}
return nil, fmt.Errorf("key exchange: %w", err)
}
s.metrics.HandshakeDone(s.clock.Now().Sub(start))
// Successful handshakes are already covered by the bounded latency metric.
// One INFO write per PFS connection turns a 10,000-account login burst into
// synchronous logger and filesystem pressure.
s.log.Debug("Key exchange completed",
zap.Int64("auth_key_id", res.Key.IntID()),
zap.Int64("server_salt", res.ServerSalt),
zap.Duration("dur", s.clock.Now().Sub(start)),
)
return nil, nil
}
// authKeyData converts a handshake result into a store record.
func authKeyData(key crypto.AuthKey, salt, createdAt int64, expiresAt int) store.AuthKeyData {
return store.AuthKeyData{
ID: key.ID,
Value: [256]byte(key.Value),
ServerSalt: salt,
CreatedAt: createdAt,
ExpiresAt: expiresAt,
}
}
// sendProtoError sends a transport-level protocol error (-code) to the client.
func (s *Server) sendProtoError(ctx context.Context, conn transport.Conn, code int32) error {
var buf bin.Buffer
buf.PutInt32(-code)
ctx, cancel := context.WithTimeout(ctx, s.writeTimeout)
defer cancel()
if err := conn.Send(ctx, &buf); err != nil {
return fmt.Errorf("send proto error %d: %w", code, err)
}
return nil
}
// sendTerminalProtoError serializes a bare transport error after the authenticated
// outbound actor has stopped. A direct write while the actor is still draining can
// otherwise interleave an encrypted update/result after -404 on the same socket.
func (s *Server) sendTerminalProtoError(ctx context.Context, c *Conn, code int32) error {
if c == nil {
return errors.New("send terminal protocol error without logical connection")
}
c.beginTerminalShutdown()
if !c.waitOutboundShutdownUntil(forceCloseBatchTimeout) {
c.closeTransport()
return errors.New("outbound writer did not stop before terminal protocol error")
}
if c.transport == nil {
return ErrConnClosed
}
return s.sendProtoError(ctx, c.transport, code)
}
// maxHandshakeReqPQ is the upper bound on the number of req_pq(_multi) frames
// allowed within one key exchange. A normal handshake sends only 1 req_pq (even
// the "fake + real" of some clients is only 2); when a client falls into a
// "receive ResPQ→immediately restart the handshake with a new nonce and resend
// req_pq" loop due to nonce desync, it sends req_pq infinitely on the same
// connection, and the gotd-delegated serverExchange blindly replies ResPQ to
// every one, never converging (see the handshake storm in
// docs/client-compat-notes.md). Exceeding this bound immediately drops the
// connection at the telesrv transport layer, prompting the client to reconnect
// with a fresh handshake (no residual phase offset) and recover. (8) leaves ample
// headroom for a few normal reconnect restarts. It complements the overall
// HandshakeMaxDuration bound (per-count / per-time).
const maxHandshakeReqPQ = 8
// errTooManyHandshakeReqPQ means the number of req_pq frames in one handshake
// exceeded maxHandshakeReqPQ (suspected client handshake restart loop). It is
// raised from bufferedConn.Recv and propagated through serverExchange back to
// handleExchange to close the connection.
var errTooManyHandshakeReqPQ = errors.New("too many req_pq frames in one handshake (client restart loop)")
// bufferedConn wraps transport.Conn and can hand already-read frames back to a
// subsequent Recv.
//
// Used for the key exchange: serveConn already read the first frame to peek the
// auth_key_id, then pushes it back for the exchange.
type bufferedConn struct {
transport.Conn
pending []bin.Buffer
reqPQCount int // req_pq(_multi) frames seen in this handshake; accessed only during the handshake (single Recv goroutine)
}
func newBufferedConn(conn transport.Conn) *bufferedConn {
return &bufferedConn{Conn: conn}
}
func (c *bufferedConn) push(b *bin.Buffer) {
if b == nil {
return
}
// serveConn is synchronously blocked in handleExchange, so the first frame's
// backing remains stable until the exchange returns. Keep a slice view instead
// of copying an attacker-sized transport frame.
buf := b.Buf
b.Buf = nil // transfer ownership; serveConn must not pin the frame after next Recv releases it
c.pending = append(c.pending, bin.Buffer{Buf: buf})
}
// Recv returns pushed frames first (FIFO), reading the underlying connection once exhausted.
func (c *bufferedConn) Recv(ctx context.Context, b *bin.Buffer) error {
for {
if len(c.pending) > 0 {
e := c.pending[0]
c.pending[0] = bin.Buffer{}
c.pending = c.pending[1:]
b.ResetTo(e.Buf)
} else {
if err := c.Conn.Recv(ctx, b); err != nil {
return err
}
}
if isUnencryptedMsgsAckFrame(b) {
// The ack is intentionally ignored during exchange. Drop its transport
// backing and shrink the retained high-water charge before the next Recv;
// otherwise a large trailing frame can consume global admission budget for
// the rest of a CPU-heavy key exchange even though no backing remains live.
b.Buf = nil
retainInboundFrameBackings(c.Conn, b)
continue
}
// req_pq count bound: only in effect during the handshake (bufferedConn is
// used only for the key exchange), and the payload id probe costs the same
// order as the msgs_ack skip above, not touching the encrypted-message hot
// path. Over the bound means an immediate drop, stopping the handshake loop.
if isUnencryptedReqPQFrame(b) {
c.reqPQCount++
if c.reqPQCount > maxHandshakeReqPQ {
return errTooManyHandshakeReqPQ
}
}
return nil
}
}
// unencryptedPayloadID returns the type id of the inner TL payload of an
// unencrypted message (auth_key_id==0). ok=false for a non-unencrypted message or
// on decode failure.
func unencryptedPayloadID(frame *bin.Buffer) (uint32, bool) {
if frame == nil || len(frame.Buf) < 24 {
return 0, false
}
if binary.LittleEndian.Uint64(frame.Buf[:8]) != 0 {
return 0, false
}
dataLen := int64(int32(binary.LittleEndian.Uint32(frame.Buf[16:20])))
if dataLen < 4 || dataLen > int64(len(frame.Buf)-20) {
return 0, false
}
return binary.LittleEndian.Uint32(frame.Buf[20:24]), true
}
func isUnencryptedMsgsAckFrame(frame *bin.Buffer) bool {
id, ok := unencryptedPayloadID(frame)
return ok && id == mt.MsgsAckTypeID
}
func isUnencryptedReqPQFrame(frame *bin.Buffer) bool {
id, ok := unencryptedPayloadID(frame)
return ok && (id == mt.ReqPqRequestTypeID || id == mt.ReqPqMultiRequestTypeID)
}