Compare commits
63
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5970579d7c | ||
|
|
8880aa2f35 | ||
|
|
f051e72f85 | ||
|
|
6ae703565c | ||
|
|
d56442cf91 | ||
|
|
dad7be8c77 | ||
|
|
7921928ec4 | ||
|
|
a71198d591 | ||
|
|
19aefb7eb8 | ||
|
|
ff2e66c0a1 | ||
|
|
d320e1df47 | ||
|
|
7b454e14ed | ||
|
|
17d44baa2f | ||
|
|
59589043db | ||
|
|
9e7c8937a6 | ||
|
|
2d0d18f0d9 | ||
|
|
ae22bda46c | ||
|
|
d82357fea4 | ||
|
|
66a39f3c83 | ||
|
|
eedf824af7 | ||
|
|
c07eb7ee52 | ||
|
|
d32a1d046a | ||
|
|
0c0ed5e72c | ||
|
|
34092ec193 | ||
|
|
726377c4c1 | ||
|
|
8e2e4abce6 | ||
|
|
870895e922 | ||
|
|
2a766cfe48 | ||
|
|
d04cf59238 | ||
|
|
0a257ad2c1 | ||
|
|
44c0361fcd | ||
|
|
907ca8c808 | ||
|
|
74fda70b67 | ||
|
|
2a397de59f | ||
|
|
bf672ee7f9 | ||
|
|
74910cfef1 | ||
|
|
41830dff71 | ||
|
|
e005c7251b | ||
|
|
11eab471f1 | ||
|
|
aa256e30c7 | ||
|
|
6f1dd39121 | ||
|
|
85d18799a4 | ||
|
|
43e781a8d3 | ||
|
|
25e8719e51 | ||
|
|
ec011f46ff | ||
|
|
aef8fcc038 | ||
|
|
96582b79fd | ||
|
|
058a286326 | ||
|
|
b1f27236d8 | ||
|
|
a5a9327b7d | ||
|
|
699ccbd3cb | ||
|
|
893ddfd259 | ||
|
|
ea2e93dbba | ||
|
|
4c6afe3e69 | ||
|
|
9f9a90cf63 | ||
|
|
0443e4ed7a | ||
|
|
c5b2292caa | ||
|
|
2ed0c6f8ec | ||
|
|
3a91aef4c9 | ||
|
|
02ebc9f328 | ||
|
|
955543b757 | ||
|
|
ecad112707 | ||
|
|
8932f28926 |
@@ -0,0 +1,23 @@
|
||||
# Changesets
|
||||
|
||||
This project uses [Changesets](https://github.com/changesets/changesets) for versioning and changelog generation.
|
||||
|
||||
## Adding a changeset
|
||||
|
||||
When you make a change that should be released, run:
|
||||
|
||||
```bash
|
||||
pnpm changeset
|
||||
```
|
||||
|
||||
This will prompt you to:
|
||||
1. Select the type of change (patch, minor, major)
|
||||
2. Write a summary of your changes
|
||||
|
||||
The changeset file will be committed with your PR.
|
||||
|
||||
## Release process
|
||||
|
||||
When changesets are merged to `main`, the release workflow will:
|
||||
1. Create a "Version Packages" PR that updates version numbers and changelogs
|
||||
2. When that PR is merged, packages are automatically published to npm
|
||||
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"$schema": "https://unpkg.com/@changesets/config@3.1.1/schema.json",
|
||||
"changelog": "@changesets/cli/changelog",
|
||||
"commit": false,
|
||||
"fixed": [],
|
||||
"linked": [],
|
||||
"access": "public",
|
||||
"baseBranch": "main",
|
||||
"updateInternalDependencies": "patch",
|
||||
"ignore": []
|
||||
}
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"$schema": "https://anthropic.com/claude-code/marketplace.schema.json",
|
||||
"name": "agent-browser",
|
||||
"description": "Browser automation for AI agents",
|
||||
"description": "Headless browser automation for AI agents",
|
||||
"owner": {
|
||||
"name": "Vercel",
|
||||
"email": "support@vercel.com"
|
||||
|
||||
+114
-37
@@ -18,6 +18,43 @@ jobs:
|
||||
- name: Check version sync
|
||||
run: node scripts/check-version-sync.js
|
||||
|
||||
typescript:
|
||||
name: TypeScript (Node ${{ matrix.node-version }})
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
node-version: [20, 22]
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v4
|
||||
with:
|
||||
version: 9
|
||||
|
||||
- name: Setup Node.js ${{ matrix.node-version }}
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: ${{ matrix.node-version }}
|
||||
cache: pnpm
|
||||
|
||||
- name: Install dependencies
|
||||
run: pnpm install
|
||||
|
||||
- name: Typecheck
|
||||
run: pnpm typecheck
|
||||
|
||||
- name: Format check
|
||||
run: pnpm format:check
|
||||
|
||||
- name: Install Playwright browsers
|
||||
run: pnpm exec playwright install --with-deps chromium
|
||||
|
||||
- name: Run tests
|
||||
run: pnpm test
|
||||
|
||||
rust:
|
||||
name: Rust
|
||||
runs-on: ubuntu-latest
|
||||
@@ -27,20 +64,12 @@ 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
|
||||
|
||||
@@ -55,7 +84,7 @@ jobs:
|
||||
target: aarch64-apple-darwin
|
||||
- os: macos-latest
|
||||
target: x86_64-apple-darwin
|
||||
- os: windows-latest
|
||||
- os: windows-latest-8-cores
|
||||
target: x86_64-pc-windows-msvc
|
||||
|
||||
steps:
|
||||
@@ -75,40 +104,27 @@ jobs:
|
||||
- name: Run Rust tests
|
||||
run: cargo test --profile ci --manifest-path cli/Cargo.toml --target ${{ matrix.target }}
|
||||
|
||||
native-e2e:
|
||||
name: Native E2E Tests
|
||||
if: github.event_name != 'pull_request'
|
||||
runs-on: ubuntu-latest
|
||||
needs: rust
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Setup Rust toolchain
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
|
||||
- name: Cache Rust build artifacts
|
||||
uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
workspaces: cli
|
||||
|
||||
- name: Install Chrome
|
||||
run: |
|
||||
cargo run --manifest-path cli/Cargo.toml -- install --with-deps
|
||||
|
||||
- name: Run e2e tests
|
||||
run: cargo test --profile ci --manifest-path cli/Cargo.toml e2e -- --ignored --test-threads=1
|
||||
|
||||
windows-integration:
|
||||
name: Windows Integration Test
|
||||
if: github.event_name != 'pull_request'
|
||||
runs-on: windows-latest
|
||||
runs-on: windows-latest-8-cores
|
||||
needs: rust-cross
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v4
|
||||
with:
|
||||
version: 9
|
||||
|
||||
- name: Setup Node.js
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: 22
|
||||
cache: pnpm
|
||||
|
||||
- name: Setup Rust toolchain
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
with:
|
||||
@@ -122,6 +138,12 @@ jobs:
|
||||
- name: Build Rust CLI
|
||||
run: cargo build --release --manifest-path cli/Cargo.toml --target x86_64-pc-windows-msvc
|
||||
|
||||
- name: Install npm dependencies
|
||||
run: pnpm install
|
||||
|
||||
- name: Build TypeScript
|
||||
run: pnpm build
|
||||
|
||||
- name: Copy CLI binary to bin directory
|
||||
run: |
|
||||
Copy-Item cli/target/x86_64-pc-windows-msvc/release/agent-browser.exe bin/agent-browser-win32-x64.exe
|
||||
@@ -139,6 +161,18 @@ jobs:
|
||||
shell: pwsh
|
||||
timeout-minutes: 10
|
||||
|
||||
- name: Verify Chromium was installed
|
||||
run: |
|
||||
$playwrightPath = "$env:LOCALAPPDATA\ms-playwright"
|
||||
if (Test-Path $playwrightPath) {
|
||||
Write-Host "Playwright browsers installed at: $playwrightPath"
|
||||
Get-ChildItem $playwrightPath -Recurse -Depth 2 | Select-Object -First 20
|
||||
} else {
|
||||
Write-Error "Playwright browsers not found!"
|
||||
exit 1
|
||||
}
|
||||
shell: pwsh
|
||||
|
||||
- name: Test daemon lifecycle (open, snapshot, close)
|
||||
run: |
|
||||
$env:PATH = "$pwd\bin;$env:PATH"
|
||||
@@ -156,6 +190,37 @@ jobs:
|
||||
shell: pwsh
|
||||
timeout-minutes: 5
|
||||
|
||||
serverless-chromium:
|
||||
name: Serverless Chromium (@sparticuz/chromium)
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v4
|
||||
with:
|
||||
version: 9
|
||||
|
||||
- name: Setup Node.js
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: 22
|
||||
cache: pnpm
|
||||
|
||||
- name: Install dependencies
|
||||
run: pnpm install
|
||||
|
||||
- name: Install @sparticuz/chromium
|
||||
run: pnpm add -D @sparticuz/chromium
|
||||
|
||||
- name: Build TypeScript
|
||||
run: pnpm build
|
||||
|
||||
- name: Run serverless integration test
|
||||
run: pnpm exec vitest run test/serverless.test.ts
|
||||
|
||||
global-install:
|
||||
name: Global Install (${{ matrix.os }})
|
||||
if: github.event_name != 'pull_request'
|
||||
@@ -170,7 +235,7 @@ jobs:
|
||||
- os: macos-latest
|
||||
target: aarch64-apple-darwin
|
||||
binary: agent-browser-darwin-arm64
|
||||
- os: windows-latest
|
||||
- os: windows-latest-8-cores
|
||||
target: x86_64-pc-windows-msvc
|
||||
binary: agent-browser-win32-x64.exe
|
||||
|
||||
@@ -178,10 +243,16 @@ jobs:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v4
|
||||
with:
|
||||
version: 9
|
||||
|
||||
- name: Setup Node.js
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: 22
|
||||
cache: pnpm
|
||||
|
||||
- name: Setup Rust toolchain
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
@@ -196,6 +267,12 @@ jobs:
|
||||
- name: Build Rust CLI
|
||||
run: cargo build --release --manifest-path cli/Cargo.toml --target ${{ matrix.target }}
|
||||
|
||||
- name: Install npm dependencies
|
||||
run: pnpm install
|
||||
|
||||
- name: Build TypeScript
|
||||
run: pnpm build
|
||||
|
||||
- name: Copy CLI binary to bin directory (Unix)
|
||||
if: runner.os != 'Windows'
|
||||
run: cp cli/target/${{ matrix.target }}/release/agent-browser bin/${{ matrix.binary }}
|
||||
@@ -222,7 +299,7 @@ jobs:
|
||||
echo "ERROR: Symlink should point to native binary, not JS wrapper"
|
||||
exit 1
|
||||
fi
|
||||
echo "Symlink correctly points to native binary"
|
||||
echo "✓ Symlink correctly points to native binary"
|
||||
shell: bash
|
||||
|
||||
- name: Verify shim points to native binary (Windows)
|
||||
@@ -237,5 +314,5 @@ jobs:
|
||||
echo "ERROR: Shim should point to native .exe, not JS wrapper"
|
||||
exit 1
|
||||
}
|
||||
echo "Shim correctly points to native binary"
|
||||
echo "✓ Shim correctly points to native binary"
|
||||
shell: pwsh
|
||||
|
||||
@@ -10,40 +10,13 @@ concurrency: ${{ github.workflow }}-${{ github.ref }}
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
pull-requests: write
|
||||
id-token: write
|
||||
|
||||
jobs:
|
||||
check-release:
|
||||
name: Check for new version
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
should_release: ${{ steps.check.outputs.should_release }}
|
||||
version: ${{ steps.check.outputs.version }}
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Compare package.json version to npm
|
||||
id: check
|
||||
run: |
|
||||
LOCAL_VERSION=$(node -p "require('./package.json').version")
|
||||
echo "Local version: $LOCAL_VERSION"
|
||||
|
||||
NPM_VERSION=$(npm view agent-browser version 2>/dev/null || echo "0.0.0")
|
||||
echo "npm version: $NPM_VERSION"
|
||||
|
||||
if [ "$LOCAL_VERSION" != "$NPM_VERSION" ]; then
|
||||
echo "Version changed: $NPM_VERSION -> $LOCAL_VERSION"
|
||||
echo "should_release=true" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "Version unchanged, skipping release"
|
||||
echo "should_release=false" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
echo "version=$LOCAL_VERSION" >> "$GITHUB_OUTPUT"
|
||||
|
||||
# Build native binaries for all platforms first
|
||||
build-binaries:
|
||||
name: Build ${{ matrix.name }}
|
||||
needs: check-release
|
||||
if: needs.check-release.outputs.should_release == 'true'
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
fail-fast: false
|
||||
@@ -59,16 +32,6 @@ jobs:
|
||||
target: aarch64-unknown-linux-gnu
|
||||
binary: agent-browser-linux-arm64
|
||||
use_zigbuild: true
|
||||
- name: Linux musl x64
|
||||
os: ubuntu-latest
|
||||
target: x86_64-unknown-linux-musl
|
||||
binary: agent-browser-linux-musl-x64
|
||||
use_zigbuild: true
|
||||
- name: Linux musl ARM64
|
||||
os: ubuntu-latest
|
||||
target: aarch64-unknown-linux-musl
|
||||
binary: agent-browser-linux-musl-arm64
|
||||
use_zigbuild: true
|
||||
- name: Windows x64
|
||||
os: ubuntu-latest
|
||||
target: x86_64-pc-windows-gnu
|
||||
@@ -165,13 +128,19 @@ jobs:
|
||||
path: artifacts/${{ matrix.binary }}
|
||||
retention-days: 7
|
||||
|
||||
publish:
|
||||
name: Publish to npm
|
||||
needs: [check-release, build-binaries]
|
||||
# Create release PR or publish to npm (with binaries)
|
||||
release:
|
||||
name: Release
|
||||
needs: build-binaries
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
published: ${{ steps.publish_metadata.outputs.published }}
|
||||
publishedPackages: ${{ steps.publish_metadata.outputs.publishedPackages }}
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
- name: Checkout Repo
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v4
|
||||
@@ -183,9 +152,8 @@ jobs:
|
||||
with:
|
||||
node-version: '22'
|
||||
cache: pnpm
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
|
||||
- name: Install dependencies
|
||||
- name: Install Dependencies
|
||||
run: pnpm install --frozen-lockfile
|
||||
|
||||
- name: Download all binary artifacts
|
||||
@@ -207,13 +175,11 @@ jobs:
|
||||
EXPECTED_BINARIES=(
|
||||
"agent-browser-linux-x64"
|
||||
"agent-browser-linux-arm64"
|
||||
"agent-browser-linux-musl-x64"
|
||||
"agent-browser-linux-musl-arm64"
|
||||
"agent-browser-win32-x64.exe"
|
||||
"agent-browser-darwin-x64"
|
||||
"agent-browser-darwin-arm64"
|
||||
)
|
||||
MIN_SIZE=100000
|
||||
MIN_SIZE=100000 # Binaries should be at least 100KB
|
||||
ERRORS=0
|
||||
for binary in "${EXPECTED_BINARIES[@]}"; do
|
||||
if [ ! -f "bin/$binary" ]; then
|
||||
@@ -233,20 +199,81 @@ jobs:
|
||||
echo "Error: $ERRORS binary issues found"
|
||||
exit 1
|
||||
fi
|
||||
echo "All 7 platform binaries present and valid"
|
||||
echo "All 5 platform binaries present and valid"
|
||||
|
||||
- name: Publish to npm
|
||||
run: pnpm publish --no-git-checks
|
||||
- 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
|
||||
title: 'chore: version packages'
|
||||
commit: 'chore: version packages'
|
||||
env:
|
||||
NODE_AUTH_TOKEN: ${{ secrets.NPM_VERCEL_TOKEN_ELEVATED }}
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
- 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:
|
||||
name: Create GitHub Release
|
||||
needs: [check-release, publish]
|
||||
needs: release
|
||||
if: needs.release.outputs.published == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
- name: Checkout Repo
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
ref: main
|
||||
|
||||
- name: Download all artifacts
|
||||
uses: actions/download-artifact@v4
|
||||
@@ -264,59 +291,28 @@ jobs:
|
||||
- name: Verify binaries exist
|
||||
run: |
|
||||
BINARY_COUNT=$(ls bin/agent-browser-* 2>/dev/null | wc -l)
|
||||
if [ "$BINARY_COUNT" -lt 7 ]; then
|
||||
echo "Error: Expected 7 binaries, found $BINARY_COUNT"
|
||||
if [ "$BINARY_COUNT" -lt 5 ]; then
|
||||
echo "Error: Expected 5 binaries, found $BINARY_COUNT"
|
||||
ls -la bin/
|
||||
exit 1
|
||||
fi
|
||||
echo "Found $BINARY_COUNT binaries"
|
||||
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v4
|
||||
with:
|
||||
version: 9
|
||||
|
||||
- name: Setup Node.js
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: '22'
|
||||
cache: pnpm
|
||||
|
||||
- name: Install dependencies
|
||||
run: pnpm install --frozen-lockfile
|
||||
|
||||
- name: Build dashboard
|
||||
run: pnpm --filter dashboard build
|
||||
|
||||
- name: Create dashboard.zip
|
||||
run: cd packages/dashboard/out && zip -r ../../../dashboard.zip .
|
||||
|
||||
- name: Extract changelog entry
|
||||
run: |
|
||||
VERSION="${{ needs.check-release.outputs.version }}"
|
||||
awk '/<!-- release:start -->/{found=1; next} /<!-- release:end -->/{found=0} found{print}' CHANGELOG.md > /tmp/release-notes.md
|
||||
|
||||
LINES=$(wc -l < /tmp/release-notes.md | tr -d ' ')
|
||||
if [ "$LINES" -lt 2 ]; then
|
||||
echo "Error: No release notes found between <!-- release:start --> and <!-- release:end --> markers in CHANGELOG.md"
|
||||
exit 1
|
||||
fi
|
||||
echo "Extracted release notes for $VERSION ($LINES lines)"
|
||||
|
||||
- name: Create GitHub Release
|
||||
run: |
|
||||
VERSION="${{ needs.check-release.outputs.version }}"
|
||||
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 assets..."
|
||||
gh release upload "$TAG" bin/agent-browser-* dashboard.zip --clobber
|
||||
echo "Release $TAG already exists, uploading binaries..."
|
||||
gh release upload "$TAG" bin/agent-browser-* --clobber
|
||||
else
|
||||
echo "Creating release $TAG..."
|
||||
gh release create "$TAG" \
|
||||
--title "$TAG" \
|
||||
--notes-file /tmp/release-notes.md \
|
||||
bin/agent-browser-* dashboard.zip
|
||||
--generate-notes \
|
||||
bin/agent-browser-*
|
||||
fi
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
+3
-11
@@ -4,11 +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/.install-method
|
||||
!bin/agent-browser
|
||||
!bin/agent-browser.cmd
|
||||
bin/agent-browser
|
||||
bin/*.d
|
||||
|
||||
# Rust build artifacts
|
||||
cli/target/
|
||||
@@ -46,9 +45,6 @@ yarn.lock
|
||||
.env
|
||||
.env.local
|
||||
|
||||
# Windows debug instance config
|
||||
scripts/windows-debug/.instance
|
||||
|
||||
# opensrc - source code for packages
|
||||
opensrc/
|
||||
|
||||
@@ -60,7 +56,3 @@ docs/package-lock.json
|
||||
|
||||
# pnpm
|
||||
.pnpm-store/
|
||||
|
||||
# next
|
||||
.next/
|
||||
out/
|
||||
|
||||
@@ -1,2 +1,3 @@
|
||||
pnpm lint-staged
|
||||
node scripts/sync-version.js
|
||||
git add cli/Cargo.toml cli/Cargo.lock
|
||||
|
||||
@@ -9,7 +9,6 @@ This project uses **pnpm**. Always use `pnpm` instead of `npm` or `yarn` for ins
|
||||
## Code Style
|
||||
|
||||
- Do not use emojis in code, output, or documentation. Unicode symbols (✓, ✗, →, ⚠) are acceptable.
|
||||
- In documentation and markdown, never use double hyphens (`--`) as a dash. Use an emdash (—) sparingly when needed. Prefer rewriting the sentence to avoid dashes entirely.
|
||||
- CLI colored output uses `cli/src/color.rs`. This module respects the `NO_COLOR` environment variable. Never use hardcoded ANSI color codes.
|
||||
- CLI flags must always use kebab-case (e.g., `--auto-connect`, `--allow-file-access`). Never use camelCase for flags (e.g., `--autoConnect` is wrong).
|
||||
|
||||
@@ -17,80 +16,35 @@ This project uses **pnpm**. Always use `pnpm` instead of `npm` or `yarn` for ins
|
||||
|
||||
When adding or changing user-facing features (new flags, commands, behaviors, environment variables, etc.), update **all** of the following:
|
||||
|
||||
1. `cli/src/output.rs` — `--help` output (flags list, examples, environment variables)
|
||||
2. `README.md` — Options table, relevant feature sections, examples
|
||||
3. `skills/agent-browser/SKILL.md` — so AI agents know about the feature
|
||||
4. `docs/src/app/` — the Next.js docs site (MDX pages)
|
||||
1. `cli/src/output.rs` -- `--help` output (flags list, examples, environment variables)
|
||||
2. `README.md` -- Options table, relevant feature sections, examples
|
||||
3. `skills/agent-browser/SKILL.md` -- so AI agents know about the feature
|
||||
4. `docs/src/app/` -- the Next.js docs site (MDX pages)
|
||||
5. Inline doc comments in the relevant source files
|
||||
|
||||
This applies to changes that either human users or AI agents would need to know about. Do not skip any of these locations.
|
||||
|
||||
In the `docs/src/app/` MDX files, always use HTML `<table>` syntax for tables (not markdown pipe tables). This matches the existing convention across the docs site.
|
||||
|
||||
## Dashboard (packages/dashboard)
|
||||
## Dual Architecture (Node.js + Native)
|
||||
|
||||
- Never use native browser dialogs (`alert`, `confirm`, `prompt`). Use shadcn/ui components (`Dialog`, `AlertDialog`, etc.) instead.
|
||||
- Use param-case (kebab-case) for all file and folder names (e.g., `session-tree.tsx`, not `SessionTree.tsx`). The `ui/` directory follows shadcn conventions which already uses param-case.
|
||||
The codebase has two daemon implementations:
|
||||
|
||||
## Releasing
|
||||
- **Node.js/Playwright** (default) -- `src/daemon.ts`, `src/actions.ts`, `src/browser.ts`, and the rest of `src/`
|
||||
- **Rust/Native** (experimental, `--native` or `AGENT_BROWSER_NATIVE=1`) -- `cli/src/native/daemon.rs`, `cli/src/native/actions.rs`, `cli/src/native/browser.rs`, and the rest of `cli/src/native/`
|
||||
|
||||
Releases are manual, single-PR affairs. There is no changesets automation. The maintainer controls the changelog voice and format.
|
||||
When modifying browser automation logic (commands, actions, protocol handling), changes **must** be made in **both** paths:
|
||||
|
||||
To prepare a release:
|
||||
| Node.js Path | Native Path |
|
||||
|---|---|
|
||||
| `src/actions.ts` | `cli/src/native/actions.rs` |
|
||||
| `src/browser.ts` | `cli/src/native/browser.rs` |
|
||||
| `src/daemon.ts` | `cli/src/native/daemon.rs` |
|
||||
| `src/protocol.ts` | `cli/src/native/cdp/client.rs` |
|
||||
| `src/snapshot.ts` | `cli/src/native/snapshot.rs` |
|
||||
| `src/state-utils.ts` | `cli/src/native/state.rs` |
|
||||
|
||||
1. Create a branch (e.g. `prepare-v0.24.0`)
|
||||
2. Bump `version` in `package.json`
|
||||
3. Run `pnpm version:sync` to update `cli/Cargo.toml`, `cli/Cargo.lock`, and `packages/dashboard/package.json`
|
||||
4. Write the changelog entry in `CHANGELOG.md` at the top, under a new `## <version>` heading, wrapped in `<!-- release:start -->` and `<!-- release:end -->` markers
|
||||
5. Add a matching entry to `docs/src/app/changelog/page.mdx` at the top (below the `# Changelog` heading)
|
||||
6. Open a PR and merge to `main`
|
||||
|
||||
When the PR merges, CI compares `package.json` version to what's on npm. If it differs, it builds all 7 platform binaries, publishes to npm, and creates the GitHub release automatically. The GitHub release body is extracted from the content between the `<!-- release:start -->` and `<!-- release:end -->` markers in `CHANGELOG.md`.
|
||||
|
||||
### Writing the changelog
|
||||
|
||||
Review the git log since the last release and write the entry in `CHANGELOG.md`. Follow the existing format and voice. Group changes under `### New Features`, `### Bug Fixes`, `### Improvements`, etc. Bold the feature/fix name, then describe it concisely. Reference PR numbers in parentheses.
|
||||
|
||||
Wrap the release notes (everything between the `## <version>` heading and the previous version) in markers so CI can extract them for the GitHub release:
|
||||
|
||||
```markdown
|
||||
## 0.24.0
|
||||
|
||||
<!-- release:start -->
|
||||
### New Features
|
||||
|
||||
- **Foo command** - Added `foo` command for bar (#1234)
|
||||
|
||||
### Bug Fixes
|
||||
|
||||
- Fixed **baz** not working when qux is enabled (#1235)
|
||||
|
||||
### Contributors
|
||||
|
||||
- @ctate
|
||||
- @somecontributor
|
||||
<!-- release:end -->
|
||||
|
||||
## 0.23.3
|
||||
```
|
||||
|
||||
Include a `### Contributors` section listing the GitHub usernames (with `@` prefix) of everyone who contributed to the release. Check the git log between the previous tag and HEAD to find them.
|
||||
|
||||
Do not prefix entries with commit hashes. Do not use the changesets `### Patch Changes` / `### Minor Changes` headings. Use descriptive section names instead.
|
||||
|
||||
### Docs changelog
|
||||
|
||||
The docs changelog at `docs/src/app/changelog/page.mdx` mirrors `CHANGELOG.md` but uses a slightly different format. Each entry uses:
|
||||
|
||||
- A `v` prefix on the version (e.g. `## v0.24.0`)
|
||||
- A date line with the full date: `<p className="text-[#888] text-sm">March 30, 2026</p>`
|
||||
- A `---` separator between entries
|
||||
|
||||
Match the existing style in that file.
|
||||
|
||||
## Architecture
|
||||
|
||||
This is a Rust codebase. The browser automation daemon lives in `cli/src/native/` (daemon, actions, browser, CDP client, snapshot, state). The `--engine` flag selects Chrome vs Lightpanda. The `install` command downloads Chrome from Chrome for Testing directly.
|
||||
New commands must be implemented in both paths, or explicitly stubbed in the native path with a clear `"Not yet implemented: {action}"` error. The goal is eventual full migration to native, but until then both paths must stay in sync.
|
||||
|
||||
## Testing
|
||||
|
||||
@@ -123,68 +77,6 @@ cd cli && cargo fmt -- --check # Check formatting
|
||||
cd cli && cargo clippy # Lint
|
||||
```
|
||||
|
||||
## Windows Debugging
|
||||
|
||||
A remote Windows Server 2022 EC2 instance is available for debugging Windows-specific issues. It uses AWS Systems Manager (SSM) with no SSH or open ports. Commands run via `aws ssm send-command` and return stdout/stderr.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
The instance must be provisioned first (one-time, by a human):
|
||||
|
||||
```bash
|
||||
./scripts/windows-debug/provision.sh
|
||||
```
|
||||
|
||||
Requires: AWS CLI v2 configured with `ec2:*`, `iam:CreateRole`, `iam:AttachRolePolicy`, `ssm:SendCommand`, `ssm:GetCommandInvocation` permissions and a default VPC.
|
||||
|
||||
### Usage
|
||||
|
||||
Start the instance (if stopped):
|
||||
|
||||
```bash
|
||||
./scripts/windows-debug/start.sh
|
||||
```
|
||||
|
||||
Run a command on Windows:
|
||||
|
||||
```bash
|
||||
./scripts/windows-debug/run.sh "<powershell-command>"
|
||||
```
|
||||
|
||||
Sync the current git branch and rebuild:
|
||||
|
||||
```bash
|
||||
./scripts/windows-debug/sync.sh
|
||||
```
|
||||
|
||||
Stop the instance when done (avoids cost):
|
||||
|
||||
```bash
|
||||
./scripts/windows-debug/stop.sh
|
||||
```
|
||||
|
||||
### Common Workflows
|
||||
|
||||
Run unit tests on Windows:
|
||||
|
||||
```bash
|
||||
./scripts/windows-debug/run.sh "cd C:\agent-browser && cargo test --manifest-path cli\Cargo.toml"
|
||||
```
|
||||
|
||||
Run e2e tests on Windows:
|
||||
|
||||
```bash
|
||||
./scripts/windows-debug/run.sh "cd C:\agent-browser && cargo test e2e --manifest-path cli\Cargo.toml -- --ignored --test-threads=1"
|
||||
```
|
||||
|
||||
Check bootstrap progress (first boot only):
|
||||
|
||||
```bash
|
||||
./scripts/windows-debug/run.sh "Get-Content C:\bootstrap.log"
|
||||
```
|
||||
|
||||
The repo lives at `C:\agent-browser` on the instance. Rust, Git, and Chrome are pre-installed. The `run.sh` wrapper automatically adds cargo and git to PATH.
|
||||
|
||||
<!-- opensrc:start -->
|
||||
|
||||
## Source Code Reference
|
||||
|
||||
+265
@@ -0,0 +1,265 @@
|
||||
# agent-browser
|
||||
|
||||
## 0.16.3-fork.1
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 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.1-fork.5
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 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-fork.0
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 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
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 7bd8ce9: Added support for chrome:// and chrome-extension:// URLs in navigation and recording commands. These special browser URLs are now preserved as-is instead of having https:// incorrectly prepended.
|
||||
|
||||
## 0.15.0
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 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
|
||||
|
||||
### Minor Changes
|
||||
|
||||
- b7665e5: - Added `keyboard` command for raw keyboard input -- type with real keystrokes, insert text, and press shortcuts at the currently focused element without needing a selector.
|
||||
- Added `--color-scheme` flag and `AGENT_BROWSER_COLOR_SCHEME` env var for persistent dark/light mode preference across browser sessions.
|
||||
- Fixed IPC EAGAIN errors (os error 35/11) by adding backpressure-aware socket writes, command serialization, and lowering the default Playwright timeout to 25s (configurable via `AGENT_BROWSER_DEFAULT_TIMEOUT`).
|
||||
- Fixed remote debugging (CDP) reconnection.
|
||||
- Fixed state load failing when no browser is running.
|
||||
- Fixed `--annotate` flag warning appearing when not explicitly passed via CLI.
|
||||
|
||||
## 0.13.0
|
||||
|
||||
### Minor Changes
|
||||
|
||||
- ebd8717: Added new diff commands for comparing snapshots, screenshots, and URLs between page states. You can now run visual pixel diffs against baseline images, compare accessibility tree snapshots with customizable depth and selectors, and diff two URLs side-by-side with optional screenshot comparison.
|
||||
|
||||
## 0.12.0
|
||||
|
||||
### Minor Changes
|
||||
|
||||
- 69ffad0: Add annotated screenshots with the new --annotate flag, which overlays numbered labels on interactive elements and prints a legend mapping each label to its element ref. This enables multimodal AI models to reason about visual layout while using the same @eN refs for subsequent interactions. The flag can also be set via the AGENT_BROWSER_ANNOTATE environment variable.
|
||||
|
||||
## 0.11.1
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- c6fc7df: Added documentation for command chaining with && across README, CLI help output, docs, and skill files, explaining how to efficiently chain multiple agent-browser commands in a single shell invocation since the browser persists via a background daemon.
|
||||
|
||||
## 0.11.0
|
||||
|
||||
### Minor Changes
|
||||
|
||||
- 5dc40b4: Added configuration file support with automatic loading from user and project directories, new profiler commands for Chrome DevTools profiling, computed styles getter, browser extension loading, storage state management, and iOS device emulation. Expanded click command with new-tab option, improved find command with additional actions and filtering options, and enhanced CDP connection to accept WebSocket URLs. Documentation has been significantly expanded with new sections for configuration, profiling, and proxy support.
|
||||
|
||||
## 0.10.0
|
||||
|
||||
### Minor Changes
|
||||
|
||||
- 1112a16: Added session persistence with automatic save/restore of cookies and localStorage across browser restarts using --session-name flag, with optional AES-256-GCM encryption for saved state data. New state management commands allow listing, showing, renaming, clearing, and cleaning up old session files. Also added --new-tab option for click commands to open links in new tabs.
|
||||
|
||||
## 0.9.4
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 323b6cd: Fix all Clippy lint warnings in the Rust CLI: remove redundant import, use `.first()` instead of `.get(0)`, use `.copied()` instead of `.map(|s| *s)`, use `.contains()` instead of `.iter().any()`, use `then_some` instead of lazy `then`, and simplify redundant match guards.
|
||||
|
||||
## 0.9.3
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- d03e238: Added support for custom executable path in CLI browser launch options. Documentation site received UI improvements including a new chat component with sheet-based interface and updated dependencies.
|
||||
|
||||
## 0.9.2
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 76d23db: Documentation site migrated to MDX for improved content authoring, added AI-powered docs chat feature, and updated README with Homebrew installation instructions for macOS users.
|
||||
|
||||
## 0.9.1
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- ae34945: Added --allow-file-access flag to enable opening and interacting with local file:// URLs (PDFs, HTML files) by passing Chromium flags that allow JavaScript access to local files. Added -C/--cursor flag for snapshots to include cursor-interactive elements like divs with onclick handlers or cursor:pointer styles, which is useful for modern web apps using custom clickable elements.
|
||||
|
||||
## 0.9.0
|
||||
|
||||
### Minor Changes
|
||||
|
||||
- 9d021bd: Add iOS Simulator and real device support for mobile Safari testing via Appium. New CLI commands include `device list` to show available simulators, `tap` and `swipe` for touch interactions, and the `--device` flag to specify which iOS device to use. Configure with `-p ios` provider flag or `AGENT_BROWSER_PROVIDER=ios` environment variable.
|
||||
|
||||
## 0.8.10
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 17dba8f: Add --stdin flag for eval command to read JavaScript from stdin, enabling heredoc usage for multiline scripts
|
||||
- daeede4: Add --stdin flag for the eval command to read JavaScript from stdin, enabling heredoc usage for multiline scripts. Also fix binary permission issues on macOS/Linux when postinstall scripts don't run (e.g., with bun).
|
||||
|
||||
## 0.8.9
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 0dc36f2: Add --stdin flag for eval command to read JavaScript from stdin, enabling heredoc usage for multiline scripts
|
||||
|
||||
## 0.8.8
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 2771588: Added base64 encoding support for the eval command with -b/--base64 flag to avoid shell escaping issues when executing JavaScript. Updated documentation with AI agent setup instructions and reorganized the docs structure by consolidating agent mode content into the installation page.
|
||||
|
||||
## 0.8.7
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- d24f753: Fixed browser launch options not being passed correctly when using persistent profiles, ensuring args, userAgent, proxy, and ignoreHTTPSErrors settings now work properly. Added pre-flight checks for socket path length limits and directory write permissions to provide clearer error messages when daemon startup fails. Improved error handling to properly exit with failure status when browser launch fails.
|
||||
|
||||
## 0.8.6
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- d75350a: Improved daemon connection reliability by adding automatic retry logic for transient errors like connection resets, broken pipes, and temporary resource unavailability. The CLI now cleans up stale socket and PID files before starting a new daemon, and includes better detection of daemon responsiveness to handle race conditions during shutdown.
|
||||
|
||||
## 0.8.5
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- cb2f8c3: Fixed version synchronization to automatically update Cargo.lock alongside Cargo.toml during releases, and made the CLI binary executable. This ensures the Rust CLI version stays in sync with the npm package version.
|
||||
|
||||
## 0.8.4
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 759302e: Fixed "Daemon not found" error when running through AI agents (e.g., Claude Code) by resolving symlinks in the executable path. Previously, npm global bin symlinks weren't being resolved correctly, causing intermittent daemon discovery failures.
|
||||
|
||||
## 0.8.3
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 4116a8a: Replaced shell-based CLI wrappers with a cross-platform Node.js wrapper to enable npx support on Windows. Added postinstall logic to patch npm's bin entry on global installs, allowing the native binary to be invoked directly with zero overhead. Added CI tests to verify global installation works correctly across all platforms.
|
||||
|
||||
## 0.8.2
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 7e6336f: Fixed the Windows CMD wrapper to use the native binary directly instead of routing through Node.js, improving startup performance and reliability. Added retry logic to the CI install command to handle transient failures during browser installation.
|
||||
|
||||
## 0.8.1
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 8eec634: Improved release workflow to validate binary file sizes and ensure binaries are executable after npm install. Updated documentation site with a new mobile navigation system and added v0.8.0 changelog entries. Reformatted CHANGELOG.md for better readability.
|
||||
|
||||
## v0.8.0
|
||||
|
||||
### New Features
|
||||
|
||||
- **Kernel cloud browser provider** - Connect to Kernel (https://kernel.sh) for remote browser infrastructure via `-p kernel` flag or `AGENT_BROWSER_PROVIDER=kernel`. Supports stealth mode, persistent profiles, and automatic profile find-or-create.
|
||||
- **Ignore HTTPS certificate errors** - New `--ignore-https-errors` flag for working with self-signed certificates and development environments
|
||||
- **Enhanced cookie management** - Extended `cookies set` command with `--url`, `--domain`, `--path`, `--httpOnly`, `--secure`, `--sameSite`, and `--expires` flags for setting cookies before page load
|
||||
|
||||
### Bug Fixes
|
||||
|
||||
- Fixed tab list command not recognizing new pages opened via clicks or `target="_blank"` links (#275)
|
||||
- Fixed `check` command hanging indefinitely (#272)
|
||||
- Fixed `set device` not applying deviceScaleFactor - HiDPI screenshots now work correctly (#270)
|
||||
- Fixed state load and profile persistence not working in v0.7.6 (#268)
|
||||
- Screenshots now save to temp directory when no path is provided (#247)
|
||||
|
||||
### Security
|
||||
|
||||
- Daemon and stream server now reject cross-origin connections (#274)
|
||||
|
||||
## 0.7.6
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- a4d0c26: Allow null values for the screenshot selector field. Previously, passing a null selector would fail validation, but now it is properly handled as an optional value.
|
||||
|
||||
## 0.7.5
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 8c2a6ec: Fix GitHub release workflow to handle existing releases. If a release already exists, binaries are uploaded to it instead of failing.
|
||||
|
||||
## 0.7.4
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 957b5e5: Fix binary permissions on install. npm doesn't preserve execute bits, so postinstall now ensures the native binary is executable.
|
||||
|
||||
## 0.7.3
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 161d8f5: Fix native binary distribution in npm package. Native binaries for all platforms (Linux x64/arm64, macOS x64/arm64, Windows x64) are now correctly included when publishing.
|
||||
|
||||
## 0.7.2
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 6afede2: Fix native binary distribution in npm package
|
||||
|
||||
Native binaries for all platforms (Linux x64/arm64, macOS x64/arm64, Windows x64) are now included in the npm package. Previously, the release workflow published to npm before building binaries, causing "No binary found" errors on installation.
|
||||
|
||||
## 0.7.1
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- Fix native binary distribution in npm package. Native binaries for all platforms (Linux x64/arm64, macOS x64/arm64, Windows x64) are now included in the npm package. Previously, the release workflow published to npm before building binaries, causing "No binary found" errors on installation.
|
||||
|
||||
## 0.7.0
|
||||
|
||||
### Minor Changes
|
||||
|
||||
- 316e649: ## New Features
|
||||
- **Cloud browser providers** - Connect to Browserbase or Browser Use for remote browser infrastructure via `-p` flag or `AGENT_BROWSER_PROVIDER` env var
|
||||
- **Persistent browser profiles** - Store cookies, localStorage, and login sessions across browser restarts with `--profile`
|
||||
- **Remote CDP WebSocket URLs** - Connect to remote browser services via WebSocket URL (e.g., `--cdp "wss://..."`)
|
||||
- **Download commands** - New `download` command and `wait --download` for file downloads with ref support
|
||||
- **Browser launch configuration** - New `--args`, `--user-agent`, and `--proxy-bypass` flags for fine-grained browser control
|
||||
- **Enhanced skills** - Hierarchical structure with references and templates for Claude Code
|
||||
|
||||
## Bug Fixes
|
||||
- Screenshot command now supports refs and has improved error messages
|
||||
- WebSocket URLs work in `connect` command
|
||||
- Fixed socket file location (uses `~/.agent-browser` instead of TMPDIR)
|
||||
- Windows binary path fix (.exe extension)
|
||||
- State load and path-based actions now show correct output messages
|
||||
|
||||
## Documentation
|
||||
- Added Claude Code marketplace plugin installation instructions
|
||||
- Updated skill documentation with references and templates
|
||||
- Improved error documentation
|
||||
Binary file not shown.
@@ -1 +0,0 @@
|
||||
/Users/leo/github.com/agent-browser/cli/target/release/agent-browser: /Users/leo/github.com/agent-browser/cli/build.rs /Users/leo/github.com/agent-browser/cli/cdp-protocol/browser_protocol.json /Users/leo/github.com/agent-browser/cli/cdp-protocol/js_protocol.json /Users/leo/github.com/agent-browser/cli/src/color.rs /Users/leo/github.com/agent-browser/cli/src/commands.rs /Users/leo/github.com/agent-browser/cli/src/connection.rs /Users/leo/github.com/agent-browser/cli/src/flags.rs /Users/leo/github.com/agent-browser/cli/src/install.rs /Users/leo/github.com/agent-browser/cli/src/main.rs /Users/leo/github.com/agent-browser/cli/src/output.rs /Users/leo/github.com/agent-browser/cli/src/validation.rs
|
||||
+2
-13
@@ -8,7 +8,7 @@
|
||||
* binary directly (zero overhead).
|
||||
*/
|
||||
|
||||
import { spawn, execSync } from 'child_process';
|
||||
import { spawn } from 'child_process';
|
||||
import { existsSync, accessSync, chmodSync, constants } from 'fs';
|
||||
import { dirname, join } from 'path';
|
||||
import { fileURLToPath } from 'url';
|
||||
@@ -16,17 +16,6 @@ import { platform, arch } from 'os';
|
||||
|
||||
const __dirname = dirname(fileURLToPath(import.meta.url));
|
||||
|
||||
// Detect if the system uses musl libc (e.g. Alpine Linux)
|
||||
function isMusl() {
|
||||
if (platform() !== 'linux') return false;
|
||||
try {
|
||||
const result = execSync('ldd --version 2>&1 || true', { encoding: 'utf8' });
|
||||
return result.toLowerCase().includes('musl');
|
||||
} catch {
|
||||
return existsSync('/lib/ld-musl-x86_64.so.1') || existsSync('/lib/ld-musl-aarch64.so.1');
|
||||
}
|
||||
}
|
||||
|
||||
// Map Node.js platform/arch to binary naming convention
|
||||
function getBinaryName() {
|
||||
const os = platform();
|
||||
@@ -38,7 +27,7 @@ function getBinaryName() {
|
||||
osKey = 'darwin';
|
||||
break;
|
||||
case 'linux':
|
||||
osKey = isMusl() ? 'linux-musl' : 'linux';
|
||||
osKey = 'linux';
|
||||
break;
|
||||
case 'win32':
|
||||
osKey = 'win32';
|
||||
|
||||
Generated
+35
-264
@@ -45,34 +45,26 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "agent-browser-stealth"
|
||||
version = "0.24.0-fork.1"
|
||||
version = "0.16.3-fork.3"
|
||||
dependencies = [
|
||||
"aes-gcm",
|
||||
"async-trait",
|
||||
"base64",
|
||||
"chrono",
|
||||
"dirs",
|
||||
"futures-util",
|
||||
"getrandom 0.2.17",
|
||||
"hex",
|
||||
"hmac",
|
||||
"image",
|
||||
"libc",
|
||||
"regex-lite",
|
||||
"reqwest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
"similar",
|
||||
"socket2",
|
||||
"time",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
"url",
|
||||
"urlencoding",
|
||||
"uuid",
|
||||
"windows-sys 0.52.0",
|
||||
"zip",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -93,15 +85,6 @@ dependencies = [
|
||||
"equator",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "android_system_properties"
|
||||
version = "0.1.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311"
|
||||
dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "anyhow"
|
||||
version = "1.0.102"
|
||||
@@ -302,19 +285,6 @@ version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724"
|
||||
|
||||
[[package]]
|
||||
name = "chrono"
|
||||
version = "0.4.44"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0"
|
||||
dependencies = [
|
||||
"iana-time-zone",
|
||||
"js-sys",
|
||||
"num-traits",
|
||||
"wasm-bindgen",
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cipher"
|
||||
version = "0.4.4"
|
||||
@@ -331,12 +301,6 @@ version = "1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3d7b894f5411737b7867f4827955924d7c254fc9f4d91a6aad6b097804b1018b"
|
||||
|
||||
[[package]]
|
||||
name = "core-foundation-sys"
|
||||
version = "0.8.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
|
||||
|
||||
[[package]]
|
||||
name = "core2"
|
||||
version = "0.4.0"
|
||||
@@ -421,15 +385,6 @@ version = "2.10.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d7a1e2f27636f116493b8b860f5546edb47c8d8f8ea73e1d2a20be88e28d1fea"
|
||||
|
||||
[[package]]
|
||||
name = "deranged"
|
||||
version = "0.5.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c"
|
||||
dependencies = [
|
||||
"powerfmt",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "digest"
|
||||
version = "0.10.7"
|
||||
@@ -438,7 +393,6 @@ checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
|
||||
dependencies = [
|
||||
"block-buffer",
|
||||
"crypto-common",
|
||||
"subtle",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -573,7 +527,6 @@ checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c"
|
||||
dependencies = [
|
||||
"crc32fast",
|
||||
"miniz_oxide",
|
||||
"zlib-rs",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -675,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",
|
||||
]
|
||||
@@ -745,21 +698,6 @@ version = "0.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
|
||||
|
||||
[[package]]
|
||||
name = "hex"
|
||||
version = "0.4.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
|
||||
|
||||
[[package]]
|
||||
name = "hmac"
|
||||
version = "0.12.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e"
|
||||
dependencies = [
|
||||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "http"
|
||||
version = "1.4.0"
|
||||
@@ -834,7 +772,7 @@ dependencies = [
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tower-service",
|
||||
"webpki-roots 1.0.5",
|
||||
"webpki-roots 1.0.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -860,30 +798,6 @@ dependencies = [
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "iana-time-zone"
|
||||
version = "0.1.65"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470"
|
||||
dependencies = [
|
||||
"android_system_properties",
|
||||
"core-foundation-sys",
|
||||
"iana-time-zone-haiku",
|
||||
"js-sys",
|
||||
"log",
|
||||
"wasm-bindgen",
|
||||
"windows-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "iana-time-zone-haiku"
|
||||
version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f"
|
||||
dependencies = [
|
||||
"cc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "icu_collections"
|
||||
version = "2.1.1"
|
||||
@@ -1129,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"
|
||||
@@ -1145,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",
|
||||
]
|
||||
|
||||
@@ -1192,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"
|
||||
@@ -1258,12 +1171,6 @@ dependencies = [
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-conv"
|
||||
version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cf97ec579c3c42f953ef76dbf8d55ac91fb219dde70e49aa4a6b7d74e9919050"
|
||||
|
||||
[[package]]
|
||||
name = "num-derive"
|
||||
version = "0.4.2"
|
||||
@@ -1386,12 +1293,6 @@ dependencies = [
|
||||
"zerovec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "powerfmt"
|
||||
version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391"
|
||||
|
||||
[[package]]
|
||||
name = "ppv-lite86"
|
||||
version = "0.2.21"
|
||||
@@ -1413,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",
|
||||
]
|
||||
@@ -1517,9 +1418,9 @@ dependencies = [
|
||||
|
||||
[[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",
|
||||
]
|
||||
@@ -1530,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"
|
||||
@@ -1670,12 +1577,6 @@ dependencies = [
|
||||
"thiserror 1.0.69",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "regex-lite"
|
||||
version = "0.1.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973"
|
||||
|
||||
[[package]]
|
||||
name = "reqwest"
|
||||
version = "0.12.28"
|
||||
@@ -1711,7 +1612,7 @@ dependencies = [
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"web-sys",
|
||||
"webpki-roots 1.0.5",
|
||||
"webpki-roots 1.0.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1943,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",
|
||||
@@ -2026,37 +1927,6 @@ dependencies = [
|
||||
"zune-jpeg 0.4.21",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "time"
|
||||
version = "0.3.47"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "743bd48c283afc0388f9b8827b976905fb217ad9e647fae3a379a9283c4def2c"
|
||||
dependencies = [
|
||||
"deranged",
|
||||
"itoa",
|
||||
"num-conv",
|
||||
"powerfmt",
|
||||
"serde_core",
|
||||
"time-core",
|
||||
"time-macros",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "time-core"
|
||||
version = "0.1.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7694e1cfe791f8d31026952abf09c69ca6f6fa4e1a1229e18988f06a04a12dca"
|
||||
|
||||
[[package]]
|
||||
name = "time-macros"
|
||||
version = "0.2.27"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2e70e4c5a0e0a8a4823ad65dfe1a6930e4f4d756dcd9dd7939022b5e8c501215"
|
||||
dependencies = [
|
||||
"num-conv",
|
||||
"time-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tinystr"
|
||||
version = "0.8.2"
|
||||
@@ -2084,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",
|
||||
@@ -2225,12 +2095,6 @@ dependencies = [
|
||||
"utf-8",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "typed-path"
|
||||
version = "0.12.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8e28f89b80c87b8fb0cf04ab448d5dd0dd0ade2f8891bae878de66a75a28600e"
|
||||
|
||||
[[package]]
|
||||
name = "typenum"
|
||||
version = "1.19.0"
|
||||
@@ -2239,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"
|
||||
@@ -2277,12 +2141,6 @@ dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "urlencoding"
|
||||
version = "2.1.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da"
|
||||
|
||||
[[package]]
|
||||
name = "utf-8"
|
||||
version = "0.7.6"
|
||||
@@ -2301,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",
|
||||
]
|
||||
@@ -2475,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",
|
||||
]
|
||||
@@ -2493,65 +2351,12 @@ version = "0.1.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a28ac98ddc8b9274cb41bb4d9d4d5c425b6020c50c46f25559911905610b4a88"
|
||||
|
||||
[[package]]
|
||||
name = "windows-core"
|
||||
version = "0.62.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb"
|
||||
dependencies = [
|
||||
"windows-implement",
|
||||
"windows-interface",
|
||||
"windows-link",
|
||||
"windows-result",
|
||||
"windows-strings",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-implement"
|
||||
version = "0.60.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-interface"
|
||||
version = "0.59.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-link"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
|
||||
|
||||
[[package]]
|
||||
name = "windows-result"
|
||||
version = "0.4.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5"
|
||||
dependencies = [
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-strings"
|
||||
version = "0.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091"
|
||||
dependencies = [
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.48.0"
|
||||
@@ -2977,43 +2782,11 @@ dependencies = [
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zip"
|
||||
version = "8.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b680f2a0cd479b4cff6e1233c483fdead418106eae419dc60200ae9850f6d004"
|
||||
dependencies = [
|
||||
"crc32fast",
|
||||
"flate2",
|
||||
"indexmap",
|
||||
"memchr",
|
||||
"typed-path",
|
||||
"zopfli",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zlib-rs"
|
||||
version = "0.6.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3be3d40e40a133f9c916ee3f9f4fa2d9d63435b5fbe1bfc6d9dae0aa0ada1513"
|
||||
|
||||
[[package]]
|
||||
name = "zmij"
|
||||
version = "1.0.12"
|
||||
version = "1.0.21"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2fc5a66a20078bf1251bde995aa2fdcc4b800c70b5d92dd2c62abc5c60f679f8"
|
||||
|
||||
[[package]]
|
||||
name = "zopfli"
|
||||
version = "0.8.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f05cd8797d63865425ff89b5c4a48804f35ba0ce8d125800027ad6017d2b5249"
|
||||
dependencies = [
|
||||
"bumpalo",
|
||||
"crc32fast",
|
||||
"log",
|
||||
"simd-adler32",
|
||||
]
|
||||
checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa"
|
||||
|
||||
[[package]]
|
||||
name = "zune-core"
|
||||
@@ -3048,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",
|
||||
]
|
||||
|
||||
+10
-16
@@ -1,27 +1,25 @@
|
||||
[package]
|
||||
name = "agent-browser-stealth"
|
||||
version = "0.24.0-fork.1"
|
||||
version = "0.16.3-fork.3"
|
||||
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"
|
||||
repository = "https://github.com/leeguooooo/agent-browser-stealth"
|
||||
homepage = "https://github.com/leeguooooo/agent-browser-stealth"
|
||||
readme = "../README.md"
|
||||
keywords = ["browser", "automation", "ai", "cdp", "chrome"]
|
||||
categories = ["command-line-utilities", "web-programming"]
|
||||
|
||||
[[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"
|
||||
regex-lite = "0.1"
|
||||
dirs = "5.0"
|
||||
base64 = "0.22"
|
||||
getrandom = "0.2"
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros", "net", "io-util", "time", "sync", "signal", "process"] }
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros", "net", "io-util", "time", "sync", "signal"] }
|
||||
tokio-tungstenite = { version = "0.24", features = ["rustls-tls-webpki-roots"] }
|
||||
futures-util = "0.3"
|
||||
url = "2"
|
||||
@@ -31,14 +29,7 @@ reqwest = { version = "0.12", default-features = false, features = ["json", "rus
|
||||
sha2 = "0.10"
|
||||
aes-gcm = "0.10"
|
||||
async-trait = "0.1"
|
||||
socket2 = "0.6"
|
||||
similar = "2"
|
||||
zip = { version = "8.2.0", default-features = false, features = ["deflate"] }
|
||||
time = { version = "0.3", features = ["formatting"] }
|
||||
hmac = "0.12"
|
||||
hex = "0.4"
|
||||
chrono = "0.4"
|
||||
urlencoding = "2"
|
||||
|
||||
[target.'cfg(unix)'.dependencies]
|
||||
libc = "0.2"
|
||||
@@ -60,3 +51,6 @@ strip = true
|
||||
inherits = "release"
|
||||
lto = "thin"
|
||||
codegen-units = 16
|
||||
|
||||
[patch.crates-io]
|
||||
zune-jpeg = { path = "vendor/zune-jpeg" }
|
||||
|
||||
+3
-3
@@ -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).is_some_and(|n| n.is_lowercase());
|
||||
let next_lower = chars.get(i + 1).map_or(false, |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)
|
||||
.is_some_and(|t| t.contains(ref_type))
|
||||
.map_or(false, |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().is_some_and(|c| c.is_ascii_digit()) {
|
||||
if variant.chars().next().map_or(false, |c| c.is_ascii_digit()) {
|
||||
variant = format!("V{}", variant);
|
||||
}
|
||||
if seen_variants.insert(variant.clone()) {
|
||||
|
||||
+294
-755
File diff suppressed because it is too large
Load Diff
+218
-200
@@ -26,8 +26,6 @@ pub struct Response {
|
||||
pub success: bool,
|
||||
pub data: Option<Value>,
|
||||
pub error: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub warning: Option<String>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
@@ -120,10 +118,14 @@ fn get_pid_path(session: &str) -> PathBuf {
|
||||
|
||||
/// Clean up stale socket and PID files for a session
|
||||
fn cleanup_stale_files(session: &str) {
|
||||
// Never delete files for a live daemon. A missing PID file can happen in
|
||||
// race scenarios, but the socket is authoritative for liveness.
|
||||
if daemon_ready(session) {
|
||||
return;
|
||||
}
|
||||
|
||||
let pid_path = get_pid_path(session);
|
||||
let _ = fs::remove_file(&pid_path);
|
||||
let stream_path = get_socket_dir().join(format!("{}.stream", session));
|
||||
let _ = fs::remove_file(&stream_path);
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
@@ -144,7 +146,7 @@ fn get_port_path(session: &str) -> PathBuf {
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
pub fn get_port_for_session(session: &str) -> u16 {
|
||||
fn get_port_for_session(session: &str) -> u16 {
|
||||
let mut hash: i32 = 0;
|
||||
for c in session.chars() {
|
||||
hash = ((hash << 5).wrapping_sub(hash)).wrapping_add(c as i32);
|
||||
@@ -154,19 +156,7 @@ pub fn get_port_for_session(session: &str) -> u16 {
|
||||
49152 + ((hash.unsigned_abs() as u32 % 16383) as u16)
|
||||
}
|
||||
|
||||
/// Read the actual daemon port from the `.port` file written by the daemon.
|
||||
/// Falls back to the hash-derived port if the file does not exist or is
|
||||
/// unreadable (e.g. daemon has not started yet).
|
||||
#[cfg(windows)]
|
||||
pub fn resolve_port(session: &str) -> u16 {
|
||||
let port_path = get_port_path(session);
|
||||
fs::read_to_string(&port_path)
|
||||
.ok()
|
||||
.and_then(|s| s.trim().parse::<u16>().ok())
|
||||
.unwrap_or_else(|| get_port_for_session(session))
|
||||
}
|
||||
|
||||
pub fn daemon_ready(session: &str) -> bool {
|
||||
fn daemon_ready(session: &str) -> bool {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let socket_path = get_socket_path(session);
|
||||
@@ -174,7 +164,7 @@ pub fn daemon_ready(session: &str) -> bool {
|
||||
}
|
||||
#[cfg(windows)]
|
||||
{
|
||||
let port = resolve_port(session);
|
||||
let port = get_port_for_session(session);
|
||||
TcpStream::connect_timeout(
|
||||
&format!("127.0.0.1:{}", port).parse().unwrap(),
|
||||
Duration::from_millis(50),
|
||||
@@ -189,131 +179,31 @@ pub struct DaemonResult {
|
||||
pub already_running: bool,
|
||||
}
|
||||
|
||||
/// Options forwarded to the daemon process as environment variables.
|
||||
/// Note: `confirm_interactive` is intentionally absent -- it is a CLI-side
|
||||
/// UX concern (prompting the user on stdin) and not a daemon configuration.
|
||||
/// The daemon only needs `confirm_actions` to gate action categories.
|
||||
pub struct DaemonOptions<'a> {
|
||||
pub headed: bool,
|
||||
pub debug: bool,
|
||||
pub executable_path: Option<&'a str>,
|
||||
pub extensions: &'a [String],
|
||||
pub args: Option<&'a str>,
|
||||
pub user_agent: Option<&'a str>,
|
||||
pub proxy: Option<&'a str>,
|
||||
pub proxy_bypass: Option<&'a str>,
|
||||
pub proxy_username: Option<&'a str>,
|
||||
pub proxy_password: Option<&'a str>,
|
||||
pub ignore_https_errors: bool,
|
||||
pub allow_file_access: bool,
|
||||
pub profile: Option<&'a str>,
|
||||
pub state: Option<&'a str>,
|
||||
pub provider: Option<&'a str>,
|
||||
pub device: Option<&'a str>,
|
||||
pub session_name: Option<&'a str>,
|
||||
pub download_path: Option<&'a str>,
|
||||
pub allowed_domains: Option<&'a [String]>,
|
||||
pub action_policy: Option<&'a str>,
|
||||
pub confirm_actions: Option<&'a str>,
|
||||
pub engine: Option<&'a str>,
|
||||
pub auto_connect: bool,
|
||||
pub force_launch: bool,
|
||||
pub idle_timeout: Option<&'a str>,
|
||||
pub cdp: Option<&'a str>,
|
||||
pub no_auto_dialog: bool,
|
||||
}
|
||||
|
||||
fn apply_daemon_env(cmd: &mut Command, session: &str, opts: &DaemonOptions) {
|
||||
cmd.env("AGENT_BROWSER_DAEMON", "1")
|
||||
.env("AGENT_BROWSER_SESSION", session);
|
||||
|
||||
if opts.headed {
|
||||
cmd.env("AGENT_BROWSER_HEADED", "1");
|
||||
}
|
||||
if opts.debug {
|
||||
cmd.env("AGENT_BROWSER_DEBUG", "1");
|
||||
}
|
||||
if let Some(path) = opts.executable_path {
|
||||
cmd.env("AGENT_BROWSER_EXECUTABLE_PATH", path);
|
||||
}
|
||||
if !opts.extensions.is_empty() {
|
||||
cmd.env("AGENT_BROWSER_EXTENSIONS", opts.extensions.join(","));
|
||||
}
|
||||
if let Some(a) = opts.args {
|
||||
cmd.env("AGENT_BROWSER_ARGS", a);
|
||||
}
|
||||
if let Some(ua) = opts.user_agent {
|
||||
cmd.env("AGENT_BROWSER_USER_AGENT", ua);
|
||||
}
|
||||
if let Some(p) = opts.proxy {
|
||||
cmd.env("AGENT_BROWSER_PROXY", p);
|
||||
}
|
||||
if let Some(pb) = opts.proxy_bypass {
|
||||
cmd.env("AGENT_BROWSER_PROXY_BYPASS", pb);
|
||||
}
|
||||
if let Some(pu) = opts.proxy_username {
|
||||
cmd.env("AGENT_BROWSER_PROXY_USERNAME", pu);
|
||||
}
|
||||
if let Some(pp) = opts.proxy_password {
|
||||
cmd.env("AGENT_BROWSER_PROXY_PASSWORD", pp);
|
||||
}
|
||||
if opts.ignore_https_errors {
|
||||
cmd.env("AGENT_BROWSER_IGNORE_HTTPS_ERRORS", "1");
|
||||
}
|
||||
if opts.allow_file_access {
|
||||
cmd.env("AGENT_BROWSER_ALLOW_FILE_ACCESS", "1");
|
||||
}
|
||||
if let Some(prof) = opts.profile {
|
||||
cmd.env("AGENT_BROWSER_PROFILE", prof);
|
||||
}
|
||||
if let Some(st) = opts.state {
|
||||
cmd.env("AGENT_BROWSER_STATE", st);
|
||||
}
|
||||
if let Some(p) = opts.provider {
|
||||
cmd.env("AGENT_BROWSER_PROVIDER", p);
|
||||
}
|
||||
if let Some(d) = opts.device {
|
||||
cmd.env("AGENT_BROWSER_IOS_DEVICE", d);
|
||||
}
|
||||
if let Some(sn) = opts.session_name {
|
||||
cmd.env("AGENT_BROWSER_SESSION_NAME", sn);
|
||||
}
|
||||
if let Some(dp) = opts.download_path {
|
||||
cmd.env("AGENT_BROWSER_DOWNLOAD_PATH", dp);
|
||||
}
|
||||
if let Some(ad) = opts.allowed_domains {
|
||||
cmd.env("AGENT_BROWSER_ALLOWED_DOMAINS", ad.join(","));
|
||||
}
|
||||
if let Some(ap) = opts.action_policy {
|
||||
cmd.env("AGENT_BROWSER_ACTION_POLICY", ap);
|
||||
}
|
||||
if let Some(ca) = opts.confirm_actions {
|
||||
cmd.env("AGENT_BROWSER_CONFIRM_ACTIONS", ca);
|
||||
}
|
||||
if let Some(engine) = opts.engine {
|
||||
cmd.env("AGENT_BROWSER_ENGINE", engine);
|
||||
}
|
||||
if opts.auto_connect {
|
||||
cmd.env("AGENT_BROWSER_AUTO_CONNECT", "1");
|
||||
}
|
||||
if opts.force_launch {
|
||||
cmd.env("AGENT_BROWSER_FORCE_LAUNCH", "1");
|
||||
}
|
||||
if let Some(idle) = opts.idle_timeout {
|
||||
cmd.env("AGENT_BROWSER_IDLE_TIMEOUT_MS", idle);
|
||||
}
|
||||
if let Some(cdp) = opts.cdp {
|
||||
cmd.env("AGENT_BROWSER_CDP", cdp);
|
||||
}
|
||||
if opts.no_auto_dialog {
|
||||
cmd.env("AGENT_BROWSER_NO_AUTO_DIALOG", "1");
|
||||
}
|
||||
}
|
||||
|
||||
pub fn ensure_daemon(session: &str, opts: &DaemonOptions) -> Result<DaemonResult, String> {
|
||||
// Socket connectivity is the sole liveness check — no PID check — so
|
||||
// callers in a different PID namespace (e.g. unshare) can still reuse
|
||||
// an existing daemon they can reach over the socket.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn ensure_daemon(
|
||||
session: &str,
|
||||
headed: bool,
|
||||
// Keep daemon resident and disable idle auto-shutdown.
|
||||
resident: bool,
|
||||
executable_path: Option<&str>,
|
||||
extensions: &[String],
|
||||
args: Option<&str>,
|
||||
user_agent: Option<&str>,
|
||||
proxy: Option<&str>,
|
||||
proxy_bypass: Option<&str>,
|
||||
ignore_https_errors: bool,
|
||||
allow_file_access: bool,
|
||||
state: Option<&str>,
|
||||
provider: Option<&str>,
|
||||
device: Option<&str>,
|
||||
session_name: Option<&str>,
|
||||
debug: bool,
|
||||
download_path: Option<&str>,
|
||||
tab_group: Option<&str>,
|
||||
tab_group_plugin_id: Option<&str>,
|
||||
) -> Result<DaemonResult, String> {
|
||||
// Socket readiness is the source of truth for a usable daemon.
|
||||
// PID files can be missing/stale under concurrent start/stop races.
|
||||
if daemon_ready(session) {
|
||||
// Double-check it's actually responsive by waiting and checking again
|
||||
// This handles the race condition where daemon is shutting down
|
||||
@@ -368,21 +258,117 @@ pub fn ensure_daemon(session: &str, opts: &DaemonOptions) -> Result<DaemonResult
|
||||
}
|
||||
|
||||
let exe_path = env::current_exe().map_err(|e| e.to_string())?;
|
||||
// Canonicalize to resolve symlinks (e.g., npm global bin symlink -> actual binary)
|
||||
let exe_path = exe_path.canonicalize().unwrap_or(exe_path);
|
||||
let exe_dir = exe_path.parent().unwrap();
|
||||
|
||||
let mut daemon_paths = vec![
|
||||
exe_dir.join("daemon.js"),
|
||||
exe_dir.join("../dist/daemon.js"),
|
||||
PathBuf::from("dist/daemon.js"),
|
||||
];
|
||||
|
||||
// Check AGENT_BROWSER_HOME environment variable
|
||||
if let Ok(home) = env::var("AGENT_BROWSER_HOME") {
|
||||
let home_path = PathBuf::from(&home);
|
||||
daemon_paths.insert(0, home_path.join("dist/daemon.js"));
|
||||
daemon_paths.insert(1, home_path.join("daemon.js"));
|
||||
}
|
||||
|
||||
let daemon_path = daemon_paths
|
||||
.iter()
|
||||
.find(|p| p.exists())
|
||||
.ok_or("Daemon not found. Set AGENT_BROWSER_HOME environment variable or run from project directory.")?;
|
||||
|
||||
// Keep handle to detect early daemon exit and surface startup errors.
|
||||
#[allow(unused_assignments)]
|
||||
let mut daemon_child: Option<std::process::Child> = None;
|
||||
|
||||
// Spawn daemon as a fully detached background process
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::process::CommandExt;
|
||||
|
||||
let mut cmd = Command::new(&exe_path);
|
||||
cmd.env("AGENT_BROWSER_DAEMON", "1");
|
||||
apply_daemon_env(&mut cmd, session, opts);
|
||||
let mut cmd = Command::new("node");
|
||||
cmd.arg(daemon_path)
|
||||
.arg(if resident {
|
||||
"--resident"
|
||||
} else {
|
||||
"--idle-auto-shutdown"
|
||||
})
|
||||
.env("AGENT_BROWSER_DAEMON", "1")
|
||||
.env("AGENT_BROWSER_SESSION", session);
|
||||
|
||||
if headed {
|
||||
cmd.env("AGENT_BROWSER_HEADED", "1");
|
||||
}
|
||||
|
||||
if let Some(path) = executable_path {
|
||||
cmd.env("AGENT_BROWSER_EXECUTABLE_PATH", path);
|
||||
}
|
||||
|
||||
if !extensions.is_empty() {
|
||||
cmd.env("AGENT_BROWSER_EXTENSIONS", extensions.join(","));
|
||||
}
|
||||
|
||||
if let Some(a) = args {
|
||||
cmd.env("AGENT_BROWSER_ARGS", a);
|
||||
}
|
||||
|
||||
if let Some(ua) = user_agent {
|
||||
cmd.env("AGENT_BROWSER_USER_AGENT", ua);
|
||||
}
|
||||
|
||||
if let Some(p) = proxy {
|
||||
cmd.env("AGENT_BROWSER_PROXY", p);
|
||||
}
|
||||
|
||||
if let Some(pb) = proxy_bypass {
|
||||
cmd.env("AGENT_BROWSER_PROXY_BYPASS", pb);
|
||||
}
|
||||
|
||||
if ignore_https_errors {
|
||||
cmd.env("AGENT_BROWSER_IGNORE_HTTPS_ERRORS", "1");
|
||||
}
|
||||
|
||||
if allow_file_access {
|
||||
cmd.env("AGENT_BROWSER_ALLOW_FILE_ACCESS", "1");
|
||||
}
|
||||
|
||||
if let Some(st) = state {
|
||||
cmd.env("AGENT_BROWSER_STATE", st);
|
||||
}
|
||||
|
||||
if let Some(p) = provider {
|
||||
cmd.env("AGENT_BROWSER_PROVIDER", p);
|
||||
}
|
||||
|
||||
if let Some(d) = device {
|
||||
cmd.env("AGENT_BROWSER_IOS_DEVICE", d);
|
||||
}
|
||||
|
||||
if let Some(sn) = session_name {
|
||||
cmd.env("AGENT_BROWSER_SESSION_NAME", sn);
|
||||
}
|
||||
|
||||
cmd.env("AGENT_BROWSER_STEALTH", "1");
|
||||
if debug {
|
||||
cmd.env("AGENT_BROWSER_DEBUG", "1");
|
||||
}
|
||||
if let Some(dp) = download_path {
|
||||
cmd.env("AGENT_BROWSER_DOWNLOAD_PATH", dp);
|
||||
}
|
||||
if let Some(tg) = tab_group {
|
||||
cmd.env("AGENT_BROWSER_TAB_GROUP", tg);
|
||||
}
|
||||
if let Some(plugin_id) = tab_group_plugin_id {
|
||||
cmd.env("AGENT_BROWSER_TAB_GROUP_PLUGIN_ID", plugin_id);
|
||||
}
|
||||
|
||||
// Create new process group and session to fully detach
|
||||
unsafe {
|
||||
cmd.pre_exec(|| {
|
||||
// Create new session (detach from terminal)
|
||||
libc::setsid();
|
||||
Ok(())
|
||||
});
|
||||
@@ -401,10 +387,85 @@ pub fn ensure_daemon(session: &str, opts: &DaemonOptions) -> Result<DaemonResult
|
||||
{
|
||||
use std::os::windows::process::CommandExt;
|
||||
|
||||
let mut cmd = Command::new(&exe_path);
|
||||
cmd.env("AGENT_BROWSER_DAEMON", "1");
|
||||
apply_daemon_env(&mut cmd, session, opts);
|
||||
// On Windows, call node directly. Command::new handles PATH resolution (node.exe or node.cmd)
|
||||
// and automatically quotes arguments containing spaces.
|
||||
let mut cmd = Command::new("node");
|
||||
cmd.arg(daemon_path)
|
||||
.arg(if resident {
|
||||
"--resident"
|
||||
} else {
|
||||
"--idle-auto-shutdown"
|
||||
})
|
||||
.env("AGENT_BROWSER_DAEMON", "1")
|
||||
.env("AGENT_BROWSER_SESSION", session);
|
||||
|
||||
if headed {
|
||||
cmd.env("AGENT_BROWSER_HEADED", "1");
|
||||
}
|
||||
|
||||
if let Some(path) = executable_path {
|
||||
cmd.env("AGENT_BROWSER_EXECUTABLE_PATH", path);
|
||||
}
|
||||
|
||||
if !extensions.is_empty() {
|
||||
cmd.env("AGENT_BROWSER_EXTENSIONS", extensions.join(","));
|
||||
}
|
||||
|
||||
if let Some(a) = args {
|
||||
cmd.env("AGENT_BROWSER_ARGS", a);
|
||||
}
|
||||
|
||||
if let Some(ua) = user_agent {
|
||||
cmd.env("AGENT_BROWSER_USER_AGENT", ua);
|
||||
}
|
||||
|
||||
if let Some(p) = proxy {
|
||||
cmd.env("AGENT_BROWSER_PROXY", p);
|
||||
}
|
||||
|
||||
if let Some(pb) = proxy_bypass {
|
||||
cmd.env("AGENT_BROWSER_PROXY_BYPASS", pb);
|
||||
}
|
||||
|
||||
if ignore_https_errors {
|
||||
cmd.env("AGENT_BROWSER_IGNORE_HTTPS_ERRORS", "1");
|
||||
}
|
||||
|
||||
if allow_file_access {
|
||||
cmd.env("AGENT_BROWSER_ALLOW_FILE_ACCESS", "1");
|
||||
}
|
||||
|
||||
if let Some(st) = state {
|
||||
cmd.env("AGENT_BROWSER_STATE", st);
|
||||
}
|
||||
|
||||
if let Some(p) = provider {
|
||||
cmd.env("AGENT_BROWSER_PROVIDER", p);
|
||||
}
|
||||
|
||||
if let Some(d) = device {
|
||||
cmd.env("AGENT_BROWSER_IOS_DEVICE", d);
|
||||
}
|
||||
|
||||
if let Some(sn) = session_name {
|
||||
cmd.env("AGENT_BROWSER_SESSION_NAME", sn);
|
||||
}
|
||||
|
||||
cmd.env("AGENT_BROWSER_STEALTH", "1");
|
||||
if debug {
|
||||
cmd.env("AGENT_BROWSER_DEBUG", "1");
|
||||
}
|
||||
if let Some(dp) = download_path {
|
||||
cmd.env("AGENT_BROWSER_DOWNLOAD_PATH", dp);
|
||||
}
|
||||
if let Some(tg) = tab_group {
|
||||
cmd.env("AGENT_BROWSER_TAB_GROUP", tg);
|
||||
}
|
||||
if let Some(plugin_id) = tab_group_plugin_id {
|
||||
cmd.env("AGENT_BROWSER_TAB_GROUP_PLUGIN_ID", plugin_id);
|
||||
}
|
||||
|
||||
// CREATE_NEW_PROCESS_GROUP | DETACHED_PROCESS
|
||||
const CREATE_NEW_PROCESS_GROUP: u32 = 0x00000200;
|
||||
const DETACHED_PROCESS: u32 = 0x00000008;
|
||||
|
||||
@@ -425,7 +486,7 @@ pub fn ensure_daemon(session: &str, opts: &DaemonOptions) -> Result<DaemonResult
|
||||
});
|
||||
}
|
||||
|
||||
// Detect early daemon exit and surface the real error from stderr
|
||||
// Surface daemon startup stderr instead of returning an opaque timeout.
|
||||
if let Some(ref mut child) = daemon_child {
|
||||
if let Ok(Some(_)) = child.try_wait() {
|
||||
let mut stderr_output = String::new();
|
||||
@@ -434,37 +495,19 @@ pub fn ensure_daemon(session: &str, opts: &DaemonOptions) -> Result<DaemonResult
|
||||
}
|
||||
let stderr_trimmed = stderr_output.trim();
|
||||
if !stderr_trimmed.is_empty() {
|
||||
let msg = if stderr_trimmed.len() > 500 {
|
||||
let mut end = 500;
|
||||
while !stderr_trimmed.is_char_boundary(end) {
|
||||
end -= 1;
|
||||
}
|
||||
&stderr_trimmed[..end]
|
||||
} else {
|
||||
stderr_trimmed
|
||||
};
|
||||
return Err(format!("Daemon process exited during startup:\n{}", msg));
|
||||
return Err(format!("Daemon failed to start: {}", stderr_trimmed));
|
||||
}
|
||||
return Err(
|
||||
"Daemon process exited during startup with no error output. \
|
||||
Re-run with --debug for more details."
|
||||
.to_string(),
|
||||
);
|
||||
return Err("Daemon failed to start: process exited during startup".to_string());
|
||||
}
|
||||
}
|
||||
|
||||
thread::sleep(Duration::from_millis(100));
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
let endpoint_info = format!(
|
||||
"socket: {}",
|
||||
Err(format!(
|
||||
"Daemon failed to start (socket: {})",
|
||||
get_socket_dir().join(format!("{}.sock", session)).display()
|
||||
);
|
||||
#[cfg(windows)]
|
||||
let endpoint_info = format!("port: 127.0.0.1:{}", resolve_port(session));
|
||||
|
||||
Err(format!("Daemon failed to start ({})", endpoint_info))
|
||||
))
|
||||
}
|
||||
|
||||
fn connect(session: &str) -> Result<Connection, String> {
|
||||
@@ -477,7 +520,7 @@ fn connect(session: &str) -> Result<Connection, String> {
|
||||
}
|
||||
#[cfg(windows)]
|
||||
{
|
||||
let port = resolve_port(session);
|
||||
let port = get_port_for_session(session);
|
||||
TcpStream::connect(format!("127.0.0.1:{}", port))
|
||||
.map(Connection::Tcp)
|
||||
.map_err(|e| format!("Failed to connect: {}", e))
|
||||
@@ -535,8 +578,6 @@ fn is_transient_error(error: &str) -> bool {
|
||||
|| error.contains("os error 2") // No such file or directory (socket gone)
|
||||
|| error.contains("os error 61") // Connection refused (macOS)
|
||||
|| error.contains("os error 111") // Connection refused (Linux)
|
||||
|| error.contains("os error 10061") // Connection refused (Windows)
|
||||
|| error.contains("os error 10054") // Connection reset by peer (Windows)
|
||||
}
|
||||
|
||||
fn send_command_once(cmd: &Value, session: &str) -> Result<Response, String> {
|
||||
@@ -712,20 +753,6 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_transient_error_connection_refused_windows() {
|
||||
assert!(is_transient_error(
|
||||
"Failed to connect: No connection could be made because the target machine actively refused it. (os error 10061)"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_transient_error_connection_reset_windows() {
|
||||
assert!(is_transient_error(
|
||||
"Failed to send: An existing connection was forcibly closed by the remote host. (os error 10054)"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_transient_error_non_transient() {
|
||||
// These should NOT be considered transient
|
||||
@@ -734,13 +761,4 @@ mod tests {
|
||||
assert!(!is_transient_error("Permission denied"));
|
||||
assert!(!is_transient_error("Daemon not found"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(windows)]
|
||||
fn test_get_port_for_session() {
|
||||
assert_eq!(get_port_for_session("default"), 50838);
|
||||
assert_eq!(get_port_for_session("my-session"), 63105);
|
||||
assert_eq!(get_port_for_session("work"), 51184);
|
||||
assert_eq!(get_port_for_session(""), 49152);
|
||||
}
|
||||
}
|
||||
|
||||
+405
-390
File diff suppressed because it is too large
Load Diff
+154
-754
@@ -1,416 +1,167 @@
|
||||
use crate::color;
|
||||
use std::fs;
|
||||
use std::io::{self, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::{exit, Command, Stdio};
|
||||
|
||||
const LAST_KNOWN_GOOD_URL: &str =
|
||||
"https://googlechromelabs.github.io/chrome-for-testing/last-known-good-versions-with-downloads.json";
|
||||
|
||||
pub fn get_browsers_dir() -> PathBuf {
|
||||
dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".agent-browser")
|
||||
.join("browsers")
|
||||
}
|
||||
|
||||
pub fn find_installed_chrome() -> Option<PathBuf> {
|
||||
let browsers_dir = get_browsers_dir();
|
||||
let debug = std::env::var("AGENT_BROWSER_DEBUG").is_ok();
|
||||
|
||||
if debug {
|
||||
let _ = writeln!(
|
||||
io::stderr(),
|
||||
"[chrome-search] home_dir={:?} browsers_dir={}",
|
||||
dirs::home_dir(),
|
||||
browsers_dir.display()
|
||||
);
|
||||
}
|
||||
|
||||
if !browsers_dir.exists() {
|
||||
if debug {
|
||||
let _ = writeln!(io::stderr(), "[chrome-search] browsers_dir does not exist");
|
||||
}
|
||||
return None;
|
||||
}
|
||||
|
||||
let entries = match fs::read_dir(&browsers_dir) {
|
||||
Ok(entries) => entries,
|
||||
Err(e) => {
|
||||
let _ = writeln!(
|
||||
io::stderr(),
|
||||
"Warning: cannot read Chrome cache directory {}: {}",
|
||||
browsers_dir.display(),
|
||||
e
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
let mut versions: Vec<_> = entries
|
||||
.filter_map(|e| e.ok())
|
||||
.filter(|e| {
|
||||
let matches = e
|
||||
.file_name()
|
||||
.to_str()
|
||||
.is_some_and(|n| n.starts_with("chrome-"));
|
||||
if debug {
|
||||
let _ = writeln!(
|
||||
io::stderr(),
|
||||
"[chrome-search] entry {:?} matches={}",
|
||||
e.file_name(),
|
||||
matches
|
||||
);
|
||||
}
|
||||
matches
|
||||
})
|
||||
.collect();
|
||||
|
||||
versions.sort_by_key(|b| std::cmp::Reverse(b.file_name()));
|
||||
|
||||
for entry in versions {
|
||||
let dir = entry.path();
|
||||
if let Some(bin) = chrome_binary_in_dir(&dir) {
|
||||
let exists = bin.exists();
|
||||
if debug {
|
||||
let _ = writeln!(
|
||||
io::stderr(),
|
||||
"[chrome-search] candidate {} exists={}",
|
||||
bin.display(),
|
||||
exists
|
||||
);
|
||||
}
|
||||
if exists {
|
||||
return Some(bin);
|
||||
}
|
||||
} else if debug {
|
||||
let _ = writeln!(
|
||||
io::stderr(),
|
||||
"[chrome-search] no binary found in {}",
|
||||
dir.display()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
if debug {
|
||||
let _ = writeln!(io::stderr(), "[chrome-search] no installed Chrome found");
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn chrome_binary_in_dir(dir: &Path) -> Option<PathBuf> {
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
let app =
|
||||
dir.join("Google Chrome for Testing.app/Contents/MacOS/Google Chrome for Testing");
|
||||
if app.exists() {
|
||||
return Some(app);
|
||||
}
|
||||
let inner = dir.join("chrome-mac-arm64/Google Chrome for Testing.app/Contents/MacOS/Google Chrome for Testing");
|
||||
if inner.exists() {
|
||||
return Some(inner);
|
||||
}
|
||||
let inner_x64 = dir.join(
|
||||
"chrome-mac-x64/Google Chrome for Testing.app/Contents/MacOS/Google Chrome for Testing",
|
||||
);
|
||||
if inner_x64.exists() {
|
||||
return Some(inner_x64);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
{
|
||||
let bin = dir.join("chrome");
|
||||
if bin.exists() {
|
||||
return Some(bin);
|
||||
}
|
||||
let inner = dir.join("chrome-linux64/chrome");
|
||||
if inner.exists() {
|
||||
return Some(inner);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
let bin = dir.join("chrome.exe");
|
||||
if bin.exists() {
|
||||
return Some(bin);
|
||||
}
|
||||
let inner = dir.join("chrome-win64/chrome.exe");
|
||||
if inner.exists() {
|
||||
return Some(inner);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(not(any(target_os = "macos", target_os = "linux", target_os = "windows")))]
|
||||
{
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn platform_key() -> &'static str {
|
||||
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
|
||||
{
|
||||
"mac-arm64"
|
||||
}
|
||||
#[cfg(all(target_os = "macos", target_arch = "x86_64"))]
|
||||
{
|
||||
"mac-x64"
|
||||
}
|
||||
#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
|
||||
{
|
||||
"linux64"
|
||||
}
|
||||
#[cfg(all(target_os = "windows", target_arch = "x86_64"))]
|
||||
{
|
||||
"win64"
|
||||
}
|
||||
#[cfg(not(any(
|
||||
all(target_os = "macos", target_arch = "aarch64"),
|
||||
all(target_os = "macos", target_arch = "x86_64"),
|
||||
all(target_os = "linux", target_arch = "x86_64"),
|
||||
all(target_os = "windows", target_arch = "x86_64"),
|
||||
)))]
|
||||
{
|
||||
// Compiles on unsupported platforms (e.g. linux aarch64) so the binary
|
||||
// can still be used for other commands like `connect`. The install path
|
||||
// guards against this at runtime before calling platform_key().
|
||||
panic!("Unsupported platform for Chrome for Testing download")
|
||||
}
|
||||
}
|
||||
|
||||
async fn fetch_download_url() -> Result<(String, String), String> {
|
||||
let resp = reqwest::get(LAST_KNOWN_GOOD_URL)
|
||||
.await
|
||||
.map_err(|e| format!("Failed to fetch version info: {}", e))?;
|
||||
|
||||
let body: serde_json::Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to parse version info: {}", e))?;
|
||||
|
||||
let channel = body
|
||||
.get("channels")
|
||||
.and_then(|c| c.get("Stable"))
|
||||
.ok_or("No Stable channel found in version info")?;
|
||||
|
||||
let version = channel
|
||||
.get("version")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or("No version string found")?
|
||||
.to_string();
|
||||
|
||||
let platform = platform_key();
|
||||
|
||||
let url = channel
|
||||
.get("downloads")
|
||||
.and_then(|d| d.get("chrome"))
|
||||
.and_then(|c| c.as_array())
|
||||
.and_then(|arr| {
|
||||
arr.iter().find_map(|entry| {
|
||||
if entry.get("platform")?.as_str()? == platform {
|
||||
Some(entry.get("url")?.as_str()?.to_string())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
})
|
||||
.ok_or_else(|| format!("No download URL found for platform: {}", platform))?;
|
||||
|
||||
Ok((version, url))
|
||||
}
|
||||
|
||||
async fn download_bytes(url: &str) -> Result<Vec<u8>, String> {
|
||||
let resp = reqwest::get(url)
|
||||
.await
|
||||
.map_err(|e| format!("Download failed: {}", e))?;
|
||||
|
||||
let total = resp.content_length();
|
||||
let mut bytes = Vec::new();
|
||||
let mut stream = resp;
|
||||
let mut downloaded: u64 = 0;
|
||||
let mut last_pct: u64 = 0;
|
||||
|
||||
loop {
|
||||
let chunk = stream
|
||||
.chunk()
|
||||
.await
|
||||
.map_err(|e| format!("Download error: {}", e))?;
|
||||
match chunk {
|
||||
Some(data) => {
|
||||
downloaded += data.len() as u64;
|
||||
bytes.extend_from_slice(&data);
|
||||
|
||||
if let Some(total) = total {
|
||||
let pct = (downloaded * 100) / total;
|
||||
if pct >= last_pct + 5 {
|
||||
last_pct = pct;
|
||||
let mb = downloaded as f64 / 1_048_576.0;
|
||||
let total_mb = total as f64 / 1_048_576.0;
|
||||
eprint!("\r {:.0}/{:.0} MB ({pct}%)", mb, total_mb);
|
||||
let _ = io::stderr().flush();
|
||||
}
|
||||
}
|
||||
}
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
|
||||
eprintln!();
|
||||
Ok(bytes)
|
||||
}
|
||||
|
||||
fn extract_zip(bytes: Vec<u8>, dest: &Path) -> Result<(), String> {
|
||||
fs::create_dir_all(dest).map_err(|e| format!("Failed to create directory: {}", e))?;
|
||||
|
||||
let cursor = io::Cursor::new(bytes);
|
||||
let mut archive =
|
||||
zip::ZipArchive::new(cursor).map_err(|e| format!("Failed to read zip archive: {}", e))?;
|
||||
|
||||
for i in 0..archive.len() {
|
||||
let mut file = archive
|
||||
.by_index(i)
|
||||
.map_err(|e| format!("Failed to read zip entry: {}", e))?;
|
||||
|
||||
let enclosed = match file.enclosed_name() {
|
||||
Some(name) => name.to_owned(),
|
||||
None => continue,
|
||||
};
|
||||
let raw_name = enclosed.to_string_lossy().to_string();
|
||||
// Strip the top-level "chrome-<platform>/" directory from zip entries.
|
||||
// On Windows, enclosed_name() normalizes paths to backslashes, so we
|
||||
// must split on either separator.
|
||||
let rel_path = raw_name
|
||||
.strip_prefix("chrome-")
|
||||
.and_then(|s| s.find(['/', '\\']).map(|i| &s[i + 1..]))
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or(raw_name.clone());
|
||||
|
||||
if rel_path.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let out_path = dest.join(&rel_path);
|
||||
|
||||
// Defense-in-depth: ensure the resolved path is inside dest
|
||||
if !out_path.starts_with(dest) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if file.is_dir() {
|
||||
fs::create_dir_all(&out_path)
|
||||
.map_err(|e| format!("Failed to create dir {}: {}", out_path.display(), e))?;
|
||||
} else {
|
||||
if let Some(parent) = out_path.parent() {
|
||||
fs::create_dir_all(parent).map_err(|e| {
|
||||
format!("Failed to create parent dir {}: {}", parent.display(), e)
|
||||
})?;
|
||||
}
|
||||
let mut out_file = fs::File::create(&out_path)
|
||||
.map_err(|e| format!("Failed to create file {}: {}", out_path.display(), e))?;
|
||||
io::copy(&mut file, &mut out_file)
|
||||
.map_err(|e| format!("Failed to write {}: {}", out_path.display(), e))?;
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
if let Some(mode) = file.unix_mode() {
|
||||
let _ = fs::set_permissions(&out_path, fs::Permissions::from_mode(mode));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn run_install(with_deps: bool) {
|
||||
if cfg!(all(target_os = "linux", target_arch = "aarch64")) {
|
||||
eprintln!(
|
||||
"{} Chrome for Testing does not provide Linux ARM64 builds.",
|
||||
color::error_indicator()
|
||||
);
|
||||
eprintln!(" Install Chromium from your system package manager instead:");
|
||||
eprintln!(" sudo apt install chromium-browser # Debian/Ubuntu");
|
||||
eprintln!(" sudo dnf install chromium # Fedora");
|
||||
eprintln!(" Then use: agent-browser --executable-path /usr/bin/chromium");
|
||||
exit(1);
|
||||
}
|
||||
|
||||
let is_linux = cfg!(target_os = "linux");
|
||||
|
||||
if is_linux {
|
||||
if with_deps {
|
||||
install_linux_deps();
|
||||
println!("{}", color::cyan("Installing system dependencies..."));
|
||||
|
||||
let (pkg_mgr, deps) = if which_exists("apt-get") {
|
||||
let libasound = if package_exists_apt("libasound2t64") {
|
||||
"libasound2t64"
|
||||
} else {
|
||||
"libasound2"
|
||||
};
|
||||
|
||||
(
|
||||
"apt-get",
|
||||
vec![
|
||||
"libxcb-shm0",
|
||||
"libx11-xcb1",
|
||||
"libx11-6",
|
||||
"libxcb1",
|
||||
"libxext6",
|
||||
"libxrandr2",
|
||||
"libxcomposite1",
|
||||
"libxcursor1",
|
||||
"libxdamage1",
|
||||
"libxfixes3",
|
||||
"libxi6",
|
||||
"libgtk-3-0",
|
||||
"libpangocairo-1.0-0",
|
||||
"libpango-1.0-0",
|
||||
"libatk1.0-0",
|
||||
"libcairo-gobject2",
|
||||
"libcairo2",
|
||||
"libgdk-pixbuf-2.0-0",
|
||||
"libxrender1",
|
||||
libasound,
|
||||
"libfreetype6",
|
||||
"libfontconfig1",
|
||||
"libdbus-1-3",
|
||||
"libnss3",
|
||||
"libnspr4",
|
||||
"libatk-bridge2.0-0",
|
||||
"libdrm2",
|
||||
"libxkbcommon0",
|
||||
"libatspi2.0-0",
|
||||
"libcups2",
|
||||
"libxshmfence1",
|
||||
"libgbm1",
|
||||
],
|
||||
)
|
||||
} else if which_exists("dnf") {
|
||||
(
|
||||
"dnf",
|
||||
vec![
|
||||
"nss",
|
||||
"nspr",
|
||||
"atk",
|
||||
"at-spi2-atk",
|
||||
"cups-libs",
|
||||
"libdrm",
|
||||
"libXcomposite",
|
||||
"libXdamage",
|
||||
"libXrandr",
|
||||
"mesa-libgbm",
|
||||
"pango",
|
||||
"alsa-lib",
|
||||
"libxkbcommon",
|
||||
"libxcb",
|
||||
"libX11-xcb",
|
||||
"libX11",
|
||||
"libXext",
|
||||
"libXcursor",
|
||||
"libXfixes",
|
||||
"libXi",
|
||||
"gtk3",
|
||||
"cairo-gobject",
|
||||
],
|
||||
)
|
||||
} else if which_exists("yum") {
|
||||
(
|
||||
"yum",
|
||||
vec![
|
||||
"nss",
|
||||
"nspr",
|
||||
"atk",
|
||||
"at-spi2-atk",
|
||||
"cups-libs",
|
||||
"libdrm",
|
||||
"libXcomposite",
|
||||
"libXdamage",
|
||||
"libXrandr",
|
||||
"mesa-libgbm",
|
||||
"pango",
|
||||
"alsa-lib",
|
||||
"libxkbcommon",
|
||||
],
|
||||
)
|
||||
} else {
|
||||
eprintln!(
|
||||
"{} No supported package manager found (apt-get, dnf, or yum)",
|
||||
color::error_indicator()
|
||||
);
|
||||
exit(1);
|
||||
};
|
||||
|
||||
let install_cmd = match pkg_mgr {
|
||||
"apt-get" => {
|
||||
format!(
|
||||
"sudo apt-get update && sudo apt-get install -y {}",
|
||||
deps.join(" ")
|
||||
)
|
||||
}
|
||||
_ => format!("sudo {} install -y {}", pkg_mgr, deps.join(" ")),
|
||||
};
|
||||
|
||||
println!("Running: {}", install_cmd);
|
||||
let status = Command::new("sh").arg("-c").arg(&install_cmd).status();
|
||||
|
||||
match status {
|
||||
Ok(s) if s.success() => {
|
||||
println!("{} System dependencies installed", color::success_indicator())
|
||||
}
|
||||
Ok(_) => eprintln!(
|
||||
"{} Failed to install some dependencies. You may need to run manually with sudo.",
|
||||
color::warning_indicator()
|
||||
),
|
||||
Err(e) => eprintln!("{} Could not run install command: {}", color::warning_indicator(), e),
|
||||
}
|
||||
} else {
|
||||
println!(
|
||||
"{} Linux detected. If browser fails to launch, run:",
|
||||
color::warning_indicator()
|
||||
);
|
||||
println!(" agent-browser install --with-deps");
|
||||
println!(" or: npx playwright install-deps chromium");
|
||||
println!();
|
||||
}
|
||||
}
|
||||
|
||||
println!("{}", color::cyan("Installing Chrome..."));
|
||||
println!("{}", color::cyan("Installing Chromium browser..."));
|
||||
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap_or_else(|e| {
|
||||
eprintln!(
|
||||
"{} Failed to create runtime: {}",
|
||||
color::error_indicator(),
|
||||
e
|
||||
);
|
||||
exit(1);
|
||||
});
|
||||
// On Windows, we need to use cmd.exe to run npx because npx is actually npx.cmd
|
||||
// and Command::new() doesn't resolve .cmd files the way the shell does.
|
||||
// Pass the entire command as a single string to /c to handle paths with spaces.
|
||||
#[cfg(windows)]
|
||||
let status = Command::new("cmd")
|
||||
.args(["/c", "npx playwright install chromium"])
|
||||
.status();
|
||||
|
||||
let (version, url) = match rt.block_on(fetch_download_url()) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
eprintln!("{} {}", color::error_indicator(), e);
|
||||
exit(1);
|
||||
}
|
||||
};
|
||||
#[cfg(not(windows))]
|
||||
let status = Command::new("npx")
|
||||
.args(["playwright", "install", "chromium"])
|
||||
.status();
|
||||
|
||||
let dest = get_browsers_dir().join(format!("chrome-{}", version));
|
||||
|
||||
if let Some(bin) = chrome_binary_in_dir(&dest) {
|
||||
if bin.exists() {
|
||||
match status {
|
||||
Ok(s) if s.success() => {
|
||||
println!(
|
||||
"{} Chrome {} is already installed",
|
||||
color::success_indicator(),
|
||||
version
|
||||
"{} Chromium installed successfully",
|
||||
color::success_indicator()
|
||||
);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
println!(" Downloading Chrome {} for {}", version, platform_key());
|
||||
println!(" {}", url);
|
||||
|
||||
let bytes = match rt.block_on(download_bytes(&url)) {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
eprintln!("{} {}", color::error_indicator(), e);
|
||||
exit(1);
|
||||
}
|
||||
};
|
||||
|
||||
match extract_zip(bytes, &dest) {
|
||||
Ok(()) => {
|
||||
println!(
|
||||
"{} Chrome {} installed successfully",
|
||||
color::success_indicator(),
|
||||
version
|
||||
);
|
||||
println!(" Location: {}", dest.display());
|
||||
|
||||
if is_linux && !with_deps {
|
||||
println!();
|
||||
println!(
|
||||
@@ -420,252 +171,22 @@ pub fn run_install(with_deps: bool) {
|
||||
println!(" agent-browser install --with-deps");
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = fs::remove_dir_all(&dest);
|
||||
eprintln!("{} {}", color::error_indicator(), e);
|
||||
Ok(_) => {
|
||||
eprintln!("{} Failed to install browser", color::error_indicator());
|
||||
if is_linux {
|
||||
println!(
|
||||
"{} Try installing system dependencies first:",
|
||||
color::yellow("Tip:")
|
||||
);
|
||||
println!(" agent-browser install --with-deps");
|
||||
}
|
||||
exit(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn report_install_status(status: io::Result<std::process::ExitStatus>) {
|
||||
match status {
|
||||
Ok(s) if s.success() => {
|
||||
println!(
|
||||
"{} System dependencies installed",
|
||||
color::success_indicator()
|
||||
)
|
||||
Err(e) => {
|
||||
eprintln!("{} Failed to run npx: {}", color::error_indicator(), e);
|
||||
eprintln!("Make sure Node.js is installed and npx is in your PATH");
|
||||
exit(1);
|
||||
}
|
||||
Ok(_) => eprintln!(
|
||||
"{} Failed to install some dependencies. You may need to run manually with sudo.",
|
||||
color::warning_indicator()
|
||||
),
|
||||
Err(e) => eprintln!(
|
||||
"{} Could not run install command: {}",
|
||||
color::warning_indicator(),
|
||||
e
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn install_linux_deps() {
|
||||
println!("{}", color::cyan("Installing system dependencies..."));
|
||||
|
||||
let (pkg_mgr, deps) = if which_exists("apt-get") {
|
||||
// On Ubuntu 24.04+, many libraries were renamed with a t64 suffix as
|
||||
// part of the 64-bit time_t transition. Using the old names can cause
|
||||
// apt to propose removing hundreds of system packages to resolve
|
||||
// conflicts. We check for the t64 variant first to avoid this.
|
||||
let apt_deps: Vec<&str> = vec![
|
||||
("libxcb-shm0", None),
|
||||
("libx11-xcb1", None),
|
||||
("libx11-6", None),
|
||||
("libxcb1", None),
|
||||
("libxext6", None),
|
||||
("libxrandr2", None),
|
||||
("libxcomposite1", None),
|
||||
("libxcursor1", None),
|
||||
("libxdamage1", None),
|
||||
("libxfixes3", None),
|
||||
("libxi6", None),
|
||||
("libgtk-3-0", Some("libgtk-3-0t64")),
|
||||
("libpangocairo-1.0-0", Some("libpangocairo-1.0-0t64")),
|
||||
("libpango-1.0-0", Some("libpango-1.0-0t64")),
|
||||
("libatk1.0-0", Some("libatk1.0-0t64")),
|
||||
("libcairo-gobject2", Some("libcairo-gobject2t64")),
|
||||
("libcairo2", Some("libcairo2t64")),
|
||||
("libgdk-pixbuf-2.0-0", Some("libgdk-pixbuf-2.0-0t64")),
|
||||
("libxrender1", None),
|
||||
("libasound2", Some("libasound2t64")),
|
||||
("libfreetype6", None),
|
||||
("libfontconfig1", None),
|
||||
("libdbus-1-3", Some("libdbus-1-3t64")),
|
||||
("libnss3", None),
|
||||
("libnspr4", None),
|
||||
("libatk-bridge2.0-0", Some("libatk-bridge2.0-0t64")),
|
||||
("libdrm2", None),
|
||||
("libxkbcommon0", None),
|
||||
("libatspi2.0-0", Some("libatspi2.0-0t64")),
|
||||
("libcups2", Some("libcups2t64")),
|
||||
("libxshmfence1", None),
|
||||
("libgbm1", None),
|
||||
// Fonts: without actual font files, pages render with missing glyphs
|
||||
// (tofu). This is especially visible for CJK and emoji characters.
|
||||
("fonts-noto-color-emoji", None),
|
||||
("fonts-noto-cjk", None),
|
||||
("fonts-freefont-ttf", None),
|
||||
]
|
||||
.into_iter()
|
||||
.map(|(base, t64_variant)| {
|
||||
if let Some(t64) = t64_variant {
|
||||
if package_exists_apt(t64) {
|
||||
return t64;
|
||||
}
|
||||
}
|
||||
base
|
||||
})
|
||||
.collect();
|
||||
|
||||
("apt-get", apt_deps)
|
||||
} else if which_exists("dnf") {
|
||||
(
|
||||
"dnf",
|
||||
vec![
|
||||
"nss",
|
||||
"nspr",
|
||||
"atk",
|
||||
"at-spi2-atk",
|
||||
"cups-libs",
|
||||
"libdrm",
|
||||
"libXcomposite",
|
||||
"libXdamage",
|
||||
"libXrandr",
|
||||
"mesa-libgbm",
|
||||
"pango",
|
||||
"alsa-lib",
|
||||
"libxkbcommon",
|
||||
"libxcb",
|
||||
"libX11-xcb",
|
||||
"libX11",
|
||||
"libXext",
|
||||
"libXcursor",
|
||||
"libXfixes",
|
||||
"libXi",
|
||||
"gtk3",
|
||||
"cairo-gobject",
|
||||
// Fonts
|
||||
"google-noto-cjk-fonts",
|
||||
"google-noto-emoji-color-fonts",
|
||||
"liberation-fonts",
|
||||
],
|
||||
)
|
||||
} else if which_exists("yum") {
|
||||
(
|
||||
"yum",
|
||||
vec![
|
||||
"nss",
|
||||
"nspr",
|
||||
"atk",
|
||||
"at-spi2-atk",
|
||||
"cups-libs",
|
||||
"libdrm",
|
||||
"libXcomposite",
|
||||
"libXdamage",
|
||||
"libXrandr",
|
||||
"mesa-libgbm",
|
||||
"pango",
|
||||
"alsa-lib",
|
||||
"libxkbcommon",
|
||||
// Fonts
|
||||
"google-noto-cjk-fonts",
|
||||
"liberation-fonts",
|
||||
],
|
||||
)
|
||||
} else {
|
||||
eprintln!(
|
||||
"{} No supported package manager found (apt-get, dnf, or yum)",
|
||||
color::error_indicator()
|
||||
);
|
||||
exit(1);
|
||||
};
|
||||
|
||||
if pkg_mgr == "apt-get" {
|
||||
// Run apt-get update first
|
||||
println!("Running: sudo apt-get update");
|
||||
let update_status = Command::new("sudo").args(["apt-get", "update"]).status();
|
||||
|
||||
match update_status {
|
||||
Ok(s) if !s.success() => {
|
||||
eprintln!(
|
||||
"{} apt-get update failed. Continuing with existing package lists.",
|
||||
color::warning_indicator()
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!(
|
||||
"{} Could not run apt-get update: {}",
|
||||
color::warning_indicator(),
|
||||
e
|
||||
);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
// Simulate the install first to detect if apt would remove any
|
||||
// packages. This prevents the catastrophic scenario where installing
|
||||
// these libraries triggers removal of hundreds of system packages
|
||||
// due to dependency conflicts (e.g. on Ubuntu 24.04 with the
|
||||
// t64 transition).
|
||||
println!("Checking for conflicts...");
|
||||
let sim_output = Command::new("sudo")
|
||||
.args(["apt-get", "install", "--simulate"])
|
||||
.args(&deps)
|
||||
.output();
|
||||
|
||||
match sim_output {
|
||||
Ok(output) => {
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
let combined = format!("{}\n{}", stdout, stderr);
|
||||
|
||||
// Count packages that would be removed
|
||||
let removals: Vec<&str> = combined
|
||||
.lines()
|
||||
.filter(|line| line.starts_with("Remv "))
|
||||
.collect();
|
||||
|
||||
if !removals.is_empty() {
|
||||
eprintln!(
|
||||
"{} Aborting: apt would remove {} package(s) to install these dependencies.",
|
||||
color::error_indicator(),
|
||||
removals.len()
|
||||
);
|
||||
eprintln!(
|
||||
" This usually means some package names have changed on your system"
|
||||
);
|
||||
eprintln!(" (e.g. Ubuntu 24.04 renamed libraries with a t64 suffix).");
|
||||
eprintln!();
|
||||
eprintln!(" Packages that would be removed:");
|
||||
for line in removals.iter().take(20) {
|
||||
eprintln!(" {}", line);
|
||||
}
|
||||
if removals.len() > 20 {
|
||||
eprintln!(" ... and {} more", removals.len() - 20);
|
||||
}
|
||||
eprintln!();
|
||||
eprintln!(" To install dependencies manually, run:");
|
||||
eprintln!(" sudo apt-get install {}", deps.join(" "));
|
||||
eprintln!();
|
||||
eprintln!(" Review the apt output carefully before confirming.");
|
||||
exit(1);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!(
|
||||
"{} Could not simulate install ({}). Proceeding with caution.",
|
||||
color::warning_indicator(),
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Safe to proceed: no removals detected
|
||||
let install_cmd = format!("sudo apt-get install -y {}", deps.join(" "));
|
||||
println!("Running: {}", install_cmd);
|
||||
let status = Command::new("sudo")
|
||||
.args(["apt-get", "install", "-y"])
|
||||
.args(&deps)
|
||||
.status();
|
||||
|
||||
report_install_status(status);
|
||||
} else {
|
||||
// dnf / yum path — these package managers do not remove packages
|
||||
// during install, so the simulate-first guard is not needed.
|
||||
let install_cmd = format!("sudo {} install -y {}", pkg_mgr, deps.join(" "));
|
||||
println!("Running: {}", install_cmd);
|
||||
let status = Command::new("sh").arg("-c").arg(&install_cmd).status();
|
||||
|
||||
report_install_status(status);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -702,124 +223,3 @@ fn package_exists_apt(pkg: &str) -> bool {
|
||||
.map(|s| s.success())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Dashboard install
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
pub fn get_dashboard_dir() -> PathBuf {
|
||||
dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".agent-browser")
|
||||
.join("dashboard")
|
||||
}
|
||||
|
||||
const DASHBOARD_VERSION: &str = env!("CARGO_PKG_VERSION");
|
||||
|
||||
fn dashboard_download_url() -> String {
|
||||
format!(
|
||||
"https://github.com/vercel-labs/agent-browser/releases/download/v{}/dashboard.zip",
|
||||
DASHBOARD_VERSION
|
||||
)
|
||||
}
|
||||
|
||||
pub fn run_dashboard_install() {
|
||||
println!("{}", color::cyan("Installing dashboard..."));
|
||||
|
||||
let dest = get_dashboard_dir();
|
||||
|
||||
if dest.join("index.html").exists() {
|
||||
println!(
|
||||
"{} Dashboard is already installed at {}",
|
||||
color::success_indicator(),
|
||||
dest.display()
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let url = dashboard_download_url();
|
||||
println!(" Downloading dashboard v{}", DASHBOARD_VERSION);
|
||||
println!(" {}", url);
|
||||
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap_or_else(|e| {
|
||||
eprintln!(
|
||||
"{} Failed to create runtime: {}",
|
||||
color::error_indicator(),
|
||||
e
|
||||
);
|
||||
exit(1);
|
||||
});
|
||||
|
||||
let bytes = match rt.block_on(download_bytes(&url)) {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
eprintln!("{} {}", color::error_indicator(), e);
|
||||
eprintln!(" The dashboard may not be available for this version yet.");
|
||||
eprintln!(" You can build it locally: cd packages/dashboard && pnpm build");
|
||||
exit(1);
|
||||
}
|
||||
};
|
||||
|
||||
match extract_dashboard_zip(bytes, &dest) {
|
||||
Ok(()) => {
|
||||
println!(
|
||||
"{} Dashboard v{} installed successfully",
|
||||
color::success_indicator(),
|
||||
DASHBOARD_VERSION
|
||||
);
|
||||
println!(" Location: {}", dest.display());
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = fs::remove_dir_all(&dest);
|
||||
eprintln!("{} {}", color::error_indicator(), e);
|
||||
exit(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_dashboard_zip(bytes: Vec<u8>, dest: &Path) -> Result<(), String> {
|
||||
fs::create_dir_all(dest).map_err(|e| format!("Failed to create directory: {}", e))?;
|
||||
|
||||
let cursor = io::Cursor::new(bytes);
|
||||
let mut archive =
|
||||
zip::ZipArchive::new(cursor).map_err(|e| format!("Failed to read zip archive: {}", e))?;
|
||||
|
||||
for i in 0..archive.len() {
|
||||
let mut file = archive
|
||||
.by_index(i)
|
||||
.map_err(|e| format!("Failed to read zip entry: {}", e))?;
|
||||
|
||||
let enclosed = match file.enclosed_name() {
|
||||
Some(name) => name.to_owned(),
|
||||
None => continue,
|
||||
};
|
||||
let rel_path = enclosed.to_string_lossy().to_string();
|
||||
|
||||
if rel_path.is_empty() || file.is_dir() {
|
||||
if file.is_dir() {
|
||||
let out_dir = dest.join(&rel_path);
|
||||
let _ = fs::create_dir_all(&out_dir);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
let out_path = dest.join(&rel_path);
|
||||
if !out_path.starts_with(dest) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Some(parent) = out_path.parent() {
|
||||
fs::create_dir_all(parent)
|
||||
.map_err(|e| format!("Failed to create parent dir {}: {}", parent.display(), e))?;
|
||||
}
|
||||
let mut out_file = fs::File::create(&out_path)
|
||||
.map_err(|e| format!("Failed to create file {}: {}", out_path.display(), e))?;
|
||||
io::copy(&mut file, &mut out_file)
|
||||
.map_err(|e| format!("Failed to write {}: {}", out_path.display(), e))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
+343
-894
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1 @@
|
||||
include!("main.rs");
|
||||
+817
-4005
File diff suppressed because it is too large
Load Diff
+70
-311
@@ -1,13 +1,11 @@
|
||||
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::io::Write;
|
||||
use std::path::PathBuf;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AuthProfile {
|
||||
pub name: String,
|
||||
pub url: String,
|
||||
@@ -19,10 +17,6 @@ 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
|
||||
@@ -54,203 +48,79 @@ fn get_profile_path(name: &str) -> PathBuf {
|
||||
get_auth_dir().join(format!("{}.json", name))
|
||||
}
|
||||
|
||||
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 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
|
||||
)
|
||||
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 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()
|
||||
))
|
||||
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()
|
||||
}
|
||||
|
||||
/// 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));
|
||||
}
|
||||
|
||||
let _ = writeln!(
|
||||
std::io::stderr(),
|
||||
"[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()?;
|
||||
fn encrypt_profile(profile: &AuthProfile) -> Result<Vec<u8>, String> {
|
||||
let key = derive_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 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())
|
||||
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())
|
||||
.map_err(|e| format!("Encryption failed: {}", e))?;
|
||||
|
||||
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,
|
||||
let mut result = Vec::with_capacity(12 + ciphertext.len());
|
||||
result.extend_from_slice(&nonce);
|
||||
result.extend_from_slice(&ciphertext);
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
fn decrypt_profile(data: &[u8]) -> Result<AuthProfile, 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));
|
||||
if data.len() < 13 {
|
||||
return Err("Encrypted data too short".to_string());
|
||||
}
|
||||
let (nonce_bytes, ciphertext) = data.split_at(12);
|
||||
|
||||
// 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())
|
||||
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))
|
||||
}
|
||||
|
||||
fn save_profile(profile: &AuthProfile) -> Result<(), String> {
|
||||
let dir = get_auth_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 _ = fs::create_dir_all(&dir);
|
||||
|
||||
let encrypted_json = encrypt_profile(profile)?;
|
||||
let encrypted = encrypt_profile(profile)?;
|
||||
let path = get_profile_path(&profile.name);
|
||||
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(())
|
||||
fs::write(&path, &encrypted).map_err(|e| format!("Failed to write profile: {}", e))
|
||||
}
|
||||
|
||||
fn load_profile(name: &str) -> Result<AuthProfile, String> {
|
||||
@@ -277,8 +147,6 @@ 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 }))
|
||||
@@ -302,8 +170,6 @@ 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 }))
|
||||
@@ -386,27 +252,10 @@ 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());
|
||||
@@ -428,8 +277,6 @@ 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();
|
||||
@@ -443,114 +290,26 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_encrypt_decrypt_roundtrip() {
|
||||
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!");
|
||||
});
|
||||
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!");
|
||||
}
|
||||
|
||||
#[test]
|
||||
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());
|
||||
});
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
+115
-844
File diff suppressed because it is too large
Load Diff
+99
-467
@@ -1,9 +1,9 @@
|
||||
use std::io::{BufRead, BufReader, Write};
|
||||
use std::io::{BufRead, BufReader};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::{Child, Command, Stdio};
|
||||
use std::time::Duration;
|
||||
|
||||
use super::discovery::discover_cdp_url;
|
||||
use super::types::BrowserVersionInfo;
|
||||
|
||||
pub struct ChromeProcess {
|
||||
child: Child,
|
||||
@@ -16,35 +16,6 @@ impl ChromeProcess {
|
||||
let _ = self.child.kill();
|
||||
let _ = self.child.wait();
|
||||
}
|
||||
|
||||
/// Returns the OS process ID of the Chrome child process.
|
||||
pub fn id(&self) -> u32 {
|
||||
self.child.id()
|
||||
}
|
||||
|
||||
/// Non-blocking check whether Chrome has exited.
|
||||
/// Returns `true` if the process has exited (and reaps it), `false` if still running.
|
||||
pub fn has_exited(&mut self) -> bool {
|
||||
matches!(self.child.try_wait(), Ok(Some(_)) | Err(_))
|
||||
}
|
||||
|
||||
/// 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 {
|
||||
@@ -58,10 +29,7 @@ impl Drop for ChromeProcess {
|
||||
std::thread::sleep(Duration::from_millis(100));
|
||||
}
|
||||
Err(e) => {
|
||||
// Use write! instead of eprintln! to avoid panicking
|
||||
// if the daemon's stderr pipe is broken (parent dropped it).
|
||||
let _ = writeln!(
|
||||
std::io::stderr(),
|
||||
eprintln!(
|
||||
"Warning: failed to clean up temp profile {}: {}",
|
||||
dir.display(),
|
||||
e
|
||||
@@ -78,8 +46,6 @@ pub struct LaunchOptions {
|
||||
pub executable_path: Option<String>,
|
||||
pub proxy: Option<String>,
|
||||
pub proxy_bypass: Option<String>,
|
||||
pub proxy_username: Option<String>,
|
||||
pub proxy_password: Option<String>,
|
||||
pub profile: Option<String>,
|
||||
pub args: Vec<String>,
|
||||
pub allow_file_access: bool,
|
||||
@@ -98,8 +64,6 @@ impl Default for LaunchOptions {
|
||||
executable_path: None,
|
||||
proxy: None,
|
||||
proxy_bypass: None,
|
||||
proxy_username: None,
|
||||
proxy_password: None,
|
||||
profile: None,
|
||||
args: Vec::new(),
|
||||
allow_file_access: false,
|
||||
@@ -115,7 +79,6 @@ impl Default for LaunchOptions {
|
||||
|
||||
struct ChromeArgs {
|
||||
args: Vec<String>,
|
||||
user_data_dir: PathBuf,
|
||||
temp_user_data_dir: Option<PathBuf>,
|
||||
}
|
||||
|
||||
@@ -124,10 +87,6 @@ fn build_chrome_args(options: &LaunchOptions) -> Result<ChromeArgs, String> {
|
||||
"--remote-debugging-port=0".to_string(),
|
||||
"--no-first-run".to_string(),
|
||||
"--no-default-browser-check".to_string(),
|
||||
// Stealth: reduce automation fingerprint surface
|
||||
"--disable-blink-features=AutomationControlled".to_string(),
|
||||
"--use-gl=angle".to_string(),
|
||||
"--use-angle=default".to_string(),
|
||||
"--disable-background-networking".to_string(),
|
||||
"--disable-backgrounding-occluded-windows".to_string(),
|
||||
"--disable-component-update".to_string(),
|
||||
@@ -136,27 +95,14 @@ fn build_chrome_args(options: &LaunchOptions) -> Result<ChromeArgs, 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(),
|
||||
];
|
||||
|
||||
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 {
|
||||
if options.headless {
|
||||
args.push("--headless=new".to_string());
|
||||
// Enable SwiftShader software rendering in headless mode. This
|
||||
// prevents silent crashes in environments where GPU drivers are
|
||||
// missing or restricted (VMs, containers, some cloud machines)
|
||||
// while preserving WebGL support. Playwright uses the same flag.
|
||||
args.push("--enable-unsafe-swiftshader".to_string());
|
||||
}
|
||||
|
||||
if let Some(ref proxy) = options.proxy {
|
||||
@@ -167,18 +113,17 @@ fn build_chrome_args(options: &LaunchOptions) -> Result<ChromeArgs, String> {
|
||||
args.push(format!("--proxy-bypass-list={}", bypass));
|
||||
}
|
||||
|
||||
let (user_data_dir, temp_user_data_dir) = if let Some(ref profile) = options.profile {
|
||||
let temp_user_data_dir = if let Some(ref profile) = options.profile {
|
||||
let expanded = expand_tilde(profile);
|
||||
let dir = PathBuf::from(&expanded);
|
||||
args.push(format!("--user-data-dir={}", expanded));
|
||||
(dir, None)
|
||||
None
|
||||
} else {
|
||||
let dir =
|
||||
std::env::temp_dir().join(format!("agent-browser-chrome-{}", uuid::Uuid::new_v4()));
|
||||
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()));
|
||||
(dir.clone(), Some(dir))
|
||||
Some(dir)
|
||||
};
|
||||
|
||||
if options.allow_file_access {
|
||||
@@ -199,7 +144,7 @@ fn build_chrome_args(options: &LaunchOptions) -> Result<ChromeArgs, String> {
|
||||
.iter()
|
||||
.any(|a| a.starts_with("--start-maximized") || a.starts_with("--window-size="));
|
||||
|
||||
if !has_window_size && options.headless && !has_extensions {
|
||||
if !has_window_size && options.headless {
|
||||
args.push("--window-size=1280,720".to_string());
|
||||
}
|
||||
|
||||
@@ -209,13 +154,8 @@ fn build_chrome_args(options: &LaunchOptions) -> Result<ChromeArgs, String> {
|
||||
args.push("--no-sandbox".to_string());
|
||||
}
|
||||
|
||||
if should_disable_dev_shm(&args) {
|
||||
args.push("--disable-dev-shm-usage".to_string());
|
||||
}
|
||||
|
||||
Ok(ChromeArgs {
|
||||
args,
|
||||
user_data_dir,
|
||||
temp_user_data_dir,
|
||||
})
|
||||
}
|
||||
@@ -223,64 +163,23 @@ fn build_chrome_args(options: &LaunchOptions) -> Result<ChromeArgs, String> {
|
||||
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_else(|| {
|
||||
let cache_dir = crate::install::get_browsers_dir();
|
||||
format!(
|
||||
"Chrome not found. Checked:\n \
|
||||
- agent-browser cache: {}\n \
|
||||
- System Chrome installations\n \
|
||||
- Puppeteer browser cache\n \
|
||||
- Playwright browser cache\n\
|
||||
Run `agent-browser install` to download Chrome, or use --executable-path.",
|
||||
cache_dir.display()
|
||||
)
|
||||
})?,
|
||||
None => {
|
||||
find_chrome().ok_or("Chrome not found. Install Chrome or use --executable-path.")?
|
||||
}
|
||||
};
|
||||
|
||||
let max_attempts = 3;
|
||||
let mut last_err = String::new();
|
||||
|
||||
for attempt in 1..=max_attempts {
|
||||
match try_launch_chrome(&chrome_path, options) {
|
||||
Ok(process) => return Ok(process),
|
||||
Err(e) => {
|
||||
last_err = e;
|
||||
if attempt < max_attempts {
|
||||
// Use write! instead of eprintln! to avoid panicking
|
||||
// if the daemon's stderr pipe is broken (parent dropped it).
|
||||
let _ = writeln!(
|
||||
std::io::stderr(),
|
||||
"[chrome] Launch attempt {}/{} failed, retrying in 500ms...",
|
||||
attempt,
|
||||
max_attempts
|
||||
);
|
||||
std::thread::sleep(Duration::from_millis(500));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(last_err)
|
||||
}
|
||||
|
||||
fn try_launch_chrome(chrome_path: &Path, options: &LaunchOptions) -> Result<ChromeProcess, String> {
|
||||
let ChromeArgs {
|
||||
args,
|
||||
user_data_dir,
|
||||
temp_user_data_dir,
|
||||
} = build_chrome_args(options)?;
|
||||
|
||||
// Mitigate stale DevToolsActivePort risk (e.g., previous crash left it behind).
|
||||
// Puppeteer does similar cleanup before spawning.
|
||||
let _ = std::fs::remove_file(user_data_dir.join("DevToolsActivePort"));
|
||||
|
||||
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)
|
||||
let mut child = Command::new(&chrome_path)
|
||||
.args(&args)
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::null())
|
||||
@@ -291,33 +190,22 @@ fn try_launch_chrome(chrome_path: &Path, options: &LaunchOptions) -> Result<Chro
|
||||
format!("Failed to launch Chrome at {:?}: {}", chrome_path, e)
|
||||
})?;
|
||||
|
||||
// Shared overall deadline so we don't double-wait (poll + stderr fallback).
|
||||
let deadline = std::time::Instant::now() + Duration::from_secs(30);
|
||||
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);
|
||||
|
||||
// Primary path: use DevToolsActivePort written into user-data-dir.
|
||||
// This is more reliable on Windows than scraping stderr for "DevTools listening on ...",
|
||||
// which can be missing/empty depending on how Chrome is launched.
|
||||
let ws_url = match wait_for_devtools_active_port(&mut child, &user_data_dir, deadline) {
|
||||
let ws_url = match wait_for_ws_url(reader) {
|
||||
Ok(url) => url,
|
||||
Err(primary_err) => {
|
||||
// Fallback: scrape stderr (legacy behavior) for better diagnostics.
|
||||
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);
|
||||
match wait_for_ws_url_until(reader, deadline) {
|
||||
Ok(url) => url,
|
||||
Err(fallback_err) => {
|
||||
let _ = child.kill();
|
||||
cleanup_temp_dir(&temp_user_data_dir);
|
||||
return Err(format!(
|
||||
"{}\n(also tried parsing stderr) {}",
|
||||
primary_err, fallback_err
|
||||
));
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = child.kill();
|
||||
cleanup_temp_dir(&temp_user_data_dir);
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -328,42 +216,8 @@ fn try_launch_chrome(chrome_path: &Path, options: &LaunchOptions) -> Result<Chro
|
||||
})
|
||||
}
|
||||
|
||||
fn wait_for_devtools_active_port(
|
||||
child: &mut Child,
|
||||
user_data_dir: &Path,
|
||||
deadline: std::time::Instant,
|
||||
) -> Result<String, String> {
|
||||
let poll_interval = Duration::from_millis(50);
|
||||
|
||||
while std::time::Instant::now() <= deadline {
|
||||
if let Ok(Some(status)) = child.try_wait() {
|
||||
// Chrome exited before writing DevToolsActivePort -- report the
|
||||
// exit code so the caller can surface it alongside stderr output.
|
||||
let code = status
|
||||
.code()
|
||||
.map(|c| format!("{}", c))
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
return Err(format!(
|
||||
"Chrome exited early (exit code: {}) without writing DevToolsActivePort",
|
||||
code
|
||||
));
|
||||
}
|
||||
|
||||
if let Some((port, ws_path)) = read_devtools_active_port(user_data_dir) {
|
||||
let ws_url = format!("ws://127.0.0.1:{}{}", port, ws_path);
|
||||
return Ok(ws_url);
|
||||
}
|
||||
|
||||
std::thread::sleep(poll_interval);
|
||||
}
|
||||
|
||||
Err("Timeout waiting for DevToolsActivePort".to_string())
|
||||
}
|
||||
|
||||
fn wait_for_ws_url_until(
|
||||
reader: BufReader<std::process::ChildStderr>,
|
||||
deadline: std::time::Instant,
|
||||
) -> Result<String, String> {
|
||||
fn wait_for_ws_url(reader: BufReader<std::process::ChildStderr>) -> Result<String, String> {
|
||||
let deadline = std::time::Instant::now() + Duration::from_secs(30);
|
||||
let prefix = "DevTools listening on ";
|
||||
let mut stderr_lines: Vec<String> = Vec::new();
|
||||
|
||||
@@ -405,10 +259,7 @@ fn chrome_launch_error(message: &str, stderr_lines: &[String]) -> String {
|
||||
|
||||
if relevant.is_empty() {
|
||||
if stderr_lines.is_empty() {
|
||||
return format!(
|
||||
"{} (no stderr output from Chrome)\nHint: try passing --args \"--no-sandbox\" if Chrome crashes silently in your environment",
|
||||
message
|
||||
);
|
||||
return format!("{} (no stderr output from Chrome)", message);
|
||||
}
|
||||
let last_lines: Vec<&String> = stderr_lines.iter().rev().take(5).collect();
|
||||
return format!(
|
||||
@@ -446,31 +297,12 @@ fn chrome_launch_error(message: &str, stderr_lines: &[String]) -> String {
|
||||
}
|
||||
|
||||
pub fn find_chrome() -> Option<PathBuf> {
|
||||
// 1. Check Chrome downloaded by `agent-browser install`
|
||||
if let Some(p) = crate::install::find_installed_chrome() {
|
||||
return Some(p);
|
||||
}
|
||||
|
||||
// If the cache directory exists but no Chrome was found, warn -- this
|
||||
// likely means the cache is corrupted or the directory layout is unexpected.
|
||||
let cache_dir = crate::install::get_browsers_dir();
|
||||
if cache_dir.exists() {
|
||||
let _ = writeln!(
|
||||
std::io::stderr(),
|
||||
"Warning: Chrome cache directory exists ({}) but no Chrome binary found inside. \
|
||||
Falling back to system Chrome. Run `agent-browser install` to re-download.",
|
||||
cache_dir.display()
|
||||
);
|
||||
}
|
||||
|
||||
// 2. Check system-installed Chrome
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
let candidates = [
|
||||
"/Applications/Google Chrome.app/Contents/MacOS/Google Chrome",
|
||||
"/Applications/Google Chrome Canary.app/Contents/MacOS/Google Chrome Canary",
|
||||
"/Applications/Chromium.app/Contents/MacOS/Chromium",
|
||||
"/Applications/Brave Browser.app/Contents/MacOS/Brave Browser",
|
||||
];
|
||||
for c in &candidates {
|
||||
let p = PathBuf::from(c);
|
||||
@@ -478,6 +310,10 @@ pub fn find_chrome() -> Option<PathBuf> {
|
||||
return Some(p);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(p) = find_playwright_chromium() {
|
||||
return Some(p);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
@@ -487,8 +323,6 @@ pub fn find_chrome() -> Option<PathBuf> {
|
||||
"google-chrome-stable",
|
||||
"chromium-browser",
|
||||
"chromium",
|
||||
"brave-browser",
|
||||
"brave-browser-stable",
|
||||
];
|
||||
for name in &candidates {
|
||||
if let Ok(output) = Command::new("which").arg(name).output() {
|
||||
@@ -500,6 +334,10 @@ pub fn find_chrome() -> Option<PathBuf> {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(p) = find_playwright_chromium() {
|
||||
return Some(p);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(target_os = "windows")]
|
||||
@@ -509,14 +347,9 @@ pub fn find_chrome() -> Option<PathBuf> {
|
||||
r"C:\Program Files (x86)\Google\Chrome\Application\chrome.exe",
|
||||
];
|
||||
if let Ok(local) = std::env::var("LOCALAPPDATA") {
|
||||
let chrome = PathBuf::from(&local).join(r"Google\Chrome\Application\chrome.exe");
|
||||
if chrome.exists() {
|
||||
return Some(chrome);
|
||||
}
|
||||
let brave =
|
||||
PathBuf::from(&local).join(r"BraveSoftware\Brave-Browser\Application\brave.exe");
|
||||
if brave.exists() {
|
||||
return Some(brave);
|
||||
let p = PathBuf::from(&local).join(r"Google\Chrome\Application\chrome.exe");
|
||||
if p.exists() {
|
||||
return Some(p);
|
||||
}
|
||||
}
|
||||
for c in &candidates {
|
||||
@@ -527,17 +360,31 @@ pub fn find_chrome() -> Option<PathBuf> {
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Fallback: check Puppeteer / Playwright browser caches
|
||||
if let Some(p) = find_puppeteer_chrome() {
|
||||
return Some(p);
|
||||
}
|
||||
if let Some(p) = find_playwright_chromium() {
|
||||
return Some(p);
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
pub async fn discover_cdp_url(port: u16) -> Result<String, String> {
|
||||
let url = format!("http://127.0.0.1:{}/json/version", port);
|
||||
|
||||
let body = tokio::time::timeout(Duration::from_secs(2), async {
|
||||
reqwest_get_string(&url).await
|
||||
})
|
||||
.await
|
||||
.map_err(|_| format!("Timeout connecting to CDP on port {}", port))?
|
||||
.map_err(|e| format!("Failed to connect to CDP on port {}: {}", port, e))?;
|
||||
|
||||
let info: BrowserVersionInfo = serde_json::from_str(&body)
|
||||
.map_err(|e| format!("Invalid /json/version response: {}", e))?;
|
||||
|
||||
info.web_socket_debugger_url
|
||||
.ok_or_else(|| format!("No webSocketDebuggerUrl in /json/version on port {}", port))
|
||||
}
|
||||
|
||||
async fn reqwest_get_string(url: &str) -> Result<String, String> {
|
||||
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)> {
|
||||
let path = user_data_dir.join("DevToolsActivePort");
|
||||
let content = std::fs::read_to_string(&path).ok()?;
|
||||
@@ -557,26 +404,18 @@ pub async fn auto_connect_cdp() -> Result<String, String> {
|
||||
for dir in &user_data_dirs {
|
||||
if let Some((port, ws_path)) = read_devtools_active_port(dir) {
|
||||
// Try HTTP endpoint first (pre-M144)
|
||||
if let Ok(ws_url) = discover_cdp_url("127.0.0.1", port, None).await {
|
||||
if let Ok(ws_url) = discover_cdp_url(port).await {
|
||||
return Ok(ws_url);
|
||||
}
|
||||
// M144+: direct WebSocket — verify the port is actually listening
|
||||
// before returning, otherwise a stale DevToolsActivePort file
|
||||
// (left behind after Chrome exits/crashes) produces a confusing
|
||||
// "connection refused" error instead of falling through.
|
||||
if is_port_reachable(port) {
|
||||
let ws_url = format!("ws://127.0.0.1:{}{}", port, ws_path);
|
||||
return Ok(ws_url);
|
||||
}
|
||||
// Port is dead — remove the stale file so future runs skip it.
|
||||
let stale = dir.join("DevToolsActivePort");
|
||||
let _ = std::fs::remove_file(&stale);
|
||||
// M144+: direct WebSocket
|
||||
let ws_url = format!("ws://127.0.0.1:{}{}", port, ws_path);
|
||||
return Ok(ws_url);
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: probe common ports
|
||||
for port in [9222u16, 9229] {
|
||||
if let Ok(ws_url) = discover_cdp_url("127.0.0.1", port, None).await {
|
||||
if let Ok(ws_url) = discover_cdp_url(port).await {
|
||||
return Ok(ws_url);
|
||||
}
|
||||
}
|
||||
@@ -584,12 +423,6 @@ pub async fn auto_connect_cdp() -> Result<String, String> {
|
||||
Err("No running Chrome instance found. Launch Chrome with --remote-debugging-port or use --cdp.".to_string())
|
||||
}
|
||||
|
||||
fn is_port_reachable(port: u16) -> bool {
|
||||
use std::net::TcpStream;
|
||||
let addr = format!("127.0.0.1:{}", port);
|
||||
TcpStream::connect_timeout(&addr.parse().unwrap(), Duration::from_millis(500)).is_ok()
|
||||
}
|
||||
|
||||
fn get_chrome_user_data_dirs() -> Vec<PathBuf> {
|
||||
let mut dirs = Vec::new();
|
||||
|
||||
@@ -597,12 +430,7 @@ fn get_chrome_user_data_dirs() -> Vec<PathBuf> {
|
||||
{
|
||||
if let Some(home) = dirs::home_dir() {
|
||||
let base = home.join("Library/Application Support");
|
||||
for name in [
|
||||
"Google/Chrome",
|
||||
"Google/Chrome Canary",
|
||||
"Chromium",
|
||||
"BraveSoftware/Brave-Browser",
|
||||
] {
|
||||
for name in ["Google/Chrome", "Google/Chrome Canary", "Chromium"] {
|
||||
dirs.push(base.join(name));
|
||||
}
|
||||
}
|
||||
@@ -612,12 +440,7 @@ fn get_chrome_user_data_dirs() -> Vec<PathBuf> {
|
||||
{
|
||||
if let Some(home) = dirs::home_dir() {
|
||||
let config = home.join(".config");
|
||||
for name in [
|
||||
"google-chrome",
|
||||
"google-chrome-unstable",
|
||||
"chromium",
|
||||
"BraveSoftware/Brave-Browser",
|
||||
] {
|
||||
for name in ["google-chrome", "google-chrome-unstable", "chromium"] {
|
||||
dirs.push(config.join(name));
|
||||
}
|
||||
}
|
||||
@@ -631,7 +454,6 @@ fn get_chrome_user_data_dirs() -> Vec<PathBuf> {
|
||||
r"Google\Chrome\User Data",
|
||||
r"Google\Chrome SxS\User Data",
|
||||
r"Chromium\User Data",
|
||||
r"BraveSoftware\Brave-Browser\User Data",
|
||||
] {
|
||||
dirs.push(base.join(name));
|
||||
}
|
||||
@@ -642,18 +464,12 @@ fn get_chrome_user_data_dirs() -> Vec<PathBuf> {
|
||||
}
|
||||
|
||||
/// Returns true if Chrome's sandbox should be disabled because the environment
|
||||
/// doesn't support it (containers, VMs, CI runners, running as root).
|
||||
/// doesn't support it (containers, VMs, running as root).
|
||||
fn should_disable_sandbox(existing_args: &[String]) -> bool {
|
||||
if existing_args.iter().any(|a| a == "--no-sandbox") {
|
||||
return false; // already set by user
|
||||
}
|
||||
|
||||
// CI environments (GitHub Actions, GitLab CI, etc.) often lack user namespace
|
||||
// support due to AppArmor or kernel restrictions.
|
||||
if std::env::var("CI").is_ok() {
|
||||
return true;
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
// Root user -- standard container default, Chrome sandbox requires non-root
|
||||
@@ -673,7 +489,10 @@ 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;
|
||||
}
|
||||
}
|
||||
@@ -682,108 +501,8 @@ fn should_disable_sandbox(existing_args: &[String]) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// Returns true if Chrome should use disk instead of /dev/shm for shared memory.
|
||||
/// On CI runners and containers, /dev/shm is often too small (64MB default),
|
||||
/// which causes Chrome to crash mid-session.
|
||||
fn should_disable_dev_shm(existing_args: &[String]) -> bool {
|
||||
if existing_args.iter().any(|a| a == "--disable-dev-shm-usage") {
|
||||
return false;
|
||||
}
|
||||
|
||||
if std::env::var("CI").is_ok() {
|
||||
return true;
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
if unsafe { libc::geteuid() } == 0 {
|
||||
return true;
|
||||
}
|
||||
if Path::new("/.dockerenv").exists() || Path::new("/run/.containerenv").exists() {
|
||||
return true;
|
||||
}
|
||||
if let Ok(cgroup) = std::fs::read_to_string("/proc/1/cgroup") {
|
||||
if cgroup.contains("docker") || cgroup.contains("kubepods") || cgroup.contains("lxc") {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
/// Search Puppeteer's browser cache for a Chrome binary.
|
||||
/// Puppeteer v19+ stores Chrome in ~/.cache/puppeteer/chrome/<platform>-<version>/
|
||||
fn find_puppeteer_chrome() -> Option<PathBuf> {
|
||||
let mut search_dirs = Vec::new();
|
||||
|
||||
if let Ok(custom) = std::env::var("PUPPETEER_CACHE_DIR") {
|
||||
search_dirs.push(PathBuf::from(custom).join("chrome"));
|
||||
}
|
||||
|
||||
if let Some(home) = dirs::home_dir() {
|
||||
search_dirs.push(home.join(".cache/puppeteer/chrome"));
|
||||
}
|
||||
|
||||
for dir in &search_dirs {
|
||||
if !dir.is_dir() {
|
||||
continue;
|
||||
}
|
||||
if let Ok(entries) = std::fs::read_dir(dir) {
|
||||
let mut matches: Vec<PathBuf> = entries
|
||||
.filter_map(|e| e.ok())
|
||||
.filter(|e| e.path().is_dir())
|
||||
.filter_map(|e| {
|
||||
let candidate = build_puppeteer_binary_path(&e.path());
|
||||
if candidate.exists() {
|
||||
Some(candidate)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
matches.sort();
|
||||
matches.reverse();
|
||||
if let Some(p) = matches.into_iter().next() {
|
||||
return Some(p);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
fn build_puppeteer_binary_path(version_dir: &Path) -> PathBuf {
|
||||
version_dir.join("chrome-linux64/chrome")
|
||||
}
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
fn build_puppeteer_binary_path(version_dir: &Path) -> PathBuf {
|
||||
// Puppeteer uses chrome-mac-arm64 or chrome-mac-x64 depending on arch
|
||||
let arm = version_dir.join(
|
||||
"chrome-mac-arm64/Google Chrome for Testing.app/Contents/MacOS/Google Chrome for Testing",
|
||||
);
|
||||
if arm.exists() {
|
||||
return arm;
|
||||
}
|
||||
version_dir.join(
|
||||
"chrome-mac-x64/Google Chrome for Testing.app/Contents/MacOS/Google Chrome for Testing",
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(target_os = "windows")]
|
||||
fn build_puppeteer_binary_path(version_dir: &Path) -> PathBuf {
|
||||
version_dir.join(r"chrome-win64\chrome.exe")
|
||||
}
|
||||
|
||||
#[cfg(not(any(target_os = "macos", target_os = "linux", target_os = "windows")))]
|
||||
fn build_puppeteer_binary_path(version_dir: &Path) -> PathBuf {
|
||||
version_dir.join("chrome")
|
||||
}
|
||||
|
||||
/// Search Playwright's browser cache for a Chromium binary.
|
||||
/// Legacy fallback for users who previously installed Chromium via Playwright.
|
||||
/// This is where `agent-browser install` (via `npx playwright install chromium`) puts it.
|
||||
fn find_playwright_chromium() -> Option<PathBuf> {
|
||||
let mut search_dirs = Vec::new();
|
||||
|
||||
@@ -861,28 +580,6 @@ mod tests {
|
||||
use super::*;
|
||||
use crate::test_utils::EnvGuard;
|
||||
|
||||
#[cfg(unix)]
|
||||
fn spawn_noop_child() -> Child {
|
||||
Command::new("/bin/sh")
|
||||
.args(["-c", "exit 0"])
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.spawn()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
fn spawn_noop_child() -> Child {
|
||||
Command::new("cmd.exe")
|
||||
.args(["/C", "exit 0"])
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.spawn()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_find_chrome_returns_some_on_host() {
|
||||
// This test only makes sense on systems with Chrome installed
|
||||
@@ -923,8 +620,6 @@ mod tests {
|
||||
fn test_chrome_launch_error_no_stderr() {
|
||||
let msg = chrome_launch_error("Chrome exited", &[]);
|
||||
assert!(msg.contains("no stderr output"));
|
||||
assert!(msg.contains("Hint:"));
|
||||
assert!(msg.contains("--no-sandbox"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -941,29 +636,18 @@ 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() {
|
||||
let guard = EnvGuard::new(&["PLAYWRIGHT_BROWSERS_PATH", "HOME", "USERPROFILE"]);
|
||||
guard.set("PLAYWRIGHT_BROWSERS_PATH", "/nonexistent/path");
|
||||
|
||||
let temp_home = std::env::temp_dir().join(format!(
|
||||
"agent-browser-test-home-{}-{}",
|
||||
std::process::id(),
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.expect("system clock should be after unix epoch")
|
||||
.as_nanos()
|
||||
));
|
||||
std::fs::create_dir_all(&temp_home).expect("temp home should be created");
|
||||
let temp_home = temp_home.to_string_lossy().to_string();
|
||||
guard.set("HOME", &temp_home);
|
||||
guard.set("USERPROFILE", &temp_home);
|
||||
|
||||
let _guard = EnvGuard::new(&["PLAYWRIGHT_BROWSERS_PATH"]);
|
||||
_guard.set("PLAYWRIGHT_BROWSERS_PATH", "/nonexistent/path");
|
||||
let result = find_playwright_chromium();
|
||||
assert!(result.is_none());
|
||||
}
|
||||
@@ -979,8 +663,7 @@ mod tests {
|
||||
assert!(result
|
||||
.args
|
||||
.iter()
|
||||
.any(|a| a == "--enable-unsafe-swiftshader"));
|
||||
assert!(result.args.iter().any(|a| a == "--window-size=1280,720"));
|
||||
.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();
|
||||
@@ -996,10 +679,6 @@ mod tests {
|
||||
};
|
||||
let result = build_chrome_args(&opts).unwrap();
|
||||
assert!(!result.args.iter().any(|a| a.contains("--headless")));
|
||||
assert!(!result
|
||||
.args
|
||||
.iter()
|
||||
.any(|a| a == "--enable-unsafe-swiftshader"));
|
||||
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());
|
||||
@@ -1043,8 +722,14 @@ mod tests {
|
||||
..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"));
|
||||
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);
|
||||
}
|
||||
@@ -1065,65 +750,6 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[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!(
|
||||
@@ -1137,7 +763,13 @@ mod tests {
|
||||
// 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 = spawn_noop_child();
|
||||
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(),
|
||||
|
||||
@@ -1,31 +1,17 @@
|
||||
use std::collections::HashMap;
|
||||
use std::io::Write;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use serde_json::Value;
|
||||
use tokio::sync::{broadcast, oneshot, Mutex};
|
||||
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
|
||||
use tokio_tungstenite::connect_async;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
|
||||
use super::types::{CdpCommand, CdpEvent, CdpMessage};
|
||||
|
||||
type PendingMap = Arc<Mutex<HashMap<u64, oneshot::Sender<CdpMessage>>>>;
|
||||
|
||||
/// Interval between WebSocket ping frames sent to keep the connection alive
|
||||
/// through intermediate proxies (reverse proxies, load balancers, service meshes).
|
||||
const WS_KEEPALIVE_INTERVAL_SECS: u64 = 30;
|
||||
|
||||
/// Raw incoming CDP message (text) broadcast to all subscribers.
|
||||
/// Used by the inspect proxy to forward responses and events to DevTools.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RawCdpMessage {
|
||||
pub text: String,
|
||||
pub session_id: Option<String>,
|
||||
}
|
||||
|
||||
pub struct CdpClient {
|
||||
ws_tx: Arc<
|
||||
Mutex<
|
||||
@@ -40,110 +26,35 @@ pub struct CdpClient {
|
||||
next_id: AtomicU64,
|
||||
pending: PendingMap,
|
||||
event_tx: broadcast::Sender<CdpEvent>,
|
||||
raw_tx: broadcast::Sender<RawCdpMessage>,
|
||||
_reader_handle: tokio::task::JoinHandle<()>,
|
||||
_keepalive_handle: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl CdpClient {
|
||||
pub async fn connect(url: &str) -> Result<Self, String> {
|
||||
Self::connect_with_headers(url, None).await
|
||||
}
|
||||
|
||||
pub async fn connect_with_headers(
|
||||
url: &str,
|
||||
headers: Option<Vec<(String, String)>>,
|
||||
) -> Result<Self, String> {
|
||||
let mut request = url
|
||||
.into_client_request()
|
||||
.map_err(|e| format!("Invalid WebSocket URL: {}", e))?;
|
||||
|
||||
if let Some(hdrs) = headers {
|
||||
let req_headers = request.headers_mut();
|
||||
for (key, value) in hdrs {
|
||||
if let (Ok(name), Ok(val)) = (
|
||||
key.parse::<tokio_tungstenite::tungstenite::http::header::HeaderName>(),
|
||||
value.parse::<tokio_tungstenite::tungstenite::http::header::HeaderValue>(),
|
||||
) {
|
||||
req_headers.insert(name, val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let ws_config = WebSocketConfig {
|
||||
max_message_size: None,
|
||||
max_frame_size: None,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let (ws_stream, _) =
|
||||
tokio_tungstenite::connect_async_with_config(request, Some(ws_config), false)
|
||||
.await
|
||||
.map_err(|e| format!("CDP WebSocket connect failed: {}", e))?;
|
||||
|
||||
enable_tcp_keepalive(ws_stream.get_ref());
|
||||
let (ws_stream, _) = connect_async(url)
|
||||
.await
|
||||
.map_err(|e| format!("CDP WebSocket connect failed: {}", e))?;
|
||||
|
||||
let (ws_tx, mut ws_rx) = ws_stream.split();
|
||||
let ws_tx = Arc::new(Mutex::new(ws_tx));
|
||||
|
||||
let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
|
||||
let (event_tx, _) = broadcast::channel(256);
|
||||
let (raw_tx, _) = broadcast::channel(512);
|
||||
|
||||
let pending_clone = pending.clone();
|
||||
let event_tx_clone = event_tx.clone();
|
||||
let raw_tx_clone = raw_tx.clone();
|
||||
|
||||
// Notify used to stop the keepalive task when the reader loop exits.
|
||||
let (cancel_tx, mut cancel_rx) = tokio::sync::watch::channel(false);
|
||||
|
||||
let reader_handle = tokio::spawn(async move {
|
||||
while let Some(msg) = ws_rx.next().await {
|
||||
// Accept both Text and Binary frames — remote CDP proxies
|
||||
// (e.g. Browserless) may send responses as Binary frames.
|
||||
let msg = match msg {
|
||||
Ok(Message::Text(text)) => text,
|
||||
Ok(Message::Binary(data)) => match String::from_utf8(data) {
|
||||
Ok(text) => text,
|
||||
Err(_) => continue,
|
||||
},
|
||||
Ok(Message::Close(frame)) => {
|
||||
if std::env::var("AGENT_BROWSER_DEBUG").is_ok() {
|
||||
let reason = frame
|
||||
.as_ref()
|
||||
.map(|f| format!("code={}, reason={}", f.code, f.reason))
|
||||
.unwrap_or_else(|| "no frame".to_string());
|
||||
let _ =
|
||||
writeln!(std::io::stderr(), "[cdp] WebSocket Close: {}", reason);
|
||||
}
|
||||
break;
|
||||
}
|
||||
Ok(Message::Pong(_)) => continue,
|
||||
Ok(Message::Close(_)) => break,
|
||||
Ok(_) => continue,
|
||||
Err(e) => {
|
||||
if std::env::var("AGENT_BROWSER_DEBUG").is_ok() {
|
||||
let _ = writeln!(std::io::stderr(), "[cdp] WebSocket Error: {}", e);
|
||||
}
|
||||
break;
|
||||
}
|
||||
Err(_) => break,
|
||||
};
|
||||
|
||||
// Broadcast raw message for inspect proxy subscribers before typed parse,
|
||||
// so messages with negative IDs (used by the inspect proxy) are still delivered.
|
||||
if raw_tx_clone.receiver_count() > 0 {
|
||||
let session_id = serde_json::from_str::<serde_json::Value>(&msg)
|
||||
.ok()
|
||||
.and_then(|v| v.get("sessionId")?.as_str().map(String::from));
|
||||
let _ = raw_tx_clone.send(RawCdpMessage {
|
||||
text: msg.clone(),
|
||||
session_id,
|
||||
});
|
||||
}
|
||||
|
||||
let parsed: CdpMessage = match serde_json::from_str(&msg) {
|
||||
Ok(m) => m,
|
||||
// Expected for inspect proxy messages with negative IDs
|
||||
// (CdpMessage.id is u64); handled via raw broadcast above.
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
@@ -163,33 +74,6 @@ impl CdpClient {
|
||||
let _ = event_tx_clone.send(event);
|
||||
}
|
||||
}
|
||||
|
||||
// Reader loop exited (connection closed or error). Drop all pending
|
||||
// command senders so callers get an immediate channel-closed error
|
||||
// instead of waiting for the 30-second timeout.
|
||||
pending_clone.lock().await.clear();
|
||||
|
||||
// Stop the keepalive task — the connection is gone.
|
||||
let _ = cancel_tx.send(true);
|
||||
});
|
||||
|
||||
// Spawn a keepalive task that sends WebSocket Ping frames at a regular
|
||||
// interval. This prevents intermediate proxies (Envoy, nginx, OpenResty,
|
||||
// cloud load balancers) from closing idle WebSocket connections. If the
|
||||
// send fails, the connection is dead and we stop pinging.
|
||||
let keepalive_tx = ws_tx.clone();
|
||||
let keepalive_handle = tokio::spawn(async move {
|
||||
let interval = std::time::Duration::from_secs(WS_KEEPALIVE_INTERVAL_SECS);
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = tokio::time::sleep(interval) => {}
|
||||
_ = cancel_rx.changed() => break,
|
||||
}
|
||||
let mut tx = keepalive_tx.lock().await;
|
||||
if tx.send(Message::Ping(Vec::new())).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(Self {
|
||||
@@ -197,9 +81,7 @@ impl CdpClient {
|
||||
next_id: AtomicU64::new(1),
|
||||
pending,
|
||||
event_tx,
|
||||
raw_tx,
|
||||
_reader_handle: reader_handle,
|
||||
_keepalive_handle: keepalive_handle,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -215,7 +97,7 @@ impl CdpClient {
|
||||
id,
|
||||
method: method.to_string(),
|
||||
params,
|
||||
session_id: session_id.filter(|s| !s.is_empty()).map(|s| s.to_string()),
|
||||
session_id: session_id.map(|s| s.to_string()),
|
||||
};
|
||||
|
||||
let json = serde_json::to_string(&cmd)
|
||||
@@ -256,21 +138,6 @@ impl CdpClient {
|
||||
self.event_tx.subscribe()
|
||||
}
|
||||
|
||||
/// Subscribe to all raw incoming CDP messages (responses + events).
|
||||
/// Used by the inspect proxy to forward traffic to the DevTools frontend.
|
||||
pub fn subscribe_raw(&self) -> broadcast::Receiver<RawCdpMessage> {
|
||||
self.raw_tx.subscribe()
|
||||
}
|
||||
|
||||
/// Create a lightweight handle for the inspect WebSocket proxy.
|
||||
/// Contains only what's needed to forward messages bidirectionally.
|
||||
pub fn inspect_handle(&self) -> InspectProxyHandle {
|
||||
InspectProxyHandle {
|
||||
ws_tx: self.ws_tx.clone(),
|
||||
raw_tx: self.raw_tx.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn send_command_typed<P: serde::Serialize, R: serde::de::DeserializeOwned>(
|
||||
&self,
|
||||
method: &str,
|
||||
@@ -293,69 +160,4 @@ impl CdpClient {
|
||||
) -> Result<Value, String> {
|
||||
self.send_command(method, None, session_id).await
|
||||
}
|
||||
|
||||
/// Send raw JSON through the WebSocket without tracking a response.
|
||||
/// Used by the inspect proxy to forward DevTools frontend messages.
|
||||
pub async fn send_raw(&self, json: String) -> Result<(), String> {
|
||||
let mut ws_tx = self.ws_tx.lock().await;
|
||||
ws_tx
|
||||
.send(Message::Text(json))
|
||||
.await
|
||||
.map_err(|e| format!("Failed to send raw CDP message: {}", e))
|
||||
}
|
||||
}
|
||||
|
||||
type WsTx = Arc<
|
||||
Mutex<
|
||||
futures_util::stream::SplitSink<
|
||||
tokio_tungstenite::WebSocketStream<
|
||||
tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
|
||||
>,
|
||||
Message,
|
||||
>,
|
||||
>,
|
||||
>;
|
||||
|
||||
/// Lightweight handle for the inspect WebSocket proxy, holding only
|
||||
/// the cloneable parts of CdpClient needed for bidirectional message forwarding.
|
||||
pub struct InspectProxyHandle {
|
||||
ws_tx: WsTx,
|
||||
raw_tx: broadcast::Sender<RawCdpMessage>,
|
||||
}
|
||||
|
||||
impl InspectProxyHandle {
|
||||
pub async fn send_raw(&self, json: String) -> Result<(), String> {
|
||||
let mut ws_tx = self.ws_tx.lock().await;
|
||||
ws_tx
|
||||
.send(Message::Text(json))
|
||||
.await
|
||||
.map_err(|e| format!("Failed to send raw CDP message: {}", e))
|
||||
}
|
||||
|
||||
pub fn subscribe_raw(&self) -> broadcast::Receiver<RawCdpMessage> {
|
||||
self.raw_tx.subscribe()
|
||||
}
|
||||
}
|
||||
|
||||
/// Enable TCP SO_KEEPALIVE on the underlying socket of a WebSocket connection.
|
||||
/// This is best-effort: failures are silently ignored since the WebSocket-level
|
||||
/// Ping keepalive provides the primary connection liveness mechanism.
|
||||
fn enable_tcp_keepalive(stream: &tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>) {
|
||||
let tcp_stream = match stream {
|
||||
tokio_tungstenite::MaybeTlsStream::Plain(s) => s,
|
||||
tokio_tungstenite::MaybeTlsStream::Rustls(s) => s.get_ref().0,
|
||||
_ => return,
|
||||
};
|
||||
|
||||
// SockRef borrows the fd without taking ownership.
|
||||
let sock = socket2::SockRef::from(tcp_stream);
|
||||
let keepalive = socket2::TcpKeepalive::new().with_time(std::time::Duration::from_secs(30));
|
||||
|
||||
// with_interval sets TCP_KEEPINTVL — the time between probes after the
|
||||
// first keepalive probe goes unanswered. Available on most platforms
|
||||
// (Linux, macOS, Windows, FreeBSD, etc.) but not OpenBSD or Haiku.
|
||||
#[cfg(not(any(target_os = "openbsd", target_os = "haiku")))]
|
||||
let keepalive = keepalive.with_interval(std::time::Duration::from_secs(10));
|
||||
|
||||
let _ = sock.set_tcp_keepalive(&keepalive);
|
||||
}
|
||||
|
||||
@@ -1,387 +0,0 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
|
||||
use super::types::BrowserVersionInfo;
|
||||
|
||||
/// Default timeout for CDP discovery HTTP requests.
|
||||
const DEFAULT_DISCOVERY_TIMEOUT: Duration = Duration::from_secs(2);
|
||||
|
||||
/// Discover the CDP WebSocket URL for the given host and port.
|
||||
///
|
||||
/// Tries three methods in order: `/json/version`, `/json/list`, and a direct
|
||||
/// WebSocket connection to `/devtools/browser`. The returned URL has its
|
||||
/// host/port rewritten to match the requested target.
|
||||
///
|
||||
/// An optional `query` string (without the leading `?`) is appended to the
|
||||
/// final WebSocket URL so that user-supplied URL parameters (e.g.
|
||||
/// `?mode=Hello`) are forwarded to the remote endpoint.
|
||||
pub async fn discover_cdp_url(
|
||||
host: &str,
|
||||
port: u16,
|
||||
query: Option<&str>,
|
||||
) -> Result<String, String> {
|
||||
discover_cdp_url_with_timeout(host, port, query, DEFAULT_DISCOVERY_TIMEOUT).await
|
||||
}
|
||||
|
||||
/// Like [`discover_cdp_url`] but with a custom request timeout.
|
||||
pub async fn discover_cdp_url_with_timeout(
|
||||
host: &str,
|
||||
port: u16,
|
||||
query: Option<&str>,
|
||||
timeout: Duration,
|
||||
) -> Result<String, String> {
|
||||
// Primary: /json/version (standard path)
|
||||
let version_err = match fetch_cdp_info(host, port, timeout).await {
|
||||
Ok(info) => {
|
||||
if let Some(ws_url) = info.web_socket_debugger_url {
|
||||
return Ok(append_query(&rewrite_ws_host(&ws_url, host, port), query));
|
||||
}
|
||||
format!(
|
||||
"No webSocketDebuggerUrl in /json/version at {}:{}",
|
||||
host, port
|
||||
)
|
||||
}
|
||||
Err(e) => e,
|
||||
};
|
||||
|
||||
// Fallback: /json/list (returns target list; look for the browser target)
|
||||
let list_err = match fetch_cdp_list(host, port, timeout).await {
|
||||
Ok(ws_url) => return Ok(append_query(&rewrite_ws_host(&ws_url, host, port), query)),
|
||||
Err(e) => e,
|
||||
};
|
||||
|
||||
// Final fallback: direct WebSocket at /devtools/browser.
|
||||
// Chrome 136+ with UI-based remote debugging (chrome://inspect) exposes
|
||||
// CDP over WebSocket but does not serve HTTP discovery endpoints.
|
||||
match discover_cdp_ws(host, port, timeout).await {
|
||||
Ok(ws_url) => Ok(append_query(&ws_url, query)),
|
||||
Err(ws_err) => Err(format!(
|
||||
"All CDP discovery methods failed for {}:{}: /json/version: {}; /json/list: {}; WebSocket: {}",
|
||||
host, port, version_err, list_err, ws_err
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Bracket an IPv6 address for use in URLs. No-op for IPv4 or already-bracketed addresses.
|
||||
fn bracket_ipv6(host: &str) -> String {
|
||||
if host.contains(':') && !host.starts_with('[') {
|
||||
format!("[{}]", host)
|
||||
} else {
|
||||
host.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
/// Fetch `/json/version` from the given host:port and parse the response.
|
||||
async fn fetch_cdp_info(
|
||||
host: &str,
|
||||
port: u16,
|
||||
timeout: Duration,
|
||||
) -> Result<BrowserVersionInfo, String> {
|
||||
let url = format!("http://{}:{}/json/version", bracket_ipv6(host), port);
|
||||
|
||||
let body = tokio::time::timeout(timeout, reqwest_get_string(&url))
|
||||
.await
|
||||
.map_err(|_| format!("Timeout connecting to CDP at {}:{}", host, port))?
|
||||
.map_err(|e| format!("Failed to connect to CDP at {}:{}: {}", host, port, e))?;
|
||||
|
||||
serde_json::from_str(&body).map_err(|e| format!("Invalid /json/version response: {}", e))
|
||||
}
|
||||
|
||||
/// Rewrite the host and port in a WebSocket URL to match the target we
|
||||
/// actually connected to. Chrome's `/json/version` always returns
|
||||
/// `ws://127.0.0.1:<local-port>/...` which is unreachable when the
|
||||
/// browser is on a remote machine or behind a port-forward.
|
||||
fn rewrite_ws_host(ws_url: &str, host: &str, port: u16) -> String {
|
||||
if let Ok(mut parsed) = url::Url::parse(ws_url) {
|
||||
let _ = parsed.set_host(Some(&bracket_ipv6(host)));
|
||||
let _ = parsed.set_port(Some(port));
|
||||
parsed.to_string()
|
||||
} else {
|
||||
ws_url.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
/// Append a query string to a URL, preserving any existing query parameters.
|
||||
fn append_query(url: &str, query: Option<&str>) -> String {
|
||||
match query {
|
||||
Some(q) if !q.is_empty() => {
|
||||
if let Ok(mut parsed) = url::Url::parse(url) {
|
||||
{
|
||||
let mut pairs = parsed.query_pairs_mut();
|
||||
pairs.extend_pairs(url::form_urlencoded::parse(q.as_bytes()));
|
||||
}
|
||||
parsed.to_string()
|
||||
} else {
|
||||
// Fallback: raw string append
|
||||
if url.contains('?') {
|
||||
format!("{}&{}", url, q)
|
||||
} else {
|
||||
format!("{}?{}", url, q)
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => url.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Fetch `/json/list` and extract the `webSocketDebuggerUrl` from the first
|
||||
/// target with `type == "browser"`, or the first target if none has that type.
|
||||
async fn fetch_cdp_list(host: &str, port: u16, timeout: Duration) -> Result<String, String> {
|
||||
let url = format!("http://{}:{}/json/list", bracket_ipv6(host), port);
|
||||
|
||||
let body = tokio::time::timeout(timeout, reqwest_get_string(&url))
|
||||
.await
|
||||
.map_err(|_| format!("Timeout connecting to /json/list at {}:{}", host, port))?
|
||||
.map_err(|e| {
|
||||
format!(
|
||||
"Failed to connect to /json/list at {}:{}: {}",
|
||||
host, port, e
|
||||
)
|
||||
})?;
|
||||
|
||||
let targets: Vec<serde_json::Value> =
|
||||
serde_json::from_str(&body).map_err(|e| format!("Invalid /json/list response: {}", e))?;
|
||||
|
||||
// Prefer targets with type "browser", fall back to first target with a ws URL
|
||||
let browser_target = targets
|
||||
.iter()
|
||||
.find(|t| t.get("type").and_then(|v| v.as_str()) == Some("browser"));
|
||||
|
||||
let target = browser_target.or_else(|| targets.first());
|
||||
|
||||
target
|
||||
.and_then(|t| t.get("webSocketDebuggerUrl"))
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| "No webSocketDebuggerUrl found in /json/list targets".to_string())
|
||||
}
|
||||
|
||||
/// Discover a CDP endpoint by connecting directly to `ws://host:port/devtools/browser`
|
||||
/// and verifying it responds to `Browser.getVersion`.
|
||||
/// Returns the WebSocket URL on success.
|
||||
async fn discover_cdp_ws(host: &str, port: u16, timeout: Duration) -> Result<String, String> {
|
||||
let ws_url = format!("ws://{}:{}/devtools/browser", bracket_ipv6(host), port);
|
||||
|
||||
tokio::time::timeout(timeout, async {
|
||||
let (mut ws_stream, _) = tokio_tungstenite::connect_async(&ws_url)
|
||||
.await
|
||||
.map_err(|e| format!("WebSocket connect failed at {}: {}", ws_url, e))?;
|
||||
|
||||
let cmd = r#"{"id":1,"method":"Browser.getVersion"}"#;
|
||||
ws_stream
|
||||
.send(Message::Text(cmd.into()))
|
||||
.await
|
||||
.map_err(|e| format!("Failed to send command: {}", e))?;
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct CdpReply {
|
||||
id: u64,
|
||||
}
|
||||
|
||||
let mut result: Result<(), String> = Err("No valid CDP response received".to_string());
|
||||
while let Some(msg) = ws_stream.next().await {
|
||||
match msg {
|
||||
Ok(Message::Text(text)) => {
|
||||
if serde_json::from_str::<CdpReply>(&text).is_ok_and(|r| r.id == 1) {
|
||||
result = Ok(());
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(Message::Close(_)) | Err(_) => break,
|
||||
_ => continue,
|
||||
}
|
||||
}
|
||||
|
||||
let _ = ws_stream.close(None).await;
|
||||
result
|
||||
})
|
||||
.await
|
||||
.map_err(|_| format!("Timeout connecting to WebSocket at {}", ws_url))?
|
||||
.map(|()| ws_url)
|
||||
}
|
||||
|
||||
async fn reqwest_get_string(url: &str) -> Result<String, String> {
|
||||
let resp = reqwest::get(url).await.map_err(|e| e.to_string())?;
|
||||
resp.text().await.map_err(|e| e.to_string())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
const HTTP_404: &str =
|
||||
"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
|
||||
|
||||
fn http_200(body: &str) -> String {
|
||||
format!(
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\nContent-Type: application/json\r\n\r\n{}",
|
||||
body.len(), body
|
||||
)
|
||||
}
|
||||
|
||||
async fn accept_http(listener: &TcpListener, response: &str) {
|
||||
let (mut s, _) = listener.accept().await.unwrap();
|
||||
let mut buf = [0u8; 1024];
|
||||
let _ = s.read(&mut buf).await;
|
||||
s.write_all(response.as_bytes()).await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn discovers_ws_url_from_json_version() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let port = listener.local_addr().unwrap().port();
|
||||
let server = tokio::spawn(async move {
|
||||
accept_http(
|
||||
&listener,
|
||||
&http_200(r#"{"webSocketDebuggerUrl":"ws://127.0.0.1:1234/"}"#),
|
||||
)
|
||||
.await;
|
||||
});
|
||||
|
||||
let ws_url = discover_cdp_url("127.0.0.1", port, None).await.unwrap();
|
||||
assert_eq!(ws_url, format!("ws://127.0.0.1:{}/", port));
|
||||
server.await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn returns_error_when_version_returns_invalid_json() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let port = listener.local_addr().unwrap().port();
|
||||
let server = tokio::spawn(async move {
|
||||
accept_http(&listener, &http_200("not-json")).await;
|
||||
// /json/list and ws fallback both fail (server closes)
|
||||
});
|
||||
|
||||
let err = discover_cdp_url("127.0.0.1", port, None).await.unwrap_err();
|
||||
assert!(err.contains("Invalid /json/version response"));
|
||||
server.await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn falls_back_to_json_list_on_version_404() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let port = listener.local_addr().unwrap().port();
|
||||
let server = tokio::spawn(async move {
|
||||
accept_http(&listener, HTTP_404).await;
|
||||
accept_http(
|
||||
&listener,
|
||||
&http_200(r#"[{"type":"browser","webSocketDebuggerUrl":"ws://127.0.0.1:1234/devtools/browser/abc"}]"#),
|
||||
).await;
|
||||
});
|
||||
|
||||
let ws_url = discover_cdp_url("127.0.0.1", port, None).await.unwrap();
|
||||
assert!(ws_url.contains("/devtools/browser/abc"));
|
||||
assert!(ws_url.contains(&port.to_string()));
|
||||
server.await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn falls_back_to_ws_when_http_returns_404() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let port = listener.local_addr().unwrap().port();
|
||||
let server = tokio::spawn(async move {
|
||||
// /json/version -> 404, /json/list -> 404
|
||||
accept_http(&listener, HTTP_404).await;
|
||||
accept_http(&listener, HTTP_404).await;
|
||||
|
||||
// WebSocket handshake + respond to Browser.getVersion
|
||||
let (stream, _) = listener.accept().await.unwrap();
|
||||
let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap();
|
||||
if let Some(Ok(Message::Text(text))) = ws.next().await {
|
||||
let req: serde_json::Value = serde_json::from_str(&text).unwrap();
|
||||
let id = req.get("id").unwrap();
|
||||
let reply = format!(
|
||||
r#"{{"id":{},"result":{{"protocolVersion":"1.3","product":"Chrome/136"}}}}"#,
|
||||
id
|
||||
);
|
||||
ws.send(Message::Text(reply)).await.unwrap();
|
||||
}
|
||||
let _ = ws.close(None).await;
|
||||
});
|
||||
|
||||
let ws_url = discover_cdp_url("127.0.0.1", port, None).await.unwrap();
|
||||
assert_eq!(ws_url, format!("ws://127.0.0.1:{}/devtools/browser", port));
|
||||
server.await.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rewrite_ws_host_replaces_host_and_port() {
|
||||
let original = "ws://127.0.0.1:9222/devtools/browser/abc";
|
||||
let rewritten = rewrite_ws_host(original, "10.211.55.12", 9223);
|
||||
assert_eq!(rewritten, "ws://10.211.55.12:9223/devtools/browser/abc");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rewrite_ws_host_handles_ipv6() {
|
||||
let original = "ws://127.0.0.1:9222/devtools/browser/abc";
|
||||
let rewritten = rewrite_ws_host(original, "::1", 9222);
|
||||
assert_eq!(rewritten, "ws://[::1]:9222/devtools/browser/abc");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn append_query_adds_params_to_url_without_query() {
|
||||
let url = "ws://127.0.0.1:9222/devtools/browser/abc";
|
||||
let result = append_query(url, Some("mode=Hello"));
|
||||
assert_eq!(
|
||||
result,
|
||||
"ws://127.0.0.1:9222/devtools/browser/abc?mode=Hello"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn append_query_merges_with_existing_query() {
|
||||
let url = "ws://127.0.0.1:9222/devtools/browser/abc?token=xyz";
|
||||
let result = append_query(url, Some("mode=Hello"));
|
||||
assert_eq!(
|
||||
result,
|
||||
"ws://127.0.0.1:9222/devtools/browser/abc?token=xyz&mode=Hello"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn append_query_noop_for_none() {
|
||||
let url = "ws://127.0.0.1:9222/devtools/browser/abc";
|
||||
let result = append_query(url, None);
|
||||
assert_eq!(result, url);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn append_query_noop_for_empty() {
|
||||
let url = "ws://127.0.0.1:9222/devtools/browser/abc";
|
||||
let result = append_query(url, Some(""));
|
||||
assert_eq!(result, url);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn append_query_handles_multiple_params() {
|
||||
let url = "ws://127.0.0.1:9222/devtools/browser/abc";
|
||||
let result = append_query(url, Some("mode=Hello&token=abc"));
|
||||
assert_eq!(
|
||||
result,
|
||||
"ws://127.0.0.1:9222/devtools/browser/abc?mode=Hello&token=abc"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn discover_preserves_query_params() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let port = listener.local_addr().unwrap().port();
|
||||
let server = tokio::spawn(async move {
|
||||
accept_http(
|
||||
&listener,
|
||||
&http_200(r#"{"webSocketDebuggerUrl":"ws://127.0.0.1:1234/"}"#),
|
||||
)
|
||||
.await;
|
||||
});
|
||||
|
||||
let ws_url = discover_cdp_url("127.0.0.1", port, Some("mode=Hello"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(ws_url, format!("ws://127.0.0.1:{}/?mode=Hello", port));
|
||||
server.await.unwrap();
|
||||
}
|
||||
}
|
||||
@@ -1,495 +0,0 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::io::{BufRead, BufReader};
|
||||
use std::net::TcpListener;
|
||||
use std::path::PathBuf;
|
||||
use std::process::{Child, Command, Stdio};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use super::discovery::discover_cdp_url_with_timeout;
|
||||
|
||||
const LIGHTPANDA_STARTUP_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
const LIGHTPANDA_POLL_INTERVAL: Duration = Duration::from_millis(100);
|
||||
const LIGHTPANDA_DISCOVERY_TIMEOUT: Duration = Duration::from_millis(500);
|
||||
const LIGHTPANDA_SESSION_TIMEOUT_SECS: u64 = 604800; // 1 week, the documented maximum
|
||||
const MAX_LOG_LINES: usize = 40;
|
||||
|
||||
pub struct LightpandaProcess {
|
||||
child: Child,
|
||||
pub ws_url: String,
|
||||
_log_drainers: Vec<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>,
|
||||
}
|
||||
|
||||
fn build_lightpanda_serve_args(port: u16, proxy: Option<&str>) -> Vec<String> {
|
||||
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(),
|
||||
LIGHTPANDA_SESSION_TIMEOUT_SECS.to_string(),
|
||||
];
|
||||
|
||||
if let Some(proxy) = proxy {
|
||||
args.push("--http_proxy".to_string());
|
||||
args.push(proxy.to_string());
|
||||
}
|
||||
|
||||
args
|
||||
}
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
struct LaunchLogBuffer {
|
||||
stdout: Arc<Mutex<VecDeque<String>>>,
|
||||
stderr: Arc<Mutex<VecDeque<String>>>,
|
||||
}
|
||||
|
||||
impl LaunchLogBuffer {
|
||||
fn push_stdout(&self, line: String) {
|
||||
push_bounded(&self.stdout, line);
|
||||
}
|
||||
|
||||
fn push_stderr(&self, line: String) {
|
||||
push_bounded(&self.stderr, line);
|
||||
}
|
||||
|
||||
fn snapshot_stdout(&self) -> Vec<String> {
|
||||
self.stdout
|
||||
.lock()
|
||||
.expect("stdout log buffer poisoned")
|
||||
.iter()
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn snapshot_stderr(&self) -> Vec<String> {
|
||||
self.stderr
|
||||
.lock()
|
||||
.expect("stderr log buffer poisoned")
|
||||
.iter()
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
fn push_bounded(buffer: &Mutex<VecDeque<String>>, line: String) {
|
||||
let mut guard = buffer.lock().expect("log buffer poisoned");
|
||||
if guard.len() >= MAX_LOG_LINES {
|
||||
guard.pop_front();
|
||||
}
|
||||
guard.push_back(line);
|
||||
}
|
||||
|
||||
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 c in &candidates {
|
||||
if c.exists() {
|
||||
return Some(c.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
pub async fn launch_lightpanda(
|
||||
options: &LightpandaLaunchOptions,
|
||||
) -> Result<LightpandaProcess, String> {
|
||||
let binary_path = match &options.executable_path {
|
||||
Some(p) => PathBuf::from(p),
|
||||
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(p) => p,
|
||||
None => TcpListener::bind("127.0.0.1:0")
|
||||
.and_then(|l| l.local_addr())
|
||||
.map(|a| a.port())
|
||||
.map_err(|e| format!("Failed to find an available port for Lightpanda: {}", e))?,
|
||||
};
|
||||
let args = build_lightpanda_serve_args(port, options.proxy.as_deref());
|
||||
|
||||
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 (log_buffer, log_drainers) = start_log_drainers(&mut child)?;
|
||||
|
||||
let ws_url =
|
||||
match wait_for_lightpanda_ready(&mut child, port, &log_buffer, LIGHTPANDA_STARTUP_TIMEOUT)
|
||||
.await
|
||||
{
|
||||
Ok(url) => url,
|
||||
Err(e) => {
|
||||
let _ = child.kill();
|
||||
let _ = child.wait();
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
|
||||
Ok(LightpandaProcess {
|
||||
child,
|
||||
ws_url,
|
||||
_log_drainers: log_drainers,
|
||||
})
|
||||
}
|
||||
|
||||
fn start_log_drainers(
|
||||
child: &mut Child,
|
||||
) -> Result<(LaunchLogBuffer, Vec<std::thread::JoinHandle<()>>), String> {
|
||||
let stdout = child.stdout.take().ok_or_else(|| {
|
||||
let _ = child.kill();
|
||||
"Failed to capture Lightpanda stdout".to_string()
|
||||
})?;
|
||||
let stderr = child.stderr.take().ok_or_else(|| {
|
||||
let _ = child.kill();
|
||||
"Failed to capture Lightpanda stderr".to_string()
|
||||
})?;
|
||||
|
||||
let logs = LaunchLogBuffer::default();
|
||||
let stdout_logs = logs.clone();
|
||||
let stderr_logs = logs.clone();
|
||||
|
||||
let stdout_handle =
|
||||
std::thread::spawn(move || drain_reader(stdout, move |line| stdout_logs.push_stdout(line)));
|
||||
let stderr_handle =
|
||||
std::thread::spawn(move || drain_reader(stderr, move |line| stderr_logs.push_stderr(line)));
|
||||
|
||||
Ok((logs, vec![stdout_handle, stderr_handle]))
|
||||
}
|
||||
|
||||
fn drain_reader<R, F>(reader: R, mut push: F)
|
||||
where
|
||||
R: std::io::Read,
|
||||
F: FnMut(String),
|
||||
{
|
||||
for line in BufReader::new(reader).lines() {
|
||||
match line {
|
||||
Ok(line) => push(line),
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_for_lightpanda_ready(
|
||||
child: &mut Child,
|
||||
port: u16,
|
||||
logs: &LaunchLogBuffer,
|
||||
startup_timeout: Duration,
|
||||
) -> Result<String, String> {
|
||||
let deadline = std::time::Instant::now() + startup_timeout;
|
||||
let mut last_probe_error = None;
|
||||
|
||||
loop {
|
||||
if let Ok(Some(status)) = child.try_wait() {
|
||||
// Give the drainer threads a brief window to flush the last log lines
|
||||
// before we snapshot them. This is best-effort: lines written just
|
||||
// before exit may still be missing, but the most useful output (early
|
||||
// startup errors) will already be in the buffer.
|
||||
tokio::time::sleep(Duration::from_millis(25)).await;
|
||||
return Err(lightpanda_launch_error(
|
||||
&format!(
|
||||
"Lightpanda exited before CDP became ready (status: {})",
|
||||
status
|
||||
),
|
||||
logs,
|
||||
last_probe_error.as_deref(),
|
||||
));
|
||||
}
|
||||
|
||||
match discover_cdp_url_with_timeout("127.0.0.1", port, None, LIGHTPANDA_DISCOVERY_TIMEOUT)
|
||||
.await
|
||||
{
|
||||
Ok(ws_url) => return Ok(ws_url),
|
||||
Err(err) => last_probe_error = Some(err),
|
||||
}
|
||||
|
||||
if std::time::Instant::now() >= deadline {
|
||||
return Err(lightpanda_launch_error(
|
||||
&format!(
|
||||
"Timed out after {}ms waiting for Lightpanda CDP endpoint on port {}",
|
||||
startup_timeout.as_millis(),
|
||||
port
|
||||
),
|
||||
logs,
|
||||
last_probe_error.as_deref(),
|
||||
));
|
||||
}
|
||||
|
||||
tokio::time::sleep(LIGHTPANDA_POLL_INTERVAL).await;
|
||||
}
|
||||
}
|
||||
|
||||
fn lightpanda_launch_error(
|
||||
message: &str,
|
||||
logs: &LaunchLogBuffer,
|
||||
last_probe_error: Option<&str>,
|
||||
) -> String {
|
||||
let stdout_lines = logs.snapshot_stdout();
|
||||
let stderr_lines = logs.snapshot_stderr();
|
||||
let mut details = Vec::new();
|
||||
|
||||
if let Some(err) = last_probe_error {
|
||||
details.push(format!("Last probe error: {}", err));
|
||||
}
|
||||
|
||||
if !stderr_lines.is_empty() {
|
||||
details.push(format!(
|
||||
"Lightpanda stderr (last {} lines):\n {}",
|
||||
stderr_lines.len(),
|
||||
stderr_lines.join("\n ")
|
||||
));
|
||||
}
|
||||
|
||||
if !stdout_lines.is_empty() {
|
||||
details.push(format!(
|
||||
"Lightpanda stdout (last {} lines):\n {}",
|
||||
stdout_lines.len(),
|
||||
stdout_lines.join("\n ")
|
||||
));
|
||||
}
|
||||
|
||||
if details.is_empty() {
|
||||
format!("{} (no stdout/stderr output from Lightpanda)", message)
|
||||
} else {
|
||||
format!("{}\n{}", message, details.join("\n"))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpListener as TokioTcpListener;
|
||||
|
||||
fn unused_port() -> u16 {
|
||||
std::net::TcpListener::bind("127.0.0.1:0")
|
||||
.unwrap()
|
||||
.local_addr()
|
||||
.unwrap()
|
||||
.port()
|
||||
}
|
||||
|
||||
async fn serve_json_version_once_after_delay(port: u16, delay_ms: u64, body: &'static str) {
|
||||
tokio::time::sleep(Duration::from_millis(delay_ms)).await;
|
||||
let listener = TokioTcpListener::bind(("127.0.0.1", port)).await.unwrap();
|
||||
let (mut socket, _) = listener.accept().await.unwrap();
|
||||
let mut buf = [0u8; 1024];
|
||||
let _ = socket.read(&mut buf).await;
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\nContent-Type: application/json\r\n\r\n{}",
|
||||
body.len(),
|
||||
body
|
||||
);
|
||||
socket.write_all(response.as_bytes()).await.unwrap();
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn waits_for_ready_without_logs() {
|
||||
let port = unused_port();
|
||||
tokio::spawn(serve_json_version_once_after_delay(
|
||||
port,
|
||||
150,
|
||||
r#"{"webSocketDebuggerUrl":"ws://127.0.0.1:9222/"}"#,
|
||||
));
|
||||
|
||||
let mut child = Command::new("/bin/sh")
|
||||
.args(["-c", "sleep 5"])
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()
|
||||
.unwrap();
|
||||
|
||||
let (logs, _drainers) = start_log_drainers(&mut child).unwrap();
|
||||
let ws_url = wait_for_lightpanda_ready(&mut child, port, &logs, LIGHTPANDA_STARTUP_TIMEOUT)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(ws_url, format!("ws://127.0.0.1:{}/", port));
|
||||
let _ = child.kill();
|
||||
let _ = child.wait();
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn child_exit_surfaces_logs() {
|
||||
let port = unused_port();
|
||||
let mut child = Command::new("/bin/sh")
|
||||
.args(["-c", "echo boom >&2; sleep 0.1; exit 23"])
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()
|
||||
.unwrap();
|
||||
|
||||
let (logs, _drainers) = start_log_drainers(&mut child).unwrap();
|
||||
let err = wait_for_lightpanda_ready(&mut child, port, &logs, LIGHTPANDA_STARTUP_TIMEOUT)
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(err.contains("Lightpanda exited before CDP became ready"));
|
||||
assert!(err.contains("boom"));
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn timeout_reports_last_probe_error() {
|
||||
let port = unused_port();
|
||||
let mut child = Command::new("/bin/sh")
|
||||
.args(["-c", "sleep 30"])
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()
|
||||
.unwrap();
|
||||
|
||||
let timeout = Duration::from_millis(300);
|
||||
let (logs, _drainers) = start_log_drainers(&mut child).unwrap();
|
||||
let err = tokio::time::timeout(
|
||||
Duration::from_secs(2),
|
||||
wait_for_lightpanda_ready(&mut child, port, &logs, timeout),
|
||||
)
|
||||
.await
|
||||
.expect("ready wait should return before outer timeout")
|
||||
.unwrap_err();
|
||||
|
||||
assert!(err.contains("Timed out after 300ms waiting for Lightpanda CDP endpoint"));
|
||||
assert!(
|
||||
err.contains("Failed to connect to CDP") || err.contains("Timeout connecting to CDP")
|
||||
);
|
||||
|
||||
let _ = child.kill();
|
||||
let _ = child.wait();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_find_lightpanda_returns_none_when_missing() {
|
||||
let _ = find_lightpanda();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lightpanda_launch_error_no_logs() {
|
||||
let logs = LaunchLogBuffer::default();
|
||||
let msg = lightpanda_launch_error("Lightpanda exited", &logs, None);
|
||||
assert!(msg.contains("no stdout/stderr output"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lightpanda_launch_error_with_lines() {
|
||||
let logs = LaunchLogBuffer::default();
|
||||
logs.push_stdout("stdout line".to_string());
|
||||
logs.push_stderr("stderr line".to_string());
|
||||
let msg = lightpanda_launch_error("Lightpanda exited", &logs, Some("connect failed"));
|
||||
assert!(msg.contains("stdout line"));
|
||||
assert!(msg.contains("stderr line"));
|
||||
assert!(msg.contains("Last probe error: connect failed"));
|
||||
}
|
||||
|
||||
#[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());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_lightpanda_serve_args_sets_explicit_session_timeout() {
|
||||
let args = build_lightpanda_serve_args(9222, None);
|
||||
|
||||
assert_eq!(
|
||||
args,
|
||||
vec![
|
||||
"serve".to_string(),
|
||||
"--host".to_string(),
|
||||
"127.0.0.1".to_string(),
|
||||
"--port".to_string(),
|
||||
"9222".to_string(),
|
||||
"--timeout".to_string(),
|
||||
"604800".to_string(),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_lightpanda_serve_args_with_proxy() {
|
||||
let args = build_lightpanda_serve_args(9333, Some("http://127.0.0.1:8080"));
|
||||
|
||||
assert_eq!(
|
||||
args,
|
||||
vec![
|
||||
"serve".to_string(),
|
||||
"--host".to_string(),
|
||||
"127.0.0.1".to_string(),
|
||||
"--port".to_string(),
|
||||
"9333".to_string(),
|
||||
"--timeout".to_string(),
|
||||
"604800".to_string(),
|
||||
"--http_proxy".to_string(),
|
||||
"http://127.0.0.1:8080".to_string(),
|
||||
]
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,3 @@
|
||||
pub mod chrome;
|
||||
pub mod client;
|
||||
pub mod discovery;
|
||||
pub mod lightpanda;
|
||||
pub mod types;
|
||||
|
||||
@@ -1,51 +1,6 @@
|
||||
use serde::{Deserialize, Deserializer, Serialize};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
/// Deserialize a value that may be either a string or an integer into a String.
|
||||
/// Lightpanda sends numeric nodeIds/childIds in AX tree responses, while Chrome
|
||||
/// sends strings. This accepts both.
|
||||
fn string_or_int<'de, D>(deserializer: D) -> Result<String, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
let v = Value::deserialize(deserializer)?;
|
||||
match v {
|
||||
Value::String(s) => Ok(s),
|
||||
Value::Number(n) => Ok(n.to_string()),
|
||||
other => Err(serde::de::Error::custom(format!(
|
||||
"expected string or integer, got {}",
|
||||
other
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Deserialize an optional Vec where each element may be a string or integer.
|
||||
fn opt_vec_string_or_int<'de, D>(deserializer: D) -> Result<Option<Vec<String>>, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
let opt: Option<Vec<Value>> = Option::deserialize(deserializer)?;
|
||||
match opt {
|
||||
None => Ok(None),
|
||||
Some(vec) => {
|
||||
let mut result = Vec::with_capacity(vec.len());
|
||||
for v in vec {
|
||||
match v {
|
||||
Value::String(s) => result.push(s),
|
||||
Value::Number(n) => result.push(n.to_string()),
|
||||
other => {
|
||||
return Err(serde::de::Error::custom(format!(
|
||||
"expected string or integer in array, got {}",
|
||||
other
|
||||
)))
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Some(result))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// CDP message envelope
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -260,7 +215,6 @@ pub struct RemoteObject {
|
||||
pub object_id: Option<String>,
|
||||
pub class_name: Option<String>,
|
||||
pub unserializable_value: Option<String>,
|
||||
pub preview: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
@@ -303,14 +257,12 @@ pub struct GetFullAXTreeResult {
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AXNode {
|
||||
#[serde(deserialize_with = "string_or_int")]
|
||||
pub node_id: String,
|
||||
pub role: Option<AXValue>,
|
||||
pub name: Option<AXValue>,
|
||||
pub value: Option<AXValue>,
|
||||
pub description: Option<AXValue>,
|
||||
pub properties: Option<Vec<AXProperty>>,
|
||||
#[serde(default, deserialize_with = "opt_vec_string_or_int")]
|
||||
pub child_ids: Option<Vec<String>>,
|
||||
pub backend_d_o_m_node_id: Option<i64>,
|
||||
pub ignored: Option<bool>,
|
||||
@@ -580,7 +532,6 @@ 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"));
|
||||
}
|
||||
|
||||
@@ -24,19 +24,6 @@ pub struct Cookie {
|
||||
pub same_site: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn get_all_cookies(client: &CdpClient, session_id: &str) -> Result<Vec<Cookie>, String> {
|
||||
let result = client
|
||||
.send_command_no_params("Network.getAllCookies", Some(session_id))
|
||||
.await?;
|
||||
|
||||
let cookies: Vec<Cookie> = result
|
||||
.get("cookies")
|
||||
.and_then(|v| serde_json::from_value(v.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
|
||||
Ok(cookies)
|
||||
}
|
||||
|
||||
pub async fn get_cookies(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
|
||||
+74
-322
@@ -1,71 +1,36 @@
|
||||
use serde_json::Value;
|
||||
use std::env;
|
||||
use std::fs;
|
||||
use std::io::Write;
|
||||
use std::path::PathBuf;
|
||||
use std::process;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::signal;
|
||||
use tokio::sync::{mpsc, RwLock};
|
||||
use tokio::sync::mpsc::{unbounded_channel, UnboundedSender};
|
||||
use tokio::time::{Duration, Instant};
|
||||
|
||||
use super::actions::{execute_command, DaemonState};
|
||||
use super::cdp::client::CdpClient;
|
||||
use super::state;
|
||||
use super::stream::StreamServer;
|
||||
|
||||
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);
|
||||
}
|
||||
|
||||
// When debug mode is on, redirect stderr to a log file so daemon
|
||||
// output can be inspected (the daemon normally has stderr piped to its
|
||||
// parent which drops the read end after startup).
|
||||
#[cfg(unix)]
|
||||
if env::var("AGENT_BROWSER_DEBUG").is_ok() {
|
||||
let log_path = socket_dir.join(format!("{}.log", session));
|
||||
if let Ok(file) = fs::File::create(&log_path) {
|
||||
use std::os::unix::io::IntoRawFd;
|
||||
let fd = file.into_raw_fd();
|
||||
unsafe {
|
||||
libc::dup2(fd, 2);
|
||||
libc::close(fd);
|
||||
}
|
||||
let _ = writeln!(
|
||||
std::io::stderr(),
|
||||
"[daemon] Debug logging started for session: {}",
|
||||
session
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let pid_path = socket_dir.join(format!("{}.pid", session));
|
||||
let _ = fs::write(&pid_path, process::id().to_string());
|
||||
|
||||
// On Unix the daemon listens on a Unix domain socket; on Windows it uses
|
||||
// TCP, so there is no .sock file — only a .port file written by the server.
|
||||
let socket_path = socket_dir.join(format!("{}.sock", session));
|
||||
|
||||
#[cfg(unix)]
|
||||
if socket_path.exists() {
|
||||
let _ = fs::remove_file(&socket_path);
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
{
|
||||
let _ = fs::remove_file(socket_dir.join(format!("{}.port", session)));
|
||||
}
|
||||
|
||||
let stream_path = socket_dir.join(format!("{}.stream", session));
|
||||
let _ = fs::remove_file(&stream_path);
|
||||
let _ = fs::remove_file(socket_dir.join(format!("{}.engine", session)));
|
||||
let _ = fs::remove_file(socket_dir.join(format!("{}.provider", session)));
|
||||
let _ = fs::remove_file(socket_dir.join(format!("{}.extensions", session)));
|
||||
|
||||
if let Ok(days_str) = env::var("AGENT_BROWSER_STATE_EXPIRE_DAYS") {
|
||||
if let Ok(days) = days_str.parse::<u64>() {
|
||||
if days > 0 {
|
||||
@@ -74,57 +39,15 @@ pub async fn run_daemon(session: &str) {
|
||||
}
|
||||
}
|
||||
|
||||
let mut stream_client: Option<Arc<RwLock<Option<Arc<CdpClient>>>>> = None;
|
||||
let mut stream_server_instance: Option<Arc<StreamServer>> = None;
|
||||
let preferred_port = env::var("AGENT_BROWSER_STREAM_PORT")
|
||||
.ok()
|
||||
.and_then(|s| s.parse::<u16>().ok())
|
||||
.unwrap_or(0);
|
||||
match StreamServer::start_without_client(preferred_port, session.to_string(), true).await {
|
||||
Ok((stream_server, client_slot)) => {
|
||||
stream_client = Some(client_slot.clone());
|
||||
if let Err(e) = fs::write(&stream_path, stream_server.port().to_string()) {
|
||||
let _ = writeln!(std::io::stderr(), "Failed to write .stream file: {}", e);
|
||||
}
|
||||
stream_server_instance = Some(Arc::new(stream_server));
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = writeln!(std::io::stderr(), "Stream server failed to start: {}", e);
|
||||
}
|
||||
}
|
||||
let result = run_socket_server(&socket_path, session, resident_mode).await;
|
||||
|
||||
// Auto-shutdown the daemon after this many ms of inactivity (no commands received).
|
||||
// Disabled when unset or 0.
|
||||
let idle_timeout_ms = env::var("AGENT_BROWSER_IDLE_TIMEOUT_MS")
|
||||
.ok()
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.filter(|&ms| ms > 0);
|
||||
|
||||
let result = run_socket_server(
|
||||
&socket_path,
|
||||
session,
|
||||
stream_client,
|
||||
stream_server_instance,
|
||||
idle_timeout_ms,
|
||||
)
|
||||
.await;
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let _ = fs::remove_file(&socket_path);
|
||||
}
|
||||
#[cfg(windows)]
|
||||
{
|
||||
let _ = fs::remove_file(socket_dir.join(format!("{}.port", session)));
|
||||
}
|
||||
let _ = fs::remove_file(&socket_path);
|
||||
let _ = fs::remove_file(&pid_path);
|
||||
let stream_path = socket_dir.join(format!("{}.stream", session));
|
||||
let _ = fs::remove_file(&stream_path);
|
||||
let _ = fs::remove_file(socket_dir.join(format!("{}.engine", session)));
|
||||
let _ = fs::remove_file(socket_dir.join(format!("{}.provider", session)));
|
||||
let _ = fs::remove_file(socket_dir.join(format!("{}.extensions", session)));
|
||||
|
||||
if let Err(e) = result {
|
||||
let _ = writeln!(std::io::stderr(), "Daemon error: {}", e);
|
||||
eprintln!("Daemon error: {}", e);
|
||||
process::exit(1);
|
||||
}
|
||||
}
|
||||
@@ -132,81 +55,49 @@ pub async fn run_daemon(session: &str) {
|
||||
#[cfg(unix)]
|
||||
async fn run_socket_server(
|
||||
socket_path: &PathBuf,
|
||||
session: &str,
|
||||
stream_client: Option<Arc<RwLock<Option<Arc<CdpClient>>>>>,
|
||||
stream_server: Option<Arc<StreamServer>>,
|
||||
idle_timeout_ms: Option<u64>,
|
||||
_session: &str,
|
||||
resident_mode: bool,
|
||||
) -> Result<(), String> {
|
||||
use tokio::net::UnixListener;
|
||||
|
||||
let listener =
|
||||
UnixListener::bind(socket_path).map_err(|e| format!("Failed to bind socket: {}", e))?;
|
||||
|
||||
let stream_file: Option<PathBuf> = if stream_server.is_some() {
|
||||
let dir = socket_path.parent().unwrap_or(std::path::Path::new("."));
|
||||
Some(dir.join(format!("{}.stream", session)))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let state: std::sync::Arc<tokio::sync::Mutex<DaemonState>> = std::sync::Arc::new(
|
||||
tokio::sync::Mutex::new(DaemonState::new_with_stream(stream_client, stream_server)),
|
||||
);
|
||||
|
||||
let (reset_tx, mut reset_rx) = mpsc::channel::<()>(64);
|
||||
let reset_tx = idle_timeout_ms.map(|_| Arc::new(reset_tx));
|
||||
|
||||
let mut drain_interval = tokio::time::interval(Duration::from_millis(500));
|
||||
drain_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
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 {
|
||||
let sleep_future = idle_timeout_ms.map(|ms| tokio::time::sleep(Duration::from_millis(ms)));
|
||||
let mut sleep_pin = sleep_future.map(Box::pin);
|
||||
|
||||
tokio::select! {
|
||||
accept_result = listener.accept() => {
|
||||
match accept_result {
|
||||
Ok((stream, _)) => {
|
||||
let state = state.clone();
|
||||
let reset_tx = reset_tx.clone();
|
||||
let sf = stream_file.clone();
|
||||
let activity_tx = activity_tx.clone();
|
||||
let active_commands = active_commands.clone();
|
||||
tokio::spawn(async move {
|
||||
handle_connection(stream, state, reset_tx, sf).await;
|
||||
handle_connection(stream, state, activity_tx, active_commands).await;
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = writeln!(std::io::stderr(), "Accept error: {}", e);
|
||||
eprintln!("Accept error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = drain_interval.tick() => {
|
||||
let mut s = state.lock().await;
|
||||
if let Some(ref mut mgr) = s.browser {
|
||||
if mgr.has_process_exited() {
|
||||
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;
|
||||
s.browser = None;
|
||||
s.screencasting = false;
|
||||
s.update_stream_client().await;
|
||||
} else {
|
||||
s.drain_cdp_events_background().await;
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
_ = async {
|
||||
if let Some(ref mut s) = sleep_pin {
|
||||
s.as_mut().await
|
||||
} else {
|
||||
std::future::pending::<()>().await
|
||||
}
|
||||
}, if idle_timeout_ms.is_some() => {
|
||||
let mut s = state.lock().await;
|
||||
if let Some(ref mut mgr) = s.browser {
|
||||
let _ = mgr.close().await;
|
||||
}
|
||||
break;
|
||||
}
|
||||
_ = reset_rx.recv(), if idle_timeout_ms.is_some() => {
|
||||
continue;
|
||||
idle_deadline = Instant::now() + Duration::from_secs(IDLE_SHUTDOWN_SECS);
|
||||
}
|
||||
_ = shutdown_signal() => {
|
||||
let mut s = state.lock().await;
|
||||
@@ -225,79 +116,55 @@ async fn run_socket_server(
|
||||
async fn run_socket_server(
|
||||
socket_path: &PathBuf,
|
||||
session: &str,
|
||||
stream_client: Option<Arc<RwLock<Option<Arc<CdpClient>>>>>,
|
||||
stream_server: Option<Arc<StreamServer>>,
|
||||
idle_timeout_ms: Option<u64>,
|
||||
resident_mode: bool,
|
||||
) -> Result<(), String> {
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
let preferred_port = get_port_for_session(session);
|
||||
// Try the hash-derived port first; if it is blocked (e.g. Windows Hyper-V
|
||||
// excluded port range), fall back to an OS-assigned ephemeral port.
|
||||
let listener = match TcpListener::bind(format!("127.0.0.1:{}", preferred_port)).await {
|
||||
Ok(l) => l,
|
||||
Err(_) => TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.map_err(|e| format!("Failed to bind TCP: {}", e))?,
|
||||
};
|
||||
let actual_port = listener
|
||||
.local_addr()
|
||||
.map_err(|e| format!("Failed to get local address: {}", e))?
|
||||
.port();
|
||||
let port = get_port_for_session(session);
|
||||
let listener = TcpListener::bind(format!("127.0.0.1:{}", port))
|
||||
.await
|
||||
.map_err(|e| format!("Failed to bind TCP: {}", e))?;
|
||||
|
||||
let socket_dir = socket_path.parent().unwrap_or(std::path::Path::new("."));
|
||||
let port_path = socket_dir.join(format!("{}.port", session));
|
||||
let _ = fs::write(&port_path, actual_port.to_string());
|
||||
let _ = fs::write(&port_path, port.to_string());
|
||||
|
||||
let stream_file: Option<PathBuf> = if stream_server.is_some() {
|
||||
Some(socket_dir.join(format!("{}.stream", session)))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let state: std::sync::Arc<tokio::sync::Mutex<DaemonState>> = std::sync::Arc::new(
|
||||
tokio::sync::Mutex::new(DaemonState::new_with_stream(stream_client, stream_server)),
|
||||
);
|
||||
|
||||
let (reset_tx, mut reset_rx) = mpsc::channel::<()>(64);
|
||||
let reset_tx = idle_timeout_ms.map(|_| Arc::new(reset_tx));
|
||||
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 {
|
||||
let sleep_future = idle_timeout_ms.map(|ms| tokio::time::sleep(Duration::from_millis(ms)));
|
||||
let mut sleep_pin = sleep_future.map(Box::pin);
|
||||
|
||||
tokio::select! {
|
||||
accept_result = listener.accept() => {
|
||||
match accept_result {
|
||||
Ok((stream, _)) => {
|
||||
let state = state.clone();
|
||||
let reset_tx = reset_tx.clone();
|
||||
let sf = stream_file.clone();
|
||||
let activity_tx = activity_tx.clone();
|
||||
let active_commands = active_commands.clone();
|
||||
tokio::spawn(async move {
|
||||
handle_connection(stream, state, reset_tx, sf).await;
|
||||
handle_connection(stream, state, activity_tx, active_commands).await;
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = writeln!(std::io::stderr(), "Accept error: {}", e);
|
||||
eprintln!("Accept error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = async {
|
||||
if let Some(ref mut s) = sleep_pin {
|
||||
s.as_mut().await
|
||||
} else {
|
||||
std::future::pending::<()>().await
|
||||
}
|
||||
}, if idle_timeout_ms.is_some() => {
|
||||
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;
|
||||
Some(_) = activity_rx.recv() => {
|
||||
idle_deadline = Instant::now() + Duration::from_secs(IDLE_SHUTDOWN_SECS);
|
||||
}
|
||||
_ = reset_rx.recv(), if idle_timeout_ms.is_some() => {
|
||||
continue;
|
||||
_ = 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;
|
||||
@@ -316,9 +183,10 @@ async fn run_socket_server(
|
||||
async fn handle_connection<S>(
|
||||
stream: S,
|
||||
state: std::sync::Arc<tokio::sync::Mutex<DaemonState>>,
|
||||
idle_reset_tx: Option<Arc<mpsc::Sender<()>>>,
|
||||
stream_file_cleanup: Option<PathBuf>,
|
||||
) where
|
||||
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);
|
||||
@@ -353,11 +221,9 @@ async fn handle_connection<S>(
|
||||
}
|
||||
};
|
||||
|
||||
if let Some(ref tx) = idle_reset_tx {
|
||||
let _ = tx.try_send(());
|
||||
}
|
||||
|
||||
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;
|
||||
@@ -367,13 +233,13 @@ async fn handle_connection<S>(
|
||||
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 {
|
||||
if let Some(ref path) = stream_file_cleanup {
|
||||
let _ = fs::remove_file(path);
|
||||
}
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
|
||||
process::exit(0);
|
||||
}
|
||||
@@ -396,25 +262,21 @@ async fn shutdown_signal() {
|
||||
let mut sigint = match signal::unix::signal(signal::unix::SignalKind::interrupt()) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
let _ = writeln!(std::io::stderr(), "Failed to install SIGINT handler: {}", e);
|
||||
eprintln!("Failed to install SIGINT handler: {}", e);
|
||||
process::exit(1);
|
||||
}
|
||||
};
|
||||
let mut sigterm = match signal::unix::signal(signal::unix::SignalKind::terminate()) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
let _ = writeln!(
|
||||
std::io::stderr(),
|
||||
"Failed to install SIGTERM handler: {}",
|
||||
e
|
||||
);
|
||||
eprintln!("Failed to install SIGTERM handler: {}", e);
|
||||
process::exit(1);
|
||||
}
|
||||
};
|
||||
let mut sighup = match signal::unix::signal(signal::unix::SignalKind::hangup()) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
let _ = writeln!(std::io::stderr(), "Failed to install SIGHUP handler: {}", e);
|
||||
eprintln!("Failed to install SIGHUP handler: {}", e);
|
||||
process::exit(1);
|
||||
}
|
||||
};
|
||||
@@ -429,7 +291,7 @@ async fn shutdown_signal() {
|
||||
#[cfg(windows)]
|
||||
{
|
||||
if let Err(e) = signal::ctrl_c().await {
|
||||
let _ = writeln!(std::io::stderr(), "Failed to install Ctrl+C handler: {}", e);
|
||||
eprintln!("Failed to install Ctrl+C handler: {}", e);
|
||||
process::exit(1);
|
||||
}
|
||||
}
|
||||
@@ -457,119 +319,9 @@ fn get_daemon_socket_dir() -> PathBuf {
|
||||
|
||||
#[cfg(windows)]
|
||||
fn get_port_for_session(session: &str) -> u16 {
|
||||
let mut hash: i32 = 0;
|
||||
for c in session.chars() {
|
||||
hash = ((hash << 5).wrapping_sub(hash)).wrapping_add(c as i32);
|
||||
}
|
||||
49152 + ((hash.unsigned_abs() as u32 % 16383) as u16)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[allow(unused_imports)]
|
||||
use super::*;
|
||||
|
||||
#[cfg(windows)]
|
||||
#[test]
|
||||
fn test_port_matches_client_algorithm() {
|
||||
assert_eq!(get_port_for_session("default"), 50838);
|
||||
assert_eq!(get_port_for_session("my-session"), 63105);
|
||||
assert_eq!(get_port_for_session("work"), 51184);
|
||||
assert_eq!(get_port_for_session(""), 49152);
|
||||
}
|
||||
|
||||
/// Guard against re-introducing `waitpid(-1)` in daemon code.
|
||||
///
|
||||
/// Issue #1035: a SIGCHLD handler that called `waitpid(-1, WNOHANG)` was
|
||||
/// added in v0.22.3 to reap zombie Chrome processes. This races with
|
||||
/// Rust's `Child::try_wait()` / `Child::wait()` because `waitpid(-1)`
|
||||
/// reaps *any* child, stealing the exit status before Rust can collect
|
||||
/// it. The result is ECHILD errors in `BrowserManager::has_process_exited()`
|
||||
/// and `ChromeProcess::kill()`, which can leave the daemon in a broken
|
||||
/// state or cause hangs on certain Linux configurations.
|
||||
///
|
||||
/// The fix uses the existing 500ms drain interval to call
|
||||
/// `has_process_exited()` (which delegates to `Child::try_wait()`)
|
||||
/// for targeted, race-free zombie detection.
|
||||
#[test]
|
||||
fn test_no_waitpid_minus_one_in_daemon() {
|
||||
let source = include_str!("daemon.rs");
|
||||
// Only check production code (everything before `#[cfg(test)]`)
|
||||
let production_code = source.split("#[cfg(test)]").next().unwrap_or(source);
|
||||
assert!(
|
||||
!production_code.contains("waitpid(-1"),
|
||||
"daemon.rs production code must not call waitpid(-1, ...). \
|
||||
Use Child::try_wait() via has_process_exited() instead. \
|
||||
See issue #1035."
|
||||
);
|
||||
}
|
||||
|
||||
/// Verify that `Child::try_wait()` correctly detects a crashed child
|
||||
/// without needing a global SIGCHLD handler or `waitpid(-1)`.
|
||||
/// This is what `has_process_exited()` uses in the fixed code.
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn test_child_try_wait_detects_exit_without_sigchld_handler() {
|
||||
use std::process::{Command, Stdio};
|
||||
|
||||
let mut child = Command::new("/bin/sh")
|
||||
.args(["-c", "exit 42"])
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.spawn()
|
||||
.expect("failed to spawn child");
|
||||
|
||||
std::thread::sleep(std::time::Duration::from_millis(200));
|
||||
|
||||
match child.try_wait() {
|
||||
Ok(Some(status)) => {
|
||||
assert!(
|
||||
!status.success(),
|
||||
"child exited with code 42, should not be success"
|
||||
);
|
||||
}
|
||||
Ok(None) => panic!("try_wait() returned None but child should have exited"),
|
||||
Err(e) => panic!("try_wait() should succeed without waitpid(-1): {}", e),
|
||||
}
|
||||
}
|
||||
|
||||
/// Verify that `ChromeProcess::has_exited()` (which uses `Child::try_wait()`)
|
||||
/// correctly detects a killed child, the same way the drain interval does
|
||||
/// in the fixed daemon code. This ensures crash detection works without
|
||||
/// a SIGCHLD handler.
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn test_has_exited_detects_killed_process() {
|
||||
use std::process::{Command, Stdio};
|
||||
|
||||
let mut child = Command::new("/bin/sh")
|
||||
.args(["-c", "sleep 60"])
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.spawn()
|
||||
.expect("failed to spawn child");
|
||||
|
||||
// Process should be running
|
||||
match child.try_wait() {
|
||||
Ok(None) => {} // expected
|
||||
other => panic!("expected Ok(None) for running process, got {:?}", other),
|
||||
}
|
||||
|
||||
// Kill it (simulates Chrome crash)
|
||||
child.kill().expect("failed to kill child");
|
||||
std::thread::sleep(std::time::Duration::from_millis(100));
|
||||
|
||||
// try_wait should detect the exit
|
||||
match child.try_wait() {
|
||||
Ok(Some(_)) => {} // expected: detected the crash
|
||||
other => panic!(
|
||||
"expected Ok(Some(_)) after kill, got {:?}. \
|
||||
Crash detection via try_wait() must work for the drain \
|
||||
interval fix (issue #1035) to function correctly.",
|
||||
other
|
||||
),
|
||||
}
|
||||
let mut hash: i64 = 0;
|
||||
for b in session.bytes() {
|
||||
hash = hash.wrapping_mul(31).wrapping_add(b as i64);
|
||||
}
|
||||
49152 + (hash.unsigned_abs() % 16383) as u16
|
||||
}
|
||||
|
||||
@@ -101,21 +101,6 @@ pub fn diff_screenshot(
|
||||
|
||||
/// Compute a snapshot diff using the Myers algorithm via the `similar` crate.
|
||||
pub fn diff_snapshots(before: &str, after: &str) -> SnapshotDiffResult {
|
||||
// Fast path: identical inputs.
|
||||
// This avoids constructing the `similar` TextDiff object and running the diff
|
||||
// iteration when agents compare a snapshot to itself (common in retry/loop
|
||||
// workloads).
|
||||
if before == after {
|
||||
let unchanged = before.lines().count();
|
||||
return SnapshotDiffResult {
|
||||
diff: String::new(),
|
||||
additions: 0,
|
||||
removals: 0,
|
||||
unchanged,
|
||||
changed: false,
|
||||
};
|
||||
}
|
||||
|
||||
let text_diff = TextDiff::from_lines(before, after);
|
||||
|
||||
let mut additions = 0usize;
|
||||
@@ -207,68 +192,4 @@ mod tests {
|
||||
assert_eq!(result.unchanged, 1);
|
||||
assert!(!result.diff.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_diff_snapshots_identical_fast_path() {
|
||||
let input = "hello\nworld\n";
|
||||
let result = diff_snapshots(input, input);
|
||||
assert!(!result.changed);
|
||||
assert_eq!(result.additions, 0);
|
||||
assert_eq!(result.removals, 0);
|
||||
assert_eq!(result.unchanged, input.lines().count());
|
||||
assert!(result.diff.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore]
|
||||
fn bench_diff_snapshots_identical_and_changed() {
|
||||
use std::hint::black_box;
|
||||
use std::time::Instant;
|
||||
|
||||
let identical_a = (0..200)
|
||||
.map(|i| format!("line {i}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
let identical_b = identical_a.clone();
|
||||
|
||||
let changed_a = identical_a.clone();
|
||||
let changed_b = (0..200)
|
||||
.map(|i| {
|
||||
if i == 123 {
|
||||
format!("line {i} changed")
|
||||
} else {
|
||||
format!("line {i}")
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
// Keep the iteration count high enough to measure, but low enough
|
||||
// to avoid long CI times when someone runs `--ignored`.
|
||||
let iters = 50_000usize;
|
||||
|
||||
let start = Instant::now();
|
||||
let mut acc_changed = 0usize;
|
||||
for _ in 0..iters {
|
||||
let r = diff_snapshots(black_box(&identical_a), black_box(&identical_b));
|
||||
acc_changed ^= r.unchanged;
|
||||
}
|
||||
let identical_ms = start.elapsed().as_secs_f64() * 1000.0;
|
||||
|
||||
let start = Instant::now();
|
||||
let mut acc_changed2 = 0usize;
|
||||
for _ in 0..iters {
|
||||
let r = diff_snapshots(black_box(&changed_a), black_box(&changed_b));
|
||||
acc_changed2 ^= r.additions;
|
||||
}
|
||||
let changed_ms = start.elapsed().as_secs_f64() * 1000.0;
|
||||
|
||||
// Prevent the compiler from optimizing everything away.
|
||||
black_box(acc_changed);
|
||||
black_box(acc_changed2);
|
||||
|
||||
println!(
|
||||
"bench_diff_snapshots_identical_and_changed: iters={iters} identical_ms={identical_ms:.2} changed_ms={changed_ms:.2}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+9
-2538
File diff suppressed because it is too large
Load Diff
+99
-496
@@ -12,7 +12,6 @@ pub struct RefEntry {
|
||||
pub name: String,
|
||||
pub nth: Option<usize>,
|
||||
pub selector: Option<String>,
|
||||
pub frame_id: Option<String>,
|
||||
}
|
||||
|
||||
pub struct RefMap {
|
||||
@@ -35,18 +34,6 @@ impl RefMap {
|
||||
role: &str,
|
||||
name: &str,
|
||||
nth: Option<usize>,
|
||||
) {
|
||||
self.add_with_frame(ref_id, backend_node_id, role, name, nth, None);
|
||||
}
|
||||
|
||||
pub fn add_with_frame(
|
||||
&mut self,
|
||||
ref_id: String,
|
||||
backend_node_id: Option<i64>,
|
||||
role: &str,
|
||||
name: &str,
|
||||
nth: Option<usize>,
|
||||
frame_id: Option<&str>,
|
||||
) {
|
||||
self.map.insert(
|
||||
ref_id,
|
||||
@@ -56,28 +43,6 @@ impl RefMap {
|
||||
name: name.to_string(),
|
||||
nth,
|
||||
selector: None,
|
||||
frame_id: frame_id.map(|s| s.to_string()),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
pub fn add_selector(
|
||||
&mut self,
|
||||
ref_id: String,
|
||||
selector: String,
|
||||
role: &str,
|
||||
name: &str,
|
||||
nth: Option<usize>,
|
||||
) {
|
||||
self.map.insert(
|
||||
ref_id,
|
||||
RefEntry {
|
||||
backend_node_id: None,
|
||||
role: role.to_string(),
|
||||
name: name.to_string(),
|
||||
nth,
|
||||
selector: Some(selector),
|
||||
frame_id: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -86,23 +51,6 @@ impl RefMap {
|
||||
self.map.get(ref_id)
|
||||
}
|
||||
|
||||
pub fn entries_sorted(&self) -> Vec<(String, RefEntry)> {
|
||||
let mut entries = self
|
||||
.map
|
||||
.iter()
|
||||
.map(|(ref_id, entry)| (ref_id.clone(), entry.clone()))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
entries.sort_by_key(|(ref_id, _)| {
|
||||
ref_id
|
||||
.strip_prefix('e')
|
||||
.and_then(|n| n.parse::<usize>().ok())
|
||||
.unwrap_or(usize::MAX)
|
||||
});
|
||||
|
||||
entries
|
||||
}
|
||||
|
||||
pub fn clear(&mut self) {
|
||||
self.map.clear();
|
||||
self.next_ref = 1;
|
||||
@@ -147,19 +95,14 @@ pub async fn resolve_element_center(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(f64, f64, String), String> {
|
||||
) -> Result<(f64, f64), String> {
|
||||
if let Some(ref_id) = parse_ref(selector_or_ref) {
|
||||
let entry = ref_map
|
||||
.get(&ref_id)
|
||||
.ok_or_else(|| format!("Unknown ref: {}", ref_id))?;
|
||||
|
||||
let effective_session_id =
|
||||
resolve_frame_session(entry.frame_id.as_deref(), session_id, iframe_sessions);
|
||||
|
||||
// Try cached backend_node_id first (fast path)
|
||||
if let Some(backend_node_id) = entry.backend_node_id {
|
||||
let result: Result<DomGetBoxModelResult, String> = client
|
||||
let result: DomGetBoxModelResult = client
|
||||
.send_command_typed(
|
||||
"DOM.getBoxModel",
|
||||
&DomGetBoxModelParams {
|
||||
@@ -167,46 +110,19 @@ pub async fn resolve_element_center(
|
||||
node_id: None,
|
||||
object_id: None,
|
||||
},
|
||||
Some(effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await;
|
||||
.await?;
|
||||
|
||||
if let Ok(r) = result {
|
||||
let (x, y) = box_model_center(&r.model);
|
||||
return Ok((x, y, effective_session_id.to_string()));
|
||||
}
|
||||
// backend_node_id is stale; re-query the accessibility tree below
|
||||
return Ok(box_model_center(&result.model));
|
||||
}
|
||||
|
||||
// Fallback: re-query the accessibility tree to find a fresh node by role/name
|
||||
let fresh_id = find_node_id_by_role_name(
|
||||
client,
|
||||
session_id,
|
||||
&entry.role,
|
||||
&entry.name,
|
||||
entry.nth,
|
||||
entry.frame_id.as_deref(),
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let result: DomGetBoxModelResult = client
|
||||
.send_command_typed(
|
||||
"DOM.getBoxModel",
|
||||
&DomGetBoxModelParams {
|
||||
backend_node_id: Some(fresh_id),
|
||||
node_id: None,
|
||||
object_id: None,
|
||||
},
|
||||
Some(effective_session_id),
|
||||
)
|
||||
.await?;
|
||||
let (x, y) = box_model_center(&result.model);
|
||||
return Ok((x, y, effective_session_id.to_string()));
|
||||
// Fallback: use role/name to find via JS
|
||||
return resolve_by_role_name(client, session_id, &entry.role, &entry.name, entry.nth).await;
|
||||
}
|
||||
|
||||
// CSS selector
|
||||
let (x, y) = resolve_by_selector(client, session_id, selector_or_ref).await?;
|
||||
Ok((x, y, session_id.to_string()))
|
||||
resolve_by_selector(client, session_id, selector_or_ref).await
|
||||
}
|
||||
|
||||
pub async fn resolve_element_object_id(
|
||||
@@ -214,19 +130,14 @@ pub async fn resolve_element_object_id(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(String, String), String> {
|
||||
) -> Result<String, String> {
|
||||
if let Some(ref_id) = parse_ref(selector_or_ref) {
|
||||
let entry = ref_map
|
||||
.get(&ref_id)
|
||||
.ok_or_else(|| format!("Unknown ref: {}", ref_id))?;
|
||||
|
||||
let effective_session_id =
|
||||
resolve_frame_session(entry.frame_id.as_deref(), session_id, iframe_sessions);
|
||||
|
||||
// Try cached backend_node_id first (fast path)
|
||||
if let Some(backend_node_id) = entry.backend_node_id {
|
||||
let result: Result<DomResolveNodeResult, String> = client
|
||||
let result: DomResolveNodeResult = client
|
||||
.send_command_typed(
|
||||
"DOM.resolveNode",
|
||||
&DomResolveNodeParams {
|
||||
@@ -234,49 +145,22 @@ pub async fn resolve_element_object_id(
|
||||
node_id: None,
|
||||
object_group: Some("agent-browser".to_string()),
|
||||
},
|
||||
Some(effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await;
|
||||
.await?;
|
||||
|
||||
if let Ok(r) = result {
|
||||
if let Some(object_id) = r.object.object_id {
|
||||
return Ok((object_id, effective_session_id.to_string()));
|
||||
}
|
||||
}
|
||||
// backend_node_id is stale; re-query the accessibility tree below
|
||||
return result
|
||||
.object
|
||||
.object_id
|
||||
.ok_or_else(|| format!("No objectId for ref {}", ref_id));
|
||||
}
|
||||
|
||||
// Fallback: re-query the accessibility tree to find a fresh node by role/name
|
||||
let fresh_id = find_node_id_by_role_name(
|
||||
client,
|
||||
session_id,
|
||||
&entry.role,
|
||||
&entry.name,
|
||||
entry.nth,
|
||||
entry.frame_id.as_deref(),
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let result: DomResolveNodeResult = client
|
||||
.send_command_typed(
|
||||
"DOM.resolveNode",
|
||||
&DomResolveNodeParams {
|
||||
backend_node_id: Some(fresh_id),
|
||||
node_id: None,
|
||||
object_group: Some("agent-browser".to_string()),
|
||||
},
|
||||
Some(effective_session_id),
|
||||
)
|
||||
.await?;
|
||||
let object_id = result
|
||||
.object
|
||||
.object_id
|
||||
.ok_or_else(|| format!("No objectId for ref {}", ref_id))?;
|
||||
return Ok((object_id, effective_session_id.to_string()));
|
||||
}
|
||||
|
||||
// Selector fallback (CSS or XPath)
|
||||
let js = build_find_element_js(selector_or_ref);
|
||||
// CSS selector fallback
|
||||
let js = format!(
|
||||
"document.querySelector({})",
|
||||
serde_json::to_string(selector_or_ref).unwrap_or_default()
|
||||
);
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
"Runtime.evaluate",
|
||||
@@ -289,149 +173,63 @@ pub async fn resolve_element_object_id(
|
||||
)
|
||||
.await?;
|
||||
|
||||
let object_id = result
|
||||
result
|
||||
.result
|
||||
.object_id
|
||||
.ok_or_else(|| format!("Element not found: {}", selector_or_ref))?;
|
||||
Ok((object_id, session_id.to_string()))
|
||||
.ok_or_else(|| format!("Element not found: {}", selector_or_ref))
|
||||
}
|
||||
|
||||
/// Determine which CDP session and parameters to use for an AX tree query.
|
||||
/// Cross-origin iframes have a dedicated session (no frameId needed);
|
||||
/// same-origin iframes use the parent session with a frameId parameter.
|
||||
pub(super) fn resolve_ax_session<'a>(
|
||||
frame_id: Option<&str>,
|
||||
session_id: &'a str,
|
||||
iframe_sessions: &'a HashMap<String, String>,
|
||||
) -> (serde_json::Value, &'a str) {
|
||||
if let Some(frame_id) = frame_id {
|
||||
if let Some(iframe_sid) = iframe_sessions.get(frame_id) {
|
||||
(serde_json::json!({}), iframe_sid.as_str())
|
||||
} else {
|
||||
(serde_json::json!({ "frameId": frame_id }), session_id)
|
||||
}
|
||||
} else {
|
||||
(serde_json::json!({}), session_id)
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve the effective CDP session for an element's frame.
|
||||
/// If the element's frame_id has a dedicated cross-origin iframe session, return it.
|
||||
/// Otherwise, return the parent session.
|
||||
fn resolve_frame_session<'a>(
|
||||
frame_id: Option<&str>,
|
||||
session_id: &'a str,
|
||||
iframe_sessions: &'a HashMap<String, String>,
|
||||
) -> &'a str {
|
||||
frame_id
|
||||
.and_then(|fid| iframe_sessions.get(fid))
|
||||
.map(|s| s.as_str())
|
||||
.unwrap_or(session_id)
|
||||
}
|
||||
|
||||
/// Re-query the accessibility tree to find a node matching role+name+nth,
|
||||
/// returning its fresh backendDOMNodeId. This uses the same data source
|
||||
/// (Accessibility.getFullAXTree) that built the ref map during snapshot,
|
||||
/// so role/name matching is guaranteed to be consistent.
|
||||
async fn find_node_id_by_role_name(
|
||||
async fn resolve_by_role_name(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
role: &str,
|
||||
name: &str,
|
||||
nth: Option<usize>,
|
||||
frame_id: Option<&str>,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<i64, String> {
|
||||
let (ax_params, effective_session_id) =
|
||||
resolve_ax_session(frame_id, session_id, iframe_sessions);
|
||||
let ax_tree: GetFullAXTreeResult = client
|
||||
.send_command_typed(
|
||||
"Accessibility.getFullAXTree",
|
||||
&ax_params,
|
||||
Some(effective_session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
) -> Result<(f64, f64), String> {
|
||||
let nth_index = nth.unwrap_or(0);
|
||||
let mut match_count: usize = 0;
|
||||
|
||||
for node in &ax_tree.nodes {
|
||||
if node.ignored.unwrap_or(false) {
|
||||
continue;
|
||||
}
|
||||
let node_role = extract_ax_string(&node.role);
|
||||
let node_name = extract_ax_string(&node.name);
|
||||
if node_role == role && node_name == name {
|
||||
if match_count == nth_index {
|
||||
return node.backend_d_o_m_node_id.ok_or_else(|| {
|
||||
format!(
|
||||
"AX node has no backendDOMNodeId for role={} name={}",
|
||||
role, name
|
||||
)
|
||||
});
|
||||
}
|
||||
match_count += 1;
|
||||
}
|
||||
}
|
||||
|
||||
Err(format!(
|
||||
"Could not locate element with role={} name={}",
|
||||
role, name
|
||||
))
|
||||
}
|
||||
|
||||
fn extract_ax_string(value: &Option<AXValue>) -> String {
|
||||
match value {
|
||||
Some(v) => match &v.value {
|
||||
Some(Value::String(s)) => s.clone(),
|
||||
Some(Value::Number(n)) => n.to_string(),
|
||||
Some(Value::Bool(b)) => b.to_string(),
|
||||
_ => String::new(),
|
||||
},
|
||||
None => String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a JS expression that finds a DOM element by CSS selector or XPath.
|
||||
fn build_find_element_js(selector: &str) -> String {
|
||||
if let Some(xpath) = selector.strip_prefix("xpath=") {
|
||||
format!(
|
||||
"document.evaluate({}, document, null, XPathResult.FIRST_ORDERED_NODE_TYPE, null).singleNodeValue",
|
||||
serde_json::to_string(xpath).unwrap_or_default()
|
||||
)
|
||||
} else {
|
||||
format!(
|
||||
"document.querySelector({})",
|
||||
serde_json::to_string(selector).unwrap_or_default()
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a JS expression that counts matching DOM elements by CSS selector or XPath.
|
||||
fn build_count_elements_js(selector: &str) -> String {
|
||||
if let Some(xpath) = selector.strip_prefix("xpath=") {
|
||||
format!(
|
||||
"document.evaluate({}, document, null, XPathResult.ORDERED_NODE_SNAPSHOT_TYPE, null).snapshotLength",
|
||||
serde_json::to_string(xpath).unwrap_or_default()
|
||||
)
|
||||
} else {
|
||||
format!(
|
||||
"document.querySelectorAll({}).length",
|
||||
serde_json::to_string(selector).unwrap_or_default()
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn build_selector_js(selector: &str) -> String {
|
||||
let find_expr = build_find_element_js(selector);
|
||||
format!(
|
||||
let js = format!(
|
||||
r#"(() => {{
|
||||
const el = {find_expr};
|
||||
const walker = document.createTreeWalker(document.body, NodeFilter.SHOW_ELEMENT);
|
||||
const matches = [];
|
||||
let node;
|
||||
while (node = walker.nextNode()) {{
|
||||
const r = node.getAttribute('role') || node.tagName.toLowerCase();
|
||||
const n = node.getAttribute('aria-label') || node.textContent.trim().slice(0, 100);
|
||||
if (r === {role} && n === {name}) matches.push(node);
|
||||
}}
|
||||
const el = matches[{nth}];
|
||||
if (!el) return null;
|
||||
const rect = el.getBoundingClientRect();
|
||||
return {{ x: rect.x + rect.width / 2, y: rect.y + rect.height / 2 }};
|
||||
}})()"#,
|
||||
)
|
||||
role = serde_json::to_string(role).unwrap_or_default(),
|
||||
name = serde_json::to_string(name).unwrap_or_default(),
|
||||
nth = nth_index,
|
||||
);
|
||||
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
"Runtime.evaluate",
|
||||
&EvaluateParams {
|
||||
expression: js,
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let val = result.result.value.unwrap_or(Value::Null);
|
||||
let x = val.get("x").and_then(|v| v.as_f64());
|
||||
let y = val.get("y").and_then(|v| v.as_f64());
|
||||
|
||||
match (x, y) {
|
||||
(Some(x), Some(y)) => Ok((x, y)),
|
||||
_ => Err(format!(
|
||||
"Could not locate element with role={} name={}",
|
||||
role, name
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn resolve_by_selector(
|
||||
@@ -439,7 +237,15 @@ async fn resolve_by_selector(
|
||||
session_id: &str,
|
||||
selector: &str,
|
||||
) -> Result<(f64, f64), String> {
|
||||
let js = build_selector_js(selector);
|
||||
let js = format!(
|
||||
r#"(() => {{
|
||||
const el = document.querySelector({sel});
|
||||
if (!el) return null;
|
||||
const rect = el.getBoundingClientRect();
|
||||
return {{ x: rect.x + rect.width / 2, y: rect.y + rect.height / 2 }};
|
||||
}})()"#,
|
||||
sel = serde_json::to_string(selector).unwrap_or_default(),
|
||||
);
|
||||
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
@@ -479,16 +285,8 @@ pub async fn get_element_text(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<String, String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
@@ -501,7 +299,7 @@ pub async fn get_element_text(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -518,16 +316,8 @@ pub async fn get_element_attribute(
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
attribute: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<Value, String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
@@ -542,7 +332,7 @@ pub async fn get_element_attribute(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -554,16 +344,8 @@ pub async fn is_element_visible(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<bool, String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
@@ -583,7 +365,7 @@ pub async fn is_element_visible(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -599,16 +381,8 @@ pub async fn is_element_enabled(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<bool, String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
@@ -620,7 +394,7 @@ pub async fn is_element_enabled(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -636,61 +410,20 @@ pub async fn is_element_checked(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<bool, String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
// Mirrors Playwright's getChecked() with follow-label retargeting:
|
||||
// 1. If element is a native checkbox/radio input, return .checked
|
||||
// 2. If element has an ARIA checked role, return aria-checked
|
||||
// 3. Follow label → input association (label.control)
|
||||
// 4. Check for nested checkbox/radio input as last resort
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
"Runtime.callFunctionOn",
|
||||
&CallFunctionOnParams {
|
||||
function_declaration: r#"function() {
|
||||
var el = this;
|
||||
// Native checkbox/radio input
|
||||
var tag = el.tagName && el.tagName.toUpperCase();
|
||||
if (tag === 'INPUT' && (el.type === 'checkbox' || el.type === 'radio')) {
|
||||
return el.checked;
|
||||
}
|
||||
// ARIA role-based checked state
|
||||
var role = el.getAttribute && el.getAttribute('role');
|
||||
var ariaCheckedRoles = ['checkbox','radio','switch','menuitemcheckbox','menuitemradio','option','treeitem'];
|
||||
if (role && ariaCheckedRoles.indexOf(role) !== -1) {
|
||||
return el.getAttribute('aria-checked') === 'true';
|
||||
}
|
||||
// Follow label association (Playwright follow-label retarget)
|
||||
var label = el;
|
||||
if (tag !== 'LABEL') {
|
||||
label = el.closest && el.closest('label');
|
||||
}
|
||||
if (label && label.tagName && label.tagName.toUpperCase() === 'LABEL' && label.control) {
|
||||
var ctrl = label.control;
|
||||
if (ctrl.type === 'checkbox' || ctrl.type === 'radio') {
|
||||
return ctrl.checked;
|
||||
}
|
||||
}
|
||||
// Check for nested native input
|
||||
var input = el.querySelector && el.querySelector('input[type="checkbox"], input[type="radio"]');
|
||||
if (input) return input.checked;
|
||||
return false;
|
||||
}"#.to_string(),
|
||||
function_declaration: "function() { return !!this.checked; }".to_string(),
|
||||
object_id: Some(object_id),
|
||||
arguments: None,
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -706,16 +439,8 @@ pub async fn get_element_inner_text(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<String, String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
@@ -727,7 +452,7 @@ pub async fn get_element_inner_text(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -743,16 +468,8 @@ pub async fn get_element_inner_html(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<String, String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
@@ -764,7 +481,7 @@ pub async fn get_element_inner_html(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -780,16 +497,8 @@ pub async fn get_element_input_value(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<String, String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
@@ -803,7 +512,7 @@ pub async fn get_element_input_value(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -820,16 +529,8 @@ pub async fn set_element_value(
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
value: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
let js = format!(
|
||||
"function() {{ this.value = {}; this.dispatchEvent(new Event('input', {{bubbles: true}})); this.dispatchEvent(new Event('change', {{bubbles: true}})); }}",
|
||||
@@ -846,7 +547,7 @@ pub async fn set_element_value(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -858,16 +559,8 @@ pub async fn get_element_bounding_box(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<Value, String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
@@ -883,7 +576,7 @@ pub async fn get_element_bounding_box(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -898,7 +591,10 @@ pub async fn get_element_count(
|
||||
session_id: &str,
|
||||
selector: &str,
|
||||
) -> Result<i64, String> {
|
||||
let js = build_count_elements_js(selector);
|
||||
let js = format!(
|
||||
"document.querySelectorAll({}).length",
|
||||
serde_json::to_string(selector).unwrap_or_default()
|
||||
);
|
||||
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
@@ -921,16 +617,8 @@ pub async fn get_element_styles(
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
properties: Option<Vec<String>>,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<Value, String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
let js = match properties {
|
||||
Some(props) => {
|
||||
@@ -968,7 +656,7 @@ pub async fn get_element_styles(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -1013,47 +701,6 @@ mod tests {
|
||||
assert!(map.get("e2").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_selector_js_css() {
|
||||
let js = build_selector_js("#submit-btn");
|
||||
assert!(js.contains("document.querySelector(\"#submit-btn\")"));
|
||||
assert!(!js.contains("document.evaluate"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_selector_js_xpath() {
|
||||
let js = build_selector_js("xpath=//button[@id='ok']");
|
||||
assert!(js.contains("document.evaluate(\"//button[@id='ok']\", document, null, XPathResult.FIRST_ORDERED_NODE_TYPE, null)"));
|
||||
assert!(!js.contains("document.querySelector"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_selector_js_xpath_empty() {
|
||||
let js = build_selector_js("xpath=");
|
||||
assert!(js.contains("document.evaluate"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_selector_js_not_xpath_prefix() {
|
||||
// "xpath" without "=" should be treated as CSS selector
|
||||
let js = build_selector_js("xpath//div");
|
||||
assert!(js.contains("document.querySelector"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_count_elements_js_css() {
|
||||
let js = build_count_elements_js(".item");
|
||||
assert!(js.contains("document.querySelectorAll(\".item\").length"));
|
||||
assert!(!js.contains("document.evaluate"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_count_elements_js_xpath() {
|
||||
let js = build_count_elements_js("xpath=//li");
|
||||
assert!(js.contains("document.evaluate(\"//li\", document, null, XPathResult.ORDERED_NODE_SNAPSHOT_TYPE, null).snapshotLength"));
|
||||
assert!(!js.contains("querySelectorAll"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_box_model_center() {
|
||||
let model = BoxModel {
|
||||
@@ -1068,48 +715,4 @@ mod tests {
|
||||
assert!((x - 60.0).abs() < 0.01);
|
||||
assert!((y - 40.0).abs() < 0.01);
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// resolve_frame_session tests (Issue #925)
|
||||
// Cross-origin iframe elements must resolve to the dedicated session.
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn test_cross_origin_element_uses_dedicated_session() {
|
||||
let mut iframe_sessions = HashMap::new();
|
||||
iframe_sessions.insert(
|
||||
"cross-origin-frame".to_string(),
|
||||
"iframe-session".to_string(),
|
||||
);
|
||||
|
||||
let session = resolve_frame_session(
|
||||
Some("cross-origin-frame"),
|
||||
"parent-session",
|
||||
&iframe_sessions,
|
||||
);
|
||||
|
||||
assert_eq!(session, "iframe-session");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_same_origin_element_uses_parent_session() {
|
||||
let iframe_sessions = HashMap::new();
|
||||
|
||||
let session = resolve_frame_session(
|
||||
Some("same-origin-frame"),
|
||||
"parent-session",
|
||||
&iframe_sessions,
|
||||
);
|
||||
|
||||
assert_eq!(session, "parent-session");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_main_frame_element_uses_parent_session() {
|
||||
let iframe_sessions = HashMap::new();
|
||||
|
||||
let session = resolve_frame_session(None, "parent-session", &iframe_sessions);
|
||||
|
||||
assert_eq!(session, "parent-session");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,362 +0,0 @@
|
||||
use std::io::Write;
|
||||
use std::sync::atomic::{AtomicI64, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
|
||||
use super::cdp::client::InspectProxyHandle;
|
||||
|
||||
/// Counter for unique attach IDs so concurrent connections don't collide.
|
||||
static ATTACH_ID: AtomicI64 = AtomicI64::new(-1000);
|
||||
|
||||
/// Lightweight HTTP + WebSocket server for `agent-browser inspect`.
|
||||
///
|
||||
/// Serves two purposes:
|
||||
/// - `GET /` redirects to Chrome's built-in DevTools frontend with `ws=` pointing to this server
|
||||
/// - WebSocket connections create a dedicated CDP session via `Target.attachToTarget` and proxy
|
||||
/// CDP messages through the daemon's existing browser-level connection, injecting/stripping
|
||||
/// `sessionId` so the DevTools frontend sees a page-level view
|
||||
pub struct InspectServer {
|
||||
port: u16,
|
||||
_handle: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl InspectServer {
|
||||
/// Start the inspect proxy server.
|
||||
///
|
||||
/// - `proxy_handle`: lightweight handle for sending/receiving raw CDP messages
|
||||
/// - `target_id`: the CDP target ID of the page to inspect
|
||||
/// - `chrome_host_port`: the Chrome debug server address (e.g. "127.0.0.1:9222")
|
||||
pub async fn start(
|
||||
proxy_handle: InspectProxyHandle,
|
||||
target_id: String,
|
||||
chrome_host_port: String,
|
||||
) -> Result<Self, String> {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.map_err(|e| format!("Failed to bind inspect server: {}", e))?;
|
||||
let port = listener
|
||||
.local_addr()
|
||||
.map_err(|e| format!("Failed to get local addr: {}", e))?
|
||||
.port();
|
||||
|
||||
let proxy = Arc::new(proxy_handle);
|
||||
|
||||
let handle = tokio::spawn(accept_loop(
|
||||
listener,
|
||||
proxy,
|
||||
target_id,
|
||||
chrome_host_port,
|
||||
port,
|
||||
));
|
||||
|
||||
Ok(Self {
|
||||
port,
|
||||
_handle: handle,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn port(&self) -> u16 {
|
||||
self.port
|
||||
}
|
||||
|
||||
pub fn shutdown(self) {
|
||||
self._handle.abort();
|
||||
}
|
||||
}
|
||||
|
||||
async fn accept_loop(
|
||||
listener: TcpListener,
|
||||
proxy: Arc<InspectProxyHandle>,
|
||||
target_id: String,
|
||||
chrome_host_port: String,
|
||||
proxy_port: u16,
|
||||
) {
|
||||
loop {
|
||||
let (stream, _) = match listener.accept().await {
|
||||
Ok(s) => s,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
let proxy = proxy.clone();
|
||||
let tid = target_id.clone();
|
||||
let chp = chrome_host_port.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = handle_connection(stream, proxy, tid, chp, proxy_port).await {
|
||||
let _ = writeln!(std::io::stderr(), "[inspect] connection error: {}", e);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_connection(
|
||||
stream: tokio::net::TcpStream,
|
||||
proxy: Arc<InspectProxyHandle>,
|
||||
target_id: String,
|
||||
chrome_host_port: String,
|
||||
proxy_port: u16,
|
||||
) -> Result<(), String> {
|
||||
// Peek at the request line to determine routing WITHOUT consuming bytes.
|
||||
// This is critical: tokio_tungstenite::accept_async needs to read the full
|
||||
// HTTP upgrade request itself, so we must not consume anything for WS paths.
|
||||
let mut peek_buf = [0u8; 32];
|
||||
let n = stream
|
||||
.peek(&mut peek_buf)
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
let peek = String::from_utf8_lossy(&peek_buf[..n]);
|
||||
|
||||
if peek.starts_with("GET /ws") {
|
||||
return handle_ws_proxy(stream, proxy, target_id).await;
|
||||
}
|
||||
|
||||
if peek.starts_with("GET / ") {
|
||||
let buf_reader = BufReader::new(stream);
|
||||
return handle_http_redirect(buf_reader, chrome_host_port, proxy_port).await;
|
||||
}
|
||||
|
||||
// Unknown request -- consume and respond 404
|
||||
let mut stream = stream;
|
||||
let mut discard = [0u8; 4096];
|
||||
let _ = stream.read(&mut discard).await;
|
||||
let resp = "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
|
||||
stream
|
||||
.write_all(resp.as_bytes())
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
const MAX_HEADER_BYTES: usize = 8192;
|
||||
|
||||
async fn handle_http_redirect(
|
||||
buf_reader: BufReader<tokio::net::TcpStream>,
|
||||
chrome_host_port: String,
|
||||
proxy_port: u16,
|
||||
) -> Result<(), String> {
|
||||
let mut br = buf_reader;
|
||||
let mut total_bytes = 0usize;
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
let n = br.read_line(&mut line).await.map_err(|e| e.to_string())?;
|
||||
total_bytes += n;
|
||||
if line == "\r\n" || line == "\n" || line.is_empty() || total_bytes > MAX_HEADER_BYTES {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let location = format!(
|
||||
"http://{}/devtools/devtools_app.html?ws=127.0.0.1:{}/ws",
|
||||
chrome_host_port, proxy_port
|
||||
);
|
||||
let body = format!(
|
||||
"<html><body>Redirecting to <a href=\"{url}\">{url}</a></body></html>",
|
||||
url = location
|
||||
);
|
||||
let resp = format!(
|
||||
"HTTP/1.1 302 Found\r\nLocation: {}\r\nContent-Type: text/html\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
|
||||
location,
|
||||
body.len(),
|
||||
body
|
||||
);
|
||||
let mut stream = br.into_inner();
|
||||
stream
|
||||
.write_all(resp.as_bytes())
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_ws_proxy(
|
||||
stream: tokio::net::TcpStream,
|
||||
proxy: Arc<InspectProxyHandle>,
|
||||
target_id: String,
|
||||
) -> Result<(), String> {
|
||||
let ws_stream = tokio_tungstenite::accept_async(stream)
|
||||
.await
|
||||
.map_err(|e| format!("WebSocket handshake failed: {}", e))?;
|
||||
|
||||
// Create a dedicated CDP session for this DevTools connection.
|
||||
// Each connection gets its own session so domain enablements (DOM.enable, etc.)
|
||||
// always trigger fresh initial state dumps from Chrome.
|
||||
let attach_id = ATTACH_ID.fetch_sub(1, Ordering::SeqCst);
|
||||
let attach_cmd = format!(
|
||||
r#"{{"id":{},"method":"Target.attachToTarget","params":{{"targetId":"{}","flatten":true}}}}"#,
|
||||
attach_id, target_id
|
||||
);
|
||||
|
||||
// Subscribe BEFORE sending so we don't miss the response (tokio broadcast
|
||||
// receivers only deliver messages to receivers that already exist).
|
||||
let mut raw_rx = proxy.subscribe_raw();
|
||||
|
||||
proxy
|
||||
.send_raw(attach_cmd)
|
||||
.await
|
||||
.map_err(|e| format!("Failed to send attachToTarget: {}", e))?;
|
||||
|
||||
// Wait for the attachToTarget response to extract the session ID
|
||||
let session_id = tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
||||
while let Ok(raw_msg) = raw_rx.recv().await {
|
||||
if let Ok(val) = serde_json::from_str::<serde_json::Value>(&raw_msg.text) {
|
||||
if val.get("id").and_then(|v| v.as_i64()) == Some(attach_id) {
|
||||
if let Some(sid) = val
|
||||
.get("result")
|
||||
.and_then(|r| r.get("sessionId"))
|
||||
.and_then(|s| s.as_str())
|
||||
{
|
||||
return Ok(sid.to_string());
|
||||
}
|
||||
return Err("attachToTarget failed".to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
Err("raw message channel closed".to_string())
|
||||
})
|
||||
.await
|
||||
.map_err(|_| "Timed out waiting for attachToTarget response".to_string())?
|
||||
.map_err(|e| format!("Failed to create DevTools session: {}", e))?;
|
||||
|
||||
let (ws_tx, mut ws_rx) = ws_stream.split();
|
||||
let ws_tx = Arc::new(Mutex::new(ws_tx));
|
||||
|
||||
let mut raw_rx = proxy.subscribe_raw();
|
||||
let ws_tx_clone = ws_tx.clone();
|
||||
let session_id_clone = session_id.clone();
|
||||
|
||||
// Chrome -> DevTools: forward messages matching our session, strip sessionId
|
||||
let mut chrome_to_devtools = tokio::spawn(async move {
|
||||
loop {
|
||||
let raw_msg = match raw_rx.recv().await {
|
||||
Ok(msg) => msg,
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
|
||||
let _ = writeln!(
|
||||
std::io::stderr(),
|
||||
"[inspect] warning: dropped {} CDP messages (channel lag)",
|
||||
n
|
||||
);
|
||||
continue;
|
||||
}
|
||||
Err(_) => break,
|
||||
};
|
||||
|
||||
if raw_msg.session_id.as_deref() != Some(&session_id_clone) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let stripped = strip_session_id(&raw_msg.text);
|
||||
|
||||
let mut tx = ws_tx_clone.lock().await;
|
||||
if tx.send(Message::Text(stripped)).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// DevTools -> Chrome: inject sessionId and forward
|
||||
let proxy_for_send = proxy.clone();
|
||||
let session_id_for_send = session_id.clone();
|
||||
let mut devtools_to_chrome = tokio::spawn(async move {
|
||||
while let Some(Ok(msg)) = ws_rx.next().await {
|
||||
let text = match msg {
|
||||
Message::Text(t) => t,
|
||||
Message::Close(_) => break,
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
let injected = inject_session_id(&text, &session_id_for_send);
|
||||
if proxy_for_send.send_raw(injected).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
tokio::select! {
|
||||
_ = &mut chrome_to_devtools => {
|
||||
devtools_to_chrome.abort();
|
||||
},
|
||||
_ = &mut devtools_to_chrome => {
|
||||
chrome_to_devtools.abort();
|
||||
},
|
||||
}
|
||||
|
||||
// Clean up the CDP session so Chrome doesn't leak attached targets
|
||||
let detach_cmd = format!(
|
||||
r#"{{"id":{},"method":"Target.detachFromTarget","params":{{"sessionId":"{}"}}}}"#,
|
||||
ATTACH_ID.fetch_sub(1, Ordering::SeqCst),
|
||||
session_id
|
||||
);
|
||||
let _ = proxy.send_raw(detach_cmd).await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn inject_session_id(json: &str, session_id: &str) -> String {
|
||||
if let Ok(mut val) = serde_json::from_str::<serde_json::Value>(json) {
|
||||
if let Some(obj) = val.as_object_mut() {
|
||||
obj.insert(
|
||||
"sessionId".to_string(),
|
||||
serde_json::Value::String(session_id.to_string()),
|
||||
);
|
||||
}
|
||||
serde_json::to_string(&val).unwrap_or_else(|_| json.to_string())
|
||||
} else {
|
||||
json.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn strip_session_id(json: &str) -> String {
|
||||
if let Ok(mut val) = serde_json::from_str::<serde_json::Value>(json) {
|
||||
if let Some(obj) = val.as_object_mut() {
|
||||
obj.remove("sessionId");
|
||||
}
|
||||
serde_json::to_string(&val).unwrap_or_else(|_| json.to_string())
|
||||
} else {
|
||||
json.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_inject_session_id() {
|
||||
let input = r#"{"id":1,"method":"DOM.getDocument"}"#;
|
||||
let result = inject_session_id(input, "abc123");
|
||||
let parsed: serde_json::Value = serde_json::from_str(&result).expect("valid JSON");
|
||||
assert_eq!(parsed["sessionId"], "abc123");
|
||||
assert_eq!(parsed["method"], "DOM.getDocument");
|
||||
assert_eq!(parsed["id"], 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_inject_session_id_empty_object() {
|
||||
let result = inject_session_id("{}", "abc");
|
||||
let parsed: serde_json::Value = serde_json::from_str(&result).expect("valid JSON");
|
||||
assert_eq!(parsed["sessionId"], "abc");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_session_id() {
|
||||
let input = r#"{"id":1,"result":{},"sessionId":"abc123"}"#;
|
||||
let result = strip_session_id(input);
|
||||
let parsed: serde_json::Value = serde_json::from_str(&result).expect("valid JSON");
|
||||
assert!(parsed.get("sessionId").is_none());
|
||||
assert_eq!(parsed["id"], 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_inject_then_strip_roundtrip() {
|
||||
let input = r#"{"id":42,"method":"Runtime.evaluate"}"#;
|
||||
let injected = inject_session_id(input, "sess1");
|
||||
let stripped = strip_session_id(&injected);
|
||||
let original: serde_json::Value = serde_json::from_str(input).unwrap();
|
||||
let result: serde_json::Value = serde_json::from_str(&stripped).unwrap();
|
||||
assert_eq!(original, result);
|
||||
}
|
||||
}
|
||||
+83
-559
@@ -1,5 +1,3 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use super::cdp::client::CdpClient;
|
||||
@@ -13,17 +11,9 @@ pub async fn click(
|
||||
selector_or_ref: &str,
|
||||
button: &str,
|
||||
click_count: i32,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let (x, y, effective_session_id) = resolve_element_center(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
dispatch_click(client, &effective_session_id, x, y, button, click_count).await
|
||||
let (x, y) = resolve_element_center(client, session_id, ref_map, selector_or_ref).await?;
|
||||
dispatch_click(client, session_id, x, y, button, click_count).await
|
||||
}
|
||||
|
||||
pub async fn dblclick(
|
||||
@@ -31,18 +21,8 @@ pub async fn dblclick(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
click(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
"left",
|
||||
2,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await
|
||||
click(client, session_id, ref_map, selector_or_ref, "left", 2).await
|
||||
}
|
||||
|
||||
pub async fn hover(
|
||||
@@ -50,16 +30,8 @@ pub async fn hover(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let (x, y, effective_session_id) = resolve_element_center(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let (x, y) = resolve_element_center(client, session_id, ref_map, selector_or_ref).await?;
|
||||
client
|
||||
.send_command_typed::<_, Value>(
|
||||
"Input.dispatchMouseEvent",
|
||||
@@ -74,7 +46,7 @@ pub async fn hover(
|
||||
delta_y: None,
|
||||
modifiers: None,
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
@@ -86,16 +58,8 @@ pub async fn fill(
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
value: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
// Focus the element
|
||||
client
|
||||
@@ -108,7 +72,7 @@ pub async fn fill(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -128,11 +92,11 @@ pub async fn fill(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Insert text (keyboard input dispatched at page level, use parent session_id)
|
||||
// Insert text
|
||||
client
|
||||
.send_command_typed::<_, Value>(
|
||||
"Input.insertText",
|
||||
@@ -146,7 +110,6 @@ pub async fn fill(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn type_text(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
@@ -155,16 +118,8 @@ pub async fn type_text(
|
||||
text: &str,
|
||||
clear: bool,
|
||||
delay_ms: Option<u64>,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
// Focus
|
||||
client
|
||||
@@ -177,7 +132,7 @@ pub async fn type_text(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -197,73 +152,50 @@ pub async fn type_text(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
|
||||
type_text_into_active_context(client, session_id, text, delay_ms).await
|
||||
}
|
||||
|
||||
pub async fn type_text_into_active_context(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
text: &str,
|
||||
delay_ms: Option<u64>,
|
||||
) -> Result<(), String> {
|
||||
let delay = delay_ms.unwrap_or(0);
|
||||
|
||||
for ch in text.chars() {
|
||||
if matches!(ch, '\n' | '\r' | '\t') {
|
||||
let (key, code, key_code) = char_to_key_info(ch);
|
||||
let text_str = key_text(&key);
|
||||
client
|
||||
.send_command_typed::<_, Value>(
|
||||
"Input.dispatchKeyEvent",
|
||||
&DispatchKeyEventParams {
|
||||
event_type: "keyDown".to_string(),
|
||||
key: Some(key.clone()),
|
||||
code: Some(code.clone()),
|
||||
text: text_str.clone(),
|
||||
unmodified_text: text_str,
|
||||
windows_virtual_key_code: Some(key_code),
|
||||
native_virtual_key_code: Some(key_code),
|
||||
modifiers: None,
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
let text_str = ch.to_string();
|
||||
let (key, code, key_code) = char_to_key_info(ch);
|
||||
|
||||
client
|
||||
.send_command_typed::<_, Value>(
|
||||
"Input.dispatchKeyEvent",
|
||||
&DispatchKeyEventParams {
|
||||
event_type: "keyUp".to_string(),
|
||||
key: Some(key),
|
||||
code: Some(code),
|
||||
text: None,
|
||||
unmodified_text: None,
|
||||
windows_virtual_key_code: Some(key_code),
|
||||
native_virtual_key_code: Some(key_code),
|
||||
modifiers: None,
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
} else {
|
||||
// VS Code/Electron webviews reject repeated dispatchKeyEvent calls
|
||||
// carrying printable `text`. Insert printable characters directly
|
||||
// and reserve key events for controls like Enter and Tab.
|
||||
client
|
||||
.send_command_typed::<_, Value>(
|
||||
"Input.insertText",
|
||||
&InsertTextParams {
|
||||
text: ch.to_string(),
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
client
|
||||
.send_command_typed::<_, Value>(
|
||||
"Input.dispatchKeyEvent",
|
||||
&DispatchKeyEventParams {
|
||||
event_type: "keyDown".to_string(),
|
||||
key: Some(key.clone()),
|
||||
code: Some(code.clone()),
|
||||
text: Some(text_str.clone()),
|
||||
unmodified_text: Some(text_str.clone()),
|
||||
windows_virtual_key_code: Some(key_code),
|
||||
native_virtual_key_code: Some(key_code),
|
||||
modifiers: None,
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
client
|
||||
.send_command_typed::<_, Value>(
|
||||
"Input.dispatchKeyEvent",
|
||||
&DispatchKeyEventParams {
|
||||
event_type: "keyUp".to_string(),
|
||||
key: Some(key),
|
||||
code: Some(code),
|
||||
text: None,
|
||||
unmodified_text: None,
|
||||
windows_virtual_key_code: Some(key_code),
|
||||
native_virtual_key_code: Some(key_code),
|
||||
modifiers: None,
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
if delay > 0 {
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(delay)).await;
|
||||
@@ -274,33 +206,8 @@ pub async fn type_text_into_active_context(
|
||||
}
|
||||
|
||||
pub async fn press_key(client: &CdpClient, session_id: &str, key: &str) -> Result<(), String> {
|
||||
press_key_with_modifiers(client, session_id, key, None).await
|
||||
}
|
||||
|
||||
/// Dispatch a keyDown+keyUp sequence for `key` with an optional CDP modifier bitmask.
|
||||
///
|
||||
/// Modifier values follow the CDP `Input.dispatchKeyEvent` spec:
|
||||
/// 1 = Alt, 2 = Control, 4 = Meta (Cmd), 8 = Shift.
|
||||
///
|
||||
/// Callers that need a platform-appropriate modifier (e.g. Cmd on macOS,
|
||||
/// Ctrl elsewhere) must choose the value themselves -- see `cfg!(target_os)`.
|
||||
pub async fn press_key_with_modifiers(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
key: &str,
|
||||
modifiers: Option<i32>,
|
||||
) -> Result<(), String> {
|
||||
let (key_name, code, key_code) = named_key_info(key);
|
||||
|
||||
// Suppress text insertion when Control (2) or Meta (4) modifiers are active,
|
||||
// since these are command chords (e.g. Ctrl+A = select-all), not text input.
|
||||
let has_command_modifier = modifiers.is_some_and(|m| m & (2 | 4) != 0);
|
||||
let text = if has_command_modifier {
|
||||
None
|
||||
} else {
|
||||
key_text(&key_name)
|
||||
};
|
||||
|
||||
client
|
||||
.send_command_typed::<_, Value>(
|
||||
"Input.dispatchKeyEvent",
|
||||
@@ -308,11 +215,11 @@ pub async fn press_key_with_modifiers(
|
||||
event_type: "keyDown".to_string(),
|
||||
key: Some(key_name.clone()),
|
||||
code: Some(code.clone()),
|
||||
text: text.clone(),
|
||||
unmodified_text: text.clone(),
|
||||
text: None,
|
||||
unmodified_text: None,
|
||||
windows_virtual_key_code: Some(key_code),
|
||||
native_virtual_key_code: Some(key_code),
|
||||
modifiers,
|
||||
modifiers: None,
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
@@ -329,7 +236,7 @@ pub async fn press_key_with_modifiers(
|
||||
unmodified_text: None,
|
||||
windows_virtual_key_code: Some(key_code),
|
||||
native_virtual_key_code: Some(key_code),
|
||||
modifiers,
|
||||
modifiers: None,
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
@@ -345,11 +252,9 @@ pub async fn scroll(
|
||||
selector_or_ref: Option<&str>,
|
||||
delta_x: f64,
|
||||
delta_y: f64,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
if let Some(sel) = selector_or_ref {
|
||||
let (object_id, effective_session_id) =
|
||||
resolve_element_object_id(client, session_id, ref_map, sel, iframe_sessions).await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, sel).await?;
|
||||
let js = "function(dx, dy) { this.scrollBy(dx, dy); }".to_string();
|
||||
client
|
||||
.send_command_typed::<_, Value>(
|
||||
@@ -370,7 +275,7 @@ pub async fn scroll(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
} else {
|
||||
@@ -396,16 +301,8 @@ pub async fn select_option(
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
values: &[String],
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
let js = r#"function(vals) {
|
||||
const options = Array.from(this.options);
|
||||
@@ -429,7 +326,7 @@ pub async fn select_option(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -441,49 +338,11 @@ pub async fn check(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let is_checked = super::element::is_element_checked(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let is_checked =
|
||||
super::element::is_element_checked(client, session_id, ref_map, selector_or_ref).await?;
|
||||
if !is_checked {
|
||||
click(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
"left",
|
||||
1,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Verify the click changed the state (Playwright parity: _setChecked re-checks).
|
||||
// If the coordinate-based click missed (e.g. hidden input, overlay), retry
|
||||
// with a JS .click() on the element and its associated input.
|
||||
if !super::element::is_element_checked(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
js_click_checkbox(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
click(client, session_id, ref_map, selector_or_ref, "left", 1).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -493,132 +352,22 @@ pub async fn uncheck(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let is_checked = super::element::is_element_checked(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let is_checked =
|
||||
super::element::is_element_checked(client, session_id, ref_map, selector_or_ref).await?;
|
||||
if is_checked {
|
||||
click(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
"left",
|
||||
1,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Same verify-and-retry as check().
|
||||
if super::element::is_element_checked(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
js_click_checkbox(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
click(client, session_id, ref_map, selector_or_ref, "left", 1).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Fallback for when the coordinate-based CDP click did not toggle the
|
||||
/// checkbox/radio state. This mirrors how Playwright dispatches clicks
|
||||
/// through the DOM rather than via raw Input.dispatchMouseEvent coordinates.
|
||||
///
|
||||
/// Uses the same follow-label resolution as `is_element_checked`:
|
||||
/// 1. If the element is a native input → `.click()` it directly.
|
||||
/// 2. If the element is inside a `<label>` → `.click()` the label's `.control`.
|
||||
/// 3. If the element has a nested `<input>` → `.click()` that input.
|
||||
/// 4. Otherwise → `.click()` the element itself (handles ARIA role controls).
|
||||
async fn js_click_checkbox(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let js = r#"function() {
|
||||
var el = this;
|
||||
var tag = el.tagName && el.tagName.toUpperCase();
|
||||
// 1. Native input — click it directly
|
||||
if (tag === 'INPUT' && (el.type === 'checkbox' || el.type === 'radio')) {
|
||||
el.click();
|
||||
return;
|
||||
}
|
||||
// 2. Follow label → control association
|
||||
var label = tag === 'LABEL' ? el : (el.closest && el.closest('label'));
|
||||
if (label && label.tagName && label.tagName.toUpperCase() === 'LABEL' && label.control) {
|
||||
label.control.click();
|
||||
return;
|
||||
}
|
||||
// 3. Nested native input
|
||||
var input = el.querySelector && el.querySelector('input[type="checkbox"], input[type="radio"]');
|
||||
if (input) {
|
||||
input.click();
|
||||
return;
|
||||
}
|
||||
// 4. ARIA role control — click the element itself
|
||||
el.click();
|
||||
}"#;
|
||||
|
||||
client
|
||||
.send_command_typed::<_, Value>(
|
||||
"Runtime.callFunctionOn",
|
||||
&CallFunctionOnParams {
|
||||
function_declaration: js.to_string(),
|
||||
object_id: Some(object_id),
|
||||
arguments: None,
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn focus(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
client
|
||||
.send_command_typed::<_, Value>(
|
||||
@@ -630,7 +379,7 @@ pub async fn focus(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -642,16 +391,8 @@ pub async fn clear(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
client
|
||||
.send_command_typed::<_, Value>(
|
||||
@@ -669,7 +410,7 @@ pub async fn clear(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -681,16 +422,8 @@ pub async fn select_all(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
client
|
||||
.send_command_typed::<_, Value>(
|
||||
@@ -714,7 +447,7 @@ pub async fn select_all(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -726,16 +459,8 @@ pub async fn scroll_into_view(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
client
|
||||
.send_command_typed::<_, Value>(
|
||||
@@ -749,7 +474,7 @@ pub async fn scroll_into_view(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -763,16 +488,8 @@ pub async fn dispatch_event(
|
||||
selector_or_ref: &str,
|
||||
event_type: &str,
|
||||
event_init: Option<&Value>,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
let init_json = event_init
|
||||
.map(|v| serde_json::to_string(v).unwrap_or("{}".to_string()))
|
||||
@@ -794,7 +511,7 @@ pub async fn dispatch_event(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -806,16 +523,8 @@ pub async fn highlight(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let (object_id, effective_session_id) = resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let object_id = resolve_element_object_id(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
client
|
||||
.send_command_typed::<_, Value>(
|
||||
@@ -836,7 +545,7 @@ pub async fn highlight(
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -848,16 +557,8 @@ pub async fn tap_touch(
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector_or_ref: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<(), String> {
|
||||
let (x, y, effective_session_id) = resolve_element_center(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector_or_ref,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
let (x, y) = resolve_element_center(client, session_id, ref_map, selector_or_ref).await?;
|
||||
|
||||
client
|
||||
.send_command(
|
||||
@@ -866,7 +567,7 @@ pub async fn tap_touch(
|
||||
"type": "touchStart",
|
||||
"touchPoints": [{ "x": x, "y": y }],
|
||||
})),
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -877,7 +578,7 @@ pub async fn tap_touch(
|
||||
"type": "touchEnd",
|
||||
"touchPoints": [],
|
||||
})),
|
||||
Some(&effective_session_id),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -965,75 +666,15 @@ fn char_to_key_info(ch: char) -> (String, String, i32) {
|
||||
' ' => (" ".to_string(), "Space".to_string(), 32),
|
||||
_ => {
|
||||
let key = ch.to_string();
|
||||
if ch.is_ascii_alphabetic() {
|
||||
// For letters the Windows VK code equals the uppercase ASCII value.
|
||||
let upper = ch.to_ascii_uppercase();
|
||||
let code = format!("Key{}", upper);
|
||||
let key_code = upper as i32;
|
||||
(key, code, key_code)
|
||||
let code = if ch.is_ascii_alphabetic() {
|
||||
format!("Key{}", ch.to_uppercase())
|
||||
} else if ch.is_ascii_digit() {
|
||||
let code = format!("Digit{}", ch);
|
||||
let key_code = ch as i32;
|
||||
(key, code, key_code)
|
||||
format!("Digit{}", ch)
|
||||
} else {
|
||||
let (code, key_code) = punctuation_key_info(ch);
|
||||
(key, code.to_string(), key_code)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Return the DOM `KeyboardEvent.code` value and Windows virtual-key code for
|
||||
/// a punctuation / symbol character assuming a US keyboard layout.
|
||||
///
|
||||
/// The Windows virtual-key codes (VK_OEM_*) differ from ASCII values for
|
||||
/// punctuation. Using the raw ASCII code would misidentify characters – e.g.
|
||||
/// '.' (ASCII 46) collides with VK_DELETE (0x2E = 46), causing the period to
|
||||
/// be swallowed.
|
||||
fn punctuation_key_info(ch: char) -> (&'static str, i32) {
|
||||
match ch {
|
||||
// VK_OEM_1 (0xBA = 186) — ";:" key on US layout
|
||||
';' | ':' => ("Semicolon", 186),
|
||||
// VK_OEM_PLUS (0xBB = 187) — "=+" key
|
||||
'=' | '+' => ("Equal", 187),
|
||||
// VK_OEM_COMMA (0xBC = 188) — ",<" key
|
||||
',' | '<' => ("Comma", 188),
|
||||
// VK_OEM_MINUS (0xBD = 189) — "-_" key
|
||||
'-' | '_' => ("Minus", 189),
|
||||
// VK_OEM_PERIOD (0xBE = 190) — ".>" key
|
||||
'.' | '>' => ("Period", 190),
|
||||
// VK_OEM_2 (0xBF = 191) — "/?" key
|
||||
'/' | '?' => ("Slash", 191),
|
||||
// VK_OEM_3 (0xC0 = 192) — "`~" key
|
||||
'`' | '~' => ("Backquote", 192),
|
||||
// VK_OEM_4 (0xDB = 219) — "[{" key
|
||||
'[' | '{' => ("BracketLeft", 219),
|
||||
// VK_OEM_5 (0xDC = 220) — "\\|" key
|
||||
'\\' | '|' => ("Backslash", 220),
|
||||
// VK_OEM_6 (0xDD = 221) — "]}" key
|
||||
']' | '}' => ("BracketRight", 221),
|
||||
// VK_OEM_7 (0xDE = 222) — "'\""" key
|
||||
'\'' | '"' => ("Quote", 222),
|
||||
_ => ("", 0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Return the `text` value that CDP `Input.dispatchKeyEvent` needs on the
|
||||
/// `keyDown` event so that Chrome performs the default action for the key.
|
||||
/// For example Enter needs `"\r"` to actually submit a form, and Tab needs
|
||||
/// `"\t"` to move focus. Non-printable / navigation keys return `None`.
|
||||
fn key_text(key_name: &str) -> Option<String> {
|
||||
match key_name {
|
||||
"Enter" => Some("\r".to_string()),
|
||||
"Tab" => Some("\t".to_string()),
|
||||
" " => Some(" ".to_string()),
|
||||
_ => {
|
||||
// Single printable characters carry themselves as text.
|
||||
if key_name.len() == 1 {
|
||||
Some(key_name.to_string())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
String::new()
|
||||
};
|
||||
let key_code = ch as i32;
|
||||
(key, code, key_code)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1064,120 +705,3 @@ fn named_key_info(key: &str) -> (String, String, i32) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Verify that `char_to_key_info` returns the correct (key, code,
|
||||
/// windowsVirtualKeyCode) triple for every character in Playwright's
|
||||
/// USKeyboardLayout. The expected values below are taken verbatim from
|
||||
/// playwright-core/lib/server/usKeyboardLayout.js so that any drift from
|
||||
/// Playwright's behaviour is caught immediately.
|
||||
#[test]
|
||||
fn test_char_to_key_info_matches_playwright_layout() {
|
||||
// (character, expected_code, expected_vk_code)
|
||||
let cases: &[(char, &str, i32)] = &[
|
||||
// Letters – VK code must equal the uppercase ASCII value.
|
||||
('a', "KeyA", 65),
|
||||
('z', "KeyZ", 90),
|
||||
('A', "KeyA", 65),
|
||||
// Digits
|
||||
('0', "Digit0", 48),
|
||||
('9', "Digit9", 57),
|
||||
// Punctuation – these are the values from Playwright's layout.
|
||||
// The bug that prompted this test sent '.' as VK 46 (= VK_DELETE).
|
||||
('.', "Period", 190),
|
||||
(',', "Comma", 188),
|
||||
('/', "Slash", 191),
|
||||
(';', "Semicolon", 186),
|
||||
('\'', "Quote", 222),
|
||||
('[', "BracketLeft", 219),
|
||||
(']', "BracketRight", 221),
|
||||
('\\', "Backslash", 220),
|
||||
('`', "Backquote", 192),
|
||||
('-', "Minus", 189),
|
||||
('=', "Equal", 187),
|
||||
// Shifted variants produced by the same physical keys.
|
||||
('>', "Period", 190),
|
||||
('<', "Comma", 188),
|
||||
('?', "Slash", 191),
|
||||
(':', "Semicolon", 186),
|
||||
('"', "Quote", 222),
|
||||
('{', "BracketLeft", 219),
|
||||
('}', "BracketRight", 221),
|
||||
('|', "Backslash", 220),
|
||||
('~', "Backquote", 192),
|
||||
('_', "Minus", 189),
|
||||
('+', "Equal", 187),
|
||||
// Whitespace / control
|
||||
(' ', "Space", 32),
|
||||
('\n', "Enter", 13),
|
||||
('\t', "Tab", 9),
|
||||
];
|
||||
|
||||
for &(ch, expected_code, expected_vk) in cases {
|
||||
let (key, code, vk) = char_to_key_info(ch);
|
||||
assert_eq!(
|
||||
code, expected_code,
|
||||
"char {:?}: expected code {:?}, got {:?}",
|
||||
ch, expected_code, code
|
||||
);
|
||||
assert_eq!(
|
||||
vk, expected_vk,
|
||||
"char {:?}: expected VK {}, got {} (ASCII would be {})",
|
||||
ch, expected_vk, vk, ch as i32
|
||||
);
|
||||
// key should be the character itself (except control chars).
|
||||
if !ch.is_control() {
|
||||
assert_eq!(key, ch.to_string(), "char {:?}: key mismatch", ch);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Regression test: period must NEVER map to VK 46 (VK_DELETE).
|
||||
#[test]
|
||||
fn test_period_is_not_vk_delete() {
|
||||
let (_, _, vk) = char_to_key_info('.');
|
||||
assert_ne!(
|
||||
vk, 46,
|
||||
"Period must not use VK code 46 (VK_DELETE); expected 190 (VK_OEM_PERIOD)"
|
||||
);
|
||||
assert_eq!(vk, 190);
|
||||
}
|
||||
|
||||
/// Characters outside the US keyboard layout should return (key, "", 0)
|
||||
/// so that `type_text` falls back to `Input.insertText`.
|
||||
#[test]
|
||||
fn test_unmapped_chars_return_zero_keycode() {
|
||||
for ch in ['@', '#', '$', '%', '^', '&', '*', '(', ')', '€', '£', '你'] {
|
||||
let (key, code, vk) = char_to_key_info(ch);
|
||||
assert_eq!(
|
||||
code, "",
|
||||
"char {:?}: unmapped char should have empty code, got {:?}",
|
||||
ch, code
|
||||
);
|
||||
assert_eq!(
|
||||
vk, 0,
|
||||
"char {:?}: unmapped char should have VK 0, got {}",
|
||||
ch, vk
|
||||
);
|
||||
assert_eq!(key, ch.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_key_text_returns_correct_text_for_special_keys() {
|
||||
assert_eq!(key_text("Enter"), Some("\r".to_string()));
|
||||
assert_eq!(key_text("Tab"), Some("\t".to_string()));
|
||||
assert_eq!(key_text(" "), Some(" ".to_string()));
|
||||
// Single printable characters carry themselves.
|
||||
assert_eq!(key_text("a"), Some("a".to_string()));
|
||||
assert_eq!(key_text("Z"), Some("Z".to_string()));
|
||||
// Non-printable named keys return None.
|
||||
assert_eq!(key_text("Escape"), None);
|
||||
assert_eq!(key_text("ArrowUp"), None);
|
||||
assert_eq!(key_text("Backspace"), None);
|
||||
assert_eq!(key_text("Delete"), None);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -15,8 +15,6 @@ pub mod diff;
|
||||
#[allow(dead_code)]
|
||||
pub mod element;
|
||||
#[allow(dead_code)]
|
||||
pub mod inspect_server;
|
||||
#[allow(dead_code)]
|
||||
pub mod interaction;
|
||||
#[allow(dead_code)]
|
||||
pub mod network;
|
||||
@@ -33,8 +31,6 @@ pub mod snapshot;
|
||||
#[allow(dead_code)]
|
||||
pub mod state;
|
||||
#[allow(dead_code)]
|
||||
pub mod stealth;
|
||||
#[allow(dead_code)]
|
||||
pub mod storage;
|
||||
#[allow(dead_code)]
|
||||
pub mod stream;
|
||||
|
||||
+14
-287
@@ -184,7 +184,7 @@ pub async fn install_domain_filter_script(
|
||||
const OrigWS = window.WebSocket;
|
||||
window.WebSocket = function(url, protocols) {{
|
||||
try {{
|
||||
const u = new URL(url, location.href);
|
||||
const u = new URL(url);
|
||||
if (!_isDomainAllowed(u.hostname)) throw new DOMException('WebSocket blocked: ' + u.hostname, 'SecurityError');
|
||||
}} catch(e) {{ if (e instanceof DOMException) throw e; }}
|
||||
return new OrigWS(url, protocols);
|
||||
@@ -233,16 +233,15 @@ pub async fn install_domain_filter_script(
|
||||
pub async fn install_domain_filter_fetch(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
handle_auth_requests: bool,
|
||||
) -> Result<(), String> {
|
||||
let mut params = json!({
|
||||
"patterns": [{ "urlPattern": "*" }]
|
||||
});
|
||||
if handle_auth_requests {
|
||||
params["handleAuthRequests"] = json!(true);
|
||||
}
|
||||
client
|
||||
.send_command("Fetch.enable", Some(params), Some(session_id))
|
||||
.send_command(
|
||||
"Fetch.enable",
|
||||
Some(json!({
|
||||
"patterns": [{ "urlPattern": "*" }]
|
||||
})),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -254,103 +253,12 @@ pub async fn install_domain_filter(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
allowed_domains: &[String],
|
||||
handle_auth_requests: bool,
|
||||
) -> Result<(), String> {
|
||||
install_domain_filter_script(client, session_id, allowed_domains).await?;
|
||||
install_domain_filter_fetch(client, session_id, handle_auth_requests).await?;
|
||||
install_domain_filter_fetch(client, session_id).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Console arg formatting (CDP RemoteObject → human-readable string)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Format a single CDP RemoteObject arg into a human-readable string.
|
||||
/// Priority: value → preview → description.
|
||||
pub fn format_console_arg(arg: &Value) -> Option<String> {
|
||||
let obj_type = arg.get("type").and_then(|v| v.as_str()).unwrap_or("");
|
||||
let subtype = arg.get("subtype").and_then(|v| v.as_str());
|
||||
|
||||
if obj_type == "undefined" {
|
||||
return Some("undefined".to_string());
|
||||
}
|
||||
|
||||
if subtype == Some("null") {
|
||||
return Some("null".to_string());
|
||||
}
|
||||
|
||||
// Primitive value
|
||||
if let Some(v) = arg.get("value") {
|
||||
return Some(match v {
|
||||
Value::String(s) => s.clone(),
|
||||
Value::Null => "null".to_string(),
|
||||
other => other.to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Skip preview for Map/Set — their description ("Map(1)", "Set(3)") is more useful
|
||||
// than their preview properties (which only show "size")
|
||||
if let Some(preview) = arg.get("preview") {
|
||||
let preview_subtype = preview.get("subtype").and_then(|v| v.as_str());
|
||||
if matches!(preview_subtype, Some("map" | "set" | "weakmap" | "weakset")) {
|
||||
return arg
|
||||
.get("description")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
}
|
||||
let is_array = subtype == Some("array") || preview_subtype == Some("array");
|
||||
if let Some(props) = preview.get("properties").and_then(|v| v.as_array()) {
|
||||
let overflow = preview
|
||||
.get("overflow")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
let formatted_props: Vec<String> = props
|
||||
.iter()
|
||||
.filter_map(|p| {
|
||||
let value_str = p.get("value").and_then(|v| v.as_str())?;
|
||||
let prop_type = p.get("type").and_then(|v| v.as_str()).unwrap_or("");
|
||||
let formatted_value = if prop_type == "string" {
|
||||
format!("\"{}\"", value_str)
|
||||
} else {
|
||||
value_str.to_string()
|
||||
};
|
||||
if is_array {
|
||||
Some(formatted_value)
|
||||
} else {
|
||||
let name = p.get("name").and_then(|v| v.as_str()).unwrap_or("?");
|
||||
Some(format!("{}: {}", name, formatted_value))
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
let inner = if overflow {
|
||||
format!("{}, ...", formatted_props.join(", "))
|
||||
} else {
|
||||
formatted_props.join(", ")
|
||||
};
|
||||
|
||||
return if is_array {
|
||||
Some(format!("[{}]", inner))
|
||||
} else {
|
||||
Some(format!("{{{}}}", inner))
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to description
|
||||
arg.get("description")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
}
|
||||
|
||||
/// Format an array of CDP RemoteObject args into a single space-separated string.
|
||||
pub fn format_console_args(args: &[Value]) -> String {
|
||||
args.iter()
|
||||
.filter_map(format_console_arg)
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ")
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Console and error tracking
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -359,7 +267,6 @@ pub fn format_console_args(args: &[Value]) -> String {
|
||||
pub struct ConsoleEntry {
|
||||
pub level: String,
|
||||
pub text: String,
|
||||
pub args: Vec<Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -385,14 +292,13 @@ impl EventTracker {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn add_console(&mut self, level: &str, text: &str, args: Vec<Value>) {
|
||||
pub fn add_console(&mut self, level: &str, text: &str) {
|
||||
if self.console_entries.len() >= self.max_entries {
|
||||
self.console_entries.remove(0);
|
||||
}
|
||||
self.console_entries.push(ConsoleEntry {
|
||||
level: level.to_string(),
|
||||
text: text.to_string(),
|
||||
args,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -414,25 +320,13 @@ impl EventTracker {
|
||||
});
|
||||
}
|
||||
|
||||
pub fn clear_console(&mut self) {
|
||||
self.console_entries.clear();
|
||||
}
|
||||
|
||||
pub fn get_console_json(&self) -> Value {
|
||||
let messages: Vec<Value> = self
|
||||
let entries: Vec<Value> = self
|
||||
.console_entries
|
||||
.iter()
|
||||
.map(|e| {
|
||||
let mut msg = json!({ "type": e.level, "text": e.text });
|
||||
if !e.args.is_empty() {
|
||||
msg.as_object_mut()
|
||||
.unwrap()
|
||||
.insert("args".to_string(), Value::Array(e.args.clone()));
|
||||
}
|
||||
msg
|
||||
})
|
||||
.map(|e| json!({ "level": e.level, "text": e.text }))
|
||||
.collect();
|
||||
json!({ "messages": messages })
|
||||
json!({ "entries": entries })
|
||||
}
|
||||
|
||||
pub fn get_errors_json(&self) -> Value {
|
||||
@@ -496,177 +390,10 @@ mod tests {
|
||||
#[test]
|
||||
fn test_event_tracker() {
|
||||
let mut tracker = EventTracker::new();
|
||||
tracker.add_console("log", "hello", vec![]);
|
||||
tracker.add_console("log", "hello");
|
||||
tracker.add_error("oops", Some("test.js"), Some(1), Some(5));
|
||||
|
||||
assert_eq!(tracker.console_entries.len(), 1);
|
||||
assert_eq!(tracker.error_entries.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_console_json_includes_args() {
|
||||
let mut tracker = EventTracker::new();
|
||||
let raw_args = vec![
|
||||
json!({"type": "string", "value": "hello"}),
|
||||
json!({"type": "number", "value": 42}),
|
||||
];
|
||||
tracker.add_console("log", "hello 42", raw_args);
|
||||
|
||||
let result = tracker.get_console_json();
|
||||
let messages = result.get("messages").unwrap().as_array().unwrap();
|
||||
assert_eq!(messages.len(), 1);
|
||||
assert_eq!(messages[0].get("text").unwrap(), "hello 42");
|
||||
let args = messages[0].get("args").unwrap().as_array().unwrap();
|
||||
assert_eq!(args.len(), 2);
|
||||
assert_eq!(args[0], json!({"type": "string", "value": "hello"}));
|
||||
assert_eq!(args[1], json!({"type": "number", "value": 42}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_console_json_empty_args_omits_field() {
|
||||
let mut tracker = EventTracker::new();
|
||||
tracker.add_console("log", "text only", vec![]);
|
||||
|
||||
let result = tracker.get_console_json();
|
||||
let messages = result.get("messages").unwrap().as_array().unwrap();
|
||||
assert!(messages[0].get("args").is_none());
|
||||
}
|
||||
|
||||
// -- format_console_arg: primitives --
|
||||
|
||||
#[test]
|
||||
fn test_format_arg_string() {
|
||||
let arg = json!({"type": "string", "value": "hello"});
|
||||
assert_eq!(format_console_arg(&arg), Some("hello".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_arg_number() {
|
||||
let arg = json!({"type": "number", "value": 42});
|
||||
assert_eq!(format_console_arg(&arg), Some("42".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_arg_null() {
|
||||
let arg = json!({"type": "object", "subtype": "null", "value": null});
|
||||
assert_eq!(format_console_arg(&arg), Some("null".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_arg_undefined() {
|
||||
let arg = json!({"type": "undefined"});
|
||||
assert_eq!(format_console_arg(&arg), Some("undefined".to_string()));
|
||||
}
|
||||
|
||||
// -- format_console_arg: objects with preview --
|
||||
|
||||
#[test]
|
||||
fn test_format_arg_object_preview() {
|
||||
let arg = json!({
|
||||
"type": "object",
|
||||
"preview": {
|
||||
"properties": [
|
||||
{"name": "userId", "type": "string", "value": "abc123"},
|
||||
{"name": "count", "type": "number", "value": "42"}
|
||||
],
|
||||
"overflow": false
|
||||
}
|
||||
});
|
||||
assert_eq!(
|
||||
format_console_arg(&arg),
|
||||
Some("{userId: \"abc123\", count: 42}".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_arg_object_preview_overflow() {
|
||||
let arg = json!({
|
||||
"type": "object",
|
||||
"preview": {
|
||||
"properties": [
|
||||
{"name": "a", "type": "number", "value": "1"}
|
||||
],
|
||||
"overflow": true
|
||||
}
|
||||
});
|
||||
assert_eq!(format_console_arg(&arg), Some("{a: 1, ...}".to_string()));
|
||||
}
|
||||
|
||||
// -- format_console_arg: arrays with preview --
|
||||
|
||||
#[test]
|
||||
fn test_format_arg_array_preview() {
|
||||
let arg = json!({
|
||||
"type": "object",
|
||||
"subtype": "array",
|
||||
"preview": {
|
||||
"subtype": "array",
|
||||
"properties": [
|
||||
{"name": "0", "type": "number", "value": "1"},
|
||||
{"name": "1", "type": "number", "value": "2"},
|
||||
{"name": "2", "type": "number", "value": "3"}
|
||||
],
|
||||
"overflow": false
|
||||
}
|
||||
});
|
||||
assert_eq!(format_console_arg(&arg), Some("[1, 2, 3]".to_string()));
|
||||
}
|
||||
|
||||
// -- format_console_arg: map/set use description --
|
||||
|
||||
#[test]
|
||||
fn test_format_arg_map_uses_description() {
|
||||
let arg = json!({
|
||||
"type": "object",
|
||||
"subtype": "map",
|
||||
"description": "Map(1)",
|
||||
"preview": {
|
||||
"subtype": "map",
|
||||
"properties": [{"name": "size", "type": "number", "value": "1"}]
|
||||
}
|
||||
});
|
||||
assert_eq!(format_console_arg(&arg), Some("Map(1)".to_string()));
|
||||
}
|
||||
|
||||
// -- format_console_arg: fallback --
|
||||
|
||||
#[test]
|
||||
fn test_format_arg_description_fallback() {
|
||||
let arg = json!({"type": "object", "description": "RegExp"});
|
||||
assert_eq!(format_console_arg(&arg), Some("RegExp".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_arg_no_value_no_preview_no_description() {
|
||||
let arg = json!({"type": "object"});
|
||||
assert_eq!(format_console_arg(&arg), None);
|
||||
}
|
||||
|
||||
// -- format_console_args --
|
||||
|
||||
#[test]
|
||||
fn test_format_console_args_join() {
|
||||
let args = vec![
|
||||
json!({"type": "string", "value": "user"}),
|
||||
json!({
|
||||
"type": "object",
|
||||
"preview": {
|
||||
"properties": [{"name": "id", "type": "number", "value": "1"}],
|
||||
"overflow": false
|
||||
}
|
||||
}),
|
||||
];
|
||||
assert_eq!(format_console_args(&args), "user {id: 1}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_console_args_filters_none() {
|
||||
// An arg that returns None should be skipped, not produce empty string
|
||||
let args = vec![
|
||||
json!({"type": "string", "value": "before"}),
|
||||
json!({"type": "object"}), // no value, preview, or description → None
|
||||
json!({"type": "string", "value": "after"}),
|
||||
];
|
||||
assert_eq!(format_console_args(&args), "before after");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -9,38 +9,6 @@ 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",
|
||||
@@ -172,7 +140,6 @@ const DOCUMENTED_ACTIONS: &[&str] = &[
|
||||
"route",
|
||||
"unroute",
|
||||
"requests",
|
||||
"request_detail",
|
||||
"credentials",
|
||||
"auth_save",
|
||||
"auth_login",
|
||||
@@ -344,7 +311,7 @@ fn minimal_command(action: &str, id: &str) -> Value {
|
||||
obj.insert("script".to_string(), json!("h => h"));
|
||||
}
|
||||
"drag" => {
|
||||
obj.insert("source".to_string(), json!("body"));
|
||||
obj.insert("selector".to_string(), json!("body"));
|
||||
obj.insert("target".to_string(), json!("body"));
|
||||
}
|
||||
"swipe" => {
|
||||
@@ -457,7 +424,6 @@ 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);
|
||||
@@ -473,7 +439,6 @@ 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",
|
||||
@@ -534,7 +499,6 @@ async fn test_daemon_state_new_defaults() {
|
||||
assert!(state.tracked_requests.is_empty());
|
||||
assert!(state.active_frame_id.is_none());
|
||||
assert!(state.webdriver_backend.is_none());
|
||||
assert!(state.stream_client.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -546,11 +510,6 @@ async fn test_tracked_request_struct() {
|
||||
headers: json!({"Accept": "text/html"}),
|
||||
timestamp: 12345,
|
||||
resource_type: "Document".to_string(),
|
||||
request_id: "1.1".to_string(),
|
||||
post_data: None,
|
||||
status: Some(200),
|
||||
response_headers: None,
|
||||
mime_type: Some("text/html".to_string()),
|
||||
};
|
||||
let serialized = serde_json::to_value(&tr).unwrap();
|
||||
assert_eq!(serialized["url"], "https://example.com/api");
|
||||
@@ -571,11 +530,6 @@ async fn test_request_tracking_state() {
|
||||
headers: json!({}),
|
||||
timestamp: 1,
|
||||
resource_type: "Document".to_string(),
|
||||
request_id: "1.1".to_string(),
|
||||
post_data: None,
|
||||
status: None,
|
||||
response_headers: None,
|
||||
mime_type: None,
|
||||
});
|
||||
state.tracked_requests.push(super::actions::TrackedRequest {
|
||||
url: "https://other.com".to_string(),
|
||||
@@ -583,11 +537,6 @@ async fn test_request_tracking_state() {
|
||||
headers: json!({}),
|
||||
timestamp: 2,
|
||||
resource_type: "XHR".to_string(),
|
||||
request_id: "1.2".to_string(),
|
||||
post_data: None,
|
||||
status: None,
|
||||
response_headers: None,
|
||||
mime_type: None,
|
||||
});
|
||||
assert_eq!(state.tracked_requests.len(), 2);
|
||||
|
||||
@@ -605,30 +554,6 @@ async fn test_request_tracking_state() {
|
||||
assert!(state.tracked_requests.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_matches_status_filter() {
|
||||
use super::actions::matches_status_filter;
|
||||
|
||||
// Exact match
|
||||
assert!(matches_status_filter(Some(200), "200"));
|
||||
assert!(!matches_status_filter(Some(201), "200"));
|
||||
|
||||
// Class match (Nxx)
|
||||
assert!(matches_status_filter(Some(200), "2xx"));
|
||||
assert!(matches_status_filter(Some(299), "2xx"));
|
||||
assert!(!matches_status_filter(Some(301), "2xx"));
|
||||
assert!(matches_status_filter(Some(404), "4xx"));
|
||||
|
||||
// Range match
|
||||
assert!(matches_status_filter(Some(400), "400-499"));
|
||||
assert!(matches_status_filter(Some(499), "400-499"));
|
||||
assert!(!matches_status_filter(Some(500), "400-499"));
|
||||
|
||||
// None status
|
||||
assert!(!matches_status_filter(None, "200"));
|
||||
assert!(!matches_status_filter(None, "2xx"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_addscript_and_addinitscript_separate_dispatch() {
|
||||
let mut state = DaemonState::new();
|
||||
|
||||
+44
-586
@@ -1,72 +1,28 @@
|
||||
//! Browser provider connections for remote CDP sessions.
|
||||
//!
|
||||
//! Supports AgentCore, Browserbase, Browserless, Browser Use, and Kernel providers.
|
||||
//! Each provider returns a CDP WebSocket URL for connecting via BrowserManager.
|
||||
//! Supports Browserbase, Browser Use, and Kernel providers. Each provider
|
||||
//! returns a CDP WebSocket URL for connecting via BrowserManager.
|
||||
|
||||
use serde_json::{json, Value};
|
||||
use std::env;
|
||||
|
||||
/// Provider session info for cleanup on failure.
|
||||
#[derive(Debug)]
|
||||
pub struct ProviderSession {
|
||||
pub provider: String,
|
||||
pub session_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct ProviderConnection {
|
||||
pub ws_url: String,
|
||||
pub session: Option<ProviderSession>,
|
||||
/// If true, the WebSocket IS the page session (no Target.* commands).
|
||||
pub direct_page: bool,
|
||||
}
|
||||
|
||||
/// Connects to the specified browser provider and returns a CDP WebSocket URL
|
||||
/// along with session info for cleanup on failure.
|
||||
pub async fn connect_provider(provider_name: &str) -> Result<ProviderConnection, String> {
|
||||
pub async fn connect_provider(
|
||||
provider_name: &str,
|
||||
) -> Result<(String, Option<ProviderSession>), String> {
|
||||
match provider_name.to_lowercase().as_str() {
|
||||
"browserbase" => {
|
||||
let (url, session) = connect_browserbase().await?;
|
||||
Ok(ProviderConnection {
|
||||
ws_url: url,
|
||||
session,
|
||||
direct_page: false,
|
||||
})
|
||||
}
|
||||
"browserless" => {
|
||||
let (url, session) = connect_browserless().await?;
|
||||
Ok(ProviderConnection {
|
||||
ws_url: url,
|
||||
session,
|
||||
direct_page: false,
|
||||
})
|
||||
}
|
||||
"browser-use" | "browseruse" => {
|
||||
let (url, session) = connect_browser_use().await?;
|
||||
Ok(ProviderConnection {
|
||||
ws_url: url,
|
||||
session,
|
||||
direct_page: false,
|
||||
})
|
||||
}
|
||||
"kernel" => {
|
||||
let (url, session) = connect_kernel().await?;
|
||||
Ok(ProviderConnection {
|
||||
ws_url: url,
|
||||
session,
|
||||
direct_page: false,
|
||||
})
|
||||
}
|
||||
"agentcore" => {
|
||||
let (url, session) = connect_agentcore().await?;
|
||||
Ok(ProviderConnection {
|
||||
ws_url: url,
|
||||
session,
|
||||
direct_page: false,
|
||||
})
|
||||
}
|
||||
"browserbase" => connect_browserbase().await,
|
||||
"browser-use" | "browseruse" => connect_browser_use().await,
|
||||
"kernel" => connect_kernel().await,
|
||||
_ => Err(format!(
|
||||
"Unknown provider '{}'. Supported: browserbase, browserless, browser-use, kernel, agentcore",
|
||||
"Unknown provider '{}'. Supported: browserbase, browser-use, kernel",
|
||||
provider_name
|
||||
)),
|
||||
}
|
||||
@@ -79,13 +35,11 @@ pub async fn close_provider_session(session: &ProviderSession) {
|
||||
"browserbase" => {
|
||||
if let Ok(api_key) = env::var("BROWSERBASE_API_KEY") {
|
||||
let _ = client
|
||||
.post(format!(
|
||||
.delete(format!(
|
||||
"https://api.browserbase.com/v1/sessions/{}",
|
||||
session.session_id
|
||||
))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("X-BB-API-Key", &api_key)
|
||||
.json(&serde_json::json!({ "status": "REQUEST_RELEASE" }))
|
||||
.send()
|
||||
.await;
|
||||
}
|
||||
@@ -104,10 +58,6 @@ pub async fn close_provider_session(session: &ProviderSession) {
|
||||
.await;
|
||||
}
|
||||
}
|
||||
"browserless" => {
|
||||
// session_id holds the stop URL for browserless
|
||||
let _ = client.delete(&session.session_id).send().await;
|
||||
}
|
||||
"kernel" => {
|
||||
if let Ok(api_key) = env::var("KERNEL_API_KEY") {
|
||||
let endpoint = env::var("KERNEL_ENDPOINT")
|
||||
@@ -123,10 +73,6 @@ pub async fn close_provider_session(session: &ProviderSession) {
|
||||
.await;
|
||||
}
|
||||
}
|
||||
"agentcore" => {
|
||||
// AgentCore session cleanup is handled via signed DELETE request
|
||||
let _ = close_agentcore_session(&session.session_id).await;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
@@ -134,13 +80,15 @@ pub async fn close_provider_session(session: &ProviderSession) {
|
||||
async fn connect_browserbase() -> Result<(String, Option<ProviderSession>), String> {
|
||||
let api_key = env::var("BROWSERBASE_API_KEY")
|
||||
.map_err(|_| "BROWSERBASE_API_KEY environment variable is not set")?;
|
||||
let project_id = env::var("BROWSERBASE_PROJECT_ID")
|
||||
.map_err(|_| "BROWSERBASE_PROJECT_ID environment variable is not set")?;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let response = client
|
||||
.post("https://api.browserbase.com/v1/sessions")
|
||||
.header("content-type", "application/json")
|
||||
.header("x-bb-api-key", &api_key)
|
||||
.body("{}")
|
||||
.header("Content-Type", "application/json")
|
||||
.header("X-BB-API-Key", &api_key)
|
||||
.json(&json!({ "projectId": project_id }))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("Browserbase request failed: {}", e))?;
|
||||
@@ -183,98 +131,62 @@ async fn connect_browserbase() -> Result<(String, Option<ProviderSession>), Stri
|
||||
))
|
||||
}
|
||||
|
||||
async fn connect_browserless() -> Result<(String, Option<ProviderSession>), String> {
|
||||
let api_key = env::var("BROWSERLESS_API_KEY")
|
||||
.map_err(|_| "BROWSERLESS_API_KEY environment variable is not set")?;
|
||||
|
||||
let api_url = env::var("BROWSERLESS_API_URL")
|
||||
.unwrap_or_else(|_| "https://production-sfo.browserless.io".to_string());
|
||||
let browser_type =
|
||||
env::var("BROWSERLESS_BROWSER_TYPE").unwrap_or_else(|_| "chromium".to_string());
|
||||
|
||||
let supported = ["chromium", "chrome"];
|
||||
if !supported.contains(&browser_type.as_str()) {
|
||||
return Err(format!(
|
||||
"BROWSERLESS_BROWSER_TYPE \"{}\" is not supported. Only {} are allowed.",
|
||||
browser_type,
|
||||
supported.join(", ")
|
||||
));
|
||||
}
|
||||
|
||||
let ttl: u64 = env::var("BROWSERLESS_TTL")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(300000);
|
||||
let stealth = env::var("BROWSERLESS_STEALTH")
|
||||
.map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
|
||||
.unwrap_or(true);
|
||||
|
||||
let url = format!("{}/session", api_url.trim_end_matches('/'));
|
||||
async fn connect_browser_use() -> Result<(String, Option<ProviderSession>), String> {
|
||||
let api_key = env::var("BROWSER_USE_API_KEY")
|
||||
.map_err(|_| "BROWSER_USE_API_KEY environment variable is not set")?;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let response = client
|
||||
.post(&url)
|
||||
.query(&[("token", &api_key)])
|
||||
.post("https://api.browser-use.com/api/v2/browsers")
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&json!({
|
||||
"ttl": ttl,
|
||||
"stealth": stealth,
|
||||
"browser": browser_type,
|
||||
}))
|
||||
.header("X-Browser-Use-API-Key", &api_key)
|
||||
.json(&json!({}))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("Browserless request failed: {}", e))?;
|
||||
.map_err(|e| format!("Browser Use request failed: {}", e))?;
|
||||
|
||||
let status = response.status();
|
||||
let body = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to read Browserless response: {}", e))?;
|
||||
.map_err(|e| format!("Failed to read Browser Use response: {}", e))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(format!(
|
||||
"Browserless API error ({}): {}",
|
||||
"Browser Use API error ({}): {}",
|
||||
status.as_u16(),
|
||||
body
|
||||
));
|
||||
}
|
||||
|
||||
let json: Value =
|
||||
serde_json::from_str(&body).map_err(|e| format!("Invalid Browserless response: {}", e))?;
|
||||
serde_json::from_str(&body).map_err(|e| format!("Invalid Browser Use response: {}", e))?;
|
||||
|
||||
let connect_url = json
|
||||
.get("connect")
|
||||
let session_id = json
|
||||
.get("id")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
|
||||
let ws_url = json
|
||||
.get("cdp_url")
|
||||
.or_else(|| json.get("cdpUrl"))
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from)
|
||||
.ok_or_else(|| "Browserless response missing 'connect' URL".to_string())?;
|
||||
|
||||
let stop_url = json
|
||||
.get("stop")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from)
|
||||
.ok_or_else(|| "Browserless response missing 'stop' URL".to_string())?;
|
||||
.ok_or_else(|| "Browser Use response missing cdp_url or cdpUrl".to_string())?;
|
||||
|
||||
Ok((
|
||||
connect_url,
|
||||
ws_url,
|
||||
Some(ProviderSession {
|
||||
provider: "browserless".to_string(),
|
||||
// Store the stop URL as the session_id for cleanup
|
||||
session_id: stop_url,
|
||||
provider: "browser-use".to_string(),
|
||||
session_id,
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
async fn connect_browser_use() -> Result<(String, Option<ProviderSession>), String> {
|
||||
let api_key = env::var("BROWSER_USE_API_KEY")
|
||||
.map_err(|_| "BROWSER_USE_API_KEY environment variable is not set")?;
|
||||
|
||||
let ws_url = format!("wss://connect.browser-use.com?apiKey={}", api_key);
|
||||
|
||||
Ok((ws_url, None))
|
||||
}
|
||||
|
||||
async fn connect_kernel() -> Result<(String, Option<ProviderSession>), String> {
|
||||
let api_key = env::var("KERNEL_API_KEY").ok();
|
||||
let api_key =
|
||||
env::var("KERNEL_API_KEY").map_err(|_| "KERNEL_API_KEY environment variable is not set")?;
|
||||
let endpoint =
|
||||
env::var("KERNEL_ENDPOINT").unwrap_or_else(|_| "https://api.onkernel.com".to_string());
|
||||
|
||||
@@ -306,11 +218,10 @@ async fn connect_kernel() -> Result<(String, Option<ProviderSession>), String> {
|
||||
}
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let mut request = client.post(&url).header("Content-Type", "application/json");
|
||||
if let Some(ref key) = api_key {
|
||||
request = request.header("Authorization", format!("Bearer {}", key));
|
||||
}
|
||||
let response = request
|
||||
let response = client
|
||||
.post(&url)
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Authorization", format!("Bearer {}", api_key))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
@@ -361,456 +272,3 @@ async fn connect_kernel() -> Result<(String, Option<ProviderSession>), String> {
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// AgentCore Provider (AWS Bedrock AgentCore Browser)
|
||||
// ============================================================================
|
||||
|
||||
mod agentcore {
|
||||
use super::*;
|
||||
|
||||
/// AgentCore-specific session info for Live View URL
|
||||
pub struct AgentCoreSessionInfo {
|
||||
pub session_id: String,
|
||||
pub browser_identifier: String,
|
||||
pub region: String,
|
||||
pub live_view_url: String,
|
||||
}
|
||||
|
||||
thread_local! {
|
||||
static AGENTCORE_INFO: std::cell::RefCell<Option<AgentCoreSessionInfo>> = const { std::cell::RefCell::new(None) };
|
||||
static AGENTCORE_WS_HEADERS: std::cell::RefCell<Option<Vec<(String, String)>>> = const { std::cell::RefCell::new(None) };
|
||||
}
|
||||
|
||||
pub fn set_agentcore_info(info: AgentCoreSessionInfo) {
|
||||
AGENTCORE_INFO.with(|cell| *cell.borrow_mut() = Some(info));
|
||||
}
|
||||
|
||||
pub fn get_agentcore_info() -> Option<AgentCoreSessionInfo> {
|
||||
AGENTCORE_INFO.with(|cell| {
|
||||
cell.borrow().as_ref().map(|i| AgentCoreSessionInfo {
|
||||
session_id: i.session_id.clone(),
|
||||
browser_identifier: i.browser_identifier.clone(),
|
||||
region: i.region.clone(),
|
||||
live_view_url: i.live_view_url.clone(),
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub fn set_agentcore_ws_headers(headers: Vec<(String, String)>) {
|
||||
AGENTCORE_WS_HEADERS.with(|cell| *cell.borrow_mut() = Some(headers));
|
||||
}
|
||||
|
||||
pub fn take_agentcore_ws_headers() -> Option<Vec<(String, String)>> {
|
||||
AGENTCORE_WS_HEADERS.with(|cell| cell.borrow_mut().take())
|
||||
}
|
||||
|
||||
pub async fn connect() -> Result<(String, Option<ProviderSession>), String> {
|
||||
let region = env::var("AGENTCORE_REGION")
|
||||
.or_else(|_| env::var("AWS_REGION"))
|
||||
.or_else(|_| env::var("AWS_DEFAULT_REGION"))
|
||||
.unwrap_or_else(|_| "us-east-1".to_string());
|
||||
let browser_id =
|
||||
env::var("AGENTCORE_BROWSER_ID").unwrap_or_else(|_| "aws.browser.v1".to_string());
|
||||
let timeout_secs: u64 = env::var("AGENTCORE_SESSION_TIMEOUT")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(3600);
|
||||
|
||||
let host = format!("bedrock-agentcore.{}.amazonaws.com", region);
|
||||
let path = format!(
|
||||
"/browsers/{}/sessions/start",
|
||||
urlencoding::encode(&browser_id)
|
||||
);
|
||||
let url = format!("https://{}{}", host, path);
|
||||
|
||||
// Generate a unique session name
|
||||
let session_name = format!("agent-browser-{}", &uuid::Uuid::new_v4().to_string()[..8]);
|
||||
|
||||
let mut body_json = json!({
|
||||
"name": session_name,
|
||||
"sessionTimeoutSeconds": timeout_secs
|
||||
});
|
||||
if let Ok(profile_id) = env::var("AGENTCORE_PROFILE_ID") {
|
||||
if !profile_id.is_empty() {
|
||||
body_json.as_object_mut().unwrap().insert(
|
||||
"profileConfiguration".to_string(),
|
||||
json!({ "profileIdentifier": profile_id }),
|
||||
);
|
||||
}
|
||||
}
|
||||
let body = serde_json::to_string(&body_json)
|
||||
.map_err(|e| format!("Failed to serialize request body: {}", e))?;
|
||||
|
||||
let signed_headers = sign_request("PUT", &url, ®ion, Some(&body)).await?;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let mut req = client.put(&url).body(body.clone());
|
||||
for (key, value) in &signed_headers {
|
||||
req = req.header(key.as_str(), value.as_str());
|
||||
}
|
||||
|
||||
let response = req
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("AgentCore request failed: {}", e))?;
|
||||
|
||||
let status = response.status();
|
||||
let resp_body = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to read AgentCore response: {}", e))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(format!(
|
||||
"AgentCore API error ({}): {}",
|
||||
status.as_u16(),
|
||||
resp_body
|
||||
));
|
||||
}
|
||||
|
||||
let json: Value = serde_json::from_str(&resp_body)
|
||||
.map_err(|e| format!("Invalid AgentCore response: {}", e))?;
|
||||
|
||||
let session_id = json
|
||||
.get("sessionId")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| "AgentCore response missing sessionId".to_string())?
|
||||
.to_string();
|
||||
|
||||
let browser_identifier = json
|
||||
.get("browserIdentifier")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or(&browser_id)
|
||||
.to_string();
|
||||
|
||||
let live_view_url = format!(
|
||||
"https://{}.console.aws.amazon.com/bedrock-agentcore/browser/{}/session/{}#",
|
||||
region, browser_identifier, session_id
|
||||
);
|
||||
|
||||
set_agentcore_info(AgentCoreSessionInfo {
|
||||
session_id: session_id.clone(),
|
||||
browser_identifier: browser_identifier.clone(),
|
||||
region: region.clone(),
|
||||
live_view_url: live_view_url.clone(),
|
||||
});
|
||||
|
||||
eprintln!("Session: {}", session_id);
|
||||
eprintln!("Live View: {}", live_view_url);
|
||||
|
||||
let ws_path = format!(
|
||||
"/browser-streams/{}/sessions/{}/automation",
|
||||
browser_identifier, session_id
|
||||
);
|
||||
let ws_url = format!("wss://{}{}", host, ws_path);
|
||||
|
||||
let ws_headers = sign_request(
|
||||
"GET",
|
||||
&format!("https://{}{}", host, ws_path),
|
||||
®ion,
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
set_agentcore_ws_headers(ws_headers);
|
||||
|
||||
Ok((
|
||||
ws_url,
|
||||
Some(ProviderSession {
|
||||
provider: "agentcore".to_string(),
|
||||
session_id,
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
/// Get AWS credentials from environment variables or AWS CLI
|
||||
fn get_aws_credentials() -> Result<(String, String, Option<String>), String> {
|
||||
// First try environment variables
|
||||
if let (Ok(access_key), Ok(secret_key)) = (
|
||||
env::var("AWS_ACCESS_KEY_ID"),
|
||||
env::var("AWS_SECRET_ACCESS_KEY"),
|
||||
) {
|
||||
return Ok((access_key, secret_key, env::var("AWS_SESSION_TOKEN").ok()));
|
||||
}
|
||||
|
||||
// Fall back to AWS CLI
|
||||
let mut cmd = std::process::Command::new("aws");
|
||||
cmd.args(["configure", "export-credentials", "--format", "env"]);
|
||||
|
||||
// Honor AWS_PROFILE
|
||||
if let Ok(profile) = env::var("AWS_PROFILE") {
|
||||
cmd.args(["--profile", &profile]);
|
||||
}
|
||||
|
||||
let output = cmd.output()
|
||||
.map_err(|e| format!("Failed to run aws CLI: {}. Install AWS CLI or set AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY", e))?;
|
||||
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
return Err(format!(
|
||||
"AWS CLI failed: {}. Run 'aws sso login' or set credentials",
|
||||
stderr.trim()
|
||||
));
|
||||
}
|
||||
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let mut access_key = None;
|
||||
let mut secret_key = None;
|
||||
let mut session_token = None;
|
||||
|
||||
for line in stdout.lines() {
|
||||
if let Some(val) = line.strip_prefix("export AWS_ACCESS_KEY_ID=") {
|
||||
access_key = Some(val.to_string());
|
||||
} else if let Some(val) = line.strip_prefix("export AWS_SECRET_ACCESS_KEY=") {
|
||||
secret_key = Some(val.to_string());
|
||||
} else if let Some(val) = line.strip_prefix("export AWS_SESSION_TOKEN=") {
|
||||
session_token = Some(val.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
match (access_key, secret_key) {
|
||||
(Some(ak), Some(sk)) => Ok((ak, sk, session_token)),
|
||||
_ => Err("Failed to parse credentials from AWS CLI output".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn sign_request(
|
||||
method: &str,
|
||||
url: &str,
|
||||
region: &str,
|
||||
body: Option<&str>,
|
||||
) -> Result<Vec<(String, String)>, String> {
|
||||
use hmac::{Hmac, Mac};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
// Get credentials from environment or AWS CLI
|
||||
let (access_key, secret_key, session_token) = get_aws_credentials()?;
|
||||
|
||||
let parsed_url = url::Url::parse(url).map_err(|e| format!("Invalid URL: {}", e))?;
|
||||
let host = parsed_url.host_str().unwrap_or("");
|
||||
|
||||
// Get current time
|
||||
let now = chrono::Utc::now();
|
||||
let amz_date = now.format("%Y%m%dT%H%M%SZ").to_string();
|
||||
let date_stamp = now.format("%Y%m%d").to_string();
|
||||
|
||||
// Create canonical request
|
||||
let payload_hash = if let Some(b) = body {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(b.as_bytes());
|
||||
hex::encode(hasher.finalize())
|
||||
} else {
|
||||
"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855".to_string()
|
||||
// empty string hash
|
||||
};
|
||||
|
||||
let canonical_uri = parsed_url.path();
|
||||
let canonical_querystring = parsed_url.query().unwrap_or("");
|
||||
|
||||
let mut signed_headers = "content-type;host;x-amz-date".to_string();
|
||||
let mut canonical_headers = format!(
|
||||
"content-type:application/json\nhost:{}\nx-amz-date:{}\n",
|
||||
host, amz_date
|
||||
);
|
||||
|
||||
if let Some(ref token) = session_token {
|
||||
signed_headers = "content-type;host;x-amz-date;x-amz-security-token".to_string();
|
||||
canonical_headers = format!(
|
||||
"content-type:application/json\nhost:{}\nx-amz-date:{}\nx-amz-security-token:{}\n",
|
||||
host, amz_date, token
|
||||
);
|
||||
}
|
||||
|
||||
let canonical_request = format!(
|
||||
"{}\n{}\n{}\n{}\n{}\n{}",
|
||||
method,
|
||||
canonical_uri,
|
||||
canonical_querystring,
|
||||
canonical_headers,
|
||||
signed_headers,
|
||||
payload_hash
|
||||
);
|
||||
|
||||
// Create string to sign
|
||||
let algorithm = "AWS4-HMAC-SHA256";
|
||||
let credential_scope = format!("{}/{}/bedrock-agentcore/aws4_request", date_stamp, region);
|
||||
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(canonical_request.as_bytes());
|
||||
let canonical_request_hash = hex::encode(hasher.finalize());
|
||||
|
||||
let string_to_sign = format!(
|
||||
"{}\n{}\n{}\n{}",
|
||||
algorithm, amz_date, credential_scope, canonical_request_hash
|
||||
);
|
||||
|
||||
// Calculate signature
|
||||
type HmacSha256 = Hmac<Sha256>;
|
||||
|
||||
let k_date = HmacSha256::new_from_slice(format!("AWS4{}", secret_key).as_bytes())
|
||||
.unwrap()
|
||||
.chain_update(date_stamp.as_bytes())
|
||||
.finalize()
|
||||
.into_bytes();
|
||||
|
||||
let k_region = HmacSha256::new_from_slice(&k_date)
|
||||
.unwrap()
|
||||
.chain_update(region.as_bytes())
|
||||
.finalize()
|
||||
.into_bytes();
|
||||
|
||||
let k_service = HmacSha256::new_from_slice(&k_region)
|
||||
.unwrap()
|
||||
.chain_update(b"bedrock-agentcore")
|
||||
.finalize()
|
||||
.into_bytes();
|
||||
|
||||
let k_signing = HmacSha256::new_from_slice(&k_service)
|
||||
.unwrap()
|
||||
.chain_update(b"aws4_request")
|
||||
.finalize()
|
||||
.into_bytes();
|
||||
|
||||
let signature = hex::encode(
|
||||
HmacSha256::new_from_slice(&k_signing)
|
||||
.unwrap()
|
||||
.chain_update(string_to_sign.as_bytes())
|
||||
.finalize()
|
||||
.into_bytes(),
|
||||
);
|
||||
|
||||
// Build authorization header
|
||||
let authorization = format!(
|
||||
"{} Credential={}/{}, SignedHeaders={}, Signature={}",
|
||||
algorithm, access_key, credential_scope, signed_headers, signature
|
||||
);
|
||||
|
||||
let mut headers = vec![
|
||||
("host".to_string(), host.to_string()),
|
||||
("content-type".to_string(), "application/json".to_string()),
|
||||
("x-amz-date".to_string(), amz_date),
|
||||
("authorization".to_string(), authorization),
|
||||
];
|
||||
|
||||
if let Some(token) = session_token {
|
||||
headers.push(("x-amz-security-token".to_string(), token));
|
||||
}
|
||||
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
pub async fn close_session(session_id: &str) -> Result<(), String> {
|
||||
let info = get_agentcore_info();
|
||||
let (region, browser_id) = match &info {
|
||||
Some(i) => (i.region.clone(), i.browser_identifier.clone()),
|
||||
None => {
|
||||
let region = env::var("AGENTCORE_REGION")
|
||||
.or_else(|_| env::var("AWS_REGION"))
|
||||
.or_else(|_| env::var("AWS_DEFAULT_REGION"))
|
||||
.unwrap_or_else(|_| "us-east-1".to_string());
|
||||
let browser_id = env::var("AGENTCORE_BROWSER_ID")
|
||||
.unwrap_or_else(|_| "aws.browser.v1".to_string());
|
||||
(region, browser_id)
|
||||
}
|
||||
};
|
||||
|
||||
let host = format!("bedrock-agentcore.{}.amazonaws.com", region);
|
||||
let path = format!(
|
||||
"/browsers/{}/sessions/stop",
|
||||
urlencoding::encode(&browser_id)
|
||||
);
|
||||
let url = format!("https://{}{}", host, path);
|
||||
|
||||
let body = serde_json::to_string(&json!({ "sessionId": session_id }))
|
||||
.map_err(|e| format!("Failed to serialize close request: {}", e))?;
|
||||
|
||||
let signed_headers = sign_request("PUT", &url, ®ion, Some(&body)).await?;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let mut req = client.put(&url).body(body);
|
||||
for (key, value) in &signed_headers {
|
||||
req = req.header(key.as_str(), value.as_str());
|
||||
}
|
||||
|
||||
let _ = req.send().await;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub use agentcore::{get_agentcore_info, take_agentcore_ws_headers};
|
||||
|
||||
async fn connect_agentcore() -> Result<(String, Option<ProviderSession>), String> {
|
||||
agentcore::connect().await
|
||||
}
|
||||
|
||||
async fn close_agentcore_session(session_id: &str) -> Result<(), String> {
|
||||
agentcore::close_session(session_id).await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_connect_provider_unknown() {
|
||||
let rt = tokio::runtime::Runtime::new().unwrap();
|
||||
let result = rt.block_on(connect_provider("unknown-provider"));
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("Unknown provider"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_agentcore_env_defaults() {
|
||||
// Test that default values are used when env vars not set
|
||||
std::env::remove_var("AGENTCORE_REGION");
|
||||
std::env::remove_var("AGENTCORE_BROWSER_ID");
|
||||
std::env::remove_var("AGENTCORE_SESSION_TIMEOUT");
|
||||
|
||||
// These would be used in connect() - just verify they don't panic
|
||||
let region = std::env::var("AGENTCORE_REGION")
|
||||
.or_else(|_| std::env::var("AWS_REGION"))
|
||||
.unwrap_or_else(|_| "us-east-1".to_string());
|
||||
assert_eq!(region, "us-east-1");
|
||||
|
||||
let browser_id =
|
||||
std::env::var("AGENTCORE_BROWSER_ID").unwrap_or_else(|_| "aws.browser.v1".to_string());
|
||||
assert_eq!(browser_id, "aws.browser.v1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_agentcore_session_info_storage() {
|
||||
let info = agentcore::AgentCoreSessionInfo {
|
||||
session_id: "test-session".to_string(),
|
||||
browser_identifier: "aws.browser.v1".to_string(),
|
||||
region: "us-east-1".to_string(),
|
||||
live_view_url: "https://example.com".to_string(),
|
||||
};
|
||||
|
||||
agentcore::set_agentcore_info(info);
|
||||
let retrieved = get_agentcore_info();
|
||||
assert!(retrieved.is_some());
|
||||
let retrieved = retrieved.unwrap();
|
||||
assert_eq!(retrieved.session_id, "test-session");
|
||||
assert_eq!(retrieved.region, "us-east-1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_agentcore_ws_headers_storage() {
|
||||
let headers = vec![
|
||||
(
|
||||
"Authorization".to_string(),
|
||||
"AWS4-HMAC-SHA256...".to_string(),
|
||||
),
|
||||
("X-Amz-Date".to_string(), "20260304T180000Z".to_string()),
|
||||
];
|
||||
|
||||
agentcore::set_agentcore_ws_headers(headers);
|
||||
let taken = take_agentcore_ws_headers();
|
||||
assert!(taken.is_some());
|
||||
assert_eq!(taken.unwrap().len(), 2);
|
||||
|
||||
// Should be None after take
|
||||
let taken_again = take_agentcore_ws_headers();
|
||||
assert!(taken_again.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
+92
-212
@@ -1,24 +1,12 @@
|
||||
use serde_json::{json, Value};
|
||||
use std::process::Stdio;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::sync::oneshot;
|
||||
|
||||
use super::cdp::client::CdpClient;
|
||||
use super::cdp::types::{CaptureScreenshotParams, CaptureScreenshotResult};
|
||||
|
||||
const CAPTURE_INTERVAL_MS: u64 = 100;
|
||||
const CAPTURE_FPS: u32 = 10;
|
||||
use std::path::PathBuf;
|
||||
use std::process::Command;
|
||||
|
||||
pub struct RecordingState {
|
||||
pub active: bool,
|
||||
pub output_path: String,
|
||||
pub temp_dir: PathBuf,
|
||||
pub frame_count: u64,
|
||||
pub capture_task: Option<tokio::task::JoinHandle<Result<(), String>>>,
|
||||
pub shared_frame_count: Option<Arc<AtomicU64>>,
|
||||
pub cancel_tx: Option<oneshot::Sender<()>>,
|
||||
}
|
||||
|
||||
impl RecordingState {
|
||||
@@ -26,10 +14,8 @@ impl RecordingState {
|
||||
Self {
|
||||
active: false,
|
||||
output_path: String::new(),
|
||||
temp_dir: PathBuf::new(),
|
||||
frame_count: 0,
|
||||
capture_task: None,
|
||||
shared_frame_count: None,
|
||||
cancel_tx: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -39,13 +25,34 @@ pub fn recording_start(state: &mut RecordingState, path: &str) -> Result<Value,
|
||||
return Err("Recording already active".to_string());
|
||||
}
|
||||
|
||||
let timestamp = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_millis();
|
||||
|
||||
let temp_dir = std::env::temp_dir().join(format!("agent-browser-recording-{}", timestamp));
|
||||
let _ = std::fs::create_dir_all(&temp_dir);
|
||||
|
||||
state.active = true;
|
||||
state.output_path = path.to_string();
|
||||
state.temp_dir = temp_dir;
|
||||
state.frame_count = 0;
|
||||
|
||||
Ok(json!({ "started": true, "path": path }))
|
||||
}
|
||||
|
||||
pub fn recording_add_frame(state: &mut RecordingState, frame_data: &[u8]) {
|
||||
if !state.active {
|
||||
return;
|
||||
}
|
||||
|
||||
let frame_path = state
|
||||
.temp_dir
|
||||
.join(format!("frame_{:06}.jpg", state.frame_count));
|
||||
let _ = std::fs::write(&frame_path, frame_data);
|
||||
state.frame_count += 1;
|
||||
}
|
||||
|
||||
pub fn recording_stop(state: &mut RecordingState) -> Result<Value, String> {
|
||||
if !state.active {
|
||||
return Err("No recording in progress".to_string());
|
||||
@@ -54,183 +61,55 @@ pub fn recording_stop(state: &mut RecordingState) -> Result<Value, String> {
|
||||
state.active = false;
|
||||
|
||||
if state.frame_count == 0 {
|
||||
let _ = std::fs::remove_dir_all(&state.temp_dir);
|
||||
return Err("No frames captured".to_string());
|
||||
}
|
||||
|
||||
Ok(json!({ "path": &state.output_path, "frames": state.frame_count }))
|
||||
}
|
||||
let frame_pattern = state
|
||||
.temp_dir
|
||||
.join("frame_%06d.jpg")
|
||||
.to_string_lossy()
|
||||
.to_string();
|
||||
|
||||
pub fn recording_restart(state: &mut RecordingState, path: &str) -> Result<Value, String> {
|
||||
let previous = if state.active {
|
||||
let stop_result = recording_stop(state);
|
||||
stop_result
|
||||
.ok()
|
||||
.and_then(|v| v.get("path").and_then(|p| p.as_str()).map(String::from))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let output = &state.output_path;
|
||||
|
||||
recording_start(state, path)?;
|
||||
|
||||
Ok(json!({
|
||||
"restarted": true,
|
||||
"previousPath": previous,
|
||||
"path": path,
|
||||
}))
|
||||
}
|
||||
|
||||
fn build_ffmpeg_command(output_path: &str) -> tokio::process::Command {
|
||||
let mut cmd = tokio::process::Command::new("ffmpeg");
|
||||
|
||||
cmd.args(["-y"])
|
||||
.args(["-avioflags", "direct"])
|
||||
// Encode with ffmpeg
|
||||
let result = Command::new("ffmpeg")
|
||||
.args([
|
||||
"-fpsprobesize",
|
||||
"0",
|
||||
"-probesize",
|
||||
"32",
|
||||
"-analyzeduration",
|
||||
"0",
|
||||
])
|
||||
.args([
|
||||
"-f",
|
||||
"image2pipe",
|
||||
"-c:v",
|
||||
"mjpeg",
|
||||
"-y",
|
||||
"-framerate",
|
||||
&CAPTURE_FPS.to_string(),
|
||||
"30",
|
||||
"-i",
|
||||
"pipe:0",
|
||||
&frame_pattern,
|
||||
"-c:v",
|
||||
"libx264",
|
||||
"-pix_fmt",
|
||||
"yuv420p",
|
||||
"-preset",
|
||||
"fast",
|
||||
output,
|
||||
])
|
||||
.args(["-vf", "pad=ceil(iw/2)*2:ceil(ih/2)*2"]);
|
||||
.output();
|
||||
|
||||
if output_path.ends_with(".webm") {
|
||||
cmd.args(["-c:v", "libvpx", "-crf", "30", "-b:v", "1M"]);
|
||||
} else {
|
||||
cmd.args(["-c:v", "libx264", "-preset", "ultrafast"]);
|
||||
}
|
||||
let _ = std::fs::remove_dir_all(&state.temp_dir);
|
||||
|
||||
cmd.args(["-pix_fmt", "yuv420p", "-threads", "1"])
|
||||
.arg(output_path)
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::piped())
|
||||
.kill_on_drop(true);
|
||||
|
||||
cmd
|
||||
}
|
||||
|
||||
/// Spawn a background task that captures screenshots at a fixed interval
|
||||
/// and pipes them to ffmpeg in real-time.
|
||||
pub fn spawn_recording_task(
|
||||
client: Arc<CdpClient>,
|
||||
session_id: String,
|
||||
output_path: String,
|
||||
shared_count: Arc<AtomicU64>,
|
||||
cancel_rx: oneshot::Receiver<()>,
|
||||
) -> tokio::task::JoinHandle<Result<(), String>> {
|
||||
tokio::spawn(async move {
|
||||
let mut cancel_rx = std::pin::pin!(cancel_rx);
|
||||
|
||||
let mut ffmpeg = build_ffmpeg_command(&output_path).spawn().map_err(|e| {
|
||||
format!(
|
||||
"ffmpeg not found or failed to execute: {}. Install ffmpeg to enable recording.",
|
||||
e
|
||||
)
|
||||
})?;
|
||||
|
||||
let mut stdin = ffmpeg
|
||||
.stdin
|
||||
.take()
|
||||
.ok_or_else(|| "Failed to open ffmpeg stdin".to_string())?;
|
||||
|
||||
let mut interval = tokio::time::interval(Duration::from_millis(CAPTURE_INTERVAL_MS));
|
||||
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
|
||||
let params = CaptureScreenshotParams {
|
||||
format: Some("jpeg".to_string()),
|
||||
quality: Some(80),
|
||||
clip: None,
|
||||
from_surface: Some(true),
|
||||
capture_beyond_viewport: None,
|
||||
};
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = &mut cancel_rx => break,
|
||||
_ = interval.tick() => {}
|
||||
match result {
|
||||
Ok(output_result) => {
|
||||
if output_result.status.success() {
|
||||
Ok(json!({ "path": output, "frames": state.frame_count }))
|
||||
} else {
|
||||
let stderr = String::from_utf8_lossy(&output_result.stderr);
|
||||
Err(format!(
|
||||
"ffmpeg failed: {}",
|
||||
stderr.chars().take(200).collect::<String>()
|
||||
))
|
||||
}
|
||||
|
||||
let result: Result<CaptureScreenshotResult, _> = client
|
||||
.send_command_typed("Page.captureScreenshot", ¶ms, Some(&session_id))
|
||||
.await;
|
||||
|
||||
let screenshot = match result {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
if e.contains("Target closed") || e.contains("not found") {
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let bytes = match base64::Engine::decode(
|
||||
&base64::engine::general_purpose::STANDARD,
|
||||
&screenshot.data,
|
||||
) {
|
||||
Ok(b) => b,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
if stdin.write_all(&bytes).await.is_err() {
|
||||
break;
|
||||
}
|
||||
shared_count.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
drop(stdin);
|
||||
|
||||
let output = ffmpeg
|
||||
.wait_with_output()
|
||||
.await
|
||||
.map_err(|e| format!("ffmpeg wait failed: {}", e))?;
|
||||
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
return Err(format!(
|
||||
"ffmpeg failed: {}",
|
||||
stderr.chars().take(300).collect::<String>()
|
||||
));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn stop_recording_task(state: &mut RecordingState) -> Result<(), String> {
|
||||
if let Some(tx) = state.cancel_tx.take() {
|
||||
let _ = tx.send(());
|
||||
Err(e) => Err(format!(
|
||||
"ffmpeg not found or failed to execute: {}. Install ffmpeg to enable recording.",
|
||||
e
|
||||
)),
|
||||
}
|
||||
|
||||
let counter = state.shared_frame_count.take();
|
||||
let handle = state.capture_task.take();
|
||||
|
||||
let result = if let Some(h) = handle {
|
||||
match h.await {
|
||||
Ok(Ok(())) => Ok(()),
|
||||
Ok(Err(e)) => Err(e),
|
||||
Err(e) => Err(format!("Recording task panicked: {}", e)),
|
||||
}
|
||||
} else {
|
||||
Ok(())
|
||||
};
|
||||
|
||||
if let Some(c) = counter {
|
||||
state.frame_count = c.load(Ordering::Relaxed);
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -253,15 +132,19 @@ mod tests {
|
||||
assert!(state.active);
|
||||
assert_eq!(state.output_path, "/tmp/test.mp4");
|
||||
assert_eq!(state.frame_count, 0);
|
||||
// Cleanup
|
||||
let _ = std::fs::remove_dir_all(&state.temp_dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recording_start_while_active() {
|
||||
let mut state = RecordingState::new();
|
||||
recording_start(&mut state, "/tmp/test1.mp4").unwrap();
|
||||
let temp_dir = state.temp_dir.clone();
|
||||
let result = recording_start(&mut state, "/tmp/test2.mp4");
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("already active"));
|
||||
let _ = std::fs::remove_dir_all(&temp_dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -283,41 +166,38 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recording_restart_while_inactive() {
|
||||
fn test_recording_add_frame_inactive() {
|
||||
let mut state = RecordingState::new();
|
||||
let result = recording_restart(&mut state, "/tmp/new.webm");
|
||||
assert!(result.is_ok());
|
||||
assert!(state.active);
|
||||
assert_eq!(state.output_path, "/tmp/new.webm");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recording_restart_while_active() {
|
||||
let mut state = RecordingState::new();
|
||||
recording_start(&mut state, "/tmp/old.webm").unwrap();
|
||||
state.frame_count = 10;
|
||||
let result = recording_restart(&mut state, "/tmp/new.webm").unwrap();
|
||||
assert!(state.active);
|
||||
assert_eq!(state.output_path, "/tmp/new.webm");
|
||||
recording_add_frame(&mut state, b"fake-frame");
|
||||
assert_eq!(state.frame_count, 0);
|
||||
assert_eq!(result["previousPath"], "/tmp/old.webm");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_ffmpeg_command_webm() {
|
||||
let cmd = build_ffmpeg_command("/tmp/out.webm");
|
||||
let args: Vec<&std::ffi::OsStr> = cmd.as_std().get_args().collect();
|
||||
let args_str: Vec<&str> = args.iter().filter_map(|a| a.to_str()).collect();
|
||||
assert!(args_str.contains(&"libvpx"));
|
||||
assert!(args_str.contains(&"/tmp/out.webm"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_ffmpeg_command_mp4() {
|
||||
let cmd = build_ffmpeg_command("/tmp/out.mp4");
|
||||
let args: Vec<&std::ffi::OsStr> = cmd.as_std().get_args().collect();
|
||||
let args_str: Vec<&str> = args.iter().filter_map(|a| a.to_str()).collect();
|
||||
assert!(args_str.contains(&"libx264"));
|
||||
assert!(args_str.contains(&"/tmp/out.mp4"));
|
||||
fn test_recording_add_frame_active() {
|
||||
let mut state = RecordingState::new();
|
||||
recording_start(&mut state, "/tmp/test.mp4").unwrap();
|
||||
recording_add_frame(&mut state, b"fake-frame-1");
|
||||
recording_add_frame(&mut state, b"fake-frame-2");
|
||||
assert_eq!(state.frame_count, 2);
|
||||
let _ = std::fs::remove_dir_all(&state.temp_dir);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn recording_restart(state: &mut RecordingState, path: &str) -> Result<Value, String> {
|
||||
let previous = if state.active {
|
||||
let stop_result = recording_stop(state);
|
||||
stop_result
|
||||
.ok()
|
||||
.and_then(|v| v.get("path").and_then(|p| p.as_str()).map(String::from))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
recording_start(state, path)?;
|
||||
|
||||
Ok(json!({
|
||||
"restarted": true,
|
||||
"previousPath": previous,
|
||||
"path": path,
|
||||
}))
|
||||
}
|
||||
|
||||
+44
-588
@@ -1,65 +1,16 @@
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use super::cdp::client::CdpClient;
|
||||
use super::cdp::types::*;
|
||||
use super::element::RefMap;
|
||||
|
||||
const ANNOTATION_OVERLAY_ID: &str = "__agent_browser_annotations__";
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct Rect {
|
||||
x: f64,
|
||||
y: f64,
|
||||
width: f64,
|
||||
height: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct RawAnnotation {
|
||||
ref_id: String,
|
||||
number: u64,
|
||||
role: String,
|
||||
name: Option<String>,
|
||||
rect: Rect,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct AnnotationBox {
|
||||
pub x: i64,
|
||||
pub y: i64,
|
||||
pub width: i64,
|
||||
pub height: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ScreenshotAnnotation {
|
||||
pub ref_id: String,
|
||||
pub number: u64,
|
||||
pub role: String,
|
||||
pub name: Option<String>,
|
||||
pub box_: AnnotationBox,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ScreenshotResult {
|
||||
pub path: String,
|
||||
pub base64: String,
|
||||
pub annotations: Vec<ScreenshotAnnotation>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ScreenshotOptions {
|
||||
pub selector: Option<String>,
|
||||
pub path: Option<String>,
|
||||
pub full_page: bool,
|
||||
pub format: String,
|
||||
pub quality: Option<i32>,
|
||||
pub annotate: bool,
|
||||
pub output_dir: Option<String>,
|
||||
}
|
||||
|
||||
impl Default for ScreenshotOptions {
|
||||
@@ -70,111 +21,16 @@ impl Default for ScreenshotOptions {
|
||||
full_page: false,
|
||||
format: "png".to_string(),
|
||||
quality: None,
|
||||
annotate: false,
|
||||
output_dir: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Serialize for ScreenshotAnnotation {
|
||||
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
|
||||
where
|
||||
S: serde::Serializer,
|
||||
{
|
||||
use serde::ser::SerializeStruct;
|
||||
|
||||
let mut state = serializer.serialize_struct("ScreenshotAnnotation", 5)?;
|
||||
state.serialize_field("ref", &self.ref_id)?;
|
||||
state.serialize_field("number", &self.number)?;
|
||||
state.serialize_field("role", &self.role)?;
|
||||
if let Some(name) = &self.name {
|
||||
state.serialize_field("name", name)?;
|
||||
}
|
||||
state.serialize_field("box", &self.box_)?;
|
||||
state.end()
|
||||
}
|
||||
}
|
||||
|
||||
/// Captures a screenshot via CDP and optionally overlays numbered annotations
|
||||
/// that mirror the Node.js screenshot `annotate` mode.
|
||||
pub async fn take_screenshot(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
options: &ScreenshotOptions,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<ScreenshotResult, String> {
|
||||
let target_rect = if options.annotate {
|
||||
match options.selector.as_deref() {
|
||||
Some(selector) => {
|
||||
get_rect_for_selector(client, session_id, ref_map, selector, iframe_sessions)
|
||||
.await?
|
||||
}
|
||||
None => None,
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let raw_annotations = if options.annotate {
|
||||
collect_annotations(client, session_id, ref_map).await?
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
let overlay_items = filter_annotations(raw_annotations, target_rect.as_ref());
|
||||
let overlay_injected = if options.annotate && !overlay_items.is_empty() {
|
||||
inject_annotation_overlay(client, session_id, &overlay_items).await?;
|
||||
true
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
let base64 =
|
||||
capture_screenshot_base64(client, session_id, ref_map, options, iframe_sessions).await;
|
||||
|
||||
if overlay_injected {
|
||||
let _ = remove_annotation_overlay(client, session_id).await;
|
||||
}
|
||||
|
||||
let base64 = base64?;
|
||||
let annotations = if options.annotate {
|
||||
let scroll = if options.full_page {
|
||||
Some(get_scroll_offsets(client, session_id).await?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
project_annotations(&overlay_items, target_rect.as_ref(), scroll)
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
|
||||
let ext = if options.format == "jpeg" {
|
||||
"jpg"
|
||||
} else {
|
||||
"png"
|
||||
};
|
||||
let path = save_screenshot(
|
||||
&base64,
|
||||
options.path.as_deref(),
|
||||
ext,
|
||||
options.output_dir.as_deref(),
|
||||
)?;
|
||||
|
||||
Ok(ScreenshotResult {
|
||||
path,
|
||||
base64,
|
||||
annotations,
|
||||
})
|
||||
}
|
||||
|
||||
async fn capture_screenshot_base64(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
options: &ScreenshotOptions,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<String, String> {
|
||||
) -> Result<(String, String), String> {
|
||||
let mut params = CaptureScreenshotParams {
|
||||
format: Some(options.format.clone()),
|
||||
quality: if options.format == "jpeg" {
|
||||
@@ -208,14 +64,40 @@ async fn capture_screenshot_base64(
|
||||
});
|
||||
}
|
||||
} else if let Some(ref selector) = options.selector {
|
||||
if let Some(rect) =
|
||||
get_rect_for_selector(client, session_id, ref_map, selector, iframe_sessions).await?
|
||||
{
|
||||
// Element screenshot via bounding box
|
||||
let object_id =
|
||||
super::element::resolve_element_object_id(client, session_id, ref_map, selector)
|
||||
.await?;
|
||||
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
"Runtime.callFunctionOn",
|
||||
&CallFunctionOnParams {
|
||||
function_declaration: r#"function() {
|
||||
const rect = this.getBoundingClientRect();
|
||||
return { x: rect.x, y: rect.y, width: rect.width, height: rect.height };
|
||||
}"#
|
||||
.to_string(),
|
||||
object_id: Some(object_id),
|
||||
arguments: None,
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
if let Some(rect) = result.result.value {
|
||||
let x = rect.get("x").and_then(|v| v.as_f64()).unwrap_or(0.0);
|
||||
let y = rect.get("y").and_then(|v| v.as_f64()).unwrap_or(0.0);
|
||||
let w = rect.get("width").and_then(|v| v.as_f64()).unwrap_or(100.0);
|
||||
let h = rect.get("height").and_then(|v| v.as_f64()).unwrap_or(100.0);
|
||||
|
||||
params.clip = Some(Viewport {
|
||||
x: rect.x,
|
||||
y: rect.y,
|
||||
width: rect.width,
|
||||
height: rect.height,
|
||||
x,
|
||||
y,
|
||||
width: w,
|
||||
height: h,
|
||||
scale: 1.0,
|
||||
});
|
||||
}
|
||||
@@ -225,345 +107,16 @@ async fn capture_screenshot_base64(
|
||||
.send_command_typed("Page.captureScreenshot", ¶ms, Some(session_id))
|
||||
.await?;
|
||||
|
||||
Ok(result.data)
|
||||
}
|
||||
let ext = if options.format == "jpeg" {
|
||||
"jpg"
|
||||
} else {
|
||||
"png"
|
||||
};
|
||||
|
||||
async fn collect_annotations(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
) -> Result<Vec<RawAnnotation>, String> {
|
||||
let entries = ref_map.entries_sorted();
|
||||
if entries.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
// Collect entries that have backend_node_ids for batch resolution.
|
||||
let with_backend_ids: Vec<(String, super::element::RefEntry, i64)> = entries
|
||||
.iter()
|
||||
.filter_map(|(ref_id, entry)| {
|
||||
entry
|
||||
.backend_node_id
|
||||
.map(|bid| (ref_id.clone(), entry.clone(), bid))
|
||||
})
|
||||
.collect();
|
||||
|
||||
if with_backend_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
// Batch-resolve all backend_node_ids to object IDs using concurrent CDP calls.
|
||||
let resolve_futures: Vec<_> = with_backend_ids
|
||||
.iter()
|
||||
.map(|(_, _, backend_node_id)| {
|
||||
client.send_command(
|
||||
"DOM.resolveNode",
|
||||
Some(serde_json::json!({
|
||||
"backendNodeId": backend_node_id,
|
||||
"objectGroup": "agent-browser-annotate"
|
||||
})),
|
||||
Some(session_id),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
|
||||
let resolve_results = futures_util::future::join_all(resolve_futures).await;
|
||||
|
||||
// Collect resolved object IDs paired with their ref info.
|
||||
let mut resolved: Vec<(String, super::element::RefEntry, String)> = Vec::new();
|
||||
for (i, result) in resolve_results.into_iter().enumerate() {
|
||||
if let Ok(val) = result {
|
||||
if let Some(oid) = val
|
||||
.get("object")
|
||||
.and_then(|o| o.get("objectId"))
|
||||
.and_then(|v| v.as_str())
|
||||
{
|
||||
let (ref_id, entry, _) = &with_backend_ids[i];
|
||||
resolved.push((ref_id.clone(), entry.clone(), oid.to_string()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if resolved.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
// Batch-get bounding rects for all resolved elements using concurrent CDP calls.
|
||||
let rect_futures: Vec<_> = resolved
|
||||
.iter()
|
||||
.map(|(_, _, object_id)| get_rect_for_object(client, session_id, object_id))
|
||||
.collect();
|
||||
|
||||
let rect_results = futures_util::future::join_all(rect_futures).await;
|
||||
|
||||
let mut annotations = Vec::new();
|
||||
for (i, rect_result) in rect_results.into_iter().enumerate() {
|
||||
let rect = match rect_result {
|
||||
Ok(Some(r)) if r.width > 0.0 && r.height > 0.0 => r,
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
let (ref_id, entry, _) = &resolved[i];
|
||||
let number = ref_id
|
||||
.strip_prefix('e')
|
||||
.and_then(|n| n.parse::<u64>().ok())
|
||||
.unwrap_or(0);
|
||||
|
||||
annotations.push(RawAnnotation {
|
||||
ref_id: ref_id.clone(),
|
||||
number,
|
||||
role: entry.role.clone(),
|
||||
name: (!entry.name.is_empty()).then_some(entry.name.clone()),
|
||||
rect,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(annotations)
|
||||
}
|
||||
|
||||
async fn get_rect_for_selector(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
ref_map: &RefMap,
|
||||
selector: &str,
|
||||
iframe_sessions: &HashMap<String, String>,
|
||||
) -> Result<Option<Rect>, String> {
|
||||
let (object_id, effective_session_id) = super::element::resolve_element_object_id(
|
||||
client,
|
||||
session_id,
|
||||
ref_map,
|
||||
selector,
|
||||
iframe_sessions,
|
||||
)
|
||||
.await?;
|
||||
get_rect_for_object(client, &effective_session_id, &object_id).await
|
||||
}
|
||||
|
||||
async fn get_rect_for_object(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
object_id: &str,
|
||||
) -> Result<Option<Rect>, String> {
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
"Runtime.callFunctionOn",
|
||||
&CallFunctionOnParams {
|
||||
function_declaration: r#"function() {
|
||||
const rect = this.getBoundingClientRect();
|
||||
return { x: rect.x, y: rect.y, width: rect.width, height: rect.height };
|
||||
}"#
|
||||
.to_string(),
|
||||
object_id: Some(object_id.to_string()),
|
||||
arguments: None,
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(result.result.value.as_ref().and_then(parse_rect))
|
||||
}
|
||||
|
||||
fn parse_rect(value: &Value) -> Option<Rect> {
|
||||
Some(Rect {
|
||||
x: value.get("x")?.as_f64()?,
|
||||
y: value.get("y")?.as_f64()?,
|
||||
width: value.get("width")?.as_f64()?,
|
||||
height: value.get("height")?.as_f64()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn filter_annotations(
|
||||
annotations: Vec<RawAnnotation>,
|
||||
target_rect: Option<&Rect>,
|
||||
) -> Vec<RawAnnotation> {
|
||||
let mut items = annotations
|
||||
.into_iter()
|
||||
.filter(|annotation| match target_rect {
|
||||
Some(target) => overlaps(&annotation.rect, target),
|
||||
None => true,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
items.sort_by_key(|annotation| annotation.number);
|
||||
items
|
||||
}
|
||||
|
||||
fn overlaps(left: &Rect, right: &Rect) -> bool {
|
||||
let left_x2 = left.x + left.width;
|
||||
let left_y2 = left.y + left.height;
|
||||
let right_x2 = right.x + right.width;
|
||||
let right_y2 = right.y + right.height;
|
||||
|
||||
left.x < right_x2 && left_x2 > right.x && left.y < right_y2 && left_y2 > right.y
|
||||
}
|
||||
|
||||
async fn inject_annotation_overlay(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
annotations: &[RawAnnotation],
|
||||
) -> Result<(), String> {
|
||||
let overlay_data = annotations
|
||||
.iter()
|
||||
.map(|annotation| {
|
||||
serde_json::json!({
|
||||
"number": annotation.number,
|
||||
"x": round(annotation.rect.x),
|
||||
"y": round(annotation.rect.y),
|
||||
"width": round(annotation.rect.width),
|
||||
"height": round(annotation.rect.height),
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let expression = format!(
|
||||
r#"(() => {{
|
||||
var items = {items};
|
||||
var id = {overlay_id};
|
||||
var existing = document.getElementById(id);
|
||||
if (existing) existing.remove();
|
||||
var sx = window.scrollX || 0;
|
||||
var sy = window.scrollY || 0;
|
||||
var c = document.createElement('div');
|
||||
c.id = id;
|
||||
c.style.cssText = 'position:absolute;top:0;left:0;width:0;height:0;pointer-events:none;z-index:2147483647;';
|
||||
for (var i = 0; i < items.length; i++) {{
|
||||
var it = items[i];
|
||||
var dx = it.x + sx;
|
||||
var dy = it.y + sy;
|
||||
var b = document.createElement('div');
|
||||
b.style.cssText = 'position:absolute;left:' + dx + 'px;top:' + dy + 'px;width:' + it.width + 'px;height:' + it.height + 'px;border:2px solid rgba(255,0,0,0.8);box-sizing:border-box;pointer-events:none;';
|
||||
var l = document.createElement('div');
|
||||
l.textContent = String(it.number);
|
||||
var labelTop = dy < 14 ? '2px' : '-14px';
|
||||
l.style.cssText = 'position:absolute;top:' + labelTop + ';left:-2px;background:rgba(255,0,0,0.9);color:#fff;font:bold 11px/14px monospace;padding:0 4px;border-radius:2px;white-space:nowrap;';
|
||||
b.appendChild(l);
|
||||
c.appendChild(b);
|
||||
}}
|
||||
document.documentElement.appendChild(c);
|
||||
return true;
|
||||
}})()"#,
|
||||
items = serde_json::to_string(&overlay_data).unwrap_or_else(|_| "[]".to_string()),
|
||||
overlay_id =
|
||||
serde_json::to_string(ANNOTATION_OVERLAY_ID).unwrap_or_else(|_| "\"\"".to_string()),
|
||||
);
|
||||
|
||||
let _: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
"Runtime.evaluate",
|
||||
&EvaluateParams {
|
||||
expression,
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn remove_annotation_overlay(client: &CdpClient, session_id: &str) -> Result<(), String> {
|
||||
let expression = format!(
|
||||
r#"(() => {{
|
||||
var el = document.getElementById({overlay_id});
|
||||
if (el) el.remove();
|
||||
return true;
|
||||
}})()"#,
|
||||
overlay_id =
|
||||
serde_json::to_string(ANNOTATION_OVERLAY_ID).unwrap_or_else(|_| "\"\"".to_string()),
|
||||
);
|
||||
|
||||
let _: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
"Runtime.evaluate",
|
||||
&EvaluateParams {
|
||||
expression,
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_scroll_offsets(client: &CdpClient, session_id: &str) -> Result<(f64, f64), String> {
|
||||
let result: EvaluateResult = client
|
||||
.send_command_typed(
|
||||
"Runtime.evaluate",
|
||||
&EvaluateParams {
|
||||
expression: "({x: window.scrollX || 0, y: window.scrollY || 0})".to_string(),
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let value = result.result.value.unwrap_or(Value::Null);
|
||||
let x = value.get("x").and_then(|v| v.as_f64()).unwrap_or(0.0);
|
||||
let y = value.get("y").and_then(|v| v.as_f64()).unwrap_or(0.0);
|
||||
Ok((x, y))
|
||||
}
|
||||
|
||||
fn project_annotations(
|
||||
annotations: &[RawAnnotation],
|
||||
target_rect: Option<&Rect>,
|
||||
scroll: Option<(f64, f64)>,
|
||||
) -> Vec<ScreenshotAnnotation> {
|
||||
annotations
|
||||
.iter()
|
||||
.map(|annotation| {
|
||||
let rect = if let Some(target) = target_rect {
|
||||
Rect {
|
||||
x: annotation.rect.x - target.x,
|
||||
y: annotation.rect.y - target.y,
|
||||
width: annotation.rect.width,
|
||||
height: annotation.rect.height,
|
||||
}
|
||||
} else if let Some((scroll_x, scroll_y)) = scroll {
|
||||
Rect {
|
||||
x: annotation.rect.x + scroll_x,
|
||||
y: annotation.rect.y + scroll_y,
|
||||
width: annotation.rect.width,
|
||||
height: annotation.rect.height,
|
||||
}
|
||||
} else {
|
||||
annotation.rect.clone()
|
||||
};
|
||||
|
||||
ScreenshotAnnotation {
|
||||
ref_id: annotation.ref_id.clone(),
|
||||
number: annotation.number,
|
||||
role: annotation.role.clone(),
|
||||
name: annotation.name.clone(),
|
||||
box_: AnnotationBox {
|
||||
x: round(rect.x),
|
||||
y: round(rect.y),
|
||||
width: round(rect.width),
|
||||
height: round(rect.height),
|
||||
},
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn save_screenshot(
|
||||
base64_data: &str,
|
||||
explicit_path: Option<&str>,
|
||||
ext: &str,
|
||||
output_dir: Option<&str>,
|
||||
) -> Result<String, String> {
|
||||
let save_path = match explicit_path {
|
||||
Some(path) => path.to_string(),
|
||||
let save_path = match &options.path {
|
||||
Some(p) => p.clone(),
|
||||
None => {
|
||||
let dir = match output_dir {
|
||||
Some(d) => PathBuf::from(d),
|
||||
None => get_screenshot_dir(),
|
||||
};
|
||||
let dir = get_screenshot_dir();
|
||||
let _ = std::fs::create_dir_all(&dir);
|
||||
let timestamp = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
@@ -574,17 +127,13 @@ fn save_screenshot(
|
||||
}
|
||||
};
|
||||
|
||||
let bytes = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, base64_data)
|
||||
let bytes = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, &result.data)
|
||||
.map_err(|e| format!("Failed to decode screenshot: {}", e))?;
|
||||
|
||||
std::fs::write(&save_path, &bytes)
|
||||
.map_err(|e| format!("Failed to save screenshot to {}: {}", save_path, e))?;
|
||||
|
||||
Ok(save_path)
|
||||
}
|
||||
|
||||
fn round(value: f64) -> i64 {
|
||||
value.round() as i64
|
||||
Ok((save_path, result.data))
|
||||
}
|
||||
|
||||
fn get_screenshot_dir() -> PathBuf {
|
||||
@@ -596,96 +145,3 @@ fn get_screenshot_dir() -> PathBuf {
|
||||
.join("screenshots")
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn filters_annotations_to_target_overlap() {
|
||||
let annotations = vec![
|
||||
RawAnnotation {
|
||||
ref_id: "e1".to_string(),
|
||||
number: 1,
|
||||
role: "button".to_string(),
|
||||
name: Some("Inside".to_string()),
|
||||
rect: Rect {
|
||||
x: 10.0,
|
||||
y: 10.0,
|
||||
width: 50.0,
|
||||
height: 20.0,
|
||||
},
|
||||
},
|
||||
RawAnnotation {
|
||||
ref_id: "e2".to_string(),
|
||||
number: 2,
|
||||
role: "button".to_string(),
|
||||
name: Some("Outside".to_string()),
|
||||
rect: Rect {
|
||||
x: 200.0,
|
||||
y: 200.0,
|
||||
width: 40.0,
|
||||
height: 20.0,
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
let target = Rect {
|
||||
x: 0.0,
|
||||
y: 0.0,
|
||||
width: 100.0,
|
||||
height: 100.0,
|
||||
};
|
||||
|
||||
let filtered = filter_annotations(annotations, Some(&target));
|
||||
assert_eq!(filtered.len(), 1);
|
||||
assert_eq!(filtered[0].ref_id, "e1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn projects_selector_annotations_relative_to_target() {
|
||||
let annotations = vec![RawAnnotation {
|
||||
ref_id: "e1".to_string(),
|
||||
number: 1,
|
||||
role: "button".to_string(),
|
||||
name: Some("Inside".to_string()),
|
||||
rect: Rect {
|
||||
x: 25.0,
|
||||
y: 35.0,
|
||||
width: 40.0,
|
||||
height: 20.0,
|
||||
},
|
||||
}];
|
||||
|
||||
let target = Rect {
|
||||
x: 10.0,
|
||||
y: 15.0,
|
||||
width: 100.0,
|
||||
height: 100.0,
|
||||
};
|
||||
|
||||
let projected = project_annotations(&annotations, Some(&target), None);
|
||||
assert_eq!(projected[0].box_.x, 15);
|
||||
assert_eq!(projected[0].box_.y, 20);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn projects_full_page_annotations_to_document_space() {
|
||||
let annotations = vec![RawAnnotation {
|
||||
ref_id: "e1".to_string(),
|
||||
number: 1,
|
||||
role: "button".to_string(),
|
||||
name: Some("Bottom".to_string()),
|
||||
rect: Rect {
|
||||
x: 5.0,
|
||||
y: 12.0,
|
||||
width: 40.0,
|
||||
height: 20.0,
|
||||
},
|
||||
}];
|
||||
|
||||
let projected = project_annotations(&annotations, None, Some((10.0, 1000.0)));
|
||||
assert_eq!(projected[0].box_.x, 15);
|
||||
assert_eq!(projected[0].box_.y, 1012);
|
||||
}
|
||||
}
|
||||
|
||||
+143
-744
File diff suppressed because it is too large
Load Diff
+41
-321
@@ -1,17 +1,12 @@
|
||||
use aes_gcm::{aead::Aead, aead::KeyInit, Aes256Gcm};
|
||||
use base64::Engine;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Value};
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::collections::HashSet;
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use super::cdp::client::CdpClient;
|
||||
use super::cdp::types::{
|
||||
AttachToTargetParams, AttachToTargetResult, CloseTargetParams, CreateTargetParams,
|
||||
CreateTargetResult, EvaluateParams,
|
||||
};
|
||||
use super::cdp::types::EvaluateParams;
|
||||
use super::cookies::{self, Cookie};
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
@@ -37,223 +32,16 @@ pub struct StorageEntry {
|
||||
pub value: String,
|
||||
}
|
||||
|
||||
fn collect_frame_origins(tree: &Value, origins: &mut HashSet<String>) {
|
||||
if let Some(frame) = tree.get("frame") {
|
||||
if let Some(url_str) = frame.get("url").and_then(|v| v.as_str()) {
|
||||
if let Ok(parsed) = url::Url::parse(url_str) {
|
||||
let origin = parsed.origin().ascii_serialization();
|
||||
if origin != "null" && !origin.is_empty() {
|
||||
origins.insert(origin);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(children) = tree.get("childFrames").and_then(|v| v.as_array()) {
|
||||
for child in children {
|
||||
collect_frame_origins(child, origins);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse the JS-evaluated origin storage data into an OriginStorage struct.
|
||||
fn parse_origin_storage(data: &Value) -> Option<OriginStorage> {
|
||||
if !data.is_object() {
|
||||
return None;
|
||||
}
|
||||
let origin = data
|
||||
.get("origin")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
if origin.is_empty() || origin == "null" {
|
||||
return None;
|
||||
}
|
||||
let local_storage: Vec<StorageEntry> = data
|
||||
.get("localStorage")
|
||||
.and_then(|v| serde_json::from_value(v.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
let session_storage: Vec<StorageEntry> = data
|
||||
.get("sessionStorage")
|
||||
.and_then(|v| serde_json::from_value(v.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
|
||||
Some(OriginStorage {
|
||||
origin,
|
||||
local_storage,
|
||||
session_storage,
|
||||
})
|
||||
}
|
||||
|
||||
/// Evaluate the storage-collection JS snippet and parse the result.
|
||||
async fn eval_origin_storage(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
origin_js: &str,
|
||||
) -> Option<OriginStorage> {
|
||||
let result = client
|
||||
.send_command_typed::<_, super::cdp::types::EvaluateResult>(
|
||||
"Runtime.evaluate",
|
||||
&EvaluateParams {
|
||||
expression: origin_js.to_string(),
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
.await
|
||||
.ok()?;
|
||||
let data = result.result.value.unwrap_or(Value::Null);
|
||||
parse_origin_storage(&data)
|
||||
}
|
||||
|
||||
/// Create a temporary CDP target, navigate it to each origin to collect localStorage,
|
||||
/// then close it. Uses Fetch interception to serve blank HTML instead of making real
|
||||
/// network requests.
|
||||
async fn collect_storage_via_temp_target(
|
||||
client: &CdpClient,
|
||||
origins: &[String],
|
||||
origin_js: &str,
|
||||
) -> Result<Vec<OriginStorage>, String> {
|
||||
let create_result: CreateTargetResult = client
|
||||
.send_command_typed(
|
||||
"Target.createTarget",
|
||||
&CreateTargetParams {
|
||||
url: "about:blank".to_string(),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let target_id = create_result.target_id;
|
||||
|
||||
// Ensure the target is closed even if attach or later steps fail
|
||||
let result = collect_storage_in_target(client, &target_id, origins, origin_js).await;
|
||||
|
||||
let _ = client
|
||||
.send_command_typed::<_, Value>(
|
||||
"Target.closeTarget",
|
||||
&CloseTargetParams { target_id },
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
async fn collect_storage_in_target(
|
||||
client: &CdpClient,
|
||||
target_id: &str,
|
||||
origins: &[String],
|
||||
origin_js: &str,
|
||||
) -> Result<Vec<OriginStorage>, String> {
|
||||
let attach_result: AttachToTargetResult = client
|
||||
.send_command_typed(
|
||||
"Target.attachToTarget",
|
||||
&AttachToTargetParams {
|
||||
target_id: target_id.to_string(),
|
||||
flatten: true,
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let temp_session = &attach_result.session_id;
|
||||
|
||||
client
|
||||
.send_command_no_params("Page.enable", Some(temp_session))
|
||||
.await?;
|
||||
client
|
||||
.send_command_no_params("Runtime.enable", Some(temp_session))
|
||||
.await?;
|
||||
|
||||
// Blank HTML response body, pre-encoded to avoid repeated base64 work per request
|
||||
let blank_html_b64 = base64::engine::general_purpose::STANDARD.encode("<html></html>");
|
||||
|
||||
let _ = client
|
||||
.send_command(
|
||||
"Fetch.enable",
|
||||
Some(json!({ "patterns": [{ "urlPattern": "*" }] })),
|
||||
Some(temp_session),
|
||||
)
|
||||
.await;
|
||||
|
||||
let mut event_rx = client.subscribe();
|
||||
let mut results = Vec::new();
|
||||
|
||||
for target_origin in origins {
|
||||
let nav_url = format!("{}/", target_origin.trim_end_matches('/'));
|
||||
if client
|
||||
.send_command(
|
||||
"Page.navigate",
|
||||
Some(json!({ "url": nav_url })),
|
||||
Some(temp_session),
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
// Fulfill intercepted requests with blank HTML until the page loads
|
||||
let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(5);
|
||||
let mut page_loaded = false;
|
||||
while tokio::time::Instant::now() < deadline {
|
||||
match tokio::time::timeout(tokio::time::Duration::from_secs(2), event_rx.recv()).await {
|
||||
Ok(Ok(evt)) if evt.session_id.as_deref() == Some(temp_session) => {
|
||||
if evt.method == "Fetch.requestPaused" {
|
||||
if let Some(request_id) =
|
||||
evt.params.get("requestId").and_then(|v| v.as_str())
|
||||
{
|
||||
let _ = client
|
||||
.send_command(
|
||||
"Fetch.fulfillRequest",
|
||||
Some(json!({
|
||||
"requestId": request_id,
|
||||
"responseCode": 200,
|
||||
"responseHeaders": [
|
||||
{ "name": "Content-Type", "value": "text/html" }
|
||||
],
|
||||
"body": &blank_html_b64
|
||||
})),
|
||||
Some(temp_session),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
} else if evt.method == "Page.loadEventFired" {
|
||||
page_loaded = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(Ok(_)) => continue, // event for a different session
|
||||
Ok(Err(_)) => continue, // lagged or closed — retry within deadline
|
||||
Err(_) => break, // outer timeout elapsed
|
||||
}
|
||||
}
|
||||
|
||||
if !page_loaded {
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Some(storage) = eval_origin_storage(client, temp_session, origin_js).await {
|
||||
if !storage.local_storage.is_empty() || !storage.session_storage.is_empty() {
|
||||
results.push(storage);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
pub async fn save_state(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
path: Option<&str>,
|
||||
session_name: Option<&str>,
|
||||
session_id_str: &str,
|
||||
visited_origins: &HashSet<String>,
|
||||
) -> Result<String, String> {
|
||||
let cookies = cookies::get_all_cookies(client, session_id).await?;
|
||||
let cookies = cookies::get_cookies(client, session_id, None).await?;
|
||||
|
||||
// Get current origin's storage
|
||||
let origin_js = r#"(() => {
|
||||
const result = { origin: location.origin, localStorage: [], sessionStorage: [] };
|
||||
try {
|
||||
@@ -271,38 +59,46 @@ pub async fn save_state(
|
||||
return result;
|
||||
})()"#;
|
||||
|
||||
// Merge visited origins with current frame tree origins
|
||||
let mut all_origins = visited_origins.clone();
|
||||
if let Ok(tree_result) = client
|
||||
.send_command_no_params("Page.getFrameTree", Some(session_id))
|
||||
.await
|
||||
{
|
||||
if let Some(tree) = tree_result.get("frameTree") {
|
||||
collect_frame_origins(tree, &mut all_origins);
|
||||
}
|
||||
}
|
||||
let origin_result: super::cdp::types::EvaluateResult = client
|
||||
.send_command_typed(
|
||||
"Runtime.evaluate",
|
||||
&EvaluateParams {
|
||||
expression: origin_js.to_string(),
|
||||
return_by_value: Some(true),
|
||||
await_promise: Some(false),
|
||||
},
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
// 1. Collect localStorage from the current page
|
||||
let mut origins = Vec::new();
|
||||
let mut current_origin = String::new();
|
||||
let origin_data = origin_result.result.value.unwrap_or(Value::Null);
|
||||
let origins = if origin_data.is_object() {
|
||||
let origin = origin_data
|
||||
.get("origin")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
let local_storage: Vec<StorageEntry> = origin_data
|
||||
.get("localStorage")
|
||||
.and_then(|v| serde_json::from_value(v.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
let session_storage: Vec<StorageEntry> = origin_data
|
||||
.get("sessionStorage")
|
||||
.and_then(|v| serde_json::from_value(v.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
|
||||
if let Some(storage) = eval_origin_storage(client, session_id, origin_js).await {
|
||||
current_origin = storage.origin.clone();
|
||||
if !storage.local_storage.is_empty() || !storage.session_storage.is_empty() {
|
||||
origins.push(storage);
|
||||
if !origin.is_empty() && origin != "null" {
|
||||
vec![OriginStorage {
|
||||
origin,
|
||||
local_storage,
|
||||
session_storage,
|
||||
}]
|
||||
} else {
|
||||
vec![]
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Collect localStorage from remaining origins via a disposable temp target
|
||||
all_origins.remove(¤t_origin);
|
||||
if !all_origins.is_empty() {
|
||||
let remaining: Vec<String> = all_origins.into_iter().collect();
|
||||
if let Ok(temp_origins) =
|
||||
collect_storage_via_temp_target(client, &remaining, origin_js).await
|
||||
{
|
||||
origins.extend(temp_origins);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
vec![]
|
||||
};
|
||||
|
||||
let state = StorageState { cookies, origins };
|
||||
let json_str = serde_json::to_string_pretty(&state)
|
||||
@@ -671,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().is_none_or(|(_, t)| modified > *t) {
|
||||
if best_path.as_ref().map_or(true, |(_, t)| modified > *t) {
|
||||
best_path = Some((path.to_string_lossy().to_string(), modified));
|
||||
}
|
||||
}
|
||||
@@ -679,41 +475,6 @@ pub fn find_auto_state_file(session_name: &str) -> Option<String> {
|
||||
best_path.map(|(p, _)| p)
|
||||
}
|
||||
|
||||
/// Dispatch a state management command from its JSON payload.
|
||||
/// Returns `Some(result)` for recognised state_* actions, `None` otherwise.
|
||||
pub fn dispatch_state_command(cmd: &Value) -> Option<Result<Value, String>> {
|
||||
let action = cmd.get("action").and_then(|v| v.as_str())?;
|
||||
match action {
|
||||
"state_list" => Some(state_list()),
|
||||
"state_show" => Some(
|
||||
cmd.get("path")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| "Missing 'path' parameter".to_string())
|
||||
.and_then(state_show),
|
||||
),
|
||||
"state_clear" => {
|
||||
let path = cmd.get("path").and_then(|v| v.as_str());
|
||||
Some(state_clear(path))
|
||||
}
|
||||
"state_clean" => {
|
||||
let days = cmd.get("days").and_then(|v| v.as_u64()).unwrap_or(30);
|
||||
Some(state_clean(days))
|
||||
}
|
||||
"state_rename" => Some(
|
||||
cmd.get("path")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| "Missing 'path' parameter".to_string())
|
||||
.and_then(|path| {
|
||||
cmd.get("name")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| "Missing 'name' parameter".to_string())
|
||||
.and_then(|name| state_rename(path, name))
|
||||
}),
|
||||
),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_sessions_dir() -> PathBuf {
|
||||
if let Some(home) = dirs::home_dir() {
|
||||
home.join(".agent-browser").join("sessions")
|
||||
@@ -843,45 +604,4 @@ mod tests {
|
||||
assert_eq!(json["secure"], true);
|
||||
assert_eq!(json["sameSite"], "Strict");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dispatch_state_command_routes_state_list() {
|
||||
let cmd = serde_json::json!({ "action": "state_list" });
|
||||
let result = dispatch_state_command(&cmd);
|
||||
assert!(result.is_some());
|
||||
assert!(result.unwrap().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dispatch_state_command_returns_none_for_unknown() {
|
||||
let cmd = serde_json::json!({ "action": "navigate" });
|
||||
assert!(dispatch_state_command(&cmd).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dispatch_state_command_returns_none_for_missing_action() {
|
||||
let cmd = serde_json::json!({});
|
||||
assert!(dispatch_state_command(&cmd).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dispatch_state_show_missing_path() {
|
||||
let cmd = serde_json::json!({ "action": "state_show" });
|
||||
let result = dispatch_state_command(&cmd).unwrap();
|
||||
assert!(result.is_err());
|
||||
assert_eq!(result.unwrap_err(), "Missing 'path' parameter");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dispatch_state_rename_missing_params() {
|
||||
let cmd = serde_json::json!({ "action": "state_rename" });
|
||||
let result = dispatch_state_command(&cmd).unwrap();
|
||||
assert!(result.is_err());
|
||||
assert_eq!(result.unwrap_err(), "Missing 'path' parameter");
|
||||
|
||||
let cmd = serde_json::json!({ "action": "state_rename", "path": "/tmp/test.json" });
|
||||
let result = dispatch_state_command(&cmd).unwrap();
|
||||
assert!(result.is_err());
|
||||
assert_eq!(result.unwrap_err(), "Missing 'name' parameter");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,237 +0,0 @@
|
||||
//! Stealth anti-detection module.
|
||||
//!
|
||||
//! Injects browser-level patches to evade bot detection (creepjs, sannysoft,
|
||||
//! Cloudflare Turnstile, etc.) by normalizing fingerprint signals that betray
|
||||
//! headless or automated Chrome instances.
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use super::cdp::client::CdpClient;
|
||||
|
||||
/// Full stealth JS payload compiled at build time (for --launch mode).
|
||||
const STEALTH_SCRIPTS_RAW: &str = include_str!("stealth_scripts.js");
|
||||
|
||||
/// Minimal stealth script for CDP-attach mode (connecting to user's real Chrome).
|
||||
/// Only removes navigator.webdriver — the browser's own fingerprint is already real.
|
||||
/// Minimal stealth script for CDP-attach mode.
|
||||
/// Emulation.setAutomationOverride handles navigator.webdriver at the native
|
||||
/// level, so no JS patching is needed in CdpAttach mode. An empty script
|
||||
/// avoids creating any detectable lie-props artifacts.
|
||||
const MINIMAL_STEALTH_SCRIPT: &str = "";
|
||||
|
||||
/// Chrome launch arguments that reduce automation fingerprint surface.
|
||||
pub const STEALTH_CHROMIUM_ARGS: &[&str] = &[
|
||||
"--disable-blink-features=AutomationControlled",
|
||||
"--use-gl=angle",
|
||||
"--use-angle=default",
|
||||
];
|
||||
|
||||
/// Connection mode determines which stealth patches to apply.
|
||||
#[derive(Clone, Copy, PartialEq)]
|
||||
pub enum StealthMode {
|
||||
/// Connected to user's real Chrome — minimal patches only (webdriver removal).
|
||||
/// The browser already has a real fingerprint; heavy patches would create detectable lies.
|
||||
CdpAttach,
|
||||
/// Launched a new Chrome instance — apply full stealth patches.
|
||||
FullLaunch,
|
||||
}
|
||||
|
||||
/// Build the stealth JS payload for the given mode and locale.
|
||||
pub fn build_stealth_script(mode: StealthMode, locale: Option<&str>) -> String {
|
||||
if mode == StealthMode::CdpAttach {
|
||||
return MINIMAL_STEALTH_SCRIPT.to_string();
|
||||
}
|
||||
|
||||
// Full launch mode: inject all patches
|
||||
let locale = locale.unwrap_or("en-US");
|
||||
let base_lang = locale.split('-').next().unwrap_or(locale);
|
||||
let languages: Vec<&str> = if base_lang == locale {
|
||||
vec![locale]
|
||||
} else {
|
||||
vec![locale, base_lang]
|
||||
};
|
||||
let config_line = format!(
|
||||
r#"const __abStealth = {{ locale: "{}", languages: {}, allowWebGLContextFallback: false }};"#,
|
||||
locale,
|
||||
serde_json::to_string(&languages).unwrap_or_else(|_| r#"["en-US","en"]"#.to_string()),
|
||||
);
|
||||
|
||||
if let Some(rest) = STEALTH_SCRIPTS_RAW.strip_prefix(
|
||||
r#"const __abStealth = { locale: "en-US", languages: ["en-US", "en"], allowWebGLContextFallback: false };"#,
|
||||
) {
|
||||
format!("{}{}", config_line, rest)
|
||||
} else {
|
||||
format!("{}\n{}", config_line, STEALTH_SCRIPTS_RAW)
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply stealth patches to a browser session.
|
||||
///
|
||||
/// In `CdpAttach` mode (user's real Chrome): only removes `navigator.webdriver`.
|
||||
/// In `FullLaunch` mode (new Chrome): injects all 32 patches + UA override.
|
||||
pub async fn apply_stealth(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
mode: StealthMode,
|
||||
locale: Option<&str>,
|
||||
) -> Result<(), String> {
|
||||
// First: disable the automation flag at the CDP protocol level.
|
||||
// This tells Chrome to natively set navigator.webdriver = false,
|
||||
// which is undetectable by lie-detection systems like CreepJS.
|
||||
// Falls back gracefully on older Chrome versions that don't support this.
|
||||
let _ = client
|
||||
.send_command(
|
||||
"Emulation.setAutomationOverride",
|
||||
Some(json!({ "enabled": false })),
|
||||
Some(session_id),
|
||||
)
|
||||
.await;
|
||||
|
||||
let script = build_stealth_script(mode, locale);
|
||||
|
||||
// Inject stealth scripts to run before page JS
|
||||
client
|
||||
.send_command(
|
||||
"Page.addScriptToEvaluateOnNewDocument",
|
||||
Some(json!({ "source": script })),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
|
||||
// In full launch mode, also override UA to remove HeadlessChrome marker
|
||||
if mode == StealthMode::FullLaunch {
|
||||
let ua = get_browser_user_agent(client, session_id).await;
|
||||
if let Some(ua) = ua {
|
||||
let cleaned = ua.replace("HeadlessChrome", "Chrome");
|
||||
if cleaned != ua {
|
||||
client
|
||||
.send_command(
|
||||
"Emulation.setUserAgentOverride",
|
||||
Some(json!({
|
||||
"userAgent": cleaned,
|
||||
"acceptLanguage": locale.unwrap_or("en-US"),
|
||||
"platform": platform_string(),
|
||||
"userAgentMetadata": build_ua_metadata(&cleaned, locale),
|
||||
})),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get the browser's User-Agent string via CDP.
|
||||
async fn get_browser_user_agent(client: &CdpClient, session_id: &str) -> Option<String> {
|
||||
let result = client
|
||||
.send_command(
|
||||
"Runtime.evaluate",
|
||||
Some(json!({ "expression": "navigator.userAgent", "returnByValue": true })),
|
||||
Some(session_id),
|
||||
)
|
||||
.await
|
||||
.ok()?;
|
||||
result
|
||||
.get("result")
|
||||
.and_then(|r| r.get("value"))
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from)
|
||||
}
|
||||
|
||||
/// Also run stealth script on the current page (for already-loaded pages after CDP attach).
|
||||
pub async fn apply_stealth_to_current_page(
|
||||
client: &CdpClient,
|
||||
session_id: &str,
|
||||
mode: StealthMode,
|
||||
locale: Option<&str>,
|
||||
) -> Result<(), String> {
|
||||
let script = build_stealth_script(mode, locale);
|
||||
client
|
||||
.send_command(
|
||||
"Runtime.evaluate",
|
||||
Some(json!({
|
||||
"expression": script,
|
||||
"returnByValue": true,
|
||||
})),
|
||||
Some(session_id),
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Strip sourceURL comments from CDP expressions to avoid leaking
|
||||
/// automation-framework identifiers in stack traces.
|
||||
pub fn strip_source_url_labels(input: &str) -> String {
|
||||
// Remove //# sourceURL=... and //@ sourceURL=...
|
||||
let re_line = regex_lite::Regex::new(r"(?i)\n?\s*//[@#]\s*sourceURL=[^\n\r]*").unwrap();
|
||||
let output = re_line.replace_all(input, "");
|
||||
// Remove /*# sourceURL=...*/ block comments
|
||||
let re_block =
|
||||
regex_lite::Regex::new(r"(?is)\n?\s*/\*[@#]\s*sourceURL=[\s\S]*?\*/").unwrap();
|
||||
re_block.replace_all(&output, "").to_string()
|
||||
}
|
||||
|
||||
fn platform_string() -> &'static str {
|
||||
if cfg!(target_os = "macos") {
|
||||
"macOS"
|
||||
} else if cfg!(target_os = "windows") {
|
||||
"Win32"
|
||||
} else {
|
||||
"Linux"
|
||||
}
|
||||
}
|
||||
|
||||
fn platform_hint() -> &'static str {
|
||||
if cfg!(target_os = "macos") {
|
||||
"macOS"
|
||||
} else if cfg!(target_os = "windows") {
|
||||
"Windows"
|
||||
} else {
|
||||
"Linux"
|
||||
}
|
||||
}
|
||||
|
||||
fn platform_version_hint() -> &'static str {
|
||||
if cfg!(target_os = "macos") {
|
||||
"14.0.0"
|
||||
} else if cfg!(target_os = "windows") {
|
||||
"10.0.0"
|
||||
} else {
|
||||
"6.5.0"
|
||||
}
|
||||
}
|
||||
|
||||
fn build_ua_metadata(ua: &str, locale: Option<&str>) -> serde_json::Value {
|
||||
// Extract Chrome version from UA string
|
||||
let chrome_version = ua
|
||||
.split("Chrome/")
|
||||
.nth(1)
|
||||
.and_then(|s| s.split_whitespace().next())
|
||||
.unwrap_or("130.0.0.0");
|
||||
let major = chrome_version.split('.').next().unwrap_or("130");
|
||||
|
||||
let _lang = locale.unwrap_or("en-US");
|
||||
|
||||
json!({
|
||||
"brands": [
|
||||
{ "brand": "Chromium", "version": major },
|
||||
{ "brand": "Google Chrome", "version": major },
|
||||
{ "brand": "Not?A_Brand", "version": "99" },
|
||||
],
|
||||
"fullVersionList": [
|
||||
{ "brand": "Chromium", "version": chrome_version },
|
||||
{ "brand": "Google Chrome", "version": chrome_version },
|
||||
{ "brand": "Not?A_Brand", "version": "99.0.0.0" },
|
||||
],
|
||||
"fullVersion": chrome_version,
|
||||
"platform": platform_hint(),
|
||||
"platformVersion": platform_version_hint(),
|
||||
"architecture": if cfg!(target_arch = "aarch64") { "arm" } else { "x86" },
|
||||
"model": "",
|
||||
"mobile": false,
|
||||
"bitness": "64",
|
||||
"wow64": false,
|
||||
})
|
||||
}
|
||||
+34
-1434
File diff suppressed because it is too large
Load Diff
@@ -1,135 +0,0 @@
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<title>Drag Probe</title>
|
||||
<style>
|
||||
body {
|
||||
margin: 0;
|
||||
font: 14px/1.4 sans-serif;
|
||||
background: #f4f4f4;
|
||||
}
|
||||
|
||||
#pad {
|
||||
position: relative;
|
||||
width: 800px;
|
||||
height: 500px;
|
||||
margin: 24px;
|
||||
border: 1px solid #999;
|
||||
background: white;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
#target {
|
||||
position: absolute;
|
||||
left: 320px;
|
||||
top: 40px;
|
||||
width: 100px;
|
||||
height: 40px;
|
||||
background: #e34c26;
|
||||
color: white;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
user-select: none;
|
||||
cursor: grab;
|
||||
}
|
||||
|
||||
#target.dragging {
|
||||
cursor: grabbing;
|
||||
background: #0d9488;
|
||||
}
|
||||
|
||||
#log {
|
||||
margin: 24px;
|
||||
white-space: pre-wrap;
|
||||
font-family: ui-monospace, monospace;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div id="pad">
|
||||
<div id="target">drag me</div>
|
||||
</div>
|
||||
<pre id="log"></pre>
|
||||
<script>
|
||||
const target = document.getElementById("target");
|
||||
const logEl = document.getElementById("log");
|
||||
|
||||
window.__dragProbe = {
|
||||
dragging: false,
|
||||
events: [],
|
||||
finalLeft: 320,
|
||||
finalTop: 40,
|
||||
};
|
||||
|
||||
let offsetX = 0;
|
||||
let offsetY = 0;
|
||||
|
||||
function pushEvent(event, extra = {}) {
|
||||
window.__dragProbe.events.push({
|
||||
type: event.type,
|
||||
button: event.button,
|
||||
buttons: event.buttons,
|
||||
x: event.clientX,
|
||||
y: event.clientY,
|
||||
target: event.target.id || event.target.tagName,
|
||||
...extra,
|
||||
});
|
||||
logEl.textContent = JSON.stringify(window.__dragProbe, null, 2);
|
||||
}
|
||||
|
||||
function onPointerLikeStart(event) {
|
||||
if (event.type === "mousedown") {
|
||||
const rect = target.getBoundingClientRect();
|
||||
offsetX = event.clientX - rect.left;
|
||||
offsetY = event.clientY - rect.top;
|
||||
window.__dragProbe.dragging = true;
|
||||
target.classList.add("dragging");
|
||||
event.preventDefault();
|
||||
}
|
||||
pushEvent(event, { phase: "start" });
|
||||
}
|
||||
|
||||
target.addEventListener("mousedown", (event) => {
|
||||
const rect = target.getBoundingClientRect();
|
||||
offsetX = event.clientX - rect.left;
|
||||
offsetY = event.clientY - rect.top;
|
||||
window.__dragProbe.dragging = true;
|
||||
target.classList.add("dragging");
|
||||
event.preventDefault();
|
||||
pushEvent(event, { phase: "start" });
|
||||
});
|
||||
target.addEventListener("pointerdown", onPointerLikeStart);
|
||||
|
||||
document.addEventListener("mousemove", (event) => {
|
||||
if (window.__dragProbe.dragging) {
|
||||
const left = event.clientX - offsetX;
|
||||
const top = event.clientY - offsetY;
|
||||
target.style.left = `${left}px`;
|
||||
target.style.top = `${top}px`;
|
||||
window.__dragProbe.finalLeft = left;
|
||||
window.__dragProbe.finalTop = top;
|
||||
}
|
||||
pushEvent(event);
|
||||
});
|
||||
document.addEventListener("pointermove", (event) => {
|
||||
pushEvent(event);
|
||||
});
|
||||
|
||||
document.addEventListener("mouseup", (event) => {
|
||||
if (window.__dragProbe.dragging) {
|
||||
window.__dragProbe.dragging = false;
|
||||
target.classList.remove("dragging");
|
||||
}
|
||||
pushEvent(event, { phase: "end" });
|
||||
});
|
||||
document.addEventListener("pointerup", (event) => {
|
||||
pushEvent(event, { phase: "end" });
|
||||
});
|
||||
target.addEventListener("dragstart", (event) => {
|
||||
pushEvent(event, { phase: "dragstart" });
|
||||
});
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -1,91 +0,0 @@
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<title>HTML5 Drag Probe</title>
|
||||
<style>
|
||||
body {
|
||||
margin: 24px;
|
||||
font: 14px/1.4 sans-serif;
|
||||
}
|
||||
|
||||
#source, #dest {
|
||||
width: 120px;
|
||||
height: 80px;
|
||||
display: inline-flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
border: 1px solid #666;
|
||||
user-select: none;
|
||||
margin-right: 40px;
|
||||
}
|
||||
|
||||
#source {
|
||||
background: #f97316;
|
||||
color: white;
|
||||
}
|
||||
|
||||
#dest {
|
||||
background: #e5e7eb;
|
||||
}
|
||||
|
||||
pre {
|
||||
margin-top: 24px;
|
||||
white-space: pre-wrap;
|
||||
font-family: ui-monospace, monospace;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div id="source" draggable="true">drag source</div>
|
||||
<div id="dest">drop zone</div>
|
||||
<pre id="log"></pre>
|
||||
<script>
|
||||
const source = document.getElementById("source");
|
||||
const dest = document.getElementById("dest");
|
||||
const logEl = document.getElementById("log");
|
||||
|
||||
window.__html5DragProbe = { events: [] };
|
||||
|
||||
function pushEvent(event, extra = {}) {
|
||||
window.__html5DragProbe.events.push({
|
||||
type: event.type,
|
||||
target: event.target.id || event.target.tagName,
|
||||
x: event.clientX,
|
||||
y: event.clientY,
|
||||
button: event.button,
|
||||
buttons: event.buttons,
|
||||
...extra,
|
||||
});
|
||||
logEl.textContent = JSON.stringify(window.__html5DragProbe, null, 2);
|
||||
}
|
||||
|
||||
for (const type of ["pointerdown", "mousedown", "dragstart", "drag", "dragend"]) {
|
||||
source.addEventListener(type, (event) => {
|
||||
if (type === "dragstart") {
|
||||
event.dataTransfer.setData("text/plain", "probe");
|
||||
}
|
||||
pushEvent(event);
|
||||
});
|
||||
}
|
||||
|
||||
for (const type of ["pointermove", "mousemove", "dragenter", "dragover", "drop", "pointerup", "mouseup"]) {
|
||||
document.addEventListener(type, (event) => {
|
||||
if (type === "dragover") {
|
||||
event.preventDefault();
|
||||
}
|
||||
if (type === "drop") {
|
||||
pushEvent(event, { dropped: event.dataTransfer.getData("text/plain") });
|
||||
return;
|
||||
}
|
||||
pushEvent(event);
|
||||
});
|
||||
}
|
||||
|
||||
dest.addEventListener("dragover", (event) => event.preventDefault());
|
||||
dest.addEventListener("drop", (event) => {
|
||||
pushEvent(event, { dropped: event.dataTransfer.getData("text/plain") });
|
||||
});
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -1,113 +0,0 @@
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<title>Pointer Capture Probe</title>
|
||||
<style>
|
||||
body {
|
||||
margin: 24px;
|
||||
font: 14px/1.4 sans-serif;
|
||||
}
|
||||
#crop {
|
||||
position: relative;
|
||||
width: 240px;
|
||||
height: 180px;
|
||||
border: 2px solid #fff;
|
||||
outline: 1px solid #555;
|
||||
background: rgba(0, 0, 0, 0.2);
|
||||
}
|
||||
#handle {
|
||||
position: absolute;
|
||||
width: 20px;
|
||||
height: 20px;
|
||||
top: -16px;
|
||||
left: -16px;
|
||||
padding-top: 13px;
|
||||
padding-left: 13px;
|
||||
box-sizing: content-box;
|
||||
background: rgba(255, 0, 0, 0.25);
|
||||
}
|
||||
#handle::after {
|
||||
content: "";
|
||||
display: block;
|
||||
width: 20px;
|
||||
height: 20px;
|
||||
border-top: 2px solid white;
|
||||
border-left: 2px solid white;
|
||||
}
|
||||
pre {
|
||||
margin-top: 24px;
|
||||
white-space: pre-wrap;
|
||||
font-family: ui-monospace, monospace;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div id="crop" aria-label="crop area">
|
||||
<div id="handle" aria-label="crop handle topLeft" data-anchor="topLeft"></div>
|
||||
</div>
|
||||
<pre id="log"></pre>
|
||||
<script>
|
||||
const crop = document.getElementById("crop");
|
||||
const handle = document.getElementById("handle");
|
||||
const logEl = document.getElementById("log");
|
||||
|
||||
const state = {
|
||||
targetAnchor: null,
|
||||
dragging: false,
|
||||
moved: false,
|
||||
events: [],
|
||||
};
|
||||
window.__pointerCaptureProbe = state;
|
||||
|
||||
function sync() {
|
||||
logEl.textContent = JSON.stringify(state, null, 2);
|
||||
}
|
||||
|
||||
function push(event, extra = {}) {
|
||||
state.events.push({
|
||||
type: event.type,
|
||||
target: event.target.id || event.target.tagName,
|
||||
currentTarget: event.currentTarget.id || event.currentTarget.tagName,
|
||||
pointerId: event.pointerId,
|
||||
button: event.button,
|
||||
buttons: event.buttons,
|
||||
hasCapture: event.currentTarget.hasPointerCapture?.(event.pointerId) ?? false,
|
||||
x: event.clientX,
|
||||
y: event.clientY,
|
||||
...extra,
|
||||
});
|
||||
sync();
|
||||
}
|
||||
|
||||
crop.addEventListener("pointerdown", (event) => {
|
||||
state.targetAnchor = event.target.getAttribute("data-anchor");
|
||||
crop.setPointerCapture(event.pointerId);
|
||||
event.preventDefault();
|
||||
push(event, { phase: "down", targetAnchor: state.targetAnchor });
|
||||
});
|
||||
|
||||
crop.addEventListener("pointermove", (event) => {
|
||||
const hasCapture = crop.hasPointerCapture(event.pointerId);
|
||||
if (hasCapture && state.targetAnchor) {
|
||||
state.dragging = true;
|
||||
state.moved = true;
|
||||
}
|
||||
push(event, { phase: hasCapture ? "drag" : "hover", targetAnchor: state.targetAnchor });
|
||||
});
|
||||
|
||||
crop.addEventListener("pointerup", (event) => {
|
||||
const hadCapture = crop.hasPointerCapture(event.pointerId);
|
||||
state.dragging = false;
|
||||
push(event, { phase: "up", targetAnchor: state.targetAnchor, hadCapture });
|
||||
state.targetAnchor = null;
|
||||
});
|
||||
|
||||
handle.addEventListener("pointerdown", (event) => push(event, { listener: "handle" }));
|
||||
handle.addEventListener("pointermove", (event) => push(event, { listener: "handle" }));
|
||||
handle.addEventListener("pointerup", (event) => push(event, { listener: "handle" }));
|
||||
|
||||
sync();
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -40,45 +40,32 @@ impl AppiumManager {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn build_ios_capabilities(
|
||||
device_udid: Option<&str>,
|
||||
device_name: Option<&str>,
|
||||
platform_version: Option<&str>,
|
||||
) -> Value {
|
||||
let mut caps = json!({
|
||||
"platformName": "iOS",
|
||||
"appium:automationName": "XCUITest",
|
||||
"browserName": "Safari",
|
||||
"appium:noReset": true,
|
||||
});
|
||||
|
||||
if let Some(name) = device_name {
|
||||
caps["appium:deviceName"] = json!(name);
|
||||
} else {
|
||||
caps["appium:deviceName"] = json!("iPhone");
|
||||
}
|
||||
|
||||
if let Some(ver) = platform_version {
|
||||
caps["appium:platformVersion"] = json!(ver);
|
||||
}
|
||||
|
||||
if let Some(udid) = device_udid {
|
||||
caps["appium:udid"] = json!(udid);
|
||||
}
|
||||
|
||||
caps
|
||||
}
|
||||
|
||||
pub async fn create_ios_session(
|
||||
&mut self,
|
||||
device_name: Option<&str>,
|
||||
platform_version: Option<&str>,
|
||||
) -> Result<Value, String> {
|
||||
let caps = Self::build_ios_capabilities(
|
||||
self.device_udid.as_deref(),
|
||||
device_name,
|
||||
platform_version,
|
||||
);
|
||||
let mut caps = json!({
|
||||
"platformName": "iOS",
|
||||
"automationName": "XCUITest",
|
||||
"browserName": "Safari",
|
||||
"noReset": true,
|
||||
});
|
||||
|
||||
if let Some(name) = device_name {
|
||||
caps["deviceName"] = json!(name);
|
||||
} else {
|
||||
caps["deviceName"] = json!("iPhone");
|
||||
}
|
||||
|
||||
if let Some(ver) = platform_version {
|
||||
caps["platformVersion"] = json!(ver);
|
||||
}
|
||||
|
||||
if let Some(ref udid) = self.device_udid {
|
||||
caps["udid"] = json!(udid);
|
||||
}
|
||||
|
||||
self.client.create_session(caps).await
|
||||
}
|
||||
|
||||
@@ -211,30 +198,4 @@ mod tests {
|
||||
assert_eq!(APPIUM_DEFAULT_PORT, 4723);
|
||||
assert_eq!(APPIUM_STARTUP_TIMEOUT_SECS, 30);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ios_capabilities_use_vendor_prefix() {
|
||||
let caps = AppiumManager::build_ios_capabilities(
|
||||
Some("TEST-UDID-123"),
|
||||
Some("iPhone 16 Pro"),
|
||||
Some("18.5"),
|
||||
);
|
||||
|
||||
// W3C standard capabilities must NOT have vendor prefix
|
||||
assert!(caps.get("platformName").is_some());
|
||||
assert!(caps.get("browserName").is_some());
|
||||
|
||||
// Non-standard capabilities MUST have appium: vendor prefix
|
||||
assert!(caps.get("appium:automationName").is_some());
|
||||
assert!(caps.get("appium:noReset").is_some());
|
||||
assert!(caps.get("appium:deviceName").is_some());
|
||||
assert!(caps.get("appium:platformVersion").is_some());
|
||||
assert!(caps.get("appium:udid").is_some());
|
||||
|
||||
// Must NOT have unprefixed non-standard capabilities
|
||||
assert!(caps.get("automationName").is_none());
|
||||
assert!(caps.get("noReset").is_none());
|
||||
assert!(caps.get("deviceName").is_none());
|
||||
assert!(caps.get("udid").is_none());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -212,6 +212,32 @@ impl WebDriverClient {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_client_new() {
|
||||
let client = WebDriverClient::new(4444);
|
||||
assert_eq!(client.base_url, "http://127.0.0.1:4444");
|
||||
assert!(client.session_id.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_session_id_none() {
|
||||
let client = WebDriverClient::new(4444);
|
||||
let result = client.session_id();
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("No active WebDriver session"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_client_custom_port() {
|
||||
let client = WebDriverClient::new(9515);
|
||||
assert_eq!(client.base_url, "http://127.0.0.1:9515");
|
||||
}
|
||||
}
|
||||
|
||||
async fn http_request(method: &str, url: &str, body: Option<&Value>) -> Result<Value, String> {
|
||||
let parsed = url::Url::parse(url).map_err(|e| format!("Invalid URL: {}", e))?;
|
||||
let host = parsed.host_str().unwrap_or("127.0.0.1");
|
||||
@@ -290,29 +316,3 @@ async fn http_request(method: &str, url: &str, body: Option<&Value>) -> Result<V
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_client_new() {
|
||||
let client = WebDriverClient::new(4444);
|
||||
assert_eq!(client.base_url, "http://127.0.0.1:4444");
|
||||
assert!(client.session_id.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_session_id_none() {
|
||||
let client = WebDriverClient::new(4444);
|
||||
let result = client.session_id();
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().contains("No active WebDriver session"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_client_custom_port() {
|
||||
let client = WebDriverClient::new(9515);
|
||||
assert_eq!(client.base_url, "http://127.0.0.1:9515");
|
||||
}
|
||||
}
|
||||
|
||||
+272
-563
File diff suppressed because it is too large
Load Diff
@@ -1,284 +0,0 @@
|
||||
use crate::color;
|
||||
use std::path::Path;
|
||||
use std::process::{exit, Command, Stdio};
|
||||
|
||||
const CURRENT_VERSION: &str = env!("CARGO_PKG_VERSION");
|
||||
const NPM_REGISTRY_URL: &str = "https://registry.npmjs.org/agent-browser/latest";
|
||||
|
||||
enum InstallMethod {
|
||||
Npm,
|
||||
Pnpm,
|
||||
Yarn,
|
||||
Bun,
|
||||
Homebrew,
|
||||
Cargo,
|
||||
Unknown,
|
||||
}
|
||||
|
||||
async fn fetch_latest_version() -> Result<String, String> {
|
||||
let resp = reqwest::get(NPM_REGISTRY_URL)
|
||||
.await
|
||||
.map_err(|e| format!("Failed to fetch version info: {}", e))?;
|
||||
|
||||
let body: serde_json::Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to parse version info: {}", e))?;
|
||||
|
||||
body.get("version")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| "No version field in registry response".to_string())
|
||||
}
|
||||
|
||||
/// Parse the `.install-method` marker written by postinstall.js.
|
||||
fn read_install_method_marker(exe_dir: &Path) -> Option<InstallMethod> {
|
||||
let contents = std::fs::read_to_string(exe_dir.join(".install-method")).ok()?;
|
||||
match contents.trim() {
|
||||
"npm" => Some(InstallMethod::Npm),
|
||||
"pnpm" => Some(InstallMethod::Pnpm),
|
||||
"yarn" => Some(InstallMethod::Yarn),
|
||||
"bun" => Some(InstallMethod::Bun),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn detect_install_method() -> InstallMethod {
|
||||
if let Ok(exe) = std::env::current_exe() {
|
||||
// Resolve symlinks to find the real binary location
|
||||
let real_path = exe.canonicalize().unwrap_or(exe);
|
||||
|
||||
// Preferred: read the marker file written at install time
|
||||
if let Some(dir) = real_path.parent() {
|
||||
if let Some(method) = read_install_method_marker(dir) {
|
||||
return method;
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: infer from executable path
|
||||
let path_str = real_path.to_string_lossy();
|
||||
|
||||
if path_str.contains("/.cargo/bin/") || path_str.contains("\\.cargo\\bin\\") {
|
||||
return InstallMethod::Cargo;
|
||||
}
|
||||
|
||||
if path_str.contains("/Cellar/agent-browser/")
|
||||
|| path_str.contains("/homebrew/")
|
||||
|| path_str.contains("/linuxbrew/")
|
||||
{
|
||||
return InstallMethod::Homebrew;
|
||||
}
|
||||
|
||||
if path_str.contains("/pnpm/") || path_str.contains("/pnpm-global/") {
|
||||
return InstallMethod::Pnpm;
|
||||
}
|
||||
|
||||
if path_str.contains("/.yarn/") || path_str.contains("/yarn/global/") {
|
||||
return InstallMethod::Yarn;
|
||||
}
|
||||
|
||||
if path_str.contains("/.bun/") {
|
||||
return InstallMethod::Bun;
|
||||
}
|
||||
|
||||
if path_str.contains("node_modules/agent-browser")
|
||||
|| path_str.contains("node_modules\\agent-browser")
|
||||
{
|
||||
return InstallMethod::Npm;
|
||||
}
|
||||
}
|
||||
|
||||
// Last resort: probe package managers via subprocess
|
||||
|
||||
#[cfg(any(target_os = "macos", target_os = "linux"))]
|
||||
{
|
||||
if command_succeeds("brew", &["list", "agent-browser"]) {
|
||||
return InstallMethod::Homebrew;
|
||||
}
|
||||
}
|
||||
|
||||
if command_output_contains(
|
||||
"pnpm",
|
||||
&["list", "-g", "agent-browser", "--depth=0"],
|
||||
"agent-browser",
|
||||
) {
|
||||
return InstallMethod::Pnpm;
|
||||
}
|
||||
|
||||
if command_output_contains("yarn", &["global", "list", "--depth=0"], "agent-browser") {
|
||||
return InstallMethod::Yarn;
|
||||
}
|
||||
|
||||
if command_output_contains("bun", &["pm", "ls", "-g"], "agent-browser") {
|
||||
return InstallMethod::Bun;
|
||||
}
|
||||
|
||||
if command_succeeds("npm", &["list", "-g", "agent-browser", "--depth=0"]) {
|
||||
return InstallMethod::Npm;
|
||||
}
|
||||
|
||||
InstallMethod::Unknown
|
||||
}
|
||||
|
||||
fn command_succeeds(cmd: &str, args: &[&str]) -> bool {
|
||||
Command::new(cmd)
|
||||
.args(args)
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.status()
|
||||
.map(|s| s.success())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn command_output_contains(cmd: &str, args: &[&str], needle: &str) -> bool {
|
||||
Command::new(cmd)
|
||||
.args(args)
|
||||
.stderr(Stdio::null())
|
||||
.output()
|
||||
.map(|o| o.status.success() && String::from_utf8_lossy(&o.stdout).contains(needle))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn run_upgrade_command(method: &InstallMethod) -> bool {
|
||||
let (cmd, args, display): (&str, &[&str], &str) = match method {
|
||||
InstallMethod::Npm => (
|
||||
"npm",
|
||||
&["install", "-g", "agent-browser@latest"],
|
||||
"npm install -g agent-browser@latest",
|
||||
),
|
||||
InstallMethod::Pnpm => (
|
||||
"pnpm",
|
||||
&["add", "-g", "agent-browser@latest"],
|
||||
"pnpm add -g agent-browser@latest",
|
||||
),
|
||||
// NOTE: `yarn global` is Yarn Classic (v1) only; Yarn Berry (v2+) removed it.
|
||||
// Users on Yarn v2+ won't reach this path — detection falls through to Unknown.
|
||||
InstallMethod::Yarn => (
|
||||
"yarn",
|
||||
&["global", "add", "agent-browser@latest"],
|
||||
"yarn global add agent-browser@latest",
|
||||
),
|
||||
InstallMethod::Bun => (
|
||||
"bun",
|
||||
&["install", "-g", "agent-browser@latest"],
|
||||
"bun install -g agent-browser@latest",
|
||||
),
|
||||
InstallMethod::Homebrew => (
|
||||
"brew",
|
||||
&["upgrade", "agent-browser"],
|
||||
"brew upgrade agent-browser",
|
||||
),
|
||||
InstallMethod::Cargo => (
|
||||
"cargo",
|
||||
&["install", "agent-browser", "--force"],
|
||||
"cargo install agent-browser --force",
|
||||
),
|
||||
InstallMethod::Unknown => return false,
|
||||
};
|
||||
|
||||
println!("Running: {}", display);
|
||||
Command::new(cmd)
|
||||
.args(args)
|
||||
.status()
|
||||
.map(|s| s.success())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
pub fn run_upgrade() {
|
||||
let current = CURRENT_VERSION;
|
||||
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap_or_else(|e| {
|
||||
eprintln!(
|
||||
"{} Failed to create runtime: {}",
|
||||
color::error_indicator(),
|
||||
e
|
||||
);
|
||||
exit(1);
|
||||
});
|
||||
|
||||
let latest = match rt.block_on(fetch_latest_version()) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
eprintln!(
|
||||
"{} Could not check latest version: {}",
|
||||
color::warning_indicator(),
|
||||
e
|
||||
);
|
||||
String::new()
|
||||
}
|
||||
};
|
||||
|
||||
if !latest.is_empty() && current == latest.as_str() {
|
||||
println!(
|
||||
"{} agent-browser is already at the latest version (v{})",
|
||||
color::success_indicator(),
|
||||
current
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let method = detect_install_method();
|
||||
|
||||
let method_name = match &method {
|
||||
InstallMethod::Npm => "npm",
|
||||
InstallMethod::Pnpm => "pnpm",
|
||||
InstallMethod::Yarn => "yarn",
|
||||
InstallMethod::Bun => "bun",
|
||||
InstallMethod::Homebrew => "Homebrew",
|
||||
InstallMethod::Cargo => "Cargo",
|
||||
InstallMethod::Unknown => "",
|
||||
};
|
||||
|
||||
if matches!(method, InstallMethod::Unknown) {
|
||||
eprintln!(
|
||||
"{} Could not detect installation method.",
|
||||
color::error_indicator()
|
||||
);
|
||||
eprintln!(" To update manually, run one of:");
|
||||
eprintln!(" npm install -g agent-browser@latest # npm");
|
||||
eprintln!(" pnpm add -g agent-browser@latest # pnpm");
|
||||
eprintln!(" yarn global add agent-browser@latest # yarn");
|
||||
eprintln!(" bun install -g agent-browser@latest # bun");
|
||||
eprintln!(" brew upgrade agent-browser # Homebrew");
|
||||
eprintln!(" cargo install agent-browser --force # Cargo");
|
||||
exit(1);
|
||||
}
|
||||
|
||||
println!("Detected installation via {}.", method_name);
|
||||
|
||||
if !latest.is_empty() {
|
||||
println!(
|
||||
"{}",
|
||||
color::cyan(&format!(
|
||||
"Upgrading agent-browser... v{} → v{}",
|
||||
current, latest
|
||||
))
|
||||
);
|
||||
} else {
|
||||
println!(
|
||||
"{}",
|
||||
color::cyan(&format!("Upgrading agent-browser (v{})...", current))
|
||||
);
|
||||
}
|
||||
|
||||
let success = run_upgrade_command(&method);
|
||||
|
||||
if success {
|
||||
if !latest.is_empty() {
|
||||
println!(
|
||||
"{} Done! v{} → v{}",
|
||||
color::success_indicator(),
|
||||
current,
|
||||
latest
|
||||
);
|
||||
} else {
|
||||
println!("{} Done!", color::success_indicator());
|
||||
}
|
||||
} else {
|
||||
eprintln!("{} Upgrade failed.", color::error_indicator());
|
||||
exit(1);
|
||||
}
|
||||
}
|
||||
Vendored
+1
@@ -0,0 +1 @@
|
||||
{"v":1}
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
{
|
||||
"git": {
|
||||
"sha1": "7cb6c7d950c040b2198da553140e1b5e8b6ac682"
|
||||
},
|
||||
"path_in_vcs": "crates/zune-jpeg"
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
/target
|
||||
Vendored
+79
@@ -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.
|
||||

|
||||
|
||||
## 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
@@ -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",
|
||||
]
|
||||
Vendored
+67
@@ -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
@@ -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" }
|
||||
Vendored
+95
@@ -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
|
||||
Vendored
+201
@@ -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.
|
||||
Vendored
+21
@@ -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.
|
||||
Vendored
+19
@@ -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.
|
||||
Vendored
+104
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
Vendored
+987
@@ -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;
|
||||
}
|
||||
}
|
||||
Vendored
+167
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
Vendored
+662
@@ -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
|
||||
}
|
||||
Vendored
+254
@@ -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(())
|
||||
}
|
||||
}
|
||||
Vendored
+206
@@ -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
@@ -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
@@ -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
@@ -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);
|
||||
}
|
||||
Vendored
+194
@@ -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;
|
||||
Vendored
+91
@@ -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))
|
||||
}
|
||||
}
|
||||
}
|
||||
Vendored
+936
@@ -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
|
||||
}
|
||||
Vendored
+688
@@ -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();
|
||||
|
||||
}
|
||||
}
|
||||
Vendored
+485
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
Vendored
+577
@@ -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(())
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
# See https://help.github.com/articles/ignoring-files/ for more about ignoring files.
|
||||
|
||||
# dependencies
|
||||
/node_modules
|
||||
/.pnp
|
||||
.pnp.*
|
||||
.yarn/*
|
||||
!.yarn/patches
|
||||
!.yarn/plugins
|
||||
!.yarn/releases
|
||||
!.yarn/versions
|
||||
|
||||
# testing
|
||||
/coverage
|
||||
|
||||
# next.js
|
||||
/.next/
|
||||
/out/
|
||||
|
||||
# production
|
||||
/build
|
||||
|
||||
# misc
|
||||
.DS_Store
|
||||
*.pem
|
||||
|
||||
# debug
|
||||
npm-debug.log*
|
||||
yarn-debug.log*
|
||||
yarn-error.log*
|
||||
.pnpm-debug.log*
|
||||
|
||||
# env files (can opt-in for committing if needed)
|
||||
.env*
|
||||
|
||||
# vercel
|
||||
.vercel
|
||||
|
||||
# typescript
|
||||
*.tsbuildinfo
|
||||
next-env.d.ts
|
||||
@@ -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. M1(1 周): `network capture` + `scroll-collect`。
|
||||
2. M2(1 周): `extract-to` + selector fallback。
|
||||
3. M3(1 周): 406 回退链路 + 文档补全。
|
||||
4. M4(1 周): 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 浏览器自动化,这次升级能直接提升流程稳定性与迭代效率。
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user