<commit_before>package influxdb

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"github.com/bmizerany/pat"
	"github.com/influxdb/influxdb/influxql"
)

// TODO: Standard response headers (see: HeaderHandler)
// TODO: Compression (see: CompressionHeaderHandler)

// TODO: Check HTTP response codes: 400, 401, 403, 409.

// getUsernameAndPassword returns the username and password encoded in
// a request. The credentials may be present as URL query params, or as
// a Basic Authentication header.
func getUsernameAndPassword(r *http.Request) (string, string, error) {
	q := r.URL.Query()
	username, password := q.Get("u"), q.Get("p")
	if username != "" && password != "" {
		return username, password, nil
	}
	auth := r.Header.Get("Authorization")
	if auth == "" {
		return "", "", nil
	}
	fields := strings.Split(auth, " ")
	if len(fields) != 2 {
		return "", "", fmt.Errorf("invalid Basic Authentication header")
	}
	bs, err := base64.StdEncoding.DecodeString(fields[1])
	if err != nil {
		return "", "", fmt.Errorf("invalid Base64 encoding")
	}
	fields = strings.Split(string(bs), ":")
	if len(fields) != 2 {
		return "", "", fmt.Errorf("invalid Basic Authentication value")
	}
	return fields[0], fields[1], nil
}

// Handler represents an HTTP handler for the InfluxDB server.
type Handler struct {
	server *Server
	mux    *pat.PatternServeMux

	// Whether endpoints require authentication.
	AuthenticationEnabled bool

	// The InfluxDB verion returned by the HTTP response header.
	Version string
}

// NewHandler returns a new instance of Handler.
func NewHandler(s *Server) *Handler {
	h := &Handler{
		server: s,
		mux:    pat.New(),
	}

	// Query serving route.
	h.mux.Get("/query", h.makeAuthenticationHandler(h.serveQuery))

	// Data-ingest route.
	h.mux.Post("/db/:db/series", h.makeAuthenticationHandler(h.serveWriteSeries))

	// Data node routes.
	h.mux.Get("/data_nodes", h.makeAuthenticationHandler(h.serveDataNodes))
	h.mux.Post("/data_nodes", h.makeAuthenticationHandler(h.serveCreateDataNode))
	h.mux.Del("/data_nodes/:id", h.makeAuthenticationHandler(h.serveDeleteDataNode))

	// Utilities
	h.mux.Get("/metastore", h.makeAuthenticationHandler(h.serveMetastore))
	h.mux.Get("/ping", h.makeAuthenticationHandler(h.servePing))

	return h
}

// ServeHTTP responds to HTTP request to the handler.
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Add("Access-Control-Allow-Origin", "*")
	w.Header().Add("Access-Control-Max-Age", "2592000")
	w.Header().Add("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE")
	w.Header().Add("Access-Control-Allow-Headers", "Origin, X-Requested-With, Content-Type, Accept")
	w.Header().Add("X-Influxdb-Version", h.Version)

	// If this is a CORS OPTIONS request then send back okie-dokie.
	if r.Method == "OPTIONS" {
		w.WriteHeader(http.StatusOK)
		return
	}

	// Otherwise handle it via pat.
	h.mux.ServeHTTP(w, r)
}

// makeAuthenticationHandler takes a custom handler and returns a standard handler, ensuring that
// if user credentials are passed in, an attempt is made to authenticate that user. If authentication
// fails, an error is returned to the user.
//
// There is one exception: if there are no users in the system, authentication is not required. This
// is to facilitate bootstrapping of a system with authentication enabled.
func (h *Handler) makeAuthenticationHandler(fn func(http.ResponseWriter, *http.Request, *User)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var user *User
		if h.AuthenticationEnabled && len(h.server.Users()) > 0 {
			username, password, err := getUsernameAndPassword(r)
			if err != nil {
				h.error(w, err.Error(), http.StatusUnauthorized)
				return
			}
			if username == "" {
				h.error(w, "username required", http.StatusUnauthorized)
				return
			}

			user, err = h.server.Authenticate(username, password)
			if err != nil {
				h.error(w, err.Error(), http.StatusUnauthorized)
				return
			}
		}
		fn(w, r, user)
	}
}

// serveQuery parses an incoming query and returns the results.
func (h *Handler) serveQuery(w http.ResponseWriter, r *http.Request, u *User) {
	q := r.URL.Query()
	p := influxql.NewParser(strings.NewReader(q.Get("q")))

	// Parse query from query string.
	query, err := p.ParseQuery()
	if err != nil {
		h.error(w, "error parsing query: "+err.Error(), http.StatusBadRequest)
		return
	}

	// Bootstrapping a secured cluster involves a special authentication policy.
	// Until an admin-level user exists on the cluster one, and only one operation
	// is permitted without authentication -- the creation of the first user.
	// This user is also automatically granted admin-level access.
	if h.AuthenticationEnabled && len(h.server.Users()) > 0 {
		if len(query.Statements) != 1 {
			_, ok := query.Statements[0].(*influxql.CreateUserStatement)
			if !ok {
				h.error(w, "initial admin user does not exist", http.StatusUnauthorized)
				return
			}
		}
	}

	for _, s := range query.Statements {
		switch c := s.(type) {
		case *influxql.CreateDatabaseStatement:
			err = h.server.CreateDatabase(c.Name)
		case *influxql.DropDatabaseStatement:
			err = h.server.DeleteDatabase(c.Name)

		case *influxql.CreateUserStatement:
			err = h.server.CreateUser(c.Name, c.Password, !h.server.AdminUserExists()) // XXX This is obviously broken.
		case *influxql.DropUserStatement:
			err = h.server.DeleteUser(c.Name)

		case *influxql.SelectStatement:
			continue
		case *influxql.DropSeriesStatement:
			continue
		case *influxql.ListSeriesStatement:
			continue
		case *influxql.ListMeasurementsStatement:
			continue
		case *influxql.ListTagKeysStatement:
			continue
		case *influxql.ListTagValuesStatement:
			continue
		case *influxql.ListFieldKeysStatement:
			continue
		case *influxql.ListFieldValuesStatement:
			continue

		case *influxql.GrantStatement:
			continue
		case *influxql.RevokeStatement:
			continue

		case *influxql.CreateRetentionPolicyStatement:
			rp := NewRetentionPolicy(c.Name)
			rp.Duration = c.Duration
			rp.ReplicaN = uint32(c.Replication)
			err = h.server.CreateRetentionPolicy(c.DB, rp)
		case *influxql.AlterRetentionPolicyStatement:
			rp := NewRetentionPolicy(c.Name)
			rp.Duration = *c.Duration // Why is this a pointer, and the next?
			rp.ReplicaN = uint32(*c.Replication)
			err = h.server.UpdateRetentionPolicy(c.DB, c.Name, rp)

		case *influxql.CreateContinuousQueryStatement:
			continue
		case *influxql.DropContinuousQueryStatement:
			continue
		case *influxql.ListContinuousQueriesStatement:
			continue
		}

		if err != nil {
			h.error(w, "error executing query: "+err.Error(), http.StatusBadRequest)
		}
	}
}

// serveWriteSeries receives incoming series data and writes it to the database.
func (h *Handler) serveWriteSeries(w http.ResponseWriter, r *http.Request, u *User) {
	// TODO: Authentication.

	/* TEMPORARILY REMOVED FOR PROTOBUFS.
	// Retrieve database from server.
	db := h.server.Database(r.URL.Query().Get(":db"))
	if db == nil {
		h.error(w, ErrDatabaseNotFound.Error(), http.StatusNotFound)
		return
	}

	// Parse time precision from query parameters.
	precision, err := parseTimePrecision(r.URL.Query().Get("time_precision"))
	if err != nil {
		h.error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Setup HTTP request reader. Wrap in a gzip reader if encoding set in header.
	reader := r.Body
	if r.Header.Get("Content-Encoding") == "gzip" {
		if reader, err = gzip.NewReader(r.Body); err != nil {
			h.error(w, err.Error(), http.StatusBadRequest)
			return
		}
	}

	// Decode series from reader.
	ss := []*serializedSeries{}
	dec := json.NewDecoder(reader)
	dec.UseNumber()
	if err := dec.Decode(&ss); err != nil {
		h.error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Convert the wire format to the internal representation of the time series.
	series, err := serializedSeriesSlice(ss).series(precision)
	if err != nil {
		h.error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Write series data to the database.
	// TODO: Allow multiple series written to DB at once.
	for _, s := range series {
		if err := db.WriteSeries(s); err != nil {
			h.error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
	*/
}

// serveMetastore returns a copy of the metastore.
func (h *Handler) serveMetastore(w http.ResponseWriter, r *http.Request, u *User) {
	// Set headers.
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("Content-Disposition", `attachment; filename="meta"`)

	// Write metastore to response body.
	if err := h.server.CopyMetastore(w); err != nil {
		h.error(w, err.Error(), http.StatusInternalServerError)
	}
}

// servePing returns a simple response to let the client know the server is running.
func (h *Handler) servePing(w http.ResponseWriter, r *http.Request, u *User) {}

// serveDataNodes returns a list of all data nodes in the cluster.
func (h *Handler) serveDataNodes(w http.ResponseWriter, r *http.Request, u *User) {
	// Generate a list of objects for encoding to the API.
	a := make([]*dataNodeJSON, 0)
	for _, n := range h.server.DataNodes() {
		a = append(a, &dataNodeJSON{
			ID:  n.ID,
			URL: n.URL.String(),
		})
	}

	w.Header().Add("content-type", "application/json")
	_ = json.NewEncoder(w).Encode(a)
}

