package e2e

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
	"testing"
	"time"

	"git.veillax.com/Veillax/blacktip"
)

// proxyBackend is a TCP server that understands the PROXY protocol the way a
// game proxy or mail server would. It replies with what it learned, or
// "plain:<first line>" when the connection did not start with a header.
func proxyBackend(t *testing.T, seen chan<- string) int {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { l.Close() })
	go func() {
		for {
			c, err := l.Accept()
			if err != nil {
				return
			}
			go func() {
				defer c.Close()
				c.SetDeadline(time.Now().Add(5 * time.Second))
				br := bufio.NewReader(c)
				first, _ := br.Peek(1)
				var head []byte
				switch {
				case len(first) == 1 && first[0] == 0x0d:
					head, _ = br.Peek(12) // a v2 header starts with a 12-byte signature
				case len(first) == 1 && first[0] == 'P':
					head, _ = br.Peek(6) // "PROXY "
				}
				var desc string
				switch {
				case bytes.HasPrefix(head, []byte("PROXY ")):
					line, _ := br.ReadString('\n')
					f := strings.Fields(strings.TrimSpace(line))
					if len(f) == 6 {
						desc = fmt.Sprintf("v1 src=%s:%s dst=%s:%s fam=%s", f[2], f[4], f[3], f[5], f[1])
					} else {
						desc = "v1 odd:" + line
					}
				case bytes.Equal(head, []byte("\r\n\r\n\x00\r\nQUIT\n")):
					fixed := make([]byte, 16)
					io.ReadFull(br, fixed)
					body := make([]byte, binary.BigEndian.Uint16(fixed[14:16]))
					io.ReadFull(br, body)
					switch fixed[13] {
					case 0x11:
						desc = fmt.Sprintf("v2 src=%s:%d dst=%s:%d fam=TCP4", net.IP(body[0:4]), binary.BigEndian.Uint16(body[8:10]), net.IP(body[4:8]), binary.BigEndian.Uint16(body[10:12]))
					case 0x21:
						desc = fmt.Sprintf("v2 src=[%s]:%d fam=TCP6", net.IP(body[0:16]), binary.BigEndian.Uint16(body[32:34]))
					default:
						desc = fmt.Sprintf("v2 fam=%#x", fixed[13])
					}
					if fixed[12] != 0x21 {
						desc += " cmd=" + strconv.Itoa(int(fixed[12]))
					}
				default:
					line, _ := br.ReadString('\n')
					desc = "plain:" + strings.TrimSpace(line)
				}
				// what comes next shows the header was exactly the first thing and nothing was lost
				c.SetReadDeadline(time.Now().Add(400 * time.Millisecond)) // don't wait forever for bytes that may never come
				if b, err := br.Peek(1); err == nil && b[0] == 0x16 {
					desc += " then=tls-hello"
				}
				if seen != nil {
					seen <- desc
				}
				fmt.Fprintf(c, "%s\n", desc)
			}()
		}
	}()
	return l.Addr().(*net.TCPAddr).Port
}

func askFrom(t *testing.T, from, addr, send string) string {
	t.Helper()
	var last error
	// The first connection through the tunnel can wait for the WireGuard handshake, which is
	// slow on a busy CI runner, so a read timeout is retried until the overall deadline.
	deadline := time.Now().Add(30 * time.Second)
	for time.Now().Before(deadline) {
		d := net.Dialer{LocalAddr: &net.TCPAddr{IP: net.ParseIP(from)}, Timeout: time.Second}
		c, err := d.Dial("tcp", addr)
		if err == nil {
			c.SetDeadline(time.Now().Add(4 * time.Second))
			fmt.Fprint(c, send)
			line, rerr := bufio.NewReader(c).ReadString('\n')
			c.Close()
			if rerr == nil {
				return strings.TrimSpace(line)
			}
			last = rerr
		} else {
			last = err
		}
		time.Sleep(150 * time.Millisecond)
	}
	t.Fatalf("no answer from %s: %v", addr, last)
	return ""
}

