<commit_before>package webui

import (
	"bytes"
	"encoding/json"
	"github.com/go-macaron/binding"
	"github.com/toni-moreno/snmpcollector/pkg/data/impexp"
	"gopkg.in/macaron.v1"
	"io/ioutil"
	"mime/multipart"
	"os"
	"strconv"
	"time"
)

type UploadForm struct {
	AutoRename bool
	ExportFile *multipart.FileHeader
}

func NewImportExport(m *macaron.Macaron) error {

	bind := binding.Bind

	m.Group("/api/cfg/export", func() {
		m.Get("/:objtype/:id", reqSignedIn, ExportObject)
		m.Post("/:objtype/:id", reqSignedIn, bind(impexp.ExportInfo{}), ExportObjectToFile)

	})
	m.Group("/api/cfg/import", func() {
		m.Post("/", reqSignedIn, binding.MultipartForm(UploadForm{}), ImportDataFile)
	})
	return nil
}

/****************/
/*IMPORT*/
/*****************/

type ImportCheck struct {
	IsOk       bool
	Message    string
	Duplicated []*impexp.ExportObject
}

func ImportDataFile(ctx *Context, uf UploadForm) {
	if (UploadForm{}) == uf {
		log.Error("Error no data in expected struct")
		ctx.JSON(404, "Error no data in expected struct")
		return
	}
	log.Debugf("Uploaded data :%+v", uf)
	if uf.ExportFile == nil {
		ctx.JSON(404, "Error no file uploaded struct")
		return
	}
	log.Debugf("Uploaded File : %+v", uf)
	file, err := uf.ExportFile.Open()
	if err != nil {
		log.Warningf("Error on Open Uploaded File: %s", err)
		ctx.JSON(404, err.Error())
	}
	buf := new(bytes.Buffer)
	buf.ReadFrom(file)
	s := buf.String()
	log.Debug("FILE DATA: %s", s)
	ImportedData := impexp.ExportData{}
	if err := json.Unmarshal(buf.Bytes(), &ImportedData); err != nil {
		log.Errorf("Error in data to struct (json-unmarshal) procces: %s", err)
		ctx.JSON(404, err.Error())
	}
	log.Debugf("IMPORTED STRUCT %+v", ImportedData)

	arrays, err := ImportedData.ImportCheck()

	if err != nil && uf.AutoRename == false {
		ctx.JSON(404, &ImportCheck{IsOk: false, Message: err.Error(), Duplicated: arrays})
		return
	}
	if err != nil && uf.AutoRename == true {
		timestamp := time.Now().Unix()
		replaced := buf.Bytes()
		for _, v := range arrays {
			oldid := v.ObjectID
			newid := v.ObjectID + "_" + strconv.FormatInt(timestamp, 10)
			log.Debugf("replacind old ID %s for new one %s", oldid, newid)
			replaced = bytes.Replace(replaced, []byte(oldid), []byte(newid), -1)
		}
		ImportedData = impexp.ExportData{}
		if err = json.Unmarshal(replaced, &ImportedData); err != nil {
			log.Errorf("Error in data to struct (json-unmarshal) procces: %s", err)
			ctx.JSON(404, err.Error())
		}
		err = ImportedData.Import()
		if err != nil {
			log.Errorf("Some Error happened on import data: %s", err)
			ctx.JSON(404, err.Error())
		}
		ctx.JSON(200, &ImportCheck{IsOk: true, Message: "all duplicated ID's have been replaced", Duplicated: arrays})
		return
	}
	err = ImportedData.Import()
	if err != nil {
		log.Errorf("Some Error happened on import data: %s", err)
		ctx.JSON(404, err.Error())
	}
	ctx.JSON(200, &ImportCheck{IsOk: true, Message: "all objects have been  imported", Duplicated: nil})
}

/****************/
/*EXPORT*/
/****************/

func ExportObject(ctx *Context) {
	id := ctx.Params(":id")
	objtype := ctx.Params(":objtype")
	exp := impexp.NewExport(&impexp.ExportInfo{FileName: "autogenerated.txt", Description: "autogenerated"})
	err := exp.Export(objtype, id)
	if err != nil {
		log.Warningf("Error on get object array for type %s with ID %s  , error: %s", objtype, id, err)
		ctx.JSON(404, err.Error())
	} else {
		ctx.JSON(200, &exp)
	}
}

func ExportObjectToFile(ctx *Context, info impexp.ExportInfo) {

	id := ctx.Params(":id")
	objtype := ctx.Params(":objtype")
	exp := impexp.NewExport(&info)
	err := exp.Export(objtype, id)
	if err != nil {
		log.Errorf("Error on create Export Data with: %s", err)
		ctx.JSON(404, err.Error())
		return
	}
	outdata, err := json.Marshal(exp)
	if err != nil {
		log.Errorf("Error on Json data formatting: %s", err)
		ctx.JSON(404, err.Error())
		return
	}
	tmpfile, err := ioutil.TempFile("", "snmpcollector")
	if err != nil {
		log.Errorf("Error on create new temporal file: %s", err)
		ctx.JSON(404, err.Error())
		return
	}
	defer os.Remove(tmpfile.Name())
	if _, err := tmpfile.Write(outdata); err != nil {
		log.Errorf("Error on write data to temporal file: %s", tmpfile.Name())
		ctx.JSON(404, err.Error())
		return
	}
	if err := tmpfile.Close(); err != nil {
		log.Errorf("Error on close temporal file: %s", tmpfile.Name())
		ctx.JSON(404, err.Error())
		return
	}

	ctx.ServeFile(tmpfile.Name())
}
<commit_msg>added Filename Info to the ImportCheck struct<commit_after>package webui

import (
	"bytes"
	"encoding/json"
	"github.com/go-macaron/binding"
	"github.com/toni-moreno/snmpcollector/pkg/data/impexp"
	"gopkg.in/macaron.v1"
	"io/ioutil"
	"mime/multipart"
	"os"
	"strconv"
	"time"
)

type UploadForm struct {
	AutoRename bool
	ExportFile *multipart.FileHeader
}

func NewImportExport(m *macaron.Macaron) error {

	bind := binding.Bind

	m.Group("/api/cfg/export", func() {
		m.Get("/:objtype/:id", reqSignedIn, ExportObject)
		m.Post("/:objtype/:id", reqSignedIn, bind(impexp.ExportInfo{}), ExportObjectToFile)

	})
	m.Group("/api/cfg/import", func() {
		m.Post("/", reqSignedIn, binding.MultipartForm(UploadForm{}), ImportDataFile)
	})
	return nil
}

/****************/
/*IMPORT*/
/*****************/

type ImportCheck struct {
	Info       *impexp.ExportInfo
	IsOk       bool
	Message    string
	Duplicated []*impexp.ExportObject
}

func ImportDataFile(ctx *Context, uf UploadForm) {
	if (UploadForm{}) == uf {
		log.Error("Error no data in expected struct")
		ctx.JSON(404, "Error no data in expected struct")
		return
	}
	log.Debugf("Uploaded data :%+v", uf)
	if uf.ExportFile == nil {
		ctx.JSON(404, "Error no file uploaded struct")
		return
	}
	log.Debugf("Uploaded File : %+v", uf)
	file, err := uf.ExportFile.Open()
	if err != nil {
		log.Warningf("Error on Open Uploaded File: %s", err)
		ctx.JSON(404, err.Error())
	}
	buf := new(bytes.Buffer)
	buf.ReadFrom(file)
	s := buf.String()
	log.Debug("FILE DATA: %s", s)
	ImportedData := impexp.ExportData{}
	if err := json.Unmarshal(buf.Bytes(), &ImportedData); err != nil {
		log.Errorf("Error in data to struct (json-unmarshal) procces: %s", err)
		ctx.JSON(404, err.Error())
	}
	log.Debugf("IMPORTED STRUCT %+v", ImportedData)

	arrays, err := ImportedData.ImportCheck()

	if err != nil && uf.AutoRename == false {
		ctx.JSON(200, &ImportCheck{Info: ImportedData.Info, IsOk: false, Message: err.Error(), Duplicated: arrays})
		return
	}
	if err != nil && uf.AutoRename == true {
		timestamp := time.Now().Unix()
		replaced := buf.Bytes()
		for _, v := range arrays {
			oldid := v.ObjectID
			newid := v.ObjectID + "_" + strconv.FormatInt(timestamp, 10)
			log.Debugf("replacind old ID %s for new one %s", oldid, newid)
			replaced = bytes.Replace(replaced, []byte(oldid), []byte(newid), -1)
		}
		ImportedData = impexp.ExportData{}
		if err = json.Unmarshal(replaced, &ImportedData); err != nil {
			log.Errorf("Error in data to struct (json-unmarshal) procces: %s", err)
			ctx.JSON(404, err.Error())
		}
		err = ImportedData.Import()
		if err != nil {
			log.Errorf("Some Error happened on import data: %s", err)
			ctx.JSON(404, err.Error())
		}
		ctx.JSON(200, &ImportCheck{Info: ImportedData.Info, IsOk: true, Message: "all duplicated ID's have been replaced", Duplicated: arrays})
		return
	}
	err = ImportedData.Import()
	if err != nil {
		log.Errorf("Some Error happened on import data: %s", err)
		ctx.JSON(404, err.Error())
	}
	ctx.JSON(200, &ImportCheck{Info: ImportedData.Info, IsOk: true, Message: "all objects have been  imported", Duplicated: nil})
}

/****************/
/*EXPORT*/
/****************/

func ExportObject(ctx *Context) {
	id := ctx.Params(":id")
	objtype := ctx.Params(":objtype")
	exp := impexp.NewExport(&impexp.ExportInfo{FileName: "autogenerated.txt", Description: "autogenerated"})
	err := exp.Export(objtype, id)
	if err != nil {
		log.Warningf("Error on get object array for type %s with ID %s  , error: %s", objtype, id, err)
		ctx.JSON(404, err.Error())
	} else {
		ctx.JSON(200, &exp)
	}
}

func ExportObjectToFile(ctx *Context, info impexp.ExportInfo) {

	id := ctx.Params(":id")
	objtype := ctx.Params(":objtype")
	exp := impexp.NewExport(&info)
	err := exp.Export(objtype, id)
	if err != nil {
		log.Errorf("Error on create Export Data with: %s", err)
		ctx.JSON(404, err.Error())
		return
	}
	outdata, err := json.Marshal(exp)
	if err != nil {
		log.Errorf("Error on Json data formatting: %s", err)
		ctx.JSON(404, err.Error())
		return
	}
	tmpfile, err := ioutil.TempFile("", "snmpcollector")
	if err != nil {
		log.Errorf("Error on create new temporal file: %s", err)
		ctx.JSON(404, err.Error())
		return
	}
	defer os.Remove(tmpfile.Name())
	if _, err := tmpfile.Write(outdata); err != nil {
		log.Errorf("Error on write data to temporal file: %s", tmpfile.Name())
		ctx.JSON(404, err.Error())
		return
	}
	if err := tmpfile.Close(); err != nil {
		log.Errorf("Error on close temporal file: %s", tmpfile.Name())
		ctx.JSON(404, err.Error())
		return
	}

	ctx.ServeFile(tmpfile.Name())
}
<|endoftext|>


