1 files changed, 69 insertions, 1 deletions

diff --git a/dns/client.go b/dns/client.go
index 525444a..f39bb4b 100644
--- a/dns/client.go
+++ b/dns/client.go
@@ -1,7 +1,9 @@
 package dns
 
 import (
+        "context"
         "fmt"
+        "time"
 
         "github.com/miekg/dns"
 
@@ -9,7 +11,9 @@ import (
 )
 
 type DNSClient struct {
-        Server string
+        Server             string
+        RecursiveResolvers []string
+        PollTimeout        time.Duration
 }
 
 type DNSTransaction struct {
@@ -142,3 +146,67 @@ func (c *DNSClient) SendQuery(t *DNSTransaction) ([]dns.RR, error) {
 
         return in.Answer, nil
 }
+
+// TODO: Copied from the letsencrypt service, merge this into existing functions
+func (c *DNSClient) sendReadQuery(ctx context.Context, fqdn string, rtype uint16, nameserver string) (*dns.Msg, error) {
+        udp := &dns.Client{Net: "udp"}
+        tcp := &dns.Client{Net: "tcp"}
+
+        m := &dns.Msg{}
+        m.SetQuestion(fqdn, rtype)
+        m.SetEdns0(4096, false)
+        m.RecursionDesired = true
+
+        in, _, err := udp.ExchangeContext(ctx, m, nameserver)
+        if in != nil && in.Truncated {
+                // If the TCP request succeeds, the err will reset to nil
+                in, _, err = tcp.ExchangeContext(ctx, m, nameserver)
+        }
+
+        if err != nil {
+                return nil, err
+        }
+
+        return in, err
+}
+
+func (c *DNSClient) WaitForDNSPropagation(ctx context.Context, fqdn, value string) error {
+        if c.RecursiveResolvers == nil {
+                return fmt.Errorf("DNSClient.WaitForDNSPropagation: RecursiveResolvers not set")
+        }
+
+        pt := c.PollTimeout
+        if pt == 0 {
+                pt = 3 * time.Second
+        }
+
+        timer := time.NewTicker(pt)
+        defer timer.Stop()
+
+        for {
+                // Give the server the initial timout to satisfy the request
+                select {
+                case <-ctx.Done():
+                        return fmt.Errorf("DNSClient.WaitForDNSPropagation: context has expired, polling terminated")
+                case <-timer.C:
+                }
+
+                ok_count := 0
+                for _, rs := range c.RecursiveResolvers {
+                        r, err := c.sendReadQuery(ctx, fqdn, dns.TypeTXT, rs)
+                        if err != nil {
+                                return err
+                        }
+
+                        if len(r.Answer) > 0 {
+                                if r.Answer[0].(*dns.TXT).Txt[0] == value {
+                                        ok_count++
+                                }
+                        }
+                }
+
+                if ok_count == len(c.RecursiveResolvers) {
+                        return nil
+                }
+        }
+}