@@ -53,6 +53,7 @@ const (
5353 DefaultStreamingResponseJsonPath = "choices.0.delta.content"
5454 DefaultDenyCode = 200
5555 DefaultDenyMessage = "很抱歉，我无法回答您的问题"
56+ DefaultTimeout = 2000
5657
5758 AliyunUserAgent = "CIPFrom/AIGateway"
5859 LengthLimit = 1800
@@ -100,6 +101,7 @@ type AISecurityConfig struct {
100101 denyMessage string
101102 protocolOriginal bool
102103 riskLevelBar string
104+ timeout uint32
103105 metrics map [string ]proxywasm.MetricCounter
104106}
105107
@@ -225,6 +227,11 @@ func parseConfig(json gjson.Result, config *AISecurityConfig, log wrapper.Log) e
225227 } else {
226228 config .riskLevelBar = HighRisk
227229 }
230+ if obj := json .Get ("timeout" ); obj .Exists () {
231+ config .timeout = uint32 (obj .Int ())
232+ } else {
233+ config .timeout = DefaultTimeout
234+ }
228235 config .client = wrapper .NewClusterClient (wrapper.FQDNCluster {
229236 FQDN : serviceName ,
230237 Port : servicePort ,
@@ -253,6 +260,7 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, config AISecurityConfig, log
253260
254261func onHttpRequestBody (ctx wrapper.HttpContext , config AISecurityConfig , body []byte , log wrapper.Log ) types.Action {
255262 log .Debugf ("checking request body..." )
263+ startTime := time .Now ().UnixMilli ()
256264 content := gjson .GetBytes (body , config .requestContentJsonPath ).String ()
257265 model := gjson .GetBytes (body , "model" ).String ()
258266 ctx .SetContext ("requestModel" , model )
@@ -279,6 +287,10 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config AISecurityConfig, body []
279287 }
280288 if riskLevelToInt (response .Data .RiskLevel ) < riskLevelToInt (config .riskLevelBar ) {
281289 if contentIndex >= len (content ) {
290+ endTime := time .Now ().UnixMilli ()
291+ ctx .SetUserAttribute ("safecheck_request_rt" , endTime - startTime )
292+ ctx .SetUserAttribute ("safecheck_status" , "request pass" )
293+ ctx .WriteUserAttributeToLogWithKey (wrapper .AILogKey )
282294 proxywasm .ResumeHttpRequest ()
283295 } else {
284296 singleCall ()
@@ -305,7 +317,9 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config AISecurityConfig, body []
305317 }
306318 ctx .DontReadResponseBody ()
307319 config .incrementCounter ("ai_sec_request_deny" , 1 )
308- ctx .SetUserAttribute ("safecheck_status" , "request deny" )
320+ endTime := time .Now ().UnixMilli ()
321+ ctx .SetUserAttribute ("safecheck_request_rt" , endTime - startTime )
322+ ctx .SetUserAttribute ("safecheck_status" , "reqeust deny" )
309323 if response .Data .Advice != nil {
310324 ctx .SetUserAttribute ("safecheck_riskLabel" , response .Data .Result [0 ].Label )
311325 ctx .SetUserAttribute ("safecheck_riskWords" , response .Data .Result [0 ].RiskWords )
@@ -345,7 +359,7 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config AISecurityConfig, body []
345359 reqParams .Add (k , v )
346360 }
347361 reqParams .Add ("Signature" , signature )
348- err := config .client .Post (fmt .Sprintf ("/?%s" , reqParams .Encode ()), [][2 ]string {{"User-Agent" , AliyunUserAgent }}, nil , callback )
362+ err := config .client .Post (fmt .Sprintf ("/?%s" , reqParams .Encode ()), [][2 ]string {{"User-Agent" , AliyunUserAgent }}, nil , callback , config . timeout )
349363 if err != nil {
350364 log .Errorf ("failed call the safe check service: %v" , err )
351365 proxywasm .ResumeHttpRequest ()
@@ -364,20 +378,6 @@ func convertHeaders(hs [][2]string) map[string][]string {
364378 return ret
365379}
366380
367- // headers: map[string][]string -> [][2]string
368- func reconvertHeaders (hs map [string ][]string ) [][2 ]string {
369- var ret [][2 ]string
370- for k , vs := range hs {
371- for _ , v := range vs {
372- ret = append (ret , [2 ]string {k , v })
373- }
374- }
375- sort .SliceStable (ret , func (i , j int ) bool {
376- return ret [i ][0 ] < ret [j ][0 ]
377- })
378- return ret
379- }
380-
381381func onHttpResponseHeaders (ctx wrapper.HttpContext , config AISecurityConfig , log wrapper.Log ) types.Action {
382382 if ! config .checkResponse {
383383 log .Debugf ("response checking is disabled" )
@@ -401,6 +401,7 @@ func onHttpResponseHeaders(ctx wrapper.HttpContext, config AISecurityConfig, log
401401
402402func onHttpResponseBody (ctx wrapper.HttpContext , config AISecurityConfig , body []byte , log wrapper.Log ) types.Action {
403403 log .Debugf ("checking response body..." )
404+ startTime := time .Now ().UnixMilli ()
404405 hdsMap := ctx .GetContext ("headers" ).(map [string ][]string )
405406 isStreamingResponse := strings .Contains (strings .Join (hdsMap ["content-type" ], ";" ), "event-stream" )
406407 model := ctx .GetStringContext ("requestModel" , "unknown" )
@@ -433,6 +434,10 @@ func onHttpResponseBody(ctx wrapper.HttpContext, config AISecurityConfig, body [
433434 }
434435 if riskLevelToInt (response .Data .RiskLevel ) < riskLevelToInt (config .riskLevelBar ) {
435436 if contentIndex >= len (content ) {
437+ endTime := time .Now ().UnixMilli ()
438+ ctx .SetUserAttribute ("safecheck_response_rt" , endTime - startTime )
439+ ctx .SetUserAttribute ("safecheck_status" , "response pass" )
440+ ctx .WriteUserAttributeToLogWithKey (wrapper .AILogKey )
436441 proxywasm .ResumeHttpResponse ()
437442 } else {
438443 singleCall ()
@@ -458,6 +463,8 @@ func onHttpResponseBody(ctx wrapper.HttpContext, config AISecurityConfig, body [
458463 proxywasm .SendHttpResponse (uint32 (config .denyCode ), [][2 ]string {{"content-type" , "application/json" }}, jsonData , - 1 )
459464 }
460465 config .incrementCounter ("ai_sec_response_deny" , 1 )
466+ endTime := time .Now ().UnixMilli ()
467+ ctx .SetUserAttribute ("safecheck_response_rt" , endTime - startTime )
461468 ctx .SetUserAttribute ("safecheck_status" , "response deny" )
462469 if response .Data .Advice != nil {
463470 ctx .SetUserAttribute ("safecheck_riskLabel" , response .Data .Result [0 ].Label )
@@ -498,7 +505,7 @@ func onHttpResponseBody(ctx wrapper.HttpContext, config AISecurityConfig, body [
498505 reqParams .Add (k , v )
499506 }
500507 reqParams .Add ("Signature" , signature )
501- err := config .client .Post (fmt .Sprintf ("/?%s" , reqParams .Encode ()), [][2 ]string {{"User-Agent" , AliyunUserAgent }}, nil , callback )
508+ err := config .client .Post (fmt .Sprintf ("/?%s" , reqParams .Encode ()), [][2 ]string {{"User-Agent" , AliyunUserAgent }}, nil , callback , config . timeout )
502509 if err != nil {
503510 log .Errorf ("failed call the safe check service: %v" , err )
504511 proxywasm .ResumeHttpResponse ()
0 commit comments