Make error response customizable via env-vars in plugin config

Return empty response if chosen by the user
Make log level of audit trail customizable
Update gatewayd_plugin.yaml to reflect changes
Rename function to prepareRequest

Author: mostafa

gatewayd_plugin.yaml

@@ -47,5 +47,17 @@ plugins:
       # False (strict): The plugin will block the request if it detects an SQL injection attack.
       #                 This greatly increases the false positive rate.
       - LIBINJECTION_PERMISSIVE_MODE=True
+      # The following env-vars are used to configure the plugin's response.
+      # Possiblel values: error, empty or terminate
+      - RESPONSE_TYPE=error
+      # Possible values: DEBUG, LOG, INFO, NOTICE, WARNING, and EXCEPTION
+      - ERROR_SEVERITY=EXCEPTION
+      # Ref: https://www.postgresql.org/docs/current/errcodes-appendix.html
+      - ERROR_NUMBER=42000
+      - ERROR_MESSAGE=SQL injection detected
+      - ERROR_DETAIL=Back off, you're not welcome here.
+      # Possible values: trace, debug, info, warn, error
+      # Other values will result in no level being set.
+      - LOG_LEVEL=error
     # Checksum hash to verify the binary before loading
     checksum: dee4aa014a722e1865d91744a4fd310772152467d9c6ab4ba17fd9dd40d3f724

main.go

@@ -46,6 +46,13 @@ func main() {
 		pluginInstance.Impl.ServingAPIAddress = cast.ToString(cfg["servingAPIAddress"])
 		pluginInstance.Impl.ModelName = cast.ToString(cfg["modelName"])
 		pluginInstance.Impl.ModelVersion = cast.ToString(cfg["modelVersion"])
+
+		pluginInstance.Impl.ResponseType = cast.ToString(cfg["responseType"])
+		pluginInstance.Impl.ErrorMessage = cast.ToString(cfg["errorMessage"])
+		pluginInstance.Impl.ErrorSeverity = cast.ToString(cfg["errorSeverity"])
+		pluginInstance.Impl.ErrorNumber = cast.ToString(cfg["errorNumber"])
+		pluginInstance.Impl.ErrorDetail = cast.ToString(cfg["errorDetail"])
+		pluginInstance.Impl.LogLevel = cast.ToString(cfg["logLevel"])
 	}
 
 	goplugin.Serve(&goplugin.ServeConfig{

plugin/metrics.go

@@ -25,9 +25,9 @@ var (
 		Name:      "detections_total",
 		Help:      "The total number of malicious requests detected",
 	}, []string{"detector"})
-	Preventions = promauto.NewCounter(prometheus.CounterOpts{
+	Preventions = promauto.NewCounterVec(prometheus.CounterOpts{
 		Namespace: metrics.Namespace,
 		Name:      "preventions_total",
 		Help:      "The total number of malicious requests prevented",
-	})
+	}, []string{"response_type"})
 )

plugin/module.go

