fix modules

7 files changed, 187 insertions, 31 deletions

diff --git a/.travis.yml b/.travis.yml
deleted file mode 100644
index ec3363d..0000000
--- a/.travis.yml
+++ /dev/null
@@ -1,24 +0,0 @@
-language: go
-
-go:
-  - 1.10.x
-  - 1.11.x
-  - 1.12.x
-  - master
-
-gobuild_args: -tags=integration -race
-
-sudo: required
-
-services:
-  - docker
-
-before_install:
-  - docker pull jcmturner/gokrb5:dns
-  - docker run -d -h kdc.test.gokrb5 -v /etc/localtime:/etc/localtime:ro -e "TEST_KDC_ADDR=127.0.0.1" -e "TEST_HTTP_ADDR=127.0.0.1" -p 53:53 -p 53:53/udp --name dns jcmturner/gokrb5:dns
-
-before_script:
-  - sudo sed -i 's/nameserver .*/nameserver 127.0.0.1/g' /etc/resolv.conf
-
-env:
-  - DNSUTILS_OVERRIDE_NS="127.0.0.1:53"
\ No newline at end of file
diff --git a/go.mod b/go.mod
deleted file mode 100644
index 0ac2665..0000000
--- a/go.mod
+++ /dev/null
@@ -1,5 +0,0 @@
-module github.com/jcmturner/dnsutils
-
-go 1.12
-
-require github.com/stretchr/testify v1.3.0
diff --git a/v2/go.mod b/v2/go.mod
new file mode 100644
index 0000000..f75ac6d
--- /dev/null
+++ b/v2/go.mod
@@ -0,0 +1,5 @@
+module github.com/jcmturner/dnsutils/v2
+
+go 1.13
+
+require github.com/stretchr/testify v1.4.0
diff --git a/go.sum b/v2/go.sum
index 4347755..e863f51 100644
--- a/go.sum
+++ b/v2/go.sum
@@ -3,5 +3,8 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
-github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q=
-github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
+github.com/stretchr/testify v1.4.0 h1:2E4SXV/wtOkTonXsotYi4li6zVWxYlZuYNCXe9XRJyk=
+github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
+gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
diff --git a/v2/srv.go b/v2/srv.go
new file mode 100644
index 0000000..15ea912
--- /dev/null
+++ b/v2/srv.go
@@ -0,0 +1,95 @@
+package dnsutils
+
+import (
+	"math/rand"
+	"net"
+	"sort"
+)
+
+// OrderedSRV returns a count of the results and a map keyed on the order they should be used.
+// This based on the records' priority and randomised selection based on their relative weighting.
+// The function's inputs are the same as those for net.LookupSRV
+// To use in the correct order:
+//
+// count, orderedSRV, err := OrderedSRV(service, proto, name)
+// i := 1
+// for  i <= count {
+//   srv := orderedSRV[i]
+//   // Do something such as dial this SRV. If fails move on the the next or break if it succeeds.
+//   i += 1
+// }
+func OrderedSRV(service, proto, name string) (int, map[int]*net.SRV, error) {
+	_, addrs, err := net.LookupSRV(service, proto, name)
+	if err != nil {
+		return 0, make(map[int]*net.SRV), err
+	}
+	index, osrv := orderSRV(addrs)
+	return index, osrv, nil
+}
+
+func orderSRV(addrs []*net.SRV) (int, map[int]*net.SRV) {
+	// Initialise the ordered map
+	var o int
+	osrv := make(map[int]*net.SRV)
+
+	prioMap := make(map[int][]*net.SRV, 0)
+	for _, srv := range addrs {
+		prioMap[int(srv.Priority)] = append(prioMap[int(srv.Priority)], srv)
+	}
+
+	priorities := make([]int, 0)
+	for p := range prioMap {
+		priorities = append(priorities, p)
+	}
+
+	var count int
+	sort.Ints(priorities)
+	for _, p := range priorities {
+		tos := weightedOrder(prioMap[p])
+		for i, s := range tos {
+			count += 1
+			osrv[o+i] = s
+		}
+		o += len(tos)
+	}
+	return count, osrv
+}
+
+func weightedOrder(srvs []*net.SRV) map[int]*net.SRV {
+	// Get the total weight
+	var tw int
+	for _, s := range srvs {
+		tw += int(s.Weight)
+	}
+
+	// Initialise the ordered map
+	o := 1
+	osrv := make(map[int]*net.SRV)
+
+	// Whilst there are still entries to be ordered
+	l := len(srvs)
+	for l > 0 {
+		i := rand.Intn(l)
+		s := srvs[i]
+		var rw int
+		if tw > 0 {
+			// Greater the weight the more likely this will be zero or less
+			rw = rand.Intn(tw) - int(s.Weight)
+		}
+		if rw <= 0 {
+			// Put entry in position
+			osrv[o] = s
+			if len(srvs) > 1 {
+				// Remove the entry from the source slice by swapping with the last entry and truncating
+				srvs[len(srvs)-1], srvs[i] = srvs[i], srvs[len(srvs)-1]
+				srvs = srvs[:len(srvs)-1]
+				l = len(srvs)
+			} else {
+				l = 0
+			}
+			o += 1
+			tw = tw - int(s.Weight)
+		}
+	}
+	return osrv
+}
diff --git a/v2/srv_integration_test.go b/v2/srv_integration_test.go
new file mode 100644
index 0000000..6ba8f51
--- /dev/null
+++ b/v2/srv_integration_test.go
@@ -0,0 +1,40 @@
+// +build integration
+// To turn on this test use -tags=integration in go test command
+
+package dnsutils
+
+import (
+	"github.com/stretchr/testify/assert"
+	"strconv"
+	"strings"
+	"testing"
+)
+
+func TestResolveKDC(t *testing.T) {
+	for i := 0; i < 100; i++ {
+		count, res, err := OrderedSRV("kerberos", "tcp", "test.gokrb5")
+		if err != nil {
+			t.Errorf("error resolving SRV DNS records: %v", err)
+		}
+		expected := []string{
+			"kdc.test.gokrb5:88",
+			"kdc1a.test.gokrb5:88",
+			"kdc2a.test.gokrb5:88",
+			"kdc1b.test.gokrb5:88",
+			"kdc2b.test.gokrb5:88",
+		}
+		assert.Equal(t, len(expected), count, "Number of SRV records not as expected: %v", res)
+		assert.Equal(t, count, len(res), "Map size does not match: %v", res)
+		for _, s := range expected {
+			var found bool
+			for _, v := range res {
+				srvStr := strings.TrimRight(v.Target, ".") + ":" + strconv.Itoa(int(v.Port))
+				if s == srvStr {
+					found = true
+					break
+				}
+			}
+			assert.True(t, found, "Record %s not found in results", s)
+		}
+	}
+}
diff --git a/v2/srv_test.go b/v2/srv_test.go
new file mode 100644
index 0000000..4596b34
--- /dev/null
+++ b/v2/srv_test.go
@@ -0,0 +1,42 @@
+package dnsutils
+
+import (
+	"github.com/stretchr/testify/assert"
+	"net"
+	"testing"
+)
+
+func TestOrderSRV(t *testing.T) {
+	srv11 := net.SRV{
+		Target:   "t11",
+		Port:     1234,
+		Priority: 1,
+		Weight:   100,
+	}
+	srv12 := net.SRV{
+		Target:   "t12",
+		Port:     1234,
+		Priority: 1,
+		Weight:   100,
+	}
+	srv13 := net.SRV{
+		Target:   "t13",
+		Port:     1234,
+		Priority: 1,
+		Weight:   20,
+	}
+	srv21 := net.SRV{
+		Target:   "t21",
+		Port:     1234,
+		Priority: 2,
+		Weight:   1,
+	}
+
+	addrs := []*net.SRV{
+		&srv11, &srv21, &srv12, &srv13,
+	}
+	count, orderedSRV := orderSRV(addrs)
+	assert.Equal(t, len(addrs), count, "Index not the expected size")
+	assert.Equal(t, len(addrs), len(orderedSRV), "orderedSRV not the expected size")
+	assert.Equal(t, uint16(2), orderedSRV[4].Priority, "Priority order not as expected")
+}