utils/network/connection.go

package network

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

// ConnectionReadOptions 从连接读取数据的选项
type ConnectionReadOptions func(conn net.Conn) error

// WithReadDeadline 读超时选项
func WithReadDeadline(duration time.Duration) ConnectionReadOptions {
	return func(conn net.Conn) error {
		err := conn.SetReadDeadline(time.Now().Add(duration))
		if err != nil {
			return err
		}

		return nil
	}
}

// ConnectionWriteOptions 从连接写入数据的选项
type ConnectionWriteOptions func(conn net.Conn) error

// WithWriteDeadline 写超时选项
func WithWriteDeadline(duration time.Duration) ConnectionWriteOptions {
	return func(conn net.Conn) error {
		err := conn.SetWriteDeadline(time.Now().Add(duration))
		if err != nil {
			return err
		}

		return nil
	}
}

// ReadUDP 读取UDP包
func ReadUDP(conn *net.UDPConn, bufferSize int, opts ...ConnectionReadOptions) ([]byte, *net.UDPAddr, error) {
	buffer := make([]byte, bufferSize)

	for _, opt := range opts {
		err := opt(conn)
		if err != nil {
			return nil, nil, err
		}
	}

	n, rAddr, err := conn.ReadFromUDP(buffer)
	if err != nil {
		return nil, nil, err
	}

	return buffer[:n], rAddr, nil
}

// WriteUDPWithRemoteAddr 指定远端(对端)地址写出UDP包
func WriteUDPWithRemoteAddr(conn *net.UDPConn, rAddr *net.UDPAddr, data []byte, opts ...ConnectionWriteOptions) error {
	for _, opt := range opts {
		err := opt(conn)
		if err != nil {
			return err
		}
	}

	_, err := conn.WriteToUDP(data, rAddr)
	if err != nil {
		return err
	}

	return nil
}

// WriteUDP 写出UDP包
func WriteUDP(conn *net.UDPConn, data []byte, opts ...ConnectionWriteOptions) error {
	for _, opt := range opts {
		err := opt(conn)
		if err != nil {
			return err
		}
	}

	_, err := conn.Write(data)
	if err != nil {
		return err
	}

	return nil
}

// ReadTCP 读取TCP数据
func ReadTCP(conn net.Conn, bufferSize int, readCallback func(data []byte) (bool, error), opts ...ConnectionReadOptions) error {
	for {
		buffer := make([]byte, bufferSize)

		for _, opt := range opts {
			err := opt(conn)
			if err != nil {
				return err
			}
		}

		n, err := conn.Read(buffer)
		if err != nil {
			return err
		}

		if n != 0 {
			readOver, err := readCallback(buffer[:n])
			if err != nil {
				return err
			}

			if readOver {
				return nil
			}
		}
	}
}

// WriteTCP 写TCP数据
func WriteTCP(conn net.Conn, data []byte, opts ...ConnectionWriteOptions) error {
	writeBytesCount := 0

	for {
		for _, opt := range opts {
			err := opt(conn)
			if err != nil {
				return err
			}
		}

		n, err := conn.Write(data[writeBytesCount:])
		if err != nil {
			return err
		}

		writeBytesCount += n
		if writeBytesCount < len(data) {
			continue
		}

		return nil
	}
}

// CloseConnection 关闭连接
func CloseConnection(conn net.Conn) {
	err := conn.Close()
	if err != nil {
		fmt.Println("Close Connection Error:", err.Error())
		return
	}

	conn = nil
}