user_resource_guard.go 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596
  1. package repository
  2. import (
  3. "context"
  4. dbent "github.com/Wei-Shaw/sub2api/ent"
  5. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  6. "github.com/Wei-Shaw/sub2api/internal/service"
  7. )
  8. type managedResourceGuardKey struct{}
  9. func needsManagedResourceGuard(ctx context.Context) bool {
  10. _, managed := service.UserManagementActor(ctx)
  11. return managed && ctx.Value(managedResourceGuardKey{}) != true
  12. }
  13. // The caller supplies a fixed table name, never request input. Role and owner
  14. // checks share the role-mutation lock and live until the resource write commits.
  15. func withManagedResourceWrite(ctx context.Context, client *dbent.Client, table string, resourceID, ownerID int64, work func(context.Context) error) error {
  16. actorID, managed := service.UserManagementActor(ctx)
  17. if !managed {
  18. return work(ctx)
  19. }
  20. if actorID <= 0 || client == nil {
  21. return infraerrors.Forbidden("USER_RESOURCE_FORBIDDEN", "User resource authorization unavailable")
  22. }
  23. var tx *dbent.Tx
  24. if current := dbent.TxFromContext(ctx); current != nil {
  25. client = current.Client()
  26. } else {
  27. var err error
  28. tx, err = client.Tx(ctx)
  29. if err != nil {
  30. return err
  31. }
  32. defer func() { _ = tx.Rollback() }()
  33. client = tx.Client()
  34. ctx = dbent.NewTxContext(ctx, tx)
  35. }
  36. release, err := lockUserHierarchy(ctx, client)
  37. if err != nil {
  38. return err
  39. }
  40. defer release()
  41. owners := []int64{actorID}
  42. if ownerID > 0 {
  43. owners = append(owners, ownerID)
  44. }
  45. if resourceID > 0 {
  46. if table != "api_keys" && table != "user_subscriptions" {
  47. return infraerrors.BadRequest("INVALID_RESOURCE", "Invalid resource")
  48. }
  49. rows, err := client.QueryContext(ctx, "SELECT user_id FROM "+table+" WHERE id = $1", resourceID)
  50. if err != nil {
  51. return err
  52. }
  53. if !rows.Next() {
  54. err := rows.Err()
  55. _ = rows.Close()
  56. if err != nil {
  57. return err
  58. }
  59. return infraerrors.NotFound("USER_RESOURCE_NOT_FOUND", "User resource not found")
  60. }
  61. var actualOwner int64
  62. err = rows.Scan(&actualOwner)
  63. _ = rows.Close()
  64. if err != nil {
  65. return err
  66. }
  67. owners = append(owners, actualOwner)
  68. }
  69. users, err := lockHierarchyUsers(ctx, client, owners...)
  70. if err != nil {
  71. return err
  72. }
  73. needsSuper := false
  74. for _, id := range owners[1:] {
  75. if users[id] == nil {
  76. return service.ErrUserNotFound
  77. }
  78. needsSuper = needsSuper || service.IsAdminRole(users[id].Role)
  79. }
  80. if err := verifyHierarchyActor(users[actorID], needsSuper); err != nil {
  81. return err
  82. }
  83. ctx = context.WithValue(ctx, managedResourceGuardKey{}, true)
  84. if err := work(ctx); err != nil {
  85. return err
  86. }
  87. if tx != nil {
  88. return tx.Commit()
  89. }
  90. return nil
  91. }