Incorporate feedback from review

Author: mattt

src/index.ts

@@ -65,8 +65,8 @@ app.use("*", async (c, next) => {
   await next()
 })
 
-const mcpServer = createMcpServer()
 app.all("/mcp", async (c) => {
+  const mcpServer = createMcpServer(c.env)
   const transport = new StreamableHTTPTransport()
   await mcpServer.connect(transport)
   return transport.handleRequest(c)

src/lib/external/policy.ts

@@ -341,10 +341,11 @@ function isPrivateIPv4(hostname: string): boolean {
     return false
   }
 
-  const [a, b] = octets.map((octet) => Number.parseInt(octet, 10))
-  if (a > 255 || b > 255) {
+  const octetNumbers = octets.map((octet) => Number.parseInt(octet, 10))
+  if (octetNumbers.some((value) => value > 255)) {
     return false
   }
+  const [a, b] = octetNumbers
 
   return (
     a === 10 ||

src/lib/mcp.ts

@@ -3,6 +3,7 @@ import { z } from "zod"
 
 import {
   assertExternalDocumentationAccess,
+  type ExternalPolicyEnv,
   extractExternalDocumentationBasePath,
   fetchExternalDocCJSON,
   validateExternalDocumentationUrl,
@@ -12,7 +13,7 @@ import { searchAppleDeveloperDocs } from "./search"
 import { generateAppleDocUrl, normalizeDocumentationPath } from "./url"
 import { fetchHIGPageData, renderHIGFromJSON } from "./hig"
 
-export function createMcpServer() {
+export function createMcpServer(externalPolicyEnv: ExternalPolicyEnv = {}) {
   const server = new McpServer({
     name: "sosumi.ai",
     version: "1.0.0",
@@ -297,7 +298,7 @@ export function createMcpServer() {
     async ({ url }) => {
       try {
         const targetUrl = validateExternalDocumentationUrl(url)
-        await assertExternalDocumentationAccess(targetUrl, {})
+        await assertExternalDocumentationAccess(targetUrl, externalPolicyEnv)
         const jsonData = await fetchExternalDocCJSON(targetUrl)
         const externalBasePath = extractExternalDocumentationBasePath(targetUrl)
         const markdown = await renderFromJSON(jsonData, targetUrl.toString(), {

tests/external.test.ts

@@ -49,6 +49,32 @@ describe("External Swift-DocC support", () => {
     ).rejects.toThrow(/not allowlisted/)
   })
 
+  it("should block local and private hosts unless explicitly allowlisted", async () => {
+    const blockedHosts = ["127.0.0.1", "10.0.0.1", "192.168.1.1", "172.16.0.1", "[::1]"]
+
+    global.fetch = vi.fn()
+    for (const host of blockedHosts) {
+      await expect(
+        assertExternalDocumentationAccess(new URL(`https://${host}/documentation/example`), {}),
+      ).rejects.toThrow(/local or private host/)
+    }
+
+    // Host policy should reject before any robots.txt request is made.
+    expect(global.fetch).not.toHaveBeenCalled()
+  })
+
+  it("should allow an explicitly allowlisted private host", async () => {
+    global.fetch = vi
+      .fn()
+      .mockResolvedValue(new Response("User-agent: *\nAllow: /", { status: 200 }))
+
+    await expect(
+      assertExternalDocumentationAccess(new URL("https://127.0.0.1/documentation/example"), {
+        EXTERNAL_DOC_HOST_ALLOWLIST: "127.0.0.1",
+      }),
+    ).resolves.toBeUndefined()
+  })
+
   it("should deny when robots.txt disallows", async () => {
     global.fetch = vi.fn().mockResolvedValue(
       new Response("User-agent: *\nDisallow: /swift-argument-parser/documentation/", {

tests/render.test.ts

@@ -290,9 +290,7 @@ describe("Render Function", () => {
                     [
                       {
                         type: "paragraph",
-                        inlineContent: [
-                          { type: "text", text: "The value for the import option." },
-                        ],
+                        inlineContent: [{ type: "text", text: "The value for the import option." }],
                       },
                     ],
                   ],