package main

import (
	"context"
	"flag"
	"io"
	"math/rand"
	"net"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"
)

var (
	requestTimeout = flag.Duration("request-timeout", 10*time.Millisecond, "Timeout for each HTTP request")
	requestCount   = flag.Int("request-count", 1000, "Number of HTTP requests to send")
	unplugInterval = flag.Int("unplug-interval", 10, "Number of HTTP requests to issue between each unplug event")
)

func TestNetsplit(t *testing.T) {
	testcase := func(
		srvConfig func(srv *httptest.Server),
		rtConfig func(rt *http.Transport),
		test func(t *testing.T, url string, client *http.Client,
			withConns func(fn func(conns []*proxyConn) []*proxyConn),
		),
	) func(t *testing.T) {
		return func(t *testing.T) {
			srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				return
			}))
			srvConfig(srv)
			srv.StartTLS()
			defer srv.Close()

			dialer := &net.Dialer{}

			var (
				pipeMu sync.Mutex
				pipes  []*proxyConn
			)

			var proxyMu sync.Mutex
			proxy, err := net.Listen("tcp", "127.0.0.1:0")
			if err != nil {
				t.Fatalf("proxy listen: %v", err)
			}
			defer proxy.Close()

			connPair := func() (net.Conn, net.Conn, error) {
				proxyMu.Lock()
				defer proxyMu.Unlock()

				c1, err := net.Dial(proxy.Addr().Network(), proxy.Addr().String())
				if err != nil {
					return nil, nil, err
				}

				c2, err := proxy.Accept()
				if err != nil {
					return nil, nil, err
				}

				return c1, c2, nil
			}

			transport := srv.Client().Transport.(*http.Transport).Clone()
			transport.ForceAttemptHTTP2 = true
			transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
				c3, err := dialer.DialContext(ctx, network, addr)
				if err != nil {
					return nil, err
				}

				c1, c2, err := connPair()
				if err != nil {
					return nil, err
				}

				pipe := newProxyConn(c2, c3)

				pipeMu.Lock()
				pipes = append(pipes, pipe)
				pipeMu.Unlock()

				return c1, nil
			}

			client := &http.Client{Transport: transport}

			withConns := func(fn func(conns []*proxyConn) []*proxyConn) {
				pipeMu.Lock()
				defer pipeMu.Unlock()
				pipes = fn(pipes)
			}

			test(t, srv.URL, client, withConns)
		}
	}

	makeRequest := func(ctx context.Context, url string, client *http.Client) error {
		req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
		if err != nil {
			return err
		}

		resp, err := client.Do(req)
		if err != nil {
			return err
		}
		defer resp.Body.Close()

		_, err = io.Copy(io.Discard, resp.Body)
		if err != nil {
			return err
		}

		return nil
	}

	run := func(t *testing.T, url string, client *http.Client, withConns func(func(conns []*proxyConn) []*proxyConn)) {
		var failures int64
		var unplugs int64
		total := *requestCount
		for i := 0; i < total; i++ {
			if interval := *unplugInterval; i%interval == interval/2 { // midway or a bit less
				withConns(func(conns []*proxyConn) []*proxyConn {
					if l := len(conns); l > 0 {
						idx := rand.Intn(l)
						pipe := conns[idx]
						pipe.unplug()
						unplugs++

						conns[idx] = conns[l-1]
						conns = conns[:l-1]
					}
					return conns
				})
			}

			err := func() error {
				ctx := context.Background()
				ctx, cancel := context.WithTimeout(ctx, *requestTimeout)
				defer cancel()
				return makeRequest(ctx, url, client)
			}()

			if err != nil {
				failures++
			}
		}

		t.Errorf("total=%d unplugs=%d failures=%d", total, unplugs, failures)
	}

	t.Run("http1", testcase(
		func(srv *httptest.Server) { srv.EnableHTTP2 = false },
		func(rt *http.Transport) {},
		run,
	))

	t.Run("http2", testcase(
		func(srv *httptest.Server) { srv.EnableHTTP2 = true },
		func(rt *http.Transport) {},
		run,
	))
}

// A proxyConn connects two net.Conns together, copying the data in each
// direction, until it is closed or unplugged. Unplugging a proxyConn causes it
// to silently stop copying the data.
type proxyConn struct {
	c1 net.Conn
	c2 net.Conn

	once sync.Once
}

func newProxyConn(c1, c2 net.Conn) *proxyConn {
	pc := &proxyConn{c1: c1, c2: c2}

	go io.Copy(c1, c2)
	go io.Copy(c2, c1)

	return pc
}

func (pc *proxyConn) unplug() {
	dl := time.Now().Add(-1 * time.Second)
	pc.c1.SetDeadline(dl)
	pc.c2.SetDeadline(dl)
}

func (pc *proxyConn) close() {
	pc.once.Do(func() {
		pc.c1.Close()
		pc.c2.Close()
	})
}