Propagate DCOverrides

go.mod

@@ -29,6 +29,7 @@ require (
 
 require (
 	github.com/gotd/td v0.139.0
+	github.com/pelletier/go-toml/v2 v2.2.4
 	github.com/txthinking/socks5 v0.0.0-20251011041537-5c31f201a10e
 	github.com/yl2chen/cidranger v1.0.2
 )

go.sum

@@ -88,6 +88,8 @@ github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaR
 github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/pelletier/go-toml v1.9.5 h1:4yBQzkHv+7BHq2PQUZF3Mx0IYxG7LsP222s7Agd3ve8=
 github.com/pelletier/go-toml v1.9.5/go.mod h1:u1nR/EPcESfeI/szUZKdtJ0xRNbUoANCkoOuaOx1Y+c=
+github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4=
+github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=

internal/cli/run_proxy.go

@@ -240,6 +240,14 @@ func runProxy(conf *config.Config, version string) error { //nolint: funlen
 		return fmt.Errorf("cannot build ip allowlist: %w", err)
 	}
 
+	dcOverrides := map[int][]string{}
+	for _, override := range conf.DCOverrides {
+		dcid := override.DC.Get()
+		for _, addr := range override.IPs {
+			dcOverrides[dcid] = append(dcOverrides[dcid], addr.Get(""))
+		}
+	}
+
 	opts := mtglib.ProxyOpts{
 		Logger:          logger,
 		Network:         ntw,
@@ -254,6 +262,7 @@ func runProxy(conf *config.Config, version string) error { //nolint: funlen
 
 		AllowFallbackOnUnknownDC: conf.AllowFallbackOnUnknownDC.Get(false),
 		TolerateTimeSkewness:     conf.TolerateTimeSkewness.Value,
+		DCOverrides:              dcOverrides,
 	}
 
 	proxy, err := mtglib.NewProxy(opts)

internal/config/config.go

@@ -64,6 +64,10 @@ type Config struct {
 			MetricPrefix TypeMetricPrefix `json:"metricPrefix"`
 		} `json:"prometheus"`
 	} `json:"stats"`
+	DCOverrides []struct {
+		DC  TypeDC
+		IPs []TypeHostPort `json:"ips"`
+	} `json:"dc_overrides"`
 }
 
 func (c *Config) Validate() error {

internal/config/parse.go

@@ -5,7 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 
-	"github.com/pelletier/go-toml"
+	"github.com/pelletier/go-toml/v2"
 )
 
 type tomlConfig struct {
@@ -59,6 +59,10 @@ type tomlConfig struct {
 			MetricPrefix string `toml:"metric-prefix" json:"metricPrefix,omitempty"`
 		} `toml:"prometheus" json:"prometheus,omitempty"`
 	} `toml:"stats" json:"stats,omitempty"`
+	DCOverrides []struct {
+		DC  int      `toml:"dc" json:"dc"`
+		IPs []string `toml:"ips" json:"ips"`
+	}
 }
 
 func Parse(rawData []byte) (*Config, error) {

internal/config/type_dc.go

@@ -0,0 +1,41 @@
+package config
+
+import (
+	"fmt"
+	"strconv"
+)
+
+type TypeDC struct {
+	Value int
+}
+
+func (t *TypeDC) Set(value string) error {
+	parsed, err := strconv.ParseInt(value, 10, 16)
+	if err != nil {
+		return fmt.Errorf("cannot parse dc: %w", err)
+	}
+
+	if parsed < 0 {
+		parsed = -parsed
+	}
+
+	t.Value = int(parsed)
+
+	return nil
+}
+
+func (t *TypeDC) UnmarshalText(data []byte) error {
+	return t.Set(string(data))
+}
+
+func (t TypeDC) MarshalText() ([]byte, error) {
+	return []byte(t.String()), nil
+}
+
+func (t TypeDC) String() string {
+	return strconv.Itoa(t.Value)
+}
+
+func (t TypeDC) Get() int {
+	return t.Value
+}

internal/config/type_dc_test.go

@@ -0,0 +1,95 @@
+package config_test
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/9seconds/mtg/v2/internal/config"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/suite"
+)
+
+type typeDCTestStruct struct {
+	Value config.TypeDC `json:"value"`
+}
+
+type TypeDCTestSuite struct {
+	suite.Suite
+}
+
+func (suite *TypeDCTestSuite) TestUnmarshalFail() {
+	testData := []string{
+		"-1s",
+		"1202002020202",
+		"xxx",
+		"-11111111111111",
+		"",
+	}
+
+	for _, v := range testData {
+		data, err := json.Marshal(map[string]string{
+			"value": v,
+		})
+		suite.NoError(err)
+
+		suite.T().Run(v, func(t *testing.T) {
+			assert.Error(t, json.Unmarshal(data, &typeDCTestStruct{}))
+		})
+	}
+}
+
+func (suite *TypeDCTestSuite) TestUnmarshalOk() {
+	testData := map[string]int{
+		"1":   1,
+		"-1":  1,
+		"203": 203,
+	}
+
+	for value, expected := range testData {
+		data, err := json.Marshal(map[string]string{
+			"value": value,
+		})
+		suite.NoError(err)
+
+		suite.T().Run(value, func(t *testing.T) {
+			testStruct := &typeDCTestStruct{}
+
+			assert.NoError(t, json.Unmarshal(data, testStruct))
+			assert.Equal(t, expected, testStruct.Value.Value)
+			assert.Equal(t, expected, testStruct.Value.Get())
+		})
+	}
+}
+
+func (suite *TypeDCTestSuite) TestMarshalOk() {
+	testData := map[string]string{
+		"1":   "1",
+		"203": "203",
+	}
+
+	for k, v := range testData {
+		value := k
+		expected := v
+
+		suite.T().Run(value, func(t *testing.T) {
+			testStruct := &typeDCTestStruct{}
+
+			assert.NoError(t, testStruct.Value.Set(value))
+
+			data, err := json.Marshal(testStruct)
+			assert.NoError(t, err)
+
+			expectedJSON, err := json.Marshal(map[string]string{
+				"value": expected,
+			})
+			assert.NoError(t, err)
+
+			assert.JSONEq(t, string(expectedJSON), string(data))
+		})
+	}
+}
+
+func TestTypeDC(t *testing.T) {
+	t.Parallel()
+	suite.Run(t, &TypeDCTestSuite{})
+}

mtglib/proxy.go

@@ -302,11 +302,7 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) {
 
 	logger := opts.getLogger("proxy")
 
-	tg, err := dc.New(
-		logger.Named("telegram"),
-		opts.getPreferIP(),
-		map[int][]string{},
-	) // TODO: propagate value
+	tg, err := dc.New(logger.Named("telegram"), opts.getPreferIP(), opts.DCOverrides)
 	if err != nil {
 		return nil, fmt.Errorf("cannot build telegram dc fetcher: %w", err)
 	}

mtglib/proxy_opts.go

@@ -113,6 +113,12 @@ type ProxyOpts struct {
 	//
 	// OBSOLETE and DEPRECATED. Ignored.
 	UseTestDCs bool
+
+	// DCOverrides defines a set of IP addresses that should be used
+	// with a higher priority to those that are calculated somehow by mtg.
+	//
+	// This is an optional setting
+	DCOverrides map[int][]string
 }
 
 func (p ProxyOpts) valid() error {