// Copyright 2018 Google LLC All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tarball_test

import (
	"archive/tar"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
	"testing"

	"github.com/google/go-containerregistry/pkg/name"
	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/compare"
	"github.com/google/go-containerregistry/pkg/v1/empty"
	"github.com/google/go-containerregistry/pkg/v1/mutate"
	"github.com/google/go-containerregistry/pkg/v1/random"
	"github.com/google/go-containerregistry/pkg/v1/tarball"
	"github.com/google/go-containerregistry/pkg/v1/types"
	"github.com/google/go-containerregistry/pkg/v1/validate"
)

func TestWrite(t *testing.T) {
	// Make a tempfile for tarball writes.
	fp, err := os.CreateTemp("", "")
	if err != nil {
		t.Fatalf("Error creating temp file.")
	}
	t.Log(fp.Name())
	defer fp.Close()
	defer os.Remove(fp.Name())

	// Make a random image
	randImage, err := random.Image(256, 8)
	if err != nil {
		t.Fatalf("Error creating random image.")
	}
	tag, err := name.NewTag("gcr.io/foo/bar:latest", name.StrictValidation)
	if err != nil {
		t.Fatalf("Error creating test tag.")
	}
	if err := tarball.WriteToFile(fp.Name(), tag, randImage); err != nil {
		t.Fatalf("Unexpected error writing tarball: %v", err)
	}

	// Make sure the image is valid and can be loaded.
	// Load it both by nil and by its name.
	for _, it := range []*name.Tag{nil, &tag} {
		tarImage, err := tarball.ImageFromPath(fp.Name(), it)
		if err != nil {
			t.Fatalf("Unexpected error reading tarball: %v", err)
		}

		if err := validate.Image(tarImage); err != nil {
			t.Errorf("validate.Image: %v", err)
		}

		if err := compare.Images(randImage, tarImage); err != nil {
			t.Errorf("compare.Images: %v", err)
		}
	}

	// Try loading a different tag, it should error.
	fakeTag, err := name.NewTag("gcr.io/notthistag:latest", name.StrictValidation)
	if err != nil {
		t.Fatalf("Error generating tag: %v", err)
	}
	if _, err := tarball.ImageFromPath(fp.Name(), &fakeTag); err == nil {
		t.Errorf("Expected error loading tag %v from image", fakeTag)
	}
}

func TestMultiWriteSameImage(t *testing.T) {
	// Make a tempfile for tarball writes.
	fp, err := os.CreateTemp("", "")
	if err != nil {
		t.Fatalf("Error creating temp file.")
	}
	t.Log(fp.Name())
	defer fp.Close()
	defer os.Remove(fp.Name())

	// Make a random image
	randImage, err := random.Image(256, 8)
	if err != nil {
		t.Fatalf("Error creating random image.")
	}

	// Make two tags that point to the random image above.
	tag1, err := name.NewTag("gcr.io/foo/bar:latest", name.StrictValidation)
	if err != nil {
		t.Fatalf("Error creating test tag1.")
	}
	tag2, err := name.NewTag("gcr.io/baz/bat:latest", name.StrictValidation)
	if err != nil {
		t.Fatalf("Error creating test tag2.")
	}
	dig3, err := name.NewDigest("gcr.io/baz/baz@sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", name.StrictValidation)
	if err != nil {
		t.Fatalf("Error creating test dig3.")
	}
	refToImage := make(map[name.Reference]v1.Image)
	refToImage[tag1] = randImage
	refToImage[tag2] = randImage
	refToImage[dig3] = randImage

	// Write the images with both tags to the tarball
	if err := tarball.MultiRefWriteToFile(fp.Name(), refToImage); err != nil {
		t.Fatalf("Unexpected error writing tarball: %v", err)
	}
	for ref := range refToImage {
		tag, ok := ref.(name.Tag)
		if !ok {
			continue
		}

		tarImage, err := tarball.ImageFromPath(fp.Name(), &tag)
		if err != nil {
			t.Fatalf("Unexpected error reading tarball: %v", err)
		}

		if err := validate.Image(tarImage); err != nil {
			t.Errorf("validate.Image: %v", err)
		}

		if err := compare.Images(randImage, tarImage); err != nil {
			t.Errorf("compare.Images: %v", err)
		}
	}
}

