From 6c65764070b24f29be322f19a4f72544cf195c4a Mon Sep 17 00:00:00 2001
From: soypat <graded.sp@gmail.com>
Date: Fri, 30 May 2025 21:33:12 -0300
Subject: [PATCH] add maximum segment size option on TCP connection init

---
 tcp/definitions.go | 45 +++++++++++++++++++++++++++++++++++++++------
 tcp/handler.go     | 16 ++++++++++++----
 2 files changed, 51 insertions(+), 10 deletions(-)

diff --git a/tcp/definitions.go b/tcp/definitions.go
index 6c8b19a..c9117e2 100644
--- a/tcp/definitions.go
+++ b/tcp/definitions.go
@@ -378,15 +378,48 @@ func (kind OptionKind) IsDefined() bool {
 	return kind <= 30 || kind == 34 || kind == 69 || kind == 172 || kind == 174
 }
 
-type OptionParser struct {
-	SkipSizeValidation bool
-	SkipObsolete       bool
+type OptionCodec struct {
+	Flags OptionFlags
 }
 
-func (op *OptionParser) ForEachOption(opts []byte, fn func(OptionKind, []byte) error) error {
+type OptionFlags uint8
+
+const (
+	OptFlagSkipSizeValidation OptionFlags = 1 << iota
+	OptFlagSkipObsolete
+)
+
+func (flags OptionFlags) HasAny(ofTheseFlags OptionFlags) bool {
+	return flags&ofTheseFlags != 0
+}
+
+func (op OptionCodec) PutOption16(dst []byte, kind OptionKind, v uint16) (int, error) {
+	return op.PutOption(dst, kind, byte(v>>8), byte(v))
+}
+
+func (op OptionCodec) PutOption32(dst []byte, kind OptionKind, v uint32) (int, error) {
+	return op.PutOption(dst, kind, byte(v>>24), byte(v>>16), byte(v>>7), byte(v))
+}
+
+func (op OptionCodec) PutOption(dst []byte, kind OptionKind, data ...byte) (int, error) {
+	putSize := 2 + len(data)
+	if len(dst) < putSize {
+		return -1, errBufferTooSmall
+	} else if putSize > 255 {
+		return -1, errors.New("option data too large")
+	} else if kind == OptNop || kind == OptEnd {
+		return -1, errors.New("cant put Nop or End option type")
+	}
+	dst[0] = byte(kind)
+	dst[1] = byte(putSize)
+	copy(dst[2:], data)
+	return putSize, nil
+}
+
+func (op OptionCodec) ForEachOption(opts []byte, fn func(OptionKind, []byte) error) error {
 	off := 0
-	skipSizeValidation := op.SkipSizeValidation
-	skipObsolete := op.SkipObsolete
+	skipSizeValidation := op.Flags.HasAny(OptFlagSkipSizeValidation)
+	skipObsolete := op.Flags.HasAny(OptFlagSkipObsolete)
 	for off < len(opts) && opts[off] != 0 {
 		kind := OptionKind(opts[off])
 		off++
diff --git a/tcp/handler.go b/tcp/handler.go
index 234b2f3..07e1eae 100644
--- a/tcp/handler.go
+++ b/tcp/handler.go
@@ -30,8 +30,9 @@ type Handler struct {
 	// connid is a conenction counter that is incremented each time a new
 	// connection is established via Open calls. This disambiguate's whether
 	// Read and Write calls belong to the current connection.
-	connid  uint16
-	closing bool
+	connid   uint16
+	optcodec OptionCodec
+	closing  bool
 }
 
 func (h *Handler) SetLoggers(handler, scb *slog.Logger) {
@@ -206,10 +207,14 @@ func (h *Handler) Send(b []byte) (int, error) {
 	if err != nil {
 		return 0, err
 	}
+
+	offset := uint8(5)
 	var segment Segment
 	if h.AwaitingSynSend() {
 		// Handling init syn segment.
 		segment = ClientSynSegment(h.scb.ISS(), h.scb.RecvWindow())
+		h.optcodec.PutOption16(b[sizeHeaderTCP:], OptMaxSegmentSize, uint16(len(b)))
+		offset++
 	} else {
 		var ok bool
 		available := min(h.bufTx.Buffered(), len(b)-sizeHeaderTCP)
@@ -225,6 +230,9 @@ func (h *Handler) Send(b []byte) (int, error) {
 			} else if n != int(segment.DATALEN) {
 				panic("expected n == available")
 			}
+		} else if segment.Flags == synack {
+			h.optcodec.PutOption16(b[sizeHeaderTCP:], OptMaxSegmentSize, uint16(len(b)))
+			offset++
 		}
 	}
 	prevState := h.scb.State()
@@ -236,9 +244,9 @@ func (h *Handler) Send(b []byte) (int, error) {
 	}
 	tfrm.SetSourcePort(h.localPort)
 	tfrm.SetDestinationPort(h.remotePort)
-	tfrm.SetSegment(segment, 5) // No TCP options.
+	tfrm.SetSegment(segment, offset) // No TCP options.
 	tfrm.SetUrgentPtr(0)
-	return sizeHeaderTCP + int(segment.DATALEN), nil
+	return int(offset)*4 + int(segment.DATALEN), nil
 }
 
 // Free returns the amount of space free in the transmit buffer. A call to [Handler.Write] with a larger buffer will fail.