feat 初始化
This commit is contained in:
143
server/status.go
Normal file
143
server/status.go
Normal file
@@ -0,0 +1,143 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"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 (
|
||||
statusWg = &sync.WaitGroup{}
|
||||
statusStopFunc context.CancelFunc
|
||||
)
|
||||
|
||||
const (
|
||||
serverStatusKey = "/server/status" // 状态key
|
||||
)
|
||||
|
||||
type StatusType string
|
||||
|
||||
const (
|
||||
StatusInit StatusType = "init" // 初始化状态
|
||||
StatusRunning StatusType = "running" // 运行中
|
||||
StatusMaintenance StatusType = "maintenance" // 维护中
|
||||
StatusWhitelist StatusType = "whitelist" // 白名单模式
|
||||
)
|
||||
|
||||
var serverStatus *Status
|
||||
|
||||
type Status struct {
|
||||
etcdClient *etcd.Client
|
||||
status StatusType
|
||||
mu sync.RWMutex
|
||||
callbacks []func(status StatusType)
|
||||
}
|
||||
|
||||
func InitStatus() {
|
||||
serverStatus = &Status{
|
||||
etcdClient: etcd.GetClient(),
|
||||
status: StatusInit,
|
||||
mu: sync.RWMutex{},
|
||||
}
|
||||
}
|
||||
|
||||
func GetServerStatus() *Status {
|
||||
return serverStatus
|
||||
}
|
||||
|
||||
func ListenStatus(ready *sync.WaitGroup) {
|
||||
var stopCtx context.Context
|
||||
stopCtx, statusStopFunc = context.WithCancel(context.Background())
|
||||
statusWg.Add(1)
|
||||
|
||||
go func(stopCtx context.Context) {
|
||||
defer statusWg.Done()
|
||||
// 监听状态变化
|
||||
serverStatusResp, _ := serverStatus.etcdClient.Get(serverStatusKey)
|
||||
for _, kv := range serverStatusResp.Kvs {
|
||||
serverStatus.onStatusChange(clientv3.EventTypePut, string(kv.Key), string(kv.Value))
|
||||
}
|
||||
serverStatusWatchChan := serverStatus.etcdClient.Watch(serverStatusKey,
|
||||
clientv3.WithRev(serverStatusResp.Header.Revision+1))
|
||||
|
||||
ready.Done()
|
||||
for {
|
||||
select {
|
||||
case msg := <-serverStatusWatchChan:
|
||||
for _, event := range msg.Events {
|
||||
serverStatus.onStatusChange(event.Type, string(event.Kv.Key), string(event.Kv.Value))
|
||||
}
|
||||
case <-stopCtx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}(stopCtx)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func CloseStatus() {
|
||||
if statusStopFunc != nil {
|
||||
statusStopFunc()
|
||||
statusWg.Wait()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Status) GetStatus() StatusType {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
return s.status
|
||||
}
|
||||
|
||||
func (s *Status) CanAccess(usn int64) bool {
|
||||
var status StatusType
|
||||
s.mu.RLock()
|
||||
status = s.status
|
||||
s.mu.RUnlock()
|
||||
switch status {
|
||||
case StatusRunning:
|
||||
return true
|
||||
case StatusMaintenance:
|
||||
return false
|
||||
case StatusWhitelist:
|
||||
return GetServerWhitelist().CheckWhitelist(usn)
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Status) RegisterChangeCallback(fc func(status StatusType)) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.callbacks = append(s.callbacks, fc)
|
||||
}
|
||||
|
||||
func (s *Status) SetStatus(status string) {
|
||||
_, err := s.etcdClient.Put(serverStatusKey, status)
|
||||
if err != nil {
|
||||
log.Errorf("server status set failed: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (s *Status) onStatusChange(t mvccpb.Event_EventType, key, value string) {
|
||||
if key != serverStatusKey {
|
||||
return
|
||||
}
|
||||
|
||||
switch t {
|
||||
case clientv3.EventTypePut:
|
||||
s.mu.Lock()
|
||||
s.status = StatusType(value)
|
||||
s.mu.Unlock()
|
||||
for _, fc := range s.callbacks {
|
||||
fc(StatusType(value))
|
||||
}
|
||||
}
|
||||
}
|
||||
215
server/whitelist.go
Normal file
215
server/whitelist.go
Normal file
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user