<commit_before>package grestclient

import (
	"errors"
	"io/ioutil"
	"net/http"
	"net/url"
)

//cloneUrl will do a deep clone of the url
//you pass in
func cloneUrl(u *url.URL) *url.URL {
	var userInfo *url.Userinfo
	if u.User != nil {
		if p, ok := u.User.Password(); ok {
			userInfo = url.UserPassword(u.User.Username(), p)
		} else {
			userInfo = url.User(u.User.Username())
		}
	}

	return &url.URL{
		Scheme:   u.Scheme,
		Opaque:   u.Opaque,
		User:     userInfo,
		Host:     u.Host,
		Path:     u.Path,
		RawQuery: u.RawQuery,
		Fragment: u.Fragment,
	}
}

type client struct {
	base        *url.URL
	reqMutators []RequestMutator
	resMutators []ResponseMutator
	headers     http.Header
	query       url.Values
	client      *http.Client
	marshaler   MarshalerFunc
	unmarshaler UnmarshalerFunc
}

func (c *client) Headers() http.Header {
	if c.headers == nil {
		c.headers = make(http.Header)
	}
	return c.headers
}

func (c *client) SetHeaders(h http.Header) {
	c.headers = h
}

func (c *client) Query() url.Values {
	if c.query == nil {
		c.query = make(url.Values)
	}
	return c.query
}

func (c *client) SetQuery(q url.Values) {
	c.query = q
}

func (c *client) AddRequestMutators(rm ...RequestMutator) Client {
	c.reqMutators = append(c.reqMutators, rm...)
	return c
}

func (c *client) AddResponseMutators(rm ...ResponseMutator) Client {
	c.resMutators = append(c.resMutators, rm...)
	return c
}

func (c *client) SetRequestMutators(rm ...RequestMutator) Client {
	c.reqMutators = rm
	return c
}

func (c *client) SetResponseMutators(rm ...ResponseMutator) Client {
	c.resMutators = rm
	return c
}

func (c *client) RequestMutators() []RequestMutator {
	return c.reqMutators
}

func (c *client) ResponseMutators() []ResponseMutator {
	return c.resMutators
}

func (c *client) Get(path string, headers http.Header, query url.Values, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	r, err := c.prepareRequest("GET", path, headers, query, nil)
	if err != nil {
		return nil, err
	}
	return c.do(r, successResult, errorResult)
}

func (c *client) Post(path string, headers http.Header, query url.Values, postBody interface{}, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	r, err := c.prepareRequest("POST", path, headers, query, postBody)
	if err != nil {
		return nil, err
	}
	return c.do(r, successResult, errorResult)
}

func (c *client) Put(path string, headers http.Header, query url.Values, putBody interface{}, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	r, err := c.prepareRequest("PUT", path, headers, query, putBody)
	if err != nil {
		return nil, err
	}
	return c.do(r, successResult, errorResult)
}

func (c *client) Patch(path string, headers http.Header, query url.Values, patchBody interface{}, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	r, err := c.prepareRequest("PATCH", path, headers, query, patchBody)
	if err != nil {
		return nil, err
	}
	return c.do(r, successResult, errorResult)
}

func (c *client) Head(path string, headers http.Header, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	r, err := c.prepareRequest("HEAD", path, headers, nil, nil)
	if err != nil {
		return nil, err
	}
	return c.do(r, successResult, errorResult)
}

func (c *client) Options(path string, headers http.Header, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	r, err := c.prepareRequest("OPTIONS", path, headers, nil, nil)
	if err != nil {
		return nil, err
	}
	return c.do(r, successResult, errorResult)
}

func (c *client) Delete(path string, headers http.Header, query url.Values, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	r, err := c.prepareRequest("DELETE", path, headers, query, nil)
	if err != nil {
		return nil, err
	}
	return c.do(r, successResult, errorResult)
}

func (c *client) do(r *http.Request, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	var err error
	if c.RequestMutators() != nil {
		for _, m := range c.RequestMutators() {
			err = m(r)
			if err != nil {
				return nil, err
			}
		}
	}
	var response *http.Response
	client := c.GetHttpClient()

	response, err = client.Do(r)

	if err != nil {
		return nil, err
	}
	defer response.Body.Close()

	if c.ResponseMutators() != nil {
		var err error
		for _, m := range c.ResponseMutators() {
			err = m(response)
			if err != nil {
				return response, err
			}
		}
	}

	if c.unmarshaler == nil {
		c.unmarshaler = StringUnmarshalerFunc
	}

	//make sure there is a body, or that there might be a body (when it is -1)
	if response.ContentLength > 0 || response.ContentLength == -1 {

		//unmarshal it depending on StatusCode
		if response.StatusCode < 300 && successResult != nil {
			//success
			err = c.unmarshaler(response.Body, successResult)
		} else if response.StatusCode < 600 && errorResult != nil {
			//error
			err = c.unmarshaler(response.Body, errorResult)
		}
	}

	if err != nil {
		//we have the response so return it even though unmarshaling might've
		//produced an error
		return response, err
	}

	return response, nil
}

func (c *client) prepareRequest(
	method string,
	path string,
	headers http.Header,
	query url.Values,
	body interface{}) (*http.Request, error) {

	var err error
	reqUrl := cloneUrl(c.base)
	reqUrl.Path += path

	//set headers
	headers = setupHeaders(c.headers, headers)
	//create query
	query = setupQuery(c.query, query)

	if c.marshaler == nil {
		c.marshaler = StringMarshalerFunc
	}

	r, err := http.NewRequest(method, reqUrl.String(), nil)
	if err != nil {
		return nil, err
	}

	r.Header = headers
	r.URL.RawQuery = query.Encode()

	var readLener ReadLener
	if body != nil {

		readLener, err = c.marshaler(body)

		if err != nil {
			return nil, err
		}
		r.ContentLength = int64(readLener.Len())
		r.Body = ioutil.NopCloser(readLener)
	}

	return r, nil
}

func setupHeaders(headers ...http.Header) http.Header {
	finalheaders := make(http.Header)
	for _, current := range headers {
		for i, v := range current {
			finalheaders[i] = v
		}
	}

	return finalheaders
}

func setupQuery(queries ...url.Values) url.Values {
	finalquery := make(url.Values)
	for _, current := range queries {
		for i, v := range current {
			finalquery[i] = v
		}
	}

	return finalquery
}

func (c *client) SetBaseUrl(u *url.URL) error {
	if u == nil {
		return errors.New("Please specify a non nil url.")
	}
	u.RawQuery = ""
	c.base = u
	return nil
}

func (c *client) BaseUrl() *url.URL {
	return c.base
}

func (c *client) Clone() Client {
	cc := &client{}
	cc.base = cloneUrl(c.base)

	cc.reqMutators = make([]RequestMutator, len(c.reqMutators))
	copy(cc.reqMutators, c.reqMutators)

	cc.resMutators = make([]ResponseMutator, len(c.resMutators))
	copy(cc.resMutators, c.resMutators)

	cc.headers = headerCopy(c.headers)
	cc.query = queryCopy(c.query)
	cc.client = c.client
	cc.marshaler = c.marshaler
	cc.unmarshaler = c.unmarshaler

	return cc
}

func headerCopy(h http.Header) http.Header {
	if h == nil {
		return nil
	}
	c := make(http.Header, len(h))
	for i, v := range h {
		c[i] = v
	}
	return c
}

func queryCopy(q url.Values) url.Values {
	if q == nil {
		return nil
	}
	c := make(url.Values, len(q))
	for i, v := range q {
		c[i] = v
	}
	return c
}

func (c *client) GetHttpClient() *http.Client {
	if c.client == nil {
		c.client = http.DefaultClient
	}
	return c.client
}

func (c *client) SetHttpClient(h *http.Client) {
	c.client = h
}

func (c *client) SetMarshaler(f MarshalerFunc) {
	c.marshaler = f
}

func (c *client) SetUnmarshaler(f UnmarshalerFunc) {
	c.unmarshaler = f
}

//New creates a new grestclient with the base url set
//to the passed in paramater.
func New(base *url.URL) (Client, error) {

	if base == nil {
		return nil, errors.New("Please specify a non nil url.")
	}
	c := &client{}
	c.base = base

	return c, nil
}
<commit_msg>300 requests no longer unmarshaled.<commit_after>package grestclient

import (
	"errors"
	"io/ioutil"
	"net/http"
	"net/url"
)

//cloneUrl will do a deep clone of the url
//you pass in
func cloneUrl(u *url.URL) *url.URL {
	var userInfo *url.Userinfo
	if u.User != nil {
		if p, ok := u.User.Password(); ok {
			userInfo = url.UserPassword(u.User.Username(), p)
		} else {
			userInfo = url.User(u.User.Username())
		}
	}

	return &url.URL{
		Scheme:   u.Scheme,
		Opaque:   u.Opaque,
		User:     userInfo,
		Host:     u.Host,
		Path:     u.Path,
		RawQuery: u.RawQuery,
		Fragment: u.Fragment,
	}
}

type client struct {
	base        *url.URL
	reqMutators []RequestMutator
	resMutators []ResponseMutator
	headers     http.Header
	query       url.Values
	client      *http.Client
	marshaler   MarshalerFunc
	unmarshaler UnmarshalerFunc
}

func (c *client) Headers() http.Header {
	if c.headers == nil {
		c.headers = make(http.Header)
	}
	return c.headers
}

func (c *client) SetHeaders(h http.Header) {
	c.headers = h
}

func (c *client) Query() url.Values {
	if c.query == nil {
		c.query = make(url.Values)
	}
	return c.query
}

func (c *client) SetQuery(q url.Values) {
	c.query = q
}

func (c *client) AddRequestMutators(rm ...RequestMutator) Client {
	c.reqMutators = append(c.reqMutators, rm...)
	return c
}

func (c *client) AddResponseMutators(rm ...ResponseMutator) Client {
	c.resMutators = append(c.resMutators, rm...)
	return c
}

func (c *client) SetRequestMutators(rm ...RequestMutator) Client {
	c.reqMutators = rm
	return c
}

func (c *client) SetResponseMutators(rm ...ResponseMutator) Client {
	c.resMutators = rm
	return c
}

func (c *client) RequestMutators() []RequestMutator {
	return c.reqMutators
}

func (c *client) ResponseMutators() []ResponseMutator {
	return c.resMutators
}

func (c *client) Get(path string, headers http.Header, query url.Values, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	r, err := c.prepareRequest("GET", path, headers, query, nil)
	if err != nil {
		return nil, err
	}
	return c.do(r, successResult, errorResult)
}

func (c *client) Post(path string, headers http.Header, query url.Values, postBody interface{}, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	r, err := c.prepareRequest("POST", path, headers, query, postBody)
	if err != nil {
		return nil, err
	}
	return c.do(r, successResult, errorResult)
}

func (c *client) Put(path string, headers http.Header, query url.Values, putBody interface{}, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	r, err := c.prepareRequest("PUT", path, headers, query, putBody)
	if err != nil {
		return nil, err
	}
	return c.do(r, successResult, errorResult)
}

func (c *client) Patch(path string, headers http.Header, query url.Values, patchBody interface{}, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	r, err := c.prepareRequest("PATCH", path, headers, query, patchBody)
	if err != nil {
		return nil, err
	}
	return c.do(r, successResult, errorResult)
}

func (c *client) Head(path string, headers http.Header, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	r, err := c.prepareRequest("HEAD", path, headers, nil, nil)
	if err != nil {
		return nil, err
	}
	return c.do(r, successResult, errorResult)
}

func (c *client) Options(path string, headers http.Header, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	r, err := c.prepareRequest("OPTIONS", path, headers, nil, nil)
	if err != nil {
		return nil, err
	}
	return c.do(r, successResult, errorResult)
}

func (c *client) Delete(path string, headers http.Header, query url.Values, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	r, err := c.prepareRequest("DELETE", path, headers, query, nil)
	if err != nil {
		return nil, err
	}
	return c.do(r, successResult, errorResult)
}

