19183140028

Committed 07 Nov 2025 10:38PM UTC coverage: 55.476% (-0.6%) from 56.094%

Build # 19183140028

Build Type

push

github

Committed by

web-flow

Commit Message

feat: ssh authorization code flow (#5873)

## Summary

Implements the oauth authorization code flow for native SSH connections
in pomerium.

## Related issues


[ENG-3056](https://linear.app/pomerium/issue/ENG-3056/ssh-authorization-code-flow-initial-implementation)

The issue has a comment explaining some known limitations.

## User Explanation

Allows user to use native SSH without requiring IDPs to support the
device code flow. Native ssh now uses the Authorization code flow, which
most if not all IDPs support - notably, any existing IDP configuration
for non-native ssh Pomerium should now work with native ssh.

## Checklist

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

---------

Co-authored-by: Joe Kralicky <joekralicky@gmail.com>

Run Details

284 of 1065 new or added lines in 19 files covered. (26.67%)

25 existing lines in 5 files now uncovered.

28737 of 51801 relevant lines covered (55.48%)

94.77 hits per line

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

77.78

/pkg/ssh/code/code_manager.go

package code

import (
        "cmp"
        "context"
        "slices"
        "sync"
        "time"

        "github.com/google/btree"
        "github.com/rs/zerolog/log"

        "github.com/pomerium/pomerium/pkg/grpc/databroker"
        "github.com/pomerium/pomerium/pkg/grpc/session"
)

type Status struct {
        Code       string
        BindingKey string
        IssuedAt   time.Time
        ExpiresAt  time.Time
        State      session.SessionBindingRequestState
}

type codeManager struct {
        client           databroker.DataBrokerServiceClient
        accessMu         *sync.RWMutex
        codeByExpiration *btree.BTreeG[Status]
}

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

func (c *codeManager) ClearRecords(_ context.Context) {
        c.accessMu.Lock()
        defer c.accessMu.Unlock()
        c.codeByExpiration.Clear(false)
}

func (c *codeManager) GetDataBrokerServiceClient() databroker.DataBrokerServiceClient {
        return c.client
}

func newCodeManager(
        client databroker.DataBrokerServiceClient,
) *codeManager {
        return &codeManager{
                client:   client,
                accessMu: &sync.RWMutex{},
                codeByExpiration: btree.NewG(2, func(a, b Status) bool {
                        return cmp.Or(
                                a.ExpiresAt.Compare(b.ExpiresAt),
                                cmp.Compare(a.Code, b.Code),
                                cmp.Compare(a.BindingKey, b.BindingKey),
                        ) < 0
                }),
        }
}

func (c *codeManager) GetByCodeID(codeID string) (Status, bool) {
        toFilter := []Status{}
        c.accessMu.RLock()
        c.codeByExpiration.AscendGreaterOrEqual(Status{
                Code: codeID,
        }, func(item Status) bool {
                if item.Code == codeID {
                        toFilter = append(toFilter, item)
                        return false
                }
                return true
        })
        c.accessMu.RUnlock()

        if len(toFilter) == 0 {
                return Status{}, false
        }

        slices.SortFunc(toFilter, func(a, b Status) int {
                return a.IssuedAt.Compare(b.IssuedAt)
        })

        n := len(toFilter) - 1
        return toFilter[n], true
}

func (c *codeManager) clearExpiredLocked() {
        toRemove := []Status{}
        c.codeByExpiration.AscendLessThan(Status{
                ExpiresAt: time.Now().Add(-DefaultCodeTTL),
        }, func(item Status) bool {
                toRemove = append(toRemove, item)
                return true
        })

        for _, el := range toRemove {
                c.codeByExpiration.Delete(el)
        }
}

func (c *codeManager) UpdateRecords(ctx context.Context, _ uint64, records []*databroker.Record) {
        c.accessMu.Lock()
        defer c.accessMu.Unlock()
        c.clearExpiredLocked()
        for _, record := range records {
                codeID := record.GetId()

                s := &session.SessionBindingRequest{}
                if err := record.GetData().UnmarshalTo(s); err != nil {
                        log.Err(err).
                                Ctx(ctx).
                                Str("component", "code-manager").
                                Msg("UpdateRecords : failed to unmarshall session binding request")
                        continue
                }
                c.codeByExpiration.ReplaceOrInsert(Status{
                        Code:       codeID,
                        BindingKey: s.GetKey(),
                        IssuedAt:   s.GetCreatedAt().AsTime(),
                        ExpiresAt:  s.GetExpiresAt().AsTime(),
                        State:      s.State,
                })
        }
}

1	package code
2
3	import (
4	"cmp"
5	"context"
6	"slices"
7	"sync"
8	"time"
9
10	"github.com/google/btree"
11	"github.com/rs/zerolog/log"
12
13	"github.com/pomerium/pomerium/pkg/grpc/databroker"
14	"github.com/pomerium/pomerium/pkg/grpc/session"
15	)
16
17	type Status struct {
18	Code string
19	BindingKey string
20	IssuedAt time.Time
21	ExpiresAt time.Time
22	State session.SessionBindingRequestState
23	}
24
25	type codeManager struct {
26	client databroker.DataBrokerServiceClient
27	accessMu *sync.RWMutex
28	codeByExpiration *btree.BTreeG[Status]
29	}
30
31	var _ databroker.SyncerHandler = (*codeManager)(nil)
32
NEW 33	func (c *codeManager) ClearRecords(_ context.Context) {	×
NEW 34	c.accessMu.Lock()	×
NEW 35	defer c.accessMu.Unlock()	×
NEW 36	c.codeByExpiration.Clear(false)	×
NEW 37	}	×
38
NEW 39	func (c *codeManager) GetDataBrokerServiceClient() databroker.DataBrokerServiceClient {	×
NEW 40	return c.client	×
NEW 41	}	×
42
43	func newCodeManager(
44	client databroker.DataBrokerServiceClient,
45	) *codeManager {	1✔
46	return &codeManager{	1✔
47	client: client,	1✔
48	accessMu: &sync.RWMutex{},	1✔
49	codeByExpiration: btree.NewG(2, func(a, b Status) bool {	15✔
50	return cmp.Or(	14✔
51	a.ExpiresAt.Compare(b.ExpiresAt),	14✔
52	cmp.Compare(a.Code, b.Code),	14✔
53	cmp.Compare(a.BindingKey, b.BindingKey),	14✔
54	) < 0	14✔
55	}),	14✔
56	}
57	}
58
59	func (c *codeManager) GetByCodeID(codeID string) (Status, bool) {	5✔
60	toFilter := []Status{}	5✔
61	c.accessMu.RLock()	5✔
62	c.codeByExpiration.AscendGreaterOrEqual(Status{	5✔
63	Code: codeID,	5✔
64	}, func(item Status) bool {	9✔
65	if item.Code == codeID {	8✔
66	toFilter = append(toFilter, item)	4✔
67	return false	4✔
68	}	4✔
NEW 69	return true	×
70	})
71	c.accessMu.RUnlock()	5✔
72		5✔
73	if len(toFilter) == 0 {	6✔
74	return Status{}, false	1✔
75	}	1✔
76
77	slices.SortFunc(toFilter, func(a, b Status) int {	4✔
NEW 78	return a.IssuedAt.Compare(b.IssuedAt)	×
NEW 79	})	×
80
81	n := len(toFilter) - 1	4✔
82	return toFilter[n], true	4✔
83	}
84
85	func (c *codeManager) clearExpiredLocked() {	6✔
86	toRemove := []Status{}	6✔
87	c.codeByExpiration.AscendLessThan(Status{	6✔
88	ExpiresAt: time.Now().Add(-DefaultCodeTTL),	6✔
89	}, func(item Status) bool {	7✔
90	toRemove = append(toRemove, item)	1✔
91	return true	1✔
92	})	1✔
93
94	for _, el := range toRemove {	7✔
95	c.codeByExpiration.Delete(el)	1✔
96	}	1✔
97	}
98
99	func (c codeManager) UpdateRecords(ctx context.Context, _ uint64, records []databroker.Record) {	4✔
100	c.accessMu.Lock()	4✔
101	defer c.accessMu.Unlock()	4✔
102	c.clearExpiredLocked()	4✔
103	for _, record := range records {	8✔
104	codeID := record.GetId()	4✔
105		4✔
106	s := &session.SessionBindingRequest{}	4✔
107	if err := record.GetData().UnmarshalTo(s); err != nil {	4✔
NEW 108	log.Err(err).	×
NEW 109	Ctx(ctx).	×
NEW 110	Str("component", "code-manager").	×
NEW 111	Msg("UpdateRecords : failed to unmarshall session binding request")	×
NEW 112	continue	×
113	}
114	c.codeByExpiration.ReplaceOrInsert(Status{	4✔
115	Code: codeID,	4✔
116	BindingKey: s.GetKey(),	4✔
117	IssuedAt: s.GetCreatedAt().AsTime(),	4✔
118	ExpiresAt: s.GetExpiresAt().AsTime(),	4✔
119	State: s.State,	4✔
120	})	4✔
121	}
122	}

pomerium / pomerium / 19183140028

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