<commit_before>// Package camoproxy provides an HTTP proxy server with content type
// restrictions as well as regex host allow list support.
package camoproxy

import (
	"github.com/gorilla/mux"
	"errors"
	"github.com/cactus/go-camo/camoproxy/encoding"
	"github.com/cactus/gologit"
	"io"
	"net"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"syscall"
	"time"
)

// Config holds configuration data used when creating a Proxy with New.
type Config struct {
	// HmacKey is a string to be used as the hmac key
	HmacKey           string
	// AllowList is a list of string represenstations of regex (not compiled
	// regex) that are used as a whitelist filter. If an AllowList is present,
	// then anything not matching is dropped. If no AllowList is present,
	// no Allow filtering is done.
	AllowList         []string
	// MaxSize is the maximum valid image size response (in bytes).
	MaxSize           int64
	// NoFollowRedirects is a boolean that specifies whether upstream redirects
	// are followed (10 depth) or not.
	NoFollowRedirects bool
	// Request timeout is a timeout for fetching upstream data.
	RequestTimeout    time.Duration
}

// Interface for Proxy to use for stats/metrics.
// This must be goroutine safe, as AddBytes and AddServed will be called from
// many goroutines.
type ProxyMetrics interface {
	AddBytes(bc int64)
	AddServed()
}

// A Proxy is a Camo like HTTP proxy, that provides content type
// restrictions as well as regex host allow list support.
type Proxy struct {
	client    *http.Client
	hmacKey   []byte
	allowList []*regexp.Regexp
	maxSize   int64
	metrics   ProxyMetrics
}

// ServerHTTP handles the client request, validates the request is validly
// HMAC signed, filters based on the Allow list, and then proxies
// valid requests to the desired endpoint. Responses are filtered for
// proper image content types.
func (p *Proxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	gologit.Debugln("Request:", req.URL)
	if p.metrics != nil {
		go p.metrics.AddServed()
	}

	w.Header().Set("Server", ServerNameVer)

	vars := mux.Vars(req)
	surl, ok := encoding.DecodeUrl(&p.hmacKey, vars["sigHash"], vars["encodedUrl"])
	if !ok {
		http.Error(w, "Bad Signature", http.StatusForbidden)
		return
	}
	gologit.Debugln("URL:", surl)

	u, err := url.Parse(surl)
	if err != nil {
		gologit.Debugln(err)
		http.Error(w, "Bad url", http.StatusBadRequest)
		return
	}

	u.Host = strings.ToLower(u.Host)
	if u.Host == "" || localhostRegex.MatchString(u.Host) {
		http.Error(w, "Bad url", http.StatusNotFound)
		return
	}

	if req.Header.Get("Via") == "ServerNameVer" {
		http.Error(w, "Request loop failure", http.StatusNotFound)
	}

	// if allowList is set, require match
	matchFound := true
	if len(p.allowList) > 0 {
		matchFound = false
		for _, rgx := range p.allowList {
			if rgx.MatchString(u.Host) {
				matchFound = true
			}
		}
	}
	if !matchFound {
		http.Error(w, "Allowlist host failure", http.StatusNotFound)
		return
	}

	// filter out rfc1918 hosts
	ip := net.ParseIP(u.Host)
	if ip != nil {
		if addr1918PrefixRegex.MatchString(ip.String()) {
			http.Error(w, "Denylist host failure", http.StatusNotFound)
			return
		}
	}

	nreq, err := http.NewRequest("GET", surl, nil)
	if err != nil {
		gologit.Debugln("Could not create NewRequest", err)
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	}

	// filter headers
	p.copyHeader(&nreq.Header, &req.Header, &ValidReqHeaders)
	if req.Header.Get("X-Forwarded-For") == "" {
		host, _, err := net.SplitHostPort(req.RemoteAddr)
		if err == nil && !addr1918PrefixRegex.MatchString(host) {
			nreq.Header.Add("X-Forwarded-For", host)
		}
	}

	// add an accept header if the client didn't send one
	if nreq.Header.Get("Accept") == "" {
		nreq.Header.Add("Accept", "image/*")
	}

	nreq.Header.Add("connection", "close")
	nreq.Header.Add("user-agent", ServerNameVer)
	nreq.Header.Add("via", ServerNameVer)

	resp, err := p.client.Do(nreq)
	if err != nil {
		gologit.Debugln("Could not connect to endpoint", err)
		if strings.Contains(err.Error(), "timeout") {
			http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		} else {
			http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		}
		return
	}
	defer resp.Body.Close()

	// check for too large a response
	if resp.ContentLength > p.maxSize {
		gologit.Debugln("Content length exceeded", surl)
		http.Error(w, "Content length exceeded", http.StatusNotFound)
		return
	}

	switch resp.StatusCode {
	case 200:
		// check content type
		ct, ok := resp.Header[http.CanonicalHeaderKey("content-type")]
		if !ok || ct[0][:6] != "image/" {
			gologit.Debugln("Non-Image content-type returned", u)
			http.Error(w, "Non-Image content-type returned",
				http.StatusBadRequest)
			return
		}
	case 300:
		gologit.Debugln("Multiple choices not supported")
		http.Error(w, "Multiple choices not supported", http.StatusNotFound)
		return
	case 301, 302, 303, 307:
		// if we get a redirect here, we either disabled following,
		// or followed until max depth and still got one (redirect loop)
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 304:
		h := w.Header()
		p.copyHeader(&h, &resp.Header, &ValidRespHeaders)
		h.Set("X-Content-Type-Options", "nosniff")
		w.WriteHeader(304)
		return
	case 404:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 500, 502, 503, 504:
		// upstream errors should probably just 502. client can try later.
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	default:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}

	h := w.Header()
	p.copyHeader(&h, &resp.Header, &ValidRespHeaders)
	h.Set("X-Content-Type-Options", "nosniff")
	h.Set("Date", formattedDate.String())
	w.WriteHeader(resp.StatusCode)

	// since this uses io.Copy from the respBody, it is streaming
	// from the request to the response. This means it will nearly
	// always end up with a chunked response.
	// Change to the following to send whole body at once, and
	// read whole body at once too:
	//    body, err := ioutil.ReadAll(resp.Body)
	//    if err != nil {
	//        gologit.Println("Error writing response:", err)
	//    }
	//    w.Write(body)
	// Might use quite a bit of memory though. Untested.
	bW, err := io.Copy(w, resp.Body)
	if err != nil {
		// only log if not broken pipe. broken pipe means the client
		// terminated conn for some reason.
		opErr, ok := err.(*net.OpError)
		if !ok || opErr.Err != syscall.EPIPE {
			gologit.Println("Error writing response:", err)
		}
		return
	}

	if p.metrics != nil {
		go p.metrics.AddBytes(bW)
	}
	gologit.Debugln(req, resp.StatusCode)
}

