Fix UserTask Status not being updated

The Status field for UserTasks was not being correctly updated when the Spec.State was not changed.
gravitational · Jan 9, 2025 · c3ae45a · c3ae45a
1 parent 5eee08d
commit c3ae45a
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 10 deletions.
diff --git a/lib/auth/usertasks/usertasksv1/service.go b/lib/auth/usertasks/usertasksv1/service.go
@@ -131,7 +131,7 @@ func (s *Service) CreateUserTask(ctx context.Context, req *usertasksv1.CreateUse
 		return nil, trace.Wrap(err)
 	}
 
-	s.updateStatus(req.UserTask)
+	s.updateStatus(req.UserTask, nil /* existing user task */)
 
 	rsp, err := s.backend.CreateUserTask(ctx, req.UserTask)
 	s.emitCreateAuditEvent(ctx, rsp, authCtx, err)
@@ -264,10 +264,7 @@ func (s *Service) UpdateUserTask(ctx context.Context, req *usertasksv1.UpdateUse
 	}
 
 	stateChanged := existingUserTask.GetSpec().GetState() != req.GetUserTask().GetSpec().GetState()
-
-	if stateChanged {
-		s.updateStatus(req.UserTask)
-	}
+	s.updateStatus(req.UserTask, existingUserTask)
 
 	rsp, err := s.backend.UpdateUserTask(ctx, req.UserTask)
 	s.emitUpdateAuditEvent(ctx, existingUserTask, req.GetUserTask(), authCtx, err)
@@ -333,9 +330,7 @@ func (s *Service) UpsertUserTask(ctx context.Context, req *usertasksv1.UpsertUse
 		stateChanged = existingUserTask.GetSpec().GetState() != req.GetUserTask().GetSpec().GetState()
 	}
 
-	if stateChanged {
-		s.updateStatus(req.UserTask)
-	}
+	s.updateStatus(req.UserTask, existingUserTask)
 
 	rsp, err := s.backend.UpsertUserTask(ctx, req.UserTask)
 	s.emitUpsertAuditEvent(ctx, existingUserTask, req.GetUserTask(), authCtx, err)
@@ -350,10 +345,21 @@ func (s *Service) UpsertUserTask(ctx context.Context, req *usertasksv1.UpsertUse
 	return rsp, nil
 }
 
-func (s *Service) updateStatus(ut *usertasksv1.UserTask) {
+func (s *Service) updateStatus(ut *usertasksv1.UserTask, existing *usertasksv1.UserTask) {
+	// Default status for UserTask.
 	ut.Status = &usertasksv1.UserTaskStatus{
 		LastStateChange: timestamppb.New(s.clock.Now()),
 	}
+
+	if existing != nil {
+		// Inherit everything from existing UserTask.
+		ut.Status = existing.Status
+
+		// Update specific values.
+		if existing.GetSpec().GetState() != ut.GetSpec().GetState() {
+			ut.Status.LastStateChange = timestamppb.New(s.clock.Now())
+		}
+	}
 }
 
 func (s *Service) emitUpsertAuditEvent(ctx context.Context, old, new *usertasksv1.UserTask, authCtx *authz.Context, err error) {

diff --git a/lib/auth/usertasks/usertasksv1/service_test.go b/lib/auth/usertasks/usertasksv1/service_test.go
@@ -153,6 +153,7 @@ func TestEvents(t *testing.T) {
 	// LastStateChange is updated.
 	require.Equal(t, timestamppb.New(fakeClock.Now()), createUserTaskResp.Status.LastStateChange)
 
+	expectedLastStateChange := createUserTaskResp.Status.LastStateChange
 	ut1.Spec.DiscoverEc2.Instances["i-345"] = &usertasksv1.DiscoverEC2Instance{
 		InstanceId:      "i-345",
 		DiscoveryConfig: "dc01",
@@ -165,7 +166,7 @@ func TestEvents(t *testing.T) {
 	require.Len(t, testReporter.emittedEvents, 1)
 	consumeAssertEvent(t, auditEventsSink.C(), auditEventFor(userTaskName, "update", "OPEN", "OPEN"))
 	// LastStateChange is not updated.
-	require.Equal(t, createUserTaskResp.Status.LastStateChange, upsertUserTaskResp.Status.LastStateChange)
+	require.Equal(t, expectedLastStateChange.AsTime(), upsertUserTaskResp.Status.LastStateChange.AsTime())
 
 	ut1.Spec.State = "RESOLVED"
 	fakeClock.Advance(1 * time.Minute)
@@ -177,6 +178,36 @@ func TestEvents(t *testing.T) {
 	// LastStateChange was updated because the state changed.
 	require.Equal(t, timestamppb.New(fakeClock.Now()), updateUserTaskResp.Status.LastStateChange)
 
+	// Updating one of the instances.
+	expectedLastStateChange = updateUserTaskResp.Status.GetLastStateChange()
+	fakeClock.Advance(1 * time.Minute)
+	ut1.Spec.DiscoverEc2.Instances["i-345"] = &usertasksv1.DiscoverEC2Instance{
+		InstanceId:      "i-345",
+		DiscoveryConfig: "dc01",
+		DiscoveryGroup:  "dg01",
+		SyncTime:        timestamppb.New(fakeClock.Now()),
+	}
+	updateUserTaskResp, err = service.UpdateUserTask(ctx, &usertasksv1.UpdateUserTaskRequest{UserTask: ut1})
+	require.NoError(t, err)
+	// Does not change the LastStateChange
+	require.Equal(t, expectedLastStateChange.AsTime(), updateUserTaskResp.Status.LastStateChange.AsTime())
+	consumeAssertEvent(t, auditEventsSink.C(), auditEventFor(userTaskName, "update", "RESOLVED", "RESOLVED"))
+
+	// Upserting one of the instances.
+	expectedLastStateChange = updateUserTaskResp.Status.GetLastStateChange()
+	fakeClock.Advance(1 * time.Minute)
+	ut1.Spec.DiscoverEc2.Instances["i-345"] = &usertasksv1.DiscoverEC2Instance{
+		InstanceId:      "i-345",
+		DiscoveryConfig: "dc01",
+		DiscoveryGroup:  "dg01",
+		SyncTime:        timestamppb.New(fakeClock.Now()),
+	}
+	upsertUserTaskResp, err = service.UpsertUserTask(ctx, &usertasksv1.UpsertUserTaskRequest{UserTask: ut1})
+	require.NoError(t, err)
+	// Does not change the LastStateChange
+	require.Equal(t, expectedLastStateChange.AsTime(), upsertUserTaskResp.Status.LastStateChange.AsTime())
+	consumeAssertEvent(t, auditEventsSink.C(), auditEventFor(userTaskName, "update", "RESOLVED", "RESOLVED"))
+
 	_, err = service.DeleteUserTask(ctx, &usertasksv1.DeleteUserTaskRequest{Name: userTaskName})
 	require.NoError(t, err)
 	// No usage report for deleted resources.