From e1b5341ebcab008433c24a9a73263f05974d8c7d Mon Sep 17 00:00:00 2001
From: Kohei Tokunaga <ktokunaga.mail@gmail.com>
Date: Thu, 18 Nov 2021 15:35:47 +0900
Subject: [PATCH] Make manifest detection stricter

Signed-off-by: Kohei Tokunaga <ktokunaga.mail@gmail.com>
---
 store/refs.go                   | 29 +++++++++++++++-----
 util/containerdutil/manifest.go | 48 +++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 6 deletions(-)

diff --git a/store/refs.go b/store/refs.go
index 66ca21f8f..fba122edc 100644
--- a/store/refs.go
+++ b/store/refs.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io/ioutil"
 	"os"
 	"path/filepath"
 	"sync"
@@ -32,6 +33,7 @@ import (
 	"github.com/containerd/containerd/remotes"
 	"github.com/containerd/containerd/remotes/docker"
 	"github.com/containerd/stargz-snapshotter/fs/source"
+	"github.com/containerd/stargz-snapshotter/util/containerdutil"
 	"github.com/containerd/stargz-snapshotter/util/lrucache"
 	digest "github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
@@ -250,10 +252,27 @@ func fetchManifestPlatform(ctx context.Context, fetcher remotes.Fetcher, desc oc
 	var manifest ocispec.Manifest
 	switch desc.MediaType {
 	case images.MediaTypeDockerSchema2Manifest, ocispec.MediaTypeImageManifest:
-		err = json.NewDecoder(r).Decode(&manifest)
+		p, err := ioutil.ReadAll(r)
+		if err != nil {
+			return ocispec.Manifest{}, err
+		}
+		if err := containerdutil.ValidateMediaType(p, desc.MediaType); err != nil {
+			return ocispec.Manifest{}, err
+		}
+		if err := json.Unmarshal(p, &manifest); err != nil {
+			return ocispec.Manifest{}, err
+		}
+		return manifest, nil
 	case images.MediaTypeDockerSchema2ManifestList, ocispec.MediaTypeImageIndex:
 		var index ocispec.Index
-		if err = json.NewDecoder(r).Decode(&index); err != nil {
+		p, err := ioutil.ReadAll(r)
+		if err != nil {
+			return ocispec.Manifest{}, err
+		}
+		if err := containerdutil.ValidateMediaType(p, desc.MediaType); err != nil {
+			return ocispec.Manifest{}, err
+		}
+		if err = json.Unmarshal(p, &index); err != nil {
 			return ocispec.Manifest{}, err
 		}
 		var target ocispec.Descriptor
@@ -273,9 +292,7 @@ func fetchManifestPlatform(ctx context.Context, fetcher remotes.Fetcher, desc oc
 		if !found {
 			return ocispec.Manifest{}, fmt.Errorf("no manifest found for platform")
 		}
-		manifest, err = fetchManifestPlatform(ctx, fetcher, target, platform)
-	default:
-		err = fmt.Errorf("unknown mediatype %q", desc.MediaType)
+		return fetchManifestPlatform(ctx, fetcher, target, platform)
 	}
-	return manifest, err
+	return ocispec.Manifest{}, fmt.Errorf("unknown mediatype %q", desc.MediaType)
 }
diff --git a/util/containerdutil/manifest.go b/util/containerdutil/manifest.go
index ec77017fa..4f9f991cc 100644
--- a/util/containerdutil/manifest.go
+++ b/util/containerdutil/manifest.go
@@ -19,6 +19,7 @@ package containerdutil
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"sort"
 
 	"github.com/containerd/containerd/content"
@@ -42,6 +43,9 @@ func ManifestDesc(ctx context.Context, provider content.Provider, image ocispec.
 			if err != nil {
 				return nil, err
 			}
+			if err := ValidateMediaType(p, desc.MediaType); err != nil {
+				return nil, err
+			}
 			var manifest ocispec.Manifest
 			if err := json.Unmarshal(p, &manifest); err != nil {
 				return nil, err
@@ -71,6 +75,9 @@ func ManifestDesc(ctx context.Context, provider content.Provider, image ocispec.
 			if err != nil {
 				return nil, err
 			}
+			if err := ValidateMediaType(p, desc.MediaType); err != nil {
+				return nil, err
+			}
 			var idx ocispec.Index
 			if err := json.Unmarshal(p, &idx); err != nil {
 				return nil, err
@@ -109,3 +116,44 @@ func ManifestDesc(ctx context.Context, provider content.Provider, image ocispec.
 	}
 	return m[0], nil
 }
+
+// Forked from github.com/containerd/containerd/images/image.go
+// commit: a776a27af54a803657d002e7574a4425b3949f56
+
+// unknownDocument represents a manifest, manifest list, or index that has not
+// yet been validated.
+type unknownDocument struct {
+	MediaType string          `json:"mediaType,omitempty"`
+	Config    json.RawMessage `json:"config,omitempty"`
+	Layers    json.RawMessage `json:"layers,omitempty"`
+	Manifests json.RawMessage `json:"manifests,omitempty"`
+	FSLayers  json.RawMessage `json:"fsLayers,omitempty"` // schema 1
+}
+
+// ValidateMediaType returns an error if the byte slice is invalid JSON or if
+// the media type identifies the blob as one format but it contains elements of
+// another format.
+func ValidateMediaType(b []byte, mt string) error {
+	var doc unknownDocument
+	if err := json.Unmarshal(b, &doc); err != nil {
+		return err
+	}
+	if len(doc.FSLayers) != 0 {
+		return fmt.Errorf("media-type: schema 1 not supported")
+	}
+	switch mt {
+	case images.MediaTypeDockerSchema2Manifest, ocispec.MediaTypeImageManifest:
+		if len(doc.Manifests) != 0 ||
+			doc.MediaType == images.MediaTypeDockerSchema2ManifestList ||
+			doc.MediaType == ocispec.MediaTypeImageIndex {
+			return fmt.Errorf("media-type: expected manifest but found index (%s)", mt)
+		}
+	case images.MediaTypeDockerSchema2ManifestList, ocispec.MediaTypeImageIndex:
+		if len(doc.Config) != 0 || len(doc.Layers) != 0 ||
+			doc.MediaType == images.MediaTypeDockerSchema2Manifest ||
+			doc.MediaType == ocispec.MediaTypeImageManifest {
+			return fmt.Errorf("media-type: expected index but found manifest (%s)", mt)
+		}
+	}
+	return nil
+}