Update version of golangci-lint

Use contexts for Conn wrapper

.dockerignore

@@ -0,0 +1,3 @@
+version.go
+vendor/
+tags

.golangci.toml

@@ -0,0 +1,12 @@
+[run]
+concurrency = 4
+deadline = "2m"
+tests = true
+skip-dirs = ["vendor"]
+skip-files = ["version.go"]
+
+[output]
+format = "colored-line-number"
+
+[linters]
+enable-all = true

.gometalinter.json

@@ -1,4 +0,0 @@
-{
-    "aggregate": true,
-    "vendor": true
-}

.travis.yml

@@ -9,10 +9,18 @@ go:
   - "1.10.x"
   - master
 
+before_script: make prepare
+
 script:
+  - make all
   - make lint
+  - make critic
   - make test
 
 cache:
   directories:
     - vendor
+
+matrix:
+  allow_failures:
+    - go: master

Dockerfile

@@ -13,11 +13,17 @@ RUN set -x \
     upx \
   && update-ca-certificates
 
-ADD . /go/src/github.com/9seconds/mtg
+COPY Gopkg.toml Gopkg.lock Makefile /go/src/github.com/9seconds/mtg/
+
+RUN set -x && \
+  cd /go/src/github.com/9seconds/mtg && \
+  make -j 4 prepare && \
+  make vendor
+
+COPY . /go/src/github.com/9seconds/mtg
 
 RUN set -x \
   && cd /go/src/github.com/9seconds/mtg \
-  && make clean \
   && make -j 4 static \
   && upx --ultra-brute -qq ./mtg
 

Gopkg.lock

@@ -3,74 +3,97 @@
 
 [[projects]]
   branch = "master"
+  digest = "1:315c5f2f60c76d89b871c73f9bd5fe689cad96597afd50fb9992228ef80bdd34"
   name = "github.com/alecthomas/template"
   packages = [
     ".",
-    "parse"
+    "parse",
   ]
+  pruneopts = "UT"
   revision = "a0175ee3bccc567396460bf5acd36800cb10c49c"
 
 [[projects]]
   branch = "master"
+  digest = "1:c198fdc381e898e8fb62b8eb62758195091c313ad18e52a3067366e1dda2fb3c"
   name = "github.com/alecthomas/units"
   packages = ["."]
+  pruneopts = "UT"
   revision = "2efee857e7cfd4f3d0138cc3cbb1b4966962b93a"
 
 [[projects]]
+  digest = "1:a8d622a8049a4aa420e1c509873bb85d4c45c5107f420d922f919bfcb8d08694"
   name = "github.com/beevik/ntp"
   packages = ["."]
+  pruneopts = "UT"
   revision = "62c80a04de2086884d8296004b6d74ee1846c582"
   version = "v0.2.0"
 
 [[projects]]
+  digest = "1:a2c1d0e43bd3baaa071d1b9ed72c27d78169b2b269f71c105ac4ba34b1be4a39"
   name = "github.com/davecgh/go-spew"
   packages = ["spew"]
+  pruneopts = "UT"
   revision = "346938d642f2ec3594ed81d874461961cd0faa76"
   version = "v1.1.0"
 
 [[projects]]
   branch = "master"
+  digest = "1:6f9339c912bbdda81302633ad7e99a28dfa5a639c864061f1929510a9a64aa74"
   name = "github.com/dustin/go-humanize"
   packages = ["."]
+  pruneopts = "UT"
   revision = "9f541cc9db5d55bce703bd99987c9d5cb8eea45e"
 
 [[projects]]
   branch = "master"
+  digest = "1:53bd4347b151fcbedcdda527f7ebf4924f0e21d672131812f857175d8c7a1051"
   name = "github.com/juju/errors"
   packages = ["."]
-  revision = "c7d06af17c68cd34c835053720b21f6549d9b0ee"
+  pruneopts = "UT"
+  revision = "812b06ada1776ad4dd95d575e18ffffe3a9ac34a"
 
 [[projects]]
+  digest = "1:0028cb19b2e4c3112225cd871870f2d9cf49b9b4276531f03438a88e94be86fe"
   name = "github.com/pmezard/go-difflib"
   packages = ["difflib"]
+  pruneopts = "UT"
   revision = "792786c7400a136282c1664665ae0a8db921c6c2"
   version = "v1.0.0"
 
 [[projects]]
+  digest = "1:274f67cb6fed9588ea2521ecdac05a6d62a8c51c074c1fccc6a49a40ba80e925"
   name = "github.com/satori/go.uuid"
   packages = ["."]
+  pruneopts = "UT"
   revision = "f58768cc1a7a7e77a3bd49e98cdd21419399b6a3"
   version = "v1.2.0"
 
 [[projects]]
+  digest = "1:18752d0b95816a1b777505a97f71c7467a8445b8ffb55631a7bf779f6ba4fa83"
   name = "github.com/stretchr/testify"
   packages = ["assert"]
+  pruneopts = "UT"
   revision = "f35b8ab0b5a2cef36673838d662e249dd9c94686"
   version = "v1.2.2"
 
 [[projects]]
+  digest = "1:3c1a69cdae3501bf75e76d0d86dc6f2b0a7421bc205c0cb7b96b19eed464a34d"
   name = "go.uber.org/atomic"
   packages = ["."]
+  pruneopts = "UT"
   revision = "1ea20fb1cbb1cc08cbd0d913a96dead89aa18289"
   version = "v1.3.2"
 
 [[projects]]
+  digest = "1:60bf2a5e347af463c42ed31a493d817f8a72f102543060ed992754e689805d1a"
   name = "go.uber.org/multierr"
   packages = ["."]
+  pruneopts = "UT"
   revision = "3c4937480c32f4c13a875a1829af76c98ca3d40a"
   version = "v1.1.0"
 
 [[projects]]
+  digest = "1:d9a420eae5f76973feeb733fbf58f6a89173255b63b27a7ae4b2124a73dc6c5b"
   name = "go.uber.org/zap"
   packages = [
     ".",
@@ -78,37 +101,54 @@
     "internal/bufferpool",
     "internal/color",
     "internal/exit",
-    "zapcore"
+    "zapcore",
   ]
-  revision = "eeedf312bc6c57391d84767a4cd413f02a917974"
+  pruneopts = "UT"
-  version = "v1.8.0"
+  revision = "4d45f9617f7d90f7a663ff21c7a4321dbe78098b"
+  version = "v1.9.0"
 
 [[projects]]
   branch = "master"
+  digest = "1:becb2131aece71c64ebca5ddfd38cf631784fbb3a678d73ead3b42107c9f41ce"
   name = "golang.org/x/net"
   packages = [
     "bpf",
     "internal/iana",
     "internal/socket",
-    "ipv4"
+    "ipv4",
   ]
-  revision = "d0887baf81f4598189d4e12a37c6da86f0bba4d0"
+  pruneopts = "UT"
+  revision = "3673e40ba22529d22c3fd7c93e97b0ce50fa7bdd"
 
 [[projects]]
+  digest = "1:c06d9e11d955af78ac3bbb26bd02e01d2f61f689e1a3bce2ef6fb683ef8a7f2d"
   name = "gopkg.in/alecthomas/kingpin.v2"
   packages = ["."]
+  pruneopts = "UT"
   revision = "947dcec5ba9c011838740e680966fd7087a71d0d"
   version = "v2.2.6"
 
 [[projects]]
+  digest = "1:38b469493eb173db9c03321d64adcad4c7991ea0a19b5edc5bdc094f0e8c7384"
   name = "gopkg.in/alexcesaro/statsd.v2"
   packages = ["."]
