NixMsg/cmd/nixmsg/healthcheck_test.go

package main

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"math/big"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"git.asio.asia/nixevol/NixMsg/internal/config"
)

func TestResolveHealthAddrFromListenAddrFile(t *testing.T) {
	dir := t.TempDir()
	if err := os.WriteFile(filepath.Join(dir, "listen.addr"), []byte("0.0.0.0:19100\n"), 0o644); err != nil {
		t.Fatal(err)
	}
	got, err := resolveHealthAddr(dir, ":7443")
	if err != nil {
		t.Fatal(err)
	}
	if got != "127.0.0.1:19100" {
		t.Fatalf("got %q", got)
	}
}

func TestResolveHealthAddrFromListen(t *testing.T) {
	got, err := resolveHealthAddr(t.TempDir(), ":9555")
	if err != nil {
		t.Fatal(err)
	}
	if got != "127.0.0.1:9555" {
		t.Fatalf("got %q", got)
	}
}

func TestCmdHealthcheckOK(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	}))
	t.Cleanup(srv.Close)

	dir := t.TempDir()
	hostPort := srv.Listener.Addr().String()
	if err := os.WriteFile(filepath.Join(dir, "listen.addr"), []byte(hostPort+"\n"), 0o644); err != nil {
		t.Fatal(err)
	}
	cfgPath := filepath.Join(dir, "config.yaml")
	if err := os.WriteFile(cfgPath, []byte("listen: \":0\"\ndata_dir: \""+filepath.ToSlash(dir)+"\"\n"), 0o644); err != nil {
		t.Fatal(err)
	}
	t.Setenv("NIXMSG_CONFIG", cfgPath)
	if err := cmdHealthcheck(nil); err != nil {
		t.Fatal(err)
	}
}

func TestCmdHealthcheckWithTLS(t *testing.T) {
	dataDir := t.TempDir()
	certPath, keyPath := writeServeTestCert(t, dataDir, "hc")
	initAdminForTest(t, dataDir)
	cfgYAML := "listen: \"127.0.0.1:0\"\ndata_dir: \"" + filepath.ToSlash(dataDir) + "\"\n" +
		"tls:\n  cert_file: \"" + filepath.ToSlash(certPath) + "\"\n" +
		"  key_file: \"" + filepath.ToSlash(keyPath) + "\"\n" +
		"  allow_plaintext: false\nlog:\n  level: error\n"
	cfgPath := filepath.Join(dataDir, "config.yaml")
	if err := os.WriteFile(cfgPath, []byte(cfgYAML), 0o644); err != nil {
		t.Fatal(err)
	}
	cfg, err := config.Load(cfgPath)
	if err != nil {
		t.Fatal(err)
	}
	if err := cfg.Validate(); err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	errCh := make(chan error, 1)
	go func() { errCh <- runServe(ctx, cfg) }()

	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		b, readErr := os.ReadFile(filepath.Join(dataDir, "listen.addr"))
		if readErr == nil && strings.TrimSpace(string(b)) != "" {
			break
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Setenv("NIXMSG_CONFIG", cfgPath)
	if err := cmdHealthcheck(nil); err != nil {
		t.Fatal(err)
	}
	cancel()
	select {
	case err := <-errCh:
		if err != nil {
			t.Fatalf("serve exit: %v", err)
		}
	case <-time.After(10 * time.Second):
		t.Fatal("serve did not stop")
	}
}

func writeServeTestCert(t *testing.T, dir, cn string) (certPath, keyPath string) {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	tmpl := &x509.Certificate{
		SerialNumber: big.NewInt(time.Now().UnixNano()),
		Subject:      pkix.Name{CommonName: cn},
		NotBefore:    time.Now().Add(-time.Hour),
		NotAfter:     time.Now().Add(24 * time.Hour),
		KeyUsage:     x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		DNSNames:     []string{"localhost"},
		IPAddresses:  []net.IP{net.ParseIP("127.0.0.1")},
	}
	der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
	if err != nil {
		t.Fatal(err)
	}
	certPath = filepath.Join(dir, cn+".crt")
	keyPath = filepath.Join(dir, cn+".key")
	certOut, err := os.Create(certPath)
	if err != nil {
		t.Fatal(err)
	}
	_ = pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: der})
	_ = certOut.Close()
	keyOut, err := os.Create(keyPath)
	if err != nil {
		t.Fatal(err)
	}
	b, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		t.Fatal(err)
	}
	_ = pem.Encode(keyOut, &pem.Block{Type: "EC PRIVATE KEY", Bytes: b})
	_ = keyOut.Close()
	return certPath, keyPath
}