From 37c10f1c07d6ad79bbcbd42caf1d802e0742a7ab Mon Sep 17 00:00:00 2001 From: NanamiAdmin Date: Sun, 5 Apr 2026 22:44:53 +0800 Subject: [PATCH] feat(watchdog): improve watchdog instance tracking - Add Close() function to watchdog for graceful shutdown - Update instance message format for better consistency - Add watchdog status tracking in StatusInfo - Include watchdog field in FrpcInstance struct and database - Change port fields from string to int in FrpcProxyInfo - Add auth_token support in InstanceInfo --- config.go | 6 ++++-- database.go | 24 +++++++++++++----------- frpAct.go | 27 +++++++++++++++++++++------ frpcProxyAct.go | 12 ++++++------ main.go | 19 +++++++++++++++---- watchdog/command.go | 18 ++++++++++++++++-- 6 files changed, 75 insertions(+), 31 deletions(-) diff --git a/config.go b/config.go index 0e96f40..ae107e1 100644 --- a/config.go +++ b/config.go @@ -27,6 +27,7 @@ type InstanceInfo struct { ServerAddr string `json:"serverAddr"` ServerPort string `json:"serverPort"` AuthMethod string `json:"auth_method"` + AuthToken string `json:"auth_token"` RunUser string `json:"runUser"` Additional map[string]interface{} `json:"additionalProperties"` } @@ -35,8 +36,8 @@ type FrpcProxyInfo struct { Name string `json:"name"` Type string `json:"type"` LocalIP string `json:"local_ip"` - LocalPort string `json:"local_port"` - RemotePort string `json:"remote_port"` + LocalPort int `json:"local_port"` + RemotePort int `json:"remote_port"` } type CreateInstanceRequest struct { @@ -238,6 +239,7 @@ func generateFrpcConfig(info InstanceInfo) string { config.Global["serverAddr"] = info.ServerAddr config.Global["serverPort"] = info.ServerPort config.Global["auth.method"] = info.AuthMethod + config.Global["auth.token"] = info.AuthToken for key, value := range info.Additional { config.Global[key] = value diff --git a/database.go b/database.go index c1a4212..7854b4f 100644 --- a/database.go +++ b/database.go @@ -36,6 +36,7 @@ type FrpcInstance struct { ConfigPath string CreatedAt time.Time CreatedBy string + Watchdog int } func InitDatabase(dbPath_data string, dbPath_log string) error { @@ -276,6 +277,7 @@ func InitFrpcDatabase(dbPath string) error { runUser TEXT NOT NULL DEFAULT 'root', configPath TEXT NOT NULL, createdAt TEXT NOT NULL, + watchdog INTEGER NOT NULL, UNIQUE(userID, name) ); ` @@ -359,7 +361,7 @@ func DBUpdateUser(userID int, username, passwd string) error { return nil } -func DBUpdateUserType (userID int, newType string) error { +func DBUpdateUserType(userID int, newType string) error { _, err := db.Exec("UPDATE userLogin SET type = ? WHERE userID = ?", newType, userID) if err != nil { return fmt.Errorf("failed to update user type: %w", err) @@ -368,8 +370,8 @@ func DBUpdateUserType (userID int, newType string) error { } func DBAddFrpcInstance(instance FrpcInstance) error { - _, err := frpcDB.Exec("INSERT INTO frpcInstances (userID, name, bootAtStart, runUser, configPath, createdAt) VALUES (?, ?, ?, ?, ?, ?)", - instance.UserID, instance.Name, instance.BootAtStart, instance.RunUser, instance.ConfigPath, time.Now().Format(time.RFC3339)) + _, err := frpcDB.Exec("INSERT INTO frpcInstances (userID, name, bootAtStart, runUser, configPath, createdAt, watchdog) VALUES (?, ?, ?, ?, ?, ?, ?)", + instance.UserID, instance.Name, instance.BootAtStart, instance.RunUser, instance.ConfigPath, time.Now().Format(time.RFC3339), instance.Watchdog) if err != nil { return fmt.Errorf("failed to insert frpc instance: %w", err) } @@ -379,8 +381,8 @@ func DBAddFrpcInstance(instance FrpcInstance) error { func DBQueryFrpcInstanceByID(instanceID int) (FrpcInstance, error) { var instance FrpcInstance var createdAtStr string - err := frpcDB.QueryRow("SELECT id, userID, name, bootAtStart, runUser, configPath, createdAt FROM frpcInstances WHERE id = ?", instanceID).Scan( - &instance.ID, &instance.UserID, &instance.Name, &instance.BootAtStart, &instance.RunUser, &instance.ConfigPath, &createdAtStr) + err := frpcDB.QueryRow("SELECT id, userID, name, bootAtStart, runUser, configPath, createdAt, watchdog FROM frpcInstances WHERE id = ?", instanceID).Scan( + &instance.ID, &instance.UserID, &instance.Name, &instance.BootAtStart, &instance.RunUser, &instance.ConfigPath, &createdAtStr, &instance.Watchdog) if err != nil { return instance, fmt.Errorf("failed to query frpc instance: %w", err) } @@ -391,8 +393,8 @@ func DBQueryFrpcInstanceByID(instanceID int) (FrpcInstance, error) { func DBQueryFrpcInstance(userID int, instanceName string) (FrpcInstance, error) { var instance FrpcInstance var createdAtStr string - err := frpcDB.QueryRow("SELECT id, userID, name, bootAtStart, runUser, configPath, createdAt FROM frpcInstances WHERE userID = ? AND name = ?", userID, instanceName).Scan( - &instance.ID, &instance.UserID, &instance.Name, &instance.BootAtStart, &instance.RunUser, &instance.ConfigPath, &createdAtStr) + err := frpcDB.QueryRow("SELECT id, userID, name, bootAtStart, runUser, configPath, createdAt, watchdog FROM frpcInstances WHERE userID = ? AND name = ?", userID, instanceName).Scan( + &instance.ID, &instance.UserID, &instance.Name, &instance.BootAtStart, &instance.RunUser, &instance.ConfigPath, &createdAtStr, &instance.Watchdog) if err != nil { return instance, fmt.Errorf("failed to query frpc instance: %w", err) } @@ -409,8 +411,8 @@ func DBRemoveFrpcInstanceByID(instanceID int) error { } func DBUpdateFrpcInstance(instance FrpcInstance) error { - _, err := frpcDB.Exec("UPDATE frpcInstances SET bootAtStart = ?, runUser = ?, configPath = ? WHERE id = ?", - instance.BootAtStart, instance.RunUser, instance.ConfigPath, instance.ID) + _, err := frpcDB.Exec("UPDATE frpcInstances SET bootAtStart = ?, runUser = ?, configPath = ?, watchdog = ? WHERE id = ?", + instance.BootAtStart, instance.RunUser, instance.ConfigPath, instance.Watchdog, instance.ID) if err != nil { return fmt.Errorf("failed to update frpc instance: %w", err) } @@ -419,7 +421,7 @@ func DBUpdateFrpcInstance(instance FrpcInstance) error { func DBListFrpcInstances() ([]FrpcInstance, error) { rows, err := frpcDB.Query(` - SELECT fi.id, fi.userID, fi.name, fi.bootAtStart, fi.runUser, fi.configPath, fi.createdAt, u.username + SELECT fi.id, fi.userID, fi.name, fi.bootAtStart, fi.runUser, fi.configPath, fi.createdAt, fi.watchdog, u.username FROM frpcInstances fi JOIN userLogin u ON fi.userID = u.userID `) @@ -432,7 +434,7 @@ func DBListFrpcInstances() ([]FrpcInstance, error) { for rows.Next() { var instance FrpcInstance var createdAtStr string - if err := rows.Scan(&instance.ID, &instance.UserID, &instance.Name, &instance.BootAtStart, &instance.RunUser, &instance.ConfigPath, &createdAtStr, &instance.CreatedBy); err != nil { + if err := rows.Scan(&instance.ID, &instance.UserID, &instance.Name, &instance.BootAtStart, &instance.RunUser, &instance.ConfigPath, &createdAtStr, &instance.Watchdog, &instance.CreatedBy); err != nil { return nil, fmt.Errorf("failed to scan frpc instance: %w", err) } instance.CreatedAt, _ = time.Parse(time.RFC3339, createdAtStr) diff --git a/frpAct.go b/frpAct.go index 3d8e886..fe42dd7 100644 --- a/frpAct.go +++ b/frpAct.go @@ -10,6 +10,7 @@ import ( "path/filepath" "strconv" "super-frpc/postLog" + "super-frpc/watchdog" "time" ) @@ -123,6 +124,7 @@ func CreateInstanceHandler(w http.ResponseWriter, r *http.Request) { BootAtStart: req.BootAtStart, RunUser: runUser, ConfigPath: configPath, + Watchdog: 0, } if err := DBAddFrpcInstance(instance); err != nil { @@ -459,7 +461,7 @@ func handleSystemConfigModify(w http.ResponseWriter, r *http.Request, instance F func GetUserInstances(userID int) ([]FrpcInstance, error) { rows, err := frpcDB.Query(` - SELECT id, userID, name, bootAtStart, runUser, configPath, createdAt + SELECT id, userID, name, bootAtStart, runUser, configPath, createdAt, watchdog FROM frpcInstances WHERE userID = ? `, userID) if err != nil { @@ -472,7 +474,7 @@ func GetUserInstances(userID int) ([]FrpcInstance, error) { var instance FrpcInstance var createdAtStr string if err := rows.Scan( - &instance.ID, &instance.UserID, &instance.Name, &instance.BootAtStart, &instance.RunUser, &instance.ConfigPath, &createdAtStr, + &instance.ID, &instance.UserID, &instance.Name, &instance.BootAtStart, &instance.RunUser, &instance.ConfigPath, &createdAtStr, &instance.Watchdog, ); err != nil { return nil, err } @@ -490,6 +492,16 @@ func getStringFromMap(m map[string]interface{}, key string) string { return "" } +func getNumFromMap(m map[string]interface{}, key string) int { + if v, ok := m[key].(int); ok { + return v + } + if v, ok := m[key].(float64); ok { + return int(v) + } + return 0 +} + func ListInstancesHandler(w http.ResponseWriter, r *http.Request) { userID, err := Auth(w, r, http.MethodGet) if err != nil { @@ -583,7 +595,7 @@ func StartInstanceHandler(w http.ResponseWriter, r *http.Request) { return } - switch initType { + switch initType { // Lunch frpc instance by init system case "windows": if err := StartWindowsService(serviceName); err != nil { postLog.Error(fmt.Sprintf("[StartInstanceHandler] Failed to start Windows service %s: %v", serviceName, err)) @@ -591,7 +603,6 @@ func StartInstanceHandler(w http.ResponseWriter, r *http.Request) { return } postLog.Info(fmt.Sprintf("[StartInstanceHandler] Windows service %s started successfully", serviceName)) - SendSuccessResponse(w, "Instance started successfully", nil) case "systemd": if err := StartSystemdService(serviceName); err != nil { @@ -600,7 +611,6 @@ func StartInstanceHandler(w http.ResponseWriter, r *http.Request) { return } postLog.Info(fmt.Sprintf("[StartInstanceHandler] Systemd service %s started successfully", serviceName)) - SendSuccessResponse(w, "Instance started successfully", nil) case "init.d": if err := StartInitDService(serviceName); err != nil { @@ -609,13 +619,18 @@ func StartInstanceHandler(w http.ResponseWriter, r *http.Request) { return } postLog.Info(fmt.Sprintf("[StartInstanceHandler] Init.d service %s started successfully", serviceName)) - SendSuccessResponse(w, "Instance started successfully", nil) default: postLog.Error(fmt.Sprintf("[StartInstanceHandler] Unsupported init system: %s", initType)) SendErrorResponse(w, http.StatusInternalServerError, fmt.Sprintf("Unsupported init system: %s", initType)) return } + + if is.watchdogConnected { + watchdog.AddInstance(serviceName) + } + + SendSuccessResponse(w, "Instance started successfully", nil) } func StopInstanceHandler(w http.ResponseWriter, r *http.Request) { diff --git a/frpcProxyAct.go b/frpcProxyAct.go index 0de8678..6c47909 100644 --- a/frpcProxyAct.go +++ b/frpcProxyAct.go @@ -53,12 +53,12 @@ func CreateProxyHandler(w http.ResponseWriter, r *http.Request) { Name: getStringFromMap(proxyInfoMap, "name"), Type: getStringFromMap(proxyInfoMap, "type"), LocalIP: getStringFromMap(proxyInfoMap, "localIP"), - LocalPort: getStringFromMap(proxyInfoMap, "localPort"), - RemotePort: getStringFromMap(proxyInfoMap, "remotePort"), + LocalPort: getNumFromMap(proxyInfoMap, "localPort"), + RemotePort: getNumFromMap(proxyInfoMap, "remotePort"), } if proxyInfo.Name == "" || proxyInfo.Type == "" || proxyInfo.LocalIP == "" || - proxyInfo.LocalPort == "" || proxyInfo.RemotePort == "" { + proxyInfo.LocalPort == 0 || proxyInfo.RemotePort == 0 { postLog.Error("[CreateProxyHandler] Missing required fields in proxyInfo") SendErrorResponse(w, http.StatusBadRequest, "Missing required fields in proxyInfo") return @@ -148,12 +148,12 @@ func ModifyProxyHandler(w http.ResponseWriter, r *http.Request) { Name: getStringFromMap(proxyInfoMap, "name"), Type: getStringFromMap(proxyInfoMap, "type"), LocalIP: getStringFromMap(proxyInfoMap, "localIP"), - LocalPort: getStringFromMap(proxyInfoMap, "localPort"), - RemotePort: getStringFromMap(proxyInfoMap, "remotePort"), + LocalPort: getNumFromMap(proxyInfoMap, "localPort"), + RemotePort: getNumFromMap(proxyInfoMap, "remotePort"), } if proxyInfo.Name == "" || proxyInfo.Type == "" || proxyInfo.LocalIP == "" || - proxyInfo.LocalPort == "" || proxyInfo.RemotePort == "" { + proxyInfo.LocalPort == 0 || proxyInfo.RemotePort == 0 { postLog.Error("[ModifyProxyHandler] Missing required fields in proxyInfo") SendErrorResponse(w, http.StatusBadRequest, "Missing required fields in proxyInfo") return diff --git a/main.go b/main.go index 8e5acfa..6ff99da 100644 --- a/main.go +++ b/main.go @@ -31,7 +31,8 @@ var softwareInfo SoftwareInfo = SoftwareInfo{ } type StatusInfo struct { - Status string + ServerStatus string + WatchdogStatus string } type Is struct { @@ -84,7 +85,7 @@ func main() { postLog.Info("Database initialized successfully") if err := InitFrpcDatabase(*dbPath_data); err != nil { - postLog.Warning(fmt.Sprintf("Failed to initialize frpc database: %v", err)) + postLog.Fatal(fmt.Sprintf("Failed to initialize frpc database: %v", err)) } frpLogger.SetDatabase(db, frpcDB) @@ -154,15 +155,25 @@ func main() { postLog.Error(fmt.Sprintf("Error closing frpc database: %v", err)) } + watchdog.Close() + postLog.Info("Server stopped") } func GetStatusHandler(w http.ResponseWriter, r *http.Request) { statusInfo := StatusInfo{ - Status: "Online", + ServerStatus: "Offline", + WatchdogStatus: "Offline", } if !is.online { - statusInfo.Status = "Offline" + statusInfo.ServerStatus = "Offline" + } else { + statusInfo.ServerStatus = "Online" + } + if !is.watchdogConnected { + statusInfo.WatchdogStatus = "Offline" + } else { + statusInfo.WatchdogStatus = "Online" } SendSuccessResponse(w, "getStatus", statusInfo) } diff --git a/watchdog/command.go b/watchdog/command.go index 20f4a35..2222a22 100644 --- a/watchdog/command.go +++ b/watchdog/command.go @@ -9,7 +9,7 @@ func AddInstance(serviceName string) bool { return false } - message := fmt.Sprintf("[addInstance] %s", serviceName) + message := fmt.Sprintf("[instance.add] %s", serviceName) response, err := sendMsg(message, 3) if err != nil { return false @@ -23,7 +23,7 @@ func RemoveInstance(serviceName string) bool { return false } - message := fmt.Sprintf("[removeInstance] %s", serviceName) + message := fmt.Sprintf("[instance.remove] %s", serviceName) response, err := sendMsg(message, 3) if err != nil { return false @@ -31,3 +31,17 @@ func RemoveInstance(serviceName string) bool { return response == "success" } + +func Close() bool { + if !IsConnected() { + return false + } + + message := "watchdog.shutdown" + response, err := sendMsg(message, 3) + if err != nil { + return false + } + + return response == "success" +} \ No newline at end of file