ai_scheduler/internal/pkg/lsxd/login.go

package lsxd

import (
	"ai_scheduler/internal/config"
	"ai_scheduler/internal/data/constants"
	"ai_scheduler/utils"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/go-kratos/kratos/v2/log"
	"github.com/redis/go-redis/v9"
)

type Login struct {
	config   *config.Config
	redisCli *redis.Client
	mu       sync.Mutex
}

func NewLogin(config *config.Config, rdb *utils.Rdb) *Login {
	return &Login{
		config:   config,
		redisCli: rdb.Rdb,
	}
}

func (l *Login) GetToken(ctx context.Context) string {
	// ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	// defer cancel()

	token, err := l.getCachedToken(ctx)
	if err != nil {
		log.Errorf("lsxd get token from redis failed, err: %v", err)
	}
	if token != "" && l.checkTokenValid(ctx, token) {
		return token
	}

	l.mu.Lock()
	defer l.mu.Unlock()

	token, err = l.getCachedToken(ctx)
	if err != nil {
		log.Errorf("lsxd get token from redis failed, err: %v", err)
	}
	if token != "" && l.checkTokenValid(ctx, token) {
		return token
	}

	token, err = l.login(ctx)
	if err != nil {
		log.Errorf("lsxd login failed, err: %v", err)
		return ""
	}
	if token == "" {
		log.Errorf("lsxd login failed, token is empty")
		return ""
	}

	if err := l.cacheToken(ctx, token); err != nil {
		log.Errorf("lsxd cache token failed, err: %v", err)
	}

	return token
}

// 校验token是否有效
func (l *Login) checkTokenValid(ctx context.Context, token string) bool {
	// 欢迎页校验token有效
	checkTokenURL := l.config.LSXD.CheckTokenURL
	if checkTokenURL == "" {
		return token != ""
	}

	status, err := l.doRequest(ctx, http.MethodGet, checkTokenURL, token, nil)
	if err != nil {
		log.Errorf("lsxd check token valid failed, err: %v", err)
		return true
	}
	if status == http.StatusOK {
		return true
	}
	if status == http.StatusUnauthorized || status == http.StatusForbidden {
		return false
	}
	return true
}

// 调用登录接口获取token
func (l *Login) login(ctx context.Context) (string, error) {
	// 1.获取配置
	loginURL := l.config.LSXD.LoginURL
	phone := l.config.LSXD.Phone
	password := l.config.LSXD.Password

	// 2.调用登录接口获取token
	if loginURL == "" {
		return "", errors.New("login url is empty")
	}
	if phone == "" || password == "" {
		return "", errors.New("phone or password is empty")
	}

	reqBody := map[string]any{
		"phone":    phone,
		"password": password,
		"code":     "456789",
	}
	bodyBytes, err := json.Marshal(reqBody)
	if err != nil {
		return "", err
	}

	status, respBody, err := l.doRequestWithBody(ctx, http.MethodPost, loginURL, "", "application/json", bodyBytes)
	if err != nil {
		return "", err
	}
	if status != http.StatusOK {
		return "", fmt.Errorf("login status code: %d", status)
	}

	type loginResp struct {
		Code        int    `json:"code"`
		Msg         string `json:"msg"`
		Message     string `json:"message"`
		AccessToken string `json:"accessToken"`
		Data        struct {
			AccessToken string `json:"accessToken"`
			Token       string `json:"token"`
		} `json:"data"`
	}

	var resp loginResp
	if err := json.Unmarshal(respBody, &resp); err != nil {
		return "", err
	}

	token := resp.AccessToken
	if token == "" {
		token = resp.Data.AccessToken
	}
	if token == "" {
		token = resp.Data.Token
	}

	if token == "" {
		return "", errors.New("token is empty")
	}

	return token, nil
}

func (l *Login) getCachedToken(ctx context.Context) (string, error) {
	token, err := l.redisCli.Get(ctx, constants.CACHE_KEY_LSXD_TOKEN).Result()
	if err == nil {
		return token, nil
	}
	if errors.Is(err, redis.Nil) {
		return "", nil
	}
	return "", err
}

func (l *Login) cacheToken(ctx context.Context, token string) error {
	if token == "" {
		return errors.New("token is empty")
	}
	return l.redisCli.Set(ctx, constants.CACHE_KEY_LSXD_TOKEN, token, constants.EXPIRE_LSXD_TOKEN).Err()
}

func (l *Login) doRequest(ctx context.Context, method string, url string, authorization string, body []byte) (int, error) {
	status, _, err := l.doRequestWithBody(ctx, method, url, authorization, "", body)
	return status, err
}

func (l *Login) doRequestWithBody(ctx context.Context, method string, url string, authorization string, contentType string, body []byte) (int, []byte, error) {
	reqCtx, cancel := context.WithTimeout(ctx, 6*time.Second)
	defer cancel()

	var reader io.Reader
	if len(body) > 0 {
		reader = bytes.NewReader(body)
	}

	req, err := http.NewRequestWithContext(reqCtx, method, url, reader)
	if err != nil {
		return 0, nil, err
	}
	if contentType != "" {
		req.Header.Set("Content-Type", contentType)
	}
	if authorization != "" {
		req.Header.Set("Authorization", authorization)
	}

	resp, err := (&http.Client{}).Do(req)
	if err != nil {
		return 0, nil, err
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return resp.StatusCode, nil, err
	}
	return resp.StatusCode, respBody, nil
}