20384304389

Committed 19 Dec 2025 10:34PM UTC coverage: 52.915% (-0.05%) from 52.963%

Build # 20384304389

Build Type

push

github

Committed by

web-flow

Commit Message

feat(ssh): custom sign in success hook (#6000)

## Summary

<!--  For example...
The existing implementation has poor numerical properties for
large arguments, so use the McGillicutty algorithm to improve
accuracy above 1e10.

The algorithm is described at
https://wikipedia.org/wiki/McGillicutty_Algorithm
-->

## Related issues

<!-- For example...
- #159
-->

## User Explanation

<!-- How would you explain this change to the user? If this
change doesn't create any user-facing changes, you can leave
this blank. If filled out, add the `docs` label -->

## Checklist

- [ ] reference any related issues
- [ ] updated unit tests
- [ ] add appropriate label (`enhancement`, `bug`, `breaking`,
`dependencies`, `ci`)
- [ ] ready for review

Run Details

12 of 65 new or added lines in 4 files covered. (18.46%)

20 existing lines in 7 files now uncovered.

29524 of 55795 relevant lines covered (52.92%)

127.25 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

86.97

/pkg/ssh/manager.go

package ssh

import (
        "cmp"
        "context"
        "errors"
        "fmt"
        "slices"
        "strconv"
        "sync"
        "time"

        corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
        envoy_config_endpoint_v3 "github.com/envoyproxy/go-control-plane/envoy/config/endpoint/v3"
        discoveryv3 "github.com/envoyproxy/go-control-plane/envoy/service/discovery/v3"
        endpointv3 "github.com/envoyproxy/go-control-plane/envoy/service/endpoint/v3"
        "github.com/envoyproxy/go-control-plane/pkg/cache/types"
        "github.com/envoyproxy/go-control-plane/pkg/cache/v3"
        "github.com/envoyproxy/go-control-plane/pkg/resource/v3"
        "github.com/envoyproxy/go-control-plane/pkg/server/delta/v3"
        "github.com/rs/zerolog"
        "golang.org/x/sync/errgroup"
        "google.golang.org/grpc/codes"
        "google.golang.org/grpc/status"
        "google.golang.org/protobuf/types/known/anypb"

        extensions_ssh "github.com/pomerium/envoy-custom/api/extensions/filters/network/ssh"
        "github.com/pomerium/pomerium/config"
        "github.com/pomerium/pomerium/internal/log"
        "github.com/pomerium/pomerium/pkg/grpc/databroker"
        "github.com/pomerium/pomerium/pkg/grpc/session"
        "github.com/pomerium/pomerium/pkg/ssh/portforward"
)

type streamClusterEndpointDiscovery struct {
        self     *StreamManager
        streamID uint64
}

type streamEndpointsUpdate struct {
        streamID uint64
        added    map[string]portforward.RoutePortForwardInfo
        removed  map[string]struct{}
}

// UpdateClusterEndpoints implements EndpointDiscoveryInterface.
func (ed *streamClusterEndpointDiscovery) UpdateClusterEndpoints(added map[string]portforward.RoutePortForwardInfo, removed map[string]struct{}) {
        // run this callback in a separate goroutine, since it can deadlock if called
        // synchronously during startup
        ed.self.endpointsUpdateQueue <- streamEndpointsUpdate{
                streamID: ed.streamID,
                added:    added,
                removed:  removed,
        }
}

// PortForwardManager implements EndpointDiscoveryInterface.
func (ed *streamClusterEndpointDiscovery) PortForwardManager() *portforward.Manager {
        return ed.self.getPortForwardManagerForStream(ed.streamID)
}

var _ EndpointDiscoveryInterface = (*streamClusterEndpointDiscovery)(nil)

var ErrReauthDone = errors.New("reauth loop done")

type activeStream struct {
        Handler            *StreamHandler
        PortForwardManager *portforward.Manager
        Session            *string
        SessionBindingID   *string
        Endpoints          map[string]struct{}
}

type StreamManager struct {
        endpointv3.UnimplementedEndpointDiscoveryServiceServer
        ready                            chan struct{}
        logger                           *zerolog.Logger
        auth                             AuthInterface
        reauthC                          chan struct{}
        initialSessionSyncDone           bool
        initialSessionBindingSyncDone    bool
        waitForInitialSessionSync        chan struct{}
        waitForInitialSessionBindingSync chan struct{}

        mu sync.Mutex

        cfg           *config.Config
        activeStreams map[uint64]*activeStream

        // Tracks stream IDs for active sessions
        sessionStreams map[string]map[uint64]struct{}

        // Tracks stream IDs per sessionBindingID for active sessions
        bindingStreams map[string]map[uint64]struct{}

        bindingSyncer *bindingSyncer
        // Tracks endpoint stream IDs for clusters
        clusterEndpoints     map[string]map[uint64]*extensions_ssh.EndpointMetadata
        endpointsUpdateQueue chan streamEndpointsUpdate
        edsCache             *cache.LinearCache
        edsServer            delta.Server
        indexer              PolicyIndexer
}

func (sm *StreamManager) getPortForwardManagerForStream(streamID uint64) *portforward.Manager {
        sm.mu.Lock()
        defer sm.mu.Unlock()
        activeStream := sm.activeStreams[streamID]
        if activeStream == nil || activeStream.PortForwardManager == nil {
                return nil
        }
        return activeStream.PortForwardManager // may be nil
}

// OnDeltaStreamClosed implements delta.Callbacks.
func (sm *StreamManager) OnDeltaStreamClosed(int64, *corev3.Node) {
}

// OnDeltaStreamOpen implements delta.Callbacks.
func (sm *StreamManager) OnDeltaStreamOpen(context.Context, int64, string) error {
        return nil
}

// OnStreamDeltaRequest implements delta.Callbacks.
func (sm *StreamManager) OnStreamDeltaRequest(_ int64, req *discoveryv3.DeltaDiscoveryRequest) error {
        if len(req.ResourceNamesSubscribe) == 0 {
                return nil
        }
        sm.mu.Lock()
        defer sm.mu.Unlock()
        initialEmptyResources := make(map[string]types.Resource)
        for _, clusterID := range req.ResourceNamesSubscribe {
                if _, ok := sm.clusterEndpoints[clusterID]; !ok {
                        sm.clusterEndpoints[clusterID] = map[uint64]*extensions_ssh.EndpointMetadata{}
                        initialEmptyResources[clusterID] = &envoy_config_endpoint_v3.ClusterLoadAssignment{
                                ClusterName: clusterID,
                        }
                }
        }
        return sm.edsCache.UpdateResources(initialEmptyResources, nil)
}

// OnStreamDeltaResponse implements delta.Callbacks.
func (sm *StreamManager) OnStreamDeltaResponse(int64, *discoveryv3.DeltaDiscoveryRequest, *discoveryv3.DeltaDiscoveryResponse) {
}

// DeltaEndpoints implements endpointv3.EndpointDiscoveryServiceServer.
func (sm *StreamManager) DeltaEndpoints(stream endpointv3.EndpointDiscoveryService_DeltaEndpointsServer) error {
        select {
        case <-stream.Context().Done():
                return context.Cause(stream.Context())
        case <-sm.ready:
        }
        log.Ctx(stream.Context()).Debug().Msg("delta endpoint stream started")
        defer log.Ctx(stream.Context()).Debug().Msg("delta endpoint stream ended")
        return sm.edsServer.DeltaStreamHandler(stream, resource.EndpointType)
}

// ClearRecords implements databroker.SyncerHandler.
func (sm *StreamManager) ClearRecords(ctx context.Context) {
        sm.mu.Lock()
        defer sm.mu.Unlock()
        if !sm.initialSessionSyncDone {
                sm.initialSessionSyncDone = true
                close(sm.waitForInitialSessionSync)
                log.Ctx(ctx).Debug().
                        Msg("ssh stream manager: initial sync done")
                return
        }
        for sessionID, streamIDs := range sm.sessionStreams {
                for streamID := range streamIDs {
                        log.Ctx(ctx).Debug().
                                Str("session-id", sessionID).
                                Uint64("stream-id", streamID).
                                Msg("terminating stream: databroker sync reset")
                        sm.terminateStreamLocked(streamID)
                }
        }
        clear(sm.sessionStreams)
}

func (sm *StreamManager) clearRecordsBinding(ctx context.Context) {
        sm.mu.Lock()
        defer sm.mu.Unlock()
        if !sm.initialSessionBindingSyncDone {
                sm.initialSessionBindingSyncDone = true
                close(sm.waitForInitialSessionBindingSync)
                log.Ctx(ctx).Debug().
                        Msg("ssh stream manager: initial sync done")
                return
        }
        for sessionID, streamIDs := range sm.bindingStreams {
                for streamID := range streamIDs {
                        log.Ctx(ctx).Debug().
                                Str("session-binding-id", sessionID).
                                Uint64("stream-id", streamID).
                                Msg("terminating stream: databroker sync reset")
                        sm.terminateStreamLocked(streamID)
                }
        }
        clear(sm.sessionStreams)
}

// GetDataBrokerServiceClient implements databroker.SyncerHandler.
func (sm *StreamManager) GetDataBrokerServiceClient() databroker.DataBrokerServiceClient {
        return sm.auth.GetDataBrokerServiceClient()
}

// UpdateRecords implements databroker.SyncerHandler.
func (sm *StreamManager) UpdateRecords(ctx context.Context, _ uint64, records []*databroker.Record) {
        sm.mu.Lock()
        defer sm.mu.Unlock()
        for _, record := range records {
                if record.DeletedAt == nil {
                        // New session
                        var s session.Session
                        if err := record.Data.UnmarshalTo(&s); err != nil {
                                log.Ctx(ctx).Err(err).Msg("invalid session object, ignoring")
                                continue
                        }
                        s.Version = strconv.FormatUint(record.GetVersion(), 10)
                        sm.indexer.OnSessionCreated(&s)
                        continue
                }
                // Session was deleted; terminate all of its associated streams
                sm.indexer.OnSessionDeleted(record.Id)
                for streamID := range sm.sessionStreams[record.Id] {
                        log.Ctx(ctx).Debug().
                                Str("session-id", record.Id).
                                Uint64("stream-id", streamID).
                                Msg("terminating stream: session revoked")
                        sm.terminateStreamLocked(streamID)
                }
                delete(sm.sessionStreams, record.Id)
        }
}

func (sm *StreamManager) updateRecordsBinding(ctx context.Context, _ uint64, records []*databroker.Record) {
        sm.mu.Lock()
        defer sm.mu.Unlock()
        for _, record := range records {
                if record.DeletedAt == nil {
                        continue
                }
                // Session binding was deleted; terminate all of its associated streams
                for streamID := range sm.bindingStreams[record.Id] {
                        log.Ctx(ctx).Debug().
                                Str("session-id", record.Id).
                                Uint64("stream-id", streamID).
                                Msg("terminating stream: session binding revoked")
                        sm.terminateStreamLocked(streamID)
                }
                delete(sm.bindingStreams, record.Id)
        }
}

func (sm *StreamManager) OnStreamAuthenticated(ctx context.Context, streamID uint64, req AuthRequest) error {
        if err := sm.waitForInitialSync(ctx); err != nil {
                return err
        }
        sm.mu.Lock()
        defer sm.mu.Unlock()
        activeStream := sm.activeStreams[streamID]
        if activeStream.Session != nil || activeStream.SessionBindingID != nil {
                return status.Errorf(codes.Internal, "stream %d already has an associated session", streamID)
        }

        if sm.sessionStreams[req.SessionID] == nil {
                sm.sessionStreams[req.SessionID] = map[uint64]struct{}{}
        }
        if sm.bindingStreams[req.SessionBindingID] == nil {
                sm.bindingStreams[req.SessionBindingID] = map[uint64]struct{}{}
        }
        sm.sessionStreams[req.SessionID][streamID] = struct{}{}
        sm.bindingStreams[req.SessionBindingID][streamID] = struct{}{}

        activeStream.Session = new(string)
        *activeStream.Session = req.SessionID
        activeStream.SessionBindingID = new(string)
        *activeStream.SessionBindingID = req.SessionBindingID

        activeStream.PortForwardManager.AddUpdateListener(activeStream.Handler)

        sm.indexer.OnStreamAuthenticated(streamID, req)
        return nil
}

type bindingSyncer struct {
        clientHandler func() databroker.DataBrokerServiceClient
        clearHandler  func(context.Context)
        updateHandler func(context.Context, uint64, []*databroker.Record)
}

var _ databroker.SyncerHandler = (*bindingSyncer)(nil)

func (sbr *bindingSyncer) ClearRecords(ctx context.Context) {
        sbr.clearHandler(ctx)
}

// GetDataBrokerServiceClient implements databroker.SyncerHandler.
func (sbr *bindingSyncer) GetDataBrokerServiceClient() databroker.DataBrokerServiceClient {
        return sbr.clientHandler()
}

// UpdateRecords implements databroker.SyncerHandler.
func (sbr *bindingSyncer) UpdateRecords(ctx context.Context, serverVersion uint64, records []*databroker.Record) {
        sbr.updateHandler(ctx, serverVersion, records)
}

func NewStreamManager(ctx context.Context, auth AuthInterface, indexer PolicyIndexer, cfg *config.Config) *StreamManager {
        sm := &StreamManager{
                logger:                           log.Ctx(ctx),
                auth:                             auth,
                ready:                            make(chan struct{}),
                waitForInitialSessionSync:        make(chan struct{}),
                waitForInitialSessionBindingSync: make(chan struct{}),
                reauthC:                          make(chan struct{}, 1),
                cfg:                              cfg,
                activeStreams:                    map[uint64]*activeStream{},
                sessionStreams:                   map[string]map[uint64]struct{}{},
                clusterEndpoints:                 map[string]map[uint64]*extensions_ssh.EndpointMetadata{},
                edsCache:                         cache.NewLinearCache(resource.EndpointType),
                endpointsUpdateQueue:             make(chan streamEndpointsUpdate, 128),
                bindingStreams:                   map[string]map[uint64]struct{}{},
                indexer:                          indexer,
        }

        bindingSyncer := &bindingSyncer{
                clientHandler: sm.GetDataBrokerServiceClient,
                clearHandler:  sm.clearRecordsBinding,
                updateHandler: sm.updateRecordsBinding,
        }

        sm.bindingSyncer = bindingSyncer

        sm.indexer.ProcessConfigUpdate(cfg)
        return sm
}

func (sm *StreamManager) waitForInitialSync(ctx context.Context) error {
        sm.mu.Lock()
        for !sm.initialSessionSyncDone || !sm.initialSessionBindingSyncDone {
                sm.mu.Unlock()
                select {
                case <-sm.waitForInitialSessionSync:
                case <-sm.waitForInitialSessionBindingSync:
                case <-time.After(10 * time.Second):
                        return errors.New("timed out waiting for initial sync")
                case <-ctx.Done():
                        return context.Cause(ctx)
                }
                sm.mu.Lock()
        }
        sm.mu.Unlock()
        return nil
}

func (sm *StreamManager) Run(ctx context.Context) error {
        sm.edsServer = delta.NewServer(ctx, sm.edsCache, sm)
        eg, eCtx := errgroup.WithContext(ctx)
        eg.Go(func() error {
                syncer := databroker.NewSyncer(
                        eCtx,
                        "ssh-auth-session-sync",
                        sm,
                        databroker.WithTypeURL("type.googleapis.com/session.Session"))
                return syncer.Run(eCtx)
        })

        eg.Go(func() error {
                syncer := databroker.NewSyncer(
                        eCtx,
                        "ssh-auth-session-binding-sync",
                        sm.bindingSyncer,
                        databroker.WithTypeURL("type.googleapis.com/session.SessionBinding"),
                )
                return syncer.Run(eCtx)
        })

        eg.Go(func() error {
                sm.reauthLoop(eCtx)
                return ErrReauthDone
        })
        eg.Go(func() error {
                sm.endpointsUpdateLoop(ctx)
                return nil
        })

        close(sm.ready)
        err := eg.Wait()
        if errors.Is(err, ErrReauthDone) {
                return nil
        }

        return err
}

func (sm *StreamManager) OnConfigChange(cfg *config.Config) {
        sm.indexer.ProcessConfigUpdate(cfg)

        // TODO: integrate the re-auth mechanism with the indexer
        select {
        case sm.reauthC <- struct{}{}:
        default:
        }
}

func (sm *StreamManager) LookupStream(streamID uint64) *StreamHandler {
        sm.mu.Lock()
        defer sm.mu.Unlock()
        if info, ok := sm.activeStreams[streamID]; ok {
                return info.Handler
        }
        return nil
}

func (sm *StreamManager) NewStreamHandler(
        _ context.Context,
        downstream *extensions_ssh.DownstreamConnectEvent,
) *StreamHandler {
        sm.mu.Lock()
        defer sm.mu.Unlock()
        streamID := downstream.StreamId

        onClose := func() {
                sm.onStreamHandlerClosed(streamID)
        }
        discovery := &streamClusterEndpointDiscovery{
                self:     sm,
                streamID: streamID,
        }
        sh := NewStreamHandler(sm.auth, discovery, sm.cfg, downstream, onClose)
        portForwardMgr := portforward.NewManager()
        sm.activeStreams[streamID] = &activeStream{
                Handler:            sh,
                Endpoints:          map[string]struct{}{},
                PortForwardManager: portForwardMgr,
        }
        sm.indexer.AddStream(streamID, portForwardMgr)
        return sh
}

func (sm *StreamManager) onStreamHandlerClosed(streamID uint64) {
        sm.mu.Lock()
        defer sm.mu.Unlock()
        info := sm.activeStreams[streamID]
        delete(sm.activeStreams, streamID)

        info.PortForwardManager.RemoveUpdateListener(info.Handler)
        sm.indexer.RemoveStream(streamID)

        if info.Session != nil {
                session := *info.Session
                delete(sm.sessionStreams[session], streamID)
                if len(sm.sessionStreams[session]) == 0 {
                        delete(sm.sessionStreams, session)
                }
        }
        if info.SessionBindingID != nil {
                bindingID := *info.SessionBindingID
                delete(sm.bindingStreams[bindingID], streamID)
                if len(sm.bindingStreams[bindingID]) == 0 {
                        delete(sm.bindingStreams, bindingID)
                }
        }

        if len(info.Endpoints) > 0 {
                sm.logger.Debug().
                        Uint64("stream-id", streamID).
                        Any("endpoints", info.Endpoints).
                        Msg("clearing endpoints for closed stream")
                sm.endpointsUpdateQueue <- streamEndpointsUpdate{
                        streamID: streamID,
                        removed:  info.Endpoints,
                }
        }
}

func (sm *StreamManager) processStreamEndpointsUpdate(update streamEndpointsUpdate) {
        sm.mu.Lock()
        defer sm.mu.Unlock()
        streamID := update.streamID

        activeStream := sm.activeStreams[streamID] // can be nil

        toUpdate := map[string]types.Resource{} // *envoy_config_endpoint_v3.LbEndpoint
        var toDelete []string
        for clusterID, info := range update.added {
                if activeStream != nil {
                        activeStream.Endpoints[clusterID] = struct{}{}
                }
                if _, ok := sm.clusterEndpoints[clusterID]; !ok {
                        sm.clusterEndpoints[clusterID] = map[uint64]*extensions_ssh.EndpointMetadata{}
                }
                sm.clusterEndpoints[clusterID][streamID] = buildEndpointMetadata(info)
                toUpdate[clusterID] = buildClusterLoadAssignment(clusterID, sm.clusterEndpoints[clusterID])
        }
        for clusterID := range update.removed {
                if activeStream != nil {
                        delete(activeStream.Endpoints, clusterID)
                }
                delete(sm.clusterEndpoints[clusterID], streamID)
                if len(sm.clusterEndpoints[clusterID]) == 0 {
                        // No more endpoints for this cluster, so delete the resource. The cluster
                        // will handle this by clearing the endpoints.
                        toDelete = append(toDelete, clusterID)
                        delete(sm.clusterEndpoints, clusterID)
                } else {
                        // There are still endpoints
                        toUpdate[clusterID] = buildClusterLoadAssignment(clusterID, sm.clusterEndpoints[clusterID])
                }
        }

        if err := sm.edsCache.UpdateResources(toUpdate, toDelete); err != nil {
                sm.logger.Err(err).Msg("error updating EDS resources")
        }
}

func buildClusterLoadAssignment(clusterID string, clusterEndpoints map[uint64]*extensions_ssh.EndpointMetadata) types.Resource {
        endpoints := []*envoy_config_endpoint_v3.LbEndpoint{}
        for streamID, metadata := range clusterEndpoints {
                endpoints = append(endpoints, buildLbEndpoint(streamID, metadata))
        }
        slices.SortFunc(endpoints, compareEndpoints)
        return &envoy_config_endpoint_v3.ClusterLoadAssignment{
                ClusterName: clusterID,
                Endpoints:   []*envoy_config_endpoint_v3.LocalityLbEndpoints{{LbEndpoints: endpoints}},
        }
}

func compareEndpoints(a, b *envoy_config_endpoint_v3.LbEndpoint) int {
        return cmp.Compare(
                a.GetEndpoint().GetAddress().GetSocketAddress().GetAddress(),
                b.GetEndpoint().GetAddress().GetSocketAddress().GetAddress())
}

func buildEndpointMetadata(info portforward.RoutePortForwardInfo) *extensions_ssh.EndpointMetadata {
        serverPort := info.Permission.ServerPort()
        return &extensions_ssh.EndpointMetadata{
                ServerPort: &extensions_ssh.ServerPort{
                        Value:     serverPort.Value,
                        IsDynamic: serverPort.IsDynamic,
                },
                MatchedPermission: &extensions_ssh.PortForwardPermission{
                        RequestedHost: info.Permission.HostMatcher.InputPattern(),
                        RequestedPort: info.Permission.RequestedPort,
                },
        }
}

func buildLbEndpoint(streamID uint64, metadata *extensions_ssh.EndpointMetadata) *envoy_config_endpoint_v3.LbEndpoint {
        endpointMdAny, _ := anypb.New(metadata)
        return &envoy_config_endpoint_v3.LbEndpoint{
                HostIdentifier: &envoy_config_endpoint_v3.LbEndpoint_Endpoint{
                        Endpoint: &envoy_config_endpoint_v3.Endpoint{
                                Address: &corev3.Address{
                                        Address: &corev3.Address_SocketAddress{
                                                SocketAddress: &corev3.SocketAddress{
                                                        Address: fmt.Sprintf("ssh:%d", streamID),
                                                        PortSpecifier: &corev3.SocketAddress_PortValue{
                                                                PortValue: metadata.ServerPort.Value,
                                                        },
                                                },
                                        },
                                },
                        },
                },
                Metadata: &corev3.Metadata{
                        TypedFilterMetadata: map[string]*anypb.Any{
                                "com.pomerium.ssh.endpoint": endpointMdAny,
                        },
                },
                HealthStatus: corev3.HealthStatus_HEALTHY,
        }
}

func (sm *StreamManager) reauthLoop(ctx context.Context) {
        for {
                select {
                case <-ctx.Done():
                        return
                case <-sm.reauthC:
                        sm.mu.Lock()
                        snapshot := make([]*activeStream, 0, len(sm.activeStreams))
                        for _, s := range sm.activeStreams {
                                snapshot = append(snapshot, s)
                        }
                        sm.mu.Unlock()

                        for _, s := range snapshot {
                                s.Handler.Reauth()
                        }
                }
        }
}

func (sm *StreamManager) endpointsUpdateLoop(ctx context.Context) {
        for {
                select {
                case <-ctx.Done():
                        return
                case update := <-sm.endpointsUpdateQueue:
                        sm.processStreamEndpointsUpdate(update)
                }
        }
}

func (sm *StreamManager) terminateStreamLocked(streamID uint64) {
        if sh, ok := sm.activeStreams[streamID]; ok {
                sh.Handler.Terminate(status.Errorf(codes.PermissionDenied, "no longer authorized"))
        }
}

pomerium / pomerium / 20384304389

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous