add whitelist check

cmd/demoworker/main.go

@@ -191,6 +191,7 @@ func withManager(idx int, endpoint string) {
 
 func main() {
 	flag.Parse()
+	log.SetFormatter(&log.TextFormatter{FullTimestamp: true, TimestampFormat: "2006-01-02 15:04:05.000"})
 	seed, err := newManagerClient(*manager)
 	if err != nil {
 		log.Fatal("connect seed manager failed", "err", err)

config/constant.go

@@ -2,6 +2,7 @@ package config
 
 const (
 	NODE_MANAGER_SET            = "node_manager_set"
+	WORKER_SETS_PREFIX          = "worker_sets_"
 	WORKER_STATUS_PREFIX        = "worker_status_"
 	WORKER_NONCE_KEY_PREFIX     = "worker_nonce_"
 	WORKER_QUEUE_PREFIX         = "worker_queue_"

server/workerManager.go

@@ -621,8 +621,20 @@ func (wm *WorkerManager) handleWorkerMsg(worker *Worker) {
 					worker.workerAddr = addr
 				}
 				worker.registed = true
+				worker.info.nodeInfo = &omanager.NodeInfoResponse{
+					MinerPubkey:    msg.RegisteMessage.MinerPubkey,
+					BenefitAddress: msg.RegisteMessage.BenefitAddress,
+					DeviceIp:       msg.RegisteMessage.DeviceIp,
+				}
 				wm.SetWorkerAddr(worker, worker.workerAddr)
 
+				if err := wm.checkWhiteList(worker, worker.info.nodeInfo.BenefitAddress); err != nil {
+					worker.quit <- err
+					return
+				} else {
+					wm.addWorkerToSets(worker, worker.info.nodeInfo.BenefitAddress)
+				}
+
 				wreg := workerRegistry{
 					worker: worker,
 					wm:     wm,

server/workerstatu.go

@@ -174,6 +174,53 @@ func (wm *WorkerManager) InActiveWorker(worker *Worker) {
 		wm.rdb.Del(context.Background(), workerUsageInfoKey(worker))
 		wm.rdb.Del(context.Background(), workerDeviceInfoKey(worker))
 		wm.rdb.Del(context.Background(), workerResourceInfoKey(worker))
+		if worker.info.nodeInfo != nil {
+			wm.rmWorkerFromSets(worker, worker.info.nodeInfo.BenefitAddress)
+		}
+	}
+}
+
+func (wm *WorkerManager) addWorkerToSets(worker *Worker, benefit string) {
+	wm.rdb.SAdd(context.Background(), workerSetsKey(benefit), worker.workerAddr)
+}
+
+func (wm *WorkerManager) rmWorkerFromSets(worker *Worker, benefit string) {
+	wm.rdb.SRem(context.Background(), workerSetsKey(benefit), worker.workerAddr)
+}
+
+func (wm *WorkerManager) getWorkerSets(benefit string) ([]string, error) {
+	list, err := wm.rdb.SMembers(context.Background(), workerSetsKey(benefit)).Result()
+	if err == redis.ErrNil {
+		return []string{}, nil
+	}
+	return list, err
+}
+
+func (wm *WorkerManager) checkWhiteList(worker *Worker, benefit string) error {
+	wh, err := wm.node.cache.GetWhWithAddr(benefit)
+	if err != nil {
+		log.WithFields(log.Fields{
+			"worker":  worker.uuid,
+			"benefit": benefit,
+		}).Error("worker not in white list")
+		return errors.New("not in white list")
+	}
+	maxNodeCount := wh.NodeNum
+	nodeList, err := wm.getWorkerSets(benefit)
+	if err != nil {
+		return errors.New("check worker white list failed")
+	}
+	nodeInSet := false
+	for _, node := range nodeList {
+		if strings.Compare(strings.ToLower(worker.workerAddr), strings.ToLower(node)) == 0 {
+			nodeInSet = true
+			break
+		}
+	}
+	if nodeInSet || len(nodeList) < int(maxNodeCount) {
+		return nil
+	} else {
+		return errors.New("reach max nodenum with the benefit")
 	}
 }
 
@@ -197,6 +244,10 @@ func workerNonceKey(w *Worker) string {
 	return config.WORKER_NONCE_KEY_PREFIX + w.workerAddr
 }
 
+func workerSetsKey(benefit string) string {
+	return config.WORKER_SETS_PREFIX + benefit
+}
+
 func workerStatusKey(w *Worker) string {
 	id := workerId(w)
 	return fmt.Sprintf("%s_%s", config.WORKER_STATUS_PREFIX, id)