Merge pull request #27 from patrickdappollonio/add-udp-support

Add support for UDP targets.

.github/workflows/release.yml

@@ -26,10 +26,10 @@ jobs:
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
       - name: Release application to Github
-        uses: goreleaser/goreleaser-action@v5
+        uses: goreleaser/goreleaser-action@v6
         with:
           distribution: goreleaser
-          version: latest
+          version: "~> 2"
           args: release --clean
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.goreleaser.yml

@@ -1,3 +1,4 @@
+version: 2
 builds:
   - env:
       - CGO_ENABLED=0

README.md

@@ -1,6 +1,6 @@
 # `wait-for`
 
-A tiny Go application with zero dependencies. Given a number of TCP `host:port` pairs, the app will wait until either all are available or a timeout is reached.
+A tiny Go application with zero dependencies. Given a number of TCP or UDP `host:port` pairs, the app will wait until either all are available or a timeout is reached. `wait-for` supports pinging TCP or UDP hosts, by prefixing the host with `tcp://` or `udp://`, respectively. If no prefix is provided, the app will default to TCP.
 
 Kudos to @vishnubob for the [original implementation in Bash](https://github.com/vishnubob/wait-for-it).
 
@@ -15,11 +15,11 @@ wait-for \
   --timeout 10s
 ```
 
-This will ping both `google.com` on port `443` and `mysql.example.com` on port `3306`. If they both start accepting connections within 10 seconds, the app will exit with a `0` exit code. If either one does not start accepting connections within 10 seconds, the app will exit with a `1` exit code, which will allow you to catch the error in CI/CD environments.
+This will ping both `google.com` on port `443` and `mysql.example.com` on port `3306` via TCP. If they both start accepting connections within 10 seconds, the app will exit with a `0` exit code. If either one does not start accepting connections within 10 seconds, the app will exit with a `1` exit code, which will allow you to catch the error in CI/CD environments.
 
 All the parameters accepted by the application are shown in the help section, as shown below.
 
-```
+```text
 wait-for allows you to wait for a TCP resource to respond to requests.
 
 It does this by performing a TCP connection to the specified host and port. If there's
@@ -36,7 +36,7 @@ Usage:
 Flags:
   -e, --every duration     time to wait between each request attempt against the host (default 1s)
   -h, --help               help for wait-for
-  -s, --host strings       hosts to connect to in the format "host:port"
+  -s, --host strings       hosts to connect to in the format "host:port" with optional protocol prefix (tcp:// or udp://)
   -t, --timeout duration   maximum time to wait for the endpoints to respond before giving up (default 10s)
   -v, --verbose            enable verbose output -- will print every time a request is made
       --version            version for wait-for

app.go

@@ -55,7 +55,7 @@ func root() *cobra.Command {
 		},
 	}
 
-	rootCommand.Flags().StringSliceVarP(&hosts, "host", "s", []string{}, "hosts to connect to in the format \"host:port\"")
+	rootCommand.Flags().StringSliceVarP(&hosts, "host", "s", []string{}, "hosts to connect to in the format \"host:port\" with optional protocol prefix (tcp:// or udp://)")
 	rootCommand.Flags().DurationVarP(&timeout, "timeout", "t", time.Second*10, "maximum time to wait for the endpoints to respond before giving up")
 	rootCommand.Flags().DurationVarP(&step, "every", "e", time.Second*1, "time to wait between each request attempt against the host")
 	rootCommand.Flags().BoolVarP(&verbose, "verbose", "v", false, "enable verbose output -- will print every time a request is made")

wait/ping.go

@@ -31,18 +31,18 @@ func (w *Wait) PingAll() error {
 	}
 }
 
-func (w *Wait) ping(startTime time.Time, host string, wg *sync.WaitGroup) {
+func (w *Wait) ping(startTime time.Time, host host, wg *sync.WaitGroup) {
 	defer wg.Done()
 
 	for {
-		conn, err := net.Dial("tcp", host)
+		conn, err := net.Dial(host.GetProtocol(), host.String())
 		if err == nil {
 			conn.Close()
-			w.log.Printf("> up:   %s (after %s)", w.pad(host), time.Since(startTime))
+			w.log.Printf("> up:   %s (after %s)", w.pad(host.String()), time.Since(startTime))
 			return
 		}
 
-		w.log.Printf("> down: %s", w.pad(host))
+		w.log.Printf("> down: %s", w.pad(host.String()))
 		time.Sleep(w.step)
 	}
 }

wait/wait.go

@@ -6,21 +6,62 @@ import (
 	"log"
 	"net"
 	"os"
+	"regexp"
 	"strings"
 	"time"
 )
 
+type proto string
+
+const (
+	tcp proto = "tcp"
+	udp proto = "udp"
+)
+
+type host struct {
+	host     string
+	port     string
+	protocol proto
+}
+
+func (h host) String() string {
+	return fmt.Sprintf("%s:%s", h.host, h.port)
+}
+
+func (h host) GetProtocol() string {
+	if h.protocol == "" {
+		return string(tcp)
+	}
+
+	return string(h.protocol)
+}
+
+func stringifyHosts(hosts []host) string {
+	var sb strings.Builder
+
+	for i, v := range hosts {
+		if i > 0 {
+			sb.WriteString(", ")
+		}
+
+		sb.WriteString(`"` + fmt.Sprintf("%s://%s:%s", v.GetProtocol(), v.host, v.port) + `"`)
+	}
+
+	return sb.String()
+}
+
 type Wait struct {
-	hosts   []string
+	hosts   []host
 	timeout time.Duration
 	step    time.Duration
 	log     *log.Logger
 	padding int
 }
 
