mtg/wrappers/conn.go

package wrappers

import (
	"net"
	"time"

	"go.uber.org/zap"
)

type ConnPurpose uint8

func (c ConnPurpose) String() string {
	switch c {
	case ConnPurposeClient:
		return "client"
	case ConnPurposeTelegram:
		return "telegram"
	}

	return ""
}

const (
	ConnPurposeClient = iota
	ConnPurposeTelegram
)

const (
	connTimeoutRead  = 5 * time.Minute
	connTimeoutWrite = 5 * time.Minute
)

type WrapConn struct {
	purpose    ConnPurpose
	connID     string
	conn       net.Conn
	logger     *zap.SugaredLogger
	publicIPv4 net.IP
	publicIPv6 net.IP
}

func (w *WrapConn) Write(p []byte) (int, error) {
	w.conn.SetWriteDeadline(time.Now().Add(connTimeoutWrite))
	n, err := w.conn.Write(p)

	w.logger.Debugw("Write to stream", "bytes", n, "error", err)

	return n, err
}

func (w *WrapConn) Read(p []byte) (int, error) {
	w.conn.SetReadDeadline(time.Now().Add(connTimeoutRead))
	n, err := w.conn.Read(p)

	w.logger.Debugw("Read from stream", "bytes", n, "error", err)

	return n, err
}

func (w *WrapConn) Close() error {
	defer w.LogDebug("Closed connection")
	return w.conn.Close()
}

func (w *WrapConn) LocalAddr() *net.TCPAddr {
	addr := w.conn.LocalAddr().(*net.TCPAddr)
	newAddr := *addr

	if w.RemoteAddr().IP.To4() != nil {
		if w.publicIPv4 != nil {
			newAddr.IP = w.publicIPv4
		}
	} else if w.publicIPv6 != nil {
		newAddr.IP = w.publicIPv6
	}

	return &newAddr
}

func (w *WrapConn) RemoteAddr() *net.TCPAddr {
	return w.conn.RemoteAddr().(*net.TCPAddr)
}

func (w *WrapConn) LogDebug(msg string, data ...interface{}) {
	w.logger.Debugw(msg, data...)
}

func (w *WrapConn) LogInfo(msg string, data ...interface{}) {
	w.logger.Infow(msg, data...)
}

func (w *WrapConn) LogWarn(msg string, data ...interface{}) {
	w.logger.Warnw(msg, data...)
}

func (w *WrapConn) LogError(msg string, data ...interface{}) {
	w.logger.Errorw(msg, data...)
}

func NewConn(connID string, purpose ConnPurpose, conn net.Conn, publicIPv4, publicIPv6 net.IP) WrapStreamReadWriteCloser {
	logger := zap.S().With(
		"connection_id", connID,
		"local_address", conn.LocalAddr(),
		"remote_address", conn.RemoteAddr(),
	)

	return &WrapConn{
		logger:     logger,
		purpose:    purpose,
		connID:     connID,
		conn:       conn,
		publicIPv4: publicIPv4,
		publicIPv6: publicIPv6,
	}
}