Move telegram to separate module

obfuscated2/frame.go

@@ -58,13 +58,7 @@ func (f Frame) DC() (n int16) {
 		n = 1
 	}
 
-	if n < 0 {
+	return
-		n = -n
-	} else if n == 0 {
-		n = 1
-	}
-
-	return n - 1
 }
 
 // Valid checks that *decrypted* frame is valid. Only magic bytes are checked.

proxy/server.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/obfuscated2"
+	"github.com/9seconds/mtg/telegram"
 	"github.com/9seconds/mtg/wrappers"
 )
 
@@ -20,6 +21,7 @@ type Server struct {
 	conf   *config.Config
 	logger *zap.SugaredLogger
 	stats  *Stats
+	tg     telegram.Telegram
 }
 
 // Serve does MTPROTO proxying.
@@ -118,23 +120,21 @@ func (s *Server) getClientStream(ctx context.Context, cancel context.CancelFunc,
 }
 
 func (s *Server) getTelegramStream(ctx context.Context, cancel context.CancelFunc, dc int16, socketID string) (io.ReadWriteCloser, error) {
-	socket, err := dialToTelegram(dc, s.conf.TimeoutRead)
+	conn, err := s.tg.Dial(dc)
 	if err != nil {
-		return nil, errors.Annotate(err, "Cannot dial")
+		return nil, errors.Annotate(err, "Cannot connect to Telegram")
-	}
-	wConn := wrappers.NewTimeoutRWC(socket, s.conf.TimeoutRead, s.conf.TimeoutWrite)
-	wConn = wrappers.NewTrafficRWC(wConn, s.stats.addIncomingTraffic, s.stats.addOutgoingTraffic)
-
-	obfs2, frame := obfuscated2.MakeTelegramObfuscated2Frame()
-	if n, err := socket.Write(frame); err != nil || n != len(frame) {
-		return nil, errors.Annotate(err, "Cannot write hadnshake frame")
 	}
 
-	wConn = wrappers.NewLogRWC(wConn, s.logger, socketID, "telegram")
+	conn = wrappers.NewTrafficRWC(conn, s.stats.addIncomingTraffic, s.stats.addOutgoingTraffic)
-	wConn = wrappers.NewStreamCipherRWC(wConn, obfs2.Encryptor, obfs2.Decryptor)
+	conn, err = s.tg.Init(conn)
-	wConn = wrappers.NewCtxRWC(ctx, cancel, wConn)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot handshake Telegram")
+	}
 
-	return wConn, nil
+	conn = wrappers.NewLogRWC(conn, s.logger, socketID, "telegram")
+	conn = wrappers.NewCtxRWC(ctx, cancel, conn)
+
+	return conn, nil
 }
 
 // NewServer creates new instance of MTPROTO proxy.
@@ -143,5 +143,6 @@ func NewServer(conf *config.Config, logger *zap.SugaredLogger, stat *Stats) *Ser
 		conf:   conf,
 		logger: logger,
 		stats:  stat,
+		tg:     telegram.NewDirectTelegram(conf),
 	}
 }

proxy/telegram.go

@@ -1,73 +0,0 @@
-package proxy
-
-import (
-	"net"
-	"time"
-
-	"github.com/juju/errors"
-)
-
-// TelegramAddress presents a pair of v4 and v6 addresses. This pairization
-// is required because we want to use DC indexes.
-type TelegramAddress struct {
-	v4 string
-	v6 string
-}
-
-// IPv4 returns v4 address.
-func (t *TelegramAddress) IPv4() string {
-	return net.JoinHostPort(t.v4, telegramPort)
-}
-
-// IPv6 returns v4 address.
-func (t *TelegramAddress) IPv6() string {
-	return net.JoinHostPort(t.v6, telegramPort)
-}
-
-// TelegramAddresses is a list of all known Telegram addresses for DC indexes.
-var TelegramAddresses = []TelegramAddress{
-	TelegramAddress{v4: "149.154.175.50", v6: "2001:b28:f23d:f001::a"},
-	TelegramAddress{v4: "149.154.167.51", v6: "2001:67c:04e8:f002::a"},
-	TelegramAddress{v4: "149.154.175.100", v6: "2001:b28:f23d:f003::a"},
-	TelegramAddress{v4: "149.154.167.91", v6: "2001:67c:04e8:f004::a"},
-	TelegramAddress{v4: "149.154.171.5", v6: "2001:b28:f23f:f005::a"},
-}
-
-const telegramPort = "443"
-
-const telegramKeepAlive = 30 * time.Second
-
-func dialToTelegram(dcIdx int16, timeout time.Duration) (net.Conn, error) {
-	if dcIdx < 0 || dcIdx >= 5 {
-		return nil, errors.New("Incorrect DC IDX")
-	}
-
-	conn, err := doDial(dcIdx, timeout)
-	if err != nil {
-		return nil, errors.Annotate(err, "Cannot dial")
-	}
-
-	if err := conn.SetKeepAlive(true); err != nil {
-		return nil, errors.Annotate(err, "Cannot establish keepalive connection")
-	}
-	if err := conn.SetKeepAlivePeriod(telegramKeepAlive); err != nil {
-		return nil, errors.Annotate(err, "Cannot set keepalive timeout")
-	}
-
-	return conn, nil
-}
-
-func doDial(dcIdx int16, timeout time.Duration) (*net.TCPConn, error) {
-	dialer := net.Dialer{Timeout: timeout}
-	addr := TelegramAddresses[dcIdx]
-
-	if conn, err := dialer.Dial("tcp", addr.IPv6()); err == nil {
-		return conn.(*net.TCPConn), nil
-	}
-
-	conn, err := dialer.Dial("tcp", addr.IPv4())
-	if err == nil {
-		return conn.(*net.TCPConn), nil
-	}
-	return nil, err
-}

