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 }