<commit_before>package problem

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/go-errors/errors"
	"github.com/stellar/horizon/context/requestid"
	"github.com/stellar/horizon/log"
	"golang.org/x/net/context"
)

var (
	errToProblemMap map[error]P = map[error]P{}
)

// RegisterError records an error -> P mapping, allowing the app to register
// specific errors that may occur in other packages to be rendered as a specific
// P instance.
//
// For example, you might want to render any db.ErrNoResults errors as a
// problem.NotFound, and you would do so by calling:
//
// problem.RegisterError(db.ErrNoResults, problem.NotFound) in you application
// initialization sequence
func RegisterError(err error, p P) {
	errToProblemMap[err] = p
}

// P is a struct that represents an error response to be rendered to a connected
// client.
type P struct {
	Type     string                 `json:"type"`
	Title    string                 `json:"title"`
	Status   int                    `json:"status"`
	Detail   string                 `json:"detail,omitempty"`
	Instance string                 `json:"instance,omitempty"`
	Extras   map[string]interface{} `json:"extras,omitempty"`
}

func (p *P) Error() string {
	return fmt.Sprintf("problem: %s", p.Type)
}

// Inflate expands a problem with contextal information.
// At present it adds the request's id as the problem's Instance, if available.
func Inflate(ctx context.Context, p *P) {
	//TODO: inflate type into full url
	//TODO: add requesting url to extra info

	p.Instance = requestid.FromContext(ctx)
}

// Render writes a http response to `w`, compliant with the "Problem
// Details for HTTP APIs" RFC:
//   https://tools.ietf.org/html/draft-ietf-appsawg-http-problem-00
//
// `p` is the problem, which may be either a concrete P struct, an implementor
// of the `HasProblem` interface, or an error.  Any other value for `p` will
// panic.
func Render(ctx context.Context, w http.ResponseWriter, p interface{}) {
	switch p := p.(type) {
	case P:
		render(ctx, w, p)
	case *P:
		render(ctx, w, *p)
	case error:
		renderErr(ctx, w, p)
	default:
		panic(fmt.Sprintf("Invalid problem: %v+", p))
	}
}

func render(ctx context.Context, w http.ResponseWriter, p P) {

	Inflate(ctx, &p)

	w.Header().Set("Content-Type", "application/problem+json")
	js, err := json.MarshalIndent(p, "", "  ")

	if err != nil {
		err := errors.Wrap(err, 1)
		log.WithStack(ctx, err).Error(err)
		http.Error(w, "error rendering problem", http.StatusInternalServerError)
		return
	}

	w.WriteHeader(p.Status)
	w.Write(js)
}

func renderErr(ctx context.Context, w http.ResponseWriter, err error) {
	origErr := err

	if err, ok := err.(*errors.Error); ok {
		origErr = err.Err
	}

	p, ok := errToProblemMap[origErr]

	if !ok {
		p = ServerError
	}

	log.WithStack(ctx, err).Error(err)
	render(ctx, w, p)
}

// Well-known and reused problems below:
var (
	// NotFound is a well-known problem type.  Use it as a shortcut
	// in your actions.
	NotFound = P{
		Type:   "not_found",
		Title:  "Resource Missing",
		Status: http.StatusNotFound,
		Detail: "The resource at the url requested was not found.  This is usually " +
			"occurs for one of two reasons:  The url requested is not valid, or no " +
			"data in our database could be found with the parameters provided.",
	}

	// ServerError is a well-known problem type.  Use it as a shortcut
	// in your actions.
	ServerError = P{
		Type:   "server_error",
		Title:  "Internal Server Error",
		Status: http.StatusInternalServerError,
		Detail: "An error occurred while processing this request.  This is usually due " +
			"to a bug within the server software.  Trying this request again may " +
			"succeed if the bug is transient, otherwise please report this issue " +
			"to the issue tracker at: https://github.com/stellar/horizon/issues." +
			" Please include this response in your issue.",
	}

	// RateLimitExceeded is a well-known problem type.  Use it as a shortcut
	// in your actions.
	RateLimitExceeded = P{
		Type:   "rate_limit_exceeded",
		Title:  "Rate limit exceeded",
		Status: 429,
		Detail: "The rate limit for the requesting IP address is over its alloted " +
			"limit.  The allowed limit and requests left per time period are " +
			"communicated to clients via the http response headers 'X-RateLimit-*' " +
			"headers.",
	}

	// NotImplemented is a well-known problem type.  Use it as a shortcut
	// in your actions.
	NotImplemented = P{
		Type:   "not_implemented",
		Title:  "Resource Not Yet Implemented",
		Status: http.StatusNotFound,
		Detail: "While the requested URL is expected to eventually point to a " +
			"valid resource, the work to implement the resource has not yet " +
			"been completed.",
	}

	// NotAcceptable is a well-known problem type.  Use it as a shortcut
	// in your actions.
	NotAcceptable = P{
		Type: "not_acceptable",
		Title: "An acceptable response content-type could not be provided for " +
			"this request",
		Status: http.StatusNotAcceptable,
	}
)
<commit_msg>Re-add HasProblem interface<commit_after>package problem

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/go-errors/errors"
	"github.com/stellar/horizon/context/requestid"
	"github.com/stellar/horizon/log"
	"golang.org/x/net/context"
)

var (
	errToProblemMap map[error]P = map[error]P{}
)

// RegisterError records an error -> P mapping, allowing the app to register
// specific errors that may occur in other packages to be rendered as a specific
// P instance.
//
// For example, you might want to render any db.ErrNoResults errors as a
// problem.NotFound, and you would do so by calling:
//
// problem.RegisterError(db.ErrNoResults, problem.NotFound) in you application
// initialization sequence
func RegisterError(err error, p P) {
	errToProblemMap[err] = p
}

// HasProblem types can be transformed into a problem.
// Implement it for custom errors.
type HasProblem interface {
	Problem() P
}

// P is a struct that represents an error response to be rendered to a connected
// client.
type P struct {
	Type     string                 `json:"type"`
	Title    string                 `json:"title"`
	Status   int                    `json:"status"`
	Detail   string                 `json:"detail,omitempty"`
	Instance string                 `json:"instance,omitempty"`
	Extras   map[string]interface{} `json:"extras,omitempty"`
}

func (p *P) Error() string {
	return fmt.Sprintf("problem: %s", p.Type)
}

// Inflate expands a problem with contextal information.
// At present it adds the request's id as the problem's Instance, if available.
func Inflate(ctx context.Context, p *P) {
	//TODO: inflate type into full url
	//TODO: add requesting url to extra info

	p.Instance = requestid.FromContext(ctx)
}

// Render writes a http response to `w`, compliant with the "Problem
// Details for HTTP APIs" RFC:
//   https://tools.ietf.org/html/draft-ietf-appsawg-http-problem-00
//
// `p` is the problem, which may be either a concrete P struct, an implementor
// of the `HasProblem` interface, or an error.  Any other value for `p` will
// panic.
func Render(ctx context.Context, w http.ResponseWriter, p interface{}) {
	switch p := p.(type) {
	case P:
		render(ctx, w, p)
	case *P:
		render(ctx, w, *p)
	case HasProblem:
		render(ctx, w, p.Problem())
	case error:
		renderErr(ctx, w, p)
	default:
		panic(fmt.Sprintf("Invalid problem: %v+", p))
	}
}

func render(ctx context.Context, w http.ResponseWriter, p P) {

	Inflate(ctx, &p)

	w.Header().Set("Content-Type", "application/problem+json")
	js, err := json.MarshalIndent(p, "", "  ")

	if err != nil {
		err := errors.Wrap(err, 1)
		log.WithStack(ctx, err).Error(err)
		http.Error(w, "error rendering problem", http.StatusInternalServerError)
		return
	}

	w.WriteHeader(p.Status)
	w.Write(js)
}

func renderErr(ctx context.Context, w http.ResponseWriter, err error) {
	origErr := err

	if err, ok := err.(*errors.Error); ok {
		origErr = err.Err
	}

	p, ok := errToProblemMap[origErr]

	if !ok {
		p = ServerError
	}

	log.WithStack(ctx, err).Error(err)
	render(ctx, w, p)
}

// Well-known and reused problems below:
var (
	// NotFound is a well-known problem type.  Use it as a shortcut
	// in your actions.
	NotFound = P{
		Type:   "not_found",
		Title:  "Resource Missing",
		Status: http.StatusNotFound,
		Detail: "The resource at the url requested was not found.  This is usually " +
			"occurs for one of two reasons:  The url requested is not valid, or no " +
			"data in our database could be found with the parameters provided.",
	}

	// ServerError is a well-known problem type.  Use it as a shortcut
	// in your actions.
	ServerError = P{
		Type:   "server_error",
		Title:  "Internal Server Error",
		Status: http.StatusInternalServerError,
		Detail: "An error occurred while processing this request.  This is usually due " +
			"to a bug within the server software.  Trying this request again may " +
			"succeed if the bug is transient, otherwise please report this issue " +
			"to the issue tracker at: https://github.com/stellar/horizon/issues." +
			" Please include this response in your issue.",
	}

	// RateLimitExceeded is a well-known problem type.  Use it as a shortcut
	// in your actions.
	RateLimitExceeded = P{
		Type:   "rate_limit_exceeded",
		Title:  "Rate limit exceeded",
		Status: 429,
		Detail: "The rate limit for the requesting IP address is over its alloted " +
			"limit.  The allowed limit and requests left per time period are " +
			"communicated to clients via the http response headers 'X-RateLimit-*' " +
			"headers.",
	}

	// NotImplemented is a well-known problem type.  Use it as a shortcut
	// in your actions.
	NotImplemented = P{
		Type:   "not_implemented",
		Title:  "Resource Not Yet Implemented",
		Status: http.StatusNotFound,
		Detail: "While the requested URL is expected to eventually point to a " +
			"valid resource, the work to implement the resource has not yet " +
			"been completed.",
	}

	// NotAcceptable is a well-known problem type.  Use it as a shortcut
	// in your actions.
	NotAcceptable = P{
		Type: "not_acceptable",
		Title: "An acceptable response content-type could not be provided for " +
			"this request",
		Status: http.StatusNotAcceptable,
	}
)
<|endoftext|>


<commit_before>// Copyright 2016 Mender Software AS
//
//    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 images

import (
	"errors"
	"net/http"
	"testing"
	"time"

	"github.com/ant0ine/go-json-rest/rest"
	"github.com/ant0ine/go-json-rest/rest/test"
	"github.com/mendersoftware/artifacts/mvc"
	"github.com/satori/go.uuid"
)

type fakeImageModeler struct {
	getImage          *SoftwareImage
	getImageError     error
	imagesList        []*SoftwareImage
	listImagesError   error
	uploadLink        *Link
	uploadLinkError   error
	downloadLink      *Link
	downloadLinkError error
	editImage         bool
	editError         error
}

func (fim *fakeImageModeler) ListImages(filters map[string]string) ([]*SoftwareImage, error) {
	return fim.imagesList, fim.listImagesError
}

func (fim *fakeImageModeler) UploadLink(imageID string, expire time.Duration) (*Link, error) {
	return fim.uploadLink, fim.uploadLinkError
}

func (fim *fakeImageModeler) DownloadLink(imageID string, expire time.Duration) (*Link, error) {
	return fim.downloadLink, fim.downloadLinkError
}

func (fim *fakeImageModeler) GetImage(id string) (*SoftwareImage, error) {
	return fim.getImage, fim.getImageError
}

func (fim *fakeImageModeler) DeleteImage(imageID string) error {
	return nil
}