// copy headers from src into dst
// empty filter map will result in no filtering being done
func (p *Proxy) copyHeader(dst, src *http.Header, filter *map[string]bool) {
	f := *filter
	filtering := false
	if len(f) > 0 {
		filtering = true
	}

	for k, vv := range *src {
		if x, ok := f[k]; filtering && (!ok || !x) {
			continue
		}
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// sets a proxy metrics (ProxyMetrics interface) for the proxy
func (p *Proxy) SetMetricsCollector(pm ProxyMetrics) {
	p.metrics = pm
}

// Returns a new Proxy. An error is returned if there was a failure
// to parse the regex from the passed Config.
func New(pc Config) (*Proxy, error) {
	tr := &http.Transport{
		Dial: func(netw, addr string) (net.Conn, error) {
			c, err := net.DialTimeout(netw, addr, pc.RequestTimeout)
			if err != nil {
				return nil, err
			}
			// also set time limit on reading
			c.SetDeadline(time.Now().Add(pc.RequestTimeout))
			return c, nil
		}}

	// spawn an idle conn trimmer
	go func() {
		// prunes every 5 minutes. this is just a guess at an
		// initial value. very busy severs may want to lower this...
		for {
			time.Sleep(5 * time.Minute)
			tr.CloseIdleConnections()
		}
	}()

	// build/compile regex
	client := &http.Client{Transport: tr}
	if pc.NoFollowRedirects {
		client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
			return errors.New("Not following redirect")
		}
	}

	allow := make([]*regexp.Regexp, 0)
	var c *regexp.Regexp
	var err error
	// compile allow list
	for _, v := range pc.AllowList {
		c, err = regexp.Compile(v)
		if err != nil {
			return nil, err
		}
		allow = append(allow, c)
	}

	return &Proxy{
		client:    client,
		hmacKey:   []byte(pc.HmacKey),
		allowList: allow,
		maxSize:   pc.MaxSize}, nil
}
<commit_msg>dump max idle conns per server<commit_after>// Package camoproxy provides an HTTP proxy server with content type
// restrictions as well as regex host allow list support.
package camoproxy

import (
	"github.com/gorilla/mux"
	"errors"
	"github.com/cactus/go-camo/camoproxy/encoding"
	"github.com/cactus/gologit"
	"io"
	"net"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"syscall"
	"time"
)

// Config holds configuration data used when creating a Proxy with New.
type Config struct {
	// HmacKey is a string to be used as the hmac key
	HmacKey           string
	// AllowList is a list of string represenstations of regex (not compiled
	// regex) that are used as a whitelist filter. If an AllowList is present,
	// then anything not matching is dropped. If no AllowList is present,
	// no Allow filtering is done.
	AllowList         []string
	// MaxSize is the maximum valid image size response (in bytes).
	MaxSize           int64
	// NoFollowRedirects is a boolean that specifies whether upstream redirects
	// are followed (10 depth) or not.
	NoFollowRedirects bool
	// Request timeout is a timeout for fetching upstream data.
	RequestTimeout    time.Duration
}

// Interface for Proxy to use for stats/metrics.
// This must be goroutine safe, as AddBytes and AddServed will be called from
// many goroutines.
type ProxyMetrics interface {
	AddBytes(bc int64)
	AddServed()
}

// A Proxy is a Camo like HTTP proxy, that provides content type
// restrictions as well as regex host allow list support.
type Proxy struct {
	client    *http.Client
	hmacKey   []byte
	allowList []*regexp.Regexp
	maxSize   int64
	metrics   ProxyMetrics
}

// ServerHTTP handles the client request, validates the request is validly
// HMAC signed, filters based on the Allow list, and then proxies
// valid requests to the desired endpoint. Responses are filtered for
// proper image content types.
func (p *Proxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	gologit.Debugln("Request:", req.URL)
	if p.metrics != nil {
		go p.metrics.AddServed()
	}

	w.Header().Set("Server", ServerNameVer)

	vars := mux.Vars(req)
	surl, ok := encoding.DecodeUrl(&p.hmacKey, vars["sigHash"], vars["encodedUrl"])
	if !ok {
		http.Error(w, "Bad Signature", http.StatusForbidden)
		return
	}
	gologit.Debugln("URL:", surl)

	u, err := url.Parse(surl)
	if err != nil {
		gologit.Debugln(err)
		http.Error(w, "Bad url", http.StatusBadRequest)
		return
	}

	u.Host = strings.ToLower(u.Host)
	if u.Host == "" || localhostRegex.MatchString(u.Host) {
		http.Error(w, "Bad url", http.StatusNotFound)
		return
	}

	if req.Header.Get("Via") == "ServerNameVer" {
		http.Error(w, "Request loop failure", http.StatusNotFound)
	}

	// if allowList is set, require match
	matchFound := true
	if len(p.allowList) > 0 {
		matchFound = false
		for _, rgx := range p.allowList {
			if rgx.MatchString(u.Host) {
				matchFound = true
			}
		}
	}
	if !matchFound {
		http.Error(w, "Allowlist host failure", http.StatusNotFound)
		return
	}

	// filter out rfc1918 hosts
	ip := net.ParseIP(u.Host)
	if ip != nil {
		if addr1918PrefixRegex.MatchString(ip.String()) {
			http.Error(w, "Denylist host failure", http.StatusNotFound)
			return
		}
	}

	nreq, err := http.NewRequest("GET", surl, nil)
	if err != nil {
		gologit.Debugln("Could not create NewRequest", err)
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	}

	// filter headers
	p.copyHeader(&nreq.Header, &req.Header, &ValidReqHeaders)
	if req.Header.Get("X-Forwarded-For") == "" {
		host, _, err := net.SplitHostPort(req.RemoteAddr)
		if err == nil && !addr1918PrefixRegex.MatchString(host) {
			nreq.Header.Add("X-Forwarded-For", host)
		}
	}

	// add an accept header if the client didn't send one
	if nreq.Header.Get("Accept") == "" {
		nreq.Header.Add("Accept", "image/*")
	}

	nreq.Header.Add("connection", "close")
	nreq.Header.Add("user-agent", ServerNameVer)
	nreq.Header.Add("via", ServerNameVer)

	resp, err := p.client.Do(nreq)
	if err != nil {
		gologit.Debugln("Could not connect to endpoint", err)
		if strings.Contains(err.Error(), "timeout") {
			http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		} else {
			http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		}
		return
	}
	defer resp.Body.Close()

	// check for too large a response
	if resp.ContentLength > p.maxSize {
		gologit.Debugln("Content length exceeded", surl)
		http.Error(w, "Content length exceeded", http.StatusNotFound)
		return
	}

	switch resp.StatusCode {
	case 200:
		// check content type
		ct, ok := resp.Header[http.CanonicalHeaderKey("content-type")]
		if !ok || ct[0][:6] != "image/" {
			gologit.Debugln("Non-Image content-type returned", u)
			http.Error(w, "Non-Image content-type returned",
				http.StatusBadRequest)
			return
		}
	case 300:
		gologit.Debugln("Multiple choices not supported")
		http.Error(w, "Multiple choices not supported", http.StatusNotFound)
		return
	case 301, 302, 303, 307:
		// if we get a redirect here, we either disabled following,
		// or followed until max depth and still got one (redirect loop)
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 304:
		h := w.Header()
		p.copyHeader(&h, &resp.Header, &ValidRespHeaders)
		h.Set("X-Content-Type-Options", "nosniff")
		w.WriteHeader(304)
		return
	case 404:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 500, 502, 503, 504:
		// upstream errors should probably just 502. client can try later.
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	default:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}

	h := w.Header()
	p.copyHeader(&h, &resp.Header, &ValidRespHeaders)
	h.Set("X-Content-Type-Options", "nosniff")
	h.Set("Date", formattedDate.String())
	w.WriteHeader(resp.StatusCode)

	// since this uses io.Copy from the respBody, it is streaming
	// from the request to the response. This means it will nearly
	// always end up with a chunked response.
	// Change to the following to send whole body at once, and
	// read whole body at once too:
	//    body, err := ioutil.ReadAll(resp.Body)
	//    if err != nil {
	//        gologit.Println("Error writing response:", err)
	//    }
	//    w.Write(body)
	// Might use quite a bit of memory though. Untested.
	bW, err := io.Copy(w, resp.Body)
	if err != nil {
		// only log if not broken pipe. broken pipe means the client
		// terminated conn for some reason.
		opErr, ok := err.(*net.OpError)
		if !ok || opErr.Err != syscall.EPIPE {
			gologit.Println("Error writing response:", err)
		}
		return
	}

	if p.metrics != nil {
		go p.metrics.AddBytes(bW)
	}
	gologit.Debugln(req, resp.StatusCode)
}

// copy headers from src into dst
// empty filter map will result in no filtering being done
func (p *Proxy) copyHeader(dst, src *http.Header, filter *map[string]bool) {
	f := *filter
	filtering := false
	if len(f) > 0 {
		filtering = true
	}

	for k, vv := range *src {
		if x, ok := f[k]; filtering && (!ok || !x) {
			continue
		}
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// sets a proxy metrics (ProxyMetrics interface) for the proxy
func (p *Proxy) SetMetricsCollector(pm ProxyMetrics) {
	p.metrics = pm
}

// Returns a new Proxy. An error is returned if there was a failure
// to parse the regex from the passed Config.
func New(pc Config) (*Proxy, error) {
	tr := &http.Transport{
		MaxIdleConnsPerHost: 8,
		Dial: func(netw, addr string) (net.Conn, error) {
			c, err := net.DialTimeout(netw, addr, pc.RequestTimeout)
			if err != nil {
				return nil, err
			}
			// also set time limit on reading
			c.SetDeadline(time.Now().Add(pc.RequestTimeout))
			return c, nil
		}}

	// spawn an idle conn trimmer
	go func() {
		// prunes every 5 minutes. this is just a guess at an
		// initial value. very busy severs may want to lower this...
		for {
			time.Sleep(5 * time.Minute)
			tr.CloseIdleConnections()
		}
	}()

	// build/compile regex
	client := &http.Client{Transport: tr}
	if pc.NoFollowRedirects {
		client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
			return errors.New("Not following redirect")
		}
	}

	allow := make([]*regexp.Regexp, 0)
	var c *regexp.Regexp
	var err error
	// compile allow list
	for _, v := range pc.AllowList {
		c, err = regexp.Compile(v)
		if err != nil {
			return nil, err
		}
		allow = append(allow, c)
	}

	return &Proxy{
		client:    client,
		hmacKey:   []byte(pc.HmacKey),
		allowList: allow,
		maxSize:   pc.MaxSize}, nil
}
<|endoftext|>


<commit_before>package warrant

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

	"github.com/pivotal-cf-experimental/warrant/internal/documents"
	"github.com/pivotal-cf-experimental/warrant/internal/network"
)

// TODO: Pagination for List

// GroupsService provides access to common group actions. Using this service,
// you can create, delete, fetch and list group resources.
type GroupsService struct {
	config Config
}

// NewGroupsService returns a GroupsService initialized with the given Config.
func NewGroupsService(config Config) GroupsService {
	return GroupsService{
		config: config,
	}
}

// Create will make a request to UAA to create a new group resource with the given
// DisplayName. A token with the "scim.write" scope is required.
func (gs GroupsService) Create(displayName, token string) (Group, error) {
	resp, err := newNetworkClient(gs.config).MakeRequest(network.Request{
		Method:        "POST",
		Path:          "/Groups",
		Authorization: network.NewTokenAuthorization(token),
		Body: network.NewJSONRequestBody(documents.CreateGroupRequest{
			DisplayName: displayName,
			Schemas:     schemas,
		}),
		AcceptableStatusCodes: []int{http.StatusCreated},
	})
	if err != nil {
		return Group{}, translateError(err)
	}

	var response documents.GroupResponse
	err = json.Unmarshal(resp.Body, &response)
	if err != nil {
		return Group{}, MalformedResponseError{err}
	}

	return newGroupFromResponse(gs.config, response), nil
}

func (gs GroupsService) AddMember(groupID, memberID, token string) error {
	_, err := newNetworkClient(gs.config).MakeRequest(network.Request{
		Method:        "POST",
		Path:          fmt.Sprintf("/Groups/%s/members", groupID),
		Authorization: network.NewTokenAuthorization(token),
		Body: network.NewJSONRequestBody(documents.Member{
			Origin: "uaa",
			Type:   "USER",
			Value:  memberID,
		}),
		AcceptableStatusCodes: []int{http.StatusCreated},
	})
	if err != nil {
		return translateError(err)
	}

	return nil
}

func (gs GroupsService) ListMembers(groupID, token string) ([]Member, error) {
	resp, err := newNetworkClient(gs.config).MakeRequest(network.Request{
		Method:                "GET",
		Path:                  fmt.Sprintf("/Groups/%s/members", groupID),
		Authorization:         network.NewTokenAuthorization(token),
		AcceptableStatusCodes: []int{http.StatusOK},
	})
	if err != nil {
		return nil, translateError(err)
	}

	var response []documents.Member
	err = json.Unmarshal(resp.Body, &response)
	if err != nil {
		return nil, MalformedResponseError{err}
	}

	var memberList []Member
	for _, m := range response {
		memberList = append(memberList, Member{
			ID: m.Value,
		})
	}

	return memberList, nil
}

func (gs GroupsService) RemoveMember(groupID, memberID, token string) error {
	_, err := newNetworkClient(gs.config).MakeRequest(network.Request{
		Method:                "DELETE",
		Path:                  fmt.Sprintf("/Groups/%s/members/%s", groupID, memberID),
		Authorization:         network.NewTokenAuthorization(token),
		AcceptableStatusCodes: []int{http.StatusOK},
	})
	if err != nil {
		return translateError(err)
	}

	return nil
}

// Get will make a request to UAA to fetch the group resource with the matching id.
// A token with the "scim.read" scope is required.
func (gs GroupsService) Get(id, token string) (Group, error) {
	resp, err := newNetworkClient(gs.config).MakeRequest(network.Request{
		Method:                "GET",
		Path:                  fmt.Sprintf("/Groups/%s", id),
		Authorization:         network.NewTokenAuthorization(token),
		AcceptableStatusCodes: []int{http.StatusOK},
	})
	if err != nil {
		return Group{}, translateError(err)
	}

	var response documents.GroupResponse
	err = json.Unmarshal(resp.Body, &response)
	if err != nil {
		return Group{}, MalformedResponseError{err}
	}

	return newGroupFromResponse(gs.config, response), nil
}

// List wil make a request to UAA to list the groups that match the given Query.
// A token with the "scim.read" scope is required.
func (gs GroupsService) List(query Query, token string) ([]Group, error) {
	requestPath := url.URL{
		Path: "/Groups",
		RawQuery: url.Values{
			"filter": []string{query.Filter},
			"sortBy": []string{query.SortBy},
		}.Encode(),
	}

	resp, err := newNetworkClient(gs.config).MakeRequest(network.Request{
		Method:                "GET",
		Path:                  requestPath.String(),
		Authorization:         network.NewTokenAuthorization(token),
		AcceptableStatusCodes: []int{http.StatusOK},
	})
	if err != nil {
		return []Group{}, translateError(err)
	}

	var response documents.GroupListResponse
	err = json.Unmarshal(resp.Body, &response)
	if err != nil {
		return []Group{}, MalformedResponseError{err}
	}

	var groupList []Group
	for _, groupResponse := range response.Resources {
		groupList = append(groupList, newGroupFromResponse(gs.config, groupResponse))
	}

	return groupList, err
}

// Delete will make a request to UAA to delete the group resource with the matching id.
// A token with the "scim.write" scope is required.
func (gs GroupsService) Delete(id, token string) error {
	_, err := newNetworkClient(gs.config).MakeRequest(network.Request{
		Method:                "DELETE",
		Path:                  fmt.Sprintf("/Groups/%s", id),
		Authorization:         network.NewTokenAuthorization(token),
		AcceptableStatusCodes: []int{http.StatusOK},
	})
	if err != nil {
		return translateError(err)
	}

	return nil
}

func newGroupFromResponse(config Config, response documents.GroupResponse) Group {
	return Group{
		ID:          response.ID,
		DisplayName: response.DisplayName,
		Version:     response.Meta.Version,
		CreatedAt:   response.Meta.Created,
		UpdatedAt:   response.Meta.LastModified,
	}
}
<commit_msg>Add comments to new group membership functions<commit_after>package warrant

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

	"github.com/pivotal-cf-experimental/warrant/internal/documents"
	"github.com/pivotal-cf-experimental/warrant/internal/network"
)

// TODO: Pagination for List

// GroupsService provides access to common group actions. Using this service,
// you can create, delete, fetch and list group resources.
type GroupsService struct {
	config Config
}

// NewGroupsService returns a GroupsService initialized with the given Config.
func NewGroupsService(config Config) GroupsService {
	return GroupsService{
		config: config,
	}
}

// Create will make a request to UAA to create a new group resource with the given
// DisplayName. A token with the "scim.write" scope is required.
func (gs GroupsService) Create(displayName, token string) (Group, error) {
	resp, err := newNetworkClient(gs.config).MakeRequest(network.Request{
		Method:        "POST",
		Path:          "/Groups",
		Authorization: network.NewTokenAuthorization(token),
		Body: network.NewJSONRequestBody(documents.CreateGroupRequest{
			DisplayName: displayName,
			Schemas:     schemas,
		}),
		AcceptableStatusCodes: []int{http.StatusCreated},
	})
	if err != nil {
		return Group{}, translateError(err)
	}

	var response documents.GroupResponse
	err = json.Unmarshal(resp.Body, &response)
	if err != nil {
		return Group{}, MalformedResponseError{err}
	}

	return newGroupFromResponse(gs.config, response), nil
}

// AddMember will make a request to UAA to add a member to the group resource with the matching id.
// A token with the "scim.write" scope is required.
func (gs GroupsService) AddMember(groupID, memberID, token string) error {
	_, err := newNetworkClient(gs.config).MakeRequest(network.Request{
		Method:        "POST",
		Path:          fmt.Sprintf("/Groups/%s/members", groupID),
		Authorization: network.NewTokenAuthorization(token),
		Body: network.NewJSONRequestBody(documents.Member{
			Origin: "uaa",
			Type:   "USER",
			Value:  memberID,
		}),
		AcceptableStatusCodes: []int{http.StatusCreated},
	})
	if err != nil {
		return translateError(err)
	}

	return nil
}