func TestMultiWriteDifferentImages(t *testing.T) {
	// Make a tempfile for tarball writes.
	fp, err := os.CreateTemp("", "")
	if err != nil {
		t.Fatalf("Error creating temp file.")
	}
	t.Log(fp.Name())
	defer fp.Close()
	defer os.Remove(fp.Name())

	// Make a random image
	randImage1, err := random.Image(256, 8)
	if err != nil {
		t.Fatalf("Error creating random image 1.")
	}

	// Make another random image
	randImage2, err := random.Image(256, 8)
	if err != nil {
		t.Fatalf("Error creating random image 2.")
	}

	// Make another random image
	randImage3, err := random.Image(256, 8)
	if err != nil {
		t.Fatalf("Error creating random image 3.")
	}

	// Create two tags, one pointing to each image created.
	tag1, err := name.NewTag("gcr.io/foo/bar:latest", name.StrictValidation)
	if err != nil {
		t.Fatalf("Error creating test tag1.")
	}
	tag2, err := name.NewTag("gcr.io/baz/bat:latest", name.StrictValidation)
	if err != nil {
		t.Fatalf("Error creating test tag2.")
	}
	dig3, err := name.NewDigest("gcr.io/baz/baz@sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", name.StrictValidation)
	if err != nil {
		t.Fatalf("Error creating test dig3.")
	}
	refToImage := make(map[name.Reference]v1.Image)
	refToImage[tag1] = randImage1
	refToImage[tag2] = randImage2
	refToImage[dig3] = randImage3

	// Write both images to the tarball.
	if err := tarball.MultiRefWriteToFile(fp.Name(), refToImage); err != nil {
		t.Fatalf("Unexpected error writing tarball: %v", err)
	}
	for ref, img := range refToImage {
		tag, ok := ref.(name.Tag)
		if !ok {
			continue
		}

		tarImage, err := tarball.ImageFromPath(fp.Name(), &tag)
		if err != nil {
			t.Fatalf("Unexpected error reading tarball: %v", err)
		}

		if err := validate.Image(tarImage); err != nil {
			t.Errorf("validate.Image: %v", err)
		}

		if err := compare.Images(img, tarImage); err != nil {
			t.Errorf("compare.Images: %v", err)
		}
	}
}

func TestWriteForeignLayers(t *testing.T) {
	// Make a tempfile for tarball writes.
	fp, err := os.CreateTemp("", "")
	if err != nil {
		t.Fatalf("Error creating temp file.")
	}
	t.Log(fp.Name())
	defer fp.Close()
	defer os.Remove(fp.Name())

	// Make a random image
	randImage, err := random.Image(256, 1)
	if err != nil {
		t.Fatalf("Error creating random image.")
	}
	tag, err := name.NewTag("gcr.io/foo/bar:latest", name.StrictValidation)
	if err != nil {
		t.Fatalf("Error creating test tag.")
	}
	randLayer, err := random.Layer(512, types.DockerForeignLayer)
	if err != nil {
		t.Fatalf("random.Layer: %v", err)
	}
	img, err := mutate.Append(randImage, mutate.Addendum{
		Layer: randLayer,
		URLs: []string{
			"example.com",
		},
	})
	if err != nil {
		t.Fatal(err)
	}
	if err := tarball.WriteToFile(fp.Name(), tag, img); err != nil {
		t.Fatalf("Unexpected error writing tarball: %v", err)
	}

	tarImage, err := tarball.ImageFromPath(fp.Name(), &tag)
	if err != nil {
		t.Fatalf("Unexpected error reading tarball: %v", err)
	}

	if err := validate.Image(tarImage); err != nil {
		t.Fatalf("validate.Image(): %v", err)
	}

	m, err := tarImage.Manifest()
	if err != nil {
		t.Fatal(err)
	}

	if got, want := m.Layers[1].MediaType, types.DockerForeignLayer; got != want {
		t.Errorf("Wrong MediaType: %s != %s", got, want)
	}
	if got, want := m.Layers[1].URLs[0], "example.com"; got != want {
		t.Errorf("Wrong URLs: %s != %s", got, want)
	}
}

