diff --git a/whois.go b/whois.go index 94ec2db..e49788a 100644 --- a/whois.go +++ b/whois.go @@ -53,6 +53,11 @@ type Client struct { disableReferral bool } +type hasTimeout struct { + Timeout time.Duration + proxy.Dialer +} + // Version returns package version func Version() string { return "1.15.4" @@ -91,6 +96,9 @@ func (c *Client) SetDialer(dialer proxy.Dialer) *Client { // SetTimeout set query timeout func (c *Client) SetTimeout(timeout time.Duration) *Client { + if d, ok := c.dialer.(*hasTimeout); ok { + d.Timeout = timeout + } c.timeout = timeout return c } diff --git a/whois_test.go b/whois_test.go index f9491cd..dae4fe3 100644 --- a/whois_test.go +++ b/whois_test.go @@ -71,6 +71,18 @@ func TestWhoisFail(t *testing.T) { assert.NotNil(t, err) } +func TestWhoisTimeout(t *testing.T) { + client := NewClient() + client.SetTimeout(1 * time.Millisecond) + _, err := client.Whois("google.com") + assert.NotNil(t, err) + assert.Contains(t, err.Error(), "timeout") + + client.SetTimeout(10 * time.Second) + _, err = client.Whois("google.com") + assert.Nil(t, err) +} + func TestWhois(t *testing.T) { tests := []string{ "com",