Update dependencies

Base prevention of replay attacks on proxy

.dockerignore

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

.gitignore

@@ -8,3 +8,4 @@
 mtg
 vendor/
 version.go
+ccbuilds/

.golangci.toml

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

.gometalinter.json

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

.travis.yml

@@ -5,14 +5,17 @@ sudo: false
 dist: trusty
 
 go:
-  - "1.9.x"
+  - "1.11.x"
-  - "1.10.x"
+  - 1.12.x
   - master
 
+before_script: make prepare
+
 script:
+  - make all
   - make lint
   - make test
 
-cache:
+matrix:
-  directories:
+  allow_failures:
-    - vendor
+    - go: master

Dockerfile

@@ -1,7 +1,7 @@
 ###############################################################################
 # BUILD STAGE
 
-FROM golang:alpine
+FROM golang:1.12-alpine
 
 RUN set -x \
   && apk --no-cache --update add \
@@ -10,14 +10,12 @@ RUN set -x \
     curl \
     git \
     make \
-    upx \
+    upx
-  && update-ca-certificates
 
-ADD . /go/src/github.com/9seconds/mtg
+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
 
@@ -27,7 +25,7 @@ RUN set -x \
 
 FROM scratch
 
-ENTRYPOINT ["/usr/local/bin/mtg"]
+ENTRYPOINT ["/mtg"]
 ENV MTG_IP=0.0.0.0 \
     MTG_PORT=3128 \
     MTG_STATS_IP=0.0.0.0 \
@@ -35,4 +33,4 @@ ENV MTG_IP=0.0.0.0 \
 EXPOSE 3128 3129
 
 COPY --from=0 /etc/ssl/certs/ca-certificates.crt /etc/ssl/certs/ca-certificates.crt
-COPY --from=0 /go/src/github.com/9seconds/mtg/mtg /usr/local/bin/mtg
+COPY --from=0 /go/src/github.com/9seconds/mtg/mtg /mtg

Gopkg.lock

@@ -1,91 +0,0 @@
-# This file is autogenerated, do not edit; changes may be undone by the next 'dep ensure'.
-
-
-[[projects]]
-  branch = "master"
-  name = "github.com/alecthomas/template"
-  packages = [
-    ".",
-    "parse"
-  ]
-  revision = "a0175ee3bccc567396460bf5acd36800cb10c49c"
-
-[[projects]]
-  branch = "master"
-  name = "github.com/alecthomas/units"
-  packages = ["."]
-  revision = "2efee857e7cfd4f3d0138cc3cbb1b4966962b93a"
-
-[[projects]]
-  name = "github.com/davecgh/go-spew"
-  packages = ["spew"]
-  revision = "346938d642f2ec3594ed81d874461961cd0faa76"
-  version = "v1.1.0"
-
-[[projects]]
-  branch = "master"
-  name = "github.com/dustin/go-humanize"
-  packages = ["."]
-  revision = "02af3965c54e8cacf948b97fef38925c4120652c"
-
-[[projects]]
-  branch = "master"
-  name = "github.com/juju/errors"
-  packages = ["."]
-  revision = "c7d06af17c68cd34c835053720b21f6549d9b0ee"
-
-[[projects]]
-  name = "github.com/pmezard/go-difflib"
-  packages = ["difflib"]
-  revision = "792786c7400a136282c1664665ae0a8db921c6c2"
-  version = "v1.0.0"
-
-[[projects]]
-  name = "github.com/satori/go.uuid"
-  packages = ["."]
-  revision = "f58768cc1a7a7e77a3bd49e98cdd21419399b6a3"
-  version = "v1.2.0"
-
-[[projects]]
-  name = "github.com/stretchr/testify"
-  packages = ["assert"]
-  revision = "f35b8ab0b5a2cef36673838d662e249dd9c94686"
-  version = "v1.2.2"
-
-[[projects]]
-  name = "go.uber.org/atomic"
-  packages = ["."]
-  revision = "1ea20fb1cbb1cc08cbd0d913a96dead89aa18289"
-  version = "v1.3.2"
-
-[[projects]]
-  name = "go.uber.org/multierr"
-  packages = ["."]
-  revision = "3c4937480c32f4c13a875a1829af76c98ca3d40a"
-  version = "v1.1.0"
-
-[[projects]]
-  name = "go.uber.org/zap"
-  packages = [
-    ".",
-    "buffer",
-    "internal/bufferpool",
-    "internal/color",
-    "internal/exit",
-    "zapcore"
-  ]
-  revision = "eeedf312bc6c57391d84767a4cd413f02a917974"
-  version = "v1.8.0"
-
-[[projects]]
-  name = "gopkg.in/alecthomas/kingpin.v2"
-  packages = ["."]
-  revision = "947dcec5ba9c011838740e680966fd7087a71d0d"
-  version = "v2.2.6"
-
-[solve-meta]
-  analyzer-name = "dep"
-  analyzer-version = 1
-  inputs-digest = "c4fdd3664f683342ad0c2509f4a8bcfe5b267a6e8cdaf36f70d39536bbf89834"
-  solver-name = "gps-cdcl"
-  solver-version = 1

Gopkg.toml

@@ -1,50 +0,0 @@
-# Gopkg.toml example
-#
-# Refer to https://golang.github.io/dep/docs/Gopkg.toml.html
-# for detailed Gopkg.toml documentation.
-#
-# required = ["github.com/user/thing/cmd/thing"]
-# ignored = ["github.com/user/project/pkgX", "bitbucket.org/user/project/pkgA/pkgY"]
-#
-# [[constraint]]
-#   name = "github.com/user/project"
-#   version = "1.0.0"
-#
-# [[constraint]]
-#   name = "github.com/user/project2"
-#   branch = "dev"
-#   source = "github.com/myfork/project2"
-#
-# [[override]]
-#   name = "github.com/x/y"
-#   version = "2.4.0"
-#
-# [prune]
-#   non-go = false
-#   go-tests = true
-#   unused-packages = true
-
-
-[prune]
-  go-tests = true
-  unused-packages = true
-
-[[constraint]]
-  name = "gopkg.in/alecthomas/kingpin.v2"
-  version = "2.2.6"
-
-[[constraint]]
-  branch = "master"
-  name = "github.com/juju/errors"
-
-[[constraint]]
-  name = "github.com/stretchr/testify"
-  version = "1.2.1"
-
-[[constraint]]
-  name = "github.com/satori/go.uuid"
-  version = "1.2.0"
-
-[[constraint]]
-  branch = "master"
-  name = "github.com/dustin/go-humanize"

Makefile

@@ -1,26 +1,31 @@
 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}")
-CC_BINARIES  := $(shell bash -c "echo -n $(APP_NAME)-{linux,windows,darwin,freebsd,openbsd}-{386,amd64} $(APP_NAME)-linux-{arm,arm64}")
-APP_DEPS     := version.go $(VENDOR_FILES)
 
-COMMON_BUILD_FLAGS := -ldflags="-s -w"
+GOLANGCI_LINT_VERSION := v1.15.0
+
+VERSION_GO         := $(shell go version)
+VERSION_DATE       := $(shell date -Ru)
+VERSION_TAG        := $(shell git describe --tags --always)
+COMMON_BUILD_FLAGS := -ldflags="-s -w -X 'main.version=$(VERSION_TAG) ($(VERSION_GO)) [$(VERSION_DATE)]'"
+
+MOD_ON  := env GO111MODULE=on
+MOD_OFF := env GO111MODULE=auto
 
 # -----------------------------------------------------------------------------
 
-$(APP_NAME): $(APP_DEPS)
+$(APP_NAME):
-	@go build $(COMMON_BUILD_FLAGS) -o "$(APP_NAME)"
+	@$(MOD_ON) go build $(COMMON_BUILD_FLAGS) -o "$(APP_NAME)"
 
-static-$(APP_NAME): $(APP_DEPS)
+static-$(APP_NAME):
-	@env CGO_ENABLED=0 GOOS=linux go build -a -installsuffix cgo $(COMMON_BUILD_FLAGS) -o "$(APP_NAME)"
+	@$(MOD_ON) env CGO_ENABLED=0 GOOS=linux go build -a -installsuffix cgo $(COMMON_BUILD_FLAGS) -o "$(APP_NAME)"
 
 $(APP_NAME)-%: GOOS=$(shell echo -n "$@" | sed 's?$(APP_NAME)-??' | cut -f1 -d-)
 $(APP_NAME)-%: GOARCH=$(shell echo -n "$@" | sed 's?$(APP_NAME)-??' | cut -f2 -d-)
-$(APP_NAME)-%: $(APP_DEPS) ccbuilds
+$(APP_NAME)-%: ccbuilds
-	@env "GOOS=$(GOOS)" "GOARCH=$(GOARCH)" \
+	@$(MOD_ON) env "GOOS=$(GOOS)" "GOARCH=$(GOARCH)" \
 		go build \
 		$(COMMON_BUILD_FLAGS) \
 		-o "./ccbuilds/$(APP_NAME)-$(GOOS)-$(GOARCH)"
@@ -28,11 +33,8 @@ $(APP_NAME)-%: $(APP_DEPS) ccbuilds
 ccbuilds:
 	@rm -rf ./ccbuilds && mkdir -p ./ccbuilds
 
-version.go:
+vendor: go.mod go.sum
-	@go generate main.go
+	@$(MOD_ON) go mod vendor
-
-vendor: Gopkg.lock Gopkg.toml install-cli
-	@dep ensure
 
 # -----------------------------------------------------------------------------
 
@@ -50,12 +52,12 @@ crosscompile-dir:
 	@rm -rf "$(CC_DIR)" && mkdir -p "$(CC_DIR)"
 
 .PHONY: test
-test: vendor install-cli version.go
+test: vendor
-	@go test -v ./...
+	@$(MOD_ON) go test -v ./...
 
 .PHONY: lint
-lint: vendor install-cli version.go
+lint: vendor
-	@$(GOMETALINTER) --deadline=2m ./...
+	@$(MOD_OFF) golangci-lint run
 
 .PHONY: clean
 clean:
@@ -67,14 +69,10 @@ clean:
 docker:
 	@docker build --pull -t "$(IMAGE_NAME)" "$(ROOT_DIR)"
 
-.PHONY: install-cli
+.PHONY: prepare
-install-cli: install-dep install-lint
+prepare: install-lint
-
-.PHONY: install-dep
-install-dep:
-	@go get 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
+		| $(MOD_OFF) bash -s -- -b $(GOPATH)/bin $(GOLANGCI_LINT_VERSION)

README.md

