Merge remote-tracking branch 'origin/master' into stable

Remove generating of version.go

Dockerfile

@@ -10,8 +10,7 @@ RUN set -x \
     curl \
     git \
     make \
-    upx \
-  && update-ca-certificates
+    upx
 
 COPY . /go/src/github.com/9seconds/mtg/
 
@@ -26,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 \
@@ -34,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

Makefile

@@ -3,26 +3,28 @@ IMAGE_NAME   := mtg
 APP_NAME     := $(IMAGE_NAME)
 
 CC_BINARIES  := $(shell bash -c "echo -n $(APP_NAME)-{linux,freebsd,openbsd}-{386,amd64} $(APP_NAME)-linux-{arm,arm64}")
-APP_DEPS     := version.go
 
-GOLANGCI_LINT_VERSION := v1.10.2
+GOLANGCI_LINT_VERSION := v1.11.2
 
-COMMON_BUILD_FLAGS := -ldflags="-s -w"
+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):
 	@$(MOD_ON) go build $(COMMON_BUILD_FLAGS) -o "$(APP_NAME)"
 
-static-$(APP_NAME): $(APP_DEPS)
+static-$(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
 	@$(MOD_ON) env "GOOS=$(GOOS)" "GOARCH=$(GOARCH)" \
 		go build \
 		$(COMMON_BUILD_FLAGS) \
@@ -31,9 +33,6 @@ $(APP_NAME)-%: $(APP_DEPS) ccbuilds
 ccbuilds:
 	@rm -rf ./ccbuilds && mkdir -p ./ccbuilds
 
-version.go:
-	@$(MOD_ON) go generate main.go
-
 vendor: go.mod go.sum
 	@$(MOD_ON) go mod vendor
 
@@ -53,15 +52,15 @@ crosscompile-dir:
 	@rm -rf "$(CC_DIR)" && mkdir -p "$(CC_DIR)"
 
 .PHONY: test
-test: vendor $(APP_DEPS)
+test: vendor
 	@$(MOD_ON) go test -v ./...
 
 .PHONY: lint
-lint: vendor $(APP_DEPS)
+lint: vendor
 	@$(MOD_OFF) golangci-lint run
 
 .PHONY: critic
-critic: vendor $(APP_DEPS)
+critic: vendor
 	@$(MOD_OFF) gocritic check-project "$(ROOT_DIR)"
 
 .PHONY: clean

README.md

@@ -3,6 +3,7 @@
 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
@@ -32,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.5 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.
 
@@ -112,11 +113,49 @@ head -c 512 /dev/urandom | md5sum | cut -f 1 -d ' '
 
 ## Secure mode
 
-If you want to support new secure mode, please prepend `dd` to the
-secret. For example, secret `cf18fa8ea0267057e2c61a5f7322a8e7` should
-be `ddcf18fa8ea0267057e2c61a5f7322a8e7`. But pay attention that some
-old clients won't support this mode. If this is not your case, I would
-suggest to go with this mode.
+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.
 
 Oneliners to generate such secrets:
 
@@ -130,6 +169,7 @@ or
 echo dd$(head -c 512 /dev/urandom | md5sum | cut -f 1 -d ' ')
 ```
 
+
 ## Environment variables
 
 It is possible to configure this tool using environment variables. You
@@ -154,8 +194,10 @@ supported environment variables:
 | `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).                                                                                                                                                                                                  |
 
 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
@@ -184,13 +226,13 @@ This tool will listen on port 3128 by default with the given secret.
 # One-line runner
 
 ```console
-docker run --name mtg --restart=unless-stopped -p 3128:3128 -p 3129:3129 -d nineseconds/mtg:stable $(openssl rand -hex 16)
+docker run --name mtg --restart=unless-stopped -p 3128:3128 -q 3129:3129 -d nineseconds/mtg:stable $(openssl rand -hex 16)
 ```
 
 or in secret mode:
 
 ```console
-docker run --name mtg --restart=unless-stopped -p 3128:3128 -p 3129:3129 -d nineseconds/mtg:stable dd$(openssl rand -hex 16)
+docker run --name mtg --restart=unless-stopped -p 3128:3128 -q 3129:3129 -d nineseconds/mtg:stable dd$(openssl rand -hex 16)
 ```
 
 You will have this tool up and running on port 3128. Now curl
