user management

Makefile

@@ -5,7 +5,7 @@ REPORTS_DIR := reports
 SERVICE_NAME ?= $(shell basename $(CURDIR))
 RUNTIME ?= alpine
 
-.PHONY: all init install-test-deps test-static test build test-binary test-self vendor image image-debug image-release reload dlv coverage clean help
+.PHONY: all init install-test-deps test-static test build test-binary test-self vendor image image-debug image-release reload dlv debug-setup coverage clean help
 
 # Default target
 all: build
@@ -96,14 +96,20 @@ image: vendor
 		-f docker/$(SERVICE_NAME)/Dockerfile \
 		.
 
-reload: build ##@ Hot-swap the running $(SERVICE_NAME) binary in place (dev only)
-              ##@ Finds the container's while-loop entrypoint via /proc, replaces /app/$(SERVICE_NAME), and kills the child so the loop respawns it
+reload: build ##@ Hot-swap the running provisioner binary in place (dev only)
+              ##@ Finds the container's while-loop entrypoint via /proc, replaces /app/provisioner, and kills the child so the loop respawns it
 	@./scripts/reload.sh
 
-dlv: ##@ Attach delve to the running $(SERVICE_NAME) process for remote debugging
+dlv: ##@ Attach delve to the running provisioner process for remote debugging
      ##@ Headless debug server on 127.0.0.1:2345 (override with DLV_LISTEN) — connect with `dlv connect` or your IDE
 	@./scripts/dlv.sh
 
+debug-setup: ##@ Set up local debug environment
+             ##@ Generates go.work (Go version taken from go.mod) and symlinks the common module for local debugging
+	@GO_VERSION=$$(grep -m1 '^go ' go.mod | awk '{print $$2}') && \
+	printf 'go %s\n\nuse (\n\t.\n\t/opt/shared/common\n)\n' "$$GO_VERSION" > go.work
+	ln -sfn /opt/shared/common common
+
 coverage:  ##@ Calculate test coverage percentage from coverage.out
 	@go tool cover -func=$(REPORTS_DIR)/coverage.out | grep total | awk '{print $$3}'
 

go.mod

