From 168cc5b488c84320beaffa526508cf36dfcc3fcf Mon Sep 17 00:00:00 2001
From: hongqiaowei <hongqiaowei@163.com>
Date: Mon, 1 Aug 2022 15:49:50 +0800
Subject: [PATCH] GrayReleasePlugin can append header

---
 fizz-core/src/main/java/we/util/WebUtils.java | 16 +++++--
 .../plugin/grayrelease/GrayReleasePlugin.java | 42 ++++++++++++++++++-
 2 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/fizz-core/src/main/java/we/util/WebUtils.java b/fizz-core/src/main/java/we/util/WebUtils.java
index 97be01a..72a6631 100644
--- a/fizz-core/src/main/java/we/util/WebUtils.java
+++ b/fizz-core/src/main/java/we/util/WebUtils.java
@@ -437,9 +437,19 @@ public abstract class WebUtils {
     }
 
     public static Map<String, String> appendHeader(ServerWebExchange exchange, String name, String value) {
-        Map<String, String> hdrs = getAppendHeaders(exchange);
-        hdrs.put(name, value);
-        return hdrs;
+        Map<String, String> appendHeaders = getAppendHeaders(exchange);
+        appendHeaders.put(name, value);
+        return appendHeaders;
+    }
+
+    public static Map<String, String> appendHeaders(ServerWebExchange exchange, Map<String, String> headers) {
+        Map<String, String> appendHeaders = getAppendHeaders(exchange);
+        headers.forEach(
+                           (name, value) -> {
+                               appendHeaders.put(name, value);
+                           }
+               );
+        return appendHeaders;
     }
 
     public static Map<String, String> getAppendHeaders(ServerWebExchange exchange) {
diff --git a/fizz-plugin/src/main/java/we/plugin/grayrelease/GrayReleasePlugin.java b/fizz-plugin/src/main/java/we/plugin/grayrelease/GrayReleasePlugin.java
index 19df1c8..e1857bb 100644
--- a/fizz-plugin/src/main/java/we/plugin/grayrelease/GrayReleasePlugin.java
+++ b/fizz-plugin/src/main/java/we/plugin/grayrelease/GrayReleasePlugin.java
@@ -385,7 +385,17 @@ public class GrayReleasePlugin extends RequestBodyPlugin {
             }
             String qry = newRouteConfig.get("query");
             if (StringUtils.isNotBlank(qry)) {
-                route.query = qry.trim();
+                String reqQry = exchange.getRequest().getURI().getQuery();
+                if (StringUtils.isBlank(reqQry)) {
+                    route.query = qry.trim();
+                } else {
+                    route.query = reqQry + Consts.S.AND + qry.trim();
+                }
+            }
+            String headerStr = newRouteConfig.get("header");
+            if (StringUtils.isNotBlank(headerStr)) {
+                Map<String, String> headers = headerStr2map(headerStr);
+                WebUtils.appendHeaders(exchange, headers);
             }
             String retryCount = newRouteConfig.get("retryCount");
             if (StringUtils.isNotBlank(retryCount)) {
@@ -455,7 +465,18 @@ public class GrayReleasePlugin extends RequestBodyPlugin {
 
         String qry = newRouteConfig.get("query");
         if (StringUtils.isNotBlank(qry)) {
-            route.query = qry.trim();
+            String reqQry = exchange.getRequest().getURI().getQuery();
+            if (StringUtils.isBlank(reqQry)) {
+                route.query = qry.trim();
+            } else {
+                route.query = reqQry + Consts.S.AND + qry.trim();
+            }
+        }
+
+        String headerStr = newRouteConfig.get("header");
+        if (StringUtils.isNotBlank(headerStr)) {
+            Map<String, String> headers = headerStr2map(headerStr);
+            WebUtils.appendHeaders(exchange, headers);
         }
 
         String timeout = newRouteConfig.get("timeout");
@@ -485,6 +506,11 @@ public class GrayReleasePlugin extends RequestBodyPlugin {
             route.backendPath = UrlTransformUtils.transform(route.path, path.trim(), WebUtils.getClientReqPath(exchange));
             WebUtils.setBackendPath(exchange, route.backendPath);
         }
+        String headerStr = newRouteConfig.get("header");
+        if (StringUtils.isNotBlank(headerStr)) {
+            Map<String, String> headers = headerStr2map(headerStr);
+            WebUtils.appendHeaders(exchange, headers);
+        }
     }
 
     private Map<String, String> routeConfig2map(String config) {
@@ -496,4 +522,16 @@ public class GrayReleasePlugin extends RequestBodyPlugin {
         }
         return result;
     }
+
+    private Map<String, String> headerStr2map(String headerStr) {
+        String[] hvs = StringUtils.split(headerStr, Consts.S.COMMA);
+        Map<String, String> headerMap = new HashMap<>(hvs.length, 1);
+        for (String hv : hvs) {
+            int eqIdx = hv.indexOf(Consts.S.EQUAL);
+            String h = hv.substring(0, eqIdx).trim();
+            String v = hv.substring(eqIdx + 1).trim();
+            headerMap.put(h, v);
+        }
+        return headerMap;
+    }
 }