Require payment to create hurlurls

Cargo.toml

@@ -1,4 +1,3 @@
-
-
 [workspace]
 members = ["urllb", "web", "shared"]
+resolver = "2"

shared/Cargo.toml

@@ -11,5 +11,6 @@ default = []
 [dependencies]
 serde = { version = "1", features = ["serde_derive"] }
 diesel = { version = "2.0.4", optional = true, features = ["network-address", "ipnet-address"] }
-validator = { version = "0.16",  features = ["derive"] }
+diesel-derive-enum = { version = "2.1.0", features = ["postgres"] }
+validator = { version = "0.16", features = ["derive"] }
 ipnet = "2"

shared/src/lib.rs

@@ -21,6 +21,9 @@ pub struct Link {
     pub fraud_reason: Option<String>,
     #[serde(skip, default)]
     pub created_by_ip: Option<ipnet::IpNet>,
+
+    pub stripe_session_id: Option<String>,
+    pub payment_status: Option<PaymentStatus>,
 }
 
 #[cfg_attr(feature = "diesel", derive(Queryable, Identifiable))]
@@ -52,6 +55,12 @@ pub struct CreateLinkDto {
     pub targets: Vec<CreateTargetDto>,
 }
 
+#[derive(Serialize, Deserialize, Clone, Debug)]
+pub enum CreateResult {
+    Link(LinkDto),
+    StripeRedirect(String),
+}
+
 #[derive(Serialize, Deserialize, Clone, Debug)]
 pub struct LinkDto {
     #[serde(flatten)]
@@ -66,3 +75,15 @@ pub struct TotalStats {
     pub redirects: i64,
     pub targets: i64,
 }
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+#[cfg_attr(feature = "diesel", derive(diesel_derive_enum::DbEnum))]
+#[cfg_attr(
+    feature = "diesel",
+    ExistingTypePath = "crate::schema::sql_types::PaymentStatus"
+)]
+pub enum PaymentStatus {
+    Pending,
+    Succeeded,
+    Failed,
+}

shared/src/schema.rs

@@ -1,6 +1,15 @@
 // @generated automatically by Diesel CLI.
 
