19475484620

Committed 18 Nov 2025 05:36PM UTC coverage: 54.799% (-0.04%) from 54.843%

Build # 19475484620

Build Type

push

github

Committed by

web-flow

Commit Message

fix: databroker client updates should propagate to ssh codes (#5935)

## Summary

SSH auth code flow is causing Pomerium to not start cleanly / update
properly.

Databroker grpc client changes now propagate to the SSH code manager

## Related issues

N/A, slack thread.

## User Explanation

N/A

## Checklist

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

Run Details

4 of 66 new or added lines in 7 files covered. (6.06%)

25 existing lines in 7 files now uncovered.

28697 of 52368 relevant lines covered (54.8%)

93.63 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 {
        clientB          databroker.ClientGetter
        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.clientB.GetDataBrokerServiceClient()
}

func newCodeManager(
        client databroker.ClientGetter,
) *codeManager {
        return &codeManager{
                clientB:  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	clientB databroker.ClientGetter
27	accessMu *sync.RWMutex
28	codeByExpiration *btree.BTreeG[Status]
29	}
30
31	var _ databroker.SyncerHandler = (*codeManager)(nil)
32
33	func (c *codeManager) ClearRecords(_ context.Context) {	×
34	c.accessMu.Lock()	×
35	defer c.accessMu.Unlock()	×
36	c.codeByExpiration.Clear(false)	×
37	}	×
38
39	func (c *codeManager) GetDataBrokerServiceClient() databroker.DataBrokerServiceClient {	×
NEW 40	return c.clientB.GetDataBrokerServiceClient()	×
41	}	×
42
43	func newCodeManager(
44	client databroker.ClientGetter,
45	) *codeManager {	1✔
46	return &codeManager{	1✔
47	clientB: 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✔
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✔
78	return a.IssuedAt.Compare(b.IssuedAt)	×
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✔
108	log.Err(err).	×
109	Ctx(ctx).	×
110	Str("component", "code-manager").	×
111	Msg("UpdateRecords : failed to unmarshall session binding request")	×
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 / 19475484620

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