mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-10-03 16:27:47 +08:00
Add Ollama Cloud support with Base URL and API Key configuration (#692)
* Add Ollama Cloud support with Base URL and API Key configuration * implemented feedback * fix: use OLLAMA_BASE_URL env fallback in validate-model endpoint * Remove dedicated Ollama configuration block * security(ollama): prevent API key leak to client-controlled URLs * added test * fix: security hardening and Ollama Cloud default URL - Add server OLLAMA_API_KEY fallback to validate-model endpoint with SSRF guard mirroring ai-providers.ts - Tighten top-level SSRF exemption: only exempt Ollama when no server OLLAMA_API_KEY is configured - Update Electron config to support OLLAMA_API_KEY env var - Change default Ollama URL from localhost:11434 to ollama.com/api (Ollama Cloud) for web UI users - Add tests for server env combo, API-key-only, and SSRF guard scenarios --------- Co-authored-by: dayuan.jiang <[email protected]>
This commit is contained in:
Marvelous Ikponmwosa
co-authored by
dayuan.jiang
dayuan.jiang
parent
e171fbcdd8
commit
a5d1554c3f
@@ -174,10 +174,21 @@ export async function POST(req: Request) {
|
||||
}
|
||||
|
||||
case "ollama": {
|
||||
const ollama = createOllama({
|
||||
baseURL: baseUrl || "http://localhost:11434",
|
||||
// SECURITY: Mirror ai-providers.ts guard — only use server
|
||||
// OLLAMA_API_KEY when the URL is also from server config.
|
||||
const ollamaApiKey = baseUrl
|
||||
? apiKey || undefined
|
||||
: apiKey || process.env.OLLAMA_API_KEY || undefined
|
||||
const ollamaProvider = createOllama({
|
||||
baseURL:
|
||||
baseUrl ||
|
||||
process.env.OLLAMA_BASE_URL ||
|
||||
"https://ollama.com/api",
|
||||
...(ollamaApiKey && {
|
||||
headers: { Authorization: `Bearer ${ollamaApiKey}` },
|
||||
}),
|
||||
})
|
||||
model = ollama(modelId)
|
||||
model = ollamaProvider(modelId)
|
||||
break
|
||||
}
|
||||
|
||||
|
||||
@@ -282,6 +282,7 @@ export function ModelConfigDialog({
|
||||
// Check credentials based on provider type
|
||||
const isBedrock = selectedProvider.provider === "bedrock"
|
||||
const isEdgeOne = selectedProvider.provider === "edgeone"
|
||||
const isOllama = selectedProvider.provider === "ollama"
|
||||
const isVertexAI = selectedProvider.provider === "vertexai"
|
||||
if (isBedrock) {
|
||||
if (
|
||||
@@ -296,7 +297,7 @@ export function ModelConfigDialog({
|
||||
if (!selectedProvider.vertexApiKey) {
|
||||
return
|
||||
}
|
||||
} else if (!isEdgeOne && !selectedProvider.apiKey) {
|
||||
} else if (!isEdgeOne && !isOllama && !selectedProvider.apiKey) {
|
||||
return
|
||||
}
|
||||
|
||||
@@ -1030,9 +1031,7 @@ export function ModelConfigDialog({
|
||||
</div>
|
||||
</>
|
||||
) : selectedProvider.provider ===
|
||||
"ollama" ||
|
||||
selectedProvider.provider ===
|
||||
"edgeone" ? (
|
||||
"edgeone" ? (
|
||||
<div className="space-y-3">
|
||||
<div className="flex items-center gap-2">
|
||||
<Button
|
||||
@@ -1100,6 +1099,9 @@ export function ModelConfigDialog({
|
||||
dict.modelConfig
|
||||
.apiKey
|
||||
}
|
||||
{selectedProvider.provider ===
|
||||
"ollama" &&
|
||||
` ${dict.modelConfig.optional}`}
|
||||
</Label>
|
||||
<div className="flex gap-2">
|
||||
<div className="relative flex-1">
|
||||
@@ -1163,7 +1165,9 @@ export function ModelConfigDialog({
|
||||
handleValidate
|
||||
}
|
||||
disabled={
|
||||
!selectedProvider.apiKey ||
|
||||
(selectedProvider.provider !==
|
||||
"ollama" &&
|
||||
!selectedProvider.apiKey) ||
|
||||
validationStatus ===
|
||||
"validating"
|
||||
}
|
||||
|
||||
@@ -359,9 +359,9 @@ const PROVIDER_ENV_MAP: Record<string, { apiKey: string; baseUrl: string }> = {
|
||||
baseUrl: "MODELSCOPE_BASE_URL",
|
||||
},
|
||||
gateway: { apiKey: "AI_GATEWAY_API_KEY", baseUrl: "AI_GATEWAY_BASE_URL" },
|
||||
// bedrock and ollama don't use API keys in the same way
|
||||
// bedrock doesn't use API keys in the same way
|
||||
bedrock: { apiKey: "", baseUrl: "" },
|
||||
ollama: { apiKey: "", baseUrl: "OLLAMA_BASE_URL" },
|
||||
ollama: { apiKey: "OLLAMA_API_KEY", baseUrl: "OLLAMA_BASE_URL" },
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
+3
-2
@@ -59,8 +59,9 @@ AI_MODEL=global.anthropic.claude-sonnet-4-5-20250929-v1:0
|
||||
# AZURE_REASONING_EFFORT=low # Optional: Azure reasoning effort (low, medium, high)
|
||||
# AZURE_REASONING_SUMMARY=detailed
|
||||
|
||||
# Ollama (Local) Configuration
|
||||
# OLLAMA_BASE_URL=http://localhost:11434/api # Optional, defaults to localhost
|
||||
# Ollama Configuration (Local or Cloud)
|
||||
# OLLAMA_BASE_URL=https://ollama.com/api # Optional, defaults to Ollama Cloud
|
||||
# OLLAMA_API_KEY=your-ollama-cloud-api-key # Optional: For Ollama Cloud or authenticated remote instances
|
||||
# OLLAMA_ENABLE_THINKING=true # Optional: Enable thinking for models that support it (e.g., qwen3)
|
||||
|
||||
# OpenRouter Configuration
|
||||
|
||||
+18
-5
@@ -596,7 +596,7 @@ function validateProviderCredentials(
|
||||
* - GOOGLE_GENERATIVE_AI_API_KEY: Google API key
|
||||
* - AZURE_RESOURCE_NAME, AZURE_API_KEY: Azure OpenAI credentials
|
||||
* - AWS_REGION, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY: AWS Bedrock credentials
|
||||
* - OLLAMA_BASE_URL: Ollama server URL (optional, defaults to http://localhost:11434)
|
||||
* - OLLAMA_BASE_URL: Ollama server URL (optional, defaults to https://ollama.com/api)
|
||||
* - OPENROUTER_API_KEY: OpenRouter API key
|
||||
* - DEEPSEEK_API_KEY: DeepSeek API key
|
||||
* - DEEPSEEK_BASE_URL: DeepSeek endpoint (optional)
|
||||
@@ -611,13 +611,15 @@ export function getAIModel(overrides?: ClientOverrides): ModelConfig {
|
||||
// SECURITY: Prevent SSRF attacks (GHSA-9qf7-mprq-9qgm)
|
||||
// If a custom baseUrl is provided, an API key MUST also be provided.
|
||||
// This prevents attackers from redirecting server API keys to malicious endpoints.
|
||||
// Exception: EdgeOne and Ollama providers don't require API keys
|
||||
// Exception: EdgeOne doesn't require API keys.
|
||||
// Ollama is exempt only when no server OLLAMA_API_KEY is configured;
|
||||
// when it IS configured, the outer guard also enforces client apiKey for custom baseUrls.
|
||||
if (
|
||||
overrides?.baseUrl &&
|
||||
!overrides?.apiKey &&
|
||||
!(overrides?.provider === "vertexai" && overrides?.vertexApiKey) &&
|
||||
overrides?.provider !== "edgeone" &&
|
||||
overrides?.provider !== "ollama"
|
||||
!(overrides?.provider === "ollama" && !process.env.OLLAMA_API_KEY)
|
||||
) {
|
||||
throw new Error(
|
||||
`API key is required when using a custom base URL. ` +
|
||||
@@ -878,8 +880,19 @@ export function getAIModel(overrides?: ClientOverrides): ModelConfig {
|
||||
|
||||
case "ollama": {
|
||||
const baseURL = overrides?.baseUrl || process.env.OLLAMA_BASE_URL
|
||||
if (baseURL) {
|
||||
const customOllama = createOllama({ baseURL })
|
||||
// SECURITY: When client provides a custom base URL, only use
|
||||
// client-provided API key. Never fall back to server OLLAMA_API_KEY
|
||||
// to prevent leaking server credentials to user-controlled endpoints.
|
||||
const apiKey = overrides?.baseUrl
|
||||
? overrides?.apiKey || undefined
|
||||
: resolveApiKey(overrides, "OLLAMA_API_KEY")
|
||||
if (baseURL || apiKey) {
|
||||
const customOllama = createOllama({
|
||||
...(baseURL && { baseURL }),
|
||||
...(apiKey && {
|
||||
headers: { Authorization: `Bearer ${apiKey}` },
|
||||
}),
|
||||
})
|
||||
model = customOllama(modelId)
|
||||
} else {
|
||||
model = ollama(modelId)
|
||||
|
||||
@@ -104,7 +104,7 @@ export const PROVIDER_INFO: Record<
|
||||
bedrock: { label: "Amazon Bedrock" },
|
||||
ollama: {
|
||||
label: "Ollama",
|
||||
defaultBaseUrl: "http://localhost:11434",
|
||||
defaultBaseUrl: "https://ollama.com/api",
|
||||
},
|
||||
openrouter: {
|
||||
label: "OpenRouter",
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import { describe, expect, it } from "vitest"
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
|
||||
import {
|
||||
getAIModel,
|
||||
resolveBaseURL,
|
||||
supportsImageInput,
|
||||
supportsPromptCaching,
|
||||
@@ -189,3 +190,120 @@ describe("supportsImageInput", () => {
|
||||
expect(supportsImageInput("gemini-pro")).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
vi.mock("ollama-ai-provider-v2", () => {
|
||||
const mockModel = { modelId: "test-model" }
|
||||
const mockProviderFn = vi.fn(() => mockModel)
|
||||
const mockCreateOllama = vi.fn(() => mockProviderFn)
|
||||
const mockOllama = vi.fn(() => mockModel)
|
||||
return { createOllama: mockCreateOllama, ollama: mockOllama }
|
||||
})
|
||||
|
||||
describe("Ollama API key security", () => {
|
||||
let createOllamaMock: ReturnType<typeof vi.fn>
|
||||
const savedEnv: Record<string, string | undefined> = {}
|
||||
|
||||
beforeEach(async () => {
|
||||
savedEnv.OLLAMA_API_KEY = process.env.OLLAMA_API_KEY
|
||||
savedEnv.OLLAMA_BASE_URL = process.env.OLLAMA_BASE_URL
|
||||
delete process.env.OLLAMA_BASE_URL
|
||||
|
||||
const mod = await import("ollama-ai-provider-v2")
|
||||
createOllamaMock = mod.createOllama as ReturnType<typeof vi.fn>
|
||||
createOllamaMock.mockClear()
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
process.env.OLLAMA_API_KEY = savedEnv.OLLAMA_API_KEY
|
||||
process.env.OLLAMA_BASE_URL = savedEnv.OLLAMA_BASE_URL
|
||||
})
|
||||
|
||||
it("applies server OLLAMA_API_KEY when no client baseUrl is provided", () => {
|
||||
process.env.OLLAMA_API_KEY = "server-secret-key"
|
||||
|
||||
getAIModel({ provider: "ollama", modelId: "llama2" })
|
||||
|
||||
expect(createOllamaMock).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
headers: { Authorization: "Bearer server-secret-key" },
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("does NOT leak server OLLAMA_API_KEY when client provides a custom baseUrl", () => {
|
||||
process.env.OLLAMA_API_KEY = "server-secret-key"
|
||||
|
||||
// When server has OLLAMA_API_KEY, the SSRF guard rejects
|
||||
// client-provided baseUrl without an apiKey outright
|
||||
expect(() =>
|
||||
getAIModel({
|
||||
provider: "ollama",
|
||||
baseUrl: "https://evil-server.com",
|
||||
modelId: "llama2",
|
||||
}),
|
||||
).toThrow("API key is required")
|
||||
})
|
||||
|
||||
it("uses client API key when client provides both baseUrl and apiKey", () => {
|
||||
process.env.OLLAMA_API_KEY = "server-secret-key"
|
||||
|
||||
getAIModel({
|
||||
provider: "ollama",
|
||||
baseUrl: "https://my-ollama.com",
|
||||
apiKey: "client-key",
|
||||
modelId: "llama2",
|
||||
})
|
||||
|
||||
expect(createOllamaMock).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
baseURL: "https://my-ollama.com",
|
||||
headers: { Authorization: "Bearer client-key" },
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("applies both server OLLAMA_BASE_URL and OLLAMA_API_KEY when no client overrides", () => {
|
||||
process.env.OLLAMA_BASE_URL = "https://cloud.ollama.com"
|
||||
process.env.OLLAMA_API_KEY = "server-key"
|
||||
|
||||
getAIModel({ provider: "ollama", modelId: "llama2" })
|
||||
|
||||
expect(createOllamaMock).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
baseURL: "https://cloud.ollama.com",
|
||||
headers: { Authorization: "Bearer server-key" },
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("works when OLLAMA_API_KEY is set but OLLAMA_BASE_URL is not", () => {
|
||||
process.env.OLLAMA_API_KEY = "server-key"
|
||||
delete process.env.OLLAMA_BASE_URL
|
||||
|
||||
getAIModel({ provider: "ollama", modelId: "llama2" })
|
||||
|
||||
expect(createOllamaMock).toHaveBeenCalledTimes(1)
|
||||
const callArgs = createOllamaMock.mock.calls[0][0]
|
||||
expect(callArgs).not.toHaveProperty("baseURL")
|
||||
expect(callArgs).toEqual(
|
||||
expect.objectContaining({
|
||||
headers: { Authorization: "Bearer server-key" },
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("allows client custom baseUrl without apiKey when no server OLLAMA_API_KEY", () => {
|
||||
delete process.env.OLLAMA_API_KEY
|
||||
|
||||
getAIModel({
|
||||
provider: "ollama",
|
||||
baseUrl: "https://my-ollama.com",
|
||||
modelId: "llama2",
|
||||
})
|
||||
|
||||
expect(createOllamaMock).toHaveBeenCalledTimes(1)
|
||||
const callArgs = createOllamaMock.mock.calls[0][0]
|
||||
expect(callArgs.baseURL).toBe("https://my-ollama.com")
|
||||
expect(callArgs).not.toHaveProperty("headers")
|
||||
})
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user