utils/network/udp_server.go

package network

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

const (
	udpServerReceiveBufferSize = 1024
)

// UDPServerRequestCallback 请求回调
// 参数：
// data: 请求数据
// 返回值:
// send: 是否发送响应
// responseBytes: 响应数据包
// err: 是否存在错误
type UDPServerRequestCallback func(data []byte) (send bool, responseBytes []byte, err error)

type UDPServerOption func(opt *UDPServerOptions)

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

func WithUDPServerWriteTimeout(timeout time.Duration) UDPServerOption {
	return func(opt *UDPServerOptions) {
		opt.writeTimeout = timeout
	}
}

func WithUDPServerReadTimeout(timeout time.Duration) UDPServerOption {
	return func(opt *UDPServerOptions) {
		opt.readTimeout = timeout
	}
}

func WithUDPServerRequestCallback(requestCallback UDPServerRequestCallback) UDPServerOption {
	return func(opt *UDPServerOptions) {
		opt.requestCallback = requestCallback
	}
}

type UDPServerOptions struct {
	// 默认1024字节，一般保证足够收取一个数据包
	receiveBufferSize int

	// 写超时，不设置就是阻塞写
	writeTimeout time.Duration

	// 读超时，不设置就是阻塞读
	readTimeout time.Duration

	// 请求数据回调
	requestCallback UDPServerRequestCallback
}

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
}

// 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)

	// 启动读取请求协程
	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() {
	dealRequestDoneChannels := make([]chan any, 0)

	for {
		select {
		case <-server.doneChan:
			for _, dealRequestDoneChan := range dealRequestDoneChannels {
				dealRequestDoneChan <- nil
				close(dealRequestDoneChan)
			}

			return
		default:
			// 读取任意客户端发来的请求，超时就是没有客户端发出请求
			data, rAddr, err := readUDP(server.conn, server.options.receiveBufferSize, WithReadDeadline(server.options.readTimeout))
			if err != nil {
				fmt.Println(err)
				continue
			}

			// 接收到请求
			dealRequestDoneChan := make(chan any)
			dealRequestDoneChannels = append(dealRequestDoneChannels, dealRequestDoneChan)
			go server.dealRequest(data, rAddr, dealRequestDoneChan)
		}
	}
}

func (server *UDPServer) dealRequest(data []byte, rAddr *net.UDPAddr, doneChan chan any) {
	for {
		select {
		case <-doneChan:
			return
		default:
			if server.options.requestCallback == nil {
				server.response(server.conn, rAddr, data)
				return
			}

			send, responseBytes, err := server.options.requestCallback(data)
			if !send {
				return
			}

			if err != nil {
				server.response(server.conn, rAddr, []byte(err.Error()))
			}

			server.response(server.conn, rAddr, responseBytes)
		}
	}
}

func (server *UDPServer) response(conn *net.UDPConn, rAddr *net.UDPAddr, data []byte) {
	err := writeUDPWithRemoteAddr(conn, rAddr, data, WithWriteDeadline(server.options.writeTimeout))
	if err != nil {
		fmt.Println("Response Error:", err)
		return
	}
}