+  pruneopts = "UT"
   revision = "7fea3f0d2fab1ad973e641e51dba45443a311a90"
   version = "v2.0.0"
 
 [solve-meta]
   analyzer-name = "dep"
   analyzer-version = 1
-  inputs-digest = "7fad0f62feb7737b064d85cc4333a1a3e9298faec2afd864b4404f515fc7f17c"
+  input-imports = [
+    "github.com/beevik/ntp",
+    "github.com/dustin/go-humanize",
+    "github.com/juju/errors",
+    "github.com/satori/go.uuid",
+    "github.com/stretchr/testify/assert",
+    "go.uber.org/zap",
+    "go.uber.org/zap/zapcore",
+    "gopkg.in/alecthomas/kingpin.v2",
+    "gopkg.in/alexcesaro/statsd.v2",
+  ]
   solver-name = "gps-cdcl"
   solver-version = 1

Makefile

@@ -1,12 +1,13 @@
 ROOT_DIR     := $(shell dirname $(realpath $(lastword $(MAKEFILE_LIST))))
 IMAGE_NAME   := mtg
 APP_NAME     := $(IMAGE_NAME)
-GOMETALINTER := gometalinter
 
 VENDOR_FILES := $(shell find "$(ROOT_DIR)/vendor" 2>/dev/null || echo -n "vendor")
 CC_BINARIES  := $(shell bash -c "echo -n $(APP_NAME)-{linux,freebsd,openbsd}-{386,amd64} $(APP_NAME)-linux-{arm,arm64}")
 APP_DEPS     := version.go $(VENDOR_FILES)
 
+GOLANGCI_LINT_VERSION := v1.9.2
+
 COMMON_BUILD_FLAGS := -ldflags="-s -w"
 
 # -----------------------------------------------------------------------------
@@ -31,8 +32,8 @@ ccbuilds:
 version.go:
 	@go generate main.go
 
-vendor: Gopkg.lock Gopkg.toml install-cli
+vendor: Gopkg.lock Gopkg.toml
-	@dep ensure
+	@dep ensure --vendor-only
 
 # -----------------------------------------------------------------------------
 
@@ -50,12 +51,16 @@ crosscompile-dir:
 	@rm -rf "$(CC_DIR)" && mkdir -p "$(CC_DIR)"
 
 .PHONY: test
-test: vendor install-cli version.go
+test: vendor version.go
 	@go test -v ./...
 
 .PHONY: lint
-lint: vendor install-cli version.go
+lint: version.go
-	@$(GOMETALINTER) --deadline=2m ./...
+	@golangci-lint run
+
+.PHONY: critic
+critic: version.go
+	@gocritic check-project "$(ROOT_DIR)"
 
 .PHONY: clean
 clean:
@@ -67,14 +72,18 @@ clean:
 docker:
 	@docker build --pull -t "$(IMAGE_NAME)" "$(ROOT_DIR)"
 
-.PHONY: install-cli
+.PHONY: prepare
-install-cli: install-dep install-lint
+prepare: install-dep install-lint install-critic
 
 .PHONY: install-dep
 install-dep:
-	@go get github.com/golang/dep/cmd/dep
+	@go get -u github.com/golang/dep/cmd/dep
 
 .PHONY: install-lint
 install-lint:
-	@go get github.com/alecthomas/gometalinter && \
+	@curl -sfL https://install.goreleaser.com/github.com/golangci/golangci-lint.sh \
-		$(GOMETALINTER) --install >/dev/null
+		| bash -s -- -b $(GOPATH)/bin $(GOLANGCI_LINT_VERSION)
+
+.PHONY: install-critic
+install-critic:
+	@go get -u github.com/go-critic/go-critic/...

README.md

