diff --git a/dcerpc/conn.go b/dcerpc/conn.go index 746bbc69..08b1195b 100644 --- a/dcerpc/conn.go +++ b/dcerpc/conn.go @@ -122,7 +122,7 @@ func Dial(ctx context.Context, addr string, opts ...Option) (Conn, error) { opts: append(opts, WithGroup(group)), } - ip, hostName, binding, err := ParseServerAddrWithDNSLookup(addr, true) + ip, hostName, binding, err := parseServerAddrWithDNSLookup(ctx, addr, settings.DNSResolver) if err != nil { return nil, fmt.Errorf("dial: %w", err) } @@ -156,6 +156,14 @@ func ParseServerAddr(addr string) (net.IP, string, *StringBinding, error) { // ParseServerAddrWithDNSLookup function parses the server address // and performs the DNS lookup if required. func ParseServerAddrWithDNSLookup(addr string, doLookup bool) (net.IP, string, *StringBinding, error) { + var resolver DNSResolver + if doLookup { + resolver = net.DefaultResolver + } + return parseServerAddrWithDNSLookup(context.Background(), addr, resolver) +} + +func parseServerAddrWithDNSLookup(ctx context.Context, addr string, resolver DNSResolver) (net.IP, string, *StringBinding, error) { if ip := net.ParseIP(addr); ip != nil { // this is an ip address. @@ -166,18 +174,18 @@ func ParseServerAddrWithDNSLookup(addr string, doLookup bool) (net.IP, string, * // this is fqdn. - if !doLookup { + if resolver == nil { return nil, addr, nil, nil } - ips, err := net.LookupIP(addr) + ias, err := resolver.LookupIPAddr(ctx, addr) if err != nil { return nil, "", nil, fmt.Errorf("lookup server address: %w", err) } - for _, ip := range ips { - if ip.To4() != nil { - return ip, addr, nil, nil + for _, ia := range ias { + if ia.IP.To4() != nil { + return ia.IP, addr, nil, nil } } diff --git a/dcerpc/transport_settings.go b/dcerpc/transport_settings.go index 26266ecf..3261b322 100644 --- a/dcerpc/transport_settings.go +++ b/dcerpc/transport_settings.go @@ -65,6 +65,8 @@ type Transport struct { // If set to `true`, new connection will be established // for every new client with matching binding. NoReuseTransport bool + // DNS resolver. + DNSResolver DNSResolver } // The transport connection option. @@ -116,6 +118,17 @@ func WithSMBDialer(dialer *smb2.Dialer) ConnectOption { return func(o *Transport) { o.SMBDialer = dialer } } +// DNSResolver interface is used to resolve the hostname to IP addresses. +type DNSResolver interface { + // LookupIPAddr looks up the given host and returns a list of IP addresses. + LookupIPAddr(ctx context.Context, host string) ([]net.IPAddr, error) +} + +// WithDNSResolver function sets the DNS resolver for hostname resolution. +func WithDNSResolver(resolver DNSResolver) ConnectOption { + return func(o *Transport) { o.DNSResolver = resolver } +} + // WithEndpointMapper option sets the endpoint mapper to find the endpoint // (port or named pipe) for the selected abstract syntax. // @@ -144,6 +157,7 @@ func NewTransport() Transport { Timeout: 10 * time.Second, Deadline: 3 * time.Second, SMBPort: 445, + DNSResolver: net.DefaultResolver, } }