internal/mtprotoedge/transport_detection_test.go
Rama main · réplica de lectura del repositorio original.
package mtprotoedge
import (
"context"
"crypto/rand"
"errors"
"io"
"net"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/mtproxy"
"github.com/iamxvbaba/td/mtproxy/obfuscator"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/transport"
)
func TestDetectTCPTransportFragmentedPrefixes(t *testing.T) {
tests := []struct {
name string
prefix []byte
mode tcpWireMode
codec detectedTCPCodec
}{
{name: "abridged", prefix: []byte{0xef}, mode: tcpWireModePlain, codec: detectedTCPCodecAbridged},
{name: "intermediate", prefix: []byte{0xee, 0xee, 0xee, 0xee}, mode: tcpWireModePlain, codec: detectedTCPCodecIntermediate},
{name: "padded", prefix: []byte{0xdd, 0xdd, 0xdd, 0xdd}, mode: tcpWireModePlain, codec: detectedTCPCodecPaddedIntermediate},
{name: "full", prefix: []byte{12, 0, 0, 0, 0, 0, 0, 0}, mode: tcpWireModePlain, codec: detectedTCPCodecFull},
{name: "obfuscated", prefix: []byte{1, 2, 3, 4, 5, 6, 7, 8}, mode: tcpWireModeObfuscated, codec: detectedTCPCodecUnknown},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
server, client := net.Pipe()
defer func() { _ = server.Close() }()
defer func() { _ = client.Close() }()
writeErr := make(chan error, 1)
go func() {
for _, b := range tt.prefix {
if _, err := client.Write([]byte{b}); err != nil {
writeErr <- err
return
}
}
writeErr <- nil
}()
probe, err := detectTCPTransport(server)
if err != nil {
t.Fatalf("detect: %v", err)
}
if probe.mode != tt.mode || probe.codec != tt.codec || probe.prefixLen != len(tt.prefix) {
t.Fatalf("probe = mode:%d codec:%d prefix:%d, want %d/%d/%d",
probe.mode, probe.codec, probe.prefixLen, tt.mode, tt.codec, len(tt.prefix))
}
if err := <-writeErr; err != nil {
t.Fatalf("write prefix: %v", err)
}
})
}
}
func TestDetectTCPTransportRejectsReservedAndInvalidFullPrefixes(t *testing.T) {
tests := []struct {
name string
prefix []byte
}{
{name: "http", prefix: []byte("GET ")},
{name: "reserved", prefix: []byte{0x16, 0x03, 0x01, 0x02}},
{name: "full_too_short", prefix: []byte{8, 0, 0, 0, 0, 0, 0, 0}},
{name: "full_unaligned", prefix: []byte{13, 0, 0, 0, 0, 0, 0, 0}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
server, client := net.Pipe()
defer func() { _ = server.Close() }()
defer func() { _ = client.Close() }()
go func() {
_, _ = client.Write(tt.prefix)
}()
_, err := detectTCPTransport(server)
if !errors.Is(err, errInvalidTCPTransportPrefix) {
t.Fatalf("detect error = %v, want invalid prefix", err)
}
})
}
}
func TestServerMixedTCPAcceptsEveryPlainCodec(t *testing.T) {
tests := []struct {
name string
protocol transport.Protocol
}{
{name: "abridged", protocol: transport.Abridged},
{name: "intermediate", protocol: transport.Intermediate},
{name: "padded_intermediate", protocol: transport.PaddedIntermediate},
{name: "full", protocol: transport.Full},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
addr, frames := startTransportFrameServer(t, Options{
Logger: zaptest.NewLogger(t),
ObfuscatedTCP: true,
})
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial: %v", err)
}
conn, err := tt.protocol.Handshake(raw)
if err != nil {
_ = raw.Close()
t.Fatalf("transport handshake: %v", err)
}
t.Cleanup(func() { _ = conn.Close() })
var payload bin.Buffer
payload.PutInt32(0x12345678)
payload.PutInt32(0x0badf00d)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := conn.Send(ctx, &payload); err != nil {
t.Fatalf("send: %v", err)
}
select {
case n := <-frames:
if n < payload.Len() || n > payload.Len()+15 {
t.Fatalf("frame len = %d, want payload %d plus at most padded-intermediate padding", n, payload.Len())
}
case <-ctx.Done():
t.Fatal("server did not receive plain frame in mixed mode")
}
})
}
}
func TestServerMixedTCPAcceptsObfuscatedCodecs(t *testing.T) {
tests := []struct {
name string
tag [4]byte
newCodec func() transport.Codec
}{
{
name: "abridged",
tag: (codec.Abridged{}).ObfuscatedTag(),
newCodec: func() transport.Codec {
return codec.NoHeader{Codec: codec.Abridged{}}
},
},
{
name: "intermediate",
tag: codec.IntermediateClientStart,
newCodec: func() transport.Codec {
return codec.NoHeader{Codec: codec.Intermediate{}}
},
},
{
name: "padded_intermediate",
tag: codec.PaddedIntermediateClientStart,
newCodec: func() transport.Codec {
return codec.NoHeader{Codec: codec.PaddedIntermediate{}}
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
addr, frames := startTransportFrameServer(t, Options{
Logger: zaptest.NewLogger(t),
ObfuscatedTCP: true,
})
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial: %v", err)
}
obfuscated := obfuscator.Obfuscated2(rand.Reader, raw)
if err := obfuscated.Handshake(tt.tag, 2, mtproxy.Secret{}); err != nil {
_ = raw.Close()
t.Fatalf("obfuscated handshake: %v", err)
}
conn, err := transport.NewProtocol(tt.newCodec).Handshake(obfuscated)
if err != nil {
_ = raw.Close()
t.Fatalf("transport handshake: %v", err)
}
t.Cleanup(func() { _ = conn.Close() })
var payload bin.Buffer
payload.PutInt32(0x12345678)
payload.PutInt32(0x0badf00d)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := conn.Send(ctx, &payload); err != nil {
t.Fatalf("send: %v", err)
}
select {
case n := <-frames:
if n < payload.Len() || n > payload.Len()+15 {
t.Fatalf("frame len = %d, want payload %d plus at most padding", n, payload.Len())
}
case <-ctx.Done():
t.Fatal("server did not receive obfuscated frame")
}
})
}
}
func TestServerMixedTCPPlainIntermediateKeyExchange(t *testing.T) {
addr, pub, _ := startTestServer(t, Options{DC: 2, ObfuscatedTCP: true})
conn, _, _ := dialHandshake(t, addr, 2, pub)
_ = conn.Close()
}
func TestServerMixedTCPResponseUsesDetectedWireMode(t *testing.T) {
plainDial := func(protocol transport.Protocol) func(*testing.T, string) transport.Conn {
return func(t *testing.T, addr string) transport.Conn {
t.Helper()
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial: %v", err)
}
conn, err := protocol.Handshake(raw)
if err != nil {
_ = raw.Close()
t.Fatalf("transport handshake: %v", err)
}
return conn
}
}
obfuscatedDial := func(
tag [4]byte,
newCodec func() transport.Codec,
) func(*testing.T, string) transport.Conn {
return func(t *testing.T, addr string) transport.Conn {
t.Helper()
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial: %v", err)
}
obfuscated := obfuscator.Obfuscated2(rand.Reader, raw)
if err := obfuscated.Handshake(tag, 2, mtproxy.Secret{}); err != nil {
_ = raw.Close()
t.Fatalf("obfuscated handshake: %v", err)
}
conn, err := transport.NewProtocol(newCodec).Handshake(obfuscated)
if err != nil {
_ = raw.Close()
t.Fatalf("transport handshake: %v", err)
}
return conn
}
}
tests := []struct {
name string
dial func(t *testing.T, addr string) transport.Conn
}{
{
name: "plain_abridged",
dial: plainDial(transport.Abridged),
},
{
name: "plain_intermediate",
dial: plainDial(transport.Intermediate),
},
{
name: "plain_padded_intermediate",
dial: plainDial(transport.PaddedIntermediate),
},
{
name: "plain_full",
dial: plainDial(transport.Full),
},
{
name: "obfuscated_abridged",
dial: obfuscatedDial(
(codec.Abridged{}).ObfuscatedTag(),
func() transport.Codec {
return codec.NoHeader{Codec: codec.Abridged{}}
},
),
},
{
name: "obfuscated_intermediate",
dial: obfuscatedDial(
codec.IntermediateClientStart,
func() transport.Codec {
return codec.NoHeader{Codec: codec.Intermediate{}}
},
),
},
{
name: "obfuscated_padded_intermediate",
dial: obfuscatedDial(
codec.PaddedIntermediateClientStart,
func() transport.Codec {
return codec.NoHeader{Codec: codec.PaddedIntermediate{}}
},
),
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
addr, _, _ := startTestServer(t, Options{DC: 2, ObfuscatedTCP: true})
conn := tt.dial(t, addr)
t.Cleanup(func() { _ = conn.Close() })
var payload bin.Buffer
payload.PutLong(0x1020304050607080) // deliberately unknown auth_key_id
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := conn.Send(ctx, &payload); err != nil {
t.Fatalf("send: %v", err)
}
var response bin.Buffer
err := conn.Recv(ctx, &response)
var protocolErr *codec.ProtocolErr
if !errors.As(err, &protocolErr) || protocolErr.Code != codec.CodeAuthKeyNotFound {
t.Fatalf("response error = %T %v, want transport -404 in the detected wire mode", err, err)
}
})
}
}
func TestServerMixedTCPRejectsInvalidPrefixesWithoutObfuscationWait(t *testing.T) {
addr, _, _ := startTestServer(t, Options{
ObfuscatedTCP: true,
HandshakeIdleTimeout: 5 * time.Second,
})
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial: %v", err)
}
defer func() { _ = raw.Close() }()
if _, err := raw.Write([]byte{8, 0, 0, 0, 0, 0, 0, 0}); err != nil {
t.Fatalf("write invalid full prefix: %v", err)
}
if err := raw.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
t.Fatalf("set read deadline: %v", err)
}
_, err = raw.Read(make([]byte, 1))
if err == nil {
t.Fatal("invalid transport prefix unexpectedly kept connection open")
}
var netErr net.Error
if errors.As(err, &netErr) && netErr.Timeout() {
t.Fatalf("invalid prefix waited for obfuscation timeout instead of failing fast: %v", err)
}
}
func TestServerMixedTCPRejectsUnknownObfuscatedProtocolTag(t *testing.T) {
addr, _, _ := startTestServer(t, Options{ObfuscatedTCP: true})
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial: %v", err)
}
defer func() { _ = raw.Close() }()
obfuscated := obfuscator.Obfuscated2(rand.Reader, raw)
if err := obfuscated.Handshake([4]byte{1, 2, 3, 4}, 2, mtproxy.Secret{}); err != nil {
t.Fatalf("obfuscated handshake: %v", err)
}
if err := raw.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
t.Fatalf("set read deadline: %v", err)
}
_, err = obfuscated.Read(make([]byte, 1))
if err == nil {
t.Fatal("unknown obfuscated protocol tag unexpectedly kept connection open")
}
var netErr net.Error
if errors.As(err, &netErr) && netErr.Timeout() {
t.Fatalf("unknown obfuscated protocol tag did not fail closed: %v", err)
}
if !errors.Is(err, io.EOF) {
// Windows may report a reset instead of EOF; any non-timeout terminal
// network error is the same fail-closed outcome.
t.Logf("terminal read error after invalid obfuscated tag: %v", err)
}
}
func startTransportFrameServer(t *testing.T, opts Options) (string, <-chan int) {
t.Helper()
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
frames := make(chan int, 1)
srv := New(opts)
srv.onFrame = func(n int) {
select {
case frames <- n:
default:
}
}
ctx, cancel := context.WithCancel(context.Background())
serveErr := make(chan error, 1)
go func() { serveErr <- srv.Serve(ctx, ln) }()
t.Cleanup(func() {
cancel()
select {
case err := <-serveErr:
if err != nil {
t.Errorf("serve: %v", err)
}
case <-time.After(5 * time.Second):
t.Error("server did not stop")
}
})
return ln.Addr().String(), frames
}
func BenchmarkDetectTCPTransport(b *testing.B) {
tests := []struct {
name string
prefix []byte
}{
{name: "plain_abridged", prefix: []byte{0xef}},
{name: "plain_intermediate", prefix: []byte{0xee, 0xee, 0xee, 0xee}},
{name: "plain_full", prefix: []byte{12, 0, 0, 0, 0, 0, 0, 0}},
{name: "obfuscated", prefix: []byte{1, 2, 3, 4, 5, 6, 7, 8}},
}
for _, tt := range tests {
b.Run(tt.name, func(b *testing.B) {
conn := &transportProbeBenchmarkConn{payload: tt.prefix}
b.ReportAllocs()
for i := 0; i < b.N; i++ {
conn.offset = 0
if _, err := detectTCPTransport(conn); err != nil {
b.Fatal(err)
}
}
})
}
}
type transportProbeBenchmarkConn struct {
payload []byte
offset int
}
func (c *transportProbeBenchmarkConn) Read(p []byte) (int, error) {
if c.offset >= len(c.payload) {
return 0, io.EOF
}
n := copy(p, c.payload[c.offset:])
c.offset += n
return n, nil
}
func (*transportProbeBenchmarkConn) Write(p []byte) (int, error) { return len(p), nil }
func (*transportProbeBenchmarkConn) Close() error { return nil }
func (*transportProbeBenchmarkConn) LocalAddr() net.Addr { return nil }
func (*transportProbeBenchmarkConn) RemoteAddr() net.Addr { return nil }
func (*transportProbeBenchmarkConn) SetDeadline(time.Time) error { return nil }
func (*transportProbeBenchmarkConn) SetReadDeadline(time.Time) error {
return nil
}
func (*transportProbeBenchmarkConn) SetWriteDeadline(time.Time) error {
return nil
}