Start to rewrite

wrappers/conn.go

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