// serveCreateDataNode creates a new data node in the cluster.
func (h *Handler) serveCreateDataNode(w http.ResponseWriter, r *http.Request, _ *User) {
	// Read in data node from request body.
	var n dataNodeJSON
	if err := json.NewDecoder(r.Body).Decode(&n); err != nil {
		h.error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Parse the URL.
	u, err := url.Parse(n.URL)
	if err != nil {
		h.error(w, "invalid data node url", http.StatusBadRequest)
		return
	}

	// Create the data node.
	if err := h.server.CreateDataNode(u); err == ErrDataNodeExists {
		h.error(w, err.Error(), http.StatusConflict)
		return
	} else if err != nil {
		h.error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// Retrieve data node reference.
	node := h.server.DataNodeByURL(u)

	// Create a new replica on the broker.
	if err := h.server.client.CreateReplica(node.ID); err != nil {
		h.error(w, err.Error(), http.StatusBadGateway)
		return
	}

	// Write new node back to client.
	w.WriteHeader(http.StatusCreated)
	w.Header().Add("content-type", "application/json")
	_ = json.NewEncoder(w).Encode(&dataNodeJSON{ID: node.ID, URL: node.URL.String()})
}

// serveDeleteDataNode removes an existing node.
func (h *Handler) serveDeleteDataNode(w http.ResponseWriter, r *http.Request, u *User) {
	// Parse node id.
	nodeID, err := strconv.ParseUint(r.URL.Query().Get(":id"), 10, 64)
	if err != nil {
		h.error(w, "invalid node id", http.StatusBadRequest)
		return
	}

	// Delete the node.
	if err := h.server.DeleteDataNode(nodeID); err == ErrDataNodeNotFound {
		h.error(w, err.Error(), http.StatusNotFound)
		return
	} else if err != nil {
		h.error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

type dataNodeJSON struct {
	ID  uint64 `json:"id"`
	URL string `json:"url"`
}

// error returns an error to the client in a standard format.
func (h *Handler) error(w http.ResponseWriter, error string, code int) {
	// TODO: Return error as JSON.
	http.Error(w, error, code)
}
<commit_msg>Some initial work<commit_after>package influxdb

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"github.com/bmizerany/pat"
	"github.com/influxdb/influxdb/influxql"
)

// TODO: Standard response headers (see: HeaderHandler)
// TODO: Compression (see: CompressionHeaderHandler)

// TODO: Check HTTP response codes: 400, 401, 403, 409.

// getUsernameAndPassword returns the username and password encoded in
// a request. The credentials may be present as URL query params, or as
// a Basic Authentication header.
func getUsernameAndPassword(r *http.Request) (string, string, error) {
	q := r.URL.Query()
	username, password := q.Get("u"), q.Get("p")
	if username != "" && password != "" {
		return username, password, nil
	}
	auth := r.Header.Get("Authorization")
	if auth == "" {
		return "", "", nil
	}
	fields := strings.Split(auth, " ")
	if len(fields) != 2 {
		return "", "", fmt.Errorf("invalid Basic Authentication header")
	}
	bs, err := base64.StdEncoding.DecodeString(fields[1])
	if err != nil {
		return "", "", fmt.Errorf("invalid Base64 encoding")
	}
	fields = strings.Split(string(bs), ":")
	if len(fields) != 2 {
		return "", "", fmt.Errorf("invalid Basic Authentication value")
	}
	return fields[0], fields[1], nil
}

// Handler represents an HTTP handler for the InfluxDB server.
type Handler struct {
	server *Server
	mux    *pat.PatternServeMux

	// Whether endpoints require authentication.
	AuthenticationEnabled bool

	// The InfluxDB verion returned by the HTTP response header.
	Version string
}

// NewHandler returns a new instance of Handler.
func NewHandler(s *Server) *Handler {
	h := &Handler{
		server: s,
		mux:    pat.New(),
	}

	// Query serving route.
	h.mux.Get("/query", h.makeAuthenticationHandler(h.serveQuery))

	// Data-ingest route.
	h.mux.Post("/db/:db/series", h.makeAuthenticationHandler(h.serveWriteSeries))

	// Data node routes.
	h.mux.Get("/data_nodes", h.makeAuthenticationHandler(h.serveDataNodes))
	h.mux.Post("/data_nodes", h.makeAuthenticationHandler(h.serveCreateDataNode))
	h.mux.Del("/data_nodes/:id", h.makeAuthenticationHandler(h.serveDeleteDataNode))

	// Utilities
	h.mux.Get("/metastore", h.makeAuthenticationHandler(h.serveMetastore))
	h.mux.Get("/ping", h.makeAuthenticationHandler(h.servePing))

	return h
}

// ServeHTTP responds to HTTP request to the handler.
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Add("Access-Control-Allow-Origin", "*")
	w.Header().Add("Access-Control-Max-Age", "2592000")
	w.Header().Add("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE")
	w.Header().Add("Access-Control-Allow-Headers", "Origin, X-Requested-With, Content-Type, Accept")
	w.Header().Add("X-Influxdb-Version", h.Version)

	// If this is a CORS OPTIONS request then send back okie-dokie.
	if r.Method == "OPTIONS" {
		w.WriteHeader(http.StatusOK)
		return
	}

	// Otherwise handle it via pat.
	h.mux.ServeHTTP(w, r)
}

// makeAuthenticationHandler takes a custom handler and returns a standard handler, ensuring that
// if user credentials are passed in, an attempt is made to authenticate that user. If authentication
// fails, an error is returned to the user.
//
// There is one exception: if there are no users in the system, authentication is not required. This
// is to facilitate bootstrapping of a system with authentication enabled.
func (h *Handler) makeAuthenticationHandler(fn func(http.ResponseWriter, *http.Request, *User)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var user *User
		if h.AuthenticationEnabled && len(h.server.Users()) > 0 {
			username, password, err := getUsernameAndPassword(r)
			if err != nil {
				h.error(w, err.Error(), http.StatusUnauthorized)
				return
			}
			if username == "" {
				h.error(w, "username required", http.StatusUnauthorized)
				return
			}

			user, err = h.server.Authenticate(username, password)
			if err != nil {
				h.error(w, err.Error(), http.StatusUnauthorized)
				return
			}
		}
		fn(w, r, user)
	}
}

// serveQuery parses an incoming query and returns the results.
func (h *Handler) serveQuery(w http.ResponseWriter, r *http.Request, u *User) {
	q := r.URL.Query()
	p := influxql.NewParser(strings.NewReader(q.Get("q")))

	// Parse query from query string.
	query, err := p.ParseQuery()
	if err != nil {
		h.error(w, "error parsing query: "+err.Error(), http.StatusBadRequest)
		return
	}

	// Bootstrapping a secured cluster involves a special authentication policy.
	// Until an admin-level user exists on the cluster one, and only one operation
	// is permitted without authentication -- the creation of the first user.
	// This user is also automatically granted admin-level access.
	if h.AuthenticationEnabled && len(h.server.Users()) > 0 {
		if len(query.Statements) != 1 {
			_, ok := query.Statements[0].(*influxql.CreateUserStatement)
			if !ok {
				h.error(w, "initial admin user does not exist", http.StatusUnauthorized)
				return
			}
		}
	}

	for _, s := range query.Statements {
		switch c := s.(type) {
		case *influxql.CreateDatabaseStatement:
			err = h.server.CreateDatabase(c.Name)
		case *influxql.DropDatabaseStatement:
			err = h.server.DeleteDatabase(c.Name)

		case *influxql.CreateUserStatement:
			err = h.server.CreateUser(c.Name, c.Password, !h.server.AdminUserExists()) // XXX This is obviously broken.
		case *influxql.DropUserStatement:
			err = h.server.DeleteUser(c.Name)

		case *influxql.SelectStatement:
			continue
		case *influxql.DropSeriesStatement:
			continue
		case *influxql.ListSeriesStatement:
			continue
		case *influxql.ListMeasurementsStatement:
			continue
		case *influxql.ListTagKeysStatement:
			continue
		case *influxql.ListTagValuesStatement:
			continue
		case *influxql.ListFieldKeysStatement:
			continue
		case *influxql.ListFieldValuesStatement:
			continue

		case *influxql.GrantStatement:
			continue
		case *influxql.RevokeStatement:
			continue

		case *influxql.CreateRetentionPolicyStatement:
			rp := NewRetentionPolicy(c.Name)
			rp.Duration = c.Duration
			rp.ReplicaN = uint32(c.Replication)
			err = h.server.CreateRetentionPolicy(c.Database, rp)
		case *influxql.AlterRetentionPolicyStatement:
			rp := NewRetentionPolicy(c.Name)
			rp.Duration = *c.Duration // Why is this a pointer, and the next?
			rp.ReplicaN = uint32(*c.Replication)
			err = h.server.UpdateRetentionPolicy(c.Database, c.Name, rp)

		case *influxql.CreateContinuousQueryStatement:
			continue
		case *influxql.DropContinuousQueryStatement:
			continue
		case *influxql.ListContinuousQueriesStatement:
			continue
		}

		if err != nil {
			h.error(w, "error executing query: "+err.Error(), http.StatusBadRequest)
		}
	}
}

// serveWriteSeries receives incoming series data and writes it to the database.
func (h *Handler) serveWriteSeries(w http.ResponseWriter, r *http.Request, u *User) {
	// TODO: Authentication.

	/* TEMPORARILY REMOVED FOR PROTOBUFS.
	// Retrieve database from server.
	db := h.server.Database(r.URL.Query().Get(":db"))
	if db == nil {
		h.error(w, ErrDatabaseNotFound.Error(), http.StatusNotFound)
		return
	}

	// Parse time precision from query parameters.
	precision, err := parseTimePrecision(r.URL.Query().Get("time_precision"))
	if err != nil {
		h.error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Setup HTTP request reader. Wrap in a gzip reader if encoding set in header.
	reader := r.Body
	if r.Header.Get("Content-Encoding") == "gzip" {
		if reader, err = gzip.NewReader(r.Body); err != nil {
			h.error(w, err.Error(), http.StatusBadRequest)
			return
		}
	}

	// Decode series from reader.
	ss := []*serializedSeries{}
	dec := json.NewDecoder(reader)
	dec.UseNumber()
	if err := dec.Decode(&ss); err != nil {
		h.error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Convert the wire format to the internal representation of the time series.
	series, err := serializedSeriesSlice(ss).series(precision)
	if err != nil {
		h.error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Write series data to the database.
	// TODO: Allow multiple series written to DB at once.
	for _, s := range series {
		if err := db.WriteSeries(s); err != nil {
			h.error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
	*/
}

// serveMetastore returns a copy of the metastore.
func (h *Handler) serveMetastore(w http.ResponseWriter, r *http.Request, u *User) {
	// Set headers.
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("Content-Disposition", `attachment; filename="meta"`)

	if err := h.server.CopyMetastore(w); err != nil {
		h.error(w, err.Error(), http.StatusInternalServerError)
	}
}

// servePing returns a simple response to let the client know the server is running.
func (h *Handler) servePing(w http.ResponseWriter, r *http.Request, u *User) {}

// serveDataNodes returns a list of all data nodes in the cluster.
func (h *Handler) serveDataNodes(w http.ResponseWriter, r *http.Request, u *User) {
	// Generate a list of objects for encoding to the API.
	a := make([]*dataNodeJSON, 0)
	for _, n := range h.server.DataNodes() {
		a = append(a, &dataNodeJSON{
			ID:  n.ID,
			URL: n.URL.String(),
		})
	}

	w.Header().Add("content-type", "application/json")
	_ = json.NewEncoder(w).Encode(a)
}

// serveCreateDataNode creates a new data node in the cluster.
func (h *Handler) serveCreateDataNode(w http.ResponseWriter, r *http.Request, _ *User) {
	// Read in data node from request body.
	var n dataNodeJSON
	if err := json.NewDecoder(r.Body).Decode(&n); err != nil {
		h.error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Parse the URL.
	u, err := url.Parse(n.URL)
	if err != nil {
		h.error(w, "invalid data node url", http.StatusBadRequest)
		return
	}

	// Create the data node.
	if err := h.server.CreateDataNode(u); err == ErrDataNodeExists {
		h.error(w, err.Error(), http.StatusConflict)
		return
	} else if err != nil {
		h.error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// Retrieve data node reference.
	node := h.server.DataNodeByURL(u)

	// Create a new replica on the broker.
	if err := h.server.client.CreateReplica(node.ID); err != nil {
		h.error(w, err.Error(), http.StatusBadGateway)
		return
	}

	// Write new node back to client.
	w.WriteHeader(http.StatusCreated)
	w.Header().Add("content-type", "application/json")
	_ = json.NewEncoder(w).Encode(&dataNodeJSON{ID: node.ID, URL: node.URL.String()})
}

// serveDeleteDataNode removes an existing node.
func (h *Handler) serveDeleteDataNode(w http.ResponseWriter, r *http.Request, u *User) {
	// Parse node id.
	nodeID, err := strconv.ParseUint(r.URL.Query().Get(":id"), 10, 64)
	if err != nil {
		h.error(w, "invalid node id", http.StatusBadRequest)
		return
	}

	// Delete the node.
	if err := h.server.DeleteDataNode(nodeID); err == ErrDataNodeNotFound {
		h.error(w, err.Error(), http.StatusNotFound)
		return
	} else if err != nil {
		h.error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

type dataNodeJSON struct {
	ID  uint64 `json:"id"`
	URL string `json:"url"`
}

// error returns an error to the client in a standard format.
func (h *Handler) error(w http.ResponseWriter, error string, code int) {
	// TODO: Return error as JSON.
	http.Error(w, error, code)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"runtime"
	"strconv"
	"time"

	cniversion "github.com/containernetworking/cni/pkg/version"
	"github.com/containers/buildah/define"
	iversion "github.com/containers/image/v5/version"
	ispecs "github.com/opencontainers/image-spec/specs-go"
	rspecs "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/spf13/cobra"
)

//Overwritten at build time
var (
	GitCommit  string
	buildInfo  string
	cniVersion string
)

type versionInfo struct {
	Version       string `json:"version"`
	GoVersion     string `json:"goVersion"`
	ImageSpec     string `json:"imageSpec"`
	RuntimeSpec   string `json:"runtimeSpec"`
	CniSpec       string `json:"cniSpec"`
	LibcniVersion string `json:"libcniVersion"`
	ImageVersion  string `json:"imageVersion"`
	GitCommit     string `json:"gitCommit"`
	Built         string `json:"built"`
	OsArch        string `json:"osArch"`
}

type versionOptions struct {
	json bool
}

func init() {
	var opts versionOptions

	//cli command to print out the version info of buildah
	versionCommand := &cobra.Command{
		Use:   "version",
		Short: "Display the Buildah version information",
		Long:  "Displays Buildah version information.",
		RunE: func(c *cobra.Command, args []string) error {
			var err error
			buildTime := int64(0)
			if buildInfo != "" {
				//converting unix time from string to int64
				buildTime, err = strconv.ParseInt(buildInfo, 10, 64)
				if err != nil {
					return err
				}
			}

			version := versionInfo{
				Version:       define.Version,
				GoVersion:     runtime.Version(),
				ImageSpec:     ispecs.Version,
				RuntimeSpec:   rspecs.Version,
				CniSpec:       cniversion.Current(),
				LibcniVersion: cniVersion,
				ImageVersion:  iversion.Version,
				GitCommit:     GitCommit,
				Built:         time.Unix(buildTime, 0).Format(time.ANSIC),
				OsArch:        runtime.GOOS + "/" + runtime.GOARCH,
			}

			if opts.json {
				data, err := json.MarshalIndent(version, "", "    ")
				if err != nil {
					return err
				}
				fmt.Printf("%s\n", data)
				return nil
			}

			fmt.Println("Version:        ", version.Version)
			fmt.Println("Go Version:     ", version.GoVersion)
			fmt.Println("Image Spec:     ", version.ImageSpec)
			fmt.Println("Runtime Spec:   ", version.RuntimeSpec)
			fmt.Println("CNI Spec:       ", version.CniSpec)
			fmt.Println("libcni Version: ", version.LibcniVersion)
			fmt.Println("image Version:  ", version.ImageVersion)
			fmt.Println("Git Commit:     ", version.GitCommit)

			//Prints out the build time in readable format
			fmt.Println("Built:          ", version.Built)
			fmt.Println("OS/Arch:        ", version.OsArch)

			return nil
		},
		Args:    cobra.NoArgs,
		Example: `buildah version`,
	}
	versionCommand.SetUsageTemplate(UsageTemplate())

	flags := versionCommand.Flags()
	flags.BoolVar(&opts.json, "json", false, "output in JSON format")

	rootCmd.AddCommand(versionCommand)
}
<commit_msg>Extract version command implementation to function<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"runtime"
	"strconv"
	"time"

	cniversion "github.com/containernetworking/cni/pkg/version"
	"github.com/containers/buildah/define"
	iversion "github.com/containers/image/v5/version"
	ispecs "github.com/opencontainers/image-spec/specs-go"
	rspecs "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/spf13/cobra"
)

//Overwritten at build time
var (
	GitCommit  string
	buildInfo  string
	cniVersion string
)

type versionInfo struct {
	Version       string `json:"version"`
	GoVersion     string `json:"goVersion"`
	ImageSpec     string `json:"imageSpec"`
	RuntimeSpec   string `json:"runtimeSpec"`
	CniSpec       string `json:"cniSpec"`
	LibcniVersion string `json:"libcniVersion"`
	ImageVersion  string `json:"imageVersion"`
	GitCommit     string `json:"gitCommit"`
	Built         string `json:"built"`
	OsArch        string `json:"osArch"`
}

type versionOptions struct {
	json bool
}

func init() {
	var opts versionOptions

	//cli command to print out the version info of buildah
	versionCommand := &cobra.Command{
		Use:   "version",
		Short: "Display the Buildah version information",
		Long:  "Displays Buildah version information.",
		RunE: func(c *cobra.Command, args []string) error {
			return versionCmd(c, args, opts)
		},
		Args:    cobra.NoArgs,
		Example: `buildah version`,
	}
	versionCommand.SetUsageTemplate(UsageTemplate())

	flags := versionCommand.Flags()
	flags.BoolVar(&opts.json, "json", false, "output in JSON format")

	rootCmd.AddCommand(versionCommand)
}

func versionCmd(c *cobra.Command, args []string, opts versionOptions) error {
	var err error
	buildTime := int64(0)
	if buildInfo != "" {
		//converting unix time from string to int64
		buildTime, err = strconv.ParseInt(buildInfo, 10, 64)
		if err != nil {
			return err
		}
	}

	version := versionInfo{
		Version:       define.Version,
		GoVersion:     runtime.Version(),
		ImageSpec:     ispecs.Version,
		RuntimeSpec:   rspecs.Version,
		CniSpec:       cniversion.Current(),
		LibcniVersion: cniVersion,
		ImageVersion:  iversion.Version,
		GitCommit:     GitCommit,
		Built:         time.Unix(buildTime, 0).Format(time.ANSIC),
		OsArch:        runtime.GOOS + "/" + runtime.GOARCH,
	}

	if opts.json {
		data, err := json.MarshalIndent(version, "", "    ")
		if err != nil {
			return err
		}
		fmt.Printf("%s\n", data)
		return nil
	}

	fmt.Println("Version:        ", version.Version)
	fmt.Println("Go Version:     ", version.GoVersion)
	fmt.Println("Image Spec:     ", version.ImageSpec)
	fmt.Println("Runtime Spec:   ", version.RuntimeSpec)
	fmt.Println("CNI Spec:       ", version.CniSpec)
	fmt.Println("libcni Version: ", version.LibcniVersion)
	fmt.Println("image Version:  ", version.ImageVersion)
	fmt.Println("Git Commit:     ", version.GitCommit)

	//Prints out the build time in readable format
	fmt.Println("Built:          ", version.Built)
	fmt.Println("OS/Arch:        ", version.OsArch)

	return nil
}
<|endoftext|>


<commit_before>package plugins

type Handler interface {
	HandleEvent(Event, Args)
}

type HandlerFunc func(Event, Args)

func (hf HandlerFunc) HandleEvent(event Event, args Args) {
	hf(event, args)
}

type Muxer interface {
	Handler
	AddHandler(Event, Handler)
	RemoveHandler(Event, Handler)
}

type mapMuxr map[Event]Handler

func (m mapMuxr) HandleEvent(event Event, args Args) {
	m[event].HandleEvent(event, args)
}

func (m mapMuxr) AddHandler(event Event, handlr Handler) {
	m[event] = handlr
}

func (m mapMuxr) RemoveHandler(event Event, handlr Handler) {
	if h, ok := m[event]; ok && h == handlr {
		delete(m, event)
	}
}
<commit_msg>Added checks to prevent a panic when recieving event not listened to by main application.<commit_after>package plugins

type Handler interface {
	HandleEvent(Event, Args)
}

type HandlerFunc func(Event, Args)

func (hf HandlerFunc) HandleEvent(event Event, args Args) {
	hf(event, args)
}

type Muxer interface {
	Handler
	AddHandler(Event, Handler)
	RemoveHandler(Event, Handler)
}

type mapMuxr map[Event]Handler

func (m mapMuxr) HandleEvent(event Event, args Args) {
	if handler, ok := m[event]; ok {
		handler.HandleEvent(event, args)
	}
}

func (m mapMuxr) AddHandler(event Event, handlr Handler) {
	m[event] = handlr
}

func (m mapMuxr) RemoveHandler(event Event, handlr Handler) {
	if h, ok := m[event]; ok && h == handlr {
		delete(m, event)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package service

import (
	libkb "github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/go-framed-msgpack-rpc/rpc"
	context "golang.org/x/net/context"
)

type transporterAndConnectionID struct {
	transporter  rpc.Transporter
	connectionID libkb.ConnectionID
}

type getObj struct {
	ui    libkb.UIKind
	retCh chan<- transporterAndConnectionID
}

type uiWrapper struct {
	cid       libkb.ConnectionID
	sessionID int
}

type setObj struct {
	cid libkb.ConnectionID
	ui  libkb.UIKind
}

type UIRouter struct {
	libkb.Contextified
	cm         *libkb.ConnectionManager
	uis        map[libkb.UIKind]libkb.ConnectionID
	setCh      chan setObj
	getCh      chan getObj
	shutdownCh chan struct{}
}

func NewUIRouter(g *libkb.GlobalContext) *UIRouter {
	ret := &UIRouter{
		Contextified: libkb.NewContextified(g),
		cm:           g.ConnectionManager,
		uis:          make(map[libkb.UIKind]libkb.ConnectionID),
		setCh:        make(chan setObj),
		getCh:        make(chan getObj),
		shutdownCh:   make(chan struct{}),
	}
	go ret.run()
	return ret
}

func (u *UIRouter) Shutdown() {
	u.shutdownCh <- struct{}{}
}

func (u *UIRouter) run() {
	for {
		select {
		case <-u.shutdownCh:
			return
		case o := <-u.setCh:
			u.uis[o.ui] = o.cid
		case o := <-u.getCh:
			var ret rpc.Transporter
			cid, ok := u.uis[o.ui]
			if ok {
				if ret = u.cm.LookupConnection(cid); ret == nil {
					u.G().Log.Debug("UIRouter: connection %v inactive, deleting registered UI %s", cid, o.ui)
					delete(u.uis, o.ui)
				}
			}
			o.retCh <- transporterAndConnectionID{transporter: ret, connectionID: cid}
		}
	}
}

func (u *UIRouter) SetUI(c libkb.ConnectionID, k libkb.UIKind) {
	u.G().Log.Debug("UIRouter: connection %v registering UI %s", c, k)
	u.setCh <- setObj{c, k}
}

func (u *UIRouter) getUI(k libkb.UIKind) (rpc.Transporter, libkb.ConnectionID) {
	retCh := make(chan transporterAndConnectionID)
	u.getCh <- getObj{k, retCh}
	ret := <-retCh
	return ret.transporter, ret.connectionID
}

func (u *UIRouter) GetIdentifyUI() (libkb.IdentifyUI, error) {
	x, _ := u.getUI(libkb.IdentifyUIKind)
	if x == nil {
		return nil, nil
	}
	cli := rpc.NewClient(x, libkb.ErrorUnwrapper{}, nil)
	iuicli := keybase1.IdentifyUiClient{Cli: cli}
	sessionID, err := iuicli.DelegateIdentifyUI(context.TODO())
	if err != nil {
		return nil, err
	}
	ret := &RemoteIdentifyUI{
		sessionID: sessionID,
		uicli:     iuicli,
		logUI: &LogUI{
			sessionID,
			&keybase1.LogUiClient{Cli: cli},
		},
		Contextified: libkb.NewContextified(u.G()),
	}
	return ret, nil
}

func (u *UIRouter) GetIdentifyUICtx(ctx context.Context) (int, libkb.IdentifyUI, error) {
	x, _ := u.getUI(libkb.IdentifyUIKind)
	if x == nil {
		return 0, nil, nil
	}
	cli := rpc.NewClient(x, libkb.ErrorUnwrapper{})
	iuicli := keybase1.IdentifyUiClient{Cli: cli}
	sessionID, err := iuicli.DelegateIdentifyUI(ctx)
	if err != nil {
		return 0, nil, err
	}
	ret := &RemoteIdentifyUI{
		sessionID: sessionID,
		uicli:     iuicli,
		logUI: &LogUI{
			sessionID,
			&keybase1.LogUiClient{Cli: cli},
		},
		Contextified: libkb.NewContextified(u.G()),
	}
	return sessionID, ret, nil
}

func (u *UIRouter) GetSecretUI(sessionID int) (ui libkb.SecretUI, err error) {
	defer u.G().Trace("UIRouter#GetSecretUI", func() error { return err })()
	x, _ := u.getUI(libkb.SecretUIKind)
	if x == nil {
		u.G().Log.Debug("| getUI(libkb.SecretUIKind) returned nil")
		return nil, nil
	}
	cli := rpc.NewClient(x, libkb.ErrorUnwrapper{}, nil)
	scli := keybase1.SecretUiClient{Cli: cli}

	u.G().Log.Debug("| returning delegated SecretUI with sessionID = %d", sessionID)
	ret := &SecretUI{
		cli:          &scli,
		sessionID:    sessionID,
		Contextified: libkb.NewContextified(u.G()),
	}
	return ret, nil
}

func (u *UIRouter) GetRekeyUI() (keybase1.RekeyUIInterface, int, error) {
	var err error
	defer u.G().Trace("UIRouter#GetRekeyUI", func() error { return err })()

	x, cid := u.getUI(libkb.RekeyUIKind)
	if x == nil {
		u.G().Log.Debug("| getUI(libkb.RekeyUIKind) returned nil")
		return nil, 0, nil
	}
	cli := rpc.NewClient(x, libkb.ErrorUnwrapper{}, nil)
	uicli := keybase1.RekeyUIClient{Cli: cli}
	sessionID, err := uicli.DelegateRekeyUI(context.TODO())
	if err != nil {
		return nil, 0, err
	}
	ret := &RekeyUI{
		Contextified: libkb.NewContextified(u.G()),
		sessionID:    sessionID,
		cli:          &uicli,
		connectionID: cid,
	}
	return ret, sessionID, nil
}

func (u *UIRouter) getOrReuseRekeyUI(prev *RekeyUI) (ret *RekeyUI, err error) {
	defer u.G().Trace("UIRouter#GetOrReuseRekeyUI", func() error { return err })()
	x, cid := u.getUI(libkb.RekeyUIKind)

	if x == nil {
		return nil, nil
	}

	if prev != nil && prev.connectionID == cid {
		return prev, nil
	}

	cli := rpc.NewClient(x, libkb.ErrorUnwrapper{}, nil)
	uicli := keybase1.RekeyUIClient{Cli: cli}
	var sessionID int
	sessionID, err = uicli.DelegateRekeyUI(context.TODO())
	if err != nil {
		return nil, err
	}
	ret = &RekeyUI{
		Contextified: libkb.NewContextified(u.G()),
		sessionID:    sessionID,
		cli:          &uicli,
		connectionID: cid,
	}

	return ret, nil
}

func (u *UIRouter) GetRekeyUINoSessionID() (ret keybase1.RekeyUIInterface, err error) {
	defer u.G().Trace("UIRouter#GetRekeyUINoSessionID", func() error { return err })()
	return u.getOrReuseRekeyUI(nil)
}
<commit_msg>Fix after rebase<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package service

import (
	libkb "github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/go-framed-msgpack-rpc/rpc"
	context "golang.org/x/net/context"
)

type transporterAndConnectionID struct {
	transporter  rpc.Transporter
	connectionID libkb.ConnectionID
}

type getObj struct {
	ui    libkb.UIKind
	retCh chan<- transporterAndConnectionID
}

type uiWrapper struct {
	cid       libkb.ConnectionID
	sessionID int
}

type setObj struct {
	cid libkb.ConnectionID
	ui  libkb.UIKind
}

type UIRouter struct {
	libkb.Contextified
	cm         *libkb.ConnectionManager
	uis        map[libkb.UIKind]libkb.ConnectionID
	setCh      chan setObj
	getCh      chan getObj
	shutdownCh chan struct{}
}

func NewUIRouter(g *libkb.GlobalContext) *UIRouter {
	ret := &UIRouter{
		Contextified: libkb.NewContextified(g),
		cm:           g.ConnectionManager,
		uis:          make(map[libkb.UIKind]libkb.ConnectionID),
		setCh:        make(chan setObj),
		getCh:        make(chan getObj),
		shutdownCh:   make(chan struct{}),
	}
	go ret.run()
	return ret
}

func (u *UIRouter) Shutdown() {
	u.shutdownCh <- struct{}{}
}

func (u *UIRouter) run() {
	for {
		select {
		case <-u.shutdownCh:
			return
		case o := <-u.setCh:
			u.uis[o.ui] = o.cid
		case o := <-u.getCh:
			var ret rpc.Transporter
			cid, ok := u.uis[o.ui]
			if ok {
				if ret = u.cm.LookupConnection(cid); ret == nil {
					u.G().Log.Debug("UIRouter: connection %v inactive, deleting registered UI %s", cid, o.ui)
					delete(u.uis, o.ui)
				}
			}
			o.retCh <- transporterAndConnectionID{transporter: ret, connectionID: cid}
		}
	}
}

func (u *UIRouter) SetUI(c libkb.ConnectionID, k libkb.UIKind) {
	u.G().Log.Debug("UIRouter: connection %v registering UI %s", c, k)
	u.setCh <- setObj{c, k}
}

func (u *UIRouter) getUI(k libkb.UIKind) (rpc.Transporter, libkb.ConnectionID) {
	retCh := make(chan transporterAndConnectionID)
	u.getCh <- getObj{k, retCh}
	ret := <-retCh
	return ret.transporter, ret.connectionID
}

func (u *UIRouter) GetIdentifyUI() (libkb.IdentifyUI, error) {
	x, _ := u.getUI(libkb.IdentifyUIKind)
	if x == nil {
		return nil, nil
	}
	cli := rpc.NewClient(x, libkb.ErrorUnwrapper{}, nil)
	iuicli := keybase1.IdentifyUiClient{Cli: cli}
	sessionID, err := iuicli.DelegateIdentifyUI(context.TODO())
	if err != nil {
		return nil, err
	}
	ret := &RemoteIdentifyUI{
		sessionID: sessionID,
		uicli:     iuicli,
		logUI: &LogUI{
			sessionID,
			&keybase1.LogUiClient{Cli: cli},
		},
		Contextified: libkb.NewContextified(u.G()),
	}
	return ret, nil
}

func (u *UIRouter) GetIdentifyUICtx(ctx context.Context) (int, libkb.IdentifyUI, error) {
	x, _ := u.getUI(libkb.IdentifyUIKind)
	if x == nil {
		return 0, nil, nil
	}
	cli := rpc.NewClient(x, libkb.ErrorUnwrapper{}, nil)
	iuicli := keybase1.IdentifyUiClient{Cli: cli}
	sessionID, err := iuicli.DelegateIdentifyUI(ctx)
	if err != nil {
		return 0, nil, err
	}
	ret := &RemoteIdentifyUI{
		sessionID: sessionID,
		uicli:     iuicli,
		logUI: &LogUI{
			sessionID,
			&keybase1.LogUiClient{Cli: cli},
		},
		Contextified: libkb.NewContextified(u.G()),
	}
	return sessionID, ret, nil
}

func (u *UIRouter) GetSecretUI(sessionID int) (ui libkb.SecretUI, err error) {
	defer u.G().Trace("UIRouter#GetSecretUI", func() error { return err })()
	x, _ := u.getUI(libkb.SecretUIKind)
	if x == nil {
		u.G().Log.Debug("| getUI(libkb.SecretUIKind) returned nil")
		return nil, nil
	}
	cli := rpc.NewClient(x, libkb.ErrorUnwrapper{}, nil)
	scli := keybase1.SecretUiClient{Cli: cli}

	u.G().Log.Debug("| returning delegated SecretUI with sessionID = %d", sessionID)
	ret := &SecretUI{
		cli:          &scli,
		sessionID:    sessionID,
		Contextified: libkb.NewContextified(u.G()),
	}
	return ret, nil
}

func (u *UIRouter) GetRekeyUI() (keybase1.RekeyUIInterface, int, error) {
	var err error
	defer u.G().Trace("UIRouter#GetRekeyUI", func() error { return err })()

	x, cid := u.getUI(libkb.RekeyUIKind)
	if x == nil {
		u.G().Log.Debug("| getUI(libkb.RekeyUIKind) returned nil")
		return nil, 0, nil
	}
	cli := rpc.NewClient(x, libkb.ErrorUnwrapper{}, nil)
	uicli := keybase1.RekeyUIClient{Cli: cli}
	sessionID, err := uicli.DelegateRekeyUI(context.TODO())
	if err != nil {
		return nil, 0, err
	}
	ret := &RekeyUI{
		Contextified: libkb.NewContextified(u.G()),
		sessionID:    sessionID,
		cli:          &uicli,
		connectionID: cid,
	}
	return ret, sessionID, nil
}

func (u *UIRouter) getOrReuseRekeyUI(prev *RekeyUI) (ret *RekeyUI, err error) {
	defer u.G().Trace("UIRouter#GetOrReuseRekeyUI", func() error { return err })()
	x, cid := u.getUI(libkb.RekeyUIKind)

	if x == nil {
		return nil, nil
	}

	if prev != nil && prev.connectionID == cid {
		return prev, nil
	}

	cli := rpc.NewClient(x, libkb.ErrorUnwrapper{}, nil)
	uicli := keybase1.RekeyUIClient{Cli: cli}
	var sessionID int
	sessionID, err = uicli.DelegateRekeyUI(context.TODO())
	if err != nil {
		return nil, err
	}
	ret = &RekeyUI{
		Contextified: libkb.NewContextified(u.G()),
		sessionID:    sessionID,
		cli:          &uicli,
		connectionID: cid,
	}

	return ret, nil
}

func (u *UIRouter) GetRekeyUINoSessionID() (ret keybase1.RekeyUIInterface, err error) {
	defer u.G().Trace("UIRouter#GetRekeyUINoSessionID", func() error { return err })()
	return u.getOrReuseRekeyUI(nil)
}
<|endoftext|>


<commit_before>package reaper_test

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

	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagertest"
	. "github.com/concourse/worker/reaper"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/ghttp"
)

var _ = Describe("Client", func() {

	Describe("It generates correct requests", func() {
		var (
			client     ReaperClient
			fakeServer *ghttp.Server
			logger     lager.Logger
		)

		BeforeEach(func() {
			fakeServer = ghttp.NewServer()
			logger = lagertest.NewTestLogger("reaper-client")
			client = NewClient(fakeServer.URL(), logger)
		})

		Describe("ping request", func() {
			It("generate correct request", func() {
				fakeServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("GET", "/ping"),
						ghttp.RespondWith(http.StatusOK, nil),
					),
				)

				err := client.Ping()
				Expect(err).ToNot(HaveOccurred())
			})

			It("returns proper error response", func() {
				fakeServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("GET", "/ping"),
						ghttp.RespondWith(http.StatusInternalServerError, nil),
					),
				)

				err := client.Ping()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(Equal("Unable to reach garden"))
			})
		})

		Describe("destroy containers request", func() {
			var buf bytes.Buffer

			BeforeEach(func() {
				handles := []string{"handle-1", "handle-2"}
				json.NewEncoder(&buf).Encode(handles)
			})

			It("generate correct requests for all container handles", func() {
				fakeServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("DELETE", "/containers/destroy"),
						ghttp.VerifyBody(buf.Bytes()),
						ghttp.RespondWithJSONEncoded(http.StatusNoContent, nil),
					),
				)
				err := client.DestroyContainers([]string{"handle-1", "handle-2"})
				Expect(err).ToNot(HaveOccurred())
			})

			It("returns proper error response", func() {
				fakeServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("DELETE", "/containers/destroy"),
						ghttp.VerifyBody(buf.Bytes()),
						ghttp.RespondWithJSONEncoded(http.StatusInternalServerError, nil),
					),
				)

				err := client.DestroyContainers([]string{"handle-1", "handle-2"})
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(Equal("received-500-response"))
			})
		})

		Describe("list containers request", func() {
			var handles []string
			It("generate correct requests for all container handles", func() {
				handles = []string{"handle-1", "handle-2"}
				fakeServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("GET", "/containers/list"),
						ghttp.RespondWithJSONEncoded(http.StatusOK, &handles),
					),
				)
				respHandles, err := client.ListHandles()
				Expect(err).ToNot(HaveOccurred())
				Expect(respHandles).To(Equal(handles))
				Expect(len(fakeServer.ReceivedRequests())).To(Equal(1))
			})

			It("returns proper error response", func() {
				fakeServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("GET", "/containers/list"),
						ghttp.RespondWithJSONEncoded(http.StatusInternalServerError, nil),
					),
				)
				_, err := client.ListHandles()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(Equal("received-500-response"))
			})
		})

		AfterEach(func() {
			fakeServer.Close()
		})
	})
})
<commit_msg>fix test flake<commit_after>package reaper_test

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

	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagertest"
	. "github.com/concourse/worker/reaper"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/ghttp"
)