func TestWriteSharedLayers(t *testing.T) {
	// Make a tempfile for tarball writes.
	fp, err := os.CreateTemp("", "")
	if err != nil {
		t.Fatalf("Error creating temp file.")
	}
	t.Log(fp.Name())
	defer fp.Close()
	defer os.Remove(fp.Name())

	// Make a random image
	randImage, err := random.Image(256, 1)
	if err != nil {
		t.Fatalf("Error creating random image.")
	}
	tag1, err := name.NewTag("gcr.io/foo/bar:latest", name.StrictValidation)
	if err != nil {
		t.Fatalf("Error creating test tag1.")
	}
	tag2, err := name.NewTag("gcr.io/baz/bat:latest", name.StrictValidation)
	if err != nil {
		t.Fatalf("Error creating test tag2.")
	}
	randLayer, err := random.Layer(512, types.DockerLayer)
	if err != nil {
		t.Fatalf("random.Layer: %v", err)
	}
	mutatedImage, err := mutate.Append(randImage, mutate.Addendum{
		Layer: randLayer,
	})
	if err != nil {
		t.Fatal(err)
	}
	refToImage := make(map[name.Reference]v1.Image)
	refToImage[tag1] = randImage
	refToImage[tag2] = mutatedImage

	// Write the images with both tags to the tarball
	if err := tarball.MultiRefWriteToFile(fp.Name(), refToImage); err != nil {
		t.Fatalf("Unexpected error writing tarball: %v", err)
	}
	for ref := range refToImage {
		tag, ok := ref.(name.Tag)
		if !ok {
			continue
		}

		tarImage, err := tarball.ImageFromPath(fp.Name(), &tag)
		if err != nil {
			t.Fatalf("Unexpected error reading tarball: %v", err)
		}

		if err := validate.Image(tarImage); err != nil {
			t.Errorf("validate.Image: %v", err)
		}

		if err := compare.Images(refToImage[tag], tarImage); err != nil {
			t.Errorf("compare.Images: %v", err)
		}
	}
	_, err = fp.Seek(0, io.SeekStart)
	if err != nil {
		t.Fatalf("Seek to start of file: %v", err)
	}
	layers, err := randImage.Layers()
	if err != nil {
		t.Fatalf("Get image layers: %v", err)
	}
	layers = append(layers, randLayer)
	wantDigests := make(map[string]struct{})
	for _, layer := range layers {
		d, err := layer.Digest()
		if err != nil {
			t.Fatalf("Get layer digest: %v", err)
		}
		wantDigests[d.Hex] = struct{}{}
	}

	const layerFileSuffix = ".tar.gz"
	r := tar.NewReader(fp)
	for {
		hdr, err := r.Next()
		if err != nil {
			if errors.Is(err, io.EOF) {
				break
			}
			t.Fatalf("Get tar header: %v", err)
		}
		if strings.HasSuffix(hdr.Name, layerFileSuffix) {
			hex := hdr.Name[:len(hdr.Name)-len(layerFileSuffix)]
			if _, ok := wantDigests[hex]; ok {
				delete(wantDigests, hex)
			} else {
				t.Errorf("Found unwanted layer with digest %q", hex)
			}
		}
	}
	if len(wantDigests) != 0 {
		for hex := range wantDigests {
			t.Errorf("Expected to find layer with digest %q but it didn't exist", hex)
		}
	}
}

func TestComputeManifest(t *testing.T) {
	var randomTag, mutatedTag = "ubuntu", "gcr.io/baz/bat:latest"

	// https://github.com/google/go-containerregistry/issues/890
	randomTagWritten := "ubuntu:latest"

	// Make a random image
	randImage, err := random.Image(256, 1)
	if err != nil {
		t.Fatalf("Error creating random image.")
	}
	randConfig, err := randImage.ConfigName()
	if err != nil {
		t.Fatalf("error getting random image config: %v", err)
	}
	tag1, err := name.NewTag(randomTag)
	if err != nil {
		t.Fatalf("Error creating test tag1.")
	}
	tag2, err := name.NewTag(mutatedTag, name.StrictValidation)
	if err != nil {
		t.Fatalf("Error creating test tag2.")
	}
	randLayer, err := random.Layer(512, types.DockerLayer)
	if err != nil {
		t.Fatalf("random.Layer: %v", err)
	}
	mutatedImage, err := mutate.Append(randImage, mutate.Addendum{
		Layer: randLayer,
	})
	if err != nil {
		t.Fatal(err)
	}
	mutatedConfig, err := mutatedImage.ConfigName()
	if err != nil {
		t.Fatalf("error getting mutated image config: %v", err)
	}
	randomLayersHashes, err := getLayersHashes(randImage)
	if err != nil {
		t.Fatalf("error getting random image hashes: %v", err)
	}
	randomLayersFilenames := getLayersFilenames(randomLayersHashes)

	mutatedLayersHashes, err := getLayersHashes(mutatedImage)
	if err != nil {
		t.Fatalf("error getting mutated image hashes: %v", err)
	}
	mutatedLayersFilenames := getLayersFilenames(mutatedLayersHashes)

	refToImage := make(map[name.Reference]v1.Image)
	refToImage[tag1] = randImage
	refToImage[tag2] = mutatedImage

	// calculate the manifest
	m, err := tarball.ComputeManifest(refToImage)
	if err != nil {
		t.Fatalf("Unexpected error calculating manifest: %v", err)
	}
	// the order of these two is based on the repo tags
	// so mutated "gcr.io/baz/bat:latest" is before random "gcr.io/foo/bar:latest"
	expected := []tarball.Descriptor{
		{
			Config:   mutatedConfig.String(),
			RepoTags: []string{mutatedTag},
			Layers:   mutatedLayersFilenames,
		},
		{
			Config:   randConfig.String(),
			RepoTags: []string{randomTagWritten},
			Layers:   randomLayersFilenames,
		},
	}
	if len(m) != len(expected) {
		t.Fatalf("mismatched manifest lengths: actual %d, expected %d", len(m), len(expected))
	}
	mBytes, err := json.Marshal(m)
	if err != nil {
		t.Fatalf("unable to marshall actual manifest to json: %v", err)
	}
	eBytes, err := json.Marshal(expected)
	if err != nil {
		t.Fatalf("unable to marshall expected manifest to json: %v", err)
	}
	if !bytes.Equal(mBytes, eBytes) {
		t.Errorf("mismatched manifests.\nActual: %s\nExpected: %s", string(mBytes), string(eBytes))
	}
}