// ListMembers will make a request to UAA to fetch the members of a group resource with the matching id.
// A token with the "scim.read" scope is required.
func (gs GroupsService) ListMembers(groupID, token string) ([]Member, error) {
	resp, err := newNetworkClient(gs.config).MakeRequest(network.Request{
		Method:                "GET",
		Path:                  fmt.Sprintf("/Groups/%s/members", groupID),
		Authorization:         network.NewTokenAuthorization(token),
		AcceptableStatusCodes: []int{http.StatusOK},
	})
	if err != nil {
		return nil, translateError(err)
	}

	var response []documents.Member
	err = json.Unmarshal(resp.Body, &response)
	if err != nil {
		return nil, MalformedResponseError{err}
	}

	var memberList []Member
	for _, m := range response {
		memberList = append(memberList, Member{
			ID: m.Value,
		})
	}

	return memberList, nil
}

// RemoveMember will make a request to UAA to remove a member from a group resource.
// A token with the "scim.write" scope is required.
func (gs GroupsService) RemoveMember(groupID, memberID, token string) error {
	_, err := newNetworkClient(gs.config).MakeRequest(network.Request{
		Method:                "DELETE",
		Path:                  fmt.Sprintf("/Groups/%s/members/%s", groupID, memberID),
		Authorization:         network.NewTokenAuthorization(token),
		AcceptableStatusCodes: []int{http.StatusOK},
	})
	if err != nil {
		return translateError(err)
	}

	return nil
}

// Get will make a request to UAA to fetch the group resource with the matching id.
// A token with the "scim.read" scope is required.
func (gs GroupsService) Get(id, token string) (Group, error) {
	resp, err := newNetworkClient(gs.config).MakeRequest(network.Request{
		Method:                "GET",
		Path:                  fmt.Sprintf("/Groups/%s", id),
		Authorization:         network.NewTokenAuthorization(token),
		AcceptableStatusCodes: []int{http.StatusOK},
	})
	if err != nil {
		return Group{}, translateError(err)
	}

	var response documents.GroupResponse
	err = json.Unmarshal(resp.Body, &response)
	if err != nil {
		return Group{}, MalformedResponseError{err}
	}

	return newGroupFromResponse(gs.config, response), nil
}

// List wil make a request to UAA to list the groups that match the given Query.
// A token with the "scim.read" scope is required.
func (gs GroupsService) List(query Query, token string) ([]Group, error) {
	requestPath := url.URL{
		Path: "/Groups",
		RawQuery: url.Values{
			"filter": []string{query.Filter},
			"sortBy": []string{query.SortBy},
		}.Encode(),
	}

	resp, err := newNetworkClient(gs.config).MakeRequest(network.Request{
		Method:                "GET",
		Path:                  requestPath.String(),
		Authorization:         network.NewTokenAuthorization(token),
		AcceptableStatusCodes: []int{http.StatusOK},
	})
	if err != nil {
		return []Group{}, translateError(err)
	}

	var response documents.GroupListResponse
	err = json.Unmarshal(resp.Body, &response)
	if err != nil {
		return []Group{}, MalformedResponseError{err}
	}

	var groupList []Group
	for _, groupResponse := range response.Resources {
		groupList = append(groupList, newGroupFromResponse(gs.config, groupResponse))
	}

	return groupList, err
}

// Delete will make a request to UAA to delete the group resource with the matching id.
// A token with the "scim.write" scope is required.
func (gs GroupsService) Delete(id, token string) error {
	_, err := newNetworkClient(gs.config).MakeRequest(network.Request{
		Method:                "DELETE",
		Path:                  fmt.Sprintf("/Groups/%s", id),
		Authorization:         network.NewTokenAuthorization(token),
		AcceptableStatusCodes: []int{http.StatusOK},
	})
	if err != nil {
		return translateError(err)
	}

	return nil
}

func newGroupFromResponse(config Config, response documents.GroupResponse) Group {
	return Group{
		ID:          response.ID,
		DisplayName: response.DisplayName,
		Version:     response.Meta.Version,
		CreatedAt:   response.Meta.Created,
		UpdatedAt:   response.Meta.LastModified,
	}
}
<|endoftext|>


<commit_before>// A circuit breaker
package circuitry

import (
	"sync"
	"time"
)

// CircuitBreaker represents a circuit breaker
type CircuitBreaker struct {
	FailCounter  int
	FailMax      int
	ResetTimeout time.Duration
	State        circuitState
	StateLock    *sync.Mutex
}

// Create a new circuit breaker with failMax failures and a resetTimeout timeout
func Breaker(failMax int, resetTimeout time.Duration) *CircuitBreaker {
	b := new(CircuitBreaker)
	b.FailCounter = 0
	b.FailMax = failMax
	b.ResetTimeout = resetTimeout
	b.StateLock = new(sync.Mutex)
	b.State = &closedCircuit{b}
	return b
}

// Reports if the circuit is closed
func (b *CircuitBreaker) IsClosed() bool {
	return b.State.BeforeCall()
}

// Reports if the circuit is open
func (b *CircuitBreaker) IsOpen() bool {
	return !b.State.BeforeCall()
}

// Pass error to the to the circuit breaker
func (b *CircuitBreaker) Error(err error) {
	if err == nil {
		b.State.HandleSuccess()
	} else {
		b.State.HandleFailure()
	}
}

// Close the circuit
func (b *CircuitBreaker) Close() {
	b.StateLock.Lock()
	b.FailCounter = 0
	b.State = &closedCircuit{b}
	b.StateLock.Unlock()
}

// Open the circuit
func (b *CircuitBreaker) Open() {
	b.StateLock.Lock()
	b.State = &openCircuit{time.Now(), b}
	b.StateLock.Unlock()
}

// Half-open the circuit
func (b *CircuitBreaker) HalfOpen() {
	b.StateLock.Lock()
	b.State = &halfopenCircuit{b}
	b.StateLock.Unlock()
}
<commit_msg>hide state machinery<commit_after>// A circuit breaker
package circuitry

import (
	"sync"
	"time"
)

// CircuitBreaker represents a circuit breaker
type CircuitBreaker struct {
	FailCounter  int
	FailMax      int
	ResetTimeout time.Duration
	state        circuitState
	lock         *sync.Mutex
}

// Create a new circuit breaker with failMax failures and a resetTimeout timeout
func Breaker(failMax int, resetTimeout time.Duration) *CircuitBreaker {
	b := new(CircuitBreaker)
	b.FailCounter = 0
	b.FailMax = failMax
	b.ResetTimeout = resetTimeout
	b.lock = new(sync.Mutex)
	b.state = &closedCircuit{b}
	return b
}

// Reports if the circuit is closed
func (b *CircuitBreaker) IsClosed() bool {
	return b.state.BeforeCall()
}

// Reports if the circuit is open
func (b *CircuitBreaker) IsOpen() bool {
	return !b.state.BeforeCall()
}

// Pass error to the to the circuit breaker
func (b *CircuitBreaker) Error(err error) {
	if err == nil {
		b.state.HandleSuccess()
	} else {
		b.state.HandleFailure()
	}
}

// Close the circuit
func (b *CircuitBreaker) Close() {
	b.lock.Lock()
	b.FailCounter = 0
	b.state = &closedCircuit{b}
	b.lock.Unlock()
}

// Open the circuit
func (b *CircuitBreaker) Open() {
	b.lock.Lock()
	b.state = &openCircuit{time.Now(), b}
	b.lock.Unlock()
}

// Half-open the circuit
func (b *CircuitBreaker) HalfOpen() {
	b.lock.Lock()
	b.state = &halfopenCircuit{b}
	b.lock.Unlock()
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/user"

	"github.com/opentable/sous/ext/docker"
	"github.com/opentable/sous/ext/git"
	"github.com/opentable/sous/ext/singularity"
	"github.com/opentable/sous/ext/storage"
	"github.com/opentable/sous/lib"
	"github.com/opentable/sous/util/cmdr"
	"github.com/opentable/sous/util/docker_registry"
	"github.com/opentable/sous/util/shell"
	"github.com/pkg/errors"
	"github.com/samsalisbury/psyringe"
	"github.com/samsalisbury/semv"
)

type (
	// Out is an output used for real data a Command returns. This should only
	// be used when a command needs to write directly to stdout, using the
	// formatting options that come with an output. Usually, you should use a
	// SuccessResult with Data to return data.
	Out struct{ *cmdr.Output }
	// ErrOut is an output used for logging from a Command. This should only be
	// used when a Command needs to write a lot of data to stderr, using the
	// formatting options that come with and Output. Usually you should use and
	// ErrorResult to return error messages.
	ErrOut struct{ *cmdr.Output }
	// SousCLIGraph is a dependency injector used to flesh out Sous commands
	// with their dependencies.
	SousCLIGraph struct{ *psyringe.Psyringe }
	// Version represents a version of Sous.

	// OutWriter is an alias on io.Writer to disguish "stderr"
	OutWriter io.Writer
	// ErrWriter is an alias on io.Writer to disguish "stderr"
	ErrWriter io.Writer
)

type (
	Version struct{ semv.Version }
	// LocalUser is the currently logged in user.
	LocalUser struct{ *User }
	// LocalSousConfig is the configuration for Sous.
	LocalSousConfig struct{ *Config }
	// LocalWorkDir is the user's current working directory when they invoke Sous.
	LocalWorkDir string
	// LocalWorkDirShell is a shell for working in the user's current working
	// directory.
	LocalWorkDirShell struct{ *shell.Sh }
	// LocalGitClient is a git client rooted in WorkdirShell.Dir.
	LocalGitClient struct{ *git.Client }
	// LocalGitRepo is the git repository containing WorkDir.
	LocalGitRepo struct{ *git.Repo }
	// GitSourceContext is the source context according to the local git repo.
	GitSourceContext struct{ *sous.SourceContext }
	// ScratchDirShell is a shell for working in the scratch area where things
	// like artefacts, and build metadata are stored. It is a new, empty
	// directory, and should be cleaned up eventually.
	ScratchDirShell struct{ *shell.Sh }
	// LocalDockerClient is a docker client object
	LocalDockerClient struct{ docker_registry.Client }
	// LocalStateReader wraps a storage.StateReader, and should be configured
	// to use the current user's local storage.
	LocalStateReader struct{ storage.StateReader }
	// LocalStateWriter wraps a storage.StateWriter, and should be configured to
	// use the current user's local storage.
	LocalStateWriter struct{ storage.StateWriter }
	// CurrentGDM is a snapshot of the GDM at application start. In a CLI
	// context, which this is, that is all we need to simply read the GDM.
	CurrentGDM struct{ *sous.Deployments }
)

// BuildGraph builds the dependency injection graph, used to populate commands
// invoked by the user.
func BuildGraph(c *cmdr.CLI, out, err io.Writer) *SousCLIGraph {
	return &SousCLIGraph{psyringe.New(
		c,
		func() OutWriter { return out },
		func() ErrWriter { return err },
		newOut,
		newErrOut,
		newLogSet,
		newLocalUser,
		newLocalSousConfig,
		newLocalWorkDir,
		newLocalWorkDirShell,
		newScratchDirShell,
		newLocalGitClient,
		newLocalGitRepo,
		newGitSourceContext,
		newSourceContext,
		newBuildContext,
		newBuildConfig,
		newBuildManager,
		newDockerClient,
		newDockerBuilder,
		newSelector,
		newLabeller,
		newRegistrar,
		newDeployer,
		newRegistry,
		newRegistryDumper,
		newLocalDiskStateManager,
		newLocalStateReader,
		newLocalStateWriter,
		newCurrentGDM,
		newCurrentState,
	)}
}

func newOut(c *cmdr.CLI) Out {
	return Out{c.Out}
}

func newErrOut(c *cmdr.CLI) ErrOut {
	return ErrOut{c.Err}
}

func newRegistryDumper(r sous.Registry) *sous.RegistryDumper {
	return sous.NewRegistryDumper(r)
}

