Support default dc idx from configs

This commit is contained in:
9seconds
2018-09-21 09:47:37 +03:00
parent 8268620615
commit 1e9bbfae55
24 changed files with 187 additions and 123 deletions
+9 -2
View File
@@ -12,6 +12,11 @@ import (
"github.com/9seconds/mtg/wrappers"
)
const (
directV4DefaultIdx = 1
directV6DefaultIdx = 1
)
var (
directV4Addresses = map[int16][]string{
0: {"149.154.175.50:443"},
@@ -65,8 +70,10 @@ func NewDirectTelegram(conf *config.Config) Telegram {
Dialer: net.Dialer{Timeout: telegramDialTimeout},
conf: conf,
},
v4Addresses: directV4Addresses,
v6Addresses: directV6Addresses,
v4DefaultIdx: directV4DefaultIdx,
v6DefaultIdx: directV6DefaultIdx,
v4Addresses: directV4Addresses,
v6Addresses: directV6Addresses,
},
}
}
+66 -29
View File
@@ -65,18 +65,20 @@ func (t *middleTelegramCaller) update() error {
return errors.Annotate(err, "Cannot get proxy secret")
}
v4Addresses, err := t.getTelegramAddresses(tgAddrProxyV4)
v4Addresses, v4DefaultIdx, err := t.getTelegramAddresses(tgAddrProxyV4)
if err != nil {
return errors.Annotate(err, "Cannot get ipv4 addresses")
}
v6Addresses, err := t.getTelegramAddresses(tgAddrProxyV6)
v6Addresses, v6DefaultIdx, err := t.getTelegramAddresses(tgAddrProxyV6)
if err != nil {
return errors.Annotate(err, "Cannot get ipv6 addresses")
}
t.dialerMutex.Lock()
t.proxySecret = secret
t.v4DefaultIdx = v4DefaultIdx
t.v6DefaultIdx = v6DefaultIdx
t.v4Addresses = v4Addresses
t.v6Addresses = v6Addresses
t.dialerMutex.Unlock()
@@ -101,52 +103,87 @@ func (t *middleTelegramCaller) getTelegramProxySecret() ([]byte, error) {
return secret, nil
}
func (t *middleTelegramCaller) getTelegramAddresses(url string) (map[int16][]string, error) {
func (t *middleTelegramCaller) getTelegramAddresses(url string) (map[int16][]string, int16, error) { // nolint: gocyclo
resp, err := t.call(url)
if err != nil {
return nil, errors.Annotate(err, "Cannot access telegram server")
return nil, 0, errors.Annotate(err, "Cannot access telegram server")
}
defer resp.Body.Close() // nolint: errcheck
scanner := bufio.NewScanner(resp.Body)
data := map[int16][]string{}
var defaultIdx int16 = 1
for scanner.Scan() {
text := strings.TrimSpace(scanner.Text())
if strings.HasPrefix(text, "#") {
switch {
case strings.HasPrefix(text, "#"):
continue
}
chunks := middleTelegramProxyConfigSplitter.Split(text, 3)
if len(chunks) != 3 || chunks[0] != "proxy_for" {
return nil, errors.Errorf("Incorrect config '%s'", text)
}
dcIdx64, err2 := strconv.ParseInt(chunks[1], 10, 16)
if err2 != nil {
return nil, errors.Errorf("Incorrect config '%s'", text)
}
dcIdx := int16(dcIdx64)
addr := strings.TrimRight(chunks[2], ";")
if _, _, err2 = net.SplitHostPort(addr); err != nil {
return nil, errors.Annotatef(err2, "Incorrect config '%s'", text)
}
if addresses, ok := data[dcIdx]; ok {
data[dcIdx] = append(addresses, addr)
} else {
data[dcIdx] = []string{addr}
case strings.HasPrefix(text, "proxy_for"):
addr, idx, err2 := t.parseProxyFor(text)
if err2 != nil {
return nil, 0, errors.Annotate(err2, "Cannot parse 'proxy_for' section")
}
if addresses, ok := data[idx]; ok {
data[idx] = append(addresses, addr)
} else {
data[idx] = []string{addr}
}
case strings.HasPrefix(text, "default"):
idx, err2 := t.parseDefault(text)
if err2 != nil {
return nil, 0, errors.Annotate(err2, "Cannot parse 'default' section")
}
defaultIdx = idx
default:
return nil, 0, errors.Errorf("Unknown config string '%s'", text)
}
}
err = scanner.Err()
if err != nil {
return nil, errors.Annotate(err, "Cannot read response from the telegram")
return nil, 0, errors.Annotate(err, "Cannot read response from the telegram")
}
return data, nil
return data, defaultIdx, nil
}
func (t *middleTelegramCaller) parseProxyFor(text string) (string, int16, error) {
chunks := middleTelegramProxyConfigSplitter.Split(text, 3)
if len(chunks) != 3 || chunks[0] != "proxy_for" {
return "", 0, errors.Errorf("Incorrect config '%s'", text)
}
dcIdx, err := strconv.ParseInt(chunks[1], 10, 16)
if err != nil {
return "", 0, errors.Annotatef(err, "Incorrect config '%s'", text)
}
addr := strings.TrimRight(chunks[2], ";")
if _, _, err = net.SplitHostPort(addr); err != nil {
return "", 0, errors.Annotatef(err, "Incorrect config '%s'", text)
}
return addr, int16(dcIdx), nil
}
func (t *middleTelegramCaller) parseDefault(text string) (int16, error) {
chunks := middleTelegramProxyConfigSplitter.Split(text, 2)
if len(chunks) != 2 || chunks[0] != "default" {
return 0, errors.Errorf("Incorrect config '%s'", text)
}
dcIdxString := strings.TrimRight(chunks[1], ";")
dcIdx, err := strconv.ParseInt(dcIdxString, 10, 16)
if err != nil {
return 0, errors.Annotatef(err, "Incorrect config '%s'", text)
}
return int16(dcIdx), nil
}
func (t *middleTelegramCaller) call(url string) (*http.Response, error) {
req, _ := http.NewRequest("GET", url, nil)
req, _ := http.NewRequest("GET", url, nil) // nolint: gosec
req.Header.Set("Accept", "text/plain")
req.Header.Set("User-Agent", tgUserAgent)
+26 -6
View File
@@ -19,8 +19,10 @@ type Telegram interface {
type baseTelegram struct {
dialer tgDialer
v4Addresses map[int16][]string
v6Addresses map[int16][]string
v4DefaultIdx int16
v6DefaultIdx int16
v4Addresses map[int16][]string
v6Addresses map[int16][]string
}
func (b *baseTelegram) dial(ctx context.Context, cancel context.CancelFunc, dcIdx int16, connID string,
@@ -28,13 +30,13 @@ func (b *baseTelegram) dial(ctx context.Context, cancel context.CancelFunc, dcId
addrs := make([]string, 2)
if proto&mtproto.ConnectionProtocolIPv6 != 0 {
if addr, ok := b.v6Addresses[dcIdx]; ok && len(addr) > 0 {
addrs = append(addrs, addr[rand.Intn(len(addr))])
if addr := b.chooseAddress(b.v6Addresses, dcIdx, b.v6DefaultIdx); addr != "" {
addrs = append(addrs, addr)
}
}
if proto&mtproto.ConnectionProtocolIPv4 != 0 {
if addr, ok := b.v4Addresses[dcIdx]; ok && len(addr) > 0 {
addrs = append(addrs, addr[rand.Intn(len(addr))])
if addr := b.chooseAddress(b.v4Addresses, dcIdx, b.v4DefaultIdx); addr != "" {
addrs = append(addrs, addr)
}
}
@@ -46,3 +48,21 @@ func (b *baseTelegram) dial(ctx context.Context, cancel context.CancelFunc, dcId
return nil, errors.New("Cannot connect to Telegram")
}
func (b *baseTelegram) chooseAddress(addresses map[int16][]string, idx, defaultIdx int16) string {
if addr, ok := addresses[idx]; ok {
return b.chooseRandomAddress(addr)
} else if addr, ok := addresses[defaultIdx]; ok {
return b.chooseRandomAddress(addr)
}
return ""
}
func (b *baseTelegram) chooseRandomAddress(addresses []string) string {
if len(addresses) > 0 {
return addresses[rand.Intn(len(addresses))]
}
return ""
}