<commit_before><commit_msg>Fix: Typo - HTML closing bracket (#96)<commit_after><|endoftext|>


<commit_before>package elements

import (
	"log"
	"strings"
)

// Flags ... Flags struct
type Flags struct {
	value string
}

// SetValue ... set the value for this element
func (flags *Flags) SetValue(value string) {
	flags.value = strings.ToUpper(strings.Replace(value, " ", "", -1))
}

// Convert ... output ASCII for this element
func (flags Flags) Convert() string {
	ret := ""
	if flags.value != "" && flags.Validate() {
		ret = " Flags: " + flags.value + " "
	}
	return ret
}

// Validate ... validate this element
func (flags Flags) Validate() bool {
	ret := true

	if ret != true {
		log.Println("Failed to validate flags '" + flags.value + "'")
	}
	return ret
}
<commit_msg>Update flags format<commit_after>package elements

import (
	"log"
	"strings"
)

// Flags ... Flags struct
type Flags struct {
	value string
}

// SetValue ... set the value for this element
func (flags *Flags) SetValue(value string) {
	flags.value = strings.ToUpper(strings.Replace(value, " ", "", -1))
}

// Convert ... output ASCII for this element
func (flags Flags) Convert() string {
	ret := ""
	if flags.value != "" && flags.Validate() {
		ret = flags.value + ": "
	}
	return ret
}

// Validate ... validate this element
func (flags Flags) Validate() bool {
	ret := true

	if ret != true {
		log.Println("Failed to validate flags '" + flags.value + "'")
	}
	return ret
}
<|endoftext|>


<commit_before>package handlers

import (
	"fmt"
	"net/http"
	"net/url"
	"os"

	"github.com/docker/distribution"
	ctxu "github.com/docker/distribution/context"
	"github.com/docker/distribution/digest"
	"github.com/docker/distribution/reference"
	"github.com/docker/distribution/registry/api/errcode"
	"github.com/docker/distribution/registry/api/v2"
	"github.com/docker/distribution/registry/storage"
	"github.com/gorilla/handlers"
)

// blobUploadDispatcher constructs and returns the blob upload handler for the
// given request context.
func blobUploadDispatcher(ctx *Context, r *http.Request) http.Handler {
	buh := &blobUploadHandler{
		Context: ctx,
		UUID:    getUploadUUID(ctx),
	}

	handler := handlers.MethodHandler{
		"GET":  http.HandlerFunc(buh.GetUploadStatus),
		"HEAD": http.HandlerFunc(buh.GetUploadStatus),
	}

	if !ctx.readOnly {
		handler["POST"] = http.HandlerFunc(buh.StartBlobUpload)
		handler["PATCH"] = http.HandlerFunc(buh.PatchBlobData)
		handler["PUT"] = http.HandlerFunc(buh.PutBlobUploadComplete)
		handler["DELETE"] = http.HandlerFunc(buh.CancelBlobUpload)
	}

	if buh.UUID != "" {
		state, err := hmacKey(ctx.Config.HTTP.Secret).unpackUploadState(r.FormValue("_state"))
		if err != nil {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				ctxu.GetLogger(ctx).Infof("error resolving upload: %v", err)
				buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadInvalid.WithDetail(err))
			})
		}
		buh.State = state

		if state.Name != ctx.Repository.Name() {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				ctxu.GetLogger(ctx).Infof("mismatched repository name in upload state: %q != %q", state.Name, buh.Repository.Name())
				buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadInvalid.WithDetail(err))
			})
		}

		if state.UUID != buh.UUID {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				ctxu.GetLogger(ctx).Infof("mismatched uuid in upload state: %q != %q", state.UUID, buh.UUID)
				buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadInvalid.WithDetail(err))
			})
		}

		blobs := ctx.Repository.Blobs(buh)
		upload, err := blobs.Resume(buh, buh.UUID)
		if err != nil {
			ctxu.GetLogger(ctx).Errorf("error resolving upload: %v", err)
			if err == distribution.ErrBlobUploadUnknown {
				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadUnknown.WithDetail(err))
				})
			}

			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
			})
		}
		buh.Upload = upload

		if state.Offset > 0 {
			// Seek the blob upload to the correct spot if it's non-zero.
			// These error conditions should be rare and demonstrate really
			// problems. We basically cancel the upload and tell the client to
			// start over.
			if nn, err := upload.Seek(buh.State.Offset, os.SEEK_SET); err != nil {
				defer upload.Close()
				ctxu.GetLogger(ctx).Infof("error seeking blob upload: %v", err)
				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadInvalid.WithDetail(err))
					upload.Cancel(buh)
				})
			} else if nn != buh.State.Offset {
				defer upload.Close()
				ctxu.GetLogger(ctx).Infof("seek to wrong offest: %d != %d", nn, buh.State.Offset)
				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadInvalid.WithDetail(err))
					upload.Cancel(buh)
				})
			}
		}

		return closeResources(handler, buh.Upload)
	}

	return handler
}

// blobUploadHandler handles the http blob upload process.
type blobUploadHandler struct {
	*Context

	// UUID identifies the upload instance for the current request. Using UUID
	// to key blob writers since this implementation uses UUIDs.
	UUID string

	Upload distribution.BlobWriter

	State blobUploadState
}

// StartBlobUpload begins the blob upload process and allocates a server-side
// blob writer session, optionally mounting the blob from a separate repository.
func (buh *blobUploadHandler) StartBlobUpload(w http.ResponseWriter, r *http.Request) {
	var options []distribution.BlobCreateOption

	fromRepo := r.FormValue("from")
	mountDigest := r.FormValue("mount")

	if mountDigest != "" && fromRepo != "" {
		opt, err := buh.createBlobMountOption(fromRepo, mountDigest)
		if err != nil {
			options = append(options, opt)
		}
	}

	blobs := buh.Repository.Blobs(buh)
	upload, err := blobs.Create(buh, options...)

	if err != nil {
		if ebm, ok := err.(distribution.ErrBlobMounted); ok {
			if err := buh.writeBlobCreatedHeaders(w, ebm.Descriptor); err != nil {
				buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
			}
		} else if err == distribution.ErrUnsupported {
			buh.Errors = append(buh.Errors, errcode.ErrorCodeUnsupported)
		} else {
			buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
		}
		return
	}

	buh.Upload = upload
	defer buh.Upload.Close()

	if err := buh.blobUploadResponse(w, r, true); err != nil {
		buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
		return
	}

	w.Header().Set("Docker-Upload-UUID", buh.Upload.ID())
	w.WriteHeader(http.StatusAccepted)
}

// GetUploadStatus returns the status of a given upload, identified by id.
func (buh *blobUploadHandler) GetUploadStatus(w http.ResponseWriter, r *http.Request) {
	if buh.Upload == nil {
		buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadUnknown)
		return
	}

	// TODO(dmcgowan): Set last argument to false in blobUploadResponse when
	// resumable upload is supported. This will enable returning a non-zero
	// range for clients to begin uploading at an offset.
	if err := buh.blobUploadResponse(w, r, true); err != nil {
		buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
		return
	}

	w.Header().Set("Docker-Upload-UUID", buh.UUID)
	w.WriteHeader(http.StatusNoContent)
}

// PatchBlobData writes data to an upload.
func (buh *blobUploadHandler) PatchBlobData(w http.ResponseWriter, r *http.Request) {
	if buh.Upload == nil {
		buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadUnknown)
		return
	}

	ct := r.Header.Get("Content-Type")
	if ct != "" && ct != "application/octet-stream" {
		buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(fmt.Errorf("Bad Content-Type")))
		// TODO(dmcgowan): encode error
		return
	}

	// TODO(dmcgowan): support Content-Range header to seek and write range

	if err := copyFullPayload(w, r, buh.Upload, buh, "blob PATCH", &buh.Errors); err != nil {
		// copyFullPayload reports the error if necessary
		return
	}

	if err := buh.blobUploadResponse(w, r, false); err != nil {
		buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
		return
	}

	w.WriteHeader(http.StatusAccepted)
}

// PutBlobUploadComplete takes the final request of a blob upload. The
// request may include all the blob data or no blob data. Any data
// provided is received and verified. If successful, the blob is linked
// into the blob store and 201 Created is returned with the canonical
// url of the blob.
func (buh *blobUploadHandler) PutBlobUploadComplete(w http.ResponseWriter, r *http.Request) {
	if buh.Upload == nil {
		buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadUnknown)
		return
	}

	dgstStr := r.FormValue("digest") // TODO(stevvooe): Support multiple digest parameters!

	if dgstStr == "" {
		// no digest? return error, but allow retry.
		buh.Errors = append(buh.Errors, v2.ErrorCodeDigestInvalid.WithDetail("digest missing"))
		return
	}

	dgst, err := digest.ParseDigest(dgstStr)
	if err != nil {
		// no digest? return error, but allow retry.
		buh.Errors = append(buh.Errors, v2.ErrorCodeDigestInvalid.WithDetail("digest parsing failed"))
		return
	}

	if err := copyFullPayload(w, r, buh.Upload, buh, "blob PUT", &buh.Errors); err != nil {
		// copyFullPayload reports the error if necessary
		return
	}

	desc, err := buh.Upload.Commit(buh, distribution.Descriptor{
		Digest: dgst,

		// TODO(stevvooe): This isn't wildly important yet, but we should
		// really set the length and mediatype. For now, we can let the
		// backend take care of this.
	})

	if err != nil {
		switch err := err.(type) {
		case distribution.ErrBlobInvalidDigest:
			buh.Errors = append(buh.Errors, v2.ErrorCodeDigestInvalid.WithDetail(err))
		default:
			switch err {
			case distribution.ErrUnsupported:
				buh.Errors = append(buh.Errors, errcode.ErrorCodeUnsupported)
			case distribution.ErrBlobInvalidLength, distribution.ErrBlobDigestUnsupported:
				buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadInvalid.WithDetail(err))
			default:
				ctxu.GetLogger(buh).Errorf("unknown error completing upload: %#v", err)
				buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
			}

		}

		// Clean up the backend blob data if there was an error.
		if err := buh.Upload.Cancel(buh); err != nil {
			// If the cleanup fails, all we can do is observe and report.
			ctxu.GetLogger(buh).Errorf("error canceling upload after error: %v", err)
		}

		return
	}
	if err := buh.writeBlobCreatedHeaders(w, desc); err != nil {
		buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
		return
	}
}

// CancelBlobUpload cancels an in-progress upload of a blob.
func (buh *blobUploadHandler) CancelBlobUpload(w http.ResponseWriter, r *http.Request) {
	if buh.Upload == nil {
		buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadUnknown)
		return
	}

	w.Header().Set("Docker-Upload-UUID", buh.UUID)
	if err := buh.Upload.Cancel(buh); err != nil {
		ctxu.GetLogger(buh).Errorf("error encountered canceling upload: %v", err)
		buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
	}

	w.WriteHeader(http.StatusNoContent)
}

// blobUploadResponse provides a standard request for uploading blobs and
// chunk responses. This sets the correct headers but the response status is
// left to the caller. The fresh argument is used to ensure that new blob
// uploads always start at a 0 offset. This allows disabling resumable push by
// always returning a 0 offset on check status.
func (buh *blobUploadHandler) blobUploadResponse(w http.ResponseWriter, r *http.Request, fresh bool) error {

	var offset int64
	if !fresh {
		var err error
		offset, err = buh.Upload.Seek(0, os.SEEK_CUR)
		if err != nil {
			ctxu.GetLogger(buh).Errorf("unable get current offset of blob upload: %v", err)
			return err
		}
	}

	// TODO(stevvooe): Need a better way to manage the upload state automatically.
	buh.State.Name = buh.Repository.Name()
	buh.State.UUID = buh.Upload.ID()
	buh.State.Offset = offset
	buh.State.StartedAt = buh.Upload.StartedAt()

	token, err := hmacKey(buh.Config.HTTP.Secret).packUploadState(buh.State)
	if err != nil {
		ctxu.GetLogger(buh).Infof("error building upload state token: %s", err)
		return err
	}

	uploadURL, err := buh.urlBuilder.BuildBlobUploadChunkURL(
		buh.Repository.Name(), buh.Upload.ID(),
		url.Values{
			"_state": []string{token},
		})
	if err != nil {
		ctxu.GetLogger(buh).Infof("error building upload url: %s", err)
		return err
	}

	endRange := offset
	if endRange > 0 {
		endRange = endRange - 1
	}

	w.Header().Set("Docker-Upload-UUID", buh.UUID)
	w.Header().Set("Location", uploadURL)
	w.Header().Set("Content-Length", "0")
	w.Header().Set("Range", fmt.Sprintf("0-%d", endRange))

	return nil
}

// mountBlob attempts to mount a blob from another repository by its digest. If
// successful, the blob is linked into the blob store and 201 Created is
// returned with the canonical url of the blob.
func (buh *blobUploadHandler) createBlobMountOption(fromRepo, mountDigest string) (distribution.BlobCreateOption, error) {
	dgst, err := digest.ParseDigest(mountDigest)
	if err != nil {
		return nil, err
	}

	ref, err := reference.ParseNamed(fromRepo)
	if err != nil {
		return nil, err
	}

	canonical, err := reference.WithDigest(ref, dgst)
	if err != nil {
		return nil, err
	}

	return storage.WithMountFrom(canonical), nil
}