func newLogSet(s *Sous, err ErrWriter) *sous.LogSet { // XXX temporary until we settle on logging
	if s.flags.Verbosity.Debug {
		if s.flags.Verbosity.Loud {
			sous.Log.Vomit.SetOutput(err)
		}
		sous.Log.Debug.SetOutput(err)
		sous.Log.Info.SetOutput(err)

	}
	if s.flags.Verbosity.Loud {
		sous.Log.Info.SetOutput(err)
	}
	if s.flags.Verbosity.Quiet {
	}
	if s.flags.Verbosity.Silent {
	}

	sous.Log.Vomit.Println("Verbose debugging enabled")
	sous.Log.Debug.Println("Regular debugging enabled")
	return &sous.Log
}

/*
func newSourceFlags(c *cmdr.CLI) (*DeployFilterFlags, error) {
	sourceFlags := &DeployFilterFlags{}
	var err error
	c.AddGlobalFlagSetFunc(func(fs *flag.FlagSet) {
		err = AddFlags(fs, sourceFlags, sourceFlagsHelp)
		if err != nil {
			panic(err)
		}
	})
	return sourceFlags, err
}
*/

func newGitSourceContext(g LocalGitRepo) (GitSourceContext, error) {
	c, err := g.SourceContext()
	return GitSourceContext{c}, initErr(err, "getting local git context")
}

func newSourceContext(g GitSourceContext, f *DeployFilterFlags) (*sous.SourceContext, error) {
	c := g.SourceContext
	if c == nil {
		c = &sous.SourceContext{}
	}

	sl, err := resolveSourceLocation(f, c)
	if err != nil {
		return nil, errors.Wrap(err, "resolving source location")
	}
	if sl.RepoURL != c.SourceLocation().RepoURL {
		// TODO: Clone the repository, and use the cloned dir as source context.
		return nil, errors.Errorf("source location %q is not the same as the remote %q",
			sl.RepoURL, c.SourceLocation().RepoURL)
	}
	return c, nil
}

func newBuildContext(wd LocalWorkDirShell, c *sous.SourceContext) *sous.BuildContext {
	return &sous.BuildContext{Sh: wd.Sh, Source: *c}
}

func newBuildConfig(f *DeployFilterFlags, p *PolicyFlags, bc *sous.BuildContext) *sous.BuildConfig {
	cfg := sous.BuildConfig{
		Repo:       f.Repo,
		Offset:     f.Offset,
		Tag:        f.Tag,
		Revision:   f.Revision,
		Strict:     p.Strict,
		ForceClone: p.ForceClone,
		Context:    bc,
	}

	return &cfg
}

func newBuildManager(bc *sous.BuildConfig, sl sous.Selector, lb sous.Labeller, rg sous.Registrar) *sous.BuildManager {
	mgr := &sous.BuildManager{
		BuildConfig: bc,
		Selector:    sl,
		Labeller:    lb,
		Registrar:   rg,
	}
	return mgr
}

func newLocalUser() (v LocalUser, err error) {
	u, err := user.Current()
	v.User = &User{u}
	return v, initErr(err, "getting current user")
}

func newLocalSousConfig(u LocalUser) (v LocalSousConfig, err error) {
	v.Config, err = newConfig(u.User)
	return v, initErr(err, "getting configuration")
}

// TODO: This should register a cleanup task with the cli, to delete the temp
// dir.
func newScratchDirShell() (v ScratchDirShell, err error) {
	const what = "getting scratch directory"
	dir, err := ioutil.TempDir("", "sous")
	if err != nil {
		return v, initErr(err, what)
	}
	v.Sh, err = shell.DefaultInDir(dir)
	v.TeeOut = os.Stdout
	v.TeeErr = os.Stderr
	return v, initErr(err, what)
}

func newLocalWorkDir() (LocalWorkDir, error) {
	s, err := os.Getwd()
	return LocalWorkDir(s), initErr(err, "determining working directory")
}

func newLocalWorkDirShell(l LocalWorkDir) (v LocalWorkDirShell, err error) {
	v.Sh, err = shell.DefaultInDir(string(l))
	v.TeeEcho = os.Stdout
	v.TeeOut = os.Stdout
	v.TeeErr = os.Stderr
	return v, initErr(err, "getting current working directory")
}

func newLocalGitClient(sh LocalWorkDirShell) (v LocalGitClient, err error) {
	v.Client, err = git.NewClient(sh.Sh)
	return v, initErr(err, "initialising git client")
}

func newLocalGitRepo(c LocalGitClient) (v LocalGitRepo, err error) {
	v.Repo, err = c.OpenRepo(".")
	return v, initErr(err, "opening local git repository")
}

func newSelector() sous.Selector {
	return &sous.EchoSelector{
		Factory: func(*sous.BuildContext) (sous.Buildpack, error) {
			return docker.NewDockerfileBuildpack(), nil
		},
	}
}

func newDockerBuilder(cfg LocalSousConfig, cl LocalDockerClient, ctx *sous.SourceContext, source LocalWorkDirShell, scratch ScratchDirShell) (*docker.Builder, error) {
	nc, err := makeDockerRegistry(cfg, cl)
	if err != nil {
		return nil, err
	}
	drh := cfg.Docker.RegistryHost
	return docker.NewBuilder(nc, drh, source.Sh, scratch.Sh)
}

func newLabeller(db *docker.Builder) sous.Labeller {
	return db
}

func newRegistrar(db *docker.Builder) sous.Registrar {
	return db
}

func newRegistry(cfg LocalSousConfig, cl LocalDockerClient) (sous.Registry, error) {
	return makeDockerRegistry(cfg, cl)
}
func newDeployer(r sous.Registry) sous.Deployer {
	// Eventually, based on configuration, we may make different decisions here.
	return singularity.NewDeployer(r, singularity.NewRectiAgent(r))
}

func newDockerClient() LocalDockerClient {
	return LocalDockerClient{docker_registry.NewClient()}
}

func newLocalDiskStateManager(c LocalSousConfig) (*storage.DiskStateManager, error) {
	sm, err := storage.NewDiskStateManager(c.StateLocation)
	return sm, initErr(err, "initialising sous state")
}

func newLocalStateReader(sm *storage.DiskStateManager) LocalStateReader {
	return LocalStateReader{sm}
}

func newLocalStateWriter(sm *storage.DiskStateManager) LocalStateWriter {
	return LocalStateWriter{sm}
}

func newCurrentState(sr LocalStateReader) (*sous.State, error) {
	state, err := sr.ReadState()
	if !os.IsNotExist(err) {
		return state, initErr(err, "reading sous state")
	}
	log.Println("error reading state: %s", err)
	log.Println("defaulting to empty state")
	return sous.NewState(), nil
}

func newCurrentGDM(state *sous.State) (CurrentGDM, error) {
	deployments, err := state.Deployments()
	return CurrentGDM{&deployments}, initErr(err, "expanding state")
}

// The funcs named makeXXX below are used to create specific implementations of
// sous native types.

// makeDockerRegistry creates a Docker version of sous.Registry
func makeDockerRegistry(cfg LocalSousConfig, cl LocalDockerClient) (*docker.NameCache, error) {
	dbCfg := cfg.Docker.DBConfig()
	db, err := docker.GetDatabase(&dbCfg)
	if err != nil {
		return nil, fmt.Errorf("unable to build name cache DB: %s", err)
	}
	return &docker.NameCache{RegistryClient: cl.Client, DB: db}, nil
}

// initErr returns nil if error is nil, otherwise an initialisation error.
// The second argument "what" should be a very short description of the
// initialisation task, e.g. "getting widget" or "reading state" etc.
func initErr(err error, what string) error {
	if err == nil {
		return nil
	}
	message := fmt.Sprintf("error %s:", what)
	if shellErr, ok := err.(shell.Error); ok {
		message += fmt.Sprintf("\ncommand failed:\nshell> %s\n%s",
			shellErr.Command.String(), shellErr.Result.Combined.String())
	} else {
		message += " " + err.Error()
	}
	return fmt.Errorf(message)
}
<commit_msg>Update CLI graph.<commit_after>package cli

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/user"

	"github.com/opentable/sous/ext/docker"
	"github.com/opentable/sous/ext/git"
	"github.com/opentable/sous/ext/singularity"
	"github.com/opentable/sous/ext/storage"
	"github.com/opentable/sous/lib"
	"github.com/opentable/sous/util/cmdr"
	"github.com/opentable/sous/util/docker_registry"
	"github.com/opentable/sous/util/shell"
	"github.com/pkg/errors"
	"github.com/samsalisbury/psyringe"
	"github.com/samsalisbury/semv"
)

type (
	// Out is an output used for real data a Command returns. This should only
	// be used when a command needs to write directly to stdout, using the
	// formatting options that come with an output. Usually, you should use a
	// SuccessResult with Data to return data.
	Out struct{ *cmdr.Output }
	// ErrOut is an output used for logging from a Command. This should only be
	// used when a Command needs to write a lot of data to stderr, using the
	// formatting options that come with and Output. Usually you should use and
	// ErrorResult to return error messages.
	ErrOut struct{ *cmdr.Output }
	// SousCLIGraph is a dependency injector used to flesh out Sous commands
	// with their dependencies.
	SousCLIGraph struct{ *psyringe.Psyringe }
	// Version represents a version of Sous.

	// OutWriter is an alias on io.Writer to disguish "stderr"
	OutWriter io.Writer
	// ErrWriter is an alias on io.Writer to disguish "stderr"
	ErrWriter io.Writer
)

type (
	Version struct{ semv.Version }
	// LocalUser is the currently logged in user.
	LocalUser struct{ *User }
	// LocalSousConfig is the configuration for Sous.
	LocalSousConfig struct{ *Config }
	// LocalWorkDir is the user's current working directory when they invoke Sous.
	LocalWorkDir string
	// LocalWorkDirShell is a shell for working in the user's current working
	// directory.
	LocalWorkDirShell struct{ *shell.Sh }
	// LocalGitClient is a git client rooted in WorkdirShell.Dir.
	LocalGitClient struct{ *git.Client }
	// LocalGitRepo is the git repository containing WorkDir.
	LocalGitRepo struct{ *git.Repo }
	// GitSourceContext is the source context according to the local git repo.
	GitSourceContext struct{ *sous.SourceContext }
	// ScratchDirShell is a shell for working in the scratch area where things
	// like artefacts, and build metadata are stored. It is a new, empty
	// directory, and should be cleaned up eventually.
	ScratchDirShell struct{ *shell.Sh }
	// LocalDockerClient is a docker client object
	LocalDockerClient struct{ docker_registry.Client }
	// LocalStateReader wraps a storage.StateReader, and should be configured
	// to use the current user's local storage.
	LocalStateReader struct{ storage.StateReader }
	// LocalStateWriter wraps a storage.StateWriter, and should be configured to
	// use the current user's local storage.
	LocalStateWriter struct{ storage.StateWriter }
	// CurrentGDM is a snapshot of the GDM at application start. In a CLI
	// context, which this is, that is all we need to simply read the GDM.
	CurrentGDM struct{ *sous.Deployments }
)

// BuildGraph builds the dependency injection graph, used to populate commands
// invoked by the user.
func BuildGraph(c *cmdr.CLI, out, err io.Writer) *SousCLIGraph {
	return &SousCLIGraph{psyringe.New(
		c,
		func() OutWriter { return out },
		func() ErrWriter { return err },
		newOut,
		newErrOut,
		newLogSet,
		newLocalUser,
		newLocalSousConfig,
		newLocalWorkDir,
		newLocalWorkDirShell,
		newScratchDirShell,
		newLocalGitClient,
		newLocalGitRepo,
		newGitSourceContext,
		newSourceContext,
		newBuildContext,
		newBuildConfig,
		newBuildManager,
		newDockerClient,
		newDockerBuilder,
		newSelector,
		newLabeller,
		newRegistrar,
		newDeployer,
		newRegistry,
		newRegistryDumper,
		newLocalDiskStateManager,
		newLocalStateReader,
		newLocalStateWriter,
		newCurrentGDM,
		newCurrentState,
	)}
}

func newOut(c *cmdr.CLI) Out {
	return Out{c.Out}
}

func newErrOut(c *cmdr.CLI) ErrOut {
	return ErrOut{c.Err}
}

func newRegistryDumper(r sous.Registry) *sous.RegistryDumper {
	return sous.NewRegistryDumper(r)
}

func newLogSet(s *Sous, err ErrWriter) *sous.LogSet { // XXX temporary until we settle on logging
	if s.flags.Verbosity.Debug {
		if s.flags.Verbosity.Loud {
			sous.Log.Vomit.SetOutput(err)
		}
		sous.Log.Debug.SetOutput(err)
		sous.Log.Info.SetOutput(err)

	}
	if s.flags.Verbosity.Loud {
		sous.Log.Info.SetOutput(err)
	}
	if s.flags.Verbosity.Quiet {
	}
	if s.flags.Verbosity.Silent {
	}

	sous.Log.Vomit.Println("Verbose debugging enabled")
	sous.Log.Debug.Println("Regular debugging enabled")
	return &sous.Log
}

