Merge pull request #9 from 9seconds/buffers

Tune socket options

client/direct.go

@@ -3,6 +3,7 @@ package client
 import (
 	"io"
 	"net"
+	"time"
 
 	"github.com/juju/errors"
 
@@ -12,10 +13,19 @@ import (
 	"github.com/9seconds/mtg/wrappers"
 )
 
+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) {
-	socket := wrappers.NewTimeoutRWC(conn, conf.TimeoutRead, conf.TimeoutWrite)
-	frame, err := obfuscated2.ExtractFrame(socket)
+	if err := config.SetSocketOptions(conn); err != nil {
+		return nil, nil, errors.Annotate(err, "Cannot set socket options")
+	}
+
+	conn.SetReadDeadline(time.Now().Add(handshakeTimeout)) // nolint: errcheck
+	frame, err := obfuscated2.ExtractFrame(conn)
+	conn.SetReadDeadline(time.Time{}) // nolint: errcheck
 	if err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot extract frame")
 	}
@@ -26,7 +36,7 @@ func DirectInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, io
 		return nil, nil, errors.Annotate(err, "Cannot parse obfuscated frame")
 	}
 
-	socket = wrappers.NewStreamCipherRWC(socket, obfs2.Encryptor, obfs2.Decryptor)
+	socket := wrappers.NewStreamCipherRWC(conn, obfs2.Encryptor, obfs2.Decryptor)
 
 	return connOpts, socket, nil
 }

config/config.go

@@ -10,6 +10,15 @@ import (
 	"github.com/juju/errors"
 )
 
