<commit_before>package server

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"html/template"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus/promhttp"

	"golang.org/x/oauth2"

	"github.com/gorilla/csrf"
	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/gorilla/sessions"
	"github.com/nsheridan/cashier/lib"
	"github.com/nsheridan/cashier/server/auth"
	"github.com/nsheridan/cashier/server/config"
	"github.com/nsheridan/cashier/server/static"
	"github.com/nsheridan/cashier/server/templates"
)

// appContext contains local context - cookiestore, authsession etc.
type appContext struct {
	cookiestore *sessions.CookieStore
	authsession *auth.Session
}

// getAuthTokenCookie retrieves a cookie from the request.
func (a *appContext) getAuthTokenCookie(r *http.Request) *oauth2.Token {
	session, _ := a.cookiestore.Get(r, "session")
	t, ok := session.Values["token"]
	if !ok {
		return nil
	}
	var tok oauth2.Token
	if err := json.Unmarshal(t.([]byte), &tok); err != nil {
		return nil
	}
	if !tok.Valid() {
		return nil
	}
	return &tok
}

// setAuthTokenCookie marshals the auth token and stores it as a cookie.
func (a *appContext) setAuthTokenCookie(w http.ResponseWriter, r *http.Request, t *oauth2.Token) {
	session, _ := a.cookiestore.Get(r, "session")
	val, _ := json.Marshal(t)
	session.Values["token"] = val
	session.Save(r, w)
}

// getAuthStateCookie retrieves the oauth csrf state value from the client request.
func (a *appContext) getAuthStateCookie(r *http.Request) string {
	session, _ := a.cookiestore.Get(r, "session")
	state, ok := session.Values["state"]
	if !ok {
		return ""
	}
	return state.(string)
}

// setAuthStateCookie saves the oauth csrf state value.
func (a *appContext) setAuthStateCookie(w http.ResponseWriter, r *http.Request, state string) {
	session, _ := a.cookiestore.Get(r, "session")
	session.Values["state"] = state
	session.Save(r, w)
}

func (a *appContext) getCurrentURL(r *http.Request) string {
	session, _ := a.cookiestore.Get(r, "session")
	path, ok := session.Values["auth_url"]
	if !ok {
		return ""
	}
	return path.(string)
}

func (a *appContext) setCurrentURL(w http.ResponseWriter, r *http.Request) {
	session, _ := a.cookiestore.Get(r, "session")
	session.Values["auth_url"] = r.URL.Path
	session.Save(r, w)
}

func (a *appContext) isLoggedIn(w http.ResponseWriter, r *http.Request) bool {
	tok := a.getAuthTokenCookie(r)
	if !tok.Valid() || !authprovider.Valid(tok) {
		return false
	}
	return true
}

func (a *appContext) login(w http.ResponseWriter, r *http.Request) (int, error) {
	a.setCurrentURL(w, r)
	http.Redirect(w, r, "/auth/login", http.StatusSeeOther)
	return http.StatusSeeOther, nil
}

// parseKey retrieves and unmarshals the signing request.
func extractKey(r *http.Request) (*lib.SignRequest, error) {
	var s lib.SignRequest
	if err := json.NewDecoder(r.Body).Decode(&s); err != nil {
		return nil, err
	}
	return &s, nil
}

// signHandler handles the "/sign" path.
// It unmarshals the client token to an oauth token, validates it and signs the provided public ssh key.
func signHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	var t string
	if ah := r.Header.Get("Authorization"); ah != "" {
		if len(ah) > 6 && strings.ToUpper(ah[0:7]) == "BEARER " {
			t = ah[7:]
		}
	}
	if t == "" {
		return http.StatusUnauthorized, errors.New(http.StatusText(http.StatusUnauthorized))
	}
	token := &oauth2.Token{
		AccessToken: t,
	}
	if !authprovider.Valid(token) {
		return http.StatusUnauthorized, errors.New(http.StatusText(http.StatusUnauthorized))
	}

	// Sign the pubkey and issue the cert.
	req, err := extractKey(r)
	if err != nil {
		return http.StatusBadRequest, errors.Wrap(err, "unable to extract key from request")
	}
	username := authprovider.Username(token)
	authprovider.Revoke(token) // We don't need this anymore.
	cert, err := keysigner.SignUserKey(req, username)
	if err != nil {
		return http.StatusInternalServerError, errors.Wrap(err, "error signing key")
	}
	if err := certstore.SetCert(cert); err != nil {
		log.Printf("Error recording cert: %v", err)
	}
	if err := json.NewEncoder(w).Encode(&lib.SignResponse{
		Status:   "ok",
		Response: string(lib.GetPublicKey(cert)),
	}); err != nil {
		return http.StatusInternalServerError, errors.Wrap(err, "error encoding response")
	}
	return http.StatusOK, nil
}

// loginHandler starts the authentication process with the provider.
func loginHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	state := newState()
	a.setAuthStateCookie(w, r, state)
	a.authsession = authprovider.StartSession(state)
	http.Redirect(w, r, a.authsession.AuthURL, http.StatusFound)
	return http.StatusFound, nil
}

// callbackHandler handles retrieving the access token from the auth provider and saves it for later use.
func callbackHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	if r.FormValue("state") != a.getAuthStateCookie(r) {
		return http.StatusUnauthorized, errors.New(http.StatusText(http.StatusUnauthorized))
	}
	code := r.FormValue("code")
	if err := a.authsession.Authorize(authprovider, code); err != nil {
		return http.StatusInternalServerError, err
	}
	a.setAuthTokenCookie(w, r, a.authsession.Token)
	http.Redirect(w, r, a.getCurrentURL(r), http.StatusFound)
	return http.StatusFound, nil
}

// rootHandler starts the auth process. If the client is authenticated it renders the token to the user.
func rootHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	if !a.isLoggedIn(w, r) {
		return a.login(w, r)
	}
	tok := a.getAuthTokenCookie(r)
	page := struct {
		Token string
	}{tok.AccessToken}

	tmpl := template.Must(template.New("token.html").Parse(templates.Token))
	tmpl.Execute(w, page)
	return http.StatusOK, nil
}

func listRevokedCertsHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	revoked, err := certstore.GetRevoked()
	if err != nil {
		return http.StatusInternalServerError, err
	}
	rl, err := keysigner.GenerateRevocationList(revoked)
	if err != nil {
		return http.StatusInternalServerError, errors.Wrap(err, "unable to generate KRL")
	}
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Write(rl)
	return http.StatusOK, nil
}

func listAllCertsHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	if !a.isLoggedIn(w, r) {
		return a.login(w, r)
	}
	tmpl := template.Must(template.New("certs.html").Parse(templates.Certs))
	tmpl.Execute(w, map[string]interface{}{
		csrf.TemplateTag: csrf.TemplateField(r),
	})
	return http.StatusOK, nil
}

func listCertsJSONHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	if !a.isLoggedIn(w, r) {
		return http.StatusUnauthorized, errors.New(http.StatusText(http.StatusUnauthorized))
	}
	includeExpired, _ := strconv.ParseBool(r.URL.Query().Get("all"))
	certs, err := certstore.List(includeExpired)
	j, err := json.Marshal(certs)
	if err != nil {
		return http.StatusInternalServerError, errors.New(http.StatusText(http.StatusInternalServerError))
	}
	w.Write(j)
	return http.StatusOK, nil
}

func revokeCertHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	if !a.isLoggedIn(w, r) {
		return a.login(w, r)
	}
	r.ParseForm()
	for _, id := range r.Form["cert_id"] {
		if err := certstore.Revoke(id); err != nil {
			return http.StatusInternalServerError, errors.Wrap(err, "unable to revoke")
		}
	}
	http.Redirect(w, r, "/admin/certs", http.StatusSeeOther)
	return http.StatusSeeOther, nil
}

// appHandler is a handler which uses appContext to manage state.
type appHandler struct {
	*appContext
	h func(*appContext, http.ResponseWriter, *http.Request) (int, error)
}

// ServeHTTP handles the request and writes responses.
func (ah appHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	status, err := ah.h(ah.appContext, w, r)
	if err != nil {
		log.Printf("HTTP %d: %q", status, err)
		http.Error(w, err.Error(), status)
	}
}

// newState generates a state identifier for the oauth process.
func newState() string {
	k := make([]byte, 32)
	if _, err := io.ReadFull(rand.Reader, k); err != nil {
		return "unexpectedstring"
	}
	return hex.EncodeToString(k)
}

func runHTTPServer(conf *config.Server, l net.Listener) {
	var err error
	ctx := &appContext{
		cookiestore: sessions.NewCookieStore([]byte(conf.CookieSecret)),
	}
	ctx.cookiestore.Options = &sessions.Options{
		MaxAge:   900,
		Path:     "/",
		Secure:   conf.UseTLS,
		HttpOnly: true,
	}

	logfile := os.Stderr
	if conf.HTTPLogFile != "" {
		logfile, err = os.OpenFile(conf.HTTPLogFile, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0640)
		if err != nil {
			log.Printf("unable to open %s for writing. logging to stdout", conf.HTTPLogFile)
			logfile = os.Stderr
		}
	}

	CSRF := csrf.Protect([]byte(conf.CSRFSecret), csrf.Secure(conf.UseTLS))
	r := mux.NewRouter()
	r.Methods("GET").Path("/").Handler(appHandler{ctx, rootHandler})
	r.Methods("GET").Path("/auth/login").Handler(appHandler{ctx, loginHandler})
	r.Methods("GET").Path("/auth/callback").Handler(appHandler{ctx, callbackHandler})
	r.Methods("POST").Path("/sign").Handler(appHandler{ctx, signHandler})
	r.Methods("GET").Path("/revoked").Handler(appHandler{ctx, listRevokedCertsHandler})
	r.Methods("POST").Path("/admin/revoke").Handler(CSRF(appHandler{ctx, revokeCertHandler}))
	r.Methods("GET").Path("/admin/certs").Handler(CSRF(appHandler{ctx, listAllCertsHandler}))
	r.Methods("GET").Path("/admin/certs.json").Handler(appHandler{ctx, listCertsJSONHandler})
	r.Methods("GET").Path("/metrics").Handler(promhttp.Handler())
	r.PathPrefix("/").Handler(http.FileServer(static.FS(false)))
	h := handlers.LoggingHandler(logfile, r)
	s := &http.Server{
		Handler: h,
	}
	s.Serve(l)
}
<commit_msg>add a healthcheck handler<commit_after>package server

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus/promhttp"

	"golang.org/x/oauth2"

	"github.com/gorilla/csrf"
	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/gorilla/sessions"
	"github.com/nsheridan/cashier/lib"
	"github.com/nsheridan/cashier/server/auth"
	"github.com/nsheridan/cashier/server/config"
	"github.com/nsheridan/cashier/server/static"
	"github.com/nsheridan/cashier/server/templates"
)

// appContext contains local context - cookiestore, authsession etc.
type appContext struct {
	cookiestore *sessions.CookieStore
	authsession *auth.Session
}

// getAuthTokenCookie retrieves a cookie from the request.
func (a *appContext) getAuthTokenCookie(r *http.Request) *oauth2.Token {
	session, _ := a.cookiestore.Get(r, "session")
	t, ok := session.Values["token"]
	if !ok {
		return nil
	}
	var tok oauth2.Token
	if err := json.Unmarshal(t.([]byte), &tok); err != nil {
		return nil
	}
	if !tok.Valid() {
		return nil
	}
	return &tok
}

// setAuthTokenCookie marshals the auth token and stores it as a cookie.
func (a *appContext) setAuthTokenCookie(w http.ResponseWriter, r *http.Request, t *oauth2.Token) {
	session, _ := a.cookiestore.Get(r, "session")
	val, _ := json.Marshal(t)
	session.Values["token"] = val
	session.Save(r, w)
}

// getAuthStateCookie retrieves the oauth csrf state value from the client request.
func (a *appContext) getAuthStateCookie(r *http.Request) string {
	session, _ := a.cookiestore.Get(r, "session")
	state, ok := session.Values["state"]
	if !ok {
		return ""
	}
	return state.(string)
}

// setAuthStateCookie saves the oauth csrf state value.
func (a *appContext) setAuthStateCookie(w http.ResponseWriter, r *http.Request, state string) {
	session, _ := a.cookiestore.Get(r, "session")
	session.Values["state"] = state
	session.Save(r, w)
}

func (a *appContext) getCurrentURL(r *http.Request) string {
	session, _ := a.cookiestore.Get(r, "session")
	path, ok := session.Values["auth_url"]
	if !ok {
		return ""
	}
	return path.(string)
}

func (a *appContext) setCurrentURL(w http.ResponseWriter, r *http.Request) {
	session, _ := a.cookiestore.Get(r, "session")
	session.Values["auth_url"] = r.URL.Path
	session.Save(r, w)
}

func (a *appContext) isLoggedIn(w http.ResponseWriter, r *http.Request) bool {
	tok := a.getAuthTokenCookie(r)
	if !tok.Valid() || !authprovider.Valid(tok) {
		return false
	}
	return true
}

func (a *appContext) login(w http.ResponseWriter, r *http.Request) (int, error) {
	a.setCurrentURL(w, r)
	http.Redirect(w, r, "/auth/login", http.StatusSeeOther)
	return http.StatusSeeOther, nil
}

// parseKey retrieves and unmarshals the signing request.
func extractKey(r *http.Request) (*lib.SignRequest, error) {
	var s lib.SignRequest
	if err := json.NewDecoder(r.Body).Decode(&s); err != nil {
		return nil, err
	}
	return &s, nil
}