// writeBlobCreatedHeaders writes the standard headers describing a newly
// created blob. A 201 Created is written as well as the canonical URL and
// blob digest.
func (buh *blobUploadHandler) writeBlobCreatedHeaders(w http.ResponseWriter, desc distribution.Descriptor) error {
	blobURL, err := buh.urlBuilder.BuildBlobURL(buh.Repository.Name(), desc.Digest)
	if err != nil {
		return err
	}

	w.Header().Set("Location", blobURL)
	w.Header().Set("Content-Length", "0")
	w.Header().Set("Docker-Content-Digest", desc.Digest.String())
	w.WriteHeader(http.StatusCreated)
	return nil
}
<commit_msg>Fixes cross-repo blob mounting in the BlobUploadHandler<commit_after>package handlers

import (
	"fmt"
	"net/http"
	"net/url"
	"os"

	"github.com/docker/distribution"
	ctxu "github.com/docker/distribution/context"
	"github.com/docker/distribution/digest"
	"github.com/docker/distribution/reference"
	"github.com/docker/distribution/registry/api/errcode"
	"github.com/docker/distribution/registry/api/v2"
	"github.com/docker/distribution/registry/storage"
	"github.com/gorilla/handlers"
)

// blobUploadDispatcher constructs and returns the blob upload handler for the
// given request context.
func blobUploadDispatcher(ctx *Context, r *http.Request) http.Handler {
	buh := &blobUploadHandler{
		Context: ctx,
		UUID:    getUploadUUID(ctx),
	}

	handler := handlers.MethodHandler{
		"GET":  http.HandlerFunc(buh.GetUploadStatus),
		"HEAD": http.HandlerFunc(buh.GetUploadStatus),
	}

	if !ctx.readOnly {
		handler["POST"] = http.HandlerFunc(buh.StartBlobUpload)
		handler["PATCH"] = http.HandlerFunc(buh.PatchBlobData)
		handler["PUT"] = http.HandlerFunc(buh.PutBlobUploadComplete)
		handler["DELETE"] = http.HandlerFunc(buh.CancelBlobUpload)
	}

	if buh.UUID != "" {
		state, err := hmacKey(ctx.Config.HTTP.Secret).unpackUploadState(r.FormValue("_state"))
		if err != nil {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				ctxu.GetLogger(ctx).Infof("error resolving upload: %v", err)
				buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadInvalid.WithDetail(err))
			})
		}
		buh.State = state

		if state.Name != ctx.Repository.Name() {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				ctxu.GetLogger(ctx).Infof("mismatched repository name in upload state: %q != %q", state.Name, buh.Repository.Name())
				buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadInvalid.WithDetail(err))
			})
		}

		if state.UUID != buh.UUID {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				ctxu.GetLogger(ctx).Infof("mismatched uuid in upload state: %q != %q", state.UUID, buh.UUID)
				buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadInvalid.WithDetail(err))
			})
		}

		blobs := ctx.Repository.Blobs(buh)
		upload, err := blobs.Resume(buh, buh.UUID)
		if err != nil {
			ctxu.GetLogger(ctx).Errorf("error resolving upload: %v", err)
			if err == distribution.ErrBlobUploadUnknown {
				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadUnknown.WithDetail(err))
				})
			}

			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
			})
		}
		buh.Upload = upload

		if state.Offset > 0 {
			// Seek the blob upload to the correct spot if it's non-zero.
			// These error conditions should be rare and demonstrate really
			// problems. We basically cancel the upload and tell the client to
			// start over.
			if nn, err := upload.Seek(buh.State.Offset, os.SEEK_SET); err != nil {
				defer upload.Close()
				ctxu.GetLogger(ctx).Infof("error seeking blob upload: %v", err)
				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadInvalid.WithDetail(err))
					upload.Cancel(buh)
				})
			} else if nn != buh.State.Offset {
				defer upload.Close()
				ctxu.GetLogger(ctx).Infof("seek to wrong offest: %d != %d", nn, buh.State.Offset)
				return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadInvalid.WithDetail(err))
					upload.Cancel(buh)
				})
			}
		}

		return closeResources(handler, buh.Upload)
	}

	return handler
}

// blobUploadHandler handles the http blob upload process.
type blobUploadHandler struct {
	*Context

	// UUID identifies the upload instance for the current request. Using UUID
	// to key blob writers since this implementation uses UUIDs.
	UUID string

	Upload distribution.BlobWriter

	State blobUploadState
}

// StartBlobUpload begins the blob upload process and allocates a server-side
// blob writer session, optionally mounting the blob from a separate repository.
func (buh *blobUploadHandler) StartBlobUpload(w http.ResponseWriter, r *http.Request) {
	var options []distribution.BlobCreateOption

	fromRepo := r.FormValue("from")
	mountDigest := r.FormValue("mount")

	if mountDigest != "" && fromRepo != "" {
		opt, err := buh.createBlobMountOption(fromRepo, mountDigest)
		if opt != nil && err == nil {
			options = append(options, opt)
		}
	}

	blobs := buh.Repository.Blobs(buh)
	upload, err := blobs.Create(buh, options...)

	if err != nil {
		if ebm, ok := err.(distribution.ErrBlobMounted); ok {
			if err := buh.writeBlobCreatedHeaders(w, ebm.Descriptor); err != nil {
				buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
			}
		} else if err == distribution.ErrUnsupported {
			buh.Errors = append(buh.Errors, errcode.ErrorCodeUnsupported)
		} else {
			buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
		}
		return
	}

	buh.Upload = upload
	defer buh.Upload.Close()

	if err := buh.blobUploadResponse(w, r, true); err != nil {
		buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
		return
	}

	w.Header().Set("Docker-Upload-UUID", buh.Upload.ID())
	w.WriteHeader(http.StatusAccepted)
}

// GetUploadStatus returns the status of a given upload, identified by id.
func (buh *blobUploadHandler) GetUploadStatus(w http.ResponseWriter, r *http.Request) {
	if buh.Upload == nil {
		buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadUnknown)
		return
	}

	// TODO(dmcgowan): Set last argument to false in blobUploadResponse when
	// resumable upload is supported. This will enable returning a non-zero
	// range for clients to begin uploading at an offset.
	if err := buh.blobUploadResponse(w, r, true); err != nil {
		buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
		return
	}

	w.Header().Set("Docker-Upload-UUID", buh.UUID)
	w.WriteHeader(http.StatusNoContent)
}

// PatchBlobData writes data to an upload.
func (buh *blobUploadHandler) PatchBlobData(w http.ResponseWriter, r *http.Request) {
	if buh.Upload == nil {
		buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadUnknown)
		return
	}

	ct := r.Header.Get("Content-Type")
	if ct != "" && ct != "application/octet-stream" {
		buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(fmt.Errorf("Bad Content-Type")))
		// TODO(dmcgowan): encode error
		return
	}

	// TODO(dmcgowan): support Content-Range header to seek and write range

	if err := copyFullPayload(w, r, buh.Upload, buh, "blob PATCH", &buh.Errors); err != nil {
		// copyFullPayload reports the error if necessary
		return
	}

	if err := buh.blobUploadResponse(w, r, false); err != nil {
		buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
		return
	}

	w.WriteHeader(http.StatusAccepted)
}

// PutBlobUploadComplete takes the final request of a blob upload. The
// request may include all the blob data or no blob data. Any data
// provided is received and verified. If successful, the blob is linked
// into the blob store and 201 Created is returned with the canonical
// url of the blob.
func (buh *blobUploadHandler) PutBlobUploadComplete(w http.ResponseWriter, r *http.Request) {
	if buh.Upload == nil {
		buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadUnknown)
		return
	}

	dgstStr := r.FormValue("digest") // TODO(stevvooe): Support multiple digest parameters!

	if dgstStr == "" {
		// no digest? return error, but allow retry.
		buh.Errors = append(buh.Errors, v2.ErrorCodeDigestInvalid.WithDetail("digest missing"))
		return
	}

	dgst, err := digest.ParseDigest(dgstStr)
	if err != nil {
		// no digest? return error, but allow retry.
		buh.Errors = append(buh.Errors, v2.ErrorCodeDigestInvalid.WithDetail("digest parsing failed"))
		return
	}

	if err := copyFullPayload(w, r, buh.Upload, buh, "blob PUT", &buh.Errors); err != nil {
		// copyFullPayload reports the error if necessary
		return
	}

	desc, err := buh.Upload.Commit(buh, distribution.Descriptor{
		Digest: dgst,

		// TODO(stevvooe): This isn't wildly important yet, but we should
		// really set the length and mediatype. For now, we can let the
		// backend take care of this.
	})

	if err != nil {
		switch err := err.(type) {
		case distribution.ErrBlobInvalidDigest:
			buh.Errors = append(buh.Errors, v2.ErrorCodeDigestInvalid.WithDetail(err))
		default:
			switch err {
			case distribution.ErrUnsupported:
				buh.Errors = append(buh.Errors, errcode.ErrorCodeUnsupported)
			case distribution.ErrBlobInvalidLength, distribution.ErrBlobDigestUnsupported:
				buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadInvalid.WithDetail(err))
			default:
				ctxu.GetLogger(buh).Errorf("unknown error completing upload: %#v", err)
				buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
			}

		}

		// Clean up the backend blob data if there was an error.
		if err := buh.Upload.Cancel(buh); err != nil {
			// If the cleanup fails, all we can do is observe and report.
			ctxu.GetLogger(buh).Errorf("error canceling upload after error: %v", err)
		}

		return
	}
	if err := buh.writeBlobCreatedHeaders(w, desc); err != nil {
		buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
		return
	}
}

// CancelBlobUpload cancels an in-progress upload of a blob.
func (buh *blobUploadHandler) CancelBlobUpload(w http.ResponseWriter, r *http.Request) {
	if buh.Upload == nil {
		buh.Errors = append(buh.Errors, v2.ErrorCodeBlobUploadUnknown)
		return
	}

	w.Header().Set("Docker-Upload-UUID", buh.UUID)
	if err := buh.Upload.Cancel(buh); err != nil {
		ctxu.GetLogger(buh).Errorf("error encountered canceling upload: %v", err)
		buh.Errors = append(buh.Errors, errcode.ErrorCodeUnknown.WithDetail(err))
	}

	w.WriteHeader(http.StatusNoContent)
}

// blobUploadResponse provides a standard request for uploading blobs and
// chunk responses. This sets the correct headers but the response status is
// left to the caller. The fresh argument is used to ensure that new blob
// uploads always start at a 0 offset. This allows disabling resumable push by
// always returning a 0 offset on check status.
func (buh *blobUploadHandler) blobUploadResponse(w http.ResponseWriter, r *http.Request, fresh bool) error {

	var offset int64
	if !fresh {
		var err error
		offset, err = buh.Upload.Seek(0, os.SEEK_CUR)
		if err != nil {
			ctxu.GetLogger(buh).Errorf("unable get current offset of blob upload: %v", err)
			return err
		}
	}

	// TODO(stevvooe): Need a better way to manage the upload state automatically.
	buh.State.Name = buh.Repository.Name()
	buh.State.UUID = buh.Upload.ID()
	buh.State.Offset = offset
	buh.State.StartedAt = buh.Upload.StartedAt()

	token, err := hmacKey(buh.Config.HTTP.Secret).packUploadState(buh.State)
	if err != nil {
		ctxu.GetLogger(buh).Infof("error building upload state token: %s", err)
		return err
	}

	uploadURL, err := buh.urlBuilder.BuildBlobUploadChunkURL(
		buh.Repository.Name(), buh.Upload.ID(),
		url.Values{
			"_state": []string{token},
		})
	if err != nil {
		ctxu.GetLogger(buh).Infof("error building upload url: %s", err)
		return err
	}

	endRange := offset
	if endRange > 0 {
		endRange = endRange - 1
	}

	w.Header().Set("Docker-Upload-UUID", buh.UUID)
	w.Header().Set("Location", uploadURL)
	w.Header().Set("Content-Length", "0")
	w.Header().Set("Range", fmt.Sprintf("0-%d", endRange))

	return nil
}

