im_server/server/kafka_file_lifecycle.go

package main

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"

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

const (
	kafkaFileLifecycleUploadCompleted = "file.upload.completed"
	kafkaFileLifecycleAttachmentsBind = "file.attachments.bind"
	kafkaFileLifecycleGroupFileCreate = "group.file.created"
	kafkaFileLifecycleGroupFileDelete = "group.file.deleted"
)

type kafkaFileLifecyclePayload struct {
	Action         string         `json:"action"`
	FileID         string         `json:"file_id,omitempty"`
	FileIDs        []string       `json:"file_ids,omitempty"`
	AttachmentURLs []string       `json:"attachment_urls,omitempty"`
	Topic          string         `json:"topic,omitempty"`
	MessageID      string         `json:"message_id,omitempty"`
	SeqID          int64          `json:"seq_id,omitempty"`
	UserID         string         `json:"user_id,omitempty"`
	GroupFileID    string         `json:"group_file_id,omitempty"`
	FileType       string         `json:"file_type,omitempty"`
	FileName       string         `json:"file_name,omitempty"`
	MIMEType       string         `json:"mime_type,omitempty"`
	FileSize       int64          `json:"file_size,omitempty"`
	Storage        map[string]any `json:"storage,omitempty"`
	OccurredAt     time.Time      `json:"occurred_at"`
}

func enqueueKafkaFileUploaded(ctx context.Context, tenantID types.TenantID, fd *types.FileDef) error {
	if fd == nil || fd.TenantID != tenantID || fd.Status != types.UploadCompleted {
		return nil
	}
	payload := kafkaFileLifecyclePayload{
		Action:     kafkaFileLifecycleUploadCompleted,
		FileID:     fd.Uid().String(),
		UserID:     types.ParseUid(fd.User).UserId(),
		MIMEType:   fd.MimeType,
		FileSize:   fd.Size,
		OccurredAt: fd.UpdatedAt,
		Storage: map[string]any{
			"location": fd.Location,
			"etag":     fd.ETag,
		},
	}
	return enqueueKafkaFileLifecycle(ctx, tenantID, kafkaFileLifecycleUploadCompleted, fd.Uid().String(), payload)
}

func enqueueKafkaFileAttachmentsBind(ctx context.Context, tenantID types.TenantID, msg *types.Message, attachmentURLs []string) error {
	if msg == nil || msg.TenantID != tenantID || len(attachmentURLs) == 0 {
		return nil
	}
	payload := kafkaFileLifecyclePayload{
		Action:         kafkaFileLifecycleAttachmentsBind,
		AttachmentURLs: append([]string(nil), attachmentURLs...),
		Topic:          msg.Topic,
		MessageID:      msg.Uid().String(),
		SeqID:          msg.SeqId,
		UserID:         types.ParseUid(msg.From).UserId(),
		FileType:       lifecycleFileTypeFromMessageHead(msg.Head),
		OccurredAt:     msg.CreatedAt,
	}
	return enqueueKafkaFileLifecycle(ctx, tenantID, kafkaFileLifecycleAttachmentsBind, fmt.Sprintf("%s:%d", msg.Topic, msg.SeqId), payload)
}

func enqueueKafkaGroupFileLifecycle(ctx context.Context, tenantID types.TenantID, eventType string, file *types.GroupFile, actor types.Uid) error {
	if file == nil || file.Topic == "" || file.FileID == "" {
		return nil
	}
	occurredAt := file.UpdatedAt
	if occurredAt.IsZero() {
		occurredAt = time.Now().UTC().Round(time.Millisecond)
	}
	payload := kafkaFileLifecyclePayload{
		Action:      eventType,
		FileID:      file.FileID,
		FileIDs:     []string{file.FileID},
		Topic:       file.Topic,
		GroupFileID: file.GroupFileID,
		SeqID:       file.SourceSeq,
		UserID:      actor.UserId(),
		FileType:    file.FileType,
		FileName:    file.FileName,
		MIMEType:    file.MIMEType,
		FileSize:    file.FileSize,
		OccurredAt:  occurredAt,
	}
	return enqueueKafkaFileLifecycle(ctx, tenantID, eventType, firstNonEmpty(file.GroupFileID, file.FileID), payload)
}

func enqueueKafkaFileLifecycle(ctx context.Context, tenantID types.TenantID, eventType, bizID string, payload kafkaFileLifecyclePayload) error {
	if payload.OccurredAt.IsZero() {
		payload.OccurredAt = time.Now().UTC().Round(time.Millisecond)
	}
	return executeKafkaFileLifecycle(tenantID, payload)
}

func executeKafkaFileLifecycle(tenantID types.TenantID, payload kafkaFileLifecyclePayload) error {
	switch payload.Action {
	case kafkaFileLifecycleAttachmentsBind:
		return executeKafkaFileAttachmentsBind(tenantID, payload)
	case kafkaFileLifecycleUploadCompleted, kafkaFileLifecycleGroupFileCreate, kafkaFileLifecycleGroupFileDelete:
		return nil
	default:
		return types.ErrMalformed
	}
}

func executeKafkaFileAttachmentsBind(tenantID types.TenantID, payload kafkaFileLifecyclePayload) error {
	if payload.Topic == "" || payload.MessageID == "" || len(payload.AttachmentURLs) == 0 {
		return types.ErrMalformed
	}
	msgID := types.ParseUid(payload.MessageID)
	if msgID.IsZero() {
		return types.ErrMalformed
	}
	if types.GetTopicCat(payload.Topic) != types.TopicCatGrp {
		return store.Files.LinkAttachments(tenantID, "", msgID, payload.AttachmentURLs)
	}
	uploader := types.ParseUserId(payload.UserID)
	if uploader.IsZero() {
		return types.ErrMalformed
	}
	for _, url := range payload.AttachmentURLs {
		fid := store.Store.GetMediaHandler().GetIdFromUrl(url)
		if fid.IsZero() {
			continue
		}
		_, err := store.GroupBusiness.CreateFile(tenantID, &types.GroupFile{
			Topic:          payload.Topic,
			FileID:         fid.String(),
			SourceTopic:    payload.Topic,
			SourceSeq:      payload.SeqID,
			UploaderID:     uploader,
			FileType:       lifecycleFileType(payload.FileType),
			IdempotencyKey: fmt.Sprintf("%s:%d:%s", payload.Topic, payload.SeqID, fid.String()),
		})
		if err != nil && !errors.Is(err, types.ErrDuplicate) {
			return err
		}
	}
	return nil
}

func lifecycleFileTypeFromMessageHead(head types.KVMap) string {
	msgType := lifecycleHeaderString(head, "x-im-type")
	if msgType == "" {
		msgType = lifecycleHeaderString(head, "type")
	}
	if msgType == "" {
		return "file"
	}
	return lifecycleFileType(msgType)
}

func lifecycleFileType(raw string) string {
	switch lifecycleCanonicalMessageType(raw) {
	case "image", "audio", "voice", "video":
		return lifecycleCanonicalMessageType(raw)
	default:
		return "file"
	}
}

func lifecycleHeaderString(head types.KVMap, key string) string {
	if head == nil {
		return ""
	}
	value, ok := head[key]
	if !ok {
		return ""
	}
	str, _ := value.(string)
	return strings.TrimSpace(str)
}

func lifecycleCanonicalMessageType(raw string) string {
	switch strings.ToLower(strings.TrimSpace(raw)) {
	case "audio":
		return "voice"
	default:
		return strings.ToLower(strings.TrimSpace(raw))
	}
}