// signHandler handles the "/sign" path.
// It unmarshals the client token to an oauth token, validates it and signs the provided public ssh key.
func signHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	var t string
	if ah := r.Header.Get("Authorization"); ah != "" {
		if len(ah) > 6 && strings.ToUpper(ah[0:7]) == "BEARER " {
			t = ah[7:]
		}
	}
	if t == "" {
		return http.StatusUnauthorized, errors.New(http.StatusText(http.StatusUnauthorized))
	}
	token := &oauth2.Token{
		AccessToken: t,
	}
	if !authprovider.Valid(token) {
		return http.StatusUnauthorized, errors.New(http.StatusText(http.StatusUnauthorized))
	}

	// Sign the pubkey and issue the cert.
	req, err := extractKey(r)
	if err != nil {
		return http.StatusBadRequest, errors.Wrap(err, "unable to extract key from request")
	}
	username := authprovider.Username(token)
	authprovider.Revoke(token) // We don't need this anymore.
	cert, err := keysigner.SignUserKey(req, username)
	if err != nil {
		return http.StatusInternalServerError, errors.Wrap(err, "error signing key")
	}
	if err := certstore.SetCert(cert); err != nil {
		log.Printf("Error recording cert: %v", err)
	}
	if err := json.NewEncoder(w).Encode(&lib.SignResponse{
		Status:   "ok",
		Response: string(lib.GetPublicKey(cert)),
	}); err != nil {
		return http.StatusInternalServerError, errors.Wrap(err, "error encoding response")
	}
	return http.StatusOK, nil
}

// loginHandler starts the authentication process with the provider.
func loginHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	state := newState()
	a.setAuthStateCookie(w, r, state)
	a.authsession = authprovider.StartSession(state)
	http.Redirect(w, r, a.authsession.AuthURL, http.StatusFound)
	return http.StatusFound, nil
}

// callbackHandler handles retrieving the access token from the auth provider and saves it for later use.
func callbackHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	if r.FormValue("state") != a.getAuthStateCookie(r) {
		return http.StatusUnauthorized, errors.New(http.StatusText(http.StatusUnauthorized))
	}
	code := r.FormValue("code")
	if err := a.authsession.Authorize(authprovider, code); err != nil {
		return http.StatusInternalServerError, err
	}
	a.setAuthTokenCookie(w, r, a.authsession.Token)
	http.Redirect(w, r, a.getCurrentURL(r), http.StatusFound)
	return http.StatusFound, nil
}

// rootHandler starts the auth process. If the client is authenticated it renders the token to the user.
func rootHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	if !a.isLoggedIn(w, r) {
		return a.login(w, r)
	}
	tok := a.getAuthTokenCookie(r)
	page := struct {
		Token string
	}{tok.AccessToken}

	tmpl := template.Must(template.New("token.html").Parse(templates.Token))
	tmpl.Execute(w, page)
	return http.StatusOK, nil
}

func listRevokedCertsHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	revoked, err := certstore.GetRevoked()
	if err != nil {
		return http.StatusInternalServerError, err
	}
	rl, err := keysigner.GenerateRevocationList(revoked)
	if err != nil {
		return http.StatusInternalServerError, errors.Wrap(err, "unable to generate KRL")
	}
	w.Header().Set("Content-Type", "application/octet-stream")
	w.Write(rl)
	return http.StatusOK, nil
}

func listAllCertsHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	if !a.isLoggedIn(w, r) {
		return a.login(w, r)
	}
	tmpl := template.Must(template.New("certs.html").Parse(templates.Certs))
	tmpl.Execute(w, map[string]interface{}{
		csrf.TemplateTag: csrf.TemplateField(r),
	})
	return http.StatusOK, nil
}

func listCertsJSONHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	if !a.isLoggedIn(w, r) {
		return http.StatusUnauthorized, errors.New(http.StatusText(http.StatusUnauthorized))
	}
	includeExpired, _ := strconv.ParseBool(r.URL.Query().Get("all"))
	certs, err := certstore.List(includeExpired)
	j, err := json.Marshal(certs)
	if err != nil {
		return http.StatusInternalServerError, errors.New(http.StatusText(http.StatusInternalServerError))
	}
	w.Write(j)
	return http.StatusOK, nil
}

func revokeCertHandler(a *appContext, w http.ResponseWriter, r *http.Request) (int, error) {
	if !a.isLoggedIn(w, r) {
		return a.login(w, r)
	}
	r.ParseForm()
	for _, id := range r.Form["cert_id"] {
		if err := certstore.Revoke(id); err != nil {
			return http.StatusInternalServerError, errors.Wrap(err, "unable to revoke")
		}
	}
	http.Redirect(w, r, "/admin/certs", http.StatusSeeOther)
	return http.StatusSeeOther, nil
}

func healthcheck(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusOK)
	fmt.Fprint(w, "ok")
}

// appHandler is a handler which uses appContext to manage state.
type appHandler struct {
	*appContext
	h func(*appContext, http.ResponseWriter, *http.Request) (int, error)
}

// ServeHTTP handles the request and writes responses.
func (ah appHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	status, err := ah.h(ah.appContext, w, r)
	if err != nil {
		log.Printf("HTTP %d: %q", status, err)
		http.Error(w, err.Error(), status)
	}
}

// newState generates a state identifier for the oauth process.
func newState() string {
	k := make([]byte, 32)
	if _, err := io.ReadFull(rand.Reader, k); err != nil {
		return "unexpectedstring"
	}
	return hex.EncodeToString(k)
}

func runHTTPServer(conf *config.Server, l net.Listener) {
	var err error
	ctx := &appContext{
		cookiestore: sessions.NewCookieStore([]byte(conf.CookieSecret)),
	}
	ctx.cookiestore.Options = &sessions.Options{
		MaxAge:   900,
		Path:     "/",
		Secure:   conf.UseTLS,
		HttpOnly: true,
	}

	logfile := os.Stderr
	if conf.HTTPLogFile != "" {
		logfile, err = os.OpenFile(conf.HTTPLogFile, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0640)
		if err != nil {
			log.Printf("unable to open %s for writing. logging to stdout", conf.HTTPLogFile)
			logfile = os.Stderr
		}
	}

	CSRF := csrf.Protect([]byte(conf.CSRFSecret), csrf.Secure(conf.UseTLS))
	r := mux.NewRouter()
	r.Methods("GET").Path("/").Handler(appHandler{ctx, rootHandler})
	r.Methods("GET").Path("/auth/login").Handler(appHandler{ctx, loginHandler})
	r.Methods("GET").Path("/auth/callback").Handler(appHandler{ctx, callbackHandler})
	r.Methods("POST").Path("/sign").Handler(appHandler{ctx, signHandler})
	r.Methods("GET").Path("/revoked").Handler(appHandler{ctx, listRevokedCertsHandler})
	r.Methods("POST").Path("/admin/revoke").Handler(CSRF(appHandler{ctx, revokeCertHandler}))
	r.Methods("GET").Path("/admin/certs").Handler(CSRF(appHandler{ctx, listAllCertsHandler}))
	r.Methods("GET").Path("/admin/certs.json").Handler(appHandler{ctx, listCertsJSONHandler})
	r.Methods("GET").Path("/metrics").Handler(promhttp.Handler())
	r.Methods("GET").Path("/healthcheck").HandlerFunc(healthcheck)
	r.PathPrefix("/").Handler(http.FileServer(static.FS(false)))
	h := handlers.LoggingHandler(logfile, r)
	s := &http.Server{
		Handler: h,
	}
	s.Serve(l)
}
<|endoftext|>


<commit_before>package pack
import (
	"text/template"
	"bytes"
	"errors"
	"runtime"
	"os/user"
	"time"
	"strings"
	"strconv"
)
const ServiceConfigFile = "service.conf"
type Service struct {
	Env          map[string]string `json:"env,omitempty"`
	RunOpts      string`json:"opts,omitempty"`
	Restart      bool`json:"restart"`
	AutoStart    bool`json:"autostart"`
	RestartDelay int`json:"restartDelay"`
	TargetInit   string`json:"target,omitempty"`

}

type Descriptor struct {
	Group              string`json:"group"`
	Name               string`json:"name"`
	Version            string`json:"version"`
	Author             string`json:"author"`
	Description        string`json:"description"`
	Depends            []string `json:"depends,omitempty"`
	Architectures      []string `json:"arch,omitempty"`
	BinName            string`json:"bin,omitempty"`
	Service            *Service`json:"service,omitempty"`
	Resources          string`json:"resources,omitempty"`
	PreInst            string `json:"preinst,omitempty"`
	PostInst           string `json:"postinst,omitempty"`
	PreRm              string `json:"prerm,omitempty"`
	TargetResourcesDir string`json:"resourcesDir,omitempty"`
	TargetBinDir       string`json:"binDir,omitempty"`
	TargetConfDir      string`json:"confDir,omitempty"`
	TargetServiceDir   string `json:"serviceDir,omitempty"`
}

func (s *Service) FillDefault() {
	if s.TargetInit == "" {
		s.TargetInit = "upstart"
	}else if s.TargetInit != "upstart" {
		panic("Only 'upstart' target init system is allowed for services")
	}
}

func (d *Descriptor) FillDefault() error {
	if d.Resources == "" {
		d.Resources = "resources"
	}
	if d.Name == "" {
		return errors.New("Name must be specified")
	}
	if d.Group == "" {
		return errors.New("Group must be specified")
	}
	if d.Author == "" {
		usr, err := user.Current()
		if err != nil {
			return err
		}
		d.Author = usr.Name
	}
	if d.Version == "" {
		d.Version = "0.0.0"
	}
	if d.Description == "" {
		d.Description = "{{.Group}} {{.Name}} built at " + time.Now().Format(time.RFC3339)
	}
	if len(d.Architectures) == 0 {
		d.Architectures = append(d.Architectures, runtime.GOARCH)
	}
	if d.BinName == "" {
		d.BinName = "{{.Group}}-{{.Name}}"
	}
	if d.Service != nil {
		d.Service.FillDefault()
	}
	if d.TargetResourcesDir == "" {
		d.TargetResourcesDir = "/usr/local/share/{{.Group}}/{{.Name}}"
	}
	if d.TargetBinDir == "" {
		d.TargetBinDir = "/usr/local/bin"
	}
	if d.TargetConfDir == "" {
		d.TargetConfDir = "/etc/{{.Group}}/{{.Name}}"
	}
	if d.TargetServiceDir == "" {
		d.TargetServiceDir = "/etc/init"
	}
	return nil
}

func (d *Descriptor) mustTemplate(field *string) {
	t, err := template.New("").Parse(*field)
	if err != nil {
		panic("Failed parse templates in " + (*field) + ": " + err.Error())
	}
	buf := new(bytes.Buffer)
	err = t.Execute(buf, (*d))
	if err != nil {
		panic("Failed execute templates in " + (*field) + ": " + err.Error())
	}
	*field = buf.String()
}

func (d *Descriptor) FillTemplates() {
	d.mustTemplate(&(d.Name))
	d.mustTemplate(&(d.Resources))
	d.mustTemplate(&(d.Group))
	d.mustTemplate(&(d.Version))
	d.mustTemplate(&(d.BinName))
	d.mustTemplate(&(d.Description))
	d.mustTemplate(&(d.Author))
	d.mustTemplate(&(d.TargetServiceDir))

	if d.Service != nil {
		d.mustTemplate(&(d.Service.TargetInit))
		for k, v := range d.Service.Env {
			tmp := v
			d.mustTemplate(&tmp)
			d.Service.Env[k] = tmp
		}
		d.Service.RunOpts = mustTemplate(d.Service.RunOpts, *d)
	}

	d.mustTemplate(&(d.TargetResourcesDir))
	d.mustTemplate(&(d.TargetBinDir))
	d.mustTemplate(&(d.TargetConfDir))
}

func normalizeArch(arch string) string {
	if _, err := strconv.ParseUint(arch, 10, 64); err != nil {
		return arch
	}
	return "i" + arch
}

func (d *Descriptor) Control() string {
	t := `Package: {{.Group}}-{{.Name}}
Version: {{.Version}}
Architecture: ` + normalizeArch(runtime.GOARCH) + `
Maintainer: {{.Author}}
`
	if len(d.Depends) != 0 {
		t += "Depends: " + strings.Join(d.Depends, ",") + "\n"
	}
	t += `Description: {{.Description}}
`
	d.mustTemplate(&t)
	return t
}

func (d *Descriptor) PreInstall() string {
	t := ""
	if d.Service != nil {
		t += mustTemplate("service {{.BinName}} stop || echo 'No installed {{.Group}} {{.Name}} instance running'", *d)
	}
	return t
}

func (d *Descriptor) PostInstall() string {
	t := ""
	if d.Service != nil {
		if d.Service.AutoStart {
			t += mustTemplate("service {{.BinName}} start || echo 'No {{.Group}} {{.Name}} instance'", *d)
		}
	}
	return t
}

func (d *Descriptor) PreRemove() string {
	return d.PreInstall()
}

func (d *Descriptor) ServiceInit() string {
	t := `# {{.Group}} {{.Name}}

description         "{{.Description}}"

start on runlevel [2345]
stop on runlevel [!2345]
{{if .Service.Restart}}
respawn
respawn limit 99999999 {{.Service.RestartDelay}}
{{end}}

script
  . {{.TargetConfDir}}/` + ServiceConfigFile + `
  exec {{.TargetBinDir}}/{{.BinName}} $RUN_OPTS | logger -t '{{.Group}}-{{.Name}}' 2>&1
end script
`
	return mustTemplate(t, *d)
}

func (d *Descriptor) ServiceFile() string {
	t := `{{.Group}}-{{.Name}}.conf`
	return mustTemplate(t, *d)
}

func (d *Descriptor) ServiceConfig() string {
	t := `#!/bin/bash
{{ range $key, $value := .Service.Env}}
{{$key}}="$value"
{{end}}
`
	return mustTemplate(t, *d) + "RUN_OPTS=\"" + d.Service.RunOpts + "\""
}