@@ -228,3 +270,14 @@ All metrics are gauges. Here is the list of metrics and their meaning:
 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/`.

config/config.go

@@ -16,6 +16,7 @@ type Config struct {
 	Debug      bool
 	Verbose    bool
 	SecureMode bool
+	SecureOnly bool
 
 	ReadBufferSize  int
 	WriteBufferSize int
@@ -37,6 +38,9 @@ type Config struct {
 		TagsFormat statsd.TagFormat
 		Enabled    bool
 	}
+	Prometheus struct {
+		Prefix string
+	}
 
 	Secret []byte
 	AdTag  []byte
@@ -115,9 +119,10 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo
 	bindIP, publicIPv4, publicIPv6, statsIP net.IP,
 	bindPort, publicIPv4Port, publicIPv6Port, statsPort, statsdPort uint16,
 	statsdIP, statsdNetwork, statsdPrefix, statsdTagsFormat string,
-	statsdTags map[string]string,
+	statsdTags map[string]string, prometheusPrefix string,
+	secureOnly bool,
 	secret, adtag []byte) (*Config, error) {
-	secureMode := false
+	secureMode := secureOnly
 	if bytes.HasPrefix(secret, []byte{0xdd}) && len(secret) == 17 {
 		secureMode = true
 		secret = bytes.TrimPrefix(secret, []byte{0xdd})
@@ -157,6 +162,7 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo
 	conf := &Config{
 		Debug:           debug,
 		Verbose:         verbose,
+		SecureOnly:      secureOnly,
 		BindIP:          bindIP,
 		BindPort:        bindPort,
 		PublicIPv4:      publicIPv4,
@@ -171,6 +177,7 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo
 		ReadBufferSize:  int(readBufferSize),
 		WriteBufferSize: int(writeBufferSize),
 	}
+	conf.Prometheus.Prefix = prometheusPrefix
 
 	if statsdIP != "" {
 		conf.StatsD.Enabled = true

go.mod

@@ -4,20 +4,29 @@ require (
 	github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc // indirect
 	github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf // indirect
 	github.com/beevik/ntp v0.2.0
+	github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/dustin/go-humanize v0.0.0-20180713052910-9f541cc9db5d
+	github.com/dustin/go-humanize v1.0.0
 	github.com/gofrs/uuid v3.1.0+incompatible
-	github.com/juju/errors v0.0.0-20180806074554-22422dad46e1
+	github.com/gogo/protobuf v1.1.1 // indirect
+	github.com/golang/protobuf v1.2.0 // indirect
+	github.com/juju/errors v0.0.0-20181012004132-a4583d0a56ea
 	github.com/juju/loggo v0.0.0-20180524022052-584905176618 // indirect
 	github.com/juju/testing v0.0.0-20180920084828-472a3e8b2073 // indirect
 	github.com/kr/pretty v0.1.0 // indirect
+	github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect
 	github.com/pkg/errors v0.8.0 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/prometheus/client_golang v0.9.0
+	github.com/prometheus/client_model v0.0.0-20180712105110-5c3871d89910 // indirect
+	github.com/prometheus/common v0.0.0-20181020173914-7e9e6cabbd39 // indirect
+	github.com/prometheus/procfs v0.0.0-20181005140218-185b4288413d // indirect
 	github.com/stretchr/testify v1.2.2
 	go.uber.org/atomic v1.3.2 // indirect
 	go.uber.org/multierr v1.1.0 // indirect
 	go.uber.org/zap v1.9.1
-	golang.org/x/net v0.0.0-20180921000356-2f5d2388922f // indirect
+	golang.org/x/net v0.0.0-20181102091132-c10e9556a7bc // indirect
+	golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f // 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

go.sum

@@ -4,14 +4,20 @@ github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf h1:qet1QNfXsQxTZq
 github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
 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/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 v0.0.0-20180713052910-9f541cc9db5d h1:lDrio3iIdNb0Gw9CgH7cQF+iuB5mOOjdJ9ERNJCBgb4=
-github.com/dustin/go-humanize v0.0.0-20180713052910-9f541cc9db5d/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
+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/gofrs/uuid v3.1.0+incompatible h1:q2rtkjaKT4YEr6E1kamy0Ha4RtepWlQBedyHx0uzKwA=
 github.com/gofrs/uuid v3.1.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
-github.com/juju/errors v0.0.0-20180806074554-22422dad46e1 h1:wnhMXidtb70kDZCeLt/EfsVtkXS5c8zLnE9y/6DIRAU=
-github.com/juju/errors v0.0.0-20180806074554-22422dad46e1/go.mod h1:W54LbzXuIE0boCoNJfwqpmkKJ1O4TCTZMetAt6jGk7Q=
+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/juju/errors v0.0.0-20181012004132-a4583d0a56ea h1:g2k+8WR7cHch4g0tBDhfiEvAp7fXxTNBiD1oC1Oxj3E=
+github.com/juju/errors v0.0.0-20181012004132-a4583d0a56ea/go.mod h1:W54LbzXuIE0boCoNJfwqpmkKJ1O4TCTZMetAt6jGk7Q=
 github.com/juju/loggo v0.0.0-20180524022052-584905176618 h1:MK144iBQF9hTSwBW/9eJm034bVoG30IshVm688T2hi8=
 github.com/juju/loggo v0.0.0-20180524022052-584905176618/go.mod h1:vgyd7OREkbtVEN/8IXZe5Ooef3LQePvuBm9UWj6ZL8U=
 github.com/juju/testing v0.0.0-20180920084828-472a3e8b2073 h1:WQM1NildKThwdP7qWrNAFGzp4ijNLw8RlgENkaI4MJs=
@@ -21,10 +27,20 @@ github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORN
 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/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/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.0 h1:tXuTFVHC03mW0D+Ua1Q2d1EAVqLTuggX50V0VLICCzY=
+github.com/prometheus/client_golang v0.9.0/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw=
+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/common v0.0.0-20181020173914-7e9e6cabbd39 h1:Cto4X6SVMWRPBkJ/3YHn1iDGDGc/Z+sW+AEMKHMVvN4=
+github.com/prometheus/common v0.0.0-20181020173914-7e9e6cabbd39/go.mod h1:daVV7qP5qjZbuso7PdcryaAu0sAZbrN9i7WWcTMWvro=
+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/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 go.uber.org/atomic v1.3.2 h1:2Oa65PReHzfn29GpvgsYwloV9AVFHPDk8tYxt2c2tr4=
@@ -33,8 +49,10 @@ 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.9.1 h1:XCJQEf3W6eZaVwhRBof6ImoYGJSITeKWsyeh3HFu/5o=
 go.uber.org/zap v1.9.1/go.mod h1:vwi/ZaCAaUcBkycHslxD9B2zi4UTXhF60s6SWpuDF0Q=
-golang.org/x/net v0.0.0-20180921000356-2f5d2388922f h1:QM2QVxvDoW9PFSPp/zy9FgxJLfaWTZlS61KEPtBwacM=
-golang.org/x/net v0.0.0-20180921000356-2f5d2388922f/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
+golang.org/x/net v0.0.0-20181102091132-c10e9556a7bc h1:ZMCWScCvS2fUVFw8LOpxyUUW5qiviqr4Dg5NdjLeiLU=
+golang.org/x/net v0.0.0-20181102091132-c10e9556a7bc/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
+golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f h1:wMNYb4v58l5UBM7MYRLPG6ZhfOqbKu7X5eyFl8ZhKvA=
+golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 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=

main.go

@@ -1,7 +1,5 @@
 package main
 
-//go:generate scripts/generate_version.sh
-
 import (
 	"encoding/json"
 	"fmt"
@@ -22,6 +20,8 @@ import (
 	"github.com/9seconds/mtg/stats"
 )
 
+var version = "dev" // this has to be set by build ld flags
+
 var (
 	app = kingpin.New("mtg", "Simple MTPROTO proxy.")
 
@@ -110,6 +110,12 @@ var (
 		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').
@@ -122,6 +128,11 @@ var (
 		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()
 
 	secret = app.Arg("secret", "Secret of this proxy.").Required().HexBytes()
 	adtag  = app.Arg("adtag", "ADTag of the proxy.").HexBytes()
@@ -146,7 +157,7 @@ func main() { // nolint: gocyclo
 		*bindIP, *publicIPv4, *publicIPv6, *statsIP,
 		*bindPort, *publicIPv4Port, *publicIPv6Port, *statsPort, *statsdPort,
 		*statsdIP, *statsdNetwork, *statsdPrefix, *statsdTagsFormat,
-		*statsdTags,
+		*statsdTags, *prometheusPrefix, *secureOnly,
 		*secret, *adtag,
 	)
 	if err != nil {
@@ -172,6 +183,7 @@ func main() { // nolint: gocyclo
 	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")
@@ -188,7 +200,7 @@ func main() { // nolint: gocyclo
 		zap.S().Infow("Use direct connection to Telegram")
 	}
 
-	if err := stats.Start(conf); err != nil {
+	if err := stats.Init(conf); err != nil {
 		panic(err)
 	}
 

proxy/proxy.go

@@ -65,6 +65,11 @@ func (p *Proxy) accept(conn net.Conn) {
 	}
 	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())
 
@@ -88,12 +93,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, p.conf.ReadBufferSize)
-		go p.directPipe(serverStream, clientStream, wait, p.conf.WriteBufferSize)
+		p.directPipe(serverStream, clientStream, wait, p.conf.WriteBufferSize)
 	}
 
 	wait.Wait()
@@ -116,13 +121,8 @@ func (p *Proxy) getTelegramConn(ctx context.Context, cancel context.CancelFunc,
 	return packetConn, nil
 }
 
-func (p *Proxy) middlePipe(src wrappers.PacketReadCloser, dst io.WriteCloser,
-	wait *sync.WaitGroup, hacks *mtproto.Hacks) {
-	defer func() {
-		src.Close() // nolint: errcheck, gosec
-		dst.Close() // nolint: errcheck, gosec
-		wait.Done()
-	}()
+func (p *Proxy) middlePipe(src wrappers.PacketReadCloser, dst io.Writer, wait *sync.WaitGroup, hacks *mtproto.Hacks) {
+	defer wait.Done()
 
 	for {
 		hacks.SimpleAck = false
@@ -140,13 +140,8 @@ func (p *Proxy) middlePipe(src wrappers.PacketReadCloser, dst io.WriteCloser,
 	}
 }
 
-func (p *Proxy) directPipe(src wrappers.StreamReadCloser, dst io.WriteCloser,
-	wait *sync.WaitGroup, bufferSize int) {
-	defer func() {
-		src.Close() // nolint: errcheck, gosec
-		dst.Close() // nolint: errcheck, gosec
-		wait.Done()
-	}()
+func (p *Proxy) directPipe(src wrappers.StreamReadCloser, dst io.Writer, wait *sync.WaitGroup, bufferSize int) {
+	defer wait.Done()
 
 	buffer := make([]byte, bufferSize)
 	if _, err := io.CopyBuffer(dst, src, buffer); err != 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 = 20
-	trafficChanLength     = 5000
+	connectionsChanLength = 10
+	trafficChanLength     = 10
 )
 
 var (
-	crashesChan     = make(chan struct{}, crashesChanLength)
-	connectionsChan = make(chan *connectionData, connectionsChanLength)
-	trafficChan     = make(chan *trafficData, trafficChanLength)
+	crashesChan     = make(chan struct{})
+	statsChan       = make(chan chan<- Stats)
+	connectionsChan = make(chan connectionData, connectionsChanLength)
+	trafficChan     = make(chan trafficData, trafficChanLength)
 )
 
 type connectionData struct {
@@ -30,81 +29,6 @@ type trafficData struct {
 	ingress bool
 }
 
-func crashManager() {
-	for range crashesChan {
-		instance.mutex.RLock()
-
-		instance.Crashes++
-
-		instance.mutex.RUnlock()
-	}
-}
-
-func connectionManager() {
-	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.Connections.Abridged.IPv4 += inc
-			} else {
-				instance.Connections.Abridged.IPv6 += inc
-			}
-		case mtproto.ConnectionTypeSecure:
-			if isIPv4 {
-				instance.Connections.Secure.IPv4 += inc
-			} else {
-				instance.Connections.Secure.IPv6 += inc
-			}
-		default:
-			if isIPv4 {
-				instance.Connections.Intermediate.IPv4 += inc
-			} else {
-				instance.Connections.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{}{}
@@ -112,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,
@@ -121,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,
@@ -130,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,
 	}
@@ -138,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,67 +3,38 @@ package stats
 import (
 	"encoding/json"
 	"net/http"
-	"sync"
-	"time"
 
 	"go.uber.org/zap"
 
 	"github.com/9seconds/mtg/config"
 )
 
-var instance *stats
-
-// Start starts new statistics server.
-func Start(conf *config.Config) error {
+func startServer(conf *config.Config, prometheusHandler http.Handler) {
 	log := zap.S().Named("stats")
 
-	instance = &stats{
-		URLs:   conf.GetURLs(),
-		Uptime: uptime(time.Now()),
-		mutex:  &sync.RWMutex{},
-	}
-
-	if conf.StatsD.Enabled {
-		client, err := newStatsd(conf)
-		if err != nil {
-			return err
-		}
-		go client.run()
-	}
-
-	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(instance)
-		instance.mutex.Unlock()
-
+		first, err := json.Marshal(GetStats())
 		if err != nil {
 			log.Errorw("Cannot encode json", "error", err)
 			http.Error(w, "Internal server error", 500)
 			return
 		}
 
-		interm := map[string]interface{}{}
-		json.Unmarshal(first, &interm) // nolint: errcheck, gosec
+		interim := map[string]interface{}{}
+		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)
 
-	go func() {
-		if err := http.ListenAndServe(conf.StatAddr(), nil); err != nil {
-			log.Fatalw("Stats server has been stopped", "error", err)
-		}
-	}()
-
-	return nil
+	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,72 +24,73 @@ 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 connectionType struct {
+	IPv6 uint32 `json:"ipv6"`
+	IPv4 uint32 `json:"ipv4"`
 }
 
-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 {
+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
 
-	value := struct {
-		All          connectionType `json:"all"`
-		Abridged     connectionType `json:"abridged"`
-		Intermediate connectionType `json:"intermediate"`
-		Secure       connectionType `json:"secure"`
-	}{
-		All:          c.All,
-		Abridged:     c.Abridged,
-		Intermediate: c.Intermediate,
-		Secure:       c.Secure,
+	return json.Marshal(c.baseConnections)
+}
+
+type traffic struct {
+	ingress uint64
+	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 connectionType struct {
-	IPv6 uint32 `json:"ipv6"`
-	IPv4 uint32 `json:"ipv4"`
-}
-
-type traffic struct {
-	Ingress trafficValue `json:"ingress"`
-	Egress  trafficValue `json:"egress"`
-}
-
 type speed struct {
-	Ingress trafficSpeedValue `json:"ingress"`
-	Egress  trafficSpeedValue `json:"egress"`
+	ingress uint64
+	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"`
 	Connections connections   `json:"connections"`
 	Traffic     traffic       `json:"traffic"`
