mirror of
https://github.com/oarkflow/mq.git
synced 2025-11-01 05:32:42 +08:00
443 lines
13 KiB
Go
443 lines
13 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"time"
|
|
|
|
"github.com/oarkflow/mq"
|
|
"github.com/oarkflow/mq/dag"
|
|
)
|
|
|
|
// User represents a user with roles
|
|
type User struct {
|
|
ID string `json:"id"`
|
|
Name string `json:"name"`
|
|
Roles []string `json:"roles"`
|
|
}
|
|
|
|
// HasRole checks if user has a specific role
|
|
func (u *User) HasRole(role string) bool {
|
|
for _, r := range u.Roles {
|
|
if r == role {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// HasAnyRole checks if user has any of the specified roles
|
|
func (u *User) HasAnyRole(roles ...string) bool {
|
|
for _, requiredRole := range roles {
|
|
if u.HasRole(requiredRole) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// LoggingMiddleware logs the start and end of task processing
|
|
func LoggingMiddleware(ctx context.Context, task *mq.Task) mq.Result {
|
|
log.Printf("Middleware: Starting processing for node %s, task %s", task.Topic, task.ID)
|
|
start := time.Now()
|
|
|
|
// For middleware, we return a successful result to continue to next middleware/processor
|
|
// The actual processing will happen after all middlewares
|
|
result := mq.Result{
|
|
Status: mq.Completed,
|
|
Ctx: ctx,
|
|
Payload: task.Payload, // Pass through the payload
|
|
}
|
|
|
|
log.Printf("Middleware: Completed in %v", time.Since(start))
|
|
return result
|
|
}
|
|
|
|
// ValidationMiddleware validates the task payload
|
|
func ValidationMiddleware(ctx context.Context, task *mq.Task) mq.Result {
|
|
log.Printf("ValidationMiddleware: Validating payload for node %s", task.Topic)
|
|
|
|
// Check if payload is empty
|
|
if len(task.Payload) == 0 {
|
|
return mq.Result{
|
|
Status: mq.Failed,
|
|
Error: fmt.Errorf("empty payload not allowed"),
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
log.Printf("ValidationMiddleware: Payload validation passed")
|
|
return mq.Result{
|
|
Status: mq.Completed,
|
|
Ctx: ctx,
|
|
Payload: task.Payload,
|
|
}
|
|
}
|
|
|
|
// RoleCheckMiddleware checks if the user has required roles for accessing a sub-DAG
|
|
func RoleCheckMiddleware(requiredRoles ...string) mq.Handler {
|
|
return func(ctx context.Context, task *mq.Task) mq.Result {
|
|
log.Printf("RoleCheckMiddleware: Checking roles %v for node %s", requiredRoles, task.Topic)
|
|
|
|
// Extract user from payload
|
|
var payload map[string]any
|
|
if err := json.Unmarshal(task.Payload, &payload); err != nil {
|
|
return mq.Result{
|
|
Status: mq.Failed,
|
|
Error: fmt.Errorf("invalid payload format: %v", err),
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
userData, exists := payload["user"]
|
|
if !exists {
|
|
return mq.Result{
|
|
Status: mq.Failed,
|
|
Error: fmt.Errorf("user information not found in payload"),
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
userBytes, err := json.Marshal(userData)
|
|
if err != nil {
|
|
return mq.Result{
|
|
Status: mq.Failed,
|
|
Error: fmt.Errorf("invalid user data: %v", err),
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
var user User
|
|
if err := json.Unmarshal(userBytes, &user); err != nil {
|
|
return mq.Result{
|
|
Status: mq.Failed,
|
|
Error: fmt.Errorf("invalid user format: %v", err),
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
if !user.HasAnyRole(requiredRoles...) {
|
|
return mq.Result{
|
|
Status: mq.Failed,
|
|
Error: fmt.Errorf("user %s does not have required roles %v. User roles: %v", user.Name, requiredRoles, user.Roles),
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
log.Printf("RoleCheckMiddleware: User %s authorized for roles %v", user.Name, requiredRoles)
|
|
return mq.Result{
|
|
Status: mq.Completed,
|
|
Ctx: ctx,
|
|
Payload: task.Payload,
|
|
}
|
|
}
|
|
}
|
|
|
|
// TimingMiddleware measures execution time
|
|
func TimingMiddleware(ctx context.Context, task *mq.Task) mq.Result {
|
|
log.Printf("TimingMiddleware: Starting timing for node %s", task.Topic)
|
|
|
|
// Add timing info to context
|
|
ctx = context.WithValue(ctx, "start_time", time.Now())
|
|
|
|
return mq.Result{
|
|
Status: mq.Completed,
|
|
Ctx: ctx,
|
|
Payload: task.Payload,
|
|
}
|
|
}
|
|
|
|
// Example processor that simulates some work
|
|
type ExampleProcessor struct {
|
|
dag.Operation
|
|
}
|
|
|
|
func (p *ExampleProcessor) ProcessTask(ctx context.Context, task *mq.Task) mq.Result {
|
|
log.Printf("Processor: Processing task %s on node %s", task.ID, task.Topic)
|
|
|
|
// Simulate some processing time
|
|
time.Sleep(100 * time.Millisecond)
|
|
|
|
// Parse the payload as JSON
|
|
var payload map[string]any
|
|
if err := json.Unmarshal(task.Payload, &payload); err != nil {
|
|
return mq.Result{
|
|
Status: mq.Failed,
|
|
Error: fmt.Errorf("invalid payload: %v", err),
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
// Add processing information
|
|
payload["processed_by"] = "ExampleProcessor"
|
|
payload["processing_time"] = time.Now().Format(time.RFC3339)
|
|
|
|
resultPayload, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return mq.Result{
|
|
Status: mq.Failed,
|
|
Error: fmt.Errorf("failed to marshal result: %v", err),
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
return mq.Result{
|
|
Status: mq.Completed,
|
|
Payload: resultPayload,
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
// AdminProcessor handles admin-specific tasks
|
|
type AdminProcessor struct {
|
|
dag.Operation
|
|
}
|
|
|
|
func (p *AdminProcessor) ProcessTask(ctx context.Context, task *mq.Task) mq.Result {
|
|
log.Printf("AdminProcessor: Processing admin task %s on node %s", task.ID, task.Topic)
|
|
|
|
// Simulate admin processing
|
|
time.Sleep(200 * time.Millisecond)
|
|
|
|
// Parse the payload as JSON
|
|
var payload map[string]any
|
|
if err := json.Unmarshal(task.Payload, &payload); err != nil {
|
|
return mq.Result{
|
|
Status: mq.Failed,
|
|
Error: fmt.Errorf("invalid payload: %v", err),
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
// Add admin-specific processing information
|
|
payload["processed_by"] = "AdminProcessor"
|
|
payload["admin_action"] = "validated_and_processed"
|
|
payload["processing_time"] = time.Now().Format(time.RFC3339)
|
|
|
|
resultPayload, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return mq.Result{
|
|
Status: mq.Failed,
|
|
Error: fmt.Errorf("failed to marshal result: %v", err),
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
return mq.Result{
|
|
Status: mq.Completed,
|
|
Payload: resultPayload,
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
// UserProcessor handles user-specific tasks
|
|
type UserProcessor struct {
|
|
dag.Operation
|
|
}
|
|
|
|
func (p *UserProcessor) ProcessTask(ctx context.Context, task *mq.Task) mq.Result {
|
|
log.Printf("UserProcessor: Processing user task %s on node %s", task.ID, task.Topic)
|
|
|
|
// Simulate user processing
|
|
time.Sleep(150 * time.Millisecond)
|
|
|
|
// Parse the payload as JSON
|
|
var payload map[string]any
|
|
if err := json.Unmarshal(task.Payload, &payload); err != nil {
|
|
return mq.Result{
|
|
Status: mq.Failed,
|
|
Error: fmt.Errorf("invalid payload: %v", err),
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
// Add user-specific processing information
|
|
payload["processed_by"] = "UserProcessor"
|
|
payload["user_action"] = "authenticated_and_processed"
|
|
payload["processing_time"] = time.Now().Format(time.RFC3339)
|
|
|
|
resultPayload, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return mq.Result{
|
|
Status: mq.Failed,
|
|
Error: fmt.Errorf("failed to marshal result: %v", err),
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
return mq.Result{
|
|
Status: mq.Completed,
|
|
Payload: resultPayload,
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
// GuestProcessor handles guest-specific tasks
|
|
type GuestProcessor struct {
|
|
dag.Operation
|
|
}
|
|
|
|
func (p *GuestProcessor) ProcessTask(ctx context.Context, task *mq.Task) mq.Result {
|
|
log.Printf("GuestProcessor: Processing guest task %s on node %s", task.ID, task.Topic)
|
|
|
|
// Simulate guest processing
|
|
time.Sleep(100 * time.Millisecond)
|
|
|
|
// Parse the payload as JSON
|
|
var payload map[string]any
|
|
if err := json.Unmarshal(task.Payload, &payload); err != nil {
|
|
return mq.Result{
|
|
Status: mq.Failed,
|
|
Error: fmt.Errorf("invalid payload: %v", err),
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
// Add guest-specific processing information
|
|
payload["processed_by"] = "GuestProcessor"
|
|
payload["guest_action"] = "limited_access_processed"
|
|
payload["processing_time"] = time.Now().Format(time.RFC3339)
|
|
|
|
resultPayload, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return mq.Result{
|
|
Status: mq.Failed,
|
|
Error: fmt.Errorf("failed to marshal result: %v", err),
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
return mq.Result{
|
|
Status: mq.Completed,
|
|
Payload: resultPayload,
|
|
Ctx: ctx,
|
|
}
|
|
}
|
|
|
|
// createAdminSubDAG creates a sub-DAG for admin operations
|
|
func createAdminSubDAG() *dag.DAG {
|
|
adminDAG := dag.NewDAG("Admin Sub-DAG", "admin-subdag", func(taskID string, result mq.Result) {
|
|
log.Printf("Admin Sub-DAG completed for task %s: %s", taskID, string(result.Payload))
|
|
})
|
|
|
|
adminDAG.AddNode(dag.Function, "Admin Validate", "admin_validate", &AdminProcessor{Operation: dag.Operation{Type: dag.Function}}, true)
|
|
adminDAG.AddNode(dag.Function, "Admin Process", "admin_process", &AdminProcessor{Operation: dag.Operation{Type: dag.Function}})
|
|
adminDAG.AddNode(dag.Function, "Admin Finalize", "admin_finalize", &AdminProcessor{Operation: dag.Operation{Type: dag.Function}})
|
|
|
|
adminDAG.AddEdge(dag.Simple, "Validate to Process", "admin_validate", "admin_process")
|
|
adminDAG.AddEdge(dag.Simple, "Process to Finalize", "admin_process", "admin_finalize")
|
|
|
|
return adminDAG
|
|
}
|
|
|
|
// createUserSubDAG creates a sub-DAG for user operations
|
|
func createUserSubDAG() *dag.DAG {
|
|
userDAG := dag.NewDAG("User Sub-DAG", "user-subdag", func(taskID string, result mq.Result) {
|
|
log.Printf("User Sub-DAG completed for task %s: %s", taskID, string(result.Payload))
|
|
})
|
|
|
|
userDAG.AddNode(dag.Function, "User Auth", "user_auth", &UserProcessor{Operation: dag.Operation{Type: dag.Function}}, true)
|
|
userDAG.AddNode(dag.Function, "User Process", "user_process", &UserProcessor{Operation: dag.Operation{Type: dag.Function}})
|
|
userDAG.AddNode(dag.Function, "User Notify", "user_notify", &UserProcessor{Operation: dag.Operation{Type: dag.Function}})
|
|
|
|
userDAG.AddEdge(dag.Simple, "Auth to Process", "user_auth", "user_process")
|
|
userDAG.AddEdge(dag.Simple, "Process to Notify", "user_process", "user_notify")
|
|
|
|
return userDAG
|
|
}
|
|
|
|
// createGuestSubDAG creates a sub-DAG for guest operations
|
|
func createGuestSubDAG() *dag.DAG {
|
|
guestDAG := dag.NewDAG("Guest Sub-DAG", "guest-subdag", func(taskID string, result mq.Result) {
|
|
log.Printf("Guest Sub-DAG completed for task %s: %s", taskID, string(result.Payload))
|
|
})
|
|
|
|
guestDAG.AddNode(dag.Function, "Guest Welcome", "guest_welcome", &GuestProcessor{Operation: dag.Operation{Type: dag.Function}}, true)
|
|
guestDAG.AddNode(dag.Function, "Guest Info", "guest_info", &GuestProcessor{Operation: dag.Operation{Type: dag.Function}})
|
|
|
|
guestDAG.AddEdge(dag.Simple, "Welcome to Info", "guest_welcome", "guest_info")
|
|
|
|
return guestDAG
|
|
}
|
|
|
|
func main() {
|
|
// Create the main DAG
|
|
flow := dag.NewDAG("Role-Based Access Control DAG", "rbac-dag", func(taskID string, result mq.Result) {
|
|
log.Printf("Main DAG completed for task %s: %s", taskID, string(result.Payload))
|
|
})
|
|
|
|
// Add entry point
|
|
flow.AddNode(dag.Function, "Entry Point", "entry", &ExampleProcessor{Operation: dag.Operation{Type: dag.Function}}, true)
|
|
|
|
// Add sub-DAGs with role-based access
|
|
flow.AddDAGNode(dag.Function, "Admin Operations", "admin_ops", createAdminSubDAG())
|
|
flow.AddDAGNode(dag.Function, "User Operations", "user_ops", createUserSubDAG())
|
|
flow.AddDAGNode(dag.Function, "Guest Operations", "guest_ops", createGuestSubDAG())
|
|
|
|
// Add edges from entry to sub-DAGs
|
|
flow.AddEdge(dag.Simple, "Entry to Admin", "entry", "admin_ops")
|
|
flow.AddEdge(dag.Simple, "Entry to User", "entry", "user_ops")
|
|
flow.AddEdge(dag.Simple, "Entry to Guest", "entry", "guest_ops")
|
|
|
|
// Add global middlewares
|
|
flow.Use(LoggingMiddleware, ValidationMiddleware)
|
|
|
|
// Add role-based middlewares for sub-DAGs
|
|
flow.UseNodeMiddlewares(
|
|
dag.NodeMiddleware{
|
|
Node: "admin_ops",
|
|
Middlewares: []mq.Handler{RoleCheckMiddleware("admin", "superuser")},
|
|
},
|
|
dag.NodeMiddleware{
|
|
Node: "user_ops",
|
|
Middlewares: []mq.Handler{RoleCheckMiddleware("user", "admin", "superuser")},
|
|
},
|
|
dag.NodeMiddleware{
|
|
Node: "guest_ops",
|
|
Middlewares: []mq.Handler{RoleCheckMiddleware("guest", "user", "admin", "superuser")},
|
|
},
|
|
)
|
|
|
|
if flow.Error != nil {
|
|
panic(flow.Error)
|
|
}
|
|
|
|
// Define test users with different roles
|
|
users := []User{
|
|
{ID: "1", Name: "Alice", Roles: []string{"admin", "superuser"}},
|
|
{ID: "2", Name: "Bob", Roles: []string{"user"}},
|
|
{ID: "3", Name: "Charlie", Roles: []string{"guest"}},
|
|
{ID: "4", Name: "Dave", Roles: []string{"user", "admin"}},
|
|
{ID: "5", Name: "Eve", Roles: []string{}}, // No roles
|
|
}
|
|
|
|
// Test each user
|
|
for _, user := range users {
|
|
log.Printf("\n=== Testing user: %s (Roles: %v) ===", user.Name, user.Roles)
|
|
|
|
// Create payload with user information
|
|
payload := map[string]any{
|
|
"user": user,
|
|
"message": fmt.Sprintf("Request from %s", user.Name),
|
|
"data": "test data",
|
|
}
|
|
|
|
payloadBytes, err := json.Marshal(payload)
|
|
if err != nil {
|
|
log.Printf("Failed to marshal payload for user %s: %v", user.Name, err)
|
|
continue
|
|
}
|
|
|
|
log.Printf("Processing request for user %s with payload: %s", user.Name, string(payloadBytes))
|
|
|
|
result := flow.Process(context.Background(), payloadBytes)
|
|
if result.Error != nil {
|
|
log.Printf("❌ DAG processing failed for user %s: %v", user.Name, result.Error)
|
|
} else {
|
|
log.Printf("✅ DAG processing completed successfully for user %s: %s", user.Name, string(result.Payload))
|
|
}
|
|
}
|
|
}
|