var _ = Describe("Client", func() {

	Describe("It generates correct requests", func() {
		var (
			client     ReaperClient
			fakeServer *ghttp.Server
			logger     lager.Logger
		)

		BeforeEach(func() {
			fakeServer = ghttp.NewServer()
			logger = lagertest.NewTestLogger("reaper-client")
			client = NewClient(fakeServer.URL(), logger)
		})

		Describe("ping request", func() {
			It("generate correct request", func() {
				fakeServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("GET", "/ping"),
						ghttp.RespondWith(http.StatusOK, nil),
					),
				)

				err := client.Ping()
				Expect(err).ToNot(HaveOccurred())
			})

			It("returns proper error response", func() {
				fakeServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("GET", "/ping"),
						ghttp.RespondWith(http.StatusInternalServerError, nil),
					),
				)

				err := client.Ping()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(Equal("Unable to reach garden"))
			})
		})

		Describe("destroy containers request", func() {
			var (
				buf     bytes.Buffer
				handles = []string{"handle-1", "handle-2"}
			)

			BeforeEach(func() {
				json.NewEncoder(&buf).Encode(handles)
			})

			It("generate correct requests for all container handles", func() {
				fakeServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("DELETE", "/containers/destroy"),
						ghttp.VerifyJSON(buf.String()),
						ghttp.RespondWithJSONEncoded(http.StatusNoContent, nil),
					),
				)
				err := client.DestroyContainers(handles)
				Expect(err).ToNot(HaveOccurred())
			})

			It("returns proper error response", func() {
				fakeServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("DELETE", "/containers/destroy"),
						ghttp.VerifyJSON(buf.String()),
						ghttp.RespondWithJSONEncoded(http.StatusInternalServerError, nil),
					),
				)

				err := client.DestroyContainers(handles)
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(Equal("received-500-response"))
			})
		})

		Describe("list containers request", func() {
			var handles []string
			It("generate correct requests for all container handles", func() {
				handles = []string{"handle-1", "handle-2"}
				fakeServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("GET", "/containers/list"),
						ghttp.RespondWithJSONEncoded(http.StatusOK, &handles),
					),
				)
				respHandles, err := client.ListHandles()
				Expect(err).ToNot(HaveOccurred())
				Expect(respHandles).To(Equal(handles))
				Expect(len(fakeServer.ReceivedRequests())).To(Equal(1))
			})

			It("returns proper error response", func() {
				fakeServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("GET", "/containers/list"),
						ghttp.RespondWithJSONEncoded(http.StatusInternalServerError, nil),
					),
				)
				_, err := client.ListHandles()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(Equal("received-500-response"))
			})
		})

		AfterEach(func() {
			fakeServer.Close()
		})
	})
})
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/golang/glog"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/api/v1"
	clientset "k8s.io/kubernetes/pkg/client/clientset_generated/release_1_4"
	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	clientcmdapi "k8s.io/kubernetes/pkg/client/unversioned/clientcmd/api"
)