// mountBlob attempts to mount a blob from another repository by its digest. If
// successful, the blob is linked into the blob store and 201 Created is
// returned with the canonical url of the blob.
func (buh *blobUploadHandler) createBlobMountOption(fromRepo, mountDigest string) (distribution.BlobCreateOption, error) {
	dgst, err := digest.ParseDigest(mountDigest)
	if err != nil {
		return nil, err
	}

	ref, err := reference.ParseNamed(fromRepo)
	if err != nil {
		return nil, err
	}

	canonical, err := reference.WithDigest(ref, dgst)
	if err != nil {
		return nil, err
	}

	return storage.WithMountFrom(canonical), nil
}

// writeBlobCreatedHeaders writes the standard headers describing a newly
// created blob. A 201 Created is written as well as the canonical URL and
// blob digest.
func (buh *blobUploadHandler) writeBlobCreatedHeaders(w http.ResponseWriter, desc distribution.Descriptor) error {
	blobURL, err := buh.urlBuilder.BuildBlobURL(buh.Repository.Name(), desc.Digest)
	if err != nil {
		return err
	}

	w.Header().Set("Location", blobURL)
	w.Header().Set("Content-Length", "0")
	w.Header().Set("Docker-Content-Digest", desc.Digest.String())
	w.WriteHeader(http.StatusCreated)
	return nil
}
<|endoftext|>


<commit_before>package volumecommands

import (
	"github.com/codegangsta/cli"
	"github.com/jrperritt/rack/handler"
	"github.com/jrperritt/rack/util"
	osVolumes "github.com/rackspace/gophercloud/openstack/blockstorage/v1/volumes"
)

var create = cli.Command{
	Name:        "create",
	Usage:       util.Usage(commandPrefix, "create", "--size <size>"),
	Description: "Creates a volume",
	Action:      actionCreate,
	Flags:       util.CommandFlags(flagsCreate, keysCreate),
	BashComplete: func(c *cli.Context) {
		util.CompleteFlags(util.CommandFlags(flagsCreate, keysCreate))
	},
}

func flagsCreate() []cli.Flag {
	return []cli.Flag{
		cli.IntFlag{
			Name:  "size",
			Usage: "[required] The size of this volume (in gigabytes). Valid values are between 75 and 1024.",
		},
		cli.StringFlag{
			Name:  "name",
			Usage: "[optional] A name for this volume.",
		},
		cli.StringFlag{
			Name:  "description",
			Usage: "[optional] A description for this volume.",
		},
		cli.StringFlag{
			Name:  "volume-type",
			Usage: "[optional] The volume type of this volume.",
		},
	}
}

var keysCreate = []string{"ID", "Name", "Description", "Size", "Volume Type", "Snapshot ID", "Attachments", "Created"}

type paramsCreate struct {
	opts *osVolumes.CreateOpts
}

type commandCreate handler.Command

func actionCreate(c *cli.Context) {
	command := &commandCreate{
		Ctx: &handler.Context{
			CLIContext: c,
		},
	}
	handler.Handle(command)
}

func (command *commandCreate) Context() *handler.Context {
	return command.Ctx
}

func (command *commandCreate) Keys() []string {
	return keysCreate
}

func (command *commandCreate) ServiceClientType() string {
	return serviceClientType
}

func (command *commandCreate) HandleFlags(resource *handler.Resource) error {
	err := command.Ctx.CheckFlagsSet([]string{"size"})
	if err != nil {
		return err
	}

	c := command.Ctx.CLIContext

	opts := &osVolumes.CreateOpts{
		Size:        c.Int("size"),
		Name:        c.String("name"),
		Description: c.String("description"),
		VolumeType:  c.String("volume-type"),
	}

	resource.Params = &paramsCreate{
		opts: opts,
	}

	return nil
}

func (command *commandCreate) Execute(resource *handler.Resource) {
	opts := resource.Params.(*paramsCreate).opts
	volume, err := osVolumes.Create(command.Ctx.ServiceClient, opts).Extract()
	if err != nil {
		resource.Err = err
		return
	}
	resource.Result = volumeSingle(volume)
}
<commit_msg>fix 'volume create' header values<commit_after>package volumecommands

import (
	"github.com/codegangsta/cli"
	"github.com/jrperritt/rack/handler"
	"github.com/jrperritt/rack/util"
	osVolumes "github.com/rackspace/gophercloud/openstack/blockstorage/v1/volumes"
)

var create = cli.Command{
	Name:        "create",
	Usage:       util.Usage(commandPrefix, "create", "--size <size>"),
	Description: "Creates a volume",
	Action:      actionCreate,
	Flags:       util.CommandFlags(flagsCreate, keysCreate),
	BashComplete: func(c *cli.Context) {
		util.CompleteFlags(util.CommandFlags(flagsCreate, keysCreate))
	},
}

func flagsCreate() []cli.Flag {
	return []cli.Flag{
		cli.IntFlag{
			Name:  "size",
			Usage: "[required] The size of this volume (in gigabytes). Valid values are between 75 and 1024.",
		},
		cli.StringFlag{
			Name:  "name",
			Usage: "[optional] A name for this volume.",
		},
		cli.StringFlag{
			Name:  "description",
			Usage: "[optional] A description for this volume.",
		},
		cli.StringFlag{
			Name:  "volume-type",
			Usage: "[optional] The volume type of this volume.",
		},
	}
}

var keysCreate = []string{"ID", "Name", "Description", "Size", "VolumeType", "SnapshotID", "Attachments", "CreatedAt"}

type paramsCreate struct {
	opts *osVolumes.CreateOpts
}

type commandCreate handler.Command

func actionCreate(c *cli.Context) {
	command := &commandCreate{
		Ctx: &handler.Context{
			CLIContext: c,
		},
	}
	handler.Handle(command)
}

func (command *commandCreate) Context() *handler.Context {
	return command.Ctx
}

func (command *commandCreate) Keys() []string {
	return keysCreate
}

func (command *commandCreate) ServiceClientType() string {
	return serviceClientType
}

func (command *commandCreate) HandleFlags(resource *handler.Resource) error {
	err := command.Ctx.CheckFlagsSet([]string{"size"})
	if err != nil {
		return err
	}

	c := command.Ctx.CLIContext

	opts := &osVolumes.CreateOpts{
		Size:        c.Int("size"),
		Name:        c.String("name"),
		Description: c.String("description"),
		VolumeType:  c.String("volume-type"),
	}

	resource.Params = &paramsCreate{
		opts: opts,
	}

	return nil
}

func (command *commandCreate) Execute(resource *handler.Resource) {
	opts := resource.Params.(*paramsCreate).opts
	volume, err := osVolumes.Create(command.Ctx.ServiceClient, opts).Extract()
	if err != nil {
		resource.Err = err
		return
	}
	resource.Result = volumeSingle(volume)
}
<|endoftext|>


<commit_before>// Copyright 2016 The Linux Foundation
//
// 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 image

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"testing"

	"github.com/opencontainers/image-spec/specs-go/v1"
)

const (
	refTag = "latest"

	layoutStr = `{"imageLayoutVersion": "1.0.0"}`

	configStr = `{
    "created": "2015-10-31T22:22:56.015925234Z",
    "author": "Alyssa P. Hacker <alyspdev@example.com>",
    "architecture": "amd64",
    "os": "linux",
    "config": {
        "User": "alice",
        "Memory": 2048,
        "MemorySwap": 4096,
        "CpuShares": 8,
        "ExposedPorts": {
            "8080/tcp": {}
        },
        "Env": [
            "PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
            "FOO=oci_is_a",
            "BAR=well_written_spec"
        ],
        "Entrypoint": [
            "/bin/my-app-binary"
        ],
        "Cmd": [
            "--foreground",
            "--config",
            "/etc/my-app.d/default.cfg"
        ],
        "Volumes": {
            "/var/job-result-data": {},
            "/var/log/my-app-logs": {}
        },
        "WorkingDir": "/home/alice"
    },
    "rootfs": {
      "diff_ids": [
        "sha256:c6f988f4874bb0add23a778f753c65efe992244e148a1d2ec2a8b664fb66bbd1",
        "sha256:5f70bf18a086007016e948b04aed3b82103a36bea41755b6cddfaf10ace3c6ef"
      ],
      "type": "layers"
    },
    "history": [
      {
        "created": "2015-10-31T22:22:54.690851953Z",
        "created_by": "/bin/sh -c #(nop) ADD file:a3bc1e842b69636f9df5256c49c5374fb4eef1e281fe3f282c65fb853ee171c5 in /"
      },
      {
        "created": "2015-10-31T22:22:55.613815829Z",
        "created_by": "/bin/sh -c #(nop) CMD [\"sh\"]",
        "empty_layer": true
      }
    ]
}
`
)

var (
	refStr = `{"digest":"<manifest_digest>","mediaType":"application/vnd.oci.image.manifest.v1+json","size":<manifest_size>}`

	manifestStr = `{
    "annotations": null,
    "config": {
        "digest": "<config_digest>",
        "mediaType": "application/vnd.oci.image.config.v1+json",
        "size": <config_size>
    },
    "layers": [
        {
            "digest": "<layer_digest>",
            "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
            "size": <layer_size>
        }
    ],
    "mediaType": "application/vnd.oci.image.manifest.v1+json",
    "schemaVersion": 2
}
 `
)

type tarContent struct {
	header *tar.Header
	b      []byte
}

type imageLayout struct {
	rootDir  string
	layout   string
	ref      string
	manifest string
	config   string
	tarList  []tarContent
}

func TestValidateLayout(t *testing.T) {
	root, err := ioutil.TempDir("", "oci-test")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(root)

	il := imageLayout{
		rootDir:  root,
		layout:   layoutStr,
		ref:      refTag,
		manifest: manifestStr,
		config:   configStr,
		tarList: []tarContent{
			tarContent{&tar.Header{Name: "test", Size: 4, Mode: 0600}, []byte("test")},
		},
	}

	// create image layout bundle
	err = createImageLayoutBundle(il)
	if err != nil {
		t.Fatal(err)
	}

	err = ValidateLayout(root, []string{refTag}, nil)
	if err != nil {
		t.Fatal(err)
	}
}

func createImageLayoutBundle(il imageLayout) error {
	err := os.MkdirAll(filepath.Join(il.rootDir, "blobs", "sha256"), 0700)
	if err != nil {
		return err
	}

	err = os.MkdirAll(filepath.Join(il.rootDir, "refs"), 0700)
	if err != nil {
		return err
	}

	// create image layout file
	err = createLayoutFile(il.rootDir)
	if err != nil {
		return err
	}

	// create image layer blob file.
	desc, err := createImageLayerFile(il.rootDir, il.tarList)
	if err != nil {
		return err
	}
	il.manifest = strings.Replace(il.manifest, "<layer_digest>", desc.Digest, 1)
	il.manifest = strings.Replace(il.manifest, "<layer_size>", strconv.FormatInt(desc.Size, 10), 1)

	desc, err = createConfigFile(il.rootDir, il.config)
	if err != nil {
		return err
	}
	il.manifest = strings.Replace(il.manifest, "<config_digest>", desc.Digest, 1)
	il.manifest = strings.Replace(il.manifest, "<config_size>", strconv.FormatInt(desc.Size, 10), 1)

	// create manifest blob file
	desc, err = createManifestFile(il.rootDir, il.manifest)
	if err != nil {
		return err
	}

	return createRefFile(il.rootDir, il.ref, desc)
}

func createLayoutFile(root string) error {
	layoutPath := filepath.Join(root, "oci-layout")
	f, err := os.Create(layoutPath)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = io.Copy(f, bytes.NewBuffer([]byte(layoutStr)))
	return err
}

func createRefFile(root, ref string, mft descriptor) error {
	refpath := filepath.Join(root, "refs", ref)
	f, err := os.Create(refpath)
	if err != nil {
		return err
	}
	defer f.Close()
	refStr = strings.Replace(refStr, "<manifest_digest>", mft.Digest, -1)
	refStr = strings.Replace(refStr, "<manifest_size>", strconv.FormatInt(mft.Size, 10), -1)
	_, err = io.Copy(f, bytes.NewBuffer([]byte(refStr)))
	return err
}

