Create & Init Project...

app/service/main/location/zlimit/BUILD

@@ -0,0 +1,56 @@
+package(default_visibility = ["//visibility:public"])
+
+load(
+    "@io_bazel_rules_go//go:def.bzl",
+    "go_test",
+    "go_library",
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = ["zlimit_test.go"],
+    embed = [":go_default_library"],
+    rundir = ".",
+    tags = ["automanaged"],
+    deps = [
+        "//library/cache/redis:go_default_library",
+        "//library/container/pool:go_default_library",
+        "//library/database/sql:go_default_library",
+        "//library/time:go_default_library",
+    ],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = [
+        "cache.go",
+        "db.go",
+        "zlimit.go",
+    ],
+    importpath = "go-common/app/service/main/location/zlimit",
+    tags = ["automanaged"],
+    deps = [
+        "//app/service/main/location/model:go_default_library",
+        "//library/cache/redis:go_default_library",
+        "//library/database/sql:go_default_library",
+        "//library/ecode:go_default_library",
+        "//library/log:go_default_library",
+        "//library/net/ip:go_default_library",
+        "//library/time:go_default_library",
+        "//library/xstr:go_default_library",
+    ],
+)
+
+filegroup(
+    name = "package-srcs",
+    srcs = glob(["**"]),
+    tags = ["automanaged"],
+    visibility = ["//visibility:private"],
+)
+
+filegroup(
+    name = "all-srcs",
+    srcs = [":package-srcs"],
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)

app/service/main/location/zlimit/cache.go

@@ -0,0 +1,88 @@
+package zlimit
+
+import (
+	"context"
+	"strconv"
+
+	"go-common/library/cache/redis"
+	"go-common/library/log"
+)
+
+func keyZlimit(aid int64) (key string) {
+	key = _prefixBlackList + strconv.FormatInt(aid, 10)
+	return
+}
+
+// existsRule if existes ruls in redis
+func (s *Service) existsRule(c context.Context, aid int64) (ok bool, err error) {
+	conn := s.redis.Get(c)
+	defer conn.Close()
+	if ok, err = redis.Bool(conn.Do("EXISTS", keyZlimit(aid))); err != nil {
+		log.Error("conn.DO(HEXISTS) error(%v)", err)
+	}
+	return
+}
+
+// rule get zone rule from redis
+func (s *Service) rule(c context.Context, aid int64, zoneids []int64) (res []int64, err error) {
+	var playauth int64
+	key := keyZlimit(aid)
+	conn := s.redis.Get(c)
+	defer conn.Close()
+	for _, v := range zoneids {
+		if err = conn.Send("HGET", key, v); err != nil {
+			log.Error("rule conn.Send(HGET, %s, %d) error(%v)", key, v, err)
+			return
+		}
+	}
+	if err = conn.Flush(); err != nil {
+		log.Error("rule conn.Flush error(%v)", err)
+		return
+	}
+	for i := 0; i < len(zoneids); i++ {
+		if playauth, err = redis.Int64(conn.Receive()); err != nil {
+			if err != redis.ErrNil {
+				log.Error("rule conn.Receive()%d error(%v)", i+1, err)
+				return
+			}
+			err = nil
+		}
+		res = append(res, playauth)
+	}
+	return
+}
+
+// addRule add zone rule from redis
+func (s *Service) addRule(c context.Context, zoneids map[int64]map[int64]int64) (err error) {
+	var key string
+	conn := s.redis.Get(c)
+	defer conn.Close()
+	count := 0
+	for aid, zids := range zoneids {
+		if key == "" {
+			key = keyZlimit(aid)
+		}
+		for zid, auth := range zids {
+			if err = conn.Send("HSET", key, zid, auth); err != nil {
+				log.Error("add conn.Send error(%v)", err)
+				return
+			}
+			count++
+		}
+	}
+	if err = conn.Send("EXPIRE", key, s.expire); err != nil {
+		log.Error("add conn.Send error(%v)", err)
+		return
+	}
+	if err = conn.Flush(); err != nil {
+		log.Error("add conn.Flush error(%v)", err)
+		return
+	}
+	for i := 0; i <= count; i++ {
+		if _, err = conn.Receive(); err != nil {
+			log.Error("add conn.Receive()%d error(%v)", i+1, err)
+			return
+		}
+	}
+	return
+}