const (
	kubeletAPIPodsURL = "http://127.0.0.1:10255/pods"
	ignorePath        = "/srv/kubernetes/manifests"
	activePath        = "/etc/kubernetes/manifests"
	manifestFilename  = "apiserver.json"
	kubeconfigPath    = "/etc/kubernetes/kubeconfig"
	secretsPath       = "/etc/kubernetes/checkpoint-secrets"
)

var (
	tempAPIServer      = []byte("temp-apiserver")
	kubeAPIServer      = []byte("kube-apiserver")
	activeManifest     = filepath.Join(activePath, manifestFilename)
	checkpointManifest = filepath.Join(ignorePath, manifestFilename)
	secureAPIAddr      = fmt.Sprintf("https://%s:%s", os.Getenv("KUBERNETES_SERVICE_HOST"), os.Getenv("KUBERNETES_SERVICE_PORT_HTTPS"))
)

var tempAPIServerManifest = v1.Pod{
	TypeMeta: unversioned.TypeMeta{
		APIVersion: "v1",
		Kind:       "Pod",
	},
	ObjectMeta: v1.ObjectMeta{
		Name:      "temp-apiserver",
		Namespace: "kube-system",
	},
}

func main() {
	glog.Info("begin apiserver checkpointing...")
	run()
}

