diff --git a/README.md b/README.md index 860f50e..db60e0e 100644 --- a/README.md +++ b/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 +``` diff --git a/main.go b/main.go index d41d9f9..6cb263c 100644 --- a/main.go +++ b/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 } diff --git a/wait/wait.go b/wait/wait.go index 597cd28..11c219e 100644 --- a/wait/wait.go +++ b/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) + } +} diff --git a/wait/wait_test.go b/wait/wait_test.go index 5ef37f4..1c4e4e3 100644 --- a/wait/wait_test.go +++ b/wait/wait_test.go @@ -1,41 +1,132 @@ package wait import ( + "io" "net" + "strings" "testing" + "time" ) -func getServer(addr string) net.Listener { - srv, err := net.Listen("tcp", addr) +func getTCPServer(proto, addr string, t *testing.T) io.Closer { + srv, err := net.Listen(proto, addr) if err != nil { panic(err) } go func() { for { - conn, _ := srv.Accept() - if conn != nil { - _ = conn.Close() + conn, err := srv.Accept() + if err != nil { + if !strings.Contains(err.Error(), "use of closed network connection") { + t.Errorf("error listening UDP %s: %s", addr, err.Error()) + } + return } + + _ = conn.Close() } }() return srv } -func TestTCP(t *testing.T) { - ok := "localhost:6432" - notok := "localhost:6431" +func getUDPServer(proto, addr string, t *testing.T) io.Closer { + udpAddr, err := net.ResolveUDPAddr(proto, addr) + if err != nil { + t.Errorf("error resolving UDP address %s: %s", udpAddr, err.Error()) + return nil + } + + conn, err := net.ListenUDP(proto, udpAddr) + if err != nil { + t.Errorf("error listening UDP %s: %s", udpAddr, err.Error()) + return nil + } + + go func() { + for { + var buf [1]byte + _, udpRemoteAddr, err := conn.ReadFromUDP(buf[0:]) + if err != nil { + if !strings.Contains(err.Error(), "use of closed network connection") { + t.Errorf("error listening UDP %s: %s", addr, err.Error()) + } + return + } + + _, err = conn.WriteToUDP(buf[0:], udpRemoteAddr) + if err != nil { + t.Errorf("error writing from server to UDP client: %s", err.Error()) + return + } + } + }() - srv := getServer(ok) - defer srv.Close() + return conn +} - e := New(WithDebug(true)) - if !e.Do([]string{ok}) { - t.FailNow() +func TestDO(t *testing.T) { + type data struct { + name string + addr string + reqAddr string + proto string + packet string + result bool } - if e.Do([]string{notok}) { - t.FailNow() + for _, row := range []data{ + { + name: "tcp success", + addr: "localhost:6432", + reqAddr: "localhost:6432", + proto: "tcp", + result: true, + }, + { + name: "tcp fail", + addr: "localhost:6432", + reqAddr: "localhost:6431", + proto: "tcp", + result: false, + }, + { + name: "udp success", + addr: "localhost:6433", + reqAddr: "localhost:6433", + proto: "udp", + packet: "1", + result: true, + }, + { + name: "udp fail", + addr: "localhost:6434", + reqAddr: "localhost:6435", + proto: "udp", + packet: "1", + result: false, + }, + } { + r := row + t.Run(row.name, func(t *testing.T) { + var srv io.Closer + if r.proto == "udp" { + srv = getUDPServer(r.proto, r.addr, t) + } else { + srv = getTCPServer(r.proto, r.addr, t) + } + defer srv.Close() + + e := New( + WithProto(r.proto), + WithUDPPacket([]byte(r.packet)), + WithDebug(false), + WithDeadline(time.Second*2), + ) + if e.Do([]string{r.reqAddr}) != r.result { + t.Errorf("%s result is not %#v", r.name, r.result) + } + }) } }