feat: add http proxy option for pull command (#34)

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

Author: gaius-qi

cmd/pull.go

@@ -46,6 +46,7 @@ var pullCmd = &cobra.Command{
 func init() {
 	flags := pullCmd.Flags()
 	flags.BoolVarP(&pullConfig.PlainHTTP, "plain-http", "p", false, "use plain HTTP instead of HTTPS")
+	flags.StringVar(&pullConfig.Proxy, "proxy", "", "use proxy for the pull operation")
 
 	if err := viper.BindPFlags(flags); err != nil {
 		panic(fmt.Errorf("bind cache pull flags to viper: %w", err))
@@ -68,6 +69,10 @@ func runPull(ctx context.Context, target string) error {
 		opts = append(opts, backend.WithPlainHTTP())
 	}
 
+	if pullConfig.Proxy != "" {
+		opts = append(opts, backend.WithProxy(pullConfig.Proxy))
+	}
+
 	if err := b.Pull(ctx, target, opts...); err != nil {
 		return err
 	}

pkg/backend/options.go

@@ -20,6 +20,7 @@ type Option func(*Options)
 
 type Options struct {
 	plainHTTP bool
+	proxy     string
 }
 
 // WithPlainHTTP sets the plain HTTP option.
@@ -28,3 +29,10 @@ func WithPlainHTTP() Option {
 		opts.plainHTTP = true
 	}
 }
+
+// WithProxy sets the proxy option.
+func WithProxy(proxy string) Option {
+	return func(opts *Options) {
+		opts.proxy = proxy
+	}
+}

pkg/backend/pull.go

@@ -21,6 +21,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"net/http"
+	"net/url"
 
 	"github.com/CloudNativeAI/modctl/pkg/storage"
 
@@ -58,9 +60,23 @@ func (b *backend) Pull(ctx context.Context, target string, opts ...Option) error
 		return fmt.Errorf("failed to create credential store: %w", err)
 	}
 
+	// create the http client.
+	httpClient := http.DefaultClient
+	if options.proxy != "" {
+		proxyURL, err := url.Parse(options.proxy)
+		if err != nil {
+			return fmt.Errorf("failed to parse the proxy URL: %w", err)
+		}
+
+		httpClient.Transport = &http.Transport{
+			Proxy: http.ProxyURL(proxyURL),
+		}
+	}
+
 	src.Client = &auth.Client{
 		Cache:      auth.NewCache(),
 		Credential: credentials.Credential(credStore),
+		Client:     httpClient,
 	}
 
 	if options.plainHTTP {

pkg/config/pull.go

@@ -18,10 +18,12 @@ package config
 
 type Pull struct {
 	PlainHTTP bool
+	Proxy     string
 }
 
 func NewPull() *Pull {
 	return &Pull{
 		PlainHTTP: false,
+		Proxy:     "",
 	}
 }