Compare commits

...
Author SHA1 Message Date
leeguooooo 8aa52d7faa fix(release): 为 fork 预发布设置默认 tag 2026-03-10 13:57:30 +09:00
leeguooooo fcd891e8a9 fix(cli): 增加 abs start 并修复 managed 9333 启动链路 2026-03-10 12:45:50 +09:00
leeguooooo 8eee9310a4 fix(cdp): 自动拉起 9333 专用浏览器 2026-03-09 17:45:13 +09:00
leeguooooo b3ed4b63be fix(native): 修复 CDP 附着卡死问题 2026-03-09 17:10:42 +09:00
leeguooooo 5d149afef1 chore(release): bump version to 0.17.0-fork.1 2026-03-09 12:02:00 +09:00
leeguooooo 356e2f5f39 feat(sync): 同步 upstream 改动并升级到 0.16.3-fork.5 2026-03-09 12:00:13 +09:00
Chris Tate 3cbc284076 fix: persist auth cookies on close in native mode (#650)
(cherry picked from commit b7e7a2548e)
2026-03-09 09:38:10 +09:00
Chris Tate bb92e08fdc Fix Chrome extensions not loading by forcing headed mode when extensions present (#652)
* Fix Chrome extensions not loading by forcing headed mode when extensions present

Fixes #640

* Restore wait_or_kill() and add tests for headless+extensions logic

Restore the ChromeProcess::wait_or_kill() method that was accidentally
removed. It is still referenced by BrowserProcess in browser.rs and is
needed for graceful shutdown / cookie persistence (PR #650).

Add unit tests verifying --headless=new is omitted when extensions are
present.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Fix window-size leak in headed+extensions mode and remove unused channel option

- Skip --window-size=1280,720 when extensions force headed mode (native)
- Remove unexplained channel: 'chromium' from extensions launch path (TS)
- Add window-size assertion to existing extension test

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: ctate <366502+ctate@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
(cherry picked from commit 68cebe5192)
2026-03-09 09:37:36 +09:00
Chris Tate 5c142b02b5 Fix: Suppress Google Translate bar in native headless mode (#649)
Fixes #617

(cherry picked from commit 492830accb)
2026-03-09 09:37:33 +09:00
Chris Tate b100870a02 fix: native auth login fails due to incompatible encryption format (#648)
* fix: native auth login fails due to incompatible encryption format

* fixes

* fixes

(cherry picked from commit 7acde7e29a)
2026-03-09 09:37:30 +09:00
leeguooooo b505d00c97 fix: escape compose shell vars to prevent stale multi-platform binaries 2026-03-05 18:55:22 +09:00
leeguooooo 5e0582cede feat: enforce default session daemon isolation and bump 0.16.3-fork.4 2026-03-05 18:44:52 +09:00
leeguooooo 5970579d7c feat: add parallel mode and idle daemon shutdown 2026-03-05 13:40:34 +09:00
leeguooooo 8880aa2f35 chore(release): 0.16.3-fork.2
- restore extension mode default to headed when headless is unspecified

- keep explicit headless override behavior
2026-03-05 11:44:47 +09:00
leeguooooo f051e72f85 chore(release): bump to 0.16.3-fork.1 2026-03-05 11:20:06 +09:00
leeguooooo 6ae703565c fix(connection): surface daemon startup stderr during launch 2026-03-05 11:18:21 +09:00
layla d56442cf91 Fix dialog dismiss command parsing (#605) 2026-03-05 11:16:14 +09:00
Li Yang dad7be8c77 fix: use reqwest for CDP port discovery instead of broken hand-rolled HTTP client (#619)
reqwest_get_string() was hand-rolling HTTP/1.1 over raw TCP despite reqwest
being an existing dependency. The hand-rolled implementation had two bugs:

1. URL path parsing: url.find('/') matched the first '/' in 'http://',
   producing path '//127.0.0.1:9222/json/version' instead of '/json/version'

2. read_to_end() hangs: Chrome's DevTools HTTP server ignores Connection: close
   and keeps the socket open, so read_to_end() waits for EOF that never comes

This caused 'agent-browser --cdp <port>' to always timeout when AGENT_BROWSER_NATIVE=1.

Fix: replace 49 lines of broken TCP code with reqwest::get(), which was
already in Cargo.toml.
2026-03-05 11:16:14 +09:00
Chris Tate 7921928ec4 headed mode (#607)
* headed mode

* fixes

* fixes

* docs

* fixes

* fixes

* fixes
2026-03-05 11:15:17 +09:00
leeguooooo a71198d591 fix(stealth): neutralize disable-devtool auto bootstrap 2026-03-05 11:11:36 +09:00
leeguooooo 19aefb7eb8 Update binaries for 0.16.1-fork.4 2026-03-04 18:46:17 +09:00
leeguooooo ff2e66c0a1 chore(release): bump to 0.16.1-fork.4 2026-03-04 18:24:10 +09:00
leeguooooo d320e1df47 refactor(cli): ignore --session and enforce default runtime session 2026-03-04 18:19:33 +09:00
leeguooooo 7b454e14ed fix: stabilize native release builds across platforms 2026-03-04 15:44:28 +09:00
leeguooooo 17d44baa2f fix(release): enforce binary version integrity pre/post publish 2026-03-04 15:01:52 +09:00
leeguooooo 59589043db fix(stealth): bypass worker wrapping in cloudflare challenge runtime 2026-03-04 13:48:16 +09:00
leeguooooo 9e7c8937a6 chore: 提交剩余文档与扩展面板资源更新 2026-03-04 11:54:55 +09:00
leeguooooo 2d0d18f0d9 fix(extension): default to high-risk mode by disabling page bridge 2026-03-04 11:46:44 +09:00
leeguooooo ae22bda46c chore(release): bump to 0.16.1-fork.3 2026-03-04 11:35:11 +09:00
leeguooooo d82357fea4 fix(release): enforce bundled binary version checks and bump to 0.16.1-fork.2 2026-03-04 10:25:04 +09:00
leeguooooo 66a39f3c83 feat: add abs alias, refresh extension panel, and bump to 0.16.1-fork.1 2026-03-04 09:57:56 +09:00
leeguooooo eedf824af7 merge: upstream v0.16.1 while preserving fork stealth behaviors 2026-03-04 09:48:58 +09:00
leeguooooo c07eb7ee52 chore: 更新 Cloudflare 及浏览器自动化攻防文章并补发 blog 链接 2026-03-04 09:39:38 +09:00
leeguooooo d32a1d046a Detail agent stealth Cloudflare fix 2026-03-03 18:04:44 +09:00
leeguooooo 0c0ed5e72c chore(release): bump version to 0.15.2-fork.2 2026-03-03 17:13:07 +09:00
leeguooooo 34092ec193 fix: prevent stale native version and skip npm-bin rewrite on pnpm 2026-03-03 17:12:17 +09:00
leeguooooo 726377c4c1 feat: add doctor diagnostics and bump to 0.15.2-fork.1 2026-03-03 17:05:42 +09:00
leeguooooo 8e2e4abce6 feat: complete agent-browser-stealth extension controls 2026-03-03 12:33:46 +09:00
leeguooooo 870895e922 feat: expand agent-browser-stealth extension capabilities 2026-03-03 12:28:34 +09:00
leeguooooo 2a766cfe48 feat: add CDP tab-group plugin handshake with silent fallback 2026-03-03 12:17:14 +09:00
leeguooooo d04cf59238 feat: auto-group agent tabs by default in local Chromium 2026-03-03 11:37:53 +09:00
leeguooooo 0a257ad2c1 merge: sync upstream/main into fork main (v0.15.2) 2026-03-03 09:57:17 +09:00
leeguooooo 44c0361fcd Update readme with stealth FAQ 2026-03-02 18:32:28 +09:00
leeguooooo 907ca8c808 Update agent-browser skill installs 2026-03-02 09:38:48 +09:00
leeguooooo 74fda70b67 fix(release): publish correct native version in fork.7 2026-02-28 11:29:21 +09:00
leeguooooo 2a397de59f merge: sync upstream/main into fork main 2026-02-28 11:19:50 +09:00
leeguooooo bf672ee7f9 fix(stealth): avoid matchMedia Illegal invocation
- bind MediaQueryList methods when proxied for prefers-color-scheme light patch

- add regression test for addEventListener/removeEventListener

- bump version to 0.14.0-fork.6 and sync Cargo metadata
2026-02-28 11:17:02 +09:00
leeguooooo 74910cfef1 Merge tag 'v0.15.0' into codex/sync-v0.15.0
v0.15.0

# Conflicts:
#	CHANGELOG.md
#	README.md
#	cli/Cargo.lock
#	cli/Cargo.toml
#	cli/src/commands.rs
#	cli/src/connection.rs
#	cli/src/flags.rs
#	cli/src/main.rs
#	docs/src/app/commands/page.mdx
#	docs/src/app/configuration/page.mdx
#	package.json
#	src/actions.ts
2026-02-27 10:49:55 +09:00
leeguooooo 41830dff71 fix(cookies): require domain and path together when url is absent 2026-02-27 10:49:12 +09:00
leeguooooo e005c7251b feat(stealth): add risk-mode signals and document stealth architecture 2026-02-25 10:27:45 +09:00
leeguooooo 11eab471f1 chore(release): bump version to 0.14.0-fork.4 2026-02-25 10:10:26 +09:00
leeguooooo aa256e30c7 Merge remote-tracking branch 'upstream/main'
# Conflicts:
#	README.md
#	cli/src/connection.rs
#	cli/src/flags.rs
#	cli/src/main.rs
#	src/browser.ts
#	src/protocol.ts
2026-02-25 10:06:20 +09:00
leeguooooo 6f1dd39121 chore(clawhub): 改为本地 pre-push 自动同步 skill 2026-02-24 18:10:05 +09:00
leeguooooo 85d18799a4 feat(skill): 新增 agent-browser-stealth 的 OpenClaw skill 与 ClawHub 自动同步 2026-02-24 18:05:29 +09:00
leeguooooo 43e781a8d3 docs(readme): 精简文档并聚焦反爬能力 2026-02-24 18:00:19 +09:00
leeguooooo 25e8719e51 docs(readme): 补充 --delay 文本转义与 stealth 行为说明 2026-02-24 17:58:13 +09:00
leeguooooo ec011f46ff fix(stealth): 修复 launch 选项与测试基线不一致问题
在 stealth 默认策略下保留自定义 user-agent,不再被 CDP 覆盖。

同步更新 protocol/browser/launch/file-access 相关测试预期,并放宽 browser.test 的 hook 超时以消除偶发超时。
2026-02-24 17:42:12 +09:00
leeguooooo aef8fcc038 ci(release): 修复 OIDC 发布认证链路
移除 setup-node 的 registry-url 注入,避免发布步骤继承无效 NODE_AUTH_TOKEN。

发布前升级 npm 到 v11,使用独立 npmrc 并启用 provenance,以匹配 npm trusted publishing。
2026-02-24 17:24:20 +09:00
leeguooooo 96582b79fd ci(release): 调整 trusted publishing 发布流程
将 changesets/action 改为仅处理 version/PR,不再由其执行 publish。

新增发布前版本检查与独立 pnpm ci:publish 步骤,避免 OIDC 发布在 action 内失败。
2026-02-24 17:20:17 +09:00
leeguooooo 058a286326 chore(release): 发布 0.14.0-fork.3
更新 package.json 版本并写入对应 changelog 条目。
2026-02-24 17:14:06 +09:00
leeguooooo b1f27236d8 fix(cli): 修复 type/keyboard 的 --delay 参数解析
将 --delay <ms> 从输入文本中剥离并写入 delay 字段,避免搜索词混入参数。

同时支持使用 -- 终止参数解析以输入字面量 --delay 文本,并补充回归测试与帮助文档。
2026-02-24 17:13:46 +09:00
leeguooooo a5a9327b7d docs(home): 补充首页能力亮点说明
- 新增自动区域检测能力描述

- 新增验证码自动重试能力描述
2026-02-24 17:06:34 +09:00
leeguooooo 699ccbd3cb feat(cli): 强制使用用户现有浏览器并移除 profile/channel
- 禁用 --profile/AGENT_BROWSER_PROFILE 与 --channel/AGENT_BROWSER_CHANNEL,并给出项目策略提示

- 默认模式强制连接 localhost:9333,连接失败直接报错,不再自动回退新开浏览器

- 同步更新 README、技能文档、docs 与 --help 输出

- 版本升级到 0.14.0-fork.2 并同步 cli/Cargo.toml 与 Cargo.lock
2026-02-24 17:05:11 +09:00
leeguooooo 893ddfd259 feat(cdp): 默认优先连接 9333 常驻 Chrome
- 无显式连接参数时先尝试 CDP 9333,失败后回退本地浏览器启动

- 修复 CDP 页选择稳定性:过滤 omnibox 系统页、无可用页时自动创建 fallback 页

- 调整页面关闭后的 active 索引维护,降低 No page found 问题

- 新增 agent-browser-stealth 二进制入口并保持与 agent-browser 行为一致

- 同步更新 CLI 帮助、README、技能文档与 docs 说明
2026-02-24 16:22:49 +09:00
leeguooooo ea2e93dbba feat(stealth): 优化隐身对抗并引入双版本发布体系
- 将 CreepJS like headless 指标优化到 0%(headless/stealth 维持 0%)

- 新增 ActiveText 与 prefers-color-scheme 探针修复

- 版本号采用 <upstream>-fork.<fork> 格式并在 --version 输出 upstream/fork

- 更新 README、SKILL 与 docs 中的版本体系说明
2026-02-24 15:29:11 +09:00
leeguooooo 4c6afe3e69 docs(config): 移除 --stealth 配置项说明 2026-02-24 14:55:07 +09:00
leeguooooo 9f9a90cf63 docs(cli): 清理过时 stealth 环境变量说明 2026-02-24 14:54:51 +09:00
leeguooooo 0443e4ed7a feat(stealth): 默认开启并收敛 chrome 指纹特征 2026-02-24 14:54:30 +09:00
leeguooooo c5b2292caa feat(stealth): 增强浏览器级 UA 覆盖并修复背景特征 2026-02-24 14:50:03 +09:00
leeguooooo 2ed0c6f8ec feat(stealth): 进一步降低 creepjs like-headless 指标 2026-02-24 14:36:56 +09:00
leeguooooo 3a91aef4c9 feat(stealth): 优化指纹信号并同步文档与包配置 2026-02-24 14:23:59 +09:00
leeguooooo 02ebc9f328 feat(stealth): 增强指纹一致性并新增 creepjs 检测脚本 2026-02-24 14:19:25 +09:00
leeguooooo 955543b757 chore: prepare fork sync and independent release setup 2026-02-24 14:14:50 +09:00
leeguooooo ecad112707 feat(cli): 默认开启 stealth 并支持 wait 区间超时 2026-02-24 12:27:04 +09:00
leeguooooo 8932f28926 fix(stealth): 修复 headed 模式下 stealth 失效并统一策略
- 修复 launch 协议未透传 stealth 导致 --headed 下补丁失效的问题\n- 在 BrowserManager 引入 StealthPolicy,统一 local/CDP/provider 能力决策\n- 增加 launch 返回 stealth 状态并在 --debug 输出连接类型与能力\n- 补充 local/CDP 回归测试与 bot.sannysoft.com 自动检查脚本\n- 同步 README、CLI help、技能文档与 CDP 文档中的 stealth 能力矩阵
2026-02-24 12:18:47 +09:00
136 changed files with 30120 additions and 3663 deletions
+8
View File
@@ -64,12 +64,20 @@ jobs:
- name: Setup Rust toolchain
uses: dtolnay/rust-toolchain@stable
with:
components: rustfmt, clippy
- name: Cache Rust build artifacts
uses: Swatinem/rust-cache@v2
with:
workspaces: cli
- name: Format check
run: cargo fmt --manifest-path cli/Cargo.toml -- --check
- name: Clippy check
run: cargo clippy --manifest-path cli/Cargo.toml -- -D warnings
- name: Run Rust tests
run: cargo test --profile ci --manifest-path cli/Cargo.toml
+56 -6
View File
@@ -11,6 +11,7 @@ concurrency: ${{ github.workflow }}-${{ github.ref }}
permissions:
contents: write
pull-requests: write
id-token: write
jobs:
# Build native binaries for all platforms first
@@ -133,8 +134,8 @@ jobs:
needs: build-binaries
runs-on: ubuntu-latest
outputs:
published: ${{ steps.changesets.outputs.published }}
publishedPackages: ${{ steps.changesets.outputs.publishedPackages }}
published: ${{ steps.publish_metadata.outputs.published }}
publishedPackages: ${{ steps.publish_metadata.outputs.publishedPackages }}
steps:
- name: Checkout Repo
uses: actions/checkout@v4
@@ -151,7 +152,6 @@ jobs:
with:
node-version: '22'
cache: pnpm
registry-url: 'https://registry.npmjs.org'
- name: Install Dependencies
run: pnpm install --frozen-lockfile
@@ -201,17 +201,67 @@ jobs:
fi
echo "All 5 platform binaries present and valid"
- name: Verify bundled binary versions
run: |
pnpm run verify:bundled-binaries
- name: Create Release Pull Request or Publish to npm
id: changesets
uses: changesets/action@v1
with:
version: pnpm ci:version
publish: pnpm ci:publish
title: 'chore: version packages'
commit: 'chore: version packages'
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
NODE_AUTH_TOKEN: ${{ secrets.NPM_VERCEL_TOKEN_ELEVATED }}
- name: Check if publish is needed
id: publish_check
if: steps.changesets.outputs.hasChangesets == 'false'
run: |
LOCAL_VERSION=$(node -p "require('./package.json').version")
REMOTE_VERSION=$(npm view agent-browser-stealth version 2>/dev/null || echo "")
echo "local_version=$LOCAL_VERSION" >> "$GITHUB_OUTPUT"
echo "remote_version=$REMOTE_VERSION" >> "$GITHUB_OUTPUT"
if [ "$LOCAL_VERSION" != "$REMOTE_VERSION" ]; then
echo "needs_publish=true" >> "$GITHUB_OUTPUT"
else
echo "needs_publish=false" >> "$GITHUB_OUTPUT"
fi
echo "Local: $LOCAL_VERSION"
echo "Remote: ${REMOTE_VERSION:-<none>}"
- name: Publish to npm (trusted publishing)
id: publish_npm
if: steps.changesets.outputs.hasChangesets == 'false' && steps.publish_check.outputs.needs_publish == 'true'
env:
NODE_AUTH_TOKEN: ''
NPM_CONFIG_USERCONFIG: /home/runner/work/_temp/trusted-npmrc
NPM_CONFIG_PROVENANCE: 'true'
run: |
npm install -g npm@^11
npm --version
printf "registry=https://registry.npmjs.org/\n" > "$NPM_CONFIG_USERCONFIG"
pnpm ci:publish
- name: Verify published npm tarball
if: steps.publish_npm.outcome == 'success'
env:
PACKAGE_NAME: agent-browser-stealth
EXPECTED_VERSION: ${{ steps.publish_check.outputs.local_version }}
run: |
pnpm run verify:registry-host-binary
- name: Set release outputs
id: publish_metadata
run: |
if [ "${{ steps.publish_npm.outcome }}" = "success" ]; then
echo "published=true" >> "$GITHUB_OUTPUT"
echo "publishedPackages=[{\"name\":\"agent-browser-stealth\",\"version\":\"${{ steps.publish_check.outputs.local_version }}\"}]" >> "$GITHUB_OUTPUT"
else
echo "published=false" >> "$GITHUB_OUTPUT"
echo "publishedPackages=[]" >> "$GITHUB_OUTPUT"
fi
# Create GitHub release with binaries after npm publish
github-release:
@@ -252,7 +302,7 @@ jobs:
run: |
VERSION=$(node -p "require('./package.json').version")
TAG="v$VERSION"
# Check if release already exists
if gh release view "$TAG" &>/dev/null; then
echo "Release $TAG already exists, uploading binaries..."
+3 -3
View File
@@ -4,10 +4,10 @@ node_modules/
# Build output
dist/
# Native binaries (keep the launcher scripts)
# Native binaries and build byproducts (keep JS launcher)
bin/agent-browser-*
!bin/agent-browser
!bin/agent-browser.cmd
bin/agent-browser
bin/*.d
# Rust build artifacts
cli/target/
+8
View File
@@ -0,0 +1,8 @@
if [ "${SKIP_CLAWHUB_SYNC:-0}" = "1" ]; then
echo "Skipping ClawHub sync (SKIP_CLAWHUB_SYNC=1)"
exit 0
fi
pnpm run clawhub:sync || {
echo "ClawHub sync failed. Push continues. Run 'pnpm run clawhub:sync' manually after fixing login/network."
}
+28 -16
View File
@@ -1,28 +1,38 @@
# agent-browser
## 0.16.1
## 0.16.3-fork.1
### Patch Changes
- c4180c8: Improved Chrome launch reliability by automatically detecting containerized environments (Docker, Podman, Kubernetes) and enabling --no-sandbox when needed. Added support for discovering Playwright-installed Chromium browsers and enhanced error messages with helpful diagnostics when Chrome fails to launch.
- Sync upstream `v0.16.2` / `v0.16.3` core fixes into the fork baseline.
- Import headed-mode behavior updates from upstream.
- Improve CDP debug-port discovery by switching to `reqwest` in native Chrome probing.
- Fix dialog dismiss command parsing consistency.
- Surface daemon startup stderr on launch failure to avoid opaque timeout-only errors.
- Keep fork stealth hardening for anti-debug self-destruct flows (`disable-devtool-auto` bootstrap neutralization).
## 0.16.0
### Minor Changes
- 05018b3: Added experimental native Rust daemon (`--native` flag, `AGENT_BROWSER_NATIVE=1` env, or `"native": true` in config). The native daemon communicates with Chrome directly via CDP, eliminating Node.js and Playwright dependencies. Supports 150+ commands with full parity to the default Node.js daemon. Includes WebDriver backend for Safari/iOS, CDP protocol codegen, request tracking, frame context management, and comprehensive e2e and parity tests.
## 0.15.3
## 0.16.1-fork.5
### Patch Changes
- 62241b5: Fixed Windows compatibility issues including proper handling of extended-length path prefixes from canonicalize(), prevention of MSYS/Git Bash path translation that could mangle arguments, and improved daemon startup reliability. Also added ARM64 Windows support in postinstall shims and expanded CI testing with a full daemon lifecycle test on Windows.
- Harden runtime stealth against anti-debug self-destruct flows on high-risk sites:
- neutralize `disable-devtool` auto bootstrap probes by hiding the `[disable-devtool-auto]` selector entry point
- preserve normal selector behavior for non-target queries to minimize side effects
- add regression tests covering the selector patch boundary
- Expand security design docs with the anti-debug execution-plane model and clarify why page self-close/redirect is a separate surface from fingerprint scoring.
## 0.15.2
## 0.15.2-fork.0
### Patch Changes
- 6aea316: Documentation site improvements and internal tooling updates including enhanced code blocks, mobile navigation, and docs chat components. CLI connection and output handling refinements. Skill creator reference documentation and scripts have been reorganized.
- Merge upstream `v0.15.2` updates, including fixes for cookies clear/tab close output, daemon EPERM liveness checks, unnamed element reference matching, and docs/skills refresh.
## 0.15.1-fork.11
### Patch Changes
- Auto-attach existing browser more reliably by trying CDP localhost:9333 first, then falling back to auto-discovery before failing.
Align daemon behavior and user-facing docs/skill guidance with the same attachment policy.
## 0.15.1
@@ -32,11 +42,13 @@
## 0.15.0
### Minor Changes
### Patch Changes
- 2e38882: - Added security hardening: authentication vault, content boundary markers, domain allowlist, action policy, action confirmation, and output length limits.
- Added `--download-path` flag (and `AGENT_BROWSER_DOWNLOAD_PATH` env / `downloadPath` config key) to set a default download directory.
- Added `--selector` flag to `scroll` command for scrolling within specific container elements.
- Fix CLI typing delay parsing so `--delay` is treated as an option instead of typed text.
- Add `--delay <ms>` parsing for `type` and `keyboard type`
- Support `--` to type literal `--delay` text
- Add regression tests for parsing and delay behavior
- Update CLI help, README, skills, and docs command references
## 0.14.0
+279 -1135
View File
File diff suppressed because it is too large Load Diff
+37 -34
View File
@@ -44,8 +44,8 @@ dependencies = [
]
[[package]]
name = "agent-browser"
version = "0.16.1"
name = "agent-browser-stealth"
version = "0.17.0-fork.2"
dependencies = [
"aes-gcm",
"async-trait",
@@ -628,20 +628,20 @@ dependencies = [
"cfg-if",
"js-sys",
"libc",
"r-efi",
"r-efi 5.3.0",
"wasip2",
"wasm-bindgen",
]
[[package]]
name = "getrandom"
version = "0.4.1"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "139ef39800118c7683f2fd3c98c1b23c09ae076556b435f8e9064ae108aaeeec"
checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555"
dependencies = [
"cfg-if",
"libc",
"r-efi",
"r-efi 6.0.0",
"wasip2",
"wasip3",
]
@@ -772,7 +772,7 @@ dependencies = [
"tokio",
"tokio-rustls",
"tower-service",
"webpki-roots 1.0.5",
"webpki-roots 1.0.6",
]
[[package]]
@@ -1043,9 +1043,9 @@ checksum = "7a79a3332a6609480d7d0c9eab957bca6b455b91bb84e66d19f5ff66294b85b8"
[[package]]
name = "libc"
version = "0.2.180"
version = "0.2.182"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bcc35a38544a891a5f7c865aca548a982ccb3b8650a5b06d0fd33a10283c56fc"
checksum = "6800badb6cb2082ffd7b6a67e6125bb39f18782f793520caee8cb8846be06112"
[[package]]
name = "libfuzzer-sys"
@@ -1059,11 +1059,10 @@ dependencies = [
[[package]]
name = "libredox"
version = "0.1.12"
version = "0.1.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3d0b95e02c851351f877147b7deea7b1afb1df71b63aa5f8270716e0c5720616"
checksum = "1744e39d1d6a9948f4f388969627434e31128196de472883b39f148769bfe30a"
dependencies = [
"bitflags",
"libc",
]
@@ -1106,9 +1105,9 @@ dependencies = [
[[package]]
name = "memchr"
version = "2.7.6"
version = "2.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273"
checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79"
[[package]]
name = "miniz_oxide"
@@ -1315,9 +1314,9 @@ dependencies = [
[[package]]
name = "proc-macro2"
version = "1.0.105"
version = "1.0.106"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "535d180e0ecab6268a3e718bb9fd44db66bbbc256257165fc699dadf70d16fe7"
checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
dependencies = [
"unicode-ident",
]
@@ -1414,14 +1413,14 @@ dependencies = [
"once_cell",
"socket2",
"tracing",
"windows-sys 0.52.0",
"windows-sys 0.60.2",
]
[[package]]
name = "quote"
version = "1.0.43"
version = "1.0.45"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dc74d9a594b72ae6656596548f56f667211f8a97b3d4c3d467150794690dc40a"
checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
dependencies = [
"proc-macro2",
]
@@ -1432,6 +1431,12 @@ version = "5.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
[[package]]
name = "r-efi"
version = "6.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
[[package]]
name = "rand"
version = "0.8.5"
@@ -1607,7 +1612,7 @@ dependencies = [
"wasm-bindgen",
"wasm-bindgen-futures",
"web-sys",
"webpki-roots 1.0.5",
"webpki-roots 1.0.6",
]
[[package]]
@@ -1839,9 +1844,9 @@ checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
[[package]]
name = "syn"
version = "2.0.114"
version = "2.0.117"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d4d107df263a3013ef9b1879b0df87d706ff80f65a86ea879bd9c31f9b307c2a"
checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99"
dependencies = [
"proc-macro2",
"quote",
@@ -1949,9 +1954,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
[[package]]
name = "tokio"
version = "1.49.0"
version = "1.50.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "72a2903cd7736441aac9df9d7688bd0ce48edccaadf181c3b90be801e81d3d86"
checksum = "27ad5e34374e03cfffefc301becb44e9dc3c17584f414349ebe29ed26661822d"
dependencies = [
"bytes",
"libc",
@@ -2098,9 +2103,9 @@ checksum = "562d481066bde0658276a35467c4af00bdc6ee726305698a55b86e61d7ad82bb"
[[package]]
name = "unicode-ident"
version = "1.0.22"
version = "1.0.24"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5"
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
[[package]]
name = "unicode-xid"
@@ -2154,7 +2159,7 @@ version = "1.21.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b672338555252d43fd2240c714dc444b8c6fb0a5c5335e65a07bba7742735ddb"
dependencies = [
"getrandom 0.4.1",
"getrandom 0.4.2",
"js-sys",
"wasm-bindgen",
]
@@ -2328,14 +2333,14 @@ version = "0.26.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9"
dependencies = [
"webpki-roots 1.0.5",
"webpki-roots 1.0.6",
]
[[package]]
name = "webpki-roots"
version = "1.0.5"
version = "1.0.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "12bed680863276c63889429bfd6cab3b99943659923822de1c8a39c49e4d722c"
checksum = "22cfaf3c063993ff62e73cb4311efde4db1efb31ab78a3e5c457939ad5cc0bed"
dependencies = [
"rustls-pki-types",
]
@@ -2779,9 +2784,9 @@ dependencies = [
[[package]]
name = "zmij"
version = "1.0.12"
version = "1.0.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2fc5a66a20078bf1251bde995aa2fdcc4b800c70b5d92dd2c62abc5c60f679f8"
checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa"
[[package]]
name = "zune-core"
@@ -2816,8 +2821,6 @@ dependencies = [
[[package]]
name = "zune-jpeg"
version = "0.5.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "410e9ecef634c709e3831c2cfdb8d9c32164fae1c67496d5b68fff728eec37fe"
dependencies = [
"zune-core 0.5.1",
]
+14 -3
View File
@@ -1,10 +1,18 @@
[package]
name = "agent-browser"
version = "0.16.1"
name = "agent-browser-stealth"
version = "0.17.0-fork.2"
edition = "2021"
description = "Fast browser automation CLI for AI agents"
description = "Stealth browser automation CLI for AI agents with anti-bot evasions"
license = "Apache-2.0"
[[bin]]
name = "agent-browser"
path = "src/main.rs"
[[bin]]
name = "agent-browser-stealth"
path = "src/main_stealth.rs"
[dependencies]
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
@@ -43,3 +51,6 @@ strip = true
inherits = "release"
lto = "thin"
codegen-units = 16
[patch.crates-io]
zune-jpeg = { path = "vendor/zune-jpeg" }
+3 -3
View File
@@ -175,7 +175,7 @@ fn to_snake_case(s: &str) -> String {
// Only insert underscore at transitions from lowercase to uppercase,
// or when an uppercase sequence ends (e.g. "DOM" -> "dom", not "d_o_m")
let prev_upper = chars[i - 1].is_uppercase();
let next_lower = chars.get(i + 1).map_or(false, |n| n.is_lowercase());
let next_lower = chars.get(i + 1).is_some_and(|n| n.is_lowercase());
if !prev_upper || next_lower {
result.push('_');
}
@@ -202,7 +202,7 @@ fn resolve_ref(
// Check if this type actually exists in the referenced domain
if domain_types
.get(ref_domain)
.map_or(false, |t| t.contains(ref_type))
.is_some_and(|t| t.contains(ref_type))
{
format!(
"super::cdp_{}::{}",
@@ -339,7 +339,7 @@ fn generate_domain(
if variant == "Self" {
variant = "SelfValue".to_string();
}
if variant.chars().next().map_or(false, |c| c.is_ascii_digit()) {
if variant.chars().next().is_some_and(|c| c.is_ascii_digit()) {
variant = format!("V{}", variant);
}
if seen_variants.insert(variant.clone()) {
+266 -177
View File
@@ -71,6 +71,62 @@ pub fn gen_id() -> String {
)
}
/// Parse free-form text arguments with optional `--delay <ms>`.
///
/// `--` can be used to stop flag parsing if text must include `--delay` literally.
fn parse_text_with_optional_delay(
args: &[&str],
context: &str,
usage: &'static str,
) -> Result<(String, Option<u64>), ParseError> {
let mut text_parts: Vec<&str> = Vec::new();
let mut delay_ms: Option<u64> = None;
let mut parse_flags = true;
let mut i = 0;
while i < args.len() {
let arg = args[i];
if parse_flags && arg == "--" {
parse_flags = false;
i += 1;
continue;
}
if parse_flags && arg == "--delay" {
let raw = args
.get(i + 1)
.ok_or_else(|| ParseError::MissingArguments {
context: format!("{} --delay", context),
usage,
})?;
let parsed = raw.parse::<u64>().map_err(|_| ParseError::InvalidValue {
message: format!(
"Invalid --delay value: {} (must be a non-negative integer in milliseconds)",
raw
),
usage,
})?;
delay_ms = Some(parsed);
i += 2;
continue;
}
text_parts.push(arg);
i += 1;
}
let text = text_parts.join(" ");
if text.is_empty() {
return Err(ParseError::MissingArguments {
context: context.to_string(),
usage,
});
}
Ok((text, delay_ms))
}
pub fn parse_command(args: &[String], flags: &Flags) -> Result<Value, ParseError> {
if args.is_empty() {
return Err(ParseError::MissingArguments {
@@ -92,7 +148,6 @@ pub fn parse_command(args: &[String], flags: &Flags) -> Result<Value, ParseError
match cmd {
// === Navigation ===
// Maps to "navigate" action in protocol; reflected in ACTION_CATEGORIES in action-policy.ts
"open" | "goto" | "navigate" => {
let url = rest.first().ok_or_else(|| ParseError::MissingArguments {
context: cmd.to_string(),
@@ -129,8 +184,45 @@ pub fn parse_command(args: &[String], flags: &Flags) -> Result<Value, ParseError
nav_cmd["iosDevice"] = json!(device);
}
}
if let Some(ref risk_mode) = flags.risk_mode {
if matches!(risk_mode.as_str(), "off" | "warn" | "block") {
nav_cmd["riskMode"] = json!(risk_mode);
} else {
return Err(ParseError::InvalidValue {
message: format!(
"Invalid --risk-mode value: {} (expected off, warn, or block)",
risk_mode
),
usage: "open <url>",
});
}
}
if let Some(ref wait_until) = flags.wait_until {
if matches!(
wait_until.as_str(),
"load" | "domcontentloaded" | "networkidle"
) {
nav_cmd["waitUntil"] = json!(wait_until);
} else {
return Err(ParseError::InvalidValue {
message: format!(
"Invalid --wait-until value: {} (expected load, domcontentloaded, or networkidle)",
wait_until
),
usage: "open <url>",
});
}
}
Ok(nav_cmd)
}
// Prepare the managed automation browser on localhost:9333 without navigating.
// This reuses the existing launch action so both Node and native daemons stay aligned.
"start" => Ok(json!({
"id": id,
"action": "launch",
"cdpPort": 9333,
"headless": !flags.headed
})),
"back" => Ok(json!({ "id": id, "action": "back" })),
"forward" => Ok(json!({ "id": id, "action": "forward" })),
"reload" => Ok(json!({ "id": id, "action": "reload" })),
@@ -168,9 +260,18 @@ pub fn parse_command(args: &[String], flags: &Flags) -> Result<Value, ParseError
"type" => {
let sel = rest.first().ok_or_else(|| ParseError::MissingArguments {
context: "type".to_string(),
usage: "type <selector> <text>",
usage: "type <selector> <text> [--delay <ms>]",
})?;
Ok(json!({ "id": id, "action": "type", "selector": sel, "text": rest[1..].join(" ") }))
let (text, delay) = parse_text_with_optional_delay(
&rest[1..],
"type",
"type <selector> <text> [--delay <ms>]",
)?;
let mut cmd = json!({ "id": id, "action": "type", "selector": sel, "text": text });
if let Some(ms) = delay {
cmd["delay"] = json!(ms);
}
Ok(cmd)
}
"hover" => {
let sel = rest.first().ok_or_else(|| ParseError::MissingArguments {
@@ -275,14 +376,16 @@ pub fn parse_command(args: &[String], flags: &Flags) -> Result<Value, ParseError
})?;
match *sub {
"type" => {
let text: String = rest[1..].join(" ");
if text.is_empty() {
return Err(ParseError::MissingArguments {
context: "keyboard type".to_string(),
usage: "keyboard type <text>",
});
let (text, delay) = parse_text_with_optional_delay(
&rest[1..],
"keyboard type",
"keyboard type <text> [--delay <ms>]",
)?;
let mut cmd = json!({ "id": id, "action": "keyboard", "subaction": "type", "text": text });
if let Some(ms) = delay {
cmd["delay"] = json!(ms);
}
Ok(json!({ "id": id, "action": "keyboard", "subaction": "type", "text": text }))
Ok(cmd)
}
"inserttext" | "insertText" => {
let text: String = rest[1..].join(" ");
@@ -429,8 +532,16 @@ pub fn parse_command(args: &[String], flags: &Flags) -> Result<Value, ParseError
return Ok(cmd);
}
// Default: selector or timeout
// Default: selector, timeout, or range (e.g. 2000-5000)
if let Some(arg) = rest.first() {
// Check for range syntax: "2000-5000"
if let Some((min_str, max_str)) = arg.split_once('-') {
if let (Ok(min), Ok(max)) = (min_str.parse::<u64>(), max_str.parse::<u64>()) {
return Ok(
json!({ "id": id, "action": "wait", "timeout": min, "timeoutMax": max }),
);
}
}
if let Ok(timeout) = arg.parse::<u64>() {
Ok(json!({ "id": id, "action": "wait", "timeout": timeout }))
} else {
@@ -439,7 +550,7 @@ pub fn parse_command(args: &[String], flags: &Flags) -> Result<Value, ParseError
} else {
Err(ParseError::MissingArguments {
context: "wait".to_string(),
usage: "wait <selector|ms|--url|--load|--fn|--text>",
usage: "wait <selector|ms|min-max|--url|--load|--fn|--text>",
})
}
}
@@ -568,149 +679,15 @@ pub fn parse_command(args: &[String], flags: &Flags) -> Result<Value, ParseError
// === Close ===
"close" | "quit" | "exit" => Ok(json!({ "id": id, "action": "close" })),
// === Authentication Vault ===
"auth" => {
let sub = rest.first().map(|s| s.as_ref());
match sub {
Some("save") => {
let name = rest.get(1).ok_or_else(|| ParseError::MissingArguments {
context: "auth save".to_string(),
usage: "agent-browser auth save <name> --url <url> --username <user> --password <pass>",
})?;
let mut url = None;
let mut username = None;
let mut password = None;
let mut password_stdin = false;
let mut username_selector = None;
let mut password_selector = None;
let mut submit_selector = None;
let mut j = 2;
while j < rest.len() {
match rest[j].as_ref() {
"--url" => {
url = rest.get(j + 1).cloned();
j += 1;
}
"--username" => {
username = rest.get(j + 1).cloned();
j += 1;
}
"--password" => {
password = rest.get(j + 1).cloned();
j += 1;
}
"--password-stdin" => {
password_stdin = true;
}
"--username-selector" => {
username_selector = rest.get(j + 1).cloned();
j += 1;
}
"--password-selector" => {
password_selector = rest.get(j + 1).cloned();
j += 1;
}
"--submit-selector" => {
submit_selector = rest.get(j + 1).cloned();
j += 1;
}
other => {
if other.starts_with("--") {
return Err(ParseError::InvalidValue {
message: format!("unknown flag '{}' for auth save", other),
usage: "agent-browser auth save <name> --url <url> --username <user> --password <pass>",
});
}
}
}
j += 1;
}
let url_val = url.ok_or_else(|| ParseError::MissingArguments {
context: "auth save".to_string(),
usage: "agent-browser auth save <name> --url <url> --username <user> --password <pass> [--password-stdin]",
})?;
let user_val = username.ok_or_else(|| ParseError::MissingArguments {
context: "auth save".to_string(),
usage: "agent-browser auth save <name> --url <url> --username <user> --password <pass> [--password-stdin]",
})?;
if !password_stdin && password.is_none() {
return Err(ParseError::MissingArguments {
context: "auth save".to_string(),
usage: "agent-browser auth save <name> --url <url> --username <user> --password <pass> [--password-stdin]",
});
}
let mut cmd = json!({
"id": id,
"action": "auth_save",
"name": name,
"url": url_val,
"username": user_val,
});
if password_stdin {
cmd["passwordStdin"] = json!(true);
}
if let Some(pass_val) = password {
cmd["password"] = json!(pass_val);
}
if let Some(us) = username_selector {
cmd["usernameSelector"] = json!(us);
}
if let Some(ps) = password_selector {
cmd["passwordSelector"] = json!(ps);
}
if let Some(ss) = submit_selector {
cmd["submitSelector"] = json!(ss);
}
Ok(cmd)
}
Some("login") => {
let name = rest.get(1).ok_or_else(|| ParseError::MissingArguments {
context: "auth login".to_string(),
usage: "agent-browser auth login <name>",
})?;
Ok(json!({ "id": id, "action": "auth_login", "name": name }))
}
Some("list") => Ok(json!({ "id": id, "action": "auth_list" })),
Some("delete") | Some("remove") => {
let name = rest.get(1).ok_or_else(|| ParseError::MissingArguments {
context: "auth delete".to_string(),
usage: "agent-browser auth delete <name>",
})?;
Ok(json!({ "id": id, "action": "auth_delete", "name": name }))
}
Some("show") => {
let name = rest.get(1).ok_or_else(|| ParseError::MissingArguments {
context: "auth show".to_string(),
usage: "agent-browser auth show <name>",
})?;
Ok(json!({ "id": id, "action": "auth_show", "name": name }))
}
_ => Err(ParseError::UnknownSubcommand {
subcommand: sub.unwrap_or("(none)").to_string(),
valid_options: &["save", "login", "list", "delete", "show"],
}),
// === Doctor ===
"doctor" => {
if !rest.is_empty() {
return Err(ParseError::InvalidValue {
message: format!("doctor does not accept arguments: {}", rest.join(" ")),
usage: "doctor",
});
}
}
// === Action Confirmation ===
"confirm" => {
let cid = rest.first().ok_or_else(|| ParseError::MissingArguments {
context: "confirm".to_string(),
usage: "agent-browser confirm <confirmation-id>",
})?;
Ok(json!({ "id": id, "action": "confirm", "confirmationId": cid }))
}
"deny" => {
let cid = rest.first().ok_or_else(|| ParseError::MissingArguments {
context: "deny".to_string(),
usage: "agent-browser deny <confirmation-id>",
})?;
Ok(json!({ "id": id, "action": "deny", "confirmationId": cid }))
Ok(json!({ "id": id, "action": "doctor" }))
}
// === Connect (CDP) ===
@@ -888,6 +865,17 @@ pub fn parse_command(args: &[String], flags: &Flags) -> Result<Value, ParseError
}
}
// Playwright requires either `url` or a complete `domain`+`path` pair.
let has_url = cookie.get("url").is_some();
let has_domain = cookie.get("domain").is_some();
let has_path = cookie.get("path").is_some();
if !has_url && (has_domain != has_path) {
return Err(ParseError::MissingArguments {
context: "cookies set".to_string(),
usage: "When not using --url, you must provide both --domain <domain> and --path <path>",
});
}
Ok(json!({ "id": id, "action": "cookies_set", "cookies": [cookie] }))
}
"clear" => Ok(json!({ "id": id, "action": "cookies_clear" })),
@@ -959,6 +947,13 @@ pub fn parse_command(args: &[String], flags: &Flags) -> Result<Value, ParseError
}
Ok(cmd)
}
Some("dismiss") => {
let mut cmd = json!({ "id": id, "action": "dialog", "response": "dismiss" });
if let Some(prompt_text) = rest.get(1) {
cmd["promptText"] = json!(prompt_text);
}
Ok(cmd)
}
Some(sub) => Err(ParseError::UnknownSubcommand {
subcommand: sub.to_string(),
valid_options: VALID,
@@ -2073,11 +2068,11 @@ mod tests {
full: false,
headed: false,
debug: false,
resident: false,
headers: None,
executable_path: None,
extensions: Vec::new(),
cdp: None,
profile: None,
state: None,
proxy: None,
proxy_bypass: None,
@@ -2088,10 +2083,12 @@ mod tests {
allow_file_access: false,
device: None,
auto_connect: false,
native: false,
engine: None,
session_name: None,
parallel: None,
cli_executable_path: false,
cli_extensions: false,
cli_profile: false,
cli_state: false,
cli_args: false,
cli_user_agent: false,
@@ -2100,16 +2097,19 @@ mod tests {
cli_allow_file_access: false,
cli_annotate: false,
cli_download_path: false,
cli_native: false,
cli_engine: false,
annotate: false,
color_scheme: None,
download_path: None,
content_boundaries: false,
max_output: None,
allowed_domains: None,
action_policy: None,
confirm_actions: None,
confirm_interactive: false,
native: false,
tab_group: None,
tab_group_plugin_id: None,
risk_mode: None,
wait_until: None,
cli_tab_group: false,
cli_tab_group_plugin_id: false,
cli_session_name: false,
cli_resident: false,
}
}
@@ -2165,28 +2165,34 @@ mod tests {
}
#[test]
fn test_cookies_set_with_domain() {
let cmd = parse_command(
fn test_cookies_set_with_domain_requires_path() {
let result = parse_command(
&args("cookies set mycookie myvalue --domain example.com"),
&default_flags(),
);
assert!(result.is_err());
}
#[test]
fn test_cookies_set_with_path_requires_domain() {
let result = parse_command(
&args("cookies set mycookie myvalue --path /api"),
&default_flags(),
);
assert!(result.is_err());
}
#[test]
fn test_cookies_set_with_domain_and_path() {
let cmd = parse_command(
&args("cookies set mycookie myvalue --domain example.com --path /api"),
&default_flags(),
)
.unwrap();
assert_eq!(cmd["action"], "cookies_set");
assert_eq!(cmd["cookies"][0]["name"], "mycookie");
assert_eq!(cmd["cookies"][0]["value"], "myvalue");
assert_eq!(cmd["cookies"][0]["domain"], "example.com");
}
#[test]
fn test_cookies_set_with_path() {
let cmd = parse_command(
&args("cookies set mycookie myvalue --path /api"),
&default_flags(),
)
.unwrap();
assert_eq!(cmd["action"], "cookies_set");
assert_eq!(cmd["cookies"][0]["name"], "mycookie");
assert_eq!(cmd["cookies"][0]["value"], "myvalue");
assert_eq!(cmd["cookies"][0]["path"], "/api");
}
@@ -2365,6 +2371,22 @@ mod tests {
assert_eq!(cmd["url"], "https://example.com");
}
#[test]
fn test_start_command_uses_managed_cdp() {
let cmd = parse_command(&args("start"), &default_flags()).unwrap();
assert_eq!(cmd["action"], "launch");
assert_eq!(cmd["cdpPort"], 9333);
assert_eq!(cmd["headless"], true);
}
#[test]
fn test_start_command_respects_headed_flag() {
let mut flags = default_flags();
flags.headed = true;
let cmd = parse_command(&args("start"), &flags).unwrap();
assert_eq!(cmd["headless"], false);
}
#[test]
fn test_navigate_with_headers() {
let mut flags = default_flags();
@@ -2375,6 +2397,22 @@ mod tests {
assert_eq!(cmd["headers"]["Authorization"], "Bearer token");
}
#[test]
fn test_navigate_with_risk_mode() {
let mut flags = default_flags();
flags.risk_mode = Some("block".to_string());
let cmd = parse_command(&args("open https://example.com"), &flags).unwrap();
assert_eq!(cmd["riskMode"], "block");
}
#[test]
fn test_navigate_with_wait_until() {
let mut flags = default_flags();
flags.wait_until = Some("domcontentloaded".to_string());
let cmd = parse_command(&args("open https://example.com"), &flags).unwrap();
assert_eq!(cmd["waitUntil"], "domcontentloaded");
}
#[test]
fn test_navigate_with_multiple_headers() {
let mut flags = default_flags();
@@ -2505,6 +2543,29 @@ mod tests {
assert_eq!(cmd["text"], "some text");
}
#[test]
fn test_type_command_with_delay() {
let cmd =
parse_command(&args("type #input some text --delay 120"), &default_flags()).unwrap();
assert_eq!(cmd["action"], "type");
assert_eq!(cmd["selector"], "#input");
assert_eq!(cmd["text"], "some text");
assert_eq!(cmd["delay"], 120);
}
#[test]
fn test_type_command_with_literal_delay_text() {
let cmd = parse_command(
&args("type #input -- --delay 120 should be typed"),
&default_flags(),
)
.unwrap();
assert_eq!(cmd["action"], "type");
assert_eq!(cmd["selector"], "#input");
assert_eq!(cmd["text"], "--delay 120 should be typed");
assert!(cmd.get("delay").is_none());
}
#[test]
fn test_select() {
let cmd = parse_command(&args("select #menu option1"), &default_flags()).unwrap();
@@ -2684,6 +2745,19 @@ mod tests {
assert_eq!(cmd["selector"], "#element");
}
#[test]
fn test_keyboard_type_with_delay() {
let cmd = parse_command(
&args("keyboard type natural typing --delay 90"),
&default_flags(),
)
.unwrap();
assert_eq!(cmd["action"], "keyboard");
assert_eq!(cmd["subaction"], "type");
assert_eq!(cmd["text"], "natural typing");
assert_eq!(cmd["delay"], 90);
}
#[test]
fn test_wait_timeout() {
let cmd = parse_command(&args("wait 5000"), &default_flags()).unwrap();
@@ -2960,6 +3034,21 @@ mod tests {
assert!(err.format().contains("Invalid base64"));
}
#[test]
fn test_doctor() {
let cmd = parse_command(&args("doctor"), &default_flags()).unwrap();
assert_eq!(cmd["action"], "doctor");
}
#[test]
fn test_doctor_rejects_arguments() {
let result = parse_command(&args("doctor extra"), &default_flags());
assert!(result.is_err());
let err = result.unwrap_err();
assert!(matches!(err, ParseError::InvalidValue { .. }));
assert!(err.format().contains("doctor does not accept arguments"));
}
#[test]
fn test_unknown_command() {
let result = parse_command(&args("unknowncommand"), &default_flags());
+684 -302
View File
File diff suppressed because it is too large Load Diff
+443 -131
View File
@@ -1,4 +1,5 @@
use crate::color;
use crate::validation::is_valid_session_name;
use serde::Deserialize;
use std::env;
use std::fs;
@@ -7,6 +8,8 @@ use std::path::{Path, PathBuf};
const CONFIG_DIR: &str = ".agent-browser";
const CONFIG_FILENAME: &str = "config.json";
const PROJECT_CONFIG_FILENAME: &str = "agent-browser.json";
const DEFAULT_TAB_GROUP: &str = "Agent Browser Stealth";
const DEFAULT_TAB_GROUP_PLUGIN_ID: &str = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
#[derive(Debug, Default, Deserialize)]
#[serde(default, rename_all = "camelCase")]
@@ -19,7 +22,6 @@ pub struct Config {
pub session_name: Option<String>,
pub executable_path: Option<String>,
pub extensions: Option<Vec<String>>,
pub profile: Option<String>,
pub state: Option<String>,
pub proxy: Option<String>,
pub proxy_bypass: Option<String>,
@@ -31,17 +33,17 @@ pub struct Config {
pub allow_file_access: Option<bool>,
pub cdp: Option<String>,
pub auto_connect: Option<bool>,
pub native: Option<bool>,
pub engine: Option<String>,
pub headers: Option<String>,
pub annotate: Option<bool>,
pub color_scheme: Option<String>,
pub download_path: Option<String>,
pub content_boundaries: Option<bool>,
pub max_output: Option<usize>,
pub allowed_domains: Option<Vec<String>>,
pub action_policy: Option<String>,
pub confirm_actions: Option<String>,
pub confirm_interactive: Option<bool>,
pub native: Option<bool>,
pub tab_group: Option<String>,
pub tab_group_plugin_id: Option<String>,
pub risk_mode: Option<String>,
pub wait_until: Option<String>,
pub parallel: Option<String>,
}
impl Config {
@@ -61,7 +63,6 @@ impl Config {
}
(a, b) => b.or(a),
},
profile: other.profile.or(self.profile),
state: other.state.or(self.state),
proxy: other.proxy.or(self.proxy),
proxy_bypass: other.proxy_bypass.or(self.proxy_bypass),
@@ -73,17 +74,17 @@ impl Config {
allow_file_access: other.allow_file_access.or(self.allow_file_access),
cdp: other.cdp.or(self.cdp),
auto_connect: other.auto_connect.or(self.auto_connect),
native: other.native.or(self.native),
engine: other.engine.or(self.engine),
headers: other.headers.or(self.headers),
annotate: other.annotate.or(self.annotate),
color_scheme: other.color_scheme.or(self.color_scheme),
download_path: other.download_path.or(self.download_path),
content_boundaries: other.content_boundaries.or(self.content_boundaries),
max_output: other.max_output.or(self.max_output),
allowed_domains: other.allowed_domains.or(self.allowed_domains),
action_policy: other.action_policy.or(self.action_policy),
confirm_actions: other.confirm_actions.or(self.confirm_actions),
confirm_interactive: other.confirm_interactive.or(self.confirm_interactive),
native: other.native.or(self.native),
tab_group: other.tab_group.or(self.tab_group),
tab_group_plugin_id: other.tab_group_plugin_id.or(self.tab_group_plugin_id),
risk_mode: other.risk_mode.or(self.risk_mode),
wait_until: other.wait_until.or(self.wait_until),
parallel: other.parallel.or(self.parallel),
}
}
}
@@ -130,11 +131,6 @@ fn parse_bool_arg(args: &[String], i: usize) -> (bool, bool) {
/// Extract --config <path> from args before full flag parsing.
/// Returns `Some(Some(path))` if --config <path> found, `Some(None)` if --config
/// was the last arg with no value, `None` if --config not present.
///
/// Only flags that consume a following argument need to be listed here.
/// Boolean flags (--content-boundaries, --confirm-interactive, etc.) are
/// intentionally absent -- they don't take a value, so they can't cause
/// the next argument to be mis-consumed.
fn extract_config_path(args: &[String]) -> Option<Option<String>> {
const FLAGS_WITH_VALUE: &[&str] = &[
"--session",
@@ -153,11 +149,14 @@ fn extract_config_path(args: &[String]) -> Option<Option<String>> {
"--device",
"--session-name",
"--color-scheme",
"--channel",
"--download-path",
"--max-output",
"--allowed-domains",
"--action-policy",
"--confirm-actions",
"--tab-group",
"--tab-group-plugin-id",
"--risk-mode",
"--wait-until",
"--parallel",
"--engine",
];
let mut i = 0;
while i < args.len() {
@@ -209,12 +208,15 @@ pub struct Flags {
pub full: bool,
pub headed: bool,
pub debug: bool,
/// Keep daemon resident and disable idle auto-shutdown.
pub resident: bool,
/// Runtime daemon session channel.
/// Defaults to `default`; when `--parallel <name>` is provided it becomes `parallel-<name>`.
pub session: String,
pub headers: Option<String>,
pub executable_path: Option<String>,
pub cdp: Option<String>,
pub extensions: Vec<String>,
pub profile: Option<String>,
pub state: Option<String>,
pub proxy: Option<String>,
pub proxy_bypass: Option<String>,
@@ -225,23 +227,30 @@ pub struct Flags {
pub allow_file_access: bool,
pub device: Option<String>,
pub auto_connect: bool,
pub native: bool,
/// Browser engine for native local launches. `chrome` is the default.
pub engine: Option<String>,
// Defaults to "default" when unset in default runtime mode.
// In --parallel mode, defaults to None unless explicitly provided on CLI.
pub session_name: Option<String>,
pub annotate: bool,
pub color_scheme: Option<String>,
pub download_path: Option<String>,
pub content_boundaries: bool,
pub max_output: Option<usize>,
pub allowed_domains: Option<Vec<String>>,
pub action_policy: Option<String>,
pub confirm_actions: Option<String>,
pub confirm_interactive: bool,
pub native: bool,
pub tab_group: Option<String>,
pub tab_group_plugin_id: Option<String>,
/// How verification/captcha detections are handled on navigation:
/// `off` (disable), `warn` (retry and warn), `block` (fail fast).
pub risk_mode: Option<String>,
/// Navigation wait strategy passed to navigate/open commands:
/// `load`, `domcontentloaded`, or `networkidle`.
pub wait_until: Option<String>,
/// Parallel execution channel name. When set, commands run in an isolated runtime session.
pub parallel: Option<String>,
// Track which launch-time options were explicitly passed via CLI
// (as opposed to being set only via environment variables)
pub cli_executable_path: bool,
pub cli_extensions: bool,
pub cli_profile: bool,
pub cli_state: bool,
pub cli_args: bool,
pub cli_user_agent: bool,
@@ -250,6 +259,12 @@ pub struct Flags {
pub cli_allow_file_access: bool,
pub cli_annotate: bool,
pub cli_download_path: bool,
pub cli_native: bool,
pub cli_engine: bool,
pub cli_tab_group: bool,
pub cli_tab_group_plugin_id: bool,
pub cli_session_name: bool,
pub cli_resident: bool,
}
pub fn parse_flags(args: &[String]) -> Flags {
@@ -277,19 +292,21 @@ pub fn parse_flags(args: &[String]) -> Flags {
let mut flags = Flags {
json: env_var_is_truthy("AGENT_BROWSER_JSON") || config.json.unwrap_or(false),
full: env_var_is_truthy("AGENT_BROWSER_FULL") || config.full.unwrap_or(false),
headed: env_var_is_truthy("AGENT_BROWSER_HEADED") || config.headed.unwrap_or(false),
headed: match env::var("AGENT_BROWSER_HEADED") {
Ok(val) => !matches!(val.to_lowercase().as_str(), "0" | "false" | "no" | ""),
Err(_) => config.headed.unwrap_or(true),
},
debug: env_var_is_truthy("AGENT_BROWSER_DEBUG") || config.debug.unwrap_or(false),
session: env::var("AGENT_BROWSER_SESSION")
.ok()
.or(config.session)
.unwrap_or_else(|| "default".to_string()),
resident: false,
// --session is disabled for users.
// Runtime session defaults to `default`, and can be isolated with `--parallel`.
session: "default".to_string(),
headers: config.headers,
executable_path: env::var("AGENT_BROWSER_EXECUTABLE_PATH")
.ok()
.or(config.executable_path),
cdp: config.cdp,
extensions,
profile: env::var("AGENT_BROWSER_PROFILE").ok().or(config.profile),
state: env::var("AGENT_BROWSER_STATE").ok().or(config.state),
proxy: env::var("AGENT_BROWSER_PROXY").ok().or(config.proxy),
proxy_bypass: env::var("AGENT_BROWSER_PROXY_BYPASS")
@@ -307,6 +324,8 @@ pub fn parse_flags(args: &[String]) -> Flags {
device: env::var("AGENT_BROWSER_IOS_DEVICE").ok().or(config.device),
auto_connect: env_var_is_truthy("AGENT_BROWSER_AUTO_CONNECT")
|| config.auto_connect.unwrap_or(false),
native: env_var_is_truthy("AGENT_BROWSER_NATIVE") || config.native.unwrap_or(false),
engine: env::var("AGENT_BROWSER_ENGINE").ok().or(config.engine),
session_name: env::var("AGENT_BROWSER_SESSION_NAME")
.ok()
.or(config.session_name),
@@ -317,33 +336,22 @@ pub fn parse_flags(args: &[String]) -> Flags {
download_path: env::var("AGENT_BROWSER_DOWNLOAD_PATH")
.ok()
.or(config.download_path),
content_boundaries: env_var_is_truthy("AGENT_BROWSER_CONTENT_BOUNDARIES")
|| config.content_boundaries.unwrap_or(false),
max_output: env::var("AGENT_BROWSER_MAX_OUTPUT")
tab_group: env::var("AGENT_BROWSER_TAB_GROUP")
.ok()
.and_then(|s| s.parse().ok())
.or(config.max_output),
allowed_domains: env::var("AGENT_BROWSER_ALLOWED_DOMAINS")
.or(config.tab_group)
.or_else(|| Some(DEFAULT_TAB_GROUP.to_string())),
tab_group_plugin_id: env::var("AGENT_BROWSER_TAB_GROUP_PLUGIN_ID")
.ok()
.map(|s| {
s.split(',')
.map(|d| d.trim().to_lowercase())
.filter(|d| !d.is_empty())
.collect()
})
.or(config.allowed_domains),
action_policy: env::var("AGENT_BROWSER_ACTION_POLICY")
.or(config.tab_group_plugin_id)
.or_else(|| Some(DEFAULT_TAB_GROUP_PLUGIN_ID.to_string())),
risk_mode: env::var("AGENT_BROWSER_RISK_MODE")
.ok()
.or(config.action_policy),
confirm_actions: env::var("AGENT_BROWSER_CONFIRM_ACTIONS")
.ok()
.or(config.confirm_actions),
confirm_interactive: env_var_is_truthy("AGENT_BROWSER_CONFIRM_INTERACTIVE")
|| config.confirm_interactive.unwrap_or(false),
native: env_var_is_truthy("AGENT_BROWSER_NATIVE") || config.native.unwrap_or(false),
.or(config.risk_mode)
.map(|s| s.to_ascii_lowercase()),
wait_until: config.wait_until.map(|s| s.to_ascii_lowercase()),
parallel: env::var("AGENT_BROWSER_PARALLEL").ok().or(config.parallel),
cli_executable_path: false,
cli_extensions: false,
cli_profile: false,
cli_state: false,
cli_args: false,
cli_user_agent: false,
@@ -352,6 +360,12 @@ pub fn parse_flags(args: &[String]) -> Flags {
cli_allow_file_access: false,
cli_annotate: false,
cli_download_path: false,
cli_native: false,
cli_engine: false,
cli_tab_group: false,
cli_tab_group_plugin_id: false,
cli_session_name: false,
cli_resident: false,
};
let mut i = 0;
@@ -385,9 +399,11 @@ pub fn parse_flags(args: &[String]) -> Flags {
i += 1;
}
}
"--session" => {
if let Some(s) = args.get(i + 1) {
flags.session = s.clone();
"--resident" => {
let (val, consumed) = parse_bool_arg(args, i);
flags.resident = val;
flags.cli_resident = true;
if consumed {
i += 1;
}
}
@@ -417,13 +433,6 @@ pub fn parse_flags(args: &[String]) -> Flags {
i += 1;
}
}
"--profile" => {
if let Some(s) = args.get(i + 1) {
flags.profile = Some(s.clone());
flags.cli_profile = true;
i += 1;
}
}
"--state" => {
if let Some(s) = args.get(i + 1) {
flags.state = Some(s.clone());
@@ -493,9 +502,31 @@ pub fn parse_flags(args: &[String]) -> Flags {
i += 1;
}
}
"--native" => {
let (val, consumed) = parse_bool_arg(args, i);
flags.native = val;
flags.cli_native = true;
if consumed {
i += 1;
}
}
"--engine" => {
if let Some(s) = args.get(i + 1) {
flags.engine = Some(s.clone());
flags.cli_engine = true;
i += 1;
}
}
"--session-name" => {
if let Some(s) = args.get(i + 1) {
flags.session_name = Some(s.clone());
flags.cli_session_name = true;
i += 1;
}
}
"--parallel" => {
if let Some(s) = args.get(i + 1) {
flags.parallel = Some(s.clone());
i += 1;
}
}
@@ -520,55 +551,29 @@ pub fn parse_flags(args: &[String]) -> Flags {
i += 1;
}
}
"--content-boundaries" => {
let (val, consumed) = parse_bool_arg(args, i);
flags.content_boundaries = val;
if consumed {
i += 1;
}
}
"--max-output" => {
"--tab-group" => {
if let Some(s) = args.get(i + 1) {
if let Ok(n) = s.parse::<usize>() {
flags.max_output = Some(n);
}
flags.tab_group = Some(s.clone());
flags.cli_tab_group = true;
i += 1;
}
}
"--allowed-domains" => {
"--tab-group-plugin-id" => {
if let Some(s) = args.get(i + 1) {
flags.allowed_domains = Some(
s.split(',')
.map(|d| d.trim().to_lowercase())
.filter(|d| !d.is_empty())
.collect(),
);
flags.tab_group_plugin_id = Some(s.clone());
flags.cli_tab_group_plugin_id = true;
i += 1;
}
}
"--action-policy" => {
"--risk-mode" => {
if let Some(s) = args.get(i + 1) {
flags.action_policy = Some(s.clone());
flags.risk_mode = Some(s.to_ascii_lowercase());
i += 1;
}
}
"--confirm-actions" => {
"--wait-until" => {
if let Some(s) = args.get(i + 1) {
flags.confirm_actions = Some(s.clone());
i += 1;
}
}
"--confirm-interactive" => {
let (val, consumed) = parse_bool_arg(args, i);
flags.confirm_interactive = val;
if consumed {
i += 1;
}
}
"--native" => {
let (val, consumed) = parse_bool_arg(args, i);
flags.native = val;
if consumed {
flags.wait_until = Some(s.to_ascii_lowercase());
i += 1;
}
}
@@ -580,6 +585,28 @@ pub fn parse_flags(args: &[String]) -> Flags {
}
i += 1;
}
if let Some(parallel_name) = &flags.parallel {
// Validate early so session id derivation cannot introduce unsafe paths.
if !is_valid_session_name(parallel_name) {
// Keep default session and let main.rs surface a user-facing validation error.
} else {
flags.session = format!("parallel-{}", parallel_name);
}
}
// Parallel mode is for isolated/stateless runs.
// Unless --session-name is explicitly provided on this invocation, disable
// auto save/restore persistence to avoid cross-flow auth leakage.
if flags.parallel.is_some() && !flags.cli_session_name {
flags.session_name = None;
}
// Keep auth/state continuity stable by default for the default runtime session.
if flags.session_name.is_none() && flags.parallel.is_none() {
flags.session_name = Some("default".to_string());
}
flags
}
@@ -593,13 +620,12 @@ pub fn clean_args(args: &[String]) -> Vec<String> {
"--full",
"--headed",
"--debug",
"--resident",
"--ignore-https-errors",
"--allow-file-access",
"--auto-connect",
"--annotate",
"--content-boundaries",
"--confirm-interactive",
"--native",
"--annotate",
];
// Global flags that always take a value (need to skip the next arg too)
const GLOBAL_FLAGS_WITH_VALUE: &[&str] = &[
@@ -608,7 +634,6 @@ pub fn clean_args(args: &[String]) -> Vec<String> {
"--executable-path",
"--cdp",
"--extension",
"--profile",
"--state",
"--proxy",
"--proxy-bypass",
@@ -620,11 +645,13 @@ pub fn clean_args(args: &[String]) -> Vec<String> {
"--session-name",
"--color-scheme",
"--download-path",
"--max-output",
"--allowed-domains",
"--action-policy",
"--confirm-actions",
"--tab-group",
"--tab-group-plugin-id",
"--risk-mode",
"--wait-until",
"--parallel",
"--config",
"--engine",
];
let mut i = 0;
@@ -658,6 +685,36 @@ pub fn clean_args(args: &[String]) -> Vec<String> {
#[cfg(test)]
mod tests {
use super::*;
use std::sync::{Mutex, MutexGuard};
static ENV_MUTEX: Mutex<()> = Mutex::new(());
struct EnvGuard<'a> {
_lock: MutexGuard<'a, ()>,
vars: Vec<(String, Option<String>)>,
}
impl<'a> EnvGuard<'a> {
fn new(var_names: &[&str]) -> Self {
let lock = ENV_MUTEX.lock().unwrap();
let vars = var_names
.iter()
.map(|&name| (name.to_string(), env::var(name).ok()))
.collect();
Self { _lock: lock, vars }
}
}
impl Drop for EnvGuard<'_> {
fn drop(&mut self) {
for (name, value) in &self.vars {
match value {
Some(v) => env::set_var(name, v),
None => env::remove_var(name),
}
}
}
}
fn args(s: &str) -> Vec<String> {
s.split_whitespace().map(String::from).collect()
@@ -765,12 +822,62 @@ mod tests {
}
#[test]
fn test_parse_flags_with_session_and_executable_path() {
fn test_parse_flags_ignores_session_flag_and_keeps_default_session() {
let flags = parse_flags(&args(
"--session test --executable-path /custom/chrome open example.com",
));
assert_eq!(flags.session, "test");
assert_eq!(flags.session, "default");
assert_eq!(flags.executable_path, Some("/custom/chrome".to_string()));
assert_eq!(flags.session_name.as_deref(), Some("default"));
}
#[test]
fn test_session_name_defaults_to_default_when_not_provided() {
let flags = parse_flags(&args("--session my-session snapshot"));
assert_eq!(flags.session_name.as_deref(), Some("default"));
}
#[test]
fn test_invalid_session_still_uses_default_session_name() {
let flags = parse_flags(&args("--session bad/session snapshot"));
assert_eq!(flags.session_name.as_deref(), Some("default"));
}
#[test]
fn test_env_session_is_ignored_and_default_session_is_used() {
let _guard = EnvGuard::new(&["AGENT_BROWSER_SESSION"]);
env::set_var("AGENT_BROWSER_SESSION", "legacy-session");
let flags = parse_flags(&args("snapshot"));
assert_eq!(flags.session, "default");
assert_eq!(flags.session_name.as_deref(), Some("default"));
}
#[test]
fn test_parallel_sets_isolated_runtime_session() {
let flags = parse_flags(&args("--parallel worker_a snapshot"));
assert_eq!(flags.parallel.as_deref(), Some("worker_a"));
assert_eq!(flags.session, "parallel-worker_a");
assert_eq!(flags.session_name, None);
}
#[test]
fn test_parallel_keeps_explicit_session_name() {
let flags = parse_flags(&args(
"--parallel worker_b --session-name keep-state snapshot",
));
assert_eq!(flags.session, "parallel-worker_b");
assert_eq!(flags.session_name.as_deref(), Some("keep-state"));
assert!(flags.cli_session_name);
}
#[test]
fn test_parallel_from_env_sets_runtime_session() {
let _guard = EnvGuard::new(&["AGENT_BROWSER_PARALLEL", "AGENT_BROWSER_SESSION_NAME"]);
env::set_var("AGENT_BROWSER_PARALLEL", "envworker");
env::set_var("AGENT_BROWSER_SESSION_NAME", "persisted");
let flags = parse_flags(&args("snapshot"));
assert_eq!(flags.session, "parallel-envworker");
assert_eq!(flags.session_name, None);
}
#[test]
@@ -795,12 +902,6 @@ mod tests {
assert!(flags.cli_extensions);
}
#[test]
fn test_cli_profile_tracking() {
let flags = parse_flags(&args("--profile /path/to/profile snapshot"));
assert!(flags.cli_profile);
}
#[test]
fn test_cli_annotate_tracking() {
let flags = parse_flags(&args("--annotate screenshot"));
@@ -814,6 +915,20 @@ mod tests {
assert!(!flags.cli_annotate);
}
#[test]
fn test_parse_resident_flag() {
let flags = parse_flags(&args("--resident open example.com"));
assert!(flags.resident);
assert!(flags.cli_resident);
}
#[test]
fn test_parse_resident_false() {
let flags = parse_flags(&args("--resident false open example.com"));
assert!(!flags.resident);
assert!(flags.cli_resident);
}
#[test]
fn test_cli_download_path_tracking() {
let flags = parse_flags(&args("--download-path /tmp/dl snapshot"));
@@ -827,13 +942,146 @@ mod tests {
assert!(!flags.cli_download_path);
}
#[test]
fn test_default_tab_group_is_enabled() {
let flags = parse_flags(&args("snapshot"));
assert_eq!(flags.tab_group.as_deref(), Some(DEFAULT_TAB_GROUP));
assert!(!flags.cli_tab_group);
}
#[test]
fn test_default_tab_group_plugin_id_is_enabled() {
let flags = parse_flags(&args("snapshot"));
assert_eq!(
flags.tab_group_plugin_id.as_deref(),
Some(DEFAULT_TAB_GROUP_PLUGIN_ID)
);
assert!(!flags.cli_tab_group_plugin_id);
}
#[test]
fn test_parse_tab_group_flag() {
let input = vec![
"--tab-group".to_string(),
"Agent Browser Stealth".to_string(),
"snapshot".to_string(),
];
let flags = parse_flags(&input);
assert_eq!(flags.tab_group.as_deref(), Some("Agent Browser Stealth"));
assert!(flags.cli_tab_group);
}
#[test]
fn test_clean_args_removes_tab_group() {
let cleaned = clean_args(&args("--tab-group AgentGroup open example.com"));
assert_eq!(cleaned, vec!["open", "example.com"]);
}
#[test]
fn test_parse_tab_group_plugin_id_flag() {
let input = vec![
"--tab-group-plugin-id".to_string(),
"cli-plugin-id".to_string(),
"snapshot".to_string(),
];
let flags = parse_flags(&input);
assert_eq!(flags.tab_group_plugin_id.as_deref(), Some("cli-plugin-id"));
assert!(flags.cli_tab_group_plugin_id);
}
#[test]
fn test_clean_args_removes_tab_group_plugin_id() {
let cleaned = clean_args(&args(
"--tab-group-plugin-id cli-plugin-id open example.com",
));
assert_eq!(cleaned, vec!["open", "example.com"]);
}
#[test]
fn test_tab_group_plugin_id_precedence_config_env_cli() {
use std::io::Write;
let _guard = EnvGuard::new(&["AGENT_BROWSER_TAB_GROUP_PLUGIN_ID"]);
let dir = std::env::temp_dir().join("ab-test-plugin-id-precedence");
let _ = fs::create_dir_all(&dir);
let config_path = dir.join("config.json");
let mut f = fs::File::create(&config_path).unwrap();
writeln!(f, r#"{{"tabGroupPluginId":"config-plugin-id"}}"#).unwrap();
env::set_var("AGENT_BROWSER_TAB_GROUP_PLUGIN_ID", "env-plugin-id");
let env_args = vec![
"--config".to_string(),
config_path.to_string_lossy().to_string(),
"snapshot".to_string(),
];
let flags_from_env = parse_flags(&env_args);
assert_eq!(
flags_from_env.tab_group_plugin_id.as_deref(),
Some("env-plugin-id")
);
let cli_args = vec![
"--config".to_string(),
config_path.to_string_lossy().to_string(),
"--tab-group-plugin-id".to_string(),
"cli-plugin-id".to_string(),
"snapshot".to_string(),
];
let flags_from_cli = parse_flags(&cli_args);
assert_eq!(
flags_from_cli.tab_group_plugin_id.as_deref(),
Some("cli-plugin-id")
);
assert!(flags_from_cli.cli_tab_group_plugin_id);
let _ = fs::remove_file(&config_path);
let _ = fs::remove_dir(&dir);
}
#[test]
fn test_parse_risk_mode_flag() {
let flags = parse_flags(&args("--risk-mode block open example.com"));
assert_eq!(flags.risk_mode.as_deref(), Some("block"));
}
#[test]
fn test_clean_args_removes_risk_mode() {
let cleaned = clean_args(&args("--risk-mode warn open example.com"));
assert_eq!(cleaned, vec!["open", "example.com"]);
}
#[test]
fn test_parse_wait_until_flag() {
let flags = parse_flags(&args("--wait-until domcontentloaded open example.com"));
assert_eq!(flags.wait_until.as_deref(), Some("domcontentloaded"));
}
#[test]
fn test_clean_args_removes_wait_until() {
let cleaned = clean_args(&args("--wait-until networkidle open example.com"));
assert_eq!(cleaned, vec!["open", "example.com"]);
}
#[test]
fn test_clean_args_removes_parallel() {
let cleaned = clean_args(&args("--parallel worker_x open example.com"));
assert_eq!(cleaned, vec!["open", "example.com"]);
}
#[test]
fn test_clean_args_removes_resident_flag() {
let cleaned = clean_args(&args("--resident open example.com"));
assert_eq!(cleaned, vec!["open", "example.com"]);
}
#[test]
fn test_cli_multiple_flags_tracking() {
let flags = parse_flags(&args(
"--executable-path /chrome --profile /profile --proxy http://proxy snapshot",
"--executable-path /chrome --proxy http://proxy snapshot",
));
assert!(flags.cli_executable_path);
assert!(flags.cli_profile);
assert!(flags.cli_proxy);
assert!(!flags.cli_extensions);
assert!(!flags.cli_state);
@@ -852,7 +1100,6 @@ mod tests {
"sessionName": "my-app",
"executablePath": "/usr/bin/chromium",
"extensions": ["/ext1", "/ext2"],
"profile": "/tmp/profile",
"state": "/tmp/state.json",
"proxy": "http://proxy:8080",
"proxyBypass": "localhost",
@@ -864,7 +1111,11 @@ mod tests {
"allowFileAccess": true,
"cdp": "9222",
"autoConnect": true,
"headers": "{\"Auth\":\"token\"}"
"headers": "{\"Auth\":\"token\"}",
"tabGroup": "Agent Browser Stealth",
"tabGroupPluginId": "tab-group-plugin-id",
"riskMode": "block",
"parallel": "worker-c"
}"#;
let config: Config = serde_json::from_str(json).unwrap();
assert_eq!(config.headed, Some(true));
@@ -878,7 +1129,6 @@ mod tests {
config.extensions,
Some(vec!["/ext1".to_string(), "/ext2".to_string()])
);
assert_eq!(config.profile.as_deref(), Some("/tmp/profile"));
assert_eq!(config.state.as_deref(), Some("/tmp/state.json"));
assert_eq!(config.proxy.as_deref(), Some("http://proxy:8080"));
assert_eq!(config.proxy_bypass.as_deref(), Some("localhost"));
@@ -891,6 +1141,13 @@ mod tests {
assert_eq!(config.cdp.as_deref(), Some("9222"));
assert_eq!(config.auto_connect, Some(true));
assert_eq!(config.headers.as_deref(), Some("{\"Auth\":\"token\"}"));
assert_eq!(config.tab_group.as_deref(), Some("Agent Browser Stealth"));
assert_eq!(
config.tab_group_plugin_id.as_deref(),
Some("tab-group-plugin-id")
);
assert_eq!(config.risk_mode.as_deref(), Some("block"));
assert_eq!(config.parallel.as_deref(), Some("worker-c"));
}
#[test]
@@ -924,7 +1181,6 @@ mod tests {
let user = Config {
headed: Some(true),
proxy: Some("http://user-proxy:8080".to_string()),
profile: Some("/user/profile".to_string()),
..Config::default()
};
let project = Config {
@@ -935,7 +1191,6 @@ mod tests {
let merged = user.merge(project);
assert_eq!(merged.headed, Some(true)); // kept from user
assert_eq!(merged.proxy.as_deref(), Some("http://project-proxy:9090")); // overridden by project
assert_eq!(merged.profile.as_deref(), Some("/user/profile")); // kept from user
assert_eq!(merged.debug, Some(true)); // added by project
}
@@ -1020,6 +1275,12 @@ mod tests {
assert_eq!(cleaned, vec!["open", "example.com"]);
}
#[test]
fn test_clean_args_removes_engine() {
let cleaned = clean_args(&args("--engine lightpanda open example.com"));
assert_eq!(cleaned, vec!["open", "example.com"]);
}
#[test]
fn test_load_config_with_config_flag() {
use std::io::Write;
@@ -1129,6 +1390,57 @@ mod tests {
assert!(!flags.auto_connect);
}
#[test]
fn test_native_false() {
let flags = parse_flags(&args("--native false open example.com"));
assert!(!flags.native);
assert!(flags.cli_native);
}
#[test]
fn test_engine_flag() {
let flags = parse_flags(&args("--engine lightpanda open example.com"));
assert_eq!(flags.engine.as_deref(), Some("lightpanda"));
assert!(flags.cli_engine);
}
#[test]
fn test_engine_from_env() {
let _guard = EnvGuard::new(&["AGENT_BROWSER_ENGINE"]);
env::set_var("AGENT_BROWSER_ENGINE", "lightpanda");
let flags = parse_flags(&args("open example.com"));
assert_eq!(flags.engine.as_deref(), Some("lightpanda"));
assert!(!flags.cli_engine);
}
#[test]
fn test_native_bare_defaults_true() {
let flags = parse_flags(&args("--native open example.com"));
assert!(flags.native);
assert!(flags.cli_native);
}
#[test]
fn test_native_from_env_sets_native_without_cli_marker() {
let _guard = EnvGuard::new(&["AGENT_BROWSER_NATIVE"]);
env::set_var("AGENT_BROWSER_NATIVE", "1");
let flags = parse_flags(&args("open example.com"));
assert!(flags.native);
assert!(!flags.cli_native);
}
#[test]
fn test_config_deserializes_native() {
let config: Config = serde_json::from_str(r#"{"native": true}"#).unwrap();
assert_eq!(config.native, Some(true));
}
#[test]
fn test_config_deserializes_engine() {
let config: Config = serde_json::from_str(r#"{"engine": "lightpanda"}"#).unwrap();
assert_eq!(config.engine.as_deref(), Some("lightpanda"));
}
#[test]
fn test_full_bare_defaults_true() {
let flags = parse_flags(&args("--full open example.com"));
+260 -345
View File
@@ -5,135 +5,21 @@ mod flags;
mod install;
mod native;
mod output;
#[cfg(test)]
mod test_utils;
mod validation;
use serde_json::json;
use std::env;
use std::fs;
use std::net::{SocketAddr, TcpStream};
use std::process::exit;
#[cfg(windows)]
use windows_sys::Win32::Foundation::CloseHandle;
#[cfg(windows)]
use windows_sys::Win32::System::Threading::{OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION};
use std::time::Duration;
use commands::{gen_id, parse_command, ParseError};
use connection::{ensure_daemon, get_socket_dir, send_command, DaemonOptions};
use flags::{clean_args, parse_flags};
use connection::{ensure_daemon, list_live_sessions, send_command, Response};
use flags::{clean_args, parse_flags, Flags};
use install::run_install;
use output::{
print_command_help, print_help, print_response_with_opts, print_version, OutputOptions,
};
use std::path::PathBuf;
use std::process::Command as ProcessCommand;
/// Run a local auth command (auth_save/list/show/delete) via node auth-cli.js.
/// These commands don't need a browser, so we handle them directly to avoid
/// sending passwords through the daemon's Unix socket channel.
fn run_auth_cli(cmd: &serde_json::Value, json_mode: bool) -> ! {
let exe_path = env::current_exe().unwrap_or_default();
let exe_path = exe_path.canonicalize().unwrap_or(exe_path);
#[cfg(windows)]
let exe_path = {
let p = exe_path.to_string_lossy();
if let Some(stripped) = p.strip_prefix(r"\\?\") {
PathBuf::from(stripped)
} else {
exe_path
}
};
let exe_dir = exe_path.parent().unwrap_or(std::path::Path::new("."));
let mut script_paths = vec![
exe_dir.join("auth-cli.js"),
exe_dir.join("../dist/auth-cli.js"),
PathBuf::from("dist/auth-cli.js"),
];
if let Ok(home) = env::var("AGENT_BROWSER_HOME") {
let home_path = PathBuf::from(&home);
script_paths.insert(0, home_path.join("dist/auth-cli.js"));
script_paths.insert(1, home_path.join("auth-cli.js"));
}
let script_path = match script_paths.iter().find(|p| p.exists()) {
Some(p) => p.clone(),
None => {
if json_mode {
println!(r#"{{"success":false,"error":"auth-cli.js not found"}}"#);
} else {
eprintln!(
"{} auth-cli.js not found. Set AGENT_BROWSER_HOME or run from project directory.",
color::error_indicator()
);
}
exit(1);
}
};
let cmd_json = serde_json::to_string(cmd).unwrap_or_default();
match ProcessCommand::new("node")
.arg(&script_path)
.arg(&cmd_json)
.output()
{
Ok(output) => {
let stderr = String::from_utf8_lossy(&output.stderr);
if !stderr.is_empty() {
eprint!("{}", stderr);
}
let stdout = String::from_utf8_lossy(&output.stdout);
let stdout = stdout.trim();
if stdout.is_empty() {
if json_mode {
println!(r#"{{"success":false,"error":"No response from auth-cli"}}"#);
} else {
eprintln!("{} No response from auth-cli", color::error_indicator());
}
exit(1);
}
if json_mode {
println!("{}", stdout);
} else {
// Parse the JSON response and use the standard output formatter
match serde_json::from_str::<connection::Response>(stdout) {
Ok(resp) => {
let action = cmd.get("action").and_then(|v| v.as_str());
let opts = OutputOptions {
json: false,
content_boundaries: false,
max_output: None,
};
print_response_with_opts(&resp, action, &opts);
if !resp.success {
exit(1);
}
}
Err(_) => {
println!("{}", stdout);
}
}
}
exit(output.status.code().unwrap_or(0));
}
Err(e) => {
if json_mode {
println!(
r#"{{"success":false,"error":"Failed to run auth-cli: {}"}}"#,
e
);
} else {
eprintln!("{} Failed to run auth-cli: {}", color::error_indicator(), e);
}
exit(1);
}
}
}
use output::{print_command_help, print_help, print_response, print_version};
fn parse_proxy(proxy_str: &str) -> serde_json::Value {
let Some(protocol_end) = proxy_str.find("://") else {
@@ -165,51 +51,35 @@ fn parse_proxy(proxy_str: &str) -> serde_json::Value {
})
}
fn should_try_default_cdp(flags: &Flags, command_name: Option<&str>) -> bool {
!matches!(command_name, Some("close"))
&& flags.cdp.is_none()
&& !flags.auto_connect
&& flags.provider.is_none()
&& flags.executable_path.is_none()
&& flags.state.is_none()
&& flags.proxy.is_none()
&& flags.args.is_none()
&& flags.user_agent.is_none()
&& !flags.ignore_https_errors
&& !flags.allow_file_access
&& flags.extensions.is_empty()
}
fn managed_cdp_port_ready() -> bool {
let addr: SocketAddr = match "127.0.0.1:9333".parse() {
Ok(addr) => addr,
Err(_) => return false,
};
TcpStream::connect_timeout(&addr, Duration::from_millis(300)).is_ok()
}
fn run_session(args: &[String], session: &str, json_mode: bool) {
let subcommand = args.get(1).map(|s| s.as_str());
match subcommand {
Some("list") => {
let socket_dir = get_socket_dir();
let mut sessions: Vec<String> = Vec::new();
if let Ok(entries) = fs::read_dir(&socket_dir) {
for entry in entries.flatten() {
let name = entry.file_name().to_string_lossy().to_string();
// Look for pid files in socket directory
if name.ends_with(".pid") {
let session_name = name.strip_suffix(".pid").unwrap_or("");
if !session_name.is_empty() {
// Check if session is actually running
let pid_path = socket_dir.join(&name);
if let Ok(pid_str) = fs::read_to_string(&pid_path) {
if let Ok(pid) = pid_str.trim().parse::<u32>() {
#[cfg(unix)]
let running = unsafe {
libc::kill(pid as i32, 0) == 0
|| std::io::Error::last_os_error().raw_os_error()
!= Some(libc::ESRCH)
};
#[cfg(windows)]
let running = unsafe {
let handle =
OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, 0, pid);
if handle != 0 {
CloseHandle(handle);
true
} else {
false
}
};
if running {
sessions.push(session_name.to_string());
}
}
}
}
}
}
}
let sessions = list_live_sessions();
if json_mode {
println!(
@@ -242,30 +112,40 @@ fn run_session(args: &[String], session: &str, json_mode: bool) {
}
fn main() {
// Ignore SIGPIPE to prevent panic when piping to head/tail
#[cfg(unix)]
unsafe {
libc::signal(libc::SIGPIPE, libc::SIG_DFL);
}
// Prevent MSYS/Git Bash path translation from mangling arguments
#[cfg(windows)]
{
env::set_var("MSYS_NO_PATHCONV", "1");
env::set_var("MSYS2_ARG_CONV_EXCL", "*");
}
// Native daemon mode: when AGENT_BROWSER_DAEMON is set, run as the daemon process
if env::var("AGENT_BROWSER_DAEMON").is_ok() {
#[cfg(unix)]
unsafe {
libc::signal(libc::SIGPIPE, libc::SIG_IGN);
}
let session = env::var("AGENT_BROWSER_SESSION").unwrap_or_else(|_| "default".to_string());
let rt = tokio::runtime::Runtime::new().expect("Failed to create tokio runtime");
rt.block_on(native::daemon::run_daemon(&session));
return;
}
// Ignore SIGPIPE to prevent panic when piping to head/tail
#[cfg(unix)]
unsafe {
libc::signal(libc::SIGPIPE, libc::SIG_DFL);
}
let args: Vec<String> = env::args().skip(1).collect();
let flags = parse_flags(&args);
let mut flags = parse_flags(&args);
let clean = clean_args(&args);
let command_name = clean.first().map(|s| s.as_str());
if flags.engine.is_some() && !flags.native {
flags.native = true;
}
let can_try_default_cdp = should_try_default_cdp(&flags, command_name);
let can_force_native_for_cdp = !matches!(command_name, Some("close"));
if !flags.native
&& can_force_native_for_cdp
&& (flags.cdp.is_some() || flags.auto_connect || can_try_default_cdp)
{
flags.native = true;
}
let has_help = args.iter().any(|a| a == "--help" || a == "-h");
let has_version = args.iter().any(|a| a == "--version" || a == "-V");
@@ -285,6 +165,80 @@ fn main() {
return;
}
if let Some(ref risk_mode) = flags.risk_mode {
if !matches!(risk_mode.as_str(), "off" | "warn" | "block") {
let msg = format!(
"Invalid --risk-mode value: {} (expected off, warn, or block)",
risk_mode
);
if flags.json {
println!(r#"{{"success":false,"error":"{}"}}"#, msg);
} else {
eprintln!("{} {}", color::error_indicator(), msg);
}
exit(1);
}
}
if let Some(ref parallel) = flags.parallel {
if !validation::is_valid_session_name(parallel) {
let msg = format!(
"Invalid --parallel value '{}'. Only alphanumeric characters, hyphens, and underscores are allowed.",
parallel
);
if flags.json {
println!(
r#"{{"success":false,"error":"{}","type":"invalid_parallel_name"}}"#,
msg.replace('"', "\\\"")
);
} else {
eprintln!("{} {}", color::error_indicator(), msg);
}
exit(1);
}
}
if args.iter().any(|a| a == "--profile") {
let msg =
"Project policy: --profile is forbidden. Use your existing browser and --session-name for state persistence.";
if flags.json {
println!(r#"{{"success":false,"error":"{}"}}"#, msg);
} else {
eprintln!("{} {}", color::error_indicator(), msg);
}
exit(1);
}
if env::var("AGENT_BROWSER_PROFILE").is_ok() {
let msg =
"Project policy: AGENT_BROWSER_PROFILE is forbidden. Remove it and use --session-name.";
if flags.json {
println!(r#"{{"success":false,"error":"{}"}}"#, msg);
} else {
eprintln!("{} {}", color::error_indicator(), msg);
}
exit(1);
}
if args.iter().any(|a| a == "--channel") {
let msg = "Project policy: --channel is forbidden. Browser selection follows your existing browser session.";
if flags.json {
println!(r#"{{"success":false,"error":"{}"}}"#, msg);
} else {
eprintln!("{} {}", color::error_indicator(), msg);
}
exit(1);
}
if env::var("AGENT_BROWSER_CHANNEL").is_ok() {
let msg =
"Project policy: AGENT_BROWSER_CHANNEL is forbidden. Remove it and use your existing browser session.";
if flags.json {
println!(r#"{{"success":false,"error":"{}"}}"#, msg);
} else {
eprintln!("{} {}", color::error_indicator(), msg);
}
exit(1);
}
if clean.is_empty() {
print_help();
return;
@@ -303,7 +257,7 @@ fn main() {
return;
}
let mut cmd = match parse_command(&clean, &flags) {
let cmd = match parse_command(&clean, &flags) {
Ok(c) => c,
Err(e) => {
if flags.json {
@@ -326,48 +280,6 @@ fn main() {
}
};
// Handle --password-stdin for auth save
if cmd.get("action").and_then(|v| v.as_str()) == Some("auth_save") {
if cmd.get("password").is_some() {
eprintln!(
"{} Passwords on the command line may be visible in process listings and shell history. Use --password-stdin instead.",
color::warning_indicator()
);
}
if cmd
.get("passwordStdin")
.and_then(|v| v.as_bool())
.unwrap_or(false)
{
let mut pass = String::new();
if std::io::stdin().read_line(&mut pass).is_err() || pass.is_empty() {
eprintln!(
"{} Failed to read password from stdin",
color::error_indicator()
);
exit(1);
}
let pass = pass.trim_end_matches('\n').trim_end_matches('\r');
if pass.is_empty() {
eprintln!("{} Password from stdin is empty", color::error_indicator());
exit(1);
}
cmd["password"] = json!(pass);
cmd.as_object_mut().unwrap().remove("passwordStdin");
}
}
// Handle local auth commands without starting the daemon.
// These don't need a browser, so we avoid sending passwords through the socket.
if let Some(action) = cmd.get("action").and_then(|v| v.as_str()) {
if matches!(
action,
"auth_save" | "auth_list" | "auth_show" | "auth_delete"
) {
run_auth_cli(&cmd, flags.json);
}
}
// Validate session name before starting daemon
if let Some(ref name) = flags.session_name {
if !validation::is_valid_session_name(name) {
@@ -384,28 +296,29 @@ fn main() {
}
}
let daemon_opts = DaemonOptions {
headed: flags.headed,
executable_path: flags.executable_path.as_deref(),
extensions: &flags.extensions,
args: flags.args.as_deref(),
user_agent: flags.user_agent.as_deref(),
proxy: flags.proxy.as_deref(),
proxy_bypass: flags.proxy_bypass.as_deref(),
ignore_https_errors: flags.ignore_https_errors,
allow_file_access: flags.allow_file_access,
profile: flags.profile.as_deref(),
state: flags.state.as_deref(),
provider: flags.provider.as_deref(),
device: flags.device.as_deref(),
session_name: flags.session_name.as_deref(),
download_path: flags.download_path.as_deref(),
allowed_domains: flags.allowed_domains.as_deref(),
action_policy: flags.action_policy.as_deref(),
confirm_actions: flags.confirm_actions.as_deref(),
native: flags.native,
};
let daemon_result = match ensure_daemon(&flags.session, &daemon_opts) {
let daemon_result = match ensure_daemon(
&flags.session,
flags.headed,
flags.resident,
flags.executable_path.as_deref(),
&flags.extensions,
flags.args.as_deref(),
flags.user_agent.as_deref(),
flags.proxy.as_deref(),
flags.proxy_bypass.as_deref(),
flags.ignore_https_errors,
flags.allow_file_access,
flags.state.as_deref(),
flags.provider.as_deref(),
flags.device.as_deref(),
flags.session_name.as_deref(),
flags.debug,
flags.native,
flags.engine.as_deref(),
flags.download_path.as_deref(),
flags.tab_group.as_deref(),
flags.tab_group_plugin_id.as_deref(),
) {
Ok(result) => result,
Err(e) => {
if flags.json {
@@ -432,11 +345,6 @@ fn main() {
} else {
None
},
if flags.cli_profile {
Some("--profile")
} else {
None
},
if flags.cli_state {
Some("--state")
} else {
@@ -461,7 +369,13 @@ fn main() {
flags.ignore_https_errors.then_some("--ignore-https-errors"),
flags.cli_allow_file_access.then_some("--allow-file-access"),
flags.cli_download_path.then_some("--download-path"),
flags.native.then_some("--native"),
flags.cli_native.then_some("--native"),
flags.cli_engine.then_some("--engine"),
flags.cli_tab_group.then_some("--tab-group"),
flags
.cli_tab_group_plugin_id
.then_some("--tab-group-plugin-id"),
flags.cli_resident.then_some("--resident"),
]
.into_iter()
.flatten()
@@ -527,6 +441,8 @@ fn main() {
exit(1);
}
let mut attached_to_existing_browser = false;
// Auto-connect to existing browser
if flags.auto_connect {
let mut launch_cmd = json!({
@@ -546,6 +462,15 @@ fn main() {
if let Some(ref dp) = flags.download_path {
launch_cmd["downloadPath"] = json!(dp);
}
if let Some(ref engine) = flags.engine {
launch_cmd["engine"] = json!(engine);
}
if let Some(ref tg) = flags.tab_group {
launch_cmd["tabGroup"] = json!(tg);
}
if let Some(ref plugin_id) = flags.tab_group_plugin_id {
launch_cmd["tabGroupPluginId"] = json!(plugin_id);
}
let err = match send_command(launch_cmd, &flags.session) {
Ok(resp) if resp.success => None,
@@ -564,6 +489,8 @@ fn main() {
}
exit(1);
}
attached_to_existing_browser = true;
}
// Connect via CDP if --cdp flag is set
@@ -636,6 +563,15 @@ fn main() {
if let Some(ref dp) = flags.download_path {
launch_cmd["downloadPath"] = json!(dp);
}
if let Some(ref engine) = flags.engine {
launch_cmd["engine"] = json!(engine);
}
if let Some(ref tg) = flags.tab_group {
launch_cmd["tabGroup"] = json!(tg);
}
if let Some(ref plugin_id) = flags.tab_group_plugin_id {
launch_cmd["tabGroupPluginId"] = json!(plugin_id);
}
let err = match send_command(launch_cmd, &flags.session) {
Ok(resp) if resp.success => None,
@@ -654,6 +590,8 @@ fn main() {
}
exit(1);
}
attached_to_existing_browser = true;
}
// Launch with cloud provider if -p flag is set
@@ -667,39 +605,55 @@ fn main() {
if let Some(ref cs) = flags.color_scheme {
launch_cmd["colorScheme"] = json!(cs);
}
if let Some(ref tg) = flags.tab_group {
launch_cmd["tabGroup"] = json!(tg);
}
if let Some(ref plugin_id) = flags.tab_group_plugin_id {
launch_cmd["tabGroupPluginId"] = json!(plugin_id);
}
let err = match send_command(launch_cmd, &flags.session) {
Ok(resp) if resp.success => None,
Ok(resp) => Some(
resp.error
.unwrap_or_else(|| "Provider connection failed".to_string()),
),
Err(e) => Some(e.to_string()),
};
if let Some(msg) = err {
if flags.json {
println!(r#"{{"success":false,"error":"{}"}}"#, msg);
} else {
eprintln!("{} {}", color::error_indicator(), msg);
match send_command(launch_cmd, &flags.session) {
Ok(resp) => {
if !resp.success {
let msg = resp
.error
.unwrap_or_else(|| "Provider connection failed".to_string());
if flags.json {
println!(r#"{{"success":false,"error":"{}"}}"#, msg);
} else {
eprintln!("{} {}", color::error_indicator(), msg);
}
exit(1);
}
}
Err(e) => {
if flags.json {
println!(r#"{{"success":false,"error":"{}"}}"#, e);
} else {
eprintln!("{} {}", color::error_indicator(), e);
}
exit(1);
}
exit(1);
}
}
// Launch headed browser or configure browser options (without CDP or provider)
if (flags.headed
|| flags.executable_path.is_some()
|| flags.profile.is_some()
|| flags.state.is_some()
|| flags.proxy.is_some()
|| flags.args.is_some()
|| flags.user_agent.is_some()
|| flags.ignore_https_errors
|| flags.allow_file_access
|| flags.debug
|| flags.color_scheme.is_some()
|| flags.download_path.is_some())
|| flags.download_path.is_some()
|| flags.engine.is_some())
&& flags.cdp.is_none()
&& flags.provider.is_none()
&& !attached_to_existing_browser
&& !can_try_default_cdp
{
let mut launch_cmd = json!({
"id": gen_id(),
@@ -716,11 +670,6 @@ fn main() {
cmd_obj.insert("executablePath".to_string(), json!(exec_path));
}
// Add profile path if specified
if let Some(ref profile_path) = flags.profile {
cmd_obj.insert("profile".to_string(), json!(profile_path));
}
// Add state path if specified
if let Some(ref state_path) = flags.state {
cmd_obj.insert("storageState".to_string(), json!(state_path));
@@ -766,23 +715,27 @@ fn main() {
if let Some(ref dp) = flags.download_path {
launch_cmd["downloadPath"] = json!(dp);
}
if let Some(ref domains) = flags.allowed_domains {
launch_cmd["allowedDomains"] = json!(domains);
if let Some(ref tg) = flags.tab_group {
launch_cmd["tabGroup"] = json!(tg);
}
if let Some(ref plugin_id) = flags.tab_group_plugin_id {
launch_cmd["tabGroupPluginId"] = json!(plugin_id);
}
match send_command(launch_cmd, &flags.session) {
Ok(resp) if !resp.success => {
// Launch command failed (e.g., invalid state file, profile error)
let error_msg = resp
.error
.unwrap_or_else(|| "Browser launch failed".to_string());
if flags.json {
println!(r#"{{"success":false,"error":"{}"}}"#, error_msg);
} else {
eprintln!("{} {}", color::error_indicator(), error_msg);
Ok(resp) => {
if !resp.success {
// Launch command failed (e.g., invalid state file)
let error_msg = resp
.error
.unwrap_or_else(|| "Browser launch failed".to_string());
if flags.json {
println!(r#"{{"success":false,"error":"{}"}}"#, error_msg);
} else {
eprintln!("{} {}", color::error_indicator(), error_msg);
}
exit(1);
}
exit(1);
}
Err(e) => {
if flags.json {
@@ -796,82 +749,32 @@ fn main() {
}
exit(1);
}
Ok(_) => {
// Launch succeeded
}
}
}
let output_opts = OutputOptions {
json: flags.json,
content_boundaries: flags.content_boundaries,
max_output: flags.max_output,
};
match send_command(cmd.clone(), &flags.session) {
Ok(resp) => {
let success = resp.success;
// Handle interactive confirmation
if flags.confirm_interactive {
if let Some(data) = &resp.data {
if data
.get("confirmation_required")
.and_then(|v| v.as_bool())
.unwrap_or(false)
{
let desc = data
.get("description")
.and_then(|v| v.as_str())
.unwrap_or("unknown action");
let category = data.get("category").and_then(|v| v.as_str()).unwrap_or("");
let cid = data
.get("confirmation_id")
.and_then(|v| v.as_str())
.unwrap_or("");
eprintln!("[agent-browser] Action requires confirmation:");
eprintln!(" {}: {}", category, desc);
eprint!(" Allow? [y/N]: ");
let mut input = String::new();
let approved = if std::io::IsTerminal::is_terminal(&std::io::stdin()) {
std::io::stdin().read_line(&mut input).is_ok()
&& matches!(input.trim().to_lowercase().as_str(), "y" | "yes")
} else {
false
};
let confirm_cmd = if approved {
json!({ "id": gen_id(), "action": "confirm", "confirmationId": cid })
} else {
json!({ "id": gen_id(), "action": "deny", "confirmationId": cid })
};
match send_command(confirm_cmd, &flags.session) {
Ok(r) => {
if !approved {
eprintln!("{} Action denied", color::error_indicator());
exit(1);
}
print_response_with_opts(&r, None, &output_opts);
}
Err(e) => {
eprintln!("{} {}", color::error_indicator(), e);
exit(1);
}
}
return;
}
}
}
// Extract action for context-specific output handling
let action = cmd.get("action").and_then(|v| v.as_str());
print_response_with_opts(&resp, action, &output_opts);
print_response(&resp, flags.json, action);
if !success {
exit(1);
}
}
Err(e) => {
let is_start = command_name == Some("start")
&& cmd.get("action").and_then(|v| v.as_str()) == Some("launch")
&& cmd.get("cdpPort").and_then(|v| v.as_u64()) == Some(9333);
if is_start && managed_cdp_port_ready() {
let resp = Response {
success: true,
data: Some(json!({ "launched": true })),
error: None,
};
print_response(&resp, flags.json, Some("launch"));
return;
}
if flags.json {
println!(r#"{{"success":false,"error":"{}"}}"#, e);
} else {
@@ -939,4 +842,16 @@ mod tests {
assert_eq!(result["username"], "user");
assert_eq!(result["password"], "p@ss:w0rd");
}
#[test]
fn test_should_try_default_cdp_for_open() {
let flags = parse_flags(&[]);
assert!(should_try_default_cdp(&flags, Some("open")));
}
#[test]
fn test_should_not_try_default_cdp_for_close() {
let flags = parse_flags(&[]);
assert!(!should_try_default_cdp(&flags, Some("close")));
}
}
+1
View File
@@ -0,0 +1 @@
include!("main.rs");
+157 -48
View File
@@ -4,7 +4,7 @@ use tokio::sync::broadcast;
use super::auth;
use super::browser::{BrowserManager, WaitUntil};
use super::cdp::chrome::LaunchOptions;
use super::cdp::chrome::{LaunchOptions, MANAGED_CDP_PORT};
use super::cdp::types::{
AttachToTargetParams, AttachToTargetResult, CdpEvent, ConsoleApiCalledEvent,
CreateTargetResult, ExceptionThrownEvent, TargetCreatedEvent, TargetDestroyedEvent,
@@ -167,13 +167,14 @@ impl DaemonState {
if let Ok(te) =
serde_json::from_value::<TargetCreatedEvent>(event.params.clone())
{
if te.target_info.target_type == "page"
if (te.target_info.target_type == "page"
|| te.target_info.target_type == "webview")
&& !te.target_info.url.is_empty()
{
let already_tracked = self
.browser
.as_ref()
.map_or(true, |b| b.has_target(&te.target_info.target_id));
.is_none_or(|b| b.has_target(&te.target_info.target_id));
if !already_tracked {
new_targets.push(te);
}
@@ -443,6 +444,7 @@ pub async fn execute_command(cmd: &Value, state: &mut DaemonState) -> Value {
session_id: attach.session_id,
url: te.target_info.url.clone(),
title: te.target_info.title.clone(),
target_type: te.target_info.target_type.clone(),
});
}
}
@@ -549,16 +551,16 @@ pub async fn execute_command(cmd: &Value, state: &mut DaemonState) -> Value {
}
// WebDriver backend: reject unsupported CDP-only actions
if matches!(state.backend_type, BackendType::WebDriver) {
if WEBDRIVER_UNSUPPORTED_ACTIONS.contains(&action) {
return error_response(
&id,
&format!(
"Action '{}' is not supported on the WebDriver backend",
action
),
);
}
if matches!(state.backend_type, BackendType::WebDriver)
&& WEBDRIVER_UNSUPPORTED_ACTIONS.contains(&action)
{
return error_response(
&id,
&format!(
"Action '{}' is not supported on the WebDriver backend",
action
),
);
}
let result = match action {
@@ -726,8 +728,13 @@ pub async fn execute_command(cmd: &Value, state: &mut DaemonState) -> Value {
async fn auto_launch(state: &mut DaemonState) -> Result<(), String> {
let options = launch_options_from_env();
let engine = env::var("AGENT_BROWSER_ENGINE").ok();
let debug_enabled = env::var("AGENT_BROWSER_DEBUG").as_deref() == Ok("1");
if let Ok(cdp) = env::var("AGENT_BROWSER_CDP") {
if debug_enabled {
eprintln!("[DEBUG] auto_launch: connecting via AGENT_BROWSER_CDP={}", cdp);
}
let mgr = BrowserManager::connect_cdp(&cdp).await?;
state.browser = Some(mgr);
state.subscribe_to_browser_events();
@@ -736,6 +743,9 @@ async fn auto_launch(state: &mut DaemonState) -> Result<(), String> {
}
if env::var("AGENT_BROWSER_AUTO_CONNECT").is_ok() {
if debug_enabled {
eprintln!("[DEBUG] auto_launch: connecting via AGENT_BROWSER_AUTO_CONNECT");
}
let mgr = BrowserManager::connect_auto().await?;
state.browser = Some(mgr);
state.subscribe_to_browser_events();
@@ -743,18 +753,48 @@ async fn auto_launch(state: &mut DaemonState) -> Result<(), String> {
return Ok(());
}
let mgr = BrowserManager::launch(options).await?;
let mgr = if should_auto_launch_managed_cdp(&options) {
if debug_enabled {
eprintln!("[DEBUG] auto_launch: launching managed Chrome on localhost:9333");
}
BrowserManager::launch_managed_cdp(options.executable_path.clone(), !options.headless)
.await?
} else {
if debug_enabled {
eprintln!("[DEBUG] auto_launch: launching local browser from env options");
}
BrowserManager::launch(options, engine.as_deref()).await?
};
state.browser = Some(mgr);
state.subscribe_to_browser_events();
try_auto_restore_state(state).await;
Ok(())
}
fn should_auto_launch_managed_cdp(options: &LaunchOptions) -> bool {
options.proxy.is_none()
&& options.proxy_bypass.is_none()
&& options.profile.is_none()
&& !options.allow_file_access
&& options.args.is_empty()
&& options.extensions.as_ref().is_none_or(|ext| ext.is_empty())
&& options.storage_state.is_none()
&& options.user_agent.is_none()
&& !options.ignore_https_errors
}
fn launch_options_from_env() -> LaunchOptions {
let headed = env::var("AGENT_BROWSER_HEADED")
.map(|v| v == "1" || v == "true")
.unwrap_or(false);
let extensions: Option<Vec<String>> = env::var("AGENT_BROWSER_EXTENSIONS").ok().map(|v| {
v.split([',', '\n'])
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect()
});
LaunchOptions {
headless: !headed,
executable_path: env::var("AGENT_BROWSER_EXECUTABLE_PATH").ok(),
@@ -772,12 +812,7 @@ fn launch_options_from_env() -> LaunchOptions {
.collect()
})
.unwrap_or_default(),
extensions: env::var("AGENT_BROWSER_EXTENSIONS").ok().map(|v| {
v.split([',', '\n'])
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect()
}),
extensions,
storage_state: env::var("AGENT_BROWSER_STATE").ok(),
user_agent: env::var("AGENT_BROWSER_USER_AGENT").ok(),
ignore_https_errors: env::var("AGENT_BROWSER_IGNORE_HTTPS_ERRORS")
@@ -785,6 +820,7 @@ fn launch_options_from_env() -> LaunchOptions {
.unwrap_or(false),
color_scheme: env::var("AGENT_BROWSER_COLOR_SCHEME").ok(),
download_path: env::var("AGENT_BROWSER_DOWNLOAD_PATH").ok(),
remote_debugging_port: None,
}
}
@@ -833,20 +869,17 @@ async fn handle_launch(cmd: &Value, state: &mut DaemonState) -> Result<Value, St
.get("autoConnect")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let engine = cmd
.get("engine")
.and_then(|v| v.as_str())
.map(String::from)
.or_else(|| env::var("AGENT_BROWSER_ENGINE").ok());
// Relaunch logic: check if we can reuse the existing connection
let needs_relaunch = if let Some(ref mgr) = state.browser {
let has_cdp_arg = cdp_url.is_some() || cdp_port.is_some();
let was_cdp = mgr.is_cdp_connection();
if has_cdp_arg != was_cdp {
true
} else if has_cdp_arg && !mgr.is_connection_alive().await {
true
} else if auto_connect && !mgr.is_connection_alive().await {
true
} else {
!mgr.is_connection_alive().await
}
has_cdp_arg != was_cdp || !mgr.is_connection_alive().await
} else {
true
};
@@ -866,6 +899,7 @@ async fn handle_launch(cmd: &Value, state: &mut DaemonState) -> Result<Value, St
.filter_map(|v| v.as_str().map(String::from))
.collect()
});
let profile = cmd.get("profile").and_then(|v| v.as_str());
let storage_state = cmd.get("storageState").and_then(|v| v.as_str());
let allow_file_access = cmd
@@ -895,7 +929,22 @@ async fn handle_launch(cmd: &Value, state: &mut DaemonState) -> Result<Value, St
}
if let Some(port) = cdp_port {
state.browser = Some(BrowserManager::connect_cdp(&port.to_string()).await?);
let headed = !headless;
let port_u16 = u16::try_from(port).map_err(|_| format!("Invalid CDP port: {}", port))?;
let browser = match BrowserManager::connect_cdp(&port.to_string()).await {
Ok(browser) => browser,
Err(err) if port_u16 == MANAGED_CDP_PORT => {
if std::env::var("AGENT_BROWSER_DEBUG").as_deref() == Ok("1") {
eprintln!(
"[DEBUG] Preferred CDP port {} unavailable ({}), launching managed Chrome profile",
MANAGED_CDP_PORT, err
);
}
BrowserManager::launch_managed_cdp(executable_path.clone(), headed).await?
}
Err(err) => return Err(err),
};
state.browser = Some(browser);
state.subscribe_to_browser_events();
return Ok(json!({ "launched": true }));
}
@@ -987,6 +1036,7 @@ async fn handle_launch(cmd: &Value, state: &mut DaemonState) -> Result<Value, St
.get("downloadPath")
.and_then(|v| v.as_str())
.map(String::from),
remote_debugging_port: None,
};
if let Some(ref domains) = cmd
@@ -997,7 +1047,7 @@ async fn handle_launch(cmd: &Value, state: &mut DaemonState) -> Result<Value, St
state.domain_filter = Some(DomainFilter::new(domains));
}
state.browser = Some(BrowserManager::launch(options).await?);
state.browser = Some(BrowserManager::launch(options, engine.as_deref()).await?);
state.subscribe_to_browser_events();
if let Some(ref filter) = state.domain_filter {
@@ -2466,6 +2516,7 @@ async fn handle_recording_start(cmd: &Value, state: &mut DaemonState) -> Result<
session_id: new_session_id.clone(),
url: nav_url.clone(),
title: String::new(),
target_type: "page".to_string(),
});
// Navigate to URL
@@ -2877,7 +2928,12 @@ async fn handle_permissions(cmd: &Value, state: &DaemonState) -> Result<Value, S
async fn handle_dialog(cmd: &Value, state: &DaemonState) -> Result<Value, String> {
let mgr = state.browser.as_ref().ok_or("Browser not launched")?;
let accept = cmd.get("accept").and_then(|v| v.as_bool()).unwrap_or(true);
let accept = cmd
.get("response")
.and_then(|v| v.as_str())
.map(|r| r == "accept")
.or_else(|| cmd.get("accept").and_then(|v| v.as_bool()))
.unwrap_or(true);
let prompt_text = cmd.get("promptText").and_then(|v| v.as_str());
mgr.handle_dialog(accept, prompt_text).await?;
@@ -3217,12 +3273,7 @@ async fn handle_frame(cmd: &Value, state: &mut DaemonState) -> Result<Value, Str
.send_command_no_params("Page.getFrameTree", Some(&session_id))
.await?;
fn find_frame(
tree: &Value,
selector: Option<&str>,
name: Option<&str>,
url: Option<&str>,
) -> Option<String> {
fn find_frame(tree: &Value, name: Option<&str>, url: Option<&str>) -> Option<String> {
let frame = tree.get("frame")?;
let frame_name = frame.get("name").and_then(|v| v.as_str()).unwrap_or("");
let frame_url = frame.get("url").and_then(|v| v.as_str()).unwrap_or("");
@@ -3241,7 +3292,7 @@ async fn handle_frame(cmd: &Value, state: &mut DaemonState) -> Result<Value, Str
if let Some(children) = tree.get("childFrames").and_then(|v| v.as_array()) {
for child in children {
if let Some(id) = find_frame(child, selector, name, url) {
if let Some(id) = find_frame(child, name, url) {
return Some(id);
}
}
@@ -3266,13 +3317,13 @@ async fn handle_frame(cmd: &Value, state: &mut DaemonState) -> Result<Value, Str
);
let result = mgr.evaluate(&js, None).await?;
let frame_name = result.as_str().ok_or("Could not find frame for selector")?;
if let Some(frame_id) = find_frame(frame_tree, None, Some(frame_name), None) {
if let Some(frame_id) = find_frame(frame_tree, Some(frame_name), None) {
state.active_frame_id = Some(frame_id);
return Ok(json!({ "frame": frame_name }));
}
}
if let Some(frame_id) = find_frame(frame_tree, selector, name, url) {
if let Some(frame_id) = find_frame(frame_tree, name, url) {
let label = name.or(url).unwrap_or("frame");
state.active_frame_id = Some(frame_id);
return Ok(json!({ "frame": label }));
@@ -4000,14 +4051,13 @@ async fn handle_waitfordownload(cmd: &Value, state: &DaemonState) -> Result<Valu
Ok(Ok(event)) => {
if event.method == "Page.downloadProgress"
&& event.session_id.as_deref() == Some(&session_id)
&& event.params.get("state").and_then(|v| v.as_str()) == Some("completed")
{
if event.params.get("state").and_then(|v| v.as_str()) == Some("completed") {
let path = cmd
.get("path")
.and_then(|v| v.as_str())
.unwrap_or("download");
return Ok(json!({ "path": path }));
}
let path = cmd
.get("path")
.and_then(|v| v.as_str())
.unwrap_or("download");
return Ok(json!({ "path": path }));
}
}
Ok(Err(_)) => return Err("Event stream closed".to_string()),
@@ -4056,6 +4106,7 @@ async fn handle_window_new(cmd: &Value, state: &mut DaemonState) -> Result<Value
session_id: attach.session_id,
url: "about:blank".to_string(),
title: String::new(),
target_type: "page".to_string(),
});
if let Some(viewport) = cmd.get("viewport") {
@@ -5124,6 +5175,39 @@ fn error_response(id: &str, error: &str) -> Value {
#[cfg(test)]
mod tests {
use super::*;
use crate::test_utils::EnvGuard;
const ENCRYPTION_KEY_ENV: &str = "AGENT_BROWSER_ENCRYPTION_KEY";
struct TestKeyGuard {
_lock: std::sync::MutexGuard<'static, ()>,
original: Option<String>,
}
impl TestKeyGuard {
fn new() -> Self {
let lock = super::auth::AUTH_TEST_MUTEX
.lock()
.unwrap_or_else(|e| e.into_inner());
let original = std::env::var(ENCRYPTION_KEY_ENV).ok();
// SAFETY: AUTH_TEST_MUTEX serializes all test access so no concurrent mutation.
unsafe { std::env::set_var(ENCRYPTION_KEY_ENV, "a".repeat(64)) };
Self {
_lock: lock,
original,
}
}
}
impl Drop for TestKeyGuard {
fn drop(&mut self) {
// SAFETY: AUTH_TEST_MUTEX is held via _lock.
match &self.original {
Some(val) => unsafe { std::env::set_var(ENCRYPTION_KEY_ENV, val) },
None => unsafe { std::env::remove_var(ENCRYPTION_KEY_ENV) },
}
}
}
#[test]
fn test_success_response_structure() {
@@ -5160,6 +5244,30 @@ mod tests {
assert!(!opts.allow_file_access);
}
#[test]
fn test_launch_options_from_env_headed_flag() {
let _guard = EnvGuard::new(&["AGENT_BROWSER_HEADED"]);
_guard.set("AGENT_BROWSER_HEADED", "1");
let opts = launch_options_from_env();
assert!(
!opts.headless,
"AGENT_BROWSER_HEADED=1 should set headless=false"
);
}
#[test]
fn test_should_auto_launch_managed_cdp_for_bare_defaults() {
let opts = launch_options_from_env();
assert!(should_auto_launch_managed_cdp(&opts));
}
#[test]
fn test_should_not_auto_launch_managed_cdp_when_proxy_is_set() {
let mut opts = launch_options_from_env();
opts.proxy = Some("http://127.0.0.1:8080".to_string());
assert!(!should_auto_launch_managed_cdp(&opts));
}
#[tokio::test]
async fn test_execute_unknown_command() {
let mut state = DaemonState::new();
@@ -5209,6 +5317,7 @@ mod tests {
#[tokio::test]
async fn test_credentials_roundtrip_via_actions() {
let _key_guard = TestKeyGuard::new();
let mut state = DaemonState::new();
let set_cmd = json!({
+310 -71
View File
@@ -1,11 +1,12 @@
use aes_gcm::{aead::Aead, aead::KeyInit, Aes256Gcm};
use base64::{engine::general_purpose::STANDARD, Engine};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::fs;
use std::path::PathBuf;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AuthProfile {
pub name: String,
pub url: String,
@@ -17,6 +18,10 @@ pub struct AuthProfile {
pub password_selector: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub submit_selector: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub created_at: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub last_login_at: Option<String>,
}
// Keep legacy Credential alias for backward compatibility
@@ -48,79 +53,202 @@ fn get_profile_path(name: &str) -> PathBuf {
get_auth_dir().join(format!("{}.json", name))
}
fn derive_encryption_key() -> Vec<u8> {
let hostname = std::env::var("HOSTNAME")
.or_else(|_| std::env::var("COMPUTERNAME"))
.unwrap_or_else(|_| {
#[cfg(unix)]
{
let mut buf = [0u8; 256];
let len = unsafe { libc::gethostname(buf.as_mut_ptr() as *mut _, buf.len()) };
if len == 0 {
let end = buf.iter().position(|&b| b == 0).unwrap_or(buf.len());
String::from_utf8_lossy(&buf[..end]).to_string()
} else {
"unknown-host".to_string()
}
}
#[cfg(not(unix))]
{
"unknown-host".to_string()
}
});
let username = std::env::var("USER")
.or_else(|_| std::env::var("USERNAME"))
.unwrap_or_else(|_| "unknown-user".to_string());
let mut hasher = Sha256::new();
hasher.update(format!("agent-browser:{}:{}", hostname, username).as_bytes());
hasher.finalize().to_vec()
const ENCRYPTION_KEY_ENV: &str = "AGENT_BROWSER_ENCRYPTION_KEY";
const KEY_FILE_NAME: &str = ".encryption-key";
fn get_agent_browser_dir() -> PathBuf {
if let Some(home) = dirs::home_dir() {
home.join(".agent-browser")
} else {
std::env::temp_dir().join("agent-browser")
}
}
fn encrypt_profile(profile: &AuthProfile) -> Result<Vec<u8>, String> {
let key = derive_encryption_key();
fn get_key_file_path() -> PathBuf {
get_agent_browser_dir().join(KEY_FILE_NAME)
}
fn parse_key_hex(hex_str: &str) -> Option<Vec<u8>> {
let hex_str = hex_str.trim();
if hex_str.len() != 64 || !hex_str.chars().all(|c| c.is_ascii_hexdigit()) {
return None;
}
let bytes: Vec<u8> = (0..32)
.map(|i| u8::from_str_radix(&hex_str[i * 2..i * 2 + 2], 16).unwrap())
.collect();
Some(bytes)
}
/// Read the encryption key from AGENT_BROWSER_ENCRYPTION_KEY env var or
/// ~/.agent-browser/.encryption-key file (matching the Node.js implementation).
fn get_encryption_key() -> Result<Vec<u8>, String> {
if let Ok(key_hex) = std::env::var(ENCRYPTION_KEY_ENV) {
return parse_key_hex(&key_hex).ok_or_else(|| {
format!(
"{} should be a 64-character hex string (256 bits). Generate one with: openssl rand -hex 32",
ENCRYPTION_KEY_ENV
)
});
}
let key_file = get_key_file_path();
if key_file.exists() {
let hex = fs::read_to_string(&key_file)
.map_err(|e| format!("Failed to read encryption key file: {}", e))?;
return parse_key_hex(&hex).ok_or_else(|| {
format!(
"Invalid encryption key in {}. Expected 64-character hex string.",
key_file.display()
)
});
}
Err(format!(
"Encryption key required. Set {} or ensure {} exists.",
ENCRYPTION_KEY_ENV,
key_file.display()
))
}
/// Ensure an encryption key exists, auto-generating one if needed.
fn ensure_encryption_key() -> Result<Vec<u8>, String> {
if let Ok(key) = get_encryption_key() {
return Ok(key);
}
let mut key = [0u8; 32];
getrandom::getrandom(&mut key).map_err(|e| format!("Failed to generate key: {}", e))?;
let key_hex = key.iter().map(|b| format!("{:02x}", b)).collect::<String>();
let dir = get_agent_browser_dir();
fs::create_dir_all(&dir).map_err(|e| format!("Failed to create directory: {}", e))?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let _ = fs::set_permissions(&dir, fs::Permissions::from_mode(0o700));
}
let key_file = get_key_file_path();
fs::write(&key_file, format!("{}\n", key_hex))
.map_err(|e| format!("Failed to write encryption key: {}", e))?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let _ = fs::set_permissions(&key_file, fs::Permissions::from_mode(0o600));
}
eprintln!(
"[agent-browser] Auto-generated encryption key at {} -- back up this file or set {}",
key_file.display(),
ENCRYPTION_KEY_ENV
);
Ok(key.to_vec())
}
/// Encrypt a profile to the JSON+base64 format compatible with Node.js.
fn encrypt_profile(profile: &AuthProfile) -> Result<String, String> {
let key = ensure_encryption_key()?;
let cipher =
Aes256Gcm::new_from_slice(&key).map_err(|e| format!("Encryption key error: {}", e))?;
let plaintext = serde_json::to_string(profile)
.map_err(|e| format!("Failed to serialize profile: {}", e))?;
let mut nonce = [0u8; 12];
getrandom::getrandom(&mut nonce).map_err(|e| format!("Failed to generate nonce: {}", e))?;
let ciphertext = cipher
.encrypt(aes_gcm::Nonce::from_slice(&nonce), plaintext.as_bytes())
let mut iv = [0u8; 12];
getrandom::getrandom(&mut iv).map_err(|e| format!("Failed to generate IV: {}", e))?;
// aes_gcm appends the 16-byte auth tag to the ciphertext
let encrypted = cipher
.encrypt(aes_gcm::Nonce::from_slice(&iv), plaintext.as_bytes())
.map_err(|e| format!("Encryption failed: {}", e))?;
let mut result = Vec::with_capacity(12 + ciphertext.len());
result.extend_from_slice(&nonce);
result.extend_from_slice(&ciphertext);
Ok(result)
let tag_offset = encrypted.len() - 16;
let ciphertext = &encrypted[..tag_offset];
let auth_tag = &encrypted[tag_offset..];
let payload = json!({
"version": 1,
"encrypted": true,
"iv": STANDARD.encode(iv),
"authTag": STANDARD.encode(auth_tag),
"data": STANDARD.encode(ciphertext),
});
serde_json::to_string_pretty(&payload)
.map_err(|e| format!("Failed to serialize payload: {}", e))
}
/// JSON envelope written by Node.js encryption (src/encryption.ts).
#[derive(Deserialize)]
struct EncryptedPayload {
#[allow(dead_code)]
version: u32,
#[allow(dead_code)]
encrypted: bool,
iv: String,
#[serde(rename = "authTag")]
auth_tag: String,
data: String,
}
fn decrypt_profile(data: &[u8]) -> Result<AuthProfile, String> {
if data.len() < 13 {
return Err("Encrypted data too short".to_string());
let text = std::str::from_utf8(data).map_err(|_| {
"Profile is not valid UTF-8 -- it may use an older incompatible binary format".to_string()
})?;
if let Ok(payload) = serde_json::from_str::<EncryptedPayload>(text) {
let key = get_encryption_key()?;
let iv = STANDARD
.decode(&payload.iv)
.map_err(|e| format!("Invalid base64 iv: {}", e))?;
let auth_tag = STANDARD
.decode(&payload.auth_tag)
.map_err(|e| format!("Invalid base64 authTag: {}", e))?;
let ciphertext = STANDARD
.decode(&payload.data)
.map_err(|e| format!("Invalid base64 data: {}", e))?;
// aes_gcm expects ciphertext || auth_tag as input to decrypt
let mut combined = Vec::with_capacity(ciphertext.len() + auth_tag.len());
combined.extend_from_slice(&ciphertext);
combined.extend_from_slice(&auth_tag);
let cipher =
Aes256Gcm::new_from_slice(&key).map_err(|e| format!("Decryption key error: {}", e))?;
let plaintext = cipher
.decrypt(aes_gcm::Nonce::from_slice(&iv), combined.as_slice())
.map_err(|e| format!("Decryption failed: {}", e))?;
let json_str = String::from_utf8(plaintext)
.map_err(|e| format!("Decrypted data is not valid UTF-8: {}", e))?;
return serde_json::from_str(&json_str).map_err(|e| format!("Invalid profile data: {}", e));
}
let (nonce_bytes, ciphertext) = data.split_at(12);
let key = derive_encryption_key();
let cipher =
Aes256Gcm::new_from_slice(&key).map_err(|e| format!("Decryption key error: {}", e))?;
let plaintext = cipher
.decrypt(aes_gcm::Nonce::from_slice(nonce_bytes), ciphertext)
.map_err(|e| format!("Decryption failed: {}", e))?;
let json_str = String::from_utf8(plaintext)
.map_err(|e| format!("Decrypted data is not valid UTF-8: {}", e))?;
serde_json::from_str(&json_str).map_err(|e| format!("Invalid profile data: {}", e))
// Fallback: try as plain unencrypted JSON profile
serde_json::from_str::<AuthProfile>(text)
.map_err(|_| "Profile is not a valid encrypted or unencrypted payload".to_string())
}
fn save_profile(profile: &AuthProfile) -> Result<(), String> {
let dir = get_auth_dir();
let _ = fs::create_dir_all(&dir);
fs::create_dir_all(&dir).map_err(|e| format!("Failed to create auth dir: {}", e))?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let _ = fs::set_permissions(&dir, fs::Permissions::from_mode(0o700));
}
let encrypted = encrypt_profile(profile)?;
let encrypted_json = encrypt_profile(profile)?;
let path = get_profile_path(&profile.name);
fs::write(&path, &encrypted).map_err(|e| format!("Failed to write profile: {}", e))
fs::write(&path, &encrypted_json).map_err(|e| format!("Failed to write profile: {}", e))?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let _ = fs::set_permissions(&path, fs::Permissions::from_mode(0o600));
}
Ok(())
}
fn load_profile(name: &str) -> Result<AuthProfile, String> {
@@ -147,6 +275,8 @@ pub fn credentials_set(
username_selector: None,
password_selector: None,
submit_selector: None,
created_at: None,
last_login_at: None,
};
save_profile(&profile)?;
Ok(json!({ "saved": name }))
@@ -170,6 +300,8 @@ pub fn auth_save(
username_selector: username_selector.map(String::from),
password_selector: password_selector.map(String::from),
submit_selector: submit_selector.map(String::from),
created_at: None,
last_login_at: None,
};
save_profile(&profile)?;
Ok(json!({ "saved": name }))
@@ -252,10 +384,27 @@ pub fn auth_show(name: &str) -> Result<Value, String> {
}))
}
#[cfg(test)]
pub(crate) static AUTH_TEST_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(());
#[cfg(test)]
mod tests {
use super::*;
fn with_test_key<F: FnOnce()>(f: F) {
let _lock = AUTH_TEST_MUTEX.lock().unwrap();
let original = std::env::var(ENCRYPTION_KEY_ENV).ok();
let test_key = "a".repeat(64);
// SAFETY: TEST_MUTEX serializes all test access so no concurrent mutation.
unsafe { std::env::set_var(ENCRYPTION_KEY_ENV, &test_key) };
f();
// SAFETY: TEST_MUTEX serializes all test access so no concurrent mutation.
match original {
Some(val) => unsafe { std::env::set_var(ENCRYPTION_KEY_ENV, val) },
None => unsafe { std::env::remove_var(ENCRYPTION_KEY_ENV) },
}
}
#[test]
fn test_validate_profile_name() {
assert!(validate_profile_name("github").is_ok());
@@ -277,6 +426,8 @@ mod tests {
username_selector: None,
password_selector: None,
submit_selector: Some("button[type=submit]".to_string()),
created_at: None,
last_login_at: None,
};
let json = serde_json::to_string(&profile).unwrap();
let parsed: AuthProfile = serde_json::from_str(&json).unwrap();
@@ -290,26 +441,114 @@ mod tests {
#[test]
fn test_encrypt_decrypt_roundtrip() {
let profile = AuthProfile {
name: "roundtrip".to_string(),
url: "https://example.com".to_string(),
username: "user".to_string(),
password: "s3cret!".to_string(),
username_selector: None,
password_selector: None,
submit_selector: None,
};
let encrypted = encrypt_profile(&profile).unwrap();
let decrypted = decrypt_profile(&encrypted).unwrap();
assert_eq!(decrypted.name, "roundtrip");
assert_eq!(decrypted.password, "s3cret!");
with_test_key(|| {
let profile = AuthProfile {
name: "roundtrip".to_string(),
url: "https://example.com".to_string(),
username: "user".to_string(),
password: "s3cret!".to_string(),
username_selector: None,
password_selector: None,
submit_selector: None,
created_at: None,
last_login_at: None,
};
let encrypted_json = encrypt_profile(&profile).unwrap();
let decrypted = decrypt_profile(encrypted_json.as_bytes()).unwrap();
assert_eq!(decrypted.name, "roundtrip");
assert_eq!(decrypted.password, "s3cret!");
});
}
#[test]
fn test_derive_encryption_key_is_stable() {
let k1 = derive_encryption_key();
let k2 = derive_encryption_key();
assert_eq!(k1, k2);
assert_eq!(k1.len(), 32);
fn test_get_encryption_key_from_env() {
with_test_key(|| {
let key = get_encryption_key().unwrap();
assert_eq!(key.len(), 32);
assert!(key.iter().all(|&b| b == 0xaa));
});
}
#[test]
fn test_parse_key_hex_valid() {
let hex = "ab".repeat(32);
let key = parse_key_hex(&hex).unwrap();
assert_eq!(key.len(), 32);
assert!(key.iter().all(|&b| b == 0xab));
}
#[test]
fn test_parse_key_hex_invalid() {
assert!(parse_key_hex("too_short").is_none());
assert!(parse_key_hex(&"g".repeat(64)).is_none());
assert!(parse_key_hex("").is_none());
}
#[test]
fn test_decrypt_json_payload_format() {
with_test_key(|| {
let key = get_encryption_key().unwrap();
let profile = AuthProfile {
name: "json-test".to_string(),
url: "https://example.com/login".to_string(),
username: "admin".to_string(),
password: "hunter2".to_string(),
username_selector: Some("#email".to_string()),
password_selector: None,
submit_selector: None,
created_at: None,
last_login_at: None,
};
// Encrypt with aes_gcm, then manually build the JSON payload
// to simulate what Node.js would produce
let cipher = Aes256Gcm::new_from_slice(&key).unwrap();
let mut iv = [0u8; 12];
getrandom::getrandom(&mut iv).unwrap();
let plaintext = serde_json::to_string(&profile).unwrap();
let encrypted = cipher
.encrypt(aes_gcm::Nonce::from_slice(&iv), plaintext.as_bytes())
.unwrap();
let tag_offset = encrypted.len() - 16;
let ciphertext = &encrypted[..tag_offset];
let auth_tag = &encrypted[tag_offset..];
let payload = format!(
r#"{{"version":1,"encrypted":true,"iv":"{}","authTag":"{}","data":"{}"}}"#,
STANDARD.encode(iv),
STANDARD.encode(auth_tag),
STANDARD.encode(ciphertext),
);
let decrypted = decrypt_profile(payload.as_bytes()).unwrap();
assert_eq!(decrypted.name, "json-test");
assert_eq!(decrypted.password, "hunter2");
assert_eq!(decrypted.username_selector, Some("#email".to_string()));
});
}
#[test]
fn test_encrypted_output_is_json_format() {
with_test_key(|| {
let profile = AuthProfile {
name: "format-check".to_string(),
url: "https://example.com".to_string(),
username: "user".to_string(),
password: "pass".to_string(),
username_selector: None,
password_selector: None,
submit_selector: None,
created_at: None,
last_login_at: None,
};
let encrypted = encrypt_profile(&profile).unwrap();
let parsed: Value = serde_json::from_str(&encrypted).unwrap();
assert_eq!(parsed["version"], 1);
assert_eq!(parsed["encrypted"], true);
assert!(parsed["iv"].is_string());
assert!(parsed["authTag"].is_string());
assert!(parsed["data"].is_string());
});
}
}
+258 -63
View File
@@ -2,11 +2,14 @@ use serde_json::{json, Value};
use std::collections::HashSet;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::time::{timeout, Duration};
use super::cdp::chrome::{
auto_connect_cdp, discover_cdp_url, launch_chrome, ChromeProcess, LaunchOptions,
auto_connect_cdp, discover_cdp_url, launch_chrome, launch_managed_chrome, ChromeProcess,
LaunchOptions,
};
use super::cdp::client::CdpClient;
use super::cdp::lightpanda::{launch_lightpanda, LightpandaLaunchOptions, LightpandaProcess};
use super::cdp::types::*;
// ---------------------------------------------------------------------------
@@ -55,6 +58,34 @@ pub fn validate_launch_options(
Ok(())
}
fn validate_lightpanda_options(options: &LaunchOptions) -> Result<(), String> {
if options
.extensions
.as_ref()
.is_some_and(|exts| !exts.is_empty())
{
return Err("Extensions are not supported with Lightpanda".to_string());
}
if options.profile.is_some() {
return Err("Profiles are not supported with Lightpanda".to_string());
}
if options.storage_state.is_some() {
return Err("Storage state is not supported with Lightpanda".to_string());
}
if options.allow_file_access {
return Err("File access is not supported with Lightpanda".to_string());
}
if !options.headless {
return Err("Headed mode is not supported with Lightpanda (headless only)".to_string());
}
if !options.args.is_empty() {
return Err(
"Custom Chrome arguments (--args) are not supported with Lightpanda".to_string(),
);
}
Ok(())
}
/// Converts common error messages into AI-friendly, actionable descriptions.
pub fn to_ai_friendly_error(error: &str) -> String {
let lower = error.to_lowercase();
@@ -80,12 +111,18 @@ pub fn to_ai_friendly_error(error: &str) -> String {
error.to_string()
}
fn is_startup_internal_page(url: &str) -> bool {
let lower = url.to_lowercase();
lower.starts_with("chrome://profile-picker")
}
#[derive(Debug, Clone)]
pub struct PageInfo {
pub target_id: String,
pub session_id: String,
pub url: String,
pub title: String,
pub target_type: String,
}
#[derive(Debug, Clone, Copy)]
@@ -105,37 +142,74 @@ impl WaitUntil {
}
}
pub enum BrowserProcess {
Chrome(ChromeProcess),
Lightpanda(LightpandaProcess),
}
pub struct BrowserManager {
pub client: CdpClient,
chrome_process: Option<ChromeProcess>,
browser_process: Option<BrowserProcess>,
cdp_connection: bool,
pages: Vec<PageInfo>,
active_page_index: usize,
default_timeout_ms: u64,
}
impl BrowserManager {
pub async fn launch(options: LaunchOptions) -> Result<Self, String> {
validate_launch_options(
options.extensions.as_deref(),
false,
options.profile.as_deref(),
options.storage_state.as_deref(),
options.allow_file_access,
options.executable_path.as_deref(),
)?;
pub async fn launch(options: LaunchOptions, engine: Option<&str>) -> Result<Self, String> {
let engine = engine.unwrap_or("chrome");
match engine {
"chrome" => validate_launch_options(
options.extensions.as_deref(),
false,
options.profile.as_deref(),
options.storage_state.as_deref(),
options.allow_file_access,
options.executable_path.as_deref(),
)?,
"lightpanda" => validate_lightpanda_options(&options)?,
_ => {
return Err(format!(
"Unknown engine '{}'. Supported engines: chrome, lightpanda",
engine
))
}
}
let ignore_https_errors = options.ignore_https_errors;
let user_agent = options.user_agent.clone();
let color_scheme = options.color_scheme.clone();
let download_path = options.download_path.clone();
let chrome = launch_chrome(&options)?;
let ws_url = chrome.ws_url.clone();
let (ws_url, process) = match engine {
"lightpanda" => {
let lp_options = LightpandaLaunchOptions {
executable_path: options.executable_path.clone(),
proxy: options.proxy.clone(),
port: None,
};
let process = tokio::task::spawn_blocking(move || launch_lightpanda(&lp_options))
.await
.map_err(|e| format!("Lightpanda launch task failed: {}", e))??;
let ws_url = process.ws_url.clone();
(ws_url, BrowserProcess::Lightpanda(process))
}
_ => {
let process = tokio::task::spawn_blocking(move || launch_chrome(&options))
.await
.map_err(|e| format!("Chrome launch task failed: {}", e))??;
let ws_url = process.ws_url.clone();
(ws_url, BrowserProcess::Chrome(process))
}
};
let client = CdpClient::connect(&ws_url).await?;
let mut manager = Self {
client,
chrome_process: Some(chrome),
browser_process: Some(process),
cdp_connection: false,
pages: Vec::new(),
active_page_index: 0,
default_timeout_ms: 25_000,
@@ -197,7 +271,32 @@ impl BrowserManager {
let client = CdpClient::connect(&ws_url).await?;
let mut manager = Self {
client,
chrome_process: None,
browser_process: None,
cdp_connection: true,
pages: Vec::new(),
active_page_index: 0,
default_timeout_ms: 10_000,
};
manager.discover_and_attach_targets().await?;
Ok(manager)
}
pub async fn launch_managed_cdp(
executable_path: Option<String>,
headed: bool,
) -> Result<Self, String> {
let process =
tokio::task::spawn_blocking(move || launch_managed_chrome(executable_path, headed))
.await
.map_err(|e| format!("Managed Chrome launch task failed: {}", e))??;
let ws_url = process.ws_url.clone();
let client = CdpClient::connect(&ws_url).await?;
let mut manager = Self {
client,
browser_process: Some(BrowserProcess::Chrome(process)),
cdp_connection: true,
pages: Vec::new(),
active_page_index: 0,
default_timeout_ms: 10_000,
@@ -212,6 +311,50 @@ impl BrowserManager {
Self::connect_cdp(&ws_url).await
}
async fn create_and_attach_blank_page(&mut self) -> Result<(), String> {
let result: CreateTargetResult = self
.client
.send_command_typed(
"Target.createTarget",
&CreateTargetParams {
url: "about:blank".to_string(),
},
None,
)
.await?;
let attach_result: AttachToTargetResult = self
.client
.send_command_typed(
"Target.attachToTarget",
&AttachToTargetParams {
target_id: result.target_id.clone(),
flatten: true,
},
None,
)
.await?;
self.enable_domains_with_timeout(&attach_result.session_id)
.await?;
self.pages = vec![PageInfo {
target_id: result.target_id,
session_id: attach_result.session_id,
url: "about:blank".to_string(),
title: String::new(),
target_type: "page".to_string(),
}];
self.active_page_index = 0;
Ok(())
}
async fn enable_domains_with_timeout(&self, session_id: &str) -> Result<(), String> {
timeout(Duration::from_secs(3), self.enable_domains(session_id))
.await
.map_err(|_| format!("Timed out enabling CDP domains for session {}", session_id))?
}
async fn discover_and_attach_targets(&mut self) -> Result<(), String> {
self.client
.send_command_typed::<_, Value>(
@@ -229,43 +372,15 @@ impl BrowserManager {
let page_targets: Vec<TargetInfo> = result
.target_infos
.into_iter()
.filter(|t| t.target_type == "page" && !t.url.is_empty())
.filter(|t| {
(t.target_type == "page" || t.target_type == "webview") && !t.url.is_empty()
})
.collect();
if page_targets.is_empty() {
// Create a new tab
let result: CreateTargetResult = self
.client
.send_command_typed(
"Target.createTarget",
&CreateTargetParams {
url: "about:blank".to_string(),
},
None,
)
.await?;
let attach_result: AttachToTargetResult = self
.client
.send_command_typed(
"Target.attachToTarget",
&AttachToTargetParams {
target_id: result.target_id.clone(),
flatten: true,
},
None,
)
.await?;
self.pages.push(PageInfo {
target_id: result.target_id,
session_id: attach_result.session_id.clone(),
url: "about:blank".to_string(),
title: String::new(),
});
self.active_page_index = 0;
self.enable_domains(&attach_result.session_id).await?;
self.create_and_attach_blank_page().await?;
} else {
let mut attached_pages = Vec::new();
for target in &page_targets {
let attach_result: AttachToTargetResult = self
.client
@@ -279,17 +394,50 @@ impl BrowserManager {
)
.await?;
self.pages.push(PageInfo {
let page_info = PageInfo {
target_id: target.target_id.clone(),
session_id: attach_result.session_id.clone(),
url: target.url.clone(),
title: target.title.clone(),
});
target_type: target.target_type.clone(),
};
match self
.enable_domains_with_timeout(&attach_result.session_id)
.await
{
Ok(()) => attached_pages.push(page_info),
Err(err) => {
if std::env::var("AGENT_BROWSER_DEBUG").as_deref() == Ok("1") {
eprintln!(
"[DEBUG] Skipping CDP target '{}' ({}): {}",
target.title, target.url, err
);
}
}
}
}
self.active_page_index = 0;
let session_id = self.pages[0].session_id.clone();
self.enable_domains(&session_id).await?;
if attached_pages.is_empty() {
self.create_and_attach_blank_page().await?;
} else {
let preferred_index = attached_pages
.iter()
.position(|page| !is_startup_internal_page(&page.url));
self.pages = attached_pages;
if let Some(index) = preferred_index {
self.active_page_index = index;
} else {
if std::env::var("AGENT_BROWSER_DEBUG").as_deref() == Ok("1") {
eprintln!(
"[DEBUG] All discovered pages were Chrome startup pages; creating a fresh about:blank target"
);
}
self.create_and_attach_blank_page().await?;
}
}
}
Ok(())
@@ -321,6 +469,7 @@ impl BrowserManager {
pub async fn navigate(&mut self, url: &str, wait_until: WaitUntil) -> Result<Value, String> {
let session_id = self.active_session_id()?.to_string();
let rx = self.client.subscribe();
let nav_result: PageNavigateResult = self
.client
@@ -338,7 +487,7 @@ impl BrowserManager {
return Err(format!("Navigation failed: {}", error_text));
}
self.wait_for_lifecycle(wait_until, &session_id).await?;
self.wait_for_lifecycle(wait_until, &session_id, rx).await?;
let page_url = self.get_url().await.unwrap_or_else(|_| url.to_string());
let title = self.get_title().await.unwrap_or_default();
@@ -355,14 +504,14 @@ impl BrowserManager {
&self,
wait_until: WaitUntil,
session_id: &str,
mut rx: tokio::sync::broadcast::Receiver<CdpEvent>,
) -> Result<(), String> {
let event_name = match wait_until {
WaitUntil::Load => "Page.loadEventFired",
WaitUntil::DomContentLoaded => "Page.domContentEventFired",
WaitUntil::NetworkIdle => return self.wait_for_network_idle(session_id).await,
WaitUntil::NetworkIdle => return self.wait_for_network_idle(session_id, rx).await,
};
let mut rx = self.client.subscribe();
let timeout = tokio::time::Duration::from_millis(self.default_timeout_ms);
tokio::time::timeout(timeout, async {
@@ -377,8 +526,11 @@ impl BrowserManager {
.map_err(|_| format!("Timeout waiting for {}", event_name))?
}
async fn wait_for_network_idle(&self, session_id: &str) -> Result<(), String> {
let mut rx = self.client.subscribe();
async fn wait_for_network_idle(
&self,
session_id: &str,
mut rx: tokio::sync::broadcast::Receiver<CdpEvent>,
) -> Result<(), String> {
let pending = Arc::new(Mutex::new(HashSet::<String>::new()));
let timeout = tokio::time::Duration::from_millis(self.default_timeout_ms);
@@ -497,7 +649,8 @@ impl BrowserManager {
wait_until: WaitUntil,
session_id: &str,
) -> Result<(), String> {
self.wait_for_lifecycle(wait_until, session_id).await
self.wait_for_lifecycle(wait_until, session_id, self.client.subscribe())
.await
}
pub async fn close(&mut self) -> Result<(), String> {
@@ -507,9 +660,13 @@ impl BrowserManager {
.send_command_no_params("Browser.close", None)
.await;
// Kill Chrome process if we own it
if let Some(ref mut chrome) = self.chrome_process {
chrome.kill();
if let Some(process) = self.browser_process.take() {
let timeout = std::time::Duration::from_secs(5);
let _ = tokio::task::spawn_blocking(move || match process {
BrowserProcess::Chrome(mut chrome) => chrome.wait_or_kill(timeout),
BrowserProcess::Lightpanda(mut lightpanda) => lightpanda.kill(),
})
.await;
}
Ok(())
@@ -538,7 +695,7 @@ impl BrowserManager {
/// Returns true if this manager was connected via CDP (as opposed to local launch).
pub fn is_cdp_connection(&self) -> bool {
self.chrome_process.is_none()
self.cdp_connection
}
/// Ensures the browser has at least one page. If `pages` is empty, creates a new
@@ -576,6 +733,7 @@ impl BrowserManager {
session_id: attach_result.session_id.clone(),
url: "about:blank".to_string(),
title: String::new(),
target_type: "page".to_string(),
});
self.active_page_index = 0;
self.enable_domains(&attach_result.session_id).await?;
@@ -608,6 +766,7 @@ impl BrowserManager {
"index": i,
"title": p.title,
"url": p.url,
"type": p.target_type,
"active": i == self.active_page_index,
})
})
@@ -648,6 +807,7 @@ impl BrowserManager {
session_id: attach.session_id,
url: target_url.to_string(),
title: String::new(),
target_type: "page".to_string(),
});
self.active_page_index = index;
@@ -1065,6 +1225,30 @@ mod tests {
assert!(validate_launch_options(None, false, None, None, false, None,).is_ok());
}
#[test]
fn test_validate_lightpanda_options_rejects_extensions() {
let opts = LaunchOptions {
extensions: Some(vec!["/tmp/ext".to_string()]),
..Default::default()
};
assert!(validate_lightpanda_options(&opts).is_err());
}
#[test]
fn test_validate_lightpanda_options_rejects_headed() {
let opts = LaunchOptions {
headless: false,
..Default::default()
};
assert!(validate_lightpanda_options(&opts).is_err());
}
#[test]
fn test_validate_lightpanda_options_valid() {
let opts = LaunchOptions::default();
assert!(validate_lightpanda_options(&opts).is_ok());
}
#[test]
fn test_to_ai_friendly_error_strict_mode() {
assert_eq!(
@@ -1110,4 +1294,15 @@ mod tests {
let msg = "Some custom error message";
assert_eq!(to_ai_friendly_error(msg), msg);
}
#[test]
fn test_is_startup_internal_page_detects_profile_picker() {
assert!(is_startup_internal_page("chrome://profile-picker/"));
}
#[test]
fn test_is_startup_internal_page_ignores_normal_pages() {
assert!(!is_startup_internal_page("https://example.com/"));
assert!(!is_startup_internal_page("about:blank"));
}
}
+340 -80
View File
@@ -8,6 +8,7 @@ use super::types::BrowserVersionInfo;
pub struct ChromeProcess {
child: Child,
pub ws_url: String,
temp_user_data_dir: Option<PathBuf>,
}
impl ChromeProcess {
@@ -15,11 +16,46 @@ impl ChromeProcess {
let _ = self.child.kill();
let _ = self.child.wait();
}
/// Wait for Chrome to exit on its own (after Browser.close CDP command),
/// falling back to kill() if it doesn't exit within the timeout.
/// This allows Chrome to flush cookies and other state to the user-data-dir.
pub fn wait_or_kill(&mut self, timeout: Duration) {
let start = std::time::Instant::now();
let poll_interval = Duration::from_millis(50);
while start.elapsed() < timeout {
match self.child.try_wait() {
Ok(Some(_)) => return,
Ok(None) => std::thread::sleep(poll_interval),
Err(_) => break,
}
}
self.kill();
}
}
impl Drop for ChromeProcess {
fn drop(&mut self) {
self.kill();
if let Some(ref dir) = self.temp_user_data_dir {
for attempt in 0..3 {
match std::fs::remove_dir_all(dir) {
Ok(()) => break,
Err(_) if attempt < 2 => {
std::thread::sleep(Duration::from_millis(100));
}
Err(e) => {
eprintln!(
"Warning: failed to clean up temp profile {}: {}",
dir.display(),
e
);
}
}
}
}
}
}
@@ -37,6 +73,7 @@ pub struct LaunchOptions {
pub ignore_https_errors: bool,
pub color_scheme: Option<String>,
pub download_path: Option<String>,
pub remote_debugging_port: Option<u16>,
}
impl Default for LaunchOptions {
@@ -55,20 +92,21 @@ impl Default for LaunchOptions {
ignore_https_errors: false,
color_scheme: None,
download_path: None,
remote_debugging_port: None,
}
}
}
pub fn launch_chrome(options: &LaunchOptions) -> Result<ChromeProcess, String> {
let chrome_path = match &options.executable_path {
Some(p) => PathBuf::from(p),
None => {
find_chrome().ok_or("Chrome not found. Install Chrome or use --executable-path.")?
}
};
struct ChromeArgs {
args: Vec<String>,
temp_user_data_dir: Option<PathBuf>,
}
fn build_chrome_args(options: &LaunchOptions) -> Result<ChromeArgs, String> {
let remote_debugging_port = options.remote_debugging_port.unwrap_or(0);
let mut args = vec![
"--remote-debugging-port=0".to_string(),
format!("--remote-debugging-port={}", remote_debugging_port),
"--remote-debugging-address=127.0.0.1".to_string(),
"--no-first-run".to_string(),
"--no-default-browser-check".to_string(),
"--disable-background-networking".to_string(),
@@ -79,13 +117,21 @@ pub fn launch_chrome(options: &LaunchOptions) -> Result<ChromeProcess, String> {
"--disable-popup-blocking".to_string(),
"--disable-prompt-on-repost".to_string(),
"--disable-sync".to_string(),
"--disable-features=Translate".to_string(),
"--enable-features=NetworkService,NetworkServiceInProcess".to_string(),
"--metrics-recording-only".to_string(),
"--password-store=basic".to_string(),
"--use-mock-keychain".to_string(),
];
if options.headless {
let has_extensions = options
.extensions
.as_ref()
.is_some_and(|exts| !exts.is_empty());
// Extensions require headed mode in native Chrome (content scripts are not
// injected in headless mode). Skip --headless when extensions are loaded.
if options.headless && !has_extensions {
args.push("--headless=new".to_string());
}
@@ -97,10 +143,18 @@ pub fn launch_chrome(options: &LaunchOptions) -> Result<ChromeProcess, String> {
args.push(format!("--proxy-bypass-list={}", bypass));
}
if let Some(ref profile) = options.profile {
let temp_user_data_dir = if let Some(ref profile) = options.profile {
let expanded = expand_tilde(profile);
args.push(format!("--user-data-dir={}", expanded));
}
None
} else {
let dir =
std::env::temp_dir().join(format!("agent-browser-chrome-{}", uuid::Uuid::new_v4()));
std::fs::create_dir_all(&dir)
.map_err(|e| format!("Failed to create temp profile dir: {}", e))?;
args.push(format!("--user-data-dir={}", dir.display()));
Some(dir)
};
if options.allow_file_access {
args.push("--allow-file-access-from-files".to_string());
@@ -115,13 +169,12 @@ pub fn launch_chrome(options: &LaunchOptions) -> Result<ChromeProcess, String> {
}
}
// Check if user args set window size (skip viewport override)
let has_window_size = options
.args
.iter()
.any(|a| a.starts_with("--start-maximized") || a.starts_with("--window-size="));
if !has_window_size && options.headless {
if !has_window_size && options.headless && !has_extensions {
args.push("--window-size=1280,720".to_string());
}
@@ -131,23 +184,104 @@ pub fn launch_chrome(options: &LaunchOptions) -> Result<ChromeProcess, String> {
args.push("--no-sandbox".to_string());
}
Ok(ChromeArgs {
args,
temp_user_data_dir,
})
}
pub const MANAGED_CDP_PORT: u16 = 9333;
pub fn managed_cdp_profile_dir() -> PathBuf {
dirs::home_dir()
.unwrap_or_else(|| std::env::temp_dir())
.join(".agent-browser")
.join("chrome-bot-profile")
}
fn cleanup_managed_profile_locks(profile_dir: &Path) {
let _ = std::fs::remove_file(profile_dir.join("DevToolsActivePort"));
if let Ok(entries) = std::fs::read_dir(profile_dir) {
for entry in entries.flatten() {
let name = entry.file_name();
if name.to_string_lossy().starts_with("Singleton") {
let _ = std::fs::remove_file(entry.path());
}
}
}
}
pub fn launch_managed_chrome(
executable_path: Option<String>,
headed: bool,
) -> Result<ChromeProcess, String> {
let profile_dir = managed_cdp_profile_dir();
std::fs::create_dir_all(&profile_dir)
.map_err(|e| format!("Failed to create managed Chrome profile dir: {}", e))?;
cleanup_managed_profile_locks(&profile_dir);
let options = LaunchOptions {
headless: !headed,
executable_path,
profile: Some(profile_dir.to_string_lossy().to_string()),
remote_debugging_port: Some(MANAGED_CDP_PORT),
..Default::default()
};
launch_chrome(&options)
}
pub fn launch_chrome(options: &LaunchOptions) -> Result<ChromeProcess, String> {
let chrome_path = match &options.executable_path {
Some(p) => PathBuf::from(p),
None => {
find_chrome().ok_or("Chrome not found. Install Chrome or use --executable-path.")?
}
};
let ChromeArgs {
args,
temp_user_data_dir,
} = build_chrome_args(options)?;
let cleanup_temp_dir = |dir: &Option<PathBuf>| {
if let Some(ref d) = dir {
let _ = std::fs::remove_dir_all(d);
}
};
let mut child = Command::new(&chrome_path)
.args(&args)
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::piped())
.spawn()
.map_err(|e| format!("Failed to launch Chrome at {:?}: {}", chrome_path, e))?;
.map_err(|e| {
cleanup_temp_dir(&temp_user_data_dir);
format!("Failed to launch Chrome at {:?}: {}", chrome_path, e)
})?;
let stderr = child
.stderr
.take()
.ok_or("Failed to capture Chrome stderr")?;
let stderr = child.stderr.take().ok_or_else(|| {
let _ = child.kill();
cleanup_temp_dir(&temp_user_data_dir);
"Failed to capture Chrome stderr".to_string()
})?;
let reader = BufReader::new(stderr);
let ws_url = wait_for_ws_url(reader)?;
let ws_url = match wait_for_ws_url(reader) {
Ok(url) => url,
Err(e) => {
let _ = child.kill();
cleanup_temp_dir(&temp_user_data_dir);
return Err(e);
}
};
Ok(ChromeProcess { child, ws_url })
Ok(ChromeProcess {
child,
ws_url,
temp_user_data_dir,
})
}
fn wait_for_ws_url(reader: BufReader<std::process::ChildStderr>) -> Result<String, String> {
@@ -315,55 +449,8 @@ pub async fn discover_cdp_url(port: u16) -> Result<String, String> {
}
async fn reqwest_get_string(url: &str) -> Result<String, String> {
let client = tokio::net::TcpStream::connect(
url.strip_prefix("http://")
.unwrap_or(url)
.split('/')
.next()
.unwrap_or("127.0.0.1:9222"),
)
.await
.map_err(|e| e.to_string())?;
let path = url
.find('/')
.and_then(|i| url[i..].find('/').map(|j| &url[i + j..]))
.unwrap_or("/json/version");
let host = url
.strip_prefix("http://")
.unwrap_or(url)
.split('/')
.next()
.unwrap_or("127.0.0.1");
let request = format!(
"GET {} HTTP/1.1\r\nHost: {}\r\nConnection: close\r\n\r\n",
path, host
);
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let mut client = client;
client
.write_all(request.as_bytes())
.await
.map_err(|e| e.to_string())?;
let mut response = Vec::new();
client
.read_to_end(&mut response)
.await
.map_err(|e| e.to_string())?;
let response_str = String::from_utf8_lossy(&response);
let body = response_str
.split("\r\n\r\n")
.nth(1)
.unwrap_or("")
.to_string();
Ok(body)
let resp = reqwest::get(url).await.map_err(|e| e.to_string())?;
resp.text().await.map_err(|e| e.to_string())
}
pub fn read_devtools_active_port(user_data_dir: &Path) -> Option<(u16, String)> {
@@ -470,10 +557,7 @@ fn should_disable_sandbox(existing_args: &[String]) -> bool {
// Generic container detection: cgroup contains docker/kubepods/lxc
if let Ok(cgroup) = std::fs::read_to_string("/proc/1/cgroup") {
if cgroup.contains("docker")
|| cgroup.contains("kubepods")
|| cgroup.contains("lxc")
{
if cgroup.contains("docker") || cgroup.contains("kubepods") || cgroup.contains("lxc") {
return true;
}
}
@@ -559,6 +643,7 @@ fn expand_tilde(path: &str) -> String {
#[cfg(test)]
mod tests {
use super::*;
use crate::test_utils::EnvGuard;
#[test]
fn test_find_chrome_returns_some_on_host() {
@@ -616,20 +701,195 @@ mod tests {
#[test]
fn test_chrome_launch_error_generic() {
let lines = vec![
"info line".to_string(),
"another info line".to_string(),
];
let lines = vec!["info line".to_string(), "another info line".to_string()];
let msg = chrome_launch_error("Chrome exited", &lines);
assert!(msg.contains("last 2 lines"));
}
#[test]
fn test_find_playwright_chromium_nonexistent() {
// With no Playwright cache, should return None
std::env::set_var("PLAYWRIGHT_BROWSERS_PATH", "/nonexistent/path");
let _guard = EnvGuard::new(&["PLAYWRIGHT_BROWSERS_PATH"]);
_guard.set("PLAYWRIGHT_BROWSERS_PATH", "/nonexistent/path");
let result = find_playwright_chromium();
std::env::remove_var("PLAYWRIGHT_BROWSERS_PATH");
assert!(result.is_none());
}
#[test]
fn test_build_args_headless_includes_headless_flag() {
let opts = LaunchOptions {
headless: true,
..Default::default()
};
let result = build_chrome_args(&opts).unwrap();
assert!(result.args.iter().any(|a| a == "--headless=new"));
assert!(result.args.iter().any(|a| a == "--window-size=1280,720"));
// Temp dir created when no profile
assert!(result.temp_user_data_dir.is_some());
let dir = result.temp_user_data_dir.unwrap();
assert!(dir.exists());
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn test_build_args_headed_no_headless_flag() {
let opts = LaunchOptions {
headless: false,
..Default::default()
};
let result = build_chrome_args(&opts).unwrap();
assert!(!result.args.iter().any(|a| a.contains("--headless")));
assert!(!result.args.iter().any(|a| a.starts_with("--window-size=")));
// Temp dir created when no profile
assert!(result.temp_user_data_dir.is_some());
let dir = result.temp_user_data_dir.unwrap();
assert!(dir.exists());
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn test_build_args_temp_user_data_dir_created() {
let opts = LaunchOptions::default();
let result = build_chrome_args(&opts).unwrap();
let dir = result.temp_user_data_dir.as_ref().unwrap();
assert!(dir.exists());
assert!(result
.args
.iter()
.any(|a| a.starts_with("--user-data-dir=")));
let _ = std::fs::remove_dir_all(dir);
}
#[test]
fn test_build_args_profile_no_temp_dir() {
let opts = LaunchOptions {
profile: Some("/tmp/my-profile".to_string()),
..Default::default()
};
let result = build_chrome_args(&opts).unwrap();
assert!(result.temp_user_data_dir.is_none());
assert!(result
.args
.iter()
.any(|a| a == "--user-data-dir=/tmp/my-profile"));
}
#[test]
fn test_build_args_custom_window_size_not_overridden() {
let opts = LaunchOptions {
headless: true,
args: vec!["--window-size=1920,1080".to_string()],
..Default::default()
};
let result = build_chrome_args(&opts).unwrap();
assert!(!result.args.iter().any(|a| a == "--window-size=1280,720"));
assert!(result.args.iter().any(|a| a == "--window-size=1920,1080"));
if let Some(ref dir) = result.temp_user_data_dir {
let _ = std::fs::remove_dir_all(dir);
}
}
#[test]
fn test_build_args_start_maximized_suppresses_default_window_size() {
let opts = LaunchOptions {
headless: true,
args: vec!["--start-maximized".to_string()],
..Default::default()
};
let result = build_chrome_args(&opts).unwrap();
assert!(!result.args.iter().any(|a| a == "--window-size=1280,720"));
assert!(result.args.iter().any(|a| a == "--start-maximized"));
if let Some(ref dir) = result.temp_user_data_dir {
let _ = std::fs::remove_dir_all(dir);
}
}
#[test]
fn test_build_args_disables_translate() {
let opts = LaunchOptions::default();
let result = build_chrome_args(&opts).unwrap();
assert!(result
.args
.iter()
.any(|a| a.contains("--disable-features") && a.contains("Translate")));
if let Some(ref dir) = result.temp_user_data_dir {
let _ = std::fs::remove_dir_all(dir);
}
}
#[test]
fn test_build_args_headless_with_extensions_skips_headless_flag() {
let opts = LaunchOptions {
headless: true,
extensions: Some(vec!["/tmp/my-ext".to_string()]),
..Default::default()
};
let result = build_chrome_args(&opts).unwrap();
assert!(
!result.args.iter().any(|a| a.contains("--headless")),
"headless flag should be omitted when extensions are present"
);
assert!(
!result.args.iter().any(|a| a.contains("--window-size")),
"window-size should be omitted when extensions force headed mode"
);
assert!(result
.args
.iter()
.any(|a| a.starts_with("--load-extension=")));
if let Some(ref dir) = result.temp_user_data_dir {
let _ = std::fs::remove_dir_all(dir);
}
}
#[test]
fn test_build_args_headed_with_extensions_no_headless_flag() {
let opts = LaunchOptions {
headless: false,
extensions: Some(vec!["/tmp/my-ext".to_string()]),
..Default::default()
};
let result = build_chrome_args(&opts).unwrap();
assert!(
!result.args.iter().any(|a| a.contains("--headless")),
"headless flag should not be present in headed mode"
);
assert!(result
.args
.iter()
.any(|a| a.starts_with("--load-extension=")));
if let Some(ref dir) = result.temp_user_data_dir {
let _ = std::fs::remove_dir_all(dir);
}
}
#[test]
fn test_chrome_process_drop_cleans_temp_dir() {
let dir = std::env::temp_dir().join(format!(
"agent-browser-chrome-drop-test-{}",
uuid::Uuid::new_v4()
));
let _ = std::fs::create_dir_all(&dir);
assert!(dir.exists());
{
// Simulate a ChromeProcess with a temp dir but a dummy child.
// We can't actually spawn Chrome here, but we can verify the Drop
// logic by creating a small helper process.
let child = Command::new("echo")
.arg("test")
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()
.unwrap();
let _process = ChromeProcess {
child,
ws_url: String::new(),
temp_user_data_dir: Some(dir.clone()),
};
// _process dropped here
}
assert!(!dir.exists(), "Temp dir should be cleaned up on drop");
}
}
+271
View File
@@ -0,0 +1,271 @@
use std::io::{BufRead, BufReader};
use std::net::TcpListener;
use std::path::PathBuf;
use std::process::{Child, Command, Stdio};
use std::time::Duration;
pub struct LightpandaProcess {
child: Child,
pub ws_url: String,
_stderr_drain: Option<std::thread::JoinHandle<()>>,
}
impl LightpandaProcess {
pub fn kill(&mut self) {
let _ = self.child.kill();
let _ = self.child.wait();
}
}
impl Drop for LightpandaProcess {
fn drop(&mut self) {
self.kill();
}
}
#[derive(Default)]
pub struct LightpandaLaunchOptions {
pub executable_path: Option<String>,
pub proxy: Option<String>,
pub port: Option<u16>,
}
pub fn find_lightpanda() -> Option<PathBuf> {
#[cfg(unix)]
{
if let Ok(output) = Command::new("which").arg("lightpanda").output() {
if output.status.success() {
let path = String::from_utf8_lossy(&output.stdout).trim().to_string();
if !path.is_empty() {
return Some(PathBuf::from(path));
}
}
}
}
#[cfg(windows)]
{
if let Ok(output) = Command::new("where").arg("lightpanda").output() {
if output.status.success() {
let path = String::from_utf8_lossy(&output.stdout)
.lines()
.next()
.unwrap_or("")
.trim()
.to_string();
if !path.is_empty() {
return Some(PathBuf::from(path));
}
}
}
}
if let Some(home) = dirs::home_dir() {
let candidates = [
home.join(".lightpanda/lightpanda"),
home.join(".local/bin/lightpanda"),
];
for candidate in &candidates {
if candidate.exists() {
return Some(candidate.clone());
}
}
}
None
}
pub fn launch_lightpanda(options: &LightpandaLaunchOptions) -> Result<LightpandaProcess, String> {
let binary_path = match &options.executable_path {
Some(path) => PathBuf::from(path),
None => find_lightpanda().ok_or(
"Lightpanda not found. Install it from https://lightpanda.io/docs/open-source/installation or use --executable-path.",
)?,
};
let port = match options.port {
Some(port) => port,
None => TcpListener::bind("127.0.0.1:0")
.and_then(|listener| listener.local_addr())
.map(|addr| addr.port())
.map_err(|e| format!("Failed to find an available port for Lightpanda: {}", e))?,
};
let mut args = vec![
"serve".to_string(),
"--host".to_string(),
"127.0.0.1".to_string(),
"--port".to_string(),
port.to_string(),
"--timeout".to_string(),
"0".to_string(),
];
if let Some(ref proxy) = options.proxy {
args.push("--http_proxy".to_string());
args.push(proxy.clone());
}
let mut child = Command::new(&binary_path)
.args(&args)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.map_err(|e| format!("Failed to launch Lightpanda at {:?}: {}", binary_path, e))?;
let stderr = child.stderr.take().ok_or_else(|| {
let _ = child.kill();
"Failed to capture Lightpanda stderr".to_string()
})?;
let reader = BufReader::new(stderr);
let (address, reader) = match wait_for_address(reader) {
Ok(result) => result,
Err(e) => {
let _ = child.kill();
return Err(e);
}
};
let ws_url = format!("ws://{}", address);
let drain = std::thread::spawn(move || {
let mut reader = reader;
let mut buf = String::new();
loop {
buf.clear();
match reader.read_line(&mut buf) {
Ok(0) | Err(_) => break,
Ok(_) => {}
}
}
});
Ok(LightpandaProcess {
child,
ws_url,
_stderr_drain: Some(drain),
})
}
fn wait_for_address(
mut reader: BufReader<std::process::ChildStderr>,
) -> Result<(String, BufReader<std::process::ChildStderr>), String> {
let deadline = std::time::Instant::now() + Duration::from_secs(30);
let mut stderr_lines: Vec<String> = Vec::new();
let mut buf = String::new();
loop {
if std::time::Instant::now() > deadline {
return Err(lightpanda_launch_error(
"Timeout waiting for Lightpanda server address",
&stderr_lines,
));
}
buf.clear();
match reader.read_line(&mut buf) {
Ok(0) => {
return Err(lightpanda_launch_error(
"Lightpanda exited before providing server address",
&stderr_lines,
));
}
Ok(_) => {
let line = buf.trim_end().to_string();
if let Some(address) = extract_address(&line) {
return Ok((address, reader));
}
stderr_lines.push(line);
}
Err(e) => {
return Err(format!("Failed to read Lightpanda stderr: {}", e));
}
}
}
}
fn extract_address(line: &str) -> Option<String> {
if let Some(idx) = line.find("address = ") {
let address = line[idx + "address = ".len()..].trim().to_string();
if !address.is_empty() {
return Some(address);
}
}
None
}
fn lightpanda_launch_error(message: &str, stderr_lines: &[String]) -> String {
if stderr_lines.is_empty() {
return format!("{} (no stderr output from Lightpanda)", message);
}
let last_lines: Vec<&String> = stderr_lines.iter().rev().take(5).collect();
format!(
"{}\nLightpanda stderr (last {} lines):\n {}",
message,
last_lines.len(),
last_lines
.into_iter()
.rev()
.map(|line| line.as_str())
.collect::<Vec<_>>()
.join("\n ")
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_extract_address_standard() {
assert_eq!(
extract_address(" address = 127.0.0.1:9222"),
Some("127.0.0.1:9222".to_string())
);
}
#[test]
fn test_extract_address_inline() {
assert_eq!(
extract_address("INFO app : server running address = 127.0.0.1:4567"),
Some("127.0.0.1:4567".to_string())
);
}
#[test]
fn test_extract_address_no_match() {
assert_eq!(extract_address("INFO app : starting up..."), None);
}
#[test]
fn test_find_lightpanda_returns_none_when_missing() {
let _ = find_lightpanda();
}
#[test]
fn test_lightpanda_launch_error_no_stderr() {
let msg = lightpanda_launch_error("Lightpanda exited", &[]);
assert!(msg.contains("no stderr output"));
}
#[test]
fn test_lightpanda_launch_error_with_lines() {
let lines = vec![
"INFO starting up".to_string(),
"ERROR bind failed: address in use".to_string(),
];
let msg = lightpanda_launch_error("Lightpanda exited", &lines);
assert!(msg.contains("bind failed"));
assert!(msg.contains("last 2 lines"));
}
#[test]
fn test_default_options() {
let opts = LightpandaLaunchOptions::default();
assert!(opts.executable_path.is_none());
assert!(opts.proxy.is_none());
assert!(opts.port.is_none());
}
}
+1
View File
@@ -1,3 +1,4 @@
pub mod chrome;
pub mod client;
pub mod lightpanda;
pub mod types;
+1
View File
@@ -532,6 +532,7 @@ pub struct BrowserVersionInfo {
/// Chromium source) into `cli/cdp-protocol/` and rebuild.
///
/// Usage: `use super::cdp::types::generated::cdp_page::*;`
#[allow(clippy::upper_case_acronyms)]
pub mod generated {
include!(concat!(env!("OUT_DIR"), "/cdp_generated.rs"));
}
+5 -7
View File
@@ -56,13 +56,11 @@ pub async fn set_cookies(
.into_iter()
.map(|mut c| {
// Auto-fill url if no domain/path/url provided
if c.get("url").is_none() && c.get("domain").is_none() && current_url.is_some() {
c.as_object_mut().map(|m| {
m.insert(
"url".to_string(),
Value::String(current_url.unwrap().to_string()),
)
});
if c.get("url").is_none() && c.get("domain").is_none() {
if let Some(url) = current_url {
c.as_object_mut()
.map(|m| m.insert("url".to_string(), Value::String(url.to_string())));
}
}
c
})
+79 -8
View File
@@ -1,16 +1,22 @@
use serde_json::Value;
use serde_json::{json, Value};
use std::env;
use std::fs;
use std::path::PathBuf;
use std::process;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::signal;
use tokio::sync::mpsc::{unbounded_channel, UnboundedSender};
use tokio::time::{Duration, Instant};
use super::actions::{execute_command, DaemonState};
use super::state;
const IDLE_SHUTDOWN_SECS: u64 = 600;
pub async fn run_daemon(session: &str) {
let resident_mode = env::args().any(|arg| arg == "--resident");
let socket_dir = get_daemon_socket_dir();
if !socket_dir.exists() {
let _ = fs::create_dir_all(&socket_dir);
@@ -18,6 +24,16 @@ pub async fn run_daemon(session: &str) {
let pid_path = socket_dir.join(format!("{}.pid", session));
let _ = fs::write(&pid_path, process::id().to_string());
let meta_path = socket_dir.join(format!("{}.meta.json", session));
if let Ok(current_exe) = env::current_exe() {
let daemon_path = current_exe.canonicalize().unwrap_or(current_exe);
let cli_version = env::var("AGENT_BROWSER_CLI_VERSION").unwrap_or_default();
let meta = json!({
"daemonPath": daemon_path.to_string_lossy(),
"cliVersion": cli_version,
});
let _ = fs::write(&meta_path, meta.to_string());
}
let socket_path = socket_dir.join(format!("{}.sock", session));
@@ -33,10 +49,11 @@ pub async fn run_daemon(session: &str) {
}
}
let result = run_socket_server(&socket_path, session).await;
let result = run_socket_server(&socket_path, session, resident_mode).await;
let _ = fs::remove_file(&socket_path);
let _ = fs::remove_file(&pid_path);
let _ = fs::remove_file(&meta_path);
let stream_path = socket_dir.join(format!("{}.stream", session));
let _ = fs::remove_file(&stream_path);
@@ -47,7 +64,11 @@ pub async fn run_daemon(session: &str) {
}
#[cfg(unix)]
async fn run_socket_server(socket_path: &PathBuf, _session: &str) -> Result<(), String> {
async fn run_socket_server(
socket_path: &PathBuf,
_session: &str,
resident_mode: bool,
) -> Result<(), String> {
use tokio::net::UnixListener;
let listener =
@@ -55,6 +76,9 @@ async fn run_socket_server(socket_path: &PathBuf, _session: &str) -> Result<(),
let state: std::sync::Arc<tokio::sync::Mutex<DaemonState>> =
std::sync::Arc::new(tokio::sync::Mutex::new(DaemonState::new()));
let active_commands = std::sync::Arc::new(AtomicUsize::new(0));
let (activity_tx, mut activity_rx) = unbounded_channel::<()>();
let mut idle_deadline = Instant::now() + Duration::from_secs(IDLE_SHUTDOWN_SECS);
loop {
tokio::select! {
@@ -62,8 +86,10 @@ async fn run_socket_server(socket_path: &PathBuf, _session: &str) -> Result<(),
match accept_result {
Ok((stream, _)) => {
let state = state.clone();
let activity_tx = activity_tx.clone();
let active_commands = active_commands.clone();
tokio::spawn(async move {
handle_connection(stream, state).await;
handle_connection(stream, state, activity_tx, active_commands).await;
});
}
Err(e) => {
@@ -71,6 +97,19 @@ async fn run_socket_server(socket_path: &PathBuf, _session: &str) -> Result<(),
}
}
}
Some(_) = activity_rx.recv() => {
idle_deadline = Instant::now() + Duration::from_secs(IDLE_SHUTDOWN_SECS);
}
_ = tokio::time::sleep_until(idle_deadline), if !resident_mode => {
if active_commands.load(Ordering::SeqCst) == 0 {
let mut s = state.lock().await;
if let Some(ref mut mgr) = s.browser {
let _ = mgr.close().await;
}
break;
}
idle_deadline = Instant::now() + Duration::from_secs(IDLE_SHUTDOWN_SECS);
}
_ = shutdown_signal() => {
let mut s = state.lock().await;
if let Some(ref mut mgr) = s.browser {
@@ -85,7 +124,11 @@ async fn run_socket_server(socket_path: &PathBuf, _session: &str) -> Result<(),
}
#[cfg(windows)]
async fn run_socket_server(socket_path: &PathBuf, session: &str) -> Result<(), String> {
async fn run_socket_server(
socket_path: &PathBuf,
session: &str,
resident_mode: bool,
) -> Result<(), String> {
use tokio::net::TcpListener;
let port = get_port_for_session(session);
@@ -99,6 +142,9 @@ async fn run_socket_server(socket_path: &PathBuf, session: &str) -> Result<(), S
let state: std::sync::Arc<tokio::sync::Mutex<DaemonState>> =
std::sync::Arc::new(tokio::sync::Mutex::new(DaemonState::new()));
let active_commands = std::sync::Arc::new(AtomicUsize::new(0));
let (activity_tx, mut activity_rx) = unbounded_channel::<()>();
let mut idle_deadline = Instant::now() + Duration::from_secs(IDLE_SHUTDOWN_SECS);
loop {
tokio::select! {
@@ -106,8 +152,10 @@ async fn run_socket_server(socket_path: &PathBuf, session: &str) -> Result<(), S
match accept_result {
Ok((stream, _)) => {
let state = state.clone();
let activity_tx = activity_tx.clone();
let active_commands = active_commands.clone();
tokio::spawn(async move {
handle_connection(stream, state).await;
handle_connection(stream, state, activity_tx, active_commands).await;
});
}
Err(e) => {
@@ -115,6 +163,20 @@ async fn run_socket_server(socket_path: &PathBuf, session: &str) -> Result<(), S
}
}
}
Some(_) = activity_rx.recv() => {
idle_deadline = Instant::now() + Duration::from_secs(IDLE_SHUTDOWN_SECS);
}
_ = tokio::time::sleep_until(idle_deadline), if !resident_mode => {
if active_commands.load(Ordering::SeqCst) == 0 {
let mut s = state.lock().await;
if let Some(ref mut mgr) = s.browser {
let _ = mgr.close().await;
}
let _ = fs::remove_file(&port_path);
break;
}
idle_deadline = Instant::now() + Duration::from_secs(IDLE_SHUTDOWN_SECS);
}
_ = shutdown_signal() => {
let mut s = state.lock().await;
if let Some(ref mut mgr) = s.browser {
@@ -129,8 +191,12 @@ async fn run_socket_server(socket_path: &PathBuf, session: &str) -> Result<(), S
Ok(())
}
async fn handle_connection<S>(stream: S, state: std::sync::Arc<tokio::sync::Mutex<DaemonState>>)
where
async fn handle_connection<S>(
stream: S,
state: std::sync::Arc<tokio::sync::Mutex<DaemonState>>,
activity_tx: UnboundedSender<()>,
active_commands: std::sync::Arc<AtomicUsize>,
) where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
let (reader, mut writer) = tokio::io::split(stream);
@@ -166,6 +232,8 @@ where
};
let is_close = cmd.get("action").and_then(|v| v.as_str()) == Some("close");
let _ = activity_tx.send(());
active_commands.fetch_add(1, Ordering::SeqCst);
let response = {
let mut s = state.lock().await;
@@ -175,8 +243,11 @@ where
let mut resp = serde_json::to_string(&response).unwrap_or_default();
resp.push('\n');
if writer.write_all(resp.as_bytes()).await.is_err() {
active_commands.fetch_sub(1, Ordering::SeqCst);
break;
}
active_commands.fetch_sub(1, Ordering::SeqCst);
let _ = activity_tx.send(());
if is_close {
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+114
View File
@@ -566,6 +566,7 @@ async fn e2e_tabs() {
let tabs = get_data(&resp)["tabs"].as_array().unwrap();
assert_eq!(tabs.len(), 1);
assert_eq!(tabs[0]["active"], true);
assert_eq!(tabs[0]["type"], "page");
// Open new tab
let resp = execute_command(
@@ -582,6 +583,7 @@ async fn e2e_tabs() {
let tabs = get_data(&resp)["tabs"].as_array().unwrap();
assert_eq!(tabs.len(), 2);
assert_eq!(tabs[1]["active"], true);
assert_eq!(tabs[1]["type"], "page");
// Switch to first tab
let resp = execute_command(
@@ -1293,3 +1295,115 @@ async fn e2e_error_handling() {
let resp = execute_command(&json!({ "id": "99", "action": "close" }), &mut state).await;
assert_success(&resp);
}
// ---------------------------------------------------------------------------
// Profile cookie persistence across restarts
// ---------------------------------------------------------------------------
#[tokio::test]
#[ignore]
async fn e2e_profile_cookie_persistence() {
let profile_dir = std::env::temp_dir().join(format!(
"agent-browser-e2e-profile-{}",
uuid::Uuid::new_v4()
));
// Session 1: launch with profile, set a cookie, close
{
let mut state = DaemonState::new();
let resp = execute_command(
&json!({
"id": "1",
"action": "launch",
"headless": true,
"profile": profile_dir.to_str().unwrap()
}),
&mut state,
)
.await;
assert_success(&resp);
let resp = execute_command(
&json!({ "id": "2", "action": "navigate", "url": "https://example.com" }),
&mut state,
)
.await;
assert_success(&resp);
let resp = execute_command(
&json!({
"id": "3",
"action": "cookies_set",
"name": "persist_test",
"value": "should_survive_restart",
"domain": ".example.com",
"path": "/",
"expires": 2000000000
}),
&mut state,
)
.await;
assert_success(&resp);
// Verify cookie is set
let resp =
execute_command(&json!({ "id": "4", "action": "cookies_get" }), &mut state).await;
assert_success(&resp);
let cookies = get_data(&resp)["cookies"].as_array().unwrap();
let found = cookies
.iter()
.any(|c| c["name"] == "persist_test" && c["value"] == "should_survive_restart");
assert!(found, "Cookie should exist before close");
let resp = execute_command(&json!({ "id": "5", "action": "close" }), &mut state).await;
assert_success(&resp);
}
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
// Session 2: reopen with the same profile, verify cookie persisted
{
let mut state = DaemonState::new();
let resp = execute_command(
&json!({
"id": "10",
"action": "launch",
"headless": true,
"profile": profile_dir.to_str().unwrap()
}),
&mut state,
)
.await;
assert_success(&resp);
let resp = execute_command(
&json!({ "id": "11", "action": "navigate", "url": "https://example.com" }),
&mut state,
)
.await;
assert_success(&resp);
let resp =
execute_command(&json!({ "id": "12", "action": "cookies_get" }), &mut state).await;
assert_success(&resp);
let cookies = get_data(&resp)["cookies"].as_array().unwrap();
let found = cookies
.iter()
.any(|c| c["name"] == "persist_test" && c["value"] == "should_survive_restart");
assert!(
found,
"Cookie should persist across restart with --profile. Cookies found: {:?}",
cookies
.iter()
.map(|c| c["name"].as_str().unwrap_or("?"))
.collect::<Vec<_>>()
);
let resp = execute_command(&json!({ "id": "99", "action": "close" }), &mut state).await;
assert_success(&resp);
}
let _ = std::fs::remove_dir_all(&profile_dir);
}
+44 -1
View File
@@ -9,6 +9,38 @@ use serde_json::{json, Value};
use super::actions::{execute_command, DaemonState};
const ENCRYPTION_KEY_ENV: &str = "AGENT_BROWSER_ENCRYPTION_KEY";
struct TestKeyGuard {
_lock: std::sync::MutexGuard<'static, ()>,
original: Option<String>,
}
impl TestKeyGuard {
fn new() -> Self {
let lock = super::auth::AUTH_TEST_MUTEX
.lock()
.unwrap_or_else(|e| e.into_inner());
let original = std::env::var(ENCRYPTION_KEY_ENV).ok();
// SAFETY: AUTH_TEST_MUTEX serializes all test access so no concurrent mutation.
unsafe { std::env::set_var(ENCRYPTION_KEY_ENV, "a".repeat(64)) };
Self {
_lock: lock,
original,
}
}
}
impl Drop for TestKeyGuard {
fn drop(&mut self) {
// SAFETY: AUTH_TEST_MUTEX is held via _lock.
match &self.original {
Some(val) => unsafe { std::env::set_var(ENCRYPTION_KEY_ENV, val) },
None => unsafe { std::env::remove_var(ENCRYPTION_KEY_ENV) },
}
}
}
/// All documented action names that should be implemented.
const DOCUMENTED_ACTIONS: &[&str] = &[
"launch",
@@ -342,13 +374,22 @@ fn minimal_command(action: &str, id: &str) -> Value {
// ---------------------------------------------------------------------------
#[tokio::test]
#[ignore]
async fn test_all_documented_actions_are_handled() {
let mut state = DaemonState::new();
for (i, action) in DOCUMENTED_ACTIONS.iter().enumerate() {
let id = format!("parity-{}", i);
let cmd = minimal_command(action, &id);
let result = execute_command(&cmd, &mut state).await;
let result = tokio::time::timeout(
tokio::time::Duration::from_millis(250),
execute_command(&cmd, &mut state),
)
.await;
let Ok(result) = result else {
continue;
};
assert!(
result.get("id").is_some(),
@@ -424,6 +465,7 @@ async fn test_credentials_list_without_browser() {
#[tokio::test]
async fn test_auth_profile_name_validation() {
use super::auth;
let _key_guard = TestKeyGuard::new();
let valid = auth::credentials_set("valid-name_123", "u", "p", None);
assert!(valid.is_ok());
let invalid = auth::credentials_set("invalid/name", "u", "p", None);
@@ -439,6 +481,7 @@ async fn test_auth_profile_name_validation() {
#[tokio::test]
async fn test_auth_save_and_show() {
use super::auth;
let _key_guard = TestKeyGuard::new();
let result = auth::auth_save(
"parity-roundtrip",
"https://example.com",
+3 -2
View File
@@ -135,6 +135,7 @@ impl ActionPolicy {
#[cfg(test)]
mod tests {
use super::*;
use crate::test_utils::EnvGuard;
#[test]
fn test_policy_allow_whitelist() {
@@ -205,12 +206,12 @@ mod tests {
#[test]
fn test_confirm_actions_from_env() {
env::set_var("AGENT_BROWSER_CONFIRM_ACTIONS", "navigate,click,fill");
let _guard = EnvGuard::new(&["AGENT_BROWSER_CONFIRM_ACTIONS"]);
_guard.set("AGENT_BROWSER_CONFIRM_ACTIONS", "navigate,click,fill");
let ca = ConfirmActions::from_env().unwrap();
assert!(ca.requires_confirmation("navigate"));
assert!(ca.requires_confirmation("click"));
assert!(ca.requires_confirmation("fill"));
assert!(!ca.requires_confirmation("screenshot"));
env::remove_var("AGENT_BROWSER_CONFIRM_ACTIONS");
}
}
+2 -14
View File
@@ -65,6 +65,7 @@ const STRUCTURAL_ROLES: &[&str] = &[
"RootWebArea",
];
#[derive(Default)]
pub struct SnapshotOptions {
pub selector: Option<String>,
pub interactive: bool,
@@ -73,18 +74,6 @@ pub struct SnapshotOptions {
pub cursor: bool,
}
impl Default for SnapshotOptions {
fn default() -> Self {
Self {
selector: None,
interactive: false,
compact: false,
depth: None,
cursor: false,
}
}
}
struct TreeNode {
role: String,
name: String,
@@ -364,8 +353,7 @@ async fn find_cursor_interactive_elements(
let escaped = text
.replace('\\', "\\\\")
.replace('"', "\\\"")
.replace('\n', " ")
.replace('\r', " ");
.replace(['\n', '\r'], " ");
lines.push(format!("[ref={}] ({}) \"{}\"", ref_id, kind, escaped));
}
+1 -1
View File
@@ -467,7 +467,7 @@ pub fn find_auto_state_file(session_name: &str) -> Option<String> {
.ok()
.and_then(|m| m.modified().ok())
.unwrap_or(std::time::UNIX_EPOCH);
if best_path.as_ref().map_or(true, |(_, t)| modified > *t) {
if best_path.as_ref().is_none_or(|(_, t)| modified > *t) {
best_path = Some((path.to_string_lossy().to_string(), modified));
}
}
+282 -86
View File
@@ -105,6 +105,42 @@ pub fn print_response_with_opts(resp: &Response, action: Option<&str>, opts: &Ou
if let Some(title) = data.get("title").and_then(|v| v.as_str()) {
println!("{} {}", color::success_indicator(), color::bold(title));
println!(" {}", color::dim(url));
if let Some(warning) = data.get("warning").and_then(|v| v.as_str()) {
println!("{} {}", color::warning_indicator(), warning);
}
if let Some(risk_signals) = data.get("riskSignals").and_then(|v| v.as_array()) {
for signal in risk_signals {
let code = signal
.get("code")
.and_then(|v| v.as_str())
.unwrap_or("unknown_risk");
let source = signal
.get("source")
.and_then(|v| v.as_str())
.unwrap_or("unknown");
let evidence = signal
.get("evidence")
.and_then(|v| v.as_str())
.unwrap_or("-");
let confidence = signal
.get("confidence")
.and_then(|v| v.as_f64())
.unwrap_or(0.0);
println!(
"{} risk-signal code={} source={} evidence={} confidence={:.2}",
color::warning_indicator(),
code,
source,
evidence,
confidence
);
}
}
if let Some(warnings) = data.get("warnings").and_then(|v| v.as_array()) {
for warning in warnings.iter().filter_map(|v| v.as_str()) {
println!("{} {}", color::warning_indicator(), warning);
}
}
return;
}
println!("{}", url);
@@ -113,6 +149,10 @@ pub fn print_response_with_opts(resp: &Response, action: Option<&str>, opts: &Ou
// Diff responses -- route by action to avoid fragile shape probing
if let Some(obj) = data.as_object() {
match action {
Some("doctor") => {
print_doctor_report(obj);
return;
}
Some("diff_snapshot") => {
print_snapshot_diff(obj);
return;
@@ -395,6 +435,12 @@ pub fn print_response_with_opts(resp: &Response, action: Option<&str>, opts: &Ou
}
return;
}
// Launch/start
if data.get("launched").is_some() {
let label = "Browser ready";
println!("{} {}", color::success_indicator(), label);
return;
}
// Closed (browser or tab)
if data.get("closed").is_some() {
let label = match action {
@@ -831,16 +877,37 @@ Aliases: goto, navigate
Global Options:
--json Output as JSON
--session <name> Use specific session
--headers <json> Set HTTP headers (scoped to this origin)
--risk-mode <mode> Risk handling for verify/captcha pages: off, warn, block
--wait-until <mode> Navigation wait strategy: load, domcontentloaded, networkidle
--headed Show browser window
Examples:
agent-browser open example.com
agent-browser --risk-mode block open example.com
agent-browser --wait-until domcontentloaded open example.com
agent-browser open https://github.com
agent-browser open localhost:3000
agent-browser open api.example.com --headers '{"Authorization": "Bearer token"}'
# ^ Headers only sent to api.example.com, not other domains
"##
}
"start" => {
r##"
agent-browser start - Start the managed automation browser on localhost:9333
Usage: agent-browser start
Starts or reuses the dedicated automation Chrome profile on localhost:9333
without navigating to a page. Use this to pre-warm the managed browser for
unattended runs.
Global Options:
--json Output as JSON
Examples:
agent-browser start
abs start
"##
}
"back" => {
@@ -854,7 +921,6 @@ the browser's back button.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser back
@@ -871,7 +937,6 @@ the browser's forward button.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser forward
@@ -888,7 +953,6 @@ the browser's reload button.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser reload
@@ -911,7 +975,6 @@ Options:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser click "#submit-button"
@@ -932,7 +995,6 @@ or triggering double-click handlers.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser dblclick "#editable-text"
@@ -950,7 +1012,6 @@ This replaces any existing content in the field.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser fill "#email" "user@example.com"
@@ -962,18 +1023,20 @@ Examples:
r##"
agent-browser type - Type text into an element
Usage: agent-browser type <selector> <text>
Usage: agent-browser type <selector> <text> [--delay <ms>]
Types text into the specified element character by character.
Unlike fill, this does not clear existing content first.
Use --delay to add per-character delay (milliseconds).
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser type "#search" "hello"
agent-browser type "#search" "iphone" --delay 120
agent-browser type @e2 "additional text"
agent-browser type @e2 -- "--delay 120 (literal text)"
See Also:
For typing into contenteditable editors (Lexical, ProseMirror, etc.)
@@ -992,7 +1055,6 @@ triggering hover states or dropdown menus.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser hover "#dropdown-trigger"
@@ -1009,7 +1071,6 @@ Sets keyboard focus to the specified element.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser focus "#input-field"
@@ -1026,7 +1087,6 @@ Checks a checkbox element. If already checked, no action is taken.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser check "#terms-checkbox"
@@ -1043,7 +1103,6 @@ Unchecks a checkbox element. If already unchecked, no action is taken.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser uncheck "#newsletter-opt-in"
@@ -1060,7 +1119,6 @@ Selects one or more options in a <select> dropdown by value.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser select "#country" "US"
@@ -1078,7 +1136,6 @@ Drags an element from source to target location.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser drag "#draggable" "#drop-zone"
@@ -1095,7 +1152,6 @@ Uploads one or more files to a file input element.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser upload "#file-input" ./document.pdf
@@ -1116,7 +1172,6 @@ Arguments:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser download "#download-btn" ./file.pdf
@@ -1147,7 +1202,6 @@ Modifiers (combine with +):
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser press Enter
@@ -1168,7 +1222,6 @@ Useful for holding modifier keys.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser keydown Shift
@@ -1185,7 +1238,6 @@ Releases a key that was pressed with keydown.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser keyup Shift
@@ -1204,7 +1256,7 @@ the current focus — essential for contenteditable editors like
Lexical, ProseMirror, CodeMirror, and Monaco.
Subcommands:
type <text> Type text character-by-character with real
type <text> [--delay <ms>] Type text character-by-character with real
key events (keydown, keypress, keyup per char)
inserttext <text> Insert text without key events (like paste)
@@ -1213,10 +1265,10 @@ directly — it already operates on the current focus.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser keyboard type "Hello, World!"
agent-browser keyboard type "human pacing" --delay 90
agent-browser keyboard type "# My Heading"
agent-browser keyboard inserttext "pasted content"
@@ -1247,7 +1299,6 @@ Options:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser scroll
@@ -1269,7 +1320,6 @@ Aliases: scrollinto
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser scrollintoview "#footer"
@@ -1282,13 +1332,14 @@ Examples:
r##"
agent-browser wait - Wait for condition
Usage: agent-browser wait <selector|ms|option>
Usage: agent-browser wait <selector|ms|min-max|option>
Waits for an element to appear, a timeout, or other conditions.
Modes:
<selector> Wait for element to appear
<ms> Wait for specified milliseconds
<min>-<max> Wait for random time between min and max ms
--url <pattern> Wait for URL to match pattern
--load <state> Wait for load state (load, domcontentloaded, networkidle)
--fn <expression> Wait for JavaScript expression to be truthy
@@ -1300,11 +1351,11 @@ Download Options (with --download):
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser wait "#loading-spinner"
agent-browser wait 2000
agent-browser wait 2000-5000 # Random wait between 2-5 seconds
agent-browser wait --url "**/dashboard"
agent-browser wait --load networkidle
agent-browser wait --fn "window.appReady === true"
@@ -1333,7 +1384,6 @@ Options:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser screenshot
@@ -1354,7 +1404,6 @@ Saves the current page as a PDF file.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser pdf ./page.pdf
@@ -1382,7 +1431,6 @@ Options:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser snapshot
@@ -1408,7 +1456,6 @@ Options:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser eval "document.title"
@@ -1437,11 +1484,9 @@ Aliases: quit, exit
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser close
agent-browser close --session mysession
"##
}
@@ -1467,7 +1512,6 @@ Subcommands:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser get text @e1
@@ -1499,7 +1543,6 @@ Subcommands:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser is visible "#modal"
@@ -1538,7 +1581,6 @@ Options:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser find role button click --name Submit
@@ -1568,7 +1610,6 @@ Subcommands:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser mouse move 100 200
@@ -1601,7 +1642,6 @@ Settings:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser set viewport 1920 1080
@@ -1635,7 +1675,6 @@ Subcommands:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser network route "**/api/*" --abort
@@ -1667,7 +1706,6 @@ Operations:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser storage local
@@ -1694,8 +1732,8 @@ Operations:
Cookie Set Options:
--url <url> URL for the cookie (allows setting before page load)
--domain <domain> Cookie domain (e.g., ".example.com")
--path <path> Cookie path (e.g., "/api")
--domain <domain> Cookie domain (use with --path, e.g., ".example.com")
--path <path> Cookie path (use with --domain, e.g., "/api")
--httpOnly Set HttpOnly flag (prevents JavaScript access)
--secure Set Secure flag (HTTPS only)
--sameSite <Strict|Lax|None> SameSite policy
@@ -1703,10 +1741,10 @@ Cookie Set Options:
Note: If --url, --domain, and --path are all omitted, the cookie will be set
for the current page URL.
When --url is omitted, --domain and --path must be provided together.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
# Simple cookie for current page
@@ -1749,7 +1787,6 @@ Operations:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser tab
@@ -1776,7 +1813,6 @@ Operations:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser window new
@@ -1798,7 +1834,6 @@ Arguments:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser frame "#embed-iframe"
@@ -1832,7 +1867,6 @@ Save Options:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
echo "pass" | agent-browser auth save github --url https://github.com/login --username user --password-stdin
@@ -1880,7 +1914,6 @@ Operations:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser dialog accept
@@ -1904,7 +1937,6 @@ Operations:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser trace start
@@ -1935,7 +1967,6 @@ Start Options:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
# Basic profiling
@@ -1974,7 +2005,6 @@ Operations:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
# Record from current page (preserves login state)
@@ -2006,7 +2036,6 @@ Options:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser console
@@ -2026,7 +2055,6 @@ Options:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser errors
@@ -2045,7 +2073,6 @@ Visually highlights an element on the page for debugging.
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser highlight "#target-element"
@@ -2072,9 +2099,12 @@ Operations:
clean --older-than <days> Delete expired state files
Automatic State Persistence:
Use --session-name to auto-save/restore state across restarts:
Use --session-name to auto-save/restore state across restarts.
If omitted in default runtime mode, it defaults to "default":
agent-browser --session-name myapp open https://example.com
Or set AGENT_BROWSER_SESSION_NAME environment variable.
Note: with --parallel <name>, persistence is disabled by default unless
--session-name is explicitly passed on the same command.
State Encryption:
Set AGENT_BROWSER_ENCRYPTION_KEY (64-char hex) for AES-256-GCM encryption.
@@ -2082,7 +2112,6 @@ State Encryption:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser state save ./auth-state.json
@@ -2102,24 +2131,18 @@ agent-browser session - Manage sessions
Usage: agent-browser session [operation]
Manage isolated browser sessions. Each session has its own browser
instance with separate cookies, storage, and state.
Show the current runtime session and active daemon state.
Operations:
(none) Show current session name
list List all active sessions
Environment:
AGENT_BROWSER_SESSION Default session name
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser session
agent-browser session list
agent-browser --session test open example.com
"##
}
@@ -2162,7 +2185,6 @@ Supported URL formats:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
# Connect to local Chrome with remote debugging
@@ -2180,6 +2202,33 @@ Examples:
agent-browser click @e1
"##
}
"doctor" => {
r##"
agent-browser doctor - Diagnose CDP, sourceURL sanitization, and tab-group plugin health
Usage: agent-browser doctor
Runs a non-destructive health check focused on:
- CDP endpoint reachability (preferred :9333 + common ports)
- DevToolsActivePort discovery from local Chrome profiles
- CDP sourceURL sanitization probe (Runtime.evaluate leakage check)
- Plugin handshake page context suitability (internal page vs http(s))
- Tab-group plugin handshake status (when connected via CDP)
Notes:
- doctor does not accept positional arguments
- If browser is not already connected, doctor will still report CDP probe results
- Plugin handshake requires CDP mode, a normal http(s) page, and the extension installed
- Extension side panel includes browser controls, developer signal panels, workflow shortcuts, and scheduled automation
Global Options:
--json Output as JSON
Examples:
agent-browser doctor
agent-browser --json doctor
"##
}
// === iOS Commands ===
"tap" => {
@@ -2193,7 +2242,6 @@ for touch-based interfaces like iOS Safari.
Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser tap "#submit-button"
@@ -2216,7 +2264,6 @@ Arguments:
Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser -p ios swipe up
@@ -2235,7 +2282,6 @@ Subcommands:
Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser device list
@@ -2290,7 +2336,6 @@ URL Diff:
Global Options:
--json Output as JSON
--session <name> Use specific session
Examples:
agent-browser diff snapshot
@@ -2314,15 +2359,17 @@ pub fn print_help() {
agent-browser - fast browser automation CLI for AI agents
Usage: agent-browser <command> [args] [options]
Aliases: agent-browser, agent-browser-stealth, abs
Core Commands:
open <url> Navigate to URL
start Start managed browser on localhost:9333
click <sel> Click element (or @ref)
dblclick <sel> Double-click element
type <sel> <text> Type into element
type <sel> <text> [--delay <ms>] Type into element
fill <sel> <text> Clear and fill
press <key> Press key (Enter, Tab, Control+a)
keyboard type <text> Type text with real keystrokes (no selector)
keyboard type <text> [--delay <ms>] Type text with real keystrokes (no selector)
keyboard inserttext <text> Insert text without key events
hover <sel> Hover element
focus <sel> Focus element
@@ -2334,7 +2381,7 @@ Core Commands:
download <sel> <path> Download file by clicking element
scroll <dir> [px] Scroll (up/down/left/right)
scrollintoview <sel> Scroll element into view
wait <sel|ms> Wait for element or time
wait <sel|ms|min-max> Wait for element, time, or random range
screenshot [path] Take screenshot
pdf <path> Save as PDF
snapshot Accessibility tree with refs (for AI)
@@ -2403,11 +2450,12 @@ Confirmation:
Sessions:
session Show current session name
session list List active sessions
session list List active sessions (stale entries are auto-cleaned)
Setup:
install Install browser binaries
install --with-deps Also install system dependencies (Linux)
doctor Diagnose CDP + sourceURL + plugin health
Snapshot Options:
-i, --interactive Only interactive elements
@@ -2416,8 +2464,7 @@ Snapshot Options:
-s, --selector <sel> Scope to CSS selector
Options:
--session <name> Isolated session (or AGENT_BROWSER_SESSION env)
--profile <path> Persistent browser profile (or AGENT_BROWSER_PROFILE env)
--session <name> Ignored (runtime uses default session unless --parallel is set)
--state <path> Load storage state from JSON file (or AGENT_BROWSER_STATE env)
--headers <json> HTTP headers scoped to URL's origin (for auth)
--executable-path <path> Custom browser executable (or AGENT_BROWSER_EXECUTABLE_PATH)
@@ -2436,22 +2483,40 @@ Options:
--json JSON output
--full, -f Full page screenshot
--annotate Annotated screenshot with numbered labels and legend
--headed Show browser window (not headless)
--headed Show browser window (not headless) (or AGENT_BROWSER_HEADED=1/true)
--cdp <port> Connect via CDP (Chrome DevTools Protocol)
--auto-connect Auto-discover and connect to running Chrome
Explicit existing-browser mode; may trigger Chrome permission prompts
--color-scheme <scheme> Color scheme: dark, light, no-preference (or AGENT_BROWSER_COLOR_SCHEME)
--download-path <path> Default download directory (or AGENT_BROWSER_DOWNLOAD_PATH)
--session-name <name> Auto-save/restore session state (cookies, localStorage)
--tab-group <name> Base title for agent tab groups (CDP plugin mode; silent no-op if plugin unavailable)
--tab-group-plugin-id <id> Expected Chrome extension ID for tab-group handshake (or AGENT_BROWSER_TAB_GROUP_PLUGIN_ID)
Extension side panel supports browser controls + console/network/DOM + workflow scheduling
--risk-mode <mode> Verify/captcha handling: off, warn, block (or AGENT_BROWSER_RISK_MODE)
--wait-until <mode> Navigation wait strategy for open/navigate: load, domcontentloaded, networkidle
--parallel <name> Isolated runtime channel for parallel AI runs (maps to parallel-<name>)
Default behavior in this mode is stateless (no auto session persistence unless --session-name is explicitly passed)
Note: starting default session reaps all non-default daemon sessions
--resident Keep daemon running; disable 10-minute idle auto-shutdown
--session-name <name> Auto-save/restore session state (defaults to "default" in non-parallel mode)
--content-boundaries Wrap page output in boundary markers (or AGENT_BROWSER_CONTENT_BOUNDARIES)
--max-output <chars> Truncate page output to N chars (or AGENT_BROWSER_MAX_OUTPUT)
--allowed-domains <list> Restrict navigation domains (or AGENT_BROWSER_ALLOWED_DOMAINS)
--action-policy <path> Action policy JSON file (or AGENT_BROWSER_ACTION_POLICY)
--confirm-actions <list> Categories requiring confirmation (or AGENT_BROWSER_CONFIRM_ACTIONS)
--confirm-interactive Interactive confirmation prompts; auto-denies if stdin is not a TTY (or AGENT_BROWSER_CONFIRM_INTERACTIVE)
--engine <name> Browser engine: chrome (default), lightpanda; implies --native (or AGENT_BROWSER_ENGINE)
--native [Experimental] Use native Rust daemon instead of Node.js (or AGENT_BROWSER_NATIVE)
--config <path> Use a custom config file (or AGENT_BROWSER_CONFIG env)
--debug Debug output
--version, -V Show version
--version, -V Show version (fork builds include upstream/fork info)
Policy:
--profile / AGENT_BROWSER_PROFILE are forbidden
--channel / AGENT_BROWSER_CHANNEL are forbidden
Daemon auto-shuts down after 10 minutes of inactivity unless --resident is set
Default mode uses localhost:9333. If 9333 is unavailable, agent-browser auto-starts a dedicated Chrome profile at ~/.agent-browser/chrome-bot-profile
Use --auto-connect only when you explicitly want to attach to an existing manual browser session
Configuration:
agent-browser looks for agent-browser.json in these locations (lowest to highest priority):
@@ -2466,21 +2531,23 @@ Configuration:
Boolean flags accept an optional true/false value to override config:
--headed (same as --headed true)
--headed false (disables "headed": true from config)
--resident false (disable resident mode for this invocation)
Extensions from user and project configs are merged (not replaced).
Example agent-browser.json:
{{"headed": true, "proxy": "http://localhost:8080", "profile": "./browser-data"}}
{{"headed": true, "proxy": "http://localhost:8080", "userAgent": "my-agent/1.0"}}
Environment:
AGENT_BROWSER_CONFIG Path to config file (or use --config)
AGENT_BROWSER_SESSION Session name (default: "default")
AGENT_BROWSER_SESSION_NAME Auto-save/restore state persistence name
AGENT_BROWSER_PARALLEL Isolated runtime channel for parallel AI runs (maps to parallel-<name>)
Best for stateless/no-login tasks where throughput matters
Note: any default-session command reaps non-default daemon sessions
AGENT_BROWSER_ENCRYPTION_KEY 64-char hex key for AES-256-GCM state encryption
AGENT_BROWSER_STATE_EXPIRE_DAYS Auto-delete states older than N days (default: 30)
AGENT_BROWSER_EXECUTABLE_PATH Custom browser executable path
AGENT_BROWSER_EXTENSIONS Comma-separated browser extension paths
AGENT_BROWSER_HEADED Show browser window (not headless)
AGENT_BROWSER_HEADED Show browser window (not headless; accepts 1 or true)
AGENT_BROWSER_JSON JSON output
AGENT_BROWSER_FULL Full page screenshot
AGENT_BROWSER_ANNOTATE Annotated screenshot with numbered labels and legend
@@ -2489,10 +2556,16 @@ Environment:
AGENT_BROWSER_PROVIDER Browser provider (ios, browserbase, kernel, browseruse)
AGENT_BROWSER_AUTO_CONNECT Auto-discover and connect to running Chrome
AGENT_BROWSER_ALLOW_FILE_ACCESS Allow file:// URLs to access local files
AGENT_BROWSER_LOCALE Override auto-detected locale (e.g., zh-TW, ja-JP)
AGENT_BROWSER_TIMEZONE Override auto-detected timezone (e.g., Asia/Taipei)
AGENT_BROWSER_COLOR_SCHEME Color scheme preference (dark, light, no-preference)
AGENT_BROWSER_DOWNLOAD_PATH Default download directory for browser downloads
AGENT_BROWSER_TAB_GROUP Base title for tab groups (default: "Agent Browser Stealth"; session suffix auto-appended)
AGENT_BROWSER_TAB_GROUP_PLUGIN_ID Expected Chrome extension ID for tab-group handshake (default: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")
AGENT_BROWSER_RISK_MODE Verify/captcha handling mode (off, warn, block)
AGENT_BROWSER_DEFAULT_TIMEOUT Default Playwright timeout in ms (default: 25000)
AGENT_BROWSER_SESSION_NAME Auto-save/load state persistence name
AGENT_BROWSER_SESSION_NAME Auto-save/load state persistence name (default: "default" when --parallel is not set)
AGENT_BROWSER_STATE_EXPIRE_DAYS Auto-delete saved states older than N days (default: 30)
AGENT_BROWSER_ENCRYPTION_KEY 64-char hex key for AES-256-GCM session encryption
AGENT_BROWSER_STREAM_PORT Enable WebSocket streaming on port (e.g., 9223)
@@ -2504,14 +2577,15 @@ Environment:
AGENT_BROWSER_ACTION_POLICY Path to action policy JSON file
AGENT_BROWSER_CONFIRM_ACTIONS Action categories requiring confirmation
AGENT_BROWSER_CONFIRM_INTERACTIVE Enable interactive confirmation prompts
AGENT_BROWSER_ENGINE Browser engine: chrome (default), lightpanda
AGENT_BROWSER_NATIVE Use native Rust daemon (experimental, no Node.js/Playwright)
Install (recommended, fastest - native Rust CLI):
npm install -g agent-browser
npm install -g agent-browser-stealth
agent-browser install # Download Chromium (first time)
Try without installing (slower, routes through Node.js):
npx agent-browser open example.com
npx agent-browser-stealth open example.com
Examples:
agent-browser open example.com
@@ -2526,8 +2600,10 @@ Examples:
agent-browser --cdp 9222 snapshot # Connect via CDP port
agent-browser --auto-connect snapshot # Auto-discover running Chrome
agent-browser --color-scheme dark open example.com # Dark mode
agent-browser --profile ~/.myapp open example.com # Persistent profile
agent-browser --risk-mode block open example.com # Block on verification/captcha pages
agent-browser --session-name myapp open example.com # Auto-save/restore state
agent-browser --parallel worker-a open example.com # Isolated runtime for parallel AI task
agent-browser --resident open example.com # Keep daemon resident until explicit close
Command Chaining:
Chain commands with && in a single shell call (browser persists via daemon):
@@ -2546,6 +2622,94 @@ iOS Simulator (requires Xcode and Appium):
);
}
pub fn print_response(resp: &Response, json: bool, action: Option<&str>) {
let opts = OutputOptions {
json,
content_boundaries: false,
max_output: None,
};
print_response_with_opts(resp, action, &opts);
}
fn status_badge(status: &str) -> String {
match status {
"pass" => color::green("PASS"),
"warn" => color::yellow("WARN"),
"fail" => color::red("FAIL"),
"skip" => color::dim("SKIP"),
_ => status.to_string(),
}
}
fn print_doctor_report(data: &serde_json::Map<String, serde_json::Value>) {
let ok = data.get("ok").and_then(|v| v.as_bool()).unwrap_or(false);
let summary = if ok {
format!("{} doctor checks passed", color::success_indicator())
} else {
format!("{} doctor found issues", color::error_indicator())
};
println!("{}", summary);
if let Some(context) = data.get("context").and_then(|v| v.as_object()) {
let launched = context
.get("launched")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let connection = context
.get("connectionKind")
.and_then(|v| v.as_str())
.unwrap_or("unknown");
let session = context
.get("session")
.and_then(|v| v.as_str())
.unwrap_or("default");
let cdp_endpoint = context
.get("cdpEndpoint")
.and_then(|v| v.as_str())
.unwrap_or("-");
println!(
" context: launched={} connection={} session={} cdp={}",
launched, connection, session, cdp_endpoint
);
}
if let Some(checks) = data.get("checks").and_then(|v| v.as_array()) {
for check in checks {
let Some(obj) = check.as_object() else {
continue;
};
let name = obj
.get("name")
.and_then(|v| v.as_str())
.unwrap_or("unknown");
let status = obj
.get("status")
.and_then(|v| v.as_str())
.unwrap_or("unknown");
let message = obj.get("message").and_then(|v| v.as_str()).unwrap_or("");
println!(" [{}] {} - {}", status_badge(status), name, message);
}
}
if let Some(plugin) = data.get("plugin").and_then(|v| v.as_object()) {
let plugin_status = plugin
.get("status")
.and_then(|v| v.as_str())
.unwrap_or("unknown");
let plugin_message = plugin.get("message").and_then(|v| v.as_str()).unwrap_or("");
let plugin_id = plugin
.get("configuredPluginId")
.and_then(|v| v.as_str())
.unwrap_or("-");
println!(
" plugin: [{}] id={} {}",
status_badge(plugin_status),
plugin_id,
plugin_message
);
}
}
fn print_snapshot_diff(data: &serde_json::Map<String, serde_json::Value>) {
let changed = data
.get("changed")
@@ -2622,6 +2786,38 @@ fn print_screenshot_diff(data: &serde_json::Map<String, serde_json::Value>) {
);
}
pub fn print_version() {
println!("agent-browser {}", env!("CARGO_PKG_VERSION"));
/// Parse fork version metadata from semver-like strings:
/// <upstream>-fork.<fork>
/// Example:
/// 0.14.0-fork.1 -> (0.14.0, 1)
fn parse_fork_version(version: &str) -> Option<(&str, &str)> {
let (upstream, fork) = version.split_once("-fork.")?;
if upstream.is_empty() || fork.is_empty() {
return None;
}
if !upstream
.chars()
.all(|c| c.is_ascii_digit() || c == '.' || c == '-')
{
return None;
}
if !fork
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '-')
{
return None;
}
Some((upstream, fork))
}
pub fn print_version() {
let version = env!("CARGO_PKG_VERSION");
if let Some((upstream, fork)) = parse_fork_version(version) {
println!(
"agent-browser {} (upstream {}, fork {})",
version, upstream, fork
);
} else {
println!("agent-browser {}", version);
}
}
+49
View File
@@ -0,0 +1,49 @@
use std::sync::{Mutex, MutexGuard};
/// Global mutex shared across all test modules to prevent parallel tests from
/// interfering with each other when mutating environment variables.
pub static ENV_MUTEX: Mutex<()> = Mutex::new(());
/// RAII guard that locks [`ENV_MUTEX`] and restores environment variables on drop.
pub struct EnvGuard<'a> {
_lock: MutexGuard<'a, ()>,
vars: Vec<(String, Option<String>)>,
}
impl<'a> EnvGuard<'a> {
pub fn new(var_names: &[&str]) -> Self {
let lock = ENV_MUTEX.lock().unwrap();
let vars = var_names
.iter()
.map(|&name| (name.to_string(), std::env::var(name).ok()))
.collect();
Self { _lock: lock, vars }
}
pub fn set(&self, name: &str, value: &str) {
debug_assert!(
self.vars.iter().any(|(n, _)| n == name),
"EnvGuard::set called with unregistered var: {name}"
);
std::env::set_var(name, value);
}
pub fn remove(&self, name: &str) {
debug_assert!(
self.vars.iter().any(|(n, _)| n == name),
"EnvGuard::remove called with unregistered var: {name}"
);
std::env::remove_var(name);
}
}
impl Drop for EnvGuard<'_> {
fn drop(&mut self) {
for (name, value) in &self.vars {
match value {
Some(v) => std::env::set_var(name, v),
None => std::env::remove_var(name),
}
}
}
}
+1
View File
@@ -0,0 +1 @@
{"v":1}
+6
View File
@@ -0,0 +1,6 @@
{
"git": {
"sha1": "7cb6c7d950c040b2198da553140e1b5e8b6ac682"
},
"path_in_vcs": "crates/zune-jpeg"
}
+1
View File
@@ -0,0 +1 @@
/target
+79
View File
@@ -0,0 +1,79 @@
# Benchmarks of popular jpeg libraries
Here I compare how long it takes popular JPEG decoders to decode the below 7680*4320 image
of (now defunct ?) [Cutefish OS](https://en.cutefishos.com/) default wallpaper.
![img](benches/images/speed_bench.jpg)
## About benchmarks
Benchmarks are weird, especially IO & multi-threaded programs. This library uses both of the above hence performance may
vary.
For best results shut down your machine, go take coffee, think about life and how it came to be and why people should
save the environment.
Then power up your machine, if it's a laptop connect it to a power supply and if there is a setting for performance
mode, tweak it.
Then run.
## Benchmarks vs real world usage
Real world usage may vary.
Notice that I'm using a large image but probably most decoding will be small to medium images.
To make the library thread safe, we do about 1.5-1.7x more allocations than libjpeg-turbo. Although, do note that the
allocations do not occur at ago, we allocate when needed and deallocate when not needed.
Do note if memory bandwidth is a limitation. This is not for you.
## Reproducibility
The benchmarks are carried out on my local machine with an AMD Ryzen 5 4500u
The benchmarks are reproducible.
To reproduce them
1. Clone this repository
2. Install rust(if you don't have it yet)
3. `cd` into the directory.
4. Run `cargo bench`
## Performance features of the three libraries
| feature | image-rs/jpeg-decoder | libjpeg-turbo | zune-jpeg |
|------------------------------|-----------------------|---------------|-----------|
| multithreaded | ✅ | ❌ | ❌ |
| platform specific intrinsics | ✅ | ✅ | ✅ |
- Image-rs/jpeg-decoder uses [rayon] under the hood but it's under a feature
flag.
- libjpeg-turbo uses hand-written asm for platform specific intrinsics, ported to
the most common architectures out there but falls back to scalar
code if it can't run in a platform.
# Finally benchmarks
[here]
## Notes
Benchmarks are ran at least once a week to catch regressions early and
are uploaded to Github pages.
Machine specs can be found on the other [landing page]
Benchmarks may not reflect real world usage(threads, other I/O machine bottlenecks)
[landing page]:https://etemesi254.github.io/posts/Zune-Benchmarks/
[here]:https://etemesi254.github.io/assets/criterion/report/index.html
[libjpeg-turbo]:https://github.com/libjpeg-turbo/libjpeg-turbo
[jpeg-decoder]:https://github.com/image-rs/jpeg-decoder
[rayon]:https://github.com/rayon-rs/rayon
+25
View File
@@ -0,0 +1,25 @@
# This file is automatically @generated by Cargo.
# It is not intended for manual editing.
version = 4
[[package]]
name = "log"
version = "0.4.29"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
[[package]]
name = "zune-core"
version = "0.5.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cb8a0807f7c01457d0379ba880ba6322660448ddebc890ce29bb64da71fb40f9"
dependencies = [
"log",
]
[[package]]
name = "zune-jpeg"
version = "0.5.12"
dependencies = [
"zune-core",
]
+67
View File
@@ -0,0 +1,67 @@
# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO
#
# When uploading crates to the registry Cargo will automatically
# "normalize" Cargo.toml files for maximal compatibility
# with all versions of Cargo and also rewrite `path` dependencies
# to registry (e.g., crates.io) dependencies.
#
# If you are reading this file be aware that the original Cargo.toml
# will likely look very different (and much more reasonable).
# See Cargo.toml.orig for the original contents.
[package]
edition = "2021"
rust-version = "1.75.0"
name = "zune-jpeg"
version = "0.5.12"
authors = ["caleb <etemesicaleb@gmail.com>"]
build = false
exclude = [
"/benches/images/*",
"/tests/*",
"/.idea/*",
"/.gradle/*",
"/test-images/*",
"fuzz/*",
]
autolib = false
autobins = false
autoexamples = false
autotests = false
autobenches = false
description = "A fast, correct and safe jpeg decoder"
readme = "README.md"
keywords = [
"jpeg",
"jpeg-decoder",
"decoder",
]
categories = ["multimedia::images"]
license = "MIT OR Apache-2.0 OR Zlib"
repository = "https://github.com/etemesi254/zune-image/tree/dev/crates/zune-jpeg"
[features]
default = [
"x86",
"neon",
"std",
]
log = ["zune-core/log"]
neon = []
portable_simd = []
std = ["zune-core/std"]
x86 = []
[lib]
name = "zune_jpeg"
path = "src/lib.rs"
[dependencies.zune-core]
version = "0.5.1"
[dev-dependencies]
[lints.rust.unexpected_cfgs]
level = "warn"
priority = 0
check-cfg = ["cfg(fuzzing)"]
+36
View File
@@ -0,0 +1,36 @@
[package]
name = "zune-jpeg"
version = "0.5.12"
rust-version = "1.75.0"
authors = ["caleb <etemesicaleb@gmail.com>"]
edition = "2021"
repository = "https://github.com/etemesi254/zune-image/tree/dev/crates/zune-jpeg"
license = "MIT OR Apache-2.0 OR Zlib"
keywords = ["jpeg", "jpeg-decoder", "decoder"]
categories = ["multimedia::images"]
exclude = ["/benches/images/*", "/tests/*", "/.idea/*", "/.gradle/*", "/test-images/*", "fuzz/*"]
description = "A fast, correct and safe jpeg decoder"
[lints.rust]
# Disable feature checker for fuzzing since it's used and cargo doesn't
# seem to recognise fuzzing
unexpected_cfgs = { level = "warn", check-cfg = ['cfg(fuzzing)'] }
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[features]
x86 = []
neon = []
std = ["zune-core/std"]
# NOTE: portable_simd requires Rust 1.87+
portable_simd = []
log = ["zune-core/log"]
default = ["x86", "neon", "std"]
[dependencies]
zune-core = { path = "../zune-core", version = "0.5.1" }
[dev-dependencies]
zune-ppm = { path = "../zune-ppm" }
+95
View File
@@ -0,0 +1,95 @@
## Version 0.5.7
- Move scalar idct to wrapping maths.
- Simd upsampling (mhils)
- Faster zero idct check (mhils)
## Version 0.5.6
- Better support for truncated images (by https://github.com/mhils)
- fix 4:1:0 chroma subsampling (by https://github.com/mhils)
- Fix some crashes
- Fix some bug on last pixel sampling
## Version 0.5.5
- Support direct conversion of Luma to RGBA
## Version 0.5.4
- Fix overriding color space when decoding Luma colorspace
## Version 0.5.3
- Fix some decoding of some images with markers in progressive segments, see https://github.com/etemesi254/zune-image/issues/295
## Version 0.5.1
- Fix decoding of particular images with a non-standard subsample, (
see https://github.com/etemesi254/zune-image/issues/291)
- Add better RGB color detection of images to match libjpeg and stb_image formats
-----
## Version 0.3.17
- Fix no-std compilation
## Version 0.3.16
- Add support for decoding to BGR and BGRA
## Version 0.3.14
- Add ability to parse exif and ICC chunk.
- Fix images with one component that were down-sampled.
### Version 0.3.13
- Allow decoding into pre-allocated buffer
- Clarify documentation
### Version 0.3.11
- Add guards for SSE and AVX code paths(allows compiling for platforms that do not support it)
### Version 0.3.0
- Overhaul to the whole decoder.
- Single threaded version
- Lightweight.
### Version 0.2.0
- New `ZuneJpegOptions` struct, this is the now recommended way to set up decoding options for
decoding
- Deprecated previous options setting functions.
- More code cleanups
- Fixed new bugs discovered by fuzzing
- Removed dependency on `num_cpu`
### Version 0.1.5
- Allow user to set memory limits in during decoding explicitly via `set_limits`
- Fixed some bugs discovered by fuzzing
- Correctly handle small images less than 16 pixels
- Gracefully handle incorrectly sampled images.
### Version 0.1.4
- Remove all `unsafe` instances except platform dependent intrinsics.
- Numerous bug fixes identified by fuzzing.
- Expose `ImageInfo` to the crate root.
### Version 0.1.3
- Fix numerous panics found by fuzzing(thanks to @[Shnatsel] for the corpus)
- Add new method `set_num_threads` that allows one to explicitly set the number of threads to use to decode the image.
### Version 0.1.2
- Add more sub checks, contributed by @[5225225]
- Privatize some modules.
### Version 0.1.1
- Fix rgba/rgbx decoding when avx optimized functions were used
- Initial support for fuzzing
- Remove `align_alloc` method which was unsound (Thanks to @[HeroicKatora] for pointing that out)
[Shnatsel]:https://github.com/Shnatsel
[HeroicKatora]:https://github.com/HeroicKatora
[5225225]:https://github.com/5225225
+201
View File
@@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) zune-image developers
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
+19
View File
@@ -0,0 +1,19 @@
zlib License
(C) zune-image developers
This software is provided 'as-is', without any express or implied
warranty. In no event will the authors be held liable for any damages
arising from the use of this software.
Permission is granted to anyone to use this software for any purpose,
including commercial applications, and to alter it and redistribute it
freely, subject to the following restrictions:
1. The origin of this software must not be misrepresented; you must not
claim that you wrote the original software. If you use this software
in a product, an acknowledgment in the product documentation would be
appreciated but is not required.
2. Altered source versions must be plainly marked as such, and must not be
misrepresented as being the original software.
3. This notice may not be removed or altered from any source distribution.
+104
View File
@@ -0,0 +1,104 @@
# Zune-JPEG
A fast, correct and safe jpeg decoder in pure Rust.
## Usage
The library provides a simple-to-use API for jpeg decoding
and an ability to add options to influence decoding.
### Example
```Rust
// Import the library
use zune_jpeg::JpegDecoder;
use std::fs::read;
fn main()->Result<(),DecoderErrors> {
// load some jpeg data
let data = read("cat.jpg").unwrap();
// create a decoder
let mut decoder = JpegDecoder::new(&data);
// decode the file
let pixels = decoder.decode()?;
}
```
The decoder supports more manipulations via `DecoderOptions`,
see additional documentation in the library.
## Goals
The implementation aims to have the following goals achieved,
in order of importance
1. Safety - Do not segfault on errors or invalid input. Panics are okay, but
should be fixed when reported. `unsafe` is only used for SIMD intrinsics,
and can be turned off entirely both at compile time and at runtime.
2. Speed - Get the data as quickly as possible, which means
1. Platform intrinsics code where justifiable
2. Carefully written platform independent code that allows the
compiler to vectorize it.
3. Regression tests.
4. Watch the memory usage of the program
3. Usability - Provide utility functions like different color conversions functions.
## Non-Goals
- Bit identical results with libjpeg/libjpeg-turbo will never be an aim of this library.
Jpeg is a lossy format with very few parts specified by the standard
(i.e it doesn't give a reference upsampling and color conversion algorithm)
## Features
- [x] A Pretty fast 8*8 integer IDCT.
- [x] Fast Huffman Decoding
- [x] Fast color convert functions.
- [x] Support for extended colorspaces like GrayScale and RGBA
- [X] Single-threaded decoding.
- [X] Support for four component JPEGs, and esoteric color schemes like CYMK
- [X] Support for `no_std`
- [X] BGR/BGRA decoding support.
## Crate Features
| feature | on | Capabilities |
|---------|-----|---------------------------------------------------------------------------------------------|
| `x86` | yes | Enables `x86` specific instructions, specifically `avx` and `sse` for accelerated decoding. |
| `std` | yes | Enable linking to the `std` crate |
Note that the `x86` features are automatically disabled on platforms that aren't x86 during compile
time hence there is no need to disable them explicitly if you are targeting such a platform.
## Using in a `no_std` environment
The crate can be used in a `no_std` environment with the `alloc` feature.
But one is required to link to a working allocator for whatever environment the decoder
will be running on
## Debug vs release
The decoder heavily relies on platform specific intrinsics, namely AVX2 and SSE to gain speed-ups in decoding,
but they [perform poorly](https://godbolt.org/z/vPq57z13b) in debug builds. To get reasonable performance even
when compiling your program in debug mode, add this to your `Cargo.toml`:
```toml
# `zune-jpeg` package will be always built with optimizations
[profile.dev.package.zune-jpeg]
opt-level = 3
```
## Benchmarks
The library tries to be at fast as [libjpeg-turbo] while being as safe as possible.
Platform specific intrinsics help get speed up intensive operations ensuring we can almost
match [libjpeg-turbo] speeds but speeds are always +- 10 ms of this library.
For more up-to-date benchmarks, see the online repo with
benchmarks [here](https://etemesi254.github.io/assets/criterion/report/index.html)
[libjpeg-turbo]:https://github.com/libjpeg-turbo/libjpeg-turbo/
[image-rs/jpeg-decoder]:https://github.com/image-rs/jpeg-decoder/tree/master/src
+811
View File
@@ -0,0 +1,811 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
#![allow(
clippy::if_not_else,
clippy::similar_names,
clippy::inline_always,
clippy::doc_markdown,
clippy::cast_sign_loss,
clippy::cast_possible_truncation
)]
//! This file exposes a single struct that can decode a huffman encoded
//! Bitstream in a JPEG file
//!
//! This code is optimized for speed.
//! It's meant to be super duper super fast, because everyone else depends on this being fast.
//! It's (annoyingly) serial hence we cant use parallel bitstreams(it's variable length coding.)
//!
//! Furthermore, on the case of refills, we have to do bytewise processing because the standard decided
//! that we want to support markers in the middle of streams(seriously few people use RST markers).
//!
//! So we pull in all optimization steps:
//! - use `inline[always]`? ✅ ,
//! - pre-execute most common cases ✅,
//! - add random comments ✅
//! - fast paths ✅.
//!
//! Speed-wise: It is probably the fastest JPEG BitStream decoder to ever sail the seven seas because of
//! a couple of optimization tricks.
//! 1. Fast refills from libjpeg-turbo
//! 2. As few as possible branches in decoder fast paths.
//! 3. Accelerated AC table decoding borrowed from stb_image.h written by Fabian Gissen (@ rygorous),
//! improved by me to handle more cases.
//! 4. Safe and extensible routines(e.g. cool ways to eliminate bounds check)
//! 5. No unsafe here
//!
//! Readability comes as a second priority(I tried with variable names this time, and we are wayy better than libjpeg).
//!
//! Anyway if you are reading this it means your cool and I hope you get whatever part of the code you are looking for
//! (or learn something cool)
//!
//! Knock yourself out.
use alloc::format;
use alloc::string::ToString;
use core::cmp::min;
use zune_core::bytestream::{ZByteReaderTrait, ZReader};
use crate::errors::DecodeErrors;
use crate::huffman::{HuffmanTable, HUFF_LOOKAHEAD};
use crate::marker::Marker;
use crate::mcu::DCT_BLOCK;
use crate::misc::UN_ZIGZAG;
macro_rules! decode_huff {
($stream:tt,$symbol:tt,$table:tt) => {
let mut code_length = $symbol >> HUFF_LOOKAHEAD;
($symbol) &= (1 << HUFF_LOOKAHEAD) - 1;
if code_length > i32::from(HUFF_LOOKAHEAD)
{
// if the symbol cannot be resolved in the first HUFF_LOOKAHEAD bits,
// we know it lies somewhere between HUFF_LOOKAHEAD and 16 bits since jpeg imposes 16 bit
// limit, we can therefore look 16 bits ahead and try to resolve the symbol
// starting from 1+HUFF_LOOKAHEAD bits.
$symbol = ($stream).peek_bits::<16>() as i32;
// (Credits to Sean T. Barrett stb library for this optimization)
// maxcode is pre-shifted 16 bytes long so that it has (16-code_length)
// zeroes at the end hence we do not need to shift in the inner loop.
while code_length < 17{
if $symbol < $table.maxcode[code_length as usize] {
break;
}
code_length += 1;
}
if code_length == 17{
// symbol could not be decoded.
//
// We may think, lets fake zeroes, noo
// panic, because Huffman codes are sensitive, probably everything
// after this will be corrupt, so no need to continue.
// panic!("Bad Huffman code length");
return Err(DecodeErrors::Format(format!("Bad Huffman Code 0x{:X}, corrupt JPEG",$symbol)))
}
$symbol >>= (16-code_length);
($symbol) = i32::from(
($table).values
[(($symbol + ($table).offset[code_length as usize]) & 0xFF) as usize],
);
}
if code_length> i32::from(($stream).bits_left){
return Err(DecodeErrors::Format(format!("Code length {code_length} more than bits left {}",($stream).bits_left)))
}
// drop bits read
($stream).drop_bits(code_length as u8);
};
}
/// A `BitStream` struct, a bit by bit reader with super powers
///
#[rustfmt::skip]
pub(crate) struct BitStream {
/// A MSB type buffer that is used for some certain operations
pub buffer: u64,
/// A TOP aligned MSB type buffer that is used to accelerate some operations like
/// peek_bits and get_bits.
///
/// By top aligned, I mean the top bit (63) represents the top bit in the buffer.
aligned_buffer: u64,
/// Tell us the bits left the two buffer
pub(crate) bits_left: u8,
/// Did we find a marker(RST/EOF) during decoding?
pub marker: Option<Marker>,
/// An i16 with the bit corresponding to successive_low set to 1, others 0.
pub successive_low_mask: i16,
spec_start: u8,
spec_end: u8,
pub eob_run: i32,
pub overread_by: usize,
/// True if we have seen end of image marker.
/// Don't read anything after that.
pub seen_eoi: bool,
}
impl BitStream {
/// Create a new BitStream
#[rustfmt::skip]
pub(crate) const fn new() -> BitStream {
BitStream {
buffer: 0,
aligned_buffer: 0,
bits_left: 0,
marker: None,
successive_low_mask: 1,
spec_start: 0,
spec_end: 0,
eob_run: 0,
overread_by: 0,
seen_eoi: false,
}
}
/// Create a new Bitstream for progressive decoding
#[allow(clippy::redundant_field_names)]
#[rustfmt::skip]
pub(crate) fn new_progressive(al: u8, spec_start: u8, spec_end: u8) -> BitStream {
BitStream {
buffer: 0,
aligned_buffer: 0,
bits_left: 0,
marker: None,
successive_low_mask: 1i16 << al,
spec_start: spec_start,
spec_end: spec_end,
eob_run: 0,
overread_by: 0,
seen_eoi: false,
}
}
/// Refill the bit buffer by (a maximum of) 32 bits
///
/// # Arguments
/// - `reader`:`&mut BufReader<R>`: A mutable reference to an underlying
/// File/Memory buffer containing a valid JPEG stream
///
/// This function will only refill if `self.count` is less than 32
#[inline(always)] // to many call sites? ( perf improvement by 4%)
pub fn refill<T>(&mut self, reader: &mut ZReader<T>) -> Result<bool, DecodeErrors>
where
T: ZByteReaderTrait
{
/// Macro version of a single byte refill.
/// Arguments
/// buffer-> our io buffer, because rust macros cannot get values from
/// the surrounding environment bits_left-> number of bits left
/// to full refill
macro_rules! refill {
($buffer:expr,$byte:expr,$bits_left:expr) => {
// read a byte from the stream
$byte = u64::from(reader.read_u8());
self.overread_by += usize::from(reader.eof()?);
// append to the buffer
// JPEG is a MSB type buffer so that means we append this
// to the lower end (0..8) of the buffer and push the rest bits above..
$buffer = ($buffer << 8) | $byte;
// Increment bits left
$bits_left += 8;
// Check for special case of OxFF, to see if it's a stream or a marker
if $byte == 0xff {
// read next byte
let mut next_byte = u64::from(reader.read_u8());
// Byte snuffing, if we encounter byte snuff, we skip the byte
if next_byte != 0x00 {
// skip that byte we read
while next_byte == 0xFF {
next_byte = u64::from(reader.read_u8());
}
if next_byte != 0x00 {
// Undo the byte append and return
$buffer >>= 8;
$bits_left -= 8;
if $bits_left != 0 {
self.aligned_buffer = $buffer << (64 - $bits_left);
}
let marker = Marker::from_u8(next_byte as u8);
self.marker = marker;
if let Some(Marker::UNKNOWN(_)) = marker{
return Err(DecodeErrors::Format("Unknown marker in bit stream".to_string()));
}
if next_byte == 0xD9 {
// special handling for eoi, fill some bytes,even if its zero,
// removes some panics
self.buffer <<= 8;
self.bits_left += 8;
self.aligned_buffer = self.buffer << (64 - self.bits_left);
}
return Ok(false);
}
}
}
};
}
// 32 bits is enough for a decode(16 bits) and receive_extend(max 16 bits)
if self.bits_left < 32 {
if self.marker.is_some() || self.overread_by > 0 || self.seen_eoi {
// found a marker, or we are in EOI
// also we are in over-reading mode, where we fill it with zeroes
// fill with zeroes
self.buffer <<= 32;
self.bits_left += 32;
self.aligned_buffer = self.buffer << (64 - self.bits_left);
return Ok(true);
}
// we optimize for the case where we don't have 255 in the stream and have 4 bytes left
// as it is the common case
//
// so we always read 4 bytes, if read_fixed_bytes errors out, the cursor is
// guaranteed not to advance in case of failure (is this true), so
// we revert the read later on (if we have 255), if this fails, we use the normal
// byte at a time read
if let Ok(bytes) = reader.read_fixed_bytes_or_error::<4>() {
// we have 4 bytes to spare, read the 4 bytes into a temporary buffer
// create buffer
let msb_buf = u32::from_be_bytes(bytes);
// check if we have 0xff
if !has_byte(msb_buf, 255) {
self.bits_left += 32;
self.buffer <<= 32;
self.buffer |= u64::from(msb_buf);
self.aligned_buffer = self.buffer << (64 - self.bits_left);
return Ok(true);
}
reader.rewind(4)?;
}
// This serves two reasons,
// 1: Make clippy shut up
// 2: Favour register reuse
let mut byte;
// 4 refills, if all succeed the stream should contain enough bits to decode a
// value
refill!(self.buffer, byte, self.bits_left);
refill!(self.buffer, byte, self.bits_left);
refill!(self.buffer, byte, self.bits_left);
refill!(self.buffer, byte, self.bits_left);
// Construct an MSB buffer whose top bits are the bitstream we are currently holding.
self.aligned_buffer = self.buffer << (64 - self.bits_left);
}
return Ok(true);
}
/// Decode the DC coefficient in a MCU block.
///
/// The decoded coefficient is written to `dc_prediction`
///
#[allow(
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
clippy::unwrap_used
)]
#[inline(always)]
fn decode_dc<T>(
&mut self, reader: &mut ZReader<T>, dc_table: &HuffmanTable, dc_prediction: &mut i32
) -> Result<bool, DecodeErrors>
where
T: ZByteReaderTrait
{
let (mut symbol, r);
if self.bits_left < 32 {
self.refill(reader)?;
};
// look a head HUFF_LOOKAHEAD bits into the bitstream
symbol = self.peek_bits::<HUFF_LOOKAHEAD>();
symbol = dc_table.lookup[symbol as usize];
decode_huff!(self, symbol, dc_table);
if symbol != 0 {
r = self.get_bits(symbol as u8);
symbol = huff_extend(r, symbol);
}
// Update DC prediction
*dc_prediction = dc_prediction.wrapping_add(symbol);
return Ok(true);
}
/// Like `decode_dc` but we do not need the result of the component, we only want to remove it
/// from the bitstream of the MCU.
fn discard_dc<T>(
&mut self, reader: &mut ZReader<T>, dc_table: &HuffmanTable
) -> Result<bool, DecodeErrors>
where
T: ZByteReaderTrait
{
let mut symbol;
if self.bits_left < 32 {
self.refill(reader)?;
};
// look a head HUFF_LOOKAHEAD bits into the bitstream
symbol = self.peek_bits::<HUFF_LOOKAHEAD>();
symbol = dc_table.lookup[symbol as usize];
decode_huff!(self, symbol, dc_table);
if symbol != 0 {
let _ = self.get_bits(symbol as u8);
}
return Ok(true);
}
/// Decode a Minimum Code Unit(MCU) as quickly as possible
///
/// # Arguments
/// - reader: The bitstream from where we read more bits.
/// - dc_table: The Huffman table used to decode the DC coefficient
/// - ac_table: The Huffman table used to decode AC values
/// - block: A memory region where we will write out the decoded values
/// - DC prediction: Last DC value for this component
///
#[allow(
clippy::many_single_char_names,
clippy::cast_possible_truncation,
clippy::cast_sign_loss
)]
#[inline(never)]
pub fn decode_mcu_block<T>(
&mut self, reader: &mut ZReader<T>, dc_table: &HuffmanTable, ac_table: &HuffmanTable,
qt_table: &[i32; DCT_BLOCK], block: &mut [i32; 64], dc_prediction: &mut i32
) -> Result<u16, DecodeErrors>
where
T: ZByteReaderTrait
{
// Get fast AC table as a reference before we enter the hot path
let ac_lookup = ac_table.ac_lookup.as_ref().unwrap();
let (mut symbol, mut r, mut fast_ac);
// Decode AC coefficients
let mut pos: usize = 1;
if self.bits_left < 1 && self.marker.is_some() {
return Err(DecodeErrors::Format(
"No more bytes left in stream before marker".to_string()
));
}
// decode DC, dc prediction will contain the value
self.decode_dc(reader, dc_table, dc_prediction)?;
// set dc to be the dc prediction.
block[0] = *dc_prediction * qt_table[0];
while pos < 64 {
self.refill(reader)?;
symbol = self.peek_bits::<HUFF_LOOKAHEAD>();
fast_ac = ac_lookup[symbol as usize];
symbol = ac_table.lookup[symbol as usize];
if fast_ac != 0 {
// FAST AC path
pos += ((fast_ac >> 4) & 15) as usize; // run
let t_pos = UN_ZIGZAG[min(pos, 63)] & 63;
block[t_pos] = i32::from(fast_ac >> 8) * (qt_table[t_pos]); // Value
self.drop_bits((fast_ac & 15) as u8);
pos += 1;
} else {
decode_huff!(self, symbol, ac_table);
r = symbol >> 4;
symbol &= 15;
if symbol != 0 {
pos += r as usize;
r = self.get_bits(symbol as u8);
symbol = huff_extend(r, symbol);
let t_pos = UN_ZIGZAG[pos & 63] & 63;
block[t_pos] = symbol * qt_table[t_pos];
pos += 1;
} else if r != 15 {
return Ok(pos as u16);
} else {
pos += 16;
}
}
}
return Ok(64);
}
/// Advance the bitstream over a block but ignore the data contained.
///
/// This updates DC prediction but we never dequantize and we never do any Zig-Zag translation
/// either. Still returns the index of the last component read.
pub fn discard_mcu_block<T>(
&mut self, reader: &mut ZReader<T>, dc_table: &HuffmanTable, ac_table: &HuffmanTable
) -> Result<u16, DecodeErrors>
where
T: ZByteReaderTrait
{
// Get fast AC table as a reference before we enter the hot path
let ac_lookup = ac_table.ac_lookup.as_ref().unwrap();
let (mut symbol, mut r, mut fast_ac);
// Decode AC coefficients
let mut pos: usize = 1;
// decode DC, dc prediction will contain the value
self.discard_dc(reader, dc_table)?;
while pos < 64 {
self.refill(reader)?;
symbol = self.peek_bits::<HUFF_LOOKAHEAD>();
fast_ac = ac_lookup[symbol as usize];
symbol = ac_table.lookup[symbol as usize];
if fast_ac != 0 {
// FAST AC path
pos += ((fast_ac >> 4) & 15) as usize; // run
self.drop_bits((fast_ac & 15) as u8);
pos += 1;
} else {
decode_huff!(self, symbol, ac_table);
r = symbol >> 4;
symbol &= 15;
if symbol != 0 {
pos += r as usize;
// Advance over bits but ignore.
let _ = self.get_bits(symbol as u8);
pos += 1;
} else if r != 15 {
return Ok(pos as u16);
} else {
pos += 16;
}
}
}
return Ok(64);
}
/// Peek `look_ahead` bits ahead without discarding them from the buffer
#[inline(always)]
#[allow(clippy::cast_possible_truncation)]
const fn peek_bits<const LOOKAHEAD: u8>(&self) -> i32 {
(self.aligned_buffer >> (64 - LOOKAHEAD)) as i32
}
/// Discard the next `N` bits without checking
#[inline]
fn drop_bits(&mut self, n: u8) {
// PS: Its a good check, but triggers fuzzer and a lot of false positives
//debug_assert!(self.bits_left >= n);
//self.bits_left -= n;
self.bits_left = self.bits_left.saturating_sub(n);
self.aligned_buffer <<= n;
}
/// Read `n_bits` from the buffer and discard them
#[inline(always)]
#[allow(clippy::cast_possible_truncation)]
fn get_bits(&mut self, n_bits: u8) -> i32 {
let mask = (1_u64 << n_bits) - 1;
self.aligned_buffer = self.aligned_buffer.rotate_left(u32::from(n_bits));
let bits = (self.aligned_buffer & mask) as i32;
self.bits_left = self.bits_left.wrapping_sub(n_bits);
bits
}
/// Decode a DC block
#[allow(clippy::cast_possible_truncation)]
#[inline]
pub(crate) fn decode_prog_dc_first<T>(
&mut self, reader: &mut ZReader<T>, dc_table: &HuffmanTable, block: &mut i16,
dc_prediction: &mut i32
) -> Result<(), DecodeErrors>
where
T: ZByteReaderTrait
{
self.decode_dc(reader, dc_table, dc_prediction)?;
*block = (*dc_prediction as i16).wrapping_mul(self.successive_low_mask);
return Ok(());
}
#[inline]
pub(crate) fn decode_prog_dc_refine<T>(
&mut self, reader: &mut ZReader<T>, block: &mut i16
) -> Result<(), DecodeErrors>
where
T: ZByteReaderTrait
{
// refinement scan
if self.bits_left < 1 {
self.refill(reader)?;
// if we find a marker, it may happens we don't refill.
// So let's confirm again that refill worked
if self.bits_left < 1 {
return Err(DecodeErrors::Format(
"Marker found where not expected in refine bit".to_string()
));
}
}
if self.get_bit() == 1 {
*block = block.wrapping_add(self.successive_low_mask);
}
Ok(())
}
/// Get a single bit from the bitstream
fn get_bit(&mut self) -> u8 {
let k = (self.aligned_buffer >> 63) as u8;
// discard a bit
self.drop_bits(1);
return k;
}
pub(crate) fn decode_mcu_ac_first<T>(
&mut self, reader: &mut ZReader<T>, ac_table: &HuffmanTable, block: &mut [i16; 64]
) -> Result<bool, DecodeErrors>
where
T: ZByteReaderTrait
{
let fast_ac = ac_table.ac_lookup.as_ref().unwrap();
let bit = self.successive_low_mask;
let mut k = self.spec_start as usize;
let (mut symbol, mut r, mut fac);
// EOB runs are handled in mcu_prog.rs
'block: loop {
self.refill(reader)?;
// Check for marker in the stream
symbol = self.peek_bits::<HUFF_LOOKAHEAD>();
fac = fast_ac[symbol as usize];
symbol = ac_table.lookup[symbol as usize];
if fac != 0 {
// fast ac path
k += ((fac >> 4) & 15) as usize; // run
block[UN_ZIGZAG[min(k, 63)] & 63] = (fac >> 8).wrapping_mul(bit); // value
self.drop_bits((fac & 15) as u8);
k += 1;
} else {
decode_huff!(self, symbol, ac_table);
r = symbol >> 4;
symbol &= 15;
if symbol != 0 {
k += r as usize;
r = self.get_bits(symbol as u8);
symbol = huff_extend(r, symbol);
block[UN_ZIGZAG[k & 63] & 63] = (symbol as i16).wrapping_mul(bit);
k += 1;
} else {
if r != 15 {
self.eob_run = 1 << r;
self.eob_run += self.get_bits(r as u8);
self.eob_run -= 1;
break;
}
k += 16;
}
}
if k > self.spec_end as usize {
break 'block;
}
}
return Ok(true);
}
#[allow(clippy::too_many_lines, clippy::op_ref)]
pub(crate) fn decode_mcu_ac_refine<T>(
&mut self, reader: &mut ZReader<T>, table: &HuffmanTable, block: &mut [i16; 64]
) -> Result<bool, DecodeErrors>
where
T: ZByteReaderTrait
{
let bit = self.successive_low_mask;
let mut k = self.spec_start;
let (mut symbol, mut r);
if self.eob_run == 0 {
'no_eob: loop {
// Decode a coefficient from the bit stream
self.refill(reader)?;
symbol = self.peek_bits::<HUFF_LOOKAHEAD>();
symbol = table.lookup[symbol as usize];
decode_huff!(self, symbol, table);
r = symbol >> 4;
symbol &= 15;
if symbol == 0 {
if r != 15 {
// EOB run is 2^r + bits
self.eob_run = 1 << r;
self.eob_run += self.get_bits(r as u8);
// EOB runs are handled by the eob logic
break 'no_eob;
}
} else {
if symbol != 1 {
return Err(DecodeErrors::HuffmanDecode(
"Bad Huffman code, corrupt JPEG?".to_string()
));
}
// get sign bit
// We assume we have enough bits, which should be correct for sane images
// since we refill by 32 above
if self.get_bit() == 1 {
symbol = i32::from(bit);
} else {
symbol = i32::from(-bit);
}
}
// Advance over already nonzero coefficients appending
// correction bits to the non-zeroes.
// A correction bit is 1 if the absolute value of the coefficient must be increased
if k <= self.spec_end {
'advance_nonzero: loop {
let coefficient = &mut block[UN_ZIGZAG[k as usize & 63] & 63];
if *coefficient != 0 {
if self.bits_left < 1 {
self.refill(reader)?;
if self.bits_left < 1 && self.marker.is_some() {
return Err(DecodeErrors::Format(
"Marker found where not expected in refine bit".to_string()
));
}
}
if self.get_bit() == 1 && (*coefficient & bit) == 0 {
if *coefficient > 0 {
*coefficient += bit;
} else {
*coefficient -= bit;
}
}
} else {
r -= 1;
if r < 0 {
// reached target zero coefficient.
break 'advance_nonzero;
}
};
if k == self.spec_end {
break 'advance_nonzero;
}
k += 1;
}
}
if symbol != 0 {
let pos = UN_ZIGZAG[k as usize & 63];
// output new non-zero coefficient.
block[pos & 63] = symbol as i16;
}
k += 1;
if k > self.spec_end {
break 'no_eob;
}
}
}
if self.eob_run > 0 {
// only run if block does not consists of purely zeroes
if &block[1..] != &[0; 63] {
self.refill(reader)?;
while k <= self.spec_end {
let coefficient = &mut block[UN_ZIGZAG[k as usize & 63] & 63];
if *coefficient != 0 && self.get_bit() == 1 {
// check if we already modified it, if so do nothing, otherwise
// append the correction bit.
if (*coefficient & bit) == 0 {
if *coefficient >= 0 {
*coefficient = coefficient.wrapping_add(bit);
} else {
*coefficient = coefficient.wrapping_sub(bit);
}
}
}
if self.bits_left < 1 {
// refill at the last possible moment
self.refill(reader)?;
}
k += 1;
}
}
// count a block completed in EOB run
self.eob_run -= 1;
}
return Ok(true);
}
pub fn update_progressive_params(&mut self, _ah: u8, al: u8, spec_start: u8, spec_end: u8) {
self.successive_low_mask = 1i16 << al;
self.spec_start = spec_start;
self.spec_end = spec_end;
}
/// Reset the stream if we have a restart marker
///
/// Restart markers indicate drop those bits in the stream and zero out
/// everything
#[cold]
pub fn reset(&mut self) {
self.bits_left = 0;
self.marker = None;
self.buffer = 0;
self.aligned_buffer = 0;
self.eob_run = 0;
}
}
/// Do the equivalent of JPEG HUFF_EXTEND
#[inline(always)]
fn huff_extend(x: i32, s: i32) -> i32 {
// if x<s return x else return x+offset[s] where offset[s] = ( (-1<<s)+1)
(x) + ((((x) - (1 << ((s) - 1))) >> 31) & (((-1) << (s)) + 1))
}
const fn has_zero(v: u32) -> bool {
// Retrieved from Stanford bithacks
// @ https://graphics.stanford.edu/~seander/bithacks.html#ZeroInWord
return !((((v & 0x7F7F_7F7F) + 0x7F7F_7F7F) | v) | 0x7F7F_7F7F) != 0;
}
const fn has_byte(b: u32, val: u8) -> bool {
// Retrieved from Stanford bithacks
// @ https://graphics.stanford.edu/~seander/bithacks.html#ZeroInWord
has_zero(b ^ ((!0_u32 / 255) * (val as u32)))
}
// mod tests {
// use zune_core::bytestream::ZCursor;
// use zune_core::colorspace::ColorSpace;
// use zune_core::options::DecoderOptions;
//
// use crate::JpegDecoder;
//
// #[test]
// fn test_image() {
// let img = "/Users/etemesi/Downloads/test_IDX_45_RAND_168601280367171438891916_minimized_837.jpg";
// let data = std::fs::read(img).unwrap();
// let options = DecoderOptions::new_cmd().jpeg_set_out_colorspace(ColorSpace::RGB);
// let mut decoder = JpegDecoder::new_with_options(ZCursor::new(&data[..]), options);
//
// decoder.decode().unwrap();
// println!("{:?}", decoder.options.jpeg_get_out_colorspace())
// }
// }
+102
View File
@@ -0,0 +1,102 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
#![allow(
clippy::many_single_char_names,
clippy::similar_names,
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
clippy::cast_possible_wrap,
clippy::too_many_arguments,
clippy::doc_markdown
)]
//! Color space conversion routines
//!
//! This files exposes functions to convert one colorspace to another in a jpeg
//! image
//!
//! Currently supported conversions are
//!
//! - `YCbCr` to `RGB,RGBA,GRAYSCALE,RGBX`.
//!
//!
//! Hey there, if your reading this it means you probably need something, so let me help you.
//!
//! There are 3 supported cpu extensions here.
//! 1. Scalar
//! 2. SSE
//! 3. AVX
//!
//! There are two types of the color convert functions
//!
//! 1. Acts on 16 pixels.
//! 2. Acts on 8 pixels.
//!
//! The reason for this is because when implementing the AVX part it occurred to me that we can actually
//! do better and process 2 MCU's if we change IDCT return type to be `i16's`, since a lot of
//! CPU's these days support AVX extensions, it becomes nice if we optimize for that path ,
//! therefore AVX routines can process 16 pixels directly and SSE and Scalar just compensate.
//!
//! By compensating, I mean I wrote the 16 pixels version operating on the 8 pixel version twice.
//!
//! Therefore if your looking to optimize some routines, probably start there.
pub use scalar::ycbcr_to_grayscale;
use zune_core::colorspace::ColorSpace;
use zune_core::options::DecoderOptions;
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[cfg(feature = "x86")]
pub use crate::color_convert::avx::{ycbcr_to_rgb_avx2, ycbcr_to_rgba_avx2};
use crate::decoder::ColorConvert16Ptr;
mod avx;
mod neon64;
mod scalar;
#[allow(unused_variables)]
pub fn choose_ycbcr_to_rgb_convert_func(
type_need: ColorSpace, options: &DecoderOptions
) -> Option<ColorConvert16Ptr> {
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[cfg(feature = "x86")]
{
use zune_core::log::debug;
if options.use_avx2() {
debug!("Using AVX optimised color conversion functions");
// I believe avx2 means sse4 is also available
// match colorspace
match type_need {
ColorSpace::RGB => return Some(ycbcr_to_rgb_avx2),
ColorSpace::RGBA => return Some(ycbcr_to_rgba_avx2),
_ => () // fall through to scalar, which has more types
};
}
}
#[cfg(all(feature = "neon", target_arch = "aarch64"))]
{
if options.use_neon() {
use crate::color_convert::neon64::{ycbcr_to_rgb_neon, ycbcr_to_rgba_neon};
match type_need {
ColorSpace::RGB => return Some(ycbcr_to_rgb_neon),
ColorSpace::RGBA => return Some(ycbcr_to_rgba_neon),
_ => () // fall through to scalar, which has more types
};
}
}
// when there is no x86 or we haven't returned by here, resort to scalar
return match type_need {
ColorSpace::RGB => Some(scalar::ycbcr_to_rgb_inner_16_scalar::<false>),
ColorSpace::RGBA => Some(scalar::ycbcr_to_rgba_inner_16_scalar::<false>),
ColorSpace::BGRA => Some(scalar::ycbcr_to_rgba_inner_16_scalar::<true>),
ColorSpace::BGR => Some(scalar::ycbcr_to_rgb_inner_16_scalar::<true>),
_ => None
};
}
+297
View File
@@ -0,0 +1,297 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
//! AVX color conversion routines
//!
//! Okay these codes are cool
//!
//! Herein lies super optimized codes to do color conversions.
//!
//!
//! 1. The YCbCr to RGB use integer approximations and not the floating point equivalent.
//! That means we may be +- 2 of pixels generated by libjpeg-turbo jpeg decoding
//! (also libjpeg uses routines like `Y = 0.29900 * R + 0.33700 * G + 0.11400 * B + 0.25000 * G`)
//!
//! Firstly, we use integers (fun fact:there is no part of this code base where were dealing with
//! floating points.., fun fact: the first fun fact wasn't even fun.)
//!
//! Secondly ,we have cool clamping code, especially for rgba , where we don't need clamping and we
//! spend our time cursing that Intel decided permute instructions to work like 2 128 bit vectors(the compiler opitmizes
//! it out to something cool).
//!
//! There isn't a lot here (not as fun as bitstream ) but I hope you find what you're looking for.
//!
//! O and ~~subscribe to my youtube channel~~
#![cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#![cfg(feature = "x86")]
#![allow(
clippy::wildcard_imports,
clippy::cast_possible_truncation,
clippy::too_many_arguments,
clippy::inline_always,
clippy::doc_markdown,
dead_code
)]
#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;
use crate::color_convert::scalar::{CB_CF, CR_CF, C_G_CB_COEF_2, C_G_CR_COEF_1, YUV_RND, Y_CF};
pub union YmmRegister {
// both are 32 when using std::mem::size_of
mm256: __m256i,
// for avx color conversion
array: [i16; 16]
}
const R_AVX_COEF: i32 = i32::from_ne_bytes([CR_CF.to_ne_bytes()[0], CR_CF.to_ne_bytes()[1], 0, 0]);
const B_AVX_COEF: i32 = i32::from_ne_bytes([0, 0, CB_CF.to_ne_bytes()[0], CB_CF.to_ne_bytes()[1]]);
const G_COEF_AVX_COEF: i32 = i32::from_ne_bytes([
C_G_CR_COEF_1.to_ne_bytes()[0],
C_G_CR_COEF_1.to_ne_bytes()[1],
C_G_CB_COEF_2.to_ne_bytes()[0],
C_G_CB_COEF_2.to_ne_bytes()[1]
]);
//--------------------------------------------------------------------------------------------------
// AVX conversion routines
//--------------------------------------------------------------------------------------------------
///
/// Convert YCBCR to RGB using AVX instructions
///
/// # Note
///**IT IS THE RESPONSIBILITY OF THE CALLER TO CALL THIS IN CPUS SUPPORTING
/// AVX2 OTHERWISE THIS IS UB**
///
/// *Peace*
///
/// This library itself will ensure that it's never called in CPU's not
/// supporting AVX2
///
/// # Arguments
/// - `y`,`cb`,`cr`: A reference of 8 i32's
/// - `out`: The output array where we store our converted items
/// - `offset`: The position from 0 where we write these RGB values
#[inline(always)]
pub fn ycbcr_to_rgb_avx2(
y: &[i16; 16], cb: &[i16; 16], cr: &[i16; 16], out: &mut [u8], offset: &mut usize
) {
// call this in another function to tell RUST to vectorize this
// storing
unsafe {
ycbcr_to_rgb_avx2_1(y, cb, cr, out, offset);
}
}
#[inline]
#[target_feature(enable = "avx2")]
unsafe fn ycbcr_to_rgb_avx2_1(
y: &[i16; 16], cb: &[i16; 16], cr: &[i16; 16], out: &mut [u8], offset: &mut usize
) {
let (mut r, mut g, mut b) = ycbcr_to_rgb_baseline_no_clamp(y, cb, cr);
r = _mm256_packus_epi16(r, _mm256_setzero_si256());
g = _mm256_packus_epi16(g, _mm256_setzero_si256());
b = _mm256_packus_epi16(b, _mm256_setzero_si256());
r = _mm256_permute4x64_epi64::<{ shuffle(3, 1, 2, 0) }>(r);
g = _mm256_permute4x64_epi64::<{ shuffle(3, 1, 2, 0) }>(g);
b = _mm256_permute4x64_epi64::<{ shuffle(3, 1, 2, 0) }>(b);
let sh_r = _mm256_setr_epi8(
0, 11, 6, 1, 12, 7, 2, 13, 8, 3, 14, 9, 4, 15, 10, 5, 0, 11, 6, 1, 12, 7, 2, 13, 8, 3, 14,
9, 4, 15, 10, 5
);
let sh_g = _mm256_setr_epi8(
5, 0, 11, 6, 1, 12, 7, 2, 13, 8, 3, 14, 9, 4, 15, 10, 5, 0, 11, 6, 1, 12, 7, 2, 13, 8, 3,
14, 9, 4, 15, 10
);
let sh_b = _mm256_setr_epi8(
10, 5, 0, 11, 6, 1, 12, 7, 2, 13, 8, 3, 14, 9, 4, 15, 10, 5, 0, 11, 6, 1, 12, 7, 2, 13, 8,
3, 14, 9, 4, 15
);
let r0 = _mm256_shuffle_epi8(r, sh_r);
let g0 = _mm256_shuffle_epi8(g, sh_g);
let b0 = _mm256_shuffle_epi8(b, sh_b);
let m0 = _mm256_setr_epi8(
0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1,
0, 0, -1, 0, 0
);
let m1 = _mm256_setr_epi8(
0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0, 0, -1, 0, 0, -1, 0, 0, -1, 0, 0,
-1, 0, 0, -1, 0
);
let p0 = _mm256_blendv_epi8(_mm256_blendv_epi8(r0, g0, m0), b0, m1);
let p1 = _mm256_blendv_epi8(_mm256_blendv_epi8(g0, b0, m0), r0, m1);
let p2 = _mm256_blendv_epi8(_mm256_blendv_epi8(b0, r0, m0), g0, m1);
let rgb0 = _mm256_permute2x128_si256::<32>(p0, p1);
let rgb1 = _mm256_permute2x128_si256::<48>(p2, p0);
_mm256_storeu_si256(out.as_mut_ptr().cast(), rgb0);
_mm_storeu_si128(out[32..].as_mut_ptr().cast(), _mm256_castsi256_si128(rgb1));
*offset += 48;
}
// Enabled avx2 automatically enables avx.
#[inline]
#[target_feature(enable = "avx2")]
/// A baseline implementation of YCbCr to RGB conversion which does not carry
/// out clamping
///
/// This is used by the `ycbcr_to_rgba_avx` and `ycbcr_to_rgbx` conversion
/// routines
unsafe fn ycbcr_to_rgb_baseline_no_clamp(
y: &[i16; 16], cb: &[i16; 16], cr: &[i16; 16]
) -> (__m256i, __m256i, __m256i) {
// Load values into a register
//
let y_c = _mm256_loadu_si256(y.as_ptr().cast());
let cb_c = _mm256_loadu_si256(cb.as_ptr().cast());
let cr_c = _mm256_loadu_si256(cr.as_ptr().cast());
// Here we want to use _mm256_madd_epi16 to perform 2 multiplications
// and one addition per instruction.
// At first, we have to pack i16 U and V that stores u8 into one u8 [U,V]
// then zero extend, and keep in mind that lanes is already been permuted.
let y_coeff = _mm256_set1_epi32(i32::from(Y_CF));
let cr_coeff = _mm256_set1_epi32(R_AVX_COEF);
let cb_coeff = _mm256_set1_epi32(B_AVX_COEF);
let cg_coeff = _mm256_set1_epi32(G_COEF_AVX_COEF);
let v_rnd = _mm256_set1_epi32(i32::from(YUV_RND));
let uv_bias = _mm256_set1_epi16(128);
// UV in memory because x86/x86_64 is always little endian
let v_0 = _mm256_slli_epi16::<8>(cb_c);
let u_v_8 = _mm256_or_si256(v_0, cr_c);
let mut u_v_lo = _mm256_unpacklo_epi8(u_v_8, _mm256_setzero_si256());
let mut u_v_hi = _mm256_unpackhi_epi8(u_v_8, _mm256_setzero_si256());
let mut y_lo = _mm256_unpacklo_epi16(y_c, _mm256_setzero_si256());
let mut y_hi = _mm256_unpackhi_epi16(y_c, _mm256_setzero_si256());
u_v_lo = _mm256_sub_epi16(u_v_lo, uv_bias);
u_v_hi = _mm256_sub_epi16(u_v_hi, uv_bias);
y_lo = _mm256_madd_epi16(y_lo, y_coeff);
y_hi = _mm256_madd_epi16(y_hi, y_coeff);
let mut r_lo = _mm256_madd_epi16(u_v_lo, cr_coeff);
let mut r_hi = _mm256_madd_epi16(u_v_hi, cr_coeff);
let mut g_lo = _mm256_madd_epi16(u_v_lo, cg_coeff);
let mut g_hi = _mm256_madd_epi16(u_v_hi, cg_coeff);
// This ordering is preferred to reduce register file pressure.
y_lo = _mm256_add_epi32(y_lo, v_rnd);
y_hi = _mm256_add_epi32(y_hi, v_rnd);
let mut b_lo = _mm256_madd_epi16(u_v_lo, cb_coeff);
let mut b_hi = _mm256_madd_epi16(u_v_hi, cb_coeff);
r_lo = _mm256_add_epi32(r_lo, y_lo);
r_hi = _mm256_add_epi32(r_hi, y_hi);
g_lo = _mm256_add_epi32(g_lo, y_lo);
g_hi = _mm256_add_epi32(g_hi, y_hi);
b_lo = _mm256_add_epi32(b_lo, y_lo);
b_hi = _mm256_add_epi32(b_hi, y_hi);
r_lo = _mm256_srai_epi32::<14>(r_lo);
r_hi = _mm256_srai_epi32::<14>(r_hi);
g_lo = _mm256_srai_epi32::<14>(g_lo);
g_hi = _mm256_srai_epi32::<14>(g_hi);
b_lo = _mm256_srai_epi32::<14>(b_lo);
b_hi = _mm256_srai_epi32::<14>(b_hi);
let r = _mm256_packus_epi32(r_lo, r_hi);
let g = _mm256_packus_epi32(g_lo, g_hi);
let b = _mm256_packus_epi32(b_lo, b_hi);
return (r, g, b);
}
#[inline(always)]
pub fn ycbcr_to_rgba_avx2(
y: &[i16; 16], cb: &[i16; 16], cr: &[i16; 16], out: &mut [u8], offset: &mut usize
) {
unsafe {
ycbcr_to_rgba_unsafe(y, cb, cr, out, offset);
}
}
#[inline]
#[target_feature(enable = "avx2")]
#[rustfmt::skip]
unsafe fn ycbcr_to_rgba_unsafe(
y: &[i16; 16], cb: &[i16; 16], cr: &[i16; 16],
out: &mut [u8],
offset: &mut usize,
)
{
// check if we have enough space to write.
let tmp:& mut [u8; 64] = out.get_mut(*offset..*offset + 64).expect("Slice to small cannot write").try_into().unwrap();
let (r, g, b) = ycbcr_to_rgb_baseline_no_clamp(y, cb, cr);
// set alpha channel to 255 for opaque
// And no these comments were not from me pressing the keyboard
// Pack the integers into u8's using unsigned saturation.
let c = _mm256_packus_epi16(r, g); //aaaaa_bbbbb_aaaaa_bbbbbb
let d = _mm256_packus_epi16(b, _mm256_set1_epi16(255)); // cccccc_dddddd_ccccccc_ddddd
// transpose_u16 and interleave channels
let e = _mm256_unpacklo_epi8(c, d); //ab_ab_ab_ab_ab_ab_ab_ab
let f = _mm256_unpackhi_epi8(c, d); //cd_cd_cd_cd_cd_cd_cd_cd
// final transpose_u16
let g = _mm256_unpacklo_epi8(e, f); //abcd_abcd_abcd_abcd_abcd
let h = _mm256_unpackhi_epi8(e, f);
// undo packus shuffling...
let i = _mm256_permute2x128_si256::<{ shuffle(3, 2, 1, 0) }>(g, h);
let j = _mm256_permute2x128_si256::<{ shuffle(1, 2, 3, 0) }>(g, h);
let k = _mm256_permute2x128_si256::<{ shuffle(3, 2, 0, 1) }>(g, h);
let l = _mm256_permute2x128_si256::<{ shuffle(0, 3, 2, 1) }>(g, h);
let m = _mm256_blend_epi32::<0b1111_0000>(i, j);
let n = _mm256_blend_epi32::<0b1111_0000>(k, l);
// Store
// Use streaming instructions to prevent polluting the cache?
_mm256_storeu_si256(tmp.as_mut_ptr().cast(), m);
_mm256_storeu_si256(tmp[32..].as_mut_ptr().cast(), n);
*offset += 64;
}
#[inline]
const fn shuffle(z: i32, y: i32, x: i32, w: i32) -> i32 {
(z << 6) | (y << 4) | (x << 2) | w
}
+144
View File
@@ -0,0 +1,144 @@
/*
* Copyright (c) 2025.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
//! Aarch64 color conversion routines
//! NEON is mandatory on aarch64.
#![cfg(all(feature = "neon", target_arch = "aarch64"))]
use core::arch::aarch64::*;
use crate::color_convert::scalar::{CB_CF, CR_CF, C_G_CB_COEF_2, C_G_CR_COEF_1, YUV_RND, Y_CF};
const C_1: u64 = u64::from_ne_bytes([
Y_CF.to_ne_bytes()[0],
Y_CF.to_ne_bytes()[1],
CR_CF.to_ne_bytes()[0],
CR_CF.to_ne_bytes()[1],
CB_CF.to_ne_bytes()[0],
CB_CF.to_ne_bytes()[1],
C_G_CR_COEF_1.to_ne_bytes()[0],
C_G_CR_COEF_1.to_ne_bytes()[1]
]);
const C_2: u64 = u64::from_ne_bytes([
C_G_CB_COEF_2.to_ne_bytes()[0],
C_G_CB_COEF_2.to_ne_bytes()[1],
0,
0,
0,
0,
0,
0
]);
#[inline(always)]
unsafe fn ycbcr_to_rgb_baseline_no_clamp(
y: &[i16; 16], cb: &[i16; 16], cr: &[i16; 16]
) -> (uint8x16_t, uint8x16_t, uint8x16_t) {
// NEON has 32 registers, so it is good idea to utilize a lot of variables at once
let cb_cr_bias = vdupq_n_s16(128);
// 0 - Y coeff, 1 - Cr, 2 - Cb, 3 - G1, 4 - G2
let coefficients = vcombine_s16(vcreate_s16(C_1), vcreate_s16(C_2));
let y0 = vld1q_s16(y.as_ptr().cast());
let y1 = vld1q_s16(y[8..].as_ptr().cast());
let mut cb0 = vld1q_s16(cb.as_ptr().cast());
let mut cb1 = vld1q_s16(cb[8..].as_ptr().cast());
let mut cr0 = vld1q_s16(cr.as_ptr().cast());
let mut cr1 = vld1q_s16(cr[8..].as_ptr().cast());
cb0 = vsubq_s16(cb0, cb_cr_bias);
cb1 = vsubq_s16(cb1, cb_cr_bias);
cr0 = vsubq_s16(cr0, cb_cr_bias);
cr1 = vsubq_s16(cr1, cb_cr_bias);
let bias = vdupq_n_s32(i32::from(YUV_RND));
let acc0 = vmlal_laneq_s16::<0>(bias, vget_low_s16(y0), coefficients);
let acc1 = vmlal_high_laneq_s16::<0>(bias, y0, coefficients);
let acc2 = vmlal_laneq_s16::<0>(bias, vget_low_s16(y1), coefficients);
let acc3 = vmlal_high_laneq_s16::<0>(bias, y1, coefficients);
let r0 = vmlal_laneq_s16::<1>(acc0, vget_low_s16(cr0), coefficients);
let r1 = vmlal_high_laneq_s16::<1>(acc1, cr0, coefficients);
let r2 = vmlal_laneq_s16::<1>(acc2, vget_low_s16(cr1), coefficients);
let r3 = vmlal_high_laneq_s16::<1>(acc3, cr1, coefficients);
let b0 = vmlal_laneq_s16::<2>(acc0, vget_low_s16(cb0), coefficients);
let b1 = vmlal_high_laneq_s16::<2>(acc1, cb0, coefficients);
let b2 = vmlal_laneq_s16::<2>(acc2, vget_low_s16(cb1), coefficients);
let b3 = vmlal_high_laneq_s16::<2>(acc3, cb1, coefficients);
// Saturating shift right with signed -> unsigned saturation
let qr0 = vqshrun_n_s32::<14>(r0);
let qr1 = vqshrun_n_s32::<14>(r1);
let qr2 = vqshrun_n_s32::<14>(r2);
let qr3 = vqshrun_n_s32::<14>(r3);
let mut g0 = vmlal_laneq_s16::<4>(acc0, vget_low_s16(cb0), coefficients);
let mut g1 = vmlal_high_laneq_s16::<4>(acc1, cb0, coefficients);
let mut g2 = vmlal_laneq_s16::<4>(acc2, vget_low_s16(cb1), coefficients);
let mut g3 = vmlal_high_laneq_s16::<4>(acc3, cb1, coefficients);
let qb0 = vqshrun_n_s32::<14>(b0);
let qb1 = vqshrun_n_s32::<14>(b1);
let qb2 = vqshrun_n_s32::<14>(b2);
let qb3 = vqshrun_n_s32::<14>(b3);
let r0 = vqmovn_u16(vcombine_u16(qr0, qr1));
let r1 = vqmovn_u16(vcombine_u16(qr2, qr3));
let b0 = vqmovn_u16(vcombine_u16(qb0, qb1));
let b1 = vqmovn_u16(vcombine_u16(qb2, qb3));
g0 = vmlal_laneq_s16::<3>(g0, vget_low_s16(cr0), coefficients);
g1 = vmlal_high_laneq_s16::<3>(g1, cr0, coefficients);
g2 = vmlal_laneq_s16::<3>(g2, vget_low_s16(cr1), coefficients);
g3 = vmlal_high_laneq_s16::<3>(g3, cr1, coefficients);
let qg0 = vqshrun_n_s32::<14>(g0);
let qg1 = vqshrun_n_s32::<14>(g1);
let qg2 = vqshrun_n_s32::<14>(g2);
let qg3 = vqshrun_n_s32::<14>(g3);
let g0 = vqmovn_u16(vcombine_u16(qg0, qg1));
let g1 = vqmovn_u16(vcombine_u16(qg2, qg3));
(
vcombine_u8(r0, r1),
vcombine_u8(g0, g1),
vcombine_u8(b0, b1)
)
}
#[inline(always)]
pub fn ycbcr_to_rgb_neon(
y: &[i16; 16], cb: &[i16; 16], cr: &[i16; 16], out: &mut [u8], offset: &mut usize
) {
// call this in another function to tell RUST to vectorize this
// storing
unsafe {
let (r, g, b) = ycbcr_to_rgb_baseline_no_clamp(y, cb, cr);
vst3q_u8(out.as_mut_ptr(), uint8x16x3_t(r, g, b));
*offset += 48;
}
}
#[inline(always)]
pub fn ycbcr_to_rgba_neon(
y: &[i16; 16], cb: &[i16; 16], cr: &[i16; 16], out: &mut [u8], offset: &mut usize
) {
unsafe {
let (r, g, b) = ycbcr_to_rgb_baseline_no_clamp(y, cb, cr);
vst4q_u8(out.as_mut_ptr(), uint8x16x4_t(r, g, b, vdupq_n_u8(255)));
*offset += 64;
}
}
+139
View File
@@ -0,0 +1,139 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
use core::convert::TryInto;
// Bt.601 Full Range inverse coefficients computed with 14 bits of precision with MPFR.
// This is important to keep them in i16.
// In most cases LLVM will detect what we're doing i16 widening to i32 math and will use
// appropriate optimizations.
pub(crate) const Y_CF: i16 = 16384;
pub(crate) const CR_CF: i16 = 22970;
pub(crate) const CB_CF: i16 = 29032;
pub(crate) const C_G_CR_COEF_1: i16 = -11700;
pub(crate) const C_G_CB_COEF_2: i16 = -5638;
pub(crate) const YUV_PREC: i16 = 14;
// Rounding const for YUV -> RGB conversion: floating equivalent 0.499(9).
pub(crate) const YUV_RND: i16 = (1 << (YUV_PREC - 1)) - 1;
/// Limit values to 0 and 255
#[inline]
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss, dead_code)]
fn clamp(a: i32) -> u8 {
a.clamp(0, 255) as u8
}
/// YCbCr to RGBA color conversion
/// Convert YCbCr to RGB/BGR
///
/// Converts to RGB if const BGRA is false
///
/// Converts to BGR if const BGRA is true
pub fn ycbcr_to_rgba_inner_16_scalar<const BGRA: bool>(
y: &[i16; 16], cb: &[i16; 16], cr: &[i16; 16], output: &mut [u8], pos: &mut usize
) {
let (_, output_position) = output.split_at_mut(*pos);
// Convert into a slice with 64 elements for Rust to see we won't go out of bounds.
let opt: &mut [u8; 64] = output_position
.get_mut(0..64)
.expect("Slice to small cannot write")
.try_into()
.unwrap();
for ((&y, (cb, cr)), out) in y
.iter()
.zip(cb.iter().zip(cr.iter()))
.zip(opt.chunks_exact_mut(4))
{
let cr = cr - 128;
let cb = cb - 128;
let y0 = i32::from(y) * i32::from(Y_CF) + i32::from(YUV_RND);
let r = (y0 + i32::from(cr) * i32::from(CR_CF)) >> YUV_PREC;
let g = (y0
+ i32::from(cr) * i32::from(C_G_CR_COEF_1)
+ i32::from(cb) * i32::from(C_G_CB_COEF_2))
>> YUV_PREC;
let b = (y0 + i32::from(cb) * i32::from(CB_CF)) >> YUV_PREC;
if BGRA {
out[0] = clamp(b);
out[1] = clamp(g);
out[2] = clamp(r);
out[3] = 255;
} else {
out[0] = clamp(r);
out[1] = clamp(g);
out[2] = clamp(b);
out[3] = 255;
}
}
*pos += 64;
}
/// Convert YCbCr to RGB/BGR
///
/// Converts to RGB if const BGRA is false
///
/// Converts to BGR if const BGRA is true
pub fn ycbcr_to_rgb_inner_16_scalar<const BGRA: bool>(
y: &[i16; 16], cb: &[i16; 16], cr: &[i16; 16], output: &mut [u8], pos: &mut usize
) {
let (_, output_position) = output.split_at_mut(*pos);
// Convert into a slice with 48 elements
let opt: &mut [u8; 48] = output_position
.get_mut(0..48)
.expect("Slice to small cannot write")
.try_into()
.unwrap();
for ((&y, (cb, cr)), out) in y
.iter()
.zip(cb.iter().zip(cr.iter()))
.zip(opt.chunks_exact_mut(3))
{
let cr = cr - 128;
let cb = cb - 128;
let y0 = i32::from(y) * i32::from(Y_CF) + i32::from(YUV_RND);
let r = (y0 + i32::from(cr) * i32::from(CR_CF)) >> YUV_PREC;
let g = (y0
+ i32::from(cr) * i32::from(C_G_CR_COEF_1)
+ i32::from(cb) * i32::from(C_G_CB_COEF_2))
>> YUV_PREC;
let b = (y0 + i32::from(cb) * i32::from(CB_CF)) >> YUV_PREC;
if BGRA {
out[0] = clamp(b);
out[1] = clamp(g);
out[2] = clamp(r);
} else {
out[0] = clamp(r);
out[1] = clamp(g);
out[2] = clamp(b);
}
}
// Increment pos
*pos += 48;
}
pub fn ycbcr_to_grayscale(y: &[i16], width: usize, padded_width: usize, output: &mut [u8]) {
for (y_in, out) in y
.chunks_exact(padded_width)
.zip(output.chunks_exact_mut(width))
{
for (y, out) in y_in.iter().zip(out.iter_mut()) {
*out = *y as u8;
}
}
}
+232
View File
@@ -0,0 +1,232 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
//! This module exports a single struct to store information about
//! JPEG image components
//!
//! The data is extracted from a SOF header.
use alloc::vec::Vec;
use alloc::{format, vec};
use zune_core::log::trace;
use crate::alloc::string::ToString;
use crate::decoder::MAX_COMPONENTS;
use crate::errors::DecodeErrors;
use crate::upsampler::upsample_no_op;
const MAX_SAMP_FACTOR: usize = 4;
/// Represents an up-sampler function, this function will be called to upsample
/// a down-sampled image
pub type UpSampler = fn(
input: &[i16],
in_near: &[i16],
in_far: &[i16],
scratch_space: &mut [i16],
output: &mut [i16]
);
/// Component Data from start of frame
#[derive(Clone)]
pub(crate) struct Components {
/// The type of component that has the metadata below, can be Y,Cb or Cr
pub component_id: ComponentID,
/// Sub-sampling ratio of this component in the x-plane
pub vertical_sample: usize,
/// Sub-sampling ratio of this component in the y-plane
pub horizontal_sample: usize,
/// DC huffman table position
pub dc_huff_table: usize,
/// AC huffman table position for this element.
pub ac_huff_table: usize,
/// Quantization table number
pub quantization_table_number: u8,
/// Specifies quantization table to use with this component
pub quantization_table: [i32; 64],
/// dc prediction for the component
pub dc_pred: i32,
/// An up-sampling function, can be basic or SSE, depending
/// on the platform
pub up_sampler: UpSampler,
/// How pixels do we need to go to get to the next line?
pub width_stride: usize,
/// Component ID for progressive
pub id: u8,
/// Whether we need to decode this image component.
pub needed: bool,
/// Upsample scanline
pub raw_coeff: Vec<i16>,
/// Upsample destination, stores a scanline worth of sub sampled data
pub upsample_dest: Vec<i16>,
/// previous row, used to handle MCU boundaries
pub row_up: Vec<i16>,
/// current row, used to handle MCU boundaries again
pub row: Vec<i16>,
pub first_row_upsample_dest: Vec<i16>,
pub idct_pos: usize,
pub x: usize,
pub w2: usize,
pub y: usize,
pub sample_ratio: SampleRatios,
// a very annoying bug
pub fix_an_annoying_bug: usize
}
impl Components {
/// Create a new instance from three bytes from the start of frame
#[inline]
pub fn from(a: [u8; 3], pos: u8) -> Result<Components, DecodeErrors> {
// it's a unique identifier.
// doesn't have to be ascending
// see tests/inputs/huge_sof_number
//
// For such cases, use the position of the component
// to determine width
let id = match pos {
0 => ComponentID::Y,
1 => ComponentID::Cb,
2 => ComponentID::Cr,
3 => ComponentID::Q,
_ => {
return Err(DecodeErrors::Format(format!(
"Unknown component id found,{pos}, expected value between 1 and 4"
)))
}
};
let horizontal_sample = (a[1] >> 4) as usize;
let vertical_sample = (a[1] & 0x0f) as usize;
// Match libjpeg turbo on checking for sampling factors
// Reject anything above 4
if horizontal_sample > MAX_SAMP_FACTOR {
return Err(DecodeErrors::Format(format!(
"Bogus Horizontal Sampling Factor {horizontal_sample}"
)));
}
if vertical_sample > MAX_SAMP_FACTOR {
return Err(DecodeErrors::Format(format!(
"Bogus Vertical Sampling Factor {vertical_sample}"
)));
}
let quantization_table_number = a[2];
// confirm quantization number is between 0 and MAX_COMPONENTS
if usize::from(quantization_table_number) >= MAX_COMPONENTS {
return Err(DecodeErrors::Format(format!(
"Too large quantization number :{quantization_table_number}, expected value between 0 and {MAX_COMPONENTS}"
)));
}
// check that upsampling ratios are powers of two
// if these fail, it's probably a corrupt image.
if !horizontal_sample.is_power_of_two() {
return Err(DecodeErrors::Format(format!(
"Horizontal sample is not a power of two({horizontal_sample}) cannot decode"
)));
}
// if !vertical_sample.is_power_of_two() {
// return Err(DecodeErrors::Format(format!(
// "Vertical sub-sample is not power of two({vertical_sample}) cannot decode"
// )));
// }
if vertical_sample == 0 {
// Check for invalid vertical sample
return Err(DecodeErrors::Format("Vertical sample is zero".to_string()));
}
trace!(
"Component ID:{:?} \tHS:{} VS:{} QT:{}",
id,
horizontal_sample,
vertical_sample,
quantization_table_number
);
Ok(Components {
component_id: id,
vertical_sample,
horizontal_sample,
quantization_table_number,
first_row_upsample_dest: vec![],
// These two will be set with sof marker
dc_huff_table: 0,
ac_huff_table: 0,
quantization_table: [0; 64],
dc_pred: 0,
up_sampler: upsample_no_op,
// set later
width_stride: horizontal_sample,
id: a[0],
needed: true,
raw_coeff: vec![],
upsample_dest: vec![],
row_up: vec![],
row: vec![],
idct_pos: 0,
x: 0,
y: 0,
w2: 0,
sample_ratio: SampleRatios::None,
fix_an_annoying_bug: 1
})
}
/// Setup space for upsampling
///
/// During upsample, we need a reference of the last row so that upsampling can
/// proceed correctly,
/// so we store the last line of every scanline and use it for the next upsampling procedure
/// to store this, but since we don't need it for 1v1 upsampling,
/// we only call this for routines that need upsampling
///
/// # Requirements
/// - width stride of this element is set for the component.
pub fn setup_upsample_scanline(&mut self) {
self.row = vec![0; self.width_stride * self.vertical_sample];
self.row_up = vec![0; self.width_stride * self.vertical_sample];
self.first_row_upsample_dest =
vec![128; self.vertical_sample * self.width_stride * self.sample_ratio.sample()];
self.upsample_dest =
vec![0; self.width_stride * self.sample_ratio.sample() * self.fix_an_annoying_bug * 8];
}
}
/// Component ID's
#[derive(Copy, Debug, Clone, PartialEq, Eq)]
pub enum ComponentID {
/// Luminance channel
Y,
/// Blue chrominance
Cb,
/// Red chrominance
Cr,
/// Q or fourth component
Q
}
#[derive(Copy, Debug, Clone, PartialEq, Eq, Default)]
pub enum SampleRatios {
HV,
V,
H,
Generic(usize, usize),
#[default]
None
}
impl SampleRatios {
pub fn sample(self) -> usize {
match self {
SampleRatios::HV => 4,
SampleRatios::V | SampleRatios::H => 2,
SampleRatios::Generic(a, b) => a * b,
SampleRatios::None => 1
}
}
}
+987
View File
@@ -0,0 +1,987 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
//! Main image logic.
#![allow(clippy::doc_markdown)]
use alloc::string::ToString;
use alloc::vec::Vec;
use alloc::{format, vec};
use zune_core::bytestream::{ZByteReaderTrait, ZReader};
use zune_core::colorspace::ColorSpace;
use zune_core::log::{error, trace, warn};
use zune_core::options::DecoderOptions;
use crate::color_convert::choose_ycbcr_to_rgb_convert_func;
use crate::components::{Components, SampleRatios};
use crate::errors::{DecodeErrors, UnsupportedSchemes};
use crate::headers::{
parse_app1, parse_app13, parse_app14, parse_app2, parse_dqt, parse_huffman, parse_sos,
parse_start_of_frame
};
use crate::huffman::HuffmanTable;
use crate::idct::{choose_idct_func, choose_idct_1x1_func, choose_idct_4x4_func};
use crate::marker::Marker;
use crate::misc::SOFMarkers;
use crate::upsampler::{
choose_horizontal_samp_function, choose_hv_samp_function, choose_v_samp_function,
generic_sampler, upsample_no_op
};
/// Maximum components
pub(crate) const MAX_COMPONENTS: usize = 4;
/// Maximum image dimensions supported.
pub(crate) const MAX_DIMENSIONS: usize = 1 << 27;
/// Color conversion function that can convert YCbCr colorspace to RGB(A/X) for
/// 16 values
///
/// The following are guarantees to the following functions
///
/// 1. The `&[i16]` slices passed contain 16 items
///
/// 2. The slices passed are in the following order
/// `y,cb,cr`
///
/// 3. `&mut [u8]` is zero initialized
///
/// 4. `&mut usize` points to the position in the array where new values should
/// be used
///
/// The pointer should
/// 1. Carry out color conversion
/// 2. Update `&mut usize` with the new position
pub type ColorConvert16Ptr = fn(&[i16; 16], &[i16; 16], &[i16; 16], &mut [u8], &mut usize);
/// IDCT function prototype
///
/// This encapsulates a dequantize and IDCT function which will carry out the
/// following functions
///
/// Multiply each 64 element block of `&mut [i16]` with `&Aligned32<[i32;64]>`
/// Carry out IDCT (type 3 dct) on ach block of 64 i16's
pub type IDCTPtr = fn(&mut [i32; 64], &mut [i16], usize);
/// An encapsulation of an ICC chunk
pub(crate) struct ICCChunk {
pub(crate) seq_no: u8,
pub(crate) num_markers: u8,
pub(crate) data: Vec<u8>
}
/// A JPEG Decoder Instance.
#[allow(clippy::upper_case_acronyms, clippy::struct_excessive_bools)]
pub struct JpegDecoder<T> {
/// Struct to hold image information from SOI
pub(crate) info: ImageInfo,
/// Quantization tables, will be set to none and the tables will
/// be moved to `components` field
pub(crate) qt_tables: [Option<[i32; 64]>; MAX_COMPONENTS],
/// DC Huffman Tables with a maximum of 4 tables for each component
pub(crate) dc_huffman_tables: [Option<HuffmanTable>; MAX_COMPONENTS],
/// AC Huffman Tables with a maximum of 4 tables for each component
pub(crate) ac_huffman_tables: [Option<HuffmanTable>; MAX_COMPONENTS],
/// Image components, holds information like DC prediction and quantization
/// tables of a component
pub(crate) components: Vec<Components>,
/// maximum horizontal component of all channels in the image
pub(crate) h_max: usize,
// maximum vertical component of all channels in the image
pub(crate) v_max: usize,
/// mcu's width (interleaved scans)
pub(crate) mcu_width: usize,
/// MCU height(interleaved scans
pub(crate) mcu_height: usize,
/// Number of MCU's in the x plane
pub(crate) mcu_x: usize,
/// Number of MCU's in the y plane
pub(crate) mcu_y: usize,
/// Is the image interleaved?
pub(crate) is_interleaved: bool,
/// Image input colorspace, should be YCbCr for a sane image, might be
/// grayscale too
pub(crate) input_colorspace: ColorSpace,
// Progressive image details
/// Is the image progressive?
pub(crate) is_progressive: bool,
/// Start of spectral scan
pub(crate) spec_start: u8,
/// End of spectral scan
pub(crate) spec_end: u8,
/// Successive approximation bit position high
pub(crate) succ_high: u8,
/// Successive approximation bit position low
pub(crate) succ_low: u8,
/// Number of components.
pub(crate) num_scans: u8,
/// For a scan, check if any component has vertical/horizontal sampling.
pub(crate) scan_subsampled: bool,
// Function pointers, for pointy stuff.
/// Dequantize and idct function
// This is determined at runtime which function to run, statically it's
// initialized to a platform independent one and during initialization
// of this struct, we check if we can switch to a faster one which
// depend on certain CPU extensions.
pub(crate) idct_func: IDCTPtr,
/// Specialized IDCT when we can guarantee only few coefficients are non-zero.
///
/// **The callee must uphold a contract**. See [`choose_idct_4x4_func`].
pub(crate) idct_4x4_func: IDCTPtr,
pub(crate) idct_1x1_func: IDCTPtr,
// Color convert function which acts on 16 YCbCr values
pub(crate) color_convert_16: ColorConvert16Ptr,
pub(crate) z_order: [usize; MAX_COMPONENTS],
/// restart markers
pub(crate) restart_interval: usize,
pub(crate) todo: usize,
// decoder options
pub(crate) options: DecoderOptions,
// byte-stream
pub(crate) stream: ZReader<T>,
// Indicate whether headers have been decoded
pub(crate) headers_decoded: bool,
pub(crate) seen_sof: bool,
// exif data, lifted from app2
pub(crate) icc_data: Vec<ICCChunk>,
pub(crate) is_mjpeg: bool,
pub(crate) coeff: usize // Solves some weird bug :)
}
impl<T> JpegDecoder<T>
where
T: ZByteReaderTrait
{
#[allow(clippy::redundant_field_names)]
fn default(options: DecoderOptions, buffer: T) -> Self {
let color_convert = choose_ycbcr_to_rgb_convert_func(ColorSpace::RGB, &options).unwrap();
JpegDecoder {
info: ImageInfo::default(),
qt_tables: [None, None, None, None],
dc_huffman_tables: [None, None, None, None],
ac_huffman_tables: [None, None, None, None],
components: vec![],
// Interleaved information
h_max: 1,
v_max: 1,
mcu_height: 0,
mcu_width: 0,
mcu_x: 0,
mcu_y: 0,
is_interleaved: false,
is_progressive: false,
spec_start: 0,
spec_end: 0,
succ_high: 0,
succ_low: 0,
num_scans: 0,
scan_subsampled: false,
idct_func: choose_idct_func(&options),
idct_4x4_func: choose_idct_4x4_func(&options),
idct_1x1_func: choose_idct_1x1_func(&options),
color_convert_16: color_convert,
input_colorspace: ColorSpace::YCbCr,
z_order: [0; MAX_COMPONENTS],
restart_interval: 0,
todo: 0x7fff_ffff,
options: options,
stream: ZReader::new(buffer),
headers_decoded: false,
seen_sof: false,
icc_data: vec![],
is_mjpeg: false,
coeff: 1
}
}
/// Decode a buffer already in memory
///
/// The buffer should be a valid jpeg file, perhaps created by the command
/// `std:::fs::read()` or a JPEG file downloaded from the internet.
///
/// # Errors
/// See DecodeErrors for an explanation
pub fn decode(&mut self) -> Result<Vec<u8>, DecodeErrors> {
self.decode_headers()?;
let size = self.output_buffer_size().unwrap();
let mut out = vec![0; size];
self.decode_into(&mut out)?;
Ok(out)
}
/// Create a new Decoder instance
///
/// # Arguments
/// - `stream`: The raw bytes of a jpeg file.
#[must_use]
#[allow(clippy::new_without_default)]
pub fn new(stream: T) -> JpegDecoder<T> {
JpegDecoder::default(DecoderOptions::default(), stream)
}
/// Returns the image information
///
/// This **must** be called after a subsequent call to [`decode`] or [`decode_headers`]
/// it will return `None`
///
/// # Returns
/// - `Some(info)`: Image information,width, height, number of components
/// - None: Indicates image headers haven't been decoded
///
/// [`decode`]: JpegDecoder::decode
/// [`decode_headers`]: JpegDecoder::decode_headers
#[must_use]
pub fn info(&self) -> Option<ImageInfo> {
// we check for fails to that call by comparing what we have to the default, if
// it's default we assume that the caller failed to uphold the
// guarantees. We can be sure that an image cannot be the default since
// its a hard panic in-case width or height are set to zero.
if !self.headers_decoded {
return None;
}
return Some(self.info.clone());
}
/// Return the number of bytes required to hold a decoded image frame
/// decoded using the given input transformations
///
/// # Returns
/// - `Some(usize)`: Minimum size for a buffer needed to decode the image
/// - `None`: Indicates the image was not decoded, or image dimensions would overflow a usize
///
#[must_use]
pub fn output_buffer_size(&self) -> Option<usize> {
return if self.headers_decoded {
Some(
usize::from(self.width())
.checked_mul(usize::from(self.height()))?
.checked_mul(self.options.jpeg_get_out_colorspace().num_components())?
)
} else {
None
};
}
/// Get an immutable reference to the decoder options
/// for the decoder instance
///
/// This can be used to modify options before actual decoding
/// but after initial creation
///
/// # Example
/// ```no_run
/// use zune_core::bytestream::ZCursor;
/// use zune_jpeg::JpegDecoder;
///
/// let mut decoder = JpegDecoder::new(ZCursor::new(&[]));
/// // get current options
/// let mut options = decoder.options();
/// // modify it
/// let new_options = options.set_max_width(10);
/// // set it back
/// decoder.set_options(new_options);
///
/// ```
#[must_use]
pub const fn options(&self) -> &DecoderOptions {
&self.options
}
/// Return the input colorspace of the image
///
/// This indicates the colorspace that is present in
/// the image, but this may be different to the colorspace that
/// the output will be transformed to
///
/// # Returns
/// -`Some(Colorspace)`: Input colorspace
/// - None : Indicates the headers weren't decoded
#[must_use]
pub fn input_colorspace(&self) -> Option<ColorSpace> {
return if self.headers_decoded { Some(self.input_colorspace) } else { None };
}
/// Set decoder options
///
/// This can be used to set new options even after initialization
/// but before decoding.
///
/// This does not bear any significance after decoding an image
///
/// # Arguments
/// - `options`: New decoder options
///
/// # Example
/// Set maximum jpeg progressive passes to be 4
///
/// ```no_run
/// use zune_core::bytestream::ZCursor;
/// use zune_jpeg::JpegDecoder;
/// let mut decoder =JpegDecoder::new(ZCursor::new(&[]));
/// // this works also because DecoderOptions implements `Copy`
/// let options = decoder.options().jpeg_set_max_scans(4);
/// // set the new options
/// decoder.set_options(options);
/// // now decode
/// decoder.decode().unwrap();
/// ```
pub fn set_options(&mut self, options: DecoderOptions) {
self.options = options;
}
/// Decode Decoder headers
///
/// This routine takes care of parsing supported headers from a Decoder
/// image
///
/// # Supported Headers
/// - APP(0)
/// - SOF(O)
/// - DQT -> Quantization tables
/// - DHT -> Huffman tables
/// - SOS -> Start of Scan
/// # Unsupported Headers
/// - SOF(n) -> Decoder images which are not baseline/progressive
/// - DAC -> Images using Arithmetic tables
/// - JPG(n)
fn decode_headers_internal(&mut self) -> Result<(), DecodeErrors> {
if self.headers_decoded {
trace!("Headers decoded!");
return Ok(());
}
// match output colorspace here
// we know this will only be called once per image
// so makes sense
// We only care for ycbcr to rgb/rgba here
// in case one is using another colorspace.
// May god help you
let out_colorspace = self.options.jpeg_get_out_colorspace();
if matches!(
out_colorspace,
ColorSpace::BGR | ColorSpace::BGRA | ColorSpace::RGB | ColorSpace::RGBA
) {
self.color_convert_16 = choose_ycbcr_to_rgb_convert_func(
self.options.jpeg_get_out_colorspace(),
&self.options
)
.unwrap();
}
// First two bytes should be jpeg soi marker
let magic_bytes = self.stream.get_u16_be_err()?;
let mut last_byte = 0;
let mut bytes_before_marker = 0;
if magic_bytes != 0xffd8 {
return Err(DecodeErrors::IllegalMagicBytes(magic_bytes));
}
loop {
// read a byte
let mut m = self.stream.read_u8_err()?;
// AND OF COURSE some images will have fill bytes in their marker
// bitstreams because why not.
//
// I am disappointed as a man.
if (m == 0xFF || m == 0) && last_byte == 0xFF {
// This handles the edge case where
// images have markers with fill bytes(0xFF)
// or byte stuffing (0)
// I.e 0xFF 0xFF 0xDA
// and
// 0xFF 0 0xDA
// It should ignore those fill bytes and take 0xDA
// I don't know why such images exist
// but they do.
// so this is for you (with love)
while m == 0xFF || m == 0x0 {
last_byte = m;
m = self.stream.read_u8_err()?;
}
}
// Last byte should be 0xFF to confirm existence of a marker since markers look
// like OxFF(some marker data)
if last_byte == 0xFF {
let marker = Marker::from_u8(m);
if let Some(n) = marker {
if bytes_before_marker > 3 {
if self.options.strict_mode()
/*No reason to use this*/
{
return Err(DecodeErrors::FormatStatic(
"[strict-mode]: Extra bytes between headers"
));
}
error!(
"Extra bytes {} before marker 0xFF{:X}",
bytes_before_marker - 3,
m
);
}
bytes_before_marker = 0;
self.parse_marker_inner(n)?;
// break after reading the start of scan.
// what follows is the image data
if n == Marker::SOS {
self.headers_decoded = true;
trace!("Input colorspace {:?}", self.input_colorspace);
// Check if image is RGB
// The check is weird, we need to check if ID
// represents R, G and B in ascii,
//
// I am not sure if this is even specified in any standard,
// but jpegli https://github.com/google/jpegli does encode
// its images that way, so this will check for that. and handle it appropriately
// It is spefified here so that on a successful header decode,we can at least
// try to attribute image colorspace correctly.
//
// It was first the issue in https://github.com/etemesi254/zune-image/issues/291
// that brought it to light
//
let mut is_rgb = self.components.len() == 3;
let chars = ['R', 'G', 'B'];
for (comp, single_char) in self.components.iter().zip(chars.iter()) {
is_rgb &= comp.id == (*single_char) as u8
}
// Image is RGB, change colorspace
if is_rgb {
self.input_colorspace = ColorSpace::RGB;
}
return Ok(());
}
} else {
bytes_before_marker = 0;
warn!("Marker 0xFF{:X} not known", m);
let length = self.stream.get_u16_be_err()?;
if length < 2 {
return Err(DecodeErrors::Format(format!(
"Found a marker with invalid length : {length}"
)));
}
warn!("Skipping {} bytes", length - 2);
self.stream.skip((length - 2) as usize)?;
}
}
last_byte = m;
bytes_before_marker += 1;
}
// Check if image is RGB
}
#[allow(clippy::too_many_lines)]
pub(crate) fn parse_marker_inner(&mut self, m: Marker) -> Result<(), DecodeErrors> {
match m {
Marker::SOF(0..=2) => {
let marker = {
// choose marker
if m == Marker::SOF(0) || m == Marker::SOF(1) {
SOFMarkers::BaselineDct
} else {
self.is_progressive = true;
SOFMarkers::ProgressiveDctHuffman
}
};
trace!("Image encoding scheme =`{:?}`", marker);
// get components
parse_start_of_frame(marker, self)?;
}
// Start of Frame Segments not supported
Marker::SOF(v) => {
let feature = UnsupportedSchemes::from_int(v);
if let Some(feature) = feature {
return Err(DecodeErrors::Unsupported(feature));
}
return Err(DecodeErrors::Format("Unsupported image format".to_string()));
}
//APP(0) segment
Marker::APP(0) => {
let mut length = self.stream.get_u16_be_err()?;
if length < 2 {
return Err(DecodeErrors::Format(format!(
"Found a marker with invalid length:{length}\n"
)));
}
// skip for now
if length > 5 {
let mut buffer = [0u8; 5];
self.stream.read_exact_bytes(&mut buffer)?;
if &buffer == b"AVI1\0" {
self.is_mjpeg = true;
}
length -= 5;
}
self.stream.skip(length.saturating_sub(2) as usize)?;
//parse_app(buf, m, &mut self.info)?;
}
Marker::APP(1) => {
parse_app1(self)?;
}
Marker::APP(2) => {
parse_app2(self)?;
}
// Quantization tables
Marker::DQT => {
parse_dqt(self)?;
}
// Huffman tables
Marker::DHT => {
parse_huffman(self)?;
}
// Start of Scan Data
Marker::SOS => {
parse_sos(self)?;
}
Marker::EOI => return Err(DecodeErrors::FormatStatic("Premature End of image")),
Marker::DAC | Marker::DNL => {
return Err(DecodeErrors::Format(format!(
"Parsing of the following header `{m:?}` is not supported,\
cannot continue"
)));
}
Marker::DRI => {
if self.stream.get_u16_be_err()? != 4 {
return Err(DecodeErrors::Format(
"Bad DRI length, Corrupt JPEG".to_string()
));
}
self.restart_interval = usize::from(self.stream.get_u16_be_err()?);
trace!("DRI marker present ({})", self.restart_interval);
self.todo = self.restart_interval;
}
Marker::APP(14) => {
parse_app14(self)?;
}
Marker::APP(13) => {
parse_app13(self)?;
}
_ => {
warn!(
"Capabilities for processing marker \"{:?}\" not implemented",
m
);
let length = self.stream.get_u16_be_err()?;
if length < 2 {
return Err(DecodeErrors::Format(format!(
"Found a marker with invalid length:{length}\n"
)));
}
warn!("Skipping {} bytes", length - 2);
self.stream.skip((length - 2) as usize)?;
}
}
Ok(())
}
/// Get the embedded ICC profile if it exists
/// and is correct
///
/// One needs not to decode the whole image to extract this,
/// calling [`decode_headers`] for an image with an ICC profile
/// allows you to decode this
///
/// # Returns
/// - `Some(Vec<u8>)`: The raw ICC profile of the image
/// - `None`: May indicate an error in the ICC profile , non-existence of
/// an ICC profile, or that the headers weren't decoded.
///
/// [`decode_headers`]:Self::decode_headers
#[must_use]
pub fn icc_profile(&self) -> Option<Vec<u8>> {
let mut marker_present: [Option<&ICCChunk>; 256] = [None; 256];
if !self.headers_decoded {
return None;
}
let num_markers = self.icc_data.len();
if num_markers == 0 || num_markers >= 255 {
return None;
}
// check validity
for chunk in &self.icc_data {
if usize::from(chunk.num_markers) != num_markers {
// all the lengths must match
return None;
}
if chunk.seq_no == 0 {
warn!("Zero sequence number in ICC, corrupt ICC chunk");
return None;
}
if marker_present[usize::from(chunk.seq_no)].is_some() {
// duplicate seq_no
warn!("Duplicate sequence number in ICC, corrupt chunk");
return None;
}
marker_present[usize::from(chunk.seq_no)] = Some(chunk);
}
let mut data = Vec::with_capacity(1000);
// assemble the data now
for chunk in marker_present.get(1..=num_markers).unwrap() {
if let Some(ch) = chunk {
data.extend_from_slice(&ch.data);
} else {
warn!("Missing icc sequence number, corrupt ICC chunk ");
return None;
}
}
Some(data)
}
/// Return the exif data for the file
///
/// This returns the raw exif data starting at the
/// TIFF header
///
/// # Returns
/// -`Some(data)`: The raw exif data, if present in the image
/// - None: May indicate the following
///
/// 1. The image doesn't have exif data
/// 2. The image headers haven't been decoded
#[must_use]
pub fn exif(&self) -> Option<&Vec<u8>> {
return self.info.exif_data.as_ref();
}
/// Return the XMP data for the file
///
/// This returns raw XMP data starting at the XML header
/// One needs an XML/XMP decoder to extract valuable metadata
///
///
/// # Returns
/// - `Some(data)`: Raw xmp data
/// - `None`: May indicate the following
/// 1. The image does not have xmp data
/// 2. The image headers have not been decoded
///
/// # Example
///
/// ```no_run
/// use zune_core::bytestream::ZCursor;
/// use zune_jpeg::JpegDecoder;
/// let mut decoder = JpegDecoder::new(ZCursor::new(&[]));
/// // decode headers to extract xmp metadata if present
/// decoder.decode_headers().unwrap();
/// if let Some(data) = decoder.xmp(){
/// let stringified = String::from_utf8_lossy(data);
/// println!("XMP")
/// } else{
/// println!("No XMP Found")
/// }
///
/// ```
pub fn xmp(&self) -> Option<&Vec<u8>> {
return self.info.xmp_data.as_ref();
}
/// Return the IPTC data for the file
///
/// This returns the raw IPTC data.
///
/// # Returns
/// -`Some(data)`: The raw IPTC data, if present in the image
/// - None: May indicate the following
///
/// 1. The image doesn't have IPTC data
/// 2. The image headers haven't been decoded
#[must_use]
pub fn iptc(&self) -> Option<&Vec<u8>> {
return self.info.iptc_data.as_ref();
}
/// Get the output colorspace the image pixels will be decoded into
///
///
/// # Note.
/// This field can only be regarded after decoding headers,
/// as markers such as Adobe APP14 may dictate different colorspaces
/// than requested.
///
/// Calling `decode_headers` is sufficient to know what colorspace the
/// output is, if this is called after `decode` it indicates the colorspace
/// the output is currently in
///
/// Additionally not all input->output colorspace mappings are supported
/// but all input colorspaces can map to RGB colorspace, so that's a safe bet
/// if one is handling image formats
///
///# Returns
/// - `Some(Colorspace)`: If headers have been decoded, the colorspace the
///output array will be in
///- `None
#[must_use]
pub fn output_colorspace(&self) -> Option<ColorSpace> {
return if self.headers_decoded {
Some(self.options.jpeg_get_out_colorspace())
} else {
None
};
}
/// Decode into a pre-allocated buffer
///
/// It is an error if the buffer size is smaller than
/// [`output_buffer_size()`](Self::output_buffer_size)
///
/// If the buffer is bigger than expected, we ignore the end padding bytes
///
/// # Example
///
/// - Read headers and then alloc a buffer big enough to hold the image
///
/// ```no_run
/// use zune_core::bytestream::ZCursor;
/// use zune_jpeg::JpegDecoder;
/// let mut decoder = JpegDecoder::new(ZCursor::new(&[]));
/// // before we get output, we must decode the headers to get width
/// // height, and input colorspace
/// decoder.decode_headers().unwrap();
///
/// let mut out = vec![0;decoder.output_buffer_size().unwrap()];
/// // write into out
/// decoder.decode_into(&mut out).unwrap();
/// ```
///
///
pub fn decode_into(&mut self, out: &mut [u8]) -> Result<(), DecodeErrors> {
self.decode_headers_internal()?;
let expected_size = self.output_buffer_size().unwrap();
if out.len() < expected_size {
// too small of a size
return Err(DecodeErrors::TooSmallOutput(expected_size, out.len()));
}
// ensure we don't touch anyone else's scratch space
let out_len = core::cmp::min(out.len(), expected_size);
let out = &mut out[0..out_len];
if self.is_progressive {
self.decode_mcu_ycbcr_progressive(out)
} else {
self.decode_mcu_ycbcr_baseline(out)
}
}
/// Read only headers from a jpeg image buffer
///
/// This allows you to extract important information like
/// image width and height without decoding the full image
///
/// # Examples
/// ```no_run
/// use zune_core::bytestream::ZCursor;
/// use zune_jpeg::{JpegDecoder};
///
/// let img_data = std::fs::read("a_valid.jpeg").unwrap();
/// let mut decoder = JpegDecoder::new(ZCursor::new(&img_data));
/// decoder.decode_headers().unwrap();
///
/// println!("Total decoder dimensions are : {:?} pixels",decoder.dimensions());
/// println!("Number of components in the image are {}", decoder.info().unwrap().components);
/// ```
/// # Errors
/// See DecodeErrors enum for list of possible errors during decoding
pub fn decode_headers(&mut self) -> Result<(), DecodeErrors> {
self.decode_headers_internal()?;
Ok(())
}
/// Create a new decoder with the specified options to be used for decoding
/// an image
///
/// # Arguments
/// - `buf`: The input buffer from where we will pull in compressed jpeg bytes from
/// - `options`: Options specific to this decoder instance
#[must_use]
pub fn new_with_options(buf: T, options: DecoderOptions) -> JpegDecoder<T> {
JpegDecoder::default(options, buf)
}
/// Set up-sampling routines in case an image is down sampled
pub(crate) fn set_upsampling(&mut self) -> Result<(), DecodeErrors> {
// no sampling, return early
// check if horizontal max ==1
if self.h_max == self.v_max && self.h_max == 1 {
return Ok(());
}
for comp in &mut self.components {
let hs = self.h_max / comp.horizontal_sample;
let vs = self.v_max / comp.vertical_sample;
let samp_factor = match (hs, vs) {
(1, 1) => {
comp.sample_ratio = SampleRatios::None;
upsample_no_op
}
(2, 1) => {
comp.sample_ratio = SampleRatios::H;
choose_horizontal_samp_function(&self.options)
}
(1, 2) => {
comp.sample_ratio = SampleRatios::V;
choose_v_samp_function(&self.options)
}
(2, 2) => {
comp.sample_ratio = SampleRatios::HV;
choose_hv_samp_function(&self.options)
}
(hs, vs) => {
comp.sample_ratio = SampleRatios::Generic(hs, vs);
generic_sampler()
}
};
comp.setup_upsample_scanline();
comp.up_sampler = samp_factor;
}
return Ok(());
}
#[must_use]
/// Get the width of the image as a u16
///
/// The width lies between 1 and 65535
pub(crate) fn width(&self) -> u16 {
self.info.width
}
/// Get the height of the image as a u16
///
/// The height lies between 1 and 65535
#[must_use]
pub(crate) fn height(&self) -> u16 {
self.info.height
}
/// Get image dimensions as a tuple of width and height
/// or `None` if the image hasn't been decoded.
///
/// # Returns
/// - `Some(width,height)`: Image dimensions
/// - None : The image headers haven't been decoded
#[must_use]
pub const fn dimensions(&self) -> Option<(usize, usize)> {
return if self.headers_decoded {
Some((self.info.width as usize, self.info.height as usize))
} else {
None
};
}
}
#[derive(Default, Clone, Eq, PartialEq, Debug)]
pub struct GainMapInfo {
pub data: Vec<u8>
}
/// A struct representing Image Information
#[derive(Default, Clone, Eq, PartialEq)]
#[allow(clippy::module_name_repetitions)]
pub struct ImageInfo {
/// Width of the image
pub width: u16,
/// Height of image
pub height: u16,
/// PixelDensity
pub pixel_density: u8,
/// Start of frame markers
pub sof: SOFMarkers,
/// Horizontal sample
pub x_density: u16,
/// Vertical sample
pub y_density: u16,
/// Number of components
pub components: u8,
/// Gain Map information, useful for
/// UHDR images
pub gain_map_info: Vec<GainMapInfo>,
/// Multi picture information, useful for
/// UHDR images
pub multi_picture_information: Option<Vec<u8>>,
/// Exif Data
pub exif_data: Option<Vec<u8>>,
/// XMP Data
pub xmp_data: Option<Vec<u8>>,
/// IPTC Data
pub iptc_data: Option<Vec<u8>>,
/// Image sub-sampling ratio
pub sample_ratio: SampleRatios
}
impl ImageInfo {
/// Set width of the image
///
/// Found in the start of frame
pub(crate) fn set_width(&mut self, width: u16) {
self.width = width;
}
/// Set height of the image
///
/// Found in the start of frame
pub(crate) fn set_height(&mut self, height: u16) {
self.height = height;
}
/// Set the image density
///
/// Found in the start of frame
pub(crate) fn set_density(&mut self, density: u8) {
self.pixel_density = density;
}
/// Set image Start of frame marker
///
/// found in the Start of frame header
pub(crate) fn set_sof_marker(&mut self, marker: SOFMarkers) {
self.sof = marker;
}
/// Set image x-density(dots per pixel)
///
/// Found in the APP(0) marker
#[allow(dead_code)]
pub(crate) fn set_x(&mut self, sample: u16) {
self.x_density = sample;
}
/// Set image y-density
///
/// Found in the APP(0) marker
#[allow(dead_code)]
pub(crate) fn set_y(&mut self, sample: u16) {
self.y_density = sample;
}
}
+167
View File
@@ -0,0 +1,167 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
//! Contains most common errors that may be encountered in decoding a Decoder
//! image
use alloc::string::String;
use core::fmt::{Debug, Display, Formatter};
use zune_core::bytestream::ZByteIoError;
use crate::misc::{
START_OF_FRAME_EXT_AR, START_OF_FRAME_EXT_SEQ, START_OF_FRAME_LOS_SEQ,
START_OF_FRAME_LOS_SEQ_AR, START_OF_FRAME_PROG_DCT_AR
};
/// Common Decode errors
#[allow(clippy::module_name_repetitions)]
pub enum DecodeErrors {
/// Any other thing we do not know
Format(String),
/// Any other thing we do not know but we
/// don't need to allocate space on the heap
FormatStatic(&'static str),
/// Illegal Magic Bytes
IllegalMagicBytes(u16),
/// problems with the Huffman Tables in a Decoder file
HuffmanDecode(String),
/// Image has zero width
ZeroError,
/// Discrete Quantization Tables error
DqtError(String),
/// Start of scan errors
SosError(String),
/// Start of frame errors
SofError(String),
/// UnsupportedImages
Unsupported(UnsupportedSchemes),
/// MCU errors
MCUError(String),
/// Exhausted data
ExhaustedData,
/// Large image dimensions(Corrupted data)?
LargeDimensions(usize),
/// Too small output for size
TooSmallOutput(usize, usize),
IoErrors(ZByteIoError)
}
#[cfg(feature = "std")]
impl std::error::Error for DecodeErrors {}
impl From<&'static str> for DecodeErrors {
fn from(data: &'static str) -> Self {
return Self::FormatStatic(data);
}
}
impl From<ZByteIoError> for DecodeErrors {
fn from(data: ZByteIoError) -> Self {
return Self::IoErrors(data);
}
}
impl Debug for DecodeErrors {
fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
match &self
{
Self::Format(ref a) => write!(f, "{a:?}"),
Self::FormatStatic(a) => write!(f, "{:?}", &a),
Self::HuffmanDecode(ref reason) =>
{
write!(f, "Error decoding huffman values: {reason}")
}
Self::ZeroError => write!(f, "Image width or height is set to zero, cannot continue"),
Self::DqtError(ref reason) => write!(f, "Error parsing DQT segment. Reason:{reason}"),
Self::SosError(ref reason) => write!(f, "Error parsing SOS Segment. Reason:{reason}"),
Self::SofError(ref reason) => write!(f, "Error parsing SOF segment. Reason:{reason}"),
Self::IllegalMagicBytes(bytes) =>
{
write!(f, "Error parsing image. Illegal start bytes:{bytes:X}")
}
Self::MCUError(ref reason) => write!(f, "Error in decoding MCU. Reason {reason}"),
Self::Unsupported(ref image_type) =>
{
write!(f, "{image_type:?}")
}
Self::ExhaustedData => write!(f, "Exhausted data in the image"),
Self::LargeDimensions(ref dimensions) => write!(
f,
"Too large dimensions {dimensions},library supports up to {}", crate::decoder::MAX_DIMENSIONS
),
Self::TooSmallOutput(expected, found) => write!(f, "Too small output, expected buffer with at least {expected} bytes but got one with {found} bytes"),
Self::IoErrors(error)=>write!(f,"I/O errors {error:?}"),
}
}
}
impl Display for DecodeErrors {
fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
write!(f, "{self:?}")
}
}
/// Contains Unsupported/Yet-to-be supported Decoder image encoding types.
#[derive(Eq, PartialEq, Copy, Clone)]
pub enum UnsupportedSchemes {
/// SOF_1 Extended sequential DCT,Huffman coding
ExtendedSequentialHuffman,
/// Lossless (sequential), huffman coding,
LosslessHuffman,
/// Extended sequential DEC, arithmetic coding
ExtendedSequentialDctArithmetic,
/// Progressive DCT, arithmetic coding,
ProgressiveDctArithmetic,
/// Lossless ( sequential), arithmetic coding
LosslessArithmetic
}
impl Debug for UnsupportedSchemes {
fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
match &self {
Self::ExtendedSequentialHuffman => {
write!(f, "The library cannot yet decode images encoded using Extended Sequential Huffman encoding scheme yet.")
}
Self::LosslessHuffman => {
write!(f, "The library cannot yet decode images encoded with Lossless Huffman encoding scheme")
}
Self::ExtendedSequentialDctArithmetic => {
write!(f,"The library cannot yet decode Images Encoded with Extended Sequential DCT Arithmetic scheme")
}
Self::ProgressiveDctArithmetic => {
write!(f,"The library cannot yet decode images encoded with Progressive DCT Arithmetic scheme")
}
Self::LosslessArithmetic => {
write!(f,"The library cannot yet decode images encoded with Lossless Arithmetic encoding scheme")
}
}
}
}
impl UnsupportedSchemes {
#[must_use]
/// Create an unsupported scheme from an integer
///
/// # Returns
/// `Some(UnsupportedScheme)` if the int refers to a specific scheme,
/// otherwise returns `None`
pub fn from_int(int: u8) -> Option<UnsupportedSchemes> {
let int = u16::from_be_bytes([0xff, int]);
match int {
START_OF_FRAME_PROG_DCT_AR => Some(Self::ProgressiveDctArithmetic),
START_OF_FRAME_LOS_SEQ => Some(Self::LosslessHuffman),
START_OF_FRAME_LOS_SEQ_AR => Some(Self::LosslessArithmetic),
START_OF_FRAME_EXT_SEQ => Some(Self::ExtendedSequentialHuffman),
START_OF_FRAME_EXT_AR => Some(Self::ExtendedSequentialDctArithmetic),
_ => None
}
}
}
+662
View File
@@ -0,0 +1,662 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
//! Decode Decoder markers/segments
//!
//! This file deals with decoding header information in a jpeg file
//!
use alloc::format;
use alloc::string::ToString;
use alloc::vec::Vec;
use zune_core::bytestream::ZByteReaderTrait;
use zune_core::colorspace::ColorSpace;
use zune_core::log::{debug, trace, warn};
use core::cmp::max;
use crate::components::{Components, SampleRatios};
use crate::decoder::{GainMapInfo, ICCChunk, JpegDecoder, MAX_COMPONENTS};
use crate::errors::DecodeErrors;
use crate::huffman::HuffmanTable;
use crate::misc::{SOFMarkers, UN_ZIGZAG};
///**B.2.4.2 Huffman table-specification syntax**
#[allow(clippy::similar_names, clippy::cast_sign_loss)]
pub(crate) fn parse_huffman<T: ZByteReaderTrait>(
decoder: &mut JpegDecoder<T>
) -> Result<(), DecodeErrors>
where
{
// Read the length of the Huffman table
let mut dht_length = i32::from(decoder.stream.get_u16_be_err()?.checked_sub(2).ok_or(
DecodeErrors::FormatStatic("Invalid Huffman length in image")
)?);
while dht_length > 16 {
// HT information
let ht_info = decoder.stream.read_u8_err()?;
// third bit indicates whether the huffman encoding is DC or AC type
let dc_or_ac = (ht_info >> 4) & 0xF;
// Indicate the position of this table, should be less than 4;
let index = (ht_info & 0xF) as usize;
// read the number of symbols
let mut num_symbols: [u8; 17] = [0; 17];
if index >= MAX_COMPONENTS {
return Err(DecodeErrors::HuffmanDecode(format!(
"Invalid DHT index {index}, expected between 0 and 3"
)));
}
if dc_or_ac > 1 {
return Err(DecodeErrors::HuffmanDecode(format!(
"Invalid DHT position {dc_or_ac}, should be 0 or 1"
)));
}
decoder.stream.read_exact_bytes(&mut num_symbols[1..17])?;
dht_length -= 1 + 16;
let symbols_sum: i32 = num_symbols.iter().map(|f| i32::from(*f)).sum();
// The sum of the number of symbols cannot be greater than 256;
if symbols_sum > 256 {
return Err(DecodeErrors::FormatStatic(
"Encountered Huffman table with excessive length in DHT"
));
}
if symbols_sum > dht_length {
return Err(DecodeErrors::HuffmanDecode(format!(
"Excessive Huffman table of length {symbols_sum} found when header length is {dht_length}"
)));
}
dht_length -= symbols_sum;
// A table containing symbols in increasing code length
let mut symbols = [0; 256];
decoder
.stream
.read_exact_bytes(&mut symbols[0..(symbols_sum as usize)])?;
// store
match dc_or_ac {
0 => {
decoder.dc_huffman_tables[index] = Some(HuffmanTable::new(
&num_symbols,
symbols,
true,
decoder.is_progressive
)?);
}
_ => {
decoder.ac_huffman_tables[index] = Some(HuffmanTable::new(
&num_symbols,
symbols,
false,
decoder.is_progressive
)?);
}
}
}
if dht_length > 0 {
return Err(DecodeErrors::FormatStatic("Bogus Huffman table definition"));
}
Ok(())
}
///**B.2.4.1 Quantization table-specification syntax**
#[allow(clippy::cast_possible_truncation, clippy::needless_range_loop)]
pub(crate) fn parse_dqt<T: ZByteReaderTrait>(img: &mut JpegDecoder<T>) -> Result<(), DecodeErrors> {
// read length
let mut qt_length =
img.stream
.get_u16_be_err()?
.checked_sub(2)
.ok_or(DecodeErrors::FormatStatic(
"Invalid DQT length. Length should be greater than 2"
))?;
// A single DQT header may have multiple QT's
while qt_length > 0 {
let qt_info = img.stream.read_u8_err()?;
// 0 = 8 bit otherwise 16 bit dqt
let precision = (qt_info >> 4) as usize;
// last 4 bits give us position
let table_position = (qt_info & 0x0f) as usize;
let precision_value = 64 * (precision + 1);
if (precision_value + 1) as u16 > qt_length {
return Err(DecodeErrors::DqtError(format!("Invalid QT table bytes left :{}. Too small to construct a valid qt table which should be {} long", qt_length, precision_value + 1)));
}
let dct_table = match precision {
0 => {
let mut qt_values = [0; 64];
img.stream.read_exact_bytes(&mut qt_values)?;
qt_length -= (precision_value as u16) + 1 /*QT BIT*/;
// carry out un zig-zag here
un_zig_zag(&qt_values)
}
1 => {
// 16 bit quantization tables
let mut qt_values = [0_u16; 64];
for i in 0..64 {
qt_values[i] = img.stream.get_u16_be_err()?;
}
qt_length -= (precision_value as u16) + 1;
un_zig_zag(&qt_values)
}
_ => {
return Err(DecodeErrors::DqtError(format!(
"Expected QT precision value of either 0 or 1, found {precision:?}"
)));
}
};
if table_position >= MAX_COMPONENTS {
return Err(DecodeErrors::DqtError(format!(
"Too large table position for QT :{table_position}, expected between 0 and 3"
)));
}
trace!("Assigning qt table {table_position} with precision {precision}");
img.qt_tables[table_position] = Some(dct_table);
}
return Ok(());
}
/// Section:`B.2.2 Frame header syntax`
pub(crate) fn parse_start_of_frame<T: ZByteReaderTrait>(
sof: SOFMarkers, img: &mut JpegDecoder<T>
) -> Result<(), DecodeErrors> {
if img.seen_sof {
return Err(DecodeErrors::SofError(
"Two Start of Frame Markers".to_string()
));
}
// Get length of the frame header
let length = img.stream.get_u16_be_err()?;
// usually 8, but can be 12 and 16, we currently support only 8
// so sorry about that 12 bit images
let dt_precision = img.stream.read_u8_err()?;
if dt_precision != 8 {
return Err(DecodeErrors::SofError(format!(
"The library can only parse 8-bit images, the image has {dt_precision} bits of precision"
)));
}
img.info.set_density(dt_precision);
// read and set the image height.
let img_height = img.stream.get_u16_be_err()?;
img.info.set_height(img_height);
// read and set the image width
let img_width = img.stream.get_u16_be_err()?;
img.info.set_width(img_width);
trace!("Image width :{}", img_width);
trace!("Image height :{}", img_height);
if usize::from(img_width) > img.options.max_width() {
return Err(DecodeErrors::Format(format!("Image width {} greater than width limit {}. If use `set_limits` if you want to support huge images", img_width, img.options.max_width())));
}
if usize::from(img_height) > img.options.max_height() {
return Err(DecodeErrors::Format(format!("Image height {} greater than height limit {}. If use `set_limits` if you want to support huge images", img_height, img.options.max_height())));
}
// Check image width or height is zero
if img_width == 0 || img_height == 0 {
return Err(DecodeErrors::ZeroError);
}
// Number of components for the image.
let num_components = img.stream.read_u8_err()?;
if num_components == 0 {
return Err(DecodeErrors::SofError(
"Number of components cannot be zero.".to_string()
));
}
let expected = 8 + 3 * u16::from(num_components);
// length should be equal to num components
if length != expected {
return Err(DecodeErrors::SofError(format!(
"Length of start of frame differs from expected {expected},value is {length}"
)));
}
trace!("Image components : {}", num_components);
if num_components == 1 {
// SOF sets the number of image components
// and that to us translates to setting input and output
// colorspaces to zero
img.input_colorspace = ColorSpace::Luma;
//img.options = img.options.jpeg_set_out_colorspace(ColorSpace::Luma);
debug!("Overriding default colorspace set to Luma");
}
if num_components == 4 && img.input_colorspace == ColorSpace::YCbCr {
trace!("Input image has 4 components, defaulting to CMYK colorspace");
// https://entropymine.wordpress.com/2018/10/22/how-is-a-jpeg-images-color-type-determined/
img.input_colorspace = ColorSpace::CMYK;
}
// set number of components
img.info.components = num_components;
let mut components = Vec::with_capacity(num_components as usize);
let mut temp = [0; 3];
for pos in 0..num_components {
// read 3 bytes for each component
img.stream.read_exact_bytes(&mut temp)?;
// create a component.
let component = Components::from(temp, pos)?;
components.push(component);
}
img.seen_sof = true;
img.info.set_sof_marker(sof);
img.components = components;
let mut h_max = 1;
let mut v_max = 1;
for comp in &img.components {
h_max = max(h_max, comp.horizontal_sample);
v_max = max(v_max, comp.vertical_sample);
}
img.info.sample_ratio = match (h_max, v_max) {
(1, 1) => SampleRatios::None,
(1, 2) => SampleRatios::V,
(2, 1) => SampleRatios::H,
(2, 2) => SampleRatios::HV,
(hs, vs) => SampleRatios::Generic(hs, vs)
};
Ok(())
}
/// Parse a start of scan data
pub(crate) fn parse_sos<T: ZByteReaderTrait>(
image: &mut JpegDecoder<T>
) -> Result<(), DecodeErrors> {
// Scan header length
let ls = usize::from(image.stream.get_u16_be_err()?);
// Number of image components in scan
let ns = image.stream.read_u8_err()?;
let mut seen: [_; 5] = [-1; { MAX_COMPONENTS + 1 }];
image.num_scans = ns;
let smallest_size = 6 + 2 * usize::from(ns);
if ls != smallest_size {
return Err(DecodeErrors::SosError(format!(
"Bad SOS length {ls},corrupt jpeg"
)));
}
// Check number of components.
if !(1..5).contains(&ns) {
return Err(DecodeErrors::SosError(format!(
"Invalid number of components in start of scan {ns}, expected in range 1..5"
)));
}
if image.info.components == 0 {
return Err(DecodeErrors::FormatStatic(
"Error decoding SOF Marker, Number of components cannot be zero."
));
}
// consume spec parameters
image.scan_subsampled = false;
for i in 0..ns {
let id = image.stream.read_u8_err()?;
if seen.contains(&i32::from(id)) {
return Err(DecodeErrors::SofError(format!(
"Duplicate ID {id} seen twice in the same component"
)));
}
seen[usize::from(i)] = i32::from(id);
// DC and AC huffman table position
// top 4 bits contain dc huffman destination table
// lower four bits contain ac huffman destination table
let y = image.stream.read_u8_err()?;
let mut j = 0;
while j < image.info.components {
if image.components[j as usize].id == id {
break;
}
j += 1;
}
if j == image.info.components {
return Err(DecodeErrors::SofError(format!(
"Invalid component id {}, expected one one of {:?}",
id,
image.components.iter().map(|c| c.id).collect::<Vec<_>>()
)));
}
let component = &mut image.components[usize::from(j)];
component.dc_huff_table = usize::from((y >> 4) & 0xF);
component.ac_huff_table = usize::from(y & 0xF);
image.z_order[i as usize] = j as usize;
if component.vertical_sample != 1 || component.horizontal_sample != 1 {
image.scan_subsampled = true;
}
trace!(
"Assigned huffman tables {}/{} to component {j}, id={}",
image.components[usize::from(j)].dc_huff_table,
image.components[usize::from(j)].ac_huff_table,
image.components[usize::from(j)].id,
);
}
// Collect the component spec parameters
// This is only needed for progressive images but I'll read
// them in order to ensure they are correct according to the spec
// Extract progressive information
// https://www.w3.org/Graphics/JPEG/itu-t81.pdf
// Page 42
// Start of spectral / predictor selection. (between 0 and 63)
image.spec_start = image.stream.read_u8_err()?;
// End of spectral selection
image.spec_end = image.stream.read_u8_err()?;
let bit_approx = image.stream.read_u8_err()?;
// successive approximation bit position high
image.succ_high = bit_approx >> 4;
if image.spec_end > 63 {
return Err(DecodeErrors::SosError(format!(
"Invalid Se parameter {}, range should be 0-63",
image.spec_end
)));
}
if image.spec_start > 63 {
return Err(DecodeErrors::SosError(format!(
"Invalid Ss parameter {}, range should be 0-63",
image.spec_start
)));
}
if image.succ_high > 13 {
return Err(DecodeErrors::SosError(format!(
"Invalid Ah parameter {}, range should be 0-13",
image.succ_low
)));
}
// successive approximation bit position low
image.succ_low = bit_approx & 0xF;
if image.succ_low > 13 {
return Err(DecodeErrors::SosError(format!(
"Invalid Al parameter {}, range should be 0-13",
image.succ_low
)));
}
// skip any bytes not read
image.stream.skip(smallest_size.saturating_sub(ls))?;
trace!(
"Ss={}, Se={} Ah={} Al={}",
image.spec_start,
image.spec_end,
image.succ_high,
image.succ_low
);
Ok(())
}
/// Parse the APP13 (IPTC) segment.
pub(crate) fn parse_app13<T: ZByteReaderTrait>(
decoder: &mut JpegDecoder<T>
) -> Result<(), DecodeErrors> {
const IPTC_PREFIX: &[u8] = b"Photoshop 3.0\0";
// skip length.
let mut length = usize::from(decoder.stream.get_u16_be());
if length < 2 {
return Err(DecodeErrors::FormatStatic("Too small APP13 length"));
}
// length bytes.
length -= 2;
if length > IPTC_PREFIX.len() && decoder.stream.peek_at(0, IPTC_PREFIX.len())? == IPTC_PREFIX {
// skip bytes we read above.
decoder.stream.skip(IPTC_PREFIX.len())?;
length -= IPTC_PREFIX.len();
let iptc_bytes = decoder.stream.peek_at(0, length)?.to_vec();
decoder.info.iptc_data = Some(iptc_bytes);
}
decoder.stream.skip(length)?;
Ok(())
}
/// Parse Adobe App14 segment
pub(crate) fn parse_app14<T: ZByteReaderTrait>(
decoder: &mut JpegDecoder<T>
) -> Result<(), DecodeErrors> {
// skip length
let mut length = usize::from(decoder.stream.get_u16_be());
if length < 2 {
return Err(DecodeErrors::FormatStatic("Too small APP14 length"));
}
if decoder.stream.peek_at(0, 5)? == b"Adobe" {
if length < 14 {
return Err(DecodeErrors::FormatStatic(
"Too short of a length for App14 segment"
));
}
// move stream 6 bytes to remove adobe id
decoder.stream.skip(6)?;
// skip version, flags0 and flags1
decoder.stream.skip(5)?;
// get color transform
let transform = decoder.stream.read_u8();
// https://exiftool.org/TagNames/JPEG.html#Adobe
match transform {
0 => decoder.input_colorspace = ColorSpace::CMYK,
1 => decoder.input_colorspace = ColorSpace::YCbCr,
2 => decoder.input_colorspace = ColorSpace::YCCK,
_ => {
return Err(DecodeErrors::Format(format!(
"Unknown Adobe colorspace {transform}"
)))
}
}
// length = 2
// adobe id = 6
// version = 5
// transform = 1
length = length.saturating_sub(14);
} else {
warn!("Not a valid Adobe APP14 Segment, skipping {} bytes", length);
length = length.saturating_sub(2);
}
// skip any proceeding lengths.
// we do not need them
decoder.stream.skip(length)?;
Ok(())
}
/// Parse the APP1 segment
///
/// This contains the exif tag
pub(crate) fn parse_app1<T: ZByteReaderTrait>(
decoder: &mut JpegDecoder<T>
) -> Result<(), DecodeErrors> {
const XMP_NAMESPACE_PREFIX: &[u8] = b"http://ns.adobe.com/xap/1.0/\0";
// contains exif data
let mut length = usize::from(decoder.stream.get_u16_be());
if length < 2 {
return Err(DecodeErrors::FormatStatic("Too small app1 length"));
}
// length bytes
length -= 2;
if length > 6 && decoder.stream.peek_at(0, 6)? == b"Exif\x00\x00" {
trace!("Exif segment present");
// skip bytes we read above
decoder.stream.skip(6)?;
length -= 6;
let exif_bytes = decoder.stream.peek_at(0, length)?.to_vec();
decoder.info.exif_data = Some(exif_bytes);
} else if length > XMP_NAMESPACE_PREFIX.len()
&& decoder.stream.peek_at(0, XMP_NAMESPACE_PREFIX.len())? == XMP_NAMESPACE_PREFIX
{
trace!("XMP Data Present");
decoder.stream.skip(XMP_NAMESPACE_PREFIX.len())?;
length -= XMP_NAMESPACE_PREFIX.len();
let xmp_data = decoder.stream.peek_at(0, length)?.to_vec();
decoder.info.xmp_data = Some(xmp_data);
} else {
warn!("Unknown format for APP1 tag, skipping");
}
decoder.stream.skip(length)?;
Ok(())
}
pub(crate) fn parse_app2<T: ZByteReaderTrait>(
decoder: &mut JpegDecoder<T>
) -> Result<(), DecodeErrors> {
static HDR_META: &[u8] = b"urn:iso:std:iso:ts:21496:-1\0";
static MPF_DATA: &[u8] = b"MPF\0";
let mut length = usize::from(decoder.stream.get_u16_be());
if length < 2 {
return Err(DecodeErrors::FormatStatic("Too small app2 segment"));
}
// length bytes
length -= 2;
if length > 14 && decoder.stream.peek_at(0, 12)? == *b"ICC_PROFILE\0" {
trace!("ICC Profile present");
// skip 12 bytes which indicate ICC profile
length -= 12;
decoder.stream.skip(12)?;
let seq_no = decoder.stream.read_u8();
let num_markers = decoder.stream.read_u8();
// deduct the two bytes we read above
length -= 2;
let data = decoder.stream.peek_at(0, length)?.to_vec();
let icc_chunk = ICCChunk {
seq_no,
num_markers,
data
};
decoder.icc_data.push(icc_chunk);
} else if length > HDR_META.len() && decoder.stream.peek_at(0, HDR_META.len())? == HDR_META {
length = length.saturating_sub(HDR_META.len());
decoder.stream.skip(HDR_META.len())?;
trace!("Gain Map metadata found");
match length {
4 => {
// If gain map metadata length == 4 then here it variables
// https://github.com/google/libultrahdr/blob/bf2aa439eea9ad5da483003fa44182f990f74091/lib/src/jpegr.cpp#L1076C1-L1077C35
// 2 bytes minimum_version: (00 00)
// 2 bytes writer_version: (00 00)
// Perhaps nothing to do with it ?
let _ = decoder.stream.get_u16_be();
let _ = decoder.stream.get_u16_be();
length -= 4;
decoder
.info
.gain_map_info
.push(GainMapInfo { data: Vec::new() });
}
n if n > 4 => {
// If there is perhaps useful gain map info
// we'll read this until end
// https://github.com/google/libultrahdr/blob/bf2aa439eea9ad5da483003fa44182f990f74091/lib/src/jpegr.cpp#L1323
let data = decoder.stream.peek_at(0, length)?.to_vec();
length -= data.len();
decoder.stream.skip(data.len())?;
decoder.info.gain_map_info.push(GainMapInfo { data });
}
_ => {}
}
} else if length > MPF_DATA.len() && decoder.stream.peek_at(0, MPF_DATA.len())? == MPF_DATA {
trace!("MPF Signature present");
length = length.saturating_sub(MPF_DATA.len());
decoder.stream.skip(MPF_DATA.len())?;
// MPF signature taken from here
// https://github.com/google/libultrahdr/blob/bf2aa439eea9ad5da483003fa44182f990f74091/lib/include/ultrahdr/multipictureformat.h#L50
// https://github.com/google/libultrahdr/blob/bf2aa439eea9ad5da483003fa44182f990f74091/lib/src/multipictureformat.cpp#L36
// More info https://www.cipa.jp/std/documents/e/DC-X007-KEY_E.pdf
let data = decoder.stream.peek_at(0, length)?.to_vec();
length -= data.len();
decoder.stream.skip(data.len())?;
decoder.info.multi_picture_information = Some(data);
}
decoder.stream.skip(length)?;
Ok(())
}
/// Small utility function to print Un-zig-zagged quantization tables
fn un_zig_zag<T>(a: &[T]) -> [i32; 64]
where
T: Default + Copy,
i32: core::convert::From<T>
{
let mut output = [i32::default(); 64];
for i in 0..64 {
output[UN_ZIGZAG[i]] = i32::from(a[i]);
}
output
}
+254
View File
@@ -0,0 +1,254 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
//! This file contains a single struct `HuffmanTable` that
//! stores Huffman tables needed during `BitStream` decoding.
#![allow(clippy::similar_names, clippy::module_name_repetitions)]
use alloc::string::ToString;
use crate::errors::DecodeErrors;
/// Determines how many bits of lookahead we have for our bitstream decoder.
pub const HUFF_LOOKAHEAD: u8 = 9;
/// A struct which contains necessary tables for decoding a JPEG
/// huffman encoded bitstream
pub struct HuffmanTable {
// element `[0]` of each array is unused
/// largest code of length k
pub(crate) maxcode: [i32; 18],
/// offset for codes of length k
/// Answers the question, where do code-lengths of length k end
/// Element 0 is unused
pub(crate) offset: [i32; 18],
/// lookup table for fast decoding
///
/// top bits above HUFF_LOOKAHEAD contain the code length.
///
/// Lower (8) bits contain the symbol in order of increasing code length.
pub(crate) lookup: [i32; 1 << HUFF_LOOKAHEAD],
/// A table which can be used to decode small AC coefficients and
/// do an equivalent of receive_extend
pub(crate) ac_lookup: Option<[i16; 1 << HUFF_LOOKAHEAD]>,
/// Directly represent contents of a JPEG DHT marker
///
/// \# number of symbols with codes of length `k` bits
// bits[0] is unused
/// Symbols in order of increasing code length
pub(crate) values: [u8; 256]
}
impl HuffmanTable {
pub fn new(
codes: &[u8; 17], values: [u8; 256], is_dc: bool, is_progressive: bool
) -> Result<HuffmanTable, DecodeErrors> {
let too_long_code = (i32::from(HUFF_LOOKAHEAD) + 1) << HUFF_LOOKAHEAD;
let mut p = HuffmanTable {
maxcode: [0; 18],
offset: [0; 18],
lookup: [too_long_code; 1 << HUFF_LOOKAHEAD],
values,
ac_lookup: None
};
p.make_derived_table(is_dc, is_progressive, codes)?;
Ok(p)
}
/// Create a new huffman tables with values that aren't fixed
/// used by fill_mjpeg_tables
pub fn new_unfilled(
codes: &[u8; 17], values: &[u8], is_dc: bool, is_progressive: bool
) -> Result<HuffmanTable, DecodeErrors> {
let mut buf = [0; 256];
buf[..values.len()].copy_from_slice(values);
HuffmanTable::new(codes, buf, is_dc, is_progressive)
}
/// Compute derived values for a Huffman table
///
/// This routine performs some validation checks on the table
#[allow(
clippy::cast_possible_truncation,
clippy::cast_possible_wrap,
clippy::cast_sign_loss,
clippy::too_many_lines,
clippy::needless_range_loop
)]
fn make_derived_table(
&mut self, is_dc: bool, _is_progressive: bool, bits: &[u8; 17]
) -> Result<(), DecodeErrors> {
// build a list of code size
let mut huff_size = [0; 257];
// Huffman code lengths
let mut huff_code: [u32; 257] = [0; 257];
// figure C.1 make table of Huffman code length for each symbol
let mut p = 0;
for l in 1..=16 {
let mut i = i32::from(bits[l]);
// table overrun is checked before ,so we dont need to check
while i != 0 {
huff_size[p] = l as u8;
p += 1;
i -= 1;
}
}
huff_size[p] = 0;
let num_symbols = p;
// Generate the codes themselves
// We also validate that the counts represent a legal Huffman code tree
let mut code = 0;
let mut si = i32::from(huff_size[0]);
p = 0;
while huff_size[p] != 0 {
while i32::from(huff_size[p]) == si {
huff_code[p] = code;
code += 1;
p += 1;
}
// maximum code of length si, pre-shifted by 16-k bits
self.maxcode[si as usize] = (code << (16 - si)) as i32;
// code is now 1 more than the last code used for code-length si; but
// it must still fit in si bits, since no code is allowed to be all ones.
if (code as i32) >= (1 << si) {
return Err(DecodeErrors::HuffmanDecode("Bad Huffman Table".to_string()));
}
code <<= 1;
si += 1;
}
// Figure F.15 generate decoding tables for bit-sequential decoding
p = 0;
for l in 0..=16 {
if bits[l] == 0 {
// -1 if no codes of this length
self.maxcode[l] = -1;
} else {
// offset[l]=codes[index of 1st symbol of code length l
// minus minimum code of length l]
self.offset[l] = (p as i32) - (huff_code[p]) as i32;
p += usize::from(bits[l]);
}
}
self.offset[17] = 0;
// we ensure that decode terminates
self.maxcode[17] = 0x000F_FFFF;
/*
* Compute lookahead tables to speed up decoding.
* First we set all the table entries to 0(left justified), indicating "too long";
* (Note too long was set during initialization)
* then we iterate through the Huffman codes that are short enough and
* fill in all the entries that correspond to bit sequences starting
* with that code.
*/
p = 0;
for l in 1..=HUFF_LOOKAHEAD {
for _ in 1..=i32::from(bits[usize::from(l)]) {
// l -> Current code length,
// p => Its index in self.code and self.values
// Generate left justified code followed by all possible bit sequences
let mut look_bits = (huff_code[p] as usize) << (HUFF_LOOKAHEAD - l);
for _ in 0..1 << (HUFF_LOOKAHEAD - l) {
self.lookup[look_bits] =
(i32::from(l) << HUFF_LOOKAHEAD) | i32::from(self.values[p]);
look_bits += 1;
}
p += 1;
}
}
// build an ac table that does an equivalent of decode and receive_extend
if !is_dc {
let mut fast = [255; 1 << HUFF_LOOKAHEAD];
// Iterate over number of symbols
for i in 0..num_symbols {
// get code size for an item
let s = huff_size[i];
if s <= HUFF_LOOKAHEAD {
// if it's lower than what we need for our lookup table create the table
let c = (huff_code[i] << (HUFF_LOOKAHEAD - s)) as usize;
let m = (1 << (HUFF_LOOKAHEAD - s)) as usize;
for j in 0..m {
fast[c + j] = i as i16;
}
}
}
// build a table that decodes both magnitude and value of small ACs in
// one go.
let mut fast_ac = [0; 1 << HUFF_LOOKAHEAD];
for i in 0..(1 << HUFF_LOOKAHEAD) {
let fast_v = fast[i];
if fast_v < 255 {
// get symbol value from AC table
let rs = self.values[fast_v as usize];
// shift by 4 to get run length
let run = i16::from((rs >> 4) & 15);
// get magnitude bits stored at the lower 3 bits
let mag_bits = i16::from(rs & 15);
// length of the bit we've read
let len = i16::from(huff_size[fast_v as usize]);
if mag_bits != 0 && (len + mag_bits) <= i16::from(HUFF_LOOKAHEAD) {
// magnitude code followed by receive_extend code
let mut k = (((i as i16) << len) & ((1 << HUFF_LOOKAHEAD) - 1))
>> (i16::from(HUFF_LOOKAHEAD) - mag_bits);
let m = 1 << (mag_bits - 1);
if k < m {
k += (!0_i16 << mag_bits) + 1;
};
// if result is small enough fit into fast ac table
if (-128..=127).contains(&k) {
fast_ac[i] = (k << 8) + (run << 4) + (len + mag_bits);
}
}
}
}
self.ac_lookup = Some(fast_ac);
}
// Validate symbols as being reasonable
// For AC tables, we make no check, but accept all byte values 0..255
// For DC tables, we require symbols to be in range 0..15
if is_dc {
for i in 0..num_symbols {
let sym = self.values[i];
if sym > 15 {
return Err(DecodeErrors::HuffmanDecode("Bad Huffman Table".to_string()));
}
}
}
Ok(())
}
}
+206
View File
@@ -0,0 +1,206 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
//! Routines for IDCT
//!
//! Essentially we provide 2 routines for IDCT, a scalar implementation and a not super optimized
//! AVX2 one, i'll talk about them here.
//!
//! There are 2 reasons why we have the avx one
//! 1. No one compiles with -C target-features=avx2 hence binaries won't probably take advantage(even
//! if it exists).
//! 2. AVX employs zero short circuit in a way the scalar code cannot employ it.
//! - AVX does this by checking for MCU's whose 63 AC coefficients are zero and if true, it writes
//! values directly, if false, it goes the long way of calculating.
//! - Although this can be trivially implemented in the scalar version, it generates code
//! I'm not happy width(scalar version that basically loops and that is too many branches for me)
//! The avx one does a better job of using bitwise or's with (`_mm256_or_si256`) which is magnitudes of faster
//! than anything I could come up with
//!
//! The AVX code also has some cool transpose_u16 instructions which look so complicated to be cool
//! (spoiler alert, i barely understand how it works, that's why I credited the owner).
//!
#![allow(
clippy::excessive_precision,
clippy::unreadable_literal,
clippy::module_name_repetitions,
unused_parens,
clippy::wildcard_imports
)]
use zune_core::log::debug;
use zune_core::options::DecoderOptions;
use crate::decoder::IDCTPtr;
use crate::idct::scalar::{idct_int, idct_int_1x1};
#[cfg(feature = "x86")]
pub mod avx2;
#[cfg(feature = "neon")]
pub mod neon;
pub mod scalar;
/// Choose an appropriate IDCT function
#[allow(unused_variables)]
pub fn choose_idct_func(options: &DecoderOptions) -> IDCTPtr {
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[cfg(feature = "x86")]
{
if options.use_avx2() {
debug!("Using vector integer IDCT");
return |a: &mut [i32; 64], b: &mut [i16], c: usize| {
// SAFETY: `options.use_avx2()` only returns true if avx2 is supported.
unsafe { avx2::idct_avx2(a,b,c) }
};
}
}
#[cfg(target_arch = "aarch64")]
#[cfg(feature = "neon")]
{
if options.use_neon() {
debug!("Using vector integer IDCT");
return |a: &mut [i32; 64], b: &mut [i16], c: usize| {
// SAFETY: `options.use_neon()` only returns true if neon is supported.
unsafe { neon::idct_neon(a,b,c) }
};
}
}
debug!("Using scalar integer IDCT");
// use generic one
return idct_int;
}
/// Choose a function to implement 4x4 IDCT.
///
/// These functions get the same input but have an extra contract: Only the first 4x4 block of
/// coefficients are non-zero. All other entries are zeroed.
///
/// **The callee must uphold that contract on return**
pub fn choose_idct_4x4_func(_options: &DecoderOptions) -> IDCTPtr {
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[cfg(feature = "x86")]
{
if _options.use_avx2() {
debug!("Using vector integer IDCT");
return |a: &mut [i32; 64], b: &mut [i16], c: usize| {
// SAFETY: `options.use_avx2()` only returns true if avx2 is supported.
unsafe { avx2::idct_avx2_4x4(a,b,c) }
};
}
}
scalar::idct4x4
}
pub fn choose_idct_1x1_func(_: &DecoderOptions) -> IDCTPtr {
// These are simple stores, no alternative implementation for now
idct_int_1x1
}
#[cfg(test)]
#[allow(unreachable_code)]
#[allow(dead_code)]
mod tests {
use super::*;
#[test]
fn idct_test0() {
let stride = 8;
let mut coeff = [10; 64];
let mut coeff2 = [10; 64];
let mut output_scalar = [0; 64];
let mut output_vector = [0; 64];
let idct_func = choose_idct_func(&DecoderOptions::new_fast());
idct_func(&mut coeff, &mut output_vector, stride);
idct_int(&mut coeff2, &mut output_scalar, stride);
assert_eq!(output_scalar, output_vector, "IDCT and scalar do not match");
}
#[test]
fn do_idct_test1() {
let stride = 8;
let mut coeff = [14; 64];
let mut coeff2 = [14; 64];
let mut output_scalar = [0; 64];
let mut output_vector = [0; 64];
let idct_func = choose_idct_func(&DecoderOptions::new_fast());
idct_func(&mut coeff, &mut output_vector, stride);
idct_int(&mut coeff2, &mut output_scalar, stride);
assert_eq!(output_scalar, output_vector, "IDCT and scalar do not match");
}
#[test]
fn do_idct_test2() {
let stride = 8;
let mut coeff = [0; 64];
coeff[0] = 255;
coeff[63] = -256;
let mut coeff2 = coeff;
let mut output_scalar = [0; 64];
let mut output_vector = [0; 64];
let idct_func = choose_idct_func(&DecoderOptions::new_fast());
idct_func(&mut coeff, &mut output_vector, stride);
idct_int(&mut coeff2, &mut output_scalar, stride);
assert_eq!(output_scalar, output_vector, "IDCT and scalar do not match");
}
#[test]
fn do_idct_zeros() {
let stride = 8;
let mut coeff = [0; 64];
let mut coeff2 = [0; 64];
let mut output_scalar = [0; 64];
let mut output_vector = [0; 64];
let idct_func = choose_idct_func(&DecoderOptions::new_fast());
idct_func(&mut coeff, &mut output_vector, stride);
idct_int(&mut coeff2, &mut output_scalar, stride);
assert_eq!(output_scalar, output_vector, "IDCT and scalar do not match");
}
#[test]
fn idct_4x4() {
#[rustfmt::skip]
const A: [i32; 32] = [
-254, -7, 0, 0, 0, 0, 0, 0,
7, 0, -30, 32, 0, 0, 0, 0,
7, 0, -30, 32, 0, 0, 0, 0,
7, 0, -30, 32, 0, 0, 0, 0,
];
let v: Vec<IDCTPtr> = vec![
choose_idct_func(&DecoderOptions::new_safe()),
choose_idct_4x4_func(&DecoderOptions::new_safe()),
choose_idct_func(&DecoderOptions::new_fast()),
choose_idct_4x4_func(&DecoderOptions::new_fast()),
];
let dct_names = vec![
"safe idct",
"safe idct 4x4",
"fast idct",
"fast idct 4x4",
];
let mut color = vec![];
for idct in v {
let mut a = [0i32; 64];
a[..32].copy_from_slice(&A);
let mut b = [0i16; 64];
idct(&mut a, &mut b, 8);
color.push(b);
}
for (wnd, name) in color.windows(2).zip(&dct_names) {
let [a, b] = wnd else { unreachable!() };
assert_eq!(a, b, "{name}");
}
}
}
+398
View File
@@ -0,0 +1,398 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
#![cfg(any(target_arch = "x86", target_arch = "x86_64"))]
//! AVX optimised IDCT.
//!
//! Okay not thaat optimised.
//!
//!
//! # The implementation
//! The implementation is neatly broken down into two operations.
//!
//! 1. Test for zeroes
//! > There is a shortcut method for idct where when all AC values are zero, we can get the answer really quickly.
//! by scaling the 1/8th of the DCT coefficient of the block to the whole block and level shifting.
//!
//! 2. If above fails, we proceed to carry out IDCT as a two pass one dimensional algorithm.
//! IT does two whole scans where it carries out IDCT on all items
//! After each successive scan, data is transposed in register(thank you x86 SIMD powers). and the second
//! pass is carried out.
//!
//! The code is not super optimized, it produces bit identical results with scalar code hence it's
//! `mm256_add_epi16`
//! and it also has the advantage of making this implementation easy to maintain.
#![cfg(feature = "x86")]
#![allow(dead_code)]
#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;
use crate::unsafe_utils::{transpose, YmmRegister};
const SCALE_BITS: i32 = 512 + 65536 + (128 << 17);
// Pack i32 to i16's,
// clamp them to be between 0-255
// Undo shuffling
// Store back to array
macro_rules! permute_store {
($x:tt,$y:tt,$index:tt,$out:tt,$stride:tt) => {
let a = _mm256_packs_epi32($x, $y);
// Clamp the values after packing, we can clamp more values at once
let b = clamp_avx(a);
// /Undo shuffling
let c = _mm256_permute4x64_epi64(b, shuffle(3, 1, 2, 0));
// store first vector
_mm_storeu_si128(
($out)
.get_mut($index..$index + 8)
.unwrap()
.as_mut_ptr()
.cast(),
_mm256_extractf128_si256::<0>(c),
);
$index += $stride;
// second vector
_mm_storeu_si128(
($out)
.get_mut($index..$index + 8)
.unwrap()
.as_mut_ptr()
.cast(),
_mm256_extractf128_si256::<1>(c),
);
$index += $stride;
};
}
#[target_feature(enable = "avx2")]
#[allow(
clippy::too_many_lines,
clippy::cast_possible_truncation,
clippy::similar_names,
clippy::op_ref,
unused_assignments,
clippy::zero_prefixed_literal
)]
pub unsafe fn idct_avx2(
in_vector: &mut [i32; 64], out_vector: &mut [i16], stride: usize,
) {
let mut pos = 0;
// load into registers
//
// We sign extend i16's to i32's and calculate them with extended precision and
// later reduce them to i16's when we are done carrying out IDCT
let rw0 = _mm256_loadu_si256(in_vector[00..].as_ptr().cast());
let rw1 = _mm256_loadu_si256(in_vector[08..].as_ptr().cast());
let rw2 = _mm256_loadu_si256(in_vector[16..].as_ptr().cast());
let rw3 = _mm256_loadu_si256(in_vector[24..].as_ptr().cast());
let rw4 = _mm256_loadu_si256(in_vector[32..].as_ptr().cast());
let rw5 = _mm256_loadu_si256(in_vector[40..].as_ptr().cast());
let rw6 = _mm256_loadu_si256(in_vector[48..].as_ptr().cast());
let rw7 = _mm256_loadu_si256(in_vector[56..].as_ptr().cast());
// Forward DCT and quantization may cause all the AC terms to be zero, for such
// cases we can try to accelerate it
// Basically the poop is that whenever the array has 63 zeroes, its idct is
// (arr[0]>>3)or (arr[0]/8) propagated to all the elements.
// We first test to see if the array contains zero elements and if it does, we go the
// short way.
//
// This reduces IDCT overhead from about 39% to 18 %, almost half
// Do another load for the first row, we don't want to check DC value, because
// we only care about AC terms
let rw8 = _mm256_loadu_si256(in_vector[1..].as_ptr().cast());
let mut bitmap = _mm256_or_si256(rw1, rw2);
bitmap = _mm256_or_si256(bitmap, rw3);
bitmap = _mm256_or_si256(bitmap, rw4);
bitmap = _mm256_or_si256(bitmap, rw5);
bitmap = _mm256_or_si256(bitmap, rw6);
bitmap = _mm256_or_si256(bitmap, rw7);
bitmap = _mm256_or_si256(bitmap, rw8);
if _mm256_testz_si256(bitmap, bitmap) == 1 {
// AC terms all zero, idct of the block is ( coeff[0] * qt[0] )/8 + 128 (bias)
// (and clamped to 255)
// Round by adding 0.5 * (1 << 3) and offset by adding (128 << 3) before scaling
let coeff = ((in_vector[0] + 4 + 1024) >> 3).clamp(0, 255) as i16;
let idct_value = _mm_set1_epi16(coeff);
macro_rules! store {
($pos:tt,$value:tt) => {
// store
_mm_storeu_si128(
out_vector
.get_mut($pos..$pos + 8)
.unwrap()
.as_mut_ptr()
.cast(),
$value,
);
$pos += stride;
};
}
store!(pos, idct_value);
store!(pos, idct_value);
store!(pos, idct_value);
store!(pos, idct_value);
store!(pos, idct_value);
store!(pos, idct_value);
store!(pos, idct_value);
store!(pos, idct_value);
return;
}
let mut row0 = YmmRegister { mm256: rw0 };
let mut row1 = YmmRegister { mm256: rw1 };
let mut row2 = YmmRegister { mm256: rw2 };
let mut row3 = YmmRegister { mm256: rw3 };
let mut row4 = YmmRegister { mm256: rw4 };
let mut row5 = YmmRegister { mm256: rw5 };
let mut row6 = YmmRegister { mm256: rw6 };
let mut row7 = YmmRegister { mm256: rw7 };
macro_rules! dct_pass {
($SCALE_BITS:tt,$scale:tt) => {
// There are a lot of ways to do this
// but to keep it simple(and beautiful), ill make a direct translation of the
// scalar code to also make this code fully transparent(this version and the non
// avx one should produce identical code.)
// even part
let p1 = (row2 + row6) * 2217;
let mut t2 = p1 + row6 * -7567;
let mut t3 = p1 + row2 * 3135;
let mut t0 = YmmRegister {
mm256: _mm256_slli_epi32((row0 + row4).mm256, 12),
};
let mut t1 = YmmRegister {
mm256: _mm256_slli_epi32((row0 - row4).mm256, 12),
};
let x0 = t0 + t3 + $SCALE_BITS;
let x3 = t0 - t3 + $SCALE_BITS;
let x1 = t1 + t2 + $SCALE_BITS;
let x2 = t1 - t2 + $SCALE_BITS;
let p3 = row7 + row3;
let p4 = row5 + row1;
let p1 = row7 + row1;
let p2 = row5 + row3;
let p5 = (p3 + p4) * 4816;
t0 = row7 * 1223;
t1 = row5 * 8410;
t2 = row3 * 12586;
t3 = row1 * 6149;
let p1 = p5 + p1 * -3685;
let p2 = p5 + (p2 * -10497);
let p3 = p3 * -8034;
let p4 = p4 * -1597;
t3 += p1 + p4;
t2 += p2 + p3;
t1 += p2 + p4;
t0 += p1 + p3;
row0.mm256 = _mm256_srai_epi32((x0 + t3).mm256, $scale);
row1.mm256 = _mm256_srai_epi32((x1 + t2).mm256, $scale);
row2.mm256 = _mm256_srai_epi32((x2 + t1).mm256, $scale);
row3.mm256 = _mm256_srai_epi32((x3 + t0).mm256, $scale);
row4.mm256 = _mm256_srai_epi32((x3 - t0).mm256, $scale);
row5.mm256 = _mm256_srai_epi32((x2 - t1).mm256, $scale);
row6.mm256 = _mm256_srai_epi32((x1 - t2).mm256, $scale);
row7.mm256 = _mm256_srai_epi32((x0 - t3).mm256, $scale);
};
}
// Process rows
dct_pass!(512, 10);
transpose(
&mut row0, &mut row1, &mut row2, &mut row3, &mut row4, &mut row5, &mut row6, &mut row7,
);
// process columns
dct_pass!(SCALE_BITS, 17);
transpose(
&mut row0, &mut row1, &mut row2, &mut row3, &mut row4, &mut row5, &mut row6, &mut row7,
);
// Pack and write the values back to the array
permute_store!((row0.mm256), (row1.mm256), pos, out_vector, stride);
permute_store!((row2.mm256), (row3.mm256), pos, out_vector, stride);
permute_store!((row4.mm256), (row5.mm256), pos, out_vector, stride);
permute_store!((row6.mm256), (row7.mm256), pos, out_vector, stride);
}
#[target_feature(enable = "avx2")]
#[allow(
clippy::too_many_lines,
clippy::cast_possible_truncation,
clippy::similar_names,
clippy::op_ref,
unused_assignments,
clippy::zero_prefixed_literal
)]
pub unsafe fn idct_avx2_4x4(
in_vector: &mut [i32; 64], out_vector: &mut [i16], stride: usize,
) {
let rw0 = _mm256_loadu_si256(in_vector[00..].as_ptr().cast());
let rw1 = _mm256_loadu_si256(in_vector[08..].as_ptr().cast());
let rw2 = _mm256_loadu_si256(in_vector[16..].as_ptr().cast());
let rw3 = _mm256_loadu_si256(in_vector[24..].as_ptr().cast());
let mut row0 = YmmRegister { mm256: rw0 };
let mut row1 = YmmRegister { mm256: rw1 };
let mut row2 = YmmRegister { mm256: rw2 };
let mut row3 = YmmRegister { mm256: rw3 };
let mut row4 = YmmRegister { mm256: rw0 };
let mut row5 = YmmRegister { mm256: rw0 };
let mut row6 = YmmRegister { mm256: rw0 };
let mut row7 = YmmRegister { mm256: rw0 };
{
row0.mm256 = _mm256_slli_epi32(row0.mm256, 12);
row0 += 512;
let i2 = row2;
let p1 = i2 * 2217;
let p3 = i2 * 5352;
let x0 = row0 + p3;
let x1 = row0 + p1;
let x2 = row0 - p1;
let x3 = row0 - p3;
// odd part
let i4 = row3;
let i3 = row1;
let p5 = (i4 + i3) * 4816;
let p1 = p5 + i3 * -3685;
let p2 = p5 + i4 * -10497;
let t3 = p5 + i3 * 867;
let t2 = p5 + i4 * -5945;
let t1 = p2 + i3 * -1597;
let t0 = p1 + i4 * -8034;
row0.mm256 = _mm256_srai_epi32((x0 + t3).mm256, 10);
row1.mm256 = _mm256_srai_epi32((x1 + t2).mm256, 10);
row2.mm256 = _mm256_srai_epi32((x2 + t1).mm256, 10);
row3.mm256 = _mm256_srai_epi32((x3 + t0).mm256, 10);
row4.mm256 = _mm256_srai_epi32((x3 - t0).mm256, 10);
row5.mm256 = _mm256_srai_epi32((x2 - t1).mm256, 10);
row6.mm256 = _mm256_srai_epi32((x1 - t2).mm256, 10);
row7.mm256 = _mm256_srai_epi32((x0 - t3).mm256, 10);
}
transpose(
&mut row0, &mut row1, &mut row2, &mut row3, &mut row4, &mut row5, &mut row6, &mut row7,
);
{
let i2 = row2;
let i0 = row0;
row0.mm256 = _mm256_slli_epi32(i0.mm256, 12);
let t0 = row0 + SCALE_BITS;
let t2 = i2 * 2217;
let t3 = i2 * 5352;
// constants scaled things up by 1<<12, plus we had 1<<2 from first
// loop, plus horizontal and vertical each scale by sqrt(8) so together
// we've got an extra 1<<3, so 1<<17 total we need to remove.
// so we want to round that, which means adding 0.5 * 1<<17,
// aka 65536. Also, we'll end up with -128 to 127 that we want
// to encode as 0..255 by adding 128, so we'll add that before the shift
// Rounding constant is already added into `t0`
let x0 = t0 + t3;
let x3 = t0 - t3;
let x1 = t0 + t2;
let x2 = t0 - t2;
// odd part
let i3 = row3;
let i1 = row1;
let p5 = (i3 + i1) * 4816;
let p1 = p5 + i1 * -3685;
let p2 = p5 + i3 * -10497;
let t3 = p5 + i1 * 867;
let t2 = p5 + i3 * -5945;
let t1 = p2 + i1 * -1597;
let t0 = p1 + i3 * -8034;
row0.mm256 = _mm256_srai_epi32((x0 + t3).mm256, 17);
row1.mm256 = _mm256_srai_epi32((x1 + t2).mm256, 17);
row2.mm256 = _mm256_srai_epi32((x2 + t1).mm256, 17);
row3.mm256 = _mm256_srai_epi32((x3 + t0).mm256, 17);
row4.mm256 = _mm256_srai_epi32((x3 - t0).mm256, 17);
row5.mm256 = _mm256_srai_epi32((x2 - t1).mm256, 17);
row6.mm256 = _mm256_srai_epi32((x1 - t2).mm256, 17);
row7.mm256 = _mm256_srai_epi32((x0 - t3).mm256, 17);
}
transpose(
&mut row0, &mut row1, &mut row2, &mut row3, &mut row4, &mut row5, &mut row6, &mut row7,
);
let mut pos = 0;
// Pack and write the values back to the array
permute_store!((row0.mm256), (row1.mm256), pos, out_vector, stride);
permute_store!((row2.mm256), (row3.mm256), pos, out_vector, stride);
permute_store!((row4.mm256), (row5.mm256), pos, out_vector, stride);
permute_store!((row6.mm256), (row7.mm256), pos, out_vector, stride);
}
#[inline]
#[target_feature(enable = "avx2")]
unsafe fn clamp_avx(reg: __m256i) -> __m256i {
let min_s = _mm256_set1_epi16(0);
let max_s = _mm256_set1_epi16(255);
let max_v = _mm256_max_epi16(reg, min_s); //max(a,0)
let min_v = _mm256_min_epi16(max_v, max_s); //min(max(a,0),255)
return min_v;
}
/// A copy of `_MM_SHUFFLE()` that doesn't require
/// a nightly compiler
#[inline]
const fn shuffle(z: i32, y: i32, x: i32, w: i32) -> i32 {
((z << 6) | (y << 4) | (x << 2) | w)
}
+280
View File
@@ -0,0 +1,280 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
#![cfg(target_arch = "aarch64")]
//! AVX optimised IDCT.
//!
//! Okay not thaat optimised.
//!
//!
//! # The implementation
//! The implementation is neatly broken down into two operations.
//!
//! 1. Test for zeroes
//! > There is a shortcut method for idct where when all AC values are zero, we can get the answer really quickly.
//! by scaling the 1/8th of the DCT coefficient of the block to the whole block and level shifting.
//!
//! 2. If above fails, we proceed to carry out IDCT as a two pass one dimensional algorithm.
//! IT does two whole scans where it carries out IDCT on all items
//! After each successive scan, data is transposed in register(thank you x86 SIMD powers). and the second
//! pass is carried out.
//!
//! The code is not super optimized, it produces bit identical results with scalar code hence it's
//! `mm256_add_epi16`
//! and it also has the advantage of making this implementation easy to maintain.
#![cfg(feature = "neon")]
use core::arch::aarch64::*;
use crate::unsafe_utils::{transpose, YmmRegister};
const SCALE_BITS: i32 = 512 + 65536 + (128 << 17);
#[inline]
#[target_feature(enable = "neon")]
unsafe fn pack_16(a: int32x4x2_t) -> int16x8_t {
vcombine_s16(vqmovn_s32(a.0), vqmovn_s32(a.1))
}
#[inline]
#[target_feature(enable = "neon")]
unsafe fn condense_bottom_16(a: int32x4x2_t, b: int32x4x2_t) -> int16x8x2_t {
int16x8x2_t(pack_16(a), pack_16(b))
}
#[target_feature(enable = "neon")]
#[allow(
clippy::too_many_lines,
clippy::cast_possible_truncation,
clippy::similar_names,
clippy::op_ref,
unused_assignments,
clippy::zero_prefixed_literal
)]
pub unsafe fn idct_neon(
in_vector: &mut [i32; 64], out_vector: &mut [i16], stride: usize
) {
let mut pos = 0;
// load into registers
//
// We sign extend i16's to i32's and calculate them with extended precision and
// later reduce them to i16's when we are done carrying out IDCT
let mut row0 = YmmRegister::load(in_vector[00..].as_ptr().cast());
let mut row1 = YmmRegister::load(in_vector[08..].as_ptr().cast());
let mut row2 = YmmRegister::load(in_vector[16..].as_ptr().cast());
let mut row3 = YmmRegister::load(in_vector[24..].as_ptr().cast());
let mut row4 = YmmRegister::load(in_vector[32..].as_ptr().cast());
let mut row5 = YmmRegister::load(in_vector[40..].as_ptr().cast());
let mut row6 = YmmRegister::load(in_vector[48..].as_ptr().cast());
let mut row7 = YmmRegister::load(in_vector[56..].as_ptr().cast());
// Forward DCT and quantization may cause all the AC terms to be zero, for such
// cases we can try to accelerate it
// Basically the poop is that whenever the array has 63 zeroes, its idct is
// (arr[0]>>3)or (arr[0]/8) propagated to all the elements.
// We first test to see if the array contains zero elements and if it does, we go the
// short way.
//
// This reduces IDCT overhead from about 39% to 18 %, almost half
// Do another load for the first row, we don't want to check DC value, because
// we only care about AC terms
// TODO this should be a shift/shuffle, not a likely unaligned load
let row8 = YmmRegister::load(in_vector[1..].as_ptr().cast());
let or_tree = (((row1 | row8) | (row2 | row3)) | ((row4 | row5) | (row6 | row7)));
if or_tree.all_zero() {
// AC terms all zero, idct of the block is ( coeff[0] * qt[0] )/8 + 128 (bias)
// (and clamped to 255)
// Round by adding 0.5 * (1 << 3) and offset by adding (128 << 3) before scaling
let coeff = ((in_vector[0] + 4 + 1024) >> 3).clamp(0, 255) as i16;
let idct_value = vdupq_n_s16(coeff);
macro_rules! store {
($pos:tt,$value:tt) => {
// store
vst1q_s16(
out_vector
.get_mut($pos..$pos + 8)
.unwrap()
.as_mut_ptr()
.cast(),
$value
);
$pos += stride;
};
}
store!(pos, idct_value);
store!(pos, idct_value);
store!(pos, idct_value);
store!(pos, idct_value);
store!(pos, idct_value);
store!(pos, idct_value);
store!(pos, idct_value);
store!(pos, idct_value);
return;
}
macro_rules! dct_pass {
($SCALE_BITS:tt,$scale:tt) => {
// There are a lot of ways to do this
// but to keep it simple(and beautiful), ill make a direct translation of the
// scalar code to also make this code fully transparent(this version and the non
// avx one should produce identical code.)
// Compiler does a pretty good job of optimizing add + mul pairs
// into multiply-acumulate pairs
// even part
let p1 = (row2 + row6) * 2217;
let mut t2 = p1 + row6 * -7567;
let mut t3 = p1 + row2 * 3135;
let mut t0 = (row0 + row4).const_shl::<12>();
let mut t1 = (row0 - row4).const_shl::<12>();
let x0 = t0 + t3 + $SCALE_BITS;
let x3 = t0 - t3 + $SCALE_BITS;
let x1 = t1 + t2 + $SCALE_BITS;
let x2 = t1 - t2 + $SCALE_BITS;
let p3 = row7 + row3;
let p4 = row5 + row1;
let p1 = row7 + row1;
let p2 = row5 + row3;
let p5 = (p3 + p4) * 4816;
t0 = row7 * 1223;
t1 = row5 * 8410;
t2 = row3 * 12586;
t3 = row1 * 6149;
let p1 = p5 + p1 * -3685;
let p2 = p5 + (p2 * -10497);
let p3 = p3 * -8034;
let p4 = p4 * -1597;
t3 += p1 + p4;
t2 += p2 + p3;
t1 += p2 + p4;
t0 += p1 + p3;
row0 = (x0 + t3).const_shra::<$scale>();
row1 = (x1 + t2).const_shra::<$scale>();
row2 = (x2 + t1).const_shra::<$scale>();
row3 = (x3 + t0).const_shra::<$scale>();
row4 = (x3 - t0).const_shra::<$scale>();
row5 = (x2 - t1).const_shra::<$scale>();
row6 = (x1 - t2).const_shra::<$scale>();
row7 = (x0 - t3).const_shra::<$scale>();
};
}
// Process rows
dct_pass!(512, 10);
transpose(
&mut row0, &mut row1, &mut row2, &mut row3, &mut row4, &mut row5, &mut row6, &mut row7
);
// process columns
dct_pass!(SCALE_BITS, 17);
transpose(
&mut row0, &mut row1, &mut row2, &mut row3, &mut row4, &mut row5, &mut row6, &mut row7
);
// Pack i32 to i16's,
// clamp them to be between 0-255
// Undo shuffling
// Store back to array
// This could potentially be reorganized to take advantage of the multi-register stores
macro_rules! permute_store {
($x:tt,$y:tt,$index:tt,$out:tt) => {
let a = condense_bottom_16($x, $y);
// Clamp the values after packing, we can clamp more values at once
let b = clamp256_neon(a);
// store first vector
vst1q_s16(
($out)
.get_mut($index..$index + 8)
.unwrap()
.as_mut_ptr()
.cast(),
b.0
);
$index += stride;
// second vector
vst1q_s16(
($out)
.get_mut($index..$index + 8)
.unwrap()
.as_mut_ptr()
.cast(),
b.1
);
$index += stride;
};
}
// Pack and write the values back to the array
permute_store!((row0.mm256), (row1.mm256), pos, out_vector);
permute_store!((row2.mm256), (row3.mm256), pos, out_vector);
permute_store!((row4.mm256), (row5.mm256), pos, out_vector);
permute_store!((row6.mm256), (row7.mm256), pos, out_vector);
}
#[inline]
#[target_feature(enable = "neon")]
unsafe fn clamp_neon(reg: int16x8_t) -> int16x8_t {
let min_s = vdupq_n_s16(0);
let max_s = vdupq_n_s16(255);
let max_v = vmaxq_s16(reg, min_s); //max(a,0)
let min_v = vminq_s16(max_v, max_s); //min(max(a,0),255)
min_v
}
#[inline]
#[target_feature(enable = "neon")]
unsafe fn clamp256_neon(reg: int16x8x2_t) -> int16x8x2_t {
int16x8x2_t(clamp_neon(reg.0), clamp_neon(reg.1))
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_neon_clamp_256() {
unsafe {
let vals: [i16; 16] = [-1, -2, -3, 4, 256, 257, 258, 240, -1, 290, 2, 3, 4, 5, 6, 7];
let loaded = vld1q_s16_x2(vals.as_ptr().cast());
let shuffled = clamp256_neon(loaded);
let mut result: [i16; 16] = [0; 16];
vst1q_s16_x2(result.as_mut_ptr().cast(), shuffled);
assert_eq!(
result,
[0, 0, 0, 4, 255, 255, 255, 240, 0, 255, 2, 3, 4, 5, 6, 7]
)
}
}
}
+293
View File
@@ -0,0 +1,293 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
//! Platform independent IDCT algorithm
//!
//! Not as fast as AVX one.
const SCALE_BITS: i32 = 512 + 65536 + (128 << 17);
#[inline(always)]
fn wa(a: i32, b: i32) -> i32 {
a.wrapping_add(b)
}
#[inline(always)]
fn ws(a: i32, b: i32) -> i32 {
a.wrapping_sub(b)
}
#[inline(always)]
fn wm(a: i32, b: i32) -> i32 {
a.wrapping_mul(b)
}
#[inline]
pub fn idct_int_1x1(in_vector: &mut [i32; 64], mut out_vector: &mut [i16], stride: usize) {
let coeff = ((wa(wa(in_vector[0], 4), 1024) >> 3).clamp(0, 255)) as i16;
out_vector[..8].fill(coeff);
for _ in 0..7 {
out_vector = &mut out_vector[stride..];
out_vector[..8].fill(coeff);
}
}
#[allow(unused_assignments)]
#[allow(
clippy::too_many_lines,
clippy::op_ref,
clippy::cast_possible_truncation
)]
pub fn idct_int(in_vector: &mut [i32; 64], out_vector: &mut [i16], stride: usize) {
let mut pos = 0;
let mut i = 0;
if &in_vector[1..] == &[0_i32; 63] {
return idct_int_1x1(in_vector, out_vector, stride);
}
// vertical pass
for ptr in 0..8 {
let p2 = in_vector[ptr + 16];
let p3 = in_vector[ptr + 48];
let p1 = wm(wa(p2, p3), 2217);
let t2 = wa(p1, wm(p3, -7567));
let t3 = wa(p1, wm(p2, 3135));
let p2 = in_vector[ptr];
let p3 = in_vector[32 + ptr];
let t0 = fsh(wa(p2, p3));
let t1 = fsh(ws(p2, p3));
let x0 = wa(wa(t0, t3), 512);
let x3 = wa(ws(t0, t3), 512);
let x1 = wa(wa(t1, t2), 512);
let x2 = wa(ws(t1, t2), 512);
let mut t0 = in_vector[ptr + 56];
let mut t1 = in_vector[ptr + 40];
let mut t2 = in_vector[ptr + 24];
let mut t3 = in_vector[ptr + 8];
let p3 = wa(t0, t2);
let p4 = wa(t1, t3);
let p1 = wa(t0, t3);
let p2 = wa(t1, t2);
let p5 = wm(wa(p3, p4), 4816);
t0 = wm(t0, 1223);
t1 = wm(t1, 8410);
t2 = wm(t2, 12586);
t3 = wm(t3, 6149);
let p1 = wa(p5, wm(p1, -3685));
let p2 = wa(p5, wm(p2, -10497));
let p3 = wm(p3, -8034);
let p4 = wm(p4, -1597);
t3 = wa(t3, wa(p1, p4));
t2 = wa(t2, wa(p2, p3));
t1 = wa(t1, wa(p2, p4));
t0 = wa(t0, wa(p1, p3));
in_vector[ptr] = ws(wa(x0, t3), 0) >> 10;
in_vector[ptr + 8] = ws(wa(x1, t2), 0) >> 10;
in_vector[ptr + 16] = ws(wa(x2, t1), 0) >> 10;
in_vector[ptr + 24] = ws(wa(x3, t0), 0) >> 10;
in_vector[ptr + 32] = ws(ws(x3, t0), 0) >> 10;
in_vector[ptr + 40] = ws(ws(x2, t1), 0) >> 10;
in_vector[ptr + 48] = ws(ws(x1, t2), 0) >> 10;
in_vector[ptr + 56] = ws(ws(x0, t3), 0) >> 10;
}
// horizontal pass
while i < 64 {
let p2 = in_vector[i + 2];
let p3 = in_vector[i + 6];
let p1 = wm(wa(p2, p3), 2217);
let t2 = wa(p1, wm(p3, -7567));
let t3 = wa(p1, wm(p2, 3135));
let p2 = in_vector[i];
let p3 = in_vector[i + 4];
let t0 = fsh(wa(p2, p3));
let t1 = fsh(ws(p2, p3));
let x0 = wa(wa(t0, t3), SCALE_BITS);
let x3 = wa(ws(t0, t3), SCALE_BITS);
let x1 = wa(wa(t1, t2), SCALE_BITS);
let x2 = wa(ws(t1, t2), SCALE_BITS);
let mut t0 = in_vector[i + 7];
let mut t1 = in_vector[i + 5];
let mut t2 = in_vector[i + 3];
let mut t3 = in_vector[i + 1];
let p3 = wa(t0, t2);
let p4 = wa(t1, t3);
let p1 = wa(t0, t3);
let p2 = wa(t1, t2);
let p5 = wm(wa(p3, p4), f2f(1.175875602));
t0 = wm(t0, 1223);
t1 = wm(t1, 8410);
t2 = wm(t2, 12586);
t3 = wm(t3, 6149);
let p1 = wa(p5, wm(p1, -3685));
let p2 = wa(p5, wm(p2, -10497));
let p3 = wm(p3, -8034);
let p4 = wm(p4, -1597);
t3 = wa(t3, wa(p1, p4));
t2 = wa(t2, wa(p2, p3));
t1 = wa(t1, wa(p2, p4));
t0 = wa(t0, wa(p1, p3));
let out: &mut [i16; 8] = out_vector
.get_mut(pos..pos + 8)
.unwrap()
.try_into()
.unwrap();
out[0] = clamp(wa(x0, t3) >> 17);
out[1] = clamp(wa(x1, t2) >> 17);
out[2] = clamp(wa(x2, t1) >> 17);
out[3] = clamp(wa(x3, t0) >> 17);
out[4] = clamp(ws(x3, t0) >> 17);
out[5] = clamp(ws(x2, t1) >> 17);
out[6] = clamp(ws(x1, t2) >> 17);
out[7] = clamp(ws(x0, t3) >> 17);
i += 8;
pos += stride;
}
}
#[inline]
#[allow(clippy::cast_possible_truncation)]
/// Multiply a number by 4096
fn f2f(x: f32) -> i32 {
(x * 4096.0 + 0.5) as i32
}
#[inline]
/// Multiply a number by 4096
fn fsh(x: i32) -> i32 {
x << 12
}
/// Clamp values between 0 and 255
#[inline]
#[allow(clippy::cast_possible_truncation)]
fn clamp(a: i32) -> i16 {
a.clamp(0, 255) as i16
}
/// IDCT assuming only the upper 4x4 is filled.
pub fn idct4x4(in_vector: &mut [i32; 64], out_vector: &mut [i16], stride: usize) {
let mut pos = 0;
// vertical pass
for ptr in 0..4 {
let i0 = wa(fsh(in_vector[ptr]), 512);
let i2 = in_vector[ptr + 16];
let p1 = wm(i2, 2217);
let p3 = wm(i2, 5352);
let x0 = wa(i0, p3);
let x1 = wa(i0, p1);
let x2 = ws(i0, p1);
let x3 = ws(i0, p3);
// odd part
let i4 = in_vector[ptr + 24];
let i3 = in_vector[ptr + 8];
let p5 = wm(wa(i4, i3), 4816);
let p1 = wa(p5, wm(i3, -3685));
let p2 = wa(p5, wm(i4, -10497));
let t3 = wa(p5, wm(i3, 867));
let t2 = wa(p5, wm(i4, -5945));
let t1 = wa(p2, wm(i3, -1597));
let t0 = wa(p1, wm(i4, -8034));
in_vector[ptr] = wa(x0, t3) >> 10;
in_vector[ptr + 8] = wa(x1, t2) >> 10;
in_vector[ptr + 16] = wa(x2, t1) >> 10;
in_vector[ptr + 24] = wa(x3, t0) >> 10;
in_vector[ptr + 32] = ws(x3, t0) >> 10;
in_vector[ptr + 40] = ws(x2, t1) >> 10;
in_vector[ptr + 48] = ws(x1, t2) >> 10;
in_vector[ptr + 56] = ws(x0, t3) >> 10;
}
// horizontal pass
for i in (0..8).map(|i| 8 * i) {
let i2 = in_vector[i + 2];
let i0 = in_vector[i];
let t0 = wa(fsh(i0), SCALE_BITS);
let t2 = wm(i2, 2217);
let t3 = wm(i2, 5352);
let x0 = wa(t0, t3);
let x3 = ws(t0, t3);
let x1 = wa(t0, t2);
let x2 = ws(t0, t2);
// odd part
let i3 = in_vector[i + 3];
let i1 = in_vector[i + 1];
let p5 = wm(wa(i3, i1), f2f(1.175875602));
let p1 = wa(p5, wm(i1, -3685));
let p2 = wa(p5, wm(i3, -10497));
let t3 = wa(p5, wm(i1, 867));
let t2 = wa(p5, wm(i3, -5945));
let t1 = wa(p2, wm(i1, -1597));
let t0 = wa(p1, wm(i3, -8034));
let out: &mut [i16; 8] = out_vector
.get_mut(pos..pos + 8)
.unwrap()
.try_into()
.unwrap();
out.copy_from_slice(&[
clamp(wa(x0, t3) >> 17),
clamp(wa(x1, t2) >> 17),
clamp(wa(x2, t1) >> 17),
clamp(wa(x3, t0) >> 17),
clamp(ws(x3, t0) >> 17),
clamp(ws(x2, t1) >> 17),
clamp(ws(x1, t2) >> 17),
clamp(ws(x0, t3) >> 17),
]);
pos += stride;
}
in_vector[32..36].fill(0);
in_vector[40..44].fill(0);
in_vector[48..52].fill(0);
in_vector[56..60].fill(0);
}
+194
View File
@@ -0,0 +1,194 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
//!This crate provides a library for decoding valid
//! ITU-T Rec. T.851 (09/2005) ITU-T T.81 (JPEG-1) or JPEG images.
//!
//!
//!
//! # Features
//! - SSE and AVX accelerated functions to speed up certain decoding operations
//! - FAST and accurate 32 bit IDCT algorithm
//! - Fast color convert functions
//! - RGBA and RGBX (4-Channel) color conversion functions
//! - YCbCr to Luma(Grayscale) conversion.
//!
//! # Usage
//! Add zune-jpeg to the dependencies in the project Cargo.toml
//!
//! ```toml
//! [dependencies]
//! zune_jpeg = "0.5"
//! ```
//! # Examples
//!
//! ## Decode a JPEG file with default arguments.
//!```no_run
//! use std::fs::read;
//! use std::io::BufReader;
//! use zune_jpeg::JpegDecoder;
//! let file_contents = BufReader::new(std::fs::File::open("a_jpeg.file").unwrap());
//! let mut decoder = JpegDecoder::new(file_contents);
//! let mut pixels = decoder.decode().unwrap();
//! ```
//!
//! ## Migrating from version 0.4--
//!
//! ### Motivation
//! zune v 0.5 reworks mainly the internal architecture of how we perform I/O
//! ,before the decoder accepted byte slices that represent the whole data as contiguous
//! but that was not ideal for all use cases, increasing memory e.g on massive files that had
//! to be read to memory.
//!
//! With v 0.5 a new I/O system is introduced, which generally introduces mechanisms to process
//! `std::io::Read + std::io::Seek` type of data feeds, (but which works in no-std), which means...
//!
//! ### What changes
//!
//! I/O code that looked like this
//!
//!```ignore
//! use zune_core::colorspace::ColorSpace;
//! use zune_jpeg::JpegDecoder;
//! // Read file into memory
//! let image = std::fs::read("image.jpg").unwrap();
//! // Make a decoder from the slice
//! let mut decoder = JpegDecoder::new(&image);
//! // decode
//! decoder.decode().unwrap();
//! ```
//!
//! Now can be rewritten in two ways.
//!
//! 1. File I/O (Using bufreader)
//!
//!```no_run
//! use std::io::BufReader;
//! use zune_core::colorspace::ColorSpace;
//! use zune_jpeg::JpegDecoder;
//!
//! let image = BufReader::new(std::fs::File::open("image.jpg").unwrap());
//! let mut decoder = JpegDecoder::new(image);
//! // decode
//! decoder.decode().unwrap();
//! ```
//!
//! 2. Reading to memory (but wrapping it in a Cursor like object)
//!```no_run
//! use zune_core::bytestream::ZCursor;
//! use zune_jpeg::JpegDecoder;
//!
//! let image_data =std::fs::read("image.jpg").unwrap();
//! // Alternatively, you can use std::io::Cursor,
//! // but it is better speed wise to use ZCursor, and it also works in
//! // no-std environments
//! let mut cursor = ZCursor::new(image_data);
//! // use the wrapped item
//! let mut decoder = JpegDecoder::new(cursor);
//! // decode
//! decoder.decode().unwrap();
//! ```
//!
//! 3. Anything that implements [ZByteReaderTrait](zune_core::bytestream::traits::ZByteReaderTrait)
//!
//! ## Decode a JPEG file to RGBA format
//!
//! - Other (limited) supported formats are and BGR, BGRA
//!
//!```no_run
//! use zune_core::bytestream::ZCursor;
//! use zune_core::colorspace::ColorSpace;
//! use zune_core::options::DecoderOptions;
//! use zune_jpeg::JpegDecoder;
//!
//! let mut options = DecoderOptions::default().jpeg_set_out_colorspace(ColorSpace::RGBA);
//!
//! let mut decoder = JpegDecoder::new_with_options(ZCursor::new(&[]),options);
//! let pixels = decoder.decode().unwrap();
//! ```
//!
//! ## Decode an image and get its width and height.
//!```no_run
//! use zune_core::bytestream::ZCursor;
//! use zune_jpeg::JpegDecoder;
//!
//! let mut decoder = JpegDecoder::new(ZCursor::new(&[]));
//! decoder.decode_headers().unwrap();
//! let image_info = decoder.info().unwrap();
//! println!("{},{}",image_info.width,image_info.height)
//! ```
//! # Crate features.
//! This crate tries to be as minimal as possible while being extensible
//! enough to handle the complexities arising from parsing different types
//! of jpeg images.
//!
//! Safety is a top concern that is why we provide both static ways to disable unsafe code,
//! disabling x86 feature, and dynamic ,by using [`DecoderOptions::set_use_unsafe(false)`],
//! both of these disable platform specific optimizations, which reduce the speed of decompression.
//!
//! Please do note that careful consideration has been taken to ensure that the unsafe paths
//! are only unsafe because they depend on platform specific intrinsics, hence no need to disable them
//!
//! The crate tries to decode as many images as possible, as a best effort, even those violating the standard
//! , this means a lot of images may get silent warnings and wrong output, but if you are sure you will be handling
//! images that follow the spec, set `ZuneJpegOptions::set_strict` to true.
//!
//![`DecoderOptions::set_use_unsafe(false)`]: https://docs.rs/zune-core/latest/zune_core/options/struct.DecoderOptions.html#method.set_use_unsafe
#![warn(
clippy::correctness,
clippy::perf,
clippy::pedantic,
clippy::inline_always,
clippy::missing_errors_doc,
clippy::panic
)]
#![allow(
clippy::needless_return,
clippy::similar_names,
clippy::inline_always,
clippy::similar_names,
clippy::doc_markdown,
clippy::module_name_repetitions,
clippy::missing_panics_doc,
clippy::missing_errors_doc
)]
// no_std compatibility
#![deny(clippy::std_instead_of_alloc, clippy::alloc_instead_of_core)]
#![cfg_attr(not(any(feature = "x86", feature = "neon")), forbid(unsafe_code))]
#![cfg_attr(not(feature = "std"), no_std)]
#![cfg_attr(feature = "portable_simd", feature(portable_simd))]
#![macro_use]
extern crate alloc;
extern crate core;
pub use zune_core;
pub use crate::components::SampleRatios;
pub use crate::decoder::{ImageInfo, JpegDecoder};
pub use crate::marker::Marker;
mod bitstream;
mod color_convert;
mod components;
mod decoder;
pub mod errors;
mod headers;
mod huffman;
#[cfg(not(fuzzing))]
mod idct;
#[cfg(fuzzing)]
pub mod idct;
mod marker;
mod mcu;
mod mcu_prog;
mod misc;
mod unsafe_utils;
mod unsafe_utils_avx2;
mod unsafe_utils_neon;
mod upsampler;
mod worker;
+91
View File
@@ -0,0 +1,91 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
#![allow(clippy::upper_case_acronyms)]
/// JPEG Markers
///
/// **NOTE** This doesn't cover all markers, just the ones zune-jpeg supports.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Marker {
/// Start Of Frame markers
///
/// - SOF(0): Baseline DCT (Huffman coding)
/// - SOF(1): Extended sequential DCT (Huffman coding)
/// - SOF(2): Progressive DCT (Huffman coding)
/// - SOF(3): Lossless (sequential) (Huffman coding)
/// - SOF(5): Differential sequential DCT (Huffman coding)
/// - SOF(6): Differential progressive DCT (Huffman coding)
/// - SOF(7): Differential lossless (sequential) (Huffman coding)
/// - SOF(9): Extended sequential DCT (arithmetic coding)
/// - SOF(10): Progressive DCT (arithmetic coding)
/// - SOF(11): Lossless (sequential) (arithmetic coding)
/// - SOF(13): Differential sequential DCT (arithmetic coding)
/// - SOF(14): Differential progressive DCT (arithmetic coding)
/// - SOF(15): Differential lossless (sequential) (arithmetic coding)
SOF(u8),
/// Define Huffman table(s)
DHT,
/// Define arithmetic coding conditioning(s)
DAC,
/// Restart with modulo 8 count `m`
RST(u8),
/// Start of image
SOI,
/// End of image
EOI,
/// Start of scan
SOS,
/// Define quantization table(s)
DQT,
/// Define number of lines
DNL,
/// Define restart interval
DRI,
/// Reserved for application segments
APP(u8),
/// Comment
COM,
/// Unknown markers
UNKNOWN(u8)
}
impl Marker {
pub fn from_u8(n: u8) -> Option<Marker> {
use self::Marker::{APP, COM, DAC, DHT, DNL, DQT, DRI, EOI, RST, SOF, SOI, SOS, UNKNOWN};
match n {
0xFE => Some(COM),
0xC0 => Some(SOF(0)),
0xC1 => Some(SOF(1)),
0xC2 => Some(SOF(2)),
0xC4 => Some(DHT),
0xCC => Some(DAC),
0xD0 => Some(RST(0)),
0xD1 => Some(RST(1)),
0xD2 => Some(RST(2)),
0xD3 => Some(RST(3)),
0xD4 => Some(RST(4)),
0xD5 => Some(RST(5)),
0xD6 => Some(RST(6)),
0xD7 => Some(RST(7)),
0xD8 => Some(SOI),
0xD9 => Some(EOI),
0xDA => Some(SOS),
0xDB => Some(DQT),
0xDC => Some(DNL),
0xDD => Some(DRI),
0xE0 => Some(APP(0)),
0xE1 => Some(APP(1)),
0xE2 => Some(APP(2)),
0xED => Some(APP(13)),
0xEE => Some(APP(14)),
_ => Some(UNKNOWN(n))
}
}
}
+936
View File
@@ -0,0 +1,936 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
use alloc::vec::Vec;
use alloc::{format, vec};
use core::cmp::min;
use zune_core::bytestream::ZByteReaderTrait;
use zune_core::colorspace::ColorSpace;
use zune_core::colorspace::ColorSpace::Luma;
use zune_core::log::{error, trace, warn};
use crate::bitstream::BitStream;
use crate::components::SampleRatios;
use crate::decoder::MAX_COMPONENTS;
use crate::errors::DecodeErrors;
use crate::marker::Marker;
use crate::mcu_prog::get_marker;
use crate::misc::{calculate_padded_width, setup_component_params};
use crate::worker::{color_convert, upsample};
use crate::JpegDecoder;
/// The size of a DC block for a MCU.
pub const DCT_BLOCK: usize = 64;
impl<T: ZByteReaderTrait> JpegDecoder<T> {
/// Check for existence of DC and AC Huffman Tables
pub(crate) fn check_tables(&self) -> Result<(), DecodeErrors> {
// check that dc and AC tables exist outside the hot path
for component in &self.components {
let _ = &self
.dc_huffman_tables
.get(component.dc_huff_table)
.as_ref()
.ok_or_else(|| {
DecodeErrors::HuffmanDecode(format!(
"No Huffman DC table for component {:?} ",
component.component_id
))
})?
.as_ref()
.ok_or_else(|| {
DecodeErrors::HuffmanDecode(format!(
"No DC table for component {:?}",
component.component_id
))
})?;
let _ = &self
.ac_huffman_tables
.get(component.ac_huff_table)
.as_ref()
.ok_or_else(|| {
DecodeErrors::HuffmanDecode(format!(
"No Huffman AC table for component {:?} ",
component.component_id
))
})?
.as_ref()
.ok_or_else(|| {
DecodeErrors::HuffmanDecode(format!(
"No AC table for component {:?}",
component.component_id
))
})?;
}
Ok(())
}
/// Decode MCUs and carry out post processing.
///
/// This is the main decoder loop for the library, the hot path.
///
/// Because of this, we pull in some very crazy optimization tricks hence readability is a pinch
/// here.
#[allow(
clippy::similar_names,
clippy::too_many_lines,
clippy::cast_possible_truncation
)]
#[inline(never)]
pub(crate) fn decode_mcu_ycbcr_baseline(
&mut self, pixels: &mut [u8]
) -> Result<(), DecodeErrors> {
setup_component_params(self)?;
// check dc and AC tables
self.check_tables()?;
let (mut mcu_width, mut mcu_height);
if self.is_interleaved {
// set upsampling functions
self.set_upsampling()?;
mcu_width = self.mcu_x;
mcu_height = self.mcu_y;
} else {
// For non-interleaved images( (1*1) subsampling)
// number of MCU's are the widths (+7 to account for paddings) divided bu 8.
mcu_width = ((self.info.width + 7) / 8) as usize;
mcu_height = ((self.info.height + 7) / 8) as usize;
}
if self.is_interleaved
&& self.input_colorspace.num_components() > 1
&& self.options.jpeg_get_out_colorspace().num_components() == 1
&& (self.info.sample_ratio == SampleRatios::V
|| self.info.sample_ratio == SampleRatios::HV)
{
// For a specific set of images, e.g interleaved,
// when converting from YcbCr to grayscale, we need to
// take into account mcu height since the MCU decoding needs to take
// it into account for padding purposes and the post processor
// parses two rows per mcu width.
//
// set coeff to be 2 to ensure that we increment two rows
// for every mcu processed also
mcu_height *= self.v_max;
mcu_height /= self.h_max;
self.coeff = 2;
}
if self.input_colorspace == ColorSpace::Luma && self.is_interleaved {
warn!("Grayscale image with down-sampled component, resetting component details");
self.reset_params();
mcu_width = ((self.info.width + 7) / 8) as usize;
mcu_height = ((self.info.height + 7) / 8) as usize;
}
let width = usize::from(self.info.width);
let padded_width = calculate_padded_width(width, self.info.sample_ratio);
let mut stream = BitStream::new();
let mut tmp = [0_i32; DCT_BLOCK];
let comp_len = self.components.len();
for (pos, comp) in self.components.iter_mut().enumerate() {
// Allocate only needed components.
//
// For special colorspaces i.e YCCK and CMYK, just allocate all of the needed
// components.
if min(
self.options.jpeg_get_out_colorspace().num_components() - 1,
pos
) == pos
|| comp_len == 4
// Special colorspace
{
// allocate enough space to hold a whole MCU width
// this means we should take into account sampling ratios
// `*8` is because each MCU spans 8 widths.
let len = comp.width_stride * comp.vertical_sample * 8;
comp.needed = true;
comp.raw_coeff = vec![0; len];
} else {
comp.needed = false;
}
}
// If all components are contained in the first scan of MCUs, then we can process into
// (upsampled) pixels immediately after each MCU, for convenience we use each row of MCUS.
// Otherwise, we must first wait until following SOS provide the remaining components.
let all_components_in_first_scan = usize::from(self.num_scans) == self.components.len();
let mut progressive_mcus: [Vec<i16>; 4] = core::array::from_fn(|_| vec![]);
if !all_components_in_first_scan {
for (component, mcu) in self.components.iter().zip(&mut progressive_mcus) {
let len = mcu_width
* component.vertical_sample
* component.horizontal_sample
* mcu_height
* 64;
*mcu = vec![0; len];
}
}
let mut pixels_written = 0;
let is_hv = usize::from(self.is_interleaved);
let upsampler_scratch_size = is_hv * self.components.iter().map(|x| x.width_stride).max().unwrap_or(0) * 8;
let mut upsampler_scratch_space = vec![0; upsampler_scratch_size];
'sos: loop {
trace!(
"Baseline decoding of components: {:?}",
&self.z_order[..usize::from(self.num_scans)]
);
trace!("Decoding MCU width: {mcu_width}, height: {mcu_height}");
for i in 0..mcu_height {
if stream.overread_by > 0 {
pixels.get_mut(pixels_written..).map(|v| v.fill(128));
if self.options.strict_mode() {
return Err(DecodeErrors::FormatStatic("Premature end of buffer"));
};
error!("Premature end of buffer");
break;
}
// decode a whole MCU width,
// this takes into account interleaved components.
let terminate = if all_components_in_first_scan {
self.decode_mcu_width::<false>(
mcu_width,
i,
&mut tmp,
&mut stream,
&mut progressive_mcus
)?
} else {
/* NB: (cae). This code was added due to the issue at https://github.com/etemesi254/zune-image/issues/277
*
* There is a particular set of images that interleave the start of scan (SOS) with the MCU,
* E.g if it's a three component image, we have SOS->MCU ->SOS->MCU ->SOS->MCU
* which presents a problem on decoding, we need to buffer the whole image before continuing since
* we won't have a row containing all the component data which will be needed e.g for color conversion.
*
* The mechanisms is that we decode the whole image upfront, which goes against the normal
* routine of decoding MCU width , so this requires more memory upfront than initial routines
* but it is a single image out of the many corpuses that exist, so its fine.
* (image in test-images/jpeg/sos_news.jpeg)
* Code contributed by Aurelia Molzer (https://github.com/197g)
*
*/
self.decode_mcu_width::<true>(
mcu_width,
i,
&mut tmp,
&mut stream,
&mut progressive_mcus
)?
};
// process that width up until it's impossible. This is faster than allocation the
// full components, which we skipped earlier.
if all_components_in_first_scan {
self.post_process(
pixels,
i,
mcu_height,
width,
padded_width,
&mut pixels_written,
&mut upsampler_scratch_space
)?;
}
match terminate {
McuContinuation::Ok => {}
McuContinuation::AnotherSos if all_components_in_first_scan => {
warn!("More than one SOS despite already having all components");
return Ok(());
}
McuContinuation::AnotherSos => continue 'sos,
McuContinuation::InterScanMarker(marker) => {
// Handle inter-scan markers (DHT/DQT/etc) uniformly here.
// This keeps all marker handling in the outer loop.
if self.advance_to_next_sos(marker, &mut stream)? {
continue 'sos;
} else {
// Hit EOI
break;
}
}
McuContinuation::Terminate => {
warn!("Got terminate signal, will not process further");
pixels.get_mut(pixels_written..).map(|v| v.fill(128));
return Ok(());
}
}
}
// Breaks if we get here, looping only if we have restarted, i.e. found another SOS and
// continued at `'sos'.
break;
}
if !all_components_in_first_scan {
self.finish_baseline_decoding(&progressive_mcus, mcu_width, pixels)?;
}
// it may happen that some images don't have the whole buffer
// so we can't panic in case of that
// assert_eq!(pixels_written, pixels.len());
// For UHD usecases that tie two images separating them with EOI and
// SOI markers, it may happen that we do not reach this image end of image
// So this ensures we reach it
// Ensure we read EOI
if !stream.seen_eoi {
let marker = get_marker(&mut self.stream, &mut stream);
match marker {
Ok(_m) => {
trace!("Found marker {:?}", _m);
}
Err(_) => {
// ignore error
}
}
}
trace!("Finished decoding image");
Ok(())
}
/// Process all MCUs when baseline decoding has been processing them component-after-component.
/// For simplicity this assembles the dequantized blocks in the order that the post processing
/// of an interleaved baseline decoding would use.
#[allow(clippy::too_many_lines)]
#[allow(clippy::cast_sign_loss)]
pub(crate) fn finish_baseline_decoding(
&mut self, block: &[Vec<i16>; MAX_COMPONENTS], _mcu_width: usize, pixels: &mut [u8]
) -> Result<(), DecodeErrors> {
let mcu_height = self.mcu_y;
// Size of our output image(width*height)
let is_hv = usize::from(self.is_interleaved);
let upsampler_scratch_size = is_hv * self.components[0].width_stride;
let width = usize::from(self.info.width);
let padded_width = calculate_padded_width(width, self.info.sample_ratio);
let mut upsampler_scratch_space = vec![0; upsampler_scratch_size];
for (pos, comp) in self.components.iter_mut().enumerate() {
// Mark only needed components for computing output colors.
if min(
self.options.jpeg_get_out_colorspace().num_components() - 1,
pos
) == pos
|| self.input_colorspace == ColorSpace::YCCK
|| self.input_colorspace == ColorSpace::CMYK
{
comp.needed = true;
} else {
comp.needed = false;
}
}
let mut pixels_written = 0;
// dequantize and idct have been performed, only color convert.
for i in 0..mcu_height {
// All the data is already in the right order, we just need to be able to pass it to
// the post_process & upsample method. That expects all the data to be stored as one
// row of MCUs in each component's `raw_coeff`.
'component: for (position, component) in &mut self.components.iter_mut().enumerate() {
if !component.needed {
continue 'component;
}
// step is the number of pixels this iteration wil be handling
// Given by the number of mcu's height and the length of the component block
// Since the component block contains the whole channel as raw pixels
// we this evenly divides the pixels into MCU blocks
//
// For interleaved images, this gives us the exact pixels comprising a whole MCU
// block
let step = block[position].len() / mcu_height;
// where we will be reading our pixels from.
let slice = &block[position][i * step..][..step];
let temp_channel = &mut component.raw_coeff;
temp_channel[..step].copy_from_slice(slice);
}
// process that whole stripe of MCUs
self.post_process(
pixels,
i,
mcu_height,
width,
padded_width,
&mut pixels_written,
&mut upsampler_scratch_space
)?;
}
return Ok(());
}
fn decode_mcu_width<const PROGRESSIVE: bool>(
&mut self, mcu_width: usize, mcu_height: usize, tmp: &mut [i32; 64],
stream: &mut BitStream, progressive: &mut [Vec<i16>; 4]
) -> Result<McuContinuation, DecodeErrors> {
let is_one_by_one = !self.scan_subsampled;
// The definition of MCU depends on the sampling factor of involved scans. When components
// have different factors then each Minimal-Coding-Unit is the least common multiple such
// that we have an integer number of blocks from each component. But the decoding of these
// components differs from it otherwise, we need an inner loop with a dynamic amount of
// coefficients per component, whereas otherwise we have exactly one block of coefficients
// encoded for each component in the bitstream order.
//
// We statically specialize on this to improve code generation of the common case a little
// bit. We could also special case common sub-sampling cases but be mindful of code bloat.
if is_one_by_one {
self.inner_decode_mcu_width::<PROGRESSIVE, false>(
mcu_width,
mcu_height,
tmp,
stream,
progressive
)
} else {
self.inner_decode_mcu_width::<PROGRESSIVE, true>(
mcu_width,
mcu_height,
tmp,
stream,
progressive
)
}
}
// Inline-never ensures we do get this function optimize on its own, into two different
// versions, without the optimizer tripping up over the complexity that comes with the
// constant folding. And constant folding is quite important for performance here as
// when `not SAMPLED` then the inner loop has exactly one iteration per component in
// the scan. The difference was ~1% or a bit more.
fn inner_decode_mcu_width<const PROGRESSIVE: bool, const SAMPLED: bool>(
&mut self, mcu_width: usize, mcu_height: usize, tmp: &mut [i32; 64],
stream: &mut BitStream, progressive: &mut [Vec<i16>; 4]
) -> Result<McuContinuation, DecodeErrors> {
let z_order = self.z_order;
let z_scans = &z_order[..usize::from(self.num_scans)];
// How much of the head of `tmp` was written by the last MCU decoding? We only check for
// two different cases and not all possible outcomes as this is only used to optimize the
// bytes written in `fill`. Since the clobber happens in UNZIGZAG order we'd be straddling
// most cache lines anyways even if we did a partial write with the exact length of the
// coefficient data which was written into `tmp`.
let mut clobber_more_than_4x4 = true;
// For non-interleaved scans (PROGRESSIVE=true), each scan contains a single component
// and we iterate over that component's actual data unit count, not the interleaved MCU
// width multiplied by sampling factor.
let scan_du_width = if PROGRESSIVE {
let k = z_scans[0];
let comp = &self.components[k];
// Calculate actual data units for this component: ceil(width / (8 * subsampling_ratio))
(self.info.width as usize * comp.horizontal_sample + self.h_max * 8 - 1)
/ (self.h_max * 8)
} else {
mcu_width
};
for j in 0..scan_du_width {
// iterate over components
for &k in z_scans {
// we made this loop body massive due to several different paths that depend on
// static conditions. Note we (potentially) call into other functions so the
// compiler will not unroll anything here anyways. The gains from separating
// differently optimized loop bodies are much greater than a single additional jump
// here.
let component = &mut self.components[k];
let dc_table = self.dc_huffman_tables[component.dc_huff_table % MAX_COMPONENTS]
.as_ref()
.ok_or(DecodeErrors::FormatStatic("DC table not found"))?;
let ac_table = self.ac_huffman_tables[component.ac_huff_table % MAX_COMPONENTS]
.as_ref()
.ok_or(DecodeErrors::FormatStatic("AC table not found"))?;
let qt_table = &component.quantization_table;
let channel = if PROGRESSIVE {
let offset =
mcu_height * component.width_stride * 8 * component.vertical_sample;
&mut progressive[k][offset..]
} else {
&mut component.raw_coeff
};
let component_samples_needed = component.needed;
// If image is interleaved iterate over scan components,
// otherwise if it-s non-interleaved, these routines iterate in
// trivial scanline order(Y,Cb,Cr)
//
// Turn the bounds into a compile time constant for a common special case. This
// allows the compiler to unroll the loop and then do a bunch of interleaving.
//
// For PROGRESSIVE (non-interleaved), we iterate data units directly so
// h_samp/v_samp loops run exactly once.
let v_step =
if SAMPLED && !PROGRESSIVE { 0..component.vertical_sample } else { 0..1 };
for v_samp in v_step {
let h_step =
if SAMPLED && !PROGRESSIVE { 0..component.horizontal_sample } else { 0..1 };
for h_samp in h_step {
let result = if component_samples_needed {
// Fill the array with zeroes, decode_mcu_block expects
// a zero based array. Clobber is in zig-zag order though.
// Writing consecutive entries is basically free in terms
// of memory throughput so we opt for a larger power of
// two which lets the compiler turn this into a repeated
// write of a zeroed vector register, which does not have
// any branches, instead of a more difficult pattern where
// we attempt to overwrite exactly one coefficient.
let clobber_len = if !clobber_more_than_4x4 { 32 } else { 64 };
tmp[..clobber_len].fill(0);
stream.decode_mcu_block(
&mut self.stream,
dc_table,
ac_table,
qt_table,
tmp,
&mut component.dc_pred
)
} else {
// We do not touch tmp so there is no need to reset it.
stream.discard_mcu_block(&mut self.stream, dc_table, ac_table)
};
// If an error occurs we can either propagate it
// as an error or print it and call terminate.
//
// This allows even corrupt images to render something,
// even if its bad, matching browsers.
//
// See example in https://github.com/etemesi254/zune-image/issues/293
let len = if let Ok(len) = result {
len
} else {
// result.is_err()
return if self.options.strict_mode() {
Err(result.err().unwrap())
} else {
error!("{}", result.err().unwrap());
Ok(McuContinuation::Terminate)
};
};
if component_samples_needed {
// tmp was only written partially, note that len is in ZigZag order.
clobber_more_than_4x4 = len > 10;
let idct_position = if PROGRESSIVE {
// For non-interleaved, j indexes data units directly
j * 8
} else {
// derived from stb and rewritten for my tastes
let c2 = v_samp * 8;
let c3 = ((j * component.horizontal_sample) + h_samp) * 8;
component.width_stride * c2 + c3
};
let idct_pos = channel.get_mut(idct_position..).unwrap();
if len <= 1 {
(self.idct_1x1_func)(tmp, idct_pos, component.width_stride);
} else if len <= 10 {
(self.idct_4x4_func)(tmp, idct_pos, component.width_stride);
} else {
// call idct.
(self.idct_func)(tmp, idct_pos, component.width_stride);
}
}
}
}
}
self.todo = self.todo.wrapping_sub(1);
if self.todo == 0 {
self.handle_rst_main(stream)?;
continue;
}
if stream.marker.is_some() && stream.bits_left == 0 {
break;
}
}
self.check_stream_marker_after_mcu_width(stream)
}
fn check_stream_marker_after_mcu_width(
&mut self, stream: &mut BitStream
) -> Result<McuContinuation, DecodeErrors> {
// After all interleaved components, that's an MCU
// handle stream markers
//
// In some corrupt images, it may occur that header markers occur in the stream.
// The spec EXPLICITLY FORBIDS this, specifically, in
// routine F.2.2.5 it says
// `The only valid marker which may occur within the Huffman coded data is the RSTm marker.`
//
// But libjpeg-turbo allows it because of some weird reason. so I'll also
// allow it because of some weird reason.
if let Some(m) = stream.marker {
if m == Marker::EOI {
// acknowledge and ignore EOI marker.
stream.marker.take();
trace!("Found EOI marker");
// Google Introduced the Ultra-HD image format which is basically
// stitching two images into one container.
// They basically separate two images via a EOI and SOI marker
// so let's just ensure if we ever see EOI, we never read past that
// ever.
// https://github.com/google/libultrahdr
stream.seen_eoi = true;
} else if let Marker::RST(_) = m {
//debug_assert_eq!(self.todo, 0);
if self.todo == 0 {
self.handle_rst(stream)?;
}
} else if let Marker::SOS = m {
self.parse_marker_inner(m)?;
stream.marker.take();
stream.reset();
trace!("Found SOS marker");
return Ok(McuContinuation::AnotherSos);
} else if matches!(m, Marker::DHT | Marker::DQT | Marker::DRI | Marker::COM)
|| matches!(m, Marker::APP(_))
{
// For non-interleaved images, setup markers can appear between scans.
// Signal the caller to handle this marker and find the next SOS.
// This keeps all marker parsing in the caller's loop.
stream.marker.take();
trace!("Found inter-scan marker {:?}", m);
return Ok(McuContinuation::InterScanMarker(m));
} else {
if self.options.strict_mode() {
return Err(DecodeErrors::Format(format!(
"Marker {m:?} found where not expected"
)));
}
error!(
"Marker `{:?}` Found within Huffman Stream, possibly corrupt jpeg",
m
);
self.parse_marker_inner(m)?;
stream.marker.take();
stream.reset();
return Ok(McuContinuation::Terminate);
}
}
Ok(McuContinuation::Ok)
}
/// Scan for the next SOS marker, parsing setup markers along the way.
///
/// This is the unified marker scanning function used after encountering an
/// inter-scan marker. It handles DHT, DQT, DRI, COM, and APP markers that
/// can appear between scans in non-interleaved images.
///
/// # Arguments
/// * `first_marker` - The first marker that was already detected (not yet parsed)
/// * `stream` - The bitstream state
///
/// # Returns
/// * `Ok(true)` - Found SOS, ready to continue decoding
/// * `Ok(false)` - Found EOI, decoding complete
/// * `Err(_)` - Error (too many markers, unexpected marker in strict mode, etc.)
fn advance_to_next_sos(
&mut self,
first_marker: Marker,
stream: &mut BitStream
) -> Result<bool, DecodeErrors> {
// Limit iterations to prevent DoS from malicious files.
const MAX_INTER_SCAN_MARKERS: usize = 64;
// Parse the first marker that triggered this call
self.parse_marker_inner(first_marker)?;
stream.reset();
for _ in 0..MAX_INTER_SCAN_MARKERS {
let marker = get_marker(&mut self.stream, stream)?;
match marker {
Marker::SOS => {
self.parse_marker_inner(Marker::SOS)?;
stream.reset();
trace!("Found SOS marker, continuing decode");
return Ok(true);
}
Marker::EOI => {
stream.seen_eoi = true;
trace!("Found EOI marker");
return Ok(false);
}
Marker::DHT | Marker::DQT | Marker::DRI | Marker::COM => {
trace!("Parsing inter-scan marker {:?}", marker);
self.parse_marker_inner(marker)?;
}
Marker::APP(_) => {
trace!("Parsing inter-scan APP marker {:?}", marker);
self.parse_marker_inner(marker)?;
}
other => {
if self.options.strict_mode() {
return Err(DecodeErrors::Format(format!(
"Unexpected marker {:?} while scanning for SOS between scans",
other
)));
}
// Non-strict: skip unknown marker
warn!("Skipping unexpected marker {:?} between scans", other);
let length = self.stream.get_u16_be_err()?;
if length >= 2 {
self.stream.skip((length - 2) as usize)?;
}
}
}
}
Err(DecodeErrors::FormatStatic(
"Too many markers between scans (exceeded limit of 64)"
))
}
// handle RST markers.
// No-op if not using restarts
// this routine is shared with mcu_prog
#[cold]
pub(crate) fn handle_rst(&mut self, stream: &mut BitStream) -> Result<(), DecodeErrors> {
self.todo = self.restart_interval;
if let Some(marker) = stream.marker {
// Found a marker
// Read stream and see what marker is stored there
match marker {
Marker::RST(_) => {
// reset stream
stream.reset();
// Initialize dc predictions to zero for all components
self.components.iter_mut().for_each(|x| x.dc_pred = 0);
// Start iterating again. from position.
}
Marker::EOI => {
// silent pass
}
_ => {
return Err(DecodeErrors::MCUError(format!(
"Marker {marker:?} found in bitstream, possibly corrupt jpeg"
)));
}
}
}
Ok(())
}
#[allow(clippy::too_many_lines, clippy::too_many_arguments)]
pub(crate) fn post_process(
&mut self, pixels: &mut [u8], i: usize, mcu_height: usize, width: usize,
padded_width: usize, pixels_written: &mut usize, upsampler_scratch_space: &mut [i16]
) -> Result<(), DecodeErrors> {
let out_colorspace_components = self.options.jpeg_get_out_colorspace().num_components();
let mut px = *pixels_written;
// indicates whether image is vertically up-sampled
let is_vertically_sampled = self
.components
.iter()
.any(|c| c.sample_ratio == SampleRatios::HV || c.sample_ratio == SampleRatios::V);
let mut comp_len = self.components.len();
// If we are moving from YCbCr -> Luma, we do not allocate storage for other components, so we
// will panic when we are trying to read samples, so for that case,
// hardcode it so that we don't panic when doing
// *samp = &samples[j][pos * padded_width..(pos + 1) * padded_width]
if out_colorspace_components < comp_len && self.options.jpeg_get_out_colorspace() == Luma {
comp_len = out_colorspace_components;
}
let mut color_conv_function =
|num_iters: usize, samples: [&[i16]; 4]| -> Result<(), DecodeErrors> {
for (pos, output) in pixels[px..]
.chunks_exact_mut(width * out_colorspace_components)
.take(num_iters)
.enumerate()
{
let mut raw_samples: [&[i16]; 4] = [&[], &[], &[], &[]];
// iterate over each line, since color-convert needs only
// one line
for (j, samp) in raw_samples.iter_mut().enumerate().take(comp_len) {
let temp = &samples[j].get(pos * padded_width..(pos + 1) * padded_width);
if temp.is_none() {
return Err(DecodeErrors::FormatStatic("Missing samples"));
}
*samp = temp.unwrap();
}
color_convert(
&raw_samples,
self.color_convert_16,
self.input_colorspace,
self.options.jpeg_get_out_colorspace(),
output,
width,
padded_width
)?;
px += width * out_colorspace_components;
}
Ok(())
};
let comps = &mut self.components[..];
if self.is_interleaved && self.options.jpeg_get_out_colorspace() != ColorSpace::Luma {
for comp in comps.iter_mut() {
upsample(
comp,
mcu_height,
i,
upsampler_scratch_space,
is_vertically_sampled
)?;
}
if is_vertically_sampled {
if i > 0 {
// write the last line, it wasn't up-sampled as we didn't have row_down
// yet
let mut samples: [&[i16]; 4] = [&[], &[], &[], &[]];
for (samp, component) in samples.iter_mut().zip(comps.iter()) {
*samp = &component.first_row_upsample_dest;
}
// ensure length matches for all samples
let _first_len = samples[0].len();
// This was a good check, but can be caused to panic, esp on invalid/corrupt images.
// See one in issue https://github.com/etemesi254/zune-image/issues/262, so for now
// we just ignore and generate invalid images at the end.
//
//
// for samp in samples.iter().take(comp_len) {
// assert_eq!(first_len, samp.len());
// }
let num_iters = self.coeff * self.v_max;
color_conv_function(num_iters, samples)?;
}
// After up-sampling the last row, save any row that can be used for
// a later up-sampling,
//
// E.g the Y sample is not sampled but we haven't finished upsampling the last row of
// the previous mcu, since we don't have the down row, so save it
for component in comps.iter_mut() {
if component.sample_ratio != SampleRatios::H {
// We don't care about H sampling factors, since it's copied in the workers function
// copy last row to be used for the next color conversion
let size = component.vertical_sample
* component.width_stride
* component.sample_ratio.sample();
let last_bytes =
component.raw_coeff.rchunks_exact_mut(size).next().unwrap();
component
.first_row_upsample_dest
.copy_from_slice(last_bytes);
}
}
}
let mut samples: [&[i16]; 4] = [&[], &[], &[], &[]];
for (samp, component) in samples.iter_mut().zip(comps.iter()) {
*samp = if component.sample_ratio == SampleRatios::None {
&component.raw_coeff
} else {
&component.upsample_dest
};
}
// we either do 7 or 8 MCU's depending on the state, this only applies to
// vertically sampled images
//
// for rows up until the last MCU, we do not upsample the last stride of the MCU
// which means that the number of iterations should take that into account is one less the
// up-sampled size
//
// For the last MCU, we upsample the last stride, meaning that if we hit the last MCU, we
// should sample full raw coeffs
let is_last_considered = is_vertically_sampled && (i != mcu_height.saturating_sub(1));
let num_iters = (8 - usize::from(is_last_considered)) * self.coeff * self.v_max;
color_conv_function(num_iters, samples)?;
} else {
let mut channels_ref: [&[i16]; MAX_COMPONENTS] = [&[]; MAX_COMPONENTS];
self.components
.iter()
.enumerate()
.for_each(|(pos, x)| channels_ref[pos] = &x.raw_coeff);
if let SampleRatios::Generic(_, v) = self.info.sample_ratio {
color_conv_function(8 * v * self.coeff, channels_ref)?;
} else {
color_conv_function(8 * self.coeff, channels_ref)?;
}
}
*pixels_written = px;
Ok(())
}
}
enum McuContinuation {
Ok,
AnotherSos,
/// Found an inter-scan marker (DHT/DQT/DRI/COM/APP) that needs handling.
/// The caller should parse it and scan for the next SOS.
InterScanMarker(Marker),
Terminate
}
+688
View File
@@ -0,0 +1,688 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
//!Routines for progressive decoding
/*
This file is needlessly complicated,
It is that way to ensure we don't burn memory anyhow
Memory is a scarce resource in some environments, I would like this to be viable
in such environments
Half of the complexity comes from the jpeg spec, because progressive decoding,
is one hell of a ride.
*/
use alloc::string::ToString;
use alloc::vec::Vec;
use alloc::{format, vec};
use core::cmp::min;
use zune_core::bytestream::{ZByteReaderTrait, ZReader};
use zune_core::colorspace::ColorSpace;
use zune_core::log::{debug, error, warn};
use crate::bitstream::BitStream;
use crate::components::SampleRatios;
use crate::decoder::{JpegDecoder, MAX_COMPONENTS};
use crate::errors::DecodeErrors;
use crate::headers::parse_sos;
use crate::marker::Marker;
use crate::mcu::DCT_BLOCK;
use crate::misc::{calculate_padded_width, setup_component_params};
impl<T: ZByteReaderTrait> JpegDecoder<T> {
/// Decode a progressive image
///
/// This routine decodes a progressive image, stopping if it finds any error.
#[allow(
clippy::needless_range_loop,
clippy::cast_sign_loss,
clippy::redundant_else,
clippy::too_many_lines
)]
#[inline(never)]
pub(crate) fn decode_mcu_ycbcr_progressive(
&mut self, pixels: &mut [u8]
) -> Result<(), DecodeErrors> {
setup_component_params(self)?;
let mut mcu_height;
// memory location for decoded pixels for components
let mut block: [Vec<i16>; MAX_COMPONENTS] = [vec![], vec![], vec![], vec![]];
let mut mcu_width;
let mut seen_scans = 1;
if self.input_colorspace == ColorSpace::Luma && self.is_interleaved {
warn!("Grayscale image with down-sampled component, resetting component details");
self.reset_params();
}
if self.is_interleaved {
// this helps us catch component errors.
self.set_upsampling()?;
}
if self.is_interleaved {
mcu_width = self.mcu_x;
mcu_height = self.mcu_y;
} else {
mcu_width = (self.info.width as usize + 7) / 8;
mcu_height = (self.info.height as usize + 7) / 8;
}
if self.is_interleaved
&& self.input_colorspace.num_components() > 1
&& self.options.jpeg_get_out_colorspace().num_components() == 1
&& (self.info.sample_ratio == SampleRatios::V
|| self.info.sample_ratio == SampleRatios::HV)
{
// For a specific set of images, e.g interleaved,
// when converting from YcbCr to grayscale, we need to
// take into account mcu height since the MCU decoding needs to take
// it into account for padding purposes and the post processor
// parses two rows per mcu width.
//
// set coeff to be 2 to ensure that we increment two rows
// for every mcu processed also
mcu_height *= self.v_max;
mcu_height /= self.h_max;
self.coeff = 2;
}
mcu_width *= 64;
for i in 0..self.input_colorspace.num_components() {
let comp = &self.components[i];
let len = mcu_width * comp.vertical_sample * comp.horizontal_sample * mcu_height;
block[i] = vec![0; len];
}
let mut stream = BitStream::new_progressive(self.succ_low, self.spec_start, self.spec_end);
// there are multiple scans in the stream, this should resolve the first scan
let result = self.parse_entropy_coded_data(&mut stream, &mut block);
if result.is_err() {
return if self.options.strict_mode() {
Err(result.err().unwrap())
} else {
error!("{}", result.err().unwrap());
// Go process it and return as much as we can, exiting here
return self.finish_progressive_decoding(&block, pixels);
};
}
// extract marker
let mut marker = stream
.marker
.take()
.ok_or(DecodeErrors::FormatStatic("Marker missing where expected"))?;
// if marker is EOI, we are done, otherwise continue scanning.
//
// In case we have a premature image, we print a warning or return
// an error, depending on the strictness of the decoder, so there
// is that logic to handle too
'eoi: while marker != Marker::EOI {
match marker {
Marker::SOS => {
parse_sos(self)?;
stream.update_progressive_params(
self.succ_high,
self.succ_low,
self.spec_start,
self.spec_end
);
// after every SOS, marker, parse data for that scan.
let result = self.parse_entropy_coded_data(&mut stream, &mut block);
// Do not error out too fast, allows the decoder to continue as much as possible
// even after errors
if result.is_err() {
return if self.options.strict_mode() {
Err(result.err().unwrap())
} else {
error!("{}", result.err().unwrap());
break 'eoi;
};
}
// extract marker, might either indicate end of image or we continue
// scanning(hence the continue statement to determine).
match get_marker(&mut self.stream, &mut stream) {
Ok(marker_n) => {
marker = marker_n;
seen_scans += 1;
if seen_scans > self.options.jpeg_get_max_scans() {
return Err(DecodeErrors::Format(format!(
"Too many scans, exceeded limit of {}",
self.options.jpeg_get_max_scans()
)));
}
stream.reset();
continue 'eoi;
}
Err(msg) => {
if self.options.strict_mode() {
return Err(msg);
}
error!("{:?}", msg);
break 'eoi;
}
}
}
Marker::RST(_n) => {
self.handle_rst(&mut stream)?;
}
_ => {
self.parse_marker_inner(marker)?;
}
}
match get_marker(&mut self.stream, &mut stream) {
Ok(marker_n) => {
marker = marker_n;
}
Err(e) => {
if self.options.strict_mode() {
return Err(e);
}
error!("{}", e);
// If we can't get the marker, just break away
// allows us to decode some corrupt images
// e.g https://github.com/etemesi254/zune-image/issues/294
break 'eoi;
}
}
}
self.finish_progressive_decoding(&block, pixels)
}
/// Reset progressive parameters
fn reset_prog_params(&mut self, stream: &mut BitStream) {
stream.reset();
self.components.iter_mut().for_each(|x| x.dc_pred = 0);
// Also reset JPEG restart intervals
self.todo = if self.restart_interval != 0 { self.restart_interval } else { usize::MAX };
}
#[allow(clippy::too_many_lines, clippy::cast_sign_loss)]
fn parse_entropy_coded_data(
&mut self, stream: &mut BitStream, buffer: &mut [Vec<i16>; MAX_COMPONENTS]
) -> Result<(), DecodeErrors> {
self.reset_prog_params(stream);
if usize::from(self.num_scans) > self.input_colorspace.num_components() {
return Err(DecodeErrors::Format(format!(
"Number of scans {} cannot be greater than number of components, {}",
self.num_scans,
self.input_colorspace.num_components()
)));
}
if self.num_scans == 1 {
// Safety checks
if self.spec_end != 0 && self.spec_start == 0 {
return Err(DecodeErrors::FormatStatic(
"Can't merge DC and AC corrupt jpeg"
));
}
// non interleaved data, process one block at a time in trivial scanline order
let k = self.z_order[0];
if k >= self.components.len() {
return Err(DecodeErrors::Format(format!(
"Cannot find component {k}, corrupt image"
)));
}
// For non-interleaved scans, iterate over the component's actual data-unit grid.
let component = &self.components[k];
let mcu_width = (self.info.width as usize * component.horizontal_sample).div_ceil(self.h_max * 8);
let mcu_height = (self.info.height as usize * component.vertical_sample).div_ceil(self.v_max * 8);
for i in 0..mcu_height {
for j in 0..mcu_width {
if self.spec_start != 0 && self.succ_high == 0 && stream.eob_run > 0 {
// handle EOB runs here.
stream.eob_run -= 1;
} else {
let start = 64 * (j + i * (self.components[k].width_stride / 8));
let data: &mut [i16; 64] = buffer
.get_mut(k)
.unwrap()
.get_mut(start..start + 64)
.ok_or(DecodeErrors::FormatStatic("Slice to Small"))?
.try_into()
.unwrap();
if self.spec_start == 0 {
let pos = self.components[k].dc_huff_table & (MAX_COMPONENTS - 1);
let dc_table = self
.dc_huffman_tables
.get(pos)
.ok_or(DecodeErrors::FormatStatic(
"No huffman table for DC component"
))?
.as_ref()
.ok_or(DecodeErrors::FormatStatic(
"Huffman table at index {} not initialized"
))?;
let dc_pred = &mut self.components[k].dc_pred;
if self.succ_high == 0 {
// first scan for this mcu
stream.decode_prog_dc_first(
&mut self.stream,
dc_table,
&mut data[0],
dc_pred
)?;
} else {
// refining scans for this MCU
stream.decode_prog_dc_refine(&mut self.stream, &mut data[0])?;
}
} else {
let pos = self.components[k].ac_huff_table;
let ac_table = self
.ac_huffman_tables
.get(pos)
.ok_or_else(|| {
DecodeErrors::Format(format!(
"No huffman table for component:{pos}"
))
})?
.as_ref()
.ok_or_else(|| {
DecodeErrors::Format(format!(
"Huffman table at index {pos} not initialized"
))
})?;
if self.succ_high == 0 {
debug_assert!(stream.eob_run == 0, "EOB run is not zero");
stream.decode_mcu_ac_first(&mut self.stream, ac_table, data)?;
} else {
// refinement scan
stream.decode_mcu_ac_refine(&mut self.stream, ac_table, data)?;
}
// Check for a marker.
// It can appear in stream CC https://github.com/etemesi254/zune-image/issues/300
// if let Some(marker) = stream.marker.take() {
// self.parse_marker_inner(marker)?;
// }
}
}
// + EOB and investigate effect.
self.todo -= 1;
self.handle_rst_main(stream)?;
}
}
} else {
if self.spec_end != 0 {
return Err(DecodeErrors::HuffmanDecode(
"Can't merge dc and AC corrupt jpeg".to_string()
));
}
// process scan n elements in order
// Do the error checking with allocs here.
// Make the one in the inner loop free of allocations.
for k in 0..self.num_scans {
let n = self.z_order[k as usize];
if n >= self.components.len() {
return Err(DecodeErrors::Format(format!(
"Cannot find component {n}, corrupt image"
)));
}
let component = &mut self.components[n];
let _ = self
.dc_huffman_tables
.get(component.dc_huff_table)
.ok_or_else(|| {
DecodeErrors::Format(format!(
"No huffman table for component:{}",
component.dc_huff_table
))
})?
.as_ref()
.ok_or_else(|| {
DecodeErrors::Format(format!(
"Huffman table at index {} not initialized",
component.dc_huff_table
))
})?;
}
// Interleaved scan
// Components shall not be interleaved in progressive mode, except for
// the DC coefficients in the first scan for each component of a progressive frame.
for i in 0..self.mcu_y {
for j in 0..self.mcu_x {
// process scan n elements in order
for k in 0..self.num_scans {
let n = self.z_order[k as usize];
let component = &mut self.components[n];
let huff_table = self
.dc_huffman_tables
.get(component.dc_huff_table)
.ok_or(DecodeErrors::FormatStatic("No huffman table for component"))?
.as_ref()
.ok_or(DecodeErrors::FormatStatic(
"Huffman table at index not initialized"
))?;
for v_samp in 0..component.vertical_sample {
for h_samp in 0..component.horizontal_sample {
let x2 = j * component.horizontal_sample + h_samp;
let y2 = i * component.vertical_sample + v_samp;
let position = 64 * (x2 + y2 * component.width_stride / 8);
let buf_n = &mut buffer[n];
let Some(data) = &mut buf_n.get_mut(position) else {
// TODO: (CAE), this is another weird sub-sampling bug, so on fix
// remove this
return Err(DecodeErrors::FormatStatic("Invalid image"));
};
if self.succ_high == 0 {
stream.decode_prog_dc_first(
&mut self.stream,
huff_table,
data,
&mut component.dc_pred
)?;
} else {
stream.decode_prog_dc_refine(&mut self.stream, data)?;
}
}
}
}
// We want wrapping subtraction here because it means
// we get a higher number in the case this underflows
self.todo -= 1;
// after every scan that's a mcu, count down restart markers.
self.handle_rst_main(stream)?;
}
}
}
return Ok(());
}
pub(crate) fn handle_rst_main(&mut self, stream: &mut BitStream) -> Result<(), DecodeErrors> {
if self.todo == 0 {
stream.refill(&mut self.stream)?;
}
if self.todo == 0
&& self.restart_interval != 0
&& stream.marker.is_none()
&& !stream.seen_eoi
{
// if no marker and we are to reset RST, look for the marker, this matches
// libjpeg-turbo behaviour and allows us to decode images in
// https://github.com/etemesi254/zune-image/issues/261
let _start = self.stream.position()?;
// skip bytes until we find marker
let marker = get_marker(&mut self.stream, stream);
// In some images, the RST marker on the last section may not be available
// as it is maybe stopped by an EOI marker, see in the case of https://github.com/etemesi254/zune-image/issues/292
// what happened was that we would go looking for the RST marker exhausting all the data
// in the image and this would return an error, so for now
// translate it to a warning, but return the image decoded up
// until that point
if let Ok(marker) = marker {
let _end = self.stream.position()?;
stream.marker = Some(marker);
// NB some warnings may be false positives.
warn!(
"{} Extraneous bytes before marker {:?}",
_end - _start,
marker
);
} else {
warn!("RST marker was not found, where expected, image may be garbled")
}
}
if self.todo == 0 {
self.handle_rst(stream)?
}
Ok(())
}
#[allow(clippy::too_many_lines)]
#[allow(clippy::needless_range_loop, clippy::cast_sign_loss)]
fn finish_progressive_decoding(
&mut self, block: &[Vec<i16>; MAX_COMPONENTS], pixels: &mut [u8]
) -> Result<(), DecodeErrors> {
// This function is complicated because we need to replicate
// the function in mcu.rs
//
// The advantage is that we do very little allocation and very lot
// channel reusing.
// The trick is to notice that we repeat the same procedure per MCU
// width.
//
// So we can set it up that we only allocate temporary storage large enough
// to store a single mcu width, then reuse it per invocation.
//
// This is advantageous to us.
//
// Remember we need to have the whole MCU buffer so we store 3 unprocessed
// channels in memory, and then we allocate the whole output buffer in memory, both of
// which are huge.
//
//
let mcu_height = if self.is_interleaved {
self.mcu_y
} else {
// For non-interleaved images( (1*1) subsampling)
// number of MCU's are the widths (+7 to account for paddings) divided by 8.
self.info.height.div_ceil(8) as usize
};
// Size of our output image(width*height)
let is_hv = usize::from(self.is_interleaved);
let upsampler_scratch_size = is_hv * self.components[0].width_stride;
let width = usize::from(self.info.width);
let padded_width = calculate_padded_width(width, self.info.sample_ratio);
let mut upsampler_scratch_space = vec![0; upsampler_scratch_size];
let mut tmp = [0_i32; DCT_BLOCK];
for (pos, comp) in self.components.iter_mut().enumerate() {
// Allocate only needed components.
//
// For special colorspaces i.e YCCK and CMYK, just allocate all of the needed
// components.
if min(
self.options.jpeg_get_out_colorspace().num_components() - 1,
pos
) == pos
|| self.input_colorspace == ColorSpace::YCCK
|| self.input_colorspace == ColorSpace::CMYK
{
// allocate enough space to hold a whole MCU width
// this means we should take into account sampling ratios
// `*8` is because each MCU spans 8 widths.
let len = comp.width_stride * comp.vertical_sample * 8;
comp.needed = true;
comp.raw_coeff = vec![0; len];
} else {
comp.needed = false;
}
}
let mut pixels_written = 0;
// dequantize, idct and color convert.
for i in 0..mcu_height {
'component: for (position, component) in &mut self.components.iter_mut().enumerate() {
if !component.needed {
continue 'component;
}
let qt_table = &component.quantization_table;
// step is the number of pixels this iteration wil be handling
// Given by the number of mcu's height and the length of the component block
// Since the component block contains the whole channel as raw pixels
// we this evenly divides the pixels into MCU blocks
//
// For interleaved images, this gives us the exact pixels comprising a whole MCU
// block
let step = block[position].len() / mcu_height;
// where we will be reading our pixels from.
let start = i * step;
let slice = &block[position][start..start + step];
let temp_channel = &mut component.raw_coeff;
// The next logical step is to iterate width wise.
// To figure out how many pixels we iterate by we use effective pixels
// Given to us by component.x
// iterate per effective pixels.
let mcu_x = component.width_stride / 8;
// iterate per every vertical sample.
for k in 0..component.vertical_sample {
for j in 0..mcu_x {
// after writing a single stride, we need to skip 8 rows.
// This does the row calculation
let width_stride = k * 8 * component.width_stride;
let start = j * 64 + width_stride;
// See https://github.com/etemesi254/zune-image/issues/262 sample 3.
let Some(qt_slice) = slice.get(start..start + 64) else {
return Err(DecodeErrors::FormatStatic(
"Invalid slice , would panic, invalid image"
));
};
// dequantize
for ((x, out), qt_val) in
qt_slice.iter().zip(tmp.iter_mut()).zip(qt_table.iter())
{
*out = i32::from(*x) * qt_val;
}
// determine where to write.
let sl = &mut temp_channel[component.idct_pos..];
component.idct_pos += 8;
// tmp now contains a dequantized block so idct it
(self.idct_func)(&mut tmp, sl, component.width_stride);
}
// after every write of 8, skip 7 since idct write stride wise 8 times.
//
// Remember each MCU is 8x8 block, so each idct will write 8 strides into
// sl
//
// and component.idct_pos is one stride long
component.idct_pos += 7 * component.width_stride;
}
component.idct_pos = 0;
}
// process that width up until it's impossible
self.post_process(
pixels,
i,
mcu_height,
width,
padded_width,
&mut pixels_written,
&mut upsampler_scratch_space
)?;
}
debug!("Finished decoding image");
return Ok(());
}
pub(crate) fn reset_params(&mut self) {
/*
Apparently, grayscale images which can be down sampled exists, which is weird in the sense
that it has one component Y, which is not usually down sampled.
This means some calculations will be wrong, so for that we explicitly reset params
for such occurrences, warn and reset the image info to appear as if it were
a non-sampled image to ensure decoding works
*/
self.h_max = 1;
self.v_max = 1;
self.info.sample_ratio = SampleRatios::None;
self.is_interleaved = false;
self.components[0].vertical_sample = 1;
self.components[0].width_stride = (((self.info.width as usize) + 7) / 8) * 8;
self.components[0].horizontal_sample = 1;
}
}
///Get a marker from the bit-stream.
///
/// This reads until it gets a marker or end of file is encountered
pub fn get_marker<T>(
reader: &mut ZReader<T>, stream: &mut BitStream
) -> Result<Marker, DecodeErrors>
where
T: ZByteReaderTrait
{
if let Some(marker) = stream.marker {
stream.marker = None;
return Ok(marker);
}
// read until we get a marker
while !reader.eof()? {
let marker = reader.read_u8_err()?;
if marker == 255 {
let mut r = reader.read_u8_err()?;
// 0xFF 0XFF(some images may be like that)
while r == 0xFF {
r = reader.read_u8_err()?;
}
if r != 0 {
return Marker::from_u8(r)
.ok_or_else(|| DecodeErrors::Format(format!("Unknown marker 0xFF{r:X}")));
}
}
}
return Err(DecodeErrors::ExhaustedData);
}
#[cfg(test)]
mod tests{
use zune_core::bytestream::ZCursor;
use crate::JpegDecoder;
#[test]
fn make_test(){
let img = "/Users/etemesi/Downloads/wrong_sampling.jpeg";
let data = ZCursor::new([255, 216, 255, 224, 0, 16, 74, 70, 73, 70, 0, 1, 0, 2, 0, 28, 0, 28, 0, 0, 255, 219, 0, 67, 0, 40, 28, 30, 20, 30, 25, 40, 35, 33, 35, 45, 43, 40, 48, 60, 100, 65, 60, 55, 55, 60, 123, 88, 93, 65, 100, 145, 128, 153, 150, 143, 128, 140, 138, 160, 180, 230, 195, 160, 170, 218, 173, 138, 140, 200, 255, 203, 218, 255, 238, 245, 255, 101, 0, 62, 8, 255, 255, 250, 255, 230, 253, 255, 17, 255, 219, 0, 67, 1, 43, 45, 45, 42, 60, 48, 60, 118, 65, 65, 118, 248, 165, 140, 165, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 241, 255, 255, 255, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 255, 192, 0, 17, 8, 0, 32, 0, 32, 3, 2, 17, 0, 1, 34, 1, 3, 17, 1, 255, 196, 0, 24, 0, 1, 1, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 5, 3, 0, 1, 4, 255, 196, 0, 37, 16, 0, 2, 2, 1, 4, 1, 3, 5, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 17, 0, 4, 18, 33, 48, 34, 65, 81, 113, 19, 20, 51, 97, 161, 255, 196, 0, 22, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 255, 196, 0, 26, 17, 1, 0, 2, 3, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 17, 18, 38, 65, 255, 218, 0, 12, 3, 1, 0, 2, 17, 3, 17, 0, 63, 0, 175, 119, 49, 197, 184, 2, 0, 0, 0, 16, 13, 129, 103, 161, 102, 178, 115, 125, 202, 68, 236, 173, 25, 42, 164, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4, 0, 38, 0, 0, 0, 0, 250, 255, 255, 255, 0, 0, 0, 0, 0, 0, 0, 67, 1, 43, 45, 45, 60, 48, 60, 118, 65, 65, 118, 248, 165, 140, 165, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 241, 255, 255, 255, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 255, 192, 0, 17, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 248, 255, 192, 0, 17, 8, 0, 32, 0, 32, 3, 1, 34, 0, 2, 17, 1, 3, 17, 1, 255, 196, 0, 24, 0, 1, 1, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 126, 0, 0, 0, 0, 0, 0, 0, 255, 255, 255, 255, 255, 255, 255, 198]);
let mut decoder = JpegDecoder::new(data);
decoder.decode().unwrap();
}
}
+485
View File
@@ -0,0 +1,485 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
//!Miscellaneous stuff
#![allow(dead_code)]
use alloc::format;
use core::cmp::max;
use core::fmt;
use core::num::NonZeroU32;
use zune_core::bytestream::ZByteReaderTrait;
use zune_core::colorspace::ColorSpace;
use zune_core::log::{trace, warn};
use crate::components::{ComponentID, SampleRatios};
use crate::errors::DecodeErrors;
use crate::huffman::HuffmanTable;
use crate::JpegDecoder;
/// Start of baseline DCT Huffman coding
pub const START_OF_FRAME_BASE: u16 = 0xffc0;
/// Start of another frame
pub const START_OF_FRAME_EXT_SEQ: u16 = 0xffc1;
/// Start of progressive DCT encoding
pub const START_OF_FRAME_PROG_DCT: u16 = 0xffc2;
/// Start of Lossless sequential Huffman coding
pub const START_OF_FRAME_LOS_SEQ: u16 = 0xffc3;
/// Start of extended sequential DCT arithmetic coding
pub const START_OF_FRAME_EXT_AR: u16 = 0xffc9;
/// Start of Progressive DCT arithmetic coding
pub const START_OF_FRAME_PROG_DCT_AR: u16 = 0xffca;
/// Start of Lossless sequential Arithmetic coding
pub const START_OF_FRAME_LOS_SEQ_AR: u16 = 0xffcb;
/// Undo run length encoding of coefficients by placing them in natural order
///
/// This is an index from position-in-bitstream to position-in-row-major-order.
#[rustfmt::skip]
pub const UN_ZIGZAG: [usize; 64 + 16] = [
0, 1, 8, 16, 9, 2, 3, 10,
17, 24, 32, 25, 18, 11, 4, 5,
12, 19, 26, 33, 40, 48, 41, 34,
27, 20, 13, 6, 7, 14, 21, 28,
35, 42, 49, 56, 57, 50, 43, 36,
29, 22, 15, 23, 30, 37, 44, 51,
58, 59, 52, 45, 38, 31, 39, 46,
53, 60, 61, 54, 47, 55, 62, 63,
// Prevent overflowing
63, 63, 63, 63, 63, 63, 63, 63,
63, 63, 63, 63, 63, 63, 63, 63
];
/// Align data to a 16 byte boundary
#[repr(align(16))]
#[derive(Clone)]
pub struct Aligned16<T: ?Sized>(pub T);
impl<T> Default for Aligned16<T>
where
T: Default
{
fn default() -> Self {
Aligned16(T::default())
}
}
/// Align data to a 32 byte boundary
#[repr(align(32))]
#[derive(Clone)]
pub struct Aligned32<T: ?Sized>(pub T);
impl<T> Default for Aligned32<T>
where
T: Default
{
fn default() -> Self {
Aligned32(T::default())
}
}
/// Markers that identify different Start of Image markers
/// They identify the type of encoding and whether the file use lossy(DCT) or
/// lossless compression and whether we use Huffman or arithmetic coding schemes
#[derive(Eq, PartialEq, Copy, Clone)]
#[allow(clippy::upper_case_acronyms)]
pub enum SOFMarkers {
/// Baseline DCT markers
BaselineDct,
/// SOF_1 Extended sequential DCT,Huffman coding
ExtendedSequentialHuffman,
/// Progressive DCT, Huffman coding
ProgressiveDctHuffman,
/// Lossless (sequential), huffman coding,
LosslessHuffman,
/// Extended sequential DEC, arithmetic coding
ExtendedSequentialDctArithmetic,
/// Progressive DCT, arithmetic coding,
ProgressiveDctArithmetic,
/// Lossless ( sequential), arithmetic coding
LosslessArithmetic
}
impl Default for SOFMarkers {
fn default() -> Self {
Self::BaselineDct
}
}
impl SOFMarkers {
/// Check if a certain marker is sequential DCT or not
pub fn is_sequential_dct(self) -> bool {
matches!(
self,
Self::BaselineDct
| Self::ExtendedSequentialHuffman
| Self::ExtendedSequentialDctArithmetic
)
}
/// Check if a marker is a Lossles type or not
pub fn is_lossless(self) -> bool {
matches!(self, Self::LosslessHuffman | Self::LosslessArithmetic)
}
/// Check whether a marker is a progressive marker or not
pub fn is_progressive(self) -> bool {
matches!(
self,
Self::ProgressiveDctHuffman | Self::ProgressiveDctArithmetic
)
}
/// Create a marker from an integer
pub fn from_int(int: u16) -> Option<SOFMarkers> {
match int {
START_OF_FRAME_BASE => Some(Self::BaselineDct),
START_OF_FRAME_PROG_DCT => Some(Self::ProgressiveDctHuffman),
START_OF_FRAME_PROG_DCT_AR => Some(Self::ProgressiveDctArithmetic),
START_OF_FRAME_LOS_SEQ => Some(Self::LosslessHuffman),
START_OF_FRAME_LOS_SEQ_AR => Some(Self::LosslessArithmetic),
START_OF_FRAME_EXT_SEQ => Some(Self::ExtendedSequentialHuffman),
START_OF_FRAME_EXT_AR => Some(Self::ExtendedSequentialDctArithmetic),
_ => None
}
}
}
impl fmt::Debug for SOFMarkers {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match &self {
Self::BaselineDct => write!(f, "Baseline DCT"),
Self::ExtendedSequentialHuffman => {
write!(f, "Extended sequential DCT, Huffman Coding")
}
Self::ProgressiveDctHuffman => write!(f, "Progressive DCT,Huffman Encoding"),
Self::LosslessHuffman => write!(f, "Lossless (sequential) Huffman encoding"),
Self::ExtendedSequentialDctArithmetic => {
write!(f, "Extended sequential DCT, arithmetic coding")
}
Self::ProgressiveDctArithmetic => write!(f, "Progressive DCT, arithmetic coding"),
Self::LosslessArithmetic => write!(f, "Lossless (sequential) arithmetic coding")
}
}
}
/// Set up component parameters.
///
/// This modifies the components in place setting up details needed by other
/// parts fo the decoder.
pub(crate) fn setup_component_params<T: ZByteReaderTrait>(
img: &mut JpegDecoder<T>
) -> Result<(), DecodeErrors> {
let img_width = img.width();
let img_height = img.height();
// in case of adobe app14 being present, zero may indicate
// either CMYK if components are 4 or RGB if components are 3,
// see https://docs.oracle.com/javase/6/docs/api/javax/imageio/metadata/doc-files/jpeg_metadata.html
// so since we may not know how many number of components
// we have when decoding app14, we have to defer that check
// until now.
//
// We know adobe app14 was present since it's the only one that can modify
// input colorspace to be CMYK
if img.components.len() == 3 && img.input_colorspace == ColorSpace::CMYK {
img.input_colorspace = ColorSpace::RGB;
}
for component in &mut img.components {
// compute interleaved image info
// h_max contains the maximum horizontal component
img.h_max = max(img.h_max, component.horizontal_sample);
// v_max contains the maximum vertical component
img.v_max = max(img.v_max, component.vertical_sample);
img.mcu_width = img.h_max * 8;
img.mcu_height = img.v_max * 8;
// Number of MCU's per width
img.mcu_x = usize::from(img.info.width).div_ceil(img.mcu_width);
// Number of MCU's per height
img.mcu_y = usize::from(img.info.height).div_ceil(img.mcu_height);
if img.h_max != 1 || img.v_max != 1 {
// interleaved images have horizontal and vertical sampling factors
// not equal to 1.
img.is_interleaved = true;
}
// Extract quantization tables from the arrays into components
let qt_table = *img.qt_tables[component.quantization_table_number as usize]
.as_ref()
.ok_or_else(|| {
DecodeErrors::DqtError(format!(
"No quantization table for component {:?}",
component.component_id
))
})?;
let x = (usize::from(img_width) * component.horizontal_sample + img.h_max - 1) / img.h_max;
let y = (usize::from(img_height) * component.horizontal_sample + img.h_max - 1) / img.v_max;
component.x = x;
component.w2 = img.mcu_x * component.horizontal_sample * 8;
// probably not needed. :)
component.y = y;
component.quantization_table = qt_table;
// initially stride contains its horizontal sub-sampling
component.width_stride *= img.mcu_x * 8;
}
{
// Sampling factors are one thing that suck
// this fixes a specific problem with images like
//
// (2 2) None
// (2 1) H
// (2 1) H
//
// The images exist in the wild, the images are not meant to exist
// but they do, it's just an annoying horizontal sub-sampling that
// I don't know why it exists.
// But it does
// So we try to cope with that.
// I am not sure of how to explain how to fix it, but it involved a debugger
// and to much coke(the legal one)
//
// If this wasn't present, self.upsample_dest would have the wrong length
let mut handle_that_annoying_bug = false;
if let Some(y_component) = img
.components
.iter()
.find(|c| c.component_id == ComponentID::Y)
{
if y_component.horizontal_sample == 2 || y_component.vertical_sample == 2 {
handle_that_annoying_bug = true;
}
}
if handle_that_annoying_bug {
for comp in &mut img.components {
if (comp.component_id != ComponentID::Y)
&& (comp.horizontal_sample != 1 || comp.vertical_sample != 1)
{
comp.fix_an_annoying_bug = 2;
}
}
}
}
if img.is_mjpeg {
fill_default_mjpeg_tables(
img.is_progressive,
&mut img.dc_huffman_tables,
&mut img.ac_huffman_tables
);
}
// check colorspace matches
if img.input_colorspace.num_components() > img.components.len() {
if img.input_colorspace == ColorSpace::YCCK {
// Some images may have YCCK format (from adobe app14 segment) which is supposed to be 4 components
// but only 3 components, see issue https://github.com/etemesi254/zune-image/issues/275
// So this is the behaviour of other decoders
// - stb_image: Treats it as YCbCr image
// - libjpeg_turbo: Does not know how to parse YCCK images (transform 2 app14) so treats
// it as YCbCr
// So I will match that to match existing ones
warn!("Treating YCCK colorspace as YCbCr as component length does not match");
img.input_colorspace = ColorSpace::YCbCr
} else {
// Note, translated this to a warning to handle valid images of the sort
// See https://github.com/etemesi254/zune-image/issues/288 where there
// was a CMYK image with two components which would be decoded to 4 components
// by the decoder.
// So with a warning that becomes supported.
//
// djpeg fails to render an image from that also probably because it does not
// understand the expected format.
if !img.options.strict_mode() {
warn!(
"Expected {} number of components but found {}",
img.input_colorspace.num_components(),
img.components.len()
);
warn!("Defaulting to multisample to decode");
// N/B: We do not post process the color of such, treating it as multiband
// is the best option since I am not aware of grayscale+alpha which is the most common
// two band format in jpeg.
if img.components.len() > 0 {
img.input_colorspace = ColorSpace::MultiBand(
NonZeroU32::new(img.components.len() as u32).unwrap()
);
}
} else {
let msg = format!(
"Expected {} number of components but found {}",
img.input_colorspace.num_components(),
img.components.len()
);
return Err(DecodeErrors::Format(msg));
}
}
}
Ok(())
}
///Calculate number of fill bytes added to the end of a JPEG image
/// to fill the image
///
/// JPEG usually inserts padding bytes if the image width cannot be evenly divided into
/// 8 , 16 or 32 chunks depending on the sub sampling ratio. So given a sub-sampling ratio,
/// and the actual width, this calculates the padded bytes that were added to the image
///
/// # Params
/// -actual_width: Actual width of the image
/// -sub_sample: Sub sampling factor of the image
///
/// # Returns
/// The padded width, this is how long the width is for a particular image
pub fn calculate_padded_width(actual_width: usize, sub_sample: SampleRatios) -> usize {
match sub_sample {
SampleRatios::None | SampleRatios::V => {
// None+V sends one MCU row, so that's a simple calculation
((actual_width + 7) / 8) * 8
}
SampleRatios::H | SampleRatios::HV => {
// sends two rows, width can be expanded by up to 15 more bytes
((actual_width + 15) / 16) * 16
}
SampleRatios::Generic(h, _) => {
((actual_width + ((h * 8).saturating_sub(1))) / (h * 8)) * (h * 8)
}
}
}
// https://www.loc.gov/preservation/digital/formats/fdd/fdd000063.shtml
// "Avery Lee, writing in the rec.video.desktop newsgroup in 2001, commented that "MJPEG, or at
// least the MJPEG in AVIs having the MJPG fourcc, is restricted JPEG with a fixed -- and
// *omitted* -- Huffman table. The JPEG must be YCbCr colorspace, it must be 4:2:2, and it must
// use basic Huffman encoding, not arithmetic or progressive.... You can indeed extract the
// MJPEG frames and decode them with a regular JPEG decoder, but you have to prepend the DHT
// segment to them, or else the decoder won't have any idea how to decompress the data.
// The exact table necessary is given in the OpenDML spec.""
pub fn fill_default_mjpeg_tables(
is_progressive: bool, dc_huffman_tables: &mut [Option<HuffmanTable>],
ac_huffman_tables: &mut [Option<HuffmanTable>]
) {
// Section K.3.3
trace!("Filling with default mjpeg tables");
if dc_huffman_tables[0].is_none() {
// Table K.3
dc_huffman_tables[0] = Some(
HuffmanTable::new_unfilled(
&[
0x00, 0x00, 0x01, 0x05, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00
],
&[
0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B
],
true,
is_progressive
)
.unwrap()
);
}
if dc_huffman_tables[1].is_none() {
// Table K.4
dc_huffman_tables[1] = Some(
HuffmanTable::new_unfilled(
&[
0x00, 0x00, 0x03, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x00,
0x00, 0x00, 0x00, 0x00
],
&[
0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B
],
true,
is_progressive
)
.unwrap()
);
}
if ac_huffman_tables[0].is_none() {
// Table K.5
ac_huffman_tables[0] = Some(
HuffmanTable::new_unfilled(
&[
0x00, 0x00, 0x02, 0x01, 0x03, 0x03, 0x02, 0x04, 0x03, 0x05, 0x05, 0x04, 0x04,
0x00, 0x00, 0x01, 0x7D
],
&[
0x01, 0x02, 0x03, 0x00, 0x04, 0x11, 0x05, 0x12, 0x21, 0x31, 0x41, 0x06, 0x13,
0x51, 0x61, 0x07, 0x22, 0x71, 0x14, 0x32, 0x81, 0x91, 0xA1, 0x08, 0x23, 0x42,
0xB1, 0xC1, 0x15, 0x52, 0xD1, 0xF0, 0x24, 0x33, 0x62, 0x72, 0x82, 0x09, 0x0A,
0x16, 0x17, 0x18, 0x19, 0x1A, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2A, 0x34, 0x35,
0x36, 0x37, 0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48, 0x49, 0x4A,
0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x63, 0x64, 0x65, 0x66, 0x67,
0x68, 0x69, 0x6A, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A, 0x83, 0x84,
0x85, 0x86, 0x87, 0x88, 0x89, 0x8A, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98,
0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6, 0xA7, 0xA8, 0xA9, 0xAA, 0xB2, 0xB3,
0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xC2, 0xC3, 0xC4, 0xC5, 0xC6, 0xC7,
0xC8, 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0xD9, 0xDA, 0xE1,
0xE2, 0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA, 0xF1, 0xF2, 0xF3, 0xF4,
0xF5, 0xF6, 0xF7, 0xF8, 0xF9, 0xFA
],
false,
is_progressive
)
.unwrap()
);
}
if ac_huffman_tables[1].is_none() {
// Table K.6
ac_huffman_tables[1] = Some(
HuffmanTable::new_unfilled(
&[
0x00, 0x00, 0x02, 0x01, 0x02, 0x04, 0x04, 0x03, 0x04, 0x07, 0x05, 0x04, 0x04,
0x00, 0x01, 0x02, 0x77
],
&[
0x00, 0x01, 0x02, 0x03, 0x11, 0x04, 0x05, 0x21, 0x31, 0x06, 0x12, 0x41, 0x51,
0x07, 0x61, 0x71, 0x13, 0x22, 0x32, 0x81, 0x08, 0x14, 0x42, 0x91, 0xA1, 0xB1,
0xC1, 0x09, 0x23, 0x33, 0x52, 0xF0, 0x15, 0x62, 0x72, 0xD1, 0x0A, 0x16, 0x24,
0x34, 0xE1, 0x25, 0xF1, 0x17, 0x18, 0x19, 0x1A, 0x26, 0x27, 0x28, 0x29, 0x2A,
0x35, 0x36, 0x37, 0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48, 0x49,
0x4A, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x63, 0x64, 0x65, 0x66,
0x67, 0x68, 0x69, 0x6A, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A, 0x82,
0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, 0x8A, 0x92, 0x93, 0x94, 0x95, 0x96,
0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6, 0xA7, 0xA8, 0xA9, 0xAA,
0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xC2, 0xC3, 0xC4, 0xC5,
0xC6, 0xC7, 0xC8, 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0xD9,
0xDA, 0xE2, 0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA, 0xF2, 0xF3, 0xF4,
0xF5, 0xF6, 0xF7, 0xF8, 0xF9, 0xFA
],
false,
is_progressive
)
.unwrap()
);
}
}
+4
View File
@@ -0,0 +1,4 @@
#[cfg(all(feature = "x86", any(target_arch = "x86", target_arch = "x86_64")))]
pub use crate::unsafe_utils_avx2::*;
#[cfg(all(feature = "neon", target_arch = "aarch64"))]
pub use crate::unsafe_utils_neon::*;
+223
View File
@@ -0,0 +1,223 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
#![cfg(all(feature = "x86", any(target_arch = "x86", target_arch = "x86_64")))]
//! This module provides unsafe ways to do some things
#![allow(clippy::wildcard_imports)]
#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;
use core::ops::{Add, AddAssign, Mul, MulAssign, Sub};
/// A copy of `_MM_SHUFFLE()` that doesn't require
/// a nightly compiler
#[inline]
const fn shuffle(z: i32, y: i32, x: i32, w: i32) -> i32 {
(z << 6) | (y << 4) | (x << 2) | w
}
/// An abstraction of an AVX ymm register that
///allows some things to not look ugly
#[derive(Clone, Copy)]
pub struct YmmRegister {
/// An AVX register
pub(crate) mm256: __m256i
}
impl Add for YmmRegister {
type Output = YmmRegister;
#[inline]
fn add(self, rhs: Self) -> Self::Output {
unsafe {
return YmmRegister {
mm256: _mm256_add_epi32(self.mm256, rhs.mm256)
};
}
}
}
impl Add<i32> for YmmRegister {
type Output = YmmRegister;
#[inline]
fn add(self, rhs: i32) -> Self::Output {
unsafe {
let tmp = _mm256_set1_epi32(rhs);
return YmmRegister {
mm256: _mm256_add_epi32(self.mm256, tmp)
};
}
}
}
impl Sub for YmmRegister {
type Output = YmmRegister;
#[inline]
fn sub(self, rhs: Self) -> Self::Output {
unsafe {
return YmmRegister {
mm256: _mm256_sub_epi32(self.mm256, rhs.mm256)
};
}
}
}
impl AddAssign for YmmRegister {
#[inline]
fn add_assign(&mut self, rhs: Self) {
unsafe {
self.mm256 = _mm256_add_epi32(self.mm256, rhs.mm256);
}
}
}
impl AddAssign<i32> for YmmRegister {
#[inline]
fn add_assign(&mut self, rhs: i32) {
unsafe {
let tmp = _mm256_set1_epi32(rhs);
self.mm256 = _mm256_add_epi32(self.mm256, tmp);
}
}
}
impl Mul for YmmRegister {
type Output = YmmRegister;
#[inline]
fn mul(self, rhs: Self) -> Self::Output {
unsafe {
YmmRegister {
mm256: _mm256_mullo_epi32(self.mm256, rhs.mm256)
}
}
}
}
impl Mul<i32> for YmmRegister {
type Output = YmmRegister;
#[inline]
fn mul(self, rhs: i32) -> Self::Output {
unsafe {
let tmp = _mm256_set1_epi32(rhs);
YmmRegister {
mm256: _mm256_mullo_epi32(self.mm256, tmp)
}
}
}
}
impl MulAssign for YmmRegister {
#[inline]
fn mul_assign(&mut self, rhs: Self) {
unsafe {
self.mm256 = _mm256_mullo_epi32(self.mm256, rhs.mm256);
}
}
}
impl MulAssign<i32> for YmmRegister {
#[inline]
fn mul_assign(&mut self, rhs: i32) {
unsafe {
let tmp = _mm256_set1_epi32(rhs);
self.mm256 = _mm256_mullo_epi32(self.mm256, tmp);
}
}
}
impl MulAssign<__m256i> for YmmRegister {
#[inline]
fn mul_assign(&mut self, rhs: __m256i) {
unsafe {
self.mm256 = _mm256_mullo_epi32(self.mm256, rhs);
}
}
}
type Reg = YmmRegister;
/// Transpose an array of 8 by 8 i32's using avx intrinsics
///
/// This was translated from [here](https://newbedev.com/transpose-an-8x8-float-using-avx-avx2)
#[allow(unused_parens, clippy::too_many_arguments)]
#[target_feature(enable = "avx2")]
#[inline]
pub unsafe fn transpose(
v0: &mut Reg, v1: &mut Reg, v2: &mut Reg, v3: &mut Reg, v4: &mut Reg, v5: &mut Reg,
v6: &mut Reg, v7: &mut Reg
) {
macro_rules! merge_epi32 {
($v0:tt,$v1:tt,$v2:tt,$v3:tt) => {
let va = _mm256_permute4x64_epi64($v0, shuffle(3, 1, 2, 0));
let vb = _mm256_permute4x64_epi64($v1, shuffle(3, 1, 2, 0));
$v2 = _mm256_unpacklo_epi32(va, vb);
$v3 = _mm256_unpackhi_epi32(va, vb);
};
}
macro_rules! merge_epi64 {
($v0:tt,$v1:tt,$v2:tt,$v3:tt) => {
let va = _mm256_permute4x64_epi64($v0, shuffle(3, 1, 2, 0));
let vb = _mm256_permute4x64_epi64($v1, shuffle(3, 1, 2, 0));
$v2 = _mm256_unpacklo_epi64(va, vb);
$v3 = _mm256_unpackhi_epi64(va, vb);
};
}
macro_rules! merge_si128 {
($v0:tt,$v1:tt,$v2:tt,$v3:tt) => {
$v2 = _mm256_permute2x128_si256($v0, $v1, shuffle(0, 2, 0, 0));
$v3 = _mm256_permute2x128_si256($v0, $v1, shuffle(0, 3, 0, 1));
};
}
let (w0, w1, w2, w3, w4, w5, w6, w7);
merge_epi32!((v0.mm256), (v1.mm256), w0, w1);
merge_epi32!((v2.mm256), (v3.mm256), w2, w3);
merge_epi32!((v4.mm256), (v5.mm256), w4, w5);
merge_epi32!((v6.mm256), (v7.mm256), w6, w7);
let (x0, x1, x2, x3, x4, x5, x6, x7);
merge_epi64!(w0, w2, x0, x1);
merge_epi64!(w1, w3, x2, x3);
merge_epi64!(w4, w6, x4, x5);
merge_epi64!(w5, w7, x6, x7);
merge_si128!(x0, x4, (v0.mm256), (v1.mm256));
merge_si128!(x1, x5, (v2.mm256), (v3.mm256));
merge_si128!(x2, x6, (v4.mm256), (v5.mm256));
merge_si128!(x3, x7, (v6.mm256), (v7.mm256));
}
+331
View File
@@ -0,0 +1,331 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
#![cfg(all(feature = "neon", target_arch = "aarch64"))]
// TODO can this be extended to armv7
//! This module provides unsafe ways to do some things
#![allow(clippy::wildcard_imports)]
use core::arch::aarch64::*;
use core::ops::{Add, AddAssign, BitOr, BitOrAssign, Mul, MulAssign, Sub};
pub type VecType = int32x4x2_t;
pub unsafe fn loadu(src: *const i32) -> VecType {
vld1q_s32_x2(src as *const _)
}
/// An abstraction of an AVX ymm register that
///allows some things to not look ugly
#[derive(Clone, Copy)]
pub struct YmmRegister {
/// An AVX register
pub(crate) mm256: VecType
}
impl YmmRegister {
#[inline]
pub unsafe fn load(src: *const i32) -> Self {
loadu(src).into()
}
#[inline]
pub fn map2(self, other: Self, f: impl Fn(int32x4_t, int32x4_t) -> int32x4_t) -> Self {
let m0 = f(self.mm256.0, other.mm256.0);
let m1 = f(self.mm256.1, other.mm256.1);
YmmRegister {
mm256: int32x4x2_t(m0, m1)
}
}
#[inline]
pub fn all_zero(self) -> bool {
unsafe {
let both = vorrq_s32(self.mm256.0, self.mm256.1);
let both_unsigned = vreinterpretq_u32_s32(both);
0 == vmaxvq_u32(both_unsigned)
}
}
#[inline]
pub fn const_shl<const N: i32>(self) -> Self {
// Ensure that we logically shift left
unsafe {
let m0 = vreinterpretq_s32_u32(vshlq_n_u32::<N>(vreinterpretq_u32_s32(self.mm256.0)));
let m1 = vreinterpretq_s32_u32(vshlq_n_u32::<N>(vreinterpretq_u32_s32(self.mm256.1)));
YmmRegister {
mm256: int32x4x2_t(m0, m1)
}
}
}
#[inline]
pub fn const_shra<const N: i32>(self) -> Self {
unsafe {
let i0 = vshrq_n_s32::<N>(self.mm256.0);
let i1 = vshrq_n_s32::<N>(self.mm256.1);
YmmRegister {
mm256: int32x4x2_t(i0, i1)
}
}
}
}
impl<T> Add<T> for YmmRegister
where
T: Into<Self>
{
type Output = YmmRegister;
#[inline]
fn add(self, rhs: T) -> Self::Output {
let rhs = rhs.into();
unsafe { self.map2(rhs, |a, b| vaddq_s32(a, b)) }
}
}
impl<T> Sub<T> for YmmRegister
where
T: Into<Self>
{
type Output = YmmRegister;
#[inline]
fn sub(self, rhs: T) -> Self::Output {
let rhs = rhs.into();
unsafe { self.map2(rhs, |a, b| vsubq_s32(a, b)) }
}
}
impl<T> AddAssign<T> for YmmRegister
where
T: Into<Self>
{
#[inline]
fn add_assign(&mut self, rhs: T) {
let rhs: Self = rhs.into();
*self = *self + rhs;
}
}
impl<T> Mul<T> for YmmRegister
where
T: Into<Self>
{
type Output = YmmRegister;
#[inline]
fn mul(self, rhs: T) -> Self::Output {
let rhs = rhs.into();
unsafe { self.map2(rhs, |a, b| vmulq_s32(a, b)) }
}
}
impl<T> MulAssign<T> for YmmRegister
where
T: Into<Self>
{
#[inline]
fn mul_assign(&mut self, rhs: T) {
let rhs: Self = rhs.into();
*self = *self * rhs;
}
}
impl<T> BitOr<T> for YmmRegister
where
T: Into<Self>
{
type Output = YmmRegister;
#[inline]
fn bitor(self, rhs: T) -> Self::Output {
let rhs = rhs.into();
unsafe { self.map2(rhs, |a, b| vorrq_s32(a, b)) }
}
}
impl<T> BitOrAssign<T> for YmmRegister
where
T: Into<Self>
{
#[inline]
fn bitor_assign(&mut self, rhs: T) {
let rhs: Self = rhs.into();
*self = *self | rhs;
}
}
impl From<i32> for YmmRegister {
#[inline]
fn from(val: i32) -> Self {
unsafe {
let dup = vdupq_n_s32(val);
YmmRegister {
mm256: int32x4x2_t(dup, dup)
}
}
}
}
impl From<VecType> for YmmRegister {
#[inline]
fn from(mm256: VecType) -> Self {
YmmRegister { mm256 }
}
}
#[allow(clippy::too_many_arguments)]
#[inline]
unsafe fn transpose4(
v0: &mut int32x4_t, v1: &mut int32x4_t, v2: &mut int32x4_t, v3: &mut int32x4_t
) {
let w0 = vtrnq_s32(
vreinterpretq_s32_s64(vtrn1q_s64(
vreinterpretq_s64_s32(*v0),
vreinterpretq_s64_s32(*v2)
)),
vreinterpretq_s32_s64(vtrn1q_s64(
vreinterpretq_s64_s32(*v1),
vreinterpretq_s64_s32(*v3)
))
);
let w1 = vtrnq_s32(
vreinterpretq_s32_s64(vtrn2q_s64(
vreinterpretq_s64_s32(*v0),
vreinterpretq_s64_s32(*v2)
)),
vreinterpretq_s32_s64(vtrn2q_s64(
vreinterpretq_s64_s32(*v1),
vreinterpretq_s64_s32(*v3)
))
);
*v0 = w0.0;
*v1 = w0.1;
*v2 = w1.0;
*v3 = w1.1;
}
/// Transpose an array of 8 by 8 i32
/// Arm has dedicated interleave/transpose instructions
/// we:
/// 1. Transpose the upper left and lower right quadrants
/// 2. Swap and transpose the upper right and lower left quadrants
#[allow(clippy::too_many_arguments)]
#[inline]
pub unsafe fn transpose(
v0: &mut YmmRegister, v1: &mut YmmRegister, v2: &mut YmmRegister, v3: &mut YmmRegister,
v4: &mut YmmRegister, v5: &mut YmmRegister, v6: &mut YmmRegister, v7: &mut YmmRegister
) {
use core::mem::swap;
let ul0 = &mut v0.mm256.0;
let ul1 = &mut v1.mm256.0;
let ul2 = &mut v2.mm256.0;
let ul3 = &mut v3.mm256.0;
let ur0 = &mut v0.mm256.1;
let ur1 = &mut v1.mm256.1;
let ur2 = &mut v2.mm256.1;
let ur3 = &mut v3.mm256.1;
let ll0 = &mut v4.mm256.0;
let ll1 = &mut v5.mm256.0;
let ll2 = &mut v6.mm256.0;
let ll3 = &mut v7.mm256.0;
let lr0 = &mut v4.mm256.1;
let lr1 = &mut v5.mm256.1;
let lr2 = &mut v6.mm256.1;
let lr3 = &mut v7.mm256.1;
swap(ur0, ll0);
swap(ur1, ll1);
swap(ur2, ll2);
swap(ur3, ll3);
transpose4(ul0, ul1, ul2, ul3);
transpose4(ur0, ur1, ur2, ur3);
transpose4(ll0, ll1, ll2, ll3);
transpose4(lr0, lr1, lr2, lr3);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_transpose() {
fn get_val(i: usize, j: usize) -> i32 {
((i * 8) / (j + 1)) as i32
}
unsafe {
let mut vals: [i32; 8 * 8] = [0; 8 * 8];
for i in 0..8 {
for j in 0..8 {
// some order-dependent value of i and j
let value = get_val(i, j);
vals[i * 8 + j] = value;
}
}
let mut regs: [YmmRegister; 8] = core::mem::transmute(vals);
let mut reg0 = regs[0];
let mut reg1 = regs[1];
let mut reg2 = regs[2];
let mut reg3 = regs[3];
let mut reg4 = regs[4];
let mut reg5 = regs[5];
let mut reg6 = regs[6];
let mut reg7 = regs[7];
transpose(
&mut reg0, &mut reg1, &mut reg2, &mut reg3, &mut reg4, &mut reg5, &mut reg6,
&mut reg7
);
regs[0] = reg0;
regs[1] = reg1;
regs[2] = reg2;
regs[3] = reg3;
regs[4] = reg4;
regs[5] = reg5;
regs[6] = reg6;
regs[7] = reg7;
let vals_from_reg: [i32; 8 * 8] = core::mem::transmute(regs);
for i in 0..8 {
for j in 0..i {
let orig = vals[i * 8 + j];
vals[i * 8 + j] = vals[j * 8 + i];
vals[j * 8 + i] = orig;
}
}
for i in 0..8 {
for j in 0..8 {
assert_eq!(vals[j * 8 + i], get_val(i, j));
assert_eq!(vals_from_reg[j * 8 + i], get_val(i, j));
}
}
assert_eq!(vals, vals_from_reg);
}
}
}
+343
View File
@@ -0,0 +1,343 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
//! Up-sampling routines
//!
//! The main upsampling method is a bi-linear interpolation or a "triangle
//! filter " or libjpeg turbo `fancy_upsampling` which is a good compromise
//! between speed and visual quality
//!
//! # The filter
//! Each output pixel is made from `(3*A+B)/4` where A is the original
//! pixel closer to the output and B is the one further.
//!
//! ```text
//!+---+---+
//! | A | B |
//! +---+---+
//! +-+-+-+-+
//! | |P| | |
//! +-+-+-+-+
//! ```
//!
//! # Horizontal Bi-linear filter
//! ```text
//! |---+-----------+---+
//! | | | |
//! | A | |p1 | p2| | B |
//! | | | |
//! |---+-----------+---+
//!
//! ```
//! For a horizontal bi-linear it's trivial to implement,
//!
//! `A` becomes the input closest to the output.
//!
//! `B` varies depending on output.
//! - For odd positions, input is the `next` pixel after A
//! - For even positions, input is the `previous` value before A.
//!
//! We iterate in a classic 1-D sliding window with a window of 3.
//! For our sliding window approach, `A` is the 1st and `B` is either the 0th term or 2nd term
//! depending on position we are writing.(see scalar code).
//!
//! For vector code see module sse for explanation.
//!
//! # Vertical bi-linear.
//! Vertical up-sampling is a bit trickier.
//!
//! ```text
//! +----+----+
//! | A1 | A2 |
//! +----+----+
//! +----+----+
//! | p1 | p2 |
//! +----+-+--+
//! +----+-+--+
//! | p3 | p4 |
//! +----+-+--+
//! +----+----+
//! | B1 | B2 |
//! +----+----+
//! ```
//!
//! For `p1`
//! - `A1` is given a weight of `3` and `B1` is given a weight of 1.
//!
//! For `p3`
//! - `B1` is given a weight of `3` and `A1` is given a weight of 1
//!
//! # Horizontal vertical downsampling/chroma quartering.
//!
//! Carry out a vertical filter in the first pass, then a horizontal filter in the second pass.
#![allow(unreachable_code)]
use zune_core::options::DecoderOptions;
use crate::components::UpSampler;
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[cfg(feature = "x86")]
mod avx2;
#[cfg(target_arch = "aarch64")]
#[cfg(feature = "neon")]
mod neon;
#[cfg(feature = "portable_simd")]
mod portable_simd;
mod scalar;
// choose the best possible implementation for this platform
#[allow(unused_variables)]
pub fn choose_horizontal_samp_function(options: &DecoderOptions) -> UpSampler {
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[cfg(feature = "x86")]
if options.use_avx2() {
return |a: &[i16], b: &[i16], c: &[i16], d: &mut [i16], e: &mut [i16]| {
// SAFETY: `options.use_avx2()` only returns true if avx2 is supported.
unsafe { avx2::upsample_horizontal_avx2(a, b, c, d, e) }
};
}
#[cfg(target_arch = "aarch64")]
#[cfg(feature = "neon")]
if options.use_neon() {
return |a: &[i16], b: &[i16], c: &[i16], d: &mut [i16], e: &mut [i16]| {
// SAFETY: `options.use_neon()` only returns true if neon is supported.
unsafe { neon::upsample_horizontal_neon(a, b, c, d, e) }
};
}
#[cfg(feature = "portable_simd")]
return portable_simd::upsample_horizontal_simd;
return scalar::upsample_horizontal;
}
#[allow(unused_variables)]
pub fn choose_hv_samp_function(options: &DecoderOptions) -> UpSampler {
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[cfg(feature = "x86")]
if options.use_avx2() {
return |a: &[i16], b: &[i16], c: &[i16], d: &mut [i16], e: &mut [i16]| {
// SAFETY: `options.use_avx2()` only returns true if avx2 is supported.
unsafe { avx2::upsample_hv_avx2(a, b, c, d, e) }
};
}
#[cfg(target_arch = "aarch64")]
#[cfg(feature = "neon")]
if options.use_neon() {
return |a: &[i16], b: &[i16], c: &[i16], d: &mut [i16], e: &mut [i16]| {
// SAFETY: `options.use_neon()` only returns true if neon is supported.
unsafe { neon::upsample_hv_neon(a, b, c, d, e) }
};
}
#[cfg(feature = "portable_simd")]
return portable_simd::upsample_hv_simd;
return scalar::upsample_hv;
}
#[allow(unused_variables)]
pub fn choose_v_samp_function(options: &DecoderOptions) -> UpSampler {
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[cfg(feature = "x86")]
if options.use_avx2() {
return |a: &[i16], b: &[i16], c: &[i16], d: &mut [i16], e: &mut [i16]| {
// SAFETY: `options.use_avx2()` only returns true if avx2 is supported.
unsafe { avx2::upsample_vertical_avx2(a, b, c, d, e) }
};
}
#[cfg(target_arch = "aarch64")]
#[cfg(feature = "neon")]
if options.use_neon() {
return |a: &[i16], b: &[i16], c: &[i16], d: &mut [i16], e: &mut [i16]| {
// SAFETY: `options.use_neon()` only returns true if neon is supported.
unsafe { neon::upsample_vertical_neon(a, b, c, d, e) }
};
}
#[cfg(feature = "portable_simd")]
return portable_simd::upsample_vertical_simd;
return scalar::upsample_vertical;
}
/// Upsample nothing
pub fn upsample_no_op(
_input: &[i16],
_in_ref: &[i16],
_in_near: &[i16],
_scratch_space: &mut [i16],
_output: &mut [i16],
) {
}
pub fn generic_sampler() -> UpSampler {
scalar::upsample_generic
}
#[cfg(test)]
mod tests {
use super::*;
#[cfg(feature = "portable_simd")]
mod portable_simd_impl {
use super::*;
#[test]
fn portable_simd_vertical() {
_test_vertical(portable_simd::upsample_vertical_simd)
}
#[test]
fn portable_simd_horizontal() {
_test_horizontal(portable_simd::upsample_horizontal_simd)
}
#[test]
fn portable_simd_hv() {
_test_hv(portable_simd::upsample_hv_simd)
}
}
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[cfg(feature = "x86")]
#[cfg(target_feature = "avx2")]
mod avx2_impl {
use super::*;
#[test]
fn avx2_vertical() {
_test_vertical(|a: &[i16], b: &[i16], c: &[i16], d: &mut [i16], e: &mut [i16]| {
// SAFETY: Test guarded behind `target_feature`
unsafe { avx2::upsample_vertical_avx2(a, b, c, d, e) }
})
}
#[test]
fn avx2_horizontal() {
_test_horizontal(|a: &[i16], b: &[i16], c: &[i16], d: &mut [i16], e: &mut [i16]| {
// SAFETY: Test guarded behind `target_feature`
unsafe { avx2::upsample_horizontal_avx2(a, b, c, d, e) }
})
}
#[test]
fn avx2_hv() {
_test_hv(|a: &[i16], b: &[i16], c: &[i16], d: &mut [i16], e: &mut [i16]| {
// SAFETY: Test guarded behind `target_feature`
unsafe { avx2::upsample_hv_avx2(a, b, c, d, e) }
})
}
}
#[cfg(target_arch = "aarch64")]
#[cfg(feature = "neon")]
#[cfg(target_feature = "neon")]
mod neon_impl {
use super::*;
#[test]
fn neon_vertical() {
_test_vertical(|a: &[i16], b: &[i16], c: &[i16], d: &mut [i16], e: &mut [i16]| {
// SAFETY: Test guarded behind `target_feature`
unsafe { neon::upsample_vertical_neon(a, b, c, d, e) }
})
}
#[test]
fn neon_horizontal() {
_test_horizontal(|a: &[i16], b: &[i16], c: &[i16], d: &mut [i16], e: &mut [i16]| {
// SAFETY: Test guarded behind `target_feature`
unsafe { neon::upsample_horizontal_neon(a, b, c, d, e) }
})
}
#[test]
fn neon_hv() {
_test_hv(|a: &[i16], b: &[i16], c: &[i16], d: &mut [i16], e: &mut [i16]| {
// SAFETY: Test guarded behind `target_feature`
unsafe { neon::upsample_hv_neon(a, b, c, d, e) }
})
}
}
fn _test_vertical(upsampler: UpSampler) {
let width = 1024;
let input: Vec<i16> = (0..width).map(|x| ((x + 10) % 256) as i16).collect();
let in_near: Vec<i16> = (0..width).map(|x| ((x + 20) % 256) as i16).collect();
let in_far: Vec<i16> = (0..width).map(|x| ((x + 30) % 256) as i16).collect();
let mut scratch = vec![0i16; width];
let mut output_scalar = vec![0i16; width * 2];
let mut output_fast = vec![0i16; width * 2];
scalar::upsample_vertical(&input, &in_near, &in_far, &mut scratch, &mut output_scalar);
upsampler(&input, &in_near, &in_far, &mut scratch, &mut output_fast);
assert_eq!(output_scalar, output_fast);
}
fn _test_horizontal(upsampler: UpSampler) {
_test_horizontal_even_width(upsampler);
_test_horizontal_odd_width(upsampler);
}
fn _test_horizontal_even_width(upsampler: UpSampler) {
let width = 1024;
let input: Vec<i16> = (0..width).map(|x| ((x + 10) % 256) as i16).collect();
let mut scratch = vec![0i16; width];
let mut output_scalar = vec![0i16; width * 2];
let mut output_fast = vec![0i16; width * 2];
scalar::upsample_horizontal(&input, &[], &[], &mut scratch, &mut output_scalar);
upsampler(&input, &[], &[], &mut scratch, &mut output_fast);
assert_eq!(output_scalar, output_fast);
}
fn _test_horizontal_odd_width(upsampler: UpSampler) {
let width = 33;
let input: Vec<i16> = (0..width).map(|x| ((x + 10) % 256) as i16).collect();
let mut scratch = vec![0i16; width];
let mut output_scalar = vec![0i16; width * 2];
let mut output_fast = vec![0i16; width * 2];
scalar::upsample_horizontal(&input, &[], &[], &mut scratch, &mut output_scalar);
upsampler(&input, &[], &[], &mut scratch, &mut output_fast);
assert_eq!(output_scalar, output_fast);
}
fn _test_hv(upsampler: UpSampler) {
let width = 512;
let input: Vec<i16> = (0..width).map(|x| ((x + 10) % 256) as i16).collect();
let in_near: Vec<i16> = (0..width).map(|x| ((x + 20) % 256) as i16).collect();
let in_far: Vec<i16> = (0..width).map(|x| ((x + 30) % 256) as i16).collect();
// Output len is width * 4 for HV (vertical * 2, then horizontal * 2 for each row)
// scratch is width * 2
let mut scratch_scalar = vec![0i16; width * 2];
let mut scratch_fast = vec![0i16; width * 2];
let mut output_scalar = vec![0i16; width * 4];
let mut output_fast = vec![0i16; width * 4];
scalar::upsample_hv(
&input,
&in_near,
&in_far,
&mut scratch_scalar,
&mut output_scalar,
);
upsampler(
&input,
&in_near,
&in_far,
&mut scratch_fast,
&mut output_fast,
);
assert_eq!(output_scalar, output_fast);
}
}
+199
View File
@@ -0,0 +1,199 @@
/*
* Copyright (c) 2025.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "avx2")]
pub unsafe fn upsample_horizontal_avx2(
input: &[i16],
in_near: &[i16],
in_far: &[i16],
scratch: &mut [i16],
output: &mut [i16],
) {
assert_eq!(input.len() * 2, output.len());
assert!(input.len() > 2);
let len = input.len();
if len < 18 {
return super::scalar::upsample_horizontal(input, in_near, in_far, scratch, output);
}
// First two pixels
output[0] = input[0];
output[1] = (input[0] * 3 + input[1] + 2) >> 2;
let v_three = _mm256_set1_epi16(3);
let v_two = _mm256_set1_epi16(2);
let upsample16 = |input: &[i16; 18], output: &mut [i16; 32]| {
let in_ptr = input.as_ptr();
let out_ptr = output.as_mut_ptr();
// SAFETY: The input is 18 * 16 bit long, so the loads are safe.
let (v_prev, v_curr, v_next) = unsafe {
(
_mm256_loadu_si256(in_ptr.add(0) as *const __m256i),
_mm256_loadu_si256(in_ptr.add(1) as *const __m256i),
_mm256_loadu_si256(in_ptr.add(2) as *const __m256i),
)
};
let v_common = _mm256_add_epi16(_mm256_mullo_epi16(v_curr, v_three), v_two);
let v_even = _mm256_srai_epi16(_mm256_add_epi16(v_common, v_prev), 2);
let v_odd = _mm256_srai_epi16(_mm256_add_epi16(v_common, v_next), 2);
let v_res_1 = _mm256_unpacklo_epi16(v_even, v_odd);
let v_res_2 = _mm256_unpackhi_epi16(v_even, v_odd);
let v_final_1 = _mm256_permute2x128_si256(v_res_1, v_res_2, 0x20);
let v_final_2 = _mm256_permute2x128_si256(v_res_1, v_res_2, 0x31);
// SAFETY: The output is 32 * 16 bit long, so the stores are safe.
unsafe {
_mm256_storeu_si256(out_ptr as *mut __m256i, v_final_1);
_mm256_storeu_si256(out_ptr.add(16) as *mut __m256i, v_final_2);
}
};
for (input, output) in input
.windows(18)
.step_by(16)
.zip(output[2..].chunks_exact_mut(32))
{
upsample16(input.try_into().unwrap(), output.try_into().unwrap());
}
// Upsample the remainder. This may have some overlap, but that's fine.
if let Some(rest_input) = input.last_chunk::<18>() {
let end = output.len() - 2;
if let Some(rest_output) = output[..end].last_chunk_mut::<32>() {
upsample16(rest_input, rest_output);
}
}
// Last two pixels.
output[output.len() - 2] = (3 * input[len - 1] + input[len - 2] + 2) >> 2;
output[output.len() - 1] = input[len - 1];
}
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "avx2")]
pub unsafe fn upsample_vertical_avx2(
input: &[i16],
in_near: &[i16],
in_far: &[i16],
scratch: &mut [i16],
output: &mut [i16],
) {
assert_eq!(input.len() * 2, output.len());
assert_eq!(in_near.len(), input.len());
assert_eq!(in_far.len(), input.len());
let len = input.len();
if len < 16 {
return super::scalar::upsample_vertical(input, in_near, in_far, scratch, output);
}
let middle = output.len() / 2;
let (out_top, out_bottom) = output.split_at_mut(middle);
let v_three = _mm256_set1_epi16(3);
let v_two = _mm256_set1_epi16(2);
let upsample16 = |input: &[i16; 16],
in_near: &[i16; 16],
in_far: &[i16; 16],
out_top: &mut [i16; 16],
out_bottom: &mut [i16; 16]| {
// SAFETY: Inputs are all 16 * 16 bit long, so the loads are safe.
let (v_in, v_near, v_far) = unsafe {
(
_mm256_loadu_si256(input.as_ptr() as *const __m256i),
_mm256_loadu_si256(in_near.as_ptr() as *const __m256i),
_mm256_loadu_si256(in_far.as_ptr() as *const __m256i),
)
};
let v_common = _mm256_add_epi16(_mm256_mullo_epi16(v_in, v_three), v_two);
let v_out_top = _mm256_srai_epi16(_mm256_add_epi16(v_common, v_near), 2);
let v_out_bottom = _mm256_srai_epi16(_mm256_add_epi16(v_common, v_far), 2);
// SAFETY: Outputs are 16 * 16 bit long, so the stores are safe.
unsafe {
_mm256_storeu_si256(out_top.as_mut_ptr() as *mut __m256i, v_out_top);
_mm256_storeu_si256(out_bottom.as_mut_ptr() as *mut __m256i, v_out_bottom);
}
};
let chunks = input
.chunks_exact(16)
.zip(in_near.chunks_exact(16))
.zip(in_far.chunks_exact(16))
.zip(out_top.chunks_exact_mut(16))
.zip(out_bottom.chunks_exact_mut(16));
for ((((input, in_near), in_far), out_top), out_bottom) in chunks {
upsample16(
input.try_into().unwrap(),
in_near.try_into().unwrap(),
in_far.try_into().unwrap(),
out_top.try_into().unwrap(),
out_bottom.try_into().unwrap(),
);
}
// Upsample the remainder. This may have some overlap, but that's fine.
// Edition upgrade will fix this nested awfulness.
if let Some(rest) = input.last_chunk::<16>() {
if let Some(rest_near) = in_near.last_chunk::<16>() {
if let Some(rest_far) = in_far.last_chunk::<16>() {
if let Some(mut rest_top) = out_top.last_chunk_mut::<16>() {
if let Some(mut rest_bottom) = out_bottom.last_chunk_mut::<16>() {
upsample16(rest, rest_near, rest_far, &mut rest_top, &mut rest_bottom);
}
}
}
}
}
}
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "avx2")]
pub unsafe fn upsample_hv_avx2(
input: &[i16],
in_near: &[i16],
in_far: &[i16],
scratch_space: &mut [i16],
output: &mut [i16],
) {
assert_eq!(input.len() * 4, output.len());
assert!(input.len() * 2 <= scratch_space.len());
let scratch_space = &mut scratch_space[..input.len() * 2];
upsample_vertical_avx2(input, in_near, in_far, &mut [], scratch_space);
let scratch_half = scratch_space.len() / 2;
let output_half = output.len() / 2;
let (scratch_top, scratch_bottom) = scratch_space.split_at_mut(scratch_half);
let (out_top, out_bottom) = output.split_at_mut(output_half);
let mut t = [0];
upsample_horizontal_avx2(scratch_top, &[], &[], &mut t, out_top);
upsample_horizontal_avx2(scratch_bottom, &[], &[], &mut t, out_bottom);
}
+191
View File
@@ -0,0 +1,191 @@
/*
* Copyright (c) 2025.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
#[cfg(target_arch = "aarch64")]
use core::arch::aarch64::*;
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
pub unsafe fn upsample_horizontal_neon(
input: &[i16], in_near: &[i16], in_far: &[i16], scratch: &mut [i16], output: &mut [i16]
) {
assert_eq!(input.len() * 2, output.len());
assert!(input.len() > 2);
let len = input.len();
if len < 10 {
return super::scalar::upsample_horizontal(input, in_near, in_far, scratch, output);
}
// First two pixels
output[0] = input[0];
output[1] = (input[0] * 3 + input[1] + 2) >> 2;
// SAFETY: NEON target feature is enabled on this function.
let v_three = unsafe { vdupq_n_s16(3) };
// SAFETY: NEON target feature is enabled on this function.
let v_two = unsafe { vdupq_n_s16(2) };
let upsample8 = |input: &[i16; 10], output: &mut [i16; 16]| {
let in_ptr = input.as_ptr();
let out_ptr = output.as_mut_ptr();
// SAFETY: The input is 10 * 16 bit long, so the loads are safe.
let (v_prev, v_curr, v_next) = unsafe {
(
vld1q_s16(in_ptr),
vld1q_s16(in_ptr.add(1)),
vld1q_s16(in_ptr.add(2))
)
};
// SAFETY: NEON target feature is enabled and vector lanes are valid.
let v_common = unsafe { vaddq_s16(vmulq_s16(v_curr, v_three), v_two) };
// SAFETY: NEON target feature is enabled and vector lanes are valid.
let v_even = unsafe { vshrq_n_s16::<2>(vaddq_s16(v_common, v_prev)) };
// SAFETY: NEON target feature is enabled and vector lanes are valid.
let v_odd = unsafe { vshrq_n_s16::<2>(vaddq_s16(v_common, v_next)) };
// SAFETY: NEON target feature is enabled and vector lanes are valid.
let v_res_1 = unsafe { vzip1q_s16(v_even, v_odd) };
// SAFETY: NEON target feature is enabled and vector lanes are valid.
let v_res_2 = unsafe { vzip2q_s16(v_even, v_odd) };
// SAFETY: The output is 16 * 16 bit long, so the stores are safe.
unsafe {
vst1q_s16(out_ptr, v_res_1);
vst1q_s16(out_ptr.add(8), v_res_2);
}
};
for (input, output) in input
.windows(10)
.step_by(8)
.zip(output[2..].chunks_exact_mut(16))
{
upsample8(input.try_into().unwrap(), output.try_into().unwrap());
}
// Upsample the remainder. This may have some overlap, but that's fine.
if let Some(rest_input) = input.last_chunk::<10>() {
let end = output.len() - 2;
if let Some(rest_output) = output[..end].last_chunk_mut::<16>() {
upsample8(rest_input, rest_output);
}
}
// Last two pixels.
output[output.len() - 2] = (3 * input[len - 1] + input[len - 2] + 2) >> 2;
output[output.len() - 1] = input[len - 1];
}
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
pub unsafe fn upsample_vertical_neon(
input: &[i16], in_near: &[i16], in_far: &[i16], scratch: &mut [i16], output: &mut [i16]
) {
assert_eq!(input.len() * 2, output.len());
assert_eq!(in_near.len(), input.len());
assert_eq!(in_far.len(), input.len());
let len = input.len();
if len < 16 {
return super::scalar::upsample_vertical(input, in_near, in_far, scratch, output);
}
let middle = output.len() / 2;
let (out_top, out_bottom) = output.split_at_mut(middle);
// SAFETY: NEON target feature is enabled on this function.
let v_three = unsafe { vdupq_n_s16(3) };
// SAFETY: NEON target feature is enabled on this function.
let v_two = unsafe { vdupq_n_s16(2) };
let upsample8 = |input: &[i16; 8],
in_near: &[i16; 8],
in_far: &[i16; 8],
out_top: &mut [i16; 8],
out_bottom: &mut [i16; 8]| {
// SAFETY: Inputs are all 8 * 16 bit long, so the loads are safe.
let (v_in, v_near, v_far) = unsafe {
(
vld1q_s16(input.as_ptr()),
vld1q_s16(in_near.as_ptr()),
vld1q_s16(in_far.as_ptr())
)
};
// SAFETY: NEON target feature is enabled and vector lanes are valid.
let v_common = unsafe { vaddq_s16(vmulq_s16(v_in, v_three), v_two) };
// SAFETY: NEON target feature is enabled and vector lanes are valid.
let v_out_top = unsafe { vshrq_n_s16::<2>(vaddq_s16(v_common, v_near)) };
// SAFETY: NEON target feature is enabled and vector lanes are valid.
let v_out_bottom = unsafe { vshrq_n_s16::<2>(vaddq_s16(v_common, v_far)) };
// SAFETY: Outputs are 8 * 16 bit long, so the stores are safe.
unsafe {
vst1q_s16(out_top.as_mut_ptr(), v_out_top);
vst1q_s16(out_bottom.as_mut_ptr(), v_out_bottom);
}
};
let chunks = input
.chunks_exact(8)
.zip(in_near.chunks_exact(8))
.zip(in_far.chunks_exact(8))
.zip(out_top.chunks_exact_mut(8))
.zip(out_bottom.chunks_exact_mut(8));
for ((((input, in_near), in_far), out_top), out_bottom) in chunks {
upsample8(
input.try_into().unwrap(),
in_near.try_into().unwrap(),
in_far.try_into().unwrap(),
out_top.try_into().unwrap(),
out_bottom.try_into().unwrap()
);
}
// Upsample the remainder.
if let Some(rest) = input.last_chunk::<8>() {
if let Some(rest_near) = in_near.last_chunk::<8>() {
if let Some(rest_far) = in_far.last_chunk::<8>() {
if let Some(mut rest_top) = out_top.last_chunk_mut::<8>() {
if let Some(mut rest_bottom) = out_bottom.last_chunk_mut::<8>() {
upsample8(rest, rest_near, rest_far, &mut rest_top, &mut rest_bottom);
}
}
}
}
}
}
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
pub unsafe fn upsample_hv_neon(
input: &[i16], in_near: &[i16], in_far: &[i16], scratch_space: &mut [i16], output: &mut [i16]
) {
assert_eq!(input.len() * 4, output.len());
assert!(input.len() * 2 <= scratch_space.len());
let scratch_space = &mut scratch_space[..input.len() * 2];
unsafe { upsample_vertical_neon(input, in_near, in_far, &mut [], scratch_space) };
let scratch_half = scratch_space.len() / 2;
let output_half = output.len() / 2;
let (scratch_top, scratch_bottom) = scratch_space.split_at_mut(scratch_half);
let (out_top, out_bottom) = output.split_at_mut(output_half);
let mut t = [0];
unsafe { upsample_horizontal_neon(scratch_top, &[], &[], &mut t, out_top) };
unsafe { upsample_horizontal_neon(scratch_bottom, &[], &[], &mut t, out_bottom) };
}
+171
View File
@@ -0,0 +1,171 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
use std::simd::prelude::*;
const LANES: usize = 16;
type V = Simd<i16, LANES>;
pub fn upsample_horizontal_simd(
input: &[i16],
in_near: &[i16],
in_far: &[i16],
scratch: &mut [i16],
output: &mut [i16],
) {
assert_eq!(input.len() * 2, output.len());
assert!(input.len() > 2);
let len = input.len();
if len < 18 {
return super::scalar::upsample_horizontal(input, in_near, in_far, scratch, output);
}
// First two pixels
output[0] = input[0];
output[1] = (input[0] * 3 + input[1] + 2) >> 2;
let v_three = V::splat(3);
let v_two = V::splat(2);
let upsample16 = |input: &[i16; 18], output: &mut [i16; 32]| {
let v_prev = V::from_slice(&input[0..LANES]);
let v_curr = V::from_slice(&input[1..LANES + 1]);
let v_next = V::from_slice(&input[2..LANES + 2]);
let v_common = v_curr * v_three + v_two;
let v_even = (v_common + v_prev) >> 2;
let v_odd = (v_common + v_next) >> 2;
let (v_res_1, v_res_2) = v_even.interleave(v_odd);
v_res_1.copy_to_slice(&mut output[0..LANES]);
v_res_2.copy_to_slice(&mut output[LANES..2 * LANES]);
};
for (input, output) in input
.windows(18)
.step_by(16)
.zip(output[2..].chunks_exact_mut(32))
{
upsample16(input.try_into().unwrap(), output.try_into().unwrap());
}
// Upsample the remainder. This may have some overlap, but that's fine.
if let Some(rest_input) = input.last_chunk::<18>() {
let end = output.len() - 2;
if let Some(rest_output) = output[..end].last_chunk_mut::<32>() {
upsample16(rest_input, rest_output);
}
}
// Last two pixels.
output[output.len() - 2] = (3 * input[len - 1] + input[len - 2] + 2) >> 2;
output[output.len() - 1] = input[len - 1];
}
pub fn upsample_vertical_simd(
input: &[i16],
in_near: &[i16],
in_far: &[i16],
_scratch_space: &mut [i16],
output: &mut [i16],
) {
assert_eq!(input.len() * 2, output.len());
assert_eq!(in_near.len(), input.len());
assert_eq!(in_far.len(), input.len());
let len = input.len();
if len < 16 {
return super::scalar::upsample_vertical(input, in_near, in_far, _scratch_space, output);
}
let middle = output.len() / 2;
let (out_top, out_bottom) = output.split_at_mut(middle);
let v_three = V::splat(3);
let v_two = V::splat(2);
let upsample16 = |input: &[i16; 16],
in_near: &[i16; 16],
in_far: &[i16; 16],
out_top: &mut [i16; 16],
out_bottom: &mut [i16; 16]| {
let v_in = V::from(*input);
let v_near = V::from(*in_near);
let v_far = V::from(*in_far);
let v_common = v_in * v_three + v_two;
let v_out_top = (v_common + v_near) >> 2;
let v_out_bottom = (v_common + v_far) >> 2;
v_out_top.copy_to_slice(out_top.as_mut_slice());
v_out_bottom.copy_to_slice(out_bottom.as_mut_slice());
};
let chunks = input
.chunks_exact(16)
.zip(in_near.chunks_exact(16))
.zip(in_far.chunks_exact(16))
.zip(out_top.chunks_exact_mut(16))
.zip(out_bottom.chunks_exact_mut(16));
for ((((input, in_near), in_far), out_top), out_bottom) in chunks {
upsample16(
input.try_into().unwrap(),
in_near.try_into().unwrap(),
in_far.try_into().unwrap(),
out_top.try_into().unwrap(),
out_bottom.try_into().unwrap(),
);
}
// Upsample the remainder. This may have some overlap, but that's fine.
// Edition upgrade will fix this nested awfulness.
if let Some(rest) = input.last_chunk::<16>() {
if let Some( rest_near) = in_near.last_chunk::<16>() {
if let Some( rest_far) = in_far.last_chunk::<16>() {
if let Some( rest_top) = out_top.last_chunk_mut::<16>() {
if let Some( rest_bottom) = out_bottom.last_chunk_mut::<16>() {
upsample16(rest, rest_near, rest_far, rest_top, rest_bottom);
}
}
}
}
}
}
pub fn upsample_hv_simd(
input: &[i16],
in_near: &[i16],
in_far: &[i16],
scratch_space: &mut [i16],
output: &mut [i16],
) {
assert_eq!(input.len() * 4, output.len());
assert!(input.len() * 2 <= scratch_space.len());
let scratch_space = &mut scratch_space[..input.len() * 2];
upsample_vertical_simd(input, in_near, in_far, &mut [], scratch_space);
let scratch_half = scratch_space.len() / 2;
let output_half = output.len() / 2;
let (scratch_top, scratch_bottom) = scratch_space.split_at_mut(scratch_half);
let (out_top, out_bottom) = output.split_at_mut(output_half);
let mut t = [0];
upsample_horizontal_simd(scratch_top, &[], &[], &mut t, out_top);
upsample_horizontal_simd(scratch_bottom, &[], &[], &mut t, out_bottom);
}
+129
View File
@@ -0,0 +1,129 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
pub fn upsample_horizontal(
input: &[i16], _ref: &[i16], _in_near: &[i16], _scratch: &mut [i16], output: &mut [i16]
) {
assert_eq!(
input.len() * 2,
output.len(),
"Input length is not half the size of the output length"
);
assert!(
output.len() > 4 && input.len() > 2,
"Too Short of a vector, cannot upsample"
);
output[0] = input[0];
output[1] = (input[0] * 3 + input[1] + 2) >> 2;
// This code is written for speed and not readability
//
// The readable code is
//
// for i in 1..input.len() - 1{
// let sample = 3 * input[i] + 2;
// out[i * 2] = (sample + input[i - 1]) >> 2;
// out[i * 2 + 1] = (sample + input[i + 1]) >> 2;
// }
//
// The output of a pixel is determined by it's surrounding neighbours but we attach more weight to it's nearest
// neighbour (input[i]) than to the next nearest neighbour.
for (output_window, input_window) in output[2..].chunks_exact_mut(2).zip(input.windows(3)) {
let sample = 3 * input_window[1] + 2;
output_window[0] = (sample + input_window[0]) >> 2;
output_window[1] = (sample + input_window[2]) >> 2;
}
// Get lengths
let out_len = output.len() - 2;
let input_len = input.len() - 2;
// slice the output vector
let f_out = &mut output[out_len..];
let i_last = &input[input_len..];
// write out manually..
f_out[0] = (3 * i_last[1] + i_last[0] + 2) >> 2;
f_out[1] = i_last[1];
}
pub fn upsample_vertical(
input: &[i16], in_near: &[i16], in_far: &[i16], _scratch_space: &mut [i16], output: &mut [i16]
) {
assert_eq!(input.len() * 2, output.len());
assert_eq!(in_near.len(), input.len());
assert_eq!(in_far.len(), input.len());
let middle = output.len() / 2;
let (out_top, out_bottom) = output.split_at_mut(middle);
// for the first row, closest row is in_near
for ((near, far), x) in input.iter().zip(in_near.iter()).zip(out_top) {
*x = (((3 * near) + 2) + far) >> 2;
}
// for the second row, the closest row to input is in_far
for ((near, far), x) in input.iter().zip(in_far.iter()).zip(out_bottom) {
*x = (((3 * near) + 2) + far) >> 2;
}
}
pub fn upsample_hv(
input: &[i16], in_near: &[i16], in_far: &[i16], scratch_space: &mut [i16], output: &mut [i16]
) {
assert_eq!(input.len() * 4, output.len());
assert!(input.len() * 2 <= scratch_space.len());
let scratch_space = &mut scratch_space[..input.len() * 2];
let mut t = [0];
upsample_vertical(input, in_near, in_far, &mut t, scratch_space);
// horizontal upsampling must be done separate for every line
// Otherwise it introduces artifacts that may cause the edge colors
// to appear on the other line.
// Since this is called for two scanlines/widths currently
// splitting the inputs and outputs into half ensures we only handle
// one scanline per iteration
let scratch_half = scratch_space.len() / 2;
let output_half = output.len() / 2;
upsample_horizontal(
&scratch_space[..scratch_half],
&[],
&[],
&mut t,
&mut output[..output_half]
);
upsample_horizontal(
&scratch_space[scratch_half..],
&[],
&[],
&mut t,
&mut output[output_half..]
);
}
pub fn upsample_generic(
input: &[i16], _in_near: &[i16], _in_far: &[i16], _scratch_space: &mut [i16],
output: &mut [i16]
) {
// use nearest sample
let difference = output.len() / input.len();
if difference > 0 {
// nearest neighbour
for (input, chunk_output) in input.iter().zip(output.chunks_exact_mut(difference)) {
chunk_output.iter_mut().for_each(|x| *x = *input);
}
}
}
+577
View File
@@ -0,0 +1,577 @@
/*
* Copyright (c) 2023.
*
* This software is free software;
*
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
*/
use alloc::format;
use core::convert::TryInto;
use core::cmp::min;
use zune_core::colorspace::ColorSpace;
use crate::color_convert::ycbcr_to_grayscale;
use crate::components::{Components, SampleRatios};
use crate::decoder::{ColorConvert16Ptr, MAX_COMPONENTS};
use crate::errors::DecodeErrors;
/// fast 0..255 * 0..255 => 0..255 rounded multiplication
///
/// Borrowed from stb
#[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
#[inline]
fn blinn_8x8(in_val: u8, y: u8) -> u8 {
let t = i32::from(in_val) * i32::from(y) + 128;
return ((t + (t >> 8)) >> 8) as u8;
}
#[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
pub(crate) fn color_convert(
unprocessed: &[&[i16]; MAX_COMPONENTS], color_convert_16: ColorConvert16Ptr,
input_colorspace: ColorSpace, output_colorspace: ColorSpace, output: &mut [u8], width: usize,
padded_width: usize
) -> Result<(), DecodeErrors> {
if input_colorspace.num_components() == 3 && input_colorspace == output_colorspace {
// sort things like RGB to RGB conversion
copy_removing_padding(unprocessed, width, padded_width, output);
return Ok(());
}
if input_colorspace.num_components() == 4 && input_colorspace == output_colorspace {
copy_removing_padding_4x(unprocessed, width, padded_width, output);
return Ok(());
}
// color convert
match (input_colorspace, output_colorspace) {
(ColorSpace::YCbCr | ColorSpace::Luma, ColorSpace::Luma) => {
ycbcr_to_grayscale(unprocessed[0], width, padded_width, output);
}
(
ColorSpace::YCbCr,
ColorSpace::RGB | ColorSpace::RGBA | ColorSpace::BGR | ColorSpace::BGRA
) => {
color_convert_ycbcr(
unprocessed,
width,
padded_width,
output_colorspace,
color_convert_16,
output
);
}
(ColorSpace::YCCK, ColorSpace::RGB) => {
color_convert_ycck_to_rgb::<3>(
unprocessed,
width,
padded_width,
output_colorspace,
color_convert_16,
output
);
}
(ColorSpace::YCCK, ColorSpace::RGBA) => {
color_convert_ycck_to_rgb::<4>(
unprocessed,
width,
padded_width,
output_colorspace,
color_convert_16,
output
);
}
(ColorSpace::CMYK, ColorSpace::RGB) => {
color_convert_cymk_to_rgb::<3>(unprocessed, width, padded_width, output);
}
(ColorSpace::CMYK, ColorSpace::RGBA) => {
color_convert_cymk_to_rgb::<4>(unprocessed, width, padded_width, output);
}
(ColorSpace::MultiBand(n), _) => {
if n.get() != 2 {
return Err(DecodeErrors::Format(format!(
"Unknown multiband sample ({n}), please share sample"
)));
}
copy_removing_padding_generic(
unprocessed,
width,
padded_width,
output,
n.get() as usize
);
}
(ColorSpace::Luma, ColorSpace::RGB) => {
// duplicate the luma channel three times to form RGB
// Note, this may assume the direct conversion
// from luma to RGB is by duplicating
//
// There may be a bit more complex ways
// of doing it but won't get onto it
convert_luma_to_rgb(unprocessed, width, padded_width, output)
}
(ColorSpace::Luma, ColorSpace::RGBA) => {
// duplicate the luma channel three times to form RGB
// add 255 as alpha
// Note, this may assume the direct conversion
// from luma to RGB is by duplicating
//
// There may be a bit more complex ways
// of doing it but won't get onto it
convert_luma_to_rgba(unprocessed, width, padded_width, output)
}
// For the other components we do nothing(currently)
_ => {
let msg = format!(
"Unimplemented colorspace mapping from {input_colorspace:?} to {output_colorspace:?}");
return Err(DecodeErrors::Format(msg));
}
}
Ok(())
}
fn convert_luma_to_rgb(
mcu_block: &[&[i16]; MAX_COMPONENTS], width: usize, padded_width: usize, output: &mut [u8]
) {
for (pix_w, y_w) in output
.chunks_exact_mut(width * 3)
.zip(mcu_block[0].chunks_exact(padded_width))
{
for (pix, c) in pix_w.chunks_exact_mut(3).zip(y_w) {
pix[0] = *c as u8;
pix[1] = *c as u8;
pix[2] = *c as u8;
}
}
}
fn convert_luma_to_rgba(
mcu_block: &[&[i16]; MAX_COMPONENTS], width: usize, padded_width: usize, output: &mut [u8]
) {
for (pix_w, y_w) in output
.chunks_exact_mut(width * 4)
.zip(mcu_block[0].chunks_exact(padded_width))
{
for (pix, c) in pix_w.chunks_exact_mut(4).zip(y_w) {
pix[0] = *c as u8;
pix[1] = *c as u8;
pix[2] = *c as u8;
pix[3] = 255;
}
}
}
/// Copy a block to output removing padding bytes from input
/// if necessary
#[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
fn copy_removing_padding(
mcu_block: &[&[i16]; MAX_COMPONENTS], width: usize, padded_width: usize, output: &mut [u8]
) {
for (((pix_w, c_w), m_w), y_w) in output
.chunks_exact_mut(width * 3)
.zip(mcu_block[0].chunks_exact(padded_width))
.zip(mcu_block[1].chunks_exact(padded_width))
.zip(mcu_block[2].chunks_exact(padded_width))
{
for (((pix, c), y), m) in pix_w.chunks_exact_mut(3).zip(c_w).zip(m_w).zip(y_w) {
pix[0] = *c as u8;
pix[1] = *y as u8;
pix[2] = *m as u8;
}
}
}
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
fn copy_removing_padding_4x(
mcu_block: &[&[i16]; MAX_COMPONENTS], width: usize, padded_width: usize, output: &mut [u8]
) {
for ((((pix_w, c_w), m_w), y_w), k_w) in output
.chunks_exact_mut(width * 4)
.zip(mcu_block[0].chunks_exact(padded_width))
.zip(mcu_block[1].chunks_exact(padded_width))
.zip(mcu_block[2].chunks_exact(padded_width))
.zip(mcu_block[3].chunks_exact(padded_width))
{
for ((((pix, c), y), m), k) in pix_w
.chunks_exact_mut(4)
.zip(c_w)
.zip(m_w)
.zip(y_w)
.zip(k_w)
{
pix[0] = *c as u8;
pix[1] = *y as u8;
pix[2] = *m as u8;
pix[3] = *k as u8;
}
}
}
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
fn copy_removing_padding_generic(
mcu_block: &[&[i16]; MAX_COMPONENTS], width: usize, padded_width: usize, output: &mut [u8],
channels: usize
) {
match channels {
// just do 2 for now
2 => {
for ((pix_w, y_w), k_w) in output
.chunks_exact_mut(width * channels)
.zip(mcu_block[0].chunks_exact(padded_width))
.zip(mcu_block[1].chunks_exact(padded_width))
{
for ((pix, c), k) in pix_w.chunks_exact_mut(2).zip(y_w).zip(k_w) {
pix[0] = *c as u8;
pix[1] = *k as u8;
}
}
}
_ => unreachable!()
}
}
/// Convert YCCK image to rgb
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
fn color_convert_ycck_to_rgb<const NUM_COMPONENTS: usize>(
mcu_block: &[&[i16]; MAX_COMPONENTS], width: usize, padded_width: usize,
output_colorspace: ColorSpace, color_convert_16: ColorConvert16Ptr, output: &mut [u8]
) {
color_convert_ycbcr(
mcu_block,
width,
padded_width,
output_colorspace,
color_convert_16,
output
);
for (pix_w, m_w) in output
.chunks_exact_mut(width * 3)
.zip(mcu_block[3].chunks_exact(padded_width))
{
for (pix, m) in pix_w.chunks_exact_mut(NUM_COMPONENTS).zip(m_w) {
let m = (*m) as u8;
pix[0] = blinn_8x8(255 - pix[0], m);
pix[1] = blinn_8x8(255 - pix[1], m);
pix[2] = blinn_8x8(255 - pix[2], m);
}
}
}
#[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
fn color_convert_cymk_to_rgb<const NUM_COMPONENTS: usize>(
mcu_block: &[&[i16]; MAX_COMPONENTS], width: usize, padded_width: usize, output: &mut [u8]
) {
for ((((pix_w, c_w), m_w), y_w), k_w) in output
.chunks_exact_mut(width * NUM_COMPONENTS)
.zip(mcu_block[0].chunks_exact(padded_width))
.zip(mcu_block[1].chunks_exact(padded_width))
.zip(mcu_block[2].chunks_exact(padded_width))
.zip(mcu_block[3].chunks_exact(padded_width))
{
for ((((pix, c), m), y), k) in pix_w
.chunks_exact_mut(3)
.zip(c_w)
.zip(m_w)
.zip(y_w)
.zip(k_w)
{
let c = *c as u8;
let m = *m as u8;
let y = *y as u8;
let k = *k as u8;
pix[0] = blinn_8x8(c, k);
pix[1] = blinn_8x8(m, k);
pix[2] = blinn_8x8(y, k);
}
}
}
/// Do color-conversion for interleaved MCU
#[allow(
clippy::similar_names,
clippy::too_many_arguments,
clippy::needless_pass_by_value,
clippy::unwrap_used
)]
fn color_convert_ycbcr(
mcu_block: &[&[i16]; MAX_COMPONENTS], width: usize, padded_width: usize,
output_colorspace: ColorSpace, color_convert_16: ColorConvert16Ptr, output: &mut [u8]
) {
let num_components = output_colorspace.num_components();
let stride = width * num_components;
// Allocate temporary buffer for small widths less than 16.
let mut temp = [0; 64];
// We need to chunk per width to ensure we can discard extra values at the end of the width.
// Since the encoder may pad bits to ensure the width is a multiple of 8.
for (((y_width, cb_width), cr_width), out) in mcu_block[0]
.chunks_exact(padded_width)
.zip(mcu_block[1].chunks_exact(padded_width))
.zip(mcu_block[2].chunks_exact(padded_width))
.zip(output.chunks_exact_mut(stride))
{
if width < 16 {
// allocate temporary buffers for the values received from idct
let mut y_out = [0; 16];
let mut cb_out = [0; 16];
let mut cr_out = [0; 16];
// copy those small widths to that buffer
// Use a min with 16 to prevent some panics, see https://github.com/etemesi254/zune-image/issues/331
y_out[0..min(y_width.len(), 16)].copy_from_slice(&y_width[0..min(y_width.len(), 16)]);
cb_out[0..min(cb_width.len(), 16)]
.copy_from_slice(&cb_width[0..min(cb_width.len(), 16)]);
cr_out[0..min(cr_width.len(), 16)]
.copy_from_slice(&cr_width[0..min(cr_width.len(), 16)]);
// we handle widths less than 16 a bit differently, allocating a temporary
// buffer and writing to that and then flushing to the out buffer
// because of the optimizations applied below,
(color_convert_16)(&y_out, &cb_out, &cr_out, &mut temp, &mut 0);
// copy to stride
out[0..width * num_components].copy_from_slice(&temp[0..width * num_components]);
// next
continue;
}
// Chunk in outputs of 16 to pass to color_convert as an array of 16 i16's.
for (((y, cb), cr), out_c) in y_width
.chunks_exact(16)
.zip(cb_width.chunks_exact(16))
.zip(cr_width.chunks_exact(16))
.zip(out.chunks_exact_mut(16 * num_components))
{
(color_convert_16)(
y.try_into().unwrap(),
cb.try_into().unwrap(),
cr.try_into().unwrap(),
out_c,
&mut 0
);
}
//we have more pixels in the end that can't be handled by the main loop.
//move pointer back a little bit to get last 16 bytes,
//color convert, and overwrite
//This means some values will be color converted twice.
for ((y, cb), cr) in y_width[width - 16..]
.chunks_exact(16)
.zip(cb_width[width - 16..].chunks_exact(16))
.zip(cr_width[width - 16..].chunks_exact(16))
.take(1)
{
(color_convert_16)(
y.try_into().unwrap(),
cb.try_into().unwrap(),
cr.try_into().unwrap(),
&mut temp,
&mut 0
);
}
let rem = out[(width - 16) * num_components..]
.chunks_exact_mut(16 * num_components)
.next()
.unwrap();
rem.copy_from_slice(&temp[0..rem.len()]);
}
}
pub(crate) fn upsample(
component: &mut Components, mcu_height: usize, i: usize, upsampler_scratch_space: &mut [i16],
has_vertical_sample: bool
) -> Result<(), DecodeErrors> {
match component.sample_ratio {
SampleRatios::V | SampleRatios::HV => {
/*
When upsampling vertically sampled images, we have a certain problem
which is that we do not have all MCU's decoded, this usually sucks at boundaries
e.g we can't upsample the last mcu row, since the row_down currently doesn't exist
To solve this we need to do two things
1. Carry over coefficients when we lack enough data to upsample
2. Upsample when we have enough data
To achieve (1), we store a previous row, and the current row in components themselves
which will later be used to make (2)
To achieve (2), we take the stored previous row(second last MCU row),
current row(last mcu row) and row down(first row of newly decoded MCU)
and upsample that and store it in first_row_upsample_dest, this contains
up-sampled coefficients for the last for the previous decoded mcu row.
The caller is then expected to process first_row_upsample_dest before processing data
in component.upsample_dest which stores the up-sampled components excluding the last row
*/
let mut dest_start = 0;
let stride_bytes_written = component.width_stride * component.sample_ratio.sample();
if i > 0 {
// Handle the last MCU of the previous row
// This wasn't up-sampled as we didn't have the row_down
// so we do it now
let stride = component.width_stride;
let dest = &mut component.first_row_upsample_dest[0..stride_bytes_written];
// get current row
let row = &component.row[..];
let row_up = &component.row_up[..];
let row_down = &component.raw_coeff[0..stride];
(component.up_sampler)(row, row_up, row_down, upsampler_scratch_space, dest);
}
// we have the Y component width stride.
// this may be higher than the actual width,(2x because vertical sampling)
//
// This will not upsample the last row
// if false, do not upsample.
// set to false on the last row of an mcu
let mut upsample = true;
let stride = component.width_stride * component.vertical_sample;
let stop_offset = component.raw_coeff.len() / component.width_stride;
if component.raw_coeff.len() != stop_offset * stride {
// slice would panic below
return Err(DecodeErrors::FormatStatic(
"Invalid component dimensions, would panic"
));
}
for (pos, curr_row) in component
.raw_coeff
.chunks_exact(component.width_stride)
.enumerate()
{
let mut dest: &mut [i16] = &mut [];
let mut row_up: &[i16] = &[];
// row below current sample
let mut row_down: &[i16] = &[];
// Order of ifs matters
if i == 0 && pos == 0 {
// first IMAGE row, row_up is the same as current row
// row_down is the row below.
row_up = &component.raw_coeff[pos * stride..(pos + 1) * stride];
row_down = &component.raw_coeff[(pos + 1) * stride..(pos + 2) * stride];
} else if i > 0 && pos == 0 {
// first row of a new mcu, previous row was copied so use that
row_up = &component.row[..];
row_down = &component.raw_coeff[(pos + 1) * stride..(pos + 2) * stride];
} else if i == mcu_height.saturating_sub(1) && pos == stop_offset - 1 {
// last IMAGE row, adjust pointer to use previous row and current row
row_up = &component.raw_coeff[(pos - 1) * stride..pos * stride];
row_down = &component.raw_coeff[pos * stride..(pos + 1) * stride];
} else if pos > 0 && pos < stop_offset - 1 {
// other rows, get row up and row down relative to our current row
// ignore last row of each mcu
row_up = &component.raw_coeff[(pos - 1) * stride..pos * stride];
row_down = &component.raw_coeff[(pos + 1) * stride..(pos + 2) * stride];
} else if pos == stop_offset - 1 {
// last MCU in a row
//
// we need a row at the next MCU but we haven't decoded that MCU yet
// so we should save this and when we have the next MCU,
// do the upsampling
// store the current row and previous row in a buffer
let prev_row = &component.raw_coeff[(pos - 1) * stride..pos * stride];
component.row_up.copy_from_slice(prev_row);
component.row.copy_from_slice(curr_row);
upsample = false;
} else {
unreachable!("Uh oh!");
}
if upsample {
dest =
&mut component.upsample_dest[dest_start..dest_start + stride_bytes_written];
dest_start += stride_bytes_written;
}
if upsample {
// upsample
(component.up_sampler)(
curr_row,
row_up,
row_down,
upsampler_scratch_space,
dest
);
}
}
}
SampleRatios::H => {
//assert_eq!(component.raw_coeff.len() * 2, component.upsample_dest.len());
// Before it was an assert, but numerous and numerous and numerous
// bug fixes and ad hoc solutions later, I have now just decided to keep it as a resize
component
.upsample_dest
.resize(component.raw_coeff.len() * 2, 0);
let raw_coeff = &component.raw_coeff;
let dest_coeff = &mut component.upsample_dest;
if has_vertical_sample {
/*
There have been images that have the following configurations.
Component ID:Y HS:2 VS:2 QT:0
Component ID:Cb HS:1 VS:1 QT:1
Component ID:Cr HS:1 VS:2 QT:1
This brings out a nasty case of misaligned sampling factors. Cr will need to save a row because
of the way we process boundaries but Cb won't since Cr is horizontally sampled while Cb is
HV sampled with respect to the image sampling factors.
So during decoding of one MCU, we could only do 7 and not 8 rows, but the SampleRatio::H never had to
save a single line, since it doesn't suffer from boundary issues.
Now this takes care of that, saving the last MCU row in case it will be needed.
We save the previous row before up-sampling this row because the boundary issue is in
the last MCU row of the previous MCU.
PS(cae): I can't add the image to the repo as it is nsfw, but can send if required
*/
let length = component.first_row_upsample_dest.len();
component
.first_row_upsample_dest
.copy_from_slice(&dest_coeff.rchunks_exact(length).next().unwrap());
}
// up-sample each row
for (single_row, output_stride) in raw_coeff
.chunks_exact(component.width_stride)
.zip(dest_coeff.chunks_exact_mut(component.width_stride * 2))
{
// upsample using the fn pointer, should only be H, so no need for
// row up and row down
(component.up_sampler)(single_row, &[], &[], &mut [], output_stride);
}
}
SampleRatios::Generic(h, v) => {
let raw_coeff = &component.raw_coeff;
let dest_coeff = &mut component.upsample_dest;
//let size = component.width_stride.div_ceil(v);
// for (single_row, output_stride) in raw_coeff
// .chunks_exact(size)
// .zip(dest_coeff.chunks_exact_mut(component.width_stride * h))
// {
// (component.up_sampler)(single_row, &[], &[], &mut [], output_stride);
//
// }
for (single_row, output_stride) in raw_coeff
.chunks_exact(component.width_stride)
.zip(dest_coeff.chunks_exact_mut(component.width_stride * h * v))
{
for row in output_stride.chunks_exact_mut(component.width_stride * h) {
(component.up_sampler)(single_row, &[], &[], &mut [], row);
}
}
}
SampleRatios::None => {}
};
Ok(())
}
+1 -1
View File
@@ -21,7 +21,7 @@ RUN rustup target add \
# Install cargo-zigbuild for easier cross-compilation (especially macOS)
RUN curl -sSL https://ziglang.org/download/0.13.0/zig-linux-x86_64-0.13.0.tar.xz | tar -xJ -C /opt \
&& ln -s /opt/zig-linux-x86_64-0.13.0/zig /usr/local/bin/zig
RUN cargo install cargo-zigbuild
RUN cargo install cargo-zigbuild --version 0.21.8
# Configure linkers for cross-compilation
RUN mkdir -p /.cargo
+7 -7
View File
@@ -20,13 +20,13 @@ services:
# Build both targets in parallel
(echo "→ Linux x64" && cargo zigbuild --release --target x86_64-unknown-linux-gnu && cp /build/target/x86_64-unknown-linux-gnu/release/agent-browser /output/agent-browser-linux-x64 && chmod +x /output/agent-browser-linux-x64 && echo "✓ Linux x64 done") &
PID1=$!
PID1=$$!
(echo "→ Linux ARM64" && cargo zigbuild --release --target aarch64-unknown-linux-gnu && cp /build/target/aarch64-unknown-linux-gnu/release/agent-browser /output/agent-browser-linux-arm64 && chmod +x /output/agent-browser-linux-arm64 && echo "✓ Linux ARM64 done") &
PID2=$!
PID2=$$!
# Wait for both to complete
wait $PID1 $PID2
wait $$PID1 $$PID2
echo ""
echo "✓ Linux platforms built successfully!"
@@ -67,8 +67,8 @@ services:
- OUTPUT_NAME=${OUTPUT_NAME:-agent-browser-linux-x64}
command: |
-c '
cargo zigbuild --release --target $TARGET
cp /build/target/$TARGET/release/agent-browser* /output/$OUTPUT_NAME
chmod +x /output/$OUTPUT_NAME 2>/dev/null || true
echo "✓ Built $OUTPUT_NAME"
cargo zigbuild --release --target $$TARGET
cp /build/target/$$TARGET/release/agent-browser* /output/$$OUTPUT_NAME
chmod +x /output/$$OUTPUT_NAME 2>/dev/null || true
echo "✓ Built $$OUTPUT_NAME"
'
+272
View File
@@ -0,0 +1,272 @@
# PRD: CLI Web 数据采集体验优化(以小红书场景为例)
- 文档版本: v0.1
- 状态: Draft
- 作者: Codex
- 日期: 2026-03-04
## 1. 背景与问题
在使用 `agent-browser` CLI 执行「小红书宠物博主采集(100 条)」时,当前流程可完成任务,但存在明显的可用性与稳定性痛点:
1. 网络层可观测性不足,响应体抓取不稳定,需注入脚本劫持。
2. 分页采集依赖手工 `scroll down + wait`,重复劳动且易漏数据。
3. 结构化导出缺少一站式命令,需要 `eval` 二次解析。
4. 页面交互依赖文本选择,页面文案变动后脆弱。
5. 反爬失败时缺少可解释的自动回退策略。
6. 用户对“可抓字段”预期不清(例如搜索接口无联系方式)。
7. 长会话缺少快照与断点续抓机制。
## 2. 目标与非目标
## 2.1 目标
1. 将常见采集链路从“脚本拼接”降为“CLI 原生命令组合”。
2. 让关键动作具备可观测性(日志)和可恢复性(快照/续跑)。
3. 降低站点轻微改版、反爬限制带来的失败率。
## 2.2 非目标
1. 不承诺绕过平台强风控或登录体系。
2. 不在本期实现完整通用爬虫 DSL。
3. 不默认抓取平台未公开展示的隐私字段。
## 3. 目标用户与核心场景
1. 增长/运营: 按关键词采集账号基础数据并导出 CSV。
2. 测试/研发: 复现抓取问题,定位请求失败原因。
3. AI Agent 工作流: 在 CLI 内稳定执行“搜索 -> 翻页 -> 提取 -> 导出”。
## 4. 需求范围与优先级
## 4.1 P0
1. `network capture` 增强模式(可过滤、可落盘 response body)。
2. `scroll-collect` 自动滚动采集(按页数或直到无新增)。
3. `extract` / `extract-to` 结构化导出(JSON/CSV)。
## 4.2 P1
1. 语义选择器与 fallback 链(role/aria/data/text)。
2. 401/403/406 智能回退(页面触发 + 回包监听)。
3. 可抓字段矩阵与二段式采集文档提示。
## 4.3 P2
1. `session snapshot` + `crawl resume` 断点续抓。
## 5. CLI 方案设计
## 5.1 网络捕获增强
命令草案:
```bash
agent-browser network capture --match '/api/sns/web/v1/search/usersearch' --save ./out.ndjson
agent-browser network capture --domain edith.xiaohongshu.com --method POST --save ./xhs_usersearch.ndjson
```
参数:
- `--match <regex>`: 按 URL 正则过滤。
- `--domain <host>`: 按域名过滤。
- `--method <GET|POST|...>`: 按方法过滤。
- `--status <code|range>`: 按状态过滤。
- `--save <path>`: NDJSON 输出文件。
- `--include-body <request|response|both>`: 控制 body 输出范围。
- `--max-body-bytes <n>`: 单条 body 截断阈值。
NDJSON 记录结构:
```json
{
"ts": "2026-03-04T10:00:00.123Z",
"session_id": "sess_abc",
"request_id": "req_123",
"method": "POST",
"url": "https://edith.xiaohongshu.com/api/sns/web/v1/search/usersearch",
"status": 200,
"duration_ms": 312,
"request_headers": {"content-type": "application/json"},
"request_body": "{...}",
"response_headers": {"content-type": "application/json"},
"response_body": "{...}",
"truncated": false
}
```
## 5.2 自动滚动采集
命令草案:
```bash
agent-browser scroll-collect --until no-new-items --max-steps 200 --idle-rounds 3
agent-browser scroll-collect --pages 20 --wait-ms 1200
```
行为:
1. 每轮执行滚动与等待。
2. 基于 DOM 项数量或网络新增请求判断“是否有新增”。
3. 达到停止条件后输出结束原因。
输出示例:
```text
step=1 new_items=15 total_items=15
step=2 new_items=15 total_items=30
...
stop_reason=no-new-items idle_rounds=3 total_items=135
```
## 5.3 结构化提取与导出
命令草案:
```bash
agent-browser extract --from network --match usersearch --fields 'name,fans,note_count,red_id'
agent-browser extract-to --from network --match usersearch --fields 'name,fans,note_count,red_id,url' --format csv --out ./users.csv
```
参数:
- `--from <network|dom|eval>`: 数据源。
- `--match <pattern>`: 来源过滤(URL/事件名)。
- `--query <JMESPath|JSONPath>`: 自定义提取表达式。
- `--fields <a,b,c>`: 字段映射快捷写法。
- `--dedupe-by <field>`: 去重键。
- `--limit <n>`: 限制条数。
- `--format <json|ndjson|csv>`: 输出格式。
- `--out <path>`: 文件输出路径。
## 5.4 语义选择器与回退链
命令草案:
```bash
agent-browser click --selector 'role=tab[name="用户"]' --fallback 'aria=用户,text=用户'
agent-browser find --selector 'data-testid=user-tab' --fallback 'role=tab[name="用户"],text=用户'
```
策略:
1. 主选择器失败后按 fallback 顺序重试。
2. 日志打印每次尝试与失败原因。
## 5.5 反爬失败自动回退
命令草案:
```bash
agent-browser request replay --on-status 401,403,406 --fallback page-action
```
策略:
1. 直接请求失败后自动回退到页面行为触发。
2. 自动复用 UA/Referer/Cookie Jar。
3. 捕获最终有效响应并给出“回退成功/失败”日志。
## 5.6 会话快照与断点续抓
命令草案:
```bash
agent-browser session snapshot save ./snapshots/xhs-20260304.json
agent-browser crawl resume --snapshot ./snapshots/xhs-20260304.json --out ./users.csv
```
快照最小字段:
- 当前 URL
- 关键词/筛选参数
- 已抓 user_id 集合摘要(可哈希分片)
- 分页进度(page/scroll step
- 导出配置(fields/format/out
## 6. 错误码设计(草案)
- `AB_NET_CAPTURE_BODY_UNAVAILABLE` (1001): 响应体不可用(被浏览器策略阻断或已释放)。
- `AB_SCROLL_TIMEOUT_NO_PROGRESS` (1101): 滚动超时且无新增。
- `AB_EXTRACT_QUERY_INVALID` (1201): 提取表达式语法错误。
- `AB_EXTRACT_OUTPUT_FAILED` (1202): 导出失败(权限/路径不可写)。
- `AB_SELECTOR_NOT_FOUND` (1301): 主选择器与 fallback 全部失败。
- `AB_REQUEST_BLOCKED_406` (1406): 请求被风控拦截,且回退链路失败。
- `AB_RESUME_SNAPSHOT_INVALID` (1501): 快照损坏或版本不兼容。
要求:
1. CLI 退出码与错误码可映射。
2. 错误输出提供 `hint`(下一步建议命令)。
## 7. 日志与可观测性
默认人类可读,开启 `--log-format json` 输出结构化日志。
JSON 日志字段:
- `ts`
- `level`
- `session_id`
- `command`
- `event`
- `step`
- `url`
- `status`
- `error_code`
- `message`
- `hint`
示例:
```json
{"ts":"2026-03-04T10:11:22.123Z","level":"INFO","command":"scroll-collect","event":"step","step":12,"new_items":15,"total_items":180}
{"ts":"2026-03-04T10:13:01.001Z","level":"WARN","command":"request replay","event":"fallback","status":406,"message":"direct request blocked, fallback to page-action"}
```
## 8. 文档与帮助信息更新要求
当功能落地时,需要同步更新以下位置(按仓库规范):
1. `cli/src/output.rs``--help`、示例、环境变量)
2. `README.md`(命令选项、样例)
3. `skills/agent-browser/SKILL.md`Agent 工作流)
4. `docs/src/app/`(新增/更新 MDX 页面,表格使用 HTML `<table>`
5. 对应源码内联注释
## 9. 验收用例(首批)
1. `network capture` 能稳定保存目标接口完整 request/response body。
2. 设置 `--max-body-bytes` 后被截断记录带 `truncated=true`
3. `scroll-collect --pages 5` 精确执行 5 轮并退出。
4. `scroll-collect --until no-new-items` 在连续空增量 N 轮后退出。
5. `extract-to ... --format csv` 产出可打开 CSV 且列名正确。
6. `extract --dedupe-by user_id` 去重结果稳定。
7. selector 主规则失败时,fallback 生效并成功点击。
8. 对 406 场景触发自动回退并成功捕获有效响应。
9. 回退失败时返回 `AB_REQUEST_BLOCKED_406` 且提供 hint。
10. `session snapshot save/load` 前后任务可恢复。
11. `crawl resume` 不重复导出已抓 ID。
12. `--log-format json` 日志字段完整,便于机器消费。
## 10. 里程碑建议
1. M11 周): `network capture` + `scroll-collect`
2. M21 周): `extract-to` + selector fallback。
3. M31 周): 406 回退链路 + 文档补全。
4. M41 周): snapshot/resume + 稳定性打磨。
## 11. 风险与缓解
1. 平台策略变化导致规则失效。
缓解: 增加站点适配层与策略开关,保留回退日志。
2. 响应体过大带来内存与 IO 压力。
缓解: 流式写入 NDJSON + 截断阈值。
3. 通用提取表达式学习成本高。
缓解: 提供字段模板与场景 presets。
## 12. 开放问题
1. `extract` 表达式标准优先 JSONPath 还是 JMESPath
2. `session snapshot` 是否需要加密(含 cookie 元信息)?
3. 是否提供站点模板(如 `preset xiaohongshu-user-search`)以降低上手成本?
@@ -0,0 +1,225 @@
# 浏览器自动化攻防方案设计:检测模型与分层控制面
本文聚焦浏览器自动化的攻防方案设计,按两个部分组织:
1. **原理**:风险评分系统如何形成结论
2. **控制面**:如何用分层设计降低风险与波动
本文不包含命令行操作与工程实现步骤。
Turnstile 专题内容见:
[Cloudflare Turnstile 攻防方案设计:系统原理与控制面](https://blog.misonote.com/zh/posts/cloudflare-turnstile-stability-principles/)
---
## 一、原理
### 1.1 风险评分不是单点命中
高风控站点的“是否挑战/是否降权”通常来自多维评分,而不是某一条规则的二元判断。
主要输入维度:
1. **一致性**:同一身份在不同表面是否互相矛盾
2. **稀有性**:低频异常组合是否出现
3. **时序性**:行为时间序列是否呈机械统计特征
4. **执行完整性**:关键链路(挑战脚本、跨域资源、worker)是否被破坏
```mermaid
flowchart LR
A["环境与行为"] --> B["一致性评分"]
A --> C["稀有性评分"]
A --> D["时序评分"]
A --> E["执行完整性评分"]
B --> F["综合风险"]
C --> F
D --> F
E --> F
F --> G{"放行/挑战/限流"}
```
### 1.2 一致性:约束集合而非单点修饰
一致性问题的本质是“同一身份在多个观测面上的约束必须同时成立”。
#### 1.2.1 约束集合示意
可以把身份一致性建模为“约束图”:
```mermaid
flowchart TD
UA["UA 字符串"] --> UACH["UA-CH / userAgentMetadata"]
UA --> LangH["Accept-Language"]
LangH --> LangJS["navigator.language(s)"]
LangJS --> Intl["Intl locale/timeZone"]
Plat["platform"] --> Rend["渲染能力/WebGL"]
Rend --> Win["窗口/屏幕参数"]
UACH --> Plat
```
图中每条边表示“两个表面必须相互一致”,否则会形成冲突分值。
#### 1.2.2 典型冲突类型
- UA 显示平台/版本与 UA-CH 不一致
- `Accept-Language``navigator.languages` 不一致
- `Intl` 时区与偏移/地区推断不一致
- 设备声明与渲染能力组合异常
工程含义:
- 修一个点可能打破另一个点
- 设计顺序应是“先定约束集合,再决定每个表面如何满足约束”
### 1.3 稀有性:组合风险而非单值风险
稀有性来自“低频组合”,其危险性来自共现而非单项。
可以将稀有性理解为“联合分布”偏离:
- 单项偏离:可被容忍
- 多项共现偏离:风险迅速累积
工程含义:
- 目标是减少低频组合在同一会话内叠加
- 目标不是拟合某个固定画像
### 1.4 时序性:统计特征而非行为语义
行为检测通常关注统计分布特征:
- 低方差:动作间隔过于稳定
- 强周期:间隔呈固定节奏
- 强同步:不同类型动作间隔一致
工程含义:
- 行为治理的目标是“分布塑形”(variance/jitter/backoff
- 行为治理不是“添加更多动作”
### 1.5 执行完整性:上游条件
执行完整性属于“系统是否能正确运行”的前置条件。
- challenge 脚本、跨域 iframe、跨域 worker 的语义被破坏时,失败率会显著上升
- 此类失败可能与“是否被识别”为不同类别的问题
工程原则:
> 执行链路保护优先于信号修饰。
### 1.6 反调试执行面:与指纹评分面并行
很多站点并不只依赖指纹评分,还会部署“主动处置型反调试”脚本。
典型路径:
```mermaid
flowchart LR
A["页面启动"] --> B["反调试探测"]
B --> C{"命中?"}
C -->|是| D["close/back/redirect"]
C -->|否| E["继续业务逻辑"]
```
该路径与指纹评分面的关系:
1. 指纹评分决定“挑战/放行/降权”
2. 反调试处置决定“页面是否继续可用”
因此,“页面自关闭”不能直接推断为“指纹被识别”,更常见是反调试链路触发。
---
## 二、控制面(分层设计)
### 2.1 控制面总览
攻防方案可以拆为四层控制面:
1. **启动控制**:治理启动早期显式风险
2. **协议控制**:治理协议层身份一致性
3. **运行时控制**:治理页面脚本可观测表面
4. **行为与会话控制**:治理时序分布与上下文漂移
```mermaid
flowchart LR
A["启动控制"] --> B["协议控制"]
B --> C["运行时控制"]
C --> D["行为与会话控制"]
D --> E["一致性与稳定性"]
```
### 2.2 启动控制
目标:降低会话早期显式风险。
设计约束:
- 只处理高置信度自动化标识
- 避免引入与协议层/运行时层不一致的改动
### 2.3 协议控制
目标:将身份约束集合落实到协议层输出。
设计要点:
- 将 UA 与 UA-CH 视为同一约束集合的不同投影
- 覆盖范围需要与目标(页面/worker/子目标)一致
### 2.4 运行时控制
目标:覆盖高频探测面,同时保证不破坏执行语义。
设计要点:
- 优先治理高频、可解释的探测路径
- 对跨域挑战链路对象设置严格注入边界
#### 2.4.1 反调试脚本治理(以 disable-devtool 类库为例)
反调试脚本通常通过固定启动入口触发(例如 `disable-devtool-auto` 标记)。
可行控制策略:
1. 仅抑制其自动启动入口,避免触发主动处置
2. 不改写通用查询/脚本加载语义,避免影响业务页面
3. 将治理范围限制在高置信度触发点,控制副作用面
这类策略的本质是“执行面隔离”,不是“伪造更多指纹”。
### 2.5 行为与会话控制
目标:塑形时间分布,减少上下文漂移。
设计要点:
- 行为治理以统计分布为目标(variance/jitter/backoff
- 会话治理以一致上下文为目标(避免身份漂移)
### 2.6 挑战场景控制面摘要(Turnstile)
Turnstile 场景下的关键控制面可抽象为:
1. 能力令牌语义:服务端验证、有限时效、单次消费
2. 作用域收缩:`hostname/action/cdata` 收缩滥用空间
3. 执行链路保护:跨域脚本/iframe/worker 语义保护
4. 摩擦与安全分离:clearance 属于体验层,不替代安全决策层
该摘要用于将 Turnstile 纳入统一控制面框架;细节见专题文章。
---
## 三、方案设计优先级
控制面设计通常按以下优先级推进:
1. 执行完整性(保证链路可运行,包含反调试触发面治理)
2. 一致性约束集合(消除跨表面矛盾)
3. 稀有性控制(避免低频组合叠加)
4. 时序分布塑形(降低机械统计特征)
5. 体验优化(降低重复挑战摩擦)
该顺序的含义是先保证“系统正确性”,再优化“稳定性与摩擦”。
@@ -0,0 +1,78 @@
# agent-browser-stealth 新功能发布
这次我们对 `agent-browser-stealth` 做了一次完整升级:
从“标签分组辅助”升级为“可控制、可观测、可编排”的 AI 浏览器控制台。
## 这次升级解决了什么问题
过去插件主要解决会话隔离和分组管理。
在真实自动化场景中,仍有三个核心缺口:
1. 缺少侧边栏直接控页能力
2. 缺少执行过程可观测性(console/network/DOM
3. 缺少可复用流程(workflow/shortcut/schedule
本次发布一次性补齐这三块。
## 新增能力一:侧边栏浏览控制
现在你可以在 side panel 直接完成页面控制:
- `open / back / forward / reload`
- 按 CSS 选择器执行 `click / fill / press`
- 标签页切换与关闭
- 直接运行快捷指令(slash shortcut
这让“流程启动”和“人工微调”可以在同一界面完成。
## 新增能力二:开发者观测能力(Developer Signals
我们新增了调试信号面板,避免黑盒执行:
- 页面 `console` 事件(含 error/warn
- `fetch/xhr` 网络事件
- 命令历史
- DOM 状态快照(文本预览、交互元素、mutation 摘要)
当流程失败时,可以快速判断是页面结构变化、网络问题还是动作配置问题。
## 新增能力三:Workflow 自动化体系
插件现在支持完整自动化闭环:
- 录制:`Start -> Stop -> Save`
- 回放:运行 workflow
- 复用:绑定 slash shortcut
- 调度:`daily / weekly / monthly / yearly`
调度任务会记录 `lastRunAt / nextRunAt`,并把执行结果写入活动流,便于排查和审计。
## 已完成实测:workflow + abs 接管流程
我们验证了一个非常实用的流程:
1. 先用 workflow 进入目标页面
2. 在关键节点停顿(checkpoint
3. 让 AI 使用 `abs` 完成上传/填写等细操作
4. 停在发布前,等待人工确认
这个模式已经在小红书发布场景中跑通:
进入发布页、上传图片、填写标题正文,且默认不自动点击“发布”。
## 兼容性与稳定性改进
这次还修复了几个关键稳定性问题:
- 修复 side panel `activeTab` 空值崩溃
- 增加面板状态归一化兜底,兼容热更新与旧返回结构
- 修复侧边栏渲染的转义问题,避免内容显示异常
## 总结
`agent-browser-stealth` 现在不只是“分组插件”,而是一个真正可用于生产流的浏览器执行层:
- 可控:侧边栏直接控页 + 分段执行
- 可观测:console/network/DOM 全链路可见
- 可复用:workflow + shortcut + schedule
如果你在做 AI 浏览器自动化,这次升级能直接提升流程稳定性与迭代效率。
@@ -0,0 +1,270 @@
# Cloudflare Turnstile 攻防方案设计:系统原理与控制面
本文聚焦 Turnstile 的攻防方案设计:
1. **系统原理**:token 的安全语义、挑战执行链路、风险评分的输入输出
2. **控制面设计**:在不同攻击面下,哪些约束是必要的、哪些约束容易引入副作用
本文不包含命令行操作与工程实现步骤。
---
## 一、系统原理
### 1.1 Turnstile 是“能力令牌”系统
Turnstile 的本质是签发一个短生命周期、单次消费的能力令牌(capability token)。
- **签发端**:浏览器端完成挑战执行后获得 token
- **消费端**:业务服务端通过 Siteverify 验证 token 并决定是否放行
```mermaid
flowchart LR
A["浏览器端挑战执行"] --> B["token"]
B --> C["业务服务端"]
C --> D["Siteverify"]
D --> E{"放行/拒绝"}
```
关键含义:
- 前端任何“通过”状态都不是业务放行条件
- 业务放行条件是“token 被正确消费”
### 1.2 Token 的三条安全语义
token 的安全语义可以抽象为三条约束:
1. **必须服务端验证**:不允许仅以前端回调作为依据
2. **有限时效**token 超过时效窗口即失效
3. **单次消费**:同一 token 重复消费应失败
这三条语义分别封装了三个常见攻击目标:
- 伪通过:绕过服务端验证
- 延迟提交:绕过时效窗口
- 重放/并发:绕过单次消费
### 1.3 挑战执行链路是“跨域执行系统”
Turnstile 的 token 产生依赖多组件协作,且跨域链路占主导:
- `api.js` 脚本
- challenge iframe
- challenge worker
- 跨域资源请求
```mermaid
flowchart TD
A["加载 api.js"] --> B["创建 iframe"]
B --> C["执行 worker"]
C --> D["收集信号 + 风险评估"]
D --> E["签发 token"]
```
该链路的工程含义:
- 任何对跨域脚本/iframe/worker 的语义改写,都可能导致 token 生成失败或质量下降
- token 失败不一定意味着“被识别”,也可能是“链路被破坏”
### 1.4 风险评分:输入不是“真假”,而是“自洽程度”
挑战执行阶段会收集环境与行为信号,形成风险评分。
- **信号输入**:环境一致性(UA/UA-CH、语言/时区、渲染能力、能力暴露)
- **行为输入**:时序分布(方差、周期性、同步性)
风险评分的关键不是“拟合某种固定画像”,而是“同一身份在多表面是否自洽”。
### 1.5 反调试脚本是独立攻击面,不等同于指纹命中
在高风控站点中,常见一类“主动防御脚本”(例如 `disable-devtool`):
1. 页面启动后检测开发者工具、控制台钩子、调试停顿特征
2. 命中后执行主动处置(`window.close` / `history.back` / 跳转错误页)
这条链路的关键点:
- 它是“反调试执行面”,不是“指纹评分面”的简单子集
- 会表现为“页面自关闭/自跳转”,但根因可能是运行时反调试触发
工程上必须将其与“指纹问题”分层处理,否则容易误判。
### 1.6 作用域绑定:hostname / action / cdata
服务端校验时提供用于绑定业务语义的字段:
- `hostname`token 允许的站点作用域
- `action`token 允许的动作作用域
- `cdata`token 允许的上下文作用域
这些字段的作用是“收缩 token 可被滥用的范围”,而不是“提高通过率”。
```mermaid
flowchart LR
A["token"] --> B["hostname 作用域"]
A --> C["action 作用域"]
A --> D["cdata 作用域"]
B --> E["降低站外盗用收益"]
C --> F["降低动作错配收益"]
D --> G["降低跨流程重放收益"]
```
### 1.7 Token 状态机(能力令牌视角)
从能力令牌视角,token 生命周期可抽象为:
```mermaid
stateDiagram-v2
[*] --> Issued: challenge ok
Issued --> Consumed: siteverify ok
Issued --> Expired: time window
Issued --> Rejected: binding mismatch
Issued --> Replayed: reused
Replayed --> Rejected
Expired --> Rejected
Consumed --> [*]
```
设计目标是让“非法路径”快速失败,并且失败类型可被服务端语义区分。
### 1.8 攻击树(高层)
Turnstile 的主要攻击目标可以抽象为:
```mermaid
flowchart TD
A["绕过业务动作门禁"] --> B["伪造或跳过服务端验证"]
A --> C["重放 token"]
A --> D["扩大 token 作用域"]
A --> E["破坏挑战执行以制造降级路径"]
C --> C1["并发提交"]
C --> C2["延迟提交"]
D --> D1["Any Hostname"]
D --> D2["action/cdata 缺失"]
```
该攻击树强调设计重点:
- 安全决策必须在服务端闭环
- token 必须被作用域收缩并按语义消费
---
## 二、控制面设计(攻防视角)
### 2.1 控制面分层
Turnstile 防线可以分为四层控制面:
1. **挑战执行控制**:保证脚本/iframe/worker 跨域链路完整
2. **服务端消费控制**:保证 token 的语义被正确消费
3. **作用域控制**:收缩 `hostname/action/cdata` 的可用范围
4. **摩擦控制**:clearance 用于降低挑战摩擦(不作为安全决策依据)
```mermaid
flowchart LR
A["挑战执行控制"] --> E["token 可生成"]
A --> F["token 质量"]
B["服务端消费控制"] --> G["安全决策闭环"]
C["作用域控制"] --> H["滥用收益收缩"]
D["摩擦控制"] --> I["挑战频率下降"]
```
### 2.2 挑战执行控制:跨域语义保护优先
挑战执行链路对跨域执行语义高度敏感。
原则:
- 跨域脚本/iframe/worker 避免语义改写
- 所有指纹修饰必须先满足“不破坏挑战执行”这一硬约束
该原则的工程含义:
- “执行完整性”是上游条件
- “信号修饰”是下游优化
补充:反调试脚本属于执行控制的一部分。
- 若页面含有主动防御脚本(如 `disable-devtool-auto` 入口),需要把“自动触发入口”与“正常业务脚本”分开处理
- 最小化目标是阻断其自动化处置链路,避免误伤挑战脚本/跨域组件
- 处理边界应是“只抑制高风险触发入口,不重写通用 DOM 语义”
### 2.3 服务端消费控制:把 token 当作能力消费
服务端消费控制的设计关键在于“放行条件定义”,而不是“接口调用细节”。
放行条件应体现三类约束:
- 真实性:校验 `success`
- 作用域:校验 `hostname`
- 语义绑定:校验 `action/cdata`
并且必须贯彻 token 的两个安全语义:
- 时效性:过期拒绝
- 单次性:重放拒绝
从攻防角度,该层解决的是“绕过与重放”。
### 2.4 作用域控制:Hostname Management 与 Any Hostname
Hostname 管理解决“站外盗用”的攻击面。
- 启用 Hostname Management:收缩 token 可用站点范围
- 启用 Any Hostname:扩大 token 可用站点范围
设计结论:
- Any Hostname 不是“更灵活”,而是“扩大攻击面”,必须用更强的服务端约束做补偿控制(来源域白名单 + 业务绑定)。
### 2.5 摩擦控制:Pre-clearance 与 cf_clearance 的边界
Pre-clearance 通过后可产生 clearance,用于后续 WAF 挑战联动。
边界定义:
- clearance 用于体验层(降低重复挑战摩擦)
- Siteverify 用于安全决策层(业务放行依据)
将两者混用会引入“体验信号替代安全信号”的设计缺陷。
### 2.6 高对抗场景:代理池与设备关联
在代理池与分布式滥用场景中,单一 IP 维度约束容易失效。
设计方向是引入更稳定的关联维度(例如设备级 ephemeral id),用于聚类与阈值策略。
该层属于平台能力与业务风控的交界:
- 平台提供关联信号
- 业务定义动作分层、阈值与处置策略
---
## 三、方案设计优先级
Turnstile 攻防设计通常按以下优先级推进:
1. 服务端消费语义闭环(真实性 + 作用域 + 绑定 + 单次性 + 时效性)
2. 挑战执行链路完整性(跨域语义保护 + 反调试触发面治理)
3. 信号一致性(减少跨字段矛盾)
4. 行为时序(降低机械分布)
5. 体验优化(clearance 等摩擦控制)
该顺序的含义是先定义“正确的安全决策”,再优化“挑战摩擦与通过率波动”。
---
## 官方参考(概念与配置)
- Widgets: <https://developers.cloudflare.com/turnstile/concepts/widget/>
- Widget configurations: <https://developers.cloudflare.com/turnstile/get-started/client-side-rendering/widget-configurations/>
- Server-side validation: <https://developers.cloudflare.com/turnstile/get-started/server-side-validation/>
- CSP: <https://developers.cloudflare.com/turnstile/reference/content-security-policy/>
- Hostname management: <https://developers.cloudflare.com/turnstile/additional-configuration/hostname-management/>
- Any Hostname: <https://developers.cloudflare.com/turnstile/additional-configuration/hostname-management/any-hostname/>
- Pre-clearance: <https://developers.cloudflare.com/turnstile/additional-configuration/hostname-management/pre-clearance/>
- Cloudflare clearance: <https://developers.cloudflare.com/cloudflare-challenges/concepts/clearance/>
- Ephemeral IDs: <https://developers.cloudflare.com/turnstile/additional-configuration/ephemeral-id/>
@@ -0,0 +1,97 @@
# agent-browser 与 agent-browser-stealth:能力差异与选型
本文给出 `agent-browser``agent-browser-stealth` 的技术差异、适用场景和升级验证步骤。
项目地址:[leeguooooo/agent-browser](https://github.com/leeguooooo/agent-browser)
---
## 1. 定位差异
- `agent-browser`:标准浏览器自动化能力
- `agent-browser-stealth`:在标准自动化能力基础上,增加反检测与高风控场景稳定性能力
---
## 2. 核心能力对比
| 维度 | agent-browser | agent-browser-stealth |
| --- | --- | --- |
| 自动化基础能力 | 支持 | 支持 |
| 指纹一致性治理 | 基础 | 多层(launch/CDP/init-script |
| 高风控站点稳定性 | 一般 | 更高 |
| 会话连续性(附着现有浏览器) | 支持 | 支持,默认附着策略更明确 |
| Cloudflare/Turnstile 回归工具 | 无专用脚本 | `check:turnstile-testkey` |
---
## 3. Cloudflare/Turnstile 相关能力(v0.15.2-fork.2+
### 3.1 挑战链路保护
- 同源 worker 注入保留
- 跨域 challenge worker 不做注入改写
- 降低 challenge worker 执行异常概率
### 3.2 导航等待策略
`open/navigate` 支持:
- `--wait-until load`
- `--wait-until domcontentloaded`
- `--wait-until networkidle`
挑战页建议优先 `domcontentloaded`,减少 `load` 阶段超时误判。
### 3.3 确定性回归
提供官方 test key 回归脚本:
```bash
pnpm run check:turnstile-testkey
```
通过特征:输出 `XXXX.DUMMY.TOKEN.XXXX`
---
## 4. 适用场景
优先使用 `agent-browser-stealth` 的场景:
1. 目标站点存在挑战页/验证码/限流
2. 自动化链路对稳定性要求高
3. 需要长期回归验证与版本门禁
使用 `agent-browser` 的场景:
1. 低风控站点
2. 以基础自动化能力验证为主
---
## 5. 升级验证步骤
```bash
# 1) 检查版本
agent-browser -V
# 2) 关闭旧 daemon,避免版本漂移
agent-browser --session default close
# 3) 运行确定性回归
pnpm run check:turnstile-testkey
# 4) 可选:真实站点回归
agent-browser --wait-until domcontentloaded open https://www.anyviewer.com/cloudflare.html
```
如果启用域名白名单(`AGENT_BROWSER_ALLOWED_DOMAINS`),需包含 `challenges.cloudflare.com`
---
## 6. 结论
`agent-browser-stealth` 适用于高风控与稳定性敏感场景;`agent-browser` 适用于标准自动化场景。
选型建议按目标站点风控强度与回归要求决定。
+2 -2
View File
@@ -14,9 +14,9 @@ const DEFAULT_MODEL = "anthropic/claude-haiku-4.5";
const SYSTEM_PROMPT = `You are a helpful documentation assistant for agent-browser, a headless browser automation CLI designed for AI agents.
GitHub repository: https://github.com/vercel-labs/agent-browser
GitHub repository: https://github.com/leeguooooo/agent-browser
Documentation: https://agent-browser.dev
npm package: agent-browser
npm package: agent-browser-stealth
You have access to the full agent-browser documentation via the bash and readFile tools. The docs are available as markdown files in the /workspace/ directory.
+153 -22
View File
@@ -1,11 +1,25 @@
import { pageMetadata } from "@/lib/page-metadata"
import { pageMetadata } from '@/lib/page-metadata';
export const metadata = pageMetadata("cdp-mode")
export const metadata = pageMetadata('cdp-mode');
# CDP Mode
Connect to an existing browser via Chrome DevTools Protocol:
Default behavior in this fork: when `--cdp` is omitted, agent-browser targets the managed automation browser on `localhost:9333`. If `:9333` is unavailable, it auto-starts Chrome with the persistent profile `~/.agent-browser/chrome-bot-profile` and retries the CDP connection.
If you want to pre-start that managed browser explicitly, run:
```bash
agent-browser start
abs start
```
Project policy:
- `--profile` / `AGENT_BROWSER_PROFILE` are forbidden
- `--channel` / `AGENT_BROWSER_CHANNEL` are forbidden
```bash
# Start Chrome with: google-chrome --remote-debugging-port=9222
@@ -52,7 +66,7 @@ AGENT_BROWSER_AUTO_CONNECT=1 agent-browser snapshot
Auto-connect discovers Chrome by:
1. Reading Chrome's `DevToolsActivePort` file from the default user data directory
2. Falling back to probing common debugging ports (9222, 9229)
2. Falling back to probing common debugging ports (9222, 9229, 9333)
This is useful when:
@@ -60,6 +74,8 @@ This is useful when:
- You want a zero-configuration connection to your existing browser
- You don't want to track which port Chrome is using
Use this mode only when you intentionally want to attach to an existing manual browser session. Recent Chrome builds may display a permission prompt before allowing remote debugging access to that session.
## Color scheme
Playwright overrides the browser's color scheme to `light` by default when connecting via CDP. Use `--color-scheme` to set a persistent preference:
@@ -75,6 +91,35 @@ Or set it globally via config or environment variable:
AGENT_BROWSER_COLOR_SCHEME=dark agent-browser --cdp 9222 open https://example.com
```
## Stealth behavior
`--stealth` is enabled by default across connection modes, but capabilities depend on how you connect:
<table>
<thead>
<tr>
<th>Connection type</th>
<th>Stealth capabilities</th>
</tr>
</thead>
<tbody>
<tr>
<td>Local launch</td>
<td>Chromium launch args + context init scripts</td>
</tr>
<tr>
<td>CDP / auto-connect</td>
<td>Context init scripts</td>
</tr>
<tr>
<td>Cloud providers</td>
<td>Context init scripts (Kernel may also apply provider-managed stealth)</td>
</tr>
</tbody>
</table>
Use `--debug` to print the active connection type and applied stealth capabilities.
## Use cases
This enables control of:
@@ -89,27 +134,113 @@ This enables control of:
<table>
<thead>
<tr><th>Option</th><th>Description</th></tr>
<tr>
<th>Option</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr><td><code>--session &lt;name&gt;</code></td><td>Use isolated session</td></tr>
<tr><td><code>--profile &lt;path&gt;</code></td><td>Persistent browser profile directory</td></tr>
<tr><td><code>-p &lt;provider&gt;</code></td><td>Cloud browser provider (<code>browserbase</code>, <code>browseruse</code>, <code>kernel</code>)</td></tr>
<tr><td><code>--headers &lt;json&gt;</code></td><td>HTTP headers scoped to origin</td></tr>
<tr><td><code>--executable-path</code></td><td>Custom browser executable</td></tr>
<tr><td><code>--args &lt;args&gt;</code></td><td>Browser launch args (comma-separated)</td></tr>
<tr><td><code>--user-agent &lt;ua&gt;</code></td><td>Custom User-Agent string</td></tr>
<tr><td><code>--proxy &lt;url&gt;</code></td><td>Proxy server URL</td></tr>
<tr><td><code>--proxy-bypass &lt;hosts&gt;</code></td><td>Hosts to bypass proxy</td></tr>
<tr><td><code>--json</code></td><td>JSON output for scripts</td></tr>
<tr><td><code>--full, -f</code></td><td>Full page screenshot</td></tr>
<tr><td><code>--name, -n</code></td><td>Locator name filter</td></tr>
<tr><td><code>--exact</code></td><td>Exact text match</td></tr>
<tr><td><code>--headed</code></td><td>Show browser window</td></tr>
<tr><td><code>{"--cdp <port|url>"}</code></td><td>CDP connection (port or WebSocket URL)</td></tr>
<tr><td><code>--auto-connect</code></td><td>Auto-discover and connect to running Chrome</td></tr>
<tr><td><code>--color-scheme &lt;scheme&gt;</code></td><td>Persistent color scheme (<code>dark</code>, <code>light</code>, <code>no-preference</code>)</td></tr>
<tr><td><code>--debug</code></td><td>Debug output</td></tr>
<tr>
<td>
<code>-p &lt;provider&gt;</code>
</td>
<td>
Cloud browser provider (<code>browserbase</code>, <code>browseruse</code>,{' '}
<code>kernel</code>)
</td>
</tr>
<tr>
<td>
<code>--headers &lt;json&gt;</code>
</td>
<td>HTTP headers scoped to origin</td>
</tr>
<tr>
<td>
<code>--executable-path</code>
</td>
<td>Custom browser executable</td>
</tr>
<tr>
<td>
<code>--args &lt;args&gt;</code>
</td>
<td>Browser launch args (comma-separated)</td>
</tr>
<tr>
<td>
<code>--user-agent &lt;ua&gt;</code>
</td>
<td>Custom User-Agent string</td>
</tr>
<tr>
<td>
<code>--proxy &lt;url&gt;</code>
</td>
<td>Proxy server URL</td>
</tr>
<tr>
<td>
<code>--proxy-bypass &lt;hosts&gt;</code>
</td>
<td>Hosts to bypass proxy</td>
</tr>
<tr>
<td>
<code>--json</code>
</td>
<td>JSON output for scripts</td>
</tr>
<tr>
<td>
<code>--full, -f</code>
</td>
<td>Full page screenshot</td>
</tr>
<tr>
<td>
<code>--name, -n</code>
</td>
<td>Locator name filter</td>
</tr>
<tr>
<td>
<code>--exact</code>
</td>
<td>Exact text match</td>
</tr>
<tr>
<td>
<code>--headed</code>
</td>
<td>Show browser window</td>
</tr>
<tr>
<td>
<code>{'--cdp <port|url>'}</code>
</td>
<td>CDP connection (port or WebSocket URL)</td>
</tr>
<tr>
<td>
<code>--auto-connect</code>
</td>
<td>Auto-discover and connect to running Chrome</td>
</tr>
<tr>
<td>
<code>--color-scheme &lt;scheme&gt;</code>
</td>
<td>
Persistent color scheme (<code>dark</code>, <code>light</code>, <code>no-preference</code>)
</td>
</tr>
<tr>
<td>
<code>--debug</code>
</td>
<td>Debug output</td>
</tr>
</tbody>
</table>
+89 -61
View File
@@ -1,19 +1,23 @@
import { pageMetadata } from "@/lib/page-metadata"
import { pageMetadata } from '@/lib/page-metadata';
export const metadata = pageMetadata("commands")
export const metadata = pageMetadata('commands');
# Commands
Executable aliases: `agent-browser`, `agent-browser-stealth`, `abs`.
## Core
```bash
agent-browser start # Start/reuse managed browser on localhost:9333
agent-browser open <url> # Navigate (aliases: goto, navigate)
agent-browser --risk-mode block open <url> # Block when verification/captcha interstitial is detected
agent-browser click <sel> # Click element (--new-tab to open in new tab)
agent-browser dblclick <sel> # Double-click
agent-browser fill <sel> <text> # Clear and fill
agent-browser type <sel> <text> # Type into element
agent-browser type <sel> <text> [--delay <ms>] # Type into element
agent-browser press <key> # Press key (Enter, Tab, Control+a) (alias: key)
agent-browser keyboard type <text> # Type at current focus (no selector needed)
agent-browser keyboard type <text> [--delay <ms>] # Type at current focus (no selector needed)
agent-browser keyboard inserttext <text> # Insert text without key events
agent-browser keydown <key> # Hold key down
agent-browser keyup <key> # Release key
@@ -32,9 +36,17 @@ agent-browser pdf <path> # Save page as PDF
agent-browser snapshot # Accessibility tree with refs
agent-browser eval <js> # Run JavaScript
agent-browser connect <port|url> # Connect to browser via CDP
agent-browser doctor # Diagnose CDP + sourceURL + tab-group plugin health
agent-browser --version # Show CLI version
agent-browser close # Close browser (aliases: quit, exit)
```
Fork builds print dual-version metadata with `--version`:
```bash
agent-browser 0.14.0-fork.1 (upstream 0.14.0, fork 1)
```
## Get info
```bash
@@ -95,6 +107,7 @@ agent-browser find nth 2 ".card" hover
```bash
agent-browser wait <selector> # Wait for element
agent-browser wait <ms> # Wait for time
agent-browser wait 2000-5000 # Random wait between 2-5 seconds
agent-browser wait --text "Welcome" # Wait for text
agent-browser wait --url "**/dash" # Wait for URL pattern
agent-browser wait --load networkidle # Wait for load state
@@ -102,6 +115,16 @@ agent-browser wait --fn "condition" # Wait for JS condition
agent-browser wait --download [path] # Wait for download
```
## Risk Mode
Control how `open`/`navigate` handles verification or captcha interstitials:
```bash
agent-browser --risk-mode warn open https://example.com # default: wait for auto-clear, then retry/warn with riskSignals
agent-browser --risk-mode block open https://example.com # fail fast on detection
agent-browser --risk-mode off open https://example.com # disable detection/retry
```
## Downloads
```bash
@@ -111,6 +134,35 @@ agent-browser wait --download [path] # Wait for any download to complete
Use `--download-path <dir>` (or `AGENT_BROWSER_DOWNLOAD_PATH` env) to set a default download directory. Without it, downloads go to a temporary directory that is deleted when the browser closes.
## Tab grouping
```bash
agent-browser open https://example.com
# CDP mode groups tabs when tab-group plugin is installed
# Override the default group title
agent-browser --tab-group "My Agent Group" open https://example.com
```
CDP mode uses a browser extension handshake to group tabs.
- Extension available: tabs are grouped by `session`.
- Extension missing/unavailable: silent no-op (commands still succeed).
- Default titles:
- `default` session: `Agent Browser Stealth`
- non-default: `Agent Browser Stealth • <session>`
- Extension side panel (`agent-browser-stealth`) also provides:
- Session window isolation and deterministic group colors.
- Browser controls (`open`, `back`, `forward`, `reload`, `click/fill/press` by CSS selector).
- Developer signals (console warnings/errors, fetch/xhr network events, DOM snapshot capture).
- Workflow automation (record actions, run workflows, assign slash shortcuts, and create schedules).
- `Keep Only This`, `Focus`, `Clean Empty Groups` quick actions.
- Toggle switches for strict isolation / activation guard / auto-clean.
- Session allowlist editing (domain fallback to `about:blank` when violated).
- Download routing to `agent-browser-stealth/<session>/...`.
- Use `--tab-group` / `AGENT_BROWSER_TAB_GROUP` for base title.
- Use `AGENT_BROWSER_TAB_GROUP_PLUGIN_ID` (or `--tab-group-plugin-id`) to override expected extension ID.
## Mouse
```bash
@@ -153,6 +205,14 @@ agent-browser storage local clear # Clear all
agent-browser storage session # Same for sessionStorage
```
For `cookies set`, use one of these patterns:
- `--url <url>`
- `--domain <domain> --path <path>`
- omit all three to scope from the current page URL
When `--url` is omitted, `--domain` and `--path` must be provided together.
## Network
```bash
@@ -199,49 +259,8 @@ agent-browser console --clear # Clear console log
agent-browser errors # View page errors
agent-browser errors --clear # Clear error log
agent-browser highlight <sel> # Highlight element
```
## Auth vault
```bash
agent-browser auth save <name> [opts] # Save auth profile
agent-browser auth login <name> # Login using saved credentials
agent-browser auth list # List saved profiles (names and URLs only)
agent-browser auth show <name> # Show profile metadata (no passwords)
agent-browser auth delete <name> # Delete a saved profile
```
Save options:
- `--url <url>` -- login page URL (required)
- `--username <user>` -- username (required)
- `--password <pass>` -- password (required unless `--password-stdin`)
- `--password-stdin` -- read password from stdin (recommended to avoid shell history exposure)
- `--username-selector <sel>` -- custom CSS selector for username field
- `--password-selector <sel>` -- custom CSS selector for password field
- `--submit-selector <sel>` -- custom CSS selector for submit button
```bash
echo "pass" | agent-browser auth save github --url https://github.com/login --username user --password-stdin
agent-browser auth login github
agent-browser auth list
```
## Confirmation
When `--confirm-actions` is set, certain action categories return a `confirmation_required` response instead of executing immediately. Use `confirm` or `deny` to approve or reject the action.
```bash
agent-browser confirm <confirmation-id> # Approve a pending action
agent-browser deny <confirmation-id> # Deny a pending action
```
Pending confirmations auto-deny after 60 seconds.
```bash
agent-browser --confirm-actions eval,download eval "document.title"
# Returns confirmation_required with ID
agent-browser confirm c_8f3a1234
agent-browser doctor # Diagnose CDP + sourceURL + plugin handshake status
pnpm run check:turnstile-testkey # Deterministic Turnstile smoke check (official test key)
```
## State management
@@ -261,7 +280,21 @@ agent-browser state clean --older-than <days> # Delete old states
```bash
agent-browser session # Show current session name
agent-browser session list # List active sessions
agent-browser session list # List active sessions (auto-cleans stale entries)
agent-browser --parallel worker-a open https://example.com # Isolated runtime for parallel AI tasks
```
When a default-session command starts, all non-default daemon sessions are reaped to avoid stale daemon reuse.
## Daemon lifetime
By default, daemon processes auto-shutdown after 10 minutes of inactivity.
Use `--resident` when you need a long-running daemon:
```bash
agent-browser --resident open https://example.com
agent-browser close
```
## Navigation
@@ -275,9 +308,7 @@ agent-browser reload # Reload page
## Global options
```bash
--session <name> # Isolated browser session
--session-name <name> # Auto-save/restore session state (cookies, localStorage)
--profile <path> # Persistent browser profile directory
--session-name <name> # Auto-save/restore session state (defaults to "default" in non-parallel mode)
--state <path> # Load storage state from JSON file
--headers <json> # HTTP headers scoped to URL's origin
--executable-path <path> # Custom browser executable
@@ -288,6 +319,7 @@ agent-browser reload # Reload page
--proxy-bypass <hosts> # Hosts to bypass proxy
--ignore-https-errors # Ignore HTTPS certificate errors
--allow-file-access # Allow file:// URLs to access local files (Chromium only)
--stealth # Stealth mode (always on by default)
-p, --provider <name> # Browser provider (ios, browserbase, kernel, browseruse)
--device <name> # iOS device name (e.g., "iPhone 15 Pro")
--json # JSON output (for scripts)
@@ -296,16 +328,12 @@ agent-browser reload # Reload page
--headed # Show browser window (not headless)
--cdp <port|url> # Connect via Chrome DevTools Protocol (port or WebSocket URL)
--auto-connect # Auto-discover and connect to running Chrome
--color-scheme <scheme> # Color scheme: dark, light, no-preference
--download-path <path> # Default download directory
--content-boundaries # Wrap page output in boundary markers for LLM safety
--max-output <chars> # Truncate page output to N characters
--allowed-domains <list> # Comma-separated allowed domain patterns
--action-policy <path> # Path to action policy JSON file
--confirm-actions <list> # Action categories requiring confirmation
--confirm-interactive # Interactive confirmation prompts (auto-denies if stdin is not a TTY)
--config <path> # Use a custom config file
--debug # Debug output
--tab-group <name> # Base title for agent tab groups (CDP plugin mode)
--tab-group-plugin-id <id> # Expected extension ID for tab-group handshake
--parallel <name> # Isolated runtime channel for parallel AI runs (maps to parallel-<name>; reaped when default session starts)
--resident # Keep daemon running; disable 10-minute idle auto-shutdown
--wait-until <mode> # Navigation wait strategy for open/navigate (load, domcontentloaded, networkidle)
--debug # Debug output (includes stealth connection type + capabilities)
```
## Command chaining
+468 -67
View File
@@ -1,24 +1,52 @@
import { pageMetadata } from "@/lib/page-metadata"
import { pageMetadata } from '@/lib/page-metadata';
export const metadata = pageMetadata("configuration")
export const metadata = pageMetadata('configuration');
# Configuration
Create an `agent-browser.json` file to set persistent defaults instead of repeating flags on every command.
In this fork, default launch behavior uses a dedicated automation browser on `localhost:9333` (CDP). If `:9333` is unavailable, agent-browser auto-starts Chrome with the persistent profile `~/.agent-browser/chrome-bot-profile` and retries the connection.
## Config File Locations
agent-browser checks two locations, merged in priority order:
<table>
<thead>
<tr><th>Priority</th><th>Location</th><th>Scope</th></tr>
<tr>
<th>Priority</th>
<th>Location</th>
<th>Scope</th>
</tr>
</thead>
<tbody>
<tr><td>1 (lowest)</td><td><code>~/.agent-browser/config.json</code></td><td>User-level defaults</td></tr>
<tr><td>2</td><td><code>./agent-browser.json</code></td><td>Project-level overrides</td></tr>
<tr><td>3</td><td><code>AGENT_BROWSER_*</code> env vars</td><td>Override config values</td></tr>
<tr><td>4 (highest)</td><td>CLI flags</td><td>Override everything</td></tr>
<tr>
<td>1 (lowest)</td>
<td>
<code>~/.agent-browser/config.json</code>
</td>
<td>User-level defaults</td>
</tr>
<tr>
<td>2</td>
<td>
<code>./agent-browser.json</code>
</td>
<td>Project-level overrides</td>
</tr>
<tr>
<td>3</td>
<td>
<code>AGENT_BROWSER_*</code> env vars
</td>
<td>Override config values</td>
</tr>
<tr>
<td>4 (highest)</td>
<td>CLI flags</td>
<td>Override everything</td>
</tr>
</tbody>
</table>
@@ -37,7 +65,6 @@ AGENT_BROWSER_CONFIG=./ci-config.json agent-browser open example.com
{
"headed": true,
"proxy": "http://localhost:8080",
"profile": "./browser-data",
"userAgent": "my-agent/1.0",
"ignoreHttpsErrors": true
}
@@ -45,46 +72,268 @@ AGENT_BROWSER_CONFIG=./ci-config.json agent-browser open example.com
## All Options
Every CLI flag can be set in the config file using its camelCase equivalent:
Most CLI flags can be set in the config file using their camelCase equivalents (`--resident` is CLI-only):
<table>
<thead>
<tr><th>Config Key</th><th>CLI Flag</th><th>Type</th></tr>
<tr>
<th>Config Key</th>
<th>CLI Flag</th>
<th>Type</th>
</tr>
</thead>
<tbody>
<tr><td><code>headed</code></td><td><code>--headed</code></td><td>boolean</td></tr>
<tr><td><code>json</code></td><td><code>--json</code></td><td>boolean</td></tr>
<tr><td><code>full</code></td><td><code>--full, -f</code></td><td>boolean</td></tr>
<tr><td><code>debug</code></td><td><code>--debug</code></td><td>boolean</td></tr>
<tr><td><code>session</code></td><td><code>--session</code></td><td>string</td></tr>
<tr><td><code>sessionName</code></td><td><code>--session-name</code></td><td>string</td></tr>
<tr><td><code>executablePath</code></td><td><code>--executable-path</code></td><td>string</td></tr>
<tr><td><code>extensions</code></td><td><code>--extension</code></td><td>string[]</td></tr>
<tr><td><code>profile</code></td><td><code>--profile</code></td><td>string</td></tr>
<tr><td><code>state</code></td><td><code>--state</code></td><td>string</td></tr>
<tr><td><code>proxy</code></td><td><code>--proxy</code></td><td>string</td></tr>
<tr><td><code>proxyBypass</code></td><td><code>--proxy-bypass</code></td><td>string</td></tr>
<tr><td><code>args</code></td><td><code>--args</code></td><td>string</td></tr>
<tr><td><code>userAgent</code></td><td><code>--user-agent</code></td><td>string</td></tr>
<tr><td><code>provider</code></td><td><code>-p, --provider</code></td><td>string</td></tr>
<tr><td><code>device</code></td><td><code>--device</code></td><td>string</td></tr>
<tr><td><code>ignoreHttpsErrors</code></td><td><code>--ignore-https-errors</code></td><td>boolean</td></tr>
<tr><td><code>allowFileAccess</code></td><td><code>--allow-file-access</code></td><td>boolean</td></tr>
<tr><td><code>cdp</code></td><td><code>--cdp</code></td><td>string</td></tr>
<tr><td><code>autoConnect</code></td><td><code>--auto-connect</code></td><td>boolean</td></tr>
<tr><td><code>colorScheme</code></td><td><code>--color-scheme</code></td><td>string (<code>dark</code>, <code>light</code>, <code>no-preference</code>)</td></tr>
<tr><td><code>downloadPath</code></td><td><code>--download-path</code></td><td>string</td></tr>
<tr><td><code>contentBoundaries</code></td><td><code>--content-boundaries</code></td><td>boolean</td></tr>
<tr><td><code>maxOutput</code></td><td><code>--max-output</code></td><td>number</td></tr>
<tr><td><code>allowedDomains</code></td><td><code>--allowed-domains</code></td><td>string[]</td></tr>
<tr><td><code>actionPolicy</code></td><td><code>--action-policy</code></td><td>string</td></tr>
<tr><td><code>confirmActions</code></td><td><code>--confirm-actions</code></td><td>string</td></tr>
<tr><td><code>confirmInteractive</code></td><td><code>--confirm-interactive</code></td><td>boolean</td></tr>
<tr><td><code>native</code></td><td><code>--native</code></td><td>boolean (experimental)</td></tr>
<tr><td><code>headers</code></td><td><code>--headers</code></td><td>string (JSON)</td></tr>
<tr>
<td>
<code>headed</code>
</td>
<td>
<code>--headed</code>
</td>
<td>boolean</td>
</tr>
<tr>
<td>
<code>json</code>
</td>
<td>
<code>--json</code>
</td>
<td>boolean</td>
</tr>
<tr>
<td>
<code>full</code>
</td>
<td>
<code>--full, -f</code>
</td>
<td>boolean</td>
</tr>
<tr>
<td>
<code>debug</code>
</td>
<td>
<code>--debug</code>
</td>
<td>boolean</td>
</tr>
<tr>
<td>
<code>sessionName</code>
</td>
<td>
<code>--session-name</code>
</td>
<td>string</td>
</tr>
<tr>
<td>
<code>parallel</code>
</td>
<td>
<code>--parallel</code>
</td>
<td>string (isolated runtime channel name)</td>
</tr>
<tr>
<td>
<code>executablePath</code>
</td>
<td>
<code>--executable-path</code>
</td>
<td>string</td>
</tr>
<tr>
<td>
<code>extensions</code>
</td>
<td>
<code>--extension</code>
</td>
<td>string[]</td>
</tr>
<tr>
<td>
<code>state</code>
</td>
<td>
<code>--state</code>
</td>
<td>string</td>
</tr>
<tr>
<td>
<code>proxy</code>
</td>
<td>
<code>--proxy</code>
</td>
<td>string</td>
</tr>
<tr>
<td>
<code>proxyBypass</code>
</td>
<td>
<code>--proxy-bypass</code>
</td>
<td>string</td>
</tr>
<tr>
<td>
<code>args</code>
</td>
<td>
<code>--args</code>
</td>
<td>string</td>
</tr>
<tr>
<td>
<code>userAgent</code>
</td>
<td>
<code>--user-agent</code>
</td>
<td>string</td>
</tr>
<tr>
<td>
<code>provider</code>
</td>
<td>
<code>-p, --provider</code>
</td>
<td>string</td>
</tr>
<tr>
<td>
<code>device</code>
</td>
<td>
<code>--device</code>
</td>
<td>string</td>
</tr>
<tr>
<td>
<code>ignoreHttpsErrors</code>
</td>
<td>
<code>--ignore-https-errors</code>
</td>
<td>boolean</td>
</tr>
<tr>
<td>
<code>allowFileAccess</code>
</td>
<td>
<code>--allow-file-access</code>
</td>
<td>boolean</td>
</tr>
<tr>
<td>
<code>cdp</code>
</td>
<td>
<code>--cdp</code>
</td>
<td>string</td>
</tr>
<tr>
<td>
<code>autoConnect</code>
</td>
<td>
<code>--auto-connect</code>
</td>
<td>boolean</td>
</tr>
<tr>
<td>
<code>engine</code>
</td>
<td>
<code>--engine</code>
</td>
<td>string (<code>chrome</code>, <code>lightpanda</code>)</td>
</tr>
<tr>
<td>
<code>colorScheme</code>
</td>
<td>
<code>--color-scheme</code>
</td>
<td>
string (<code>dark</code>, <code>light</code>, <code>no-preference</code>)
</td>
</tr>
<tr>
<td>
<code>downloadPath</code>
</td>
<td>
<code>--download-path</code>
</td>
<td>string</td>
</tr>
<tr>
<td>
<code>tabGroup</code>
</td>
<td>
<code>--tab-group</code>
</td>
<td>string (base title for session tab grouping via CDP plugin handshake)</td>
</tr>
<tr>
<td>
<code>tabGroupPluginId</code>
</td>
<td>
<code>--tab-group-plugin-id</code>
</td>
<td>string (expected extension ID for tab-group plugin handshake)</td>
</tr>
<tr>
<td>
<code>riskMode</code>
</td>
<td>
<code>--risk-mode</code>
</td>
<td>
string (<code>off</code>, <code>warn</code>, <code>block</code>)
</td>
</tr>
<tr>
<td>
<code>headers</code>
</td>
<td>
<code>--headers</code>
</td>
<td>string (JSON)</td>
</tr>
</tbody>
</table>
`riskMode` defaults to `warn` when unset.
`engine` defaults to `chrome`. `lightpanda` implies native mode and is headless-only.
For tab grouping in CDP mode, grouping is best-effort through the extension handshake:
extension available => grouped by session; extension missing/unavailable => silent no-op.
With the `agent-browser-stealth` extension installed, the side panel also exposes
session window isolation controls, activation guard toggles, empty-group cleanup, per-session allowlist policy editing, browser action controls, developer signal panels (console/network/DOM), and workflow automation (recording, shortcuts, and scheduled runs).
## Common Configurations
### Local Development
@@ -92,7 +341,7 @@ Every CLI flag can be set in the config file using its camelCase equivalent:
```json
{
"headed": true,
"profile": "./browser-data"
"sessionName": "local-dev"
}
```
@@ -124,17 +373,26 @@ Every CLI flag can be set in the config file using its camelCase equivalent:
}
```
### AI Agent Security
### Parallel Stateless Worker
```json
{
"contentBoundaries": true,
"maxOutput": 50000,
"allowedDomains": ["your-app.com", "*.your-app.com"],
"actionPolicy": "./policy.json"
"parallel": "worker-a"
}
```
Use this for stateless throughput tasks. For authenticated flows, prefer a stable `sessionName`.
When a default-session command runs, non-default daemon sessions are reaped.
## CLI-only daemon lifecycle flag
`--resident` is a CLI-only flag (not a config/env key). It keeps the daemon alive and disables the default 10-minute idle auto-shutdown.
```bash
agent-browser --resident open example.com
agent-browser close
```
## Overriding Boolean Options
Boolean flags accept an optional `true`/`false` value to override config settings:
@@ -150,7 +408,9 @@ agent-browser --headed open example.com # same as --headed true
agent-browser --headed true open example.com # explicit
```
This applies to all boolean flags: `--headed`, `--debug`, `--json`, `--ignore-https-errors`, `--allow-file-access`, `--auto-connect`, `--content-boundaries`, `--confirm-interactive`, `--native`.
This applies to all boolean flags: `--headed`, `--debug`, `--json`, `--ignore-https-errors`, `--allow-file-access`, `--auto-connect`, `--resident`.
For environment variables, headed mode accepts either `AGENT_BROWSER_HEADED=1` or `AGENT_BROWSER_HEADED=true`.
## Extensions Merging
@@ -158,34 +418,175 @@ Extensions from user-level and project-level configs are **concatenated**, not r
The `AGENT_BROWSER_EXTENSIONS` environment variable and CLI `--extension` flags follow the standard priority rules (env replaces config, CLI appends).
In this fork, local launches and extension launches remain headed by default unless headless is explicitly requested.
## Environment Variables
These environment variables configure additional daemon and runtime behavior:
<table>
<thead>
<tr><th>Variable</th><th>Description</th><th>Default</th></tr>
<tr>
<th>Variable</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr><td><code>AGENT_BROWSER_AUTO_CONNECT</code></td><td>Auto-discover and connect to a running Chrome instance.</td><td>(disabled)</td></tr>
<tr><td><code>AGENT_BROWSER_ALLOW_FILE_ACCESS</code></td><td>Allow <code>file://</code> URLs to access local files.</td><td>(disabled)</td></tr>
<tr><td><code>AGENT_BROWSER_COLOR_SCHEME</code></td><td>Color scheme preference (<code>dark</code>, <code>light</code>, <code>no-preference</code>).</td><td>(none)</td></tr>
<tr><td><code>AGENT_BROWSER_DOWNLOAD_PATH</code></td><td>Default directory for browser downloads.</td><td>(temp directory)</td></tr>
<tr><td><code>AGENT_BROWSER_DEFAULT_TIMEOUT</code></td><td>Default Playwright timeout in ms. Keep below 30000 to avoid IPC timeouts.</td><td><code>25000</code></td></tr>
<tr><td><code>AGENT_BROWSER_SESSION_NAME</code></td><td>Auto-save/load state persistence name.</td><td>(none)</td></tr>
<tr><td><code>AGENT_BROWSER_STATE_EXPIRE_DAYS</code></td><td>Auto-delete saved session states older than N days.</td><td><code>30</code></td></tr>
<tr><td><code>AGENT_BROWSER_ENCRYPTION_KEY</code></td><td>64-char hex key for AES-256-GCM session encryption.</td><td>(none)</td></tr>
<tr><td><code>AGENT_BROWSER_STREAM_PORT</code></td><td>Enable WebSocket streaming on the specified port (e.g., <code>9223</code>).</td><td>(disabled)</td></tr>
<tr><td><code>AGENT_BROWSER_IOS_DEVICE</code></td><td>Default iOS device name for the <code>ios</code> provider.</td><td>(none)</td></tr>
<tr><td><code>AGENT_BROWSER_IOS_UDID</code></td><td>Default iOS device UDID for the <code>ios</code> provider.</td><td>(none)</td></tr>
<tr><td><code>AGENT_BROWSER_DEBUG</code></td><td>Enable debug output (<code>1</code> to enable).</td><td>(disabled)</td></tr>
<tr><td><code>AGENT_BROWSER_CONTENT_BOUNDARIES</code></td><td>Wrap page output in boundary markers for LLM safety.</td><td>(disabled)</td></tr>
<tr><td><code>AGENT_BROWSER_MAX_OUTPUT</code></td><td>Max characters for page output (truncates beyond limit).</td><td>(unlimited)</td></tr>
<tr><td><code>AGENT_BROWSER_ALLOWED_DOMAINS</code></td><td>Comma-separated allowed domain patterns (e.g., <code>example.com,*.example.com</code>).</td><td>(unrestricted)</td></tr>
<tr><td><code>AGENT_BROWSER_ACTION_POLICY</code></td><td>Path to action policy JSON file.</td><td>(none)</td></tr>
<tr><td><code>AGENT_BROWSER_CONFIRM_ACTIONS</code></td><td>Comma-separated action categories requiring confirmation.</td><td>(none)</td></tr>
<tr><td><code>AGENT_BROWSER_CONFIRM_INTERACTIVE</code></td><td>Enable interactive confirmation prompts (auto-denies if stdin is not a TTY).</td><td>(disabled)</td></tr>
<tr><td><code>AGENT_BROWSER_NATIVE</code></td><td>Use the experimental native Rust daemon instead of Node.js/Playwright.</td><td>(disabled)</td></tr>
<tr>
<td>
<code>AGENT_BROWSER_AUTO_CONNECT</code>
</td>
<td>Auto-discover and connect to a running Chrome instance.</td>
<td>(disabled)</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_ALLOW_FILE_ACCESS</code>
</td>
<td>
Allow <code>file://</code> URLs to access local files.
</td>
<td>(disabled)</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_COLOR_SCHEME</code>
</td>
<td>
Color scheme preference (<code>dark</code>, <code>light</code>, <code>no-preference</code>).
</td>
<td>(none)</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_DOWNLOAD_PATH</code>
</td>
<td>Default directory for browser downloads.</td>
<td>(temp directory)</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_ENGINE</code>
</td>
<td>
Browser engine to use: <code>chrome</code> (default), <code>lightpanda</code>. Implies native mode.
</td>
<td>
<code>chrome</code>
</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_TAB_GROUP</code>
</td>
<td>Base title for tab grouping. Session suffix is appended automatically in CDP mode.</td>
<td>
<code>Agent Browser Stealth</code>
</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_TAB_GROUP_PLUGIN_ID</code>
</td>
<td>Expected extension ID for CDP tab-group plugin handshake.</td>
<td>
<code>aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa</code>
</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_RISK_MODE</code>
</td>
<td>
Verification/captcha handling mode (<code>off</code>, <code>warn</code>, <code>block</code>
).
</td>
<td>
<code>warn</code>
</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_DEFAULT_TIMEOUT</code>
</td>
<td>Default Playwright timeout in ms. Keep below 30000 to avoid IPC timeouts.</td>
<td>
<code>25000</code>
</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_SESSION_NAME</code>
</td>
<td>
Auto-save/load state persistence name (defaults to <code>default</code> when unset).
</td>
<td>
<code>default</code>
</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_PARALLEL</code>
</td>
<td>
Isolated runtime channel name for parallel AI runs (maps to <code>parallel-&lt;name&gt;</code>). Non-default daemons are reaped when a default-session command starts.
</td>
<td>(none)</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_STATE_EXPIRE_DAYS</code>
</td>
<td>Auto-delete saved session states older than N days.</td>
<td>
<code>30</code>
</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_ENCRYPTION_KEY</code>
</td>
<td>64-char hex key for AES-256-GCM session encryption.</td>
<td>(none)</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_STREAM_PORT</code>
</td>
<td>
Enable WebSocket streaming on the specified port (e.g., <code>9223</code>).
</td>
<td>(disabled)</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_IOS_DEVICE</code>
</td>
<td>
Default iOS device name for the <code>ios</code> provider.
</td>
<td>(none)</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_IOS_UDID</code>
</td>
<td>
Default iOS device UDID for the <code>ios</code> provider.
</td>
<td>(none)</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_DEBUG</code>
</td>
<td>
Enable debug output (<code>1</code> to enable).
</td>
<td>(disabled)</td>
</tr>
</tbody>
</table>
+105
View File
@@ -0,0 +1,105 @@
import { pageMetadata } from "@/lib/page-metadata"
export const metadata = pageMetadata("engines/chrome")
# Chrome
Chrome (and Chromium) is the default browser engine. agent-browser discovers, launches, and manages the Chrome process automatically via the Chrome DevTools Protocol (CDP).
## Binary Discovery
When no `--executable-path` is provided, agent-browser searches for Chrome in this order:
<table>
<thead>
<tr><th>Platform</th><th>Locations checked</th></tr>
</thead>
<tbody>
<tr>
<td>macOS</td>
<td>
<code>/Applications/Google Chrome.app</code>,
<code>/Applications/Google Chrome Canary.app</code>,
<code>/Applications/Chromium.app</code>,
Playwright Chromium cache
</td>
</tr>
<tr>
<td>Linux</td>
<td>
<code>google-chrome</code>,
<code>google-chrome-stable</code>,
<code>chromium-browser</code>,
<code>chromium</code> in PATH,
Playwright Chromium cache
</td>
</tr>
<tr>
<td>Windows</td>
<td>
<code>%LOCALAPPDATA%\Google\Chrome\Application\chrome.exe</code>,
<code>C:\Program Files\Google\Chrome\Application\chrome.exe</code>,
<code>C:\Program Files (x86)\...\chrome.exe</code>
</td>
</tr>
</tbody>
</table>
If Chrome is not found, run `agent-browser install` to download Chromium via Playwright.
## Usage
Chrome is the default engine. No `--engine` flag is needed:
```bash
agent-browser open example.com
```
To be explicit:
```bash
agent-browser --engine chrome open example.com
```
## Custom Binary
Point to any Chromium-based browser with `--executable-path`:
```bash
agent-browser --executable-path /path/to/chromium open example.com
```
Or via environment variable:
```bash
export AGENT_BROWSER_EXECUTABLE_PATH=/path/to/chromium
agent-browser open example.com
```
## Chrome-Specific Features
These features are available only with Chrome:
<table>
<thead>
<tr><th>Feature</th><th>Flag</th></tr>
</thead>
<tbody>
<tr><td>Browser extensions</td><td><code>--extension &lt;path&gt;</code></td></tr>
<tr><td>Persistent profiles</td><td><code>--profile &lt;path&gt;</code></td></tr>
<tr><td>Storage state</td><td><code>--state &lt;path&gt;</code></td></tr>
<tr><td>File URL access</td><td><code>--allow-file-access</code></td></tr>
<tr><td>Headed mode</td><td><code>--headed</code></td></tr>
<tr><td>Custom launch args</td><td><code>--args &lt;args&gt;</code></td></tr>
</tbody>
</table>
## Containers and CI
In Docker, CI runners, or other sandboxed environments, Chrome's user namespace sandbox may need to be disabled:
```bash
agent-browser --args "--no-sandbox" open example.com
```
agent-browser automatically adds `--no-sandbox` when it detects a container environment (Docker, Podman, or root execution).
+97
View File
@@ -0,0 +1,97 @@
import { pageMetadata } from "@/lib/page-metadata"
export const metadata = pageMetadata("engines/lightpanda")
# Lightpanda
[Lightpanda](https://lightpanda.io/) is a headless browser engine built from scratch in Zig. It is intended for machine-driven workloads where fast startup and low memory use matter more than full Chrome compatibility.
agent-browser manages Lightpanda the same way it manages Chrome: spawn the process, connect via CDP, and drive the same downstream commands (`snapshot`, `click`, `fill`, `screenshot`, and so on).
## Installation
Install the Lightpanda binary before using it with agent-browser:
<table>
<thead>
<tr><th>Platform</th><th>Command</th></tr>
</thead>
<tbody>
<tr>
<td>macOS (Apple Silicon)</td>
<td><code>curl -L -o lightpanda https://github.com/lightpanda-io/browser/releases/download/nightly/lightpanda-aarch64-macos &amp;&amp; chmod a+x ./lightpanda</code></td>
</tr>
<tr>
<td>Linux (x86_64)</td>
<td><code>curl -L -o lightpanda https://github.com/lightpanda-io/browser/releases/download/nightly/lightpanda-x86_64-linux &amp;&amp; chmod a+x ./lightpanda</code></td>
</tr>
</tbody>
</table>
Move the binary somewhere in your `PATH` such as `/usr/local/bin/lightpanda` or `~/.local/bin/lightpanda`.
See the [Lightpanda installation docs](https://lightpanda.io/docs/open-source/installation) for more options.
## Usage
Use `--engine` to select Lightpanda:
```bash
agent-browser --engine lightpanda open example.com
agent-browser --engine lightpanda snapshot
agent-browser --engine lightpanda screenshot
```
Or set it as the default via environment variable:
```bash
export AGENT_BROWSER_ENGINE=lightpanda
agent-browser open example.com
```
Or in `agent-browser.json`:
```json
{
"engine": "lightpanda"
}
```
## Custom Binary Path
If the `lightpanda` binary is not in your `PATH`, use `--executable-path`:
```bash
agent-browser --engine lightpanda --executable-path /path/to/lightpanda open example.com
```
## Differences From Chrome
Lightpanda is headless-only and does not support several Chrome-specific features:
<table>
<thead>
<tr><th>Feature</th><th>Status</th></tr>
</thead>
<tbody>
<tr><td>Extensions (<code>--extension</code>)</td><td>Not supported</td></tr>
<tr><td>Persistent profiles (<code>--profile</code>)</td><td>Not supported</td></tr>
<tr><td>Storage state (<code>--state</code>)</td><td>Not supported</td></tr>
<tr><td>File access (<code>--allow-file-access</code>)</td><td>Not supported</td></tr>
<tr><td>Headed mode (<code>--headed</code>)</td><td>Not applicable</td></tr>
<tr><td>Screenshots</td><td>Depends on Lightpanda CDP support</td></tr>
</tbody>
</table>
agent-browser returns a clear error if you combine `--engine lightpanda` with unsupported flags.
## When To Use Lightpanda
Lightpanda is a good fit for:
- Fast scraping and extraction jobs
- AI agent workflows where speed and low memory matter
- CI environments with constrained resources
- High-volume parallel automation
Use Chrome when you need full browser fidelity, extensions, or persistent profiles.
+18 -9
View File
@@ -9,7 +9,7 @@ export const metadata = pageMetadata("installation")
Installs the native Rust binary for maximum performance:
```bash
npm install -g agent-browser
npm install -g agent-browser-stealth
agent-browser install # Download Chromium
```
@@ -20,8 +20,8 @@ This is the fastest option -- commands run through the native Rust CLI directly
Run directly with `npx` if you want to try it without installing globally:
```bash
npx agent-browser install # Download Chromium (first time only)
npx agent-browser open example.com
npx agent-browser-stealth install # Download Chromium (first time only)
npx agent-browser-stealth open example.com
```
> **Note:** `npx` routes through Node.js before reaching the Rust CLI, so it is noticeably slower than a global install. For regular use, install globally.
@@ -31,14 +31,14 @@ npx agent-browser open example.com
For projects that want to pin the version in `package.json`:
```bash
npm install agent-browser
npx agent-browser install
npm install agent-browser-stealth
npx agent-browser-stealth install
```
Then use via `npx` or `package.json` scripts:
```bash
npx agent-browser open example.com
npx agent-browser-stealth open example.com
```
## Homebrew (macOS)
@@ -51,7 +51,7 @@ agent-browser install # Download Chromium
## From source
```bash
git clone https://github.com/vercel-labs/agent-browser
git clone https://github.com/leeguooooo/agent-browser
cd agent-browser
pnpm install
pnpm build
@@ -60,6 +60,15 @@ pnpm build:native
pnpm link --global
```
## Fork versioning
Fork releases use a dual-version format:
- `<upstream>-fork.<fork>`
- Example: `0.14.0-fork.1`
`agent-browser --version` prints the full version and also shows upstream and fork parts for fork builds.
## Linux dependencies
On Linux, install system dependencies:
@@ -89,7 +98,7 @@ AGENT_BROWSER_EXECUTABLE_PATH=/path/to/chromium agent-browser open example.com
```typescript
import chromium from '@sparticuz/chromium';
import { BrowserManager } from 'agent-browser';
import { BrowserManager } from 'agent-browser-stealth';
export async function handler() {
const browser = new BrowserManager();
@@ -110,7 +119,7 @@ agent-browser works with any AI agent out of the box. For richer context:
Install the skill for your AI coding assistant:
```bash
npx skills add vercel-labs/agent-browser
npx skills add leeguooooo/agent-browser
```
This works with Claude Code, Codex, Cursor, Gemini CLI, GitHub Copilot, Goose, OpenCode, and Windsurf. The skill is fetched from the repository and stays up to date automatically.
+6 -2
View File
@@ -7,13 +7,15 @@ export const metadata = pageMetadata("")
Browser automation CLI designed for AI agents. Compact text output minimizes context usage. Fast Rust CLI with Node.js fallback.
```bash
npm install -g agent-browser # all platforms (fastest, native Rust CLI)
npm install -g agent-browser-stealth # all platforms (fastest, native Rust CLI)
brew install agent-browser # macOS
# or try without installing
npx agent-browser open example.com
npx agent-browser-stealth open example.com
```
Executable aliases after install: `agent-browser`, `agent-browser-stealth`, and `abs`.
## Features
- **Agent-first** - Compact text output uses fewer tokens than JSON, designed for AI context efficiency
@@ -22,6 +24,8 @@ npx agent-browser open example.com
- **Complete** - 50+ commands for navigation, forms, screenshots, network, storage
- **Sessions** - Multiple isolated browser instances with separate auth
- **Cross-platform** - macOS, Linux, Windows with native binaries
- **Auto region detection** - Locale, timezone, and Accept-Language automatically match the target site's TLD
- **Captcha auto-retry** - Detects captcha/verification pages and retries with randomized backoff
## Works with
+66 -43
View File
@@ -1,61 +1,57 @@
import { pageMetadata } from "@/lib/page-metadata"
import { pageMetadata } from '@/lib/page-metadata';
export const metadata = pageMetadata("sessions")
export const metadata = pageMetadata('sessions');
# Sessions
Run multiple isolated browser instances:
Use the default runtime session or an isolated parallel runtime channel, plus optional named persistence:
```bash
# Different sessions
agent-browser --session agent1 open site-a.com
agent-browser --session agent2 open site-b.com
# Or via environment variable
AGENT_BROWSER_SESSION=agent1 agent-browser click "#btn"
# List active sessions
agent-browser session list
# Output:
# Active sessions:
# -> default
# agent1
# Show current session
# Show current runtime session
agent-browser session
# Output: default
# Isolated runtime channel for parallel AI flow
agent-browser --parallel worker-a session
# Output: parallel-worker-a
# Show active daemon sessions
agent-browser session list
# stale entries are auto-cleaned during listing
```
## Session isolation
Each session has its own:
Runtime session defaults to `default`. Use `--parallel <name>` when you need isolated concurrent runtime channels, and `--session-name` to isolate persisted state files per workflow.
- Browser instance
- Cookies and storage
- Navigation history
- Cookies and storage snapshots
- Authentication state
- Saved state lifecycle
## Persistent profiles
Daemons auto-shutdown after 10 minutes of inactivity by default. Use `--resident` to keep a daemon alive until explicit `close`.
By default, browser state is lost when the browser closes. Use `--profile` to persist state across restarts:
## Parallel runtime channels
Use `--parallel <name>` to isolate runtime channels for concurrent AI execution:
```bash
# Use a persistent profile directory
agent-browser --profile ~/.myapp-profile open myapp.com
# Login once, then reuse the authenticated session
agent-browser --profile ~/.myapp-profile open myapp.com/dashboard
# Or via environment variable
AGENT_BROWSER_PROFILE=~/.myapp-profile agent-browser open myapp.com
agent-browser --parallel worker-a open https://example.com
agent-browser --parallel worker-b open https://example.org
```
The profile directory stores:
`--parallel` is intended for stateless throughput tasks (navigation/extraction/checks). For authenticated flows, use a stable `--session-name`.
- Cookies and localStorage
- IndexedDB data
- Service workers
- Browser cache
- Login sessions
Default session isolation policy:
- Running a default-session command reaps all non-default daemon sessions (`parallel-*` and legacy named channels).
- This keeps the primary `default` runtime channel free from stale daemon reuse.
For long-running workers, add `--resident` to disable idle auto-shutdown:
```bash
agent-browser --parallel worker-a --resident open https://example.com
agent-browser --parallel worker-a close
```
## Session persistence
@@ -73,6 +69,10 @@ export AGENT_BROWSER_SESSION_NAME=twitter
agent-browser open twitter.com
```
If `--session-name` is omitted, it defaults to `default`.
When `--parallel` is enabled, auto persistence is disabled by default unless `--session-name` is explicitly passed on that command.
State files are stored in `~/.agent-browser/sessions/` and automatically loaded on daemon start.
### Session name rules
@@ -185,12 +185,35 @@ agent-browser set headers '{"X-Custom-Header": "value"}'
<table>
<thead>
<tr><th>Variable</th><th>Description</th></tr>
<tr>
<th>Variable</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr><td><code>AGENT_BROWSER_SESSION</code></td><td>Browser session ID (default: "default")</td></tr>
<tr><td><code>AGENT_BROWSER_SESSION_NAME</code></td><td>Auto-save/load state persistence name</td></tr>
<tr><td><code>AGENT_BROWSER_ENCRYPTION_KEY</code></td><td>64-char hex key for AES-256-GCM encryption</td></tr>
<tr><td><code>AGENT_BROWSER_STATE_EXPIRE_DAYS</code></td><td>Auto-delete states older than N days (default: 30)</td></tr>
<tr>
<td>
<code>AGENT_BROWSER_SESSION_NAME</code>
</td>
<td>Auto-save/load state persistence name</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_PARALLEL</code>
</td>
<td>Isolated runtime channel name for parallel AI runs (maps to <code>parallel-&lt;name&gt;</code>). Non-default daemons are reaped when a default-session command starts.</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_ENCRYPTION_KEY</code>
</td>
<td>64-char hex key for AES-256-GCM encryption</td>
</tr>
<tr>
<td>
<code>AGENT_BROWSER_STATE_EXPIRE_DAYS</code>
</td>
<td>Auto-delete states older than N days (default: 30)</td>
</tr>
</tbody>
</table>
+1 -1
View File
@@ -176,7 +176,7 @@ Send input events to control the browser remotely.
For advanced use, control streaming directly via the TypeScript API:
```typescript
import { BrowserManager } from 'agent-browser';
import { BrowserManager } from 'agent-browser-stealth';
const browser = new BrowserManager();
await browser.launch({ headless: true });
+2 -2
View File
@@ -53,7 +53,7 @@ export function Header() {
</div>
<nav className="flex items-center gap-4">
<a
href="https://github.com/vercel-labs/agent-browser"
href="https://github.com/leeguooooo/agent-browser"
target="_blank"
rel="noopener noreferrer"
className="flex items-center gap-1.5 text-sm text-neutral-500 hover:text-neutral-900 transition-colors dark:text-neutral-400 dark:hover:text-neutral-100"
@@ -69,7 +69,7 @@ export function Header() {
<span>16k</span>
</a>
<a
href="https://www.npmjs.com/package/agent-browser"
href="https://www.npmjs.com/package/agent-browser-stealth"
target="_blank"
rel="noopener noreferrer"
className="text-sm text-neutral-500 hover:text-neutral-900 transition-colors dark:text-neutral-400 dark:hover:text-neutral-100"
+7
View File
@@ -40,6 +40,13 @@ export const navigation: NavSection[] = [
{ name: "Native Mode (Experimental)", href: "/native-mode" },
],
},
{
title: "Engines",
items: [
{ name: "Chrome", href: "/engines/chrome" },
{ name: "Lightpanda", href: "/engines/lightpanda" },
],
},
{
title: null,
items: [{ name: "Changelog", href: "/changelog" }],
+2
View File
@@ -14,6 +14,8 @@ export const PAGE_TITLES: Record<string, string> = {
profiler: "Profiler",
ios: "iOS Simulator",
security: "Security",
"engines/chrome": "Chrome",
"engines/lightpanda": "Lightpanda",
"native-mode": "Native Mode (Experimental)",
changelog: "Changelog",
};
+148
View File
@@ -0,0 +1,148 @@
# Upstream Sync Audit (2026-03-09)
Scope: compare current `main` plus the local in-progress sync worktree with `upstream/main`.
## Already Synced
- `de5ea1d` `fix: use reqwest for CDP port discovery instead of broken hand-rolled HTTP client (#619)`
- `8f6ad81` `Fix dialog dismiss command parsing (#605)`
- `7acde7e` `fix: native auth login fails due to incompatible encryption format (#648)`
- `492830a` `Fix: Suppress Google Translate bar in native headless mode (#649)`
- `68cebe5` `Fix Chrome extensions not loading by forcing headed mode when extensions present (#652)`
- `b7e7a25` `fix: persist auth cookies on close in native mode (#650)`
## Absorbed Locally (Not Exact Cherry-Picks)
- `eaa968e` `fix: suppress spurious --native warning when set via env var (#611)`
- Covered by the local native CLI restoration in:
- [cli/src/flags.rs](/Users/leo/github.com/agent-browser/cli/src/flags.rs)
- [cli/src/main.rs](/Users/leo/github.com/agent-browser/cli/src/main.rs)
- [cli/src/connection.rs](/Users/leo/github.com/agent-browser/cli/src/connection.rs)
- [cli/src/native/daemon.rs](/Users/leo/github.com/agent-browser/cli/src/native/daemon.rs)
- `788ad0e` `chore: add cargo fmt check to Rust CI and fix existing violations (#620)`
- The Rust CI `fmt` check is already present in [.github/workflows/ci.yml](/Users/leo/github.com/agent-browser/.github/workflows/ci.yml).
- `aba2353` `Fix clippy warnings across CLI codebase (#654)`
- The current worktree already carries the relevant CLI cleanup needed for `cargo clippy -- -D warnings` to pass.
- `d9387aa` `ci: add clippy check to Rust CI workflow (#675)`
- The Rust CI `clippy` check is already present in [.github/workflows/ci.yml](/Users/leo/github.com/agent-browser/.github/workflows/ci.yml).
- `f262ff1` `docs: improve snapshot usage guidance and add reproducibility check (#630)`
- Safe docs-only sync. Applied locally in [skills/dogfood/SKILL.md](/Users/leo/github.com/agent-browser/skills/dogfood/SKILL.md).
- `a0bd0c2` `Add webview support for Electron apps in native mode (#671)`
- Applied locally in:
- [cli/src/native/actions.rs](/Users/leo/github.com/agent-browser/cli/src/native/actions.rs)
- [cli/src/native/browser.rs](/Users/leo/github.com/agent-browser/cli/src/native/browser.rs)
- Broadens native target discovery from `page` to `page | webview` and adds `type` to native `tab_list` output.
- Does not alter the fork's Node.js stealth launch defaults.
- `36c2e06` `add benchmarks (#637)`
- Applied locally in:
- [package.json](/Users/leo/github.com/agent-browser/package.json)
- [test/benchmarks/run.ts](/Users/leo/github.com/agent-browser/test/benchmarks/run.ts)
- [test/benchmarks/scenarios.ts](/Users/leo/github.com/agent-browser/test/benchmarks/scenarios.ts)
- Adds developer benchmark scripts only. No runtime or stealth launch behavior changes.
- `0da54c7` `lightpanda (#646)` core feature set
- Applied locally in:
- [cli/src/flags.rs](/Users/leo/github.com/agent-browser/cli/src/flags.rs)
- [cli/src/main.rs](/Users/leo/github.com/agent-browser/cli/src/main.rs)
- [cli/src/connection.rs](/Users/leo/github.com/agent-browser/cli/src/connection.rs)
- [cli/src/native/actions.rs](/Users/leo/github.com/agent-browser/cli/src/native/actions.rs)
- [cli/src/native/browser.rs](/Users/leo/github.com/agent-browser/cli/src/native/browser.rs)
- [cli/src/native/cdp/lightpanda.rs](/Users/leo/github.com/agent-browser/cli/src/native/cdp/lightpanda.rs)
- [src/protocol.ts](/Users/leo/github.com/agent-browser/src/protocol.ts)
- [src/types.ts](/Users/leo/github.com/agent-browser/src/types.ts)
- [src/actions.ts](/Users/leo/github.com/agent-browser/src/actions.ts)
- [docs/src/app/engines/chrome/page.mdx](/Users/leo/github.com/agent-browser/docs/src/app/engines/chrome/page.mdx)
- [docs/src/app/engines/lightpanda/page.mdx](/Users/leo/github.com/agent-browser/docs/src/app/engines/lightpanda/page.mdx)
- [docs/src/lib/docs-navigation.ts](/Users/leo/github.com/agent-browser/docs/src/lib/docs-navigation.ts)
- [docs/src/lib/page-titles.ts](/Users/leo/github.com/agent-browser/docs/src/lib/page-titles.ts)
- [test/benchmarks/run.ts](/Users/leo/github.com/agent-browser/test/benchmarks/run.ts)
- [test/benchmarks/engine-scenarios.ts](/Users/leo/github.com/agent-browser/test/benchmarks/engine-scenarios.ts)
- [test/benchmarks/pages/article.html](/Users/leo/github.com/agent-browser/test/benchmarks/pages/article.html)
- [test/benchmarks/pages/dashboard.html](/Users/leo/github.com/agent-browser/test/benchmarks/pages/dashboard.html)
- [test/benchmarks/pages/ecommerce.html](/Users/leo/github.com/agent-browser/test/benchmarks/pages/ecommerce.html)
- Shared launch protocol now accepts `engine`. The Node path still rejects `engine=lightpanda` with a clear `--native` requirement, while the native path can launch either `chrome` or `lightpanda`.
- This preserves the current Node.js/Chrome stealth path while adding the native-only alternative engine surface and its supporting docs/benchmarks.
## Remaining Upstream Commits
Current status: there are no remaining upstream feature commits that are both codeful and safe to port directly into this fork. What remains is either release metadata or the stealth-sensitive `#607` launch-policy batch.
### Low Risk / Independent Of Stealth
- `94521e7` `chore: add minor changeset for release (#683)`
- Release metadata only.
- `2bab729` `chore: version packages (#684)`
- Release/version bump only.
- `01ac557` `chore: add patch changeset for release (#609)`
- Release metadata only.
- `7d2c895` `chore: add patch changeset for release (#612)`
- Release metadata only.
- `7edc5d5` `chore: version packages (#610)`
- Release/version bump only.
- `794a77e` `chore: version packages (#613)`
- Release/version bump only.
### Needs Manual Review Because It Touches Stealth-Sensitive Launch Behavior
- `e5fd26e` `headed mode (#607)`
- Overlaps with our fork-modified launch path:
- `src/browser.ts`
- `src/daemon.ts`
- `cli/src/native/cdp/chrome.rs`
- `cli/src/connection.rs`
- Upstream intent:
- honor `AGENT_BROWSER_HEADED`
- support headed launch in more places
- add temp profile cleanup and tests
- Fork-specific risk:
- upstream changes persistent extension launch from `headless: false` to `headless: options.headless ?? true` in `src/browser.ts`
- our fork intentionally keeps extension launches headed by default via [src/browser.ts](/Users/leo/github.com/agent-browser/src/browser.ts#L2131)
- our daemon auto-launch path already honors `AGENT_BROWSER_HEADED=1` and `AGENT_BROWSER_HEADED=true` in [src/daemon.ts](/Users/leo/github.com/agent-browser/src/daemon.ts#L523)
- the native temp-profile cleanup and extension-headed logic from upstream are already present in [cli/src/native/cdp/chrome.rs](/Users/leo/github.com/agent-browser/cli/src/native/cdp/chrome.rs)
- blindly reapplying the upstream Node hunk would move extension launch defaults back toward upstream headless behavior and would change current stealth assumptions
- Recommendation:
- do not cherry-pick this commit directly
- keep fork ownership of headed/headless defaults in the Node.js path
- extract only test-only utilities or assertions that do not alter launch policy
- local regression tests now lock the fork policy in [src/browser.test.ts](/Users/leo/github.com/agent-browser/src/browser.test.ts), including default local headed launch and extension launches remaining headed by default
- Node daemon env parsing is also locked in [src/daemon.test.ts](/Users/leo/github.com/agent-browser/src/daemon.test.ts), including `AGENT_BROWSER_HEADED=true` and comma/newline parsing for extensions and args
- treat headless/headed defaults as a fork-owned policy decision
### Already Partly Reimplemented In Fork
- `139dd0e` `fix: surface daemon startup errors instead of opaque timeout message (#614)`
- Current fork already captures daemon stderr with `Stdio::piped()` and checks `try_wait()` during startup polling in [cli/src/connection.rs](/Users/leo/github.com/agent-browser/cli/src/connection.rs#L478) and [cli/src/connection.rs](/Users/leo/github.com/agent-browser/cli/src/connection.rs#L685).
- `AGENT_BROWSER_DEBUG` forwarding is already present in [cli/src/connection.rs](/Users/leo/github.com/agent-browser/cli/src/connection.rs#L550) and [cli/src/connection.rs](/Users/leo/github.com/agent-browser/cli/src/connection.rs#L651).
- Re-review on 2026-03-09 confirms the local implementation is functionally equivalent or stronger than upstream, with the same stderr surfacing and early-exit detection but fork-specific daemon spawn logic.
- Recommendation: treat `#614` as absorbed locally and do not cherry-pick it.
## Fork-Specific Blockers Found During Audit
- Native CLI wiring was missing during the initial audit, but has since been restored locally.
- Remaining blocker is no longer the `--native` switch itself.
- The real decision point is whether this fork wants to expose new native features (`--engine`, Lightpanda, Electron webview) that do not help stealth directly but do expand the maintained surface area.
- That decision has now been made in favor of exposing them locally, so the blocker section is effectively closed for the current sync round.
## Current Verification
- `cd /Users/leo/github.com/agent-browser/cli && cargo fmt -- --check`
- `cd /Users/leo/github.com/agent-browser/cli && cargo clippy -- -D warnings`
- `cd /Users/leo/github.com/agent-browser/cli && cargo test`
- `cd /Users/leo/github.com/agent-browser && pnpm build`
- `cd /Users/leo/github.com/agent-browser && pnpm exec tsx test/benchmarks/run.ts --node-only --iterations 1 --warmup 0`
- `cd /Users/leo/github.com/agent-browser && pnpm exec vitest run src/actions.test.ts test/keyboard.test.ts test/launch-options.test.ts`
All checks pass against the current local sync worktree.
## Recommended Migration Order
1. CI hygiene batch
- Already absorbed locally via the current worktree.
- No stealth behavior change.
2. Docs-only batch
- Safe to keep following `#630`-style guidance updates.
- No runtime behavior change.
3. Headed-mode audit
- Reconcile upstream `#607` against fork-owned stealth launch defaults instead of cherry-picking it.
4. Release metadata
- Keep fork-owned release/versioning flow.
- Do not mirror upstream changesets or version bumps unless this fork explicitly decides to realign its release train.
+425
View File
@@ -0,0 +1,425 @@
(() => {
const REQUEST_TYPE = 'AB_TAB_GROUP_REQUEST';
const RESPONSE_TYPE = 'AB_TAB_GROUP_RESPONSE';
const CONTENT_EVENT_TYPE = 'AB_CONTENT_EVENT';
const CONTENT_EXECUTE_ACTION = 'AB_CONTENT_EXECUTE_ACTION';
const CONTENT_GET_DOM_STATE = 'AB_CONTENT_GET_DOM_STATE';
const CONTENT_PING = 'AB_CONTENT_PING';
const PAGE_BRIDGE_EVENT = 'AB_PAGE_BRIDGE_EVENT';
const STORAGE_OPTIONS_KEY = 'abExtensionOptionsV1';
const mutationState = {
total: 0,
recent: [],
observerReady: false,
};
function pushMutationSummary(entry) {
mutationState.total += 1;
mutationState.recent.push({
...entry,
timestamp: Date.now(),
});
if (mutationState.recent.length > 40) {
mutationState.recent.splice(0, mutationState.recent.length - 40);
}
}
function serializeValue(value, depth = 0) {
if (value === null || typeof value === 'undefined') return value;
if (typeof value === 'string') return value.slice(0, 300);
if (typeof value === 'number' || typeof value === 'boolean') return value;
if (value instanceof Error) return `${value.name}: ${value.message}`;
if (depth > 2) return '[depth-limit]';
if (Array.isArray(value)) {
return value.slice(0, 10).map((item) => serializeValue(item, depth + 1));
}
if (typeof value === 'object') {
const out = {};
for (const [key, entry] of Object.entries(value).slice(0, 15)) {
out[key] = serializeValue(entry, depth + 1);
}
return out;
}
return String(value).slice(0, 300);
}
function sendRuntimeEvent(kind, payload) {
try {
chrome.runtime.sendMessage({
type: CONTENT_EVENT_TYPE,
kind,
payload: serializeValue(payload),
url: window.location.href,
title: document.title,
timestamp: Date.now(),
});
} catch {
// Ignore runtime channel errors.
}
}
function getPageBridgeEnabled() {
return new Promise((resolve) => {
try {
chrome.storage.local.get([STORAGE_OPTIONS_KEY], (result) => {
if (chrome.runtime.lastError) {
resolve(false);
return;
}
const rawOptions = result?.[STORAGE_OPTIONS_KEY];
resolve(Boolean(rawOptions && typeof rawOptions === 'object' && rawOptions.pageBridgeEnabled === true));
});
} catch {
resolve(false);
}
});
}
async function installPageBridge() {
// Receives events emitted by the injected page-world hook script.
const bridgeListener = (event) => {
if (event.source !== window) return;
const data = event.data;
if (!data || data.type !== PAGE_BRIDGE_EVENT) return;
sendRuntimeEvent(data.kind || 'page-event', data.payload || {});
};
window.addEventListener('message', bridgeListener);
const parent = document.documentElement || document.head || document.body;
if (!parent) return;
if (!(await getPageBridgeEnabled())) {
sendRuntimeEvent('lifecycle', {
event: 'bridge-disabled-default',
});
return;
}
// Use external extension script instead of inline text to reduce CSP conflicts.
const script = document.createElement('script');
script.src = chrome.runtime.getURL('page-bridge.js');
script.async = false;
script.dataset.abBridgeEvent = PAGE_BRIDGE_EVENT;
script.onload = () => script.remove();
script.onerror = () => {
sendRuntimeEvent('lifecycle', {
event: 'bridge-load-failed',
host: window.location.hostname,
});
script.remove();
};
parent.appendChild(script);
}
function ensureMutationObserver() {
if (mutationState.observerReady) return;
if (!document.documentElement) return;
const observer = new MutationObserver((records) => {
const summary = {
records: records.length,
addedNodes: 0,
removedNodes: 0,
};
for (const record of records.slice(0, 40)) {
summary.addedNodes += record.addedNodes?.length || 0;
summary.removedNodes += record.removedNodes?.length || 0;
}
pushMutationSummary(summary);
});
observer.observe(document.documentElement, {
childList: true,
subtree: true,
attributes: true,
attributeFilter: ['class', 'style', 'hidden', 'disabled', 'aria-hidden'],
});
mutationState.observerReady = true;
}
function toSimpleNode(element) {
if (!element || typeof element !== 'object') return null;
const node = {
tag: element.tagName?.toLowerCase() || 'unknown',
id: element.id || undefined,
className: typeof element.className === 'string' ? element.className.slice(0, 120) : '',
role: element.getAttribute?.('role') || undefined,
name:
element.getAttribute?.('aria-label') ||
element.getAttribute?.('name') ||
element.getAttribute?.('placeholder') ||
'',
text: (element.textContent || '').trim().replace(/\s+/g, ' ').slice(0, 160),
disabled: element.disabled === true,
hidden: element.hidden === true,
};
return node;
}
function collectInteractiveElements(root, limit = 80) {
const selector = [
'a[href]',
'button',
'input',
'select',
'textarea',
'summary',
'[role="button"]',
'[role="link"]',
'[tabindex]'
].join(',');
const out = [];
const nodes = root.querySelectorAll(selector);
for (const element of nodes) {
if (out.length >= limit) break;
out.push(toSimpleNode(element));
}
return out.filter(Boolean);
}
function collectDomState(options = {}) {
const selector = typeof options.selector === 'string' ? options.selector.trim() : '';
const root = selector ? document.querySelector(selector) : document.body || document.documentElement;
if (!root) {
return {
ok: false,
error: selector ? `selector-not-found: ${selector}` : 'root-not-found',
};
}
const textPreview = (root.textContent || '').replace(/\s+/g, ' ').trim().slice(0, 1000);
const interactiveOnly = options.interactiveOnly === true;
const interactiveElements = collectInteractiveElements(root, options.maxNodes || 80);
const dom = {
href: window.location.href,
title: document.title,
readyState: document.readyState,
selector: selector || null,
rootTag: root.tagName?.toLowerCase() || 'unknown',
textPreview,
interactiveCount: interactiveElements.length,
interactiveElements,
mutation: {
total: mutationState.total,
recent: mutationState.recent.slice(-10),
},
capturedAt: Date.now(),
};
if (interactiveOnly) {
dom.textPreview = '';
}
return {
ok: true,
state: dom,
};
}
function queryElement(selector) {
if (typeof selector !== 'string' || selector.trim().length === 0) {
throw new Error('selector is required');
}
const element = document.querySelector(selector);
if (!element) {
throw new Error(`Element not found: ${selector}`);
}
return element;
}
function focusElement(element) {
if (typeof element.focus === 'function') {
element.focus({ preventScroll: false });
}
}
function dispatchInputEvents(element) {
element.dispatchEvent(new Event('input', { bubbles: true }));
element.dispatchEvent(new Event('change', { bubbles: true }));
}
async function executeAction(command, args = {}) {
switch (command) {
case 'click': {
const element = queryElement(args.selector);
focusElement(element);
element.click();
return { ok: true, action: command, selector: args.selector };
}
case 'fill': {
const element = queryElement(args.selector);
if (!('value' in element)) {
throw new Error(`Element is not fillable: ${args.selector}`);
}
focusElement(element);
element.value = typeof args.value === 'string' ? args.value : String(args.value || '');
dispatchInputEvents(element);
return { ok: true, action: command, selector: args.selector, valueLength: element.value.length };
}
case 'press': {
const key = typeof args.key === 'string' && args.key.trim().length > 0 ? args.key.trim() : 'Enter';
let target;
if (typeof args.selector === 'string' && args.selector.trim().length > 0) {
target = queryElement(args.selector);
focusElement(target);
} else {
target = document.activeElement || document.body;
}
const down = new KeyboardEvent('keydown', { key, bubbles: true });
const up = new KeyboardEvent('keyup', { key, bubbles: true });
target.dispatchEvent(down);
target.dispatchEvent(up);
return { ok: true, action: command, key };
}
case 'eval': {
if (typeof args.expression !== 'string' || args.expression.trim().length === 0) {
throw new Error('expression is required');
}
const fn = new Function(`return (${args.expression});`);
const result = fn();
return { ok: true, action: command, result: serializeValue(result) };
}
case 'snapshot': {
return {
ok: true,
action: command,
...collectDomState({
selector: args.selector,
interactiveOnly: args.interactiveOnly === true,
maxNodes: args.maxNodes,
}),
};
}
default:
throw new Error(`Unknown content action: ${command}`);
}
}
ensureMutationObserver();
installPageBridge();
window.addEventListener('message', (event) => {
if (event.source !== window) {
return;
}
const data = event.data;
if (!data || data.type !== REQUEST_TYPE) {
return;
}
const request = {
type: REQUEST_TYPE,
nonce: data.nonce,
session: data.session,
groupTitle: data.groupTitle,
pluginId: data.pluginId,
allowedDomains: Array.isArray(data.allowedDomains) ? data.allowedDomains : undefined,
};
try {
chrome.runtime.sendMessage(request, (response) => {
const lastError = chrome.runtime.lastError;
if (lastError) {
window.postMessage(
{
type: RESPONSE_TYPE,
nonce: request.nonce,
ok: false,
error: lastError.message,
},
'*'
);
return;
}
const payload = response && typeof response === 'object' ? response : { ok: false };
window.postMessage(
{
type: RESPONSE_TYPE,
nonce: request.nonce,
ok: payload.ok === true,
extensionId:
typeof payload.extensionId === 'string' && payload.extensionId.length > 0
? payload.extensionId
: chrome.runtime.id,
groupId: typeof payload.groupId === 'number' ? payload.groupId : undefined,
windowId: typeof payload.windowId === 'number' ? payload.windowId : undefined,
color: typeof payload.color === 'string' ? payload.color : undefined,
collapsed: payload.collapsed === true,
policy:
payload.policy && typeof payload.policy === 'object'
? {
enforced: payload.policy.enforced === true,
blocked: payload.policy.blocked === true,
reason:
typeof payload.policy.reason === 'string' ? payload.policy.reason : undefined,
}
: undefined,
riskHints: Array.isArray(payload.riskHints) ? payload.riskHints : undefined,
error: typeof payload.error === 'string' ? payload.error : undefined,
},
'*'
);
});
} catch (error) {
const errorMessage = error instanceof Error ? error.message : String(error);
window.postMessage(
{
type: RESPONSE_TYPE,
nonce: request.nonce,
ok: false,
error: errorMessage,
},
'*'
);
}
});
chrome.runtime.onMessage.addListener((message, _sender, sendResponse) => {
if (!message || typeof message !== 'object') return;
if (message.type === CONTENT_PING) {
sendResponse({
ok: true,
href: window.location.href,
title: document.title,
readyState: document.readyState,
});
return;
}
if (message.type === CONTENT_GET_DOM_STATE) {
sendResponse(collectDomState(message.options || {}));
return;
}
if (message.type === CONTENT_EXECUTE_ACTION) {
executeAction(message.command, message.args || {})
.then((result) => sendResponse(result))
.catch((error) => {
sendResponse({
ok: false,
action: message.command,
error: error instanceof Error ? error.message : String(error),
});
});
return true;
}
});
})();
+7
View File
@@ -0,0 +1,7 @@
<svg width="128" height="128" viewBox="0 0 128 128" fill="none" xmlns="http://www.w3.org/2000/svg">
<rect width="128" height="128" rx="32" fill="#1A73E8"/>
<rect x="30" y="34" width="68" height="10" rx="2" fill="white"/>
<rect x="30" y="54" width="48" height="10" rx="2" fill="white" fill-opacity="0.8"/>
<rect x="30" y="74" width="28" height="10" rx="2" fill="white" fill-opacity="0.6"/>
<circle cx="94" cy="90" r="10" fill="#34A853" stroke="#1A73E8" stroke-width="4"/>
</svg>

After

Width:  |  Height:  |  Size: 488 B

+34
View File
@@ -0,0 +1,34 @@
{
"manifest_version": 3,
"name": "agent-browser-stealth",
"version": "0.2.0",
"description": "Session-aware tab grouping and coordination for CDP-driven agent-browser workflows.",
"icons": {
"128": "icons/icon.svg"
},
"permissions": ["tabs", "tabGroups", "downloads", "storage", "sidePanel", "alarms"],
"host_permissions": ["<all_urls>"],
"background": {
"service_worker": "service-worker.js"
},
"action": {
"default_title": "agent-browser-stealth"
},
"side_panel": {
"default_path": "sidepanel.html"
},
"content_scripts": [
{
"matches": ["<all_urls>"],
"js": ["content-script.js"],
"run_at": "document_start",
"match_about_blank": true
}
],
"web_accessible_resources": [
{
"resources": ["page-bridge.js"],
"matches": ["<all_urls>"]
}
]
}
+133
View File
@@ -0,0 +1,133 @@
(() => {
if (window.__AB_STEALTH_BRIDGE_INSTALLED__) return;
window.__AB_STEALTH_BRIDGE_INSTALLED__ = true;
const currentScript = document.currentScript;
const TYPE = currentScript?.dataset?.abBridgeEvent || 'AB_PAGE_BRIDGE_EVENT';
const post = (kind, payload) => {
try {
window.postMessage({ type: TYPE, kind, payload, timestamp: Date.now() }, '*');
} catch {
// Ignore post failures.
}
};
const serializeArg = (value, depth = 0) => {
if (value === null || typeof value === 'undefined') return value;
if (typeof value === 'string') return value.slice(0, 250);
if (typeof value === 'number' || typeof value === 'boolean') return value;
if (value instanceof Error) return `${value.name}: ${value.message}`;
if (depth > 2) return '[depth-limit]';
if (Array.isArray(value)) return value.slice(0, 10).map((item) => serializeArg(item, depth + 1));
if (typeof value === 'object') {
const out = {};
const entries = Object.entries(value).slice(0, 12);
for (const [k, v] of entries) {
out[k] = serializeArg(v, depth + 1);
}
return out;
}
return String(value).slice(0, 250);
};
const patchConsoleMethod = (name) => {
const original = console[name];
if (typeof original !== 'function') return;
console[name] = function patchedConsole(...args) {
post('console', {
level: name,
args: args.map((arg) => serializeArg(arg)),
});
return original.apply(this, args);
};
};
patchConsoleMethod('error');
patchConsoleMethod('warn');
window.addEventListener('error', (event) => {
post('console', {
level: 'error',
message: event.message,
source: event.filename,
line: event.lineno,
column: event.colno,
});
});
window.addEventListener('unhandledrejection', (event) => {
post('console', {
level: 'error',
message: 'Unhandled rejection',
reason: serializeArg(event.reason),
});
});
if (typeof window.fetch === 'function') {
const originalFetch = window.fetch.bind(window);
window.fetch = async (...args) => {
const startedAt = Date.now();
const requestInfo = args[0];
const requestInit = args[1] || {};
const method = requestInit.method || 'GET';
const url = typeof requestInfo === 'string' ? requestInfo : requestInfo?.url || '';
try {
const response = await originalFetch(...args);
post('network', {
transport: 'fetch',
method,
url,
status: response.status,
ok: response.ok,
durationMs: Date.now() - startedAt,
});
return response;
} catch (error) {
post('network', {
transport: 'fetch',
method,
url,
error: serializeArg(error),
durationMs: Date.now() - startedAt,
});
throw error;
}
};
}
if (typeof window.XMLHttpRequest === 'function') {
const originalOpen = XMLHttpRequest.prototype.open;
const originalSend = XMLHttpRequest.prototype.send;
XMLHttpRequest.prototype.open = function patchedOpen(method, url, ...rest) {
this.__abRequestMeta = {
method: typeof method === 'string' ? method : 'GET',
url: typeof url === 'string' ? url : String(url || ''),
startedAt: Date.now(),
};
return originalOpen.call(this, method, url, ...rest);
};
XMLHttpRequest.prototype.send = function patchedSend(...args) {
this.addEventListener('loadend', () => {
const meta = this.__abRequestMeta || {};
post('network', {
transport: 'xhr',
method: meta.method || 'GET',
url: meta.url || '',
status: this.status,
ok: this.status >= 200 && this.status < 400,
durationMs: Date.now() - (meta.startedAt || Date.now()),
});
});
return originalSend.apply(this, args);
};
}
post('lifecycle', {
event: 'bridge-installed',
href: location.href,
});
})();
File diff suppressed because it is too large Load Diff
+258
View File
@@ -0,0 +1,258 @@
:root {
--bg: #f3f5f7;
--surface: #ffffff;
--surface-alt: #f6f8fb;
--primary: #1769e0;
--primary-hover: #0f58c0;
--border: #d8dde4;
--text-main: #18212f;
--text-secondary: #4a5568;
--text-muted: #667287;
--success: #117a3d;
--warning: #b86d00;
--danger: #bd1e24;
--radius: 10px;
--mono: "SFMono-Regular", Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace;
}
* {
box-sizing: border-box;
}
body {
margin: 0;
padding: 14px;
background: var(--bg);
color: var(--text-main);
font-family: "SF Pro Text", -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
font-size: 13px;
line-height: 1.4;
}
header {
display: flex;
flex-direction: column;
gap: 10px;
margin-bottom: 14px;
}
h1 {
margin: 0;
font-size: 19px;
font-weight: 650;
}
h3 {
margin: 0 0 12px 0;
font-size: 11px;
letter-spacing: 0.08em;
text-transform: uppercase;
color: var(--text-secondary);
}
h4 {
margin: 0;
font-size: 14px;
}
.card {
background: var(--surface);
border: 1px solid var(--border);
border-radius: var(--radius);
padding: 14px;
margin-bottom: 12px;
}
.actions {
display: flex;
gap: 8px;
}
button {
all: unset;
display: inline-flex;
align-items: center;
justify-content: center;
border: 1px solid var(--border);
border-radius: 6px;
padding: 6px 10px;
background: var(--surface);
color: var(--primary);
cursor: pointer;
font-size: 12px;
font-weight: 600;
}
button:hover {
border-color: var(--primary);
background: #eef4ff;
}
button.primary {
background: var(--primary);
border-color: var(--primary);
color: #fff;
}
button.primary:hover {
background: var(--primary-hover);
}
button.danger {
color: var(--danger);
}
.status-line {
min-height: 18px;
color: var(--text-secondary);
font-size: 12px;
}
.status-line.ok {
color: var(--success);
}
.status-line.warn {
color: var(--warning);
}
.status-line.error {
color: var(--danger);
}
.row {
display: flex;
gap: 8px;
align-items: center;
}
.row.wrap {
flex-wrap: wrap;
}
.row + .row {
margin-top: 8px;
}
input,
select,
textarea {
width: 100%;
padding: 7px 9px;
border-radius: 6px;
border: 1px solid var(--border);
background: var(--surface-alt);
color: var(--text-main);
font-size: 12px;
}
input.mono,
textarea.mono,
code,
pre {
font-family: var(--mono);
}
pre {
margin: 0;
background: #0f172a;
color: #dce6fb;
border-radius: 8px;
padding: 10px;
max-height: 220px;
overflow: auto;
font-size: 11px;
}
.tags {
display: flex;
gap: 6px;
flex-wrap: wrap;
margin-top: 8px;
}
.tag {
border-radius: 999px;
border: 1px solid var(--border);
padding: 2px 8px;
font-size: 11px;
color: var(--text-secondary);
background: var(--surface-alt);
}
.stack {
display: flex;
flex-direction: column;
gap: 10px;
}
.list {
display: flex;
flex-direction: column;
gap: 6px;
}
.item {
border: 1px solid var(--border);
border-radius: 8px;
padding: 8px;
background: var(--surface-alt);
}
.item-title {
font-weight: 600;
color: var(--text-main);
}
.item-url {
font-size: 11px;
color: var(--text-muted);
white-space: nowrap;
overflow: hidden;
text-overflow: ellipsis;
}
.caption {
font-size: 11px;
color: var(--text-muted);
}
hr {
border: none;
border-top: 1px solid var(--border);
margin: 12px 0;
}
.section-title {
display: flex;
justify-content: space-between;
align-items: center;
margin-bottom: 8px;
}
.grid-2 {
display: grid;
grid-template-columns: 1fr 1fr;
gap: 8px;
}
.event-pill {
display: inline-flex;
align-items: center;
border-radius: 999px;
border: 1px solid var(--border);
padding: 1px 7px;
font-size: 10px;
color: var(--text-secondary);
background: #f0f4fa;
}
.muted {
color: var(--text-muted);
}
.empty {
text-align: center;
color: var(--text-muted);
padding: 14px 0;
}
+28
View File
@@ -0,0 +1,28 @@
<!doctype html>
<html lang="en">
<head>
<meta charset="UTF-8" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
<title>agent-browser-stealth panel</title>
<link rel="stylesheet" href="sidepanel.css" />
</head>
<body>
<header>
<h1>agent-browser-stealth</h1>
<div class="actions">
<button id="refresh-btn" type="button">Refresh</button>
<button id="cleanup-btn" type="button">Clean Empty Groups</button>
</div>
<div id="status-line" class="status-line"></div>
</header>
<section id="control" class="card"></section>
<section id="summary" class="card"></section>
<section id="automation" class="card"></section>
<section id="developer" class="card"></section>
<section id="sessions" class="stack"></section>
<section id="downloads" class="card"></section>
<script src="sidepanel.js"></script>
</body>
</html>
File diff suppressed because it is too large Load Diff
+35 -11
View File
@@ -1,7 +1,7 @@
{
"name": "agent-browser",
"version": "0.16.1",
"description": "Headless browser automation CLI for AI agents",
"name": "agent-browser-stealth",
"version": "0.17.0-fork.2",
"description": "Stealth browser automation CLI for AI agents with anti-bot evasions",
"type": "module",
"main": "dist/daemon.js",
"files": [
@@ -11,20 +11,23 @@
"skills"
],
"bin": {
"agent-browser": "./bin/agent-browser.js"
"agent-browser-stealth": "./bin/agent-browser.js",
"agent-browser": "./bin/agent-browser.js",
"abs": "./bin/agent-browser.js"
},
"scripts": {
"prepare": "husky",
"version:sync": "node scripts/sync-version.js",
"version": "npm run version:sync && git add cli/Cargo.toml",
"native:clean": "cargo clean --manifest-path cli/Cargo.toml -p agent-browser-stealth",
"build": "tsc",
"build:native": "npm run version:sync && cargo build --release --manifest-path cli/Cargo.toml && node scripts/copy-native.js",
"build:native": "npm run version:sync && npm run native:clean && cargo build --release --manifest-path cli/Cargo.toml && node scripts/copy-native.js",
"build:linux": "npm run version:sync && docker compose -f docker/docker-compose.yml run --rm build-linux",
"build:macos": "npm run version:sync && (cargo build --release --manifest-path cli/Cargo.toml --target aarch64-apple-darwin & cargo build --release --manifest-path cli/Cargo.toml --target x86_64-apple-darwin & wait) && cp cli/target/aarch64-apple-darwin/release/agent-browser bin/agent-browser-darwin-arm64 && cp cli/target/x86_64-apple-darwin/release/agent-browser bin/agent-browser-darwin-x64",
"build:macos": "npm run version:sync && npm run native:clean && (cargo build --release --manifest-path cli/Cargo.toml --target aarch64-apple-darwin & cargo build --release --manifest-path cli/Cargo.toml --target x86_64-apple-darwin & wait) && node scripts/copy-native.js cli/target/aarch64-apple-darwin/release/agent-browser bin/agent-browser-darwin-arm64 && node scripts/copy-native.js cli/target/x86_64-apple-darwin/release/agent-browser bin/agent-browser-darwin-x64",
"build:windows": "npm run version:sync && docker compose -f docker/docker-compose.yml run --rm build-windows",
"build:all-platforms": "npm run version:sync && (npm run build:linux & npm run build:windows & wait) && npm run build:macos",
"build:docker": "docker build -t agent-browser-builder -f docker/Dockerfile.build .",
"release": "npm run version:sync && npm run build && npm run build:all-platforms && npm publish",
"release": "pnpm run version:sync && pnpm run build && pnpm run build:all-platforms && pnpm run verify:bundled-binaries && pnpm run verify:native-version && pnpm publish",
"start": "node dist/daemon.js",
"dev": "tsx src/daemon.ts",
"typecheck": "tsc --noEmit",
@@ -33,15 +36,33 @@
"test": "vitest run",
"test:watch": "vitest",
"test:e2e:dogfood": "vitest run test/e2e/dogfood.eval.ts",
"bench": "pnpm build && tsx test/benchmarks/run.ts",
"bench:node": "pnpm build && tsx test/benchmarks/run.ts --node-only",
"bench:native": "pnpm build && tsx test/benchmarks/run.ts --native-only",
"bench:engine": "pnpm build && tsx test/benchmarks/run.ts --engine",
"check:daemon-pid-recovery": "node scripts/check-daemon-pid-recovery.js",
"check:stealth-regression": "node scripts/check-stealth-regression.js",
"check:turnstile-testkey": "pnpm exec tsx scripts/check-turnstile-testkey.ts",
"postinstall": "node scripts/postinstall.js",
"verify:native-version": "node scripts/verify-native-version.js",
"verify:bundled-binaries": "node scripts/verify-bundled-binaries.js",
"verify:packed-host-binary": "node scripts/verify-packed-host-binary.js",
"verify:registry-host-binary": "node scripts/verify-registry-host-binary.js",
"prepublishOnly": "pnpm run verify:bundled-binaries && pnpm run verify:native-version && pnpm run verify:packed-host-binary",
"clawhub:sync": "bash scripts/clawhub-sync.sh",
"sync:upstream": "bash scripts/sync-upstream.sh",
"sync:upstream:push": "bash scripts/sync-upstream.sh --push",
"changeset": "changeset",
"ci:version": "changeset version && pnpm run version:sync && pnpm install --no-frozen-lockfile",
"ci:publish": "pnpm run version:sync && pnpm run build && changeset publish"
"ci:publish": "pnpm run version:sync && pnpm run build && pnpm run build:native && pnpm run verify:bundled-binaries && pnpm run verify:native-version && pnpm run verify:packed-host-binary && changeset publish"
},
"keywords": [
"browser",
"automation",
"headless",
"stealth",
"anti-bot",
"anti-detection",
"playwright",
"cli",
"agent"
@@ -49,12 +70,15 @@
"license": "Apache-2.0",
"repository": {
"type": "git",
"url": "git+https://github.com/vercel-labs/agent-browser.git"
"url": "git+https://github.com/leeguooooo/agent-browser.git"
},
"publishConfig": {
"tag": "fork"
},
"bugs": {
"url": "https://github.com/vercel-labs/agent-browser/issues"
"url": "https://github.com/leeguooooo/agent-browser/issues"
},
"homepage": "https://github.com/vercel-labs/agent-browser#readme",
"homepage": "https://github.com/leeguooooo/agent-browser#readme",
"dependencies": {
"node-simctl": "^7.4.0",
"playwright-core": "^1.57.0",

Some files were not shown because too many files have changed in this diff Show More