junhong_cmp_fiber/internal/infrastructure/wecom/token_provider.go

package wecom

import (
	"context"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/bytedance/sonic"
	"github.com/google/uuid"
	"github.com/redis/go-redis/v9"
	"go.uber.org/zap"

	"github.com/break/junhong_cmp_fiber/internal/infrastructure/integrationlog"
	"github.com/break/junhong_cmp_fiber/internal/model"
	"github.com/break/junhong_cmp_fiber/pkg/constants"
	"github.com/break/junhong_cmp_fiber/pkg/errors"
	"github.com/break/junhong_cmp_fiber/pkg/middleware"
)

const releaseTokenLockScript = `
if redis.call("GET", KEYS[1]) == ARGV[1] then
    return redis.call("DEL", KEYS[1])
end
return 0
`

// TokenApplicationRepository 定义 access_token Provider 所需的应用配置查询边界。
type TokenApplicationRepository interface {
	GetEnabled(ctx context.Context, applicationID uint) (*model.WeComApplication, error)
	MarkConnected(ctx context.Context, applicationID uint, connectedAt time.Time) error
}

// TokenIntegrationLog 定义企微外呼的 Integration Log 边界。
type TokenIntegrationLog interface {
	Start(ctx context.Context, input integrationlog.Attempt) (*model.IntegrationLog, error)
	Complete(ctx context.Context, integrationID string, completion integrationlog.Completion) (*model.IntegrationLog, error)
}

// TokenProvider 按应用缓存企业微信 access_token，并用 Redis 短锁抑制并发回源。
type TokenProvider struct {
	repo        TokenApplicationRepository
	redis       *redis.Client
	integration TokenIntegrationLog
	httpClient  *http.Client
	baseURL     string
	logger      *zap.Logger
	now         func() time.Time
}

// NewTokenProvider 创建企业微信 access_token Provider。
func NewTokenProvider(
	repo TokenApplicationRepository,
	redisClient *redis.Client,
	integration TokenIntegrationLog,
	baseURL string,
	timeout time.Duration,
	logger *zap.Logger,
) *TokenProvider {
	if timeout <= 0 {
		timeout = constants.WeComDefaultHTTPTimeout
	}
	return &TokenProvider{
		repo: repo, redis: redisClient, integration: integration,
		httpClient: &http.Client{Timeout: timeout}, baseURL: strings.TrimRight(baseURL, "/"),
		logger: logger, now: time.Now,
	}
}

// GetAccessToken 优先读取应用缓存，未命中时串行调用企业微信 Token 接口。
func (p *TokenProvider) GetAccessToken(ctx context.Context, applicationID uint) (string, error) {
	if p == nil || p.repo == nil || p.integration == nil || p.httpClient == nil || applicationID == 0 {
		return "", errors.New(errors.CodeWeComCredentialInvalid)
	}
	cacheKey := constants.RedisWeComAccessTokenKey(applicationID)
	if token, found, err := p.getCachedToken(ctx, cacheKey); err == nil && found {
		return token, nil
	} else if err != nil {
		p.warnRedis("读取企业微信 access_token 缓存失败，改为受控直连", err, applicationID)
		return p.fetchAndCache(ctx, applicationID, cacheKey)
	}
	if p.redis == nil {
		return p.fetchAndCache(ctx, applicationID, cacheKey)
	}
	lockKey := constants.RedisWeComAccessTokenLockKey(applicationID)
	lockValue := uuid.NewString()
	acquired, err := p.redis.SetNX(ctx, lockKey, lockValue, constants.WeComTokenLockTTL).Result()
	if err != nil {
		p.warnRedis("获取企业微信 access_token 回源锁失败，改为受控直连", err, applicationID)
		return p.fetchAndCache(ctx, applicationID, cacheKey)
	}
	if !acquired {
		return p.waitForToken(ctx, applicationID, cacheKey)
	}
	defer p.releaseLock(ctx, lockKey, lockValue, applicationID)
	if token, found, err := p.getCachedToken(ctx, cacheKey); err == nil && found {
		return token, nil
	}
	return p.fetchAndCache(ctx, applicationID, cacheKey)
}

