diff --git a/conntypes/packet.go b/conntypes/packet.go new file mode 100644 index 0000000..fa72022 --- /dev/null +++ b/conntypes/packet.go @@ -0,0 +1,3 @@ +package conntypes + +type Packet []byte diff --git a/wrappers/interfaces.go b/conntypes/wrappers.go similarity index 97% rename from wrappers/interfaces.go rename to conntypes/wrappers.go index fd06f59..b26e245 100644 --- a/wrappers/interfaces.go +++ b/conntypes/wrappers.go @@ -1,4 +1,4 @@ -package wrappers +package conntypes import ( "io" @@ -8,8 +8,6 @@ import ( "go.uber.org/zap" ) -type Packet []byte - // Wrap is a base interface for all wrappers in this package. type Wrap interface { Conn() net.Conn diff --git a/go.mod b/go.mod index abc3da9..2179870 100644 --- a/go.mod +++ b/go.mod @@ -9,6 +9,8 @@ require ( github.com/allegro/bigcache v1.2.1 github.com/beevik/ntp v0.2.0 github.com/cespare/xxhash v1.1.0 + github.com/dustin/go-humanize v1.0.0 + github.com/gammazero/deque v0.0.0-20190521012701-46e4ffb7a622 github.com/juju/errors v0.0.0-20190806202954-0232dcc7464d github.com/kr/pretty v0.1.0 // indirect github.com/pkg/errors v0.8.1 diff --git a/go.sum b/go.sum index 9b41aba..94907f7 100644 --- a/go.sum +++ b/go.sum @@ -24,6 +24,10 @@ github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghf github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo= +github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk= +github.com/gammazero/deque v0.0.0-20190521012701-46e4ffb7a622 h1:lxbhOGZ9pU3Kf8P6lFluUcE82yVZn2EqEf4+mWRNPV0= +github.com/gammazero/deque v0.0.0-20190521012701-46e4ffb7a622/go.mod h1:D90+MBHVc9Sk1lJAbEVgws0eYEurY4mv2TDso3Nxh3w= github.com/go-kit/kit v0.8.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as= github.com/go-logfmt/logfmt v0.3.0/go.mod h1:Qt1PoO58o5twSAckw1HlFXLmHsOX5/0LbT9GBnD5lWE= github.com/go-logfmt/logfmt v0.4.0/go.mod h1:3RMwSq7FuexP4Kalkev3ejPJsZTpXXBr9+V4qmtdjCk= diff --git a/hub/closeable_channel.go b/hub/closeable_channel.go new file mode 100644 index 0000000..5042ef7 --- /dev/null +++ b/hub/closeable_channel.go @@ -0,0 +1,59 @@ +package hub + +import ( + "context" + "errors" + "time" + + "github.com/9seconds/mtg/conntypes" +) + +const closeableChannelReadTimeout = 2 * time.Minute + +type ChannelReadCloser interface { + Read() (conntypes.Packet, error) + Close() +} + +type closeableChannel struct { + channel chan conntypes.Packet + ctx context.Context + cancel context.CancelFunc +} + +func (c *closeableChannel) Read() (conntypes.Packet, error) { + timer := time.NewTimer(closeableChannelReadTimeout) + defer timer.Stop() + + select { + case <-timer.C: + return nil, errors.New("timeout") + case <-c.ctx.Done(): + return nil, errors.New("channel was closed") + case packet := <-c.channel: + return packet, nil + } +} + +func (c *closeableChannel) write(packet conntypes.Packet) error { + select { + case <-c.ctx.Done(): + return errors.New("channel was closed") + case c.channel <- packet: + return nil + } +} + +func (c *closeableChannel) Close() { + c.cancel() + c.channel = nil +} + +func newCloseableChannel(ctx context.Context) *closeableChannel { + ctx, cancel := context.WithCancel(ctx) + return &closeableChannel{ + channel: make(chan conntypes.Packet), + ctx: ctx, + cancel: cancel, + } +} diff --git a/hub/connection.go b/hub/connection.go new file mode 100644 index 0000000..8879f40 --- /dev/null +++ b/hub/connection.go @@ -0,0 +1,100 @@ +package hub + +import ( + "fmt" + "math/rand" + "sync" + + "github.com/9seconds/mtg/conntypes" + "github.com/9seconds/mtg/mtproto" + "github.com/9seconds/mtg/protocol" +) + +type connectionID int + +type connection struct { + conn conntypes.PacketReadWriteCloser + mutex sync.RWMutex + id connectionID + hub *connectionHub + pending uint + closing bool +} + +func (c *connection) Write(packet conntypes.Packet) error { + c.mutex.Lock() + defer c.mutex.Unlock() + + 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.pending = 0 + } + return err +} + +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) Stats() (bool, uint) { + c.mutex.RLock() + defer c.mutex.RUnlock() + + return c.closing, c.pending +} + +func (c *connection) Close() error { + c.mutex.Lock() + defer c.mutex.Unlock() + + c.closing = true + return c.conn.Close() +} + +func (c *connection) run() { + for { + packet, err := c.conn.Read() + if err != nil { + c.Close() + c.hub.brokenSocketsChan <- c.id + c.hub = nil + return + } + + // TODO + if channel, ok := Registry.getChannel(conntypes.ConnID{}); ok { + go channel.write(packet) // nolint: errcheck + } + } +} + +func newConnection(hub *connectionHub, req *protocol.TelegramRequest) (*connection, error) { + conn, err := mtproto.TelegramProtocol(req) + if err != nil { + return nil, fmt.Errorf("cannot create a new connection: %w", err) + } + + rv := &connection{ + conn: conn, + hub: hub, + id: connectionID(rand.Int()), + } + go rv.run() + + return rv, nil +} diff --git a/hub/connection_hub.go b/hub/connection_hub.go new file mode 100644 index 0000000..44bda44 --- /dev/null +++ b/hub/connection_hub.go @@ -0,0 +1,84 @@ +package hub + +import "time" + +const hubGCEvery = time.Minute + +type connectionHub struct { + sockets map[connectionID]*connection + + brokenSocketsChan chan connectionID + connectionRequestsChan chan *connectionHubRequest + returnConnectionsChan chan *connection +} + +func (h *connectionHub) run() { + gcTicker := time.NewTicker(hubGCEvery) + defer gcTicker.Stop() + + for { + select { + case <-gcTicker.C: + h.runGC() + case id := <-h.brokenSocketsChan: + h.runBrokenConnection(id) + case request := <-h.connectionRequestsChan: + h.runConnectionRequest(request) + case conn := <-h.returnConnectionsChan: + h.runReturnConnection(conn) + } + } +} + +func (h *connectionHub) runBrokenConnection(id connectionID) { + delete(h.sockets, id) +} + +func (h *connectionHub) runGC() { + for key, conn := range h.sockets { + closing, pending := conn.Stats() + switch { + case closing: + delete(h.sockets, key) + case pending == 0: + conn.Close() + delete(h.sockets, key) + return + } + + } +} + +func (h *connectionHub) runConnectionRequest(req *connectionHubRequest) { + for key, conn := range h.sockets { + closing, _ := conn.Stats() + delete(h.sockets, key) + + if !closing { + req.responseChan <- conn + return + } + } + + newConn, err := newConnection(h, req.req) + if err != nil { + close(req.responseChan) + return + } + + req.responseChan <- newConn +} + +func (h *connectionHub) runReturnConnection(conn *connection) { + h.sockets[conn.id] = conn +} + +func newConnectionHub() *connectionHub { + return &connectionHub{ + sockets: map[connectionID]*connection{}, + + brokenSocketsChan: make(chan connectionID, 1), + connectionRequestsChan: make(chan *connectionHubRequest), + returnConnectionsChan: make(chan *connection, 1), + } +} diff --git a/hub/connection_hub_request.go b/hub/connection_hub_request.go new file mode 100644 index 0000000..eef2c7c --- /dev/null +++ b/hub/connection_hub_request.go @@ -0,0 +1,8 @@ +package hub + +import "github.com/9seconds/mtg/protocol" + +type connectionHubRequest struct { + req *protocol.TelegramRequest + responseChan chan<- *connection +} diff --git a/hub/hub.go b/hub/hub.go new file mode 100644 index 0000000..ef160d0 --- /dev/null +++ b/hub/hub.go @@ -0,0 +1,48 @@ +package hub + +import ( + "errors" + "sync" + + "github.com/9seconds/mtg/conntypes" + "github.com/9seconds/mtg/protocol" +) + +type Concentrator struct { + hubs sync.Map +} + +func (c *Concentrator) Write(packet conntypes.Packet, req *protocol.TelegramRequest) error { + hub := c.getHub(req) + connectionChan := make(chan *connection) + hub.connectionRequestsChan <- &connectionHubRequest{ + req: req, + responseChan: connectionChan, + } + + conn, ok := <-connectionChan + if !ok { + return errors.New("cannot establish connection to telegram") + } +} + +func (c *Concentrator) getHub(req *protocol.TelegramRequest) *connectionHub { + dcMapRaw, ok := c.hubs.Load(req.ClientProtocol.DC()) + if !ok { + dcMapRaw, _ = c.hubs.LoadOrStore(req.ClientProtocol.DC(), &sync.Map{}) + } + dcMap := dcMapRaw.(*sync.Map) + + loaded := true + hubRaw, ok := dcMap.Load(req.ClientProtocol.ConnectionProtocol()) + if !ok { + hubRaw, loaded = dcMap.LoadOrStore(req.ClientProtocol.ConnectionProtocol(), + newConnectionHub()) + } + hub := hubRaw.(*connectionHub) + if !loaded { + go hub.run() + } + + return hub +} diff --git a/hub/registry.go b/hub/registry.go new file mode 100644 index 0000000..0891e45 --- /dev/null +++ b/hub/registry.go @@ -0,0 +1,53 @@ +package hub + +import ( + "context" + "sync" + + "github.com/9seconds/mtg/conntypes" +) + +var Registry *RegistryStruct + +type RegistryStruct struct { + conns map[string]*closeableChannel + ctx context.Context + mutex sync.RWMutex +} + +func (r *RegistryStruct) Register(id conntypes.ConnID) ChannelReadCloser { + channel := newCloseableChannel(r.ctx) + + r.mutex.Lock() + r.conns[string(id[:])] = channel + r.mutex.Unlock() + + return channel +} + +func (r *RegistryStruct) 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 *RegistryStruct) getChannel(id conntypes.ConnID) (*closeableChannel, bool) { + r.mutex.RLock() + defer r.mutex.RUnlock() + + if value, ok := r.conns[string(id[:])]; ok { + return value, true + } + return nil, false +} + +func InitRegistry(ctx context.Context) { + Registry = &RegistryStruct{ + ctx: ctx, + conns: map[string]*closeableChannel{}, + } +} diff --git a/mtproto/protocol.go b/mtproto/protocol.go index bca7366..d6a188d 100644 --- a/mtproto/protocol.go +++ b/mtproto/protocol.go @@ -3,16 +3,15 @@ package mtproto import ( "fmt" + "github.com/9seconds/mtg/conntypes" "github.com/9seconds/mtg/mtproto/rpc" "github.com/9seconds/mtg/protocol" "github.com/9seconds/mtg/telegram" "github.com/9seconds/mtg/wrappers" ) -func TelegramProtocol(req *protocol.TelegramRequest) (wrappers.Wrap, error) { - conn, err := telegram.Middle.Dial(req.Ctx, - req.Cancel, - req.ClientProtocol.DC(), +func TelegramProtocol(req *protocol.TelegramRequest) (conntypes.PacketReadWriteCloser, error) { + conn, err := telegram.Middle.Dial(req.ClientProtocol.DC(), req.ClientProtocol.ConnectionProtocol()) if err != nil { return nil, fmt.Errorf("cannot connect to telegram: %w", err) @@ -42,7 +41,7 @@ func TelegramProtocol(req *protocol.TelegramRequest) (wrappers.Wrap, error) { return frameConn, nil } -func doRPCNonceRequest(conn wrappers.PacketWriter) (*rpc.NonceRequest, error) { +func doRPCNonceRequest(conn conntypes.PacketWriter) (*rpc.NonceRequest, error) { rpcNonceReq, err := rpc.NewNonceRequest(telegram.Middle.Secret()) if err != nil { panic(err) @@ -54,7 +53,7 @@ func doRPCNonceRequest(conn wrappers.PacketWriter) (*rpc.NonceRequest, error) { return rpcNonceReq, nil } -func getRPCNonceResponse(conn wrappers.PacketReader, req *rpc.NonceRequest) (*rpc.NonceResponse, error) { +func getRPCNonceResponse(conn conntypes.PacketReader, req *rpc.NonceRequest) (*rpc.NonceResponse, error) { packet, err := conn.Read() if err != nil { return nil, fmt.Errorf("cannot read from connection: %w", err) @@ -71,14 +70,14 @@ func getRPCNonceResponse(conn wrappers.PacketReader, req *rpc.NonceRequest) (*rp return resp, nil } -func doRPCHandshakeRequest(conn wrappers.PacketWriter) error { +func doRPCHandshakeRequest(conn conntypes.PacketWriter) error { if err := conn.Write(rpc.HandshakeRequest); err != nil { return fmt.Errorf("cannot make a request: %w", err) } return nil } -func getRPCHandshakeResponse(conn wrappers.PacketReader) error { +func getRPCHandshakeResponse(conn conntypes.PacketReader) error { packet, err := conn.Read() if err != nil { return fmt.Errorf("cannot read a response: %w", err) diff --git a/obfuscated2/client_protocol.go b/obfuscated2/client_protocol.go index 37d9800..028f868 100644 --- a/obfuscated2/client_protocol.go +++ b/obfuscated2/client_protocol.go @@ -37,7 +37,7 @@ func (c *ClientProtocol) DC() conntypes.DC { return c.dc } -func (c *ClientProtocol) Handshake(socket wrappers.StreamReadWriteCloser) (wrappers.StreamReadWriteCloser, error) { +func (c *ClientProtocol) Handshake(socket conntypes.StreamReadWriteCloser) (conntypes.StreamReadWriteCloser, error) { fm, err := c.ReadFrame(socket) if err != nil { return nil, fmt.Errorf("cannot make a client handshake: %w", err) @@ -88,7 +88,7 @@ func (c *ClientProtocol) Handshake(socket wrappers.StreamReadWriteCloser) (wrapp return wrappers.NewObfuscated2(socket, encryptor, decryptor), nil } -func (c *ClientProtocol) ReadFrame(socket wrappers.StreamReader) (fm Frame, err error) { +func (c *ClientProtocol) ReadFrame(socket conntypes.StreamReader) (fm Frame, err error) { if _, err = io.ReadFull(handshakeReader{socket}, fm.Bytes()); err != nil { err = fmt.Errorf("cannot extract obfuscated2 frame: %w", err) } @@ -96,7 +96,7 @@ func (c *ClientProtocol) ReadFrame(socket wrappers.StreamReader) (fm Frame, err } type handshakeReader struct { - parent wrappers.StreamReader + parent conntypes.StreamReader } func (h handshakeReader) Read(p []byte) (int, error) { diff --git a/obfuscated2/telegram_protocol.go b/obfuscated2/telegram_protocol.go index eda4460..441ecd9 100644 --- a/obfuscated2/telegram_protocol.go +++ b/obfuscated2/telegram_protocol.go @@ -4,20 +4,21 @@ import ( "crypto/rand" "fmt" + "github.com/9seconds/mtg/conntypes" "github.com/9seconds/mtg/protocol" "github.com/9seconds/mtg/telegram" "github.com/9seconds/mtg/utils" "github.com/9seconds/mtg/wrappers" ) -func TelegramProtocol(req *protocol.TelegramRequest) (wrappers.Wrap, error) { - socket, err := telegram.Direct.Dial(req.Ctx, - req.Cancel, - req.ClientProtocol.DC(), +func TelegramProtocol(req *protocol.TelegramRequest) (conntypes.StreamReadWriteCloser, error) { + conn, err := telegram.Direct.Dial(req.ClientProtocol.DC(), req.ClientProtocol.ConnectionProtocol()) if err != nil { return nil, fmt.Errorf("cannot dial to telegram: %w", err) } + conn = wrappers.NewTimeout(conn) + conn = wrappers.NewCtx(req.Ctx, req.Cancel, conn) fm := generateFrame(req.ClientProtocol) data := fm.Bytes() @@ -30,11 +31,11 @@ func TelegramProtocol(req *protocol.TelegramRequest) (wrappers.Wrap, error) { encryptor.XORKeyStream(data, data) copy(data[:frameOffsetIV], copyFrame[:frameOffsetIV]) - if _, err := socket.Write(data); err != nil { + if _, err := conn.Write(data); err != nil { return nil, fmt.Errorf("cannot write handshake frame to telegram: %w", err) } - return wrappers.NewObfuscated2(socket, encryptor, decryptor), nil + return wrappers.NewObfuscated2(conn, encryptor, decryptor), nil } func generateFrame(cp protocol.ClientProtocol) (fm Frame) { diff --git a/protocol/interfaces.go b/protocol/interfaces.go index a392de1..2e27fc3 100644 --- a/protocol/interfaces.go +++ b/protocol/interfaces.go @@ -1,16 +1,12 @@ package protocol -import ( - "github.com/9seconds/mtg/conntypes" - "github.com/9seconds/mtg/wrappers" -) +import "github.com/9seconds/mtg/conntypes" type ClientProtocol interface { - Handshake(wrappers.StreamReadWriteCloser) (wrappers.StreamReadWriteCloser, error) + Handshake(conntypes.StreamReadWriteCloser) (conntypes.StreamReadWriteCloser, error) ConnectionType() conntypes.ConnectionType ConnectionProtocol() conntypes.ConnectionProtocol DC() conntypes.DC } -type TelegramProtocol func(*TelegramRequest) (wrappers.Wrap, error) type ClientProtocolMaker func() ClientProtocol diff --git a/protocol/request.go b/protocol/request.go index e73489d..00a33f8 100644 --- a/protocol/request.go +++ b/protocol/request.go @@ -6,12 +6,11 @@ import ( "go.uber.org/zap" "github.com/9seconds/mtg/conntypes" - "github.com/9seconds/mtg/wrappers" ) type TelegramRequest struct { Logger *zap.SugaredLogger - ClientConn wrappers.StreamReadWriteCloser + ClientConn conntypes.StreamReadWriteCloser ConnID conntypes.ConnID Ctx context.Context Cancel context.CancelFunc diff --git a/proxy/proxy.go b/proxy/proxy.go index b42cd60..501b721 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -63,25 +63,26 @@ func (p *Proxy) accept(conn net.Conn) { ctx, cancel := context.WithCancel(p.Context) defer cancel() - wrappedConn := wrappers.NewClientConn(ctx, cancel, conn, connID) - wrappedConn = wrappers.NewTraffic(wrappedConn) - defer wrappedConn.Close() + clientConn := wrappers.NewClientConn(conn, connID) + clientConn = wrappers.NewCtx(ctx, cancel, clientConn) + clientConn = wrappers.NewTimeout(clientConn) + clientConn = wrappers.NewTraffic(clientConn) + defer clientConn.Close() clientProtocol := p.ClientProtocolMaker() - wrappedConn, err := clientProtocol.Handshake(wrappedConn) + clientConn, err := clientProtocol.Handshake(clientConn) if err != nil { logger.Warnw("Cannot perform client handshake", "error", err) return } - defer wrappedConn.Close() - stats.S.ClientConnected(clientProtocol.ConnectionType(), wrappedConn.RemoteAddr()) - defer stats.S.ClientDisconnected(clientProtocol.ConnectionType(), wrappedConn.RemoteAddr()) + stats.S.ClientConnected(clientProtocol.ConnectionType(), clientConn.RemoteAddr()) + defer stats.S.ClientDisconnected(clientProtocol.ConnectionType(), clientConn.RemoteAddr()) logger.Infow("Client connected", "addr", conn.RemoteAddr()) req := &protocol.TelegramRequest{ Logger: logger, - ClientConn: wrappedConn, + ClientConn: clientConn, ConnID: connID, Ctx: ctx, Cancel: cancel, @@ -102,7 +103,7 @@ func (p *Proxy) acceptDirectConnection(request *protocol.TelegramRequest) error if err != nil { return err } - telegramConn := telegramConnRaw.(wrappers.StreamReadWriteCloser) + telegramConn := telegramConnRaw.(conntypes.StreamReadWriteCloser) defer telegramConn.Close() wg := &sync.WaitGroup{} diff --git a/telegram/base.go b/telegram/base.go index bc3319e..bee996a 100644 --- a/telegram/base.go +++ b/telegram/base.go @@ -1,7 +1,6 @@ package telegram import ( - "context" "fmt" "math/rand" "net" @@ -28,25 +27,8 @@ func (b *baseTelegram) Secret() []byte { return b.secret } -func (b *baseTelegram) dialToAddress(ctx context.Context, - cancel context.CancelFunc, - addr string) (wrappers.StreamReadWriteCloser, error) { - conn, err := b.dialer.Dial("tcp", addr) - if err != nil { - return nil, fmt.Errorf("dial has failed: %w", err) - } - - if err := utils.InitTCP(conn); err != nil { - return nil, fmt.Errorf("cannot initialize tcp socket: %w", err) - } - - return wrappers.NewTelegramConn(ctx, cancel, conn), nil -} - -func (b *baseTelegram) dial(ctx context.Context, - cancel context.CancelFunc, - dc conntypes.DC, - protocol conntypes.ConnectionProtocol) (wrappers.StreamReadWriteCloser, error) { +func (b *baseTelegram) dial(dc conntypes.DC, + protocol conntypes.ConnectionProtocol) (conntypes.StreamReadWriteCloser, error) { addr := "" switch protocol { @@ -56,7 +38,16 @@ func (b *baseTelegram) dial(ctx context.Context, addr = b.chooseAddress(b.v6Addresses, dc, b.V6DefaultDC) } - return b.dialToAddress(ctx, cancel, addr) + conn, err := b.dialer.Dial("tcp", addr) + if err != nil { + return nil, fmt.Errorf("dial has failed: %w", err) + } + + if err := utils.InitTCP(conn); err != nil { + return nil, fmt.Errorf("cannot initialize tcp socket: %w", err) + } + + return wrappers.NewTelegramConn(conn), nil } func (b *baseTelegram) chooseAddress(addresses map[conntypes.DC][]string, @@ -66,7 +57,10 @@ func (b *baseTelegram) chooseAddress(addresses map[conntypes.DC][]string, addrs, _ = addresses[defaultDC] } - if len(addrs) > 0 { + switch { + case len(addrs) == 1: + return addrs[0] + case len(addrs) > 1: return addrs[rand.Intn(len(addrs))] } diff --git a/telegram/direct.go b/telegram/direct.go index 153143f..b69075e 100644 --- a/telegram/direct.go +++ b/telegram/direct.go @@ -1,15 +1,11 @@ package telegram import ( - "context" "net" "github.com/9seconds/mtg/conntypes" - "github.com/9seconds/mtg/wrappers" ) -var Direct = newDirectTelegram() - const ( directV4DefaultIdx conntypes.DC = 1 directV6DefaultIdx conntypes.DC = 1 @@ -36,10 +32,8 @@ type directTelegram struct { baseTelegram } -func (d *directTelegram) Dial(ctx context.Context, - cancel context.CancelFunc, - dc conntypes.DC, - protocol conntypes.ConnectionProtocol) (wrappers.StreamReadWriteCloser, error) { +func (d *directTelegram) Dial(dc conntypes.DC, + protocol conntypes.ConnectionProtocol) (conntypes.StreamReadWriteCloser, error) { switch { case dc < 0: dc = -dc @@ -47,17 +41,15 @@ func (d *directTelegram) Dial(ctx context.Context, dc = conntypes.DCDefaultIdx } - return d.baseTelegram.dial(ctx, cancel, dc-1, protocol) + return d.baseTelegram.dial(dc-1, protocol) } -func newDirectTelegram() Telegram { - return &directTelegram{ - baseTelegram: baseTelegram{ - dialer: net.Dialer{Timeout: telegramDialTimeout}, - v4DefaultDC: directV4DefaultIdx, - V6DefaultDC: directV6DefaultIdx, - v4Addresses: directV4Addresses, - v6Addresses: directV6Addresses, - }, - } +var Direct = &directTelegram{ + baseTelegram: baseTelegram{ + dialer: net.Dialer{Timeout: telegramDialTimeout}, + v4DefaultDC: directV4DefaultIdx, + V6DefaultDC: directV6DefaultIdx, + v4Addresses: directV4Addresses, + v6Addresses: directV6Addresses, + }, } diff --git a/telegram/interfaces.go b/telegram/interfaces.go index 9dc1bb8..c780a62 100644 --- a/telegram/interfaces.go +++ b/telegram/interfaces.go @@ -1,16 +1,8 @@ package telegram -import ( - "context" - - "github.com/9seconds/mtg/conntypes" - "github.com/9seconds/mtg/wrappers" -) +import "github.com/9seconds/mtg/conntypes" type Telegram interface { - Dial(context.Context, - context.CancelFunc, - conntypes.DC, - conntypes.ConnectionProtocol) (wrappers.StreamReadWriteCloser, error) + Dial(conntypes.DC, conntypes.ConnectionProtocol) (conntypes.StreamReadWriteCloser, error) Secret() []byte } diff --git a/telegram/middle.go b/telegram/middle.go index f6c6b2f..6dc0279 100644 --- a/telegram/middle.go +++ b/telegram/middle.go @@ -1,7 +1,6 @@ package telegram import ( - "context" "fmt" "net" "sync" @@ -11,7 +10,6 @@ import ( "github.com/9seconds/mtg/conntypes" "github.com/9seconds/mtg/telegram/api" - "github.com/9seconds/mtg/wrappers" ) const middleTelegramBackgroundUpdateEvery = time.Hour @@ -67,10 +65,8 @@ func (m *middleTelegram) backgroundUpdate() { } } -func (m *middleTelegram) Dial(ctx context.Context, - cancel context.CancelFunc, - dc conntypes.DC, - protocol conntypes.ConnectionProtocol) (wrappers.StreamReadWriteCloser, error) { +func (m *middleTelegram) Dial(dc conntypes.DC, + protocol conntypes.ConnectionProtocol) (conntypes.StreamReadWriteCloser, error) { if dc == 0 { dc = conntypes.DCDefaultIdx } @@ -78,7 +74,7 @@ func (m *middleTelegram) Dial(ctx context.Context, m.mutex.RLock() defer m.mutex.RUnlock() - return m.baseTelegram.dial(ctx, cancel, dc, protocol) + return m.baseTelegram.dial(dc, protocol) } func MiddleInit() { diff --git a/wrappers/blockcipher.go b/wrappers/blockcipher.go index 4554515..12c015f 100644 --- a/wrappers/blockcipher.go +++ b/wrappers/blockcipher.go @@ -10,6 +10,8 @@ import ( "time" "go.uber.org/zap" + + "github.com/9seconds/mtg/conntypes" ) const blockCipherReadCurrentDataBufferSize = 1024 + 1 // +1 because telegram operates with blocks mod 4 @@ -17,7 +19,7 @@ const blockCipherReadCurrentDataBufferSize = 1024 + 1 // +1 because telegram ope type wrapperBlockCipher struct { buf bytes.Buffer - parent StreamReadWriteCloser + parent conntypes.StreamReadWriteCloser encryptor cipher.BlockMode decryptor cipher.BlockMode } @@ -47,7 +49,8 @@ func (w *wrapperBlockCipher) ReadTimeout(p []byte, timeout time.Duration) (int, return w.read(p, readAllTimeout(timeout)) } -func (w *wrapperBlockCipher) read(p []byte, reader func(StreamReadWriteCloser) ([]byte, error)) (int, error) { +func (w *wrapperBlockCipher) read(p []byte, + reader func(conntypes.StreamReadWriteCloser) ([]byte, error)) (int, error) { if w.buf.Len() > 0 { return w.flush(p) } @@ -90,7 +93,7 @@ func (w *wrapperBlockCipher) encrypt(p []byte) ([]byte, error) { return encrypted, nil } -func readAll(src StreamReadWriteCloser) (rv []byte, err error) { +func readAll(src conntypes.StreamReadWriteCloser) (rv []byte, err error) { buf := make([]byte, blockCipherReadCurrentDataBufferSize) n := blockCipherReadCurrentDataBufferSize @@ -105,8 +108,8 @@ func readAll(src StreamReadWriteCloser) (rv []byte, err error) { return rv, nil } -func readAllTimeout(timeout time.Duration) func(StreamReadWriteCloser) ([]byte, error) { - return func(src StreamReadWriteCloser) (rv []byte, err error) { +func readAllTimeout(timeout time.Duration) func(conntypes.StreamReadWriteCloser) ([]byte, error) { + return func(src conntypes.StreamReadWriteCloser) (rv []byte, err error) { tmo := timeout buf := make([]byte, blockCipherReadCurrentDataBufferSize) n := blockCipherReadCurrentDataBufferSize @@ -148,7 +151,8 @@ func (w *wrapperBlockCipher) RemoteAddr() *net.TCPAddr { return w.parent.RemoteAddr() } -func newBlockCipher(parent StreamReadWriteCloser, encryptor, decryptor cipher.BlockMode) StreamReadWriteCloser { +func newBlockCipher(parent conntypes.StreamReadWriteCloser, + encryptor, decryptor cipher.BlockMode) conntypes.StreamReadWriteCloser { return &wrapperBlockCipher{ parent: parent, encryptor: encryptor, diff --git a/wrappers/conn.go b/wrappers/conn.go index 1649cd1..2d2f982 100644 --- a/wrappers/conn.go +++ b/wrappers/conn.go @@ -1,7 +1,6 @@ package wrappers import ( - "context" "fmt" "net" "time" @@ -19,15 +18,8 @@ const ( connPurposeTelegram ) -const ( - connTimeoutRead = 2 * time.Minute - connTimeoutWrite = 2 * time.Minute -) - type wrapperConn struct { parent net.Conn - ctx context.Context - cancel context.CancelFunc connID conntypes.ConnID logger *zap.SugaredLogger localAddr *net.TCPAddr @@ -35,61 +27,45 @@ type wrapperConn struct { } func (w *wrapperConn) WriteTimeout(p []byte, timeout time.Duration) (int, error) { - select { - case <-w.ctx.Done(): + if err := w.parent.SetWriteDeadline(time.Now().Add(timeout)); err != nil { w.Close() - return 0, fmt.Errorf("cannot write because context was closed: %w", w.ctx.Err()) - - default: - if err := w.parent.SetWriteDeadline(time.Now().Add(timeout)); err != nil { - w.Close() // nolint: gosec - return 0, fmt.Errorf("cannot set write deadline to the socket: %w", err) - } - - n, err := w.parent.Write(p) - w.logger.Debugw("Write to stream", "bytes", n, "error", err) - if err != nil { - w.Close() // nolint: gosec - } - - return n, err + return 0, fmt.Errorf("cannot set write deadline to the socket: %w", err) } + + return w.Write(p) } func (w *wrapperConn) Write(p []byte) (int, error) { - return w.WriteTimeout(p, connTimeoutWrite) + n, err := w.parent.Write(p) + w.logger.Debugw("write to stream", "bytes", n, "error", err) + if err != nil { + w.Close() // nolint: gosec + } + + return n, err } func (w *wrapperConn) ReadTimeout(p []byte, timeout time.Duration) (int, error) { - select { - case <-w.ctx.Done(): + if err := w.parent.SetReadDeadline(time.Now().Add(timeout)); err != nil { w.Close() - return 0, fmt.Errorf("cannot read because context was closed: %w", w.ctx.Err()) - - default: - if err := w.parent.SetReadDeadline(time.Now().Add(timeout)); err != nil { - w.Close() - return 0, fmt.Errorf("cannot set read deadline to the socket: %w", err) - } - - n, err := w.parent.Read(p) - w.logger.Debugw("Read from stream", "bytes", n, "error", err) - if err != nil { - w.Close() - } - - return n, err + return 0, fmt.Errorf("cannot set read deadline to the socket: %w", err) } + + return w.Read(p) } func (w *wrapperConn) Read(p []byte) (int, error) { - return w.ReadTimeout(p, connTimeoutRead) + n, err := w.parent.Read(p) + w.logger.Debugw("Read from stream", "bytes", n, "error", err) + if err != nil { + w.Close() + } + + return n, err } func (w *wrapperConn) Close() error { w.logger.Debugw("Close connection") - w.cancel() - return w.parent.Close() } @@ -109,11 +85,9 @@ func (w *wrapperConn) RemoteAddr() *net.TCPAddr { return w.remoteAddr } -func newConn(ctx context.Context, - cancel context.CancelFunc, - parent net.Conn, +func newConn(parent net.Conn, connID conntypes.ConnID, - purpose connPurpose) StreamReadWriteCloser { + purpose connPurpose) conntypes.StreamReadWriteCloser { localAddr := *parent.LocalAddr().(*net.TCPAddr) if parent.RemoteAddr().(*net.TCPAddr).IP.To4() != nil { @@ -135,8 +109,6 @@ func newConn(ctx context.Context, return &wrapperConn{ parent: parent, - ctx: ctx, - cancel: cancel, connID: connID, logger: logger, remoteAddr: parent.RemoteAddr().(*net.TCPAddr), @@ -144,15 +116,11 @@ func newConn(ctx context.Context, } } -func NewClientConn(ctx context.Context, - cancel context.CancelFunc, - parent net.Conn, - connID conntypes.ConnID) StreamReadWriteCloser { - return newConn(ctx, cancel, parent, connID, connPurposeClient) +func NewClientConn(parent net.Conn, + connID conntypes.ConnID) conntypes.StreamReadWriteCloser { + return newConn(parent, connID, connPurposeClient) } -func NewTelegramConn(ctx context.Context, - cancel context.CancelFunc, - parent net.Conn) StreamReadWriteCloser { - return newConn(ctx, cancel, parent, conntypes.ConnID{}, connPurposeTelegram) +func NewTelegramConn(parent net.Conn) conntypes.StreamReadWriteCloser { + return newConn(parent, conntypes.ConnID{}, connPurposeTelegram) } diff --git a/wrappers/ctx.go b/wrappers/ctx.go new file mode 100644 index 0000000..46f4242 --- /dev/null +++ b/wrappers/ctx.go @@ -0,0 +1,89 @@ +package wrappers + +import ( + "context" + "fmt" + "net" + "time" + + "go.uber.org/zap" + + "github.com/9seconds/mtg/conntypes" +) + +type wrapperCtx struct { + parent conntypes.StreamReadWriteCloser + ctx context.Context + cancel context.CancelFunc +} + +func (w *wrapperCtx) WriteTimeout(p []byte, timeout time.Duration) (int, error) { + select { + case <-w.ctx.Done(): + w.Close() + return 0, fmt.Errorf("cannot write because context was closed: %w", w.ctx.Err()) + default: + return w.parent.WriteTimeout(p, timeout) + } +} + +func (w *wrapperCtx) Write(p []byte) (int, error) { + select { + case <-w.ctx.Done(): + w.Close() + return 0, fmt.Errorf("cannot write because context was closed: %w", w.ctx.Err()) + default: + return w.parent.Write(p) + } +} + +func (w *wrapperCtx) ReadTimeout(p []byte, timeout time.Duration) (int, error) { + select { + case <-w.ctx.Done(): + w.Close() + return 0, fmt.Errorf("cannot write because context was closed: %w", w.ctx.Err()) + default: + return w.parent.ReadTimeout(p, timeout) + } +} + +func (w *wrapperCtx) Read(p []byte) (int, error) { + select { + case <-w.ctx.Done(): + w.Close() + return 0, fmt.Errorf("cannot write because context was closed: %w", w.ctx.Err()) + default: + return w.parent.Read(p) + } +} + +func (w *wrapperCtx) Close() error { + w.cancel() + return w.parent.Close() +} + +func (w *wrapperCtx) Conn() net.Conn { + return w.parent.Conn() +} + +func (w *wrapperCtx) Logger() *zap.SugaredLogger { + return w.parent.Logger().Named("ctx") +} + +func (w *wrapperCtx) LocalAddr() *net.TCPAddr { + return w.parent.LocalAddr() +} + +func (w *wrapperCtx) RemoteAddr() *net.TCPAddr { + return w.parent.RemoteAddr() +} + +func NewCtx(ctx context.Context, + cancel context.CancelFunc, + parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser { + return &wrapperCtx{ + parent: parent, + ctx: ctx, + cancel: cancel, + } +} diff --git a/wrappers/mtproto_cipher.go b/wrappers/mtproto_cipher.go index e297d31..01ec5c5 100644 --- a/wrappers/mtproto_cipher.go +++ b/wrappers/mtproto_cipher.go @@ -9,6 +9,7 @@ import ( "encoding/binary" "net" + "github.com/9seconds/mtg/conntypes" "github.com/9seconds/mtg/mtproto/rpc" "github.com/9seconds/mtg/utils" ) @@ -22,10 +23,10 @@ const ( var mtprotoEmptyIP = [4]byte{0x00, 0x00, 0x00, 0x00} -func NewMiddleProxyCipher(parent StreamReadWriteCloser, +func NewMiddleProxyCipher(parent conntypes.StreamReadWriteCloser, req *rpc.NonceRequest, resp *rpc.NonceResponse, - secret []byte) StreamReadWriteCloser { + secret []byte) conntypes.StreamReadWriteCloser { localAddr := parent.LocalAddr() remoteAddr := parent.RemoteAddr() diff --git a/wrappers/mtproto_frame.go b/wrappers/mtproto_frame.go index 3fecc56..071aec6 100644 --- a/wrappers/mtproto_frame.go +++ b/wrappers/mtproto_frame.go @@ -11,6 +11,8 @@ import ( "net" "go.uber.org/zap" + + "github.com/9seconds/mtg/conntypes" ) const ( @@ -33,13 +35,13 @@ var mtprotoFramePadding = []byte{0x04, 0x00, 0x00, 0x00} // PADDING is custom padding schema to complete frame length to such that // len(frame) % 16 == 0 type wrapperMtprotoFrame struct { - parent StreamReadWriteCloser + parent conntypes.StreamReadWriteCloser logger *zap.SugaredLogger readSeqNo int32 writeSeqNo int32 } -func (w *wrapperMtprotoFrame) Read() (Packet, error) { +func (w *wrapperMtprotoFrame) Read() (conntypes.Packet, error) { buf := &bytes.Buffer{} sum := crc32.NewIEEE() writer := io.MultiWriter(buf, sum) @@ -101,7 +103,7 @@ func (w *wrapperMtprotoFrame) Read() (Packet, error) { return data, nil } -func (w *wrapperMtprotoFrame) Write(p Packet) error { +func (w *wrapperMtprotoFrame) Write(p conntypes.Packet) error { messageLength := 4 + 4 + len(p) + 4 paddingLength := (aes.BlockSize - messageLength%aes.BlockSize) % aes.BlockSize @@ -149,7 +151,7 @@ func (w *wrapperMtprotoFrame) RemoteAddr() *net.TCPAddr { return w.parent.RemoteAddr() } -func NewMtprotoFrame(parent StreamReadWriteCloser, seqNo int32) PacketReadWriteCloser { +func NewMtprotoFrame(parent conntypes.StreamReadWriteCloser, seqNo int32) conntypes.PacketReadWriteCloser { return &wrapperMtprotoFrame{ parent: parent, logger: parent.Logger().Named("mtproto-frame"), diff --git a/wrappers/obfuscated2.go b/wrappers/obfuscated2.go index 4ef601a..d6f7a69 100644 --- a/wrappers/obfuscated2.go +++ b/wrappers/obfuscated2.go @@ -7,12 +7,14 @@ import ( "time" "go.uber.org/zap" + + "github.com/9seconds/mtg/conntypes" ) type wrapperObfuscated2 struct { encryptor cipher.Stream decryptor cipher.Stream - parent StreamReadWriteCloser + parent conntypes.StreamReadWriteCloser } func (w *wrapperObfuscated2) ReadTimeout(p []byte, timeout time.Duration) (int, error) { @@ -71,7 +73,8 @@ func (w *wrapperObfuscated2) Close() error { return w.parent.Close() } -func NewObfuscated2(socket StreamReadWriteCloser, encryptor, decryptor cipher.Stream) StreamReadWriteCloser { +func NewObfuscated2(socket conntypes.StreamReadWriteCloser, + encryptor, decryptor cipher.Stream) conntypes.StreamReadWriteCloser { return &wrapperObfuscated2{ parent: socket, encryptor: encryptor, diff --git a/wrappers/stats.go b/wrappers/stats.go index 4f33eee..33f57be 100644 --- a/wrappers/stats.go +++ b/wrappers/stats.go @@ -6,11 +6,12 @@ import ( "go.uber.org/zap" + "github.com/9seconds/mtg/conntypes" "github.com/9seconds/mtg/stats" ) type wrapperStats struct { - parent StreamReadWriteCloser + parent conntypes.StreamReadWriteCloser } func (w *wrapperStats) Write(p []byte) (int, error) { @@ -61,6 +62,6 @@ func (w *wrapperStats) Close() error { return w.parent.Close() } -func NewTraffic(parent StreamReadWriteCloser) StreamReadWriteCloser { +func NewTraffic(parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser { return &wrapperStats{parent} } diff --git a/wrappers/timeout.go b/wrappers/timeout.go new file mode 100644 index 0000000..f9c0302 --- /dev/null +++ b/wrappers/timeout.go @@ -0,0 +1,61 @@ +package wrappers + +import ( + "net" + "time" + + "go.uber.org/zap" + + "github.com/9seconds/mtg/conntypes" +) + +const ( + timeoutRead = 2 * time.Minute + timeoutWrite = 2 * time.Minute +) + +type wrapperTimeout struct { + parent conntypes.StreamReadWriteCloser +} + +func (w *wrapperTimeout) WriteTimeout(p []byte, timeout time.Duration) (int, error) { + return w.parent.WriteTimeout(p, timeout) +} + +func (w *wrapperTimeout) Write(p []byte) (int, error) { + return w.parent.WriteTimeout(p, timeoutWrite) +} + +func (w *wrapperTimeout) ReadTimeout(p []byte, timeout time.Duration) (int, error) { + return w.parent.ReadTimeout(p, timeout) +} + +func (w *wrapperTimeout) Read(p []byte) (int, error) { + return w.parent.ReadTimeout(p, timeoutRead) +} + +func (w *wrapperTimeout) Close() error { + return w.parent.Close() +} + +func (w *wrapperTimeout) Conn() net.Conn { + return w.parent.Conn() +} + +func (w *wrapperTimeout) Logger() *zap.SugaredLogger { + return w.parent.Logger().Named("timeout") +} + +func (w *wrapperTimeout) LocalAddr() *net.TCPAddr { + return w.parent.LocalAddr() +} + +func (w *wrapperTimeout) RemoteAddr() *net.TCPAddr { + return w.parent.RemoteAddr() +} + +func NewTimeout(parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser { + return &wrapperTimeout{ + parent: parent, + } +}