16997273510

Committed 15 Aug 2025 07:03PM UTC coverage: 52.593% (+0.6%) from 51.998%

Build # 16997273510

Build Type

push

github

Committed by

web-flow

Commit Message

databroker: add server info method (#5784)

## Summary
Add a new server info method that returns the server version, earliest
record version and latest record version.

Also clean up some of the databroker code and allow gRPC reflection for
the control plane.

## Related issues
-
[ENG-2712](https://linear.app/pomerium/issue/ENG-2712/databroker-add-method-to-retrieve-server-info)

## Checklist

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

Coverage Stats

122 of 264 new or added lines in 10 files covered. (46.21%)

469 existing lines in 25 files now uncovered.

23936 of 45512 relevant lines covered (52.59%)

76.42 hits per line

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

0.0

/internal/zero/controller/controller.go

// Package controller implements Pomerium managed mode
package controller

import (
        "context"
        "fmt"
        "net"
        "net/url"
        "time"

        "github.com/rs/zerolog"
        "golang.org/x/sync/errgroup"

        "github.com/pomerium/pomerium/internal/log"
        "github.com/pomerium/pomerium/internal/retry"
        sdk "github.com/pomerium/pomerium/internal/zero/api"
        "github.com/pomerium/pomerium/internal/zero/bootstrap"
        "github.com/pomerium/pomerium/internal/zero/bootstrap/writers"
        connect_mux "github.com/pomerium/pomerium/internal/zero/connect-mux"
        "github.com/pomerium/pomerium/internal/zero/controller/usagereporter"
        "github.com/pomerium/pomerium/internal/zero/healthcheck"
        "github.com/pomerium/pomerium/internal/zero/reconciler"
        "github.com/pomerium/pomerium/internal/zero/telemetry"
        "github.com/pomerium/pomerium/internal/zero/telemetry/sessions"
        "github.com/pomerium/pomerium/pkg/cmd/pomerium"
        "github.com/pomerium/pomerium/pkg/grpc/databroker"
)

// Run runs Pomerium is managed mode using the provided token.
func Run(ctx context.Context, opts ...Option) error {
        c := controller{cfg: newControllerConfig(opts...)}

        err := c.initAPI(ctx)
        if err != nil {
                return fmt.Errorf("init api: %w", err)
        }

        var writer writers.ConfigWriter
        if c.cfg.bootstrapConfigFileName != nil {
                var err error
                var uri string
                if c.cfg.bootstrapConfigWritebackURI != nil {
                        // if there is an explicitly configured writeback URI, use it
                        uri = *c.cfg.bootstrapConfigWritebackURI
                } else {
                        // otherwise, default to "file://<filename>"
                        uri = "file://" + *c.cfg.bootstrapConfigFileName
                }
                writer, err = writers.NewForURI(uri)
                if err != nil {
                        return fmt.Errorf("error creating bootstrap config writer: %w", err)
                }
        }

        src, err := bootstrap.New([]byte(c.cfg.apiToken), c.cfg.bootstrapConfigFileName, writer, c.api)
        if err != nil {
                return fmt.Errorf("error creating bootstrap config: %w", err)
        }
        c.bootstrapConfig = src

        eg, ctx := errgroup.WithContext(ctx)
        eg.Go(func() error { return run(ctx, "connect", c.runConnect) })
        eg.Go(func() error { return run(ctx, "connect-log", c.RunConnectLog) })
        eg.Go(func() error { return run(ctx, "zero-bootstrap", c.runBootstrap) })
        eg.Go(func() error { return run(ctx, "pomerium-core", c.runPomeriumCore) })
        eg.Go(func() error { return run(ctx, "zero-control-loop", c.runZeroControlLoop) })
        eg.Go(func() error {
                <-ctx.Done()
                log.Ctx(ctx).Info().Msgf("shutting down: %v", context.Cause(ctx))
                return nil
        })
        return eg.Wait()
}

type controller struct {
        cfg *controllerConfig

        api *sdk.API

        bootstrapConfig *bootstrap.Source
}

func (c *controller) initAPI(ctx context.Context) error {
        api, err := sdk.NewAPI(ctx,
                sdk.WithClusterAPIEndpoint(c.cfg.clusterAPIEndpoint),
                sdk.WithAPIToken(c.cfg.apiToken),
                sdk.WithConnectAPIEndpoint(c.cfg.connectAPIEndpoint),
                sdk.WithOTELEndpoint(c.cfg.otelEndpoint),
        )
        if err != nil {
                return fmt.Errorf("error initializing cloud api: %w", err)
        }

        c.api = api
        return nil
}

func run(ctx context.Context, name string, runFn func(context.Context) error) error {
        log.Ctx(ctx).Debug().Str("name", name).Msg("starting")
        defer log.Ctx(ctx).Debug().Str("name", name).Msg("stopped")
        err := runFn(ctx)
        if err != nil && ctx.Err() == nil {
                return fmt.Errorf("%s: %w", name, err)
        }
        return nil
}

func (c *controller) runBootstrap(ctx context.Context) error {
        ctx = log.WithContext(ctx, func(c zerolog.Context) zerolog.Context {
                return c.Str("service", "zero-bootstrap")
        })
        return c.bootstrapConfig.Run(ctx)
}

func (c *controller) runPomeriumCore(ctx context.Context) error {
        err := c.bootstrapConfig.WaitReady(ctx)
        if err != nil {
                return fmt.Errorf("waiting for config source to be ready: %w", err)
        }
        return pomerium.Run(ctx, c.bootstrapConfig)
}

func (c *controller) runConnect(ctx context.Context) error {
        ctx = log.WithContext(ctx, func(c zerolog.Context) zerolog.Context {
                return c.Str("service", "zero-connect")
        })

        return c.api.Connect(ctx)
}

func (c *controller) runZeroControlLoop(ctx context.Context) error {
        ctx = log.WithContext(ctx, func(c zerolog.Context) zerolog.Context {
                return c.Str("control-group", "zero-cluster")
        })

        err := c.bootstrapConfig.WaitReady(ctx)
        if err != nil {
                return fmt.Errorf("waiting for config source to be ready: %w", err)
        }

        r := NewDatabrokerRestartRunner(ctx, c.bootstrapConfig)
        defer r.Close()

        var leaseStatus LeaseStatus
        tm, err := telemetry.New(ctx, c.api,
                r.GetDatabrokerClient,
                leaseStatus.HasLease,
                c.getEnvoyScrapeURL(),
        )
        if err != nil {
                return fmt.Errorf("init telemetry: %w", err)
        }
        defer c.shutdownWithTimeout(ctx, "telemetry", tm.Shutdown)

        eg, ctx := errgroup.WithContext(ctx)
        eg.Go(func() error { return tm.Run(ctx) })
        eg.Go(func() error {
                return r.Run(ctx,
                        WithLease(
                                c.runReconcilerLeased,
                                c.runSessionAnalyticsLeased,
                                c.runHealthChecksLeased,
                                leaseStatus.MonitorLease,
                                c.runUsageReporter,
                        ),
                )
        })
        return eg.Wait()
}

func (c *controller) shutdownWithTimeout(ctx context.Context, name string, fn func(context.Context) error) {
        ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), c.cfg.shutdownTimeout)
        defer cancel()

        log.Ctx(ctx).Debug().Str("timeout", c.cfg.shutdownTimeout.String()).Msgf("shutting down %s ...", name)
        err := fn(ctx)
        if err != nil {
                log.Ctx(ctx).Error().Err(err).Msgf("error shutting down %s", name)
        } else {
                log.Ctx(ctx).Debug().Msgf("%s shutdown complete", name)
        }
}

