Files
inp2025/utils/udp.go

86 lines
1.5 KiB
Go

package utils
import (
"fmt"
"net"
"github.com/vmihailenco/msgpack/v5"
)
type UDPPayload struct {
MagicNumber int `msgpack:"magicNumber"`
Data string `msgpack:"data"`
}
const (
BUFFER_SIZE int = 1024
MAGIC_NUMBER int = 114514
)
func ListenUDP(
port string,
doneCh chan struct{},
) (string, error) {
addr, err := net.ResolveUDPAddr("udp", ":"+port)
if err != nil {
return "", fmt.Errorf("failed to resolve address, %w", err)
}
conn, err := net.ListenUDP("udp", addr)
if err != nil {
return "", fmt.Errorf("failed to listen on addr '%v', %w", addr, err)
}
defer conn.Close()
dataCh := make(chan 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 {
dataCh <- payload.Data
return
}
}
}()
select {
case <-doneCh:
return "", nil
case data := <-dataCh:
return data, nil
}
return "", fmt.Errorf("failed to receive data")
}
func SendPayload(endpoint, data string) 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(UDPPayload{
MagicNumber: MAGIC_NUMBER,
Data: data,
})
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
}