From c1efd9f91b1fd9b6271bcf20fa697b4c1df0e4ab Mon Sep 17 00:00:00 2001
From: Josh Deprez <josh.deprez@gmail.com>
Date: Fri, 19 Apr 2024 12:21:33 +1000
Subject: [PATCH] Fix ZIP ATP replies

---
 README.md                |  5 +++--
 atalk/atp/atp.go         | 10 ++++++++++
 atalk/raw.go             |  9 ++++++++-
 atalk/zip/getzonelist.go |  4 ++++
 zip.go                   | 35 +++++++++++++++++++++++++++++++----
 5 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 4a95ee5..2d1c0ac 100644
--- a/README.md
+++ b/README.md
@@ -22,8 +22,9 @@ Things I plan to fix Real Soon Now:
   interface. I haven't tested it with netatalk or tashrouter on the same
   host, but I think using a distinct Ethernet address would help them coexist.
 * It doesn't do any of the required packet splitting to keep packets under the
-  AppleTalk size limits. In particular ZIP GetZoneList Replies are incorrect
-  when the zone list would exceed the limit.
+  AppleTalk size limits. ~~In particular ZIP GetZoneList Replies are incorrect
+  when the zone list would exceed the limit.~~ GetZoneList is now fixed, but
+  various others need fixing.
 * It logs a lot and has no other monitoring or observability capability. I plan
   to add a Prometheus metrics endpoint and at least add log levels / verbosity
   flag.
diff --git a/atalk/atp/atp.go b/atalk/atp/atp.go
index aed07c1..4ecf8ed 100644
--- a/atalk/atp/atp.go
+++ b/atalk/atp/atp.go
@@ -20,6 +20,13 @@ import (
 	"bytes"
 	"encoding/binary"
 	"fmt"
+
+	"gitea.drjosh.dev/josh/jrouter/atalk"
+)
+
+const (
+	HeaderSize  = 8
+	MaxDataSize = atalk.DDPMaxDataSize - HeaderSize // 578
 )
 
 const (
@@ -50,6 +57,9 @@ func (p *TReq) Marshal() ([]byte, error) {
 	if p.TRelTimeoutIndicator > 4 {
 		return nil, fmt.Errorf("invalid TRel timeout indicator [%d > 4]", p.TRelTimeoutIndicator)
 	}
+	if len(p.Data) > MaxDataSize {
+		return nil, fmt.Errorf("%w [%d > %d]", atalk.ErrDataTooBig, len(p.Data), MaxDataSize)
+	}
 
 	b := bytes.NewBuffer(nil)
 	ci := byte(ciFuncTReq)
diff --git a/atalk/raw.go b/atalk/raw.go
index ab05b44..7b70637 100644
--- a/atalk/raw.go
+++ b/atalk/raw.go
@@ -16,4 +16,11 @@
 
 package atalk
 
-const DDPExtHeaderSize = 13
+import "errors"
+
+const (
+	DDPExtHeaderSize = 13
+	DDPMaxDataSize   = 586
+)
+
+var ErrDataTooBig = errors.New("data too big")
diff --git a/atalk/zip/getzonelist.go b/atalk/zip/getzonelist.go
index 2e86d64..04bf541 100644
--- a/atalk/zip/getzonelist.go
+++ b/atalk/zip/getzonelist.go
@@ -21,6 +21,7 @@ import (
 	"encoding/binary"
 	"fmt"
 
+	"gitea.drjosh.dev/josh/jrouter/atalk"
 	"gitea.drjosh.dev/josh/jrouter/atalk/atp"
 )
 
@@ -78,6 +79,9 @@ func (p *GetZonesReplyPacket) MarshalTResp() (*atp.TResp, error) {
 		b.WriteString(z)
 	}
 	r.Data = b.Bytes()
+	if len(r.Data) > atp.MaxDataSize {
+		return nil, fmt.Errorf("%w [%d > %d]", atalk.ErrDataTooBig, len(r.Data), atp.MaxDataSize)
+	}
 	return r, nil
 }
 
diff --git a/zip.go b/zip.go
index 0354458..a4e026f 100644
--- a/zip.go
+++ b/zip.go
@@ -32,7 +32,7 @@ import (
 
 func handleZIP(pcapHandle *pcap.Handle, srcHWAddr, myHWAddr ethernet.Addr, myAddr aarp.AddrPair, cfg *config, zones *ZoneTable, ddpkt *ddp.ExtPacket) error {
 	switch ddpkt.Proto {
-	case 3: // ATP
+	case ddp.ProtoATP:
 		atpkt, err := atp.UnmarshalPacket(ddpkt.Data)
 		if err != nil {
 			return err
@@ -43,11 +43,15 @@ func handleZIP(pcapHandle *pcap.Handle, srcHWAddr, myHWAddr ethernet.Addr, myAdd
 			if err != nil {
 				return err
 			}
-			// TODO: handle this in a more transactiony way
+			if gzl.StartIndex == 0 {
+				return fmt.Errorf("ZIP ATP: received request with StartIndex = 0 (invalid)")
+			}
+
 			resp := &zip.GetZonesReplyPacket{
 				TID:      gzl.TID,
-				LastFlag: true, // TODO: support multiple response packets
+				LastFlag: true,
 			}
+
 			switch gzl.Function {
 			case zip.FunctionGetZoneList:
 				resp.Zones = zones.AllNames()
@@ -59,6 +63,29 @@ func handleZIP(pcapHandle *pcap.Handle, srcHWAddr, myHWAddr ethernet.Addr, myAdd
 				resp.Zones = []string{cfg.EtherTalk.ZoneName}
 			}
 
+			// Inside AppleTalk SE, pp 8-8
+			if int(gzl.StartIndex) > len(resp.Zones) {
+				// "Note: A 0-byte response will be returned by a router if the
+				// index specified in the request is greater than the index of
+				// the last zone in the list (and the user bytes field will
+				// indicate no more zones)."
+				resp.Zones = nil
+			} else {
+				// Trim the zones list
+				// "zone names in the router are assumed to be numbered starting
+				// with 1"
+				resp.Zones = resp.Zones[gzl.StartIndex-1:]
+				size := 0
+				for i, z := range resp.Zones {
+					size += 1 + len(z) // length prefix plus string
+					if size > atp.MaxDataSize {
+						resp.LastFlag = false
+						resp.Zones = resp.Zones[:i]
+						break
+					}
+				}
+			}
+
 			respATP, err := resp.MarshalTResp()
 			if err != nil {
 				return err
@@ -99,7 +126,7 @@ func handleZIP(pcapHandle *pcap.Handle, srcHWAddr, myHWAddr ethernet.Addr, myAdd
 			return fmt.Errorf("unsupported ATP packet type %T for ZIP", atpkt)
 		}
 
-	case 6: // ZIP
+	case ddp.ProtoZIP:
 		zipkt, err := zip.UnmarshalPacket(ddpkt.Data)
 		if err != nil {
 			return err