| 1 | package netclient |
| 2 | |
| 3 | import ( |
| 4 | "fmt" |
| 5 | "net" |
| 6 | "net/http" |
| 7 | "net/url" |
| 8 | "strings" |
| 9 | ) |
| 10 | |
| 11 | // ModelProxyOriginalURLHeader is private to the authenticated loopback hop. |
| 12 | // The desktop validates it against the route's frozen provider origins and |
| 13 | // strips it before sending the model request upstream. |
| 14 | const ModelProxyOriginalURLHeader = "X-Reasonix-Model-Proxy-URL" |
| 15 | |
| 16 | type modelCredentialProxyTransport struct { |
| 17 | base *http.Transport |
| 18 | endpoint *url.URL |
| 19 | token string |
| 20 | } |
| 21 | |
| 22 | func NewModelCredentialProxyClient(endpoint, token string) (*http.Client, error) { |
| 23 | u, err := url.Parse(endpoint) |
| 24 | if err != nil || u.Scheme != "http" || u.User != nil || u.RawQuery != "" || u.Fragment != "" || (u.Path != "" && u.Path != "/") { |
| 25 | return nil, fmt.Errorf("invalid model credential proxy endpoint") |
| 26 | } |
| 27 | ip := net.ParseIP(u.Hostname()) |
| 28 | if ip == nil || !ip.IsLoopback() || strings.TrimSpace(token) == "" { |
| 29 | return nil, fmt.Errorf("model credential proxy requires a loopback endpoint and a token") |
| 30 | } |
| 31 | base := http.DefaultTransport.(*http.Transport).Clone() |
| 32 | base.Proxy = nil |
| 33 | return &http.Client{Transport: &modelCredentialProxyTransport{base, u, token}, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}, nil |
| 34 | } |
| 35 | |
| 36 | func (t *modelCredentialProxyTransport) RoundTrip(request *http.Request) (*http.Response, error) { |
| 37 | proxied := request.Clone(request.Context()) |
| 38 | original := request.URL.String() |
| 39 | proxied.URL.Scheme, proxied.URL.Host = t.endpoint.Scheme, t.endpoint.Host |
| 40 | proxied.Host = "" |
| 41 | proxied.Header.Set(ModelProxyOriginalURLHeader, original) |
| 42 | proxied.Header.Del("x-api-key") |
| 43 | proxied.Header.Set("Authorization", "Bearer "+t.token) |
| 44 | return t.base.RoundTrip(proxied) |
| 45 | } |
| 46 | |
| 47 | func (t *modelCredentialProxyTransport) CloseIdleConnections() { t.base.CloseIdleConnections() } |
| 48 |