/*
func newSourceFlags(c *cmdr.CLI) (*DeployFilterFlags, error) {
	sourceFlags := &DeployFilterFlags{}
	var err error
	c.AddGlobalFlagSetFunc(func(fs *flag.FlagSet) {
		err = AddFlags(fs, sourceFlags, sourceFlagsHelp)
		if err != nil {
			panic(err)
		}
	})
	return sourceFlags, err
}
*/

func newGitSourceContext(g LocalGitRepo) (GitSourceContext, error) {
	c, err := g.SourceContext()
	return GitSourceContext{c}, initErr(err, "getting local git context")
}

func newSourceContext(g GitSourceContext, f *DeployFilterFlags) (*sous.SourceContext, error) {
	c := g.SourceContext
	if c == nil {
		c = &sous.SourceContext{}
	}

	sl, err := resolveSourceLocation(f, c)
	if err != nil {
		return nil, errors.Wrap(err, "resolving source location")
	}
	if sl.RepoURL != c.SourceLocation().RepoURL {
		// TODO: Clone the repository, and use the cloned dir as source context.
		return nil, errors.Errorf("source location %q is not the same as the remote %q",
			sl.RepoURL, c.SourceLocation().RepoURL)
	}
	return c, nil
}

func newBuildContext(wd LocalWorkDirShell, c *sous.SourceContext) *sous.BuildContext {
	return &sous.BuildContext{Sh: wd.Sh, Source: *c}
}

func newBuildConfig(f *DeployFilterFlags, p *PolicyFlags, bc *sous.BuildContext) *sous.BuildConfig {
	cfg := sous.BuildConfig{
		Repo:       f.Repo,
		Offset:     f.Offset,
		Tag:        f.Tag,
		Revision:   f.Revision,
		Strict:     p.Strict,
		ForceClone: p.ForceClone,
		Context:    bc,
	}

	return &cfg
}

func newBuildManager(bc *sous.BuildConfig, sl sous.Selector, lb sous.Labeller, rg sous.Registrar) *sous.BuildManager {
	mgr := &sous.BuildManager{
		BuildConfig: bc,
		Selector:    sl,
		Labeller:    lb,
		Registrar:   rg,
	}
	return mgr
}

func newLocalUser() (v LocalUser, err error) {
	u, err := user.Current()
	v.User = &User{u}
	return v, initErr(err, "getting current user")
}

func newLocalSousConfig(u LocalUser) (v LocalSousConfig, err error) {
	v.Config, err = newConfig(u.User)
	return v, initErr(err, "getting configuration")
}

// TODO: This should register a cleanup task with the cli, to delete the temp
// dir.
func newScratchDirShell() (v ScratchDirShell, err error) {
	const what = "getting scratch directory"
	dir, err := ioutil.TempDir("", "sous")
	if err != nil {
		return v, initErr(err, what)
	}
	v.Sh, err = shell.DefaultInDir(dir)
	v.TeeOut = os.Stdout
	v.TeeErr = os.Stderr
	return v, initErr(err, what)
}

func newLocalWorkDir() (LocalWorkDir, error) {
	s, err := os.Getwd()
	return LocalWorkDir(s), initErr(err, "determining working directory")
}

func newLocalWorkDirShell(l LocalWorkDir) (v LocalWorkDirShell, err error) {
	v.Sh, err = shell.DefaultInDir(string(l))
	v.TeeEcho = os.Stdout
	v.TeeOut = os.Stdout
	v.TeeErr = os.Stderr
	return v, initErr(err, "getting current working directory")
}

func newLocalGitClient(sh LocalWorkDirShell) (v LocalGitClient, err error) {
	v.Client, err = git.NewClient(sh.Sh)
	return v, initErr(err, "initialising git client")
}

func newLocalGitRepo(c LocalGitClient) (v LocalGitRepo, err error) {
	v.Repo, err = c.OpenRepo(".")
	return v, initErr(err, "opening local git repository")
}

func newSelector() sous.Selector {
	return &sous.EchoSelector{
		Factory: func(*sous.BuildContext) (sous.Buildpack, error) {
			return docker.NewDockerfileBuildpack(), nil
		},
	}
}

func newDockerBuilder(cfg LocalSousConfig, cl LocalDockerClient, ctx *sous.SourceContext, source LocalWorkDirShell, scratch ScratchDirShell) (*docker.Builder, error) {
	nc, err := makeDockerRegistry(cfg, cl)
	if err != nil {
		return nil, err
	}
	drh := cfg.Docker.RegistryHost
	return docker.NewBuilder(nc, drh, source.Sh, scratch.Sh)
}

func newLabeller(db *docker.Builder) sous.Labeller {
	return db
}

func newRegistrar(db *docker.Builder) sous.Registrar {
	return db
}

func newRegistry(cfg LocalSousConfig, cl LocalDockerClient) (sous.Registry, error) {
	return makeDockerRegistry(cfg, cl)
}
func newDeployer(r sous.Registry) sous.Deployer {
	// Eventually, based on configuration, we may make different decisions here.
	return singularity.NewDeployer(r, singularity.NewRectiAgent(r))
}

func newDockerClient() LocalDockerClient {
	return LocalDockerClient{docker_registry.NewClient()}
}

func newLocalDiskStateManager(c LocalSousConfig) *storage.DiskStateManager {
	return storage.NewDiskStateManager(c.StateLocation)
}

func newLocalStateReader(sm *storage.DiskStateManager) LocalStateReader {
	return LocalStateReader{sm}
}

func newLocalStateWriter(sm *storage.DiskStateManager) LocalStateWriter {
	return LocalStateWriter{sm}
}

func newCurrentState(sr LocalStateReader) (*sous.State, error) {
	state, err := sr.ReadState()
	if !os.IsNotExist(err) {
		return state, initErr(err, "reading sous state")
	}
	log.Println("error reading state: %s", err)
	log.Println("defaulting to empty state")
	return sous.NewState(), nil
}

func newCurrentGDM(state *sous.State) (CurrentGDM, error) {
	deployments, err := state.Deployments()
	return CurrentGDM{&deployments}, initErr(err, "expanding state")
}

// The funcs named makeXXX below are used to create specific implementations of
// sous native types.

// makeDockerRegistry creates a Docker version of sous.Registry
func makeDockerRegistry(cfg LocalSousConfig, cl LocalDockerClient) (*docker.NameCache, error) {
	dbCfg := cfg.Docker.DBConfig()
	db, err := docker.GetDatabase(&dbCfg)
	if err != nil {
		return nil, fmt.Errorf("unable to build name cache DB: %s", err)
	}
	return &docker.NameCache{RegistryClient: cl.Client, DB: db}, nil
}

// initErr returns nil if error is nil, otherwise an initialisation error.
// The second argument "what" should be a very short description of the
// initialisation task, e.g. "getting widget" or "reading state" etc.
func initErr(err error, what string) error {
	if err == nil {
		return nil
	}
	message := fmt.Sprintf("error %s:", what)
	if shellErr, ok := err.(shell.Error); ok {
		message += fmt.Sprintf("\ncommand failed:\nshell> %s\n%s",
			shellErr.Command.String(), shellErr.Result.Combined.String())
	} else {
		message += " " + err.Error()
	}
	return fmt.Errorf(message)
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"log"
	//	"gopkg.in/yaml.v2"
	"os"
	//	"text/template"
	"github.com/qadium/plumber/bindata"
	"github.com/qadium/plumber/graph"
	"github.com/qadium/plumber/shell"
	"os/exec"
	"path/filepath"
	"text/template"
	// "golang.org/x/oauth2/google"
	// "golang.org/x/oauth2"
	// "google.golang.org/cloud"
	// "google.golang.org/cloud/container"
	// kubectl "github.com/GoogleCloudPlatform/kubernetes/pkg/kubectl/cmd"
	// cmdutil "github.com/GoogleCloudPlatform/kubernetes/pkg/kubectl/cmd/util"
	// "github.com/GoogleCloudPlatform/kubernetes/pkg/client/clientcmd"
	// "os"
)

type pipelineInfo struct {
	path           string
	name           string
	commit         string
	plumberVersion string
	plumberCommit  string
}

type kubeData struct {
	BundleName     string
	ExternalFacing bool
	PipelineName   string
	PipelineCommit string
	PlumberVersion string
	PlumberCommit  string
	ImageName      string
	Args           []string
}

func contextsToGraph(ctxs []*Context) []*graph.Node {
	nodes := make([]*graph.Node, len(ctxs))
	m := make(map[string]int) // this map maps inputs to the index of the
	// node that uses it
	// build a map to create the DAG
	for i, ctx := range ctxs {
		nodes[i] = graph.NewNode(ctx.Name)
		for _, input := range ctx.Inputs {
			m[input.Name] = i
		}
	}

	for i, ctx := range ctxs {
		for _, output := range ctx.Outputs {
			if v, ok := m[output.Name]; ok {
				nodes[i].Children = append(nodes[i].Children, nodes[v])
			}
		}
	}
	return nodes
}

func localStart(sortedPipeline []string) error {
	log.Printf(" |  Starting bundles...")
	managerDockerArgs := []string{"run", "-p", "9800:9800", "--rm", "plumber/manager"}
	// walk through the reverse sorted bundles and start them up
	for i := len(sortedPipeline) - 1; i >= 0; i-- {
		bundleName := sortedPipeline[i]
		log.Printf("    Starting: '%s'", bundleName)
		cmd := exec.Command("docker", "run", "-d", "-P", fmt.Sprintf("plumber/%s", bundleName))
		containerId, err := cmd.Output()
		if err != nil {
			return err
		}

		defer func() {
			log.Printf("    Stopping: '%s'", bundleName)
			cmd := exec.Command("docker", "rm", "-f", string(containerId)[0:4])
			_, err := cmd.Output()
			if err != nil {
				panic(err)
			}
			log.Printf("    Stopped.")
		}()

		log.Printf("    Started: %s", string(containerId))
		cmd = exec.Command("docker", "inspect", "--format='{{(index (index .NetworkSettings.Ports \"9800/tcp\") 0).HostPort}}'", string(containerId)[0:4])
		portNum, err := cmd.Output()
		if err != nil {
			return err
		}
		// should use docker host IP (for local deploy)
		// should use "names" for kubernetes deploy
		managerDockerArgs = append(managerDockerArgs, fmt.Sprintf("http://172.17.42.1:%s", string(portNum[:len(portNum)-1])))
	}
	log.Printf("    Done.")
	log.Printf("    Args passed to 'docker': %v", managerDockerArgs)

	log.Printf(" |  Running manager. CTRL-C to quit.")
	err := shell.RunAndLog("docker", managerDockerArgs...)
	if err != nil {
		return err
	}
	log.Printf("    Done.")
	return nil
}

func writeKubernetesTemplate(tmplType string, destFilename string, templateData kubeData) error {
	tmpl, err := bindata.Asset(fmt.Sprintf("templates/%s.yaml", tmplType))
	if err != nil {
		return err
	}

	file, err := os.Create(destFilename)
	if err != nil {
		return err
	}
	defer func() {
		if err := file.Close(); err != nil {
			panic(err)
		}
	}()

	tmplFile, err := template.New("template").Parse(string(tmpl))
	if err != nil {
		return err
	}
	if err := tmplFile.Execute(file, templateData); err != nil {
		return err
	}
	return nil
}

func writeKubernetesFiles(templateData kubeData) error {
	log.Printf(" |     Writing '%s'", templateData.BundleName)
	k8s, err := KubernetesPath(templateData.PipelineName)
	if err != nil {
		return err
	}

	log.Printf("       Creating service file.")
	err = writeKubernetesTemplate("service", fmt.Sprintf("%s/%s.yaml", k8s, templateData.BundleName), templateData)
	if err != nil {
		return err
	}
	log.Printf("       Created.")

	log.Printf("       Creating replication controller file.")
	err = writeKubernetesTemplate("replication-controller", fmt.Sprintf("%s/%s-rc.yaml", k8s, templateData.BundleName), templateData)
	if err != nil {
		return err
	}
	log.Printf("       Created.")
	log.Printf("       Done.")
	return nil
}