func TestComputeManifest_FailsOnNoRefs(t *testing.T) {
	_, err := tarball.ComputeManifest(nil)
	if err == nil || !strings.Contains(err.Error(), "set of images is empty") {
		t.Error("expected calculateManifest to fail with nil input")
	}

	_, err = tarball.ComputeManifest(map[name.Reference]v1.Image{})
	if err == nil || !strings.Contains(err.Error(), "set of images is empty") {
		t.Error("expected calculateManifest to fail with empty input")
	}
}

func getLayersHashes(img v1.Image) ([]string, error) {
	hashes := []string{}
	layers, err := img.Layers()
	if err != nil {
		return nil, fmt.Errorf("error getting image layers: %w", err)
	}
	for i, l := range layers {
		hash, err := l.Digest()
		if err != nil {
			return nil, fmt.Errorf("error getting digest for layer %d: %w", i, err)
		}
		hashes = append(hashes, hash.Hex)
	}
	return hashes, nil
}

func getLayersFilenames(hashes []string) []string {
	filenames := make([]string, 0, len(hashes))
	for _, h := range hashes {
		filenames = append(filenames, fmt.Sprintf("%s.tar.gz", h))
	}
	return filenames
}

// TestWriteClosesLayerBeforeOpeningNext is a regression test for the deadlock
// that occurs when Write does not close each layer reader before opening the
// next one. Remote layer readers returned by remote.Get(...).Image() hold a
// pull-job slot until Close (see remote.WithJobs / pkg/v1/remote.limiter),
// so leaving readers open here deadlocks the write loop once defaultJobs
// readers are outstanding.
func TestWriteClosesLayerBeforeOpeningNext(t *testing.T) {
	tokens := make(chan struct{}, 4)
	img, err := random.Image(256, int64(cap(tokens)+1))
	if err != nil {
		t.Fatalf("random.Image: %v", err)
	}
	ref, err := name.NewTag("test/limited:latest")
	if err != nil {
		t.Fatalf("name.NewTag: %v", err)
	}
	if err := tarball.MultiRefWrite(map[name.Reference]v1.Image{ref: &tokenLimitedImage{Image: img, tokens: tokens}}, io.Discard); err != nil {
		t.Fatalf("tarball.MultiRefWrite: %v", err)
	}
	if got := len(tokens); got != 0 {
		t.Fatalf("open layer readers after write: got %d, want 0", got)
	}
}

// tokenLimitedImage wraps img so each layer's Compressed() reader takes one
// of cap(tokens) slots, modeling remote.WithJobs slot behavior. The slot is
// released only when Close() is called on the returned reader.
type tokenLimitedImage struct {
	v1.Image
	tokens chan struct{}
}

func (i *tokenLimitedImage) Layers() ([]v1.Layer, error) {
	ls, err := i.Image.Layers()
	if err != nil {
		return nil, err
	}
	out := make([]v1.Layer, len(ls))
	for j, l := range ls {
		out[j] = &tokenLimitedLayer{Layer: l, tokens: i.tokens}
	}
	return out, nil
}

type tokenLimitedLayer struct {
	v1.Layer
	tokens chan struct{}
}

func (l *tokenLimitedLayer) Compressed() (io.ReadCloser, error) {
	select {
	case l.tokens <- struct{}{}:
	default:
		return nil, errors.New("layer reader opened before previous readers were closed")
	}
	rc, err := l.Layer.Compressed()
	if err != nil {
		<-l.tokens
		return nil, err
	}
	return &tokenReleasingReadCloser{ReadCloser: rc, release: func() { <-l.tokens }}, nil
}

