Add client package

client/client.go

@@ -0,0 +1,10 @@
+package client
+
+import (
+	"io"
+	"net"
+
+	"github.com/9seconds/mtg/config"
+)
+
+type Init func(net.Conn, *config.Config) (int16, io.ReadWriteCloser, error)

client/direct.go

@@ -0,0 +1,29 @@
+package client
+
+import (
+	"io"
+	"net"
+
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/obfuscated2"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+func DirectInit(conn net.Conn, conf *config.Config) (int16, io.ReadWriteCloser, error) {
+	socket := wrappers.NewTimeoutRWC(conn, conf.TimeoutRead, conf.TimeoutWrite)
+	frame, err := obfuscated2.ExtractFrame(socket)
+	if err != nil {
+		return 0, nil, errors.Annotate(err, "Cannot extract frame")
+	}
+
+	obfs2, dc, err := obfuscated2.ParseObfuscated2ClientFrame(conf.Secret, frame)
+	if err != nil {
+		return 0, nil, errors.Annotate(err, "Cannot parse obfuscated frame")
+	}
+
+	socket = wrappers.NewStreamCipherRWC(socket, obfs2.Encryptor, obfs2.Decryptor)
+
+	return dc, socket, nil
+}

obfuscated2/frame_test.go

@@ -34,7 +34,7 @@ func TestFrameMagic(t *testing.T) {
 }
 
 func TestFrameDC(t *testing.T) {
-	assert.Equal(t, int16(770), makeFrame().DC())
+	assert.Equal(t, int16(771), makeFrame().DC())
 }
 
 func TestFrameValid(t *testing.T) {

obfuscated2/obfuscated2.go

@@ -19,9 +19,7 @@ type Obfuscated2 struct {
 // details: http://telegra.ph/telegram-blocks-wtf-05-26
 //
 // Beware, link above is in russian.
-func ParseObfuscated2ClientFrame(secret, data []byte) (*Obfuscated2, int16, error) {
-	frame := Frame(data)
-
+func ParseObfuscated2ClientFrame(secret []byte, frame Frame) (*Obfuscated2, int16, error) {
 	decHasher := sha256.New()
 	decHasher.Write(frame.Key()) // nolint: errcheck
 	decHasher.Write(secret)      // nolint: errcheck

proxy/server.go

@@ -10,8 +10,8 @@ import (
 	uuid "github.com/satori/go.uuid"
 	"go.uber.org/zap"
 
+	"github.com/9seconds/mtg/client"
 	"github.com/9seconds/mtg/config"
-	"github.com/9seconds/mtg/obfuscated2"
 	"github.com/9seconds/mtg/telegram"
 	"github.com/9seconds/mtg/wrappers"
 )
@@ -22,6 +22,7 @@ type Server struct {
 	logger     *zap.SugaredLogger
 	stats      *Stats
 	tg         telegram.Telegram
+	clientInit client.Init
 }
 
 // Serve does MTPROTO proxying.
@@ -59,7 +60,7 @@ func (s *Server) accept(conn net.Conn) {
 		"socketid", socketID,
 	)
 
-	clientConn, dc, err := s.getClientStream(ctx, cancel, conn, socketID)
+	dc, clientConn, err := s.getClientStream(ctx, cancel, conn, socketID)
 	if err != nil {
 		s.logger.Warnw("Cannot initialize client connection",
 			"addr", conn.RemoteAddr().String(),
@@ -99,24 +100,17 @@ func (s *Server) accept(conn net.Conn) {
 	)
 }
 
-func (s *Server) getClientStream(ctx context.Context, cancel context.CancelFunc, conn net.Conn, socketID string) (io.ReadWriteCloser, int16, error) {
-	wConn := wrappers.NewTimeoutRWC(conn, s.conf.TimeoutRead, s.conf.TimeoutWrite)
-	wConn = wrappers.NewTrafficRWC(wConn, s.stats.addIncomingTraffic, s.stats.addOutgoingTraffic)
-	frame, err := obfuscated2.ExtractFrame(wConn)
+func (s *Server) getClientStream(ctx context.Context, cancel context.CancelFunc, conn net.Conn, socketID string) (int16, io.ReadWriteCloser, error) {
+	dc, socket, err := s.clientInit(conn, s.conf)
 	if err != nil {
-		return nil, 0, errors.Annotate(err, "Cannot create client stream")
+		return 0, nil, errors.Annotate(err, "Cannot init client connection")
 	}
 
-	obfs2, dc, err := obfuscated2.ParseObfuscated2ClientFrame(s.conf.Secret, frame)
-	if err != nil {
-		return nil, 0, errors.Annotate(err, "Cannot create client stream")
-	}
+	socket = wrappers.NewTrafficRWC(socket, s.stats.addIncomingTraffic, s.stats.addOutgoingTraffic)
+	socket = wrappers.NewLogRWC(socket, s.logger, socketID, "client")
+	socket = wrappers.NewCtxRWC(ctx, cancel, socket)
 
-	wConn = wrappers.NewLogRWC(wConn, s.logger, socketID, "client")
-	wConn = wrappers.NewStreamCipherRWC(wConn, obfs2.Encryptor, obfs2.Decryptor)
-	wConn = wrappers.NewCtxRWC(ctx, cancel, wConn)
-
-	return wConn, dc, nil
+	return dc, socket, nil
 }
 
 func (s *Server) getTelegramStream(ctx context.Context, cancel context.CancelFunc, dc int16, socketID string) (io.ReadWriteCloser, error) {
@@ -144,5 +138,6 @@ func NewServer(conf *config.Config, logger *zap.SugaredLogger, stat *Stats) *Ser
 		logger:     logger,
 		stats:      stat,
 		tg:         telegram.NewDirectTelegram(conf),
+		clientInit: client.DirectInit,
 	}
 }