at-container-registry/cmd/s3-test/main.go

// Command s3-test is a diagnostic tool that tests S3 connectivity using both
// AWS SDK v1 (used by distribution's storage driver) and AWS SDK v2 (used by
// ATCR's presigned URL service). It helps diagnose signature compatibility
// issues with S3-compatible storage providers.
package main

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	awsv1 "github.com/aws/aws-sdk-go/aws"
	credentialsv1 "github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	s3v1 "github.com/aws/aws-sdk-go/service/s3"

	awsv2 "github.com/aws/aws-sdk-go-v2/aws"
	configv2 "github.com/aws/aws-sdk-go-v2/config"
	credentialsv2 "github.com/aws/aws-sdk-go-v2/credentials"
	s3v2 "github.com/aws/aws-sdk-go-v2/service/s3"
)

func main() {
	var (
		envFile   = flag.String("env-file", "", "Load environment variables from file (KEY=VALUE format)")
		accessKey = flag.String("access-key", "", "S3 access key (env: AWS_ACCESS_KEY_ID)")
		secretKey = flag.String("secret-key", "", "S3 secret key (env: AWS_SECRET_ACCESS_KEY)")
		region    = flag.String("region", "", "S3 region (env: S3_REGION)")
		bucket    = flag.String("bucket", "", "S3 bucket name (env: S3_BUCKET)")
		endpoint  = flag.String("endpoint", "", "S3 endpoint URL (env: S3_ENDPOINT)")
		pullZone  = flag.String("pull-zone", "", "CDN pull zone URL for presigned reads (env: PULL_ZONE)")
		prefix    = flag.String("prefix", "docker/registry/v2/blobs", "Key prefix for list operations")
		verbose   = flag.Bool("verbose", false, "Enable SDK debug signing logs")
	)
	flag.Parse()

	// Load env file first, then let flags and real env vars override
	if *envFile != "" {
		if err := loadEnvFile(*envFile); err != nil {
			fmt.Fprintf(os.Stderr, "Error loading env file: %v\n", err)
			os.Exit(1)
		}
	}

	// Resolve: flag > env var > default
	if *accessKey == "" {
		*accessKey = os.Getenv("AWS_ACCESS_KEY_ID")
	}
	if *secretKey == "" {
		*secretKey = os.Getenv("AWS_SECRET_ACCESS_KEY")
	}
	if *region == "" {
		*region = envOr("S3_REGION", "us-east-1")
	}
	if *bucket == "" {
		*bucket = os.Getenv("S3_BUCKET")
	}
	if *endpoint == "" {
		*endpoint = os.Getenv("S3_ENDPOINT")
	}
	if *pullZone == "" {
		*pullZone = os.Getenv("PULL_ZONE")
	}

	if *accessKey == "" || *secretKey == "" || *bucket == "" {
		fmt.Fprintln(os.Stderr, "Usage: s3-test [--env-file FILE] [--access-key KEY] [--secret-key KEY] [--bucket BUCKET] [--endpoint URL] [--region REGION] [--prefix PREFIX] [--verbose]")
		fmt.Fprintln(os.Stderr, "Env vars: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, S3_BUCKET, S3_REGION, S3_ENDPOINT")
		os.Exit(1)
	}

	fmt.Println("S3 Connectivity Diagnostic")
	fmt.Println("==========================")
	fmt.Printf("Endpoint:  %s\n", valueOr(*endpoint, "(default AWS)"))
	fmt.Printf("Pull Zone: %s\n", valueOr(*pullZone, "(none)"))
	fmt.Printf("Region:    %s\n", *region)
	fmt.Printf("AccessKey: %s...%s (%d chars)\n", (*accessKey)[:3], (*accessKey)[len(*accessKey)-3:], len(*accessKey))
	fmt.Printf("SecretKey: %s...%s (%d chars)\n", (*secretKey)[:3], (*secretKey)[len(*secretKey)-3:], len(*secretKey))
	fmt.Printf("Bucket:    %s\n", *bucket)
	fmt.Printf("Prefix:    %s\n", *prefix)
	fmt.Println()

	ctx := context.Background()
	results := make([]result, 0, 6)

	// Build SDK v1 client (SigV4) — matches distribution driver's New()
	v1Client := buildV1Client(*accessKey, *secretKey, *region, *endpoint, *verbose)

	// Test 1: SDK v1 SigV4 HeadBucket
	results = append(results, runTest("SDK v1 / SigV4 / HeadBucket", func() error {
		_, err := v1Client.HeadBucketWithContext(ctx, &s3v1.HeadBucketInput{
			Bucket: awsv1.String(*bucket),
		})
		return err
	}))

	// Test 2: SDK v1 SigV4 ListObjectsV2
	results = append(results, runTest("SDK v1 / SigV4 / ListObjectsV2", func() error {
		_, err := v1Client.ListObjectsV2WithContext(ctx, &s3v1.ListObjectsV2Input{
			Bucket:  awsv1.String(*bucket),
			Prefix:  awsv1.String(*prefix),
			MaxKeys: awsv1.Int64(5),
		})
		return err
	}))

	// Test 3: SDK v1 SigV4 ListObjectsV2Pages (paginated, matches doWalk)
	results = append(results, runTest("SDK v1 / SigV4 / ListObjectsV2Pages", func() error {
		return v1Client.ListObjectsV2PagesWithContext(ctx, &s3v1.ListObjectsV2Input{
			Bucket:  awsv1.String(*bucket),
			Prefix:  awsv1.String(*prefix),
			MaxKeys: awsv1.Int64(5),
		}, func(page *s3v1.ListObjectsV2Output, lastPage bool) bool {
			return false // stop after first page
		})
	}))

	// Build SDK v2 client — matches NewS3Service()
	v2Client := buildV2Client(ctx, *accessKey, *secretKey, *region, *endpoint)

	// Test 5: SDK v2 SigV4 HeadBucket
	results = append(results, runTest("SDK v2 / SigV4 / HeadBucket", func() error {
		_, err := v2Client.HeadBucket(ctx, &s3v2.HeadBucketInput{
			Bucket: awsv2.String(*bucket),
		})
		return err
	}))

	// Test 6: SDK v2 SigV4 ListObjectsV2
	results = append(results, runTest("SDK v2 / SigV4 / ListObjectsV2", func() error {
		_, err := v2Client.ListObjectsV2(ctx, &s3v2.ListObjectsV2Input{
			Bucket:  awsv2.String(*bucket),
			Prefix:  awsv2.String(*prefix),
			MaxKeys: awsv2.Int32(5),
		})
		return err
	}))

	// Find a real object key for GetObject / presigned URL tests
	var testKey string
	listOut, err := v2Client.ListObjectsV2(ctx, &s3v2.ListObjectsV2Input{
		Bucket:  awsv2.String(*bucket),
		Prefix:  awsv2.String(*prefix),
		MaxKeys: awsv2.Int32(1),
	})
	if err == nil && len(listOut.Contents) > 0 {
		testKey = *listOut.Contents[0].Key
	}

	if testKey == "" {
		fmt.Printf("\n  (Skipping GetObject/Presigned tests — no objects found under prefix %q)\n", *prefix)
	} else {
		fmt.Printf("\n  Test object: %s\n\n", testKey)

		// Test 7: SDK v1 GetObject (HEAD only)
		results = append(results, runTest("SDK v1 / SigV4 / HeadObject", func() error {
			_, err := v1Client.HeadObjectWithContext(ctx, &s3v1.HeadObjectInput{
				Bucket: awsv1.String(*bucket),
				Key:    awsv1.String(testKey),
			})
			return err
		}))

		// Test 8: SDK v2 GetObject (HEAD only)
		results = append(results, runTest("SDK v2 / SigV4 / HeadObject", func() error {
			_, err := v2Client.HeadObject(ctx, &s3v2.HeadObjectInput{
				Bucket: awsv2.String(*bucket),
				Key:    awsv2.String(testKey),
			})
			return err
		}))

		// Test 9: SDK v2 Presigned GET URL (generate + fetch)
		presignClient := s3v2.NewPresignClient(v2Client)
		results = append(results, runTest("SDK v2 / Presigned GET URL", func() error {
			presigned, err := presignClient.PresignGetObject(ctx, &s3v2.GetObjectInput{
				Bucket: awsv2.String(*bucket),
				Key:    awsv2.String(testKey),
			}, func(opts *s3v2.PresignOptions) {
				opts.Expires = 5 * time.Minute
			})
			if err != nil {
				return fmt.Errorf("presign: %w", err)
			}
			if *verbose {
				// Show host + query params (no path to avoid leaking key structure)
				u, _ := url.Parse(presigned.URL)
				fmt.Printf("\n    Presigned host: %s\n", u.Host)
				fmt.Printf("    Signed headers: %s\n", presigned.SignedHeader)
			}
			resp, err := http.Get(presigned.URL)
			if err != nil {
				return fmt.Errorf("fetch: %w", err)
			}
			body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
			resp.Body.Close()
			if resp.StatusCode != 200 {
				return fmt.Errorf("presigned URL returned %d: %s", resp.StatusCode, string(body))
			}
			return nil
		}))

		// Pull zone presigned tests — sign against real endpoint, swap host to pull zone
		if *pullZone != "" {
			results = append(results, runTest("SDK v2 / Presigned GET via Pull Zone", func() error {
				presigned, err := presignClient.PresignGetObject(ctx, &s3v2.GetObjectInput{
					Bucket: awsv2.String(*bucket),
					Key:    awsv2.String(testKey),
				}, func(opts *s3v2.PresignOptions) {
					opts.Expires = 5 * time.Minute
				})
				if err != nil {
					return fmt.Errorf("presign: %w", err)
				}
				pzURL := swapHost(presigned.URL, *pullZone)
				if *verbose {
					fmt.Printf("\n    Signed against: %s\n", presigned.URL[:40]+"...")
					fmt.Printf("    Fetching from:  %s\n", pzURL[:40]+"...")
				}
				resp, err := http.Get(pzURL)
				if err != nil {
					return fmt.Errorf("fetch: %w", err)
				}
				body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
				resp.Body.Close()
				if resp.StatusCode != 200 {
					return fmt.Errorf("pull zone GET returned %d: %s", resp.StatusCode, string(body))
				}
				return nil
			}))

		}

		// Test 10: SDK v2 Presigned PUT URL (generate + upload empty)
		results = append(results, runTest("SDK v2 / Presigned PUT URL", func() error {
			putKey := *prefix + "/_s3-test-probe"
			presigned, err := presignClient.PresignPutObject(ctx, &s3v2.PutObjectInput{
				Bucket: awsv2.String(*bucket),
				Key:    awsv2.String(putKey),
			}, func(opts *s3v2.PresignOptions) {
				opts.Expires = 5 * time.Minute
			})
			if err != nil {
				return fmt.Errorf("presign: %w", err)
			}
			req, err := http.NewRequestWithContext(ctx, http.MethodPut, presigned.URL, strings.NewReader(""))
			if err != nil {
				return fmt.Errorf("build request: %w", err)
			}
			req.Header.Set("Content-Length", "0")
			resp, err := http.DefaultClient.Do(req)
			if err != nil {
				return fmt.Errorf("fetch: %w", err)
			}
			resp.Body.Close()
			if resp.StatusCode != 200 {
				return fmt.Errorf("presigned PUT returned %d", resp.StatusCode)
			}
			// Clean up
			_, _ = v2Client.DeleteObject(ctx, &s3v2.DeleteObjectInput{
				Bucket: awsv2.String(*bucket),
				Key:    awsv2.String(putKey),
			})
			return nil
		}))
	}

	// Print summary
	fmt.Println()
	fmt.Println("Summary")
	fmt.Println("=======")

	allPass := true
	for _, r := range results {
		status := "PASS"
		if !r.ok {
			status = "FAIL"
			allPass = false
		}
		fmt.Printf("  [%s] %s (%s)\n", status, r.name, r.duration.Round(time.Millisecond))
		if !r.ok {
			fmt.Printf("         Error: %s\n", r.err)
		}
	}

	fmt.Println()
	if allPass {
		fmt.Println("Diagnosis: All tests passed. S3 connectivity is working with both SDKs.")
	} else {
		fmt.Println("Diagnosis: Some tests failed. Review errors above.")
	}
}

