From 4f0355fcf75d4a98fe5b4d002bda9553e6c2055e Mon Sep 17 00:00:00 2001
From: Eric Muellenbach <emuellenbach@gmx.de>
Date: Mon, 14 Aug 2017 12:34:56 +0200
Subject: [PATCH] Add Support for Registry V2 Manifest Format in Clair module

---
 clair/vulns.go | 48 ++++++++++++++++++++++++++++++++++++++++++++++-
 main.go        | 51 ++++++++++++++++++++++++++++++++++++++------------
 2 files changed, 86 insertions(+), 13 deletions(-)

diff --git a/clair/vulns.go b/clair/vulns.go
index 2a3da094..3edaffff 100644
--- a/clair/vulns.go
+++ b/clair/vulns.go
@@ -6,6 +6,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/docker/distribution"
 	"github.com/docker/distribution/manifest/schema1"
 	"github.com/jessfraz/reg/registry"
 )
@@ -79,7 +80,7 @@ func (c *Clair) Vulnerabilities(r *registry.Registry, repo, tag string, m schema
 	return report, nil
 }
 
-// NewClairLayer will form a layer struct required for a clar scan
+// NewClairLayer will form a layer struct required for a clair scan
 func (c *Clair) NewClairLayer(r *registry.Registry, image string, fsLayers []schema1.FSLayer, index int) (*Layer, error) {
 	var parentName string
 	if index < len(fsLayers)-1 {
@@ -123,3 +124,48 @@ func (c *Clair) NewClairLayer(r *registry.Registry, image string, fsLayers []sch
 		Headers:    h,
 	}, nil
 }
+
+// NewClairLayerV2 will form a layer struct required for a clair scan
+func (c *Clair) NewClairLayerV2(r *registry.Registry, image string, fsLayers []distribution.Descriptor, index int) (*Layer, error) {
+	var parentName string
+	if index < len(fsLayers)-1 {
+		parentName = fsLayers[index+1].Digest.String()
+	}
+
+	// form the path
+	p := strings.Join([]string{r.URL, "v2", image, "blobs", fsLayers[index].Digest.String()}, "/")
+
+	useBasicAuth := false
+
+	// get the token
+	token, err := r.Token(p)
+	if err != nil {
+		// if we get an error here of type: malformed auth challenge header: 'Basic realm="Registry Realm"'
+		// we need to use basic auth for the registry
+		if !strings.Contains(err.Error(), `malformed auth challenge header: 'Basic realm="Registry Realm"'`) {
+			return nil, err
+		}
+		useBasicAuth = true
+	}
+
+	h := make(map[string]string)
+	if token != "" && !useBasicAuth {
+		h = map[string]string{
+			"Authorization": fmt.Sprintf("Bearer %s", token),
+		}
+	}
+
+	if useBasicAuth {
+		h = map[string]string{
+			"Authorization": fmt.Sprintf("Basic %s", base64.StdEncoding.EncodeToString([]byte(r.Username+":"+r.Password))),
+		}
+	}
+
+	return &Layer{
+		Name:       fsLayers[index].Digest.String(),
+		Path:       p,
+		ParentName: parentName,
+		Format:     "Docker",
+		Headers:    h,
+	}, nil
+}
diff --git a/main.go b/main.go
index 16ce95cb..af01aa0e 100644
--- a/main.go
+++ b/main.go
@@ -10,7 +10,8 @@ import (
 	"text/tabwriter"
 
 	"github.com/Sirupsen/logrus"
-	"github.com/docker/distribution/manifest/schema1"
+	"github.com/docker/distribution"
+	"github.com/docker/distribution/manifest/schema2"
 	"github.com/docker/docker/api/types"
 	"github.com/jessfraz/reg/clair"
 	"github.com/jessfraz/reg/registry"
@@ -267,20 +268,46 @@ func main() {
 				// FIXME use clair.Vulnerabilities
 
 				// get the manifest
-				m, err := r.ManifestV1(repo, ref)
+				m2, err := r.Manifest(repo, ref)
+
 				if err != nil {
 					return err
 				}
 
-				// filter out the empty layers
-				var filteredLayers []schema1.FSLayer
-				for _, layer := range m.FSLayers {
-					if !clair.IsEmptyLayer(layer.BlobSum) {
-						filteredLayers = append(filteredLayers, layer)
+				mf, ok := m2.(schema2.Manifest)
+
+				var filteredLayers []distribution.Descriptor
+
+				if ok {
+
+					for _, layer := range mf.Layers {
+						if !clair.IsEmptyLayer(layer.Digest) {
+							filteredLayers = append(filteredLayers, layer)
+						}
 					}
+
+				} else {
+
+					fmt.Println("Couldn't retrieve manifest V2, fallback to v1")
+					m, err := r.ManifestV1(repo, ref)
+					if err != nil {
+						return err
+					}
+
+					for _, layer := range m.FSLayers {
+						if !clair.IsEmptyLayer(layer.BlobSum) {
+
+							newLayer := distribution.Descriptor{
+								Digest: layer.BlobSum,
+							}
+
+							filteredLayers = append(filteredLayers, newLayer)
+						}
+					}
+
 				}
-				m.FSLayers = filteredLayers
-				if len(m.FSLayers) == 0 {
+
+				if len(filteredLayers) == 0 {
 					fmt.Printf("No need to analyse image %s:%s as there is no non-emtpy layer", repo, ref)
 					return nil
 				}
@@ -291,9 +318,9 @@ func main() {
 					return err
 				}
 
-				for i := len(m.FSLayers) - 1; i >= 0; i-- {
+				for i := len(filteredLayers) - 1; i >= 0; i-- {
 					// form the clair layer
-					l, err := cr.NewClairLayer(r, repo, m.FSLayers, i)
+					l, err := cr.NewClairLayerV2(r, repo, filteredLayers, i)
 					if err != nil {
 						return err
 					}
@@ -304,7 +331,7 @@ func main() {
 					}
 				}
 
-				vl, err := cr.GetLayer(m.FSLayers[0].BlobSum.String(), false, true)
+				vl, err := cr.GetLayer(filteredLayers[0].Digest.String(), false, true)
 				if err != nil {
 					return err
 				}