Refactor checker logic to integrate context handling in MTProxy checks. Update error handling to utilize tgquick for response verification, replacing previous waitPostPayload function. Simplify error messages for closed connections.

2026-04-11 01:44:53 +07:00
parent e3049028cd
commit 7d13c599c4
5 changed files with 287 additions and 46 deletions
@@ -0,0 +1,225 @@
+// Package tgquick performs a minimal MTProto step (req_pq → resPQ) through an established MTProxy tunnel,
+// matching Telethon's RandomizedIntermediate + MTProxyIO framing.
+package tgquick
+
+import (
+	"bufio"
+	"context"
+	"crypto/cipher"
+	crand "crypto/rand"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"time"
+
+	"mtproxy_checker/internal/faketls"
+)
+
+const (
+	tlReqPQ = 0xd712e4be
+	tlResPQ = 0x05162463
+	maxRI   = 1 << 20
+)
+
+// ErrNoResPQ means no valid resPQ was received from Telegram through the tunnel.
+var ErrNoResPQ = errors.New("no resPQ from telegram through proxy (tunnel may be broken)")
+
+// ErrPeerClosed is returned when the remote side closes TCP during the MTProto probe.
+var ErrPeerClosed = errors.New("connection closed by peer during mtproto probe")
+
+func newMessageID() int64 {
+	t := time.Now().UnixNano()
+	return (t / 4) &^ 3
+}
+
+func buildReqPQ() ([]byte, error) {
+	nonce := make([]byte, 16)
+	if _, err := io.ReadFull(crand.Reader, nonce); err != nil {
+		return nil, err
+	}
+	body := make([]byte, 4+16)
+	binary.LittleEndian.PutUint32(body, tlReqPQ)
+	copy(body[4:], nonce)
+	return wrapUnencrypted(body), nil
+}
+
+func wrapUnencrypted(msgData []byte) []byte {
+	out := make([]byte, 8+8+4+len(msgData))
+	// auth_key_id = 0
+	binary.LittleEndian.PutUint64(out[8:], uint64(newMessageID()))
+	binary.LittleEndian.PutUint32(out[16:], uint32(len(msgData)))
+	copy(out[20:], msgData)
+	return out
+}
+
+// randomizedIntermediateEncode packs payload with 0–3 random padding bytes (Telethon RandomizedIntermediatePacketCodec).
+func randomizedIntermediateEncode(payload []byte) ([]byte, error) {
+	var rnd [1]byte
+	if _, err := crand.Read(rnd[:]); err != nil {
+		return nil, err
+	}
+	pad := int(rnd[0] % 4)
+	padding := make([]byte, pad)
+	if _, err := crand.Read(padding); err != nil {
+		return nil, err
+	}
+	body := append(append([]byte{}, payload...), padding...)
+	out := make([]byte, 4+len(body))
+	binary.LittleEndian.PutUint32(out, uint32(len(body)))
+	copy(out[4:], body)
+	return out, nil
+}
+
+func popRI(acc *[]byte) ([]byte, bool) {
+	a := *acc
+	if len(a) < 4 {
+		return nil, false
+	}
+	L := int(binary.LittleEndian.Uint32(a[0:4]))
+	if L < 0 || L > maxRI || len(a) < 4+L {
+		return nil, false
+	}
+	frame := a[4 : 4+L]
+	*acc = append([]byte(nil), a[4+L:]...)
+	pad := len(frame) % 4
+	if pad > 0 {
+		frame = frame[:len(frame)-pad]
+	}
+	return frame, true
+}
+
+func isResPQ(mtInner []byte) bool {
+	if len(mtInner) < 20 {
+		return false
+	}
+	if binary.LittleEndian.Uint64(mtInner[0:8]) != 0 {
+		return false
+	}
+	ml := int(binary.LittleEndian.Uint32(mtInner[16:20]))
+	if ml < 4 || ml > maxRI || 20+ml > len(mtInner) {
+		return false
+	}
+	body := mtInner[20 : 20+ml]
+	return binary.LittleEndian.Uint32(body[0:4]) == tlResPQ
+}
+
+// VerifyResPQ sends req_pq and waits for resPQ on the MTProxy byte stream (decrypt with dec, encrypt with enc).
+// If eeTLS is true, each MTProxy chunk is wrapped in TLS 1.2 application records (0x17) as in checkEE.
+func VerifyResPQ(ctx context.Context, conn net.Conn, enc, dec cipher.Stream, eeTLS bool) error {
+	plain, err := buildReqPQ()
+	if err != nil {
+		return fmt.Errorf("build req_pq: %w", err)
+	}
+	ri, err := randomizedIntermediateEncode(plain)
+	if err != nil {
+		return fmt.Errorf("randomized intermediate: %w", err)
+	}
+	wire := append([]byte(nil), ri...)
+	enc.XORKeyStream(wire, wire)
+	if eeTLS {
+		if err := faketls.WriteTLSApplicationData(conn, wire); err != nil {
+			return fmt.Errorf("write tls app (req_pq): %w", err)
+		}
+	} else {
+		if _, err := conn.Write(wire); err != nil {
+			return fmt.Errorf("write req_pq: %w", err)
+		}
+	}
+
+	deadline, ok := ctx.Deadline()
+	if !ok {
+		deadline = time.Now().Add(15 * time.Second)
+	}
+
+	var acc []byte
+	var br *bufio.Reader
+	if eeTLS {
+		br = bufio.NewReader(conn)
+	}
+
+	const maxAcc = 256 << 10
+	for time.Now().Before(deadline) {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+		chunkDeadline := time.Now().Add(3 * time.Second)
+		if chunkDeadline.After(deadline) {
+			chunkDeadline = deadline
+		}
+		_ = conn.SetReadDeadline(chunkDeadline)
+		var chunk []byte
+		var rerr error
+		if eeTLS {
+			chunk, rerr = readOneTLSApplicationRecord(br)
+		} else {
+			chunk, rerr = readAtMost(conn, 65536)
+		}
+		if rerr != nil {
+			if ne, ok := rerr.(net.Error); ok && ne.Timeout() {
+				if time.Now().Before(deadline) {
+					continue
+				}
+				break
+			}
+			if errors.Is(rerr, io.EOF) {
+				return fmt.Errorf("%w", ErrPeerClosed)
+			}
+			return fmt.Errorf("read mtproxy stream: %w", rerr)
+		}
+		if len(chunk) == 0 {
+			continue
+		}
+		dec.XORKeyStream(chunk, chunk)
+		acc = append(acc, chunk...)
+		if len(acc) > maxAcc {
+			return fmt.Errorf("mtproxy read buffer overflow")
+		}
+		for {
+			frame, ok := popRI(&acc)
+			if !ok {
+				break
+			}
+			if isResPQ(frame) {
+				_ = conn.SetReadDeadline(time.Time{})
+				return nil
+			}
+		}
+	}
+	_ = conn.SetReadDeadline(time.Time{})
+	return ErrNoResPQ
+}
+
+func readOneTLSApplicationRecord(br *bufio.Reader) ([]byte, error) {
+	h := make([]byte, 5)
+	if _, err := io.ReadFull(br, h); err != nil {
+		return nil, err
+	}
+	if h[0] != 0x17 {
+		return nil, fmt.Errorf("unexpected tls record type 0x%02x", h[0])
+	}
+	n := int(binary.BigEndian.Uint16(h[3:5]))
+	if n <= 0 || n > 1<<20 {
+		return nil, fmt.Errorf("invalid tls app length %d", n)
+	}
+	p := make([]byte, n)
+	if _, err := io.ReadFull(br, p); err != nil {
+		return nil, err
+	}
+	return p, nil
+}
+
+func readAtMost(conn net.Conn, max int) ([]byte, error) {
+	buf := make([]byte, max)
+	n, err := conn.Read(buf)
+	if n > 0 {
+		if err == io.EOF {
+			err = nil
+		}
+		return buf[:n], err
+	}
+	return nil, err
+}
@@ -0,0 +1,36 @@
+package tgquick
+
+import (
+	"encoding/binary"
+	"testing"
+)
+
+func TestPopRI_RoundTrip(t *testing.T) {
+	payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
+	enc, err := randomizedIntermediateEncode(payload)
+	if err != nil {
+		t.Fatal(err)
+	}
+	var acc []byte
+	acc = append(acc, enc...)
+	got, ok := popRI(&acc)
+	if !ok || len(acc) != 0 {
+		t.Fatalf("popRI: ok=%v rest=%d", ok, len(acc))
+	}
+	if string(got) != string(payload) {
+		t.Fatalf("payload mismatch: %v vs %v", got, payload)
+	}
+}
+
+func TestIsResPQ(t *testing.T) {
+	nonce := make([]byte, 16)
+	// resPQ: constructor + nonce(16) + server_nonce(16) + pq bytes + Vector<long>
+	msgData := make([]byte, 4+16+16+4+4)
+	binary.LittleEndian.PutUint32(msgData, tlResPQ)
+	copy(msgData[4:20], nonce)
+	copy(msgData[20:36], nonce)
+	// pq empty, fingerprints empty
+	if !isResPQ(wrapUnencrypted(msgData)) {
+		t.Fatal("expected resPQ")
+	}
+}