Merge pull request #6872 from ferrous-systems/pa-rate-limits-part1

First batch of rate limiter changes

Author: Turbo87

migrations/2021-07-18-125813_add_rate_limit_action/down.sql

@@ -0,0 +1,9 @@
+DELETE FROM publish_limit_buckets WHERE action != 0;
+ALTER TABLE publish_limit_buckets DROP CONSTRAINT publish_limit_buckets_pkey;
+ALTER TABLE publish_limit_buckets ADD CONSTRAINT publish_limit_buckets_pkey PRIMARY KEY (user_id);
+ALTER TABLE publish_limit_buckets DROP COLUMN action;
+
+DELETE FROM publish_rate_overrides WHERE action != 0;
+ALTER TABLE publish_rate_overrides DROP CONSTRAINT publish_rate_overrides_pkey;
+ALTER TABLE publish_rate_overrides ADD CONSTRAINT publish_rate_overrides_pkey PRIMARY KEY (user_id);
+ALTER TABLE publish_rate_overrides DROP COLUMN action;

migrations/2021-07-18-125813_add_rate_limit_action/up.sql

@@ -0,0 +1,7 @@
+ALTER TABLE publish_limit_buckets ADD COLUMN action INTEGER NOT NULL DEFAULT 0;
+ALTER TABLE publish_limit_buckets DROP CONSTRAINT publish_limit_buckets_pkey;
+ALTER TABLE publish_limit_buckets ADD CONSTRAINT publish_limit_buckets_pkey PRIMARY KEY (user_id, action);
+
+ALTER TABLE publish_rate_overrides ADD COLUMN action INTEGER NOT NULL DEFAULT 0;
+ALTER TABLE publish_rate_overrides DROP CONSTRAINT publish_rate_overrides_pkey;
+ALTER TABLE publish_rate_overrides ADD CONSTRAINT publish_rate_overrides_pkey PRIMARY KEY (user_id, action);

src/config/server.rs

@@ -2,7 +2,7 @@ use anyhow::{anyhow, Context};
 use ipnetwork::IpNetwork;
 use oauth2::{ClientId, ClientSecret};
 
-use crate::publish_rate_limit::PublishRateLimit;
+use crate::rate_limiter::RateLimiter;
 use crate::{env, env_optional, Env};
 
 use super::base::Base;
@@ -30,7 +30,7 @@ pub struct Server {
     pub gh_client_secret: ClientSecret,
     pub max_upload_size: u64,
     pub max_unpack_size: u64,
-    pub publish_rate_limit: PublishRateLimit,
+    pub rate_limiter: RateLimiter,
     pub new_version_rate_limit: Option<u32>,
     pub blocked_traffic: Vec<(String, Vec<String>)>,
     pub max_allowed_page_offset: u32,
@@ -153,7 +153,7 @@ impl Default for Server {
             gh_client_secret: ClientSecret::new(env("GH_CLIENT_SECRET")),
             max_upload_size: 10 * 1024 * 1024, // 10 MB default file upload size limit
             max_unpack_size: 512 * 1024 * 1024, // 512 MB max when decompressed
-            publish_rate_limit: Default::default(),
+            rate_limiter: Default::default(),
             new_version_rate_limit: env_optional("MAX_NEW_VERSIONS_DAILY"),
             blocked_traffic: blocked_traffic(),
             max_allowed_page_offset: env_optional("WEB_MAX_ALLOWED_PAGE_OFFSET").unwrap_or(200),

src/controllers/krate/publish.rs

@@ -137,8 +137,7 @@ pub async fn publish(app: AppState, req: BytesRequest) -> AppResult<Json<GoodCra
             };
 
             let license_file = new_crate.license_file.as_deref();
-            let krate =
-                persist.create_or_update(conn, user.id, Some(&app.config.publish_rate_limit))?;
+            let krate = persist.create_or_update(conn, user.id, Some(&app.config.rate_limiter))?;
 
             let owners = krate.owners(conn)?;
             if user.rights(&app, &owners)? < Rights::Publish {

src/lib.rs

@@ -46,8 +46,9 @@ pub mod github;
 pub mod headers;
 pub mod metrics;
 pub mod middleware;
-mod publish_rate_limit;
+mod rate_limiter;
 pub mod schema;
+#[macro_use]
 pub mod sql;
 pub mod ssh;
 pub mod swirl;

src/models/action.rs

@@ -1,22 +1,14 @@
-use chrono::NaiveDateTime;
-use diesel::prelude::*;
-use diesel::{
-    deserialize::{self, FromSql},
-    pg::Pg,
-    serialize::{self, Output, ToSql},
-    sql_types::Integer,
-};
-
 use crate::models::{ApiToken, User, Version};
 use crate::schema::*;
+use chrono::NaiveDateTime;
+use diesel::prelude::*;
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, FromSqlRow, AsExpression)]
-#[repr(i32)]
-#[diesel(sql_type = Integer)]
-pub enum VersionAction {
-    Publish = 0,
-    Yank = 1,
-    Unyank = 2,
+pg_enum! {
+    pub enum VersionAction {
+        Publish = 0,
+        Yank = 1,
+        Unyank = 2,
+    }
 }
 
 impl From<VersionAction> for &'static str {
@@ -37,23 +29,6 @@ impl From<VersionAction> for String {
     }
 }
 
