upkg/tool/oci/push_test.go

package oci

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"gitea.loveuer.com/loveuer/upkg/tool"
)

func TestParseImageAddress(t *testing.T) {
	tests := []struct {
		name     string
		address  string
		wantReg  string
		wantRepo string
		wantTag  string
		wantErr  bool
	}{
		{
			name:     "valid address with tag",
			address:  "localhost:5000/myapp:latest",
			wantReg:  "localhost:5000",
			wantRepo: "myapp",
			wantTag:  "latest",
			wantErr:  false,
		},
		{
			name:     "valid address without tag",
			address:  "registry.example.com/library/nginx",
			wantReg:  "registry.example.com",
			wantRepo: "library/nginx",
			wantTag:  "latest",
			wantErr:  false,
		},
		{
			name:     "valid address with IP",
			address:  "192.168.1.1:5000/library/nginx:1.20",
			wantReg:  "192.168.1.1:5000",
			wantRepo: "library/nginx",
			wantTag:  "1.20",
			wantErr:  false,
		},
		{
			name:    "invalid address - no slash",
			address: "localhost:5000",
			wantErr: true,
		},
		{
			name:    "empty address",
			address: "",
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			gotReg, gotRepo, gotTag, err := parseImageAddress(tt.address)
			if (err != nil) != tt.wantErr {
				t.Errorf("parseImageAddress() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !tt.wantErr {
				if gotReg != tt.wantReg {
					t.Errorf("parseImageAddress() registry = %v, want %v", gotReg, tt.wantReg)
				}
				if gotRepo != tt.wantRepo {
					t.Errorf("parseImageAddress() repository = %v, want %v", gotRepo, tt.wantRepo)
				}
				if gotTag != tt.wantTag {
					t.Errorf("parseImageAddress() tag = %v, want %v", gotTag, tt.wantTag)
				}
			}
		})
	}
}

func TestComputeDigest(t *testing.T) {
	data := []byte("test data")
	digest := computeDigest(data)

	if !strings.HasPrefix(digest, "sha256:") {
		t.Errorf("computeDigest() should start with 'sha256:', got %s", digest)
	}

	// Test with same data should produce same digest
	digest2 := computeDigest(data)
	if digest != digest2 {
		t.Errorf("computeDigest() should produce same digest for same data, got %s and %s", digest, digest2)
	}

	// Test with different data should produce different digest
	differentData := []byte("different data")
	digest3 := computeDigest(differentData)
	if digest == digest3 {
		t.Errorf("computeDigest() should produce different digest for different data")
	}
}

func TestExtractImageFromTar(t *testing.T) {
	// Create a mock tar file with OCI image structure
	var buf bytes.Buffer
	tw := tar.NewWriter(&buf)

	// Create mock config
	configData := []byte(`{"architecture":"amd64","config":{"Env":["PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"]}}`)
	configDigest := computeDigest(configData)

	// Create mock layer
	layerData := []byte("mock layer data")
	layerDigest := computeDigest(layerData)

	// Create manifest.json
	dockerManifest := []struct {
		Config   string   `json:"Config"`
		RepoTags []string `json:"RepoTags"`
		Layers   []string `json:"Layers"`
	}{
		{
			Config:   "config.json",
			RepoTags: []string{"test:latest"},
			Layers:   []string{"layer.tar"},
		},
	}
	manifestData, _ := json.Marshal(dockerManifest)

	// Write files to tar
	tw.WriteHeader(&tar.Header{
		Name: "manifest.json",
		Size: int64(len(manifestData)),
		Mode: 0644,
	})
	tw.Write(manifestData)

	tw.WriteHeader(&tar.Header{
		Name: "config.json",
		Size: int64(len(configData)),
		Mode: 0644,
	})
	tw.Write(configData)

	tw.WriteHeader(&tar.Header{
		Name: "layer.tar",
		Size: int64(len(layerData)),
		Mode: 0644,
	})
	tw.Write(layerData)

	tw.Close()

	// Test extraction
	manifest, config, layers, err := extractImageFromTar(&buf)
	if err != nil {
		t.Fatalf("extractImageFromTar() error = %v", err)
	}

	if len(manifest) == 0 {
		t.Error("extractImageFromTar() should return non-empty manifest")
	}

	if config.digest != configDigest {
		t.Errorf("extractImageFromTar() config digest = %v, want %v", config.digest, configDigest)
	}

	if len(layers) != 1 {
		t.Errorf("extractImageFromTar() should return 1 layer, got %d", len(layers))
	}

	if layers[0].digest != layerDigest {
		t.Errorf("extractImageFromTar() layer digest = %v, want %v", layers[0].digest, layerDigest)
	}
}

