Fix missing risk challenge response

server/db/mysql/tenant_adapter.go

@@ -5028,6 +5028,9 @@ func (a *adapter) riskChallengeByIDTx(ctx context.Context, q sqlx.QueryerContext
 		query += " FOR UPDATE"
 	}
 	if err := sqlx.GetContext(ctx, q, &row, query, challengeID); err != nil {
+		if err == sql.ErrNoRows {
+			return nil, t.ErrNotFound
+		}
 		return nil, err
 	}
 	challenge := &t.RiskChallenge{

server/im_http.go

@@ -53,6 +53,7 @@ const (
 	imReasonUnsupportedMedia    = "UNSUPPORTED_MEDIA_TYPE"
 	imReasonFileUnavailable     = "FILE_UNAVAILABLE"
 	imReasonRateLimited         = "RATE_LIMITED"
+	imReasonResourceNotFound    = "RESOURCE_NOT_FOUND"
 )
 
 const (
@@ -1257,6 +1258,10 @@ func serveIMRiskChallengeConfirm(wrt http.ResponseWriter, req *http.Request) {
 	}
 	challenge, err := store.Safety.ConfirmRiskChallenge(principal.TenantID, principal.Uid, challengeID, decision, tokenID, time.Now().UTC().Round(time.Millisecond))
 	if err != nil {
+		if errors.Is(err, types.ErrNotFound) {
+			writeIMError(wrt, req, http.StatusNotFound, 30004, imReasonResourceNotFound, "risk challenge not found", nil)
+			return
+		}
 		writeIMRelationError(wrt, req, err)
 		return
 	}

server/im_http_test.go

@@ -307,6 +307,7 @@ type testDevicePushTokenStore struct {
 
 type testSafetyStore struct {
 	riskChallenge *types.RiskChallenge
+	riskErr       error
 	reasons       []types.AbuseReason
 	abuseReport   *types.AbuseReport
 	feedback      *types.FeedbackSubmission
@@ -856,6 +857,9 @@ func (s *testSafetyStore) CreateRiskChallenge(tenantID types.TenantID, challenge
 }
 
 func (s *testSafetyStore) ConfirmRiskChallenge(tenantID types.TenantID, uid types.Uid, challengeID, decision, actionTokenID string, now time.Time) (*types.RiskChallenge, error) {
+	if s.riskErr != nil {
+		return nil, s.riskErr
+	}
 	s.tenantID = tenantID
 	s.uid = uid
 	s.decision = decision
@@ -4261,6 +4265,28 @@ func TestServeIMRiskChallengeConfirm(t *testing.T) {
 	}
 }
 
+func TestServeIMRiskChallengeConfirmNotFound(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+	expectIMBearer(t, ctrl, "token-1", 7, 42)
+	oldSafety := store.Safety
+	store.Safety = &testSafetyStore{riskErr: types.ErrNotFound}
+	defer func() { store.Safety = oldSafety }()
+
+	req := httptest.NewRequest(http.MethodPost, "/v1/im/security/risk-challenges/risk_fake/confirm", bytes.NewBufferString(`{"decision":"cancel"}`))
+	req.Header.Set("Authorization", "Bearer token-1")
+	rec := httptest.NewRecorder()
+	serveIMRiskChallengeConfirm(rec, req)
+
+	if rec.Code != http.StatusNotFound {
+		t.Fatalf("status: got %d want %d body=%s", rec.Code, http.StatusNotFound, rec.Body.String())
+	}
+	resp := decodeIMHTTPResponse(t, rec)
+	if resp.Reason != imReasonResourceNotFound {
+		t.Fatalf("reason: got %q want %q", resp.Reason, imReasonResourceNotFound)
+	}
+}
+
 func TestServeIMAbuseReportsAndReasons(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()