Support custom header

core-impl/filter/src/main/java/com/alipay/sofa/rpc/filter/ConsumerCustomHeaderFilter.java

@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alipay.sofa.rpc.filter;
+
+import com.alipay.sofa.common.utils.Ordered;
+import com.alipay.sofa.rpc.common.utils.CommonUtils;
+import com.alipay.sofa.rpc.context.RpcInternalContext;
+import com.alipay.sofa.rpc.core.exception.SofaRpcException;
+import com.alipay.sofa.rpc.core.request.SofaRequest;
+import com.alipay.sofa.rpc.core.response.SofaResponse;
+import com.alipay.sofa.rpc.ext.Extension;
+
+import java.util.Map;
+
+/**
+ * Set customHeader to  SofaRequest.requestProps
+ *
+ * @author zhaowang
+ * @version : ConsumerCustomHeaderFilter.java, v 0.1 2021年06月23日 2:05 下午 zhaowang
+ */
+@AutoActive(consumerSide = true)
+@Extension(value = "consumerException", order = Ordered.LOWEST_PRECEDENCE)
+public class ConsumerCustomHeaderFilter extends Filter {
+
+    @Override
+    public SofaResponse invoke(FilterInvoker invoker, SofaRequest request) throws SofaRpcException {
+        setCustomHeader(request);
+        return invoker.invoke(request);
+    }
+
+    private void setCustomHeader(SofaRequest sofaRequest) {
+        RpcInternalContext context = RpcInternalContext.getContext();
+        Map customHeader = context.getCustomHeader();
+        if (CommonUtils.isNotEmpty(customHeader)) {
+            sofaRequest.addRequestProps(customHeader);
+        }
+        context.clearCustomHeader();
+    }
+}

core-impl/filter/src/main/resources/META-INF/services/sofa-rpc/com.alipay.sofa.rpc.filter.Filter

@@ -1,3 +1,4 @@
 # name                                                         # order
 com.alipay.sofa.rpc.filter.ProviderExceptionFilter             # -20000
-com.alipay.sofa.rpc.filter.ConsumerExceptionFilter             # -20000
+com.alipay.sofa.rpc.filter.ConsumerExceptionFilter             # -20000
+com.alipay.sofa.rpc.filter.ConsumerCustomHeaderFilter          # Ordered.LOWEST_PRECEDENCE

core-impl/filter/src/test/java/com/alipay/sofa/rpc/filter/CustomHeaderFilterTest.java

@@ -0,0 +1,63 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alipay.sofa.rpc.filter;
+
+import com.alipay.sofa.rpc.config.AbstractInterfaceConfig;
+import com.alipay.sofa.rpc.context.RpcInternalContext;
+import com.alipay.sofa.rpc.core.exception.SofaRpcException;
+import com.alipay.sofa.rpc.core.request.SofaRequest;
+import com.alipay.sofa.rpc.core.response.SofaResponse;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.Map;
+
+/**
+ * @author zhaowang
+ * @version : CustomHeaderFilterTest.java, v 0.1 2021年06月23日 2:26 下午 zhaowang
+ */
+public class CustomHeaderFilterTest {
+
+    @Test
+    public void testCustomFilter() {
+        RpcInternalContext.getContext().addCustomHeader("a", "b");
+        ConsumerCustomHeaderFilter filter = new ConsumerCustomHeaderFilter();
+        SofaRequest request = new SofaRequest();
+        filter.invoke(new EmptyInvoker(null), request);
+        Assert.assertTrue(RpcInternalContext.getContext().getCustomHeader().isEmpty());
+        Assert.assertEquals("b", request.getRequestProp("a"));
+    }
+
+    static class EmptyInvoker extends FilterInvoker {
+
+        public Map<String, String> getMetaHolder() {
+            return metaHolder;
+        }
+
+        private Map<String, String> metaHolder;
+
+        protected EmptyInvoker(AbstractInterfaceConfig config) {
+            super(config);
+        }
+
+        @Override
+        public SofaResponse invoke(SofaRequest request) throws SofaRpcException {
+            return null;
+        }
+    }
+
+}

core/api/src/main/java/com/alipay/sofa/rpc/context/RpcInternalContext.java

@@ -28,6 +28,7 @@
 import java.net.InetSocketAddress;
 import java.util.ArrayDeque;
 import java.util.Deque;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -184,6 +185,13 @@ protected RpcInternalContext() {
      */
     private ProviderInfo        providerInfo;
 
+
+    /**
+     * 自定义 header ，用完一次即删
+     */
+    protected HashMap<String, String> customHeader = new HashMap<>();
+
+
     /**
      * Is provider side.
      *
@@ -499,4 +507,25 @@ static boolean isHiddenParamKey(String key) {
         char c = key.charAt(0);
         return c == RpcConstants.HIDE_KEY_PREFIX;
     }
+
+
+    public Map<String, String> getCustomHeader() {
+        return new HashMap<>(customHeader);
+    }
+
+    /**
+     * 设置请求头，与 RequestBaggage 相比
+     * 1. 不受 enable baggage 开关影响，始终生效
+     * 2. 仅对一次调用生效，调用完成之会被清空
+     *
+     * @param key header key
+     * @param value header value
+     */
+    public void addCustomHeader(String key, String value) {
+        customHeader.put(key, value);
+    }
+
+    public void clearCustomHeader() {
+        customHeader.clear();
+    }
 }