func TestExtractImageFromTarMissingManifest(t *testing.T) {
	// Create a tar file without manifest.json
	var buf bytes.Buffer
	tw := tar.NewWriter(&buf)

	tw.WriteHeader(&tar.Header{
		Name: "somefile.txt",
		Size: 9,
		Mode: 0644,
	})
	tw.Write([]byte("test data"))
	tw.Close()

	// Test extraction should fail
	_, _, _, err := extractImageFromTar(bytes.NewReader(buf.Bytes()))
	if err == nil {
		t.Error("extractImageFromTar() should return error for missing manifest.json")
	}

	if !strings.Contains(err.Error(), "manifest.json not found") {
		t.Errorf("extractImageFromTar() should return manifest not found error, got %v", err)
	}
}

func TestWithPushOptions(t *testing.T) {
	opt := &ociUploadOpt{}

	// Test WithPushPlainHTTP
	WithPushPlainHTTP(true)(opt)
	if !opt.PlainHTTP {
		t.Error("WithPushPlainHTTP() should set PlainHTTP to true")
	}

	// Test WithPushSkipTLSVerify
	WithPushSkipTLSVerify(true)(opt)
	if !opt.SkipTLSVerify {
		t.Error("WithPushSkipTLSVerify() should set SkipTLSVerify to true")
	}

	// Test WithPushAuth
	WithPushAuth("user", "pass")(opt)
	if opt.Username != "user" || opt.Password != "pass" {
		t.Error("WithPushAuth() should set username and password")
	}
}

func TestUploadBlob(t *testing.T) {
	// Create a mock HTTP server
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.Method == http.MethodHead && strings.Contains(r.URL.Path, "/blobs/"):
			// Blob exists check
			w.WriteHeader(http.StatusOK)
		case r.Method == http.MethodPost && strings.Contains(r.URL.Path, "/blobs/uploads/"):
			// Start upload session
			w.Header().Set("Location", "/v2/test/repo/blobs/upload/session")
			w.WriteHeader(http.StatusAccepted)
		case r.Method == http.MethodPut && strings.Contains(r.URL.Path, "/blobs/upload/"):
			// Complete upload
			w.WriteHeader(http.StatusCreated)
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	}))
	defer server.Close()

	client := tool.NewClient(false, "")
	ctx := context.Background()

	data := []byte("test blob data")
	digest := computeDigest(data)

	err := uploadBlob(ctx, client, strings.TrimPrefix(server.URL, "http://"), "test/repo", data, digest, &ociUploadOpt{PlainHTTP: true})
	if err != nil {
		t.Errorf("uploadBlob() error = %v", err)
	}
}

func TestUploadBlobWithAuth(t *testing.T) {
	// Create a mock HTTP server with auth
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Check for basic auth
		username, password, ok := r.BasicAuth()
		if !ok || username != "testuser" || password != "testpass" {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}

		switch {
		case r.Method == http.MethodHead && strings.Contains(r.URL.Path, "/blobs/"):
			w.WriteHeader(http.StatusOK)
		case r.Method == http.MethodPost && strings.Contains(r.URL.Path, "/blobs/uploads/"):
			w.Header().Set("Location", "/v2/test/repo/blobs/upload/session")
			w.WriteHeader(http.StatusAccepted)
		case r.Method == http.MethodPut && strings.Contains(r.URL.Path, "/blobs/upload/"):
			w.WriteHeader(http.StatusCreated)
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	}))
	defer server.Close()

	client := tool.NewClient(false, "")
	ctx := context.Background()

	data := []byte("test blob data")
	digest := computeDigest(data)

	opt := &ociUploadOpt{
		Username:  "testuser",
		Password:  "testpass",
		PlainHTTP: true,
	}

	err := uploadBlob(ctx, client, strings.TrimPrefix(server.URL, "http://"), "test/repo", data, digest, opt)
	if err != nil {
		t.Errorf("uploadBlob() with auth error = %v", err)
	}
}

