package tests

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

	"github.com/VictoriaMetrics/VictoriaMetrics/apptest"
)

func TestSingleVMAuthRouterWithAuth(t *testing.T) {
	tc := apptest.NewTestCase(t)
	defer tc.Stop()

	var authorizedRequestsCount, unauthorizedRequestsCount int
	backendWithAuth := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
		authorizedRequestsCount++
	}))
	defer backendWithAuth.Close()
	backend := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
		unauthorizedRequestsCount++
	}))
	defer backend.Close()

	authConfig := fmt.Sprintf(`
users:
- name: user1
  username: ba-username
  password: ba-password
  url_prefix: %s
unauthorized_user:
   url_map:
   - src_paths:
     - /backend/health
     - /backend/ready
     url_prefix: %s
  `, backendWithAuth.URL, backend.URL)

	vmauth := tc.MustStartVmauth("vmauth", nil, authConfig)

	makeGetRequestExpectCode := func(prepareRequest func(*http.Request), expectCode int) {
		t.Helper()
		req, err := http.NewRequest("GET", fmt.Sprintf("http://%s", vmauth.GetHTTPListenAddr()), nil)
		if err != nil {
			t.Fatalf("cannot build http.Request: %s", err)
		}
		prepareRequest(req)
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			t.Fatalf("cannot make http.Get request for target=%q: %s", req.URL, err)
		}
		responseText, err := io.ReadAll(resp.Body)
		if err != nil {
			t.Fatalf("cannot read response body: %s", err)
		}
		resp.Body.Close()
		if resp.StatusCode != expectCode {
			t.Fatalf("unexpected http response code: %d, want: %d, response text: %s", resp.StatusCode, expectCode, responseText)
		}
	}
	assertBackendsRequestsCount := func(expectAuthorized, expectUnauthorized int) {
		t.Helper()
		if expectAuthorized != authorizedRequestsCount {
			t.Fatalf("expected to have %d authorized proxied requests, got: %d", expectAuthorized, authorizedRequestsCount)
		}

		if expectUnauthorized != unauthorizedRequestsCount {
			t.Fatalf("expected to have %d unauthorized proxied requests, got: %d", expectUnauthorized, unauthorizedRequestsCount)
		}

	}

	makeGetRequestExpectCode(func(r *http.Request) {
		r.URL.Path = "/backend/api"
		r.URL.User = url.UserPassword("ba-username", "ba-password")
	}, http.StatusOK)
	assertBackendsRequestsCount(1, 0)

	makeGetRequestExpectCode(func(r *http.Request) {
		r.URL.Path = "/backend/health"
	}, http.StatusOK)
	assertBackendsRequestsCount(1, 1)

	// remove unauthorized section and proxy only specified path for authorized
	vmauth.UpdateConfiguration(t, fmt.Sprintf(`
users:
- name: user1
  username: ba-username
  password: ba-password
  url_map:
  - src_paths:
    - /backend/health
    url_prefix: %s
`, backendWithAuth.URL))

	// ensure unauthorized requests no longer served
	makeGetRequestExpectCode(func(r *http.Request) {
		r.URL.Path = "/backend/health"
	}, http.StatusUnauthorized)
	assertBackendsRequestsCount(1, 1)

	makeGetRequestExpectCode(func(r *http.Request) {
		r.URL.User = url.UserPassword("ba-username", "ba-password")
		r.URL.Path = "/backend/health"
	}, http.StatusOK)
	assertBackendsRequestsCount(2, 1)

	// url path is missing at proxy configuration
	makeGetRequestExpectCode(func(r *http.Request) {
		r.URL.User = url.UserPassword("ba-username", "ba-password")
		r.URL.Path = "/backend"
	}, http.StatusBadRequest)
	assertBackendsRequestsCount(2, 1)

}

