fix xnet.Pipe

This commit is contained in:
ginuerzh
2025-08-01 23:00:50 +08:00
parent a5309eee97
commit db21de831a
22 changed files with 122 additions and 51 deletions
+21 -2
View File
@@ -1,10 +1,15 @@
package io
import (
"errors"
"io"
"time"
)
var (
ErrUnsupported = errors.New("unsupported")
)
type CloseRead interface {
CloseRead() error
}
@@ -29,14 +34,14 @@ func (rw *readWriter) CloseRead() error {
if sc, ok := rw.Writer.(CloseRead); ok {
return sc.CloseRead()
}
return nil
return ErrUnsupported
}
func (rw *readWriter) CloseWrite() error {
if sc, ok := rw.Writer.(CloseWrite); ok {
return sc.CloseWrite()
}
return nil
return ErrUnsupported
}
type readWriteCloser struct {
@@ -53,6 +58,20 @@ func NewReadWriteCloser(r io.Reader, w io.Writer, c io.Closer) io.ReadWriteClose
}
}
func (rwc *readWriteCloser) CloseRead() error {
if sc, ok := rwc.Writer.(CloseRead); ok {
return sc.CloseRead()
}
return ErrUnsupported
}
func (rwc *readWriteCloser) CloseWrite() error {
if sc, ok := rwc.Writer.(CloseWrite); ok {
return sc.CloseWrite()
}
return ErrUnsupported
}
type setReadDeadline interface {
SetReadDeadline(t time.Time) error
}
+2 -2
View File
@@ -129,12 +129,12 @@ func (c *readWriteConn) CloseRead() error {
if sc, ok := c.Conn.(xio.CloseRead); ok {
return sc.CloseRead()
}
return nil
return xio.ErrUnsupported
}
func (c *readWriteConn) CloseWrite() error {
if sc, ok := c.Conn.(xio.CloseWrite); ok {
return sc.CloseWrite()
}
return nil
return xio.ErrUnsupported
}
+8 -3
View File
@@ -15,7 +15,7 @@ const (
tcpWaitTimeout = 10 * time.Second
)
func Pipe(ctx context.Context, rw1, rw2 io.ReadWriter) error {
func Pipe(ctx context.Context, rw1, rw2 io.ReadWriteCloser) error {
wg := sync.WaitGroup{}
wg.Add(2)
@@ -55,7 +55,7 @@ func Pipe(ctx context.Context, rw1, rw2 io.ReadWriter) error {
return nil
}
func pipeBuffer(dst io.ReadWriter, src io.ReadWriter, bufferSize int) error {
func pipeBuffer(dst io.ReadWriteCloser, src io.ReadWriteCloser, bufferSize int) error {
buf := bufpool.Get(bufferSize)
defer bufpool.Put(buf)
@@ -65,8 +65,13 @@ func pipeBuffer(dst io.ReadWriter, src io.ReadWriter, bufferSize int) error {
if cr, ok := src.(xio.CloseRead); ok {
cr.CloseRead()
}
if cw, ok := dst.(xio.CloseWrite); ok {
cw.CloseWrite()
if e := cw.CloseWrite(); e == xio.ErrUnsupported {
dst.Close()
}
} else {
dst.Close()
}
// Set TCP half-close timeout.
+29
View File
@@ -3,9 +3,38 @@ package proxyproto
import (
"net"
xio "github.com/go-gost/x/internal/io"
proxyproto "github.com/pires/go-proxyproto"
)
type serverConn struct {
net.Conn
}
func (c *serverConn) CloseRead() error {
if sc, ok := c.Conn.(xio.CloseRead); ok {
return sc.CloseRead()
}
if conn, ok := c.Conn.(*proxyproto.Conn); ok {
if tcpConn, ok := conn.TCPConn(); ok {
return tcpConn.CloseRead()
}
}
return xio.ErrUnsupported
}
func (c *serverConn) CloseWrite() error {
if sc, ok := c.Conn.(xio.CloseWrite); ok {
return sc.CloseWrite()
}
if conn, ok := c.Conn.(*proxyproto.Conn); ok {
if tcpConn, ok := conn.TCPConn(); ok {
return tcpConn.CloseWrite()
}
}
return xio.ErrUnsupported
}
func WrapClientConn(ppv int, src, dst net.Addr, c net.Conn) net.Conn {
if ppv <= 0 {
return c
+17 -3
View File
@@ -7,13 +7,27 @@ import (
proxyproto "github.com/pires/go-proxyproto"
)
type listener struct {
net.Listener
}
func (ln *listener) Accept() (net.Conn, error) {
conn, err := ln.Listener.Accept()
if err != nil {
return nil, err
}
return &serverConn{Conn: conn}, nil
}
func WrapListener(ppv int, ln net.Listener, readHeaderTimeout time.Duration) net.Listener {
if ppv <= 0 {
return ln
}
return &proxyproto.Listener{
Listener: ln,
ReadHeaderTimeout: readHeaderTimeout,
return &listener{
Listener: &proxyproto.Listener{
Listener: ln,
ReadHeaderTimeout: readHeaderTimeout,
},
}
}
+4 -4
View File
@@ -191,7 +191,7 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO
ro.Dst = cc.RemoteAddr().String()
ro.Time = time.Time{}
shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriter(br, conn), cc, node, req, &pStats, &ho)
shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriteCloser(br, conn, conn), cc, node, req, &pStats, &ho)
if err != nil || shouldClose {
return err
}
@@ -212,7 +212,7 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO
log.Trace(string(dump))
}
if shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriter(br, conn), cc, node, req, &pStats, &ho); err != nil || shouldClose {
if shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriteCloser(br, conn, conn), cc, node, req, &pStats, &ho); err != nil || shouldClose {
return err
}
}
@@ -365,7 +365,7 @@ func (h *Sniffer) serveH2(ctx context.Context, conn net.Conn, ho *HandleOptions)
return nil
}
func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriter, node *chain.Node, req *http.Request, pStats stats.Stats, ho *HandleOptions) (close bool, err error) {
func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, node *chain.Node, req *http.Request, pStats stats.Stats, ho *HandleOptions) (close bool, err error) {
close = true
log := ho.Log
@@ -586,7 +586,7 @@ func upgradeType(h http.Header) string {
return h.Get("Upgrade")
}
func (h *Sniffer) handleUpgradeResponse(ctx context.Context, rw io.ReadWriter, cc io.ReadWriter, req *http.Request, res *http.Response, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
func (h *Sniffer) handleUpgradeResponse(ctx context.Context, rw, cc io.ReadWriteCloser, req *http.Request, res *http.Response, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
reqUpType := upgradeType(req.Header)
resUpType := upgradeType(res.Header)
if !strings.EqualFold(reqUpType, resUpType) {
+4 -4
View File
@@ -198,7 +198,7 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO
ro.Dst = cc.RemoteAddr().String()
ro.Time = time.Time{}
shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriter(br, conn), cc, req, ro, &pStats, log)
shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriteCloser(br, conn, conn), cc, req, ro, &pStats, log)
if err != nil || shouldClose {
return err
}
@@ -219,7 +219,7 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO
log.Trace(string(dump))
}
if shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriter(br, conn), cc, req, ro, &pStats, log); err != nil || shouldClose {
if shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriteCloser(br, conn, conn), cc, req, ro, &pStats, log); err != nil || shouldClose {
return err
}
}
@@ -277,7 +277,7 @@ func (h *Sniffer) serveH2(ctx context.Context, conn net.Conn, ho *HandleOptions)
return nil
}
func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriter, req *http.Request, ro *xrecorder.HandlerRecorderObject, pStats stats.Stats, log logger.Logger) (close bool, err error) {
func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, req *http.Request, ro *xrecorder.HandlerRecorderObject, pStats stats.Stats, log logger.Logger) (close bool, err error) {
close = true
ro2 := &xrecorder.HandlerRecorderObject{}
@@ -433,7 +433,7 @@ func upgradeType(h http.Header) string {
return h.Get("Upgrade")
}
func (h *Sniffer) handleUpgradeResponse(ctx context.Context, rw io.ReadWriter, cc io.ReadWriter, req *http.Request, res *http.Response, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
func (h *Sniffer) handleUpgradeResponse(ctx context.Context, rw, cc io.ReadWriteCloser, req *http.Request, res *http.Response, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
reqUpType := upgradeType(req.Header)
resUpType := upgradeType(res.Header)
if !strings.EqualFold(reqUpType, resUpType) {
+2 -2
View File
@@ -46,12 +46,12 @@ func (c *shadowConn) CloseRead() error {
if sc, ok := c.Conn.(xio.CloseRead); ok {
return sc.CloseRead()
}
return nil
return xio.ErrUnsupported
}
func (c *shadowConn) CloseWrite() error {
if sc, ok := c.Conn.(xio.CloseWrite); ok {
return sc.CloseWrite()
}
return nil
return xio.ErrUnsupported
}