func TestUploadManifest(t *testing.T) {
	// Create a mock HTTP server
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodPut && strings.Contains(r.URL.Path, "/manifests/") {
			w.WriteHeader(http.StatusCreated)
		} else {
			w.WriteHeader(http.StatusNotFound)
		}
	}))
	defer server.Close()

	client := tool.NewClient(false, "")
	ctx := context.Background()

	manifest := []byte(`{"schemaVersion":2,"mediaType":"application/vnd.oci.image.manifest.v1+json"}`)

	err := uploadManifest(ctx, client, strings.TrimPrefix(server.URL, "http://"), "test/repo", "latest", manifest, &ociUploadOpt{PlainHTTP: true})
	if err != nil {
		t.Errorf("uploadManifest() error = %v", err)
	}
}

func TestPushImageGzipDetection(t *testing.T) {
	// Test that gzip detection works by creating a tar file with gzip compression
	// and verifying it can be processed through the gzip detection path

	// Create a tar file with OCI image structure
	var tarBuf bytes.Buffer
	tw := tar.NewWriter(&tarBuf)

	// Create mock config
	configData := []byte(`{"architecture":"amd64"}`)

	// Create mock layer
	layerData := []byte("mock layer data")

	// Create manifest.json
	dockerManifest := []struct {
		Config   string   `json:"Config"`
		RepoTags []string `json:"RepoTags"`
		Layers   []string `json:"Layers"`
	}{
		{
			Config:   "config.json",
			RepoTags: []string{"test:latest"},
			Layers:   []string{"layer.tar"},
		},
	}
	manifestData, _ := json.Marshal(dockerManifest)

	// Write files to tar
	tw.WriteHeader(&tar.Header{
		Name: "manifest.json",
		Size: int64(len(manifestData)),
		Mode: 0644,
	})
	tw.Write(manifestData)

	tw.WriteHeader(&tar.Header{
		Name: "config.json",
		Size: int64(len(configData)),
		Mode: 0644,
	})
	tw.Write(configData)

	tw.WriteHeader(&tar.Header{
		Name: "layer.tar",
		Size: int64(len(layerData)),
		Mode: 0644,
	})
	tw.Write(layerData)

	tw.Close()

	// Now gzip the tar data
	var gzippedTar bytes.Buffer
	gz := gzip.NewWriter(&gzippedTar)
	gz.Write(tarBuf.Bytes())
	gz.Close()

	// Create a mock server
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusNotFound)
	}))
	defer server.Close()

	ctx := context.Background()

	// Test with gzipped tar data (should fail at server, but gzip detection should work)
	err := PushImage(ctx, bytes.NewReader(gzippedTar.Bytes()), "localhost:5000/test:latest",
		WithPushPlainHTTP(true),
		WithPushSkipTLSVerify(true))

	// Should fail due to server not being a real registry, but not due to gzip detection
	if err == nil {
		t.Error("PushImage() should fail with mock server")
	}

	// The error should not be related to gzip detection
	if strings.Contains(err.Error(), "gzip reader failed") {
		t.Errorf("PushImage() should not fail with gzip reader error, got %v", err)
	}
}

func TestPushImageInvalidAddress(t *testing.T) {
	ctx := context.Background()
	data := bytes.NewReader([]byte("test data"))

	err := PushImage(ctx, data, "invalid-address", WithPushPlainHTTP(true))
	if err == nil {
		t.Error("PushImage() should return error for invalid address")
	}

	if !strings.Contains(err.Error(), "parse image address failed") {
		t.Errorf("PushImage() should return address parsing error, got %v", err)
	}
}

func TestPushImageOptions(t *testing.T) {
	// Test that options are applied correctly
	opt := &ociUploadOpt{}
	for _, fn := range []OCIUploadOpt{
		WithPushPlainHTTP(true),
		WithPushSkipTLSVerify(true),
		WithPushAuth("user", "pass"),
	} {
		fn(opt)
	}

	if !opt.PlainHTTP {
		t.Error("PlainHTTP option should be set to true")
	}
	if !opt.SkipTLSVerify {
		t.Error("SkipTLSVerify option should be set to true")
	}
	if opt.Username != "user" || opt.Password != "pass" {
		t.Error("Auth option should be set correctly")
	}
}

// Benchmark tests
func BenchmarkComputeDigest(b *testing.B) {
	data := make([]byte, 1024) // 1KB of data
	for i := 0; i < b.N; i++ {
		computeDigest(data)
	}
}

func BenchmarkParseImageAddress(b *testing.B) {
	address := "registry.example.com:5000/library/nginx:1.20"
	for i := 0; i < b.N; i++ {
		parseImageAddress(address)
	}
}