GtDataStore/app/cmd/dtgateway/internal/middleware/datasignMiddleware.go

package middleware

import (
	"GtDataStore/app/cmd/dtgateway/internal/config"
	"GtDataStore/app/model"
	"GtDataStore/common/crypto/md5"
	"GtDataStore/common/utils"
	"GtDataStore/common/xerr"
	"bytes"
	"context"
	"fmt"
	"github.com/zeromicro/go-zero/core/stores/sqlx"
	"github.com/zeromicro/go-zero/core/validation"
	"github.com/zeromicro/go-zero/rest/httpx"
	"io/ioutil"
	"k8s.io/apimachinery/pkg/util/json"
	"math"
	"net/http"
	"net/url"
	"sync"
	"sync/atomic"
	"time"
)

const (
	TIME_OUT                = 10 // ts时间窗口大小, 单位: 秒
	S2_MIN_LENGTH           = 2048
	S2_HEAD_LENGTH          = 200
	S2_TAIL_LENGTH          = 200
	APP_SECRET_CACHE_EXPIRE = 600
)

var (
	appSecrets     = make(map[string]secretExpire)
	appSecretsLock sync.Mutex
	validator      atomic.Value
)

type (
	secretExpire struct {
		Secret string
		Status int64
		Expire time.Time
	}

	DataSignMiddleware struct {
		AppInfo model.DcAppInfoModel
	}

	CommonParams struct {
		Ts        int64  `form:"ts"`
		ProjectId int64  `form:"project_id"`
		Sign      string `form:"sign"`
		SignFlag  uint8  `form:"sign_flag,optional"`
		AppName   string `header:"APP-NAME"`
	}
)

func NewDataSignMiddleware(conf config.Config) *DataSignMiddleware {
	mysql := sqlx.NewMysql(conf.DtDataStoreDB.DataSource)
	return &DataSignMiddleware{
		AppInfo: model.NewDcAppInfoModel(mysql),
	}
}

func (m *DataSignMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		cps := CommonParams{}
		err := m.parseCommonParams(r, &cps)

		if err != nil {
			httpx.WriteJson(w, http.StatusOK, map[string]interface{}{
				"code": xerr.REUQEST_PARAM_ERROR,
				"msg":  err.Error(),
			})
			return
		}

		// 验证时间窗口
		if m.checkTs(cps.Ts) == false {
			httpx.WriteJson(w, http.StatusOK, map[string]interface{}{
				"code": xerr.REUQEST_PARAM_TS_ERROR,
				"msg":  "请求参数错误",
			})
			return
		}

		// 得到secret, 如果发生错误, 则响应失败
		appSecret, err := m.getAppSecret(cps.AppName)
		if err != nil {
			httpx.WriteJson(w, http.StatusOK, map[string]interface{}{
				"code": xerr.REUQEST_PARAM_APP_NAME_ERROR,
				"msg":  err.Error(),
			})
			return
		}

		// 解析query, 得到sign和s1
		s1, err := m.parseQuery(r)
		if err != nil {
			fmt.Printf("m.parseQuery(r) error :s\n", err.Error())
		}
		s2, err := m.parseBody(r)
		if err != nil {
			fmt.Printf("m.parseBody(r) error :s\n", err.Error())
		}

		s3 := r.ContentLength

		csign := m.calcSign(s1, s2, appSecret, s3, cps.SignFlag)
		if csign != cps.Sign {
			httpx.WriteJson(w, http.StatusOK, map[string]interface{}{
				"code": xerr.REUQEST_PARAM_DATA_SIGN_ERROR,
				"msg":  "未通过数据认证",
			})
			return
		}

		next(w, r)
	}
}

func (m *DataSignMiddleware) calcSign(s1, s2, appSecret string, s3 int64, signFlag uint8) string {
	// 处理s2
	if signFlag == 1 {
		s2 = m.cutS2(s2)
	} else {
		s2, _ = m.sortS2(s2)
	}

	//fmt.Printf("sign data: %s\n", fmt.Sprintf("%s_%s_%d_%s", s1, s2, s3, appSecret))
	//fmt.Printf("md5: %s\n", md5.Md5([]byte(fmt.Sprintf("%s%s%d%s", s1, s2, s3, appSecret))))

	return md5.Md5([]byte(fmt.Sprintf("%s%s%d%s", s1, s2, s3, appSecret)))
}

func (m *DataSignMiddleware) cutS2(s2 string) string {
	s2l := len(s2)
	if s2l <= S2_MIN_LENGTH {
		return s2
	}
	return fmt.Sprintf("%s%s", s2[:S2_HEAD_LENGTH], s2[s2l-S2_TAIL_LENGTH:])
}

func (m *DataSignMiddleware) sortS2(s2 string) (string, error) {
	var mi map[string]interface{}
	if err := json.Unmarshal([]byte(s2), &mi); err != nil {
		return "", err
	}

	smi := utils.SortMapByKey(mi)

	if bs, err := json.Marshal(smi); err != nil {
		return "", err
	} else {
		return string(bs), nil
	}
}

func (m *DataSignMiddleware) parseCommonParams(r *http.Request, v any) error {
	if err := httpx.ParsePath(r, v); err != nil {
		return err
	}

	if err := httpx.ParseForm(r, v); err != nil {
		return err
	}

	if err := httpx.ParseHeaders(r, v); err != nil {
		return err
	}

	if valid, ok := v.(validation.Validator); ok {
		return valid.Validate()
	} else if val := validator.Load(); val != nil {
		return val.(httpx.Validator).Validate(r, v)
	}
	return nil
}

func (m *DataSignMiddleware) parseBody(r *http.Request) (body string, err error) {
	if r.Method == http.MethodGet {
		return "", nil
	}

	cnt, _ := ioutil.ReadAll(r.Body)
	r.Body = ioutil.NopCloser(bytes.NewReader(cnt))

	return string(cnt), nil
}

func (m *DataSignMiddleware) parseQuery(r *http.Request) (query string, err error) {
	if vs, err := url.ParseQuery(r.URL.RawQuery); err != nil {
		return "", err
	} else {
		vs.Del("sign")
		return vs.Encode(), nil
	}
}

func (m *DataSignMiddleware) getAppSecret(appName string) (string, error) {
	if secret, ok := appSecrets[appName]; ok {
		if secret.Expire.After(time.Now()) == false {
			return secret.Secret, nil
		}
	}

	// 从数据库中读取appSecret
	if appInfo, err := m.AppInfo.FindOneByAppName(context.Background(), appName); err == nil {
		appSecretsLock.Lock()
		defer appSecretsLock.Unlock()
		appSecrets[appName] = secretExpire{
			Secret: appInfo.Secret,
			Status: appInfo.Status,
			Expire: time.Now().Add(APP_SECRET_CACHE_EXPIRE * time.Second),
		}
		if appInfo.Status == 0 {
			return appInfo.Secret, nil
		}
	}

	return "", xerr.NewErrCodeMsg(xerr.REUQEST_PARAM_APP_NAME_ERROR, "该app name错误, 无法进行数据验签")
}

func (m *DataSignMiddleware) checkTs(ts int64) bool {
	return math.Abs(float64(time.Now().Unix()-ts)) <= TIME_OUT
}