UDP support/documentation, more tests

README.md

@@ -63,6 +63,7 @@ $ net-wait-go
 ## 1 service check
 ```bash
 net-wait-go -addrs ya.ru:443 -debug true
+
 2020/06/30 18:07:38 ya.ru:443 is OK
 
 # return code is 0
@@ -71,6 +72,7 @@ net-wait-go -addrs ya.ru:443 -debug true
 ## 2 services check
 ```bash
 net-wait-go -addrs ya.ru:443,yandex.ru:443 -debug true
+
 2020/06/30 18:09:24 yandex.ru:443 is OK
 2020/06/30 18:09:24 ya.ru:443 is OK
 
@@ -80,10 +82,38 @@ net-wait-go -addrs ya.ru:443,yandex.ru:443 -debug true
 ## 2 services check (fail)
 ```bash
 net-wait-go -addrs ya.ru:445,yandex.ru:445 -debug true
+
 2020/06/30 18:09:24 yandex.ru:445 is FAILED
 2020/06/30 18:09:24 ya.ru:445 is is FAILED
 ...
 # return code is 1 (if at least 1 service is failed)
 ```
 
+# UDP support
+Since UDP as protocol does not provide connection between a server and clients,
+it is not supported in the most of popular
+utilities:
+ - `wait-for-it` issue - https://github.com/vishnubob/wait-for-it/issues/29)
+ - `netcat` (`nc`) has following note in its manual page:
+     ```
+    CAVEATS
+           UDP port scans will always succeed (i.e. report the port as open)
+    ```
+   
+`net-wait-go` provides UDP support, working following way:
+ - sends a meaningful packet to the server
+ - waits for a message back from the server (1 byte at least)
+ 
+## UDP packet example
+Counter Strike game server is accessible via UDP.
+Let's check random Counter Strike server
+ by sending A2S_INFO packet (https://developer.valvesoftware.com/wiki/Server_queries#A2S_INFO)
+ 
+```bash
+net-wait-go -proto udp -addrs 46.174.53.245:27015,185.158.113.136:27015 -packet '/////1RTb3VyY2UgRW5naW5lIFF1ZXJ5AA==' -debug true
+ 
+2020/07/12 15:13:25 udp 185.158.113.136:27015 is OK
+2020/07/12 15:13:25 udp 46.174.53.245:27015 is OK
 
+# return code is 0
+```

main.go

@@ -1,7 +1,9 @@
 package main
 
 import (
+	"encoding/base64"
 	"flag"
+	"fmt"
 	"log"
 	"os"
 	"strings"
@@ -29,6 +31,9 @@ func main() {
 	var debug bool
 	flag.BoolVar(&debug, "debug", false, "debug messages toggler")
 
+	var packetBase64 string
+	flag.StringVar(&packetBase64, "packet", "", "UDP packet to be sent")
+
 	flag.Parse()
 
 	addrsSlice := strings.FieldsFunc(addrs, func(c rune) bool {
@@ -42,12 +47,21 @@ func main() {
 		os.Exit(2)
 	}
 
+	packetBytes, err := base64.StdEncoding.DecodeString(packetBase64)
+	if err != nil {
+		fmt.Println("packet base64 decode error:", err)
+		flag.Usage()
+
+		os.Exit(2)
+	}
+
 	if wait.New(
-		wait.WithProto("tcp"),
+		wait.WithProto(proto),
 		wait.WithWait(time.Duration(delayMS)*time.Millisecond),
 		wait.WithBreak(time.Duration(breakMS)*time.Millisecond),
 		wait.WithDeadline(time.Duration(deadlineMS)*time.Millisecond),
 		wait.WithDebug(debug),
+		wait.WithUDPPacket(packetBytes),
 	).Do(addrsSlice) {
 		return
 	}

wait/wait.go

@@ -8,31 +8,33 @@ import (
 )
 
 type Executor struct {
-	Proto    string
-	Addrs    []string
-	Wait     time.Duration
-	Break    time.Duration
-	Deadline time.Duration
-	Debug    bool
+	Proto     string
+	Wait      time.Duration
+	Break     time.Duration
+	Deadline  time.Duration
+	Debug     bool
+	UDPPacket []byte
 }
 
 type Option func(*Executor)
 
 func New(opts ...Option) *Executor {
 	const (
-		defaultProto    = "tcp"
-		defaultWait     = 200 * time.Millisecond
-		defaultBreak    = 50 * time.Millisecond
-		defaultDeadline = 15 * time.Second
-		defaultDebug    = false
+		defaultProto     = "tcp"
+		defaultWait      = 200 * time.Millisecond
+		defaultBreak     = 50 * time.Millisecond
+		defaultDeadline  = 15 * time.Second
+		defaultDebug     = false
+		defaultUDPPacket = ""
 	)
 
 	e := &Executor{
-		Proto:    defaultProto,
-		Wait:     defaultWait,
-		Break:    defaultBreak,
-		Deadline: defaultDeadline,
-		Debug:    defaultDebug,
+		Proto:     defaultProto,
+		Wait:      defaultWait,
+		Break:     defaultBreak,
+		Deadline:  defaultDeadline,
+		Debug:     defaultDebug,
+		UDPPacket: []byte(defaultUDPPacket),
 	}
 
 	for _, opt := range opts {
@@ -72,6 +74,12 @@ func WithDebug(debug bool) Option {
 	}
 }
 
+func WithUDPPacket(packet []byte) Option {
+	return func(h *Executor) {
+		h.UDPPacket = packet
+	}
+}
+
 func (e *Executor) Do(addrs []string) bool {
 	deadlineCh := time.After(e.Deadline)
 	successCh := make(chan struct{})
@@ -85,26 +93,24 @@ func (e *Executor) Do(addrs []string) bool {
 				defer wg.Done()
 
 				for {
-					conn, err := net.DialTimeout(e.Proto, addr, e.Wait)
-					if err != nil {
-						if e.Debug {
-							log.Printf("%s is FAILED", addr)
+					select {
+					case <-deadlineCh:
+						return
+					default:
+						if e.Proto == "udp" {
+							if !e.doUDP(addr) {
+								continue
+							}
+						} else if !e.doTCP(addr) {
+							continue
 						}
 
-						if e.Break > 0 {
-							time.Sleep(e.Break)
+						if e.Debug {
+							log.Printf("%s %s is OK", e.Proto, addr)
 						}
 
-						continue
-					}
-
-					if e.Debug {
-						log.Printf("%s is OK", addr)
+						return
 					}
-
-					_ = conn.Close()
-
-					return
 				}
 			}(addr)
 		}
@@ -120,3 +126,63 @@ func (e *Executor) Do(addrs []string) bool {
 		return true
 	}
 }
+
+func (e *Executor) doTCP(addr string) bool {
+	conn, err := net.DialTimeout(e.Proto, addr, e.Wait)
+	if err != nil {
+		e.processFail(addr)
+
+		return false
+	}
+
+	defer conn.Close()
+
+	return true
+}
+
+func (e *Executor) doUDP(addr string) bool {
+	udpAddr, err := net.ResolveUDPAddr("udp", addr)
+	if err != nil {
+		return false
+	}
+
+	conn, err := net.DialTimeout(e.Proto, udpAddr.String(), e.Wait)
+	if err != nil {
+		e.processFail(addr)
+
+		return false
+	}
+
+	defer conn.Close()
+
+	// If UDP packet is set - send it
+	if len(e.UDPPacket) > 0 {
+		_, err = conn.Write(e.UDPPacket)
+		if err != nil {
+			e.processFail(addr)
+
+			return false
+		}
+	}
+
+	// Wait for at least 1 byte response
+	d := make([]byte, 1)
+	_, err = conn.Read(d)
+	if err != nil {
+		e.processFail(addr)
+
+		return false
+	}
+
+	return true
+}
+
+func (e *Executor) processFail(addr string) {
+	if e.Debug {
+		log.Printf("%s %s is FAILED", e.Proto, addr)
+	}
+
+	if e.Break > 0 {
+		time.Sleep(e.Break)
+	}
+}