creator-hub/cmd/docker-gateway/douyin_test.go

package main

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"git.ipao.vip/rogee/creator-hub/internal/douyin"
	"golang.org/x/net/websocket"
)

type fakeRestrictedBrowser struct {
	cookies  []douyin.Cookie
	urls     []string
	response restrictedBrowserResponse
	after    func()
}

func (browser *fakeRestrictedBrowser) SetCookies(_ context.Context, _ string, cookies []douyin.Cookie) error {
	browser.cookies = cookies
	if browser.after != nil {
		browser.after()
	}
	return nil
}

func (browser *fakeRestrictedBrowser) Get(_ context.Context, _ string, target string) (restrictedBrowserResponse, error) {
	browser.urls = append(browser.urls, target)
	if browser.after != nil {
		browser.after()
	}
	return browser.response, nil
}

func TestGatewayRestrictedDouyinContract(t *testing.T) {
	labels := map[string]string{
		managedLabel: "true", idLabel: "account-a", bindingVersionLabel: "2", networkIDLabel: "network-a", networkExitLabel: "exit-a",
	}
	self, _ := os.Hostname()
	var stateMu sync.Mutex
	containerNetworks := map[string]string{"creatorhub_browser-account-a": "network-a"}
	runtimeAttached := true
	server := httptest.NewServer(withAliasReservations(self, func(response http.ResponseWriter, request *http.Request) {
		switch request.URL.Path {
		case "/containers/" + namePrefix + "account-a/json":
			stateMu.Lock()
			labelCopy, networkCopy := map[string]string{}, map[string]any{}
			for key, value := range labels {
				labelCopy[key] = value
			}
			for name, id := range containerNetworks {
				networkCopy[name] = map[string]string{"NetworkID": id}
			}
			stateMu.Unlock()
			_ = json.NewEncoder(response).Encode(map[string]any{"Id": "runtime-a", "Config": map[string]any{"Labels": labelCopy},
				"NetworkSettings": map[string]any{"Networks": networkCopy}})
		case "/networks/network-a":
			stateMu.Lock()
			members := map[string]any{self: map[string]string{"Name": self, "IPv4Address": "127.0.0.1/8"}}
			if runtimeAttached {
				members["runtime-a"] = map[string]string{"Name": namePrefix + "account-a", "IPv4Address": "127.0.0.2/8"}
			}
			stateMu.Unlock()
			_ = json.NewEncoder(response).Encode(map[string]any{
				"Id": "network-a", "Name": "creatorhub_browser-account-a", "Driver": "bridge", "Internal": false, "Attachable": false, "Ingress": false,
				"Labels":     map[string]string{managedLabel: "true", networkRoleLabel: browserNetworkRole, idLabel: "account-a", bindingVersionLabel: "2"},
				"Containers": members,
			})
		default:
			response.WriteHeader(http.StatusNotFound)
		}
	}))
	defer server.Close()
	docker := dockerClient{baseURL: server.URL, client: server.Client(), slow: server.Client()}
	browser := &fakeRestrictedBrowser{response: restrictedBrowserResponse{Status: 200, Body: `{"status_code":0}`, Challenge: douyin.ChallengeNone}}
	app := newGatewayWithBrowser(docker, "creatorhub_browser", testToken, self, browser)
	generation := `"binding_version":2,"runtime_id":"runtime-a","network_id":"network-a","network_exit_id":"exit-a"`

	cookieBody := `{` + generation + `,"cookies":[{"name":"sessionid","value":"private-session","domain":".douyin.com","path":"/"}]}`
	response, err := app.Test(authed(http.MethodPost, "/v1/browsers/account-a/douyin/cookies", strings.NewReader(cookieBody)))
	if err != nil || response.StatusCode != http.StatusNoContent || len(browser.cookies) != 1 || browser.cookies[0].Value != "private-session" {
		t.Fatalf("set cookies failed: status=%d cookies=%#v err=%v", response.StatusCode, browser.cookies, err)
	}
	response.Body.Close()

	identityURL := "https://www.douyin.com" + douyinIdentityPath
	getBody := `{` + generation + `,"url":"` + identityURL + `"}`
	response, err = app.Test(authed(http.MethodPost, "/v1/browsers/account-a/douyin/get", strings.NewReader(getBody)))
	body, _ := io.ReadAll(response.Body)
	response.Body.Close()
	if err != nil || response.StatusCode != http.StatusOK || len(browser.urls) != 1 || browser.urls[0] != identityURL ||
		!strings.Contains(string(body), `\"status_code\":0`) || strings.Contains(string(body), "private-session") {
		t.Fatalf("get failed or leaked cookies: status=%d urls=%#v body=%s err=%v", response.StatusCode, browser.urls, body, err)
	}

	for name, request := range map[string]*http.Request{
		"unauthenticated":  httptest.NewRequest(http.MethodPost, "/v1/browsers/account-a/douyin/get", strings.NewReader(getBody)),
		"invalid account":  authed(http.MethodPost, "/v1/browsers/AccountA/douyin/get", strings.NewReader(getBody)),
		"generic URL":      authed(http.MethodPost, "/v1/browsers/account-a/douyin/get", strings.NewReader(`{`+generation+`,"url":"https://example.com/"}`)),
		"generic CDP":      authed(http.MethodPost, "/v1/browsers/account-a/douyin/get", strings.NewReader(`{`+generation+`,"url":"`+identityURL+`","method":"Runtime.evaluate"}`)),
		"trailing null":    authed(http.MethodPost, "/v1/browsers/account-a/douyin/get", strings.NewReader(getBody+`null`)),
		"stale generation": authed(http.MethodPost, "/v1/browsers/account-a/douyin/get", strings.NewReader(`{"binding_version":1,"runtime_id":"runtime-a","network_id":"network-a","network_exit_id":"exit-a","url":"`+identityURL+`"}`)),
	} {
		t.Run(name, func(t *testing.T) {
			before := len(browser.urls)
			response, err := app.Test(request)
			if err != nil {
				t.Fatal(err)
			}
			response.Body.Close()
			want := http.StatusBadRequest
			if name == "unauthenticated" {
				want = http.StatusUnauthorized
			} else if name == "stale generation" {
				want = http.StatusConflict
			}
			if response.StatusCode != want || len(browser.urls) != before {
				t.Fatalf("status=%d want=%d calls=%d want=%d", response.StatusCode, want, len(browser.urls), before)
			}
		})
	}

	stateMu.Lock()
	runtimeAttached = false
	stateMu.Unlock()
	before := len(browser.urls)
	response, err = app.Test(authed(http.MethodPost, "/v1/browsers/account-a/douyin/get", strings.NewReader(getBody)))
	if err != nil || response.StatusCode != http.StatusConflict || len(browser.urls) != before {
		t.Fatalf("wrong network membership reached browser: status=%d calls=%d want=%d err=%v", response.StatusCode, len(browser.urls), before, err)
	}
	response.Body.Close()
	stateMu.Lock()
	runtimeAttached = true
	stateMu.Unlock()

	browser.after = func() {
		stateMu.Lock()
		containerNetworks["other-tenant"] = "network-b"
		stateMu.Unlock()
	}
	response, err = app.Test(authed(http.MethodPost, "/v1/browsers/account-a/douyin/cookies", strings.NewReader(cookieBody)))
	if err != nil || response.StatusCode != http.StatusConflict {
		t.Fatalf("post-operation cross-network attachment was accepted: status=%d err=%v", response.StatusCode, err)
	}
	response.Body.Close()
	stateMu.Lock()
	delete(containerNetworks, "other-tenant")
	stateMu.Unlock()

	browser.after = func() {
		stateMu.Lock()
		labels[networkIDLabel] = "network-replaced"
		stateMu.Unlock()
	}
	response, err = app.Test(authed(http.MethodPost, "/v1/browsers/account-a/douyin/cookies", strings.NewReader(cookieBody)))
	if err != nil || response.StatusCode != http.StatusConflict {
		t.Fatalf("post-operation generation replacement was accepted: status=%d err=%v", response.StatusCode, err)
	}
	response.Body.Close()
}