func (c *controller) runReconcilerLeased(ctx context.Context, client databroker.DataBrokerServiceClient) error {
        return retry.WithBackoff(ctx, "zero-reconciler", func(ctx context.Context) error {
                return reconciler.Run(ctx,
                        reconciler.WithAPI(c.api),
                        reconciler.WithDataBrokerClient(client),
                        reconciler.WithTracerProvider(c.cfg.tracerProvider),
                )
        })
}

func (c *controller) runSessionAnalyticsLeased(ctx context.Context, client databroker.DataBrokerServiceClient) error {
        return retry.WithBackoff(ctx, "zero-analytics", func(ctx context.Context) error {
                return sessions.Collect(ctx, client, time.Hour)
        })
}

func (c *controller) runHealthChecksLeased(ctx context.Context, client databroker.DataBrokerServiceClient) error {
        return retry.WithBackoff(ctx, "zero-healthcheck", func(ctx context.Context) error {
                checker := healthcheck.NewChecker(c.bootstrapConfig, client)
                eg, ctx := errgroup.WithContext(ctx)
                eg.Go(func() error { return checker.Run(ctx) })
                eg.Go(func() error {
                        return c.api.Watch(ctx, connect_mux.WithOnRunHealthChecks(func(_ context.Context) {
                                checker.ForceCheck()
                        }))
                })
                return eg.Wait()
        })
}

