diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index f5157c5..a6b762e 100644 --- a/.github/workflows/release.yml +++ b/.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 diff --git a/main.go b/main.go index 740ae6d..76f8e56 100644 --- a/main.go +++ b/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{ diff --git a/plugin/constants.go b/plugin/constants.go index 2897925..8186021 100644 --- a/plugin/constants.go +++ b/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" diff --git a/plugin/plugin.go b/plugin/plugin.go index 2be7da1..970fe44 100644 --- a/plugin/plugin.go +++ b/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(), diff --git a/plugin/plugin_test.go b/plugin/plugin_test.go index fdcb19f..6fd4dd2 100644 --- a/plugin/plugin_test.go +++ b/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,26 +224,22 @@ 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) } @@ -206,37 +247,39 @@ func Test_OnTrafficFromClinetFailedPrediction(t *testing.T) { ) 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) }