func createManifestFile(root, str string) (descriptor, error) {
	name := filepath.Join(root, "blobs", "sha256", "test-manifest")
	f, err := os.Create(name)
	if err != nil {
		return descriptor{}, err
	}
	defer f.Close()

	_, err = io.Copy(f, bytes.NewBuffer([]byte(str)))
	if err != nil {
		return descriptor{}, err
	}

	return createHashedBlob(name)
}

func createConfigFile(root, config string) (descriptor, error) {
	name := filepath.Join(root, "blobs", "sha256", "test-config")
	f, err := os.Create(name)
	if err != nil {
		return descriptor{}, err
	}
	defer f.Close()

	_, err = io.Copy(f, bytes.NewBuffer([]byte(config)))
	if err != nil {
		return descriptor{}, err
	}

	return createHashedBlob(name)
}

func createImageLayerFile(root string, list []tarContent) (descriptor, error) {
	name := filepath.Join(root, "blobs", "sha256", "test-layer")
	err := createTarBlob(name, list)
	if err != nil {
		return descriptor{}, err
	}

	desc, err := createHashedBlob(name)
	if err != nil {
		return descriptor{}, err
	}

	desc.MediaType = v1.MediaTypeImageLayer
	return desc, nil
}

func createTarBlob(name string, list []tarContent) error {
	file, err := os.Create(name)
	if err != nil {
		return err
	}
	defer file.Close()
	gzipWriter := gzip.NewWriter(file)
	defer gzipWriter.Close()
	tarWriter := tar.NewWriter(gzipWriter)
	defer tarWriter.Close()

	for _, content := range list {
		if err = tarWriter.WriteHeader(content.header); err != nil {
			return err
		}
		if _, err = io.Copy(tarWriter, bytes.NewReader(content.b)); err != nil {
			return err
		}
	}
	return nil
}

func createHashedBlob(name string) (descriptor, error) {
	desc, err := newDescriptor(name)
	if err != nil {
		return descriptor{}, err
	}

	// Rename the file to hashed-digest name.
	err = os.Rename(name, filepath.Join(filepath.Dir(name), desc.Digest))
	if err != nil {
		return descriptor{}, err
	}

	//Normalize the hashed digest.
	desc.Digest = "sha256:" + desc.Digest

	return desc, nil
}

func newDescriptor(name string) (descriptor, error) {
	file, err := os.Open(name)
	if err != nil {
		return descriptor{}, err
	}
	defer file.Close()

	// generate sha256 hash
	hash := sha256.New()
	size, err := io.Copy(hash, file)
	if err != nil {
		return descriptor{}, err
	}

	return descriptor{
		Digest: fmt.Sprintf("%x", hash.Sum(nil)),
		Size:   size,
	}, nil
}
<commit_msg>image_test.go: Improve the test content<commit_after>// Copyright 2016 The Linux Foundation
//
// 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 image

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"testing"

	"github.com/opencontainers/image-spec/specs-go/v1"
)

const (
	refTag = "latest"

	layoutStr = `{"imageLayoutVersion": "1.0.0"}`

	configStr = `{
    "created": "2015-10-31T22:22:56.015925234Z",
    "author": "Alyssa P. Hacker <alyspdev@example.com>",
    "architecture": "amd64",
    "os": "linux",
    "config": {
        "User": "alice",
        "Memory": 2048,
        "MemorySwap": 4096,
        "CpuShares": 8,
        "ExposedPorts": {
            "8080/tcp": {}
        },
        "Env": [
            "PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
            "FOO=oci_is_a",
            "BAR=well_written_spec"
        ],
        "Entrypoint": [
            "/bin/my-app-binary"
        ],
        "Cmd": [
            "--foreground",
            "--config",
            "/etc/my-app.d/default.cfg"
        ],
        "Volumes": {
            "/var/job-result-data": {},
            "/var/log/my-app-logs": {}
        },
        "WorkingDir": "/home/alice",
        "Labels": {
            "com.example.project.git.url": "https://example.com/project.git",
            "com.example.project.git.commit": "45a939b2999782a3f005621a8d0f29aa387e1d6b"
        }
    },
    "rootfs": {
      "diff_ids": [
        "sha256:c6f988f4874bb0add23a778f753c65efe992244e148a1d2ec2a8b664fb66bbd1",
        "sha256:5f70bf18a086007016e948b04aed3b82103a36bea41755b6cddfaf10ace3c6ef"
      ],
      "type": "layers"
    },
    "history": [
      {
        "created": "2015-10-31T22:22:54.690851953Z",
        "created_by": "/bin/sh -c #(nop) ADD file:a3bc1e842b69636f9df5256c49c5374fb4eef1e281fe3f282c65fb853ee171c5 in /"
      },
      {
        "created": "2015-10-31T22:22:55.613815829Z",
        "created_by": "/bin/sh -c #(nop) CMD [\"sh\"]",
        "empty_layer": true
      }
    ]
}
`
)

var (
	refStr = `{"digest":"<manifest_digest>","mediaType":"application/vnd.oci.image.manifest.v1+json","size":<manifest_size>}`

	manifestStr = `{
    "annotations": null,
    "config": {
        "digest": "<config_digest>",
        "mediaType": "application/vnd.oci.image.config.v1+json",
        "size": <config_size>
    },
    "layers": [
        {
            "digest": "<layer_digest>",
            "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip",
            "size": <layer_size>
        }
    ],
    "mediaType": "application/vnd.oci.image.manifest.v1+json",
    "schemaVersion": 2
}
 `
)

type tarContent struct {
	header *tar.Header
	b      []byte
}

type imageLayout struct {
	rootDir  string
	layout   string
	ref      string
	manifest string
	config   string
	tarList  []tarContent
}

func TestValidateLayout(t *testing.T) {
	root, err := ioutil.TempDir("", "oci-test")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(root)

	il := imageLayout{
		rootDir:  root,
		layout:   layoutStr,
		ref:      refTag,
		manifest: manifestStr,
		config:   configStr,
		tarList: []tarContent{
			tarContent{&tar.Header{Name: "test", Size: 4, Mode: 0600}, []byte("test")},
		},
	}

	// create image layout bundle
	err = createImageLayoutBundle(il)
	if err != nil {
		t.Fatal(err)
	}

	err = ValidateLayout(root, []string{refTag}, nil)
	if err != nil {
		t.Fatal(err)
	}
}

func createImageLayoutBundle(il imageLayout) error {
	err := os.MkdirAll(filepath.Join(il.rootDir, "blobs", "sha256"), 0700)
	if err != nil {
		return err
	}

	err = os.MkdirAll(filepath.Join(il.rootDir, "refs"), 0700)
	if err != nil {
		return err
	}

	// create image layout file
	err = createLayoutFile(il.rootDir)
	if err != nil {
		return err
	}

	// create image layer blob file.
	desc, err := createImageLayerFile(il.rootDir, il.tarList)
	if err != nil {
		return err
	}
	il.manifest = strings.Replace(il.manifest, "<layer_digest>", desc.Digest, 1)
	il.manifest = strings.Replace(il.manifest, "<layer_size>", strconv.FormatInt(desc.Size, 10), 1)

	desc, err = createConfigFile(il.rootDir, il.config)
	if err != nil {
		return err
	}
	il.manifest = strings.Replace(il.manifest, "<config_digest>", desc.Digest, 1)
	il.manifest = strings.Replace(il.manifest, "<config_size>", strconv.FormatInt(desc.Size, 10), 1)

	// create manifest blob file
	desc, err = createManifestFile(il.rootDir, il.manifest)
	if err != nil {
		return err
	}

	return createRefFile(il.rootDir, il.ref, desc)
}

func createLayoutFile(root string) error {
	layoutPath := filepath.Join(root, "oci-layout")
	f, err := os.Create(layoutPath)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = io.Copy(f, bytes.NewBuffer([]byte(layoutStr)))
	return err
}

func createRefFile(root, ref string, mft descriptor) error {
	refpath := filepath.Join(root, "refs", ref)
	f, err := os.Create(refpath)
	if err != nil {
		return err
	}
	defer f.Close()
	refStr = strings.Replace(refStr, "<manifest_digest>", mft.Digest, -1)
	refStr = strings.Replace(refStr, "<manifest_size>", strconv.FormatInt(mft.Size, 10), -1)
	_, err = io.Copy(f, bytes.NewBuffer([]byte(refStr)))
	return err
}

func createManifestFile(root, str string) (descriptor, error) {
	name := filepath.Join(root, "blobs", "sha256", "test-manifest")
	f, err := os.Create(name)
	if err != nil {
		return descriptor{}, err
	}
	defer f.Close()

	_, err = io.Copy(f, bytes.NewBuffer([]byte(str)))
	if err != nil {
		return descriptor{}, err
	}

	return createHashedBlob(name)
}

func createConfigFile(root, config string) (descriptor, error) {
	name := filepath.Join(root, "blobs", "sha256", "test-config")
	f, err := os.Create(name)
	if err != nil {
		return descriptor{}, err
	}
	defer f.Close()

	_, err = io.Copy(f, bytes.NewBuffer([]byte(config)))
	if err != nil {
		return descriptor{}, err
	}

	return createHashedBlob(name)
}

func createImageLayerFile(root string, list []tarContent) (descriptor, error) {
	name := filepath.Join(root, "blobs", "sha256", "test-layer")
	err := createTarBlob(name, list)
	if err != nil {
		return descriptor{}, err
	}

	desc, err := createHashedBlob(name)
	if err != nil {
		return descriptor{}, err
	}

	desc.MediaType = v1.MediaTypeImageLayer
	return desc, nil
}

func createTarBlob(name string, list []tarContent) error {
	file, err := os.Create(name)
	if err != nil {
		return err
	}
	defer file.Close()
	gzipWriter := gzip.NewWriter(file)
	defer gzipWriter.Close()
	tarWriter := tar.NewWriter(gzipWriter)
	defer tarWriter.Close()

	for _, content := range list {
		if err = tarWriter.WriteHeader(content.header); err != nil {
			return err
		}
		if _, err = io.Copy(tarWriter, bytes.NewReader(content.b)); err != nil {
			return err
		}
	}
	return nil
}

func createHashedBlob(name string) (descriptor, error) {
	desc, err := newDescriptor(name)
	if err != nil {
		return descriptor{}, err
	}

	// Rename the file to hashed-digest name.
	err = os.Rename(name, filepath.Join(filepath.Dir(name), desc.Digest))
	if err != nil {
		return descriptor{}, err
	}

	//Normalize the hashed digest.
	desc.Digest = "sha256:" + desc.Digest

	return desc, nil
}

func newDescriptor(name string) (descriptor, error) {
	file, err := os.Open(name)
	if err != nil {
		return descriptor{}, err
	}
	defer file.Close()

	// generate sha256 hash
	hash := sha256.New()
	size, err := io.Copy(hash, file)
	if err != nil {
		return descriptor{}, err
	}

	return descriptor{
		Digest: fmt.Sprintf("%x", hash.Sum(nil)),
		Size:   size,
	}, nil
}
<|endoftext|>


<commit_before>package consensus

import (
	"bytes"
	"errors"
	"fmt"

	"github.com/rivine/rivine/build"
	"github.com/rivine/rivine/crypto"
	"github.com/rivine/rivine/encoding"
	"github.com/rivine/rivine/types"

	"github.com/coreos/bbolt"
)

// manageErr handles an error detected by the consistency checks.
func manageErr(tx *bolt.Tx, err error) {
	markInconsistency(tx)
	if build.DEBUG {
		panic(err)
	} else {
		fmt.Println(err)
	}
}

