<commit_before>package order

import (
	"errors"
	"net/url"
	"strconv"

	stripe "github.com/stripe/stripe-go"
)

// Client is used to invoke /orders APIs.
type Client struct {
	B   stripe.Backend
	Key string
}

// New POSTs a new order.
// For more details see https://stripe.com/docs/api#create_order.
func New(params *stripe.OrderParams) (*stripe.Order, error) {
	return getC().New(params)
}

// New POSTs a new order.
// For more details see https://stripe.com/docs/api#create_order.
func (c Client) New(params *stripe.OrderParams) (*stripe.Order, error) {
	var body *url.Values
	var commonParams *stripe.Params

	if params != nil {
		body = &url.Values{}
		commonParams = &params.Params
		// Required fields
		body.Add("currency", string(params.Currency))

		if params.Customer != "" {
			body.Add("customer", params.Customer)
		}

		if params.Email != "" {
			body.Add("email", params.Email)
		}

		if params.Email != "" {
			body.Add("email", params.Email)
		}

		if len(params.Items) > 0 {
			for _, item := range params.Items {
				body.Add("items[][description]", item.Description)
				body.Add("items[][type]", string(item.Type))
				body.Add("items[][amount]", strconv.FormatInt(item.Amount, 10))
				if item.Currency != "" {
					body.Add("items[][currency]", string(item.Currency))
				}
				if item.Parent != "" {
					body.Add("items[][parent]", item.Parent)
				}
				if item.Quantity != nil {
					body.Add("items[][quantity]", strconv.FormatInt(*item.Quantity, 10))
				}
			}
		}

		body.Add("shipping[address][line1]", params.Shipping.Address.Line1)
		if params.Shipping.Address.Line2 != "" {
			body.Add("shipping[address][line2]", params.Shipping.Address.Line2)
		}
		if params.Shipping.Address.City != "" {
			body.Add("shipping[address][city]", params.Shipping.Address.City)
		}
		if params.Shipping.Address.Country != "" {
			body.Add("shipping[address][country]", params.Shipping.Address.Country)
		}
		if params.Shipping.Address.PostalCode != "" {
			body.Add("shipping[address][postal_code]", params.Shipping.Address.PostalCode)
		}
		if params.Shipping.Address.State != "" {
			body.Add("shipping[address][state]", params.Shipping.Address.State)
		}

		if params.Shipping.Name != "" {
			body.Add("shipping[name]", params.Shipping.Name)
		}
		if params.Shipping.Phone != "" {
			body.Add("shipping[phone]", params.Shipping.Phone)
		}

		params.AppendTo(body)
	}

	p := &stripe.Order{}
	err := c.B.Call("POST", "/orders", c.Key, body, commonParams, p)

	return p, err
}

// Update updates an order's properties.
// For more details see https://stripe.com/docs/api#update_order.
func Update(id string, params *stripe.OrderUpdateParams) (*stripe.Order, error) {
	return getC().Update(id, params)
}

// Update updates an order's properties.
// For more details see https://stripe.com/docs/api#update_order.
func (c Client) Update(id string, params *stripe.OrderUpdateParams) (*stripe.Order, error) {
	var body *url.Values
	var commonParams *stripe.Params

	if params != nil {
		body = &url.Values{}

		if params.Coupon != "" {
			body.Add("coupon", params.Coupon)
		}

		if params.SelectedShippingMethod != "" {
			body.Add("selected_shipping_method", params.SelectedShippingMethod)
		}

		if params.Status != "" {
			body.Add("status", string(params.Status))
		}

		params.AppendTo(body)
	}

	o := &stripe.Order{}
	err := c.B.Call("POST", "/orders/"+id, c.Key, body, commonParams, o)

	return o, err
}

// Pay pays an order
// For more details see https://stripe.com/docs/api#pay_order.
func Pay(id string, params *stripe.OrderPayParams) (*stripe.Order, error) {
	return getC().Pay(id, params)
}

// Pay pays an order
// For more details see https://stripe.com/docs/api#pay_order.
func (c Client) Pay(id string, params *stripe.OrderPayParams) (*stripe.Order, error) {
	var body *url.Values
	var commonParams *stripe.Params

	if params != nil {
		body = &url.Values{}

		if params.Source == nil && len(params.Customer) == 0 {
			err := errors.New("Invalid order pay params: either customer or a source must be set")
			return nil, err
		}

		// We can't use `AppendDetails` since that nests under `card`.
		if params.Source != nil {
			if len(params.Source.Token) > 0 {
				body.Add("source", params.Source.Token)
			} else if params.Source.Card != nil {
				c := params.Source.Card

				body.Add("source[object]", "card")
				body.Add("source[number]", c.Number)
				body.Add("source[exp_month]", c.Month)
				body.Add("source[exp_year]", c.Year)

				if len(c.CVC) > 0 {
					body.Add("source[cvc]", c.CVC)
				}

				body.Add("source[name]", c.Name)

				if len(c.Address1) > 0 {
					body.Add("source[address_line1]", c.Address1)
				}

				if len(c.Address2) > 0 {
					body.Add("source[address_line2]", c.Address2)
				}
				if len(c.City) > 0 {
					body.Add("source[address_city]", c.City)
				}

				if len(c.State) > 0 {
					body.Add("source[address_state]", c.State)
				}
				if len(c.Zip) > 0 {
					body.Add("source[address_zip]", c.Zip)
				}
				if len(c.Country) > 0 {
					body.Add("source[address_country]", c.Country)
				}
			}
		}

		if len(params.Customer) > 0 {
			body.Add("customer", params.Customer)
		}

		if params.ApplicationFee > 0 {
			body.Add("application_fee", strconv.FormatInt(params.ApplicationFee, 10))
		}

		if params.Email != "" {
			body.Add("email", params.Email)
		}

		params.AppendTo(body)
	}

	o := &stripe.Order{}
	err := c.B.Call("POST", "/orders/"+id+"/pay", c.Key, body, commonParams, o)

	return o, err
}

// Get returns the details of an order
// For more details see https://stripe.com/docs/api#retrieve_order.
func Get(id string) (*stripe.Order, error) {
	return getC().Get(id)
}

func (c Client) Get(id string) (*stripe.Order, error) {
	order := &stripe.Order{}
	err := c.B.Call("GET", "/orders/"+id, c.Key, nil, nil, order)

	return order, err
}

// List returns a list of orders.
// For more details see https://stripe.com/docs/api#list_orders
func List(params *stripe.OrderListParams) *Iter {
	return getC().List(params)
}

func (c Client) List(params *stripe.OrderListParams) *Iter {
	type orderList struct {
		stripe.ListMeta
		Values []*stripe.Order `json:"data"`
	}

	var body *url.Values
	var lp *stripe.ListParams

	if params != nil {
		body = &url.Values{}

		for _, id := range params.IDs {
			params.Filters.AddFilter("ids[]", "", id)
		}

		if params.Status != "" {
			params.Filters.AddFilter("status", "", string(params.Status))
		}

		params.AppendTo(body)
		lp = &params.ListParams
	}

	return &Iter{stripe.GetIter(lp, body, func(b url.Values) ([]interface{}, stripe.ListMeta, error) {
		list := &orderList{}
		err := c.B.Call("GET", "/orders", c.Key, &b, nil, list)

		ret := make([]interface{}, len(list.Values))
		for i, v := range list.Values {
			ret[i] = v
		}

		return ret, list.ListMeta, err
	})}
}

// Iter is an iterator for lists of Orders.
// The embedded Iter carries methods with it;
// see its documentation for details.
type Iter struct {
	*stripe.Iter
}

// Order returns the most recent Order
// visited by a call to Next.
func (i *Iter) Order() *stripe.Order {
	return i.Current().(*stripe.Order)
}

func getC() Client {
	return Client{stripe.GetBackend(stripe.APIBackend), stripe.Key}
}
<commit_msg>Fix Pay() commonParams issue<commit_after>package order

import (
	"errors"
	"net/url"
	"strconv"

	stripe "github.com/stripe/stripe-go"
)

// Client is used to invoke /orders APIs.
type Client struct {
	B   stripe.Backend
	Key string
}

// New POSTs a new order.
// For more details see https://stripe.com/docs/api#create_order.
func New(params *stripe.OrderParams) (*stripe.Order, error) {
	return getC().New(params)
}

// New POSTs a new order.
// For more details see https://stripe.com/docs/api#create_order.
func (c Client) New(params *stripe.OrderParams) (*stripe.Order, error) {
	var body *url.Values
	var commonParams *stripe.Params

	if params != nil {
		body = &url.Values{}
		commonParams = &params.Params
		// Required fields
		body.Add("currency", string(params.Currency))

		if params.Customer != "" {
			body.Add("customer", params.Customer)
		}

		if params.Email != "" {
			body.Add("email", params.Email)
		}

		if params.Email != "" {
			body.Add("email", params.Email)
		}

		if len(params.Items) > 0 {
			for _, item := range params.Items {
				body.Add("items[][description]", item.Description)
				body.Add("items[][type]", string(item.Type))
				body.Add("items[][amount]", strconv.FormatInt(item.Amount, 10))
				if item.Currency != "" {
					body.Add("items[][currency]", string(item.Currency))
				}
				if item.Parent != "" {
					body.Add("items[][parent]", item.Parent)
				}
				if item.Quantity != nil {
					body.Add("items[][quantity]", strconv.FormatInt(*item.Quantity, 10))
				}
			}
		}

		body.Add("shipping[address][line1]", params.Shipping.Address.Line1)
		if params.Shipping.Address.Line2 != "" {
			body.Add("shipping[address][line2]", params.Shipping.Address.Line2)
		}
		if params.Shipping.Address.City != "" {
			body.Add("shipping[address][city]", params.Shipping.Address.City)
		}
		if params.Shipping.Address.Country != "" {
			body.Add("shipping[address][country]", params.Shipping.Address.Country)
		}
		if params.Shipping.Address.PostalCode != "" {
			body.Add("shipping[address][postal_code]", params.Shipping.Address.PostalCode)
		}
		if params.Shipping.Address.State != "" {
			body.Add("shipping[address][state]", params.Shipping.Address.State)
		}

		if params.Shipping.Name != "" {
			body.Add("shipping[name]", params.Shipping.Name)
		}
		if params.Shipping.Phone != "" {
			body.Add("shipping[phone]", params.Shipping.Phone)
		}

		params.AppendTo(body)
	}

	p := &stripe.Order{}
	err := c.B.Call("POST", "/orders", c.Key, body, commonParams, p)

	return p, err
}

// Update updates an order's properties.
// For more details see https://stripe.com/docs/api#update_order.
func Update(id string, params *stripe.OrderUpdateParams) (*stripe.Order, error) {
	return getC().Update(id, params)
}

// Update updates an order's properties.
// For more details see https://stripe.com/docs/api#update_order.
func (c Client) Update(id string, params *stripe.OrderUpdateParams) (*stripe.Order, error) {
	var body *url.Values
	var commonParams *stripe.Params

	if params != nil {
		body = &url.Values{}

		if params.Coupon != "" {
			body.Add("coupon", params.Coupon)
		}

		if params.SelectedShippingMethod != "" {
			body.Add("selected_shipping_method", params.SelectedShippingMethod)
		}

		if params.Status != "" {
			body.Add("status", string(params.Status))
		}

		params.AppendTo(body)
	}

	o := &stripe.Order{}
	err := c.B.Call("POST", "/orders/"+id, c.Key, body, commonParams, o)

	return o, err
}

// Pay pays an order
// For more details see https://stripe.com/docs/api#pay_order.
func Pay(id string, params *stripe.OrderPayParams) (*stripe.Order, error) {
	return getC().Pay(id, params)
}

// Pay pays an order
// For more details see https://stripe.com/docs/api#pay_order.
func (c Client) Pay(id string, params *stripe.OrderPayParams) (*stripe.Order, error) {
	var body *url.Values
	var commonParams *stripe.Params

	if params != nil {
		body = &url.Values{}
		commonParams = &params.Params
		if params.Source == nil && len(params.Customer) == 0 {
			err := errors.New("Invalid order pay params: either customer or a source must be set")
			return nil, err
		}
		// We can't use `AppendDetails` since that nests under `card`.
		if params.Source != nil {
			if len(params.Source.Token) > 0 {
				body.Add("source", params.Source.Token)
			} else if params.Source.Card != nil {
				c := params.Source.Card

				body.Add("source[object]", "card")
				body.Add("source[number]", c.Number)
				body.Add("source[exp_month]", c.Month)
				body.Add("source[exp_year]", c.Year)

				if len(c.CVC) > 0 {
					body.Add("source[cvc]", c.CVC)
				}

				body.Add("source[name]", c.Name)

				if len(c.Address1) > 0 {
					body.Add("source[address_line1]", c.Address1)
				}

				if len(c.Address2) > 0 {
					body.Add("source[address_line2]", c.Address2)
				}
				if len(c.City) > 0 {
					body.Add("source[address_city]", c.City)
				}

				if len(c.State) > 0 {
					body.Add("source[address_state]", c.State)
				}
				if len(c.Zip) > 0 {
					body.Add("source[address_zip]", c.Zip)
				}
				if len(c.Country) > 0 {
					body.Add("source[address_country]", c.Country)
				}
			}
		}

		if len(params.Customer) > 0 {
			body.Add("customer", params.Customer)
		}

		if params.ApplicationFee > 0 {
			body.Add("application_fee", strconv.FormatInt(params.ApplicationFee, 10))
		}

		if params.Email != "" {
			body.Add("email", params.Email)
		}

		params.AppendTo(body)
	}

	o := &stripe.Order{}
	err := c.B.Call("POST", "/orders/"+id+"/pay", c.Key, body, commonParams, o)

	return o, err
}

