fix xnet.Pipe
This commit is contained in:
+21
-2
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user