// consensusChecksum grabs a checksum of the consensus set by pushing all of
// the elements in sorted order into a merkle tree and taking the root. All
// consensus sets with the same current block should have identical consensus
// checksums.
func consensusChecksum(tx *bolt.Tx) crypto.Hash {
	// Create a checksum tree.
	tree := crypto.NewTree()

	// For all of the constant buckets, push every key and every value. Buckets
	// are sorted in byte-order, therefore this operation is deterministic.
	consensusSetBuckets := []*bolt.Bucket{
		tx.Bucket(BlockPath),
		tx.Bucket(CoinOutputs),
		tx.Bucket(BlockStakeOutputs),
	}
	for i := range consensusSetBuckets {
		err := consensusSetBuckets[i].ForEach(func(k, v []byte) error {
			tree.Push(k)
			tree.Push(v)
			return nil
		})
		if err != nil {
			manageErr(tx, err)
		}
	}

	// Iterate through all the buckets looking for buckets prefixed with
	// prefixDCO or prefixFCEX. Buckets are presented in byte-sorted order by
	// name.
	err := tx.ForEach(func(name []byte, b *bolt.Bucket) error {
		// If the bucket is not a delayed coin output bucket or a file
		// contract expiration bucket, skip.
		if !bytes.HasPrefix(name, prefixDCO) && !bytes.HasPrefix(name, prefixFCEX) {
			return nil
		}

		// The bucket is a prefixed bucket - add all elements to the tree.
		return b.ForEach(func(k, v []byte) error {
			tree.Push(k)
			tree.Push(v)
			return nil
		})
	})
	if err != nil {
		manageErr(tx, err)
	}

	return tree.Root()
}

// checkCoinCount checks that the number of coins countable within the
// consensus set equal the expected number of coins for the block height.
func checkCoinCount(tx *bolt.Tx, bh types.BlockHeight) {

	// Add all of the siacoin outputs.
	var scocoins types.Currency
	err := tx.Bucket(CoinOutputs).ForEach(func(_, scoBytes []byte) error {
		var sco types.CoinOutput
		err := encoding.Unmarshal(scoBytes, &sco)
		if err != nil {
			manageErr(tx, err)
		}
		scocoins = scocoins.Add(sco.Value)
		return nil
	})
	if err != nil {
		manageErr(tx, err)
	}

	//special check if maturity is bigger than the current block height
	delay := types.MaturityDelay - 1
	if bh < types.MaturityDelay {
		delay = bh
	}

	expectedcoins := types.GenesisCoinCount.Add(types.BlockCreatorFee.Mul64(uint64(bh - delay)))
	totalcoins := scocoins
	if !totalcoins.Equals(expectedcoins) {
		diagnostics := fmt.Sprintf("Wrong number of coins: %v\n", scocoins)
		if totalcoins.Cmp(expectedcoins) < 0 {
			diagnostics += fmt.Sprintf("total: %v\nexpected: %v\n expected is bigger: %v", totalcoins, expectedcoins, expectedcoins.Sub(totalcoins))
		} else {
			diagnostics += fmt.Sprintf("total: %v\nexpected: %v\n expected is smaler: %v", totalcoins, expectedcoins, totalcoins.Sub(expectedcoins))
		}
		manageErr(tx, errors.New(diagnostics))
	}
}

// checkBlockStakeCount checks that the number of siafunds countable within the
// consensus set equal the expected number of BlockStakeOutputs for the block height.
func checkBlockStakeCount(tx *bolt.Tx) {
	var total types.Currency
	err := tx.Bucket(BlockStakeOutputs).ForEach(func(_, siafundOutputBytes []byte) error {
		var sfo types.BlockStakeOutput
		err := encoding.Unmarshal(siafundOutputBytes, &sfo)
		if err != nil {
			manageErr(tx, err)
		}
		total = total.Add(sfo.Value)
		return nil
	})
	if err != nil {
		manageErr(tx, err)
	}
	if !total.Equals(types.GenesisBlockStakeCount) {
		manageErr(tx, errors.New("wrong number if blockstakes in the consensus set"))
	}
}

// checkRevertApply reverts the most recent block, checking to see that the
// consensus set hash matches the hash obtained for the previous block. Then it
// applies the block again and checks that the consensus set hash matches the
// original consensus set hash.
func (cs *ConsensusSet) checkRevertApply(tx *bolt.Tx) {
	current := currentProcessedBlock(tx)
	// Don't perform the check if this block is the genesis block.
	if current.Block.ID() == cs.blockRoot.Block.ID() {
		return
	}

	parent, err := getBlockMap(tx, current.Block.ParentID)
	if err != nil {
		manageErr(tx, err)
	}
	if current.Height != parent.Height+1 {
		manageErr(tx, errors.New("parent structure of a block is incorrect"))
	}
	_, _, err = cs.forkBlockchain(tx, parent)
	if err != nil {
		manageErr(tx, err)
	}
	if consensusChecksum(tx) != parent.ConsensusChecksum {
		//manageErr(tx, errors.New("consensus checksum mismatch after reverting")) //this fails issue #28
	}
	_, _, err = cs.forkBlockchain(tx, current)
	if err != nil {
		manageErr(tx, err)
	}
	if consensusChecksum(tx) != current.ConsensusChecksum {
		manageErr(tx, errors.New("consensus checksum mismatch after re-applying"))
	}
}

// checkConsistency runs a series of checks to make sure that the consensus set
// is consistent with some rules that should always be true.
func (cs *ConsensusSet) checkConsistency(tx *bolt.Tx) {
	if cs.checkingConsistency {
		return
	}
	cs.checkingConsistency = true
	checkCoinCount(tx, cs.Height())
	checkBlockStakeCount(tx)
	if build.DEBUG {
		cs.checkRevertApply(tx)
	}
	cs.checkingConsistency = false
}

// maybeCheckConsistency runs a consistency check with a small probability.
// Useful for detecting database corruption in production without needing to go
// through the extremely slow process of running a consistency check every
// block.
func (cs *ConsensusSet) maybeCheckConsistency(tx *bolt.Tx) {
	n, err := crypto.RandIntn(1000)
	if err != nil {
		manageErr(tx, err)
	}
	if n == 0 {
		cs.checkConsistency(tx)
	}
}

// TODO: Check that every file contract has an expiration too, and that the
// number of file contracts + the number of expirations is equal.
<commit_msg>Remove checkCoinCount function<commit_after>package consensus

import (
	"bytes"
	"errors"
	"fmt"

	"github.com/rivine/rivine/build"
	"github.com/rivine/rivine/crypto"
	"github.com/rivine/rivine/encoding"
	"github.com/rivine/rivine/types"

	"github.com/coreos/bbolt"
)

// manageErr handles an error detected by the consistency checks.
func manageErr(tx *bolt.Tx, err error) {
	markInconsistency(tx)
	if build.DEBUG {
		panic(err)
	} else {
		fmt.Println(err)
	}
}

// consensusChecksum grabs a checksum of the consensus set by pushing all of
// the elements in sorted order into a merkle tree and taking the root. All
// consensus sets with the same current block should have identical consensus
// checksums.
func consensusChecksum(tx *bolt.Tx) crypto.Hash {
	// Create a checksum tree.
	tree := crypto.NewTree()

	// For all of the constant buckets, push every key and every value. Buckets
	// are sorted in byte-order, therefore this operation is deterministic.
	consensusSetBuckets := []*bolt.Bucket{
		tx.Bucket(BlockPath),
		tx.Bucket(CoinOutputs),
		tx.Bucket(BlockStakeOutputs),
	}
	for i := range consensusSetBuckets {
		err := consensusSetBuckets[i].ForEach(func(k, v []byte) error {
			tree.Push(k)
			tree.Push(v)
			return nil
		})
		if err != nil {
			manageErr(tx, err)
		}
	}

	// Iterate through all the buckets looking for buckets prefixed with
	// prefixDCO or prefixFCEX. Buckets are presented in byte-sorted order by
	// name.
	err := tx.ForEach(func(name []byte, b *bolt.Bucket) error {
		// If the bucket is not a delayed coin output bucket or a file
		// contract expiration bucket, skip.
		if !bytes.HasPrefix(name, prefixDCO) && !bytes.HasPrefix(name, prefixFCEX) {
			return nil
		}

		// The bucket is a prefixed bucket - add all elements to the tree.
		return b.ForEach(func(k, v []byte) error {
			tree.Push(k)
			tree.Push(v)
			return nil
		})
	})
	if err != nil {
		manageErr(tx, err)
	}

	return tree.Root()
}

// checkBlockStakeCount checks that the number of siafunds countable within the
// consensus set equal the expected number of BlockStakeOutputs for the block height.
func checkBlockStakeCount(tx *bolt.Tx) {
	var total types.Currency
	err := tx.Bucket(BlockStakeOutputs).ForEach(func(_, siafundOutputBytes []byte) error {
		var sfo types.BlockStakeOutput
		err := encoding.Unmarshal(siafundOutputBytes, &sfo)
		if err != nil {
			manageErr(tx, err)
		}
		total = total.Add(sfo.Value)
		return nil
	})
	if err != nil {
		manageErr(tx, err)
	}
	if !total.Equals(types.GenesisBlockStakeCount) {
		manageErr(tx, errors.New("wrong number if blockstakes in the consensus set"))
	}
}

// checkRevertApply reverts the most recent block, checking to see that the
// consensus set hash matches the hash obtained for the previous block. Then it
// applies the block again and checks that the consensus set hash matches the
// original consensus set hash.
func (cs *ConsensusSet) checkRevertApply(tx *bolt.Tx) {
	current := currentProcessedBlock(tx)
	// Don't perform the check if this block is the genesis block.
	if current.Block.ID() == cs.blockRoot.Block.ID() {
		return
	}

	parent, err := getBlockMap(tx, current.Block.ParentID)
	if err != nil {
		manageErr(tx, err)
	}
	if current.Height != parent.Height+1 {
		manageErr(tx, errors.New("parent structure of a block is incorrect"))
	}
	_, _, err = cs.forkBlockchain(tx, parent)
	if err != nil {
		manageErr(tx, err)
	}
	if consensusChecksum(tx) != parent.ConsensusChecksum {
		//manageErr(tx, errors.New("consensus checksum mismatch after reverting")) //this fails issue #28
	}
	_, _, err = cs.forkBlockchain(tx, current)
	if err != nil {
		manageErr(tx, err)
	}
	if consensusChecksum(tx) != current.ConsensusChecksum {
		manageErr(tx, errors.New("consensus checksum mismatch after re-applying"))
	}
}

// checkConsistency runs a series of checks to make sure that the consensus set
// is consistent with some rules that should always be true.
func (cs *ConsensusSet) checkConsistency(tx *bolt.Tx) {
	if cs.checkingConsistency {
		return
	}
	cs.checkingConsistency = true
	checkBlockStakeCount(tx)
	if build.DEBUG {
		cs.checkRevertApply(tx)
	}
	cs.checkingConsistency = false
}

// maybeCheckConsistency runs a consistency check with a small probability.
// Useful for detecting database corruption in production without needing to go
// through the extremely slow process of running a consistency check every
// block.
func (cs *ConsensusSet) maybeCheckConsistency(tx *bolt.Tx) {
	n, err := crypto.RandIntn(1000)
	if err != nil {
		manageErr(tx, err)
	}
	if n == 0 {
		cs.checkConsistency(tx)
	}
}

// TODO: Check that every file contract has an expiration too, and that the
// number of file contracts + the number of expirations is equal.
<|endoftext|>


<commit_before>package parse

import (
	"github.com/demizer/go-elog"
	"testing"
	"bufio"
	"os"
	"bytes"
	"strings"
	"flag"
	"fmt"
)

type LexTest struct {
	name           string
	description    string
	data           string
	items          string
	expect         string
	collectedItems []item
}

type LexTests []LexTest

func (l LexTests) SearchByName(name string) *LexTest {
	for _, test := range l {
		if test.name == name {
			return &test
		}
	}
	return nil
}

var Tests LexTests

