Skip to content

Commit db16e4c

Browse files
authored
fix: add the directory info to the tar header (#91)
Signed-off-by: chlins <[email protected]>
1 parent c4ea1c7 commit db16e4c

File tree

4 files changed

+252
-66
lines changed

4 files changed

+252
-66
lines changed

pkg/archiver/archiver.go

Lines changed: 133 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -25,63 +25,115 @@ import (
2525
"strings"
2626
)
2727

28-
// Tar tars the target file and return the content by stream.
29-
func Tar(path string) (io.Reader, error) {
28+
// Tar creates a tar archive of the specified path (file or directory)
29+
// and returns the content as a stream. For individual files, it preserves
30+
// the directory structure relative to the working directory.
31+
func Tar(srcPath string, workDir string) (io.Reader, error) {
3032
pr, pw := io.Pipe()
33+
3134
go func() {
3235
defer pw.Close()
33-
// create the tar writer.
3436
tw := tar.NewWriter(pw)
3537
defer tw.Close()
3638

37-
file, err := os.Open(path)
39+
info, err := os.Stat(srcPath)
3840
if err != nil {
39-
pw.CloseWithError(fmt.Errorf("failed to open file: %w", err))
41+
pw.CloseWithError(fmt.Errorf("failed to stat source path: %w", err))
4042
return
4143
}
4244

43-
defer file.Close()
44-
info, err := file.Stat()
45-
if err != nil {
46-
pw.CloseWithError(fmt.Errorf("failed to stat file: %w", err))
47-
return
48-
}
45+
// Handle directories and files differently.
46+
if info.IsDir() {
47+
// For directories, walk through and add all files/subdirs.
48+
err = filepath.Walk(srcPath, func(path string, info os.FileInfo, err error) error {
49+
if err != nil {
50+
return err
51+
}
52+
53+
// Create a relative path for the tar file header.
54+
relPath, err := filepath.Rel(workDir, path)
55+
if err != nil {
56+
return fmt.Errorf("failed to get relative path: %w", err)
57+
}
58+
59+
header, err := tar.FileInfoHeader(info, "")
60+
if err != nil {
61+
return fmt.Errorf("failed to create tar header: %w", err)
62+
}
63+
64+
// Set the header name to preserve directory structure.
65+
header.Name = relPath
66+
if err := tw.WriteHeader(header); err != nil {
67+
return fmt.Errorf("failed to write header: %w", err)
68+
}
69+
70+
if !info.IsDir() {
71+
file, err := os.Open(path)
72+
if err != nil {
73+
return fmt.Errorf("failed to open file %s: %w", path, err)
74+
}
75+
defer file.Close()
76+
77+
if _, err := io.Copy(tw, file); err != nil {
78+
return fmt.Errorf("failed to write file %s to tar: %w", path, err)
79+
}
80+
}
81+
82+
return nil
83+
})
4984

50-
header, err := tar.FileInfoHeader(info, info.Name())
51-
if err != nil {
52-
pw.CloseWithError(fmt.Errorf("failed to create tar file info header: %w", err))
53-
return
54-
}
85+
if err != nil {
86+
pw.CloseWithError(fmt.Errorf("failed to walk directory: %w", err))
87+
return
88+
}
89+
} else {
90+
// For a single file, include the directory structure.
91+
file, err := os.Open(srcPath)
92+
if err != nil {
93+
pw.CloseWithError(fmt.Errorf("failed to open file: %w", err))
94+
return
95+
}
96+
defer file.Close()
5597

56-
if err := tw.WriteHeader(header); err != nil {
57-
pw.CloseWithError(fmt.Errorf("failed to write header to tar writer: %w", err))
58-
return
59-
}
98+
header, err := tar.FileInfoHeader(info, "")
99+
if err != nil {
100+
pw.CloseWithError(fmt.Errorf("failed to create tar header: %w", err))
101+
return
102+
}
60103

61-
_, err = io.Copy(tw, file)
62-
if err != nil {
63-
pw.CloseWithError(fmt.Errorf("failed to copy file to tar writer: %w", err))
64-
return
104+
// Use relative path as the header name to preserve directory structure
105+
// This keeps the directory structure as part of the file path in the tar.
106+
relPath, err := filepath.Rel(workDir, srcPath)
107+
if err != nil {
108+
pw.CloseWithError(fmt.Errorf("failed to get relative path: %w", err))
109+
return
110+
}
111+
112+
// Use the relative path (including directories) as the header name.
113+
header.Name = relPath
114+
if err := tw.WriteHeader(header); err != nil {
115+
pw.CloseWithError(fmt.Errorf("failed to write header: %w", err))
116+
return
117+
}
118+
119+
if _, err := io.Copy(tw, file); err != nil {
120+
pw.CloseWithError(fmt.Errorf("failed to copy file to tar: %w", err))
121+
return
122+
}
65123
}
66124
}()
67125

68126
return pr, nil
69127
}
70128

71-
// Untar untars the target stream to the destination path.
129+
// Untar extracts the contents of a tar archive from the provided reader
130+
// to the specified destination path.
72131
func Untar(reader io.Reader, destPath string) error {
73-
// uncompress gzip if it is a .tar.gz file
74-
// gzipReader, err := gzip.NewReader(reader)
75-
// if err != nil {
76-
// return err
77-
// }
78-
// defer gzipReader.Close()
79-
// tarReader := tar.NewReader(gzipReader)
80-
81132
tarReader := tar.NewReader(reader)
82133

134+
// Ensure destination directory exists.
83135
if err := os.MkdirAll(destPath, 0755); err != nil {
84-
return err
136+
return fmt.Errorf("failed to create destination directory: %w", err)
85137
}
86138

87139
for {
@@ -90,39 +142,74 @@ func Untar(reader io.Reader, destPath string) error {
90142
break
91143
}
92144
if err != nil {
93-
return err
145+
return fmt.Errorf("error reading tar: %w", err)
94146
}
95147

96-
// sanitize filepaths to prevent directory traversal.
148+
// Sanitize file paths to prevent directory traversal.
97149
cleanPath := filepath.Clean(header.Name)
98-
if strings.Contains(cleanPath, "..") {
150+
if strings.Contains(cleanPath, "..") || strings.HasPrefix(cleanPath, "/") || strings.HasPrefix(cleanPath, ":\\") {
99151
return fmt.Errorf("tar file contains invalid path: %s", cleanPath)
100152
}
101153

102-
path := filepath.Join(destPath, cleanPath)
103-
// check the file type.
154+
targetPath := filepath.Join(destPath, cleanPath)
155+
156+
// Create directories for all path components.
157+
dirPath := filepath.Dir(targetPath)
158+
if err := os.MkdirAll(dirPath, 0755); err != nil {
159+
return fmt.Errorf("failed to create directory %s: %w", dirPath, err)
160+
}
161+
104162
switch header.Typeflag {
105163
case tar.TypeDir:
106-
if err := os.MkdirAll(path, 0755); err != nil {
107-
return err
164+
if err := os.MkdirAll(targetPath, os.FileMode(header.Mode)); err != nil {
165+
return fmt.Errorf("failed to create directory %s: %w", targetPath, err)
108166
}
167+
109168
case tar.TypeReg:
110-
file, err := os.Create(path)
169+
file, err := os.OpenFile(
170+
targetPath,
171+
os.O_CREATE|os.O_RDWR|os.O_TRUNC,
172+
os.FileMode(header.Mode),
173+
)
111174
if err != nil {
112-
return err
175+
return fmt.Errorf("failed to create file %s: %w", targetPath, err)
113176
}
114177

115178
if _, err := io.Copy(file, tarReader); err != nil {
116179
file.Close()
117-
return err
180+
return fmt.Errorf("failed to write to file %s: %w", targetPath, err)
118181
}
119182
file.Close()
120183

121-
if err := os.Chmod(path, os.FileMode(header.Mode)); err != nil {
122-
return err
184+
case tar.TypeSymlink:
185+
if isRel(header.Linkname, destPath) && isRel(header.Name, destPath) {
186+
if err := os.Symlink(header.Linkname, targetPath); err != nil {
187+
return fmt.Errorf("failed to create symlink %s -> %s: %w", targetPath, header.Linkname, err)
188+
}
189+
} else {
190+
return fmt.Errorf("symlink %s -> %s points outside of destination directory", targetPath, header.Linkname)
123191
}
192+
193+
default:
194+
// Skip other types.
195+
continue
124196
}
125197
}
126198

127199
return nil
128200
}
201+
202+
// isRel checks if the candidate path is within the target directory after resolving symbolic links.
203+
func isRel(candidate, target string) bool {
204+
if filepath.IsAbs(candidate) {
205+
return false
206+
}
207+
208+
realpath, err := filepath.EvalSymlinks(filepath.Join(target, candidate))
209+
if err != nil {
210+
return false
211+
}
212+
213+
relpath, err := filepath.Rel(target, realpath)
214+
return err == nil && !strings.HasPrefix(filepath.Clean(relpath), "..")
215+
}

pkg/archiver/archiver_test.go

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
/*
2+
* Copyright 2024 The CNAI Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package archiver
18+
19+
import (
20+
"bytes"
21+
"io"
22+
"os"
23+
"path/filepath"
24+
"testing"
25+
)
26+
27+
func TestTar(t *testing.T) {
28+
tmpDir, err := os.MkdirTemp("", "archiver_test")
29+
if err != nil {
30+
t.Fatal(err)
31+
}
32+
defer os.RemoveAll(tmpDir)
33+
34+
filePath := filepath.Join(tmpDir, "testfile.txt")
35+
if err := os.WriteFile(filePath, []byte("hello"), 0644); err != nil {
36+
t.Fatalf("write file error: %v", err)
37+
}
38+
39+
tarReader, err := Tar(filePath, tmpDir)
40+
if err != nil {
41+
t.Fatalf("Tar error: %v", err)
42+
}
43+
44+
var buf bytes.Buffer
45+
if _, err := io.Copy(&buf, tarReader); err != nil {
46+
t.Fatalf("copy tar error: %v", err)
47+
}
48+
49+
if buf.Len() == 0 {
50+
t.Fatal("tar archive is empty")
51+
}
52+
}
53+
54+
func TestUntar(t *testing.T) {
55+
tmpDir, err := os.MkdirTemp("", "archiver_test")
56+
if err != nil {
57+
t.Fatal(err)
58+
}
59+
defer os.RemoveAll(tmpDir)
60+
61+
filePath := filepath.Join(tmpDir, "testfile.txt")
62+
if err := os.WriteFile(filePath, []byte("hello"), 0644); err != nil {
63+
t.Fatalf("write file error: %v", err)
64+
}
65+
66+
tarReader, err := Tar(filePath, tmpDir)
67+
if err != nil {
68+
t.Fatalf("Tar error: %v", err)
69+
}
70+
71+
var buf bytes.Buffer
72+
if _, err := io.Copy(&buf, tarReader); err != nil {
73+
t.Fatalf("copy tar error: %v", err)
74+
}
75+
76+
extractDir, err := os.MkdirTemp("", "archiver_extracted")
77+
if err != nil {
78+
t.Fatal(err)
79+
}
80+
defer os.RemoveAll(extractDir)
81+
82+
if err := Untar(bytes.NewReader(buf.Bytes()), extractDir); err != nil {
83+
t.Fatalf("Untar error: %v", err)
84+
}
85+
86+
extractedFile := filepath.Join(extractDir, filepath.Base(filePath))
87+
data, err := os.ReadFile(extractedFile)
88+
if err != nil {
89+
t.Fatalf("read extracted file error: %v", err)
90+
}
91+
92+
if string(data) != "hello" {
93+
t.Errorf("expected 'hello', got '%s'", string(data))
94+
}
95+
}

pkg/backend/build/build.go

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -21,37 +21,49 @@ import (
2121
"context"
2222
"encoding/json"
2323
"fmt"
24+
"os"
2425
"path/filepath"
2526
"time"
2627

2728
"github.com/CloudNativeAI/modctl/pkg/archiver"
2829
"github.com/CloudNativeAI/modctl/pkg/modelfile"
2930
"github.com/CloudNativeAI/modctl/pkg/storage"
30-
modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
3131

32+
modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
3233
godigest "github.com/opencontainers/go-digest"
3334
spec "github.com/opencontainers/image-spec/specs-go"
3435
ocispec "github.com/opencontainers/image-spec/specs-go/v1"
3536
)
3637

3738
// BuildLayer converts the file to the image blob and push it to the storage.
3839
func BuildLayer(ctx context.Context, store storage.Storage, mediaType, workDir, repo, path string) (ocispec.Descriptor, error) {
39-
reader, err := archiver.Tar(path)
40+
info, err := os.Stat(path)
4041
if err != nil {
41-
return ocispec.Descriptor{}, fmt.Errorf("failed to tar file: %w", err)
42+
return ocispec.Descriptor{}, fmt.Errorf("failed to get file info: %w", err)
4243
}
4344

44-
digest, size, err := store.PushBlob(ctx, repo, reader)
45-
if err != nil {
46-
return ocispec.Descriptor{}, fmt.Errorf("failed to push blob to storage: %w", err)
45+
if info.IsDir() {
46+
return ocispec.Descriptor{}, fmt.Errorf("%s is a directory and not supported yet", path)
4747
}
4848

49-
absPath, err := filepath.Abs(workDir)
49+
workDirPath, err := filepath.Abs(workDir)
5050
if err != nil {
5151
return ocispec.Descriptor{}, fmt.Errorf("failed to get absolute path of workDir: %w", err)
5252
}
5353

54-
filePath, err := filepath.Rel(absPath, path)
54+
reader, err := archiver.Tar(path, workDirPath)
55+
if err != nil {
56+
return ocispec.Descriptor{}, fmt.Errorf("failed to tar file: %w", err)
57+
}
58+
59+
digest, size, err := store.PushBlob(ctx, repo, reader)
60+
if err != nil {
61+
return ocispec.Descriptor{}, fmt.Errorf("failed to push blob to storage: %w", err)
62+
}
63+
64+
// Gets the relative path of the file as annotation.
65+
//nolint:typecheck
66+
relPath, err := filepath.Rel(workDirPath, path)
5567
if err != nil {
5668
return ocispec.Descriptor{}, fmt.Errorf("failed to get relative path: %w", err)
5769
}
@@ -61,7 +73,7 @@ func BuildLayer(ctx context.Context, store storage.Storage, mediaType, workDir,
6173
Digest: godigest.Digest(digest),
6274
Size: size,
6375
Annotations: map[string]string{
64-
modelspec.AnnotationFilepath: filePath,
76+
modelspec.AnnotationFilepath: relPath,
6577
},
6678
}, nil
6779
}

0 commit comments

Comments
 (0)