nezha_domains/cmd/dashboard/controller/agentcompat_capability_contract.go

//go:build agentcompat

package controller

import (
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"strconv"
	"strings"

	"github.com/gin-gonic/gin"

	"github.com/nezhahq/nezha/service/rpc"
)

const (
	agentcompatCapabilityRequestMaxBytes = 512

	agentcompatCapabilityRegisterPath   = "/agentcompat/io-stream-capability/register"
	agentcompatCapabilityWaitPath       = "/agentcompat/io-stream-capability/wait"
	agentcompatCapabilityCancelPath     = "/agentcompat/io-stream-capability/cancel"
	agentcompatCapabilityUnregisterPath = "/agentcompat/io-stream-capability/unregister"
)

const (
	agentcompatCapabilityPurposeTerminal    = "terminal"
	agentcompatCapabilityPurposeFileManager = "file_manager"
	agentcompatCapabilityPurposeNAT         = "nat"
)

var (
	errAgentcompatCapabilityInvalid     = errors.New("agentcompat capability request is invalid")
	errAgentcompatCapabilityUnavailable = errors.New("agentcompat capability is not available")
	errAgentcompatCapabilityConflict    = errors.New("agentcompat capability is active")
	errAgentcompatCapabilityCleanup     = errors.New("agentcompat capability cleanup failed")
)

type agentcompatCapabilityRegisterRequest struct {
	Purpose    string `json:"purpose"`
	ServerID   uint64 `json:"server_id"`
	ResourceID uint64 `json:"resource_id,omitempty"`
}

type agentcompatCapabilityRegisterResponse struct {
	Capability string `json:"capability"`
}

type agentcompatCapabilityAccessRequest struct {
	Capability string `json:"capability"`
	Purpose    string `json:"purpose"`
	ServerID   uint64 `json:"server_id"`
	ResourceID uint64 `json:"resource_id,omitempty"`
}

type agentcompatCapabilityWaitRequest = agentcompatCapabilityAccessRequest
type agentcompatCapabilityCancelRequest = agentcompatCapabilityAccessRequest
type agentcompatCapabilityUnregisterRequest = agentcompatCapabilityAccessRequest

type agentcompatCapabilityWaitResponse struct {
	StreamID string `json:"stream_id"`
}

type agentcompatCapabilityEmptyResponse struct{}

func decodeAgentcompatCapabilityRequest[T any](context *gin.Context, destination *T) error {
	context.Request.Body = http.MaxBytesReader(context.Writer, context.Request.Body, agentcompatCapabilityRequestMaxBytes)
	decoder := json.NewDecoder(context.Request.Body)
	fields, err := decodeAgentcompatCapabilityObject(decoder)
	if err != nil {
		return errAgentcompatCapabilityInvalid
	}
	switch request := any(destination).(type) {
	case *agentcompatCapabilityRegisterRequest:
		request.Purpose, err = decodeAgentcompatCapabilityString(fields.purpose)
		if err == nil {
			request.ServerID, err = decodeAgentcompatCapabilityUint64(fields.serverID)
		}
		if err == nil && fields.resourceID != nil {
			request.ResourceID, err = decodeAgentcompatCapabilityUint64(fields.resourceID)
		}
		if err != nil || fields.purpose == nil || fields.serverID == nil {
			return errAgentcompatCapabilityInvalid
		}
	case *agentcompatCapabilityAccessRequest:
		request.Capability, err = decodeAgentcompatCapabilityString(fields.capability)
		if err == nil {
			request.Purpose, err = decodeAgentcompatCapabilityString(fields.purpose)
		}
		if err == nil {
			request.ServerID, err = decodeAgentcompatCapabilityUint64(fields.serverID)
		}
		if err == nil && fields.resourceID != nil {
			request.ResourceID, err = decodeAgentcompatCapabilityUint64(fields.resourceID)
		}
		if err != nil || fields.capability == nil || fields.purpose == nil || fields.serverID == nil {
			return errAgentcompatCapabilityInvalid
		}
	default:
		return errAgentcompatCapabilityInvalid
	}
	return nil
}

type agentcompatCapabilityFields struct {
	purpose    json.RawMessage
	serverID   json.RawMessage
	resourceID json.RawMessage
	capability json.RawMessage
	seen       uint8
}

func decodeAgentcompatCapabilityObject(decoder *json.Decoder) (agentcompatCapabilityFields, error) {
	var fields agentcompatCapabilityFields
	token, err := decoder.Token()
	if err != nil || token != json.Delim('{') {
		return fields, errAgentcompatCapabilityInvalid
	}
	for decoder.More() {
		keyToken, err := decoder.Token()
		key, validKey := keyToken.(string)
		if err != nil || !validKey {
			return fields, errAgentcompatCapabilityInvalid
		}
		var bit uint8
		var target *json.RawMessage
		switch key {
		case "purpose":
			bit, target = 1, &fields.purpose
		case "server_id":
			bit, target = 2, &fields.serverID
		case "resource_id":
			bit, target = 4, &fields.resourceID
		case "capability":
			bit, target = 8, &fields.capability
		default:
			return fields, errAgentcompatCapabilityInvalid
		}
		if fields.seen&bit != 0 || decoder.Decode(target) != nil {
			return fields, errAgentcompatCapabilityInvalid
		}
		fields.seen |= bit
	}
	if token, err = decoder.Token(); err != nil || token != json.Delim('}') {
		return fields, errAgentcompatCapabilityInvalid
	}
	if _, err = decoder.Token(); !errors.Is(err, io.EOF) {
		return fields, errAgentcompatCapabilityInvalid
	}
	return fields, nil
}

func decodeAgentcompatCapabilityString(raw json.RawMessage) (string, error) {
	if strings.TrimSpace(string(raw)) == "null" {
		return "", errAgentcompatCapabilityInvalid
	}
	var value string
	if err := json.Unmarshal(raw, &value); err != nil {
		return "", errAgentcompatCapabilityInvalid
	}
	return value, nil
}

func decodeAgentcompatCapabilityUint64(raw json.RawMessage) (uint64, error) {
	return strconv.ParseUint(strings.TrimSpace(string(raw)), 10, 64)
}

func parseAgentcompatCapabilityPurpose(value string) (rpc.AgentCompatCapabilityPurpose, error) {
	switch value {
	case agentcompatCapabilityPurposeTerminal:
		return rpc.AgentCompatCapabilityTerminal, nil
	case agentcompatCapabilityPurposeFileManager:
		return rpc.AgentCompatCapabilityFileManager, nil
	case agentcompatCapabilityPurposeNAT:
		return rpc.AgentCompatCapabilityNAT, nil
	default:
		return 0, errAgentcompatCapabilityInvalid
	}
}

func validateAgentcompatCapabilityIdentity(purpose rpc.AgentCompatCapabilityPurpose, serverID, resourceID uint64) error {
	if serverID == 0 {
		return errAgentcompatCapabilityInvalid
	}
	switch purpose {
	case rpc.AgentCompatCapabilityTerminal, rpc.AgentCompatCapabilityFileManager:
		if resourceID != 0 {
			return errAgentcompatCapabilityInvalid
		}
	case rpc.AgentCompatCapabilityNAT:
		if resourceID == 0 {
			return errAgentcompatCapabilityInvalid
		}
	default:
		return errAgentcompatCapabilityInvalid
	}
	return nil
}