JCS-pub/coordinator/internal/repl/user.go

package repl

import (
	"context"
	"encoding/hex"
	"fmt"
	"os"

	"github.com/spf13/cobra"
	"gitlink.org.cn/cloudream/common/pkgs/logger"
	stgglb "gitlink.org.cn/cloudream/jcs-pub/common/globals"
	hubrpc "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/rpc/hub"
	jcstypes "gitlink.org.cn/cloudream/jcs-pub/common/types"
	"gitlink.org.cn/cloudream/jcs-pub/coordinator/internal/accesstoken"
	"gitlink.org.cn/cloudream/jcs-pub/coordinator/internal/db"
	"golang.org/x/crypto/bcrypt"
	"golang.org/x/term"
	"gorm.io/gorm"
)

func init() {
	userCmd := &cobra.Command{
		Use:   "user",
		Short: "user command",
	}
	RootCmd.AddCommand(userCmd)

	createCmd := &cobra.Command{
		Use:   "create [account] [nickName]",
		Short: "create a new user account",
		Args:  cobra.ExactArgs(2),
		Run: func(cmd *cobra.Command, args []string) {
			userCreate(GetCmdCtx(cmd), args[0], args[1])
		},
	}
	userCmd.AddCommand(createCmd)

	logoutCmd := &cobra.Command{
		Use:   "logout [account] [tokenID]",
		Short: "logout from a user account",
		Args:  cobra.ExactArgs(2),
		Run: func(cmd *cobra.Command, args []string) {
			userLogout(GetCmdCtx(cmd), args[0], jcstypes.AccessTokenID(args[1]))
		},
	}
	userCmd.AddCommand(logoutCmd)
}

func userCreate(ctx *CommandContext, account string, nickName string) {
	_, err := ctx.repl.db.User().GetByAccount(ctx.repl.db.DefCtx(), account)
	if err == nil {
		fmt.Printf("user %s already exists\n", account)
		return
	}

	fmt.Printf("input account password: ")
	pass, err := term.ReadPassword(int(os.Stdin.Fd()))
	if err != nil {
		fmt.Println("error reading password:", err)
		return
	}

	passHash, err := bcrypt.GenerateFromPassword(pass, bcrypt.DefaultCost)
	if err != nil {
		fmt.Println("error hashing password:", err)
		return
	}

	user, err := db.DoTx01(ctx.repl.db, func(tx db.SQLContext) (jcstypes.User, error) {
		return ctx.repl.db.User().Create(tx, account, hex.EncodeToString(passHash), nickName)
	})
	if err != nil {
		fmt.Println("error creating user:", err)
		return
	}

	fmt.Printf("user %s created\n", user.Account)
}

func userLogout(ctx *CommandContext, account string, tokenID jcstypes.AccessTokenID) {
	acc, err := ctx.repl.db.User().GetByAccount(ctx.repl.db.DefCtx(), account)
	if err != nil {
		fmt.Printf("user %s not found\n", account)
		return
	}

	log := logger.WithField("UserID", acc.UserID).WithField("TokenID", tokenID)

	d := ctx.repl.db
	loaded, err := db.DoTx01(d, func(tx db.SQLContext) ([]jcstypes.LoadedAccessToken, error) {
		token, err := d.UserAccessToken().GetByID(tx, acc.UserID, tokenID)
		if err != nil {
			return nil, err
		}

		err = d.UserAccessToken().DeleteByID(tx, token.UserID, token.TokenID)
		if err != nil {
			return nil, err
		}

		loaded, err := d.LoadedAccessToken().GetByUserIDAndTokenID(tx, token.UserID, token.TokenID)
		if err != nil {
			return nil, err
		}

		err = d.LoadedAccessToken().DeleteAllByUserIDAndTokenID(tx, token.UserID, token.TokenID)
		if err != nil {
			return nil, err
		}

		return loaded, nil
	})
	if err != nil {
		log.Warnf("delete access token: %v", err)
		if err == gorm.ErrRecordNotFound {
			return
		}

		return
	}

	ctx.repl.accessToken.NotifyTokenInvalid(accesstoken.CacheKey{
		UserID:  acc.UserID,
		TokenID: tokenID,
	})

	var loadedHubIDs []jcstypes.HubID
	for _, l := range loaded {
		loadedHubIDs = append(loadedHubIDs, l.HubID)
	}

	notifyLoadedHubs(ctx, acc.UserID, tokenID, loadedHubIDs)
}

func notifyLoadedHubs(ctx *CommandContext, userID jcstypes.UserID, tokenID jcstypes.AccessTokenID, loadedHubIDs []jcstypes.HubID) {
	log := logger.WithField("UserID", userID).WithField("TokenID", tokenID)

	d := ctx.repl.db

	loadedHubs, err := d.Hub().BatchGetByID(d.DefCtx(), loadedHubIDs)
	if err != nil {
		log.Warnf("getting hubs: %v", err)
		return
	}

	for _, l := range loadedHubs {
		addr, ok := l.Address.(*jcstypes.GRPCAddressInfo)
		if !ok {
			continue
		}

		cli := stgglb.HubRPCPool.Get(addr.ExternalIP, addr.ExternalGRPCPort)
		// 不关心返回值
		cli.NotifyUserAccessTokenInvalid(context.Background(), &hubrpc.NotifyUserAccessTokenInvalid{
			UserID:  userID,
			TokenID: tokenID,
		})
		cli.Release()
	}
}