// Get returns the details of an order
// For more details see https://stripe.com/docs/api#retrieve_order.
func Get(id string) (*stripe.Order, error) {
	return getC().Get(id)
}

func (c Client) Get(id string) (*stripe.Order, error) {
	order := &stripe.Order{}
	err := c.B.Call("GET", "/orders/"+id, c.Key, nil, nil, order)

	return order, err
}

// List returns a list of orders.
// For more details see https://stripe.com/docs/api#list_orders
func List(params *stripe.OrderListParams) *Iter {
	return getC().List(params)
}

func (c Client) List(params *stripe.OrderListParams) *Iter {
	type orderList struct {
		stripe.ListMeta
		Values []*stripe.Order `json:"data"`
	}

	var body *url.Values
	var lp *stripe.ListParams

	if params != nil {
		body = &url.Values{}

		for _, id := range params.IDs {
			params.Filters.AddFilter("ids[]", "", id)
		}

		if params.Status != "" {
			params.Filters.AddFilter("status", "", string(params.Status))
		}

		params.AppendTo(body)
		lp = &params.ListParams
	}

	return &Iter{stripe.GetIter(lp, body, func(b url.Values) ([]interface{}, stripe.ListMeta, error) {
		list := &orderList{}
		err := c.B.Call("GET", "/orders", c.Key, &b, nil, list)

		ret := make([]interface{}, len(list.Values))
		for i, v := range list.Values {
			ret[i] = v
		}

		return ret, list.ListMeta, err
	})}
}

// Iter is an iterator for lists of Orders.
// The embedded Iter carries methods with it;
// see its documentation for details.
type Iter struct {
	*stripe.Iter
}

// Order returns the most recent Order
// visited by a call to Next.
func (i *Iter) Order() *stripe.Order {
	return i.Current().(*stripe.Order)
}

func getC() Client {
	return Client{stripe.GetBackend(stripe.APIBackend), stripe.Key}
}
<|endoftext|>


<commit_before>package driver

import (
	"time"
)

// HTTP のキャッシュみたいなことができるように。
// 取得操作の場合、対象の更新日時がキャッシュの日時以降、または、ダイジェストがキャッシュと異なる場合のみ、現在の対象が返る。
//
// 対象が            返り値 返りスタンプ
// 無い              nil    nil
// キャッシュと同じ  nil    あり
// キャッシュと違う  あり   あり
//
// 更新操作の場合、対象がキャッシュと等しい場合のみ操作が行われ、新しいスタンプが返る。

// キャッシュの情報。
type Stamp struct {
	Date     time.Time `json:"date"                      bson:"date"`                      // 元データの更新日時。
	ExpiDate time.Time `json:"expiration_date,omitempty" bson:"expiration_date,omitempty"` // 有効期限。
	Digest   string    `json:"digest"                    bson:"digest"`                    // ハッシュ値とか。
}

type DatedIdProviderLister interface {
	StampedIdProviders(caStmp *Stamp) ([]*IdProvider, *Stamp, error)
}

type JsBackend interface {
	// オブジェクトのソースを取得する。
	StampedObject(dir, objName string, caStmp *Stamp) (*Object, *Stamp, error)
}

type JsBackendRegistry interface {
	JsRegistry
	JsBackend
}
<commit_msg>キャッシュ用タイムスタンプの処理を整理<commit_after>package driver

import (
	"github.com/realglobe-Inc/go-lib-rg/erro"
	"net/http"
	"time"
)

// HTTP のキャッシュみたいなことができるように。
// 取得操作の場合、対象の更新日時がキャッシュの日時以降、または、ダイジェストがキャッシュと異なる場合のみ、現在の対象が返る。
//
// 対象が            返り値 返りスタンプ
// 無い              nil    nil
// キャッシュと同じ  nil    あり
// キャッシュと違う  あり   あり
//
// 更新操作の場合、対象がキャッシュと等しい場合のみ操作が行われ、新しいスタンプが返る。

// キャッシュの情報。
type Stamp struct {
	Date     time.Time `json:"date"                      bson:"date"`                      // 元データの更新日時。
	ExpiDate time.Time `json:"expiration_date,omitempty" bson:"expiration_date,omitempty"` // 有効期限。
	Digest   string    `json:"digest"                    bson:"digest"`                    // ハッシュ値とか。
}

func ParseStampFromRequestHeader(h http.Header) (*Stamp, error) {
	dig, dateStr := h.Get("If-None-Match"), h.Get("If-Modified-Since")
	if dig == "" && dateStr == "" {
		return nil, nil
	}
	stmp := &Stamp{}
	stmp.Digest = dig
	if dateStr != "" {
		date, err := time.Parse(time.RFC1123, dateStr)
		if err != nil {
			return nil, erro.Wrap(err)
		}
		stmp.Date = date
	}
	return stmp, nil
}
func WriteStampToResponseHeader(stmp *Stamp, h http.Header) {
	h.Set("Last-Modified", stmp.Date.Format(time.RFC1123))
	h.Set("Expires", stmp.ExpiDate.Format(time.RFC1123))
	h.Set("ETag", stmp.Digest)
}

type DatedIdProviderLister interface {
	StampedIdProviders(caStmp *Stamp) ([]*IdProvider, *Stamp, error)
}

type JsBackend interface {
	// オブジェクトのソースを取得する。
	StampedObject(dir, objName string, caStmp *Stamp) (*Object, *Stamp, error)
}

type JsBackendRegistry interface {
	JsRegistry
	JsBackend
}
<|endoftext|>


<commit_before>package rproxy

import (
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
)

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

func CreateReverseProxy(target *url.URL) *httputil.ReverseProxy {
	targetQuery := target.RawQuery
	director := func(req *http.Request) {
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		req.URL.Path = singleJoiningSlash(target.Path, req.URL.Path)
		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}
	}
	return &httputil.ReverseProxy{Director: director}
}

type BackendNode struct {
	targetUrl string
	server    *httputil.ReverseProxy
}

type Backend struct {
	// map on targetUrl
	nodes []*BackendNode
}

func NewBackend() *Backend {
	b := &Backend{}
	b.nodes = make([]*BackendNode, 0)
	return b
}

func (b *Backend) deleteNode(targetUrl string) error {
	found := -1
	for i, n := range b.nodes {
		if n.targetUrl == targetUrl {
			found = i
			break
		}
	}

	if found > -1 {
		if found == len(b.nodes)-1 {
			b.nodes = append(b.nodes[:found])
		} else {
			b.nodes = append(b.nodes[:found], b.nodes[found+1:]...)
		}
		return nil
	}

	return fmt.Errorf("Could not found node %s", targetUrl)
}

func (b *Backend) addNode(targetUrl string) error {
	target, err := url.Parse(targetUrl)
	if err != nil {
		return err
	}
	// find if there is a node for this targetUrl
	n := b.getNode(targetUrl)
	if n == nil {
		n = &BackendNode{}
	}
	n.targetUrl = targetUrl
	n.server = CreateReverseProxy(target)
	b.nodes = append(b.nodes, n)
	return nil
}

func (b *Backend) getNode(targetUrl string) *BackendNode {
	for _, n := range b.nodes {
		if n.targetUrl == targetUrl {
			return n
		}
	}
	return nil
}

// nextNode returns a node candidate for serving the request
func (b *Backend) nextNode() *BackendNode {
	// Pickup one backend proxy
	// TODO round rubin
	for _, p := range b.nodes {
		return p
	}

	return nil
}

func (b *Backend) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	node := b.nextNode()
	if node != nil {
		// TODO next node if this node fails to serve
		node.server.ServeHTTP(rw, req)
	} else {
		http.Error(rw, http.StatusText(http.StatusServiceUnavailable), http.StatusServiceUnavailable)
	}

}
<commit_msg>We now rotate<commit_after>package rproxy

import (
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
)

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

func CreateReverseProxy(target *url.URL) *httputil.ReverseProxy {
	targetQuery := target.RawQuery
	director := func(req *http.Request) {
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		req.URL.Path = singleJoiningSlash(target.Path, req.URL.Path)
		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}
	}
	return &httputil.ReverseProxy{Director: director}
}

type BackendNode struct {
	targetUrl string
	server    *httputil.ReverseProxy
}

type Backend struct {
	// map on targetUrl
	nodes []*BackendNode
	next  int
}

func NewBackend() *Backend {
	b := &Backend{}
	b.nodes = make([]*BackendNode, 0)
	return b
}

func (b *Backend) deleteNode(targetUrl string) error {
	found := -1
	for i, n := range b.nodes {
		if n.targetUrl == targetUrl {
			found = i
			break
		}
	}

	if found > -1 {
		if found == len(b.nodes)-1 {
			b.nodes = append(b.nodes[:found])
		} else {
			b.nodes = append(b.nodes[:found], b.nodes[found+1:]...)
		}
		return nil
	}

	return fmt.Errorf("Could not found node %s", targetUrl)
}

func (b *Backend) addNode(targetUrl string) error {
	target, err := url.Parse(targetUrl)
	if err != nil {
		return err
	}
	// find if there is a node for this targetUrl
	n := b.getNode(targetUrl)
	if n == nil {
		n = &BackendNode{}
	}
	n.targetUrl = targetUrl
	n.server = CreateReverseProxy(target)
	b.nodes = append(b.nodes, n)
	return nil
}

func (b *Backend) getNode(targetUrl string) *BackendNode {
	for _, n := range b.nodes {
		if n.targetUrl == targetUrl {
			return n
		}
	}
	return nil
}

// nextNode returns a node candidate for serving the request
func (b *Backend) nextNode() *BackendNode {
	// Pickup one backend proxy
	if len(b.nodes) == 0 {
		return nil
	}
	p := b.nodes[b.next]
	b.next++
	if b.next >= len(b.nodes) {
		b.next = 0
	}

	return p
}

func (b *Backend) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	node := b.nextNode()
	if node != nil {
		// TODO next node if this node fails to serve
		node.server.ServeHTTP(rw, req)
	} else {
		http.Error(rw, http.StatusText(http.StatusServiceUnavailable), http.StatusServiceUnavailable)
	}

}
<|endoftext|>


<commit_before>package kapacitor

import (
	"errors"
	"time"

	"sync"
	"sync/atomic"

	"github.com/influxdata/kapacitor/edge"
	"github.com/influxdata/kapacitor/models"
	"github.com/influxdata/kapacitor/pipeline"
)

type BarrierNode struct {
	node
	b              *pipeline.BarrierNode
	barrierStopper map[models.GroupID]func()
}

// Create a new  BarrierNode, which emits a barrier if data traffic has been idle for the configured amount of time.
func newBarrierNode(et *ExecutingTask, n *pipeline.BarrierNode, d NodeDiagnostic) (*BarrierNode, error) {
	if n.Idle == 0 && n.Period == 0 {
		return nil, errors.New("barrier node must have either a non zero idle or a non zero period")
	}
	bn := &BarrierNode{
		node:           node{Node: n, et: et, diag: d},
		b:              n,
		barrierStopper: map[models.GroupID]func(){},
	}
	bn.node.runF = bn.runBarrierEmitter
	return bn, nil
}

func (n *BarrierNode) runBarrierEmitter([]byte) error {
	defer n.stopBarrierEmitter()
	consumer := edge.NewGroupedConsumer(n.ins[0], n)
	n.statMap.Set(statCardinalityGauge, consumer.CardinalityVar())
	return consumer.Consume()
}

func (n *BarrierNode) stopBarrierEmitter() {
	for _, stopF := range n.barrierStopper {
		stopF()
	}
}

func (n *BarrierNode) NewGroup(group edge.GroupInfo, first edge.PointMeta) (edge.Receiver, error) {
	r, stopF, err := n.newBarrier(group, first)
	if err != nil {
		return nil, err
	}
	n.barrierStopper[group.ID] = stopF
	return edge.NewReceiverFromForwardReceiverWithStats(
		n.outs,
		edge.NewTimedForwardReceiver(n.timer, r),
	), nil
}

