Skip to content

Commit

Permalink
improve proxysupport, get rid of countrychecker
Browse files Browse the repository at this point in the history
  • Loading branch information
arriven committed Jun 4, 2023
1 parent acc9c83 commit d8bb046
Show file tree
Hide file tree
Showing 6 changed files with 82 additions and 161 deletions.
4 changes: 1 addition & 3 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,6 @@ func main() {
runnerConfigOptions := job.NewConfigOptionsWithFlags()
jobsGlobalConfig := job.NewGlobalConfigWithFlags()
otaConfig := ota.NewConfigWithFlags()
countryCheckerConfig := utils.NewCountryCheckerConfigWithFlags()
updaterMode, destinationPath := config.NewUpdaterOptionsWithFlags()
prometheusOn, prometheusListenAddress := metrics.NewOptionsWithFlags()
pprof := flag.String("pprof", utils.GetEnvStringDefault("GO_PPROF_ENDPOINT", ""), "enable pprof")
Expand Down Expand Up @@ -100,8 +99,7 @@ func main() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()

metrics.InitOrFail(ctx, logger, *prometheusOn, *prometheusListenAddress, jobsGlobalConfig.ClientID,
utils.CheckCountryOrFail(ctx, logger, countryCheckerConfig, jobsGlobalConfig.GetProxyParams(logger, nil)))
metrics.InitOrFail(ctx, logger, *prometheusOn, *prometheusListenAddress, jobsGlobalConfig.ClientID, "")
job.NewRunner(runnerConfigOptions, jobsGlobalConfig, newReporter(*logFormat, *lessStats, logger)).Run(ctx, logger)
}

