guard/internal/usecase/opa.go

package usecase

import (
	"ark-permission/internal/domain"
	"ark-permission/internal/domain/usecase"
	ers "code.30cm.net/wanderland/library-go/errors"
	"context"
	_ "embed"
	"fmt"
	"github.com/open-policy-agent/opa/rego"
	"github.com/zeromicro/go-zero/core/logx"
	"strings"
)

//go:embed "rule.rego"
var policy []byte

type OpaUseCaseParam struct{}

type opaUseCase struct {
	// 查詢這個角色是否可用
	allowQuery rego.PreparedEvalQuery

	policies []map[string]any
}

func (o *opaUseCase) GetPolicy(ctx context.Context) []map[string]any {
	return o.policies
}

func (o *opaUseCase) CheckRBACPermission(ctx context.Context, req usecase.CheckReq) (usecase.CheckOPAResp, error) {
	results, err := o.allowQuery.Eval(ctx, rego.EvalInput(map[string]any{
		"roles":    req.Roles,
		"path":     req.Path,
		"method":   req.Method,
		"policies": o.policies,
	}))

	if err != nil {
		return usecase.CheckOPAResp{}, domain.PermissionGetDataError(fmt.Sprintf("failed to evaluate policy: %v", err))
	}

	if len(results) == 0 {
		logx.WithCallerSkip(1).WithFields(
			logx.Field("roles", req.Roles),
			logx.Field("path", req.Path),
			logx.Field("method", req.Method),
			logx.Field("policies", o.policies),
		).Error("empty RBAC policy result, possibly due to an incorrect query string or policy")
		return usecase.CheckOPAResp{}, domain.PermissionGetDataError("no results returned from policy evaluation")
	}

	data, ok := results[0].Expressions[0].Value.(map[string]any)
	if !ok {
		return usecase.CheckOPAResp{}, domain.PermissionGetDataError("unexpected data format in policy evaluation result")
	}
	resp, err := convertToCheckOPAResp(data)
	if !ok {
		return usecase.CheckOPAResp{}, domain.PermissionGetDataError(err.Error())
	}

	return resp, nil
}

// LoadPolicy 逐一處理 Policy 並且處理超時
func (o *opaUseCase) LoadPolicy(ctx context.Context, input []usecase.Policy) error {
	mapped := make([]map[string]any, 0, len(input))

	for i, policy := range input {
		select {
		case <-ctx.Done(): // 監控是否超時或取消
			logx.WithCallerSkip(1).WithFields(
				logx.Field("input", input),
			).Error("LoadPolicy context time out")
			// TODO 部分完成後處理，記錄日誌並返回成功的部分，或應該要重新 Loading....
			o.policies = append(o.policies, mapped...)
			return ers.SystemTimeoutError(fmt.Sprintf("operation timed out after processing %d policies: %v", i, ctx.Err()))
		default:
			// 繼續處理
			mapped = append(mapped, policiesToMap(policy))
		}
	}

	// 完成所有更新後紀錄，整個取代 policies
	o.policies = mapped
	return nil
}

func NewOpaUseCase(param OpaUseCaseParam) (usecase.OpaUseCase, error) {
	module := rego.Module("policy", string(policy))
	ctx := context.Background()
	var allowQueryErr error
	uc := &opaUseCase{}
	uc.allowQuery, allowQueryErr = rego.New(
		rego.Query("data.rbac"), // 要尋找的話 data 必帶， rbac = rego package ， allow 是要query 啥
		module,
	).PrepareForEval(ctx)
	if allowQueryErr != nil {
		return &opaUseCase{}, domain.PermissionGetDataError(allowQueryErr.Error())
	}

	return uc, nil
}

// 內部使用
func policiesToMap(policy usecase.Policy) map[string]any {
	return map[string]any{
		"methods": policy.Methods,
		"name":    policy.Name,
		"path":    policy.Path,
		"role":    policy.Role,
	}
}

func convertToCheckOPAResp(data map[string]any) (usecase.CheckOPAResp, error) {
	var response usecase.CheckOPAResp

	// 解析 allow 欄位
	allow, ok := data["allow"].(bool)
	if !ok {
		return usecase.CheckOPAResp{}, fmt.Errorf("missing or invalid 'allow' field")
	}
	response.Allow = allow

	// 解析 policy_name 欄位
	policyData, ok := data["policy_name"].(map[string]any)
	if ok {
		if name, ok := policyData["name"].(string); ok {
			response.PolicyName = name
			response.PlainCode = strings.HasSuffix(name, ".plan_code")
		}
	} else {
		return usecase.CheckOPAResp{}, fmt.Errorf("missing or invalid 'policy_name' field")
	}

	// 解析 request 欄位
	requestData, ok := data["request"].(map[string]any)
	if !ok {
		return usecase.CheckOPAResp{}, fmt.Errorf("missing or invalid 'request' field")
	}

	// 解析 method 和 path
	response.Request.Method, ok = requestData["method"].(string)
	if !ok {
		return usecase.CheckOPAResp{}, fmt.Errorf("missing or invalid 'method' field")
	}
	response.Request.Path, ok = requestData["path"].(string)
	if !ok {
		return usecase.CheckOPAResp{}, fmt.Errorf("missing or invalid 'path' field")
	}

	// 解析 policies 欄位
	policiesData, ok := requestData["policies"].([]any)
	if !ok {
		return usecase.CheckOPAResp{}, fmt.Errorf("missing or invalid 'policies' field")
	}
	response.Request.Policies = make([]usecase.Policy, len(policiesData))
	for i, policyData := range policiesData {
		policyMap, ok := policyData.(map[string]any)
		if !ok {
			return usecase.CheckOPAResp{}, fmt.Errorf("invalid policy format")
		}
		// 解析 methods
		methodsData, ok := policyMap["methods"].([]any)
		if !ok {
			return usecase.CheckOPAResp{}, fmt.Errorf("missing or invalid 'methods' field in policy")
		}
		methods := make([]string, len(methodsData))
		for j, m := range methodsData {
			methods[j], ok = m.(string)
			if !ok {
				return usecase.CheckOPAResp{}, fmt.Errorf("invalid method format in policy")
			}
		}
		// 組裝 policy
		response.Request.Policies[i] = usecase.Policy{
			Methods: methods,
			Name:    policyMap["name"].(string),
			Path:    policyMap["path"].(string),
			Role:    policyMap["role"].(string),
		}
	}

	// 解析 roles 欄位
	rolesData, ok := requestData["roles"].([]any)
	if !ok {
		return usecase.CheckOPAResp{}, fmt.Errorf("missing or invalid 'roles' field")
	}
	response.Request.Roles = make([]string, len(rolesData))
	for i, r := range rolesData {
		response.Request.Roles[i], ok = r.(string)
		if !ok {
			return usecase.CheckOPAResp{}, fmt.Errorf("invalid role format")
		}
	}

	return response, nil
}