add origin check

routes/generate/index.js

@@ -33,6 +33,23 @@ const IMAGE_SIZES = {
   },
 };
 
+const ALLOWED_ORIGINS = [
+  "huggingface.co",
+  "enzostvs-deepsite.hf.space",
+  "hf.co",
+  "localhost:3001",
+];
+
+function isOriginAllowed(request) {
+  const origin = request.headers.origin ?? request.headers.host;
+  const referer = request.headers.referer;
+
+  return (
+    ALLOWED_ORIGINS.some((allowed) => origin.includes(allowed)) ||
+    ALLOWED_ORIGINS.some((allowed) => referer?.startsWith(allowed))
+  );
+}
+
 module.exports = async function (fastify, opts) {
   fastify.get("/:inputs", async function (request, reply) {
     let { inputs } = request.params;
@@ -51,6 +68,10 @@ module.exports = async function (fastify, opts) {
       return reply.header("Content-Type", "image/jpeg").send(file);
     }
 
+    if (!isOriginAllowed(request)) {
+      return reply.code(403).send({ error: "Forbidden" });
+    }
+
     const { height, width } =
       IMAGE_SIZES[format ?? "square"] ?? IMAGE_SIZES["square"];
 

routes/root.js

@@ -200,9 +200,7 @@ module.exports = async function (fastify, opts) {
       <div class="code-block">
         <code>GET /generate/<span class="param">{your-prompt}</span>?format=<span class="param">{format}</span></code>
       </div>
-      
-      <a href="/generate/a-cute-robot-painting-a-sunset" class="example-link">Try an Example →</a>
-    </section>
+          </section>
     
     <section>
       <h2>Available Formats</h2>