diff --git a/.env.example b/.env.example index 8a556a2..e8e0793 100644 --- a/.env.example +++ b/.env.example @@ -24,6 +24,15 @@ ARTEX_LLM_PROXY= # 154 bytes per tool result plus one system note. # ARTEX_UNTRUSTED_BOUNDARY=0 +# Noise profile of the requests ARTEX sends to targets (WebFetch, http custom tools, enrich +# probe). Default off: no limit, requests unchanged. quiet = 2 req/s per host, burst 4, +# 1 in flight; balanced = 5 req/s, burst 10, 2 in flight; aggressive = 20 req/s, burst 40, +# 4 in flight. A profile also adds jitter, backs off on 429/503 and merges identical GETs +# of one task (short cache). Bash-launched tools are not throttled. +# ARTEX_NOISE_PROFILE=balanced +# User-Agent set on target requests that have none, only under a profile. Empty = unchanged. +# ARTEX_TARGET_USER_AGENT= + # Update banner (optional). ARTEX_UPDATE_CHECK=1 lets the backend check the public GitHub # releases API for a newer ARTEX release at most once per 24 h and show a dismissible banner. # Off by default: unset, no update request is made. diff --git a/README.md b/README.md index ad5d048..07cff9b 100644 --- a/README.md +++ b/README.md @@ -264,6 +264,8 @@ Optional: `ARTEX_LLM_PROVIDER` (anthropic | openai | openai-responses | ollama) **Untrusted-output boundary (on by default)**: tool results for the worker, planner, chat and main agent are framed between nonce-tagged markers so the model treats them as data, not instructions. Unset or empty, the boundary is on; set `ARTEX_UNTRUSTED_BOUNDARY=0` (or `false` / `off` / `no`) to disable it. Measured cost: about 154 bytes per tool result plus one system note. +**Noise profile (off by default)**: `ARTEX_NOISE_PROFILE` paces the requests ARTEX itself sends to targets (WebFetch, http custom tools, the enrich probe). Unset or `off`, nothing changes: no limit, no cache. Set `quiet` (2 req/s per host, burst 4, 1 in flight), `balanced` (5 req/s, burst 10, 2 in flight) or `aggressive` (20 req/s, burst 40, 4 in flight), e.g. `ARTEX_NOISE_PROFILE=balanced` in `.env` or the environment. A profile also adds a small jitter, backs off on 429/503 (honouring `Retry-After`, bounded retries, GET/HEAD/OPTIONS only), merges identical GETs of one task and caches them briefly (60/30/10 s; never POST, PUT, PATCH or DELETE), and logs a per-host noise line each minute. `ARTEX_TARGET_USER_AGENT` sets a User-Agent on target requests that carry none, under a profile; empty keeps the current one. Tools launched through Bash are not paced yet. + **Common flags**: `./start.sh -addr :8787 -proxy :8788` (`-addr` frontend+API, `-proxy` traffic recording proxy). The startup script passes arguments through to `artex` verbatim. ### Reverse proxy deployment (HTTPS / only exposing 443) diff --git a/agent/webfetch.go b/agent/webfetch.go index 323bd5e..ad323c8 100644 --- a/agent/webfetch.go +++ b/agent/webfetch.go @@ -2,7 +2,11 @@ package agent import ( "context" + "encoding/json" + "net/url" + "slices" + "github.com/Autumn-27/artex/egress" "github.com/Autumn-27/artex/intercept" actool "github.com/Autumn-27/norma/tool" ) @@ -17,11 +21,43 @@ import ( // on the direct retry is turned off for the run, since it would leave without // the proxy. The tool is added after the session's own tools, where agentcore // adds it, so the tool list keeps its order. caCert is the recording proxy's -// CA, trusted to verify the HTTPS it re-signs. tools is not modified. +// CA, trusted to verify the HTTPS it re-signs. Each call goes through the shared +// egress governor (see governedFetch). tools is not modified. func withWebFetch(ctx context.Context, tools []actool.CoreTool, proxy, caCert string) []actool.CoreTool { scope, scoped := intercept.TargetScopeFromContext(ctx) gated := scoped && scope.Enabled != nil && scope.Enabled() webFetch := actool.NewWebFetch(actool.WebFetchConfig{Proxy: proxy, CACert: caCert, CheckRedirect: intercept.RedirectCheck("WebFetch"), DialContext: intercept.PinnedDialer("WebFetch", scoped), DisableDirectRetry: gated}) - return append(tools[:len(tools):len(tools)], webFetch) + return append(tools[:len(tools):len(tools)], governedFetch{webFetch, egress.Shared()}) +} + +// governedFetch runs each WebFetch call as one send under the egress governor +// (off by default; then the call runs unchanged), under a noise profile: +// the rate, jitter and concurrency of the URL's host, and, within one task, +// identical calls (WebFetch only sends GET) coalesced and cached. A failed call +// is not cached. WebFetch ignores tc, so a shared result fits every caller. +// Norma owns the HTTP client, so redirect hops are not paced separately and a +// 429/503 is not retried. +type governedFetch struct { + actool.CoreTool + g *egress.Governor +} + +func (f governedFetch) Call(ctx context.Context, input json.RawMessage, tc *actool.ToolContext) (actool.Result, error) { + var in struct { + URL string `json:"url"` + } + if json.Unmarshal(input, &in) != nil { + return f.CoreTool.Call(ctx, input, tc) + } + u, err := url.Parse(in.URL) + if err != nil || u.Host == "" { + return f.CoreTool.Call(ctx, input, tc) // the tool reports the bad URL + } + r, err := egress.Do(f.g, ctx, u, string(input), func() (actool.Result, bool, error) { + r, err := f.CoreTool.Call(ctx, input, tc) + return r, err == nil && !r.IsError, err + }) + r.Content = slices.Clone(r.Content) // each caller gets its own slice + return r, err } diff --git a/docker-compose.yml b/docker-compose.yml index c8da3e6..a008023 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -40,6 +40,9 @@ services: # disables it. Cost: about 154 bytes per tool result plus one system note. ARTEX_UNTRUSTED_BOUNDARY: ${ARTEX_UNTRUSTED_BOUNDARY:-} ARTEX_UPDATE_CHECK: ${ARTEX_UPDATE_CHECK:-} + # Noise profile for target requests: off (default) | quiet | balanced | aggressive. + ARTEX_NOISE_PROFILE: ${ARTEX_NOISE_PROFILE:-} + ARTEX_TARGET_USER_AGENT: ${ARTEX_TARGET_USER_AGENT:-} ports: - "${ARTEX_BIND_ADDR:-127.0.0.1}:8787:8787" - "127.0.0.1:8788:8788" diff --git a/egress/egress.go b/egress/egress.go new file mode 100644 index 0000000..94c3926 --- /dev/null +++ b/egress/egress.go @@ -0,0 +1,453 @@ +// Package egress is the governor the platform-owned HTTP requests to a target +// pass through (WebFetch, the http custom tools, the enrich probe). A noise +// profile (ARTEX_NOISE_PROFILE) sets every lever: a token bucket, a jitter and a +// concurrency cap per host, a global concurrency cap, backoff on 429/503, and, +// for the GET requests of one task, coalescing and a short cache. The default +// profile, off, is a strict passthrough: requests and timing are unchanged. +package egress + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "log" + "math/rand/v2" + "net" + "net/http" + "net/url" + "os" + "sort" + "strconv" + "strings" + "sync" + "time" + "unicode" + + "github.com/Autumn-27/artex/intercept" +) + +// Profile is one noise level. The zero Profile is off. Later levers (scanner +// flags, proxy throttle) are new fields here. +type Profile struct { + Name string + RPS float64 // token refill per host, requests per second; 0 = off + Burst int // bucket size per host + MaxPerHost int // requests in flight per host + MaxGlobal int // requests in flight across all hosts + Jitter float64 // each send waits a uniform [0, min(Jitter/RPS, MaxJitter)] + CacheTTL time.Duration // per-task cache of identical GETs; 0 = coalescing only + MaxRetries int // retries after a 429/503 (bodiless GET/HEAD/OPTIONS only) + BaseDelay time.Duration // backoff without Retry-After, doubled per retry + MaxDelay time.Duration // cap of one backoff; a longer Retry-After is not retried + UserAgent string // ARTEX_TARGET_USER_AGENT, set on requests that have none; "" = unchanged +} + +// MaxJitter caps the jitter of one send. +const MaxJitter = 250 * time.Millisecond + +// Profiles are the values of ARTEX_NOISE_PROFILE besides off (the default). +var Profiles = map[string]Profile{ + "quiet": {Name: "quiet", RPS: 2, Burst: 4, MaxPerHost: 1, MaxGlobal: 8, Jitter: 0.5, CacheTTL: time.Minute, MaxRetries: 3, BaseDelay: 2 * time.Second, MaxDelay: time.Minute}, + "balanced": {Name: "balanced", RPS: 5, Burst: 10, MaxPerHost: 2, MaxGlobal: 16, Jitter: 0.2, CacheTTL: 30 * time.Second, MaxRetries: 3, BaseDelay: time.Second, MaxDelay: 30 * time.Second}, + "aggressive": {Name: "aggressive", RPS: 20, Burst: 40, MaxPerHost: 4, MaxGlobal: 32, Jitter: 0.1, CacheTTL: 10 * time.Second, MaxRetries: 2, BaseDelay: 500 * time.Millisecond, MaxDelay: 10 * time.Second}, +} + +// ProfileFromEnv returns the profile ARTEX_NOISE_PROFILE names; unset, "off" or +// unknown is off. +func ProfileFromEnv() Profile { + name := strings.ToLower(strings.TrimSpace(os.Getenv("ARTEX_NOISE_PROFILE"))) + p, ok := Profiles[name] + if !ok { + if name != "" && name != "off" { + log.Printf("[egress] ARTEX_NOISE_PROFILE=%q is not off, quiet, balanced or aggressive; using off", name) + } + return Profile{Name: "off"} + } + if ua := strings.TrimSpace(os.Getenv("ARTEX_TARGET_USER_AGENT")); strings.ContainsFunc(ua, unicode.IsControl) { + log.Printf("[egress] ARTEX_TARGET_USER_AGENT contains a control character; ignored") + } else { + p.UserAgent = ua + } + return p +} + +// Shared is the one governor of the process, built from the environment on +// first use; WebFetch, the http custom tools and the enrich probe all use it. +// Under a profile it logs the noise counter of each host once a minute. +var Shared = sync.OnceValue(func() *Governor { + g := New(ProfileFromEnv()) + log.Printf("[egress] noise profile: %s", g.p.Name) + if !g.off() { + go func() { + for range time.Tick(time.Minute) { + for _, l := range g.report() { + log.Print(l) + } + } + }() + } + return g +}) + +// Governor holds the per-host state and the per-task response cache. +type Governor struct { + p Profile + global chan struct{} + mu sync.Mutex + hosts map[string]*hostState + flights map[string]*flight +} + +type hostState struct { + slots chan struct{} + tokens float64 + last time.Time + notBefore time.Time // end of a 429/503 backoff + // noise counter of the current minute, and the peak per minute so far + sent, dupes, inflight, peakInflight, peakRPM int +} + +type flight struct { + done chan struct{} + val any + err error + exp time.Time // zero until cached +} + +// New builds a Governor for p. +func New(p Profile) *Governor { + return &Governor{p: p, global: make(chan struct{}, max(p.MaxGlobal, 1)), hosts: map[string]*hostState{}, flights: map[string]*flight{}} +} + +func (g *Governor) off() bool { return g.p.RPS <= 0 } + +// HostKey is the lowercased host, with the port when it is not the default. +func HostKey(u *url.URL) string { + h, p := strings.ToLower(u.Hostname()), u.Port() + if p == "" || (u.Scheme == "http" && p == "80") || (u.Scheme == "https" && p == "443") { + return h + } + return net.JoinHostPort(h, p) +} + +func (g *Governor) jitterBound() time.Duration { + return min(time.Duration(g.p.Jitter/g.p.RPS*float64(time.Second)), MaxJitter) +} + +func (g *Governor) jitter() time.Duration { return rand.N(g.jitterBound() + 1) } + +func (g *Governor) host(key string) *hostState { + g.mu.Lock() + defer g.mu.Unlock() + h := g.hosts[key] + if h == nil { + h = &hostState{slots: make(chan struct{}, max(g.p.MaxPerHost, 1)), tokens: float64(g.p.Burst), last: time.Now()} + g.hosts[key] = h + } + return h +} + +// acquire waits for a host slot, a token of the host's bucket (or the end of +// its backoff) plus the jitter, then a global slot, taken last so a waiting +// host holds no global slot. release frees both; it may be called twice. +func (g *Governor) acquire(ctx context.Context, key string) (release func(), err error) { + h := g.host(key) + select { + case h.slots <- struct{}{}: + case <-ctx.Done(): + return nil, ctx.Err() + } + if err = sleep(ctx, g.reserve(h)+g.jitter()); err == nil { + select { + case g.global <- struct{}{}: + case <-ctx.Done(): + err = ctx.Err() + } + } + if err != nil { + <-h.slots + return nil, err + } + g.mu.Lock() + h.sent++ + h.inflight++ + h.peakInflight = max(h.peakInflight, h.inflight) + g.mu.Unlock() + var once sync.Once + return func() { + once.Do(func() { g.mu.Lock(); h.inflight--; g.mu.Unlock(); <-g.global; <-h.slots }) + }, nil +} + +// reserve takes a token from h and returns how long to wait for it. +func (g *Governor) reserve(h *hostState) time.Duration { + g.mu.Lock() + defer g.mu.Unlock() + now := time.Now() + h.tokens = min(float64(g.p.Burst), h.tokens+now.Sub(h.last).Seconds()*g.p.RPS) - 1 + h.last = now + return max(time.Duration(-h.tokens/g.p.RPS*float64(time.Second)), h.notBefore.Sub(now), 0) +} + +// backoff is the wait after a 429/503: Retry-After (seconds or HTTP date), else +// BaseDelay doubled per retry, capped at MaxDelay. ok is false when Retry-After +// asks for more than MaxDelay. +func (g *Governor) backoff(h http.Header, attempt int) (d time.Duration, ok bool) { + v := strings.TrimSpace(h.Get("Retry-After")) + if n, err := strconv.Atoi(v); err == nil && n >= 0 { + return time.Duration(n) * time.Second, time.Duration(n)*time.Second <= g.p.MaxDelay + } + if t, err := http.ParseTime(v); err == nil { + return max(time.Until(t), 0), time.Until(t) <= g.p.MaxDelay + } + if d = g.p.BaseDelay << attempt; d <= 0 || d > g.p.MaxDelay { + d = g.p.MaxDelay + } + return d, true +} + +func sleep(ctx context.Context, d time.Duration) error { + if d <= 0 { + return ctx.Err() + } + t := time.NewTimer(d) + defer t.Stop() + select { + case <-t.C: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// report returns one noise line per active host and starts a new minute. +func (g *Governor) report() []string { + g.mu.Lock() + defer g.mu.Unlock() + var out []string + for key, h := range g.hosts { + if h.sent == 0 && h.dupes == 0 { + continue + } + h.peakRPM = max(h.peakRPM, h.sent) + out = append(out, fmt.Sprintf("[egress] noise profile=%s host=%s rpm=%d peak_rpm=%d peak_inflight=%d duplicates_avoided=%d", + g.p.Name, key, h.sent, h.peakRPM, h.peakInflight, h.dupes)) + h.sent, h.dupes, h.peakInflight = 0, 0, h.inflight + } + sort.Strings(out) + return out +} + +var errNotShared = errors.New("egress: response not shared") + +// share runs fn once for the concurrent callers of key and keeps its value for +// the ttl fn returns. A caller that waited on a flight that failed for a reason +// of the leader's own (its context, a response too large to share) runs fn +// itself, uncached. +func (g *Governor) share(ctx context.Context, host, key string, fn func() (any, time.Duration, error)) (any, error) { + g.mu.Lock() + f := g.flights[key] + if f != nil && !f.exp.IsZero() && time.Now().After(f.exp) { + f = nil + } + if f != nil { + g.mu.Unlock() + select { + case <-f.done: + case <-ctx.Done(): + return nil, ctx.Err() + } + if errors.Is(f.err, errNotShared) || errors.Is(f.err, context.Canceled) || errors.Is(f.err, context.DeadlineExceeded) { + v, _, err := fn() + return v, err + } + h := g.host(host) + g.mu.Lock() + h.dupes++ + g.mu.Unlock() + return f.val, f.err + } + f = &flight{done: make(chan struct{})} + g.flights[key] = f + for k, o := range g.flights { + if !o.exp.IsZero() && time.Now().After(o.exp) { + delete(g.flights, k) + } + } + g.mu.Unlock() + v, ttl, err := fn() + f.val, f.err = v, err + g.mu.Lock() + if err == nil && ttl > 0 { + f.exp = time.Now().Add(ttl) + } else if g.flights[key] == f { + delete(g.flights, key) + } + close(f.done) + g.mu.Unlock() + return v, err +} + +// taskKey identifies the task (or conversation) of ctx; "" when there is none, +// and then nothing is coalesced or cached. +func taskKey(ctx context.Context) string { + s, _ := intercept.TargetScopeFromContext(ctx) + if s.TaskID != "" { + return "task:" + s.TaskID + } + if s.ConversationID != 0 { + return "conv:" + strconv.FormatInt(s.ConversationID, 10) + } + return "" +} + +// Do runs fn as one send to u for a caller whose HTTP client is not ours +// (WebFetch): the rate, jitter and concurrency of u's host and, when id is not +// empty and ctx has a task, coalescing and caching under id within the task +// (keep reports whether the value may be cached). fn hides the response, so a +// 429/503 is not retried here. Off, fn just runs. +func Do[T any](g *Governor, ctx context.Context, u *url.URL, id string, fn func() (val T, keep bool, err error)) (T, error) { + if g.off() { + v, _, err := fn() + return v, err + } + host := HostKey(u) + run := func() (any, time.Duration, error) { + release, err := g.acquire(ctx, host) + if err != nil { + return nil, 0, err + } + defer release() + v, keep, err := fn() + if !keep { + return v, 0, err + } + return v, g.p.CacheTTL, err + } + var v any + var err error + if task := taskKey(ctx); task != "" && id != "" { + v, err = g.share(ctx, host, task+"\x00"+id, run) + } else { + v, _, err = run() + } + t, _ := v.(T) + return t, err +} + +// Transport returns base (nil: http.DefaultTransport) governed by g; off, base +// itself. +func (g *Governor) Transport(base http.RoundTripper) http.RoundTripper { + if base == nil { + base = http.DefaultTransport + } + if g.off() { + return base + } + return &transport{g, base} +} + +type transport struct { + g *Governor + base http.RoundTripper +} + +const maxSharedBody = 2 << 20 // a larger response is neither coalesced nor cached + +func (t *transport) RoundTrip(req *http.Request) (*http.Response, error) { + if ua := t.g.p.UserAgent; ua != "" && req.Header.Get("User-Agent") == "" { + req = req.Clone(req.Context()) + req.Header.Set("User-Agent", ua) + } + task := taskKey(req.Context()) + if task == "" || req.Method != http.MethodGet || (req.Body != nil && req.Body != http.NoBody) { + return t.send(req) + } + // The key holds the headers: a request with other credentials is its own. + var hdr bytes.Buffer + _ = req.Header.Write(&hdr) + var own *http.Response // the leader's response when too large to share + v, err := t.g.share(req.Context(), HostKey(req.URL), task+"\x00GET "+req.URL.String()+"\x00"+hdr.String(), func() (any, time.Duration, error) { + resp, err := t.send(req) + if err != nil { + return nil, 0, err + } + body, err := io.ReadAll(io.LimitReader(resp.Body, maxSharedBody+1)) + if err != nil || len(body) > maxSharedBody { + resp.Body = struct { + io.Reader + io.Closer + }{io.MultiReader(bytes.NewReader(body), resp.Body), resp.Body} + own = resp + return nil, 0, errNotShared + } + resp.Body.Close() + ttl := t.g.p.CacheTTL + if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500 { + ttl = 0 // coalesced, never cached + } + return &stored{resp, body}, ttl, nil + }) + if own != nil || err != nil { + return own, err + } + s := v.(*stored) + r := *s.resp + r.Header, r.Body, r.ContentLength, r.Request = s.resp.Header.Clone(), io.NopCloser(bytes.NewReader(s.body)), int64(len(s.body)), req + return &r, nil +} + +// stored is a shared response; each caller gets a copy. +type stored struct { + resp *http.Response + body []byte +} + +// send makes one request within the host's limits. A 429/503 pauses the host +// and, for a bodiless GET/HEAD/OPTIONS, is retried up to MaxRetries times; a +// request that may change state is never retried. The slots are held until +// the response body is closed. +func (t *transport) send(req *http.Request) (*http.Response, error) { + host := HostKey(req.URL) + for attempt := 0; ; attempt++ { + release, err := t.g.acquire(req.Context(), host) + if err != nil { + return nil, err + } + resp, err := t.base.RoundTrip(req) + if err != nil { + release() + return nil, err + } + resp.Body = &releaseBody{resp.Body, release} + if resp.StatusCode != http.StatusTooManyRequests && resp.StatusCode != http.StatusServiceUnavailable { + return resp, nil + } + d, ok := t.g.backoff(resp.Header, attempt) + t.g.mu.Lock() + h := t.g.hosts[host] + if until := time.Now().Add(min(d, t.g.p.MaxDelay)); until.After(h.notBefore) { + h.notBefore = until + } + t.g.mu.Unlock() + idempotent := req.Method == http.MethodGet || req.Method == http.MethodHead || req.Method == http.MethodOptions + if !ok || !idempotent || (req.Body != nil && req.Body != http.NoBody) || attempt >= t.g.p.MaxRetries { + return resp, nil + } + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 64<<10)) + resp.Body.Close() + } +} + +type releaseBody struct { + io.ReadCloser + release func() +} + +func (b *releaseBody) Close() error { + err := b.ReadCloser.Close() + b.release() + return err +} diff --git a/egress/egress_test.go b/egress/egress_test.go new file mode 100644 index 0000000..d28e65c --- /dev/null +++ b/egress/egress_test.go @@ -0,0 +1,244 @@ +package egress + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/Autumn-27/artex/intercept" +) + +// open is a profile that does not limit, for the tests of one lever. +var open = Profile{Name: "test", RPS: 1000, Burst: 1000, MaxPerHost: 100, MaxGlobal: 100, CacheTTL: time.Minute, MaxRetries: 2, BaseDelay: 10 * time.Millisecond, MaxDelay: 2 * time.Second} + +func task(id string) context.Context { + return intercept.WithTargetScope(context.Background(), intercept.TargetScope{TaskID: id}) +} + +// server counts its hits and the peak of requests in flight; h, if set, answers. +type server struct { + *httptest.Server + hits, inflight, peak atomic.Int64 +} + +func newServer(t *testing.T, delay time.Duration, h func(n int64, w http.ResponseWriter, r *http.Request)) *server { + s := &server{} + s.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + n, cur := s.hits.Add(1), s.inflight.Add(1) + for p := s.peak.Load(); cur > p && !s.peak.CompareAndSwap(p, cur); p = s.peak.Load() { + } + time.Sleep(delay) + s.inflight.Add(-1) + if h != nil { + h(n, w, r) + return + } + _, _ = io.WriteString(w, "ok") + })) + t.Cleanup(s.Close) + return s +} + +func do(t *testing.T, g *Governor, ctx context.Context, method, u string) (int, string) { + req, _ := http.NewRequestWithContext(ctx, method, u, nil) + resp, err := (&http.Client{Transport: g.Transport(nil)}).Do(req) + if err != nil { + t.Errorf("%s %s: %v", method, u, err) + return 0, "" + } + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + return resp.StatusCode, string(b) +} + +func parallel(n int, f func(i int)) { + var wg sync.WaitGroup + for i := range n { + wg.Add(1) + go func() { defer wg.Done(); f(i) }() + } + wg.Wait() +} + +func TestOffIsStrictPassthrough(t *testing.T) { + g := New(Profile{Name: "off"}) + if g.Transport(http.DefaultTransport) != http.DefaultTransport { + t.Fatal("off: Transport must return the base transport itself") + } + s := newServer(t, 50*time.Millisecond, nil) + start := time.Now() + parallel(20, func(int) { do(t, g, task("t1"), "GET", s.URL) }) + if d := time.Since(start); s.hits.Load() != 20 || s.peak.Load() < 10 || d > time.Second { + t.Fatalf("off: hits %d, peak %d, %s; want 20 parallel hits, no delay", s.hits.Load(), s.peak.Load(), d) + } + var calls atomic.Int64 + u, _ := url.Parse(s.URL) + for range 2 { + _, _ = Do(g, task("t1"), u, "x", func() (int, bool, error) { calls.Add(1); return 0, true, nil }) + } + if calls.Load() != 2 || len(g.report()) != 0 { + t.Fatalf("off: Do calls = %d, report %v; want 2 calls, no counter", calls.Load(), g.report()) + } +} + +// TestProfilesRespectRateBurstConcurrency sends burst+extra requests at once: +// the burst leaves at once, the rest at RPS, never more than MaxPerHost at a time. +func TestProfilesRespectRateBurstConcurrency(t *testing.T) { + for _, tc := range []struct { + name string + extra int + }{{"quiet", 2}, {"balanced", 5}} { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + p := Profiles[tc.name] + g := New(p) + s := newServer(t, 20*time.Millisecond, nil) + n := p.Burst + tc.extra + start := time.Now() + parallel(n, func(int) { do(t, g, context.Background(), "GET", s.URL) }) + d := time.Since(start) + if rate := float64(tc.extra) / d.Seconds(); rate > p.RPS*1.1 || d > 5*time.Second { + t.Errorf("%d requests in %s: %.1f req/s after the burst, want <= %.0f", n, d, rate, p.RPS) + } + if s.hits.Load() != int64(n) || s.peak.Load() > int64(p.MaxPerHost) { + t.Errorf("hits %d, peak in flight %d; want %d, <= %d", s.hits.Load(), s.peak.Load(), n, p.MaxPerHost) + } + }) + } +} + +func TestBackoff(t *testing.T) { + g := New(open) + s := newServer(t, 0, func(n int64, w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/once" && n == 1: + w.Header().Set("Retry-After", "1") + w.WriteHeader(http.StatusTooManyRequests) + case r.URL.Path == "/far": + w.Header().Set("Retry-After", "3600") + w.WriteHeader(http.StatusTooManyRequests) + case r.URL.Path == "/always": + w.WriteHeader(http.StatusServiceUnavailable) + } + }) + start := time.Now() + if code, _ := do(t, g, context.Background(), "GET", s.URL+"/once"); code != 200 || s.hits.Load() != 2 || time.Since(start) < 950*time.Millisecond { + t.Fatalf("Retry-After 1: status %d, hits %d, %s", code, s.hits.Load(), time.Since(start)) + } + if code, _ := do(t, g, context.Background(), "GET", s.URL+"/always"); code != 503 || s.hits.Load() != 5 { + t.Fatalf("bounded: status %d, hits %d; want 503 after 1+2 more hits", code, s.hits.Load()) + } + if code, _ := do(t, g, context.Background(), "POST", s.URL+"/always"); code != 503 || s.hits.Load() != 6 { + t.Fatalf("POST: status %d, hits %d; want no retry", code, s.hits.Load()) + } + if code, _ := do(t, g, context.Background(), "GET", s.URL+"/far"); code != 429 || s.hits.Load() != 7 { + t.Fatalf("Retry-After beyond MaxDelay: status %d, hits %d; want no retry", code, s.hits.Load()) + } + h := http.Header{"Retry-After": {time.Now().Add(2 * time.Second).UTC().Format(http.TimeFormat)}} + if d, ok := g.backoff(h, 0); !ok || d <= 0 || d > 2*time.Second { + t.Fatalf("HTTP-date Retry-After: %s, %v", d, ok) + } +} + +func TestDedupeAndCache(t *testing.T) { + p := open + p.CacheTTL = 200 * time.Millisecond + g := New(p) + s := newServer(t, 50*time.Millisecond, nil) + parallel(5, func(int) { + if _, body := do(t, g, task("t1"), "GET", s.URL); body != "ok" { + t.Errorf("body = %q", body) + } + }) + do(t, g, task("t1"), "GET", s.URL) // cached + if s.hits.Load() != 1 { + t.Fatalf("coalesced and cached GETs: hits %d, want 1", s.hits.Load()) + } + do(t, g, task("t2"), "GET", s.URL) // another task + do(t, g, context.Background(), "GET", s.URL) // no task + time.Sleep(250 * time.Millisecond) + do(t, g, task("t1"), "GET", s.URL) // expired + if s.hits.Load() != 4 { + t.Fatalf("other task, no task, expired: hits %d, want 4", s.hits.Load()) + } + methods := []string{"POST", "PUT", "PATCH", "DELETE"} + parallel(8, func(i int) { do(t, g, task("t1"), methods[i/2], s.URL) }) + if s.hits.Load() != 12 { + t.Fatalf("mutating methods: hits %d, want 12", s.hits.Load()) + } + if r := g.report(); len(r) != 1 || !strings.Contains(r[0], "rpm=12 ") || !strings.Contains(r[0], "duplicates_avoided=5") { + t.Fatalf("noise counter: %v", r) + } + // WebFetch path: Do coalesces and caches what fn keeps. + var calls atomic.Int64 + u, _ := url.Parse(s.URL) + fn := func(keep bool) func() (string, bool, error) { + return func() (string, bool, error) { calls.Add(1); time.Sleep(50 * time.Millisecond); return "r", keep, nil } + } + parallel(3, func(int) { _, _ = Do(g, task("t1"), u, "a", fn(true)) }) + _, _ = Do(g, task("t1"), u, "a", fn(true)) + _, _ = Do(g, task("t1"), u, "b", fn(false)) + _, _ = Do(g, task("t1"), u, "b", fn(false)) + if calls.Load() != 3 { + t.Fatalf("Do calls = %d, want 3", calls.Load()) + } +} + +func TestJitterBounded(t *testing.T) { + for name, want := range map[string]time.Duration{"quiet": 250 * time.Millisecond, "balanced": 40 * time.Millisecond, "aggressive": 5 * time.Millisecond} { + g := New(Profiles[name]) + if b := g.jitterBound(); b != want { + t.Fatalf("%s: bound %s, want %s", name, b, want) + } + var nonzero bool + for range 1000 { + d := g.jitter() + if d < 0 || d > want { + t.Fatalf("%s: jitter %s outside [0, %s]", name, d, want) + } + nonzero = nonzero || d > 0 + } + if !nonzero { + t.Fatalf("%s: jitter always 0", name) + } + } +} + +func TestUserAgentAndEnv(t *testing.T) { + var got atomic.Value + s := newServer(t, 0, func(_ int64, _ http.ResponseWriter, r *http.Request) { got.Store(r.UserAgent()) }) + p := open + do(t, New(p), context.Background(), "GET", s.URL) + if ua := got.Load().(string); !strings.HasPrefix(ua, "Go-http-client/") { + t.Fatalf("unset: User-Agent = %q, want the Go default", ua) + } + p.UserAgent = "engagement-ua" + do(t, New(p), context.Background(), "GET", s.URL) + if ua := got.Load().(string); ua != "engagement-ua" { + t.Fatalf("set: User-Agent = %q", ua) + } + if p := ProfileFromEnv(); p.RPS != 0 || p.Name != "off" { + t.Fatalf("default profile: %+v, want off", p) + } + t.Setenv("ARTEX_NOISE_PROFILE", "Balanced") + t.Setenv("ARTEX_TARGET_USER_AGENT", "x/1") + if p := ProfileFromEnv(); p.Name != "balanced" || p.RPS != 5 || p.UserAgent != "x/1" { + t.Fatalf("balanced from env: %+v", p) + } + t.Setenv("ARTEX_NOISE_PROFILE", "loud") + if p := ProfileFromEnv(); p.RPS != 0 { + t.Fatalf("unknown profile: %+v, want off", p) + } + for in, want := range map[string]string{"http://Example.COM:80/x": "example.com", "https://a.test:8443/": "a.test:8443"} { + if u, _ := url.Parse(in); HostKey(u) != want { + t.Errorf("HostKey(%s) = %q, want %q", in, HostKey(u), want) + } + } +} diff --git a/enrich/enrich.go b/enrich/enrich.go index 0d86f9b..e3821cf 100644 --- a/enrich/enrich.go +++ b/enrich/enrich.go @@ -21,6 +21,7 @@ import ( "time" "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/egress" "github.com/Autumn-27/artex/intercept" "github.com/miekg/dns" @@ -97,7 +98,8 @@ func New(as *db.AssetStore, proxy func() string, workers int) *Engine { // intercept.PinnedDialer, under the target scope of the request's context: a // probe without a scope is refused. When the scope's gate is on, the probe // connects to the target itself, without the proxy, so that the dialer decides -// on the target and connects only to the addresses a pin covers. +// on the target and connects only to the addresses a pin covers. Every request +// goes through the shared egress governor (the noise profile; off by default). func buildClient(proxy func() string) *http.Client { tr := &http.Transport{ TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, @@ -118,7 +120,7 @@ func buildClient(proxy func() string) *http.Client { } } return &http.Client{ - Transport: tr, + Transport: egress.Shared().Transport(tr), Timeout: httpTimeout, // cap redirects; keep them within scope by re-checking at probe time CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }, @@ -245,7 +247,6 @@ func (e *Engine) doHTTP(ctx context.Context, id int64, rawURL string) { if err != nil { return } - req.Header.Set("User-Agent", "artex-enrich/1.0") resp, err := e.client.Do(req) if err != nil { return diff --git a/server/customtool.go b/server/customtool.go index d12c7d5..d39acfa 100644 --- a/server/customtool.go +++ b/server/customtool.go @@ -18,6 +18,7 @@ import ( "time" "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/egress" "github.com/Autumn-27/artex/guard" "github.com/Autumn-27/artex/intercept" "github.com/Autumn-27/norma/permission" @@ -457,9 +458,13 @@ func (s *Server) runHTTPTool(ctx context.Context, key string, execRaw json.RawMe client := &http.Client{Timeout: timeoutOr(spec.TimeoutMs, 30000)} // Every redirect hop goes through the target gate, as the rendered URL did. client.CheckRedirect = intercept.RedirectCheck(key) + // The shared egress governor paces every request and redirect hop per host + // under the noise profile (off by default: the base transport itself). + var base http.RoundTripper = http.DefaultTransport if tr := s.httpToolTransport(ctx, key, spec); tr != nil { - client.Transport = tr + base = tr } + client.Transport = egress.Shared().Transport(base) resp, err := client.Do(req) if err != nil { return actool.Errorf("request failed: " + err.Error()), nil