func (c *client) do(r *http.Request, successResult interface{}, errorResult interface{}) (*http.Response, error) {
	var err error
	if c.RequestMutators() != nil {
		for _, m := range c.RequestMutators() {
			err = m(r)
			if err != nil {
				return nil, err
			}
		}
	}
	var response *http.Response
	client := c.GetHttpClient()

	response, err = client.Do(r)

	if err != nil {
		return nil, err
	}
	defer response.Body.Close()

	if c.ResponseMutators() != nil {
		var err error
		for _, m := range c.ResponseMutators() {
			err = m(response)
			if err != nil {
				return response, err
			}
		}
	}

	if c.unmarshaler == nil {
		c.unmarshaler = StringUnmarshalerFunc
	}

	//make sure there is a body, or that there might be a body (when it is -1)
	if response.ContentLength > 0 || response.ContentLength == -1 {

		//unmarshal it depending on StatusCode
		if response.StatusCode < 300 && successResult != nil {
			//success
			err = c.unmarshaler(response.Body, successResult)
		} else if response.StatusCode >= 400 && errorResult != nil {
			//error
			err = c.unmarshaler(response.Body, errorResult)
		}
	}

	if err != nil {
		//we have the response so return it even though unmarshaling might've
		//produced an error
		return response, err
	}

	return response, nil
}

func (c *client) prepareRequest(
	method string,
	path string,
	headers http.Header,
	query url.Values,
	body interface{}) (*http.Request, error) {

	var err error
	reqUrl := cloneUrl(c.base)
	reqUrl.Path += path

	//set headers
	headers = setupHeaders(c.headers, headers)
	//create query
	query = setupQuery(c.query, query)

	if c.marshaler == nil {
		c.marshaler = StringMarshalerFunc
	}

	r, err := http.NewRequest(method, reqUrl.String(), nil)
	if err != nil {
		return nil, err
	}

	r.Header = headers
	r.URL.RawQuery = query.Encode()

	var readLener ReadLener
	if body != nil {

		readLener, err = c.marshaler(body)

		if err != nil {
			return nil, err
		}
		r.ContentLength = int64(readLener.Len())
		r.Body = ioutil.NopCloser(readLener)
	}

	return r, nil
}

func setupHeaders(headers ...http.Header) http.Header {
	finalheaders := make(http.Header)
	for _, current := range headers {
		for i, v := range current {
			finalheaders[i] = v
		}
	}

	return finalheaders
}

func setupQuery(queries ...url.Values) url.Values {
	finalquery := make(url.Values)
	for _, current := range queries {
		for i, v := range current {
			finalquery[i] = v
		}
	}

	return finalquery
}

func (c *client) SetBaseUrl(u *url.URL) error {
	if u == nil {
		return errors.New("Please specify a non nil url.")
	}
	u.RawQuery = ""
	c.base = u
	return nil
}

func (c *client) BaseUrl() *url.URL {
	return c.base
}

func (c *client) Clone() Client {
	cc := &client{}
	cc.base = cloneUrl(c.base)

	cc.reqMutators = make([]RequestMutator, len(c.reqMutators))
	copy(cc.reqMutators, c.reqMutators)

	cc.resMutators = make([]ResponseMutator, len(c.resMutators))
	copy(cc.resMutators, c.resMutators)

	cc.headers = headerCopy(c.headers)
	cc.query = queryCopy(c.query)
	cc.client = c.client
	cc.marshaler = c.marshaler
	cc.unmarshaler = c.unmarshaler

	return cc
}

func headerCopy(h http.Header) http.Header {
	if h == nil {
		return nil
	}
	c := make(http.Header, len(h))
	for i, v := range h {
		c[i] = v
	}
	return c
}

func queryCopy(q url.Values) url.Values {
	if q == nil {
		return nil
	}
	c := make(url.Values, len(q))
	for i, v := range q {
		c[i] = v
	}
	return c
}

func (c *client) GetHttpClient() *http.Client {
	if c.client == nil {
		c.client = http.DefaultClient
	}
	return c.client
}

func (c *client) SetHttpClient(h *http.Client) {
	c.client = h
}

func (c *client) SetMarshaler(f MarshalerFunc) {
	c.marshaler = f
}

func (c *client) SetUnmarshaler(f UnmarshalerFunc) {
	c.unmarshaler = f
}

//New creates a new grestclient with the base url set
//to the passed in paramater.
func New(base *url.URL) (Client, error) {

	if base == nil {
		return nil, errors.New("Please specify a non nil url.")
	}
	c := &client{}
	c.base = base

	return c, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/sha256"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"text/template"
	"time"

	"github.com/spf13/cobra"
)

var (
	EXHIBIT_PREFIX = "exhibit."
	SIGN_PREFIX    = "sign."
	VALUE_SUFFIX   = ".value"

	VARIABLE_REGEX = regexp.MustCompile(`{{\s*([\w\._-]+)\s*}}`)
)

type ContractTemplate struct {
	TemplateHash string
	Vars         []string
	Exhibits     []string
	Signing      []string
}

// return list of all variables and exhibits found in contract template
func parseTemplate(b []byte) ContractTemplate {
	var contractTemplate ContractTemplate
	matches := VARIABLE_REGEX.FindAllStringSubmatch(string(b), -1)
	for _, m := range matches {
		match := m[1]

		if strings.HasPrefix(match, EXHIBIT_PREFIX) {
			exhibitName := strings.TrimPrefix(match, EXHIBIT_PREFIX)
			// only append if it doesn't have .value
			if !strings.HasSuffix(exhibitName, VALUE_SUFFIX) {
				contractTemplate.Exhibits = appendNew(contractTemplate.Exhibits, exhibitName)
			}
		} else if strings.HasPrefix(match, SIGN_PREFIX) {
			signName := strings.TrimPrefix(match, SIGN_PREFIX)
			contractTemplate.Signing = appendNew(contractTemplate.Signing, signName)
		} else {
			contractTemplate.Vars = appendNew(contractTemplate.Vars, match)
		}
	}
	return contractTemplate
}

// copy the template into a new dir and instantiate params file with empty values
func newContract(cmd *cobra.Command, args []string) error {
	if len(args) != 2 {
		return fmt.Errorf("new expects two args: directory for new engagement and template path")
	}

	engagementPath, tmplPath := args[0], args[1]

	b, err := ioutil.ReadFile(tmplPath)
	if err != nil {
		return err
	}

	// get list of variables for a blank config file
	tmpl := parseTemplate(b)

	// compute hash of template
	h := sha256.New()
	h.Write(b)
	templateHash := h.Sum(nil)
	tmpl.TemplateHash = fmt.Sprintf("%X", templateHash)

	// make the directory and copy over the template
	if err := os.MkdirAll(engagementPath, 0700); err != nil {
		return err
	}

	if err := ioutil.WriteFile(filepath.Join(engagementPath, "template.md"), b, 0600); err != nil {
		return err
	}

	// generate history file
	hfString := "// Do not modify this file yourself under any circumstances!"

	if err := ioutil.WriteFile(filepath.Join(engagementPath, "history.toml"), []byte(hfString), 0600); err != nil {
		return err
	}

	paramsFile := generateParamsFile(tmpl)

	// write the params file
	return ioutil.WriteFile(filepath.Join(engagementPath, "params.toml"), paramsFile, 0600)
}

func generateParamsFile(tmpl ContractTemplate) []byte {
	var paramsFileTemplate *template.Template
	var err error
	var buffer bytes.Buffer

	paramsFileTemplate, err = template.New("paramsFile").Parse(paramsFileDefault)
	if err != nil {
		panic(err)
	}

	if err := paramsFileTemplate.Execute(&buffer, tmpl); err != nil {
		panic(err)
	}
	return buffer.Bytes()
}

const paramsFileDefault = `# This is a TOML file containing parameters for this contract

[meta]
# This hash must match the hash of the local template.md file. DO NOT MAKE CHANGES TO THIS HASH.
template = "{{ .TemplateHash}}"

[var]
{{range .Vars}}{{.}} = ""
{{end}}

[exhibit]
{{range .Exhibits}}{{.}} = ""
{{end}}

[sign]
{{range .Signing}}{{.}} = ""
{{end}}
`

//-----------------------------------------

func compileContract(cmd *cobra.Command, args []string) error {

	// load the params from toml file
	params, err := loadConfig()
	if err != nil {
		return err
	}

	// read the contract template
	b, err := ioutil.ReadFile("template.md")
	if err != nil {
		return err
	}

	tmpl := parseTemplate(b)
	exhibits := tmpl.Exhibits

	// substitute params into template holes
	var missingParams []string
	markdownOutput := VARIABLE_REGEX.ReplaceAllFunc(b, func(in []byte) []byte {
		paramName := strings.TrimSuffix(strings.TrimPrefix(string(in), "{{"), "}}")

		// if its an exhibit, we replace it with the exhibit number.
		// if its a var, we replace it with its value
		if strings.HasPrefix(paramName, EXHIBIT_PREFIX) {
			exhibitName := strings.TrimPrefix(paramName, EXHIBIT_PREFIX)

			if strings.HasSuffix(exhibitName, VALUE_SUFFIX) {
				exhibitName = strings.TrimSuffix(exhibitName, VALUE_SUFFIX)
				for _, e := range exhibits {
					if exhibitName == e {
						exhibitValue := params.GetString("exhibit." + exhibitName)
						if exhibitValue != "" {
							return []byte(exhibitValue)
						}
					}
				}
			} else {
				for i, e := range exhibits {
					if exhibitName == e {
						return []byte(fmt.Sprintf("Exhibit %d", i+1))
					}
				}
			}
		} else if strings.HasPrefix(paramName, SIGN_PREFIX) {
			signName := strings.TrimPrefix(paramName, SIGN_PREFIX)
			paramVal := params.GetString("sign." + signName)
			if paramVal != "" {
				return []byte(paramVal)
			}
		} else {
			paramVal := params.GetString("var." + paramName)
			if paramVal != "" {
				return []byte(paramVal)
			}
		}

		missingParams = append(missingParams, paramName)
		return []byte("----")
	})

	// error if params is missing anything
	if len(missingParams) > 0 {
		return fmt.Errorf("Missing params: %v", missingParams)
	}

	switch outputType {
	case "md":
		if err := ioutil.WriteFile("contract.md", markdownOutput, 0600); err != nil {
			return err
		}
	case "html":
		htmlOutput := markdown2html(markdownOutput)
		if err := ioutil.WriteFile("contract.html", htmlOutput, 0600); err != nil {
			return err
		}
	case "pdf":
		// requires the md to be written
		mdPath := "contract.md"
		if err := ioutil.WriteFile(mdPath, markdownOutput, 0600); err != nil {
			return err
		}
		cmd := exec.Command("pandoc", mdPath, "--latex-engine=xelatex", "-o", "contract.pdf")
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		return cmd.Run()
	default:
		return fmt.Errorf("Unknown output format: %v", outputType)
	}

	return nil
}

func reviseContract(cmd *cobra.Command, args []string) error {

	// read the params file
	paramsB, err := ioutil.ReadFile("params.toml")
	if err != nil {
		return err
	}

	// read the contract template
	templateB, err := ioutil.ReadFile("template.md")
	if err != nil {
		return err
	}

	// read history file
	historyB, err := ioutil.ReadFile("history.toml")
	if err != nil {
		return err
	}

	// combine params, history, and template data
	allB := [3][]byte{paramsB, templateB, historyB}
	byteArray := make([]byte, 3, 3)

	for _, element := range allB {
		byteArray = append(byteArray, element...)
	}

	// hash params, history, and template data
	h := sha256.New()
	h.Write(byteArray)
	t := time.Now()
	hashTime := fmt.Sprintf("\n%s: '%X'", t, h.Sum(nil))

	// write hash to history file
	hFile, err := os.OpenFile("history.toml", os.O_RDWR|os.O_APPEND, 0600)
	if err != nil {
		return err
	}

	defer hFile.Close()

	if _, err = hFile.WriteString(hashTime); err != nil {
		return err
	}

	fmt.Println("A revision hash has been added to your history file; your changes are secure.")

	return nil

}
<commit_msg>closes #18<commit_after>package main

import (
	"bytes"
	"crypto/sha256"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"text/template"
	"time"

	"github.com/spf13/cobra"
)

var (
	EXHIBIT_PREFIX = "exhibit."
	SIGN_PREFIX    = "sign."
	VALUE_SUFFIX   = ".value"

	VARIABLE_REGEX = regexp.MustCompile(`{{\s*([\w\._-]+)\s*}}`)
)

type ContractTemplate struct {
	TemplateHash string
	Vars         []string
	Exhibits     []string
	Signing      []string
}

