From e87a10c08fbef60a761ae431966cd2f840746791 Mon Sep 17 00:00:00 2001
From: fengyarnom <fengyarnom@gmail.com>
Date: Thu, 15 May 2025 12:17:41 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=96=B0=E5=A2=9E=20UDP=20=E8=A7=A3?=
 =?UTF-8?q?=E6=9E=90?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 main.go       | 65 +++++++++++++++++++++++++++++-------------
 udp_server.go | 78 +++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 123 insertions(+), 20 deletions(-)
 create mode 100644 udp_server.go

diff --git a/main.go b/main.go
index 4f91d3d..97528f1 100644
--- a/main.go
+++ b/main.go
@@ -1,51 +1,76 @@
 package main
 
 import (
+	"flag"
 	"fmt"
 	"log"
 	"sync"
 )
 
 func main() {
-    // 初始化日志
-    if err := InitLogger(); err != nil {
-        log.Fatalf("初始化日志系统失败: %v", err)
-    }
-    defer CloseLogger() 
-    // 启动日志轮转
-     StartLogRotation()
+	// 命令行参数解析
+	serverType := flag.String("server", "tcp", "服务器类型: tcp 或 udp")
+	tcpPort := flag.String("tcpport", "10002", "TCP服务器端口")
+	udpPort := flag.String("udpport", "10002", "UDP服务器端口")
+	httpPort := flag.String("httpport", "10001", "HTTP服务器端口")
+	flag.Parse()
+
+	// 初始化日志
+	if err := InitLogger(); err != nil {
+		log.Fatalf("初始化日志系统失败: %v", err)
+	}
+	defer CloseLogger()
+	// 启动日志轮转
+	StartLogRotation()
+
 	// 初始化数据库
 	err := InitDB()
 	if err != nil {
 		log.Fatalf("初始化数据库失败: %v", err)
 	}
 	defer CloseDB()
-	
+
 	fmt.Println("iProbe 斜侧仪监控系统启动中...")
-	
+
 	var wg sync.WaitGroup
-	
-	// 启动TCP服务器
+
+	// 根据选择启动TCP或UDP服务器
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		if err := StartTCPServer(":10002"); err != nil {
-			log.Fatalf("TCP服务器启动失败: %v", err)
+		if *serverType == "tcp" {
+			Logger.Printf("使用TCP服务器模式")
+			if err := StartTCPServer(":" + *tcpPort); err != nil {
+				log.Fatalf("TCP服务器启动失败: %v", err)
+			}
+		} else if *serverType == "udp" {
+			Logger.Printf("使用UDP服务器模式")
+			if err := StartUDPServer(":" + *udpPort); err != nil {
+				log.Fatalf("UDP服务器启动失败: %v", err)
+			}
+		} else {
+			log.Fatalf("未知的服务器类型: %s", *serverType)
 		}
 	}()
-	
+
 	// 启动HTTP服务器
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		if err := StartHTTPServer(":10001"); err != nil {
+		if err := StartHTTPServer(":" + *httpPort); err != nil {
 			log.Fatalf("HTTP服务器启动失败: %v", err)
 		}
 	}()
-	
-	fmt.Println("服务器已启动成功")
-	fmt.Println("- HTTP接口: http://localhost:10001")
-	fmt.Println("- TCP接口: localhost:10002")
-	
+
+	if *serverType == "tcp" {
+		fmt.Println("服务器已启动成功")
+		fmt.Printf("- HTTP接口: http://localhost:%s\n", *httpPort)
+		fmt.Printf("- TCP接口: localhost:%s\n", *tcpPort)
+	} else {
+		fmt.Println("服务器已启动成功")
+		fmt.Printf("- HTTP接口: http://localhost:%s\n", *httpPort)
+		fmt.Printf("- UDP接口: localhost:%s\n", *udpPort)
+	}
+
 	wg.Wait()
 }
diff --git a/udp_server.go b/udp_server.go
new file mode 100644
index 0000000..b7dc6ac
--- /dev/null
+++ b/udp_server.go
@@ -0,0 +1,78 @@
+// udp_server.go
+package main
+
+import (
+	"fmt"
+	"net"
+)
+
+// StartUDPServer 启动UDP服务器
+func StartUDPServer(address string) error {
+	// 解析UDP地址
+	addr, err := net.ResolveUDPAddr("udp", address)
+	if err != nil {
+		return fmt.Errorf("解析UDP地址失败: %v", err)
+	}
+
+	// 创建UDP连接
+	conn, err := net.ListenUDP("udp", addr)
+	if err != nil {
+		return fmt.Errorf("监听UDP地址失败: %v", err)
+	}
+	defer conn.Close()
+
+	// 启动客户端清理
+	startClientCleanup()
+
+	Logger.Printf("UDP服务器已启动，正在监听 %s\n", address)
+
+	buffer := make([]byte, 1024)
+
+	for {
+		n, remoteAddr, err := conn.ReadFromUDP(buffer)
+		if err != nil {
+			Logger.Printf("从UDP客户端读取失败: %v", err)
+			continue
+		}
+
+		// 处理数据包
+		go handleUDPPacket(conn, remoteAddr, buffer[:n])
+	}
+}
+
+// handleUDPPacket 处理UDP数据包
+func handleUDPPacket(conn *net.UDPConn, addr *net.UDPAddr, data []byte) {
+	// 获取客户端信息
+	remoteAddr := addr.String()
+
+	// 添加到在线客户端列表
+	addClient(remoteAddr)
+
+	// 将字节数据转换为字符串，并记录原始数据
+	rawData := string(data)
+	TCPDataLogger.Printf("从UDP客户端 %s 接收到原始数据: %s", remoteAddr, rawData)
+
+	// 尝试解析数据
+	sensorID, x, y, z, err := parseData(rawData)
+
+	if err == nil {
+		TCPDataLogger.Printf("解析成功 - UDP客户端: %s, 传感器ID: %d, 值: X=%.3f, Y=%.3f, Z=%.3f",
+			remoteAddr, sensorID, x, y, z)
+
+		// 保存数据到数据库
+		if err := SaveSensorData(sensorID, x, y, z); err != nil {
+			Logger.Printf("保存传感器数据失败: %v", err)
+		}
+	} else {
+		TCPDataLogger.Printf("无法解析从UDP客户端 %s 接收到的数据: %s, 错误: %v", remoteAddr, rawData, err)
+	}
+
+	// 发送响应
+	resp := "OK\n"
+	if _, err := conn.WriteToUDP([]byte(resp), addr); err != nil {
+		Logger.Printf("发送响应到UDP客户端 %s 失败: %v", remoteAddr, err)
+	}
+
+	// 更新客户端最后活跃时间
+	updateClientLastSeen(remoteAddr)
+}