From 84ee48335b8a644c3597fb47299280a30fd7c4e6 Mon Sep 17 00:00:00 2001
From: Zexi Li <lizexi@yunionyun.com>
Date: Thu, 11 Oct 2018 19:04:11 +0800
Subject: [PATCH] fix: region task notify not work

---
 cmd/climc/shell/disks.go            | 15 +++++++++++----
 pkg/cloudcommon/db/taskman/tasks.go |  2 +-
 pkg/mcclient/auth/middleware.go     |  8 +++++++-
 pkg/mcclient/session.go             |  2 ++
 4 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/cmd/climc/shell/disks.go b/cmd/climc/shell/disks.go
index e4f7d7bed4..716d557c08 100644
--- a/cmd/climc/shell/disks.go
+++ b/cmd/climc/shell/disks.go
@@ -160,10 +160,11 @@ func init() {
 	})
 
 	type DiskCreateOptions struct {
-		STORAGE  string `help:"ID or name of storage where the disk is created"`
-		NAME     string `help:"Name of the disk"`
-		DISKDESC string `help:"Image size or size of virtual disk"`
-		Desc     string `help:"Description" metavar:"Description"`
+		STORAGE    string `help:"ID or name of storage where the disk is created"`
+		NAME       string `help:"Name of the disk"`
+		DISKDESC   string `help:"Image size or size of virtual disk"`
+		Desc       string `help:"Description" metavar:"Description"`
+		TaskNotify bool   `help:"Setup task notify"`
 	}
 	R(&DiskCreateOptions{}, "disk-create", "Create a virtual disk", func(s *mcclient.ClientSession, args *DiskCreateOptions) error {
 		params := jsonutils.NewDict()
@@ -172,11 +173,17 @@ func init() {
 		if len(args.Desc) > 0 {
 			params.Add(jsonutils.NewString(args.Desc), "description")
 		}
+		if args.TaskNotify {
+			s.PrepareTask()
+		}
 		disk, err := modules.Disks.CreateInContext(s, params, &modules.Storages, args.STORAGE)
 		if err != nil {
 			return err
 		}
 		printObject(disk)
+		if args.TaskNotify {
+			s.WaitTaskNotify()
+		}
 		return nil
 	})
 
diff --git a/pkg/cloudcommon/db/taskman/tasks.go b/pkg/cloudcommon/db/taskman/tasks.go
index 70cbca04bf..dbda4a7bd0 100644
--- a/pkg/cloudcommon/db/taskman/tasks.go
+++ b/pkg/cloudcommon/db/taskman/tasks.go
@@ -156,7 +156,7 @@ func fetchTaskParams(ctx context.Context, taskName string, taskData *jsonutils.J
 		}
 	} else {
 		if !reqContext.IsZero() {
-			if len(reqContext.TaskId) > 0 {
+			if len(reqContext.TaskId) > 0 && len(reqContext.TaskNotifyUrl) == 0 {
 				data.Add(jsonutils.NewString(reqContext.TaskId), PARENT_TASK_ID_KEY)
 			}
 			if len(reqContext.TaskNotifyUrl) > 0 {
diff --git a/pkg/mcclient/auth/middleware.go b/pkg/mcclient/auth/middleware.go
index d48d4f693b..f779b132ff 100644
--- a/pkg/mcclient/auth/middleware.go
+++ b/pkg/mcclient/auth/middleware.go
@@ -18,7 +18,7 @@ const (
 
 func Authenticate(f appsrv.FilterHandler) appsrv.FilterHandler {
 	return func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
-		tokenStr := r.Header.Get("X-Auth-Token")
+		tokenStr := r.Header.Get(mcclient.AUTH_TOKEN)
 		if len(tokenStr) == 0 {
 			httperrors.UnauthorizedError(w, "Unauthorized")
 			return
@@ -30,6 +30,12 @@ func Authenticate(f appsrv.FilterHandler) appsrv.FilterHandler {
 			return
 		}
 		ctx = context.WithValue(ctx, AUTH_TOKEN, token)
+		if taskId := r.Header.Get(mcclient.TASK_ID); taskId != "" {
+			ctx = context.WithValue(ctx, appctx.APP_CONTEXT_KEY_TASK_ID, taskId)
+		}
+		if taskNotifyUrl := r.Header.Get(mcclient.TASK_NOTIFY_URL); taskNotifyUrl != "" {
+			ctx = context.WithValue(ctx, appctx.APP_CONTEXT_KEY_TASK_NOTIFY_URL, taskNotifyUrl)
+		}
 		f(ctx, w, r)
 	}
 }
diff --git a/pkg/mcclient/session.go b/pkg/mcclient/session.go
index 9da75aa2ff..624badf80c 100644
--- a/pkg/mcclient/session.go
+++ b/pkg/mcclient/session.go
@@ -16,7 +16,9 @@ import (
 )
 
 const (
+	TASK_ID         = "X-Task-Id"
 	TASK_NOTIFY_URL = "X-Task-Notify-Url"
+	AUTH_TOKEN      = "X-Auth-Token"
 
 	DEFAULT_API_VERSION = "v1"
 )