[frontend-gray] Increase gray types according to the ratio-weight gray (#1291)

plugins/wasm-go/extensions/frontend-gray/main.go

@@ -28,6 +28,7 @@ func main() {
 func parseConfig(json gjson.Result, grayConfig *config.GrayConfig, log wrapper.Log) error {
 	// 解析json 为GrayConfig
 	config.JsonToGrayConfig(json, grayConfig)
+	log.Infof("Rewrite: %v, GrayDeployments: %v", json.Get("rewrite"), json.Get("grayDeployments"))
 	return nil
 }
 
@@ -40,10 +41,12 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, grayConfig config.GrayConfig,
 	path, _ := proxywasm.GetHttpRequestHeader(":path")
 	fetchMode, _ := proxywasm.GetHttpRequestHeader("sec-fetch-mode")
 
-	isIndex := util.IsIndexRequest(fetchMode, path)
+	isPageRequest := util.IsPageRequest(fetchMode, path)
 	hasRewrite := len(grayConfig.Rewrite.File) > 0 || len(grayConfig.Rewrite.Index) > 0
-	grayKeyValue := util.GetGrayKey(util.ExtractCookieValueByKey(cookies, grayConfig.GrayKey), grayConfig.GraySubKey)
-
+	grayKeyValueByCookie := util.ExtractCookieValueByKey(cookies, grayConfig.GrayKey)
+	grayKeyValueByHeader, _ := proxywasm.GetHttpRequestHeader(grayConfig.GrayKey)
+	// 优先从cookie中获取，否则从header中获取
+	grayKeyValue := util.GetGrayKey(grayKeyValueByCookie, grayKeyValueByHeader, grayConfig.GraySubKey)
 	// 如果有重写的配置，则进行重写
 	if hasRewrite {
 		// 禁止重新路由，要在更改Header之前操作，否则会失效
@@ -53,22 +56,34 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, grayConfig config.GrayConfig,
 	// 删除Accept-Encoding，避免压缩， 如果是压缩的内容，后续插件就没法处理了
 	_ = proxywasm.RemoveHttpRequestHeader("Accept-Encoding")
 	_ = proxywasm.RemoveHttpRequestHeader("Content-Length")
+	deployment := &config.Deployment{}
 
-	grayDeployment := util.FilterGrayRule(&grayConfig, grayKeyValue, log.Infof)
-	frontendVersion := util.GetVersion(grayConfig.BaseDeployment.Version, cookies, isIndex)
-	backendVersion := ""
-
-	// 命中灰度规则
-	if grayDeployment != nil {
-		frontendVersion = util.GetVersion(grayDeployment.Version, cookies, isIndex)
-		backendVersion = grayDeployment.BackendVersion
+	preVersion, preUniqueClientId := util.GetXPreHigressVersion(cookies)
+	// 客户端唯一ID，用于在按照比率灰度时候 客户访问黏贴
+	uniqueClientId := grayKeyValue
+	if uniqueClientId == "" {
+		xForwardedFor, _ := proxywasm.GetHttpRequestHeader("X-Forwarded-For")
+		uniqueClientId = util.GetRealIpFromXff(xForwardedFor)
 	}
 
-	proxywasm.AddHttpRequestHeader(config.XHigressTag, frontendVersion)
+	// 如果没有配置比例，则进行灰度规则匹配
+	if isPageRequest {
+		log.Infof("grayConfig.TotalGrayWeight==== %v", grayConfig.TotalGrayWeight)
+		if grayConfig.TotalGrayWeight > 0 {
+			deployment = util.FilterGrayWeight(&grayConfig, preVersion, preUniqueClientId, uniqueClientId)
+		} else {
+			deployment = util.FilterGrayRule(&grayConfig, grayKeyValue)
+		}
+		log.Infof("index deployment: %v, path: %v, backend: %v, xPreHigressVersion: %s,%s", deployment, path, deployment.BackendVersion, preVersion, preUniqueClientId)
+	} else {
+		deployment = util.GetVersion(grayConfig, deployment, preVersion, isPageRequest)
+	}
+	proxywasm.AddHttpRequestHeader(config.XHigressTag, deployment.Version)
 
-	ctx.SetContext(config.XPreHigressTag, frontendVersion)
-	ctx.SetContext(config.XMseTag, backendVersion)
-	ctx.SetContext(config.IsIndex, isIndex)
+	ctx.SetContext(config.XPreHigressTag, deployment.Version)
+	ctx.SetContext(grayConfig.BackendGrayTag, deployment.BackendVersion)
+	ctx.SetContext(config.IsPageRequest, isPageRequest)
+	ctx.SetContext(config.XUniqueClientId, uniqueClientId)
 
 	rewrite := grayConfig.Rewrite
 	if rewrite.Host != "" {
@@ -77,12 +92,12 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, grayConfig config.GrayConfig,
 
 	if hasRewrite {
 		rewritePath := path
-		if isIndex {
-			rewritePath = util.IndexRewrite(path, frontendVersion, grayConfig.Rewrite.Index)
+		if isPageRequest {
+			rewritePath = util.IndexRewrite(path, deployment.Version, grayConfig.Rewrite.Index)
 		} else {
-			rewritePath = util.PrefixFileRewrite(path, frontendVersion, grayConfig.Rewrite.File)
+			rewritePath = util.PrefixFileRewrite(path, deployment.Version, grayConfig.Rewrite.File)
 		}
-		log.Infof("rewrite path: %s %s %v", path, frontendVersion, rewritePath)
+		log.Infof("rewrite path: %s %s %v", path, deployment.Version, rewritePath)
 		proxywasm.ReplaceHttpRequestHeader(":path", rewritePath)
 	}
 
@@ -95,15 +110,34 @@ func onHttpResponseHeader(ctx wrapper.HttpContext, grayConfig config.GrayConfig,
 	}
 	status, err := proxywasm.GetHttpResponseHeader(":status")
 	contentType, _ := proxywasm.GetHttpResponseHeader("Content-Type")
+
+	if grayConfig.Rewrite != nil && grayConfig.Rewrite.Host != "" {
+		// 删除Content-Disposition，避免自动下载文件
+		proxywasm.RemoveHttpResponseHeader("Content-Disposition")
+	}
+
+	isPageRequest, ok := ctx.GetContext(config.IsPageRequest).(bool)
+	if !ok {
+		isPageRequest = false // 默认值
+	}
+
 	if err != nil || status != "200" {
-		isIndex := ctx.GetContext(config.IsIndex)
 		if status == "404" {
-			if grayConfig.Rewrite.NotFound != "" && isIndex != nil && isIndex.(bool) {
-				ctx.SetContext(config.NotFound, true)
+			if grayConfig.Rewrite.NotFound != "" && isPageRequest {
+				ctx.SetContext(config.IsNotFound, true)
 				responseHeaders, _ := proxywasm.GetHttpResponseHeaders()
 				headersMap := util.ConvertHeaders(responseHeaders)
-				headersMap[":status"][0] = "200"
-				headersMap["content-type"][0] = "text/html"
+				if _, ok := headersMap[":status"]; !ok {
+					headersMap[":status"] = []string{"200"} // 如果没有初始化，设定默认值
+				} else {
+					headersMap[":status"][0] = "200" // 修改现有值
+				}
+				if _, ok := headersMap["content-type"]; !ok {
+					headersMap["content-type"] = []string{"text/html"} // 如果没有初始化，设定默认值
+				} else {
+					headersMap["content-type"][0] = "text/html" // 修改现有值
+				}
+				// 删除 content-length 键
 				delete(headersMap, "content-length")
 				proxywasm.ReplaceHttpResponseHeaders(util.ReconvertHeaders(headersMap))
 				ctx.BufferResponseBody()
@@ -119,24 +153,22 @@ func onHttpResponseHeader(ctx wrapper.HttpContext, grayConfig config.GrayConfig,
 	// 删除content-length，可能要修改Response返回值
 	proxywasm.RemoveHttpResponseHeader("Content-Length")
 
-	// 删除Content-Disposition，避免自动下载文件
-	proxywasm.RemoveHttpResponseHeader("Content-Disposition")
-
-	if strings.HasPrefix(contentType, "text/html") {
-		ctx.SetContext(config.IsHTML, true)
+	if strings.HasPrefix(contentType, "text/html") || isPageRequest {
 		// 不会进去Streaming 的Body处理
 		ctx.BufferResponseBody()
 
-		// 添加Cache-Control 头部，禁止缓存
-		proxywasm.ReplaceHttpRequestHeader("Cache-Control", "no-cache, no-store")
+		proxywasm.ReplaceHttpResponseHeader("Cache-Control", "no-cache, no-store")
 
 		frontendVersion := ctx.GetContext(config.XPreHigressTag).(string)
-		backendVersion := ctx.GetContext(config.XMseTag).(string)
+		xUniqueClient := ctx.GetContext(config.XUniqueClientId).(string)
 
-		// 设置当前的前端版本
-		proxywasm.AddHttpResponseHeader("Set-Cookie", fmt.Sprintf("%s=%s; Path=/;", config.XPreHigressTag, frontendVersion))
-		// 设置后端的前端版本
-		proxywasm.AddHttpResponseHeader("Set-Cookie", fmt.Sprintf("%s=%s; Path=/;", config.XMseTag, backendVersion))
+		// 设置前端的版本
+		proxywasm.AddHttpResponseHeader("Set-Cookie", fmt.Sprintf("%s=%s,%s; Max-Age=%s; Path=/;", config.XPreHigressTag, frontendVersion, xUniqueClient, grayConfig.UserStickyMaxAge))
+		// 设置后端的版本
+		if util.IsBackendGrayEnabled(grayConfig) {
+			backendVersion := ctx.GetContext(grayConfig.BackendGrayTag).(string)
+			proxywasm.AddHttpResponseHeader("Set-Cookie", fmt.Sprintf("%s=%s; Max-Age=%s; Path=/;", grayConfig.BackendGrayTag, backendVersion, grayConfig.UserStickyMaxAge))
+		}
 	}
 	return types.ActionContinue
 }
@@ -145,26 +177,52 @@ func onHttpResponseBody(ctx wrapper.HttpContext, grayConfig config.GrayConfig, b
 	if !util.IsGrayEnabled(grayConfig) {
 		return types.ActionContinue
 	}
-	backendVersion := ctx.GetContext(config.XMseTag)
-	isHtml := ctx.GetContext(config.IsHTML)
-	isIndex := ctx.GetContext(config.IsIndex)
-	notFoundUri := ctx.GetContext(config.NotFound)
-	if isIndex != nil && isIndex.(bool) && notFoundUri != nil && notFoundUri.(bool) && grayConfig.Rewrite.Host != "" && grayConfig.Rewrite.NotFound != "" {
+	isPageRequest, ok := ctx.GetContext(config.IsPageRequest).(bool)
+	if !ok {
+		isPageRequest = false // 默认值
+	}
+	frontendVersion := ctx.GetContext(config.XPreHigressTag).(string)
+
+	isNotFound, ok := ctx.GetContext(config.IsNotFound).(bool)
+	if !ok {
+		isNotFound = false // 默认值
+	}
+
+	if isPageRequest && isNotFound && grayConfig.Rewrite.Host != "" && grayConfig.Rewrite.NotFound != "" {
 		client := wrapper.NewClusterClient(wrapper.RouteCluster{Host: grayConfig.Rewrite.Host})
-		client.Get(grayConfig.Rewrite.NotFound, nil, func(statusCode int, responseHeaders http.Header, responseBody []byte) {
+
+		client.Get(strings.Replace(grayConfig.Rewrite.NotFound, "{version}", frontendVersion, -1), nil, func(statusCode int, responseHeaders http.Header, responseBody []byte) {
 			proxywasm.ReplaceHttpResponseBody(responseBody)
 			proxywasm.ResumeHttpResponse()
 		}, 1500)
 		return types.ActionPause
 	}
 
-	// 以text/html 开头，将 cookie转到cookie
-	if isHtml != nil && isHtml.(bool) && backendVersion != nil && backendVersion.(string) != "" {
-		newText := strings.ReplaceAll(string(body), "</head>", `<script>
-				!function(e,t){function n(e){var n="; "+t.cookie,r=n.split("; "+e+"=");return 2===r.length?r.pop().split(";").shift():null}var r=n("x-mse-tag");if(!r)return null;var s=XMLHttpRequest.prototype.open;XMLHttpRequest.prototype.open=function(e,t,n,a,i){return this._XHR=!0,this.addEventListener("readystatechange",function(){1===this.readyState&&r&&this.setRequestHeader("x-mse-tag",r)}),s.apply(this,arguments)};var a=e.fetch;e.fetch=function(e,t){return"undefined"==typeof t&&(t={}),"undefined"==typeof t.headers&&(t.headers={}),r&&(t.headers["x-mse-tag"]=r),a.apply(this,[e,t])}}(window,document);
-			</script>
-		</head>`)
-		if err := proxywasm.ReplaceHttpResponseBody([]byte(newText)); err != nil {
+	if isPageRequest {
+		// 将原始字节转换为字符串
+		newBody := string(body)
+
+		// 收集需要插入的内容
+		headInjection := strings.Join(grayConfig.Injection.Head, "\n")
+		bodyFirstInjection := strings.Join(grayConfig.Injection.Body.First, "\n")
+		bodyLastInjection := strings.Join(grayConfig.Injection.Body.Last, "\n")
+
+		// 使用 strings.Builder 来提高性能
+		var sb strings.Builder
+		// 预分配内存，避免多次内存分配
+		sb.Grow(len(newBody) + len(headInjection) + len(bodyFirstInjection) + len(bodyLastInjection))
+		sb.WriteString(newBody)
+
+		// 进行替换
+		content := sb.String()
+		content = strings.ReplaceAll(content, "</head>", fmt.Sprintf("%s\n</head>", headInjection))
+		content = strings.ReplaceAll(content, "<body>", fmt.Sprintf("<body>\n%s", bodyFirstInjection))
+		content = strings.ReplaceAll(content, "</body>", fmt.Sprintf("%s\n</body>", bodyLastInjection))
+
+		// 最终结果
+		newBody = content
+
+		if err := proxywasm.ReplaceHttpResponseBody([]byte(newBody)); err != nil {
 			return types.ActionContinue
 		}
 	}