Fix bugs, improve error handling, and expand test coverage

.github/workflows/release.yml

@@ -1,4 +1,4 @@
-name: Build and release gatewayd-plugin-cache
+name: Build and release gatewayd-plugin-sql-ids-ips
 
 on:
   push:
@@ -13,11 +13,11 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - name: Set up Go 1.24
-        uses: actions/setup-go@v3
+        uses: actions/setup-go@v5
         with:
           go-version: "1.24"
           cache: true
@@ -27,7 +27,7 @@ jobs:
         run: |
           make build-release
       - name: Create release and add artifacts
-        uses: softprops/action-gh-release@v1
+        uses: softprops/action-gh-release@v2
         with:
           files: |
             dist/*.tar.gz

main.go

@@ -4,6 +4,7 @@ import (
 	"flag"
 	"log"
 	"os"
+	"time"
 
 	sdkConfig "github.com/gatewayd-io/gatewayd-plugin-sdk/config"
 	"github.com/gatewayd-io/gatewayd-plugin-sdk/logging"
@@ -62,6 +63,8 @@ func main() {
 		pluginInstance.Impl.ErrorNumber = cast.ToString(cfg["errorNumber"])
 		pluginInstance.Impl.ErrorDetail = cast.ToString(cfg["errorDetail"])
 		pluginInstance.Impl.LogLevel = cast.ToString(cfg["logLevel"])
+		pluginInstance.Impl.PredictionTimeout = time.Duration(
+			cast.ToInt(cfg["predictionTimeout"])) * time.Second
 	}
 
 	goplugin.Serve(&goplugin.ServeConfig{

plugin/constants.go

@@ -1,6 +1,10 @@
 package plugin
 
+import "time"
+
 const (
+	DefaultPredictionTimeout time.Duration = 10 * time.Second
+
 	DecodedQueryField string = "decodedQuery"
 	DetectorField     string = "detector"
 	QueryField        string = "query"

plugin/plugin.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/base64"
 	"encoding/json"
+	"time"
 
 	"github.com/carlmjohnson/requests"
 	"github.com/corazawaf/libinjection-go"
@@ -34,6 +35,7 @@ type Plugin struct {
 	ErrorNumber                string
 	ErrorDetail                string
 	LogLevel                   string
+	PredictionTimeout          time.Duration
 }
 
 type InjectionDetectionPlugin struct {
@@ -107,6 +109,13 @@ func (p *Plugin) OnTrafficFromClient(ctx context.Context, req *v1.Struct) (*v1.S
 	}
 	queryString := cast.ToString(queryMap[StringField])
 
+	timeout := p.PredictionTimeout
+	if timeout == 0 {
+		timeout = DefaultPredictionTimeout
+	}
+	reqCtx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
 	var output map[string]any
 	err = requests.
 		URL(p.PredictionAPIAddress).
@@ -115,7 +124,7 @@ func (p *Plugin) OnTrafficFromClient(ctx context.Context, req *v1.Struct) (*v1.S
 			QueryField: queryString,
 		}).
 		ToJSON(&output).
-		Fetch(context.Background())
+		Fetch(reqCtx)
 	if err != nil {
 		p.Logger.Error("Failed to make POST request", ErrorField, err)
 		if p.isSQLi(queryString) && !p.LibinjectionPermissiveMode {
@@ -194,16 +203,24 @@ func (p *Plugin) prepareResponse(req *v1.Struct, fields map[string]any) *v1.Stru
 		encapsulatedResponse = postgres.ErrorResponse(
 			p.ErrorMessage,
 			p.ErrorSeverity,
-			ErrorNumber,
-			ErrorDetail,
+			p.ErrorNumber,
+			p.ErrorDetail,
 		)
 	} else {
 		// Create a PostgreSQL empty query response.
-		encapsulatedResponse, _ = (&pgproto3.EmptyQueryResponse{}).Encode(nil)
+		var encErr error
+		encapsulatedResponse, encErr = (&pgproto3.EmptyQueryResponse{}).Encode(nil)
+		if encErr != nil {
+			p.Logger.Error("Failed to encode empty query response", ErrorField, encErr)
+			return req
+		}
 	}
 
-	// Create and encode a ready for query response.
-	response, _ := (&pgproto3.ReadyForQuery{TxStatus: 'I'}).Encode(encapsulatedResponse)
+	response, encErr := (&pgproto3.ReadyForQuery{TxStatus: 'I'}).Encode(encapsulatedResponse)
+	if encErr != nil {
+		p.Logger.Error("Failed to encode ready for query response", ErrorField, encErr)
+		return req
+	}
 
 	signals, err := v1.NewList([]any{
 		sdkAct.Terminate().ToMap(),

plugin/plugin_test.go

@@ -68,7 +68,52 @@ func Test_errorResponse(t *testing.T) {
 	assert.Len(t, resp.Fields[sdkAct.Signals].GetListValue().AsSlice(), 2)
 }
 
-func Test_OnTrafficFromClinet(t *testing.T) {
+func Test_emptyResponse(t *testing.T) {
+	p := &Plugin{
+		Logger:       hclog.NewNullLogger(),
+		ResponseType: "empty",
+	}
+
+	query := pgproto3.Query{String: "SELECT * FROM users WHERE id = 1 OR 1=1"}
+	queryBytes, err := query.Encode(nil)
+	require.NoError(t, err)
+
+	req := map[string]any{
+		"request": queryBytes,
+	}
+	reqJSON, err := v1.NewStruct(req)
+	require.NoError(t, err)
+	assert.NotNil(t, reqJSON)
+
+	resp := p.prepareResponse(
+		reqJSON,
+		map[string]any{
+			"score":    0.9999,
+			"detector": "deep_learning_model",
+		},
+	)
+	assert.Equal(t, reqJSON, resp)
+	assert.Len(t, resp.GetFields(), 3)
+	assert.Contains(t, resp.GetFields(), "request")
+	assert.Contains(t, resp.GetFields(), "response")
+	assert.Contains(t, resp.GetFields(), sdkAct.Signals)
+	assert.Len(t, resp.Fields[sdkAct.Signals].GetListValue().AsSlice(), 2)
+}
+
+func Test_GetPluginConfig(t *testing.T) {
+	p := &Plugin{
+		Logger: hclog.NewNullLogger(),
+	}
+
+	result, err := p.GetPluginConfig(context.Background(), nil)
+	require.NoError(t, err)
+	assert.NotNil(t, result)
+	assert.Contains(t, result.GetFields(), "id")
+	assert.Contains(t, result.GetFields(), "hooks")
+	assert.Contains(t, result.GetFields(), "config")
+}
+
+func Test_OnTrafficFromClient(t *testing.T) {
 	p := &Plugin{
 		Logger: hclog.NewNullLogger(),
 	}
@@ -116,7 +161,7 @@ func Test_OnTrafficFromClinet(t *testing.T) {
 	assert.Len(t, resp.Fields[sdkAct.Signals].GetListValue().AsSlice(), 2)
 }
 
-func Test_OnTrafficFromClinetFailedTokenization(t *testing.T) {
+func Test_OnTrafficFromClientFailedTokenization(t *testing.T) {
 	plugins := []*Plugin{
 		{
 			Logger: hclog.NewNullLogger(),
@@ -179,64 +224,62 @@ func Test_OnTrafficFromClinetFailedTokenization(t *testing.T) {
 	}
 }
 
-func Test_OnTrafficFromClinetFailedPrediction(t *testing.T) {
-	plugins := []*Plugin{
-		{
-			Logger: hclog.NewNullLogger(),
-			// If libinjection is disabled, the response should not contain the "response" field,
-			// and the "signals" field, which means the plugin will not terminate the request.
-			EnableLibinjection: false,
-		},
-		{
-			Logger: hclog.NewNullLogger(),
-			// If libinjection is enabled, the response should contain the "response" field,
-			// and the "signals" field, which means the plugin will terminate the request.
-			EnableLibinjection: true,
-		},
+func Test_OnTrafficFromClientBelowThreshold(t *testing.T) {
+	p := &Plugin{
+		Logger:    hclog.NewNullLogger(),
+		Threshold: 0.8,
 	}
+
 	server := httptest.NewServer(
 		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			switch r.URL.Path {
 			case PredictPath:
-				w.WriteHeader(http.StatusInternalServerError)
+				w.WriteHeader(http.StatusOK)
+				w.Header().Set("Content-Type", "application/json")
+				resp := map[string]any{"confidence": 0.1}
+				data, _ := json.Marshal(resp)
+				_, err := w.Write(data)
+				require.NoError(t, err)
 			default:
 				w.WriteHeader(http.StatusNotFound)
 			}
 		}),
 	)
 	defer server.Close()
 
-	for i := range plugins {
-		plugins[i].PredictionAPIAddress = server.URL
+	p.PredictionAPIAddress = server.URL
 
-		query := pgproto3.Query{String: "SELECT * FROM users WHERE id = 1 OR 1=1"}
-		queryBytes, err := query.Encode(nil)
-		require.NoError(t, err)
+	query := pgproto3.Query{String: "SELECT name FROM products WHERE id = 42"}
+	queryBytes, err := query.Encode(nil)
+	require.NoError(t, err)
 
-		req := map[string]any{
-			"request": queryBytes,
-		}
-		reqJSON, err := v1.NewStruct(req)
-		require.NoError(t, err)
-		assert.NotNil(t, reqJSON)
+	req := map[string]any{
+		"request": queryBytes,
+	}
+	reqJSON, err := v1.NewStruct(req)
+	require.NoError(t, err)
 
-		resp, err := plugins[i].OnTrafficFromClient(context.Background(), reqJSON)
-		require.NoError(t, err)
-		assert.NotNil(t, resp)
-		if plugins[i].EnableLibinjection {
-			assert.Len(t, resp.GetFields(), 4)
-			assert.Contains(t, resp.GetFields(), "request")
-			assert.Contains(t, resp.GetFields(), "query")
-			assert.Contains(t, resp.GetFields(), "response")
-			assert.Contains(t, resp.GetFields(), sdkAct.Signals)
-			// 2 signals: Terminate and Log.
-			assert.Len(t, resp.Fields[sdkAct.Signals].GetListValue().AsSlice(), 2)
-		} else {
-			assert.Len(t, resp.GetFields(), 2)
-			assert.Contains(t, resp.GetFields(), "request")
-			assert.Contains(t, resp.GetFields(), "query")
-			assert.NotContains(t, resp.GetFields(), "response")
-			assert.NotContains(t, resp.GetFields(), sdkAct.Signals)
-		}
+	resp, err := p.OnTrafficFromClient(context.Background(), reqJSON)
+	require.NoError(t, err)
+	assert.NotNil(t, resp)
+	assert.NotContains(t, resp.GetFields(), "response")
+	assert.NotContains(t, resp.GetFields(), sdkAct.Signals)
+}
+
+func Test_OnTrafficFromClientEmptyQuery(t *testing.T) {
+	p := &Plugin{
+		Logger: hclog.NewNullLogger(),
 	}
+
+	req := map[string]any{
+		"request": []byte{},
+	}
+	reqJSON, err := v1.NewStruct(req)
+	require.NoError(t, err)
+
+	resp, err := p.OnTrafficFromClient(context.Background(), reqJSON)
+	assert.NoError(t, err)
+	assert.NotNil(t, resp)
+	assert.NotContains(t, resp.GetFields(), "response")
+	assert.NotContains(t, resp.GetFields(), sdkAct.Signals)
 }