Add client abridged protocol

conntypes/acks.go

@@ -0,0 +1,6 @@
+package conntypes
+
+type ConnectionAcks struct {
+	Simple bool
+	Quick  bool
+}

conntypes/wrap_interfaces.go

@@ -0,0 +1,14 @@
+package conntypes
+
+import (
+	"net"
+
+	"go.uber.org/zap"
+)
+
+type Wrap interface {
+	Conn() net.Conn
+	Logger() *zap.SugaredLogger
+	LocalAddr() *net.TCPAddr
+	RemoteAddr() *net.TCPAddr
+}

conntypes/wrap_packet_ack_interfaces.go

@@ -0,0 +1,41 @@
+package conntypes
+
+import "io"
+
+type PacketAckReader interface {
+	Read(*ConnectionAcks) (Packet, error)
+}
+
+type PacketAckWriter interface {
+	Write(Packet, *ConnectionAcks) error
+}
+
+type PacketAckCloser interface {
+	io.Closer
+}
+
+type PacketAckReadCloser interface {
+	PacketAckReader
+	PacketAckCloser
+}
+
+type PacketAckWriteCloser interface {
+	PacketAckWriter
+	PacketAckCloser
+}
+
+type PacketAckReadWriter interface {
+	PacketAckReader
+	PacketAckWriter
+}
+
+type PacketAckReadWriteCloser interface {
+	PacketAckReader
+	PacketAckWriter
+	PacketAckCloser
+}
+
+type PacketAckFullReadWriteCloser interface {
+	Wrap
+	PacketAckReadWriteCloser
+}

conntypes/wrap_packet_interfaces.go

@@ -0,0 +1,51 @@
+package conntypes
+
+import "io"
+
+type BasePacketReader interface {
+	Read() (Packet, error)
+}
+
+type BasePacketWriter interface {
+	Write(Packet) error
+}
+
+type PacketReader interface {
+	Wrap
+	BasePacketReader
+}
+
+type PacketWriter interface {
+	Wrap
+	BasePacketWriter
+}
+
+type PacketCloser interface {
+	Wrap
+	io.Closer
+}
+
+type PacketReadCloser interface {
+	Wrap
+	BasePacketReader
+	io.Closer
+}
+
+type PacketWriteCloser interface {
+	Wrap
+	BasePacketWriter
+	io.Closer
+}
+
+type PacketReadWriter interface {
+	Wrap
+	BasePacketWriter
+	BasePacketReader
+}
+
+type PacketReadWriteCloser interface {
+	Wrap
+	BasePacketWriter
+	BasePacketReader
+	io.Closer
+}

conntypes/wrap_stream_interfaces.go

@@ -0,0 +1,56 @@
+package conntypes
+
+import (
+	"io"
+	"time"
+)
+
+type BaseStreamReaderWithTimeout interface {
+	ReadTimeout([]byte, time.Duration) (int, error)
+}
+
+type BaseStreamWriterWithTimeout interface {
+	WriteTimeout([]byte, time.Duration) (int, error)
+}
+
+type StreamReader interface {
+	Wrap
+	io.Reader
+	BaseStreamReaderWithTimeout
+}
+
+type StreamWriter interface {
+	Wrap
+	io.Writer
+	BaseStreamWriterWithTimeout
+}
+
+type StreamCloser interface {
+	Wrap
+	io.Closer
+}
+
+type StreamReadCloser interface {
+	Wrap
+	io.ReadCloser
+	BaseStreamReaderWithTimeout
+}
+
+type StreamWriteCloser interface {
+	Wrap
+	io.WriteCloser
+	BaseStreamWriterWithTimeout
+}
+
+type StreamReadWriter interface {
+	Wrap
+	io.ReadWriter
+	BaseStreamReaderWithTimeout
+}
+
+type StreamReadWriteCloser interface {
+	Wrap
+	io.ReadWriteCloser
+	BaseStreamReaderWithTimeout
+	BaseStreamWriterWithTimeout
+}

conntypes/wrappers.go

@@ -1,115 +1,2 @@
 package conntypes
 
