fix race condition

handler/forward/local/handler.go

@@ -2,7 +2,6 @@ package local
 
 import (
 	"bufio"
-	"bytes"
 	"context"
 	"crypto/tls"
 	"errors"
@@ -195,7 +194,7 @@ func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriter, log l
 		err = func() error {
 			req, err := http.ReadRequest(br)
 			if err != nil {
-				log.Errorf("read http request: %v", err)
+				// log.Errorf("read http request: %v", err)
 				return err
 			}
 
@@ -268,10 +267,14 @@ func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriter, log l
 				})
 			}
 
+			if err := req.Write(cc); err != nil {
+				cc.Close()
+				log.Warnf("send request to node %s(%s): %v", target.Name, target.Addr, err)
+				return resp.Write(rw)
+			}
+
 			if req.Header.Get("Upgrade") == "websocket" {
-				var buf bytes.Buffer
+				err := xnet.Transport(cc, xio.NewReadWriter(br, rw))
-				req.Write(&buf)
-				err := xnet.Transport(cc, xio.NewReadWriter(io.MultiReader(&buf, br), rw))
 				if err == nil {
 					err = io.EOF
 				}
@@ -281,12 +284,6 @@ func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriter, log l
 			go func() {
 				defer cc.Close()
 
-				if err := req.Write(cc); err != nil {
-					log.Warnf("send request to node %s(%s): %v", target.Name, target.Addr, err)
-					resp.Write(rw)
-					return
-				}
-
 				res, err := http.ReadResponse(bufio.NewReader(cc), req)
 				if err != nil {
 					log.Warnf("read response from node %s(%s): %v", target.Name, target.Addr, err)

handler/forward/remote/handler.go

@@ -2,7 +2,6 @@ package remote
 
 import (
 	"bufio"
-	"bytes"
 	"context"
 	"crypto/tls"
 	"errors"
@@ -197,7 +196,7 @@ func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriter, remot
 		err = func() error {
 			req, err := http.ReadRequest(br)
 			if err != nil {
-				log.Errorf("read http request: %v", err)
+				// log.Errorf("read http request: %v", err)
 				return err
 			}
 
@@ -272,10 +271,14 @@ func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriter, remot
 
 			cc = proxyproto.WrapClientConn(h.md.proxyProtocol, remoteAddr, localAddr, cc)
 
+			if err := req.Write(cc); err != nil {
+				cc.Close()
+				log.Warnf("send request to node %s(%s): %v", target.Name, target.Addr, err)
+				return resp.Write(rw)
+			}
+
 			if req.Header.Get("Upgrade") == "websocket" {
-				var buf bytes.Buffer
+				err := xnet.Transport(cc, xio.NewReadWriter(br, rw))
-				req.Write(&buf)
-				err := xnet.Transport(cc, xio.NewReadWriter(io.MultiReader(&buf, br), rw))
 				if err == nil {
 					err = io.EOF
 				}
@@ -285,12 +288,6 @@ func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriter, remot
 			go func() {
 				defer cc.Close()
 
-				if err := req.Write(cc); err != nil {
-					log.Warnf("send request to node %s(%s): %v", target.Name, target.Addr, err)
-					resp.Write(rw)
-					return
-				}
-
 				res, err := http.ReadResponse(bufio.NewReader(cc), req)
 				if err != nil {
 					log.Warnf("read response from node %s(%s): %v", target.Name, target.Addr, err)

listener/rtcp/listener.go

@@ -3,6 +3,7 @@ package rtcp
 import (
 	"context"
 	"net"
+	"sync"
 
 	"github.com/go-gost/core/chain"
 	"github.com/go-gost/core/listener"
@@ -27,6 +28,7 @@ type rtcpListener struct {
 	logger  logger.Logger
 	closed  chan struct{}
 	options listener.Options
+	mu      sync.Mutex
 }
 
 func NewListener(opts ...listener.Option) listener.Listener {
@@ -72,23 +74,25 @@ func (l *rtcpListener) Accept() (conn net.Conn, err error) {
 	default:
 	}
 
-	if l.ln == nil {
+	ln := l.getListener()
-		l.ln, err = l.router.Bind(
+	if ln == nil {
+		ln, err = l.router.Bind(
 			context.Background(), "tcp", l.laddr.String(),
 			chain.MuxBindOption(true),
 		)
 		if err != nil {
 			return nil, listener.NewAcceptError(err)
 		}
-		l.ln = metrics.WrapListener(l.options.Service, l.ln)
+		ln = metrics.WrapListener(l.options.Service, ln)
-		l.ln = admission.WrapListener(l.options.Admission, l.ln)
+		ln = admission.WrapListener(l.options.Admission, ln)
-		l.ln = limiter.WrapListener(l.options.TrafficLimiter, l.ln)
+		ln = limiter.WrapListener(l.options.TrafficLimiter, ln)
-		l.ln = climiter.WrapListener(l.options.ConnLimiter, l.ln)
+		ln = climiter.WrapListener(l.options.ConnLimiter, ln)
+		l.setListener(ln)
 	}
 	conn, err = l.ln.Accept()
 	if err != nil {
-		l.ln.Close()
+		ln.Close()
-		l.ln = nil
+		l.setListener(nil)
 		return nil, listener.NewAcceptError(err)
 	}
 	return
@@ -103,8 +107,9 @@ func (l *rtcpListener) Close() error {
 	case <-l.closed:
 	default:
 		close(l.closed)
-		if l.ln != nil {
+		ln := l.getListener()
-			l.ln.Close()
+		if ln != nil {
+			ln.Close()
 			// l.ln = nil
 		}
 	}
@@ -112,6 +117,18 @@ func (l *rtcpListener) Close() error {
 	return nil
 }
 
+func (l *rtcpListener) setListener(ln net.Listener) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	l.ln = ln
+}
+
+func (l *rtcpListener) getListener() net.Listener {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	return l.ln
+}
+
 type bindAddr struct {
 	addr string
 }