From d45678b81245077e8cb4b4a7104444e3548df671 Mon Sep 17 00:00:00 2001
From: Josh Deprez <josh.deprez@gmail.com>
Date: Mon, 1 Apr 2024 15:14:31 +1100
Subject: [PATCH] Tweak ZI-Rsp subcode logic

---
 aurp/transport.go | 23 ++++++++++++++---------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/aurp/transport.go b/aurp/transport.go
index a2a4949..40ede57 100644
--- a/aurp/transport.go
+++ b/aurp/transport.go
@@ -182,15 +182,20 @@ func (tr *Transport) NewRIAckPacket(connID, seq uint16, szi RoutingFlag) *RIAckP
 // zone information. It automatically chooses between subcodes 1 or 2 depending
 // on whether there is one network ID or more than one network ID.
 func (tr *Transport) NewZIRspPacket(zones ZoneTuples) *ZIRspPacket {
-	nns := make(map[uint16]struct{})
-	for _, z := range zones {
-		nns[z.Network] = struct{}{}
-	}
-	// Only one network: use extended
-	// More than one network: use non-extended
-	subcode := SubcodeZoneInfoExt
-	if len(nns) != 1 {
-		subcode = SubcodeZoneInfoNonExt
+	// Only one zone: use non-extended
+	subcode := SubcodeZoneInfoNonExt
+	if len(zones) > 1 {
+		// Count distinct networks
+		nns := make(map[uint16]struct{})
+		for _, z := range zones {
+			nns[z.Network] = struct{}{}
+		}
+
+		// Only one network: use extended format
+		// More than one network: use non-extended
+		if len(nns) == 1 {
+			subcode = SubcodeZoneInfoExt
+		}
 	}
 
 	return &ZIRspPacket{