add way to bind to specific interface

arriven · Apr 28, 2022 · b220cf2 · b220cf2
1 parent 8ac3e4e
commit b220cf2
Show file tree

Hide file tree

Showing 6 changed files with 68 additions and 27 deletions.
diff --git a/src/core/http/http.go b/src/core/http/http.go
@@ -101,6 +101,7 @@ type ClientConfig struct {
 	IdleTimeout     *time.Duration
 	MaxIdleConns    *int
 	ProxyURLs       string
+	LocalAddr       string
 }
 
 // NewClient creates a fasthttp client based on the config.
@@ -119,7 +120,9 @@ func NewClient(ctx context.Context, clientConfig ClientConfig, logger *zap.Logge
 		tlsConfig = clientConfig.TLSClientConfig
 	}
 
-	proxyFunc := utils.GetProxyFunc(templates.ParseAndExecute(logger, clientConfig.ProxyURLs, ctx), timeout, true)
+	proxyURLs := templates.ParseAndExecute(logger, clientConfig.ProxyURLs, ctx)
+	localAddr := utils.ResolveAddr("tcp", clientConfig.LocalAddr)
+	proxyFunc := utils.GetProxyFunc(proxyURLs, localAddr, timeout, true)
 
 	if clientConfig.StaticHost != nil {
 		makeHostClient := func(tls bool) *fasthttp.HostClient {

diff --git a/src/core/packetgen/connection.go b/src/core/packetgen/connection.go
@@ -115,6 +115,7 @@ type netConnConfig struct {
 	Address         string
 	Timeout         time.Duration
 	ProxyURLs       string
+	LocalAddr       string
 	Reader          *netReaderConfig
 	TLSClientConfig *tls.Config
 }
@@ -149,7 +150,7 @@ func readStub(ctx context.Context, conn net.Conn, c *netReaderConfig) {
 }
 
 func openNetConn(ctx context.Context, c netConnConfig) (*netConn, error) {
-	conn, err := utils.GetProxyFunc(c.ProxyURLs, c.Timeout, false)(c.Protocol, c.Address)
+	conn, err := utils.GetProxyFunc(c.ProxyURLs, utils.ResolveAddr(c.Protocol, c.LocalAddr), c.Timeout, false)(c.Protocol, c.Address)
 
 	switch {
 	case err != nil:

diff --git a/src/job/base.go b/src/job/base.go
@@ -41,6 +41,7 @@ type GlobalConfig struct {
 	ClientID string
 
 	ProxyURLs           string
+	LocalAddr           string
 	SkipEncrypted       bool
 	EnablePrimitiveJobs bool
 	ScaleFactor         int
@@ -56,6 +57,8 @@ func NewGlobalConfigWithFlags() *GlobalConfig {
 
 	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", ""),
+		"specify ip address of local interface to use")
 	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")
 	flag.BoolVar(&res.EnablePrimitiveJobs, "enable-primitive", utils.GetEnvBoolDefault("ENABLE_PRIMITIVE", true),

diff --git a/src/job/http.go b/src/job/http.go
@@ -210,6 +210,10 @@ func getHTTPJobConfigs(ctx context.Context, args config.Args, global GlobalConfi
 		clientConfig.ProxyURLs = templates.ParseAndExecute(logger, global.ProxyURLs, ctx)
 	}
 
+	if global.LocalAddr != "" {
+		clientConfig.LocalAddr = templates.ParseAndExecute(logger, global.LocalAddr, ctx)
+	}
+
 	requestTpl, err = templates.ParseMapStruct(jobConfig.Request)
 	if err != nil {
 		return nil, nil, nil, fmt.Errorf("error parsing request config: %w", err)

diff --git a/src/job/packetgen.go b/src/job/packetgen.go
@@ -138,30 +138,16 @@ func getNextPacket(ctx context.Context, logger *zap.Logger, packetsChan chan *te
 	}
 }
 
-func parsePacketgenArgs(ctx context.Context, args config.Args, globalConfig *GlobalConfig, a *metrics.Accumulator, logger *zap.Logger) (
-	tpl *packetgenJobConfig, err error,
-) {
-	type packetDescriptor struct {
-		Packet map[string]any
-		Count  int
-	}
-
-	var jobConfig struct {
-		BasicJobConfig
-		StaticPacket bool
-		Packet       map[string]any
-		Packets      []packetDescriptor
-		Connection   packetgen.ConnectionConfig
-	}
-
-	if err = ParseConfig(&jobConfig, args, *globalConfig); err != nil {
-		return nil, fmt.Errorf("error parsing job config: %w", err)
-	}
+type packetDescriptor struct {
+	Packet map[string]any
+	Count  int
+}
 
-	packetTpls := make([]*templates.MapStruct, 0, len(jobConfig.Packets)+1)
+func parsePackets(packets []packetDescriptor, dflt map[string]any) ([]*templates.MapStruct, error) {
+	packetTpls := make([]*templates.MapStruct, 0, len(packets)+1)
 
-	if len(jobConfig.Packets) > 0 {
-		for _, descriptor := range jobConfig.Packets {
+	if len(packets) > 0 {
+		for _, descriptor := range packets {
 			if descriptor.Count < 1 {
 				descriptor.Count = 1
 			}
@@ -176,18 +162,45 @@ func parsePacketgenArgs(ctx context.Context, args config.Args, globalConfig *Glo
 			}
 		}
 	} else {
-		packetTpl, err := templates.ParseMapStruct(jobConfig.Packet)
+		packetTpl, err := templates.ParseMapStruct(dflt)
 		if err != nil {
 			return nil, fmt.Errorf("error parsing packet: %w", err)
 		}
 
 		packetTpls = append(packetTpls, packetTpl)
 	}
 
+	return packetTpls, nil
+}
+
+func parsePacketgenArgs(ctx context.Context, args config.Args, globalConfig *GlobalConfig, a *metrics.Accumulator, logger *zap.Logger) (
+	tpl *packetgenJobConfig, err error,
+) {
+	var jobConfig struct {
+		BasicJobConfig
+		StaticPacket bool
+		Packet       map[string]any
+		Packets      []packetDescriptor
+		Connection   packetgen.ConnectionConfig
+	}
+
+	if err = ParseConfig(&jobConfig, args, *globalConfig); err != nil {
+		return nil, fmt.Errorf("error parsing job config: %w", err)
+	}
+
+	packetTpls, err := parsePackets(jobConfig.Packets, jobConfig.Packet)
+	if err != nil {
+		return nil, err
+	}
+
 	if globalConfig.ProxyURLs != "" && jobConfig.Connection.Args["protocol"] == "tcp" {
 		jobConfig.Connection.Args["proxy_urls"] = templates.ParseAndExecute(logger, globalConfig.ProxyURLs, ctx)
 	}
 
+	if globalConfig.LocalAddr != "" {
+		jobConfig.Connection.Args["local_addr"] = templates.ParseAndExecute(logger, globalConfig.LocalAddr, ctx)
+	}
+
 	return &packetgenJobConfig{
 		BasicJobConfig: jobConfig.BasicJobConfig,
 		StaticPacket:   jobConfig.StaticPacket,

diff --git a/src/utils/proxy.go b/src/utils/proxy.go
@@ -15,8 +15,8 @@ import (
 
 type ProxyFunc func(network, addr string) (net.Conn, error)
 
-func GetProxyFunc(proxyURLs string, timeout time.Duration, httpEnabled bool) ProxyFunc {
-	direct := &net.Dialer{Timeout: timeout}
+func GetProxyFunc(proxyURLs string, localAddr net.Addr, timeout time.Duration, httpEnabled bool) ProxyFunc {
+	direct := &net.Dialer{Timeout: timeout, LocalAddr: localAddr}
 	if proxyURLs == "" {
 		return proxy.FromEnvironmentUsing(direct).Dial
 	}
@@ -49,3 +49,20 @@ func GetProxyFunc(proxyURLs string, timeout time.Duration, httpEnabled bool) Pro
 		}
 	}
 }
+
+func ResolveAddr(network, addr string) net.Addr {
+	if addr == "" {
+		return nil
+	}
+
+	ip := net.ParseIP(addr)
+
+	switch network {
+	case "tcp", "tcp4", "tcp6":
+		return &net.TCPAddr{IP: ip}
+	case "udp", "udp4", "udp6":
+		return &net.UDPAddr{IP: ip}
+	default:
+		return &net.IPAddr{IP: ip}
+	}
+}