@@ -3,7 +3,7 @@ module github.com/k8shell-io/ssh-proxy
 go 1.24.5
 
 require (
-	github.com/k8shell-io/common v0.31.0
+	github.com/k8shell-io/common v0.34.3
 	github.com/nats-io/nats.go v1.47.0
 	github.com/rs/zerolog v1.34.0
 	golang.org/x/crypto v0.43.0
@@ -20,7 +20,6 @@ require (
 	github.com/go-playground/universal-translator v0.18.1 // indirect
 	github.com/go-playground/validator/v10 v10.27.0 // indirect
 	github.com/golang-jwt/jwt/v5 v5.3.1 // indirect
-	github.com/golang/protobuf v1.5.4 // indirect
 	github.com/klauspost/compress v1.18.0 // indirect
 	github.com/kr/pretty v0.3.1 // indirect
 	github.com/leodido/go-urn v1.4.0 // indirect

go.sum

@@ -31,10 +31,8 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
-github.com/k8shell-io/common v0.30.10 h1:sIL7pjx38YE/KtRB2VXK+BHL7D+NcM88Lzop/xJfBT8=
-github.com/k8shell-io/common v0.30.10/go.mod h1:40c5GkpS7Y0/aOFa37Lq8z/mLUn3k3GV/AHtFJFL28k=
-github.com/k8shell-io/common v0.31.0 h1:iFo6IUdNQCU2cKnUPd2u2vJ14UKR+rodTPj9uax6Ay8=
-github.com/k8shell-io/common v0.31.0/go.mod h1:40c5GkpS7Y0/aOFa37Lq8z/mLUn3k3GV/AHtFJFL28k=
+github.com/k8shell-io/common v0.34.3 h1:wTwCHjPSH4T8Mj4GIwXlekXVQwyt0jw5a7ZE4Vw5KtY=
+github.com/k8shell-io/common v0.34.3/go.mod h1:40c5GkpS7Y0/aOFa37Lq8z/mLUn3k3GV/AHtFJFL28k=
 github.com/k8shell-io/crypto v0.41.1-ssh-proxy h1:8+q6Ofc2ky23Oc9iNyiq8aeiQBIP+y3+O6zzHqe1f48=
 github.com/k8shell-io/crypto v0.41.1-ssh-proxy/go.mod h1:RVZeOJCpqtogniULztSXQESKJCfcI8WCxsS0FagMA8U=
 github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo=

internal/server/auth.go

@@ -351,6 +351,7 @@ func (s *Server) resolveAuthMethods(ctx context.Context, connInfo *Connection) (
 	req, err := authz.NewUserAuthEvalRequest(connInfo.user.Username).
 		WithIDP(connInfo.user.Source).
 		WithOrg(connInfo.user.Organization).
+		WithSurface(authz.AuthSurfaceSSH).
 		Build()
 	if err != nil {
 		return nil, fmt.Errorf("failed to build user:auth request: %w", err)

internal/server/connection.go

@@ -41,8 +41,10 @@ type Connection struct {
 	ctx              context.Context               // context for managing the connection
 	seqNumberGen     int64                         // sequence number for exec commands
 	connId           string                        // session key for the connection
+	connKey          string                        // key under which this Connection is stored in connStates
 	cancel           context.CancelFunc            // function to cancel the context
 	log              *zerolog.Logger               // logger instance, reused from server
+	transport        io.Closer                     // underlying SSH transport; set once the handshake completes, used to force-terminate the session
 	userStr          *userstr.UserStr              // user string information
 	clientIP         string                        // client IP address (detected from proxy protocol if available)
 	clientPort       int                           // client port (detected from proxy protocol if available)
@@ -162,7 +164,37 @@ func GetConnectionByAddress(remoteAddr string) *Connection {
 func RemoveState(state *Connection) {
 	connStatesMutex.Lock()
 	defer connStatesMutex.Unlock()
-	delete(connStates, fmt.Sprintf("connid-%s", state.userStr.Username()))
+	delete(connStates, state.connKey)
+}
+
+// closeAllSSHConnectionsForUser force-closes every live SSH connection owned
+// by username in this process. A user can hold more than one concurrent
+// session (multiple terminals, port-forwards, etc.), so every match is
+// closed, not just the first.
+func closeAllSSHConnectionsForUser(username string) {
+	connStatesMutex.RLock()
+	matches := make([]*Connection, 0, 1)
+	for _, c := range connStates {
+		if connectionBelongsToUser(c, username) {
+			matches = append(matches, c)
+		}
+	}
+	connStatesMutex.RUnlock()
+
+	for _, c := range matches {
+		c.log.Info().Msgf("Closing SSH connection for locked user %s", username)
+		c.ForceClose()
+	}
+}
+
+// connectionBelongsToUser reports whether c belongs to username, preferring
+// the identity-resolved username and falling back to the raw SSH login name
+// for connections that haven't completed authentication yet.
+func connectionBelongsToUser(c *Connection, username string) bool {
+	if c.user != nil {
+		return c.user.Username == username
+	}
+	return c.userStr != nil && c.userStr.Username() == username
 }
 
 // GetConnInfo retrieves or creates a Connection object for the given ssh.ConnMetadata
@@ -186,6 +218,7 @@ func (s *Server) GetConnInfo(conn ssh.ConnMetadata) (*Connection, error) {
 
 			connInfo = &Connection{
 				connId:        fmt.Sprintf("%s-%d-%s", GetProxyID(), os.Getpid(), strings.ToLower(rand.Text()[:2])),
+				connKey:       connID,
 				log:           s.log,
 				identity:      s.identity,
 				sessionClient: s.sessionClient,
@@ -256,10 +289,32 @@ func (c *Connection) Close() error {
 		c.reportWg.Wait()
 	}
 
+	RemoveState(c)
+
 	//c.cancel()
 	return nil
 }
 
+// SetTransport records the underlying SSH transport for the connection, so
+// it can be force-closed later (e.g. ForceClose) independently of the
+// per-connection context.
+func (c *Connection) SetTransport(transport io.Closer) {
+	c.transport = transport
+}
+
+// ForceClose terminates the connection's underlying SSH transport, e.g. when
+// the user's account becomes locked mid session. Closing the transport drops
+// every open SSH channel (unblocking any in-flight reads/writes with EOF)
+// and causes the channel loop to observe a closed channels chan, unwinding
+// through the normal Close/RemoveState path. It deliberately does not cancel
+// c.ctx: Close() reuses c.ctx to send the final session-delete update, and a
+// canceled context would make that call fail immediately.
+func (c *Connection) ForceClose() {
+	if c.transport != nil {
+		_ = c.transport.Close()
+	}
+}
+
 // reportSessionData periodically reports session data
 func (c *Connection) reportSessionData() {
 	t := time.NewTicker(SESSION_UPDATE_INTERVAL)

internal/server/ssh.go

@@ -133,8 +133,9 @@ func NewServer(configPath string) (*Server, error) {
 			if err != nil {
 				return nil, fmt.Errorf("create userstr jetstream cache: %w", err)
 			}
+			go server.startUserLockWatcher()
 		} else {
-			server.log.Warn().Msg("NATS client is not configured, userstr cache disabled")
+			server.log.Warn().Msg("NATS client is not configured, userstr cache and user lock watching disabled")
 		}
 
 		if config.Session.IsEnabled() {
@@ -257,8 +258,9 @@ func HandleConnectionChildProcess(configPath string) error {
 		if err != nil {
 			return fmt.Errorf("create userstr jetstream cache: %w", err)
 		}
+		go server.startUserLockWatcher()
 	} else {
-		logger.Warn().Msg("NATS client is not configured, userstr cache disabled")
+		logger.Warn().Msg("NATS client is not configured, userstr cache and user lock watching disabled")
 	}
 
 	if config.Session.IsEnabled() {
@@ -365,7 +367,6 @@ func (s *Server) handleConnection(netConn net.Conn, isDirect bool) {
 			}
 
 			connInfo.Close()
-			RemoveState(connInfo)
 		} else if s.fpub != nil {
 			pubErr := s.fpub.PublishFailure(ip, port, "", []string{err.Error()})
 			if pubErr != nil {
@@ -389,6 +390,7 @@ func (s *Server) handleConnection(netConn net.Conn, isDirect bool) {
 	}
 	connInfo.clientIP = ip
 	connInfo.clientPort = port
+	connInfo.SetTransport(sshConn)
 	defer connInfo.Close()
 
 	go s.handleGlobalRequests(requests)

internal/server/userlock.go

@@ -0,0 +1,84 @@
+// Copyright 2025 The K8shell Authors. All rights reserved.
+// Use of this source code is governed by a AGPLv3
+// license that can be found in the LICENSE file.
+
+package server
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"os"
+	"time"
+
+	natsc "github.com/k8shell-io/common/pkg/nats"
+)
+
+// userLockWatcherRetryInterval is how long to wait before retrying
+// subscription setup, e.g. when the locked-users bucket's underlying stream
+// doesn't exist yet because api-server hasn't created it.
+const userLockWatcherRetryInterval = 10 * time.Second
+
+// startUserLockWatcher subscribes to account lock/unlock events published by
+// api-server and force-closes every live SSH connection for a user as soon
+// as their account is locked. Unlocking is a no-op here: it does not
+// reconnect anyone, the user simply logs in again normally on their next
+// attempt. It blocks until s.ctx is canceled, so callers should run it in a
+// goroutine.
+//
+// The durable consumer name is unique per process, not just per replica:
+// in forking mode each connection is handled by its own child process, each
+// with its own local connection state, so every process that can hold live
+// connections must see every lock event. A durable name shared across
+// subscribers would instead load-balance events between them, so any single
+// subscriber would only observe some lock events instead of all of them.
+func (s *Server) startUserLockWatcher() {
+	if s.nats == nil {
+		return
+	}
+
+	durable := fmt.Sprintf("ssh-proxy-%s-%d", GetProxyID(), os.Getpid())
+
+	for {
+		if s.ctx.Err() != nil {
+			return
+		}
+
+		sub, err := s.nats.NewKVSubscriber(natsc.LOCKED_USERS_BUCKET, durable, natsc.KVSubscriberOptions{})
+		if err == nil {
+			err = sub.StartPull(s.ctx, s.handleUserLockEvent)
+		}
+
+		if s.ctx.Err() != nil {
+			return
+		}
+		if err != nil {
+			s.log.Warn().Msgf("locked-users subscription unavailable, retrying in %s: %v",
+				userLockWatcherRetryInterval, err)
+		}
+
+		select {
+		case <-s.ctx.Done():
+			return
+		case <-time.After(userLockWatcherRetryInterval):
+		}
+	}
+}
+
+// handleUserLockEvent processes a single locked-users KV event.
+func (s *Server) handleUserLockEvent(_ context.Context, ev *natsc.KVEvent) error {
+	// Deletes/purges of the key carry no value; there is nothing to act on.
+	if len(ev.Value) == 0 {
+		return nil
+	}
+
+	var state natsc.UserLockState
+	if err := json.Unmarshal(ev.Value, &state); err != nil {
+		return fmt.Errorf("failed to decode lock state for user %s: %w", ev.Key, err)
+	}
+
+	if state.Locked {
+		closeAllSSHConnectionsForUser(ev.Key)
+	}
+	return nil
+}

internal/workspace/provision.go

@@ -19,6 +19,7 @@
 	"github.com/k8shell-io/common/pkg/gapi"
 	"github.com/k8shell-io/common/pkg/models"
 	"github.com/k8shell-io/common/pkg/userstr"
+	"google.golang.org/grpc"
 )
 
 const PROVISION_TIMEOUT = 120 * time.Second
@@ -271,6 +272,15 @@
 				return gapi.ProtoToWorkspaceDetails(wsDetails), nil
 			}
 		}
+		for _, wsDetails := range ws.Workspaces {
+			if wsDetails.GetWorkspaceStatus().GetStatus() == string(models.WorkspaceStatusStopped) {
+				wsname, err := startWorkspace(ctx, wsDetails.GetName(), writer, backends)
+				if err != nil {
+					return nil, fmt.Errorf("%w: failed to start workspace", err)
+				}
+				return finalizeWorkspace(ctx, wsname, userStr.Username(), backends)
+			}
+		}
 		return nil, fmt.Errorf("%w: no running workspaces found. Please try again later.", ErrWorkspaceNotFound)
 	}
 
@@ -283,23 +293,54 @@
 		return nil, fmt.Errorf("%w: failed to provision workspace", err)
 	}
 
+	return finalizeWorkspace(ctx, wsname, userStr.Username(), backends)
+}
+
+// finalizeWorkspace looks up the workspace by name after it has been
+// provisioned or resumed and returns its details once confirmed running.
+func finalizeWorkspace(ctx context.Context, wsname string, username string, backends Backends) (*models.WorkspaceDetails, error) {
 	wsStatus, err := backends.Provisioner().FindWorkspace(ctx,
 		&provisionerv1.FindWorkspaceRequest{Workspace: wsname})
 	if err != nil {
-		return nil, fmt.Errorf("failed to get workspace status for user %s: %w", userStr.Username(), err)
+		return nil, fmt.Errorf("failed to get workspace status for user %s: %w", username, err)
 	}
 
 	if wsStatus.GetWorkspaceStatus().GetStatus() == string(models.WorkspaceStatusRunning) {
 		return gapi.ProtoToWorkspaceDetails(wsStatus), nil
 	}
 
 	return nil, fmt.Errorf("failed to ensure workspace for user %s: workspace status is %q",
-		userStr.Username(), wsStatus.GetWorkspaceStatus().GetStatus())
+		username, wsStatus.GetWorkspaceStatus().GetStatus())
 }
 
 // provisionWorkspace provisions a new workspace for the user.
 func provisionWorkspace(ctx context.Context, userStr *userstr.UserStr, writer *InfoWriter,
 	backends Backends) (string, error) {
+	timeout := int32(PROVISION_TIMEOUT / time.Second)
+	_, _, stream, err := backends.Provisioner().ProvisionHandshake(ctx, *userStr, timeout)
+	if err != nil {
+		return "", fmt.Errorf("handshake failed for user %s: %w", userStr.Username(), err)
+	}
+
+	return consumeProvisionStream(ctx, stream, writer)
+}
+
+// startWorkspace resumes a previously stopped workspace.
+func startWorkspace(ctx context.Context, workspaceName string, writer *InfoWriter,
+	backends Backends) (string, error) {
+	timeout := int32(PROVISION_TIMEOUT / time.Second)
+	_, _, stream, err := backends.Provisioner().StartHandshake(ctx, workspaceName, timeout)
+	if err != nil {
+		return "", fmt.Errorf("start handshake failed for workspace %s: %w", workspaceName, err)
+	}
+
+	return consumeProvisionStream(ctx, stream, writer)
+}
+
+// consumeProvisionStream drains a provisioning/start stream, reporting
+// progress via writer, until the workspace reaches a terminal state.
+func consumeProvisionStream(ctx context.Context,
+	stream grpc.ServerStreamingClient[provisionerv1.ProvisionWorkspaceResponse], writer *InfoWriter) (string, error) {
 	var (
 		name     string
 		eventErr error
@@ -310,12 +351,6 @@
 		writer.EndProvisioning(eventErr != nil)
 	}()
 
-	timeout := int32(PROVISION_TIMEOUT / time.Second)
-	_, _, stream, err := backends.Provisioner().ProvisionHandshake(ctx, *userStr, timeout)
-	if err != nil {
-		return "", fmt.Errorf("handshake failed for user %s: %w", userStr.Username(), err)
-	}
-
 loop:
 	for {
 		select {