Merge pull request #17 from iamelevich/feature/custom-skipper

Add `SetSkipper` method to `Plugin`
iamelevich · Apr 1, 2023 · aabfaea · aabfaea
2 parents 76d8ad6 + 15c9c8f
commit aabfaea
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 13 deletions.
diff --git a/plugin.go b/plugin.go
@@ -12,6 +12,11 @@ import (
 	"github.com/pocketbase/pocketbase/core"
 )
 
+// DefaultSkipper skip proxy middleware for requests, where path starts with /_/ or /api/.
+func DefaultSkipper(c echo.Context) bool {
+	return strings.HasPrefix(c.Request().URL.Path, "/_/") || strings.HasPrefix(c.Request().URL.Path, "/api/")
+}
+
 // Options defines optional struct to customize the default plugin behavior.
 type Options struct {
 	// Enabled defines if proxy should be enabled.
@@ -35,6 +40,9 @@ type Plugin struct {
 
 	// parsedUrl from options.Url
 	parsedUrl *url.URL
+
+	// Skipper function for proxy middleware
+	skipper middleware.Skipper
 }
 
 // Validate plugin options. Return error if some option is invalid.
@@ -67,22 +75,34 @@ func (p *Plugin) Validate() error {
 	return nil
 }
 
-func (p *Plugin) enableProxy(e *core.ServeEvent) error {
-	if p.options.Enabled {
-		// Skip PocketBase routes
-		skipperFunc := func(c echo.Context) bool {
-			return strings.HasPrefix(c.Request().URL.Path, "/_/") || strings.HasPrefix(c.Request().URL.Path, "/api/")
-		}
+/*
+SetSkipper set skipper function that should return true if that route shouldn't be proxied.
+
+If not set, the DefaultSkipper is used:
+
+If set - you should also control the middleware behavior for /_/ and /api/ routes.
+
+Example:
 
+	plugin.SetSkipper(func(c echo.Context) bool {
+		return c.Request().URL.Path == "/my-super-secret-route"
+	})
+*/
+func (p *Plugin) SetSkipper(skipper middleware.Skipper) {
+	p.skipper = skipper
+}
+
+func (p *Plugin) enableProxy(e *core.ServeEvent) {
+	if p.options.Enabled {
 		if p.options.ProxyLogsEnabled {
 			e.Router.Use(middleware.LoggerWithConfig(middleware.LoggerConfig{
-				Skipper: skipperFunc,
+				Skipper: p.skipper,
 			}))
 		} else {
 			log.Println("Proxy logs are disabled")
 		}
 		e.Router.Use(middleware.ProxyWithConfig(middleware.ProxyConfig{
-			Skipper: skipperFunc,
+			Skipper: p.skipper,
 			Balancer: middleware.NewRoundRobinBalancer([]*middleware.ProxyTarget{
 				{
 					URL: p.parsedUrl,
@@ -100,7 +120,6 @@ func (p *Plugin) enableProxy(e *core.ServeEvent) error {
 			color.CyanString("%s", p.parsedUrl.String()),
 		)
 	}
-	return nil
 }
 
 // MustRegister is a helper function that registers plugin and panics if error occurred.
@@ -114,7 +133,10 @@ func MustRegister(app core.App, options *Options) *Plugin {
 
 // Register registers plugin.
 func Register(app core.App, options *Options) (*Plugin, error) {
-	p := &Plugin{app: app}
+	p := &Plugin{
+		app:     app,
+		skipper: DefaultSkipper,
+	}
 
 	// Set default options
 	if options != nil {
@@ -129,9 +151,7 @@ func Register(app core.App, options *Options) (*Plugin, error) {
 	}
 
 	app.OnBeforeServe().Add(func(e *core.ServeEvent) error {
-		if err := p.enableProxy(e); err != nil {
-			return err
-		}
+		p.enableProxy(e)
 		return nil
 	})
 

diff --git a/plugin_test.go b/plugin_test.go
@@ -4,6 +4,7 @@ import (
 	"net/http"
 	"testing"
 
+	"github.com/labstack/echo/v5"
 	"github.com/pocketbase/pocketbase"
 	"github.com/pocketbase/pocketbase/core"
 	"github.com/pocketbase/pocketbase/tests"
@@ -114,6 +115,13 @@ func TestPlugin_Validate(t *testing.T) {
 	}
 }
 
+func TestPlugin_Register(t *testing.T) {
+	_, err := Register(nil, nil)
+	if err == nil {
+		t.Errorf("Register() should fail when app is nil")
+	}
+}
+
 func TestPlugin_MustRegister(t *testing.T) {
 	// setup the test ApiScenario app instance
 	setupTestApp := func(options *Options) func() (*tests.TestApp, error) {
@@ -151,6 +159,18 @@ func TestPlugin_MustRegister(t *testing.T) {
 				Url:     "http://localhost:1234",
 			}),
 		},
+		{
+			Name:            "/ request should be proxied when enabled and ProxyLogsEnabled",
+			Method:          http.MethodPost,
+			Url:             "/",
+			ExpectedStatus:  200,
+			ExpectedContent: []string{`OK from /`},
+			TestAppFactory: setupTestApp(&Options{
+				Enabled:          true,
+				Url:              "http://localhost:1234",
+				ProxyLogsEnabled: true,
+			}),
+		},
 		{
 			Name:            "/ shouldn be proxied when options nil",
 			Method:          http.MethodPost,
@@ -192,6 +212,30 @@ func TestPlugin_MustRegister(t *testing.T) {
 				Url:     "http://localhost:1234",
 			}),
 		},
+		{
+			Name:            "/my-super-api-path request should not be proxied when enabled with custom skipper",
+			Method:          http.MethodPost,
+			Url:             "/my-super-api-path",
+			ExpectedStatus:  404,
+			ExpectedContent: []string{`"data":{}`},
+			TestAppFactory: func() (*tests.TestApp, error) {
+				testApp, err := tests.NewTestApp()
+				if err != nil {
+					return nil, err
+				}
+
+				p := MustRegister(testApp, &Options{
+					Enabled: true,
+					Url:     "http://localhost:1234",
+				})
+
+				p.SetSkipper(func(c echo.Context) bool {
+					return c.Request().URL.Path == "/my-super-api-path"
+				})
+
+				return testApp, nil
+			},
+		},
 	}
 
 	for _, scenario := range scenarios {