From 34d9aaefb9642b7a01b9906f6fdc16c7245fa392 Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Sat, 18 Jan 2025 19:59:33 -0800 Subject: [PATCH] Add more new relic transactions for auth. (#147) --- integration-tests/comfy_node_test.go | 2 +- integration-tests/node_ban_test.go | 5 ++-- integration-tests/node_test.go | 2 +- integration-tests/node_version_test.go | 13 +++++----- .../personal_access_token_test.go | 7 +++--- integration-tests/publisher_ban_test.go | 2 +- integration-tests/publisher_test.go | 7 +++--- .../authorization/authorization_manager.go | 25 +++++++++++++++---- .../authorization_middleware_test.go | 3 ++- server/server.go | 2 +- 10 files changed, 44 insertions(+), 24 deletions(-) diff --git a/integration-tests/comfy_node_test.go b/integration-tests/comfy_node_test.go index bdde927..eaa16b4 100644 --- a/integration-tests/comfy_node_test.go +++ b/integration-tests/comfy_node_test.go @@ -22,7 +22,7 @@ func TestRegistryComfyNode(t *testing.T) { // Initialize server implementation and authorization middleware impl := NewStrictServerImplementationWithMocks(client, &config.Config{}) - authz := authorization.NewAuthorizationManager(client, impl.RegistryService).AuthorizationMiddleware() + authz := authorization.NewAuthorizationManager(client, impl.RegistryService, impl.NewRelicApp).AuthorizationMiddleware() // Setup test user context and publisher ctx, _ := setupTestUser(client) diff --git a/integration-tests/node_ban_test.go b/integration-tests/node_ban_test.go index 5d6c2ed..24416db 100644 --- a/integration-tests/node_ban_test.go +++ b/integration-tests/node_ban_test.go @@ -2,7 +2,6 @@ package integration import ( "context" - "google.golang.org/protobuf/proto" "net/http" "registry-backend/config" "registry-backend/drip" @@ -10,6 +9,8 @@ import ( drip_authorization "registry-backend/server/middleware/authorization" "testing" + "google.golang.org/protobuf/proto" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -21,7 +22,7 @@ func TestNodeBan(t *testing.T) { // Setup the mock services and server impl := NewStrictServerImplementationWithMocks(client, &config.Config{}) - authz := drip_authorization.NewAuthorizationManager(client, impl.RegistryService).AuthorizationMiddleware() + authz := drip_authorization.NewAuthorizationManager(client, impl.RegistryService, impl.NewRelicApp).AuthorizationMiddleware() t.Run("Node Ban Tests", func(t *testing.T) { userCtx, _ := setupTestUser(client) diff --git a/integration-tests/node_test.go b/integration-tests/node_test.go index 303b3d6..8c9fa2f 100644 --- a/integration-tests/node_test.go +++ b/integration-tests/node_test.go @@ -19,7 +19,7 @@ func TestRegistryNode(t *testing.T) { // Initialize server implementation and authorization middleware impl := NewStrictServerImplementationWithMocks(client, &config.Config{}) - authz := authorization.NewAuthorizationManager(client, impl.RegistryService).AuthorizationMiddleware() + authz := authorization.NewAuthorizationManager(client, impl.RegistryService, impl.NewRelicApp).AuthorizationMiddleware() // Setup user context and publisher ctx, _ := setupTestUser(client) diff --git a/integration-tests/node_version_test.go b/integration-tests/node_version_test.go index 75b3419..d006024 100644 --- a/integration-tests/node_version_test.go +++ b/integration-tests/node_version_test.go @@ -4,11 +4,6 @@ import ( "context" "encoding/json" "fmt" - "github.com/labstack/echo/v4" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/mock" - "github.com/stretchr/testify/require" - "google.golang.org/protobuf/proto" "net/http" "net/http/httptest" "registry-backend/config" @@ -20,6 +15,12 @@ import ( registry "registry-backend/services/registry" "testing" "time" + + "github.com/labstack/echo/v4" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" ) func TestRegistryNodeVersion(t *testing.T) { @@ -28,7 +29,7 @@ func TestRegistryNodeVersion(t *testing.T) { // Initialize server implementation and authorization middleware impl := NewStrictServerImplementationWithMocks(client, &config.Config{}) - authz := authorization.NewAuthorizationManager(client, impl.RegistryService).AuthorizationMiddleware() + authz := authorization.NewAuthorizationManager(client, impl.RegistryService, impl.NewRelicApp).AuthorizationMiddleware() // Setup user context and publisher ctx, _ := setupTestUser(client) diff --git a/integration-tests/personal_access_token_test.go b/integration-tests/personal_access_token_test.go index 491139b..e361132 100644 --- a/integration-tests/personal_access_token_test.go +++ b/integration-tests/personal_access_token_test.go @@ -2,13 +2,14 @@ package integration import ( "context" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" "registry-backend/config" "registry-backend/drip" authorization "registry-backend/server/middleware/authorization" "strings" "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestRegistryPersonalAccessToken(t *testing.T) { @@ -18,7 +19,7 @@ func TestRegistryPersonalAccessToken(t *testing.T) { // Initialize server implementation and authorization middleware impl := NewStrictServerImplementationWithMocks(client, &config.Config{}) - authz := authorization.NewAuthorizationManager(client, impl.RegistryService).AuthorizationMiddleware() + authz := authorization.NewAuthorizationManager(client, impl.RegistryService, impl.NewRelicApp).AuthorizationMiddleware() // Setup test user and publisher ctx, _ := setupTestUser(client) diff --git a/integration-tests/publisher_ban_test.go b/integration-tests/publisher_ban_test.go index 55caeeb..fe91af7 100644 --- a/integration-tests/publisher_ban_test.go +++ b/integration-tests/publisher_ban_test.go @@ -20,7 +20,7 @@ func TestPublisherBan(t *testing.T) { // Setup the mock services and server impl := NewStrictServerImplementationWithMocks(client, &config.Config{}) - authz := authorization.NewAuthorizationManager(client, impl.RegistryService).AuthorizationMiddleware() + authz := authorization.NewAuthorizationManager(client, impl.RegistryService, impl.NewRelicApp).AuthorizationMiddleware() t.Run("Publisher Ban Tests", func(t *testing.T) { userCtx, _ := setupTestUser(client) diff --git a/integration-tests/publisher_test.go b/integration-tests/publisher_test.go index 7f4f898..cb201fb 100644 --- a/integration-tests/publisher_test.go +++ b/integration-tests/publisher_test.go @@ -2,12 +2,13 @@ package integration import ( "context" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" "registry-backend/config" "registry-backend/drip" authorization "registry-backend/server/middleware/authorization" "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestRegistryPublisher(t *testing.T) { @@ -16,7 +17,7 @@ func TestRegistryPublisher(t *testing.T) { defer cleanup() impl := NewStrictServerImplementationWithMocks(client, &config.Config{}) - authz := authorization.NewAuthorizationManager(client, impl.RegistryService).AuthorizationMiddleware() + authz := authorization.NewAuthorizationManager(client, impl.RegistryService, impl.NewRelicApp).AuthorizationMiddleware() // Create a test user and a random publisher for testing ctx, testUser := setupTestUser(client) diff --git a/server/middleware/authorization/authorization_manager.go b/server/middleware/authorization/authorization_manager.go index 74cf1cd..3339d56 100644 --- a/server/middleware/authorization/authorization_manager.go +++ b/server/middleware/authorization/authorization_manager.go @@ -11,6 +11,7 @@ import ( "github.com/google/uuid" "github.com/labstack/echo/v4" + "github.com/newrelic/go-agent/v3/newrelic" strictecho "github.com/oapi-codegen/runtime/strictmiddleware/echo" "github.com/rs/zerolog/log" ) @@ -34,16 +35,18 @@ type Assertor interface { // AuthorizationManager manages authorization-related tasks type AuthorizationManager struct { - EntClient *ent.Client - Assertor Assertor + EntClient *ent.Client + Assertor Assertor + NewRelicApp *newrelic.Application } // NewAuthorizationManager creates a new instance of AuthorizationManager func NewAuthorizationManager( - entClient *ent.Client, assertor Assertor) *AuthorizationManager { + entClient *ent.Client, assertor Assertor, newRelicApp *newrelic.Application) *AuthorizationManager { return &AuthorizationManager{ - EntClient: entClient, - Assertor: assertor, + EntClient: entClient, + Assertor: assertor, + NewRelicApp: newRelicApp, } } @@ -52,6 +55,8 @@ func (m *AuthorizationManager) assertUserBanned() drip.StrictMiddlewareFunc { return func(f strictecho.StrictEchoHandlerFunc, operationID string) strictecho.StrictEchoHandlerFunc { return func(c echo.Context, request interface{}) (response interface{}, err error) { ctx := c.Request().Context() + txn := m.NewRelicApp.StartTransaction("AuthorizationManager.assertUserBanned") + defer txn.End() v := ctx.Value(drip_authentication.UserContextKey) userDetails, ok := v.(*drip_authentication.UserDetails) if !ok { @@ -78,6 +83,8 @@ func (m *AuthorizationManager) assertPublisherPermission( return func(f strictecho.StrictEchoHandlerFunc, operationID string) strictecho.StrictEchoHandlerFunc { return func(c echo.Context, request interface{}) (response interface{}, err error) { ctx := c.Request().Context() + txn := m.NewRelicApp.StartTransaction("AuthorizationManager.assertPublisherPermission") + defer txn.End() v := ctx.Value(drip_authentication.UserContextKey) userDetails, ok := v.(*drip_authentication.UserDetails) if !ok { @@ -114,6 +121,8 @@ func (m *AuthorizationManager) assertNodeBanned(extractor func(req interface{}) return func(f strictecho.StrictEchoHandlerFunc, operationID string) strictecho.StrictEchoHandlerFunc { return func(c echo.Context, request interface{}) (response interface{}, err error) { ctx := c.Request().Context() + txn := m.NewRelicApp.StartTransaction("AuthorizationManager.assertNodeBanned") + defer txn.End() nodeID := extractor(request) err = m.Assertor.AssertNodeBanned(ctx, m.EntClient, nodeID) switch { @@ -136,6 +145,8 @@ func (m *AuthorizationManager) assertPublisherBanned(extractor func(req interfac return func(f strictecho.StrictEchoHandlerFunc, operationID string) strictecho.StrictEchoHandlerFunc { return func(c echo.Context, request interface{}) (response interface{}, err error) { ctx := c.Request().Context() + txn := m.NewRelicApp.StartTransaction("AuthorizationManager.assertPublisherBanned") + defer txn.End() publisherID := extractor(request) switch err = m.Assertor.AssertPublisherBanned(ctx, m.EntClient, publisherID); { @@ -161,6 +172,8 @@ func (m *AuthorizationManager) assertPersonalAccessTokenValid( return func(f strictecho.StrictEchoHandlerFunc, operationID string) strictecho.StrictEchoHandlerFunc { return func(c echo.Context, request interface{}) (response interface{}, err error) { ctx := c.Request().Context() + txn := m.NewRelicApp.StartTransaction("AuthorizationManager.assertPersonalAccessTokenValid") + defer txn.End() pubID := extractorPublsherID(request) pat := extractorPAT(request) tokenValid, err := m.Assertor.IsPersonalAccessTokenValidForPublisher( @@ -187,6 +200,8 @@ func (m *AuthorizationManager) assertNodeBelongsToPublisher( return func(f strictecho.StrictEchoHandlerFunc, operationID string) strictecho.StrictEchoHandlerFunc { return func(c echo.Context, request interface{}) (response interface{}, err error) { ctx := c.Request().Context() + txn := m.NewRelicApp.StartTransaction("AuthorizationManager.assertNodeBelongsToPublisher") + defer txn.End() pubID := extractorPublsherID(request) nodeID := extractorNodeID(request) diff --git a/server/middleware/authorization/authorization_middleware_test.go b/server/middleware/authorization/authorization_middleware_test.go index 40f055c..e9c2bbf 100644 --- a/server/middleware/authorization/authorization_middleware_test.go +++ b/server/middleware/authorization/authorization_middleware_test.go @@ -12,6 +12,7 @@ import ( "github.com/google/uuid" "github.com/labstack/echo/v4" + "github.com/newrelic/go-agent/v3/newrelic" "github.com/stretchr/testify/assert" ) @@ -51,7 +52,7 @@ func (m mockAlwayErrorAssertor) IsPersonalAccessTokenValidForPublisher(ctx conte } func TestNoAuthz(t *testing.T) { - mw := NewAuthorizationManager(&ent.Client{}, mockAlwayErrorAssertor{}).AuthorizationMiddleware() + mw := NewAuthorizationManager(&ent.Client{}, mockAlwayErrorAssertor{}, &newrelic.Application{}).AuthorizationMiddleware() req, res := httptest.NewRequest(http.MethodGet, "/", nil), httptest.NewRecorder() ctx := echo.New().NewContext(req, res) diff --git a/server/server.go b/server/server.go index 7a78574..e1a09a9 100644 --- a/server/server.go +++ b/server/server.go @@ -138,7 +138,7 @@ func (s *Server) Start() error { s.Dependencies.DiscordService, s.Dependencies.AlgoliaService, s.NewRelicApp) // Define middleware for authorization - authorizationManager := drip_authorization.NewAuthorizationManager(s.Client, impl.RegistryService) + authorizationManager := drip_authorization.NewAuthorizationManager(s.Client, impl.RegistryService, s.NewRelicApp) middlewares := []generated.StrictMiddlewareFunc{ authorizationManager.AuthorizationMiddleware(), }