type result struct {
	name     string
	ok       bool
	err      error
	duration time.Duration
}

func runTest(name string, fn func() error) result {
	fmt.Printf("  Testing: %s ... ", name)
	start := time.Now()
	err := fn()
	d := time.Since(start)
	if err != nil {
		fmt.Printf("FAIL (%s)\n", d.Round(time.Millisecond))
		return result{name: name, ok: false, err: err, duration: d}
	}
	fmt.Printf("PASS (%s)\n", d.Round(time.Millisecond))
	return result{name: name, ok: true, duration: d}
}

func loadEnvFile(path string) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		line = strings.TrimPrefix(line, "export ")
		k, v, ok := strings.Cut(line, "=")
		if !ok {
			continue
		}
		v = strings.Trim(v, `"'`)
		os.Setenv(strings.TrimSpace(k), strings.TrimSpace(v))
	}
	return scanner.Err()
}

func envOr(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}

func swapHost(presignedURL, pullZone string) string {
	parsed, err := url.Parse(presignedURL)
	if err != nil {
		return presignedURL
	}
	pz, err := url.Parse(pullZone)
	if err != nil {
		return presignedURL
	}
	parsed.Scheme = pz.Scheme
	parsed.Host = pz.Host
	return parsed.String()
}

func valueOr(s, fallback string) string {
	if s == "" {
		return fallback
	}
	return s
}

