package main

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strconv"
	"strings"
	"testing"
	"time"
)

func TestMemoryProxyUsesSOCKS5Credentials(t *testing.T) {
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer listener.Close()
	done := make(chan error, 1)
	go func() {
		connection, err := listener.Accept()
		if err != nil {
			done <- err
			return
		}
		defer connection.Close()
		greeting := make([]byte, 3)
		if _, err := io.ReadFull(connection, greeting); err != nil {
			done <- err
			return
		}
		_, _ = connection.Write([]byte{5, 2})
		authHeader := make([]byte, 2)
		_, _ = io.ReadFull(connection, authHeader)
		username := make([]byte, int(authHeader[1]))
		_, _ = io.ReadFull(connection, username)
		var passwordLength [1]byte
		_, _ = io.ReadFull(connection, passwordLength[:])
		password := make([]byte, int(passwordLength[0]))
		_, _ = io.ReadFull(connection, password)
		if string(username) != "operator" || string(password) != "ephemeral" {
			done <- io.ErrUnexpectedEOF
			return
		}
		_, _ = connection.Write([]byte{1, 0})
		requestHeader := make([]byte, 5)
		_, _ = io.ReadFull(connection, requestHeader)
		host := make([]byte, int(requestHeader[4]))
		_, _ = io.ReadFull(connection, host)
		port := make([]byte, 2)
		_, _ = io.ReadFull(connection, port)
		if string(host) != "example.com" || binary.BigEndian.Uint16(port) != 443 {
			done <- io.ErrUnexpectedEOF
			return
		}
		if _, err = connection.Write([]byte{5, 0, 0, 1, 127, 0, 0, 1, 0, 0}); err != nil {
			done <- err
			return
		}
		var tunneled [1]byte
		_, err = io.ReadFull(connection, tunneled[:])
		if err == nil && tunneled[0] != 'x' {
			err = io.ErrUnexpectedEOF
		}
		done <- err
	}()

	host, portText, _ := net.SplitHostPort(listener.Addr().String())
	port, _ := net.LookupPort("tcp", portText)
	registry := newMemoryProxyRegistry()
	proxyURL, cleanup, err := registry.configure("account-a", 1, "127.0.0.1", 0, gatewayProxyExit{
		Protocol: "socks5", Host: host, Port: port, Username: "operator", Password: "ephemeral",
	})
	if err != nil {
		t.Fatal(err)
	}
	defer cleanup()
	parsed, _ := url.Parse(proxyURL)
	connection, err := net.Dial("tcp", strings.Replace(parsed.Host, browserProxyHost, "127.0.0.1", 1))
	if err != nil {
		t.Fatal(err)
	}
	if _, err := fmt.Fprint(connection, "CONNECT example.com:443 HTTP/1.1\r\nHost: example.com:443\r\n\r\n"); err != nil {
		t.Fatal(err)
	}
	response, err := http.ReadResponse(bufio.NewReader(connection), &http.Request{Method: http.MethodConnect})
	if err != nil || response.StatusCode != http.StatusOK {
		t.Fatalf("memory proxy CONNECT failed: response=%v err=%v", response, err)
	}
	if _, err := connection.Write([]byte{'x'}); err != nil {
		t.Fatal(err)
	}
	_ = connection.Close()
	if err := <-done; err != nil {
		t.Fatal(err)
	}
}

func TestMemoryProxyUsesAbsoluteFormForHTTPUpstream(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
		if request.Method == http.MethodConnect {
			http.Error(response, "CONNECT forbidden", http.StatusMethodNotAllowed)
			return
		}
		if !request.URL.IsAbs() || request.URL.String() != "http://example.com/plain" {
			t.Fatalf("expected absolute-form request, got %q", request.URL.String())
		}
		if request.Header.Get("Proxy-Authorization") == "" {
			t.Fatal("upstream proxy credentials were not applied")
		}
		_, _ = response.Write([]byte("forwarded"))
	}))
	defer upstream.Close()
	address, _ := url.Parse(upstream.URL)
	port, _ := strconv.Atoi(address.Port())
	registry := newMemoryProxyRegistry()
	proxyURL, cleanup, err := registry.configure("account-a", 1, "127.0.0.1", 0, gatewayProxyExit{
		Protocol: "http", Host: address.Hostname(), Port: port, Username: "operator", Password: "ephemeral",
	})
	if err != nil {
		t.Fatal(err)
	}
	defer cleanup()
	proxyAddress := strings.Replace(strings.TrimPrefix(proxyURL, "http://"), browserProxyHost, "127.0.0.1", 1)
	client := &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(&url.URL{Scheme: "http", Host: proxyAddress})}}
	response, err := client.Get("http://example.com/plain")
	if err != nil {
		t.Fatal(err)
	}
	defer response.Body.Close()
	body, _ := io.ReadAll(response.Body)
	if response.StatusCode != http.StatusOK || string(body) != "forwarded" {
		t.Fatalf("plain HTTP was not forwarded: status=%d body=%s", response.StatusCode, body)
	}
}

func TestMemoryProxyRejectsCrossAliasAddress(t *testing.T) {
	registry := newMemoryProxyRegistry()
	proxyURL, cleanup, err := registry.configure("account-a", 1, "127.0.0.1", 0, gatewayProxyExit{Protocol: "http", Host: "127.0.0.1", Port: 1})
	if err != nil {
		t.Fatal(err)
	}
	defer cleanup()
	parsed, _ := url.Parse(proxyURL)
	if connection, err := net.DialTimeout("tcp", net.JoinHostPort("127.0.0.2", parsed.Port()), 100*time.Millisecond); err == nil {
		_ = connection.Close()
		t.Fatal("another tenant address could reach account-a proxy")
	}
}

func TestMemoryProxyRemoveRequiresMatchingGeneration(t *testing.T) {
	registry := newMemoryProxyRegistry()
	proxyURL, cleanup, err := registry.configure("account-a", 2, "127.0.0.1", 0, gatewayProxyExit{Protocol: "http", Host: "127.0.0.1", Port: 1})
	if err != nil {
		t.Fatal(err)
	}
	defer cleanup()
	if !registry.bind("account-a", 2, proxyURL, "container-c2") {
		t.Fatal("bind proxy generation")
	}
	proxy := registry.proxies["account-a"]
	if registry.remove("account-a", 2, "") || registry.remove("account-a", 2, "container-c1") || registry.proxies["account-a"] != proxy {
		t.Fatal("stale generation removed the current proxy")
	}
}

func TestMemoryProxyReplacesStaleGenerationWithoutOldCleanup(t *testing.T) {
	registry := newMemoryProxyRegistry()
	exit := gatewayProxyExit{Protocol: "http", Host: "127.0.0.1", Port: 1}
	oldURL, oldCleanup, err := registry.configure("account-a", 1, "127.0.0.1", 0, exit, "network-n1")
	if err != nil {
		t.Fatal(err)
	}
	oldPort := registry.proxies["account-a"].listener.Addr().(*net.TCPAddr).Port
	newURL, newCleanup, err := registry.configure("account-a", 2, "127.0.0.1", oldPort, exit, "network-n2")
	if err != nil {
		t.Fatal(err)
	}
	defer newCleanup()
	if newURL != oldURL || registry.proxies["account-a"].bindingVersion != 2 {
		t.Fatalf("stale proxy was not replaced: old=%q new=%q proxy=%#v", oldURL, newURL, registry.proxies["account-a"])
	}
	oldCleanup()
	if registry.proxies["account-a"].bindingVersion != 2 {
		t.Fatal("old cleanup removed the replacement proxy")
	}
}