From 3e40267cfa3342e1b715633ba2acd1c7191a9269 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Sun, 20 Oct 2024 18:33:05 -0700
Subject: [PATCH] block urls pointing to internal addresses (#1012)

---
 skyvern/config.py                    |  1 +
 skyvern/exceptions.py                |  8 ++++++++
 skyvern/forge/sdk/core/validators.py | 18 ++++++++++++++++++
 skyvern/forge/sdk/schemas/tasks.py   | 16 ++++++++++++++--
 4 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/skyvern/config.py b/skyvern/config.py
index 56442aad..fb6356f8 100644
--- a/skyvern/config.py
+++ b/skyvern/config.py
@@ -34,6 +34,7 @@ class Settings(BaseSettings):
     LOG_LEVEL: str = "INFO"
     PORT: int = 8000
     ALLOWED_ORIGINS: list[str] = ["*"]
+    BLOCKED_HOSTS: list[str] = ["localhost"]
 
     # Secret key for JWT. Please generate your own secret key in production
     SECRET_KEY: str = "PLACEHOLDER"
diff --git a/skyvern/exceptions.py b/skyvern/exceptions.py
index e5b976a1..a389a114 100644
--- a/skyvern/exceptions.py
+++ b/skyvern/exceptions.py
@@ -502,3 +502,11 @@ class InvalidUrl(SkyvernHTTPException):
         super().__init__(
             f"Invalid URL: {url}. Skyvern supports HTTP and HTTPS urls.", status_code=status.HTTP_400_BAD_REQUEST
         )
+
+
+class BlockedHost(SkyvernHTTPException):
+    def __init__(self, host: str) -> None:
+        super().__init__(
+            f"The host in your url is blocked: {host}",
+            status_code=status.HTTP_400_BAD_REQUEST,
+        )
diff --git a/skyvern/forge/sdk/core/validators.py b/skyvern/forge/sdk/core/validators.py
index c2107fb0..080e79c9 100644
--- a/skyvern/forge/sdk/core/validators.py
+++ b/skyvern/forge/sdk/core/validators.py
@@ -1,5 +1,8 @@
+import ipaddress
+
 from pydantic import HttpUrl, ValidationError, parse_obj_as
 
+from skyvern.config import settings
 from skyvern.exceptions import InvalidUrl
 
 
@@ -11,3 +14,18 @@ def validate_url(url: str) -> str:
     except ValidationError:
         # Handle the validation error
         raise InvalidUrl(url=url)
+
+
+def is_blocked_host(host: str) -> bool:
+    try:
+        ip = ipaddress.ip_address(host)
+        # Check if the IP is private, link-local, loopback, or reserved
+        return ip.is_private or ip.is_link_local or ip.is_loopback or ip.is_reserved
+    except ValueError:
+        # If the host is not a valid IP address (e.g., it's a domain name like localhost), handle it here
+        for blocked_host in settings.BLOCKED_HOSTS:
+            if blocked_host == host:
+                return True
+        return False
+    except Exception:
+        return False
diff --git a/skyvern/forge/sdk/schemas/tasks.py b/skyvern/forge/sdk/schemas/tasks.py
index 11e5bb45..df0b5104 100644
--- a/skyvern/forge/sdk/schemas/tasks.py
+++ b/skyvern/forge/sdk/schemas/tasks.py
@@ -4,9 +4,10 @@ from datetime import datetime
 from enum import StrEnum
 from typing import Any
 
-from pydantic import BaseModel, Field, HttpUrl
+from pydantic import BaseModel, Field, HttpUrl, field_validator
 
-from skyvern.exceptions import InvalidTaskStatusTransition, TaskAlreadyCanceled
+from skyvern.exceptions import BlockedHost, InvalidTaskStatusTransition, TaskAlreadyCanceled
+from skyvern.forge.sdk.core.validators import is_blocked_host
 
 
 class ProxyLocation(StrEnum):
@@ -89,6 +90,17 @@ class TaskRequest(TaskBase):
     )
     totp_verification_url: HttpUrl | None = None
 
+    @field_validator("url", "webhook_callback_url", "totp_verification_url")
+    @classmethod
+    def validate_urls(cls, v: HttpUrl | None) -> HttpUrl | None:
+        if not v or not v.host:
+            return None
+        host = v.host
+        blocked = is_blocked_host(host)
+        if blocked:
+            raise BlockedHost(host=host)
+        return v
+
 
 class TaskStatus(StrEnum):
     created = "created"