Skip to content

chore: add test for workspace proxy derp meshing #12220

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Mar 5, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 35 additions & 14 deletions enterprise/coderd/coderdenttest/proxytest.go
Original file line number Diff line number Diff line change
Expand Up @@ -38,15 +38,29 @@ type ProxyOptions struct {
// ProxyURL is optional
ProxyURL *url.URL

// Token is optional. If specified, a new workspace proxy region will not be
// created, and the proxy will become a replica of the existing proxy
// region.
Token string

// FlushStats is optional
FlushStats chan chan<- struct{}
}

// NewWorkspaceProxy will configure a wsproxy.Server with the given options.
// The new wsproxy will register itself with the given coderd.API instance.
// The first user owner client is required to create the wsproxy on the coderd
// api server.
func NewWorkspaceProxy(t *testing.T, coderdAPI *coderd.API, owner *codersdk.Client, options *ProxyOptions) *wsproxy.Server {
type WorkspaceProxy struct {
*wsproxy.Server

ServerURL *url.URL
}

// NewWorkspaceProxyReplica will configure a wsproxy.Server with the given
// options. The new wsproxy replica will register itself with the given
// coderd.API instance.
//
// If a token is not provided, a new workspace proxy region is created using the
// owner client. If a token is provided, the proxy will become a replica of the
// existing proxy region.
func NewWorkspaceProxyReplica(t *testing.T, coderdAPI *coderd.API, owner *codersdk.Client, options *ProxyOptions) WorkspaceProxy {
ctx, cancelFunc := context.WithCancel(context.Background())
t.Cleanup(cancelFunc)

Expand Down Expand Up @@ -107,11 +121,15 @@ func NewWorkspaceProxy(t *testing.T, coderdAPI *coderd.API, owner *codersdk.Clie
options.Name = namesgenerator.GetRandomName(1)
}

proxyRes, err := owner.CreateWorkspaceProxy(ctx, codersdk.CreateWorkspaceProxyRequest{
Name: options.Name,
Icon: "/emojis/flag.png",
})
require.NoError(t, err, "failed to create workspace proxy")
token := options.Token
if token == "" {
proxyRes, err := owner.CreateWorkspaceProxy(ctx, codersdk.CreateWorkspaceProxyRequest{
Name: options.Name,
Icon: "/emojis/flag.png",
})
require.NoError(t, err, "failed to create workspace proxy")
token = proxyRes.ProxyToken
}

// Inherit collector options from coderd, but keep the wsproxy reporter.
statsCollectorOptions := coderdAPI.Options.WorkspaceAppsStatsCollectorOptions
Expand All @@ -121,7 +139,7 @@ func NewWorkspaceProxy(t *testing.T, coderdAPI *coderd.API, owner *codersdk.Clie
}

wssrv, err := wsproxy.New(ctx, &wsproxy.Options{
Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug),
Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug).With(slog.F("server_url", serverURL.String())),
Experiments: options.Experiments,
DashboardURL: coderdAPI.AccessURL,
AccessURL: accessURL,
Expand All @@ -131,14 +149,14 @@ func NewWorkspaceProxy(t *testing.T, coderdAPI *coderd.API, owner *codersdk.Clie
Tracing: coderdAPI.TracerProvider,
APIRateLimit: coderdAPI.APIRateLimit,
SecureAuthCookie: coderdAPI.SecureAuthCookie,
ProxySessionToken: proxyRes.ProxyToken,
ProxySessionToken: token,
DisablePathApps: options.DisablePathApps,
// We need a new registry to not conflict with the coderd internal
// proxy metrics.
PrometheusRegistry: prometheus.NewRegistry(),
DERPEnabled: !options.DerpDisabled,
DERPOnly: options.DerpOnly,
DERPServerRelayAddress: accessURL.String(),
DERPServerRelayAddress: serverURL.String(),
StatsCollectorOptions: statsCollectorOptions,
})
require.NoError(t, err)
Expand All @@ -151,5 +169,8 @@ func NewWorkspaceProxy(t *testing.T, coderdAPI *coderd.API, owner *codersdk.Clie
handler = wssrv.Handler
mutex.Unlock()

return wssrv
return WorkspaceProxy{
Server: wssrv,
ServerURL: serverURL,
}
}
4 changes: 2 additions & 2 deletions enterprise/coderd/workspaceproxy_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ func TestRegions(t *testing.T) {
require.NoError(t, err)

const proxyName = "hello"
_ = coderdenttest.NewWorkspaceProxy(t, api, client, &coderdenttest.ProxyOptions{
_ = coderdenttest.NewWorkspaceProxyReplica(t, api, client, &coderdenttest.ProxyOptions{
Name: proxyName,
AppHostname: appHostname + ".proxy",
})
Expand Down Expand Up @@ -734,7 +734,7 @@ func TestReconnectingPTYSignedToken(t *testing.T) {
proxyURL, err := url.Parse(fmt.Sprintf("https://%s.com", namesgenerator.GetRandomName(1)))
require.NoError(t, err)

_ = coderdenttest.NewWorkspaceProxy(t, api, client, &coderdenttest.ProxyOptions{
_ = coderdenttest.NewWorkspaceProxyReplica(t, api, client, &coderdenttest.ProxyOptions{
Name: namesgenerator.GetRandomName(1),
ProxyURL: proxyURL,
AppHostname: "*.sub.example.com",
Expand Down
26 changes: 14 additions & 12 deletions enterprise/wsproxy/wsproxy.go
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,7 @@ type Server struct {
ctx context.Context
cancel context.CancelFunc
derpCloseFunc func()
registerDone <-chan struct{}
registerLoop *wsproxysdk.RegisterWorkspaceProxyLoop
}

// New creates a new workspace proxy server. This requires a primary coderd
Expand Down Expand Up @@ -210,7 +210,7 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
// goroutine to periodically re-register.
replicaID := uuid.New()
osHostname := cliutil.Hostname()
regResp, registerDone, err := client.RegisterWorkspaceProxyLoop(ctx, wsproxysdk.RegisterWorkspaceProxyLoopOpts{
registerLoop, regResp, err := client.RegisterWorkspaceProxyLoop(ctx, wsproxysdk.RegisterWorkspaceProxyLoopOpts{
Logger: opts.Logger,
Request: wsproxysdk.RegisterWorkspaceProxyRequest{
AccessURL: opts.AccessURL.String(),
Expand All @@ -230,12 +230,13 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
if err != nil {
return nil, xerrors.Errorf("register proxy: %w", err)
}
s.registerDone = registerDone
err = s.handleRegister(ctx, regResp)
s.registerLoop = registerLoop

derpServer.SetMeshKey(regResp.DERPMeshKey)
err = s.handleRegister(regResp)
if err != nil {
return nil, xerrors.Errorf("handle register: %w", err)
}
derpServer.SetMeshKey(regResp.DERPMeshKey)

secKey, err := workspaceapps.KeyFromString(regResp.AppSecurityKey)
if err != nil {
Expand Down Expand Up @@ -409,16 +410,16 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
return s, nil
}

func (s *Server) RegisterNow() error {
_, err := s.registerLoop.RegisterNow()
return err
}

func (s *Server) Close() error {
s.cancel()

var err error
registerDoneWaitTicker := time.NewTicker(11 * time.Second) // the attempt timeout is 10s
select {
case <-registerDoneWaitTicker.C:
err = multierror.Append(err, xerrors.New("timed out waiting for registerDone"))
case <-s.registerDone:
}
s.registerLoop.Close()
s.derpCloseFunc()
appServerErr := s.AppServer.Close()
if appServerErr != nil {
Expand All @@ -437,11 +438,12 @@ func (*Server) mutateRegister(_ *wsproxysdk.RegisterWorkspaceProxyRequest) {
// package in the primary and update req.ReplicaError accordingly.
}

func (s *Server) handleRegister(_ context.Context, res wsproxysdk.RegisterWorkspaceProxyResponse) error {
func (s *Server) handleRegister(res wsproxysdk.RegisterWorkspaceProxyResponse) error {
addresses := make([]string, len(res.SiblingReplicas))
for i, replica := range res.SiblingReplicas {
addresses[i] = replica.RelayAddress
}
s.Logger.Debug(s.ctx, "setting DERP mesh sibling addresses", slog.F("addresses", addresses))
s.derpMesh.SetAddresses(addresses, false)

s.latestDERPMap.Store(res.DERPMap)
Expand Down
Loading