netbird/client/internal/certproof/collect_test.go

package certproof

import (
	"context"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/netbirdio/netbird/shared/management/certposture"
	"github.com/netbirdio/netbird/shared/management/certposture/certtest"
	"github.com/netbirdio/netbird/shared/management/proto"
)

var peerKey = []byte("peer-public-key-aaaaaaaaaaaaaaaa")

func TestCollect_ProvesOneMatchingCertificatePerChallenge(t *testing.T) {
	corpCA := certtest.NewCA(t, "corp-root")
	otherCA := certtest.NewCA(t, "other-root")
	unrelatedCA := certtest.NewCA(t, "unrelated-root")

	dir := t.TempDir()
	deviceKey := certtest.ECDSAKey(t)
	device := corpCA.Issue(t, deviceKey, "device")
	writeFile(t, dir, "device.pem", certtest.CertPEM(device)+certtest.KeyPEM(t, deviceKey))

	otherKey := certtest.RSAKey(t)
	writeFile(t, dir, "other.crt", certtest.CertPEM(otherCA.Issue(t, otherKey, "other")))
	writeFile(t, dir, "other.key", certtest.KeyPEM(t, otherKey))

	writeFile(t, dir, "keyless.crt", certtest.CertPEM(corpCA.Issue(t, certtest.ECDSAKey(t), "keyless")))
	writeFile(t, dir, "notes.txt", "ignored")

	challenger := certposture.NewChallenger([]byte("secret"))
	nonce := challenger.Nonce(peerKey, time.Now())
	challenge := func(cas ...string) *proto.Checks {
		return &proto.Checks{CertificateChallenge: &proto.CertificateChallenge{Nonce: nonce, CaCertificates: cas}}
	}
	checks := []*proto.Checks{
		{Files: []string{"/usr/bin/agent"}},
		challenge(corpCA.PEM),
		challenge(corpCA.PEM),
		challenge(otherCA.PEM),
		challenge(unrelatedCA.PEM),
		challenge("not a pem"),
	}

	proofs := Collect(context.Background(), NewFileStore(dir), checks, peerKey)

	require.Len(t, proofs, 2)
	var subjects []string
	for _, p := range proofs {
		chain, err := challenger.Verify(p, peerKey, time.Now())
		require.NoError(t, err)
		subjects = append(subjects, chain[0].Subject.CommonName)
	}
	assert.ElementsMatch(t, []string{"device", "other"}, subjects)
}

func TestCollect_NothingToProve(t *testing.T) {
	dir := t.TempDir()
	key := certtest.ECDSAKey(t)
	ca := certtest.NewCA(t, "root")
	writeFile(t, dir, "device.pem", certtest.CertPEM(ca.Issue(t, key, "device"))+certtest.KeyPEM(t, key))

	tests := []struct {
		name   string
		store  Store
		checks []*proto.Checks
	}{
		{"no checks", NewFileStore(dir), nil},
		{"files only", NewFileStore(dir), []*proto.Checks{{Files: []string{"/bin/x"}}}},
		{"challenge without nonce", NewFileStore(dir), []*proto.Checks{{CertificateChallenge: &proto.CertificateChallenge{CaCertificates: []string{ca.PEM}}}}},
		{"missing store dir", NewFileStore(filepath.Join(dir, "missing")), []*proto.Checks{{CertificateChallenge: &proto.CertificateChallenge{Nonce: []byte{1}, CaCertificates: []string{ca.PEM}}}}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			assert.Nil(t, Collect(context.Background(), tt.store, tt.checks, peerKey))
		})
	}
}

func TestFileStore_ChainWithIntermediate(t *testing.T) {
	root := certtest.NewCA(t, "root")
	intermediate := certtest.NewIntermediate(t, root, "intermediate")
	key := certtest.ECDSAKey(t)
	leaf := intermediate.Issue(t, key, "device")

	dir := t.TempDir()
	writeFile(t, dir, "device.pem", certtest.CertPEM(leaf)+certtest.CertPEM(intermediate.Cert)+certtest.KeyPEM(t, key))

	candidates, err := NewFileStore(dir).Candidates(context.Background())
	require.NoError(t, err)
	require.Len(t, candidates, 1)
	require.Len(t, candidates[0].Chain, 2)

	roots, err := certposture.ParseCAs([]string{root.PEM})
	require.NoError(t, err)
	assert.NoError(t, certposture.VerifyChain(candidates[0].Chain, roots, time.Now()))
}

func writeFile(t *testing.T, dir, name, content string) {
	t.Helper()
	require.NoError(t, os.WriteFile(filepath.Join(dir, name), []byte(content), 0o600))
}

func TestCollectChallenges_RefusesMalformedInput(t *testing.T) {
	ca := certtest.NewCA(t, "corp-root")
	dir := t.TempDir()
	key := certtest.ECDSAKey(t)
	writeFile(t, dir, "device.pem", certtest.CertPEM(ca.Issue(t, key, "device"))+certtest.KeyPEM(t, key))
	store := NewFileStore(dir)
	nonce := certposture.NewChallenger([]byte("secret")).Nonce(peerKey, time.Now())

	tests := []struct {
		name    string
		nonce   []byte
		peerKey []byte
		want    int
	}{
		{"issued nonce and peer key are signed", nonce, peerKey, 1},
		{"short nonce is not signed", nonce[:8], peerKey, 0},
		{"oversized nonce is not signed", append(append([]byte{}, nonce...), 0), peerKey, 0},
		{"short peer key is not signed", nonce, peerKey[:16], 0},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			challenges := []*proto.CertificateChallenge{{Nonce: tt.nonce, CaCertificates: []string{ca.PEM}}}
			assert.Len(t, CollectChallenges(context.Background(), store, challenges, tt.peerKey), tt.want,
				"the device key signs only what management could have issued")
		})
	}
}

func TestFileStore_SkipsKeyOfAnotherCertificate(t *testing.T) {
	ca := certtest.NewCA(t, "corp-root")
	dir := t.TempDir()

	// A stale key next to a renewed certificate, sorted before the good pair, must not
	// produce a proof that management rejects and stop the search there.
	writeFile(t, dir, "a-renewed.crt", certtest.CertPEM(ca.Issue(t, certtest.ECDSAKey(t), "renewed")))
	writeFile(t, dir, "a-renewed.key", certtest.KeyPEM(t, certtest.ECDSAKey(t)))
	goodKey := certtest.ECDSAKey(t)
	good := ca.Issue(t, goodKey, "good")
	writeFile(t, dir, "b-good.pem", certtest.CertPEM(good)+certtest.KeyPEM(t, goodKey))

	candidates, err := NewFileStore(dir).Candidates(context.Background())
	require.NoError(t, err)
	require.Len(t, candidates, 1, "only the certificate whose key matches is a candidate")
	assert.True(t, good.Equal(candidates[0].Chain[0]), "the matching pair is kept")

	challenger := certposture.NewChallenger([]byte("secret"))
	now := time.Now()
	nonce := challenger.Nonce(peerKey, now)
	proofs := CollectChallenges(context.Background(), NewFileStore(dir), []*proto.CertificateChallenge{{Nonce: nonce, CaCertificates: []string{ca.PEM}}}, peerKey)
	require.Len(t, proofs, 1)
	_, err = challenger.Verify(proofs[0], peerKey, now)
	assert.NoError(t, err, "the proof sent is one management accepts")
}