app/service/main/location/zlimit/db.go

@@ -0,0 +1,96 @@
+package zlimit
+
+import (
+	"context"
+
+	"go-common/library/database/sql"
+	"go-common/library/log"
+	"go-common/library/xstr"
+)
+
+const (
+	_getPolicySQL       = "SELECT id, play_auth, down_auth, zone_id FROM policy_item WHERE zone_id <> '' AND state=1"
+	_getRelationSQL     = "SELECT policy_id FROM archive_relation WHERE aid=?"
+	_getGolbalPolicySQL = "select group_id,group_concat(id) from policy_item WHERE zone_id <> '' AND state=1 GROUP BY group_id"
+)
+
+// policies get policy data from db
+func (s *Service) policies(c context.Context) (res map[int64]map[int64]int64, err error) {
+	var (
+		tmpres map[int64]int64
+		ok     bool
+	)
+	rows, err := s.db.Query(c, _getPolicySQL)
+	if err != nil {
+		log.Error("db.Query error(%v)", err)
+		return
+	}
+	defer rows.Close()
+	res = make(map[int64]map[int64]int64)
+	for rows.Next() {
+		var (
+			pid, playAuth, downAuth int64
+			zoneID                  string
+			zoneIDs                 []int64
+		)
+		if err = rows.Scan(&pid, &playAuth, &downAuth, &zoneID); err != nil {
+			log.Error("rows.Scan error(%v)", err)
+			return
+		}
+		if zoneIDs, err = xstr.SplitInts(zoneID); err != nil {
+			log.Error("xstr.SplitInts(%s) error(%v)", zoneID, err)
+			continue
+		}
+		for _, zoneid := range zoneIDs {
+			if tmpres, ok = res[pid]; !ok {
+				tmpres = make(map[int64]int64)
+				res[pid] = tmpres
+			}
+			resCode := playAuth<<8 | downAuth
+			tmpres[zoneid] = resCode
+		}
+	}
+	return
+}
+
+// groupPolicies get policy data from db group by group_id
+func (s *Service) groupPolicies(c context.Context) (res map[int64][]int64, err error) {
+	rows, err := s.db.Query(c, _getGolbalPolicySQL)
+	if err != nil {
+		log.Error("db.Query error(%v)", err)
+		return
+	}
+	defer rows.Close()
+	res = make(map[int64][]int64)
+	for rows.Next() {
+		var (
+			groupID int64
+			pids    string
+			zoneIDs []int64
+		)
+		if err = rows.Scan(&groupID, &pids); err != nil {
+			log.Error("rows.Scan error(%v)", err)
+			return
+		}
+		if zoneIDs, err = xstr.SplitInts(pids); err != nil {
+			log.Error("xstr.SplitInts(%s) error(%v)", pids, err)
+			continue
+		}
+		res[groupID] = zoneIDs
+	}
+	return
+}
+
+// policy get pids from db by aid
+func (s *Service) groupid(c context.Context, aid int64) (gid int64, err error) {
+	row := s.getRelationStmt.QueryRow(c, aid)
+	if err = row.Scan(&gid); err != nil {
+		if err == sql.ErrNoRows {
+			gid = 0
+			err = nil
+		} else {
+			log.Error("rows.Scan error(%v)", err)
+		}
+	}
+	return
+}

app/service/main/location/zlimit/zlimit.go