func mustTemplate(pattern string, params interface{}) string {
	t, err := template.New("").Parse(pattern)
	if err != nil {
		panic("Failed parse templates in " + (pattern) + ": " + err.Error())
	}
	buf := new(bytes.Buffer)
	err = t.Execute(buf, params)
	if err != nil {
		panic("Failed execute templates in " + (pattern) + ": " + err.Error())
	}
	return buf.String()
}
<commit_msg>fix environment variables template<commit_after>package pack
import (
	"text/template"
	"bytes"
	"errors"
	"runtime"
	"os/user"
	"time"
	"strings"
	"strconv"
)
const ServiceConfigFile = "service.conf"
type Service struct {
	Env          map[string]string `json:"env,omitempty"`
	RunOpts      string`json:"opts,omitempty"`
	Restart      bool`json:"restart"`
	AutoStart    bool`json:"autostart"`
	RestartDelay int`json:"restartDelay"`
	TargetInit   string`json:"target,omitempty"`

}

type Descriptor struct {
	Group              string`json:"group"`
	Name               string`json:"name"`
	Version            string`json:"version"`
	Author             string`json:"author"`
	Description        string`json:"description"`
	Depends            []string `json:"depends,omitempty"`
	Architectures      []string `json:"arch,omitempty"`
	BinName            string`json:"bin,omitempty"`
	Service            *Service`json:"service,omitempty"`
	Resources          string`json:"resources,omitempty"`
	PreInst            string `json:"preinst,omitempty"`
	PostInst           string `json:"postinst,omitempty"`
	PreRm              string `json:"prerm,omitempty"`
	TargetResourcesDir string`json:"resourcesDir,omitempty"`
	TargetBinDir       string`json:"binDir,omitempty"`
	TargetConfDir      string`json:"confDir,omitempty"`
	TargetServiceDir   string `json:"serviceDir,omitempty"`
}

func (s *Service) FillDefault() {
	if s.TargetInit == "" {
		s.TargetInit = "upstart"
	}else if s.TargetInit != "upstart" {
		panic("Only 'upstart' target init system is allowed for services")
	}
}

func (d *Descriptor) FillDefault() error {
	if d.Resources == "" {
		d.Resources = "resources"
	}
	if d.Name == "" {
		return errors.New("Name must be specified")
	}
	if d.Group == "" {
		return errors.New("Group must be specified")
	}
	if d.Author == "" {
		usr, err := user.Current()
		if err != nil {
			return err
		}
		d.Author = usr.Name
	}
	if d.Version == "" {
		d.Version = "0.0.0"
	}
	if d.Description == "" {
		d.Description = "{{.Group}} {{.Name}} built at " + time.Now().Format(time.RFC3339)
	}
	if len(d.Architectures) == 0 {
		d.Architectures = append(d.Architectures, runtime.GOARCH)
	}
	if d.BinName == "" {
		d.BinName = "{{.Group}}-{{.Name}}"
	}
	if d.Service != nil {
		d.Service.FillDefault()
	}
	if d.TargetResourcesDir == "" {
		d.TargetResourcesDir = "/usr/local/share/{{.Group}}/{{.Name}}"
	}
	if d.TargetBinDir == "" {
		d.TargetBinDir = "/usr/local/bin"
	}
	if d.TargetConfDir == "" {
		d.TargetConfDir = "/etc/{{.Group}}/{{.Name}}"
	}
	if d.TargetServiceDir == "" {
		d.TargetServiceDir = "/etc/init"
	}
	return nil
}

func (d *Descriptor) mustTemplate(field *string) {
	t, err := template.New("").Parse(*field)
	if err != nil {
		panic("Failed parse templates in " + (*field) + ": " + err.Error())
	}
	buf := new(bytes.Buffer)
	err = t.Execute(buf, (*d))
	if err != nil {
		panic("Failed execute templates in " + (*field) + ": " + err.Error())
	}
	*field = buf.String()
}

func (d *Descriptor) FillTemplates() {
	d.mustTemplate(&(d.Name))
	d.mustTemplate(&(d.Resources))
	d.mustTemplate(&(d.Group))
	d.mustTemplate(&(d.Version))
	d.mustTemplate(&(d.BinName))
	d.mustTemplate(&(d.Description))
	d.mustTemplate(&(d.Author))
	d.mustTemplate(&(d.TargetServiceDir))

	if d.Service != nil {
		d.mustTemplate(&(d.Service.TargetInit))
		for k, v := range d.Service.Env {
			tmp := v
			d.mustTemplate(&tmp)
			d.Service.Env[k] = tmp
		}
		d.Service.RunOpts = mustTemplate(d.Service.RunOpts, *d)
	}

	d.mustTemplate(&(d.TargetResourcesDir))
	d.mustTemplate(&(d.TargetBinDir))
	d.mustTemplate(&(d.TargetConfDir))
}

func normalizeArch(arch string) string {
	if _, err := strconv.ParseUint(arch, 10, 64); err != nil {
		return arch
	}
	return "i" + arch
}

func (d *Descriptor) Control() string {
	t := `Package: {{.Group}}-{{.Name}}
Version: {{.Version}}
Architecture: ` + normalizeArch(runtime.GOARCH) + `
Maintainer: {{.Author}}
`
	if len(d.Depends) != 0 {
		t += "Depends: " + strings.Join(d.Depends, ",") + "\n"
	}
	t += `Description: {{.Description}}
`
	d.mustTemplate(&t)
	return t
}

func (d *Descriptor) PreInstall() string {
	t := ""
	if d.Service != nil {
		t += mustTemplate("service {{.BinName}} stop || echo 'No installed {{.Group}} {{.Name}} instance running'", *d)
	}
	return t
}

func (d *Descriptor) PostInstall() string {
	t := ""
	if d.Service != nil {
		if d.Service.AutoStart {
			t += mustTemplate("service {{.BinName}} start || echo 'No {{.Group}} {{.Name}} instance'", *d)
		}
	}
	return t
}

func (d *Descriptor) PreRemove() string {
	return d.PreInstall()
}

func (d *Descriptor) ServiceInit() string {
	t := `# {{.Group}} {{.Name}}

description         "{{.Description}}"

start on runlevel [2345]
stop on runlevel [!2345]
{{if .Service.Restart}}
respawn
respawn limit 99999999 {{.Service.RestartDelay}}
{{end}}

script
  . {{.TargetConfDir}}/` + ServiceConfigFile + `
  exec {{.TargetBinDir}}/{{.BinName}} $RUN_OPTS | logger -t '{{.Group}}-{{.Name}}' 2>&1
end script
`
	return mustTemplate(t, *d)
}

func (d *Descriptor) ServiceFile() string {
	t := `{{.Group}}-{{.Name}}.conf`
	return mustTemplate(t, *d)
}

func (d *Descriptor) ServiceConfig() string {
	t := `#!/bin/bash
{{ range $key, $value := .Service.Env}}
{{$key}}="{{$value}}"{{end}}
`
	return mustTemplate(t, *d) + "RUN_OPTS=\"" + d.Service.RunOpts + "\""
}

func mustTemplate(pattern string, params interface{}) string {
	t, err := template.New("").Parse(pattern)
	if err != nil {
		panic("Failed parse templates in " + (pattern) + ": " + err.Error())
	}
	buf := new(bytes.Buffer)
	err = t.Execute(buf, params)
	if err != nil {
		panic("Failed execute templates in " + (pattern) + ": " + err.Error())
	}
	return buf.String()
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"strings"

	"github.com/sirupsen/logrus"
	"sigs.k8s.io/yaml"

	prowapi "k8s.io/test-infra/prow/apis/prowjobs/v1"
	"k8s.io/test-infra/prow/config"
	"k8s.io/test-infra/prow/config/secret"
	prowflagutil "k8s.io/test-infra/prow/flagutil"
	"k8s.io/test-infra/prow/github"
	"k8s.io/test-infra/prow/pjutil"
)

type options struct {
	jobName       string
	configPath    string
	jobConfigPath string

	baseRef    string
	baseSha    string
	pullNumber int
	pullSha    string
	pullAuthor string
	org        string
	repo       string

	local bool

	github       prowflagutil.GitHubOptions
	githubClient githubClient
	pullRequest  *github.PullRequest
}

func (o *options) getPullRequest() (*github.PullRequest, error) {
	if o.pullRequest != nil {
		return o.pullRequest, nil
	}
	pr, err := o.githubClient.GetPullRequest(o.org, o.repo, o.pullNumber)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch PullRequest from GitHub: %v", err)
	}
	o.pullRequest = pr
	return pr, nil
}

func (o *options) defaultPR(pjs *prowapi.ProwJobSpec) error {
	if pjs.Refs.Pulls[0].Number == 0 {
		fmt.Fprint(os.Stderr, "PR Number: ")
		var pullNumber int
		fmt.Scanln(&pullNumber)
		pjs.Refs.Pulls[0].Number = pullNumber
		o.pullNumber = pullNumber
	}
	if pjs.Refs.Pulls[0].Author == "" {
		pr, err := o.getPullRequest()
		if err != nil {
			return err
		}
		pjs.Refs.Pulls[0].Author = pr.User.Login
	}
	if pjs.Refs.Pulls[0].SHA == "" {
		pr, err := o.getPullRequest()
		if err != nil {
			return err
		}
		pjs.Refs.Pulls[0].SHA = pr.Head.SHA
	}
	return nil
}

func (o *options) defaultBaseRef(pjs *prowapi.ProwJobSpec) error {
	if pjs.Refs.BaseRef == "" {
		if o.pullNumber != 0 {
			pr, err := o.getPullRequest()
			if err != nil {
				return err
			}
			pjs.Refs.BaseRef = pr.Base.Ref
		} else {
			fmt.Fprint(os.Stderr, "Base ref (e.g. master): ")
			fmt.Scanln(&pjs.Refs.BaseRef)
		}
	}
	if pjs.Refs.BaseSHA == "" {
		if o.pullNumber != 0 {
			pr, err := o.getPullRequest()
			if err != nil {
				return err
			}
			pjs.Refs.BaseSHA = pr.Base.SHA
		} else {
			baseSHA, err := o.githubClient.GetRef(o.org, o.repo, fmt.Sprintf("heads/%s", pjs.Refs.BaseRef))
			if err != nil {
				logrus.Fatalf("failed to get base sha: %v", err)
				return err
			}
			pjs.Refs.BaseSHA = baseSHA
		}
	}
	return nil
}

type githubClient interface {
	GetPullRequest(org, repo string, number int) (*github.PullRequest, error)
	GetRef(org, repo, ref string) (string, error)
}

func (o *options) Validate() error {
	if o.jobName == "" {
		return errors.New("required flag --job was unset")
	}

	if o.configPath == "" {
		return errors.New("required flag --config-path was unset")
	}

	if err := o.github.Validate(false); err != nil {
		return err
	}

	return nil
}

func gatherOptions() options {
	var o options
	fs := flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	fs.StringVar(&o.jobName, "job", "", "Job to run.")
	fs.BoolVar(&o.local, "local", false, "Print help for running locally")
	fs.StringVar(&o.configPath, "config-path", "", "Path to config.yaml.")
	fs.StringVar(&o.jobConfigPath, "job-config-path", "", "Path to prow job configs.")
	fs.StringVar(&o.baseRef, "base-ref", "", "Git base ref under test")
	fs.StringVar(&o.baseSha, "base-sha", "", "Git base SHA under test")
	fs.IntVar(&o.pullNumber, "pull-number", 0, "Git pull number under test")
	fs.StringVar(&o.pullSha, "pull-sha", "", "Git pull SHA under test")
	fs.StringVar(&o.pullAuthor, "pull-author", "", "Git pull author under test")
	o.github.AddFlagsWithoutDefaultGitHubTokenPath(fs)
	fs.Parse(os.Args[1:])
	return o
}

func main() {
	o := gatherOptions()
	if err := o.Validate(); err != nil {
		logrus.WithError(err).Fatalf("Bad flags")
	}

	conf, err := config.Load(o.configPath, o.jobConfigPath)
	if err != nil {
		logrus.WithError(err).Fatal("Error loading config")
	}

	var secretAgent *secret.Agent
	if o.github.TokenPath != "" {
		secretAgent = &secret.Agent{}
		if err := secretAgent.Start([]string{o.github.TokenPath}); err != nil {
			logrus.WithError(err).Fatal("Failed to start secret agent")
		}
	}
	o.githubClient, err = o.github.GitHubClient(secretAgent, false)
	if err != nil {
		logrus.WithError(err).Fatal("Failed to get GitHub client")
	}

	var pjs prowapi.ProwJobSpec
	var labels map[string]string
	var annotations map[string]string
	var found bool
	for fullRepoName, ps := range conf.Presubmits {
		org, repo, err := splitRepoName(fullRepoName)
		if err != nil {
			logrus.WithError(err).Warnf("Invalid repo name %s.", fullRepoName)
			continue
		}
		for _, p := range ps {
			if p.Name == o.jobName {
				pjs = pjutil.PresubmitSpec(p, prowapi.Refs{
					Org:     org,
					Repo:    repo,
					BaseRef: o.baseRef,
					BaseSHA: o.baseSha,
					Pulls: []prowapi.Pull{{
						Author: o.pullAuthor,
						Number: o.pullNumber,
						SHA:    o.pullSha,
					}},
				})
				labels = p.Labels
				annotations = p.Annotations
				found = true
			}
		}
	}
	for fullRepoName, ps := range conf.Postsubmits {
		org, repo, err := splitRepoName(fullRepoName)
		if err != nil {
			logrus.WithError(err).Warnf("Invalid repo name %s.", fullRepoName)
			continue
		}
		for _, p := range ps {
			if p.Name == o.jobName {
				pjs = pjutil.PostsubmitSpec(p, prowapi.Refs{
					Org:     org,
					Repo:    repo,
					BaseRef: o.baseRef,
					BaseSHA: o.baseSha,
				})
				labels = p.Labels
				annotations = p.Annotations
				found = true
			}
		}
	}
	for _, p := range conf.Periodics {
		if p.Name == o.jobName {
			pjs = pjutil.PeriodicSpec(p)
			labels = p.Labels
			annotations = p.Annotations
			found = true
		}
	}
	if !found {
		logrus.Fatalf("Job %s not found.", o.jobName)
	}
	if pjs.Refs != nil {
		o.org = pjs.Refs.Org
		o.repo = pjs.Refs.Repo
		if len(pjs.Refs.Pulls) != 0 {
			if err := o.defaultPR(&pjs); err != nil {
				logrus.WithError(err).Fatal("Failed to default PR")
			}
		}
		if err := o.defaultBaseRef(&pjs); err != nil {
			logrus.WithError(err).Fatal("Failed to default base ref")
		}
	}
	pj := pjutil.NewProwJob(pjs, labels, annotations)
	b, err := yaml.Marshal(&pj)
	if err != nil {
		logrus.WithError(err).Fatal("Error marshalling YAML.")
	}
	fmt.Print(string(b))
	if o.local {
		logrus.Info("Use 'bazel run //prow/cmd/phaino' to run this job locally in docker")
	}
}

