wechat-api/internal/middleware/openauthority_middleware.go

package middleware

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"

	"wechat-api/ent"
	"wechat-api/ent/apikey"
	"wechat-api/ent/predicate"
	"wechat-api/internal/utils/compapi"
	"wechat-api/internal/utils/contextkey"

	"wechat-api/internal/config"

	"github.com/redis/go-redis/v9"
	"github.com/suyuan32/simple-admin-common/utils/jwt"
	"github.com/zeromicro/go-zero/core/errorx"
	"github.com/zeromicro/go-zero/rest/httpx"
)

/*
//"wechat-api/internal/types/payload"
var p types.payload.SendWxPayload
*/

type OpenAuthorityMiddleware struct {
	DB     *ent.Client
	Rds    redis.UniversalClient
	Config config.Config
}

func NewOpenAuthorityMiddleware(db *ent.Client, rds redis.UniversalClient, c config.Config) *OpenAuthorityMiddleware {
	return &OpenAuthorityMiddleware{
		DB:     db,
		Rds:    rds,
		Config: c,
	}
}

func (m *OpenAuthorityMiddleware) checkTokenUserInfo(ctx context.Context, authToken string) (*ent.ApiKey, int, error) {
	var (
		rc        int
		err       error
		apiKeyObj *ent.ApiKey
		fromId    = -1
	)
	_ = fromId

	//首先从redis取数据
	apiKeyObj, rc, err = m.getTokenUserInfoByRds(ctx, authToken)
	if rc <= 0 || err != nil { //无法获得后再从数据库获得

		rc = 0
		err = nil
		apiKeyObj, rc, err = m.getTokenUserInfoByDb(ctx, authToken)
		if err == nil {
			//get apiKeyObj from db succ
			fromId = 1
			err = m.saveTokenUserInfoToRds(ctx, authToken, apiKeyObj)
		}
	} else {
		fromId = 2
	}

	if err == nil {

		/*
			fromStr := ""
			switch fromId {
			case 1:
				fromStr = "DB"
			case 2:
				fromStr = "RDS"
			}

			fmt.Println("=========================================>>>")
			fmt.Printf("In checkTokenUserInfo Get Token Info From %s\n", fromStr)
			fmt.Printf("Auth Token:'%s'\n", apiKeyObj.Key)
			fmt.Printf("AgentID:%d\n", apiKeyObj.AgentID)
			fmt.Printf("APIBase:'%s'\n", apiKeyObj.Edges.Agent.APIBase)
			fmt.Printf("APIKey:'%s'\n", apiKeyObj.Edges.Agent.APIKey)
			fmt.Printf("Type:%d\n", apiKeyObj.Edges.Agent.Type)
			fmt.Printf("Model:'%s'\n", apiKeyObj.Edges.Agent.Model)
			//fmt.Println(typekit.PrettyPrint(apiKeyObj))
			//fmt.Printf("OpenaiBase:'%s'\n", apiKeyObj.OpenaiBase)
			//fmt.Printf("OpenaiKey:'%s'\n", apiKeyObj.OpenaiKey)
			fmt.Println("<<<=========================================")
		*/

	}
	return apiKeyObj, rc, err
}

func (m *OpenAuthorityMiddleware) saveTokenUserInfoToRds(ctx context.Context, authToken string, saveInfo *ent.ApiKey) error {
	bs, err := json.Marshal(saveInfo)
	if err == nil {

		_, err = m.Rds.HSet(ctx, compapi.APIAuthInfoKey, authToken, string(bs)).Result()
	}
	return err
}
func (m *OpenAuthorityMiddleware) getTokenUserInfoByRds(ctx context.Context, authToken string) (*ent.ApiKey, int, error) {

	rcode := -1
	result := ent.ApiKey{}
	m.Rds.Del(ctx, compapi.APIAuthInfoKey) //for debug
	jsonStr, err := m.Rds.HGet(ctx, compapi.APIAuthInfoKey, authToken).Result()
	if errors.Is(err, redis.Nil) {
		rcode = 0 //key not exist
	} else if err == nil { //find key

		err := json.Unmarshal([]byte(jsonStr), &result)
		if err == nil {
			rcode = 1
		}
	}
	return &result, rcode, err
}

func (m *OpenAuthorityMiddleware) getTokenUserInfoByDb(ctx context.Context, loginToken string) (*ent.ApiKey, int, error) {
	rcode := -1

	var predicates []predicate.ApiKey
	predicates = append(predicates, apikey.KeyEQ(loginToken))
	val, err := m.DB.ApiKey.Query().Where(predicates...).WithAgent().Only(ctx)
	//val, err := m.DB.ApiKey.Query().Where(predicates...).First(ctx)
	if err != nil {
		return nil, rcode, err
	}
	if val.Edges.Agent == nil {
		return nil, rcode, fmt.Errorf("edge get agent info is nil by agentid:%d", val.AgentID)
	}
	return val, 1, nil
}

func (m *OpenAuthorityMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {

		ctx := r.Context()
		ctx = contextkey.HttpResponseWriterKey.WithValue(ctx, w) //context存入http.ResponseWriter
		authToken := jwt.StripBearerPrefixFromToken(r.Header.Get("Authorization"))
		if len(authToken) == 0 {
			httpx.Error(w, errorx.NewApiError(http.StatusForbidden, "无法获取token"))
			return
		}
		apiKeyObj, _, err := m.checkTokenUserInfo(ctx, authToken)
		if err != nil {
			httpx.Error(w, errorx.NewApiError(http.StatusForbidden, "无法获取合适的授权信息"))
			return
		}
		if len(apiKeyObj.OpenaiBase) == 0 {
			httpx.Error(w, errorx.NewApiError(http.StatusForbidden, "缺失OpenaiBase"))
			return
		}

		ctx = contextkey.AuthTokenInfoKey.WithValue(ctx, apiKeyObj)

		newReq := r.WithContext(ctx)
		// Passthrough to next handler if need
		next(w, newReq)
	}
}