package main

import (
	"context"
	"flag"
	"log"
	"net"
	"net/http"
	"net/url"
	"time"

	"golang.org/x/sync/errgroup"
)

const alphabeticChars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"

var (
	listenAddr = flag.String("listen", ":8081", "address to listen on")
	useProxy   = flag.Bool("use-proxy", true, "whether to use the bad proxy")
	proxyURL   = flag.String("proxy", "http://localhost:8081", "proxy url")
)

const (
	targetURL = "https://example.com"
)

func mainServer() error {
	l, err := net.Listen("tcp", *listenAddr)
	if err != nil {
		return err
	}
	defer l.Close()

	log.Printf("mainServer: listening on %s", *listenAddr)
	for {
		conn, err := l.Accept()
		if err != nil {
			return err
		}
		go handleConnection(conn)
	}
}

func handleConnection(conn net.Conn) {
	log.Printf("Client connected: %s", conn.RemoteAddr().String())
	defer func() {
		log.Printf("Client disconnected: %s", conn.RemoteAddr().String())
		conn.Close()
	}()

	total := 0
	for {
		t, err := conn.Write([]byte(alphabeticChars))
		if err != nil {
			if err.Error() == "write: broken pipe" || err.Error() == "write: connection reset by peer" {
				log.Printf("Client %s disconnected gracefully, totalBytes: %d", conn.RemoteAddr().String(), total)
			} else {
				log.Printf("Error writing to client %s: %v, totalBytes: %d", conn.RemoteAddr().String(), err, total)
			}
			return
		}
		total += t
		time.Sleep(10 * time.Nanosecond)
	}
}

func mainClient() error {
	pu, err := url.Parse(*proxyURL)
	if err != nil {
		return err
	}

	tr := &http.Transport{
		MaxResponseHeaderBytes: 100000, // !! Not honored by proxy CONNECT request.
		Proxy:                  http.ProxyURL(pu),
		DialContext: (&net.Dialer{
			Timeout: 5 * time.Second,
		}).DialContext,
	}

	client := &http.Client{
		Transport: tr,
		Timeout:   600 * time.Second,
	}

	log.Printf("mainClient: Attempting to make a GET request to %s through proxy %s...", targetURL, pu.String())
	resp, err := client.Get(targetURL)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	log.Printf("mainClient: Request status: %d", resp.StatusCode)
	return nil
}

func mainClientDirect() error {
	tr := &http.Transport{
		// Default MaxResponseHeaderBytes of 10 << 20 is honored.
	}

	client := &http.Client{
		Transport: tr,
	}

	log.Printf("mainClientDirect: Attempting to make a direct GET request to proxyURL %s", *proxyURL)
	resp, err := client.Get(*proxyURL)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	log.Printf("mainClientDirect: Request status: %d", resp.StatusCode)
	return nil
}

func main() {
	flag.Parse()
	eg, _ := errgroup.WithContext(context.Background())

	// writes infinite alphabets any all clients.
	eg.Go(mainServer)
	if *useProxy {
		// exhaust memory.
		eg.Go(mainClient)
	} else {
		// direct (non-proxy) request to the bad server.
		eg.Go(mainClientDirect)
	}
	log.Printf("status %#v", eg.Wait())
}