Merge pull request #31 from 9seconds/propagate-buffer-sizes

Propagate buffer sizes

client/direct.go

@@ -12,11 +12,7 @@ import (
 	"github.com/9seconds/mtg/wrappers"
 )
 
-const (
+const handshakeTimeout = 10 * time.Second
-	handshakeTimeout = 10 * time.Second
-	readBufferSize   = 64 * 1024
-	writeBufferSize  = 64 * 1024
-)
 
 // DirectInit initializes client connection for proxy which connects to
 // Telegram directly.
@@ -25,10 +21,10 @@ func DirectInit(socket net.Conn, connID string, conf *config.Config) (wrappers.W
 	if err := tcpSocket.SetNoDelay(false); err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot disable NO_DELAY to client socket")
 	}
-	if err := tcpSocket.SetReadBuffer(readBufferSize); err != nil {
+	if err := tcpSocket.SetReadBuffer(conf.ReadBufferSize); err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot set read buffer size of client socket")
 	}
-	if err := tcpSocket.SetWriteBuffer(writeBufferSize); err != nil {
+	if err := tcpSocket.SetWriteBuffer(conf.WriteBufferSize); err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot set write buffer size of client socket")
 	}
 

config/config.go

@@ -1,28 +1,25 @@
 package config
 
 import (
+	"bytes"
 	"encoding/hex"
 	"fmt"
 	"net"
 	"strconv"
-	"strings"
 
 	"github.com/juju/errors"
 	statsd "gopkg.in/alexcesaro/statsd.v2"
 )
 