// Invalidate 删除指定应用的 access_token 缓存。
func (p *TokenProvider) Invalidate(ctx context.Context, applicationID uint) {
	if p == nil || p.redis == nil || applicationID == 0 {
		return
	}
	if err := p.redis.Del(ctx, constants.RedisWeComAccessTokenKey(applicationID)).Err(); err != nil {
		p.warnRedis("删除企业微信 access_token 缓存失败", err, applicationID)
	}
}

func (p *TokenProvider) getCachedToken(ctx context.Context, cacheKey string) (string, bool, error) {
	if p.redis == nil {
		return "", false, nil
	}
	token, err := p.redis.Get(ctx, cacheKey).Result()
	if err == redis.Nil {
		return "", false, nil
	}
	if err != nil {
		return "", false, err
	}
	return token, token != "", nil
}

func (p *TokenProvider) waitForToken(ctx context.Context, applicationID uint, cacheKey string) (string, error) {
	timer := time.NewTimer(constants.WeComTokenWaitTimeout)
	defer timer.Stop()
	ticker := time.NewTicker(constants.WeComTokenWaitPollInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return "", errors.Wrap(errors.CodeServiceUnavailable, ctx.Err(), "等待企业微信 access_token 已取消")
		case <-timer.C:
			return "", errors.New(errors.CodeServiceUnavailable, "企业微信 access_token 正在刷新，请稍后重试")
		case <-ticker.C:
			token, found, err := p.getCachedToken(ctx, cacheKey)
			if err != nil {
				p.warnRedis("等待企业微信 access_token 时 Redis 不可用，改为受控直连", err, applicationID)
				return p.fetchAndCache(ctx, applicationID, cacheKey)
			}
			if found {
				return token, nil
			}
		}
	}
}

func (p *TokenProvider) fetchAndCache(ctx context.Context, applicationID uint, cacheKey string) (string, error) {
	application, err := p.repo.GetEnabled(ctx, applicationID)
	if err != nil {
		return "", err
	}
	request, err := p.newTokenRequest(ctx, application.CorpID, application.Secret)
	if err != nil {
		return "", err
	}
	resourceID := strconv.FormatUint(uint64(applicationID), 10)
	requestID := middleware.GetRequestIDFromContext(ctx)
	integrationID, triggerSeries, correlationID := singleIntegrationLinkage(requestID)
	attempt, err := p.integration.Start(ctx, integrationlog.Attempt{
		IntegrationID: integrationID,
		Provider:      constants.IntegrationProviderWeCom, Direction: constants.IntegrationDirectionOutbound,
		Operation: constants.IntegrationOperationWeComAccessToken, ResourceType: constants.WeComApplicationResourceType,
		ResourceID: &resourceID, RequestSummary: map[string]any{
			"application_id": applicationID, "corp_id": application.CorpID, "agent_id": application.AgentID,
		}, RequestID: requestID, CorrelationID: correlationID, TriggerSeries: triggerSeries,
	})
	if err != nil {
		return "", err
	}
	startedAt := p.now()
	response, err := p.httpClient.Do(request)
	if err != nil {
		return "", p.completeFailed(ctx, attempt.IntegrationID, 0, "request_failed", "企业微信 Token 请求失败", startedAt)
	}
	defer response.Body.Close()
	tokenResponse, err := p.readTokenResponse(response)
	if err != nil {
		return "", p.completeFailed(ctx, attempt.IntegrationID, response.StatusCode, "invalid_response", "企业微信 Token 响应无效", startedAt)
	}
	if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices || tokenResponse.ErrCode != 0 || tokenResponse.AccessToken == "" {
		providerCode := strconv.FormatInt(tokenResponse.ErrCode, 10)
		if tokenResponse.ErrCode == 0 {
			providerCode = strconv.Itoa(response.StatusCode)
		}
		return "", p.completeFailed(ctx, attempt.IntegrationID, response.StatusCode, providerCode, tokenResponse.ErrMsg, startedAt)
	}
	ttl := time.Duration(tokenResponse.ExpiresIn)*time.Second - constants.WeComTokenRefreshAdvance
	if ttl <= 0 {
		return "", p.completeFailed(ctx, attempt.IntegrationID, response.StatusCode, "invalid_expires_in", "企业微信 access_token 有效期无效", startedAt)
	}
	if err := p.completeSuccess(ctx, attempt.IntegrationID, response.StatusCode, tokenResponse, startedAt); err != nil {
		return "", err
	}
	if p.redis != nil {
		if err := p.redis.Set(ctx, cacheKey, tokenResponse.AccessToken, ttl).Err(); err != nil {
			p.warnRedis("缓存企业微信 access_token 失败，本次继续使用直连结果", err, applicationID)
		}
	}
	if err := p.repo.MarkConnected(ctx, applicationID, p.now()); err != nil && p.logger != nil {
		p.logger.Error("记录企业微信最近连接时间失败，本次 token 仍可使用", zap.Uint("application_id", applicationID), zap.Error(err))
	}
	return tokenResponse.AccessToken, nil
}