// return list of all variables and exhibits found in contract template
func parseTemplate(b []byte) ContractTemplate {
	var contractTemplate ContractTemplate
	matches := VARIABLE_REGEX.FindAllStringSubmatch(string(b), -1)
	for _, m := range matches {
		match := m[1]

		if strings.HasPrefix(match, EXHIBIT_PREFIX) {
			exhibitName := strings.TrimPrefix(match, EXHIBIT_PREFIX)
			// only append if it doesn't have .value
			if !strings.HasSuffix(exhibitName, VALUE_SUFFIX) {
				contractTemplate.Exhibits = appendNew(contractTemplate.Exhibits, exhibitName)
			}
		} else if strings.HasPrefix(match, SIGN_PREFIX) {
			signName := strings.TrimPrefix(match, SIGN_PREFIX)
			contractTemplate.Signing = appendNew(contractTemplate.Signing, signName)
		} else {
			contractTemplate.Vars = appendNew(contractTemplate.Vars, match)
		}
	}
	return contractTemplate
}

// copy the template into a new dir and instantiate params file with empty values
func newContract(cmd *cobra.Command, args []string) error {
	if len(args) != 2 {
		return fmt.Errorf("new expects two args: directory for new engagement and template path")
	}

	engagementPath, tmplPath := args[0], args[1]

	b, err := ioutil.ReadFile(tmplPath)
	if err != nil {
		return err
	}

	// get list of variables for a blank config file
	tmpl := parseTemplate(b)

	// compute hash of template
	h := sha256.New()
	h.Write(b)
	templateHash := h.Sum(nil)
	tmpl.TemplateHash = fmt.Sprintf("%X", templateHash)

	// make the directory and copy over the template
	if err := os.MkdirAll(engagementPath, 0700); err != nil {
		return err
	}

	if err := ioutil.WriteFile(filepath.Join(engagementPath, "template.md"), b, 0600); err != nil {
		return err
	}

	// generate history file
	hfString := "// Do not modify this file yourself under any circumstances!"

	if err := ioutil.WriteFile(filepath.Join(engagementPath, "history.toml"), []byte(hfString), 0600); err != nil {
		return err
	}

	paramsFile := generateParamsFile(tmpl)

	// write the params file
	return ioutil.WriteFile(filepath.Join(engagementPath, "params.toml"), paramsFile, 0600)
}

func generateParamsFile(tmpl ContractTemplate) []byte {
	var paramsFileTemplate *template.Template
	var err error
	var buffer bytes.Buffer

	paramsFileTemplate, err = template.New("paramsFile").Parse(paramsFileDefault)
	if err != nil {
		panic(err)
	}

	if err := paramsFileTemplate.Execute(&buffer, tmpl); err != nil {
		panic(err)
	}
	return buffer.Bytes()
}

const paramsFileDefault = `# This is a TOML file containing parameters for this contract

[meta]
# This hash must match the hash of the local template.md file. DO NOT MAKE CHANGES TO THIS HASH.
template = "{{ .TemplateHash}}"

[var]
{{range .Vars}}{{.}} = ""
{{end}}

[exhibit]
{{range .Exhibits}}{{.}} = ""
{{end}}

[sign]
{{range .Signing}}{{.}} = ""
{{end}}
`

//-----------------------------------------

func compileContract(cmd *cobra.Command, args []string) error {

	// load the params from toml file
	params, err := loadConfig()
	if err != nil {
		return err
	}

	// read the contract template
	b, err := ioutil.ReadFile("template.md")
	if err != nil {
		return err
	}

	tmpl := parseTemplate(b)
	exhibits := tmpl.Exhibits

	// substitute params into template holes
	var missingParams []string
	markdownOutput := VARIABLE_REGEX.ReplaceAllFunc(b, func(in []byte) []byte {
		paramName := strings.TrimSuffix(strings.TrimPrefix(string(in), "{{"), "}}")

		// if its an exhibit, we replace it with the exhibit number.
		// if its a var, we replace it with its value
		if strings.HasPrefix(paramName, EXHIBIT_PREFIX) {
			exhibitName := strings.TrimPrefix(paramName, EXHIBIT_PREFIX)

			if strings.HasSuffix(exhibitName, VALUE_SUFFIX) {
				exhibitName = strings.TrimSuffix(exhibitName, VALUE_SUFFIX)
				for _, e := range exhibits {
					if exhibitName == e {
						exhibitValue := params.GetString("exhibit." + exhibitName)
						if exhibitValue != "" {
							return []byte(exhibitValue)
						}
					}
				}
			} else {
				for i, e := range exhibits {
					if exhibitName == e {
						return []byte(fmt.Sprintf("Exhibit %d", i+1))
					}
				}
			}
		} else if strings.HasPrefix(paramName, SIGN_PREFIX) {
			signName := strings.TrimPrefix(paramName, SIGN_PREFIX)
			paramVal := params.GetString("sign." + signName)
			if paramVal != "" {
				return []byte(paramVal)
			}
		} else {
			paramVal := params.GetString("var." + paramName)
			if paramVal != "" {
				return []byte(paramVal)
			}
		}

		missingParams = append(missingParams, paramName)
		return []byte("----")
	})

	// error if params is missing anything
	if len(missingParams) > 0 {
		return fmt.Errorf("Missing params: %v", missingParams)
	}

	switch outputType {
	case "md":
		if err := ioutil.WriteFile("contract.md", markdownOutput, 0600); err != nil {
			return err
		}
	case "html":
		htmlOutput := markdown2html(markdownOutput)
		if err := ioutil.WriteFile("contract.html", htmlOutput, 0600); err != nil {
			return err
		}
	case "pdf":
		// requires the md to be written
		mdPath := "contract.md"
		if err := ioutil.WriteFile(mdPath, markdownOutput, 0600); err != nil {
			return err
		}
		cmd := exec.Command("pandoc", mdPath, "--pdf-engine=xelatex", "-o", "contract.pdf")
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		return cmd.Run()
	default:
		return fmt.Errorf("Unknown output format: %v", outputType)
	}

	return nil
}

func reviseContract(cmd *cobra.Command, args []string) error {

	// read the params file
	paramsB, err := ioutil.ReadFile("params.toml")
	if err != nil {
		return err
	}

	// read the contract template
	templateB, err := ioutil.ReadFile("template.md")
	if err != nil {
		return err
	}

	// read history file
	historyB, err := ioutil.ReadFile("history.toml")
	if err != nil {
		return err
	}

	// combine params, history, and template data
	allB := [3][]byte{paramsB, templateB, historyB}
	byteArray := make([]byte, 3, 3)

	for _, element := range allB {
		byteArray = append(byteArray, element...)
	}

	// hash params, history, and template data
	h := sha256.New()
	h.Write(byteArray)
	t := time.Now()
	hashTime := fmt.Sprintf("\n%s: '%X'", t, h.Sum(nil))

	// write hash to history file
	hFile, err := os.OpenFile("history.toml", os.O_RDWR|os.O_APPEND, 0600)
	if err != nil {
		return err
	}

	defer hFile.Close()

	if _, err = hFile.WriteString(hashTime); err != nil {
		return err
	}

	fmt.Println("A revision hash has been added to your history file; your changes are secure.")

	return nil

}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"os"
	"os/exec"
	"path/filepath"
	"testing"

	"github.com/pachyderm/pachyderm/src/client"
	pfsclient "github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/require"
)

func TestTensorFlow(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)

	cwd, err := os.Getwd()
	require.NoError(t, err)
	exampleDir := filepath.Join(cwd, "../../examples/tensor_flow")
	cmd := exec.Command("make", "all")
	cmd.Dir = exampleDir
	_, err = cmd.CombinedOutput()
	require.NoError(t, err)

	commitInfos, err := c.ListCommit([]string{"GoT_scripts"}, nil, client.CommitTypeRead, false, client.CommitStatusAll, nil)
	require.NoError(t, err)
	require.Equal(t, 1, len(commitInfos))
	inputCommitID := commitInfos[0].Commit.ID

	// Wait until the GoT_generate job has finished
	commitInfos, err = c.FlushCommit([]*pfsclient.Commit{client.NewCommit("GoT_scripts", inputCommitID)}, nil)
	require.NoError(t, err)
	require.Equal(t, 2, len(commitInfos))

	repos := []interface{}{"GoT_train", "GoT_generate"}
	var generateCommitID string
	for _, commitInfo := range commitInfos {
		require.EqualOneOf(t, repos, commitInfo.Commit.Repo.Name)
		if commitInfo.Commit.Repo.Name == "GoT_generate" {
			generateCommitID = commitInfo.Commit.ID
		}
	}

	// Make sure the final output is non zero
	var buffer bytes.Buffer
	require.NoError(t, c.GetFile("GoT_generate", generateCommitID, "new_script.txt", 0, 0, "", false, nil, &buffer))
	if buffer.Len() < 100 {
		t.Fatalf("Output GoT script is too small (has len=%v)", buffer.Len())
	}
	require.NoError(t, c.DeleteAll())
}
<commit_msg>Make TestTensorFlow parallelizable<commit_after>package server

import (
	"bytes"
	"os"
	"os/exec"
	"path/filepath"
	"testing"

	"github.com/pachyderm/pachyderm/src/client"
	pfsclient "github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/require"
)

func TestTensorFlow(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	c := getPachClient(t)
	t.Parallel()

	cwd, err := os.Getwd()
	require.NoError(t, err)
	exampleDir := filepath.Join(cwd, "../../examples/tensor_flow")
	cmd := exec.Command("make", "all")
	cmd.Dir = exampleDir
	_, err = cmd.CombinedOutput()
	require.NoError(t, err)

	commitInfos, err := c.ListCommit([]string{"GoT_scripts"}, nil, client.CommitTypeRead, false, client.CommitStatusAll, nil)
	require.NoError(t, err)
	require.Equal(t, 1, len(commitInfos))
	inputCommitID := commitInfos[0].Commit.ID

	// Wait until the GoT_generate job has finished
	commitInfos, err = c.FlushCommit([]*pfsclient.Commit{client.NewCommit("GoT_scripts", inputCommitID)}, nil)
	require.NoError(t, err)
	require.Equal(t, 2, len(commitInfos))

	repos := []interface{}{"GoT_train", "GoT_generate"}
	var generateCommitID string
	for _, commitInfo := range commitInfos {
		require.EqualOneOf(t, repos, commitInfo.Commit.Repo.Name)
		if commitInfo.Commit.Repo.Name == "GoT_generate" {
			generateCommitID = commitInfo.Commit.ID
		}
	}

	// Make sure the final output is non zero
	var buffer bytes.Buffer
	require.NoError(t, c.GetFile("GoT_generate", generateCommitID, "new_script.txt", 0, 0, "", false, nil, &buffer))
	if buffer.Len() < 100 {
		t.Fatalf("Output GoT script is too small (has len=%v)", buffer.Len())
	}
	require.NoError(t, c.DeleteRepo("GoT_generate", false))
	require.NoError(t, c.DeleteRepo("GoT_train", false))
	require.NoError(t, c.DeleteRepo("GoT_scripts", false))
}
<|endoftext|>


<commit_before>package fuse

import (
	"io/ioutil"
	"os"
	"os/signal"
	pathpkg "path"
	"path/filepath"
	"strings"

	"github.com/hanwen/go-fuse/v2/fs"

	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pkg/errors"
	"github.com/pachyderm/pachyderm/src/server/pkg/uuid"
)