func (n *BarrierNode) newBarrier(group edge.GroupInfo, first edge.PointMeta) (edge.ForwardReceiver, func(), error) {
	switch {
	case n.b.Idle != 0:
		idleBarrier := newIdleBarrier(
			first.Name(),
			group,
			n.b.Idle,
			n.outs,
		)
		return idleBarrier, idleBarrier.Stop, nil
	case n.b.Period != 0:
		periodicBarrier := newPeriodicBarrier(
			first.Name(),
			group,
			n.b.Period,
			n.outs,
		)
		return periodicBarrier, periodicBarrier.Stop, nil
	default:
		return nil, nil, errors.New("unreachable code, barrier node should have non-zero idle or non-zero period")
	}
}

type idleBarrier struct {
	name  string
	group edge.GroupInfo

	idle  time.Duration
	lastT atomic.Value
	timer *time.Timer
	wg    sync.WaitGroup
	outs  []edge.StatsEdge
	stopC chan struct{}
}

func newIdleBarrier(name string, group edge.GroupInfo, idle time.Duration, outs []edge.StatsEdge) *idleBarrier {
	r := &idleBarrier{
		name:  name,
		group: group,
		idle:  idle,
		lastT: atomic.Value{},
		timer: time.NewTimer(idle),
		wg:    sync.WaitGroup{},
		outs:  outs,
		stopC: make(chan struct{}, 1),
	}

	r.Init()

	return r
}

func (n *idleBarrier) Init() {
	n.lastT.Store(time.Time{})
	n.wg.Add(1)

	go n.idleHandler()
}

func (n *idleBarrier) Stop() {
	close(n.stopC)
	n.timer.Stop()
	n.wg.Wait()
}

func (n *idleBarrier) BeginBatch(m edge.BeginBatchMessage) (edge.Message, error) {
	return m, nil
}
func (n *idleBarrier) BatchPoint(m edge.BatchPointMessage) (edge.Message, error) {
	if !m.Time().Before(n.lastT.Load().(time.Time)) {
		n.resetTimer()
		return m, nil
	}
	return nil, nil
}
func (n *idleBarrier) EndBatch(m edge.EndBatchMessage) (edge.Message, error) {
	return m, nil
}
func (n *idleBarrier) Barrier(m edge.BarrierMessage) (edge.Message, error) {
	if !m.Time().Before(n.lastT.Load().(time.Time)) {
		n.resetTimer()
		return m, nil
	}
	return nil, nil
}
func (n *idleBarrier) DeleteGroup(m edge.DeleteGroupMessage) (edge.Message, error) {
	if m.GroupID() == n.group.ID {
		n.Stop()
	}
	return m, nil
}

func (n *idleBarrier) Point(m edge.PointMessage) (edge.Message, error) {
	if !m.Time().Before(n.lastT.Load().(time.Time)) {
		n.resetTimer()
		return m, nil
	}
	return nil, nil
}

func (n *idleBarrier) resetTimer() {
	n.timer.Reset(n.idle)
}

func (n *idleBarrier) emitBarrier() error {
	nowT := time.Now()
	n.lastT.Store(nowT)
	return edge.Forward(n.outs, edge.NewBarrierMessage(n.group, nowT))
}

func (n *idleBarrier) idleHandler() {
	defer n.wg.Done()
	for {
		select {
		case <-n.timer.C:
			n.emitBarrier()
			n.resetTimer()
		case <-n.stopC:
			return
		}
	}
}

type periodicBarrier struct {
	name  string
	group edge.GroupInfo

	lastT  atomic.Value
	ticker *time.Ticker
	wg     sync.WaitGroup
	outs   []edge.StatsEdge
	stopC  chan bool
}

func newPeriodicBarrier(name string, group edge.GroupInfo, period time.Duration, outs []edge.StatsEdge) *periodicBarrier {
	r := &periodicBarrier{
		name:   name,
		group:  group,
		lastT:  atomic.Value{},
		ticker: time.NewTicker(period),
		wg:     sync.WaitGroup{},
		outs:   outs,
		stopC:  make(chan bool, 1),
	}

	r.Init()

	return r
}

func (n *periodicBarrier) Init() {
	n.lastT.Store(time.Time{})
	n.wg.Add(1)

	go n.periodicEmitter()
}

func (n *periodicBarrier) Stop() {
	n.stopC <- true
	n.ticker.Stop()
	n.wg.Wait()
}

func (n *periodicBarrier) BeginBatch(m edge.BeginBatchMessage) (edge.Message, error) {
	return m, nil
}
func (n *periodicBarrier) BatchPoint(m edge.BatchPointMessage) (edge.Message, error) {
	if !m.Time().Before(n.lastT.Load().(time.Time)) {
		return m, nil
	}
	return nil, nil
}
func (n *periodicBarrier) EndBatch(m edge.EndBatchMessage) (edge.Message, error) {
	return m, nil
}
func (n *periodicBarrier) Barrier(m edge.BarrierMessage) (edge.Message, error) {
	if !m.Time().Before(n.lastT.Load().(time.Time)) {
		return m, nil
	}
	return nil, nil
}
func (n *periodicBarrier) DeleteGroup(m edge.DeleteGroupMessage) (edge.Message, error) {
	if m.GroupID() == n.group.ID {
		n.Stop()
	}
	return m, nil
}

func (n *periodicBarrier) Point(m edge.PointMessage) (edge.Message, error) {
	if !m.Time().Before(n.lastT.Load().(time.Time)) {
		return m, nil
	}
	return nil, nil
}

func (n *periodicBarrier) emitBarrier() error {
	nowT := time.Now()
	n.lastT.Store(nowT)
	return edge.Forward(n.outs, edge.NewBarrierMessage(n.group, nowT))
}

func (n *periodicBarrier) periodicEmitter() {
	defer n.wg.Done()
	for {
		select {
		case <-n.ticker.C:
			n.emitBarrier()
		case <-n.stopC:
			return
		}
	}
}
<commit_msg>adding explicit timer.Stop before timer.Reset<commit_after>package kapacitor

import (
	"errors"
	"time"

	"sync"
	"sync/atomic"

	"github.com/influxdata/kapacitor/edge"
	"github.com/influxdata/kapacitor/models"
	"github.com/influxdata/kapacitor/pipeline"
)

type BarrierNode struct {
	node
	b              *pipeline.BarrierNode
	barrierStopper map[models.GroupID]func()
}

// Create a new  BarrierNode, which emits a barrier if data traffic has been idle for the configured amount of time.
func newBarrierNode(et *ExecutingTask, n *pipeline.BarrierNode, d NodeDiagnostic) (*BarrierNode, error) {
	if n.Idle == 0 && n.Period == 0 {
		return nil, errors.New("barrier node must have either a non zero idle or a non zero period")
	}
	bn := &BarrierNode{
		node:           node{Node: n, et: et, diag: d},
		b:              n,
		barrierStopper: map[models.GroupID]func(){},
	}
	bn.node.runF = bn.runBarrierEmitter
	return bn, nil
}

func (n *BarrierNode) runBarrierEmitter([]byte) error {
	defer n.stopBarrierEmitter()
	consumer := edge.NewGroupedConsumer(n.ins[0], n)
	n.statMap.Set(statCardinalityGauge, consumer.CardinalityVar())
	return consumer.Consume()
}

func (n *BarrierNode) stopBarrierEmitter() {
	for _, stopF := range n.barrierStopper {
		stopF()
	}
}

func (n *BarrierNode) NewGroup(group edge.GroupInfo, first edge.PointMeta) (edge.Receiver, error) {
	r, stopF, err := n.newBarrier(group, first)
	if err != nil {
		return nil, err
	}
	n.barrierStopper[group.ID] = stopF
	return edge.NewReceiverFromForwardReceiverWithStats(
		n.outs,
		edge.NewTimedForwardReceiver(n.timer, r),
	), nil
}

func (n *BarrierNode) newBarrier(group edge.GroupInfo, first edge.PointMeta) (edge.ForwardReceiver, func(), error) {
	switch {
	case n.b.Idle != 0:
		idleBarrier := newIdleBarrier(
			first.Name(),
			group,
			n.b.Idle,
			n.outs,
		)
		return idleBarrier, idleBarrier.Stop, nil
	case n.b.Period != 0:
		periodicBarrier := newPeriodicBarrier(
			first.Name(),
			group,
			n.b.Period,
			n.outs,
		)
		return periodicBarrier, periodicBarrier.Stop, nil
	default:
		return nil, nil, errors.New("unreachable code, barrier node should have non-zero idle or non-zero period")
	}
}

type idleBarrier struct {
	name  string
	group edge.GroupInfo

	idle  time.Duration
	lastT atomic.Value
	timer *time.Timer
	wg    sync.WaitGroup
	outs  []edge.StatsEdge
	stopC chan struct{}
}

func newIdleBarrier(name string, group edge.GroupInfo, idle time.Duration, outs []edge.StatsEdge) *idleBarrier {
	r := &idleBarrier{
		name:  name,
		group: group,
		idle:  idle,
		lastT: atomic.Value{},
		timer: time.NewTimer(idle),
		wg:    sync.WaitGroup{},
		outs:  outs,
		stopC: make(chan struct{}, 1),
	}

	r.Init()

	return r
}

func (n *idleBarrier) Init() {
	n.lastT.Store(time.Time{})
	n.wg.Add(1)

	go n.idleHandler()
}

func (n *idleBarrier) Stop() {
	close(n.stopC)
	n.timer.Stop()
	n.wg.Wait()
}

func (n *idleBarrier) BeginBatch(m edge.BeginBatchMessage) (edge.Message, error) {
	return m, nil
}
func (n *idleBarrier) BatchPoint(m edge.BatchPointMessage) (edge.Message, error) {
	if !m.Time().Before(n.lastT.Load().(time.Time)) {
		n.resetTimer()
		return m, nil
	}
	return nil, nil
}
func (n *idleBarrier) EndBatch(m edge.EndBatchMessage) (edge.Message, error) {
	return m, nil
}
func (n *idleBarrier) Barrier(m edge.BarrierMessage) (edge.Message, error) {
	if !m.Time().Before(n.lastT.Load().(time.Time)) {
		n.resetTimer()
		return m, nil
	}
	return nil, nil
}
func (n *idleBarrier) DeleteGroup(m edge.DeleteGroupMessage) (edge.Message, error) {
	if m.GroupID() == n.group.ID {
		n.Stop()
	}
	return m, nil
}

func (n *idleBarrier) Point(m edge.PointMessage) (edge.Message, error) {
	if !m.Time().Before(n.lastT.Load().(time.Time)) {
		n.resetTimer()
		return m, nil
	}
	return nil, nil
}

func (n *idleBarrier) resetTimer() {
	n.timer.Stop()
	n.timer.Reset(n.idle)
}

func (n *idleBarrier) emitBarrier() error {
	nowT := time.Now()
	n.lastT.Store(nowT)
	return edge.Forward(n.outs, edge.NewBarrierMessage(n.group, nowT))
}

func (n *idleBarrier) idleHandler() {
	defer n.wg.Done()
	for {
		select {
		case <-n.timer.C:
			n.emitBarrier()
			n.resetTimer()
		case <-n.stopC:
			return
		}
	}
}

type periodicBarrier struct {
	name  string
	group edge.GroupInfo

	lastT  atomic.Value
	ticker *time.Ticker
	wg     sync.WaitGroup
	outs   []edge.StatsEdge
	stopC  chan bool
}

func newPeriodicBarrier(name string, group edge.GroupInfo, period time.Duration, outs []edge.StatsEdge) *periodicBarrier {
	r := &periodicBarrier{
		name:   name,
		group:  group,
		lastT:  atomic.Value{},
		ticker: time.NewTicker(period),
		wg:     sync.WaitGroup{},
		outs:   outs,
		stopC:  make(chan bool, 1),
	}

	r.Init()

	return r
}

func (n *periodicBarrier) Init() {
	n.lastT.Store(time.Time{})
	n.wg.Add(1)

	go n.periodicEmitter()
}

func (n *periodicBarrier) Stop() {
	n.stopC <- true
	n.ticker.Stop()
	n.wg.Wait()
}

func (n *periodicBarrier) BeginBatch(m edge.BeginBatchMessage) (edge.Message, error) {
	return m, nil
}
func (n *periodicBarrier) BatchPoint(m edge.BatchPointMessage) (edge.Message, error) {
	if !m.Time().Before(n.lastT.Load().(time.Time)) {
		return m, nil
	}
	return nil, nil
}
func (n *periodicBarrier) EndBatch(m edge.EndBatchMessage) (edge.Message, error) {
	return m, nil
}
func (n *periodicBarrier) Barrier(m edge.BarrierMessage) (edge.Message, error) {
	if !m.Time().Before(n.lastT.Load().(time.Time)) {
		return m, nil
	}
	return nil, nil
}
func (n *periodicBarrier) DeleteGroup(m edge.DeleteGroupMessage) (edge.Message, error) {
	if m.GroupID() == n.group.ID {
		n.Stop()
	}
	return m, nil
}