+var reLooksLikeProtocol = regexp.MustCompile(`^(\w+)://`)
+
 func New(hosts []string, step, timeout time.Duration, verbose bool) (*Wait, error) {
 	w := &Wait{
-		hosts:   hosts,
 		timeout: timeout,
 		step:    step,
 	}
@@ -29,16 +70,41 @@ func New(hosts []string, step, timeout time.Duration, verbose bool) (*Wait, erro
 		return nil, fmt.Errorf("no hosts specified")
 	}
 
+	full := make([]host, 0, len(hosts))
 	for _, v := range hosts {
 		if len(v) > w.padding {
 			w.padding = len(v)
 		}
 
-		if _, _, err := net.SplitHostPort(v); err != nil {
-			return nil, fmt.Errorf("invalid host format: %q -- must be in the format \"host:port\"", v)
+		var proto proto
+
+		if strings.HasPrefix(v, "tcp://") {
+			proto = tcp
+			v = strings.TrimPrefix(v, "tcp://")
 		}
+
+		if strings.HasPrefix(v, "udp://") {
+			proto = udp
+			v = strings.TrimPrefix(v, "udp://")
+		}
+
+		if proto == "" && reLooksLikeProtocol.MatchString(v) {
+			return nil, fmt.Errorf("invalid protocol specified: %q -- only \"tcp\" and \"udp\" are supported", v)
+		}
+
+		parsedHost, parsedPort, err := net.SplitHostPort(v)
+		if err != nil {
+			return nil, fmt.Errorf("invalid host format: %q -- must be in the format \"host:port\" or \"(tcp|udp)://host:port\"", v)
+		}
+
+		full = append(full, host{
+			host:     parsedHost,
+			port:     parsedPort,
+			protocol: proto,
+		})
 	}
 
+	w.hosts = full
 	w.log = log.New(io.Discard, "", 0)
 
 	if verbose {
@@ -51,7 +117,7 @@ func New(hosts []string, step, timeout time.Duration, verbose bool) (*Wait, erro
 func (w *Wait) String() string {
 	return fmt.Sprintf(
 		"Waiting for hosts: %s (timeout: %s, attempting every %s)",
-		strings.Join(w.hosts, ", "),
+		stringifyHosts(w.hosts),
 		w.timeout,
 		w.step,
 	)