package http2

import (
	"context"
	"crypto/tls"
	"io"
	"net"
	"net/http"
	"sync/atomic"
	"testing"
	"time"

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

func TestTLSConnectionLimitRecoversAfterHandshakeTimeout(t *testing.T) {
	certificate, roots := testCertificate(t)
	var handled atomic.Int32
	server := New(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		handled.Add(1)
		_, _ = io.WriteString(writer, "ok")
	}))
	server.MaxConnections = 1
	server.TLSHandshakeTimeout = time.Second
	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,
		})
	}()
	t.Cleanup(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		if err := server.Shutdown(ctx); err != nil {
			t.Errorf("shutdown: %v", err)
		}
		_ = listener.Close()
		if err := serveResult(t, done); err != nil {
			t.Errorf("serve: %v", err)
		}
	})

	// Occupy the only connection slot without sending a TLS ClientHello.
	held, err := net.DialTimeout("tcp", listener.Addr().String(), time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer held.Close()
	waitHTTP2Stats(t, server, func(stats Stats) bool { return stats.ActiveConnections == 1 })

	rejected, err := net.DialTimeout("tcp", listener.Addr().String(), time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer rejected.Close()
	if err := rejected.SetReadDeadline(time.Now().Add(250 * time.Millisecond)); err != nil {
		t.Fatal(err)
	}
	var buffer [1]byte
	if _, err := rejected.Read(buffer[:]); err == nil {
		t.Fatal("connection above the cap was not closed")
	} else if timeout, ok := err.(net.Error); ok && timeout.Timeout() {
		t.Fatal("connection above the cap was not rejected promptly")
	}
	if stats := server.Stats(); stats.RejectedConnections != 1 || handled.Load() != 0 {
		t.Fatalf("unexpected admission state: %+v handled=%d", stats, handled.Load())
	}

	// The handshake deadline, rather than a client disconnect, must release it.
	waitHTTP2Stats(t, server, func(stats Stats) bool { return stats.ActiveConnections == 0 })
	if err := held.SetReadDeadline(time.Now().Add(250 * time.Millisecond)); err != nil {
		t.Fatal(err)
	}
	if _, err := held.Read(buffer[:]); err == nil {
		t.Fatal("timed-out TLS connection remained open")
	} else if timeout, ok := err.(net.Error); ok && timeout.Timeout() {
		t.Fatal("timed-out TLS connection was not closed")
	}

	transport := &xhttp2.Transport{TLSClientConfig: &tls.Config{
		RootCAs: roots, ServerName: "localhost", MinVersion: tls.VersionTLS12,
		NextProtos: []string{"h2"},
	}}
	defer transport.CloseIdleConnections()
	client := &http.Client{Transport: transport, Timeout: 3 * time.Second}
	response, err := client.Get("https://" + listener.Addr().String() + "/health")
	if err != nil {
		t.Fatal(err)
	}
	body, readErr := io.ReadAll(response.Body)
	response.Body.Close()
	if readErr != nil || response.StatusCode != http.StatusOK || response.ProtoMajor != 2 || string(body) != "ok" {
		t.Fatalf("TLS admission did not recover: status=%d protocol=%s body=%q err=%v", response.StatusCode, response.Proto, body, readErr)
	}
	transport.CloseIdleConnections()
	waitHTTP2Stats(t, server, func(stats Stats) bool { return stats.ActiveConnections == 0 && stats.ActiveRequests == 0 })
	if stats := server.Stats(); stats.AcceptedConnections != 3 || stats.RejectedConnections != 1 || stats.ProtocolErrors != 0 || handled.Load() != 1 {
		t.Fatalf("unexpected recovered admission state: %+v handled=%d", stats, handled.Load())
	}
}

func waitHTTP2Stats(t *testing.T, server *Server, predicate func(Stats) bool) {
	t.Helper()
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if predicate(server.Stats()) {
			return
		}
		time.Sleep(time.Millisecond)
	}
	t.Fatalf("HTTP/2 stats did not reach expected state: %+v", server.Stats())
}