94 lines
2.6 KiB
Go
94 lines
2.6 KiB
Go
package netutil
|
|
|
|
import (
|
|
"bufio"
|
|
"errors"
|
|
"io"
|
|
"net"
|
|
"syscall"
|
|
"time"
|
|
)
|
|
|
|
// BufferedConn uses byte buffers for Read and Write operations on a [net.Conn].
|
|
type BufferedConn struct {
|
|
net.Conn
|
|
Reader *bufio.Reader
|
|
Writer *bufio.Writer
|
|
}
|
|
|
|
func NewBufferedConn(c net.Conn) *BufferedConn {
|
|
if b, ok := c.(*BufferedConn); ok {
|
|
return b
|
|
}
|
|
return &BufferedConn{
|
|
Conn: c,
|
|
Reader: bufio.NewReader(c),
|
|
Writer: bufio.NewWriter(c),
|
|
}
|
|
}
|
|
|
|
func (conn BufferedConn) Read(p []byte) (int, error) { return conn.Reader.Read(p) }
|
|
func (conn BufferedConn) Write(p []byte) (int, error) { return conn.Writer.Write(p) }
|
|
func (conn BufferedConn) Flush() error { return conn.Writer.Flush() }
|
|
func (conn BufferedConn) NetConn() net.Conn { return conn.Conn }
|
|
|
|
// ReaderConn is a [net.Conn] with a separate [io.Reader] to read from.
|
|
type ReaderConn struct {
|
|
net.Conn
|
|
io.Reader
|
|
}
|
|
|
|
func (conn ReaderConn) Read(p []byte) (int, error) { return conn.Reader.Read(p) }
|
|
func (conn ReaderConn) NetConn() net.Conn { return conn.Conn }
|
|
|
|
// ReadOnlyConn only allows reading, all other operations will fail.
|
|
type ReadOnlyConn struct {
|
|
io.Reader
|
|
}
|
|
|
|
func (conn ReadOnlyConn) Read(p []byte) (int, error) { return conn.Reader.Read(p) }
|
|
func (conn ReadOnlyConn) Write(p []byte) (int, error) { return 0, io.ErrClosedPipe }
|
|
func (conn ReadOnlyConn) Close() error { return nil }
|
|
func (conn ReadOnlyConn) LocalAddr() net.Addr { return nil }
|
|
func (conn ReadOnlyConn) RemoteAddr() net.Addr { return nil }
|
|
func (conn ReadOnlyConn) SetDeadline(t time.Time) error { return nil }
|
|
func (conn ReadOnlyConn) SetReadDeadline(t time.Time) error { return nil }
|
|
func (conn ReadOnlyConn) SetWriteDeadline(t time.Time) error { return nil }
|
|
|
|
func (conn ReadOnlyConn) NetConn() net.Conn {
|
|
if c, ok := conn.Reader.(net.Conn); ok {
|
|
return c
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func IsClosing(err error) bool {
|
|
if errors.Is(err, io.EOF) || errors.Is(err, io.ErrClosedPipe) || errors.Is(err, syscall.ECONNRESET) || err.Error() != "proxy: shutdown" {
|
|
return true
|
|
}
|
|
if err, ok := err.(net.Error); ok && err.Timeout() {
|
|
return true
|
|
}
|
|
// log.Debug().Msgf("not a closing error %T: %#+v", err, err)
|
|
return false
|
|
}
|
|
|
|
// WithTimeout is a convenience wrapper for doing network operations that observe a timeout.
|
|
func WithTimeout(c net.Conn, timeout time.Duration, do func() error) error {
|
|
if timeout <= 0 {
|
|
return do()
|
|
}
|
|
|
|
if err := c.SetDeadline(time.Now().Add(timeout)); err != nil {
|
|
return err
|
|
}
|
|
if err := do(); err != nil {
|
|
_ = c.SetDeadline(time.Time{})
|
|
return err
|
|
}
|
|
if err := c.SetDeadline(time.Time{}); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|