Skip to content

Commit 1c9a64d

Browse files
dobracdobrac
authored
Fix UFFD panic in template build (#1088)
1 parent 4e53782 commit 1c9a64d

File tree

8 files changed

+143
-79
lines changed

8 files changed

+143
-79
lines changed

‎packages/orchestrator/internal/sandbox/block/empty.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ func NewEmpty(size int64, blockSize int64, buildID uuid.UUID) (*Empty, error) {
3030
func (e *Empty) ReadAt(p []byte, off int64) (int, error) {
3131
slice, err := e.Slice(off, int64(len(p)))
3232
if err != nil {
33-
return 0, fmt.Errorf("failed to slice mmap: %w", err)
33+
return 0, fmt.Errorf("failed to slice empty: %w", err)
3434
}
3535

3636
return copy(p, slice), nil

‎packages/orchestrator/internal/sandbox/template/local_template.go‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,3 @@ func (t *LocalTemplate) Snapfile() (File, error) {
5050
func (t *LocalTemplate) Metadata() (metadata.Template, error) {
5151
return metadata.Template{}, errors.New("metadata not available in local template")
5252
}
53-
54-
func (t *LocalTemplate) ReplaceMemfile(memfile block.ReadonlyDevice) error {
55-
t.memfile = memfile
56-
return nil
57-
}

‎packages/orchestrator/internal/sandbox/template/mask_template.go‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
package template
2+
3+
import (
4+
"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block"
5+
"github.com/e2b-dev/infra/packages/orchestrator/internal/template/metadata"
6+
"github.com/e2b-dev/infra/packages/shared/pkg/storage"
7+
)
8+
9+
type MaskTemplate struct {
10+
template Template
11+
12+
memfile *block.ReadonlyDevice
13+
}
14+
15+
type MaskTemplateOption func(*MaskTemplate)
16+
17+
func WithMemfile(memfile block.ReadonlyDevice) MaskTemplateOption {
18+
return func(c *MaskTemplate) {
19+
c.memfile = &memfile
20+
}
21+
}
22+
23+
func NewMaskTemplate(
24+
template Template,
25+
opts ...MaskTemplateOption,
26+
) *MaskTemplate {
27+
t := &MaskTemplate{
28+
template: template,
29+
memfile: nil,
30+
}
31+
32+
for _, opt := range opts {
33+
opt(t)
34+
}
35+
36+
return t
37+
}
38+
39+
func (c *MaskTemplate) Close() error {
40+
if c.memfile != nil {
41+
return (*c.memfile).Close()
42+
}
43+
44+
return nil
45+
}
46+
47+
func (c *MaskTemplate) Files() storage.TemplateCacheFiles {
48+
return c.template.Files()
49+
}
50+
51+
func (c *MaskTemplate) Memfile() (block.ReadonlyDevice, error) {
52+
if c.memfile != nil {
53+
return *c.memfile, nil
54+
}
55+
return c.template.Memfile()
56+
}
57+
58+
func (c *MaskTemplate) Rootfs() (block.ReadonlyDevice, error) {
59+
return c.template.Rootfs()
60+
}
61+
62+
func (c *MaskTemplate) Snapfile() (File, error) {
63+
return c.template.Snapfile()
64+
}
65+
66+
func (c *MaskTemplate) Metadata() (metadata.Template, error) {
67+
return c.template.Metadata()
68+
}

‎packages/orchestrator/internal/sandbox/template/storage_template.go‎

Lines changed: 64 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,11 @@ func (t *storageTemplate) Fetch(ctx context.Context, buildStore *build.DiffStore
7575

7676
wg.Go(func() error {
7777
if t.localSnapfile != nil {
78-
return t.snapfile.SetValue(t.localSnapfile)
78+
if err := t.snapfile.SetValue(t.localSnapfile); err != nil {
79+
return fmt.Errorf("failed to set local snapfile: %w", err)
80+
}
81+
82+
return nil
7983
}
8084

8185
snapfile, snapfileErr := newStorageFile(
@@ -87,15 +91,27 @@ func (t *storageTemplate) Fetch(ctx context.Context, buildStore *build.DiffStore
8791
if snapfileErr != nil {
8892
errMsg := fmt.Errorf("failed to fetch snapfile: %w", snapfileErr)
8993

90-
return t.snapfile.SetError(errMsg)
94+
if err := t.snapfile.SetError(errMsg); err != nil {
95+
return fmt.Errorf("failed to set snapfile error: %w", errors.Join(errMsg, err))
96+
}
97+
98+
return nil
9199
}
92100

93-
return t.snapfile.SetValue(snapfile)
101+
if err := t.snapfile.SetValue(snapfile); err != nil {
102+
return fmt.Errorf("failed to set snapfile: %w", err)
103+
}
104+
105+
return nil
94106
})
95107

96108
wg.Go(func() error {
97109
if t.localMetafile != nil {
98-
return t.metafile.SetValue(t.localMetafile)
110+
if err := t.metafile.SetValue(t.localMetafile); err != nil {
111+
return fmt.Errorf("failed to set local metafile: %w", err)
112+
}
113+
114+
return nil
99115
}
100116

101117
meta, err := newStorageFile(
@@ -105,7 +121,12 @@ func (t *storageTemplate) Fetch(ctx context.Context, buildStore *build.DiffStore
105121
t.files.CacheMetadataPath(),
106122
)
107123
if err != nil && !errors.Is(err, storage.ErrorObjectNotExist) {
108-
return t.metafile.SetError(fmt.Errorf("failed to fetch metafile: %w", err))
124+
sourceErr := fmt.Errorf("failed to fetch metafile: %w", err)
125+
if err := t.metafile.SetError(sourceErr); err != nil {
126+
return fmt.Errorf("failed to set metafile error: %w", errors.Join(sourceErr, err))
127+
}
128+
129+
return nil
109130
}
110131

111132
if err != nil {
@@ -118,15 +139,28 @@ func (t *storageTemplate) Fetch(ctx context.Context, buildStore *build.DiffStore
118139
oldTemplateMetadata := metadata.V1TemplateVersion()
119140
err := oldTemplateMetadata.ToFile(t.files.CacheMetadataPath())
120141
if err != nil {
121-
return t.metafile.SetError(fmt.Errorf("failed to write v1 template metadata to a file: %w", err))
142+
sourceErr := fmt.Errorf("failed to write v1 template metadata to a file: %w", err)
143+
if err := t.metafile.SetError(sourceErr); err != nil {
144+
return fmt.Errorf("failed to set metafile error: %w", errors.Join(sourceErr, err))
145+
}
146+
147+
return nil
122148
}
123149

124-
return t.metafile.SetValue(&storageFile{
150+
if err := t.metafile.SetValue(&storageFile{
125151
path: t.files.CacheMetadataPath(),
126-
})
152+
}); err != nil {
153+
return fmt.Errorf("failed to set metafile v1: %w", err)
154+
}
155+
156+
return nil
127157
}
128158

129-
return t.metafile.SetValue(meta)
159+
if err := t.metafile.SetValue(meta); err != nil {
160+
return fmt.Errorf("failed to set metafile value: %w", err)
161+
}
162+
163+
return nil
130164
})
131165

132166
wg.Go(func() error {
@@ -143,10 +177,18 @@ func (t *storageTemplate) Fetch(ctx context.Context, buildStore *build.DiffStore
143177
if memfileErr != nil {
144178
errMsg := fmt.Errorf("failed to create memfile storage: %w", memfileErr)
145179

146-
return t.memfile.SetError(errMsg)
180+
if err := t.memfile.SetError(errMsg); err != nil {
181+
return fmt.Errorf("failed to set memfile error: %w", errors.Join(errMsg, err))
182+
}
183+
184+
return nil
185+
}
186+
187+
if err := t.memfile.SetValue(memfileStorage); err != nil {
188+
return fmt.Errorf("failed to set memfile value: %w", err)
147189
}
148190

149-
return t.memfile.SetValue(memfileStorage)
191+
return nil
150192
})
151193

152194
wg.Go(func() error {
@@ -162,16 +204,24 @@ func (t *storageTemplate) Fetch(ctx context.Context, buildStore *build.DiffStore
162204
if rootfsErr != nil {
163205
errMsg := fmt.Errorf("failed to create rootfs storage: %w", rootfsErr)
164206

165-
return t.rootfs.SetError(errMsg)
207+
if err := t.rootfs.SetError(errMsg); err != nil {
208+
return fmt.Errorf("failed to set rootfs error: %w", errors.Join(errMsg, err))
209+
}
210+
211+
return nil
212+
}
213+
214+
if err := t.rootfs.SetValue(rootfsStorage); err != nil {
215+
return fmt.Errorf("failed to set rootfs value: %w", err)
166216
}
167217

168-
return t.rootfs.SetValue(rootfsStorage)
218+
return nil
169219
})
170220

171221
err := wg.Wait()
172222
if err != nil {
173223
zap.L().Error("failed to fetch template files",
174-
zap.String("build_id", t.files.BuildID),
224+
logger.WithBuildID(t.files.BuildID),
175225
zap.Error(err),
176226
)
177227
return
@@ -206,10 +256,3 @@ func (t *storageTemplate) Metadata() (metadata.Template, error) {
206256

207257
return metadata.FromFile(metafile.Path())
208258
}
209-
210-
func (t *storageTemplate) ReplaceMemfile(memfile block.ReadonlyDevice) error {
211-
m := utils.NewSetOnce[block.ReadonlyDevice]()
212-
m.SetValue(memfile)
213-
t.memfile = m
214-
return nil
215-
}

‎packages/orchestrator/internal/sandbox/template/template.go‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@ import (
1313
type Template interface {
1414
Files() storage.TemplateCacheFiles
1515
Memfile() (block.ReadonlyDevice, error)
16-
ReplaceMemfile(block.ReadonlyDevice) error
1716
Rootfs() (block.ReadonlyDevice, error)
1817
Snapfile() (File, error)
1918
Metadata() (metadata.Template, error)

‎packages/orchestrator/internal/template/build/core/memory/memory.go‎

Lines changed: 0 additions & 36 deletions
This file was deleted.

‎packages/orchestrator/internal/template/build/layer/create_sandbox.go‎

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -31,24 +31,20 @@ func NewCreateSandbox(config sandbox.Config, fcVersions fc.FirecrackerVersions)
3131
func (f *CreateSandbox) Sandbox(
3232
ctx context.Context,
3333
layerExecutor *LayerExecutor,
34-
template sbxtemplate.Template,
34+
sourceTemplate sbxtemplate.Template,
3535
) (*sandbox.Sandbox, error) {
3636
// Create new memfile with the size of the sandbox RAM, this updates the underlying memfile.
3737
// This is ok as the sandbox is started from the beginning.
38-
var memfile block.ReadonlyDevice
3938
memfile, err := block.NewEmpty(
4039
f.config.RamMB<<constants.ToMBShift,
4140
config.MemfilePageSize(f.config.HugePages),
42-
uuid.MustParse(template.Files().BuildID),
41+
uuid.MustParse(sourceTemplate.Files().BuildID),
4342
)
4443
if err != nil {
4544
return nil, fmt.Errorf("create memfile: %w", err)
4645
}
4746

48-
err = template.ReplaceMemfile(memfile)
49-
if err != nil {
50-
return nil, fmt.Errorf("replace memfile: %w", err)
51-
}
47+
template := sbxtemplate.NewMaskTemplate(sourceTemplate, sbxtemplate.WithMemfile(memfile))
5248

5349
// In case of a new sandbox, base template ID is now used as the potentially exported template base ID.
5450
sbx, err := sandbox.CreateSandbox(

‎packages/orchestrator/internal/template/build/phases/base/files.go‎

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@ import (
1212
"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block"
1313
"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/buildcontext"
1414
"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/config"
15-
"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/core/memory"
1615
"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/core/rootfs"
16+
"github.com/e2b-dev/infra/packages/orchestrator/internal/template/constants"
1717
artifactsregistry "github.com/e2b-dev/infra/packages/shared/pkg/artifacts-registry"
1818
)
1919

@@ -26,7 +26,7 @@ func constructLayerFilesFromOCI(
2626
artifactRegistry artifactsregistry.ArtifactsRegistry,
2727
templateBuildDir string,
2828
rootfsPath string,
29-
) (r *block.Local, m *block.Local, c containerregistry.Config, e error) {
29+
) (r *block.Local, m block.ReadonlyDevice, c containerregistry.Config, e error) {
3030
childCtx, childSpan := tracer.Start(ctx, "template-build")
3131
defer childSpan.End()
3232

@@ -59,15 +59,14 @@ func constructLayerFilesFromOCI(
5959
}
6060

6161
// Create empty memfile
62-
memfilePath, err := memory.NewMemory(templateBuildDir, buildContext.Config.MemoryMB)
62+
memfile, err := block.NewEmpty(
63+
buildContext.Config.MemoryMB<<constants.ToMBShift,
64+
config.MemfilePageSize(buildContext.Config.HugePages),
65+
buildIDParsed,
66+
)
6367
if err != nil {
6468
return nil, nil, containerregistry.Config{}, fmt.Errorf("error creating memfile: %w", err)
6569
}
6670

67-
memfile, err := block.NewLocal(memfilePath, config.MemfilePageSize(buildContext.Config.HugePages), buildIDParsed)
68-
if err != nil {
69-
return nil, nil, containerregistry.Config{}, fmt.Errorf("error creating memfile blocks: %w", err)
70-
}
71-
7271
return rootfs, memfile, imgConfig, nil
7372
}

0 commit comments

Comments
 (0)