From 83aa87af4a6d0802b13ff58680ed74eed995ff93 Mon Sep 17 00:00:00 2001
From: childe <rmself@qq.com>
Date: Mon, 31 Aug 2020 11:54:52 +0800
Subject: [PATCH] fix #99
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

建议添加KV模板的include_keys方法
---
 filter/kv.go      | 47 +++++++++++++++++++++++++++++++++++++----------
 filter/kv_test.go | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 10 deletions(-)

diff --git a/filter/kv.go b/filter/kv.go
index 53a0094c..0116bdc1 100644
--- a/filter/kv.go
+++ b/filter/kv.go
@@ -9,14 +9,16 @@ import (
 )
 
 type KVFilter struct {
-	config      map[interface{}]interface{}
-	fields      map[field_setter.FieldSetter]value_render.ValueRender
-	src         value_render.ValueRender
-	target      string
-	field_split string
-	value_split string
-	trim        string
-	trim_key    string
+	config       map[interface{}]interface{}
+	fields       map[field_setter.FieldSetter]value_render.ValueRender
+	src          value_render.ValueRender
+	target       string
+	field_split  string
+	value_split  string
+	trim         string
+	trim_key     string
+	include_keys map[string]bool
+	exclude_keys map[string]bool
 }
 
 func (l *MethodLibrary) NewKVFilter(config map[interface{}]interface{}) *KVFilter {
@@ -60,6 +62,21 @@ func (l *MethodLibrary) NewKVFilter(config map[interface{}]interface{}) *KVFilte
 	} else {
 		plugin.trim_key = ""
 	}
+
+	plugin.include_keys = make(map[string]bool)
+	if include_keys, ok := config["include_keys"]; ok {
+		for _, k := range include_keys.([]interface{}) {
+			plugin.include_keys[k.(string)] = true
+		}
+	}
+
+	plugin.exclude_keys = make(map[string]bool)
+	if exclude_keys, ok := config["exclude_keys"]; ok {
+		for _, k := range exclude_keys.([]interface{}) {
+			plugin.exclude_keys[k.(string)] = true
+		}
+	}
+
 	return plugin
 }
 
@@ -79,14 +96,24 @@ func (p *KVFilter) Filter(event map[string]interface{}) (map[string]interface{},
 		event[p.target] = o
 	}
 
-	success := true
+	var success bool = true
+	var key string
 	for _, kv := range A {
 		a := strings.SplitN(kv, p.value_split, 2)
 		if len(a) != 2 {
 			success = false
 			continue
 		}
-		o[strings.Trim(a[0], p.trim_key)] = strings.Trim(a[1], p.trim)
+
+		key = strings.Trim(a[0], p.trim_key)
+
+		if _, ok := p.exclude_keys[key]; ok {
+			continue
+		}
+
+		if _, ok := p.include_keys[key]; len(p.include_keys) == 0 || ok {
+			o[key] = strings.Trim(a[1], p.trim)
+		}
 	}
 	return event, success
 }
diff --git a/filter/kv_test.go b/filter/kv_test.go
index 57e68c7c..6ca5fb33 100644
--- a/filter/kv_test.go
+++ b/filter/kv_test.go
@@ -2,6 +2,42 @@ package filter
 
 import "testing"
 
+func TestIncludeKeys(t *testing.T) {
+	config := make(map[interface{}]interface{})
+	config["field_split"] = " "
+	config["value_split"] = "="
+	config["src"] = "message"
+	config["include_keys"] = []interface{}{"a", "b", "c", "xyz"}
+	config["exclude_keys"] = []interface{}{"c"}
+	f := methodLibrary.NewKVFilter(config)
+
+	event := make(map[string]interface{})
+	event["message"] = "a=aaa b=bbb c=ccc xyz=\txyzxyz\t d=ddd"
+	t.Log(event)
+
+	event, ok := f.Filter(event)
+	if !ok {
+		t.Error("kv failed")
+	}
+	t.Log(event)
+
+	if event["a"] != "aaa" {
+		t.Error("kv failed")
+	}
+	if event["b"] != "bbb" {
+		t.Error("kv failed")
+	}
+	if _, ok := event["c"]; ok {
+		t.Error("c is excluded")
+	}
+	if event["xyz"] != "\txyzxyz\t" {
+		t.Error("kv failed")
+	}
+	if _, ok := event["d"]; ok {
+		t.Error("d is excluded")
+	}
+}
+
 func TestKVFilter(t *testing.T) {
 	config := make(map[interface{}]interface{})
 	config["field_split"] = " "