Listen UDP for probing

README.md

@@ -22,6 +22,8 @@ FROM <from_ip2> CONTENT <content2>
 ...
 ```
 
+此外，还可以通过 UDP 广播获得服务端地址。服务器监听 65433/udp 端口，你可以发送 `PROBE`，服务器会返回 `HERE`。
+
 ## Known Issues
 
 好像对 IPv6 支持还有点问题。

internal/udp.go

@@ -0,0 +1,25 @@
+package internal
+
+import (
+	"net"
+	"regexp"
+
+	"go.uber.org/zap"
+)
+
+var udpRegex = regexp.MustCompile(`^PROBE`)
+
+// UDPHandler 处理广播的消息
+func UDPHandler(ln *net.UDPConn) {
+	data := make([]byte, 1024)
+	for {
+		n, addr, err := ln.ReadFromUDP(data)
+		if err != nil {
+			Logger.Fatal("Read error", zap.Error(err))
+		}
+		Logger.Info("Received", zap.String("addr", addr.String()), zap.String("msg", string(data[:n])))
+		if udpRegex.MatchString(string(data[:n])) {
+			ln.WriteToUDP([]byte("HERE"), addr)
+		}
+	}
+}

main.go

@@ -1,20 +1,33 @@
 package main
 
 import (
-	"log"
 	"net"
 	"socket-server/internal"
+	"sync"
 
 	"go.uber.org/zap"
 )
 
 func main() {
-	ln, err := net.Listen("tcp", ":65432")
-	internal.Logger.Info("Listening on :65432")
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go listenTCP(&wg)
+	go listenUDP(&wg)
+	wg.Wait() // 等待下面两个函数协程执行完毕
+}
+
+func listenTCP(wg *sync.WaitGroup) {
+	defer wg.Done()
+	ln, err := net.ListenTCP("tcp", &net.TCPAddr{
+		IP:   net.IPv4(0, 0, 0, 0),
+		Port: 65432,
+	})
 	if err != nil {
-		log.Fatal(err)
+		internal.Logger.Fatal("Listen error", zap.Error(err))
 		panic(err)
 	}
+	internal.Logger.Info("Listening messaging service on 0.0.0.0:65432/tcp")
+
 	for {
 		conn, err := ln.Accept()
 		if err != nil {
@@ -23,3 +36,16 @@ func main() {
 		go internal.TCPHandler(conn)
 	}
 }
+
+func listenUDP(wg *sync.WaitGroup) {
+	defer wg.Done()
+	ln, err := net.ListenUDP("udp", &net.UDPAddr{
+		IP:   net.IPv4(0, 0, 0, 0),
+		Port: 65432,
+	})
+	if err != nil {
+		internal.Logger.Fatal("Listen error", zap.Error(err))
+	}
+	internal.Logger.Info("Listening probe on 0.0.0.0:65432/udp")
+	internal.UDPHandler(ln)
+}