func run() {
	client := newAPIClient()
	for {
		var podList v1.PodList
		if err := json.Unmarshal(getPodsFromKubeletAPI(), &podList); err != nil {
			glog.Fatal(err)
		}
		switch {
		case bothAPIServersRunning(podList):
			glog.Info("both temp and kube apiserver running, removing temp apiserver")
			// Both the self-hosted API Server and the temp API Server are running.
			// Remove the temp API Server manifest from the config dir so that the
			// kubelet will stop it.
			if err := os.Remove(activeManifest); err != nil {
				glog.Error(err)
			}
		case kubeSystemAPIServerRunning(podList, client):
			glog.Info("kube-apiserver found, creating temp-apiserver manifest")
			// The self-hosted API Server is running. Let's snapshot the pod,
			// clean it up a bit, and then save it to the ignore path for
			// later use.
			tempAPIServerManifest.Spec = parseAPIPodSpec(podList)
			convertSecretsToVolumeMounts(client, &tempAPIServerManifest)
			writeManifest(tempAPIServerManifest)
			glog.Infof("finished creating temp-apiserver manifest at %s\n", checkpointManifest)

		default:
			glog.Info("no apiserver running, installing temp apiserver static manifest")
			b, err := ioutil.ReadFile(checkpointManifest)
			if err != nil {
				glog.Error(err)
			} else {
				if err := ioutil.WriteFile(activeManifest, b, 0644); err != nil {
					glog.Error(err)
				}
			}
		}
		time.Sleep(1 * time.Second)
	}
}

func stripNonessentialInfo(p *v1.Pod) {
	p.Spec.ServiceAccountName = ""
	p.Spec.DeprecatedServiceAccount = ""
	p.Status.Reset()
}

func getPodsFromKubeletAPI() []byte {
	var pods []byte
	res, err := http.Get(kubeletAPIPodsURL)
	if err != nil {
		glog.Error(err)
		return pods
	}
	pods, err = ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		glog.Error(err)
	}
	return pods
}

func bothAPIServersRunning(pods v1.PodList) bool {
	var kubeAPISeen, tempAPISeen bool
	for _, p := range pods.Items {
		kubeAPISeen = kubeAPISeen || isKubeAPI(p)
		tempAPISeen = tempAPISeen || isTempAPI(p)
		if kubeAPISeen && tempAPISeen {
			return true
		}
	}
	return false
}

func kubeSystemAPIServerRunning(pods v1.PodList, client clientset.Interface) bool {
	for _, p := range pods.Items {
		if isKubeAPI(p) {
			// Make sure it's actually running. Sometimes we get that
			// pod manifest back, but the server is not actually running.
			_, err := client.Discovery().ServerVersion()
			return err == nil
		}
	}
	return false
}

func isKubeAPI(pod v1.Pod) bool {
	return strings.Contains(pod.Name, "kube-apiserver") && pod.Namespace == api.NamespaceSystem
}

func isTempAPI(pod v1.Pod) bool {
	return strings.Contains(pod.Name, "temp-apiserver") && pod.Namespace == api.NamespaceSystem
}

// cleanVolumes will sanitize the list of volumes and volume mounts
// to remove the default service account token.
func cleanVolumes(p *v1.Pod) {
	volumes := make([]v1.Volume, 0, len(p.Spec.Volumes))
	for _, v := range p.Spec.Volumes {
		if !strings.HasPrefix(v.Name, "default-token") {
			volumes = append(volumes, v)
		}
	}
	p.Spec.Volumes = volumes
	for i := range p.Spec.Containers {
		c := &p.Spec.Containers[i]
		volumeMounts := make([]v1.VolumeMount, 0, len(c.VolumeMounts))
		for _, vm := range c.VolumeMounts {
			if !strings.HasPrefix(vm.Name, "default-token") {
				volumeMounts = append(volumeMounts, vm)
			}
		}
		c.VolumeMounts = volumeMounts
	}
}

// writeManifest will write the manifest to the ignore path.
// It first writes the file to a temp file, and then atomically moves it into
// the actual ignore path and correct file name.
func writeManifest(manifest v1.Pod) {
	m, err := json.Marshal(manifest)
	if err != nil {
		glog.Fatal(err)
	}
	writeAndAtomicCopy(m, checkpointManifest)
}

func parseAPIPodSpec(podList v1.PodList) v1.PodSpec {
	var apiPod v1.Pod
	for _, p := range podList.Items {
		if isKubeAPI(p) {
			apiPod = p
			break
		}
	}
	cleanVolumes(&apiPod)
	stripNonessentialInfo(&apiPod)
	return apiPod.Spec
}

func newAPIClient() clientset.Interface {
	kubeConfig, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(
		&clientcmd.ClientConfigLoadingRules{ExplicitPath: kubeconfigPath},
		&clientcmd.ConfigOverrides{ClusterInfo: clientcmdapi.Cluster{Server: secureAPIAddr}}).ClientConfig()
	if err != nil {
		glog.Fatal(err)
	}
	return clientset.NewForConfigOrDie(kubeConfig)
}

func convertSecretsToVolumeMounts(client clientset.Interface, pod *v1.Pod) {
	glog.Info("converting secrets to volume mounts")
	spec := pod.Spec
	for i := range spec.Volumes {
		v := &spec.Volumes[i]
		if v.Secret != nil {
			secretName := v.Secret.SecretName
			basePath := filepath.Join(secretsPath, pod.Name, v.Secret.SecretName)
			v.HostPath = &v1.HostPathVolumeSource{
				Path: basePath,
			}
			copySecretsToDisk(client, secretName, basePath)
			v.Secret = nil
		}
	}
}

func copySecretsToDisk(client clientset.Interface, secretName, basePath string) {
	glog.Info("copying secrets to disk")
	if err := os.MkdirAll(basePath, 0755); err != nil {
		glog.Fatal(err)
	}
	glog.Infof("created directory %s", basePath)
	s, err := client.Core().Secrets(api.NamespaceSystem).Get(secretName)
	if err != nil {
		glog.Fatal(err)
	}
	for name, value := range s.Data {
		path := filepath.Join(basePath, name)
		writeAndAtomicCopy(value, path)
	}
}

func writeAndAtomicCopy(data []byte, path string) {
	// First write a "temp" file.
	tmpfile := filepath.Join(filepath.Dir(path), "."+filepath.Base(path))
	if err := ioutil.WriteFile(tmpfile, data, 0644); err != nil {
		glog.Fatal(err)
	}
	// Finally, copy that file to the correct location.
	if err := os.Rename(tmpfile, path); err != nil {
		glog.Fatal(err)
	}
}
<commit_msg>checkpoint: make run func generic<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/golang/glog"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/api/v1"
	clientset "k8s.io/kubernetes/pkg/client/clientset_generated/release_1_4"
	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	clientcmdapi "k8s.io/kubernetes/pkg/client/unversioned/clientcmd/api"
)

const (
	kubeletAPIPodsURL = "http://127.0.0.1:10255/pods"
	ignorePath        = "/srv/kubernetes/manifests"
	activePath        = "/etc/kubernetes/manifests"
	kubeconfigPath    = "/etc/kubernetes/kubeconfig"
	secretsPath       = "/etc/kubernetes/checkpoint-secrets"

	tempAPIServer = "temp-apiserver"
	kubeAPIServer = "kube-apiserver"
)

var (
	secureAPIAddr = fmt.Sprintf("https://%s:%s", os.Getenv("KUBERNETES_SERVICE_HOST"), os.Getenv("KUBERNETES_SERVICE_PORT_HTTPS"))
)

var tempAPIServerManifest = v1.Pod{
	TypeMeta: unversioned.TypeMeta{
		APIVersion: "v1",
		Kind:       "Pod",
	},
	ObjectMeta: v1.ObjectMeta{
		Name:      "temp-apiserver",
		Namespace: api.NamespaceSystem,
	},
}

var tempPodSpecMap = map[string]v1.Pod{
	tempAPIServer: tempAPIServerManifest,
}

func main() {
	glog.Info("begin pods checkpointing...")
	run(kubeAPIServer, tempAPIServer, api.NamespaceSystem)
}

func run(actualPodName, tempPodName, namespace string) {
	client := newAPIClient()
	for {
		var podList v1.PodList
		if err := json.Unmarshal(getPodsFromKubeletAPI(), &podList); err != nil {
			glog.Fatal(err)
		}
		switch {
		case bothRunning(podList, actualPodName, tempPodName, namespace):
			glog.Infof("both temp %v and actual %v pods running, removing temp pod", actualPodName, tempPodName)
			// Both the temp and actual pods are running.
			// Remove the temp manifest from the config dir so that the
			// kubelet will stop it.
			if err := os.Remove(activeManifest(tempPodName)); err != nil {
				glog.Error(err)
			}
		case isPodRunning(podList, client, actualPodName, namespace):
			glog.Infof("actual pod %v found, creating temp pod manifest", actualPodName)
			// The actual is running. Let's snapshot the pod,
			// clean it up a bit, and then save it to the ignore path for
			// later use.
			tempSpec, ok := tempPodSpecMap[tempPodName]
			if !ok {
				glog.Fatalf("cannot find pod spec for %v", tempPodName)
			}
			tempSpec.Spec = parseAPIPodSpec(podList, actualPodName, namespace)
			convertSecretsToVolumeMounts(client, &tempSpec)
			writeManifest(tempSpec, tempPodName)
			glog.Infof("finished creating temp pod %v manifest at %s\n", tempPodName, checkpointManifest(tempPodName))

		default:
			glog.Info("no actual pod running, installing temp pod static manifest")
			b, err := ioutil.ReadFile(checkpointManifest(tempPodName))
			if err != nil {
				glog.Error(err)
			} else {
				if err := ioutil.WriteFile(activeManifest(tempPodName), b, 0644); err != nil {
					glog.Error(err)
				}
			}
		}
		time.Sleep(1 * time.Second)
	}
}

func stripNonessentialInfo(p *v1.Pod) {
	p.Spec.ServiceAccountName = ""
	p.Spec.DeprecatedServiceAccount = ""
	p.Status.Reset()
}

func getPodsFromKubeletAPI() []byte {
	var pods []byte
	res, err := http.Get(kubeletAPIPodsURL)
	if err != nil {
		glog.Error(err)
		return pods
	}
	pods, err = ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		glog.Error(err)
	}
	return pods
}

func bothRunning(pods v1.PodList, an, tn, ns string) bool {
	var actualPodSeen, tempPodSeen bool
	for _, p := range pods.Items {
		actualPodSeen = actualPodSeen || isPod(p, an, ns)
		tempPodSeen = tempPodSeen || isPod(p, tn, ns)
		if actualPodSeen && tempPodSeen {
			return true
		}
	}
	return false
}

