net: implement windows LookupMX and LookupAddr

Also sort SRV records before returning from LookupSRV.

R=rsc
CC=golang-dev
https://golang.org/cl/4817049

5 files changed, 93 insertions, 28 deletions

diff --git a/src/pkg/net/dnsclient.go b/src/pkg/net/dnsclient.go
index 280b19453e..93c04f6b59 100644
--- a/src/pkg/net/dnsclient.go
+++ b/src/pkg/net/dnsclient.go
@@ -9,6 +9,7 @@ import (
 	"fmt"
 	"os"
 	"rand"
+	"sort"
 )
 
 // DNSError represents a DNS lookup error.
@@ -182,9 +183,9 @@ func (s byPriorityWeight) Less(i, j int) bool {
 		(s[i].Priority == s[j].Priority && s[i].Weight < s[j].Weight)
 }
 
-// shuffleSRVByWeight shuffles SRV records by weight using the algorithm
+// shuffleByWeight shuffles SRV records by weight using the algorithm
 // described in RFC 2782.  
-func shuffleSRVByWeight(addrs []*SRV) {
+func (addrs byPriorityWeight) shuffleByWeight() {
 	sum := 0
 	for _, addr := range addrs {
 		sum += int(addr.Weight)
@@ -208,6 +209,19 @@ func shuffleSRVByWeight(addrs []*SRV) {
 	}
 }
 
+// sort reorders SRV records as specified in RFC 2782.
+func (addrs byPriorityWeight) sort() {
+	sort.Sort(addrs)
+	i := 0
+	for j := 1; j < len(addrs); j++ {
+		if addrs[i].Priority != addrs[j].Priority {
+			addrs[i:j].shuffleByWeight()
+			i = j
+		}
+	}
+	addrs[i:].shuffleByWeight()
+}
+
 // An MX represents a single DNS MX record.
 type MX struct {
 	Host string
@@ -222,3 +236,12 @@ func (s byPref) Len() int { return len(s) }
 func (s byPref) Less(i, j int) bool { return s[i].Pref < s[j].Pref }
 
 func (s byPref) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
+
+// sort reorders MX records as specified in RFC 5321.
+func (s byPref) sort() {
+	for i := range s {
+		j := rand.Intn(i + 1)
+		s[i], s[j] = s[j], s[i]
+	}
+	sort.Sort(s)
+}
diff --git a/src/pkg/net/srv_test.go b/src/pkg/net/lookup_test.go
index f1c7a0ab49..995ab03d09 100644
--- a/src/pkg/net/srv_test.go
+++ b/src/pkg/net/lookup_test.go
@@ -27,3 +27,31 @@ func TestGoogleSRV(t *testing.T) {
 		t.Errorf("no results")
 	}
 }
+
+func TestGmailMX(t *testing.T) {
+	if testing.Short() || avoidMacFirewall {
+		t.Logf("skipping test to avoid external network")
+		return
+	}
+	mx, err := LookupMX("gmail.com")
+	if err != nil {
+		t.Errorf("failed: %s", err)
+	}
+	if len(mx) == 0 {
+		t.Errorf("no results")
+	}
+}
+
+func TestGoogleDNSAddr(t *testing.T) {
+	if testing.Short() || avoidMacFirewall {
+		t.Logf("skipping test to avoid external network")
+		return
+	}
+	names, err := LookupAddr("8.8.8.8")
+	if err != nil {
+		t.Errorf("failed: %s", err)
+	}
+	if len(names) == 0 {
+		t.Errorf("no results")
+	}
+}
diff --git a/src/pkg/net/lookup_unix.go b/src/pkg/net/lookup_unix.go
index 168d3fa6d4..8f5e66212b 100644
--- a/src/pkg/net/lookup_unix.go
+++ b/src/pkg/net/lookup_unix.go
@@ -6,8 +6,6 @@ package net
 
 import (
 	"os"
-	"rand"
-	"sort"
 )
 
 // LookupHost looks up the given host using the local resolver.
@@ -68,15 +66,7 @@ func LookupSRV(service, proto, name string) (cname string, addrs []*SRV, err os.
 		r := rr.(*dnsRR_SRV)
 		addrs[i] = &SRV{r.Target, r.Port, r.Priority, r.Weight}
 	}
-	sort.Sort(byPriorityWeight(addrs))
-	i := 0
-	for j := 1; j < len(addrs); j++ {
-		if addrs[i].Priority != addrs[j].Priority {
-			shuffleSRVByWeight(addrs[i:j])
-			i = j
-		}
-	}
-	shuffleSRVByWeight(addrs[i:len(addrs)])
+	byPriorityWeight(addrs).sort()
 	return
 }
 
@@ -91,12 +81,7 @@ func LookupMX(name string) (mx []*MX, err os.Error) {
 		r := rr[i].(*dnsRR_MX)
 		mx[i] = &MX{r.Mx, r.Pref}
 	}
-	// Shuffle the records to match RFC 5321 when sorted
-	for i := range mx {
-		j := rand.Intn(i + 1)
-		mx[i], mx[j] = mx[j], mx[i]
-	}
-	sort.Sort(byPref(mx))
+	byPref(mx).sort()
 	return
 }
 
