diff --git a/main.go b/main.go index 42b9206..a83e7af 100644 --- a/main.go +++ b/main.go @@ -87,8 +87,17 @@ func init() { if len(agentConfig.DNS) > 0 { dnsServers = agentConfig.DNS } - dnsServer := dnsServers[time.Now().Unix()%int64(len(dnsServers))] - return d.DialContext(ctx, "udp", dnsServer) + index := int(time.Now().Unix()) % int(len(dnsServers)) + queue := generateQueue(index, len(dnsServers)) + var conn net.Conn + var err error + for i := 0; i < len(queue); i++ { + conn, err = d.DialContext(ctx, "udp", dnsServers[queue[i]]) + if err == nil { + return conn, nil + } + } + return nil, err } flag.CommandLine.ParseErrorsWhitelist.UnknownFlags = true @@ -636,3 +645,15 @@ func println(v ...interface{}) { fmt.Println(v...) } } + +func generateQueue(start int, size int) []int { + var result []int + for i := start; i < start+size; i++ { + if i < size { + result = append(result, i) + } else { + result = append(result, i-size) + } + } + return result +} diff --git a/main_test.go b/main_test.go new file mode 100644 index 0000000..2d801e7 --- /dev/null +++ b/main_test.go @@ -0,0 +1,25 @@ +package main + +import ( + "reflect" + "testing" +) + +func Test(t *testing.T) { + cases := []struct { + start, size int + want []int + }{ + {0, 2, []int{0, 1}}, + {1, 2, []int{1, 0}}, + {0, 3, []int{0, 1, 2}}, + {1, 3, []int{1, 2, 0}}, + {2, 3, []int{2, 0, 1}}, + } + + for _, c := range cases { + if !reflect.DeepEqual(c.want, generateQueue(c.start, c.size)) { + t.Errorf("generateQueue(%d, %d) == %d, want %d", c.start, c.size, generateQueue(c.start, c.size), c.want) + } + } +}