func (fim *fakeImageModeler) CreateImage(constructorData *SoftwareImageConstructor) (string, error) {
	return "", nil
}

func (fim *fakeImageModeler) EditImage(id string, constructorData *SoftwareImageConstructor) (bool, error) {
	return fim.editImage, fim.editError
}

func setUpRestTest(route string, handler func(w rest.ResponseWriter, r *rest.Request)) *rest.Api {
	router, _ := rest.MakeRouter(rest.Get(route, handler))
	api := rest.NewApi()
	api.SetApp(router)

	return api
}

func TestControllerGetImage(t *testing.T) {
	imagesModel := new(fakeImageModeler)
	controller := NewSoftwareImagesController(imagesModel, mvc.RESTViewDefaults{})

	api := setUpRestTest("/api/0.0.1/images/:id", controller.GetImage)

	//no uuid provided
	recorded := test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/123", nil))
	recorded.CodeIs(http.StatusBadRequest)

	//have correct id, but no image
	id := uuid.NewV4().String()
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id, nil))
	recorded.CodeIs(http.StatusNotFound)

	//have correct id, but error getting image
	imagesModel.getImageError = errors.New("error")
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id, nil))
	recorded.CodeIs(http.StatusInternalServerError)

	// have image, get OK
	image := NewSoftwareImageConstructor()
	constructorImage := NewSoftwareImageFromConstructor(image)
	imagesModel.getImageError = nil
	imagesModel.getImage = constructorImage
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id, nil))
	recorded.CodeIs(http.StatusOK)
	recorded.ContentTypeIsJson()
}

func TestControllerListImages(t *testing.T) {
	imagesModel := new(fakeImageModeler)
	controller := NewSoftwareImagesController(imagesModel, mvc.RESTViewDefaults{})

	api := setUpRestTest("/api/0.0.1/images", controller.ListImages)

	//getting list error
	imagesModel.listImagesError = errors.New("error")
	recorded := test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images", nil))
	recorded.CodeIs(http.StatusInternalServerError)

	//getting list OK
	imagesModel.listImagesError = nil
	image := NewSoftwareImageConstructor()
	constructorImage := NewSoftwareImageFromConstructor(image)
	imagesModel.imagesList = append(imagesModel.imagesList, constructorImage)
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images", nil))
	recorded.CodeIs(http.StatusOK)
	recorded.ContentTypeIsJson()
}

func TestControllerUploadLink(t *testing.T) {
	imagesModel := new(fakeImageModeler)
	controller := NewSoftwareImagesController(imagesModel, mvc.RESTViewDefaults{})

	api := setUpRestTest("/api/0.0.1/images/:id/upload", controller.UploadLink)

	// wrong id
	recorded := test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/wrong_id/upload", nil))
	recorded.CodeIs(http.StatusBadRequest)

	// correct id; no upload link
	id := uuid.NewV4().String()
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id+"/upload", nil))
	recorded.CodeIs(http.StatusNotFound)

	// correct id; error generating link
	imagesModel.uploadLinkError = errors.New("error")
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id+"/upload", nil))
	recorded.CodeIs(http.StatusInternalServerError)

	// upload link OK
	imagesModel.uploadLinkError = nil
	link := NewLink("uri", time.Now())
	imagesModel.uploadLink = link
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id+"/upload", nil))
	recorded.CodeIs(http.StatusOK)
	recorded.ContentTypeIsJson()

}

func TestControllerDownloadLink(t *testing.T) {
	imagesModel := new(fakeImageModeler)
	controller := NewSoftwareImagesController(imagesModel, mvc.RESTViewDefaults{})

	api := setUpRestTest("/api/0.0.1/images/:id/download", controller.DownloadLink)

	// wrong id
	recorded := test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/wrong_id/download", nil))
	recorded.CodeIs(http.StatusBadRequest)

	// correct id; no upload link
	id := uuid.NewV4().String()
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id+"/download", nil))
	recorded.CodeIs(http.StatusNotFound)

	// correct id; error generating link
	imagesModel.downloadLinkError = errors.New("error")
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id+"/download", nil))
	recorded.CodeIs(http.StatusInternalServerError)

	// download link OK
	imagesModel.downloadLinkError = nil
	link := NewLink("uri", time.Now())
	imagesModel.downloadLink = link
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id+"/download", nil))
	recorded.CodeIs(http.StatusOK)
	recorded.ContentTypeIsJson()
}

func TestControllerDeleteImage(t *testing.T) {
	imagesModel := new(fakeImageModeler)
	controller := NewSoftwareImagesController(imagesModel, mvc.RESTViewDefaults{})

	api := setUpRestTest("/api/0.0.1/images/:id", controller.DeleteImage)

	// wrong id
	recorded := test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/wrong_id", nil))
	recorded.CodeIs(http.StatusBadRequest)

	// correct id; delete OK
	id := uuid.NewV4().String()
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id, nil))
	recorded.CodeIs(http.StatusNoContent)
}

func TestControllerEditImage(t *testing.T) {
	imagesModel := new(fakeImageModeler)
	controller := NewSoftwareImagesController(imagesModel, mvc.RESTViewDefaults{})

	api := setUpRestTest("/api/0.0.1/images/:id", controller.EditImage)

	// wrong id
	recorded := test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/wrong_id", nil))
	recorded.CodeIs(http.StatusBadRequest)

	// correct id; no payload
	id := uuid.NewV4().String()
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id, nil))
	recorded.CodeIs(http.StatusBadRequest)

	// correct id; invalid payload
	//image := NewSoftwareImageConstructor()
	//constructorImage := NewSoftwareImageFromConstructor(image)
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id, map[string]string{"image": "bad_image"}))
	recorded.CodeIs(http.StatusBadRequest)

	// correct id; correct payload; edit error
	imagesModel.editError = errors.New("error")
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id,
			map[string]string{"yocto_id": "1234-1234", "name": "myImage", "device_type": "myDevice"}))
	recorded.CodeIs(http.StatusInternalServerError)

	// correct id; correct payload; edit no image
	imagesModel.editError = nil
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id,
			map[string]string{"yocto_id": "1234-1234", "name": "myImage", "device_type": "myDevice"}))
	recorded.CodeIs(http.StatusNotFound)

	// correct id; correct payload; have image
	imagesModel.editImage = true
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id,
			map[string]string{"yocto_id": "1234-1234", "name": "myImage", "device_type": "myDevice"}))
	recorded.CodeIs(http.StatusNoContent)
}
<commit_msg>Add option for checking returned bodies in image controller tests.<commit_after>// Copyright 2016 Mender Software AS
//
//    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 images

import (
	"errors"
	"net/http"
	"testing"
	"time"

	"github.com/ant0ine/go-json-rest/rest"
	"github.com/ant0ine/go-json-rest/rest/test"
	"github.com/mendersoftware/artifacts/mvc"
	"github.com/satori/go.uuid"
)

type fakeImageModeler struct {
	getImage          *SoftwareImage
	getImageError     error
	imagesList        []*SoftwareImage
	listImagesError   error
	uploadLink        *Link
	uploadLinkError   error
	downloadLink      *Link
	downloadLinkError error
	editImage         bool
	editError         error
}

func (fim *fakeImageModeler) ListImages(filters map[string]string) ([]*SoftwareImage, error) {
	return fim.imagesList, fim.listImagesError
}

func (fim *fakeImageModeler) UploadLink(imageID string, expire time.Duration) (*Link, error) {
	return fim.uploadLink, fim.uploadLinkError
}

func (fim *fakeImageModeler) DownloadLink(imageID string, expire time.Duration) (*Link, error) {
	return fim.downloadLink, fim.downloadLinkError
}

func (fim *fakeImageModeler) GetImage(id string) (*SoftwareImage, error) {
	return fim.getImage, fim.getImageError
}

func (fim *fakeImageModeler) DeleteImage(imageID string) error {
	return nil
}

func (fim *fakeImageModeler) CreateImage(constructorData *SoftwareImageConstructor) (string, error) {
	return "", nil
}

func (fim *fakeImageModeler) EditImage(id string, constructorData *SoftwareImageConstructor) (bool, error) {
	return fim.editImage, fim.editError
}

func setUpRestTest(route string, handler func(w rest.ResponseWriter, r *rest.Request)) *rest.Api {
	router, _ := rest.MakeRouter(rest.Get(route, handler))
	api := rest.NewApi()
	api.SetApp(router)

	return api
}

func TestControllerGetImage(t *testing.T) {
	imagesModel := new(fakeImageModeler)
	controller := NewSoftwareImagesController(imagesModel, mvc.RESTViewDefaults{})

	api := setUpRestTest("/api/0.0.1/images/:id", controller.GetImage)

	//no uuid provided
	recorded := test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/123", nil))
	recorded.CodeIs(http.StatusBadRequest)

	//have correct id, but no image
	id := uuid.NewV4().String()
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id, nil))
	recorded.CodeIs(http.StatusNotFound)

	//have correct id, but error getting image
	imagesModel.getImageError = errors.New("error")
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id, nil))
	recorded.CodeIs(http.StatusInternalServerError)

	// have image, get OK
	image := NewSoftwareImageConstructor()
	constructorImage := NewSoftwareImageFromConstructor(image)
	imagesModel.getImageError = nil
	imagesModel.getImage = constructorImage
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id, nil))
	recorded.CodeIs(http.StatusOK)
	recorded.ContentTypeIsJson()

	var receivedImage SoftwareImage
	if err := recorded.DecodeJsonPayload(&receivedImage); err != nil {
		t.FailNow()
	}
}

func TestControllerListImages(t *testing.T) {
	imagesModel := new(fakeImageModeler)
	controller := NewSoftwareImagesController(imagesModel, mvc.RESTViewDefaults{})

	api := setUpRestTest("/api/0.0.1/images", controller.ListImages)

	//getting list error
	imagesModel.listImagesError = errors.New("error")
	recorded := test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images", nil))
	recorded.CodeIs(http.StatusInternalServerError)

	//getting list OK
	imagesModel.listImagesError = nil
	image := NewSoftwareImageConstructor()
	constructorImage := NewSoftwareImageFromConstructor(image)
	imagesModel.imagesList = append(imagesModel.imagesList, constructorImage)
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images", nil))
	recorded.CodeIs(http.StatusOK)
	recorded.ContentTypeIsJson()
}

func TestControllerUploadLink(t *testing.T) {
	imagesModel := new(fakeImageModeler)
	controller := NewSoftwareImagesController(imagesModel, mvc.RESTViewDefaults{})

	api := setUpRestTest("/api/0.0.1/images/:id/upload", controller.UploadLink)

	// wrong id
	recorded := test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/wrong_id/upload", nil))
	recorded.CodeIs(http.StatusBadRequest)

	// correct id; no upload link
	id := uuid.NewV4().String()
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id+"/upload", nil))
	recorded.CodeIs(http.StatusNotFound)

	// correct id; error generating link
	imagesModel.uploadLinkError = errors.New("error")
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id+"/upload", nil))
	recorded.CodeIs(http.StatusInternalServerError)

	// upload link OK
	imagesModel.uploadLinkError = nil
	link := NewLink("uri", time.Now())
	imagesModel.uploadLink = link
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id+"/upload", nil))
	recorded.CodeIs(http.StatusOK)
	recorded.ContentTypeIsJson()

	var receivedLink Link
	// check if returned body has Link structure
	if err := recorded.DecodeJsonPayload(&receivedLink); err != nil {
		t.FailNow()
	}
	// check if 'uri' is set correctly
	if receivedLink.Uri != "uri" {
		t.FailNow()
	}

}

