fix: sync same-port transport detection

internal/mtprotoedge/transport_detection.go

@@ -0,0 +1,234 @@
+package mtprotoedge
+
+import (
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+
+	"github.com/iamxvbaba/td/mtproxy/obfuscated2"
+	"github.com/iamxvbaba/td/proto/codec"
+	"github.com/iamxvbaba/td/transport"
+)
+
+var errInvalidTCPTransportPrefix = errors.New("invalid MTProto TCP transport prefix")
+var errInvalidObfuscatedProtocol = errors.New("invalid obfuscated MTProto protocol tag")
+
+type tcpWireMode uint8
+
+const (
+	tcpWireModeUnknown tcpWireMode = iota
+	tcpWireModePlain
+	tcpWireModeObfuscated
+)
+
+type detectedTCPCodec uint8
+
+const (
+	detectedTCPCodecUnknown detectedTCPCodec = iota
+	detectedTCPCodecAbridged
+	detectedTCPCodecIntermediate
+	detectedTCPCodecPaddedIntermediate
+	detectedTCPCodecFull
+)
+
+type tcpTransportProbe struct {
+	raw       net.Conn
+	prefix    [8]byte
+	prefixLen int
+	mode      tcpWireMode
+	codec     detectedTCPCodec
+}
+
+// detectTCPTransport reads at most the first eight bytes once per physical TCP
+// connection. Telegram's obfuscated2 nonce generation explicitly excludes all
+// plaintext codec tags and requires bytes 4..8 to be non-zero, while the first
+// Full frame has transport sequence number zero. Those disjoint invariants let
+// one public port admit both wire modes without probabilistic guessing.
+func detectTCPTransport(raw net.Conn) (tcpTransportProbe, error) {
+	probe := tcpTransportProbe{raw: raw}
+	if _, err := io.ReadFull(raw, probe.prefix[:1]); err != nil {
+		return probe, fmt.Errorf("read first transport byte: %w", err)
+	}
+	probe.prefixLen = 1
+	if probe.prefix[0] == codec.AbridgedClientStart[0] {
+		probe.mode = tcpWireModePlain
+		probe.codec = detectedTCPCodecAbridged
+		return probe, nil
+	}
+
+	if _, err := io.ReadFull(raw, probe.prefix[1:4]); err != nil {
+		return probe, fmt.Errorf("read transport prefix: %w", err)
+	}
+	probe.prefixLen = 4
+	var firstFour [4]byte
+	copy(firstFour[:], probe.prefix[:4])
+	switch firstFour {
+	case codec.IntermediateClientStart:
+		probe.mode = tcpWireModePlain
+		probe.codec = detectedTCPCodecIntermediate
+		return probe, nil
+	case codec.PaddedIntermediateClientStart:
+		probe.mode = tcpWireModePlain
+		probe.codec = detectedTCPCodecPaddedIntermediate
+		return probe, nil
+	}
+
+	first := binary.LittleEndian.Uint32(firstFour[:])
+	if isHTTPHeaderPrefix(firstFour) || first == 0x02010316 {
+		return probe, fmt.Errorf("%w: reserved prefix %x", errInvalidTCPTransportPrefix, firstFour)
+	}
+
+	if _, err := io.ReadFull(raw, probe.prefix[4:8]); err != nil {
+		return probe, fmt.Errorf("read transport discriminator: %w", err)
+	}
+	probe.prefixLen = 8
+	if binary.LittleEndian.Uint32(probe.prefix[4:8]) == 0 {
+		length := binary.LittleEndian.Uint32(probe.prefix[:4])
+		if length < 3*4 || length > maxTransportMessageSize || length%4 != 0 {
+			return probe, fmt.Errorf("%w: invalid full header length %d", errInvalidTCPTransportPrefix, length)
+		}
+		probe.mode = tcpWireModePlain
+		probe.codec = detectedTCPCodecFull
+		return probe, nil
+	}
+
+	probe.mode = tcpWireModeObfuscated
+	return probe, nil
+}
+
+func (p tcpTransportProbe) replayConn() net.Conn {
+	return newReplayNetConn(p.raw, p.prefix[:p.prefixLen])
+}
+
+func (p tcpTransportProbe) plainFrameConn() net.Conn {
+	if p.codec == detectedTCPCodecFull {
+		// Full has no standalone codec tag: the first eight bytes are already the
+		// length and sequence number of its first frame and must be replayed.
+		return p.replayConn()
+	}
+	// Abridged/intermediate/padded-intermediate tags are client-only connection
+	// headers. They were consumed by detection and are not part of the first frame.
+	return p.raw
+}
+
+type replayNetConn struct {
+	net.Conn
+	prefix [8]byte
+	n      uint8
+	offset uint8
+}
+
+func newReplayNetConn(conn net.Conn, prefix []byte) *replayNetConn {
+	if len(prefix) > 8 {
+		panic("mtprotoedge: replay prefix exceeds fixed transport discriminator")
+	}
+	replayed := &replayNetConn{Conn: conn, n: uint8(len(prefix))}
+	copy(replayed.prefix[:], prefix)
+	return replayed
+}
+
+func (c *replayNetConn) Read(p []byte) (int, error) {
+	if c.offset < c.n {
+		n := copy(p, c.prefix[c.offset:c.n])
+		c.offset += uint8(n)
+		return n, nil
+	}
+	return c.Conn.Read(p)
+}
+
+type serverObfuscatedConn struct {
+	net.Conn
+	rw io.ReadWriter
+}
+
+func (c *serverObfuscatedConn) Read(p []byte) (int, error) {
+	return c.rw.Read(p)
+}
+
+func (c *serverObfuscatedConn) Write(p []byte) (int, error) {
+	return c.rw.Write(p)
+}
+
+func detectedCompatCodec(kind detectedTCPCodec) (transport.Codec, error) {
+	switch kind {
+	case detectedTCPCodecAbridged:
+		return &quickAckAbridgedCodec{}, nil
+	case detectedTCPCodecIntermediate:
+		return &quickAckIntermediateCodec{}, nil
+	case detectedTCPCodecPaddedIntermediate:
+		return &quickAckPaddedIntermediateCodec{}, nil
+	case detectedTCPCodecFull:
+		return transport.Full.Codec(), nil
+	default:
+		return nil, fmt.Errorf("%w: unknown detected codec %d", errInvalidTCPTransportPrefix, kind)
+	}
+}
+
+func detectObfuscatedProtocol(tag [4]byte) (detectedTCPCodec, int, error) {
+	switch tag {
+	case (codec.Abridged{}).ObfuscatedTag():
+		return detectedTCPCodecAbridged, 1, nil
+	case codec.IntermediateClientStart:
+		return detectedTCPCodecIntermediate, 4, nil
+	case codec.PaddedIntermediateClientStart:
+		return detectedTCPCodecPaddedIntermediate, 4, nil
+	default:
+		return detectedTCPCodecUnknown, 0, fmt.Errorf("%w: %x", errInvalidObfuscatedProtocol, tag)
+	}
+}
+
+func (s *Server) promoteMixedTCP(raw net.Conn) (transport.Conn, string, error) {
+	probe, err := detectTCPTransport(raw)
+	if err != nil {
+		return nil, "tcp_auto", err
+	}
+	if probe.mode == tcpWireModePlain {
+		if s.codec != nil {
+			conn, err := newCompatTransportConn(s.codec, probe.replayConn(), s.frameBudget)
+			return conn, "tcp", err
+		}
+		connCodec, err := detectedCompatCodec(probe.codec)
+		if err != nil {
+			return nil, "tcp", err
+		}
+		conn, err := newCompatTransportConnWithCodec(
+			probe.plainFrameConn(),
+			connCodec,
+			s.frameBudget,
+			false,
+		)
+		return conn, "tcp", err
+	}
+
+	replayed := probe.replayConn()
+	rw, metadata, err := obfuscated2.Accept(replayed, nil)
+	if err != nil {
+		return nil, "obfuscated_tcp", fmt.Errorf("accept obfuscated2: %w", err)
+	}
+	obfuscated := &serverObfuscatedConn{Conn: replayed, rw: rw}
+	detected, tagLen, err := detectObfuscatedProtocol(metadata.Protocol)
+	if err != nil {
+		return nil, "obfuscated_tcp", err
+	}
+	if s.codec != nil {
+		conn, err := newCompatTransportConn(
+			s.codec,
+			newReplayNetConn(obfuscated, metadata.Protocol[:tagLen]),
+			s.frameBudget,
+		)
+		return conn, "obfuscated_tcp", err
+	}
+	connCodec, err := detectedCompatCodec(detected)
+	if err != nil {
+		return nil, "obfuscated_tcp", err
+	}
+	conn, err := newCompatTransportConnWithCodec(
+		obfuscated,
+		connCodec,
+		s.frameBudget,
+		false,
+	)
+	return conn, "obfuscated_tcp", err
+}