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