feat: RS erasure-coded handshake chunking for PQ key exchange (#30)

* feat(header): add HandshakeIXPSK0Chunked subtype and ChunkHeader

Add RS erasure-coded chunk header types for oversized PQ handshakes.
The 8-byte ChunkHeader carries handshake_id, noise message number,
chunk index, total chunks, and data shard count for reconstruction.

* feat: RS erasure-coded chunking in handshake send paths

Add Reed-Solomon encoding for oversized handshake messages (PQ
handshakes ~9KB). Messages exceeding 1200 bytes are automatically
split into k+m chunks where k=ceil(payload/1200) and m=3 parity.
Any k of k+m chunks arriving suffices for reconstruction.

Send paths modified: handleOutbound (initiator), ixHandshakeStage1
(responder direct + relay), and ErrAlreadySeen cached resend.
Non-PQ handshakes below threshold bypass chunking entirely.

* feat: RS reassembly buffer in handshake receive path

Add ReassemblyManager for reconstructing chunked handshake messages.
Chunks are buffered by (handshakeID, noiseMsgNum) key and RS-decoded
when k shards arrive. Buffers are bounded (256 max) and expired
(5s timeout) for DoS mitigation. HandleIncoming dispatches chunked
packets to reassembly and re-injects completed messages.

* feat: add length-prefix framing and fix e2e PQ handshake test

RS encode now prepends a 4-byte big-endian length prefix before
splitting into shards, allowing the decoder to strip RS padding
that was corrupting protobuf unmarshal. Updated all unit tests to
account for the +4 byte prefix in data shard count calculations.
Rewrote TestGoodHandshakePQ to use router-based assertTunnel
approach since chunked handshakes produce multiple UDP packets.

All unit tests, reassembly tests, PQ e2e test, and non-PQ e2e
tests pass -- backward compatibility confirmed.

* fix: resolve testifylint CI failures

Use assert.LessOrEqual instead of assert.True for comparison,
and assert.Empty instead of assert.Len(0) per golangci-lint
testifylint rules.

---------

Co-authored-by: privsim <excaliberswake@pm.me>

Author: disentangle-network

e2e/handshakes_test.go

@@ -135,42 +135,30 @@ func TestGoodHandshake(t *testing.T) {
 }
 
 func TestGoodHandshakePQ(t *testing.T) {
-	// Post-quantum handshake test: ML-DSA-87 certificates + hybrid X25519/ML-KEM-1024 key exchange
+	// Post-quantum handshake test: ML-DSA-87 certificates + hybrid X25519/ML-KEM-1024 key exchange.
+	// PQ handshakes produce ~9KB messages which are RS-chunked into multiple UDP packets.
+	// The router handles forwarding all chunks automatically.
 	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)
+	theirControl.InjectLightHouseAddr(myVpnIpNet[0].Addr(), myUdpAddr)
 
 	// 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))
+	r := router.NewR(t, myControl, theirControl)
+	defer r.RenderFlow()
 
-	t.Log("PQ: Wait until we see my cached packet come through")
-	myControl.WaitForType(1, 0, theirControl)
+	t.Log("PQ: Do a bidirectional tunnel test (includes RS-chunked handshake)")
+	assertTunnel(t, myVpnIpNet[0].Addr(), theirVpnIpNet[0].Addr(), myControl, theirControl, r)
 
 	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()

go.mod

@@ -45,6 +45,8 @@ require (
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/google/btree v1.1.2 // indirect
+	github.com/klauspost/cpuid/v2 v2.3.0 // indirect
+	github.com/klauspost/reedsolomon v1.13.2 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/prometheus/client_model v0.6.2 // indirect

go.sum

@@ -76,6 +76,10 @@ github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
 github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo=
 github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ=
+github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y=
+github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
+github.com/klauspost/reedsolomon v1.13.2 h1:9qtQy2tKEVpVB8Pfq87ZljHZb60/LbeTQ1OxV8EGzdE=
+github.com/klauspost/reedsolomon v1.13.2/go.mod h1:ggJT9lc71Vu+cSOPBlxGvBN6TfAS77qB4fp8vJ05NSA=
 github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 github.com/konsorten/go-windows-terminal-sequences v1.0.3/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc=

handshake_ix.go

@@ -402,17 +402,35 @@ func ixHandshakeStage1(f *Interface, via ViaSender, packet []byte, h *header.H)
 			}
 
 			msg = existing.HandshakePacket[2]
-			f.messageMetrics.Tx(header.Handshake, header.MessageSubType(msg[1]), 1)
+			cachedChunked := needsChunking(msg)
+			if cachedChunked {
+				f.messageMetrics.Tx(header.Handshake, header.HandshakeIXPSK0Chunked, 1)
+			} else {
+				f.messageMetrics.Tx(header.Handshake, header.MessageSubType(msg[1]), 1)
+			}
 			if !via.IsRelayed {
-				err := f.outside.WriteTo(msg, via.UdpAddr)
-				if err != nil {
-					f.l.WithField("vpnAddrs", existing.vpnAddrs).WithField("from", via).
-						WithField("handshake", m{"stage": 2, "style": "ix_psk0"}).WithField("cached", true).
-						WithError(err).Error("Failed to send handshake message")
+				if cachedChunked {
+					err := sendHandshakeChunked(f.l, f.outside, msg, existing.remoteIndexId, 1, via.UdpAddr)
+					if err != nil {
+						f.l.WithField("vpnAddrs", existing.vpnAddrs).WithField("from", via).
+							WithField("handshake", m{"stage": 2, "style": "ix_psk0_chunked"}).WithField("cached", true).
+							WithError(err).Error("Failed to send chunked handshake message")
+					} else {
+						f.l.WithField("vpnAddrs", existing.vpnAddrs).WithField("from", via).
+							WithField("handshake", m{"stage": 2, "style": "ix_psk0_chunked"}).WithField("cached", true).
+							Info("Handshake message sent")
+					}
 				} else {
-					f.l.WithField("vpnAddrs", existing.vpnAddrs).WithField("from", via).
-						WithField("handshake", m{"stage": 2, "style": "ix_psk0"}).WithField("cached", true).
-						Info("Handshake message sent")
+					err := f.outside.WriteTo(msg, via.UdpAddr)
+					if err != nil {
+						f.l.WithField("vpnAddrs", existing.vpnAddrs).WithField("from", via).
+							WithField("handshake", m{"stage": 2, "style": "ix_psk0"}).WithField("cached", true).
+							WithError(err).Error("Failed to send handshake message")
+					} else {
+						f.l.WithField("vpnAddrs", existing.vpnAddrs).WithField("from", via).
+							WithField("handshake", m{"stage": 2, "style": "ix_psk0"}).WithField("cached", true).
+							Info("Handshake message sent")
+					}
 				}
 				return
 			} else {
@@ -421,9 +439,17 @@ func ixHandshakeStage1(f *Interface, via ViaSender, packet []byte, h *header.H)
 					return
 				}
 				hostinfo.relayState.InsertRelayTo(via.relayHI.vpnAddrs[0])
-				f.SendVia(via.relayHI, via.relay, msg, make([]byte, 12), make([]byte, mtu), false)
+				if cachedChunked {
+					sendHandshakeChunkedVia(f, via.relayHI, via.relay, msg, existing.remoteIndexId, 1)
+				} else {
+					f.SendVia(via.relayHI, via.relay, msg, make([]byte, 12), make([]byte, mtu), false)
+				}
+				hsStyle := "ix_psk0"
+				if cachedChunked {
+					hsStyle = "ix_psk0_chunked"
+				}
 				f.l.WithField("vpnAddrs", existing.vpnAddrs).WithField("relay", via.relayHI.vpnAddrs[0]).
-					WithField("handshake", m{"stage": 2, "style": "ix_psk0"}).WithField("cached", true).
+					WithField("handshake", m{"stage": 2, "style": hsStyle}).WithField("cached", true).
 					Info("Handshake message sent")
 				return
 			}
@@ -471,16 +497,29 @@ func ixHandshakeStage1(f *Interface, via ViaSender, packet []byte, h *header.H)
 	}
 
 	// Do the send
