diff --git a/ws/client.go b/ws/client.go index 686b040c..5136f5a2 100644 --- a/ws/client.go +++ b/ws/client.go @@ -34,6 +34,7 @@ type Client struct { domain string headers http.Header source string + httpClient *http.Client conn *ws.Conn connUrl *url.URL serviceID string @@ -114,6 +115,15 @@ func WithClientAssertionProvider(provider larkcore.ClientAssertionProvider) Clie cli.clientAssertionProvider = provider } } + +func WithHTTPClient(httpClient *http.Client) ClientOption { + return func(cli *Client) { + if httpClient != nil { + cli.httpClient = httpClient + } + } +} + func WithOnReady(f func()) ClientOption { return func(cli *Client) { cli.onReady = f @@ -190,6 +200,7 @@ func NewClient(appId, appSecret string, opts ...ClientOption) *Client { pingInterval: 2 * time.Minute, cache: larkcache.New(30 * time.Second), domain: lark.FeishuBaseUrl, + httpClient: bootstrapHTTPClient, } for _, opt := range opts { @@ -420,7 +431,7 @@ func (c *Client) getConnURL(ctx context.Context) (url string, err error) { } } req.Header.Set("User-Agent", larkcore.UserAgent(c.source)) - resp, err := bootstrapHTTPClient.Do(req) + resp, err := c.httpClient.Do(req) if err != nil { return } diff --git a/ws/client_test.go b/ws/client_test.go index 29bb732d..e7c6361d 100644 --- a/ws/client_test.go +++ b/ws/client_test.go @@ -12,6 +12,16 @@ import ( larkcore "github.com/larksuite/oapi-sdk-go/v3/core" ) +type headerRoundTripper struct { + base http.RoundTripper +} + +func (t headerRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + req = req.Clone(req.Context()) + req.Header.Set("X-Bootstrap-Client", "custom") + return t.base.RoundTrip(req) +} + type wsMockClientAssertionProvider struct { mu sync.Mutex tokens []*larkcore.Token @@ -32,6 +42,28 @@ func (p *wsMockClientAssertionProvider) RetrieveToken(ctx context.Context, aud s return token, nil } +func TestGetConnURLUsesConfiguredHTTPClient(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("X-Bootstrap-Client") != "custom" { + t.Fatalf("bootstrap request did not use configured HTTP client") + } + _ = json.NewEncoder(w).Encode(&EndpointResp{Code: OK, Data: &Endpoint{Url: "wss://example.com/ws"}}) + })) + defer server.Close() + + httpClient := server.Client() + httpClient.Transport = headerRoundTripper{base: httpClient.Transport} + + client := NewClient("app-id", "app-secret", WithDomain(server.URL), WithHTTPClient(httpClient)) + connURL, err := client.getConnURL(context.Background()) + if err != nil { + t.Fatalf("get conn url failed: %v", err) + } + if connURL != "wss://example.com/ws" { + t.Fatalf("unexpected conn url: %s", connURL) + } +} + func TestGetConnURLWithAppSecret(t *testing.T) { originalClient := bootstrapHTTPClient defer func() { bootstrapHTTPClient = originalClient }()