From 0e2898f6be9ea71980ba8d13aaf28a9c1c94c9e2 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Mon, 2 Jul 2018 08:09:22 +0300
Subject: [PATCH] Return back timeoutrwc

---
 client/direct.go       |  7 +++----
 config/config.go       |  3 +++
 telegram/dialer.go     |  7 +++----
 wrappers/timeoutrwc.go | 31 +++++++++++++++++++++++++++++++
 4 files changed, 40 insertions(+), 8 deletions(-)
 create mode 100644 wrappers/timeoutrwc.go

diff --git a/client/direct.go b/client/direct.go
index ed40dbe..23e9613 100644
--- a/client/direct.go
+++ b/client/direct.go
@@ -13,9 +13,7 @@ import (
 	"github.com/9seconds/mtg/wrappers"
 )
 
-const (
-	handshakeTimeout = 10 * time.Second
-)
+const handshakeTimeout = 10 * time.Second
 
 // DirectInit initializes client to access Telegram bypassing middleproxies.
 func DirectInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, io.ReadWriteCloser, error) {
@@ -36,7 +34,8 @@ func DirectInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, io
 	}
 	connOpts.ConnectionProto = mtproto.ConnectionProtocolAny
 
-	socket := wrappers.NewStreamCipherRWC(conn, obfs2.Encryptor, obfs2.Decryptor)
+	socket := wrappers.NewTimeoutRWC(conn)
+	socket = wrappers.NewStreamCipherRWC(socket, obfs2.Encryptor, obfs2.Decryptor)
 
 	return connOpts, socket, nil
 }
diff --git a/config/config.go b/config/config.go
index 54f4cba..9935fa8 100644
--- a/config/config.go
+++ b/config/config.go
@@ -16,6 +16,9 @@ const (
 	BufferReadSize  = 32 * 1024
 	BufferSizeCopy  = 32 * 1024
 
+	TimeoutRead  = time.Minute
+	TimeoutWrite = time.Minute
+
 	keepAlivePeriod = 20 * time.Second
 )
 
diff --git a/telegram/dialer.go b/telegram/dialer.go
index ef51f50..f6b2cf7 100644
--- a/telegram/dialer.go
+++ b/telegram/dialer.go
@@ -8,11 +8,10 @@ import (
 	"github.com/juju/errors"
 
 	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/wrappers"
 )
 
-const (
-	telegramDialTimeout = 10 * time.Second
-)
+const telegramDialTimeout = 10 * time.Second
 
 type tgDialer struct {
 	net.Dialer
@@ -36,5 +35,5 @@ func (t *tgDialer) dialRWC(addr string) (io.ReadWriteCloser, error) {
 		return nil, err
 	}
 
-	return conn, nil
+	return wrappers.NewTimeoutRWC(conn), nil
 }
diff --git a/wrappers/timeoutrwc.go b/wrappers/timeoutrwc.go
new file mode 100644
index 0000000..e4ea9b0
--- /dev/null
+++ b/wrappers/timeoutrwc.go
@@ -0,0 +1,31 @@
+package wrappers
+
+import (
+	"io"
+	"net"
+	"time"
+
+	"github.com/9seconds/mtg/config"
+)
+
+type TimeoutReadWriteCloser struct {
+	conn net.Conn
+}
+
+func (t *TimeoutReadWriteCloser) Read(p []byte) (int, error) {
+	t.conn.SetReadDeadline(time.Now().Add(config.TimeoutRead))
+	return t.conn.Read(p)
+}
+
+func (t *TimeoutReadWriteCloser) Write(p []byte) (int, error) {
+	t.conn.SetWriteDeadline(time.Now().Add(config.TimeoutWrite))
+	return t.conn.Write(p)
+}
+
+func (t *TimeoutReadWriteCloser) Close() error {
+	return t.conn.Close()
+}
+
+func NewTimeoutRWC(conn net.Conn) io.ReadWriteCloser {
+	return &TimeoutReadWriteCloser{conn}
+}