+pub mod sql_types {
+    #[derive(diesel::sql_types::SqlType)]
+    #[diesel(postgres_type(name = "payment_status"))]
+    pub struct PaymentStatus;
+}
+
 diesel::table! {
+    use diesel::sql_types::*;
+    use super::sql_types::PaymentStatus;
+
     links (id) {
         id -> Int4,
         url -> Varchar,
@@ -9,6 +18,8 @@ diesel::table! {
         fraud -> Bool,
         fraud_reason -> Nullable<Text>,
         created_by_ip -> Nullable<Inet>,
+        stripe_session_id -> Nullable<Text>,
+        payment_status -> Nullable<PaymentStatus>,
     }
 }
 

urllb/Cargo.toml

@@ -20,6 +20,7 @@ tracing = "0.1"
 diesel_migrations = "2"
 diesel = { version = "2.0.4", features = ["postgres"] }
 diesel-async = { version = "0.2", features = ["tokio-postgres", "postgres", "tokio", "bb8"] }
+diesel-derive-enum = { version = "2.1.0", features = ["postgres"] }
 ipnet = "2"
 bb8 = "0.8"
 bb8-postgres = "0.8"
@@ -37,3 +38,6 @@ envy = "0.4"
 dotenvy = "0.15"
 
 shared = { path = "../shared", features = ["diesel"] }
+
+async-stripe = { version = "0.37", features = ["runtime-tokio-hyper-rustls-webpki"] }
+thiserror = "1"

urllb/migrations/2024-07-02-221141_add_payment/down.sql

@@ -0,0 +1,7 @@
+-- This file should undo anything in `up.sql`
+
+alter table links
+    drop column stripe_session_id;
+
+alter table links
+    drop column payment_status;

urllb/migrations/2024-07-02-221141_add_payment/up.sql

@@ -0,0 +1,9 @@
+-- Your SQL goes here
+
+create type payment_status as enum ('pending', 'succeeded', 'failed');
+
+alter table links
+    add column stripe_session_id text;
+
+alter table links
+    add column payment_status payment_status;

urllb/src/error.rs

@@ -0,0 +1,24 @@
+use axum::http::StatusCode;
+use diesel_async::pooled_connection::PoolError;
+use tracing::error;
+
+pub type Result<T> = std::result::Result<T, Error>;
+
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("Database error: {0}")]
+    DieselError(#[from] diesel::result::Error),
+    #[error("Stripe Error: {0}")]
+    StripeError(#[from] stripe::StripeError),
+    #[error("Pool error: {0}")]
+    PoolError(#[from] bb8::RunError<PoolError>),
+}
+
+impl From<Error> for StatusCode {
+    fn from(err: Error) -> Self {
+        error!("Internal server error: {:?}", err);
+        match err {
+            _ => StatusCode::INTERNAL_SERVER_ERROR,
+        }
+    }
+}

urllb/src/main.rs

@@ -1,8 +1,13 @@
 #[macro_use]
 extern crate diesel;
 
-use std::net::SocketAddr;
-
+use crate::db::Pool;
+use crate::db::{old_connection, run_migrations};
+use crate::error::Error;
+use crate::models::{CreateLinkDto, LinkDto};
+use crate::service::{
+    create_link, get_link_and_targets, increase_redirect_count, set_link_payment_status,
+};
 use axum::body::{Empty, Full};
 use axum::extract::{Path, State};
 use axum::http::{header, HeaderValue};
@@ -13,24 +18,29 @@ use axum::{
     http::StatusCode,
     response::IntoResponse,
     routing::{get, post},
-    Json, Router,
+    Extension, Json, Router,
 };
 use axum_client_ip::{SecureClientIp, SecureClientIpSource};
 use diesel_async::pooled_connection::AsyncDieselConnectionManager;
 use include_dir::{include_dir, Dir};
 use lazy_static::lazy_static;
+use nanoid::nanoid;
 use rand::seq::SliceRandom;
 use serde::{Deserialize, Serialize};
+use shared::{CreateResult, PaymentStatus};
+use std::net::SocketAddr;
+use std::str::FromStr;
+use stripe::{
+    CheckoutSession, CheckoutSessionId, CheckoutSessionMode, CreateCheckoutSession,
+    CreateCheckoutSessionCustomText, CreateCheckoutSessionCustomTextSubmit,
+    CreateCheckoutSessionLineItems,
+};
 use tokio::io;
 use tower_http::services::ServeDir;
 use validator::Validate;
 
-use crate::db::Pool;
-use crate::db::{old_connection, run_migrations};
-use crate::models::{CreateLinkDto, LinkDto};
-use crate::service::{create_link, get_link_and_targets, increase_redirect_count};
-
 mod db;
+mod error;
 mod models;
 mod schema;
 mod service;
@@ -42,6 +52,7 @@ static STATIC_DIR: Dir<'_> = include_dir!("$CARGO_MANIFEST_DIR/../web/dist");
 struct Config {
     ip_source: SecureClientIpSource,
     database_url: String,
+    stripe_secret_key: String,
 }
 
 #[tokio::main]
@@ -56,6 +67,8 @@ async fn main() {
 
     tracing_subscriber::fmt::init();
 
+    let stripe_client = stripe::Client::new(config.stripe_secret_key);
+
     let manager = AsyncDieselConnectionManager::new(config.database_url);
 
     let pool = Pool::builder().build(manager).await.unwrap();
@@ -83,6 +96,7 @@ async fn main() {
         .nest("/static", static_router)
         .route("/:link", get(link).post(post_link))
         .with_state(pool)
+        .layer(Extension(stripe_client))
         .layer(config.ip_source.into_extension());
 
     let addr = SocketAddr::from(([0, 0, 0, 0], 3000));
@@ -110,6 +124,15 @@ async fn link(
         .await
         .map_err(|_| StatusCode::NOT_FOUND)?;
 
+    if let Some(payment_status) = &link.payment_status {
+        match payment_status {
+            PaymentStatus::Pending | PaymentStatus::Failed => {
+                return Err(StatusCode::NOT_FOUND);
+            }
+            PaymentStatus::Succeeded => {}
+        }
+    }
+
     if link.fraud {
         return Ok(Response::builder()
             .status(StatusCode::UNAVAILABLE_FOR_LEGAL_REASONS)
@@ -154,8 +177,16 @@ lazy_static! {
         .collect();
 }
 
+lazy_static! {
+    static ref WHITELIST: Vec<String> = include_str!("whitelist.txt")
+        .lines()
+        .map(|s| s.to_string())
+        .collect();
+}
+
 async fn post_link(
     State(pool): State<Pool>,
+    Extension(stripe): Extension<stripe::Client>,
     SecureClientIp(ip): SecureClientIp,
     Json(body): Json<CreateLinkDto>,
 ) -> Result<impl IntoResponse, StatusCode> {
@@ -170,34 +201,112 @@ async fn post_link(
         return Err(StatusCode::FORBIDDEN);
     }
 
-    let mut connection = pool
-        .get()
-        .await
-        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+    let whitelisted = WHITELIST.iter().any(|w| {
+        body.targets
+            .iter()
+            .map(|t| &t.target_url)
+            .any(|t| t.starts_with(w))
+    });
 
-    let (link, target_results) = create_link(&mut connection, &body, ip.into())
-        .await
-        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+    let url = nanoid!(5);
 
-    Ok(Json(LinkDto {
-        link,
-        targets: target_results,
-    }))
+    let success_url = format!("https://hurlurl.com/info/{}", url);
+
+    let session = if !whitelisted {
+        let mut create_session = CreateCheckoutSession {
+            line_items: Some(vec![CreateCheckoutSessionLineItems {
+                price: Some("price_1PYEggFEynvp7vAemBawXewH".to_string()),
+                quantity: Some(1),
+                ..Default::default()
+            }]),
+            mode: Some(CheckoutSessionMode::Payment),
+            success_url: Some(&success_url),
+            cancel_url: Some("https://hurlurl.com"),
+            ..Default::default()
+        };
+
+        let session = CheckoutSession::create(&stripe, create_session)
+            .await
+            .map_err(error::Error::StripeError)?;
+        Some(session)
+    } else {
+        None
+    };
+
+    let mut connection = pool.get().await.map_err(Error::PoolError)?;
+    let (link, target_results) = create_link(
+        &mut connection,
+        &body,
+        &url,
+        ip.into(),
+        session.as_ref().map(|s| s.id.to_string()),
+    )
+    .await?;
+
+    match session {
+        None => Ok(Json(CreateResult::Link(LinkDto {
+            link,
+            targets: target_results,
+        }))),
+        Some(session) => Ok(Json(CreateResult::StripeRedirect(
+            session.url.ok_or(StatusCode::INTERNAL_SERVER_ERROR)?,
+        ))),
+    }
 }
 
 async fn link_info(
     Path(params): Path<Params>,
     State(pool): State<Pool>,
+    Extension(stripe): Extension<stripe::Client>,
 ) -> Result<impl IntoResponse, StatusCode> {
     let mut connection = pool
         .get()
         .await
         .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
 
-    let (link, results) = get_link_and_targets(&mut connection, &params.link)
+    let (mut link, results) = get_link_and_targets(&mut connection, &params.link)
         .await
         .map_err(|_| StatusCode::NOT_FOUND)?;
 
+    if let Some(status) = &link.payment_status {
+        match status {
+            PaymentStatus::Pending => {
+                if let Some(id) = &link.stripe_session_id {
+                    let session = CheckoutSession::retrieve(
+                        &stripe,
+                        &CheckoutSessionId::from_str(id).unwrap(),
+                        &[],
+                    )
+                    .await
+                    .map_err(Error::StripeError)?;
+
+                    if session.status == Some(stripe::CheckoutSessionStatus::Complete) {
+                        set_link_payment_status(
+                            &mut connection,
+                            &params.link,
+                            PaymentStatus::Succeeded,
+                        )
+                        .await?;
+                    } else {
+                        set_link_payment_status(
+                            &mut connection,
+                            &params.link,
+                            PaymentStatus::Failed,
+                        )
+                        .await?;
+                        return Err(StatusCode::NOT_FOUND);
+                    }
+                }
+            }
+            PaymentStatus::Failed => {
+                return Err(StatusCode::NOT_FOUND);
+            }
+            PaymentStatus::Succeeded => {}
+        }
+    }
+
+    link.stripe_session_id = None;
+
     Ok(Json(LinkDto {
         link,
         targets: results,

urllb/src/models.rs

@@ -9,6 +9,8 @@ pub struct NewLink<'a> {
     pub url: &'a str,
     pub permanent_redirect: bool,
     pub created_by_ip: Option<ipnet::IpNet>,
+    pub stripe_session_id: Option<&'a str>,
+    pub payment_status: Option<PaymentStatus>,
 }
 
 #[derive(Insertable)]