@@ -45,6 +45,21 @@ var (
 			"threshold":                  sdkConfig.GetEnv("THRESHOLD", "0.8"),
 			"enableLibinjection":         sdkConfig.GetEnv("ENABLE_LIBINJECTION", "true"),
 			"libinjectionPermissiveMode": sdkConfig.GetEnv("LIBINJECTION_MODE", "true"),
+
+			// Possible values: error or empty
+			"responseType": sdkConfig.GetEnv("RESPONSE_TYPE", ResponseType),
+
+			// This is part of the error response and the audit trail
+			"errorMessage": sdkConfig.GetEnv("ERROR_MESSAGE", ErrorMessage),
+
+			// Response type: error
+			// Possible severity values: DEBUG, LOG, INFO, NOTICE, WARNING, and EXCEPTION
+			"errorSeverity": sdkConfig.GetEnv("ERROR_SEVERITY", ErrorSeverity),
+			"errorNumber":   sdkConfig.GetEnv("ERROR_NUMBER", ErrorNumber),
+			"errorDetail":   sdkConfig.GetEnv("ERROR_DETAIL", ErrorDetail),
+
+			// Log an audit trail
+			"logLevel": sdkConfig.GetEnv("LOG_LEVEL", LogLevel),
 		},
 		"hooks": []interface{}{
 			// Converting HookName to int32 is required because the plugin

plugin/plugin.go

@@ -15,6 +15,7 @@ import (
 	"github.com/hashicorp/go-hclog"
 	goplugin "github.com/hashicorp/go-plugin"
 	"github.com/jackc/pgx/v5/pgproto3"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/spf13/cast"
 	"google.golang.org/grpc"
 )
@@ -30,15 +31,17 @@ const (
 	OutputsField      string = "outputs"
 	TokensField       string = "tokens"
 	StringField       string = "String"
+	ResponseTypeField string = "response_type"
 
 	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."
+	ResponseType  string = "error"
+	ErrorSeverity string = "EXCEPTION"
+	ErrorNumber   string = "42000"
+	ErrorMessage  string = "SQL injection detected"
+	ErrorDetail   string = "Back off, you're not welcome here."
+	LogLevel      string = "error"
 
 	TokenizeAndSequencePath string = "/tokenize_and_sequence"
 	PredictPath             string = "/v1/models/%s/versions/%s:predict"
@@ -55,6 +58,12 @@ type Plugin struct {
 	ServingAPIAddress          string
 	ModelName                  string
 	ModelVersion               string
+	ResponseType               string
+	ErrorMessage               string
+	ErrorSeverity              string
+	ErrorNumber                string
+	ErrorDetail                string
+	LogLevel                   string
 }
 
 type InjectionDetectionPlugin struct {
@@ -139,7 +148,7 @@ func (p *Plugin) OnTrafficFromClient(ctx context.Context, req *v1.Struct) (*v1.S
 	if err != nil {
 		p.Logger.Error("Failed to make POST request", ErrorField, err)
 		if p.isSQLi(queryString) && !p.LibinjectionPermissiveMode {
-			return p.errorResponse(
+			return p.prepareResponse(
 				req,
 				map[string]any{
 					QueryField:    queryString,
@@ -163,7 +172,7 @@ func (p *Plugin) OnTrafficFromClient(ctx context.Context, req *v1.Struct) (*v1.S
 	if err != nil {
 		p.Logger.Error("Failed to make POST request", ErrorField, err)
 		if p.isSQLi(queryString) && !p.LibinjectionPermissiveMode {
-			return p.errorResponse(
+			return p.prepareResponse(
 				req,
 				map[string]any{
 					QueryField:    queryString,
@@ -189,8 +198,8 @@ func (p *Plugin) OnTrafficFromClient(ctx context.Context, req *v1.Struct) (*v1.S
 		}
 
 		Detections.With(map[string]string{DetectorField: DeepLearningModel}).Inc()
-		p.Logger.Warn(DetectionMessage, ScoreField, score, DetectorField, DeepLearningModel)
-		return p.errorResponse(
+		p.Logger.Warn(p.ErrorMessage, ScoreField, score, DetectorField, DeepLearningModel)
+		return p.prepareResponse(
 			req,
 			map[string]any{
 				QueryField:    queryString,
@@ -200,8 +209,8 @@ func (p *Plugin) OnTrafficFromClient(ctx context.Context, req *v1.Struct) (*v1.S
 		), nil
 	} else if p.EnableLibinjection && injection && !p.LibinjectionPermissiveMode {
 		Detections.With(map[string]string{DetectorField: Libinjection}).Inc()
-		p.Logger.Warn(DetectionMessage, DetectorField, Libinjection)
-		return p.errorResponse(
+		p.Logger.Warn(p.ErrorMessage, DetectorField, Libinjection)
+		return p.prepareResponse(
 			req,
 			map[string]any{
 				QueryField:    queryString,
@@ -224,35 +233,36 @@ 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(DetectionMessage, DetectorField, Libinjection)
+		p.Logger.Warn(p.ErrorMessage, DetectorField, Libinjection)
 	}
 	p.Logger.Trace("SQLInjection", IsInjectionField, cast.ToString(injection))
 	return injection
 }
 
-func (p *Plugin) errorResponse(req *v1.Struct, fields map[string]any) *v1.Struct {
-	Preventions.Inc()
+func (p *Plugin) prepareResponse(req *v1.Struct, fields map[string]any) *v1.Struct {
+	Preventions.With(prometheus.Labels{ResponseTypeField: p.ResponseType}).Inc()
 
-	// Create a PostgreSQL error response.
-	errResp := postgres.ErrorResponse(
-		DetectionMessage,
-		ExceptionLevel,
-		ErrorNumber,
-		ErrorResponseMessage,
-	)
+	var encapsulatedResponse []byte
 
-	// 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", ErrorField, err)
-		return req
+	if p.ResponseType == "error" {
+		// Create a PostgreSQL error response.
+		encapsulatedResponse = postgres.ErrorResponse(
+			p.ErrorMessage,
+			p.ErrorSeverity,
+			ErrorNumber,
+			ErrorDetail,
+		)
+	} else {
+		// Create a PostgreSQL empty query response.
+		encapsulatedResponse, _ = (&pgproto3.EmptyQueryResponse{}).Encode(nil)
 	}
 
+	// Create and encode a ready for query response.
+	response, _ := (&pgproto3.ReadyForQuery{TxStatus: 'I'}).Encode(encapsulatedResponse)
+
 	signals, err := v1.NewList([]any{
 		sdkAct.Terminate().ToMap(),
-		sdkAct.Log(ErrorLevel, DetectionMessage, fields).ToMap(),
+		sdkAct.Log(p.LogLevel, p.ErrorMessage, fields).ToMap(),
 	})
 	if err != nil {
 		p.Logger.Error("Failed to create signals", ErrorField, err)

plugin/plugin_test.go

@@ -52,7 +52,7 @@ func Test_errorResponse(t *testing.T) {
 	require.NoError(t, err)
 	assert.NotNil(t, reqJSON)
 
-	resp := p.errorResponse(
+	resp := p.prepareResponse(
 		reqJSON,
 		map[string]any{
 			"score":    0.9999,