Improve code organization and naming consistency

audit/logging_auditor.go

@@ -2,7 +2,7 @@ package audit
 
 import "log/slog"
 
-// LoggingAuditor implements Auditor by logging to slog
+// LoggingAuditor implements proxy.Auditor by logging to slog
 type LoggingAuditor struct {
 	logger *slog.Logger
 }

audit/request.go

@@ -1,6 +1,8 @@
 package audit
 
-import "net/http"
+type Auditor interface {
+	AuditRequest(req Request)
+}
 
 // Request represents information about an HTTP request for auditing
 type Request struct {
@@ -9,11 +11,3 @@ type Request struct {
 	Allowed bool
 	Rule    string // The rule that matched (if any)
 }
-
-// HTTPRequestToAuditRequest converts an http.Request to an audit.Request
-func HTTPRequestToAuditRequest(httpReq *http.Request) *Request {
-	return &Request{
-		Method: httpReq.Method,
-		URL:    httpReq.URL.String(),
-	}
-}

cli/cli.go

@@ -59,7 +59,7 @@ Examples:
 			},
 		},
 		Handler: func(inv *serpent.Invocation) error {
-			return Run(config, inv.Args)
+			return Run(inv.Context(), config, inv.Args)
 		},
 	}
 }
@@ -89,7 +89,9 @@ func setupLogging(logLevel string) *slog.Logger {
 }
 
 // Run executes the jail command with the given configuration and arguments
-func Run(config Config, args []string) error {
+func Run(ctx context.Context, config Config, args []string) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
 	logger := setupLogging(config.LogLevel)
 
 	// Get command arguments
@@ -113,7 +115,7 @@ func Run(config Config, args []string) error {
 	ruleEngine := rules.NewRuleEngine(allowRules, logger)
 
 	// Create auditor
-	// auditor := audit.NewLoggingAuditor(logger)
+	auditor := audit.NewLoggingAuditor(logger)
 
 	// Create certificate manager
 	certManager, err := tls.NewCertificateManager(logger)
@@ -123,53 +125,32 @@ func Run(config Config, args []string) error {
 	}
 
 	// Create jail instance
-	jailInstance, err := jail.New(context.Background(), jail.Config{
+	jailInstance, err := jail.New(ctx, jail.Config{
 		RuleEngine:  ruleEngine,
-		Auditor:     audit.NewLoggingAuditor(logger),
-		Logger:      logger,
+		Auditor:     auditor,
 		CertManager: certManager,
+		Logger:      logger,
 	})
 	if err != nil {
-		logger.Error("Failed to create jail instance", "error", err)
 		return fmt.Errorf("failed to create jail instance: %v", err)
 	}
 
 	// Setup signal handling BEFORE any setup
 	sigChan := make(chan os.Signal, 1)
 	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
 
-	// Handle signals immediately in background
-	go func() {
-		sig := <-sigChan
-		logger.Info("Received signal during setup, cleaning up...", "signal", sig)
-		err := jailInstance.Close()
-		if err != nil {
-			logger.Error("Emergency cleanup failed", "error", err)
-		}
-		os.Exit(1)
-	}()
-
-	// Ensure cleanup happens no matter what
-	defer func() {
-		logger.Debug("Starting cleanup process")
-		err := jailInstance.Close()
-		if err != nil {
-			logger.Error("Failed to cleanup jail", "error", err)
-		} else {
-			logger.Debug("Cleanup completed successfully")
-		}
-	}()
-
 	// Open jail (starts network namespace and proxy server)
 	err = jailInstance.Start()
 	if err != nil {
-		logger.Error("Failed to open jail", "error", err)
 		return fmt.Errorf("failed to open jail: %v", err)
 	}
-
-	// Create context for graceful shutdown
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
+	defer func() {
+		logger.Info("Closing jail...")
+		err := jailInstance.Close()
+		if err != nil {
+			logger.Error("Failed to close jail", "error", err)
+		}
+	}()
 
 	// Execute command in jail
 	go func() {

jail.go

@@ -2,37 +2,29 @@ package jail
 
 import (
 	"context"
-	cryptotls "crypto/tls"
 	"fmt"
 	"log/slog"
 	"os/exec"
 	"runtime"
 	"time"
 
+	"github.com/coder/jail/audit"
 	"github.com/coder/jail/namespace"
 	"github.com/coder/jail/proxy"
+	"github.com/coder/jail/rules"
+	"github.com/coder/jail/tls"
 )
 
-type Commander interface {
-	Start() error
-	Command(command []string) *exec.Cmd
-	Close() error
-}
-
-type CertificateManager interface {
-	SetupTLSAndWriteCACert() (*cryptotls.Config, string, string, error)
-}
-
 type Config struct {
-	RuleEngine  proxy.RuleEvaluator
-	Auditor     proxy.Auditor
-	CertManager CertificateManager
+	RuleEngine  rules.Evaluator
+	Auditor     audit.Auditor
+	CertManager tls.Manager
 	Logger      *slog.Logger
 }
 
 type Jail struct {
-	commander   Commander
-	proxyServer *proxy.ProxyServer
+	commander   namespace.Commander
+	proxyServer *proxy.Server
 	logger      *slog.Logger
 	ctx         context.Context
 	cancel      context.CancelFunc
@@ -56,7 +48,7 @@ func New(ctx context.Context, config Config) (*Jail, error) {
 	})
 
 	// Create commander
-	commander, err := newCommander(namespace.Config{
+	commander, err := newNamespaceCommander(namespace.Config{
 		Logger:         config.Logger,
 		HttpProxyPort:  8080,
 		HttpsProxyPort: 8443,
@@ -125,8 +117,8 @@ func (j *Jail) Close() error {
 	return j.commander.Close()
 }
 
-// newCommander creates a new NetJail instance for the current platform
-func newCommander(config namespace.Config) (Commander, error) {
+// newNamespaceCommander creates a new namespace instance for the current platform
+func newNamespaceCommander(config namespace.Config) (namespace.Commander, error) {
 	switch runtime.GOOS {
 	case "darwin":
 		return namespace.NewMacOS(config)

namespace/namespace.go

@@ -3,13 +3,20 @@ package namespace
 import (
 	"fmt"
 	"log/slog"
+	"os/exec"
 	"time"
 )
 
 const (
 	namespacePrefix = "coder_jail"
 )
 
+type Commander interface {
+	Start() error
+	Command(command []string) *exec.Cmd
+	Close() error
+}
+
 // JailConfig holds configuration for network jail
 type Config struct {
 	Logger         *slog.Logger
@@ -18,18 +25,6 @@ type Config struct {
 	Env            map[string]string
 }
 
-// // NewJail creates a new NetJail instance for the current platform
-// func New(config Config) (jail.Commander, error) {
-// 	switch runtime.GOOS {
-// 	case "darwin":
-// 		return NewMacOS(config)
-// 	case "linux":
-// 		return NewLinux(config)
-// 	default:
-// 		return nil, fmt.Errorf("unsupported platform: %s", runtime.GOOS)
-// 	}
-// }
-
 func newNamespaceName() string {
 	return fmt.Sprintf("%s_%d", namespacePrefix, time.Now().UnixNano()%10000000)
 }

namespace/noop.go

@@ -6,10 +6,6 @@ import (
 
 type noop struct{}
 
-func newNoop(_ Config) (*noop, error) {
-	return &noop{}, nil
-}
-
 func (n *noop) Command(_ []string) *exec.Cmd {
 	return exec.Command("true")
 }