Refactor into an application

1 files changed, 44 insertions, 34 deletions

diff --git a/dns/cilent.go b/dns/cilent.go
index 63e3734..989bb4d 100644
--- a/dns/cilent.go
+++ b/dns/cilent.go
@@ -1,17 +1,44 @@
 package dns
 
 import (
-        "code.crute.me/mcrute/go_ddns_manager/bind"
         "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()}
+        t := &dns.Transfer{TsigSecret: k.AsMap()} // Always uses tcp
 
         m := &dns.Msg{}
         m.SetAxfr(zone.Name)
@@ -20,48 +47,31 @@ func (c *DNSClient) AXFR(zone *bind.Zone) (chan *dns.Envelope, error) {
         return t.In(m, c.Server)
 }
 
-func (c *DNSClient) Insert(zone *bind.Zone, rrs ...RR) error {
-        k := zone.Keys()[0]
-        dc := &dns.Client{TsigSecret: k.AsMap()}
-
+func (c *DNSClient) StartUpdate(zone *bind.Zone) *DNSTransaction {
         m := &dns.Msg{}
         m.SetUpdate(zone.Name)
-        m.Insert(toRRSet(zone, rrs...))
-        k.Sign(m)
 
-        if _, _, err := dc.Exchange(m, c.Server); err != nil {
-                return err
+        return &DNSTransaction{
+                zone: zone,
+                key:  zone.Keys()[0],
+                msg:  m,
         }
-
-        return nil
 }
 
-func (c *DNSClient) Remove(zone *bind.Zone, rrs ...RR) error {
-        k := zone.Keys()[0]
-        dc := &dns.Client{TsigSecret: k.AsMap()}
+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()}
 
-        m := &dns.Msg{}
-        m.SetUpdate(zone.Name)
-        m.Remove(toRRSet(zone, rrs...))
-        k.Sign(m)
+        t.msg.SetEdns0(4096, false)
+        t.key.Sign(t.msg)
 
-        if _, _, err := dc.Exchange(m, c.Server); err != nil {
-                return err
+        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)
         }
 
-        return nil
-}
-
-func (c *DNSClient) RemoveAll(zone *bind.Zone, rrs ...RR) error {
-        k := zone.Keys()[0]
-        dc := &dns.Client{TsigSecret: k.AsMap()}
-
-        m := &dns.Msg{}
-        m.SetUpdate(zone.Name)
-        m.RemoveRRset(toRRSet(zone, rrs...))
-        k.Sign(m)
-
-        if _, _, err := dc.Exchange(m, c.Server); err != nil {
+        if err != nil {
                 return err
         }