package main

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"sort"
	"strings"
	"time"
)

const (
	originHost        = "example.test"
	sameSiteHost      = "leak.example.test"
	crossSiteHost     = "attacker.test"
	originAuthSecret  = "Bearer ORIGIN_SECRET_FROM_INITIAL_REQUEST"
	cookieSecret      = "session=COOKIE_SECRET_FROM_INITIAL_REQUEST"
	callerProxySecret = "Basic CALLER_SUPPLIED_PROXY_SECRET"
	proxyURLUser      = "proxyuser"
	proxyURLPass      = "proxy-pass"
)

type capture struct {
	Label      string
	Method     string
	RequestURI string
	Host       string
	Headers    http.Header
}

type scenario struct {
	name            string
	initialURL      string
	redirectURL     string
	directHost      string
	headerProxyAuth bool
	proxyURLAuth    bool
	expectDirectPA  bool
	expectDirectOA  bool
	expectDirectC   bool
	expectProxyPA   bool
	expectError     bool
}

func main() {
	failures := 0
	if err := runEnvNoProxyScenario(); err != nil {
		failures++
		fmt.Printf("SCENARIO env_http_proxy_no_proxy_same_site_leaks FAIL: %v\n", err)
	}

	scenarios := []scenario{
		{
			name:            "caller_proxy_authorization_same_site_proxy_to_direct_leaks",
			initialURL:      "http://" + originHost + "/start",
			redirectURL:     "http://" + sameSiteHost + "/sink",
			directHost:      sameSiteHost,
			headerProxyAuth: true,
			expectDirectPA:  true,
			expectDirectOA:  true,
			expectDirectC:   true,
			expectProxyPA:   true,
		},
		{
			name:           "proxy_url_authorization_same_site_proxy_to_direct_safe",
			initialURL:     "http://" + originHost + "/start",
			redirectURL:    "http://" + sameSiteHost + "/sink",
			directHost:     sameSiteHost,
			proxyURLAuth:   true,
			expectDirectOA: true,
			expectDirectC:  true,
			expectProxyPA:  true,
		},
		{
			name:            "caller_proxy_authorization_cross_site_proxy_to_direct_stripped",
			initialURL:      "http://" + originHost + "/start",
			redirectURL:     "http://" + crossSiteHost + "/sink",
			directHost:      crossSiteHost,
			headerProxyAuth: true,
			expectProxyPA:   true,
		},
		{
			name:           "proxy_url_authorization_cross_site_proxy_to_direct_safe",
			initialURL:     "http://" + originHost + "/start",
			redirectURL:    "http://" + crossSiteHost + "/sink",
			directHost:     crossSiteHost,
			proxyURLAuth:   true,
			expectProxyPA:  true,
			expectDirectOA: false,
			expectDirectC:  false,
		},
	}

	for _, sc := range scenarios {
		if err := runScenario(sc); err != nil {
			failures++
			fmt.Printf("SCENARIO %s FAIL: %v\n", sc.name, err)
		}
	}

	fmt.Printf("overall=%s failures=%d scenarios=%d\n", passFail(failures == 0), failures, len(scenarios)+1)
	if failures != 0 {
		os.Exit(1)
	}
}