func splitRepoName(repo string) (string, string, error) {
	s := strings.SplitN(repo, "/", 2)
	if len(s) != 2 {
		return "", "", fmt.Errorf("repo %s cannot be split into org/repo", repo)
	}
	return s[0], s[1], nil
}
<commit_msg>mkpj: reduce number of local variables<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"strings"

	"github.com/sirupsen/logrus"
	"sigs.k8s.io/yaml"

	prowapi "k8s.io/test-infra/prow/apis/prowjobs/v1"
	"k8s.io/test-infra/prow/config"
	"k8s.io/test-infra/prow/config/secret"
	prowflagutil "k8s.io/test-infra/prow/flagutil"
	"k8s.io/test-infra/prow/github"
	"k8s.io/test-infra/prow/pjutil"
)

type options struct {
	jobName       string
	configPath    string
	jobConfigPath string

	baseRef    string
	baseSha    string
	pullNumber int
	pullSha    string
	pullAuthor string
	org        string
	repo       string

	local bool

	github       prowflagutil.GitHubOptions
	githubClient githubClient
	pullRequest  *github.PullRequest
}

func (o *options) getPullRequest() (*github.PullRequest, error) {
	if o.pullRequest != nil {
		return o.pullRequest, nil
	}
	pr, err := o.githubClient.GetPullRequest(o.org, o.repo, o.pullNumber)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch PullRequest from GitHub: %v", err)
	}
	o.pullRequest = pr
	return pr, nil
}

func (o *options) defaultPR(pjs *prowapi.ProwJobSpec) error {
	if pjs.Refs.Pulls[0].Number == 0 {
		fmt.Fprint(os.Stderr, "PR Number: ")
		var pullNumber int
		fmt.Scanln(&pullNumber)
		pjs.Refs.Pulls[0].Number = pullNumber
		o.pullNumber = pullNumber
	}
	if pjs.Refs.Pulls[0].Author == "" {
		pr, err := o.getPullRequest()
		if err != nil {
			return err
		}
		pjs.Refs.Pulls[0].Author = pr.User.Login
	}
	if pjs.Refs.Pulls[0].SHA == "" {
		pr, err := o.getPullRequest()
		if err != nil {
			return err
		}
		pjs.Refs.Pulls[0].SHA = pr.Head.SHA
	}
	return nil
}

func (o *options) defaultBaseRef(pjs *prowapi.ProwJobSpec) error {
	if pjs.Refs.BaseRef == "" {
		if o.pullNumber != 0 {
			pr, err := o.getPullRequest()
			if err != nil {
				return err
			}
			pjs.Refs.BaseRef = pr.Base.Ref
		} else {
			fmt.Fprint(os.Stderr, "Base ref (e.g. master): ")
			fmt.Scanln(&pjs.Refs.BaseRef)
		}
	}
	if pjs.Refs.BaseSHA == "" {
		if o.pullNumber != 0 {
			pr, err := o.getPullRequest()
			if err != nil {
				return err
			}
			pjs.Refs.BaseSHA = pr.Base.SHA
		} else {
			baseSHA, err := o.githubClient.GetRef(o.org, o.repo, fmt.Sprintf("heads/%s", pjs.Refs.BaseRef))
			if err != nil {
				logrus.Fatalf("failed to get base sha: %v", err)
				return err
			}
			pjs.Refs.BaseSHA = baseSHA
		}
	}
	return nil
}

type githubClient interface {
	GetPullRequest(org, repo string, number int) (*github.PullRequest, error)
	GetRef(org, repo, ref string) (string, error)
}

func (o *options) Validate() error {
	if o.jobName == "" {
		return errors.New("required flag --job was unset")
	}

	if o.configPath == "" {
		return errors.New("required flag --config-path was unset")
	}

	if err := o.github.Validate(false); err != nil {
		return err
	}

	return nil
}

func gatherOptions() options {
	var o options
	fs := flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	fs.StringVar(&o.jobName, "job", "", "Job to run.")
	fs.BoolVar(&o.local, "local", false, "Print help for running locally")
	fs.StringVar(&o.configPath, "config-path", "", "Path to config.yaml.")
	fs.StringVar(&o.jobConfigPath, "job-config-path", "", "Path to prow job configs.")
	fs.StringVar(&o.baseRef, "base-ref", "", "Git base ref under test")
	fs.StringVar(&o.baseSha, "base-sha", "", "Git base SHA under test")
	fs.IntVar(&o.pullNumber, "pull-number", 0, "Git pull number under test")
	fs.StringVar(&o.pullSha, "pull-sha", "", "Git pull SHA under test")
	fs.StringVar(&o.pullAuthor, "pull-author", "", "Git pull author under test")
	o.github.AddFlagsWithoutDefaultGitHubTokenPath(fs)
	fs.Parse(os.Args[1:])
	return o
}

func main() {
	o := gatherOptions()
	if err := o.Validate(); err != nil {
		logrus.WithError(err).Fatalf("Bad flags")
	}

	conf, err := config.Load(o.configPath, o.jobConfigPath)
	if err != nil {
		logrus.WithError(err).Fatal("Error loading config")
	}

	var secretAgent *secret.Agent
	if o.github.TokenPath != "" {
		secretAgent = &secret.Agent{}
		if err := secretAgent.Start([]string{o.github.TokenPath}); err != nil {
			logrus.WithError(err).Fatal("Failed to start secret agent")
		}
	}
	o.githubClient, err = o.github.GitHubClient(secretAgent, false)
	if err != nil {
		logrus.WithError(err).Fatal("Failed to get GitHub client")
	}

	var job config.JobBase
	var pjs prowapi.ProwJobSpec
	for fullRepoName, ps := range conf.Presubmits {
		org, repo, err := splitRepoName(fullRepoName)
		if err != nil {
			logrus.WithError(err).Warnf("Invalid repo name %s.", fullRepoName)
			continue
		}
		for _, p := range ps {
			if p.Name == o.jobName {
				job = p.JobBase
				pjs = pjutil.PresubmitSpec(p, prowapi.Refs{
					Org:     org,
					Repo:    repo,
					BaseRef: o.baseRef,
					BaseSHA: o.baseSha,
					Pulls: []prowapi.Pull{{
						Author: o.pullAuthor,
						Number: o.pullNumber,
						SHA:    o.pullSha,
					}},
				})
			}
		}
	}
	for fullRepoName, ps := range conf.Postsubmits {
		org, repo, err := splitRepoName(fullRepoName)
		if err != nil {
			logrus.WithError(err).Warnf("Invalid repo name %s.", fullRepoName)
			continue
		}
		for _, p := range ps {
			if p.Name == o.jobName {
				job = p.JobBase
				pjs = pjutil.PostsubmitSpec(p, prowapi.Refs{
					Org:     org,
					Repo:    repo,
					BaseRef: o.baseRef,
					BaseSHA: o.baseSha,
				})
			}
		}
	}
	for _, p := range conf.Periodics {
		if p.Name == o.jobName {
			job = p.JobBase
			pjs = pjutil.PeriodicSpec(p)
		}
	}
	if job.Name == "" {
		logrus.Fatalf("Job %s not found.", o.jobName)
	}
	if pjs.Refs != nil {
		o.org = pjs.Refs.Org
		o.repo = pjs.Refs.Repo
		if len(pjs.Refs.Pulls) != 0 {
			if err := o.defaultPR(&pjs); err != nil {
				logrus.WithError(err).Fatal("Failed to default PR")
			}
		}
		if err := o.defaultBaseRef(&pjs); err != nil {
			logrus.WithError(err).Fatal("Failed to default base ref")
		}
	}
	pj := pjutil.NewProwJob(pjs, job.Labels, job.Annotations)
	b, err := yaml.Marshal(&pj)
	if err != nil {
		logrus.WithError(err).Fatal("Error marshalling YAML.")
	}
	fmt.Print(string(b))
	if o.local {
		logrus.Info("Use 'bazel run //prow/cmd/phaino' to run this job locally in docker")
	}
}

func splitRepoName(repo string) (string, string, error) {
	s := strings.SplitN(repo, "/", 2)
	if len(s) != 2 {
		return "", "", fmt.Errorf("repo %s cannot be split into org/repo", repo)
	}
	return s[0], s[1], nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Pulcy.
//
// 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 main

import (
	"os"
)

const (
	defaultInstanceCount  = 3
	defaultGluonImage     = "pulcy/gluon:0.14.0"
	defaultRebootStrategy = "etcd-lock"
)

func defaultDomain() string {
	return os.Getenv("QUARK_DOMAIN")
}

func defaultPrivateRegistryUrl() string {
	return os.Getenv("QUARK_REGISTRY_URL")
}

func defaultPrivateRegistryUserName() string {
	return os.Getenv("QUARK_REGISTRY_USERNAME")
}

func defaultPrivateRegistryPassword() string {
	return os.Getenv("QUARK_REGISTRY_PASSWORD")
}

func defaultSshKeys() []string {
	return []string{os.Getenv("QUARK_SSH_KEY")}
}

func defaultSshKeyGithubAccount() string {
	return os.Getenv("QUARK_SSH_KEY_GITHUB_ACCOUNT")
}

func defaultVagrantFolder() string {
	return os.Getenv("QUARK_VAGRANT_FOLDER")
}

func defaultVaultAddr() string {
	return os.Getenv("VAULT_ADDR")
}

func defaultVaultCACert() string {
	return os.Getenv("VAULT_CACERT")
}
<commit_msg>Updated gluon to 0.14.1<commit_after>// Copyright (c) 2016 Pulcy.
//
// 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 main

import (
	"os"
)

const (
	defaultInstanceCount  = 3
	defaultGluonImage     = "pulcy/gluon:0.14.1"
	defaultRebootStrategy = "etcd-lock"
)

func defaultDomain() string {
	return os.Getenv("QUARK_DOMAIN")
}

func defaultPrivateRegistryUrl() string {
	return os.Getenv("QUARK_REGISTRY_URL")
}

func defaultPrivateRegistryUserName() string {
	return os.Getenv("QUARK_REGISTRY_USERNAME")
}

func defaultPrivateRegistryPassword() string {
	return os.Getenv("QUARK_REGISTRY_PASSWORD")
}

func defaultSshKeys() []string {
	return []string{os.Getenv("QUARK_SSH_KEY")}
}

func defaultSshKeyGithubAccount() string {
	return os.Getenv("QUARK_SSH_KEY_GITHUB_ACCOUNT")
}

func defaultVagrantFolder() string {
	return os.Getenv("QUARK_VAGRANT_FOLDER")
}

func defaultVaultAddr() string {
	return os.Getenv("VAULT_ADDR")
}

func defaultVaultCACert() string {
	return os.Getenv("VAULT_CACERT")
}
<|endoftext|>


<commit_before>package main

import (
	"gobot.io/x/gobot"
	"gobot.io/x/gobot/platforms/intel-iot/edison"
	"gobot.io/x/gobot/drivers/i2c"
	"gobot.io/x/gobot/drivers/aio"
	"time"
	"fmt"
	"math"
)

const (
	sensitivity = 25
)

func average(slc[]int)int {
	total := 0
	for _, v := range slc {
		total +=v
	}
	return total/len(slc)
}

func calibrateLighting(lightSensor *aio.GroveLightSensorDriver) int {
	completeTime := time.Now().Add(3 * time.Second)
	readings := []int{}
	for time.Now().Unix() < completeTime.Unix() {
		reading, _ := lightSensor.Read()
		readings = append(readings, reading)
	}
	return average(readings)
}

func main() {
	e := edison.NewAdaptor()
	lightSensor := aio.NewGroveLightSensorDriver(e, "0", 250)
	screen := i2c.NewGroveLcdDriver(e)

	fmt.Println(calibrateLighting(lightSensor))
	averageLight := calibrateLighting(lightSensor)

	work := func() {
		for {
			lightScore, _ := lightSensor.Read()
			delta := averageLight - lightScore
			if math.Abs(float64(delta)) > sensitivity {
				screen.SetRGB(255, 0, 0)
				time.Sleep(1 * time.Second)
			}
			screen.SetRGB(0, 255, 0)
		}
	}

	robot := gobot.NewRobot("buttonBot",
		[]gobot.Connection{e},
		[]gobot.Device{screen, lightSensor},
		work,
	)

	robot.Start()
}<commit_msg>Added basic temperature readings<commit_after>package main

import (
	"gobot.io/x/gobot"
	"gobot.io/x/gobot/platforms/intel-iot/edison"
	"gobot.io/x/gobot/drivers/i2c"
	"gobot.io/x/gobot/drivers/aio"
	"time"
	"fmt"
	"math"
)

const (
	lightTolerance = 50
)

func average(slc[]int)int {
	total := 0
	for _, v := range slc {
		total +=v
	}
	return total/len(slc)
}

func calibrateLighting(lightSensor *aio.GroveLightSensorDriver) int {
	completeTime := time.Now().Add(3 * time.Second)
	readings := []int{}
	for time.Now().Unix() < completeTime.Unix() {
		reading, _ := lightSensor.Read()
		readings = append(readings, reading)
	}
	return average(readings)
}

func main() {
	e := edison.NewAdaptor()
	lightSensor := aio.NewGroveLightSensorDriver(e, "0", 250)
	temperatureSensor := aio.NewGroveTemperatureSensorDriver(e, "3")
	screen := i2c.NewGroveLcdDriver(e)

	fmt.Println(calibrateLighting(lightSensor))
	averageLight := calibrateLighting(lightSensor)

	work := func() {
		for {
			lightScore, _ := lightSensor.Read()
			temperature := temperatureSensor.Temperature()
			fmt.Println(temperature)
			delta := averageLight - lightScore
			if math.Abs(float64(delta)) > lightTolerance {
				screen.SetRGB(255, 0, 0)
			} else {
				screen.SetRGB(0, 255, 0)
			}

		}
	}

	robot := gobot.NewRobot("buttonBot",
		[]gobot.Connection{e},
		[]gobot.Device{screen, lightSensor, temperatureSensor},
		work,
	)

	robot.Start()
}<|endoftext|>


<commit_before>package core

import (
	"errors"
	"reflect"
	"strings"
	"syscall"
	"testing"
	"time"

	"github.com/elves/elvish/edit/tty"
	"github.com/elves/elvish/edit/ui"
	"github.com/elves/elvish/styled"
	"github.com/elves/elvish/sys"
)

func TestReadCode_AbortsOnSetupError(t *testing.T) {
	terminal := newFakeTTY()
	terminal.setupErr = errors.New("a fake error")

	ed := NewEditor(terminal, nil)
	_, err := ed.ReadCode()

	if err != terminal.setupErr {
		t.Errorf("ReadCode returns error %v, want %v", err, terminal.setupErr)
	}
}

func TestReadCode_CallsRestore(t *testing.T) {
	restoreCalled := 0
	terminal := newFakeTTY()
	terminal.restoreFunc = func() { restoreCalled++ }
	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}

	ed := NewEditor(terminal, nil)
	ed.ReadCode()

	if restoreCalled != 1 {
		t.Errorf("Restore callback called %d times, want once", restoreCalled)
	}
}

