Skip to content

Commit ed24982

Browse files
LordofAvernusLordofAvernus
committed
refactor: optimization generate token
1 parent 64fa7c7 commit ed24982

File tree

2 files changed

+59
-33
lines changed

2 files changed

+59
-33
lines changed

‎sqle/api/app.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,7 @@ func StartApi(net *gracenet.Net, exitChan chan struct{}, config *config.SqleOpti
257257
v1ProjectOpRouter.DELETE("/:project_name/instance_audit_plans/:instance_audit_plan_id/", v1.DeleteInstanceAuditPlan)
258258
v1ProjectOpRouter.PUT("/:project_name/instance_audit_plans/:instance_audit_plan_id/", v1.UpdateInstanceAuditPlan)
259259
v1ProjectOpRouter.PATCH("/:project_name/instance_audit_plans/:instance_audit_plan_id/", v1.UpdateInstanceAuditPlanStatus)
260-
v1ProjectOpRouter.PATCH("/:project_name/instance_audit_plans/:instance_audit_plan_id/token", v1.GenerateAuditPlanToken)
260+
v1ProjectOpRouter.PATCH("/:project_name/instance_audit_plans/:instance_audit_plan_id/token", v1.RefreshAuditPlanToken)
261261

262262
// audit plan; 智能扫描任务
263263
v1ProjectOpRouter.DELETE("/:project_name/instance_audit_plans/:instance_audit_plan_id/audit_plans/:audit_plan_id/", v1.DeleteAuditPlanById)

‎sqle/api/controller/v1/instance_audit_plan.go

Lines changed: 58 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -200,7 +200,7 @@ func CreateInstanceAuditPlan(c echo.Context) error {
200200
return controller.JSONBaseErrorReq(c, err)
201201
}
202202
// generate token , 生成ID后根据ID生成token
203-
err = GenerateInstanceAuditPlanToken(ap.GetIDStr())
203+
err = HandleAuditPlanToken(ap.GetIDStr())
204204
if err != nil {
205205
return controller.JSONBaseErrorReq(c, err)
206206
}
@@ -214,7 +214,7 @@ func CreateInstanceAuditPlan(c echo.Context) error {
214214
})
215215
}
216216