func (p *TokenProvider) newTokenRequest(ctx context.Context, corpID, secret string) (*http.Request, error) {
	endpoint, err := url.Parse(p.baseURL + "/cgi-bin/gettoken")
	if err != nil || endpoint.Scheme == "" || endpoint.Host == "" {
		return nil, errors.New(errors.CodeWeComCredentialInvalid, "企业微信 API 地址配置无效")
	}
	query := endpoint.Query()
	query.Set("corpid", corpID)
	query.Set("corpsecret", secret)
	endpoint.RawQuery = query.Encode()
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint.String(), nil)
	if err != nil {
		return nil, errors.Wrap(errors.CodeInternalError, err, "创建企业微信 Token 请求失败")
	}
	return request, nil
}

type tokenResponse struct {
	ErrCode     int64  `json:"errcode"`
	ErrMsg      string `json:"errmsg"`
	AccessToken string `json:"access_token"`
	ExpiresIn   int64  `json:"expires_in"`
}

func (p *TokenProvider) readTokenResponse(response *http.Response) (tokenResponse, error) {
	var result tokenResponse
	body, err := io.ReadAll(io.LimitReader(response.Body, constants.WeComMaxResponseBodyBytes+1))
	if err != nil {
		return result, errors.Wrap(errors.CodeServiceUnavailable, err, "读取企业微信 Token 响应失败")
	}
	if int64(len(body)) > constants.WeComMaxResponseBodyBytes {
		return result, errors.New(errors.CodeServiceUnavailable, "企业微信 Token 响应过大")
	}
	if err := sonic.Unmarshal(body, &result); err != nil {
		return result, errors.Wrap(errors.CodeServiceUnavailable, err, "解析企业微信 Token 响应失败")
	}
	return result, nil
}

func (p *TokenProvider) completeSuccess(ctx context.Context, integrationID string, status int, result tokenResponse, startedAt time.Time) error {
	_, err := p.integration.Complete(ctx, integrationID, integrationlog.Completion{
		Result: constants.IntegrationResultSuccess, HTTPStatus: status,
		ProviderCode: strconv.FormatInt(result.ErrCode, 10), ProviderMessage: result.ErrMsg,
		ResponseSummary: map[string]any{"errcode": result.ErrCode, "expires_in": result.ExpiresIn, "token_present": result.AccessToken != ""},
		DurationMS:      p.now().Sub(startedAt).Milliseconds(),
	})
	return err
}

func (p *TokenProvider) completeFailed(ctx context.Context, integrationID string, status int, providerCode, providerMessage string, startedAt time.Time) error {
	if providerMessage == "" {
		providerMessage = "企业微信 Token 接口返回失败"
	}
	_, err := p.integration.Complete(ctx, integrationID, integrationlog.Completion{
		Result: constants.IntegrationResultFailed, HTTPStatus: status, ProviderCode: providerCode,
		ProviderMessage: providerMessage, ResponseSummary: map[string]any{"success": false},
		DurationMS: p.now().Sub(startedAt).Milliseconds(),
	})
	if err != nil {
		return err
	}
	return errors.New(errors.CodeServiceUnavailable, "企业微信连接失败，请检查应用配置和可信 IP")
}

func (p *TokenProvider) releaseLock(ctx context.Context, lockKey, lockValue string, applicationID uint) {
	if p.redis == nil {
		return
	}
	if err := p.redis.Eval(ctx, releaseTokenLockScript, []string{lockKey}, lockValue).Err(); err != nil {
		p.warnRedis("释放企业微信 access_token 回源锁失败", err, applicationID)
	}
}

func (p *TokenProvider) warnRedis(message string, err error, applicationID uint) {
	if p.logger != nil {
		p.logger.Warn(message, zap.Uint("application_id", applicationID), zap.Error(err))
	}
}