@@ -97,6 +98,78 @@ type stats struct {
 	Uptime      uptime        `json:"uptime"`
 	Crashes     uint32        `json:"crashes"`
 
-	speedCurrent speed
-	mutex        *sync.RWMutex
+	previousTraffic traffic
+}
+
+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

@@ -36,7 +36,7 @@ type statsdExporter struct {
 
 func (s *statsdExporter) run() {
 	for range time.Tick(statsdPollTime) {
-		instance.mutex.Lock()
+		instance := GetStats()
 
 		s.client.Gauge(statsdConnectionsAbridgedV4, instance.Connections.Abridged.IPv4)
 		s.client.Gauge(statsdConnectionsAbridgedV6, instance.Connections.Abridged.IPv6)
@@ -44,13 +44,11 @@ func (s *statsdExporter) run() {
 		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, uint64(instance.Traffic.Ingress))
-		s.client.Gauge(statsdTrafficEgress, uint64(instance.Traffic.Egress))
-		s.client.Gauge(statsdSpeedIngress, uint64(instance.Speed.Ingress))
-		s.client.Gauge(statsdSpeedEgress, uint64(instance.Speed.Egress))
+		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)
-
-		instance.mutex.Unlock()
 	}
 }
 

wrappers/conn.go

@@ -38,13 +38,6 @@ const (
 	connTimeoutWrite = 2 * time.Minute
 )
 
