Fix TrustTunnel

transport/trusttunnel/packet.go

@@ -94,18 +94,7 @@ func (u *clientPacketConn) writePacketToServer(buffer *buf.Buffer, source M.Sock
 	common.Must(binary.Write(header, binary.BigEndian, source.Port))
 	common.Must(binary.Write(header, binary.BigEndian, uint8(len(appName))))
 	common.Must1(header.WriteString(appName))
-	_, err := u.writer.Write(header.Bytes())
+	return u.writeChunks(header.Bytes(), buffer.Bytes())
-	if err != nil {
-		return err
-	}
-	_, err = u.writer.Write(buffer.Bytes())
-	if err != nil {
-		return err
-	}
-	if u.flusher != nil {
-		u.flusher.Flush()
-	}
-	return nil
 }
 
 var (
@@ -195,16 +184,5 @@ func (u *serverPacketConn) writePacketToClient(buffer *buf.Buffer, source M.Sock
 	common.Must1(header.Write(sourceAddress[:]))
 	common.Must(binary.Write(header, binary.BigEndian, source.Port))
 	common.Must(header.WriteZeroN(16 + 2))
-	_, err := u.writer.Write(header.Bytes())
+	return u.writeChunks(header.Bytes(), buffer.Bytes())
-	if err != nil {
-		return err
-	}
-	_, err = u.writer.Write(buffer.Bytes())
-	if err != nil {
-		return err
-	}
-	if u.flusher != nil {
-		u.flusher.Flush()
-	}
-	return nil
 }

transport/trusttunnel/protocol.go

@@ -71,6 +71,9 @@ type httpConn struct {
 	localAddr  net.Addr
 	deadline   *time.Timer
 	done       chan struct{}
+	closed     bool
+
+	mtx sync.Mutex
 }
 
 func (h *httpConn) setup(body io.ReadCloser, err error) {
@@ -93,6 +96,9 @@ func (h *httpConn) waitCreated() error {
 }
 
 func (h *httpConn) Close() error {
+	h.mtx.Lock()
+	h.closed = true
+	h.mtx.Unlock()
 	h.setup(nil, net.ErrClosed)
 	if closer, ok := h.writer.(io.Closer); ok {
 		_ = closer.Close()
@@ -117,11 +123,28 @@ func (h *httpConn) Close() error {
 }
 
 func (h *httpConn) writeFlush(p []byte) (n int, err error) {
-	n, err = h.writer.Write(p)
+	err = h.writeChunks(p)
+	if err != nil {
+		return 0, err
+	}
+	return len(p), nil
+}
+
+func (h *httpConn) writeChunks(chunks ...[]byte) error {
+	h.mtx.Lock()
+	defer h.mtx.Unlock()
+	if h.closed {
+		return net.ErrClosed
+	}
+	for _, chunk := range chunks {
+		if _, err := h.writer.Write(chunk); err != nil {
+			return err
+		}
+	}
 	if h.flusher != nil {
 		h.flusher.Flush()
 	}
-	return n, err
+	return nil
 }
 
 func (h *httpConn) RemoteAddr() net.Addr {