From 9a1f68e567c4b897712d3f6821ea093fea042806 Mon Sep 17 00:00:00 2001
From: Connor Prussin <connor@prussin.net>
Date: Wed, 11 Sep 2024 08:52:05 -0700
Subject: [PATCH] feat(staking): add region blocking

---
 apps/staking/src/app/blocked/page.tsx         |  1 +
 apps/staking/src/components/Blocked/index.tsx | 19 +++++++++++++++
 apps/staking/src/components/Error/index.tsx   |  2 +-
 apps/staking/src/components/Header/index.tsx  |  2 --
 .../staking/src/components/NotFound/index.tsx |  2 +-
 .../src/components/WalletButton/index.tsx     |  8 +++++++
 apps/staking/src/config/isomorphic.ts         | 10 ++++++++
 apps/staking/src/config/server.ts             |  7 ++++++
 apps/staking/src/middleware.ts                | 24 +++++++++++++++++++
 9 files changed, 71 insertions(+), 4 deletions(-)
 create mode 100644 apps/staking/src/app/blocked/page.tsx
 create mode 100644 apps/staking/src/components/Blocked/index.tsx
 create mode 100644 apps/staking/src/middleware.ts

diff --git a/apps/staking/src/app/blocked/page.tsx b/apps/staking/src/app/blocked/page.tsx
new file mode 100644
index 0000000000..0de6e61b29
--- /dev/null
+++ b/apps/staking/src/app/blocked/page.tsx
@@ -0,0 +1 @@
+export { Blocked as default } from "../../components/Blocked";
diff --git a/apps/staking/src/components/Blocked/index.tsx b/apps/staking/src/components/Blocked/index.tsx
new file mode 100644
index 0000000000..2fe9a8372c
--- /dev/null
+++ b/apps/staking/src/components/Blocked/index.tsx
@@ -0,0 +1,19 @@
+import { LinkButton } from "../Button";
+
+export const Blocked = () => (
+  <main className="grid size-full place-content-center py-20 text-center">
+    <h1 className="mb-8 text-4xl font-semibold text-pythpurple-400">
+      {"We're sorry"}
+    </h1>
+    <p className="mb-20 text-lg">
+      This program is currently unavailable to users in your region
+    </p>
+    <LinkButton
+      className="place-self-center px-24 py-3"
+      href="https://www.pyth.network"
+      target="_blank"
+    >
+      Read More About Pyth
+    </LinkButton>
+  </main>
+);
diff --git a/apps/staking/src/components/Error/index.tsx b/apps/staking/src/components/Error/index.tsx
index 5baed81814..dd465bc477 100644
--- a/apps/staking/src/components/Error/index.tsx
+++ b/apps/staking/src/components/Error/index.tsx
@@ -17,7 +17,7 @@ export const Error = ({ error, reset }: Props) => {
   }, [error, logger]);
 
   return (
-    <main className="flex size-full flex-col items-center justify-center gap-2 text-center">
+    <main className="flex size-full flex-col items-center justify-center gap-2 py-20 text-center">
       <div className="mb-10 flex flex-row items-center gap-6">
         <ExclamationTriangleIcon className="size-16 text-red-700" />
         <h1 className="text-3xl font-light">Uh oh!</h1>
diff --git a/apps/staking/src/components/Header/index.tsx b/apps/staking/src/components/Header/index.tsx
index 92d1e184cc..77cc549708 100644
--- a/apps/staking/src/components/Header/index.tsx
+++ b/apps/staking/src/components/Header/index.tsx
@@ -1,5 +1,3 @@
-"use client";
-
 import clsx from "clsx";
 import type { HTMLAttributes } from "react";
 
diff --git a/apps/staking/src/components/NotFound/index.tsx b/apps/staking/src/components/NotFound/index.tsx
index c7aa5012e6..21147010d1 100644
--- a/apps/staking/src/components/NotFound/index.tsx
+++ b/apps/staking/src/components/NotFound/index.tsx
@@ -1,7 +1,7 @@
 import { LinkButton } from "../Button";
 
 export const NotFound = () => (
-  <main className="grid size-full place-content-center text-center">
+  <main className="grid size-full place-content-center py-20 text-center">
     <h1 className="mb-8 text-4xl font-semibold text-pythpurple-400">
       Not Found
     </h1>
diff --git a/apps/staking/src/components/WalletButton/index.tsx b/apps/staking/src/components/WalletButton/index.tsx
index 3eb729c332..6c27f1fe09 100644
--- a/apps/staking/src/components/WalletButton/index.tsx
+++ b/apps/staking/src/components/WalletButton/index.tsx
@@ -14,6 +14,7 @@ import { useWallet } from "@solana/wallet-adapter-react";
 import { useWalletModal } from "@solana/wallet-adapter-react-ui";
 import type { PublicKey } from "@solana/web3.js";
 import clsx from "clsx";
+import { useSelectedLayoutSegment } from "next/navigation";
 import {
   type ComponentProps,
   type ComponentType,
@@ -33,6 +34,7 @@ import {
   SubmenuTrigger,
 } from "react-aria-components";
 
+import { BLOCKED_SEGMENT } from "../../config/isomorphic";
 import {
   StateType as ApiStateType,
   type States,
@@ -47,6 +49,12 @@ import { ModalDialog } from "../ModalDialog";
 type Props = Omit<ComponentProps<typeof Button>, "onClick" | "children">;
 
 export const WalletButton = (props: Props) => {
+  const segment = useSelectedLayoutSegment();
+  // eslint-disable-next-line unicorn/no-null
+  return segment === BLOCKED_SEGMENT ? null : <WalletButtonImpl {...props} />;
+};
+
+const WalletButtonImpl = (props: Props) => {
   const api = useApi();
 
   switch (api.type) {
diff --git a/apps/staking/src/config/isomorphic.ts b/apps/staking/src/config/isomorphic.ts
index db2626f0dc..935e0a2221 100644
--- a/apps/staking/src/config/isomorphic.ts
+++ b/apps/staking/src/config/isomorphic.ts
@@ -11,3 +11,13 @@
  * with the optimized React build, etc.
  */
 export const IS_PRODUCTION_BUILD = process.env.NODE_ENV === "production";
+
+/**
+ * Region-blocked requests will be redirected here.  This is used in the
+ * middleware to implement the block, and also consumed in any components that
+ * are part of the page layout but need to know if the request is blocked from
+ * accessing the app, such as the WalletButton in the app header.
+ *
+ * Don't change unless you also change the relevant app route path to match.
+ */
+export const BLOCKED_SEGMENT = "blocked";
diff --git a/apps/staking/src/config/server.ts b/apps/staking/src/config/server.ts
index 4643f40a4d..b3f58c9ee9 100644
--- a/apps/staking/src/config/server.ts
+++ b/apps/staking/src/config/server.ts
@@ -36,6 +36,13 @@ export const WALLETCONNECT_PROJECT_ID = demandInProduction(
 );
 export const RPC = process.env.RPC;
 export const IS_MAINNET = process.env.IS_MAINNET !== undefined;
+export const BLOCKED_REGIONS =
+  process.env.BLOCKED_REGIONS === undefined ||
+  process.env.BLOCKED_REGIONS === ""
+    ? []
+    : process.env.BLOCKED_REGIONS.split(",").map((region) =>
+        region.toLowerCase().trim(),
+      );
 
 class MissingEnvironmentError extends Error {
   constructor(name: string) {
diff --git a/apps/staking/src/middleware.ts b/apps/staking/src/middleware.ts
new file mode 100644
index 0000000000..a51d0acc1e
--- /dev/null
+++ b/apps/staking/src/middleware.ts
@@ -0,0 +1,24 @@
+import { type NextRequest, NextResponse } from "next/server";
+
+import { BLOCKED_SEGMENT } from "./config/isomorphic";
+import { BLOCKED_REGIONS } from "./config/server";
+
+export const middleware = (request: NextRequest) => {
+  if (blockRequest(request)) {
+    return NextResponse.rewrite(new URL(`/${BLOCKED_SEGMENT}`, request.url));
+  } else if (request.nextUrl.pathname.startsWith("/blocked")) {
+    return NextResponse.rewrite(new URL("/not-found", request.url));
+  } else {
+    return;
+  }
+};
+
+const blockRequest = (request: NextRequest) =>
+  request.geo?.country !== undefined &&
+  BLOCKED_REGIONS.includes(request.geo.country.toLowerCase());
+
+export const config = {
+  matcher: [
+    "/((?!_next/static|_next/image|favicon.ico|sitemap.xml|robots.txt).*)",
+  ],
+};