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 {
		if duration == 0 {
			return nil
		}

		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 {
		if duration == 0 {
			return nil
		}

		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
}