func TestSingleVMAuthRouterWithInternalAddr(t *testing.T) {
	tc := apptest.NewTestCase(t)
	defer tc.Stop()

	var proxiedRequestsCount int
	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
		proxiedRequestsCount++
	}))
	defer backend.Close()

	authConfig := fmt.Sprintf(`
unauthorized_user:
   url_prefix: %s
  `, backend.URL)

	const (
		// it's not possible to use random ports
		// since it makes test flaky
		listenPortPublic  = "50127"
		listenPortPrivate = "50126"
	)

	vmauthFlags := []string{
		fmt.Sprintf("-httpListenAddr=127.0.0.1:%s", listenPortPublic),
		fmt.Sprintf("-httpInternalListenAddr=127.0.0.1:%s", listenPortPrivate),
		"-flagsAuthKey=protected",
	}
	vmauth := tc.MustStartVmauth("vmauth", vmauthFlags, authConfig)

	makeGetRequestExpectCode := func(targetURL string, expectCode int) {
		t.Helper()
		resp, err := http.Get(targetURL)
		if err != nil {
			t.Fatalf("cannot make http.Get request for target=%q: %s", targetURL, err)
		}
		responseText, err := io.ReadAll(resp.Body)
		if err != nil {
			t.Fatalf("cannot read response body: %s", err)
		}
		resp.Body.Close()
		if resp.StatusCode != expectCode {
			t.Fatalf("unexpected http response code: %d, want: %d, response text: %s", resp.StatusCode, expectCode, responseText)
		}
	}
	assertBackendRequestsCount := func(expected int) {
		t.Helper()
		if proxiedRequestsCount != expected {
			t.Fatalf("expected to have %d proxied requests, got: %d", expected, proxiedRequestsCount)
		}
	}
	// built-in http server must reject request, since it protected with authKey
	makeGetRequestExpectCode(fmt.Sprintf("http://127.0.0.1:%s/flags", listenPortPrivate), http.StatusUnauthorized)
	assertBackendRequestsCount(0)

	makeGetRequestExpectCode(fmt.Sprintf("http://127.0.0.1:%s/flags", listenPortPublic), http.StatusOK)
	assertBackendRequestsCount(1)

	// reload config and ensure that vmauth no longer proxies requests to the backend
	vmauth.UpdateConfiguration(t, "")
	makeGetRequestExpectCode(fmt.Sprintf("http://127.0.0.1:%s/flags", listenPortPrivate), http.StatusUnauthorized)
	assertBackendRequestsCount(1)
}

func TestSingleVMAuthHTTPServerAuthKeys(t *testing.T) {
	tc := apptest.NewTestCase(t)
	defer tc.Stop()

	var authorizedRequestsCount int
	backend := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
		authorizedRequestsCount++
	}))
	defer backend.Close()

	authConfig := fmt.Sprintf(`
unauthorized_user:
   url_map:
   - src_paths:
     - /backend/health
     - /backend/ready
     url_prefix: %s
  `, backend.URL)

	const (
		authKey  = "key"
		username = "user"
		password = "password"
	)
	flags := []string{
		"--reloadAuthKey=" + authKey,
		"--pprofAuthKey=" + authKey,
		"--metricsAuthKey=" + authKey,
		"--httpAuth.username=" + username,
		"--httpAuth.password=" + password}
	vmauth := tc.MustStartVmauth("vmauth", flags, authConfig)

	makeGetRequestExpectCode := func(prepareRequest func(*http.Request), expectCode int) {
		t.Helper()
		req, err := http.NewRequest("GET", fmt.Sprintf("http://%s", vmauth.GetHTTPListenAddr()), nil)
		if err != nil {
			t.Fatalf("cannot build http.Request: %s", err)
		}
		prepareRequest(req)
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			t.Fatalf("cannot make http.Get request for target=%q: %s", req.URL, err)
		}
		responseText, err := io.ReadAll(resp.Body)
		if err != nil {
			t.Fatalf("cannot read response body: %s", err)
		}
		resp.Body.Close()
		if resp.StatusCode != expectCode {
			t.Fatalf("unexpected http response code: %d, want: %d, response text: %s", resp.StatusCode, expectCode, responseText)
		}
	}
	assertBackendsRequestsCount := func(expectAuthorized int) {
		t.Helper()
		if expectAuthorized != authorizedRequestsCount {
			t.Fatalf("expected to have %d authorized proxied requests, got: %d", expectAuthorized, authorizedRequestsCount)
		}
	}

	// authKey overrides basic auth
	makeGetRequestExpectCode(func(r *http.Request) {
		r.URL.Path = "/metrics"
		q := r.URL.Query()
		q.Add("authKey", authKey)
		r.URL.RawQuery = q.Encode()
	}, http.StatusOK)
	assertBackendsRequestsCount(0)

	// authKey overrides basic auth at pprof handler
	makeGetRequestExpectCode(func(r *http.Request) {
		r.URL.Path = "/debug/pprof/heap"
		q := r.URL.Query()
		q.Add("authKey", authKey)
		r.URL.RawQuery = q.Encode()
	}, http.StatusOK)
	assertBackendsRequestsCount(0)

	// authKey overrides basic auth at app internal router
	makeGetRequestExpectCode(func(r *http.Request) {
		r.URL.Path = "/-/reload"
		q := r.URL.Query()
		q.Add("authKey", authKey)
		r.URL.RawQuery = q.Encode()
	}, http.StatusOK)
	assertBackendsRequestsCount(0)

	// no auth request fails
	makeGetRequestExpectCode(func(r *http.Request) {
		r.URL.Path = "/metrics"
	}, http.StatusUnauthorized)
	assertBackendsRequestsCount(0)

	// no auth request fails at app internal router
	makeGetRequestExpectCode(func(r *http.Request) {
		r.URL.Path = "/-/reload"
	}, http.StatusUnauthorized)
	assertBackendsRequestsCount(0)

	// basic auth request ok
	makeGetRequestExpectCode(func(r *http.Request) {
		r.URL.Path = "/metrics"
		r.URL.User = url.UserPassword(username, password)
	}, http.StatusUnauthorized)
	assertBackendsRequestsCount(0)

	// basic auth request ok at app internal router
	makeGetRequestExpectCode(func(r *http.Request) {
		r.URL.Path = "/-/reload"
		r.URL.User = url.UserPassword(username, password)
	}, http.StatusUnauthorized)
	assertBackendsRequestsCount(0)

	// basic auth to backend ok
	makeGetRequestExpectCode(func(r *http.Request) {
		r.URL.Path = "/backend/health"
		r.URL.User = url.UserPassword(username, password)
	}, http.StatusOK)
	assertBackendsRequestsCount(1)

}