func TestCDPBrowserUsesOnlyNarrowCommands(t *testing.T) {
	var server *httptest.Server
	var mu sync.Mutex
	methods := []string{}
	secretSeen := false
	cookieNames := []string{}
	cookiesHostOnly := true
	setCookieCalls := 0
	pageOrigin := douyinOrigin
	onlyOldLoader := false
	fetchMode := "ok"
	fetchExpression := ""
	mux := http.NewServeMux()
	mux.HandleFunc("/json/list", func(response http.ResponseWriter, request *http.Request) {
		wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/devtools/page/one"
		_ = json.NewEncoder(response).Encode([]map[string]string{{"type": "page", "webSocketDebuggerUrl": wsURL}})
	})
	mux.Handle("/devtools/page/one", websocket.Server{
		Handshake: func(*websocket.Config, *http.Request) error { return nil },
		Handler: func(connection *websocket.Conn) {
			for {
				var command struct {
					ID     int             `json:"id"`
					Method string          `json:"method"`
					Params json.RawMessage `json:"params"`
				}
				if websocket.JSON.Receive(connection, &command) != nil {
					return
				}
				mu.Lock()
				methods = append(methods, command.Method)
				secretSeen = secretSeen || strings.Contains(string(command.Params), "private-session")
				result := any(map[string]any{})
				var beforeReply, afterReply []map[string]any
				switch command.Method {
				case "Network.clearBrowserCookies":
					cookieNames = nil
				case "Network.setCookies":
					var params struct {
						Cookies []map[string]any `json:"cookies"`
					}
					_ = json.Unmarshal(command.Params, &params)
					setCookieCalls++
					cookieNames = cookieNames[:0]
					for _, cookie := range params.Cookies {
						cookieNames = append(cookieNames, cookie["name"].(string))
						_, hasDomain := cookie["domain"]
						cookiesHostOnly = cookiesHostOnly && !hasDomain && cookie["url"] == douyinOriginURL
					}
				case "Page.navigate":
					result = map[string]any{"frameId": "frame-new", "loaderId": "loader-new"}
					beforeReply = append(beforeReply, map[string]any{"method": "Page.lifecycleEvent", "params": map[string]any{
						"frameId": "frame-old", "loaderId": "loader-old", "name": "load",
					}})
					if !onlyOldLoader {
						afterReply = append(afterReply, map[string]any{"method": "Page.lifecycleEvent", "params": map[string]any{
							"frameId": "frame-new", "loaderId": "loader-new", "name": "load",
						}})
					}
				case "Runtime.evaluate":
					var params struct {
						Expression string `json:"expression"`
					}
					_ = json.Unmarshal(command.Params, &params)
					if params.Expression == "location.origin" {
						result = map[string]any{"result": map[string]any{"value": pageOrigin}}
					} else {
						fetchExpression = params.Expression
						value := map[string]any{"status": 412, "body": `{"captcha":true}`, "too_large": false}
						if fetchMode == "redirect" {
							value = map[string]any{"status": 302, "body": "", "too_large": false}
						} else if fetchMode == "too_large" {
							value = map[string]any{"too_large": true}
						}
						result = map[string]any{"result": map[string]any{"value": map[string]any{
							"status": value["status"], "body": value["body"], "too_large": value["too_large"],
						}}}
					}
				}
				mu.Unlock()
				for _, event := range beforeReply {
					_ = websocket.JSON.Send(connection, event)
				}
				_ = websocket.JSON.Send(connection, map[string]any{"id": command.ID, "result": result})
				for _, event := range afterReply {
					_ = websocket.JSON.Send(connection, event)
				}
			}
		},
	})
	server = httptest.NewServer(mux)
	defer server.Close()
	browser := cdpBrowser{endpoint: func(string) string { return server.URL }, client: server.Client()}
	if err := browser.SetCookies(context.Background(), "account-a", []douyin.Cookie{{
		Name: "old_auth", Value: "old-session", Domain: ".douyin.com", Path: "/",
	}}); err != nil {
		t.Fatal(err)
	}
	if err := browser.SetCookies(context.Background(), "account-a", []douyin.Cookie{{
		Name: "sessionid", Value: "private-session", Domain: ".douyin.com", Path: "/",
	}}); err != nil {
		t.Fatal(err)
	}
	result, err := browser.Get(context.Background(), "account-a", "https://www.douyin.com"+douyinIdentityPath)
	if err != nil || result.Status != 412 || result.Challenge != douyin.ChallengeCaptcha {
		t.Fatalf("unexpected CDP response: %#v err=%v", result, err)
	}
	mu.Lock()
	if !secretSeen || !cookiesHostOnly || strings.Join(cookieNames, ",") != "sessionid" ||
		strings.Join(methods, ",") != "Network.enable,Network.clearBrowserCookies,Page.enable,Page.setLifecycleEventsEnabled,Page.navigate,Runtime.evaluate,Network.setCookies,"+
			"Network.enable,Network.clearBrowserCookies,Page.enable,Page.setLifecycleEventsEnabled,Page.navigate,Runtime.evaluate,Network.setCookies,Runtime.evaluate,Runtime.evaluate" ||
		!strings.Contains(fetchExpression, `redirect:"error"`) || !strings.Contains(fetchExpression, "getReader()") ||
		!strings.Contains(fetchExpression, "q.cancel()") || !strings.Contains(fetchExpression, ">=1048576") || strings.Contains(fetchExpression, "r.text()") {
		mu.Unlock()
		t.Fatalf("unexpected CDP contract: methods=%#v cookies=%#v secret_seen=%v expression=%s", methods, cookieNames, secretSeen, fetchExpression)
	}
	pageOrigin = "https://login.douyin.com"
	setCookiesBeforeRedirect := setCookieCalls
	mu.Unlock()
	if err := browser.SetCookies(context.Background(), "account-a", []douyin.Cookie{{
		Name: "sessionid", Value: "private-session", Domain: ".douyin.com", Path: "/",
	}}); err == nil {
		t.Fatal("accepted navigation redirected to a Douyin subdomain")
	}
	mu.Lock()
	if setCookieCalls != setCookiesBeforeRedirect {
		mu.Unlock()
		t.Fatal("set cookies before rejecting redirected navigation")
	}
	pageOrigin, fetchMode = douyinOrigin, "redirect"
	mu.Unlock()
	if _, err := browser.Get(context.Background(), "account-a", "https://www.douyin.com"+douyinIdentityPath); err == nil {
		t.Fatal("accepted a redirected fetch")
	}
	mu.Lock()
	fetchMode = "too_large"
	mu.Unlock()
	if _, err := browser.Get(context.Background(), "account-a", "https://www.douyin.com"+douyinIdentityPath); err == nil {
		t.Fatal("accepted a response at the 1 MiB limit")
	}
	mu.Lock()
	onlyOldLoader = true
	setCookiesBeforeOldLoader := setCookieCalls
	mu.Unlock()
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()
	if err := browser.SetCookies(ctx, "account-a", []douyin.Cookie{{
		Name: "sessionid", Value: "private-session", Domain: ".douyin.com", Path: "/",
	}}); err == nil {
		t.Fatal("accepted an old page load event for the new navigation")
	}
	mu.Lock()
	defer mu.Unlock()
	if setCookieCalls != setCookiesBeforeOldLoader {
		t.Fatal("set cookies before the new loader completed")
	}
}

