github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/utils/relay/relay.go (about)

     1  package relay
     2  
     3  import (
     4  	"errors"
     5  	"io"
     6  	"net"
     7  
     8  	"github.com/Asutorufa/yuhaiin/pkg/log"
     9  	"github.com/Asutorufa/yuhaiin/pkg/utils/pool"
    10  )
    11  
    12  var ignoreError = []error{
    13  	io.EOF,
    14  	// os.ErrDeadlineExceeded,
    15  	net.ErrClosed,
    16  }
    17  
    18  func logE(msg string, err error) {
    19  	if err == nil {
    20  		return
    21  	}
    22  
    23  	for _, e := range ignoreError {
    24  		if errors.Is(err, e) {
    25  			return
    26  		}
    27  	}
    28  
    29  	log.Error(msg, "err", err)
    30  }
    31  
    32  func AppendIgnoreError(err error) {
    33  	ignoreError = append(ignoreError, err)
    34  }
    35  
    36  // Relay pipe
    37  func Relay(rw1, rw2 io.ReadWriteCloser) {
    38  	wait := make(chan struct{})
    39  	go func() {
    40  		defer close(wait)
    41  		_, err := Copy(rw2, rw1)
    42  		logE("relay rw1 -> rw2", err)
    43  		closeWrite(rw2) // make another Copy exit
    44  		closeRead(rw1)
    45  	}()
    46  
    47  	_, err := Copy(rw1, rw2)
    48  	logE("relay rw2 -> rw1", err)
    49  	closeWrite(rw1)
    50  	closeRead(rw2)
    51  
    52  	<-wait
    53  }
    54  
    55  func closeRead(rw io.ReadWriteCloser) {
    56  	if cr, ok := rw.(interface{ CloseRead() error }); ok {
    57  		_ = cr.CloseRead()
    58  	}
    59  }
    60  
    61  func closeWrite(rw io.ReadWriteCloser) {
    62  	if r, ok := rw.(interface{ CloseWrite() error }); ok {
    63  		_ = r.CloseWrite()
    64  		return
    65  	}
    66  
    67  	_ = rw.Close()
    68  }
    69  
    70  func Copy(dst io.Writer, src io.Reader) (n int64, err error) {
    71  	buf := pool.GetBytes(pool.DefaultSize)
    72  	defer pool.PutBytes(buf)
    73  	// to avoid using (*net.TCPConn).ReadFrom that will make new none-zero buf
    74  	return io.CopyBuffer(WriteOnlyWriter{dst}, ReadOnlyReader{src}, buf)
    75  }
    76  
    77  func CopyN(dst io.Writer, src io.Reader, n int64) (written int64, err error) {
    78  	if n <= 0 {
    79  		return 0, nil
    80  	}
    81  
    82  	written, err = Copy(dst, io.LimitReader(src, n))
    83  	if written == n {
    84  		return n, nil
    85  	}
    86  	if written < n && err == nil {
    87  		// src stopped early; must have been EOF.
    88  		err = io.EOF
    89  	}
    90  	return
    91  }
    92  
    93  type ReadOnlyReader struct{ io.Reader }
    94  type WriteOnlyWriter struct{ io.Writer }