diff --git a/cmd/control-plane/douyin.go b/cmd/control-plane/douyin.go index 299b20d..cd826c8 100644 --- a/cmd/control-plane/douyin.go +++ b/cmd/control-plane/douyin.go @@ -37,8 +37,8 @@ const maxCreatorLoginQRBytes = 8 << 20 func (browser douyinGatewayBrowser) LoginQR(ctx context.Context) (douyinLoginQRResponse, error) { payload := gatewayGenerationPayload(browser.environment) - status, body, err := gatewayCall(ctx, browser.gateway, http.MethodPost, - "/v1/browsers/"+url.PathEscape(browser.environment.Alias)+"/douyin/login-qr", payload, 30*time.Second) + status, body, err := gatewayCallWithLimit(ctx, browser.gateway, http.MethodPost, + "/v1/browsers/"+url.PathEscape(browser.environment.Alias)+"/douyin/login-qr", payload, 30*time.Second, largeGatewayResponseLimit) if err != nil { return douyinLoginQRResponse{}, err } diff --git a/cmd/control-plane/douyin_test.go b/cmd/control-plane/douyin_test.go index f8af67e..58c7868 100644 --- a/cmd/control-plane/douyin_test.go +++ b/cmd/control-plane/douyin_test.go @@ -45,6 +45,7 @@ func TestDouyinGatewayBrowserFencesAccountGeneration(t *testing.T) { } func TestDouyinGatewayBrowserCapturesLoginScreen(t *testing.T) { + encoded := strings.Repeat("cG5n", 300_000) server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { if request.URL.Path != "/v1/browsers/account-a/douyin/login-qr" { t.Fatalf("unexpected path: %s", request.URL.Path) @@ -54,14 +55,14 @@ func TestDouyinGatewayBrowserCapturesLoginScreen(t *testing.T) { } _ = json.NewEncoder(response).Encode(map[string]any{ "content_type": "image/png", - "body_base64": "cG5n", + "body_base64": encoded, "qr_detected": true, }) })) defer server.Close() browser := douyinGatewayBrowser{gateway: hub.Gateway{Endpoint: server.URL, Token: "gateway-token-1"}, environment: readyDouyinEnvironment()} result, err := browser.LoginQR(context.Background()) - if err != nil || result.ContentType != "image/png" || result.BodyBase64 != "cG5n" || !result.QRDetected { + if err != nil || result.ContentType != "image/png" || result.BodyBase64 != encoded || !result.QRDetected { t.Fatalf("unexpected login screen: %#v err=%v", result, err) } } diff --git a/cmd/control-plane/hub.go b/cmd/control-plane/hub.go index db29b7e..7463779 100644 --- a/cmd/control-plane/hub.go +++ b/cmd/control-plane/hub.go @@ -79,8 +79,17 @@ const ( var gatewayGenerationIDPattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._/-]{0,127}$`) +const ( + defaultGatewayResponseLimit = 1 << 20 + largeGatewayResponseLimit = 16 << 20 +) + // gatewayCall 调用某个网关的 /v1 路由;ok 为 false 时 status/body 携带网关错误。 func gatewayCall(ctx context.Context, target hub.Gateway, method, path string, body any, timeout time.Duration) (status int, responseBody []byte, err error) { + return gatewayCallWithLimit(ctx, target, method, path, body, timeout, defaultGatewayResponseLimit) +} + +func gatewayCallWithLimit(ctx context.Context, target hub.Gateway, method, path string, body any, timeout time.Duration, responseLimit int) (status int, responseBody []byte, err error) { callCtx, cancel := context.WithTimeout(ctx, timeout) defer cancel() var payload io.Reader @@ -104,7 +113,7 @@ func gatewayCall(ctx context.Context, target hub.Gateway, method, path string, b return 0, nil, requestErr } defer response.Body.Close() - responseBody, err = io.ReadAll(io.LimitReader(response.Body, 1<<20)) + responseBody, err = io.ReadAll(io.LimitReader(response.Body, int64(responseLimit))) return response.StatusCode, responseBody, err } diff --git a/cmd/docker_gateway/douyin.py b/cmd/docker_gateway/douyin.py index 9b029b8..7cae881 100644 --- a/cmd/docker_gateway/douyin.py +++ b/cmd/docker_gateway/douyin.py @@ -31,6 +31,8 @@ LOGIN_ORIGINS = frozenset( {ORIGIN, "https://sso.douyin.com", "https://verify.snssdk.com", "https://verify.bytedance.com"} ) LOGIN_SCREENSHOT_LIMIT = 8 << 20 +LOGIN_RENDER_TIMEOUT = 12.0 +LOGIN_RENDER_MIN_BYTES = 20 << 10 IDENTITY_URL = ( ORIGIN + "/aweme/v1/web/user/profile/self/?aid=6383&device_platform=webapp" ) @@ -396,6 +398,22 @@ class DouyinBrowser: raise DouyinError("restricted browser fetch returned invalid body") return BrowserResponse(status, body, detect_challenge(status, body)) + def _wait_for_login_render(self, cdp: CDPConnection) -> None: + deadline = time.monotonic() + LOGIN_RENDER_TIMEOUT + while True: + screenshot = cdp.command( + "Page.captureScreenshot", {"format": "png", "fromSurface": False} + ) + if ( + isinstance(screenshot, dict) + and isinstance(screenshot.get("data"), str) + and len(screenshot["data"]) >= LOGIN_RENDER_MIN_BYTES + ): + return + if time.monotonic() >= deadline: + raise DouyinError("Douyin login page did not finish rendering") + time.sleep(0.5) + def login_qr(self, alias: str) -> BrowserLoginQRResponse: with self.connection(alias) as cdp: navigation = cdp.command("Page.navigate", {"url": ORIGIN_URL}) @@ -405,7 +423,7 @@ class DouyinBrowser: or navigation.get("errorText") ): raise DouyinError("Douyin login page navigation failed") - time.sleep(0.5) + self._wait_for_login_render(cdp) opened = cdp.evaluate( """(() => { const text = value => String(value || '').replace(/\\s+/g, ''); @@ -456,7 +474,7 @@ class DouyinBrowser: or not isinstance(page.get("qr_detected"), bool) ): raise DouyinError("Douyin login page origin is not allowed") - screenshot_params = {"format": "png", "fromSurface": True} + screenshot_params = {"format": "png", "fromSurface": False} if isinstance(page.get("clip"), dict): screenshot_params["clip"] = page["clip"] screenshot = cdp.command("Page.captureScreenshot", screenshot_params) diff --git a/cmd/docker_gateway/test_gateway.py b/cmd/docker_gateway/test_gateway.py index cf229cd..485bd7f 100644 --- a/cmd/docker_gateway/test_gateway.py +++ b/cmd/docker_gateway/test_gateway.py @@ -730,7 +730,7 @@ class BrowserTests(unittest.TestCase): self.assertNotIn("Network.setCookies", [method for method, _ in cdp.commands]) def test_login_qr_captures_a_browser_screen_without_credentials(self) -> None: - screenshot = base64.b64encode(b"png-bytes").decode("ascii") + screenshot = base64.b64encode(b"png-bytes" * 3000).decode("ascii") cdp = Mock() cdp.evaluate.side_effect = [ True, @@ -739,6 +739,7 @@ class BrowserTests(unittest.TestCase): cdp.command.side_effect = [ {"frameId": "frame-1"}, {"data": screenshot}, + {"data": screenshot}, ] browser = DouyinBrowser() self._with_connection(browser, cast(BrowserCDP, cdp)) @@ -749,7 +750,11 @@ class BrowserTests(unittest.TestCase): self.assertTrue(response.qr_detected) self.assertEqual( [call.args[0] for call in cdp.command.call_args_list], - ["Page.navigate", "Page.captureScreenshot"], + [ + "Page.navigate", + "Page.captureScreenshot", + "Page.captureScreenshot", + ], ) self.assertFalse(any("cookie" in expression.lower() for expression in cdp.evaluate.call_args.args))