From 7c477840773d028e9838e73cf79e8a890faaaeac Mon Sep 17 00:00:00 2001
From: Tobias Krischer <tobias.krischer@elyxon.de>
Date: Sun, 1 Nov 2020 16:40:22 +0100
Subject: [PATCH] add tests for self-serve

---
 setup_test.go |  7 +++++-
 wgsd_test.go  | 62 ++++++++++++++++++++++++++++++++++++++++++++-------
 2 files changed, 60 insertions(+), 9 deletions(-)

diff --git a/setup_test.go b/setup_test.go
index 39e7410..db36228 100644
--- a/setup_test.go
+++ b/setup_test.go
@@ -22,9 +22,14 @@ func TestSetup(t *testing.T) {
 			"wgsd example.com.",
 			true,
 		},
+		{
+			"with local ip",
+			"wgsd example.com. wg0 127.0.0.1",
+			false,
+		},
 		{
 			"too many tokens",
-			"wgsd example.com. wg0 extra",
+			"wgsd example.com. wg0 127.0.0.1 extra",
 			true,
 		},
 	}
diff --git a/wgsd_test.go b/wgsd_test.go
index 88aea11..7bc527f 100644
--- a/wgsd_test.go
+++ b/wgsd_test.go
@@ -16,13 +16,16 @@ import (
 )
 
 type mockClient struct {
+	host  wgtypes.Peer
 	peers []wgtypes.Peer
 }
 
 func (m *mockClient) Device(d string) (*wgtypes.Device, error) {
 	return &wgtypes.Device{
-		Name:  d,
-		Peers: m.peers,
+		Name:       d,
+		PublicKey:  m.host.PublicKey,
+		ListenPort: m.host.Endpoint.Port,
+		Peers:      m.peers,
 	}, nil
 }
 
@@ -35,22 +38,32 @@ func base32EqualsPubKey(t *testing.T, b32 string, pubKey wgtypes.Key) bool {
 }
 
 func TestWGSD(t *testing.T) {
+	hostkey := [32]byte{}
+	hostkey[0] = 1
+	host := wgtypes.Peer{
+		Endpoint: &net.UDPAddr{
+			IP:   net.ParseIP("127.0.0.1"),
+			Port: 1,
+		},
+		PublicKey: hostkey,
+	}
+	hostb32 := strings.ToLower(base32.StdEncoding.EncodeToString(host.PublicKey[:]))
 	key1 := [32]byte{}
-	key1[0] = 1
+	key1[0] = 2
 	peer1 := wgtypes.Peer{
 		Endpoint: &net.UDPAddr{
 			IP:   net.ParseIP("1.1.1.1"),
-			Port: 1,
+			Port: 2,
 		},
 		PublicKey: key1,
 	}
 	peer1b32 := strings.ToLower(base32.StdEncoding.EncodeToString(peer1.PublicKey[:]))
 	key2 := [32]byte{}
-	key2[0] = 2
+	key2[0] = 3
 	peer2 := wgtypes.Peer{
 		Endpoint: &net.UDPAddr{
 			IP:   net.ParseIP("::2"),
-			Port: 2,
+			Port: 3,
 		},
 		PublicKey: key2,
 	}
@@ -58,10 +71,12 @@ func TestWGSD(t *testing.T) {
 	p := &WGSD{
 		Next: test.ErrorHandler(),
 		client: &mockClient{
+			host:  host,
 			peers: []wgtypes.Peer{peer1, peer2},
 		},
 		zone:   "example.com.",
 		device: "wg0",
+		wgIP:   host.Endpoint.IP,
 	}
 
 	testCases := []test.Case{
@@ -70,16 +85,39 @@ func TestWGSD(t *testing.T) {
 			Qtype: dns.TypePTR,
 			Rcode: dns.RcodeSuccess,
 			Answer: []dns.RR{
+				test.PTR(fmt.Sprintf("_wireguard._udp.example.com. 0 IN PTR %s._wireguard._udp.example.com.", hostb32)),
 				test.PTR(fmt.Sprintf("_wireguard._udp.example.com. 0 IN PTR %s._wireguard._udp.example.com.", peer1b32)),
 				test.PTR(fmt.Sprintf("_wireguard._udp.example.com. 0 IN PTR %s._wireguard._udp.example.com.", peer2b32)),
 			},
 		},
+		test.Case{
+			Qname: fmt.Sprintf("_wireguard._udp.example.com."),
+			Qtype: dns.TypeSRV,
+			Rcode: dns.RcodeSuccess,
+			Answer: []dns.RR{
+				test.SRV(fmt.Sprintf("_wireguard._udp.example.com. 0 IN SRV 0 0 1 %s.example.com.", hostb32)),
+			},
+			Extra: []dns.RR{
+				test.A(fmt.Sprintf("%s.example.com. 0 IN A %s", hostb32, host.Endpoint.IP.String())),
+			},
+		},
+		test.Case{
+			Qname: fmt.Sprintf("%s._wireguard._udp.example.com.", hostb32),
+			Qtype: dns.TypeSRV,
+			Rcode: dns.RcodeSuccess,
+			Answer: []dns.RR{
+				test.SRV(fmt.Sprintf("%s._wireguard._udp.example.com. 0 IN SRV 0 0 1 %s.example.com.", hostb32, hostb32)),
+			},
+			Extra: []dns.RR{
+				test.A(fmt.Sprintf("%s.example.com. 0 IN A %s", hostb32, host.Endpoint.IP.String())),
+			},
+		},
 		test.Case{
 			Qname: fmt.Sprintf("%s._wireguard._udp.example.com.", peer1b32),
 			Qtype: dns.TypeSRV,
 			Rcode: dns.RcodeSuccess,
 			Answer: []dns.RR{
-				test.SRV(fmt.Sprintf("%s._wireguard._udp.example.com. 0 IN SRV 0 0 1 %s.example.com.", peer1b32, peer1b32)),
+				test.SRV(fmt.Sprintf("%s._wireguard._udp.example.com. 0 IN SRV 0 0 2 %s.example.com.", peer1b32, peer1b32)),
 			},
 			Extra: []dns.RR{
 				test.A(fmt.Sprintf("%s.example.com. 0 IN A %s", peer1b32, peer1.Endpoint.IP.String())),
@@ -90,12 +128,20 @@ func TestWGSD(t *testing.T) {
 			Qtype: dns.TypeSRV,
 			Rcode: dns.RcodeSuccess,
 			Answer: []dns.RR{
-				test.SRV(fmt.Sprintf("%s._wireguard._udp.example.com. 0 IN SRV 0 0 2 %s.example.com.", peer2b32, peer2b32)),
+				test.SRV(fmt.Sprintf("%s._wireguard._udp.example.com. 0 IN SRV 0 0 3 %s.example.com.", peer2b32, peer2b32)),
 			},
 			Extra: []dns.RR{
 				test.AAAA(fmt.Sprintf("%s.example.com. 0 IN AAAA %s", peer2b32, peer2.Endpoint.IP.String())),
 			},
 		},
+		test.Case{
+			Qname: fmt.Sprintf("%s.example.com.", hostb32),
+			Qtype: dns.TypeA,
+			Rcode: dns.RcodeSuccess,
+			Answer: []dns.RR{
+				test.A(fmt.Sprintf("%s.example.com. 0 IN A %s", hostb32, host.Endpoint.IP.String())),
+			},
+		},
 		test.Case{
 			Qname: fmt.Sprintf("%s.example.com.", peer1b32),
 			Qtype: dns.TypeA,