Checkpoint

internal/netutil/conn.go

@@ -0,0 +1,93 @@
+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
+}