func (n *periodicBarrier) Point(m edge.PointMessage) (edge.Message, error) {
	if !m.Time().Before(n.lastT.Load().(time.Time)) {
		return m, nil
	}
	return nil, nil
}

func (n *periodicBarrier) emitBarrier() error {
	nowT := time.Now()
	n.lastT.Store(nowT)
	return edge.Forward(n.outs, edge.NewBarrierMessage(n.group, nowT))
}

func (n *periodicBarrier) periodicEmitter() {
	defer n.wg.Done()
	for {
		select {
		case <-n.ticker.C:
			n.emitBarrier()
		case <-n.stopC:
			return
		}
	}
}
<|endoftext|>


<commit_before>package stl

import (
	"bytes"
	"strings"
	"testing"
)

func TestPoint(t *testing.T) {
	p := Point{0, 0, 1}
	if !p.IsNormal() {
		t.Error("should be normal")
	}
	if p.Abs() != 1.0 {
		t.Error("abs is wrong")
	}
}

func TestStlReader(t *testing.T) {
	data := []byte{
		// header
		'C', 'O', 'L', 'O', 'R', '=', 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 0
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 16
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 32
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 64
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 80
		1, 0, 0, 0, // 80: 4 bytes, little endian, how many triangles
		// 84: triangle
	}

	tri := Triangle{
		Vector{1, 1, 1},
		[3]Point{Point{0, 0, 0},
			Point{1, 2, 3},
			Point{3, 2, -1},
		}, 0,
	}
	data = append(data, formatTriangle(&tri)...)

	s, err := Decode(bytes.NewReader(data[0:10]))
	if err == nil {
		t.Error("should have gotten an error")
	}

	s, err = Decode(bytes.NewReader(data))
	if err != nil {
		t.Error("got err", err)
	}
	if !strings.Contains(string(s.Header[:]), "COLOR=") {
		t.Error("header is missing color=")
	}
	if len(s.Triangles) != 1 {
		t.Errorf("expected 1 triangle, got %d", len(s.Triangles))
	}

	if !s.Bounds.Min.Equals(Point{0, 0, -1}) {
		t.Error("bad bounds min", s.Bounds.Min)
	}
	if !s.Bounds.Max.Equals(Point{3, 2, 3}) {
		t.Error("bad bounds max", s.Bounds.Max)
	}
}
<commit_msg>Fix the build.<commit_after>package stl

import (
	"bytes"
	"strings"
	"testing"
)

/*
func TestPoint(t *testing.T) {
	p := Point{0, 0, 1}
	if !p.IsNormal() {
		t.Error("should be normal")
	}
	if p.Abs() != 1.0 {
		t.Error("abs is wrong")
	}
}
*/

func TestStlReader(t *testing.T) {
	data := []byte{
		// header
		'C', 'O', 'L', 'O', 'R', '=', 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 0
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 16
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 32
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 64
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // 80
		1, 0, 0, 0, // 80: 4 bytes, little endian, how many triangles
		// 84: triangle
	}

	tri := Triangle{
		Vector{1, 1, 1},
		[3]Point{Point{0, 0, 0},
			Point{1, 2, 3},
			Point{3, 2, -1},
		}, 0,
	}
	data = append(data, formatTriangle(&tri)...)

	s, err := Decode(bytes.NewReader(data[0:10]))
	if err == nil {
		t.Error("should have gotten an error")
	}

	s, err = Decode(bytes.NewReader(data))
	if err != nil {
		t.Error("got err", err)
	}
	if !strings.Contains(string(s.Header[:]), "COLOR=") {
		t.Error("header is missing color=")
	}
	if len(s.Triangles) != 1 {
		t.Errorf("expected 1 triangle, got %d", len(s.Triangles))
	}

	if !s.Bounds.Min.Equals(Point{0, 0, -1}) {
		t.Error("bad bounds min", s.Bounds.Min)
	}
	if !s.Bounds.Max.Equals(Point{3, 2, 3}) {
		t.Error("bad bounds max", s.Bounds.Max)
	}
}
<|endoftext|>


<commit_before>package sudoku

import (
	"testing"
)

func TestRandomWeightedIndex(t *testing.T) {
	result := randomIndexWithNormalizedWeights([]float64{1.0, 0.0})
	if result != 0 {
		t.Log("Got wrong result with random weights")
		t.Fail()
	}
	result = randomIndexWithNormalizedWeights([]float64{0.5, 0.0, 0.5})
	if result != 0 && result != 2 {
		t.Log("Didn't get one of two legal weights")
		t.Fail()
	}
	result = randomIndexWithNormalizedWeights([]float64{0.0, 0.0, 1.0})
	if result != 2 {
		t.Log("Should have gotten last item in random weights; we didn't")
		t.Fail()
	}
	if weightsNormalized([]float64{1.0, 0.000001}) {
		t.Log("thought weights were normalized when they weren't")
		t.Fail()
	}
	if !weightsNormalized([]float64{0.5, 0.25, 0.25}) {
		t.Log("Didn't think weights were normalized but they were")
		t.Fail()
	}

	if weightsNormalized([]float64{0.5, -0.25, 0.25}) {
		t.Error("A negative weight was considered normal.")
	}

	result = randomIndexWithInvertedWeights([]float64{0.0, 0.0, 1.0})
	if result == 2 {
		t.Log("Got the wrong index for inverted weights")
		t.Fail()
	}

	weightResult := normalizedWeights([]float64{2.0, 1.0, 1.0})
	if weightResult[0] != 0.5 || weightResult[1] != 0.25 || weightResult[2] != 0.25 {
		t.Log("Nomralized weights came back wrong")
		t.Fail()
	}

	weightResult = normalizedWeights([]float64{1.0, 1.0, -0.5})
	if weightResult[0] != 0.5 || weightResult[1] != 0.5 || weightResult[2] != 0 {
		t.Error("Normalized weights with a negative came back wrong: ", weightResult)
	}

	weightResult = normalizedWeights([]float64{-0.25, -0.5, 0.25})
	if weightResult[0] != 0.25 || weightResult[1] != 0 || weightResult[2] != 0.75 {
		t.Error("Normalized weights with two different negative numbers came back wrong: ", weightResult)
	}

	result = randomIndexWithWeights([]float64{1.0, 0.0})
	if result != 0 {
		t.Log("Got wrong result with random weights")
		t.Fail()
	}
	result = randomIndexWithWeights([]float64{5.0, 0.0, 5.0})
	if result != 0 && result != 2 {
		t.Log("Didn't get one of two legal weights")
		t.Fail()
	}
	result = randomIndexWithWeights([]float64{0.0, 0.0, 5.0})
	if result != 2 {
		t.Log("Should have gotten last item in random weights; we didn't")
		t.Fail()
	}
}<commit_msg>TESTS FAIL. Added a test that shows that invertedRandomIndex is broken.<commit_after>package sudoku

import (
	"testing"
)

func TestRandomWeightedIndex(t *testing.T) {
	result := randomIndexWithNormalizedWeights([]float64{1.0, 0.0})
	if result != 0 {
		t.Log("Got wrong result with random weights")
		t.Fail()
	}
	result = randomIndexWithNormalizedWeights([]float64{0.5, 0.0, 0.5})
	if result != 0 && result != 2 {
		t.Log("Didn't get one of two legal weights")
		t.Fail()
	}
	result = randomIndexWithNormalizedWeights([]float64{0.0, 0.0, 1.0})
	if result != 2 {
		t.Log("Should have gotten last item in random weights; we didn't")
		t.Fail()
	}
	if weightsNormalized([]float64{1.0, 0.000001}) {
		t.Log("thought weights were normalized when they weren't")
		t.Fail()
	}
	if !weightsNormalized([]float64{0.5, 0.25, 0.25}) {
		t.Log("Didn't think weights were normalized but they were")
		t.Fail()
	}

	if weightsNormalized([]float64{0.5, -0.25, 0.25}) {
		t.Error("A negative weight was considered normal.")
	}

	result = randomIndexWithInvertedWeights([]float64{0.0, 0.0, 1.0})
	if result == 2 {
		t.Log("Got the wrong index for inverted weights")
		t.Fail()
	}

	//TODO: this test is inherently flaky
	for i := 0; i < 10; i++ {
		result = randomIndexWithInvertedWeights([]float64{0.5, 1.0, 0.0})
		if result == 2 {
			t.Error("RandominzedIndexWithInvertedWeights returned wrong result")
			break
		}
	}

	weightResult := normalizedWeights([]float64{2.0, 1.0, 1.0})
	if weightResult[0] != 0.5 || weightResult[1] != 0.25 || weightResult[2] != 0.25 {
		t.Log("Nomralized weights came back wrong")
		t.Fail()
	}

	weightResult = normalizedWeights([]float64{1.0, 1.0, -0.5})
	if weightResult[0] != 0.5 || weightResult[1] != 0.5 || weightResult[2] != 0 {
		t.Error("Normalized weights with a negative came back wrong: ", weightResult)
	}

	weightResult = normalizedWeights([]float64{-0.25, -0.5, 0.25})
	if weightResult[0] != 0.25 || weightResult[1] != 0 || weightResult[2] != 0.75 {
		t.Error("Normalized weights with two different negative numbers came back wrong: ", weightResult)
	}

	result = randomIndexWithWeights([]float64{1.0, 0.0})
	if result != 0 {
		t.Log("Got wrong result with random weights")
		t.Fail()
	}
	result = randomIndexWithWeights([]float64{5.0, 0.0, 5.0})
	if result != 0 && result != 2 {
		t.Log("Didn't get one of two legal weights")
		t.Fail()
	}
	result = randomIndexWithWeights([]float64{0.0, 0.0, 5.0})
	if result != 2 {
		t.Log("Should have gotten last item in random weights; we didn't")
		t.Fail()
	}
}
<|endoftext|>


<commit_before>//go:build gofuzz
// +build gofuzz

