mindxdl--common--head_handler.go

// Copyright (c) 2021. Huawei Technologies Co., Ltd. All rights reserved.

// Package common the common package of the program
package common

import (
   "errors"
   "github.com/gin-gonic/gin"
   "golang.org/x/net/context"
   "huawei.com/npu-exporter/hwlog"
   "net/http"
   "strconv"
)

const (
   base    = 10
   bitSize = 64

   // HeaderUserIDKey HeaderUserIDKey
   HeaderUserIDKey = "UserID"
   // HeaderGroupIDKey HeaderGroupIDKey
   HeaderGroupIDKey = "GroupID"
   // HeaderRoleIDKey HeaderRoleIDKey
   HeaderRoleIDKey = "RoleID"
   // HeaderRequestIDKey HeaderRequestIDKey
   HeaderRequestIDKey = "RequestID"
   // HeaderRoleCodeKey HeaderRoleCodeKey
   HeaderRoleCodeKey = "RoleCode"
)

// RespMsg response msg
type RespMsg struct {
   Status string      `json:"status"`
   Msg    string      `json:"msg"`
   Data   interface{} `json:"data,omitempty"`
}

// HeaderInfo header struct
type HeaderInfo struct {
   TraceID  string `sql:"-" json:"requestID,omitempty"`
   UserID   uint64 `gorm:"" json:"userId,omitempty"`
   GroupID  uint64 `gorm:"" json:"groupId,omitempty"`
   RoleID   uint64 `sql:"-" json:"roleId,omitempty"`
   RoleCode string `sql:"-" json:"roleCode,omitempty"`
}

// ConstructResp construct response
func ConstructResp(c *gin.Context, errorCode string, msg string, data interface{}) {
   ConstructRespWithStatus(c, errorCode, msg, data, http.StatusOK)
}

// ConstructRespWithStatus construct response
func ConstructRespWithStatus(c *gin.Context, errorCode string, msg string, data interface{}, status int) {
   if msg == "" {
      msg = ErrorMap[errorCode]
   }
   result := RespMsg{
      Status: errorCode,
      Msg:    msg,
      Data:   data,
   }
   c.JSON(status, result)
}

// GetHeaderInfo get header info
func GetHeaderInfo(c *gin.Context) (HeaderInfo, string, error) {
   var hInfo HeaderInfo
   uuid := c.GetHeader(HeaderRequestIDKey)
   hInfo.TraceID = uuid
   userIDStr := c.GetHeader(HeaderUserIDKey)
   if userIDStr == "" {
      return HeaderInfo{}, ErrorNotFoundUserID, errors.New("no login")
   }
   userID, err := strconv.ParseUint(userIDStr, base, bitSize)
   if err != nil {
      return HeaderInfo{}, ErrorConvertStr2Uint, errors.New("convert userID failed")
   }
   hInfo.UserID = userID
   groupIDStr := c.GetHeader(HeaderGroupIDKey)
   if groupIDStr == "" {
      return hInfo, ErrorNotFoundGroupID, errors.New("not found groupID")
   }
   groupID, err := strconv.ParseUint(groupIDStr, base, bitSize)
   if err != nil {
      return HeaderInfo{}, ErrorConvertStr2Uint, errors.New("convert groupID failed")
   }
   hInfo.GroupID = groupID
   roleIDStr := c.GetHeader(HeaderRoleIDKey)
   if roleIDStr == "" {
      return hInfo, ErrorNotFoundRoleID, errors.New("not found roleID")
   }
   roleID, err := strconv.ParseUint(roleIDStr, base, bitSize)
   if err != nil {
      return HeaderInfo{}, ErrorConvertStr2Uint, errors.New("convert groupID failed")
   }
   hInfo.RoleID = roleID
   code := c.GetHeader(HeaderRoleCodeKey)
   hInfo.RoleCode = code
   return hInfo, "", nil
}

// BaseCtx the base of all controller
type BaseCtx struct {
   Ctx    context.Context
   HdInfo HeaderInfo
}

// WrapHeader check and wrap the header parameter
func WrapHeader() gin.HandlerFunc {
   return func(c *gin.Context) {
      hInfo, errorCode, err := GetHeaderInfo(c)
      if err != nil {
         hwlog.OpLog.Error(err)
         c.Abort()
         ConstructRespWithStatus(c, errorCode, err.Error(), nil, http.StatusUnauthorized)
         return
      }
      ctx := context.WithValue(context.TODO(), hwlog.ReqID, hInfo.TraceID)
      ctx = context.WithValue(ctx, hwlog.UserID, hInfo.UserID)
      baseCtx := BaseCtx{
         Ctx:    ctx,
         HdInfo: hInfo,
      }
      c.Set("Ctx", baseCtx)
   }
}

// Convert extract the baseCtx from gin.context and use as a parameter
func Convert(f func(b BaseCtx, c *gin.Context)) gin.HandlerFunc {
   return func(c *gin.Context) {
      baseCtx, ok := c.Get("Ctx")
      if !ok {
         ConstructResp(c, ErrNoLoginInfo, "", nil)
      }
      b, err := baseCtx.(BaseCtx)
      if !err {
         ConstructResp(c, ErrTypeConvert, "", nil)
      }
      f(b, c)
   }
}