doctor: deepen DC verification with MTProto handshake probe

internal/cli/doctor.go

@@ -18,6 +18,7 @@ import (
 	"github.com/9seconds/mtg/v2/internal/config"
 	"github.com/9seconds/mtg/v2/internal/utils"
 	"github.com/9seconds/mtg/v2/mtglib"
+	"github.com/9seconds/mtg/v2/mtglib/dcprobe"
 	"github.com/9seconds/mtg/v2/network/v2"
 	"github.com/beevik/ntp"
 )
@@ -46,7 +47,7 @@ var (
 	)
 
 	tplODCConnect = template.Must(
-		template.New("").Parse("  ✅ DC {{ .dc }}\n"),
+		template.New("").Parse("  ✅ DC {{ .dc }} (rpc {{ .rtt }})\n"),
 	)
 	tplEDCConnect = template.Must(
 		template.New("").Parse("  ❌ DC {{ .dc }}: {{ .error }}\n"),
@@ -238,32 +239,37 @@ func (d *Doctor) checkNetwork(ntw mtglib.Network) bool {
 	dcs := slices.Collect(maps.Keys(essentials.TelegramCoreAddresses))
 	slices.Sort(dcs)
 
-	errs := make([]error, len(dcs))
+	type dcResult struct {
+		rtt time.Duration
+		err error
+	}
+	results := make([]dcResult, len(dcs))
 
 	var wg sync.WaitGroup
 	for i, dc := range dcs {
 		wg.Go(func() {
 			defer func() {
 				if r := recover(); r != nil {
-					errs[i] = fmt.Errorf("panic: %v", r)
+					results[i].err = fmt.Errorf("panic: %v", r)
 				}
 			}()
-			errs[i] = d.checkNetworkAddresses(ntw, essentials.TelegramCoreAddresses[dc])
+			results[i].rtt, results[i].err = d.checkNetworkAddresses(ntw, dc, essentials.TelegramCoreAddresses[dc])
 		})
 	}
 	wg.Wait()
 
 	ok := true
 
 	for i, dc := range dcs {
-		if errs[i] == nil {
+		if results[i].err == nil {
 			tplODCConnect.Execute(os.Stdout, map[string]any{ //nolint: errcheck
-				"dc": dc,
+				"dc":  dc,
+				"rtt": results[i].rtt.Round(time.Microsecond),
 			})
 		} else {
 			tplEDCConnect.Execute(os.Stdout, map[string]any{ //nolint: errcheck
 				"dc":    dc,
-				"error": errs[i],
+				"error": results[i].err,
 			})
 			ok = false
 		}
@@ -272,7 +278,7 @@ func (d *Doctor) checkNetwork(ntw mtglib.Network) bool {
 	return ok
 }
 
-func (d *Doctor) checkNetworkAddresses(ntw mtglib.Network, addresses []string) error {
+func (d *Doctor) checkNetworkAddresses(ntw mtglib.Network, dc int, addresses []string) (time.Duration, error) {
 	checkAddresses := []string{}
 
 	switch d.conf.PreferIP.Get("prefer-ip4") {
@@ -303,29 +309,33 @@ func (d *Doctor) checkNetworkAddresses(ntw mtglib.Network, addresses []string) e
 	}
 
 	if len(checkAddresses) == 0 {
-		return fmt.Errorf("no suitable addresses after IP version filtering")
+		return 0, fmt.Errorf("no suitable addresses after IP version filtering")
 	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
-	var (
-		conn net.Conn
-		err  error
-	)
+	var lastErr error
 
 	for _, addr := range checkAddresses {
-		conn, err = ntw.DialContext(ctx, "tcp", addr)
+		conn, err := ntw.DialContext(ctx, "tcp", addr)
 		if err != nil {
+			lastErr = fmt.Errorf("tcp connect to %s: %w", addr, err)
 			continue
 		}
 
+		rtt, err := dcprobe.Probe(ctx, conn, dc)
 		conn.Close() //nolint: errcheck
 
-		return nil
+		if err != nil {
+			lastErr = fmt.Errorf("rpc handshake to %s: %w", addr, err)
+			continue
+		}
+
+		return rtt, nil
 	}
 
-	return err
+	return 0, lastErr
 }
 
 func (d *Doctor) checkFrontingDomain(ntw mtglib.Network) bool {

mtglib/dcprobe/probe.go

@@ -0,0 +1,181 @@
+// Package dcprobe verifies that a TCP endpoint is a real Telegram DC by
+// performing the unauthenticated first step of the MTProto handshake
+// (req_pq_multi -> resPQ) on top of mtg's existing obfuscated2 transport.
+//
+// No auth_key is generated; no long-lived state is introduced. Two TL
+// messages, one round-trip. A generic listener cannot fake the reply
+// because it must echo back our random nonce in resPQ.
+//
+// References:
+//   - https://core.telegram.org/mtproto/auth_key      (handshake step 1)
+//   - https://core.telegram.org/schema/mtproto        (TL schema)
+//   - https://core.telegram.org/mtproto/mtproto-transports#padded-intermediate
+package dcprobe
+
+import (
+	"bytes"
+	"context"
+	"crypto/rand"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"time"
+
+	"github.com/9seconds/mtg/v2/essentials"
+	"github.com/9seconds/mtg/v2/mtglib/obfuscation"
+)
+
+// MTProto wire constants (https://core.telegram.org/schema/mtproto).
+//
+//	req_pq_multi#be7e8ef1 nonce:int128 = ResPQ;
+//	resPQ#05162463 nonce:int128 server_nonce:int128 pq:string
+//	    server_public_key_fingerprints:Vector<long> = ResPQ;
+const (
+	ctorReqPQMulti uint32 = 0xbe7e8ef1
+	ctorResPQ      uint32 = 0x05162463
+
+	// Minimum legal resPQ frame: 20-byte unencrypted-message envelope +
+	// 4-byte ctor + 16-byte nonce echo. Anything below cannot be a resPQ.
+	minResPQFrame = 20 + 4 + 16
+	// Upper bound: real resPQ replies are ~84 bytes (envelope + ~64-byte
+	// payload). 256 is comfortable headroom; anything beyond is hostile or
+	// not Telegram.
+	maxResPQFrame = 256
+)
+
+// Probe sends req_pq_multi over an obfuscated2 + padded-intermediate transport
+// and verifies that the peer replies with a matching resPQ.
+//
+// conn must be a freshly opened reliable byte stream (typically TCP) to a
+// Telegram DC, but a SOCKS/proxy-wrapped net.Conn works just as well — Probe
+// adapts whatever it gets to the half-close interface mtg's obfuscator
+// requires. Probe does NOT close conn — the caller does. dc is the DC number
+// (1..5) that gets baked into the obfuscated2 handshake frame.
+//
+// The returned duration is the round-trip from "first byte sent after the
+// obfs handshake" to "resPQ frame fully read".
+func Probe(ctx context.Context, conn net.Conn, dc int) (time.Duration, error) {
+	if deadline, ok := ctx.Deadline(); ok {
+		_ = conn.SetDeadline(deadline)
+		defer func() { _ = conn.SetDeadline(time.Time{}) }()
+	}
+
+	// Honour ctx cancellation as well as its deadline: a parent ctx that is
+	// canceled (without an earlier deadline expiring) would otherwise let
+	// Probe block on an in-flight Read until the deadline. Forcing the
+	// deadline to "now" makes the next syscall return an i/o timeout error
+	// that Probe wraps and surfaces.
+	stop := context.AfterFunc(ctx, func() {
+		_ = conn.SetDeadline(time.Now())
+	})
+	defer stop()
+
+	// 1. obfuscated2 handshake. Empty Secret = no MTProxy secret mixing,
+	// which is how mtg itself talks to a DC (see mtglib/proxy.go).
+	obfsConn, err := obfuscation.Obfuscator{}.SendHandshake(adaptConn(conn), dc)
+	if err != nil {
+		return 0, fmt.Errorf("obfuscated2 handshake: %w", err)
+	}
+
+	// 2. build req_pq_multi TL payload: 4-byte LE constructor + 16-byte nonce.
+	var nonce [16]byte
+	if _, err := rand.Read(nonce[:]); err != nil {
+		return 0, fmt.Errorf("read nonce: %w", err)
+	}
+	tlBody := make([]byte, 4+16)
+	binary.LittleEndian.PutUint32(tlBody[:4], ctorReqPQMulti)
+	copy(tlBody[4:], nonce[:])
+
+	// 3. wrap in an MTProto unencrypted message envelope (per
+	// https://core.telegram.org/mtproto/description#unencrypted-message):
+	//   auth_key_id:long(=0) | message_id:long | message_data_length:int | message_data:bytes
+	// Without this envelope the DC silently drops the connection.
+	msg := make([]byte, 8+8+4+len(tlBody))
+	// auth_key_id = 0 (already zeroed by make)
+	binary.LittleEndian.PutUint64(msg[8:16], generateMessageID())
+	binary.LittleEndian.PutUint32(msg[16:20], uint32(len(tlBody)))
+	copy(msg[20:], tlBody)
+
+	// 4. wrap in a padded-intermediate frame: length(LE) + msg.
+	// Padding is allowed [0..15] but not required when len(msg) % 4 == 0.
+	frame := make([]byte, 4+len(msg))
+	binary.LittleEndian.PutUint32(frame[:4], uint32(len(msg)))
+	copy(frame[4:], msg)
+
+	start := time.Now()
+	if _, err := obfsConn.Write(frame); err != nil {
+		return 0, fmt.Errorf("write req_pq_multi: %w", err)
+	}
+
+	// 5. read padded-intermediate reply: length, then that many bytes.
+	// The reply is itself an MTProto unencrypted message (same envelope as
+	// what we sent), so we must skip 20 bytes to get to the resPQ TL.
+	var lenBuf [4]byte
+	if _, err := io.ReadFull(obfsConn, lenBuf[:]); err != nil {
+		return 0, fmt.Errorf("read frame length: %w", err)
+	}
+	respLen := binary.LittleEndian.Uint32(lenBuf[:])
+	if respLen < minResPQFrame {
+		return 0, fmt.Errorf("%w: resPQ frame too short (%d bytes)", ErrNotTelegram, respLen)
+	}
+	if respLen > maxResPQFrame {
+		return 0, fmt.Errorf("%w: resPQ frame too large (%d bytes, max %d)", ErrNotTelegram, respLen, maxResPQFrame)
+	}
+	resp := make([]byte, respLen)
+	if _, err := io.ReadFull(obfsConn, resp); err != nil {
+		return 0, fmt.Errorf("read resPQ frame: %w", err)
+	}
+	rtt := time.Since(start)
+
+	// 6. unwrap the MTProto envelope: skip auth_key_id(8) + message_id(8) +
+	// message_data_length(4) = 20 bytes.
+	tlResp := resp[20:]
+
+	// 7. verify constructor and nonce echo. We deliberately do not parse
+	// server_nonce, pq, or fingerprints — they are not needed to prove
+	// the peer can speak MTProto.
+	if got := binary.LittleEndian.Uint32(tlResp[:4]); got != ctorResPQ {
+		return rtt, fmt.Errorf("%w: got constructor 0x%08x, want resPQ 0x%08x", ErrNotTelegram, got, ctorResPQ)
+	}
+	if !bytes.Equal(tlResp[4:4+16], nonce[:]) {
+		return rtt, fmt.Errorf("%w: nonce echo mismatch", ErrNotTelegram)
+	}
+
+	return rtt, nil
+}
+
+// generateMessageID returns an MTProto message_id roughly synchronized with
+// server time, with the lower 2 bits cleared (client-to-server requests).
+// See https://core.telegram.org/mtproto/description#message-identifier-msg-id.
+func generateMessageID() uint64 {
+	nano := uint64(time.Now().UnixNano())
+	sec := nano / 1_000_000_000
+	nsInSec := nano % 1_000_000_000
+	subsec := (nsInSec << 32) / 1_000_000_000
+	id := (sec << 32) | subsec
+	return id &^ 3
+}
+
+// ErrNotTelegram is returned (wrapped) when the peer's reply is not a
+// well-formed resPQ matching our nonce. Use errors.Is to distinguish
+// "the TCP connection was OK but the peer is not a Telegram DC" from
+// transport errors.
+var ErrNotTelegram = errors.New("peer did not respond with a matching resPQ")
+
+// adaptConn returns conn as essentials.Conn if it already satisfies the
+// interface (typically *net.TCPConn), otherwise wraps it with no-op
+// CloseRead/CloseWrite. mtg's obfuscator only ever calls Read/Write/Close,
+// so the no-ops are safe.
+func adaptConn(conn net.Conn) essentials.Conn {
+	if ec, ok := conn.(essentials.Conn); ok {
+		return ec
+	}
+	return halfCloseShim{Conn: conn}
+}
+
+type halfCloseShim struct{ net.Conn }
+
+func (halfCloseShim) CloseRead() error  { return nil }
+func (halfCloseShim) CloseWrite() error { return nil }