From ad6dd98befab200f4909c66d8421d6acfcdcec69 Mon Sep 17 00:00:00 2001
From: yivlad <jacenko.vlad@gmail.com>
Date: Mon, 17 Jun 2024 10:34:27 +0200
Subject: [PATCH] Add IP status endpoint response schema

---
 .../compliance/logic/useBlockedPages.ts         |  2 +-
 .../compliance/logic/useIPAndAddressCheck.ts    |  4 ++--
 .../features/compliance/logic/useVpnCheck.ts    | 17 +++++++++--------
 3 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/packages/app/src/features/compliance/logic/useBlockedPages.ts b/packages/app/src/features/compliance/logic/useBlockedPages.ts
index 78a824524..d5a6e0154 100644
--- a/packages/app/src/features/compliance/logic/useBlockedPages.ts
+++ b/packages/app/src/features/compliance/logic/useBlockedPages.ts
@@ -11,7 +11,7 @@ export function useBlockedPages(): (keyof typeof paths)[] {
   // biome-ignore lint/correctness/useHookAtTopLevel: <explanation>
   const vpnCheck = useVpnCheck({ authUrl: apiUrl })
 
-  if (vpnCheck.data && blockedPagesByCountryCode[vpnCheck.data.countryCode] !== undefined) {
+  if (vpnCheck.data?.countryCode && blockedPagesByCountryCode[vpnCheck.data.countryCode] !== undefined) {
     return blockedPagesByCountryCode[vpnCheck.data.countryCode]!
   }
 
diff --git a/packages/app/src/features/compliance/logic/useIPAndAddressCheck.ts b/packages/app/src/features/compliance/logic/useIPAndAddressCheck.ts
index 84a636ac7..b1e3d7858 100644
--- a/packages/app/src/features/compliance/logic/useIPAndAddressCheck.ts
+++ b/packages/app/src/features/compliance/logic/useIPAndAddressCheck.ts
@@ -34,7 +34,7 @@ export function useIPAndAddressCheck(): UseIPAndAddressCheck {
     return { blocked: true, reason: 'vpn-detected' }
   }
 
-  if (vpnCheck.data && blockedCountryCodes.includes(vpnCheck.data.countryCode)) {
+  if (vpnCheck.data?.countryCode && blockedCountryCodes.includes(vpnCheck.data.countryCode)) {
     return { blocked: true, reason: 'region-blocked', data: { countryCode: vpnCheck.data.countryCode } }
   }
 
@@ -42,7 +42,7 @@ export function useIPAndAddressCheck(): UseIPAndAddressCheck {
     return { blocked: true, reason: 'address-not-allowed' }
   }
 
-  if (vpnCheck.data && isCurrentPageBlocked) {
+  if (vpnCheck.data?.countryCode && isCurrentPageBlocked) {
     return {
       blocked: true,
       reason: 'page-not-available-in-region',
diff --git a/packages/app/src/features/compliance/logic/useVpnCheck.ts b/packages/app/src/features/compliance/logic/useVpnCheck.ts
index a2b1f4292..1703fab82 100644
--- a/packages/app/src/features/compliance/logic/useVpnCheck.ts
+++ b/packages/app/src/features/compliance/logic/useVpnCheck.ts
@@ -1,19 +1,23 @@
 import { useQuery } from '@tanstack/react-query'
 
+import { z } from 'zod'
 import { ReadHookParams } from './types'
 
 type VpnResponse = {
   isConnectedToVpn: boolean
-  countryCode: string
+  countryCode: string | null | undefined
 }
 
+const apiResponseSchema = z.object({
+  is_vpn: z.boolean(),
+  country_code: z.string().nullable().optional(),
+})
+
 async function checkVpn(authUrl: string): Promise<VpnResponse> {
   if (!authUrl) {
     throw new Error('Missing auth URL')
   }
 
-  let isConnectedToVpn = false
-  let countryCode = ''
   // TODO is this the best way to get a user's IP address?
   const ipRes = await fetch('https://api.ipify.org/?format=json')
   if (!ipRes.ok) {
@@ -26,12 +30,9 @@ async function checkVpn(authUrl: string): Promise<VpnResponse> {
     throw new Error('Could not fetch VPN status')
   }
 
-  const { country_code, is_vpn } = (await vpnRes.json()) as any
-
-  countryCode = country_code
-  isConnectedToVpn = is_vpn
+  const { country_code, is_vpn } = apiResponseSchema.parse(await vpnRes.json())
 
-  return { countryCode, isConnectedToVpn }
+  return { countryCode: country_code, isConnectedToVpn: is_vpn }
 }
 
 type Props = ReadHookParams<VpnResponse> & { authUrl: string }