// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package nfs4server

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"math/big"
	"net"
	"testing"
	"time"

	"sourcedock.dev/petrbalvin/nfs/internal/nfs4"
	"sourcedock.dev/petrbalvin/nfs/internal/nfsclient"
	"sourcedock.dev/petrbalvin/nfs/internal/server"
)

// testCertificate builds a self signed certificate for the TLS tests.
func testCertificate(t *testing.T) *tls.Certificate {
	t.Helper()
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	tmpl := x509.Certificate{
		SerialNumber:          big.NewInt(1),
		Subject:               pkix.Name{CommonName: "nfs.test"},
		NotBefore:             time.Now().Add(-time.Hour),
		NotAfter:              time.Now().Add(time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		DNSNames:              []string{"localhost"},
		IPAddresses:           []net.IP{net.ParseIP("127.0.0.1")},
	}
	der, err := x509.CreateCertificate(rand.Reader, &tmpl, &tmpl, &priv.PublicKey, priv)
	if err != nil {
		t.Fatal(err)
	}
	return &tls.Certificate{Certificate: [][]byte{der}, PrivateKey: priv}
}

// testRootPool trusts the self signed certificate.
func testRootPool(t *testing.T, cert *tls.Certificate) *x509.CertPool {
	t.Helper()
	pool := x509.NewCertPool()
	pool.AddCert(leafOf(t, cert))
	return pool
}

func leafOf(t *testing.T, cert *tls.Certificate) *x509.Certificate {
	t.Helper()
	leaf, err := x509.ParseCertificate(cert.Certificate[0])
	if err != nil {
		t.Fatal(err)
	}
	return leaf
}

// The AUTH_TLS probe upgrades the connection and every further call
// rides the encrypted channel; a probe inside an upgraded session and a
// probe on a non NULL procedure are refused.
func TestRPCWithTLS(t *testing.T) {
	h := testTree(t)
	cert := testCertificate(t)
	h.TLSConfig = &tls.Config{Certificates: []tls.Certificate{*cert}}

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	srv := &server.Server{Handle: h.HandleConn}
	go srv.Serve(t.Context(), ln)
	defer ln.Close()

	cfg := &tls.Config{ServerName: "localhost", RootCAs: testRootPool(t, cert)}
	cl, err := nfsclient.Dial(ln.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer cl.Close()
	if err := cl.EnableTLS(cfg); err != nil {
		t.Fatalf("enable tls: %v", err)
	}

	// The session and the compounds ride the encrypted channel.
	if err := cl.Establish("tls"); err != nil {
		t.Fatalf("establish: %v", err)
	}
	res, _, err := cl.Compound("tls", [][]byte{
		nfs4.AppendPutRootfh(nil),
		nfs4.AppendLookup(nil, "a.txt"),
		nfs4.AppendGetattr(nil, nfs4.OfBits(nfs4.AttrSize)),
	})
	if err != nil || res.Status != nfs4.ErrOK {
		t.Fatalf("compound: status %d %v", res.Status, err)
	}

	// A second probe inside the upgraded session is refused: the client
	// reports the rejection instead of succeeding silently.
	if err := cl.EnableTLS(cfg); err == nil {
		t.Fatal("second probe inside a tls session succeeded")
	}
}

// A handler without a TLSConfig refuses the probe.
func TestTLSRefusedWithoutConfig(t *testing.T) {
	h := testTree(t)
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	srv := &server.Server{Handle: h.HandleConn}
	go srv.Serve(t.Context(), ln)
	defer ln.Close()

	cl, err := nfsclient.Dial(ln.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer cl.Close()
	if err := cl.EnableTLS(&tls.Config{InsecureSkipVerify: true}); err == nil {
		t.Fatal("the upgrade succeeded without a server side config")
	}
}