@@ -0,0 +1,300 @@
+package zlimit
+
+import (
+	"context"
+	"time"
+
+	"go-common/app/service/main/location/model"
+	"go-common/library/cache/redis"
+	"go-common/library/database/sql"
+	"go-common/library/ecode"
+	"go-common/library/log"
+	xip "go-common/library/net/ip"
+	xtime "go-common/library/time"
+	"go-common/library/xstr"
+)
+
+const (
+	_prefixBlackList = "zl_"
+)
+
+// Config default struct
+type Config struct {
+	DB        *sql.Config
+	Redis     *Redis
+	IPFile    string
+	FlushTime xtime.Duration
+}
+
+// Redis redis struct
+type Redis struct {
+	*redis.Config
+	Expire xtime.Duration
+}
+
+// Service zlimit service struct
+type Service struct {
+	// mysql
+	db                 *sql.DB
+	getPolicyStmt      *sql.Stmt
+	getRelationStmt    *sql.Stmt
+	getGroupPolicyStmt *sql.Stmt
+	// redis
+	redis     *redis.Pool
+	expire    int32
+	flushTime time.Duration
+	// cache
+	policy      map[int64]map[int64]int64
+	groupPolicy map[int64][]int64
+	missch      chan interface{}
+	// xip
+	list *xip.List
+}
+
+// New new zlimit service
+func New(c *Config) (s *Service) {
+	var err error
+	s = &Service{
+		db:          sql.NewMySQL(c.DB),
+		redis:       redis.NewPool(c.Redis.Config),
+		expire:      int32(time.Duration(c.Redis.Expire) / time.Second),
+		missch:      make(chan interface{}, 1024),
+		policy:      make(map[int64]map[int64]int64),
+		groupPolicy: make(map[int64][]int64),
+		flushTime:   time.Duration(c.FlushTime),
+	}
+	s.getPolicyStmt = s.db.Prepared(_getPolicySQL)
+	s.getRelationStmt = s.db.Prepared(_getRelationSQL)
+	s.getGroupPolicyStmt = s.db.Prepared(_getGolbalPolicySQL)
+	s.load()
+	s.list, err = xip.New(c.IPFile)
+	if err != nil {
+		log.Error("xip.New(%s) error(%v)", c.IPFile, err)
+		panic(err)
+	}
+	go s.reloadproc()
+	go s.cacheproc()
+	return
+}
+
+func (s *Service) load() {
+	var (
+		tmpPolicy      map[int64]map[int64]int64
+		tmpGroupPolicy map[int64][]int64
+		err            error
+	)
+	if tmpPolicy, err = s.policies(context.TODO()); err != nil {
+		log.Error("s.policies error(%v)", err)
+	} else if len(tmpPolicy) > 0 {
+		s.policy = tmpPolicy
+	}
+	if tmpGroupPolicy, err = s.groupPolicies(context.TODO()); err != nil {
+		log.Error("s.groupPolicies error(%v)", err)
+	} else if len(tmpGroupPolicy) > 0 {
+		s.groupPolicy = tmpGroupPolicy
+	}
+}
+
+// reloadproc reload data from db
+func (s *Service) reloadproc() {
+	for {
+		s.load()
+		time.Sleep(s.flushTime)
+	}
+}
+
+// Find redio rule by aid and ipaddr
+func (s *Service) Find(c context.Context, aid int64, ipaddr, cdnip string) (ret, retdown int64, err error) {
+	var (
+		ok                  bool
+		auth, pid, zid, gid int64
+		rules, pids         []int64
+		zids                map[int64]int64
+		ipInfo              *xip.Zone
+	)
+	ipInfo = s.list.Zone(ipaddr)
+	if (ipInfo != nil) && (ipInfo.Province == "共享地址" || ipInfo.City == "共享地址") && cdnip != "" {
+		ipInfo = s.list.Zone(cdnip)
+	}
+	if ipInfo == nil {
+		ret = model.Allow
+		retdown = model.AllowDown
+		return
+	}
+	uz := s.zoneids(ipInfo) // country, state, city
+	if ok, err = s.existsRule(c, aid); err != nil {
+		log.Error("s.existsRule error(%v)", err)
+		err = nil
+	} else if ok {
+		if rules, err = s.rule(c, aid, uz); err != nil {
+			log.Error("s.rule(%d) error(%v) ", aid, err)
+			err = nil
+		} else {
+			for _, auth = range rules {
+				retdown = 0xff & auth
+				ret = auth >> 8
+				if ret != 0 {
+					break
+				}
+			}
+			if ret == 0 {
+				ret = model.Allow
+				retdown = model.AllowDown
+			}
+			return
+		}
+	}
+	if gid, err = s.groupid(c, aid); err != nil {
+		return
+	} else if gid != 0 {
+		if pids, ok = s.groupPolicy[gid]; ok {
+			for _, pid = range pids {
+				if zids, ok = s.policy[pid]; !ok {
+					continue
+				}
+				if ret == 0 {
+					//  ret already set skip check
+					for _, zid = range uz {
+						if auth, ok = zids[zid]; ok {
+							if ret == 0 {
+								retdown = 0xff & auth
+								ret = auth >> 8 // ret must not be zero
+								break
+							}
+						}
+					}
+				}
+				tmpZids := map[int64]map[int64]int64{
+					aid: zids,
+				}
+				s.addCache(tmpZids)
+			}
+			if ret == 0 {
+				ret = model.Allow
+				retdown = model.AllowDown
+			}
+			return
+		}
+	}
+	ret = model.Allow
+	retdown = model.AllowDown
+	zids = make(map[int64]int64)
+	zids[0] = ret<<8 | retdown
+	tmpZids := map[int64]map[int64]int64{
+		aid: zids,
+	}
+	s.addCache(tmpZids)
+	return
+}
+
+// Forbid check ip is forbid or not.
+func (s *Service) Forbid(c context.Context, pstr string, ipaddr string) (err error) {
+	if pstr == "" {
+		return
+	}
+	var (
+		ret  int64
+		pids []int64
+	)
+	if pids, err = xstr.SplitInts(pstr); err != nil {
+		log.Error("xstr.SplitInts(%s) error(%v)", pstr, err)
+		return
+	}
+	if ret, _ = s.FindByPid(c, pids, ipaddr); ret == model.Forbidden {
+		err = ecode.ZlimitForbidden
+	}
+	return
+}
+
+// FindByPid redio rule by policy_id and ipaddr
+func (s *Service) FindByPid(c context.Context, pids []int64, ipaddr string) (ret, retdown int64) {
+	var (
+		ok      bool
+		auth    int64
+		zoneids []int64
+	)
+	ret = model.Allow
+	retdown = model.AllowDown
+	ipInfo := s.list.Zone(ipaddr)
+	if ipInfo == nil {
+		return
+	}
+	zoneids = s.zoneids(ipInfo)
+	for _, pid := range pids {
+		if _, ok = s.policy[pid]; !ok {
+			continue
+		}
+		for _, zoneid := range zoneids {
+			if auth, ok = s.policy[pid][zoneid]; ok {
+				retdown = 0xff & auth
+				ret = auth >> 8
+				break
+			}
+		}
+	}
+	return
+}
+
+// FindByGid redio rule by group_id and ipaddr(or cdnip)
+func (s *Service) FindByGid(c context.Context, gid int64, ipaddr, cdnip string) (ret, retdown int64) {
+	var ipInfo *xip.Zone
+	ret = model.Allow
+	retdown = model.AllowDown
+	ipInfo = s.list.Zone(ipaddr)
+	if (ipInfo != nil) && (ipInfo.Province == "共享地址" || ipInfo.City == "共享地址") && cdnip != "" {
+		ipInfo = s.list.Zone(cdnip)
+	}
+	if ipInfo == nil {
+		return
+	}
+	zoneids := s.zoneids(ipInfo)
+	if pids, ok := s.groupPolicy[gid]; ok {
+		for _, pid := range pids {
+			if _, ok := s.policy[pid]; !ok {
+				continue
+			}
+			for _, zoneid := range zoneids {
+				if auth, ok := s.policy[pid][zoneid]; ok {
+					retdown = 0xff & auth
+					ret = auth >> 8
+					break
+				}
+			}
+		}
+	}
+	return
+}
+
+// zoneids make zoneids
+func (s *Service) zoneids(ipinfos *xip.Zone) []int64 {
+	cZid := xip.ZoneID(ipinfos.Country, "", "")
+	cpZid := xip.ZoneID(ipinfos.Country, ipinfos.Province, "")
+	cpcZid := xip.ZoneID(ipinfos.Country, ipinfos.Province, ipinfos.City)
+	zoneids := []int64{0, cZid, cpZid, cpcZid}
+	return zoneids
+}
+
+func (s *Service) addCache(d interface{}) {
+	// asynchronous add rules to redis
+	select {
+	case s.missch <- d:
+	default:
+		log.Warn("cacheproc chan full")
+	}
+}
+
+// cacheproc is a routine for add rules into redis.
+func (s *Service) cacheproc() {
+	for {
+		d := <-s.missch
+		switch d.(type) {
+		case map[int64]map[int64]int64:
+			v := d.(map[int64]map[int64]int64)
+			if err := s.addRule(context.TODO(), v); err != nil {
+				log.Error("s.addRule error(%v) error(%v)", v, err)
+			}
+		default:
+			log.Warn("cacheproc can't process the type")
+		}
+	}
+}

