package main

import (
	"bytes"
	"encoding/binary"
	"io"
	"net"
	"testing"
	"time"
)

type proxyRecord struct {
	header  []byte
	payload []byte
	peer    net.IP
}

func TestMailTarget(t *testing.T) {
	tunnel, err := parseCIDR("10.153.238.225/30")
	if err != nil {
		t.Fatal(err)
	}
	target, upstream, err := mailProxyTarget("stalwart-relay.stalwart.svc.cluster.local", tunnel)
	if err != nil {
		t.Fatal(err)
	}
	if !target.address.Equal(tunnel.IP) || upstream != "stalwart-relay.stalwart.svc.cluster.local" {
		t.Fatalf("unexpected proxy target: %#v %q", target, upstream)
	}
	if len(target.ports) != 3 || target.ports[0] != 25 || target.ports[1] != 465 || target.ports[2] != 993 {
		t.Fatalf("unexpected proxy ports: %v", target.ports)
	}
}

func TestMailProxyTCP(t *testing.T) {
	upstream := listenTCP(t)
	listener := listenTCP(t)
	failures := make(chan error, 1)
	proxy := newMailProxy(listener, upstream.Addr().String(), net.ParseIP("127.0.0.2"), failures, make(chan struct{}, mailConnectionLimit))
	go proxy.serve()
	defer proxy.close()
	records := make(chan proxyRecord, 1)
	go func() {
		connection, err := upstream.Accept()
		if err != nil {
			return
		}
		defer connection.Close()
		header := make([]byte, 28)
		if _, err := io.ReadFull(connection, header); err != nil {
			return
		}
		payload := make([]byte, 4)
		if _, err := io.ReadFull(connection, payload); err != nil {
			return
		}
		records <- proxyRecord{header: header, payload: payload, peer: connection.RemoteAddr().(*net.TCPAddr).IP}
		_, _ = connection.Write([]byte("pong"))
	}()
	client, err := net.DialTCP("tcp4", nil, listener.Addr().(*net.TCPAddr))
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()
	if _, err := client.Write([]byte("ping")); err != nil {
		t.Fatal(err)
	}
	reply := make([]byte, 4)
	if _, err := io.ReadFull(client, reply); err != nil {
		t.Fatal(err)
	}
	if string(reply) != "pong" {
		t.Fatalf("reply = %q", reply)
	}
	select {
	case record := <-records:
		assertProxyRecord(t, record, client.LocalAddr().(*net.TCPAddr), listener.Addr().(*net.TCPAddr))
	case <-time.After(time.Second):
		t.Fatal("upstream did not receive the proxied stream")
	}
	select {
	case err := <-failures:
		t.Fatalf("unexpected proxy failure: %v", err)
	default:
	}
}

func TestMailDialFailure(t *testing.T) {
	unavailable := listenTCP(t)
	upstream := unavailable.Addr().String()
	_ = unavailable.Close()
	listener := listenTCP(t)
	failures := make(chan error, 1)
	proxy := newMailProxy(listener, upstream, nil, failures, make(chan struct{}, mailConnectionLimit))
	go proxy.serve()
	defer proxy.close()
	client, err := net.Dial("tcp4", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()
	if err := client.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
		t.Fatal(err)
	}
	if _, err := client.Read(make([]byte, 1)); err == nil {
		t.Fatal("client remained open after the upstream refusal")
	}
	select {
	case err := <-failures:
		t.Fatalf("connection-level dial failure was reported as relay-fatal: %v", err)
	default:
	}
}

func TestMailProxyClose(t *testing.T) {
	upstream := listenTCP(t)
	listener := listenTCP(t)
	proxy := newMailProxy(listener, upstream.Addr().String(), nil, make(chan error, 1), make(chan struct{}, mailConnectionLimit))
	go proxy.serve()
	accepted := make(chan net.Conn, 1)
	go func() {
		connection, err := upstream.Accept()
		if err == nil {
			accepted <- connection
		}
	}()
	client, err := net.Dial("tcp4", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()
	var server net.Conn
	select {
	case server = <-accepted:
		defer server.Close()
	case <-time.After(time.Second):
		t.Fatal("proxy did not open the upstream stream")
	}
	closed := make(chan struct{})
	go func() {
		proxy.close()
		close(closed)
	}()
	select {
	case <-closed:
	case <-time.After(time.Second):
		t.Fatal("proxy shutdown did not close active streams")
	}
}

func TestMailProxyConnectionLimit(t *testing.T) {
	upstream := listenTCP(t)
	listener := listenTCP(t)
	slots := make(chan struct{}, 1)
	proxy := newMailProxy(listener, upstream.Addr().String(), nil, make(chan error, 1), slots)
	go proxy.serve()
	defer proxy.close()
	accepted := make(chan net.Conn, 2)
	go func() {
		for {
			connection, err := upstream.Accept()
			if err != nil {
				return
			}
			accepted <- connection
		}
	}()
	first, err := net.Dial("tcp4", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer first.Close()
	firstUpstream := waitForAccepted(t, accepted)
	defer firstUpstream.Close()
	second, err := net.Dial("tcp4", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer second.Close()
	if err := second.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
		t.Fatal(err)
	}
	if _, err := second.Read(make([]byte, 1)); err == nil {
		t.Fatal("over-limit connection remained open")
	}
	select {
	case connection := <-accepted:
		_ = connection.Close()
		t.Fatal("over-limit connection reached the upstream")
	case <-time.After(100 * time.Millisecond):
	}
	_ = first.Close()
	_ = firstUpstream.Close()
	deadline := time.Now().Add(time.Second)
	for len(slots) != 0 && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if len(slots) != 0 {
		t.Fatal("connection slot was not released")
	}
	third, err := net.Dial("tcp4", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer third.Close()
	thirdUpstream := waitForAccepted(t, accepted)
	_ = thirdUpstream.Close()
}

func listenTCP(t *testing.T) *net.TCPListener {
	t.Helper()
	listener, err := net.ListenTCP("tcp4", &net.TCPAddr{IP: net.ParseIP("127.0.0.1")})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = listener.Close() })
	return listener
}

func waitForAccepted(t *testing.T, connections <-chan net.Conn) net.Conn {
	t.Helper()
	select {
	case connection := <-connections:
		return connection
	case <-time.After(time.Second):
		t.Fatal("proxy did not open the upstream stream")
		return nil
	}
}

func assertProxyRecord(t *testing.T, record proxyRecord, source, destination *net.TCPAddr) {
	t.Helper()
	if len(record.header) != 28 || !bytes.Equal(record.header[:12], proxySignature) || record.header[12] != 0x21 || record.header[13] != 0x11 || binary.BigEndian.Uint16(record.header[14:16]) != 12 {
		t.Fatalf("invalid PROXY header: %x", record.header)
	}
	if !bytes.Equal(record.header[16:20], source.IP.To4()) || !bytes.Equal(record.header[20:24], destination.IP.To4()) || binary.BigEndian.Uint16(record.header[24:26]) != uint16(source.Port) || binary.BigEndian.Uint16(record.header[26:28]) != uint16(destination.Port) {
		t.Fatalf("unexpected PROXY addresses: %x", record.header)
	}
	if string(record.payload) != "ping" {
		t.Fatalf("payload = %q", record.payload)
	}
	if !record.peer.Equal(net.ParseIP("127.0.0.2")) {
		t.Fatalf("upstream peer = %s", record.peer)
	}
}