@@ -3,9 +3,9 @@
 Bullshit-free MTPROTO proxy for Telegram
 
 [![Build Status](https://travis-ci.org/9seconds/mtg.svg?branch=master)](https://travis-ci.org/9seconds/mtg)
+[![Go Report Card](https://goreportcard.com/badge/github.com/9seconds/mtg)](https://goreportcard.com/report/github.com/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
@@ -14,7 +14,6 @@ are the most notable:
 * [Official](https://github.com/TelegramMessenger/MTProxy)
 * [Python](https://github.com/alexbers/mtprotoproxy)
 * [Erlang](https://github.com/seriyps/mtproto_proxy)
-* [JS](https://github.com/FreedomPrevails/JSMTProxy)
 
 Almost all of them follow the way how official proxy was build. This
 includes support of multiple secrets, support of promoted channels etc.
@@ -34,7 +33,7 @@ mtg is an implementation in golang which is intended to be:
   software. I also believe that in case of throwout proxies, this feature
   is useless luxury.
 * **Minimum docker image size**
-  Official image is less than 2 megabytes. Literally.
+  Official image is less than 3 megabytes. Literally.
 * **No management WebUI**
   This is an implementation of simple lightweight proxy. I won't do that.
 
@@ -51,30 +50,56 @@ To get promoted channel, please contact
 second parameter.
 
 
+# Source code organization
+
+There are 2 main branches:
+
+1. `master` branch contains potentially unstable features
+2. `stable` branch contains stable version. Usually you want to use this branch.
+
 # 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
 
+Docker follows the same policy as the source code organization:
+
+- `latest` mirrors the master branch
+- `stable` mirrors the stable branch
+- tags are for tagged releases
+
 ```console
-$ docker pull nineseconds/mtg
+docker pull nineseconds/mtg:latest
 ```
 
+```console
+docker pull nineseconds/mtg:stable
+```
+
+```console
+docker pull nineseconds/mtg:0.10
+```
+
+# Ansible role
+
+You can find unofficial Ansible role for mtg here: https://github.com/rlex/ansible-role-mtg
+Also, there is another project on Ansible Galaxy: https://galaxy.ansible.com/ivansible/lin_mtproxy
+
 # Configuration
 
 Basically, to run this tool you need to configure as less as possible.
@@ -82,27 +107,157 @@ 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
+
+_tl;dr - use secret mode for all new installation of proxy; only clients
+with dd-secrets will be able to connect. This mode abuses attempts to
+DPI MTPROTO traffic._
+
+Secure mode is not the best name and of course, it creates a lot of
+confusion. To explain what it means, we need to tell you some bits on
+dd-secrets.
+
+MTPROTO proxy protocol requires 16-byte secret. You usually
+propagate it as a 32 characters hexadecimal string like
+`282831900f371ca182feb0e4e1e1aeef` (if you decode this string
+to bytes, you will get a real secret which is used in the
+protocol). Everything went quite good until the moment when
+developers found an evidence that [protocol is quite weak to
+DPI](https://github.com/TelegramMessenger/MTProxy/issues/35) and some
+enthusiasts even created simple proofs of concepts on [detecting MTPROTO
+traffic](https://github.com/darkk/poormansmtproto).
+
+Telegram team has introduced a patch called dd-secrets. If you have
+a secret `282831900f371ca182feb0e4e1e1aeef` then your dd-secret is
+`dd282831900f371ca182feb0e4e1e1aeef`. That is, you just add dd prefix
+to the secret, prepend it with dd. In that case, original secret
+`282831900f371ca182feb0e4e1e1aeef` is used but client and server start
+to act a little bit different: they start to add random noise to the
+packets so they can't be detected by their length. In order to keep
+backward compatibility, all proxies a quite liberal to the secrets to
+use: if the client uses plain secret, without dd prefix, they fall back
+to the normal behavior. If dd-secret is used (proxy can extract this
+information on the handshake), then more secured, the hardened behavior
+is used.
+
+Yes, it can look like a hack but it is as it is.
+
+Now going back to the secure mode: if you do not pass `-s` flag to the
+mtg, then it checks what mode is requested by the client. If the client
+uses plain secret, without dd prefix, then proxy falls back to the
+original behavior and do not play with paddings. If dd-secret is used
+and client demands this mode, then proxy start to add that random noise
+to the packets. But if you pass `-s`, then only clients with dd-secrets
+can connect. How to migrate existing clients then? If a client is new
+enough, you can just prepend the secret with dd string in the settings.
+If it is an old guy, then nothing to do, sorry.
+
+Why this mode matters? We do not have evidence but there is quite a big
+suspicion that some ISPs start to filter MTPROTO traffic. If they detect
+the IP address which acts as a proxy, they block it and no clients can
+use this proxy. This is an attempt to prevent such a situation.
+
+General rule of thumb: with all new installation of proxies I would
+advise to go with secure mode by default. But please do remember that it
+means that clients, which do not pass dd-prefix to their secrets, will
+not be able to connect. *Secure mode works only with dd-prefixes!*
+
+Oneliners to generate such secrets:
+
+```console
+echo dd$(openssl rand -hex 16)
+```
+
+or
+
+```console
+echo dd$(head -c 512 /dev/urandom | md5sum | cut -f 1 -d ' ')
+```
+
+
+## Antireplay cache
+
+In order to prevent replay attacks, we have internal storage of first
+frames messages for connected clients. These frames are generated
+randomly by design and we have negligible possibility of duplication
+(probability is 1/(2^64)) but it could be quite effective in order to
+prevent replays.
+
+
+## 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_PROMETHEUS_PREFIX`       | `--prometheus-prefix`       | `mtg`                             | Which namespace should be used for prometheus metrics.                                                                                                                                                                                                                     |
+| `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.                                                                                                                                                                      |
+| `MTG_SECURE_ONLY`             | `-s`, `--secure-only`       | `false`                           | Support only clients with secure mode (i.e only clients with dd-secrets).                                                                                                                                                                                                  |
+| `MTG_ANTIREPLAY_MAXSIZE`      | `anti-replay-max-size`      | `128`                             | Max size of antireplay cache in megabytes.                                                                                                                                                                                                                                 |
+| `MTG_ANTIREPLAY_EVICTIONTIME` | `anti-replay-eviction-time` | `168h`                            | Eviction time for antireplay cache entries.                                                                                                                                                                                                                                |
+
+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>
 ```
 
 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 $(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)
 ```
 
 You will have this tool up and running on port 3128. Now curl
@@ -110,3 +265,44 @@ You will have this tool up and running on port 3128. Now curl
 port 3129 will show you some statistics if you are interested in.
 
 Also, you can use [run-mtg.sh](https://github.com/9seconds/mtg/blob/master/run-mtg.sh) script
+
+
+# statsd integration
+
+mtg provides an integration with statsd, you can enable it with command
+line interface. To enable it, you have to provide IP address of statsd
+service.
+
+Out of the box, mtg supports 2 additional dialects: [InfluxDB](https://www.influxdata.com/blog/getting-started-with-sending-statsd-metrics-to-telegraf-influxdb/)
+and [Datadog](https://docs.datadoghq.com/developers/dogstatsd/).
+
+All metrics are gauges. Here is the list of metrics and their meaning:
+
+| Metric name                     | Unit    | Description                                               |
+|---------------------------------|---------|-----------------------------------------------------------|
+| `connections.abridged.ipv4`     | number  | The number of active abridged IPv4 connections            |
+| `connections.abridged.ipv6`     | number  | The number of active abridged IPv6 connections            |
+| `connections.intermediate.ipv4` | number  | The number of active intermediate IPv4 connections        |
+| `connections.intermediate.ipv6` | number  | The number of active intermediate IPv6 connections        |
+| `connections.secure.ipv4`       | number  | The number of active secure intermediate IPv4 connections |
+| `connections.secure.ipv6`       | number  | The number of active secure intermediate IPv6 connections |
+| `crashes`                       | number  | An amount of crashes in client handlers                   |
+| `traffic.ingress`               | bytes   | Ingress traffic from the start of application (incoming)  |
+| `traffic.egress`                | bytes   | Egress traffic from the start of application (outgoing)   |
+| `speed.ingress`                 | bytes/s | Ingress bandwidth of the latest second (incoming traffic) |
+| `speed.egress`                  | bytes/s | Egress bandwidth of the latest second (outgoing traffic)  |
+
+All metrics are prefixed with given prefix. Default prefix is `mtg`.
+With such prefix metric name `traffic.ingress`, for example, would be
+`mtg.traffic.ingress`.
+
+
+# Prometheus integration
+
+[Prometheus](https://prometheus.io) integration comes out of
+the box, you do not need to setup anything special. Prometheus
+scrape endpoint lives on the same IP/port where generic stats
+service (`http://${MTG_STATS_IP}:${MTG_STATS_PORT}`) but on
+`/prometheus` path. So, if you access http stats service as `curl
+http://localhost:3129/`, then your prometheus endpoint is `curl
+http://localhost:3129/prometheus/`.

antireplay/cache.go

@@ -0,0 +1,37 @@
+package antireplay
+
+import (
+	"github.com/allegro/bigcache"
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/config"
+)
+
+// Cache defines storage for obfuscated2 handshake frames.
+type Cache struct {
+	cache *bigcache.BigCache
+}
+
+func (a Cache) Add(frame []byte) {
+	a.cache.Set(string(frame), nil) // nolint: errcheck
+}
+
+func (a Cache) Has(frame []byte) bool {
+	_, err := a.cache.Get(string(frame))
+
+	return err == nil
+}
+
+func NewCache(config *config.Config) (Cache, error) {
+	cache, err := bigcache.NewBigCache(bigcache.Config{
+		Shards:           1024,
+		LifeWindow:       config.AntiReplayEvictionTime,
+		Hasher:           hasher{},
+		HardMaxCacheSize: config.AntiReplayMaxSize,
+	})
+	if err != nil {
+		return Cache{}, errors.Annotate(err, "Cannot make cache")
+	}
+
+	return Cache{cache}, nil
+}

antireplay/hasher.go

@@ -0,0 +1,9 @@
+package antireplay
+
+import "github.com/cespare/xxhash"
+
+type hasher struct{}
+
+func (h hasher) Sum64(value string) uint64 {
+	return xxhash.Sum64String(value)
+}

client/client.go

@@ -1,12 +1,15 @@
 package client
 
 import (
+	"context"
 	"net"
 
+	"github.com/9seconds/mtg/antireplay"
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
 	"github.com/9seconds/mtg/wrappers"
 )
 
 // 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,
+	antireplay.Cache, *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error)

client/direct.go

@@ -1,49 +1,55 @@
 package client
 
 import (
+	"context"
 	"net"
 	"time"
 
 	"github.com/juju/errors"
 
+	"github.com/9seconds/mtg/antireplay"
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
 	"github.com/9seconds/mtg/obfuscated2"
 	"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, antiReplayCache antireplay.Cache,
+	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")
 	}
 
-	socket.SetReadDeadline(time.Now().Add(handshakeTimeout)) // nolint: errcheck
+	socket.SetReadDeadline(time.Now().Add(handshakeTimeout)) // nolint: errcheck, gosec
 	frame, err := obfuscated2.ExtractFrame(socket)
 	if err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot extract frame")
 	}
-	socket.SetReadDeadline(time.Time{}) // nolint: errcheck
+	socket.SetReadDeadline(time.Time{}) // nolint: errcheck, gosec
 
-	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")
 	}
+
+	if antiReplayCache.Has([]byte(frame)) {
+		return nil, nil, errors.New("Replay attack is detected")
+	}
+	antiReplayCache.Add([]byte(frame))
+
 	connOpts.ConnectionProto = mtproto.ConnectionProtocolAny
 	connOpts.ClientAddr = conn.RemoteAddr()
 

client/middle.go

@@ -1,8 +1,10 @@
 package client
 
 import (
+	"context"
 	"net"
 
+	"github.com/9seconds/mtg/antireplay"
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
 	"github.com/9seconds/mtg/wrappers"
@@ -10,16 +12,25 @@ 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, antiReplayCache antireplay.Cache,
+	conf *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) {
+	conn, opts, err := DirectInit(ctx, cancel, socket, connID, antiReplayCache, conf)
 	if err != nil {
 		return nil, nil, err
 	}
 	connStream := conn.(wrappers.StreamReadWriteCloser)
 
-	newConn := wrappers.NewMTProtoAbridged(connStream, opts)
+	var newConn wrappers.PacketReadWriteCloser
-	if opts.ConnectionType != mtproto.ConnectionTypeAbridged {
+	switch opts.ConnectionType {
+	case mtproto.ConnectionTypeAbridged:
+		newConn = wrappers.NewMTProtoAbridged(connStream, opts)
+	case mtproto.ConnectionTypeIntermediate:
 		newConn = wrappers.NewMTProtoIntermediate(connStream, opts)
+	case mtproto.ConnectionTypeSecure:
+		newConn = wrappers.NewMTProtoIntermediateSecure(connStream, opts)
+	default:
+		panic("Unknown connection type")
 	}
 
 	opts.ConnectionProto = mtproto.ConnectionProtocolIPv4

config/config.go

@@ -1,24 +1,26 @@
 package config
 
 import (
+	"bytes"
 	"encoding/hex"
 	"fmt"
 	"net"
 	"strconv"
+	"time"
 
 	"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
+	SecureOnly bool
+
+	ReadBufferSize  int
+	WriteBufferSize int
 
 	BindPort       uint16
 	PublicIPv4Port uint16
@@ -30,6 +32,20 @@ type Config struct {
 	PublicIPv6 net.IP
 	StatsIP    net.IP
 
+	AntiReplayMaxSize      int
+	AntiReplayEvictionTime time.Duration
+
+	StatsD struct {
+		Addr       net.Addr
+		Prefix     string
+		Tags       map[string]string
+		TagsFormat statsd.TagFormat
+		Enabled    bool
+	}
+	Prometheus struct {
+		Prefix string
+	}
+
 	Secret []byte
 	AdTag  []byte
 }
@@ -46,6 +62,7 @@ type URLs struct {
 type IPURLs struct {
 	IPv4      URLs   `json:"ipv4"`
 	IPv6      URLs   `json:"ipv6"`
+	BotSecret string `json:"secret_for_mtproxybot"`
 }
 
 // BindAddr returns connection for this server to bind to.
@@ -64,15 +81,32 @@ func (c *Config) UseMiddleProxy() bool {
 	return len(c.AdTag) > 0
 }
 
+// BotSecretString returns secret string which should work with MTProxybot.
+func (c *Config) BotSecretString() string {
+	return hex.EncodeToString(c.Secret)
+}
+
+// SecretString returns a secret in a form entered on the start of the
+// application.
+func (c *Config) SecretString() string {
+	secret := c.BotSecretString()
+	if c.SecureMode {
+		return "dd" + secret
+	}
+	return secret
+}
+
 // GetURLs returns configured IPURLs instance with links to this server.
 func (c *Config) GetURLs() IPURLs {
 	urls := IPURLs{}
+	secret := c.SecretString()
 	if c.PublicIPv4 != nil {
-		urls.IPv4 = getURLs(c.PublicIPv4, c.PublicIPv4Port, c.Secret)
+		urls.IPv4 = getURLs(c.PublicIPv4, c.PublicIPv4Port, secret)
 	}
 	if c.PublicIPv6 != nil {
-		urls.IPv6 = getURLs(c.PublicIPv6, c.PublicIPv6Port, c.Secret)
+		urls.IPv6 = getURLs(c.PublicIPv6, c.PublicIPv6Port, secret)
 	}
+	urls.BotSecret = c.BotSecretString()
 
 	return urls
 }
@@ -85,27 +119,23 @@ 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) (*Config, error) {
+	statsdTags map[string]string, prometheusPrefix string,
-	if len(secret) != 32 {
+	secureOnly bool,
+	antiReplayMaxSize int, antiReplayEvictionTime time.Duration,
+	secret, adtag []byte) (*Config, error) {
+	secureMode := secureOnly
+	if bytes.HasPrefix(secret, []byte{0xdd}) && len(secret) == 17 {
+		secureMode = true
+		secret = bytes.TrimPrefix(secret, []byte{0xdd})
+	} 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 {
@@ -114,8 +144,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 {
@@ -137,16 +167,57 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo
 	conf := &Config{
 		Debug:                  debug,
 		Verbose:                verbose,
+		SecureOnly:             secureOnly,
 		BindIP:                 bindIP,
 		BindPort:               bindPort,
 		PublicIPv4:             publicIPv4,
-		PublicIPv4Port: PublicIPv4Port,
+		PublicIPv4Port:         publicIPv4Port,
 		PublicIPv6:             publicIPv6,
 		PublicIPv6Port:         publicIPv6Port,
 		StatsIP:                statsIP,
 		StatsPort:              statsPort,
-		Secret:         secretBytes,
+		Secret:                 secret,
-		AdTag:          adTagBytes,
+		AdTag:                  adtag,
+		SecureMode:             secureMode,
+		ReadBufferSize:         int(readBufferSize),
+		WriteBufferSize:        int(writeBufferSize),
+		AntiReplayMaxSize:      antiReplayMaxSize,
+		AntiReplayEvictionTime: antiReplayEvictionTime,
+	}
+	conf.Prometheus.Prefix = prometheusPrefix
+
+	if statsdIP != "" {
+		conf.StatsD.Enabled = true
+		conf.StatsD.Prefix = statsdPrefix
+		conf.StatsD.Tags = statsdTags
+
+		var (
+			addr net.Addr
+			err  error
+		)
+		hostPort := net.JoinHostPort(statsdIP, strconv.Itoa(int(statsdPort)))
+		switch statsdNetwork {
+		case "tcp":
+			addr, err = net.ResolveTCPAddr("tcp", hostPort)
+		case "udp":
+			addr, err = net.ResolveUDPAddr("udp", hostPort)
+		default:
+			err = errors.Errorf("Unknown network %s", statsdNetwork)
+		}
+		if err != nil {
+			return nil, errors.Annotate(err, "Cannot resolve statsd address")
+		}
+		conf.StatsD.Addr = addr
+
+		switch statsdTagsFormat {
+		case "datadog":
+			conf.StatsD.TagsFormat = statsd.Datadog
+		case "influxdb":
+			conf.StatsD.TagsFormat = statsd.InfluxDB
+		case "":
+		default:
+			return nil, errors.Errorf("Unknown tags format %s", statsdTagsFormat)
+		}
 	}
 
 	return conf, nil

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{FallbackDelay: -1}
+	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

@@ -1,17 +1,16 @@
 package config
 
 import (
-	"encoding/hex"
 	"net"
 	"net/url"
 	"strconv"
 )
 
-func getURLs(addr net.IP, port uint16, secret []byte) (urls URLs) {
+func getURLs(addr net.IP, port uint16, secret string) (urls URLs) {
 	values := url.Values{}
 	values.Set("server", addr.String())
 	values.Set("port", strconv.Itoa(int(port)))
-	values.Set("secret", hex.EncodeToString(secret))
+	values.Set("secret", secret)
 
 	urls.TG = makeTGURL(values)
 	urls.TMe = makeTMeURL(values)
@@ -43,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",
@@ -53,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()
 }

go.mod

@@ -0,0 +1,30 @@
+module github.com/9seconds/mtg
+
+replace github.com/golang/lint => github.com/golang/lint v0.0.0-20190227174305-8f45f776aaf1
+
+require (
+	github.com/OneOfOne/xxhash v1.2.5 // indirect
+	github.com/allegro/bigcache v1.2.0
+	github.com/beevik/ntp v0.2.0
+	github.com/cespare/xxhash v1.1.0
+	github.com/dustin/go-humanize v1.0.0
+	github.com/gofrs/uuid v3.2.0+incompatible
+	github.com/juju/errors v0.0.0-20190207033735-e65537c515d7
+	github.com/juju/loggo v0.0.0-20190526231331-6e530bcce5d8 // indirect
+	github.com/juju/testing v0.0.0-20190429233213-dfc56b8c09fc // indirect
+	github.com/kr/pretty v0.1.0 // indirect
+	github.com/pkg/errors v0.8.1 // indirect
+	github.com/prometheus/client_golang v0.9.4
+	github.com/spaolacci/murmur3 v1.1.0 // indirect
+	github.com/stretchr/testify v1.3.0
+	go.uber.org/atomic v1.4.0 // indirect
+	go.uber.org/multierr v1.1.0 // indirect
+	go.uber.org/zap v1.10.0
+	golang.org/x/net v0.0.0-20190607181551-461777fb6f67 // indirect
+	golang.org/x/sys v0.0.0-20190610200419-93c9922d18ae // indirect
+	gopkg.in/alecthomas/kingpin.v2 v2.2.6
+	gopkg.in/alexcesaro/statsd.v2 v2.0.0
+	gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect
+	gopkg.in/mgo.v2 v2.0.0-20180705113604-9856a29383ce // indirect
+	gopkg.in/yaml.v2 v2.2.2 // indirect
+)

go.sum

@@ -0,0 +1,115 @@
+github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAElWljhcU=
+github.com/OneOfOne/xxhash v1.2.5 h1:zl/OfRA6nftbBK9qTohYBJ5xvw6C/oNKizR7cZGl3cI=
+github.com/OneOfOne/xxhash v1.2.5/go.mod h1:eZbhyaAYD41SGSSsnmcpxVoRiQ/MPUTjUdIIOT9Um7Q=
+github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc h1:cAKDfWh5VpdgMhJosfJnn5/FoN2SRZ4p7fJNX58YPaU=
+github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
+github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf h1:qet1QNfXsQxTZqLG4oE62mJzwPIB8+Tee4RNCL9ulrY=
+github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
+github.com/allegro/bigcache v1.2.0 h1:qDaE0QoF29wKBb3+pXFrJFy1ihe5OT9OiXhg1t85SxM=
+github.com/allegro/bigcache v1.2.0/go.mod h1:Cb/ax3seSYIx7SuZdm2G2xzfwmv3TPSk2ucNfQESPXM=
+github.com/beevik/ntp v0.2.0 h1:sGsd+kAXzT0bfVfzJfce04g+dSRfrs+tbQW8lweuYgw=
+github.com/beevik/ntp v0.2.0/go.mod h1:hIHWr+l3+/clUnF44zdK+CWW7fO8dR5cIylAQ76NRpg=
+github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973 h1:xJ4a3vCFaGF/jqvzLMYoU8P317H5OQ+Via4RmuPwCS0=
+github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24TaqPxmxbtue+5NUziq4I4S80YR8gNf3Q=
+github.com/beorn7/perks v1.0.0 h1:HWo1m869IqiPhD389kmkxeTalrjNbbJTC8LXupb+sl0=
+github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8=
+github.com/cespare/xxhash v1.1.0 h1:a6HrQnmkObjyL+Gs60czilIUGqrzKutQD6XZog3p+ko=
+github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo=
+github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
+github.com/go-kit/kit v0.8.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as=
+github.com/go-logfmt/logfmt v0.3.0/go.mod h1:Qt1PoO58o5twSAckw1HlFXLmHsOX5/0LbT9GBnD5lWE=
+github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
+github.com/gofrs/uuid v3.2.0+incompatible h1:y12jRkkFxsd7GpqdSZ+/KCs/fJbqpEXSGd4+jfEaewE=
+github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
+github.com/gogo/protobuf v1.1.1 h1:72R+M5VuhED/KujmZVcIquuo8mBgX4oVda//DQb3PXo=
+github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
+github.com/golang/protobuf v1.2.0 h1:P3YflyNX/ehuJFLhxviNdFxQPkGK5cDcApsge1SqnvM=
+github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
+github.com/golang/protobuf v1.3.1 h1:YF8+flBXS5eO826T4nzqPrxfhQThhXl0YzfuUPu4SBg=
+github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
+github.com/json-iterator/go v1.1.6/go.mod h1:+SdeFBvtyEkXs7REEP0seUULqWtbJapLOCVDaaPEHmU=
+github.com/juju/errors v0.0.0-20190207033735-e65537c515d7 h1:dMIPRDg6gi7CUp0Kj2+HxqJ5kTr1iAdzsXYIrLCNSmU=
+github.com/juju/errors v0.0.0-20190207033735-e65537c515d7/go.mod h1:W54LbzXuIE0boCoNJfwqpmkKJ1O4TCTZMetAt6jGk7Q=
+github.com/juju/loggo v0.0.0-20190526231331-6e530bcce5d8 h1:UUHMLvzt/31azWTN/ifGWef4WUqvXk0iRqdhdy/2uzI=
+github.com/juju/loggo v0.0.0-20190526231331-6e530bcce5d8/go.mod h1:vgyd7OREkbtVEN/8IXZe5Ooef3LQePvuBm9UWj6ZL8U=
+github.com/juju/testing v0.0.0-20190429233213-dfc56b8c09fc h1:5xUWujf6ES9tEpFHFzI34vcHm8U07lGjxAuJML3qwqM=
+github.com/juju/testing v0.0.0-20190429233213-dfc56b8c09fc/go.mod h1:63prj8cnj0tU0S9OHjGJn+b1h0ZghCndfnbQolrYTwA=
+github.com/julienschmidt/httprouter v1.2.0/go.mod h1:SYymIcj16QtmaHHD7aYtjjsJG7VTCxuUUipMqKk8s4w=
+github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
+github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc=
+github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
+github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
+github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
+github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
+github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
+github.com/matttproud/golang_protobuf_extensions v1.0.1 h1:4hp9jkHxhMHkqkrB3Ix0jegS5sx/RkqARlsWZ6pIwiU=
+github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
+github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
+github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
+github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
+github.com/pkg/errors v0.8.0 h1:WdK/asTD0HN+q6hsWO3/vpuAkAr+tw6aNJNDFFf0+qw=
+github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I=
+github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/prometheus/client_golang v0.9.1/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw=
+github.com/prometheus/client_golang v0.9.4 h1:Y8E/JaaPbmFSW2V81Ab/d8yZFYQQGbni1b1jPcG9Y6A=
+github.com/prometheus/client_golang v0.9.4/go.mod h1:oCXIBxdI62A4cR6aTRJCgetEjecSIYzOEaeAn4iYEpM=
+github.com/prometheus/client_model v0.0.0-20180712105110-5c3871d89910 h1:idejC8f05m9MGOsuEi1ATq9shN03HrxNkD/luQvxCv8=
+github.com/prometheus/client_model v0.0.0-20180712105110-5c3871d89910/go.mod h1:MbSGuTsp3dbXC40dX6PRTWyKYBIrTGTE9sqQNg2J8bo=
+github.com/prometheus/client_model v0.0.0-20190129233127-fd36f4220a90 h1:S/YWwWx/RA8rT8tKFRuGUZhuA90OyIBpPCXkcbwU8DE=
+github.com/prometheus/client_model v0.0.0-20190129233127-fd36f4220a90/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
+github.com/prometheus/common v0.4.1 h1:K0MGApIoQvMw27RTdJkPbr3JZ7DNbtxQNyi5STVM6Kw=
+github.com/prometheus/common v0.4.1/go.mod h1:TNfzLD0ON7rHzMJeJkieUDPYmFC7Snx/y86RQel1bk4=
+github.com/prometheus/procfs v0.0.0-20181005140218-185b4288413d h1:GoAlyOgbOEIFdaDqxJVlbOQ1DtGmZWs/Qau0hIlk+WQ=
+github.com/prometheus/procfs v0.0.0-20181005140218-185b4288413d/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk=
+github.com/prometheus/procfs v0.0.2 h1:6LJUbpNm42llc4HRCuvApCSWB/WfhuNo9K98Q9sNGfs=
+github.com/prometheus/procfs v0.0.2/go.mod h1:TjEm7ze935MbeOT/UhFTIMYKhuLP4wbCsTZCD3I8kEA=
+github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=
+github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=
+github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI=
+github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w=
+github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
+github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q=
+github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
+go.uber.org/atomic v1.4.0 h1:cxzIVoETapQEqDhQu3QfnvXAV4AlzcvUCxkVUFw3+EU=
+go.uber.org/atomic v1.4.0/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE=
+go.uber.org/multierr v1.1.0 h1:HoEmRHQPVSqub6w2z2d2EOVs2fjyFRGyofhKuyDq0QI=
+go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
+go.uber.org/zap v1.10.0 h1:ORx85nbTijNz8ljznvCMR1ZBIPKFn3jQrag10X2AsuM=
+go.uber.org/zap v1.10.0/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q=
+golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
+golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
+golang.org/x/net v0.0.0-20181114220301-adae6a3d119a/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
+golang.org/x/net v0.0.0-20190607181551-461777fb6f67 h1:rJJxsykSlULwd2P2+pg/rtnwN2FrWp4IuCxOSyS0V00=
+golang.org/x/net v0.0.0-20190607181551-461777fb6f67/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
+golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4 h1:YUO/7uOKsKeq9UokNS62b8FYywz3ker1l1vDZRCRefw=
+golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20181116152217-5ac8a444bdc5/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20190610200419-93c9922d18ae h1:xiXzMMEQdQcric9hXtr1QU98MHunKK7OTtsoU6bYWs4=
+golang.org/x/sys v0.0.0-20190610200419-93c9922d18ae/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
+gopkg.in/alecthomas/kingpin.v2 v2.2.6 h1:jMFz6MfLP0/4fUyZle81rXUoxOBFi19VUFKVDOQfozc=
+gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
+gopkg.in/alexcesaro/statsd.v2 v2.0.0 h1:FXkZSCZIH17vLCO5sO2UucTHsH9pc+17F6pl3JVCwMc=
+gopkg.in/alexcesaro/statsd.v2 v2.0.0/go.mod h1:i0ubccKGzBVNBpdGV5MocxyA/XlLUJzA7SLonnE4drU=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
+gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/mgo.v2 v2.0.0-20180705113604-9856a29383ce h1:xcEWjVhvbDy+nHP67nPDDpbYrY+ILlfndk4bRioVHaU=
+gopkg.in/mgo.v2 v2.0.0-20180705113604-9856a29383ce/go.mod h1:yeKp02qBN3iKW1OzL3MGk2IdtZzaj7SFntXj72NppTA=
+gopkg.in/yaml.v2 v2.2.1 h1:mUhvW9EsL+naU5Q3cakzfE91YhliOondGd6ZrsDBHQE=
+gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
+gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=

main.go

@@ -1,86 +1,159 @@
 package main
 
-//go:generate scripts/generate_version.sh
-
 import (
 	"encoding/json"
+	"fmt"
 	"io"
 	"math/rand"
 	"os"
 	"syscall"
 	"time"
 
+	"github.com/juju/errors"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 	kingpin "gopkg.in/alecthomas/kingpin.v2"
 
 	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/ntp"
 	"github.com/9seconds/mtg/proxy"
 	"github.com/9seconds/mtg/stats"
-	"github.com/juju/errors"
 )
 
+var version = "dev" // this has to be set by build ld flags
+
 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",
+		"Which port to bind to.").
 		Short('p').
 		Envar("MTG_PORT").
 		Default("3128").
 		Uint16()
 
-	publicIPv4 = app.Flag("public-ipv4", "Which IPv4 address is public.").
+	publicIPv4 = app.Flag("public-ipv4",
+		"Which IPv4 address is public.").
 		Short('4').
 		Envar("MTG_IPV4").
 		IP()
-	publicIPv4Port = app.Flag("public-ipv4-port", "Which IPv4 port is public. Default is 'bind-port' value.").
+	publicIPv4Port = app.Flag("public-ipv4-port",
+		"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",
+		"Which IPv6 address is public.").
 		Short('6').
 		Envar("MTG_IPV6").
 		IP()
-	publicIPv6Port = app.Flag("public-ipv6-port", "Which IPv6 port is public. Default is 'bind-port' value.").
+	publicIPv6Port = app.Flag("public-ipv6-port",
+		"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",
+		"Which port bind stats to.").
 		Short('q').
 		Envar("MTG_STATS_PORT").
 		Default("3129").
 		Uint16()
 
-	secret = app.Arg("secret", "Secret of this proxy.").Required().String()
+	statsdIP = app.Flag("statsd-ip",
-	adtag  = app.Arg("adtag", "ADTag of the proxy.").String()
+		"Which IP should we use for working with statsd.").
+		Envar("MTG_STATSD_IP").
+		String()
+	statsdPort = app.Flag("statsd-port",
+		"Which port should we use for working with statsd.").
+		Envar("MTG_STATSD_PORT").
+		Default("8125").
+		Uint16()
+	statsdNetwork = app.Flag("statsd-network",
+		"Which network is used to work with statsd. Only 'tcp' and 'udp' are supported.").
+		Envar("MTG_STATSD_NETWORK").
+		Default("udp").
+		String()
+	statsdPrefix = app.Flag("statsd-prefix",
+		"Which bucket prefix should we use for sending stats to statsd.").
+		Envar("MTG_STATSD_PREFIX").
+		Default("mtg").
+		String()
+	statsdTagsFormat = app.Flag("statsd-tags-format",
+		"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()
+
+	prometheusPrefix = app.Flag("prometheus-prefix",
+		"Which namespace to use to send stats to Prometheus.").
+		Envar("MTG_PROMETHEUS_PREFIX").
+		Default("mtg").
+		String()
+
+	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()
+	secureOnly = app.Flag("secure-only",
+		"Support clients with dd-secrets only.").
+		Short('s').
+		Envar("MTG_SECURE_ONLY").
+		Bool()
+
+	antiReplayMaxSize = app.Flag("anti-replay-max-size",
+		"Max size of antireplay cache in megabytes.").
+		Envar("MTG_ANTIREPLAY_MAXSIZE").
+		Default("128").
+		Int()
+	antiReplayEvictionTime = app.Flag("anti-replay-eviction-time",
+		"Eviction time period for obfuscated2 handshakes").
+		Envar("MTG_ANTIREPLAY_EVICTIONTIME").
+		Default("168h").
+		Duration()
+
+	secret = app.Arg("secret", "Secret of this proxy.").Required().HexBytes()
+	adtag  = app.Arg("adtag", "ADTag of the proxy.").HexBytes()
 )
 
-func init() {
+func main() { // nolint: gocyclo
 	rand.Seed(time.Now().UTC().UnixNano())
 	app.Version(version)
+	app.HelpFlag.Short('h')
 
-}
-
-func main() {
 	kingpin.MustParse(app.Parse(os.Args[1:]))
 
 	err := setRLimit()
@@ -89,10 +162,12 @@ 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, *prometheusPrefix, *secureOnly,
+		*antiReplayMaxSize, *antiReplayEvictionTime,
 		*secret, *adtag,
 	)
 	if err != nil {
@@ -100,11 +175,12 @@ func main() {
 	}
 
 	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()
@@ -117,16 +193,31 @@ func main() {
 	defer logger.Sync() // nolint: errcheck
 
 	printURLs(conf.GetURLs())
+	zap.S().Debugw("Configuration", "config", conf)
 
 	if conf.UseMiddleProxy() {
 		zap.S().Infow("Use middle proxy connection to Telegram")
+		if diff, err := ntp.Fetch(); err != nil {
+			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))
+			}
+			go ntp.AutoUpdate()
+		}
 	} else {
 		zap.S().Infow("Use direct connection to Telegram")
 	}
 
-	go stats.Start(conf)
+	if err := stats.Init(conf); err != nil {
+		panic(err)
+	}
 
-	server := proxy.NewProxy(conf)
+	server, err := proxy.NewProxy(conf)
+	if err != nil {
+		panic(err)
+	}
 	if err := server.Serve(); err != nil {
 		zap.S().Fatalw("Server stopped", "error", err)
 	}
@@ -161,6 +252,6 @@ func printURLs(data interface{}) {
 }
 
 func usage(msg string) {
-	io.WriteString(os.Stderr, msg+"\n") // nolint: errcheck
+	io.WriteString(os.Stderr, msg+"\n") // nolint: errcheck, gosec
 	os.Exit(1)
 }

mtproto/connection_options.go

@@ -39,6 +39,7 @@ const (
 	ConnectionTypeUnknown ConnectionType = iota
 	ConnectionTypeAbridged
 	ConnectionTypeIntermediate
+	ConnectionTypeSecure
 )
 
 // ConnectionProtocol* define which connection protocols to use.
@@ -53,6 +54,7 @@ const (
 var (
 	ConnectionTagAbridged     = []byte{0xef, 0xef, 0xef, 0xef}
 	ConnectionTagIntermediate = []byte{0xee, 0xee, 0xee, 0xee}
+	ConnectionTagSecure       = []byte{0xdd, 0xdd, 0xdd, 0xdd}
 )
 
 // Tag maps connection type to the corresponding handshake tag.
@@ -62,6 +64,8 @@ func (t ConnectionType) Tag() ([]byte, error) {
 		return ConnectionTagAbridged, nil
 	case ConnectionTypeIntermediate:
 		return ConnectionTagIntermediate, nil
+	case ConnectionTypeSecure:
+		return ConnectionTagSecure, nil
 	default:
 		return nil, errors.Errorf("Unknown connection type %d", t)
 	}
@@ -75,6 +79,9 @@ func ConnectionTagFromHandshake(magic []byte) (ConnectionType, error) {
 	if bytes.Equal(magic, ConnectionTagAbridged) {
 		return ConnectionTypeAbridged, nil
 	}
+	if bytes.Equal(magic, ConnectionTagSecure) {
+		return ConnectionTypeSecure, nil
+	}
 
 	return ConnectionTypeUnknown, errors.New("Unknown handshake protocol")
 }

mtproto/rpc/handshake_request.go

@@ -12,10 +12,10 @@ func (r *HandshakeRequest) Bytes() []byte {
 	buf := &bytes.Buffer{}
 	buf.Grow(len(TagHandshake) + len(HandshakeFlags) + len(HandshakeSenderPID) + len(HandshakePeerPID))
 
-	buf.Write(TagHandshake)
+	buf.Write(TagHandshake)       // nolint: gosec
-	buf.Write(HandshakeFlags)
+	buf.Write(HandshakeFlags)     // nolint: gosec
-	buf.Write(HandshakeSenderPID)
+	buf.Write(HandshakeSenderPID) // nolint: gosec
-	buf.Write(HandshakePeerPID)
+	buf.Write(HandshakePeerPID)   // nolint: gosec
 
 	return buf.Bytes()
 }

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)      // nolint: gosec
-	buf.Write(r.Flags[:])
+	buf.Write(r.Flags)     // nolint: gosec
-	buf.Write(r.SenderPID[:])
+	buf.Write(r.SenderPID) // nolint: gosec
-	buf.Write(r.PeerPID[:])
+	buf.Write(r.PeerPID)   // nolint: gosec
 
 	return buf.Bytes()
 }

mtproto/rpc/nonce_request.go

@@ -21,11 +21,11 @@ type NonceRequest struct {
 func (r *NonceRequest) Bytes() []byte {
 	buf := &bytes.Buffer{}
 
-	buf.Write(TagNonce)
+	buf.Write(TagNonce)       // nolint: gosec
-	buf.Write(r.KeySelector)
+	buf.Write(r.KeySelector)  // nolint: gosec
-	buf.Write(NonceCryptoAES)
+	buf.Write(NonceCryptoAES) // nolint: gosec
-	buf.Write(r.CryptoTS)
+	buf.Write(r.CryptoTS)     // nolint: gosec
-	buf.Write(r.Nonce)
+	buf.Write(r.Nonce)        // nolint: gosec
 
 	return buf.Bytes()
 }

mtproto/rpc/nonce_response.go

@@ -18,11 +18,11 @@ type NonceResponse struct {
 func (r *NonceResponse) Bytes() []byte {
 	buf := &bytes.Buffer{}
 
-	buf.Write(r.Type)
+	buf.Write(r.Type)        // nolint: gosec
-	buf.Write(r.KeySelector)
+	buf.Write(r.KeySelector) // nolint: gosec
-	buf.Write(r.Crypto)
+	buf.Write(r.Crypto)      // nolint: gosec
-	buf.Write(r.CryptoTS)
+	buf.Write(r.CryptoTS)    // nolint: gosec
-	buf.Write(r.Nonce)
+	buf.Write(r.Nonce)       // nolint: gosec
 
 	return buf.Bytes()
 }

mtproto/rpc/proxy_flags.go

@@ -9,12 +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          proxyRequestFlags = 0x8000000
 )
 
 var proxyRequestFlagsEncryptedPrefix [8]byte
@@ -50,6 +51,9 @@ func (r proxyRequestFlags) String() string {
 	if r&proxyRequestFlagsQuickAck != 0 {
 		flags = append(flags, "QUICK_ACK")
 	}
+	if r&proxyRequestFlagsPad != 0 {
+		flags = append(flags, "PAD")
+	}
 
 	return strings.Join(flags, " | ")
 }

mtproto/rpc/proxy_request.go

@@ -49,22 +49,23 @@ func (r *ProxyRequest) MakeHeader(message []byte) (*bytes.Buffer, fmt.Stringer)
 		flags |= proxyRequestFlagsEncrypted
 	}
 
-	buf.Write(TagProxyRequest)
+	buf.Write(TagProxyRequest)                 // nolint: gosec
-	buf.Write(flags.Bytes())
+	buf.Write(flags.Bytes())                   // nolint: gosec
-	buf.Write(r.ConnectionID)
+	buf.Write(r.ConnectionID)                  // nolint: gosec
-	buf.Write(r.ClientIPPort)
+	buf.Write(r.ClientIPPort)                  // nolint: gosec
-	buf.Write(r.OurIPPort)
+	buf.Write(r.OurIPPort)                     // nolint: gosec
-	buf.Write(ProxyRequestExtraSize)
+	buf.Write(ProxyRequestExtraSize)           // nolint: gosec
-	buf.Write(ProxyRequestProxyTag)
+	buf.Write(ProxyRequestProxyTag)            // nolint: gosec
-	buf.WriteByte(byte(len(r.ADTag)))
+	buf.WriteByte(byte(len(r.ADTag)))          // nolint: gosec
-	buf.Write(r.ADTag)
+	buf.Write(r.ADTag)                         // nolint: gosec
-	buf.Write(make([]byte, (4-buf.Len()%4)%4))
+	buf.Write(make([]byte, (4-buf.Len()%4)%4)) // nolint: gosec
 
 	return buf, flags
 }
 
 // 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 {
@@ -72,6 +73,10 @@ func NewProxyRequest(clientAddr, ownAddr *net.TCPAddr, opts *mtproto.ConnectionO
 		flags |= proxyRequestFlagsAbdridged
 	case mtproto.ConnectionTypeIntermediate:
 		flags |= proxyRequestFlagsIntermediate
+	case mtproto.ConnectionTypeSecure:
+		flags |= proxyRequestFlagsIntermediate | proxyRequestFlagsPad
+	default:
+		panic("Unknown connection type")
 	}
 
 	request := &ProxyRequest{

mtproto/rpc/rpc.go

@@ -23,11 +23,6 @@ var (
 	ProxyRequestExtraSize = []byte{0x18, 0x00, 0x00, 0x00}
 	ProxyRequestProxyTag  = []byte{0xae, 0x26, 0x1e, 0xdb}
 
-	HandshakeSenderPID []byte
-	HandshakePeerPID   []byte
-)
-
-func init() {
 	HandshakeSenderPID = []byte("IPIPPRPDTIME")
 	HandshakePeerPID   = []byte("IPIPPRPDTIME")
-}
+)

ntp/ntp.go

@@ -0,0 +1,60 @@
+package ntp
+
+import (
+	"math/rand"
+	"time"
+
+	"github.com/beevik/ntp"
+	"github.com/juju/errors"
+	"go.uber.org/zap"
+)
+
+const autoUpdatePeriod = time.Minute
+
+var ntpEndpoints = []string{
+	"0.pool.ntp.org",
+	"1.pool.ntp.org",
+	"2.pool.ntp.org",
+	"3.pool.ntp.org",
+}
+
+// Fetch fetches the data on time drift.
+func Fetch() (time.Duration, error) {
+	url := ntpEndpoints[rand.Intn(len(ntpEndpoints))]
+	resp, err := ntp.Query(url)
+	if err != nil {
+		return 0, errors.Annotatef(err, "Cannot fetch NTP server %s", url)
+	}
+
+	offsetInt := int64(resp.ClockOffset)
+	if offsetInt < 0 {
+		offsetInt = -offsetInt
+	}
+	offset := time.Duration(offsetInt)
+
+	return offset, nil
+}
+
+// AutoUpdate runs periodic check of current time .drift state.
+func AutoUpdate() {
+	logger := zap.S().Named("ntp")
+
+	for range time.Tick(autoUpdatePeriod) {
+		diff, err := Fetch()
+		if err != nil {
+			logger.Debugw("Cannot fetch time from NTP", "error", err)
+			continue
+		}
+
+		switch {
+		case diff < 400*time.Millisecond:
+			logger.Debugw("NTP time drift", "value", diff.String())
+		case diff < 600*time.Millisecond:
+			logger.Infow("NTP time drift", "value", diff.String())
+		case diff < 800*time.Millisecond:
+			logger.Warnw("NTP time drift", "value", diff.String())
+		default:
+			logger.Errorw("NTP time drift", "value", diff.String())
+		}
+	}
+}

obfuscated2/frame.go

@@ -113,7 +113,7 @@ func generateFrame(connectionType mtproto.ConnectionType) Frame {
 		}
 
 		// error has to be checked before calling this function
-		tag, _ := connectionType.Tag() // nolint: errcheck
+		tag, _ := connectionType.Tag() // nolint: errcheck, gosec
 		copy(frame.Magic(), tag)
 
 		return frame

obfuscated2/frame_test.go

@@ -78,10 +78,10 @@ func TestFrameGenerateValid(t *testing.T) {
 	}
 	for _, test := range validTests {
 		t.Run(strconv.Itoa(int(test)), func(tt *testing.T) {
-			frame := generateFrame(test)
+			frame := generateFrame(test) // nolint: scopelint
 			conType, err := frame.ConnectionType()
-			assert.Nil(t, err)
+			assert.Nil(tt, err)
-			assert.Equal(t, conType, test)
+			assert.Equal(tt, conType, test) // nolint: scopelint
 		})
 	}
 }

obfuscated2/obfuscated2.go

@@ -23,14 +23,14 @@ type Obfuscated2 struct {
 // Beware, link above is in russian.
 func ParseObfuscated2ClientFrame(secret []byte, frame Frame) (*Obfuscated2, *mtproto.ConnectionOpts, error) {
 	decHasher := sha256.New()
-	decHasher.Write(frame.Key()) // nolint: errcheck
+	decHasher.Write(frame.Key()) // nolint: errcheck, gosec
-	decHasher.Write(secret)      // nolint: errcheck
+	decHasher.Write(secret)      // nolint: errcheck, gosec
 	decryptor := makeStreamCipher(decHasher.Sum(nil), frame.IV())
 
 	invertedFrame := frame.Invert()
 	encHasher := sha256.New()
-	encHasher.Write(invertedFrame.Key()) // nolint: errcheck
+	encHasher.Write(invertedFrame.Key()) // nolint: errcheck, gosec
-	encHasher.Write(secret)              // nolint: errcheck
+	encHasher.Write(secret)              // nolint: errcheck, gosec
 	encryptor := makeStreamCipher(encHasher.Sum(nil), invertedFrame.IV())
 
 	decryptedFrame := make(Frame, FrameLen)
@@ -76,6 +76,6 @@ func MakeTelegramObfuscated2Frame(opts *mtproto.ConnectionOpts) (*Obfuscated2, F
 }
 
 func makeStreamCipher(key, iv []byte) cipher.Stream {
-	block, _ := aes.NewCipher(key)
+	block, _ := aes.NewCipher(key) // nolint: gosec
 	return cipher.NewCTR(block, iv)
 }

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, gosec
-	clientHasher.Write(secret)
+	clientHasher.Write(secret)            // nolint: errcheck, gosec
 	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, gosec
-	clientHasher.Write(secret)
+	clientHasher.Write(secret)                    // nolint: errcheck, gosec
 	invertedClientKey := clientHasher.Sum(nil)
 	clientDecryptor := makeStreamCipher(invertedClientKey, invertedClientFrame.IV())
 

proxy/proxy.go

@@ -1,14 +1,16 @@
 package proxy
 
 import (
+	"context"
 	"io"
 	"net"
 	"sync"
 
+	"github.com/gofrs/uuid"
 	"github.com/juju/errors"
-	uuid "github.com/satori/go.uuid"
 	"go.uber.org/zap"
 
+	"github.com/9seconds/mtg/antireplay"
 	"github.com/9seconds/mtg/client"
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
@@ -19,6 +21,7 @@ import (
 
 // Proxy is a core of this program.
 type Proxy struct {
+	antiReplayCache antireplay.Cache
 	clientInit      client.Init
 	tg              telegram.Telegram
 	conf            *config.Config
@@ -41,11 +44,13 @@ func (p *Proxy) Serve() error {
 }
 
 func (p *Proxy) accept(conn net.Conn) {
-	connID := uuid.NewV4().String()
+	connID := uuid.Must(uuid.NewV4()).String()
 	log := zap.S().With("connection_id", connID).Named("main")
+	ctx, cancel := context.WithCancel(context.Background())
 
 	defer func() {
-		conn.Close() // nolint: errcheck
+		cancel()
+		conn.Close() // nolint: errcheck, gosec
 
 		if err := recover(); err != nil {
 			stats.NewCrash()
@@ -55,23 +60,34 @@ 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.antiReplayCache, p.conf)
 	if err != nil {
 		log.Errorw("Cannot initialize client connection", "error", err)
 		return
 	}
 	defer clientConn.(io.Closer).Close() // nolint: errcheck
 
+	if p.conf.SecureOnly && opts.ConnectionType != mtproto.ConnectionTypeSecure {
+		log.Errorw("Proxy supports only secure connections", "connection_type", opts.ConnectionType)
+		return
+	}
+
 	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() // nolint: gosec
+		clientConn.(io.Closer).Close() // nolint: gosec
+	}()
+
 	wait := &sync.WaitGroup{}
 	wait.Add(2)
 
@@ -79,12 +95,12 @@ func (p *Proxy) accept(conn net.Conn) {
 		clientPacket := clientConn.(wrappers.PacketReadWriteCloser)
 		serverPacket := serverConn.(wrappers.PacketReadWriteCloser)
 		go p.middlePipe(clientPacket, serverPacket, wait, &opts.ReadHacks)
-		go p.middlePipe(serverPacket, clientPacket, wait, &opts.WriteHacks)
+		p.middlePipe(serverPacket, clientPacket, wait, &opts.WriteHacks)
 	} 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)
+		p.directPipe(serverStream, clientStream, wait, p.conf.WriteBufferSize)
 	}
 
 	wait.Wait()
@@ -92,8 +108,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,12 +123,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.Writer, wait *sync.WaitGroup, hacks *mtproto.Hacks) {
-	defer func() {
+	defer wait.Done()
-		src.Close() // nolint: errcheck
-		dst.Close() // nolint: errcheck
-		wait.Done()
-	}()
 
 	for {
 		hacks.SimpleAck = false
@@ -129,23 +142,25 @@ 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.Writer, wait *sync.WaitGroup, bufferSize int) {
-	defer func() {
+	defer wait.Done()
-		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)
 	}
 }
 
 // NewProxy returns new proxy instance.
-func NewProxy(conf *config.Config) *Proxy {
+func NewProxy(conf *config.Config) (*Proxy, error) {
 	var clientInit client.Init
 	var tg telegram.Telegram
 
+	cache, err := antireplay.NewCache(conf)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot make proxy")
+	}
+
 	if conf.UseMiddleProxy() {
 		clientInit = client.MiddleInit
 		tg = telegram.NewMiddleTelegram(conf)
@@ -155,8 +170,9 @@ func NewProxy(conf *config.Config) *Proxy {
 	}
 
 	return &Proxy{
+		antiReplayCache: cache,
 		conf:            conf,
 		clientInit:      clientInit,
 		tg:              tg,
-	}
+	}, nil
 }

scripts/generate_version.sh

@@ -1,12 +0,0 @@
-#!/bin/sh
-set -eu
-
-PROJECT_DIR="$(git rev-parse --show-toplevel)"
-OUTPUT_FILE="${PROJECT_DIR}/version.go"
-
-cat > "$OUTPUT_FILE" <<EOF
-package main
-// autogenerated by $(basename "$0") on $(date -Ru)
-
-const version = "$(git describe --long --always) ($(go version)) [$(date -Ru)]"
-EOF

stats/channels.go

@@ -2,21 +2,20 @@ package stats
 
 import (
 	"net"
-	"time"
 
 	"github.com/9seconds/mtg/mtproto"
 )
 
 const (
-	crashesChanLength     = 1
+	connectionsChanLength = 10
-	connectionsChanLength = 20
+	trafficChanLength     = 10
-	trafficChanLength     = 5000
 )
 
 var (
-	crashesChan     = make(chan struct{}, crashesChanLength)
+	crashesChan     = make(chan struct{})
-	connectionsChan = make(chan *connectionData, connectionsChanLength)
+	statsChan       = make(chan chan<- Stats)
-	trafficChan     = make(chan *trafficData, trafficChanLength)
+	connectionsChan = make(chan connectionData, connectionsChanLength)
+	trafficChan     = make(chan trafficData, trafficChanLength)
 )
 
 type connectionData struct {
@@ -30,87 +29,6 @@ type trafficData struct {
 	ingress bool
 }
 
-func crashManager() {
-	for range crashesChan {
-		instance.mutex.RLock()
-
-		instance.Crashes++
-
-		instance.mutex.RUnlock()
-	}
-}
-
-func connectionManager() { // nolint: gocyclo
-	for event := range connectionsChan {
-		instance.mutex.RLock()
-
-		isIPv4 := event.addr.IP.To4() != nil
-		var inc uint32 = 1
-		if !event.connected {
-			inc = ^uint32(0)
-		}
-
-		switch event.connectionType {
-		case mtproto.ConnectionTypeAbridged:
-			if isIPv4 {
-				instance.ActiveConnections.Abridged.IPv4 += inc
-				if event.connected {
-					instance.AllConnections.Abridged.IPv4 += inc
-				}
-			} else {
-				instance.ActiveConnections.Abridged.IPv6 += inc
-				if event.connected {
-					instance.AllConnections.Abridged.IPv6 += inc
-				}
-			}
-		default:
-			if isIPv4 {
-				instance.ActiveConnections.Intermediate.IPv4 += inc
-				if event.connected {
-					instance.AllConnections.Intermediate.IPv4 += inc
-				}
-			} else {
-				instance.ActiveConnections.Intermediate.IPv6 += inc
-				if event.connected {
-					instance.AllConnections.Intermediate.IPv6 += inc
-				}
-			}
-		}
-
-		instance.mutex.RUnlock()
-	}
-}
-
-func trafficManager() {
-	speedChan := time.Tick(time.Second)
-
-	for {
-		select {
-		case event := <-trafficChan:
-			instance.mutex.RLock()
-
-			if event.ingress {
-				instance.Traffic.Ingress += trafficValue(event.traffic)
-				instance.speedCurrent.Ingress += trafficSpeedValue(event.traffic)
-			} else {
-				instance.Traffic.Egress += trafficValue(event.traffic)
-				instance.speedCurrent.Egress += trafficSpeedValue(event.traffic)
-			}
-
-			instance.mutex.RUnlock()
-		case <-speedChan:
-			instance.mutex.RLock()
-
-			instance.Speed.Ingress = instance.speedCurrent.Ingress
-			instance.Speed.Egress = instance.speedCurrent.Egress
-			instance.speedCurrent.Ingress = trafficSpeedValue(0)
-			instance.speedCurrent.Egress = trafficSpeedValue(0)
-
-			instance.mutex.RUnlock()
-		}
-	}
-}
-
 // NewCrash indicates new crash.
 func NewCrash() {
 	crashesChan <- struct{}{}
@@ -118,7 +36,7 @@ func NewCrash() {
 
 // ClientConnected indicates that new client was connected.
 func ClientConnected(connectionType mtproto.ConnectionType, addr *net.TCPAddr) {
-	connectionsChan <- &connectionData{
+	connectionsChan <- connectionData{
 		connectionType: connectionType,
 		addr:           addr,
 		connected:      true,
@@ -127,7 +45,7 @@ func ClientConnected(connectionType mtproto.ConnectionType, addr *net.TCPAddr) {
 
 // ClientDisconnected indicates that client was disconnected.
 func ClientDisconnected(connectionType mtproto.ConnectionType, addr *net.TCPAddr) {
-	connectionsChan <- &connectionData{
+	connectionsChan <- connectionData{
 		connectionType: connectionType,
 		addr:           addr,
 		connected:      false,
@@ -136,7 +54,7 @@ func ClientDisconnected(connectionType mtproto.ConnectionType, addr *net.TCPAddr
 
 // IngressTraffic accounts new ingress traffic.
 func IngressTraffic(traffic int) {
-	trafficChan <- &trafficData{
+	trafficChan <- trafficData{
 		traffic: traffic,
 		ingress: true,
 	}
@@ -144,8 +62,15 @@ func IngressTraffic(traffic int) {
 
 // EgressTraffic accounts new ingress traffic.
 func EgressTraffic(traffic int) {
-	trafficChan <- &trafficData{
+	trafficChan <- trafficData{
 		traffic: traffic,
 		ingress: false,
 	}
 }
+
+// GetStats returns a snapshot of Stats instance.
+func GetStats() Stats {
+	rpcChan := make(chan Stats)
+	statsChan <- rpcChan
+	return <-rpcChan
+}

stats/init.go

@@ -0,0 +1,28 @@
+package stats
+
+import (
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/config"
+)
+
+// Init initializes stats subsystem.
+func Init(conf *config.Config) error {
+	if conf.StatsD.Enabled {
+		client, err := newStatsd(conf)
+		if err != nil {
+			return errors.Annotate(err, "Cannot initialize statsd client")
+		}
+		go client.run()
+	}
+	prometheus, err := newPrometheus(conf)
+	if err != nil {
+		return errors.Annotate(err, "Cannot initialize prometheus client")
+	}
+	go prometheus.run()
+
+	go NewStats(conf).start()
+	go startServer(conf, prometheus.getHTTPHandler())
+
+	return nil
+}

stats/prometheus.go

@@ -0,0 +1,91 @@
+package stats
+
+import (
+	"net/http"
+	"time"
+
+	"github.com/juju/errors"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+
+	"github.com/9seconds/mtg/config"
+)
+
+const prometheusPollTime = time.Second
+
+type prometheusExporter struct {
+	registry prometheus.Gatherer
+
+	connections *prometheus.GaugeVec
+	traffic     *prometheus.GaugeVec
+	speed       *prometheus.GaugeVec
+	crashes     prometheus.Gauge
+}
+
+func (p *prometheusExporter) run() {
+	for range time.Tick(prometheusPollTime) {
+		instance := GetStats()
+
+		p.connections.WithLabelValues("abridged", "v4").Set(float64(instance.Connections.Abridged.IPv4))
+		p.connections.WithLabelValues("abridged", "v6").Set(float64(instance.Connections.Abridged.IPv6))
+		p.connections.WithLabelValues("intermediate", "v4").Set(float64(instance.Connections.Intermediate.IPv4))
+		p.connections.WithLabelValues("intermediate", "v6").Set(float64(instance.Connections.Intermediate.IPv6))
+		p.connections.WithLabelValues("secure", "v4").Set(float64(instance.Connections.Secure.IPv4))
+		p.connections.WithLabelValues("secure", "v6").Set(float64(instance.Connections.Secure.IPv6))
+		p.traffic.WithLabelValues("ingress").Set(float64(instance.Traffic.ingress))
+		p.traffic.WithLabelValues("egress").Set(float64(instance.Traffic.egress))
+		p.speed.WithLabelValues("ingress").Set(float64(instance.Speed.ingress))
+		p.speed.WithLabelValues("egress").Set(float64(instance.Speed.egress))
+		p.crashes.Set(float64(instance.Crashes))
+	}
+}
+
+func (p *prometheusExporter) getHTTPHandler() http.Handler {
+	return promhttp.HandlerFor(p.registry, promhttp.HandlerOpts{})
+}
+
+func newPrometheus(conf *config.Config) (*prometheusExporter, error) {
+	registry := prometheus.NewRegistry()
+
+	connections := prometheus.NewGaugeVec(prometheus.GaugeOpts{
+		Namespace: conf.Prometheus.Prefix,
+		Name:      "connections",
+		Help:      "Current number of connections to the proxy.",
+	}, []string{"type", "protocol"})
+	traffic := prometheus.NewGaugeVec(prometheus.GaugeOpts{
+		Namespace: conf.Prometheus.Prefix,
+		Name:      "traffic",
+		Help:      "Traffic passed through the proxy in bytes.",
+	}, []string{"direction"})
+	speed := prometheus.NewGaugeVec(prometheus.GaugeOpts{
+		Namespace: conf.Prometheus.Prefix,
+		Name:      "speed",
+		Help:      "Current throughput in bytes per second.",
+	}, []string{"direction"})
+	crashes := prometheus.NewGauge(prometheus.GaugeOpts{
+		Namespace: conf.Prometheus.Prefix,
+		Name:      "crashes",
+		Help:      "How many crashes happened.",
+	})
+
+	if err := registry.Register(connections); err != nil {
+		return nil, errors.Annotate(err, "Cannot register connections collector")
+	}
+	if err := registry.Register(traffic); err != nil {
+		return nil, errors.Annotate(err, "cannot register traffic collector")
+	}
+	if err := registry.Register(speed); err != nil {
+		return nil, errors.Annotate(err, "cannot register speed collector")
+	}
+	if err := registry.Register(crashes); err != nil {
+		return nil, errors.Annotate(err, "cannot register crashes collector")
+	}
+
+	return &prometheusExporter{
+		registry:    registry,
+		connections: connections,
+		traffic:     traffic,
+		speed:       speed,
+		crashes:     crashes,
+	}, nil
+}

stats/server.go

@@ -3,53 +3,36 @@ package stats
 import (
 	"encoding/json"
 	"net/http"
-	"sync"
-	"time"
 
 	"go.uber.org/zap"
 
 	"github.com/9seconds/mtg/config"
 )
 
-var instance *stats
+func startServer(conf *config.Config, prometheusHandler http.Handler) {
-
-// Start starts new statisitcs server.
-func Start(conf *config.Config) {
 	log := zap.S().Named("stats")
 
-	instance = &stats{
+	http.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
-		URLs:   conf.GetURLs(),
-		Uptime: uptime(time.Now()),
-		mutex:  &sync.RWMutex{},
-	}
-
-	go crashManager()
-	go connectionManager()
-	go trafficManager()
-
-	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "application/json")
 
-		instance.mutex.Lock()
+		first, err := json.Marshal(GetStats())
-		first, err := json.Marshal(instance)
-		instance.mutex.Unlock()
-
 		if err != nil {
 			log.Errorw("Cannot encode json", "error", err)
 			http.Error(w, "Internal server error", 500)
 			return
 		}
 
-		interm := map[string]interface{}{}
+		interim := map[string]interface{}{}
-		json.Unmarshal(first, &interm) // nolint: errcheck
+		json.Unmarshal(first, &interim) // nolint: errcheck, gosec
 
 		encoder := json.NewEncoder(w)
 		encoder.SetEscapeHTML(false)
 		encoder.SetIndent("", "  ")
-		if err = encoder.Encode(interm); err != nil {
+		if err = encoder.Encode(interim); err != nil {
 			log.Errorw("Cannot encode json", "error", err)
 		}
 	})
+	http.Handle("/prometheus/", prometheusHandler)
 
 	if err := http.ListenAndServe(conf.StatAddr(), nil); err != nil {
 		log.Fatalw("Stats server has been stopped", "error", err)

stats/stats.go

@@ -4,12 +4,12 @@ import (
 	"encoding/json"
 	"fmt"
 	"strconv"
-	"sync"
 	"time"
 
 	humanize "github.com/dustin/go-humanize"
 
 	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/mtproto"
 )
 
 type uptime time.Time
@@ -24,77 +24,152 @@ func (u uptime) MarshalJSON() ([]byte, error) {
 	return json.Marshal(value)
 }
 
-type trafficValue uint64
-
-func (t trafficValue) MarshalJSON() ([]byte, error) {
-	tv := uint64(t)
-	value := map[string]interface{}{
-		"bytes": tv,
-		"human": humanize.Bytes(tv),
-	}
-
-	return json.Marshal(value)
-}
-
-type trafficSpeedValue uint64
-
-func (t trafficSpeedValue) MarshalJSON() ([]byte, error) {
-	speed := uint64(t)
-	value := map[string]interface{}{
-		"bytes/s": speed,
-		"human":   fmt.Sprintf("%s/S", humanize.Bytes(speed)),
-	}
-
-	return json.Marshal(value)
-}
-
-type connections struct {
-	All          connectionType `json:"all"`
-	Abridged     connectionType `json:"abridged"`
-	Intermediate connectionType `json:"intermediate"`
-}
-
-func (c connections) MarshalJSON() ([]byte, error) {
-	c.All.IPv4 = c.Abridged.IPv4 + c.Intermediate.IPv4
-	c.All.IPv6 = c.Abridged.IPv6 + c.Intermediate.IPv6
-
-	value := struct {
-		All          connectionType `json:"all"`
-		Abridged     connectionType `json:"abridged"`
-		Intermediate connectionType `json:"intermediate"`
-	}{
-		All:          c.All,
-		Abridged:     c.Abridged,
-		Intermediate: c.Intermediate,
-	}
-
-	return json.Marshal(value)
-}
-
 type connectionType struct {
 	IPv6 uint32 `json:"ipv6"`
 	IPv4 uint32 `json:"ipv4"`
 }
 
+type baseConnections struct {
+	All          connectionType `json:"all"`
+	Abridged     connectionType `json:"abridged"`
+	Intermediate connectionType `json:"intermediate"`
+	Secure       connectionType `json:"secure"`
+}
+
+type connections struct {
+	baseConnections
+}
+
+func (c connections) MarshalJSON() ([]byte, error) {
+	c.All.IPv4 = c.Abridged.IPv4 + c.Intermediate.IPv4 + c.Secure.IPv4
+	c.All.IPv6 = c.Abridged.IPv6 + c.Intermediate.IPv6 + c.Secure.IPv6
+
+	return json.Marshal(c.baseConnections)
+}
+
 type traffic struct {
-	Ingress trafficValue `json:"ingress"`
+	ingress uint64
-	Egress  trafficValue `json:"egress"`
+	egress  uint64
+}
+
+func (t *traffic) dumpValue(value uint64) map[string]interface{} {
+	return map[string]interface{}{
+		"bytes": value,
+		"human": humanize.Bytes(value),
+	}
+}
+
+func (t traffic) MarshalJSON() ([]byte, error) {
+	value := map[string]map[string]interface{}{
+		"ingress": t.dumpValue(t.ingress),
+		"egress":  t.dumpValue(t.egress),
+	}
+
+	return json.Marshal(value)
 }
 
 type speed struct {
-	Ingress trafficSpeedValue `json:"ingress"`
+	ingress uint64
-	Egress  trafficSpeedValue `json:"egress"`
+	egress  uint64
 }
 
-type stats struct {
+func (s *speed) dumpValue(value uint64) map[string]interface{} {
+	return map[string]interface{}{
+		"bytes/s": value,
+		"human":   fmt.Sprintf("%s/s", humanize.Bytes(value)),
+	}
+}
+
+func (s speed) MarshalJSON() ([]byte, error) {
+	value := map[string]map[string]interface{}{
+		"ingress": s.dumpValue(s.ingress),
+		"egress":  s.dumpValue(s.egress),
+	}
+
+	return json.Marshal(value)
+}
+
+// Stats represents a statistics of the proxy.
+type Stats struct {
 	URLs        config.IPURLs `json:"urls"`
-	ActiveConnections connections   `json:"active_connections"`
+	Connections connections   `json:"connections"`
-	AllConnections    connections   `json:"all_connections"`
 	Traffic     traffic       `json:"traffic"`
 	Speed       speed         `json:"speed"`
 	Uptime      uptime        `json:"uptime"`
 	Crashes     uint32        `json:"crashes"`
 
-	speedCurrent speed
+	previousTraffic traffic
-	mutex        *sync.RWMutex
+}
+
+func (s *Stats) start() {
+	speedChan := time.Tick(time.Second)
+
+	for {
+		select {
+		case <-speedChan:
+			s.handleSpeed()
+		case event := <-trafficChan:
+			s.handleTraffic(event)
+		case event := <-connectionsChan:
+			s.handleConnection(event)
+		case getStatsChan := <-statsChan:
+			s.handleGetStats(getStatsChan)
+		case <-crashesChan:
+			s.handleCrash()
+		}
+	}
+}
+
+func (s *Stats) handleTraffic(evt trafficData) {
+	if evt.ingress {
+		s.Traffic.ingress += uint64(evt.traffic)
+	} else {
+		s.Traffic.egress += uint64(evt.traffic)
+	}
+}
+
+func (s *Stats) handleSpeed() {
+	s.Speed.ingress = s.Traffic.ingress - s.previousTraffic.ingress
+	s.Speed.egress = s.Traffic.egress - s.previousTraffic.egress
+	s.previousTraffic.ingress = s.Traffic.ingress
+	s.previousTraffic.egress = s.Traffic.egress
+}
+
+func (s *Stats) handleConnection(evt connectionData) {
+	var inc uint32 = 1
+	if !evt.connected {
+		inc = ^uint32(0)
+	}
+
+	var conn *connectionType
+	switch evt.connectionType {
+	case mtproto.ConnectionTypeAbridged:
+		conn = &s.Connections.Abridged
+	case mtproto.ConnectionTypeSecure:
+		conn = &s.Connections.Secure
+	default:
+		conn = &s.Connections.Intermediate
+	}
+
+	if evt.addr.IP.To4() != nil {
+		conn.IPv4 += inc
+	} else {
+		conn.IPv6 += inc
+	}
+}
+
+func (s *Stats) handleGetStats(getStatsChan chan<- Stats) {
+	getStatsChan <- *s
+}
+
+func (s *Stats) handleCrash() {
+	s.Crashes++
+}
+
+// NewStats creates a new instance of Stats structure.
+func NewStats(conf *config.Config) *Stats {
+	return &Stats{
+		URLs:   conf.GetURLs(),
+		Uptime: uptime(time.Now()),
+	}
 }

stats/statsd.go

@@ -0,0 +1,77 @@
+package stats
+
+import (
+	"time"
+
+	"github.com/juju/errors"
+	statsd "gopkg.in/alexcesaro/statsd.v2"
+
+	"github.com/9seconds/mtg/config"
+)
+
+const (
+	statsdConnectionsAbridgedV4 = "connections.abridged.ipv4"
+	statsdConnectionsAbridgedV6 = "connections.abridged.ipv6"
+
+	statsdConnectionsIntermediateV4 = "connections.intermediate.ipv4"
+	statsdConnectionsIntermediateV6 = "connections.intermediate.ipv6"
+
+	statsdConnectionsSecureV4 = "connections.secure.ipv4"
+	statsdConnectionsSecureV6 = "connections.secure.ipv6"
+
+	statsdTrafficIngress = "traffic.ingress"
+	statsdTrafficEgress  = "traffic.egress"
+
+	statsdSpeedIngress = "speed.ingress"
+	statsdSpeedEgress  = "speed.egress"
+
+	statsdCrashes = "crashes"
+)
+
+const statsdPollTime = time.Second
+
+type statsdExporter struct {
+	client *statsd.Client
+}
+
+func (s *statsdExporter) run() {
+	for range time.Tick(statsdPollTime) {
+		instance := GetStats()
+
+		s.client.Gauge(statsdConnectionsAbridgedV4, instance.Connections.Abridged.IPv4)
+		s.client.Gauge(statsdConnectionsAbridgedV6, instance.Connections.Abridged.IPv6)
+		s.client.Gauge(statsdConnectionsIntermediateV4, instance.Connections.Intermediate.IPv4)
+		s.client.Gauge(statsdConnectionsIntermediateV6, instance.Connections.Intermediate.IPv6)
+		s.client.Gauge(statsdConnectionsSecureV4, instance.Connections.Secure.IPv4)
+		s.client.Gauge(statsdConnectionsSecureV6, instance.Connections.Secure.IPv6)
+		s.client.Gauge(statsdTrafficIngress, instance.Traffic.ingress)
+		s.client.Gauge(statsdTrafficEgress, instance.Traffic.egress)
+		s.client.Gauge(statsdSpeedIngress, instance.Speed.ingress)
+		s.client.Gauge(statsdSpeedEgress, instance.Speed.egress)
+		s.client.Gauge(statsdCrashes, instance.Crashes)
+	}
+}
+
+func newStatsd(conf *config.Config) (*statsdExporter, error) {
+	options := []statsd.Option{
+		statsd.Network(conf.StatsD.Addr.Network()),
+		statsd.Address(conf.StatsD.Addr.String()),
+		statsd.Prefix(conf.StatsD.Prefix),
+	}
+
+	if conf.StatsD.TagsFormat > 0 {
+		options = append(options, statsd.TagsFormat(conf.StatsD.TagsFormat))
+		tags := make([]string, len(conf.StatsD.Tags)*2)
+		for k, v := range conf.StatsD.Tags {
+			tags = append(tags, k, v)
+		}
+		options = append(options, statsd.Tags(tags...))
+	}
+
+	client, err := statsd.New(options...)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot create statsd client")
+	}
+
+	return &statsdExporter{client: client}, nil
+}

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"
@@ -11,20 +12,25 @@ import (
 	"github.com/9seconds/mtg/wrappers"
 )
 
+const (
+	directV4DefaultIdx = 1
+	directV6DefaultIdx = 1
+)
+
 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 +38,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 +47,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 +64,16 @@ 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{
+		baseTelegram: baseTelegram{
 			dialer: tgDialer{
 				Dialer: net.Dialer{Timeout: telegramDialTimeout},
 				conf:   conf,
 			},
+			v4DefaultIdx: directV4DefaultIdx,
+			v6DefaultIdx: directV6DefaultIdx,
 			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() {
@@ -63,18 +65,20 @@ func (t *middleTelegramCaller) update() error {
 		return errors.Annotate(err, "Cannot get proxy secret")
 	}
 
-	v4Addresses, err := t.getTelegramAddresses(tgAddrProxyV4)
+	v4Addresses, v4DefaultIdx, err := t.getTelegramAddresses(tgAddrProxyV4)
 	if err != nil {
 		return errors.Annotate(err, "Cannot get ipv4 addresses")
 	}
 
-	v6Addresses, err := t.getTelegramAddresses(tgAddrProxyV6)
+	v6Addresses, v6DefaultIdx, err := t.getTelegramAddresses(tgAddrProxyV6)
 	if err != nil {
 		return errors.Annotate(err, "Cannot get ipv6 addresses")
 	}
 
 	t.dialerMutex.Lock()
 	t.proxySecret = secret
+	t.v4DefaultIdx = v4DefaultIdx
+	t.v6DefaultIdx = v6DefaultIdx
 	t.v4Addresses = v4Addresses
 	t.v6Addresses = v6Addresses
 	t.dialerMutex.Unlock()
@@ -99,52 +103,87 @@ func (t *middleTelegramCaller) getTelegramProxySecret() ([]byte, error) {
 	return secret, nil
 }
 
-func (t *middleTelegramCaller) getTelegramAddresses(url string) (map[int16][]string, error) {
+func (t *middleTelegramCaller) getTelegramAddresses(url string) (map[int16][]string, int16, error) { // nolint: gocyclo
 	resp, err := t.call(url)
 	if err != nil {
-		return nil, errors.Annotate(err, "Cannot access telegram server")
+		return nil, 0, errors.Annotate(err, "Cannot access telegram server")
 	}
 	defer resp.Body.Close() // nolint: errcheck
 
 	scanner := bufio.NewScanner(resp.Body)
 	data := map[int16][]string{}
+
+	var defaultIdx int16 = 1
 	for scanner.Scan() {
 		text := strings.TrimSpace(scanner.Text())
-		if strings.HasPrefix(text, "#") {
+		switch {
+		case strings.HasPrefix(text, "#"):
 			continue
-		}
+		case strings.HasPrefix(text, "proxy_for"):
-
+			addr, idx, err2 := t.parseProxyFor(text)
-		chunks := middleTelegramProxyConfigSplitter.Split(text, 3)
-		if len(chunks) != 3 || chunks[0] != "proxy_for" {
-			return nil, errors.Errorf("Incorrect config '%s'", text)
-		}
-		dcIdx64, err2 := strconv.ParseInt(chunks[1], 10, 16)
 			if err2 != nil {
-			return nil, errors.Errorf("Incorrect config '%s'", text)
+				return nil, 0, errors.Annotate(err2, "Cannot parse 'proxy_for' section")
 			}
-		dcIdx := int16(dcIdx64)
+			if addresses, ok := data[idx]; ok {
-
+				data[idx] = append(addresses, addr)
-		addr := strings.TrimRight(chunks[2], ";")
-		if _, _, err2 = net.SplitHostPort(addr); err != nil {
-			return nil, errors.Annotatef(err2, "Incorrect config '%s'", text)
-		}
-
-		if addresses, ok := data[dcIdx]; ok {
-			data[dcIdx] = append(addresses, addr)
 			} else {
-			data[dcIdx] = []string{addr}
+				data[idx] = []string{addr}
+			}
+		case strings.HasPrefix(text, "default"):
+			idx, err2 := t.parseDefault(text)
+			if err2 != nil {
+				return nil, 0, errors.Annotate(err2, "Cannot parse 'default' section")
+			}
+			defaultIdx = idx
+		default:
+			return nil, 0, errors.Errorf("Unknown config string '%s'", text)
 		}
 	}
+
 	err = scanner.Err()
 	if err != nil {
-		return nil, errors.Annotate(err, "Cannot read response from the telegram")
+		return nil, 0, errors.Annotate(err, "Cannot read response from the telegram")
 	}
 
-	return data, nil
+	return data, defaultIdx, nil
+}
+
+func (t *middleTelegramCaller) parseProxyFor(text string) (string, int16, error) {
+	chunks := middleTelegramProxyConfigSplitter.Split(text, 3)
+	if len(chunks) != 3 || chunks[0] != "proxy_for" {
+		return "", 0, errors.Errorf("Incorrect config '%s'", text)
+	}
+
+	dcIdx, err := strconv.ParseInt(chunks[1], 10, 16)
+	if err != nil {
+		return "", 0, errors.Annotatef(err, "Incorrect config '%s'", text)
+	}
+
+	addr := strings.TrimRight(chunks[2], ";")
+	if _, _, err = net.SplitHostPort(addr); err != nil {
+		return "", 0, errors.Annotatef(err, "Incorrect config '%s'", text)
+	}
+
+	return addr, int16(dcIdx), nil
+}
+
+func (t *middleTelegramCaller) parseDefault(text string) (int16, error) {
+	chunks := middleTelegramProxyConfigSplitter.Split(text, 2)
+	if len(chunks) != 2 || chunks[0] != "default" {
+		return 0, errors.Errorf("Incorrect config '%s'", text)
+	}
+
+	dcIdxString := strings.TrimRight(chunks[1], ";")
+	dcIdx, err := strconv.ParseInt(dcIdxString, 10, 16)
+	if err != nil {
+		return 0, errors.Annotatef(err, "Incorrect config '%s'", text)
+	}
+
+	return int16(dcIdx), nil
 }
 
 func (t *middleTelegramCaller) call(url string) (*http.Response, error) {
-	req, _ := http.NewRequest("GET", url, nil)
+	req, _ := http.NewRequest("GET", url, nil) // nolint: gosec
 	req.Header.Set("Accept", "text/plain")
 	req.Header.Set("User-Agent", tgUserAgent)
 

telegram/telegram.go

@@ -1,6 +1,7 @@
 package telegram
 
 import (
+	"context"
 	"math/rand"
 
 	"github.com/juju/errors"
@@ -11,36 +12,57 @@ 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)
 }
 
 type baseTelegram struct {
 	dialer tgDialer
 
+	v4DefaultIdx int16
+	v6DefaultIdx int16
 	v4Addresses  map[int16][]string
 	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 {
-		if addr, ok := b.v6Addresses[dcIdx]; ok && len(addr) > 0 {
+		if addr := b.chooseAddress(b.v6Addresses, dcIdx, b.v6DefaultIdx); addr != "" {
-			addrs = append(addrs, addr[rand.Intn(len(addr))])
+			addrs = append(addrs, addr)
 		}
 	}
 	if proto&mtproto.ConnectionProtocolIPv4 != 0 {
-		if addr, ok := b.v4Addresses[dcIdx]; ok && len(addr) > 0 {
+		if addr := b.chooseAddress(b.v4Addresses, dcIdx, b.v4DefaultIdx); addr != "" {
-			addrs = append(addrs, addr[rand.Intn(len(addr))])
+			addrs = append(addrs, addr)
 		}
 	}
 
 	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
 		}
 	}
 
 	return nil, errors.New("Cannot connect to Telegram")
 }
+
+func (b *baseTelegram) chooseAddress(addresses map[int16][]string, idx, defaultIdx int16) string {
+	if addr, ok := addresses[idx]; ok {
+		return b.chooseRandomAddress(addr)
+	} else if addr, ok := addresses[defaultIdx]; ok {
+		return b.chooseRandomAddress(addr)
+	}
+
+	return ""
+}
+
+func (b *baseTelegram) chooseRandomAddress(addresses []string) string {
+	if len(addresses) > 0 {
+		return addresses[rand.Intn(len(addresses))]
+	}
+
+	return ""
+}

wrappers/blockcipher.go

@@ -40,7 +40,7 @@ func (b *BlockCipher) Read(p []byte) (int, error) {
 	}
 
 	b.decryptor.CryptBlocks(buf, buf)
-	b.buf.Write(buf)
+	b.buf.Write(buf) // nolint: gosec
 
 	return b.flush(p)
 }

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
@@ -39,8 +41,10 @@ const (
 // 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 +52,54 @@ 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():
+		c.Close() // nolint: gosec
+		return 0, errors.Annotate(c.ctx.Err(), "Cannot write because context was closed")
+	default:
+		if err := c.conn.SetWriteDeadline(time.Now().Add(connTimeoutWrite)); err != nil {
+			c.Close() // nolint: gosec
+			return 0, errors.Annotate(err, "Cannot set write deadline to the socket")
+		}
 
+		n, err := c.conn.Write(p)
 		c.logger.Debugw("Write to stream", "bytes", n, "error", err)
 		stats.EgressTraffic(n)
+		if err != nil {
+			c.Close() // nolint: gosec
+		}
 
 		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():
+		c.Close() // nolint: gosec
+		return 0, errors.Annotate(c.ctx.Err(), "Cannot read because context was closed")
+	default:
+		if err := c.conn.SetReadDeadline(time.Now().Add(connTimeoutRead)); err != nil {
+			c.Close() // nolint: gosec
+			return 0, errors.Annotate(err, "Cannot set read deadline to the socket")
+		}
 
+		n, err := c.conn.Read(p)
 		c.logger.Debugw("Read from stream", "bytes", n, "error", err)
 		stats.IngressTraffic(n)
+		if err != nil {
+			c.Close() // nolint: gosec
+		}
 
 		return n, err
+	}
 }
 
 // Close closes underlying net.Conn instance.
 func (c *Conn) Close() error {
-	defer c.logger.Debugw("Close connection")
+	c.logger.Debugw("Close connection")
+	c.cancel()
+
 	return c.conn.Close()
 }
 
@@ -100,7 +130,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 +140,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_abridged.go

@@ -119,9 +119,9 @@ func (m *MTProtoAbridged) Write(p []byte) (int, error) {
 		buf := &bytes.Buffer{}
 		buf.Grow(1 + 3 + len(p))
 
-		buf.WriteByte(byte(mtprotoAbridgedSmallPacketLength))
+		buf.WriteByte(byte(mtprotoAbridgedSmallPacketLength)) // nolint: gosec
-		buf.Write(length24[:])
+		buf.Write(length24[:])                                // nolint: gosec
-		buf.Write(p)
+		buf.Write(p)                                          // nolint: gosec
 
 		return m.conn.Write(buf.Bytes())
 	}

wrappers/mtproto_cipher.go

@@ -5,7 +5,7 @@ import (
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/md5"  // nolint: gas
-	"crypto/sha1"
+	"crypto/sha1" // nolint: gosec
 	"encoding/binary"
 	"net"
 
@@ -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)   // nolint: gosec
-	message.Write(req.Nonce[:])
+	message.Write(req.Nonce)    // nolint: gosec
-	message.Write(req.CryptoTS[:])
+	message.Write(req.CryptoTS) // nolint: gosec
 
 	clientIPv4 := emptyIP[:]
 	serverIPv4 := emptyIP[:]
@@ -49,36 +51,36 @@ func deriveKeys(purpose cipherPurpose, req *rpc.NonceRequest, resp *rpc.NonceRes
 		clientIPv4 = utils.ReverseBytes(client.IP.To4())
 		serverIPv4 = utils.ReverseBytes(remote.IP.To4())
 	}
-	message.Write(serverIPv4)
+	message.Write(serverIPv4) // nolint: gosec
 
 	var port [2]byte
 	binary.LittleEndian.PutUint16(port[:], uint16(client.Port))
-	message.Write(port[:])
+	message.Write(port[:]) // nolint: gosec
 
 	switch purpose {
 	case cipherPurposeClient:
-		message.WriteString("CLIENT")
+		message.WriteString("CLIENT") // nolint: gosec
 	case cipherPurposeServer:
-		message.WriteString("SERVER")
+		message.WriteString("SERVER") // nolint: gosec
 	default:
 		panic("Unexpected cipher purpose")
 	}
 
-	message.Write(clientIPv4)
+	message.Write(clientIPv4) // nolint: gosec
 	binary.LittleEndian.PutUint16(port[:], uint16(remote.Port))
-	message.Write(port[:])
+	message.Write(port[:])    // nolint: gosec
-	message.Write(secret)
+	message.Write(secret)     // nolint: gosec
-	message.Write(resp.Nonce[:])
+	message.Write(resp.Nonce) // nolint: gosec
 
 	if client.IP.To4() == nil {
-		message.Write(client.IP.To16())
+		message.Write(client.IP.To16()) // nolint: gosec
-		message.Write(remote.IP.To16())
+		message.Write(remote.IP.To16()) // nolint: gosec
 	}
-	message.Write(req.Nonce[:])
+	message.Write(req.Nonce) // nolint: gosec
 
 	data := message.Bytes()
 	md5sum := md5.Sum(data[1:]) // nolint: gas
-	sha1sum := sha1.Sum(data)
+	sha1sum := sha1.Sum(data)   // nolint: gosec
 
 	key := append(md5sum[:12], sha1sum[:]...)
 	iv := md5.Sum(data[2:]) // 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)
 	}
 
@@ -72,12 +73,12 @@ func (m *MTProtoFrame) Read() ([]byte, error) { // nolint: gocyclo
 	}
 
 	var seqNo int32
-	binary.Read(buf, binary.LittleEndian, &seqNo) // nolint: errcheck
+	binary.Read(buf, binary.LittleEndian, &seqNo) // nolint: errcheck, gosec
 	if seqNo != m.readSeqNo {
 		return nil, errors.Errorf("Unexpected sequence number %d (wait for %d)", seqNo, m.readSeqNo)
 	}
 
-	data, _ := ioutil.ReadAll(buf)
+	data, _ := ioutil.ReadAll(buf) // nolint: gosec
 	buf.Reset()
 	// write to buf, not to writer. This is because we are going to fetch
 	// crc32 checksum.
@@ -108,13 +109,13 @@ func (m *MTProtoFrame) Write(p []byte) (int, error) {
 	buf := &bytes.Buffer{}
 	buf.Grow(messageLength + paddingLength)
 
-	binary.Write(buf, binary.LittleEndian, uint32(messageLength)) // nolint: errcheck
+	binary.Write(buf, binary.LittleEndian, uint32(messageLength)) // nolint: errcheck, gosec
-	binary.Write(buf, binary.LittleEndian, m.writeSeqNo)          // nolint: errcheck
+	binary.Write(buf, binary.LittleEndian, m.writeSeqNo)          // nolint: errcheck, gosec
-	buf.Write(p)
+	buf.Write(p)                                                  // nolint: gosec
 
 	checksum := crc32.ChecksumIEEE(buf.Bytes())
-	binary.Write(buf, binary.LittleEndian, checksum) // nolint: errcheck
+	binary.Write(buf, binary.LittleEndian, checksum)              // nolint: errcheck, gosec
-	buf.Write(bytes.Repeat(mtprotoFramePadding, paddingLength/4))
+	buf.Write(bytes.Repeat(mtprotoFramePadding, paddingLength/4)) // nolint: gosec
 
 	m.logger.Debugw("Write MTProto frame",
 		"length", len(p),

wrappers/mtproto_intermediate.go

@@ -62,10 +62,6 @@ func (m *MTProtoIntermediate) Read() ([]byte, error) {
 		return nil, errors.Annotate(err, "Cannot read the message")
 	}
 
-	if length%4 != 0 {
-		length -= length % 4
-	}
-
 	return buf.Bytes()[:length], nil
 }
 
@@ -80,7 +76,7 @@ func (m *MTProtoIntermediate) Write(p []byte) (int, error) {
 		"counter", m.writeCounter,
 	)
 
-	if m.opts.ReadHacks.SimpleAck {
+	if m.opts.WriteHacks.SimpleAck {
 		return m.conn.Write(p)
 	}
 

wrappers/mtproto_intermediate_secure.go

@@ -0,0 +1,74 @@
+package wrappers
+
+import (
+	"bytes"
+	"encoding/binary"
+	"math/rand"
+
+	"github.com/9seconds/mtg/mtproto"
+)
+
+// MTProtoIntermediateSecure is an extension of MTProtoIntermediate
+// mode which supports random paddings (socalled 'secure mode' or
+// 'dd-secrets').
+type MTProtoIntermediateSecure struct {
+	MTProtoIntermediate
+}
+
+func (m *MTProtoIntermediateSecure) Read() ([]byte, error) {
+	data, err := m.MTProtoIntermediate.Read()
+	if err != nil {
+		return nil, err
+	}
+	length := len(data) - (len(data) % 4)
+
+	return data[:length], nil
+}
+
+func (m *MTProtoIntermediateSecure) Write(p []byte) (int, error) {
+	defer func() {
+		m.writeCounter++
+	}()
+
+	m.logger.Debugw("Write packet",
+		"simple_ack", m.opts.WriteHacks.SimpleAck,
+		"quick_ack", m.opts.WriteHacks.QuickAck,
+		"counter", m.writeCounter,
+	)
+
+	if m.opts.WriteHacks.SimpleAck {
+		return m.conn.Write(p)
+	}
+
+	buf := &bytes.Buffer{}
+	paddingLength := rand.Intn(4)
+	buf.Grow(4 + len(p) + paddingLength)
+
+	binary.Write(buf, binary.LittleEndian, uint32(len(p)+paddingLength)) // nolint: errcheck, gosec
+	buf.Write(p)                                                         // nolint: gosec
+	buf.Write(make([]byte, paddingLength))                               // nolint: gosec
+
+	m.logger.Debugw("Write packet with padding",
+		"simple_ack", m.opts.WriteHacks.SimpleAck,
+		"quick_ack", m.opts.WriteHacks.QuickAck,
+		"counter", m.writeCounter,
+		"padding_length", paddingLength,
+		"length", len(p),
+	)
+
+	_, err := m.conn.Write(buf.Bytes())
+
+	return len(p), err
+}
+
+// NewMTProtoIntermediateSecure create new instance of
+// MTProtoIntermediateSecure instance.
+func NewMTProtoIntermediateSecure(conn StreamReadWriteCloser, opts *mtproto.ConnectionOpts) PacketReadWriteCloser {
+	return &MTProtoIntermediateSecure{
+		MTProtoIntermediate: MTProtoIntermediate{
+			conn:   conn,
+			logger: conn.Logger().Named("mtproto-intermediate-secure"),
+			opts:   opts,
+		},
+	}
+}

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")
@@ -120,7 +120,7 @@ func (m *MTProtoProxy) Write(p []byte) (int, error) {
 			zap.Stringer("flags", flags),
 		)
 	}
-	header.Write(p)
+	header.Write(p) // nolint: gosec
 
 	if _, err := m.conn.Write(header.Bytes()); err != nil {
 		return 0, err
@@ -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) // nolint: gosec
+
+	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,14 @@
+package wrappers
+
+import (
+	"bytes"
+	"sync"
+)
+
+var (
+	streamCipherBufferPool = sync.Pool{
+		New: func() interface{} {
+			return &bytes.Buffer{}
+		},
+	}
+)