feat: support interceptor for build (#151)

Signed-off-by: chlins <[email protected]>

Author: chlins

pkg/backend/attach.go

@@ -32,6 +32,7 @@ import (
 	"github.com/CloudNativeAI/modctl/pkg/backend/build"
 	buildconfig "github.com/CloudNativeAI/modctl/pkg/backend/build/config"
 	"github.com/CloudNativeAI/modctl/pkg/backend/build/hooks"
+	"github.com/CloudNativeAI/modctl/pkg/backend/build/interceptor"
 	"github.com/CloudNativeAI/modctl/pkg/backend/processor"
 	"github.com/CloudNativeAI/modctl/pkg/backend/remote"
 	"github.com/CloudNativeAI/modctl/pkg/config"
@@ -299,6 +300,10 @@ func (b *backend) getBuilder(reference string, cfg *config.Attach) (build.Builde
 		build.WithPlainHTTP(cfg.PlainHTTP),
 		build.WithInsecure(cfg.Insecure),
 	}
+	if cfg.Nydusify {
+		opts = append(opts, build.WithInterceptor(interceptor.NewNydus()))
+	}
+
 	builder, err := build.NewBuilder(outputType, b.store, repo, tag, opts...)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create builder: %w", err)

pkg/backend/build.go

@@ -26,6 +26,7 @@ import (
 	"github.com/CloudNativeAI/modctl/pkg/backend/build"
 	buildconfig "github.com/CloudNativeAI/modctl/pkg/backend/build/config"
 	"github.com/CloudNativeAI/modctl/pkg/backend/build/hooks"
+	"github.com/CloudNativeAI/modctl/pkg/backend/build/interceptor"
 	"github.com/CloudNativeAI/modctl/pkg/backend/processor"
 	"github.com/CloudNativeAI/modctl/pkg/config"
 	"github.com/CloudNativeAI/modctl/pkg/modelfile"
@@ -67,6 +68,10 @@ func (b *backend) Build(ctx context.Context, modelfilePath, workDir, target stri
 		build.WithPlainHTTP(cfg.PlainHTTP),
 		build.WithInsecure(cfg.Insecure),
 	}
+	if cfg.Nydusify {
+		opts = append(opts, build.WithInterceptor(interceptor.NewNydus()))
+	}
+
 	builder, err := build.NewBuilder(outputType, b.store, repo, tag, opts...)
 	if err != nil {
 		return fmt.Errorf("failed to create builder: %w", err)

pkg/backend/build/builder.go

@@ -24,6 +24,7 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"sync"
 	"time"
 
 	modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
@@ -34,6 +35,7 @@ import (
 
 	buildconfig "github.com/CloudNativeAI/modctl/pkg/backend/build/config"
 	"github.com/CloudNativeAI/modctl/pkg/backend/build/hooks"
+	"github.com/CloudNativeAI/modctl/pkg/backend/build/interceptor"
 	"github.com/CloudNativeAI/modctl/pkg/codec"
 	"github.com/CloudNativeAI/modctl/pkg/storage"
 )
@@ -96,10 +98,11 @@ func NewBuilder(outputType OutputType, store storage.Storage, repo, tag string,
 	}
 
 	return &abstractBuilder{
-		store:    store,
-		repo:     repo,
-		tag:      tag,
-		strategy: strategy,
+		store:       store,
+		repo:        repo,
+		tag:         tag,
+		strategy:    strategy,
+		interceptor: cfg.interceptor,
 	}, nil
 }
 
@@ -110,6 +113,8 @@ type abstractBuilder struct {
 	tag   string
 	// strategy is the output strategy used to output the blob.
 	strategy OutputStrategy
+	// interceptor is the interceptor used to intercept the build process.
+	interceptor interceptor.Interceptor
 }
 
 func (ab *abstractBuilder) BuildLayer(ctx context.Context, mediaType, workDir, path string, hooks hooks.Hooks) (ocispec.Descriptor, error) {
@@ -167,7 +172,39 @@ func (ab *abstractBuilder) BuildLayer(ctx context.Context, mediaType, workDir, p
 		}
 	}
 
-	return ab.strategy.OutputLayer(ctx, mediaType, relPath, digest, size, reader, hooks)
+	var (
+		wg        sync.WaitGroup
+		itErr     error
+		applyDesc interceptor.ApplyDescriptorFn
+	)
+	// Intercept the reader if needed.
+	if ab.interceptor != nil {
+		var itReader io.Reader
+		reader, itReader = splitReader(reader)
+
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			applyDesc, itErr = ab.interceptor.Intercept(ctx, mediaType, relPath, codec.Type(), itReader)
+		}()
+	}
+
+	desc, err := ab.strategy.OutputLayer(ctx, mediaType, relPath, digest, size, reader, hooks)
+	if err != nil {
+		return desc, err
+	}
+
+	// Wait for the interceptor to finish.
+	wg.Wait()
+	if itErr != nil {
+		return desc, itErr
+	}
+
+	if applyDesc != nil {
+		applyDesc(&desc)
+	}
+
+	return desc, nil
 }
 
 func (ab *abstractBuilder) BuildConfig(ctx context.Context, layers []ocispec.Descriptor, modelConfig *buildconfig.Model, hooks hooks.Hooks) (ocispec.Descriptor, error) {
@@ -247,3 +284,23 @@ func buildModelConfig(modelConfig *buildconfig.Model, layers []ocispec.Descripto
 		ModelFS:    fs,
 	}, nil
 }
+
+// splitReader splits the original reader into two readers.
+func splitReader(original io.Reader) (io.Reader, io.Reader) {
+	r1, w1 := io.Pipe()
+	r2, w2 := io.Pipe()
+	multiWriter := io.MultiWriter(w1, w2)
+
+	go func() {
+		defer w1.Close()
+		defer w2.Close()
+
+		_, err := io.Copy(multiWriter, original)
+		if err != nil {
+			w1.CloseWithError(err)
+			w2.CloseWithError(err)
+		}
+	}()
+
+	return r1, r2
+}

pkg/backend/build/config.go

@@ -16,12 +16,17 @@
 
 package build
 
+import (
+	"github.com/CloudNativeAI/modctl/pkg/backend/build/interceptor"
+)
+
 type Option func(*config)
 
 // config is the configuration for the building.
 type config struct {
-	plainHTTP bool
-	insecure  bool
+	plainHTTP   bool
+	insecure    bool
+	interceptor interceptor.Interceptor
 }
 
 func WithPlainHTTP(plainHTTP bool) Option {
@@ -35,3 +40,9 @@ func WithInsecure(insecure bool) Option {
 		c.insecure = insecure
 	}
 }
+
+func WithInterceptor(interceptor interceptor.Interceptor) Option {
+	return func(c *config) {
+		c.interceptor = interceptor
+	}
+}

pkg/backend/build/interceptor/interceptor.go

@@ -0,0 +1,33 @@
+/*
+ *     Copyright 2025 The CNAI Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package interceptor
+
+import (
+	"context"
+	"io"
+
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+)
+
+// ApplyDescriptorFn is a function that applies changes to the descriptor.
+type ApplyDescriptorFn func(desc *ocispec.Descriptor)
+
+// Interceptor is an interface that defines the interceptor for the building stream.
+type Interceptor interface {
+	// Intercept intercepts the building stream for some customized logic, readerType is the original stream type, such as raw or tar.
+	Intercept(ctx context.Context, mediaType string, filepath string, readerType string, reader io.Reader) (ApplyDescriptorFn, error)
+}

pkg/backend/build/interceptor/nydus.go

@@ -0,0 +1,33 @@
+/*
+ *     Copyright 2025 The CNAI Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package interceptor
+
+import (
+	"context"
+	"io"
+)
+
+type nydus struct{}
+
+func NewNydus() *nydus {
+	return &nydus{}
+}
+
+func (n *nydus) Intercept(ctx context.Context, mediaType string, filepath string, readerType string, reader io.Reader) (ApplyDescriptorFn, error) {
+	// TODO: Implement nydus interceptor
+	return nil, nil
+}

pkg/codec/codec.go

@@ -34,6 +34,9 @@ const (
 
 // Codec is an interface for encoding and decoding the data.
 type Codec interface {
+	// Type returns the type of the codec.
+	Type() Type
+
 	// Encode encodes the target file into a reader.
 	Encode(targetFilePath, workDirPath string) (io.Reader, error)
 

pkg/codec/raw.go

@@ -30,6 +30,11 @@ func newRaw() *raw {
 	return &raw{}
 }
 
+// Type returns the type of the codec.
+func (r *raw) Type() string {
+	return Raw
+}
+
 // Encode reads the target file into a reader.
 func (r *raw) Encode(targetFilePath, workDirPath string) (io.Reader, error) {
 	return os.Open(targetFilePath)

pkg/codec/tar.go

@@ -30,6 +30,11 @@ func newTar() *tar {
 	return &tar{}
 }
 
+// Type returns the type of the codec.
+func (t *tar) Type() string {
+	return Tar
+}
+
 // Encode tars the target file into a reader.
 func (t *tar) Encode(targetFilePath, workDirPath string) (io.Reader, error) {
 	return archiver.Tar(targetFilePath, workDirPath)