fix(ws): support custom bootstrap HTTP client

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
 	}

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 }()