func TestReadCode_ResetsStateBeforeReturn(t *testing.T) {
	terminal := newFakeTTY()
	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}

	ed := NewEditor(terminal, nil)
	ed.State.Raw.Code = "some code"
	ed.ReadCode()

	if code := ed.State.Raw.Code; code != "" {
		t.Errorf("Editor state has code %q, want empty", code)
	}
}

func TestReadCode_PassesInputEventsToMode(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)
	m := &fakeMode{maxKeys: 3}
	ed.State.Raw.Mode = m

	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'b'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'c'}

	ed.ReadCode()

	wantKeysHandled := []ui.Key{
		ui.Key{Rune: 'a'}, ui.Key{Rune: 'b'}, ui.Key{Rune: 'c'},
	}
	if !reflect.DeepEqual(m.keysHandled, wantKeysHandled) {
		t.Errorf("Mode gets keys %v, want %v", m.keysHandled, wantKeysHandled)
	}
}

func TestReadCode_CallsBeforeReadlineOnce(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)

	called := 0
	ed.Config.Raw.BeforeReadline = []func(){func() { called++ }}

	// Causes basicMode to quit
	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}

	ed.ReadCode()

	if called != 1 {
		t.Errorf("BeforeReadline hook called %d times, want 1", called)
	}
}

func TestReadCode_CallsAfterReadlineOnceWithCode(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)

	called := 0
	code := ""
	ed.Config.Raw.AfterReadline = []func(string){func(s string) {
		called++
		code = s
	}}

	// Causes basicMode to write state.Code and then quit
	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'b'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'c'}
	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}

	ed.ReadCode()

	if called != 1 {
		t.Errorf("AfterReadline hook called %d times, want 1", called)
	}
	if code != "abc" {
		t.Errorf("AfterReadline hook called with %q, want %q", code, "abc")
	}
}

func TestReadCode_RespectsMaxHeight(t *testing.T) {
	maxHeight := 5

	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)
	// Will fill more than maxHeight but less than terminal height
	ed.State.Raw.Code = strings.Repeat("a", 80*10)
	ed.State.Raw.Dot = len(ed.State.Raw.Code)

	codeCh, _ := ed.readCodeAsync()

	buf1 := <-terminal.bufCh
	// Make sure that normally the height does exceed maxHeight.
	if h := len(buf1.Lines); h <= maxHeight {
		t.Errorf("Buffer height is %d, should > %d", h, maxHeight)
	}

	ed.Config.Mutex.Lock()
	ed.Config.Raw.MaxHeight = maxHeight
	ed.Config.Mutex.Unlock()

	ed.loop.Redraw(false)
	buf2 := <-terminal.bufCh
	if h := len(buf2.Lines); h > maxHeight {
		t.Errorf("Buffer height is %d, should <= %d", h, maxHeight)
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}

var bufChTimeout = 1 * time.Second

func TestReadCode_RendersHighlightedCode(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)
	ed.Config.Raw.Highlighter = func(code string) (styled.Text, []error) {
		return styled.Text{
			&styled.Segment{styled.Style{Foreground: "red"}, code}}, nil
	}

	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'b'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'c'}
	codeCh, _ := ed.readCodeAsync()

	wantBuf := ui.NewBufferBuilder(80).
		WriteString("abc", "31" /* SGR for red foreground */).
		SetDotToCursor().Buffer()
	if !terminal.checkBuffer(wantBuf) {
		t.Errorf("Did not see buffer containing highlighted code")
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}

func TestReadCode_RendersErrorFromHighlighter(t *testing.T) {
	// TODO
}

func TestReadCode_RendersPrompt(t *testing.T) {
	// TODO
}

func TestReadCode_RendersRPrompt(t *testing.T) {
	// TODO
}

func TestReadCode_SupportsPersistentRPrompt(t *testing.T) {
	// TODO
}

func TestReadCode_UsesFinalStateInFinalRedraw(t *testing.T) {
	terminal := newFakeTTY()

	ed := NewEditor(terminal, nil)
	ed.State.Raw.Code = "some code"
	// We use the dot as a signal for distinguishing non-final and final state.
	// In the final state, the dot will be set to the length of the code (9).
	ed.State.Raw.Dot = 1

	codeCh, _ := ed.readCodeAsync()
	// Wait until a non-final state is drawn.
	wantBuf := ui.NewBufferBuilder(80).WriteUnstyled("s").SetDotToCursor().
		WriteUnstyled("ome code").Buffer()
	if !terminal.checkBuffer(wantBuf) {
		t.Errorf("did not get expected buffer before sending Enter")
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh

	// Last element in bufs is nil
	finalBuf := terminal.bufs[len(terminal.bufs)-2]
	wantFinalBuf := ui.NewBufferBuilder(80).WriteUnstyled("some code").
		SetDotToCursor().Buffer()
	if !reflect.DeepEqual(finalBuf, wantFinalBuf) {
		t.Errorf("final buffer is %v, want %v", finalBuf, wantFinalBuf)
	}
}

func TestReadCode_QuitsOnSIGHUP(t *testing.T) {
	terminal := newFakeTTY()
	sigs := newFakeSignalSource()
	ed := NewEditor(terminal, sigs)

	codeCh, _ := ed.readCodeAsync()
	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	wantBuf := ui.NewBufferBuilder(80).WriteUnstyled("a").
		SetDotToCursor().Buffer()
	if !terminal.checkBuffer(wantBuf) {
		t.Errorf("did not get expected buffer before sending SIGHUP")
	}

	sigs.ch <- syscall.SIGHUP

	select {
	case <-codeCh:
		// TODO: Test that ReadCode returns with io.EOF
	case <-time.After(time.Second):
		t.Errorf("SIGHUP did not cause ReadCode to return")
	}
}

func TestReadCode_ResetsOnSIGHUP(t *testing.T) {
	terminal := newFakeTTY()
	sigs := newFakeSignalSource()
	ed := NewEditor(terminal, sigs)

	codeCh, _ := ed.readCodeAsync()
	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	wantBuf := ui.NewBufferBuilder(80).WriteUnstyled("a").
		SetDotToCursor().Buffer()
	if !terminal.checkBuffer(wantBuf) {
		t.Errorf("did not get expected buffer before sending SIGINT")
	}

	sigs.ch <- syscall.SIGINT

	wantBuf = ui.NewBufferBuilder(80).Buffer()
	if !terminal.checkBuffer(wantBuf) {
		t.Errorf("Terminal state is not reset after SIGINT")
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}

func TestReadCode_RedrawsOnSIGWINCH(t *testing.T) {
	terminal := newFakeTTY()
	sigs := newFakeSignalSource()
	ed := NewEditor(terminal, sigs)

	ed.State.Raw.Code = "1234567890"
	ed.State.Raw.Dot = len(ed.State.Raw.Code)

	codeCh, _ := ed.readCodeAsync()
	wantBuf := ui.NewBufferBuilder(80).WriteUnstyled("1234567890").
		SetDotToCursor().Buffer()
	if !terminal.checkBuffer(wantBuf) {
		t.Errorf("did not get expected buffer before sending SIGWINCH")
	}

	terminal.setSize(24, 4)
	sigs.ch <- sys.SIGWINCH

	wantBuf = ui.NewBufferBuilder(4).WriteUnstyled("1234567890").
		SetDotToCursor().Buffer()
	if !terminal.checkBuffer(wantBuf) {
		t.Errorf("Terminal is not redrawn after SIGWINCH")
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}
<commit_msg>newedit/core: Trivial cleanup.<commit_after>package core

import (
	"errors"
	"reflect"
	"strings"
	"syscall"
	"testing"
	"time"

	"github.com/elves/elvish/edit/tty"
	"github.com/elves/elvish/edit/ui"
	"github.com/elves/elvish/styled"
	"github.com/elves/elvish/sys"
)

func TestReadCode_AbortsOnSetupError(t *testing.T) {
	terminal := newFakeTTY()
	terminal.setupErr = errors.New("a fake error")

	ed := NewEditor(terminal, nil)
	_, err := ed.ReadCode()

	if err != terminal.setupErr {
		t.Errorf("ReadCode returns error %v, want %v", err, terminal.setupErr)
	}
}

func TestReadCode_CallsRestore(t *testing.T) {
	restoreCalled := 0
	terminal := newFakeTTY()
	terminal.restoreFunc = func() { restoreCalled++ }
	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}

	ed := NewEditor(terminal, nil)
	ed.ReadCode()

	if restoreCalled != 1 {
		t.Errorf("Restore callback called %d times, want once", restoreCalled)
	}
}

func TestReadCode_ResetsStateBeforeReturn(t *testing.T) {
	terminal := newFakeTTY()
	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}

	ed := NewEditor(terminal, nil)
	ed.State.Raw.Code = "some code"
	ed.ReadCode()

	if code := ed.State.Raw.Code; code != "" {
		t.Errorf("Editor state has code %q, want empty", code)
	}
}

func TestReadCode_PassesInputEventsToMode(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)
	m := &fakeMode{maxKeys: 3}
	ed.State.Raw.Mode = m

	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'b'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'c'}

	ed.ReadCode()

	wantKeysHandled := []ui.Key{
		ui.Key{Rune: 'a'}, ui.Key{Rune: 'b'}, ui.Key{Rune: 'c'},
	}
	if !reflect.DeepEqual(m.keysHandled, wantKeysHandled) {
		t.Errorf("Mode gets keys %v, want %v", m.keysHandled, wantKeysHandled)
	}
}

func TestReadCode_CallsBeforeReadlineOnce(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)

	called := 0
	ed.Config.Raw.BeforeReadline = []func(){func() { called++ }}

	// Causes basicMode to quit
	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}

	ed.ReadCode()

	if called != 1 {
		t.Errorf("BeforeReadline hook called %d times, want 1", called)
	}
}

func TestReadCode_CallsAfterReadlineOnceWithCode(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)

	called := 0
	code := ""
	ed.Config.Raw.AfterReadline = []func(string){func(s string) {
		called++
		code = s
	}}

	// Causes basicMode to write state.Code and then quit
	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'b'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'c'}
	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}

	ed.ReadCode()

	if called != 1 {
		t.Errorf("AfterReadline hook called %d times, want 1", called)
	}
	if code != "abc" {
		t.Errorf("AfterReadline hook called with %q, want %q", code, "abc")
	}
}

func TestReadCode_RespectsMaxHeight(t *testing.T) {
	maxHeight := 5

	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)
	// Will fill more than maxHeight but less than terminal height
	ed.State.Raw.Code = strings.Repeat("a", 80*10)
	ed.State.Raw.Dot = len(ed.State.Raw.Code)

	codeCh, _ := ed.readCodeAsync()

	buf1 := <-terminal.bufCh
	// Make sure that normally the height does exceed maxHeight.
	if h := len(buf1.Lines); h <= maxHeight {
		t.Errorf("Buffer height is %d, should > %d", h, maxHeight)
	}

	ed.Config.Mutex.Lock()
	ed.Config.Raw.MaxHeight = maxHeight
	ed.Config.Mutex.Unlock()

	ed.Redraw(false)
	buf2 := <-terminal.bufCh
	if h := len(buf2.Lines); h > maxHeight {
		t.Errorf("Buffer height is %d, should <= %d", h, maxHeight)
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}

var bufChTimeout = 1 * time.Second

