fix: 网元跟踪任务优化socket接收

This commit is contained in:
TsMask
2024-10-15 12:11:05 +08:00
parent 06a9a3986b
commit 318fc6cc20
3 changed files with 38 additions and 27 deletions

View File

@@ -52,16 +52,18 @@ func (s *SocketTCP) Close() {
}
// Resolve 处理消息
func (s *SocketTCP) Resolve(callback func(conn *net.Conn)) error {
func (s *SocketTCP) Resolve(callback func(conn *net.Conn, err error)) {
if s.Listener == nil {
return fmt.Errorf("tcp service not created")
callback(nil, fmt.Errorf("tcp service not created"))
return
}
listener := *s.Listener
for {
select {
case <-s.StopChan:
return fmt.Errorf("udp service stop")
callback(nil, fmt.Errorf("udp service stop"))
return
default:
conn, err := listener.Accept()
if err != nil {
@@ -69,14 +71,7 @@ func (s *SocketTCP) Resolve(callback func(conn *net.Conn)) error {
continue
}
defer conn.Close()
// 处理连接
callback(&conn)
// 发送响应
if _, err = conn.Write([]byte("tcp>")); err != nil {
fmt.Println("Error sending response:", err)
}
callback(&conn, nil)
}
}
}

View File

@@ -50,22 +50,18 @@ func (s *SocketUDP) Close() {
}
// Resolve 处理消息
func (s *SocketUDP) Resolve(callback func(*net.UDPConn)) error {
func (s *SocketUDP) Resolve(callback func(*net.UDPConn, error)) {
if s.Conn == nil {
return fmt.Errorf("udp service not created")
callback(nil, fmt.Errorf("udp service not created"))
return
}
for {
select {
case <-s.StopChan:
return fmt.Errorf("udp service stop")
callback(nil, fmt.Errorf("udp service not created"))
default:
callback(s.Conn)
// 发送响应
if _, err := s.Conn.WriteTo([]byte("udp>"), s.Conn.RemoteAddr()); err != nil {
fmt.Println("Error sending response:", err)
}
callback(s.Conn, nil)
}
}
}

View File

@@ -60,12 +60,17 @@ func (r *TraceTask) CreateUDP() error {
}
// 接收处理UDP数据
go r.udpService.Resolve(func(conn *net.UDPConn) {
go r.udpService.Resolve(func(conn *net.UDPConn, err error) {
if err != nil {
logger.Errorf("UDP Resolve %s", err.Error())
return
}
// 读取数据
buf := make([]byte, 2048)
n, err := conn.Read(buf)
n, addr, err := conn.ReadFromUDPAddrPort(buf)
if err != nil {
logger.Errorf("error reading from UDP connection: %s", err.Error())
logger.Errorf("UDP Resolve ReadFromUDPAddrPort Error: %s", err.Error())
return
}
@@ -73,7 +78,7 @@ func (r *TraceTask) CreateUDP() error {
// logger.Infof("socket UDP Base64: %s", base64.StdEncoding.EncodeToString(buf[:n]))
mData, err := UDPDataHandler(buf, n)
if err != nil {
logger.Errorf("udp resolve data fail: %s", err.Error())
logger.Errorf("UDP Resolve UDPDataHandler Error: %s", err.Error())
return
}
taskId := parse.Number(mData["taskId"])
@@ -98,6 +103,11 @@ func (r *TraceTask) CreateUDP() error {
logger.Infof("pcapFile: %s", v)
wsService.NewWSSend.ByGroupID(fmt.Sprintf("%s%d", wsService.GROUP_TRACE_NE, taskId), taskId)
}
// 发送响应
if _, err := conn.WriteToUDPAddrPort([]byte("udp>"), addr); err != nil {
logger.Errorf("UDP Resolve WriteToUDPAddrPort Error: %s", err.Error())
}
})
// ============ 测试接收网元UDP发过来的数据
@@ -107,13 +117,18 @@ func (r *TraceTask) CreateUDP() error {
return err
}
// 接收处理TCP数据
go tcpService.Resolve(func(conn *net.Conn) {
go tcpService.Resolve(func(conn *net.Conn, err error) {
if err != nil {
logger.Errorf("TCP Resolve %s", err.Error())
return
}
c := (*conn)
// 读取数据
buf := make([]byte, 2048)
n, err := c.Read(buf)
if err != nil {
logger.Errorf("error reading from TCP connection: %s", err.Error())
logger.Errorf("TCP Resolve Read Error: %s", err.Error())
return
}
@@ -122,7 +137,7 @@ func (r *TraceTask) CreateUDP() error {
logger.Infof("socket TCP Base64: %s", deData)
mData, err := UDPDataHandler(deData, len(deData))
if err != nil {
logger.Errorf("tcp resolve data fail: %s", err.Error())
logger.Errorf("TCP Resolve UDPDataHandler Error: %s", err.Error())
return
}
taskId := parse.Number(mData["taskId"])
@@ -147,6 +162,11 @@ func (r *TraceTask) CreateUDP() error {
logger.Infof("pcapFile: %s", v)
wsService.NewWSSend.ByGroupID(fmt.Sprintf("%s%d", wsService.GROUP_TRACE_NE, taskId), taskId)
}
// 发送响应
if _, err = c.Write([]byte("tcp>")); err != nil {
logger.Errorf("TCP Resolve Write Error: %s", err.Error())
}
})
return nil
}