fix rate limit being global instead of per-IP

.env.example

@@ -1,5 +1,6 @@
 # bnuy-api
 POST_LIMIT=20
+SECRET=generate-me
 LOG_LEVEL=INFO
 FORWARDED_ALLOW_IPS=172.16.0.0/12
 

collector.py

@@ -73,7 +73,7 @@ async def save_picture(pool):
                     async with pool.acquire() as conn:
                         async with conn.cursor() as cursor:
                             await cursor.execute(
-                                "INSERT INTO images (url, filename, subreddit) VALUES (%s, %s, %s)",
+                                "INSERT INTO images (url, filename, source) VALUES (%s, %s, %s)",
                                 (url, generate_filename, subreddit)
                             )
                             await conn.commit()

main.py

@@ -1,8 +1,10 @@
+import uuid
 import aiohttp
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException, Request, UploadFile
 from fastapi.responses import FileResponse
 from fastapi import Depends, FastAPI
-from pyrate_limiter import Duration, Limiter, Rate
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+from pyrate_limiter import AbstractBucket, BucketFactory, Duration, InMemoryBucket, Limiter, MonotonicClock, Rate, RateItem
 from fastapi_limiter.depends import RateLimiter
 from fastapi import Depends
 import asyncmy
@@ -31,6 +33,9 @@ logging.basicConfig(
     handlers=[file_handler, console_handler],
 )
 
+security = HTTPBearer()
+SECRET = os.getenv("SECRET")
+
 @asynccontextmanager
 async def connect_db(app: FastAPI):
     app.state.pool = await asyncmy.create_pool(
@@ -55,8 +60,27 @@ async def connect_db(app: FastAPI):
         app.state.pool.close()
         await app.state.pool.wait_closed()
 
+class MultiBucketFactory(BucketFactory):
+    def __init__(self, rates, clock):
+        self.clock = clock
+        self.rates = rates
+        self.buckets = {}
+
+    def wrap_item(self, name: str, weight: int = 1) -> RateItem:
+        """Time-stamping item, return a RateItem"""
+        now = self.clock.now()
+        return RateItem(name, now, weight=weight)
+
+    def get(self, item: RateItem) -> AbstractBucket:
+        if item.name not in self.buckets:
+            new_bucket = self.create(InMemoryBucket, self.rates)
+            self.buckets.update({item.name: new_bucket})
+
+        return self.buckets[item.name]
+
 app = FastAPI(lifespan=connect_db)
-limiter = Limiter(Rate(50, Duration.MINUTE))
+rates = [Rate(50, Duration.MINUTE)]
+limiter = Limiter(MultiBucketFactory(rates,MonotonicClock()))
 
 async def create_tables(pool):
     async with pool.acquire() as conn:
@@ -66,7 +90,7 @@ async def create_tables(pool):
                 id INT AUTO_INCREMENT PRIMARY KEY,
                 url VARCHAR(255) NOT NULL,
                 filename VARCHAR(255) NOT NULL,
-                subreddit VARCHAR(255) NOT NULL,
+                source VARCHAR(255) NOT NULL,
                 created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP    
                 )
             """)
@@ -80,12 +104,12 @@ async def root():
 async def get_random_bnuy(request: Request):
     async with app.state.pool.acquire() as conn:
         async with conn.cursor() as cursor:
-            await cursor.execute("SELECT filename, subreddit, url FROM images ORDER BY RAND() LIMIT 1;")
+            await cursor.execute("SELECT filename, source, url FROM images ORDER BY RAND() LIMIT 1;")
             result = await cursor.fetchone()
             if result:
                 filepath = os.path.join("data/images", result[0])
                 if os.path.exists(filepath):
-                    return {"url": f"{request.base_url}images/{result[0]}", "source": f"https://www.reddit.com/r/{result[1]}/", "original_url": result[2]}
+                    return {"url": f"{request.base_url}images/{result[0]}", "source": result[1], "original_url": result[2]}
                 else:
                     raise HTTPException(status_code=404, detail="Image file not found")
             else:
@@ -115,4 +139,33 @@ async def fetch_images():
             logging.info("Image collection completed. Sleeping for 1 hour...")
         except Exception as e:
             logging.error(f"Error during image collection: {e}")
-        await asyncio.sleep(86400)  # Sleep for 24 hours
+        await asyncio.sleep(86400)  # Sleep for 24 hours
+
+@app.post("/upload", dependencies=[Depends(RateLimiter(limiter=limiter))])
+async def upload_image(file: UploadFile, request: Request, credentials: HTTPAuthorizationCredentials = Depends(security)):
+    # if SECRET isn't set, return Unauthrorized
+    if credentials.credentials != SECRET or credentials.credentials is None:
+        raise HTTPException(status_code=401, detail="Unauthorized")
+    if not file:
+        raise HTTPException(status_code=400, detail="No file uploaded")
+    if not file.filename.lower().endswith(('.jpg', '.jpeg', '.png', '.gif')) or file.content_type not in ["image/jpeg", "image/png", "image/gif"]:
+        raise HTTPException(status_code=400, detail="Unsupported file type")
+    try:
+        content = await file.read()
+        generate_filename = str(uuid.uuid4()) + os.path.splitext(file.filename)[1]
+        filename = os.path.join("data/images", generate_filename)
+        with open(filename, "wb") as f:
+            f.write(content)
+        
+        logging.info(f"Saved uploaded image to {filename}")
+
+        async with app.state.pool.acquire() as conn:
+            async with conn.cursor() as cursor:
+                await cursor.execute(
+                    "INSERT INTO images (url, filename, source) VALUES (%s, %s, %s)",
+                    (f"{request.base_url}images/{generate_filename}", generate_filename, "user_upload")
+                )
+                await conn.commit()
+    except Exception as e:
+        logging.error(f"Error saving uploaded image: {e}")
+        raise HTTPException(status_code=500, detail="Failed to save image")

requirements.txt

@@ -1,4 +1,5 @@
 fastapi[standard]
 fastapi_limiter
 asyncmy
-aiohttp
+aiohttp
+python-multipart