func runEnvNoProxyScenario() error {
	var directCap *capture
	direct := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		directCap = captureRequest("direct", r)
		w.Header().Set("X-Direct-Reached", "true")
		_, _ = io.WriteString(w, "direct-ok")
	}))
	defer direct.Close()

	var proxyCaps []capture
	proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		proxyCaps = append(proxyCaps, *captureRequest("proxy", r))
		http.Redirect(w, r, "http://"+sameSiteHost+"/sink", http.StatusFound)
	}))
	defer proxy.Close()

	for _, k := range []string{"HTTP_PROXY", "http_proxy", "HTTPS_PROXY", "https_proxy", "NO_PROXY", "no_proxy", "REQUEST_METHOD"} {
		_ = os.Unsetenv(k)
	}
	_ = os.Setenv("HTTP_PROXY", proxy.URL)
	_ = os.Setenv("NO_PROXY", sameSiteHost)

	directURL, err := url.Parse(direct.URL)
	if err != nil {
		return err
	}
	directAddr := directURL.Host
	proxyURL, err := url.Parse(proxy.URL)
	if err != nil {
		return err
	}
	proxyAddr := proxyURL.Host

	tr := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		DialContext: func(ctx context.Context, network, address string) (net.Conn, error) {
			host, _, err := net.SplitHostPort(address)
			if err != nil {
				return nil, err
			}
			switch {
			case address == proxyAddr:
				return (&net.Dialer{Timeout: 5 * time.Second}).DialContext(ctx, network, address)
			case host == sameSiteHost:
				return (&net.Dialer{Timeout: 5 * time.Second}).DialContext(ctx, network, directAddr)
			default:
				return nil, fmt.Errorf("unexpected dial address %q", address)
			}
		},
	}
	defer tr.CloseIdleConnections()

	req, err := http.NewRequest(http.MethodGet, "http://"+originHost+"/start", nil)
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", originAuthSecret)
	req.Header.Set("Cookie", cookieSecret)
	req.Header.Set("Proxy-Authorization", callerProxySecret)

	resp, err := (&http.Client{Transport: tr, Timeout: 10 * time.Second}).Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, resp.Body)

	if len(proxyCaps) != 1 {
		return fmt.Errorf("proxy request count=%d want 1", len(proxyCaps))
	}
	if directCap == nil {
		return fmt.Errorf("direct server was not reached")
	}
	if directCap.Headers.Get("Proxy-Authorization") == "" {
		return fmt.Errorf("direct server did not receive Proxy-Authorization\nproxy=%s\ndirect=%s", formatCapture(proxyCaps[0]), formatCapture(*directCap))
	}

	fmt.Printf("SCENARIO env_http_proxy_no_proxy_same_site_leaks PASS\n")
	fmt.Printf("  env=HTTP_PROXY=%s NO_PROXY=%s\n", proxy.URL, sameSiteHost)
	fmt.Printf("  proxy_request=%s\n", formatCapture(proxyCaps[0]))
	fmt.Printf("  direct_request=%s\n", formatCapture(*directCap))
	return nil
}

