feat: kill active connections when allow rule is deleted or denied (#3)

Track active proxy connections by the rule ID that authorized them using
a new ConnTracker. When a rule is deleted or changed to deny, all
connections authorized by that rule are cancelled immediately by closing
both sides of the tunnel.

Key changes:
- ConnTracker maps rule IDs to cancellable connections
- Bypass plugin writes authorizing rule ID into context
- SOCKS5 and HTTP handlers register connections with ConnTracker
- Pipe actively closes both connections on context cancellation
- API and HTMX handlers call CancelByRule on rule delete/deny

Author: tito

cmd/greyproxy/program.go

@@ -262,6 +262,7 @@ func (p *program) buildGreyproxyService() error {
 	shared.Cache = tmpSvc.Cache
 	shared.Bus = tmpSvc.Bus
 	shared.Waiters = tmpSvc.Waiters
+	shared.ConnTracker = greyproxy.NewConnTracker()
 	shared.Version = version
 
 	// Collect listening ports for the health endpoint
@@ -288,7 +289,7 @@ func (p *program) buildGreyproxyService() error {
 	// Create and register gost plugins
 	autherPlugin := greyproxy_plugins.NewAuther()
 	admissionPlugin := greyproxy_plugins.NewAdmission()
-	bypassPlugin := greyproxy_plugins.NewBypass(shared.DB, shared.Cache, shared.Bus, shared.Waiters)
+	bypassPlugin := greyproxy_plugins.NewBypass(shared.DB, shared.Cache, shared.Bus, shared.Waiters, shared.ConnTracker)
 	resolverPlugin := greyproxy_plugins.NewResolver(shared.Cache)
 
 	registry.AutherRegistry().Register(gaCfg.Auther, autherPlugin)
@@ -302,7 +303,7 @@ func (p *program) buildGreyproxyService() error {
 	// Build HTTP router with REST API + HTMX UI + WebSocket
 	router, g := greyproxy_api.NewRouter(shared, gaCfg.PathPrefix)
 	greyproxy_ui.RegisterPageRoutes(g, shared.DB, shared.Bus)
-	greyproxy_ui.RegisterHTMXRoutes(g, shared.DB, shared.Bus, shared.Waiters)
+	greyproxy_ui.RegisterHTMXRoutes(g, shared.DB, shared.Bus, shared.Waiters, shared.ConnTracker)
 
 	// Create the actual service
 	svc := &greyproxy.Service{}

internal/gostx/ctx/value.go

@@ -90,3 +90,30 @@ func ClientIDFromContext(ctx context.Context) ClientID {
 	v, _ := ctx.Value(clientIDKey{}).(ClientID)
 	return v
 }
+
+// ConnCanceller allows registering connection cancel functions by rule ID.
+// Implemented by greyproxy.ConnTracker.
+type ConnCanceller interface {
+	Register(ruleID int64, cancel context.CancelFunc) uint64
+	Unregister(ruleID int64, id uint64)
+}
+
+// BypassResult is a mutable container placed in the context before calling
+// bypass.Contains. The bypass plugin fills in the RuleID and Tracker when
+// a connection is allowed by a rule, so the handler can register the
+// connection for cancellation if the rule is later deleted.
+type BypassResult struct {
+	RuleID  int64
+	Tracker ConnCanceller
+}
+
+type bypassResultKey struct{}
+
+func ContextWithBypassResult(ctx context.Context, result *BypassResult) context.Context {
+	return context.WithValue(ctx, bypassResultKey{}, result)
+}
+
+func BypassResultFromContext(ctx context.Context) *BypassResult {
+	v, _ := ctx.Value(bypassResultKey{}).(*BypassResult)
+	return v
+}

internal/gostx/handler/http/handler.go

@@ -297,17 +297,32 @@ func (h *httpHandler) handleRequest(ctx context.Context, conn net.Conn, req *htt
 
 	ctx = xctx.ContextWithClientID(ctx, xctx.ClientID(clientID))
 
-	if h.options.Bypass != nil &&
-		h.options.Bypass.Contains(ctx, network, addr, bypass.WithService(h.options.Service)) {
-		resp.StatusCode = http.StatusForbidden
+	var bypassResult *xctx.BypassResult
+	if h.options.Bypass != nil {
+		bypassResult = &xctx.BypassResult{}
+		checkCtx := xctx.ContextWithBypassResult(ctx, bypassResult)
+		if h.options.Bypass.Contains(checkCtx, network, addr, bypass.WithService(h.options.Service)) {
+			resp.StatusCode = http.StatusForbidden
+
+			if log.IsLevelEnabled(logger.TraceLevel) {
+				dump, _ := httputil.DumpResponse(resp, false)
+				log.Trace(string(dump))
+			}
+			log.Debug("bypass: ", addr)
+			resp.Write(conn)
+			return xbypass.ErrBypass
+		}
 
-		if log.IsLevelEnabled(logger.TraceLevel) {
-			dump, _ := httputil.DumpResponse(resp, false)
-			log.Trace(string(dump))
+		// Register this connection for cancellation if the rule is later deleted.
+		if bypassResult.RuleID != 0 && bypassResult.Tracker != nil {
+			var pipeCancel context.CancelFunc
+			ctx, pipeCancel = context.WithCancel(ctx)
+			connID := bypassResult.Tracker.Register(bypassResult.RuleID, pipeCancel)
+			defer func() {
+				bypassResult.Tracker.Unregister(bypassResult.RuleID, connID)
+				pipeCancel()
+			}()
 		}
-		log.Debug("bypass: ", addr)
-		resp.Write(conn)
-		return xbypass.ErrBypass
 	}
 
 	if network == "udp" {

internal/gostx/handler/socks/v5/connect.go

@@ -56,7 +56,9 @@ func (h *socks5Handler) handleConnect(ctx context.Context, conn net.Conn, networ
 	}
 
 	if h.options.Bypass != nil {
-		bypassCtx, bypassCancel := context.WithCancel(ctx)
+		bypassResult := &xctx.BypassResult{}
+		resultCtx := xctx.ContextWithBypassResult(ctx, bypassResult)
+		bypassCtx, bypassCancel := context.WithCancel(resultCtx)
 
 		// Monitor the client TCP connection for close during the bypass check.
 		// During SOCKS5 CONNECT, the client waits for the server reply before
@@ -96,6 +98,17 @@ func (h *socks5Handler) handleConnect(ctx context.Context, conn net.Conn, networ
 			log.Debug("bypass: ", address)
 			return resp.Write(conn)
 		}
+
+		// Register this connection for cancellation if the rule is later deleted.
+		if bypassResult.RuleID != 0 && bypassResult.Tracker != nil {
+			var pipeCancel context.CancelFunc
+			ctx, pipeCancel = context.WithCancel(ctx)
+			connID := bypassResult.Tracker.Register(bypassResult.RuleID, pipeCancel)
+			defer func() {
+				bypassResult.Tracker.Unregister(bypassResult.RuleID, connID)
+				pipeCancel()
+			}()
+		}
 	}
 
 	switch h.md.hash {

internal/gostx/internal/net/pipe.go

@@ -43,7 +43,11 @@ func Pipe(ctx context.Context, rw1, rw2 io.ReadWriteCloser) error {
 	select {
 	case <-done:
 	case <-ctx.Done():
-		return nil
+		// Close both sides to unblock the pipeBuffer goroutines.
+		rw1.Close()
+		rw2.Close()
+		<-done
+		return ctx.Err()
 	}
 
 	select {

internal/greyproxy/api/router.go

@@ -10,12 +10,13 @@ import (
 
 // Shared holds shared state passed to all handlers.
 type Shared struct {
-	DB      *greyproxy.DB
-	Cache   *greyproxy.DNSCache
-	Bus     *greyproxy.EventBus
-	Waiters *greyproxy.WaiterTracker
-	Version string
-	Ports   map[string]int
+	DB          *greyproxy.DB
+	Cache       *greyproxy.DNSCache
+	Bus         *greyproxy.EventBus
+	Waiters     *greyproxy.WaiterTracker
+	ConnTracker *greyproxy.ConnTracker
+	Version     string
+	Ports       map[string]int
 }
 
 // NewRouter creates the Gin router with all routes.

internal/greyproxy/api/rules.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"log/slog"
 	"net/http"
 	"strconv"
 
@@ -81,6 +82,11 @@ func RulesUpdateHandler(s *Shared) gin.HandlerFunc {
 			return
 		}
 
+		// If the rule was changed to deny, cancel connections that relied on it.
+		if rule.Action == "deny" && s.ConnTracker != nil {
+			s.ConnTracker.CancelByRule(id)
+		}
+
 		c.JSON(http.StatusOK, rule.ToJSON())
 	}
 }
@@ -111,12 +117,19 @@ func RulesDeleteHandler(s *Shared) gin.HandlerFunc {
 			return
 		}
 
+		slog.Info("api: deleting rule", "rule_id", id)
+
 		deleted, err := greyproxy.DeleteRule(s.DB, id)
 		if err != nil {
 			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
 			return
 		}
 
+		slog.Info("api: rule deleted, cancelling connections", "rule_id", id, "deleted", deleted)
+		if deleted && s.ConnTracker != nil {
+			s.ConnTracker.CancelByRule(id)
+		}
+
 		c.JSON(http.StatusOK, gin.H{"status": "ok", "deleted": deleted})
 	}
 }

internal/greyproxy/conn_tracker.go

@@ -0,0 +1,76 @@
+package greyproxy
+
+import (
+	"context"
+	"log/slog"
+	"sync"
+	"sync/atomic"
+)
+
+var nextConnID atomic.Uint64
+
+// ConnTracker tracks active proxy connections by the rule ID that authorized
+// them. When a rule is deleted or changed to deny, all connections that were
+// allowed by that rule can be cancelled immediately.
+type ConnTracker struct {
+	mu    sync.Mutex
+	conns map[int64]map[uint64]context.CancelFunc
+}
+
+func NewConnTracker() *ConnTracker {
+	return &ConnTracker{
+		conns: make(map[int64]map[uint64]context.CancelFunc),
+	}
+}
+
+// Register associates a cancel function with a rule ID and returns an ID
+// that can be used to unregister later.
+func (ct *ConnTracker) Register(ruleID int64, cancel context.CancelFunc) uint64 {
+	id := nextConnID.Add(1)
+
+	ct.mu.Lock()
+	defer ct.mu.Unlock()
+
+	if ct.conns[ruleID] == nil {
+		ct.conns[ruleID] = make(map[uint64]context.CancelFunc)
+	}
+	ct.conns[ruleID][id] = cancel
+
+	slog.Info("conn_tracker: registered", "conn_id", id, "rule_id", ruleID, "total_for_rule", len(ct.conns[ruleID]))
+	return id
+}
+
+// Unregister removes a previously registered connection.
+// Called when a connection ends naturally.
+func (ct *ConnTracker) Unregister(ruleID int64, id uint64) {
+	ct.mu.Lock()
+	defer ct.mu.Unlock()
+
+	if m, ok := ct.conns[ruleID]; ok {
+		delete(m, id)
+		if len(m) == 0 {
+			delete(ct.conns, ruleID)
+		}
+		slog.Info("conn_tracker: unregistered", "conn_id", id, "rule_id", ruleID)
+	}
+}
+
+// CancelByRule cancels all active connections that were authorized by the
+// given rule ID and removes them from tracking.
+func (ct *ConnTracker) CancelByRule(ruleID int64) {
+	ct.mu.Lock()
+	cancels := ct.conns[ruleID]
+	delete(ct.conns, ruleID)
+	ct.mu.Unlock()
+
+	if len(cancels) == 0 {
+		slog.Info("conn_tracker: cancel by rule, no active connections", "rule_id", ruleID)
+		return
+	}
+
+	slog.Info("conn_tracker: cancel by rule, killing connections", "rule_id", ruleID, "count", len(cancels))
+	for id, cancel := range cancels {
+		slog.Info("conn_tracker: cancelling conn", "conn_id", id, "rule_id", ruleID)
+		cancel()
+	}
+}