Fix invalid memory dereference in insertTasks when no step config exists for a task

pkg/repository/v1/task.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"hash/fnv"
 	"sort"
+	"strings"
 	"time"
 
 	"github.com/google/uuid"
@@ -1663,6 +1664,74 @@ func (r *sharedRepository) createTasks(
 	return r.insertTasks(ctx, tx, tenantId, tasks, stepIdsToConfig)
 }
 
+func partitionInsertTasksByStepConfig(
+	tasks []CreateTaskOpts,
+	stepIdsToConfig map[string]*sqlcv1.ListStepsByIdsRow,
+) ([]CreateTaskOpts, []string) {
+	if len(tasks) == 0 {
+		return nil, nil
+	}
+
+	validTasks := make([]CreateTaskOpts, 0, len(tasks))
+	missing := make(map[string]struct{})
+
+	for _, task := range tasks {
+		stepConfig, ok := stepIdsToConfig[task.StepId]
+		if !ok || stepConfig == nil {
+			missing[task.StepId] = struct{}{}
+			continue
+		}
+
+		validTasks = append(validTasks, task)
+	}
+
+	if len(missing) == 0 {
+		return validTasks, nil
+	}
+
+	missingIds := make([]string, 0, len(missing))
+
+	for id := range missing {
+		missingIds = append(missingIds, id)
+	}
+
+	return validTasks, missingIds
+}
+
+func partitionReplayTasksByStepConfig(
+	tasks []ReplayTaskOpts,
+	stepIdsToConfig map[string]*sqlcv1.ListStepsByIdsRow,
+) ([]ReplayTaskOpts, []string) {
+	if len(tasks) == 0 {
+		return nil, nil
+	}
+
+	validTasks := make([]ReplayTaskOpts, 0, len(tasks))
+	missing := make(map[string]struct{})
+
+	for _, task := range tasks {
+		stepConfig, ok := stepIdsToConfig[task.StepId]
+		if !ok || stepConfig == nil {
+			missing[task.StepId] = struct{}{}
+			continue
+		}
+
+		validTasks = append(validTasks, task)
+	}
+
+	if len(missing) == 0 {
+		return validTasks, nil
+	}
+
+	missingIds := make([]string, 0, len(missing))
+
+	for id := range missing {
+		missingIds = append(missingIds, id)
+	}
+
+	return validTasks, missingIds
+}
+
 // insertTasks inserts new tasks into the database. note that we're using Postgres rules to automatically insert the created
 // tasks into the queue_items table.
 func (r *sharedRepository) insertTasks(
@@ -1676,6 +1745,26 @@ func (r *sharedRepository) insertTasks(
 		return nil, nil
 	}
 
+	validTasks, missingStepIds := partitionInsertTasksByStepConfig(tasks, stepIdsToConfig)
+
+	if len(validTasks) == 0 {
+		if len(missingStepIds) > 0 {
+			return nil, fmt.Errorf("missing step config for steps: %s", strings.Join(missingStepIds, ", "))
+		}
+
+		return nil, nil
+	}
+
+	if len(missingStepIds) > 0 && r.l != nil {
+		r.l.Warn().
+			Str("tenant_id", tenantId).
+			Strs("missing_step_ids", missingStepIds).
+			Int("skipped_tasks", len(tasks)-len(validTasks)).
+			Msg("skipping insert tasks without step config")
+	}
+
+	tasks = validTasks
+
 	expressions, err := r.getStepExpressions(ctx, tx, stepIdsToConfig)
 
 	if err != nil {
@@ -2222,6 +2311,26 @@ func (r *sharedRepository) replayTasks(
 		stepIdsToConfig[sqlchelpers.UUIDToStr(step.ID)] = step
 	}
 
+	validTasks, missingStepIds := partitionReplayTasksByStepConfig(tasks, stepIdsToConfig)
+
+	if len(validTasks) == 0 {
+		if len(missingStepIds) > 0 {
+			return nil, fmt.Errorf("missing step config for steps: %s", strings.Join(missingStepIds, ", "))
+		}
+
+		return nil, nil
+	}
+
+	if len(missingStepIds) > 0 && r.l != nil {
+		r.l.Warn().
+			Str("tenant_id", tenantId).
+			Strs("missing_step_ids", missingStepIds).
+			Int("skipped_tasks", len(tasks)-len(validTasks)).
+			Msg("skipping replay tasks without step config")
+	}
+
+	tasks = validTasks
+
 	concurrencyStrats, err := r.getConcurrencyExpressions(ctx, tx, tenantId, stepIdsToConfig)
 
 	if err != nil {

pkg/repository/v1/task_test.go

@@ -0,0 +1,98 @@
+package v1
+
+import (
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+
+	"github.com/hatchet-dev/hatchet/pkg/repository/v1/sqlcv1"
+)
+
+func TestPartitionInsertTasksByStepConfigMissingStep(t *testing.T) {
+	stepA := uuid.NewString()
+	stepB := uuid.NewString()
+
+	tasks := []CreateTaskOpts{
+		{StepId: stepA},
+		{StepId: stepB},
+	}
+
+	stepConfig := map[string]*sqlcv1.ListStepsByIdsRow{
+		stepA: {},
+	}
+
+	valid, missing := partitionInsertTasksByStepConfig(tasks, stepConfig)
+	require.Len(t, valid, 1)
+	require.Equal(t, stepA, valid[0].StepId)
+	require.Equal(t, []string{stepB}, missing)
+}
+
+func TestPartitionInsertTasksByStepConfigAllPresent(t *testing.T) {
+	stepA := uuid.NewString()
+	stepB := uuid.NewString()
+
+	tasks := []CreateTaskOpts{
+		{StepId: stepA},
+		{StepId: stepB},
+	}
+
+	stepConfig := map[string]*sqlcv1.ListStepsByIdsRow{
+		stepA: {},
+		stepB: {},
+	}
+
+	valid, missing := partitionInsertTasksByStepConfig(tasks, stepConfig)
+	require.Len(t, valid, 2)
+	require.Empty(t, missing)
+}
+
+func TestPartitionInsertTasksByStepConfigEmptyInput(t *testing.T) {
+	valid, missing := partitionInsertTasksByStepConfig(nil, nil)
+	require.Nil(t, valid)
+	require.Nil(t, missing)
+}
+
+func TestPartitionReplayTasksByStepConfigMissingStep(t *testing.T) {
+	stepA := uuid.NewString()
+	stepB := uuid.NewString()
+
+	tasks := []ReplayTaskOpts{
+		{StepId: stepA},
+		{StepId: stepB},
+	}
+
+	stepConfig := map[string]*sqlcv1.ListStepsByIdsRow{
+		stepA: {},
+	}
+
+	valid, missing := partitionReplayTasksByStepConfig(tasks, stepConfig)
+	require.Len(t, valid, 1)
+	require.Equal(t, stepA, valid[0].StepId)
+	require.Equal(t, []string{stepB}, missing)
+}
+
+func TestPartitionReplayTasksByStepConfigAllPresent(t *testing.T) {
+	stepA := uuid.NewString()
+	stepB := uuid.NewString()
+
+	tasks := []ReplayTaskOpts{
+		{StepId: stepA},
+		{StepId: stepB},
+	}
+
+	stepConfig := map[string]*sqlcv1.ListStepsByIdsRow{
+		stepA: {},
+		stepB: {},
+	}
+
+	valid, missing := partitionReplayTasksByStepConfig(tasks, stepConfig)
+	require.Len(t, valid, 2)
+	require.Empty(t, missing)
+}
+
+func TestPartitionReplayTasksByStepConfigEmptyInput(t *testing.T) {
+	valid, missing := partitionReplayTasksByStepConfig(nil, nil)
+	require.Nil(t, valid)
+	require.Nil(t, missing)
+}