@@ -5,7 +5,6 @@ Bullshit-free MTPROTO proxy for Telegram
 [![Build Status](https://travis-ci.org/9seconds/mtg.svg?branch=master)](https://travis-ci.org/9seconds/mtg)
 [![Docker Build Status](https://img.shields.io/docker/build/nineseconds/mtg.svg)](https://hub.docker.com/r/nineseconds/mtg/)
 
-
 # Rationale
 
 There are several available proxies for Telegram MTPROTO available. Here
@@ -60,19 +59,19 @@ There are 2 main branches:
 # How to build
 
 ```console
-$ make
+make
 ```
 
 If you want to build for another platform:
 
 ```console
-$ make crosscompile
+make crosscompile
 ```
 
 If you want to build Docker image (called `mtg`):
 
 ```console
-$ make docker
+make docker
 ```
 
 # Docker image
@@ -84,15 +83,15 @@ Docker follows the same policy as the source code organization:
 - tags are for tagged releases
 
 ```console
-$ docker pull nineseconds/mtg:latest
+docker pull nineseconds/mtg:latest
 ```
 
 ```console
-$ docker pull nineseconds/mtg:stable
+docker pull nineseconds/mtg:stable
 ```
 
 ```console
-$ docker pull nineseconds/mtg:0.10
+docker pull nineseconds/mtg:0.10
 ```
 
 # Configuration
@@ -102,13 +101,13 @@ Basically, to run this tool you need to configure as less as possible.
 First, you need to generate a secret:
 
 ```console
-$ openssl rand -hex 16
+openssl rand -hex 16
 ```
 
 or
 
 ```console
-$ head -c 512 /dev/urandom | md5sum | cut -f 1 -d ' '
+head -c 512 /dev/urandom | md5sum | cut -f 1 -d ' '
 ```
 
 ## Secure mode
@@ -122,28 +121,62 @@ suggest to go with this mode.
 Oneliners to generate such secrets:
 
 ```console
-$ echo dd$(openssl rand -hex 16)
+echo dd$(openssl rand -hex 16)
 ```
 
 or
 
 ```console
-$ echo dd$(head -c 512 /dev/urandom | md5sum | cut -f 1 -d ' ')
+echo dd$(head -c 512 /dev/urandom | md5sum | cut -f 1 -d ' ')
 ```
 
+## Environment variables
+
+It is possible to configure this tool using environment variables. You
+can configure any flag but not secret or adtag. Here is the list of
+supported environment variables:
+
+| Environment variable     | Corresponding flags    | Default value                     | Description                                                                                                                                                                                                                                                                |
+|--------------------------|------------------------|-----------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `MTG_DEBUG`              | `-d`, `--debug`        | `false`                           | Run in debug mode. Usually, you need to run in this mode  only if you develop this tool or its maintainer is asking you to provide  logs with such verbosity.                                                                                                              |
+| `MTG_VERBOSE`            | `-v`, `--verbose`      | `false`                           | Run in verbose mode. This is way less chatty than debug mode.                                                                                                                                                                                                              |
+| `MTG_IP`                 | `-b`, `--bind-ip`      | `127.0.0.1`                       | Which IP should we bind to. As usual, `0.0.0.0` means that we want to listen on all interfaces. Also, 4 zeroes will bind to both IPv4 and IPv6.                                                                                                                            |
+| `MTG_PORT`               | `-p`, `--bind-port`    | `3128`                            | Which port should we bind to (listen on).                                                                                                                                                                                                                                  |
+| `MTG_IPV4`               | `-4`, `--public-ipv4`  | [Autodetect](https://ifconfig.co) | IPv4 address of this proxy. This is required if you NAT your proxy or run it in a docker container. In that case, you absolutely need to specify public IPv4 address of the proxy, otherwise either URLs will be broken or proxy could not access Telegram middle proxies. |
+| `MTG_IPV4_PORT`          | `--public-ipv4-port`   | Value of `--bind-port`            | Which port should be public of IPv4 interface. This affects only generated links and should be changed only if you NAT your proxy or run it in a docker container.                                                                                                         |
+| `MTG_IPV6`               | `-6`, `--public-ipv6`  | [Autodetect](https://ifconfig.co) | IPv6 address of this proxy. This is required if you NAT your proxy or run it in a docker container. In that case, you absolutely need to specify public IPv6 address of the proxy, otherwise either URLs will be broken or proxy could not access Telegram middle proxies. |
+| `MTG_IPV6_PORT`          | `--public-ipv6-port`   | Value of `--bind-port`            | Which port should be public of IPv6 interface. This affects only generated links and should be changed only if you NAT your proxy or run it in a docker container.                                                                                                         |
+| `MTG_STATS_IP`           | `-t`, `--stats-ip`     | `127.0.0.1`                       | Which IP should we bind the internal statistics HTTP server.                                                                                                                                                                                                               |
+| `MTG_STATS_PORT`         | `-q`, `--stats-port`   | `3129`                            | Which port should we bind the internal statistics HTTP server.                                                                                                                                                                                                             |
+| `MTG_STATSD_IP`          | `--statsd-ip`          |                                   | IP/host addresses of statsd service. No defaults, by defaults we do not send anything there.                                                                                                                                                                               |
+| `MTG_STATSD_PORT`        | `--statsd-port`        | `8125`                            | Which port should we use to work with statsd.                                                                                                                                                                                                                              |
+| `MTG_STATSD_NETWORK`     | `--statsd-network`     | `udp`                             | Which protocol should we use to work with statsd. Possible options are `udp` and `tcp`.                                                                                                                                                                                    |
+| `MTG_STATSD_PREFIX`      | `--statsd-prefix`      | `mtg`                             | Which bucket prefix we should use. For example, if you set `mtg`, then metric `traffic.ingress` would be send as `mtg.traffic.ingress`.                                                                                                                                    |
+| `MTG_STATSD_TAGS_FORMAT` | `--statsd-tags-format` |                                   | Which tags format we should use. By default, we are using default vanilla statsd tags format but if you want to send directly to InfluxDB or Datadog, please specify it there. Possible options are `influxdb` and `datadog`.                                              |
+| `MTG_STATSD_TAGS`        | `--statsd-tags`        |                                   | Which tags should we send to statsd with our metrics. Please specify them as `key=value` pairs.                                                                                                                                                                            |
+| `MTG_BUFFER_WRITE`       | `-w`, `--write-buffer` | `65536`                           | The size of TCP write buffer in bytes. Write buffer is the buffer for messages which are going from client to Telegram.                                                                                                                                                    |
+| `MTG_BUFFER_READ`        | `-r`, `--read-buffer`  | `131072`                          | The size of TCP read buffer in bytes. Read buffer is the buffer for messages from Telegram to client.                                                                                                                                                                      |
+
+Usually you want to modify only read/write buffer sizes. If you feel
+that proxy is slow, try to increase both sizes giving more priority to
+read buffer.
+
+Unfortunately, MTPROTO proxy protocol does not allow us to use splice
+or any other neat tricks how to eliminate the need of copying data into
+userspace.
 
 # How to run the tool
 
 Now run the tool:
 
 ```console
-$ mtg <secret>
+mtg <secret>
 ```
 
 How to run the tool with ADTag:
 
 ```console
-$ mtg <secret> <adtag>
+mtg <secret> <adtag>
 ```
 
 This tool will listen on port 3128 by default with the given secret.
@@ -151,13 +184,13 @@ This tool will listen on port 3128 by default with the given secret.
 # One-line runner
 
 ```console
-$ docker run --name mtg --restart=unless-stopped -p 3128:3128 -p 3129:3129 -d nineseconds/mtg:stable $(openssl rand -hex 16)
+docker run --name mtg --restart=unless-stopped -p 3128:3128 -p 3129:3129 -d nineseconds/mtg:stable $(openssl rand -hex 16)
 ```
 
 or in secret mode:
 
 ```console
-$ docker run --name mtg --restart=unless-stopped -p 3128:3128 -p 3129:3129 -d nineseconds/mtg:stable dd$(openssl rand -hex 16)
+docker run --name mtg --restart=unless-stopped -p 3128:3128 -p 3129:3129 -d nineseconds/mtg:stable dd$(openssl rand -hex 16)
 ```
 
 You will have this tool up and running on port 3128. Now curl

client/client.go

@@ -1,6 +1,7 @@
 package client
 
 import (
+	"context"
 	"net"
 
 	"github.com/9seconds/mtg/config"
@@ -9,4 +10,5 @@ import (
 )
 
 // Init defines common method for initializing client connections.
-type Init func(net.Conn, string, *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error)
+type Init func(context.Context, context.CancelFunc, net.Conn, string,
+	*config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error)

client/direct.go

@@ -1,6 +1,7 @@
 package client
 
 import (
+	"context"
 	"net"
 	"time"
 
@@ -12,23 +13,20 @@ 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.
-func DirectInit(socket net.Conn, connID string, conf *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) {
+func DirectInit(ctx context.Context, cancel context.CancelFunc, socket net.Conn,
+	connID string, conf *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) {
 	tcpSocket := socket.(*net.TCPConn)
 	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")
 	}
 
@@ -39,7 +37,7 @@ func DirectInit(socket net.Conn, connID string, conf *config.Config) (wrappers.W
 	}
 	socket.SetReadDeadline(time.Time{}) // nolint: errcheck
 
-	conn := wrappers.NewConn(socket, connID, wrappers.ConnPurposeClient, conf.PublicIPv4, conf.PublicIPv6)
+	conn := wrappers.NewConn(ctx, cancel, socket, connID, wrappers.ConnPurposeClient, conf.PublicIPv4, conf.PublicIPv6)
 	obfs2, connOpts, err := obfuscated2.ParseObfuscated2ClientFrame(conf.Secret, frame)
 	if err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot parse obfuscated frame")

client/middle.go

@@ -1,6 +1,7 @@
 package client
 
 import (
+	"context"
 	"net"
 
 	"github.com/9seconds/mtg/config"
@@ -10,8 +11,9 @@ import (
 
 // MiddleInit initializes client connection for proxy which has to
 // support promoted channels, connect to Telegram middle proxies etc.
-func MiddleInit(socket net.Conn, connID string, conf *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) {
+func MiddleInit(ctx context.Context, cancel context.CancelFunc, socket net.Conn,
-	conn, opts, err := DirectInit(socket, connID, conf)
+	connID string, conf *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) {
+	conn, opts, err := DirectInit(ctx, cancel, socket, connID, conf)
 	if err != nil {
 		return nil, nil, err
 	}

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,33 +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
-	bindIP net.IP, bindPort uint16,
+	writeBufferSize, readBufferSize uint32,
-	publicIPv4 net.IP, PublicIPv4Port uint16,
+	bindIP, publicIPv4, publicIPv6, statsIP net.IP,
-	publicIPv6 net.IP, publicIPv6Port uint16,
+	bindPort, publicIPv4Port, publicIPv6Port, statsPort, statsdPort uint16,
-	statsIP net.IP, statsPort uint16,
+	statsdIP, statsdNetwork, statsdPrefix, statsdTagsFormat string,
-	secret, adtag string,
+	statsdTags map[string]string,
-	statsdIP string, statsdPort uint16, statsdNetwork string, statsdPrefix string,
+	secret, adtag []byte) (*Config, error) {
-	statsdTagsFormat string, statsdTags map[string]string) (*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 {
@@ -149,8 +134,8 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo
 			return nil, errors.Errorf("IP %s is not IPv4", publicIPv4.String())
 		}
 	}
-	if PublicIPv4Port == 0 {
+	if publicIPv4Port == 0 {
-		PublicIPv4Port = bindPort
+		publicIPv4Port = bindPort
 	}
 
 	if publicIPv6 == nil {
@@ -170,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 != "" {
@@ -190,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":

config/global_ips.go

@@ -1,6 +1,7 @@
 package config
 
 import (
+	"context"
 	"io/ioutil"
 	"net"
 	"net/http"
@@ -9,16 +10,28 @@ import (
 	"github.com/juju/errors"
 )
 
+const ifconfigAddress = "https://ifconfig.co/ip"
+
 func getGlobalIPv4() (net.IP, error) {
-	return fetchIP("https://v4.ifconfig.co/ip")
+	return fetchIP("tcp4")
 }
 
 func getGlobalIPv6() (net.IP, error) {
-	return fetchIP("https://v6.ifconfig.co/ip")
+	return fetchIP("tcp6")
 }
 
-func fetchIP(url string) (net.IP, error) {
+func fetchIP(network string) (net.IP, error) {
-	resp, err := http.Get(url)
+	dialer := &net.Dialer{DualStack: false}
+	client := &http.Client{
+		Jar: nil,
+		Transport: &http.Transport{
+			DialContext: func(ctx context.Context, _, addr string) (net.Conn, error) {
+				return dialer.DialContext(ctx, network, addr)
+			},
+		},
+	}
+
+	resp, err := client.Get(ifconfigAddress)
 	if err != nil {
 		return nil, err
 	}

config/urls.go

@@ -42,7 +42,7 @@ func makeTMeURL(values url.Values) string {
 }
 
 func makeQRCodeURL(data string) string {
-	QRURL := url.URL{
+	qr := url.URL{
 		Scheme: "https",
 		Host:   "api.qrserver.com",
 		Path:   "v1/create-qr-code",
@@ -52,7 +52,7 @@ func makeQRCodeURL(data string) string {
 	values.Set("qzone", "4")
 	values.Set("format", "svg")
 	values.Set("data", data)
-	QRURL.RawQuery = values.Encode()
+	qr.RawQuery = values.Encode()
 
-	return QRURL.String()
+	return qr.String()
 }

main.go

@@ -25,86 +25,115 @@ import (
 var (
 	app = kingpin.New("mtg", "Simple MTPROTO proxy.")
 
-	debug = app.Flag("debug", "Run in debug mode.").
+	debug = app.Flag("debug",
+		"Run in debug mode.").
 		Short('d').
 		Envar("MTG_DEBUG").
 		Bool()
-	verbose = app.Flag("verbose", "Run in verbose mode.").
+	verbose = app.Flag("verbose",
+		"Run in verbose mode.").
 		Short('v').
 		Envar("MTG_VERBOSE").
 		Bool()
 
-	bindIP = app.Flag("bind-ip", "Which IP to bind to.").
+	bindIP = app.Flag("bind-ip",
+		"Which IP to bind to.").
 		Short('b').
 		Envar("MTG_IP").
 		Default("127.0.0.1").
 		IP()
-	bindPort = app.Flag("bind-port", "Which port to bind to.").
+	bindPort = app.Flag("bind-port",
-			Short('p').
+		"Which port to bind to.").
-			Envar("MTG_PORT").
+		Short('p').
-			Default("3128").
+		Envar("MTG_PORT").
-			Uint16()
+		Default("3128").
+		Uint16()
 
-	publicIPv4 = app.Flag("public-ipv4", "Which IPv4 address is public.").
+	publicIPv4 = app.Flag("public-ipv4",
-			Short('4').
+		"Which IPv4 address is public.").
-			Envar("MTG_IPV4").
+		Short('4').
-			IP()
+		Envar("MTG_IPV4").
-	publicIPv4Port = app.Flag("public-ipv4-port", "Which IPv4 port is public. Default is 'bind-port' value.").
+		IP()
-			Envar("MTG_IPV4_PORT").
+	publicIPv4Port = app.Flag("public-ipv4-port",
-			Uint16()
+		"Which IPv4 port is public. Default is 'bind-port' value.").
+		Envar("MTG_IPV4_PORT").
+		Uint16()
 
-	publicIPv6 = app.Flag("public-ipv6", "Which IPv6 address is public.").
+	publicIPv6 = app.Flag("public-ipv6",
-			Short('6').
+		"Which IPv6 address is public.").
-			Envar("MTG_IPV6").
+		Short('6').
-			IP()
+		Envar("MTG_IPV6").
-	publicIPv6Port = app.Flag("public-ipv6-port", "Which IPv6 port is public. Default is 'bind-port' value.").
+		IP()
-			Envar("MTG_IPV6_PORT").
+	publicIPv6Port = app.Flag("public-ipv6-port",
-			Uint16()
+		"Which IPv6 port is public. Default is 'bind-port' value.").
+		Envar("MTG_IPV6_PORT").
+		Uint16()
 
-	statsIP = app.Flag("stats-ip", "Which IP bind stats server to.").
+	statsIP = app.Flag("stats-ip",
+		"Which IP bind stats server to.").
 		Short('t').
 		Envar("MTG_STATS_IP").
 		Default("127.0.0.1").
 		IP()
-	statsPort = app.Flag("stats-port", "Which port bind stats to.").
+	statsPort = app.Flag("stats-port",
-			Short('q').
+		"Which port bind stats to.").
-			Envar("MTG_STATS_PORT").
+		Short('q').
-			Default("3129").
+		Envar("MTG_STATS_PORT").
-			Uint16()
+		Default("3129").
+		Uint16()
 
-	statsdIP = app.Flag("statsd-ip", "Which IP should we use for working with statsd.").
+	statsdIP = app.Flag("statsd-ip",
-			Envar("MTG_STATSD_IP").
+		"Which IP should we use for working with statsd.").
-			String()
+		Envar("MTG_STATSD_IP").
-	statsdPort = app.Flag("statsd-port", "Which port should we use for working with statsd.").
+		String()
-			Envar("MTG_STATSD_PORT").
+	statsdPort = app.Flag("statsd-port",
-			Default("8125").
+		"Which port should we use for working with statsd.").
-			Uint16()
+		Envar("MTG_STATSD_PORT").
-	statsdNetwork = app.Flag("statsd-network", "Which network is used to work with statsd. Only 'tcp' and 'udp' are supported.").
+		Default("8125").
-			Envar("MTG_STATSD_NETWORK").
+		Uint16()
-			Default("udp").
+	statsdNetwork = app.Flag("statsd-network",
-			String()
+		"Which network is used to work with statsd. Only 'tcp' and 'udp' are supported.").
-	statsdPrefix = app.Flag("statsd-prefix", "Which bucket prefix should we use for sending stats to statsd.").
+		Envar("MTG_STATSD_NETWORK").
-			Envar("MTG_STATSD_PREFIX").
+		Default("udp").
-			Default("mtg").
+		String()
-			String()
+	statsdPrefix = app.Flag("statsd-prefix",
-	statsdTagsFormat = app.Flag("statsd-tags-format", "Which tag format should we use to send stats metrics. Valid options are 'datadog' and 'influxdb'.").
+		"Which bucket prefix should we use for sending stats to statsd.").
-				Envar("MTG_STATSD_TAGS_FORMAT").
+		Envar("MTG_STATSD_PREFIX").
-				String()
+		Default("mtg").
-	statsdTags = app.Flag("statsd-tags", "Tags to use for working with statsd (specified as 'key=value').").
+		String()
-			Envar("MTG_STATSD_TAGS").
+	statsdTagsFormat = app.Flag("statsd-tags-format",
-			StringMap()
+		"Which tag format should we use to send stats metrics. Valid options are 'datadog' and 'influxdb'.").
+		Envar("MTG_STATSD_TAGS_FORMAT").
+		String()
+	statsdTags = app.Flag("statsd-tags",
+		"Tags to use for working with statsd (specified as 'key=value').").
+		Envar("MTG_STATSD_TAGS").
+		StringMap()
 
-	secret = app.Arg("secret", "Secret of this proxy.").Required().String()
+	writeBufferSize = app.Flag("write-buffer",
-	adtag  = app.Arg("adtag", "ADTag of the proxy.").String()
+		"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()
+
+	secret = app.Arg("secret", "Secret of this proxy.").Required().HexBytes()
+	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()
@@ -113,24 +142,24 @@ func main() {
 	}
 
 	conf, err := config.NewConfig(*debug, *verbose,
-		*bindIP, *bindPort,
+		*writeBufferSize, *readBufferSize,
-		*publicIPv4, *publicIPv4Port,
+		*bindIP, *publicIPv4, *publicIPv6, *statsIP,
-		*publicIPv6, *publicIPv6Port,
+		*bindPort, *publicIPv4Port, *publicIPv6Port, *statsPort, *statsdPort,
-		*statsIP, *statsPort,
+		*statsdIP, *statsdNetwork, *statsdPrefix, *statsdTagsFormat,
+		*statsdTags,
 		*secret, *adtag,
-		*statsdIP, *statsdPort, *statsdNetwork, *statsdPrefix,
-		*statsdTagsFormat, *statsdTags,
 	)
 	if err != nil {
 		usage(err.Error())
 	}
 
 	atom := zap.NewAtomicLevel()
-	if conf.Debug {
+	switch {
+	case conf.Debug:
 		atom.SetLevel(zapcore.DebugLevel)
-	} else if conf.Verbose {
+	case conf.Verbose:
 		atom.SetLevel(zapcore.InfoLevel)
-	} else {
+	default:
 		atom.SetLevel(zapcore.ErrorLevel)
 	}
 	encoderCfg := zap.NewProductionEncoderConfig()
@@ -150,7 +179,8 @@ func main() {
 			zap.S().Warnw("Could not fetch time data from NTP")
 		} else {
 			if diff >= time.Second {
-				usage(fmt.Sprintf("You choose to use middle proxy but your clock drift (%s) is bigger than 1 second. Please, sync your time", diff))
+				usage(fmt.Sprintf("You choose to use middle proxy but your clock drift (%s) "+
+					"is bigger than 1 second. Please, sync your time", diff))
 			}
 			go ntp.AutoUpdate()
 		}

mtproto/rpc/handshake_response.go

@@ -19,10 +19,10 @@ type HandshakeResponse struct {
 func (r *HandshakeResponse) Bytes() []byte {
 	buf := &bytes.Buffer{}
 
-	buf.Write(r.Type[:])
+	buf.Write(r.Type)
-	buf.Write(r.Flags[:])
+	buf.Write(r.Flags)
-	buf.Write(r.SenderPID[:])
+	buf.Write(r.SenderPID)
-	buf.Write(r.PeerPID[:])
+	buf.Write(r.PeerPID)
 
 	return buf.Bytes()
 }

mtproto/rpc/proxy_flags.go

@@ -9,13 +9,13 @@ type proxyRequestFlags uint32
 
 const (
 	proxyRequestFlagsHasAdTag     proxyRequestFlags = 0x8
-	proxyRequestFlagsEncrypted                      = 0x2
+	proxyRequestFlagsEncrypted    proxyRequestFlags = 0x2
-	proxyRequestFlagsMagic                          = 0x1000
+	proxyRequestFlagsMagic        proxyRequestFlags = 0x1000
-	proxyRequestFlagsExtMode2                       = 0x20000
+	proxyRequestFlagsExtMode2     proxyRequestFlags = 0x20000
-	proxyRequestFlagsIntermediate                   = 0x20000000
+	proxyRequestFlagsIntermediate proxyRequestFlags = 0x20000000
-	proxyRequestFlagsAbdridged                      = 0x40000000
+	proxyRequestFlagsAbdridged    proxyRequestFlags = 0x40000000
-	proxyRequestFlagsQuickAck                       = 0x80000000
+	proxyRequestFlagsQuickAck     proxyRequestFlags = 0x80000000
-	proxyRequestFlagsPad                            = 0x8000000
+	proxyRequestFlagsPad          proxyRequestFlags = 0x8000000
 )
 
 var proxyRequestFlagsEncryptedPrefix [8]byte

mtproto/rpc/proxy_request.go

@@ -64,7 +64,8 @@ func (r *ProxyRequest) MakeHeader(message []byte) (*bytes.Buffer, fmt.Stringer)
 }
 
 // NewProxyRequest build new ProxyRequest data structure.
-func NewProxyRequest(clientAddr, ownAddr *net.TCPAddr, opts *mtproto.ConnectionOpts, adTag []byte) (*ProxyRequest, error) {
+func NewProxyRequest(clientAddr, ownAddr *net.TCPAddr,
+	opts *mtproto.ConnectionOpts, adTag []byte) (*ProxyRequest, error) {
 	flags := proxyRequestFlagsHasAdTag | proxyRequestFlagsMagic | proxyRequestFlagsExtMode2
 
 	switch opts.ConnectionType {

obfuscated2/obfuscated2_test.go

@@ -47,8 +47,8 @@ func TestObfs2Full(t *testing.T) {
 
 	clientFrame := generateFrame(mtproto.ConnectionTypeIntermediate)
 	clientHasher := sha256.New()
-	clientHasher.Write(clientFrame.Key())
+	clientHasher.Write(clientFrame.Key()) // nolint: errcheck
-	clientHasher.Write(secret)
+	clientHasher.Write(secret)            // nolint: errcheck
 	clientKey := clientHasher.Sum(nil)
 
 	encryptor := makeStreamCipher(clientKey, clientFrame.IV())
@@ -58,8 +58,8 @@ func TestObfs2Full(t *testing.T) {
 
 	invertedClientFrame := clientFrame.Invert()
 	clientHasher = sha256.New()
-	clientHasher.Write(invertedClientFrame.Key())
+	clientHasher.Write(invertedClientFrame.Key()) // nolint: errcheck
-	clientHasher.Write(secret)
+	clientHasher.Write(secret)                    // nolint: errcheck
 	invertedClientKey := clientHasher.Sum(nil)
 	clientDecryptor := makeStreamCipher(invertedClientKey, invertedClientFrame.IV())
 

proxy/proxy.go

@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"context"
 	"io"
 	"net"
 	"sync"
@@ -43,8 +44,10 @@ func (p *Proxy) Serve() error {
 func (p *Proxy) accept(conn net.Conn) {
 	connID := uuid.NewV4().String()
 	log := zap.S().With("connection_id", connID).Named("main")
+	ctx, cancel := context.WithCancel(context.Background())
 
 	defer func() {
+		cancel()
 		conn.Close() // nolint: errcheck
 
 		if err := recover(); err != nil {
@@ -55,7 +58,7 @@ func (p *Proxy) accept(conn net.Conn) {
 
 	log.Infow("Client connected", "addr", conn.RemoteAddr())
 
-	clientConn, opts, err := p.clientInit(conn, connID, p.conf)
+	clientConn, opts, err := p.clientInit(ctx, cancel, conn, connID, p.conf)
 	if err != nil {
 		log.Errorw("Cannot initialize client connection", "error", err)
 		return
@@ -65,13 +68,19 @@ func (p *Proxy) accept(conn net.Conn) {
 	stats.ClientConnected(opts.ConnectionType, clientConn.RemoteAddr())
 	defer stats.ClientDisconnected(opts.ConnectionType, clientConn.RemoteAddr())
 
-	serverConn, err := p.getTelegramConn(opts, connID)
+	serverConn, err := p.getTelegramConn(ctx, cancel, opts, connID)
 	if err != nil {
 		log.Errorw("Cannot initialize server connection", "error", err)
 		return
 	}
 	defer serverConn.(io.Closer).Close() // nolint: errcheck
 
+	go func() {
+		<-ctx.Done()
+		serverConn.(io.Closer).Close()
+		clientConn.(io.Closer).Close()
+	}()
+
 	wait := &sync.WaitGroup{}
 	wait.Add(2)
 
@@ -83,8 +92,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()
@@ -92,8 +101,9 @@ func (p *Proxy) accept(conn net.Conn) {
 	log.Infow("Client disconnected", "addr", conn.RemoteAddr())
 }
 
-func (p *Proxy) getTelegramConn(opts *mtproto.ConnectionOpts, connID string) (wrappers.Wrap, error) {
+func (p *Proxy) getTelegramConn(ctx context.Context, cancel context.CancelFunc,
-	streamConn, err := p.tg.Dial(connID, opts)
+	opts *mtproto.ConnectionOpts, connID string) (wrappers.Wrap, error) {
+	streamConn, err := p.tg.Dial(ctx, cancel, connID, opts)
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot dial to Telegram")
 	}
@@ -106,7 +116,8 @@ func (p *Proxy) getTelegramConn(opts *mtproto.ConnectionOpts, connID string) (wr
 	return packetConn, nil
 }
 
-func (p *Proxy) middlePipe(src wrappers.PacketReadCloser, dst io.WriteCloser, wait *sync.WaitGroup, hacks *mtproto.Hacks) {
+func (p *Proxy) middlePipe(src wrappers.PacketReadCloser, dst io.WriteCloser,
+	wait *sync.WaitGroup, hacks *mtproto.Hacks) {
 	defer func() {
 		src.Close() // nolint: errcheck
 		dst.Close() // nolint: errcheck
@@ -129,14 +140,16 @@ func (p *Proxy) middlePipe(src wrappers.PacketReadCloser, dst io.WriteCloser, wa
 	}
 }
 
-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)
 	}
 }

stats/server.go

@@ -13,7 +13,7 @@ import (
 
 var instance *stats
 
-// Start starts new statisitcs server.
+// Start starts new statistics server.
 func Start(conf *config.Config) error {
 	log := zap.S().Named("stats")
 

telegram/dialer.go

@@ -1,6 +1,7 @@
 package telegram
 
 import (
+	"context"
 	"net"
 	"time"
 
@@ -10,11 +11,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,22 +29,24 @@ 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")
 	}
 
 	return conn, nil
 }
 
-func (t *tgDialer) dialRWC(addr, connID string) (wrappers.StreamReadWriteCloser, error) {
+func (t *tgDialer) dialRWC(ctx context.Context, cancel context.CancelFunc,
+	addr, connID string) (wrappers.StreamReadWriteCloser, error) {
 	conn, err := t.dial(addr)
 	if err != nil {
 		return nil, err
 	}
-	tgConn := wrappers.NewConn(conn, connID, wrappers.ConnPurposeTelegram, t.conf.PublicIPv4, t.conf.PublicIPv6)
+	tgConn := wrappers.NewConn(ctx, cancel, conn, connID,
+		wrappers.ConnPurposeTelegram, t.conf.PublicIPv4, t.conf.PublicIPv6)
 
 	return tgConn, nil
 }

telegram/direct.go

@@ -1,6 +1,7 @@
 package telegram
 
 import (
+	"context"
 	"net"
 
 	"github.com/juju/errors"
@@ -13,18 +14,18 @@ import (
 
 var (
 	directV4Addresses = map[int16][]string{
-		0: []string{"149.154.175.50:443"},
+		0: {"149.154.175.50:443"},
-		1: []string{"149.154.167.51:443"},
+		1: {"149.154.167.51:443"},
-		2: []string{"149.154.175.100:443"},
+		2: {"149.154.175.100:443"},
-		3: []string{"149.154.167.91:443"},
+		3: {"149.154.167.91:443"},
-		4: []string{"149.154.171.5:443"},
+		4: {"149.154.171.5:443"},
 	}
 	directV6Addresses = map[int16][]string{
-		0: []string{"[2001:b28:f23d:f001::a]:443"},
+		0: {"[2001:b28:f23d:f001::a]:443"},
-		1: []string{"[2001:67c:04e8:f002::a]:443"},
+		1: {"[2001:67c:04e8:f002::a]:443"},
-		2: []string{"[2001:b28:f23d:f003::a]:443"},
+		2: {"[2001:b28:f23d:f003::a]:443"},
-		3: []string{"[2001:67c:04e8:f004::a]:443"},
+		3: {"[2001:67c:04e8:f004::a]:443"},
-		4: []string{"[2001:b28:f23f:f005::a]:443"},
+		4: {"[2001:b28:f23f:f005::a]:443"},
 	}
 )
 
@@ -32,7 +33,8 @@ type directTelegram struct {
 	baseTelegram
 }
 
-func (t *directTelegram) Dial(connID string, connOpts *mtproto.ConnectionOpts) (wrappers.StreamReadWriteCloser, error) {
+func (t *directTelegram) Dial(ctx context.Context, cancel context.CancelFunc,
+	connID string, connOpts *mtproto.ConnectionOpts) (wrappers.StreamReadWriteCloser, error) {
 	dc := connOpts.DC
 	if dc < 0 {
 		dc = -dc
@@ -40,10 +42,11 @@ func (t *directTelegram) Dial(connID string, connOpts *mtproto.ConnectionOpts) (
 		dc = 1
 	}
 
-	return t.baseTelegram.dial(dc-1, connID, connOpts.ConnectionProto)
+	return t.baseTelegram.dial(ctx, cancel, dc-1, connID, connOpts.ConnectionProto)
 }
 
-func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.StreamReadWriteCloser) (wrappers.Wrap, error) {
+func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts,
+	conn wrappers.StreamReadWriteCloser) (wrappers.Wrap, error) {
 	obfs2, frame := obfuscated2.MakeTelegramObfuscated2Frame(connOpts)
 
 	if _, err := conn.Write(frame); err != nil {
@@ -56,12 +59,14 @@ func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.St
 // NewDirectTelegram returns Telegram instance which connects directly
 // to Telegram bypassing middleproxies.
 func NewDirectTelegram(conf *config.Config) Telegram {
-	return &directTelegram{baseTelegram{
+	return &directTelegram{
-		dialer: tgDialer{
+		baseTelegram: baseTelegram{
-			Dialer: net.Dialer{Timeout: telegramDialTimeout},
+			dialer: tgDialer{
-			conf:   conf,
+				Dialer: net.Dialer{Timeout: telegramDialTimeout},
+				conf:   conf,
+			},
+			v4Addresses: directV4Addresses,
+			v6Addresses: directV6Addresses,
 		},
-		v4Addresses: directV4Addresses,
+	}
-		v6Addresses: directV6Addresses,
-	}}
 }

telegram/middle.go

@@ -20,7 +20,8 @@ type middleTelegram struct {
 	conf *config.Config
 }
 
-func (t *middleTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.StreamReadWriteCloser) (wrappers.Wrap, error) {
+func (t *middleTelegram) Init(connOpts *mtproto.ConnectionOpts,
+	conn wrappers.StreamReadWriteCloser) (wrappers.Wrap, error) {
 	rpcNonceConn := wrappers.NewMTProtoFrame(conn, rpc.SeqNoNonce)
 
 	rpcNonceReq, err := t.sendRPCNonceRequest(rpcNonceConn)
@@ -65,7 +66,8 @@ func (t *middleTelegram) sendRPCNonceRequest(conn io.Writer) (*rpc.NonceRequest,
 	return rpcNonceReq, nil
 }
 
-func (t *middleTelegram) receiveRPCNonceResponse(conn wrappers.PacketReader, req *rpc.NonceRequest) (*rpc.NonceResponse, error) {
+func (t *middleTelegram) receiveRPCNonceResponse(conn wrappers.PacketReader,
+	req *rpc.NonceRequest) (*rpc.NonceResponse, error) {
 	packet, err := conn.Read()
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot read RPC nonce response")
@@ -91,7 +93,8 @@ func (t *middleTelegram) sendRPCHandshakeRequest(conn io.Writer) (*rpc.Handshake
 	return req, nil
 }
 
-func (t *middleTelegram) receiveRPCHandshakeResponse(conn wrappers.PacketReader, req *rpc.HandshakeRequest) (*rpc.HandshakeResponse, error) {
+func (t *middleTelegram) receiveRPCHandshakeResponse(conn wrappers.PacketReader,
+	req *rpc.HandshakeRequest) (*rpc.HandshakeResponse, error) {
 	packet, err := conn.Read()
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot read RPC handshake response")

telegram/middle_caller.go

@@ -2,6 +2,7 @@ package telegram
 
 import (
 	"bufio"
+	"context"
 	"io/ioutil"
 	"net"
 	"net/http"
@@ -38,7 +39,8 @@ type middleTelegramCaller struct {
 	httpClient  *http.Client
 }
 
-func (t *middleTelegramCaller) Dial(connID string, connOpts *mtproto.ConnectionOpts) (wrappers.StreamReadWriteCloser, error) {
+func (t *middleTelegramCaller) Dial(ctx context.Context, cancel context.CancelFunc, connID string,
+	connOpts *mtproto.ConnectionOpts) (wrappers.StreamReadWriteCloser, error) {
 	dc := connOpts.DC
 	if dc == 0 {
 		dc = 1
@@ -46,7 +48,7 @@ func (t *middleTelegramCaller) Dial(connID string, connOpts *mtproto.ConnectionO
 	t.dialerMutex.RLock()
 	defer t.dialerMutex.RUnlock()
 
-	return t.baseTelegram.dial(dc, connID, connOpts.ConnectionProto)
+	return t.baseTelegram.dial(ctx, cancel, dc, connID, connOpts.ConnectionProto)
 }
 
 func (t *middleTelegramCaller) autoUpdate() {

telegram/telegram.go

@@ -1,6 +1,7 @@
 package telegram
 
 import (
+	"context"
 	"math/rand"
 
 	"github.com/juju/errors"
@@ -11,7 +12,7 @@ import (
 
 // Telegram is an interface for different Telegram work modes.
 type Telegram interface {
-	Dial(string, *mtproto.ConnectionOpts) (wrappers.StreamReadWriteCloser, error)
+	Dial(context.Context, context.CancelFunc, string, *mtproto.ConnectionOpts) (wrappers.StreamReadWriteCloser, error)
 	Init(*mtproto.ConnectionOpts, wrappers.StreamReadWriteCloser) (wrappers.Wrap, error)
 }
 
@@ -22,7 +23,8 @@ type baseTelegram struct {
 	v6Addresses map[int16][]string
 }
 
-func (b *baseTelegram) dial(dcIdx int16, connID string, proto mtproto.ConnectionProtocol) (wrappers.StreamReadWriteCloser, error) {
+func (b *baseTelegram) dial(ctx context.Context, cancel context.CancelFunc, dcIdx int16, connID string,
+	proto mtproto.ConnectionProtocol) (wrappers.StreamReadWriteCloser, error) {
 	addrs := make([]string, 2)
 
 	if proto&mtproto.ConnectionProtocolIPv6 != 0 {
@@ -37,7 +39,7 @@ func (b *baseTelegram) dial(dcIdx int16, connID string, proto mtproto.Connection
 	}
 
 	for _, addr := range addrs {
-		if conn, err := b.dialer.dialRWC(addr, connID); err == nil {
+		if conn, err := b.dialer.dialRWC(ctx, cancel, addr, connID); err == nil {
 			return conn, err
 		}
 	}

wrappers/conn.go

@@ -1,15 +1,17 @@
 package wrappers
 
 import (
+	"context"
 	"net"
 	"time"
 
 	"go.uber.org/zap"
 
 	"github.com/9seconds/mtg/stats"
+	"github.com/juju/errors"
 )
 
-// ConnPurpose is intented to be identifier of connection purpose. We
+// ConnPurpose is intended to be identifier of connection purpose. We
 // sometimes want to treat client/telegram connection differently (for
 // logging for example).
 type ConnPurpose uint8
@@ -36,11 +38,20 @@ const (
 	connTimeoutWrite = 2 * time.Minute
 )
 
+type ioResult struct {
+	n   int
+	err error
+}
+
+type ioFunc func([]byte) (int, error)
+
 // Conn is a basic wrapper for net.Conn providing the most low-level
 // logic and management as possible.
 type Conn struct {
-	connID string
 	conn   net.Conn
+	ctx    context.Context
+	cancel context.CancelFunc
+	connID string
 	logger *zap.SugaredLogger
 
 	publicIPv4 net.IP
@@ -48,28 +59,64 @@ type Conn struct {
 }
 
 func (c *Conn) Write(p []byte) (int, error) {
-	c.conn.SetWriteDeadline(time.Now().Add(connTimeoutWrite)) // nolint: errcheck
+	select {
-	n, err := c.conn.Write(p)
+	case <-c.ctx.Done():
+		return 0, errors.Annotate(c.ctx.Err(), "Cannot write because context was closed")
+	default:
+		n, err := c.doIO(c.conn.Write, p, connTimeoutWrite)
 
-	c.logger.Debugw("Write to stream", "bytes", n, "error", err)
+		c.logger.Debugw("Write to stream", "bytes", n, "error", err)
-	stats.EgressTraffic(n)
+		stats.EgressTraffic(n)
 
-	return n, err
+		return n, err
+	}
 }
 
 func (c *Conn) Read(p []byte) (int, error) {
-	c.conn.SetReadDeadline(time.Now().Add(connTimeoutRead)) // nolint: errcheck
+	select {
-	n, err := c.conn.Read(p)
+	case <-c.ctx.Done():
+		return 0, errors.Annotate(c.ctx.Err(), "Cannot read because context was closed")
+	default:
+		n, err := c.doIO(c.conn.Read, p, connTimeoutRead)
 
-	c.logger.Debugw("Read from stream", "bytes", n, "error", err)
+		c.logger.Debugw("Read from stream", "bytes", n, "error", err)
-	stats.IngressTraffic(n)
+		stats.IngressTraffic(n)
 
-	return n, err
+		return n, err
+	}
+}
+
+func (c *Conn) doIO(callback ioFunc, p []byte, timeout time.Duration) (int, error) {
+	resChan := make(chan ioResult, 1)
+	timer := time.NewTimer(timeout)
+
+	go func() {
+		n, err := callback(p)
+		resChan <- ioResult{n: n, err: err}
+	}()
+
+	select {
+	case res := <-resChan:
+		timer.Stop()
+		if res.err != nil {
+			c.Close()
+		}
+		return res.n, res.err
+	case <-c.ctx.Done():
+		timer.Stop()
+		c.Close()
+		return 0, errors.Annotate(c.ctx.Err(), "Cannot do IO because context is closed")
+	case <-timer.C:
+		c.Close()
+		return 0, errors.Annotate(c.ctx.Err(), "Timeout on IO operation")
+	}
 }
 
 // Close closes underlying net.Conn instance.
 func (c *Conn) Close() error {
 	defer c.logger.Debugw("Close connection")
+
+	c.cancel()
 	return c.conn.Close()
 }
 
@@ -100,7 +147,8 @@ func (c *Conn) RemoteAddr() *net.TCPAddr {
 }
 
 // NewConn initializes Conn wrapper for net.Conn.
-func NewConn(conn net.Conn, connID string, purpose ConnPurpose, publicIPv4, publicIPv6 net.IP) StreamReadWriteCloser {
+func NewConn(ctx context.Context, cancel context.CancelFunc, conn net.Conn,
+	connID string, purpose ConnPurpose, publicIPv4, publicIPv6 net.IP) StreamReadWriteCloser {
 	logger := zap.S().With(
 		"connection_id", connID,
 		"local_address", conn.LocalAddr(),
@@ -109,9 +157,11 @@ func NewConn(conn net.Conn, connID string, purpose ConnPurpose, publicIPv4, publ
 	).Named("conn")
 
 	wrapper := Conn{
-		logger:     logger,
-		connID:     connID,
 		conn:       conn,
+		ctx:        ctx,
+		cancel:     cancel,
+		connID:     connID,
+		logger:     logger,
 		publicIPv4: publicIPv4,
 		publicIPv6: publicIPv6,
 	}

wrappers/mtproto_cipher.go

@@ -24,7 +24,8 @@ var emptyIP = [4]byte{0x00, 0x00, 0x00, 0x00}
 
 // NewMiddleProxyCipher creates new block cipher to proxy<->telegram
 // connection.
-func NewMiddleProxyCipher(conn StreamReadWriteCloser, req *rpc.NonceRequest, resp *rpc.NonceResponse, secret []byte) StreamReadWriteCloser {
+func NewMiddleProxyCipher(conn StreamReadWriteCloser,
+	req *rpc.NonceRequest, resp *rpc.NonceResponse, secret []byte) StreamReadWriteCloser {
 	localAddr := conn.LocalAddr()
 	remoteAddr := conn.RemoteAddr()
 
@@ -37,11 +38,12 @@ func NewMiddleProxyCipher(conn StreamReadWriteCloser, req *rpc.NonceRequest, res
 	return NewBlockCipher(conn, enc, dec)
 }
 
-func deriveKeys(purpose cipherPurpose, req *rpc.NonceRequest, resp *rpc.NonceResponse, client *net.TCPAddr, remote *net.TCPAddr, secret []byte) ([]byte, []byte) {
+func deriveKeys(purpose cipherPurpose, req *rpc.NonceRequest, resp *rpc.NonceResponse,
+	client, remote *net.TCPAddr, secret []byte) ([]byte, []byte) {
 	message := bytes.Buffer{}
-	message.Write(resp.Nonce[:])
+	message.Write(resp.Nonce)
-	message.Write(req.Nonce[:])
+	message.Write(req.Nonce)
-	message.Write(req.CryptoTS[:])
+	message.Write(req.CryptoTS)
 
 	clientIPv4 := emptyIP[:]
 	serverIPv4 := emptyIP[:]
@@ -68,13 +70,13 @@ func deriveKeys(purpose cipherPurpose, req *rpc.NonceRequest, resp *rpc.NonceRes
 	binary.LittleEndian.PutUint16(port[:], uint16(remote.Port))
 	message.Write(port[:])
 	message.Write(secret)
-	message.Write(resp.Nonce[:])
+	message.Write(resp.Nonce)
 
 	if client.IP.To4() == nil {
 		message.Write(client.IP.To16())
 		message.Write(remote.IP.To16())
 	}
-	message.Write(req.Nonce[:])
+	message.Write(req.Nonce)
 
 	data := message.Bytes()
 	md5sum := md5.Sum(data[1:]) // nolint: gas

wrappers/mtproto_frame.go

@@ -61,7 +61,8 @@ func (m *MTProtoFrame) Read() ([]byte, error) { // nolint: gocyclo
 		"messageLength", messageLength,
 		"sequence_number", m.readSeqNo,
 	)
-	if messageLength%4 != 0 || messageLength < mtprotoFrameMinMessageLength || messageLength > mtprotoFrameMaxMessageLength {
+	if messageLength%4 != 0 || messageLength < mtprotoFrameMinMessageLength ||
+		messageLength > mtprotoFrameMaxMessageLength {
 		return nil, errors.Errorf("Incorrect frame message length %d", messageLength)
 	}
 

wrappers/mtproto_proxy.go

@@ -56,7 +56,7 @@ func (m *MTProtoProxy) Read() ([]byte, error) {
 	case bytes.Equal(tag, rpc.TagSimpleAck):
 		return m.readSimpleAck(packet)
 	case bytes.Equal(tag, rpc.TagCloseExt):
-		return m.readCloseExt(packet)
+		return m.readCloseExt()
 	}
 
 	return nil, errors.Errorf("Unknown RPC answer %v", tag)
@@ -91,7 +91,7 @@ func (m *MTProtoProxy) readSimpleAck(data []byte) ([]byte, error) {
 	return data, nil
 }
 
-func (m *MTProtoProxy) readCloseExt(data []byte) ([]byte, error) {
+func (m *MTProtoProxy) readCloseExt() ([]byte, error) {
 	m.logger.Debugw("Read RPC_CLOSE_EXT", "counter", m.readCounter)
 
 	return nil, errors.New("Connection has been closed remotely by RPC call")
@@ -150,7 +150,8 @@ func (m *MTProtoProxy) Close() error {
 }
 
 // NewMTProtoProxy creates new RPC wrapper.
-func NewMTProtoProxy(conn PacketReadWriteCloser, connOpts *mtproto.ConnectionOpts, adTag []byte) (PacketReadWriteCloser, error) {
+func NewMTProtoProxy(conn PacketReadWriteCloser, connOpts *mtproto.ConnectionOpts,
+	adTag []byte) (PacketReadWriteCloser, error) {
 	req, err := rpc.NewProxyRequest(connOpts.ClientAddr, conn.LocalAddr(), connOpts, adTag)
 	if err != nil {
 		return nil, errors.Annotate(err, "Cannot create new RPC proxy request")

wrappers/streamcipher.go

@@ -1,6 +1,7 @@
 package wrappers
 
 import (
+	"bytes"
 	"crypto/cipher"
 	"net"
 
@@ -28,10 +29,17 @@ func (s *StreamCipher) Read(p []byte) (int, error) {
 }
 
 func (s *StreamCipher) Write(p []byte) (int, error) {
-	encrypted := make([]byte, len(p))
+	buf := streamCipherBufferPool.Get().(*bytes.Buffer)
-	s.encryptor.XORKeyStream(encrypted, p)
+	defer streamCipherBufferPool.Put(buf)
 
-	return s.conn.Write(encrypted)
+	buf.Reset()
+	buf.Grow(len(p))
+	buf.Write(p)
+
+	data := buf.Bytes()
+	s.encryptor.XORKeyStream(data, data)
+
+	return s.conn.Write(data)
 }
 
 // Logger returns an instance of the logger for this wrapper.

wrappers/streamcipher_pool.go

@@ -0,0 +1,16 @@
+package wrappers
+
+import (
+	"bytes"
+	"sync"
+)
+
+var streamCipherBufferPool sync.Pool
+
+func init() {
+	streamCipherBufferPool = sync.Pool{
+		New: func() interface{} {
+			return &bytes.Buffer{}
+		},
+	}
+}