-import (
-	"io"
-	"net"
-	"time"
-
-	"go.uber.org/zap"
-)
-
-// Wrap is a base interface for all wrappers in this package.
-type Wrap interface {
-	Conn() net.Conn
-	Logger() *zap.SugaredLogger
-	LocalAddr() *net.TCPAddr
-	RemoteAddr() *net.TCPAddr
-}
-
-type BaseReaderWithTimeout interface {
-	ReadTimeout([]byte, time.Duration) (int, error)
-}
-
-type BaseWriterWithTimeout interface {
-	WriteTimeout([]byte, time.Duration) (int, error)
-}
-
-type BasePacketReader interface {
-	Read() (Packet, error)
-}
-
-type BasePacketWriter interface {
-	Write(Packet) error
-}
-
-type StreamReader interface {
-	Wrap
-	io.Reader
-	BaseReaderWithTimeout
-}
-
-type StreamWriter interface {
-	Wrap
-	io.Writer
-	BaseWriterWithTimeout
-}
-
-type StreamCloser interface {
-	Wrap
-	io.Closer
-}
-
-type StreamReadCloser interface {
-	Wrap
-	io.ReadCloser
-	BaseReaderWithTimeout
-}
-
-type StreamWriteCloser interface {
-	Wrap
-	io.WriteCloser
-	BaseWriterWithTimeout
-}
-
-type StreamReadWriter interface {
-	Wrap
-	io.ReadWriter
-	BaseReaderWithTimeout
-}
-
-type StreamReadWriteCloser interface {
-	Wrap
-	io.ReadWriteCloser
-	BaseReaderWithTimeout
-	BaseWriterWithTimeout
-}
-
-type PacketReader interface {
-	Wrap
-	BasePacketReader
-}
-
-type PacketWriter interface {
-	Wrap
-	BasePacketWriter
-}
-
-type PacketCloser interface {
-	Wrap
-	io.Closer
-}
-
-type PacketReadCloser interface {
-	Wrap
-	BasePacketReader
-	io.Closer
-}
-
-type PacketWriteCloser interface {
-	Wrap
-	BasePacketWriter
-	io.Closer
-}
-
-type PacketReadWriter interface {
-	Wrap
-	BasePacketWriter
-	BasePacketReader
-}
-
-type PacketReadWriteCloser interface {
-	Wrap
-	BasePacketWriter
-	BasePacketReader
-	io.Closer
-}

go.mod

@@ -11,7 +11,10 @@ require (
 	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/gofrs/uuid v3.2.0+incompatible
 	github.com/juju/errors v0.0.0-20190806202954-0232dcc7464d
+	github.com/juju/loggo v0.0.0-20190526231331-6e530bcce5d8 // indirect
+	github.com/juju/testing v0.0.0-20191001232224-ce9dec17d28b // indirect
 	github.com/kr/pretty v0.1.0 // indirect
 	github.com/pkg/errors v0.8.1
 	github.com/prometheus/client_golang v1.1.0
@@ -26,6 +29,7 @@ require (
 	gopkg.in/alecthomas/kingpin.v2 v2.2.6
 	gopkg.in/alexcesaro/statsd.v2 v2.0.0
 	gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect
+	gopkg.in/mgo.v2 v2.0.0-20190816093944-a6b53ec6cb22 // indirect
 )
 
 go 1.13

go.sum

@@ -32,6 +32,8 @@ github.com/go-kit/kit v0.8.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2
 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=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
+github.com/gofrs/uuid v3.2.0+incompatible h1:y12jRkkFxsd7GpqdSZ+/KCs/fJbqpEXSGd4+jfEaewE=
+github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
 github.com/gogo/protobuf v1.1.1 h1:72R+M5VuhED/KujmZVcIquuo8mBgX4oVda//DQb3PXo=
 github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
 github.com/golang/protobuf v1.2.0 h1:P3YflyNX/ehuJFLhxviNdFxQPkGK5cDcApsge1SqnvM=
@@ -46,6 +48,10 @@ github.com/json-iterator/go v1.1.6/go.mod h1:+SdeFBvtyEkXs7REEP0seUULqWtbJapLOCV
 github.com/json-iterator/go v1.1.7/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4=
 github.com/juju/errors v0.0.0-20190806202954-0232dcc7464d h1:hJXjZMxj0SWlMoQkzeZDLi2cmeiWKa7y1B8Rg+qaoEc=
 github.com/juju/errors v0.0.0-20190806202954-0232dcc7464d/go.mod h1:W54LbzXuIE0boCoNJfwqpmkKJ1O4TCTZMetAt6jGk7Q=
+github.com/juju/loggo v0.0.0-20190526231331-6e530bcce5d8 h1:UUHMLvzt/31azWTN/ifGWef4WUqvXk0iRqdhdy/2uzI=
+github.com/juju/loggo v0.0.0-20190526231331-6e530bcce5d8/go.mod h1:vgyd7OREkbtVEN/8IXZe5Ooef3LQePvuBm9UWj6ZL8U=
+github.com/juju/testing v0.0.0-20191001232224-ce9dec17d28b h1:Rrp0ByJXEjhREMPGTt3aWYjoIsUGCbt21ekbeJcTWv0=
+github.com/juju/testing v0.0.0-20191001232224-ce9dec17d28b/go.mod h1:63prj8cnj0tU0S9OHjGJn+b1h0ZghCndfnbQolrYTwA=
 github.com/julienschmidt/httprouter v1.2.0/go.mod h1:SYymIcj16QtmaHHD7aYtjjsJG7VTCxuUUipMqKk8s4w=
 github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc=
@@ -128,6 +134,8 @@ gopkg.in/alexcesaro/statsd.v2 v2.0.0/go.mod h1:i0ubccKGzBVNBpdGV5MocxyA/XlLUJzA7
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/mgo.v2 v2.0.0-20190816093944-a6b53ec6cb22 h1:VpOs+IwYnYBaFnrNAeB8UUWtL3vEUnzSCL1nVjPhqrw=
+gopkg.in/mgo.v2 v2.0.0-20190816093944-a6b53ec6cb22/go.mod h1:yeKp02qBN3iKW1OzL3MGk2IdtZzaj7SFntXj72NppTA=
 gopkg.in/yaml.v2 v2.2.1 h1:mUhvW9EsL+naU5Q3cakzfE91YhliOondGd6ZrsDBHQE=
 gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=

hub/connection.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/9seconds/mtg/conntypes"
 	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/mtproto/rpc"
 	"github.com/9seconds/mtg/protocol"
 )
 