+// Buffer sizes define internal socket buffer sizes.
+const (
+	BufferWriteSize = 32 * 1024
+	BufferReadSize  = 32 * 1024
+	BufferSizeCopy  = 32 * 1024
+
+	keepAlivePeriod = 20 * time.Second
+)
+
 // Config represents common configuration of mtg.
 type Config struct {
 	Debug   bool
@@ -20,9 +29,6 @@ type Config struct {
 	PublicIPv6Port uint16
 	StatsPort      uint16
 
-	TimeoutRead  time.Duration
-	TimeoutWrite time.Duration
-
 	BindIP     net.IP
 	PublicIPv4 net.IP
 	PublicIPv6 net.IP
@@ -85,7 +91,6 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo
 	publicIPv4 net.IP, PublicIPv4Port uint16,
 	publicIPv6 net.IP, publicIPv6Port uint16,
 	statsIP net.IP, statsPort uint16,
-	timeoutRead, timeoutWrite time.Duration,
 	secret string) (*Config, error) {
 	if len(secret) != 32 {
 		return nil, errors.New("Telegram demands secret of length 32")
@@ -136,10 +141,31 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo
 		PublicIPv6Port: publicIPv6Port,
 		StatsIP:        statsIP,
 		StatsPort:      statsPort,
-		TimeoutRead:    timeoutRead,
-		TimeoutWrite:   timeoutWrite,
 		Secret:         secretBytes,
 	}
 
 	return conf, nil
 }
+
+// SetSocketOptions makes socket keepalive, sets buffer sizes
+func SetSocketOptions(conn net.Conn) error {
+	socket := conn.(*net.TCPConn)
+
+	if err := socket.SetReadBuffer(BufferReadSize); err != nil {
+		return errors.Annotate(err, "Cannot set read buffer size")
+	}
+	if err := socket.SetWriteBuffer(BufferWriteSize); err != nil {
+		return errors.Annotate(err, "Cannot set write buffer size")
+	}
+	if err := socket.SetKeepAlive(true); err != nil {
+		return errors.Annotate(err, "Cannot make socket keepalive")
+	}
+	if err := socket.SetKeepAlivePeriod(keepAlivePeriod); err != nil {
+		return errors.Annotate(err, "Cannot set keepalive period")
+	}
+	if err := socket.SetNoDelay(true); err != nil {
+		return errors.Annotate(err, "Cannot activate nodelay for the socket")
+	}
+
+	return nil
+}

main.go

@@ -69,17 +69,6 @@ var (
 			Default("3129").
 			Uint16()
 
-	readTimeout = app.Flag("read-timeout", "Socket read timeout.").
-			Short('r').
-			Envar("MTG_READ_TIMEOUT").
-			Default("30s").
-			Duration()
-	writeTimeout = app.Flag("write-timeout", "Socket write timeout.").
-			Short('w').
-			Envar("MTG_WRITE_TIMEOUT").
-			Default("30s").
-			Duration()
-
 	secret = app.Arg("secret", "Secret of this proxy.").Required().String()
 )
 
@@ -102,7 +91,6 @@ func main() {
 		*publicIPv4, *publicIPv4Port,
 		*publicIPv6, *publicIPv6Port,
 		*statsIP, *statsPort,
-		*readTimeout, *writeTimeout,
 		*secret,
 	)
 	if err != nil {

proxy/copy_pool.go

@@ -1,15 +1,17 @@
 package proxy
 
-import "sync"
+import (
+	"sync"
 
-const copyBufferSize = 30 * 1024
+	"github.com/9seconds/mtg/config"
+)
 
 var copyPool sync.Pool
 
 func init() {
 	copyPool = sync.Pool{
 		New: func() interface{} {
-			data := make([]byte, copyBufferSize)
+			data := make([]byte, config.BufferSizeCopy)
 			return &data
 		},
 	}

telegram/dialer.go

@@ -8,29 +8,23 @@ import (
 	"github.com/juju/errors"
 
 	"github.com/9seconds/mtg/config"
-	"github.com/9seconds/mtg/wrappers"
 )
 
-const telegramKeepAlive = 30 * time.Second
+const (
+	telegramDialTimeout = 10 * time.Second
+)
 
 type tgDialer struct {
 	net.Dialer
-
-	conf *config.Config
 }
 
 func (t *tgDialer) dial(addr string) (net.Conn, error) {
-	connRaw, err := t.Dialer.Dial("tcp", addr)
+	conn, err := t.Dialer.Dial("tcp", addr)
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot connect to Telegram")
 	}
-	conn := connRaw.(*net.TCPConn)
-
-	if err = conn.SetKeepAlive(true); err != nil {
-		return nil, errors.Annotate(err, "Cannot establish keepalive connection")
-	}
-	if err = conn.SetKeepAlivePeriod(telegramKeepAlive); err != nil {
-		return nil, errors.Annotate(err, "Cannot set keepalive timeout")
+	if err = config.SetSocketOptions(conn); err != nil {
+		return nil, errors.Annotate(err, "Cannot set socket options")
 	}
 
 	return conn, nil
@@ -42,12 +36,5 @@ func (t *tgDialer) dialRWC(addr string) (io.ReadWriteCloser, error) {
 		return nil, err
 	}
 
-	return wrappers.NewTimeoutRWC(conn, t.conf.TimeoutRead, t.conf.TimeoutWrite), nil
-}
-
-func newDialer(conf *config.Config) tgDialer {
-	return tgDialer{
-		Dialer: net.Dialer{Timeout: conf.TimeoutRead},
-		conf:   conf,
-	}
+	return conn, nil
 }

telegram/direct.go

@@ -2,6 +2,7 @@ package telegram
 
 import (
 	"io"
+	"net"
 
 	"github.com/juju/errors"
 
@@ -58,7 +59,7 @@ func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts, conn io.ReadWrit
 // to Telegram bypassing middleproxies.
 func NewDirectTelegram(conf *config.Config) Telegram {
 	return &directTelegram{baseTelegram{
-		dialer:      newDialer(conf),
+		dialer:      tgDialer{net.Dialer{Timeout: telegramDialTimeout}},
 		v4Addresses: directV4Addresses,
 		v6Addresses: directV6Addresses,
 	}}

wrappers/timeoutrwc.go

@@ -1,42 +0,0 @@
-package wrappers
-
-import (
-	"io"
-	"net"
-	"time"
-)
-
-// TimeoutReadWriteCloser sets timeouts for read/write into underlying
-// network connection.
-type TimeoutReadWriteCloser struct {
-	conn         net.Conn
-	readTimeout  time.Duration
-	writeTimeout time.Duration
-}
-
-// Read reads from connection
-func (t *TimeoutReadWriteCloser) Read(p []byte) (int, error) {
-	t.conn.SetReadDeadline(time.Now().Add(t.readTimeout)) // nolint: errcheck, gas
-	return t.conn.Read(p)
-}
-
-// Write writes into connection.
-func (t *TimeoutReadWriteCloser) Write(p []byte) (int, error) {
-	t.conn.SetWriteDeadline(time.Now().Add(t.writeTimeout)) // nolint: errcheck, gas
-	return t.conn.Write(p)
-}
-
-// Close closes underlying connection.
-func (t *TimeoutReadWriteCloser) Close() error {
-	return t.conn.Close()
-}
-
-// NewTimeoutRWC returns wrapper over net.Conn which sets deadlines for
-// every wrapped Read/Write.
-func NewTimeoutRWC(conn net.Conn, readTimeout, writeTimeout time.Duration) io.ReadWriteCloser {
-	return &TimeoutReadWriteCloser{
-		conn:         conn,
-		readTimeout:  readTimeout,
-		writeTimeout: writeTimeout,
-	}
-}