diff --git a/.gitignore b/.gitignore index c2658d7..8d2f53b 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ node_modules/ +app.yaml diff --git a/README.md b/README.md index 74afd12..e03f086 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -[English](README_EN.md) | 中文 +[English](README_EN.md) | [한국어](README_KR.md) | 中文 # kiro-proxy @@ -25,6 +25,57 @@ npx kiro-proxy | `PORT` | `3456` | 监听端口 | | `PROXY_API_KEY` | 无 | 设置后所有请求需携带此 key 进行鉴权,未设置则不校验 | | `HTTPS_PROXY` | 无 | HTTP/HTTPS 代理地址,如 `http://127.0.0.1:7890` | +| `MULTI_USER` | `false` | `true` 启用多用户模式(等同 `--multi-user` 参数) | +| `TOKEN_DB_PATH` | `~/.kiro-proxy/tokens.db` | 多用户模式下的 Token 数据库路径 | +| `DATABRICKS_APP_PORT` | - | 设置后优先于 PORT | + +## 多用户模式 + +支持多客户端各自使用自己的 Kiro token 同时使用代理。服务器将 token 缓存到 SQLite DB,过期时自动刷新。 + +```bash +# CLI 参数 +node server.js --multi-user + +# 或环境变量 +MULTI_USER=true node server.js +``` + +客户端需传递以下请求头: + +| 请求头 | 必需 | 说明 | +|--------|------|------| +| `X-Kiro-Access-Token` | 是* | 当前 access token | +| `X-Kiro-Refresh-Token` | 是* | refresh token(服务器用来自动续期) | +| `X-Kiro-Auth-Method` | 否 | `social` 或 `IdC` | +| `X-Kiro-Profile-Arn` | 否 | profile ARN | +| `X-Kiro-Region` | 否 | AWS region | +| `X-Kiro-Provider` | 否 | provider 类型 | + +\* 至少提供其中一个。建议同时传两个以启用自动续期。 + +提取 token: + +```bash +# 使用内置脚本 +./scripts/extract-token.sh headers # 输出 -H 参数 +./scripts/extract-token.sh env # 输出 export 语句 +./scripts/extract-token.sh curl # 输出单行 curl headers + +# 配合 curl 使用 +eval curl http://localhost:3456/v1/messages \ + -H "Content-Type: application/json" \ + $(./scripts/extract-token.sh curl) \ + -d '{"model": "claude-sonnet-4.6", "max_tokens": 1024, "messages": [{"role": "user", "content": "Hello"}]}' +``` + +Token 流程: +1. 首次请求:验证 token 后存入 DB +2. 后续请求:使用 DB 中的缓存 token +3. 过期时:服务器自动刷新并更新 DB +4. refresh token 本身过期:返回 401,客户端需提交新 token + +不使用 `--multi-user` 时行为与原版完全一致(读取本地 token 文件)。 ## API diff --git a/README_EN.md b/README_EN.md index 52acc9e..064edd9 100644 --- a/README_EN.md +++ b/README_EN.md @@ -1,4 +1,4 @@ -English | [中文](README.md) +[한국어](README_KR.md) | [中文](README.md) | English # kiro-proxy @@ -25,6 +25,57 @@ Server listens on `http://localhost:3456` by default. | `PORT` | `3456` | Listen port | | `PROXY_API_KEY` | None | When set, all requests must include this key for authentication. No validation when unset | | `HTTPS_PROXY` | None | HTTP/HTTPS proxy URL, e.g. `http://127.0.0.1:7890` | +| `MULTI_USER` | `false` | `true` enables multi-user mode (same as `--multi-user` flag) | +| `TOKEN_DB_PATH` | `~/.kiro-proxy/tokens.db` | Token database path for multi-user mode | +| `DATABRICKS_APP_PORT` | - | Overrides PORT when set | + +## Multi-user Mode + +Allows multiple clients to use the proxy simultaneously with their own Kiro tokens. The server caches tokens in a SQLite DB and auto-refreshes on expiry. + +```bash +# CLI flag +node server.js --multi-user + +# Or environment variable +MULTI_USER=true node server.js +``` + +Clients must include the following headers: + +| Header | Required | Description | +|--------|----------|-------------| +| `X-Kiro-Access-Token` | Yes* | Current access token | +| `X-Kiro-Refresh-Token` | Yes* | Refresh token (used by server to auto-renew) | +| `X-Kiro-Auth-Method` | No | `social` or `IdC` | +| `X-Kiro-Profile-Arn` | No | Profile ARN | +| `X-Kiro-Region` | No | AWS region | +| `X-Kiro-Provider` | No | Provider type | + +\* At least one is required. Provide both to enable auto-renewal. + +Extracting tokens: + +```bash +# Using the built-in script +./scripts/extract-token.sh headers # Output -H flags +./scripts/extract-token.sh env # Output export statements +./scripts/extract-token.sh curl # Output single-line curl headers + +# Use with curl +eval curl http://localhost:3456/v1/messages \ + -H "Content-Type: application/json" \ + $(./scripts/extract-token.sh curl) \ + -d '{"model": "claude-sonnet-4.6", "max_tokens": 1024, "messages": [{"role": "user", "content": "Hello"}]}' +``` + +Token flow: +1. First request: validate token and store in DB +2. Subsequent requests: use cached token from DB +3. On expiry: server auto-refreshes and updates DB +4. If refresh token itself expires: returns 401, client must submit a fresh token + +Without `--multi-user`, behaves exactly as before (reads local token file). ## API diff --git a/README_KR.md b/README_KR.md new file mode 100644 index 0000000..49669e6 --- /dev/null +++ b/README_KR.md @@ -0,0 +1,180 @@ +[English](README_EN.md) | [中文](README.md) | 한국어 + +# kiro-proxy + +[Kiro](https://kiro.dev) 구독에 포함된 Claude 모델을 Claude Code에서 사용할 수 있게 해주는 프록시. + +Kiro 인증 토큰을 읽어서 Amazon Q Developer로 요청을 프록시하고, OpenAI 및 Anthropic 호환 API 엔드포인트를 제공합니다. + +## 전제조건 + +Kiro를 설치하고 로그인해서 `~/.aws/sso/cache/kiro-auth-token.json`이 존재하고 유효해야 합니다. + +## 빠른 시작 + +```bash +npx kiro-proxy +``` + +서버 기본 포트: `http://localhost:3456` + +## 설정 + +| 환경변수 | 기본값 | 설명 | +|----------|--------|------| +| `PORT` | `3456` | 수신 포트 | +| `PROXY_API_KEY` | 없음 | 설정 시 모든 요청에 이 키로 인증 필요. 미설정 시 검증 안 함 | +| `HTTPS_PROXY` | 없음 | HTTP/HTTPS 프록시 주소, 예: `http://127.0.0.1:7890` | +| `MULTI_USER` | `false` | `true`이면 멀티유저 모드 (`--multi-user` 플래그와 동일) | +| `TOKEN_DB_PATH` | `~/.kiro-proxy/tokens.db` | 멀티유저 모드 토큰 DB 경로 | +| `DATABRICKS_APP_PORT` | - | 설정 시 PORT보다 우선 | + +## 멀티유저 모드 + +여러 클라이언트가 각자의 Kiro 토큰으로 프록시를 동시에 사용 가능. 서버가 토큰을 SQLite DB에 캐시하고 만료 시 자동 refresh. + +```bash +# CLI 플래그 +node server.js --multi-user + +# 또는 환경변수 +MULTI_USER=true node server.js +``` + +클라이언트는 다음 헤더를 포함해야 합니다: + +| 헤더 | 필수 | 설명 | +|------|------|------| +| `X-Kiro-Access-Token` | 예* | 현재 access token | +| `X-Kiro-Refresh-Token` | 예* | refresh token (서버가 자동 갱신에 사용) | +| `X-Kiro-Auth-Method` | 아니오 | `social` 또는 `IdC` | +| `X-Kiro-Profile-Arn` | 아니오 | profile ARN | +| `X-Kiro-Region` | 아니오 | AWS region | +| `X-Kiro-Provider` | 아니오 | provider 타입 | + +\* 둘 중 하나는 필수. 자동 갱신을 위해 둘 다 보내는 것을 권장. + +토큰 추출: + +```bash +# 내장 스크립트 사용 +./scripts/extract-token.sh headers # -H 플래그 출력 +./scripts/extract-token.sh env # export 구문 출력 +./scripts/extract-token.sh curl # 한 줄 curl headers 출력 + +# curl과 함께 사용 +eval curl http://localhost:3456/v1/messages \ + -H "Content-Type: application/json" \ + $(./scripts/extract-token.sh curl) \ + -d '{"model": "claude-sonnet-4.6", "max_tokens": 1024, "messages": [{"role": "user", "content": "Hello"}]}' +``` + +토큰 흐름: +1. 첫 요청: 토큰 검증 후 DB에 저장 +2. 이후 요청: DB 캐시 사용 +3. 만료 시: 서버가 자동 refresh → DB 갱신 +4. refresh token 자체 만료: 401 반환 → 클라이언트가 새 토큰 제출 + +`--multi-user` 없이 실행하면 기존과 완전히 동일하게 동작 (로컬 토큰 파일 읽기). + +## API + +### GET /v1/models — 사용 가능한 모델 조회 + +```bash +curl http://localhost:3456/v1/models +``` + +### POST /v1/messages — Anthropic 호환 + +```bash +# 비스트리밍 +curl http://localhost:3456/v1/messages \ + -H "Content-Type: application/json" \ + -H "x-api-key: any" \ + -d '{"model": "claude-sonnet-4.6", "max_tokens": 1024, "messages": [{"role": "user", "content": "Hello"}]}' + +# 스트리밍 +curl http://localhost:3456/v1/messages \ + -H "Content-Type: application/json" \ + -H "x-api-key: any" \ + -d '{"model": "claude-sonnet-4.6", "max_tokens": 1024, "messages": [{"role": "user", "content": "Hello"}], "stream": true}' +``` + +### POST /v1/chat/completions — OpenAI 호환 + +```bash +# 비스트리밍 +curl http://localhost:3456/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{"model": "claude-sonnet-4.6", "messages": [{"role": "user", "content": "Hello"}]}' + +# 스트리밍 +curl http://localhost:3456/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{"model": "claude-sonnet-4.6", "messages": [{"role": "user", "content": "Hello"}], "stream": true}' +``` + +### GET /health + +토큰 상태 및 만료 시간 확인. + +### GET /credits + +크레딧 사용량 통계. `period` 파라미터 지원: + +```bash +# 오늘 사용량 (기본) +curl http://localhost:3456/credits + +# 최근 7일 +curl http://localhost:3456/credits?period=7d + +# 최근 30일 +curl http://localhost:3456/credits?period=30d + +# 전체 +curl http://localhost:3456/credits?period=all +``` + +## Claude Code 연동 + +Claude Code는 기본적으로 Anthropic 공식 model ID를 사용합니다. 환경변수로 Q Developer model ID에 매핑해야 합니다. + +`~/.claude/settings.json`에 추가: + +```json +{ + "env": { + "ANTHROPIC_AUTH_TOKEN": "any", + "ANTHROPIC_BASE_URL": "http://localhost:3456", + "ANTHROPIC_DEFAULT_SONNET_MODEL": "claude-sonnet-4.6", + "ANTHROPIC_DEFAULT_OPUS_MODEL": "claude-opus-4.6", + "ANTHROPIC_DEFAULT_HAIKU_MODEL": "claude-haiku-4.5" + }, + "model": "sonnet" +} +``` + +`model` 옵션: `sonnet`, `opus`, `haiku`. `[1m]` 접미사로 1M 컨텍스트 윈도우 활성화 (예: `"opus[1m]"`). + +> 주의: `ANTHROPIC_MODEL` 환경변수를 설정하지 마세요. `model` 필드를 덮어써서 컨텍스트 윈도우 설정이 무효화됩니다. + +## 프록시 설정 + +2026년 5월 1일부터 Kiro의 Claude 모델은 중국 대륙 및 홍콩/마카오/대만에서 사용할 수 없습니다. `Invalid model` 에러가 발생하면 프록시를 설정하세요. + +> 주의: 프록시 노드는 다른 지역(싱가포르, 태국, 한국 등)을 선택해야 합니다. + +환경변수로 HTTP 프록시 설정: + +```bash +# 프록시 설정 후 시작 +HTTPS_PROXY=http://127.0.0.1:7890 npx kiro-proxy +``` + +지원 환경변수: `HTTPS_PROXY`, `https_proxy`, `HTTP_PROXY`, `http_proxy` (왼쪽부터 우선순위). + +## 관련 프로젝트 + +- [kiro-web-search](https://github.com/Colin3191/kiro-web-search) — Kiro 내장 웹 검색을 MCP server로 래핑, Claude Code 등에서 사용 가능 diff --git a/package.json b/package.json index 52d6fea..2ab8266 100644 --- a/package.json +++ b/package.json @@ -1,11 +1,14 @@ { - "name": "@colin3191/kiro-proxy", - "version": "0.2.4", - "description": "Kiro API proxy with OpenAI and Anthropic compatible endpoints", + "name": "@leecoder/kiro-proxy", + "version": "0.3.0", + "description": "Kiro API proxy with OpenAI and Anthropic compatible endpoints — multi-user edition", "type": "module", "bin": { "kiro-proxy": "./server.js" }, + "scripts": { + "start": "node server.js" + }, "engines": { "node": ">=18" }, @@ -14,16 +17,18 @@ "q-client.js", "token-reader.js", "token-counter.js", + "token-store.js", "usage-tracker.js", "logger.js", "proxy-config.js" ], "repository": { "type": "git", - "url": "https://github.com/Colin3191/kiro-proxy.git" + "url": "https://github.com/leecoder/kiro-proxy.git" }, "dependencies": { "@aws/codewhisperer-streaming-client": "^1.0.34", + "better-sqlite3": "^11.0.0", "express": "^4.21.0", "https-proxy-agent": "^7.0.0", "undici": "^6.19.0" diff --git a/scripts/extract-token.sh b/scripts/extract-token.sh new file mode 100755 index 0000000..c517181 --- /dev/null +++ b/scripts/extract-token.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash +set -euo pipefail + +TOKEN_FILE="${KIRO_TOKEN_FILE:-$HOME/.aws/sso/cache/kiro-auth-token.json}" + +if [ ! -f "$TOKEN_FILE" ]; then + echo "Error: $TOKEN_FILE not found. Login to Kiro first." >&2 + exit 1 +fi + +ACCESS_TOKEN=$(python3 -c "import sys,json;print(json.load(sys.stdin)['accessToken'])" < "$TOKEN_FILE") +REFRESH_TOKEN=$(python3 -c "import sys,json;print(json.load(sys.stdin).get('refreshToken',''))" < "$TOKEN_FILE") +AUTH_METHOD=$(python3 -c "import sys,json;print(json.load(sys.stdin).get('authMethod',''))" < "$TOKEN_FILE") +PROFILE_ARN=$(python3 -c "import sys,json;print(json.load(sys.stdin).get('profileArn',''))" < "$TOKEN_FILE") +REGION=$(python3 -c "import sys,json;print(json.load(sys.stdin).get('region',''))" < "$TOKEN_FILE") + +case "${1:-headers}" in + headers) + echo "-H \"X-Kiro-Access-Token: $ACCESS_TOKEN\"" + [ -n "$REFRESH_TOKEN" ] && echo "-H \"X-Kiro-Refresh-Token: $REFRESH_TOKEN\"" + [ -n "$AUTH_METHOD" ] && echo "-H \"X-Kiro-Auth-Method: $AUTH_METHOD\"" + [ -n "$PROFILE_ARN" ] && echo "-H \"X-Kiro-Profile-Arn: $PROFILE_ARN\"" + [ -n "$REGION" ] && echo "-H \"X-Kiro-Region: $REGION\"" + ;; + env) + echo "export X_KIRO_ACCESS_TOKEN=\"$ACCESS_TOKEN\"" + [ -n "$REFRESH_TOKEN" ] && echo "export X_KIRO_REFRESH_TOKEN=\"$REFRESH_TOKEN\"" + [ -n "$AUTH_METHOD" ] && echo "export X_KIRO_AUTH_METHOD=\"$AUTH_METHOD\"" + [ -n "$PROFILE_ARN" ] && echo "export X_KIRO_PROFILE_ARN=\"$PROFILE_ARN\"" + [ -n "$REGION" ] && echo "export X_KIRO_REGION=\"$REGION\"" + ;; + curl) + HEADERS="-H \"X-Kiro-Access-Token: $ACCESS_TOKEN\"" + [ -n "$REFRESH_TOKEN" ] && HEADERS="$HEADERS -H \"X-Kiro-Refresh-Token: $REFRESH_TOKEN\"" + [ -n "$AUTH_METHOD" ] && HEADERS="$HEADERS -H \"X-Kiro-Auth-Method: $AUTH_METHOD\"" + [ -n "$PROFILE_ARN" ] && HEADERS="$HEADERS -H \"X-Kiro-Profile-Arn: $PROFILE_ARN\"" + [ -n "$REGION" ] && HEADERS="$HEADERS -H \"X-Kiro-Region: $REGION\"" + echo "$HEADERS" + ;; + *) + echo "Usage: $(basename "$0") [headers|env|curl]" >&2 + echo " headers — print -H flags (default)" >&2 + echo " env — print export statements (eval-able)" >&2 + echo " curl — print single-line curl headers" >&2 + exit 1 + ;; +esac diff --git a/server.js b/server.js index 5e4a315..3fed341 100755 --- a/server.js +++ b/server.js @@ -1,20 +1,26 @@ #!/usr/bin/env node import express from 'express'; import crypto from 'crypto'; -import { getAccessToken } from './token-reader.js'; +import { getAccessToken, getAccessTokenFromRequest } from './token-reader.js'; import { createClient, chat, chatStream, listAvailableModels } from './q-client.js'; import { c, log, tagLog, logSummary, reqId, tagError } from './logger.js'; import { countMessages, countContent } from './token-counter.js'; import { recordUsage, queryUsage, todaySummary } from './usage-tracker.js'; import { initGlobalProxy } from './proxy-config.js'; +import { initTokenStore } from './token-store.js'; const proxyUrl = initGlobalProxy(); if (proxyUrl) tagLog('proxy', `Using proxy: ${proxyUrl}`); +const args = process.argv.slice(2); +const MULTI_USER = args.includes('--multi-user') || process.env.MULTI_USER === '1' || process.env.MULTI_USER === 'true'; +if (MULTI_USER) initTokenStore(); + const app = express(); app.use(express.json({ limit: '10mb' })); -const PORT = process.env.PORT || 3456; +const PORT = process.env.DATABRICKS_APP_PORT || process.env.PORT || 3456; +const HOST = '0.0.0.0'; const PROXY_API_KEY = process.env.PROXY_API_KEY; function authMiddleware(req, res, next) { @@ -29,9 +35,47 @@ app.use(authMiddleware); let cachedClient = null; let cachedToken = null; +const clientCache = new Map(); + +function extractKiroHeaders(req) { + return { + accessToken: req.headers['x-kiro-access-token'], + refreshToken: req.headers['x-kiro-refresh-token'], + authMethod: req.headers['x-kiro-auth-method'], + profileArn: req.headers['x-kiro-profile-arn'], + region: req.headers['x-kiro-region'], + provider: req.headers['x-kiro-provider'], + clientIdHash: req.headers['x-kiro-client-id-hash'], + }; +} + +async function getClient(req) { + let tokenData; + + if (MULTI_USER) { + const headers = extractKiroHeaders(req); + if (!headers.accessToken && !headers.refreshToken) { + throw new Error('X-Kiro-Access-Token or X-Kiro-Refresh-Token header required'); + } + tokenData = await getAccessTokenFromRequest(headers); + + if (clientCache.has(tokenData.accessToken)) { + return { client: clientCache.get(tokenData.accessToken), tokenData }; + } + const client = createClient(tokenData.accessToken, { + authMethod: tokenData.authMethod, + profileArn: tokenData.profileArn, + provider: tokenData.provider, + }); + clientCache.set(tokenData.accessToken, client); + if (clientCache.size > 50) { + const oldest = clientCache.keys().next().value; + clientCache.delete(oldest); + } + return { client, tokenData }; + } -async function getClient() { - const tokenData = await getAccessToken(); + tokenData = await getAccessToken(); if (!cachedClient || cachedToken !== tokenData.accessToken) { cachedClient = createClient(tokenData.accessToken, { authMethod: tokenData.authMethod, @@ -57,7 +101,7 @@ app.post('/v1/messages', async (req, res) => { return res.status(400).json({ type: 'error', error: { type: 'invalid_request_error', message: 'messages required' } }); } - const { client, tokenData } = await getClient(); + const { client, tokenData } = await getClient(req); const opts = { messages, system, tools, profileArn: tokenData.profileArn, modelId: model }; const rid = reqId(); const start = Date.now(); @@ -226,7 +270,7 @@ app.post('/v1/messages', async (req, res) => { } } catch (err) { tagError('anthropic', err.message || err); - const status = err.message?.includes('expired') ? 401 : 500; + const status = err.message?.includes('expired') ? 401 : err.message?.includes('X-Kiro-') ? 401 : 500; res.status(status).json({ type: 'error', error: { type: status === 401 ? 'authentication_error' : 'api_error', message: err.message } }); } }); @@ -247,7 +291,7 @@ app.post('/v1/chat/completions', async (req, res) => { messages.push({ role: m.role, content: m.content }); } - const { client, tokenData } = await getClient(); + const { client, tokenData } = await getClient(req); const opts = { messages, system, profileArn: tokenData.profileArn, modelId: model }; const rid = reqId(); const start = Date.now(); @@ -317,9 +361,16 @@ app.post('/v1/chat/completions', async (req, res) => { // ============================================================ // GET /v1/models // ============================================================ -app.get('/v1/models', async (_req, res) => { +app.get('/v1/models', async (req, res) => { try { - const tokenData = await getAccessToken(); + let tokenData; + if (MULTI_USER) { + const headers = extractKiroHeaders(req); + if (!headers.accessToken && !headers.refreshToken) return res.status(401).json({ error: { message: 'X-Kiro-Access-Token or X-Kiro-Refresh-Token header required' } }); + tokenData = await getAccessTokenFromRequest(headers); + } else { + tokenData = await getAccessToken(); + } const { models, defaultModel } = await listAvailableModels(tokenData.accessToken, { profileArn: tokenData.profileArn, authMethod: tokenData.authMethod, provider: tokenData.provider, }); @@ -336,9 +387,16 @@ app.get('/v1/models', async (_req, res) => { } }); -app.get('/q/models', async (_req, res) => { +app.get('/q/models', async (req, res) => { try { - const tokenData = await getAccessToken(); + let tokenData; + if (MULTI_USER) { + const headers = extractKiroHeaders(req); + if (!headers.accessToken && !headers.refreshToken) return res.status(401).json({ error: { message: 'X-Kiro-Access-Token or X-Kiro-Refresh-Token header required' } }); + tokenData = await getAccessTokenFromRequest(headers); + } else { + tokenData = await getAccessToken(); + } const result = await listAvailableModels(tokenData.accessToken, { profileArn: tokenData.profileArn, authMethod: tokenData.authMethod, provider: tokenData.provider, }); @@ -348,11 +406,18 @@ app.get('/q/models', async (_req, res) => { } }); -app.get('/health', async (_req, res) => { +app.get('/health', async (req, res) => { try { - const tokenData = await getAccessToken(); + let tokenData; + if (MULTI_USER) { + const headers = extractKiroHeaders(req); + if (!headers.accessToken && !headers.refreshToken) return res.json({ status: 'ok', mode: 'multi-user', message: 'Send X-Kiro-Access-Token header for token health' }); + tokenData = await getAccessTokenFromRequest(headers); + } else { + tokenData = await getAccessToken(); + } const expired = tokenData.expiresAt && new Date(tokenData.expiresAt) < new Date(); - res.json({ status: expired ? 'token_expired' : 'ok', provider: tokenData.provider || 'unknown', expiresAt: tokenData.expiresAt }); + res.json({ status: expired ? 'token_expired' : 'ok', mode: MULTI_USER ? 'multi-user' : 'local', provider: tokenData.provider || 'unknown', expiresAt: tokenData.expiresAt }); } catch (err) { res.status(503).json({ status: 'error', message: err.message }); } @@ -366,18 +431,21 @@ app.get('/credits', (_req, res) => { res.json(queryUsage(period)); }); -app.listen(PORT, async () => { - console.log(`${c.cyan}Kiro Proxy${c.reset} running on ${c.green}http://localhost:${PORT}${c.reset}`); - console.log(` ${c.gray}Anthropic:${c.reset} http://localhost:${PORT}/v1/messages`); - console.log(` ${c.gray}OpenAI: ${c.reset} http://localhost:${PORT}/v1/chat/completions`); - console.log(` ${c.gray}Models: ${c.reset} http://localhost:${PORT}/v1/models`); - console.log(` ${c.gray}Credits: ${c.reset} http://localhost:${PORT}/credits`); +app.listen(PORT, HOST, async () => { + const modeLabel = MULTI_USER ? `${c.magenta}multi-user${c.reset} (token DB)` : `${c.green}local${c.reset}`; + console.log(`${c.cyan}Kiro Proxy${c.reset} running on ${c.green}http://${HOST}:${PORT}${c.reset} [${modeLabel}]`); + console.log(` ${c.gray}Anthropic:${c.reset} http://${HOST}:${PORT}/v1/messages`); + console.log(` ${c.gray}OpenAI: ${c.reset} http://${HOST}:${PORT}/v1/chat/completions`); + console.log(` ${c.gray}Models: ${c.reset} http://${HOST}:${PORT}/v1/models`); + console.log(` ${c.gray}Credits: ${c.reset} http://${HOST}:${PORT}/credits`); console.log(` ${c.gray}Auth: ${c.reset} ${PROXY_API_KEY ? `${c.green}enabled${c.reset} (PROXY_API_KEY)` : `${c.yellow}disabled${c.reset} (no PROXY_API_KEY set)`}`); - try { - const t = await getAccessToken(); - console.log(` ${c.gray}Provider: ${c.yellow}${t.provider || 'unknown'}${c.reset}, Expires: ${c.dim}${t.expiresAt || 'unknown'}${c.reset}`); - } catch (err) { - console.warn(` ${c.yellow}Warning:${c.reset} ${err.message}`); + if (!MULTI_USER) { + try { + const t = await getAccessToken(); + console.log(` ${c.gray}Provider: ${c.yellow}${t.provider || 'unknown'}${c.reset}, Expires: ${c.dim}${t.expiresAt || 'unknown'}${c.reset}`); + } catch (err) { + console.warn(` ${c.yellow}Warning:${c.reset} ${err.message}`); + } } }); diff --git a/token-reader.js b/token-reader.js index 0b86b44..e921455 100644 --- a/token-reader.js +++ b/token-reader.js @@ -2,6 +2,7 @@ import fs from 'fs'; import path from 'path'; import os from 'os'; import { tagLog, tagWarn, tagError } from './logger.js'; +import { hashToken, getStoredToken, upsertToken, deleteToken } from './token-store.js'; const SSO_CACHE_DIR = path.join(os.homedir(), '.aws', 'sso', 'cache'); const KIRO_TOKEN_FILE = 'kiro-auth-token.json'; @@ -244,3 +245,71 @@ function enrichWithProfile(tokenData) { } return tokenData; } + +// ============================================================ +// Remote mode: 클라이언트가 보낸 kiro token으로 동작 +// ============================================================ + +const refreshLocks = new Map(); + +/** + * @param {object} headers — { accessToken, refreshToken, ?authMethod, ?profileArn, ?region, ?provider } + */ +export async function getAccessTokenFromRequest(headers) { + const { accessToken, refreshToken: clientRefreshToken, authMethod, profileArn, region, provider, clientIdHash } = headers; + + if (!accessToken && !clientRefreshToken) { + throw new Error('X-Kiro-Access-Token or X-Kiro-Refresh-Token required'); + } + + const keySource = clientRefreshToken || accessToken; + const keyHash = hashToken(keySource); + + const stored = getStoredToken(keyHash); + + if (stored && !isTokenExpired(stored)) { + return stored; + } + + if (refreshLocks.has(keyHash)) { + tagLog('token', `[multi] Waiting for ongoing refresh (${keyHash.slice(0, 8)}...)`); + return refreshLocks.get(keyHash); + } + + const promise = (async () => { + try { + const tokenToRefresh = stored || { accessToken, refreshToken: clientRefreshToken, authMethod, profileArn, region, provider, clientIdHash }; + + if (!tokenToRefresh.refreshToken) { + if (tokenToRefresh.accessToken) { + upsertToken(keyHash, tokenToRefresh); + return tokenToRefresh; + } + throw new Error('Token expired and no refreshToken available. Client must re-login in Kiro.'); + } + + // 첫 요청: expiresAt 없으면 accessToken이 유효하다고 가정하고 저장만 + if (!stored && accessToken) { + const expiresAt = new Date(Date.now() + 3600 * 1000).toISOString(); + const tokenToStore = { accessToken, refreshToken: clientRefreshToken, authMethod, profileArn, region, provider, clientIdHash, expiresAt }; + upsertToken(keyHash, tokenToStore); + return tokenToStore; + } + + tagLog('token', `[multi] Refreshing token (${keyHash.slice(0, 8)}...)`); + const refreshed = await refreshToken(tokenToRefresh); + upsertToken(keyHash, refreshed); + tagLog('token', `[multi] Token refreshed, new expiry: ${refreshed.expiresAt}`); + return refreshed; + } catch (err) { + tagError('token', `[multi] Refresh failed (${keyHash.slice(0, 8)}...):`, err.message); + deleteToken(keyHash); + throw err; + } finally { + refreshLocks.delete(keyHash); + } + })(); + + refreshLocks.set(keyHash, promise); + return promise; +} diff --git a/token-store.js b/token-store.js new file mode 100644 index 0000000..bd89b9f --- /dev/null +++ b/token-store.js @@ -0,0 +1,107 @@ +import crypto from 'crypto'; +import fs from 'fs'; +import path from 'path'; +import os from 'os'; +import Database from 'better-sqlite3'; +import { tagLog } from './logger.js'; + +// DB 경로: 환경변수 또는 기본값 +const DB_PATH = process.env.TOKEN_DB_PATH || path.join(os.homedir(), '.kiro-proxy', 'tokens.db'); + +let db; + +function getDb() { + if (db) return db; + const dir = path.dirname(DB_PATH); + fs.mkdirSync(dir, { recursive: true }); + db = new Database(DB_PATH); + db.pragma('journal_mode = WAL'); + db.exec(` + CREATE TABLE IF NOT EXISTS tokens ( + key_hash TEXT PRIMARY KEY, + access_token TEXT NOT NULL, + refresh_token TEXT, + expires_at TEXT NOT NULL, + auth_method TEXT, + profile_arn TEXT, + region TEXT, + provider TEXT, + client_id_hash TEXT, + updated_at INTEGER NOT NULL + ) + `); + return db; +} + +export function hashToken(rawToken) { + return crypto.createHash('sha256').update(rawToken).digest('hex'); +} + +export function getStoredToken(keyHash) { + const row = getDb().prepare('SELECT * FROM tokens WHERE key_hash = ?').get(keyHash); + if (!row) return null; + return { + accessToken: row.access_token, + refreshToken: row.refresh_token, + expiresAt: row.expires_at, + authMethod: row.auth_method, + profileArn: row.profile_arn, + region: row.region, + provider: row.provider, + clientIdHash: row.client_id_hash, + }; +} + +export function upsertToken(keyHash, tokenData) { + const stmt = getDb().prepare(` + INSERT INTO tokens (key_hash, access_token, refresh_token, expires_at, auth_method, profile_arn, region, provider, client_id_hash, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(key_hash) DO UPDATE SET + access_token = excluded.access_token, + refresh_token = excluded.refresh_token, + expires_at = excluded.expires_at, + auth_method = excluded.auth_method, + profile_arn = excluded.profile_arn, + region = excluded.region, + provider = excluded.provider, + client_id_hash = excluded.client_id_hash, + updated_at = excluded.updated_at + `); + stmt.run( + keyHash, + tokenData.accessToken, + tokenData.refreshToken || null, + tokenData.expiresAt, + tokenData.authMethod || null, + tokenData.profileArn || null, + tokenData.region || null, + tokenData.provider || null, + tokenData.clientIdHash || null, + Date.now(), + ); +} + +export function deleteToken(keyHash) { + getDb().prepare('DELETE FROM tokens WHERE key_hash = ?').run(keyHash); +} + +/** + * 오래된 토큰 정리 — 기본 7일 미갱신 엔트리 삭제 + */ +export function purgeExpired(maxAgeMs = 7 * 24 * 60 * 60 * 1000) { + const cutoff = Date.now() - maxAgeMs; + const result = getDb().prepare('DELETE FROM tokens WHERE updated_at < ?').run(cutoff); + if (result.changes > 0) { + tagLog('token-store', `Purged ${result.changes} stale token(s)`); + } + return result.changes; +} + +/** + * 서버 시작 시 1회 purge 실행 + */ +export function initTokenStore() { + getDb(); // ensure table exists + purgeExpired(); + tagLog('token-store', `Initialized at ${DB_PATH}`); +}