mtproxy_checker/internal/checker/checker.go

package checker

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"syscall"
	"time"

	"mtproxy_checker/internal/faketls"
	"mtproxy_checker/internal/mtproxy"
	"mtproxy_checker/internal/secret"
	"mtproxy_checker/internal/tgquick"
)

// 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 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, o)
	case secret.KindDD:
		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, o *Options) error {
	// ee-секрет хранит домен как «сырой» хвост (часто 0xd0 + ASCII hostname). В TLS SNI нужен только hostname,
	// как в официальном клиенте Telegram — иначе прокси сбрасывает соединение до ServerHello.
	sni := faketls.SNIDomain(p.Domain)
	if len(sni) == 0 {
		sni = p.Domain
	}
	ch, err := faketls.BuildTdesktopClientHello(p.Key, sni)
	if err != nil {
		return fmt.Errorf("fake-tls client hello: %w", err)
	}
	if _, err := conn.Write(ch.Record); err != nil {
		return fmt.Errorf("write client hello: %w", err)
	}
	// ServerHello читаем с сырого conn (ровно длина записи), как в tdesktop Read — без readahead.
	resp, err := faketls.ReadServerHello(conn)
	if err != nil {
		return fmt.Errorf("read server hello: %w", err)
	}
	if err := faketls.VerifyServerHelloTdesktop(resp, p.Key, ch.RandomField); err != nil {
		return fmt.Errorf("verify server hello: %w", err)
	}

	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)
	}
	// После init один bufio на чтение TLS 0x17 (Telethon/telemt: не смешивать conn.Read и буфер).
	br := bufio.NewReader(conn)
	if o.Probe == ProbeFast {
		return waitPostPayload(br, conn)
	}
	if err := tgquick.DrainPostInitEE(ctx, br, conn, dec, 5*time.Second); err != nil {
		return err
	}
	strict := o.Probe == ProbeDeepStrict
	if err := tgquick.VerifyResPQ(ctx, conn, enc, dec, br, strict); err != nil {
		if errors.Is(err, tgquick.ErrPeerClosed) {
			return fmt.Errorf("%w", ErrProxyClosed)
		}
		return err
	}
	return nil
}

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)
	}
	if _, err := conn.Write(hdr); err != nil {
		return fmt.Errorf("write mtproxy header: %w", err)
	}
	if o.Probe == ProbeFast {
		return waitPostPayload(conn, conn)
	}
	strict := o.Probe == ProbeDeepStrict
	if err := tgquick.VerifyResPQ(ctx, conn, enc, dec, nil, strict); err != nil {
		if errors.Is(err, tgquick.ErrPeerClosed) {
			return fmt.Errorf("%w", ErrProxyClosed)
		}
		return err
	}
	return nil
}

// isPeerClosedReadErr reports RST/half-close style errors from Read (not timeouts).
func isPeerClosedReadErr(err error) bool {
	if err == nil {
		return false
	}
	var ne net.Error
	if errors.As(err, &ne) && ne.Timeout() {
		return false
	}
	return errors.Is(err, syscall.ECONNRESET) ||
		errors.Is(err, syscall.ECONNABORTED) ||
		errors.Is(err, syscall.EPIPE) ||
		errors.Is(err, syscall.ENOTCONN)
}

// waitPostPayload после init повторяет логику Telethon TcpMTProxy._connect (#1134):
// ждём появления данных или таймаут, затем считаем успехом только если соединение не закрыто сразу после payload.
// Наличие входящих байтов не обязательно — часть прокси молчит до первого MTProto от клиента.
// После окна idle дополнительно читаем короткими интервалами: часть прокси шлёт FIN/RST сразу после init,
// но не в первых 200ms цикла — без этого fast давал бы ложный OK.
func waitPostPayload(r io.Reader, c net.Conn) error {
	deadline := time.Now().Add(2 * time.Second)
	buf := make([]byte, 4096)
	for time.Now().Before(deadline) {
		_ = c.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
		n, err := r.Read(buf)
		if n > 0 {
			_ = c.SetReadDeadline(time.Time{})
			return nil
		}
		if err != nil {
			if errors.Is(err, io.EOF) {
				_ = c.SetReadDeadline(time.Time{})
				return ErrProxyClosed
			}
			if isPeerClosedReadErr(err) {
				_ = c.SetReadDeadline(time.Time{})
				return ErrProxyClosed
			}
			var ne net.Error
			if errors.As(err, &ne) && ne.Timeout() {
				continue
			}
			_ = c.SetReadDeadline(time.Time{})
			return err
		}
	}
	// Отложенное закрытие со стороны прокси (после того как «молчали» в основном окне).
	const probeEach = 400 * time.Millisecond
	const probeRounds = 4
	for probes := 0; probes < probeRounds; probes++ {
		_ = c.SetReadDeadline(time.Now().Add(probeEach))
		n, err := r.Read(buf)
		if n > 0 {
			_ = c.SetReadDeadline(time.Time{})
			return nil
		}
		if err != nil {
			if errors.Is(err, io.EOF) {
				_ = c.SetReadDeadline(time.Time{})
				return ErrProxyClosed
			}
			if isPeerClosedReadErr(err) {
				_ = c.SetReadDeadline(time.Time{})
				return ErrProxyClosed
			}
			var ne net.Error
			if errors.As(err, &ne) && ne.Timeout() {
				continue
			}
			_ = c.SetReadDeadline(time.Time{})
			return err
		}
	}
	_ = c.SetReadDeadline(time.Time{})
	return nil
}