app/service/main/location/zlimit/zlimit_test.go

@@ -0,0 +1,71 @@
+package zlimit
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"go-common/library/cache/redis"
+	"go-common/library/container/pool"
+	"go-common/library/database/sql"
+	xtime "go-common/library/time"
+)
+
+func TestZlimit(t *testing.T) {
+	c := &Config{
+		DB: &sql.Config{
+			Addr:   "172.16.33.54:3306",
+			DSN:    "test:test@tcp(172.16.33.54:3306)/bilibili_oversea?timeout=5s&readTimeout=5s&writeTimeout=5s&parseTime=true&loc=Local&charset=utf8,utf8mb4",
+			Active: 5,
+			Idle:   2,
+		},
+		Redis: &Redis{
+			Config: &redis.Config{
+				Config: &pool.Config{
+					Active:      20,
+					Idle:        10,
+					IdleTimeout: xtime.Duration(80 * time.Millisecond),
+				},
+				Name:         "business/zlimit",
+				Proto:        "tcp",
+				Addr:         "172.16.33.54:6380",
+				DialTimeout:  xtime.Duration(50 * time.Millisecond),
+				ReadTimeout:  xtime.Duration(100 * time.Millisecond),
+				WriteTimeout: xtime.Duration(100 * time.Millisecond),
+			},
+			Expire: xtime.Duration(1 * time.Second),
+		},
+		IPFile:    "/data/conf/iprepo/iprepo.txt",
+		FlushTime: xtime.Duration(1 * time.Second),
+	}
+	s := New(c)
+	testFind(t, s, 0, "", "")
+	testForbid(t, s, "", "")
+	testFindByPid(t, s, []int64{0}, "")
+	testFindByGid(t, s, 0, "", "")
+}
+
+func testFind(t *testing.T, s *Service, aid int64, ipaddr, cdnip string) {
+	ret, retdown, err := s.Find(context.TODO(), aid, ipaddr, cdnip)
+	if err != nil {
+		t.Errorf("Service: find err: %v", err)
+	} else {
+		t.Logf("Service: find %d,%d", ret, retdown)
+	}
+}
+
+func testForbid(t *testing.T, s *Service, pstr, ipaddr string) {
+	if err := s.Forbid(context.TODO(), pstr, ipaddr); err != nil {
+		t.Errorf("Service: forbid err: %v", err)
+	}
+}
+
+func testFindByPid(t *testing.T, s *Service, pids []int64, ipaddr string) {
+	ret, retdown := s.FindByPid(context.TODO(), pids, ipaddr)
+	t.Logf("Service: findByPid %d,%d", ret, retdown)
+}
+
+func testFindByGid(t *testing.T, s *Service, gid int64, ipaddr, cdnip string) {
+	ret, retdown := s.FindByGid(context.TODO(), gid, ipaddr, cdnip)
+	t.Logf("Service: findByGid %d,%d", ret, retdown)
+}