Skip to content

Commit 33ad3e4

Browse files
ilopezlunailopezluna
committed
feat(registry): add support for custom default and insecure registry options
1 parent b51ca01 commit 33ad3e4

File tree

5 files changed

+76
-14
lines changed

5 files changed

+76
-14
lines changed

‎cmd/cli/commands/package.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -414,7 +414,7 @@ func newModelRunnerTarget(client *desktop.Client, tag string) (*modelRunnerTarge
414414
var err error
415415
// Normalize the tag to add default namespace (ai/) and tag (:latest) if missing
416416
normalizedTag := models.NormalizeModelName(tag)
417-
target.tag, err = name.NewTag(normalizedTag)
417+
target.tag, err = name.NewTag(normalizedTag, getDefaultRegistryOptions()...)
418418
if err != nil {
419419
return nil, fmt.Errorf("invalid tag: %w", err)
420420
}

‎cmd/cli/commands/tag.go‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package commands
22

33
import (
44
"fmt"
5+
"os"
56
"strings"
67

78
"github.com/docker/model-runner/cmd/cli/commands/completion"
@@ -11,6 +12,21 @@ import (
1112
"github.com/spf13/cobra"
1213
)
1314

15+
// getDefaultRegistryOptions returns name.Option slice with custom default registry
16+
// and insecure flag if the corresponding environment variables are set.
17+
// - DEFAULT_REGISTRY: Override the default registry (index.docker.io)
18+
// - INSECURE_REGISTRY: Set to "true" to allow HTTP connections
19+
func getDefaultRegistryOptions() []name.Option {
20+
var opts []name.Option
21+
if defaultReg := os.Getenv("DEFAULT_REGISTRY"); defaultReg != "" {
22+
opts = append(opts, name.WithDefaultRegistry(defaultReg))
23+
}
24+
if os.Getenv("INSECURE_REGISTRY") == "true" {
25+
opts = append(opts, name.Insecure)
26+
}
27+
return opts
28+
}
29+
1430
func newTagCmd() *cobra.Command {
1531
c := &cobra.Command{
1632
Use: "tag SOURCE TARGET",
@@ -42,7 +58,7 @@ func tagModel(cmd *cobra.Command, desktopClient *desktop.Client, source, target
4258
// Normalize target model name to add default org and tag if missing
4359
target = models.NormalizeModelName(target)
4460
// Ensure tag is valid
45-
tag, err := name.NewTag(target)
61+
tag, err := name.NewTag(target, getDefaultRegistryOptions()...)
4662
if err != nil {
4763
return fmt.Errorf("invalid tag: %w", err)
4864
}

‎pkg/distribution/internal/store/index.go‎

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,28 @@ import (
1010
"github.com/google/go-containerregistry/pkg/name"
1111
)
1212

13+
// getDefaultRegistryOptions returns name.Option slice with custom default registry
14+
// and insecure flag if the corresponding environment variables are set.
15+
// - DEFAULT_REGISTRY: Override the default registry (index.docker.io)
16+
// - INSECURE_REGISTRY: Set to "true" to allow HTTP connections
17+
func getDefaultRegistryOptions() []name.Option {
18+
var opts []name.Option
19+
if defaultReg := os.Getenv("DEFAULT_REGISTRY"); defaultReg != "" {
20+
opts = append(opts, name.WithDefaultRegistry(defaultReg))
21+
}
22+
if os.Getenv("INSECURE_REGISTRY") == "true" {
23+
opts = append(opts, name.Insecure)
24+
}
25+
return opts
26+
}
27+
1328
// Index represents the index of all models in the store
1429
type Index struct {
1530
Models []IndexEntry `json:"models"`
1631
}
1732

1833
func (i Index) Tag(reference string, tag string) (Index, error) {
19-
tagRef, err := name.NewTag(tag)
34+
tagRef, err := name.NewTag(tag, getDefaultRegistryOptions()...)
2035
if err != nil {
2136
return Index{}, fmt.Errorf("invalid tag: %w", err)
2237
}
@@ -39,7 +54,7 @@ func (i Index) Tag(reference string, tag string) (Index, error) {
3954
}
4055

4156
func (i Index) UnTag(tag string) (name.Tag, Index, error) {
42-
tagRef, err := name.NewTag(tag)
57+
tagRef, err := name.NewTag(tag, getDefaultRegistryOptions()...)
4358
if err != nil {
4459
return name.Tag{}, Index{}, err
4560
}
@@ -141,12 +156,12 @@ type IndexEntry struct {
141156
}
142157

143158
func (e IndexEntry) HasTag(tag string) bool {
144-
ref, err := name.NewTag(tag)
159+
ref, err := name.NewTag(tag, getDefaultRegistryOptions()...)
145160
if err != nil {
146161
return false
147162
}
148163
for _, t := range e.Tags {
149-
tr, err := name.ParseReference(t)
164+
tr, err := name.ParseReference(t, getDefaultRegistryOptions()...)
150165
if err != nil {
151166
continue
152167
}
@@ -159,7 +174,7 @@ func (e IndexEntry) HasTag(tag string) bool {
159174

160175
func (e IndexEntry) hasTag(tag name.Tag) bool {
161176
for _, t := range e.Tags {
162-
tr, err := name.ParseReference(t)
177+
tr, err := name.ParseReference(t, getDefaultRegistryOptions()...)
163178
if err != nil {
164179
continue
165180
}
@@ -174,7 +189,7 @@ func (e IndexEntry) MatchesReference(reference string) bool {
174189
if e.ID == reference {
175190
return true
176191
}
177-
ref, err := name.ParseReference(reference)
192+
ref, err := name.ParseReference(reference, getDefaultRegistryOptions()...)
178193
if err != nil {
179194
return false
180195
}
@@ -200,7 +215,7 @@ func (e IndexEntry) Tag(tag name.Tag) IndexEntry {
200215
func (e IndexEntry) UnTag(tag name.Tag) IndexEntry {
201216
var tags []string
202217
for i, t := range e.Tags {
203-
tr, err := name.ParseReference(t)
218+
tr, err := name.ParseReference(t, getDefaultRegistryOptions()...)
204219
if err != nil {
205220
continue
206221
}

‎pkg/distribution/registry/client.go‎

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"fmt"
66
"io"
77
"net/http"
8+
"os"
89
"strings"
910

1011
"github.com/google/go-containerregistry/pkg/authn"
@@ -21,6 +22,21 @@ const (
2122
DefaultUserAgent = "model-distribution"
2223
)
2324

25+
// getDefaultRegistryOptions returns name.Option slice with custom default registry
26+
// and insecure flag if the corresponding environment variables are set.
27+
// - DEFAULT_REGISTRY: Override the default registry (index.docker.io)
28+
// - INSECURE_REGISTRY: Set to "true" to allow HTTP connections
29+
func getDefaultRegistryOptions() []name.Option {
30+
var opts []name.Option
31+
if defaultReg := os.Getenv("DEFAULT_REGISTRY"); defaultReg != "" {
32+
opts = append(opts, name.WithDefaultRegistry(defaultReg))
33+
}
34+
if os.Getenv("INSECURE_REGISTRY") == "true" {
35+
opts = append(opts, name.Insecure)
36+
}
37+
return opts
38+
}
39+
2440
var (
2541
DefaultTransport = remote.DefaultTransport
2642
)
@@ -75,7 +91,7 @@ func NewClient(opts ...ClientOption) *Client {
7591

7692
func (c *Client) Model(ctx context.Context, reference string) (types.ModelArtifact, error) {
7793
// Parse the reference
78-
ref, err := name.ParseReference(reference)
94+
ref, err := name.ParseReference(reference, getDefaultRegistryOptions()...)
7995
if err != nil {
8096
return nil, NewReferenceError(reference, err)
8197
}
@@ -115,7 +131,7 @@ func (c *Client) Model(ctx context.Context, reference string) (types.ModelArtifa
115131

116132
func (c *Client) BlobURL(reference string, digest v1.Hash) (string, error) {
117133
// Parse the reference
118-
ref, err := name.ParseReference(reference)
134+
ref, err := name.ParseReference(reference, getDefaultRegistryOptions()...)
119135
if err != nil {
120136
return "", NewReferenceError(reference, err)
121137
}
@@ -129,7 +145,7 @@ func (c *Client) BlobURL(reference string, digest v1.Hash) (string, error) {
129145

130146
func (c *Client) BearerToken(ctx context.Context, reference string) (string, error) {
131147
// Parse the reference
132-
ref, err := name.ParseReference(reference)
148+
ref, err := name.ParseReference(reference, getDefaultRegistryOptions()...)
133149
if err != nil {
134150
return "", NewReferenceError(reference, err)
135151
}
@@ -165,7 +181,7 @@ type Target struct {
165181
}
166182

167183
func (c *Client) NewTarget(tag string) (*Target, error) {
168-
ref, err := name.NewTag(tag)
184+
ref, err := name.NewTag(tag, getDefaultRegistryOptions()...)
169185
if err != nil {
170186
return nil, fmt.Errorf("invalid tag: %q: %w", tag, err)
171187
}

‎pkg/metrics/metrics.go‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,21 @@ import (
1515
"github.com/sirupsen/logrus"
1616
)
1717

18+
// getDefaultRegistryOptions returns name.Option slice with custom default registry
19+
// and insecure flag if the corresponding environment variables are set.
20+
// - DEFAULT_REGISTRY: Override the default registry (index.docker.io)
21+
// - INSECURE_REGISTRY: Set to "true" to allow HTTP connections
22+
func getDefaultRegistryOptions() []name.Option {
23+
var opts []name.Option
24+
if defaultReg := os.Getenv("DEFAULT_REGISTRY"); defaultReg != "" {
25+
opts = append(opts, name.WithDefaultRegistry(defaultReg))
26+
}
27+
if os.Getenv("INSECURE_REGISTRY") == "true" {
28+
opts = append(opts, name.Insecure)
29+
}
30+
return opts
31+
}
32+
1833
type Tracker struct {
1934
doNotTrack bool
2035
transport http.RoundTripper
@@ -86,7 +101,7 @@ func (t *Tracker) trackModel(model types.Model, userAgent, action string) {
86101
}
87102
ua := strings.Join(parts, " ")
88103
for _, tag := range tags {
89-
ref, err := name.ParseReference(tag)
104+
ref, err := name.ParseReference(tag, getDefaultRegistryOptions()...)
90105
if err != nil {
91106
t.log.Errorf("Error parsing reference: %v\n", err)
92107
return

0 commit comments

Comments
 (0)