func ParseTestData(t *testing.T, filepath string) ([]LexTest, error) {
	testData, err := os.Open(filepath)
	defer testData.Close()
	if err != nil {
		return nil, err
	}

	var LexTests []LexTest
	var curTest = new(LexTest)
	var buffer bytes.Buffer

	scanner := bufio.NewScanner(testData)

	for scanner.Scan() {
		switch scanner.Text() {
		case "#name":
			// buffer = bytes.NewBuffer(buffer.Bytes())
			// name starts a new section
			if buffer.Len() > 0 {
				// Apend the last section to the array and
				// reset
				curTest.expect = buffer.String()
				LexTests = append(LexTests, *curTest)
			}
			curTest = new(LexTest)
			buffer.Reset()
		case "#description":
			curTest.name = strings.TrimRight(buffer.String(), "\n")
			buffer.Reset()
		case "#data":
			curTest.description = strings.TrimRight(buffer.String(), "\n")
			buffer.Reset()
		case "#items":
			curTest.data = strings.TrimRight(buffer.String(), "\n")
			buffer.Reset()
		case "#parse-expect":
			curTest.items = buffer.String()
			buffer.Reset()
		default:
			// Collect the text in between sections
			buffer.WriteString(fmt.Sprintln(scanner.Text()))
		}
	}

	if err := scanner.Err(); err != nil {
		t.Error(err)
	}

	if buffer.Len() > 0 {
		// Apend the last section to the array and
		curTest.expect = buffer.String()
		LexTests = append(LexTests, *curTest)
	}

	return LexTests, nil
}

// SetDebug is typically called from the init() function in a test file. SetDebug parses debug flags
// passed to the test binary and also sets the template for logging output.
func SetDebug() {
	var debug bool
	flag.BoolVar(&debug, "debug", false, "Enable debug output.")
	flag.Parse()
	if debug {
		log.SetLevel(log.LEVEL_DEBUG)
	}

	log.SetTemplate("{{if .Date}}{{.Date}} {{end}}" +
		"{{if .Prefix}}{{.Prefix}} {{end}}" +
		"{{if .LogLabel}}{{.LogLabel}} {{end}}" +
		"{{if .FileName}}{{.FileName}}: {{end}}" +
		"{{if .FunctionName}}{{.FunctionName}}{{end}}" +
		"{{if .LineNumber}}#{{.LineNumber}}: {{end}}" +
		"{{if .Text}}{{.Text}}{{end}}")

	log.SetFlags(log.Lansi | log.LnoPrefix | log.LfunctionName |
		log.LlineNumber)
}

<commit_msg>util_test.go: Add documentation<commit_after>// go-rst - A reStructuredText parser for Go
// 2014 (c) The go-rst Authors
// MIT Licensed. See LICENSE for details.

package parse

import (
	"github.com/demizer/go-elog"
	"bufio"
	"os"
	"bytes"
	"strings"
	"flag"
	"fmt"
)

func init() { SetDebug() }

// LexTest is the structure that contains parsed test data from the *.dat files in the testdata
// directory.
type LexTest struct {
	name           string
	description    string
	data           string  // The input data to be parsed
	items          string  // The expected lex items output in json
	expect         string  // The expected parsed output in json
}

type LexTests []LexTest

// Search l by name for a specific test.
func (l LexTests) SearchByName(name string) *LexTest {
	for _, test := range l {
		if test.name == name {
			return &test
		}
	}
	return nil
}

// ParseTestData parses testdata contained it dat files in the testdata directory. The testdata was
// contained to these files because it became to large to be included legibly inside the *_test.go
// files. ParseTestData is a simple parser for the testdata files and stores the result of the parse
// into the first return variable.
func ParseTestData(filepath string) ([]LexTest, error) {
	testData, err := os.Open(filepath)
	defer testData.Close()
	if err != nil {
		return nil, err
	}

	var LexTests []LexTest
	var curTest = new(LexTest)
	var buffer bytes.Buffer

	scanner := bufio.NewScanner(testData)

	for scanner.Scan() {
		switch scanner.Text() {
		case "#name":
			// buffer = bytes.NewBuffer(buffer.Bytes())
			// name starts a new section
			if buffer.Len() > 0 {
				// Apend the last section to the array and
				// reset
				curTest.expect = buffer.String()
				LexTests = append(LexTests, *curTest)
			}
			curTest = new(LexTest)
			buffer.Reset()
		case "#description":
			curTest.name = strings.TrimRight(buffer.String(), "\n")
			buffer.Reset()
		case "#data":
			curTest.description = strings.TrimRight(buffer.String(), "\n")
			buffer.Reset()
		case "#items":
			curTest.data = strings.TrimRight(buffer.String(), "\n")
			buffer.Reset()
		case "#parse-expect":
			curTest.items = buffer.String()
			buffer.Reset()
		default:
			// Collect the text in between sections
			buffer.WriteString(fmt.Sprintln(scanner.Text()))
		}
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	if buffer.Len() > 0 {
		// Apend the last section to the array and
		curTest.expect = buffer.String()
		LexTests = append(LexTests, *curTest)
	}

	return LexTests, nil
}

// SetDebug is typically called from the init() function in a test file. SetDebug parses debug flags
// passed to the test binary and also sets the template for logging output.
func SetDebug() {
	var debug bool

	flag.BoolVar(&debug, "debug", false, "Enable debug output.")
	flag.Parse()

	if debug {
		log.SetLevel(log.LEVEL_DEBUG)
	}

	log.SetTemplate("{{if .Date}}{{.Date}} {{end}}" +
		"{{if .Prefix}}{{.Prefix}} {{end}}" +
		"{{if .LogLabel}}{{.LogLabel}} {{end}}" +
		"{{if .FileName}}{{.FileName}}: {{end}}" +
		"{{if .FunctionName}}{{.FunctionName}}{{end}}" +
		"{{if .LineNumber}}#{{.LineNumber}}: {{end}}" +
		"{{if .Text}}{{.Text}}{{end}}")

	log.SetFlags(log.Lansi | log.LnoPrefix | log.LfunctionName |
		log.LlineNumber)
}

<|endoftext|>


<commit_before>/*
Copyright 2021 The Vitess Authors.

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 vtgate

import (
	"context"
	"fmt"
	"testing"

	"github.com/google/go-cmp/cmp"

	"vitess.io/vitess/go/test/endtoend/cluster"

	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
)

func TestFunctionInDefault(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	// set the sql mode ALLOW_INVALID_DATES
	exec(t, conn, `SET sql_mode = 'ALLOW_INVALID_DATES'`)

	exec(t, conn, `create table function_default (x varchar(25) DEFAULT (TRIM(" check ")))`)
	exec(t, conn, "drop table function_default")

	exec(t, conn, `create table function_default (
ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
dt DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
ts2 TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
dt2 DATETIME DEFAULT CURRENT_TIMESTAMP,
ts3 TIMESTAMP DEFAULT 0,
dt3 DATETIME DEFAULT 0,
ts4 TIMESTAMP DEFAULT 0 ON UPDATE CURRENT_TIMESTAMP,
dt4 DATETIME DEFAULT 0 ON UPDATE CURRENT_TIMESTAMP,
ts5 TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
ts6 TIMESTAMP NULL ON UPDATE CURRENT_TIMESTAMP,
dt5 DATETIME ON UPDATE CURRENT_TIMESTAMP,
dt6 DATETIME NOT NULL ON UPDATE CURRENT_TIMESTAMP,
ts7 TIMESTAMP(6) DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6),
ts8 TIMESTAMP DEFAULT NOW(),
ts9 TIMESTAMP DEFAULT LOCALTIMESTAMP,
ts10 TIMESTAMP DEFAULT LOCALTIME,
ts11 TIMESTAMP DEFAULT LOCALTIMESTAMP(),
ts12 TIMESTAMP DEFAULT LOCALTIME()
)`)
	exec(t, conn, "drop table function_default")

	// this query works because utc_timestamp will get parenthesised before reaching MySQL. However, this syntax is not supported in MySQL 8.0
	exec(t, conn, `create table function_default (ts TIMESTAMP DEFAULT UTC_TIMESTAMP)`)
	exec(t, conn, "drop table function_default")

	exec(t, conn, `create table function_default (x varchar(25) DEFAULT "check")`)
	exec(t, conn, "drop table function_default")
}

// TestCheckConstraint test check constraints on CREATE TABLE
// This feature is supported from MySQL 8.0.16 and MariaDB 10.2.1.
func TestCheckConstraint(t *testing.T) {
	conn, err := mysql.Connect(context.Background(), &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	query := `CREATE TABLE t7 (CHECK (c1 <> c2), c1 INT CHECK (c1 > 10), c2 INT CONSTRAINT c2_positive CHECK (c2 > 0), c3 INT CHECK (c3 < 100), CONSTRAINT c1_nonzero CHECK (c1 <> 0), CHECK (c1 > c3));`
	exec(t, conn, query)

	checkQuery := `SELECT CONSTRAINT_NAME FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS WHERE TABLE_NAME = 't7' order by CONSTRAINT_NAME;`
	expected := `[[VARCHAR("c1_nonzero")] [VARCHAR("c2_positive")] [VARCHAR("t7_chk_1")] [VARCHAR("t7_chk_2")] [VARCHAR("t7_chk_3")] [VARCHAR("t7_chk_4")]]`

	assertMatches(t, conn, checkQuery, expected)

	cleanup := `DROP TABLE t7`
	exec(t, conn, cleanup)
}

func TestValueDefault(t *testing.T) {
	vtParams := mysql.ConnParams{
		Host: "localhost",
		Port: clusterInstance.VtgateMySQLPort,
	}
	conn, err := mysql.Connect(context.Background(), &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	exec(t, conn, `create table test_float_default (pos_f float default 2.1, neg_f float default -2.1,b blob default ('abc'));`)
	defer exec(t, conn, `drop table test_float_default`)
	assertMatches(t, conn, "select table_name, column_name, column_default from information_schema.columns where table_name = 'test_float_default' order by column_name", `[[VARCHAR("test_float_default") VARCHAR("b") BLOB("_utf8mb3\\'abc\\'")] [VARCHAR("test_float_default") VARCHAR("neg_f") BLOB("-2.1")] [VARCHAR("test_float_default") VARCHAR("pos_f") BLOB("2.1")]]`)
}

func TestVersionCommentWorks(t *testing.T) {
	conn, err := mysql.Connect(context.Background(), &vtParams)
	require.NoError(t, err)
	defer conn.Close()
	exec(t, conn, "/*!80000 SET SESSION information_schema_stats_expiry=0 */")
}

func assertMatches(t *testing.T, conn *mysql.Conn, query, expected string) {
	t.Helper()
	qr := exec(t, conn, query)
	got := fmt.Sprintf("%v", qr.Rows)
	diff := cmp.Diff(expected, got)
	if diff != "" {
		t.Errorf("Query: %s (-want +got):\n%s", query, diff)
	}
}