func TestCDPDiscoveryDoesNotFollowRedirects(t *testing.T) {
	redirected := 0
	server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
		if request.URL.Path == "/json/list" {
			http.Redirect(response, request, "/redirected", http.StatusFound)
			return
		}
		redirected++
		response.WriteHeader(http.StatusInternalServerError)
	}))
	defer server.Close()
	browser := cdpBrowser{endpoint: func(string) string { return server.URL }, client: server.Client()}
	if _, err := browser.connect(context.Background(), "account-a"); err == nil || redirected != 0 {
		t.Fatalf("discovery redirect was followed: redirected=%d err=%v", redirected, err)
	}
}

func TestCDPDiscoveryRequiresOnePage(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
		_ = json.NewEncoder(response).Encode([]map[string]string{
			{"type": "page", "webSocketDebuggerUrl": "ws://localhost/devtools/page/one"},
			{"type": "page", "webSocketDebuggerUrl": "ws://localhost/devtools/page/two"},
		})
	}))
	defer server.Close()
	browser := cdpBrowser{endpoint: func(string) string { return server.URL }, client: server.Client()}
	if _, err := browser.connect(context.Background(), "account-a"); err == nil {
		t.Fatal("accepted a profile with multiple page targets")
	}
}

func TestCDPDiscoveryRequiresOneJSONValue(t *testing.T) {
	for name, suffix := range map[string]string{
		"null": "null", "other value": `{}`, "garbage": "garbage", "oversized": strings.Repeat(" ", 64<<10), "whitespace": " \n\t",
	} {
		t.Run(name, func(t *testing.T) {
			websocketAttempts := 0
			var server *httptest.Server
			mux := http.NewServeMux()
			mux.HandleFunc("/json/list", func(response http.ResponseWriter, request *http.Request) {
				wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/devtools/page/one"
				_, _ = response.Write([]byte(`[{"type":"page","webSocketDebuggerUrl":"` + wsURL + `"}]` + suffix))
			})
			mux.Handle("/devtools/page/one", websocket.Server{
				Handshake: func(*websocket.Config, *http.Request) error {
					websocketAttempts++
					return nil
				},
				Handler: func(connection *websocket.Conn) {},
			})
			server = httptest.NewServer(mux)
			defer server.Close()

			connection, err := (cdpBrowser{endpoint: func(string) string { return server.URL }, client: server.Client()}).connect(context.Background(), "account-a")
			if connection != nil {
				connection.Close()
			}
			valid := name == "whitespace"
			wantAttempts := 0
			if valid {
				wantAttempts = 1
			}
			if (err == nil) != valid || websocketAttempts != wantAttempts {
				t.Fatalf("err=%v websocket attempts=%d", err, websocketAttempts)
			}
		})
	}
}

func TestDouyinURLContract(t *testing.T) {
	for target, want := range map[string]bool{
		"https://www.douyin.com" + douyinIdentityPath:                                                          true,
		"https://www.douyin.com" + douyinWorksPath + "?sec_user_id=sec-a&count=20&max_cursor=0":                true,
		"https://www.douyin.com" + douyinWorksPath + "?sec_user_id=sec-a&count=20&max_cursor=1":                false,
		"https://www.douyin.com" + douyinWorksPath + "?sec_user_id=sec-a&count=20&max_cursor=0&method=publish": false,
		"https://www.douyin.com/aweme/v1/web/commit/item/":                                                     false,
	} {
		if got := validDouyinURL(target); got != want {
			t.Fatalf("validDouyinURL(%q)=%v want=%v", target, got, want)
		}
	}
}