func isPodRunning(pods v1.PodList, client clientset.Interface, n, ns string) bool {
	for _, p := range pods.Items {
		if isPod(p, n, ns) {
			if n == kubeAPIServer {
				// Make sure it's actually running. Sometimes we get that
				// pod manifest back, but the server is not actually running.
				_, err := client.Discovery().ServerVersion()
				return err == nil
			}
			return true
		}
	}
	return false
}

func isPod(pod v1.Pod, n, ns string) bool {
	return strings.Contains(pod.Name, n) && pod.Namespace == ns
}

// cleanVolumes will sanitize the list of volumes and volume mounts
// to remove the default service account token.
func cleanVolumes(p *v1.Pod) {
	volumes := make([]v1.Volume, 0, len(p.Spec.Volumes))
	for _, v := range p.Spec.Volumes {
		if !strings.HasPrefix(v.Name, "default-token") {
			volumes = append(volumes, v)
		}
	}
	p.Spec.Volumes = volumes
	for i := range p.Spec.Containers {
		c := &p.Spec.Containers[i]
		volumeMounts := make([]v1.VolumeMount, 0, len(c.VolumeMounts))
		for _, vm := range c.VolumeMounts {
			if !strings.HasPrefix(vm.Name, "default-token") {
				volumeMounts = append(volumeMounts, vm)
			}
		}
		c.VolumeMounts = volumeMounts
	}
}

// writeManifest will write the manifest to the ignore path.
// It first writes the file to a temp file, and then atomically moves it into
// the actual ignore path and correct file name.
func writeManifest(manifest v1.Pod, name string) {
	m, err := json.Marshal(manifest)
	if err != nil {
		glog.Fatal(err)
	}
	writeAndAtomicCopy(m, checkpointManifest(name))
}

func parseAPIPodSpec(podList v1.PodList, n, ns string) v1.PodSpec {
	var apiPod v1.Pod
	for _, p := range podList.Items {
		if isPod(p, n, ns) {
			apiPod = p
			break
		}
	}
	cleanVolumes(&apiPod)
	stripNonessentialInfo(&apiPod)
	return apiPod.Spec
}

func newAPIClient() clientset.Interface {
	kubeConfig, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(
		&clientcmd.ClientConfigLoadingRules{ExplicitPath: kubeconfigPath},
		&clientcmd.ConfigOverrides{ClusterInfo: clientcmdapi.Cluster{Server: secureAPIAddr}}).ClientConfig()
	if err != nil {
		glog.Fatal(err)
	}
	return clientset.NewForConfigOrDie(kubeConfig)
}

func convertSecretsToVolumeMounts(client clientset.Interface, pod *v1.Pod) {
	glog.Info("converting secrets to volume mounts")
	spec := pod.Spec
	for i := range spec.Volumes {
		v := &spec.Volumes[i]
		if v.Secret != nil {
			secretName := v.Secret.SecretName
			basePath := filepath.Join(secretsPath, pod.Name, v.Secret.SecretName)
			v.HostPath = &v1.HostPathVolumeSource{
				Path: basePath,
			}
			copySecretsToDisk(client, secretName, basePath)
			v.Secret = nil
		}
	}
}

func copySecretsToDisk(client clientset.Interface, secretName, basePath string) {
	glog.Info("copying secrets to disk")
	if err := os.MkdirAll(basePath, 0755); err != nil {
		glog.Fatal(err)
	}
	glog.Infof("created directory %s", basePath)
	s, err := client.Core().Secrets(api.NamespaceSystem).Get(secretName)
	if err != nil {
		glog.Fatal(err)
	}
	for name, value := range s.Data {
		path := filepath.Join(basePath, name)
		writeAndAtomicCopy(value, path)
	}
}

func writeAndAtomicCopy(data []byte, path string) {
	// First write a "temp" file.
	tmpfile := filepath.Join(filepath.Dir(path), "."+filepath.Base(path))
	if err := ioutil.WriteFile(tmpfile, data, 0644); err != nil {
		glog.Fatal(err)
	}
	// Finally, copy that file to the correct location.
	if err := os.Rename(tmpfile, path); err != nil {
		glog.Fatal(err)
	}
}

func activeManifest(name string) string {
	return filepath.Join(activePath, name+".json")
}

func checkpointManifest(name string) string {
	return filepath.Join(ignorePath, name+".json")
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package wrangler

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"
	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/concurrency"
	"vitess.io/vitess/go/vt/key"
	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
	querypb "vitess.io/vitess/go/vt/proto/query"
	tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vtgate/engine"
	"vitess.io/vitess/go/vt/vttablet/tabletmanager/vreplication"
)

type vdiff struct {
	wr       *Wrangler
	workflow string
	differs  map[string]*tableDiffer
	sources  map[string]*dfParams
	targets  map[string]*dfParams
}

type tableDiffer struct {
	targetTable      string
	targetExpression string
	sourceExpression string
	compareCols      []int
	orderBy          []engine.OrderbyParams
}

type dfParams struct {
	master   *topo.TabletInfo
	tablet   *topo.TabletInfo
	position mysql.Position
}

// VDiff reports differences between the sources and targets of a vreplication workflow.
func (wr *Wrangler) VDiff(ctx context.Context, targetKeyspace, workflow string, filteredReplicationWaitTime time.Duration) error {
	mi, err := wr.buildMigrater(ctx, targetKeyspace, workflow)
	if err != nil {
		wr.Logger().Errorf("buildMigrater failed: %v", err)
		return err
	}
	if err := mi.validate(ctx, false /* isWrite */); err != nil {
		mi.wr.Logger().Errorf("validate failed: %v", err)
		return err
	}
	df := &vdiff{
		wr:       wr,
		workflow: workflow,
		sources:  make(map[string]*dfParams),
		targets:  make(map[string]*dfParams),
	}
	for shard, source := range mi.sources {
		df.sources[shard] = &dfParams{
			master: source.master,
		}
	}
	var oneTarget *miTarget
	for shard, target := range mi.targets {
		df.targets[shard] = &dfParams{
			master: target.master,
		}
		oneTarget = target
	}
	var oneFilter *binlogdatapb.Filter
	for _, bls := range oneTarget.sources {
		oneFilter = bls.Filter
		break
	}
	schm, err := wr.GetSchema(ctx, oneTarget.master.Alias, nil, nil, false)
	if err != nil {
		return err
	}
	df.differs, err = buildVDiffPlan(ctx, oneFilter, schm)
	return err
}

func buildVDiffPlan(ctx context.Context, filter *binlogdatapb.Filter, schm *tabletmanagerdatapb.SchemaDefinition) (map[string]*tableDiffer, error) {
	differs := make(map[string]*tableDiffer)
	for _, table := range schm.TableDefinitions {
		rule, err := vreplication.MatchTable(table.Name, filter)
		if err != nil {
			return nil, err
		}
		if rule == nil {
			continue
		}
		query := rule.Filter
		if rule.Filter == "" || key.IsKeyRange(rule.Filter) {
			buf := sqlparser.NewTrackedBuffer(nil)
			buf.Myprintf("select * from %v", sqlparser.NewTableIdent(table.Name))
			query = buf.String()
		}
		differs[table.Name], err = buildDifferPlan(table, query)
		if err != nil {
			return nil, err
		}
	}
	return differs, nil
}

func buildDifferPlan(table *tabletmanagerdatapb.TableDefinition, query string) (*tableDiffer, error) {
	statement, err := sqlparser.Parse(query)
	if err != nil {
		return nil, err
	}
	sel, ok := statement.(*sqlparser.Select)
	if !ok {
		return nil, fmt.Errorf("unexpected: %v", sqlparser.String(statement))
	}
	td := &tableDiffer{
		targetTable: table.Name,
	}
	sourceSelect := &sqlparser.Select{}
	targetSelect := &sqlparser.Select{}
	for _, selExpr := range sel.SelectExprs {
		switch selExpr := selExpr.(type) {
		case *sqlparser.StarExpr:
			for _, fld := range table.Fields {
				aliased := &sqlparser.AliasedExpr{Expr: &sqlparser.ColName{Name: sqlparser.NewColIdent(fld.Name)}}
				sourceSelect.SelectExprs = append(sourceSelect.SelectExprs, aliased)
				targetSelect.SelectExprs = append(targetSelect.SelectExprs, aliased)
			}
		case *sqlparser.AliasedExpr:
			var targetCol *sqlparser.ColName
			if !selExpr.As.IsEmpty() {
				targetCol = &sqlparser.ColName{Name: selExpr.As}
			} else {
				if colAs, ok := selExpr.Expr.(*sqlparser.ColName); ok {
					targetCol = colAs
				} else {
					return nil, fmt.Errorf("expression needs an alias: %v", sqlparser.String(selExpr))
				}
			}
			sourceSelect.SelectExprs = append(sourceSelect.SelectExprs, selExpr)
			targetSelect.SelectExprs = append(targetSelect.SelectExprs, &sqlparser.AliasedExpr{Expr: targetCol})
		}
	}
	fields := make(map[string]querypb.Type)
	for _, field := range table.Fields {
		fields[strings.ToLower(field.Name)] = field.Type
	}

	td.compareCols = make([]int, len(sourceSelect.SelectExprs))
	for i := range td.compareCols {
		colname := sourceSelect.SelectExprs[i].(*sqlparser.AliasedExpr).Expr.(*sqlparser.ColName).Name.Lowered()
		typ, ok := fields[colname]
		if !ok {
			return nil, fmt.Errorf("column %v not found in table %v", colname, table.Name)
		}
		td.compareCols[i] = i
		if sqltypes.IsText(typ) {
			sourceSelect.SelectExprs = append(sourceSelect.SelectExprs, wrapWeightString(sourceSelect.SelectExprs[i]))
			targetSelect.SelectExprs = append(targetSelect.SelectExprs, wrapWeightString(targetSelect.SelectExprs[i]))
			td.compareCols[i] = len(sourceSelect.SelectExprs) - 1
		}
	}

	sourceSelect.From = sel.From
	targetSelect.From = sqlparser.TableExprs{
		&sqlparser.AliasedTableExpr{
			Expr: &sqlparser.TableName{
				Name: sqlparser.NewTableIdent(table.Name),
			},
		},
	}

	var orderby sqlparser.OrderBy
	for _, pk := range table.PrimaryKeyColumns {
		found := false
		for i, selExpr := range targetSelect.SelectExprs {
			colname := selExpr.(*sqlparser.AliasedExpr).Expr.(*sqlparser.ColName).Name.Lowered()
			if pk == colname {
				td.orderBy = append(td.orderBy, engine.OrderbyParams{Col: td.compareCols[i]})
				found = true
				break
			}
		}
		if !found {
			// Unreachable.
			return nil, fmt.Errorf("column %v not found in table %v", pk, table.Name)
		}
		orderby = append(orderby, &sqlparser.Order{Expr: &sqlparser.ColName{Name: sqlparser.NewColIdent(pk)}})
	}
	targetSelect.OrderBy = orderby

	sourceSelect.Where = removeKeyrange(sel.Where)
	sourceSelect.GroupBy = sel.GroupBy
	sourceSelect.OrderBy = orderby

	td.sourceExpression = sqlparser.String(sourceSelect)
	td.targetExpression = sqlparser.String(targetSelect)
	return td, nil
}

func (df *vdiff) stopTargetStreams(ctx context.Context) error {
	var mu sync.Mutex

	err := df.forAll(df.targets, func(target *dfParams) error {
		query := fmt.Sprintf("update _vt.vreplication set state='Stopped', message='for vdiff' where db_name=%s and workflow=%s", encodeString(target.master.DbName()), encodeString(df.workflow))
		_, err := df.wr.tmc.VReplicationExec(ctx, target.master.Tablet, query)
		if err != nil {
			return err
		}
		query = fmt.Sprintf("select source, pos from _vt.vreplication where db_name=%s and workflow=%s", encodeString(target.master.DbName()), encodeString(df.workflow))
		p3qr, err := df.wr.tmc.VReplicationExec(ctx, target.master.Tablet, query)
		if err != nil {
			return err
		}
		qr := sqltypes.Proto3ToResult(p3qr)

		for _, row := range qr.Rows {
			var bls binlogdatapb.BinlogSource
			if err := proto.UnmarshalText(row[0].ToString(), &bls); err != nil {
				return err
			}
			pos, err := mysql.DecodePosition(row[1].ToString())
			if err != nil {
				return err
			}
			func() {
				mu.Lock()
				defer mu.Unlock()

				source, ok := df.sources[bls.Shard]
				if !ok {
					// Unreachable.
					return
				}
				if !source.position.IsZero() && source.position.AtLeast(pos) {
					return
				}
				source.position = pos
			}()
		}
		return nil
	})
	if err != nil {
		return err
	}
	return nil
}

