Add client abridged protocol

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"