reexamine socket and start/stop handling

Author: brianm

pkg/agent/agent.go

@@ -22,7 +22,6 @@ var errAgentStopped = errors.New("agent has been stopped")
 type Agent struct {
 	keyring  agent.Agent
 	caClient *caclient.Client
-	ctx      context.Context
 
 	agentSocketPath string
 	agentListener   net.Listener
@@ -35,40 +34,39 @@ type Agent struct {
 	closeOnce sync.Once
 }
 
-// Start creates and starts an SSH agent that listens on the specified socket path.
-// If agentSocketPath is empty, a temporary socket will be created.
-// The agent will automatically close when the context is cancelled.
-func Start(ctx context.Context, caClient *caclient.Client, agentSocketPath string) (*Agent, error) {
-	keyring := agent.NewKeyring()
-	a := &Agent{
-		ctx:             ctx,
-		agentSocketPath: agentSocketPath,
-		keyring:         keyring,
-		caClient:        caClient,
-		done:            make(chan struct{}),
-	}
-
+// New creates a new SSH agent. This does not start listening - call Serve() to begin accepting connections.
+// If agentSocketPath is empty, a temporary socket will be created when Serve() is called.
+func New(caClient *caclient.Client, agentSocketPath string) (*Agent, error) {
 	pub, priv, err := sshcert.GenerateKeys()
 	if err != nil {
 		return nil, err
 	}
-	a.privateKey = priv
-	a.publicKey = pub
 
-	if a.ctx == nil {
-		a.ctx = context.Background()
-	}
+	return &Agent{
+		agentSocketPath: agentSocketPath,
+		keyring:         agent.NewKeyring(),
+		caClient:        caClient,
+		publicKey:       pub,
+		privateKey:      priv,
+		done:            make(chan struct{}),
+	}, nil
+}
 
-	err = a.startAgentListener()
-	if err != nil {
-		return nil, err
+// Serve starts the agent listening on the configured socket and blocks until the context is cancelled.
+// Returns an error if the listener cannot be started, otherwise returns ctx.Err() when shutdown completes.
+func (a *Agent) Serve(ctx context.Context) error {
+	if err := a.startAgentListener(); err != nil {
+		return err
 	}
 
-	context.AfterFunc(ctx, func() {
-		a.Close()
-	})
+	// Serve connections in background
+	go a.serve(ctx)
+
+	// Block until context cancelled
+	<-ctx.Done()
+	a.Close()
 
-	return a, nil
+	return ctx.Err()
 }
 
 // Credential contains the private key and certificate in PEM format
@@ -143,23 +141,35 @@ func (a *Agent) startAgentListener() error {
 		return fmt.Errorf("unable to set permissions on agent socket: %w", err)
 	}
 	a.agentListener = agentListener
-	go a.listenAndServeAgent()
 	return nil
 }
 
-func (a *Agent) listenAndServeAgent() {
-	for a.Running() {
+func (a *Agent) serve(ctx context.Context) {
+	for {
+		// Check if context is done
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
+
 		conn, err := a.agentListener.Accept()
 		if err != nil {
 			if conn != nil {
 				conn.Close()
 			}
-			if !a.Running() {
-				// Agent is shutting down
+			// Check if error is from listener being closed
+			if errors.Is(err, net.ErrClosed) {
+				return
+			}
+			// Check context again before logging
+			select {
+			case <-ctx.Done():
 				return
+			default:
+				log.Warnf("error on accept from SSH_AUTH_SOCK listener: %v", err)
+				continue
 			}
-			log.Warnf("error on accept from SSH_AUTH_SOCK listener: %v", err)
-			continue
 		}
 		go a.serveAgent(conn)
 	}

pkg/agent/agent_test.go

@@ -30,10 +30,22 @@ func TestBasics(t *testing.T) {
 	require.NoError(t, err)
 
 	ctx, cancel := context.WithCancel(t.Context())
+	defer cancel()
 
-	a, err := agent.Start(ctx, nil, "")
+	a, err := agent.New(nil, "")
 	require.NoError(t, err)
 
+	// Serve in background
+	go func() {
+		err := a.Serve(ctx)
+		if err != nil && !errors.Is(err, context.Canceled) {
+			t.Errorf("agent.Serve error: %v", err)
+		}
+	}()
+
+	// Give agent time to start listening
+	time.Sleep(10 * time.Millisecond)
+
 	server, err := sshd.Start(caPub)
 	require.NoError(t, err)
 	defer server.Close()

pkg/broker/broker.go

@@ -40,23 +40,31 @@ type MatchResponse struct {
 	Error string
 }
 
-func New(ctx context.Context, log slog.Logger, socketPath string, authCommand string) (*Broker, error) {
-	b := Broker{
+// New creates a new Broker instance. This does not start listening - call Serve() to begin accepting connections.
+func New(log slog.Logger, socketPath string, authCommand string) *Broker {
+	return &Broker{
 		auth:             NewAuth(authCommand),
 		brokerSocketPath: socketPath,
 		done:             make(chan struct{}),
 		log:              log,
 	}
+}
 
-	err := b.startBrokerListener()
-	if err != nil {
-		return nil, fmt.Errorf("Unable to start broker socket: %w", err)
+// Serve starts the broker listening on the configured socket and blocks until the context is cancelled.
+// Returns an error if the listener cannot be started, otherwise returns ctx.Err() when shutdown completes.
+func (b *Broker) Serve(ctx context.Context) error {
+	if err := b.startBrokerListener(); err != nil {
+		return fmt.Errorf("unable to start broker socket: %w", err)
 	}
 
-	context.AfterFunc(ctx, func() {
-		b.Close()
-	})
-	return &b, nil
+	// Serve connections in background
+	go b.serve(ctx)
+
+	// Block until context cancelled
+	<-ctx.Done()
+	b.Close()
+
+	return ctx.Err()
 }
 
 func (b *Broker) startBrokerListener() error {
@@ -67,7 +75,6 @@ func (b *Broker) startBrokerListener() error {
 	}
 
 	b.brokerListener = brokerListener
-	go b.listenAndServe()
 	return nil
 }
 
@@ -77,26 +84,37 @@ func (b *Broker) Match(input MatchRequest, output *MatchResponse) error {
 	return nil
 }
 
-func (b *Broker) listenAndServe() {
+func (b *Broker) serve(ctx context.Context) {
 	server := rpc.NewServer()
 	server.Register(b)
 	for {
+		// Check if context is done
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
+
 		conn, err := b.brokerListener.Accept()
 		if err != nil {
 			// Check if error is from listener being closed
 			if errors.Is(err, net.ErrClosed) {
 				// Listener closed, exit gracefully
 				return
 			}
-			// Log other errors only if still running
-			if b.Running() {
+			// Check context again before logging
+			select {
+			case <-ctx.Done():
+				return
+			default:
 				b.log.Warn("Unable to accept connection", "error", err)
 				continue
 			}
-			return
 		}
-		defer conn.Close()
-		go server.ServeConn(conn)
+		go func() {
+			defer conn.Close()
+			server.ServeConn(conn)
+		}()
 	}
 }
 

pkg/broker/broker_test.go

@@ -4,6 +4,7 @@ import (
 	"log/slog"
 	"net/rpc"
 	"testing"
+	"time"
 
 	"github.com/lmittmann/tint"
 	"github.com/stretchr/testify/require"
@@ -13,14 +14,21 @@ func Test_RpcBasics(t *testing.T) {
 	ctx := t.Context()
 	authCommand := "echo '6:thello,'"
 	socketPath := t.TempDir() + "/broker.sock"
-	b, err := New(
-		ctx,
-		*testLogger(t),
-		socketPath,
-		authCommand)
-	require.NoError(t, err)
+
+	b := New(*testLogger(t), socketPath, authCommand)
+
+	// Serve in background
+	go func() {
+		err := b.Serve(ctx)
+		if err != nil && err != ctx.Err() {
+			t.Errorf("broker.Serve error: %v", err)
+		}
+	}()
 	defer b.Close()
 
+	// Give broker time to start listening
+	time.Sleep(10 * time.Millisecond)
+
 	client, err := rpc.Dial("unix", socketPath)
 	require.NoError(t, err)
 
@@ -35,14 +43,21 @@ func Test_MatchRequestFields(t *testing.T) {
 	ctx := t.Context()
 	authCommand := "echo '6:thello,'"
 	socketPath := t.TempDir() + "/broker.sock"
-	b, err := New(
-		ctx,
-		*testLogger(t),
-		socketPath,
-		authCommand)
-	require.NoError(t, err)
+
+	b := New(*testLogger(t), socketPath, authCommand)
+
+	// Serve in background
+	go func() {
+		err := b.Serve(ctx)
+		if err != nil && err != ctx.Err() {
+			t.Errorf("broker.Serve error: %v", err)
+		}
+	}()
 	defer b.Close()
 
+	// Give broker time to start listening
+	time.Sleep(10 * time.Millisecond)
+
 	client, err := rpc.Dial("unix", socketPath)
 	require.NoError(t, err)