Skip to content

Commit 72e6ea1

Browse files
bduffanybduffany
authored
Move ImageCacheAuthenticator to env (#4869)
1 parent 3f22530 commit 72e6ea1

File tree

17 files changed

+188
-163
lines changed

17 files changed

+188
-163
lines changed

‎enterprise/server/cmd/executor/executor.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,9 @@ func main() {
199199
}
200200
executorID := executorUUID.String()
201201

202+
imageCacheAuth := container.NewImageCacheAuthenticator(container.ImageCacheAuthenticatorOpts{})
203+
env.SetImageCacheAuthenticator(imageCacheAuth)
204+
202205
runnerPool, err := runner.NewPool(env, &runner.PoolOptions{})
203206
if err != nil {
204207
log.Fatalf("Failed to initialize runner pool: %s", err)

‎enterprise/server/remote_execution/container/container.go‎

Lines changed: 20 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,8 @@ var (
4242

4343
containerRegistries = flag.Slice("executor.container_registries", []ContainerRegistry{}, "")
4444
debugUseLocalImagesOnly = flag.Bool("debug_use_local_images_only", false, "Do not pull OCI images and only used locally cached images. This can be set to test local image builds during development without needing to push to a container registry. Not intended for production use.")
45+
46+
slowPullWarnOnce sync.Once
4547
)
4648

4749
type ContainerRegistry struct {
@@ -231,13 +233,17 @@ type Stdio struct {
231233

232234
// PullImageIfNecessary pulls the image configured for the container if it
233235
// is not cached locally.
234-
func PullImageIfNecessary(ctx context.Context, env environment.Env, cacheAuth *ImageCacheAuthenticator, ctr CommandContainer, creds PullCredentials, imageRef string) error {
236+
func PullImageIfNecessary(ctx context.Context, env environment.Env, ctr CommandContainer, creds PullCredentials, imageRef string) error {
235237
if *debugUseLocalImagesOnly {
236238
return nil
237239
}
238-
if env.GetAuthenticator() == nil {
240+
cacheAuth := env.GetImageCacheAuthenticator()
241+
if cacheAuth == nil || env.GetAuthenticator() == nil {
239242
// If we don't have an authenticator available, fall back to
240243
// authenticating the creds with the image registry on every request.
244+
slowPullWarnOnce.Do(func() {
245+
log.CtxWarningf(ctx, "Authentication is not properly configured; this will result in slower image pulls.")
246+
})
241247
return ctr.PullImage(ctx, creds)
242248
}
243249
isCached, err := ctr.IsImageCached(ctx)
@@ -279,39 +285,33 @@ func (p PullCredentials) String() string {
279285
return p.Username + ":" + p.Password
280286
}
281287

282-
// ImageCacheToken is a claim to be able to access a locally cached image.
283-
type ImageCacheToken struct {
284-
GroupID string
285-
ImageRef string
286-
}
287-
288288
// NewImageCacheToken returns the token representing the authenticated group ID,
289289
// pull credentials, and image ref. For the same sets of those values, the
290290
// same token is always returned.
291-
func NewImageCacheToken(ctx context.Context, env environment.Env, creds PullCredentials, imageRef string) (ImageCacheToken, error) {
291+
func NewImageCacheToken(ctx context.Context, env environment.Env, creds PullCredentials, imageRef string) (interfaces.ImageCacheToken, error) {
292292
groupID := ""
293293
u, err := perms.AuthenticatedUser(ctx, env)
294294
if err != nil {
295295
if !authutil.IsAnonymousUserError(err) {
296-
return ImageCacheToken{}, err
296+
return interfaces.ImageCacheToken{}, err
297297
}
298298
} else {
299299
groupID = u.GetGroupID()
300300
}
301-
return ImageCacheToken{
301+
return interfaces.ImageCacheToken{
302302
GroupID: groupID,
303303
ImageRef: imageRef,
304304
}, nil
305305
}
306306

307-
// ImageCacheAuthenticator grants access to short-lived tokens for accessing
307+
// imageCacheAuthenticator grants access to short-lived tokens for accessing
308308
// locally cached images without needing to re-authenticate with the remote
309309
// registry (which can be slow).
310-
type ImageCacheAuthenticator struct {
310+
type imageCacheAuthenticator struct {
311311
opts ImageCacheAuthenticatorOpts
312312

313313
mu sync.Mutex // protects(tokenExpireTimes)
314-
tokenExpireTimes map[ImageCacheToken]time.Time
314+
tokenExpireTimes map[interfaces.ImageCacheToken]time.Time
315315
}
316316

317317
type ImageCacheAuthenticatorOpts struct {
@@ -320,31 +320,31 @@ type ImageCacheAuthenticatorOpts struct {
320320
TokenTTL time.Duration
321321
}
322322

323-
func NewImageCacheAuthenticator(opts ImageCacheAuthenticatorOpts) *ImageCacheAuthenticator {
323+
func NewImageCacheAuthenticator(opts ImageCacheAuthenticatorOpts) *imageCacheAuthenticator {
324324
if opts.TokenTTL == 0 {
325325
opts.TokenTTL = defaultImageCacheTokenTTL
326326
}
327-
return &ImageCacheAuthenticator{
327+
return &imageCacheAuthenticator{
328328
opts: opts,
329-
tokenExpireTimes: map[ImageCacheToken]time.Time{},
329+
tokenExpireTimes: map[interfaces.ImageCacheToken]time.Time{},
330330
}
331331
}
332332

333-
func (a *ImageCacheAuthenticator) IsAuthorized(token ImageCacheToken) bool {
333+
func (a *imageCacheAuthenticator) IsAuthorized(token interfaces.ImageCacheToken) bool {
334334
a.mu.Lock()
335335
defer a.mu.Unlock()
336336
a.purgeExpiredTokens()
337337
_, ok := a.tokenExpireTimes[token]
338338
return ok
339339
}
340340

341-
func (a *ImageCacheAuthenticator) Refresh(token ImageCacheToken) {
341+
func (a *imageCacheAuthenticator) Refresh(token interfaces.ImageCacheToken) {
342342
a.mu.Lock()
343343
defer a.mu.Unlock()
344344
a.tokenExpireTimes[token] = time.Now().Add(a.opts.TokenTTL)
345345
}
346346

347-
func (a *ImageCacheAuthenticator) purgeExpiredTokens() {
347+
func (a *imageCacheAuthenticator) purgeExpiredTokens() {
348348
for token, expireTime := range a.tokenExpireTimes {
349349
if time.Now().After(expireTime) {
350350
delete(a.tokenExpireTimes, token)

‎enterprise/server/remote_execution/container/container_test.go‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ func TestPullImageIfNecessary_ValidCredentials(t *testing.T) {
6161
env := testenv.GetTestEnv(t)
6262
ta := testauth.NewTestAuthenticator(testauth.TestUsers("US1", "GR1", "US2", "GR2"))
6363
env.SetAuthenticator(ta)
64-
cacheAuth := container.NewImageCacheAuthenticator(container.ImageCacheAuthenticatorOpts{})
64+
env.SetImageCacheAuthenticator(container.NewImageCacheAuthenticator(container.ImageCacheAuthenticatorOpts{}))
6565
imageRef := "docker.io/some-org/some-image:v1.0.0"
6666
ctx := userCtx(t, ta, "US1")
6767
goodCreds1 := container.PullCredentials{
@@ -76,17 +76,17 @@ func TestPullImageIfNecessary_ValidCredentials(t *testing.T) {
7676

7777
assert.Equal(t, 0, c.PullCount, "sanity check: pull count should be 0 initially")
7878

79-
err := container.PullImageIfNecessary(ctx, env, cacheAuth, c, goodCreds1, imageRef)
79+
err := container.PullImageIfNecessary(ctx, env, c, goodCreds1, imageRef)
8080

8181
require.NoError(t, err)
8282
assert.Equal(t, 1, c.PullCount, "should pull the image if credentials are valid")
8383

84-
err = container.PullImageIfNecessary(ctx, env, cacheAuth, c, goodCreds1, imageRef)
84+
err = container.PullImageIfNecessary(ctx, env, c, goodCreds1, imageRef)
8585

8686
require.NoError(t, err)
8787
assert.Equal(t, 1, c.PullCount, "should not need to immediately re-authenticate with the remote registry")
8888

89-
err = container.PullImageIfNecessary(ctx, env, cacheAuth, c, goodCreds2, imageRef)
89+
err = container.PullImageIfNecessary(ctx, env, c, goodCreds2, imageRef)
9090

9191
require.NoError(t, err)
9292
assert.Equal(
@@ -98,7 +98,7 @@ func TestPullImageIfNecessary_InvalidCredentials_PermissionDenied(t *testing.T)
9898
env := testenv.GetTestEnv(t)
9999
ta := testauth.NewTestAuthenticator(testauth.TestUsers("US1", "GR1", "US2", "GR2"))
100100
env.SetAuthenticator(ta)
101-
cacheAuth := container.NewImageCacheAuthenticator(container.ImageCacheAuthenticatorOpts{})
101+
env.SetImageCacheAuthenticator(container.NewImageCacheAuthenticator(container.ImageCacheAuthenticatorOpts{}))
102102
imageRef := "docker.io/some-org/some-image:v1.0.0"
103103
ctx := userCtx(t, ta, "US1")
104104
goodCreds := container.PullCredentials{
@@ -111,15 +111,15 @@ func TestPullImageIfNecessary_InvalidCredentials_PermissionDenied(t *testing.T)
111111
Password: "trying-to-guess-the-real-secret",
112112
}
113113

114-
err := container.PullImageIfNecessary(ctx, env, cacheAuth, c, badCreds, imageRef)
114+
err := container.PullImageIfNecessary(ctx, env, c, badCreds, imageRef)
115115

116116
require.True(t, status.IsPermissionDeniedError(err), "should return PermissionDenied if credentials are valid")
117117

118-
err = container.PullImageIfNecessary(ctx, env, cacheAuth, c, badCreds, imageRef)
118+
err = container.PullImageIfNecessary(ctx, env, c, badCreds, imageRef)
119119

120120
require.True(t, status.IsPermissionDeniedError(err), "should return PermissionDenied on subsequent attempts as well")
121121

122-
err = container.PullImageIfNecessary(ctx, env, cacheAuth, c, goodCreds, imageRef)
122+
err = container.PullImageIfNecessary(ctx, env, c, goodCreds, imageRef)
123123

124124
require.NoError(t, err, "good creds should still work after previous incorrect attempts")
125125
}

‎enterprise/server/remote_execution/containers/docker/docker.go‎

Lines changed: 16 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -60,10 +60,9 @@ var (
6060
)
6161

6262
type Provider struct {
63-
env environment.Env
64-
client *dockerclient.Client
65-
imageCacheAuth *container.ImageCacheAuthenticator
66-
buildRoot string
63+
env environment.Env
64+
client *dockerclient.Client
65+
buildRoot string
6766
}
6867

6968
var (
@@ -100,17 +99,16 @@ func NewClient() (*dockerclient.Client, error) {
10099
return dockerClient, initErr
101100
}
102101

103-
func NewProvider(env environment.Env, imageCacheAuth *container.ImageCacheAuthenticator, hostBuildRoot string) (*Provider, error) {
102+
func NewProvider(env environment.Env, hostBuildRoot string) (*Provider, error) {
104103
client, err := NewClient()
105104
if err != nil {
106105
return nil, status.FailedPreconditionErrorf("Failed to create docker client: %s", err)
107106
}
108107

109108
return &Provider{
110-
env: env,
111-
client: client,
112-
imageCacheAuth: imageCacheAuth,
113-
buildRoot: hostBuildRoot,
109+
env: env,
110+
client: client,
111+
buildRoot: hostBuildRoot,
114112
}, nil
115113
}
116114

@@ -130,7 +128,7 @@ func (p *Provider) New(ctx context.Context, props *platform.Properties, task *re
130128
Volumes: *dockerVolumes,
131129
InheritUserIDs: *dockerInheritUserIDs,
132130
}
133-
return NewDockerContainer(p.env, p.imageCacheAuth, p.client, props.ContainerImage, p.buildRoot, opts), nil
131+
return NewDockerContainer(p.env, p.client, props.ContainerImage, p.buildRoot, opts), nil
134132
}
135133

136134
type DockerOptions struct {
@@ -151,8 +149,7 @@ type DockerOptions struct {
151149

152150
// dockerCommandContainer containerizes a command's execution using a Docker container.
153151
type dockerCommandContainer struct {
154-
env environment.Env
155-
imageCacheAuth *container.ImageCacheAuthenticator
152+
env environment.Env
156153

157154
image string
158155
// hostRootDir is the path on the _host_ machine ("node", in k8s land) of the
@@ -173,14 +170,13 @@ type dockerCommandContainer struct {
173170
removed bool
174171
}
175172

176-
func NewDockerContainer(env environment.Env, imageCacheAuth *container.ImageCacheAuthenticator, client *dockerclient.Client, image, hostRootDir string, options *DockerOptions) *dockerCommandContainer {
173+
func NewDockerContainer(env environment.Env, client *dockerclient.Client, image, hostRootDir string, options *DockerOptions) *dockerCommandContainer {
177174
return &dockerCommandContainer{
178-
env: env,
179-
imageCacheAuth: imageCacheAuth,
180-
image: image,
181-
hostRootDir: hostRootDir,
182-
client: client,
183-
options: options,
175+
env: env,
176+
image: image,
177+
hostRootDir: hostRootDir,
178+
client: client,
179+
options: options,
184180
}
185181
}
186182

@@ -206,7 +202,7 @@ func (r *dockerCommandContainer) Run(ctx context.Context, command *repb.Command,
206202

207203
// explicitly pull the image before running to avoid the
208204
// pull output logs spilling into the execution logs.
209-
if err := container.PullImageIfNecessary(ctx, r.env, r.imageCacheAuth, r, creds, r.image); err != nil {
205+
if err := container.PullImageIfNecessary(ctx, r.env, r, creds, r.image); err != nil {
210206
result.Error = wrapDockerErr(err, fmt.Sprintf("failed to pull docker image %q", r.image))
211207
return result
212208
}

‎enterprise/server/remote_execution/containers/docker/docker_test.go‎

Lines changed: 12 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -59,8 +59,8 @@ func TestDockerRun(t *testing.T) {
5959
}
6060
env := testenv.GetTestEnv(t)
6161
env.SetAuthenticator(testauth.NewTestAuthenticator(testauth.TestUsers("US1", "GR1")))
62-
cacheAuth := container.NewImageCacheAuthenticator(container.ImageCacheAuthenticatorOpts{})
63-
c := docker.NewDockerContainer(env, cacheAuth, dc, "mirror.gcr.io/library/busybox", rootDir, cfg)
62+
env.SetImageCacheAuthenticator(container.NewImageCacheAuthenticator(container.ImageCacheAuthenticatorOpts{}))
63+
c := docker.NewDockerContainer(env, dc, "mirror.gcr.io/library/busybox", rootDir, cfg)
6464

6565
res := c.Run(ctx, cmd, workDir, container.PullCredentials{})
6666

@@ -95,8 +95,8 @@ func TestDockerLifecycleControl(t *testing.T) {
9595
}
9696
env := testenv.GetTestEnv(t)
9797
env.SetAuthenticator(testauth.NewTestAuthenticator(testauth.TestUsers("US1", "GR1")))
98-
cacheAuth := container.NewImageCacheAuthenticator(container.ImageCacheAuthenticatorOpts{})
99-
c := docker.NewDockerContainer(env, cacheAuth, dc, "mirror.gcr.io/library/busybox", rootDir, cfg)
98+
env.SetImageCacheAuthenticator(container.NewImageCacheAuthenticator(container.ImageCacheAuthenticatorOpts{}))
99+
c := docker.NewDockerContainer(env, dc, "mirror.gcr.io/library/busybox", rootDir, cfg)
100100

101101
isContainerRunning := false
102102
t.Cleanup(func() {
@@ -114,7 +114,7 @@ func TestDockerLifecycleControl(t *testing.T) {
114114
})
115115

116116
err = container.PullImageIfNecessary(
117-
ctx, env, cacheAuth, c, container.PullCredentials{},
117+
ctx, env, c, container.PullCredentials{},
118118
"mirror.gcr.io/library/busybox",
119119
)
120120
require.NoError(t, err)
@@ -184,11 +184,10 @@ func TestDockerRun_Timeout_StdoutStderrStillVisible(t *testing.T) {
184184
}}
185185
env := testenv.GetTestEnv(t)
186186
env.SetAuthenticator(testauth.NewTestAuthenticator(testauth.TestUsers("US1", "GR1")))
187-
cacheAuth := container.NewImageCacheAuthenticator(container.ImageCacheAuthenticatorOpts{})
188-
c := docker.NewDockerContainer(env, cacheAuth, dc, "mirror.gcr.io/library/busybox", rootDir, cfg)
187+
c := docker.NewDockerContainer(env, dc, "mirror.gcr.io/library/busybox", rootDir, cfg)
189188
// Ensure the image is cached
190189
err = container.PullImageIfNecessary(
191-
ctx, env, cacheAuth, c, container.PullCredentials{}, "mirror.gcr.io/library/busybox")
190+
ctx, env, c, container.PullCredentials{}, "mirror.gcr.io/library/busybox")
192191
require.NoError(t, err)
193192

194193
ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
@@ -238,11 +237,10 @@ func TestDockerExec_Timeout_StdoutStderrStillVisible(t *testing.T) {
238237
}}
239238
env := testenv.GetTestEnv(t)
240239
env.SetAuthenticator(testauth.NewTestAuthenticator(testauth.TestUsers("US1", "GR1")))
241-
cacheAuth := container.NewImageCacheAuthenticator(container.ImageCacheAuthenticatorOpts{})
242-
c := docker.NewDockerContainer(env, cacheAuth, dc, "mirror.gcr.io/library/busybox", rootDir, cfg)
240+
c := docker.NewDockerContainer(env, dc, "mirror.gcr.io/library/busybox", rootDir, cfg)
243241
// Ensure the image is cached
244242
err = container.PullImageIfNecessary(
245-
ctx, env, cacheAuth, c, container.PullCredentials{}, "mirror.gcr.io/library/busybox")
243+
ctx, env, c, container.PullCredentials{}, "mirror.gcr.io/library/busybox")
246244
require.NoError(t, err)
247245

248246
err = c.Create(ctx, workDir)
@@ -298,10 +296,9 @@ func TestDockerExec_Stdio(t *testing.T) {
298296
}
299297
env := testenv.GetTestEnv(t)
300298
env.SetAuthenticator(testauth.NewTestAuthenticator(testauth.TestUsers("US1", "GR1")))
301-
cacheAuth := container.NewImageCacheAuthenticator(container.ImageCacheAuthenticatorOpts{})
302-
c := docker.NewDockerContainer(env, cacheAuth, dc, "mirror.gcr.io/library/busybox", rootDir, cfg)
299+
c := docker.NewDockerContainer(env, dc, "mirror.gcr.io/library/busybox", rootDir, cfg)
303300
err = container.PullImageIfNecessary(
304-
ctx, env, cacheAuth, c, container.PullCredentials{},
301+
ctx, env, c, container.PullCredentials{},
305302
"mirror.gcr.io/library/busybox",
306303
)
307304
require.NoError(t, err)
@@ -347,8 +344,7 @@ func TestDockerRun_LongRunningProcess_CanGetAllLogs(t *testing.T) {
347344
}
348345
env := testenv.GetTestEnv(t)
349346
env.SetAuthenticator(testauth.NewTestAuthenticator(testauth.TestUsers("US1", "GR1")))
350-
cacheAuth := container.NewImageCacheAuthenticator(container.ImageCacheAuthenticatorOpts{})
351-
c := docker.NewDockerContainer(env, cacheAuth, dc, "mirror.gcr.io/library/busybox", rootDir, cfg)
347+
c := docker.NewDockerContainer(env, dc, "mirror.gcr.io/library/busybox", rootDir, cfg)
352348

353349
res := c.Run(ctx, cmd, workDir, container.PullCredentials{})
354350

‎enterprise/server/remote_execution/containers/firecracker/BUILD‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ go_test(
9090
"//enterprise/server/remote_execution/runner",
9191
"//enterprise/server/remote_execution/snaploader",
9292
"//enterprise/server/remote_execution/workspace",
93+
"//enterprise/server/testutil/testcontainer",
9394
"//enterprise/server/util/ext4",
9495
"//proto:firecracker_go_proto",
9596
"//proto:remote_execution_go_proto",
@@ -148,6 +149,7 @@ go_test(
148149
"//enterprise/server/remote_execution/runner",
149150
"//enterprise/server/remote_execution/snaploader",
150151
"//enterprise/server/remote_execution/workspace",
152+
"//enterprise/server/testutil/testcontainer",
151153
"//enterprise/server/util/ext4",
152154
"//proto:firecracker_go_proto",
153155
"//proto:remote_execution_go_proto",

0 commit comments

Comments
 (0)