/*
   Copyright The containerd Authors.
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at
       http://www.apache.org/licenses/LICENSE-2.0
   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package fuzz

import (
	"bytes"
	"context"
	"os"

	fuzz "github.com/AdaLogics/go-fuzz-headers"
	"github.com/sirupsen/logrus"

	"github.com/containerd/containerd/archive"
	"github.com/containerd/containerd/content/local"
	imageArchive "github.com/containerd/containerd/images/archive"
	"github.com/containerd/containerd/log"
)

// FuzzApply implements a fuzzer that applies
// a fuzzed tar archive on a directory
func FuzzApply(data []byte) int {
	ctx := context.Background()

	// Apply() is logging the below message, which is too noisy and not really useful
	// if the input is random.
	//
	// level=warning msg="ignored xattr ... in archive" error="operation not supported"
	log.G(ctx).Logger.SetLevel(logrus.PanicLevel)

	f := fuzz.NewConsumer(data)
	iters, err := f.GetInt()
	if err != nil {
		return 0
	}
	maxIters := 20
	tmpDir, err := os.MkdirTemp("", "prefix-test")
	if err != nil {
		return 0
	}
	defer os.RemoveAll(tmpDir)
	for i := 0; i < iters%maxIters; i++ {
		rBytes, err := f.TarBytes()
		if err != nil {
			return 0
		}
		r := bytes.NewReader(rBytes)
		_, _ = archive.Apply(ctx, tmpDir, r)
	}
	return 1
}

// FuzzImportIndex implements a fuzzer
// that targets archive.ImportIndex()
func FuzzImportIndex(data []byte) int {
	f := fuzz.NewConsumer(data)
	tarBytes, err := f.GetBytes()
	if err != nil {
		return 0
	}
	ctx := context.Background()
	r := bytes.NewReader(tarBytes)
	tmpdir, err := os.MkdirTemp("", "fuzzing-")
	if err != nil {
		return 0
	}
	cs, err := local.NewStore(tmpdir)
	if err != nil {
		return 0
	}
	_, _ = imageArchive.ImportIndex(ctx, cs, r)
	return 1
}
<commit_msg>fuzzing: create structured tar bytes in archive fuzzer<commit_after>//go:build gofuzz
// +build gofuzz

/*
   Copyright The containerd Authors.
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at
       http://www.apache.org/licenses/LICENSE-2.0
   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package fuzz

import (
	"bytes"
	"context"
	"os"

	fuzz "github.com/AdaLogics/go-fuzz-headers"
	"github.com/sirupsen/logrus"

	"github.com/containerd/containerd/archive"
	"github.com/containerd/containerd/content/local"
	imageArchive "github.com/containerd/containerd/images/archive"
	"github.com/containerd/containerd/log"
)

// FuzzApply implements a fuzzer that applies
// a fuzzed tar archive on a directory
func FuzzApply(data []byte) int {
	ctx := context.Background()

	// Apply() is logging the below message, which is too noisy and not really useful
	// if the input is random.
	//
	// level=warning msg="ignored xattr ... in archive" error="operation not supported"
	log.G(ctx).Logger.SetLevel(logrus.PanicLevel)

	f := fuzz.NewConsumer(data)
	iters, err := f.GetInt()
	if err != nil {
		return 0
	}
	maxIters := 20
	tmpDir, err := os.MkdirTemp("", "prefix-test")
	if err != nil {
		return 0
	}
	defer os.RemoveAll(tmpDir)
	for i := 0; i < iters%maxIters; i++ {
		rBytes, err := f.TarBytes()
		if err != nil {
			return 0
		}
		r := bytes.NewReader(rBytes)
		_, _ = archive.Apply(ctx, tmpDir, r)
	}
	return 1
}

// FuzzImportIndex implements a fuzzer
// that targets archive.ImportIndex()
func FuzzImportIndex(data []byte) int {
	f := fuzz.NewConsumer(data)
	tarBytes, err := f.TarBytes()
	if err != nil {
		return 0
	}
	ctx := context.Background()
	r := bytes.NewReader(tarBytes)
	tmpdir, err := os.MkdirTemp("", "fuzzing-")
	if err != nil {
		return 0
	}
	cs, err := local.NewStore(tmpdir)
	if err != nil {
		return 0
	}
	_, _ = imageArchive.ImportIndex(ctx, cs, r)
	return 1
}
<|endoftext|>


<commit_before>// layer_0 contains a test that brings up layer 0 (a single-node Kubernetes
// cluster) on a GCP VM. To run it on your workstation, first create a service
// account:
//   https://console.cloud.google.com/apis/credentials/serviceaccountkey
// make it a "Compute Admin", and download the JSON key. Go to:
//   https://console.cloud.google.com/iam-admin/serviceaccounts
// tick the box by ######-compute@, and make your new SA a "Service Account
// User". You might need to wait ~10 minutes for IAM to propagate. Then replace
// your-project-id below and run:
//   bazel test \
//     --test_env GCP_PROJECT_ID=your-project-id --test_env GCP_ZONE=europe-west1-c \
//     --test_env GOOGLE_APPLICATION_CREDENTIALS=path/to/credentials.json \
//     --test_tag_filters=external //src/go/tests/layer-0:all

package layer_0

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"log"
	mrand "math/rand"
	"os"
	"path"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/apcera/libretto/ssh"
	"github.com/apcera/libretto/virtualmachine/gcp"
	cssh "golang.org/x/crypto/ssh"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/compute/v1"
)

const (
	alphanum = "abcdefghijklmnopqrstuvwxyz0123456789"

	// Note: to use a workspace-relative path, the go_test() must have `rundir = "."`.
	clusterInstallScript = "./src/bootstrap/robot/install_k8s_on_robot.sh"
)

var (
	GCP_PROJECT_ID                 = mustGetenv("GCP_PROJECT_ID")
	GCP_ZONE                       = mustGetenv("GCP_ZONE")
	GOOGLE_APPLICATION_CREDENTIALS = mustGetenv("GOOGLE_APPLICATION_CREDENTIALS")
)

func init() {
	mrand.Seed(time.Now().UnixNano())
}

func mustGetenv(key string) string {
	value := os.Getenv(key)
	if value == "" {
		log.Fatalf("ERROR: environment variable %s must be non-empty", key)
	}
	return value
}

func randomString(n int) string {
	b := make([]byte, n)
	for i := range b {
		b[i] = alphanum[mrand.Intn(len(alphanum))]
	}
	return string(b)
}

// generateKeypair generates an SSH keypair that will be used to connect to the VM.
func generateKeypair() ([]byte, []byte, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to generate SSH private key: %v", err)
	}
	privatePEM := pem.EncodeToMemory(&pem.Block{
		Type:    "RSA PRIVATE KEY",
		Headers: nil,
		Bytes:   x509.MarshalPKCS1PrivateKey(privateKey),
	})
	publicKey, err := cssh.NewPublicKey(&privateKey.PublicKey)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to generate SSH public key: %v", err)
	}
	publicPEM := cssh.MarshalAuthorizedKey(publicKey)
	return privatePEM, publicPEM, nil
}

// getImageName gets the name of the latest image in a given family.
func getImageName(project, family string) (string, error) {
	// Use oauth2.NoContext if there isn't a good context to pass in.
	ctx := context.Background()

	client, err := google.DefaultClient(ctx, compute.ComputeScope)
	if err != nil {
		return "", err
	}
	computeService, err := compute.New(client)
	if err != nil {
		return "", err
	}
	image, err := computeService.Images.GetFromFamily(project, family).Do()
	return image.Name, err
}

func newVM(privatePEM []byte, publicPEM []byte, imageProject, imageFamily string) (*gcp.VM, error) {
	sourceImage, err := getImageName(imageProject, imageFamily)
	if err != nil {
		return nil, err
	}

	return &gcp.VM{
		Name:          "layer-zero-test-vm-" + randomString(8),
		Zone:          GCP_ZONE,
		MachineType:   "n1-standard-4",
		SourceImage:   sourceImage,
		ImageProjects: []string{imageProject},
		Disks: []gcp.Disk{
			{
				DiskType:   "pd-standard",
				DiskSizeGb: 200,
				AutoDelete: true,
			},
		},
		Network:     "default",
		Subnetwork:  "default",
		Project:     GCP_PROJECT_ID,
		Scopes:      []string{"https://www.googleapis.com/auth/cloud-platform"},
		AccountFile: GOOGLE_APPLICATION_CREDENTIALS,
		SSHCreds: ssh.Credentials{
			SSHUser:       "robot",
			SSHPrivateKey: string(privatePEM),
		},
		SSHPublicKey: string(publicPEM),
		Tags:         []string{"delete-after-one-day"},
	}, nil
}

// isStartedByKokoro checks whether the test was started by Kokoro by checking
// the hostname. This is used to identify whether the VM's private IP should be
// used for SSH connections (valid from within the same cloud project as the
// target VM), or whether to use the public IP (valid from a dev workstation).
func isStartedByKokoro() (bool, error) {
	hostname, err := os.Hostname()
	if err != nil {
		return false, fmt.Errorf("failed to get hostname: %v", err)
	}
	return strings.HasPrefix(hostname, "kokoro"), nil
}

// getSSH gets a connected SSH client, connecting to the public or private IP
// depending on how the test was started.
func getSSH(t *testing.T, vm *gcp.VM) (ssh.Client, error) {
	ips, err := vm.GetIPs()
	if err != nil {
		return nil, err
	}

	client := &ssh.SSHClient{
		Creds: &vm.SSHCreds,
		Port:  22,
	}

	if onKokoro, err := isStartedByKokoro(); err != nil {
		return nil, err
	} else if onKokoro {
		client.IP = ips[gcp.PrivateIP]
		t.Logf("Detected that VM is started by Kokoro, so using private IP: %v", client.IP)
	} else {
		client.IP = ips[gcp.PublicIP]
		t.Logf("Using public VM IP: %v", client.IP)
	}

	if err := client.WaitForSSH(5 * time.Minute); err != nil {
		return nil, err
	}

	return client, nil
}

// uploadLocalFile uploads a local file over SSH.
func uploadLocalFile(client ssh.Client, localPath string, remotePath string, mode uint32) error {
	s, err := os.Stat(localPath)
	if err != nil {
		return fmt.Errorf("failed to stat %s: %v", localPath, err)
	}
	f, err := os.Open(localPath)
	if err != nil {
		return fmt.Errorf("failed to open %s: %v", localPath, err)
	}
	if err := client.Upload(f, remotePath, int(s.Size()), mode); err != nil {
		return fmt.Errorf("failed to upload to %s: %v", remotePath, err)
	}
	return nil
}

// uploadAndRun uploads a local executable script over SSH and runs it.
func uploadAndRun(t *testing.T, client ssh.Client, localPath string, args ...string) error {
	remotePath := "~/" + path.Base(localPath)
	if err := uploadLocalFile(client, localPath, remotePath, 0755); err != nil {
		return err
	}
	quotedArgs := make([]string, len(args))
	for i, arg := range args {
		quotedArgs[i] = strconv.Quote(arg)
	}
	// To avoid mixing the output streams of multiple commands running in
	// parallel tests, we combine stdout and stderr on the remote host, and
	// capture stdout and write it to the test log here.
	command := fmt.Sprintf("%s %s 2>&1", remotePath, strings.Join(quotedArgs, " "))
	stdout := new(bytes.Buffer)
	if err := client.Run(command, stdout, os.Stderr); err != nil {
		t.Logf("`%s` failed with %v:\n%s", command, err, stdout.String())
		return fmt.Errorf("failed to run `%s` on VM: %v", remotePath, err)
	}
	return nil
}

// TestInstallLayerZero creates a new VM and installs a Kubernetes cluster on it.
func TestInstallLayerZero(t *testing.T) {
	privatePEM, publicPEM, err := generateKeypair()
	if err != nil {
		t.Fatal(err)
	}

	// Together, these identify the public image used for the VM, without
	// pinning a specific version. We test each image in parallel as the tests
	// are independent and take 3-4 minutes.
	testCases := []struct {
		imageProject string
		imageFamily  string
	}{
		{"ubuntu-os-cloud", "ubuntu-1404-lts"},
		{"ubuntu-os-cloud", "ubuntu-1604-lts"},
		{"ubuntu-os-cloud", "ubuntu-1804-lts"},
	}
	for _, tc := range testCases {
		vm, err := newVM(privatePEM, publicPEM, tc.imageProject, tc.imageFamily)
		if err != nil {
			t.Fatal(err)
		}
		t.Run(tc.imageFamily, func(t *testing.T) {
			t.Parallel()
			defer func() {
				t.Log("Deleting VM...")
				if err := vm.Destroy(); err != nil {
					t.Error("Failed to delete VM:", err)
				}
			}()

			t.Logf("Provisioning VM %s...", vm.Name)
			if err := vm.Provision(); err != nil {
				t.Fatal("Failed to provision VM:", err)
			}

			t.Log("Establishing SSH connection to VM...")
			client, err := getSSH(t, vm)
			if err != nil {
				t.Fatal("Failed to SSH to VM:", err)
			}

			t.Log("Installing cluster on VM...")
			if err := uploadAndRun(t, client, clusterInstallScript); err != nil {
				t.Fatal("Failed to install cluster:", err)
			}
		})
	}
}
<commit_msg>Test install_k8s_on_robot.sh with xtrace<commit_after>// layer_0 contains a test that brings up layer 0 (a single-node Kubernetes
// cluster) on a GCP VM. To run it on your workstation, first create a service
// account:
//   https://console.cloud.google.com/apis/credentials/serviceaccountkey
// make it a "Compute Admin", and download the JSON key. Go to:
//   https://console.cloud.google.com/iam-admin/serviceaccounts
// tick the box by ######-compute@, and make your new SA a "Service Account
// User". You might need to wait ~10 minutes for IAM to propagate. Then replace
// your-project-id below and run:
//   bazel test \
//     --test_env GCP_PROJECT_ID=your-project-id --test_env GCP_ZONE=europe-west1-c \
//     --test_env GOOGLE_APPLICATION_CREDENTIALS=path/to/credentials.json \
//     --test_tag_filters=external //src/go/tests/layer-0:all

package layer_0

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"log"
	mrand "math/rand"
	"os"
	"path"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/apcera/libretto/ssh"
	"github.com/apcera/libretto/virtualmachine/gcp"
	cssh "golang.org/x/crypto/ssh"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/compute/v1"
)

const (
	alphanum = "abcdefghijklmnopqrstuvwxyz0123456789"

	// Note: to use a workspace-relative path, the go_test() must have `rundir = "."`.
	clusterInstallScript = "./src/bootstrap/robot/install_k8s_on_robot.sh"
)

var (
	GCP_PROJECT_ID                 = mustGetenv("GCP_PROJECT_ID")
	GCP_ZONE                       = mustGetenv("GCP_ZONE")
	GOOGLE_APPLICATION_CREDENTIALS = mustGetenv("GOOGLE_APPLICATION_CREDENTIALS")
)

func init() {
	mrand.Seed(time.Now().UnixNano())
}

func mustGetenv(key string) string {
	value := os.Getenv(key)
	if value == "" {
		log.Fatalf("ERROR: environment variable %s must be non-empty", key)
	}
	return value
}

func randomString(n int) string {
	b := make([]byte, n)
	for i := range b {
		b[i] = alphanum[mrand.Intn(len(alphanum))]
	}
	return string(b)
}

// generateKeypair generates an SSH keypair that will be used to connect to the VM.
func generateKeypair() ([]byte, []byte, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to generate SSH private key: %v", err)
	}
	privatePEM := pem.EncodeToMemory(&pem.Block{
		Type:    "RSA PRIVATE KEY",
		Headers: nil,
		Bytes:   x509.MarshalPKCS1PrivateKey(privateKey),
	})
	publicKey, err := cssh.NewPublicKey(&privateKey.PublicKey)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to generate SSH public key: %v", err)
	}
	publicPEM := cssh.MarshalAuthorizedKey(publicKey)
	return privatePEM, publicPEM, nil
}

// getImageName gets the name of the latest image in a given family.
func getImageName(project, family string) (string, error) {
	// Use oauth2.NoContext if there isn't a good context to pass in.
	ctx := context.Background()

	client, err := google.DefaultClient(ctx, compute.ComputeScope)
	if err != nil {
		return "", err
	}
	computeService, err := compute.New(client)
	if err != nil {
		return "", err
	}
	image, err := computeService.Images.GetFromFamily(project, family).Do()
	return image.Name, err
}

func newVM(privatePEM []byte, publicPEM []byte, imageProject, imageFamily string) (*gcp.VM, error) {
	sourceImage, err := getImageName(imageProject, imageFamily)
	if err != nil {
		return nil, err
	}

	return &gcp.VM{
		Name:          "layer-zero-test-vm-" + randomString(8),
		Zone:          GCP_ZONE,
		MachineType:   "n1-standard-4",
		SourceImage:   sourceImage,
		ImageProjects: []string{imageProject},
		Disks: []gcp.Disk{
			{
				DiskType:   "pd-standard",
				DiskSizeGb: 200,
				AutoDelete: true,
			},
		},
		Network:     "default",
		Subnetwork:  "default",
		Project:     GCP_PROJECT_ID,
		Scopes:      []string{"https://www.googleapis.com/auth/cloud-platform"},
		AccountFile: GOOGLE_APPLICATION_CREDENTIALS,
		SSHCreds: ssh.Credentials{
			SSHUser:       "robot",
			SSHPrivateKey: string(privatePEM),
		},
		SSHPublicKey: string(publicPEM),
		Tags:         []string{"delete-after-one-day"},
	}, nil
}

// isStartedByKokoro checks whether the test was started by Kokoro by checking
// the hostname. This is used to identify whether the VM's private IP should be
// used for SSH connections (valid from within the same cloud project as the
// target VM), or whether to use the public IP (valid from a dev workstation).
func isStartedByKokoro() (bool, error) {
	hostname, err := os.Hostname()
	if err != nil {
		return false, fmt.Errorf("failed to get hostname: %v", err)
	}
	return strings.HasPrefix(hostname, "kokoro"), nil
}

// getSSH gets a connected SSH client, connecting to the public or private IP
// depending on how the test was started.
func getSSH(t *testing.T, vm *gcp.VM) (ssh.Client, error) {
	ips, err := vm.GetIPs()
	if err != nil {
		return nil, err
	}

	client := &ssh.SSHClient{
		Creds: &vm.SSHCreds,
		Port:  22,
	}

	if onKokoro, err := isStartedByKokoro(); err != nil {
		return nil, err
	} else if onKokoro {
		client.IP = ips[gcp.PrivateIP]
		t.Logf("Detected that VM is started by Kokoro, so using private IP: %v", client.IP)
	} else {
		client.IP = ips[gcp.PublicIP]
		t.Logf("Using public VM IP: %v", client.IP)
	}

	if err := client.WaitForSSH(5 * time.Minute); err != nil {
		return nil, err
	}

	return client, nil
}

// uploadLocalFile uploads a local file over SSH.
func uploadLocalFile(client ssh.Client, localPath string, remotePath string, mode uint32) error {
	s, err := os.Stat(localPath)
	if err != nil {
		return fmt.Errorf("failed to stat %s: %v", localPath, err)
	}
	f, err := os.Open(localPath)
	if err != nil {
		return fmt.Errorf("failed to open %s: %v", localPath, err)
	}
	if err := client.Upload(f, remotePath, int(s.Size()), mode); err != nil {
		return fmt.Errorf("failed to upload to %s: %v", remotePath, err)
	}
	return nil
}

// uploadAndRun uploads a local bash script over SSH and runs it.
func uploadAndRun(t *testing.T, client ssh.Client, localPath string, args ...string) error {
	remotePath := "~/" + path.Base(localPath)
	if err := uploadLocalFile(client, localPath, remotePath, 0755); err != nil {
		return err
	}
	quotedArgs := make([]string, len(args))
	for i, arg := range args {
		quotedArgs[i] = strconv.Quote(arg)
	}
	// To avoid mixing the output streams of multiple commands running in
	// parallel tests, we combine stdout and stderr on the remote host, and
	// capture stdout and write it to the test log here.
	command := fmt.Sprintf(`bash -x %s %s 2>&1`, remotePath, strings.Join(quotedArgs, " "))
	stdout := new(bytes.Buffer)
	if err := client.Run(command, stdout, os.Stderr); err != nil {
		t.Logf("`%s` failed with %v:\n%s", command, err, stdout.String())
		return fmt.Errorf("failed to run `%s` on VM: %v", remotePath, err)
	}
	return nil
}

// TestInstallLayerZero creates a new VM and installs a Kubernetes cluster on it.
func TestInstallLayerZero(t *testing.T) {
	privatePEM, publicPEM, err := generateKeypair()
	if err != nil {
		t.Fatal(err)
	}

	// Together, these identify the public image used for the VM, without
	// pinning a specific version. We test each image in parallel as the tests
	// are independent and take 3-4 minutes.
	testCases := []struct {
		imageProject string
		imageFamily  string
	}{
		{"ubuntu-os-cloud", "ubuntu-1404-lts"},
		{"ubuntu-os-cloud", "ubuntu-1604-lts"},
		{"ubuntu-os-cloud", "ubuntu-1804-lts"},
	}
	for _, tc := range testCases {
		vm, err := newVM(privatePEM, publicPEM, tc.imageProject, tc.imageFamily)
		if err != nil {
			t.Fatal(err)
		}
		t.Run(tc.imageFamily, func(t *testing.T) {
			t.Parallel()
			defer func() {
				t.Log("Deleting VM...")
				if err := vm.Destroy(); err != nil {
					t.Error("Failed to delete VM:", err)
				}
			}()

			t.Logf("Provisioning VM %s...", vm.Name)
			if err := vm.Provision(); err != nil {
				t.Fatal("Failed to provision VM:", err)
			}

			t.Log("Establishing SSH connection to VM...")
			client, err := getSSH(t, vm)
			if err != nil {
				t.Fatal("Failed to SSH to VM:", err)
			}

			t.Log("Installing cluster on VM...")
			if err := uploadAndRun(t, client, clusterInstallScript); err != nil {
				t.Fatal("Failed to install cluster:", err)
			}
		})
	}
}
<|endoftext|>


<commit_before>package packer

import (
	"fmt"
	"log"
	"sync"
)

// This is the key in configurations that is set to "true" when Packer
// debugging is enabled.
const DebugConfigKey = "packer_debug"

// A Build represents a single job within Packer that is responsible for
// building some machine image artifact. Builds are meant to be parallelized.
type Build interface {
	// Name is the name of the build. This is unique across a single template,
	// but not absolutely unique. This is meant more to describe to the user
	// what is being built rather than being a unique identifier.
	Name() string

	// Prepare configures the various components of this build and reports
	// any errors in doing so (such as syntax errors, validation errors, etc.)
	Prepare() error

	// Run runs the actual builder, returning an artifact implementation
	// of what is built. If anything goes wrong, an error is returned.
	Run(Ui, Cache) ([]Artifact, error)

	// Cancel will cancel a running build. This will block until the build
	// is actually completely cancelled.
	Cancel()

	// SetDebug will enable/disable debug mode. Debug mode is always
	// enabled by adding the additional key "packer_debug" to boolean
	// true in the configuration of the various components. This must
	// be called prior to Prepare.
	//
	// When SetDebug is set to true, parallelism between builds is
	// strictly prohibited.
	SetDebug(bool)
}

// A build struct represents a single build job, the result of which should
// be a single machine image artifact. This artifact may be comprised of
// multiple files, of course, but it should be for only a single provider
// (such as VirtualBox, EC2, etc.).
type coreBuild struct {
	name           string
	builder        Builder
	builderConfig  interface{}
	hooks          map[string][]Hook
	postProcessors [][]coreBuildPostProcessor
	provisioners   []coreBuildProvisioner

	debug         bool
	l             sync.Mutex
	prepareCalled bool
}

// Keeps track of the post-processor and the configuration of the
// post-processor used within a build.
type coreBuildPostProcessor struct {
	processor         PostProcessor
	config            interface{}
	keepInputArtifact bool
}

// Keeps track of the provisioner and the configuration of the provisioner
// within the build.
type coreBuildProvisioner struct {
	provisioner Provisioner
	config      []interface{}
}

// Returns the name of the build.
func (b *coreBuild) Name() string {
	return b.name
}

// Prepare prepares the build by doing some initialization for the builder
// and any hooks. This _must_ be called prior to Run.
func (b *coreBuild) Prepare() (err error) {
	b.l.Lock()
	defer b.l.Unlock()

	if b.prepareCalled {
		panic("prepare already called")
	}

	b.prepareCalled = true

	debugConfig := map[string]interface{}{
		DebugConfigKey: b.debug,
	}

	// Prepare the builder
	err = b.builder.Prepare(b.builderConfig, debugConfig)
	if err != nil {
		log.Printf("Build '%s' prepare failure: %s\n", b.name, err)
		return
	}

	// Prepare the provisioners
	for _, coreProv := range b.provisioners {
		configs := make([]interface{}, len(coreProv.config), len(coreProv.config)+1)
		copy(configs, coreProv.config)
		configs = append(configs, debugConfig)

		if err = coreProv.provisioner.Prepare(configs...); err != nil {
			return
		}
	}

	// Prepare the post-processors
	for _, ppSeq := range b.postProcessors {
		for _, corePP := range ppSeq {
			if err = corePP.processor.Configure(corePP.config); err != nil {
				return
			}
		}
	}

	return
}

// Runs the actual build. Prepare must be called prior to running this.
func (b *coreBuild) Run(ui Ui, cache Cache) ([]Artifact, error) {
	if !b.prepareCalled {
		panic("Prepare must be called first")
	}

	// Copy the hooks
	hooks := make(map[string][]Hook)
	for hookName, hookList := range b.hooks {
		hooks[hookName] = make([]Hook, len(hookList))
		copy(hooks[hookName], hookList)
	}

	// Add a hook for the provisioners if we have provisioners
	if len(b.provisioners) > 0 {
		provisioners := make([]Provisioner, len(b.provisioners))
		for i, p := range b.provisioners {
			provisioners[i] = p.provisioner
		}

		if _, ok := hooks[HookProvision]; !ok {
			hooks[HookProvision] = make([]Hook, 0, 1)
		}

		hooks[HookProvision] = append(hooks[HookProvision], &ProvisionHook{provisioners})
	}

	hook := &DispatchHook{hooks}
	artifacts := make([]Artifact, 0, 1)

	builderArtifact, err := b.builder.Run(ui, hook, cache)
	if err != nil {
		return nil, err
	}

	errors := make([]error, 0)
	keepOriginalArtifact := len(b.postProcessors) == 0

	// Run the post-processors
PostProcessorRunSeqLoop:
	for _, ppSeq := range b.postProcessors {
		artifact := builderArtifact

		var priorArtifact Artifact
		for i, corePP := range ppSeq {
			if i == 0 {
				// This is the first post-processor. We handle deleting
				// previous artifacts a bit different because multiple
				// post-processors may be using the original and need it.
				if !keepOriginalArtifact {
					keepOriginalArtifact = corePP.keepInputArtifact
				}
			} else {
				if priorArtifact == nil {
					errors = append(errors, fmt.Errorf("Post-processor returned nil artifact mid-chain."))
					continue PostProcessorRunSeqLoop
				}

				// We have a prior artifact. If we want to keep it, we append
				// it to the results list. Otherwise, we destroy it.
				if corePP.keepInputArtifact {
					artifacts = append(artifacts, priorArtifact)
				} else {
					if err := priorArtifact.Destroy(); err != nil {
						errors = append(errors, fmt.Errorf("Failed cleaning up prior artifact: %s", err))
					}
				}
			}

			var err error
			artifact, err = corePP.processor.PostProcess(ui, artifact)
			if err != nil {
				errors = append(errors, fmt.Errorf("Post-processor failed: %s", err))
				continue PostProcessorRunSeqLoop
			}

			priorArtifact = artifact
		}

		// Add on the last artifact to the results
		if priorArtifact != nil {
			artifacts = append(artifacts, priorArtifact)
		}
	}

	if keepOriginalArtifact {
		artifacts = append(artifacts, nil)
		copy(artifacts[1:], artifacts)
		artifacts[0] = builderArtifact
	} else {
		log.Printf("Deleting original artifact for build '%s'", b.name)
		if err := builderArtifact.Destroy(); err != nil {
			errors = append(errors, fmt.Errorf("Error destroying builder artifact: %s", err))
		}
	}

	if len(errors) > 0 {
		err = &MultiError{errors}
	}

	return artifacts, err
}

func (b *coreBuild) SetDebug(val bool) {
	if b.prepareCalled {
		panic("prepare has already been called")
	}

	b.debug = val
}

// Cancels the build if it is running.
func (b *coreBuild) Cancel() {
	b.builder.Cancel()
}
<commit_msg>packer: Post-process chain works properly<commit_after>package packer

import (
	"fmt"
	"log"
	"sync"
)

// This is the key in configurations that is set to "true" when Packer
// debugging is enabled.
const DebugConfigKey = "packer_debug"

// A Build represents a single job within Packer that is responsible for
// building some machine image artifact. Builds are meant to be parallelized.
type Build interface {
	// Name is the name of the build. This is unique across a single template,
	// but not absolutely unique. This is meant more to describe to the user
	// what is being built rather than being a unique identifier.
	Name() string

	// Prepare configures the various components of this build and reports
	// any errors in doing so (such as syntax errors, validation errors, etc.)
	Prepare() error

	// Run runs the actual builder, returning an artifact implementation
	// of what is built. If anything goes wrong, an error is returned.
	Run(Ui, Cache) ([]Artifact, error)

	// Cancel will cancel a running build. This will block until the build
	// is actually completely cancelled.
	Cancel()

	// SetDebug will enable/disable debug mode. Debug mode is always
	// enabled by adding the additional key "packer_debug" to boolean
	// true in the configuration of the various components. This must
	// be called prior to Prepare.
	//
	// When SetDebug is set to true, parallelism between builds is
	// strictly prohibited.
	SetDebug(bool)
}

// A build struct represents a single build job, the result of which should
// be a single machine image artifact. This artifact may be comprised of
// multiple files, of course, but it should be for only a single provider
// (such as VirtualBox, EC2, etc.).
type coreBuild struct {
	name           string
	builder        Builder
	builderConfig  interface{}
	hooks          map[string][]Hook
	postProcessors [][]coreBuildPostProcessor
	provisioners   []coreBuildProvisioner

	debug         bool
	l             sync.Mutex
	prepareCalled bool
}

// Keeps track of the post-processor and the configuration of the
// post-processor used within a build.
type coreBuildPostProcessor struct {
	processor         PostProcessor
	config            interface{}
	keepInputArtifact bool
}

// Keeps track of the provisioner and the configuration of the provisioner
// within the build.
type coreBuildProvisioner struct {
	provisioner Provisioner
	config      []interface{}
}

// Returns the name of the build.
func (b *coreBuild) Name() string {
	return b.name
}

// Prepare prepares the build by doing some initialization for the builder
// and any hooks. This _must_ be called prior to Run.
func (b *coreBuild) Prepare() (err error) {
	b.l.Lock()
	defer b.l.Unlock()

	if b.prepareCalled {
		panic("prepare already called")
	}

	b.prepareCalled = true

	debugConfig := map[string]interface{}{
		DebugConfigKey: b.debug,
	}

	// Prepare the builder
	err = b.builder.Prepare(b.builderConfig, debugConfig)
	if err != nil {
		log.Printf("Build '%s' prepare failure: %s\n", b.name, err)
		return
	}

	// Prepare the provisioners
	for _, coreProv := range b.provisioners {
		configs := make([]interface{}, len(coreProv.config), len(coreProv.config)+1)
		copy(configs, coreProv.config)
		configs = append(configs, debugConfig)

		if err = coreProv.provisioner.Prepare(configs...); err != nil {
			return
		}
	}

	// Prepare the post-processors
	for _, ppSeq := range b.postProcessors {
		for _, corePP := range ppSeq {
			if err = corePP.processor.Configure(corePP.config); err != nil {
				return
			}
		}
	}

	return
}

// Runs the actual build. Prepare must be called prior to running this.
func (b *coreBuild) Run(ui Ui, cache Cache) ([]Artifact, error) {
	if !b.prepareCalled {
		panic("Prepare must be called first")
	}

	// Copy the hooks
	hooks := make(map[string][]Hook)
	for hookName, hookList := range b.hooks {
		hooks[hookName] = make([]Hook, len(hookList))
		copy(hooks[hookName], hookList)
	}

	// Add a hook for the provisioners if we have provisioners
	if len(b.provisioners) > 0 {
		provisioners := make([]Provisioner, len(b.provisioners))
		for i, p := range b.provisioners {
			provisioners[i] = p.provisioner
		}

		if _, ok := hooks[HookProvision]; !ok {
			hooks[HookProvision] = make([]Hook, 0, 1)
		}

		hooks[HookProvision] = append(hooks[HookProvision], &ProvisionHook{provisioners})
	}

	hook := &DispatchHook{hooks}
	artifacts := make([]Artifact, 0, 1)

	builderArtifact, err := b.builder.Run(ui, hook, cache)
	if err != nil {
		return nil, err
	}

	errors := make([]error, 0)
	keepOriginalArtifact := len(b.postProcessors) == 0

	// Run the post-processors
PostProcessorRunSeqLoop:
	for _, ppSeq := range b.postProcessors {
		priorArtifact := builderArtifact
		for i, corePP := range ppSeq {
			artifact, err := corePP.processor.PostProcess(ui, priorArtifact)
			if err != nil {
				errors = append(errors, fmt.Errorf("Post-processor failed: %s", err))
				continue PostProcessorRunSeqLoop
			}

			if artifact == nil {
				log.Println("Nil artifact, halting post-processor chain.")
				continue PostProcessorRunSeqLoop
			}

			if i == 0 {
				// This is the first post-processor. We handle deleting
				// previous artifacts a bit different because multiple
				// post-processors may be using the original and need it.
				if !keepOriginalArtifact {
					keepOriginalArtifact = corePP.keepInputArtifact
				}
			} else {
				// We have a prior artifact. If we want to keep it, we append
				// it to the results list. Otherwise, we destroy it.
				if corePP.keepInputArtifact {
					artifacts = append(artifacts, priorArtifact)
				} else {
					if err := priorArtifact.Destroy(); err != nil {
						errors = append(errors, fmt.Errorf("Failed cleaning up prior artifact: %s", err))
					}
				}
			}

			priorArtifact = artifact
		}

		// Add on the last artifact to the results
		if priorArtifact != nil {
			artifacts = append(artifacts, priorArtifact)
		}
	}

	if keepOriginalArtifact {
		artifacts = append(artifacts, nil)
		copy(artifacts[1:], artifacts)
		artifacts[0] = builderArtifact
	} else {
		log.Printf("Deleting original artifact for build '%s'", b.name)
		if err := builderArtifact.Destroy(); err != nil {
			errors = append(errors, fmt.Errorf("Error destroying builder artifact: %s", err))
		}
	}

	if len(errors) > 0 {
		err = &MultiError{errors}
	}

	return artifacts, err
}

func (b *coreBuild) SetDebug(val bool) {
	if b.prepareCalled {
		panic("prepare has already been called")
	}

	b.debug = val
}

// Cancels the build if it is running.
func (b *coreBuild) Cancel() {
	b.builder.Cancel()
}
<|endoftext|>


<commit_before>package storage

import (
	"errors"
	"net/http"
	"net/url"
	"strconv"
	"time"
)

const (
	authURL          = "https://auth.selcdn.ru/"
	authUserHeader   = "X-Auth-User"
	authKeyHeader    = "X-Auth-Key"
	authExpireHeader = "X-Expire-Auth-Token"
	storageURLHeader = "X-Storage-Url"
	// tokenDurationAdd used to reduce duration of token
	// to re-auth before token gets expired
	tokenDurationAdd = 10 * time.Second
)

var (
	// ErrorAuth occurs when client is unable to authenticate
	ErrorAuth = errors.New("Authentication error")
	// ErrorBadCredentials occurs when incorrect user/key provided
	ErrorBadCredentials = errors.New("Bad auth credentials provided")
)

// Token returns current auth token
func (c *Client) Token() string {
	return c.token
}

// Auth performs authentication to selectel and stores token and storage url
func (c *Client) Auth(user, key string) error {
	if blank(user) || blank(key) {
		return ErrorBadCredentials
	}

	request, _ := http.NewRequest(getMethod, authURL, nil)
	request.Header.Add(authUserHeader, user)
	request.Header.Add(authKeyHeader, key)

	res, err := c.do(request)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusNoContent {
		return ErrorAuth
	}
	expire, err := strconv.Atoi(res.Header.Get(authExpireHeader))
	if err != nil {
		return err
	}

	c.tokenExpire = expire
	c.token = res.Header.Get(authTokenHeader)
	if blank(c.token) {
		return ErrorAuth
	}
	c.storageURL, err = url.Parse(res.Header.Get(storageURLHeader))
	if err != nil || blank(c.storageURL.String()) {
		return ErrorAuth
	}

	c.user, c.key = user, key
	now := time.Now()
	c.expireFrom = &now

	return nil
}

// Expired returns true if token is expired or does not exist
func (c *Client) Expired() bool {
	if c.expireFrom == nil {
		return true
	}
	duration := time.Duration(c.tokenExpire) * time.Second
	expiredFrom := c.expireFrom.Add(duration).Add(tokenDurationAdd)
	return expiredFrom.Before(time.Now())
}
<commit_msg>added condition to "token does not exists"<commit_after>package storage

import (
	"errors"
	"net/http"
	"net/url"
	"strconv"
	"time"
)

const (
	authURL          = "https://auth.selcdn.ru/"
	authUserHeader   = "X-Auth-User"
	authKeyHeader    = "X-Auth-Key"
	authExpireHeader = "X-Expire-Auth-Token"
	storageURLHeader = "X-Storage-Url"
	// tokenDurationAdd used to reduce duration of token
	// to re-auth before token gets expired
	tokenDurationAdd = 10 * time.Second
)

var (
	// ErrorAuth occurs when client is unable to authenticate
	ErrorAuth = errors.New("Authentication error")
	// ErrorBadCredentials occurs when incorrect user/key provided
	ErrorBadCredentials = errors.New("Bad auth credentials provided")
)

// Token returns current auth token
func (c *Client) Token() string {
	return c.token
}

// Auth performs authentication to selectel and stores token and storage url
func (c *Client) Auth(user, key string) error {
	if blank(user) || blank(key) {
		return ErrorBadCredentials
	}

	request, _ := http.NewRequest(getMethod, authURL, nil)
	request.Header.Add(authUserHeader, user)
	request.Header.Add(authKeyHeader, key)

	res, err := c.do(request)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusNoContent {
		return ErrorAuth
	}
	expire, err := strconv.Atoi(res.Header.Get(authExpireHeader))
	if err != nil {
		return err
	}

	c.tokenExpire = expire
	c.token = res.Header.Get(authTokenHeader)
	if blank(c.token) {
		return ErrorAuth
	}
	c.storageURL, err = url.Parse(res.Header.Get(storageURLHeader))
	if err != nil || blank(c.storageURL.String()) {
		return ErrorAuth
	}

	c.user, c.key = user, key
	now := time.Now()
	c.expireFrom = &now

	return nil
}

// Expired returns true if token is expired or does not exist
func (c *Client) Expired() bool {
	if c.expireFrom == nil || blank(c.token) {
		return true
	}
	duration := time.Duration(c.tokenExpire) * time.Second
	expiredFrom := c.expireFrom.Add(duration).Add(tokenDurationAdd)
	return expiredFrom.Before(time.Now())
}
<|endoftext|>


<commit_before>package gridscale

import (
	"os"
	"testing"
	"time"
)

func TestCreateRenameDeleteStorage(t *testing.T) {
	userID := os.Getenv("GRIDSCALE_USERID")
	apiToken := os.Getenv("GRIDSCALE_APITOKEN")
	endpoint := "https://api.gridscale.io"
	c, _ := NewClient(userID, apiToken, endpoint)

	s, err := c.CreateStorage("45ed677b-3702-4b36-be2a-a2eab9827950", "foobar", 10, nil, nil)
	if err != nil {
		t.Errorf("Could not create storage: %s", err)
		return
	}

	defer c.DeleteStorage(s.ID)

	for s.Status == "in-provisioning" {
		time.Sleep(time.Second)
		s, err = c.GetStorage(s.ID)
	}
	if s.Status != "active" {
		t.Errorf("Unexpected storage status after creation: %s expected %s", s.Status, "active")
	}

	err = c.UpdateStorageName(s.ID, "bambaz")
	if err != nil {
		t.Errorf("Could not update storage name (%s): %s", s.ID, err)
		return
	}

	strgs, err := c.GetStorages()
	if err != nil {
		t.Errorf("Could not get list of storages: %s\n", err)
		return
	}

	found := false
	for _, strg := range strgs {
		if strg.ID == s.ID {
			found = true
			if strg.Name != "bambaz" {
				t.Errorf("Changing storage name from 'foobar' to 'bambaz' did not work!")
				return
			}
		}
	}
	if !found {
		t.Errorf("Could not find a newly created storage '%s' in returned storage list!", s.ID)
		return
	}

	for s.Status == "in-provisioning" {
		time.Sleep(time.Second)
		s, err = c.GetStorage(s.ID)
	}

	err = c.DeleteStorage(s.ID)
	if err != nil {
		t.Errorf("Could not delete storage: %s", err)
		return
	}

	strgs, err = c.GetStorages()
	if err != nil {
		t.Errorf("Could not get list of storage: %s\n", err)
		return
	}

	for _, strg := range strgs {
		if strg.ID == s.ID && strg.Status != "to-be-deleted" {
			t.Errorf("Could not delete storage %s: %s", s.ID, strg.Status)
		}
	}
}
<commit_msg>fixing test<commit_after>package gridscale

import (
	"os"
	"testing"
	"time"
)

func TestCreateRenameDeleteStorage(t *testing.T) {
	userID := os.Getenv("GRIDSCALE_USERID")
	apiToken := os.Getenv("GRIDSCALE_APITOKEN")
	endpoint := "https://api.gridscale.io"
	c, _ := NewClient(userID, apiToken, endpoint)

	s, err := c.CreateStorage("45ed677b-3702-4b36-be2a-a2eab9827950", "foobar", 10, nil, nil)
	if err != nil {
		t.Errorf("Could not create storage: %s", err)
		return
	}

	defer c.DeleteStorage(s.ID)

	for s.Status == "in-provisioning" {
		time.Sleep(time.Second)
		s, err = c.GetStorage(s.ID)
	}
	if s.Status != "active" {
		t.Errorf("Unexpected storage status after creation: %s expected %s", s.Status, "active")
	}

	err = c.UpdateStorageName(s.ID, "bambaz")
	if err != nil {
		t.Errorf("Could not update storage name (%s): %s", s.ID, err)
		return
	}

	strgs, err := c.GetStorages()
	if err != nil {
		t.Errorf("Could not get list of storages: %s\n", err)
		return
	}

	found := false
	for _, strg := range strgs {
		if strg.ID == s.ID {
			found = true
			if strg.Name != "bambaz" {
				t.Errorf("Changing storage name from 'foobar' to 'bambaz' did not work!")
				return
			}
		}
	}
	if !found {
		t.Errorf("Could not find a newly created storage '%s' in returned storage list!", s.ID)
		return
	}

	s, err = c.GetStorage(s.ID)
	for s.Status == "in-provisioning" {
		time.Sleep(time.Second)
		s, err = c.GetStorage(s.ID)
	}

	err = c.DeleteStorage(s.ID)
	if err != nil {
		t.Errorf("Could not delete storage: %s", err)
		return
	}

	strgs, err = c.GetStorages()
	if err != nil {
		t.Errorf("Could not get list of storage: %s\n", err)
		return
	}

	for _, strg := range strgs {
		if strg.ID == s.ID && strg.Status != "to-be-deleted" {
			t.Errorf("Could not delete storage %s: %s", s.ID, strg.Status)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"testing"
)

func TestRandomID(t *testing.T) {
	for _, c := range []struct {
		available func(ID) bool
		wantErr   bool
	}{
		{func(id ID) bool { return true; }, false},
		{func(id ID) bool { return false; }, true},
	} {
		_, err := randomID(c.available)
		if c.wantErr {
			if err == nil {
				t.Errorf(`randomID() didn't error as expected`)
			}
		} else if err != nil {
			t.Errorf(`randomID() errored unexpectedly`)
		}
	}
}
<commit_msg>Run go fmt<commit_after>package main

import (
	"testing"
)

func TestRandomID(t *testing.T) {
	for _, c := range []struct {
		available func(ID) bool
		wantErr   bool
	}{
		{func(id ID) bool { return true }, false},
		{func(id ID) bool { return false }, true},
	} {
		_, err := randomID(c.available)
		if c.wantErr {
			if err == nil {
				t.Errorf(`randomID() didn't error as expected`)
			}
		} else if err != nil {
			t.Errorf(`randomID() errored unexpectedly`)
		}
	}
}
<|endoftext|>


<commit_before>// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// Copyright 2013 The Go-MySQL-Driver Authors. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package mysql

import (
	"fmt"
	"io"
	"os"
	"path"
	"strings"
)

var (
	fileRegister   map[string]bool
	readerRegister map[string]func() io.Reader
)

// RegisterLocalFile adds the given file to the file whitelist,
// so that it can be used by "LOAD DATA LOCAL INFILE <filepath>".
// Alternatively you can allow the use of all local files with
// the DSN parameter 'allowAllFiles=true'
//
//  filePath := "/home/gopher/data.csv"
//  mysql.RegisterLocalFile(filePath)
//  err := db.Exec("LOAD DATA LOCAL INFILE '" + filePath + "' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterLocalFile(filePath string) {
	// lazy map init
	if fileRegister == nil {
		fileRegister = make(map[string]bool)
	}

	fileRegister[strings.Trim(filePath, `"`)] = true
}

// DeregisterLocalFile removes the given filepath from the whitelist.
func DeregisterLocalFile(filePath string) {
	delete(fileRegister, strings.Trim(filePath, `"`))
}

// RegisterReaderHandler registers a handler function which is used
// to receive a io.Reader.
// The Reader can be used by "LOAD DATA LOCAL INFILE Reader::<name>".
// If the handler returns a io.ReadCloser Close() is called when the
// request is finished.
//
//  mysql.RegisterReaderHandler("data", func() io.Reader {
//  	var csvReader io.Reader // Some Reader that returns CSV data
//  	... // Open Reader here
//  	return csvReader
//  })
//  err := db.Exec("LOAD DATA LOCAL INFILE 'Reader::data' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterReaderHandler(name string, handler func() io.Reader) {
	// lazy map init
	if readerRegister == nil {
		readerRegister = make(map[string]func() io.Reader)
	}

	readerRegister[name] = handler
}

// DeregisterReaderHandler removes the ReaderHandler function with
// the given name from the registry.
func DeregisterReaderHandler(name string) {
	delete(readerRegister, name)
}

func deferredClose(err *error, closer io.Closer) {
	closeErr := closer.Close()
	if *err == nil {
		*err = closeErr
	}
}

func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
	var rdr io.Reader
	var data []byte

	// The server might return an an absolute path. See issue #355.
	base := path.Base(name)

	if strings.HasPrefix(base, "Reader::") { // io.Reader
		name = base[8:]
		if handler, inMap := readerRegister[name]; inMap {
			rdr = handler()
			if rdr != nil {
				data = make([]byte, 4+mc.maxWriteSize)

				if cl, ok := rdr.(io.Closer); ok {
					defer deferredClose(&err, cl)
				}
			} else {
				err = fmt.Errorf("Reader '%s' is <nil>", name)
			}
		} else {
			err = fmt.Errorf("Reader '%s' is not registered", name)
		}
	} else { // File
		name = strings.Trim(name, `"`)
		if mc.cfg.allowAllFiles || fileRegister[name] {
			var file *os.File
			var fi os.FileInfo

			if file, err = os.Open(name); err == nil {
				defer deferredClose(&err, file)

				// get file size
				if fi, err = file.Stat(); err == nil {
					rdr = file
					if fileSize := int(fi.Size()); fileSize <= mc.maxWriteSize {
						data = make([]byte, 4+fileSize)
					} else if fileSize <= mc.maxPacketAllowed {
						data = make([]byte, 4+mc.maxWriteSize)
					} else {
						err = fmt.Errorf("Local File '%s' too large: Size: %d, Max: %d", name, fileSize, mc.maxPacketAllowed)
					}
				}
			}
		} else {
			err = fmt.Errorf("Local File '%s' is not registered. Use the DSN parameter 'allowAllFiles=true' to allow all files", name)
		}
	}

	// send content packets
	if err == nil {
		var n int
		for err == nil {
			n, err = rdr.Read(data[4:])
			if n > 0 {
				if ioErr := mc.writePacket(data[:4+n]); ioErr != nil {
					return ioErr
				}
			}
		}
		if err == io.EOF {
			err = nil
		}
	}

	// send empty packet (termination)
	if data == nil {
		data = make([]byte, 4)
	}
	if ioErr := mc.writePacket(data[:4]); ioErr != nil {
		return ioErr
	}

	// read OK packet
	if err == nil {
		return mc.readResultOK()
	} else {
		mc.readPacket()
	}
	return err
}
<commit_msg>infile: Allow '/' in Reader name<commit_after>// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// Copyright 2013 The Go-MySQL-Driver Authors. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package mysql

import (
	"fmt"
	"io"
	"os"
	"strings"
)

var (
	fileRegister   map[string]bool
	readerRegister map[string]func() io.Reader
)

// RegisterLocalFile adds the given file to the file whitelist,
// so that it can be used by "LOAD DATA LOCAL INFILE <filepath>".
// Alternatively you can allow the use of all local files with
// the DSN parameter 'allowAllFiles=true'
//
//  filePath := "/home/gopher/data.csv"
//  mysql.RegisterLocalFile(filePath)
//  err := db.Exec("LOAD DATA LOCAL INFILE '" + filePath + "' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterLocalFile(filePath string) {
	// lazy map init
	if fileRegister == nil {
		fileRegister = make(map[string]bool)
	}

	fileRegister[strings.Trim(filePath, `"`)] = true
}

// DeregisterLocalFile removes the given filepath from the whitelist.
func DeregisterLocalFile(filePath string) {
	delete(fileRegister, strings.Trim(filePath, `"`))
}

// RegisterReaderHandler registers a handler function which is used
// to receive a io.Reader.
// The Reader can be used by "LOAD DATA LOCAL INFILE Reader::<name>".
// If the handler returns a io.ReadCloser Close() is called when the
// request is finished.
//
//  mysql.RegisterReaderHandler("data", func() io.Reader {
//  	var csvReader io.Reader // Some Reader that returns CSV data
//  	... // Open Reader here
//  	return csvReader
//  })
//  err := db.Exec("LOAD DATA LOCAL INFILE 'Reader::data' INTO TABLE foo")
//  if err != nil {
//  ...
//
func RegisterReaderHandler(name string, handler func() io.Reader) {
	// lazy map init
	if readerRegister == nil {
		readerRegister = make(map[string]func() io.Reader)
	}

	readerRegister[name] = handler
}

// DeregisterReaderHandler removes the ReaderHandler function with
// the given name from the registry.
func DeregisterReaderHandler(name string) {
	delete(readerRegister, name)
}

func deferredClose(err *error, closer io.Closer) {
	closeErr := closer.Close()
	if *err == nil {
		*err = closeErr
	}
}

func (mc *mysqlConn) handleInFileRequest(name string) (err error) {
	var rdr io.Reader
	var data []byte

	if idx := strings.Index(name, "Reader::"); idx == 0 || (idx > 0 && name[idx-1] == '/') { // io.Reader
		// The server might return an an absolute path. See issue #355.
		name = name[idx+8:]

		if handler, inMap := readerRegister[name]; inMap {
			rdr = handler()
			if rdr != nil {
				data = make([]byte, 4+mc.maxWriteSize)

				if cl, ok := rdr.(io.Closer); ok {
					defer deferredClose(&err, cl)
				}
			} else {
				err = fmt.Errorf("Reader '%s' is <nil>", name)
			}
		} else {
			err = fmt.Errorf("Reader '%s' is not registered", name)
		}
	} else { // File
		name = strings.Trim(name, `"`)
		if mc.cfg.allowAllFiles || fileRegister[name] {
			var file *os.File
			var fi os.FileInfo

			if file, err = os.Open(name); err == nil {
				defer deferredClose(&err, file)

				// get file size
				if fi, err = file.Stat(); err == nil {
					rdr = file
					if fileSize := int(fi.Size()); fileSize <= mc.maxWriteSize {
						data = make([]byte, 4+fileSize)
					} else if fileSize <= mc.maxPacketAllowed {
						data = make([]byte, 4+mc.maxWriteSize)
					} else {
						err = fmt.Errorf("Local File '%s' too large: Size: %d, Max: %d", name, fileSize, mc.maxPacketAllowed)
					}
				}
			}
		} else {
			err = fmt.Errorf("Local File '%s' is not registered. Use the DSN parameter 'allowAllFiles=true' to allow all files", name)
		}
	}

	// send content packets
	if err == nil {
		var n int
		for err == nil {
			n, err = rdr.Read(data[4:])
			if n > 0 {
				if ioErr := mc.writePacket(data[:4+n]); ioErr != nil {
					return ioErr
				}
			}
		}
		if err == io.EOF {
			err = nil
		}
	}

	// send empty packet (termination)
	if data == nil {
		data = make([]byte, 4)
	}
	if ioErr := mc.writePacket(data[:4]); ioErr != nil {
		return ioErr
	}

	// read OK packet
	if err == nil {
		return mc.readResultOK()
	} else {
		mc.readPacket()
	}
	return err
}
<|endoftext|>