mc-proxy/internal/l7/prefixconn_test.go

package l7

import (
	"io"
	"net"
	"testing"
	"time"
)

func TestPrefixConnRead(t *testing.T) {
	// Create a TCP pair.
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	defer func() { _ = ln.Close() }()

	go func() {
		conn, err := ln.Accept()
		if err != nil {
			return
		}
		defer func() { _ = conn.Close() }()
		_, _ = conn.Write([]byte("WORLD"))
	}()

	conn, err := net.Dial("tcp", ln.Addr().String())
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	defer func() { _ = conn.Close() }()

	pc := NewPrefixConn(conn, []byte("HELLO"))

	// Read all data: should get "HELLOWORLD".
	_ = conn.SetReadDeadline(time.Now().Add(5 * time.Second))
	all, err := io.ReadAll(pc)
	if err != nil {
		t.Fatalf("ReadAll: %v", err)
	}
	if string(all) != "HELLOWORLD" {
		t.Fatalf("got %q, want %q", all, "HELLOWORLD")
	}
}

func TestPrefixConnSmallReads(t *testing.T) {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	defer func() { _ = ln.Close() }()

	go func() {
		conn, err := ln.Accept()
		if err != nil {
			return
		}
		defer func() { _ = conn.Close() }()
		_, _ = conn.Write([]byte("CD"))
	}()

	conn, err := net.Dial("tcp", ln.Addr().String())
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	defer func() { _ = conn.Close() }()

	pc := NewPrefixConn(conn, []byte("AB"))

	// Read 1 byte at a time from the prefix.
	buf := make([]byte, 1)
	n, err := pc.Read(buf)
	if err != nil || n != 1 || buf[0] != 'A' {
		t.Fatalf("first read: n=%d, err=%v, buf=%q", n, err, buf[:n])
	}
	n, err = pc.Read(buf)
	if err != nil || n != 1 || buf[0] != 'B' {
		t.Fatalf("second read: n=%d, err=%v, buf=%q", n, err, buf[:n])
	}

	// Now reads come from the underlying conn.
	_ = conn.SetReadDeadline(time.Now().Add(5 * time.Second))
	rest, err := io.ReadAll(pc)
	if err != nil {
		t.Fatalf("ReadAll: %v", err)
	}
	if string(rest) != "CD" {
		t.Fatalf("got %q, want %q", rest, "CD")
	}
}

func TestPrefixConnEmptyPrefix(t *testing.T) {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	defer func() { _ = ln.Close() }()

	go func() {
		conn, err := ln.Accept()
		if err != nil {
			return
		}
		defer func() { _ = conn.Close() }()
		_, _ = conn.Write([]byte("DATA"))
	}()

	conn, err := net.Dial("tcp", ln.Addr().String())
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	defer func() { _ = conn.Close() }()

	pc := NewPrefixConn(conn, nil)
	_ = conn.SetReadDeadline(time.Now().Add(5 * time.Second))
	all, err := io.ReadAll(pc)
	if err != nil {
		t.Fatalf("ReadAll: %v", err)
	}
	if string(all) != "DATA" {
		t.Fatalf("got %q, want %q", all, "DATA")
	}
}

func TestPrefixConnDelegates(t *testing.T) {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	defer func() { _ = ln.Close() }()

	go func() {
		conn, _ := ln.Accept()
		if conn != nil {
			_ = conn.Close()
		}
	}()

	conn, err := net.Dial("tcp", ln.Addr().String())
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	defer func() { _ = conn.Close() }()

	pc := NewPrefixConn(conn, []byte("X"))

	// RemoteAddr, LocalAddr should delegate.
	if pc.RemoteAddr() == nil {
		t.Fatal("RemoteAddr returned nil")
	}
	if pc.LocalAddr() == nil {
		t.Fatal("LocalAddr returned nil")
	}
}