admin/ARTEX
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(tasks): add runtime controls and inherited context

Browse Source
This commit is contained in:
pphua committed 2026-08-24 19:29:00 +08:00
1 parent 18926c129f
commit 4392b4fee1
50 files changed
+8291 -663
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

traffic/traffic.go
+23 -42
View File
@@ -11,6 +11,7 @@ import (
"database/sql"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"io/fs"
@@ -28,9 +29,9 @@ import (
"time"
"unicode/utf8"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/norma/permission"
actool "github.com/Autumn-27/norma/tool"
"github.com/Autumn-27/artex/db"
mproxy "github.com/lqqyt2423/go-mitmproxy/proxy"
_ "modernc.org/sqlite"
)
@@ -307,15 +308,6 @@ VALUES(?,?,?,?,?,?,?)`,
return
}
<<<<<<< Updated upstream
reqTxt := fmt.Sprintf("%s %s %s\n%s\n%s", method, f.Request.URL.RequestURI(), f.Request.Proto, headerLines(f.Request.Header), reqBody)
respTxt := fmt.Sprintf("HTTP %d\n%s\n%s", f.Response.StatusCode, headerLines(f.Response.Header), respBody)
_ = os.WriteFile(filepath.Join(exDir, "request.http"), []byte(reqTxt), 0o644)
_ = os.WriteFile(filepath.Join(exDir, "response.http"), []byte(respTxt), 0o644)
meta := fmt.Sprintf(`{"id":%q,"host":%q,"method":%q,"url":%q,"template":%q,"status":%d,"content_type":%q}`,
id, host, method, f.Request.URL.String(), tmpl, f.Response.StatusCode, ct)
_ = os.WriteFile(filepath.Join(exDir, "meta.json"), []byte(meta), 0o644)
=======
for _, h := range []string{reqB.hash, respB.hash} {
if h == "" {
continue
@@ -325,7 +317,6 @@ VALUES(?,?,?,?,?,?,?)`,
return
}
}
>>>>>>> Stashed changes
if t.fts {
// The index is fed from memory, so a body that spilled to _blobs is still
@@ -480,6 +471,26 @@ func headerLines(h map[string][]string) string {
return b.String()
}
// requestHeaderLines restores the HTTP Host header, which net/http stores on
// Request.Host rather than in Header. Keeping it in request.http makes the raw
// capture complete and directly replayable.
func requestHeaderLines(req *mproxy.Request) string {
headers := req.Header.Clone()
headers.Del("Host")
host := req.URL.Host
if raw := req.Raw(); raw != nil && strings.TrimSpace(raw.Host) != "" {
host = raw.Host
}
var b strings.Builder
if host = strings.TrimSpace(host); host != "" {
b.WriteString("Host: ")
b.WriteString(host)
b.WriteByte('\n')
}
b.WriteString(headerLines(headers))
return b.String()
}
func sanitize(s string) string {
r := strings.NewReplacer("/", "_", "\\", "_", ":", "_", "?", "_", "*", "_", "\"", "_", "<", "_", ">", "_", "|", "_")
out := r.Replace(s)
@@ -944,14 +955,6 @@ func (t *Traffic) reapStage(stageDir string) {
// so picking "api.example.com" never sweeps "api.example.com.cn". Returns rows
// deleted. Duplicate hosts are harmless (idempotent deletes, single tree pass).
func (t *Traffic) DeleteHostsExact(hosts []string) (int64, error) {
<<<<<<< Updated upstream
t.wmu.Lock()
defer t.wmu.Unlock()
var deleted int64
removed := make(map[string]bool)
for _, h := range hosts {
res, err := t.db.Exec(`DELETE FROM exchanges WHERE host=?`, h)
=======
stage, err := t.StageDeleteHostsExact(hosts)
if err != nil {
return 0, err
@@ -1028,30 +1031,9 @@ func (t *Traffic) StageDeleteHostsExact(hosts []string) (*HostDeleteStage, error
stage.tx = tx
for _, h := range unique {
n, err := t.deleteWhere(tx, `host=?`, h)
>>>>>>> Stashed changes
if err != nil {
return deleted, err
return fail(err)
}
<<<<<<< Updated upstream
n, _ := res.RowsAffected()
deleted += n
if !removed[h] {
os.RemoveAll(filepath.Join(t.dir, sanitize(h)))
removed[h] = true
}
}
if deleted > 0 {
_ = t.gcBlobs()
}
return deleted, nil
}
// gcBlobs removes content-addressed blobs in _blobs that no remaining exchange
// tree references. References appear as "@blob sha256:<hex>" lines inside the
// tree's request.http/response.http files (see bodyOrBlob). Best-effort: walk
// errors only skip cleanup of the affected files. Callers must hold wmu so this
// never races a concurrent record() writing a fresh blob + tree.
=======
stage.deleted += n
}
@@ -1122,7 +1104,6 @@ func (s *HostDeleteStage) Commit() error {
// left the buckets behind permanently. Best-effort: walk errors only skip the
// affected entries. Callers must hold wmu so this never races a concurrent
// record() writing a fresh blob + reference.
>>>>>>> Stashed changes
func (t *Traffic) gcBlobs() error {
refs := make(map[string]struct{})
rows, err := t.db.Query(`SELECT DISTINCT hash FROM blob_refs`)
traffic/traffic_test.go
+164
View File
@@ -2,12 +2,30 @@ package traffic
import (
"fmt"
"net/http"
"net/url"
"os"
"path/filepath"
"strings"
"testing"
mproxy "github.com/lqqyt2423/go-mitmproxy/proxy"
)
func TestRequestHeaderLinesIncludesHost(t *testing.T) {
req := &mproxy.Request{
URL: &url.URL{Host: "target.example:8443"},
Header: http.Header{"Accept": []string{"application/json"}},
}
got := requestHeaderLines(req)
if !strings.Contains(got, "Host: target.example:8443\n") {
t.Fatalf("request headers missing Host: %q", got)
}
if !strings.Contains(got, "Accept: application/json\n") {
t.Fatalf("request headers missing regular header: %q", got)
}
}
// TestDeleteHost verifies the delete contract: rows for hosts containing the
// substring are removed together with their file trees, non-matching hosts are
// untouched, and the count is right.
@@ -162,6 +180,152 @@ VALUES(?,?,?,?,?,?,?,?,?,?,?)`,
}
}
func TestDeleteHostsExactReportsTreeRemovalFailure(t *testing.T) {
dir := t.TempDir()
tr, err := Open(dir, "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
defer tr.Close()
const host = "api.example.com"
if _, err := tr.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path)
VALUES(?,?,?,?,?,?,?,?,?,?,?)`,
"1-0001", 1, host, "GET", "/", "http://"+host+"/", 200, "text/html", 0, 0, host+"/GET/1-0001"); err != nil {
t.Fatal(err)
}
notDir := filepath.Join(dir, "not-a-directory")
if err := os.WriteFile(notDir, []byte("x"), 0o644); err != nil {
t.Fatal(err)
}
tr.dir = notDir
n, err := tr.DeleteHostsExact([]string{host})
if err == nil {
t.Fatal("DeleteHostsExact returned nil after traffic tree removal failed")
}
if n != 0 {
t.Fatalf("deleted=%d, want 0 after atomic rollback", n)
}
var count int
if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges WHERE host=?`, host).Scan(&count); err != nil || count != 1 {
t.Fatalf("rolled-back index count=%d err=%v, want 1", count, err)
}
}
func TestDeleteHostsExactRollsBackWholeIndexBatch(t *testing.T) {
dir := t.TempDir()
tr, err := Open(dir, "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
defer tr.Close()
for i, host := range []string{"a.example.com", "b.example.com"} {
id := fmt.Sprintf("1-%04d", i+1)
if err := os.MkdirAll(filepath.Join(dir, host, "GET", id), 0o755); err != nil {
t.Fatal(err)
}
if _, err := tr.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path)
VALUES(?,?,?,?,?,?,?,?,?,?,?)`, id, 1, host, "GET", "/", "http://"+host+"/", 200, "text/html", 0, 0, host+"/GET/"+id); err != nil {
t.Fatal(err)
}
}
if _, err := tr.DB().Exec(`CREATE TRIGGER fail_second_host BEFORE DELETE ON exchanges
WHEN OLD.host='b.example.com' BEGIN SELECT RAISE(ABORT, 'forced delete failure'); END`); err != nil {
t.Fatal(err)
}
if n, err := tr.DeleteHostsExact([]string{"a.example.com", "b.example.com"}); err == nil || n != 0 {
t.Fatalf("DeleteHostsExact failure=(%d,%v), want (0,error)", n, err)
}
var count int
if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges`).Scan(&count); err != nil || count != 2 {
t.Fatalf("rolled-back index count=%d err=%v, want 2", count, err)
}
for _, host := range []string{"a.example.com", "b.example.com"} {
if _, err := os.Stat(filepath.Join(dir, host)); err != nil {
t.Fatalf("tree %s changed despite index rollback: %v", host, err)
}
}
}
func TestStageDeleteHostsExactRollbackRestoresIndexAndTree(t *testing.T) {
dir := t.TempDir()
tr, err := Open(dir, "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
defer tr.Close()
const host = "rollback.example.com"
const id = "1-0001"
tree := filepath.Join(dir, host, "GET", id)
if err := os.MkdirAll(tree, 0o755); err != nil {
t.Fatal(err)
}
marker := filepath.Join(tree, "request.http")
if err := os.WriteFile(marker, []byte("original"), 0o644); err != nil {
t.Fatal(err)
}
if _, err := tr.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path)
VALUES(?,?,?,?,?,?,?,?,?,?,?)`, id, 1, host, "GET", "/", "http://"+host+"/", 200, "text/html", 0, 0, host+"/GET/"+id); err != nil {
t.Fatal(err)
}
stage, err := tr.StageDeleteHostsExact([]string{host})
if err != nil {
t.Fatal(err)
}
if stage.Deleted() != 1 {
t.Fatalf("staged deleted=%d, want 1", stage.Deleted())
}
if _, err := os.Stat(filepath.Join(dir, host)); !os.IsNotExist(err) {
t.Fatalf("host tree was not staged: %v", err)
}
if err := stage.Rollback(); err != nil {
t.Fatal(err)
}
if got, err := os.ReadFile(marker); err != nil || string(got) != "original" {
t.Fatalf("restored tree content=%q err=%v", got, err)
}
var count int
if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges WHERE host=?`, host).Scan(&count); err != nil || count != 1 {
t.Fatalf("restored index count=%d err=%v, want 1", count, err)
}
}
func TestStageDeleteHostsExactRollbackReportsRestoreFailure(t *testing.T) {
dir := t.TempDir()
tr, err := Open(dir, "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
defer tr.Close()
const host = "conflict.example.com"
hostDir := filepath.Join(dir, host)
if err := os.MkdirAll(hostDir, 0o755); err != nil {
t.Fatal(err)
}
stage, err := tr.StageDeleteHostsExact([]string{host})
if err != nil {
t.Fatal(err)
}
// Simulate an out-of-band conflicting destination. Rollback must surface the
// failed rename instead of claiming the external data was restored.
if err := os.WriteFile(hostDir, []byte("conflict"), 0o644); err != nil {
t.Fatal(err)
}
if err := stage.Rollback(); err == nil || !strings.Contains(err.Error(), "restore") {
t.Fatalf("rollback err=%v, want restore failure", err)
}
if _, err := os.Stat(stage.stageDir); err != nil {
t.Fatalf("staging was removed after failed restore: %v", err)
}
}
// TestDeleteHostGCBlobs verifies blob garbage collection: after a host's trees
// are removed, blobs referenced by no remaining exchange are deleted, while
// blobs still referenced (including shared ones) survive.