Expand Down
72 changes: 63 additions & 9 deletions src/job/base.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,12 @@ package job
import (
"context"
"flag"
"io"
"math/rand"
"net/http"
"net/url"
"os"
"path/filepath"
"time"

"github.com/google/uuid"
Expand All @@ -43,9 +48,11 @@ type GlobalConfig struct {
ClientID string
UserID string

ProxyURLs string
LocalAddr string
Interface string
proxyURLs string
proxylist string
defaultProxyProto string
localAddr string
iface string
SkipEncrypted bool
EnablePrimitiveJobs bool
ScaleFactor float64
Expand All @@ -62,11 +69,13 @@ func NewGlobalConfigWithFlags() *GlobalConfig {

flag.StringVar(&res.UserID, "user-id", utils.GetEnvStringDefault("USER_ID", ""),
"user id for optional metrics")
flag.StringVar(&res.ProxyURLs, "proxy", utils.GetEnvStringDefault("SYSTEM_PROXY", ""),
flag.StringVar(&res.proxyURLs, "proxy", utils.GetEnvStringDefault("SYSTEM_PROXY", ""),
"system proxy to set by default (can be a comma-separated list or a template)")
flag.StringVar(&res.LocalAddr, "local-address", utils.GetEnvStringDefault("LOCAL_ADDRESS", ""),
flag.StringVar(&res.proxylist, "proxylist", "", "file or url to read a list of proxies from")
flag.StringVar(&res.defaultProxyProto, "default-proxy-proto", "socks5", "protocol to fallback to if proxy contains only address")
flag.StringVar(&res.localAddr, "local-address", utils.GetEnvStringDefault("LOCAL_ADDRESS", ""),
"specify ip address of local interface to use")
flag.StringVar(&res.Interface, "interface", utils.GetEnvStringDefault("NETWORK_INTERFACE", ""),
flag.StringVar(&res.iface, "interface", utils.GetEnvStringDefault("NETWORK_INTERFACE", ""),
"specify which interface to bind to for attacks (ignored on windows)")
flag.BoolVar(&res.SkipEncrypted, "skip-encrypted", utils.GetEnvBoolDefault("SKIP_ENCRYPTED", false),
"set to true if you want to only run plaintext jobs from the config for security considerations")
Expand All @@ -91,16 +100,61 @@ func NewGlobalConfigWithFlags() *GlobalConfig {

func (g GlobalConfig) GetProxyParams(logger *zap.Logger, data any) utils.ProxyParams {
return utils.ProxyParams{
URLs: templates.ParseAndExecute(logger, g.ProxyURLs, data),
LocalAddr: templates.ParseAndExecute(logger, g.LocalAddr, data),
Interface: templates.ParseAndExecute(logger, g.Interface, data),
URLs: templates.ParseAndExecute(logger, g.proxyURLs, data),
DefaultProto: g.defaultProxyProto,
LocalAddr: templates.ParseAndExecute(logger, g.localAddr, data),
Interface: templates.ParseAndExecute(logger, g.iface, data),
}
}

func (g *GlobalConfig) initProxylist(ctx context.Context) error {
if g.proxyURLs != "" || g.proxylist == "" {
return nil
}
proxylist, err := readProxylist(ctx, g.proxylist)
if err != nil {
return err
}
g.proxyURLs = string(proxylist)
return nil
}

func readProxylist(ctx context.Context, path string) ([]byte, error) {
proxylistURL, err := url.ParseRequestURI(path)
// absolute paths can be interpreted as a URL with no schema, need to check for that explicitly
if err != nil || filepath.IsAbs(path) {
res, err := os.ReadFile(path)
if err != nil {
return nil, err
}

return res, nil
}

const requestTimeout = 20 * time.Second

ctx, cancel := context.WithTimeout(ctx, requestTimeout)
defer cancel()

req, err := http.NewRequestWithContext(ctx, http.MethodGet, proxylistURL.String(), nil)
if err != nil {
return nil, err
}

resp, err := http.DefaultClient.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()

return io.ReadAll(resp.Body)
}

// Job comment for linter
type Job = func(ctx context.Context, args config.Args, globalConfig *GlobalConfig, a *metrics.Accumulator, logger *zap.Logger) (data any, err error)

// Get job by type name
//
//nolint:cyclop // The string map alternative is orders of magnitude slower
func Get(t string) Job {
switch t {
Expand Down
1 change: 1 addition & 0 deletions src/job/config/defaultconfig.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
)

// DefaultConfig is the config embedded into the app that it will use if not able to fetch any other config
//
//nolint:lll // Makes no sense splitting this into multiple lines
var DefaultConfig = ``

Expand Down
4 changes: 4 additions & 0 deletions src/job/runner.go
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,10 @@ func NewRunner(cfgOptions *ConfigOptions, globalJobsCfg *GlobalConfig, reporter

// Run the runner and block until Stop() is called
func (r *Runner) Run(ctx context.Context, logger *zap.Logger) {
if err := r.globalJobsCfg.initProxylist(ctx); err != nil {
logger.Warn("failed to init proxylist", zap.Error(err))
}

ctx = context.WithValue(ctx, templates.ContextKey("goos"), runtime.GOOS)
ctx = context.WithValue(ctx, templates.ContextKey("goarch"), runtime.GOARCH)
ctx = context.WithValue(ctx, templates.ContextKey("version"), ota.Version)
Expand Down
142 changes: 0 additions & 142 deletions src/utils/countrychecker.go

This file was deleted.

20 changes: 13 additions & 7 deletions src/utils/proxy.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,11 @@ import (
type ProxyFunc func(network, addr string) (net.Conn, error)

type ProxyParams struct {
URLs string
LocalAddr string
Interface string
Timeout time.Duration
URLs string
DefaultProto string
LocalAddr string
Interface string
Timeout time.Duration
}

// this won't work for udp payloads but if people use proxies they might not want to have their ip exposed
Expand All @@ -31,13 +32,18 @@ func GetProxyFunc(ctx context.Context, params ProxyParams, protocol string) Prox
return proxy.FromEnvironmentUsing(direct).Dial
}

proxies := strings.Split(params.URLs, ",")
proxies := strings.Fields(strings.ReplaceAll(params.URLs, ",", " "))

// We need to dial new proxy on each call
return func(network, addr string) (net.Conn, error) {
u, err := url.Parse(proxies[rand.Intn(len(proxies))]) //nolint:gosec // Cryptographically secure random not required
selected := proxies[rand.Intn(len(proxies))] //nolint:gosec // Cryptographically secure random not required
u, err := url.Parse(selected)
if err != nil {
return nil, fmt.Errorf("error building proxy %v: %w", u.String(), err)
selected = params.DefaultProto + "://" + selected
u, err = url.Parse(selected)
if err != nil {
return nil, fmt.Errorf("error building proxy %v: %w", selected, err)
}
}

switch u.Scheme {
Expand Down

0 comments on commit d8bb046

Please sign in to comment.