sonet/internal/auth/logic/auth_server.go

package logic

import (
	"context"
	"encoding/base64"
	"errors"
	"github.com/bytedance/sonic"
	"google.golang.org/grpc"
	"google.golang.org/grpc/reflection"
	"net"
	"sonet/api/gen/auth"
	"sonet/internal/auth/data"
	"sonet/pkg/config"
	"sonet/pkg/grpc/discovery"
	"sonet/pkg/utils/logger"
	"sonet/pkg/utils/security"
	"time"
)

type AuthServer struct {
	auth.UnimplementedAuthServer
	aesTokenKey []byte
	userDao     *data.AuthUserDao
}

func NewAuthServer(aesTokenKey string, userDao *data.AuthUserDao) *AuthServer {
	keyBytes, err := base64.StdEncoding.DecodeString(aesTokenKey)
	if err != nil {
		panic(err)
	}
	return &AuthServer{
		aesTokenKey: keyBytes,
		userDao:     userDao,
	}
}

func (s *AuthServer) Run(conf config.GrpcConfig, register *discovery.Register) (err error) {
	server := grpc.NewServer(
		config.GetGrpcOptions(conf)...,
	)
	if !conf.NoReflection {
		// 注册反射服务
		reflection.Register(server)
	}
	auth.RegisterAuthServer(server, s)
	listen, err := net.Listen("tcp", conf.Address)
	if err != nil {
		return
	}

	// registry discovery
	reg := conf.Register
	if reg.Name == "" {
		reg.Name = auth.Auth_ServiceDesc.ServiceName
	}
	if reg.Addr == "" {
		reg.Addr, err = discovery.RegisterAddress(conf.Address)
		if err != nil {
			return
		}
	}
	if err = register.Register(reg); err != nil {
		return
	}

	// run serve
	logger.Infof("%s grpc server running %s\n", reg.Name, listen.Addr().String())
	err = server.Serve(listen)
	return
}

func (s *AuthServer) Login(ctx context.Context, req *auth.ReqLogin) (*auth.ResLogin, error) {
	user, err := s.userDao.FindByAccount(req.Account)
	if err != nil {
		return nil, err
	}
	if user == nil {
		// 不存在注册
		now := time.Now()
		uid, err := s.userDao.NextUid()
		if err != nil {
			return nil, err
		}
		user = &data.User{
			Uid:      uid,
			Account:  req.Account,
			Username: req.Account,
			Password: req.Password,
			CreateAt: &now,
			UpdateAt: &now,
		}
		err = s.userDao.Create(user)
		if err != nil {
			return nil, err
		}
		// return nil, errors.New("not found account " + req.Account)
	}
	// verify password login
	if req.Password != user.Password {
		return nil, errors.New("account or password error")
	}

	// generate token
	subject := &auth.Subject{Uid: user.Uid, Username: user.Username, Time: time.Now().UnixMilli()}
	bytes, err := sonic.Marshal(subject)
	if err != nil {
		return nil, err
	}
	encode, err := security.EncryptAesCBC(bytes, s.aesTokenKey)
	if err != nil {
		return nil, err
	}
	token := base64.URLEncoding.EncodeToString(encode)

	res := &auth.ResLogin{Token: token, Subject: subject}
	return res, nil
}

func (s *AuthServer) Verify(ctx context.Context, req *auth.ReqVerify) (*auth.Subject, error) {
	bytes, err := base64.URLEncoding.DecodeString(req.Token)
	if err != nil {
		return nil, errors.New("token decode fail: " + err.Error())
	}
	decode, err := security.DecryptAesCBC(bytes, s.aesTokenKey)
	if err != nil {
		return nil, errors.New("invalidate token: " + err.Error())
	}
	subject := &auth.Subject{}
	err = sonic.Unmarshal(decode, subject)
	if err != nil {
		return nil, errors.New("token payload decode fail: " + err.Error())
	}
	return subject, nil
}

func (s *AuthServer) Registry(ctx context.Context, req *auth.ReqRegistry) (*auth.Subject, error) {
	return nil, errors.New("can not registry")
}

func (s *AuthServer) FindByUid(ctx context.Context, req *auth.ReqFindByUid) (subject *auth.Subject, err error) {
	sub, err := s.userDao.FindByUid(req.Uid)
	if err != nil || sub == nil {
		return
	}

	subject = &auth.Subject{
		Uid:      sub.Uid,
		Username: sub.Username,
		Time:     time.Now().UnixMilli(),
		Extra:    map[string]string{"avatar": sub.Avatar},
	}
	return
}