package main

import (
	"bufio"
	"context"
	"crypto/tls"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"time"
)

const browserProxyHost = "docker-gateway"

type memoryProxyRegistry struct {
	mu      sync.Mutex
	proxies map[string]*memoryProxy
}

type memoryProxy struct {
	mu             sync.RWMutex
	exit           gatewayProxyExit
	bindingVersion int64
	runtimeID      string
	networkID      string
	bindHost       string
	listener       net.Listener
	server         *http.Server
	tunnels        map[net.Conn]net.Conn
	url            string
}

func newMemoryProxyRegistry() *memoryProxyRegistry {
	return &memoryProxyRegistry{proxies: map[string]*memoryProxy{}}
}

func (registry *memoryProxyRegistry) configure(alias string, bindingVersion int64, bindHost string, port int, exit gatewayProxyExit,
	networkIDs ...string) (string, func(), error) {
	registry.mu.Lock()
	defer registry.mu.Unlock()
	networkID := ""
	if len(networkIDs) == 1 {
		networkID = networkIDs[0]
	}
	if proxy := registry.proxies[alias]; proxy != nil {
		if proxy.bindingVersion == bindingVersion && proxy.bindHost == bindHost &&
			(port == 0 || proxy.listener.Addr().(*net.TCPAddr).Port == port) && proxy.exit == exit && proxy.networkID == networkID {
			return proxy.url, func() { registry.removeObject(alias, proxy) }, nil
		}
		delete(registry.proxies, alias)
		// A gateway replica may still hold the previous runtime generation. Close
		// it before rebinding a restored listener, while removeObject's identity
		// check keeps old cleanup callbacks from deleting the replacement.
		closeMemoryProxy(proxy)
	}
	listener, err := net.Listen("tcp4", net.JoinHostPort(bindHost, strconv.Itoa(port)))
	if err != nil {
		return "", nil, err
	}
	actualPort := listener.Addr().(*net.TCPAddr).Port
	proxy := &memoryProxy{exit: exit, bindingVersion: bindingVersion, networkID: networkID, bindHost: bindHost, listener: listener,
		tunnels: make(map[net.Conn]net.Conn), url: "http://" + net.JoinHostPort(browserProxyHost, strconv.Itoa(actualPort))}
	proxy.server = &http.Server{Handler: proxy, ReadHeaderTimeout: 10 * time.Second, IdleTimeout: 60 * time.Second}
	registry.proxies[alias] = proxy
	go func() { _ = proxy.server.Serve(listener) }()
	undo := func() {
		registry.removeObject(alias, proxy)
	}
	return proxy.url, undo, nil
}

func (registry *memoryProxyRegistry) removeObject(alias string, proxy *memoryProxy) {
	registry.mu.Lock()
	if registry.proxies[alias] == proxy {
		delete(registry.proxies, alias)
	} else {
		proxy = nil
	}
	registry.mu.Unlock()
	if proxy != nil {
		closeMemoryProxy(proxy)
	}
}

func closeMemoryProxy(proxy *memoryProxy) {
	proxy.mu.Lock()
	tunnels := proxy.tunnels
	proxy.tunnels = nil
	proxy.mu.Unlock()
	_ = proxy.listener.Close()
	_ = proxy.server.Close()
	for client, upstream := range tunnels {
		_ = client.Close()
		_ = upstream.Close()
	}
}

func (registry *memoryProxyRegistry) ready(alias string, port int, runtimeID string, networkIDs ...string) bool {
	registry.mu.Lock()
	defer registry.mu.Unlock()
	proxy := registry.proxies[alias]
	return proxy != nil && runtimeID != "" && proxy.runtimeID == runtimeID && proxy.listener.Addr().(*net.TCPAddr).Port == port &&
		(len(networkIDs) == 0 || proxy.networkID == networkIDs[0])
}

