From fb6e955f0003e3bb6eb5b34a0d393718495badc0 Mon Sep 17 00:00:00 2001
From: riccardom <riccardomanfrin@gmail.com>
Date: Tue, 20 Oct 2026 17:53:03 +0200
Subject: [PATCH] [client] pqkem: cover the env flags, signal port, offer
 marker and status flag

These branches decide whether ML-KEM runs, which port a peer is reached
on, and what the status reports, and none of them had a test.
---
 client/internal/peer/conn_pq_test.go | 31 +++++++++++++++
 client/internal/pqkem/env_test.go    | 58 ++++++++++++++++++++++++++++
 client/internal/pqkem_signal_test.go | 35 +++++++++++++++++
 3 files changed, 124 insertions(+)
 create mode 100644 client/internal/pqkem/env_test.go
 create mode 100644 client/internal/pqkem_signal_test.go

diff --git a/client/internal/peer/conn_pq_test.go b/client/internal/peer/conn_pq_test.go
index 3aa7cfa46..ae11f70e8 100644
--- a/client/internal/peer/conn_pq_test.go
+++ b/client/internal/peer/conn_pq_test.go
@@ -164,3 +164,34 @@ func TestConn_RecoverPQKeyOverSignalOnlyWhenConnected(t *testing.T) {
 	c.RecoverPQKeyOverSignal()
 	assert.True(t, offerAttempted(hook), "a connected peer must get a recovery offer")
 }
+
+// TestConn_QuantumResistant covers the status flag: Rosenpass counts only when both sides
+// run it, and an established ML-KEM PSK counts on its own.
+func TestConn_QuantumResistant(t *testing.T) {
+	localRP := []byte("local-rosenpass-key")
+	remoteRP := []byte("remote-rosenpass-key")
+
+	cases := []struct {
+		name          string
+		localRP       []byte
+		remoteRP      []byte
+		pqEstablished bool
+		want          bool
+	}{
+		{"nothing", nil, nil, false, false},
+		{"rosenpass on both sides", localRP, remoteRP, false, true},
+		{"rosenpass only locally", localRP, nil, false, false},
+		{"rosenpass only remotely", nil, remoteRP, false, false},
+		{"ml-kem established", nil, nil, true, true},
+		{"ml-kem established, rosenpass only remotely", nil, remoteRP, true, true},
+		{"ml-kem established, rosenpass on both sides", localRP, remoteRP, true, true},
+	}
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			conn := &Conn{config: ConnConfig{RosenpassConfig: RosenpassConfig{PubKey: c.localRP}}}
+			assert.Equal(t, c.want, conn.quantumResistant(c.remoteRP, c.pqEstablished),
+				"local rosenpass %t, remote rosenpass %t, ml-kem %t",
+				c.localRP != nil, c.remoteRP != nil, c.pqEstablished)
+		})
+	}
+}
diff --git a/client/internal/pqkem/env_test.go b/client/internal/pqkem/env_test.go
new file mode 100644
index 000000000..6af943eb8
--- /dev/null
+++ b/client/internal/pqkem/env_test.go
@@ -0,0 +1,58 @@
+package pqkem
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestEnvFlags(t *testing.T) {
+	cases := []struct {
+		value string
+		want  bool
+	}{
+		{"", false},
+		{"on", true},
+		{"ON", true},
+		{"  on  ", true},
+		{"off", false},
+		{"Off", false},
+		{"true", true},
+		{"TRUE", true},
+		{"1", true},
+		{"t", true},
+		{"false", false},
+		{"0", false},
+		{" 1 ", true},
+		{"yes", false},
+		{"garbage", false},
+	}
+
+	flags := []struct {
+		env string
+		get func() bool
+	}{
+		{EnvEnabled, Enabled},
+		{EnvStrict, Strict},
+	}
+
+	for _, f := range flags {
+		for _, c := range cases {
+			t.Run(f.env+"="+c.value, func(t *testing.T) {
+				t.Setenv(f.env, c.value)
+				assert.Equal(t, c.want, f.get(), "%s=%q", f.env, c.value)
+			})
+		}
+	}
+}
+
+func TestEnvFlagsUnset(t *testing.T) {
+	for _, env := range []string{EnvEnabled, EnvStrict} {
+		t.Setenv(env, "")
+		require.NoError(t, os.Unsetenv(env))
+	}
+	assert.False(t, Enabled(), "an unset %s must leave ML-KEM disabled", EnvEnabled)
+	assert.False(t, Strict(), "an unset %s must leave strict mode off", EnvStrict)
+}
diff --git a/client/internal/pqkem_signal_test.go b/client/internal/pqkem_signal_test.go
new file mode 100644
index 000000000..9118802e9
--- /dev/null
+++ b/client/internal/pqkem_signal_test.go
@@ -0,0 +1,35 @@
+package internal
+
+import (
+	"math"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	sProto "github.com/netbirdio/netbird/shared/signal/proto"
+)
+
+func signalOfferWithMlkemPort(port uint32) *sProto.Message {
+	return &sProto.Message{Body: &sProto.Body{
+		Type:      sProto.Body_OFFER,
+		Payload:   "ufrag:pwd",
+		MlkemPort: &port,
+	}}
+}
+
+// TestConvertToOfferAnswer_MlkemPort checks that the proto's 32-bit ML-KEM port is narrowed
+// to uint16 only when it fits: a larger value is a malformed message, not a port to
+// truncate.
+func TestConvertToOfferAnswer_MlkemPort(t *testing.T) {
+	for _, port := range []uint32{0, 51833, math.MaxUint16} {
+		oa, err := convertToOfferAnswer(signalOfferWithMlkemPort(port))
+		require.NoError(t, err, "port %d", port)
+		assert.Equal(t, uint16(port), oa.MlkemPort, "port %d must survive the conversion", port)
+	}
+
+	for _, port := range []uint32{math.MaxUint16 + 1, math.MaxUint32} {
+		_, err := convertToOfferAnswer(signalOfferWithMlkemPort(port))
+		assert.Error(t, err, "port %d does not fit in uint16 and must be rejected", port)
+	}
+}