service-common/server/whitelist.go

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
}