[优化] 增加自动更新功能，支持更新请求和版本管理

2026-10-03 23:06:36 +08:00 · 2026-06-01 21:50:22 +08:00
parent cc50cc695e
commit d6e7e2baa2
11 changed files with 1524 additions and 211 deletions
@@ -40,8 +40,13 @@ type RelayConfig struct {

 // RelaySettings contains runtime settings for the Relay.
 type RelaySettings struct {
-	HeartbeatInterval       int  `json:"heartbeat_interval"`
-	WebsocketUpgradeEnabled bool `json:"websocket_upgrade_enabled"`
+	HeartbeatInterval       int    `json:"heartbeat_interval"`
+	WebsocketUpgradeEnabled bool   `json:"websocket_upgrade_enabled"`
+	AutoUpdate              bool   `json:"auto_update"`
+	UpdateRepo              string `json:"update_repo"`
+	UpdateNow               bool   `json:"update_now"`
+	UpdateChannel           string `json:"update_channel"`
+	UpdateTag               string `json:"update_tag"`
 }

 // RelayHeartbeatResponse is the response returned to the Relay from a heartbeat.
@@ -63,6 +68,15 @@ func HeartbeatRelay(node *model.Node, payload RelayHeartbeatPayload) (*RelayHear
 	payload.Name = strings.TrimSpace(payload.Name)
 	payload.IP = strings.TrimSpace(payload.IP)

+	previous := *node
+	updateNow := node.UpdateRequested
+	updateChannel := normalizeReleaseChannel(node.UpdateChannel)
+	updateTag := strings.TrimSpace(node.UpdateTag)
+
+	node.UpdateRequested = false
+	node.UpdateChannel = ReleaseChannelStable.String()
+	node.UpdateTag = ""
+
 	changes := make(map[string]any)
 	appendRelayChange := func(key string, before any, after any) {
 		if before != after {
@@ -74,6 +88,16 @@ func HeartbeatRelay(node *model.Node, payload RelayHeartbeatPayload) (*RelayHear
 	appendRelayChange("ext_version", node.ExtVersion, payload.ExtVersion)
 	appendRelayChange("relay_status", node.RelayStatus, payload.RelayStatus)

+	if previous.UpdateRequested {
+		appendRelayChange("update_requested", previous.UpdateRequested, false)
+	}
+	if previous.UpdateChannel != ReleaseChannelStable.String() {
+		appendRelayChange("update_channel", previous.UpdateChannel, ReleaseChannelStable.String())
+	}
+	if previous.UpdateTag != "" {
+		appendRelayChange("update_tag", previous.UpdateTag, "")
+	}
+
 	if payload.Name != "" && strings.TrimSpace(node.Name) == "" {
 		appendRelayChange("name", node.Name, payload.Name)
 		node.Name = payload.Name
@@ -113,7 +137,7 @@ func HeartbeatRelay(node *model.Node, payload RelayHeartbeatPayload) (*RelayHear

 	return &RelayHeartbeatResponse{
 		RelayConfig:   buildRelayConfig(node),
-		RelaySettings: buildRelaySettings(),
+		RelaySettings: buildRelaySettings(node, updateNow, updateChannel.String(), updateTag),
 	}, nil
 }

@@ -145,10 +169,22 @@ func buildRelayConfig(node *model.Node) *RelayConfig {
 	}
 }

-func buildRelaySettings() *RelaySettings {
+func buildRelaySettings(node *model.Node, updateNow bool, updateChannel string, updateTag string) *RelaySettings {
+	autoUpdate := false
+	if node != nil {
+		autoUpdate = node.AutoUpdateEnabled
+	}
+	if strings.TrimSpace(updateChannel) == "" {
+		updateChannel = ReleaseChannelStable.String()
+	}
 	return &RelaySettings{
 		HeartbeatInterval:       common.AgentHeartbeatInterval,
 		WebsocketUpgradeEnabled: common.AgentWebsocketUpgradeEnabled,
+		AutoUpdate:              autoUpdate,
+		UpdateRepo:              common.AgentUpdateRepo,
+		UpdateNow:               updateNow,
+		UpdateChannel:           updateChannel,
+		UpdateTag:               strings.TrimSpace(updateTag),
 	}
 }

@@ -236,6 +272,13 @@ func HeartbeatFlared(node *model.Node, payload FlaredHeartbeatPayload) (*FlaredH

 	now := time.Now()
 	previous := *node
+	updateNow := node.UpdateRequested
+	updateChannel := normalizeReleaseChannel(node.UpdateChannel)
+	updateTag := strings.TrimSpace(node.UpdateTag)
+
+	node.UpdateRequested = false
+	node.UpdateChannel = ReleaseChannelStable.String()
+	node.UpdateTag = ""

 	changes := make(map[string]any)
 	if previous.Version != payload.ClientVersion {
@@ -258,6 +301,16 @@ func HeartbeatFlared(node *model.Node, payload FlaredHeartbeatPayload) (*FlaredH
 	node.LastSeenAt = now
 	node.Status = NodeStatusOnline

+	if previous.UpdateRequested {
+		changes["update_requested"] = false
+	}
+	if previous.UpdateChannel != ReleaseChannelStable.String() {
+		changes["update_channel"] = ReleaseChannelStable.String()
+	}
+	if previous.UpdateTag != "" {
+		changes["update_tag"] = ""
+	}
+
 	if !node.IPManualOverride && payload.IP != "" && previous.IP != payload.IP {
 		changes["ip"] = payload.IP
 		node.IP = payload.IP
@@ -290,7 +343,7 @@ func HeartbeatFlared(node *model.Node, payload FlaredHeartbeatPayload) (*FlaredH
 	}
 	return &FlaredHeartbeatResponse{
 		ActiveConfig:   activeConfig,
-		TunnelSettings: buildRelaySettings(),
+		TunnelSettings: buildRelaySettings(node, updateNow, updateChannel.String(), updateTag),
 	}, nil
 }

@@ -401,3 +401,127 @@ func TestTunnelRoutePublishAndFlaredConfigUseRelayPorts(t *testing.T) {
 		t.Fatalf("unexpected proxy domains: %+v", proxy.CustomDomains)
 	}
 }
+
+func TestHeartbeatRelaySelfUpdatePropagationAndReset(t *testing.T) {
+	setupServiceTestDB(t)
+
+	node := &model.Node{
+		NodeID:            "relay-update-node",
+		Name:              "relay-u",
+		IP:                "1.1.1.1",
+		AccessToken:       "relay-update-token",
+		Status:            NodeStatusPending,
+		NodeType:          "tunnel_relay",
+		AutoUpdateEnabled: true,
+		UpdateRequested:   true,
+		UpdateChannel:     "preview",
+		UpdateTag:         "v1.2.3",
+	}
+	if err := node.Insert(); err != nil {
+		t.Fatalf("failed to seed relay node: %v", err)
+	}
+
+	resp, err := HeartbeatRelay(node, RelayHeartbeatPayload{
+		Version:     "v1.0.0",
+		ExtVersion:  "0.61.0",
+		RelayStatus: "healthy",
+	})
+	if err != nil {
+		t.Fatalf("HeartbeatRelay failed: %v", err)
+	}
+
+	if resp == nil || resp.RelaySettings == nil {
+		t.Fatal("expected non-nil response with RelaySettings")
+	}
+
+	settings := resp.RelaySettings
+	if !settings.AutoUpdate {
+		t.Error("expected AutoUpdate to be true")
+	}
+	if !settings.UpdateNow {
+		t.Error("expected UpdateNow to be true")
+	}
+	if settings.UpdateChannel != "preview" {
+		t.Errorf("expected UpdateChannel to be preview, got %q", settings.UpdateChannel)
+	}
+	if settings.UpdateTag != "v1.2.3" {
+		t.Errorf("expected UpdateTag to be v1.2.3, got %q", settings.UpdateTag)
+	}
+
+	// Verify that the requested update was cleared in the DB
+	updated, err := model.GetNodeByNodeID(node.NodeID)
+	if err != nil {
+		t.Fatalf("failed to reload node: %v", err)
+	}
+	if updated.UpdateRequested {
+		t.Error("expected UpdateRequested to be reset to false in the database")
+	}
+	if updated.UpdateChannel != "stable" {
+		t.Errorf("expected UpdateChannel to be reset to stable, got %q", updated.UpdateChannel)
+	}
+	if updated.UpdateTag != "" {
+		t.Errorf("expected UpdateTag to be reset to empty, got %q", updated.UpdateTag)
+	}
+}
+
+func TestHeartbeatFlaredSelfUpdatePropagationAndReset(t *testing.T) {
+	setupServiceTestDB(t)
+
+	node := &model.Node{
+		NodeID:            "flared-update-node",
+		Name:              "flared-u",
+		IP:                "1.1.1.2",
+		AccessToken:       "flared-update-token",
+		Status:            NodeStatusPending,
+		NodeType:          "tunnel_client",
+		AutoUpdateEnabled: true,
+		UpdateRequested:   true,
+		UpdateChannel:     "stable",
+		UpdateTag:         "v2.3.4",
+	}
+	if err := node.Insert(); err != nil {
+		t.Fatalf("failed to seed flared node: %v", err)
+	}
+
+	resp, err := HeartbeatFlared(node, FlaredHeartbeatPayload{
+		ClientVersion: "v1.0.0",
+		FrpVersion:    "0.61.0",
+		TunnelStatus:  "running",
+	})
+	if err != nil {
+		t.Fatalf("HeartbeatFlared failed: %v", err)
+	}
+
+	if resp == nil || resp.TunnelSettings == nil {
+		t.Fatal("expected non-nil response with TunnelSettings")
+	}
+
+	settings := resp.TunnelSettings
+	if !settings.AutoUpdate {
+		t.Error("expected AutoUpdate to be true")
+	}
+	if !settings.UpdateNow {
+		t.Error("expected UpdateNow to be true")
+	}
+	if settings.UpdateChannel != "stable" {
+		t.Errorf("expected UpdateChannel to be stable, got %q", settings.UpdateChannel)
+	}
+	if settings.UpdateTag != "v2.3.4" {
+		t.Errorf("expected UpdateTag to be v2.3.4, got %q", settings.UpdateTag)
+	}
+
+	// Verify that the requested update was cleared in the DB
+	updated, err := model.GetNodeByNodeID(node.NodeID)
+	if err != nil {
+		t.Fatalf("failed to reload node: %v", err)
+	}
+	if updated.UpdateRequested {
+		t.Error("expected UpdateRequested to be reset to false in the database")
+	}
+	if updated.UpdateChannel != "stable" {
+		t.Errorf("expected UpdateChannel to be reset to stable, got %q", updated.UpdateChannel)
+	}
+	if updated.UpdateTag != "" {
+		t.Errorf("expected UpdateTag to be reset to empty, got %q", updated.UpdateTag)
+	}
+}