-	f.messageMetrics.Tx(header.Handshake, header.MessageSubType(msg[1]), 1)
+	chunked := needsChunking(msg)
+	if chunked {
+		f.messageMetrics.Tx(header.Handshake, header.HandshakeIXPSK0Chunked, 1)
+	} else {
+		f.messageMetrics.Tx(header.Handshake, header.MessageSubType(msg[1]), 1)
+	}
 	if !via.IsRelayed {
-		err = f.outside.WriteTo(msg, via.UdpAddr)
+		if chunked {
+			err = sendHandshakeChunked(f.l, f.outside, msg, hs.Details.InitiatorIndex, 1, via.UdpAddr)
+		} else {
+			err = f.outside.WriteTo(msg, via.UdpAddr)
+		}
+		hsStyle := "ix_psk0"
+		if chunked {
+			hsStyle = "ix_psk0_chunked"
+		}
 		log := f.l.WithField("vpnAddrs", vpnAddrs).WithField("from", via).
 			WithField("certName", certName).
 			WithField("certVersion", certVersion).
 			WithField("fingerprint", fingerprint).
 			WithField("issuer", issuer).
 			WithField("initiatorIndex", hs.Details.InitiatorIndex).WithField("responderIndex", hs.Details.ResponderIndex).
-			WithField("remoteIndex", h.RemoteIndex).WithField("handshake", m{"stage": 2, "style": "ix_psk0"})
+			WithField("remoteIndex", h.RemoteIndex).WithField("handshake", m{"stage": 2, "style": hsStyle})
 		if err != nil {
 			log.WithError(err).Error("Failed to send handshake")
 		} else {
@@ -495,14 +534,22 @@ func ixHandshakeStage1(f *Interface, via ViaSender, packet []byte, h *header.H)
 		// I successfully received a handshake. Just in case I marked this tunnel as 'Disestablished', ensure
 		// it's correctly marked as working.
 		via.relayHI.relayState.UpdateRelayForByIdxState(via.remoteIdx, Established)
-		f.SendVia(via.relayHI, via.relay, msg, make([]byte, 12), make([]byte, mtu), false)
+		if chunked {
+			sendHandshakeChunkedVia(f, via.relayHI, via.relay, msg, hs.Details.InitiatorIndex, 1)
+		} else {
+			f.SendVia(via.relayHI, via.relay, msg, make([]byte, 12), make([]byte, mtu), false)
+		}
+		hsStyle := "ix_psk0"
+		if chunked {
+			hsStyle = "ix_psk0_chunked"
+		}
 		f.l.WithField("vpnAddrs", vpnAddrs).WithField("relay", via.relayHI.vpnAddrs[0]).
 			WithField("certName", certName).
 			WithField("certVersion", certVersion).
 			WithField("fingerprint", fingerprint).
 			WithField("issuer", issuer).
 			WithField("initiatorIndex", hs.Details.InitiatorIndex).WithField("responderIndex", hs.Details.ResponderIndex).
-			WithField("remoteIndex", h.RemoteIndex).WithField("handshake", m{"stage": 2, "style": "ix_psk0"}).
+			WithField("remoteIndex", h.RemoteIndex).WithField("handshake", m{"stage": 2, "style": hsStyle}).
 			Info("Handshake message sent")
 	}
 

handshake_manager.go

@@ -60,6 +60,7 @@ type HandshakeManager struct {
 	metricTimedOut         metrics.Counter
 	f                      *Interface
 	l                      *logrus.Logger
+	reassembly             *ReassemblyManager
 
 	// can be used to trigger outbound handshake for the given vpnIp
 	trigger chan netip.Addr
@@ -117,6 +118,7 @@ func NewHandshakeManager(l *logrus.Logger, mainHostMap *HostMap, lightHouse *Lig
 		metricInitiated:        metrics.GetOrRegisterCounter("handshake_manager.initiated", nil),
 		metricTimedOut:         metrics.GetOrRegisterCounter("handshake_manager.timed_out", nil),
 		l:                      l,
+		reassembly:             NewReassemblyManager(l),
 	}
 }
 
@@ -158,6 +160,19 @@ func (hm *HandshakeManager) HandleIncoming(via ViaSender, packet []byte, h *head
 				hm.DeleteHostInfo(newHostinfo.hostinfo)
 			}
 		}
