Merge pull request #7 from disentangle-network/feature/pq-handshake

disentangle-network · web-flow · commit 2227d9c221bd · 2026-02-19T07:02:09.000-07:00
Merge PQ handshake, CLI, and e2e test to master
diff --git a/cert/cert.go b/cert/cert.go
@@ -158,6 +158,20 @@ func Recombine(v Version, rawCertBytes, publicKey []byte, curve Curve) (Certific
 	return c, nil
 }
 
+// UnmarshalCertificateFromBytes unmarshals a certificate from raw bytes when the full certificate
+// (including public key) is already present. This is used for PQ certificates where the public key
+// is not stripped during handshake marshaling.
+func UnmarshalCertificateFromBytes(b []byte, v Version, curve Curve) (Certificate, error) {
+	switch v {
+	case VersionPre1, Version1:
+		return unmarshalCertificateV1(b, nil)
+	case Version2:
+		return unmarshalCertificateV2(b, nil, curve)
+	default:
+		return nil, ErrUnknownVersion
+	}
+}
+
 // CalculateAlternateFingerprint calculates a 2nd fingerprint representation for P256 certificates
 // CAPool blocklist testing through `VerifyCertificate` and `VerifyCachedCertificate` automatically performs this step.
 func CalculateAlternateFingerprint(c Certificate) (string, error) {
diff --git a/cert/cert_v2.go b/cert/cert_v2.go
@@ -239,13 +239,14 @@ func (c *certificateV2) VerifyPrivateKey(curve Curve, key []byte) error {
 		pub = privkey.PublicKey().Bytes()
 	case Curve_PQ:
 		// Host certs use ML-KEM-1024 key agreement keys
-		if len(key) != mlkem1024.PrivateKeySize {
+		var sk mlkem1024.PrivateKey
+		if err := sk.Unpack(key); err != nil {
+			return ErrInvalidPrivateKey
+		}
+		derivedPub, err := sk.Public().(*mlkem1024.PublicKey).MarshalBinary()
+		if err != nil {
 			return ErrInvalidPrivateKey
 		}
-		_, sk := mlkem1024.NewKeyFromSeed(key[:mlkem1024.KeySeedSize])
-		pub, _ = sk.MarshalBinary()
-		// Compare only the public key portion
-		derivedPub := pub[:mlkem1024.PublicKeySize]
 		if !bytes.Equal(derivedPub, c.publicKey) {
 			return ErrPublicPrivateKeyMismatch
 		}
diff --git a/cert/pem.go b/cert/pem.go
@@ -15,25 +15,25 @@ const ( //cert banners
 )
 
 const ( //key-agreement-key banners
-	X25519PrivateKeyBanner  = "NEBULA X25519 PRIVATE KEY"
-	X25519PublicKeyBanner   = "NEBULA X25519 PUBLIC KEY"
-	P256PrivateKeyBanner    = "NEBULA P256 PRIVATE KEY"
-	P256PublicKeyBanner     = "NEBULA P256 PUBLIC KEY"
+	X25519PrivateKeyBanner    = "NEBULA X25519 PRIVATE KEY"
+	X25519PublicKeyBanner     = "NEBULA X25519 PUBLIC KEY"
+	P256PrivateKeyBanner      = "NEBULA P256 PRIVATE KEY"
+	P256PublicKeyBanner       = "NEBULA P256 PUBLIC KEY"
 	MLKEM1024PrivateKeyBanner = "NEBULA MLKEM1024 PRIVATE KEY"
 	MLKEM1024PublicKeyBanner  = "NEBULA MLKEM1024 PUBLIC KEY"
 )
 
 /* including "ECDSA" in the P256 banners is a clue that these keys should be used only for signing */
 const ( //signing key banners
-	EncryptedECDSAP256PrivateKeyBanner   = "NEBULA ECDSA P256 ENCRYPTED PRIVATE KEY"
-	ECDSAP256PrivateKeyBanner            = "NEBULA ECDSA P256 PRIVATE KEY"
-	ECDSAP256PublicKeyBanner             = "NEBULA ECDSA P256 PUBLIC KEY"
-	EncryptedEd25519PrivateKeyBanner     = "NEBULA ED25519 ENCRYPTED PRIVATE KEY"
-	Ed25519PrivateKeyBanner              = "NEBULA ED25519 PRIVATE KEY"
-	Ed25519PublicKeyBanner               = "NEBULA ED25519 PUBLIC KEY"
-	EncryptedMLDSA87PrivateKeyBanner     = "NEBULA MLDSA87 ENCRYPTED PRIVATE KEY"
-	MLDSA87PrivateKeyBanner              = "NEBULA MLDSA87 PRIVATE KEY"
-	MLDSA87PublicKeyBanner               = "NEBULA MLDSA87 PUBLIC KEY"
+	EncryptedECDSAP256PrivateKeyBanner = "NEBULA ECDSA P256 ENCRYPTED PRIVATE KEY"
+	ECDSAP256PrivateKeyBanner          = "NEBULA ECDSA P256 PRIVATE KEY"
+	ECDSAP256PublicKeyBanner           = "NEBULA ECDSA P256 PUBLIC KEY"
+	EncryptedEd25519PrivateKeyBanner   = "NEBULA ED25519 ENCRYPTED PRIVATE KEY"
+	Ed25519PrivateKeyBanner            = "NEBULA ED25519 PRIVATE KEY"
+	Ed25519PublicKeyBanner             = "NEBULA ED25519 PUBLIC KEY"
+	EncryptedMLDSA87PrivateKeyBanner   = "NEBULA MLDSA87 ENCRYPTED PRIVATE KEY"
+	MLDSA87PrivateKeyBanner            = "NEBULA MLDSA87 PRIVATE KEY"
+	MLDSA87PublicKeyBanner             = "NEBULA MLDSA87 PUBLIC KEY"
 )
 
 // UnmarshalCertificateFromPEM will try to unmarshal the first pem block in a byte array, returning any non consumed
diff --git a/cert/sign_test.go b/cert/sign_test.go
@@ -173,7 +173,7 @@ func TestCertificateV2_SignPQ(t *testing.T) {
 	assert.Equal(t, pubBytes, c.PublicKey())
 
 	// Verify signature length matches ML-DSA-87
-	assert.Equal(t, mldsa87.SignatureSize, len(c.Signature()))
+	assert.Len(t, c.Signature(), mldsa87.SignatureSize)
 
 	// Marshal and unmarshal roundtrip
 	b, err := c.Marshal()
diff --git a/cert_test/cert.go b/cert_test/cert.go
@@ -9,6 +9,8 @@ import (
 	"net/netip"
 	"time"
 
+	"github.com/cloudflare/circl/kem/mlkem/mlkem1024"
+	"github.com/cloudflare/circl/sign/mldsa/mldsa87"
 	"github.com/slackhq/nebula/cert"
 	"golang.org/x/crypto/curve25519"
 	"golang.org/x/crypto/ed25519"
@@ -30,6 +32,13 @@ func NewTestCaCert(version cert.Version, curve cert.Curve, before, after time.Ti
 
 		pub = elliptic.Marshal(elliptic.P256(), privk.PublicKey.X, privk.PublicKey.Y)
 		priv = privk.D.FillBytes(make([]byte, 32))
+	case cert.Curve_PQ:
+		pk, sk, err := mldsa87.GenerateKey(rand.Reader)
+		if err != nil {
+			panic(err)
+		}
+		pub = pk.Bytes()
+		priv = sk.Bytes()
 	default:
 		// There is no default to allow the underlying lib to respond with an error
 	}
@@ -84,6 +93,8 @@ func NewTestCert(v cert.Version, curve cert.Curve, ca cert.Certificate, key []by
 		pub, priv = X25519Keypair()
 	case cert.Curve_P256:
 		pub, priv = P256Keypair()
+	case cert.Curve_PQ:
+		pub, priv = PQKeypair()
 	default:
 		panic("unknown curve")
 	}
@@ -163,3 +174,13 @@ func P256Keypair() ([]byte, []byte) {
 	pubkey := privkey.PublicKey()
 	return pubkey.Bytes(), privkey.Bytes()
 }
+
+func PQKeypair() ([]byte, []byte) {
+	pk, sk, err := mlkem1024.GenerateKeyPair(rand.Reader)
+	if err != nil {
+		panic(err)
+	}
+	pubBytes, _ := pk.MarshalBinary()
+	privBytes, _ := sk.MarshalBinary()
+	return pubBytes, privBytes
+}
diff --git a/cmd/nebula-cert/ca.go b/cmd/nebula-cert/ca.go
@@ -13,6 +13,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/cloudflare/circl/sign/mldsa/mldsa87"
 	"github.com/skip2/go-qrcode"
 	"github.com/slackhq/nebula/cert"
 	"github.com/slackhq/nebula/pkclient"
@@ -59,7 +60,7 @@ func newCaFlags() *caFlags {
 	cf.argonParallelism = cf.set.Uint("argon-parallelism", 4, "Optional: Argon2 parallelism parameter used for encrypted private key passphrase")
 	cf.argonIterations = cf.set.Uint("argon-iterations", 1, "Optional: Argon2 iterations parameter used for encrypted private key passphrase")
 	cf.encryption = cf.set.Bool("encrypt", false, "Optional: prompt for passphrase and write out-key in an encrypted format")
-	cf.curve = cf.set.String("curve", "25519", "EdDSA/ECDSA Curve (25519, P256)")
+	cf.curve = cf.set.String("curve", "25519", "EdDSA/ECDSA/PQ Curve (25519, P256, PQ)")
 	cf.p11url = p11Flag(cf.set)
 
 	cf.ips = cf.set.String("ips", "", "Deprecated, see -networks")
@@ -243,11 +244,23 @@ func ca(args []string, out io.Writer, errOut io.Writer, pr PasswordReader) error
 			}
 			rawPriv = eKey.Bytes()
 			pub = eKey.PublicKey().Bytes()
+		case "PQ":
+			curve = cert.Curve_PQ
+			pk, sk, err := mldsa87.GenerateKey(rand.Reader)
+			if err != nil {
+				return fmt.Errorf("error while generating ML-DSA-87 keys: %s", err)
+			}
+			pub = pk.Bytes()
+			rawPriv = sk.Bytes()
 		default:
 			return fmt.Errorf("invalid curve: %s", *cf.curve)
 		}
 	}
 
+	if curve == cert.Curve_PQ && version != cert.Version2 {
+		return fmt.Errorf("PQ curve requires certificate version 2")
+	}
+
 	t := &cert.TBSCertificate{
 		Version:        version,
 		Name:           *cf.name,
diff --git a/cmd/nebula-cert/ca_test.go b/cmd/nebula-cert/ca_test.go
@@ -34,7 +34,7 @@ func Test_caHelp(t *testing.T) {
 			"  -argon-parallelism uint\n"+
 			"    \tOptional: Argon2 parallelism parameter used for encrypted private key passphrase (default 4)\n"+
 			"  -curve string\n"+
-			"    \tEdDSA/ECDSA Curve (25519, P256) (default \"25519\")\n"+
+			"    \tEdDSA/ECDSA/PQ Curve (25519, P256, PQ) (default \"25519\")\n"+
 			"  -duration duration\n"+
 			"    \tOptional: amount of time the certificate should be valid for. Valid time units are seconds: \"s\", minutes: \"m\", hours: \"h\" (default 8760h0m0s)\n"+
 			"  -encrypt\n"+
diff --git a/cmd/nebula-cert/keygen.go b/cmd/nebula-cert/keygen.go
@@ -24,7 +24,7 @@ func newKeygenFlags() *keygenFlags {
 	cf.set.Usage = func() {}
 	cf.outPubPath = cf.set.String("out-pub", "", "Required: path to write the public key to")
 	cf.outKeyPath = cf.set.String("out-key", "", "Required: path to write the private key to")
-	cf.curve = cf.set.String("curve", "25519", "ECDH Curve (25519, P256)")
+	cf.curve = cf.set.String("curve", "25519", "ECDH/KEM Curve (25519, P256, PQ)")
 	cf.p11url = p11Flag(cf.set)
 	return &cf
 }
@@ -64,6 +64,9 @@ func keygen(args []string, out io.Writer, errOut io.Writer) error {
 		case "P256":
 			pub, rawPriv = p256Keypair()
 			curve = cert.Curve_P256
+		case "PQ":
+			pub, rawPriv = pqKeypair()
+			curve = cert.Curve_PQ
 		default:
 			return fmt.Errorf("invalid curve: %s", *cf.curve)
 		}
diff --git a/cmd/nebula-cert/keygen_test.go b/cmd/nebula-cert/keygen_test.go
@@ -21,7 +21,7 @@ func Test_keygenHelp(t *testing.T) {
 		t,
 		"Usage of "+os.Args[0]+" keygen <flags>: create a public/private key pair. the public key can be passed to `nebula-cert sign`\n"+
 			"  -curve string\n"+
-			"    \tECDH Curve (25519, P256) (default \"25519\")\n"+
+			"    \tECDH/KEM Curve (25519, P256, PQ) (default \"25519\")\n"+
 			"  -out-key string\n"+
 			"    \tRequired: path to write the private key to\n"+
 			"  -out-pub string\n"+
diff --git a/cmd/nebula-cert/sign.go b/cmd/nebula-cert/sign.go
@@ -14,6 +14,7 @@ import (
 
 	"github.com/skip2/go-qrcode"
 	"github.com/slackhq/nebula/cert"
+	"github.com/slackhq/nebula/noiseutil"
 	"github.com/slackhq/nebula/pkclient"
 	"golang.org/x/crypto/curve25519"
 )
@@ -405,6 +406,8 @@ func newKeypair(curve cert.Curve) ([]byte, []byte) {
 		return x25519Keypair()
 	case cert.Curve_P256:
 		return p256Keypair()
+	case cert.Curve_PQ:
+		return pqKeypair()
 	default:
 		return nil, nil
 	}
@@ -433,6 +436,14 @@ func p256Keypair() ([]byte, []byte) {
 	return pubkey.Bytes(), privkey.Bytes()
 }
 
+func pqKeypair() ([]byte, []byte) {
+	pub, priv, err := noiseutil.PQKEMKeypair()
+	if err != nil {
+		panic(err)
+	}
+	return pub, priv
+}
+
 func signSummary() string {
 	return "sign <flags>: create and sign a certificate"
 }
diff --git a/connection_state.go b/connection_state.go
@@ -25,10 +25,17 @@ type ConnectionState struct {
 	messageCounter atomic.Uint64
 	window         *Bits
 	writeLock      sync.Mutex
+
+	// Post-quantum KEM state (only used for Curve_PQ)
+	pqKemPubKey  []byte // Our ephemeral ML-KEM-1024 public key
+	pqKemPrivKey []byte // Our ephemeral ML-KEM-1024 private key
+	pqKemSS      []byte // Shared secret from KEM exchange
 }
 
 func NewConnectionState(l *logrus.Logger, cs *CertState, crt cert.Certificate, initiator bool, pattern noise.HandshakePattern) (*ConnectionState, error) {
 	var dhFunc noise.DHFunc
+	var pqKemPub, pqKemPriv []byte
+
 	switch crt.Curve() {
 	case cert.Curve_CURVE25519:
 		dhFunc = noise.DH25519
@@ -38,6 +45,19 @@ func NewConnectionState(l *logrus.Logger, cs *CertState, crt cert.Certificate, i
 		} else {
 			dhFunc = noiseutil.DHP256
 		}
+	case cert.Curve_PQ:
+		// Hybrid mode: X25519 DH for classical security + ML-KEM-1024 for PQ security.
+		// The Noise IX handshake uses X25519 for the DH tokens. The ML-KEM-1024
+		// exchange is layered on top via the handshake payload (KemPublicKey/KemCiphertext).
+		// Both shared secrets are mixed into the final symmetric keys via HKDF.
+		dhFunc = noise.DH25519
+
+		// Generate ephemeral ML-KEM-1024 keypair for this handshake
+		var err error
+		pqKemPub, pqKemPriv, err = noiseutil.PQKEMKeypair()
+		if err != nil {
+			return nil, fmt.Errorf("NewConnectionState: ML-KEM-1024 keygen failed: %s", err)
+		}
 	default:
 		return nil, fmt.Errorf("invalid curve: %s", crt.Curve())
 	}
@@ -49,7 +69,24 @@ func NewConnectionState(l *logrus.Logger, cs *CertState, crt cert.Certificate, i
 		ncs = noise.NewCipherSuite(dhFunc, noiseutil.CipherAESGCM, noise.HashSHA256)
 	}
 
-	static := noise.DHKey{Private: cs.privateKey, Public: crt.PublicKey()}
+	// For PQ hybrid mode, use a fresh X25519 keypair as the Noise "static" key.
+	// The actual identity authentication comes from the ML-DSA-87 cert signature,
+	// not from the DH static key. This is safe because:
+	// 1. The cert is verified against the CA in the handshake payload
+	// 2. The KEM exchange provides the PQ-secure key agreement
+	// 3. The X25519 DH provides defense-in-depth
+	var static noise.DHKey
+	if crt.Curve() == cert.Curve_PQ {
+		// Generate ephemeral X25519 keypair (cert's public key is ML-KEM, not X25519)
+		var err error
+		static, err = noise.DH25519.GenerateKeypair(rand.Reader)
+		if err != nil {
+			return nil, fmt.Errorf("NewConnectionState: X25519 keygen failed: %s", err)
+		}
+	} else {
+		static = noise.DHKey{Private: cs.privateKey, Public: crt.PublicKey()}
+	}
+
 	hs, err := noise.NewHandshakeState(noise.Config{
 		CipherSuite:   ncs,
 		Random:        rand.Reader,
@@ -67,10 +104,12 @@ func NewConnectionState(l *logrus.Logger, cs *CertState, crt cert.Certificate, i
 	// The queue and ready params prevent a counter race that would happen when
 	// sending stored packets and simultaneously accepting new traffic.
 	ci := &ConnectionState{
-		H:         hs,
-		initiator: initiator,
-		window:    NewBits(ReplayWindow),
-		myCert:    crt,
+		H:            hs,
+		initiator:    initiator,
+		window:       NewBits(ReplayWindow),
+		myCert:       crt,
+		pqKemPubKey:  pqKemPub,
+		pqKemPrivKey: pqKemPriv,
 	}
 	// always start the counter from 2, as packet 1 and packet 2 are handshake packets.
 	ci.messageCounter.Add(2)
diff --git a/e2e/handshakes_test.go b/e2e/handshakes_test.go
@@ -134,6 +134,48 @@ func TestGoodHandshake(t *testing.T) {
 	theirControl.Stop()
 }
 
+func TestGoodHandshakePQ(t *testing.T) {
+	// Post-quantum handshake test: ML-DSA-87 certificates + hybrid X25519/ML-KEM-1024 key exchange
+	ca, _, caKey, _ := cert_test.NewTestCaCert(cert.Version2, cert.Curve_PQ, time.Now(), time.Now().Add(10*time.Minute), nil, nil, []string{})
+	myControl, myVpnIpNet, myUdpAddr, _ := newSimpleServer(cert.Version2, ca, caKey, "me-pq", "10.128.0.1/24", nil)
+	theirControl, theirVpnIpNet, theirUdpAddr, _ := newSimpleServer(cert.Version2, ca, caKey, "them-pq", "10.128.0.2/24", nil)
+
+	// Put their info in our lighthouse
+	myControl.InjectLightHouseAddr(theirVpnIpNet[0].Addr(), theirUdpAddr)
+
+	// Start the servers
+	myControl.Start()
+	theirControl.Start()
+
+	t.Log("PQ: Send a udp packet through to begin standing up the tunnel")
+	myControl.InjectTunUDPPacket(theirVpnIpNet[0].Addr(), 80, myVpnIpNet[0].Addr(), 80, []byte("Hi from PQ me"))
+
+	t.Log("PQ: Have them consume my stage 0 packet. They have a tunnel now")
+	theirControl.InjectUDPPacket(myControl.GetFromUDP(true))
+
+	t.Log("PQ: Have me consume their stage 1 packet. I have a tunnel now")
+	myControl.InjectUDPPacket(theirControl.GetFromUDP(true))
+
+	t.Log("PQ: Wait until we see my cached packet come through")
+	myControl.WaitForType(1, 0, theirControl)
+
+	t.Log("PQ: Make sure our host infos are correct")
+	assertHostInfoPair(t, myUdpAddr, theirUdpAddr, myVpnIpNet, theirVpnIpNet, myControl, theirControl)
+
+	t.Log("PQ: Get that cached packet and make sure it looks right")
+	myCachedPacket := theirControl.GetFromTun(true)
+	assertUdpPacket(t, []byte("Hi from PQ me"), myCachedPacket, myVpnIpNet[0].Addr(), theirVpnIpNet[0].Addr(), 80, 80)
+
+	t.Log("PQ: Do a bidirectional tunnel test")
+	r := router.NewR(t, myControl, theirControl)
+	defer r.RenderFlow()
+	assertTunnel(t, myVpnIpNet[0].Addr(), theirVpnIpNet[0].Addr(), myControl, theirControl, r)
+
+	r.RenderHostmaps("PQ Final hostmaps", myControl, theirControl)
+	myControl.Stop()
+	theirControl.Stop()
+}
+
 func TestGoodHandshakeNoOverlap(t *testing.T) {
 	ca, _, caKey, _ := cert_test.NewTestCaCert(cert.Version2, cert.Curve_CURVE25519, time.Now(), time.Now().Add(10*time.Minute), nil, nil, []string{})
 	myControl, myVpnIpNet, myUdpAddr, _ := newSimpleServer(cert.Version2, ca, caKey, "me", "10.128.0.1/24", nil)
diff --git a/e2e/helpers_test.go b/e2e/helpers_test.go
@@ -102,7 +102,7 @@ func newSimpleServerWithUdpAndUnsafeNetworks(v cert.Version, caCrt cert.Certific
 		}
 	}
 
-	_, _, myPrivKey, myPEM := cert_test.NewTestCert(v, cert.Curve_CURVE25519, caCrt, caKey, name, time.Now(), time.Now().Add(5*time.Minute), vpnNetworks, unsafeNetworks, []string{})
+	_, _, myPrivKey, myPEM := cert_test.NewTestCert(v, caCrt.Curve(), caCrt, caKey, name, time.Now(), time.Now().Add(5*time.Minute), vpnNetworks, unsafeNetworks, []string{})
 
 	caB, err := caCrt.MarshalPEM()
 	if err != nil {
diff --git a/handshake_ix.go b/handshake_ix.go
diff --git a/nebula.pb.go b/nebula.pb.go
diff --git a/nebula.proto b/nebula.proto
diff --git a/noiseutil/pq.go b/noiseutil/pq.go
diff --git a/noiseutil/pq_test.go b/noiseutil/pq_test.go
diff --git a/pki.go b/pki.go
