package http2

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"errors"
	"math/big"
	"net"
	"net/http"
	"sync"
	"testing"
	"time"

	xhttp2 "golang.org/x/net/http2"
)

func TestServerHandlesH2CRequest(t *testing.T) {
	server := New(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if request.ProtoMajor != 2 {
			http.Error(writer, "not HTTP/2", http.StatusInternalServerError)
			return
		}
		_, _ = writer.Write([]byte("ok"))
	}))
	address, stop := startServer(t, server)
	defer stop()

	transport := &xhttp2.Transport{
		AllowHTTP: true,
		DialTLSContext: func(ctx context.Context, network, address string, _ *tls.Config) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, network, address)
		},
	}
	client := &http.Client{Transport: transport, Timeout: 2 * time.Second}
	defer transport.CloseIdleConnections()
	response, err := client.Get("http://" + address + "/health")
	if err != nil {
		t.Fatal(err)
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		t.Fatalf("unexpected status: %d", response.StatusCode)
	}
	if stats := server.Stats(); stats.AcceptedConnections != 1 {
		t.Fatalf("unexpected HTTP/2 stats: %+v", stats)
	}
}

func TestServerEnforcesGlobalRequestLimitAndRecovers(t *testing.T) {
	started := make(chan struct{})
	release := make(chan struct{})
	var once sync.Once
	server := New(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if request.URL.Path == "/hold" {
			once.Do(func() { close(started) })
			<-release
		}
		_, _ = writer.Write([]byte("ok"))
	}))
	server.MaxConcurrentRequests = 1
	address, stop := startServer(t, server)
	defer stop()
	transport := &xhttp2.Transport{
		AllowHTTP: true,
		DialTLSContext: func(ctx context.Context, network, address string, _ *tls.Config) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, network, address)
		},
	}
	client := &http.Client{Transport: transport, Timeout: 3 * time.Second}
	defer transport.CloseIdleConnections()

	firstDone := make(chan *http.Response, 1)
	firstErr := make(chan error, 1)
	go func() {
		response, err := client.Get("http://" + address + "/hold")
		if err != nil {
			firstErr <- err
			return
		}
		firstDone <- response
	}()
	select {
	case <-started:
	case err := <-firstErr:
		t.Fatal(err)
	case <-time.After(time.Second):
		t.Fatal("first HTTP/2 stream did not start")
	}

	response, err := client.Get("http://" + address + "/busy")
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("expected 503 at stream limit, got %d", response.StatusCode)
	}
	if stats := server.Stats(); stats.RejectedRequests != 1 {
		t.Fatalf("expected one rejected stream, got %+v", stats)
	}

	close(release)
	select {
	case response = <-firstDone:
		response.Body.Close()
	case err := <-firstErr:
		t.Fatal(err)
	case <-time.After(time.Second):
		t.Fatal("held HTTP/2 stream did not finish")
	}
	response, err = client.Get("http://" + address + "/recovered")
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusOK {
		t.Fatalf("request limit did not recover, got %d", response.StatusCode)
	}
}

func TestServerTLSUsesH2ALPN(t *testing.T) {
	certificate, roots := testCertificate(t)
	server := New(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if request.TLS == nil || request.TLS.NegotiatedProtocol != "h2" {
			http.Error(writer, "missing h2 ALPN", http.StatusInternalServerError)
			return
		}
		writer.WriteHeader(http.StatusNoContent)
	}))
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() {
		done <- server.ServeTLS(context.Background(), listener, &tls.Config{
			Certificates: []tls.Certificate{certificate}, MinVersion: tls.VersionTLS12,
		})
	}()
	stop := func() {
		t.Helper()
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		if err := server.Shutdown(ctx); err != nil && !errors.Is(err, context.DeadlineExceeded) {
			t.Errorf("shutdown: %v", err)
		}
		if err := <-done; err != nil {
			t.Errorf("serve: %v", err)
		}
	}
	defer stop()

	transport := &xhttp2.Transport{TLSClientConfig: &tls.Config{
		RootCAs: roots, ServerName: "localhost", MinVersion: tls.VersionTLS12,
		NextProtos: []string{"h2"},
	}}
	client := &http.Client{Transport: transport, Timeout: 2 * time.Second}
	defer transport.CloseIdleConnections()
	response, err := client.Get("https://" + listener.Addr().String() + "/health")
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusNoContent {
		t.Fatalf("unexpected TLS status: %d", response.StatusCode)
	}
}

func startServer(t *testing.T, server *Server) (string, func()) {
	t.Helper()
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() { done <- server.Serve(context.Background(), listener) }()
	stop := func() {
		t.Helper()
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		if err := server.Shutdown(ctx); err != nil && !errors.Is(err, context.DeadlineExceeded) {
			t.Errorf("shutdown: %v", err)
		}
		if err := <-done; err != nil {
			t.Errorf("serve: %v", err)
		}
	}
	return listener.Addr().String(), stop
}

func testCertificate(t *testing.T) (tls.Certificate, *x509.CertPool) {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 120))
	if err != nil {
		t.Fatal(err)
	}
	template := &x509.Certificate{
		SerialNumber: serial,
		Subject:      pkix.Name{CommonName: "localhost"},
		DNSNames:     []string{"localhost"},
		NotBefore:    time.Now().Add(-time.Minute),
		NotAfter:     time.Now().Add(time.Hour),
		KeyUsage:     x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}
	der, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
	if err != nil {
		t.Fatal(err)
	}
	leaf, err := x509.ParseCertificate(der)
	if err != nil {
		t.Fatal(err)
	}
	certificate := tls.Certificate{
		Certificate: [][]byte{der},
		PrivateKey:  key,
	}
	certificate.Leaf = leaf
	roots := x509.NewCertPool()
	roots.AddCert(leaf)
	return certificate, roots
}