+
+	case header.HandshakeIXPSK0Chunked:
+		// RS-coded chunk: buffer and attempt reassembly
+		reassembled, ok := hm.reassembly.HandleChunk(packet, h)
+		if ok {
+			// Reassembly complete: re-parse the reconstructed header and process as normal
+			var reassembledH header.H
+			if err := reassembledH.Parse(reassembled); err != nil {
+				hm.l.WithError(err).Error("Failed to parse reassembled handshake header")
+				return
+			}
+			hm.HandleIncoming(via, reassembled, &reassembledH)
+		}
 	}
 }
 
@@ -236,18 +251,32 @@ func (hm *HandshakeManager) handleOutbound(vpnIp netip.Addr, lighthouseTriggered
 	}
 
 	// Send the handshake to all known ips, stage 2 takes care of assigning the hostinfo.remote based on the first to reply
+	chunked := needsChunking(hostinfo.HandshakePacket[0])
 	var sentTo []netip.AddrPort
 	hostinfo.remotes.ForEach(hm.mainHostMap.GetPreferredRanges(), func(addr netip.AddrPort, _ bool) {
-		hm.messageMetrics.Tx(header.Handshake, header.MessageSubType(hostinfo.HandshakePacket[0][1]), 1)
-		err := hm.outside.WriteTo(hostinfo.HandshakePacket[0], addr)
-		if err != nil {
-			hostinfo.logger(hm.l).WithField("udpAddr", addr).
-				WithField("initiatorIndex", hostinfo.localIndexId).
-				WithField("handshake", m{"stage": 1, "style": "ix_psk0"}).
-				WithError(err).Error("Failed to send handshake message")
-
+		if chunked {
+			hm.messageMetrics.Tx(header.Handshake, header.HandshakeIXPSK0Chunked, 1)
+			err := sendHandshakeChunked(hm.l, hm.outside, hostinfo.HandshakePacket[0],
+				hostinfo.localIndexId, 0, addr)
+			if err != nil {
+				hostinfo.logger(hm.l).WithField("udpAddr", addr).
+					WithField("initiatorIndex", hostinfo.localIndexId).
+					WithField("handshake", m{"stage": 1, "style": "ix_psk0_chunked"}).
+					WithError(err).Error("Failed to send chunked handshake message")
+			} else {
+				sentTo = append(sentTo, addr)
+			}
 		} else {
-			sentTo = append(sentTo, addr)
+			hm.messageMetrics.Tx(header.Handshake, header.MessageSubType(hostinfo.HandshakePacket[0][1]), 1)
+			err := hm.outside.WriteTo(hostinfo.HandshakePacket[0], addr)
+			if err != nil {
+				hostinfo.logger(hm.l).WithField("udpAddr", addr).
+					WithField("initiatorIndex", hostinfo.localIndexId).
+					WithField("handshake", m{"stage": 1, "style": "ix_psk0"}).
+					WithError(err).Error("Failed to send handshake message")
+			} else {
+				sentTo = append(sentTo, addr)
+			}
 		}
 	})
 
@@ -345,7 +374,12 @@ func (hm *HandshakeManager) handleOutbound(vpnIp netip.Addr, lighthouseTriggered
 			switch existingRelay.State {
 			case Established:
 				hostinfo.logger(hm.l).WithField("relay", relay.String()).Info("Send handshake via relay")
-				hm.f.SendVia(relayHostInfo, existingRelay, hostinfo.HandshakePacket[0], make([]byte, 12), make([]byte, mtu), false)
+				if chunked {
+					sendHandshakeChunkedVia(hm.f, relayHostInfo, existingRelay, hostinfo.HandshakePacket[0],
+						hostinfo.localIndexId, 0)
+				} else {
+					hm.f.SendVia(relayHostInfo, existingRelay, hostinfo.HandshakePacket[0], make([]byte, 12), make([]byte, mtu), false)
+				}
 			case Disestablished:
 				// Mark this relay as 'requested'
 				relayHostInfo.relayState.UpdateRelayForByIpState(vpnIp, Requested)

header/header.go

@@ -60,8 +60,9 @@ const (
 )
 
 const (
-	HandshakeIXPSK0 MessageSubType = 0
-	HandshakeXXPSK0 MessageSubType = 1
+	HandshakeIXPSK0        MessageSubType = 0
+	HandshakeXXPSK0        MessageSubType = 1
+	HandshakeIXPSK0Chunked MessageSubType = 2
 )
 
 var ErrHeaderTooShort = errors.New("header is too short")
@@ -83,7 +84,8 @@ var subTypeMap = map[MessageType]*map[MessageSubType]string{
 	Test:        &subTypeTestMap,
 	CloseTunnel: &subTypeNoneMap,
 	Handshake: {
-		HandshakeIXPSK0: "ix_psk0",
+		HandshakeIXPSK0:        "ix_psk0",
+		HandshakeIXPSK0Chunked: "ix_psk0_chunked",
 	},
 	Control: &subTypeNoneMap,
 }
@@ -193,3 +195,67 @@ func NewHeader(b []byte) (*H, error) {
 	}
 	return h, nil
 }
+
+// ChunkHeader is an 8-byte header prepended to each RS-coded chunk of an
+// oversized handshake message. It sits between the Nebula header and the
+// chunk payload.
+//
+//	0       8      16      24      32
+//	|-------|-------|-------|-------|
+//	|        handshake_id (uint32) |
+//	|-------|-------|-------|-------|
+//	|msg_num|chunk_i| total | data |
+//	|-------|-------|-------|-------|
+const ChunkHeaderLen = 8
+
+type ChunkHeader struct {
+	HandshakeID uint32 // Matches InitiatorIndex for session disambiguation
+	NoiseMsgNum uint8  // 0 = message 1 (initiator), 1 = message 2 (responder)
+	ChunkIdx    uint8  // 0..TotalChunks-1
+	TotalChunks uint8  // Total number of chunks (k + m)
+	DataShards  uint8  // Number of data shards (k), reconstruction threshold
+}
+
+var ErrChunkHeaderTooShort = errors.New("chunk header is too short")
+
+// EncodeChunkHeader writes an 8-byte chunk header into the provided slice.
+// The slice must have capacity for at least ChunkHeaderLen bytes.
+func EncodeChunkHeader(b []byte, handshakeID uint32, noiseMsgNum, chunkIdx, totalChunks, dataShards uint8) []byte {
+	b = b[:ChunkHeaderLen]
+	binary.BigEndian.PutUint32(b[0:4], handshakeID)
+	b[4] = noiseMsgNum
+	b[5] = chunkIdx
+	b[6] = totalChunks
+	b[7] = dataShards
+	return b
+}
+
+// ParseChunkHeader reads an 8-byte chunk header from the provided bytes.
+func (ch *ChunkHeader) Parse(b []byte) error {
+	if len(b) < ChunkHeaderLen {
+		return ErrChunkHeaderTooShort
+	}
+	ch.HandshakeID = binary.BigEndian.Uint32(b[0:4])
+	ch.NoiseMsgNum = b[4]
+	ch.ChunkIdx = b[5]
+	ch.TotalChunks = b[6]
+	ch.DataShards = b[7]
+	return nil
+}
+
+// Encode writes the chunk header into the provided byte slice.
+func (ch *ChunkHeader) Encode(b []byte) ([]byte, error) {
+	if ch == nil {
+		return nil, errors.New("nil chunk header")
+	}
+	return EncodeChunkHeader(b, ch.HandshakeID, ch.NoiseMsgNum, ch.ChunkIdx, ch.TotalChunks, ch.DataShards), nil
+}
+
+// String creates a readable string representation of a chunk header.
+func (ch *ChunkHeader) String() string {
+	if ch == nil {
+		return "<nil>"
+	}
+	return fmt.Sprintf("handshake_id=%d noise_msg=%d chunk=%d/%d data_shards=%d",
+		ch.HandshakeID, ch.NoiseMsgNum, ch.ChunkIdx, ch.TotalChunks, ch.DataShards)
+}