diff --git a/Makefile b/Makefile index fae81f3..778fafe 100644 --- a/Makefile +++ b/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}' diff --git a/bin/ssh-proxy b/bin/ssh-proxy index 92d6ba1..f2eeedd 100755 Binary files a/bin/ssh-proxy and b/bin/ssh-proxy differ diff --git a/go.mod b/go.mod index e8c76a0..6653d59 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 182c816..47f3cff 100644 --- a/go.sum +++ b/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= diff --git a/internal/server/auth.go b/internal/server/auth.go index feb8e37..925cdaa 100644 --- a/internal/server/auth.go +++ b/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) diff --git a/internal/server/connection.go b/internal/server/connection.go index 85d09c7..3e4258c 100644 --- a/internal/server/connection.go +++ b/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) diff --git a/internal/server/ssh.go b/internal/server/ssh.go index ac2fa4e..718a48c 100644 --- a/internal/server/ssh.go +++ b/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) diff --git a/internal/server/userlock.go b/internal/server/userlock.go new file mode 100644 index 0000000..e1a0880 --- /dev/null +++ b/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 +} diff --git a/internal/workspace/provision.go b/internal/workspace/provision.go index 6b2643e..b6dac4e 100644 --- a/internal/workspace/provision.go +++ b/internal/workspace/provision.go @@ -19,6 +19,7 @@ import ( "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 @@ func EnsureWorkspace(ctx context.Context, userStr *userstr.UserStr, writer *Info 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,10 +293,16 @@ func EnsureWorkspace(ctx context.Context, userStr *userstr.UserStr, writer *Info 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) { @@ -294,12 +310,37 @@ func EnsureWorkspace(ctx context.Context, userStr *userstr.UserStr, writer *Info } 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 @@ func provisionWorkspace(ctx context.Context, userStr *userstr.UserStr, writer *I 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 {