func TestSingleVMAuthUseProxyProtocol(t *testing.T) {
	tc := apptest.NewTestCase(t)
	defer tc.Stop()

	var requestsCount int
	var actualForwardedForHeader string
	backend := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
		actualForwardedForHeader = r.Header.Get("X-Forwarded-For")
		requestsCount++
	}))
	defer backend.Close()

	authConfig := fmt.Sprintf(`
unauthorized_user:
  url_prefix: %s
  `, backend.URL)

	vmauth := tc.MustStartVmauth("vmauth", []string{
		"-httpListenAddr.useProxyProtocol=true",
	}, authConfig)

	req, err := http.NewRequest("GET", fmt.Sprintf("http://%s/backend", vmauth.GetHTTPListenAddr()), nil)
	if err != nil {
		t.Fatalf("cannot build http.Request: %s", err)
	}

	// make request using proxy protocol
	c := &http.Client{
		Transport: &http.Transport{
			DialContext: func(_ context.Context, network, addr string) (net.Conn, error) {
				conn, err := net.Dial(network, addr)
				if err != nil {
					return nil, err
				}

				// Write a proxy protocol header to the connection
				if _, err := conn.Write([]byte{
					0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A, // signature
					0x21,       // version 2
					0x11,       // family IPv4
					0x00, 0x0C, // length: 12 bytes (IPv4 + ports)
					192, 168, 1, 100, // source IP
					10, 0, 0, 1, // destination IP
					0x1F, 0x90, // source port 8080
					0x00, 0x50, // destination port 80
				}); err != nil {
					t.Fatalf("cannot send proxy protocol header: %s", err)
				}

				return conn, nil
			},
		},
	}

	resp, err := c.Do(req)
	if err != nil {
		t.Fatalf("cannot make http.Get request for target=%q: %s", req.URL, err)
	}
	responseText, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("cannot read response body: %s", err)
	}
	resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("unexpected http response code: %d, want: %d, response text: %s", resp.StatusCode, http.StatusOK, responseText)
	}

	// ensure that request was proxied
	if requestsCount != 1 {
		t.Fatalf("expected to have %d unauthorized proxied requests, got: %d", 1, requestsCount)
	}

	// ensure that X-Forwarded-For header is set to the source IP from proxy protocol
	expectedForwardedForHeader := "192.168.1.100"
	if actualForwardedForHeader != expectedForwardedForHeader {
		t.Fatalf("expected X-Forwarded-For header to be equal to proxy source IP, got: %s, want: %s'", actualForwardedForHeader, expectedForwardedForHeader)
	}
}