func (c *controller) runUsageReporter(ctx context.Context, client databroker.DataBrokerServiceClient) error {
        ur := usagereporter.New(c.api, c.bootstrapConfig.GetConfig().ZeroPseudonymizationKey, time.Minute)
        return retry.WithBackoff(ctx, "zero-usage-reporter", func(ctx context.Context) error {
                // start the usage reporter
                return ur.Run(ctx, client)
        })
}

func (c *controller) getEnvoyScrapeURL() string {
        return (&url.URL{
                Scheme: "http",
                Host:   net.JoinHostPort("localhost", c.bootstrapConfig.GetConfig().OutboundPort),
                Path:   "/envoy/stats/prometheus",
        }).String()
}

1	// Package controller implements Pomerium managed mode
2	package controller
3
4	import (
5	"context"
6	"fmt"
7	"net"
8	"net/url"
9	"time"
10
11	"github.com/rs/zerolog"
12	"golang.org/x/sync/errgroup"
13
14	"github.com/pomerium/pomerium/internal/log"
15	"github.com/pomerium/pomerium/internal/retry"
16	sdk "github.com/pomerium/pomerium/internal/zero/api"
17	"github.com/pomerium/pomerium/internal/zero/bootstrap"
18	"github.com/pomerium/pomerium/internal/zero/bootstrap/writers"
19	connect_mux "github.com/pomerium/pomerium/internal/zero/connect-mux"
20	"github.com/pomerium/pomerium/internal/zero/controller/usagereporter"
21	"github.com/pomerium/pomerium/internal/zero/healthcheck"
22	"github.com/pomerium/pomerium/internal/zero/reconciler"
23	"github.com/pomerium/pomerium/internal/zero/telemetry"
24	"github.com/pomerium/pomerium/internal/zero/telemetry/sessions"
25	"github.com/pomerium/pomerium/pkg/cmd/pomerium"
26	"github.com/pomerium/pomerium/pkg/grpc/databroker"
27	)
28
29	// Run runs Pomerium is managed mode using the provided token.
30	func Run(ctx context.Context, opts ...Option) error {	×
31	c := controller{cfg: newControllerConfig(opts...)}	×
32		×
33	err := c.initAPI(ctx)	×
34	if err != nil {	×
35	return fmt.Errorf("init api: %w", err)	×
36	}	×
37
38	var writer writers.ConfigWriter	×
39	if c.cfg.bootstrapConfigFileName != nil {	×
40	var err error	×
41	var uri string	×
42	if c.cfg.bootstrapConfigWritebackURI != nil {	×
43	// if there is an explicitly configured writeback URI, use it	×
44	uri = *c.cfg.bootstrapConfigWritebackURI	×
45	} else {	×
46	// otherwise, default to "file://<filename>"	×
47	uri = "file://" + *c.cfg.bootstrapConfigFileName	×
48	}	×
49	writer, err = writers.NewForURI(uri)	×
50	if err != nil {	×
51	return fmt.Errorf("error creating bootstrap config writer: %w", err)	×
52	}	×
53	}
54
55	src, err := bootstrap.New([]byte(c.cfg.apiToken), c.cfg.bootstrapConfigFileName, writer, c.api)	×
56	if err != nil {	×
57	return fmt.Errorf("error creating bootstrap config: %w", err)	×
58	}	×
59	c.bootstrapConfig = src	×
60		×
61	eg, ctx := errgroup.WithContext(ctx)	×
62	eg.Go(func() error { return run(ctx, "connect", c.runConnect) })	×
63	eg.Go(func() error { return run(ctx, "connect-log", c.RunConnectLog) })	×
64	eg.Go(func() error { return run(ctx, "zero-bootstrap", c.runBootstrap) })	×
65	eg.Go(func() error { return run(ctx, "pomerium-core", c.runPomeriumCore) })	×
66	eg.Go(func() error { return run(ctx, "zero-control-loop", c.runZeroControlLoop) })	×
67	eg.Go(func() error {	×
68	<-ctx.Done()	×
69	log.Ctx(ctx).Info().Msgf("shutting down: %v", context.Cause(ctx))	×
70	return nil	×
71	})	×
72	return eg.Wait()	×
73	}
74
75	type controller struct {
76	cfg *controllerConfig
77
78	api *sdk.API
79
80	bootstrapConfig *bootstrap.Source
81	}
82
83	func (c *controller) initAPI(ctx context.Context) error {	×
84	api, err := sdk.NewAPI(ctx,	×
85	sdk.WithClusterAPIEndpoint(c.cfg.clusterAPIEndpoint),	×
86	sdk.WithAPIToken(c.cfg.apiToken),	×
87	sdk.WithConnectAPIEndpoint(c.cfg.connectAPIEndpoint),	×
88	sdk.WithOTELEndpoint(c.cfg.otelEndpoint),	×
89	)	×
90	if err != nil {	×
91	return fmt.Errorf("error initializing cloud api: %w", err)	×
92	}	×
93
94	c.api = api	×
95	return nil	×
96	}
97
98	func run(ctx context.Context, name string, runFn func(context.Context) error) error {	×
99	log.Ctx(ctx).Debug().Str("name", name).Msg("starting")	×
100	defer log.Ctx(ctx).Debug().Str("name", name).Msg("stopped")	×
101	err := runFn(ctx)	×
102	if err != nil && ctx.Err() == nil {	×
103	return fmt.Errorf("%s: %w", name, err)	×
104	}	×
105	return nil	×
106	}
107
108	func (c *controller) runBootstrap(ctx context.Context) error {	×
109	ctx = log.WithContext(ctx, func(c zerolog.Context) zerolog.Context {	×
110	return c.Str("service", "zero-bootstrap")	×
111	})	×
112	return c.bootstrapConfig.Run(ctx)	×
113	}
114
115	func (c *controller) runPomeriumCore(ctx context.Context) error {	×
116	err := c.bootstrapConfig.WaitReady(ctx)	×
117	if err != nil {	×
118	return fmt.Errorf("waiting for config source to be ready: %w", err)	×
119	}	×
120	return pomerium.Run(ctx, c.bootstrapConfig)	×
121	}
122
123	func (c *controller) runConnect(ctx context.Context) error {	×
124	ctx = log.WithContext(ctx, func(c zerolog.Context) zerolog.Context {	×
125	return c.Str("service", "zero-connect")	×
126	})	×
127
128	return c.api.Connect(ctx)	×
129	}
130
131	func (c *controller) runZeroControlLoop(ctx context.Context) error {	×
132	ctx = log.WithContext(ctx, func(c zerolog.Context) zerolog.Context {	×
133	return c.Str("control-group", "zero-cluster")	×
134	})	×
135
136	err := c.bootstrapConfig.WaitReady(ctx)	×
137	if err != nil {	×
138	return fmt.Errorf("waiting for config source to be ready: %w", err)	×
139	}	×
140
141	r := NewDatabrokerRestartRunner(ctx, c.bootstrapConfig)	×
142	defer r.Close()	×
143		×
144	var leaseStatus LeaseStatus	×
145	tm, err := telemetry.New(ctx, c.api,	×
146	r.GetDatabrokerClient,	×
147	leaseStatus.HasLease,	×
148	c.getEnvoyScrapeURL(),	×
149	)	×
150	if err != nil {	×
151	return fmt.Errorf("init telemetry: %w", err)	×
152	}	×
153	defer c.shutdownWithTimeout(ctx, "telemetry", tm.Shutdown)	×
154		×
155	eg, ctx := errgroup.WithContext(ctx)	×
156	eg.Go(func() error { return tm.Run(ctx) })	×
157	eg.Go(func() error {	×
158	return r.Run(ctx,	×
159	WithLease(	×
160	c.runReconcilerLeased,	×
161	c.runSessionAnalyticsLeased,	×
162	c.runHealthChecksLeased,	×
163	leaseStatus.MonitorLease,	×
164	c.runUsageReporter,	×
165	),	×
166	)	×
167	})	×
168	return eg.Wait()	×
169	}
170
171	func (c *controller) shutdownWithTimeout(ctx context.Context, name string, fn func(context.Context) error) {	×
172	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), c.cfg.shutdownTimeout)	×
173	defer cancel()	×
174		×
175	log.Ctx(ctx).Debug().Str("timeout", c.cfg.shutdownTimeout.String()).Msgf("shutting down %s ...", name)	×
176	err := fn(ctx)	×
177	if err != nil {	×
178	log.Ctx(ctx).Error().Err(err).Msgf("error shutting down %s", name)	×
179	} else {	×
180	log.Ctx(ctx).Debug().Msgf("%s shutdown complete", name)	×
181	}	×
182	}
183
184	func (c *controller) runReconcilerLeased(ctx context.Context, client databroker.DataBrokerServiceClient) error {	×
185	return retry.WithBackoff(ctx, "zero-reconciler", func(ctx context.Context) error {	×
186	return reconciler.Run(ctx,	×
187	reconciler.WithAPI(c.api),	×
188	reconciler.WithDataBrokerClient(client),	×
189	reconciler.WithTracerProvider(c.cfg.tracerProvider),	×
190	)	×
UNCOV 191	})	×
192	}
193
194	func (c *controller) runSessionAnalyticsLeased(ctx context.Context, client databroker.DataBrokerServiceClient) error {	×
195	return retry.WithBackoff(ctx, "zero-analytics", func(ctx context.Context) error {	×
196	return sessions.Collect(ctx, client, time.Hour)	×
UNCOV 197	})	×
198	}
199
200	func (c *controller) runHealthChecksLeased(ctx context.Context, client databroker.DataBrokerServiceClient) error {	×
201	return retry.WithBackoff(ctx, "zero-healthcheck", func(ctx context.Context) error {	×
202	checker := healthcheck.NewChecker(c.bootstrapConfig, client)	×
203	eg, ctx := errgroup.WithContext(ctx)	×
204	eg.Go(func() error { return checker.Run(ctx) })	×
205	eg.Go(func() error {	×
206	return c.api.Watch(ctx, connect_mux.WithOnRunHealthChecks(func(_ context.Context) {	×
207	checker.ForceCheck()	×
UNCOV 208	}))	×
209	})
UNCOV 210	return eg.Wait()	×
211	})
212	}
213
214	func (c *controller) runUsageReporter(ctx context.Context, client databroker.DataBrokerServiceClient) error {	×
215	ur := usagereporter.New(c.api, c.bootstrapConfig.GetConfig().ZeroPseudonymizationKey, time.Minute)	×
216	return retry.WithBackoff(ctx, "zero-usage-reporter", func(ctx context.Context) error {	×
217	// start the usage reporter	×
218	return ur.Run(ctx, client)	×
UNCOV 219	})	×
220	}
221
222	func (c *controller) getEnvoyScrapeURL() string {	×
223	return (&url.URL{	×
224	Scheme: "http",	×
225	Host: net.JoinHostPort("localhost", c.bootstrapConfig.GetConfig().OutboundPort),	×
226	Path: "/envoy/stats/prometheus",	×
227	}).String()	×
UNCOV 228	}	×

pomerium / pomerium / 16997273510

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