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.

internal/checker/checker.go

@@ -4,23 +4,18 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"io"
 	"net"
-	"time"
 
 	"mtproxy_checker/internal/faketls"
 	"mtproxy_checker/internal/mtproxy"
 	"mtproxy_checker/internal/secret"
+	"mtproxy_checker/internal/tgquick"
 )
 
-// ErrProxyClosed indicates the MTProxy dropped the TCP connection right after the probe (Telethon #1134 style).
+// ErrProxyClosed indicates the peer closed the TCP connection during the check (Telethon #1134 style).
 var ErrProxyClosed = errors.New("mtproxy closed connection after initial payload")
 
-// ErrNoDataAfterHeader indicates the proxy accepted the probe but sent no application data back within the wait window.
-// Treating silence as success matched some servers but often disagrees with Telegram client "unavailable".
-var ErrNoDataAfterHeader = errors.New("no data from proxy after mtproxy header (timeout)")
-
-// Check runs the handshake probe for the given parsed secret.
+// 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 {
 	conn, err := dialTCP(ctx, host, port)
 	if err != nil {
@@ -30,15 +25,15 @@ func Check(ctx context.Context, host string, port int, parsed *secret.Parsed, dc
 
 	switch parsed.Kind {
 	case secret.KindEE:
-		return checkEE(conn, parsed, dcID)
+		return checkEE(ctx, conn, parsed, dcID)
 	case secret.KindDD:
-		return checkDD(conn, parsed, dcID)
+		return checkDD(ctx, conn, parsed, dcID)
 	default:
 		return fmt.Errorf("unknown secret kind")
 	}
 }
 
-func checkEE(conn net.Conn, p *secret.Parsed, dcID int16) error {
+func checkEE(ctx context.Context, conn net.Conn, p *secret.Parsed, dcID int16) error {
 	// ee-секрет хранит домен как «сырой» хвост (часто 0xd0 + ASCII hostname). В TLS SNI нужен только hostname,
 	// как в официальном клиенте Telegram — иначе прокси сбрасывает соединение до ServerHello.
 	sni := faketls.SNIDomain(p.Domain)
@@ -60,50 +55,35 @@ func checkEE(conn net.Conn, p *secret.Parsed, dcID int16) error {
 		return fmt.Errorf("verify server hello: %w", err)
 	}
 
-	hdr, _, _, err := mtproxy.InitHeader(p.Key, dcID)
+	hdr, enc, dec, err := mtproxy.InitHeader(p.Key, dcID)
 	if err != nil {
 		return fmt.Errorf("mtproxy header: %w", err)
 	}
 	if err := faketls.WriteTLSApplicationData(conn, hdr); err != nil {
 		return fmt.Errorf("write mtproxy header: %w", err)
 	}
-	return waitPostPayload(conn)
+	if err := tgquick.VerifyResPQ(ctx, conn, enc, dec, true); err != nil {
+		if errors.Is(err, tgquick.ErrPeerClosed) {
+			return fmt.Errorf("%w", ErrProxyClosed)
+		}
+		return err
+	}
+	return nil
 }
 
-func checkDD(conn net.Conn, p *secret.Parsed, dcID int16) error {
-	hdr, _, _, err := mtproxy.InitHeader(p.Key, dcID)
+func checkDD(ctx context.Context, conn net.Conn, p *secret.Parsed, dcID int16) error {
+	hdr, enc, dec, err := mtproxy.InitHeader(p.Key, dcID)
 	if err != nil {
 		return fmt.Errorf("mtproxy header: %w", err)
 	}
 	if _, err := conn.Write(hdr); err != nil {
 		return fmt.Errorf("write mtproxy header: %w", err)
 	}
-	return waitPostPayload(conn)
-}
-
-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
+	if err := tgquick.VerifyResPQ(ctx, conn, enc, dec, false); err != nil {
+		if errors.Is(err, tgquick.ErrPeerClosed) {
+			return fmt.Errorf("%w", ErrProxyClosed)
 		}
+		return err
 	}
-	_ = conn.SetReadDeadline(time.Time{})
-	return ErrNoDataAfterHeader
+	return nil
 }