func TestControllerDownloadLink(t *testing.T) {
	imagesModel := new(fakeImageModeler)
	controller := NewSoftwareImagesController(imagesModel, mvc.RESTViewDefaults{})

	api := setUpRestTest("/api/0.0.1/images/:id/download", controller.DownloadLink)

	// wrong id
	recorded := test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/wrong_id/download", nil))
	recorded.CodeIs(http.StatusBadRequest)

	// correct id; no upload link
	id := uuid.NewV4().String()
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id+"/download", nil))
	recorded.CodeIs(http.StatusNotFound)

	// correct id; error generating link
	imagesModel.downloadLinkError = errors.New("error")
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id+"/download", nil))
	recorded.CodeIs(http.StatusInternalServerError)

	// download link OK
	imagesModel.downloadLinkError = nil
	time := time.Now()
	link := NewLink("uri", time)
	imagesModel.downloadLink = link
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id+"/download", nil))
	recorded.CodeIs(http.StatusOK)
	recorded.ContentTypeIsJson()

	var receivedLink Link
	// check if returned body has Link structure
	if err := recorded.DecodeJsonPayload(&receivedLink); err != nil {
		t.FailNow()
	}
	// check if 'uri' is set correctly
	if receivedLink.Uri != "uri" {
		t.FailNow()
	}
}

func TestControllerDeleteImage(t *testing.T) {
	imagesModel := new(fakeImageModeler)
	controller := NewSoftwareImagesController(imagesModel, mvc.RESTViewDefaults{})

	api := setUpRestTest("/api/0.0.1/images/:id", controller.DeleteImage)

	// wrong id
	recorded := test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/wrong_id", nil))
	recorded.CodeIs(http.StatusBadRequest)

	// correct id; delete OK
	id := uuid.NewV4().String()
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id, nil))
	recorded.CodeIs(http.StatusNoContent)
	recorded.BodyIs("")
}

func TestControllerEditImage(t *testing.T) {
	imagesModel := new(fakeImageModeler)
	controller := NewSoftwareImagesController(imagesModel, mvc.RESTViewDefaults{})

	api := setUpRestTest("/api/0.0.1/images/:id", controller.EditImage)

	// wrong id
	recorded := test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/wrong_id", nil))
	recorded.CodeIs(http.StatusBadRequest)

	// correct id; no payload
	id := uuid.NewV4().String()
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id, nil))
	recorded.CodeIs(http.StatusBadRequest)

	// correct id; invalid payload
	//image := NewSoftwareImageConstructor()
	//constructorImage := NewSoftwareImageFromConstructor(image)
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id, map[string]string{"image": "bad_image"}))
	recorded.CodeIs(http.StatusBadRequest)

	// correct id; correct payload; edit error
	imagesModel.editError = errors.New("error")
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id,
			map[string]string{"yocto_id": "1234-1234", "name": "myImage", "device_type": "myDevice"}))
	recorded.CodeIs(http.StatusInternalServerError)

	// correct id; correct payload; edit no image
	imagesModel.editError = nil
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id,
			map[string]string{"yocto_id": "1234-1234", "name": "myImage", "device_type": "myDevice"}))
	recorded.CodeIs(http.StatusNotFound)

	// correct id; correct payload; have image
	imagesModel.editImage = true
	recorded = test.RunRequest(t, api.MakeHandler(),
		test.MakeSimpleRequest("GET", "http://localhost/api/0.0.1/images/"+id,
			map[string]string{"yocto_id": "1234-1234", "name": "myImage", "device_type": "myDevice"}))
	recorded.CodeIs(http.StatusNoContent)
	recorded.BodyIs("")
}
<|endoftext|>


<commit_before>package template

import (
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"log"
)

type template struct {
	Context  string
	Skeleton []map[string]string
}

func New(filepath string) *template {
	t := template{}
	data, err := ioutil.ReadFile(filepath)

	if err != nil {
		log.Fatal(err)
		return &template{}
	}

	if err = yaml.Unmarshal(data, &t); err != nil {
		log.Fatal(err)
		return &template{}
	}

	return &t
}
<commit_msg>Make template struct public<commit_after>package template

import (
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"log"
)

type Template struct {
	Context  string
	Skeleton []map[string]interface{}
}

