ScuroNeko/mtg
1
0
Fork 0
mirror of https://github.com/ScuroNeko/mtg.git synced 2026-09-01 03:54:02 +03:00
Code Issues Packages Projects Releases Wiki Activity

Create internal DC package

Browse Source
This commit is contained in:
9seconds
2026-02-16 14:15:38 +01:00
parent 852ca713c8
commit 0a5a45b32d
13 changed files with 440 additions and 452 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mtglib/internal/dc/addr.go
+10
View File
@@ -0,0 +1,10 @@
package dc
type Addr struct {
Network string
Address string
}
func (d Addr) String() string {
return d.Address
}
mtglib/internal/dc/addr_set.go
+33
View File
@@ -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
}
mtglib/internal/dc/addr_test.go
+16
View File
@@ -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())
}
mtglib/internal/dc/init.go
+79
View File
@@ -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"},
},
},
}
)
mtglib/internal/dc/telegram.go
+160
View File
@@ -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
}
mtglib/internal/dc/view.go
+24
View File
@@ -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
}
mtglib/internal/dc/view_test.go
+92
View File
@@ -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{})
}