utils/network/udp_server.go

package network

import (
	"fmt"
	"net"
	"time"
)

const (
	udpServerReceiveBufferSize = 1024
)

type UDPServerOption func(opt *UDPServerOptions)

func WithReceiveBufferSize(receiveBufferSize int) UDPServerOption {
	return func(opt *UDPServerOptions) {
		opt.ReceiveBufferSize = receiveBufferSize
	}
}

func WithWriteTimeout(timeout time.Duration) UDPServerOption {
	return func(opt *UDPServerOptions) {
		opt.WriteTimeout = timeout
	}
}

func WithReadTimeout(timeout time.Duration) UDPServerOption {
	return func(opt *UDPServerOptions) {
		opt.ReadTimeout = timeout
	}
}

type UDPServerOptions struct {
	ReceiveBufferSize int
	WriteTimeout      time.Duration
	ReadTimeout       time.Duration
}

func NewUDPServerOptions(opts ...UDPServerOption) *UDPServerOptions {
	options := new(UDPServerOptions)

	for _, opt := range opts {
		opt(options)
	}

	if options.ReceiveBufferSize == 0 {
		options.ReceiveBufferSize = udpServerReceiveBufferSize
	}

	return options
}

type UDPServer struct {
	options             *UDPServerOptions
	conn                *net.UDPConn
	doneChan            chan any
	dealRequestChan     chan *remoteData
	dealRequestDoneChan chan any
}

type remoteData struct {
	data  []byte
	rAddr *net.UDPAddr
}

// Connect 建立连接
func (server *UDPServer) Connect(address string, options *UDPServerOptions) error {
	addr, err := net.ResolveUDPAddr("udp", address)
	if err != nil {
		return err
	}

	// 监听端口
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		return err
	}

	server.options = options
	server.conn = conn
	server.doneChan = make(chan any)

	// 启动读取请求协程
	go server.readRequest()

	return nil
}

// Disconnect 断开连接
func (server *UDPServer) Disconnect() {
	server.doneChan <- nil
	close(server.doneChan)
	server.doneChan = nil

	CloseConnection(server.conn)
	server.conn = nil
}

func (server *UDPServer) readRequest() {
	server.dealRequestChan = make(chan *remoteData)
	server.dealRequestDoneChan = make(chan any)

	go server.dealRequestAndResponse()

	for {
		select {
		case <-server.doneChan:
			server.dealRequestDoneChan <- nil
			close(server.dealRequestDoneChan)
			server.dealRequestDoneChan = nil

			close(server.dealRequestChan)
			server.dealRequestChan = nil

			return
		default:
			// 读取任意客户端发来的请求
			data, rAddr, err := readUDP(server.conn, server.options.ReceiveBufferSize, WithReadDeadline(server.options.ReadTimeout))
			if err != nil {
				fmt.Println(err)
				continue
			}

			server.dealRequestChan <- &remoteData{
				data:  data,
				rAddr: rAddr,
			}
		}
	}
}

func (server *UDPServer) dealRequestAndResponse() {
	// 回调上层
}