Correct multiplexing

faketls/client_protocol.go

@@ -11,12 +11,10 @@ import (
 	"sync"
 	"time"
 
-	"github.com/9seconds/mtg/antireplay"
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/conntypes"
 	"github.com/9seconds/mtg/obfuscated2"
 	"github.com/9seconds/mtg/protocol"
-	"github.com/9seconds/mtg/stats"
 	"github.com/9seconds/mtg/tlstypes"
 	"github.com/9seconds/mtg/wrappers/stream"
 )
@@ -84,13 +82,6 @@ func (c *ClientProtocol) tlsHandshake(conn io.ReadWriter) error {
 		return errBadTime
 	}
 
-	if antireplay.Cache.HasTLS(clientHello.Random[:]) {
-		stats.Stats.AntiReplayDetected()
-		return errors.New("antireplay detected")
-	}
-
-	antireplay.Cache.AddTLS(clientHello.Random[:])
-
 	hostCert, err := connectionServerInstance.get()
 	if err != nil {
 		return fmt.Errorf("cannot get host certificate: %w", err)

hub/connection.go

@@ -5,6 +5,8 @@ import (
 	"math/rand"
 	"sync"
 
+	"go.uber.org/zap"
+
 	"github.com/9seconds/mtg/conntypes"
 	"github.com/9seconds/mtg/mtproto"
 	"github.com/9seconds/mtg/mtproto/rpc"
@@ -12,108 +14,149 @@ import (
 )
 
 type connection struct {
-	conn         conntypes.PacketReadWriteCloser
-	mutex        sync.RWMutex
-	shutdownOnce sync.Once
-	hub          *connectionHub
-	id           int
-	pending      uint
-	done         chan struct{}
+	conn            conntypes.PacketReadWriteCloser
+	proxyConns      map[string]*ProxyConn
+	closeOnce       sync.Once
+	proxyConnsMutex sync.RWMutex
+	id              int
+	logger          *zap.SugaredLogger
+
+	channelDone       chan struct{}
+	channelWrite      chan conntypes.Packet
+	channelRead       chan *rpc.ProxyResponse
+	channelConnAttach chan *ProxyConn
+	channelConnDetach chan conntypes.ConnID
 }
 
-func (c *connection) read() (conntypes.Packet, error) {
-	packet, err := c.conn.Read()
+func (c *connection) run() {
+	defer c.Close()
 
-	c.mutex.Lock()
-	if err != nil {
-		c.pending--
-	} else {
-		c.pending = 0
-	}
-	c.mutex.Unlock()
+	for {
+		select {
+		case <-c.channelDone:
+			for _, v := range c.proxyConns {
+				v.Close()
+			}
 
-	return packet, err
+			return
+		case resp := <-c.channelRead:
+			if channel, ok := c.proxyConns[string(resp.ConnID[:])]; ok {
+				if resp.Type == rpc.ProxyResponseTypeCloseExt {
+					channel.Close()
+				} else {
+					channel.put(resp)
+				}
+			}
+		case packet := <-c.channelWrite:
+			if err := c.conn.Write(packet); err != nil {
+				c.logger.Debugw("Cannot write packet", "error", err)
+				c.Close()
+			}
+		case conn := <-c.channelConnAttach:
+			c.proxyConnsMutex.Lock()
+			c.proxyConns[string(conn.req.ConnID[:])] = conn
+			c.proxyConnsMutex.Unlock()
+			conn.channelWrite = c.channelWrite
+		case connID := <-c.channelConnDetach:
+			if conn, ok := c.proxyConns[string(connID[:])]; ok {
+				c.proxyConnsMutex.Lock()
+				delete(c.proxyConns, string(connID[:]))
+				c.proxyConnsMutex.Unlock()
+				conn.Close()
+			}
+		}
+	}
 }
 
-func (c *connection) write(packet conntypes.Packet) error {
-	err := c.conn.Write(packet)
-	if err != nil {
-		// if we tried to write into a socket and it was broken, it is
-		// a time to reconsider the prescence of this socket at all.
-		//
-		// probably we need to remove it completely because it seems
-		// that connection is broken.
-		c.mutex.Lock()
-		c.pending = 0
-		c.mutex.Unlock()
-	}
+func (c *connection) readLoop() {
+	for {
+		packet, err := c.conn.Read()
+		if err != nil {
+			c.logger.Debugw("Cannot read packet", "error", err)
+			c.Close()
 
-	return err
+			return
+		}
+
+		response, err := rpc.ParseProxyResponse(packet)
+		if err != nil {
+			c.logger.Debugw("Failed response", "error", err)
+			continue
+		}
+
+		select {
+		case <-c.channelDone:
+			return
+		case c.channelRead <- response:
+		}
+	}
 }
 
-func (c *connection) shutdown() {
-	c.shutdownOnce.Do(func() {
+func (c *connection) Close() {
+	c.closeOnce.Do(func() {
+		c.logger.Debugw("Closing connection")
+
+		close(c.channelDone)
 		c.conn.Close()
-		close(c.done)
-		c.hub.channelBrokenSockets <- c.id
 	})
 }
 
-func (c *connection) closed() bool {
+func (c *connection) Done() bool {
 	select {
-	case <-c.done:
+	case <-c.channelDone:
 		return true
 	default:
-		return false
+		return c.Len() == 0
 	}
 }
 
-func (c *connection) idle() bool {
-	c.mutex.RLock()
-	defer c.mutex.RUnlock()
+func (c *connection) Len() int {
+	c.proxyConnsMutex.RLock()
+	defer c.proxyConnsMutex.RUnlock()
 
-	return c.pending == 0
+	return len(c.proxyConns)
 }
 
-func (c *connection) run() {
-	logger := c.hub.logger.Named("connection").With("id", c.id)
-
-	for {
-		packet, err := c.read()
-		if err != nil {
-			c.shutdown()
-			return
-		}
-
-		response, err := rpc.ParseProxyResponse(packet)
-		if err != nil {
-			logger.Debugw("Failed response", "error", err)
-			continue
-		}
-
-		if response.Type == rpc.ProxyResponseTypeCloseExt {
-			logger.Debugw("Proxy has closed connection")
-			return
-		}
+func (c *connection) Attach(conn *ProxyConn) error {
+	select {
+	case <-c.channelDone:
+		return ErrClosed
+	case c.channelConnAttach <- conn:
+		return nil
+	}
+}
 
-		if channel, ok := Registry.getChannel(response.ConnID); ok {
-			go channel.sendBack(response) // nolint: errcheck
-		}
+func (c *connection) Detach(connID conntypes.ConnID) {
+	select {
+	case <-c.channelDone:
+	case c.channelConnDetach <- connID:
 	}
 }
 
-func newConnection(req *protocol.TelegramRequest, hub *connectionHub) (*connection, error) {
+func newConnection(req *protocol.TelegramRequest) (*connection, error) {
 	conn, err := mtproto.TelegramProtocol(req)
 	if err != nil {
 		return nil, fmt.Errorf("cannot create a new connection: %w", err)
 	}
 
+	id := rand.Int() // nolint: gosec
 	rv := &connection{
 		conn: conn,
-		hub:  hub,
-		id:   rand.Int(), // nolint: gosec
-		done: make(chan struct{}),
+		id:   id,
+		logger: zap.S().Named("hub-connection").With("id", id,
+			"dc", req.ClientProtocol.DC(),
+			"protocol", req.ClientProtocol.ConnectionProtocol()),
+		proxyConns: make(map[string]*ProxyConn),
+
+		channelRead:       make(chan *rpc.ProxyResponse, 1),
+		channelDone:       make(chan struct{}),
+		channelWrite:      make(chan conntypes.Packet),
+		channelConnAttach: make(chan *ProxyConn),
+		channelConnDetach: make(chan conntypes.ConnID),
 	}
+
+	go rv.readLoop()
+
 	go rv.run()
 
 	return rv, nil

hub/connection_hub.go

@@ -1,114 +0,0 @@
-package hub
-
-import (
-	"time"
-
-	"go.uber.org/zap"
-
-	"github.com/9seconds/mtg/protocol"
-)
-
-const hubGCEvery = time.Minute
-
-type connectionHubRequest struct {
-	request  *protocol.TelegramRequest
-	response chan<- *connection
-}
-
-type connectionHub struct {
-	sockets map[int]*connection
-	logger  *zap.SugaredLogger
-
-	channelBrokenSockets      chan int
-	channelConnectionRequests chan *connectionHubRequest
-	channelReturnConnections  chan *connection
-}
-
-func (c *connectionHub) run() {
-	ticker := time.NewTicker(hubGCEvery)
-	defer ticker.Stop()
-
-	for {
-		select {
-		case <-ticker.C:
-			c.runGC()
-		case request := <-c.channelConnectionRequests:
-			c.runConnectionRequest(request)
-		case id := <-c.channelBrokenSockets:
-			c.runBrokenSocket(id)
-		case conn := <-c.channelReturnConnections:
-			c.runReturnConnection(conn)
-		}
-	}
-}
-
-func (c *connectionHub) runGC() {
-	logger := c.logger.Named("gc")
-
-	for key, conn := range c.sockets {
-		switch {
-		case conn.closed():
-			logger.Debugw("Delete closed socket", "key", key)
-			delete(c.sockets, key)
-		case conn.idle():
-			logger.Debugw("Delete idle socket", "key", key)
-			conn.shutdown()
-			delete(c.sockets, key)
-
-			return
-		}
-	}
-}
-
-func (c *connectionHub) runConnectionRequest(req *connectionHubRequest) {
-	logger := c.logger.Named("request").With("connection-id", req.request.ConnID)
-
-	for key, conn := range c.sockets {
-		delete(c.sockets, key)
-
-		if !conn.closed() {
-			logger.Debugw("Choose connection",
-				"id", conn.id,
-				"remote_addr", conn.conn.RemoteAddr())
-			req.response <- conn
-			close(req.response)
-
-			return
-		}
-	}
-
-	if conn, err := newConnection(req.request, c); err == nil {
-		logger.Debugw("New connection",
-			"id", conn.id,
-			"remote_addr", conn.conn.RemoteAddr())
-		req.response <- conn
-	}
-
-	close(req.response)
-}
-
-func (c *connectionHub) runBrokenSocket(id int) {
-	c.logger.Named("broken-socket").Debugw("Delete broken socket", "id", id)
-	delete(c.sockets, id)
-}
-
-func (c *connectionHub) runReturnConnection(conn *connection) {
-	c.logger.Named("return-connection").Debugw("Return connection",
-		"id", conn.id,
-		"remote_addr", conn.conn.RemoteAddr())
-
-	c.sockets[conn.id] = conn
-}
-
-func newConnectionHub(logger *zap.SugaredLogger) *connectionHub {
-	rv := &connectionHub{
-		logger:                    logger.Named("connection-hub"),
-		sockets:                   map[int]*connection{},
-		channelBrokenSockets:      make(chan int, 1),
-		channelConnectionRequests: make(chan *connectionHubRequest),
-		channelReturnConnections:  make(chan *connection, 1),
-	}
-	go rv.run()
-
-	return rv
-}

hub/connection_list.go

@@ -0,0 +1,70 @@
+package hub
+
+import (
+	"fmt"
+	"sort"
+)
+
+const connectionListMaxClientsPerConnection = 2
+
+type connectionList struct {
+	connections []*connection
+}
+
+func (c *connectionList) Get(conn *ProxyConn) (*connection, error) {
+	if len(c.connections) > 0 {
+		c.gc()
+	}
+
+	if len(c.connections) > 0 && c.connections[0].Len() < connectionListMaxClientsPerConnection {
+		if err := c.connections[0].Attach(conn); err == nil {
+			return c.connections[0], nil
+		}
+	}
+
+	newConn, err := newConnection(conn.req)
+	if err != nil {
+		return nil, fmt.Errorf("cannot allocate a new connection: %w", err)
+	}
+
+	if err = newConn.Attach(conn); err != nil {
+		newConn.Close()
+		return nil, fmt.Errorf("cannot attach to the newly created connection: %w", err)
+	}
+
+	c.connections = append(c.connections, newConn)
+	lastIndex := len(c.connections) - 1
+	c.connections[0], c.connections[lastIndex] = c.connections[lastIndex], c.connections[0]
+
+	return newConn, nil
+}
+
+func (c *connectionList) gc() {
+	prevLen := len(c.connections)
+
+	for i := len(c.connections) - 1; i >= 0; i-- {
+		lastIndex := len(c.connections) - 1
+
+		if c.connections[i].Done() {
+			c.connections[i].Close()
+
+			if len(c.connections)-1 == i {
+				c.connections = c.connections[:lastIndex]
+			} else {
+				c.connections[i], c.connections[lastIndex] = c.connections[lastIndex], c.connections[i]
+			}
+		}
+	}
+
+	if prevLen != len(c.connections) {
+		c.sort()
+	}
+}
+
+func (c *connectionList) sort() {
+	if len(c.connections) > 1 {
+		sort.Slice(c.connections, func(i, j int) bool {
+			return c.connections[i].Len() < c.connections[j].Len()
+		})
+	}
+}

hub/ctx_channel.go

@@ -1,61 +0,0 @@
-package hub
-
-import (
-	"context"
-	"time"
-
-	"github.com/9seconds/mtg/mtproto/rpc"
-)
-
-const closeableChannelReadTimeout = 2 * time.Minute
-
-type ChannelReadCloser interface {
-	Read() (*rpc.ProxyResponse, error)
-	Close() error
-}
-
-type ctxChannel struct {
-	channel chan *rpc.ProxyResponse
-	ctx     context.Context
-	cancel  context.CancelFunc
-}
-
-func (c *ctxChannel) Read() (*rpc.ProxyResponse, error) {
-	timer := time.NewTimer(closeableChannelReadTimeout)
-	defer timer.Stop()
-
-	select {
-	case <-timer.C:
-		return nil, ErrTimeout
-	case <-c.ctx.Done():
-		return nil, ErrClosed
-	case packet := <-c.channel:
-		return packet, nil
-	}
-}
-
-func (c *ctxChannel) sendBack(response *rpc.ProxyResponse) error {
-	select {
-	case <-c.ctx.Done():
-		return ErrClosed
-	case c.channel <- response:
-		return nil
-	}
-}
-
-func (c *ctxChannel) Close() error {
-	c.cancel()
-	c.channel = nil
-
-	return nil
-}
-
-func newCtxChannel(ctx context.Context) *ctxChannel {
-	ctx, cancel := context.WithCancel(ctx)
-
-	return &ctxChannel{
-		channel: make(chan *rpc.ProxyResponse),
-		ctx:     ctx,
-		cancel:  cancel,
-	}
-}

hub/hub.go

@@ -1,73 +1,40 @@
 package hub
 
 import (
-	"encoding/binary"
-	"fmt"
-	"strings"
+	"context"
 	"sync"
 
-	"go.uber.org/zap"
-
-	"github.com/9seconds/mtg/conntypes"
 	"github.com/9seconds/mtg/protocol"
 )
 
 type hub struct {
-	logger *zap.SugaredLogger
-	subs   map[string]*connectionHub
-	mutex  sync.RWMutex
+	muxes map[int32]*mux
+	mutex sync.RWMutex
+	ctx   context.Context
 }
 
-func (h *hub) Write(packet conntypes.Packet, req *protocol.TelegramRequest) error {
-	sub := h.getHub(req)
-	connections := make(chan *connection)
-	sub.channelConnectionRequests <- &connectionHubRequest{
-		request:  req,
-		response: connections,
-	}
-
-	conn, ok := <-connections
-	if !ok {
-		return ErrCannotCreateConnection
-	}
-
-	if err := conn.write(packet); err != nil {
-		conn.shutdown()
-		return fmt.Errorf("cannot send packet: %w", err)
-	}
-	sub.channelReturnConnections <- conn
-
-	return nil
+func (h *hub) Register(req *protocol.TelegramRequest) (*ProxyConn, error) {
+	return h.getMux(req).Get(req)
 }
 
-func (h *hub) getHub(req *protocol.TelegramRequest) *connectionHub {
-	keyBuilder := strings.Builder{}
-	binary.Write(&keyBuilder, binary.LittleEndian, int16(req.ClientProtocol.DC())) // nolint: errcheck
-	keyBuilder.WriteRune('_')
-	binary.Write(&keyBuilder, binary.LittleEndian, uint8(req.ClientProtocol.ConnectionProtocol())) // nolint: errcheck
-	key := keyBuilder.String()
+func (h *hub) getMux(req *protocol.TelegramRequest) *mux {
+	var key int32 = 32767 + int32(req.ClientProtocol.DC()) + 100000*int32(req.ClientProtocol.ConnectionProtocol())
 
 	h.mutex.RLock()
-	rv, ok := h.subs[key]
+	m, ok := h.muxes[key]
 	h.mutex.RUnlock()
 
 	if !ok {
 		h.mutex.Lock()
-		defer h.mutex.Unlock()
+		m, ok = h.muxes[key]
 
-		rv, ok = h.subs[key]
 		if !ok {
-			h.logger.Debugw("Create new connection hub",
-				"dc", req.ClientProtocol.DC(),
-				"protocol", req.ClientProtocol.ConnectionProtocol())
-
-			rv = newConnectionHub(h.logger.With(
-				"dc", req.ClientProtocol.DC(),
-				"protocol", req.ClientProtocol.ConnectionProtocol(),
-			))
-			h.subs[key] = rv
+			m = newMux(h.ctx)
+			h.muxes[key] = m
 		}
+
+		h.mutex.Unlock()
 	}
 
-	return rv
+	return m
 }

hub/init.go

@@ -4,30 +4,21 @@ import (
 	"context"
 	"errors"
 	"sync"
-
-	"go.uber.org/zap"
 )
 
 var (
-	Registry *registry
-	Hub      *hub
-
-	ErrTimeout                = errors.New("timeout")
-	ErrClosed                 = errors.New("channel was closed")
-	ErrCannotCreateConnection = errors.New("cannot create connection")
+	ErrTimeout = errors.New("timeout")
+	ErrClosed  = errors.New("context is closed")
 
+	Hub      Interface
 	initOnce sync.Once
 )
 
 func Init(ctx context.Context) {
 	initOnce.Do(func() {
-		Registry = &registry{
-			conns: map[string]*ctxChannel{},
-			ctx:   ctx,
-		}
 		Hub = &hub{
-			subs:   map[string]*connectionHub{},
-			logger: zap.S().Named("hub"),
+			muxes: make(map[int32]*mux),
+			ctx:   ctx,
 		}
 	})
 }

hub/interface.go

@@ -0,0 +1,7 @@
+package hub
+
+import "github.com/9seconds/mtg/protocol"
+
+type Interface interface {
+	Register(*protocol.TelegramRequest) (*ProxyConn, error)
+}

hub/mux.go

@@ -0,0 +1,80 @@
+package hub
+
+import (
+	"context"
+
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/protocol"
+)
+
+type muxNewRequest struct {
+	req  *protocol.TelegramRequest
+	resp chan<- muxNewResponse
+}
+
+type muxNewResponse struct {
+	conn *ProxyConn
+	err  error
+}
+
+type mux struct {
+	connections   connectionList
+	clients       map[string]*connection
+	ctx           context.Context
+	channelClosed chan conntypes.ConnID
+	channelNew    chan muxNewRequest
+}
+
+func (m *mux) run() {
+	for {
+		select {
+		case <-m.ctx.Done():
+			for _, v := range m.clients {
+				v.Close()
+			}
+
+			return
+		case req := <-m.channelNew:
+			proxyConn := newProxyConn(req.req, m.channelClosed)
+			conn, err := m.connections.Get(proxyConn)
+
+			if err == nil {
+				m.clients[string(req.req.ConnID[:])] = conn
+			}
+
+			req.resp <- muxNewResponse{
+				conn: proxyConn,
+				err:  err,
+			}
+			close(req.resp)
+		case connID := <-m.channelClosed:
+			if conn, ok := m.clients[string(connID[:])]; ok {
+				conn.Detach(connID)
+				delete(m.clients, string(connID[:]))
+			}
+		}
+	}
+}
+
+func (m *mux) Get(req *protocol.TelegramRequest) (*ProxyConn, error) {
+	resp := make(chan muxNewResponse)
+	m.channelNew <- muxNewRequest{
+		req:  req,
+		resp: resp,
+	}
+	rv := <-resp
+
+	return rv.conn, rv.err
+}
+
+func newMux(ctx context.Context) *mux {
+	m := &mux{
+		ctx:           ctx,
+		clients:       make(map[string]*connection),
+		channelClosed: make(chan conntypes.ConnID, 1),
+		channelNew:    make(chan muxNewRequest),
+	}
+	go m.run()
+
+	return m
+}

hub/proxy_conn.go

@@ -0,0 +1,77 @@
+package hub
+
+import (
+	"sync"
+	"time"
+
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/mtproto/rpc"
+	"github.com/9seconds/mtg/protocol"
+)
+
+const (
+	proxyConnWriteTimeout = 2 * time.Minute
+	proxyConnReadTimeout  = 2 * time.Minute
+)
+
+type ProxyConn struct {
+	closeOnce       sync.Once
+	req             *protocol.TelegramRequest
+	channelResponse chan *rpc.ProxyResponse
+	channelClosed   chan<- conntypes.ConnID
+	channelWrite    chan<- conntypes.Packet
+	channelDone     chan struct{}
+}
+
+func (p *ProxyConn) Read() (*rpc.ProxyResponse, error) {
+	timer := time.NewTimer(proxyConnReadTimeout)
+	defer timer.Stop()
+
+	select {
+	case <-timer.C:
+		return nil, ErrTimeout
+	case <-p.channelDone:
+		return nil, ErrClosed
+	case packet := <-p.channelResponse:
+		return packet, nil
+	}
+}
+
+func (p *ProxyConn) Write(packet conntypes.Packet) error {
+	timer := time.NewTimer(proxyConnWriteTimeout)
+	defer timer.Stop()
+
+	select {
+	case <-timer.C:
+		return ErrTimeout
+	case <-p.channelDone:
+		return ErrClosed
+	case p.channelWrite <- packet:
+		return nil
+	}
+}
+
+func (p *ProxyConn) put(response *rpc.ProxyResponse) {
+	select {
+	case <-p.channelDone:
+	case p.channelResponse <- response:
+	}
+}
+
+func (p *ProxyConn) Close() {
+	p.closeOnce.Do(func() {
+		close(p.channelDone)
+		go func() {
+			p.channelClosed <- p.req.ConnID
+		}()
+	})
+}
+
+func newProxyConn(req *protocol.TelegramRequest, channelClosed chan<- conntypes.ConnID) *ProxyConn {
+	return &ProxyConn{
+		channelResponse: make(chan *rpc.ProxyResponse),
+		channelDone:     make(chan struct{}),
+		channelClosed:   channelClosed,
+		req:             req,
+	}
+}

hub/registry.go

@@ -1,45 +0,0 @@
-package hub
-
-import (
-	"context"
-	"sync"
-
-	"github.com/9seconds/mtg/conntypes"
-)
-
-type registry struct {
-	conns map[string]*ctxChannel
-	ctx   context.Context
-	mutex sync.RWMutex
-}
-
-func (r *registry) Register(id conntypes.ConnID) ChannelReadCloser {
-	channel := newCtxChannel(r.ctx)
-
-	r.mutex.Lock()
-	r.conns[string(id[:])] = channel
-	r.mutex.Unlock()
-
-	return channel
-}
-
-func (r *registry) Unregister(id conntypes.ConnID) {
-	r.mutex.Lock()
-	defer r.mutex.Unlock()
-
-	if channel, ok := r.conns[string(id[:])]; ok {
-		channel.Close()
-		delete(r.conns, string(id[:]))
-	}
-}
-
-func (r *registry) getChannel(id conntypes.ConnID) (*ctxChannel, bool) {
-	r.mutex.RLock()
-	defer r.mutex.RUnlock()
-
-	if value, ok := r.conns[string(id[:])]; ok {
-		return value, true
-	}
-
-	return nil, false
-}

main.go

@@ -29,92 +29,92 @@ var (
 		Required().
 		Enum("simple", "secured", "tls")
 
-	proxyCommand = app.Command("proxy",
+	runCommand = app.Command("run",
 		"Run new proxy instance")
-	proxyDebug = proxyCommand.Flag("debug",
+	runDebug = runCommand.Flag("debug",
 		"Run in debug mode.").
 		Short('d').
 		Envar("MTG_DEBUG").
 		Bool()
-	proxyVerbose = proxyCommand.Flag("verbose",
+	runVerbose = runCommand.Flag("verbose",
 		"Run in verbose mode.").
 		Short('v').
 		Envar("MTG_VERBOSE").
 		Bool()
-	proxyBind = proxyCommand.Flag("bind",
+	runBind = runCommand.Flag("bind",
 		"Host:Port to bind proxy to.").
 		Short('b').
 		Envar("MTG_BIND").
 		Default("0.0.0.0:3128").
 		TCP()
-	proxyPublicIPv4 = proxyCommand.Flag("public-ipv4",
+	runPublicIPv4 = runCommand.Flag("public-ipv4",
 		"Which IPv4 host:port to use.").
 		Short('4').
 		Envar("MTG_IPV4").
 		TCP()
-	proxyPublicIPv6 = proxyCommand.Flag("public-ipv6",
+	runPublicIPv6 = runCommand.Flag("public-ipv6",
 		"Which IPv6 host:port to use.").
 		Short('6').
 		Envar("MTG_IPV6").
 		TCP()
-	proxyStatsBind = proxyCommand.Flag("stats-bind",
+	runStatsBind = runCommand.Flag("stats-bind",
 		"Which Host:Port to bind stats server to.").
 		Short('t').
 		Envar("MTG_STATS_BIND").
 		Default("127.0.0.1:3129").
 		TCP()
-	proxyStatsNamespace = proxyCommand.Flag("stats-namespace",
+	runStatsNamespace = runCommand.Flag("stats-namespace",
 		"Which namespace to use for Prometheus.").
 		Envar("MTG_STATS_NAMESPACE").
 		Default("mtg").
 		String()
-	proxyStatsdAddress = proxyCommand.Flag("statsd-addr",
+	runStatsdAddress = runCommand.Flag("statsd-addr",
 		"Host:port of statsd server").
 		Envar("MTG_STATSD_ADDR").
 		TCP()
-	proxyStatsdNetwork = proxyCommand.Flag("statsd-network",
+	runStatsdNetwork = runCommand.Flag("statsd-network",
 		"Which network is used to work with statsd. Only 'tcp' and 'udp' are supported.").
 		Envar("MTG_STATSD_NETWORK").
 		Default("udp").
 		Enum("udp", "tcp")
-	proxyStatsdTagsFormat = proxyCommand.Flag("statsd-tags-format",
+	runStatsdTagsFormat = runCommand.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").
 		Default("influxdb").
 		Enum("datadog", "influxdb")
-	proxyStatsdTags = proxyCommand.Flag("statsd-tags",
+	runStatsdTags = runCommand.Flag("statsd-tags",
 		"Tags to use for working with statsd (specified as 'key=value').").
 		Envar("MTG_STATSD_TAGS").
 		StringMap()
-	proxyWriteBufferSize = proxyCommand.Flag("write-buffer",
+	runWriteBufferSize = runCommand.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("65536KB").
 		Bytes()
-	proxyReadBufferSize = proxyCommand.Flag("read-buffer",
+	runReadBufferSize = runCommand.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("131072KB").
 		Bytes()
-	proxyTLSCloakPort = proxyCommand.Flag("cloak-port",
+	runTLSCloakPort = runCommand.Flag("cloak-port",
 		"Port which should be used for host cloaking.").
 		Envar("MTG_CLOAK_PORT").
 		Default("443").
 		Uint16()
-	proxyAntiReplayMaxSize = proxyCommand.Flag("anti-replay-max-size",
+	runAntiReplayMaxSize = runCommand.Flag("anti-replay-max-size",
 		"Max size of antireplay cache in megabytes.").
 		Envar("MTG_ANTIREPLAY_MAXSIZE").
 		Default("128").
 		Int()
-	proxyAntiReplayEvictionTime = proxyCommand.Flag("anti-replay-eviction-time",
+	runAntiReplayEvictionTime = runCommand.Flag("anti-replay-eviction-time",
 		"Eviction time period for obfuscated2 handshakes").
 		Envar("MTG_ANTIREPLAY_EVICTIONTIME").
 		Default("168h").
 		Duration()
-	proxySecret = proxyCommand.Arg("secret", "Secret of this proxy.").Required().HexBytes()
-	proxyAdtag  = proxyCommand.Arg("adtag", "ADTag of the proxy.").HexBytes()
+	runSecret = runCommand.Arg("secret", "Secret of this proxy.").Required().HexBytes()
+	runAdtag  = runCommand.Arg("adtag", "ADTag of the proxy.").HexBytes()
 )
 
 func main() {
@@ -129,26 +129,26 @@ func main() {
 	switch kingpin.MustParse(app.Parse(os.Args[1:])) {
 	case generateSecretCommand.FullCommand():
 		cli.Generate(*generateSecretType, *generateCloakHost)
-	case proxyCommand.FullCommand():
+	case runCommand.FullCommand():
 		err := config.Init(
-			config.Opt{Option: config.OptionTypeDebug, Value: *proxyDebug},
-			config.Opt{Option: config.OptionTypeVerbose, Value: *proxyVerbose},
-			config.Opt{Option: config.OptionTypeBind, Value: *proxyBind},
-			config.Opt{Option: config.OptionTypePublicIPv4, Value: *proxyPublicIPv4},
-			config.Opt{Option: config.OptionTypePublicIPv6, Value: *proxyPublicIPv6},
-			config.Opt{Option: config.OptionTypeStatsBind, Value: *proxyStatsBind},
-			config.Opt{Option: config.OptionTypeStatsNamespace, Value: *proxyStatsNamespace},
-			config.Opt{Option: config.OptionTypeStatsdAddress, Value: *proxyStatsdAddress},
-			config.Opt{Option: config.OptionTypeStatsdNetwork, Value: *proxyStatsdNetwork},
-			config.Opt{Option: config.OptionTypeStatsdTagsFormat, Value: *proxyStatsdTagsFormat},
-			config.Opt{Option: config.OptionTypeStatsdTags, Value: *proxyStatsdTags},
-			config.Opt{Option: config.OptionTypeWriteBufferSize, Value: *proxyWriteBufferSize},
-			config.Opt{Option: config.OptionTypeReadBufferSize, Value: *proxyReadBufferSize},
-			config.Opt{Option: config.OptionTypeCloakPort, Value: *proxyTLSCloakPort},
-			config.Opt{Option: config.OptionTypeAntiReplayMaxSize, Value: *proxyAntiReplayMaxSize},
-			config.Opt{Option: config.OptionTypeAntiReplayEvictionTime, Value: *proxyAntiReplayEvictionTime},
-			config.Opt{Option: config.OptionTypeSecret, Value: *proxySecret},
-			config.Opt{Option: config.OptionTypeAdtag, Value: *proxyAdtag},
+			config.Opt{Option: config.OptionTypeDebug, Value: *runDebug},
+			config.Opt{Option: config.OptionTypeVerbose, Value: *runVerbose},
+			config.Opt{Option: config.OptionTypeBind, Value: *runBind},
+			config.Opt{Option: config.OptionTypePublicIPv4, Value: *runPublicIPv4},
+			config.Opt{Option: config.OptionTypePublicIPv6, Value: *runPublicIPv6},
+			config.Opt{Option: config.OptionTypeStatsBind, Value: *runStatsBind},
+			config.Opt{Option: config.OptionTypeStatsNamespace, Value: *runStatsNamespace},
+			config.Opt{Option: config.OptionTypeStatsdAddress, Value: *runStatsdAddress},
+			config.Opt{Option: config.OptionTypeStatsdNetwork, Value: *runStatsdNetwork},
+			config.Opt{Option: config.OptionTypeStatsdTagsFormat, Value: *runStatsdTagsFormat},
+			config.Opt{Option: config.OptionTypeStatsdTags, Value: *runStatsdTags},
+			config.Opt{Option: config.OptionTypeWriteBufferSize, Value: *runWriteBufferSize},
+			config.Opt{Option: config.OptionTypeReadBufferSize, Value: *runReadBufferSize},
+			config.Opt{Option: config.OptionTypeCloakPort, Value: *runTLSCloakPort},
+			config.Opt{Option: config.OptionTypeAntiReplayMaxSize, Value: *runAntiReplayMaxSize},
+			config.Opt{Option: config.OptionTypeAntiReplayEvictionTime, Value: *runAntiReplayEvictionTime},
+			config.Opt{Option: config.OptionTypeSecret, Value: *runSecret},
+			config.Opt{Option: config.OptionTypeAdtag, Value: *runAdtag},
 		)
 		if err != nil {
 			cli.Fatal(err)

obfuscated2/client_protocol.go

@@ -81,13 +81,13 @@ func (c *ClientProtocol) Handshake(socket conntypes.StreamReadWriteCloser) (conn
 		c.dc = conntypes.DCDefaultIdx
 	}
 
-	antiReplayKey := decryptedFrame.Unique()
-	if antireplay.Cache.HasObfuscated2(antiReplayKey) {
-		stats.Stats.AntiReplayDetected()
+	replayKeys := decryptedFrame.Unique()
+	if antireplay.Cache.HasObfuscated2(replayKeys) {
+		stats.Stats.ReplayDetected()
 		return nil, errors.New("replay attack is detected")
 	}
 
-	antireplay.Cache.AddObfuscated2(antiReplayKey)
+	antireplay.Cache.AddObfuscated2(replayKeys)
 
 	return stream.NewObfuscated2(socket, encryptor, decryptor), nil
 }

proxy/middle.go

@@ -11,7 +11,11 @@ import (
 )
 
 func middleConnection(request *protocol.TelegramRequest) {
-	telegramConn := packetack.NewProxy(request)
+	telegramConn, err := packetack.NewProxy(request)
+	if err != nil {
+		request.Logger.Debugw("Cannot dial to Telegram", "error", err)
+		return
+	}
 	defer telegramConn.Close()
 
 	var clientConn conntypes.PacketAckFullReadWriteCloser

stats/interfaces.go

@@ -34,8 +34,8 @@ type CrashInterface interface {
 	Crash()
 }
 
-type AntiReplayDetectedInterface interface {
-	AntiReplayDetected()
+type ReplayDetectedInterface interface {
+	ReplayDetected()
 }
 
 type Interface interface {
@@ -46,5 +46,5 @@ type Interface interface {
 	TelegramConnectedInterface
 	TelegramDisconnectedInterface
 	CrashInterface
-	AntiReplayDetectedInterface
+	ReplayDetectedInterface
 }

stats/multi_stats.go

@@ -50,8 +50,8 @@ func (m multiStats) Crash() {
 	}
 }
 
-func (m multiStats) AntiReplayDetected() {
+func (m multiStats) ReplayDetected() {
 	for i := range m {
-		go m[i].AntiReplayDetected()
+		go m[i].ReplayDetected()
 	}
 }

stats/stats_prometheus.go

@@ -18,7 +18,7 @@ type statsPrometheus struct {
 	telegramConnections *prometheus.GaugeVec
 	traffic             *prometheus.GaugeVec
 	crashes             prometheus.Gauge
-	antiReplays         prometheus.Counter
+	replayAttacks       prometheus.Counter
 }
 
 func (s *statsPrometheus) IngressTraffic(traffic int) {
@@ -84,8 +84,8 @@ func (s *statsPrometheus) Crash() {
 	s.crashes.Inc()
 }
 
-func (s *statsPrometheus) AntiReplayDetected() {
-	s.antiReplays.Inc()
+func (s *statsPrometheus) ReplayDetected() {
+	s.replayAttacks.Inc()
 }
 
 func newStatsPrometheus(mux *http.ServeMux) (Interface, error) {
@@ -112,10 +112,10 @@ func newStatsPrometheus(mux *http.ServeMux) (Interface, error) {
 			Name:      "crashes",
 			Help:      "How many crashes happened.",
 		}),
-		antiReplays: prometheus.NewCounter(prometheus.CounterOpts{
+		replayAttacks: prometheus.NewCounter(prometheus.CounterOpts{
 			Namespace: config.C.StatsNamespace,
-			Name:      "anti_replays",
-			Help:      "How many anti replay attacks were prevented.",
+			Name:      "replay_attacks",
+			Help:      "How many replay attacks were prevented.",
 		}),
 	}
 
@@ -135,8 +135,8 @@ func newStatsPrometheus(mux *http.ServeMux) (Interface, error) {
 		return nil, fmt.Errorf("cannot register metrics for crashes: %w", err)
 	}
 
-	if err := registry.Register(instance.antiReplays); err != nil {
-		return nil, fmt.Errorf("cannot register metrics for anti replays: %w", err)
+	if err := registry.Register(instance.replayAttacks); err != nil {
+		return nil, fmt.Errorf("cannot register metrics for replays: %w", err)
 	}
 
 	handler := promhttp.HandlerFor(registry, promhttp.HandlerOpts{})

stats/stats_statsd.go

@@ -79,8 +79,8 @@ func (s *statsStatsd) Crash() {
 	s.client.Increment("crashes")
 }
 
-func (s *statsStatsd) AntiReplayDetected() {
-	s.client.Increment("anti_replays")
+func (s *statsStatsd) ReplayDetected() {
+	s.client.Increment("replay_attacks")
 }
 
 func newStatsStatsd() (Interface, error) {

wrappers/packetack/proxy.go

@@ -5,7 +5,6 @@ import (
 	"encoding/binary"
 	"fmt"
 	"net"
-	"sync"
 
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/conntypes"
@@ -16,10 +15,9 @@ import (
 
 type wrapperProxy struct {
 	request      *protocol.TelegramRequest
+	proxy        *hub.ProxyConn
 	clientIPPort []byte
 	ourIPPort    []byte
-	channelRead  hub.ChannelReadCloser
-	closeOnce    sync.Once
 	flags        rpc.ProxyRequestFlags
 }
 
@@ -47,11 +45,11 @@ func (w *wrapperProxy) Write(packet conntypes.Packet, acks *conntypes.Connection
 	buf.Write(make([]byte, (4-buf.Len()%4)%4))
 	buf.Write(packet)
 
-	return hub.Hub.Write(buf.Bytes(), w.request)
+	return w.proxy.Write(buf.Bytes())
 }
 
 func (w *wrapperProxy) Read(acks *conntypes.ConnectionAcks) (conntypes.Packet, error) {
-	resp, err := w.channelRead.Read()
+	resp, err := w.proxy.Read()
 	if err != nil {
 		return nil, fmt.Errorf("cannot read a response: %w", err)
 	}
@@ -64,15 +62,11 @@ func (w *wrapperProxy) Read(acks *conntypes.ConnectionAcks) (conntypes.Packet, e
 }
 
 func (w *wrapperProxy) Close() error {
-	w.closeOnce.Do(func() {
-		w.channelRead.Close()
-		hub.Registry.Unregister(w.request.ConnID)
-	})
-
+	w.proxy.Close()
 	return nil
 }
 
-func NewProxy(request *protocol.TelegramRequest) conntypes.PacketAckReadWriteCloser {
+func NewProxy(request *protocol.TelegramRequest) (conntypes.PacketAckReadWriteCloser, error) {
 	flags := rpc.ProxyRequestFlagsHasAdTag | rpc.ProxyRequestFlagsMagic | rpc.ProxyRequestFlagsExtMode2
 
 	switch request.ClientProtocol.ConnectionType() {
@@ -86,13 +80,18 @@ func NewProxy(request *protocol.TelegramRequest) conntypes.PacketAckReadWriteClo
 		panic("unknown connection type")
 	}
 
+	proxy, err := hub.Hub.Register(request)
+	if err != nil {
+		return nil, fmt.Errorf("cannot make a new proxy wrapper: %w", err)
+	}
+
 	return &wrapperProxy{
 		flags:        flags,
 		request:      request,
-		channelRead:  hub.Registry.Register(request.ConnID),
+		proxy:        proxy,
 		clientIPPort: proxyGetIPPort(request.ClientConn.RemoteAddr()),
 		ourIPPort:    proxyGetIPPort(request.ClientConn.LocalAddr()),
-	}
+	}, nil
 }
 
 func proxyGetIPPort(addr *net.TCPAddr) []byte {