diff --git a/src/framework/socket/tcp_server.go b/src/framework/socket/tcp_server.go index 857af870..0957438f 100644 --- a/src/framework/socket/tcp_server.go +++ b/src/framework/socket/tcp_server.go @@ -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) } } } diff --git a/src/framework/socket/udp_server.go b/src/framework/socket/udp_server.go index 33a45a60..d93001b0 100644 --- a/src/framework/socket/udp_server.go +++ b/src/framework/socket/udp_server.go @@ -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) } } } diff --git a/src/modules/trace/service/trace_task.go b/src/modules/trace/service/trace_task.go index 98e8d7c0..15f09976 100644 --- a/src/modules/trace/service/trace_task.go +++ b/src/modules/trace/service/trace_task.go @@ -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 }