im_server/server/filesafety/aliyun/green.go

// Package aliyun wraps Alibaba Cloud Green file moderation APIs.
package aliyun

import (
	"context"
	"encoding/json"
	"errors"
	"strconv"
	"strings"

	openapi "github.com/alibabacloud-go/darabonba-openapi/v2/client"
	green "github.com/alibabacloud-go/green-20220302/v3/client"
	"github.com/alibabacloud-go/tea/dara"

	"github.com/tinode/chat/server/store/types"
)

const defaultDocumentService = "document_detection"

type GreenConfig struct {
	AccessKeyID     string
	AccessKeySecret string
	RegionID        string
	Endpoint        string
	DocumentService string
}

type OSSObject struct {
	BucketName string
	ObjectName string
	RegionID   string
	DataID     string
}

type SubmitResult struct {
	ProviderTaskID string
	RequestID      string
}

type PollResult struct {
	Status     string
	RiskLevel  string
	RiskLabels []string
	Raw        map[string]any
	RequestID  string
}

type GreenClient struct {
	service string
	client  *green.Client
}

func NewGreenClient(conf GreenConfig) (*GreenClient, error) {
	if strings.TrimSpace(conf.AccessKeyID) == "" {
		return nil, errors.New("missing Access Key ID")
	}
	if strings.TrimSpace(conf.AccessKeySecret) == "" {
		return nil, errors.New("missing Access Key Secret")
	}
	if strings.TrimSpace(conf.RegionID) == "" {
		return nil, errors.New("missing Region ID")
	}
	service := strings.TrimSpace(conf.DocumentService)
	if service == "" {
		service = defaultDocumentService
	}
	clientConf := &openapi.Config{
		AccessKeyId:     dara.String(conf.AccessKeyID),
		AccessKeySecret: dara.String(conf.AccessKeySecret),
		RegionId:        dara.String(conf.RegionID),
		Protocol:        dara.String("https"),
	}
	if endpoint := strings.TrimSpace(conf.Endpoint); endpoint != "" {
		clientConf.Endpoint = dara.String(endpoint)
	}
	client, err := green.NewClient(clientConf)
	if err != nil {
		return nil, err
	}
	return &GreenClient{service: service, client: client}, nil
}

func (c *GreenClient) SubmitOSSFile(ctx context.Context, object OSSObject) (*SubmitResult, error) {
	params, err := buildOSSServiceParameters(object)
	if err != nil {
		return nil, err
	}
	resp, err := c.client.FileModerationWithContext(ctx, &green.FileModerationRequest{
		Service:           dara.String(c.service),
		ServiceParameters: dara.String(params),
	}, &dara.RuntimeOptions{})
	if err != nil {
		return nil, err
	}
	if resp == nil || resp.Body == nil || resp.Body.Code == nil || *resp.Body.Code != 200 || resp.Body.Data == nil || resp.Body.Data.TaskId == nil {
		return nil, responseError("file moderation submit failed", responseCode(resp), responseMessage(resp))
	}
	return &SubmitResult{ProviderTaskID: *resp.Body.Data.TaskId, RequestID: responseRequestID(resp)}, nil
}

func (c *GreenClient) PollOSSFile(ctx context.Context, providerTaskID string) (*PollResult, error) {
	params, err := json.Marshal(map[string]string{"taskId": strings.TrimSpace(providerTaskID)})
	if err != nil {
		return nil, err
	}
	resp, err := c.client.DescribeFileModerationResultWithContext(ctx, &green.DescribeFileModerationResultRequest{
		Service:           dara.String(c.service),
		ServiceParameters: dara.String(string(params)),
	}, &dara.RuntimeOptions{})
	if err != nil {
		return nil, err
	}
	if resp == nil || resp.Body == nil || resp.Body.Code == nil || *resp.Body.Code != 200 {
		return nil, responseError("file moderation poll failed", describeResponseCode(resp), describeResponseMessage(resp))
	}
	raw := map[string]any{}
	if data := resp.Body.Data; data != nil {
		encoded, _ := json.Marshal(data)
		json.Unmarshal(encoded, &raw)
		riskLevel := strings.TrimSpace(dara.StringValue(data.RiskLevel))
		return &PollResult{
			Status:     statusFromRiskLevel(riskLevel),
			RiskLevel:  riskLevel,
			RiskLabels: labelsFromFileResult(data),
			Raw:        raw,
			RequestID:  describeResponseRequestID(resp),
		}, nil
	}
	return &PollResult{Status: types.FileReviewStatusRunning, Raw: raw, RequestID: describeResponseRequestID(resp)}, nil
}

