Skip to content

Commit

Permalink
Add setting new resolver
Browse files Browse the repository at this point in the history
  • Loading branch information
matsuyoshi30 committed Nov 24, 2020
1 parent 0109868 commit a01027f
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 0 deletions.
3 changes: 3 additions & 0 deletions attacker/attacker.go
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,9 @@ func Attack(ctx context.Context, target string, resCh chan<- *Result, metricsCh
if opts.LocalAddr.IP == nil {
opts.LocalAddr = DefaultLocalAddr
}
if len(opts.Resolvers) > 0 {
net.DefaultResolver = NewResolver(opts.Resolvers)
}
if opts.Attacker == nil {
opts.Attacker = vegeta.NewAttacker(
vegeta.Timeout(opts.Timeout),
Expand Down
31 changes: 31 additions & 0 deletions attacker/resolver.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
package attacker

import (
"context"
"net"
"sync/atomic"
"time"
)

type resolver struct {
addrs []string
idx uint64
}

func NewResolver(addrs []string) *net.Resolver {
r := &resolver{addrs: addrs}

return &net.Resolver{
PreferGo: true,
Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
d := net.Dialer{
Timeout: time.Millisecond * time.Duration(2000),
}
return d.DialContext(ctx, network, r.address())
},
}
}

func (r *resolver) address() string {
return r.addrs[atomic.AddUint64(&r.idx, 1)%uint64(len(r.addrs))]
}

0 comments on commit a01027f

Please sign in to comment.