From 836a481026fec0c8de8a217448ae8a78a6fb0931 Mon Sep 17 00:00:00 2001 From: 9seconds Date: Mon, 16 Feb 2026 14:48:30 +0100 Subject: [PATCH] Propagate DCOverrides --- go.mod | 1 + go.sum | 2 + internal/cli/run_proxy.go | 9 ++++ internal/config/config.go | 4 ++ internal/config/parse.go | 6 ++- internal/config/type_dc.go | 41 ++++++++++++++ internal/config/type_dc_test.go | 95 +++++++++++++++++++++++++++++++++ mtglib/proxy.go | 6 +-- mtglib/proxy_opts.go | 6 +++ 9 files changed, 164 insertions(+), 6 deletions(-) create mode 100644 internal/config/type_dc.go create mode 100644 internal/config/type_dc_test.go diff --git a/go.mod b/go.mod index 55929b9..cc30b92 100644 --- a/go.mod +++ b/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 ) diff --git a/go.sum b/go.sum index cf02420..7477545 100644 --- a/go.sum +++ b/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= diff --git a/internal/cli/run_proxy.go b/internal/cli/run_proxy.go index c853212..5357309 100644 --- a/internal/cli/run_proxy.go +++ b/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) diff --git a/internal/config/config.go b/internal/config/config.go index 9b69a94..c330a25 100644 --- a/internal/config/config.go +++ b/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 { diff --git a/internal/config/parse.go b/internal/config/parse.go index 591e6bf..8d056f3 100644 --- a/internal/config/parse.go +++ b/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) { diff --git a/internal/config/type_dc.go b/internal/config/type_dc.go new file mode 100644 index 0000000..d3b1257 --- /dev/null +++ b/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 +} diff --git a/internal/config/type_dc_test.go b/internal/config/type_dc_test.go new file mode 100644 index 0000000..b51229c --- /dev/null +++ b/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{}) +} diff --git a/mtglib/proxy.go b/mtglib/proxy.go index ba80d99..c037a39 100644 --- a/mtglib/proxy.go +++ b/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) } diff --git a/mtglib/proxy_opts.go b/mtglib/proxy_opts.go index a966d48..53f3434 100644 --- a/mtglib/proxy_opts.go +++ b/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 {