diff --git a/faketls/client_protocol.go b/faketls/client_protocol.go index 9acdbd3..f0b82b2 100644 --- a/faketls/client_protocol.go +++ b/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) diff --git a/hub/connection.go b/hub/connection.go index c3edeb8..60ee738 100644 --- a/hub/connection.go +++ b/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 -func (c *connection) read() (conntypes.Packet, error) { - packet, err := c.conn.Read() - - c.mutex.Lock() - if err != nil { - c.pending-- - } else { - c.pending = 0 - } - c.mutex.Unlock() - - return packet, err -} - -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() - } - - return err -} - -func (c *connection) shutdown() { - c.shutdownOnce.Do(func() { - c.conn.Close() - close(c.done) - c.hub.channelBrokenSockets <- c.id - }) -} - -func (c *connection) closed() bool { - select { - case <-c.done: - return true - default: - return false - } -} - -func (c *connection) idle() bool { - c.mutex.RLock() - defer c.mutex.RUnlock() - - return c.pending == 0 + channelDone chan struct{} + channelWrite chan conntypes.Packet + channelRead chan *rpc.ProxyResponse + channelConnAttach chan *ProxyConn + channelConnDetach chan conntypes.ConnID } func (c *connection) run() { - logger := c.hub.logger.Named("connection").With("id", c.id) + defer c.Close() for { - packet, err := c.read() + select { + case <-c.channelDone: + for _, v := range c.proxyConns { + v.Close() + } + + 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) readLoop() { + for { + packet, err := c.conn.Read() if err != nil { - c.shutdown() + c.logger.Debugw("Cannot read packet", "error", err) + c.Close() + return } response, err := rpc.ParseProxyResponse(packet) if err != nil { - logger.Debugw("Failed response", "error", err) + c.logger.Debugw("Failed response", "error", err) continue } - if response.Type == rpc.ProxyResponseTypeCloseExt { - logger.Debugw("Proxy has closed connection") + select { + case <-c.channelDone: return - } - - if channel, ok := Registry.getChannel(response.ConnID); ok { - go channel.sendBack(response) // nolint: errcheck + case c.channelRead <- response: } } } -func newConnection(req *protocol.TelegramRequest, hub *connectionHub) (*connection, error) { +func (c *connection) Close() { + c.closeOnce.Do(func() { + c.logger.Debugw("Closing connection") + + close(c.channelDone) + c.conn.Close() + }) +} + +func (c *connection) Done() bool { + select { + case <-c.channelDone: + return true + default: + return c.Len() == 0 + } +} + +func (c *connection) Len() int { + c.proxyConnsMutex.RLock() + defer c.proxyConnsMutex.RUnlock() + + return len(c.proxyConns) +} + +func (c *connection) Attach(conn *ProxyConn) error { + select { + case <-c.channelDone: + return ErrClosed + case c.channelConnAttach <- conn: + return nil + } +} + +func (c *connection) Detach(connID conntypes.ConnID) { + select { + case <-c.channelDone: + case c.channelConnDetach <- connID: + } +} + +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 diff --git a/hub/connection_hub.go b/hub/connection_hub.go deleted file mode 100644 index 8ba9027..0000000 --- a/hub/connection_hub.go +++ /dev/null @@ -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 -} diff --git a/hub/connection_list.go b/hub/connection_list.go new file mode 100644 index 0000000..7bb4c48 --- /dev/null +++ b/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() + }) + } +} diff --git a/hub/ctx_channel.go b/hub/ctx_channel.go deleted file mode 100644 index e9e8b86..0000000 --- a/hub/ctx_channel.go +++ /dev/null @@ -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, - } -} diff --git a/hub/hub.go b/hub/hub.go index 57c2875..edfc442 100644 --- a/hub/hub.go +++ b/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 } diff --git a/hub/init.go b/hub/init.go index 35a0fee..19f5289 100644 --- a/hub/init.go +++ b/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 = ®istry{ - conns: map[string]*ctxChannel{}, - ctx: ctx, - } Hub = &hub{ - subs: map[string]*connectionHub{}, - logger: zap.S().Named("hub"), + muxes: make(map[int32]*mux), + ctx: ctx, } }) } diff --git a/hub/interface.go b/hub/interface.go new file mode 100644 index 0000000..a1efcdc --- /dev/null +++ b/hub/interface.go @@ -0,0 +1,7 @@ +package hub + +import "github.com/9seconds/mtg/protocol" + +type Interface interface { + Register(*protocol.TelegramRequest) (*ProxyConn, error) +} diff --git a/hub/mux.go b/hub/mux.go new file mode 100644 index 0000000..7aa82f6 --- /dev/null +++ b/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 +} diff --git a/hub/proxy_conn.go b/hub/proxy_conn.go new file mode 100644 index 0000000..a3e4686 --- /dev/null +++ b/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, + } +} diff --git a/hub/registry.go b/hub/registry.go deleted file mode 100644 index c8652a1..0000000 --- a/hub/registry.go +++ /dev/null @@ -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 -} diff --git a/main.go b/main.go index 62336ac..169f834 100644 --- a/main.go +++ b/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) diff --git a/obfuscated2/client_protocol.go b/obfuscated2/client_protocol.go index 9f940af..ec96807 100644 --- a/obfuscated2/client_protocol.go +++ b/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 } diff --git a/proxy/middle.go b/proxy/middle.go index d7d3d42..c0995dc 100644 --- a/proxy/middle.go +++ b/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 diff --git a/stats/interfaces.go b/stats/interfaces.go index 9a27580..0672982 100644 --- a/stats/interfaces.go +++ b/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 } diff --git a/stats/multi_stats.go b/stats/multi_stats.go index c67a4ef..f87ee75 100644 --- a/stats/multi_stats.go +++ b/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() } } diff --git a/stats/stats_prometheus.go b/stats/stats_prometheus.go index 69a3f6c..e2cefb2 100644 --- a/stats/stats_prometheus.go +++ b/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{}) diff --git a/stats/stats_statsd.go b/stats/stats_statsd.go index e68ede7..e6baa5d 100644 --- a/stats/stats_statsd.go +++ b/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) { diff --git a/wrappers/packetack/proxy.go b/wrappers/packetack/proxy.go index 768eddb..62ac7d2 100644 --- a/wrappers/packetack/proxy.go +++ b/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 {