mirror of
https://github.com/ScuroNeko/mtg.git
synced 2026-08-31 12:04:03 +03:00
Create internal DC package
This commit is contained in:
9seconds
@@ -0,0 +1,10 @@
|
||||
package dc
|
||||
|
||||
type Addr struct {
|
||||
Network string
|
||||
Address string
|
||||
}
|
||||
|
||||
func (d Addr) String() string {
|
||||
return d.Address
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
package dc
|
||||
|
||||
import "math/rand/v2"
|
||||
|
||||
type dcAddrSet struct {
|
||||
v4 map[int][]Addr
|
||||
v6 map[int][]Addr
|
||||
}
|
||||
|
||||
func (d dcAddrSet) getV4(dc int) []Addr {
|
||||
if d.v4 == nil {
|
||||
return nil
|
||||
}
|
||||
return d.get(d.v4[dc])
|
||||
}
|
||||
|
||||
func (d dcAddrSet) getV6(dc int) []Addr {
|
||||
if d.v6 == nil {
|
||||
return nil
|
||||
}
|
||||
return d.get(d.v6[dc])
|
||||
}
|
||||
|
||||
func (d dcAddrSet) get(addrs []Addr) []Addr {
|
||||
otherSet := make([]Addr, 0, len(addrs))
|
||||
otherSet = append(otherSet, addrs...)
|
||||
|
||||
rand.Shuffle(len(otherSet), func(i, j int) {
|
||||
otherSet[i], otherSet[j] = otherSet[j], otherSet[i]
|
||||
})
|
||||
|
||||
return otherSet
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package dc_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/9seconds/mtg/v2/mtglib/internal/dc"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestAddr(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
addr := dc.Addr{Network: "tcp4", Address: "127.0.0.1:443"}
|
||||
|
||||
assert.Equal(t, "127.0.0.1:443", addr.String())
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
package dc
|
||||
|
||||
import "time"
|
||||
|
||||
type preferIP uint8
|
||||
|
||||
const (
|
||||
preferIPOnlyIPv4 preferIP = iota
|
||||
preferIPOnlyIPv6
|
||||
preferIPPreferIPv4
|
||||
preferIPPreferIPv6
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultDC = 2
|
||||
DefaultUpdateDCAddressesEach = time.Hour
|
||||
|
||||
defaultAppID = 123456
|
||||
defaultAppHash = ""
|
||||
)
|
||||
|
||||
type Logger interface {
|
||||
Info(msg string)
|
||||
WarningError(msg string, err error)
|
||||
}
|
||||
|
||||
var (
|
||||
// https://github.com/telegramdesktop/tdesktop/blob/master/Telegram/SourceFiles/mtproto/mtproto_dc_options.cpp#L30
|
||||
defaultDCAddrSet = dcAddrSet{
|
||||
v4: map[int][]Addr{
|
||||
1: {
|
||||
{Network: "tcp4", Address: "149.154.175.50:443"},
|
||||
},
|
||||
2: {
|
||||
{Network: "tcp4", Address: "149.154.167.51:443"},
|
||||
{Network: "tcp4", Address: "95.161.76.100:443"},
|
||||
},
|
||||
3: {
|
||||
{Network: "tcp4", Address: "149.154.175.100:443"},
|
||||
},
|
||||
4: {
|
||||
{Network: "tcp4", Address: "149.154.167.91:443"},
|
||||
},
|
||||
5: {
|
||||
{Network: "tcp4", Address: "149.154.171.5:443"},
|
||||
},
|
||||
},
|
||||
v6: map[int][]Addr{
|
||||
1: {
|
||||
{Network: "tcp6", Address: "[2001:b28:f23d:f001::a]:443"},
|
||||
},
|
||||
2: {
|
||||
{Network: "tcp6", Address: "[2001:67c:04e8:f002::a]:443"},
|
||||
},
|
||||
3: {
|
||||
{Network: "tcp6", Address: "[2001:b28:f23d:f003::a]:443"},
|
||||
},
|
||||
4: {
|
||||
{Network: "tcp6", Address: "[2001:67c:04e8:f004::a]:443"},
|
||||
},
|
||||
5: {
|
||||
{Network: "tcp6", Address: "[2001:b28:f23f:f005::a]:443"},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
defaultDCOverridesAddrSet = dcAddrSet{
|
||||
v4: map[int][]Addr{
|
||||
203: {
|
||||
{Network: "tcp4", Address: "91.105.192.100:443"},
|
||||
},
|
||||
},
|
||||
v6: map[int][]Addr{
|
||||
203: {
|
||||
{Network: "tcp6", Address: "[2a0a:f280:0203:000a:5000:0000:0000:0100]:443"},
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,160 @@
|
||||
package dc
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gotd/td/telegram"
|
||||
)
|
||||
|
||||
type Telegram struct {
|
||||
logger Logger
|
||||
lock sync.RWMutex
|
||||
view dcView
|
||||
preferIP preferIP
|
||||
client *telegram.Client
|
||||
}
|
||||
|
||||
func (t *Telegram) GetAddresses(dc int) []Addr {
|
||||
t.lock.RLock()
|
||||
defer t.lock.RUnlock()
|
||||
|
||||
switch t.preferIP {
|
||||
case preferIPOnlyIPv4:
|
||||
return t.view.getV4(dc)
|
||||
case preferIPOnlyIPv6:
|
||||
return t.view.getV4(dc)
|
||||
case preferIPPreferIPv4:
|
||||
return append(t.view.getV4(dc), t.view.getV6(dc)...)
|
||||
}
|
||||
|
||||
return append(t.view.getV6(dc), t.view.getV4(dc)...)
|
||||
}
|
||||
|
||||
func (t *Telegram) Run(ctx context.Context, updateEach time.Duration) {
|
||||
if updateEach == 0 {
|
||||
updateEach = DefaultUpdateDCAddressesEach
|
||||
}
|
||||
|
||||
t.update(ctx)
|
||||
|
||||
ticker := time.NewTicker(updateEach)
|
||||
defer func() {
|
||||
ticker.Stop()
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-ticker.C:
|
||||
default:
|
||||
}
|
||||
}()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
t.update(ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Telegram) update(ctx context.Context) {
|
||||
collected := dcAddrSet{
|
||||
v4: map[int][]Addr{},
|
||||
v6: map[int][]Addr{},
|
||||
}
|
||||
|
||||
err := t.client.Run(ctx, func(tgctx context.Context) error {
|
||||
conf, err := t.client.API().HelpGetConfig(tgctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, opt := range conf.DCOptions {
|
||||
addr := net.JoinHostPort(opt.IPAddress, strconv.Itoa(opt.Port))
|
||||
|
||||
if opt.Ipv6 {
|
||||
collected.v6[opt.ID] = append(collected.v6[opt.ID], Addr{
|
||||
Network: "tcp6",
|
||||
Address: addr,
|
||||
})
|
||||
} else {
|
||||
collected.v4[opt.ID] = append(collected.v4[opt.ID], Addr{
|
||||
Network: "tcp4",
|
||||
Address: addr,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.logger.WarningError("update has failed", err)
|
||||
return
|
||||
}
|
||||
|
||||
t.lock.Lock()
|
||||
t.view.collected = collected
|
||||
t.lock.Unlock()
|
||||
|
||||
t.logger.Info(fmt.Sprintf("updated DC list: %v", collected))
|
||||
}
|
||||
|
||||
func New(logger Logger, ipPreference string, userOverrides map[int][]string) (*Telegram, error) {
|
||||
var pref preferIP
|
||||
|
||||
switch strings.ToLower(ipPreference) {
|
||||
case "prefer-ipv4":
|
||||
pref = preferIPPreferIPv4
|
||||
case "prefer-ipv6":
|
||||
pref = preferIPPreferIPv6
|
||||
case "only-ipv4":
|
||||
pref = preferIPOnlyIPv4
|
||||
case "only-ipv6":
|
||||
pref = preferIPOnlyIPv6
|
||||
default:
|
||||
return nil, fmt.Errorf("unknown ip preference %s", ipPreference)
|
||||
}
|
||||
|
||||
overrides := dcAddrSet{}
|
||||
for dc, addrs := range userOverrides {
|
||||
for _, addr := range addrs {
|
||||
host, _, err := net.SplitHostPort(addr)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("incorrect host %s: %w", addr, err)
|
||||
}
|
||||
|
||||
parsed := net.ParseIP(host)
|
||||
if parsed == nil {
|
||||
return nil, fmt.Errorf("incorrect host %s", addr)
|
||||
}
|
||||
|
||||
if parsed.To4() != nil {
|
||||
overrides.v4[dc] = append(overrides.v4[dc], Addr{
|
||||
Network: "tcp4",
|
||||
Address: addr,
|
||||
})
|
||||
} else {
|
||||
overrides.v6[dc] = append(overrides.v6[dc], Addr{
|
||||
Network: "tcp6",
|
||||
Address: addr,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return &Telegram{
|
||||
view: dcView{
|
||||
overrides: overrides,
|
||||
},
|
||||
logger: logger,
|
||||
client: telegram.NewClient(defaultAppID, defaultAppHash, telegram.Options{}),
|
||||
preferIP: pref,
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package dc
|
||||
|
||||
type dcView struct {
|
||||
overrides dcAddrSet
|
||||
collected dcAddrSet
|
||||
}
|
||||
|
||||
func (d dcView) getV4(dc int) []Addr {
|
||||
addrs := d.overrides.getV4(dc)
|
||||
addrs = append(addrs, defaultDCOverridesAddrSet.getV4(dc)...)
|
||||
addrs = append(addrs, d.collected.getV4(dc)...)
|
||||
addrs = append(addrs, defaultDCAddrSet.getV4(dc)...)
|
||||
|
||||
return addrs
|
||||
}
|
||||
|
||||
func (d dcView) getV6(dc int) []Addr {
|
||||
addrs := d.overrides.getV6(dc)
|
||||
addrs = append(addrs, defaultDCOverridesAddrSet.getV6(dc)...)
|
||||
addrs = append(addrs, d.collected.getV6(dc)...)
|
||||
addrs = append(addrs, defaultDCAddrSet.getV6(dc)...)
|
||||
|
||||
return addrs
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
package dc
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
type ViewTestSuite struct {
|
||||
suite.Suite
|
||||
|
||||
view dcView
|
||||
}
|
||||
|
||||
func (suite *ViewTestSuite) SetupSuite() {
|
||||
suite.view = dcView{
|
||||
overrides: dcAddrSet{
|
||||
v4: map[int][]Addr{
|
||||
111: {
|
||||
{Network: "tcp4", Address: "127.0.0.1:443"},
|
||||
},
|
||||
203: {
|
||||
{Network: "tcp4", Address: "127.0.0.2:443"},
|
||||
},
|
||||
},
|
||||
v6: map[int][]Addr{
|
||||
203: {
|
||||
{Network: "tcp6", Address: "xxx"},
|
||||
},
|
||||
},
|
||||
},
|
||||
collected: dcAddrSet{
|
||||
v4: map[int][]Addr{
|
||||
1: {
|
||||
{Network: "tcp4", Address: "127.1.0.1:443"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (suite *ViewTestSuite) TestGetV4() {
|
||||
testData := map[int][]Addr{
|
||||
111: {
|
||||
{"tcp4", "127.0.0.1:443"},
|
||||
},
|
||||
203: {
|
||||
{"tcp4", "127.0.0.2:443"},
|
||||
{"tcp4", "91.105.192.100:443"},
|
||||
},
|
||||
2: {
|
||||
{"tcp4", "149.154.167.51:443"},
|
||||
{"tcp4", "95.161.76.100:443"},
|
||||
},
|
||||
1: {
|
||||
{"tcp4", "127.1.0.1:443"},
|
||||
{"tcp4", "149.154.175.50:443"},
|
||||
},
|
||||
}
|
||||
|
||||
for dc, addresses := range testData {
|
||||
suite.T().Run(fmt.Sprintf("dc%d", dc), func(t *testing.T) {
|
||||
assert.ElementsMatch(t, addresses, suite.view.getV4(dc))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (suite *ViewTestSuite) TestGetV6() {
|
||||
testData := map[int][]Addr{
|
||||
111: {},
|
||||
203: {
|
||||
{"tcp6", "xxx"},
|
||||
{"tcp6", "[2a0a:f280:0203:000a:5000:0000:0000:0100]:443"},
|
||||
},
|
||||
1: {
|
||||
{"tcp6", "[2001:b28:f23d:f001::a]:443"},
|
||||
},
|
||||
}
|
||||
|
||||
for dc, addresses := range testData {
|
||||
suite.T().Run(fmt.Sprintf("dc%d", dc), func(t *testing.T) {
|
||||
assert.ElementsMatch(t, addresses, suite.view.getV6(dc))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestView(t *testing.T) {
|
||||
t.Parallel()
|
||||
suite.Run(t, &ViewTestSuite{})
|
||||
}
|
||||
@@ -1,19 +0,0 @@
|
||||
package telegram
|
||||
|
||||
type dcAddresses struct {
|
||||
v4 map[int][]tgAddr
|
||||
v6 map[int][]tgAddr
|
||||
}
|
||||
|
||||
func (a dcAddresses) getV4(dc int) []tgAddr {
|
||||
return a.v4[dc]
|
||||
}
|
||||
|
||||
func (a dcAddresses) getV6(dc int) []tgAddr {
|
||||
return a.v6[dc]
|
||||
}
|
||||
|
||||
func (a dcAddresses) isValidDC(dc int) bool {
|
||||
_, ok := a.v4[dc]
|
||||
return ok
|
||||
}
|
||||
@@ -1,80 +0,0 @@
|
||||
package telegram
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/9seconds/mtg/v2/essentials"
|
||||
)
|
||||
|
||||
var errNoAddresses = errors.New("no addresses")
|
||||
|
||||
type preferIP uint8
|
||||
|
||||
const (
|
||||
preferIPOnlyIPv4 preferIP = iota
|
||||
preferIPOnlyIPv6
|
||||
preferIPPreferIPv4
|
||||
preferIPPreferIPv6
|
||||
)
|
||||
|
||||
const (
|
||||
defaultDC = 2
|
||||
defaultUpdateDCAddressesEach = time.Hour
|
||||
defaultAppID = 123456
|
||||
defaultAppHash = ""
|
||||
)
|
||||
|
||||
type loggerInterface interface {
|
||||
Info(msg string)
|
||||
WarningError(msg string, err error)
|
||||
}
|
||||
|
||||
type tgAddr struct {
|
||||
network string
|
||||
address string
|
||||
}
|
||||
|
||||
// https://github.com/telegramdesktop/tdesktop/blob/master/Telegram/SourceFiles/mtproto/mtproto_dc_options.cpp#L30
|
||||
var (
|
||||
defaultV4Addresses = map[int][]tgAddr{
|
||||
1: {
|
||||
{network: "tcp4", address: "149.154.175.50:443"},
|
||||
},
|
||||
2: {
|
||||
{network: "tcp4", address: "149.154.167.51:443"},
|
||||
{network: "tcp4", address: "95.161.76.100:443"},
|
||||
},
|
||||
3: {
|
||||
{network: "tcp4", address: "149.154.175.100:443"},
|
||||
},
|
||||
4: {
|
||||
{network: "tcp4", address: "149.154.167.91:443"},
|
||||
},
|
||||
5: {
|
||||
{network: "tcp4", address: "149.154.171.5:443"},
|
||||
},
|
||||
}
|
||||
defaultV6Addresses = map[int][]tgAddr{
|
||||
1: {
|
||||
{network: "tcp6", address: "[2001:b28:f23d:f001::a]:443"},
|
||||
},
|
||||
2: {
|
||||
{network: "tcp6", address: "[2001:67c:04e8:f002::a]:443"},
|
||||
},
|
||||
3: {
|
||||
{network: "tcp6", address: "[2001:b28:f23d:f003::a]:443"},
|
||||
},
|
||||
4: {
|
||||
{network: "tcp6", address: "[2001:67c:04e8:f004::a]:443"},
|
||||
},
|
||||
5: {
|
||||
{network: "tcp6", address: "[2001:b28:f23f:f005::a]:443"},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
type Dialer interface {
|
||||
DialContext(ctx context.Context, network, address string) (essentials.Conn, error)
|
||||
}
|
||||
@@ -1,42 +0,0 @@
|
||||
package telegram
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"strconv"
|
||||
|
||||
"github.com/gotd/td/telegram"
|
||||
)
|
||||
|
||||
type rpcClient struct {
|
||||
*telegram.Client
|
||||
}
|
||||
|
||||
func (r rpcClient) getDCAddresses(logger loggerInterface, ctx context.Context) (dcAddresses, error) {
|
||||
addrs := dcAddresses{
|
||||
v4: map[int][]tgAddr{},
|
||||
v6: map[int][]tgAddr{},
|
||||
}
|
||||
|
||||
err := r.Client.Run(ctx, func(_ context.Context) error {
|
||||
for _, opt := range r.Client.Config().DCOptions {
|
||||
addr := net.JoinHostPort(opt.IPAddress, strconv.Itoa(opt.Port))
|
||||
|
||||
if opt.Ipv6 {
|
||||
addrs.v6[opt.ID] = append(addrs.v6[opt.ID], tgAddr{
|
||||
network: "tcp6",
|
||||
address: addr,
|
||||
})
|
||||
} else {
|
||||
addrs.v4[opt.ID] = append(addrs.v4[opt.ID], tgAddr{
|
||||
network: "tcp4",
|
||||
address: addr,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
|
||||
return addrs, err
|
||||
}
|
||||
@@ -1,139 +0,0 @@
|
||||
package telegram
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/9seconds/mtg/v2/essentials"
|
||||
"github.com/gotd/td/telegram"
|
||||
)
|
||||
|
||||
type Telegram struct {
|
||||
ctx context.Context
|
||||
ctxCancel context.CancelFunc
|
||||
lock sync.RWMutex
|
||||
|
||||
dialer Dialer
|
||||
preferIP preferIP
|
||||
addresses dcAddresses
|
||||
rpc rpcClient
|
||||
}
|
||||
|
||||
func (t *Telegram) Dial(ctx context.Context, dc int) (essentials.Conn, error) {
|
||||
var addresses []tgAddr
|
||||
|
||||
t.lock.RLock()
|
||||
switch t.preferIP {
|
||||
case preferIPOnlyIPv4:
|
||||
addresses = t.addresses.getV4(dc)
|
||||
case preferIPOnlyIPv6:
|
||||
addresses = t.addresses.getV6(dc)
|
||||
case preferIPPreferIPv4:
|
||||
addresses = append(t.addresses.getV4(dc), t.addresses.getV6(dc)...)
|
||||
case preferIPPreferIPv6:
|
||||
addresses = append(t.addresses.getV6(dc), t.addresses.getV4(dc)...)
|
||||
}
|
||||
t.lock.RUnlock()
|
||||
|
||||
var conn essentials.Conn
|
||||
|
||||
err := errNoAddresses
|
||||
|
||||
for _, v := range addresses {
|
||||
conn, err = t.dialer.DialContext(ctx, v.network, v.address)
|
||||
if err == nil {
|
||||
return conn, nil
|
||||
}
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("cannot dial to %d dc: %w", dc, err)
|
||||
}
|
||||
|
||||
func (t *Telegram) IsKnownDC(dc int) bool {
|
||||
return t.addresses.isValidDC(dc)
|
||||
}
|
||||
|
||||
func (t *Telegram) GetFallbackDC() int {
|
||||
return defaultDC
|
||||
}
|
||||
|
||||
func (t *Telegram) Shutdown() {
|
||||
t.ctxCancel()
|
||||
}
|
||||
|
||||
func (t *Telegram) Run(logger loggerInterface, updateEach time.Duration) {
|
||||
if updateEach == 0 {
|
||||
updateEach = defaultUpdateDCAddressesEach
|
||||
}
|
||||
|
||||
t.update(logger)
|
||||
|
||||
ticker := time.NewTicker(updateEach)
|
||||
defer func() {
|
||||
ticker.Stop()
|
||||
|
||||
select {
|
||||
case <-ticker.C:
|
||||
default:
|
||||
}
|
||||
}()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-t.ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
t.update(logger)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Telegram) update(logger loggerInterface) {
|
||||
otherAddresses, err := t.rpc.getDCAddresses(logger, t.ctx)
|
||||
if err != nil {
|
||||
logger.WarningError("Cannot update DC list", err)
|
||||
return
|
||||
}
|
||||
|
||||
t.lock.Lock()
|
||||
t.addresses = otherAddresses
|
||||
t.lock.Unlock()
|
||||
|
||||
logger.Info(fmt.Sprintf("DC are updated: %v", t.addresses))
|
||||
}
|
||||
|
||||
func New(dialer Dialer, ipPreference string) (*Telegram, error) {
|
||||
var pref preferIP
|
||||
|
||||
switch strings.ToLower(ipPreference) {
|
||||
case "prefer-ipv4":
|
||||
pref = preferIPPreferIPv4
|
||||
case "prefer-ipv6":
|
||||
pref = preferIPPreferIPv6
|
||||
case "only-ipv4":
|
||||
pref = preferIPOnlyIPv4
|
||||
case "only-ipv6":
|
||||
pref = preferIPOnlyIPv6
|
||||
default:
|
||||
return nil, fmt.Errorf("unknown ip preference %s", ipPreference)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
return &Telegram{
|
||||
ctx: ctx,
|
||||
ctxCancel: cancel,
|
||||
dialer: dialer,
|
||||
preferIP: pref,
|
||||
addresses: dcAddresses{
|
||||
v4: defaultV4Addresses,
|
||||
v6: defaultV6Addresses,
|
||||
},
|
||||
rpc: rpcClient{
|
||||
Client: telegram.NewClient(defaultAppID, defaultAppHash, telegram.Options{}),
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
@@ -1,159 +0,0 @@
|
||||
package telegram
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net"
|
||||
"strconv"
|
||||
"testing"
|
||||
|
||||
"github.com/9seconds/mtg/v2/internal/testlib"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
type TelegramTestSuite struct {
|
||||
suite.Suite
|
||||
|
||||
dialerMock *testlib.MtglibNetworkMock
|
||||
t *Telegram
|
||||
}
|
||||
|
||||
func (suite *TelegramTestSuite) SetupTest() {
|
||||
suite.dialerMock = &testlib.MtglibNetworkMock{}
|
||||
suite.t, _ = New(suite.dialerMock, "prefer-ipv4", false)
|
||||
}
|
||||
|
||||
func (suite *TelegramTestSuite) TearDownTest() {
|
||||
suite.dialerMock.AssertExpectations(suite.T())
|
||||
}
|
||||
|
||||
func (suite *TelegramTestSuite) TestUnknownDC() {
|
||||
testData := []int{
|
||||
-1,
|
||||
0,
|
||||
6,
|
||||
100,
|
||||
}
|
||||
|
||||
for _, v := range testData {
|
||||
value := v
|
||||
|
||||
suite.T().Run(strconv.Itoa(value), func(t *testing.T) {
|
||||
_, err := suite.t.Dial(context.Background(), value)
|
||||
assert.Error(t, err)
|
||||
assert.False(t, suite.t.IsKnownDC(value))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (suite *TelegramTestSuite) TestDialToCorrectIPs() {
|
||||
testData := map[int][]tgAddr{}
|
||||
|
||||
for i := 1; i <= 5; i++ {
|
||||
testData[i] = []tgAddr{}
|
||||
testData[i] = append(testData[i], productionV4Addresses[i-1]...)
|
||||
testData[i] = append(testData[i], productionV6Addresses[i-1]...)
|
||||
}
|
||||
|
||||
for i, v := range testData {
|
||||
idx := i
|
||||
addresses := v
|
||||
|
||||
suite.T().Run(strconv.Itoa(idx), func(t *testing.T) {
|
||||
for _, addr := range addresses {
|
||||
suite.dialerMock.
|
||||
On("DialContext", mock.Anything, addr.network, addr.address).
|
||||
Once().
|
||||
Return((*net.TCPConn)(nil), io.EOF)
|
||||
}
|
||||
|
||||
_, err := suite.t.Dial(context.Background(), idx)
|
||||
assert.True(t, errors.Is(err, io.EOF))
|
||||
assert.True(t, suite.t.IsKnownDC(idx))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (suite *TelegramTestSuite) TestDialPreferIPRange() {
|
||||
testData := map[string][]tgAddr{
|
||||
"prefer-ipv4": {testV4Addresses[0][0], testV6Addresses[0][0]},
|
||||
"prefer-ipv6": {testV6Addresses[0][0], testV4Addresses[0][0]},
|
||||
"only-ipv4": {testV4Addresses[0][0]},
|
||||
"only-ipv6": {testV6Addresses[0][0]},
|
||||
}
|
||||
|
||||
for k, v := range testData {
|
||||
name := k
|
||||
addresses := v
|
||||
|
||||
suite.T().Run(name, func(t *testing.T) {
|
||||
for _, addr := range addresses {
|
||||
suite.dialerMock.
|
||||
On("DialContext", mock.Anything, addr.network, addr.address).
|
||||
Once().
|
||||
Return((*net.TCPConn)(nil), io.EOF)
|
||||
}
|
||||
|
||||
tg, _ := New(suite.dialerMock, name, true)
|
||||
_, err := tg.Dial(context.Background(), 1)
|
||||
|
||||
assert.True(t, errors.Is(err, io.EOF))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (suite *TelegramTestSuite) TestDialPreferIPPriority() {
|
||||
testData := map[string]tgAddr{
|
||||
"prefer-ipv4": productionV4Addresses[0][0],
|
||||
"prefer-ipv6": productionV6Addresses[0][0],
|
||||
}
|
||||
|
||||
for k, v := range testData {
|
||||
name := k
|
||||
addr := v
|
||||
|
||||
suite.T().Run(name, func(t *testing.T) {
|
||||
conn := &net.TCPConn{}
|
||||
|
||||
suite.dialerMock.
|
||||
On("DialContext", mock.Anything, addr.network, addr.address).
|
||||
Once().
|
||||
Return(conn, nil)
|
||||
|
||||
tg, _ := New(suite.dialerMock, name, false)
|
||||
|
||||
res, err := tg.Dial(context.Background(), 1)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, conn, res)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (suite *TelegramTestSuite) TestUnknownPreferIP() {
|
||||
_, err := New(suite.dialerMock, "xxx", false)
|
||||
suite.Error(err)
|
||||
}
|
||||
|
||||
func (suite *TelegramTestSuite) TestFallbackDC() {
|
||||
dcs := make([]int, 10)
|
||||
|
||||
for i := 0; i < len(dcs); i++ {
|
||||
dcs[i] = suite.t.GetFallbackDC()
|
||||
}
|
||||
|
||||
for _, v := range dcs {
|
||||
value := v
|
||||
|
||||
suite.T().Run(strconv.Itoa(value), func(t *testing.T) {
|
||||
assert.True(t, suite.t.IsKnownDC(value))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTelegram(t *testing.T) {
|
||||
t.Parallel()
|
||||
suite.Run(t, &TelegramTestSuite{})
|
||||
}
|
||||
+26
-13
@@ -10,11 +10,11 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/9seconds/mtg/v2/essentials"
|
||||
"github.com/9seconds/mtg/v2/mtglib/internal/dc"
|
||||
"github.com/9seconds/mtg/v2/mtglib/internal/faketls"
|
||||
"github.com/9seconds/mtg/v2/mtglib/internal/faketls/record"
|
||||
"github.com/9seconds/mtg/v2/mtglib/internal/obfuscated2"
|
||||
"github.com/9seconds/mtg/v2/mtglib/internal/relay"
|
||||
"github.com/9seconds/mtg/v2/mtglib/internal/telegram"
|
||||
"github.com/panjf2000/ants/v2"
|
||||
)
|
||||
|
||||
@@ -28,7 +28,7 @@ type Proxy struct {
|
||||
tolerateTimeSkewness time.Duration
|
||||
domainFrontingPort int
|
||||
workerPool *ants.PoolWithFunc
|
||||
telegram *telegram.Telegram
|
||||
telegram *dc.Telegram
|
||||
|
||||
secret Secret
|
||||
network Network
|
||||
@@ -144,7 +144,6 @@ func (p *Proxy) Shutdown() {
|
||||
p.ctxCancel()
|
||||
p.streamWaitGroup.Wait()
|
||||
p.workerPool.Release()
|
||||
p.telegram.Shutdown()
|
||||
|
||||
p.allowlist.Shutdown()
|
||||
p.blocklist.Shutdown()
|
||||
@@ -220,18 +219,26 @@ func (p *Proxy) doObfuscated2Handshake(ctx *streamContext) error {
|
||||
}
|
||||
|
||||
func (p *Proxy) doTelegramCall(ctx *streamContext) error {
|
||||
dc := ctx.dc
|
||||
|
||||
if p.allowFallbackOnUnknownDC && !p.telegram.IsKnownDC(dc) {
|
||||
dc = p.telegram.GetFallbackDC()
|
||||
ctx.logger = ctx.logger.BindInt("fallback_dc", dc)
|
||||
dcid := ctx.dc
|
||||
|
||||
addresses := p.telegram.GetAddresses(dcid)
|
||||
if len(addresses) == 0 && p.allowFallbackOnUnknownDC {
|
||||
ctx.logger = ctx.logger.BindInt("fallback_dc", dc.DefaultDC)
|
||||
ctx.logger.Warning("unknown DC, fallbacks")
|
||||
addresses = p.telegram.GetAddresses(dc.DefaultDC)
|
||||
}
|
||||
|
||||
conn, err := p.telegram.Dial(ctx, dc)
|
||||
var conn essentials.Conn
|
||||
var err error
|
||||
|
||||
for _, addr := range addresses {
|
||||
conn, err = p.network.Dial(addr.Network, addr.Address)
|
||||
if err == nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("cannot dial to Telegram: %w", err)
|
||||
return fmt.Errorf("no addresses to call: %w", err)
|
||||
}
|
||||
|
||||
encryptor, decryptor, err := obfuscated2.ServerHandshake(conn)
|
||||
@@ -293,9 +300,15 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) {
|
||||
return nil, fmt.Errorf("invalid settings: %w", err)
|
||||
}
|
||||
|
||||
tg, err := telegram.New(opts.Network, opts.getPreferIP())
|
||||
logger := opts.getLogger("proxy")
|
||||
|
||||
tg, err := dc.New(
|
||||
logger.Named("telegram"),
|
||||
opts.getPreferIP(),
|
||||
map[int][]string{},
|
||||
) // TODO: propagate value
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot build telegram dialer: %w", err)
|
||||
return nil, fmt.Errorf("cannot build telegram dc fetcher: %w", err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
@@ -315,7 +328,7 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) {
|
||||
telegram: tg,
|
||||
}
|
||||
|
||||
go tg.Run(proxy.logger.Named("telegram"), 0)
|
||||
go tg.Run(ctx, 0) // TODO: propagate value
|
||||
|
||||
pool, err := ants.NewPoolWithFunc(opts.getConcurrency(),
|
||||
func(arg interface{}) {
|
||||
|
||||
Reference in New Issue
Block a user