From 9268e6752909c039ee72765b002e3b4d6664d3c3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=B2=B3?= <cnliuhe@126.com>
Date: Sat, 7 Dec 2019 20:44:32 +0800
Subject: [PATCH] fix port mux bug

---
 lib/mux/pconn.go | 24 ++++++++++++++++--------
 lib/mux/pmux.go  |  4 +++-
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/lib/mux/pconn.go b/lib/mux/pconn.go
index b351e7c..35af3cc 100644
--- a/lib/mux/pconn.go
+++ b/lib/mux/pconn.go
@@ -6,15 +6,17 @@ import (
 )
 
 type PortConn struct {
-	Conn  net.Conn
-	rs    []byte
-	start int
+	Conn     net.Conn
+	rs       []byte
+	readMore bool
+	start    int
 }
 
-func newPortConn(conn net.Conn, rs []byte) *PortConn {
+func newPortConn(conn net.Conn, rs []byte, readMore bool) *PortConn {
 	return &PortConn{
-		Conn: conn,
-		rs:   rs,
+		Conn:     conn,
+		rs:       rs,
+		readMore: readMore,
 	}
 }
 
@@ -29,9 +31,15 @@ func (pConn *PortConn) Read(b []byte) (n int, err error) {
 		defer func() {
 			pConn.start = len(pConn.rs)
 		}()
-		return copy(b, pConn.rs[pConn.start:]), nil
+		n = copy(b, pConn.rs[pConn.start:])
+		if !pConn.readMore {
+			return
+		}
 	}
-	return pConn.Conn.Read(b)
+	var n2 = 0
+	n2, err = pConn.Conn.Read(b[n:])
+	n = n + n2
+	return
 }
 
 func (pConn *PortConn) Write(b []byte) (n int, err error) {
diff --git a/lib/mux/pmux.go b/lib/mux/pmux.go
index e593bb0..f4e9d4a 100644
--- a/lib/mux/pmux.go
+++ b/lib/mux/pmux.go
@@ -89,6 +89,7 @@ func (pMux *PortMux) process(conn net.Conn) {
 	var ch chan *PortConn
 	var rs []byte
 	var buffer bytes.Buffer
+	var readMore = false
 	switch common.BytesToNum(buf) {
 	case HTTP_CONNECT, HTTP_DELETE, HTTP_GET, HTTP_HEAD, HTTP_OPTIONS, HTTP_POST, HTTP_PUT, HTTP_TRACE: //http and manager
 		buffer.Reset()
@@ -123,6 +124,7 @@ func (pMux *PortMux) process(conn net.Conn) {
 	case CLIENT: // client connection
 		ch = pMux.clientConn
 	default: // https
+		readMore = true
 		ch = pMux.httpsConn
 	}
 	if len(rs) == 0 {
@@ -131,7 +133,7 @@ func (pMux *PortMux) process(conn net.Conn) {
 	timer := time.NewTimer(ACCEPT_TIME_OUT)
 	select {
 	case <-timer.C:
-	case ch <- newPortConn(conn, rs):
+	case ch <- newPortConn(conn, rs, readMore):
 	}
 }