func (df *vdiff) forAll(participants map[string]*dfParams, f func(*dfParams) error) error {
	var wg sync.WaitGroup
	allErrors := &concurrency.AllErrorRecorder{}
	for _, participant := range participants {
		wg.Add(1)
		go func(participant *dfParams) {
			defer wg.Done()

			if err := f(participant); err != nil {
				allErrors.RecordError(err)
			}
		}(participant)
	}
	wg.Wait()
	return allErrors.AggrError(vterrors.Aggregate)
}

func removeKeyrange(where *sqlparser.Where) *sqlparser.Where {
	if where == nil {
		return nil
	}
	if isFuncKeyrange(where.Expr) {
		return nil
	}
	where.Expr = removeExprKeyrange(where.Expr)
	return where
}

func removeExprKeyrange(node sqlparser.Expr) sqlparser.Expr {
	switch node := node.(type) {
	case *sqlparser.AndExpr:
		if isFuncKeyrange(node.Left) {
			return removeExprKeyrange(node.Right)
		}
		if isFuncKeyrange(node.Right) {
			return removeExprKeyrange(node.Left)
		}
		return &sqlparser.AndExpr{
			Left:  removeExprKeyrange(node.Left),
			Right: removeExprKeyrange(node.Right),
		}
	case *sqlparser.ParenExpr:
		return &sqlparser.ParenExpr{
			Expr: removeExprKeyrange(node.Expr),
		}
	}
	return node
}

func isFuncKeyrange(expr sqlparser.Expr) bool {
	funcExpr, ok := expr.(*sqlparser.FuncExpr)
	return ok && funcExpr.Name.EqualString("in_keyrange")
}

func wrapWeightString(expr sqlparser.SelectExpr) *sqlparser.AliasedExpr {
	return &sqlparser.AliasedExpr{
		Expr: &sqlparser.FuncExpr{
			Name: sqlparser.NewColIdent("weight_string"),
			Exprs: []sqlparser.SelectExpr{
				&sqlparser.AliasedExpr{
					Expr: expr.(*sqlparser.AliasedExpr).Expr,
				},
			},
		},
	}
}
<commit_msg>vdiff: selectTablets<commit_after>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package wrangler

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"
	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/concurrency"
	"vitess.io/vitess/go/vt/discovery"
	"vitess.io/vitess/go/vt/key"
	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
	querypb "vitess.io/vitess/go/vt/proto/query"
	tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vtgate/engine"
	"vitess.io/vitess/go/vt/vttablet/tabletmanager/vreplication"
)

type vdiff struct {
	wr             *Wrangler
	workflow       string
	sourceKeyspace string
	targetKeyspace string
	sourceCell     string
	targetCell     string
	tabletTypesStr string
	differs        map[string]*tableDiffer
	sources        map[string]*dfParams
	targets        map[string]*dfParams
}

type tableDiffer struct {
	targetTable      string
	targetExpression string
	sourceExpression string
	compareCols      []int
	orderBy          []engine.OrderbyParams
}

type dfParams struct {
	master   *topo.TabletInfo
	tablet   *topodatapb.Tablet
	position mysql.Position
}

// VDiff reports differences between the sources and targets of a vreplication workflow.
func (wr *Wrangler) VDiff(ctx context.Context, targetKeyspace, workflow, sourceCell, targetCell, tabletTypesStr string, filteredReplicationWaitTime time.Duration) error {
	mi, err := wr.buildMigrater(ctx, targetKeyspace, workflow)
	if err != nil {
		wr.Logger().Errorf("buildMigrater failed: %v", err)
		return err
	}
	if err := mi.validate(ctx, false /* isWrite */); err != nil {
		mi.wr.Logger().Errorf("validate failed: %v", err)
		return err
	}
	df := &vdiff{
		wr:             wr,
		workflow:       workflow,
		sourceKeyspace: mi.sourceKeyspace,
		targetKeyspace: mi.targetKeyspace,
		sourceCell:     sourceCell,
		targetCell:     targetCell,
		tabletTypesStr: tabletTypesStr,
		sources:        make(map[string]*dfParams),
		targets:        make(map[string]*dfParams),
	}
	for shard, source := range mi.sources {
		df.sources[shard] = &dfParams{
			master: source.master,
		}
	}
	var oneTarget *miTarget
	for shard, target := range mi.targets {
		df.targets[shard] = &dfParams{
			master: target.master,
		}
		oneTarget = target
	}
	var oneFilter *binlogdatapb.Filter
	for _, bls := range oneTarget.sources {
		oneFilter = bls.Filter
		break
	}
	schm, err := wr.GetSchema(ctx, oneTarget.master.Alias, nil, nil, false)
	if err != nil {
		return err
	}
	df.differs, err = buildVDiffPlan(ctx, oneFilter, schm)
	return err
}

func buildVDiffPlan(ctx context.Context, filter *binlogdatapb.Filter, schm *tabletmanagerdatapb.SchemaDefinition) (map[string]*tableDiffer, error) {
	differs := make(map[string]*tableDiffer)
	for _, table := range schm.TableDefinitions {
		rule, err := vreplication.MatchTable(table.Name, filter)
		if err != nil {
			return nil, err
		}
		if rule == nil {
			continue
		}
		query := rule.Filter
		if rule.Filter == "" || key.IsKeyRange(rule.Filter) {
			buf := sqlparser.NewTrackedBuffer(nil)
			buf.Myprintf("select * from %v", sqlparser.NewTableIdent(table.Name))
			query = buf.String()
		}
		differs[table.Name], err = buildDifferPlan(table, query)
		if err != nil {
			return nil, err
		}
	}
	return differs, nil
}

func buildDifferPlan(table *tabletmanagerdatapb.TableDefinition, query string) (*tableDiffer, error) {
	statement, err := sqlparser.Parse(query)
	if err != nil {
		return nil, err
	}
	sel, ok := statement.(*sqlparser.Select)
	if !ok {
		return nil, fmt.Errorf("unexpected: %v", sqlparser.String(statement))
	}
	td := &tableDiffer{
		targetTable: table.Name,
	}
	sourceSelect := &sqlparser.Select{}
	targetSelect := &sqlparser.Select{}
	for _, selExpr := range sel.SelectExprs {
		switch selExpr := selExpr.(type) {
		case *sqlparser.StarExpr:
			for _, fld := range table.Fields {
				aliased := &sqlparser.AliasedExpr{Expr: &sqlparser.ColName{Name: sqlparser.NewColIdent(fld.Name)}}
				sourceSelect.SelectExprs = append(sourceSelect.SelectExprs, aliased)
				targetSelect.SelectExprs = append(targetSelect.SelectExprs, aliased)
			}
		case *sqlparser.AliasedExpr:
			var targetCol *sqlparser.ColName
			if !selExpr.As.IsEmpty() {
				targetCol = &sqlparser.ColName{Name: selExpr.As}
			} else {
				if colAs, ok := selExpr.Expr.(*sqlparser.ColName); ok {
					targetCol = colAs
				} else {
					return nil, fmt.Errorf("expression needs an alias: %v", sqlparser.String(selExpr))
				}
			}
			sourceSelect.SelectExprs = append(sourceSelect.SelectExprs, selExpr)
			targetSelect.SelectExprs = append(targetSelect.SelectExprs, &sqlparser.AliasedExpr{Expr: targetCol})
		}
	}
	fields := make(map[string]querypb.Type)
	for _, field := range table.Fields {
		fields[strings.ToLower(field.Name)] = field.Type
	}

	td.compareCols = make([]int, len(sourceSelect.SelectExprs))
	for i := range td.compareCols {
		colname := sourceSelect.SelectExprs[i].(*sqlparser.AliasedExpr).Expr.(*sqlparser.ColName).Name.Lowered()
		typ, ok := fields[colname]
		if !ok {
			return nil, fmt.Errorf("column %v not found in table %v", colname, table.Name)
		}
		td.compareCols[i] = i
		if sqltypes.IsText(typ) {
			sourceSelect.SelectExprs = append(sourceSelect.SelectExprs, wrapWeightString(sourceSelect.SelectExprs[i]))
			targetSelect.SelectExprs = append(targetSelect.SelectExprs, wrapWeightString(targetSelect.SelectExprs[i]))
			td.compareCols[i] = len(sourceSelect.SelectExprs) - 1
		}
	}

	sourceSelect.From = sel.From
	targetSelect.From = sqlparser.TableExprs{
		&sqlparser.AliasedTableExpr{
			Expr: &sqlparser.TableName{
				Name: sqlparser.NewTableIdent(table.Name),
			},
		},
	}

	var orderby sqlparser.OrderBy
	for _, pk := range table.PrimaryKeyColumns {
		found := false
		for i, selExpr := range targetSelect.SelectExprs {
			colname := selExpr.(*sqlparser.AliasedExpr).Expr.(*sqlparser.ColName).Name.Lowered()
			if pk == colname {
				td.orderBy = append(td.orderBy, engine.OrderbyParams{Col: td.compareCols[i]})
				found = true
				break
			}
		}
		if !found {
			// Unreachable.
			return nil, fmt.Errorf("column %v not found in table %v", pk, table.Name)
		}
		orderby = append(orderby, &sqlparser.Order{Expr: &sqlparser.ColName{Name: sqlparser.NewColIdent(pk)}})
	}
	targetSelect.OrderBy = orderby

	sourceSelect.Where = removeKeyrange(sel.Where)
	sourceSelect.GroupBy = sel.GroupBy
	sourceSelect.OrderBy = orderby

	td.sourceExpression = sqlparser.String(sourceSelect)
	td.targetExpression = sqlparser.String(targetSelect)
	return td, nil
}

func (df *vdiff) stopTargetStreams(ctx context.Context) error {
	var mu sync.Mutex

	err := df.forAll(df.targets, func(shard string, target *dfParams) error {
		query := fmt.Sprintf("update _vt.vreplication set state='Stopped', message='for vdiff' where db_name=%s and workflow=%s", encodeString(target.master.DbName()), encodeString(df.workflow))
		_, err := df.wr.tmc.VReplicationExec(ctx, target.master.Tablet, query)
		if err != nil {
			return err
		}
		query = fmt.Sprintf("select source, pos from _vt.vreplication where db_name=%s and workflow=%s", encodeString(target.master.DbName()), encodeString(df.workflow))
		p3qr, err := df.wr.tmc.VReplicationExec(ctx, target.master.Tablet, query)
		if err != nil {
			return err
		}
		qr := sqltypes.Proto3ToResult(p3qr)

		for _, row := range qr.Rows {
			var bls binlogdatapb.BinlogSource
			if err := proto.UnmarshalText(row[0].ToString(), &bls); err != nil {
				return err
			}
			pos, err := mysql.DecodePosition(row[1].ToString())
			if err != nil {
				return err
			}
			func() {
				mu.Lock()
				defer mu.Unlock()

				source, ok := df.sources[bls.Shard]
				if !ok {
					// Unreachable.
					return
				}
				if !source.position.IsZero() && source.position.AtLeast(pos) {
					return
				}
				source.position = pos
			}()
		}
		return nil
	})
	if err != nil {
		return err
	}
	return nil
}

func (df *vdiff) selectTablets(ctx context.Context) error {
	var wg sync.WaitGroup
	var err1, err2 error

	// Parallelize all discovery.
	wg.Add(1)
	go func() {
		defer wg.Done()
		err1 = df.forAll(df.sources, func(shard string, source *dfParams) error {
			tp, err := discovery.NewTabletPicker(ctx, df.wr.ts, df.targetCell, df.targetKeyspace, shard, df.tabletTypesStr)
			if err != nil {
				return err
			}
			defer tp.Close()

			tablet, err := tp.PickForStreaming(ctx)
			if err != nil {
				return err
			}
			df.sources[shard].tablet = tablet
			return nil
		})
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		err2 = df.forAll(df.targets, func(shard string, target *dfParams) error {
			tp, err := discovery.NewTabletPicker(ctx, df.wr.ts, df.targetCell, df.targetKeyspace, shard, df.tabletTypesStr)
			if err != nil {
				return err
			}
			defer tp.Close()

			tablet, err := tp.PickForStreaming(ctx)
			if err != nil {
				return err
			}
			df.targets[shard].tablet = tablet
			return nil
		})
	}()

	wg.Wait()
	if err1 != nil {
		return err1
	}
	return err2
}