func exec(t *testing.T, conn *mysql.Conn, query string) *sqltypes.Result {
	t.Helper()
	qr, err := conn.ExecuteFetch(query, 1000, true)
	require.NoError(t, err, "for query: "+query)
	return qr
}
<commit_msg>fixed test output<commit_after>/*
Copyright 2021 The Vitess Authors.

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 vtgate

import (
	"context"
	"fmt"
	"testing"

	"github.com/google/go-cmp/cmp"

	"vitess.io/vitess/go/test/endtoend/cluster"

	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
)

func TestFunctionInDefault(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	// set the sql mode ALLOW_INVALID_DATES
	exec(t, conn, `SET sql_mode = 'ALLOW_INVALID_DATES'`)

	exec(t, conn, `create table function_default (x varchar(25) DEFAULT (TRIM(" check ")))`)
	exec(t, conn, "drop table function_default")

	exec(t, conn, `create table function_default (
ts TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
dt DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
ts2 TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
dt2 DATETIME DEFAULT CURRENT_TIMESTAMP,
ts3 TIMESTAMP DEFAULT 0,
dt3 DATETIME DEFAULT 0,
ts4 TIMESTAMP DEFAULT 0 ON UPDATE CURRENT_TIMESTAMP,
dt4 DATETIME DEFAULT 0 ON UPDATE CURRENT_TIMESTAMP,
ts5 TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
ts6 TIMESTAMP NULL ON UPDATE CURRENT_TIMESTAMP,
dt5 DATETIME ON UPDATE CURRENT_TIMESTAMP,
dt6 DATETIME NOT NULL ON UPDATE CURRENT_TIMESTAMP,
ts7 TIMESTAMP(6) DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6),
ts8 TIMESTAMP DEFAULT NOW(),
ts9 TIMESTAMP DEFAULT LOCALTIMESTAMP,
ts10 TIMESTAMP DEFAULT LOCALTIME,
ts11 TIMESTAMP DEFAULT LOCALTIMESTAMP(),
ts12 TIMESTAMP DEFAULT LOCALTIME()
)`)
	exec(t, conn, "drop table function_default")

	// this query works because utc_timestamp will get parenthesised before reaching MySQL. However, this syntax is not supported in MySQL 8.0
	exec(t, conn, `create table function_default (ts TIMESTAMP DEFAULT UTC_TIMESTAMP)`)
	exec(t, conn, "drop table function_default")

	exec(t, conn, `create table function_default (x varchar(25) DEFAULT "check")`)
	exec(t, conn, "drop table function_default")
}

// TestCheckConstraint test check constraints on CREATE TABLE
// This feature is supported from MySQL 8.0.16 and MariaDB 10.2.1.
func TestCheckConstraint(t *testing.T) {
	conn, err := mysql.Connect(context.Background(), &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	query := `CREATE TABLE t7 (CHECK (c1 <> c2), c1 INT CHECK (c1 > 10), c2 INT CONSTRAINT c2_positive CHECK (c2 > 0), c3 INT CHECK (c3 < 100), CONSTRAINT c1_nonzero CHECK (c1 <> 0), CHECK (c1 > c3));`
	exec(t, conn, query)

	checkQuery := `SELECT CONSTRAINT_NAME FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS WHERE TABLE_NAME = 't7' order by CONSTRAINT_NAME;`
	expected := `[[VARCHAR("c1_nonzero")] [VARCHAR("c2_positive")] [VARCHAR("t7_chk_1")] [VARCHAR("t7_chk_2")] [VARCHAR("t7_chk_3")] [VARCHAR("t7_chk_4")]]`

	assertMatches(t, conn, checkQuery, expected)

	cleanup := `DROP TABLE t7`
	exec(t, conn, cleanup)
}

func TestValueDefault(t *testing.T) {
	vtParams := mysql.ConnParams{
		Host: "localhost",
		Port: clusterInstance.VtgateMySQLPort,
	}
	conn, err := mysql.Connect(context.Background(), &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	exec(t, conn, `create table test_float_default (pos_f float default 2.1, neg_f float default -2.1,b blob default ('abc'));`)
	defer exec(t, conn, `drop table test_float_default`)
	assertMatches(t, conn, "select table_name, column_name, column_default from information_schema.columns where table_name = 'test_float_default' order by column_name", `[[VARBINARY("test_float_default") VARCHAR("b") BLOB("_utf8mb3\\'abc\\'")] [VARBINARY("test_float_default") VARCHAR("neg_f") BLOB("-2.1")] [VARBINARY("test_float_default") VARCHAR("pos_f") BLOB("2.1")]]`)
}

func TestVersionCommentWorks(t *testing.T) {
	conn, err := mysql.Connect(context.Background(), &vtParams)
	require.NoError(t, err)
	defer conn.Close()
	exec(t, conn, "/*!80000 SET SESSION information_schema_stats_expiry=0 */")
}

func assertMatches(t *testing.T, conn *mysql.Conn, query, expected string) {
	t.Helper()
	qr := exec(t, conn, query)
	got := fmt.Sprintf("%v", qr.Rows)
	diff := cmp.Diff(expected, got)
	if diff != "" {
		t.Errorf("Query: %s (-want +got):\n%s", query, diff)
	}
}

func exec(t *testing.T, conn *mysql.Conn, query string) *sqltypes.Result {
	t.Helper()
	qr, err := conn.ExecuteFetch(query, 1000, true)
	require.NoError(t, err, "for query: "+query)
	return qr
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Vitess Authors.

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 discovery

import (
	"context"
	"errors"
	"fmt"

	"github.com/spf13/pflag"

	vtadminpb "vitess.io/vitess/go/vt/proto/vtadmin"
)

var (
	// ErrImplementationNotRegistered is returned from discovery.New if it is
	// called with an unknown implementation.
	ErrImplementationNotRegistered = errors.New("no discovery factory registered for implementation")
	// ErrNoVTGates should be returned from DiscoverVTGate* methods when they
	// are unable to find any vtgates for the given filter/query/tags.
	ErrNoVTGates = errors.New("no vtgates found")
	// ErrNoVtctlds should be returned from DiscoverVtctld* methods when they
	// are unable to find any vtctlds for the given filter/query/tags.
	ErrNoVtctlds = errors.New("no vtctlds found")
)

// Discovery defines the interface that service discovery plugins must
// implement. See ConsulDiscovery for an example implementation.
type Discovery interface {
	// DiscoverVTGate returns a vtgate found in the discovery service.
	// Tags can optionally be used to filter the set of potential gates further.
	// Which gate in a set of found gates is returned is not specified by the
	// interface, and can be implementation-specific.
	DiscoverVTGate(ctx context.Context, tags []string) (*vtadminpb.VTGate, error)
	// DiscoverVTGateAddr returns the address of a of vtgate found in the
	// discovery service. Tags can optionally be used to filter the set of
	// potential gates further. Which gate in a set of found gates is used to
	// return an address is not specified by the interface, and can be
	// implementation-specific.
	DiscoverVTGateAddr(ctx context.Context, tags []string) (string, error)
	// DiscoverVTGates returns a list of vtgates found in the discovery service.
	// Tags can optionally be used to filter gates. Order of the gates is not
	// specified by the interface, and can be implementation-specific.
	DiscoverVTGates(ctx context.Context, tags []string) ([]*vtadminpb.VTGate, error)
	// DiscoverVtctld returns a vtctld found in the discovery service.
	// Tags can optionally be used to filter the set of potential vtctlds
	// further. Which vtctld in a set of found vtctlds is returned is not
	// specified by the interface, and can be implementation-specific.
	DiscoverVtctld(ctx context.Context, tags []string) (*vtadminpb.Vtctld, error)
	// DiscoverVtctldAddr returns the address of a vtctld found in the discovery
	// service. Tags can optionally be used to filter the set of potential
	// vtctlds further. Which vtctld in a set of potential vtctld is used to
	// return an address is not specified by the interface, and can be
	// implementation-specific.
	DiscoverVtctldAddr(ctx context.Context, tags []string) (string, error)
	// DiscoverVtctlds returns a list of vtctlds found in the discovery service.
	// Tags can optionally be used to filter vtctlds. Order of the vtctlds is
	// not specified by the interface, and can be implementation-specific.
	DiscoverVtctlds(ctx context.Context, tags []string) ([]*vtadminpb.Vtctld, error)
}

// Factory represents a function that can create a Discovery implementation.
// This package will provide several implementations and register them for use.
// The flags FlagSet is provided for convenience, but also to hint to plugin
// developers that they should expect the args to be in a format compatible with
// pflag.
type Factory func(cluster *vtadminpb.Cluster, flags *pflag.FlagSet, args []string) (Discovery, error)

// nolint:gochecknoglobals
var registry = map[string]Factory{}

// Register registers a factory for the given implementation name. Attempting
// to register multiple factories for the same implementation name causes a
// panic.
func Register(name string, factory Factory) {
	_, ok := registry[name]
	if ok {
		panic("[discovery] factory already registered for " + name)
	}

	registry[name] = factory
}

// New returns a Discovery implementation using the registered factory for the
// implementation. Usage of the args slice is dependent on the implementation's
// factory.
func New(impl string, cluster *vtadminpb.Cluster, args []string) (Discovery, error) {
	factory, ok := registry[impl]
	if !ok {
		return nil, fmt.Errorf("%w %s", ErrImplementationNotRegistered, impl)
	}

	return factory(cluster, pflag.NewFlagSet("discovery:"+impl, pflag.ContinueOnError), args)
}

func init() { // nolint:gochecknoinits
	Register("consul", NewConsul)
	Register("staticFile", NewStaticFile)
}
<commit_msg>Rename staticFile=>staticfile to match the go/vitess _style_<commit_after>/*
Copyright 2020 The Vitess Authors.

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 discovery

import (
	"context"
	"errors"
	"fmt"

	"github.com/spf13/pflag"

	vtadminpb "vitess.io/vitess/go/vt/proto/vtadmin"
)

var (
	// ErrImplementationNotRegistered is returned from discovery.New if it is
	// called with an unknown implementation.
	ErrImplementationNotRegistered = errors.New("no discovery factory registered for implementation")
	// ErrNoVTGates should be returned from DiscoverVTGate* methods when they
	// are unable to find any vtgates for the given filter/query/tags.
	ErrNoVTGates = errors.New("no vtgates found")
	// ErrNoVtctlds should be returned from DiscoverVtctld* methods when they
	// are unable to find any vtctlds for the given filter/query/tags.
	ErrNoVtctlds = errors.New("no vtctlds found")
)

// Discovery defines the interface that service discovery plugins must
// implement. See ConsulDiscovery for an example implementation.
type Discovery interface {
	// DiscoverVTGate returns a vtgate found in the discovery service.
	// Tags can optionally be used to filter the set of potential gates further.
	// Which gate in a set of found gates is returned is not specified by the
	// interface, and can be implementation-specific.
	DiscoverVTGate(ctx context.Context, tags []string) (*vtadminpb.VTGate, error)
	// DiscoverVTGateAddr returns the address of a of vtgate found in the
	// discovery service. Tags can optionally be used to filter the set of
	// potential gates further. Which gate in a set of found gates is used to
	// return an address is not specified by the interface, and can be
	// implementation-specific.
	DiscoverVTGateAddr(ctx context.Context, tags []string) (string, error)
	// DiscoverVTGates returns a list of vtgates found in the discovery service.
	// Tags can optionally be used to filter gates. Order of the gates is not
	// specified by the interface, and can be implementation-specific.
	DiscoverVTGates(ctx context.Context, tags []string) ([]*vtadminpb.VTGate, error)
	// DiscoverVtctld returns a vtctld found in the discovery service.
	// Tags can optionally be used to filter the set of potential vtctlds
	// further. Which vtctld in a set of found vtctlds is returned is not
	// specified by the interface, and can be implementation-specific.
	DiscoverVtctld(ctx context.Context, tags []string) (*vtadminpb.Vtctld, error)
	// DiscoverVtctldAddr returns the address of a vtctld found in the discovery
	// service. Tags can optionally be used to filter the set of potential
	// vtctlds further. Which vtctld in a set of potential vtctld is used to
	// return an address is not specified by the interface, and can be
	// implementation-specific.
	DiscoverVtctldAddr(ctx context.Context, tags []string) (string, error)
	// DiscoverVtctlds returns a list of vtctlds found in the discovery service.
	// Tags can optionally be used to filter vtctlds. Order of the vtctlds is
	// not specified by the interface, and can be implementation-specific.
	DiscoverVtctlds(ctx context.Context, tags []string) ([]*vtadminpb.Vtctld, error)
}

// Factory represents a function that can create a Discovery implementation.
// This package will provide several implementations and register them for use.
// The flags FlagSet is provided for convenience, but also to hint to plugin
// developers that they should expect the args to be in a format compatible with
// pflag.
type Factory func(cluster *vtadminpb.Cluster, flags *pflag.FlagSet, args []string) (Discovery, error)

// nolint:gochecknoglobals
var registry = map[string]Factory{}

// Register registers a factory for the given implementation name. Attempting
// to register multiple factories for the same implementation name causes a
// panic.
func Register(name string, factory Factory) {
	_, ok := registry[name]
	if ok {
		panic("[discovery] factory already registered for " + name)
	}

	registry[name] = factory
}

// New returns a Discovery implementation using the registered factory for the
// implementation. Usage of the args slice is dependent on the implementation's
// factory.
func New(impl string, cluster *vtadminpb.Cluster, args []string) (Discovery, error) {
	factory, ok := registry[impl]
	if !ok {
		return nil, fmt.Errorf("%w %s", ErrImplementationNotRegistered, impl)
	}

	return factory(cluster, pflag.NewFlagSet("discovery:"+impl, pflag.ContinueOnError), args)
}

func init() { // nolint:gochecknoinits
	Register("consul", NewConsul)
	Register("staticfile", NewStaticFile)
}
<|endoftext|>