// Copyright (c) Tailscale Inc & contributors
// SPDX-License-Identifier: BSD-3-Clause

package main

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strconv"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

func TestManagedNaviFunnelRuntimeAppliesHTTPPayload(t *testing.T) {
	t.Parallel()

	listenPort := freeManagedNaviFunnelPort(t)
	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = io.WriteString(w, fmt.Sprintf("path=%s query=%s host=%s", r.URL.Path, r.URL.RawQuery, r.Host))
	}))
	defer backend.Close()

	backendURL, err := url.Parse(backend.URL)
	if err != nil {
		t.Fatalf("url.Parse(backend): %v", err)
	}
	backendPort, err := strconv.Atoi(backendURL.Port())
	if err != nil {
		t.Fatalf("Atoi(backend port): %v", err)
	}

	rt, err := newManagedNaviFunnelRuntime(context.Background(), t.Logf, []string{"127.0.0.1"}, nil, managedNaviFunnelTLSConfig{})
	if err != nil {
		t.Fatalf("newManagedNaviFunnelRuntime(): %v", err)
	}
	defer rt.close()

	if err := rt.applyPayload(&managedNaviFunnelPayload{
		Services: []managedNaviFunnelSyncRecord{
			{
				OrgID:   1,
				Service: map[string]any{"id": "101"},
				Public: managedNaviFunnelPublic{
					Host:        "remote-http.example.test",
					ListenProto: managedNaviFunnelListenProtoHTTP,
					ListenPort:  listenPort,
					MountPath:   "/app",
				},
				Backend: managedNaviFunnelBackend{
					Network: "tcp",
					Address: backendURL.Hostname(),
					Port:    backendPort,
					Scheme:  "http",
				},
			},
		},
	}); err != nil {
		t.Fatalf("applyPayload(http): %v", err)
	}

	client := &http.Client{Timeout: 2 * time.Second}
	var (
		respBody string
		lastErr  error
	)
	for attempt := 0; attempt < 30; attempt++ {
		req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("http://127.0.0.1:%d/app/hello?x=1", listenPort), nil)
		if err != nil {
			t.Fatalf("http.NewRequest(): %v", err)
		}
		req.Host = "remote-http.example.test"
		resp, err := client.Do(req)
		if err != nil {
			lastErr = err
			time.Sleep(100 * time.Millisecond)
			continue
		}
		body, err := io.ReadAll(resp.Body)
		resp.Body.Close()
		if err != nil {
			t.Fatalf("ReadAll(response): %v", err)
		}
		if resp.StatusCode != http.StatusOK {
			lastErr = fmt.Errorf("status=%d body=%s", resp.StatusCode, string(body))
			time.Sleep(100 * time.Millisecond)
			continue
		}
		respBody = string(body)
		lastErr = nil
		break
	}
	if lastErr != nil {
		t.Fatalf("request through managed runtime failed: %v", lastErr)
	}
	if want := "path=/hello query=x=1 host=remote-http.example.test"; respBody != want {
		t.Fatalf("runtime response = %q, want %q", respBody, want)
	}
}

func TestManagedNaviFunnelRuntimeWrapHandlerUsesExternalPort(t *testing.T) {
	t.Parallel()

	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = io.WriteString(w, fmt.Sprintf("wrapped-path=%s host=%s", r.URL.Path, r.Host))
	}))
	defer backend.Close()

	backendURL, err := url.Parse(backend.URL)
	if err != nil {
		t.Fatalf("url.Parse(backend): %v", err)
	}
	backendPort, err := strconv.Atoi(backendURL.Port())
	if err != nil {
		t.Fatalf("Atoi(backend port): %v", err)
	}

	externalKey := managedNaviFunnelListenerKey{Port: 443, TLS: true}
	rt, err := newManagedNaviFunnelRuntime(context.Background(), t.Logf, []string{"127.0.0.1"}, []managedNaviFunnelListenerKey{externalKey}, managedNaviFunnelTLSConfig{})
	if err != nil {
		t.Fatalf("newManagedNaviFunnelRuntime(): %v", err)
	}
	defer rt.close()

	if err := rt.applyPayload(&managedNaviFunnelPayload{
		Services: []managedNaviFunnelSyncRecord{
			{
				OrgID:   2,
				Service: map[string]any{"id": "202"},
				Public: managedNaviFunnelPublic{
					Host:        "external-https.example.test",
					ListenProto: managedNaviFunnelListenProtoHTTPS,
					ListenPort:  443,
					MountPath:   "/secure",
				},
				Backend: managedNaviFunnelBackend{
					Network: "tcp",
					Address: backendURL.Hostname(),
					Port:    backendPort,
					Scheme:  "http",
				},
			},
		},
	}); err != nil {
		t.Fatalf("applyPayload(https): %v", err)
	}

	handler := rt.wrapHandler(externalKey, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "fallback", http.StatusTeapot)
	}))
	srv := httptest.NewServer(handler)
	defer srv.Close()

	req, err := http.NewRequest(http.MethodGet, srv.URL+"/secure/hello", nil)
	if err != nil {
		t.Fatalf("http.NewRequest(): %v", err)
	}
	req.Host = "external-https.example.test"
	resp, err := srv.Client().Do(req)
	if err != nil {
		t.Fatalf("client.Do(proxy): %v", err)
	}
	body, err := io.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		t.Fatalf("ReadAll(proxy): %v", err)
	}
	if got, want := resp.StatusCode, http.StatusOK; got != want {
		t.Fatalf("proxy status = %d, want %d body=%s", got, want, string(body))
	}
	if want := "wrapped-path=/hello host=external-https.example.test"; string(body) != want {
		t.Fatalf("proxy body = %q, want %q", string(body), want)
	}

	req, err = http.NewRequest(http.MethodGet, srv.URL+"/unmatched", nil)
	if err != nil {
		t.Fatalf("http.NewRequest(fallback): %v", err)
	}
	req.Host = "no-match.example.test"
	resp, err = srv.Client().Do(req)
	if err != nil {
		t.Fatalf("client.Do(fallback): %v", err)
	}
	resp.Body.Close()
	if got, want := resp.StatusCode, http.StatusTeapot; got != want {
		t.Fatalf("fallback status = %d, want %d", got, want)
	}
}

