From 46279102381447ff29e7d639a955ead00fc556d6 Mon Sep 17 00:00:00 2001
From: Alexey Dolotov <alexey@dolotov.com>
Date: Mon, 30 Mar 2026 10:34:16 +0300
Subject: [PATCH] fix: use shared idle tracker for relay connections
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

connIdleTimeout previously set per-direction deadlines independently.
During media downloads the client→telegram direction can be idle at the
application level while telegram→client is actively streaming data.
After IdleTimeout (default 1 min) the idle direction's ReadDeadline
fires, tearing down the entire relay and breaking media transfers.

Replace the per-direction timeout with a shared atomic timestamp that
both pump goroutines update on any successful Read or Write. When a
ReadDeadline fires on the idle direction, we check the shared tracker:
if the other direction was recently active, we retry instead of closing.
The connection is only torn down when both directions are idle for the
full timeout period.

This matches the documented IdleTimeout contract: "if we have any
message which will pass to either direction, a timer is reset."

Overhead: one atomic.Int64 (8 bytes) per connection pair, one
atomic.Store (~1 ns) per Read/Write with data, zero extra goroutines.

Fixes #423
---
 mtglib/conns.go | 58 ++++++++++++++++++++++++++++++++++++++++++++-----
 mtglib/proxy.go | 12 ++++++----
 2 files changed, 61 insertions(+), 9 deletions(-)

diff --git a/mtglib/conns.go b/mtglib/conns.go
index 12f19d1..55d57ee 100644
--- a/mtglib/conns.go
+++ b/mtglib/conns.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"sync/atomic"
 	"time"
 
 	"github.com/9seconds/mtg/v2/essentials"
@@ -97,20 +98,67 @@ func newConnProxyProtocol(source, target essentials.Conn) *connProxyProtocol {
 	}
 }
 
+// idleTracker is a shared idle tracker for a pair of relay connections.
+// Both directions update the same timestamp so that activity in one direction
+// prevents the other (idle) direction from timing out.
+type idleTracker struct {
+	lastActive atomic.Int64 // unix nanos
+	timeout    time.Duration
+}
+
+func newIdleTracker(timeout time.Duration) *idleTracker {
+	t := &idleTracker{timeout: timeout}
+	t.touch()
+
+	return t
+}
+
+func (t *idleTracker) touch() {
+	t.lastActive.Store(time.Now().UnixNano())
+}
+
+func (t *idleTracker) isIdle() bool {
+	last := time.Unix(0, t.lastActive.Load())
+
+	return time.Since(last) >= t.timeout
+}
+
 type connIdleTimeout struct {
 	essentials.Conn
 
-	timeout time.Duration
+	tracker *idleTracker
 }
 
 func (c connIdleTimeout) Read(b []byte) (int, error) {
-	c.SetReadDeadline(time.Now().Add(c.timeout)) //nolint: errcheck
+	for {
+		c.SetReadDeadline(time.Now().Add(c.tracker.timeout)) //nolint: errcheck
 
-	return c.Conn.Read(b) //nolint: wrapcheck
+		n, err := c.Conn.Read(b)
+		if n > 0 {
+			c.tracker.touch()
+
+			return n, err //nolint: wrapcheck
+		}
+
+		if err != nil {
+			if netErr, ok := err.(net.Error); ok && netErr.Timeout() && !c.tracker.isIdle() { //nolint: errorlint
+				continue
+			}
+
+			return 0, err //nolint: wrapcheck
+		}
+
+		return 0, nil
+	}
 }
 
 func (c connIdleTimeout) Write(b []byte) (int, error) {
-	c.SetWriteDeadline(time.Now().Add(c.timeout)) //nolint: errcheck
+	c.SetWriteDeadline(time.Now().Add(c.tracker.timeout)) //nolint: errcheck
 
-	return c.Conn.Write(b) //nolint: wrapcheck
+	n, err := c.Conn.Write(b)
+	if n > 0 {
+		c.tracker.touch()
+	}
+
+	return n, err //nolint: wrapcheck
 }
diff --git a/mtglib/proxy.go b/mtglib/proxy.go
index 8a904ee..7acf8b7 100644
--- a/mtglib/proxy.go
+++ b/mtglib/proxy.go
@@ -102,11 +102,13 @@ func (p *Proxy) ServeConn(conn essentials.Conn) {
 		return
 	}
 
+	tracker := newIdleTracker(p.idleTimeout)
+
 	relay.Relay(
 		ctx,
 		ctx.logger.Named("relay"),
-		connIdleTimeout{Conn: ctx.telegramConn, timeout: p.idleTimeout},
-		connIdleTimeout{Conn: ctx.clientConn, timeout: p.idleTimeout},
+		connIdleTimeout{Conn: ctx.telegramConn, tracker: tracker},
+		connIdleTimeout{Conn: ctx.clientConn, tracker: tracker},
 	)
 }
 
@@ -305,11 +307,13 @@ func (p *Proxy) doDomainFronting(ctx *streamContext, conn *connRewind) {
 		stream:   p.eventStream,
 	}
 
+	tracker := newIdleTracker(p.idleTimeout)
+
 	relay.Relay(
 		ctx,
 		ctx.logger.Named("domain-fronting"),
-		connIdleTimeout{Conn: frontConn, timeout: p.idleTimeout},
-		connIdleTimeout{Conn: conn, timeout: p.idleTimeout},
+		connIdleTimeout{Conn: frontConn, tracker: tracker},
+		connIdleTimeout{Conn: conn, tracker: tracker},
 	)
 }