type tokenReleasingReadCloser struct {
	io.ReadCloser
	release func()
	closed  bool
}

func (rc *tokenReleasingReadCloser) Close() error {
	err := rc.ReadCloser.Close()
	if !rc.closed {
		rc.closed = true
		rc.release()
	}
	return err
}

func TestWriteDynamicExtensions(t *testing.T) {
	// Make a tempfile for tarball writes.
	fp, err := os.CreateTemp("", "")
	if err != nil {
		t.Fatalf("Error creating temp file.")
	}
	t.Log(fp.Name())
	defer fp.Close()
	defer os.Remove(fp.Name())

	// Create layers with different media types
	gzipLayer, err := random.Layer(128, types.OCILayer)
	if err != nil {
		t.Fatalf("random.Layer (gzip): %v", err)
	}
	zstdLayer, err := random.Layer(128, types.OCILayerZStd)
	if err != nil {
		t.Fatalf("random.Layer (zstd): %v", err)
	}
	tarLayer, err := random.Layer(128, types.OCIUncompressedLayer)
	if err != nil {
		t.Fatalf("random.Layer (tar): %v", err)
	}

	gzipDigest, err := gzipLayer.Digest()
	if err != nil {
		t.Fatal(err)
	}
	zstdDigest, err := zstdLayer.Digest()
	if err != nil {
		t.Fatal(err)
	}
	tarDigest, err := tarLayer.Digest()
	if err != nil {
		t.Fatal(err)
	}

	// Append layers to an empty image
	img, err := mutate.Append(empty.Image, mutate.Addendum{Layer: gzipLayer}, mutate.Addendum{Layer: zstdLayer}, mutate.Addendum{Layer: tarLayer})
	if err != nil {
		t.Fatalf("mutate.Append: %v", err)
	}

	tag, err := name.NewTag("gcr.io/foo/bar:latest", name.StrictValidation)
	if err != nil {
		t.Fatalf("Error creating test tag: %v", err)
	}

	if err := tarball.WriteToFile(fp.Name(), tag, img); err != nil {
		t.Fatalf("Unexpected error writing tarball: %v", err)
	}

	// 1. Verify files in tar archive have correct names
	_, err = fp.Seek(0, io.SeekStart)
	if err != nil {
		t.Fatalf("Seek to start of file: %v", err)
	}

	expectedFiles := map[string]bool{
		gzipDigest.Hex + ".tar.gz":  false,
		zstdDigest.Hex + ".tar.zst": false,
		tarDigest.Hex + ".tar":      false,
	}

	var manifestContent []byte

	r := tar.NewReader(fp)
	for {
		hdr, err := r.Next()
		if err != nil {
			if errors.Is(err, io.EOF) {
				break
			}
			t.Fatalf("Get tar header: %v", err)
		}
		if _, ok := expectedFiles[hdr.Name]; ok {
			expectedFiles[hdr.Name] = true
		}
		if hdr.Name == "manifest.json" {
			manifestContent, err = io.ReadAll(r)
			if err != nil {
				t.Fatalf("reading manifest.json: %v", err)
			}
		}
	}

	for filename, found := range expectedFiles {
		if !found {
			t.Errorf("Expected to find file %q in tarball, but it was not found", filename)
		}
	}

	// 2. Verify manifest.json contains correct file names
	if len(manifestContent) == 0 {
		t.Fatal("manifest.json not found or empty")
	}

	var manifest []struct {
		Config   string   `json:"Config"`
		RepoTags []string `json:"RepoTags"`
		Layers   []string `json:"Layers"`
	}
	if err := json.Unmarshal(manifestContent, &manifest); err != nil {
		t.Fatalf("unmarshalling manifest: %v", err)
	}

	if len(manifest) != 1 {
		t.Fatalf("expected 1 entry in manifest, got %d", len(manifest))
	}

	mEntry := manifest[0]
	expectedLayers := []string{
		gzipDigest.Hex + ".tar.gz",
		zstdDigest.Hex + ".tar.zst",
		tarDigest.Hex + ".tar",
	}

	if len(mEntry.Layers) != len(expectedLayers) {
		t.Fatalf("expected %d layers in manifest, got %d", len(expectedLayers), len(mEntry.Layers))
	}

	for i, wantLayer := range expectedLayers {
		if got := mEntry.Layers[i]; got != wantLayer {
			t.Errorf("layer %d: got manifest layer name %q, want %q", i, got, wantLayer)
		}
	}
}
