mirror of
https://github.com/BaaRai/ARTEX.git
synced 2026-11-14 20:20:54 +08:00
Compare commits
17
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
939f8d4976 | ||
|
|
b5c0bb9f33 | ||
|
|
9d75df335f | ||
|
|
fe9814c5f6 | ||
|
|
53fc8478da | ||
|
|
edc2f4da6d | ||
|
|
01a9dae49e | ||
|
|
4e126a8524 | ||
|
|
abb6cb177d | ||
|
|
5122baa77f | ||
|
|
fcee6af45e | ||
|
|
b74ff21e11 | ||
|
|
7899662aee | ||
|
|
68976a7680 | ||
|
|
5aab3d3df8 | ||
|
|
dc169b074f | ||
|
|
37e5b69877 |
@@ -2,6 +2,9 @@ name: release
|
||||
on:
|
||||
push:
|
||||
tags: ["v*"]
|
||||
# Dry run: builds, SBOM, attestations and workflow artifacts only. The GitHub
|
||||
# Release and the image push are gated on github.event_name == 'push'.
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: write # required to create the Release
|
||||
@@ -81,6 +84,10 @@ jobs:
|
||||
release:
|
||||
needs: binaries
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write # create the Release (softprops/action-gh-release)
|
||||
id-token: write # Sigstore signing of the provenance attestation
|
||||
attestations: write # store the attestation on the repository
|
||||
steps:
|
||||
- uses: actions/checkout@v4 # CHANGELOG.md and the extractor script
|
||||
- uses: actions/download-artifact@v4
|
||||
@@ -91,6 +98,39 @@ jobs:
|
||||
- name: Generate checksums
|
||||
working-directory: dist
|
||||
run: sha256sum *.zip > SHA256SUMS
|
||||
# SBOM of what ships: the Go modules embedded in the linux/amd64 binary, plus
|
||||
# the frontend lockfile and the Go module files.
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: bin-linux-amd64
|
||||
path: ${{ runner.temp }}/sbom-src/bin
|
||||
- uses: anchore/sbom-action/download-syft@66cbf4bc1f1c0d2edc94016e65bc221b6bb0ad6c # v0.24.3
|
||||
id: syft
|
||||
with:
|
||||
syft-version: v1.54.0
|
||||
- name: Generate CycloneDX SBOM
|
||||
env:
|
||||
SYFT: ${{ steps.syft.outputs.cmd }}
|
||||
SRC: ${{ runner.temp }}/sbom-src
|
||||
TAG: ${{ github.ref_name }}
|
||||
run: |
|
||||
mkdir -p "$SRC/web"
|
||||
cp go.mod go.sum "$SRC/"
|
||||
cp web/package-lock.json "$SRC/web/"
|
||||
sbom="dist/artex-${TAG}.cdx.json"
|
||||
"$SYFT" scan "dir:$SRC" --source-name artex --source-version "$TAG" -o "cyclonedx-json=$sbom"
|
||||
grep -q '"pkg:golang/' "$sbom" || { echo "::error::SBOM has no pkg:golang entries"; exit 1; }
|
||||
grep -q '"pkg:npm/' "$sbom" || { echo "::error::SBOM has no pkg:npm entries"; exit 1; }
|
||||
echo "SBOM_PATH=$sbom" >> "$GITHUB_ENV"
|
||||
- uses: actions/attest-build-provenance@4d101475d8b20a2381f78447822ac1eab6504dd8 # v4.2.2
|
||||
with:
|
||||
subject-path: dist/*.zip
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: release-metadata
|
||||
path: |
|
||||
dist/SHA256SUMS
|
||||
${{ env.SBOM_PATH }}
|
||||
- name: Release body from CHANGELOG.md
|
||||
env:
|
||||
TAG: ${{ github.ref_name }}
|
||||
@@ -103,10 +143,12 @@ jobs:
|
||||
echo "Release $TAG. See CHANGELOG.md for the list of changes in this release." > "$RUNNER_TEMP/release-body.md"
|
||||
fi
|
||||
- uses: softprops/action-gh-release@v2
|
||||
if: github.event_name == 'push'
|
||||
with:
|
||||
files: |
|
||||
dist/*.zip
|
||||
dist/SHA256SUMS
|
||||
${{ env.SBOM_PATH }}
|
||||
generate_release_notes: false
|
||||
body_path: ${{ runner.temp }}/release-body.md
|
||||
|
||||
@@ -118,7 +160,7 @@ jobs:
|
||||
# and provide DOCKERHUB_USERNAME / DOCKERHUB_TOKEN secrets plus DOCKERHUB_IMAGE).
|
||||
docker:
|
||||
needs: binaries
|
||||
if: ${{ vars.DOCKERHUB_ENABLED == 'true' }}
|
||||
if: ${{ vars.DOCKERHUB_ENABLED == 'true' && github.event_name == 'push' }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4 # fetch the Dockerfile and skills/
|
||||
|
||||
@@ -134,6 +134,15 @@ cp config.example.json config.json # fill in the database connection
|
||||
> Please start with `start.sh` / `start.bat` instead of running `./artex` directly. It is a daemon script: if the process crashes, it restarts it after a backoff.
|
||||
> To keep it running in the background: `nohup ./start.sh >artex.log 2>&1 &`.
|
||||
|
||||
#### Verifying a release
|
||||
|
||||
Each release attaches `SHA256SUMS` and a CycloneDX SBOM, `artex-<tag>.cdx.json` (Go modules embedded in the linux/amd64 binary, the frontend `package-lock.json`, and `go.mod`/`go.sum`), next to the zips. The zips also carry a GitHub build provenance attestation. Releases published before this was added have neither the SBOM nor the attestation.
|
||||
|
||||
```bash
|
||||
sha256sum -c SHA256SUMS # add --ignore-missing if you downloaded only some zips
|
||||
gh attestation verify <zip> --repo BaaRai/ARTEX # provenance: built by this repository's release workflow
|
||||
```
|
||||
|
||||
### Option 4: build the single binary from source
|
||||
|
||||
```bash
|
||||
@@ -181,8 +190,9 @@ gunzip -c backups/artex-db-<timestamp>.sql.gz | docker compose exec -T postgres
|
||||
psql -U "$POSTGRES_USER" -d "$POSTGRES_DB"
|
||||
```
|
||||
|
||||
Stored secrets (LLM API keys, MCP env, task credentials, notification channel
|
||||
config) are encrypted in the database with the key in `keys/secret.key`
|
||||
Stored secrets (LLM API keys and proxies, MCP env, task credentials,
|
||||
notification channel config, the Brave and Tavily search keys, the web search
|
||||
and global proxies, the execution spec of custom tools) are encrypted in the database with the key in `keys/secret.key`
|
||||
(generated at first start, mode 0600, next to `jwt.key`). **A database backup
|
||||
without that key can no longer restore those secrets**: back up `keys/` with the
|
||||
dump, separately. Deployments without a persistent disk set `ARTEX_SECRET_KEY`
|
||||
|
||||
@@ -114,6 +114,7 @@ func bbotTyped(raw bbotRawEvent) bbotEvent {
|
||||
}
|
||||
if u, err := url.Parse(strings.TrimSpace(raw.URL)); err == nil && ev.Host != "" &&
|
||||
(u.Scheme == "http" || u.Scheme == "https") && strings.EqualFold(u.Hostname(), ev.Host) {
|
||||
u.User = nil // credentials never reach an asset
|
||||
ev.URL = u.String()
|
||||
}
|
||||
return ev
|
||||
|
||||
@@ -71,3 +71,15 @@ func TestParseBbotJSONLTypedOnly(t *testing.T) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The typed url of an in-scope URL event loses its userinfo, keeps its query.
|
||||
func TestParseBbotURLUserinfo(t *testing.T) {
|
||||
res, err := parseBbotJSONL(strings.NewReader(`{"type":"URL","host":"www.example.com","url":"https://scanuser:s3cr3t@www.example.com/app?id=7","scope_distance":0}` + "\n"))
|
||||
if err != nil || len(res.InScope) != 1 {
|
||||
t.Fatalf("parse: %v, %+v", err, res)
|
||||
}
|
||||
want := []assetInputItem{{Type: "service", URL: "https://www.example.com/app?id=7"}}
|
||||
if got := res.assetItems(); !reflect.DeepEqual(got, want) {
|
||||
t.Errorf("assetItems = %+v, want %+v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
+27
-5
@@ -83,8 +83,8 @@ func scanHost(h string) string {
|
||||
return ""
|
||||
}
|
||||
|
||||
// scanURL returns raw re-serialized and its host when it is an http(s) URL
|
||||
// with a valid host, "" otherwise.
|
||||
// scanURL returns raw re-serialized without its userinfo, and its host, when
|
||||
// it is an http(s) URL with a valid host, "" otherwise.
|
||||
func scanURL(raw string) (string, string) {
|
||||
u, err := url.Parse(strings.TrimSpace(raw))
|
||||
if err != nil || (u.Scheme != "http" && u.Scheme != "https") {
|
||||
@@ -94,9 +94,31 @@ func scanURL(raw string) (string, string) {
|
||||
if h == "" {
|
||||
return "", ""
|
||||
}
|
||||
u.User = nil // credentials never reach an asset, a note or a fact
|
||||
return u.String(), h
|
||||
}
|
||||
|
||||
// stripUserinfo removes the userinfo (user:pass@) from raw, a URL or a bare
|
||||
// host[:port] as a scanner prints it, and keeps every other byte. It is for
|
||||
// the raw values kept in a fact detail: unlike scanURL it never rejects, and a
|
||||
// value with no userinfo, or a path with no authority, is returned unchanged.
|
||||
func stripUserinfo(raw string) string {
|
||||
start := 0 // the authority follows "scheme://" or "//"; a bare value is all authority
|
||||
if i := strings.Index(raw, "//"); i == 0 || i > 0 && raw[i-1] == ':' && !strings.ContainsAny(raw[:i-1], "/?#@") {
|
||||
start = i + 2
|
||||
} else if strings.ContainsAny(raw, "/?#") {
|
||||
return raw
|
||||
}
|
||||
auth := raw[start:]
|
||||
if j := strings.IndexAny(auth, "/?#"); j >= 0 {
|
||||
auth = auth[:j]
|
||||
}
|
||||
if at := strings.LastIndex(auth, "@"); at >= 0 {
|
||||
return raw[:start] + raw[start+at+1:]
|
||||
}
|
||||
return raw
|
||||
}
|
||||
|
||||
// scanHostOf reads a host from a URL, a host:port or a bare host.
|
||||
func scanHostOf(s string) string {
|
||||
if _, h := scanURL(s); h != "" {
|
||||
@@ -342,7 +364,7 @@ func (s *scanIngest) addNuclei(line []byte) {
|
||||
}
|
||||
s.note("results", tid+" ("+sev+") on "+host)
|
||||
s.fact(fmt.Sprintf("nuclei: template %s (severity %s) matched on %s; unverified scanner result", tid, sev, host),
|
||||
map[string]any{"name": r.Info.Name, "host": r.Host, "matched_at": r.MatchedAt,
|
||||
map[string]any{"name": r.Info.Name, "host": stripUserinfo(r.Host), "matched_at": stripUserinfo(r.MatchedAt),
|
||||
"matcher_name": r.MatcherName, "extracted_results": r.Extracted})
|
||||
}
|
||||
|
||||
@@ -394,7 +416,7 @@ func (s *scanIngest) addHttpx(line []byte) {
|
||||
s.note("urls", fmt.Sprintf("%s %d", u, r.StatusCode))
|
||||
s.fact(fmt.Sprintf("httpx: %s responded with status %d", u, r.StatusCode),
|
||||
map[string]any{"title": r.Title, "tech": r.Tech, "webserver": r.Webserver, "content_length": r.ContentLength,
|
||||
"host": r.Host, "ip": r.IP, "tls": r.TLS})
|
||||
"host": stripUserinfo(r.Host), "ip": r.IP, "tls": r.TLS})
|
||||
}
|
||||
|
||||
// ---- ffuf -of json ----
|
||||
@@ -480,5 +502,5 @@ func (s *scanIngest) addFfuf(raw json.RawMessage) {
|
||||
s.Assets = append(s.Assets, item)
|
||||
s.note("urls", fmt.Sprintf("%s %d", u, r.Status))
|
||||
s.fact(fmt.Sprintf("ffuf: %s returned status %d, %d bytes, %d words, %d lines", u, r.Status, r.Length, r.Words, r.Lines),
|
||||
map[string]any{"input": r.Input, "redirectlocation": r.Redirect})
|
||||
map[string]any{"input": r.Input, "redirectlocation": stripUserinfo(r.Redirect)})
|
||||
}
|
||||
|
||||
@@ -248,6 +248,137 @@ func TestScanIngestJSONLLineLimits(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// scanURL drops the userinfo and keeps scheme, host, port, path and query; an
|
||||
// @ in the path is not userinfo; a URL it refused before is still refused.
|
||||
func TestScanURLUserinfo(t *testing.T) {
|
||||
cases := []struct{ raw, url, host string }{
|
||||
{"http://scanuser:s3cr3t@scan-ingest.example.com:8080/app?id=7&x=y", "http://scan-ingest.example.com:8080/app?id=7&x=y", "scan-ingest.example.com"},
|
||||
{"https://scanuser@scan-ingest.example.com/app", "https://scan-ingest.example.com/app", "scan-ingest.example.com"},
|
||||
{"https://scan-ingest.example.com/a@b/c?x=@y", "https://scan-ingest.example.com/a@b/c?x=@y", "scan-ingest.example.com"},
|
||||
{"http://scanuser:s3cr3t@scan-ingest.example.com/%zz", "", ""},
|
||||
{"javascript:alert(1)", "", ""},
|
||||
{"http://scanuser:s3cr3t@bad host/", "", ""},
|
||||
{"not a url", "", ""},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if u, h := scanURL(c.raw); u != c.url || h != c.host {
|
||||
t.Errorf("scanURL(%q) = %q, %q; want %q, %q", c.raw, u, h, c.url, c.host)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestStripUserinfo(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"https://scanuser:s3cr3t@scan-ingest.example.com:443/login?next=1": "https://scan-ingest.example.com:443/login?next=1",
|
||||
"http://scanuser:s3cr3t@scan-ingest.example.com/%zz": "http://scan-ingest.example.com/%zz",
|
||||
"//scanuser:s3cr3t@scan-ingest.example.com/": "//scan-ingest.example.com/",
|
||||
"scanuser:s3cr3t@scan-ingest.example.com:443": "scan-ingest.example.com:443",
|
||||
"https://scan-ingest.example.com/a@b?x=@y": "https://scan-ingest.example.com/a@b?x=@y",
|
||||
"https://scan-ingest.example.com/?r=http://scanuser:s3cr3t@x/": "https://scan-ingest.example.com/?r=http://scanuser:s3cr3t@x/",
|
||||
"/admin@x/": "/admin@x/",
|
||||
"scan-ingest.example.com:443": "scan-ingest.example.com:443",
|
||||
"198.51.100.77": "198.51.100.77",
|
||||
"": "",
|
||||
}
|
||||
for raw, want := range cases {
|
||||
if got := stripUserinfo(raw); got != want {
|
||||
t.Errorf("stripUserinfo(%q) = %q, want %q", raw, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Scanner lines carrying user:pass: no asset, url, notable item or fact keeps
|
||||
// the credentials, and the query survives.
|
||||
const (
|
||||
nucleiUserinfoFixture = `{"template-id":"tech-detect","info":{"severity":"info"},"host":"https://scanuser:s3cr3t@scan-ingest.example.com","matched-at":"https://scanuser:s3cr3t@scan-ingest.example.com/login?next=1"}
|
||||
{"template-id":"tech-detect","info":{"severity":"info"},"host":"scanuser:s3cr3t@scan-ingest.example.com:443","matched-at":"https://scanuser:s3cr3t@scan-ingest.example.com:443/x"}
|
||||
`
|
||||
httpxUserinfoFixture = `{"url":"https://scanuser:s3cr3t@scan-ingest.example.com/app?id=7","status_code":200,"host":"198.51.100.77"}
|
||||
`
|
||||
ffufUserinfoFixture = `{"results":[{"input":{"FUZZ":"app"},"status":301,"url":"https://scanuser:s3cr3t@scan-ingest.example.com/app?id=7","redirectlocation":"https://scanuser:s3cr3t@scan-ingest.example.com/app/"}]}
|
||||
`
|
||||
)
|
||||
|
||||
// hasScanCreds reports whether v holds a user or password of the fixtures.
|
||||
func hasScanCreds(v string) bool {
|
||||
for _, c := range []string{"s3cr3t", "scanuser", "alice", "bob"} {
|
||||
if strings.Contains(v, c) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func checkNoUserinfo(t *testing.T, format string, s *scanIngest) {
|
||||
t.Helper()
|
||||
var all []string
|
||||
for _, a := range s.Assets {
|
||||
all = append(all, a.URL)
|
||||
}
|
||||
for u := range s.urls {
|
||||
all = append(all, u)
|
||||
}
|
||||
for _, items := range s.Notable {
|
||||
all = append(all, items...)
|
||||
}
|
||||
for _, f := range s.Facts {
|
||||
d, _ := strconv.Unquote(f.Detail)
|
||||
all = append(all, f.Summary, d)
|
||||
}
|
||||
for _, v := range all {
|
||||
if hasScanCreds(v) {
|
||||
t.Fatalf("%s: credentials kept in %q", format, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestScanIngestStripsUserinfo(t *testing.T) {
|
||||
cases := []struct {
|
||||
format, content string
|
||||
want string // events urls facts assets
|
||||
url string // first asset URL
|
||||
}{
|
||||
{"nuclei-jsonl", nucleiUserinfoFixture, "[2 2 2 2]", "https://scan-ingest.example.com"},
|
||||
{"httpx-jsonl", httpxUserinfoFixture, "[1 1 1 1]", "https://scan-ingest.example.com/app?id=7"},
|
||||
{"ffuf-json", ffufUserinfoFixture, "[1 1 1 1]", "https://scan-ingest.example.com/app?id=7"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
s := parseScan(t, c.format, c.content)
|
||||
if got := fmt.Sprint([]int{s.Events, len(s.urls), len(s.Facts), len(s.Assets)}); got != c.want || s.Skipped != 0 {
|
||||
t.Fatalf("%s: events,urls,facts,assets = %s, want %s (skipped %d)", c.format, got, c.want, s.Skipped)
|
||||
}
|
||||
if s.Assets[0].URL != c.url {
|
||||
t.Fatalf("%s: asset URL %q, want %q", c.format, s.Assets[0].URL, c.url)
|
||||
}
|
||||
checkNoUserinfo(t, c.format, s)
|
||||
}
|
||||
// The raw values kept in the fact detail lose only their userinfo.
|
||||
s := parseScan(t, "nuclei-jsonl", nucleiUserinfoFixture)
|
||||
if d, _ := strconv.Unquote(s.Facts[1].Detail); !strings.Contains(d, `"host":"scan-ingest.example.com:443"`) ||
|
||||
!strings.Contains(d, `"matched_at":"https://scan-ingest.example.com:443/x"`) {
|
||||
t.Fatalf("nuclei detail %s", d)
|
||||
}
|
||||
s = parseScan(t, "ffuf-json", ffufUserinfoFixture)
|
||||
if d, _ := strconv.Unquote(s.Facts[0].Detail); !strings.Contains(d, `"redirectlocation":"https://scan-ingest.example.com/app/"`) {
|
||||
t.Fatalf("ffuf detail %s", d)
|
||||
}
|
||||
}
|
||||
|
||||
// The same URL with different credentials is one URL after normalisation:
|
||||
// both events are counted, both assets carry the same URL.
|
||||
func TestScanIngestUserinfoDedup(t *testing.T) {
|
||||
s := parseScan(t, "httpx-jsonl", httpxUserinfoDedupFixture)
|
||||
if s.Events != 2 || len(s.urls) != 1 || len(s.Assets) != 2 || s.Assets[0].URL != s.Assets[1].URL ||
|
||||
s.Assets[0].URL != "https://scan-ingest.example.com/app?id=7" {
|
||||
t.Fatalf("events=%d urls=%d assets=%+v", s.Events, len(s.urls), s.Assets)
|
||||
}
|
||||
checkNoUserinfo(t, "httpx-jsonl", s)
|
||||
}
|
||||
|
||||
const httpxUserinfoDedupFixture = `{"url":"https://alice:s3cr3t-one@scan-ingest.example.com/app?id=7","status_code":200}
|
||||
{"url":"https://bob:s3cr3t-two@scan-ingest.example.com/app?id=7","status_code":200}
|
||||
`
|
||||
|
||||
// Past scanMaxItems assets and facts, records are counted, not ingested.
|
||||
func TestScanIngestItemCap(t *testing.T) {
|
||||
var b strings.Builder
|
||||
|
||||
@@ -199,3 +199,78 @@ func TestIngestScanOutputJSONPG(t *testing.T) {
|
||||
t.Fatalf("re-ingestion: %d assets and %d facts, first run %d and %d", a, f, assets, facts)
|
||||
}
|
||||
}
|
||||
|
||||
// ingest_scan_output (PG), credentials in URLs: no asset row and no fact keeps
|
||||
// them, and the same URL with different credentials is one asset row and one
|
||||
// fact.
|
||||
func TestIngestScanOutputUserinfoPG(t *testing.T) {
|
||||
d := testDB(t)
|
||||
t.Cleanup(func() { _ = d.Close() })
|
||||
task, err := d.CreateTaskWithOptions("ingest_scan_output_userinfo", "scan-ingest.example.com", db.TaskCreateOptions{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
const assetsWhere = `domain LIKE '%scan-ingest.example.com' OR url LIKE '%scan-ingest.example.com%'`
|
||||
t.Cleanup(func() {
|
||||
_ = d.DeleteTask(task.ID)
|
||||
_, _ = d.Exec(`DELETE FROM assets WHERE ` + assetsWhere)
|
||||
})
|
||||
store := d.Exploration(task.ExplorationID)
|
||||
tsx := NewToolSet(store, "worker-1")
|
||||
tsx.SetTaskID(task.ID)
|
||||
tsx.SetAssetStore(d.Assets(), d.Companies())
|
||||
dir := t.TempDir()
|
||||
tool := tsx.ingestScanOutput()
|
||||
ingest := func(format, content string) map[string]any {
|
||||
t.Helper()
|
||||
if err := os.WriteFile(filepath.Join(dir, "scan.out"), []byte(content), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
in, _ := json.Marshal(map[string]string{"format": format, "path": "scan.out"})
|
||||
res, err := tool.Call(context.Background(), in, &actool.ToolContext{WorkingDir: dir})
|
||||
var out map[string]any
|
||||
if err != nil || res.IsError || json.Unmarshal([]byte(res.Flatten()), &out) != nil ||
|
||||
out["asset_errors"] != nil || out["fact_errors"] != nil {
|
||||
t.Fatalf("%s: %v %s", format, err, res.Flatten())
|
||||
}
|
||||
return out
|
||||
}
|
||||
count := func(q string) int {
|
||||
t.Helper()
|
||||
var n int
|
||||
if err := d.QueryRow(q).Scan(&n); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
out := ingest("httpx-jsonl", httpxUserinfoDedupFixture)
|
||||
if out["events"] != 2.0 || out["urls"] != 1.0 || out["facts_written"] != 1.0 || out["facts_existing"] != 1.0 {
|
||||
t.Fatalf("same URL, different credentials: %v", out)
|
||||
}
|
||||
const appURL = `SELECT count(*) FROM assets WHERE type='service' AND url='https://scan-ingest.example.com/app?id=7'`
|
||||
if n := count(appURL); n != 1 {
|
||||
t.Fatalf("%d asset rows for the normalised URL, want 1", n)
|
||||
}
|
||||
ingest("httpx-jsonl", httpxUserinfoFixture)
|
||||
ingest("ffuf-json", ffufUserinfoFixture)
|
||||
ingest("nuclei-jsonl", nucleiUserinfoFixture)
|
||||
if n := count(appURL); n != 1 {
|
||||
t.Fatalf("%d asset rows for the normalised URL after httpx and ffuf, want 1", n)
|
||||
}
|
||||
if n := count(`SELECT count(*) FROM assets a WHERE (` + assetsWhere + `) AND row_to_json(a)::text ~ 's3cr3t|scanuser|alice|bob'`); n != 0 {
|
||||
t.Fatalf("%d asset rows keep credentials", n)
|
||||
}
|
||||
facts, err := store.ListByKind(db.KindFact, 500)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(facts) == 0 {
|
||||
t.Fatal("no facts written")
|
||||
}
|
||||
for _, f := range facts {
|
||||
if hasScanCreds(string(f.Payload)) {
|
||||
t.Fatalf("fact keeps credentials: %s", f.Payload)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+3
-1
@@ -160,7 +160,8 @@ func calcCSegment(ipStr string) string {
|
||||
return ipnet.String()
|
||||
}
|
||||
|
||||
// normalizeURL lowercases scheme and host, strips trailing slash from bare roots.
|
||||
// normalizeURL lowercases scheme and host, drops userinfo (credentials must
|
||||
// never reach assets.url), and strips trailing slash from bare roots.
|
||||
func normalizeURL(raw string) string {
|
||||
u, err := url.Parse(raw)
|
||||
if err != nil {
|
||||
@@ -168,6 +169,7 @@ func normalizeURL(raw string) string {
|
||||
}
|
||||
u.Scheme = strings.ToLower(u.Scheme)
|
||||
u.Host = strings.ToLower(u.Host)
|
||||
u.User = nil
|
||||
s := u.String()
|
||||
if strings.HasSuffix(s, "/") && u.Path == "/" && u.RawQuery == "" {
|
||||
s = strings.TrimSuffix(s, "/")
|
||||
|
||||
+144
-3
@@ -893,6 +893,10 @@ func TestNormalizeURL(t *testing.T) {
|
||||
{"HTTPS://Example.COM/path", "https://example.com/path"},
|
||||
{"http://example.com/", "http://example.com"},
|
||||
{"http://example.com/page", "http://example.com/page"},
|
||||
{"https://u:p@Example.com/app?x=1", "https://example.com/app?x=1"},
|
||||
{"https://user@example.com/", "https://example.com"},
|
||||
{"https://u:p@example.com:8443/users/@alice?x=1", "https://example.com:8443/users/@alice?x=1"},
|
||||
{"https://example.com/@me?r=http://a:b@x/", "https://example.com/@me?r=http://a:b@x/"},
|
||||
}
|
||||
for _, tc := range tests {
|
||||
got := normalizeURL(tc.raw)
|
||||
@@ -902,6 +906,139 @@ func TestNormalizeURL(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// storedURLRow reads back the url, domain and port of one asset row.
|
||||
func storedURLRow(t *testing.T, d *DB, id int64) (string, string, int) {
|
||||
t.Helper()
|
||||
var u, domain *string
|
||||
var port *int
|
||||
if err := d.QueryRow(`SELECT url, domain, port FROM assets WHERE id = $1`, id).Scan(&u, &domain, &port); err != nil {
|
||||
t.Fatalf("read asset %d: %v", id, err)
|
||||
}
|
||||
var us, ds string
|
||||
var p int
|
||||
if u != nil {
|
||||
us = *u
|
||||
}
|
||||
if domain != nil {
|
||||
ds = *domain
|
||||
}
|
||||
if port != nil {
|
||||
p = *port
|
||||
}
|
||||
return us, ds, p
|
||||
}
|
||||
|
||||
func TestUpsertHTTPServiceStripsUserinfo(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
cleanup := func() {
|
||||
d.Exec(`DELETE FROM assets WHERE url LIKE '%storeuserinfo-svc.com%' OR domain LIKE '%storeuserinfo-svc.com' OR root_domain = 'storeuserinfo-svc.com'`)
|
||||
}
|
||||
cleanup()
|
||||
defer cleanup()
|
||||
|
||||
id1, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{URL: "https://alice:s3cr3t-one@App.storeuserinfo-svc.com/app?x=1"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
u, domain, port := storedURLRow(t, d, id1)
|
||||
if u != "https://app.storeuserinfo-svc.com/app?x=1" || domain != "app.storeuserinfo-svc.com" || port != 443 {
|
||||
t.Fatalf("stored url=%q domain=%q port=%d", u, domain, port)
|
||||
}
|
||||
|
||||
// Same URL, different credentials: deduplicated onto the same row.
|
||||
id2, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{URL: "https://bob:s3cr3t-two@app.storeuserinfo-svc.com/app?x=1"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id2 != id1 {
|
||||
t.Errorf("dedup across credentials failed: %d vs %d", id2, id1)
|
||||
}
|
||||
var n int
|
||||
d.QueryRow(`SELECT count(*) FROM assets WHERE type = 'service' AND service_type = 'http' AND url LIKE '%storeuserinfo-svc.com/app%'`).Scan(&n)
|
||||
if n != 1 {
|
||||
t.Errorf("want 1 http service row, got %d", n)
|
||||
}
|
||||
d.QueryRow(`SELECT count(*) FROM assets WHERE url LIKE '%s3cr3t%' AND url LIKE '%storeuserinfo-svc.com%'`).Scan(&n)
|
||||
if n != 0 {
|
||||
t.Errorf("credentials stored in %d asset url(s)", n)
|
||||
}
|
||||
|
||||
// @ in the path is kept; only the userinfo goes.
|
||||
id3, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{URL: "https://u:p@app.storeuserinfo-svc.com:8443/users/@alice?x=1"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
u, domain, port = storedURLRow(t, d, id3)
|
||||
if u != "https://app.storeuserinfo-svc.com:8443/users/@alice?x=1" || domain != "app.storeuserinfo-svc.com" || port != 8443 {
|
||||
t.Errorf("@ path: stored url=%q domain=%q port=%d", u, domain, port)
|
||||
}
|
||||
|
||||
// Regression: a URL without userinfo is stored unchanged.
|
||||
plain := "https://plain.storeuserinfo-svc.com/@me?q=1"
|
||||
id4, err := av2.UpsertHTTPService(UpsertHTTPServiceReq{URL: plain})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if u, _, _ = storedURLRow(t, d, id4); u != plain {
|
||||
t.Errorf("plain url: want %q, got %q", plain, u)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpsertEndpointStripsUserinfo(t *testing.T) {
|
||||
d, av2, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
cleanup := func() {
|
||||
d.Exec(`DELETE FROM assets WHERE url LIKE '%storeuserinfo-ep.com%' OR domain LIKE '%storeuserinfo-ep.com' OR root_domain = 'storeuserinfo-ep.com'`)
|
||||
}
|
||||
cleanup()
|
||||
defer cleanup()
|
||||
|
||||
id1, err := av2.UpsertEndpoint(UpsertEndpointReq{URL: "https://alice:s3cr3t-one@api.storeuserinfo-ep.com/users?id=1", Method: "GET"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
u, domain, port := storedURLRow(t, d, id1)
|
||||
if u != "https://api.storeuserinfo-ep.com/users?id=1" || domain != "api.storeuserinfo-ep.com" || port != 443 {
|
||||
t.Fatalf("stored url=%q domain=%q port=%d", u, domain, port)
|
||||
}
|
||||
|
||||
id2, err := av2.UpsertEndpoint(UpsertEndpointReq{URL: "https://bob:s3cr3t-two@api.storeuserinfo-ep.com/users?id=1", Method: "GET"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id2 != id1 {
|
||||
t.Errorf("endpoint dedup across credentials failed: %d vs %d", id2, id1)
|
||||
}
|
||||
var n int
|
||||
d.QueryRow(`SELECT count(*) FROM assets WHERE type = 'endpoint' AND url LIKE '%storeuserinfo-ep.com/users%'`).Scan(&n)
|
||||
if n != 1 {
|
||||
t.Errorf("want 1 endpoint row, got %d", n)
|
||||
}
|
||||
d.QueryRow(`SELECT count(*) FROM assets WHERE url LIKE '%s3cr3t%' AND url LIKE '%storeuserinfo-ep.com%'`).Scan(&n)
|
||||
if n != 0 {
|
||||
t.Errorf("credentials stored in %d asset url(s)", n)
|
||||
}
|
||||
|
||||
id3, err := av2.UpsertEndpoint(UpsertEndpointReq{URL: "http://u:p@api.storeuserinfo-ep.com/users/@alice?x=1", Method: "POST"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
u, _, port = storedURLRow(t, d, id3)
|
||||
if u != "http://api.storeuserinfo-ep.com/users/@alice?x=1" || port != 80 {
|
||||
t.Errorf("@ path: stored url=%q port=%d", u, port)
|
||||
}
|
||||
|
||||
plain := "https://plain.storeuserinfo-ep.com/@me?q=1"
|
||||
id4, err := av2.UpsertEndpoint(UpsertEndpointReq{URL: plain, Method: "GET"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if u, _, _ = storedURLRow(t, d, id4); u != plain {
|
||||
t.Errorf("plain url: want %q, got %q", plain, u)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssetPaginationUsesStableIDTieBreaker(t *testing.T) {
|
||||
d, assets, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
@@ -1010,7 +1147,8 @@ func TestAssetIPRejectsHostname(t *testing.T) {
|
||||
// are excluded. Direct source tasks' scope is included.
|
||||
func TestQueryDSLInScopeMembership(t *testing.T) {
|
||||
d, assets, _ := testSetup(t)
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanups below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
task, err := d.CreateTask("scope membership", "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
@@ -1020,7 +1158,10 @@ func TestQueryDSLInScopeMembership(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(task.ID); _ = d.DeleteTask(src.ID) })
|
||||
t.Cleanup(func() {
|
||||
cleanupErr(t, "task", d.DeleteTask(task.ID))
|
||||
cleanupErr(t, "source task", d.DeleteTask(src.ID))
|
||||
})
|
||||
// task ← src is a direct source relation.
|
||||
if _, err := d.Exec(`INSERT INTO task_relations(task_id, source_task_id) VALUES ($1,$2)`, task.ID, src.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
@@ -1079,7 +1220,7 @@ func TestQueryDSLInScopeMembership(t *testing.T) {
|
||||
}
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = d.Exec(`DELETE FROM assets WHERE id=ANY($1::bigint[]) OR id=ANY($2::bigint[])`, inScopeIDs, outIDs)
|
||||
cleanupExec(t, d, `DELETE FROM assets WHERE id=ANY($1::bigint[]) OR id=ANY($2::bigint[])`, inScopeIDs, outIDs)
|
||||
})
|
||||
|
||||
got, err := assets.QueryDSLInScope(marker, "", task.ID, 50, 0)
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// These helpers report a failed test cleanup instead of swallowing it: a silent
|
||||
// failure leaves the test's rows behind for the next run on the same database.
|
||||
|
||||
// cleanupErr reports err from the cleanup step described by what.
|
||||
func cleanupErr(t *testing.T, what string, err error) {
|
||||
t.Helper()
|
||||
if err != nil {
|
||||
t.Errorf("cleanup %s: %v", what, err)
|
||||
}
|
||||
}
|
||||
|
||||
// cleanupExec runs one cleanup statement. Matching no row is not an error.
|
||||
func cleanupExec(t *testing.T, d *DB, query string, args ...any) {
|
||||
t.Helper()
|
||||
_, err := d.Exec(query, args...)
|
||||
cleanupErr(t, query, err)
|
||||
}
|
||||
|
||||
// cleanupProfile deletes a profile through DeleteProfile. A profile the test
|
||||
// already deleted is not an error.
|
||||
func cleanupProfile(t *testing.T, d *DB, id int64) {
|
||||
t.Helper()
|
||||
if err := d.DeleteProfile(id); !errors.Is(err, ErrLLMProfileNotFound) {
|
||||
cleanupErr(t, fmt.Sprintf("profile %d", id), err)
|
||||
}
|
||||
}
|
||||
@@ -10,7 +10,8 @@ import (
|
||||
|
||||
func TestAssetUpsertWaitsForCompanyScopeMutation(t *testing.T) {
|
||||
d, assets, companies := testSetup(t)
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
stamp := time.Now().UnixNano()
|
||||
domain := fmt.Sprintf("scope-lock-%d.example", stamp)
|
||||
@@ -23,8 +24,8 @@ func TestAssetUpsertWaitsForCompanyScopeMutation(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = d.Exec(`DELETE FROM assets WHERE domain = $1`, domain)
|
||||
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, oldCompany, newCompany)
|
||||
cleanupExec(t, d, `DELETE FROM assets WHERE domain = $1`, domain)
|
||||
cleanupExec(t, d, `DELETE FROM companies WHERE id IN ($1,$2)`, oldCompany, newCompany)
|
||||
})
|
||||
if added, _, invalid, validationErrors, err := companies.AddScopeInputsChecked(oldCompany, []ScopeInput{
|
||||
{Kind: "domain", Value: domain},
|
||||
@@ -90,7 +91,8 @@ func TestAssetUpsertWaitsForCompanyScopeMutation(t *testing.T) {
|
||||
|
||||
func TestResolveAndRecomputeUseStableCompanyIDTieBreak(t *testing.T) {
|
||||
d, assets, companies := testSetup(t)
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
stamp := time.Now().UnixNano()
|
||||
first, _, err := companies.UpsertCompany(fmt.Sprintf("Tie Break First %d", stamp), "")
|
||||
@@ -109,8 +111,8 @@ func TestResolveAndRecomputeUseStableCompanyIDTieBreak(t *testing.T) {
|
||||
icp := fmt.Sprintf("ICP-TIE-%d", stamp)
|
||||
appName := fmt.Sprintf("Tie Break App %d", stamp)
|
||||
t.Cleanup(func() {
|
||||
_, _ = d.Exec(`DELETE FROM assets WHERE domain=$1 OR ip=$2 OR app_name=$3`, domain, ip, appName)
|
||||
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2)`, first, second)
|
||||
cleanupExec(t, d, `DELETE FROM assets WHERE domain=$1 OR ip=$2 OR app_name=$3`, domain, ip, appName)
|
||||
cleanupExec(t, d, `DELETE FROM companies WHERE id IN ($1,$2)`, first, second)
|
||||
})
|
||||
|
||||
rules := []ScopeInput{
|
||||
|
||||
@@ -2,6 +2,7 @@ package db
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math/rand/v2"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
@@ -331,19 +332,36 @@ func TestCompanyScopeAttribution(t *testing.T) {
|
||||
d.Close()
|
||||
t.Fatalf("startMax: %v", err)
|
||||
}
|
||||
// Remove everything the test creates, or the next pass on the same database
|
||||
// finds this company's scope and attributes its assets to it. Assets first
|
||||
// (they reference the company), then the company; its scope rows cascade.
|
||||
var cid int64
|
||||
t.Cleanup(func() {
|
||||
_, _ = d.Exec(`DELETE FROM assets WHERE id > $1`, startMax)
|
||||
d.Close()
|
||||
defer d.Close()
|
||||
if _, err := d.Exec(`DELETE FROM assets WHERE id > $1`, startMax); err != nil {
|
||||
t.Errorf("cleanup assets: %v", err)
|
||||
}
|
||||
if cid != 0 {
|
||||
if _, err := d.Exec(`DELETE FROM companies WHERE id = $1`, cid); err != nil {
|
||||
t.Errorf("cleanup company %d: %v", cid, err)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
uniq := startMax + 1
|
||||
root := fmt.Sprintf("scopetest%d.com", uniq)
|
||||
// Unique per run: MAX(assets.id) comes back to the same value after cleanup,
|
||||
// so it cannot tell one pass from the next. The in-scope /24 is picked per
|
||||
// run from 198.18.0.0/15 so a scope left by an interrupted run cannot claim
|
||||
// this run's IP.
|
||||
uniq := fmt.Sprintf("%d%d", time.Now().UnixNano(), rand.Uint32())
|
||||
n := rand.IntN(512)
|
||||
cidr := fmt.Sprintf("198.%d.%d.0/24", 18+n/256, n%256)
|
||||
root := fmt.Sprintf("scopetest%s.com", uniq)
|
||||
sub := "api." + root
|
||||
ipIn := "198.51.100.9"
|
||||
ipIn := fmt.Sprintf("198.%d.%d.9", 18+n/256, n%256)
|
||||
ipOut := "203.0.113.9"
|
||||
outDomain := fmt.Sprintf("other%d.net", uniq)
|
||||
outDomain := fmt.Sprintf("other%s.net", uniq)
|
||||
|
||||
cid, _, err := cs.UpsertCompany(fmt.Sprintf("ScopeCo %d", uniq), "")
|
||||
cid, _, err = cs.UpsertCompany(fmt.Sprintf("ScopeCo %s", uniq), "")
|
||||
if err != nil {
|
||||
t.Fatalf("UpsertCompany: %v", err)
|
||||
}
|
||||
@@ -359,7 +377,7 @@ func TestCompanyScopeAttribution(t *testing.T) {
|
||||
t.Fatalf("pre-scope asset should be unattributed, got %v", *preCompanyID)
|
||||
}
|
||||
|
||||
cs.AddScope(cid, []string{root, "198.51.100.0/24"}, "unit test")
|
||||
cs.AddScope(cid, []string{root, cidr}, "unit test")
|
||||
|
||||
mustCid := func(id int64, want int64, label string) {
|
||||
var cID *int64
|
||||
|
||||
+11
-4
@@ -107,7 +107,10 @@ func scanProfile(sc interface{ Scan(...any) error }, into *string, p *LLMProfile
|
||||
&p.Retry.Stream.Attempts, &p.Retry.Stream.IntervalMS,
|
||||
&p.PriceInputPerM, &p.PriceOutputPerM, &p.PriceCacheReadPerM, &p.PriceCacheWritePerM,
|
||||
&p.UserAgent)
|
||||
if err != nil || into != &p.APIKey {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if p.Proxy, err = openSecret(colLLMProxy, p.Proxy); err != nil || into != &p.APIKey {
|
||||
return err
|
||||
}
|
||||
p.APIKey, err = openSecret(colLLMKey, p.APIKey)
|
||||
@@ -191,12 +194,16 @@ func (d *DB) SaveProfile(p *LLMProfile) (int64, error) {
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
proxy, err := sealSecret(colLLMProxy, p.Proxy)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
r := p.Retry.Clamped()
|
||||
if p.ID == 0 {
|
||||
var id int64
|
||||
err := d.QueryRow(`INSERT INTO llm_profiles(name,format,base_url,proxy,model,api_key,api_key_hint,rate_per_second,rate_per_minute,context_window_k,reasoning_effort,priority,pool_exclude,thinking_type,streaming,max_tokens,max_tokens_field,session_header_key,retry_connect_attempts,retry_connect_interval_ms,retry_empty_attempts,retry_empty_interval_ms,retry_stream_attempts,retry_stream_interval_ms,user_agent)
|
||||
VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22,$23,$24,$25) RETURNING id`,
|
||||
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, key, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ThinkingType, p.Streaming, p.MaxTokens, p.MaxTokensField, p.SessionHeaderKey,
|
||||
p.Name, p.Format, p.BaseURL, proxy, p.Model, key, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ThinkingType, p.Streaming, p.MaxTokens, p.MaxTokensField, p.SessionHeaderKey,
|
||||
r.Connect.Attempts, r.Connect.IntervalMS, r.Empty.Attempts, r.Empty.IntervalMS, r.Stream.Attempts, r.Stream.IntervalMS, p.UserAgent).Scan(&id)
|
||||
if err == nil {
|
||||
_, err = d.saveProfilePrices(id, p)
|
||||
@@ -205,7 +212,7 @@ VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$
|
||||
}
|
||||
if p.APIKey == "" {
|
||||
_, err := d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,rate_per_second=$6,rate_per_minute=$7,context_window_k=$8,reasoning_effort=$9,priority=$10,pool_exclude=$11,thinking_type=$12,streaming=$13,max_tokens=$14,max_tokens_field=$15,session_header_key=$16,retry_connect_attempts=$17,retry_connect_interval_ms=$18,retry_empty_attempts=$19,retry_empty_interval_ms=$20,retry_stream_attempts=$21,retry_stream_interval_ms=$22,user_agent=$23 WHERE id=$24`,
|
||||
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ThinkingType, p.Streaming, p.MaxTokens, p.MaxTokensField, p.SessionHeaderKey,
|
||||
p.Name, p.Format, p.BaseURL, proxy, p.Model, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ThinkingType, p.Streaming, p.MaxTokens, p.MaxTokensField, p.SessionHeaderKey,
|
||||
r.Connect.Attempts, r.Connect.IntervalMS, r.Empty.Attempts, r.Empty.IntervalMS, r.Stream.Attempts, r.Stream.IntervalMS, p.UserAgent, p.ID)
|
||||
if err == nil {
|
||||
_, err = d.saveProfilePrices(p.ID, p)
|
||||
@@ -213,7 +220,7 @@ VALUES ($1,$2,NULLIF($3,''),NULLIF($4,''),$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$
|
||||
return p.ID, err
|
||||
}
|
||||
_, err = d.Exec(`UPDATE llm_profiles SET name=$1,format=$2,base_url=NULLIF($3,''),proxy=NULLIF($4,''),model=$5,api_key=$6,api_key_hint=$7,rate_per_second=$8,rate_per_minute=$9,context_window_k=$10,reasoning_effort=$11,priority=$12,pool_exclude=$13,thinking_type=$14,streaming=$15,max_tokens=$16,max_tokens_field=$17,session_header_key=$18,retry_connect_attempts=$19,retry_connect_interval_ms=$20,retry_empty_attempts=$21,retry_empty_interval_ms=$22,retry_stream_attempts=$23,retry_stream_interval_ms=$24,user_agent=$25 WHERE id=$26`,
|
||||
p.Name, p.Format, p.BaseURL, p.Proxy, p.Model, key, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ThinkingType, p.Streaming, p.MaxTokens, p.MaxTokensField, p.SessionHeaderKey,
|
||||
p.Name, p.Format, p.BaseURL, proxy, p.Model, key, hint, p.RatePerSecond, p.RatePerMinute, p.ContextWindowK, p.ReasoningEffort, p.Priority, p.PoolExclude, p.ThinkingType, p.Streaming, p.MaxTokens, p.MaxTokensField, p.SessionHeaderKey,
|
||||
r.Connect.Attempts, r.Connect.IntervalMS, r.Empty.Attempts, r.Empty.IntervalMS, r.Stream.Attempts, r.Stream.IntervalMS, p.UserAgent, p.ID)
|
||||
if err == nil {
|
||||
_, err = d.saveProfilePrices(p.ID, p)
|
||||
|
||||
+22
-6
@@ -4,7 +4,9 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestPoolProfilesOrder pins the failover chain query: keyless profiles can't
|
||||
@@ -17,17 +19,22 @@ func TestPoolProfilesOrder(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Close via Cleanup, registered FIRST so it runs LAST: a plain defer would
|
||||
// close the pool before the row-deleting cleanups below run.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
// llm_profiles.name is unique, so a fixed name collides with a row left on
|
||||
// the same database by an earlier run.
|
||||
sfx := fmt.Sprintf("-%d", time.Now().UnixNano())
|
||||
mk := func(name string, priority int, exclude bool, key string) int64 {
|
||||
id, err := d.SaveProfile(&LLMProfile{
|
||||
Name: name, Format: "openai", Model: "m", APIKey: key,
|
||||
Name: name + sfx, Format: "openai", Model: "m", APIKey: key,
|
||||
Priority: priority, PoolExclude: exclude,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, id) })
|
||||
t.Cleanup(func() { deleteTestProfile(t, d, id) })
|
||||
return id
|
||||
}
|
||||
lo := mk("t-pool-lo", 1, false, "k1")
|
||||
@@ -42,11 +49,11 @@ func TestPoolProfilesOrder(t *testing.T) {
|
||||
var got []int64
|
||||
for _, p := range chain {
|
||||
switch p.Name {
|
||||
case "t-pool-lo", "t-pool-hi":
|
||||
case "t-pool-lo" + sfx, "t-pool-hi" + sfx:
|
||||
got = append(got, p.ID)
|
||||
case "t-pool-excluded":
|
||||
case "t-pool-excluded" + sfx:
|
||||
t.Fatal("pool_exclude profile entered the failover chain")
|
||||
case "t-pool-nokey":
|
||||
case "t-pool-nokey" + sfx:
|
||||
t.Fatal("keyless profile entered the failover chain")
|
||||
}
|
||||
}
|
||||
@@ -55,6 +62,15 @@ func TestPoolProfilesOrder(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// deleteTestProfile removes a profile a test created and reports a failed delete
|
||||
// instead of silently leaving the row for the next run.
|
||||
func deleteTestProfile(t *testing.T, d *DB, id int64) {
|
||||
t.Helper()
|
||||
if _, err := d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, id); err != nil {
|
||||
t.Errorf("cleanup profile %d: %v", id, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteProfileContextHonorsCancellation(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
@@ -11,7 +11,8 @@ func TestConversationPinOrderingAndPatch(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
suffix := time.Now().UnixNano()
|
||||
first, err := d.CreateConversation("mainagent", fmt.Sprintf("pin-first-%d", suffix), nil)
|
||||
@@ -24,8 +25,8 @@ func TestConversationPinOrderingAndPatch(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = d.DeleteConversation(first.ID)
|
||||
_ = d.DeleteConversation(second.ID)
|
||||
cleanupErr(t, "first conversation", d.DeleteConversation(first.ID))
|
||||
cleanupErr(t, "second conversation", d.DeleteConversation(second.ID))
|
||||
})
|
||||
|
||||
pinned := true
|
||||
|
||||
@@ -102,12 +102,13 @@ func TestRecordCoverageUpsertHistory(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
task, err := d.CreateTaskWithOptions("coverage ledger", "g", TaskCreateOptions{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "task", d.DeleteTask(task.ID)) })
|
||||
as := d.Assets()
|
||||
store := d.Exploration(task.ExplorationID)
|
||||
for _, vc := range []string{"SQLi", "XSS"} {
|
||||
@@ -158,12 +159,13 @@ func TestCoverageCompletenessBudgetThenReviveDone(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
task, err := d.CreateTaskWithOptions("coverage revive", "g", TaskCreateOptions{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "task", d.DeleteTask(task.ID)) })
|
||||
as, store := d.Assets(), d.Exploration(task.ExplorationID)
|
||||
state := func() string {
|
||||
t.Helper()
|
||||
@@ -221,12 +223,13 @@ func TestRecordCoverageRowCap(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
task, err := d.CreateTaskWithOptions("coverage cap", "g", TaskCreateOptions{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "task", d.DeleteTask(task.ID)) })
|
||||
if _, err := d.Exec(`INSERT INTO coverage_ledger(task_id, surface, risk_area, outcome)
|
||||
SELECT $1, '/s' || g, 'xss', 'needs_follow_up' FROM generate_series(1, $2::int) g`, task.ID, MaxCoverageRowsPerTask); err != nil {
|
||||
t.Fatal(err)
|
||||
|
||||
@@ -12,7 +12,8 @@ func TestCustomToolCRUD(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
key := "ct_test_tool"
|
||||
_ = d.DeleteCustomTool(key) // clean slate
|
||||
@@ -30,7 +31,7 @@ func TestCustomToolCRUD(t *testing.T) {
|
||||
if err := d.CreateCustomTool(in); err != nil {
|
||||
t.Fatalf("create: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteCustomTool(key) })
|
||||
t.Cleanup(func() { cleanupErr(t, "custom tool", d.DeleteCustomTool(key)) })
|
||||
|
||||
got, err := d.GetTool(key)
|
||||
if err != nil || got == nil {
|
||||
|
||||
@@ -0,0 +1,245 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
const (
|
||||
evCookie = "session=ev-cookie-secret"
|
||||
evBearer = "Bearer ev-auth-secret"
|
||||
)
|
||||
|
||||
// evidenceFixture records a finding. With legacy=false the snapshot goes through
|
||||
// the normal insert path; with legacy=true it is written in clear by raw SQL.
|
||||
func evidenceFixture(t *testing.T, d *DB, tag string, legacy bool) (*Task, int64, TrafficEvidenceSnapshot) {
|
||||
t.Helper()
|
||||
task, err := d.CreateTask("evidence heads at rest "+tag, "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s := TrafficEvidenceSnapshot{SourceTrafficID: "t-sec-ev-" + tag, CapturedAt: 1, URL: "http://ev.example/" + tag, Method: "GET", Status: 200,
|
||||
ReqHead: "GET /" + tag + " HTTP/1.1\nCookie: " + evCookie + "\nAuthorization: " + evBearer + "\n",
|
||||
RespHead: "HTTP/1.1 200 OK\nSet-Cookie: " + evCookie + "\n", ReqHash: "rh-" + tag, RespHash: "ph-" + tag}
|
||||
s.ID = TrafficSnapshotID(s)
|
||||
t.Cleanup(func() {
|
||||
d.Exec(`DELETE FROM task_archives WHERE task_id=$1`, task.ID)
|
||||
d.DeleteTaskCascade(task.ID, false, true) // findings and their bindings too
|
||||
if _, err := d.Exec(`DELETE FROM traffic_evidence_snapshots WHERE id=$1`, s.ID); err != nil {
|
||||
t.Errorf("cleanup snapshot: %v", err)
|
||||
}
|
||||
})
|
||||
f, err := d.Exploration(task.ExplorationID).RecordFinding(t.Context(), RecordFindingInput{TaskID: task.ID, ExplorationID: task.ExplorationID, Summary: "evidence " + tag, Severity: "low"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if legacy {
|
||||
_, err = d.Exec(`INSERT INTO traffic_evidence_snapshots(id,source_traffic_id,captured_at,url,method,status,req_head,resp_head,req_hash,resp_hash,req_len,resp_len)
|
||||
VALUES($1,$2,1,$3,'GET',200,$4,$5,$6,$7,0,0)`, s.ID, s.SourceTrafficID, s.URL, s.ReqHead, s.RespHead, s.ReqHash, s.RespHash)
|
||||
if err == nil {
|
||||
_, err = d.Exec(`INSERT INTO finding_traffic_bindings(finding_id,snapshot_id,position) VALUES($1,$2,0)`, f.FindingID, s.ID)
|
||||
}
|
||||
} else {
|
||||
err = d.WithEvidenceTx(t.Context(), func(tx *sql.Tx) error {
|
||||
return AddFindingTrafficTx(tx, f.FindingID, []PreparedTrafficEvidence{{Ref: TrafficRef{TrafficID: s.SourceTrafficID, Role: "proof"}, Snapshot: s}})
|
||||
})
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return task, f.FindingID, s
|
||||
}
|
||||
|
||||
// rawHeads reads the stored heads; both must be sealed and hide the secrets.
|
||||
func rawHeads(t *testing.T, d *DB, id string) (req, resp string) {
|
||||
t.Helper()
|
||||
if err := d.QueryRow(`SELECT req_head,resp_head FROM traffic_evidence_snapshots WHERE id=$1`, id).Scan(&req, &resp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, v := range []string{req, resp} {
|
||||
if !isEncrypted(v) || strings.Contains(v, evCookie) || strings.Contains(v, evBearer) {
|
||||
t.Fatalf("head stored in clear: %q", v)
|
||||
}
|
||||
}
|
||||
return req, resp
|
||||
}
|
||||
|
||||
// assertClearHeads checks the application read and that the stored value is
|
||||
// sealed exactly once.
|
||||
func assertClearHeads(t *testing.T, d *DB, findingID int64, want TrafficEvidenceSnapshot) {
|
||||
t.Helper()
|
||||
got, err := d.GetFindingTraffic(t.Context(), findingID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(got.Bindings) != 1 || got.Bindings[0].Snapshot.ReqHead != want.ReqHead || got.Bindings[0].Snapshot.RespHead != want.RespHead || got.Bindings[0].SnapshotID != want.ID {
|
||||
t.Fatalf("read back: %+v", got.Bindings)
|
||||
}
|
||||
req, resp := rawHeads(t, d, want.ID)
|
||||
box := activeSecretBox.Load()
|
||||
if p, err := box.decrypt(colEvReqHead.aad(), req); err != nil || p != want.ReqHead {
|
||||
t.Fatalf("req_head not sealed once: %q %v", p, err)
|
||||
}
|
||||
if p, err := box.decrypt(colEvRespHead.aad(), resp); err != nil || p != want.RespHead {
|
||||
t.Fatalf("resp_head not sealed once: %q %v", p, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvidenceHeadsSealedAtInsert(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
t.Cleanup(func() { d.Close() })
|
||||
_, fid, s := evidenceFixture(t, d, "insert", false)
|
||||
assertClearHeads(t, d, fid, s)
|
||||
var url, method string
|
||||
if err := d.QueryRow(`SELECT url,method FROM traffic_evidence_snapshots WHERE id=$1`, s.ID).Scan(&url, &method); err != nil || url != s.URL || method != "GET" {
|
||||
t.Fatalf("metadata must stay in clear: %q %q %v", url, method, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvidenceHeadsBackfillAndKeyChecks(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
t.Cleanup(func() { d.Close() })
|
||||
_, fid, s := evidenceFixture(t, d, "legacy", true)
|
||||
t.Setenv(secretKeyEnv, base64.StdEncoding.EncodeToString(testSecretKey))
|
||||
if err := d.InitSecrets(t.TempDir()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertClearHeads(t, d, fid, s)
|
||||
req, resp := rawHeads(t, d, s.ID)
|
||||
if err := d.InitSecrets(t.TempDir()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if r2, p2 := rawHeads(t, d, s.ID); r2 != req || p2 != resp {
|
||||
t.Fatal("second backfill must leave the ciphertexts untouched")
|
||||
}
|
||||
|
||||
// A head sealed under another key: the read and the startup check fail and
|
||||
// name the column. The startup check samples one sealed row per column, so
|
||||
// this row must be the only sealed snapshot.
|
||||
var sealedRows int
|
||||
if err := d.QueryRow(`SELECT count(*) FROM traffic_evidence_snapshots WHERE req_head LIKE 'enc:v1:%' OR resp_head LIKE 'enc:v1:%'`).Scan(&sealedRows); err != nil || sealedRows != 1 {
|
||||
t.Fatalf("want only this fixture's snapshot sealed, got %d (%v)", sealedRows, err)
|
||||
}
|
||||
other, _ := newSecretBox([]byte("ffffffffffffffffffffffffffffffff"))
|
||||
for _, c := range []secretColumn{colEvReqHead, colEvRespHead} {
|
||||
bad, _ := other.encrypt(c.aad(), "GET / HTTP/1.1\nCookie: other\n")
|
||||
if _, err := d.Exec(`UPDATE traffic_evidence_snapshots SET `+c.col+`=$1 WHERE id=$2`, bad, s.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := d.GetFindingTraffic(t.Context(), fid); err == nil || !strings.Contains(err.Error(), c.aad()) {
|
||||
t.Fatalf("want a read failure naming %s, got %v", c.aad(), err)
|
||||
}
|
||||
if err := d.InitSecrets(t.TempDir()); err == nil || !strings.Contains(err.Error(), "secret key check") || !strings.Contains(err.Error(), c.aad()) {
|
||||
t.Fatalf("want a key-check failure naming %s, got %v", c.aad(), err)
|
||||
}
|
||||
if _, err := d.Exec(`UPDATE traffic_evidence_snapshots SET req_head=$1,resp_head=$2 WHERE id=$3`, req, resp, s.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
assertClearHeads(t, d, fid, s)
|
||||
|
||||
// No key with sealed heads: refuse, and each head counts. One head stays
|
||||
// sealed so the refusal never turns into a new key.
|
||||
t.Setenv(secretKeyEnv, "")
|
||||
dir := t.TempDir()
|
||||
withHeads := encryptedCount(t, d, dir)
|
||||
for _, c := range []secretColumn{colEvReqHead, colEvRespHead} {
|
||||
if _, err := d.Exec(`UPDATE traffic_evidence_snapshots SET req_head=$1,resp_head=$2 WHERE id=$3`, req, resp, s.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := d.Exec(`UPDATE traffic_evidence_snapshots SET `+c.col+`='' WHERE id=$1`, s.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n := encryptedCount(t, d, dir); withHeads-n != 1 {
|
||||
t.Fatalf("no-key refusal counted %d then %d encrypted values; %s must count", withHeads, n, c.aad())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Export keeps the heads sealed; import opens them to check the snapshot ID,
|
||||
// then the insert seals them once. A legacy clear archive is sealed on import.
|
||||
func TestEvidenceHeadsArchiveRoundTrip(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
t.Cleanup(func() { d.Close() })
|
||||
if err = d.EnsureLLMRecordsTable(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = d.EnsureLLMUsageTable(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, legacy := range []bool{false, true} {
|
||||
t.Run(fmt.Sprint("legacy=", legacy), func(t *testing.T) {
|
||||
task, fid, s := evidenceFixture(t, d, fmt.Sprint("archive-", legacy), false)
|
||||
if err := d.SetPaused(task.ID, true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
archive, err := d.QueueTaskArchive(task.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = d.ClaimTaskArchiveJob(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
snapshot, err := d.SnapshotTaskArchive(task.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
exported := string(snapshot.Tables["traffic_evidence_snapshots"])
|
||||
if !strings.Contains(exported, secretPrefix) || strings.Contains(exported, evCookie) || strings.Contains(exported, evBearer) {
|
||||
t.Fatalf("archive must keep the heads sealed: %s", exported)
|
||||
}
|
||||
if err = d.CompleteTaskArchive(archive.ID, snapshot, "/tmp/evidence-heads-test.tar.zst", "fixture", 1, 1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var rows []map[string]any
|
||||
if err = json.Unmarshal(snapshot.Tables["traffic_evidence_snapshots"], &rows); err != nil || len(rows) != 1 {
|
||||
t.Fatal(rows, err)
|
||||
}
|
||||
if legacy { // an archive written before the heads were sealed
|
||||
rows[0]["req_head"], rows[0]["resp_head"] = s.ReqHead, s.RespHead
|
||||
snapshot.Tables["traffic_evidence_snapshots"], _ = json.Marshal(rows)
|
||||
} else { // a head sealed under another key fails the import, naming the column
|
||||
other, _ := newSecretBox([]byte("ffffffffffffffffffffffffffffffff"))
|
||||
bad, _ := other.encrypt(colEvReqHead.aad(), s.ReqHead)
|
||||
good := rows[0]["req_head"]
|
||||
rows[0]["req_head"] = bad
|
||||
raw, _ := json.Marshal(rows)
|
||||
if _, err := ArchiveEvidenceSnapshots(&TaskArchiveSnapshot{FormatVersion: snapshot.FormatVersion, Tables: map[string]json.RawMessage{"traffic_evidence_snapshots": raw}}); err == nil || !strings.Contains(err.Error(), colEvReqHead.aad()) {
|
||||
t.Fatalf("want an import failure naming %s, got %v", colEvReqHead.aad(), err)
|
||||
}
|
||||
rows[0]["req_head"] = good
|
||||
}
|
||||
list, err := ArchiveEvidenceSnapshots(snapshot)
|
||||
if err != nil || len(list) != 1 || list[0].ReqHead != s.ReqHead || list[0].RespHead != s.RespHead {
|
||||
t.Fatalf("archive snapshots: %+v %v", list, err)
|
||||
}
|
||||
// The hot row is gone (GC), so the restore inserts it again.
|
||||
if _, err = d.Exec(`DELETE FROM traffic_evidence_snapshots WHERE id=$1`, s.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = d.QueueTaskArchiveRestore(archive.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = d.ClaimTaskArchiveJob(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = d.RestoreTaskArchive(archive.ID, snapshot, 0); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertClearHeads(t, d, fid, s)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -10,13 +10,14 @@ func TestInheritedActivityReadsRequireTerminalIntent(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanups below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
expID, err := d.CreateExploration("source activity boundary", "source activity boundary")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _, _ = d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) })
|
||||
t.Cleanup(func() { cleanupExec(t, d, `DELETE FROM explorations WHERE id=$1`, expID) })
|
||||
store := d.Exploration(expID)
|
||||
intentID, err := store.AddIntent(map[string]any{"summary": "source work"}, 1, nil, "planner")
|
||||
if err != nil {
|
||||
@@ -93,7 +94,8 @@ func TestExplorationDirectSourceReadView(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanups below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
grand, err := d.CreateTask("grand source", "grand goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
@@ -112,9 +114,9 @@ func TestExplorationDirectSourceReadView(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = d.DeleteTask(current.ID)
|
||||
_ = d.DeleteTask(source.ID)
|
||||
_ = d.DeleteTask(grand.ID)
|
||||
cleanupErr(t, "current task", d.DeleteTask(current.ID))
|
||||
cleanupErr(t, "source task", d.DeleteTask(source.ID))
|
||||
cleanupErr(t, "grand task", d.DeleteTask(grand.ID))
|
||||
})
|
||||
|
||||
grandStore := d.Exploration(grand.ExplorationID)
|
||||
@@ -315,7 +317,8 @@ func TestTaskAssetContextUsesDirectSourceScopeAndAnchors(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanups below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
grand, err := d.CreateTask("grand assets", "grand goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
@@ -351,10 +354,11 @@ func TestTaskAssetContextUsesDirectSourceScopeAndAnchors(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = d.DeleteTask(current.ID)
|
||||
_ = d.DeleteTask(source.ID)
|
||||
_ = d.DeleteTask(grand.ID)
|
||||
_, _ = assets.DeleteByIDs([]int64{sourceTestedID, sourceUntestedID, sourceAnchoredOnlyID, grandID})
|
||||
cleanupErr(t, "current task", d.DeleteTask(current.ID))
|
||||
cleanupErr(t, "source task", d.DeleteTask(source.ID))
|
||||
cleanupErr(t, "grand task", d.DeleteTask(grand.ID))
|
||||
_, err := assets.DeleteByIDs([]int64{sourceTestedID, sourceUntestedID, sourceAnchoredOnlyID, grandID})
|
||||
cleanupErr(t, "assets", err)
|
||||
})
|
||||
if _, err := assets.AddAgentScope(source.ID, "root_domain", sourceTestedDomain, "test", "agent"); err != nil {
|
||||
t.Fatal(err)
|
||||
|
||||
+24
-2
@@ -195,10 +195,29 @@ func InsertEvidenceSnapshotTx(tx *sql.Tx, s TrafficEvidenceSnapshot) error {
|
||||
if s.CreatedAt.IsZero() {
|
||||
s.CreatedAt = time.Now().UTC()
|
||||
}
|
||||
_, err := tx.Exec(`INSERT INTO traffic_evidence_snapshots
|
||||
// The ID hashes the clear heads; only the stored heads are sealed.
|
||||
reqHead, err := sealSecret(colEvReqHead, s.ReqHead)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
respHead, err := sealSecret(colEvRespHead, s.RespHead)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = tx.Exec(`INSERT INTO traffic_evidence_snapshots
|
||||
(id,source_traffic_id,captured_at,url,method,status,content_type,req_head,resp_head,req_hash,resp_hash,req_len,resp_len,created_at)
|
||||
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14) ON CONFLICT(id) DO NOTHING`,
|
||||
s.ID, s.SourceTrafficID, s.CapturedAt, s.URL, s.Method, s.Status, s.ContentType, s.ReqHead, s.RespHead, s.ReqHash, s.RespHash, s.ReqLen, s.RespLen, s.CreatedAt)
|
||||
s.ID, s.SourceTrafficID, s.CapturedAt, s.URL, s.Method, s.Status, s.ContentType, reqHead, respHead, s.ReqHash, s.RespHash, s.ReqLen, s.RespLen, s.CreatedAt)
|
||||
return err
|
||||
}
|
||||
|
||||
// openHeads decrypts the heads of a snapshot read from the table or an archive.
|
||||
// Clear legacy values pass through; a value that fails to decrypt is an error.
|
||||
func (s *TrafficEvidenceSnapshot) openHeads() (err error) {
|
||||
if s.ReqHead, err = openSecret(colEvReqHead, s.ReqHead); err != nil {
|
||||
return err
|
||||
}
|
||||
s.RespHead, err = openSecret(colEvRespHead, s.RespHead)
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -263,6 +282,9 @@ FROM finding_traffic_bindings b JOIN traffic_evidence_snapshots s ON s.id=b.snap
|
||||
if err := json.Unmarshal(raw, &b.Snapshot); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := b.Snapshot.openHeads(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out.Bindings = append(out.Bindings, b)
|
||||
}
|
||||
return out, rows.Err()
|
||||
|
||||
@@ -15,8 +15,17 @@ func ArchiveEvidenceSnapshots(snapshot *TaskArchiveSnapshot) ([]TrafficEvidenceS
|
||||
if snapshot.FormatVersion < 3 {
|
||||
return nil, ErrTaskArchiveFormatMismatch
|
||||
}
|
||||
err := json.Unmarshal(snapshot.Tables["traffic_evidence_snapshots"], &out)
|
||||
return out, err
|
||||
// The archive keeps the heads as stored (sealed). They are opened here so the
|
||||
// restore can check each content-addressed ID; the insert then seals them once.
|
||||
if err := json.Unmarshal(snapshot.Tables["traffic_evidence_snapshots"], &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for i := range out {
|
||||
if err := out[i].openHeads(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func restoreFindingTrafficTx(tx *sql.Tx, snapshot *TaskArchiveSnapshot) error {
|
||||
|
||||
+17
-9
@@ -1,6 +1,7 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
@@ -14,7 +15,8 @@ func TestProfileRetryRoundTrip(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
want := RetryOverride{
|
||||
Connect: RetryRule{Attempts: 5, IntervalMS: 2000},
|
||||
@@ -22,12 +24,12 @@ func TestProfileRetryRoundTrip(t *testing.T) {
|
||||
Stream: RetryRule{IntervalMS: 250},
|
||||
}
|
||||
id, err := d.SaveProfile(&LLMProfile{
|
||||
Name: "t-retry-roundtrip", Format: "openai", Model: "m", APIKey: "k", Retry: want,
|
||||
Name: fmt.Sprintf("t-retry-roundtrip-%d", time.Now().UnixNano()), Format: "openai", Model: "m", APIKey: "k", Retry: want,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, id) })
|
||||
t.Cleanup(func() { deleteTestProfile(t, d, id) })
|
||||
|
||||
got, err := d.ProfileByID(id)
|
||||
if err != nil || got == nil {
|
||||
@@ -71,16 +73,17 @@ func TestProfileRetryClampedOnSave(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
id, err := d.SaveProfile(&LLMProfile{
|
||||
Name: "t-retry-clamp", Format: "openai", Model: "m", APIKey: "k",
|
||||
Name: fmt.Sprintf("t-retry-clamp-%d", time.Now().UnixNano()), Format: "openai", Model: "m", APIKey: "k",
|
||||
Retry: RetryOverride{Connect: RetryRule{Attempts: -99, IntervalMS: -5}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, id) })
|
||||
t.Cleanup(func() { deleteTestProfile(t, d, id) })
|
||||
got, err := d.ProfileByID(id)
|
||||
if err != nil || got == nil {
|
||||
t.Fatalf("ProfileByID: %v", err)
|
||||
@@ -97,17 +100,22 @@ func TestLLMRetryPolicyRoundTrip(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the restore below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
before, hadBefore, err := d.GetSetting(settingLLMRetryPolicy)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
var err error
|
||||
if hadBefore {
|
||||
d.SetSetting(settingLLMRetryPolicy, before)
|
||||
err = d.SetSetting(settingLLMRetryPolicy, before)
|
||||
} else {
|
||||
d.Exec(`DELETE FROM settings WHERE key=$1`, settingLLMRetryPolicy)
|
||||
_, err = d.Exec(`DELETE FROM settings WHERE key=$1`, settingLLMRetryPolicy)
|
||||
}
|
||||
if err != nil {
|
||||
t.Errorf("restore %s: %v", settingLLMRetryPolicy, err)
|
||||
}
|
||||
})
|
||||
|
||||
|
||||
+71
-15
@@ -19,7 +19,8 @@ import (
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// Secrets at rest: AES-256-GCM, a random nonce per value, AAD "<table>.<column>".
|
||||
// Secrets at rest: AES-256-GCM, a random nonce per value, AAD "<table>.<column>"
|
||||
// ("settings.<key>" for the secret settings).
|
||||
// TEXT stores "enc:v1:<base64(nonce||ciphertext)>"; JSONB stores
|
||||
// {"__artex_enc":"enc:v1:..."} over the canonical JSON. Empty values stay empty.
|
||||
// A value that fails to decrypt is an error, never returned as-is.
|
||||
@@ -34,9 +35,36 @@ const (
|
||||
type secretColumn struct {
|
||||
table, col string
|
||||
json bool
|
||||
setting string // a settings row: its key; "" for a table column
|
||||
}
|
||||
|
||||
func (c secretColumn) aad() string { return c.table + "." + c.col }
|
||||
func (c secretColumn) aad() string {
|
||||
if c.setting != "" {
|
||||
return "settings." + c.setting
|
||||
}
|
||||
return c.table + "." + c.col
|
||||
}
|
||||
|
||||
// rowFilter limits a settings entry to its own row; the key is a code-owned
|
||||
// constant from secretSettingKeys, never input. tools.exec covers the
|
||||
// user-created tools only: the seeds rewrite the built-in rows in clear.
|
||||
func (c secretColumn) rowFilter() string {
|
||||
if c.setting != "" {
|
||||
return ` AND key='` + c.setting + `'`
|
||||
}
|
||||
if c == colToolExec {
|
||||
return ` AND system=false`
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// idCol names the row key: settings and tools rows are keyed by name, the rest by id.
|
||||
func (c secretColumn) idCol() string {
|
||||
if c.setting != "" || c == colToolExec {
|
||||
return "key"
|
||||
}
|
||||
return "id"
|
||||
}
|
||||
|
||||
// valueExpr is the SQL text that holds the ciphertext, if any.
|
||||
func (c secretColumn) valueExpr() string {
|
||||
@@ -47,13 +75,41 @@ func (c secretColumn) valueExpr() string {
|
||||
}
|
||||
|
||||
var (
|
||||
colLLMKey = secretColumn{"llm_profiles", "api_key", false}
|
||||
colMCPEnv = secretColumn{"mcp_servers", "env", true}
|
||||
colCredPass = secretColumn{"task_credentials", "password", false}
|
||||
colNotifConfig = secretColumn{"notification_channels", "config", true}
|
||||
secretColumns = []secretColumn{colLLMKey, colMCPEnv, colCredPass, colNotifConfig}
|
||||
colLLMKey = secretColumn{"llm_profiles", "api_key", false, ""}
|
||||
colLLMProxy = secretColumn{"llm_profiles", "proxy", false, ""} // may carry user:pass
|
||||
colMCPEnv = secretColumn{"mcp_servers", "env", true, ""}
|
||||
colCredPass = secretColumn{"task_credentials", "password", false, ""}
|
||||
colNotifConfig = secretColumn{"notification_channels", "config", true, ""}
|
||||
colToolExec = secretColumn{"tools", "exec", true, ""} // system=false rows: headers, proxy, body, code, command
|
||||
// Evidence head blocks carry the captured Cookie and Authorization headers.
|
||||
colEvReqHead = secretColumn{"traffic_evidence_snapshots", "req_head", false, ""}
|
||||
colEvRespHead = secretColumn{"traffic_evidence_snapshots", "resp_head", false, ""}
|
||||
secretColumns = append([]secretColumn{colLLMKey, colLLMProxy, colMCPEnv, colCredPass, colNotifConfig, colToolExec, colEvReqHead, colEvRespHead}, secretSettingColumns()...)
|
||||
)
|
||||
|
||||
// secretSettingKeys are the settings stored encrypted, AAD "settings.<key>".
|
||||
// Every other setting stays in clear.
|
||||
var secretSettingKeys = []string{"brave_search_api_key", "tavily_search_api_key", "web_search_proxy", "global_proxy"}
|
||||
|
||||
func secretSettingColumns() []secretColumn {
|
||||
out := make([]secretColumn, 0, len(secretSettingKeys))
|
||||
for _, k := range secretSettingKeys {
|
||||
out = append(out, secretColumn{"settings", "value", false, k})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// secretSetting returns the secret column of a settings key, ok=false when the
|
||||
// key is stored in clear.
|
||||
func secretSetting(key string) (secretColumn, bool) {
|
||||
for _, c := range secretColumns {
|
||||
if c.setting != "" && c.setting == key {
|
||||
return c, true
|
||||
}
|
||||
}
|
||||
return secretColumn{}, false
|
||||
}
|
||||
|
||||
type secretBox struct{ aead cipher.AEAD }
|
||||
|
||||
var activeSecretBox atomic.Pointer[secretBox]
|
||||
@@ -203,7 +259,7 @@ func (d *DB) InitSecrets(keyDir string) error {
|
||||
// A key other than the one that sealed the data fails here, at startup.
|
||||
for _, c := range secretColumns {
|
||||
var v string
|
||||
err := conn.QueryRowContext(ctx, `SELECT `+c.valueExpr()+` FROM `+c.table+` WHERE `+c.valueExpr()+` LIKE 'enc:v1:%' LIMIT 1`).Scan(&v)
|
||||
err := conn.QueryRowContext(ctx, `SELECT `+c.valueExpr()+` FROM `+c.table+` WHERE `+c.valueExpr()+` LIKE 'enc:v1:%'`+c.rowFilter()+` LIMIT 1`).Scan(&v)
|
||||
if err == nil {
|
||||
_, err = box.decrypt(c.aad(), v)
|
||||
}
|
||||
@@ -241,7 +297,7 @@ func loadSecretKey(ctx context.Context, conn *sql.Conn, keyDir string) ([]byte,
|
||||
}
|
||||
var parts []string
|
||||
for _, c := range secretColumns {
|
||||
parts = append(parts, `(SELECT count(*) FROM `+c.table+` WHERE `+c.valueExpr()+` LIKE 'enc:v1:%')`)
|
||||
parts = append(parts, `(SELECT count(*) FROM `+c.table+` WHERE `+c.valueExpr()+` LIKE 'enc:v1:%'`+c.rowFilter()+`)`)
|
||||
}
|
||||
var n int64
|
||||
if err := conn.QueryRowContext(ctx, `SELECT `+strings.Join(parts, "+")).Scan(&n); err != nil {
|
||||
@@ -279,14 +335,14 @@ func loadSecretKey(ctx context.Context, conn *sql.Conn, keyDir string) ([]byte,
|
||||
func backfillSecrets(ctx context.Context, conn *sql.Conn, box *secretBox) (int, error) {
|
||||
total := 0
|
||||
for _, c := range secretColumns {
|
||||
var ids []int64
|
||||
var ids []any // int64 ids, or the settings key
|
||||
var olds, news []string
|
||||
rows, err := conn.QueryContext(ctx, `SELECT id, `+c.col+`::text FROM `+c.table+` WHERE `+c.col+` IS NOT NULL`)
|
||||
rows, err := conn.QueryContext(ctx, `SELECT `+c.idCol()+`, `+c.col+`::text FROM `+c.table+` WHERE `+c.col+` IS NOT NULL`+c.rowFilter())
|
||||
if err != nil {
|
||||
return total, err
|
||||
}
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
var id any
|
||||
var old string
|
||||
if err = rows.Scan(&id, &old); err != nil {
|
||||
break
|
||||
@@ -296,7 +352,7 @@ func backfillSecrets(ctx context.Context, conn *sql.Conn, box *secretBox) (int,
|
||||
}
|
||||
nv, serr := box.seal(c, old)
|
||||
if serr != nil {
|
||||
err = fmt.Errorf("backfill %s id=%d: %w", c.aad(), id, serr)
|
||||
err = fmt.Errorf("backfill %s id=%v: %w", c.aad(), id, serr)
|
||||
break
|
||||
}
|
||||
if isSealed(c, nv) { // empty values stay as they are
|
||||
@@ -315,9 +371,9 @@ func backfillSecrets(ctx context.Context, conn *sql.Conn, box *secretBox) (int,
|
||||
cast = "::jsonb"
|
||||
}
|
||||
for i, id := range ids {
|
||||
res, err := conn.ExecContext(ctx, `UPDATE `+c.table+` SET `+c.col+`=$1`+cast+` WHERE id=$2 AND `+c.col+`::text=$3`, news[i], id, olds[i])
|
||||
res, err := conn.ExecContext(ctx, `UPDATE `+c.table+` SET `+c.col+`=$1`+cast+` WHERE `+c.idCol()+`=$2 AND `+c.col+`::text=$3`, news[i], id, olds[i])
|
||||
if err != nil {
|
||||
return total, fmt.Errorf("backfill %s id=%d: %w", c.aad(), id, err)
|
||||
return total, fmt.Errorf("backfill %s id=%v: %w", c.aad(), id, err)
|
||||
}
|
||||
if n, _ := res.RowsAffected(); n > 0 {
|
||||
total++
|
||||
|
||||
+140
-10
@@ -5,6 +5,7 @@ import (
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
@@ -54,7 +55,20 @@ func TestSecretBoxCrypto(t *testing.T) {
|
||||
|
||||
type secretRows struct{ llm, mcp, cred, ch, task int64 }
|
||||
|
||||
var plainSecrets = []string{"sk-plain-llm", "mcp-plain-tok", "plain-pass", "hook-plain-secret"}
|
||||
var plainSecrets = []string{"sk-plain-llm", "mcp-plain-tok", "plain-pass", "hook-plain-secret", "llmproxy-plain-pass",
|
||||
"brave-plain", "tavily-plain", "wsproxy-plain-pass", "gproxy-plain-pass"}
|
||||
|
||||
// secretSettingValues are the four secret settings written by the fixture.
|
||||
var secretSettingValues = map[string]string{
|
||||
"brave_search_api_key": "brave-plain",
|
||||
"tavily_search_api_key": "tavily-plain",
|
||||
"web_search_proxy": "socks5://u:wsproxy-plain-pass@ws.example:1080",
|
||||
"global_proxy": "http://u:gproxy-plain-pass@gp.example:3128",
|
||||
}
|
||||
|
||||
const llmProxyPlain = "http://u:llmproxy-plain-pass@llm-proxy.example:3128"
|
||||
|
||||
const deleteSecretSettings = `DELETE FROM settings WHERE key IN ('brave_search_api_key','tavily_search_api_key','web_search_proxy','global_proxy')`
|
||||
|
||||
// secretFixture writes one secret to each location, through the db API or, for
|
||||
// a legacy database, as raw plaintext SQL.
|
||||
@@ -70,6 +84,7 @@ func secretFixture(t *testing.T, d *DB, legacy bool) (r secretRows) {
|
||||
d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, r.llm)
|
||||
d.Exec(`DELETE FROM mcp_servers WHERE id=$1`, r.mcp)
|
||||
d.Exec(`DELETE FROM notification_channels WHERE id=$1`, r.ch)
|
||||
d.Exec(deleteSecretSettings)
|
||||
})
|
||||
env, cfg := `{"Authorization":"mcp-plain-tok"}`, `{"url":"https://h.example","headers":{"X":"hook-plain-secret"}}`
|
||||
if legacy {
|
||||
@@ -78,13 +93,23 @@ func secretFixture(t *testing.T, d *DB, legacy bool) (r secretRows) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
ins(&r.llm, `INSERT INTO llm_profiles(name,format,model,api_key) VALUES ('t-sec','openai','m','sk-plain-llm') RETURNING id`)
|
||||
ins(&r.llm, `INSERT INTO llm_profiles(name,format,model,api_key,proxy) VALUES ('t-sec','openai','m','sk-plain-llm',$1) RETURNING id`, llmProxyPlain)
|
||||
for k, v := range secretSettingValues {
|
||||
if _, err := d.Exec(`INSERT INTO settings(key,value) VALUES ($1,$2)`, k, v); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
ins(&r.mcp, `INSERT INTO mcp_servers(name,transport,env) VALUES ('t-sec','http',$1) RETURNING id`, env)
|
||||
ins(&r.cred, `INSERT INTO task_credentials(task_id,domain,username,password) VALUES ($1,'x.example','u','plain-pass') RETURNING id`, r.task)
|
||||
ins(&r.ch, `INSERT INTO notification_channels(name,kind,config) VALUES ('t-sec','webhook',$1) RETURNING id`, cfg)
|
||||
return r
|
||||
}
|
||||
r.llm, err = d.SaveProfile(&LLMProfile{Name: "t-sec", Format: "openai", Model: "m", APIKey: "sk-plain-llm"})
|
||||
r.llm, err = d.SaveProfile(&LLMProfile{Name: "t-sec", Format: "openai", Model: "m", APIKey: "sk-plain-llm", Proxy: llmProxyPlain})
|
||||
for k, v := range secretSettingValues {
|
||||
if err == nil {
|
||||
err = d.SetSetting(k, v)
|
||||
}
|
||||
}
|
||||
if err == nil {
|
||||
r.mcp, err = d.SaveMCP(&MCPServer{Name: "t-sec", Transport: "http", Env: json.RawMessage(env)})
|
||||
}
|
||||
@@ -102,19 +127,28 @@ func secretFixture(t *testing.T, d *DB, legacy bool) (r secretRows) {
|
||||
return r
|
||||
}
|
||||
|
||||
// rawSecrets reads the four columns as stored in PostgreSQL.
|
||||
// rawSecrets reads the secret columns and settings as stored in PostgreSQL.
|
||||
func rawSecrets(t *testing.T, d *DB, r secretRows) string {
|
||||
t.Helper()
|
||||
var a, b, c, e string
|
||||
var a, b, c, e, px string
|
||||
if err := d.QueryRow(`SELECT (SELECT api_key FROM llm_profiles WHERE id=$1), (SELECT env::text FROM mcp_servers WHERE id=$2),
|
||||
(SELECT password FROM task_credentials WHERE id=$3), (SELECT config::text FROM notification_channels WHERE id=$4)`,
|
||||
r.llm, r.mcp, r.cred, r.ch).Scan(&a, &b, &c, &e); err != nil {
|
||||
(SELECT password FROM task_credentials WHERE id=$3), (SELECT config::text FROM notification_channels WHERE id=$4),
|
||||
(SELECT proxy FROM llm_profiles WHERE id=$1)`,
|
||||
r.llm, r.mcp, r.cred, r.ch).Scan(&a, &b, &c, &e, &px); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !isEncrypted(a) || !isSealed(colMCPEnv, b) || !isEncrypted(c) || !isSealed(colNotifConfig, e) {
|
||||
if !isEncrypted(a) || !isSealed(colMCPEnv, b) || !isEncrypted(c) || !isSealed(colNotifConfig, e) || !isEncrypted(px) {
|
||||
t.Fatal("a stored secret is not encrypted")
|
||||
}
|
||||
all := strings.Join([]string{a, b, c, e}, "|")
|
||||
vals := []string{a, b, c, e, px}
|
||||
for _, k := range secretSettingKeys {
|
||||
var v string
|
||||
if err := d.QueryRow(`SELECT value FROM settings WHERE key=$1`, k).Scan(&v); err != nil || !isEncrypted(v) {
|
||||
t.Fatalf("setting %s not encrypted at rest (%v)", k, err)
|
||||
}
|
||||
vals = append(vals, v)
|
||||
}
|
||||
all := strings.Join(vals, "|")
|
||||
for _, p := range plainSecrets {
|
||||
if strings.Contains(all, p) {
|
||||
t.Fatalf("plaintext %q found at rest", p)
|
||||
@@ -128,7 +162,20 @@ func assertPlainReads(t *testing.T, d *DB, r secretRows) {
|
||||
t.Helper()
|
||||
got := ""
|
||||
if p, err := d.ProfileByID(r.llm); err == nil && p != nil {
|
||||
got += p.APIKey
|
||||
got += p.APIKey + p.Proxy
|
||||
}
|
||||
if ps, err := d.ListProfiles(); err == nil {
|
||||
for _, p := range ps {
|
||||
if p.ID == r.llm && p.Proxy != llmProxyPlain {
|
||||
t.Fatalf("ListProfiles proxy %q not decrypted", p.Proxy)
|
||||
}
|
||||
}
|
||||
}
|
||||
for k, want := range secretSettingValues {
|
||||
if v, ok, err := d.GetSetting(k); err != nil || !ok || v != want {
|
||||
t.Fatalf("GetSetting(%s) = %q %v %v", k, v, ok, err)
|
||||
}
|
||||
got += want
|
||||
}
|
||||
ms, _ := d.ListMCP()
|
||||
for _, m := range ms {
|
||||
@@ -194,4 +241,87 @@ func TestSecretsBackfillLegacyAndKeyChecks(t *testing.T) {
|
||||
t.Fatalf("want wrong-key refusal, got %v", err)
|
||||
}
|
||||
assertPlainReads(t, d, r) // the active key is unchanged
|
||||
|
||||
// The no-key refusal counts the four encrypted settings.
|
||||
t.Setenv(secretKeyEnv, "")
|
||||
withSettings := encryptedCount(t, d, dir)
|
||||
if _, err := d.Exec(deleteSecretSettings); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n := encryptedCount(t, d, dir); withSettings-n != len(secretSettingKeys) {
|
||||
t.Fatalf("no-key refusal counted %d then %d encrypted values; the settings must count", withSettings, n)
|
||||
}
|
||||
|
||||
// A setting sealed under another key fails the startup key check.
|
||||
other, _ := newSecretBox([]byte("ffffffffffffffffffffffffffffffff"))
|
||||
v, _ := other.encrypt("settings.global_proxy", "http://other.example:3128")
|
||||
if _, err := d.Exec(`INSERT INTO settings(key,value) VALUES ('global_proxy',$1)`, v); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Setenv(secretKeyEnv, base64.StdEncoding.EncodeToString(testSecretKey))
|
||||
if err := d.InitSecrets(t.TempDir()); err == nil || !strings.Contains(err.Error(), "settings.global_proxy") {
|
||||
t.Fatalf("want a key-check failure on settings.global_proxy, got %v", err)
|
||||
}
|
||||
if _, err := d.Exec(deleteSecretSettings); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// encryptedCount reads the encrypted-value count from the no-key refusal.
|
||||
func encryptedCount(t *testing.T, d *DB, dir string) int {
|
||||
t.Helper()
|
||||
err := d.InitSecrets(dir)
|
||||
var n int
|
||||
if err == nil {
|
||||
t.Fatal("want a no-key refusal")
|
||||
}
|
||||
if _, serr := fmt.Sscanf(err.Error(), "the database holds %d encrypted", &n); serr != nil {
|
||||
t.Fatalf("refusal %q: %v", err, serr)
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// Secret settings: empty stays empty, InsertSettingIfAbsent encrypts, other
|
||||
// settings stay in clear; an empty LLM proxy stays empty.
|
||||
func TestSecretSettingsEmptyAndClear(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
d.Exec(deleteSecretSettings)
|
||||
d.Exec(`DELETE FROM settings WHERE key='t_sec_clear'`)
|
||||
d.Exec(`DELETE FROM llm_profiles WHERE name='t-sec-noproxy'`)
|
||||
d.Close()
|
||||
})
|
||||
raw := func(k string) string {
|
||||
var v string
|
||||
if err := d.QueryRow(`SELECT value FROM settings WHERE key=$1`, k).Scan(&v); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return v
|
||||
}
|
||||
if err := d.SetSetting("global_proxy", ""); err != nil || raw("global_proxy") != "" {
|
||||
t.Fatalf("empty secret setting must stay empty (%v)", err)
|
||||
}
|
||||
if ok, err := d.InsertSettingIfAbsent("tavily_search_api_key", "tavily-insert"); err != nil || !ok || !isEncrypted(raw("tavily_search_api_key")) {
|
||||
t.Fatalf("InsertSettingIfAbsent must encrypt a secret setting (%v %v)", ok, err)
|
||||
}
|
||||
if v, _, err := d.GetSetting("tavily_search_api_key"); err != nil || v != "tavily-insert" {
|
||||
t.Fatalf("read back %q %v", v, err)
|
||||
}
|
||||
if err := d.SetSetting("t_sec_clear", "visible"); err != nil || raw("t_sec_clear") != "visible" {
|
||||
t.Fatalf("a non-secret setting must stay in clear (%v)", err)
|
||||
}
|
||||
id, err := d.SaveProfile(&LLMProfile{Name: "t-sec-noproxy", Format: "openai", Model: "m", APIKey: "k"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var px *string
|
||||
if err := d.QueryRow(`SELECT proxy FROM llm_profiles WHERE id=$1`, id).Scan(&px); err != nil || px != nil {
|
||||
t.Fatalf("empty proxy stored as %v (%v), want NULL", px, err)
|
||||
}
|
||||
if p, err := d.ProfileByID(id); err != nil || p.Proxy != "" {
|
||||
t.Fatalf("empty proxy read back as %q (%v)", p.Proxy, err)
|
||||
}
|
||||
}
|
||||
|
||||
+23
-2
@@ -5,7 +5,8 @@ import "database/sql"
|
||||
// Settings is a tiny key-value store for global app config the UI toggles at
|
||||
// runtime (e.g. traffic_capture). Missing keys fall back to caller defaults.
|
||||
|
||||
// GetSetting returns the stored value and ok=false when the key is unset.
|
||||
// GetSetting returns the stored value and ok=false when the key is unset. A
|
||||
// secret setting (secretSettingKeys) is decrypted.
|
||||
func (d *DB) GetSetting(key string) (value string, ok bool, err error) {
|
||||
err = d.QueryRow(`SELECT value FROM settings WHERE key=$1`, key).Scan(&value)
|
||||
if err == sql.ErrNoRows {
|
||||
@@ -14,12 +15,29 @@ func (d *DB) GetSetting(key string) (value string, ok bool, err error) {
|
||||
if err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
if c, secret := secretSetting(key); secret {
|
||||
if value, err = openSecret(c, value); err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
}
|
||||
return value, true, nil
|
||||
}
|
||||
|
||||
// sealSetting encrypts the value of a secret setting; other values pass as is.
|
||||
func sealSetting(key, value string) (string, error) {
|
||||
if c, secret := secretSetting(key); secret {
|
||||
return sealSecret(c, value)
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
// SetSetting upserts a setting value.
|
||||
func (d *DB) SetSetting(key, value string) error {
|
||||
_, err := d.Exec(`
|
||||
value, err := sealSetting(key, value)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = d.Exec(`
|
||||
INSERT INTO settings(key, value) VALUES ($1, $2)
|
||||
ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, updated_at = now()`, key, value)
|
||||
return err
|
||||
@@ -32,6 +50,9 @@ ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, updated_at = now()`, key
|
||||
// fast-fail path — on a read error or a concurrent race, the existing value will not be
|
||||
// silently overwritten by ON CONFLICT DO UPDATE.
|
||||
func (d *DB) InsertSettingIfAbsent(key, value string) (inserted bool, err error) {
|
||||
if value, err = sealSetting(key, value); err != nil {
|
||||
return false, err
|
||||
}
|
||||
res, err := d.Exec(`
|
||||
INSERT INTO settings(key, value) VALUES ($1, $2)
|
||||
ON CONFLICT (key) DO NOTHING`, key, value)
|
||||
|
||||
+24
-9
@@ -12,13 +12,14 @@ func TestRegisterTaskAssetScopesCreatesAssetsAndPersistsTextScope(t *testing.T)
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanups below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
task, err := d.CreateTask("manual scope registration", "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "task", d.DeleteTask(task.ID)) })
|
||||
suffix := time.Now().UnixNano()
|
||||
domain := fmt.Sprintf("manual-scope-%d.example.test", suffix)
|
||||
ip := fmt.Sprintf("203.0.%d.%d", (suffix/250)%250, suffix%250+1)
|
||||
@@ -48,7 +49,10 @@ func TestRegisterTaskAssetScopesCreatesAssetsAndPersistsTextScope(t *testing.T)
|
||||
t.Fatalf("unexpected task asset provenance: %+v", asset)
|
||||
}
|
||||
}
|
||||
t.Cleanup(func() { _, _ = d.Assets().DeleteByIDs(assetIDs) })
|
||||
t.Cleanup(func() {
|
||||
_, err := d.Assets().DeleteByIDs(assetIDs)
|
||||
cleanupErr(t, "assets", err)
|
||||
})
|
||||
scopes, err := d.Assets().ListTaskScope(task.ID)
|
||||
if err != nil || len(scopes) != 5 {
|
||||
t.Fatalf("task scope=%+v err=%v", scopes, err)
|
||||
@@ -91,18 +95,22 @@ func TestTaskAssetAttachDetachPreservesGlobalAssetAndAnchors(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanups below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
task, err := d.CreateTask("asset editing", "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "task", d.DeleteTask(task.ID)) })
|
||||
assetID, err := d.Assets().UpsertRootDomain(UpsertRootDomainReq{Domain: fmt.Sprintf("manual-%d.example.test", time.Now().UnixNano())})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _, _ = d.Assets().DeleteByIDs([]int64{assetID}) })
|
||||
t.Cleanup(func() {
|
||||
_, err := d.Assets().DeleteByIDs([]int64{assetID})
|
||||
cleanupErr(t, "asset", err)
|
||||
})
|
||||
|
||||
mutation, err := d.Assets().AttachAssetsToTask(task.ID, []int64{assetID, assetID}, "authorized asset list item 3")
|
||||
if err != nil {
|
||||
@@ -144,7 +152,8 @@ func TestIntentAssetsIncludesDirectSourceProvenance(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanups below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
source, err := d.CreateTask("source assets", "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
@@ -154,7 +163,10 @@ func TestIntentAssetsIncludesDirectSourceProvenance(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(current.ID); _ = d.DeleteTask(source.ID) })
|
||||
t.Cleanup(func() {
|
||||
cleanupErr(t, "current task", d.DeleteTask(current.ID))
|
||||
cleanupErr(t, "source task", d.DeleteTask(source.ID))
|
||||
})
|
||||
|
||||
assetID, err := d.Assets().UpsertHTTPService(UpsertHTTPServiceReq{
|
||||
URL: fmt.Sprintf("https://intent-%d.example.test", time.Now().UnixNano()), TaskID: source.ID,
|
||||
@@ -162,7 +174,10 @@ func TestIntentAssetsIncludesDirectSourceProvenance(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _, _ = d.Assets().DeleteByIDs([]int64{assetID}) })
|
||||
t.Cleanup(func() {
|
||||
_, err := d.Assets().DeleteByIDs([]int64{assetID})
|
||||
cleanupErr(t, "asset", err)
|
||||
})
|
||||
intentID, err := d.Exploration(source.ExplorationID).AddIntent(map[string]any{"summary": "source worker"}, 5, []int64{assetID}, "planner")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
|
||||
@@ -13,7 +13,8 @@ func TestTaskCategoryCRUDAndTaskAssignment(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
name := fmt.Sprintf("Category %d", time.Now().UnixNano())
|
||||
category, err := d.CreateTaskCategory(" " + strings.ReplaceAll(name, " ", " ") + " ")
|
||||
@@ -31,7 +32,7 @@ func TestTaskCategoryCRUDAndTaskAssignment(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "task", d.DeleteTask(task.ID)) })
|
||||
loaded, err := d.GetTask(task.ID)
|
||||
if err != nil || loaded == nil || loaded.CategoryID == nil || *loaded.CategoryID != category.ID || loaded.CategoryName != name {
|
||||
t.Fatalf("categorized task=%+v err=%v", loaded, err)
|
||||
@@ -95,14 +96,18 @@ func TestSetTasksCategoryBatch(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanups below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
stamp := time.Now().UnixNano()
|
||||
category, err := d.CreateTaskCategory(fmt.Sprintf("Batch %d", stamp))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _, _ = d.DeleteTaskCategory(category.ID) })
|
||||
t.Cleanup(func() {
|
||||
_, err := d.DeleteTaskCategory(category.ID)
|
||||
cleanupErr(t, "category", err)
|
||||
})
|
||||
|
||||
ids := make([]int64, 0, 3)
|
||||
for i := range 3 {
|
||||
@@ -110,7 +115,7 @@ func TestSetTasksCategoryBatch(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "task", d.DeleteTask(task.ID)) })
|
||||
ids = append(ids, task.ID)
|
||||
}
|
||||
|
||||
@@ -152,14 +157,15 @@ func TestSetTasksCategoryRejectsBadInputAtomically(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
stamp := time.Now().UnixNano()
|
||||
task, err := d.CreateTaskWithOptions(fmt.Sprintf("batch-reject-%d", stamp), "goal", TaskCreateOptions{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "task", d.DeleteTask(task.ID)) })
|
||||
|
||||
missingCategory := int64(1 << 62)
|
||||
if _, _, err := d.SetTasksCategory([]int64{task.ID}, &missingCategory); !errors.Is(err, ErrTaskCategoryNotFound) {
|
||||
|
||||
@@ -13,7 +13,8 @@ func TestTaskLLMProfileMutationsLockTaskBeforeProfile(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
suffix := time.Now().UnixNano()
|
||||
first, err := d.SaveProfile(&LLMProfile{
|
||||
@@ -38,9 +39,9 @@ func TestTaskLLMProfileMutationsLockTaskBeforeProfile(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = d.DeleteTask(task.ID)
|
||||
_ = d.DeleteProfile(first)
|
||||
_ = d.DeleteProfile(second)
|
||||
cleanupErr(t, "task", d.DeleteTask(task.ID))
|
||||
cleanupProfile(t, d, first)
|
||||
cleanupProfile(t, d, second)
|
||||
})
|
||||
|
||||
replaceDB, replacePID := openSingleConnectionTestDB(t, dsn)
|
||||
@@ -112,7 +113,8 @@ func TestDeleteProfileLocksNonTaskReferencesBeforeProfile(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
suffix := time.Now().UnixNano()
|
||||
profileID, err := d.SaveProfile(&LLMProfile{
|
||||
@@ -132,8 +134,8 @@ func TestDeleteProfileLocksNonTaskReferencesBeforeProfile(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = d.DeleteAgent(agent.Key)
|
||||
_ = d.DeleteProfile(profileID)
|
||||
cleanupErr(t, "agent", d.DeleteAgent(agent.Key))
|
||||
cleanupProfile(t, d, profileID)
|
||||
})
|
||||
|
||||
blocker, err := d.Begin()
|
||||
@@ -169,7 +171,8 @@ func TestDeleteProfileLocksNonTaskReferencesBeforeProfile(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
suffix := time.Now().UnixNano()
|
||||
profileID, err := d.SaveProfile(&LLMProfile{
|
||||
@@ -184,8 +187,8 @@ func TestDeleteProfileLocksNonTaskReferencesBeforeProfile(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = d.DeleteConversation(conversation.ID)
|
||||
_ = d.DeleteProfile(profileID)
|
||||
cleanupErr(t, "conversation", d.DeleteConversation(conversation.ID))
|
||||
cleanupProfile(t, d, profileID)
|
||||
})
|
||||
|
||||
blocker, err := d.Begin()
|
||||
@@ -225,7 +228,8 @@ func TestNonTaskProfileMutationsLockReferenceBeforeProfile(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
suffix := time.Now().UnixNano()
|
||||
profileID, err := d.SaveProfile(&LLMProfile{
|
||||
@@ -240,8 +244,8 @@ func TestNonTaskProfileMutationsLockReferenceBeforeProfile(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = d.DeleteAgent(agent.Key)
|
||||
_ = d.DeleteProfile(profileID)
|
||||
cleanupErr(t, "agent", d.DeleteAgent(agent.Key))
|
||||
cleanupProfile(t, d, profileID)
|
||||
})
|
||||
|
||||
blocker, err := d.Begin()
|
||||
@@ -273,7 +277,8 @@ func TestNonTaskProfileMutationsLockReferenceBeforeProfile(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
suffix := time.Now().UnixNano()
|
||||
profileID, err := d.SaveProfile(&LLMProfile{
|
||||
@@ -288,8 +293,8 @@ func TestNonTaskProfileMutationsLockReferenceBeforeProfile(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = d.DeleteConversation(conversation.ID)
|
||||
_ = d.DeleteProfile(profileID)
|
||||
cleanupErr(t, "conversation", d.DeleteConversation(conversation.ID))
|
||||
cleanupProfile(t, d, profileID)
|
||||
})
|
||||
|
||||
blocker, err := d.Begin()
|
||||
@@ -323,7 +328,8 @@ func TestCreateConversationAndDeleteProfileDoNotDeadlock(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
suffix := time.Now().UnixNano()
|
||||
profileID, err := d.SaveProfile(&LLMProfile{
|
||||
@@ -334,8 +340,8 @@ func TestCreateConversationAndDeleteProfileDoNotDeadlock(t *testing.T) {
|
||||
}
|
||||
title := fmt.Sprintf("conversation create race %d", suffix)
|
||||
t.Cleanup(func() {
|
||||
_, _ = d.Exec(`DELETE FROM conversations WHERE title=$1`, title)
|
||||
_ = d.DeleteProfile(profileID)
|
||||
cleanupExec(t, d, `DELETE FROM conversations WHERE title=$1`, title)
|
||||
cleanupProfile(t, d, profileID)
|
||||
})
|
||||
|
||||
profileBlocker, err := d.Begin()
|
||||
@@ -403,7 +409,8 @@ func TestDeleteProfileRetriesReferenceCommittedAfterInitialScan(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
suffix := time.Now().UnixNano()
|
||||
profileID, err := d.SaveProfile(&LLMProfile{
|
||||
@@ -418,8 +425,8 @@ func TestDeleteProfileRetriesReferenceCommittedAfterInitialScan(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = d.DeleteTask(task.ID)
|
||||
_ = d.DeleteProfile(profileID)
|
||||
cleanupErr(t, "task", d.DeleteTask(task.ID))
|
||||
cleanupProfile(t, d, profileID)
|
||||
})
|
||||
|
||||
// Keep the new reference uncommitted while deletion takes its initial
|
||||
|
||||
@@ -11,7 +11,8 @@ func TestListTasksBulkHydratesTaskContext(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanups below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
stamp := time.Now().UnixNano()
|
||||
profileID, err := d.SaveProfile(&LLMProfile{
|
||||
@@ -20,17 +21,17 @@ func TestListTasksBulkHydratesTaskContext(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteProfile(profileID) })
|
||||
t.Cleanup(func() { cleanupProfile(t, d, profileID) })
|
||||
source, err := d.CreateTask(fmt.Sprintf("bulk-list-source-%d", stamp), "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(source.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "source task", d.DeleteTask(source.ID)) })
|
||||
companyID, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Bulk List Company %d", stamp), "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.Companies().DeleteCompany(companyID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "company", d.Companies().DeleteCompany(companyID)) })
|
||||
child, err := d.CreateTaskWithOptions("bulk-list-child", "goal", TaskCreateOptions{
|
||||
LLMProfileIDs: []int64{profileID},
|
||||
SourceTaskIDs: []int64{source.ID},
|
||||
@@ -39,7 +40,7 @@ func TestListTasksBulkHydratesTaskContext(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(child.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "child task", d.DeleteTask(child.ID)) })
|
||||
|
||||
tasks, err := d.ListTasks()
|
||||
if err != nil {
|
||||
@@ -72,13 +73,14 @@ func TestTaskListMetricsAll(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
task, err := d.CreateTask("task-list-metrics", "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "task", d.DeleteTask(task.ID)) })
|
||||
|
||||
resultAt := time.Now().Add(-time.Second).Truncate(time.Second)
|
||||
latestAt := resultAt.Add(time.Second)
|
||||
|
||||
@@ -14,7 +14,8 @@ func TestTaskTemplateCRUDAndNormalizedUniqueness(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
suffix := time.Now().UnixNano()
|
||||
name := fmt.Sprintf("Template %d", suffix)
|
||||
@@ -26,7 +27,10 @@ func TestTaskTemplateCRUDAndNormalizedUniqueness(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _, _ = d.DeleteTaskTemplate(created.ID) })
|
||||
t.Cleanup(func() {
|
||||
_, err := d.DeleteTaskTemplate(created.ID)
|
||||
cleanupErr(t, "template", err)
|
||||
})
|
||||
if created.Name != name || created.Description != "initial description" || created.Goal != "initial goal" {
|
||||
t.Fatalf("template was not normalized: %+v", created)
|
||||
}
|
||||
@@ -85,7 +89,8 @@ func TestTaskTemplateDisjointPatchesCompose(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanup below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
created, err := d.CreateTaskTemplate(TaskTemplateInput{
|
||||
Name: fmt.Sprintf("Concurrent template %d", time.Now().UnixNano()),
|
||||
@@ -95,7 +100,10 @@ func TestTaskTemplateDisjointPatchesCompose(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _, _ = d.DeleteTaskTemplate(created.ID) })
|
||||
t.Cleanup(func() {
|
||||
_, err := d.DeleteTaskTemplate(created.ID)
|
||||
cleanupErr(t, "template", err)
|
||||
})
|
||||
|
||||
description := "description from concurrent patch"
|
||||
goal := "goal from concurrent patch"
|
||||
|
||||
+30
-21
@@ -13,7 +13,7 @@ func TestTaskLifecycleAndDeleteCascade(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
tk, err := d.CreateTask("migration test", "goal X", nil, 0, 0)
|
||||
if err != nil {
|
||||
@@ -87,7 +87,8 @@ func TestTaskDeleteCascadeAssets(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanups below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
first, err := d.CreateTask("cascade delete test", "goal A", nil, 0, 0)
|
||||
if err != nil {
|
||||
@@ -138,9 +139,10 @@ func TestTaskDeleteCascadeAssets(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = d.DeleteTask(first.ID)
|
||||
_ = d.DeleteTask(second.ID)
|
||||
_, _ = assets.DeleteByIDs([]int64{exclusiveID, sharedID, anchorOnlyID, protectedID})
|
||||
cleanupErr(t, "first task", d.DeleteTask(first.ID))
|
||||
cleanupErr(t, "second task", d.DeleteTask(second.ID))
|
||||
_, err := assets.DeleteByIDs([]int64{exclusiveID, sharedID, anchorOnlyID, protectedID})
|
||||
cleanupErr(t, "assets", err)
|
||||
})
|
||||
|
||||
hosts, err := assets.HostsByTask(first.ID)
|
||||
@@ -211,7 +213,8 @@ func TestTaskRelationsAndLLMFailoverChain(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanups below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
suffix := time.Now().UnixNano()
|
||||
profileIDs := make([]int64, 0, 3)
|
||||
@@ -242,11 +245,11 @@ func TestTaskRelationsAndLLMFailoverChain(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = d.DeleteTask(child.ID)
|
||||
_ = d.DeleteTask(sourceA.ID)
|
||||
_ = d.DeleteTask(sourceB.ID)
|
||||
cleanupErr(t, "child task", d.DeleteTask(child.ID))
|
||||
cleanupErr(t, "source task A", d.DeleteTask(sourceA.ID))
|
||||
cleanupErr(t, "source task B", d.DeleteTask(sourceB.ID))
|
||||
for _, id := range profileIDs {
|
||||
_ = d.DeleteProfile(id)
|
||||
cleanupProfile(t, d, id)
|
||||
}
|
||||
})
|
||||
|
||||
@@ -469,13 +472,14 @@ func TestTaskContextRejectsDuplicatesAndAllowsTerminalLLMEdits(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanups below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
source, err := d.CreateTask("source", "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(source.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "source task", d.DeleteTask(source.ID)) })
|
||||
if _, err := d.CreateTaskWithOptions("bad", "goal", TaskCreateOptions{SourceTaskIDs: []int64{source.ID, source.ID}}); err == nil {
|
||||
t.Fatal("duplicate source task ids should be rejected")
|
||||
}
|
||||
@@ -490,13 +494,13 @@ func TestTaskContextRejectsDuplicatesAndAllowsTerminalLLMEdits(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteProfile(profileID) })
|
||||
t.Cleanup(func() { cleanupProfile(t, d, profileID) })
|
||||
|
||||
task, err := d.CreateTask("terminal", "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "terminal task", d.DeleteTask(task.ID)) })
|
||||
if err := d.SetStatus(task.ID, "done"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -523,7 +527,8 @@ func TestCreateTaskWithCompanyScopes(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanups below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
suffix := time.Now().UnixNano()
|
||||
companyA, _, err := d.Companies().UpsertCompany(fmt.Sprintf("Task Company A %d", suffix), "")
|
||||
@@ -539,14 +544,14 @@ func TestCreateTaskWithCompanyScopes(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = d.Exec(`DELETE FROM companies WHERE id IN ($1,$2,$3)`, companyA, companyB, emptyCompany)
|
||||
cleanupExec(t, d, `DELETE FROM companies WHERE id IN ($1,$2,$3)`, companyA, companyB, emptyCompany)
|
||||
})
|
||||
|
||||
existingTask, err := d.CreateTask("existing company asset owner", "goal", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(existingTask.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "existing task", d.DeleteTask(existingTask.ID)) })
|
||||
var companyAssetA, companyAssetB int64
|
||||
domainA := fmt.Sprintf("task-company-a-%d.example.test", suffix)
|
||||
domainB := fmt.Sprintf("task-company-b-%d.example.test", suffix)
|
||||
@@ -562,7 +567,10 @@ VALUES ('root_domain',$1,$1,$2,'explicit')
|
||||
RETURNING id`, domainB, companyB).Scan(&companyAssetB); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _, _ = d.Assets().DeleteByIDs([]int64{companyAssetA, companyAssetB}) })
|
||||
t.Cleanup(func() {
|
||||
_, err := d.Assets().DeleteByIDs([]int64{companyAssetA, companyAssetB})
|
||||
cleanupErr(t, "company assets", err)
|
||||
})
|
||||
|
||||
task, err := d.CreateTaskWithOptions("company-scoped task", "use company scope", TaskCreateOptions{
|
||||
CompanyIDs: []int64{companyA, companyB},
|
||||
@@ -570,7 +578,7 @@ RETURNING id`, domainB, companyB).Scan(&companyAssetB); err != nil {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(task.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "company-scoped task", d.DeleteTask(task.ID)) })
|
||||
if fmt.Sprint(task.CompanyIDs) != fmt.Sprint([]int64{companyA, companyB}) {
|
||||
t.Fatalf("creation result company IDs=%v", task.CompanyIDs)
|
||||
}
|
||||
@@ -674,7 +682,8 @@ func TestListTasksOrderByIDDesc(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer d.Close()
|
||||
// Registered FIRST so it runs LAST, after the row cleanups below.
|
||||
t.Cleanup(func() { d.Close() })
|
||||
|
||||
stamp := time.Now().UnixNano()
|
||||
var ids []int64
|
||||
@@ -683,7 +692,7 @@ func TestListTasksOrderByIDDesc(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = d.DeleteTask(tk.ID) })
|
||||
t.Cleanup(func() { cleanupErr(t, "order task", d.DeleteTask(tk.ID)) })
|
||||
ids = append(ids, tk.ID)
|
||||
}
|
||||
|
||||
|
||||
+23
-8
@@ -110,6 +110,13 @@ func scanTool(rows interface{ Scan(...any) error }) (*Tool, error) {
|
||||
if err := rows.Scan(&t.Key, &t.System, &t.Description, &t.Schema, &agents, &t.Enabled, &t.Kind, &t.Exec, &t.Deferred); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !t.System { // a custom tool's exec is sealed at rest
|
||||
plain, err := openSecret(colToolExec, string(t.Exec))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
t.Exec = json.RawMessage(plain)
|
||||
}
|
||||
if len(agents) > 0 {
|
||||
_ = json.Unmarshal(agents, &t.Agents)
|
||||
}
|
||||
@@ -154,15 +161,15 @@ func (d *DB) CreateCustomTool(t *Tool) error {
|
||||
if len(schema) == 0 {
|
||||
schema = json.RawMessage("{}")
|
||||
}
|
||||
exec := t.Exec
|
||||
if len(exec) == 0 {
|
||||
exec = json.RawMessage("{}")
|
||||
exec, err := sealToolExec(t.Exec)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
agents, _ := json.Marshal(t.Agents)
|
||||
if len(agents) == 0 {
|
||||
agents = json.RawMessage("[]")
|
||||
}
|
||||
_, err := d.Exec(`
|
||||
_, err = d.Exec(`
|
||||
INSERT INTO tools(key, system, description, schema, agents, enabled, kind, exec, deferred)
|
||||
VALUES ($1, false, $2, $3, $4, $5, $6, $7, $8)`,
|
||||
t.Key, t.Description, schema, agents, t.Enabled, t.Kind, exec, t.Deferred)
|
||||
@@ -176,21 +183,29 @@ func (d *DB) UpdateCustomTool(t *Tool) error {
|
||||
if len(schema) == 0 {
|
||||
schema = json.RawMessage("{}")
|
||||
}
|
||||
exec := t.Exec
|
||||
if len(exec) == 0 {
|
||||
exec = json.RawMessage("{}")
|
||||
exec, err := sealToolExec(t.Exec)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
agents, _ := json.Marshal(t.Agents)
|
||||
if len(agents) == 0 {
|
||||
agents = json.RawMessage("[]")
|
||||
}
|
||||
_, err := d.Exec(`
|
||||
_, err = d.Exec(`
|
||||
UPDATE tools SET description=$2, schema=$3, agents=$4, enabled=$5, kind=$6, exec=$7, deferred=$8
|
||||
WHERE key=$1 AND system=false`,
|
||||
t.Key, t.Description, schema, agents, t.Enabled, t.Kind, exec, t.Deferred)
|
||||
return err
|
||||
}
|
||||
|
||||
// sealToolExec encrypts a custom tool's exec spec for storage; empty means {}.
|
||||
func sealToolExec(exec json.RawMessage) (string, error) {
|
||||
if len(exec) == 0 {
|
||||
exec = json.RawMessage("{}")
|
||||
}
|
||||
return sealSecret(colToolExec, string(exec))
|
||||
}
|
||||
|
||||
// DeleteCustomTool removes a custom tool (system=false only; built-ins protected).
|
||||
func (d *DB) DeleteCustomTool(key string) error {
|
||||
_, err := d.Exec(`DELETE FROM tools WHERE key=$1 AND system=false`, key)
|
||||
|
||||
@@ -0,0 +1,184 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// toolExecPlain is a custom http tool's exec carrying a secret in every field
|
||||
// the phase 3 encryption covers.
|
||||
const toolExecPlain = `{"method":"GET","url":"https://api.example/x","headers":{"Authorization":"Bearer tool-hdr-plain"},` +
|
||||
`"body":"tool-body-plain","proxy":"http://u:tool-proxy-plain@p.example:3128"}`
|
||||
|
||||
var toolPlainSecrets = []string{"tool-hdr-plain", "tool-body-plain", "tool-proxy-plain", "tool-cmd-plain"}
|
||||
|
||||
func rawToolExec(t *testing.T, d *DB, key string) string {
|
||||
t.Helper()
|
||||
var v string
|
||||
if err := d.QueryRow(`SELECT exec::text FROM tools WHERE key=$1`, key).Scan(&v); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
// assertToolExecSealed checks a custom tool's exec is sealed at rest and read
|
||||
// back in clear, equal to want, by every db getter.
|
||||
func assertToolExecSealed(t *testing.T, d *DB, key, want string) {
|
||||
t.Helper()
|
||||
raw := rawToolExec(t, d, key)
|
||||
if !isSealed(colToolExec, raw) {
|
||||
t.Fatalf("tools.exec of %s not sealed at rest: %s", key, raw)
|
||||
}
|
||||
for _, p := range toolPlainSecrets {
|
||||
if strings.Contains(raw, p) {
|
||||
t.Fatalf("plaintext %q found at rest", p)
|
||||
}
|
||||
}
|
||||
var got []json.RawMessage
|
||||
if tl, err := d.GetTool(key); err != nil || tl == nil {
|
||||
t.Fatalf("GetTool: %v", err)
|
||||
} else {
|
||||
got = append(got, tl.Exec)
|
||||
}
|
||||
for _, list := range []func() ([]*Tool, error){d.ListTools, d.ListCustomTools} {
|
||||
ts, err := list()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, tl := range ts {
|
||||
if tl.Key == key {
|
||||
got = append(got, tl.Exec)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(got) != 3 {
|
||||
t.Fatalf("tool %s read by %d getters, want 3", key, len(got))
|
||||
}
|
||||
for _, g := range got {
|
||||
if !jsonEqual(t, string(g), want) {
|
||||
t.Fatalf("exec read back as %s, want %s", g, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func jsonEqual(t *testing.T, a, b string) bool {
|
||||
t.Helper()
|
||||
var x, y any
|
||||
if json.Unmarshal([]byte(a), &x) != nil || json.Unmarshal([]byte(b), &y) != nil {
|
||||
return false
|
||||
}
|
||||
ca, _ := json.Marshal(x)
|
||||
cb, _ := json.Marshal(y)
|
||||
return string(ca) == string(cb)
|
||||
}
|
||||
|
||||
func TestToolExecSealedAtRest(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
d.Exec(`DELETE FROM tools WHERE key LIKE 't_secexec%'`)
|
||||
d.Close()
|
||||
})
|
||||
if err := d.CreateCustomTool(&Tool{Key: "t_secexec", Kind: "http", Exec: json.RawMessage(toolExecPlain), Enabled: true}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertToolExecSealed(t, d, "t_secexec", toolExecPlain)
|
||||
first := rawToolExec(t, d, "t_secexec")
|
||||
|
||||
cmd := `{"command":"curl -H 'X-Key: tool-cmd-plain' {args}","timeout_ms":5000}`
|
||||
if err := d.UpdateCustomTool(&Tool{Key: "t_secexec", Kind: "command", Exec: json.RawMessage(cmd), Enabled: true}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertToolExecSealed(t, d, "t_secexec", cmd)
|
||||
if err := d.UpdateCustomTool(&Tool{Key: "t_secexec", Kind: "http", Exec: json.RawMessage(toolExecPlain), Enabled: true}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if rawToolExec(t, d, "t_secexec") == first {
|
||||
t.Fatal("the same exec must seal under a fresh nonce")
|
||||
}
|
||||
|
||||
// A shell tool has no exec: {} stays {}.
|
||||
if err := d.CreateCustomTool(&Tool{Key: "t_secexec_shell", Kind: "shell", Enabled: true}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if raw := rawToolExec(t, d, "t_secexec_shell"); raw != "{}" {
|
||||
t.Fatalf("empty exec stored as %s, want {}", raw)
|
||||
}
|
||||
if err := d.CreateCustomTool(&Tool{Key: "t_secexec_resv", Kind: "command", Exec: json.RawMessage(`{"__artex_enc":"x"}`)}); err == nil {
|
||||
t.Fatal("reserved key must be refused on write")
|
||||
}
|
||||
|
||||
// A value sealed under another key fails the read (fail closed).
|
||||
other, _ := newSecretBox([]byte("ffffffffffffffffffffffffffffffff"))
|
||||
foreign, _ := other.seal(colToolExec, toolExecPlain)
|
||||
if _, err := d.Exec(`UPDATE tools SET exec=$1::jsonb WHERE key='t_secexec'`, foreign); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := d.GetTool("t_secexec"); err == nil || !strings.Contains(err.Error(), "tools.exec") {
|
||||
t.Fatalf("a foreign ciphertext must fail the read, got %v", err)
|
||||
}
|
||||
if _, err := d.ListCustomTools(); err == nil {
|
||||
t.Fatal("ListCustomTools must fail on a foreign ciphertext")
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolExecBackfillAndKeyChecks(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
d.Exec(`DELETE FROM tools WHERE key LIKE 't_secexec%'`)
|
||||
d.Close()
|
||||
})
|
||||
sysExec := `{"command":"echo system-plain"}`
|
||||
for _, q := range []struct {
|
||||
key, exec string
|
||||
system bool
|
||||
}{{"t_secexec_legacy", toolExecPlain, false}, {"t_secexec_sys", sysExec, true}, {"t_secexec_empty", `{}`, false}} {
|
||||
if _, err := d.Exec(`INSERT INTO tools(key, system, kind, exec) VALUES ($1, $2, 'http', $3::jsonb)`, q.key, q.system, q.exec); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
t.Setenv(secretKeyEnv, base64.StdEncoding.EncodeToString(testSecretKey))
|
||||
if err := d.InitSecrets(t.TempDir()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertToolExecSealed(t, d, "t_secexec_legacy", toolExecPlain)
|
||||
first := rawToolExec(t, d, "t_secexec_legacy")
|
||||
if raw := rawToolExec(t, d, "t_secexec_sys"); !jsonEqual(t, raw, sysExec) {
|
||||
t.Fatalf("a system tool's exec must stay in clear, got %s", raw)
|
||||
}
|
||||
if raw := rawToolExec(t, d, "t_secexec_empty"); raw != "{}" {
|
||||
t.Fatalf("an empty exec must stay {}, got %s", raw)
|
||||
}
|
||||
if err := d.InitSecrets(t.TempDir()); err != nil || rawToolExec(t, d, "t_secexec_legacy") != first {
|
||||
t.Fatalf("second backfill must leave the ciphertext untouched (%v)", err)
|
||||
}
|
||||
|
||||
// Missing key: the sealed exec counts in the refusal.
|
||||
t.Setenv(secretKeyEnv, "")
|
||||
dir := t.TempDir()
|
||||
withTool := encryptedCount(t, d, dir)
|
||||
if _, err := d.Exec(`UPDATE tools SET exec='{}'::jsonb WHERE key='t_secexec_legacy'`); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n := encryptedCount(t, d, dir); withTool-n != 1 {
|
||||
t.Fatalf("no-key refusal counted %d then %d encrypted values; tools.exec must count", withTool, n)
|
||||
}
|
||||
|
||||
// Wrong key: an exec sealed under another key fails the startup key check.
|
||||
other, _ := newSecretBox([]byte("ffffffffffffffffffffffffffffffff"))
|
||||
foreign, _ := other.seal(colToolExec, toolExecPlain)
|
||||
if _, err := d.Exec(`UPDATE tools SET exec=$1::jsonb WHERE key='t_secexec_legacy'`, foreign); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Setenv(secretKeyEnv, base64.StdEncoding.EncodeToString(testSecretKey))
|
||||
if err := d.InitSecrets(t.TempDir()); err == nil || !strings.Contains(err.Error(), "tools.exec") {
|
||||
t.Fatalf("want a key-check failure on tools.exec, got %v", err)
|
||||
}
|
||||
}
|
||||
@@ -34,5 +34,10 @@ func runEvidenceSuite(m *testing.M) int {
|
||||
panic(err)
|
||||
}
|
||||
defer conn.ExecContext(context.Background(), `SELECT pg_advisory_unlock(7337741002)`)
|
||||
// Snapshot inserts seal the heads: load a fixed test key (env, no key file).
|
||||
os.Setenv("ARTEX_SECRET_KEY", "MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY=")
|
||||
if err = pg.InitSecrets(""); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return m.Run()
|
||||
}
|
||||
|
||||
@@ -21,6 +21,7 @@ import (
|
||||
"github.com/Autumn-27/artex/egress"
|
||||
"github.com/Autumn-27/artex/guard"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
"github.com/Autumn-27/artex/notify"
|
||||
"github.com/Autumn-27/norma/permission"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
@@ -72,6 +73,11 @@ func (s *Server) pgCreateCustomTool(w http.ResponseWriter, r *http.Request) {
|
||||
writeErr(w, 409, "this key already exists (built-in or custom tool)")
|
||||
return
|
||||
}
|
||||
var err error
|
||||
if req.Exec, err = unmaskToolExec("", nil, req.Kind, req.Exec); err != nil { // nothing stored: a placeholder is refused
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if err := pg.CreateCustomTool(&db.Tool{
|
||||
Key: req.Key, Description: req.Description, Schema: req.Schema, Agents: req.Agents,
|
||||
Enabled: req.Enabled, Kind: req.Kind, Exec: req.Exec, Deferred: req.Deferred,
|
||||
@@ -110,6 +116,10 @@ func (s *Server) pgUpdateCustomTool(w http.ResponseWriter, r *http.Request) {
|
||||
writeErr(w, 400, "http tools must provide a parameter JSON Schema (cannot be left empty)")
|
||||
return
|
||||
}
|
||||
if req.Exec, err = unmaskToolExec(existing.Kind, existing.Exec, req.Kind, req.Exec); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if err := pg.UpdateCustomTool(&db.Tool{
|
||||
Key: key, Description: req.Description, Schema: req.Schema, Agents: req.Agents,
|
||||
Enabled: req.Enabled, Kind: req.Kind, Exec: req.Exec, Deferred: req.Deferred,
|
||||
@@ -138,6 +148,7 @@ func (s *Server) pgDeleteCustomTool(w http.ResponseWriter, r *http.Request) {
|
||||
// real tool call — it runs arbitrary command/script/http on the server, which the
|
||||
// custom-tool feature already allows, so no new capability is granted.
|
||||
type testToolReq struct {
|
||||
Key string `json:"key"` // the saved tool being edited, if any: its stored values replace placeholders
|
||||
Kind string `json:"kind"` // command | script | http
|
||||
Exec json.RawMessage `json:"exec"`
|
||||
Params map[string]any `json:"params"`
|
||||
@@ -156,6 +167,23 @@ func (s *Server) pgTestCustomTool(w http.ResponseWriter, r *http.Request) {
|
||||
writeErr(w, 400, "invalid request body")
|
||||
return
|
||||
}
|
||||
var storedKind string
|
||||
var stored json.RawMessage
|
||||
if req.Key != "" {
|
||||
t, err := pg.GetTool(req.Key)
|
||||
if err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if t != nil && !t.System {
|
||||
storedKind, stored = t.Kind, t.Exec
|
||||
}
|
||||
}
|
||||
var err error
|
||||
if req.Exec, err = unmaskToolExec(storedKind, stored, req.Kind, req.Exec); err != nil {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
params := req.Params
|
||||
if params == nil {
|
||||
params = map[string]any{}
|
||||
@@ -236,6 +264,86 @@ type httpExec struct {
|
||||
UseRecordingProxy bool `json:"use_recording_proxy"`
|
||||
}
|
||||
|
||||
// toolExecFlags are the response-only *_set flags of maskToolExec; never stored.
|
||||
var toolExecFlags = []string{"headers_set", "proxy_set", "body_set"}
|
||||
|
||||
// maskToolExec is what the API shows of a custom tool's exec: every header value,
|
||||
// the proxy and the body become the placeholder, each with a *_set flag.
|
||||
func maskToolExec(raw json.RawMessage) json.RawMessage {
|
||||
var e map[string]any
|
||||
if json.Unmarshal(raw, &e) != nil || e == nil {
|
||||
return raw
|
||||
}
|
||||
if h, ok := e["headers"].(map[string]any); ok {
|
||||
set := false
|
||||
for k, v := range h {
|
||||
if v != "" {
|
||||
h[k], set = notify.MaskedPrefix, true
|
||||
}
|
||||
}
|
||||
e["headers_set"] = set
|
||||
}
|
||||
for _, k := range []string{"proxy", "body"} {
|
||||
if s, ok := e[k].(string); ok {
|
||||
e[k], e[k+"_set"] = maskedSecret(s), s != ""
|
||||
}
|
||||
}
|
||||
out, _ := json.Marshal(e)
|
||||
return out
|
||||
}
|
||||
|
||||
// unmaskToolExec puts the stored value back for every header, proxy or body sent
|
||||
// as the placeholder, so a client echoing the list response keeps the stored
|
||||
// secret; an explicit value replaces it and "" clears it. A placeholder with
|
||||
// nothing stored to keep (stored nil on create) is refused. As for the MCP env,
|
||||
// a stored secret is only kept for the destination it was configured for: when
|
||||
// the url (exact string) or the kind changes, a placeholder is refused.
|
||||
func unmaskToolExec(storedKind string, stored json.RawMessage, kind string, in json.RawMessage) (json.RawMessage, error) {
|
||||
var e, old map[string]any
|
||||
if json.Unmarshal(in, &e) != nil || e == nil {
|
||||
return in, nil // not an object: stored or refused as before
|
||||
}
|
||||
_ = json.Unmarshal(stored, &old)
|
||||
changed, kept := false, false
|
||||
for _, k := range toolExecFlags {
|
||||
if _, ok := e[k]; ok {
|
||||
delete(e, k)
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
oldHeaders, _ := old["headers"].(map[string]any)
|
||||
headers, _ := e["headers"].(map[string]any)
|
||||
for k, v := range headers {
|
||||
if s, ok := v.(string); !ok || !isMaskedSecret(s) {
|
||||
continue
|
||||
}
|
||||
prev, ok := oldHeaders[k]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("header %q: masked value but nothing stored to keep", k)
|
||||
}
|
||||
headers[k], changed, kept = prev, true, true
|
||||
}
|
||||
for _, k := range []string{"proxy", "body"} {
|
||||
if s, ok := e[k].(string); !ok || !isMaskedSecret(s) {
|
||||
continue
|
||||
}
|
||||
prev, _ := old[k].(string)
|
||||
if prev == "" {
|
||||
return nil, fmt.Errorf("%s: masked value but nothing stored to keep", k)
|
||||
}
|
||||
e[k], changed, kept = prev, true, true
|
||||
}
|
||||
oldURL, _ := old["url"].(string)
|
||||
newURL, _ := e["url"].(string)
|
||||
if kept && (oldURL != newURL || storedKind != kind) {
|
||||
return nil, fmt.Errorf("destination changed (url or kind) while a stored secret was kept: re-enter the masked header/proxy/body or clear it; stored secrets are not sent to a new destination")
|
||||
}
|
||||
if !changed {
|
||||
return in, nil
|
||||
}
|
||||
return json.Marshal(e)
|
||||
}
|
||||
|
||||
func timeoutOr(ms, def int) time.Duration {
|
||||
if ms <= 0 {
|
||||
return time.Duration(def) * time.Millisecond
|
||||
|
||||
+4
-2
@@ -603,7 +603,8 @@ type LLMProfileDTO struct {
|
||||
Name string `json:"name"`
|
||||
Format string `json:"format"`
|
||||
BaseURL string `json:"base_url,omitempty"`
|
||||
Proxy string `json:"proxy,omitempty"`
|
||||
Proxy string `json:"proxy,omitempty"` // masked: the placeholder when set
|
||||
ProxySet bool `json:"proxy_set"`
|
||||
Model string `json:"model"`
|
||||
APIKeyHint string `json:"api_key_hint,omitempty"`
|
||||
RatePerSecond float64 `json:"rate_per_second"`
|
||||
@@ -641,7 +642,8 @@ func llmProfileDTO(p *db.LLMProfile) LLMProfileDTO {
|
||||
Name: p.Name,
|
||||
Format: p.Format,
|
||||
BaseURL: p.BaseURL,
|
||||
Proxy: p.Proxy,
|
||||
Proxy: maskedSecret(p.Proxy),
|
||||
ProxySet: p.Proxy != "",
|
||||
Model: p.Model,
|
||||
APIKeyHint: p.APIKeyHint,
|
||||
RatePerSecond: p.RatePerSecond,
|
||||
|
||||
@@ -172,6 +172,11 @@ func TestFindingTrafficAPIAndExport(t *testing.T) {
|
||||
if w = req("GET", fmt.Sprintf("%s/%d", base, first), ""); w.Code != 200 || !strings.Contains(w.Body.String(), `"binary":true`) {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
// The viewer reads the heads in clear; the table stores them sealed.
|
||||
var rawHead string
|
||||
if err := s.m.pg.QueryRow(`SELECT req_head FROM traffic_evidence_snapshots WHERE id=$1`, list.Bindings[0].SnapshotID).Scan(&rawHead); err != nil || !strings.HasPrefix(rawHead, "enc:v1:") || !strings.Contains(w.Body.String(), `"req_head":"POST / HTTP/1.1"`) {
|
||||
t.Fatalf("head at rest %q (%v), viewer %s", rawHead, err, w.Body)
|
||||
}
|
||||
if w = req("GET", fmt.Sprintf("%s/%d/body?side=response&offset=20001", base, first), ""); w.Code != 422 {
|
||||
t.Fatal(w.Code, w.Body)
|
||||
}
|
||||
|
||||
@@ -216,6 +216,10 @@ func (s *Server) toolCreateCustomTool() actool.CoreTool {
|
||||
if exist, _ := s.m.pg.GetTool(a.Key); exist != nil {
|
||||
return actool.Errorf("this key already exists: " + a.Key), nil
|
||||
}
|
||||
var err error
|
||||
if a.Exec, err = unmaskToolExec("", nil, a.Kind, a.Exec); err != nil { // nothing stored: a placeholder is refused
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if err := s.m.pg.CreateCustomTool(toDBTool(a)); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
@@ -239,6 +243,10 @@ func (s *Server) toolUpdateCustomTool() actool.CoreTool {
|
||||
if a.Kind == "http" && !hasSchemaProps(a.Schema) {
|
||||
return actool.Errorf("http tools must provide a parameter JSON Schema (cannot be empty)"), nil
|
||||
}
|
||||
var err error
|
||||
if a.Exec, err = unmaskToolExec(existing.Kind, existing.Exec, a.Kind, a.Exec); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
if err := s.m.pg.UpdateCustomTool(toDBTool(a)); err != nil {
|
||||
return actool.Errorf(err.Error()), nil
|
||||
}
|
||||
|
||||
@@ -144,3 +144,118 @@ func TestMCPMaskedEnvDestinationGuard(t *testing.T) {
|
||||
t.Fatal("explicit env with a new url was not saved")
|
||||
}
|
||||
}
|
||||
|
||||
// Secret settings and the LLM profile proxy: stored encrypted, never returned
|
||||
// in clear; on write the placeholder keeps, an explicit value replaces, ""
|
||||
// clears.
|
||||
func TestProxySecretsAPI(t *testing.T) {
|
||||
s, _, request := trafficEvidenceServer(t)
|
||||
pg := s.m.PG()
|
||||
t.Cleanup(func() {
|
||||
pg.Exec(`DELETE FROM settings WHERE key IN ('brave_search_api_key','tavily_search_api_key','web_search_proxy','global_proxy')`)
|
||||
pg.Exec(`DELETE FROM llm_profiles WHERE name='t-proxy'`)
|
||||
})
|
||||
secrets := []string{"brave-api-0007", "tavily-api-0008", "wsp-pass-0009", "gp-pass-0010", "llmp-pass-0011"}
|
||||
put := func(body string) {
|
||||
t.Helper()
|
||||
resp := request("PUT", "/api/settings", body)
|
||||
for _, sec := range secrets {
|
||||
if resp.Code != 200 || strings.Contains(resp.Body.String(), sec) {
|
||||
t.Fatalf("PUT /api/settings %s: HTTP %d or a secret returned in clear", body, resp.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
setting := func(k string) string {
|
||||
t.Helper()
|
||||
var raw string
|
||||
if err := pg.QueryRow(`SELECT value FROM settings WHERE key=$1`, k).Scan(&raw); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
v, _, err := pg.GetSetting(k)
|
||||
if err != nil || (raw != "" && !strings.HasPrefix(raw, "enc:v1:")) {
|
||||
t.Fatalf("setting %s stored as %q (%v)", k, raw, err)
|
||||
}
|
||||
return v
|
||||
}
|
||||
put(`{"brave_search_api_key":"brave-api-0007","tavily_search_api_key":"tavily-api-0008",
|
||||
"web_search_proxy":"socks5://u:wsp-pass-0009@ws.example:1080","global_proxy":"http://u:gp-pass-0010@gp.example:3128"}`)
|
||||
var got map[string]any
|
||||
resp := request("GET", "/api/settings", "")
|
||||
_ = json.Unmarshal(resp.Body.Bytes(), &got)
|
||||
for _, sec := range secrets {
|
||||
if strings.Contains(resp.Body.String(), sec) {
|
||||
t.Fatal("GET /api/settings returned a secret in clear")
|
||||
}
|
||||
}
|
||||
if got["global_proxy"] != "__masked__" || got["global_proxy_set"] != true || got["web_search_proxy"] != "__masked__" || got["web_search_proxy_set"] != true {
|
||||
t.Fatalf("proxies not masked with a presence flag: %v", got)
|
||||
}
|
||||
|
||||
put(`{"web_search_proxy":"__masked__","global_proxy":"__masked__","brave_search_api_key":"__masked__"}`)
|
||||
if setting("web_search_proxy") != "socks5://u:wsp-pass-0009@ws.example:1080" || setting("global_proxy") != "http://u:gp-pass-0010@gp.example:3128" ||
|
||||
setting("brave_search_api_key") != "brave-api-0007" || setting("tavily_search_api_key") != "tavily-api-0008" {
|
||||
t.Fatal("a masked placeholder replaced a stored secret")
|
||||
}
|
||||
put(`{"global_proxy":"http://gp2.example:3128","web_search_proxy":""}`)
|
||||
if setting("global_proxy") != "http://gp2.example:3128" || setting("web_search_proxy") != "" {
|
||||
t.Fatal("explicit value must replace and \"\" must clear")
|
||||
}
|
||||
_ = json.Unmarshal(request("GET", "/api/settings", "").Body.Bytes(), &got)
|
||||
if got["web_search_proxy"] != "" || got["web_search_proxy_set"] != false {
|
||||
t.Fatalf("cleared proxy shown as %v / %v", got["web_search_proxy"], got["web_search_proxy_set"])
|
||||
}
|
||||
|
||||
// LLM profile proxy.
|
||||
id, err := pg.SaveProfile(&db.LLMProfile{Name: "t-proxy", Format: "openai", Model: "m", APIKey: "k", Proxy: "http://u:llmp-pass-0011@llm.example:3128"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
proxy := func() string {
|
||||
var raw *string
|
||||
if err := pg.QueryRow(`SELECT proxy FROM llm_profiles WHERE id=$1`, id).Scan(&raw); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if raw != nil && !strings.HasPrefix(*raw, "enc:v1:") {
|
||||
t.Fatalf("profile proxy stored in clear: %q", *raw)
|
||||
}
|
||||
p, err := pg.ProfileByID(id)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return p.Proxy
|
||||
}
|
||||
for _, path := range []string{"/api/llm/profiles", "/api/llm"} {
|
||||
if resp := request("GET", path, ""); resp.Code != 200 || strings.Contains(resp.Body.String(), "llmp-pass-0011") {
|
||||
t.Fatalf("GET %s: HTTP %d or the proxy returned in clear", path, resp.Code)
|
||||
}
|
||||
}
|
||||
var list struct{ Profiles []LLMProfileDTO }
|
||||
_ = json.Unmarshal(request("GET", "/api/llm/profiles", "").Body.Bytes(), &list)
|
||||
for _, p := range list.Profiles {
|
||||
if p.Name == "t-proxy" && (p.Proxy != "__masked__" || !p.ProxySet) {
|
||||
t.Fatalf("profile proxy not masked: %+v", p)
|
||||
}
|
||||
}
|
||||
save := func(proxyJSON string, want int) {
|
||||
t.Helper()
|
||||
body := `{"id":` + strconv.FormatInt(id, 10) + `,"name":"t-proxy","format":"openai","model":"m","proxy":` + proxyJSON + `}`
|
||||
if resp := request("POST", "/api/llm/profiles", body); resp.Code != want {
|
||||
t.Fatalf("save proxy %s: want HTTP %d, got %d %s", proxyJSON, want, resp.Code, resp.Body.String())
|
||||
}
|
||||
}
|
||||
save(`"__masked__"`, 200)
|
||||
if proxy() != "http://u:llmp-pass-0011@llm.example:3128" {
|
||||
t.Fatal("the placeholder replaced the stored profile proxy")
|
||||
}
|
||||
save(`"socks5://llm2.example:1080"`, 200)
|
||||
if proxy() != "socks5://llm2.example:1080" {
|
||||
t.Fatal("explicit proxy not saved")
|
||||
}
|
||||
save(`""`, 200)
|
||||
if proxy() != "" {
|
||||
t.Fatal("empty proxy must clear")
|
||||
}
|
||||
if resp := request("POST", "/api/llm/profiles", `{"name":"t-proxy","format":"openai","model":"m","proxy":"__masked__"}`); resp.Code != 400 {
|
||||
t.Fatalf("a placeholder on a new profile must be refused, got HTTP %d", resp.Code)
|
||||
}
|
||||
}
|
||||
|
||||
+27
-4
@@ -1355,7 +1355,8 @@ func (s *Server) getLLM(w http.ResponseWriter, r *http.Request) {
|
||||
"provider": s.llmCfg.Provider(),
|
||||
"model": s.llmCfg.Model,
|
||||
"base_url": s.llmCfg.BaseURL,
|
||||
"proxy": s.llmCfg.Proxy,
|
||||
"proxy": maskedSecret(s.llmCfg.Proxy),
|
||||
"proxy_set": s.llmCfg.Proxy != "",
|
||||
"key_set": s.llmCfg.APIKey != "",
|
||||
"rate_per_second": s.llmCfg.RatePerSecond,
|
||||
"rate_per_minute": s.llmCfg.RatePerMinute,
|
||||
@@ -1383,6 +1384,11 @@ func (s *Server) setLLM(w http.ResponseWriter, r *http.Request) {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
if isMaskedSecret(req.Proxy) { // getLLM returns the proxy masked: keep the current one
|
||||
s.cfgMu.Lock()
|
||||
req.Proxy = s.llmCfg.Proxy
|
||||
s.cfgMu.Unlock()
|
||||
}
|
||||
cfg := agent.ConfigFrom(req.Provider, req.Model, req.BaseURL, req.APIKey, req.Proxy)
|
||||
cfg.RatePerSecond, cfg.RatePerMinute = req.RatePerSecond, req.RatePerMinute
|
||||
cfg.ThinkingType = req.ThinkingType
|
||||
@@ -1455,6 +1461,14 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) {
|
||||
stored = p
|
||||
}
|
||||
}
|
||||
if isMaskedSecret(req.Proxy) { // the form echoes the masked proxy: test with the stored one
|
||||
req.Proxy = ""
|
||||
if stored == nil {
|
||||
s.cfgMu.Lock()
|
||||
req.Proxy = s.llmCfg.Proxy
|
||||
s.cfgMu.Unlock()
|
||||
}
|
||||
}
|
||||
if stored != nil {
|
||||
if strings.TrimSpace(req.Provider) == "" {
|
||||
req.Provider = stored.Format
|
||||
@@ -3487,8 +3501,10 @@ func (s *Server) settingsPayload() map[string]any {
|
||||
"web_search_backend": backend,
|
||||
"brave_key_set": strings.TrimSpace(braveKey) != "",
|
||||
"tavily_key_set": strings.TrimSpace(tavilyKey) != "",
|
||||
"web_search_proxy": proxy, // dedicated egress proxy (http/https/socks5), empty = direct
|
||||
"global_proxy": s.m.GlobalProxy(), // global egress proxy (http/https/socks5), all target traffic goes through it, empty = direct
|
||||
"web_search_proxy": maskedSecret(proxy), // proxies may carry user:pass: masked when set, "" = direct
|
||||
"web_search_proxy_set": proxy != "",
|
||||
"global_proxy": maskedSecret(s.m.GlobalProxy()), // global egress proxy, masked when set, "" = direct
|
||||
"global_proxy_set": s.m.GlobalProxy() != "",
|
||||
"python_interpreter": strings.TrimSpace(pyStored), // user/auto-set value (empty = use runtime detection)
|
||||
"workers": s.m.Workers(), // number of concurrent worker agents (default 3); applies to tasks started afterwards
|
||||
"task_concurrency_enabled": concOn, // task concurrency cap switch (default off)
|
||||
@@ -3591,6 +3607,9 @@ func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
// A masked placeholder echoed back keeps the stored secret, like null.
|
||||
req.BraveKey, req.TavilyKey = keepIfMasked(req.BraveKey), keepIfMasked(req.TavilyKey)
|
||||
req.WebSearchProxy, req.GlobalProxy = keepIfMasked(req.WebSearchProxy), keepIfMasked(req.GlobalProxy)
|
||||
if req.ConstraintsInjectPlanner != nil {
|
||||
if err := s.m.pg.SetBool(settingConstraintsInjectPlanner, *req.ConstraintsInjectPlanner); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
@@ -3788,13 +3807,17 @@ func (s *Server) testWebSearch(w http.ResponseWriter, r *http.Request) {
|
||||
writeErr(w, 400, err.Error())
|
||||
return
|
||||
}
|
||||
_, backend, storedBraveKey, storedTavilyKey, _ := s.m.WebSearch()
|
||||
_, backend, storedBraveKey, storedTavilyKey, storedProxy := s.m.WebSearch()
|
||||
if strings.TrimSpace(req.Backend) != "" {
|
||||
backend = req.Backend
|
||||
}
|
||||
// Proxy is taken from the form as-is (empty = direct), so testing reflects exactly
|
||||
// what's shown — including an intentional "clear proxy to test direct" before saving.
|
||||
// The masked placeholder the form shows for a saved proxy means the saved one.
|
||||
proxy := strings.TrimSpace(req.Proxy)
|
||||
if isMaskedSecret(proxy) {
|
||||
proxy = storedProxy
|
||||
}
|
||||
// API keys are secrets the form omits when already saved, so fall back to stored.
|
||||
braveKey := storedBraveKey
|
||||
if strings.TrimSpace(req.BraveKey) != "" {
|
||||
|
||||
@@ -754,6 +754,11 @@ func (s *Server) pgListTools(w http.ResponseWriter, r *http.Request) {
|
||||
if ts == nil {
|
||||
ts = []*db.Tool{}
|
||||
}
|
||||
for _, tool := range ts {
|
||||
if !tool.System {
|
||||
tool.Exec = maskToolExec(tool.Exec)
|
||||
}
|
||||
}
|
||||
// Usage is best-effort decoration. Runtime tool resolution keeps using the plain
|
||||
// catalog query, so agent assembly never pays for this aggregate.
|
||||
counts, countErr := pg.ToolUsageCounts()
|
||||
@@ -864,6 +869,26 @@ func (s *Server) pgListMCP(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, 200, map[string]any{"servers": ms})
|
||||
}
|
||||
|
||||
// maskedSecret is what a response shows for a stored secret such as a proxy
|
||||
// URL: the whole value is replaced by the placeholder, "" stays "".
|
||||
func maskedSecret(v string) string {
|
||||
if v == "" {
|
||||
return ""
|
||||
}
|
||||
return notify.MaskedPrefix
|
||||
}
|
||||
|
||||
// isMaskedSecret reports a placeholder sent back on write: keep the stored value.
|
||||
func isMaskedSecret(v string) bool { return notify.IsMasked(strings.TrimSpace(v)) }
|
||||
|
||||
// keepIfMasked turns a placeholder in an optional field into nil (keep).
|
||||
func keepIfMasked(v *string) *string {
|
||||
if v != nil && isMaskedSecret(*v) {
|
||||
return nil
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
// maskMCPEnv replaces every non-empty env value (API keys, auth headers) with
|
||||
// notify.MaskedPrefix so the API never returns it in clear.
|
||||
func maskMCPEnv(raw json.RawMessage) json.RawMessage {
|
||||
@@ -1949,6 +1974,16 @@ func (s *Server) pgSaveProfile(w http.ResponseWriter, r *http.Request) {
|
||||
p := body.LLMProfile
|
||||
p.APIKey = body.APIKey
|
||||
p.Streaming = body.Streaming == nil || *body.Streaming
|
||||
// The list returns the proxy masked: the placeholder sent back keeps the
|
||||
// stored proxy and is never stored itself.
|
||||
if isMaskedSecret(p.Proxy) {
|
||||
stored, err := pg.ProfileByID(p.ID)
|
||||
if p.ID == 0 || err != nil || stored == nil {
|
||||
writeErr(w, 400, "proxy is a masked placeholder but no stored proxy to keep; enter the proxy or clear it")
|
||||
return
|
||||
}
|
||||
p.Proxy = stored.Proxy
|
||||
}
|
||||
// Output cap: negative values are meaningless, clamp to zero (= do not send the
|
||||
// field). The field-name switch only matters for Chat Completions — anthropic and
|
||||
// openai-responses hard-code their field names; storing it would only mislead later readers, so it is cleared for any non-openai format. Unknown values are cleared too, to avoid throwing DB CHECK errors at the user.
|
||||
@@ -2150,6 +2185,14 @@ func (s *Server) pgListModels(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
}
|
||||
}
|
||||
if isMaskedSecret(req.Proxy) { // the form echoes the masked proxy: use the stored one
|
||||
req.Proxy = ""
|
||||
if req.ProfileID != nil {
|
||||
if p, err := s.m.pg.ProfileByID(*req.ProfileID); err == nil && p != nil {
|
||||
req.Proxy = p.Proxy
|
||||
}
|
||||
}
|
||||
}
|
||||
provider := strings.TrimSpace(req.Provider)
|
||||
// Ollama serves its model list keyless on the local server: no key gate.
|
||||
if apiKey == "" && provider != "ollama" {
|
||||
|
||||
@@ -0,0 +1,248 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// A placeholder keeps the stored secret only for the url and kind it was saved
|
||||
// with: a changed destination is refused before any write or request.
|
||||
func TestCustomToolExecDestinationGuard(t *testing.T) {
|
||||
s, _, request := trafficEvidenceServer(t)
|
||||
pg := s.m.PG()
|
||||
var hits atomic.Int32
|
||||
echo := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
hits.Add(1)
|
||||
io.WriteString(w, "auth="+r.Header.Get("Authorization"))
|
||||
}))
|
||||
t.Cleanup(echo.Close)
|
||||
t.Cleanup(func() { pg.Exec(`DELETE FROM tools WHERE key LIKE 't_secguard%'`) })
|
||||
const secret = "Bearer guard-plain-0001"
|
||||
schema := `{"type":"object","properties":{"q":{"type":"string"}}}`
|
||||
body := func(kind, url, hdr string) string {
|
||||
e, _ := json.Marshal(map[string]any{"method": "GET", "url": url, "headers": map[string]string{"Authorization": hdr}})
|
||||
return `{"key":"t_secguard","kind":"` + kind + `","enabled":true,"schema":` + schema + `,"exec":` + string(e) + `}`
|
||||
}
|
||||
storedExec := func() (url, auth string) {
|
||||
tl, err := pg.GetTool("t_secguard")
|
||||
if err != nil || tl == nil {
|
||||
t.Fatalf("GetTool: %v", err)
|
||||
}
|
||||
var e struct {
|
||||
URL string `json:"url"`
|
||||
Headers map[string]string `json:"headers"`
|
||||
}
|
||||
_ = json.Unmarshal(tl.Exec, &e)
|
||||
return e.URL, e.Headers["Authorization"]
|
||||
}
|
||||
if resp := request("POST", "/api/tools/custom", body("http", echo.URL, secret)); resp.Code != 200 {
|
||||
t.Fatalf("create: HTTP %d %s", resp.Code, resp.Body.String())
|
||||
}
|
||||
moved := echo.URL + "/elsewhere"
|
||||
|
||||
// 1. Masked header, changed url: refused, nothing written.
|
||||
if resp := request("PUT", "/api/tools/custom/t_secguard", body("http", moved, "__masked__")); resp.Code != 400 {
|
||||
t.Fatalf("masked header with a changed url must be refused, got HTTP %d", resp.Code)
|
||||
}
|
||||
if url, auth := storedExec(); url != echo.URL || auth != secret {
|
||||
t.Fatalf("a refused update changed the stored exec: %s %s", url, auth)
|
||||
}
|
||||
// An empty url against a stored one is a change too.
|
||||
if resp := request("PUT", "/api/tools/custom/t_secguard", body("http", "", "__masked__")); resp.Code != 400 {
|
||||
t.Fatalf("masked header with an emptied url must be refused, got HTTP %d", resp.Code)
|
||||
}
|
||||
|
||||
// 2. Masked header, same url: kept.
|
||||
if resp := request("PUT", "/api/tools/custom/t_secguard", body("http", echo.URL, "__masked__")); resp.Code != 200 {
|
||||
t.Fatalf("masked header with the same url: HTTP %d %s", resp.Code, resp.Body.String())
|
||||
}
|
||||
if url, auth := storedExec(); url != echo.URL || auth != secret {
|
||||
t.Fatalf("the stored secret was not kept: %s %s", url, auth)
|
||||
}
|
||||
|
||||
// 3. Test endpoint, masked header, changed url: refused, no request sent.
|
||||
e, _ := json.Marshal(map[string]any{"method": "GET", "url": moved, "headers": map[string]string{"Authorization": "__masked__"}})
|
||||
if resp := request("POST", "/api/tools/custom/test", `{"key":"t_secguard","kind":"http","exec":`+string(e)+`}`); resp.Code != 400 {
|
||||
t.Fatalf("test endpoint with a changed url must be refused, got HTTP %d %s", resp.Code, resp.Body.String())
|
||||
}
|
||||
if n := hits.Load(); n != 0 {
|
||||
t.Fatalf("the echo server received %d request(s); the stored secret left", n)
|
||||
}
|
||||
|
||||
// 4. Kind change while keeping a placeholder: refused.
|
||||
if resp := request("PUT", "/api/tools/custom/t_secguard", body("command", echo.URL, "__masked__")); resp.Code != 400 {
|
||||
t.Fatalf("a kind change with a placeholder must be refused, got HTTP %d", resp.Code)
|
||||
}
|
||||
if tl, _ := pg.GetTool("t_secguard"); tl == nil || tl.Kind != "http" {
|
||||
t.Fatalf("a refused kind change was written: %v", tl)
|
||||
}
|
||||
|
||||
// 5. An explicit value with a changed url is allowed.
|
||||
if resp := request("PUT", "/api/tools/custom/t_secguard", body("http", moved, "Bearer typed-0002")); resp.Code != 200 {
|
||||
t.Fatalf("explicit header with a changed url: HTTP %d %s", resp.Code, resp.Body.String())
|
||||
}
|
||||
if url, auth := storedExec(); url != moved || auth != "Bearer typed-0002" {
|
||||
t.Fatalf("explicit replace not stored: %s %s", url, auth)
|
||||
}
|
||||
}
|
||||
|
||||
// A custom tool's exec: headers, proxy and body are masked by GET /api/tools;
|
||||
// the placeholder keeps the stored value on update and in the test endpoint, a
|
||||
// create with it is refused, and the executor runs on the clear exec.
|
||||
func TestCustomToolExecMasking(t *testing.T) {
|
||||
s, _, request := trafficEvidenceServer(t)
|
||||
pg := s.m.PG()
|
||||
echo := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
io.WriteString(w, "auth="+r.Header.Get("Authorization")+" body="+string(b))
|
||||
}))
|
||||
t.Cleanup(echo.Close)
|
||||
t.Cleanup(func() { pg.Exec(`DELETE FROM tools WHERE key LIKE 't_secexec%'`) })
|
||||
secrets := []string{"hdr-plain-0001", "body-plain-0002", "proxy-plain-0003"}
|
||||
schema := `{"type":"object","properties":{"q":{"type":"string"}}}`
|
||||
execOf := func(hdr, body, proxy string) string {
|
||||
e, _ := json.Marshal(map[string]any{"method": "POST", "url": echo.URL,
|
||||
"headers": map[string]string{"Authorization": hdr, "X-Empty": ""}, "body": body, "proxy": proxy})
|
||||
return string(e)
|
||||
}
|
||||
toolBody := func(key, exec string) string {
|
||||
return `{"key":"` + key + `","kind":"http","enabled":true,"schema":` + schema + `,"exec":` + exec + `}`
|
||||
}
|
||||
stored := func() map[string]any {
|
||||
tl, err := pg.GetTool("t_secexec")
|
||||
if err != nil || tl == nil {
|
||||
t.Fatalf("GetTool: %v", err)
|
||||
}
|
||||
var e map[string]any
|
||||
_ = json.Unmarshal(tl.Exec, &e)
|
||||
return e
|
||||
}
|
||||
|
||||
plain := execOf("Bearer "+secrets[0], secrets[1], "http://u:"+secrets[2]+"@127.0.0.1:1")
|
||||
if resp := request("POST", "/api/tools/custom", toolBody("t_secexec", plain)); resp.Code != 200 {
|
||||
t.Fatalf("create: HTTP %d %s", resp.Code, resp.Body.String())
|
||||
}
|
||||
var raw string
|
||||
if err := pg.QueryRow(`SELECT exec::text FROM tools WHERE key='t_secexec'`).Scan(&raw); err != nil || !strings.Contains(raw, "__artex_enc") {
|
||||
t.Fatalf("exec not sealed at rest: %s %v", raw, err)
|
||||
}
|
||||
|
||||
// List: masked with *_set flags, nothing in clear.
|
||||
resp := request("GET", "/api/tools", "")
|
||||
for _, sec := range secrets {
|
||||
if resp.Code != 200 || strings.Contains(resp.Body.String(), sec) {
|
||||
t.Fatalf("GET /api/tools: HTTP %d or a secret returned in clear", resp.Code)
|
||||
}
|
||||
}
|
||||
var list struct {
|
||||
Tools []struct {
|
||||
Key string `json:"key"`
|
||||
Exec json.RawMessage `json:"exec"`
|
||||
}
|
||||
}
|
||||
_ = json.Unmarshal(resp.Body.Bytes(), &list)
|
||||
var masked json.RawMessage
|
||||
for _, tl := range list.Tools {
|
||||
if tl.Key == "t_secexec" {
|
||||
masked = tl.Exec
|
||||
}
|
||||
}
|
||||
var m map[string]any
|
||||
_ = json.Unmarshal(masked, &m)
|
||||
h, _ := m["headers"].(map[string]any)
|
||||
if h["Authorization"] != "__masked__" || h["X-Empty"] != "" || m["body"] != "__masked__" || m["proxy"] != "__masked__" ||
|
||||
m["headers_set"] != true || m["body_set"] != true || m["proxy_set"] != true || m["url"] != echo.URL {
|
||||
t.Fatalf("exec not masked as expected: %s", masked)
|
||||
}
|
||||
|
||||
// Update echoing the masked exec: every stored value is kept, no flag stored.
|
||||
if resp := request("PUT", "/api/tools/custom/t_secexec", toolBody("t_secexec", string(masked))); resp.Code != 200 {
|
||||
t.Fatalf("update with placeholders: HTTP %d %s", resp.Code, resp.Body.String())
|
||||
}
|
||||
e := stored()
|
||||
if e["headers"].(map[string]any)["Authorization"] != "Bearer "+secrets[0] || e["body"] != secrets[1] ||
|
||||
!strings.Contains(e["proxy"].(string), secrets[2]) || e["body_set"] != nil {
|
||||
t.Fatalf("placeholders did not keep the stored exec: %v", e)
|
||||
}
|
||||
|
||||
// Test endpoint: the saved tool's key resolves the placeholders.
|
||||
noProxy := execOf("__masked__", "__masked__", "")
|
||||
if resp := request("POST", "/api/tools/custom/test", `{"kind":"http","exec":`+noProxy+`}`); resp.Code != 400 {
|
||||
t.Fatalf("a placeholder without a saved tool must be refused, got HTTP %d", resp.Code)
|
||||
}
|
||||
resp = request("POST", "/api/tools/custom/test", `{"key":"t_secexec","kind":"http","exec":`+noProxy+`}`)
|
||||
if resp.Code != 200 || !strings.Contains(resp.Body.String(), "auth=Bearer "+secrets[0]+" body="+secrets[1]) {
|
||||
t.Fatalf("test endpoint did not use the stored values: HTTP %d %s", resp.Code, resp.Body.String())
|
||||
}
|
||||
|
||||
// Replace the header, keep the body, clear the proxy.
|
||||
if resp := request("PUT", "/api/tools/custom/t_secexec", toolBody("t_secexec", execOf("Bearer new-0004", "__masked__", ""))); resp.Code != 200 {
|
||||
t.Fatalf("update: HTTP %d %s", resp.Code, resp.Body.String())
|
||||
}
|
||||
e = stored()
|
||||
if e["headers"].(map[string]any)["Authorization"] != "Bearer new-0004" || e["body"] != secrets[1] || e["proxy"] != "" {
|
||||
t.Fatalf("replace/keep/clear: %v", e)
|
||||
}
|
||||
if resp := request("GET", "/api/tools", ""); strings.Contains(resp.Body.String(), "new-0004") {
|
||||
t.Fatal("a replaced secret is returned in clear")
|
||||
}
|
||||
|
||||
// Create with a placeholder: refused, nothing written.
|
||||
for _, ex := range []string{execOf("__masked__", "", ""), execOf("", "__masked__", ""), execOf("", "", "__masked__")} {
|
||||
if resp := request("POST", "/api/tools/custom", toolBody("t_secexec_new", ex)); resp.Code != 400 {
|
||||
t.Fatalf("create with a placeholder must be refused, got HTTP %d", resp.Code)
|
||||
}
|
||||
}
|
||||
if tl, _ := pg.GetTool("t_secexec_new"); tl != nil {
|
||||
t.Fatal("a refused create wrote a row")
|
||||
}
|
||||
|
||||
// The executor receives the clear exec from the sealed row.
|
||||
tools, err := s.customTools()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, tl := range tools {
|
||||
if tl.Name() != "t_secexec" {
|
||||
continue
|
||||
}
|
||||
res, err := tl.Call(context.Background(), json.RawMessage(`{"q":"x"}`), nil)
|
||||
if err != nil || !strings.Contains(res.Flatten(), "auth=Bearer new-0004 body="+secrets[1]) {
|
||||
t.Fatalf("executor output %q %v", res.Flatten(), err)
|
||||
}
|
||||
return
|
||||
}
|
||||
t.Fatal("t_secexec not built by customTools")
|
||||
}
|
||||
|
||||
// The platform tools follow the API: create refuses a placeholder, update keeps the stored value.
|
||||
func TestPlatformCustomToolPlaceholders(t *testing.T) {
|
||||
s, _, _ := trafficEvidenceServer(t)
|
||||
pg := s.m.PG()
|
||||
t.Cleanup(func() { pg.Exec(`DELETE FROM tools WHERE key LIKE 't_secexec%'`) })
|
||||
in := func(exec string) json.RawMessage {
|
||||
return json.RawMessage(`{"key":"t_secexec_pt","kind":"command","exec":` + exec + `}`)
|
||||
}
|
||||
create, update := s.toolCreateCustomTool(), s.toolUpdateCustomTool()
|
||||
if res, _ := create.Call(context.Background(), in(`{"command":"echo","proxy":"__masked__"}`), nil); !res.IsError {
|
||||
t.Fatalf("create with a placeholder must fail: %s", res.Flatten())
|
||||
}
|
||||
if res, _ := create.Call(context.Background(), in(`{"command":"echo","body":"pt-body-plain"}`), nil); res.IsError {
|
||||
t.Fatalf("create: %s", res.Flatten())
|
||||
}
|
||||
if res, _ := update.Call(context.Background(), in(`{"command":"echo 2","body":"__masked__"}`), nil); res.IsError {
|
||||
t.Fatalf("update: %s", res.Flatten())
|
||||
}
|
||||
tl, err := pg.GetTool("t_secexec_pt")
|
||||
if err != nil || tl == nil || !jsonHas(tl.Exec, "pt-body-plain") || !jsonHas(tl.Exec, "echo 2") {
|
||||
t.Fatalf("update with a placeholder must keep the stored body: %v %v", tl, err)
|
||||
}
|
||||
}
|
||||
|
||||
func jsonHas(raw json.RawMessage, s string) bool { return strings.Contains(string(raw), s) }
|
||||
@@ -662,7 +662,8 @@ function CustomToolDialog({
|
||||
setTesting(true);
|
||||
setTestResult(null);
|
||||
try {
|
||||
const r = await api.testCustomTool({ kind, exec: buildExec(), params });
|
||||
// key: a saved tool's stored secrets replace the masked placeholders.
|
||||
const r = await api.testCustomTool({ key: tool?.key, kind, exec: buildExec(), params });
|
||||
setTestResult(r);
|
||||
} catch (e) {
|
||||
setTestResult({ output: (e as Error).message, is_error: true });
|
||||
|
||||
+1
-1
@@ -1113,7 +1113,7 @@ export const api = {
|
||||
t: Pick<Tool, "description" | "schema" | "agents" | "enabled" | "kind" | "exec" | "deferred">,
|
||||
) => put<{ ok: boolean }>(`/tools/custom/${key}`, t),
|
||||
deleteCustomTool: (key: string) => del<{ deleted: string }>(`/tools/custom/${key}`),
|
||||
testCustomTool: (body: { kind: string; exec: Record<string, unknown>; params: Record<string, unknown> }) =>
|
||||
testCustomTool: (body: { key?: string; kind: string; exec: Record<string, unknown>; params: Record<string, unknown> }) =>
|
||||
post<{ output: string; is_error: boolean }>("/tools/custom/test", body),
|
||||
detectPython: () => post<{ python_interpreter: string }>("/settings/python/detect", {}),
|
||||
|
||||
|
||||
Reference in New Issue
Block a user