diff --git a/api/api.go b/api/api.go index b86d863d..f611d9d7 100644 --- a/api/api.go +++ b/api/api.go @@ -18,6 +18,12 @@ type CallContext interface { SetData(interface{}) // Get the data field, can be used to pass information between OnEnter&OnExit GetData() interface{} + // Get the map data field by key + GetKeyData(key string) interface{} + // Set the map data field by key + SetKeyData(key string, val interface{}) + // Has the map data field by key + HasKeyData(key string) bool // Get the original function parameter at index idx GetParam(idx int) interface{} // Change the original function parameter at index idx diff --git a/tool/instrument/template.go b/tool/instrument/template.go index a47727cf..78bf7f31 100644 --- a/tool/instrument/template.go +++ b/tool/instrument/template.go @@ -17,6 +17,27 @@ func (c *CallContextImpl) SetSkipCall(skip bool) { c.SkipCall = skip } func (c *CallContextImpl) IsSkipCall() bool { return c.SkipCall } func (c *CallContextImpl) SetData(data interface{}) { c.Data = data } func (c *CallContextImpl) GetData() interface{} { return c.Data } +func (c *CallContextImpl) GetKeyData(key string) interface{} { + if data == nil { + return nil + } + return data.(map[string]interface{})[key] +} +func (c *CallContextImpl) SetKeyData(key string, val interface{}) { + if data == nil { + data = make(map[string]interface{}) + } + data.(map[string]interface{})[key] = val +} + +func (c *CallContextImpl) HasKeyData(key string) bool { + if data == nil { + return false + } + _, ok := data.(map[string]interface{})[key] + return ok +} + func (c *CallContextImpl) GetParam(idx int) interface{} { switch idx { }