返回 DeepSeek-Reasonix
cred_proxy.go
根目录 / desktop / cred_proxy.go
1 package main
2
3 import (
4 "bytes"
5 "crypto/hmac"
6 "crypto/rand"
7 "crypto/sha256"
8 "encoding/binary"
9 "encoding/hex"
10 "encoding/json"
11 "fmt"
12 "io"
13 "log"
14 "maps"
15 "net"
16 "net/http"
17 "net/http/httputil"
18 "net/url"
19 "strconv"
20 "strings"
21 "sync"
22 "time"
23
24 "reasonix/internal/config"
25 "reasonix/internal/netclient"
26 )
27
28 // Local-proxy mode tunnels model calls to this desktop, which swaps a scoped
29 // virtual token for the real provider key. The real key never leaves desktop.
30
31 // credentialProxyProviderName is the provider entry the bootstrap installs in
32 // the remote config; the serve launches with --model <name>.
33 const credentialProxyProviderName = "reasonix-desktop-proxy"
34
35 type credProxyRoute struct {
36 modelSnapshot *config.Config
37 proxy *httputil.ReverseProxy
38 model string
39 ref string
40 apiKeyEnv string
41 provider string
42 origins map[string]bool
43 scope string
44 revision string
45 active int
46 retired bool
47 extraBody map[string]any
48 host string
49 workspace string
50 holds map[string]bool
51 }
52
53 // credentialProxy is the desktop-side key holder: a loopback HTTP endpoint
54 // that authenticates requests by virtual token and forwards them to the real
55 // provider with the real key. One instance serves the whole app.
56 type credentialProxy struct {
57 mu sync.Mutex
58 updateMu sync.Mutex
59 ln net.Listener
60 server *http.Server
61 port int
62 routes map[string]*credProxyRoute
63 ownership map[string]*credentialProxyOwnership
64 modelSettingsSource func(http.ResponseWriter, *http.Request, *credProxyRoute)
65 }
66
67 func (p *credentialProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
68 // Unauthenticated liveness endpoint for the desktop's reverse-tunnel
69 // probe: the listener only exists behind the SSH reverse forward, so a
70 // 204 here proves serve → remote loopback → tunnel → desktop end to end.
71 if r.URL.Path == "/healthz" {
72 w.WriteHeader(http.StatusNoContent)
73 return
74 }
75 token := bearerToken(r.Header.Get("Authorization"))
76 p.mu.Lock()
77 route := p.routes[token]
78 if route != nil && !route.retired {
79 route.active++
80 } else {
81 route = nil
82 }
83 routeCount := len(p.routes)
84 p.mu.Unlock()
85 if route == nil {
86 log.Printf("[remote] credProxy: rejected %s %s routeCount=%d", r.Method, r.URL.Path, routeCount)
87 http.Error(w, "invalid credential proxy token", http.StatusUnauthorized)
88 return
89 }
90 defer func() {
91 p.mu.Lock()
92 route.active--
93 if route.retired && route.active == 0 && p.routes[token] == route {
94 delete(p.routes, token)
95 }
96 p.mu.Unlock()
97 }()
98 if r.URL.Path == "/model-settings-source" && p.modelSettingsSource != nil {
99 p.modelSettingsSource(w, r, route)
100 return
101 }
102 if original := r.Header.Get(netclient.ModelProxyOriginalURLHeader); original != "" {
103 u, err := url.Parse(original)
104 if err != nil || u.User != nil || u.Fragment != "" || !route.origins[u.Scheme+"://"+u.Host] {
105 http.Error(w, "invalid model credential proxy destination", http.StatusForbidden)
106 return
107 }
108 }
109 if route.model != "" && r.Body != nil && (r.Method == http.MethodPost || r.Method == http.MethodPut) {
110 const rewriteLimit = 64 << 20
111 if r.ContentLength > rewriteLimit {
112 http.Error(w, "credential proxy request body is too large", http.StatusRequestEntityTooLarge)
113 return
114 }
115 buffered, err := io.ReadAll(io.LimitReader(r.Body, rewriteLimit+1))
116 switch {
117 case err != nil:
118 _ = r.Body.Close()
119 http.Error(w, "credential proxy could not read request body", http.StatusBadRequest)
120 return
121 case int64(len(buffered)) > rewriteLimit:
122 _ = r.Body.Close()
123 http.Error(w, "credential proxy request body is too large", http.StatusRequestEntityTooLarge)
124 return
125 default:
126 _ = r.Body.Close()
127 body := rewriteJSONModel(buffered, route.model)
128 if len(route.extraBody) > 0 {
129 var payload map[string]any
130 if json.Unmarshal(body, &payload) == nil && payload != nil {
131 maps.Copy(payload, route.extraBody)
132 if encoded, err := json.Marshal(payload); err == nil {
133 body = encoded
134 }
135 }
136 }
137 r.Body = io.NopCloser(bytes.NewReader(body))
138 r.ContentLength = int64(len(body))
139 r.Header.Set("Content-Length", strconv.Itoa(len(body)))
140 }
141 }
142 route.proxy.ServeHTTP(w, r)
143 }
144
145 func rewriteJSONModel(body []byte, model string) []byte {
146 if model == "" || len(body) == 0 {
147 return body
148 }
149 var payload map[string]any
150 if err := json.Unmarshal(body, &payload); err != nil || payload == nil {
151 // Unparseable or a literal null body ("null" decodes into a nil
152 // map): assigning into nil would panic, and there is nothing to
153 // rewrite — pass the body through untouched.
154 return body
155 }
156 if current, ok := payload["model"].(string); ok && current == model {
157 return body
158 }
159 payload["model"] = model
160 out, err := json.Marshal(payload)
161 if err != nil {
162 return body
163 }
164 return out
165 }
166
167 func (p *credentialProxy) setRoute(token, ref string, upstream *url.URL, apiKey, model, kind string) {
168 p.updateMu.Lock()
169 defer p.updateMu.Unlock()
170 p.setRouteLocked(token, ref, proxyUpstream{url: upstream, apiKey: apiKey, model: model, kind: kind})
171 }
172
173 func (p *credentialProxy) resolveAndSetRoute(token, ref string, resolve func() (proxyUpstream, error)) (proxyUpstream, error) {
174 p.updateMu.Lock()
175 defer p.updateMu.Unlock()
176 up, err := resolve()
177 if err != nil {
178 return proxyUpstream{}, err
179 }
180 if err := p.validateModelSettingsOfferCapacity(up); err != nil {
181 return proxyUpstream{}, err
182 }
183 p.setRouteLocked(token, ref, up)
184 return up, nil
185 }
186
187 func (p *credentialProxy) setRouteLocked(token, ref string, up proxyUpstream) {
188 if up.kind == "" {
189 up.kind = "openai"
190 }
191 proxy := &httputil.ReverseProxy{FlushInterval: -1}
192 proxy.Rewrite = func(req *httputil.ProxyRequest) {
193 req.SetURL(up.url)
194 if original := req.In.Header.Get(netclient.ModelProxyOriginalURLHeader); original != "" {
195 // ServeHTTP validated the destination against this frozen route.
196 req.Out.URL, _ = url.Parse(original)
197 } else if up.requestURL != nil {
198 req.Out.URL = new(url.URL)
199 *req.Out.URL = *up.requestURL
200 }
201 req.Out.Header.Del(netclient.ModelProxyOriginalURLHeader)
202 for _, header := range []string{"Forwarded", "X-Forwarded-For", "X-Forwarded-Host", "X-Forwarded-Proto", "X-Real-IP", "Via"} {
203 req.Out.Header.Del(header)
204 }
205 if up.kind == "anthropic" && !up.authHeader {
206 req.Out.Header.Del("Authorization")
207 req.Out.Header.Set("x-api-key", up.apiKey)
208 req.Out.Header.Set("anthropic-version", "2023-06-01")
209 } else {
210 req.Out.Header.Del("x-api-key")
211 req.Out.Header.Set("Authorization", "Bearer "+up.apiKey)
212 }
213 for name, value := range up.headers {
214 req.Out.Header.Set(name, value)
215 }
216 }
217 p.mu.Lock()
218 defer p.mu.Unlock()
219 // A token is a connection version. Re-registration must never redirect
220 // requests already accepted by a runtime holding that token.
221 if route := p.routes[token]; route != nil {
222 if up.offerID != "" {
223 if route.holds == nil {
224 route.holds = map[string]bool{}
225 }
226 route.holds[up.offerID] = true
227 route.retired = false
228 }
229 return
230 }
231 origins := map[string]bool{up.url.Scheme + "://" + up.url.Host: true}
232 if up.requestURL != nil {
233 origins[up.requestURL.Scheme+"://"+up.requestURL.Host] = true
234 }
235 p.routes[token] = &credProxyRoute{
236 modelSnapshot: up.modelSnapshot,
237 proxy: proxy, model: up.model, ref: ref,
238 apiKeyEnv: strings.TrimSpace(up.apiKeyEnv), provider: strings.TrimSpace(up.provider),
239 origins: origins,
240 scope: up.scope, revision: up.revision, extraBody: up.extraBody,
241 host: up.host, workspace: up.workspace, holds: map[string]bool{},
242 }
243 if up.offerID != "" {
244 p.routes[token].holds[up.offerID] = true
245 }
246 }
247
248 func (p *credentialProxy) close() {
249 p.mu.Lock()
250 server, listener := p.server, p.ln
251 p.server, p.ln = nil, nil
252 p.mu.Unlock()
253 if server != nil {
254 _ = server.Close()
255 }
256 if listener != nil {
257 _ = listener.Close()
258 }
259 }
260
261 func bearerToken(header string) string {
262 prefix, value, ok := strings.Cut(strings.TrimSpace(header), " ")
263 if !ok || !strings.EqualFold(prefix, "Bearer") {
264 return ""
265 }
266 return strings.TrimSpace(value)
267 }
268
269 // credentialProxyPort returns the proxy's loopback port, starting the proxy
270 // on first use.
271 func (a *App) credentialProxyPort() (int, error) {
272 a.credProxyMu.Lock()
273 defer a.credProxyMu.Unlock()
274 if a.credProxy != nil {
275 return a.credProxy.port, nil
276 }
277 ln, err := net.Listen("tcp", "127.0.0.1:0")
278 if err != nil {
279 return 0, fmt.Errorf("credential proxy: listen: %w", err)
280 }
281 p := &credentialProxy{ln: ln, port: ln.Addr().(*net.TCPAddr).Port, routes: map[string]*credProxyRoute{}}
282 p.modelSettingsSource = a.serveModelSettingsSource
283 server := &http.Server{
284 Handler: p,
285 ReadHeaderTimeout: 10 * time.Second,
286 IdleTimeout: 2 * time.Minute,
287 MaxHeaderBytes: 1 << 20,
288 }
289 p.server = server
290 a.credProxy = p
291 a.goSafe("credentialProxy", func() { _ = server.Serve(ln) })
292 return p.port, nil
293 }
294
295 func (a *App) closeCredentialProxy() {
296 a.credProxyMu.Lock()
297 defer a.credProxyMu.Unlock()
298 if a.credProxy != nil {
299 a.credProxy.close()
300 a.credProxy = nil
301 }
302 }
303
304 // credentialProxySecret loads (creating on first use) the persisted random
305 // secret every virtual token derives from. Rotating it revokes all tokens.
306 func (a *App) credentialProxySecret() (string, error) {
307 remotePrefsMu.Lock()
308 defer remotePrefsMu.Unlock()
309 p, err := updateRemotePrefsLocked(func(p *remotePrefs) (bool, error) {
310 if p.CredentialProxySecret != "" {
311 return false, nil
312 }
313 buf := make([]byte, 32)
314 if _, err := rand.Read(buf); err != nil {
315 return false, fmt.Errorf("credential proxy: generate secret: %w", err)
316 }
317 p.CredentialProxySecret = hex.EncodeToString(buf)
318 return true, nil
319 })
320 if err != nil {
321 return "", fmt.Errorf("credential proxy: persist secret: %w", err)
322 }
323 return p.CredentialProxySecret, nil
324 }
325
326 // credentialProxyModelTokenFor gives each staged model an immutable route.
327 // A controller already running with the previous virtual token therefore keeps
328 // its old upstream for the whole turn while Serve builds and publishes the new
329 // controller. This is the cross-process half of failure-atomic model switches.
330 func credentialProxyModelTokenFor(secret, hostID, workspace, modelRef string, revisions ...string) string {
331 mac := hmac.New(sha256.New, []byte(secret))
332 _, _ = mac.Write([]byte("reasonix-credential-proxy-model:v3"))
333 for _, field := range append([]string{hostID, workspace, modelRef}, revisions...) {
334 var size [8]byte
335 binary.BigEndian.PutUint64(size[:], uint64(len(field)))
336 _, _ = mac.Write(size[:])
337 _, _ = mac.Write([]byte(field))
338 }
339 return hex.EncodeToString(mac.Sum(nil))[:32]
340 }
341
342 // credentialProxyRouteInfo is everything a serve bootstrap needs to install
343 // the desktop hop on the remote: the virtual token, the model name and
344 // provider kind the remote provider entry should carry, and the proxy's
345 // loopback port.
346 type credentialProxyRouteInfo struct {
347 token string
348 model string
349 kind string
350 port int
351 revision string
352 }
353
354 // proxyUpstream is the resolved desktop-side provider a route forwards to.
355 type proxyUpstream struct {
356 modelSnapshot *config.Config
357 host, workspace, offerID string
358 apiKey string
359 url *url.URL
360 model string
361 kind string
362 apiKeyEnv string
363 provider string
364 requestURL *url.URL
365 headers map[string]string
366 extraBody map[string]any
367 authHeader bool
368 scope string
369 revision string
370 }
371
372 // resolveProxyProvider resolves a desktop model ref into the upstream the
373 // credential proxy should forward to, including the auth-header shape its
374 // provider kind expects.
375 func resolveProxyProvider(cfg *config.Config, ref string) (proxyUpstream, error) {
376 entry, ok := cfg.ResolveModel(ref)
377 if !ok {
378 return proxyUpstream{}, fmt.Errorf("credential proxy: model %q has no provider", ref)
379 }
380 apiKey := entry.APIKey()
381 if apiKey == "" {
382 return proxyUpstream{}, fmt.Errorf("credential proxy: the local provider credential is not configured")
383 }
384 base := strings.TrimSpace(entry.BaseURL)
385 if base == "" {
386 base = "https://api.openai.com"
387 }
388 upstream, err := url.Parse(strings.TrimRight(base, "/") + "/")
389 if err != nil {
390 return proxyUpstream{}, fmt.Errorf("credential proxy: provider base_url: %w", err)
391 }
392 if (upstream.Scheme != "http" && upstream.Scheme != "https") || upstream.Host == "" || upstream.User != nil || upstream.Fragment != "" {
393 return proxyUpstream{}, fmt.Errorf("credential proxy: provider base_url must be an http(s) URL without credentials or a fragment")
394 }
395 kind := strings.TrimSpace(entry.Kind)
396 if kind == "" {
397 kind = "openai"
398 }
399 var exactURL *url.URL
400 if exact := config.ProviderEffectiveRequestURL(entry); exact != "" {
401 exactURL, err = url.Parse(exact)
402 if err != nil || (exactURL.Scheme != "http" && exactURL.Scheme != "https") || exactURL.Host == "" || exactURL.User != nil || exactURL.Fragment != "" {
403 return proxyUpstream{}, fmt.Errorf("credential proxy: invalid request URL")
404 }
405 }
406 return proxyUpstream{
407 modelSnapshot: cfg,
408 apiKey: apiKey, url: upstream, model: entry.Model, kind: kind,
409 apiKeyEnv: entry.APIKeyEnv, provider: entry.Name,
410 requestURL: exactURL, headers: entry.Headers, extraBody: entry.ExtraBody, authHeader: entry.AuthHeader,
411 }, nil
412 }
413
414 // registerCredentialProxyRoute binds one workspace token to the current
415 // desktop default provider without exposing its real key to the remote.
416 func (a *App) registerCredentialProxyRoute(hostID, workspace string) (credentialProxyRouteInfo, error) {
417 cfg, err := config.Load()
418 if err != nil {
419 return credentialProxyRouteInfo{}, err
420 }
421 ref := strings.TrimSpace(cfg.DefaultModel)
422 if workspaceModel := a.desktopModelForWorkspace(hostID, workspace); workspaceModel != "" {
423 ref = workspaceModel
424 }
425 return a.applyCredentialProxyModel(hostID, workspace, ref)
426 }
427
428 // desktopModelForWorkspace deterministically selects the newest tab-owned
429 // model for a workspace; map iteration order must never choose a route.
430 func (a *App) desktopModelForWorkspace(hostID, workspace string) string {
431 a.remoteTabMu.Lock()
432 defer a.remoteTabMu.Unlock()
433 var selected string
434 var selectedSeq uint64
435 for _, tab := range a.remoteTabs {
436 if tab == nil || tab.ref.HostID != hostID || tab.ref.Workspace != workspace || strings.TrimSpace(tab.model) == "" {
437 continue
438 }
439 if tab.modelSeq >= selectedSeq {
440 selected, selectedSeq = tab.model, tab.modelSeq
441 }
442 }
443 return selected
444 }
445
446 func (a *App) applyCredentialProxyModel(hostID, workspace, ref string) (credentialProxyRouteInfo, error) {
447 cfg, err := config.LoadModelRuntimeSnapshot(".")
448 if err != nil {
449 return credentialProxyRouteInfo{}, err
450 }
451 return a.applyCredentialProxySnapshot(hostID, workspace, ref, cfg)
452 }
453
454 func (a *App) applyCredentialProxySnapshot(hostID, workspace, ref string, cfg *config.Config, generation ...string) (credentialProxyRouteInfo, error) {
455 port, err := a.credentialProxyPort()
456 if err != nil {
457 return credentialProxyRouteInfo{}, err
458 }
459 a.credProxyMu.Lock()
460 proxy := a.credProxy
461 a.credProxyMu.Unlock()
462 if proxy == nil {
463 return credentialProxyRouteInfo{}, fmt.Errorf("credential proxy: not running")
464 }
465 // Route tokens include the canonical desktop model ref. Never mutate the
466 // route held by an in-flight controller during a model switch.
467 secret, err := a.credentialProxySecret()
468 if err != nil {
469 return credentialProxyRouteInfo{}, err
470 }
471 revision := cfg.ModelRuntimeFingerprint(ref)
472 if len(generation) > 0 {
473 revision = generation[0]
474 }
475 token := credentialProxyModelTokenFor(secret, hostID, workspace, ref, revision)
476 up, err := proxy.resolveAndSetRoute(token, ref, func() (proxyUpstream, error) {
477 up, err := resolveProxyProvider(cfg, ref)
478 up.scope, up.revision = credentialProxyScope(hostID, workspace), revision
479 up.host, up.workspace = hostID, workspace
480 if len(generation) > 1 {
481 up.offerID = generation[1]
482 }
483 return up, err
484 })
485 if err != nil {
486 return credentialProxyRouteInfo{}, err
487 }
488 return credentialProxyRouteInfo{token: token, model: up.model, kind: up.kind, port: port, revision: revision}, nil
489 }
490
491 // saveProviderCredential writes only the credential store. Existing routes own
492 // their frozen upstream until the corresponding remote runtime is retired.
493 func (a *App) saveProviderCredential(apiKeyEnv, value string) (string, error) {
494 apiKeyEnv = strings.TrimSpace(apiKeyEnv)
495 value = strings.TrimSpace(value)
496 if err := upsertDotEnv(apiKeyEnv, value); err != nil {
497 return "", err
498 }
499 return providerCredentialSourceNotice(apiKeyEnv, value), nil
500 }
501
502 // credentialModeView returns the host entry's normalized credential mode for
503 // views ("" reads as "remote" — the default).
504 func credentialModeView(h config.RemoteHostEntry) string {
505 if h.CredentialProxyEnabled() {
506 return "local-proxy"
507 }
508 return "remote"
509 }
510
511 // normalizeCredentialMode validates an input credential mode.
512 func normalizeCredentialMode(mode string) string {
513 switch strings.ToLower(strings.TrimSpace(mode)) {
514 case "local-proxy":
515 return "local-proxy"
516 default:
517 return ""
518 }
519 }
520
520 lines GO