Base prevention of replay attacks on proxy

antireplay/cache.go

+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

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

client/client.go

 	"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(context.Context, context.CancelFunc, net.Conn, string,
-	*config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error)
+	antireplay.Cache, *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error)

client/direct.go

 
 	"github.com/juju/errors"
 
+	"github.com/9seconds/mtg/antireplay"
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
 	"github.com/9seconds/mtg/obfuscated2"
 // DirectInit initializes client connection for proxy which connects to
 // Telegram directly.
 func DirectInit(ctx context.Context, cancel context.CancelFunc, socket net.Conn,
-	connID string, conf *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) {
+	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 != 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

 	"context"
 	"net"
 
+	"github.com/9seconds/mtg/antireplay"
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
 	"github.com/9seconds/mtg/wrappers"
 // MiddleInit initializes client connection for proxy which has to
 // support promoted channels, connect to Telegram middle proxies etc.
 func MiddleInit(ctx context.Context, cancel context.CancelFunc, socket net.Conn,
-	connID string, conf *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) {
-	conn, opts, err := DirectInit(ctx, cancel, 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
 	}

config/config.go

 	"fmt"
 	"net"
 	"strconv"
+	"time"
 
 	"github.com/juju/errors"
 	statsd "gopkg.in/alexcesaro/statsd.v2"
 	PublicIPv6 net.IP
 	StatsIP    net.IP
 
+	AntiReplayMaxSize      int
+	AntiReplayEvictionTime time.Duration
+
 	StatsD struct {
 		Addr       net.Addr
 		Prefix     string
 	statsdIP, statsdNetwork, statsdPrefix, statsdTagsFormat string,
 	statsdTags map[string]string, prometheusPrefix string,
 	secureOnly bool,
+	antiReplayMaxSize int, antiReplayEvictionTime time.Duration,
 	secret, adtag []byte) (*Config, error) {
 	secureMode := secureOnly
 	if bytes.HasPrefix(secret, []byte{0xdd}) && len(secret) == 17 {
 	}
 
 	conf := &Config{
-		Debug:           debug,
-		Verbose:         verbose,
-		SecureOnly:      secureOnly,
-		BindIP:          bindIP,
-		BindPort:        bindPort,
-		PublicIPv4:      publicIPv4,
-		PublicIPv4Port:  publicIPv4Port,
-		PublicIPv6:      publicIPv6,
-		PublicIPv6Port:  publicIPv6Port,
-		StatsIP:         statsIP,
-		StatsPort:       statsPort,
-		Secret:          secret,
-		AdTag:           adtag,
-		SecureMode:      secureMode,
-		ReadBufferSize:  int(readBufferSize),
-		WriteBufferSize: int(writeBufferSize),
+		Debug:                  debug,
+		Verbose:                verbose,
+		SecureOnly:             secureOnly,
+		BindIP:                 bindIP,
+		BindPort:               bindPort,
+		PublicIPv4:             publicIPv4,
+		PublicIPv4Port:         publicIPv4Port,
+		PublicIPv6:             publicIPv6,
+		PublicIPv6Port:         publicIPv6Port,
+		StatsIP:                statsIP,
+		StatsPort:              statsPort,
+		Secret:                 secret,
+		AdTag:                  adtag,
+		SecureMode:             secureMode,
+		ReadBufferSize:         int(readBufferSize),
+		WriteBufferSize:        int(writeBufferSize),
+		AntiReplayMaxSize:      antiReplayMaxSize,
+		AntiReplayEvictionTime: antiReplayEvictionTime,
 	}
 	conf.Prometheus.Prefix = prometheusPrefix
 

go.mod

 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/beorn7/perks v1.0.0 // indirect
+	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/golang/protobuf v1.3.1 // indirect
 	github.com/prometheus/client_model v0.0.0-20190129233127-fd36f4220a90 // indirect
 	github.com/prometheus/common v0.3.0 // indirect
 	github.com/prometheus/procfs v0.0.0-20190425082905-87a4384529e0 // indirect
+	github.com/spaolacci/murmur3 v1.1.0 // indirect
 	github.com/stretchr/testify v1.3.0
 	go.uber.org/atomic v1.3.2 // indirect
 	go.uber.org/multierr v1.1.0 // indirect

go.sum

+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/prometheus/procfs v0.0.0-20190425082905-87a4384529e0 h1:c8R11WC8m7KNMkTv/0+Be8vvwo4I3/Ut9AC2FW8fX3U=
 github.com/prometheus/procfs v0.0.0-20190425082905-87a4384529e0/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=

main.go

 		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()
 )
 		*bindPort, *publicIPv4Port, *publicIPv6Port, *statsPort, *statsdPort,
 		*statsdIP, *statsdNetwork, *statsdPrefix, *statsdTagsFormat,
 		*statsdTags, *prometheusPrefix, *secureOnly,
+		*antiReplayMaxSize, *antiReplayEvictionTime,
 		*secret, *adtag,
 	)
 	if 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)
 	}

proxy/proxy.go

 	"github.com/juju/errors"
 	"go.uber.org/zap"
 
+	"github.com/9seconds/mtg/antireplay"
 	"github.com/9seconds/mtg/client"
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
 
 // Proxy is a core of this program.
 type Proxy struct {
-	clientInit client.Init
-	tg         telegram.Telegram
-	conf       *config.Config
+	antiReplayCache antireplay.Cache
+	clientInit      client.Init
+	tg              telegram.Telegram
+	conf            *config.Config
 }
 
 // Serve runs TCP proxy server.
 
 	log.Infow("Client connected", "addr", conn.RemoteAddr())
 
-	clientConn, opts, err := p.clientInit(ctx, cancel, 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
 }
 
 // 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)
 	}
 
 	return &Proxy{
-		conf:       conf,
-		clientInit: clientInit,
-		tg:         tg,
-	}
+		antiReplayCache: cache,
+		conf:            conf,
+		clientInit:      clientInit,
+		tg:              tg,
+	}, nil
 }