Enhance MTProxy checker functionality by introducing a new -probe flag for selecting between fast and deep probing modes. Update README and Docker documentation to reflect this change, including details on timeout handling and exit codes for improved clarity.

internal/checker/checker.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"time"
 
@@ -17,25 +18,29 @@ import (
 // ErrProxyClosed indicates the peer closed the TCP connection during the check (Telethon #1134 style).
 var ErrProxyClosed = errors.New("mtproxy closed connection after initial payload")
 
-// Check runs Fake-TLS/dd handshake, MTProxy init, then a minimal MTProto req_pq and expects resPQ from Telegram DC (same idea as the mobile client path).
-func Check(ctx context.Context, host string, port int, parsed *secret.Parsed, dcID int16) error {
+// ErrNoDataAfterHeader is returned in ProbeFast when the proxy sends nothing after the init payload within the wait window.
+var ErrNoDataAfterHeader = errors.New("no data from proxy after mtproxy header (timeout)")
+
+// Check runs Fake-TLS/dd handshake and MTProxy init; further steps depend on opts.Probe (see Options). opts nil => ProbeFast.
+func Check(ctx context.Context, host string, port int, parsed *secret.Parsed, dcID int16, opts *Options) error {
 	conn, err := dialTCP(ctx, host, port)
 	if err != nil {
 		return fmt.Errorf("tcp dial: %w", err)
 	}
 	defer conn.Close()
 
+	o := effectiveOpts(opts)
 	switch parsed.Kind {
 	case secret.KindEE:
-		return checkEE(ctx, conn, parsed, dcID)
+		return checkEE(ctx, conn, parsed, dcID, o)
 	case secret.KindDD:
-		return checkDD(ctx, conn, parsed, dcID)
+		return checkDD(ctx, conn, parsed, dcID, o)
 	default:
 		return fmt.Errorf("unknown secret kind")
 	}
 }
 
-func checkEE(ctx context.Context, conn net.Conn, p *secret.Parsed, dcID int16) error {
+func checkEE(ctx context.Context, conn net.Conn, p *secret.Parsed, dcID int16, o *Options) error {
 	// ee-секрет хранит домен как «сырой» хвост (часто 0xd0 + ASCII hostname). В TLS SNI нужен только hostname,
 	// как в официальном клиенте Telegram — иначе прокси сбрасывает соединение до ServerHello.
 	sni := faketls.SNIDomain(p.Domain)
@@ -64,6 +69,9 @@ func checkEE(ctx context.Context, conn net.Conn, p *secret.Parsed, dcID int16) e
 	if err := faketls.WriteTLSApplicationData(conn, hdr); err != nil {
 		return fmt.Errorf("write mtproxy header: %w", err)
 	}
+	if o.Probe == ProbeFast {
+		return waitPostPayload(conn)
+	}
 	br := bufio.NewReader(conn)
 	if err := tgquick.DrainPostInitEE(ctx, br, conn, dec, 2*time.Second); err != nil {
 		return err
@@ -77,7 +85,7 @@ func checkEE(ctx context.Context, conn net.Conn, p *secret.Parsed, dcID int16) e
 	return nil
 }
 
-func checkDD(ctx context.Context, conn net.Conn, p *secret.Parsed, dcID int16) error {
+func checkDD(ctx context.Context, conn net.Conn, p *secret.Parsed, dcID int16, o *Options) error {
 	hdr, enc, dec, err := mtproxy.InitHeader(p.Key, dcID)
 	if err != nil {
 		return fmt.Errorf("mtproxy header: %w", err)
@@ -85,6 +93,9 @@ func checkDD(ctx context.Context, conn net.Conn, p *secret.Parsed, dcID int16) e
 	if _, err := conn.Write(hdr); err != nil {
 		return fmt.Errorf("write mtproxy header: %w", err)
 	}
+	if o.Probe == ProbeFast {
+		return waitPostPayload(conn)
+	}
 	if err := tgquick.VerifyResPQ(ctx, conn, enc, dec, nil); err != nil {
 		if errors.Is(err, tgquick.ErrPeerClosed) {
 			return fmt.Errorf("%w", ErrProxyClosed)
@@ -93,3 +104,30 @@ func checkDD(ctx context.Context, conn net.Conn, p *secret.Parsed, dcID int16) e
 	}
 	return nil
 }
+
+func waitPostPayload(conn net.Conn) error {
+	deadline := time.Now().Add(2 * time.Second)
+	buf := make([]byte, 4096)
+	for time.Now().Before(deadline) {
+		_ = conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
+		n, err := conn.Read(buf)
+		if n > 0 {
+			_ = conn.SetReadDeadline(time.Time{})
+			return nil
+		}
+		if err != nil {
+			if errors.Is(err, io.EOF) {
+				_ = conn.SetReadDeadline(time.Time{})
+				return ErrProxyClosed
+			}
+			var ne net.Error
+			if errors.As(err, &ne) && ne.Timeout() {
+				continue
+			}
+			_ = conn.SetReadDeadline(time.Time{})
+			return err
+		}
+	}
+	_ = conn.SetReadDeadline(time.Time{})
+	return ErrNoDataAfterHeader
+}

internal/checker/options.go

@@ -0,0 +1,35 @@
+package checker
+
+import "strings"
+
+// ProbeMode selects how strictly we validate after MTProxy init.
+type ProbeMode int
+
+const (
+	// ProbeFast: handshake + init, then any inbound data within a short window (legacy behavior; fast).
+	ProbeFast ProbeMode = iota
+	// ProbeDeep: MTProto req_pq and expect resPQ from Telegram DC through the tunnel (stricter; slower).
+	ProbeDeep
+)
+
+// Options configures Check. Nil or zero value uses ProbeFast.
+type Options struct {
+	Probe ProbeMode
+}
+
+// ParseProbe maps env/flag strings: "", "fast" -> ProbeFast; "deep" -> ProbeDeep.
+func ParseProbe(s string) ProbeMode {
+	switch strings.ToLower(strings.TrimSpace(s)) {
+	case "deep", "respq", "dc":
+		return ProbeDeep
+	default:
+		return ProbeFast
+	}
+}
+
+func effectiveOpts(opts *Options) *Options {
+	if opts == nil {
+		return &Options{Probe: ProbeFast}
+	}
+	return opts
+}

internal/checkresult/checkresult.go

@@ -16,6 +16,9 @@ func Classify(err error) (exitCode int, message string) {
 	if errors.Is(err, checker.ErrProxyClosed) {
 		return 3, err.Error()
 	}
+	if errors.Is(err, checker.ErrNoDataAfterHeader) {
+		return 1, err.Error()
+	}
 	if errors.Is(err, context.DeadlineExceeded) {
 		return 4, "timeout"
 	}