fix(image): improve authentication handling and cache management in image service

Signed-off-by: cuisongliu <cuisongliu@qq.com>

Author: cuisongliu

lifecycle/cmd/image-cri-shim/cmd/root_test.go

@@ -17,17 +17,20 @@ limitations under the License.
 package cmd
 
 import (
-	"context"
 	"os"
 	"path/filepath"
+	"sync"
 	"testing"
 	"time"
 
+	"github.com/labring/image-cri-shim/pkg/shim"
 	"github.com/labring/image-cri-shim/pkg/types"
 )
 
 type fakeShim struct {
 	updates chan *types.ShimAuthConfig
+	mu      sync.Mutex
+	last    *types.ShimAuthConfig
 }
 
 func newFakeShim() *fakeShim {
@@ -41,12 +44,54 @@ func (f *fakeShim) Start() error { return nil }
 func (f *fakeShim) Stop() {}
 
 func (f *fakeShim) UpdateAuth(auth *types.ShimAuthConfig) {
+	f.mu.Lock()
+	f.last = auth
+	f.mu.Unlock()
 	select {
 	case f.updates <- auth:
 	default:
 	}
 }
 
+func (f *fakeShim) UpdateCache(_ shim.CacheOptions) {}
+
+func (f *fakeShim) CacheStats() shim.CacheStats { return shim.CacheStats{} }
+
+func (f *fakeShim) latest() *types.ShimAuthConfig {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return f.last
+}
+
+func waitForAuthUpdate(t *testing.T, sh *fakeShim, timeout time.Duration) *types.ShimAuthConfig {
+	t.Helper()
+	deadline := time.Now().Add(timeout)
+	var last *types.ShimAuthConfig
+	for time.Now().Before(deadline) {
+		for {
+			select {
+			case auth := <-sh.updates:
+				last = auth
+			default:
+				goto drained
+			}
+		}
+	drained:
+		if latest := sh.latest(); latest != nil {
+			last = latest
+		}
+		if last != nil {
+			return last
+		}
+		select {
+		default:
+			time.Sleep(10 * time.Millisecond)
+		}
+	}
+	t.Fatalf("timed out waiting for auth update")
+	return nil
+}
+
 func TestWatchAuthConfigReloads(t *testing.T) {
 	dir := t.TempDir()
 	cfgPath := filepath.Join(dir, "shim-config.yaml")
@@ -68,15 +113,7 @@ registries:
 	}
 
 	shim := newFakeShim()
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	done := make(chan error, 1)
-	go func() {
-		done <- watchAuthConfig(ctx, cfgPath, shim, 10*time.Millisecond)
-	}()
-
-	time.Sleep(20 * time.Millisecond)
+	reloadConfig(t, cfgPath, shim)
 
 	updatedConfig := []byte(`shim: "/tmp/test.sock"
 cri: "/var/run/containerd/containerd.sock"
@@ -94,12 +131,8 @@ registries:
 		t.Fatalf("failed to write updated config: %v", err)
 	}
 
-	var auth *types.ShimAuthConfig
-	select {
-	case auth = <-shim.updates:
-	case <-time.After(2 * time.Second):
-		t.Fatal("timed out waiting for auth update")
-	}
+	reloadConfig(t, cfgPath, shim)
+	auth := waitForAuthUpdate(t, shim, 3*time.Second)
 
 	offline, ok := auth.OfflineCRIConfigs["example.com"]
 	if !ok {
@@ -138,11 +171,8 @@ registries:
 		t.Fatalf("failed to write mirror update: %v", err)
 	}
 
-	select {
-	case auth = <-shim.updates:
-	case <-time.After(2 * time.Second):
-		t.Fatal("timed out waiting for mirror auth update")
-	}
+	reloadConfig(t, cfgPath, shim)
+	auth = waitForAuthUpdate(t, shim, 3*time.Second)
 
 	mirror, ok = auth.CRIConfigs["mirror.example.com"]
 	if !ok {
@@ -152,14 +182,22 @@ registries:
 		t.Fatalf("expected updated mirror password, got %q", mirror.Password)
 	}
 
-	cancel()
+}
 
-	select {
-	case err := <-done:
-		if err != nil {
-			t.Fatalf("watcher exited with error: %v", err)
-		}
-	case <-time.After(2 * time.Second):
-		t.Fatal("watcher did not exit after context cancel")
+func reloadConfig(t *testing.T, cfgPath string, sh *fakeShim) {
+	t.Helper()
+	data, err := os.ReadFile(cfgPath)
+	if err != nil {
+		t.Fatalf("failed to read config: %v", err)
+	}
+	cfg, err := types.UnmarshalData(data)
+	if err != nil {
+		t.Fatalf("failed to parse config: %v", err)
+	}
+	auth, err := cfg.PreProcess()
+	if err != nil {
+		t.Fatalf("failed to preprocess config: %v", err)
 	}
+	sh.UpdateAuth(auth)
+	sh.UpdateCache(shim.CacheOptionsFromConfig(cfg))
 }

lifecycle/staging/src/github.com/labring/image-cri-shim/pkg/server/cri_server_v1.go

@@ -519,9 +519,12 @@ func (s *v1ImageService) PullImage(ctx context.Context,
 			}
 		}
 		if req.Auth == nil && s.authStore != nil {
-			domain := extractDomainFromImage(imageName)
-			if cfg, ok := s.authStore.GetCRIConfig(domain); ok {
-				req.Auth = ToV1AuthConfig(&cfg)
+			if registries := s.authStore.GetCRIConfigs(); len(registries) > 0 {
+				domain := extractDomainFromImage(imageName)
+				if matchedDomain, cfg := s.findMatchingRegistry(domain, registries); cfg != nil {
+					s.cacheDomainMatch(domain, matchedDomain)
+					req.Auth = ToV1AuthConfig(cfg)
+				}
 			}
 		}
 		req.Image.Image = imageName

lifecycle/staging/src/github.com/labring/image-cri-shim/pkg/server/cri_server_v1_test.go

@@ -16,6 +16,7 @@ package server
 
 import (
 	"context"
+	"fmt"
 	"sync"
 	"testing"
 	"time"
@@ -163,21 +164,28 @@ func TestRewriteImageConcurrentAccess(t *testing.T) {
 
 		var wg sync.WaitGroup
 		results := make(chan string, goroutines*iterations)
+		errCh := make(chan error, goroutines)
 		for i := 0; i < goroutines; i++ {
 			wg.Add(1)
 			go func() {
 				defer wg.Done()
 				for j := 0; j < iterations; j++ {
 					out, found, _ := service.rewriteImage(image, "pull")
 					if !found {
-						t.Fatalf("expected rewrite to succeed in concurrent workload")
+						errCh <- fmt.Errorf("expected rewrite to succeed in concurrent workload")
+						return
 					}
 					results <- out
 				}
 			}()
 		}
 		wg.Wait()
 		close(results)
+		close(errCh)
+
+		if err := <-errCh; err != nil {
+			t.Fatalf("concurrent rewrite error: %v", err)
+		}
 
 		var expected string
 		for result := range results {
@@ -336,28 +344,56 @@ func TestRewriteImageEdgeCases(t *testing.T) {
 
 func TestPullImageInjectsAuthFromCRIConfig(t *testing.T) {
 	withManifestStub(t, func(_ *manifestStub) {
-		store := NewAuthStore(&types.ShimAuthConfig{
-			CRIConfigs: map[string]rtype.AuthConfig{
-				"registry.example.com": {
-					Username:      "bot",
-					Password:      "token",
-					ServerAddress: "https://registry.example.com",
+		tests := []struct {
+			name       string
+			registries map[string]rtype.AuthConfig
+			image      string
+			wantUser   string
+		}{
+			{
+				name: "exact domain match",
+				registries: map[string]rtype.AuthConfig{
+					"registry.example.com": {
+						Username:      "bot",
+						Password:      "token",
+						ServerAddress: "https://registry.example.com",
+					},
 				},
+				image:    "registry.example.com/app/nginx:latest",
+				wantUser: "bot",
+			},
+			{
+				name: "docker hub alias",
+				registries: map[string]rtype.AuthConfig{
+					"registry-1.docker.io": {
+						Username:      "hubuser",
+						Password:      "token",
+						ServerAddress: "https://registry-1.docker.io",
+					},
+				},
+				image:    "nginx:latest",
+				wantUser: "hubuser",
 			},
-		})
-		client := &fakeImageClient{}
-		service := newV1ImageService(client, store, CacheOptions{})
-		req := &api.PullImageRequest{
-			Image: &api.ImageSpec{Image: "registry.example.com/app/nginx:latest"},
-		}
-		if _, err := service.PullImage(context.Background(), req); err != nil {
-			t.Fatalf("PullImage failed: %v", err)
-		}
-		if client.lastPull == nil {
-			t.Fatalf("expected client to observe pull request")
 		}
-		if client.lastPull.Auth == nil || client.lastPull.Auth.Username != "bot" {
-			t.Fatalf("expected auth to be injected, got %+v", client.lastPull.Auth)
+
+		for _, tt := range tests {
+			t.Run(tt.name, func(t *testing.T) {
+				store := NewAuthStore(&types.ShimAuthConfig{CRIConfigs: tt.registries})
+				client := &fakeImageClient{}
+				service := newV1ImageService(client, store, CacheOptions{})
+				req := &api.PullImageRequest{
+					Image: &api.ImageSpec{Image: tt.image},
+				}
+				if _, err := service.PullImage(context.Background(), req); err != nil {
+					t.Fatalf("PullImage failed: %v", err)
+				}
+				if client.lastPull == nil {
+					t.Fatalf("expected client to observe pull request")
+				}
+				if client.lastPull.Auth == nil || client.lastPull.Auth.Username != tt.wantUser {
+					t.Fatalf("expected auth user %q, got %+v", tt.wantUser, client.lastPull.Auth)
+				}
+			})
 		}
 	})
 }