telegram/dialer.go

@@ -0,0 +1,53 @@
+package telegram
+
+import (
+	"io"
+	"net"
+	"time"
+
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+const telegramKeepAlive = 30 * time.Second
+
+type tgDialer struct {
+	net.Dialer
+
+	conf *config.Config
+}
+
+func (t *tgDialer) dial(addr string) (net.Conn, error) {
+	connRaw, err := t.Dialer.Dial("tcp", addr)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot connect to Telegram")
+	}
+	conn := connRaw.(*net.TCPConn)
+
+	if err = conn.SetKeepAlive(true); err != nil {
+		return nil, errors.Annotate(err, "Cannot establish keepalive connection")
+	}
+	if err = conn.SetKeepAlivePeriod(telegramKeepAlive); err != nil {
+		return nil, errors.Annotate(err, "Cannot set keepalive timeout")
+	}
+
+	return conn, nil
+}
+
+func (t *tgDialer) dialRWC(addr string) (io.ReadWriteCloser, error) {
+	conn, err := t.dial(addr)
+	if err != nil {
+		return nil, err
+	}
+
+	return wrappers.NewTimeoutRWC(conn, t.conf.TimeoutRead, t.conf.TimeoutWrite), nil
+}
+
+func newDialer(conf *config.Config) *tgDialer {
+	return &tgDialer{
+		Dialer: net.Dialer{Timeout: conf.TimeoutRead},
+		conf:   conf,
+	}
+}

telegram/direct.go

@@ -0,0 +1,59 @@
+package telegram
+
+import (
+	"io"
+
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/obfuscated2"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+var (
+	directV4Addresses = map[int16][]string{
+		0: []string{"149.154.175.50:443"},
+		1: []string{"149.154.167.51:443"},
+		2: []string{"149.154.175.100:443"},
+		3: []string{"149.154.167.91:443"},
+		4: []string{"149.154.171.5:443"},
+	}
+	directV6Addresses = map[int16][]string{
+		0: []string{"[2001:b28:f23d:f001::a]:443"},
+		1: []string{"[2001:67c:04e8:f002::a]:443"},
+		2: []string{"[2001:b28:f23d:f003::a]:443"},
+		3: []string{"[2001:67c:04e8:f004::a]:443"},
+		4: []string{"[2001:b28:f23f:f005::a]:443"},
+	}
+)
+
+type directTelegram struct {
+	baseTelegram
+}
+
+func (t *directTelegram) Dial(dcIdx int16) (io.ReadWriteCloser, error) {
+	if dcIdx < 0 {
+		dcIdx = -dcIdx
+	} else if dcIdx == 0 {
+		dcIdx = 1
+	}
+
+	return t.baseTelegram.Dial(dcIdx - 1)
+}
+
+func (t *directTelegram) Init(conn io.ReadWriteCloser) (io.ReadWriteCloser, error) {
+	obfs2, frame := obfuscated2.MakeTelegramObfuscated2Frame()
+	if n, err := conn.Write(frame); err != nil || n != len(frame) {
+		return nil, errors.Annotate(err, "Cannot write hadnshake frame")
+	}
+
+	return wrappers.NewStreamCipherRWC(conn, obfs2.Encryptor, obfs2.Decryptor), nil
+}
+
+func NewDirectTelegram(conf *config.Config) Telegram {
+	return &directTelegram{baseTelegram{
+		dialer:      newDialer(conf),
+		v4Addresses: directV4Addresses,
+		v6Addresses: directV6Addresses,
+	}}
+}

telegram/telegram.go

@@ -0,0 +1,38 @@
+package telegram
+
+import (
+	"io"
+	"math/rand"
+
+	"github.com/juju/errors"
+)
+
+type Telegram interface {
+	Dial(int16) (io.ReadWriteCloser, error)
+	Init(io.ReadWriteCloser) (io.ReadWriteCloser, error)
+}
+
+type baseTelegram struct {
+	dialer *tgDialer
+
+	v4Addresses map[int16][]string
+	v6Addresses map[int16][]string
+}
+
+func (b *baseTelegram) Dial(dcIdx int16) (io.ReadWriteCloser, error) {
+	addrs := make([]string, 2)
+	if addr, ok := b.v6Addresses[dcIdx]; ok && len(addr) > 0 {
+		addrs = append(addrs, addr[rand.Intn(len(addr))])
+	}
+	if addr, ok := b.v4Addresses[dcIdx]; ok && len(addr) > 0 {
+		addrs = append(addrs, addr[rand.Intn(len(addr))])
+	}
+
+	for _, addr := range addrs {
+		if conn, err := b.dialer.dialRWC(addr); err == nil {
+			return conn, err
+		}
+	}
+
+	return nil, errors.New("Cannot connect to Telegram")
+}