217-
func GenerateInstanceAuditPlanToken(instanceAuditPlanID string) error {
217+
func HandleAuditPlanToken(instanceAuditPlanID string) error {
218218
s := model.GetStorage()
219219

220220
ap, exist, err := s.GetInstanceAuditPlanDetail(instanceAuditPlanID)
@@ -224,36 +224,50 @@ func GenerateInstanceAuditPlanToken(instanceAuditPlanID string) error {
224224
if !exist {
225225
return errors.NewInstanceAuditPlanNotExistErr()
226226
}
227-
return generateAndUpdateAuditPlanToken(ap, tokenExpire)
227+
228+
return UpdateInstanceAuditPlanToken(ap, tokenExpire)
228229
}
229230

230-
func generateAndUpdateAuditPlanToken(ap *model.InstanceAuditPlan, tokenExpire time.Duration) (err error) {
231-
var needGenerateToken bool
232-
for _, v := range ap.AuditPlans {
233-
if _, ok := auditplan.GetSupportedScannerAuditPlanType()[v.Type]; ok {
234-
needGenerateToken = true
235-
}
236-
}
237-
var t string
238-
if needGenerateToken {
239-
t, err = dmsCommonJwt.GenJwtToken(dmsCommonJwt.WithExpiredTime(tokenExpire), dmsCommonJwt.WithAuditPlanName(utils.Md5(ap.GetIDStr())))
231+
func UpdateInstanceAuditPlanToken(ap *model.InstanceAuditPlan, tokenExpire time.Duration) error {
232+
// 存在scanner依赖的任务类型时候，重新生成token
233+
needGenerate := HasScannerTypeSubPlans(ap)
234+
// 当前token是否为为空
235+
currentTokenEmpty := ap.Token == ""
236+
237+
var token string
238+
var err error
239+
if needGenerate {
240+
token, err = newAuditPlanToken(ap, tokenExpire)
240241
if err != nil {
241242
return errors.New(errors.DataConflict, err)
242243
}
243244
}
244245

245-
// 更新token的条件
246-
// 1. 需要生成token，并且现在token为空
247-
// 2. 不需要生成token（移除token），并且现在token不为空
248-
if needGenerateToken == (ap.Token == "") {
249-
err = model.GetStorage().UpdateInstanceAuditPlanByID(ap.ID, map[string]interface{}{"token": t})
250-
if err != nil {
251-
return err
252-
}
246+
// 1. 添加token: 存在scanner类型任务并且原本token为空
247+
// 2. 删除token: 不存在scanner类型任务并且原本token不为空
248+
if needGenerate == currentTokenEmpty {
249+
return model.GetStorage().UpdateInstanceAuditPlanByID(ap.ID, map[string]interface{}{"token": token})
253250
}
254251
return nil
255252
}
256253

254+
func HasScannerTypeSubPlans(ap *model.InstanceAuditPlan) bool {
255+
supportedTypes := auditplan.GetSupportedScannerAuditPlanType()
256+
for _, plan := range ap.AuditPlans {
257+
if _, ok := supportedTypes[plan.Type]; ok {
258+
return true
259+
}
260+
}
261+
return false
262+
}
263+
264+
func newAuditPlanToken(ap *model.InstanceAuditPlan, tokenExpire time.Duration) (string, error) {
265+
return dmsCommonJwt.GenJwtToken(
266+
dmsCommonJwt.WithExpiredTime(tokenExpire),
267+
dmsCommonJwt.WithAuditPlanName(utils.Md5(ap.GetIDStr())),
268+
)
269+
}
270+
257271
// @Summary 删除实例扫描任务
258272
// @Description delete instance audit plan
259273
// @Id deleteInstanceAuditPlanV1
@@ -413,7 +427,7 @@ func UpdateInstanceAuditPlan(c echo.Context) error {
413427
if err != nil {
414428
return controller.JSONBaseErrorReq(c, err)
415429
}
416-
err = GenerateInstanceAuditPlanToken(instanceAuditPlanID)
430+
err = HandleAuditPlanToken(instanceAuditPlanID)
417431
if err != nil {
418432
return controller.JSONBaseErrorReq(c, err)
419433
}
@@ -792,7 +806,7 @@ func DeleteAuditPlanById(c echo.Context) error {
792806
if err != nil {
793807
return controller.JSONBaseErrorReq(c, err)
794808
}
795-
err = GenerateInstanceAuditPlanToken(instanceAuditPlanID)
809+
err = HandleAuditPlanToken(instanceAuditPlanID)
796810
if err != nil {
797811
return controller.JSONBaseErrorReq(c, err)
798812
}
@@ -1334,22 +1348,22 @@ func AuditPlanTriggerSqlAudit(c echo.Context) error {
13341348
return controller.JSONBaseErrorReq(c, nil)
13351349
}
13361350

1337-
type GenerateAuditPlanTokenReqV1 struct {
1351+
type RefreshAuditPlanTokenReqV1 struct {
13381352
ExpiresInDays *int `json:"expires_in_days"`
13391353
}
13401354

1341-
// @Summary 生成扫描任务token
1342-
// @Description generate audit plan token
1343-
// @Id generateAuditPlanTokenV1
1355+
// @Summary 重置扫描任务token
1356+
// @Description refresh audit plan token
1357+
// @Id refreshAuditPlanTokenV1
13441358
// @Tags instance_audit_plan
13451359
// @Security ApiKeyAuth
1346-
// @param audit_plan body v1.GenerateAuditPlanTokenReqV1 false "update instance audit plan token"
1360+
// @param audit_plan body v1.RefreshAuditPlanTokenReqV1 false "update instance audit plan token"
13471361
// @Param project_name path string true "project name"
13481362
// @Param instance_audit_plan_id path string true "instance audit plan id"
13491363
// @Success 200 {object} controller.BaseRes
13501364
// @router /v1/projects/{project_name}/instance_audit_plans/{instance_audit_plan_id}/token [patch]
1351-
func GenerateAuditPlanToken(c echo.Context) error {
1352-
req := new(GenerateAuditPlanTokenReqV1)
1365+
func RefreshAuditPlanToken(c echo.Context) error {
1366+
req := new(RefreshAuditPlanTokenReqV1)
13531367
if err := controller.BindAndValidateReq(c, req); err != nil {
13541368
return controller.JSONBaseErrorReq(c, err)
13551369
}
@@ -1374,10 +1388,22 @@ func GenerateAuditPlanToken(c echo.Context) error {
13741388
expireDuration = time.Duration(expiresInDays) * 24 * time.Hour
13751389
}
13761390
}
1377-
instanceAuditPlan.Token = ""
1378-
err = generateAndUpdateAuditPlanToken(instanceAuditPlan, expireDuration)
1391+
1392+
err = RefreshInstanceAuditPlanToken(instanceAuditPlan, expireDuration)
13791393
if err != nil {
13801394
return controller.JSONBaseErrorReq(c, err)
13811395
}
13821396
return controller.JSONBaseErrorReq(c, nil)
13831397
}
1398+
1399+
func RefreshInstanceAuditPlanToken(ap *model.InstanceAuditPlan, tokenExpire time.Duration) error {
1400+
var token string
1401+
var err error
1402+
if HasScannerTypeSubPlans(ap) {
1403+
token, err = newAuditPlanToken(ap, tokenExpire)
1404+
if err != nil {
1405+
return errors.New(errors.DataConflict, err)
1406+
}
1407+
}
1408+
return model.GetStorage().UpdateInstanceAuditPlanByID(ap.ID, map[string]interface{}{"token": token})
1409+
}

0 commit comments

Comments
 (0)