// Mount pfs to target, opts may be left nil.
func Mount(c *client.APIClient, target string, opts *Options) (retErr error) {
	if err := opts.validate(); err != nil {
		return err
	}
	commits := make(map[string]string)
	for repo, branch := range opts.getBranches() {
		if uuid.IsUUIDWithoutDashes(branch) {
			commits[repo] = branch
		}
	}
	// branches := opts.getBranches()
	rootDir, err := ioutil.TempDir("", "pfs")
	if err != nil {
		return err
	}
	defer func() {
		if err := os.RemoveAll(rootDir); err != nil && retErr == nil {
			retErr = err
		}
	}()
	root, err := newLoopbackRoot(rootDir, target, c, opts)
	if err != nil {
		return err
	}
	server, err := fs.Mount(target, root, opts.getFuse())
	if err != nil {
		return errors.Wrap(err, "fs.Mount")
	}
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)
	go func() {
		select {
		case <-sigChan:
		case <-opts.getUnmount():
		}
		server.Unmount()
	}()
	server.Serve()
	pfcs := make(map[string]client.PutFileClient)
	pfc := func(repo string) (client.PutFileClient, error) {
		if pfc, ok := pfcs[repo]; ok {
			return pfc, nil
		}
		pfc, err := c.NewPutFileClient()
		if err != nil {
			return nil, err
		}
		pfcs[repo] = pfc
		return pfc, nil
	}
	defer func() {
		for _, pfc := range pfcs {
			if err := pfc.Close(); err != nil && retErr == nil {
				retErr = err
			}
		}
	}()
	for path, state := range root.files {
		if state != dirty {
			continue
		}
		parts := strings.Split(path, "/")
		pfc, err := pfc(parts[0])
		if err != nil {
			return err
		}
		if err := func() (retErr error) {
			f, err := os.Open(filepath.Join(root.rootPath, path))
			if err != nil {
				if os.IsNotExist(err) {
					return pfc.DeleteFile(parts[0], root.branch(parts[0]), pathpkg.Join(parts[1:]...))
				}
				return errors.Wrap(err, "os.Open")
			}
			defer func() {
				if err := f.Close(); err != nil && retErr == nil {
					retErr = errors.Wrap(err, "f.Close")
				}
			}()
			if _, err := pfc.PutFileOverwrite(parts[0], root.branch(parts[0]),
				pathpkg.Join(parts[1:]...), f, 0); err != nil {
				return err
			}
			return nil
		}(); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>Adds progress bars for fuse upload.<commit_after>package fuse

import (
	"io/ioutil"
	"os"
	"os/signal"
	pathpkg "path"
	"path/filepath"
	"strings"

	"github.com/hanwen/go-fuse/v2/fs"

	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pkg/errors"
	"github.com/pachyderm/pachyderm/src/server/pkg/progress"
	"github.com/pachyderm/pachyderm/src/server/pkg/uuid"
)

// Mount pfs to target, opts may be left nil.
func Mount(c *client.APIClient, target string, opts *Options) (retErr error) {
	if err := opts.validate(); err != nil {
		return err
	}
	commits := make(map[string]string)
	for repo, branch := range opts.getBranches() {
		if uuid.IsUUIDWithoutDashes(branch) {
			commits[repo] = branch
		}
	}
	// branches := opts.getBranches()
	rootDir, err := ioutil.TempDir("", "pfs")
	if err != nil {
		return err
	}
	defer func() {
		if err := os.RemoveAll(rootDir); err != nil && retErr == nil {
			retErr = err
		}
	}()
	root, err := newLoopbackRoot(rootDir, target, c, opts)
	if err != nil {
		return err
	}
	server, err := fs.Mount(target, root, opts.getFuse())
	if err != nil {
		return errors.Wrap(err, "fs.Mount")
	}
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)
	go func() {
		select {
		case <-sigChan:
		case <-opts.getUnmount():
		}
		server.Unmount()
	}()
	server.Serve()
	pfcs := make(map[string]client.PutFileClient)
	pfc := func(repo string) (client.PutFileClient, error) {
		if pfc, ok := pfcs[repo]; ok {
			return pfc, nil
		}
		pfc, err := c.NewPutFileClient()
		if err != nil {
			return nil, err
		}
		pfcs[repo] = pfc
		return pfc, nil
	}
	defer func() {
		for _, pfc := range pfcs {
			if err := pfc.Close(); err != nil && retErr == nil {
				retErr = err
			}
		}
	}()
	for path, state := range root.files {
		if state != dirty {
			continue
		}
		parts := strings.Split(path, "/")
		pfc, err := pfc(parts[0])
		if err != nil {
			return err
		}
		if err := func() (retErr error) {
			f, err := progress.Open(filepath.Join(root.rootPath, path))
			if err != nil {
				if os.IsNotExist(err) {
					return pfc.DeleteFile(parts[0], root.branch(parts[0]), pathpkg.Join(parts[1:]...))
				}
				return errors.Wrap(err, "os.Open")
			}
			defer func() {
				if err := f.Close(); err != nil && retErr == nil {
					retErr = errors.Wrap(err, "f.Close")
				}
			}()
			if _, err := pfc.PutFileOverwrite(parts[0], root.branch(parts[0]),
				pathpkg.Join(parts[1:]...), f, 0); err != nil {
				return err
			}
			return nil
		}(); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Package sync provides utility functions similar to `git pull/push` for PFS
package sync

import (
	"io"
	"os"
	"path"
	"path/filepath"
	"sync"
	"sync/atomic"
	"syscall"

	pachclient "github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/limit"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/server/pkg/hashtree"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"

	"golang.org/x/sync/errgroup"
)

// Puller as a struct for managing a Pull operation.
type Puller struct {
	sync.Mutex
	// errCh contains an error from the pipe goros
	errCh chan error
	// pipes is a set containing all pipes that are currently blocking
	pipes map[string]bool
	// cleaned signals if the cleanup goroutine has been started
	cleaned bool
	// wg is used to wait for all goroutines associated with this Puller
	// to complete.
	wg sync.WaitGroup
	// size is the total amount this puller has pulled
	size int64
}

// NewPuller creates a new Puller struct.
func NewPuller() *Puller {
	return &Puller{
		errCh: make(chan error, 1),
		pipes: make(map[string]bool),
	}
}

type sizeWriter struct {
	w    io.Writer
	size int64
}

func (s *sizeWriter) Write(p []byte) (int, error) {
	n, err := s.w.Write(p)
	s.size += int64(n)
	return n, err
}

func (p *Puller) makePipe(path string, f func(io.Writer) error) error {
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		return err
	}
	if err := syscall.Mkfifo(path, 0666); err != nil {
		return err
	}
	func() {
		p.Lock()
		defer p.Unlock()
		p.pipes[path] = true
	}()
	// This goro will block until the user's code opens the
	// fifo.  That means we need to "abandon" this goro so that
	// the function can return and the caller can execute the
	// user's code. Waiting for this goro to return would
	// produce a deadlock. This goro will exit (if it hasn't already)
	// when CleanUp is called.
	p.wg.Add(1)
	go func() {
		defer p.wg.Done()
		if err := func() (retErr error) {
			file, err := os.OpenFile(path, os.O_WRONLY, os.ModeNamedPipe)
			if err != nil {
				return err
			}
			defer func() {
				if err := file.Close(); err != nil && retErr == nil {
					retErr = err
				}
			}()
			// If the CleanUp routine has already run, then there's
			// no point in downloading and sending the file, so we
			// exit early.
			if func() bool {
				p.Lock()
				defer p.Unlock()
				delete(p.pipes, path)
				return p.cleaned
			}() {
				return nil
			}
			w := &sizeWriter{w: file}
			if err := f(w); err != nil {
				return err
			}
			atomic.AddInt64(&p.size, w.size)
			return nil
		}(); err != nil {
			select {
			case p.errCh <- err:
			default:
			}
		}
	}()
	return nil
}

func (p *Puller) makeFile(path string, f func(io.Writer) error) (retErr error) {
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		return err
	}
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer func() {
		if err := file.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	w := &sizeWriter{w: file}
	if err := f(w); err != nil {
		return err
	}
	atomic.AddInt64(&p.size, w.size)
	return nil
}

// Pull clones an entire repo at a certain commit.
// root is the local path you want to clone to.
// fileInfo is the file/dir we are puuling.
// pipes causes the function to create named pipes in place of files, thus
// lazily downloading the data as it's needed.
// tree is a hashtree to mirror the pulled content into (it may be left nil)
// treeRoot is the root the data is mirrored to within tree
func (p *Puller) Pull(client *pachclient.APIClient, root string, repo, commit, file string,
	pipes bool, concurrency int, tree hashtree.OpenHashTree, treeRoot string) error {
	limiter := limit.New(concurrency)
	var eg errgroup.Group
	if err := client.Walk(repo, commit, file, func(fileInfo *pfs.FileInfo) error {
		basepath, err := filepath.Rel(file, fileInfo.File.Path)
		if err != nil {
			return err
		}
		if tree != nil {
			treePath := path.Join(treeRoot, basepath)
			if err := tree.PutFile(treePath, fileInfo.Objects, int64(fileInfo.SizeBytes)); err != nil {
				return err
			}
		}
		path := filepath.Join(root, basepath)
		if fileInfo.FileType == pfs.FileType_DIR {
			return os.MkdirAll(path, 0700)
		}
		if pipes {
			return p.makePipe(path, func(w io.Writer) error {
				return client.GetFile(repo, commit, fileInfo.File.Path, 0, 0, w)
			})
		}
		eg.Go(func() (retErr error) {
			limiter.Acquire()
			defer limiter.Release()
			return p.makeFile(path, func(w io.Writer) error {
				return client.GetFile(repo, commit, fileInfo.File.Path, 0, 0, w)
			})
		})
		return nil
	}); err != nil {
		return err
	}
	return eg.Wait()
}

// PullDiff is like Pull except that it materializes a Diff of the content
// rather than a the actual content. If newOnly is true then only new files
// will be downloaded and they will be downloaded under root. Otherwise new and
// old files will be downloaded under root/new and root/old respectively.
func (p *Puller) PullDiff(client *pachclient.APIClient, root string, newRepo, newCommit, newPath, oldRepo, oldCommit, oldPath string,
	newOnly bool, pipes bool, concurrency int, tree hashtree.OpenHashTree, treeRoot string) error {
	limiter := limit.New(concurrency)
	var eg errgroup.Group
	newFiles, oldFiles, err := client.DiffFile(newRepo, newCommit, newPath, oldRepo, oldCommit, oldPath)
	if err != nil {
		return err
	}
	for _, newFile := range newFiles {
		basepath, err := filepath.Rel(newPath, newFile.File.Path)
		if err != nil {
			return err
		}
		if tree != nil {
			treePath := path.Join(treeRoot, "new", basepath)
			if newOnly {
				treePath = path.Join(treeRoot, basepath)
			}
			if err := tree.PutFile(treePath, newFile.Objects, int64(newFile.SizeBytes)); err != nil {
				return err
			}
		}
		path := filepath.Join(root, "new", basepath)
		if newOnly {
			path = filepath.Join(root, basepath)
		}
		if pipes {
			if err := p.makePipe(path, func(w io.Writer) error {
				return client.GetFile(newFile.File.Commit.Repo.Name, newFile.File.Commit.ID, newFile.File.Path, 0, 0, w)
			}); err != nil {
				return err
			}
		} else {
			newFile := newFile
			limiter.Acquire()
			eg.Go(func() error {
				defer limiter.Release()
				return p.makeFile(path, func(w io.Writer) error {
					return client.GetFile(newFile.File.Commit.Repo.Name, newFile.File.Commit.ID, newFile.File.Path, 0, 0, w)
				})
			})
		}
	}
	if !newOnly {
		for _, oldFile := range oldFiles {
			basepath, err := filepath.Rel(oldPath, oldFile.File.Path)
			if err != nil {
				return err
			}
			if tree != nil {
				treePath := path.Join(treeRoot, "old", basepath)
				if err := tree.PutFile(treePath, oldFile.Objects, int64(oldFile.SizeBytes)); err != nil {
					return err
				}
			}
			path := filepath.Join(root, "old", basepath)
			if pipes {
				if err := p.makePipe(path, func(w io.Writer) error {
					return client.GetFile(oldFile.File.Commit.Repo.Name, oldFile.File.Commit.ID, oldFile.File.Path, 0, 0, w)
				}); err != nil {
					return err
				}
			} else {
				oldFile := oldFile
				limiter.Acquire()
				eg.Go(func() error {
					defer limiter.Release()
					return p.makeFile(path, func(w io.Writer) error {
						return client.GetFile(oldFile.File.Commit.Repo.Name, oldFile.File.Commit.ID, oldFile.File.Path, 0, 0, w)
					})
				})
			}
		}
	}
	return eg.Wait()
}

// PullTree pulls from a raw HashTree rather than a repo.
func (p *Puller) PullTree(client *pachclient.APIClient, root string, tree hashtree.HashTree, pipes bool, concurrency int) error {
	limiter := limit.New(concurrency)
	var eg errgroup.Group
	if err := tree.Walk(func(path string, node *hashtree.NodeProto) error {
		if node.FileNode != nil {
			path := filepath.Join(root, path)
			var hashes []string
			for _, object := range node.FileNode.Objects {
				hashes = append(hashes, object.Hash)
			}
			if pipes {
				return p.makePipe(path, func(w io.Writer) error {
					return client.GetObjects(hashes, 0, 0, w)
				})
			}
			limiter.Acquire()
			eg.Go(func() (retErr error) {
				defer limiter.Release()
				return p.makeFile(path, func(w io.Writer) error {
					return client.GetObjects(hashes, 0, 0, w)
				})
			})
		}
		return nil
	}); err != nil {
		return err
	}
	return eg.Wait()
}

// CleanUp cleans up blocked syscalls for pipes that were never opened. And
// returns the total number of bytes that have been pulled/pushed. It also
// returns any errors that might have been encountered while trying to read
// data for the pipes. CleanUp should be called after all code that might
// access pipes has completed running, it should not be called concurrently.
func (p *Puller) CleanUp() (int64, error) {
	var result error
	select {
	case result = <-p.errCh:
	default:
	}

	// Open all the pipes to unblock the goros
	var pipes []io.Closer
	func() {
		p.Lock()
		defer p.Unlock()
		p.cleaned = true
		for path := range p.pipes {
			f, err := os.OpenFile(path, syscall.O_NONBLOCK+os.O_RDONLY, os.ModeNamedPipe)
			if err != nil && result == nil {
				result = err
			}
			pipes = append(pipes, f)
		}
		p.pipes = make(map[string]bool)
	}()

	// Wait for all goros to exit
	p.wg.Wait()

	// Close the pipes
	for _, pipe := range pipes {
		if err := pipe.Close(); err != nil && result == nil {
			result = err
		}
	}
	size := p.size
	p.size = 0
	return size, result
}

// Push puts files under root into an open commit.
func Push(client *pachclient.APIClient, root string, commit *pfs.Commit, overwrite bool) error {
	var g errgroup.Group
	if err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		g.Go(func() (retErr error) {
			if path == root || info.IsDir() {
				return nil
			}

			f, err := os.Open(path)
			if err != nil {
				return err
			}
			defer func() {
				if err := f.Close(); err != nil && retErr == nil {
					retErr = err
				}
			}()

			relPath, err := filepath.Rel(root, path)
			if err != nil {
				return err
			}

			if overwrite {
				if err := client.DeleteFile(commit.Repo.Name, commit.ID, relPath); err != nil {
					return err
				}
			}

			_, err = client.PutFile(commit.Repo.Name, commit.ID, relPath, f)
			return err
		})
		return nil
	}); err != nil {
		return err
	}

	return g.Wait()
}

