Files
WeKnora/internal/utils/taskid.go
T
wizardchen a17a2250ed fix(security): enforce KB access on wiki reads and related IDOR paths
Gate wiki GET routes and initialization config reads with KBAccessRead,
and bind async task progress endpoints to the caller tenant via task ID
parsing. Closes cross-tenant read paths reported in GHSA-ph8f-4qwx-89rr
and related initialization/progress leaks.
2026-07-02 12:06:22 +08:00

130 lines
4.1 KiB
Go

package utils
import (
"fmt"
"strconv"
"strings"
"time"
"github.com/google/uuid"
)
// GenerateTaskID generates a unique task ID with multiple collision-resistant elements.
// The format is: <taskType>_<tenantID>_<timestamp>_<uuid>_<businessID>
//
// Parameters:
// - taskType: Type of task (e.g., "faq_import", "kb_clone")
// - tenantID: Tenant ID for multi-tenancy isolation
// - businessID: Optional business-specific ID (e.g., knowledge base ID)
//
// Returns a task ID like: "faq_import_12345_1704628851692_a1b2c3d4_kb789"
func GenerateTaskID(taskType string, tenantID uint64, businessID ...string) string {
// Use current timestamp in milliseconds for temporal uniqueness
timestamp := time.Now().UnixMilli()
// Generate a short UUID (first 8 characters for brevity)
shortUUID := strings.ReplaceAll(uuid.New().String()[:8], "-", "")
// Build the task ID components
components := []string{
sanitizeTaskType(taskType),
strconv.FormatUint(tenantID, 10),
strconv.FormatInt(timestamp, 10),
shortUUID,
}
// Add business ID if provided
if len(businessID) > 0 && businessID[0] != "" {
components = append(components, sanitizeBusinessID(businessID[0]))
}
return strings.Join(components, "_")
}
// GenerateTaskIDWithPrefix generates a task ID with a custom prefix.
// This is useful when you want more control over the task ID format.
func GenerateTaskIDWithPrefix(prefix string, tenantID uint64, businessID ...string) string {
timestamp := time.Now().UnixMilli()
shortUUID := strings.ReplaceAll(uuid.New().String()[:8], "-", "")
components := []string{
sanitizeTaskType(prefix),
strconv.FormatUint(tenantID, 10),
strconv.FormatInt(timestamp, 10),
shortUUID,
}
if len(businessID) > 0 && businessID[0] != "" {
components = append(components, sanitizeBusinessID(businessID[0]))
}
return strings.Join(components, "_")
}
// ParseTaskID parses a task ID generated by GenerateTaskID and returns its components.
// Returns taskType, tenantID, timestamp, uuid, businessID, and error.
//
// Task types may contain underscores (e.g. "faq_import", "kb_clone"), so the
// parser locates the tenant/timestamp pair rather than assuming parts[0] is
// the full task type.
func ParseTaskID(taskID string) (taskType string, tenantID uint64, timestamp int64, uuidPart string, businessID string, err error) {
parts := strings.Split(taskID, "_")
if len(parts) < 4 {
err = fmt.Errorf("invalid task ID format: %s", taskID)
return
}
tenantIdx := -1
for i := 1; i < len(parts)-2; i++ {
candidateTenant, parseErr := strconv.ParseUint(parts[i], 10, 64)
if parseErr != nil || candidateTenant == 0 {
continue
}
candidateTS, parseErr := strconv.ParseInt(parts[i+1], 10, 64)
if parseErr != nil || candidateTS < 1_000_000_000_000 {
continue
}
tenantID = candidateTenant
timestamp = candidateTS
tenantIdx = i
break
}
if tenantIdx < 0 {
err = fmt.Errorf("invalid task ID format: %s", taskID)
return
}
taskType = strings.Join(parts[:tenantIdx], "_")
uuidPart = parts[tenantIdx+2]
if len(parts) > tenantIdx+3 {
businessID = strings.Join(parts[tenantIdx+3:], "_")
}
return
}
// TaskTenantID extracts the tenant ID embedded in a GenerateTaskID task ID.
func TaskTenantID(taskID string) (uint64, error) {
_, tenantID, _, _, _, err := ParseTaskID(taskID)
return tenantID, err
}
// sanitizeTaskType ensures task type is safe for use in task ID
func sanitizeTaskType(taskType string) string {
// Replace colons and other special characters with underscores
taskType = strings.ReplaceAll(taskType, ":", "_")
taskType = strings.ReplaceAll(taskType, "-", "_")
taskType = strings.ReplaceAll(taskType, " ", "_")
return strings.ToLower(taskType)
}
// sanitizeBusinessID ensures business ID is safe for use in task ID
func sanitizeBusinessID(businessID string) string {
// Take first 12 characters and replace special characters
if len(businessID) > 12 {
businessID = businessID[:12]
}
businessID = strings.ReplaceAll(businessID, "-", "")
businessID = strings.ReplaceAll(businessID, "_", "")
businessID = strings.ReplaceAll(businessID, ":", "")
return businessID
}