func buildOSSServiceParameters(object OSSObject) (string, error) {
	if strings.TrimSpace(object.BucketName) == "" || strings.TrimSpace(object.ObjectName) == "" || strings.TrimSpace(object.RegionID) == "" {
		return "", types.ErrMalformed
	}
	params := map[string]string{
		"ossBucketName": strings.TrimSpace(object.BucketName),
		"ossObjectName": strings.TrimSpace(object.ObjectName),
		"ossRegionId":   strings.TrimSpace(object.RegionID),
	}
	if dataID := strings.TrimSpace(object.DataID); dataID != "" {
		params["dataId"] = dataID
	}
	encoded, err := json.Marshal(params)
	return string(encoded), err
}

func statusFromRiskLevel(riskLevel string) string {
	switch strings.ToLower(strings.TrimSpace(riskLevel)) {
	case "":
		return types.FileReviewStatusRunning
	case "none", "low":
		return types.FileReviewStatusClean
	default:
		return types.FileReviewStatusRisky
	}
}

func labelsFromFileResult(data *green.DescribeFileModerationResultResponseBodyData) []string {
	seen := map[string]bool{}
	var labels []string
	for _, page := range data.PageResult {
		if page == nil {
			continue
		}
		for _, image := range page.ImageResult {
			if image == nil {
				continue
			}
			for _, result := range image.LabelResult {
				if result == nil {
					continue
				}
				if label := strings.TrimSpace(dara.StringValue(result.Label)); label != "" && !seen[label] {
					seen[label] = true
					labels = append(labels, label)
				}
			}
		}
		for _, text := range page.TextResult {
			if text == nil {
				continue
			}
			for _, label := range strings.Split(dara.StringValue(text.Labels), ",") {
				label = strings.TrimSpace(label)
				if label != "" && !seen[label] {
					seen[label] = true
					labels = append(labels, label)
				}
			}
		}
	}
	return labels
}

func responseError(prefix string, code int32, msg string) error {
	if msg == "" {
		return errors.New(prefix)
	}
	return errors.New(prefix + ": " + msg + " (" + strconv.FormatInt(int64(code), 10) + ")")
}

func responseCode(resp *green.FileModerationResponse) int32 {
	if resp != nil && resp.Body != nil && resp.Body.Code != nil {
		return *resp.Body.Code
	}
	return 0
}

func responseMessage(resp *green.FileModerationResponse) string {
	if resp != nil && resp.Body != nil && resp.Body.Message != nil {
		return *resp.Body.Message
	}
	return ""
}

func responseRequestID(resp *green.FileModerationResponse) string {
	if resp != nil && resp.Body != nil && resp.Body.RequestId != nil {
		return *resp.Body.RequestId
	}
	return ""
}

func describeResponseCode(resp *green.DescribeFileModerationResultResponse) int32 {
	if resp != nil && resp.Body != nil && resp.Body.Code != nil {
		return *resp.Body.Code
	}
	return 0
}

func describeResponseMessage(resp *green.DescribeFileModerationResultResponse) string {
	if resp != nil && resp.Body != nil && resp.Body.Message != nil {
		return *resp.Body.Message
	}
	return ""
}

func describeResponseRequestID(resp *green.DescribeFileModerationResultResponse) string {
	if resp != nil && resp.Body != nil && resp.Body.RequestId != nil {
		return *resp.Body.RequestId
	}
	return ""
}