func TestProxyProtocolCarriesTheRealClientAddress(t *testing.T) {
	w := setup(t, blacktip.TraefikConfig{})
	be := proxyBackend(t, nil)
	base := "/api/orgs/" + w.org + "/mappings"
	open := func(ip string, site string, proxy string) string {
		port := freePort(t)
		w.admin.must("POST", base, map[string]any{"proto": "tcp", "edgeId": w.edge, "bindIp": ip, "port": port, "targets": tgt(site, be),
			"enabled": true, "proxyProtocol": proxy}, nil)
		return ip + ":" + strconv.Itoa(port)
	}

	// through the WireGuard tunnel and straight to the edge's own machine, in both versions
	for _, tc := range []struct{ name, site, proxy string }{
		{"v1 via tunnel", w.site, "v1"}, {"v2 via tunnel", w.site, "v2"}, {"v1 local", w.local, "v1"}, {"v2 local", w.local, "v2"},
	} {
		addr := open("127.0.0.1", tc.site, tc.proxy)
		t.Logf("case %s on %s", tc.name, addr)
		got := askFrom(t, "127.0.0.9", addr, "hello\n")
		if !strings.HasPrefix(got, tc.proxy) || !strings.Contains(got, "src=127.0.0.9:") || !strings.Contains(got, "dst="+addr) || !strings.Contains(got, "fam=TCP4") {
			t.Errorf("%s: the target learned %q, want the real client 127.0.0.9 and the public address %s", tc.name, got, addr)
		}
	}

	// off by default: the target gets exactly what the client sent, no header
	addr := open("127.0.0.1", w.site, "")
	if got := askFrom(t, "127.0.0.9", addr, "hello\n"); got != "plain:hello" {
		t.Errorf("without PROXY protocol the target saw %q", got)
	}

	// the header comes before the client's own bytes, even when they are a TLS hello
	seen := make(chan string, 4)
	tbe := proxyBackend(t, seen)
	port := freePort(t)
	w.admin.must("POST", "/api/orgs/"+w.org+"/domains", map[string]any{"name": "example.test", "edgeId": w.edge}, nil)
	w.admin.must("POST", base, map[string]any{"proto": "tls", "edgeId": w.edge, "bindIp": "127.0.0.1", "port": port, "hostname": "pass.example.test",
		"targets": tgt(w.site, tbe), "enabled": true, "proxyProtocol": "v2"}, nil)
	for i := 0; i < 50; i++ {
		d := net.Dialer{LocalAddr: &net.TCPAddr{IP: net.ParseIP("127.0.0.9")}, Timeout: time.Second}
		c, err := d.Dial("tcp", "127.0.0.1:"+strconv.Itoa(port))
		if err == nil {
			tls.Client(c, &tls.Config{ServerName: "pass.example.test", InsecureSkipVerify: true}).Handshake() // fails: the backend isn't TLS
			c.Close()
		}
		select {
		case got := <-seen:
			if !strings.Contains(got, "v2 src=127.0.0.9:") || !strings.Contains(got, "then=tls-hello") {
				t.Errorf("tls passthrough: %q", got)
			}
			goto validated
		case <-time.After(200 * time.Millisecond):
		}
	}
	t.Fatal("tls passthrough never reached the target")
validated:

	// only tcp and tls ports can carry it
	w.admin.expect(400, "POST", base, map[string]any{"proto": "udp", "edgeId": w.edge, "bindIp": "127.0.0.1", "port": freePort(t), "targets": tgt(w.site, be), "proxyProtocol": "v1"})
	w.admin.expect(400, "POST", base, map[string]any{"proto": "http", "edgeId": w.edge, "bindIp": "127.0.0.1", "port": freePort(t), "hostname": "w.example.test", "targets": tgt(w.site, be), "proxyProtocol": "v2"})
	w.admin.expect(400, "POST", base, map[string]any{"proto": "tcp", "edgeId": w.edge, "bindIp": "127.0.0.1", "port": freePort(t), "targets": tgt(w.site, be), "proxyProtocol": "v3"})
}