1 files changed, 144 insertions, 0 deletions

diff --git a/dns/client.go b/dns/client.go
new file mode 100644
index 0000000..f322b45
--- /dev/null
+++ b/dns/client.go
@@ -0,0 +1,144 @@
+package dns
+
+import (
+        "fmt"
+
+        "github.com/miekg/dns"
+
+        "code.crute.me/mcrute/go_ddns_manager/bind"
+)
+
+type DNSClient struct {
+        Server string
+}
+
+type DNSTransaction struct {
+        zone *bind.Zone
+        key  *bind.Key
+        msg  *dns.Msg
+}
+
+func (t *DNSTransaction) Upsert(rrs ...RR) *DNSTransaction {
+        t.RemoveAll(rrs...).Insert(rrs...)
+        return t
+}
+
+func (t *DNSTransaction) Insert(rrs ...RR) *DNSTransaction {
+        t.msg.Insert(toRRSet(t.zone, rrs...))
+        return t
+}
+
+func (t *DNSTransaction) Remove(rrs ...RR) *DNSTransaction {
+        t.msg.Remove(toRRSet(t.zone, rrs...))
+        return t
+}
+
+func (t *DNSTransaction) RemoveAll(rrs ...RR) *DNSTransaction {
+        t.msg.RemoveRRset(toRRSet(t.zone, rrs...))
+        return t
+}
+
+func (c *DNSClient) AXFR(zone *bind.Zone) (chan *dns.Envelope, error) {
+        k := zone.Keys()[0]
+        t := &dns.Transfer{TsigSecret: k.AsMap()} // Always uses tcp
+
+        m := &dns.Msg{}
+        m.SetAxfr(zone.Name)
+        k.Sign(m)
+
+        return t.In(m, fmt.Sprintf("%s:53", c.Server))
+}
+
+func (c *DNSClient) ReadRemoteZone(zone *bind.Zone) ([]RR, error) {
+        rrs := []RR{}
+        seenSoa := false
+
+        data, err := c.AXFR(zone)
+        if err != nil {
+                return nil, err
+        }
+
+        for rd := range data {
+                for _, r := range rd.RR {
+                        switch dr := FromDNS(r).(type) {
+                        case *SOA:
+                                // Transfers have 2 SOA records, exclude the last one
+                                if !seenSoa {
+                                        rrs = append(rrs, dr)
+                                        seenSoa = true
+                                }
+                        case RR:
+                                rrs = append(rrs, dr)
+                        default:
+                                // This should only be possible if we somehow are
+                                // missing generated DNS data types.
+                                return nil, fmt.Errorf("Invalid return type")
+                        }
+                }
+        }
+
+        return rrs, nil
+}
+
+func (c *DNSClient) StartUpdate(zone *bind.Zone) *DNSTransaction {
+        m := &dns.Msg{}
+        m.SetUpdate(zone.Name)
+
+        return &DNSTransaction{
+                zone: zone,
+                key:  zone.Keys()[0],
+                msg:  m,
+        }
+}
+
+func (c *DNSClient) QueryRecursive(zone *bind.Zone, fqdn string, rtype uint16) *DNSTransaction {
+        m := &dns.Msg{}
+        m.RecursionDesired = true
+        m.SetQuestion(fqdn, rtype)
+
+        return &DNSTransaction{
+                zone: zone,
+                key:  zone.Keys()[0],
+                msg:  m,
+        }
+}
+
+func (c *DNSClient) SendUpdate(t *DNSTransaction) error {
+        udp := &dns.Client{Net: "udp", TsigSecret: t.key.AsMap()}
+        tcp := &dns.Client{Net: "tcp", TsigSecret: t.key.AsMap()}
+
+        t.msg.SetEdns0(4096, false)
+        t.key.Sign(t.msg)
+
+        in, _, err := udp.Exchange(t.msg, c.Server)
+        if in != nil && in.Truncated {
+                // If the TCP request succeeds, the err will reset to nil
+                in, _, err = tcp.Exchange(t.msg, c.Server)
+        }
+
+        if err != nil {
+                return err
+        }
+
+        return nil
+}
+
+func (c *DNSClient) SendQuery(t *DNSTransaction) ([]dns.RR, error) {
+        udp := &dns.Client{Net: "udp", TsigSecret: t.key.AsMap()}
+        tcp := &dns.Client{Net: "tcp", TsigSecret: t.key.AsMap()}
+
+        t.msg.SetEdns0(4096, false)
+        t.key.Sign(t.msg)
+
+        in, _, err := udp.Exchange(t.msg, c.Server)
+        if in != nil && in.Truncated {
+                // If the TCP request succeeds, the err will reset to nil
+                in, _, err = tcp.Exchange(t.msg, c.Server)
+        }
+
+        if err != nil {
+                return nil, err
+        }
+
+        return in.Answer, nil
+}