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

0.0

/pkg/ssh/code/issuer.go

package code

import (
        "context"
        "crypto/rand"
        "encoding/base64"
        "fmt"
        "sync"
        "sync/atomic"
        "time"

        "github.com/cenkalti/backoff/v4"
        "golang.org/x/sync/errgroup"
        "google.golang.org/grpc/codes"
        "google.golang.org/grpc/status"

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

type issuer struct {
        client databroker.DataBrokerServiceClient
        done   chan struct{}

        setupDone *atomic.Uint32
        setupF    *sync.Once
        // CodeAcessor

        mgr *codeManager
        Reader
        Revoker
}

func NewIssuer(ctx context.Context, client databroker.DataBrokerServiceClient) Issuer {
        doneC := make(chan struct{})
        initVal := &atomic.Uint32{}
        initVal.Store(0)
        i := &issuer{
                client:    client,
                done:      doneC,
                setupDone: initVal,
                setupF:    &sync.Once{},
                mgr:       newCodeManager(client),
                Reader:    NewReader(client),
                Revoker:   NewRevoker(client),
        }

        eg, ctxca := errgroup.WithContext(ctx)

        eg.Go(func() error {
                syncer := databroker.NewSyncer(
                        ctxca,
                        "session-biding-request-mgr",
                        i.mgr,
                        databroker.WithTypeURL("type.googleapis.com/session.SessionBindingRequest"),
                )
                return syncer.Run(ctxca)
        })
        go func() {
                defer close(i.done)
                _ = eg.Wait()
        }()
        return i
}

var _ Issuer = (*issuer)(nil)

func (i *issuer) waitForSetup() error {
        // FIXME: this needs to run once everywhere we query SessionBindingRequest and SessionBinding's
        // we want to avoid sharing a sync.Once for coordination across packages, and run this only once
        // per pomerium instance.
        i.setupF.Do(func() {
                ctxT, ca := context.WithTimeout(context.Background(), 5*time.Minute)
                defer ca()
                if err := i.setup(ctxT); err != nil {
                        panic(err)
                }
        })

        if i.setupDone.Load() == 0 {
                return fmt.Errorf("not yet initialized")
        }
        return nil
}

func (i *issuer) IssueCode() CodeID {
        code := [16]byte{}
        _, _ = rand.Read(code[:])
        codeStr := base64.RawURLEncoding.EncodeToString(code[:])
        return CodeID(codeStr)
}

func (i *issuer) OnCodeDecision(ctx context.Context, code CodeID) <-chan Status {
        ret := make(chan Status, 1)

        go func() {
                defer close(ret)
                t := time.NewTicker(time.Millisecond * 150)
                defer t.Stop()
                id := string(code)

                for {
                RETRY:
                        select {
                        case <-ctx.Done():
                                return
                        case <-t.C:
                                st, ok := i.mgr.GetByCodeID(id)
                                if !ok {
                                        goto RETRY
                                }
                                if st.ExpiresAt.Before(time.Now()) {
                                        return
                                }
                                if st.State != session.SessionBindingRequestState_InFlight {
                                        ret <- st
                                        return
                                }
                        }
                }
        }()
        return ret
}

func (i *issuer) setup(ctx context.Context) error {
        reqCap := uint64(50000)

        b := backoff.WithContext(backoff.NewExponentialBackOff(), ctx)
        if err := backoff.Retry(func() error {
                _, err := i.client.SetOptions(ctx, &databroker.SetOptionsRequest{
                        Type: "type.googleapis.com/session.SessionBindingRequest",
                        Options: &databroker.Options{
                                Capacity:        &reqCap,
                                IndexableFields: []string{"key"},
                        },
                })
                return err
        }, b); err != nil {
                return err
        }

        if err := backoff.Retry(func() error {
                _, err := i.client.SetOptions(ctx, &databroker.SetOptionsRequest{
                        Type: "type.googleapis.com/session.SessionBinding",
                        Options: &databroker.Options{
                                IndexableFields: []string{
                                        "session_id",
                                        "user_id",
                                },
                        },
                })
                return err
        }, b); err != nil {
                return err
        }

        if err := backoff.Retry(func() error {
                _, err := i.client.SetOptions(ctx, &databroker.SetOptionsRequest{
                        Type: "type.googleapis.com/session.IdentityBinding",
                        Options: &databroker.Options{
                                IndexableFields: []string{
                                        "user_id",
                                },
                        },
                })
                return err
        }, b); err != nil {
                return err
        }
        i.setupDone.CompareAndSwap(0, 1)
        return nil
}

func (i *issuer) AssociateCode(
        ctx context.Context,
        code CodeID,
        sbr *session.SessionBindingRequest,
) (CodeID, error) {
        if err := i.waitForSetup(); err != nil {
                return "", err
        }
        b := backoff.WithContext(backoff.NewExponentialBackOff(), ctx)
        maybeCode, err := backoff.RetryWithData(func() (CodeID, error) {
                maybeCode, err := getCodeByBindingKey(ctx, i.client, sbr.Key)
                if st, ok := status.FromError(err); ok && st.Code() == codes.NotFound {
                        return "", nil
                }
                return maybeCode, nil
        }, b)
        if err != nil {
                return "", err
        }
        if maybeCode == "" {
                if _, err := i.client.Put(ctx, &databroker.PutRequest{
                        Records: []*databroker.Record{
                                {
                                        Type: grpcutil.GetTypeURL(sbr),
                                        Id:   string(code),
                                        Data: protoutil.NewAny(sbr),
                                },
                        },
                }); err != nil {
                        return "", err
                }
        }
        maybeCode, err = backoff.RetryWithData(func() (CodeID, error) {
                maybeCode, err := getCodeByBindingKey(ctx, i.client, sbr.Key)
                if st, ok := status.FromError(err); ok && st.Code() == codes.NotFound {
                        return "", nil
                }
                if err != nil {
                        return "", err
                }
                if maybeCode == "" {
                        return "", fmt.Errorf("failed to resolve code")
                }
                return maybeCode, nil
        }, b)
        if err != nil {
                return "", err
        }
        return maybeCode, nil
}

func (i *issuer) Done() chan struct{} {
        return i.done
}

1	package code
2
3	import (
4	"context"
5	"crypto/rand"
6	"encoding/base64"
7	"fmt"
8	"sync"
9	"sync/atomic"
10	"time"
11
12	"github.com/cenkalti/backoff/v4"
13	"golang.org/x/sync/errgroup"
14	"google.golang.org/grpc/codes"
15	"google.golang.org/grpc/status"
16
17	"github.com/pomerium/pomerium/pkg/grpc/databroker"
18	"github.com/pomerium/pomerium/pkg/grpc/session"
19	"github.com/pomerium/pomerium/pkg/grpcutil"
20	"github.com/pomerium/pomerium/pkg/protoutil"
21	)
22
23	type issuer struct {
24	client databroker.DataBrokerServiceClient
25	done chan struct{}
26
27	setupDone *atomic.Uint32
28	setupF *sync.Once
29	// CodeAcessor
30
31	mgr *codeManager
32	Reader
33	Revoker
34	}
35
NEW 36	func NewIssuer(ctx context.Context, client databroker.DataBrokerServiceClient) Issuer {	×
NEW 37	doneC := make(chan struct{})	×
NEW 38	initVal := &atomic.Uint32{}	×
NEW 39	initVal.Store(0)	×
NEW 40	i := &issuer{	×
NEW 41	client: client,	×
NEW 42	done: doneC,	×
NEW 43	setupDone: initVal,	×
NEW 44	setupF: &sync.Once{},	×
NEW 45	mgr: newCodeManager(client),	×
NEW 46	Reader: NewReader(client),	×
NEW 47	Revoker: NewRevoker(client),	×
NEW 48	}	×
NEW 49		×
NEW 50	eg, ctxca := errgroup.WithContext(ctx)	×
NEW 51		×
NEW 52	eg.Go(func() error {	×
NEW 53	syncer := databroker.NewSyncer(	×
NEW 54	ctxca,	×
NEW 55	"session-biding-request-mgr",	×
NEW 56	i.mgr,	×
NEW 57	databroker.WithTypeURL("type.googleapis.com/session.SessionBindingRequest"),	×
NEW 58	)	×
NEW 59	return syncer.Run(ctxca)	×
NEW 60	})	×
NEW 61	go func() {	×
NEW 62	defer close(i.done)	×
NEW 63	_ = eg.Wait()	×
NEW 64	}()	×
NEW 65	return i	×
66	}
67
68	var _ Issuer = (*issuer)(nil)
69
NEW 70	func (i *issuer) waitForSetup() error {	×
NEW 71	// FIXME: this needs to run once everywhere we query SessionBindingRequest and SessionBinding's	×
NEW 72	// we want to avoid sharing a sync.Once for coordination across packages, and run this only once	×
NEW 73	// per pomerium instance.	×
NEW 74	i.setupF.Do(func() {	×
NEW 75	ctxT, ca := context.WithTimeout(context.Background(), 5*time.Minute)	×
NEW 76	defer ca()	×
NEW 77	if err := i.setup(ctxT); err != nil {	×
NEW 78	panic(err)	×
79	}
80	})
81
NEW 82	if i.setupDone.Load() == 0 {	×
NEW 83	return fmt.Errorf("not yet initialized")	×
NEW 84	}	×
NEW 85	return nil	×
86	}
87
NEW 88	func (i *issuer) IssueCode() CodeID {	×
NEW 89	code := [16]byte{}	×
NEW 90	_, _ = rand.Read(code[:])	×
NEW 91	codeStr := base64.RawURLEncoding.EncodeToString(code[:])	×
NEW 92	return CodeID(codeStr)	×
NEW 93	}	×
94
NEW 95	func (i *issuer) OnCodeDecision(ctx context.Context, code CodeID) <-chan Status {	×
NEW 96	ret := make(chan Status, 1)	×
NEW 97		×
NEW 98	go func() {	×
NEW 99	defer close(ret)	×
NEW 100	t := time.NewTicker(time.Millisecond * 150)	×
NEW 101	defer t.Stop()	×
NEW 102	id := string(code)	×
NEW 103		×
NEW 104	for {	×
NEW 105	RETRY:	×
NEW 106	select {	×
NEW 107	case <-ctx.Done():	×
NEW 108	return	×
NEW 109	case <-t.C:	×
NEW 110	st, ok := i.mgr.GetByCodeID(id)	×
NEW 111	if !ok {	×
NEW 112	goto RETRY	×
113	}
NEW 114	if st.ExpiresAt.Before(time.Now()) {	×
NEW 115	return	×
NEW 116	}	×
NEW 117	if st.State != session.SessionBindingRequestState_InFlight {	×
NEW 118	ret <- st	×
NEW 119	return	×
NEW 120	}	×
121	}
122	}
123	}()
NEW 124	return ret	×
125	}
126
NEW 127	func (i *issuer) setup(ctx context.Context) error {	×
NEW 128	reqCap := uint64(50000)	×
NEW 129		×
NEW 130	b := backoff.WithContext(backoff.NewExponentialBackOff(), ctx)	×
NEW 131	if err := backoff.Retry(func() error {	×
NEW 132	_, err := i.client.SetOptions(ctx, &databroker.SetOptionsRequest{	×
NEW 133	Type: "type.googleapis.com/session.SessionBindingRequest",	×
NEW 134	Options: &databroker.Options{	×
NEW 135	Capacity: &reqCap,	×
NEW 136	IndexableFields: []string{"key"},	×
NEW 137	},	×
NEW 138	})	×
NEW 139	return err	×
NEW 140	}, b); err != nil {	×
NEW 141	return err	×
NEW 142	}	×
143
NEW 144	if err := backoff.Retry(func() error {	×
NEW 145	_, err := i.client.SetOptions(ctx, &databroker.SetOptionsRequest{	×
NEW 146	Type: "type.googleapis.com/session.SessionBinding",	×
NEW 147	Options: &databroker.Options{	×
NEW 148	IndexableFields: []string{	×
NEW 149	"session_id",	×
NEW 150	"user_id",	×
NEW 151	},	×
NEW 152	},	×
NEW 153	})	×
NEW 154	return err	×
NEW 155	}, b); err != nil {	×
NEW 156	return err	×
NEW 157	}	×
158
NEW 159	if err := backoff.Retry(func() error {	×
NEW 160	_, err := i.client.SetOptions(ctx, &databroker.SetOptionsRequest{	×
NEW 161	Type: "type.googleapis.com/session.IdentityBinding",	×
NEW 162	Options: &databroker.Options{	×
NEW 163	IndexableFields: []string{	×
NEW 164	"user_id",	×
NEW 165	},	×
NEW 166	},	×
NEW 167	})	×
NEW 168	return err	×
NEW 169	}, b); err != nil {	×
NEW 170	return err	×
NEW 171	}	×
NEW 172	i.setupDone.CompareAndSwap(0, 1)	×
NEW 173	return nil	×
174	}
175
176	func (i *issuer) AssociateCode(
177	ctx context.Context,
178	code CodeID,
179	sbr *session.SessionBindingRequest,
NEW 180	) (CodeID, error) {	×
NEW 181	if err := i.waitForSetup(); err != nil {	×
NEW 182	return "", err	×
NEW 183	}	×
NEW 184	b := backoff.WithContext(backoff.NewExponentialBackOff(), ctx)	×
NEW 185	maybeCode, err := backoff.RetryWithData(func() (CodeID, error) {	×
NEW 186	maybeCode, err := getCodeByBindingKey(ctx, i.client, sbr.Key)	×
NEW 187	if st, ok := status.FromError(err); ok && st.Code() == codes.NotFound {	×
NEW 188	return "", nil	×
NEW 189	}	×
NEW 190	return maybeCode, nil	×
191	}, b)
NEW 192	if err != nil {	×
NEW 193	return "", err	×
NEW 194	}	×
NEW 195	if maybeCode == "" {	×
NEW 196	if _, err := i.client.Put(ctx, &databroker.PutRequest{	×
NEW 197	Records: []*databroker.Record{	×
NEW 198	{	×
NEW 199	Type: grpcutil.GetTypeURL(sbr),	×
NEW 200	Id: string(code),	×
NEW 201	Data: protoutil.NewAny(sbr),	×
NEW 202	},	×
NEW 203	},	×
NEW 204	}); err != nil {	×
NEW 205	return "", err	×
NEW 206	}	×
207	}
NEW 208	maybeCode, err = backoff.RetryWithData(func() (CodeID, error) {	×
NEW 209	maybeCode, err := getCodeByBindingKey(ctx, i.client, sbr.Key)	×
NEW 210	if st, ok := status.FromError(err); ok && st.Code() == codes.NotFound {	×
NEW 211	return "", nil	×
NEW 212	}	×
NEW 213	if err != nil {	×
NEW 214	return "", err	×
NEW 215	}	×
NEW 216	if maybeCode == "" {	×
NEW 217	return "", fmt.Errorf("failed to resolve code")	×
NEW 218	}	×
NEW 219	return maybeCode, nil	×
220	}, b)
NEW 221	if err != nil {	×
NEW 222	return "", err	×
NEW 223	}	×
NEW 224	return maybeCode, nil	×
225	}
226
NEW 227	func (i *issuer) Done() chan struct{} {	×
NEW 228	return i.done	×
NEW 229	}	×

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