func TestManagedNaviFunnelRuntimeRejectsTCPPayload(t *testing.T) {
	t.Parallel()

	rt, err := newManagedNaviFunnelRuntime(context.Background(), t.Logf, []string{"127.0.0.1"}, nil, managedNaviFunnelTLSConfig{})
	if err != nil {
		t.Fatalf("newManagedNaviFunnelRuntime(): %v", err)
	}
	defer rt.close()

	err = rt.applyPayload(&managedNaviFunnelPayload{
		Services: []managedNaviFunnelSyncRecord{
			{
				Public: managedNaviFunnelPublic{
					Host:        "remote-tcp.example.test",
					ListenProto: "tcp",
					ListenPort:  freeManagedNaviFunnelPort(t),
				},
				Backend: managedNaviFunnelBackend{
					Network: "tcp",
					Address: "100.64.0.10",
					Port:    5432,
				},
			},
		},
	})
	if err == nil {
		t.Fatal("expected tcp payload to be rejected")
	}
}

func TestManagedNaviFunnelRuntimeProxiesWebsocket(t *testing.T) {
	t.Parallel()

	listenPort := freeManagedNaviFunnelPort(t)
	upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			t.Errorf("backend websocket upgrade failed: %v", err)
			return
		}
		defer conn.Close()

		mt, payload, err := conn.ReadMessage()
		if err != nil {
			t.Errorf("backend websocket read failed: %v", err)
			return
		}
		if err := conn.WriteMessage(mt, append([]byte("echo:"), payload...)); err != nil {
			t.Errorf("backend websocket write failed: %v", err)
		}
	}))
	defer backend.Close()

	backendURL, err := url.Parse(backend.URL)
	if err != nil {
		t.Fatalf("url.Parse(backend): %v", err)
	}
	backendPort, err := strconv.Atoi(backendURL.Port())
	if err != nil {
		t.Fatalf("Atoi(backend port): %v", err)
	}

	rt, err := newManagedNaviFunnelRuntime(context.Background(), t.Logf, []string{"127.0.0.1"}, nil, managedNaviFunnelTLSConfig{})
	if err != nil {
		t.Fatalf("newManagedNaviFunnelRuntime(): %v", err)
	}
	defer rt.close()

	if err := rt.applyPayload(&managedNaviFunnelPayload{
		Services: []managedNaviFunnelSyncRecord{
			{
				OrgID:   3,
				Service: map[string]any{"id": "303"},
				Public: managedNaviFunnelPublic{
					Host:        "remote-ws.example.test",
					ListenProto: managedNaviFunnelListenProtoWS,
					ListenPort:  listenPort,
					MountPath:   "/socket",
				},
				Backend: managedNaviFunnelBackend{
					Network: "tcp",
					Address: backendURL.Hostname(),
					Port:    backendPort,
					Scheme:  "http",
				},
			},
		},
	}); err != nil {
		t.Fatalf("applyPayload(ws): %v", err)
	}

	dialer := websocket.Dialer{
		NetDialContext: func(ctx context.Context, network, _ string) (net.Conn, error) {
			var d net.Dialer
			return d.DialContext(ctx, network, net.JoinHostPort("127.0.0.1", strconv.Itoa(listenPort)))
		},
		HandshakeTimeout: 2 * time.Second,
	}
	conn, _, err := dialer.Dial(fmt.Sprintf("ws://remote-ws.example.test:%d/socket/echo", listenPort), nil)
	if err != nil {
		t.Fatalf("websocket dial through runtime failed: %v", err)
	}
	defer conn.Close()

	if err := conn.WriteMessage(websocket.TextMessage, []byte("hello")); err != nil {
		t.Fatalf("websocket write failed: %v", err)
	}
	mt, payload, err := conn.ReadMessage()
	if err != nil {
		t.Fatalf("websocket read failed: %v", err)
	}
	if mt != websocket.TextMessage {
		t.Fatalf("message type = %d, want %d", mt, websocket.TextMessage)
	}
	if got, want := string(payload), "echo:hello"; got != want {
		t.Fatalf("payload = %q, want %q", got, want)
	}
}

func freeManagedNaviFunnelPort(t *testing.T) int {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("Listen(): %v", err)
	}
	defer ln.Close()
	return ln.Addr().(*net.TCPAddr).Port
}
