Merge pull request #18 from dnnrly/feature/dns

Feature/dns

.gitignore

@@ -12,6 +12,7 @@ dist/
 /wait-for
 
 # IDEs
+.vscode/
 .idea/
 *.iml
 

Makefile

@@ -49,6 +49,10 @@ deps: ./bin/tparse
 	go get -v ./...
 	go mod tidy
 
+.PHONY: mocks
+mocks: ## generate mocks for interfaces
+	mockgen -source=waitfor.go -package=waitfor > waitfor_mock_test.go
+
 .PHONY: build
 build: ## build the application
 	go build -o wait-for ./cmd/wait-for

README.md

@@ -7,6 +7,12 @@ or environment.
 Typically, you would use this to wait on another resource (such as an HTTP resource)
 to become available before continuing - or timeout and exit with an error.
 
+At the moment, you can wait for a few different kinds of thing. They are:
+
+* HTTP or HTTPS success response
+* TCP or GRPC connection
+* DNS IP resolve address change
+
 [![GitHub release (latest SemVer)](https://img.shields.io/github/v/release/dnnrly/wait-for)](https://github.com/dnnrly/wait-for/releases/latest)
 [![GitHub Workflow Status](https://img.shields.io/github/workflow/status/dnnrly/wait-for/Release%20workflow)](https://github.com/dnnrly/wait-for/actions?query=workflow%3A%22Release+workflow%22)
 [![codecov](https://codecov.io/gh/dnnrly/wait-for/branch/main/graph/badge.svg?token=s0OfKkTFuI)](https://codecov.io/gh/dnnrly/wait-for)
@@ -28,14 +34,14 @@ go install github.com/dnnrly/wait-for/cmd/wait-for@latest
 If you don't have Go installed (in a Docker container, for example) then you can take advantage of the pre-built versions. Check out the [releases](https://github.com/dnnrly/wait-for/releases) and check out the links for direct downloads. You can download and unpack a release like so:
 
 ```shell
-wget https://github.com/dnnrly/wait-for/releases/download/v0.0.1/wait-for_0.0.1_linux_386.tar.gz
-gunzip wait-for_0.0.1_linux_386.tar.gz
-tar -xfv wait-for_0.0.1_linux_386.tar
+wget https://github.com/dnnrly/wait-for/releases/download/v0.0.5/wait-for_0.0.5_linux_386.tar.gz
+gunzip wait-for_0.0.5_linux_386.tar.gz
+tar -xfv wait-for_0.0.5_linux_386.tar
 ```
 
 In your Dockerfile, you can do this:
 ```docker
-ADD https://github.com/dnnrly/wait-for/releases/download/v0.0.1/wait-for_0.0.1_linux_386.tar.gz wait-for.tar.gz
+ADD https://github.com/dnnrly/wait-for/releases/download/v0.0.1/wait-for_0.0.5_linux_386.tar.gz wait-for.tar.gz
 RUN gunzip wait-for.tar.gz && tar -xf wait-for.tar
 ```
 
@@ -50,10 +56,21 @@ $ wait-for http://your-service-here:8080/health https://another-service/
 ``` 
 
 ### Waiting for gRPC services
+
 ```shell script
 $ wait-for grpc-server:8092 other-grpc-server:9091
 ```
 
+### Waiting for DNS changes
+
+```shell script
+$ wait-for dns:google.com
+```
+
+This will wait for the list of IP addresses bound to that DNS name to be
+updated, regardless of order. You can use this to wait for a DNS update
+such as failover or other similar operations.
+
 ### Preconfiguring services to connect to
 
 ```shell script
@@ -81,6 +98,9 @@ wait-for:
   snmp-service:
     type: tcp
     target: snmp-trap-dns:514
+  dns-thing:
+    type: dns
+    target: your.r53-entry.com
 ```
 
 ### Using `wait-for` in Docker Compose

cmd/wait-for/main.go

@@ -4,6 +4,7 @@ import (
 	"flag"
 	"fmt"
 	"log"
+	"net"
 	"os"
 
 	waitfor "github.com/dnnrly/wait-for"
@@ -40,6 +41,13 @@ func main() {
 		os.Exit(1)
 	}
 
+	waitfor.SupportedWaiters = map[string]waitfor.Waiter{
+		"http": waitfor.WaiterFunc(waitfor.HTTPWaiter),
+		"tcp":  waitfor.WaiterFunc(waitfor.TCPWaiter),
+		"grpc": waitfor.WaiterFunc(waitfor.GRPCWaiter),
+		"dns":  waitfor.NewDNSWaiter(net.LookupIP, logger),
+	}
+
 	err = waitfor.WaitOn(config, logger, flag.Args(), waitfor.SupportedWaiters)
 	if err != nil {
 		_, _ = fmt.Fprintf(os.Stderr, "%v", err)

config.go

@@ -96,6 +96,15 @@ func (c *Config) AddFromString(t string) error {
 		return nil
 	}
 
+	if strings.HasPrefix(t, "dns:") {
+		c.Targets[t] = TargetConfig{
+			Target:  strings.Replace(t, "dns:", "", 1),
+			Type:    "dns",
+			Timeout: c.DefaultTimeout,
+		}
+		return nil
+	}
+
 	return errors.New("unable to understand target " + t)
 }
 

config_test.go

@@ -85,9 +85,10 @@ func TestConfig_AddFromString(t *testing.T) {
 	assert.NoError(t, config.AddFromString("https://some-host/endpoint"))
 	assert.NoError(t, config.AddFromString("http://another-host/endpoint"))
 	assert.NoError(t, config.AddFromString("tcp:listener-tcp:9090"))
+	assert.NoError(t, config.AddFromString("dns:some.dns.com"))
 	assert.Error(t, config.AddFromString("udp:some-listener:9090"))
 
-	assert.Equal(t, 4, len(config.Targets))
+	assert.Equal(t, 5, len(config.Targets))
 
 	assert.Equal(t, "http://some-host/endpoint", config.Targets["http://some-host/endpoint"].Target)
 	assert.Equal(t, "http", config.Targets["http://some-host/endpoint"].Type)
@@ -107,6 +108,10 @@ func TestConfig_AddFromString(t *testing.T) {
 	assert.Equal(t, "listener-tcp:9090", config.Targets["tcp:listener-tcp:9090"].Target)
 	assert.Equal(t, "tcp", config.Targets["tcp:listener-tcp:9090"].Type)
 	assert.Equal(t, time.Second*5, config.Targets["tcp:listener-tcp:9090"].Timeout)
+
+	assert.Equal(t, "some.dns.com", config.Targets["dns:some.dns.com"].Target)
+	assert.Equal(t, "dns", config.Targets["dns:some.dns.com"].Type)
+	assert.Equal(t, time.Second*5, config.Targets["dns:some.dns.com"].Timeout)
 }
 
 func TestConfig_Filters(t *testing.T) {

waitfor.go

@@ -3,35 +3,43 @@ package waitfor
 import (
 	"context"
 	"fmt"
-	"google.golang.org/grpc/credentials/insecure"
 	"net"
 	"net/http"
+	"sort"
+	"strings"
 	"time"
 
+	"google.golang.org/grpc/credentials/insecure"
+
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
 
 	"github.com/spf13/afero"
 )
 
+type Waiter interface {
+	Wait(name string, target *TargetConfig) error
+}
+
 // WaiterFunc is used to implement waiting for a specific type of target.
 // The name is used in the error and target is the actual destination being tested.
 type WaiterFunc func(name string, target *TargetConfig) error
+
+func (w WaiterFunc) Wait(name string, target *TargetConfig) error {
+	return w(name, target)
+}
+
 type Logger func(string, ...interface{})
 
 // NullLogger can be used in place of a real logging function
 var NullLogger = func(f string, a ...interface{}) {}
 
 // SupportedWaiters is a mapping of known protocol names to waiter implementations
-var SupportedWaiters = map[string]WaiterFunc{
-	"http": HTTPWaiter,
-	"tcp":  TCPWaiter,
-	"grpc": GRPCWaiter,
-}
+var SupportedWaiters map[string]Waiter
 
 // WaitOn implements waiting for many targets, using the location of config file provided with named targets to wait until
 // all of those targets are responding as expected
-func WaitOn(config *Config, logger Logger, targets []string, waiters map[string]WaiterFunc) error {
+func WaitOn(config *Config, logger Logger, targets []string, waiters map[string]Waiter) error {
 
 	for _, target := range targets {
 		if !config.GotTarget(target) {
@@ -80,7 +88,7 @@ func OpenConfig(configFile, defaultTimeout, defaultHTTPTimeout string, fs afero.
 	return config, nil
 }
 
-func waitOnTargets(logger Logger, targets map[string]TargetConfig, waiters map[string]WaiterFunc) error {
+func waitOnTargets(logger Logger, targets map[string]TargetConfig, waiters map[string]Waiter) error {
 	var eg errgroup.Group
 
 	for name, target := range targets {
@@ -108,14 +116,14 @@ func waitOnTargets(logger Logger, targets map[string]TargetConfig, waiters map[s
 	return nil
 }
 
-func waitOnSingleTarget(name string, logger Logger, target TargetConfig, waiter WaiterFunc) error {
+func waitOnSingleTarget(name string, logger Logger, target TargetConfig, waiter Waiter) error {
 	end := time.Now().Add(target.Timeout)
 
-	err := waiter(name, &target)
+	err := waiter.Wait(name, &target)
 	for err != nil && end.After(time.Now()) {
 		logger("error while waiting for %s: %v", name, err)
 		time.Sleep(time.Second)
-		err = waiter(name, &target)
+		err = waiter.Wait(name, &target)
 	}
 
 	if err != nil {
@@ -182,3 +190,59 @@ func isSuccess(code int) bool {
 
 	return true
 }
+
+type DNSLookup func(host string) ([]net.IP, error)
+
+type DNSWaiter struct {
+	lookup DNSLookup
+	logger Logger
+}
+
+func NewDNSWaiter(lookup DNSLookup, logger Logger) *DNSWaiter {
+	return &DNSWaiter{
+		lookup: lookup,
+		logger: logger,
+	}
+}
+
+type IPList []net.IP
+
+func (l IPList) Equals(r IPList) bool {
+	return l.String() == r.String()
+}
+
+func (l IPList) Len() int {
+	return len(l)
+}
+func (l IPList) Swap(i, j int)      { l[i], l[j] = l[j], l[i] }
+func (l IPList) Less(i, j int) bool { return strings.Compare(l[i].String(), l[j].String()) < 0 }
+func (l IPList) String() string {
+	sort.Sort(l)
+	var s []string
+	for _, v := range l {
+		s = append(s, v.String())
+	}
+	return strings.Join(s, ",")
+}
+
+func (w *DNSWaiter) Wait(host string, target *TargetConfig) error {
+	in, _ := w.lookup(target.Target)
+	initial := IPList(in)
+	last := initial
+
+	start := time.Now()
+	now := start
+
+	for now.Sub(start) < target.Timeout {
+		w.logger("got DNS result %s", last)
+		time.Sleep(time.Second)
+		l, _ := w.lookup(target.Target)
+		last = IPList(l)
+
+		if !initial.Equals(last) {
+			return nil
+		}
+		now = time.Now()
+	}
+	return fmt.Errorf("timed out waiting for DNS update to %s", host)
+}