radosgw-admin: can set user's bucket quota

Signed-off-by: Yehuda Sadeh <yehuda@inktank.com>

3 files changed, 83 insertions, 33 deletions

diff --git a/src/rgw/rgw_admin.cc b/src/rgw/rgw_admin.cc
index 6c02c57262e..5058d11de74 100644
--- a/src/rgw/rgw_admin.cc
+++ b/src/rgw/rgw_admin.cc
@@ -678,6 +678,64 @@ static bool dump_string(const char *field_name, bufferlist& bl, Formatter *f)
   return true;
 }
 
+void set_quota_info(RGWQuotaInfo& quota, int opt_cmd, int64_t max_size, int64_t max_objects)
+{
+  switch (opt_cmd) {
+    case OPT_QUOTA_ENABLE:
+      quota.enabled = true;
+
+      // falling through on purpose
+
+    case OPT_QUOTA_SET:
+      if (max_objects >= 0) {
+        quota.max_objects = max_objects;
+      }
+      if (max_size >= 0) {
+        quota.max_size_kb = rgw_rounded_kb(max_size);
+      }
+      break;
+    case OPT_QUOTA_DISABLE:
+      quota.enabled = false;
+      break;
+  }
+}
+
+int set_bucket_quota(RGWRados *store, int opt_cmd, string& bucket_name, int64_t max_size, int64_t max_objects)
+{
+  RGWBucketInfo bucket_info;
+  map<string, bufferlist> attrs;
+  int r = store->get_bucket_info(NULL, bucket_name, bucket_info, NULL, &attrs);
+  if (r < 0) {
+    cerr << "could not get bucket info for bucket=" << bucket_name << ": " << cpp_strerror(-r) << std::endl;
+    return -r;
+  }
+
+  set_quota_info(bucket_info.quota, opt_cmd, max_size, max_objects);
+
+   r = store->put_bucket_instance_info(bucket_info, false, 0, &attrs);
+  if (r < 0) {
+    cerr << "ERROR: failed writing bucket instance info: " << cpp_strerror(-r) << std::endl;
+    return -r;
+  }
+  return 0;
+}
+
+int set_user_bucket_quota(int opt_cmd, RGWUser& user, RGWUserAdminOpState& op_state, int64_t max_size, int64_t max_objects)
+{
+  RGWUserInfo& user_info = op_state.get_user_info();
+
+  set_quota_info(user_info.bucket_quota, opt_cmd, max_size, max_objects);
+
+  op_state.set_bucket_quota(user_info.bucket_quota);
+
+  string err;
+  int r = user.modify(op_state, &err);
+  if (r < 0) {
+    cerr << "ERROR: failed updating user info: " << cpp_strerror(-r) << ": " << err << std::endl;
+    return -r;
+  }
+  return 0;
+}
 
 int main(int argc, char **argv) 
 {
@@ -2257,42 +2315,15 @@ next:
   bool quota_op = (opt_cmd == OPT_QUOTA_SET || opt_cmd == OPT_QUOTA_ENABLE || opt_cmd == OPT_QUOTA_DISABLE);
 
   if (quota_op) {
-    if (bucket_name.empty()) {
-      cerr << "ERROR: bucket name was not specified" << std::endl;
+    if (bucket_name.empty() && user_id.empty()) {
+      cerr << "ERROR: bucket name or uid is required for quota operation" << std::endl;
       return EINVAL;
     }
 
-    RGWBucketInfo bucket_info;
-    map<string, bufferlist> attrs;
-    int r = store->get_bucket_info(NULL, bucket_name, bucket_info, NULL, &attrs);
-    if (r < 0) {
-      cerr << "could not get bucket info for bucket=" << bucket_name << ": " << cpp_strerror(-r) << std::endl;
-      return -r;
-    }
-
-    switch (opt_cmd) {
-      case OPT_QUOTA_ENABLE:
-        bucket_info.quota.enabled = true;
-
-        // falling through on purpose
-
-      case OPT_QUOTA_SET:
-        if (max_objects >= 0) {
-          bucket_info.quota.max_objects = max_objects;
-        }
-        if (max_size >= 0) {
-          bucket_info.quota.max_size_kb = rgw_rounded_kb(max_size);
-        }
-        break;
-      case OPT_QUOTA_DISABLE:
-        bucket_info.quota.enabled = false;
-        break;
-    }
-
-    r = store->put_bucket_instance_info(bucket_info, false, 0, &attrs);
-    if (r < 0) {
-      cerr << "ERROR: failed writing bucket instance info: " << cpp_strerror(-r) << std::endl;
-      return -r;
+    if (!bucket_name.empty()) {
+      set_bucket_quota(store, opt_cmd, bucket_name, max_size, max_objects);
+    } else if (!user_id.empty()) {
+      set_user_bucket_quota(opt_cmd, user, user_op, max_size, max_objects);
     }
   }
   return 0;
diff --git a/src/rgw/rgw_user.cc b/src/rgw/rgw_user.cc
index 5e5b5c564bb..dc529e3d48d 100644
--- a/src/rgw/rgw_user.cc
+++ b/src/rgw/rgw_user.cc
@@ -1682,6 +1682,9 @@ int RGWUser::execute_add(RGWUserAdminOpState& op_state, std::string *err_msg)
   if (op_state.op_mask_specified)
     user_info.op_mask = op_state.get_op_mask();
 
+  if (op_state.has_bucket_quota())
+    user_info.bucket_quota = op_state.get_bucket_quota();
+
   // update the request
   op_state.set_user_info(user_info);
   op_state.set_populated();
@@ -1884,6 +1887,9 @@ int RGWUser::execute_modify(RGWUserAdminOpState& op_state, std::string *err_msg)
   if (op_state.op_mask_specified)
     user_info.op_mask = op_state.get_op_mask();
 
+  if (op_state.has_bucket_quota())
+    user_info.bucket_quota = op_state.get_bucket_quota();
+
   if (op_state.has_suspension_op()) {
     __u8 suspended = op_state.get_suspension_status();
     user_info.suspended = suspended;
diff --git a/src/rgw/rgw_user.h b/src/rgw/rgw_user.h
index 32bcf199001..e71b8f81778 100644
--- a/src/rgw/rgw_user.h
+++ b/src/rgw/rgw_user.h
@@ -172,6 +172,10 @@ struct RGWUserAdminOpState {
   bool subuser_params_checked;
   bool user_params_checked;
 
+  bool bucket_quota_specified;
+
+  RGWQuotaInfo bucket_quota;
+
   void set_access_key(std::string& access_key) {
     if (access_key.empty())
       return;
@@ -285,6 +289,12 @@ struct RGWUserAdminOpState {
     key_op = true;
   }
 
+  void set_bucket_quota(RGWQuotaInfo& quota)
+  {
+    bucket_quota = quota;
+    bucket_quota_specified = true;
+  }
+
   bool is_populated() { return populated; };
   bool is_initialized() { return initialized; };
   bool has_existing_user() { return existing_user; };
@@ -303,6 +313,7 @@ struct RGWUserAdminOpState {
   bool will_purge_keys() { return purge_keys; };
   bool will_purge_data() { return purge_data; };
   bool will_generate_subuser() { return gen_subuser; };
+  bool has_bucket_quota() { return bucket_quota_specified; }
   void set_populated() { populated = true; };
   void clear_populated() { populated = false; };
   void set_initialized() { initialized = true; };
@@ -317,6 +328,7 @@ struct RGWUserAdminOpState {
   uint32_t get_subuser_perm() { return perm_mask; };
   uint32_t get_max_buckets() { return max_buckets; };
   uint32_t get_op_mask() { return op_mask; };
+  RGWQuotaInfo& get_bucket_quota() { return bucket_quota; }
 
   std::string get_user_id() { return user_id; };
   std::string get_subuser() { return subuser; };
@@ -403,6 +415,7 @@ struct RGWUserAdminOpState {
     key_params_checked = false;
     subuser_params_checked = false;
     user_params_checked = false;
+    bucket_quota_specified = false;
   }
 };