From e59351f8faf3ac4172a46e4fd9aa895626cab706 Mon Sep 17 00:00:00 2001
From: soypat <graded.sp@gmail.com>
Date: Sat, 18 Nov 2023 19:28:22 -0300
Subject: [PATCH] need to fix broken Send/Recv test

---
 control.go          |  1 +
 stack/socket_tcp.go | 23 ++++++++++++++++++++++-
 stack/stack_test.go | 18 ++++++++++++++++++
 3 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/control.go b/control.go
index 014b3c6..c0e7e2a 100644
--- a/control.go
+++ b/control.go
@@ -105,6 +105,7 @@ func (seg *Segment) Last() Value {
 }
 
 // PendingSegment calculates a suitable next segment to send from a payload length.
+// It does not modify the ControlBlock state.
 func (tcb *ControlBlock) PendingSegment(payloadLen int) (_ Segment, ok bool) {
 	if (payloadLen == 0 && tcb.pending == 0) || (payloadLen > 0 && tcb.state != StateEstablished) {
 		return Segment{}, false // No pending segment.
diff --git a/stack/socket_tcp.go b/stack/socket_tcp.go
index c72f9ac..acb0c87 100644
--- a/stack/socket_tcp.go
+++ b/stack/socket_tcp.go
@@ -10,6 +10,8 @@ import (
 	"github.com/soypat/seqs"
 )
 
+const defaultSocketSize = 2048
+
 type tcp struct {
 	stack     *PortStack
 	scb       seqs.ControlBlock
@@ -30,7 +32,7 @@ func (t *tcp) State() seqs.State {
 func (t *tcp) Send(b []byte) error {
 	if t.tx.buf == nil {
 		t.tx = ring{
-			buf: make([]byte, max(2048, len(b))),
+			buf: make([]byte, max(defaultSocketSize, len(b))),
 		}
 	}
 	_, err := t.tx.Write(b)
@@ -40,6 +42,11 @@ func (t *tcp) Send(b []byte) error {
 	return nil
 }
 
+func (t *tcp) Recv(b []byte) (int, error) {
+	n, err := t.rx.Read(b)
+	return n, err
+}
+
 // DialTCP opens an active TCP connection to the given remote address.
 func DialTCP(stack *PortStack, localPort uint16, remoteMAC [6]byte, remote netip.AddrPort, iss seqs.Value, window seqs.Size) (*tcp, error) {
 	t := tcp{
@@ -110,6 +117,20 @@ func (t *tcp) handleRecv(response []byte, pkt *TCPPacket) (n int, err error) {
 	// if segIncoming.SEQ != t.scb.RecvNext() {
 	// 	return 0, ErrDroppedPacket // SCB does not admit out-of-order packets.
 	// }
+	if segIncoming.Flags.HasAny(seqs.FlagPSH) {
+		if len(payload) != int(segIncoming.DATALEN) {
+			return 0, errors.New("segment data length does not match payload length")
+		}
+		if t.rx.buf == nil {
+			t.rx = ring{
+				buf: make([]byte, defaultSocketSize),
+			}
+		}
+		_, err = t.rx.Write(payload)
+		if err != nil {
+			return 0, err
+		}
+	}
 	err = t.scb.Recv(segIncoming)
 	if err != nil {
 		return 0, err
diff --git a/stack/stack_test.go b/stack/stack_test.go
index ead4661..545b1fc 100644
--- a/stack/stack_test.go
+++ b/stack/stack_test.go
@@ -117,6 +117,24 @@ func TestStackSendReceive(t *testing.T) {
 		t.Fatal("not established")
 	}
 
+	// Send data from client to server.
+	const data = "hello world"
+	err = clientTCP.Send([]byte(data))
+	if err != nil {
+		t.Fatal(err)
+	}
+	txStacks(t, 1, Client, Server)
+	if clientTCP.State() != seqs.StateEstablished || serverTCP.State() != seqs.StateEstablished {
+		t.Fatal("not established")
+	}
+	var buf [len(data)]byte
+	n, err := serverTCP.Recv(buf[:])
+	if err != nil {
+		t.Fatal(err)
+	}
+	if string(buf[:n]) != data {
+		t.Error("got", string(buf[:n]), "want", data)
+	}
 }
 
 func isDroppedPacket(err error) bool {