// PushObj pushes data from commit to an object store.
func PushObj(pachClient pachclient.APIClient, commit *pfs.Commit, objClient obj.Client, root string) error {
	var eg errgroup.Group
	sem := make(chan struct{}, 200)
	if err := pachClient.Walk(commit.Repo.Name, commit.ID, "", func(fileInfo *pfs.FileInfo) error {
		if fileInfo.FileType != pfs.FileType_FILE {
			return nil
		}
		eg.Go(func() (retErr error) {
			sem <- struct{}{}
			defer func() { <-sem }()
			w, err := objClient.Writer(filepath.Join(root, fileInfo.File.Path))
			if err != nil {
				return err
			}
			defer func() {
				if err := w.Close(); err != nil && retErr == nil {
					retErr = err
				}
			}()
			return pachClient.GetFile(commit.Repo.Name, commit.ID, fileInfo.File.Path, 0, 0, w)
		})
		return nil
	}); err != nil {
		return err
	}
	return eg.Wait()
}
<commit_msg>Don't pull directories as files to stats branch.<commit_after>// Package sync provides utility functions similar to `git pull/push` for PFS
package sync

import (
	"io"
	"os"
	"path"
	"path/filepath"
	"sync"
	"sync/atomic"
	"syscall"

	pachclient "github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/limit"
	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/server/pkg/hashtree"
	"github.com/pachyderm/pachyderm/src/server/pkg/obj"

	"golang.org/x/sync/errgroup"
)

// Puller as a struct for managing a Pull operation.
type Puller struct {
	sync.Mutex
	// errCh contains an error from the pipe goros
	errCh chan error
	// pipes is a set containing all pipes that are currently blocking
	pipes map[string]bool
	// cleaned signals if the cleanup goroutine has been started
	cleaned bool
	// wg is used to wait for all goroutines associated with this Puller
	// to complete.
	wg sync.WaitGroup
	// size is the total amount this puller has pulled
	size int64
}

// NewPuller creates a new Puller struct.
func NewPuller() *Puller {
	return &Puller{
		errCh: make(chan error, 1),
		pipes: make(map[string]bool),
	}
}

type sizeWriter struct {
	w    io.Writer
	size int64
}

func (s *sizeWriter) Write(p []byte) (int, error) {
	n, err := s.w.Write(p)
	s.size += int64(n)
	return n, err
}

func (p *Puller) makePipe(path string, f func(io.Writer) error) error {
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		return err
	}
	if err := syscall.Mkfifo(path, 0666); err != nil {
		return err
	}
	func() {
		p.Lock()
		defer p.Unlock()
		p.pipes[path] = true
	}()
	// This goro will block until the user's code opens the
	// fifo.  That means we need to "abandon" this goro so that
	// the function can return and the caller can execute the
	// user's code. Waiting for this goro to return would
	// produce a deadlock. This goro will exit (if it hasn't already)
	// when CleanUp is called.
	p.wg.Add(1)
	go func() {
		defer p.wg.Done()
		if err := func() (retErr error) {
			file, err := os.OpenFile(path, os.O_WRONLY, os.ModeNamedPipe)
			if err != nil {
				return err
			}
			defer func() {
				if err := file.Close(); err != nil && retErr == nil {
					retErr = err
				}
			}()
			// If the CleanUp routine has already run, then there's
			// no point in downloading and sending the file, so we
			// exit early.
			if func() bool {
				p.Lock()
				defer p.Unlock()
				delete(p.pipes, path)
				return p.cleaned
			}() {
				return nil
			}
			w := &sizeWriter{w: file}
			if err := f(w); err != nil {
				return err
			}
			atomic.AddInt64(&p.size, w.size)
			return nil
		}(); err != nil {
			select {
			case p.errCh <- err:
			default:
			}
		}
	}()
	return nil
}

func (p *Puller) makeFile(path string, f func(io.Writer) error) (retErr error) {
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		return err
	}
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer func() {
		if err := file.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	w := &sizeWriter{w: file}
	if err := f(w); err != nil {
		return err
	}
	atomic.AddInt64(&p.size, w.size)
	return nil
}

// Pull clones an entire repo at a certain commit.
// root is the local path you want to clone to.
// fileInfo is the file/dir we are puuling.
// pipes causes the function to create named pipes in place of files, thus
// lazily downloading the data as it's needed.
// tree is a hashtree to mirror the pulled content into (it may be left nil)
// treeRoot is the root the data is mirrored to within tree
func (p *Puller) Pull(client *pachclient.APIClient, root string, repo, commit, file string,
	pipes bool, concurrency int, tree hashtree.OpenHashTree, treeRoot string) error {
	limiter := limit.New(concurrency)
	var eg errgroup.Group
	if err := client.Walk(repo, commit, file, func(fileInfo *pfs.FileInfo) error {
		basepath, err := filepath.Rel(file, fileInfo.File.Path)
		if err != nil {
			return err
		}
		if tree != nil {
			treePath := path.Join(treeRoot, basepath)
			if fileInfo.FileType == pfs.FileType_DIR {
				if err := tree.PutDir(treePath); err != nil {
					return err
				}
			} else {
				if err := tree.PutFile(treePath, fileInfo.Objects, int64(fileInfo.SizeBytes)); err != nil {
					return err
				}
			}
		}
		path := filepath.Join(root, basepath)
		if fileInfo.FileType == pfs.FileType_DIR {
			return os.MkdirAll(path, 0700)
		}
		if pipes {
			return p.makePipe(path, func(w io.Writer) error {
				return client.GetFile(repo, commit, fileInfo.File.Path, 0, 0, w)
			})
		}
		eg.Go(func() (retErr error) {
			limiter.Acquire()
			defer limiter.Release()
			return p.makeFile(path, func(w io.Writer) error {
				return client.GetFile(repo, commit, fileInfo.File.Path, 0, 0, w)
			})
		})
		return nil
	}); err != nil {
		return err
	}
	return eg.Wait()
}

// PullDiff is like Pull except that it materializes a Diff of the content
// rather than a the actual content. If newOnly is true then only new files
// will be downloaded and they will be downloaded under root. Otherwise new and
// old files will be downloaded under root/new and root/old respectively.
func (p *Puller) PullDiff(client *pachclient.APIClient, root string, newRepo, newCommit, newPath, oldRepo, oldCommit, oldPath string,
	newOnly bool, pipes bool, concurrency int, tree hashtree.OpenHashTree, treeRoot string) error {
	limiter := limit.New(concurrency)
	var eg errgroup.Group
	newFiles, oldFiles, err := client.DiffFile(newRepo, newCommit, newPath, oldRepo, oldCommit, oldPath)
	if err != nil {
		return err
	}
	for _, newFile := range newFiles {
		basepath, err := filepath.Rel(newPath, newFile.File.Path)
		if err != nil {
			return err
		}
		if tree != nil {
			treePath := path.Join(treeRoot, "new", basepath)
			if newOnly {
				treePath = path.Join(treeRoot, basepath)
			}
			if err := tree.PutFile(treePath, newFile.Objects, int64(newFile.SizeBytes)); err != nil {
				return err
			}
		}
		path := filepath.Join(root, "new", basepath)
		if newOnly {
			path = filepath.Join(root, basepath)
		}
		if pipes {
			if err := p.makePipe(path, func(w io.Writer) error {
				return client.GetFile(newFile.File.Commit.Repo.Name, newFile.File.Commit.ID, newFile.File.Path, 0, 0, w)
			}); err != nil {
				return err
			}
		} else {
			newFile := newFile
			limiter.Acquire()
			eg.Go(func() error {
				defer limiter.Release()
				return p.makeFile(path, func(w io.Writer) error {
					return client.GetFile(newFile.File.Commit.Repo.Name, newFile.File.Commit.ID, newFile.File.Path, 0, 0, w)
				})
			})
		}
	}
	if !newOnly {
		for _, oldFile := range oldFiles {
			basepath, err := filepath.Rel(oldPath, oldFile.File.Path)
			if err != nil {
				return err
			}
			if tree != nil {
				treePath := path.Join(treeRoot, "old", basepath)
				if err := tree.PutFile(treePath, oldFile.Objects, int64(oldFile.SizeBytes)); err != nil {
					return err
				}
			}
			path := filepath.Join(root, "old", basepath)
			if pipes {
				if err := p.makePipe(path, func(w io.Writer) error {
					return client.GetFile(oldFile.File.Commit.Repo.Name, oldFile.File.Commit.ID, oldFile.File.Path, 0, 0, w)
				}); err != nil {
					return err
				}
			} else {
				oldFile := oldFile
				limiter.Acquire()
				eg.Go(func() error {
					defer limiter.Release()
					return p.makeFile(path, func(w io.Writer) error {
						return client.GetFile(oldFile.File.Commit.Repo.Name, oldFile.File.Commit.ID, oldFile.File.Path, 0, 0, w)
					})
				})
			}
		}
	}
	return eg.Wait()
}

// PullTree pulls from a raw HashTree rather than a repo.
func (p *Puller) PullTree(client *pachclient.APIClient, root string, tree hashtree.HashTree, pipes bool, concurrency int) error {
	limiter := limit.New(concurrency)
	var eg errgroup.Group
	if err := tree.Walk(func(path string, node *hashtree.NodeProto) error {
		if node.FileNode != nil {
			path := filepath.Join(root, path)
			var hashes []string
			for _, object := range node.FileNode.Objects {
				hashes = append(hashes, object.Hash)
			}
			if pipes {
				return p.makePipe(path, func(w io.Writer) error {
					return client.GetObjects(hashes, 0, 0, w)
				})
			}
			limiter.Acquire()
			eg.Go(func() (retErr error) {
				defer limiter.Release()
				return p.makeFile(path, func(w io.Writer) error {
					return client.GetObjects(hashes, 0, 0, w)
				})
			})
		}
		return nil
	}); err != nil {
		return err
	}
	return eg.Wait()
}