-type ioResult struct {
-	n   int
-	err error
-}
-
-type ioFunc func([]byte) (int, error)
-
 // Conn is a basic wrapper for net.Conn providing the most low-level
 // logic and management as possible.
 type Conn struct {
@@ -61,12 +54,20 @@ type Conn struct {
 func (c *Conn) Write(p []byte) (int, error) {
 	select {
 	case <-c.ctx.Done():
+		c.Close() // nolint: gosec
 		return 0, errors.Annotate(c.ctx.Err(), "Cannot write because context was closed")
 	default:
-		n, err := c.doIO(c.conn.Write, p, connTimeoutWrite)
+		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
 	}
@@ -75,48 +76,30 @@ func (c *Conn) Write(p []byte) (int, error) {
 func (c *Conn) Read(p []byte) (int, error) {
 	select {
 	case <-c.ctx.Done():
+		c.Close() // nolint: gosec
 		return 0, errors.Annotate(c.ctx.Err(), "Cannot read because context was closed")
 	default:
-		n, err := c.doIO(c.conn.Read, p, connTimeoutRead)
+		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
 	}
 }
 
-func (c *Conn) doIO(callback ioFunc, p []byte, timeout time.Duration) (int, error) {
-	resChan := make(chan ioResult, 1)
-	timer := time.NewTimer(timeout)
-
-	go func() {
-		n, err := callback(p)
-		resChan <- ioResult{n: n, err: err}
-	}()
-
-	select {
-	case res := <-resChan:
-		timer.Stop()
-		if res.err != nil {
-			c.Close() // nolint: gosec
-		}
-		return res.n, res.err
-	case <-c.ctx.Done():
-		timer.Stop()
-		c.Close() // nolint: gosec
-		return 0, errors.Annotate(c.ctx.Err(), "Cannot do IO because context is closed")
-	case <-timer.C:
-		c.Close() // nolint: gosec
-		return 0, errors.Annotate(c.ctx.Err(), "Timeout on IO operation")
-	}
-}
-
 // Close closes underlying net.Conn instance.
 func (c *Conn) Close() error {
-	defer c.logger.Debugw("Close connection")
-
+	c.logger.Debugw("Close connection")
 	c.cancel()
+
 	return c.conn.Close()
 }