From 2f604928c38f644fd5a59bdeac0afd62f2d76e7d Mon Sep 17 00:00:00 2001
From: rkonfj <rkonfj@gmail.com>
Date: Sun, 9 Jun 2024 13:59:12 +0800
Subject: [PATCH] rdt: do not panic

---
 cmd/pgcli/share/filemanager.go |  1 +
 rdt/rdt.go                     | 13 +++++++++++++
 2 files changed, 14 insertions(+)

diff --git a/cmd/pgcli/share/filemanager.go b/cmd/pgcli/share/filemanager.go
index c41b64a..83d59d9 100644
--- a/cmd/pgcli/share/filemanager.go
+++ b/cmd/pgcli/share/filemanager.go
@@ -54,6 +54,7 @@ func (fm *FileManager) Open(index int) (*os.File, error) {
 }
 
 func (fm *FileManager) HandleRequest(peerID string, conn net.Conn) {
+	defer conn.Close()
 	header := make([]byte, 4)
 	_, err := io.ReadFull(conn, header)
 	if err != nil || !bytes.Equal(header[:2], []byte{0, 0}) {
diff --git a/rdt/rdt.go b/rdt/rdt.go
index fce46eb..7050c9f 100644
--- a/rdt/rdt.go
+++ b/rdt/rdt.go
@@ -253,6 +253,11 @@ func (c *rdtConn) buildFrame(cmd byte, no uint32, length uint16, data []byte) []
 }
 
 func (c *rdtConn) recv(pkt []byte) {
+	defer func() {
+		if err := recover(); err != nil {
+			slog.Debug("Recv", "recover", err)
+		}
+	}()
 	no := binary.BigEndian.Uint32(pkt[1:5])
 	l := binary.BigEndian.Uint16(pkt[5:7])
 	slog.Debug("RDTRecv", "cmd", pkt[0], "no", no, "peer", c.remoteAddr, "len", len(pkt))
@@ -348,6 +353,11 @@ func (c *rdtConn) sendNCK(no uint32) {
 func (c *rdtConn) sendFIN() error {
 	exit := make(chan struct{})
 	go func() {
+		defer func() {
+			if err := recover(); err != nil {
+				slog.Debug("SendFIN", "recover", err)
+			}
+		}()
 		for range 5 {
 			select {
 			case <-exit:
@@ -533,6 +543,9 @@ func (l *RDTListener) recvPacket(pkt []byte, addr net.Addr) {
 		conn.recv(pkt)
 		return
 	}
+	if pkt[0] != 0 {
+		return
+	}
 	conn = l.newConn(addr)
 	conn.server = true
 	l.acceptConnMap[addr.String()] = conn