// CleanUp cleans up blocked syscalls for pipes that were never opened. And
// returns the total number of bytes that have been pulled/pushed. It also
// returns any errors that might have been encountered while trying to read
// data for the pipes. CleanUp should be called after all code that might
// access pipes has completed running, it should not be called concurrently.
func (p *Puller) CleanUp() (int64, error) {
	var result error
	select {
	case result = <-p.errCh:
	default:
	}

	// Open all the pipes to unblock the goros
	var pipes []io.Closer
	func() {
		p.Lock()
		defer p.Unlock()
		p.cleaned = true
		for path := range p.pipes {
			f, err := os.OpenFile(path, syscall.O_NONBLOCK+os.O_RDONLY, os.ModeNamedPipe)
			if err != nil && result == nil {
				result = err
			}
			pipes = append(pipes, f)
		}
		p.pipes = make(map[string]bool)
	}()

	// Wait for all goros to exit
	p.wg.Wait()

	// Close the pipes
	for _, pipe := range pipes {
		if err := pipe.Close(); err != nil && result == nil {
			result = err
		}
	}
	size := p.size
	p.size = 0
	return size, result
}

// Push puts files under root into an open commit.
func Push(client *pachclient.APIClient, root string, commit *pfs.Commit, overwrite bool) error {
	var g errgroup.Group
	if err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		g.Go(func() (retErr error) {
			if path == root || info.IsDir() {
				return nil
			}

			f, err := os.Open(path)
			if err != nil {
				return err
			}
			defer func() {
				if err := f.Close(); err != nil && retErr == nil {
					retErr = err
				}
			}()

			relPath, err := filepath.Rel(root, path)
			if err != nil {
				return err
			}

			if overwrite {
				if err := client.DeleteFile(commit.Repo.Name, commit.ID, relPath); err != nil {
					return err
				}
			}

			_, err = client.PutFile(commit.Repo.Name, commit.ID, relPath, f)
			return err
		})
		return nil
	}); err != nil {
		return err
	}

	return g.Wait()
}

// PushObj pushes data from commit to an object store.
func PushObj(pachClient pachclient.APIClient, commit *pfs.Commit, objClient obj.Client, root string) error {
	var eg errgroup.Group
	sem := make(chan struct{}, 200)
	if err := pachClient.Walk(commit.Repo.Name, commit.ID, "", func(fileInfo *pfs.FileInfo) error {
		if fileInfo.FileType != pfs.FileType_FILE {
			return nil
		}
		eg.Go(func() (retErr error) {
			sem <- struct{}{}
			defer func() { <-sem }()
			w, err := objClient.Writer(filepath.Join(root, fileInfo.File.Path))
			if err != nil {
				return err
			}
			defer func() {
				if err := w.Close(); err != nil && retErr == nil {
					retErr = err
				}
			}()
			return pachClient.GetFile(commit.Repo.Name, commit.ID, fileInfo.File.Path, 0, 0, w)
		})
		return nil
	}); err != nil {
		return err
	}
	return eg.Wait()
}
<|endoftext|>


<commit_before>package cors

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

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
)

func init() {
	gin.SetMode(gin.TestMode)
}

func performRequest(r http.Handler, method, path string) *httptest.ResponseRecorder {
	req, _ := http.NewRequest(method, path, nil)
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}

func TestBadConfig(t *testing.T) {
	assert.Panics(t, func() { New(Config{}) })
	assert.Panics(t, func() {
		New(Config{
			AllowAllOrigins: true,
			AllowedOrigins:  []string{"http://google.com"},
		})
	})
	assert.Panics(t, func() {
		New(Config{
			AllowAllOrigins: true,
			AllowOriginFunc: func(origin string) bool { return false },
		})
	})
	assert.Panics(t, func() {
		New(Config{
			AllowedOrigins: []string{"google.com"},
		})
	})
}

func TestNormalize(t *testing.T) {
	values := normalize([]string{
		"http-Access ", "Post", "POST", " poSt  ",
		"HTTP-Access", "",
	})
	assert.Equal(t, values, []string{"http-access", "post", ""})

	values = normalize(nil)
	assert.Nil(t, values)

	values = normalize([]string{})
	assert.Equal(t, values, []string{})
}