func remoteStart(sortedPipeline []string, projectId string, pipeline pipelineInfo) error {
	// we can probably get the project name with google cloud SDK
	log.Printf("   Creating 'k8s' directory...")
	k8s, err := KubernetesPath(pipeline.name)
	if err != nil {
		return err
	}

	if err := os.MkdirAll(k8s, 0755); err != nil {
		return err
	}
	log.Printf("   Created.")

	args := []string{}

	for i := len(sortedPipeline) - 1; i >= 0; i-- {
		bundleName := sortedPipeline[i]
		localDockerTag := fmt.Sprintf("plumber/%s", bundleName)
		remoteDockerTag := fmt.Sprintf("gcr.io/%s/plumber-%s", projectId, bundleName)
		data := kubeData{
			BundleName:     bundleName,
			ImageName:      remoteDockerTag,
			PlumberVersion: pipeline.plumberVersion,
			PlumberCommit:  pipeline.plumberCommit,
			PipelineName:   pipeline.name,
			PipelineCommit: pipeline.commit,
			ExternalFacing: false,
			Args:           []string{},
		}

		// step 1. re-tag local containers to gcr.io/$GCE/$pipeline-$bundlename
		log.Printf("    Retagging: '%s'", bundleName)
		err := shell.RunAndLog("docker", "tag", "-f", localDockerTag, remoteDockerTag)
		if err != nil {
			return err
		}
		// step 2. push them to gce
		log.Printf("    Submitting: '%s'", remoteDockerTag)
		err = shell.RunAndLog("gcloud", "preview", "docker", "push", remoteDockerTag)
		if err != nil {
			return err
		}
		// step 3. generate k8s files in pipelinePath
		if err := writeKubernetesFiles(data); err != nil {
			return err
		}

		// append to arglist (args now in sorted order)
		args = append(args, fmt.Sprintf("http://%s:9800", bundleName))
	}
	// create the manager service
	data := kubeData{
		BundleName:     "manager",
		ImageName:      fmt.Sprintf("gcr.io/%s/plumber-manager", projectId),
		PlumberVersion: pipeline.plumberVersion,
		PlumberCommit:  pipeline.plumberCommit,
		PipelineName:   pipeline.name,
		PipelineCommit: pipeline.commit,
		ExternalFacing: true,
		Args:           args,
	}
	// step 1. re-tag local containers to gcr.io/$GCE/$pipeline-$bundlename
	log.Printf("    Retagging: '%s'", data.BundleName)
	err = shell.RunAndLog("docker", "tag", "-f", "plumber/manager", data.ImageName)
	if err != nil {
		return err
	}
	// step 2. push them to gce
	log.Printf("    Submitting: '%s'", data.ImageName)
	err = shell.RunAndLog("gcloud", "preview", "docker", "push", data.ImageName)
	if err != nil {
		return err
	}
	// step 3. generate k8s file in pipeline
	if err := writeKubernetesFiles(data); err != nil {
		return err
	}

	// step 4. launch all the services
	err = shell.RunAndLog("kubectl", "create", "-f", fmt.Sprintf("%s/k8s", pipeline.path))
	if err != nil {
		return err
	}
	// step 5: open up the firewall?
	return nil
}

func Start(pipeline, gce, plumberVersion, plumberGitCommit string) error {
	log.Printf("==> Starting '%s' pipeline", pipeline)
	defer log.Printf("<== '%s' finished.", pipeline)

	log.Printf(" |  Building dependency graph.")
	path, err := GetPipeline(pipeline)
	if err != nil {
		return err
	}

	configs, err := filepath.Glob(fmt.Sprintf("%s/*.yml", path))
	if err != nil {
		return err
	}

	ctxs := make([]*Context, len(configs))
	for i, config := range configs {
		ctxs[i], err = ParseConfig(config)
		if err != nil {
			return err
		}
	}

	// graph with diamond (test case)
	// n1 := graph.NewNode("foo")
	// n2 := graph.NewNode("bar")
	// n3 := graph.NewNode("joe")
	// n4 := graph.NewNode("bob")
	// n1.Children = append(n1.Children, n2, n3)
	// n2.Children = append(n2.Children, n4)
	// n3.Children = append(n3.Children, n4)

	g := contextsToGraph(ctxs)
	sortedPipeline, err := graph.ReverseTopoSort(g)
	if err != nil {
		return err
	}
	log.Printf("    Reverse sorted: %v", sortedPipeline)
	log.Printf("    Completed.")

	if gce != "" {
		// start GOOGLE experiments?
		// when start is invoked with --gce PROJECT_ID, this piece of code
		// should be run
		// client, err := google.DefaultClient(oauth2.NoContext, "https://www.googleapis.com/auth/compute")
		// if err != nil {
		// 	return err
		// }
		// cloudCtx := cloud.NewContext("kubernetes-fun", client)
		//
		// resources, err := container.Clusters(cloudCtx, "")
		// if err != nil {
		// 	return err
		// }
		// for _, op := range resources {
		// 	log.Printf("%v", op)
		// }
		//
		// loadingRules := clientcmd.NewDefaultClientConfigLoadingRules()
		// log.Printf("loading rules: %v", *loadingRules)
		// configOverrides := &clientcmd.ConfigOverrides{}
		// kubeConfig := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(loadingRules, configOverrides)
		// cfg, err := kubeConfig.ClientConfig()
		// if err != nil {
		// 	return err
		// }

		// well, we just shell out!
		// f := cmdutil.NewFactory(nil)
		// cmd := kubectl.NewCmdCreate(f, os.Stdout)
		// f.BindFlags(cmd.PersistentFlags())
		// cmd.Flags().Set("filename", "/Users/echu/.plumber/foo/k8s")
		// cmd.Run(cmd, []string{})

		// end GOOGLE experiments

		info := pipelineInfo{
			name:           pipeline,
			path:           path,
			commit:         "",
			plumberVersion: plumberVersion,
			plumberCommit:  plumberGitCommit,
		}
		log.Printf(" |  Running remote pipeline.")
		return remoteStart(sortedPipeline, gce, info)
	} else {
		log.Printf(" |  Running local pipeline.")
		return localStart(sortedPipeline)
	}
	return nil
}
<commit_msg>updated start<commit_after>package cli

import (
	"fmt"
	"log"
	//	"gopkg.in/yaml.v2"
	"os"
	//	"text/template"
	"github.com/qadium/plumber/bindata"
	"github.com/qadium/plumber/graph"
	"github.com/qadium/plumber/shell"
	"os/exec"
	"path/filepath"
	"text/template"
	// "golang.org/x/oauth2/google"
	// "golang.org/x/oauth2"
	// "google.golang.org/cloud"
	// "google.golang.org/cloud/container"
	// kubectl "github.com/GoogleCloudPlatform/kubernetes/pkg/kubectl/cmd"
	// cmdutil "github.com/GoogleCloudPlatform/kubernetes/pkg/kubectl/cmd/util"
	// "github.com/GoogleCloudPlatform/kubernetes/pkg/client/clientcmd"
	// "os"
)

type pipelineInfo struct {
	path           string
	name           string
	commit         string
	plumberVersion string
	plumberCommit  string
}

type kubeData struct {
	BundleName     string
	ExternalFacing bool
	PipelineName   string
	PipelineCommit string
	PlumberVersion string
	PlumberCommit  string
	ImageName      string
	Args           []string
}

func contextsToGraph(ctxs []*Context) []*graph.Node {
	nodes := make([]*graph.Node, len(ctxs))
	m := make(map[string]int) // this map maps inputs to the index of the
	// node that uses it
	// build a map to create the DAG
	for i, ctx := range ctxs {
		nodes[i] = graph.NewNode(ctx.Name)
		for _, input := range ctx.Inputs {
			m[input.Name] = i
		}
	}

	for i, ctx := range ctxs {
		for _, output := range ctx.Outputs {
			if v, ok := m[output.Name]; ok {
				nodes[i].AddChildren(nodes[v])
			}
		}
	}
	return nodes
}

func localStart(sortedPipeline []string) error {
	log.Printf(" |  Starting bundles...")
	managerDockerArgs := []string{"run", "-p", "9800:9800", "--rm", "plumber/manager"}
	// walk through the reverse sorted bundles and start them up
	for i := len(sortedPipeline) - 1; i >= 0; i-- {
		bundleName := sortedPipeline[i]
		log.Printf("    Starting: '%s'", bundleName)
		cmd := exec.Command("docker", "run", "-d", "-P", fmt.Sprintf("plumber/%s", bundleName))
		containerId, err := cmd.Output()
		if err != nil {
			return err
		}

		defer func() {
			log.Printf("    Stopping: '%s'", bundleName)
			cmd := exec.Command("docker", "rm", "-f", string(containerId)[0:4])
			_, err := cmd.Output()
			if err != nil {
				panic(err)
			}
			log.Printf("    Stopped.")
		}()

		log.Printf("    Started: %s", string(containerId))
		cmd = exec.Command("docker", "inspect", "--format='{{(index (index .NetworkSettings.Ports \"9800/tcp\") 0).HostPort}}'", string(containerId)[0:4])
		portNum, err := cmd.Output()
		if err != nil {
			return err
		}
		// should use docker host IP (for local deploy)
		// should use "names" for kubernetes deploy
		managerDockerArgs = append(managerDockerArgs, fmt.Sprintf("http://172.17.42.1:%s", string(portNum[:len(portNum)-1])))
	}
	log.Printf("    Done.")
	log.Printf("    Args passed to 'docker': %v", managerDockerArgs)

	log.Printf(" |  Running manager. CTRL-C to quit.")
	err := shell.RunAndLog("docker", managerDockerArgs...)
	if err != nil {
		return err
	}
	log.Printf("    Done.")
	return nil
}

func writeKubernetesTemplate(tmplType string, destFilename string, templateData kubeData) error {
	tmpl, err := bindata.Asset(fmt.Sprintf("templates/%s.yaml", tmplType))
	if err != nil {
		return err
	}

	file, err := os.Create(destFilename)
	if err != nil {
		return err
	}
	defer func() {
		if err := file.Close(); err != nil {
			panic(err)
		}
	}()

	tmplFile, err := template.New("template").Parse(string(tmpl))
	if err != nil {
		return err
	}
	if err := tmplFile.Execute(file, templateData); err != nil {
		return err
	}
	return nil
}

func writeKubernetesFiles(templateData kubeData) error {
	log.Printf(" |     Writing '%s'", templateData.BundleName)
	k8s, err := KubernetesPath(templateData.PipelineName)
	if err != nil {
		return err
	}

	log.Printf("       Creating service file.")
	err = writeKubernetesTemplate("service", fmt.Sprintf("%s/%s.yaml", k8s, templateData.BundleName), templateData)
	if err != nil {
		return err
	}
	log.Printf("       Created.")

	log.Printf("       Creating replication controller file.")
	err = writeKubernetesTemplate("replication-controller", fmt.Sprintf("%s/%s-rc.yaml", k8s, templateData.BundleName), templateData)
	if err != nil {
		return err
	}
	log.Printf("       Created.")
	log.Printf("       Done.")
	return nil
}

