Merge branch 'feature_DSCP' into 'rel'

TSG-19340： add Feature dscp

See merge request tango/shaping-engine!78

14 files changed, 266 insertions, 8 deletions

diff --git a/common/include/raw_packet.h b/common/include/raw_packet.h
index 55e4b46..c06c371 100644
--- a/common/include/raw_packet.h
+++ b/common/include/raw_packet.h
@@ -93,6 +93,7 @@ int raw_packet_parser_get_most_outer_address(struct raw_pkt_parser *handler, str
 uint64_t raw_packet_parser_get_hash_value(struct raw_pkt_parser *handler, enum ldbc_method method, int dir_is_internal);
 
 int raw_packet_tcp_payload_len_get(struct raw_pkt_parser *handler);
+void raw_packet_set_dscp(struct ethhdr *eth_hdr, uint8_t dscp);
 
 #ifdef __cpluscplus
 }
diff --git a/common/src/raw_packet.cpp b/common/src/raw_packet.cpp
index 03be2b1..83dd0f0 100644
--- a/common/src/raw_packet.cpp
+++ b/common/src/raw_packet.cpp
@@ -1,5 +1,6 @@
 #include <assert.h>
 #include <errno.h>
+#include <linux/if_ether.h>
 #include <string.h>
 #include <stdlib.h>
 
@@ -1056,4 +1057,51 @@ static const void *parse_mpls(struct raw_pkt_parser *handler, const void *data,
         // TESTED
         return parse_mpls(handler, data_next_layer, data_next_length, LAYER_TYPE_MPLS);
     }
+}
+
+static unsigned short ip_checksum(struct iphdr *iph) {
+    // 初始化检验和为0
+    unsigned short checksum = 0;
+    // 将IP头的指针转换为无符号短整型的指针
+    unsigned short *p = (unsigned short *)iph;
+    // 计算IP头的长度，单位为字节
+    int len = iph->ihl * 4;
+    // 逐个累加IP头的每个16位字段
+    while (len > 1) {
+        checksum += *p++;
+        len -= 2;
+    }
+    // 如果IP头的长度为奇数，将最后一个字节视为高8位，低8位补0，再累加
+    if (len == 1) {
+        checksum += *(unsigned char *)p << 8;
+    }
+    // 将检验和的高16位和低16位相加，直到没有进位
+    while (checksum >> 16) {
+        checksum = (checksum & 0xFFFF) + (checksum >> 16);
+    }
+    // 将检验和取反
+    checksum = ~checksum;
+    // 返回检验和
+    return checksum;
+}
+
+void raw_packet_set_dscp(struct ethhdr *eth_hdr, uint8_t dscp)
+{
+    if (eth_hdr->h_proto == htons(ETH_P_IP))
+    {
+        struct iphdr *ip_hdr = (struct iphdr *)(eth_hdr + 1);
+        ip_hdr->tos = (ip_hdr->tos & 0x03) | (dscp << 2);
+        //recalculate checksum
+        ip_hdr->check = ip_checksum(ip_hdr);
+    } 
+    else if (eth_hdr->h_proto == htons(ETH_P_IPV6))
+    {
+        struct ip6_hdr *ip6_hdr = (struct ip6_hdr *)(eth_hdr + 1);
+        unsigned int flow = ntohl(ip6_hdr->ip6_flow);
+        flow &= 0xF03FFFFF;
+        flow |= (dscp << 22);
+        ip6_hdr->ip6_flow = htonl(flow);
+    }
+
+    return;
 }
\ No newline at end of file
diff --git a/common/test/CMakeLists.txt b/common/test/CMakeLists.txt
index 00dd805..fcc2225 100644
--- a/common/test/CMakeLists.txt
+++ b/common/test/CMakeLists.txt
@@ -20,7 +20,7 @@ target_link_libraries(gtest_session_table common pthread gtest)
 
 add_executable(gtest_raw_packet gtest_raw_packet.cpp)
 target_include_directories(gtest_raw_packet PUBLIC ${CMAKE_SOURCE_DIR}/common/include)
-target_link_libraries(gtest_raw_packet common pthread gtest)
+target_link_libraries(gtest_raw_packet common pthread pcap gtest)
 
 ###############################################################################
 # gtest_ctrl_packet