func TestGenerateNormalHeaders_AllowAllOrigins(t *testing.T) {
	header := generateNormalHeaders(Config{
		AllowAllOrigins: false,
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Origin"), "")
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 1)

	header = generateNormalHeaders(Config{
		AllowAllOrigins: true,
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Origin"), "*")
	assert.Equal(t, header.Get("Vary"), "")
	assert.Len(t, header, 1)
}

func TestGenerateNormalHeaders_AllowCredentials(t *testing.T) {
	header := generateNormalHeaders(Config{
		AllowCredentials: true,
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Credentials"), "true")
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 2)
}

func TestGenerateNormalHeaders_ExposedHeaders(t *testing.T) {
	header := generateNormalHeaders(Config{
		ExposedHeaders: []string{"X-user", "xPassword"},
	})
	assert.Equal(t, header.Get("Access-Control-Expose-Headers"), "x-user, xpassword")
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 2)
}

func TestGeneratePreflightHeaders(t *testing.T) {
	header := generatePreflightHeaders(Config{
		AllowAllOrigins: false,
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Origin"), "")
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 1)

	header = generateNormalHeaders(Config{
		AllowAllOrigins: true,
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Origin"), "*")
	assert.Equal(t, header.Get("Vary"), "")
	assert.Len(t, header, 1)
}

func TestGeneratePreflightHeaders_AllowCredentials(t *testing.T) {
	header := generatePreflightHeaders(Config{
		AllowCredentials: true,
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Credentials"), "true")
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 2)
}

func TestGeneratePreflightHeaders_AllowedMethods(t *testing.T) {
	header := generatePreflightHeaders(Config{
		AllowedMethods: []string{"GET ", "post", "PUT", " put  "},
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Methods"), "get, post, put")
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 2)
}

func TestGeneratePreflightHeaders_AllowedHeaders(t *testing.T) {
	header := generatePreflightHeaders(Config{
		AllowedHeaders: []string{"X-user", "Content-Type"},
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Headers"), "x-user, content-type")
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 2)
}

func TestGeneratePreflightHeaders_MaxAge(t *testing.T) {
	header := generatePreflightHeaders(Config{
		MaxAge: 12 * time.Hour,
	})
	assert.Equal(t, header.Get("Access-Control-Max-Age"), "43200") // 12*60*60
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 2)
}

func TestValidateOrigin(t *testing.T) {
	cors := newCors(Config{
		AllowAllOrigins: true,
	})
	assert.True(t, cors.validateOrigin("http://google.com"))
	assert.True(t, cors.validateOrigin("https://google.com"))
	assert.True(t, cors.validateOrigin("example.com"))

	cors = newCors(Config{
		AllowedOrigins: []string{"https://google.com", "https://github.com"},
		AllowOriginFunc: func(origin string) bool {
			return (origin == "http://news.ycombinator.com")
		},
	})
	assert.False(t, cors.validateOrigin("http://google.com"))
	assert.True(t, cors.validateOrigin("https://google.com"))
	assert.True(t, cors.validateOrigin("https://github.com"))
	assert.True(t, cors.validateOrigin("http://news.ycombinator.com"))
	assert.False(t, cors.validateOrigin("http://example.com"))
	assert.False(t, cors.validateOrigin("google.com"))
}

func TestPasses0(t *testing.T) {
	called := false
	router := gin.New()
	router.Use(New(Config{
		AllowedOrigins:   []string{"http://google.com"},
		AllowedMethods:   []string{" GeT ", "get", "post", "PUT  ", "Head", "POST"},
		AllowedHeaders:   []string{"Content-type", "timeStamp "},
		ExposedHeaders:   []string{"Data", "x-User"},
		AllowCredentials: true,
		MaxAge:           12 * time.Hour,
		AllowOriginFunc: func(origin string) bool {
			return origin == "http://github.com"
		},
	}))
	router.GET("/", func(c *gin.Context) {
		called = true
	})

	w := httptest.NewRecorder()
	req, _ := http.NewRequest("GET", "/", nil)
	router.ServeHTTP(w, req)
	assert.True(t, called)
	assert.NotContains(t, w.Header(), "Access-Control-Allow-Origin")
	assert.NotContains(t, w.Header(), "Access-Control-Allow-Credentials")
	assert.NotContains(t, w.Header(), "Access-Control-Expose-Headers")

	called = false
	w = httptest.NewRecorder()
	req, _ = http.NewRequest("GET", "/", nil)
	req.Header.Set("Origin", "http://google.com")
	router.ServeHTTP(w, req)
	assert.True(t, called)
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Origin"), "http://google.com")
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Credentials"), "true")
	assert.Equal(t, w.Header().Get("Access-Control-Expose-Headers"), "data, x-user")

	called = false
	w = httptest.NewRecorder()
	req, _ = http.NewRequest("GET", "/", nil)
	req.Header.Set("Origin", "https://google.com")
	router.ServeHTTP(w, req)
	assert.False(t, called)
	assert.NotContains(t, w.Header(), "Access-Control-Allow-Origin")
	assert.NotContains(t, w.Header(), "Access-Control-Allow-Credentials")
	assert.NotContains(t, w.Header(), "Access-Control-Expose-Headers")

	called = false
	w = httptest.NewRecorder()
	req, _ = http.NewRequest("OPTIONS", "/", nil)
	req.Header.Set("Origin", "http://github.com")
	router.ServeHTTP(w, req)
	assert.False(t, called)
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Origin"), "http://github.com")
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Credentials"), "true")
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Methods"), "get, post, put, head")
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Headers"), "content-type, timestamp")
	assert.Equal(t, w.Header().Get("Access-Control-Max-Age"), "43200")

	called = false
	w = httptest.NewRecorder()
	req, _ = http.NewRequest("OPTIONS", "/", nil)
	req.Header.Set("Origin", "http://example.com")
	router.ServeHTTP(w, req)
	assert.False(t, called)
	assert.NotContains(t, w.Header(), "Access-Control-Allow-Origin")
	assert.NotContains(t, w.Header(), "Access-Control-Allow-Credentials")
	assert.NotContains(t, w.Header(), "Access-Control-Allow-Methods")
	assert.NotContains(t, w.Header(), "Access-Control-Allow-Headers")
	assert.NotContains(t, w.Header(), "Access-Control-Max-Age")
}

func TestPasses1(t *testing.T) {

}

func TestPasses2(t *testing.T) {

}
<commit_msg>Finishing unit tests<commit_after>package cors

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

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
)

func init() {
	gin.SetMode(gin.TestMode)
}

func newTestRouter(config Config) *gin.Engine {
	router := gin.New()
	router.Use(New(config))
	router.GET("/", func(c *gin.Context) {
		c.String(200, "get")
	})
	router.POST("/", func(c *gin.Context) {
		c.String(200, "post")
	})
	router.PATCH("/", func(c *gin.Context) {
		c.String(200, "patch")
	})
	return router
}

func performRequest(r http.Handler, method, origin string) *httptest.ResponseRecorder {
	req, _ := http.NewRequest(method, "/", nil)
	if len(origin) > 0 {
		req.Header.Set("Origin", origin)
	}
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}

func TestBadConfig(t *testing.T) {
	assert.Panics(t, func() { New(Config{}) })
	assert.Panics(t, func() {
		New(Config{
			AllowAllOrigins: true,
			AllowedOrigins:  []string{"http://google.com"},
		})
	})
	assert.Panics(t, func() {
		New(Config{
			AllowAllOrigins: true,
			AllowOriginFunc: func(origin string) bool { return false },
		})
	})
	assert.Panics(t, func() {
		New(Config{
			AllowedOrigins: []string{"google.com"},
		})
	})
}

func TestNormalize(t *testing.T) {
	values := normalize([]string{
		"http-Access ", "Post", "POST", " poSt  ",
		"HTTP-Access", "",
	})
	assert.Equal(t, values, []string{"http-access", "post", ""})

	values = normalize(nil)
	assert.Nil(t, values)

	values = normalize([]string{})
	assert.Equal(t, values, []string{})
}

func TestGenerateNormalHeaders_AllowAllOrigins(t *testing.T) {
	header := generateNormalHeaders(Config{
		AllowAllOrigins: false,
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Origin"), "")
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 1)

	header = generateNormalHeaders(Config{
		AllowAllOrigins: true,
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Origin"), "*")
	assert.Equal(t, header.Get("Vary"), "")
	assert.Len(t, header, 1)
}

func TestGenerateNormalHeaders_AllowCredentials(t *testing.T) {
	header := generateNormalHeaders(Config{
		AllowCredentials: true,
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Credentials"), "true")
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 2)
}

func TestGenerateNormalHeaders_ExposedHeaders(t *testing.T) {
	header := generateNormalHeaders(Config{
		ExposedHeaders: []string{"X-user", "xPassword"},
	})
	assert.Equal(t, header.Get("Access-Control-Expose-Headers"), "x-user, xpassword")
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 2)
}

func TestGeneratePreflightHeaders(t *testing.T) {
	header := generatePreflightHeaders(Config{
		AllowAllOrigins: false,
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Origin"), "")
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 1)

	header = generateNormalHeaders(Config{
		AllowAllOrigins: true,
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Origin"), "*")
	assert.Equal(t, header.Get("Vary"), "")
	assert.Len(t, header, 1)
}

func TestGeneratePreflightHeaders_AllowCredentials(t *testing.T) {
	header := generatePreflightHeaders(Config{
		AllowCredentials: true,
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Credentials"), "true")
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 2)
}

func TestGeneratePreflightHeaders_AllowedMethods(t *testing.T) {
	header := generatePreflightHeaders(Config{
		AllowedMethods: []string{"GET ", "post", "PUT", " put  "},
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Methods"), "get, post, put")
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 2)
}

func TestGeneratePreflightHeaders_AllowedHeaders(t *testing.T) {
	header := generatePreflightHeaders(Config{
		AllowedHeaders: []string{"X-user", "Content-Type"},
	})
	assert.Equal(t, header.Get("Access-Control-Allow-Headers"), "x-user, content-type")
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 2)
}

func TestGeneratePreflightHeaders_MaxAge(t *testing.T) {
	header := generatePreflightHeaders(Config{
		MaxAge: 12 * time.Hour,
	})
	assert.Equal(t, header.Get("Access-Control-Max-Age"), "43200") // 12*60*60
	assert.Equal(t, header.Get("Vary"), "Origin")
	assert.Len(t, header, 2)
}

func TestValidateOrigin(t *testing.T) {
	cors := newCors(Config{
		AllowAllOrigins: true,
	})
	assert.True(t, cors.validateOrigin("http://google.com"))
	assert.True(t, cors.validateOrigin("https://google.com"))
	assert.True(t, cors.validateOrigin("example.com"))

	cors = newCors(Config{
		AllowedOrigins: []string{"https://google.com", "https://github.com"},
		AllowOriginFunc: func(origin string) bool {
			return (origin == "http://news.ycombinator.com")
		},
	})
	assert.False(t, cors.validateOrigin("http://google.com"))
	assert.True(t, cors.validateOrigin("https://google.com"))
	assert.True(t, cors.validateOrigin("https://github.com"))
	assert.True(t, cors.validateOrigin("http://news.ycombinator.com"))
	assert.False(t, cors.validateOrigin("http://example.com"))
	assert.False(t, cors.validateOrigin("google.com"))
}

func TestPassesAllowedOrigins(t *testing.T) {
	router := newTestRouter(Config{
		AllowedOrigins:   []string{"http://google.com"},
		AllowedMethods:   []string{" GeT ", "get", "post", "PUT  ", "Head", "POST"},
		AllowedHeaders:   []string{"Content-type", "timeStamp "},
		ExposedHeaders:   []string{"Data", "x-User"},
		AllowCredentials: true,
		MaxAge:           12 * time.Hour,
		AllowOriginFunc: func(origin string) bool {
			return origin == "http://github.com"
		},
	})

	// no CORS request, origin == ""
	w := performRequest(router, "GET", "")
	assert.Equal(t, w.Body.String(), "get")
	assert.Empty(t, w.Header().Get("Access-Control-Allow-Origin"))
	assert.Empty(t, w.Header().Get("Access-Control-Allow-Credentials"))
	assert.Empty(t, w.Header().Get("Access-Control-Expose-Headers"))

	// allowed CORS request
	w = performRequest(router, "GET", "http://google.com")
	assert.Equal(t, w.Body.String(), "get")
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Origin"), "http://google.com")
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Credentials"), "true")
	assert.Equal(t, w.Header().Get("Access-Control-Expose-Headers"), "data, x-user")

	// deny CORS request
	w = performRequest(router, "GET", "https://google.com")
	assert.Equal(t, w.Code, 403)
	assert.Empty(t, w.Header().Get("Access-Control-Allow-Origin"))
	assert.Empty(t, w.Header().Get("Access-Control-Allow-Credentials"))
	assert.Empty(t, w.Header().Get("Access-Control-Expose-Headers"))

	// allowed CORS prefligh request
	w = performRequest(router, "OPTIONS", "http://github.com")
	assert.Equal(t, w.Code, 200)
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Origin"), "http://github.com")
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Credentials"), "true")
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Methods"), "get, post, put, head")
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Headers"), "content-type, timestamp")
	assert.Equal(t, w.Header().Get("Access-Control-Max-Age"), "43200")

	// deny CORS prefligh request
	w = performRequest(router, "OPTIONS", "http://example.com")
	assert.Equal(t, w.Code, 403)
	assert.Empty(t, w.Header().Get("Access-Control-Allow-Origin"))
	assert.Empty(t, w.Header().Get("Access-Control-Allow-Credentials"))
	assert.Empty(t, w.Header().Get("Access-Control-Allow-Methods"))
	assert.Empty(t, w.Header().Get("Access-Control-Allow-Headers"))
	assert.Empty(t, w.Header().Get("Access-Control-Max-Age"))
}

func TestPassesAllowedAllOrigins(t *testing.T) {
	router := newTestRouter(Config{
		AllowAllOrigins:  true,
		AllowedMethods:   []string{" Patch ", "get", "post", "POST"},
		AllowedHeaders:   []string{"Content-type", "  testheader "},
		ExposedHeaders:   []string{"Data2", "x-User2"},
		AllowCredentials: false,
		MaxAge:           10 * time.Hour,
	})

	// no CORS request, origin == ""
	w := performRequest(router, "GET", "")
	assert.Equal(t, w.Body.String(), "get")
	assert.Empty(t, w.Header().Get("Access-Control-Allow-Origin"))
	assert.Empty(t, w.Header().Get("Access-Control-Allow-Credentials"))
	assert.Empty(t, w.Header().Get("Access-Control-Expose-Headers"))

	// allowed CORS request
	w = performRequest(router, "POST", "example.com")
	assert.Equal(t, w.Body.String(), "post")
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Origin"), "*")
	assert.Equal(t, w.Header().Get("Access-Control-Expose-Headers"), "data2, x-user2")
	assert.Empty(t, w.Header().Get("Access-Control-Allow-Credentials"))

	// allowed CORS prefligh request
	w = performRequest(router, "OPTIONS", "https://facebook.com")
	assert.Equal(t, w.Code, 200)
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Origin"), "*")
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Methods"), "patch, get, post")
	assert.Equal(t, w.Header().Get("Access-Control-Allow-Headers"), "content-type, testheader")
	assert.Equal(t, w.Header().Get("Access-Control-Max-Age"), "36000")
	assert.Empty(t, w.Header().Get("Access-Control-Allow-Credentials"))
}
<|endoftext|>


<commit_before>// Copyright 2014 Google Inc. All Rights Reserved.
//
// 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 info

import (
	"fmt"
	"sort"
	"time"
)

type CpuSpecMask struct {
	Data []uint64 `json:"data,omitempty"`
}

type CpuSpec struct {
	Limit    uint64      `json:"limit"`
	MaxLimit uint64      `json:"max_limit"`
	Mask     CpuSpecMask `json:"mask,omitempty"`
}

type MemorySpec struct {
	// The amount of memory requested. Default is unlimited (-1).
	// Units: bytes.
	Limit uint64 `json:"limit,omitempty"`

	// The amount of guaranteed memory.  Default is 0.
	// Units: bytes.
	Reservation uint64 `json:"reservation,omitempty"`

	// The amount of swap space requested. Default is unlimited (-1).
	// Units: bytes.
	SwapLimit uint64 `json:"swap_limit,omitempty"`
}

type ContainerSpec struct {
	Cpu    *CpuSpec    `json:"cpu,omitempty"`
	Memory *MemorySpec `json:"memory,omitempty"`
}

type ContainerInfo struct {
	// The absolute name of the container.
	Name string `json:"name"`

	// The direct subcontainers of the current container.
	Subcontainers []string `json:"subcontainers,omitempty"`

	// The isolation used in the container.
	Spec *ContainerSpec `json:"spec,omitempty"`

	// Historical statistics gathered from the container.
	Stats []*ContainerStats `json:"stats,omitempty"`

	StatsSummary *ContainerStatsSummary `json:"stats_summary,omitempty"`
}

func (self *ContainerInfo) StatsAfter(ref time.Time) []*ContainerStats {
	n := len(self.Stats) + 1
	for i, s := range self.Stats {
		if s.Timestamp.After(ref) {
			n = i
			break
		}
	}
	if n > len(self.Stats) {
		return nil
	}
	return self.Stats[n:]
}

func (self *ContainerInfo) StatsStartTime() time.Time {
	var ret time.Time
	for _, s := range self.Stats {
		if s.Timestamp.Before(ret) || ret.IsZero() {
			ret = s.Timestamp
		}
	}
	return ret
}

func (self *ContainerInfo) StatsEndTime() time.Time {
	var ret time.Time
	for i := len(self.Stats) - 1; i >= 0; i-- {
		s := self.Stats[i]
		if s.Timestamp.After(ret) {
			ret = s.Timestamp
		}
	}
	return ret
}

type CpuStats struct {
	Usage struct {
		// Number of nanoseconds of CPU time used by the container
		// since the beginning. This is a ccumulative
		// value, not an instantaneous value.
		Total uint64 `json:"total"`

		// Per CPU/core usage of the container.
		// Unit: nanoseconds.
		PerCpu []uint64 `json:"per_cpu,omitempty"`

		// How much time was spent in user space since beginning.
		// Unit: nanoseconds
		User uint64 `json:"user"`

		// How much time was spent in kernel space since beginning.
		// Unit: nanoseconds
		System uint64 `json:"system"`
	} `json:"usage"`
	Load int32 `json:"load"`
}

type MemoryStats struct {
	// Memory limit, equivalent to "limit" in MemorySpec.
	// Units: Bytes.
	Limit uint64 `json:"limit,omitempty"`

	// Usage statistics.

	// Current memory usage, this includes all memory regardless of when it was
	// accessed.
	// Units: Bytes.
	Usage uint64 `json:"usage,omitempty"`

	// The amount of working set memory, this includes recently accessed memory,
	// dirty memory, and kernel memmory. Working set is <= "usage".
	// Units: Bytes.
	WorkingSet uint64 `json:"working_set,omitempty"`

	ContainerData    MemoryStatsMemoryData `json:"container_data,omitempty"`
	HierarchicalData MemoryStatsMemoryData `json:"hierarchical_data,omitempty"`
}

type MemoryStatsMemoryData struct {
	Pgfault    uint64 `json:"pgfault,omitempty"`
	Pgmajfault uint64 `json:"pgmajfault,omitempty"`
}

type ContainerStats struct {
	// The time of this stat point.
	Timestamp time.Time    `json:"timestamp"`
	Cpu       *CpuStats    `json:"cpu,omitempty"`
	Memory    *MemoryStats `json:"memory,omitempty"`
}

type ContainerStatsSample struct {
	// Timetamp of the end of the sample period
	Timestamp time.Time `json:"timestamp"`
	// Duration of the sample period
	Duration time.Duration `json:"duration"`
	Cpu      struct {
		// number of nanoseconds of CPU time used by the container
		// within one second.
		Usage uint64 `json:"usage"`
	} `json:"cpu"`
	Memory struct {
		// Units: Bytes.
		Usage uint64 `json:"usage"`
	} `json:"memory"`
}

// This is not exported.
// Use FillPercentile to calculate percentiles
type percentile struct {
	Percentage int    `json:"percentage"`
	Value      uint64 `json:"value"`
}

type ContainerStatsSummary struct {
	// TODO(dengnan): More things?
	MaxMemoryUsage         uint64                  `json:"max_memory_usage,omitempty"`
	AvgMemoryUsage         uint64                  `json:"avg_memory_usage,omitempty"`
	Samples                []*ContainerStatsSample `json:"samples,omitempty"`
	MemoryUsagePercentiles []percentile            `json:"memory_usage_percentiles,omitempty"`
	CpuUsagePercentiles    []percentile            `json:"cpu_usage_percentiles,omitempty"`
}

// Each sample needs two stats because the cpu usage in ContainerStats is
// cumulative.
// prev should be an earlier observation than current.
// This method is not thread/goroutine safe.
func NewSample(prev, current *ContainerStats) (*ContainerStatsSample, error) {
	if prev == nil || current == nil {
		return nil, fmt.Errorf("empty stats")
	}
	// Ignore this sample if it is incomplete
	if prev.Cpu == nil || prev.Memory == nil || current.Cpu == nil || current.Memory == nil {
		return nil, fmt.Errorf("incomplete stats")
	}
	// prev must be an early observation
	if !current.Timestamp.After(prev.Timestamp) {
		return nil, fmt.Errorf("wrong stats order")
	}
	sample := new(ContainerStatsSample)
	// Caculate the diff to get the CPU usage within the time interval.
	sample.Cpu.Usage = current.Cpu.Usage.Total - prev.Cpu.Usage.Total
	// Memory usage is current memory usage
	sample.Memory.Usage = current.Memory.Usage
	sample.Timestamp = current.Timestamp
	sample.Duration = current.Timestamp.Sub(prev.Timestamp)

	return sample, nil
}

type uint64Slice []uint64

func (self uint64Slice) Len() int {
	return len(self)
}

func (self uint64Slice) Less(i, j int) bool {
	return self[i] < self[j]
}

func (self uint64Slice) Swap(i, j int) {
	self[i], self[j] = self[j], self[i]
}

func (self uint64Slice) Percentiles(ps ...int) []uint64 {
	if len(self) == 0 {
		return nil
	}
	ret := make([]uint64, 0, len(ps))
	sort.Sort(self)
	for _, p := range ps {
		idx := (float64(p) / 100.0) * float64(len(self)+1)
		if idx > float64(len(self)-1) {
			ret = append(ret, self[len(self)-1])
		} else {
			ret = append(ret, self[int(idx)])
		}
	}
	return ret
}

// len(bs) <= len(as)
func intZipuint64(as []int, bs []uint64) []percentile {
	if len(bs) == 0 {
		return nil
	}
	ret := make([]percentile, len(bs))
	for i, b := range bs {
		a := as[i]
		ret[i] = percentile{
			Percentage: a,
			Value:      b,
		}
	}
	return ret
}

func (self *ContainerStatsSummary) FillPercentiles(cpuPercentages, memoryPercentages []int) {
	if len(self.Samples) == 0 {
		return
	}
	cpuUsages := make([]uint64, 0, len(self.Samples))
	memUsages := make([]uint64, 0, len(self.Samples))

	for _, sample := range self.Samples {
		if sample == nil {
			continue
		}
		cpuUsages = append(cpuUsages, sample.Cpu.Usage)
		memUsages = append(memUsages, sample.Memory.Usage)
	}

	cpuPercentiles := uint64Slice(cpuUsages).Percentiles(cpuPercentages...)
	memPercentiles := uint64Slice(memUsages).Percentiles(memoryPercentages...)
	self.CpuUsagePercentiles = intZipuint64(cpuPercentages, cpuPercentiles)
	self.MemoryUsagePercentiles = intZipuint64(memoryPercentages, memPercentiles)
}
<commit_msg>ignore sample if it leads cpu utilization < 0<commit_after>// Copyright 2014 Google Inc. All Rights Reserved.
//
// 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 info

import (
	"fmt"
	"sort"
	"time"
)

type CpuSpecMask struct {
	Data []uint64 `json:"data,omitempty"`
}

type CpuSpec struct {
	Limit    uint64      `json:"limit"`
	MaxLimit uint64      `json:"max_limit"`
	Mask     CpuSpecMask `json:"mask,omitempty"`
}

type MemorySpec struct {
	// The amount of memory requested. Default is unlimited (-1).
	// Units: bytes.
	Limit uint64 `json:"limit,omitempty"`

	// The amount of guaranteed memory.  Default is 0.
	// Units: bytes.
	Reservation uint64 `json:"reservation,omitempty"`

	// The amount of swap space requested. Default is unlimited (-1).
	// Units: bytes.
	SwapLimit uint64 `json:"swap_limit,omitempty"`
}

type ContainerSpec struct {
	Cpu    *CpuSpec    `json:"cpu,omitempty"`
	Memory *MemorySpec `json:"memory,omitempty"`
}

type ContainerInfo struct {
	// The absolute name of the container.
	Name string `json:"name"`

	// The direct subcontainers of the current container.
	Subcontainers []string `json:"subcontainers,omitempty"`

	// The isolation used in the container.
	Spec *ContainerSpec `json:"spec,omitempty"`

	// Historical statistics gathered from the container.
	Stats []*ContainerStats `json:"stats,omitempty"`

	StatsSummary *ContainerStatsSummary `json:"stats_summary,omitempty"`
}

func (self *ContainerInfo) StatsAfter(ref time.Time) []*ContainerStats {
	n := len(self.Stats) + 1
	for i, s := range self.Stats {
		if s.Timestamp.After(ref) {
			n = i
			break
		}
	}
	if n > len(self.Stats) {
		return nil
	}
	return self.Stats[n:]
}

func (self *ContainerInfo) StatsStartTime() time.Time {
	var ret time.Time
	for _, s := range self.Stats {
		if s.Timestamp.Before(ret) || ret.IsZero() {
			ret = s.Timestamp
		}
	}
	return ret
}

func (self *ContainerInfo) StatsEndTime() time.Time {
	var ret time.Time
	for i := len(self.Stats) - 1; i >= 0; i-- {
		s := self.Stats[i]
		if s.Timestamp.After(ret) {
			ret = s.Timestamp
		}
	}
	return ret
}

type CpuStats struct {
	Usage struct {
		// Number of nanoseconds of CPU time used by the container
		// since the beginning. This is a ccumulative
		// value, not an instantaneous value.
		Total uint64 `json:"total"`

		// Per CPU/core usage of the container.
		// Unit: nanoseconds.
		PerCpu []uint64 `json:"per_cpu,omitempty"`

		// How much time was spent in user space since beginning.
		// Unit: nanoseconds
		User uint64 `json:"user"`

		// How much time was spent in kernel space since beginning.
		// Unit: nanoseconds
		System uint64 `json:"system"`
	} `json:"usage"`
	Load int32 `json:"load"`
}

type MemoryStats struct {
	// Memory limit, equivalent to "limit" in MemorySpec.
	// Units: Bytes.
	Limit uint64 `json:"limit,omitempty"`

	// Usage statistics.

	// Current memory usage, this includes all memory regardless of when it was
	// accessed.
	// Units: Bytes.
	Usage uint64 `json:"usage,omitempty"`

	// The amount of working set memory, this includes recently accessed memory,
	// dirty memory, and kernel memmory. Working set is <= "usage".
	// Units: Bytes.
	WorkingSet uint64 `json:"working_set,omitempty"`

	ContainerData    MemoryStatsMemoryData `json:"container_data,omitempty"`
	HierarchicalData MemoryStatsMemoryData `json:"hierarchical_data,omitempty"`
}

type MemoryStatsMemoryData struct {
	Pgfault    uint64 `json:"pgfault,omitempty"`
	Pgmajfault uint64 `json:"pgmajfault,omitempty"`
}

type ContainerStats struct {
	// The time of this stat point.
	Timestamp time.Time    `json:"timestamp"`
	Cpu       *CpuStats    `json:"cpu,omitempty"`
	Memory    *MemoryStats `json:"memory,omitempty"`
}

type ContainerStatsSample struct {
	// Timetamp of the end of the sample period
	Timestamp time.Time `json:"timestamp"`
	// Duration of the sample period
	Duration time.Duration `json:"duration"`
	Cpu      struct {
		// number of nanoseconds of CPU time used by the container
		// within one second.
		Usage uint64 `json:"usage"`
	} `json:"cpu"`
	Memory struct {
		// Units: Bytes.
		Usage uint64 `json:"usage"`
	} `json:"memory"`
}

// This is not exported.
// Use FillPercentile to calculate percentiles
type percentile struct {
	Percentage int    `json:"percentage"`
	Value      uint64 `json:"value"`
}

type ContainerStatsSummary struct {
	// TODO(dengnan): More things?
	MaxMemoryUsage         uint64                  `json:"max_memory_usage,omitempty"`
	AvgMemoryUsage         uint64                  `json:"avg_memory_usage,omitempty"`
	Samples                []*ContainerStatsSample `json:"samples,omitempty"`
	MemoryUsagePercentiles []percentile            `json:"memory_usage_percentiles,omitempty"`
	CpuUsagePercentiles    []percentile            `json:"cpu_usage_percentiles,omitempty"`
}

// Each sample needs two stats because the cpu usage in ContainerStats is
// cumulative.
// prev should be an earlier observation than current.
// This method is not thread/goroutine safe.
func NewSample(prev, current *ContainerStats) (*ContainerStatsSample, error) {
	if prev == nil || current == nil {
		return nil, fmt.Errorf("empty stats")
	}
	// Ignore this sample if it is incomplete
	if prev.Cpu == nil || prev.Memory == nil || current.Cpu == nil || current.Memory == nil {
		return nil, fmt.Errorf("incomplete stats")
	}
	// prev must be an early observation
	if !current.Timestamp.After(prev.Timestamp) {
		return nil, fmt.Errorf("wrong stats order")
	}
	// This data is invalid.
	if current.Cpu.Usage.Total < prev.Cpu.Usage.Total {
		return nil, fmt.Errorf("current CPU usage is less than prev CPU usage (cumulative).")
	}
	sample := new(ContainerStatsSample)
	// Caculate the diff to get the CPU usage within the time interval.
	sample.Cpu.Usage = current.Cpu.Usage.Total - prev.Cpu.Usage.Total
	// Memory usage is current memory usage
	sample.Memory.Usage = current.Memory.Usage
	sample.Timestamp = current.Timestamp
	sample.Duration = current.Timestamp.Sub(prev.Timestamp)

	return sample, nil
}

type uint64Slice []uint64

func (self uint64Slice) Len() int {
	return len(self)
}

func (self uint64Slice) Less(i, j int) bool {
	return self[i] < self[j]
}

func (self uint64Slice) Swap(i, j int) {
	self[i], self[j] = self[j], self[i]
}

func (self uint64Slice) Percentiles(ps ...int) []uint64 {
	if len(self) == 0 {
		return nil
	}
	ret := make([]uint64, 0, len(ps))
	sort.Sort(self)
	for _, p := range ps {
		idx := (float64(p) / 100.0) * float64(len(self)+1)
		if idx > float64(len(self)-1) {
			ret = append(ret, self[len(self)-1])
		} else {
			ret = append(ret, self[int(idx)])
		}
	}
	return ret
}

// len(bs) <= len(as)
func intZipuint64(as []int, bs []uint64) []percentile {
	if len(bs) == 0 {
		return nil
	}
	ret := make([]percentile, len(bs))
	for i, b := range bs {
		a := as[i]
		ret[i] = percentile{
			Percentage: a,
			Value:      b,
		}
	}
	return ret
}

func (self *ContainerStatsSummary) FillPercentiles(cpuPercentages, memoryPercentages []int) {
	if len(self.Samples) == 0 {
		return
	}
	cpuUsages := make([]uint64, 0, len(self.Samples))
	memUsages := make([]uint64, 0, len(self.Samples))

	for _, sample := range self.Samples {
		if sample == nil {
			continue
		}
		cpuUsages = append(cpuUsages, sample.Cpu.Usage)
		memUsages = append(memUsages, sample.Memory.Usage)
	}

	cpuPercentiles := uint64Slice(cpuUsages).Percentiles(cpuPercentages...)
	memPercentiles := uint64Slice(memUsages).Percentiles(memoryPercentages...)
	self.CpuUsagePercentiles = intZipuint64(cpuPercentages, cpuPercentiles)
	self.MemoryUsagePercentiles = intZipuint64(memoryPercentages, memPercentiles)
}
<|endoftext|>