Include more fields in the Act log

Refactor constants

Author: mostafa

plugin/metrics.go

@@ -20,11 +20,11 @@ var (
 		Name:      "on_traffic_from_client_total",
 		Help:      "The total number of calls to the onTrafficFromClient method",
 	})
-	Detections = promauto.NewCounter(prometheus.CounterOpts{
+	Detections = promauto.NewCounterVec(prometheus.CounterOpts{
 		Namespace: metrics.Namespace,
 		Name:      "detections_total",
 		Help:      "The total number of malicious requests detected",
-	})
+	}, []string{"detector"})
 	Preventions = promauto.NewCounter(prometheus.CounterOpts{
 		Namespace: metrics.Namespace,
 		Name:      "preventions_total",

plugin/plugin.go

@@ -19,6 +19,32 @@ import (
 	"google.golang.org/grpc"
 )
 
+const (
+	DecodedQueryField string = "decodedQuery"
+	DetectorField     string = "detector"
+	ScoreField        string = "score"
+	QueryField        string = "query"
+	ErrorField        string = "error"
+	IsInjectionField  string = "is_injection"
+	ResponseField     string = "response"
+	OutputsField      string = "outputs"
+	TokensField       string = "tokens"
+	RequestField      string = "request"
+	StringField       string = "String"
+
+	DeepLearningModel string = "deep_learning_model"
+	Libinjection      string = "libinjection"
+
+	ErrorLevel           string = "error"
+	ExceptionLevel       string = "EXCEPTION"
+	ErrorNumber          string = "42000"
+	DetectionMessage     string = "SQL injection detected"
+	ErrorResponseMessage string = "Back off, you're not welcome here."
+
+	TokenizeAndSequencePath string = "/tokenize_and_sequence"
+	PredictPath             string = "/v1/models/%s/versions/%s:predict"
+)
+
 type Plugin struct {
 	goplugin.GRPCPlugin
 	v1.GatewayDPluginServiceServer
@@ -44,7 +70,9 @@ func (p *InjectionDetectionPlugin) GRPCServer(b *goplugin.GRPCBroker, s *grpc.Se
 }
 
 // GRPCClient returns the plugin client.
-func (p *InjectionDetectionPlugin) GRPCClient(ctx context.Context, b *goplugin.GRPCBroker, c *grpc.ClientConn) (interface{}, error) {
+func (p *InjectionDetectionPlugin) GRPCClient(
+	ctx context.Context, b *goplugin.GRPCBroker, c *grpc.ClientConn,
+) (any, error) {
 	return v1.NewGatewayDPluginServiceClient(c), nil
 }
 
@@ -73,92 +101,119 @@ func (p *Plugin) OnTrafficFromClient(ctx context.Context, req *v1.Struct) (*v1.S
 	// Handle the client message.
 	req, err := postgres.HandleClientMessage(req, p.Logger)
 	if err != nil {
-		p.Logger.Debug("Failed to handle client message", "error", err)
+		p.Logger.Debug("Failed to handle client message", ErrorField, err)
 	}
 
 	// Get the client request from the GatewayD request.
-	request := cast.ToString(sdkPlugin.GetAttr(req, "request", ""))
+	request := cast.ToString(sdkPlugin.GetAttr(req, RequestField, ""))
 	if request == "" {
 		return req, nil
 	}
 
 	// Get the query from the request.
-	query := cast.ToString(sdkPlugin.GetAttr(req, "query", ""))
+	query := cast.ToString(sdkPlugin.GetAttr(req, QueryField, ""))
 	if query == "" {
 		p.Logger.Debug("Failed to get query from request, possibly not a SQL query request")
 		return req, nil
 	}
-	p.Logger.Trace("Query", "query", query)
+	p.Logger.Trace("Query", QueryField, query)
 
 	// Decode the query.
 	decodedQuery, err := base64.StdEncoding.DecodeString(query)
 	if err != nil {
 		return req, err
 	}
-	p.Logger.Trace("Decoded Query", "decodedQuery", decodedQuery)
+	p.Logger.Trace("Decoded Query", DecodedQueryField, decodedQuery)
 
 	// Unmarshal query into a map.
-	var queryMap map[string]interface{}
+	var queryMap map[string]any
 	if err := json.Unmarshal(decodedQuery, &queryMap); err != nil {
-		p.Logger.Error("Failed to unmarshal query", "error", err)
+		p.Logger.Error("Failed to unmarshal query", ErrorField, err)
 		return req, nil
 	}
-	queryString := cast.ToString(queryMap["String"])
+	queryString := cast.ToString(queryMap[StringField])
 
-	var tokens map[string]interface{}
+	var tokens map[string]any
 	err = requests.
 		URL(p.TokenizerAPIAddress).
-		Path("/tokenize_and_sequence").
-		BodyJSON(map[string]interface{}{
-			"query": queryString,
+		Path(TokenizeAndSequencePath).
+		BodyJSON(map[string]any{
+			QueryField: queryString,
 		}).
 		ToJSON(&tokens).
 		Fetch(context.Background())
 	if err != nil {
-		p.Logger.Error("Failed to make POST request", "error", err)
+		p.Logger.Error("Failed to make POST request", ErrorField, err)
 		if p.isSQLi(queryString) && !p.LibinjectionPermissiveMode {
-			return p.errorResponse(req, queryString), nil
+			return p.errorResponse(
+				req,
+				map[string]any{
+					QueryField:    queryString,
+					DetectorField: Libinjection,
+					ErrorField:    "Failed to make POST request to tokenizer API",
+				},
+			), nil
 		}
 		return req, nil
 	}
 
-	var output map[string]interface{}
+	var output map[string]any
 	err = requests.
 		URL(p.ServingAPIAddress).
-		Path(fmt.Sprintf("/v1/models/%s/versions/%s:predict", p.ModelName, p.ModelVersion)).
-		BodyJSON(map[string]interface{}{
-			"inputs": []interface{}{cast.ToSlice(tokens["tokens"])},
+		Path(fmt.Sprintf(PredictPath, p.ModelName, p.ModelVersion)).
+		BodyJSON(map[string]any{
+			"inputs": []any{cast.ToSlice(tokens[TokensField])},
 		}).
 		ToJSON(&output).
 		Fetch(context.Background())
 	if err != nil {
-		p.Logger.Error("Failed to make POST request", "error", err)
+		p.Logger.Error("Failed to make POST request", ErrorField, err)
 		if p.isSQLi(queryString) && !p.LibinjectionPermissiveMode {
-			return p.errorResponse(req, queryString), nil
+			return p.errorResponse(
+				req,
+				map[string]any{
+					QueryField:    queryString,
+					DetectorField: Libinjection,
+					ErrorField:    "Failed to make POST request to serving API",
+				},
+			), nil
 		}
 		return req, nil
 	}
 
-	predictions := cast.ToSlice(output["outputs"])
+	predictions := cast.ToSlice(output[OutputsField])
 	scores := cast.ToSlice(predictions[0])
 	score := cast.ToFloat32(scores[0])
-	p.Logger.Trace("Deep learning model prediction", "score", score)
+	p.Logger.Trace("Deep learning model prediction", ScoreField, score)
 
 	// Check the prediction against the threshold,
 	// otherwise check if the query is an SQL injection using libinjection.
 	injection := p.isSQLi(queryString)
 	if score >= p.Threshold {
 		if p.EnableLibinjection && !injection {
-			p.Logger.Debug("False positive detected by libinjection")
+			p.Logger.Debug("False positive detected", DetectorField, Libinjection)
 		}
 
-		Detections.Inc()
-		p.Logger.Warn("SQL injection detected by deep learning model", "score", score)
-		return p.errorResponse(req, queryString), nil
+		Detections.With(map[string]string{DetectorField: DeepLearningModel}).Inc()
+		p.Logger.Warn(DetectionMessage, ScoreField, score, DetectorField, DeepLearningModel)
+		return p.errorResponse(
+			req,
+			map[string]any{
+				QueryField:    queryString,
+				ScoreField:    score,
+				DetectorField: DeepLearningModel,
+			},
+		), nil
 	} else if p.EnableLibinjection && injection && !p.LibinjectionPermissiveMode {
-		Detections.Inc()
-		p.Logger.Warn("SQL injection detected by libinjection")
-		return p.errorResponse(req, queryString), nil
+		Detections.With(map[string]string{DetectorField: Libinjection}).Inc()
+		p.Logger.Warn(DetectionMessage, DetectorField, Libinjection)
+		return p.errorResponse(
+			req,
+			map[string]any{
+				QueryField:    queryString,
+				DetectorField: Libinjection,
+			},
+		), nil
 	} else {
 		p.Logger.Trace("No SQL injection detected")
 	}
@@ -175,45 +230,43 @@ func (p *Plugin) isSQLi(query string) bool {
 	// Check if the query is an SQL injection using libinjection.
 	injection, _ := libinjection.IsSQLi(query)
 	if injection {
-		p.Logger.Warn("SQL injection detected by libinjection")
+		p.Logger.Warn(DetectionMessage, DetectorField, Libinjection)
 	}
-	p.Logger.Trace("SQLInjection", "is_injection", cast.ToString(injection))
+	p.Logger.Trace("SQLInjection", IsInjectionField, cast.ToString(injection))
 	return injection
 }
 
-func (p *Plugin) errorResponse(req *v1.Struct, queryString string) *v1.Struct {
+func (p *Plugin) errorResponse(req *v1.Struct, fields map[string]any) *v1.Struct {
 	Preventions.Inc()
 
 	// Create a PostgreSQL error response.
 	errResp := postgres.ErrorResponse(
-		"SQL injection detected",
-		"EXCEPTION",
-		"42000",
-		"Back off, you're not welcome here.",
+		DetectionMessage,
+		ExceptionLevel,
+		ErrorNumber,
+		ErrorResponseMessage,
 	)
 
 	// Create a ready for query response.
 	readyForQuery := &pgproto3.ReadyForQuery{TxStatus: 'I'}
 	// TODO: Decide whether to terminate the connection.
 	response, err := readyForQuery.Encode(errResp)
 	if err != nil {
-		p.Logger.Error("Failed to encode ready for query response", "error", err)
+		p.Logger.Error("Failed to encode ready for query response", ErrorField, err)
 		return req
 	}
 
 	signals, err := v1.NewList([]any{
 		sdkAct.Terminate().ToMap(),
-		sdkAct.Log("error", "SQL injection detected", map[string]any{
-			"query": queryString,
-		}).ToMap(),
+		sdkAct.Log(ErrorLevel, DetectionMessage, fields).ToMap(),
 	})
 	if err != nil {
-		p.Logger.Error("Failed to create signals", "error", err)
+		p.Logger.Error("Failed to create signals", ErrorField, err)
 		return req
 	}
 
 	// Create a response to send back to the client.
 	req.Fields[sdkAct.Signals] = v1.NewListValue(signals)
-	req.Fields["response"] = v1.NewBytesValue(response)
+	req.Fields[ResponseField] = v1.NewBytesValue(response)
 	return req
 }