diff --git a/src/pkg/net/lookup_windows.go b/src/pkg/net/lookup_windows.go
index 16b37f56cb..fa3ad7c7f4 100644
--- a/src/pkg/net/lookup_windows.go
+++ b/src/pkg/net/lookup_windows.go
@@ -85,23 +85,46 @@ func LookupSRV(service, proto, name string) (cname string, addrs []*SRV, err os.
 		return "", nil, os.NewSyscallError("LookupSRV", int(e))
 	}
 	defer syscall.DnsRecordListFree(r, 1)
-	addrs = make([]*SRV, 100)
-	i := 0
+	addrs = make([]*SRV, 0, 10)
 	for p := r; p != nil && p.Type == syscall.DNS_TYPE_SRV; p = p.Next {
 		v := (*syscall.DNSSRVData)(unsafe.Pointer(&p.Data[0]))
-		addrs[i] = &SRV{syscall.UTF16ToString((*[256]uint16)(unsafe.Pointer(v.Target))[:]), v.Port, v.Priority, v.Weight}
-		i++
+		addrs = append(addrs, &SRV{syscall.UTF16ToString((*[256]uint16)(unsafe.Pointer(v.Target))[:]), v.Port, v.Priority, v.Weight})
 	}
-	addrs = addrs[0:i]
+	byPriorityWeight(addrs).sort()
 	return name, addrs, nil
 }
 
-// TODO(brainman): implement LookupMX and LookupAddr.
-
 func LookupMX(name string) (mx []*MX, err os.Error) {
-	return nil, os.NewSyscallError("LookupMX", syscall.EWINDOWS)
+	var r *syscall.DNSRecord
+	e := syscall.DnsQuery(name, syscall.DNS_TYPE_MX, 0, nil, &r, nil)
+	if int(e) != 0 {
+		return nil, os.NewSyscallError("LookupMX", int(e))
+	}
+	defer syscall.DnsRecordListFree(r, 1)
+	mx = make([]*MX, 0, 10)
+	for p := r; p != nil && p.Type == syscall.DNS_TYPE_MX; p = p.Next {
+		v := (*syscall.DNSMXData)(unsafe.Pointer(&p.Data[0]))
+		mx = append(mx, &MX{syscall.UTF16ToString((*[256]uint16)(unsafe.Pointer(v.NameExchange))[:]) + ".", v.Preference})
+	}
+	byPref(mx).sort()
+	return mx, nil
 }
 
 func LookupAddr(addr string) (name []string, err os.Error) {
-	return nil, os.NewSyscallError("LookupAddr", syscall.EWINDOWS)
+	arpa, err := reverseaddr(addr)
+	if err != nil {
+		return nil, err
+	}
+	var r *syscall.DNSRecord
+	e := syscall.DnsQuery(arpa, syscall.DNS_TYPE_PTR, 0, nil, &r, nil)
+	if int(e) != 0 {
+		return nil, os.NewSyscallError("LookupAddr", int(e))
+	}
+	defer syscall.DnsRecordListFree(r, 1)
+	name = make([]string, 0, 10)
+	for p := r; p != nil && p.Type == syscall.DNS_TYPE_PTR; p = p.Next {
+		v := (*syscall.DNSPTRData)(unsafe.Pointer(&p.Data[0]))
+		name = append(name, syscall.UTF16ToString((*[256]uint16)(unsafe.Pointer(v.Host))[:]))
+	}
+	return name, nil
 }
diff --git a/src/pkg/syscall/ztypes_windows.go b/src/pkg/syscall/ztypes_windows.go
index 10780f7672..07f2b85f07 100644
--- a/src/pkg/syscall/ztypes_windows.go
+++ b/src/pkg/syscall/ztypes_windows.go
@@ -478,6 +478,12 @@ type DNSPTRData struct {
 	Host *uint16
 }
 
+type DNSMXData struct {
+	NameExchange *uint16
+	Preference   uint16
+	Pad          uint16
+}
+
 type DNSRecord struct {
 	Next     *DNSRecord
 	Name     *uint16