From 44b33b44aa1445bad0bcb5852c7a1606e30620f7 Mon Sep 17 00:00:00 2001
From: "Guillaume J. Charmes" <guillaume.charmes@dotcloud.com>
Date: Wed, 8 May 2013 22:37:33 -0700
Subject: [PATCH] Factorize the checksums functions

---
 graph.go    | 24 ++++++++++++++++++++++++
 image.go    | 37 ++++++++++++++++---------------------
 registry.go | 42 ++++++++++++++++++++++++++----------------
 3 files changed, 66 insertions(+), 37 deletions(-)

diff --git a/graph.go b/graph.go
index 24b8975dbd..5bf8623c62 100644
--- a/graph.go
+++ b/graph.go
@@ -1,6 +1,7 @@
 package docker
 
 import (
+	"encoding/json"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -300,3 +301,26 @@ func (graph *Graph) Heads() (map[string]*Image, error) {
 func (graph *Graph) imageRoot(id string) string {
 	return path.Join(graph.Root, id)
 }
+
+func (graph *Graph) getStoredChecksums() (map[string]string, error) {
+	checksums := make(map[string]string)
+	// FIXME: Store the checksum in memory
+
+	if checksumDict, err := ioutil.ReadFile(path.Join(graph.Root, "checksums")); err == nil {
+		if err := json.Unmarshal(checksumDict, &checksums); err != nil {
+			return nil, err
+		}
+	}
+	return checksums, nil
+}
+
+func (graph *Graph) storeChecksums(checksums map[string]string) error {
+	checksumJson, err := json.Marshal(checksums)
+	if err != nil {
+		return err
+	}
+	if err := ioutil.WriteFile(path.Join(graph.Root, "checksums"), checksumJson, 0600); err != nil {
+		return err
+	}
+	return nil
+}
diff --git a/image.go b/image.go
index 55be632880..413d95673b 100644
--- a/image.go
+++ b/image.go
@@ -295,16 +295,12 @@ func (img *Image) Checksum() (string, error) {
 		return "", err
 	}
 
-	checksumDictPth := path.Join(root, "..", "..", "checksums")
-	checksums := make(map[string]string)
-
-	if checksumDict, err := ioutil.ReadFile(checksumDictPth); err == nil {
-		if err := json.Unmarshal(checksumDict, &checksums); err != nil {
-			return "", err
-		}
-		if checksum, ok := checksums[img.Id]; ok {
-			return checksum, nil
-		}
+	checksums, err := img.graph.getStoredChecksums()
+	if err != nil {
+		return "", err
+	}
+	if checksum, ok := checksums[img.Id]; ok {
+		return checksum, nil
 	}
 
 	layer, err := img.layer()
@@ -343,24 +339,23 @@ func (img *Image) Checksum() (string, error) {
 	if _, err := io.Copy(h, layerData); err != nil {
 		return "", err
 	}
-
 	hash := "sha256:" + hex.EncodeToString(h.Sum(nil))
-	checksums[img.Id] = hash
 
 	// Reload the json file to make sure not to overwrite faster sums
 	img.graph.lockSumFile.Lock()
 	defer img.graph.lockSumFile.Unlock()
-	if checksumDict, err := ioutil.ReadFile(checksumDictPth); err == nil {
-		if err := json.Unmarshal(checksumDict, &checksums); err != nil {
-			return "", err
-		}
-	}
-	checksumJson, err := json.Marshal(checksums)
+
+	checksums, err = img.graph.getStoredChecksums()
 	if err != nil {
+		return "", err
+	}
+
+	checksums[img.Id] = hash
+
+	// Dump the checksums to disc
+	if err := img.graph.storeChecksums(checksums); err != nil {
 		return hash, err
 	}
-	if err := ioutil.WriteFile(checksumDictPth, checksumJson, 0600); err != nil {
-		return hash, err
-	}
+
 	return hash, nil
 }
diff --git a/registry.go b/registry.go
index 004befa5df..40127ad41a 100644
--- a/registry.go
+++ b/registry.go
@@ -578,6 +578,24 @@ func (graph *Graph) PushRepository(stdout io.Writer, remote string, localRepo Re
 	return filteredRepo, nil
 }
 
+// Retrieve the checksum of an image
+// Priority:
+// - Check on the stored checksums
+// - Check if the archive is exists, if it does not, ask the registry
+// - If the archive does exists, process the checksum from it
+// - If the archive does not exists and not found on registry, process checksum from layer
+func (graph *Graph) getChecksum(imageId string) (string, error) {
+	img, err := graph.Get(imageId)
+	if err != nil {
+		return "", err
+	}
+	checksum, err := img.Checksum()
+	if err != nil {
+		return "", err
+	}
+	return checksum, nil
+}
+
 type ImgListJson struct {
 	Id       string `json:"id"`
 	Checksum string `json:"checksum,omitempty"`
@@ -592,7 +610,14 @@ func (graph *Graph) PushRepository(stdout io.Writer, remote string, localRepo Re
 	var imgList []*ImgListJson
 
 	for _, id := range localRepo {
-		imgList = append(imgList, &ImgListJson{Id: id})
+		checksum, err := graph.getChecksum(id)
+		if err != nil {
+			return err
+		}
+		imgList = append(imgList, &ImgListJson{
+			Id:       id,
+			Checksum: checksum,
+		})
 	}
 
 	imgListJson, err := json.Marshal(imgList)
@@ -666,21 +691,6 @@ func (graph *Graph) PushRepository(stdout io.Writer, remote string, localRepo Re
 		}
 	}
 
-	for _, elem := range imgList {
-		img, err := graph.Get(elem.Id)
-		if err != nil {
-			return err
-		}
-		if elem.Checksum, err = img.Checksum(); err != nil {
-			return err
-		}
-	}
-	imgListJson, err = json.Marshal(imgList)
-	if err != nil {
-		return err
-	}
-	Debugf("json sent: %s\n", imgListJson)
-
 	req2, err := http.NewRequest("PUT", INDEX_ENDPOINT+"/repositories/"+remote+"/images", bytes.NewReader(imgListJson))
 	if err != nil {
 		return err