144 lines
2.7 KiB
Go
144 lines
2.7 KiB
Go
package utils
|
|
|
|
import (
|
|
"fmt"
|
|
"net"
|
|
|
|
"gitea.konchin.com/ytshih/inp2025/types"
|
|
"github.com/vmihailenco/msgpack/v5"
|
|
)
|
|
|
|
const (
|
|
BUFFER_SIZE int = 1024
|
|
MAGIC_NUMBER int = 114514
|
|
)
|
|
|
|
type UDPReqType int
|
|
|
|
const (
|
|
UDPReqTypeData UDPReqType = iota
|
|
UDPReqTypePingRequest
|
|
UDPReqTypePingReply
|
|
)
|
|
|
|
type UDPPayload struct {
|
|
MagicNumber int `msgpack:"magicNumber"`
|
|
Endpoint string `msgpack:"endpoint"`
|
|
Type UDPReqType `msgpack:"type"`
|
|
|
|
Data string `msgpack:"data"`
|
|
}
|
|
|
|
func ListenUDPData(
|
|
port int,
|
|
dataCh chan string,
|
|
) (string, types.ShutdownFunc, error) {
|
|
return ListenUDP(port, dataCh, nil)
|
|
}
|
|
|
|
func Ping(endpoint string) error {
|
|
pingCh := make(chan string)
|
|
local, shutdown, err := ListenUDP(0, nil, pingCh)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer shutdown()
|
|
|
|
err = SendRawPayload(endpoint, UDPPayload{
|
|
MagicNumber: MAGIC_NUMBER,
|
|
Endpoint: local,
|
|
Type: UDPReqTypePingRequest,
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
<-pingCh
|
|
return nil
|
|
}
|
|
|
|
func ListenUDP(
|
|
port int,
|
|
dataCh chan string,
|
|
pingCh chan string,
|
|
) (string, types.ShutdownFunc, error) {
|
|
conn, err := net.ListenUDP("udp", &net.UDPAddr{Port: port})
|
|
if err != nil {
|
|
return "", nil, fmt.Errorf("failed to listen udp, %w", err)
|
|
}
|
|
local := conn.LocalAddr().String()
|
|
|
|
go func() {
|
|
for {
|
|
buffer := make([]byte, BUFFER_SIZE)
|
|
|
|
n, _, err := conn.ReadFromUDP(buffer)
|
|
if err != nil {
|
|
continue
|
|
}
|
|
|
|
var payload UDPPayload
|
|
err = msgpack.Unmarshal(buffer[:n], &payload)
|
|
if err == nil && payload.MagicNumber == MAGIC_NUMBER {
|
|
switch payload.Type {
|
|
case UDPReqTypeData:
|
|
if dataCh != nil {
|
|
dataCh <- payload.Data
|
|
}
|
|
case UDPReqTypePingRequest:
|
|
SendRawPayload(payload.Endpoint, UDPPayload{
|
|
MagicNumber: MAGIC_NUMBER,
|
|
Endpoint: local,
|
|
Type: UDPReqTypePingReply,
|
|
})
|
|
case UDPReqTypePingReply:
|
|
if pingCh != nil {
|
|
pingCh <- payload.Endpoint
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
|
|
return local, func() { conn.Close() }, nil
|
|
}
|
|
|
|
func SendPayload(
|
|
local, remote string,
|
|
data any,
|
|
) error {
|
|
sdata, err := msgpack.Marshal(data)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to marshal data, %w", err)
|
|
}
|
|
|
|
return SendRawPayload(remote, UDPPayload{
|
|
MagicNumber: MAGIC_NUMBER,
|
|
Endpoint: local,
|
|
Type: UDPReqTypeData,
|
|
Data: string(sdata),
|
|
})
|
|
}
|
|
|
|
func SendRawPayload(
|
|
endpoint string,
|
|
payload UDPPayload,
|
|
) error {
|
|
conn, err := net.Dial("udp", endpoint)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to dial endpoint, %w", err)
|
|
}
|
|
defer conn.Close()
|
|
|
|
b, err := msgpack.Marshal(payload)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to marshal payload, %w", err)
|
|
}
|
|
|
|
_, err = conn.Write(b)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to send payload, %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|