OpenFlare/Wavelet/internal/apps/openflare/agent/routers.go

// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0

package agent

import (
	"net/http"
	"strconv"

	"github.com/Rain-kl/Wavelet/internal/apps/openflare/apiutil"
	"github.com/Rain-kl/Wavelet/internal/apps/openflare/pages"
	"github.com/Rain-kl/Wavelet/internal/apps/openflare/websocket"
	"github.com/Rain-kl/Wavelet/internal/common/response"
	"github.com/gin-gonic/gin"
)

// RegisterHandler registers or discovers an agent node.
func RegisterHandler(c *gin.Context) {
	var payload NodePayload
	if !apiutil.BindJSON(c, &payload) {
		return
	}
	payload.IP = resolveReportedNodeIP(payload.IP, c.Request.RemoteAddr)

	var (
		result *RegistrationResponse
		err    error
	)
	if authNode, ok := AgentNodeFromContext(c); ok {
		result, err = RegisterWithAccessToken(c.Request.Context(), authNode, payload)
	} else {
		result, err = RegisterWithDiscovery(c.Request.Context(), payload)
	}
	if apiutil.AbortBadRequestOnError(c, err) {
		return
	}
	c.JSON(http.StatusOK, response.OK(result))
}

// HeartbeatHandler records agent heartbeat state.
func HeartbeatHandler(c *gin.Context) {
	var payload NodePayload
	if !apiutil.BindJSON(c, &payload) {
		return
	}
	payload.IP = resolveReportedNodeIP(payload.IP, c.Request.RemoteAddr)

	authNode, ok := AgentNodeFromContext(c)
	if !ok {
		response.AbortUnauthorized(c, errInvalidAgentToken)
		return
	}

	heartbeat, err := HeartbeatNode(c.Request.Context(), authNode, payload)
	if apiutil.AbortBadRequestOnError(c, err) {
		return
	}
	c.JSON(http.StatusOK, response.OK(heartbeat))
}

// GetActiveConfigHandler returns the active configuration version.
func GetActiveConfigHandler(c *gin.Context) {
	if _, ok := AgentNodeFromContext(c); !ok {
		response.AbortUnauthorized(c, errNodeMissingFromContext)
		return
	}
	config, err := GetActiveConfig(c.Request.Context())
	if apiutil.AbortBadRequestOnError(c, err) {
		return
	}
	c.JSON(http.StatusOK, response.OK(config))
}

// SyncWAFIPGroupsHandler syncs WAF IP groups for an agent.
func SyncWAFIPGroupsHandler(c *gin.Context) {
	var input WAFIPGroupSyncInput
	if !apiutil.BindJSON(c, &input) {
		return
	}
	result, err := SyncWAFIPGroups(c.Request.Context(), input)
	if apiutil.AbortBadRequestOnError(c, err) {
		return
	}
	c.JSON(http.StatusOK, response.OK(result))
}

// ReportApplyLogHandler records an agent apply log entry.
func ReportApplyLogHandler(c *gin.Context) {
	var payload ApplyLogPayload
	if !apiutil.BindJSON(c, &payload) {
		return
	}
	if authNode, ok := AgentNodeFromContext(c); ok {
		payload.NodeID = authNode.NodeID
	}
	log, err := ReportApplyLog(c.Request.Context(), payload)
	if apiutil.AbortBadRequestOnError(c, err) {
		return
	}
	c.JSON(http.StatusOK, response.OK(log))
}

// DownloadPagesPackageHandler streams the Pages deployment artifact to an authenticated agent.
func DownloadPagesPackageHandler(c *gin.Context) {
	deploymentID, ok := pagesDeploymentIDParam(c)
	if !ok {
		return
	}
	packageObj, fileName, err := pages.OpenDeploymentPackage(c.Request.Context(), deploymentID)
	if apiutil.AbortBadRequestOnError(c, err) {
		return
	}
	defer packageObj.Body.Close()
	c.Header("Content-Disposition", "attachment; filename="+fileName)
	if packageObj.ContentType != "" {
		c.Header("Content-Type", packageObj.ContentType)
	}
	c.DataFromReader(http.StatusOK, packageObj.ContentLength, packageObj.ContentType, packageObj.Body, nil)
}

func pagesDeploymentIDParam(c *gin.Context) (uint, bool) {
	raw := c.Param("deployment_id")
	if raw == "" {
		response.AbortBadRequest(c, "无效的 ID")
		return 0, false
	}
	id64, err := strconv.ParseUint(raw, 10, 64)
	if err != nil || id64 == 0 {
		response.AbortBadRequest(c, "无效的 ID")
		return 0, false
	}
	return uint(id64), true
}

// AgentWebSocketHandler upgrades an authenticated agent websocket connection.
func AgentWebSocketHandler(c *gin.Context) {
	authNode, ok := AgentNodeFromContext(c)
	if !ok {
		response.AbortUnauthorized(c, errInvalidAgentToken)
		return
	}
	websocket.ServeAgent(c, authNode.NodeID, HandleWSStatus)
}