func remoteStart(sortedPipeline []string, projectId string, pipeline pipelineInfo) error {
	// we can probably get the project name with google cloud SDK
	log.Printf("   Creating 'k8s' directory...")
	k8s, err := KubernetesPath(pipeline.name)
	if err != nil {
		return err
	}

	if err := os.MkdirAll(k8s, 0755); err != nil {
		return err
	}
	log.Printf("   Created.")

	args := []string{}

	for i := len(sortedPipeline) - 1; i >= 0; i-- {
		bundleName := sortedPipeline[i]
		localDockerTag := fmt.Sprintf("plumber/%s", bundleName)
		remoteDockerTag := fmt.Sprintf("gcr.io/%s/plumber-%s", projectId, bundleName)
		data := kubeData{
			BundleName:     bundleName,
			ImageName:      remoteDockerTag,
			PlumberVersion: pipeline.plumberVersion,
			PlumberCommit:  pipeline.plumberCommit,
			PipelineName:   pipeline.name,
			PipelineCommit: pipeline.commit,
			ExternalFacing: false,
			Args:           []string{},
		}

		// step 1. re-tag local containers to gcr.io/$GCE/$pipeline-$bundlename
		log.Printf("    Retagging: '%s'", bundleName)
		err := shell.RunAndLog("docker", "tag", "-f", localDockerTag, remoteDockerTag)
		if err != nil {
			return err
		}
		// step 2. push them to gce
		log.Printf("    Submitting: '%s'", remoteDockerTag)
		err = shell.RunAndLog("gcloud", "preview", "docker", "push", remoteDockerTag)
		if err != nil {
			return err
		}
		// step 3. generate k8s files in pipelinePath
		if err := writeKubernetesFiles(data); err != nil {
			return err
		}

		// append to arglist (args now in sorted order)
		args = append(args, fmt.Sprintf("http://%s:9800", bundleName))
	}
	// create the manager service
	data := kubeData{
		BundleName:     "manager",
		ImageName:      fmt.Sprintf("gcr.io/%s/plumber-manager", projectId),
		PlumberVersion: pipeline.plumberVersion,
		PlumberCommit:  pipeline.plumberCommit,
		PipelineName:   pipeline.name,
		PipelineCommit: pipeline.commit,
		ExternalFacing: true,
		Args:           args,
	}
	// step 1. re-tag local containers to gcr.io/$GCE/$pipeline-$bundlename
	log.Printf("    Retagging: '%s'", data.BundleName)
	err = shell.RunAndLog("docker", "tag", "-f", "plumber/manager", data.ImageName)
	if err != nil {
		return err
	}
	// step 2. push them to gce
	log.Printf("    Submitting: '%s'", data.ImageName)
	err = shell.RunAndLog("gcloud", "preview", "docker", "push", data.ImageName)
	if err != nil {
		return err
	}
	// step 3. generate k8s file in pipeline
	if err := writeKubernetesFiles(data); err != nil {
		return err
	}

	// step 4. launch all the services
	err = shell.RunAndLog("kubectl", "create", "-f", fmt.Sprintf("%s/k8s", pipeline.path))
	if err != nil {
		return err
	}
	// step 5: open up the firewall?
	return nil
}

func Start(pipeline, gce, plumberVersion, plumberGitCommit string) error {
	log.Printf("==> Starting '%s' pipeline", pipeline)
	defer log.Printf("<== '%s' finished.", pipeline)

	log.Printf(" |  Building dependency graph.")
	path, err := GetPipeline(pipeline)
	if err != nil {
		return err
	}

	configs, err := filepath.Glob(fmt.Sprintf("%s/*.yml", path))
	if err != nil {
		return err
	}

	ctxs := make([]*Context, len(configs))
	for i, config := range configs {
		ctxs[i], err = ParseConfig(config)
		if err != nil {
			return err
		}
	}

	g := contextsToGraph(ctxs)
	sortedPipeline, err := graph.ReverseTopoSort(g)
	if err != nil {
		return err
	}
	log.Printf("    Reverse sorted: %v", sortedPipeline)
	log.Printf("    Completed.")

	if gce != "" {
		// start GOOGLE experiments?
		// when start is invoked with --gce PROJECT_ID, this piece of code
		// should be run
		// client, err := google.DefaultClient(oauth2.NoContext, "https://www.googleapis.com/auth/compute")
		// if err != nil {
		// 	return err
		// }
		// cloudCtx := cloud.NewContext("kubernetes-fun", client)
		//
		// resources, err := container.Clusters(cloudCtx, "")
		// if err != nil {
		// 	return err
		// }
		// for _, op := range resources {
		// 	log.Printf("%v", op)
		// }
		//
		// loadingRules := clientcmd.NewDefaultClientConfigLoadingRules()
		// log.Printf("loading rules: %v", *loadingRules)
		// configOverrides := &clientcmd.ConfigOverrides{}
		// kubeConfig := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(loadingRules, configOverrides)
		// cfg, err := kubeConfig.ClientConfig()
		// if err != nil {
		// 	return err
		// }

		// well, we just shell out!
		// f := cmdutil.NewFactory(nil)
		// cmd := kubectl.NewCmdCreate(f, os.Stdout)
		// f.BindFlags(cmd.PersistentFlags())
		// cmd.Flags().Set("filename", "/Users/echu/.plumber/foo/k8s")
		// cmd.Run(cmd, []string{})

		// end GOOGLE experiments

		info := pipelineInfo{
			name:           pipeline,
			path:           path,
			commit:         "",
			plumberVersion: plumberVersion,
			plumberCommit:  plumberGitCommit,
		}
		log.Printf(" |  Running remote pipeline.")
		return remoteStart(sortedPipeline, gce, info)
	} else {
		log.Printf(" |  Running local pipeline.")
		return localStart(sortedPipeline)
	}
	return nil
}
<|endoftext|>


<commit_before>// A memcached binary protocol client.
package memcached

import (
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"net"
	"strings"
	"sync"

	"github.com/dustin/gomemcached"
)

const bufsize = 1024

// The Client itself.
type Client struct {
	conn    io.ReadWriteCloser
	healthy bool

	hdrBuf []byte
}

// Connect to a memcached server.
func Connect(prot, dest string) (rv *Client, err error) {
	conn, err := net.Dial(prot, dest)
	if err != nil {
		return nil, err
	}
	return &Client{
		conn:    conn,
		healthy: true,
		hdrBuf:  make([]byte, gomemcached.HDR_LEN),
	}, nil
}

// Close the connection when you're done.
func (c *Client) Close() {
	c.conn.Close()
}

// Return false if this client has had issues communicating.
//
// This is useful for connection pools where we want to
// non-destructively determine that a connection may be reused.
func (c Client) IsHealthy() bool {
	return c.healthy
}

// Send a custom request and get the response.
func (client *Client) Send(req *gomemcached.MCRequest) (rv *gomemcached.MCResponse, err error) {
	err = transmitRequest(client.conn, req)
	if err != nil {
		client.healthy = false
		return
	}
	resp, err := getResponse(client.conn, client.hdrBuf)
	if err != nil {
		client.healthy = false
	}
	return resp, err
}

// Send a request, but do not wait for a response.
func (client *Client) Transmit(req *gomemcached.MCRequest) error {
	err := transmitRequest(client.conn, req)
	if err != nil {
		client.healthy = false
	}
	return err
}

// Receive a response
func (client *Client) Receive() (*gomemcached.MCResponse, error) {
	resp, err := getResponse(client.conn, client.hdrBuf)
	if err != nil {
		client.healthy = false
	}
	return resp, err
}

// Get the value for a key.
func (client *Client) Get(vb uint16, key string) (*gomemcached.MCResponse, error) {
	return client.Send(&gomemcached.MCRequest{
		Opcode:  gomemcached.GET,
		VBucket: vb,
		Key:     []byte(key),
		Cas:     0,
		Opaque:  0,
		Extras:  []byte{},
		Body:    []byte{}})
}

// Delete a key.
func (client *Client) Del(vb uint16, key string) (*gomemcached.MCResponse, error) {
	return client.Send(&gomemcached.MCRequest{
		Opcode:  gomemcached.DELETE,
		VBucket: vb,
		Key:     []byte(key),
		Cas:     0,
		Opaque:  0,
		Extras:  []byte{},
		Body:    []byte{}})
}

// List auth mechanisms
func (client *Client) AuthList() (*gomemcached.MCResponse, error) {
	return client.Send(&gomemcached.MCRequest{
		Opcode:  gomemcached.SASL_LIST_MECHS,
		VBucket: 0,
		Key:     []byte{},
		Cas:     0,
		Opaque:  0,
		Extras:  []byte{},
		Body:    []byte{}})
}

func (client *Client) Auth(user, pass string) (*gomemcached.MCResponse, error) {
	res, err := client.AuthList()

	if err != nil {
		return res, err
	}

	authMech := string(res.Body)
	if strings.Index(authMech, "PLAIN") != -1 {
		return client.Send(&gomemcached.MCRequest{
			Opcode:  gomemcached.SASL_AUTH,
			VBucket: 0,
			Key:     []byte("PLAIN"),
			Cas:     0,
			Opaque:  0,
			Extras:  []byte{},
			Body:    []byte(fmt.Sprintf("\x00%s\x00%s", user, pass))})
	}
	return res, fmt.Errorf("Auth mechanism PLAIN not supported")
}

func (client *Client) store(opcode gomemcached.CommandCode, vb uint16,
	key string, flags int, exp int, body []byte) (*gomemcached.MCResponse, error) {

	req := &gomemcached.MCRequest{
		Opcode:  opcode,
		VBucket: vb,
		Key:     []byte(key),
		Cas:     0,
		Opaque:  0,
		Extras:  []byte{0, 0, 0, 0, 0, 0, 0, 0},
		Body:    body}

	binary.BigEndian.PutUint64(req.Extras, uint64(flags)<<32|uint64(exp))
	return client.Send(req)
}

// Increment a value.
func (client *Client) Incr(vb uint16, key string,
	amt, def uint64, exp int) (uint64, error) {

	req := &gomemcached.MCRequest{
		Opcode:  gomemcached.INCREMENT,
		VBucket: vb,
		Key:     []byte(key),
		Cas:     0,
		Opaque:  0,
		Extras:  make([]byte, 8+8+4),
		Body:    []byte{}}
	binary.BigEndian.PutUint64(req.Extras[:8], amt)
	binary.BigEndian.PutUint64(req.Extras[8:16], def)
	binary.BigEndian.PutUint32(req.Extras[16:20], uint32(exp))

	resp, err := client.Send(req)
	if err != nil {
		return 0, err
	}

	return binary.BigEndian.Uint64(resp.Body), nil
}

// Add a value for a key (store if not exists).
func (client *Client) Add(vb uint16, key string, flags int, exp int,
	body []byte) (*gomemcached.MCResponse, error) {
	return client.store(gomemcached.ADD, vb, key, flags, exp, body)
}

// Set the value for a key.
func (client *Client) Set(vb uint16, key string, flags int, exp int,
	body []byte) (*gomemcached.MCResponse, error) {
	return client.store(gomemcached.SET, vb, key, flags, exp, body)
}

// Get keys in bulk
func (client *Client) GetBulk(vb uint16, keys []string) (map[string]*gomemcached.MCResponse, error) {
	terminalOpaque := uint32(len(keys) + 5)
	rv := map[string]*gomemcached.MCResponse{}
	wg := sync.WaitGroup{}
	going := true

	defer func() {
		going = false
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		for going {
			res, err := client.Receive()
			if err != nil {
				return
			}
			if res.Opaque == terminalOpaque {
				return
			}
			if res.Opcode != gomemcached.GETQ {
				log.Panicf("Unexpected opcode in GETQ response: %+v",
					res)
			}
			rv[keys[res.Opaque]] = res
		}
	}()

	for i, k := range keys {
		err := client.Transmit(&gomemcached.MCRequest{
			Opcode:  gomemcached.GETQ,
			VBucket: vb,
			Key:     []byte(k),
			Cas:     0,
			Opaque:  uint32(i),
			Extras:  []byte{},
			Body:    []byte{}})
		if err != nil {
			return rv, err
		}
	}

	err := client.Transmit(&gomemcached.MCRequest{
		Opcode: gomemcached.NOOP,
		Key:    []byte{},
		Cas:    0,
		Extras: []byte{},
		Body:   []byte{},
		Opaque: terminalOpaque})
	if err != nil {
		return rv, err
	}

	wg.Wait()

	return rv, nil
}

// Operation to perform on this CAS loop.
type CasOp uint8

const (
	// Store the new value normally
	CASStore = CasOp(iota)
	// Stop attempting to CAS, leave value untouched
	CASQuit
	// Delete the current value
	CASDelete
)

// User specified termination is returned as an error.
func (c CasOp) Error() string {
	switch c {
	case CASStore:
		return "CAS store"
	case CASQuit:
		return "CAS quit"
	case CASDelete:
		return "CAS delete"
	}
	panic("Unhandled value")
}

// A function to perform a CAS transform
type CasFunc func(current []byte) ([]byte, CasOp)

