refactor: make rbac package generic and reusable

api/v1/server/authz/middleware.go

@@ -8,7 +8,7 @@ import (
 	"github.com/rs/zerolog"
 
 	"github.com/hatchet-dev/hatchet/api/v1/server/middleware"
-	"github.com/hatchet-dev/hatchet/api/v1/server/rbac"
+	"github.com/hatchet-dev/hatchet/pkg/auth/rbac"
 	"github.com/hatchet-dev/hatchet/pkg/config/server"
 	"github.com/hatchet-dev/hatchet/pkg/repository/sqlcv1"
 )
@@ -20,7 +20,7 @@ type AuthZ struct {
 }
 
 func NewAuthZ(config *server.ServerConfig) (*AuthZ, error) {
-	rbacAuthorizer, err := rbac.NewAuthorizer()
+	rbacAuthorizer, err := newHatchetAuthorizer()
 	if err != nil {
 		return nil, err
 	}
@@ -197,7 +197,7 @@ func (a *AuthZ) authorizeTenantOperations(tenantMemberRole sqlcv1.TenantMemberRo
 	}
 
 	// at the moment, tenant members are only restricted from creating other tenant users.
-	if !a.rbac.IsAuthorized(tenantMemberRole, r.OperationID) {
+	if !a.rbac.IsAuthorized(string(tenantMemberRole), r.OperationID) {
 		return echo.NewHTTPError(http.StatusUnauthorized, "Not authorized to perform this operation")
 	}
 

api/v1/server/authz/rbac.go

@@ -0,0 +1,25 @@
+package authz
+
+import (
+	_ "embed"
+
+	"github.com/hatchet-dev/hatchet/api/v1/server/oas/gen"
+
+	"github.com/hatchet-dev/hatchet/pkg/auth/rbac"
+)
+
+//go:embed rbac.yaml
+var yamlFile []byte
+
+func newHatchetAuthorizer() (*rbac.Authorizer, error) {
+	permMap, err := rbac.LoadPermissionMap(yamlFile)
+	if err != nil {
+		return nil, err
+	}
+	spec, err := gen.GetSwagger()
+	if err != nil {
+		return nil, err
+	}
+
+	return rbac.NewAuthorizer(permMap, spec)
+}

api/v1/server/rbac/rbac_test.go → api/v1/server/authz/rbac_test.go

@@ -1,9 +1,10 @@
-package rbac
+package authz
 
 import (
 	"testing"
 
 	"github.com/hatchet-dev/hatchet/api/v1/server/oas/gen"
+	"github.com/hatchet-dev/hatchet/pkg/auth/rbac"
 	"github.com/hatchet-dev/hatchet/pkg/repository/sqlcv1"
 
 	"github.com/stretchr/testify/assert"
@@ -34,21 +35,21 @@ func operationIdsFromSpec() []string {
 }
 
 func TestAuthorizeTenantOperations(t *testing.T) {
-	r, err := NewAuthorizer()
+	r, err := newHatchetAuthorizer()
 	assert.Nil(t, err)
 	allOperations := operationIdsFromSpec()
 	for _, operationId := range allOperations {
-		assert.Equal(t, r.IsAuthorized(sqlcv1.TenantMemberRoleADMIN, operationId), true)
-		assert.Equal(t, r.IsAuthorized(sqlcv1.TenantMemberRoleOWNER, operationId), true)
-		if OperationIn(operationId, adminAndOwnerOnly) {
-			assert.Equal(t, r.IsAuthorized(sqlcv1.TenantMemberRoleMEMBER, operationId), false)
+		assert.Equal(t, r.IsAuthorized(string(sqlcv1.TenantMemberRoleADMIN), operationId), true)
+		assert.Equal(t, r.IsAuthorized(string(sqlcv1.TenantMemberRoleOWNER), operationId), true)
+		if rbac.OperationIn(operationId, adminAndOwnerOnly) {
+			assert.Equal(t, r.IsAuthorized(string(sqlcv1.TenantMemberRoleMEMBER), operationId), false)
 		} else {
-			assert.Equal(t, r.IsAuthorized(sqlcv1.TenantMemberRoleMEMBER, operationId), true)
+			assert.Equal(t, r.IsAuthorized(string(sqlcv1.TenantMemberRoleMEMBER), operationId), true)
 		}
 	}
 }
 
 func TestValidateSpec(t *testing.T) {
-	_, err := NewAuthorizer()
+	_, err := newHatchetAuthorizer()
 	assert.Nil(t, err)
 }

internal/msgqueue/rabbitmq/rabbitmq.go

@@ -754,13 +754,10 @@ func (t *MessageQueueImpl) subscribe(
 			return err
 		}
 
-		wg.Add(1) // we add an extra delta for the deliveries channel to be closed
-		defer wg.Done()
-
 		for rabbitMsg := range deliveries {
 			wg.Add(1)
 
-			go func(rabbitMsg amqp.Delivery) {
+			go func(rabbitMsg amqp.Delivery, session int) {
 				defer wg.Done()
 
 				msg := &msgqueue.Message{}
@@ -831,7 +828,7 @@ func (t *MessageQueueImpl) subscribe(
 					}
 				}
 
-				t.l.Debug().Msgf("(session: %d) got msg", sessionCount)
+				t.l.Debug().Msgf("(session: %d) got msg", session)
 
 				if err := preAck(msg); err != nil {
 					if isPermanentPreAckError(err) {
@@ -868,15 +865,17 @@ func (t *MessageQueueImpl) subscribe(
 					t.l.Error().Msgf("error in post-ack: %v", err)
 					return
 				}
-			}(rabbitMsg)
+			}(rabbitMsg, sessionCount)
 		}
 
 		t.l.Info().Msg("deliveries channel closed")
 
 		return nil
 	}
 
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		retryCount := 0
 		lastRetry := time.Now()
 

internal/services/partition/partition.go

@@ -284,6 +284,8 @@ func (p *Partition) StartSchedulerPartition(ctx context.Context) (func() error,
 
 	p.schedulerCron.Start()
 
+	rebalanceInactiveSchedulerPartitions(ctx, p.l, p.repo) // nolint: errcheck
+
 	return cleanup, nil
 }
 

api/v1/server/rbac/rbac.go → pkg/auth/rbac/rbac.go

@@ -1,17 +1,12 @@
 package rbac
 
 import (
-	_ "embed"
 	"fmt"
 	"maps"
 	"strings"
 
 	"github.com/getkin/kin-openapi/openapi3"
 	"go.yaml.in/yaml/v3"
-
-	"github.com/hatchet-dev/hatchet/api/v1/server/oas/gen"
-
-	"github.com/hatchet-dev/hatchet/pkg/repository/sqlcv1"
 )
 
 func OperationIn(operationId string, operationIds []string) bool {
@@ -28,16 +23,11 @@ type Authorizer struct {
 	permissionMap PermissionMap
 }
 
-func NewAuthorizer() (*Authorizer, error) {
-	permMap, err := LoadYaml()
-	if err != nil {
-		return nil, err
-	}
-	spec, err := gen.GetSwagger()
-	if err != nil {
-		return nil, err
-	}
-	err = permMap.ValidateSpec(*spec)
+func NewAuthorizer(
+	permMap *PermissionMap,
+	spec *openapi3.T,
+) (*Authorizer, error) {
+	err := permMap.ValidateSpec(*spec)
 	if err != nil {
 		return nil, err
 	}
@@ -46,8 +36,8 @@ func NewAuthorizer() (*Authorizer, error) {
 	}, nil
 }
 
-func (a *Authorizer) IsAuthorized(role sqlcv1.TenantMemberRole, operation string) bool {
-	return a.permissionMap.HasPermission(string(role), operation)
+func (a *Authorizer) IsAuthorized(role string, operation string) bool {
+	return a.permissionMap.HasPermission(role, operation)
 }
 
 type Role struct {
@@ -146,12 +136,9 @@ func (p *PermissionMap) ValidateSpec(spec openapi3.T) error {
 	return nil
 }
 
-//go:embed rbac.yaml
-var yamlFile []byte
-
-func LoadYaml() (*PermissionMap, error) {
+func LoadPermissionMap(yamlBytes []byte) (*PermissionMap, error) {
 	var yamlContents PermissionMap
-	err := yaml.Unmarshal(yamlFile, &yamlContents)
+	err := yaml.Unmarshal(yamlBytes, &yamlContents)
 	if err != nil {
 		return nil, err
 	}

pkg/config/loader/loader.go

@@ -512,7 +512,7 @@ func createControllerLayer(dc *database.Layer, cf *server.ServerConfigFile, vers
 			Version:  version,
 		}, dc.V1.SecurityCheck())
 
-		defer securityCheck.Check()
+		go securityCheck.Check()
 	}
 
 	var analyticsEmitter analytics.Analytics

pkg/security/security.go

@@ -1,9 +1,11 @@
 package security
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"net/http"
+	"time"
 
 	v1 "github.com/hatchet-dev/hatchet/pkg/repository"
 
@@ -51,8 +53,17 @@ func (a DefaultSecurityCheck) Check() {
 		return
 	}
 
-	req := fmt.Sprintf("%s/check?version=%s&tag=%s", a.Endpoint, a.Version, ident)
-	resp, err := http.Get(req) // #nosec
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	reqURL := fmt.Sprintf("%s/check?version=%s&tag=%s", a.Endpoint, a.Version, ident)
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil)
+	if err != nil {
+		a.Logger.Debug().Msgf("Error creating security check request: %s", err)
+		return
+	}
+
+	resp, err := http.DefaultClient.Do(req) // #nosec
 	if err != nil {
 		a.Logger.Debug().Msgf("Error making request to security endpoint: %s", err)
 		return

pkg/testing/harness/engine.go

@@ -147,6 +147,7 @@ func startEngine() func() {
 	os.Setenv("SERVER_LOGGER_FORMAT", "console")
 	os.Setenv("DATABASE_LOGGER_LEVEL", "error")
 	os.Setenv("DATABASE_LOGGER_FORMAT", "console")
+	os.Setenv("SERVER_SECURITY_CHECK_ENABLED", "false")
 	os.Setenv("SERVER_ADDITIONAL_LOGGERS_QUEUE_LEVEL", "error")
 	os.Setenv("SERVER_ADDITIONAL_LOGGERS_QUEUE_FORMAT", "console")
 	os.Setenv("SERVER_ADDITIONAL_LOGGERS_PGXSTATS_LEVEL", "error")