refactor: optimize code robustness

tanxiao1990 · Oct 21, 2024 · f89ef04 · f89ef04
1 parent f55cd9b
commit f89ef04
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 6 deletions.
diff --git a/alert/process/process.go b/alert/process/process.go
@@ -170,7 +170,9 @@ func (p *Processor) Handle(anomalyPoints []common.AnomalyPoint, from string, inh
 		p.handleEvent(events)
 	}
 
-	p.HandleRecover(alertingKeys, now, inhibit)
+	if from == "inner" {
+		p.HandleRecover(alertingKeys, now, inhibit)
+	}
 }
 
 func (p *Processor) BuildEvent(anomalyPoint common.AnomalyPoint, from string, now int64) *models.AlertCurEvent {

diff --git a/alert/sender/ibex.go b/alert/sender/ibex.go
@@ -185,6 +185,11 @@ func canDoIbex(username string, tpl *models.TaskTpl, host string, targetCache *m
 }
 
 func TaskAdd(f models.TaskForm, authUser string, isCenter bool) (int64, error) {
+	if storage.Cache == nil {
+		logger.Warning("event_callback_ibex: redis cache is nil")
+		return 0, fmt.Errorf("redis cache is nil")
+	}
+
 	hosts := cleanHosts(f.Hosts)
 	if len(hosts) == 0 {
 		return 0, fmt.Errorf("arg(hosts) empty")

diff --git a/models/migrate/migrate.go b/models/migrate/migrate.go
@@ -28,7 +28,7 @@ func MigrateIbexTables(db *gorm.DB) {
 		db = db.Set("gorm:table_options", tableOptions)
 	}
 
-	dts := []interface{}{&imodels.TaskMeta{}, &imodels.TaskScheduler{}, &imodels.TaskSchedulerHealth{}, &imodels.TaskHostDoing{}, &imodels.TaskAction{}}
+	dts := []interface{}{&imodels.TaskMeta{}, &imodels.TaskScheduler{}, &imodels.TaskSchedulerHealth{}, &TaskHostDoing{}, &imodels.TaskAction{}}
 	for _, dt := range dts {
 		err := db.AutoMigrate(dt)
 		if err != nil {
@@ -280,3 +280,15 @@ type BuiltinPayloads struct {
 	UUID        int64 `json:"uuid" gorm:"type:bigint;not null;index:idx_uuid;comment:'uuid of payload'"`
 	ComponentID int64 `json:"component_id" gorm:"type:bigint;index:idx_component,sort:asc;not null;default:0;comment:'component_id of payload'"`
 }
+
+type TaskHostDoing struct {
+	Id             int64  `gorm:"column:id;index;primaryKey:false"`
+	Host           string `gorm:"column:host;size:128;not null;index"`
+	Clock          int64  `gorm:"column:clock;not null;default:0"`
+	Action         string `gorm:"column:action;size:16;not null"`
+	AlertTriggered bool   `gorm:"-"`
+}
+
+func (TaskHostDoing) TableName() string {
+	return "task_host_doing"
+}
diff --git a/pkg/oauth2x/oauth2x.go b/pkg/oauth2x/oauth2x.go
@@ -208,7 +208,7 @@ func (s *SsoClient) exchangeUser(code string) (*CallbackOutput, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to exchange token: %s", err)
 	}
-	userInfo, err := s.getUserInfo(s.UserInfoAddr, oauth2Token.AccessToken, s.TranTokenMethod)
+	userInfo, err := s.getUserInfo(s.Config.ClientID, s.UserInfoAddr, oauth2Token.AccessToken, s.TranTokenMethod)
 	if err != nil {
 		logger.Errorf("failed to get user info: %s", err)
 		return nil, fmt.Errorf("failed to get user info: %s", err)
@@ -223,18 +223,18 @@ func (s *SsoClient) exchangeUser(code string) (*CallbackOutput, error) {
 	}, nil
 }
 
-func (s *SsoClient) getUserInfo(UserInfoAddr, accessToken string, TranTokenMethod string) ([]byte, error) {
+func (s *SsoClient) getUserInfo(ClientId, UserInfoAddr, accessToken string, TranTokenMethod string) ([]byte, error) {
 	var req *http.Request
 	if TranTokenMethod == "formdata" {
-		body := bytes.NewBuffer([]byte("access_token=" + accessToken))
+		body := bytes.NewBuffer([]byte("access_token=" + accessToken + "&client_id=" + ClientId))
 		r, err := http.NewRequest("POST", UserInfoAddr, body)
 		if err != nil {
 			return nil, err
 		}
 		r.Header.Add("Content-Type", "application/x-www-form-urlencoded")
 		req = r
 	} else if TranTokenMethod == "querystring" {
-		r, err := http.NewRequest("GET", UserInfoAddr+"?access_token="+accessToken, nil)
+		r, err := http.NewRequest("GET", UserInfoAddr+"?access_token="+accessToken+"&client_id="+ClientId, nil)
 		if err != nil {
 			return nil, err
 		}
@@ -246,6 +246,7 @@ func (s *SsoClient) getUserInfo(UserInfoAddr, accessToken string, TranTokenMetho
 			return nil, err
 		}
 		r.Header.Add("Authorization", "Bearer "+accessToken)
+		r.Header.Add("client_id", ClientId)
 		req = r
 	}