// Perform a CAS transform with the given function.
//
// If the value does not exist, an empty byte string will be sent to f
func (client *Client) CAS(vb uint16, k string, f CasFunc,
	initexp int) (rv *gomemcached.MCResponse, err error) {

	flags := 0
	exp := 0

	for {
		orig, err := client.Get(vb, k)
		if err != nil && (orig == nil || orig.Status != gomemcached.KEY_ENOENT) {
			return rv, err
		}

		if orig.Status == gomemcached.KEY_ENOENT {
			init, operation := f([]byte{})
			if operation == CASQuit || operation == CASDelete {
				return nil, operation
			}
			// If it doesn't exist, add it
			resp, err := client.Add(vb, k, 0, initexp, init)
			if err == nil && resp.Status != gomemcached.KEY_EEXISTS {
				return rv, err
			}
			// Copy the body into this response.
			resp.Body = init
			return resp, err
		} else {
			var req *gomemcached.MCRequest
			newValue, operation := f(orig.Body)

			switch operation {
			case CASQuit:
				return nil, operation
			case CASStore:
				req = &gomemcached.MCRequest{
					Opcode:  gomemcached.SET,
					VBucket: vb,
					Key:     []byte(k),
					Cas:     orig.Cas,
					Opaque:  0,
					Extras:  []byte{0, 0, 0, 0, 0, 0, 0, 0},
					Body:    newValue}

				binary.BigEndian.PutUint64(req.Extras,
					uint64(flags)<<32|uint64(exp))
			case CASDelete:
				req = &gomemcached.MCRequest{
					Opcode:  gomemcached.DELETE,
					VBucket: vb,
					Key:     []byte(k),
					Cas:     orig.Cas}
			}
			resp, err := client.Send(req)
			if err == nil {
				return resp, nil
			}
		}
	}
	panic("Unreachable")
}

// Stats returns a slice of these.
type StatValue struct {
	// The stat key
	Key string
	// The stat value
	Val string
}

// Get stats from the server
// use "" as the stat key for toplevel stats.
func (client *Client) Stats(key string) ([]StatValue, error) {
	rv := make([]StatValue, 0, 128)

	req := &gomemcached.MCRequest{
		Opcode:  gomemcached.STAT,
		VBucket: 0,
		Key:     []byte(key),
		Cas:     0,
		Opaque:  918494,
		Extras:  []byte{}}

	err := transmitRequest(client.conn, req)
	if err != nil {
		return rv, err
	}

	for {
		res, err := getResponse(client.conn, client.hdrBuf)
		if err != nil {
			return rv, err
		}
		k := string(res.Key)
		if k == "" {
			break
		}
		rv = append(rv, StatValue{
			Key: k,
			Val: string(res.Body),
		})
	}

	return rv, nil
}

// Get the stats from the server as a map
func (client *Client) StatsMap(key string) (map[string]string, error) {
	rv := make(map[string]string)
	st, err := client.Stats(key)
	if err != nil {
		return rv, err
	}
	for _, sv := range st {
		rv[sv.Key] = sv.Val
	}
	return rv, nil
}
<commit_msg>Don't return success on add failure.<commit_after>// A memcached binary protocol client.
package memcached

import (
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"net"
	"strings"
	"sync"

	"github.com/dustin/gomemcached"
)

const bufsize = 1024

// The Client itself.
type Client struct {
	conn    io.ReadWriteCloser
	healthy bool

	hdrBuf []byte
}

// Connect to a memcached server.
func Connect(prot, dest string) (rv *Client, err error) {
	conn, err := net.Dial(prot, dest)
	if err != nil {
		return nil, err
	}
	return &Client{
		conn:    conn,
		healthy: true,
		hdrBuf:  make([]byte, gomemcached.HDR_LEN),
	}, nil
}

// Close the connection when you're done.
func (c *Client) Close() {
	c.conn.Close()
}

// Return false if this client has had issues communicating.
//
// This is useful for connection pools where we want to
// non-destructively determine that a connection may be reused.
func (c Client) IsHealthy() bool {
	return c.healthy
}

// Send a custom request and get the response.
func (client *Client) Send(req *gomemcached.MCRequest) (rv *gomemcached.MCResponse, err error) {
	err = transmitRequest(client.conn, req)
	if err != nil {
		client.healthy = false
		return
	}
	resp, err := getResponse(client.conn, client.hdrBuf)
	if err != nil {
		client.healthy = false
	}
	return resp, err
}

// Send a request, but do not wait for a response.
func (client *Client) Transmit(req *gomemcached.MCRequest) error {
	err := transmitRequest(client.conn, req)
	if err != nil {
		client.healthy = false
	}
	return err
}

// Receive a response
func (client *Client) Receive() (*gomemcached.MCResponse, error) {
	resp, err := getResponse(client.conn, client.hdrBuf)
	if err != nil {
		client.healthy = false
	}
	return resp, err
}

// Get the value for a key.
func (client *Client) Get(vb uint16, key string) (*gomemcached.MCResponse, error) {
	return client.Send(&gomemcached.MCRequest{
		Opcode:  gomemcached.GET,
		VBucket: vb,
		Key:     []byte(key),
		Cas:     0,
		Opaque:  0,
		Extras:  []byte{},
		Body:    []byte{}})
}

// Delete a key.
func (client *Client) Del(vb uint16, key string) (*gomemcached.MCResponse, error) {
	return client.Send(&gomemcached.MCRequest{
		Opcode:  gomemcached.DELETE,
		VBucket: vb,
		Key:     []byte(key),
		Cas:     0,
		Opaque:  0,
		Extras:  []byte{},
		Body:    []byte{}})
}

// List auth mechanisms
func (client *Client) AuthList() (*gomemcached.MCResponse, error) {
	return client.Send(&gomemcached.MCRequest{
		Opcode:  gomemcached.SASL_LIST_MECHS,
		VBucket: 0,
		Key:     []byte{},
		Cas:     0,
		Opaque:  0,
		Extras:  []byte{},
		Body:    []byte{}})
}

func (client *Client) Auth(user, pass string) (*gomemcached.MCResponse, error) {
	res, err := client.AuthList()

	if err != nil {
		return res, err
	}

	authMech := string(res.Body)
	if strings.Index(authMech, "PLAIN") != -1 {
		return client.Send(&gomemcached.MCRequest{
			Opcode:  gomemcached.SASL_AUTH,
			VBucket: 0,
			Key:     []byte("PLAIN"),
			Cas:     0,
			Opaque:  0,
			Extras:  []byte{},
			Body:    []byte(fmt.Sprintf("\x00%s\x00%s", user, pass))})
	}
	return res, fmt.Errorf("Auth mechanism PLAIN not supported")
}

func (client *Client) store(opcode gomemcached.CommandCode, vb uint16,
	key string, flags int, exp int, body []byte) (*gomemcached.MCResponse, error) {

	req := &gomemcached.MCRequest{
		Opcode:  opcode,
		VBucket: vb,
		Key:     []byte(key),
		Cas:     0,
		Opaque:  0,
		Extras:  []byte{0, 0, 0, 0, 0, 0, 0, 0},
		Body:    body}

	binary.BigEndian.PutUint64(req.Extras, uint64(flags)<<32|uint64(exp))
	return client.Send(req)
}

// Increment a value.
func (client *Client) Incr(vb uint16, key string,
	amt, def uint64, exp int) (uint64, error) {

	req := &gomemcached.MCRequest{
		Opcode:  gomemcached.INCREMENT,
		VBucket: vb,
		Key:     []byte(key),
		Cas:     0,
		Opaque:  0,
		Extras:  make([]byte, 8+8+4),
		Body:    []byte{}}
	binary.BigEndian.PutUint64(req.Extras[:8], amt)
	binary.BigEndian.PutUint64(req.Extras[8:16], def)
	binary.BigEndian.PutUint32(req.Extras[16:20], uint32(exp))

	resp, err := client.Send(req)
	if err != nil {
		return 0, err
	}

	return binary.BigEndian.Uint64(resp.Body), nil
}

// Add a value for a key (store if not exists).
func (client *Client) Add(vb uint16, key string, flags int, exp int,
	body []byte) (*gomemcached.MCResponse, error) {
	return client.store(gomemcached.ADD, vb, key, flags, exp, body)
}

// Set the value for a key.
func (client *Client) Set(vb uint16, key string, flags int, exp int,
	body []byte) (*gomemcached.MCResponse, error) {
	return client.store(gomemcached.SET, vb, key, flags, exp, body)
}

// Get keys in bulk
func (client *Client) GetBulk(vb uint16, keys []string) (map[string]*gomemcached.MCResponse, error) {
	terminalOpaque := uint32(len(keys) + 5)
	rv := map[string]*gomemcached.MCResponse{}
	wg := sync.WaitGroup{}
	going := true

	defer func() {
		going = false
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		for going {
			res, err := client.Receive()
			if err != nil {
				return
			}
			if res.Opaque == terminalOpaque {
				return
			}
			if res.Opcode != gomemcached.GETQ {
				log.Panicf("Unexpected opcode in GETQ response: %+v",
					res)
			}
			rv[keys[res.Opaque]] = res
		}
	}()

	for i, k := range keys {
		err := client.Transmit(&gomemcached.MCRequest{
			Opcode:  gomemcached.GETQ,
			VBucket: vb,
			Key:     []byte(k),
			Cas:     0,
			Opaque:  uint32(i),
			Extras:  []byte{},
			Body:    []byte{}})
		if err != nil {
			return rv, err
		}
	}

	err := client.Transmit(&gomemcached.MCRequest{
		Opcode: gomemcached.NOOP,
		Key:    []byte{},
		Cas:    0,
		Extras: []byte{},
		Body:   []byte{},
		Opaque: terminalOpaque})
	if err != nil {
		return rv, err
	}

	wg.Wait()

	return rv, nil
}

// Operation to perform on this CAS loop.
type CasOp uint8

const (
	// Store the new value normally
	CASStore = CasOp(iota)
	// Stop attempting to CAS, leave value untouched
	CASQuit
	// Delete the current value
	CASDelete
)

// User specified termination is returned as an error.
func (c CasOp) Error() string {
	switch c {
	case CASStore:
		return "CAS store"
	case CASQuit:
		return "CAS quit"
	case CASDelete:
		return "CAS delete"
	}
	panic("Unhandled value")
}

// A function to perform a CAS transform
type CasFunc func(current []byte) ([]byte, CasOp)

// Perform a CAS transform with the given function.
//
// If the value does not exist, an empty byte string will be sent to f
func (client *Client) CAS(vb uint16, k string, f CasFunc,
	initexp int) (rv *gomemcached.MCResponse, err error) {

	flags := 0
	exp := 0

	for {
		orig, err := client.Get(vb, k)
		if err != nil && (orig == nil || orig.Status != gomemcached.KEY_ENOENT) {
			return rv, err
		}

		if orig.Status == gomemcached.KEY_ENOENT {
			init, operation := f([]byte{})
			if operation == CASQuit || operation == CASDelete {
				return nil, operation
			}
			// If it doesn't exist, add it
			resp, err := client.Add(vb, k, 0, initexp, init)
			if err == nil && resp.Status != gomemcached.KEY_EEXISTS {
				return rv, err
			}
		} else {
			var req *gomemcached.MCRequest
			newValue, operation := f(orig.Body)

			switch operation {
			case CASQuit:
				return nil, operation
			case CASStore:
				req = &gomemcached.MCRequest{
					Opcode:  gomemcached.SET,
					VBucket: vb,
					Key:     []byte(k),
					Cas:     orig.Cas,
					Opaque:  0,
					Extras:  []byte{0, 0, 0, 0, 0, 0, 0, 0},
					Body:    newValue}

				binary.BigEndian.PutUint64(req.Extras,
					uint64(flags)<<32|uint64(exp))
			case CASDelete:
				req = &gomemcached.MCRequest{
					Opcode:  gomemcached.DELETE,
					VBucket: vb,
					Key:     []byte(k),
					Cas:     orig.Cas}
			}
			resp, err := client.Send(req)
			if err == nil {
				return resp, nil
			}
		}
	}
	panic("Unreachable")
}

// Stats returns a slice of these.
type StatValue struct {
	// The stat key
	Key string
	// The stat value
	Val string
}

// Get stats from the server
// use "" as the stat key for toplevel stats.
func (client *Client) Stats(key string) ([]StatValue, error) {
	rv := make([]StatValue, 0, 128)

	req := &gomemcached.MCRequest{
		Opcode:  gomemcached.STAT,
		VBucket: 0,
		Key:     []byte(key),
		Cas:     0,
		Opaque:  918494,
		Extras:  []byte{}}

	err := transmitRequest(client.conn, req)
	if err != nil {
		return rv, err
	}

	for {
		res, err := getResponse(client.conn, client.hdrBuf)
		if err != nil {
			return rv, err
		}
		k := string(res.Key)
		if k == "" {
			break
		}
		rv = append(rv, StatValue{
			Key: k,
			Val: string(res.Body),
		})
	}

	return rv, nil
}

// Get the stats from the server as a map
func (client *Client) StatsMap(key string) (map[string]string, error) {
	rv := make(map[string]string)
	st, err := client.Stats(key)
	if err != nil {
		return rv, err
	}
	for _, sv := range st {
		rv[sv.Key] = sv.Val
	}
	return rv, nil
}
<|endoftext|>