feat 初始化

server/whitelist.go

@@ -0,0 +1,215 @@
+package server
+
+import (
+	"context"
+	"encoding/json"
+	"git.hlsq.asia/alex/service-common/db/etcd"
+	"git.hlsq.asia/alex/service-common/log"
+	"go.etcd.io/etcd/api/v3/mvccpb"
+	clientv3 "go.etcd.io/etcd/client/v3"
+	"sync"
+)
+
+var (
+	whiteListenWg       = &sync.WaitGroup{}
+	whiteListenStopFunc context.CancelFunc
+)
+
+const (
+	serverWhitelistKey        = "/server/whitelist"         // 白名单列表key
+	serverFactoryWhitelistKey = "/server/factory_whitelist" // 工厂白名单列表key
+)
+
+const whitelistUsnList = "usn_list" //  usn白名单列表字段名
+
+type Whitelist struct {
+	etcdClient       *etcd.Client
+	mu               sync.RWMutex
+	whitelist        map[int64]struct{}
+	factoryWhitelist map[int64]struct{}
+}
+
+var serverWhitelist *Whitelist
+
+func InitWhitelist() {
+	serverWhitelist = &Whitelist{
+		etcdClient:       etcd.GetClient(),
+		mu:               sync.RWMutex{},
+		whitelist:        make(map[int64]struct{}),
+		factoryWhitelist: make(map[int64]struct{}),
+	}
+}
+
+func ListenWhitelist(ready *sync.WaitGroup) {
+	var stopCtx context.Context
+	stopCtx, whiteListenStopFunc = context.WithCancel(context.Background())
+	whiteListenWg.Add(1)
+
+	go func(stopCtx context.Context) {
+		defer whiteListenWg.Done()
+		// 监听工厂白名单变化
+		serverFactoryWhitelistResp, _ := serverWhitelist.etcdClient.Get(serverFactoryWhitelistKey)
+		for _, kv := range serverFactoryWhitelistResp.Kvs {
+			serverWhitelist.onFactoryWhitelistChange(clientv3.EventTypePut, string(kv.Key), string(kv.Value))
+		}
+		serverFactoryWhitelistWatchChan := serverWhitelist.etcdClient.Watch(serverFactoryWhitelistKey,
+			clientv3.WithRev(serverFactoryWhitelistResp.Header.Revision+1))
+
+		// 监听白名单变化
+		serverWhitelistResp, _ := serverWhitelist.etcdClient.Get(serverWhitelistKey)
+		for _, kv := range serverWhitelistResp.Kvs {
+			serverWhitelist.onWhitelistChange(clientv3.EventTypePut, string(kv.Key), string(kv.Value))
+		}
+		serverWhitelistWatchChan := serverWhitelist.etcdClient.Watch(serverWhitelistKey,
+			clientv3.WithRev(serverWhitelistResp.Header.Revision+1))
+
+		ready.Done()
+		for {
+			select {
+			case msg := <-serverFactoryWhitelistWatchChan:
+				for _, event := range msg.Events {
+					serverWhitelist.onFactoryWhitelistChange(event.Type, string(event.Kv.Key), string(event.Kv.Value))
+				}
+			case msg := <-serverWhitelistWatchChan:
+				for _, event := range msg.Events {
+					serverWhitelist.onWhitelistChange(event.Type, string(event.Kv.Key), string(event.Kv.Value))
+				}
+			case <-stopCtx.Done():
+				return
+			}
+		}
+	}(stopCtx)
+
+	return
+}
+
+func CloseWhitelist() {
+	if whiteListenStopFunc != nil {
+		whiteListenStopFunc()
+		whiteListenWg.Wait()
+	}
+}
+
+func GetServerWhitelist() *Whitelist {
+	return serverWhitelist
+}
+
+func (s *Whitelist) GetWhitelist() map[int64]struct{} {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	return s.whitelist
+}
+
+func (s *Whitelist) SetWhitelist(whitelist []int64) {
+	white := make(map[string]interface{})
+	white[whitelistUsnList] = whitelist
+
+	whiteStr, err := json.Marshal(white)
+	if err != nil {
+		log.Errorf("server whitelist set failed: %v", err)
+		return
+	}
+
+	_, err = s.etcdClient.Put(serverWhitelistKey, string(whiteStr))
+	if err != nil {
+		log.Errorf("server whitelist set failed: %v", err)
+		return
+	}
+}
+
+func (s *Whitelist) CheckWhitelist(usn int64) bool {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	_, ok := s.whitelist[usn]
+	return ok
+}
+
+func (s *Whitelist) GetFactoryWhitelist() map[int64]struct{} {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	return s.factoryWhitelist
+}
+
+func (s *Whitelist) SetFactoryWhitelist(whitelist []int64) {
+	white := make(map[string]interface{})
+	white[whitelistUsnList] = whitelist
+
+	whiteStr, err := json.Marshal(white)
+	if err != nil {
+		log.Errorf("server factory whitelist set failed: %v", err)
+		return
+	}
+
+	_, err = s.etcdClient.Put(serverFactoryWhitelistKey, string(whiteStr))
+	if err != nil {
+		log.Errorf("server factory whitelist set failed: %v", err)
+		return
+	}
+}
+
+func (s *Whitelist) CheckFactoryWhitelist(usn int64) bool {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	_, ok := s.factoryWhitelist[usn]
+	return ok
+}
+
+func (s *Whitelist) onFactoryWhitelistChange(t mvccpb.Event_EventType, key, value string) {
+	if key != serverFactoryWhitelistKey {
+		return
+	}
+
+	switch t {
+	case clientv3.EventTypePut:
+		s.mu.Lock()
+		defer s.mu.Unlock()
+
+		userSns := s.handlerWhitelistChange(value)
+
+		s.factoryWhitelist = make(map[int64]struct{})
+		for _, usn := range userSns {
+			s.factoryWhitelist[usn] = struct{}{}
+		}
+	}
+}
+
+func (s *Whitelist) onWhitelistChange(t mvccpb.Event_EventType, key, value string) {
+	if key != serverWhitelistKey {
+		return
+	}
+
+	switch t {
+	case clientv3.EventTypePut:
+		s.mu.Lock()
+		defer s.mu.Unlock()
+
+		userSns := s.handlerWhitelistChange(value)
+
+		s.whitelist = make(map[int64]struct{})
+		for _, usn := range userSns {
+			s.whitelist[usn] = struct{}{}
+		}
+	}
+}
+
+func (s *Whitelist) handlerWhitelistChange(value string) []int64 {
+	userSns := make([]int64, 0)
+
+	white := make(map[string]interface{})
+	if err := json.Unmarshal([]byte(value), &white); err != nil {
+		log.Errorf("server whitelist set failed: %v", err)
+		return userSns
+	}
+
+	if usnList, ok := white[whitelistUsnList].([]interface{}); ok {
+		for _, usn := range usnList {
+			userSns = append(userSns, int64(usn.(float64)))
+		}
+	}
+
+	return userSns
+}