task isolation (vesoft-inc#152)

server/Makefile

@@ -9,7 +9,7 @@ build:
 	go build -o nebula-studio-server main.go
 
 run:
-	go run main.go
+	go run main.go -studio-config="./config/example-config.yaml"
 
 fmt:
 	go fmt ./...

server/README.md

@@ -36,7 +36,6 @@ $ ./nebula-studio-server -studio-config="./config/example-config.yaml"
 | DownloadErrLog | /api/import-tasks/{id:string}/err-logs       | GET |
 | ReadLog            | /api/import-tasks/logs                       | GET      |
 | ReadErrLog | /api/import-tasks/err-logs                   | GET |
-| Callback           | /api/import-tasks/finish                     | POST     |
 | GetWorkingDir      | /api/import-tasks/working-dir                | GET      |
 | GetTaskDir    | /api/import-tasks/task-dir                   | GET |
 | GetTaskLogNames | /api/import-tasks/{id:string}/task-log-names | GET    |
@@ -69,6 +68,10 @@ Response:
 
 #### ConnectDB  API
 
+Header:
+
+Authorization Basic cm9vdDoxMjM=
+
 The request json body:
 
 ```json
@@ -82,8 +85,7 @@ Response:
     "code": 0,
     "message": "",
     "data": {
-        "nsid": "e870674d-6ebc-4d9d-a1f7-bf59fdca24e8",
-        "version": "v2.6"
+        "version": "v3.0"
     }
 }
 ```
@@ -305,26 +307,6 @@ Response:
 }
 ```
 
-#### Callback API
-
-The request json body:
-
-```json
-{
-  "task_id": "123456"
-}
-```
-
-Response:
-
-```json
-{
-    "code": 0,
-    "message": "",
-    "data": ""
-}
-```
-
 #### GetWorkingDir API
 
 Response:

server/go.mod

@@ -4,6 +4,7 @@ go 1.17
 
 require (
 	github.com/axgle/mahonia v0.0.0-20180208002826-3358181d7394
+	github.com/iris-contrib/middleware/jwt v0.0.0-20210110101738-6d0a4d799b5d
 	github.com/kataras/iris/v12 v12.2.0-alpha4
 	github.com/mattn/go-sqlite3 v2.0.3+incompatible
 	github.com/saintfish/chardet v0.0.0-20120816061221-3af4cd4741ca
@@ -24,6 +25,7 @@ require (
 	github.com/andybalholm/brotli v1.0.3 // indirect
 	github.com/aymerick/douceur v0.2.0 // indirect
 	github.com/aymerick/raymond v2.0.3-0.20180322193309-b565731e1464+incompatible // indirect
+	github.com/dgrijalva/jwt-go v3.2.0+incompatible // indirect
 	github.com/eknkc/amber v0.0.0-20171010120322-cdade1c07385 // indirect
 	github.com/facebook/fbthrift v0.31.1-0.20211129061412-801ed7f9f295 // indirect
 	github.com/fatih/structs v1.1.0 // indirect

server/pkg/api/routes/import.go

@@ -41,10 +41,6 @@ var ImportRoute = base.Route{
 			Path: "err-logs",
 			GET:  controller.ReadErrLog,
 		},
-		{
-			Path: "finish",
-			POST: controller.Callback,
-		},
 		{
 			Path: "working-dir",
 			GET:  controller.GetWorkingDir,

server/pkg/webserver/app.go

@@ -3,6 +3,7 @@ package webserver
 import (
 	"github.com/vesoft-inc/nebula-studio/server/pkg/api/routes"
 	"github.com/vesoft-inc/nebula-studio/server/pkg/webserver/base"
+	"github.com/vesoft-inc/nebula-studio/server/pkg/webserver/middleware"
 
 	"github.com/kataras/iris/v12"
 	"github.com/kataras/iris/v12/middleware/logger"
@@ -16,6 +17,9 @@ func InitApp() *iris.Application {
 
 	APIRoute := &base.Route{
 		Path: "",
+		Middlewares: []base.Hook{
+			middleware.AuthenticatedLoginHandler,
+		},
 		SubRoutes: []base.Route{
 			routes.FilesRoute,
 			routes.GatewayRoute,

server/pkg/webserver/base/scheme.go

@@ -1,6 +1,7 @@
 package base
 
 import (
+	"net/http"
 	"strings"
 
 	"github.com/kataras/iris/v12"
@@ -12,6 +13,7 @@ import (
 type Result interface{}
 
 type Handler func(iris.Context) Result
+type Hook func(ctx iris.Context) error
 
 type Method struct {
 	Register func(path string, handlers ...context.Handler) *router.Route
@@ -20,7 +22,7 @@ type Method struct {
 
 type Route struct {
 	Path                   string
-	MiddleWares            []iris.Handler
+	Middlewares            []Hook
 	GET, POST, PUT, DELETE Handler
 	Desc                   string
 	SubRoutes              []Route
@@ -32,7 +34,7 @@ type Response struct {
 	Data    interface{} `json:"data"`
 }
 
-func WrapHandler(handler Handler) iris.Handler {
+func WrapHandler(handler Handler, desc string) iris.Handler {
 	return func(ctx iris.Context) {
 		result := handler(ctx)
 		ctx.StatusCode(iris.StatusOK)
@@ -52,6 +54,10 @@ func SetRoute(r router.Party, route *Route) {
 		routePath = "/" + routePath
 	}
 
+	if len(route.Middlewares) > 0 {
+		r.Use(route.Wrap(route.Middlewares)...)
+	}
+
 	methods := []Method{
 		{r.Get, route.GET},
 		{r.Post, route.POST},
@@ -62,7 +68,7 @@ func SetRoute(r router.Party, route *Route) {
 	for _, method := range methods {
 		if method.Handler != nil {
 			zap.L().Info(r.GetRelPath())
-			middleWares = append(middleWares, WrapHandler(method.Handler))
+			middleWares = append(middleWares, WrapHandler(method.Handler, route.Desc))
 			_ = method.Register(routePath, middleWares...)
 		}
 	}
@@ -75,3 +81,34 @@ func SetRoute(r router.Party, route *Route) {
 		}
 	}
 }
+
+func (rt *Route) Wrap(hooks []Hook) []iris.Handler {
+	handlers := make([]iris.Handler, 0)
+	for _, hook := range hooks {
+		handlers = append(handlers, hookHandler(hook, rt.Desc))
+	}
+	return handlers
+}
+
+func hookHandler(h Hook, desc string) iris.Handler {
+	return func(ctx iris.Context) {
+		err := h(ctx)
+		if err == nil {
+			return
+		}
+
+		zap.L().Warn(err.Error())
+		resp := &Response{
+			Code:    Error,
+			Message: err.Error(),
+		}
+
+		ctx.StatusCode(http.StatusOK)
+		_, err = ctx.JSON(resp)
+		if err != nil {
+			zap.L().Warn("response error",
+				zap.Any("body", resp),
+				zap.Error(err))
+		}
+	}
+}

server/pkg/webserver/controller/gateway.go

@@ -2,11 +2,13 @@ package controller
 
 import (
 	"encoding/base64"
+	"fmt"
 	"strings"
 
 	"github.com/vesoft-inc/nebula-http-gateway/ccore/nebula/gateway/dao"
 	"github.com/vesoft-inc/nebula-http-gateway/ccore/nebula/types"
 	"github.com/vesoft-inc/nebula-studio/server/pkg/webserver/base"
+	"github.com/vesoft-inc/nebula-studio/server/pkg/webserver/middleware"
 
 	"github.com/kataras/iris/v12"
 	"go.uber.org/zap"
@@ -149,10 +151,19 @@ func ConnectDB(ctx iris.Context) base.Result {
 			Message: err.Error(),
 		}
 	}
+	nebulaAddress := fmt.Sprintf("%s:%d", params.Address, params.Port)
+	loginToken, err := middleware.GetLoginToken(nebulaAddress, username)
+	if err != nil {
+		zap.L().Warn("connect DB fail", zap.Error(err))
+		return base.Response{
+			Code:    base.Error,
+			Message: err.Error(),
+		}
+	}
 	data := make(map[string]string)
-	data["nsid"] = clientInfo.ClientID
 	data["version"] = string(clientInfo.NebulaVersion)
 	ctx.SetCookieKV("nsid", data["nsid"])
+	ctx.SetCookieKV("token", loginToken)
 	return base.Response{
 		Code: base.Success,
 		Data: data,

server/pkg/webserver/controller/import.go

@@ -156,7 +156,9 @@ func HandleImportAction(ctx iris.Context) base.Result {
 			Message: err.Error(),
 		}
 	}
-	data, err := importer.ImportAction(params.TaskId, importer.NewTaskAction(params.TaskAction))
+	nebulaAddress := ctx.Values().GetString("nebulaAddress")
+	username := ctx.Values().GetString("username")
+	data, err := importer.ImportAction(params.TaskId, nebulaAddress, username, importer.NewTaskAction(params.TaskAction))
 	if err != nil {
 		zap.L().Warn("importAction fail", zap.Error(err))
 		return base.Response{
@@ -425,71 +427,6 @@ func readFile(path string, offset int64, limit int64) ([]string, error) {
 	return res, nil
 }
 
-func Callback(ctx iris.Context) base.Result {
-	type Params struct {
-		TaskId string `json:"taskId"`
-	}
-	params := new(Params)
-	err := ctx.ReadJSON(params)
-	if err != nil {
-		zap.L().Warn("taskId get fail", zap.Error(err))
-		return base.Response{
-			Code:    base.Error,
-			Message: err.Error(),
-		}
-	}
-	taskId := params.TaskId
-
-	muTaskId.RLock()
-	taskIdBytes, err := ioutil.ReadFile(config.Cfg.Web.TaskIdPath)
-	muTaskId.RUnlock()
-	if err != nil {
-		zap.L().Warn("read taskId file error", zap.Error(err))
-		return base.Response{
-			Code:    base.Error,
-			Message: err.Error(),
-		}
-	}
-	taskIdJSON := make(map[string]bool)
-	if len(taskIdBytes) != 0 {
-		err := json.Unmarshal(taskIdBytes, &taskIdJSON)
-		if err != nil {
-			zap.L().Warn("parse taskId file error", zap.Error(err))
-			return base.Response{
-				Code:    base.Error,
-				Message: err.Error(),
-			}
-		}
-	}
-
-	taskIdJSON[taskId] = true
-	jsonStr, err := json.Marshal(taskIdJSON)
-	if err != nil {
-		zap.L().Warn("map to json error", zap.Error(err))
-		return base.Response{
-			Code:    base.Error,
-			Message: err.Error(),
-		}
-	}
-
-	muTaskId.Lock()
-	err = os.WriteFile(config.Cfg.Web.TaskIdPath, jsonStr, 0644)
-	muTaskId.Unlock()
-	if err != nil {
-		zap.L().Warn("write jsonId file error", zap.Error(err))
-		return base.Response{
-			Code:    base.Error,
-			Message: err.Error(),
-		}
-	}
-
-	return base.Response{
-		Code:    base.Success,
-		Data:    "",
-		Message: "",
-	}
-}
-
 func GetWorkingDir(ctx iris.Context) base.Result {
 	data := dirResponse{
 		UploadDir: config.Cfg.Web.UploadDir,

server/pkg/webserver/middleware/auth.go

@@ -0,0 +1,68 @@
+package middleware
+
+import (
+	"time"
+
+	"github.com/iris-contrib/middleware/jwt"
+	"github.com/kataras/iris/v12"
+)
+
+var (
+	// TODO: Make it configurable
+	mySecret     = []byte("login secret")
+	WhiteListMap = map[string]struct{}{
+		"POST/api-nebula/db/connect": {},
+	}
+)
+
+func GetLoginToken(nebulaAddress string, username string) (string, error) {
+	now := time.Now()
+	token := jwt.NewTokenWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
+		"nebulaAddress": nebulaAddress,
+		"username":      username,
+		"iat":           now.Unix(),
+		"exp":           now.Add(24 * time.Hour).Unix(),
+	})
+	tokenString, err := token.SignedString(mySecret)
+	if err != nil {
+		return "", err
+	}
+
+	return tokenString, nil
+}
+
+func AuthenticatedLoginHandler(ctx iris.Context) error {
+	url := ctx.RouteName()
+	//HACK: Whitelisted urls do not require JWT authentication
+	if _, ok := WhiteListMap[url]; ok {
+		ctx.Next()
+		return nil
+	}
+
+	j := jwt.New(jwt.Config{
+		ValidationKeyGetter: func(token *jwt.Token) (interface{}, error) {
+			return mySecret, nil
+		},
+		Expiration: true,
+		Extractor: func(ctx iris.Context) (string, error) {
+			cookie, err := ctx.Request().Cookie("token")
+			if err != nil {
+				return "", err
+			}
+			return cookie.Value, nil
+		},
+		SigningMethod: jwt.SigningMethodHS256,
+	})
+	if err := j.CheckJWT(ctx); err != nil {
+		return err
+	}
+
+	token := ctx.Values().Get("jwt").(*jwt.Token)
+
+	userInfo := token.Claims.(jwt.MapClaims)
+	for key, value := range userInfo {
+		ctx.Values().Set(key, value)
+	}
+	ctx.Next()
+	return nil
+}