// buildV1Client constructs an SDK v1 S3 client identically to
// distribution/distribution's s3-aws driver New() function.
func buildV1Client(accessKey, secretKey, region, endpoint string, verbose bool) *s3v1.S3 {
	awsConfig := awsv1.NewConfig()

	if verbose {
		awsConfig.WithLogLevel(awsv1.LogDebugWithSigning)
	}

	awsConfig.WithCredentials(credentialsv1.NewStaticCredentials(accessKey, secretKey, ""))
	awsConfig.WithRegion(region)

	if endpoint != "" {
		awsConfig.WithEndpoint(endpoint)
		awsConfig.WithS3ForcePathStyle(true)
	}

	sess, err := session.NewSession(awsConfig)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to create SDK v1 session: %v\n", err)
		os.Exit(1)
	}

	return s3v1.New(sess)
}

// buildV2Client constructs an SDK v2 S3 client identically to
// ATCR's NewS3Service() in pkg/s3/types.go.
func buildV2Client(ctx context.Context, accessKey, secretKey, region, endpoint string) *s3v2.Client {
	cfg, err := configv2.LoadDefaultConfig(ctx,
		configv2.WithRegion(region),
		configv2.WithCredentialsProvider(
			credentialsv2.NewStaticCredentialsProvider(accessKey, secretKey, ""),
		),
	)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to load SDK v2 config: %v\n", err)
		os.Exit(1)
	}

	return s3v2.NewFromConfig(cfg, func(o *s3v2.Options) {
		if endpoint != "" {
			o.BaseEndpoint = awsv2.String(endpoint)
			o.UsePathStyle = true
		}
	})
}