func runScenario(sc scenario) error {
	var directCap *capture
	direct := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		directCap = captureRequest("direct", r)
		w.Header().Set("X-Direct-Reached", "true")
		_, _ = io.WriteString(w, "direct-ok")
	}))
	defer direct.Close()

	var proxyCaps []capture
	proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		proxyCaps = append(proxyCaps, *captureRequest("proxy", r))
		http.Redirect(w, r, sc.redirectURL, http.StatusFound)
	}))
	defer proxy.Close()

	proxyURL, err := url.Parse(proxy.URL)
	if err != nil {
		return err
	}
	if sc.proxyURLAuth {
		proxyURL.User = url.UserPassword(proxyURLUser, proxyURLPass)
	}
	directURL, err := url.Parse(direct.URL)
	if err != nil {
		return err
	}
	directAddr := directURL.Host
	proxyAddr := proxyURL.Host

	tr := &http.Transport{
		Proxy: func(req *http.Request) (*url.URL, error) {
			switch req.URL.Hostname() {
			case originHost:
				return proxyURL, nil
			case sc.directHost:
				return nil, nil
			default:
				return nil, fmt.Errorf("unexpected proxy decision host %q", req.URL.Host)
			}
		},
		DialContext: func(ctx context.Context, network, address string) (net.Conn, error) {
			host, _, err := net.SplitHostPort(address)
			if err != nil {
				return nil, err
			}
			switch {
			case address == proxyAddr:
				return (&net.Dialer{Timeout: 5 * time.Second}).DialContext(ctx, network, address)
			case host == sc.directHost:
				return (&net.Dialer{Timeout: 5 * time.Second}).DialContext(ctx, network, directAddr)
			default:
				return nil, fmt.Errorf("unexpected dial address %q", address)
			}
		},
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	defer tr.CloseIdleConnections()

	client := &http.Client{
		Transport: tr,
		Timeout:   10 * time.Second,
	}
	req, err := http.NewRequest(http.MethodGet, sc.initialURL, nil)
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", originAuthSecret)
	req.Header.Set("Cookie", cookieSecret)
	if sc.headerProxyAuth {
		req.Header.Set("Proxy-Authorization", callerProxySecret)
	}

	resp, err := client.Do(req)
	if err != nil {
		if sc.expectError {
			fmt.Printf("SCENARIO %s PASS: expected error=%v\n", sc.name, err)
			return nil
		}
		return err
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, resp.Body)
	if sc.expectError {
		return fmt.Errorf("expected error but got status %s", resp.Status)
	}

	if len(proxyCaps) != 1 {
		return fmt.Errorf("proxy request count=%d want 1", len(proxyCaps))
	}
	if directCap == nil {
		return fmt.Errorf("direct server was not reached")
	}

	proxyPA := proxyCaps[0].Headers.Get("Proxy-Authorization")
	directPA := directCap.Headers.Get("Proxy-Authorization")
	directOA := directCap.Headers.Get("Authorization")
	directCookie := directCap.Headers.Get("Cookie")
	checks := []struct {
		name string
		got  bool
		want bool
		val  string
	}{
		{"proxy_has_proxy_authorization", proxyPA != "", sc.expectProxyPA, proxyPA},
		{"direct_has_proxy_authorization", directPA != "", sc.expectDirectPA, directPA},
		{"direct_has_origin_authorization", directOA != "", sc.expectDirectOA, directOA},
		{"direct_has_cookie", directCookie != "", sc.expectDirectC, directCookie},
	}
	for _, check := range checks {
		if check.got != check.want {
			return fmt.Errorf("%s=%v want %v value=%q\nproxy=%s\ndirect=%s", check.name, check.got, check.want, redact(check.val), formatCapture(proxyCaps[0]), formatCapture(*directCap))
		}
	}

	fmt.Printf("SCENARIO %s PASS\n", sc.name)
	fmt.Printf("  proxy_request=%s\n", formatCapture(proxyCaps[0]))
	fmt.Printf("  direct_request=%s\n", formatCapture(*directCap))
	return nil
}

func captureRequest(label string, r *http.Request) *capture {
	return &capture{
		Label:      label,
		Method:     r.Method,
		RequestURI: r.RequestURI,
		Host:       r.Host,
		Headers:    r.Header.Clone(),
	}
}

func formatCapture(c capture) string {
	keys := make([]string, 0, len(c.Headers))
	for k := range c.Headers {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	parts := []string{
		fmt.Sprintf("label=%s", c.Label),
		fmt.Sprintf("method=%s", c.Method),
		fmt.Sprintf("requestURI=%q", c.RequestURI),
		fmt.Sprintf("host=%q", c.Host),
	}
	for _, k := range keys {
		if strings.EqualFold(k, "User-Agent") || strings.EqualFold(k, "Accept-Encoding") {
			continue
		}
		parts = append(parts, fmt.Sprintf("%s=%q", k, redact(strings.Join(c.Headers.Values(k), ","))))
	}
	return strings.Join(parts, " ")
}

func redact(s string) string {
	replacer := strings.NewReplacer(
		originAuthSecret, "Bearer ORIGIN_SECRET_FROM_INITIAL_REQUEST",
		cookieSecret, "session=COOKIE_SECRET_FROM_INITIAL_REQUEST",
		callerProxySecret, "Basic CALLER_SUPPLIED_PROXY_SECRET",
	)
	return replacer.Replace(s)
}

func passFail(ok bool) string {
	if ok {
		return "PASS"
	}
	return "FAIL"
}