diff --git a/transport/trusttunnel/packet.go b/transport/trusttunnel/packet.go index 068823de..8610d8ef 100644 --- a/transport/trusttunnel/packet.go +++ b/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()) - 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 + return u.writeChunks(header.Bytes(), buffer.Bytes()) } 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()) - 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 + return u.writeChunks(header.Bytes(), buffer.Bytes()) } diff --git a/transport/trusttunnel/protocol.go b/transport/trusttunnel/protocol.go index cea14fcd..047c96c8 100644 --- a/transport/trusttunnel/protocol.go +++ b/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 {