Fix: fuck udp

utils/udp.go

@@ -1,16 +1,22 @@
 package utils
 
 import (
+	"errors"
 	"fmt"
 	"net"
+	"time"
 
 	"gitea.konchin.com/ytshih/inp2025/types"
+	"github.com/spf13/viper"
 	"github.com/vmihailenco/msgpack/v5"
+	"go.uber.org/zap"
 )
 
 const (
 	BUFFER_SIZE  int = 1024
 	MAGIC_NUMBER int = 114514
+
+	LISTEN_TIMEOUT = 200 * time.Millisecond
 )
 
 type UDPReqType int
@@ -36,24 +42,36 @@ func ListenUDPData(
 	return ListenUDP(port, dataCh, nil)
 }
 
-func Ping(endpoint string) error {
+func Ping(endpoints []string) ([]string, error) {
 	pingCh := make(chan string)
 	local, shutdown, err := ListenUDP(0, nil, pingCh)
 	if err != nil {
-		return err
+		return []string{}, err
 	}
 	defer shutdown()
 
-	err = SendRawPayload(endpoint, UDPPayload{
-		MagicNumber: MAGIC_NUMBER,
-		Endpoint:    local,
-		Type:        UDPReqTypePingRequest,
-	})
-	if err != nil {
-		return err
+	for _, endpoint := range endpoints {
+		SendRawPayload(endpoint, UDPPayload{
+			MagicNumber: MAGIC_NUMBER,
+			Endpoint:    local,
+			Type:        UDPReqTypePingRequest,
+		})
+	}
+
+	doneCh := make(chan struct{})
+	go func() {
+		time.Sleep(LISTEN_TIMEOUT)
+		doneCh <- struct{}{}
+	}()
+	ret := []string{}
+	for {
+		select {
+		case <-doneCh:
+			return ret, nil
+		case endpoint := <-pingCh:
+			ret = append(ret, endpoint)
+		}
 	}
-	<-pingCh
-	return nil
 }
 
 func ListenUDP(
@@ -61,11 +79,12 @@ func ListenUDP(
 	dataCh chan string,
 	pingCh chan string,
 ) (string, types.ShutdownFunc, error) {
-	conn, err := net.ListenUDP("udp", &net.UDPAddr{Port: port})
+	conn, err := net.ListenUDP("udp4", &net.UDPAddr{Port: port})
 	if err != nil {
 		return "", nil, fmt.Errorf("failed to listen udp, %w", err)
 	}
-	local := conn.LocalAddr().String()
+	addr, _ := net.ResolveUDPAddr("udp4", conn.LocalAddr().String())
+	local := fmt.Sprintf("%s:%d", viper.GetString("host"), addr.Port)
 
 	go func() {
 		for {
@@ -73,6 +92,11 @@ func ListenUDP(
 
 			n, _, err := conn.ReadFromUDP(buffer)
 			if err != nil {
+				zap.L().Error("fuck udp",
+					zap.Error(err))
+				if errors.Is(err, net.ErrClosed) {
+					return
+				}
 				continue
 			}
 
@@ -123,7 +147,7 @@ func SendRawPayload(
 	endpoint string,
 	payload UDPPayload,
 ) error {
-	conn, err := net.Dial("udp", endpoint)
+	conn, err := net.Dial("udp4", endpoint)
 	if err != nil {
 		return fmt.Errorf("failed to dial endpoint, %w", err)
 	}