Add intermediate secure wrapper

client/middle.go

@@ -17,9 +17,16 @@ func MiddleInit(socket net.Conn, connID string, conf *config.Config) (wrappers.W
 	}
 	connStream := conn.(wrappers.StreamReadWriteCloser)
 
-	newConn := wrappers.NewMTProtoAbridged(connStream, opts)
-	if opts.ConnectionType != mtproto.ConnectionTypeAbridged {
+	var newConn wrappers.PacketReadWriteCloser
+	switch opts.ConnectionType {
+	case mtproto.ConnectionTypeAbridged:
+		newConn = wrappers.NewMTProtoAbridged(connStream, opts)
+	case mtproto.ConnectionTypeIntermediate:
 		newConn = wrappers.NewMTProtoIntermediate(connStream, opts)
+	case mtproto.ConnectionTypeSecure:
+		newConn = wrappers.NewMTProtoIntermediateSecure(connStream, opts)
+	default:
+		panic("Unknown connection type")
 	}
 
 	opts.ConnectionProto = mtproto.ConnectionProtocolIPv4

wrappers/mtproto_intermediate_secure.go

@@ -0,0 +1,69 @@
+package wrappers
+
+import (
+	"bytes"
+	"encoding/binary"
+	"math/rand"
+
+	"github.com/9seconds/mtg/mtproto"
+)
+
+type MTProtoIntermediateSecure struct {
+	MTProtoIntermediate
+}
+
+func (m *MTProtoIntermediateSecure) Read() ([]byte, error) {
+	data, err := m.MTProtoIntermediate.Read()
+	if err != nil {
+		return nil, err
+	}
+	length := len(data) - (len(data) % 4)
+
+	return data[:length], nil
+}
+
+func (m *MTProtoIntermediateSecure) Write(p []byte) (int, error) {
+	defer func() {
+		m.writeCounter++
+	}()
+
+	m.logger.Debugw("Write packet",
+		"simple_ack", m.opts.WriteHacks.SimpleAck,
+		"quick_ack", m.opts.WriteHacks.QuickAck,
+		"counter", m.writeCounter,
+	)
+
+	if m.opts.ReadHacks.SimpleAck {
+		return m.conn.Write(p)
+	}
+
+	buf := &bytes.Buffer{}
+	paddingLength := rand.Intn(4)
+	buf.Grow(4 + len(p) + paddingLength)
+
+	binary.Write(buf, binary.LittleEndian, uint32(len(p)+paddingLength))
+	buf.Write(p)
+	buf.Write(make([]byte, paddingLength))
+
+	m.logger.Debugw("Write packet with padding",
+		"simple_ack", m.opts.WriteHacks.SimpleAck,
+		"quick_ack", m.opts.WriteHacks.QuickAck,
+		"counter", m.writeCounter,
+		"padding_length", paddingLength,
+		"length", len(p),
+	)
+
+	_, err := m.conn.Write(buf.Bytes())
+
+	return len(p), err
+}
+
+func NewMTProtoIntermediateSecure(conn StreamReadWriteCloser, opts *mtproto.ConnectionOpts) PacketReadWriteCloser {
+	return &MTProtoIntermediateSecure{
+		MTProtoIntermediate: MTProtoIntermediate{
+			conn:   conn,
+			logger: conn.Logger().Named("mtproto-intermediate-secure"),
+			opts:   opts,
+		},
+	}
+}