blob: c19841ebddc7ffeb2c502c0e006177052370fac7 [file] [log] [blame]
Mohammed Naser274554f2025-01-26 12:50:42 -05001package main
2
3import (
4 "context"
5 "fmt"
6 "os"
7 "strings"
8 "sync"
9
10 "github.com/containers/image/v5/docker"
11 "github.com/containers/image/v5/manifest"
12 "github.com/containers/image/v5/types"
13 "github.com/goccy/go-yaml"
14 "github.com/goccy/go-yaml/ast"
15 "github.com/goccy/go-yaml/parser"
16 "github.com/opencontainers/go-digest"
17 log "github.com/sirupsen/logrus"
18 "golang.org/x/sync/singleflight"
19)
20
21var digestGroup singleflight.Group
22
23// GetImageDigest fetches the digest for a given image reference.
24func GetImageDigest(ctx context.Context, reference string) (digest.Digest, error) {
25 ref, err := docker.ParseReference(reference)
26 if err != nil {
27 return "", fmt.Errorf("error parsing reference '%s': %w", reference, err)
28 }
29
30 sysCtx := &types.SystemContext{}
31 imageSource, err := ref.NewImageSource(ctx, sysCtx)
32 if err != nil {
33 return "", fmt.Errorf("error creating image source for '%s': %w", reference, err)
34 }
35 defer imageSource.Close()
36
37 rawManifest, _, err := imageSource.GetManifest(ctx, nil)
38 if err != nil {
39 return "", fmt.Errorf("error getting manifest for '%s': %w", reference, err)
40 }
41
42 dgst, err := manifest.Digest(rawManifest)
43 if err != nil {
44 return "", fmt.Errorf("error getting digest for '%s': %w", reference, err)
45 }
46
47 return dgst, nil
48}
49
50// GetImageNameToPull normalizes the image name by replacing variables and adding necessary prefixes.
51func GetImageNameToPull(image string, release string) string {
52 // Replace Jinja2 variables with actual values
53 image = strings.ReplaceAll(image, "{{ atmosphere_image_prefix }}", "")
54 image = strings.ReplaceAll(image, "{{ atmosphere_release }}", release)
55
56 // Add mirror if the image is not hosted with us
57 if !strings.HasPrefix(image, "registry.atmosphere.dev") {
58 image = fmt.Sprintf("harbor.atmosphere.dev/%s", image)
59 }
60
61 // Switch out of the CDN since we are in CI
62 if strings.HasPrefix(image, "registry.atmosphere.dev") {
63 image = strings.ReplaceAll(image, "registry.atmosphere.dev", "harbor.atmosphere.dev")
64 }
65
66 return image
67}
68
69// AppendDigestToImage appends the digest to the original image reference.
70func AppendDigestToImage(image string, dgst digest.Digest) string {
71 if strings.Contains(image, "@") {
72 // Replace existing digest if present
73 parts := strings.Split(image, "@")
74 return parts[0] + "@" + dgst.String()
75 }
76 // Append digest
77 return image + "@" + dgst.String()
78}
79
80func main() {
81 varsFilePath := "roles/defaults/vars/main.yml"
82
83 file, err := parser.ParseFile(varsFilePath, parser.ParseComments)
84 if err != nil {
85 log.WithError(err).Fatal("error parsing yaml file")
86 }
87
88 if len(file.Docs) != 1 {
89 log.Fatal("expected exactly one yaml document")
90 }
91
92 doc := file.Docs[0]
93 body := doc.Body.(*ast.MappingNode)
94
95 var release string
96 var images *ast.MappingNode
97
98 for _, item := range body.Values {
99 switch item.Key.(*ast.StringNode).Value {
100 case "atmosphere_release":
101 release = item.Value.(*ast.StringNode).Value
102 case "_atmosphere_images":
103 images = item.Value.(*ast.MappingNode)
104 }
105 }
106
107 if release == "" {
108 log.Fatalf("atmosphere_release not found")
109 }
110
111 if images == nil {
112 log.Fatalf("_atmosphere_images not found")
113 }
114
115 type imageInfo struct {
116 Key string
117 Value string
118 Normalized string
119 Digest digest.Digest
120 }
121
122 var imageInfos []imageInfo
123 uniqueImages := make(map[string][]int)
124
125 for i, item := range images.Values {
126 normalized := GetImageNameToPull(item.Value.(*ast.StringNode).Value, release)
127 info := imageInfo{
128 Key: item.Key.(*ast.StringNode).Value,
129 Value: item.Value.(*ast.StringNode).Value,
130 Normalized: normalized,
131 }
132 imageInfos = append(imageInfos, info)
133 uniqueImages[normalized] = append(uniqueImages[normalized], i)
134 }
135
136 digestMap := make(map[string]digest.Digest)
137 var mapMutex sync.Mutex
138 var wg sync.WaitGroup
139
140 for normImg := range uniqueImages {
141 wg.Add(1)
142 go func(normImg string) {
143 defer wg.Done()
144
145 result, err, _ := digestGroup.Do(normImg, func() (interface{}, error) {
146 dgst, err := GetImageDigest(context.TODO(), "//"+normImg)
147 if err != nil {
148 return nil, err
149 }
150 return dgst, nil
151 })
152
153 if err != nil {
154 log.WithError(err).WithFields(log.Fields{
155 "image": normImg,
156 }).Error("Error fetching digest")
157 return
158 }
159
160 dgst := result.(digest.Digest)
161
162 mapMutex.Lock()
163 digestMap[normImg] = dgst
164 mapMutex.Unlock()
165
166 log.WithFields(log.Fields{
167 "image": normImg,
168 "digest": dgst,
169 }).Info("Fetched image digest")
170 }(normImg)
171 }
172
173 wg.Wait()
174
175 // Update the image references with digests
176 for normImg, indices := range uniqueImages {
177 dgst, exists := digestMap[normImg]
178 if !exists {
179 log.WithField("image", normImg).Error("Digest not found, skipping update")
180 continue
181 }
182 for _, idx := range indices {
183 updatedImage, err := yaml.ValueToNode(AppendDigestToImage(imageInfos[idx].Value, dgst))
184 if err != nil {
185 log.WithError(err).Fatal("error converting value to node")
186 }
187
188 images.Values[idx].Value = updatedImage
189 }
190 }
191
192 if err := os.WriteFile(varsFilePath, []byte(file.String()), 0644); err != nil {
193 log.WithError(err).Fatal("error writing updated yaml file")
194 }
195
196 log.Info("Successfully updated YAML file with image digests")
197}