diff --git a/telegram/base.go b/telegram/base.go index 2f0a693..bb27ff0 100644 --- a/telegram/base.go +++ b/telegram/base.go @@ -7,6 +7,7 @@ import ( "go.uber.org/zap" + "github.com/9seconds/mtg/config" "github.com/9seconds/mtg/conntypes" "github.com/9seconds/mtg/utils" "github.com/9seconds/mtg/wrappers/stream" @@ -29,17 +30,7 @@ func (b *baseTelegram) Secret() []byte { func (b *baseTelegram) dial(dc conntypes.DC, protocol conntypes.ConnectionProtocol) (conntypes.StreamReadWriteCloser, error) { - addresses := make([]string, 0, 2) - - if protocol&conntypes.ConnectionProtocolIPv6 != 0 { - addresses = append(addresses, b.chooseAddress(b.v6Addresses, dc, b.v6DefaultDC)) - } - - if protocol&conntypes.ConnectionProtocolIPv4 != 0 { - addresses = append(addresses, b.chooseAddress(b.v4Addresses, dc, b.v4DefaultDC)) - } - - for _, addr := range addresses { + for _, addr := range b.getAddresses(dc, protocol) { conn, err := b.dialer.Dial("tcp", addr) if err != nil { b.logger.Infow("Cannot dial to Telegram", "address", addr, "error", err) @@ -57,6 +48,30 @@ func (b *baseTelegram) dial(dc conntypes.DC, return nil, errors.New("cannot dial to the chosen DC") } +func (b *baseTelegram) getAddresses(dc conntypes.DC, protocol conntypes.ConnectionProtocol) []string { + addresses := make([]string, 0, 2) + protos := []conntypes.ConnectionProtocol{ + conntypes.ConnectionProtocolIPv6, + conntypes.ConnectionProtocolIPv4, + } + + if config.C.PreferIP == config.PreferIPv4 { + protos[0], protos[1] = protos[1], protos[0] + } + + for _, proto := range protos { + switch { + case proto&protocol == 0: + case proto&conntypes.ConnectionProtocolIPv6 != 0: + addresses = append(addresses, b.chooseAddress(b.v6Addresses, dc, b.v6DefaultDC)) + case proto&conntypes.ConnectionProtocolIPv4 != 0: + addresses = append(addresses, b.chooseAddress(b.v4Addresses, dc, b.v4DefaultDC)) + } + } + + return addresses +} + func (b *baseTelegram) chooseAddress(addresses map[conntypes.DC][]string, dc, defaultDC conntypes.DC) string { addrs, ok := addresses[dc]