benew/internal/service/white_list_service.go

package service

import (
	"context"
	"sync"
	"time"

	"github.com/rs/zerolog/log"
	"sikey.com.cn/benew/ext/internal/model"
)

type WhiteIplistService struct {
	iplist      []string
	iplistMutex sync.RWMutex

	whiteIplistModel *model.WhiteIplistModel
}

func NewWhiteIplistService(whiteIplistModel *model.WhiteIplistModel) *WhiteIplistService {
	srv := &WhiteIplistService{
		iplist: make([]string, 0),

		whiteIplistModel: whiteIplistModel,
	}

	srv.updateIplistTask()
	ticker := time.NewTicker(60 * time.Second)
	go func() {
		for {
			select {
			case <-ticker.C:
				srv.updateIplistTask()
			}
		}
	}()
	return srv
}

func (s *WhiteIplistService) updateIplistTask() {
	whiteIplist, err := s.whiteIplistModel.Select(context.TODO())
	if err != nil {
		log.Error().Stack().Err(err).Msg("updateIplistTask")
		return
	}

	s.iplistMutex.Lock()
	defer s.iplistMutex.Unlock()

	s.iplist = s.iplist[:0]
	for _, ip := range whiteIplist {
		s.iplist = append(s.iplist, ip.IP)
	}
}

// IsWhiteIplisted Check if the IP is in the whitelist
func (s *WhiteIplistService) IsWhiteIplisted(ip string) bool {
	s.iplistMutex.RLock()
	defer s.iplistMutex.RUnlock()

	for _, wip := range s.iplist {
		if wip == ip {
			return true
		}
	}
	return false
}