avoid race potential when calling installSession for a lane

lanes.go

@@ -327,7 +327,8 @@ func lanePortOffset(myAddr, peerAddr netip.Addr, peerPortCount uint16) uint16 {
 // can name any lane, and installing on sight would let them make us hold a
 // replay window and two cipher states per lane without authenticating anything.
 // The caller must install with installSession once the packet decrypts, which is
-// the first moment the lane is known to be real.
+// the first moment the lane is known to be real, and must use the session
+// installSession hands back from then on.
 func (i *HostInfo) laneSession(s uint8) (ci *ConnectionState, cached bool, err error) {
 	ls := i.lanes
 	if ls == nil || s == 0 || int(s) >= len(ls.sessions) {
@@ -346,18 +347,21 @@ func (i *HostInfo) laneSession(s uint8) (ci *ConnectionState, cached bool, err e
 }
 
 // installSession publishes a session derived by laneSession, so the next packet
-// on the lane doesn't have to derive it again. cs must have already decrypted the
-// packet at messageCounter.
+// on the lane doesn't have to derive it again, and returns the session that is
+// now installed. cs must have already decrypted the packet at messageCounter.
 //
-// Two routines can race on a lane's first packet and derive a session each. The
-// loser's is dropped, and with it the replay-window entry for the packet it just
-// accepted, so hand that counter to the session that survives — the keys are
-// identical, so it is the same window in every respect that matters.
-func (ls *laneSet) installSession(l *slog.Logger, s uint8, cs *ConnectionState, messageCounter uint64) {
+// Two routines can race on a lane's first packet and derive a session each, and
+// only one of them can be installed. Every session derived for a lane holds the
+// same key, so a second one that ever encrypted would restart the message
+// counter and repeat nonces under it. Returning the winner solves this.
+func (ls *laneSet) installSession(l *slog.Logger, s uint8, cs *ConnectionState, messageCounter uint64) *ConnectionState {
 	if ls.sessions[s].CompareAndSwap(nil, cs) {
-		return
+		return cs
 	}
-	ls.sessions[s].Load().noteSeen(l, messageCounter)
+
+	winner := ls.sessions[s].Load()
+	winner.noteSeen(l, messageCounter)
+	return winner
 }
 
 // session returns lane s's session for our own use, deriving and installing it

lanes_test.go

@@ -238,7 +238,7 @@ func TestLaneSessionRxDerivation(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, []byte("real lane traffic"), pt)
 
-	respLS.installSession(test.NewLogger(), 2, ci, 1)
+	assert.Same(t, ci, respLS.installSession(test.NewLogger(), 2, ci, 1))
 	assert.Same(t, ci, respLS.sessions[2].Load())
 
 	// Now it is a hit, and the replay window the decrypt above advanced is the
@@ -255,7 +255,8 @@ func TestLaneSessionRxDerivation(t *testing.T) {
 	other, err := newLaneConnectionState(&respLS.material, 2)
 	require.NoError(t, err)
 	require.True(t, ci.window.Check(test.NewLogger(), 7), "counter 7 seen before the race")
-	respLS.installSession(test.NewLogger(), 2, other, 7)
+	assert.Same(t, ci, respLS.installSession(test.NewLogger(), 2, other, 7),
+		"the loser must be handed the installed session, never left holding its own")
 	assert.Same(t, ci, respLS.sessions[2].Load())
 	assert.False(t, ci.window.Check(test.NewLogger(), 7),
 		"the loser's counter was not carried to the surviving window")

outside.go

@@ -175,9 +175,8 @@ func (f *Interface) readOutsidePackets(via ViaSender, packet []byte, rxc *rxCont
 	}
 
 	if !laneCached {
-		// The packet decrypted, so the peer really is using this lane and the
-		// session we derived for it is worth keeping.
-		hostinfo.lanes.installSession(f.l, lane, ci, h.MessageCounter)
+		// The packet decrypted, so the peer really is using this lane and the session we derived for it is valid
+		ci = hostinfo.lanes.installSession(f.l, lane, ci, h.MessageCounter)
 	}
 
 	// Roam before we respond, but only on the base tunnel: a lane's source
@@ -192,7 +191,7 @@ func (f *Interface) readOutsidePackets(via ViaSender, packet []byte, rxc *rxCont
 	case header.Message:
 		switch h.Subtype {
 		case header.MessageNone:
-			f.handleOutsideMessagePacket(hostinfo, ci, h.MessageCounter, out, rxc)
+			f.handleOutsideMessagePacket(hostinfo, h.MessageCounter, out, rxc)
 		default:
 			hostinfo.logger(f.l).Error("IsValidSubType was true, but unexpected message subtype seen", "from", via, "header", h)
 			return
@@ -522,7 +521,7 @@ func parseV4(data []byte, incoming bool, fp *firewall.ParsedPacket) error {
 	return nil
 }
 
-func (f *Interface) handleOutsideMessagePacket(hostinfo *HostInfo, ci *ConnectionState, messageCounter uint64, out []byte, rxc *rxContext) {
+func (f *Interface) handleOutsideMessagePacket(hostinfo *HostInfo, messageCounter uint64, out []byte, rxc *rxContext) {
 	err := newPacket(out, true, rxc.fwPacket)
 	if err != nil {
 		hostinfo.logger(f.l).Warn("Error while validating inbound packet", "error", err, "packet", out)
@@ -540,7 +539,7 @@ func (f *Interface) handleOutsideMessagePacket(hostinfo *HostInfo, ci *Connectio
 		return
 	}
 
-	err = f.batchers[rxc.q].Commit(out, batch.SortKey{Epoch: ci.epoch, Counter: messageCounter}, rxc.fwPacket)
+	err = f.batchers[rxc.q].Commit(out, batch.SortKey{Epoch: hostinfo.ConnectionState.epoch, Counter: messageCounter}, rxc.fwPacket)
 	if err != nil {
 		f.l.Error("Failed to write to tun", "error", err)
 	}