func (registry *memoryProxyRegistry) bind(alias string, bindingVersion int64, proxyURL, runtimeID string, networkIDs ...string) bool {
	registry.mu.Lock()
	defer registry.mu.Unlock()
	proxy := registry.proxies[alias]
	if proxy == nil || runtimeID == "" || proxy.bindingVersion != bindingVersion || proxy.url != proxyURL {
		return false
	}
	proxy.runtimeID = runtimeID
	if len(networkIDs) == 1 {
		proxy.networkID = networkIDs[0]
	}
	return true
}

func (registry *memoryProxyRegistry) remove(alias string, bindingVersion int64, runtimeID string, networkIDs ...string) bool {
	registry.mu.Lock()
	proxy := registry.proxies[alias]
	if proxy != nil && runtimeID != "" && proxy.bindingVersion == bindingVersion && proxy.runtimeID == runtimeID &&
		(len(networkIDs) == 0 || proxy.networkID == networkIDs[0]) {
		delete(registry.proxies, alias)
	} else if proxy != nil {
		registry.mu.Unlock()
		return false
	} else {
		proxy = nil
	}
	registry.mu.Unlock()
	if proxy != nil {
		closeMemoryProxy(proxy)
	}
	return true
}

func (proxy *memoryProxy) ServeHTTP(response http.ResponseWriter, request *http.Request) {
	if request.Method == http.MethodConnect {
		proxy.tunnel(response, request)
		return
	}
	proxy.mu.RLock()
	exit := proxy.exit
	proxy.mu.RUnlock()
	transport := &http.Transport{DisableKeepAlives: true}
	if exit.Protocol == "http" || exit.Protocol == "https" {
		upstream := &url.URL{Scheme: exit.Protocol, Host: net.JoinHostPort(exit.Host, strconv.Itoa(exit.Port))}
		if exit.Username != "" {
			upstream.User = url.UserPassword(exit.Username, exit.Password)
		}
		transport.Proxy = http.ProxyURL(upstream)
	} else {
		transport.DialContext = proxy.dialContext
	}
	defer transport.CloseIdleConnections()
	outbound := request.Clone(request.Context())
	outbound.RequestURI = ""
	outbound.Header.Del("Proxy-Authorization")
	result, err := transport.RoundTrip(outbound)
	if err != nil {
		http.Error(response, "proxy connection failed", http.StatusBadGateway)
		return
	}
	defer result.Body.Close()
	for key, values := range result.Header {
		for _, value := range values {
			response.Header().Add(key, value)
		}
	}
	response.WriteHeader(result.StatusCode)
	_, _ = io.Copy(response, result.Body)
}

func (proxy *memoryProxy) tunnel(response http.ResponseWriter, request *http.Request) {
	upstream, err := proxy.dialContext(request.Context(), "tcp", request.Host)
	if err != nil {
		http.Error(response, "proxy connection failed", http.StatusBadGateway)
		return
	}
	client, buffered, err := http.NewResponseController(response).Hijack()
	if err != nil {
		_ = upstream.Close()
		http.Error(response, "proxy tunnel unavailable", http.StatusInternalServerError)
		return
	}
	if _, err := buffered.WriteString("HTTP/1.1 200 Connection Established\r\n\r\n"); err != nil || buffered.Flush() != nil {
		_ = client.Close()
		_ = upstream.Close()
		return
	}
	proxy.mu.Lock()
	if proxy.tunnels == nil {
		proxy.mu.Unlock()
		_ = client.Close()
		_ = upstream.Close()
		return
	}
	proxy.tunnels[client] = upstream
	proxy.mu.Unlock()
	defer func() {
		proxy.mu.Lock()
		delete(proxy.tunnels, client)
		proxy.mu.Unlock()
		_ = client.Close()
		_ = upstream.Close()
	}()
	done := make(chan struct{}, 2)
	go func() { _, _ = io.Copy(upstream, client); done <- struct{}{} }()
	go func() { _, _ = io.Copy(client, upstream); done <- struct{}{} }()
	<-done
}

