chore: refactored saved query id parameter to always expect int64 (#792)

SpecterOps · Aug 14, 2024 · 17dd2e2 · 17dd2e2
1 parent a2a391f
commit 17dd2e2
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 13 deletions.
diff --git a/cmd/api/src/api/v2/saved_queries.go b/cmd/api/src/api/v2/saved_queries.go
@@ -178,7 +178,7 @@ func (s Resources) UpdateSavedQuery(response http.ResponseWriter, request *http.
 	} else if err := api.ReadJSONRequestPayloadLimited(&updateRequest, request); err != nil {
 		api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, err.Error(), request), response)
 		return
-	} else if savedQueryID, err := strconv.Atoi(rawSavedQueryID); err != nil {
+	} else if savedQueryID, err := strconv.ParseInt(rawSavedQueryID, 10, 64); err != nil {
 		api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, api.ErrorResponseDetailsIDMalformed, request), response)
 		return
 	} else if savedQuery, err = s.DB.GetSavedQuery(request.Context(), savedQueryID); err != nil {
@@ -223,7 +223,7 @@ func (s Resources) DeleteSavedQuery(response http.ResponseWriter, request *http.
 
 	if user, isUser := auth.GetUserFromAuthCtx(ctx2.FromRequest(request).AuthCtx); !isUser {
 		api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, "No associated user found", request), response)
-	} else if savedQueryID, err := strconv.Atoi(rawSavedQueryID); err != nil {
+	} else if savedQueryID, err := strconv.ParseInt(rawSavedQueryID, 10, 64); err != nil {
 		api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, api.ErrorResponseDetailsIDMalformed, request), response)
 	} else if savedQueryBelongsToUser, err := s.DB.SavedQueryBelongsToUser(request.Context(), user.ID, savedQueryID); errors.Is(err, database.ErrNotFound) {
 		api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusNotFound, "query does not exist", request), response)

diff --git a/cmd/api/src/database/mocks/db.go b/cmd/api/src/database/mocks/db.go
diff --git a/cmd/api/src/database/saved_queries.go b/cmd/api/src/database/saved_queries.go
@@ -18,26 +18,27 @@ package database
 
 import (
 	"context"
+
 	"github.com/gofrs/uuid"
 	"github.com/specterops/bloodhound/src/model"
 	"gorm.io/gorm"
 )
 
 type SavedQueriesData interface {
-	GetSavedQuery(ctx context.Context, queryID int) (model.SavedQuery, error)
+	GetSavedQuery(ctx context.Context, savedQueryID int64) (model.SavedQuery, error)
 	ListSavedQueries(ctx context.Context, userID uuid.UUID, order string, filter model.SQLFilter, skip, limit int) (model.SavedQueries, int, error)
 	CreateSavedQuery(ctx context.Context, userID uuid.UUID, name string, query string, description string) (model.SavedQuery, error)
 	UpdateSavedQuery(ctx context.Context, savedQuery model.SavedQuery) (model.SavedQuery, error)
-	DeleteSavedQuery(ctx context.Context, id int) error
-	SavedQueryBelongsToUser(ctx context.Context, userID uuid.UUID, savedQueryID int) (bool, error)
+	DeleteSavedQuery(ctx context.Context, savedQueryID int64) error
+	SavedQueryBelongsToUser(ctx context.Context, userID uuid.UUID, savedQueryID int64) (bool, error)
 	GetSharedSavedQueries(ctx context.Context, userID uuid.UUID) (model.SavedQueries, error)
 	GetPublicSavedQueries(ctx context.Context) (model.SavedQueries, error)
 	IsSavedQueryPublic(ctx context.Context, savedQueryID int64) (bool, error)
 }
 
-func (s *BloodhoundDB) GetSavedQuery(ctx context.Context, queryID int) (model.SavedQuery, error) {
+func (s *BloodhoundDB) GetSavedQuery(ctx context.Context, savedQueryID int64) (model.SavedQuery, error) {
 	savedQuery := model.SavedQuery{}
-	result := s.db.WithContext(ctx).First(&savedQuery, queryID)
+	result := s.db.WithContext(ctx).First(&savedQuery, savedQueryID)
 	return savedQuery, CheckError(result)
 }
 
@@ -84,11 +85,11 @@ func (s *BloodhoundDB) UpdateSavedQuery(ctx context.Context, savedQuery model.Sa
 	return savedQuery, CheckError(s.db.WithContext(ctx).Save(&savedQuery))
 }
 
-func (s *BloodhoundDB) DeleteSavedQuery(ctx context.Context, id int) error {
-	return CheckError(s.db.WithContext(ctx).Delete(&model.SavedQuery{}, id))
+func (s *BloodhoundDB) DeleteSavedQuery(ctx context.Context, savedQueryID int64) error {
+	return CheckError(s.db.WithContext(ctx).Delete(&model.SavedQuery{}, savedQueryID))
 }
 
-func (s *BloodhoundDB) SavedQueryBelongsToUser(ctx context.Context, userID uuid.UUID, savedQueryID int) (bool, error) {
+func (s *BloodhoundDB) SavedQueryBelongsToUser(ctx context.Context, userID uuid.UUID, savedQueryID int64) (bool, error) {
 	var savedQuery model.SavedQuery
 	if result := s.db.WithContext(ctx).First(&savedQuery, savedQueryID); result.Error != nil {
 		return false, CheckError(result)