func New(filepath string) *Template {
	t := Template{}
	data, err := ioutil.ReadFile(filepath)

	if err != nil {
		log.Fatal(err)
		return &Template{}
	}

	if err = yaml.Unmarshal(data, &t); err != nil {
		log.Fatal(err)
		return &Template{}
	}

	return &t
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes 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 metrics

import (
	"fmt"
	"testing"
	"time"

	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apiserver/pkg/admission"
)

var (
	kind     = schema.GroupVersionKind{Group: "kgroup", Version: "kversion", Kind: "kind"}
	resource = schema.GroupVersionResource{Group: "rgroup", Version: "rversion", Resource: "resource"}
	attr     = admission.NewAttributesRecord(nil, nil, kind, "ns", "name", resource, "subresource", admission.Create, nil)
)

func TestObserveAdmissionStep(t *testing.T) {
	Metrics.reset()
	handler := WithStepMetrics(&mutatingAndValidatingFakeHandler{admission.NewHandler(admission.Create), true, true})
	handler.(admission.MutationInterface).Admit(attr)
	handler.(admission.ValidationInterface).Validate(attr)
	wantLabels := map[string]string{
		"operation":   string(admission.Create),
		"group":       resource.Group,
		"version":     resource.Version,
		"resource":    resource.Resource,
		"subresource": "subresource",
		"type":        "admit",
		"rejected":    "false",
	}
	expectHistogramCountTotal(t, "apiserver_admission_step_admission_latencies_seconds", wantLabels, 1)
	expectFindMetric(t, "apiserver_admission_step_admission_latencies_seconds_summary", wantLabels)

	wantLabels["type"] = "validate"
	expectHistogramCountTotal(t, "apiserver_admission_step_admission_latencies_seconds", wantLabels, 1)
	expectFindMetric(t, "apiserver_admission_step_admission_latencies_seconds_summary", wantLabels)
}

func TestObserveAdmissionController(t *testing.T) {
	Metrics.reset()
	handler := WithControllerMetrics(&mutatingAndValidatingFakeHandler{admission.NewHandler(admission.Create), true, true}, "a")
	handler.(admission.MutationInterface).Admit(attr)
	handler.(admission.ValidationInterface).Validate(attr)
	wantLabels := map[string]string{
		"name":        "a",
		"operation":   string(admission.Create),
		"group":       resource.Group,
		"version":     resource.Version,
		"resource":    resource.Resource,
		"subresource": "subresource",
		"type":        "validate",
		"rejected":    "false",
	}
	expectHistogramCountTotal(t, "apiserver_admission_controller_admission_latencies_seconds", wantLabels, 1)

	wantLabels["type"] = "validate"
	expectHistogramCountTotal(t, "apiserver_admission_controller_admission_latencies_seconds", wantLabels, 1)
}

func TestObserveWebhook(t *testing.T) {
	Metrics.reset()
	Metrics.ObserveWebhook(2*time.Second, false, attr, stepAdmit, "x")
	wantLabels := map[string]string{
		"name":        "x",
		"operation":   string(admission.Create),
		"group":       resource.Group,
		"version":     resource.Version,
		"resource":    resource.Resource,
		"subresource": "subresource",
		"type":        "admit",
		"rejected":    "false",
	}
	expectHistogramCountTotal(t, "apiserver_admission_webhook_admission_latencies_seconds", wantLabels, 1)
}

func TestWithMetrics(t *testing.T) {
	Metrics.reset()

	type Test struct {
		name            string
		ns              string
		operation       admission.Operation
		handler         admission.Interface
		admit, validate bool
	}
	for _, test := range []Test{
		{
			"both-interfaces-admit-and-validate",
			"some-ns",
			admission.Create,
			&mutatingAndValidatingFakeHandler{admission.NewHandler(admission.Create, admission.Update), true, true},
			true, true,
		},
		{
			"both-interfaces-dont-admit",
			"some-ns",
			admission.Create,
			&mutatingAndValidatingFakeHandler{admission.NewHandler(admission.Create, admission.Update), false, true},
			false, true,
		},
		{
			"both-interfaces-admit-dont-validate",
			"some-ns",
			admission.Create,
			&mutatingAndValidatingFakeHandler{admission.NewHandler(admission.Create, admission.Update), true, false},
			true, false,
		},
		{
			"validate-interfaces-validate",
			"some-ns",
			admission.Create,
			&validatingFakeHandler{admission.NewHandler(admission.Create, admission.Update), true},
			true, true,
		},
		{
			"validate-interfaces-dont-validate",
			"some-ns",
			admission.Create,
			&validatingFakeHandler{admission.NewHandler(admission.Create, admission.Update), true},
			true, false,
		},
		{
			"mutating-interfaces-admit",
			"some-ns",
			admission.Create,
			&mutatingFakeHandler{admission.NewHandler(admission.Create, admission.Update), true},
			true, true,
		},
		{
			"mutating-interfaces-dont-admit",
			"some-ns",
			admission.Create,
			&mutatingFakeHandler{admission.NewHandler(admission.Create, admission.Update), false},
			true, false,
		},
	} {
		Metrics.reset()

		h := WithMetrics(test.handler, Metrics.ObserveAdmissionController, test.name)

		// test mutation
		err := h.(admission.MutationInterface).Admit(admission.NewAttributesRecord(nil, nil, schema.GroupVersionKind{}, test.ns, "", schema.GroupVersionResource{}, "", test.operation, nil))
		if test.admit && err != nil {
			t.Errorf("expected admit to succeed, but failed: %v", err)
			continue
		} else if !test.admit && err == nil {
			t.Errorf("expected admit to fail, but it succeeded")
			continue
		}

		filter := map[string]string{"rejected": "false"}
		if !test.admit {
			filter["rejected"] = "true"
		}
		if _, mutating := test.handler.(admission.MutationInterface); mutating {
			expectHistogramCountTotal(t, "apiserver_admission_controller_admission_latencies_seconds", filter, 1)
		} else {
			expectHistogramCountTotal(t, "apiserver_admission_controller_admission_latencies_seconds", filter, 0)
		}

		if err == nil {
			// skip validation step if mutation failed
			continue
		}

		// test validation
		err = h.(admission.ValidationInterface).Validate(admission.NewAttributesRecord(nil, nil, schema.GroupVersionKind{}, test.ns, "", schema.GroupVersionResource{}, "", test.operation, nil))
		if test.validate && err != nil {
			t.Errorf("expected admit to succeed, but failed: %v", err)
			continue
		} else if !test.validate && err == nil {
			t.Errorf("expected validation to fail, but it succeeded")
			continue
		}

		filter = map[string]string{"rejected": "false"}
		if !test.admit {
			filter["rejected"] = "true"
		}
		if _, validating := test.handler.(admission.ValidationInterface); validating {
			expectHistogramCountTotal(t, "apiserver_admission_controller_admission_latencies_seconds", filter, 1)
		} else {
			expectHistogramCountTotal(t, "apiserver_admission_controller_admission_latencies_seconds", filter, 0)
		}
	}
}

type mutatingAndValidatingFakeHandler struct {
	*admission.Handler
	admit    bool
	validate bool
}

func (h *mutatingAndValidatingFakeHandler) Admit(a admission.Attributes) (err error) {
	if h.admit {
		return nil
	}
	return fmt.Errorf("don't admit")
}

func (h *mutatingAndValidatingFakeHandler) Validate(a admission.Attributes) (err error) {
	if h.validate {
		return nil
	}
	return fmt.Errorf("don't validate")
}

type validatingFakeHandler struct {
	*admission.Handler
	validate bool
}

func (h *validatingFakeHandler) Validate(a admission.Attributes) (err error) {
	if h.validate {
		return nil
	}
	return fmt.Errorf("don't validate")
}

type mutatingFakeHandler struct {
	*admission.Handler
	admit bool
}

func (h *mutatingFakeHandler) Amit(a admission.Attributes) (err error) {
	if h.admit {
		return nil
	}
	return fmt.Errorf("don't admit")
}
<commit_msg>Fix admission metrics tests<commit_after>/*
Copyright 2017 The Kubernetes 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 metrics

import (
	"fmt"
	"testing"
	"time"

	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apiserver/pkg/admission"
)

var (
	kind     = schema.GroupVersionKind{Group: "kgroup", Version: "kversion", Kind: "kind"}
	resource = schema.GroupVersionResource{Group: "rgroup", Version: "rversion", Resource: "resource"}
	attr     = admission.NewAttributesRecord(nil, nil, kind, "ns", "name", resource, "subresource", admission.Create, nil)
)

func TestObserveAdmissionStep(t *testing.T) {
	Metrics.reset()
	handler := WithStepMetrics(&mutatingAndValidatingFakeHandler{admission.NewHandler(admission.Create), true, true})
	handler.(admission.MutationInterface).Admit(attr)
	handler.(admission.ValidationInterface).Validate(attr)
	wantLabels := map[string]string{
		"operation":   string(admission.Create),
		"group":       resource.Group,
		"version":     resource.Version,
		"resource":    resource.Resource,
		"subresource": "subresource",
		"type":        "admit",
		"rejected":    "false",
	}
	expectHistogramCountTotal(t, "apiserver_admission_step_admission_latencies_seconds", wantLabels, 1)
	expectFindMetric(t, "apiserver_admission_step_admission_latencies_seconds_summary", wantLabels)

	wantLabels["type"] = "validate"
	expectHistogramCountTotal(t, "apiserver_admission_step_admission_latencies_seconds", wantLabels, 1)
	expectFindMetric(t, "apiserver_admission_step_admission_latencies_seconds_summary", wantLabels)
}

func TestObserveAdmissionController(t *testing.T) {
	Metrics.reset()
	handler := WithControllerMetrics(&mutatingAndValidatingFakeHandler{admission.NewHandler(admission.Create), true, true}, "a")
	handler.(admission.MutationInterface).Admit(attr)
	handler.(admission.ValidationInterface).Validate(attr)
	wantLabels := map[string]string{
		"name":        "a",
		"operation":   string(admission.Create),
		"group":       resource.Group,
		"version":     resource.Version,
		"resource":    resource.Resource,
		"subresource": "subresource",
		"type":        "admit",
		"rejected":    "false",
	}
	expectHistogramCountTotal(t, "apiserver_admission_controller_admission_latencies_seconds", wantLabels, 1)

	wantLabels["type"] = "validate"
	expectHistogramCountTotal(t, "apiserver_admission_controller_admission_latencies_seconds", wantLabels, 1)
}

func TestObserveWebhook(t *testing.T) {
	Metrics.reset()
	Metrics.ObserveWebhook(2*time.Second, false, attr, stepAdmit, "x")
	wantLabels := map[string]string{
		"name":        "x",
		"operation":   string(admission.Create),
		"group":       resource.Group,
		"version":     resource.Version,
		"resource":    resource.Resource,
		"subresource": "subresource",
		"type":        "admit",
		"rejected":    "false",
	}
	expectHistogramCountTotal(t, "apiserver_admission_webhook_admission_latencies_seconds", wantLabels, 1)
}

func TestWithMetrics(t *testing.T) {
	Metrics.reset()

	type Test struct {
		name            string
		ns              string
		operation       admission.Operation
		handler         admission.Interface
		admit, validate bool
	}
	for _, test := range []Test{
		{
			"both-interfaces-admit-and-validate",
			"some-ns",
			admission.Create,
			&mutatingAndValidatingFakeHandler{admission.NewHandler(admission.Create, admission.Update), true, true},
			true, true,
		},
		{
			"both-interfaces-dont-admit",
			"some-ns",
			admission.Create,
			&mutatingAndValidatingFakeHandler{admission.NewHandler(admission.Create, admission.Update), false, true},
			false, true,
		},
		{
			"both-interfaces-admit-dont-validate",
			"some-ns",
			admission.Create,
			&mutatingAndValidatingFakeHandler{admission.NewHandler(admission.Create, admission.Update), true, false},
			true, false,
		},
		{
			"validate-interfaces-validate",
			"some-ns",
			admission.Create,
			&validatingFakeHandler{admission.NewHandler(admission.Create, admission.Update), true},
			true, true,
		},
		{
			"validate-interfaces-dont-validate",
			"some-ns",
			admission.Create,
			&validatingFakeHandler{admission.NewHandler(admission.Create, admission.Update), false},
			true, false,
		},
		{
			"mutating-interfaces-admit",
			"some-ns",
			admission.Create,
			&mutatingFakeHandler{admission.NewHandler(admission.Create, admission.Update), true},
			true, true,
		},
		{
			"mutating-interfaces-dont-admit",
			"some-ns",
			admission.Create,
			&mutatingFakeHandler{admission.NewHandler(admission.Create, admission.Update), false},
			false, true,
		},
	} {
		Metrics.reset()

		h := WithMetrics(test.handler, Metrics.ObserveAdmissionController, test.name)

		// test mutation
		err := h.(admission.MutationInterface).Admit(admission.NewAttributesRecord(nil, nil, schema.GroupVersionKind{}, test.ns, "", schema.GroupVersionResource{}, "", test.operation, nil))
		if test.admit && err != nil {
			t.Errorf("expected admit to succeed, but failed: %v", err)
			continue
		} else if !test.admit && err == nil {
			t.Errorf("expected admit to fail, but it succeeded")
			continue
		}

		filter := map[string]string{"type": "admit", "rejected": "false"}
		if !test.admit {
			filter["rejected"] = "true"
		}
		if _, mutating := test.handler.(admission.MutationInterface); mutating {
			expectHistogramCountTotal(t, "apiserver_admission_controller_admission_latencies_seconds", filter, 1)
		} else {
			expectHistogramCountTotal(t, "apiserver_admission_controller_admission_latencies_seconds", filter, 0)
		}

		if err != nil {
			// skip validation step if mutation failed
			continue
		}

		// test validation
		err = h.(admission.ValidationInterface).Validate(admission.NewAttributesRecord(nil, nil, schema.GroupVersionKind{}, test.ns, "", schema.GroupVersionResource{}, "", test.operation, nil))
		if test.validate && err != nil {
			t.Errorf("expected admit to succeed, but failed: %v", err)
			continue
		} else if !test.validate && err == nil {
			t.Errorf("expected validation to fail, but it succeeded")
			continue
		}

		filter = map[string]string{"type": "validate", "rejected": "false"}
		if !test.validate {
			filter["rejected"] = "true"
		}
		if _, validating := test.handler.(admission.ValidationInterface); validating {
			expectHistogramCountTotal(t, "apiserver_admission_controller_admission_latencies_seconds", filter, 1)
		} else {
			expectHistogramCountTotal(t, "apiserver_admission_controller_admission_latencies_seconds", filter, 0)
		}
	}
}

type mutatingAndValidatingFakeHandler struct {
	*admission.Handler
	admit    bool
	validate bool
}

func (h *mutatingAndValidatingFakeHandler) Admit(a admission.Attributes) (err error) {
	if h.admit {
		return nil
	}
	return fmt.Errorf("don't admit")
}

func (h *mutatingAndValidatingFakeHandler) Validate(a admission.Attributes) (err error) {
	if h.validate {
		return nil
	}
	return fmt.Errorf("don't validate")
}

type validatingFakeHandler struct {
	*admission.Handler
	validate bool
}

func (h *validatingFakeHandler) Validate(a admission.Attributes) (err error) {
	if h.validate {
		return nil
	}
	return fmt.Errorf("don't validate")
}

type mutatingFakeHandler struct {
	*admission.Handler
	admit bool
}

func (h *mutatingFakeHandler) Admit(a admission.Attributes) (err error) {
	if h.admit {
		return nil
	}
	return fmt.Errorf("don't admit")
}
<|endoftext|>


<commit_before>package httpconv

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"reflect"
	"strings"

	"github.com/fission/fission-workflows/pkg/types"
	"github.com/fission/fission-workflows/pkg/types/typedvalues"
	"github.com/golang/protobuf/proto"
	"github.com/sirupsen/logrus"
)

const (
	inputContentType    = "content-type"
	headerContentType   = "Content-Type"
	contentTypeJSON     = "application/json"
	contentTypeBytes    = "application/octet-stream"
	contentTypeText     = "text/plain"
	contentTypeTask     = "application/vnd.fission.workflows.workflow" // Default format: protobuf, +json for json
	contentTypeWorkflow = "application/vnd.fission.workflows.task"     // Default format: protobuf, +json for json
	contentTypeProtobuf = "application/protobuf"                       // Default format: protobuf, +json for json
	contentTypeDefault  = contentTypeText
	methodDefault       = http.MethodGet
)

// ParseRequest maps a HTTP request to a target map.
func ParseRequest(r *http.Request) (map[string]*types.TypedValue, error) {
	target := map[string]*types.TypedValue{}
	// Content-Type is a common problem, so log this for every request
	contentType := r.Header.Get(headerContentType)

	// Map body to "main" input
	bodyInput, err := ParseBody(r.Body, contentType)
	defer r.Body.Close()
	if err != nil {
		return nil, fmt.Errorf("failed to parse request: %v", err)
	}
	target[types.InputMain] = &bodyInput

	// Map query to "query.x"
	query := ParseQuery(r)
	target[types.InputQuery] = &query

	// Map headers to "headers.x"
	headers := ParseHeaders(r)
	target[types.InputHeaders] = &headers

	// Map http method to "method"
	method := ParseMethod(r)
	target[types.InputMethod] = &method

	return target, nil
}

func ParseBody(data io.Reader, contentType string) (types.TypedValue, error) {
	if len(contentType) == 0 {
		contentType = contentTypeDefault
	}

	tv := types.TypedValue{}
	tv.SetLabel(headerContentType, contentType)

	bs, err := ioutil.ReadAll(data)
	if err != nil {
		return tv, err
	}

	// Attempt to parse body according to provided ContentType
	switch normalizeContentType(contentType) {
	case contentTypeJSON:
		var i interface{}
		err := json.Unmarshal(bs, &i)
		if err != nil {
			logrus.Warnf("Failed to parse JSON data (len: %v, data: '%.50s' cause: %v), skipping further parsing.",
				len(bs), string(bs), err)
			tv = *typedvalues.ParseBytes(bs)
		} else {
			tv = *typedvalues.MustParse(i)
		}
	case contentTypeText:
		tv = *typedvalues.ParseString(string(bs))
	case contentTypeProtobuf:
		fallthrough
	case contentTypeTask:
		fallthrough
	case contentTypeWorkflow:
		// TODO support json-encoded workflow/task
		var m proto.Message
		err := proto.Unmarshal(bs, m)
		if err != nil {
			return tv, err
		}
		t, err := typedvalues.Parse(m)
		if err != nil {
			return tv, err
		}
		tv = *t
	default:
		// In other cases do not attempt to interpret the data
		fallthrough
	case contentTypeBytes:
		tv = *typedvalues.ParseBytes(bs)
	}

	return tv, nil
}

// ParseMethod maps the method param from a request to a TypedValue
func ParseMethod(r *http.Request) types.TypedValue {
	return *typedvalues.ParseString(r.Method)
}

// ParseHeaders maps the headers from a request to the "headers" key in the target map
func ParseHeaders(r *http.Request) types.TypedValue {
	// For now we do not support multi-valued headers
	headers := flattenMultimap(r.Header)

	tv := typedvalues.MustParse(headers)
	return *tv
}

// ParseQuery maps the query params from a request to the "query" key in the target map
func ParseQuery(r *http.Request) types.TypedValue {
	// For now we do not support multi-valued query params
	query := flattenMultimap(r.URL.Query())

	tv := typedvalues.MustParse(query)
	return *tv
}

// formatting logic
func FormatResponse(w http.ResponseWriter, output *types.TypedValue, outputErr *types.Error) {
	if w == nil {
		panic("cannot format response to nil")
	}

	if outputErr != nil {
		// TODO provide different http codes based on error
		http.Error(w, outputErr.Error(), http.StatusInternalServerError)
		return
	}

	if output == nil {
		w.WriteHeader(http.StatusNoContent)
		output = typedvalues.ParseNil()
		return
	}

	w.WriteHeader(http.StatusOK)
	contentType := DetermineContentType(output)
	w.Header().Set(headerContentType, contentType)
	bs, err := FormatBody(*output, contentType)
	if err != nil {
		FormatResponse(w, nil, &types.Error{
			Message: fmt.Sprintf("Failed to format response body: %v", err),
		})
	}
	w.Write(bs)
	return
}

func FormatRequest(source map[string]*types.TypedValue, target *http.Request) error {
	if target == nil {
		panic("cannot format request to nil")
	}

	// Map content-type to the request's content-type
	contentType := DetermineContentTypeInputs(source)

	// Map main input to body
	mainInput, ok := source[types.InputMain]
	if ok && mainInput != nil {
		bs, err := FormatBody(*mainInput, contentType)
		if err != nil {
			return err
		}
		target.Body = ioutil.NopCloser(bytes.NewReader(bs))
	}

	// Map method input to HTTP method
	method := FormatMethod(source)
	target.Method = method

	// Map query input to URL query
	query := FormatQuery(source)
	if query != nil {
		if target.URL == nil {
			panic("request has no URL")
		}
		target.URL.RawQuery = query.Encode()
	}

	// Map headers input to HTTP headers
	headers := FormatHeaders(source)
	if target.Header == nil {
		target.Header = headers
	} else {
		for k, v := range headers {
			if len(v) > 0 {
				target.Header.Set(k, v[0])
			}
		}
	}
	target.Header.Set(headerContentType, contentType)

	return nil
}

func FormatMethod(inputs map[string]*types.TypedValue) string {
	tv, ok := inputs[types.InputMethod]
	if ok && tv != nil {
		contentType, err := typedvalues.FormatString(tv)
		if err == nil {
			return contentType
		}
		logrus.Error("Invalid method in inputs: %+v", tv)
	}
	return methodDefault
}

// TODO support multivalued query params at some point
func FormatQuery(inputs map[string]*types.TypedValue) url.Values {
	queryInput := inputs[types.InputQuery]
	if queryInput == nil {
		return nil
	}
	target := url.Values{}

	i, err := typedvalues.Format(queryInput)
	if err != nil {
		logrus.Errorf("Failed to format headers: %v", err)
	}

	switch i.(type) {
	case map[string]interface{}:
		for k, v := range i.(map[string]interface{}) {
			target.Add(k, fmt.Sprintf("%v", v))
		}
	default:
		logrus.Warnf("Ignoring invalid type of query input (expected map[string]interface{}, was %v)",
			reflect.TypeOf(i))
	}
	return target
}

func FormatBody(value types.TypedValue, contentType string) ([]byte, error) {
	if len(contentType) == 0 {
		contentType = contentTypeDefault
	}

	i, err := typedvalues.Format(&value)
	if err != nil {
		return nil, err
	}

	// Attempt to parse body according to provided ContentType
	var bs []byte
	switch normalizeContentType(contentType) {
	case contentTypeJSON:
		bs, err = json.Marshal(i)
		if err != nil {
			return nil, err
		}
	case contentTypeText:
		switch t := i.(type) {
		case string:
			bs = []byte(t)
		case []byte:
			bs = t
		default:
			bs = []byte(fmt.Sprintf("%v", t))
		}
	case contentTypeProtobuf:
		fallthrough
	case contentTypeTask:
		fallthrough
	case contentTypeWorkflow:
		// TODO support json
		m, ok := i.(proto.Message)
		if !ok {
			return nil, fmt.Errorf("illegal content type '%T', should be protobuf", i)
		}
		bs, err = proto.Marshal(m)
		if err != nil {
			return nil, err
		}
	default:
		fallthrough
	case contentTypeBytes:
		var ok bool
		bs, ok = i.([]byte)
		if !ok {
			return nil, fmt.Errorf("illegal content type '%T', should be []byte", i)
		}
	}
	return bs, nil
}

func DetermineContentType(value *types.TypedValue) string {
	if value == nil {
		return contentTypeBytes
	}

	ct, ok := value.GetLabel(headerContentType)
	if ok && len(ct) > 0 {
		return ct
	}

	// Otherwise, check for primitive types of the main input
	switch typedvalues.ValueType(value.Type) {
	// TODO task and workflow
	case typedvalues.TypeMap:
		fallthrough
	case typedvalues.TypeList:
		return contentTypeJSON
	case typedvalues.TypeNumber:
		fallthrough
	case typedvalues.TypeExpression:
		fallthrough
	case typedvalues.TypeString:
		return contentTypeText
	default:
		return contentTypeBytes
	}
}

func DetermineContentTypeInputs(inputs map[string]*types.TypedValue) string {
	// Check for forced contentType in inputs
	ctTv, ok := inputs[inputContentType]
	if ok && ctTv != nil {
		contentType, err := typedvalues.FormatString(ctTv)
		if err == nil {
			return contentType
		}
		logrus.Error("Invalid content type in inputs: %+v", ctTv)
	}

	// Otherwise, check for label on body input
	return DetermineContentType(inputs[types.InputMain])
}

// TODO support multi-headers at some point
func FormatHeaders(inputs map[string]*types.TypedValue) http.Header {
	headers := http.Header{}
	rawHeaders, ok := inputs[types.InputHeaders]
	if !ok || rawHeaders == nil {
		return headers
	}

	// TODO handle partial map
	i, err := typedvalues.Format(rawHeaders)
	if err != nil {
		logrus.Errorf("Failed to format headers: %v", err)
	}

	switch i.(type) {
	case map[string]interface{}:
		for k, v := range i.(map[string]interface{}) {
			headers.Add(k, fmt.Sprintf("%v", v))
		}
	default:
		logrus.Warnf("Ignoring invalid type of headers input (expected map[string]interface{}, was %v)",
			reflect.TypeOf(i))
	}
	return headers
}

// Util

func flattenMultimap(mm map[string][]string) map[string]interface{} {
	target := map[string]interface{}{}
	for k, v := range mm {
		target[k] = v[0]
	}
	return target
}

func normalizeContentType(contentType string) string {
	contentType = strings.ToLower(strings.TrimSpace(contentType))
	matchContentType := contentType
	// Heuristics, because everything to do with HTTP is ambiguous...
	if strings.Contains(contentType, "json") { // TODO exclude JSON representation of protobuf objects
		matchContentType = contentTypeJSON
	}
	if strings.HasPrefix(contentType, "text") {
		matchContentType = contentTypeText
	}
	if strings.Contains(contentType, "protobuf") {
		matchContentType = contentTypeProtobuf
	}
	return matchContentType
}
<commit_msg>Set content-length when setting new body (#156)<commit_after>package httpconv

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"reflect"
	"strings"

	"github.com/fission/fission-workflows/pkg/types"
	"github.com/fission/fission-workflows/pkg/types/typedvalues"
	"github.com/golang/protobuf/proto"
	"github.com/sirupsen/logrus"
)

const (
	inputContentType    = "content-type"
	headerContentType   = "Content-Type"
	contentTypeJSON     = "application/json"
	contentTypeBytes    = "application/octet-stream"
	contentTypeText     = "text/plain"
	contentTypeTask     = "application/vnd.fission.workflows.workflow" // Default format: protobuf, +json for json
	contentTypeWorkflow = "application/vnd.fission.workflows.task"     // Default format: protobuf, +json for json
	contentTypeProtobuf = "application/protobuf"                       // Default format: protobuf, +json for json
	contentTypeDefault  = contentTypeText
	methodDefault       = http.MethodPost
)

// ParseRequest maps a HTTP request to a target map.
func ParseRequest(r *http.Request) (map[string]*types.TypedValue, error) {
	target := map[string]*types.TypedValue{}
	// Content-Type is a common problem, so log this for every request
	contentType := r.Header.Get(headerContentType)

	// Map body to "main" input
	bodyInput, err := ParseBody(r.Body, contentType)
	defer r.Body.Close()
	if err != nil {
		return nil, fmt.Errorf("failed to parse request: %v", err)
	}
	target[types.InputMain] = &bodyInput

	// Map query to "query.x"
	query := ParseQuery(r)
	target[types.InputQuery] = &query

	// Map headers to "headers.x"
	headers := ParseHeaders(r)
	target[types.InputHeaders] = &headers

	// Map http method to "method"
	method := ParseMethod(r)
	target[types.InputMethod] = &method

	return target, nil
}

func ParseBody(data io.Reader, contentType string) (types.TypedValue, error) {
	if len(contentType) == 0 {
		contentType = contentTypeDefault
	}

	tv := types.TypedValue{}
	tv.SetLabel(headerContentType, contentType)

	bs, err := ioutil.ReadAll(data)
	if err != nil {
		return tv, err
	}

	// Attempt to parse body according to provided ContentType
	switch normalizeContentType(contentType) {
	case contentTypeJSON:
		var i interface{}
		err := json.Unmarshal(bs, &i)
		if err != nil {
			logrus.Warnf("Failed to parse JSON data (len: %v, data: '%.50s' cause: %v), skipping further parsing.",
				len(bs), string(bs), err)
			tv = *typedvalues.ParseBytes(bs)
		} else {
			tv = *typedvalues.MustParse(i)
		}
	case contentTypeText:
		tv = *typedvalues.ParseString(string(bs))
	case contentTypeProtobuf:
		fallthrough
	case contentTypeTask:
		fallthrough
	case contentTypeWorkflow:
		// TODO support json-encoded workflow/task
		var m proto.Message
		err := proto.Unmarshal(bs, m)
		if err != nil {
			return tv, err
		}
		t, err := typedvalues.Parse(m)
		if err != nil {
			return tv, err
		}
		tv = *t
	default:
		// In other cases do not attempt to interpret the data
		fallthrough
	case contentTypeBytes:
		tv = *typedvalues.ParseBytes(bs)
	}

	return tv, nil
}

// ParseMethod maps the method param from a request to a TypedValue
func ParseMethod(r *http.Request) types.TypedValue {
	return *typedvalues.ParseString(r.Method)
}

// ParseHeaders maps the headers from a request to the "headers" key in the target map
func ParseHeaders(r *http.Request) types.TypedValue {
	// For now we do not support multi-valued headers
	headers := flattenMultimap(r.Header)

	tv := typedvalues.MustParse(headers)
	return *tv
}

// ParseQuery maps the query params from a request to the "query" key in the target map
func ParseQuery(r *http.Request) types.TypedValue {
	// For now we do not support multi-valued query params
	query := flattenMultimap(r.URL.Query())

	tv := typedvalues.MustParse(query)
	return *tv
}

// formatting logic
func FormatResponse(w http.ResponseWriter, output *types.TypedValue, outputErr *types.Error) {
	if w == nil {
		panic("cannot format response to nil")
	}

	if outputErr != nil {
		// TODO provide different http codes based on error
		http.Error(w, outputErr.Error(), http.StatusInternalServerError)
		return
	}

	if output == nil {
		w.WriteHeader(http.StatusNoContent)
		output = typedvalues.ParseNil()
		return
	}

	w.WriteHeader(http.StatusOK)
	contentType := DetermineContentType(output)
	w.Header().Set(headerContentType, contentType)
	bs, err := FormatBody(*output, contentType)
	if err != nil {
		FormatResponse(w, nil, &types.Error{
			Message: fmt.Sprintf("Failed to format response body: %v", err),
		})
	}
	w.Write(bs)
	return
}

func FormatRequest(source map[string]*types.TypedValue, target *http.Request) error {
	if target == nil {
		panic("cannot format request to nil")
	}

	// Map content-type to the request's content-type
	contentType := DetermineContentTypeInputs(source)

	// Map main input to body
	mainInput, ok := source[types.InputMain]
	if ok && mainInput != nil {
		bs, err := FormatBody(*mainInput, contentType)
		if err != nil {
			return err
		}
		target.Body = ioutil.NopCloser(bytes.NewReader(bs))
		target.ContentLength = int64(len(bs))
	}

	// Map method input to HTTP method
	method := FormatMethod(source)
	target.Method = method

	// Map query input to URL query
	query := FormatQuery(source)
	if query != nil {
		if target.URL == nil {
			panic("request has no URL")
		}
		target.URL.RawQuery = query.Encode()
	}

	// Map headers input to HTTP headers
	headers := FormatHeaders(source)
	if target.Header == nil {
		target.Header = headers
	} else {
		for k, v := range headers {
			if len(v) > 0 {
				target.Header.Set(k, v[0])
			}
		}
	}
	target.Header.Set(headerContentType, contentType)

	return nil
}

func FormatMethod(inputs map[string]*types.TypedValue) string {
	tv, ok := inputs[types.InputMethod]
	if ok && tv != nil {
		contentType, err := typedvalues.FormatString(tv)
		if err == nil {
			return contentType
		}
		logrus.Error("Invalid method in inputs: %+v", tv)
	}
	return methodDefault
}

// TODO support multivalued query params at some point
func FormatQuery(inputs map[string]*types.TypedValue) url.Values {
	queryInput := inputs[types.InputQuery]
	if queryInput == nil {
		return nil
	}
	target := url.Values{}

	i, err := typedvalues.Format(queryInput)
	if err != nil {
		logrus.Errorf("Failed to format headers: %v", err)
	}

	switch i.(type) {
	case map[string]interface{}:
		for k, v := range i.(map[string]interface{}) {
			target.Add(k, fmt.Sprintf("%v", v))
		}
	default:
		logrus.Warnf("Ignoring invalid type of query input (expected map[string]interface{}, was %v)",
			reflect.TypeOf(i))
	}
	return target
}

func FormatBody(value types.TypedValue, contentType string) ([]byte, error) {
	if len(contentType) == 0 {
		contentType = contentTypeDefault
	}

	i, err := typedvalues.Format(&value)
	if err != nil {
		return nil, err
	}

	// Attempt to parse body according to provided ContentType
	var bs []byte
	switch normalizeContentType(contentType) {
	case contentTypeJSON:
		bs, err = json.Marshal(i)
		if err != nil {
			return nil, err
		}
	case contentTypeText:
		switch t := i.(type) {
		case string:
			bs = []byte(t)
		case []byte:
			bs = t
		default:
			bs = []byte(fmt.Sprintf("%v", t))
		}
	case contentTypeProtobuf:
		fallthrough
	case contentTypeTask:
		fallthrough
	case contentTypeWorkflow:
		// TODO support json
		m, ok := i.(proto.Message)
		if !ok {
			return nil, fmt.Errorf("illegal content type '%T', should be protobuf", i)
		}
		bs, err = proto.Marshal(m)
		if err != nil {
			return nil, err
		}
	default:
		fallthrough
	case contentTypeBytes:
		var ok bool
		bs, ok = i.([]byte)
		if !ok {
			return nil, fmt.Errorf("illegal content type '%T', should be []byte", i)
		}
	}
	return bs, nil
}

func DetermineContentType(value *types.TypedValue) string {
	if value == nil {
		return contentTypeBytes
	}

	ct, ok := value.GetLabel(headerContentType)
	if ok && len(ct) > 0 {
		return ct
	}

	// Otherwise, check for primitive types of the main input
	switch typedvalues.ValueType(value.Type) {
	// TODO task and workflow
	case typedvalues.TypeMap:
		fallthrough
	case typedvalues.TypeList:
		return contentTypeJSON
	case typedvalues.TypeNumber:
		fallthrough
	case typedvalues.TypeExpression:
		fallthrough
	case typedvalues.TypeString:
		return contentTypeText
	default:
		return contentTypeBytes
	}
}

func DetermineContentTypeInputs(inputs map[string]*types.TypedValue) string {
	// Check for forced contentType in inputs
	ctTv, ok := inputs[inputContentType]
	if ok && ctTv != nil {
		contentType, err := typedvalues.FormatString(ctTv)
		if err == nil {
			return contentType
		}
		logrus.Error("Invalid content type in inputs: %+v", ctTv)
	}

	// Otherwise, check for label on body input
	return DetermineContentType(inputs[types.InputMain])
}

// TODO support multi-headers at some point
func FormatHeaders(inputs map[string]*types.TypedValue) http.Header {
	headers := http.Header{}
	rawHeaders, ok := inputs[types.InputHeaders]
	if !ok || rawHeaders == nil {
		return headers
	}

	// TODO handle partial map
	i, err := typedvalues.Format(rawHeaders)
	if err != nil {
		logrus.Errorf("Failed to format headers: %v", err)
	}

	switch i.(type) {
	case map[string]interface{}:
		for k, v := range i.(map[string]interface{}) {
			headers.Add(k, fmt.Sprintf("%v", v))
		}
	default:
		logrus.Warnf("Ignoring invalid type of headers input (expected map[string]interface{}, was %v)",
			reflect.TypeOf(i))
	}
	return headers
}

// Util

func flattenMultimap(mm map[string][]string) map[string]interface{} {
	target := map[string]interface{}{}
	for k, v := range mm {
		target[k] = v[0]
	}
	return target
}

func normalizeContentType(contentType string) string {
	contentType = strings.ToLower(strings.TrimSpace(contentType))
	matchContentType := contentType
	// Heuristics, because everything to do with HTTP is ambiguous...
	if strings.Contains(contentType, "json") { // TODO exclude JSON representation of protobuf objects
		matchContentType = contentTypeJSON
	}
	if strings.HasPrefix(contentType, "text") {
		matchContentType = contentTypeText
	}
	if strings.Contains(contentType, "protobuf") {
		matchContentType = contentTypeProtobuf
	}
	return matchContentType
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

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 attestation defines methods to attest a message using Pgp Private and
// Public Key pair.
package attestation

import (
	"bytes"
	"crypto"
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/golang/glog"
	"github.com/grafeas/kritis/pkg/kritis/secrets"
	"github.com/pkg/errors"

	"golang.org/x/crypto/openpgp"
	"golang.org/x/crypto/openpgp/armor"
	"golang.org/x/crypto/openpgp/packet"
	"golang.org/x/crypto/openpgp/s2k"
)

const (
	RSABits = 4096
)

var pgpConfig = packet.Config{
	// Use Sha256
	DefaultHash:            crypto.SHA256,
	DefaultCipher:          packet.CipherAES256,
	DefaultCompressionAlgo: packet.CompressionZLIB,
	CompressionConfig: &packet.CompressionConfig{
		Level: packet.DefaultCompression,
	},
	RSABits: RSABits,
}

// VerifyMessageAttestation verifies if the image is attested using the PEM
// encoded public key.
func VerifyMessageAttestation(pubKey string, sig string, message string) error {
	text, err := GetPlainMessage(pubKey, sig)
	if err != nil {
		return err
	}
	// Finally, make sure the signature is over the right message.
	if string(text) != message {
		return fmt.Errorf("signature could not be verified. got: %s, want: %s", text, message)
	}
	return nil
}

// GetPlainMessage verifies if the image is attested using the PEM
// encoded public key and returns the plain text in bytes
func GetPlainMessage(pubKey string, sig string) ([]byte, error) {
	keyring, err := openpgp.ReadArmoredKeyRing(strings.NewReader(pubKey))
	if err != nil {
		return nil, errors.Wrap(err, "read armored key ring")
	}
	buf := bytes.NewBufferString(sig)
	armorBlock, err := armor.Decode(buf)
	if err != nil {
		return nil, errors.Wrap(err, "could not decode armor signature")
	}
	md, err := openpgp.ReadMessage(armorBlock.Body, keyring, nil, &pgpConfig)
	if err != nil {
		return nil, errors.Wrap(err, "could not read armor signature")
	}

	// MessageDetails.UnverifiedBody signature is not verified until we read it.
	// This will call PublicKey.VerifySignature for the keys in the keyring.
	plaintext, err := ioutil.ReadAll(md.UnverifiedBody)
	if err != nil {
		return nil, errors.Wrap(err, "could not verify armor signature")
	}
	// Make sure after reading the UnverifiedBody above, there is no signature error.
	if md.SignatureError != nil {
		return nil, fmt.Errorf("bad signature found: %s", md.SignatureError)
	}
	if md.Signature == nil {
		return nil, fmt.Errorf("no signature found for given key")
	}

	return plaintext, nil
}

// CreateMessageAttestation attests the message using the given PGP key.
// pgpKey: PGP key
// message: Message to attest
func CreateMessageAttestation(pgpKey *secrets.PgpKey, message string) (string, error) {
	// First Create a signer Entity from public and private keys.
	signer, err := createEntityFromKeys(pgpKey.PublicKey(), pgpKey.PrivateKey())
	if err != nil {
		return "", errors.Wrap(err, "creating entity keys")
	}

	b := new(bytes.Buffer)
	// Armor Encode it.
	armorWriter, errEncode := armor.Encode(b, openpgp.SignatureType, make(map[string]string))
	if errEncode != nil {
		return "", errors.Wrap(err, "encoding data")
	}
	// Finally Sign the Text.
	w, err := openpgp.Sign(armorWriter, signer, nil, &pgpConfig)
	if err != nil {
		return "", errors.Wrap(err, "opengpg signing")
	}

	_, err = w.Write([]byte(message))
	if err != nil {
		return "", errors.Wrap(err, "writing signed data")
	}
	w.Close()
	armorWriter.Close()
	return string(b.Bytes()), nil
}

func createEntityFromKeys(pubKey *packet.PublicKey, privKey *packet.PrivateKey) (*openpgp.Entity, error) {
	currentTime := pgpConfig.Now()
	uid := packet.NewUserId("", "", "")
	if uid == nil {
		return nil, errors.New("user id field contained invalid characters")
	}

	e := &openpgp.Entity{
		PrimaryKey: pubKey,
		PrivateKey: privKey,
		Identities: make(map[string]*openpgp.Identity),
	}
	isPrimaryID := true
	e.Identities[uid.Id] = &openpgp.Identity{
		Name:   uid.Id,
		UserId: uid,
		SelfSignature: &packet.Signature{
			CreationTime: currentTime,
			SigType:      packet.SigTypePositiveCert,
			PubKeyAlgo:   packet.PubKeyAlgoRSA,
			Hash:         pgpConfig.Hash(),
			IsPrimaryId:  &isPrimaryID,
			FlagsValid:   true,
			FlagSign:     true,
			FlagCertify:  true,
			IssuerKeyId:  &e.PrimaryKey.KeyId,
		},
	}
	err := e.Identities[uid.Id].SelfSignature.SignUserId(uid.Id, e.PrimaryKey, e.PrivateKey, &pgpConfig)
	if err != nil {
		return nil, err
	}

	// Set Config Hash from Config
	hashID, ok := s2k.HashToHashId(pgpConfig.DefaultHash)
	if !ok {
		return nil, fmt.Errorf("tried to convert unknown hash %d", pgpConfig.DefaultHash)
	}
	e.Identities[uid.Id].SelfSignature.PreferredHash = []uint8{hashID}

	// Set Config Cipher from Config
	e.Identities[uid.Id].SelfSignature.PreferredSymmetric = []uint8{uint8(pgpConfig.DefaultCipher)}

	return e, nil
}
<commit_msg>Remove unused import<commit_after>/*
Copyright 2018 Google LLC

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 attestation defines methods to attest a message using Pgp Private and
// Public Key pair.
package attestation

import (
	"bytes"
	"crypto"
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/grafeas/kritis/pkg/kritis/secrets"
	"github.com/pkg/errors"

	"golang.org/x/crypto/openpgp"
	"golang.org/x/crypto/openpgp/armor"
	"golang.org/x/crypto/openpgp/packet"
	"golang.org/x/crypto/openpgp/s2k"
)

const (
	RSABits = 4096
)

var pgpConfig = packet.Config{
	// Use Sha256
	DefaultHash:            crypto.SHA256,
	DefaultCipher:          packet.CipherAES256,
	DefaultCompressionAlgo: packet.CompressionZLIB,
	CompressionConfig: &packet.CompressionConfig{
		Level: packet.DefaultCompression,
	},
	RSABits: RSABits,
}

// VerifyMessageAttestation verifies if the image is attested using the PEM
// encoded public key.
func VerifyMessageAttestation(pubKey string, sig string, message string) error {
	text, err := GetPlainMessage(pubKey, sig)
	if err != nil {
		return err
	}
	// Finally, make sure the signature is over the right message.
	if string(text) != message {
		return fmt.Errorf("signature could not be verified. got: %s, want: %s", text, message)
	}
	return nil
}

// GetPlainMessage verifies if the image is attested using the PEM
// encoded public key and returns the plain text in bytes
func GetPlainMessage(pubKey string, sig string) ([]byte, error) {
	keyring, err := openpgp.ReadArmoredKeyRing(strings.NewReader(pubKey))
	if err != nil {
		return nil, errors.Wrap(err, "read armored key ring")
	}
	buf := bytes.NewBufferString(sig)
	armorBlock, err := armor.Decode(buf)
	if err != nil {
		return nil, errors.Wrap(err, "could not decode armor signature")
	}
	md, err := openpgp.ReadMessage(armorBlock.Body, keyring, nil, &pgpConfig)
	if err != nil {
		return nil, errors.Wrap(err, "could not read armor signature")
	}

	// MessageDetails.UnverifiedBody signature is not verified until we read it.
	// This will call PublicKey.VerifySignature for the keys in the keyring.
	plaintext, err := ioutil.ReadAll(md.UnverifiedBody)
	if err != nil {
		return nil, errors.Wrap(err, "could not verify armor signature")
	}
	// Make sure after reading the UnverifiedBody above, there is no signature error.
	if md.SignatureError != nil {
		return nil, fmt.Errorf("bad signature found: %s", md.SignatureError)
	}
	if md.Signature == nil {
		return nil, fmt.Errorf("no signature found for given key")
	}

	return plaintext, nil
}

// CreateMessageAttestation attests the message using the given PGP key.
// pgpKey: PGP key
// message: Message to attest
func CreateMessageAttestation(pgpKey *secrets.PgpKey, message string) (string, error) {
	// First Create a signer Entity from public and private keys.
	signer, err := createEntityFromKeys(pgpKey.PublicKey(), pgpKey.PrivateKey())
	if err != nil {
		return "", errors.Wrap(err, "creating entity keys")
	}

	b := new(bytes.Buffer)
	// Armor Encode it.
	armorWriter, errEncode := armor.Encode(b, openpgp.SignatureType, make(map[string]string))
	if errEncode != nil {
		return "", errors.Wrap(err, "encoding data")
	}
	// Finally Sign the Text.
	w, err := openpgp.Sign(armorWriter, signer, nil, &pgpConfig)
	if err != nil {
		return "", errors.Wrap(err, "opengpg signing")
	}

	_, err = w.Write([]byte(message))
	if err != nil {
		return "", errors.Wrap(err, "writing signed data")
	}
	w.Close()
	armorWriter.Close()
	return string(b.Bytes()), nil
}

func createEntityFromKeys(pubKey *packet.PublicKey, privKey *packet.PrivateKey) (*openpgp.Entity, error) {
	currentTime := pgpConfig.Now()
	uid := packet.NewUserId("", "", "")
	if uid == nil {
		return nil, errors.New("user id field contained invalid characters")
	}

	e := &openpgp.Entity{
		PrimaryKey: pubKey,
		PrivateKey: privKey,
		Identities: make(map[string]*openpgp.Identity),
	}
	isPrimaryID := true
	e.Identities[uid.Id] = &openpgp.Identity{
		Name:   uid.Id,
		UserId: uid,
		SelfSignature: &packet.Signature{
			CreationTime: currentTime,
			SigType:      packet.SigTypePositiveCert,
			PubKeyAlgo:   packet.PubKeyAlgoRSA,
			Hash:         pgpConfig.Hash(),
			IsPrimaryId:  &isPrimaryID,
			FlagsValid:   true,
			FlagSign:     true,
			FlagCertify:  true,
			IssuerKeyId:  &e.PrimaryKey.KeyId,
		},
	}
	err := e.Identities[uid.Id].SelfSignature.SignUserId(uid.Id, e.PrimaryKey, e.PrivateKey, &pgpConfig)
	if err != nil {
		return nil, err
	}

	// Set Config Hash from Config
	hashID, ok := s2k.HashToHashId(pgpConfig.DefaultHash)
	if !ok {
		return nil, fmt.Errorf("tried to convert unknown hash %d", pgpConfig.DefaultHash)
	}
	e.Identities[uid.Id].SelfSignature.PreferredHash = []uint8{hashID}

	// Set Config Cipher from Config
	e.Identities[uid.Id].SelfSignature.PreferredSymmetric = []uint8{uint8(pgpConfig.DefaultCipher)}

	return e, nil
}
<|endoftext|>


<commit_before>// Copyright 2018 The OpenPitrix Authors. All rights reserved.
// Use of this source code is governed by a Apache license
// that can be found in the LICENSE file.

package drone

import (
	"bytes"
	"context"
	"fmt"
	"os/exec"
	"path/filepath"
	"reflect"
	"runtime"
	"strings"
	"time"

	"openpitrix.io/openpitrix/pkg/libconfd"
	"openpitrix.io/openpitrix/pkg/logger"
	"openpitrix.io/openpitrix/pkg/pb/drone"
	"openpitrix.io/openpitrix/pkg/pb/types"
	"openpitrix.io/openpitrix/pkg/util/funcutil"
)

var (
	_ pbdrone.DroneServiceServer = (*Server)(nil)
)

func (p *Server) GetDroneConfig(context.Context, *pbtypes.Empty) (*pbtypes.DroneConfig, error) {
	logger.Info(funcutil.CallerName(1))

	return p.cfg.Get(), nil
}
func (p *Server) SetDroneConfig(ctx context.Context, cfg *pbtypes.DroneConfig) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))

	if reflect.DeepEqual(cfg, p.cfg.Get()) {
		return &pbtypes.Empty{}, nil
	}

	if err := p.cfg.Set(cfg); err != nil {
		logger.Warn("%+v", err)
		return &pbtypes.Empty{}, err
	}

	if err := p.cfg.Save(); err != nil {
		logger.Warn("%+v", err)
		return &pbtypes.Empty{}, err
	}

	return &pbtypes.Empty{}, nil
}

func (p *Server) GetConfdConfig(ctx context.Context, arg *pbtypes.Empty) (*pbtypes.ConfdConfig, error) {
	logger.Info(funcutil.CallerName(1))

	return p.confd.GetConfig(), nil
}

func (p *Server) SetConfdConfig(ctx context.Context, arg *pbtypes.ConfdConfig) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))

	if err := p.confd.SetConfig(arg); err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}
	if err := p.confd.Save(); err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	return &pbtypes.Empty{}, nil
}

func (p *Server) GetFrontgateConfig(context.Context, *pbtypes.Empty) (*pbtypes.FrontgateConfig, error) {
	logger.Info(funcutil.CallerName(1))

	cfg, err := p.fg.GetConfig()
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}
	return cfg, nil
}

func (p *Server) SetFrontgateConfig(ctx context.Context, cfg *pbtypes.FrontgateConfig) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))

	err := p.fg.SetConfig(cfg)
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}
	return &pbtypes.Empty{}, nil
}

func (p *Server) IsConfdRunning(ctx context.Context, arg *pbtypes.Empty) (*pbtypes.Bool, error) {
	logger.Info(funcutil.CallerName(1))

	return &pbtypes.Bool{Value: p.confd.IsRunning()}, nil
}

func (p *Server) StartConfd(ctx context.Context, arg *pbtypes.Empty) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))

	cfg := p.cfg.Get()
	logger.Info("StartConfd:", cfg)

	err := p.confd.Start(func(opt *libconfd.Config) {
		opt.HookAbsKeyAdjuster = func(absKey string) (realKey string) {
			if absKey == "/self" {
				return "/" + cfg.ConfdSelfHost
			}
			if strings.HasPrefix(absKey, "/self/") {
				return "/" + cfg.ConfdSelfHost + absKey[len("/self/")-1:]
			}
			return absKey
		}
		opt.HookOnCheckCmdDone = func(trName, cmd string, err error) {
			if err != nil {
				logger.Warn("%+v", err)
				return
			}
			if trName == "cmd.info" {
				go func() {
					status, err := LoadLastCmdStatus(cfg.CmdInfoLogPath)
					if err != nil {
						logger.Warn("%+v", err)
					} else {
						p.fg.ReportSubTaskStatus(&pbtypes.SubTaskStatus{
							TaskId: status.SubtaskId,
							Status: status.Status,
						})
					}
				}()
			}
		}
		opt.HookOnReloadCmdDone = func(trName, cmd string, err error) {
			if err != nil {
				logger.Warn("%+v", err)
				return
			}
			if trName == "cmd.info" {
				go func() {
					status, err := LoadLastCmdStatus(cfg.CmdInfoLogPath)
					if err != nil {
						logger.Warn("%+v", err)
					} else {
						p.fg.ReportSubTaskStatus(&pbtypes.SubTaskStatus{
							TaskId: status.SubtaskId,
							Status: status.Status,
						})
					}
				}()
			}
		}
		opt.HookOnUpdateDone = func(trName string, err error) {
			if err != nil {
				logger.Warn("%+v", err)
				return
			}
		}
	})

	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	return &pbtypes.Empty{}, nil
}

func (p *Server) StopConfd(ctx context.Context, arg *pbtypes.Empty) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))

	if err := p.confd.Stop(); err != nil {
		logger.Error("StopConfd:", err)
		return nil, err
	}

	return &pbtypes.Empty{}, nil
}

func (p *Server) GetTemplateFiles(ctx context.Context, arg *pbtypes.Empty) (*pbtypes.StringList, error) {
	logger.Info(funcutil.CallerName(1))

	if !p.confd.IsRunning() {
		err := fmt.Errorf("drone: confd is not started")
		logger.Warn("%+v", err)
		return nil, err
	}

	cfg := p.confd.GetConfig()
	confdir := cfg.GetProcessorConfig().GetConfdir()
	if confdir == "" {
		err := fmt.Errorf("drone: invaid confdir: %q", confdir)
		logger.Warn("%+v", err)
		return nil, err
	}

	_, paths, err := libconfd.ListTemplateResource(filepath.Join(confdir, "conf.d"))
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	reply := &pbtypes.StringList{}
	for _, s := range paths {
		reply.ValueList = append(reply.ValueList, filepath.Base(s))
	}

	return reply, nil
}

func (p *Server) GetValues(ctx context.Context, arg *pbtypes.StringList) (*pbtypes.StringMap, error) {
	logger.Info(funcutil.CallerName(1))

	if !p.confd.IsRunning() {
		err := fmt.Errorf("drone: confd is not started")
		logger.Warn("%+v", err)
		return nil, err
	}

	client := p.confd.GetBackendClient()
	if client == nil {
		return nil, fmt.Errorf("drone: confd is not started")
	}

	keys := arg.GetValueList()
	m, err := client.GetValues(keys)
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	reply := &pbtypes.StringMap{
		ValueMap: make(map[string]string),
	}

	for _, k := range keys {
		reply.ValueMap[k] = m[k]
	}

	return reply, nil
}

func (p *Server) PingPilot(ctx context.Context, arg *pbtypes.FrontgateEndpoint) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))

	c, err := p.fg.getClient()
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	_, err = c.PingPilot(&pbtypes.Empty{})
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	return &pbtypes.Empty{}, nil
}

func (p *Server) PingFrontgate(ctx context.Context, arg *pbtypes.FrontgateEndpoint) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))

	c, err := p.fg.getClient()
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	_, err = c.PingFrontgate(&pbtypes.Empty{})
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	return &pbtypes.Empty{}, nil
}

func (p *Server) PingDrone(ctx context.Context, arg *pbtypes.Empty) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))
	return &pbtypes.Empty{}, nil
}

func (p *Server) RunCommand(ctx context.Context, arg *pbtypes.RunCommandOnDroneRequest) (*pbtypes.String, error) {
	logger.Info(funcutil.CallerName(1))

	var c *exec.Cmd
	if runtime.GOOS == "windows" {
		c = exec.Command("cmd", "/C", arg.GetCommand())
	} else {
		c = exec.Command("/bin/sh", "-c", arg.GetCommand())
	}

	var b bytes.Buffer
	c.Stdout = &b
	c.Stderr = &b

	if err := c.Start(); err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	var timeout = time.Second * 3
	if x := arg.GetTimeoutSeconds(); x > 0 {
		timeout = time.Duration(x) * time.Second
	}

	timer := time.AfterFunc(timeout, func() {
		c.Process.Kill()
	})

	err := c.Wait()
	timer.Stop()

	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	return &pbtypes.String{Value: string(b.Bytes())}, nil
}
<commit_msg>metadata/drone: keep all confd keys in /self name space<commit_after>// Copyright 2018 The OpenPitrix Authors. All rights reserved.
// Use of this source code is governed by a Apache license
// that can be found in the LICENSE file.

package drone

import (
	"bytes"
	"context"
	"fmt"
	"os/exec"
	"path/filepath"
	"reflect"
	"runtime"
	"strings"
	"time"

	"openpitrix.io/openpitrix/pkg/libconfd"
	"openpitrix.io/openpitrix/pkg/logger"
	"openpitrix.io/openpitrix/pkg/pb/drone"
	"openpitrix.io/openpitrix/pkg/pb/types"
	"openpitrix.io/openpitrix/pkg/util/funcutil"
)

var (
	_ pbdrone.DroneServiceServer = (*Server)(nil)
)

func (p *Server) GetDroneConfig(context.Context, *pbtypes.Empty) (*pbtypes.DroneConfig, error) {
	logger.Info(funcutil.CallerName(1))

	return p.cfg.Get(), nil
}
func (p *Server) SetDroneConfig(ctx context.Context, cfg *pbtypes.DroneConfig) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))

	if reflect.DeepEqual(cfg, p.cfg.Get()) {
		return &pbtypes.Empty{}, nil
	}

	if err := p.cfg.Set(cfg); err != nil {
		logger.Warn("%+v", err)
		return &pbtypes.Empty{}, err
	}

	if err := p.cfg.Save(); err != nil {
		logger.Warn("%+v", err)
		return &pbtypes.Empty{}, err
	}

	return &pbtypes.Empty{}, nil
}

func (p *Server) GetConfdConfig(ctx context.Context, arg *pbtypes.Empty) (*pbtypes.ConfdConfig, error) {
	logger.Info(funcutil.CallerName(1))

	return p.confd.GetConfig(), nil
}

func (p *Server) SetConfdConfig(ctx context.Context, arg *pbtypes.ConfdConfig) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))

	if err := p.confd.SetConfig(arg); err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}
	if err := p.confd.Save(); err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	return &pbtypes.Empty{}, nil
}

func (p *Server) GetFrontgateConfig(context.Context, *pbtypes.Empty) (*pbtypes.FrontgateConfig, error) {
	logger.Info(funcutil.CallerName(1))

	cfg, err := p.fg.GetConfig()
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}
	return cfg, nil
}

func (p *Server) SetFrontgateConfig(ctx context.Context, cfg *pbtypes.FrontgateConfig) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))

	err := p.fg.SetConfig(cfg)
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}
	return &pbtypes.Empty{}, nil
}

func (p *Server) IsConfdRunning(ctx context.Context, arg *pbtypes.Empty) (*pbtypes.Bool, error) {
	logger.Info(funcutil.CallerName(1))

	return &pbtypes.Bool{Value: p.confd.IsRunning()}, nil
}

func (p *Server) StartConfd(ctx context.Context, arg *pbtypes.Empty) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))

	cfg := p.cfg.Get()
	logger.Info("StartConfd:", cfg)

	err := p.confd.Start(func(opt *libconfd.Config) {
		opt.HookAbsKeyAdjuster = func(absKey string) (realKey string) {
			if absKey == "/self" {
				return "/" + cfg.ConfdSelfHost
			}
			if strings.HasPrefix(absKey, "/self/") {
				return "/" + cfg.ConfdSelfHost + absKey[len("/self/")-1:]
			} else {
				return "/" + cfg.ConfdSelfHost + absKey
			}
			return absKey
		}
		opt.HookOnCheckCmdDone = func(trName, cmd string, err error) {
			if err != nil {
				logger.Warn("%+v", err)
				return
			}
			if trName == "cmd.info" {
				go func() {
					status, err := LoadLastCmdStatus(cfg.CmdInfoLogPath)
					if err != nil {
						logger.Warn("%+v", err)
					} else {
						p.fg.ReportSubTaskStatus(&pbtypes.SubTaskStatus{
							TaskId: status.SubtaskId,
							Status: status.Status,
						})
					}
				}()
			}
		}
		opt.HookOnReloadCmdDone = func(trName, cmd string, err error) {
			if err != nil {
				logger.Warn("%+v", err)
				return
			}
			if trName == "cmd.info" {
				go func() {
					status, err := LoadLastCmdStatus(cfg.CmdInfoLogPath)
					if err != nil {
						logger.Warn("%+v", err)
					} else {
						p.fg.ReportSubTaskStatus(&pbtypes.SubTaskStatus{
							TaskId: status.SubtaskId,
							Status: status.Status,
						})
					}
				}()
			}
		}
		opt.HookOnUpdateDone = func(trName string, err error) {
			if err != nil {
				logger.Warn("%+v", err)
				return
			}
		}
	})

	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	return &pbtypes.Empty{}, nil
}

func (p *Server) StopConfd(ctx context.Context, arg *pbtypes.Empty) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))

	if err := p.confd.Stop(); err != nil {
		logger.Error("StopConfd:", err)
		return nil, err
	}

	return &pbtypes.Empty{}, nil
}

func (p *Server) GetTemplateFiles(ctx context.Context, arg *pbtypes.Empty) (*pbtypes.StringList, error) {
	logger.Info(funcutil.CallerName(1))

	if !p.confd.IsRunning() {
		err := fmt.Errorf("drone: confd is not started")
		logger.Warn("%+v", err)
		return nil, err
	}

	cfg := p.confd.GetConfig()
	confdir := cfg.GetProcessorConfig().GetConfdir()
	if confdir == "" {
		err := fmt.Errorf("drone: invaid confdir: %q", confdir)
		logger.Warn("%+v", err)
		return nil, err
	}

	_, paths, err := libconfd.ListTemplateResource(filepath.Join(confdir, "conf.d"))
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	reply := &pbtypes.StringList{}
	for _, s := range paths {
		reply.ValueList = append(reply.ValueList, filepath.Base(s))
	}

	return reply, nil
}

func (p *Server) GetValues(ctx context.Context, arg *pbtypes.StringList) (*pbtypes.StringMap, error) {
	logger.Info(funcutil.CallerName(1))

	if !p.confd.IsRunning() {
		err := fmt.Errorf("drone: confd is not started")
		logger.Warn("%+v", err)
		return nil, err
	}

	client := p.confd.GetBackendClient()
	if client == nil {
		return nil, fmt.Errorf("drone: confd is not started")
	}

	keys := arg.GetValueList()
	m, err := client.GetValues(keys)
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	reply := &pbtypes.StringMap{
		ValueMap: make(map[string]string),
	}

	for _, k := range keys {
		reply.ValueMap[k] = m[k]
	}

	return reply, nil
}

func (p *Server) PingPilot(ctx context.Context, arg *pbtypes.FrontgateEndpoint) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))

	c, err := p.fg.getClient()
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	_, err = c.PingPilot(&pbtypes.Empty{})
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	return &pbtypes.Empty{}, nil
}

func (p *Server) PingFrontgate(ctx context.Context, arg *pbtypes.FrontgateEndpoint) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))

	c, err := p.fg.getClient()
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	_, err = c.PingFrontgate(&pbtypes.Empty{})
	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	return &pbtypes.Empty{}, nil
}

func (p *Server) PingDrone(ctx context.Context, arg *pbtypes.Empty) (*pbtypes.Empty, error) {
	logger.Info(funcutil.CallerName(1))
	return &pbtypes.Empty{}, nil
}

func (p *Server) RunCommand(ctx context.Context, arg *pbtypes.RunCommandOnDroneRequest) (*pbtypes.String, error) {
	logger.Info(funcutil.CallerName(1))

	var c *exec.Cmd
	if runtime.GOOS == "windows" {
		c = exec.Command("cmd", "/C", arg.GetCommand())
	} else {
		c = exec.Command("/bin/sh", "-c", arg.GetCommand())
	}

	var b bytes.Buffer
	c.Stdout = &b
	c.Stderr = &b

	if err := c.Start(); err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	var timeout = time.Second * 3
	if x := arg.GetTimeoutSeconds(); x > 0 {
		timeout = time.Duration(x) * time.Second
	}

	timer := time.AfterFunc(timeout, func() {
		c.Process.Kill()
	})

	err := c.Wait()
	timer.Stop()

	if err != nil {
		logger.Warn("%+v", err)
		return nil, err
	}

	return &pbtypes.String{Value: string(b.Bytes())}, nil
}
<|endoftext|>