@@ -73,6 +74,8 @@ func (c *connection) idle() bool {
 }
 
 func (c *connection) run() {
+	logger := c.hub.logger.Named("connection").With("id", c.id)
+
 	for {
 		packet, err := c.read()
 		if err != nil {
@@ -80,8 +83,18 @@ func (c *connection) run() {
 			return
 		}
 
-		if channel, ok := Registry.getChannel(conntypes.ConnID{}); ok {
+		response, err := rpc.ParseProxyResponse(packet)
-			go channel.write(packet) // nolint: errcheck
+		if err != nil {
+			logger.Debugw("Failed response", "error", err)
+			continue
+		}
+		if response.Type == rpc.ProxyResponseTypeCloseExt {
+			logger.Debugw("Proxy has closed connection")
+			return
+		}
+
+		if channel, ok := Registry.getChannel(response.ConnID); ok {
+			go channel.sendBack(response) // nolint: errcheck
 		}
 	}
 }

hub/connection_hub.go

@@ -3,6 +3,8 @@ package hub
 import (
 	"time"
 
+	"go.uber.org/zap"
+
 	"github.com/9seconds/mtg/protocol"
 )
 
@@ -15,6 +17,7 @@ type connectionHubRequest struct {
 
 type connectionHub struct {
 	sockets map[int]*connection
+	logger  *zap.SugaredLogger
 
 	channelBrokenSockets      chan int
 	channelConnectionRequests chan *connectionHubRequest
@@ -76,8 +79,9 @@ func (c *connectionHub) runReturnConnection(conn *connection) {
 	c.sockets[conn.id] = conn
 }
 
-func newConnectionHub() *connectionHub {
+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),

hub/ctx_channel.go

@@ -4,23 +4,23 @@ import (
 	"context"
 	"time"
 
-	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/mtproto/rpc"
 )
 
 const closeableChannelReadTimeout = 2 * time.Minute
 
 type ChannelReadCloser interface {
-	Read() (conntypes.Packet, error)
+	Read() (*rpc.ProxyResponse, error)
 	Close() error
 }
 
 type ctxChannel struct {
-	channel chan conntypes.Packet
+	channel chan *rpc.ProxyResponse
 	ctx     context.Context
 	cancel  context.CancelFunc
 }
 
-func (c *ctxChannel) Read() (conntypes.Packet, error) {
+func (c *ctxChannel) Read() (*rpc.ProxyResponse, error) {
 	timer := time.NewTimer(closeableChannelReadTimeout)
 	defer timer.Stop()
 
@@ -34,11 +34,11 @@ func (c *ctxChannel) Read() (conntypes.Packet, error) {
 	}
 }
 
-func (c *ctxChannel) write(packet conntypes.Packet) error {
+func (c *ctxChannel) sendBack(response *rpc.ProxyResponse) error {
 	select {
 	case <-c.ctx.Done():
 		return ErrClosed
-	case c.channel <- packet:
+	case c.channel <- response:
 		return nil
 	}
 }
@@ -52,7 +52,7 @@ func (c *ctxChannel) Close() error {
 func newCtxChannel(ctx context.Context) *ctxChannel {
 	ctx, cancel := context.WithCancel(ctx)
 	return &ctxChannel{
-		channel: make(chan conntypes.Packet),
+		channel: make(chan *rpc.ProxyResponse),
 		ctx:     ctx,
 		cancel:  cancel,
 	}

hub/hub.go

@@ -6,13 +6,16 @@ import (
 	"strings"
 	"sync"
 
+	"go.uber.org/zap"
+
 	"github.com/9seconds/mtg/conntypes"
 	"github.com/9seconds/mtg/protocol"
 )
 
 type hub struct {
-	subs  map[string]*connectionHub
+	logger *zap.SugaredLogger
-	mutex sync.RWMutex
+	subs   map[string]*connectionHub
+	mutex  sync.RWMutex
 }
 
 func (h *hub) Write(packet conntypes.Packet, req *protocol.TelegramRequest) error {
@@ -51,7 +54,10 @@ func (h *hub) getHub(req *protocol.TelegramRequest) *connectionHub {
 
 		rv, ok = h.subs[key]
 		if !ok {
-			rv = newConnectionHub()
+			rv = newConnectionHub(h.logger.With(
+				"dc", req.ClientProtocol.DC(),
+				"protocol", req.ClientProtocol.ConnectionProtocol(),
+			))
 			h.subs[key] = rv
 		}
 	}

hub/init.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"errors"
 	"sync"
+
+	"go.uber.org/zap"
 )
 
 var (
@@ -24,7 +26,8 @@ func Init(ctx context.Context) {
 			ctx:   ctx,
 		}
 		Hub = &hub{
-			subs: map[string]*connectionHub{},
+			subs:   map[string]*connectionHub{},
+			logger: zap.S().Named("hub"),
 		}
 	})
 }

mtproto/protocol.go

@@ -7,7 +7,8 @@ import (
 	"github.com/9seconds/mtg/mtproto/rpc"
 	"github.com/9seconds/mtg/protocol"
 	"github.com/9seconds/mtg/telegram"
-	"github.com/9seconds/mtg/wrappers"
+	"github.com/9seconds/mtg/wrappers/packet"
+	"github.com/9seconds/mtg/wrappers/stream"
 )
 
 func TelegramProtocol(req *protocol.TelegramRequest) (conntypes.PacketReadWriteCloser, error) {
@@ -17,7 +18,7 @@ func TelegramProtocol(req *protocol.TelegramRequest) (conntypes.PacketReadWriteC
 		return nil, fmt.Errorf("cannot connect to telegram: %w", err)
 	}
 
-	rpcNonceConn := wrappers.NewMtprotoFrame(conn, rpc.SeqNoNonce)
+	rpcNonceConn := packet.NewMtprotoFrame(conn, rpc.SeqNoNonce)
 	rpcNonceReq, err := doRPCNonceRequest(rpcNonceConn)
 	if err != nil {
 		return nil, fmt.Errorf("cannot do nonce request: %w", err)
@@ -28,8 +29,8 @@ func TelegramProtocol(req *protocol.TelegramRequest) (conntypes.PacketReadWriteC
 		return nil, fmt.Errorf("cannot get nonce response: %w", err)
 	}
 
-	secureConn := wrappers.NewMiddleProxyCipher(conn, rpcNonceReq, rpcNonceResp, telegram.Middle.Secret())
+	secureConn := stream.NewMiddleProxyCipher(conn, rpcNonceReq, rpcNonceResp, telegram.Middle.Secret())
-	frameConn := wrappers.NewMtprotoFrame(secureConn, rpc.SeqNoHandshake)
+	frameConn := packet.NewMtprotoFrame(secureConn, rpc.SeqNoHandshake)
 
 	if err := doRPCHandshakeRequest(frameConn); err != nil {
 		return nil, fmt.Errorf("cannot do handshake request: %w", err)

mtproto/rpc/proxy_flags.go

@@ -5,53 +5,53 @@ import (
 	"strings"
 )
 
-type proxyRequestFlags uint32
+type ProxyRequestFlags uint32
 
 const (
-	proxyRequestFlagsHasAdTag     proxyRequestFlags = 0x8
+	ProxyRequestFlagsHasAdTag     ProxyRequestFlags = 0x8
-	proxyRequestFlagsEncrypted    proxyRequestFlags = 0x2
+	ProxyRequestFlagsEncrypted    ProxyRequestFlags = 0x2
-	proxyRequestFlagsMagic        proxyRequestFlags = 0x1000
+	ProxyRequestFlagsMagic        ProxyRequestFlags = 0x1000
-	proxyRequestFlagsExtMode2     proxyRequestFlags = 0x20000
+	ProxyRequestFlagsExtMode2     ProxyRequestFlags = 0x20000
-	proxyRequestFlagsIntermediate proxyRequestFlags = 0x20000000
+	ProxyRequestFlagsIntermediate ProxyRequestFlags = 0x20000000
-	proxyRequestFlagsAbdridged    proxyRequestFlags = 0x40000000
+	ProxyRequestFlagsAbdridged    ProxyRequestFlags = 0x40000000
-	proxyRequestFlagsQuickAck     proxyRequestFlags = 0x80000000
+	ProxyRequestFlagsQuickAck     ProxyRequestFlags = 0x80000000
-	proxyRequestFlagsPad          proxyRequestFlags = 0x8000000
+	ProxyRequestFlagsPad          ProxyRequestFlags = 0x8000000
 )
 
-var proxyRequestFlagsEncryptedPrefix [8]byte
+var ProxyRequestFlagsEncryptedPrefix [8]byte
 
-func (r proxyRequestFlags) Bytes() []byte {
+func (r ProxyRequestFlags) Bytes() []byte {
 	converted := make([]byte, 4)
 	binary.LittleEndian.PutUint32(converted, uint32(r))
 
 	return converted
 }
 
-func (r proxyRequestFlags) String() string {
+func (r ProxyRequestFlags) String() string {
 	flags := make([]string, 0, 7)
 
-	if r&proxyRequestFlagsHasAdTag != 0 {
+	if r&ProxyRequestFlagsHasAdTag != 0 {
 		flags = append(flags, "HAS_AD_TAG")
 	}
-	if r&proxyRequestFlagsEncrypted != 0 {
+	if r&ProxyRequestFlagsEncrypted != 0 {
 		flags = append(flags, "ENCRYPTED")
 	}
-	if r&proxyRequestFlagsMagic != 0 {
+	if r&ProxyRequestFlagsMagic != 0 {
 		flags = append(flags, "MAGIC")
 	}
-	if r&proxyRequestFlagsExtMode2 != 0 {
+	if r&ProxyRequestFlagsExtMode2 != 0 {
 		flags = append(flags, "EXT_MODE_2")
 	}
-	if r&proxyRequestFlagsIntermediate != 0 {
+	if r&ProxyRequestFlagsIntermediate != 0 {
 		flags = append(flags, "INTERMEDIATE")
 	}
-	if r&proxyRequestFlagsAbdridged != 0 {
+	if r&ProxyRequestFlagsAbdridged != 0 {
 		flags = append(flags, "ABRIDGED")
 	}
-	if r&proxyRequestFlagsQuickAck != 0 {
+	if r&ProxyRequestFlagsQuickAck != 0 {
 		flags = append(flags, "QUICK_ACK")
 	}
-	if r&proxyRequestFlagsPad != 0 {
+	if r&ProxyRequestFlagsPad != 0 {
 		flags = append(flags, "PAD")
 	}
 

mtproto/rpc/proxy_response.go

@@ -0,0 +1,51 @@
+package rpc
+
+import (
+	"bytes"
+	"fmt"
+
+	"github.com/9seconds/mtg/conntypes"
+)
+
+type ProxyResponseType uint8
+
+const (
+	ProxyResponseTypeAns ProxyResponseType = iota
+	ProxyResponseTypeSimpleAck
+	ProxyResponseTypeCloseExt
+)
+
+type ProxyResponse struct {
+	Type    ProxyResponseType
+	ConnID  conntypes.ConnID
+	Payload conntypes.Packet
+}
+
+func ParseProxyResponse(packet conntypes.Packet) (*ProxyResponse, error) {
+	var response ProxyResponse
+
+	if len(packet) < 4 {
+		return nil, fmt.Errorf("incorrect packet length: %d", len(packet))
+	}
+
+	tag := packet[:4]
+	switch {
+	case bytes.Equal(tag, TagProxyAns):
+		response.Type = ProxyResponseTypeAns
+		copy(response.ConnID[:], packet[8:16])
+		response.Payload = packet[16:]
+		return &response, nil
+
+	case bytes.Equal(tag, TagSimpleAck):
+		response.Type = ProxyResponseTypeSimpleAck
+		copy(response.ConnID[:], packet[4:12])
+		response.Payload = packet[12:]
+		return &response, nil
+
+	case bytes.Equal(tag, TagCloseExt):
+		response.Type = ProxyResponseTypeCloseExt
+		return &response, nil
+	}
+
+	return nil, fmt.Errorf("unknown response type %x", tag)
+}

obfuscated2/client_protocol.go

@@ -14,7 +14,7 @@ import (
 	"github.com/9seconds/mtg/conntypes"
 	"github.com/9seconds/mtg/protocol"
 	"github.com/9seconds/mtg/utils"
-	"github.com/9seconds/mtg/wrappers"
+	"github.com/9seconds/mtg/wrappers/stream"
 )
 
 const clientProtocolHandshakeTimeout = 10 * time.Second
@@ -85,7 +85,7 @@ func (c *ClientProtocol) Handshake(socket conntypes.StreamReadWriteCloser) (conn
 	}
 	antireplay.Cache.Add(antiReplayKey)
 
-	return wrappers.NewObfuscated2(socket, encryptor, decryptor), nil
+	return stream.NewObfuscated2(socket, encryptor, decryptor), nil
 }
 
 func (c *ClientProtocol) ReadFrame(socket conntypes.StreamReader) (fm Frame, err error) {

obfuscated2/telegram_protocol.go

@@ -8,7 +8,7 @@ import (
 	"github.com/9seconds/mtg/protocol"
 	"github.com/9seconds/mtg/telegram"
 	"github.com/9seconds/mtg/utils"
-	"github.com/9seconds/mtg/wrappers"
+	"github.com/9seconds/mtg/wrappers/stream"
 )
 
 func TelegramProtocol(req *protocol.TelegramRequest) (conntypes.StreamReadWriteCloser, error) {
@@ -17,8 +17,8 @@ func TelegramProtocol(req *protocol.TelegramRequest) (conntypes.StreamReadWriteC
 	if err != nil {
 		return nil, fmt.Errorf("cannot dial to telegram: %w", err)
 	}
-	conn = wrappers.NewTimeout(conn)
+	conn = stream.NewTimeout(conn)
-	conn = wrappers.NewCtx(req.Ctx, req.Cancel, conn)
+	conn = stream.NewCtx(req.Ctx, req.Cancel, conn)
 	fm := generateFrame(req.ClientProtocol)
 	data := fm.Bytes()
 
@@ -35,7 +35,7 @@ func TelegramProtocol(req *protocol.TelegramRequest) (conntypes.StreamReadWriteC
 		return nil, fmt.Errorf("cannot write handshake frame to telegram: %w", err)
 	}
 
-	return wrappers.NewObfuscated2(conn, encryptor, decryptor), nil
+	return stream.NewObfuscated2(conn, encryptor, decryptor), nil
 }
 
 func generateFrame(cp protocol.ClientProtocol) (fm Frame) {

proxy/proxy.go

@@ -14,7 +14,7 @@ import (
 	"github.com/9seconds/mtg/protocol"
 	"github.com/9seconds/mtg/stats"
 	"github.com/9seconds/mtg/utils"
-	"github.com/9seconds/mtg/wrappers"
+	"github.com/9seconds/mtg/wrappers/stream"
 )
 
 const directPipeBufferSize = 1024 * 1024
@@ -63,9 +63,9 @@ func (p *Proxy) accept(conn net.Conn) {
 	ctx, cancel := context.WithCancel(p.Context)
 	defer cancel()
 
-	clientConn := wrappers.NewClientConn(conn, connID)
+	clientConn := stream.NewClientConn(conn, connID)
-	clientConn = wrappers.NewCtx(ctx, cancel, clientConn)
+	clientConn = stream.NewCtx(ctx, cancel, clientConn)
-	clientConn = wrappers.NewTimeout(clientConn)
+	clientConn = stream.NewTimeout(clientConn)
 	defer clientConn.Close()
 
 	clientProtocol := p.ClientProtocolMaker()

telegram/base.go

@@ -7,7 +7,7 @@ import (
 
 	"github.com/9seconds/mtg/conntypes"
 	"github.com/9seconds/mtg/utils"
-	"github.com/9seconds/mtg/wrappers"
+	"github.com/9seconds/mtg/wrappers/stream"
 )
 
 type baseTelegram struct {
@@ -44,7 +44,7 @@ func (b *baseTelegram) dial(dc conntypes.DC,
 		return nil, fmt.Errorf("cannot initialize tcp socket: %w", err)
 	}
 
-	return wrappers.NewTelegramConn(dc, conn), nil
+	return stream.NewTelegramConn(dc, conn), nil
 }
 
 func (b *baseTelegram) chooseAddress(addresses map[conntypes.DC][]string,

wrappers/packet/mtproto_frame.go

@@ -1,4 +1,4 @@
-package wrappers
+package packet
 
 import (
 	"bytes"

wrappers/packetack/client_abridged.go

@@ -0,0 +1,119 @@
+package packetack
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"net"
+
+	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/utils"
+)
+
+const (
+	clientAbridgedSmallPacketLength = 0x7f
+	clientAbridgedQuickAckLength    = 0x80
+	clientAbridgedLargePacketLength = 16777216 // 256 ^ 3
+)
+
+type wrapperClientAbridged struct {
+	parent conntypes.StreamReadWriteCloser
+}
+
+func (w *wrapperClientAbridged) Read(acks *conntypes.ConnectionAcks) (conntypes.Packet, error) {
+	buf := bytes.Buffer{}
+
+	if _, err := io.CopyN(&buf, w.parent, 1); err != nil {
+		return nil, fmt.Errorf("cannot read message length: %w", err)
+	}
+	msgLength := uint32(buf.Bytes()[0])
+	buf.Reset()
+
+	if msgLength >= clientAbridgedQuickAckLength {
+		acks.Quick = true
+		msgLength -= clientAbridgedQuickAckLength
+	}
+
+	if msgLength == clientAbridgedSmallPacketLength {
+		if _, err := io.CopyN(&buf, w.parent, 3); err != nil {
+			return nil, fmt.Errorf("cannot read correct message length: %w", err)
+		}
+		number := utils.Uint24{}
+		copy(number[:], buf.Bytes())
+		msgLength = utils.FromUint24(number)
+	}
+	msgLength *= 4
+
+	buf.Reset()
+	if _, err := io.CopyN(&buf, w.parent, int64(msgLength)); err != nil {
+		return nil, fmt.Errorf("cannot read message: %w", err)
+	}
+
+	return conntypes.Packet(buf.Bytes()), nil
+}
+
+func (w *wrapperClientAbridged) Write(packet conntypes.Packet, acks *conntypes.ConnectionAcks) error {
+	if len(packet)%4 != 0 {
+		return fmt.Errorf("incorrect packet length %d", len(packet))
+	}
+
+	if acks.Simple {
+		if _, err := w.parent.Write(utils.ReverseBytes(packet)); err != nil {
+			return fmt.Errorf("cannot send a simpleacked packet: %w", err)
+		}
+		return nil
+	}
+
+	packetLength := len(packet) / 4
+	switch {
+	case packetLength < clientAbridgedSmallPacketLength:
+		data := append([]byte{byte(packetLength)}, packet...)
+		if _, err := w.parent.Write(data); err != nil {
+			return fmt.Errorf("cannot send small packet: %w", err)
+		}
+		return nil
+
+	case packetLength < clientAbridgedLargePacketLength:
+		length24 := utils.ToUint24(uint32(packetLength))
+		buf := bytes.Buffer{}
+
+		buf.WriteByte(byte(clientAbridgedSmallPacketLength))
+		buf.Write(length24[:])
+		buf.Write(packet)
+
+		if _, err := w.parent.Write(buf.Bytes()); err != nil {
+			return fmt.Errorf("cannot send large packet: %w", err)
+		}
+		return nil
+	}
+
+	return fmt.Errorf("packet is too big: %d", len(packet))
+}
+
+func (w *wrapperClientAbridged) Close() error {
+	return w.parent.Close()
+}
+
+func (w *wrapperClientAbridged) Conn() net.Conn {
+	return w.parent.Conn()
+}
+
+func (w *wrapperClientAbridged) LocalAddr() *net.TCPAddr {
+	return w.parent.LocalAddr()
+}
+
+func (w *wrapperClientAbridged) RemoteAddr() *net.TCPAddr {
+	return w.parent.RemoteAddr()
+}
+
+func (w *wrapperClientAbridged) Logger() *zap.SugaredLogger {
+	return w.parent.Logger().Named("client-abridged")
+}
+
+func NewClientAbridged(parent conntypes.StreamReadWriteCloser) conntypes.PacketAckFullReadWriteCloser {
+	return &wrapperClientAbridged{
+		parent: parent,
+	}
+}

wrappers/packetack/proxy.go

@@ -0,0 +1,99 @@
+package packetack
+
+import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
+	"net"
+
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/hub"
+	"github.com/9seconds/mtg/mtproto/rpc"
+	"github.com/9seconds/mtg/protocol"
+)
+
+type wrapperProxy struct {
+	flags        rpc.ProxyRequestFlags
+	request      *protocol.TelegramRequest
+	clientIPPort []byte
+	ourIPPort    []byte
+	channelRead  hub.ChannelReadCloser
+}
+
+func (w *wrapperProxy) Write(packet conntypes.Packet, acks *conntypes.ConnectionAcks) error {
+	buf := bytes.Buffer{}
+
+	flags := w.flags
+	if acks.Quick {
+		flags |= rpc.ProxyRequestFlagsQuickAck
+	}
+	if bytes.HasPrefix(packet, rpc.ProxyRequestFlagsEncryptedPrefix[:]) {
+		flags |= rpc.ProxyRequestFlagsEncrypted
+	}
+
+	buf.Write(rpc.TagProxyRequest)
+	buf.Write(flags.Bytes())
+	buf.Write(w.request.ConnID[:])
+	buf.Write(w.clientIPPort)
+	buf.Write(w.ourIPPort)
+	buf.Write(rpc.ProxyRequestExtraSize)
+	buf.Write(rpc.ProxyRequestProxyTag)
+	buf.WriteByte(byte(len(config.C.AdTag)))
+	buf.Write(config.C.AdTag)
+	buf.Write(make([]byte, (4-buf.Len()%4)%4))
+	buf.Write(packet)
+
+	return hub.Hub.Write(buf.Bytes(), w.request)
+}
+
+func (w *wrapperProxy) Read(acks *conntypes.ConnectionAcks) (conntypes.Packet, error) {
+	resp, err := w.channelRead.Read()
+	if err != nil {
+		return nil, fmt.Errorf("cannot read a response: %w", err)
+	}
+
+	if resp.Type == rpc.ProxyResponseTypeSimpleAck {
+		acks.Simple = true
+	}
+
+	return resp.Payload, nil
+}
+
+func (w *wrapperProxy) Close() error {
+	return w.channelRead.Close()
+}
+
+func NewProxy(request *protocol.TelegramRequest) conntypes.PacketAckReadWriteCloser {
+	flags := rpc.ProxyRequestFlagsHasAdTag | rpc.ProxyRequestFlagsMagic | rpc.ProxyRequestFlagsExtMode2
+
+	switch request.ClientProtocol.ConnectionType() {
+	case conntypes.ConnectionTypeAbridged:
+		flags |= rpc.ProxyRequestFlagsAbdridged
+	case conntypes.ConnectionTypeIntermediate:
+		flags |= rpc.ProxyRequestFlagsIntermediate
+	case conntypes.ConnectionTypeSecure:
+		flags |= rpc.ProxyRequestFlagsIntermediate | rpc.ProxyRequestFlagsPad
+	default:
+		panic("unknown connection type")
+	}
+
+	return &wrapperProxy{
+		flags:        flags,
+		request:      request,
+		channelRead:  hub.Registry.Register(request.ConnID),
+		clientIPPort: proxyGetIPPort(request.ClientConn.RemoteAddr()),
+		ourIPPort:    proxyGetIPPort(request.ClientConn.LocalAddr()),
+	}
+}
+
+func proxyGetIPPort(addr *net.TCPAddr) []byte {
+	rv := [16 + 4]byte{}
+	port := [4]byte{}
+
+	copy(rv[:16], addr.IP.To16())
+	binary.LittleEndian.PutUint32(port[:], uint32(addr.Port))
+	copy(rv[16:], port[:])
+
+	return rv[:]
+}

wrappers/stream/base.go

@@ -1,4 +1,4 @@
-package wrappers
+package stream
 
 import (
 	"net"

wrappers/stream/blockcipher.go

@@ -1,4 +1,4 @@
-package wrappers
+package stream
 
 import (
 	"bytes"

wrappers/stream/conn.go

@@ -1,4 +1,4 @@
-package wrappers
+package stream
 
 import (
 	"fmt"

wrappers/stream/ctx.go

@@ -1,4 +1,4 @@
-package wrappers
+package stream
 
 import (
 	"context"

wrappers/stream/mtproto_cipher.go

@@ -1,4 +1,4 @@
-package wrappers
+package stream
 
 import (
 	"bytes"

wrappers/stream/obfuscated2.go

@@ -1,4 +1,4 @@
-package wrappers
+package stream
 
 import (
 	"crypto/cipher"

wrappers/stream/stats_telegram.go

@@ -1,13 +1,14 @@
-package wrappers
+package stream
 
 import (
 	"net"
 	"sync"
 	"time"
 
+	"go.uber.org/zap"
+
 	"github.com/9seconds/mtg/conntypes"
 	"github.com/9seconds/mtg/stats"
-	"go.uber.org/zap"
 )
 
 type wrapperTelegramStats struct {

wrappers/stream/stats_traffic.go

@@ -1,4 +1,4 @@
-package wrappers
+package stream
 
 import (
 	"net"

wrappers/stream/timeout.go

@@ -1,4 +1,4 @@
-package wrappers
+package stream
 
 import (
 	"net"