修改目录结构

db/db.go

@@ -2,10 +2,10 @@ package db
 
 import (
 	"fmt"
-	"log"
 	"sync"
 
 	"github.com/astaxie/beego/orm"
+	beeLogger "github.com/beego/bee/logger"
 	"github.com/go-redis/redis/v8"
 
 	_ "github.com/go-sql-driver/mysql"
@@ -44,6 +44,7 @@ func newRedisClient() {
 		Password: redis_pass, // no password set
 		DB:       redis_db,   // use default DB
 	})
+	beeLogger.Log.Success("Redis connect success")
 }
 
 func newMysqlClient() {
@@ -52,12 +53,11 @@ func newMysqlClient() {
 	_ = orm.RegisterDriver(mysql_driver, orm.DRMySQL)
 	err := orm.RegisterDataBase("default", mysql_driver, dbConfStr, 10, 100)
 	if err != nil {
-		log.Println("db connect error => ", err)
+		beeLogger.Log.Fatalf("db connect error => %v", err)
 		return
 	}
-
 	orm.RegisterModel(new(CompanyApp), new(ImBase), new(ImChatRoom), new(ImUserRelationship), new(ImUser))
 	orm.Debug = true
 	MysqlClient = orm.NewOrm()
-	log.Println("MySQL connect success")
+	beeLogger.Log.Success("MySQL connect success")
 }

go.mod

@@ -4,6 +4,7 @@ go 1.12
 
 require (
 	github.com/astaxie/beego v1.12.1
+	github.com/beego/bee v1.11.1
 	github.com/go-playground/universal-translator v0.17.0 // indirect
 	github.com/go-redis/redis/v8 v8.0.0-beta.5
 	github.com/go-sql-driver/mysql v1.5.0
@@ -13,6 +14,7 @@ require (
 	github.com/rs/xid v1.2.1
 	github.com/segmentio/ksuid v1.0.2
 	golang.org/x/net v0.0.0-20191021144547-ec77196f6094
+	golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e
 	google.golang.org/grpc v1.30.0
 	google.golang.org/protobuf v1.24.0
 	gopkg.in/go-playground/validator.v9 v9.31.0

helper/auth.go

+package helper
+
+import (
+	"context"
+	"crypto/sha1"
+	"encoding/hex"
+	"fmt"
+	"regexp"
+	"strconv"
+	"time"
+
+	"golang.org/x/time/rate"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+const (
+	health_path = "/pb.Health/Check"
+)
+
+func permissionError(msg string) error {
+	return status.Error(codes.InvalidArgument, msg)
+}
+
+func checksum(req map[string]string) error {
+
+	err, AppSecret := GetSecretByKey(req["Appkey"])
+	if err != nil {
+		return status.Error(codes.NotFound, "appkey不存在")
+	}
+
+	server_time := time.Now().Unix()
+	client_time, err := strconv.Atoi(req["Curtime"])
+	if err != nil {
+		return status.Error(codes.Aborted, err.Error())
+	}
+	if server_time > int64(client_time) {
+		return permissionError("当前客户端时间不能小于服务端时间")
+	}
+
+	sha11 := sha1.New()
+	sha11.Write([]byte(AppSecret + req["Nonce"] + req["Curtime"]))
+	sevice_checksum := hex.EncodeToString(sha11.Sum([]byte(nil)))
+
+	if sevice_checksum != req["Checksum"] {
+		return permissionError("加密串不正确")
+	}
+	return nil
+}
+
+func Auth(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
+	handler grpc.UnaryHandler) (interface{}, error) {
+
+	if info.FullMethod != health_path {
+		limiter := rate.NewLimiter(rate.Every(100*time.Millisecond), 1)
+		if limiter.Allow() {
+			request_map := make(map[string]string)
+			s := fmt.Sprintf("%v", req)
+
+			key_r, _ := regexp.Compile(`Appkey:"(.*?)"`)
+			nonce_r, _ := regexp.Compile(`Nonce:"(.*?)"`)
+			c_r, _ := regexp.Compile(`Curtime:"(.*?)"`)
+			cs_r, _ := regexp.Compile(`Checksum:"(.*?)"`)
+
+			appkey := key_r.FindAllStringSubmatch(s, -1)
+			nonce := nonce_r.FindAllStringSubmatch(s, -1)
+			curtime := c_r.FindAllStringSubmatch(s, -1)
+			check_sum := cs_r.FindAllStringSubmatch(s, -1)
+
+			if len(appkey) < 1 {
+				return nil, permissionError("缺少appkey")
+			}
+			if len(nonce) < 1 {
+				return nil, permissionError("缺少nonce")
+			}
+			if len(curtime) < 1 {
+				return nil, permissionError("缺少curtime")
+			}
+			if len(check_sum) < 1 {
+				return nil, permissionError("缺少checksum")
+			}
+			request_map["Appkey"] = appkey[0][1]
+			request_map["Nonce"] = nonce[0][1]
+			request_map["Curtime"] = curtime[0][1]
+			request_map["Checksum"] = check_sum[0][1]
+			if err := checksum(request_map); err != nil {
+				return nil, err
+			}
+		} else {
+			return nil, status.Error(codes.Unauthenticated, "请求超过限制,请稍后再试")
+		}
+
+	}
+	return handler(ctx, req)
+}

initialization/Registration_services.go

+package initialization
+
+import (
+	"im-microservice/helper"
+	"im-microservice/pb"
+	"im-microservice/sevice/health"
+	"im-microservice/sevice/im_chat_room"
+	ic "im-microservice/sevice/im_configure"
+	"im-microservice/sevice/im_user"
+	iur "im-microservice/sevice/im_user_relationship"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/reflection"
+)
+
+var RpcServer *grpc.Server
+
+func init() {
+	var opts []grpc.ServerOption
+	opts = append(opts, grpc.UnaryInterceptor(helper.Auth))
+	s := grpc.NewServer(opts...)
+	srv := health.NewServer()
+	pb.RegisterHealthServer(s, srv)
+	pb.RegisterConfigureSeviceServer(s, &ic.ConfigureSevice{})
+	pb.RegisterChatRoomServiceServer(s, &im_chat_room.ImChatRoomService{})
+	pb.RegisterImUserServer(s, &im_user.ImUserServer{})
+	pb.RegisterUserRelationshipServiceServer(s, &iur.UserRelationshipService{})
+	reflection.Register(s)
+	RpcServer = s
+}

main.go

 package main
 
 import (
-	"crypto/sha1"
-	"encoding/hex"
-	"fmt"
-	"log"
 	"net"
-	"regexp"
-	"strconv"
-	"sync"
-	"time"
 
-	"im-microservice/helper"
-	"im-microservice/pb"
-	"im-microservice/sevice/im_chat_room"
-	ic "im-microservice/sevice/im_configure"
-	"im-microservice/sevice/im_user"
-	iur "im-microservice/sevice/im_user_relationship"
-
-	"golang.org/x/net/context"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/reflection"
-	"google.golang.org/grpc/status"
+	beeLogger "github.com/beego/bee/logger"
+	initialize "im-microservice/initialization"
 )
 
-type Health interface {
-	// Check returns if server is healthy or not
-	Check(c context.Context) (bool, error)
-}
-
-// 健康检测
-type Server struct {
-	mu        sync.Mutex
-	statusMap map[string]pb.HealthCheckResponse_ServingStatus
-}
-
-func NewServer() *Server {
-	return &Server{
-		statusMap: make(map[string]pb.HealthCheckResponse_ServingStatus),
-	}
-}
-func permissionError(msg string) error {
-	return status.Error(codes.InvalidArgument, msg)
-}
-
-func (s *Server) Check(ctx context.Context, in *pb.HealthCheckRequest) (*pb.HealthCheckResponse, error) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	if in.Service == "" {
-		return &pb.HealthCheckResponse{Status: pb.HealthCheckResponse_SERVING}, nil
-	}
-	if status, ok := s.statusMap[in.Service]; ok {
-		return &pb.HealthCheckResponse{Status: status}, nil
-	}
-	return nil, status.Error(codes.NotFound, "unkonw service")
-}
-
 const (
 	port = ":50051"
-	health_path = "/pb.Health/Check"
 )
 
-func checksum(req map[string]string) error {
-
-	err, AppSecret := helper.GetSecretByKey(req["Appkey"])
-	if err != nil {
-		return status.Error(codes.NotFound, "appkey不存在")
-	}
-
-	server_time := time.Now().Unix()
-	client_time, err := strconv.Atoi(req["Curtime"])
-	if err != nil {
-		return status.Error(codes.Aborted, err.Error())
-	}
-	if server_time > int64(client_time) {
-		return permissionError("当前客户端时间不能小于服务端时间")
-	}
-
-	sha11 := sha1.New()
-	sha11.Write([]byte(AppSecret + req["Nonce"] + req["Curtime"]))
-	sevice_checksum := hex.EncodeToString(sha11.Sum([]byte(nil)))
-
-	if sevice_checksum != req["Checksum"] {
-		return permissionError("加密串不正确")
-	}
-	return nil
-}
-
-func auth(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
-
-	if info.FullMethod != health_path {
-		request_map := make(map[string]string)
-		s := fmt.Sprintf("%v", req)
-
-		key_r, _ := regexp.Compile(`Appkey:"(.*?)"`)
-		nonce_r, _ := regexp.Compile(`Nonce:"(.*?)"`)
-		c_r, _ := regexp.Compile(`Curtime:"(.*?)"`)
-		cs_r, _ := regexp.Compile(`Checksum:"(.*?)"`)
-
-		appkey := key_r.FindAllStringSubmatch(s, -1)
-		nonce := nonce_r.FindAllStringSubmatch(s, -1)
-		curtime := c_r.FindAllStringSubmatch(s, -1)
-		check_sum := cs_r.FindAllStringSubmatch(s, -1)
-
-		if len(appkey) < 1 {
-			return nil, permissionError("缺少appkey")
-		}
-		if len(nonce) < 1 {
-			return nil, permissionError("缺少nonce")
-		}
-		if len(curtime) < 1 {
-			return nil, permissionError("缺少curtime")
-		}
-		if len(check_sum) < 1 {
-			return nil, permissionError("缺少checksum")
-		}
-		request_map["Appkey"] = appkey[0][1]
-		request_map["Nonce"] = nonce[0][1]
-		request_map["Curtime"] = curtime[0][1]
-		request_map["Checksum"] = check_sum[0][1]
-		if err := checksum(request_map); err != nil {
-			return nil, err
-		}
-	}
-	return handler(ctx, req)
-}
-
 func main() {
 
 	lis, err := net.Listen("tcp", port)
 	if err != nil {
-		log.Fatalf("failed to listen: %v", err)
+		beeLogger.Log.Fatalf("failed to listen: %v", err)
 	}
 
-	var opts []grpc.ServerOption
-	opts = append(opts, grpc.UnaryInterceptor(auth))
-	s := grpc.NewServer(opts...)
-
-	srv := NewServer()
-	pb.RegisterHealthServer(s, srv)
-	pb.RegisterConfigureSeviceServer(s, &ic.ConfigureSevice{})
-	pb.RegisterChatRoomServiceServer(s, &im_chat_room.ImChatRoomService{})
-	pb.RegisterImUserServer(s, &im_user.ImUserServer{})
-	pb.RegisterUserRelationshipServiceServer(s, &iur.UserRelationshipService{})
-	reflection.Register(s)
-	
-	log.Println("gRPC server is running on " + port + " port.")
-	if err := s.Serve(lis); err != nil {
-		log.Fatalf("failed to serve: %v", err)
+	beeLogger.Log.Success("gRPC server is running on " + port + " port.")
+	if err := initialize.RpcServer.Serve(lis); err != nil {
+		beeLogger.Log.Fatalf("failed to serve: %v", err)
 	}
 
 }

sevice/health/main.go

+package health
+
+import (
+	"context"
+	"sync"
+
+	"im-microservice/pb"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+type Health interface {
+	// Check returns if server is healthy or not
+	Check(c context.Context) (bool, error)
+}
+
+// 健康检测
+type Server struct {
+	mu        sync.Mutex
+	statusMap map[string]pb.HealthCheckResponse_ServingStatus
+}
+
+func (s *Server) Check(ctx context.Context, in *pb.HealthCheckRequest) (*pb.HealthCheckResponse, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if in.Service == "" {
+		return &pb.HealthCheckResponse{Status: pb.HealthCheckResponse_SERVING}, nil
+	}
+	if status, ok := s.statusMap[in.Service]; ok {
+		return &pb.HealthCheckResponse{Status: status}, nil
+	}
+	return nil, status.Error(codes.NotFound, "unkonw service")
+}
+
+func NewServer() *Server {
+	return &Server{
+		statusMap: make(map[string]pb.HealthCheckResponse_ServingStatus),
+	}
+}