func (proxy *memoryProxy) dialContext(ctx context.Context, _, target string) (net.Conn, error) {
	ctx, cancel := context.WithTimeout(ctx, 20*time.Second)
	defer cancel()
	proxy.mu.RLock()
	exit := proxy.exit
	proxy.mu.RUnlock()
	switch exit.Protocol {
	case "http", "https":
		return dialHTTPProxy(ctx, exit, target)
	case "socks4":
		return dialSOCKS4Proxy(ctx, exit, target)
	case "socks5":
		return dialSOCKS5Proxy(ctx, exit, target)
	default:
		return nil, errors.New("unsupported proxy protocol")
	}
}

func dialHTTPProxy(ctx context.Context, exit gatewayProxyExit, target string) (net.Conn, error) {
	address := net.JoinHostPort(exit.Host, strconv.Itoa(exit.Port))
	connection, err := (&net.Dialer{}).DialContext(ctx, "tcp", address)
	if err != nil {
		return nil, err
	}
	if exit.Protocol == "https" {
		tlsConnection := tls.Client(connection, &tls.Config{ServerName: exit.Host, MinVersion: tls.VersionTLS12})
		if err := tlsConnection.HandshakeContext(ctx); err != nil {
			_ = connection.Close()
			return nil, err
		}
		connection = tlsConnection
	}
	request := &http.Request{Method: http.MethodConnect, URL: &url.URL{Opaque: target}, Host: target, Header: make(http.Header)}
	if exit.Username != "" {
		request.Header.Set("Proxy-Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(exit.Username+":"+exit.Password)))
	}
	if deadline, ok := ctx.Deadline(); ok {
		_ = connection.SetDeadline(deadline)
	}
	if err := request.Write(connection); err != nil {
		_ = connection.Close()
		return nil, err
	}
	result, err := http.ReadResponse(bufio.NewReader(connection), request)
	if err != nil {
		_ = connection.Close()
		return nil, err
	}
	if result.StatusCode != http.StatusOK {
		_ = result.Body.Close()
		_ = connection.Close()
		return nil, fmt.Errorf("upstream proxy returned %s", result.Status)
	}
	_ = connection.SetDeadline(time.Time{})
	return connection, nil
}

func dialSOCKS4Proxy(ctx context.Context, exit gatewayProxyExit, target string) (net.Conn, error) {
	connection, err := (&net.Dialer{}).DialContext(ctx, "tcp", net.JoinHostPort(exit.Host, strconv.Itoa(exit.Port)))
	if err != nil {
		return nil, err
	}
	host, portText, err := net.SplitHostPort(target)
	if err != nil {
		_ = connection.Close()
		return nil, err
	}
	port, err := strconv.Atoi(portText)
	if err != nil || port < 1 || port > 65535 {
		_ = connection.Close()
		return nil, errors.New("invalid SOCKS4 target")
	}
	payload := []byte{4, 1, byte(port >> 8), byte(port), 0, 0, 0, 1}
	if ip := net.ParseIP(host).To4(); ip != nil {
		copy(payload[4:8], ip)
	}
	payload = append(payload, exit.Username...)
	payload = append(payload, 0)
	if net.ParseIP(host).To4() == nil {
		payload = append(payload, host...)
		payload = append(payload, 0)
	}
	if err := exchangeSOCKS(ctx, connection, payload, 8); err != nil {
		return nil, err
	}
	return connection, nil
}

func dialSOCKS5Proxy(ctx context.Context, exit gatewayProxyExit, target string) (net.Conn, error) {
	connection, err := (&net.Dialer{}).DialContext(ctx, "tcp", net.JoinHostPort(exit.Host, strconv.Itoa(exit.Port)))
	if err != nil {
		return nil, err
	}
	methods := []byte{5, 1, 0}
	if exit.Username != "" {
		methods = []byte{5, 1, 2}
	}
	if deadline, ok := ctx.Deadline(); ok {
		_ = connection.SetDeadline(deadline)
	}
	if _, err := connection.Write(methods); err != nil {
		_ = connection.Close()
		return nil, err
	}
	selection := make([]byte, 2)
	if _, err := io.ReadFull(connection, selection); err != nil || selection[0] != 5 || selection[1] == 0xff {
		_ = connection.Close()
		return nil, errors.New("SOCKS5 authentication method rejected")
	}
	if selection[1] == 2 {
		if len(exit.Username) > 255 || len(exit.Password) > 255 {
			_ = connection.Close()
			return nil, errors.New("SOCKS5 credentials too long")
		}
		auth := append([]byte{1, byte(len(exit.Username))}, exit.Username...)
		auth = append(auth, byte(len(exit.Password)))
		auth = append(auth, exit.Password...)
		if _, err := connection.Write(auth); err != nil {
			_ = connection.Close()
			return nil, err
		}
		result := make([]byte, 2)
		if _, err := io.ReadFull(connection, result); err != nil || result[1] != 0 {
			_ = connection.Close()
			return nil, errors.New("SOCKS5 authentication rejected")
		}
	} else if exit.Username != "" {
		_ = connection.Close()
		return nil, errors.New("SOCKS5 proxy skipped required authentication")
	}
	host, portText, err := net.SplitHostPort(target)
	if err != nil {
		_ = connection.Close()
		return nil, err
	}
	port, err := strconv.Atoi(portText)
	if err != nil || port < 1 || port > 65535 {
		_ = connection.Close()
		return nil, errors.New("invalid SOCKS5 target")
	}
	request := []byte{5, 1, 0}
	if ip := net.ParseIP(host); ip != nil && ip.To4() != nil {
		request = append(request, 1)
		request = append(request, ip.To4()...)
	} else if ip != nil {
		request = append(request, 4)
		request = append(request, ip.To16()...)
	} else {
		if len(host) > 255 {
			_ = connection.Close()
			return nil, errors.New("SOCKS5 target too long")
		}
		request = append(request, 3, byte(len(host)))
		request = append(request, host...)
	}
	portBytes := make([]byte, 2)
	binary.BigEndian.PutUint16(portBytes, uint16(port))
	request = append(request, portBytes...)
	if _, err := connection.Write(request); err != nil {
		_ = connection.Close()
		return nil, err
	}
	header := make([]byte, 4)
	if _, err := io.ReadFull(connection, header); err != nil || header[0] != 5 || header[1] != 0 {
		_ = connection.Close()
		return nil, errors.New("SOCKS5 proxy rejected connection")
	}
	addressLength := 0
	switch header[3] {
	case 1:
		addressLength = 4
	case 4:
		addressLength = 16
	case 3:
		var length [1]byte
		if _, err := io.ReadFull(connection, length[:]); err != nil {
			_ = connection.Close()
			return nil, err
		}
		addressLength = int(length[0])
	default:
		_ = connection.Close()
		return nil, errors.New("invalid SOCKS5 response")
	}
	if _, err := io.CopyN(io.Discard, connection, int64(addressLength+2)); err != nil {
		_ = connection.Close()
		return nil, err
	}
	_ = connection.SetDeadline(time.Time{})
	return connection, nil
}

func exchangeSOCKS(ctx context.Context, connection net.Conn, request []byte, responseBytes int) error {
	if deadline, ok := ctx.Deadline(); ok {
		_ = connection.SetDeadline(deadline)
	}
	if _, err := connection.Write(request); err != nil {
		_ = connection.Close()
		return err
	}
	if responseBytes > 0 {
		response := make([]byte, responseBytes)
		if _, err := io.ReadFull(connection, response); err != nil {
			_ = connection.Close()
			return err
		}
		if responseBytes == 8 && response[1] != 90 {
			_ = connection.Close()
			return errors.New("SOCKS4 proxy rejected connection")
		}
	}
	_ = connection.SetDeadline(time.Time{})
	return nil
}