From 6c0dde5a0093db1dafa4af7cd708f7a06184dbf3 Mon Sep 17 00:00:00 2001
From: lwch <lwch748@gmail.com>
Date: Tue, 24 Aug 2021 18:26:57 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E9=87=8D=E8=AF=95=E9=80=BB?=
 =?UTF-8?q?=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 code/client/pool/conn.go      | 24 +++++++++++++++++++-----
 code/server/handler/client.go | 12 +++++++++++-
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/code/client/pool/conn.go b/code/client/pool/conn.go
index 17d786a..a70f77f 100644
--- a/code/client/pool/conn.go
+++ b/code/client/pool/conn.go
@@ -2,6 +2,7 @@ package pool
 
 import (
 	"context"
+	"errors"
 	"natpass/code/network"
 	"strings"
 	"sync"
@@ -126,6 +127,20 @@ func (conn *Conn) loopRead(cancel context.CancelFunc) {
 	}
 }
 
+func loopWrite(conn *network.Conn, msg *network.Msg, timeout time.Duration) error {
+	var err error
+	for i := 0; i < 10; i++ {
+		err = conn.WriteMessage(msg, timeout)
+		if err != nil {
+			if strings.Contains(err.Error(), "i/o timeout") {
+				continue
+			}
+			return err
+		}
+	}
+	return errors.New("too many retry")
+}
+
 func (conn *Conn) loopWrite(cancel context.CancelFunc) {
 	defer conn.Close()
 	defer cancel()
@@ -135,12 +150,11 @@ func (conn *Conn) loopWrite(cancel context.CancelFunc) {
 			return
 		}
 		msg.From = conn.ID
-		err := conn.conn.WriteMessage(msg, conn.parent.cfg.WriteTimeout)
-		if err == nil {
-			continue
+		err := loopWrite(conn.conn, msg, conn.parent.cfg.WriteTimeout)
+		if err != nil {
+			logging.Error("write message error on %s: %v", conn.ID, err)
+			return
 		}
-		logging.Error("write message error on %s: %v", conn.ID, err)
-		return
 	}
 }
 
diff --git a/code/server/handler/client.go b/code/server/handler/client.go
index 0844d1a..4c4f7ff 100644
--- a/code/server/handler/client.go
+++ b/code/server/handler/client.go
@@ -1,6 +1,7 @@
 package handler
 
 import (
+	"errors"
 	"natpass/code/network"
 	"strings"
 	"sync"
@@ -57,7 +58,16 @@ func (c *client) run() {
 }
 
 func (c *client) writeMessage(msg *network.Msg) error {
-	return c.c.WriteMessage(msg, c.parent.cfg.WriteTimeout)
+	for i := 0; i < 10; i++ {
+		err := c.c.WriteMessage(msg, c.parent.cfg.WriteTimeout)
+		if err != nil {
+			if strings.Contains(err.Error(), "i/o timeout") {
+				continue
+			}
+			return err
+		}
+	}
+	return errors.New("too many retry")
 }
 
 func (c *client) addLink(id string) {