Merge pull request #198 from BaaRai/fix/enrich-http-gate

enrich: route the HTTP probe through the pinned dialer (F13)

agent/tools.go

@@ -221,7 +221,7 @@ func (t *ToolSet) SetNotifyFinding(fn func(int64, string)) { t.notifyFinding = f
 // enrich). Kept as an interface here to avoid coupling agent → enrich.
 type EnrichTrigger interface {
 	ResolveDomain(id int64, host string)
-	ProbeSite(id int64, url string)
+	ProbeSite(ctx context.Context, id int64, url string)
 }
 
 // WriteCounts breaks down what a worker persisted this run, by node kind, so the

enrich/enrich.go

@@ -7,6 +7,7 @@
 package enrich
 
 import (
+	"context"
 	"crypto/tls"
 	"html"
 	"io"
@@ -20,6 +21,7 @@ import (
 	"time"
 
 	"github.com/Autumn-27/artex/db"
+	"github.com/Autumn-27/artex/intercept"
 
 	"github.com/miekg/dns"
 	"github.com/projectdiscovery/dnsx/libs/dnsx"
@@ -34,8 +36,9 @@ const (
 
 type job struct {
 	kind jobKind
-	id   int64  // asset id (domain for DNS, site for HTTP)
-	arg  string // host (DNS) or url (HTTP)
+	id   int64           // asset id (domain for DNS, site for HTTP)
+	arg  string          // host (DNS) or url (HTTP)
+	ctx  context.Context // HTTP: carries the target scope the probe dials under
 }
 
 // Engine owns the resolver, the proxy-routed HTTP client, and the worker pool.
@@ -90,15 +93,23 @@ func New(as *db.AssetStore, proxy func() string, workers int) *Engine {
 // buildClient returns an HTTP client that dials via the recording proxy (resolved
 // per-request via proxy(), so the traffic-capture toggle applies live) and skips
 // TLS verification (the proxy re-signs with its MITM CA; targets are often
-// self-signed — this is a pentest probe).
+// self-signed — this is a pentest probe). Every connection goes through
+// 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.
 func buildClient(proxy func() string) *http.Client {
 	tr := &http.Transport{
 		TLSClientConfig:     &tls.Config{InsecureSkipVerify: true},
 		DisableKeepAlives:   true,
 		TLSHandshakeTimeout: httpTimeout,
+		DialContext:         intercept.PinnedDialer("enrich", true),
 	}
 	if proxy != nil {
-		tr.Proxy = func(*http.Request) (*url.URL, error) {
+		tr.Proxy = func(req *http.Request) (*url.URL, error) {
+			if intercept.GateOn(req.Context()) {
+				return nil, nil // direct: the dialer decides on the target
+			}
 			p := proxy()
 			if p == "" {
 				return nil, nil // direct
@@ -116,10 +127,14 @@ func buildClient(proxy func() string) *http.Client {
 
 // ResolveDomain enqueues a DNS resolution for a domain asset (ungated). host is the
 // FQDN. No-op for an empty engine.
-func (e *Engine) ResolveDomain(id int64, host string) { e.enqueue(job{jobDNS, id, host}) }
+func (e *Engine) ResolveDomain(id int64, host string) { e.enqueue(job{jobDNS, id, host, nil}) }
 
 // ProbeSite enqueues an HTTP probe for a web-asset (site). rawURL is the site URL.
-func (e *Engine) ProbeSite(id int64, rawURL string) { e.enqueue(job{jobHTTP, id, rawURL}) }
+// ctx carries the target scope of the call that asks for the probe: the probe
+// connects under that scope (see buildClient).
+func (e *Engine) ProbeSite(ctx context.Context, id int64, rawURL string) {
+	e.enqueue(job{jobHTTP, id, rawURL, ctx})
+}
 
 func (e *Engine) enqueue(j job) {
 	if e == nil || id0(j.id) || j.arg == "" {
@@ -155,7 +170,7 @@ func (e *Engine) worker() {
 			case jobDNS:
 				e.doDNS(j.id, j.arg)
 			case jobHTTP:
-				e.doHTTP(j.id, j.arg)
+				e.doHTTP(j.ctx, j.id, j.arg)
 			}
 		}
 	}
@@ -221,12 +236,12 @@ func (e *Engine) doDNS(id int64, host string) {
 
 var reTitle = regexp.MustCompile(`(?is)<title[^>]*>(.*?)</title>`)
 
-func (e *Engine) doHTTP(id int64, rawURL string) {
+func (e *Engine) doHTTP(ctx context.Context, id int64, rawURL string) {
 	host := hostOf(rawURL)
 	if host == "" {
 		return
 	}
-	req, err := http.NewRequest("GET", rawURL, nil)
+	req, err := http.NewRequestWithContext(ctx, "GET", rawURL, nil)
 	if err != nil {
 		return
 	}

intercept/enrich_probe_test.go

@@ -0,0 +1,56 @@
+package intercept_test
+
+import (
+	"context"
+	"net"
+	"slices"
+	"testing"
+	"time"
+
+	"github.com/Autumn-27/artex/enrich"
+	"github.com/Autumn-27/artex/intercept"
+)
+
+// F13: the HTTP probe of the enrich engine connects through the pinned dialer.
+// A probe without a target scope is refused and dials nothing; a probe under a
+// gated scope dials the address the pin holds, never the name, and not the
+// recording proxy. Without a database: no address reached here is served, so
+// no probe writes to the asset store.
+func TestEnrichProbeDialsThroughPinnedDialer(t *testing.T) {
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	proxyAddr := l.Addr().String() // closed: a dial to it is refused
+	l.Close()
+
+	for _, tc := range []struct {
+		name  string
+		proxy func() string
+	}{
+		{"direct", nil},
+		{"recording proxy set", func() string { return "http://" + proxyAddr }},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			g := intercept.NewDialGate(t)
+			g.Resolve("app.example.com", "203.0.113.7")
+			g.Pin("app.example.com", 8080, "203.0.113.7")
+			g.Exempt("http://" + proxyAddr)
+			e := enrich.New(nil, tc.proxy, 1)
+			t.Cleanup(e.Close)
+
+			// One worker: the probe without a scope runs to its end first.
+			e.ProbeSite(context.Background(), 1, "http://app.example.com:8080/")
+			e.ProbeSite(g.Ctx(), 2, "http://app.example.com:8080/")
+
+			deadline := time.Now().Add(5 * time.Second)
+			for len(g.Calls()) == 0 && time.Now().Before(deadline) {
+				time.Sleep(10 * time.Millisecond)
+			}
+			time.Sleep(100 * time.Millisecond)
+			if got := g.Calls(); !slices.Equal(got, []string{"203.0.113.7:8080"}) {
+				t.Fatalf("dialled %v, want the pinned address only, through the pinned dialer", got)
+			}
+		})
+	}
+}

intercept/export_test.go

@@ -0,0 +1,33 @@
+package intercept
+
+import (
+	"context"
+	"testing"
+)
+
+// DialGate exposes the pinned-dial fixture of target_dial_test.go to the tests
+// of package intercept_test, which may import the packages that import
+// intercept (enrich).
+type DialGate struct {
+	g   *gateFixture
+	rec *dialRecorder
+}
+
+func NewDialGate(t *testing.T) *DialGate {
+	g, rec := dialGate(t)
+	return &DialGate{g: g, rec: rec}
+}
+
+// Ctx returns a context carrying the fixture's target scope, gate on.
+func (d *DialGate) Ctx() context.Context { return d.g.ctx() }
+
+func (d *DialGate) Pin(host string, port int, ips ...string) { d.g.pin(host, port, ips...) }
+
+func (d *DialGate) Resolve(host string, ips ...string) { d.g.res.set(host, ips...) }
+
+func (d *DialGate) Exempt(urls ...string) {
+	d.g.ic.SetDerivedExemptions(func() []string { return urls })
+}
+
+// Calls returns every address the base dialer was asked to connect to.
+func (d *DialGate) Calls() []string { return d.rec.calls() }