func TestReadCode_RendersHighlightedCode(t *testing.T) {
	terminal := newFakeTTY()
	ed := NewEditor(terminal, nil)
	ed.Config.Raw.Highlighter = func(code string) (styled.Text, []error) {
		return styled.Text{
			&styled.Segment{styled.Style{Foreground: "red"}, code}}, nil
	}

	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'b'}
	terminal.eventCh <- tty.KeyEvent{Rune: 'c'}
	codeCh, _ := ed.readCodeAsync()

	wantBuf := ui.NewBufferBuilder(80).
		WriteString("abc", "31" /* SGR for red foreground */).
		SetDotToCursor().Buffer()
	if !terminal.checkBuffer(wantBuf) {
		t.Errorf("Did not see buffer containing highlighted code")
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}

func TestReadCode_RendersErrorFromHighlighter(t *testing.T) {
	// TODO
}

func TestReadCode_RendersPrompt(t *testing.T) {
	// TODO
}

func TestReadCode_RendersRPrompt(t *testing.T) {
	// TODO
}

func TestReadCode_SupportsPersistentRPrompt(t *testing.T) {
	// TODO
}

func TestReadCode_UsesFinalStateInFinalRedraw(t *testing.T) {
	terminal := newFakeTTY()

	ed := NewEditor(terminal, nil)
	ed.State.Raw.Code = "some code"
	// We use the dot as a signal for distinguishing non-final and final state.
	// In the final state, the dot will be set to the length of the code (9).
	ed.State.Raw.Dot = 1

	codeCh, _ := ed.readCodeAsync()
	// Wait until a non-final state is drawn.
	wantBuf := ui.NewBufferBuilder(80).WriteUnstyled("s").SetDotToCursor().
		WriteUnstyled("ome code").Buffer()
	if !terminal.checkBuffer(wantBuf) {
		t.Errorf("did not get expected buffer before sending Enter")
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh

	// Last element in bufs is nil
	finalBuf := terminal.bufs[len(terminal.bufs)-2]
	wantFinalBuf := ui.NewBufferBuilder(80).WriteUnstyled("some code").
		SetDotToCursor().Buffer()
	if !reflect.DeepEqual(finalBuf, wantFinalBuf) {
		t.Errorf("final buffer is %v, want %v", finalBuf, wantFinalBuf)
	}
}

func TestReadCode_QuitsOnSIGHUP(t *testing.T) {
	terminal := newFakeTTY()
	sigs := newFakeSignalSource()
	ed := NewEditor(terminal, sigs)

	codeCh, _ := ed.readCodeAsync()
	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	wantBuf := ui.NewBufferBuilder(80).WriteUnstyled("a").
		SetDotToCursor().Buffer()
	if !terminal.checkBuffer(wantBuf) {
		t.Errorf("did not get expected buffer before sending SIGHUP")
	}

	sigs.ch <- syscall.SIGHUP

	select {
	case <-codeCh:
		// TODO: Test that ReadCode returns with io.EOF
	case <-time.After(time.Second):
		t.Errorf("SIGHUP did not cause ReadCode to return")
	}
}

func TestReadCode_ResetsOnSIGHUP(t *testing.T) {
	terminal := newFakeTTY()
	sigs := newFakeSignalSource()
	ed := NewEditor(terminal, sigs)

	codeCh, _ := ed.readCodeAsync()
	terminal.eventCh <- tty.KeyEvent{Rune: 'a'}
	wantBuf := ui.NewBufferBuilder(80).WriteUnstyled("a").
		SetDotToCursor().Buffer()
	if !terminal.checkBuffer(wantBuf) {
		t.Errorf("did not get expected buffer before sending SIGINT")
	}

	sigs.ch <- syscall.SIGINT

	wantBuf = ui.NewBufferBuilder(80).Buffer()
	if !terminal.checkBuffer(wantBuf) {
		t.Errorf("Terminal state is not reset after SIGINT")
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}

func TestReadCode_RedrawsOnSIGWINCH(t *testing.T) {
	terminal := newFakeTTY()
	sigs := newFakeSignalSource()
	ed := NewEditor(terminal, sigs)

	ed.State.Raw.Code = "1234567890"
	ed.State.Raw.Dot = len(ed.State.Raw.Code)

	codeCh, _ := ed.readCodeAsync()
	wantBuf := ui.NewBufferBuilder(80).WriteUnstyled("1234567890").
		SetDotToCursor().Buffer()
	if !terminal.checkBuffer(wantBuf) {
		t.Errorf("did not get expected buffer before sending SIGWINCH")
	}

	terminal.setSize(24, 4)
	sigs.ch <- sys.SIGWINCH

	wantBuf = ui.NewBufferBuilder(4).WriteUnstyled("1234567890").
		SetDotToCursor().Buffer()
	if !terminal.checkBuffer(wantBuf) {
		t.Errorf("Terminal is not redrawn after SIGWINCH")
	}

	terminal.eventCh <- tty.KeyEvent{Rune: '\n'}
	<-codeCh
}
<|endoftext|>


<commit_before>// +build V7

package rpc

import (
	"bytes"
	"encoding/gob"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
	"sync"

	"code.cloudfoundry.org/cli/command"
	plugin "code.cloudfoundry.org/cli/plugin/v7"
	plugin_models "code.cloudfoundry.org/cli/plugin/v7/models"
	"code.cloudfoundry.org/cli/util/command_parser"
	"code.cloudfoundry.org/cli/version"
	"github.com/blang/semver"
)

type CliRpcCmd struct {
	PluginMetadata *plugin.PluginMetadata
	MetadataMutex  *sync.RWMutex
	Config         command.Config
	PluginActor    PluginActor
	outputCapture  OutputCapture
	outputBucket   *bytes.Buffer
	stdout         io.Writer
}

func (cmd *CliRpcCmd) ApiEndpoint(args string, retVal *string) error {
	*retVal = cmd.Config.Target()

	return nil
}

func (cmd *CliRpcCmd) IsMinCliVersion(passedVersion string, retVal *bool) error {
	if version.VersionString() == version.DefaultVersion {
		*retVal = true
		return nil
	}

	actualVersion, err := semver.Make(version.VersionString())
	if err != nil {
		return err
	}

	requiredVersion, err := semver.Make(passedVersion)
	if err != nil {
		return err
	}

	*retVal = actualVersion.GTE(requiredVersion)

	return nil
}

func (cmd *CliRpcCmd) SetPluginMetadata(pluginMetadata plugin.PluginMetadata, retVal *bool) error {
	cmd.MetadataMutex.Lock()
	defer cmd.MetadataMutex.Unlock()

	cmd.PluginMetadata = &pluginMetadata
	*retVal = true
	return nil
}

func (cmd *CliRpcCmd) CliCommand(args []string, retVal *[]string) error {
	outBuffer := bytes.Buffer{}
	errBuffer := bytes.Buffer{}

	p, err := command_parser.NewCommandParserForPlugins(os.Stdout, &outBuffer, &errBuffer)
	if err != nil {
		return err
	}
	exitCode := p.ParseCommandFromArgs(args)

	if exitCode == command_parser.UnknownCommandCode {
		return errors.New("UnknownCommandCode")
	} else if exitCode > 0 {
		return errors.New("Some other error")
	}

	*retVal = strings.Split(strings.TrimSuffix(outBuffer.String(), "\n"), "\n")

	return nil
}

func (cmd *CliRpcCmd) GetOutputAndReset(args bool, retVal *[]string) error {
	return errors.New("unimplemented")
}

func (cmd *CliRpcCmd) GetApp(appName string, retVal *plugin_models.DetailedApplicationSummary) error {
	if !cmd.Config.HasTargetedOrganization() {
		return errors.New("no organization targeted")
	}
	spaceGUID := cmd.Config.TargetedSpace().GUID
	if spaceGUID == "" {
		return errors.New("no space targeted")
	}

	app, _, err := cmd.PluginActor.GetDetailedAppSummary(appName, spaceGUID, true)

	if err != nil {
		return err
	}
	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	dec := gob.NewDecoder(&b)

	err = enc.Encode(app)
	if err != nil {
		panic(err)
	}

	err = dec.Decode(retVal)
	if err != nil {
		panic(err)
	}
	return nil
}

func (cmd *CliRpcCmd) GetApps(appName string, retVal *[]plugin_models.Application) error {
	spaceGUID := cmd.Config.TargetedSpace().GUID
	if spaceGUID == "" {
		return errors.New("no space targeted")
	}
	applications, warnings, err := cmd.PluginActor.GetApplicationsBySpace(spaceGUID)
	if err != nil {
		return err
	}
	for _, warning := range warnings {
		fmt.Printf("GetApps warning: %s\n", warning)
	}
	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	dec := gob.NewDecoder(&b)

	err = enc.Encode(applications)
	if err != nil {
		panic(err)
	}

	err = dec.Decode(retVal)
	if err != nil {
		panic(err)
	}
	return nil
}

func (cmd *CliRpcCmd) GetCurrentOrg(args string, retVal *plugin_models.Org) error {
	org := cmd.Config.TargetedOrganization()
	if org.Name == "" {
		return errors.New("no organization targeted")
	}

	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	dec := gob.NewDecoder(&b)

	err := enc.Encode(org)
	if err != nil {
		panic(err)
	}

	err = dec.Decode(retVal)
	if err != nil {
		panic(err)
	}
	return nil
}

func (cmd *CliRpcCmd) GetCurrentSpace(args string, retVal *plugin_models.Space) error {
	orgGUID := cmd.Config.TargetedOrganization().GUID
	if orgGUID == "" {
		return errors.New("no organization targeted")
	}
	spaceName := cmd.Config.TargetedSpace().Name
	if spaceName == "" {
		return errors.New("no space targeted")
	}
	space, _, err := cmd.PluginActor.GetSpaceByNameAndOrganization(spaceName, orgGUID)

	if err != nil {
		return err
	}
	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	dec := gob.NewDecoder(&b)

	err = enc.Encode(space)
	if err != nil {
		panic(err)
	}

	err = dec.Decode(retVal)
	if err != nil {
		panic(err)
	}
	return nil
}

func (cmd *CliRpcCmd) GetOrg(orgName string, retVal *plugin_models.OrgSummary) error {
	org, _, err := cmd.PluginActor.GetOrganizationByName(orgName)
	if err != nil {
		return err
	}

	spaces, _, err := cmd.PluginActor.GetOrganizationSpaces(org.GUID)
	if err != nil {
		return err
	}

	domains, _, err := cmd.PluginActor.GetOrganizationDomains(org.GUID, "")
	if err != nil {
		return err
	}

	for _, space := range spaces {
		newSpace := plugin_models.Space{
			GUID: space.GUID,
			Name: space.Name,
		}
		retVal.Spaces = append(retVal.Spaces, newSpace)
	}

	for _, domain := range domains {
		newDomain := plugin_models.Domain{GUID: domain.GUID, Name: domain.Name}
		retVal.Domains = append(retVal.Domains, newDomain)
	}

	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	dec := gob.NewDecoder(&b)

	err = enc.Encode(org.Metadata)
	if err != nil {
		panic(err)
	}

	var metadata plugin_models.Metadata

	err = dec.Decode(&metadata)
	if err != nil {
		panic(err)
	}

	retVal.Name = org.Name
	retVal.GUID = org.GUID
	retVal.Metadata = metadata

	return nil
}

func (cmd *CliRpcCmd) GetSpace(spaceName string, retVal *plugin_models.Space) error {
	orgGUID := cmd.Config.TargetedOrganization().GUID
	if orgGUID == "" {
		return errors.New("no organization targeted")
	}
	space, _, err := cmd.PluginActor.GetSpaceByNameAndOrganization(spaceName, orgGUID)
	if err != nil {
		return err
	}

	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	dec := gob.NewDecoder(&b)

	err = enc.Encode(space)
	if err != nil {
		panic(err)
	}

	err = dec.Decode(retVal)
	if err != nil {
		panic(err)
	}

	return nil
}

func (cmd *CliRpcCmd) GetSpaces(not_used string, retVal *[]plugin_models.Space) error {
	orgGUID := cmd.Config.TargetedOrganization().GUID
	if orgGUID == "" {
		return errors.New("no organization targeted")
	}
	spaces, _, err := cmd.PluginActor.GetOrganizationSpaces(orgGUID)
	if err != nil {
		return err
	}

	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	dec := gob.NewDecoder(&b)

	err = enc.Encode(spaces)
	if err != nil {
		return err
	}

	err = dec.Decode(retVal)
	if err != nil {
		return err
	}

	return nil
}

func (cmd *CliRpcCmd) IsLoggedIn(args string, retVal *bool) error {
	*retVal = cmd.Config.AccessToken() != ""
	return nil
}

func (cmd *CliRpcCmd) IsSkipSSLValidation(args string, retVal *bool) error {
	*retVal = cmd.Config.SkipSSLValidation()
	return nil
}

func (cmd *CliRpcCmd) AccessToken(args string, retVal *string) error {
	accessToken, err := cmd.PluginActor.RefreshAccessToken()
	if err != nil {
		panic(err)
	}
	*retVal = accessToken
	return nil
}

func (cmd *CliRpcCmd) Username(args string, retVal *string) error {
	username, err := cmd.Config.CurrentUserName()
	if err != nil {
		return fmt.Errorf("error processing config: %s", err.Error())
	}
	if username == "" {
		return errors.New("not logged in")
	}
	*retVal = username

	return nil
}
<commit_msg>Remove unused GetOutputAndReset v7 RPC method<commit_after>// +build V7

package rpc

import (
	"bytes"
	"encoding/gob"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
	"sync"

	"code.cloudfoundry.org/cli/command"
	plugin "code.cloudfoundry.org/cli/plugin/v7"
	plugin_models "code.cloudfoundry.org/cli/plugin/v7/models"
	"code.cloudfoundry.org/cli/util/command_parser"
	"code.cloudfoundry.org/cli/version"
	"github.com/blang/semver"
)

type CliRpcCmd struct {
	PluginMetadata *plugin.PluginMetadata
	MetadataMutex  *sync.RWMutex
	Config         command.Config
	PluginActor    PluginActor
	outputCapture  OutputCapture
	outputBucket   *bytes.Buffer
	stdout         io.Writer
}

func (cmd *CliRpcCmd) ApiEndpoint(args string, retVal *string) error {
	*retVal = cmd.Config.Target()

	return nil
}

func (cmd *CliRpcCmd) IsMinCliVersion(passedVersion string, retVal *bool) error {
	if version.VersionString() == version.DefaultVersion {
		*retVal = true
		return nil
	}

	actualVersion, err := semver.Make(version.VersionString())
	if err != nil {
		return err
	}

	requiredVersion, err := semver.Make(passedVersion)
	if err != nil {
		return err
	}

	*retVal = actualVersion.GTE(requiredVersion)

	return nil
}

func (cmd *CliRpcCmd) SetPluginMetadata(pluginMetadata plugin.PluginMetadata, retVal *bool) error {
	cmd.MetadataMutex.Lock()
	defer cmd.MetadataMutex.Unlock()

	cmd.PluginMetadata = &pluginMetadata
	*retVal = true
	return nil
}

func (cmd *CliRpcCmd) CliCommand(args []string, retVal *[]string) error {
	outBuffer := bytes.Buffer{}
	errBuffer := bytes.Buffer{}

	p, err := command_parser.NewCommandParserForPlugins(os.Stdout, &outBuffer, &errBuffer)
	if err != nil {
		return err
	}
	exitCode := p.ParseCommandFromArgs(args)

	if exitCode == command_parser.UnknownCommandCode {
		return errors.New("UnknownCommandCode")
	} else if exitCode > 0 {
		return errors.New("Some other error")
	}

	*retVal = strings.Split(strings.TrimSuffix(outBuffer.String(), "\n"), "\n")

	return nil
}

func (cmd *CliRpcCmd) GetApp(appName string, retVal *plugin_models.DetailedApplicationSummary) error {
	if !cmd.Config.HasTargetedOrganization() {
		return errors.New("no organization targeted")
	}
	spaceGUID := cmd.Config.TargetedSpace().GUID
	if spaceGUID == "" {
		return errors.New("no space targeted")
	}

	app, _, err := cmd.PluginActor.GetDetailedAppSummary(appName, spaceGUID, true)

	if err != nil {
		return err
	}
	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	dec := gob.NewDecoder(&b)

	err = enc.Encode(app)
	if err != nil {
		panic(err)
	}

	err = dec.Decode(retVal)
	if err != nil {
		panic(err)
	}
	return nil
}

func (cmd *CliRpcCmd) GetApps(appName string, retVal *[]plugin_models.Application) error {
	spaceGUID := cmd.Config.TargetedSpace().GUID
	if spaceGUID == "" {
		return errors.New("no space targeted")
	}
	applications, warnings, err := cmd.PluginActor.GetApplicationsBySpace(spaceGUID)
	if err != nil {
		return err
	}
	for _, warning := range warnings {
		fmt.Printf("GetApps warning: %s\n", warning)
	}
	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	dec := gob.NewDecoder(&b)

	err = enc.Encode(applications)
	if err != nil {
		panic(err)
	}

	err = dec.Decode(retVal)
	if err != nil {
		panic(err)
	}
	return nil
}

func (cmd *CliRpcCmd) GetCurrentOrg(args string, retVal *plugin_models.Org) error {
	org := cmd.Config.TargetedOrganization()
	if org.Name == "" {
		return errors.New("no organization targeted")
	}

	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	dec := gob.NewDecoder(&b)

	err := enc.Encode(org)
	if err != nil {
		panic(err)
	}

	err = dec.Decode(retVal)
	if err != nil {
		panic(err)
	}
	return nil
}

func (cmd *CliRpcCmd) GetCurrentSpace(args string, retVal *plugin_models.Space) error {
	orgGUID := cmd.Config.TargetedOrganization().GUID
	if orgGUID == "" {
		return errors.New("no organization targeted")
	}
	spaceName := cmd.Config.TargetedSpace().Name
	if spaceName == "" {
		return errors.New("no space targeted")
	}
	space, _, err := cmd.PluginActor.GetSpaceByNameAndOrganization(spaceName, orgGUID)

	if err != nil {
		return err
	}
	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	dec := gob.NewDecoder(&b)

	err = enc.Encode(space)
	if err != nil {
		panic(err)
	}

	err = dec.Decode(retVal)
	if err != nil {
		panic(err)
	}
	return nil
}

func (cmd *CliRpcCmd) GetOrg(orgName string, retVal *plugin_models.OrgSummary) error {
	org, _, err := cmd.PluginActor.GetOrganizationByName(orgName)
	if err != nil {
		return err
	}

	spaces, _, err := cmd.PluginActor.GetOrganizationSpaces(org.GUID)
	if err != nil {
		return err
	}

	domains, _, err := cmd.PluginActor.GetOrganizationDomains(org.GUID, "")
	if err != nil {
		return err
	}

	for _, space := range spaces {
		newSpace := plugin_models.Space{
			GUID: space.GUID,
			Name: space.Name,
		}
		retVal.Spaces = append(retVal.Spaces, newSpace)
	}

	for _, domain := range domains {
		newDomain := plugin_models.Domain{GUID: domain.GUID, Name: domain.Name}
		retVal.Domains = append(retVal.Domains, newDomain)
	}

	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	dec := gob.NewDecoder(&b)

	err = enc.Encode(org.Metadata)
	if err != nil {
		panic(err)
	}

	var metadata plugin_models.Metadata

	err = dec.Decode(&metadata)
	if err != nil {
		panic(err)
	}

	retVal.Name = org.Name
	retVal.GUID = org.GUID
	retVal.Metadata = metadata

	return nil
}

func (cmd *CliRpcCmd) GetSpace(spaceName string, retVal *plugin_models.Space) error {
	orgGUID := cmd.Config.TargetedOrganization().GUID
	if orgGUID == "" {
		return errors.New("no organization targeted")
	}
	space, _, err := cmd.PluginActor.GetSpaceByNameAndOrganization(spaceName, orgGUID)
	if err != nil {
		return err
	}

	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	dec := gob.NewDecoder(&b)

	err = enc.Encode(space)
	if err != nil {
		panic(err)
	}

	err = dec.Decode(retVal)
	if err != nil {
		panic(err)
	}

	return nil
}

func (cmd *CliRpcCmd) GetSpaces(not_used string, retVal *[]plugin_models.Space) error {
	orgGUID := cmd.Config.TargetedOrganization().GUID
	if orgGUID == "" {
		return errors.New("no organization targeted")
	}
	spaces, _, err := cmd.PluginActor.GetOrganizationSpaces(orgGUID)
	if err != nil {
		return err
	}

	var b bytes.Buffer
	enc := gob.NewEncoder(&b)
	dec := gob.NewDecoder(&b)

	err = enc.Encode(spaces)
	if err != nil {
		return err
	}

	err = dec.Decode(retVal)
	if err != nil {
		return err
	}

	return nil
}

func (cmd *CliRpcCmd) IsLoggedIn(args string, retVal *bool) error {
	*retVal = cmd.Config.AccessToken() != ""
	return nil
}

func (cmd *CliRpcCmd) IsSkipSSLValidation(args string, retVal *bool) error {
	*retVal = cmd.Config.SkipSSLValidation()
	return nil
}

func (cmd *CliRpcCmd) AccessToken(args string, retVal *string) error {
	accessToken, err := cmd.PluginActor.RefreshAccessToken()
	if err != nil {
		panic(err)
	}
	*retVal = accessToken
	return nil
}

func (cmd *CliRpcCmd) Username(args string, retVal *string) error {
	username, err := cmd.Config.CurrentUserName()
	if err != nil {
		return fmt.Errorf("error processing config: %s", err.Error())
	}
	if username == "" {
		return errors.New("not logged in")
	}
	*retVal = username

	return nil
}
<|endoftext|>


<commit_before>// Package db exposes a lightweight abstraction over the SQLite code.
// It performs some basic mapping of lower-level types to rqlite types.
package db

import (
	"database/sql/driver"
	"expvar"
	"fmt"
	"io"
	"time"

	"github.com/mattn/go-sqlite3"
)

const bkDelay = 250

const (
	fkChecks         = "PRAGMA foreign_keys"
	fkChecksEnabled  = "PRAGMA foreign_keys=ON"
	fkChecksDisabled = "PRAGMA foreign_keys=OFF"

	numExecutions   = "executions"
	numQueries      = "queries"
	numTransactions = "transactions"
)

// DBVersion is the SQLite version.
var DBVersion string

// stats captures stats for the DB layer.
var stats *expvar.Map

func init() {
	DBVersion, _, _ = sqlite3.Version()
	stats = expvar.NewMap("db")
	stats.Add(numExecutions, 0)
	stats.Add(numQueries, 0)
	stats.Add(numTransactions, 0)
}

// DB is the SQL database.
type DB struct {
	sqlite3conn *sqlite3.SQLiteConn // Driver connection to database.
	path        string              // Path to database file.
	dsn         string              // DSN, if any.
	memory      bool                // In-memory only.
}

// Result represents the outcome of an operation that changes rows.
type Result struct {
	LastInsertID int64   `json:"last_insert_id,omitempty"`
	RowsAffected int64   `json:"rows_affected,omitempty"`
	Error        string  `json:"error,omitempty"`
	Time         float64 `json:"time,omitempty"`
}

// Rows represents the outcome of an operation that returns query data.
type Rows struct {
	Columns []string        `json:"columns,omitempty"`
	Types   []string        `json:"types,omitempty"`
	Values  [][]interface{} `json:"values,omitempty"`
	Error   string          `json:"error,omitempty"`
	Time    float64         `json:"time,omitempty"`
}

// Open opens a file-based database, creating it if it does not exist.
func Open(dbPath string) (*DB, error) {
	return open(fqdsn(dbPath, ""))
}

// OpenWithDSN opens a file-based database, creating it if it does not exist.
func OpenWithDSN(dbPath, dsn string) (*DB, error) {
	return open(fqdsn(dbPath, dsn))
}

// OpenInMemory opens an in-memory database.
func OpenInMemory() (*DB, error) {
	return open(fqdsn(":memory:", ""))
}

// OpenInMemoryWithDSN opens an in-memory database with a specific DSN.
func OpenInMemoryWithDSN(dsn string) (*DB, error) {
	return open(fqdsn(":memory:", dsn))
}

// LoadInMemoryWithDSN loads an in-memory database with that at the path,
// with the specified DSN
func LoadInMemoryWithDSN(dbPath, dsn string) (*DB, error) {
	db, err := OpenInMemoryWithDSN(dsn)
	if err != nil {
		return nil, err
	}

	srcDB, err := Open(dbPath)
	if err != nil {
		return nil, err
	}

	bk, err := db.sqlite3conn.Backup("main", srcDB.sqlite3conn, "main")
	if err != nil {
		return nil, err
	}

	for {
		done, err := bk.Step(-1)
		if err != nil {
			bk.Finish()
			return nil, err
		}
		if done {
			break
		}
		time.Sleep(bkDelay * time.Millisecond)
	}

	if err := bk.Finish(); err != nil {
		return nil, err
	}
	if err := srcDB.Close(); err != nil {
		return nil, err
	}

	return db, nil
}

// Close closes the underlying database connection.
func (db *DB) Close() error {
	return db.sqlite3conn.Close()
}

func open(dbPath string) (*DB, error) {
	d := sqlite3.SQLiteDriver{}
	dbc, err := d.Open(dbPath)
	if err != nil {
		return nil, err
	}

	return &DB{
		sqlite3conn: dbc.(*sqlite3.SQLiteConn),
		path:        dbPath,
	}, nil
}

// EnableFKConstraints allows control of foreign key constraint checks.
func (db *DB) EnableFKConstraints(e bool) error {
	q := fkChecksEnabled
	if !e {
		q = fkChecksDisabled
	}
	_, err := db.sqlite3conn.Exec(q, nil)
	return err
}

// FKConstraints returns whether FK constraints are set or not.
func (db *DB) FKConstraints() (bool, error) {
	r, err := db.sqlite3conn.Query(fkChecks, nil)
	if err != nil {
		return false, err
	}

	dest := make([]driver.Value, len(r.Columns()))
	types := r.(*sqlite3.SQLiteRows).DeclTypes()
	if err := r.Next(dest); err != nil {
		return false, err
	}

	values := normalizeRowValues(dest, types)
	if values[0] == int64(1) {
		return true, nil
	}
	return false, nil
}

// Execute executes queries that modify the database.
func (db *DB) Execute(queries []string, tx, xTime bool) ([]*Result, error) {
	stats.Add(numQueries, 1)
	stats.Add(numExecutions, int64(len(queries)))

	type Execer interface {
		Exec(query string, args []driver.Value) (driver.Result, error)
	}

	var allResults []*Result
	err := func() error {
		var execer Execer
		var rollback bool
		var t driver.Tx
		var err error

		// Check for the err, if set rollback.
		defer func() {
			if t != nil {
				if rollback {
					t.Rollback()
					return
				}
				t.Commit()
			}
		}()

		// handleError sets the error field on the given result. It returns
		// whether the caller should continue processing or break.
		handleError := func(result *Result, err error) bool {
			result.Error = err.Error()
			allResults = append(allResults, result)
			if tx {
				rollback = true // Will trigger the rollback.
				return false
			}
			return true
		}

		execer = db.sqlite3conn

		// Create the correct execution object, depending on whether a
		// transaction was requested.
		if tx {
			t, err = db.sqlite3conn.Begin()
			if err != nil {
				return err
			}
		}

		// Execute each query.
		for _, q := range queries {
			if q == "" {
				continue
			}

			result := &Result{}
			start := time.Now()

			r, err := execer.Exec(q, nil)
			if err != nil {
				if handleError(result, err) {
					continue
				}
				break
			}

			lid, err := r.LastInsertId()
			if err != nil {
				if handleError(result, err) {
					continue
				}
				break
			}
			result.LastInsertID = lid

			ra, err := r.RowsAffected()
			if err != nil {
				if handleError(result, err) {
					continue
				}
				break
			}
			result.RowsAffected = ra
			if xTime {
				result.Time = time.Now().Sub(start).Seconds()
			}
			allResults = append(allResults, result)
		}

		return nil
	}()

	return allResults, err
}

// Query executes queries that return rows, but don't modify the database.
func (db *DB) Query(queries []string, tx, xTime bool) ([]*Rows, error) {
	stats.Add(numQueries, int64(len(queries)))

	type Queryer interface {
		Query(query string, args []driver.Value) (driver.Rows, error)
	}

	var allRows []*Rows
	err := func() (err error) {
		var queryer Queryer
		var t driver.Tx
		defer func() {
			// XXX THIS DOESN'T ACTUALLY WORK! Might as WELL JUST COMMIT?
			if t != nil {
				if err != nil {
					t.Rollback()
					return
				}
				t.Commit()
			}
		}()

		queryer = db.sqlite3conn

		// Create the correct query object, depending on whether a
		// transaction was requested.
		if tx {
			t, err = db.sqlite3conn.Begin()
			if err != nil {
				return err
			}
		}

		for _, q := range queries {
			if q == "" {
				continue
			}

			rows := &Rows{}
			start := time.Now()

			rs, err := queryer.Query(q, nil)
			if err != nil {
				rows.Error = err.Error()
				allRows = append(allRows, rows)
				continue
			}
			defer rs.Close()
			columns := rs.Columns()

			rows.Columns = columns
			rows.Types = rs.(*sqlite3.SQLiteRows).DeclTypes()
			dest := make([]driver.Value, len(rows.Columns))
			for {
				err := rs.Next(dest)
				if err != nil {
					if err != io.EOF {
						rows.Error = err.Error()
					}
					break
				}

				values := normalizeRowValues(dest, rows.Types)
				rows.Values = append(rows.Values, values)
			}
			if xTime {
				rows.Time = time.Now().Sub(start).Seconds()
			}
			allRows = append(allRows, rows)
		}

		return nil
	}()

	return allRows, err
}

// Backup writes a consistent snapshot of the database to the given file.
func (db *DB) Backup(path string) error {
	dstDB, err := Open(path)
	if err != nil {
		return err
	}

	bk, err := dstDB.sqlite3conn.Backup("main", db.sqlite3conn, "main")
	if err != nil {
		return err
	}

	for {
		done, err := bk.Step(-1)
		if err != nil {
			bk.Finish()
			return err
		}
		if done {
			break
		}
		time.Sleep(bkDelay * time.Millisecond)
	}

	if err := bk.Finish(); err != nil {
		return err
	}
	if err := dstDB.Close(); err != nil {
		return err
	}

	return nil
}

// normalizeRowValues performs some normalization of values in the returned rows.
// Text values come over (from sqlite-go) as []byte instead of strings
// for some reason, so we have explicitly convert (but only when type
// is "text" so we don't affect BLOB types)
func normalizeRowValues(row []driver.Value, types []string) []interface{} {
	values := make([]interface{}, len(types))
	for i, v := range row {
		if types[i] == "text" {
			switch val := v.(type) {
			case []byte:
				values[i] = string(val)
			default:
				values[i] = val
			}
		} else {
			values[i] = v
		}
	}
	return values
}

// fqdsn returns the fully-qualified datasource name.
func fqdsn(path, dsn string) string {
	if dsn != "" {
		return fmt.Sprintf("file:%s?%s", path, dsn)
	}
	return path
}
<commit_msg>More database-level stats<commit_after>// Package db exposes a lightweight abstraction over the SQLite code.
// It performs some basic mapping of lower-level types to rqlite types.
package db

import (
	"database/sql/driver"
	"expvar"
	"fmt"
	"io"
	"time"

	"github.com/mattn/go-sqlite3"
)

const bkDelay = 250

const (
	fkChecks         = "PRAGMA foreign_keys"
	fkChecksEnabled  = "PRAGMA foreign_keys=ON"
	fkChecksDisabled = "PRAGMA foreign_keys=OFF"

	numExecutions      = "executions"
	numExecutionErrors = "execution_errors"
	numQueries         = "queries"
	numETx             = "execute_transactions"
	numQTx             = "query_transactions"
)

// DBVersion is the SQLite version.
var DBVersion string

// stats captures stats for the DB layer.
var stats *expvar.Map

func init() {
	DBVersion, _, _ = sqlite3.Version()
	stats = expvar.NewMap("db")
	stats.Add(numExecutions, 0)
	stats.Add(numExecutionErrors, 0)
	stats.Add(numQueries, 0)
	stats.Add(numETx, 0)
	stats.Add(numQTx, 0)

}

// DB is the SQL database.
type DB struct {
	sqlite3conn *sqlite3.SQLiteConn // Driver connection to database.
	path        string              // Path to database file.
	dsn         string              // DSN, if any.
	memory      bool                // In-memory only.
}

// Result represents the outcome of an operation that changes rows.
type Result struct {
	LastInsertID int64   `json:"last_insert_id,omitempty"`
	RowsAffected int64   `json:"rows_affected,omitempty"`
	Error        string  `json:"error,omitempty"`
	Time         float64 `json:"time,omitempty"`
}

// Rows represents the outcome of an operation that returns query data.
type Rows struct {
	Columns []string        `json:"columns,omitempty"`
	Types   []string        `json:"types,omitempty"`
	Values  [][]interface{} `json:"values,omitempty"`
	Error   string          `json:"error,omitempty"`
	Time    float64         `json:"time,omitempty"`
}

// Open opens a file-based database, creating it if it does not exist.
func Open(dbPath string) (*DB, error) {
	return open(fqdsn(dbPath, ""))
}

// OpenWithDSN opens a file-based database, creating it if it does not exist.
func OpenWithDSN(dbPath, dsn string) (*DB, error) {
	return open(fqdsn(dbPath, dsn))
}

// OpenInMemory opens an in-memory database.
func OpenInMemory() (*DB, error) {
	return open(fqdsn(":memory:", ""))
}

// OpenInMemoryWithDSN opens an in-memory database with a specific DSN.
func OpenInMemoryWithDSN(dsn string) (*DB, error) {
	return open(fqdsn(":memory:", dsn))
}

// LoadInMemoryWithDSN loads an in-memory database with that at the path,
// with the specified DSN
func LoadInMemoryWithDSN(dbPath, dsn string) (*DB, error) {
	db, err := OpenInMemoryWithDSN(dsn)
	if err != nil {
		return nil, err
	}

	srcDB, err := Open(dbPath)
	if err != nil {
		return nil, err
	}

	bk, err := db.sqlite3conn.Backup("main", srcDB.sqlite3conn, "main")
	if err != nil {
		return nil, err
	}

	for {
		done, err := bk.Step(-1)
		if err != nil {
			bk.Finish()
			return nil, err
		}
		if done {
			break
		}
		time.Sleep(bkDelay * time.Millisecond)
	}

	if err := bk.Finish(); err != nil {
		return nil, err
	}
	if err := srcDB.Close(); err != nil {
		return nil, err
	}

	return db, nil
}

// Close closes the underlying database connection.
func (db *DB) Close() error {
	return db.sqlite3conn.Close()
}

func open(dbPath string) (*DB, error) {
	d := sqlite3.SQLiteDriver{}
	dbc, err := d.Open(dbPath)
	if err != nil {
		return nil, err
	}

	return &DB{
		sqlite3conn: dbc.(*sqlite3.SQLiteConn),
		path:        dbPath,
	}, nil
}

// EnableFKConstraints allows control of foreign key constraint checks.
func (db *DB) EnableFKConstraints(e bool) error {
	q := fkChecksEnabled
	if !e {
		q = fkChecksDisabled
	}
	_, err := db.sqlite3conn.Exec(q, nil)
	return err
}

// FKConstraints returns whether FK constraints are set or not.
func (db *DB) FKConstraints() (bool, error) {
	r, err := db.sqlite3conn.Query(fkChecks, nil)
	if err != nil {
		return false, err
	}

	dest := make([]driver.Value, len(r.Columns()))
	types := r.(*sqlite3.SQLiteRows).DeclTypes()
	if err := r.Next(dest); err != nil {
		return false, err
	}

	values := normalizeRowValues(dest, types)
	if values[0] == int64(1) {
		return true, nil
	}
	return false, nil
}

// Execute executes queries that modify the database.
func (db *DB) Execute(queries []string, tx, xTime bool) ([]*Result, error) {
	stats.Add(numExecutions, int64(len(queries)))
	if tx {
		stats.Add(numETx, 1)
	}

	type Execer interface {
		Exec(query string, args []driver.Value) (driver.Result, error)
	}

	var allResults []*Result
	err := func() error {
		var execer Execer
		var rollback bool
		var t driver.Tx
		var err error

		// Check for the err, if set rollback.
		defer func() {
			if t != nil {
				if rollback {
					t.Rollback()
					return
				}
				t.Commit()
			}
		}()

		// handleError sets the error field on the given result. It returns
		// whether the caller should continue processing or break.
		handleError := func(result *Result, err error) bool {
			stats.Add(numExecutionErrors, 1)

			result.Error = err.Error()
			allResults = append(allResults, result)
			if tx {
				rollback = true // Will trigger the rollback.
				return false
			}
			return true
		}

		execer = db.sqlite3conn

		// Create the correct execution object, depending on whether a
		// transaction was requested.
		if tx {
			t, err = db.sqlite3conn.Begin()
			if err != nil {
				return err
			}
		}

		// Execute each query.
		for _, q := range queries {
			if q == "" {
				continue
			}

			result := &Result{}
			start := time.Now()

			r, err := execer.Exec(q, nil)
			if err != nil {
				if handleError(result, err) {
					continue
				}
				break
			}

			lid, err := r.LastInsertId()
			if err != nil {
				if handleError(result, err) {
					continue
				}
				break
			}
			result.LastInsertID = lid

			ra, err := r.RowsAffected()
			if err != nil {
				if handleError(result, err) {
					continue
				}
				break
			}
			result.RowsAffected = ra
			if xTime {
				result.Time = time.Now().Sub(start).Seconds()
			}
			allResults = append(allResults, result)
		}

		return nil
	}()

	return allResults, err
}

// Query executes queries that return rows, but don't modify the database.
func (db *DB) Query(queries []string, tx, xTime bool) ([]*Rows, error) {
	stats.Add(numQueries, int64(len(queries)))
	if tx {
		stats.Add(numQTx, 1)
	}

	type Queryer interface {
		Query(query string, args []driver.Value) (driver.Rows, error)
	}

	var allRows []*Rows
	err := func() (err error) {
		var queryer Queryer
		var t driver.Tx
		defer func() {
			// XXX THIS DOESN'T ACTUALLY WORK! Might as WELL JUST COMMIT?
			if t != nil {
				if err != nil {
					t.Rollback()
					return
				}
				t.Commit()
			}
		}()

		queryer = db.sqlite3conn

		// Create the correct query object, depending on whether a
		// transaction was requested.
		if tx {
			t, err = db.sqlite3conn.Begin()
			if err != nil {
				return err
			}
		}

		for _, q := range queries {
			if q == "" {
				continue
			}

			rows := &Rows{}
			start := time.Now()

			rs, err := queryer.Query(q, nil)
			if err != nil {
				rows.Error = err.Error()
				allRows = append(allRows, rows)
				continue
			}
			defer rs.Close()
			columns := rs.Columns()

			rows.Columns = columns
			rows.Types = rs.(*sqlite3.SQLiteRows).DeclTypes()
			dest := make([]driver.Value, len(rows.Columns))
			for {
				err := rs.Next(dest)
				if err != nil {
					if err != io.EOF {
						rows.Error = err.Error()
					}
					break
				}

				values := normalizeRowValues(dest, rows.Types)
				rows.Values = append(rows.Values, values)
			}
			if xTime {
				rows.Time = time.Now().Sub(start).Seconds()
			}
			allRows = append(allRows, rows)
		}

		return nil
	}()

	return allRows, err
}

// Backup writes a consistent snapshot of the database to the given file.
func (db *DB) Backup(path string) error {
	dstDB, err := Open(path)
	if err != nil {
		return err
	}

	bk, err := dstDB.sqlite3conn.Backup("main", db.sqlite3conn, "main")
	if err != nil {
		return err
	}

	for {
		done, err := bk.Step(-1)
		if err != nil {
			bk.Finish()
			return err
		}
		if done {
			break
		}
		time.Sleep(bkDelay * time.Millisecond)
	}

	if err := bk.Finish(); err != nil {
		return err
	}
	if err := dstDB.Close(); err != nil {
		return err
	}

	return nil
}

// normalizeRowValues performs some normalization of values in the returned rows.
// Text values come over (from sqlite-go) as []byte instead of strings
// for some reason, so we have explicitly convert (but only when type
// is "text" so we don't affect BLOB types)
func normalizeRowValues(row []driver.Value, types []string) []interface{} {
	values := make([]interface{}, len(types))
	for i, v := range row {
		if types[i] == "text" {
			switch val := v.(type) {
			case []byte:
				values[i] = string(val)
			default:
				values[i] = val
			}
		} else {
			values[i] = v
		}
	}
	return values
}

// fqdsn returns the fully-qualified datasource name.
func fqdsn(path, dsn string) string {
	if dsn != "" {
		return fmt.Sprintf("file:%s?%s", path, dsn)
	}
	return path
}
<|endoftext|>