udp-server.go

package main

import (
	"context"
	"fmt"
	"log"
	"net"
	"os"

	"golang.org/x/sync/errgroup"
)

func main() {
	log.SetFlags(0)

	err := run(context.Background())
	if err != nil {
		log.Fatalln(err)
	}
}

func run(ctx context.Context) error {
	port, ok := os.LookupEnv("PORT")
	if !ok {
		return fmt.Errorf("PORT is required")
	}

	srcAddr, err := net.ResolveUDPAddr("udp", ":"+port)
	if err != nil {
		return fmt.Errorf("error parsing source address: %w", err)
	}

	log.Println("listen", "on", srcAddr.String())
	conn, err := net.ListenUDP("udp", srcAddr)
	if err != nil {
		return fmt.Errorf("failed to start udp connection: %w", err)
	}

	type Message struct {
		Addr *net.UDPAddr
		Data []byte
	}

	msgs := make(chan Message, 1)

	eg, ctx := errgroup.WithContext(ctx)
	eg.Go(func() error {
		for {
			data := make([]byte, 2<<15)
			n, dstAddr, err := conn.ReadFromUDP(data)
			if err != nil {
				return fmt.Errorf("error reading packet: %w", err)
			}

			log.Println("recv", string(data[:n]), "from", dstAddr.String())

			select {
			case <-ctx.Done():
				return context.Cause(ctx)
			case msgs <- Message{Addr: dstAddr, Data: data[:n]}:
			}
		}
	})
	eg.Go(func() error {
		for {
			var msg Message
			select {
			case <-ctx.Done():
				return context.Cause(ctx)
			case msg = <-msgs:
			}

			log.Println("send", string(msg.Data), "to", msg.Addr.String())

			_, err := conn.WriteToUDP(msg.Data, msg.Addr)
			if err != nil {
				return fmt.Errorf("error writing packet: %w", err)
			}
		}
	})
	return eg.Wait()
}