package main

import (
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"
)

const maxBodyBytes = 64 << 20

var hopHeaders = map[string]struct{}{
	"Connection":          {},
	"Keep-Alive":          {},
	"Proxy-Authenticate":  {},
	"Proxy-Authorization": {},
	"Te":                  {},
	"Trailer":             {},
	"Transfer-Encoding":   {},
	"Upgrade":             {},
}

var clientIPHeaders = map[string]struct{}{
	"Cf-Connecting-Ip": {},
	"Forwarded":        {},
	"X-Forwarded-For":  {},
	"X-Real-Ip":        {},
}

type proxy struct {
	tokenHash [32]byte
	client    *http.Client
}

func main() {
	token := os.Getenv("BRIDGE_TOKEN")
	if token == "" {
		log.Fatal("BRIDGE_TOKEN is required")
	}

	addr := env("ADDR", ":8080")
	p := &proxy{
		tokenHash: sha256.Sum256([]byte(token)),
		client: &http.Client{
			Timeout: 60 * time.Second,
			Transport: &http.Transport{
				Proxy:               http.ProxyFromEnvironment,
				DialContext:         (&net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
				ForceAttemptHTTP2:   true,
				DisableCompression:  true,
				MaxIdleConns:        100,
				IdleConnTimeout:     90 * time.Second,
				TLSHandshakeTimeout: 10 * time.Second,
			},
		},
	}

	server := &http.Server{
		Addr:              addr,
		Handler:           p,
		ReadHeaderTimeout: 10 * time.Second,
	}

	go func() {
		log.Printf("bridge proxy listening on %s", addr)
		if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			log.Fatal(err)
		}
	}()

	stop := make(chan os.Signal, 1)
	signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)
	<-stop

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_ = server.Shutdown(ctx)
}

func (p *proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/healthz" {
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok\n"))
		return
	}

	if !p.validToken(strings.Trim(r.URL.Path, "/")) {
		http.Error(w, "not found", http.StatusNotFound)
		return
	}

	target, err := parseTarget(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	resp, err := p.fetch(r, target)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

	copyHeader(w.Header(), resp.Header)
	if resp.ContentLength >= 0 {
		w.Header().Set("Content-Length", strconv.FormatInt(resp.ContentLength, 10))
	}
	w.WriteHeader(resp.StatusCode)
	_, _ = io.Copy(w, resp.Body)
}

func (p *proxy) validToken(got string) bool {
	gotHash := sha256.Sum256([]byte(got))
	return subtle.ConstantTimeCompare(gotHash[:], p.tokenHash[:]) == 1
}

func parseTarget(r *http.Request) (string, error) {
	target := r.URL.Query().Get("url")
	if target == "" {
		return "", errors.New("missing url")
	}

	u, err := url.Parse(target)
	if err != nil || u.Host == "" {
		return "", errors.New("invalid url")
	}
	if u.Scheme != "http" && u.Scheme != "https" {
		return "", errors.New("url must start with http:// or https://")
	}

	return u.String(), nil
}

func (p *proxy) fetch(in *http.Request, target string) (*http.Response, error) {
	body := http.MaxBytesReader(nil, in.Body, maxBodyBytes)
	out, err := http.NewRequestWithContext(in.Context(), in.Method, target, body)
	if err != nil {
		return nil, fmt.Errorf("create upstream request: %w", err)
	}

	copyHeader(out.Header, in.Header)
	stripRequestHeaders(out.Header)

	return p.client.Do(out)
}

func copyHeader(dst, src http.Header) {
	for key, values := range src {
		if _, ok := hopHeaders[http.CanonicalHeaderKey(key)]; ok {
			continue
		}
		for _, value := range values {
			dst.Add(key, value)
		}
	}
}

func stripRequestHeaders(h http.Header) {
	for key := range hopHeaders {
		h.Del(key)
	}
	for key := range clientIPHeaders {
		h.Del(key)
	}
}

func env(key, fallback string) string {
	if value := os.Getenv(key); value != "" {
		return value
	}
	return fallback
}