-// Buffer sizes define internal socket buffer sizes.
-const (
-	BufferWriteSize = 32 * 1024
-	BufferReadSize  = 32 * 1024
-)
-
 // Config represents common configuration of mtg.
 type Config struct {
 	Debug      bool
 	Verbose    bool
 	SecureMode bool
 
+	ReadBufferSize  int
+	WriteBufferSize int
+
 	BindPort       uint16
 	PublicIPv4Port uint16
 	PublicIPv6Port uint16
@@ -114,30 +111,21 @@ func getAddr(host fmt.Stringer, port uint16) string {
 // fetches data from external sources. Parameters passed to this
 // function, should come from command line arguments.
 func NewConfig(debug, verbose bool, // nolint: gocyclo
+	writeBufferSize, readBufferSize uint32,
 	bindIP, publicIPv4, publicIPv6, statsIP net.IP,
 	bindPort, publicIPv4Port, publicIPv6Port, statsPort, statsdPort uint16,
-	secret, adtag, statsdIP, statsdNetwork, statsdPrefix, statsdTagsFormat string,
+	statsdIP, statsdNetwork, statsdPrefix, statsdTagsFormat string,
-	statsdTags map[string]string) (*Config, error) {
+	statsdTags map[string]string,
+	secret, adtag []byte) (*Config, error) {
 	secureMode := false
-	if strings.HasPrefix(secret, "dd") && len(secret) == 34 {
+	if bytes.HasPrefix(secret, []byte{0xdd}) && len(secret) == 17 {
 		secureMode = true
-		secret = strings.TrimPrefix(secret, "dd")
+		secret = bytes.TrimPrefix(secret, []byte{0xdd})
-	} else if len(secret) != 32 {
+	} else if len(secret) != 16 {
 		return nil, errors.New("Telegram demands secret of length 32")
 	}
-	secretBytes, err := hex.DecodeString(secret)
-	if err != nil {
-		return nil, errors.Annotate(err, "Cannot create config")
-	}
-
-	var adTagBytes []byte
-	if len(adtag) != 0 {
-		adTagBytes, err = hex.DecodeString(adtag)
-		if err != nil {
-			return nil, errors.Annotate(err, "Cannot create config")
-		}
-	}
 
+	var err error
 	if publicIPv4 == nil {
 		publicIPv4, err = getGlobalIPv4()
 		if err != nil {
@@ -167,19 +155,21 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo
 	}
 
 	conf := &Config{
-		Debug:          debug,
+		Debug:           debug,
-		Verbose:        verbose,
+		Verbose:         verbose,
-		BindIP:         bindIP,
+		BindIP:          bindIP,
-		BindPort:       bindPort,
+		BindPort:        bindPort,
-		PublicIPv4:     publicIPv4,
+		PublicIPv4:      publicIPv4,
-		PublicIPv4Port: publicIPv4Port,
+		PublicIPv4Port:  publicIPv4Port,
-		PublicIPv6:     publicIPv6,
+		PublicIPv6:      publicIPv6,
-		PublicIPv6Port: publicIPv6Port,
+		PublicIPv6Port:  publicIPv6Port,
-		StatsIP:        statsIP,
+		StatsIP:         statsIP,
-		StatsPort:      statsPort,
+		StatsPort:       statsPort,
-		Secret:         secretBytes,
+		Secret:          secret,
-		AdTag:          adTagBytes,
+		AdTag:           adtag,
-		SecureMode:     secureMode,
+		SecureMode:      secureMode,
+		ReadBufferSize:  int(readBufferSize),
+		WriteBufferSize: int(writeBufferSize),
 	}
 
 	if statsdIP != "" {
@@ -187,7 +177,10 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo
 		conf.StatsD.Prefix = statsdPrefix
 		conf.StatsD.Tags = statsdTags
 
-		var addr net.Addr
+		var (
+			addr net.Addr
+			err  error
+		)
 		hostPort := net.JoinHostPort(statsdIP, strconv.Itoa(int(statsdPort)))
 		switch statsdNetwork {
 		case "tcp":

main.go

@@ -49,6 +49,19 @@ var (
 		Default("3128").
 		Uint16()
 
+	writeBufferSize = app.Flag("write-buffer",
+		"Write buffer size in bytes. You can think about it as a buffer from client to Telegram.").
+		Short('w').
+		Envar("MTG_BUFFER_WRITE").
+		Default("65536").
+		Uint32()
+	readBufferSize = app.Flag("read-buffer",
+		"Read buffer size in bytes. You can think about it as a buffer from Telegram to client.").
+		Short('r').
+		Envar("MTG_BUFFER_READ").
+		Default("131072").
+		Uint32()
+
 	publicIPv4 = app.Flag("public-ipv4",
 		"Which IPv4 address is public.").
 		Short('4').
@@ -110,17 +123,17 @@ var (
 		Envar("MTG_STATSD_TAGS").
 		StringMap()
 
-	secret = app.Arg("secret", "Secret of this proxy.").Required().String()
+	secret = app.Arg("secret", "Secret of this proxy.").Required().HexBytes()
-	adtag  = app.Arg("adtag", "ADTag of the proxy.").String()
+	adtag  = app.Arg("adtag", "ADTag of the proxy.").HexBytes()
 )
 
 func init() {
 	rand.Seed(time.Now().UTC().UnixNano())
 	app.Version(version)
-
+	app.HelpFlag.Short('h')
 }
 
-func main() {
+func main() { // nolint: gocyclo
 	kingpin.MustParse(app.Parse(os.Args[1:]))
 
 	err := setRLimit()
@@ -129,10 +142,12 @@ func main() {
 	}
 
 	conf, err := config.NewConfig(*debug, *verbose,
+		*writeBufferSize, *readBufferSize,
 		*bindIP, *publicIPv4, *publicIPv6, *statsIP,
 		*bindPort, *publicIPv4Port, *publicIPv6Port, *statsPort, *statsdPort,
-		*secret, *adtag, *statsdIP, *statsdNetwork, *statsdPrefix, *statsdTagsFormat,
+		*statsdIP, *statsdNetwork, *statsdPrefix, *statsdTagsFormat,
 		*statsdTags,
+		*secret, *adtag,
 	)
 	if err != nil {
 		usage(err.Error())

proxy/proxy.go

@@ -83,8 +83,8 @@ func (p *Proxy) accept(conn net.Conn) {
 	} else {
 		clientStream := clientConn.(wrappers.StreamReadWriteCloser)
 		serverStream := serverConn.(wrappers.StreamReadWriteCloser)
-		go p.directPipe(clientStream, serverStream, wait)
+		go p.directPipe(clientStream, serverStream, wait, p.conf.ReadBufferSize)
-		go p.directPipe(serverStream, clientStream, wait)
+		go p.directPipe(serverStream, clientStream, wait, p.conf.WriteBufferSize)
 	}
 
 	wait.Wait()
@@ -130,14 +130,16 @@ func (p *Proxy) middlePipe(src wrappers.PacketReadCloser, dst io.WriteCloser,
 	}
 }
 
-func (p *Proxy) directPipe(src wrappers.StreamReadCloser, dst io.WriteCloser, wait *sync.WaitGroup) {
+func (p *Proxy) directPipe(src wrappers.StreamReadCloser, dst io.WriteCloser,
+	wait *sync.WaitGroup, bufferSize int) {
 	defer func() {
 		src.Close() // nolint: errcheck
 		dst.Close() // nolint: errcheck
 		wait.Done()
 	}()
 
-	if _, err := io.Copy(dst, src); err != nil {
+	buffer := make([]byte, bufferSize)
+	if _, err := io.CopyBuffer(dst, src, buffer); err != nil {
 		src.Logger().Warnw("Cannot pump sockets", "error", err)
 	}
 }

telegram/dialer.go

@@ -10,11 +10,7 @@ import (
 	"github.com/9seconds/mtg/wrappers"
 )
 
-const (
+const telegramDialTimeout = 10 * time.Second
-	telegramDialTimeout = 10 * time.Second
-	readBufferSize      = 64 * 1024
-	writeBufferSize     = 64 * 1024
-)
 
 type tgDialer struct {
 	net.Dialer
@@ -32,10 +28,10 @@ func (t *tgDialer) dial(addr string) (net.Conn, error) {
 	if err = tcpSocket.SetNoDelay(true); err != nil {
 		return nil, errors.Annotate(err, "Cannot set NO_DELAY to Telegram")
 	}
-	if err = tcpSocket.SetReadBuffer(readBufferSize); err != nil {
+	if err = tcpSocket.SetReadBuffer(t.conf.WriteBufferSize); err != nil {
 		return nil, errors.Annotate(err, "Cannot set read buffer size on telegram socket")
 	}
-	if err = tcpSocket.SetWriteBuffer(writeBufferSize); err != nil {
+	if err = tcpSocket.SetWriteBuffer(t.conf.ReadBufferSize); err != nil {
 		return nil, errors.Annotate(err, "Cannot set write buffer size on telegram socket")
 	}