feat: 添加socket工具
This commit is contained in:
TsMask
96
src/framework/socket/tcp_client.go
Normal file
96
src/framework/socket/tcp_client.go
Normal file
@@ -0,0 +1,96 @@
|
||||
package socket
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"net"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ConnTCP 连接TCP客户端
|
||||
type ConnTCP struct {
|
||||
Addr string `json:"addr"` // 主机地址
|
||||
Port int64 `json:"port"` // 端口
|
||||
|
||||
DialTimeOut time.Duration `json:"dialTimeOut"` // 连接超时断开
|
||||
|
||||
Client *net.Conn `json:"client"`
|
||||
LastResult string `json:"lastResult"` // 记最后一次发送消息的结果
|
||||
}
|
||||
|
||||
// New 创建TCP客户端
|
||||
func (c *ConnTCP) New() (*ConnTCP, error) {
|
||||
// IPV6地址协议
|
||||
proto := "tcp"
|
||||
if strings.Contains(c.Addr, ":") {
|
||||
proto = "tcp6"
|
||||
c.Addr = fmt.Sprintf("[%s]", c.Addr)
|
||||
}
|
||||
address := fmt.Sprintf("%s:%d", c.Addr, c.Port)
|
||||
|
||||
// 默认等待5s
|
||||
if c.DialTimeOut == 0 {
|
||||
c.DialTimeOut = 5 * time.Second
|
||||
}
|
||||
|
||||
// 连接到服务端
|
||||
client, err := net.DialTimeout(proto, address, c.DialTimeOut)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
c.Client = &client
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// Close 关闭当前TCP客户端
|
||||
func (c *ConnTCP) Close() {
|
||||
if c.Client != nil {
|
||||
(*c.Client).Close()
|
||||
}
|
||||
}
|
||||
|
||||
// Send 发送消息
|
||||
func (c *ConnTCP) Send(msg []byte, timer time.Duration) (string, error) {
|
||||
if c.Client == nil {
|
||||
return "", fmt.Errorf("tcp client not connected")
|
||||
}
|
||||
conn := *c.Client
|
||||
|
||||
// 写入信息
|
||||
if len(msg) > 0 {
|
||||
if _, err := conn.Write(msg); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
||||
var buf bytes.Buffer
|
||||
defer buf.Reset()
|
||||
|
||||
tmp := make([]byte, 1024)
|
||||
for {
|
||||
select {
|
||||
case <-time.After(timer):
|
||||
c.LastResult = buf.String()
|
||||
return c.LastResult, fmt.Errorf("timeout")
|
||||
default:
|
||||
// 读取命令消息
|
||||
n, err := conn.Read(tmp)
|
||||
if n == 0 || err != nil {
|
||||
tmp = nil
|
||||
break
|
||||
}
|
||||
|
||||
tmpStr := string(tmp[:n])
|
||||
buf.WriteString(tmpStr)
|
||||
|
||||
// 是否有终止符
|
||||
if strings.HasSuffix(tmpStr, ">") || strings.HasSuffix(tmpStr, "> ") || strings.HasSuffix(tmpStr, "# ") {
|
||||
tmp = nil
|
||||
c.LastResult = buf.String()
|
||||
return c.LastResult, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
83
src/framework/socket/tcp_server.go
Normal file
83
src/framework/socket/tcp_server.go
Normal file
@@ -0,0 +1,83 @@
|
||||
package socket
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"strings"
|
||||
|
||||
"be.ems/src/framework/logger"
|
||||
)
|
||||
|
||||
// SocketTCP TCP服务端
|
||||
type SocketTCP struct {
|
||||
Addr string `json:"addr"` // 主机地址
|
||||
Port int64 `json:"port"` // 端口
|
||||
Listen *net.Listener `json:"listen"`
|
||||
StopChan chan struct{} `json:"stop"` // 停止信号
|
||||
}
|
||||
|
||||
// New 创建TCP服务端
|
||||
func (s *SocketTCP) New() (*SocketTCP, error) {
|
||||
// IPV6地址协议
|
||||
proto := "tcp"
|
||||
if strings.Contains(s.Addr, ":") {
|
||||
proto = "tcp6"
|
||||
s.Addr = fmt.Sprintf("[%s]", s.Addr)
|
||||
}
|
||||
address := fmt.Sprintf("%s:%d", s.Addr, s.Port)
|
||||
|
||||
ln, err := net.Listen(proto, address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
s.Listen = &ln
|
||||
s.StopChan = make(chan struct{}, 1)
|
||||
return s, nil
|
||||
}
|
||||
|
||||
// Close 关闭当前TCP服务端
|
||||
func (s *SocketTCP) Close() {
|
||||
if s.Listen != nil {
|
||||
s.StopChan <- struct{}{}
|
||||
(*s.Listen).Close()
|
||||
}
|
||||
}
|
||||
|
||||
// Resolve 处理消息
|
||||
func (s *SocketTCP) Resolve(bufferSize int, callback func([]byte, int)) error {
|
||||
if s.Listen == nil {
|
||||
return fmt.Errorf("tcp service not created")
|
||||
}
|
||||
|
||||
ln := *s.Listen
|
||||
buffer := make([]byte, bufferSize)
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-s.StopChan:
|
||||
return fmt.Errorf("udp service stop")
|
||||
default:
|
||||
conn, err := ln.Accept()
|
||||
if err != nil {
|
||||
logger.Errorf("Error accepting connection: %v ", err)
|
||||
continue
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
// 读取数据
|
||||
n, err := conn.Read(buffer)
|
||||
if err != nil {
|
||||
fmt.Println("Error reading from TCP connection:", err)
|
||||
continue
|
||||
}
|
||||
|
||||
callback(buffer, n)
|
||||
|
||||
// 发送响应
|
||||
if _, err = conn.Write([]byte("tcp>")); err != nil {
|
||||
fmt.Println("Error sending response:", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
96
src/framework/socket/udp_client.go
Normal file
96
src/framework/socket/udp_client.go
Normal file
@@ -0,0 +1,96 @@
|
||||
package socket
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"net"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ConnUDP 连接UDP客户端
|
||||
type ConnUDP struct {
|
||||
Addr string `json:"addr"` // 主机地址
|
||||
Port int64 `json:"port"` // 端口
|
||||
|
||||
DialTimeOut time.Duration `json:"dialTimeOut"` // 连接超时断开
|
||||
|
||||
Client *net.Conn `json:"client"`
|
||||
LastResult string `json:"lastResult"` // 记最后一次发送消息的结果
|
||||
}
|
||||
|
||||
// New 创建UDP客户端
|
||||
func (c *ConnUDP) New() (*ConnUDP, error) {
|
||||
// IPV6地址协议
|
||||
proto := "udp"
|
||||
if strings.Contains(c.Addr, ":") {
|
||||
proto = "udp6"
|
||||
c.Addr = fmt.Sprintf("[%s]", c.Addr)
|
||||
}
|
||||
address := fmt.Sprintf("%s:%d", c.Addr, c.Port)
|
||||
|
||||
// 默认等待5s
|
||||
if c.DialTimeOut == 0 {
|
||||
c.DialTimeOut = 5 * time.Second
|
||||
}
|
||||
|
||||
// 连接到服务端
|
||||
client, err := net.DialTimeout(proto, address, c.DialTimeOut)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
c.Client = &client
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// Close 关闭当前UDP客户端
|
||||
func (c *ConnUDP) Close() {
|
||||
if c.Client != nil {
|
||||
(*c.Client).Close()
|
||||
}
|
||||
}
|
||||
|
||||
// Send 发送消息
|
||||
func (c *ConnUDP) Send(msg []byte, ms int) (string, error) {
|
||||
if c.Client == nil {
|
||||
return "", fmt.Errorf("udp client not connected")
|
||||
}
|
||||
conn := *c.Client
|
||||
|
||||
// 写入信息
|
||||
if len(msg) > 0 {
|
||||
if _, err := conn.Write(msg); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
||||
var buf bytes.Buffer
|
||||
defer buf.Reset()
|
||||
|
||||
tmp := make([]byte, 1024)
|
||||
for {
|
||||
select {
|
||||
case <-time.After(time.Duration(time.Duration(ms).Milliseconds())):
|
||||
c.LastResult = buf.String()
|
||||
return c.LastResult, fmt.Errorf("timeout")
|
||||
default:
|
||||
// 读取命令消息
|
||||
n, err := conn.Read(tmp)
|
||||
if n == 0 || err != nil {
|
||||
tmp = nil
|
||||
break
|
||||
}
|
||||
|
||||
tmpStr := string(tmp[:n])
|
||||
buf.WriteString(tmpStr)
|
||||
|
||||
// 是否有终止符
|
||||
if strings.HasSuffix(tmpStr, ">") || strings.HasSuffix(tmpStr, "> ") || strings.HasSuffix(tmpStr, "# ") {
|
||||
tmp = nil
|
||||
c.LastResult = buf.String()
|
||||
return c.LastResult, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
80
src/framework/socket/udp_server.go
Normal file
80
src/framework/socket/udp_server.go
Normal file
@@ -0,0 +1,80 @@
|
||||
package socket
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// SocketUDP UDP服务端
|
||||
type SocketUDP struct {
|
||||
Addr string `json:"addr"` // 主机地址
|
||||
Port int64 `json:"port"` // 端口
|
||||
Conn *net.UDPConn `json:"conn"`
|
||||
StopChan chan struct{} `json:"stop"` // 停止信号
|
||||
}
|
||||
|
||||
// New 创建UDP服务端
|
||||
func (s *SocketUDP) New() (*SocketUDP, error) {
|
||||
// IPV6地址协议
|
||||
proto := "udp"
|
||||
if strings.Contains(s.Addr, ":") {
|
||||
proto = "udp6"
|
||||
s.Addr = fmt.Sprintf("[%s]", s.Addr)
|
||||
}
|
||||
address := fmt.Sprintf("%s:%d", s.Addr, s.Port)
|
||||
|
||||
// 解析 UDP 地址
|
||||
udpAddr, err := net.ResolveUDPAddr(proto, address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 监听 UDP 地址
|
||||
conn, err := net.ListenUDP("udp", udpAddr)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
s.Conn = conn
|
||||
s.StopChan = make(chan struct{}, 1)
|
||||
return s, nil
|
||||
}
|
||||
|
||||
// CloseService 关闭当前UDP服务端
|
||||
func (s *SocketUDP) Close() {
|
||||
if s.Conn != nil {
|
||||
s.StopChan <- struct{}{}
|
||||
(*s.Conn).Close()
|
||||
}
|
||||
}
|
||||
|
||||
// Resolve 处理消息
|
||||
func (s *SocketUDP) Resolve(bufferSize int, callback func([]byte, int)) error {
|
||||
if s.Conn == nil {
|
||||
return fmt.Errorf("udp service not created")
|
||||
}
|
||||
|
||||
buffer := make([]byte, bufferSize)
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-s.StopChan:
|
||||
return fmt.Errorf("udp service stop")
|
||||
default:
|
||||
// 读取数据
|
||||
n, addr, err := s.Conn.ReadFromUDP(buffer)
|
||||
if err != nil {
|
||||
fmt.Println("Error reading from UDP connection:", err)
|
||||
continue
|
||||
}
|
||||
|
||||
callback(buffer, n)
|
||||
|
||||
// 发送响应
|
||||
if _, err = s.Conn.WriteToUDP([]byte("udp>"), addr); err != nil {
|
||||
fmt.Println("Error sending response:", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user