-impl FromSql<Integer, Pg> for VersionAction {
-    fn from_sql(bytes: diesel::pg::PgValue<'_>) -> deserialize::Result<Self> {
-        match <i32 as FromSql<Integer, Pg>>::from_sql(bytes)? {
-            0 => Ok(VersionAction::Publish),
-            1 => Ok(VersionAction::Yank),
-            2 => Ok(VersionAction::Unyank),
-            n => Err(format!("unknown version action: {n}").into()),
-        }
-    }
-}
-
-impl ToSql<Integer, Pg> for VersionAction {
-    fn to_sql(&self, out: &mut Output<'_, '_, Pg>) -> serialize::Result {
-        ToSql::<Integer, Pg>::to_sql(&(*self as i32), &mut out.reborrow())
-    }
-}
-
 #[derive(Debug, Clone, Copy, Queryable, Identifiable, Associations)]
 #[diesel(
     table_name = version_owner_actions,

src/models/dependency.rs

@@ -1,5 +1,3 @@
-use diesel::deserialize::{self, FromSql};
-use diesel::pg::Pg;
 use diesel::sql_types::{Integer, Text};
 
 use crate::models::{Crate, Version};
@@ -36,14 +34,12 @@ pub struct ReverseDependency {
     pub name: String,
 }
 
-#[derive(Copy, Clone, Serialize, Deserialize, Debug, FromSqlRow)]
-#[serde(rename_all = "lowercase")]
-#[repr(u32)]
-pub enum DependencyKind {
-    Normal = 0,
-    Build = 1,
-    Dev = 2,
-    // if you add a kind here, be sure to update `from_row` below.
+pg_enum! {
+    pub enum DependencyKind {
+        Normal = 0,
+        Build = 1,
+        Dev = 2,
+    }
 }
 
 impl From<IndexDependencyKind> for DependencyKind {
@@ -65,14 +61,3 @@ impl From<DependencyKind> for IndexDependencyKind {
         }
     }
 }
-
-impl FromSql<Integer, Pg> for DependencyKind {
-    fn from_sql(bytes: diesel::pg::PgValue<'_>) -> deserialize::Result<Self> {
-        match <i32 as FromSql<Integer, Pg>>::from_sql(bytes)? {
-            0 => Ok(DependencyKind::Normal),
-            1 => Ok(DependencyKind::Build),
-            2 => Ok(DependencyKind::Dev),
-            n => Err(format!("unknown kind: {n}").into()),
-        }
-    }
-}

src/models/krate.rs

@@ -17,7 +17,7 @@ use crate::models::{
 use crate::util::errors::{cargo_err, AppResult};
 
 use crate::models::helpers::with_count::*;
-use crate::publish_rate_limit::PublishRateLimit;
+use crate::rate_limiter::RateLimiter;
 use crate::schema::*;
 use crate::sql::canon_crate_name;
 
@@ -107,7 +107,7 @@ impl<'a> NewCrate<'a> {
         self,
         conn: &mut PgConnection,
         uploader: i32,
-        rate_limit: Option<&PublishRateLimit>,
+        rate_limit: Option<&RateLimiter>,
     ) -> AppResult<Crate> {
         use diesel::update;