Refactored all the things!

client/client.go

@@ -1,12 +0,0 @@
-package client
-
-import (
-	"net"
-
-	"github.com/9seconds/mtg/config"
-	"github.com/9seconds/mtg/mtproto"
-	"github.com/9seconds/mtg/wrappers"
-)
-
-// Init has to initialize client connection based on given config.
-type Init func(net.Conn, string, *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error)

client/direct.go

@@ -1,6 +1,7 @@
 package client
 
 import (
+	"context"
 	"net"
 	"time"
 
@@ -14,28 +15,29 @@ import (
 
 const handshakeTimeout = 10 * time.Second
 
-// DirectInit initializes client to access Telegram bypassing middleproxies.
-func DirectInit(conn net.Conn, socketID string, conf *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error) {
-	if err := config.SetSocketOptions(conn); err != nil {
+func DirectInit(ctx context.Context, cancel context.CancelFunc, socket net.Conn, connID string,
+	conf *config.Config) (wrappers.WrapStreamReadWriteCloser, *mtproto.ConnectionOpts, error) {
+	if err := config.SetSocketOptions(socket); err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot set socket options")
 	}
 
-	conn.SetReadDeadline(time.Now().Add(handshakeTimeout)) // nolint: errcheck
-	frame, err := obfuscated2.ExtractFrame(conn)
-	conn.SetReadDeadline(time.Time{}) // nolint: errcheck
+	socket.SetReadDeadline(time.Now().Add(handshakeTimeout))
+	frame, err := obfuscated2.ExtractFrame(socket)
 	if err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot extract frame")
 	}
+	socket.SetReadDeadline(time.Time{})
+	conn := wrappers.NewConn(socket, connID, wrappers.ConnPurposeClient, conf.PublicIPv4, conf.PublicIPv6)
 
 	obfs2, connOpts, err := obfuscated2.ParseObfuscated2ClientFrame(conf.Secret, frame)
 	if err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot parse obfuscated frame")
 	}
 	connOpts.ConnectionProto = mtproto.ConnectionProtocolAny
-	connOpts.ClientAddr = conn.RemoteAddr().(*net.TCPAddr)
+	connOpts.ClientAddr = conn.RemoteAddr()
 
-	socket := wrappers.NewTimeoutRWC(conn, socketID, conf.PublicIPv4, conf.PublicIPv6)
-	socket = wrappers.NewStreamCipherRWC(socket, obfs2.Encryptor, obfs2.Decryptor)
+	conn = wrappers.NewCtx(ctx, cancel, conn)
+	conn = wrappers.NewStreamCipher(conn, obfs2.Encryptor, obfs2.Decryptor)
 
-	return socket, connOpts, nil
+	return conn, connOpts, nil
 }

client/middle.go

@@ -1,30 +1,30 @@
 package client
 
 import (
+	"context"
 	"net"
 
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
-	mtwrappers "github.com/9seconds/mtg/mtproto/wrappers"
 	"github.com/9seconds/mtg/wrappers"
 )
 
-func MiddleInit(conn net.Conn, socketID string, conf *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error) {
-	newConn, opts, err := DirectInit(conn, socketID, conf)
+func MiddleInit(ctx context.Context, cancel context.CancelFunc, socket net.Conn, connID string,
+	conf *config.Config) (wrappers.WrapPacketReadWriteCloser, *mtproto.ConnectionOpts, error) {
+	conn, opts, err := DirectInit(ctx, cancel, socket, connID, conf)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	if opts.ConnectionType == mtproto.ConnectionTypeAbridged {
-		newConn = mtwrappers.NewAbridgedRWC(newConn, opts)
-	} else {
-		newConn = mtwrappers.NewIntermediateRWC(newConn, opts)
+	newConn := wrappers.NewMTProtoAbridged(conn, opts)
+	if opts.ConnectionType != mtproto.ConnectionTypeAbridged {
+		newConn = wrappers.NewMTProtoIntermediate(conn, opts)
 	}
 
 	opts.ConnectionProto = mtproto.ConnectionProtocolIPv4
-	if conn.LocalAddr().(*net.TCPAddr).IP.To4() == nil {
+	if socket.LocalAddr().(*net.TCPAddr).IP.To4() == nil {
 		opts.ConnectionProto = mtproto.ConnectionProtocolIPv6
 	}
 
-	return newConn, opts, nil
+	return newConn, opts, err
 }