diff --git a/common/test/gtest_raw_packet.cpp b/common/test/gtest_raw_packet.cpp
index 47b4af6..02422da 100644
--- a/common/test/gtest_raw_packet.cpp
+++ b/common/test/gtest_raw_packet.cpp
@@ -1,5 +1,10 @@
+#include <cstdint>
 #include <gtest/gtest.h>
 #include <arpa/inet.h>
+#include <linux/if_ether.h>
+#include <pcap/pcap.h>
+#include <netinet/ip.h>
+#include <netinet/ip6.h>
 
 #include "raw_packet.h"
 #include "addr_tuple4.h"
@@ -1366,6 +1371,54 @@ TEST(RAW_PACKET, GET_HASH_VALUE)
 	EXPECT_TRUE(raw_packet_parser_get_hash_value(&handler, LDBC_METHOD_HASH_INNERMOST_INT_IP, 1) == raw_packet_parser_get_hash_value(&handler, LDBC_METHOD_HASH_INNERMOST_EXT_IP, 0));
 }
 
+TEST(RAW_PACKET, SET_IPV4_DSCP_VALUE)
+{
+	char error[100];
+	pcap_t *handle = pcap_open_offline("../../../common/test/pcap/ipv4_single_pkt.pcap", error);
+	ASSERT_NE(handle, nullptr);
+
+	struct pcap_pkthdr *header;
+	const u_char *packet;
+	int ret = pcap_next_ex(handle, &header, &packet);
+	ASSERT_EQ(ret, 1);
+
+	struct ethhdr *eth = (struct ethhdr *)packet;
+	struct iphdr *ip = (struct iphdr *)(eth + 1);
+
+	unsigned char dscp_value = ip->tos >> 2;
+	EXPECT_EQ(dscp_value, 0);
+
+	raw_packet_set_dscp(eth, 40);//40 is CS5
+	dscp_value = ip->tos >> 2;
+	EXPECT_EQ(dscp_value, 40);
+	pcap_close(handle);
+}
+
+TEST(RAW_PACKET, SET_IPV6_DSCP_VALUE)
+{
+	char error[100];
+	pcap_t *handle = pcap_open_offline("../../../common/test/pcap/ipv6_single_pkt.pcap", error);
+	ASSERT_NE(handle, nullptr);
+
+	struct pcap_pkthdr *header;
+	const u_char *packet;
+	int ret = pcap_next_ex(handle, &header, &packet);
+	ASSERT_EQ(ret, 1);
+
+	struct ethhdr *eth = (struct ethhdr *)packet;
+	struct ip6_hdr *ip = (struct ip6_hdr *)(eth + 1);
+
+	uint32_t flow_label = ntohl(ip->ip6_flow);
+	unsigned char dscp_value = (flow_label & 0x0fc00000) >> 22;
+	EXPECT_EQ(dscp_value, 0);
+
+	raw_packet_set_dscp(eth, 40);//40 is CS5
+	flow_label = ntohl(ip->ip6_flow);
+	dscp_value = (flow_label & 0x0fc00000) >> 22;
+	EXPECT_EQ(dscp_value, 40);
+	pcap_close(handle);
+}
+
 int main(int argc, char **argv)
 {
 	::testing::InitGoogleTest(&argc, argv);
diff --git a/common/test/pcap/ipv4_single_pkt.pcap b/common/test/pcap/ipv4_single_pkt.pcap
new file mode 100644
index 0000000..8823b94
--- /dev/null
+++ b/common/test/pcap/ipv4_single_pkt.pcap
diff --git a/common/test/pcap/ipv6_single_pkt.pcap b/common/test/pcap/ipv6_single_pkt.pcap
new file mode 100644
index 0000000..cf1ee6b
--- /dev/null
+++ b/common/test/pcap/ipv6_single_pkt.pcap
diff --git a/shaping/include/shaper.h b/shaping/include/shaper.h
index 215f6c4..9276c98 100644
--- a/shaping/include/shaper.h
+++ b/shaping/include/shaper.h
@@ -188,6 +188,8 @@ struct shaping_flow {
     int rule_num;
     struct shaping_rule_info matched_rule_infos[SHAPING_RULE_NUM_MAX];
     int priority;
+    unsigned char dscp_enable;
+    unsigned char dscp_value;
     struct addr_tuple4 tuple4;
     int anchor;//rule_idx
     int ref_cnt;
diff --git a/shaping/include/shaper_maat.h b/shaping/include/shaper_maat.h
index 1c8cb62..dc170e5 100644
--- a/shaping/include/shaper_maat.h
+++ b/shaping/include/shaper_maat.h
@@ -13,6 +13,8 @@ struct shaping_rule {
     int borrow_pf_id_array[SHAPING_REF_PROFILE_NUM_MAX];
     int borrow_pf_num;
     int fair_factor;
+    unsigned char dscp_enable;
+    unsigned char dscp_value;
 };
 
 struct shaping_profile {
diff --git a/shaping/src/main.cpp b/shaping/src/main.cpp
index 54e2f60..63b7a49 100644
--- a/shaping/src/main.cpp
+++ b/shaping/src/main.cpp
@@ -84,7 +84,7 @@ int main(int argc, char **argv)
 
     if (signal(SIGQUIT, sig_handler) == SIG_ERR || signal(SIGTERM, sig_handler) == SIG_ERR)
     {
-        LOG_ERROR("%s: unable to register SIGQUIT signal handler, error %d: %s", LOG_TAG_SHAPING, errno, strerror(errno));
+        LOG_ERROR("%s: unable to register SIGQUIT or SIGTERM signal handler, error %d: %s", LOG_TAG_SHAPING, errno, strerror(errno));
         LOG_CLOSE();
         return -1;
     }
diff --git a/shaping/src/shaper.cpp b/shaping/src/shaper.cpp
index cb4c382..59dbf53 100644
--- a/shaping/src/shaper.cpp
+++ b/shaping/src/shaper.cpp
@@ -1062,6 +1062,11 @@ void shaping_packet_process(struct shaping_thread_ctx *ctx, marsio_buff_t *rx_bu
     struct shaping_marsio_info *marsio_info = ctx->marsio_info;
 
     sf->processed_pkts++;
+
+    if (sf->dscp_enable) {
+        struct ethhdr *eth_hdr = (struct ethhdr*)marsio_buff_mtod(rx_buff);
+        raw_packet_set_dscp(eth_hdr, sf->dscp_value);
+    }
     
     if (meta->is_tcp_pure_ctrl) {
         shaper_token_consume_force(sf, meta);
diff --git a/shaping/src/shaper_maat.cpp b/shaping/src/shaper_maat.cpp
index 54d59c5..e7de972 100644
--- a/shaping/src/shaper_maat.cpp
+++ b/shaping/src/shaper_maat.cpp
@@ -7,6 +7,7 @@
 #include <MESA/maat.h>
 #include <MESA/MESA_handle_logger.h>
 
+#include "raw_packet.h"
 #include "log.h"
 #include "shaper.h"
 #include "shaper_maat.h"
@@ -36,6 +37,45 @@ struct shaper_maat_config {
     enum log_level log_level;
 };
 
+enum dscp_class
+{
+    DSCP_CLASS_DF = 0,
+    DSCP_CLASS_CS1,
+    DSCP_CLASS_AF13,
+    DSCP_CLASS_AF12,
+    DSCP_CLASS_AF11,
+    DSCP_CLASS_CS2,
+    DSCP_CLASS_AF23,
+    DSCP_CLASS_AF22,
+    DSCP_CLASS_AF21,
+    DSCP_CLASS_CS3,
+    DSCP_CLASS_AF33,
+    DSCP_CLASS_AF32,
+    DSCP_CLASS_AF31,
+    DSCP_CLASS_CS4,
+    DSCP_CLASS_AF43,
+    DSCP_CLASS_AF42,
+    DSCP_CLASS_AF41,
+    DSCP_CLASS_CS5,
+    DSCP_CLASS_EF,
+    DSCP_CLASS_CS6,
+    DSCP_CLASS_CS7,
+    DSCP_CLASS_MAX,
+};
+#define DSCP_VALUE_MAX 64
+                                                    //   0                1               2               3               4               5                6              7              8               9             10
+static int dscp_value_to_priority[DSCP_VALUE_MAX] = {DSCP_CLASS_DF, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_CS1, DSCP_CLASS_MAX, DSCP_CLASS_AF11,
+                                                     //  11              12              13              14              15              16              17              18              19              20             21
+                                                     DSCP_CLASS_MAX, DSCP_CLASS_AF12, DSCP_CLASS_MAX, DSCP_CLASS_AF13, DSCP_CLASS_MAX, DSCP_CLASS_CS2, DSCP_CLASS_MAX, DSCP_CLASS_AF21, DSCP_CLASS_MAX, DSCP_CLASS_AF22, DSCP_CLASS_MAX,
+                                                     //  22              23              24              25              26              27              28              29              30              31             32
+                                                     DSCP_CLASS_AF23, DSCP_CLASS_MAX, DSCP_CLASS_CS3, DSCP_CLASS_MAX, DSCP_CLASS_AF31, DSCP_CLASS_MAX, DSCP_CLASS_AF32, DSCP_CLASS_MAX, DSCP_CLASS_AF33, DSCP_CLASS_MAX, DSCP_CLASS_CS4,
+                                                     //  33              34              35              36              37              38              39              40              41              42             43
+                                                     DSCP_CLASS_MAX, DSCP_CLASS_AF41, DSCP_CLASS_MAX, DSCP_CLASS_AF42, DSCP_CLASS_MAX, DSCP_CLASS_AF43, DSCP_CLASS_MAX, DSCP_CLASS_CS5, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_MAX,
+                                                     //  44              45              46              47              48              49              50              51              52              53             54
+                                                     DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_EF, DSCP_CLASS_MAX, DSCP_CLASS_CS6, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_MAX,
+                                                     //  55              56              57              58              59              60              61              62              63
+                                                     DSCP_CLASS_MAX, DSCP_CLASS_CS7, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_MAX, DSCP_CLASS_MAX};
+
 struct maat *g_maat_instance = NULL;
 
 void shaper_rule_ex_new(const char *table_name, int table_id, const char *key, const char *table_line, void **ad, long argl, void *argp)
@@ -44,6 +84,7 @@ void shaper_rule_ex_new(const char *table_name, int table_id, const char *key, c
     cJSON *json=NULL;
     cJSON *tmp_obj = NULL;
     cJSON *tmp_array_obj = NULL;
+    cJSON *dscp_obj = NULL;
     int array_size;
     char user_region[1024] = {0};
     int i, ret;
@@ -87,6 +128,25 @@ void shaper_rule_ex_new(const char *table_name, int table_id, const char *key, c
     }
     s_rule->fair_factor = tmp_obj->valueint;
 
+    //dscp_marking
+    tmp_obj = cJSON_GetObjectItem(json, "dscp_marking");
+    if (!tmp_obj) {
+        LOG_ERROR("%s: json parse dscp_marking failed for table line %s", LOG_TAG_MAAT, table_line);
+        goto END;
+    }
+    dscp_obj = cJSON_GetObjectItem(tmp_obj, "enabled");
+    if (dscp_obj && dscp_obj->valueint == 1) {
+        dscp_obj = cJSON_GetObjectItem(tmp_obj, "dscp_value");
+        if (dscp_obj && dscp_obj->valueint < DSCP_VALUE_MAX && dscp_value_to_priority[dscp_obj->valueint] != DSCP_CLASS_MAX) {
+            s_rule->dscp_enable = 1;
+            s_rule->dscp_value = dscp_obj->valueint;
+        } else {
+            LOG_ERROR("%s: json parse dscp_value wrong for table line %s", LOG_TAG_MAAT, table_line);
+            goto END;
+        }
+    }
+
+    //profile_chain
     tmp_obj = cJSON_GetObjectItem(json, "profile_chain");
     if (!tmp_obj) {//required
         LOG_ERROR("%s: json parse profile_chain failed for table line %s", LOG_TAG_MAAT, table_line);
@@ -149,7 +209,6 @@ void shaper_profile_ex_new(const char *table_name, int table_id, const char *key
     char type_arg[64] = {0};
     char limits[128] = {0};
     char aqm_options[64] = {0};
-    char dscp_marking[64] = {0};
     char volume_based_shaping[64] = {0};
     int limit_bandwidth;
     int array_size, i;
@@ -161,9 +220,9 @@ void shaper_profile_ex_new(const char *table_name, int table_id, const char *key
 
     s_pf = (struct shaping_profile*)calloc(1, sizeof(struct shaping_profile));
 
-    ret = sscanf(table_line, "%d\t%63s\t%63s\t%127s\t%63s\t%63s\t%63s\t%d",
-            &s_pf->id, profile_type, type_arg, limits, aqm_options, dscp_marking, volume_based_shaping, &s_pf->valid);
-    if (ret != 8) {
+    ret = sscanf(table_line, "%d\t%63s\t%63s\t%127s\t%63s\t%63s\t%d",
+            &s_pf->id, profile_type, type_arg, limits, aqm_options, volume_based_shaping, &s_pf->valid);
+    if (ret != 7) {
         LOG_ERROR("%s: sscanf parse failed for profile line %s", LOG_TAG_MAAT, table_line);
         goto END;
     }
@@ -332,6 +391,13 @@ static int shaper_rule_update(struct shaping_thread_ctx *ctx, struct shaping_flo
         }
     }
 
+    if (s_rule->dscp_enable) {
+        sf->dscp_enable = 1;
+        if (dscp_value_to_priority[s_rule->dscp_value] > dscp_value_to_priority[sf->dscp_value]) {
+            sf->dscp_value = s_rule->dscp_value;
+        }
+    }
+
     if (s_rule->borrow_pf_num == 0) {
         return 0;
     }
diff --git a/shaping/test/gtest_shaper_maat.cpp b/shaping/test/gtest_shaper_maat.cpp
index 770c937..96071e2 100644
--- a/shaping/test/gtest_shaper_maat.cpp
+++ b/shaping/test/gtest_shaper_maat.cpp
@@ -13,7 +13,7 @@ TEST(shaping_rule, parse)
                         0\t\
                         1\t\
                         {}\t\
-                        {\"vsys_id\":2333,\"priority\":1,\"fair_factor\":10,\"profile_chain\":[1,2,3]}\t\
+                        {\"vsys_id\":2333,\"priority\":1,\"fair_factor\":10,\"dscp_marking\":{\"enabled\":1,\"dscp_value\":10},\"profile_chain\":[1,2,3]}\t\
                         0\t\
                         1\t\
                         }";
@@ -29,6 +29,8 @@ TEST(shaping_rule, parse)
     EXPECT_EQ(s_rule->borrow_pf_id_array[0], 2);
     EXPECT_EQ(s_rule->borrow_pf_id_array[1], 3);
     EXPECT_EQ(s_rule->priority, 1);
+    EXPECT_EQ(s_rule->dscp_enable, 1);
+    EXPECT_EQ(s_rule->dscp_value, 10);
 
     shaper_rule_ex_dup(0, (void**)&s_rule_dup, (void**)&s_rule, 0, NULL);
     EXPECT_EQ(s_rule, s_rule_dup);
@@ -43,7 +45,6 @@ TEST(shaping_profile, parse)
                         max_min_host_fairness\t\
                         [{\"direction\":\"incoming\",\"bandwidth\":1024},{\"direction\":\"outgoing\",\"bandwidth\":2048}]\t\
                         {\"enabled\":1,\"algorithm\":\"codel\"}\t\
-                        null\t\
                         {}\t\
                         1";
     struct shaping_profile *s_pf = NULL;
@@ -239,6 +240,78 @@ TEST(shaping_flow, update_rule_after_priority_confirmed)
     shaping_engine_destroy(ctx);
 }
 
+TEST(shaping_flow, update_rule_dscp)
+{
+    struct shaping_ctx *ctx = NULL;
+    struct shaping_flow sf;
+    struct shaping_rule_info *rule_info;
+    long long rule_ids[] = {1, 2, 3, 4};
+    int prioritys[] = {1, 2, 3, 4};
+    int profile_nums[] = {1, 1, 1, 1};
+    int profile_ids[][MAX_REF_PROFILE] = {{1}, {2}, {3}, {4}};
+
+    stub_init();
+
+    ctx = shaping_engine_init();
+
+    stub_set_matched_shaping_rules(4, rule_ids, prioritys, profile_nums, profile_ids);
+    stub_set_shaping_rule_dscp_value(1, 10);//AF11
+    stub_set_shaping_rule_dscp_value(2, 12);//AF12
+    stub_set_shaping_rule_dscp_value(3, 14);//AF13
+    stub_set_shaping_rule_dscp_value(4, 40);//CS5
+
+    ctx->maat_info->rule_table_id = STUB_MAAT_SHAPING_RULE_TABLE_ID;
+    ctx->maat_info->profile_table_id = STUB_MAAT_SHAPING_PROFILE_TABLE_ID;
+
+    memset(&sf, 0, sizeof(sf));
+    sf.priority = SHAPING_PRIORITY_NUM_MAX;
+
+    long long rule_id1[] = {1, 2};
+    shaper_rules_update(&ctx->thread_ctx[0], &sf, rule_id1, 2);
+
+    EXPECT_EQ(sf.rule_num, 2);
+    EXPECT_EQ(sf.dscp_enable, 1);
+    EXPECT_EQ(sf.dscp_value, 10);
+
+    rule_info = &sf.matched_rule_infos[0];
+    EXPECT_EQ(rule_info->id, 1);
+    EXPECT_EQ(rule_info->primary.id, 1);
+    EXPECT_EQ(rule_info->primary.priority, 1);
+    EXPECT_EQ(rule_info->borrowing_num, 0);
+
+    rule_info = &sf.matched_rule_infos[1];
+    EXPECT_EQ(rule_info->id, 2);
+    EXPECT_EQ(rule_info->primary.id, 2);
+    EXPECT_EQ(rule_info->primary.priority, 1);
+    EXPECT_EQ(rule_info->borrowing_num, 0);
+
+    long long rule_id2[] = {3};
+    shaper_rules_update(&ctx->thread_ctx[0], &sf, rule_id2, 1);
+    EXPECT_EQ(sf.rule_num, 3);
+    EXPECT_EQ(sf.dscp_enable, 1);
+    EXPECT_EQ(sf.dscp_value, 10);
+
+    rule_info = &sf.matched_rule_infos[2];
+    EXPECT_EQ(rule_info->id, 3);
+    EXPECT_EQ(rule_info->primary.id, 3);
+    EXPECT_EQ(rule_info->primary.priority, 1);
+    EXPECT_EQ(rule_info->borrowing_num, 0);
+
+    long long rule_id3[] = {4};
+    shaper_rules_update(&ctx->thread_ctx[0], &sf, rule_id3, 3);
+    EXPECT_EQ(sf.rule_num, 4);
+    EXPECT_EQ(sf.dscp_enable, 1);
+    EXPECT_EQ(sf.dscp_value, 40);
+
+    rule_info = &sf.matched_rule_infos[3];
+    EXPECT_EQ(rule_info->id, 4);
+    EXPECT_EQ(rule_info->primary.id, 4);
+    EXPECT_EQ(rule_info->primary.priority, 1);
+    EXPECT_EQ(rule_info->borrowing_num, 0);
+
+    shaping_engine_destroy(ctx);
+}
+
 int main(int argc, char **argv)
 {
 	::testing::InitGoogleTest(&argc, argv);
diff --git a/shaping/test/stub.cpp b/shaping/test/stub.cpp
index 8ba25f9..02b43ea 100644
--- a/shaping/test/stub.cpp
+++ b/shaping/test/stub.cpp
@@ -142,6 +142,13 @@ void stub_set_matched_shaping_rules(int rule_num, long long *rule_id, const int
     return;
 }
 
+void stub_set_shaping_rule_dscp_value(int rule_id, int dscp_value)
+{
+    matched_rules.rules[rule_id].dscp_enable = 1;
+    matched_rules.rules[rule_id].dscp_value = dscp_value;
+    return;
+}
+
 void stub_clear_matched_shaping_rules()
 {
     memset(&matched_rules, 0, sizeof(struct stub_matched_rules));
diff --git a/shaping/test/stub.h b/shaping/test/stub.h
index a00bd01..d466944 100644
--- a/shaping/test/stub.h
+++ b/shaping/test/stub.h
@@ -36,6 +36,7 @@ void stub_set_profile_type(int profile_id, enum shaping_profile_type type);
 void stub_set_async_token_get_times(int profile_id, int times);
 
 void stub_set_matched_shaping_rules(int rule_num, long long *rule_id, const int *priority, const int *profile_num, int profile_id[][MAX_REF_PROFILE]);
+void stub_set_shaping_rule_dscp_value(int rule_id, int dscp_value);
 void stub_clear_matched_shaping_rules();
 
 void stub_send_packet(struct stub_packet *packet);