func (df *vdiff) forAll(participants map[string]*dfParams, f func(string, *dfParams) error) error {
	var wg sync.WaitGroup
	allErrors := &concurrency.AllErrorRecorder{}
	for shard, participant := range participants {
		wg.Add(1)
		go func(shard string, participant *dfParams) {
			defer wg.Done()

			if err := f(shard, participant); err != nil {
				allErrors.RecordError(err)
			}
		}(shard, participant)
	}
	wg.Wait()
	return allErrors.AggrError(vterrors.Aggregate)
}

func removeKeyrange(where *sqlparser.Where) *sqlparser.Where {
	if where == nil {
		return nil
	}
	if isFuncKeyrange(where.Expr) {
		return nil
	}
	where.Expr = removeExprKeyrange(where.Expr)
	return where
}

func removeExprKeyrange(node sqlparser.Expr) sqlparser.Expr {
	switch node := node.(type) {
	case *sqlparser.AndExpr:
		if isFuncKeyrange(node.Left) {
			return removeExprKeyrange(node.Right)
		}
		if isFuncKeyrange(node.Right) {
			return removeExprKeyrange(node.Left)
		}
		return &sqlparser.AndExpr{
			Left:  removeExprKeyrange(node.Left),
			Right: removeExprKeyrange(node.Right),
		}
	case *sqlparser.ParenExpr:
		return &sqlparser.ParenExpr{
			Expr: removeExprKeyrange(node.Expr),
		}
	}
	return node
}

func isFuncKeyrange(expr sqlparser.Expr) bool {
	funcExpr, ok := expr.(*sqlparser.FuncExpr)
	return ok && funcExpr.Name.EqualString("in_keyrange")
}

func wrapWeightString(expr sqlparser.SelectExpr) *sqlparser.AliasedExpr {
	return &sqlparser.AliasedExpr{
		Expr: &sqlparser.FuncExpr{
			Name: sqlparser.NewColIdent("weight_string"),
			Exprs: []sqlparser.SelectExpr{
				&sqlparser.AliasedExpr{
					Expr: expr.(*sqlparser.AliasedExpr).Expr,
				},
			},
		},
	}
}
<|endoftext|>


<commit_before>package control

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	log "github.com/Sirupsen/logrus"

	"github.com/codegangsta/cli"
	machineUtil "github.com/docker/machine/utils"
	"github.com/rancher/os/config"
)

const (
	NAME string = "rancher"
	BITS int    = 2048
)

func tlsConfCommands() []cli.Command {
	return []cli.Command{
		{
			Name:      "generate",
			ShortName: "gen",
			Usage:     "generates new set of TLS configuration certs",
			Action:    tlsConfCreate,
			Flags: []cli.Flag{
				cli.StringSliceFlag{
					Name:  "hostname, H",
					Usage: "the hostname for which you want to generate the certificate",
					Value: &cli.StringSlice{"localhost"},
				},
				cli.BoolFlag{
					Name:  "server, s",
					Usage: "generate the server keys instead of client keys",
				},
				cli.StringFlag{
					Name:  "dir, d",
					Usage: "the directory to save/read the certs to/from",
					Value: "",
				},
			},
		},
	}
}

func writeCerts(generateServer bool, hostname []string, cfg *config.CloudConfig, certPath, keyPath, caCertPath, caKeyPath string) error {
	if !generateServer {
		return machineUtil.GenerateCert([]string{""}, certPath, keyPath, caCertPath, caKeyPath, NAME, BITS)
	}

	if cfg.Rancher.Docker.ServerKey == "" || cfg.Rancher.Docker.ServerCert == "" {
		err := machineUtil.GenerateCert(hostname, certPath, keyPath, caCertPath, caKeyPath, NAME, BITS)
		if err != nil {
			return err
		}

		cert, err := ioutil.ReadFile(certPath)
		if err != nil {
			return err
		}

		key, err := ioutil.ReadFile(keyPath)
		if err != nil {
			return err
		}

		cfg, err = cfg.Merge(map[interface{}]interface{}{
			"rancher": map[interface{}]interface{}{
				"docker": map[interface{}]interface{}{
					"server_cert": string(cert),
					"server_key":  string(key),
				},
			},
		})
		if err != nil {
			return err
		}

		return cfg.Save() // certPath, keyPath are already written to by machineUtil.GenerateCert()
	}

	if err := ioutil.WriteFile(certPath, []byte(cfg.Rancher.Docker.ServerCert), 0400); err != nil {
		return err
	}

	return ioutil.WriteFile(keyPath, []byte(cfg.Rancher.Docker.ServerKey), 0400)

}

func writeCaCerts(cfg *config.CloudConfig, caCertPath, caKeyPath string) (*config.CloudConfig, error) {
	if cfg.Rancher.Docker.CACert == "" {
		if err := machineUtil.GenerateCACertificate(caCertPath, caKeyPath, NAME, BITS); err != nil {
			return nil, err
		}

		caCert, err := ioutil.ReadFile(caCertPath)
		if err != nil {
			return nil, err
		}

		caKey, err := ioutil.ReadFile(caKeyPath)
		if err != nil {
			return nil, err
		}

		cfg, err = cfg.Merge(map[interface{}]interface{}{
			"rancher": map[interface{}]interface{}{
				"docker": map[interface{}]interface{}{
					"ca_key":  string(caKey),
					"ca_cert": string(caCert),
				},
			},
		})
		if err != nil {
			return nil, err
		}

		if err = cfg.Save(); err != nil {
			return nil, err
		}

		return cfg, nil // caCertPath, caKeyPath are already written to by machineUtil.GenerateCACertificate()
	}

	if err := ioutil.WriteFile(caCertPath, []byte(cfg.Rancher.Docker.CACert), 0400); err != nil {
		return nil, err
	}

	if err := ioutil.WriteFile(caKeyPath, []byte(cfg.Rancher.Docker.CAKey), 0400); err != nil {
		return nil, err
	}

	return cfg, nil
}

func tlsConfCreate(c *cli.Context) {
	err := generate(c)
	if err != nil {
		log.Fatal(err)
	}
}

func generate(c *cli.Context) error {
	generateServer := c.Bool("server")
	outDir := c.String("dir")
	hostnames := c.StringSlice("hostname")

	return Generate(generateServer, outDir, hostnames)
}

func Generate(generateServer bool, outDir string, hostnames []string) error {
	cfg, err := config.LoadConfig()
	if err != nil {
		return err
	}

	if outDir == "" {
		return fmt.Errorf("out directory (-d, --dir) not specified")
	}
	caCertPath := filepath.Join(outDir, "ca.pem")
	caKeyPath := filepath.Join(outDir, "ca-key.pem")
	certPath := filepath.Join(outDir, "cert.pem")
	keyPath := filepath.Join(outDir, "key.pem")

	if generateServer {
		certPath = filepath.Join(outDir, "server-cert.pem")
		keyPath = filepath.Join(outDir, "server-key.pem")
	}

	if _, err := os.Stat(outDir); os.IsNotExist(err) {
		if err := os.MkdirAll(outDir, 0700); err != nil {
			return err
		}
	}

	cfg, err = writeCaCerts(cfg, caCertPath, caKeyPath)
	if err != nil {
		return err
	}

	return writeCerts(generateServer, hostnames, cfg, certPath, keyPath, caCertPath, caKeyPath)
}
<commit_msg>Default dirs for TLS cert generation<commit_after>package control

import (
	"io/ioutil"
	"os"
	"path/filepath"

	log "github.com/Sirupsen/logrus"

	"github.com/codegangsta/cli"
	machineUtil "github.com/docker/machine/utils"
	"github.com/rancher/os/config"
)

const (
	NAME string = "rancher"
	BITS int    = 2048
)

func tlsConfCommands() []cli.Command {
	return []cli.Command{
		{
			Name:      "generate",
			ShortName: "gen",
			Usage:     "generates new set of TLS configuration certs",
			Action:    tlsConfCreate,
			Flags: []cli.Flag{
				cli.StringSliceFlag{
					Name:  "hostname, H",
					Usage: "the hostname for which you want to generate the certificate",
					Value: &cli.StringSlice{"localhost"},
				},
				cli.BoolFlag{
					Name:  "server, s",
					Usage: "generate the server keys instead of client keys",
				},
				cli.StringFlag{
					Name:  "dir, d",
					Usage: "the directory to save/read the certs to/from",
					Value: "",
				},
			},
		},
	}
}

func writeCerts(generateServer bool, hostname []string, cfg *config.CloudConfig, certPath, keyPath, caCertPath, caKeyPath string) error {
	if !generateServer {
		return machineUtil.GenerateCert([]string{""}, certPath, keyPath, caCertPath, caKeyPath, NAME, BITS)
	}

	if cfg.Rancher.Docker.ServerKey == "" || cfg.Rancher.Docker.ServerCert == "" {
		err := machineUtil.GenerateCert(hostname, certPath, keyPath, caCertPath, caKeyPath, NAME, BITS)
		if err != nil {
			return err
		}

		cert, err := ioutil.ReadFile(certPath)
		if err != nil {
			return err
		}

		key, err := ioutil.ReadFile(keyPath)
		if err != nil {
			return err
		}

		cfg, err = cfg.Merge(map[interface{}]interface{}{
			"rancher": map[interface{}]interface{}{
				"docker": map[interface{}]interface{}{
					"server_cert": string(cert),
					"server_key":  string(key),
				},
			},
		})
		if err != nil {
			return err
		}

		return cfg.Save() // certPath, keyPath are already written to by machineUtil.GenerateCert()
	}

	if err := ioutil.WriteFile(certPath, []byte(cfg.Rancher.Docker.ServerCert), 0400); err != nil {
		return err
	}

	return ioutil.WriteFile(keyPath, []byte(cfg.Rancher.Docker.ServerKey), 0400)

}

func writeCaCerts(cfg *config.CloudConfig, caCertPath, caKeyPath string) (*config.CloudConfig, error) {
	if cfg.Rancher.Docker.CACert == "" {
		if err := machineUtil.GenerateCACertificate(caCertPath, caKeyPath, NAME, BITS); err != nil {
			return nil, err
		}

		caCert, err := ioutil.ReadFile(caCertPath)
		if err != nil {
			return nil, err
		}

		caKey, err := ioutil.ReadFile(caKeyPath)
		if err != nil {
			return nil, err
		}

		cfg, err = cfg.Merge(map[interface{}]interface{}{
			"rancher": map[interface{}]interface{}{
				"docker": map[interface{}]interface{}{
					"ca_key":  string(caKey),
					"ca_cert": string(caCert),
				},
			},
		})
		if err != nil {
			return nil, err
		}

		if err = cfg.Save(); err != nil {
			return nil, err
		}

		return cfg, nil // caCertPath, caKeyPath are already written to by machineUtil.GenerateCACertificate()
	}

	if err := ioutil.WriteFile(caCertPath, []byte(cfg.Rancher.Docker.CACert), 0400); err != nil {
		return nil, err
	}

	if err := ioutil.WriteFile(caKeyPath, []byte(cfg.Rancher.Docker.CAKey), 0400); err != nil {
		return nil, err
	}

	return cfg, nil
}

func tlsConfCreate(c *cli.Context) {
	err := generate(c)
	if err != nil {
		log.Fatal(err)
	}
}

func generate(c *cli.Context) error {
	generateServer := c.Bool("server")
	outDir := c.String("dir")
	hostnames := c.StringSlice("hostname")

	return Generate(generateServer, outDir, hostnames)
}

func Generate(generateServer bool, outDir string, hostnames []string) error {
	cfg, err := config.LoadConfig()
	if err != nil {
		return err
	}

	if outDir == "" {
		if generateServer {
			outDir = "/etc/docker/tls"
		} else {
			outDir = "/home/rancher/.docker"
		}
		log.Infof("Out directory (-d, --dir) not specified, using default: %s", outDir)
	}
	caCertPath := filepath.Join(outDir, "ca.pem")
	caKeyPath := filepath.Join(outDir, "ca-key.pem")
	certPath := filepath.Join(outDir, "cert.pem")
	keyPath := filepath.Join(outDir, "key.pem")

	if generateServer {
		certPath = filepath.Join(outDir, "server-cert.pem")
		keyPath = filepath.Join(outDir, "server-key.pem")
	}

	if _, err := os.Stat(outDir); os.IsNotExist(err) {
		if err := os.MkdirAll(outDir, 0700); err != nil {
			return err
		}
	}

	cfg, err = writeCaCerts(cfg, caCertPath, caKeyPath)
	if err != nil {
		return err
	}
	if err := writeCerts(generateServer, hostnames, cfg, certPath, keyPath, caCertPath, caKeyPath); err != nil {
		return err
	}

	if !generateServer {
		if err := filepath.Walk(outDir, func(path string, info os.FileInfo, err error) error {
			return os.Chown(path, 1100, 1100) // rancher:rancher
		}); err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>