<commit_before>/*

For keeping a minimum running, perhaps when doing a routing table update, if destination hosts are all
 expired or about to expire we start more.

*/

package main

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	"github.com/iron-io/iron_go/cache"
	"github.com/iron-io/iron_go/worker"
	"github.com/iron-io/common"
	"log"
	"math/rand"
	// "net"
	"net/http"
	"net/url"
	"reflect"
	"strings"
	"time"
	"runtime"
	"flag"
//	"io/ioutil"
)

var config struct {
	Iron struct {
	Token      string `json:"token"`
	ProjectId  string `json:"project_id"`
} `json:"iron"`
	Logging       struct {
	To     string `json:"to"`
	Level  string `json:"level"`
	Prefix string `json:"prefix"`
}
}

//var routingTable = map[string]*Route{}
var icache = cache.New("routing-table")

func init() {

}

type Route struct {
	// TODO: Change destinations to a simple cache so it can expire entries after 55 minutes (the one we use in common?)
	Host         string `json:"host"`
	Destinations []string  `json:"destinations"`
	ProjectId    string  `json:"project_id"`
	Token        string  `json:"token"` // store this so we can queue up new workers on demand
	CodeName     string  `json:"code_name"`
}

// for adding new hosts
type Route2 struct {
	Host string `json:"host"`
	Dest string `json:"dest"`
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Println("Running on", runtime.NumCPU(), "CPUs")

	var configFile string
	var env string
	flag.StringVar(&configFile, "c", "", "Config file name")
	// when this was e, it was erroring out.
	flag.StringVar(&env, "e", "development", "environment")

	flag.Parse() // Scans the arg list and sets up flags

	// Deployer is now passing -c in since we're using upstart and it doesn't know what directory to run in
	if configFile == "" {
		configFile = "config_" + env + ".json"
	}

	common.LoadConfig("iron_mq", configFile, &config)
	common.SetLogLevel(config.Logging.Level)
	common.SetLogLocation(config.Logging.To, config.Logging.Prefix)

	icache.Settings.UseConfigMap(map[string]interface{}{"token": config.Iron.Token, "project_id": config.Iron.ProjectId})

	r := mux.NewRouter()
	s := r.Headers("Iron-Router", "").Subrouter()
	s.HandleFunc("/", AddWorker)
	r.HandleFunc("/addworker", AddWorker)
	r.HandleFunc("/ping", Ping) // for health

	r.HandleFunc("/", ProxyFunc)

	http.Handle("/", r)
	port := 80
	fmt.Println("listening and serving on port", port)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%v", port), nil))
}

func ProxyFunc(w http.ResponseWriter, req *http.Request) {
	fmt.Println("HOST:", req.Host)
	host := strings.Split(req.Host, ":")[0]

	// We look up the destinations in the routing table and there can be 3 possible scenarios:
	// 1) This host was never registered so we return 404
	// 2) This host has active workers so we do the proxy
	// 3) This host has no active workers so we queue one (or more) up and return a 503 or something with message that says "try again in a minute"
	//	route := routingTable[host]
 fmt.Println("getting route for host:", host)
	route, err := getRoute(host)
	// choose random dest
	if err != nil {
		common.SendError(w, 400, fmt.Sprintln("Host not registered or error!", err))
		return
	}
	//	if route == nil { // route.Host == "" {
	//		common.SendError(w, 400, fmt.Sprintln(w, "Host not configured!"))
	//		return
	//	}
	dlen := len(route.Destinations)
	if dlen == 0 {
		fmt.Println("No workers running, starting new task.")
		startNewWorker(route)
		common.SendError(w, 500, fmt.Sprintln("No workers running, starting them up..."))
		return
	}
	if dlen == 1 {
		fmt.Println("Only one worker running, starting a new task.")
		startNewWorker(route)
	}
	destIndex := rand.Intn(dlen)
	destUrlString := route.Destinations[destIndex]
	// todo: should check if http:// already exists.
	destUrlString2 := "http://" + destUrlString
	destUrl, err := url.Parse(destUrlString2)
	if err != nil {
		fmt.Println("error!", err)
		panic(err)
	}
	fmt.Println("proxying to", destUrl)
	proxy := NewSingleHostReverseProxy(destUrl)
	err = proxy.ServeHTTP(w, req)
	if err != nil {
		fmt.Println("Error proxying!", err)
		etype := reflect.TypeOf(err)
		fmt.Println("err type:", etype)
		w.WriteHeader(http.StatusInternalServerError)
		// can't figure out how to compare types so comparing strings.... lame.
		if strings.Contains(etype.String(), "net.OpError") { // == reflect.TypeOf(net.OpError{}) { // couldn't figure out a better way to do this
			if len(route.Destinations) > 2 { // always want at least two running
				fmt.Println("It's a network error, removing this destination from routing table.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				err := putRoute(route)
				if err != nil {
					fmt.Println("couldn't update routing table 1", err)
					common.SendError(w, 500, fmt.Sprintln("couldn't update routing table 1", err))
					return
				}
				fmt.Println("New route:", route)
				return
			} else {
				fmt.Println("It's a network error and less than two other workers available so we're going to remove it and start new task.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				err := putRoute(route)
				if err != nil {
					fmt.Println("couldn't update routing table:", err)
					common.SendError(w, 500, fmt.Sprintln("couldn't update routing table", err))
					return
				}
				fmt.Println("New route:", route)
			}
			// start new worker if it's a connection error
			startNewWorker(route)
		}
		return
	}
	fmt.Println("Served!")
	// todo: how to handle destination failures. I got this in log output when testing a bad endpoint:
	// 2012/12/26 23:22:08 http: proxy error: dial tcp 127.0.0.1:8082: connection refused
}

func startNewWorker(route *Route) (error) {
	fmt.Println("Starting a new worker")
	// start new worker
	payload := map[string]interface{}{
		"token":      route.Token,
		"project_id": route.ProjectId,
		"code_name":  route.CodeName,
	}
	workerapi := worker.New()
	workerapi.Settings.UseConfigMap(payload)
	jsonPayload, err := json.Marshal(payload)
	if err != nil {
		fmt.Println("Couldn't marshal json!", err)
		return err
	}
	timeout := time.Second * time.Duration(1800 + rand.Intn(600)) // a little random factor in here to spread out worker deaths
	task := worker.Task{
		CodeName: route.CodeName,
		Payload:  string(jsonPayload),
		Timeout:  &timeout, // let's have these die quickly while testing
	}
	tasks := make([]worker.Task, 1)
	tasks[0] = task
	taskIds, err := workerapi.TaskQueue(tasks...)
	fmt.Println("Tasks queued.", taskIds)
	if err != nil {
		fmt.Println("Couldn't queue up worker!", err)
		return err
	}
	return err
}

// When a worker starts up, it calls this
func AddWorker(w http.ResponseWriter, req *http.Request) {
	log.Println("AddWorker called!")

//	s, err := ioutil.ReadAll(req.Body)
//	fmt.Println("req.body:", err, string(s))

	// get project id and token
	projectId := req.FormValue("project_id")
	token := req.FormValue("token")
	codeName := req.FormValue("code_name")
	fmt.Println("project_id:", projectId, "token:", token, "code_name:", codeName)

	// check header for what operation to perform
	routerHeader := req.Header.Get("Iron-Router")
	if routerHeader == "register" {
		route := Route{}
		if !common.ReadJSON(w, req, &route) {
			return
		}
		fmt.Println("body read into route:", route)
		route.ProjectId = projectId
		route.Token = token
		route.CodeName = codeName
		// todo: do we need to close body?
		err := putRoute(&route)
		if err != nil {
			fmt.Println("couldn't register host:", err)
			common.SendError(w, 400, fmt.Sprintln("Could not register host!", err))
			return
		}
		fmt.Println("registered route:", route)
		fmt.Fprintln(w, "Host registered successfully.")

	} else {
		r2 := Route2{}
		if !common.ReadJSON(w, req, &r2) {
			return
		}
		// todo: do we need to close body?
		fmt.Println("DECODED:", r2)
		route, err := getRoute(r2.Host)
		//		route := routingTable[r2.Host]
		if err != nil {
			common.SendError(w, 400, fmt.Sprintln("This host is not registered!", err))
			return
			//			route = &Route{}
		}
		fmt.Println("ROUTE:", route)
		route.Destinations = append(route.Destinations, r2.Dest)
		fmt.Println("ROUTE new:", route)
		err = putRoute(route)
		if err != nil {
			fmt.Println("couldn't register host:", err)
			common.SendError(w, 400, fmt.Sprintln("Could not register host!", err))
			return
		}
		//		routingTable[r2.Host] = route
		//		fmt.Println("New routing table:", routingTable)
		fmt.Fprintln(w, "Worker added")
	}
}

func getRoute(host string) (*Route, error) {
	rx, err := icache.Get(host)
	if err != nil {
		return nil, err
	}
	rx2 := []byte(rx.(string))
	route := Route{}
	err = json.Unmarshal(rx2, &route)
	if err != nil {
		return nil, err
	}
	return &route, err
}

func putRoute(route *Route) (error) {
	item := cache.Item{}
	v, err := json.Marshal(route)
	if err != nil {
		return err
	}
	item.Value = string(v)
	err = icache.Put(route.Host, &item)
	return err
}

func Ping(w http.ResponseWriter, req *http.Request) {
	fmt.Fprintln(w, "pong")
}
<commit_msg>Made 3 the minimum number of workers.<commit_after>/*

For keeping a minimum running, perhaps when doing a routing table update, if destination hosts are all
 expired or about to expire we start more.

*/

package main

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	"github.com/iron-io/iron_go/cache"
	"github.com/iron-io/iron_go/worker"
	"github.com/iron-io/common"
	"log"
	"math/rand"
	// "net"
	"net/http"
	"net/url"
	"reflect"
	"strings"
	"time"
	"runtime"
	"flag"
//	"io/ioutil"
)

var config struct {
	Iron struct {
	Token      string `json:"token"`
	ProjectId  string `json:"project_id"`
} `json:"iron"`
	Logging       struct {
	To     string `json:"to"`
	Level  string `json:"level"`
	Prefix string `json:"prefix"`
}
}

//var routingTable = map[string]*Route{}
var icache = cache.New("routing-table")

func init() {

}

type Route struct {
	// TODO: Change destinations to a simple cache so it can expire entries after 55 minutes (the one we use in common?)
	Host         string `json:"host"`
	Destinations []string  `json:"destinations"`
	ProjectId    string  `json:"project_id"`
	Token        string  `json:"token"` // store this so we can queue up new workers on demand
	CodeName     string  `json:"code_name"`
}

// for adding new hosts
type Route2 struct {
	Host string `json:"host"`
	Dest string `json:"dest"`
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Println("Running on", runtime.NumCPU(), "CPUs")

	var configFile string
	var env string
	flag.StringVar(&configFile, "c", "", "Config file name")
	// when this was e, it was erroring out.
	flag.StringVar(&env, "e", "development", "environment")

	flag.Parse() // Scans the arg list and sets up flags

	// Deployer is now passing -c in since we're using upstart and it doesn't know what directory to run in
	if configFile == "" {
		configFile = "config_" + env + ".json"
	}

	common.LoadConfig("iron_mq", configFile, &config)
	common.SetLogLevel(config.Logging.Level)
	common.SetLogLocation(config.Logging.To, config.Logging.Prefix)

	icache.Settings.UseConfigMap(map[string]interface{}{"token": config.Iron.Token, "project_id": config.Iron.ProjectId})

	r := mux.NewRouter()
	s := r.Headers("Iron-Router", "").Subrouter()
	s.HandleFunc("/", AddWorker)
	r.HandleFunc("/addworker", AddWorker)
	r.HandleFunc("/ping", Ping) // for health

	r.HandleFunc("/", ProxyFunc)

	http.Handle("/", r)
	port := 80
	fmt.Println("listening and serving on port", port)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%v", port), nil))
}

func ProxyFunc(w http.ResponseWriter, req *http.Request) {
	fmt.Println("HOST:", req.Host)
	host := strings.Split(req.Host, ":")[0]

	// We look up the destinations in the routing table and there can be 3 possible scenarios:
	// 1) This host was never registered so we return 404
	// 2) This host has active workers so we do the proxy
	// 3) This host has no active workers so we queue one (or more) up and return a 503 or something with message that says "try again in a minute"
	//	route := routingTable[host]
 fmt.Println("getting route for host:", host)
	route, err := getRoute(host)
	// choose random dest
	if err != nil {
		common.SendError(w, 400, fmt.Sprintln("Host not registered or error!", err))
		return
	}
	//	if route == nil { // route.Host == "" {
	//		common.SendError(w, 400, fmt.Sprintln(w, "Host not configured!"))
	//		return
	//	}
	dlen := len(route.Destinations)
	if dlen == 0 {
		fmt.Println("No workers running, starting new task.")
		startNewWorker(route)
		common.SendError(w, 500, fmt.Sprintln("No workers running, starting them up..."))
		return
	}
	if dlen < 3 {
		fmt.Println("Only one worker running, starting a new task.")
		startNewWorker(route)
	}
	destIndex := rand.Intn(dlen)
	destUrlString := route.Destinations[destIndex]
	// todo: should check if http:// already exists.
	destUrlString2 := "http://" + destUrlString
	destUrl, err := url.Parse(destUrlString2)
	if err != nil {
		fmt.Println("error!", err)
		panic(err)
	}
	fmt.Println("proxying to", destUrl)
	proxy := NewSingleHostReverseProxy(destUrl)
	err = proxy.ServeHTTP(w, req)
	if err != nil {
		fmt.Println("Error proxying!", err)
		etype := reflect.TypeOf(err)
		fmt.Println("err type:", etype)
		w.WriteHeader(http.StatusInternalServerError)
		// can't figure out how to compare types so comparing strings.... lame.
		if strings.Contains(etype.String(), "net.OpError") { // == reflect.TypeOf(net.OpError{}) { // couldn't figure out a better way to do this
			if len(route.Destinations) > 3 { // always want at least two running
				fmt.Println("It's a network error, removing this destination from routing table.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				err := putRoute(route)
				if err != nil {
					fmt.Println("couldn't update routing table 1", err)
					common.SendError(w, 500, fmt.Sprintln("couldn't update routing table 1", err))
					return
				}
				fmt.Println("New route:", route)
				return
			} else {
				fmt.Println("It's a network error and less than two other workers available so we're going to remove it and start new task.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				err := putRoute(route)
				if err != nil {
					fmt.Println("couldn't update routing table:", err)
					common.SendError(w, 500, fmt.Sprintln("couldn't update routing table", err))
					return
				}
				fmt.Println("New route:", route)
			}
			// start new worker if it's a connection error
			startNewWorker(route)
		}
		return
	}
	fmt.Println("Served!")
	// todo: how to handle destination failures. I got this in log output when testing a bad endpoint:
	// 2012/12/26 23:22:08 http: proxy error: dial tcp 127.0.0.1:8082: connection refused
}

func startNewWorker(route *Route) (error) {
	fmt.Println("Starting a new worker")
	// start new worker
	payload := map[string]interface{}{
		"token":      route.Token,
		"project_id": route.ProjectId,
		"code_name":  route.CodeName,
	}
	workerapi := worker.New()
	workerapi.Settings.UseConfigMap(payload)
	jsonPayload, err := json.Marshal(payload)
	if err != nil {
		fmt.Println("Couldn't marshal json!", err)
		return err
	}
	timeout := time.Second * time.Duration(1800 + rand.Intn(600)) // a little random factor in here to spread out worker deaths
	task := worker.Task{
		CodeName: route.CodeName,
		Payload:  string(jsonPayload),
		Timeout:  &timeout, // let's have these die quickly while testing
	}
	tasks := make([]worker.Task, 1)
	tasks[0] = task
	taskIds, err := workerapi.TaskQueue(tasks...)
	fmt.Println("Tasks queued.", taskIds)
	if err != nil {
		fmt.Println("Couldn't queue up worker!", err)
		return err
	}
	return err
}

// When a worker starts up, it calls this
func AddWorker(w http.ResponseWriter, req *http.Request) {
	log.Println("AddWorker called!")

//	s, err := ioutil.ReadAll(req.Body)
//	fmt.Println("req.body:", err, string(s))

	// get project id and token
	projectId := req.FormValue("project_id")
	token := req.FormValue("token")
	codeName := req.FormValue("code_name")
	fmt.Println("project_id:", projectId, "token:", token, "code_name:", codeName)

	// check header for what operation to perform
	routerHeader := req.Header.Get("Iron-Router")
	if routerHeader == "register" {
		route := Route{}
		if !common.ReadJSON(w, req, &route) {
			return
		}
		fmt.Println("body read into route:", route)
		route.ProjectId = projectId
		route.Token = token
		route.CodeName = codeName
		// todo: do we need to close body?
		err := putRoute(&route)
		if err != nil {
			fmt.Println("couldn't register host:", err)
			common.SendError(w, 400, fmt.Sprintln("Could not register host!", err))
			return
		}
		fmt.Println("registered route:", route)
		fmt.Fprintln(w, "Host registered successfully.")

	} else {
		r2 := Route2{}
		if !common.ReadJSON(w, req, &r2) {
			return
		}
		// todo: do we need to close body?
		fmt.Println("DECODED:", r2)
		route, err := getRoute(r2.Host)
		//		route := routingTable[r2.Host]
		if err != nil {
			common.SendError(w, 400, fmt.Sprintln("This host is not registered!", err))
			return
			//			route = &Route{}
		}
		fmt.Println("ROUTE:", route)
		route.Destinations = append(route.Destinations, r2.Dest)
		fmt.Println("ROUTE new:", route)
		err = putRoute(route)
		if err != nil {
			fmt.Println("couldn't register host:", err)
			common.SendError(w, 400, fmt.Sprintln("Could not register host!", err))
			return
		}
		//		routingTable[r2.Host] = route
		//		fmt.Println("New routing table:", routingTable)
		fmt.Fprintln(w, "Worker added")
	}
}

func getRoute(host string) (*Route, error) {
	rx, err := icache.Get(host)
	if err != nil {
		return nil, err
	}
	rx2 := []byte(rx.(string))
	route := Route{}
	err = json.Unmarshal(rx2, &route)
	if err != nil {
		return nil, err
	}
	return &route, err
}

func putRoute(route *Route) (error) {
	item := cache.Item{}
	v, err := json.Marshal(route)
	if err != nil {
		return err
	}
	item.Value = string(v)
	err = icache.Put(route.Host, &item)
	return err
}

func Ping(w http.ResponseWriter, req *http.Request) {
	fmt.Fprintln(w, "pong")
}
<|endoftext|>


<commit_before>/*

For keeping a minimum running, perhaps when doing a routing table update, if destination hosts are all
 expired or about to expire we start more.

*/

package main

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	"github.com/iron-io/iron_go/cache"
	"github.com/iron-io/iron_go/worker"
	"github.com/iron-io/common"
	"log"
	"math/rand"
	// "net"
	"net/http"
	"net/url"
	"reflect"
	"strings"
	"time"
	"runtime"
	"flag"
	"io/ioutil"
)

var config struct {
	Iron struct {
	Token      string `json:"token"`
	ProjectId  string `json:"project_id"`
} `json:"iron"`
	Logging       struct {
	To     string `json:"to"`
	Level  string `json:"level"`
	Prefix string `json:"prefix"`
}
}

//var routingTable = map[string]*Route{}
var icache = cache.New("routing-table")

func init() {

}

type Route struct {
	// TODO: Change destinations to a simple cache so it can expire entries after 55 minutes (the one we use in common?)
	Host         string `json:"host"`
	Destinations []string  `json:"destinations"`
	ProjectId    string  `json:"project_id"`
	Token        string  `json:"token"` // store this so we can queue up new workers on demand
	CodeName     string  `json:"code_name"`
}

// for adding new hosts
type Route2 struct {
	Host string `json:"host"`
	Dest string `json:"dest"`
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Println("Running on", runtime.NumCPU(), "CPUs")

	var configFile string
	var env string
	flag.StringVar(&configFile, "c", "", "Config file name")
	// when this was e, it was erroring out.
	flag.StringVar(&env, "e2", "development", "environment")

	flag.Parse() // Scans the arg list and sets up flags

	// Deployer is now passing -c in since we're using upstart and it doesn't know what directory to run in
	if configFile == "" {
		configFile = "config_" + env + ".json"
	}

	common.LoadConfig("iron_mq", configFile, &config)
	common.SetLogLevel(config.Logging.Level)
	common.SetLogLocation(config.Logging.To, config.Logging.Prefix)

	icache.Settings.UseConfigMap(map[string]interface{}{"token": config.Iron.Token, "project_id": config.Iron.ProjectId})

	r := mux.NewRouter()
	s := r.Headers("Iron-Router", "").Subrouter()
	s.HandleFunc("/", AddWorker)
	r.HandleFunc("/addworker", AddWorker)

	r.HandleFunc("/", ProxyFunc)

	http.Handle("/", r)
	port := 80
	fmt.Println("listening and serving on port", port)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%v", port), nil))
}

func ProxyFunc(w http.ResponseWriter, req *http.Request) {
	fmt.Println("HOST:", req.Host)
	host := strings.Split(req.Host, ":")[0]

	// We look up the destinations in the routing table and there can be 3 possible scenarios:
	// 1) This host was never registered so we return 404
	// 2) This host has active workers so we do the proxy
	// 3) This host has no active workers so we queue one (or more) up and return a 503 or something with message that says "try again in a minute"
	//	route := routingTable[host]
	route, err := getRoute(host)
	// choose random dest
	if err != nil {
		common.SendError(w, 400, fmt.Sprintln(w, "Host not configured or error!", err))
		return
	}
	//	if route == nil { // route.Host == "" {
	//		common.SendError(w, 400, fmt.Sprintln(w, "Host not configured!"))
	//		return
	//	}
	destIndex := rand.Intn(len(route.Destinations))
	destUrlString := route.Destinations[destIndex]
	// todo: should check if http:// already exists.
	destUrlString2 := "http://" + destUrlString
	destUrl, err := url.Parse(destUrlString2)
	if err != nil {
		fmt.Println("error!", err)
		panic(err)
	}
	fmt.Println("proxying to", destUrl)
	proxy := NewSingleHostReverseProxy(destUrl)
	err = proxy.ServeHTTP(w, req)
	if err != nil {
		fmt.Println("Error proxying!", err)
		etype := reflect.TypeOf(err)
		fmt.Println("err type:", etype)
		w.WriteHeader(http.StatusInternalServerError)
		// can't figure out how to compare types so comparing strings.... lame.
		if strings.Contains(etype.String(), "net.OpError") { // == reflect.TypeOf(net.OpError{}) { // couldn't figure out a better way to do this
			if len(route.Destinations) > 1 {
				fmt.Println("It's a network error, removing this destination from routing table.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				putRoute(route)
				fmt.Println("New route:", route)
				return
			} else {
				fmt.Println("It's a network error and no other destinations available so we're going to remove it and start new task.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				putRoute(route)
				fmt.Println("New route:", route)
			}
			// start new worker
			payload := map[string]interface{}{
				"token":      route.Token,
				"project_id": route.ProjectId,
				"code_name":  route.CodeName,
			}
			workerapi := worker.New()
			workerapi.Settings.UseConfigMap(payload)
			jsonPayload, err := json.Marshal(payload)
			if err != nil {
				fmt.Println("Couldn't marshal json!", err)
				return
			}
			timeout := time.Second*120
			task := worker.Task{
				CodeName: route.CodeName,
				Payload:  string(jsonPayload),
				Timeout:  &timeout, // let's have these die quickly while testing
			}
			tasks := make([]worker.Task, 1)
			tasks[0] = task
			taskIds, err := workerapi.TaskQueue(tasks...)
			fmt.Println("Tasks queued.", taskIds)
			if err != nil {
				fmt.Println("Couldn't queue up worker!", err)
				return
			}
		}
		// start new worker if it's a connection error
		return
	}
	fmt.Println("Served!")
	// todo: how to handle destination failures. I got this in log output when testing a bad endpoint:
	// 2012/12/26 23:22:08 http: proxy error: dial tcp 127.0.0.1:8082: connection refused
}

// When a worker starts up, it calls this
func AddWorker(w http.ResponseWriter, req *http.Request) {
	log.Println("AddWorker called!")

	s, err := ioutil.ReadAll(req.Body)
	fmt.Println("req.body:", err, s)

	// get project id and token
	projectId := req.FormValue("project_id")
	token := req.FormValue("token")
	codeName := req.FormValue("code_name")
	fmt.Println("project_id:", projectId, "token:", token, "code_name:", codeName)

	// check header for what operation to perform
	routerHeader := req.Header.Get("Iron-Router")
	if routerHeader == "register" {
		route := Route{}
		decoder := json.NewDecoder(req.Body)
		decoder.Decode(&route)
		route.ProjectId = projectId
		route.Token = token
		route.CodeName = codeName
		// todo: do we need to close body?
		fmt.Println("registered route:", route)

		putRoute(route)

	} else {
		r2 := Route2{}
		decoder := json.NewDecoder(req.Body)
		decoder.Decode(&r2)
		// todo: do we need to close body?
		fmt.Println("DECODED:", r2)

		route, err := getRoute(r2.Host)
		//		route := routingTable[r2.Host]
		if err != nil {
			common.SendError(w, 400, fmt.Sprintln(w, "This host is not registered!", err))
			return
			//			route = &Route{}
		}
		fmt.Println("ROUTE:", route)
		route.Destinations = append(route.Destinations, r2.Dest)
		fmt.Println("ROUTE new:", route)

		putRoute(route)
		//		routingTable[r2.Host] = route
		//		fmt.Println("New routing table:", routingTable)
		fmt.Fprintln(w, "Worker added")
	}
}

func getRoute(host string) (Route, error) {
	rx, err := icache.Get(host)
	route := Route{}
	if err == nil {
		route = rx.(Route)
	}
	return route, err
}

func putRoute(route Route) {
	item := cache.Item{}
	item.Value = route
	icache.Put(route.Host, &item)
}
<commit_msg>printing body as string<commit_after>/*

For keeping a minimum running, perhaps when doing a routing table update, if destination hosts are all
 expired or about to expire we start more.

*/

package main

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	"github.com/iron-io/iron_go/cache"
	"github.com/iron-io/iron_go/worker"
	"github.com/iron-io/common"
	"log"
	"math/rand"
	// "net"
	"net/http"
	"net/url"
	"reflect"
	"strings"
	"time"
	"runtime"
	"flag"
	"io/ioutil"
)

var config struct {
	Iron struct {
	Token      string `json:"token"`
	ProjectId  string `json:"project_id"`
} `json:"iron"`
	Logging       struct {
	To     string `json:"to"`
	Level  string `json:"level"`
	Prefix string `json:"prefix"`
}
}

//var routingTable = map[string]*Route{}
var icache = cache.New("routing-table")

func init() {

}

type Route struct {
	// TODO: Change destinations to a simple cache so it can expire entries after 55 minutes (the one we use in common?)
	Host         string `json:"host"`
	Destinations []string  `json:"destinations"`
	ProjectId    string  `json:"project_id"`
	Token        string  `json:"token"` // store this so we can queue up new workers on demand
	CodeName     string  `json:"code_name"`
}

// for adding new hosts
type Route2 struct {
	Host string `json:"host"`
	Dest string `json:"dest"`
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Println("Running on", runtime.NumCPU(), "CPUs")

	var configFile string
	var env string
	flag.StringVar(&configFile, "c", "", "Config file name")
	// when this was e, it was erroring out.
	flag.StringVar(&env, "e2", "development", "environment")

	flag.Parse() // Scans the arg list and sets up flags

	// Deployer is now passing -c in since we're using upstart and it doesn't know what directory to run in
	if configFile == "" {
		configFile = "config_" + env + ".json"
	}

	common.LoadConfig("iron_mq", configFile, &config)
	common.SetLogLevel(config.Logging.Level)
	common.SetLogLocation(config.Logging.To, config.Logging.Prefix)

	icache.Settings.UseConfigMap(map[string]interface{}{"token": config.Iron.Token, "project_id": config.Iron.ProjectId})

	r := mux.NewRouter()
	s := r.Headers("Iron-Router", "").Subrouter()
	s.HandleFunc("/", AddWorker)
	r.HandleFunc("/addworker", AddWorker)

	r.HandleFunc("/", ProxyFunc)

	http.Handle("/", r)
	port := 80
	fmt.Println("listening and serving on port", port)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%v", port), nil))
}

func ProxyFunc(w http.ResponseWriter, req *http.Request) {
	fmt.Println("HOST:", req.Host)
	host := strings.Split(req.Host, ":")[0]

	// We look up the destinations in the routing table and there can be 3 possible scenarios:
	// 1) This host was never registered so we return 404
	// 2) This host has active workers so we do the proxy
	// 3) This host has no active workers so we queue one (or more) up and return a 503 or something with message that says "try again in a minute"
	//	route := routingTable[host]
	route, err := getRoute(host)
	// choose random dest
	if err != nil {
		common.SendError(w, 400, fmt.Sprintln(w, "Host not configured or error!", err))
		return
	}
	//	if route == nil { // route.Host == "" {
	//		common.SendError(w, 400, fmt.Sprintln(w, "Host not configured!"))
	//		return
	//	}
	destIndex := rand.Intn(len(route.Destinations))
	destUrlString := route.Destinations[destIndex]
	// todo: should check if http:// already exists.
	destUrlString2 := "http://" + destUrlString
	destUrl, err := url.Parse(destUrlString2)
	if err != nil {
		fmt.Println("error!", err)
		panic(err)
	}
	fmt.Println("proxying to", destUrl)
	proxy := NewSingleHostReverseProxy(destUrl)
	err = proxy.ServeHTTP(w, req)
	if err != nil {
		fmt.Println("Error proxying!", err)
		etype := reflect.TypeOf(err)
		fmt.Println("err type:", etype)
		w.WriteHeader(http.StatusInternalServerError)
		// can't figure out how to compare types so comparing strings.... lame.
		if strings.Contains(etype.String(), "net.OpError") { // == reflect.TypeOf(net.OpError{}) { // couldn't figure out a better way to do this
			if len(route.Destinations) > 1 {
				fmt.Println("It's a network error, removing this destination from routing table.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				putRoute(route)
				fmt.Println("New route:", route)
				return
			} else {
				fmt.Println("It's a network error and no other destinations available so we're going to remove it and start new task.")
				route.Destinations = append(route.Destinations[:destIndex], route.Destinations[destIndex + 1:]...)
				putRoute(route)
				fmt.Println("New route:", route)
			}
			// start new worker
			payload := map[string]interface{}{
				"token":      route.Token,
				"project_id": route.ProjectId,
				"code_name":  route.CodeName,
			}
			workerapi := worker.New()
			workerapi.Settings.UseConfigMap(payload)
			jsonPayload, err := json.Marshal(payload)
			if err != nil {
				fmt.Println("Couldn't marshal json!", err)
				return
			}
			timeout := time.Second*120
			task := worker.Task{
				CodeName: route.CodeName,
				Payload:  string(jsonPayload),
				Timeout:  &timeout, // let's have these die quickly while testing
			}
			tasks := make([]worker.Task, 1)
			tasks[0] = task
			taskIds, err := workerapi.TaskQueue(tasks...)
			fmt.Println("Tasks queued.", taskIds)
			if err != nil {
				fmt.Println("Couldn't queue up worker!", err)
				return
			}
		}
		// start new worker if it's a connection error
		return
	}
	fmt.Println("Served!")
	// todo: how to handle destination failures. I got this in log output when testing a bad endpoint:
	// 2012/12/26 23:22:08 http: proxy error: dial tcp 127.0.0.1:8082: connection refused
}

// When a worker starts up, it calls this
func AddWorker(w http.ResponseWriter, req *http.Request) {
	log.Println("AddWorker called!")

	s, err := ioutil.ReadAll(req.Body)
	fmt.Println("req.body:", err, string(s))

	// get project id and token
	projectId := req.FormValue("project_id")
	token := req.FormValue("token")
	codeName := req.FormValue("code_name")
	fmt.Println("project_id:", projectId, "token:", token, "code_name:", codeName)

	// check header for what operation to perform
	routerHeader := req.Header.Get("Iron-Router")
	if routerHeader == "register" {
		route := Route{}
		decoder := json.NewDecoder(req.Body)
		decoder.Decode(&route)
		route.ProjectId = projectId
		route.Token = token
		route.CodeName = codeName
		// todo: do we need to close body?
		fmt.Println("registered route:", route)
		putRoute(route)

	} else {
		r2 := Route2{}
		decoder := json.NewDecoder(req.Body)
		decoder.Decode(&r2)
		// todo: do we need to close body?
		fmt.Println("DECODED:", r2)
		route, err := getRoute(r2.Host)
		//		route := routingTable[r2.Host]
		if err != nil {
			common.SendError(w, 400, fmt.Sprintln(w, "This host is not registered!", err))
			return
			//			route = &Route{}
		}
		fmt.Println("ROUTE:", route)
		route.Destinations = append(route.Destinations, r2.Dest)
		fmt.Println("ROUTE new:", route)
		putRoute(route)
		//		routingTable[r2.Host] = route
		//		fmt.Println("New routing table:", routingTable)
		fmt.Fprintln(w, "Worker added")
	}
}

func getRoute(host string) (Route, error) {
	rx, err := icache.Get(host)
	route := Route{}
	if err == nil {
		route = rx.(Route)
	}
	return route, err
}

func putRoute(route Route) {
	item := cache.Item{}
	item.Value = route
	icache.Put(route.Host, &item)
}
<|endoftext|>


<commit_before>package encodes

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"net"
	"time"
)

const (
	IVLength = 1400
)

type AESConn struct {
	conn net.Conn
	r    *cipher.StreamReader
	w    *cipher.StreamWriter
}

func (ac *AESConn) Read(dst []byte) (n int, err error) {
	n, err = ac.r.Read(dst)
	return n, err
}

func (ac *AESConn) Write(src []byte) (n int, err error) {
	n, err = ac.w.Write(src)
	return n, err
}

func (ac *AESConn) Close() error {
	return ac.conn.Close()
}

func (ac *AESConn) LocalAddr() net.Addr {
	return ac.conn.LocalAddr()
}

func (ac *AESConn) RemoteAddr() net.Addr {
	return ac.conn.RemoteAddr()
}

func (ac *AESConn) SetDeadline(t time.Time) error {
	return ac.conn.SetDeadline(t)
}

func (ac *AESConn) SetReadDeadline(t time.Time) error {
	return ac.conn.SetReadDeadline(t)
}

func (ac *AESConn) SetWriteDeadline(t time.Time) error {
	return ac.conn.SetWriteDeadline(t)
}

func NewAESConn(key string, iv [aes.BlockSize]byte, conn net.Conn) (aesconn *AESConn, err error) {
	keybin := []byte(key)
	t := time.Now().UTC()

	hashkey := sha256.Sum256(keybin)
	salt := sha256.Sum256(
		[]byte(fmt.Sprintf("%d-%02d-%02d", t.Year(), t.Month(), t.Day())))

	hashsalt := append(hashkey[:], salt[:]...)

	aeskey := sha256.Sum256(hashsalt)

	rblock, err := aes.NewCipher(aeskey[:])
	wblock, err := aes.NewCipher(aeskey[:])
	if err != nil {
		return nil, err
	}

	rstream := cipher.NewOFB(rblock, iv[:])
	wstream := cipher.NewOFB(wblock, iv[:])

	aesconn = &AESConn{
		conn: conn,
		r:    &cipher.StreamReader{S: rstream, R: conn},
		w:    &cipher.StreamWriter{S: wstream, W: conn},
	}
	return aesconn, err
}

type Listener struct {
	netl net.Listener
	key  string
}

func (l *Listener) Accept() (aesconn net.Conn, err error) {
	conn, err := l.netl.Accept()
	if err != nil {
		return nil, err
	}

	conn.SetReadDeadline(time.Now().Add(60 * time.Second))

	ivvector := make([]byte, IVLength)
	var ri = 0
	for ri < IVLength {
		r, err := conn.Read(ivvector[ri:])
		if err != nil {
			return nil, err
		}
		ri += r
	}

	passwd := sha256.Sum256([]byte(l.key))
	longiv := append(passwd[:], ivvector...)

	hash := sha256.Sum256(longiv)

	var iv [aes.BlockSize]byte
	copy(iv[:], hash[:aes.BlockSize])

	aesconn, err = NewAESConn(l.key, iv, conn)

	return aesconn, err
}

func (l *Listener) Close() error {
	return l.netl.Close()
}

func (l *Listener) Addr() net.Addr {
	return l.netl.Addr()
}

func Listen(laddr string, key string) (l *Listener, err error) {
	netl, err := net.Listen("tcp", laddr)
	if err != nil {
		return nil, err
	}

	l = &Listener{
		netl: netl,
		key:  key,
	}
	return l, nil
}

func Dial(address string, key string) (aesconn net.Conn, err error) {
	conn, err := net.Dial("tcp", address)
	if err != nil {
		return nil, err
	}

	ivvector := make([]byte, IVLength)
	_, err = rand.Read(ivvector[:])

	if err != nil {
		return nil, err
	}

	var wi = 0
	for wi < IVLength {
		r, err := conn.Write(ivvector[wi:])
		if err != nil {
			return nil, err
		}
		wi += r
	}

	passwd := sha256.Sum256([]byte(key))
	longiv := append(passwd[:], ivvector...)

	hash := sha256.Sum256(longiv)

	var iv [aes.BlockSize]byte
	copy(iv[:], hash[:aes.BlockSize])

	aesconn, err = NewAESConn(key, iv, conn)
	return aesconn, err
}
<commit_msg>clear code<commit_after>package encodes

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"net"
	"time"
)

const (
	ivLength = 3200
)

type AESConn struct {
	conn net.Conn
	r    *cipher.StreamReader
	w    *cipher.StreamWriter
}

func (ac *AESConn) Read(dst []byte) (int, error) {
	n, err := ac.r.Read(dst)
	return n, err
}

func (ac *AESConn) Write(src []byte) (int, error) {
	n, err := ac.w.Write(src)
	return n, err
}

func (ac *AESConn) Close() error {
	return ac.conn.Close()
}

func (ac *AESConn) LocalAddr() net.Addr {
	return ac.conn.LocalAddr()
}

func (ac *AESConn) RemoteAddr() net.Addr {
	return ac.conn.RemoteAddr()
}

func (ac *AESConn) SetDeadline(t time.Time) error {
	return ac.conn.SetDeadline(t)
}

func (ac *AESConn) SetReadDeadline(t time.Time) error {
	return ac.conn.SetReadDeadline(t)
}

func (ac *AESConn) SetWriteDeadline(t time.Time) error {
	return ac.conn.SetWriteDeadline(t)
}

func NewAESConn(key string, iv [aes.BlockSize]byte, conn net.Conn) (*AESConn, error) {
	keybin := []byte(key)
	t := time.Now().UTC()

	hashkey := sha256.Sum256(keybin)
	salt := sha256.Sum256(
		[]byte(fmt.Sprintf("%d-%02d-%02d", t.Year(), t.Month(), t.Day())))

	hashsalt := append(hashkey[:], salt[:]...)

	aeskey := sha256.Sum256(hashsalt)

	rblock, err := aes.NewCipher(aeskey[:])
	if err != nil {
		return nil, err
	}

	wblock, err := aes.NewCipher(aeskey[:])
	if err != nil {
		return nil, err
	}

	rstream := cipher.NewOFB(rblock, iv[:])
	wstream := cipher.NewOFB(wblock, iv[:])

	aesconn := &AESConn{
		conn: conn,
		r:    &cipher.StreamReader{S: rstream, R: conn},
		w:    &cipher.StreamWriter{S: wstream, W: conn},
	}
	return aesconn, err
}

type Listener struct {
	netl net.Listener
	key  string
}

func (l *Listener) Accept() (net.Conn, error) {
	conn, err := l.netl.Accept()
	if err != nil {
		return nil, err
	}

	conn.SetReadDeadline(time.Now().Add(30 * time.Second))

	ivvector := make([]byte, ivLength)
	var ri = 0
	for ri < ivLength {
		r, _err := conn.Read(ivvector[ri:])
		if _err != nil {
			return nil, _err
		}
		ri += r
	}

	passwd := sha256.Sum256([]byte(l.key))
	longiv := append(passwd[:], ivvector...)

	hash := sha256.Sum256(longiv)

	var iv [aes.BlockSize]byte
	copy(iv[:], hash[:aes.BlockSize])

	aesconn, err := NewAESConn(l.key, iv, conn)

	return aesconn, err
}

func (l *Listener) Close() error {
	return l.netl.Close()
}

func (l *Listener) Addr() net.Addr {
	return l.netl.Addr()
}

func Listen(laddr string, key string) (*Listener, error) {
	netl, err := net.Listen("tcp", laddr)
	if err != nil {
		return nil, err
	}

	l := &Listener{
		netl: netl,
		key:  key,
	}
	return l, nil
}

func Dial(address string, key string) (net.Conn, error) {
	conn, err := net.Dial("tcp", address)
	if err != nil {
		return nil, err
	}

	ivvector := make([]byte, ivLength)
	_, err = rand.Read(ivvector[:])

	if err != nil {
		return nil, err
	}

	var wi = 0
	for wi < ivLength {
		r, _err := conn.Write(ivvector[wi:])
		if _err != nil {
			return nil, _err
		}
		wi += r
	}

	passwd := sha256.Sum256([]byte(key))
	longiv := append(passwd[:], ivvector...)

	hash := sha256.Sum256(longiv)

	var iv [aes.BlockSize]byte
	copy(iv[:], hash[:aes.BlockSize])

	aesconn, err := NewAESConn(key, iv, conn)
	return aesconn, err
}
<|endoftext|>


<commit_before>package layercake_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"syscall"
	"time"

	quotaedaufs "code.cloudfoundry.org/garden-shed/docker_drivers/aufs"
	"code.cloudfoundry.org/garden-shed/layercake"
	"code.cloudfoundry.org/lager/lagertest"
	"github.com/docker/docker/daemon/graphdriver"
	"github.com/docker/docker/graph"
	"github.com/docker/docker/image"
	"github.com/docker/docker/pkg/archive"
	"github.com/eapache/go-resiliency/retrier"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"

	_ "github.com/docker/docker/daemon/graphdriver/aufs"
	_ "github.com/docker/docker/daemon/graphdriver/vfs"
	_ "github.com/docker/docker/pkg/chrootarchive" // allow reexec of docker-applyLayer
	"github.com/docker/docker/pkg/reexec"
)

func init() {
	reexec.Init()
}

var _ = Describe("Docker", func() {
	var (
		root   string
		driver graphdriver.Driver
		cake   *layercake.Docker
	)

	BeforeEach(func() {
		var err error

		root, err = ioutil.TempDir("", "cakeroot")
		Expect(err).NotTo(HaveOccurred())

		Expect(syscall.Mount("tmpfs", root, "tmpfs", 0, "")).To(Succeed())
		driver, err := graphdriver.GetDriver("vfs", root, nil)
		Expect(err).NotTo(HaveOccurred())

		graph, err := graph.NewGraph(root, driver)
		Expect(err).NotTo(HaveOccurred())

		cake = &layercake.Docker{
			Graph:  graph,
			Driver: driver,
		}
	})

	AfterEach(func() {
		Expect(syscall.Unmount(root, 0)).To(Succeed())
		Expect(os.RemoveAll(root)).To(Succeed())
	})

	Describe("Register", func() {
		Context("after registering a layer", func() {
			var id layercake.ID
			var parent layercake.ID

			BeforeEach(func() {
				id = layercake.ContainerID("")
				parent = layercake.ContainerID("")
			})

			ItCanReadWriteTheLayer := func() {
				It("can read and write files", func() {
					p, err := cake.Path(id)
					Expect(err).NotTo(HaveOccurred())
					Expect(ioutil.WriteFile(path.Join(p, "foo"), []byte("hi"), 0700)).To(Succeed())

					p, err = cake.Path(id)
					Expect(err).NotTo(HaveOccurred())
					Expect(path.Join(p, "foo")).To(BeAnExistingFile())
				})

				It("can get back the image", func() {
					img, err := cake.Get(id)
					Expect(err).NotTo(HaveOccurred())
					Expect(img.ID).To(Equal(id.GraphID()))
					Expect(img.Parent).To(Equal(parent.GraphID()))
				})
			}

			Context("when the new layer is a docker image", func() {
				JustBeforeEach(func() {
					id = layercake.DockerImageID("70d8f0edf5c9008eb61c7c52c458e7e0a831649dbb238b93dde0854faae314a8")
					registerImageLayer(cake, &image.Image{
						ID:     id.GraphID(),
						Parent: parent.GraphID(),
					})
				})

				Context("without a parent", func() {
					ItCanReadWriteTheLayer()

					It("can read the files in the image", func() {
						p, err := cake.Path(id)
						Expect(err).NotTo(HaveOccurred())

						Expect(path.Join(p, id.GraphID())).To(BeAnExistingFile())
					})

					It("can be deleted", func() {
						cake.Remove(id)

						filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
							Expect(path).To(BeADirectory())
							return nil
						})
					})
				})

				Context("with a parent", func() {
					BeforeEach(func() {
						parent = layercake.DockerImageID("07d8fe0df5c9008eb16c7c52c548e7e0a831649dbb238b93dde0854faae3148a")
						registerImageLayer(cake, &image.Image{
							ID:     parent.GraphID(),
							Parent: "",
						})
					})

					ItCanReadWriteTheLayer()

					It("inherits files from the parent layer", func() {
						p, err := cake.Path(id)
						Expect(err).NotTo(HaveOccurred())

						Expect(path.Join(p, parent.GraphID())).To(BeAnExistingFile())
					})

					It("can read the files in the image", func() {
						p, err := cake.Path(id)
						Expect(err).NotTo(HaveOccurred())

						Expect(path.Join(p, id.GraphID())).To(BeAnExistingFile())
					})
				})
			})

			Context("when the new layer is a container", func() {
				Context("with a parent", func() {
					BeforeEach(func() {
						parent = layercake.DockerImageID("70d8f0edf5c9008eb61c7c52c458e7e0a831649dbb238b93dde0854faae314a8")
						registerImageLayer(cake, &image.Image{
							ID:     parent.GraphID(),
							Parent: "",
						})

						id = layercake.ContainerID("abc")
						createContainerLayer(cake, id, parent, "potato")
					})

					ItCanReadWriteTheLayer()

					It("inherits files from the parent layer", func() {
						p, err := cake.Path(id)
						Expect(err).NotTo(HaveOccurred())

						Expect(path.Join(p, parent.GraphID())).To(BeAnExistingFile())
					})

					It("saves the container ID in the graph", func() {
						p, err := cake.Get(id)
						Expect(err).NotTo(HaveOccurred())

						Expect(p.Container).To(Equal("potato"))
					})
				})
			})
		})
	})

	Describe("All", func() {
		BeforeEach(func() {
			createContainerLayer(cake, layercake.ContainerID("def"), layercake.DockerImageID(""), "")
			createContainerLayer(cake, layercake.ContainerID("abc"), layercake.ContainerID("def"), "")
			createContainerLayer(cake, layercake.ContainerID("child2"), layercake.ContainerID("def"), "")
		})

		It("returns all the layers in the graph", func() {
			Expect(cake.All()).To(HaveLen(3))

			var ids []string
			for _, layer := range cake.All() {
				ids = append(ids, layer.ID)
			}

			Expect(ids).To(ContainElement(layercake.ContainerID("def").GraphID()))
			Expect(ids).To(ContainElement(layercake.ContainerID("abc").GraphID()))
			Expect(ids).To(ContainElement(layercake.ContainerID("child2").GraphID()))
		})
	})

	Describe("IsLeaf", func() {
		BeforeEach(func() {
			createContainerLayer(cake, layercake.ContainerID("def"), layercake.DockerImageID(""), "")
			createContainerLayer(cake, layercake.ContainerID("abc"), layercake.ContainerID("def"), "")
			createContainerLayer(cake, layercake.ContainerID("child2"), layercake.ContainerID("def"), "")
		})

		Context("when an image has no children", func() {
			It("is a leaf", func() {
				Expect(cake.IsLeaf(layercake.ContainerID("abc"))).To(Equal(true))
			})
		})

		Context("when an image has children", func() {
			It("is not a leaf", func() {
				Expect(cake.IsLeaf(layercake.ContainerID("def"))).To(Equal(false))
			})
		})

		Context("when an image's final child is removed", func() {
			It("is becomes a leaf", func() {
				Expect(cake.IsLeaf(layercake.ContainerID("def"))).To(Equal(false))

				Expect(cake.Remove(layercake.ContainerID("abc"))).To(Succeed())
				Expect(cake.IsLeaf(layercake.ContainerID("def"))).To(Equal(false))

				Expect(cake.Remove(layercake.ContainerID("child2"))).To(Succeed())
				Expect(cake.IsLeaf(layercake.ContainerID("def"))).To(Equal(true))
			})
		})
	})

	Describe("GetAllLeaves", func() {
		BeforeEach(func() {
			createContainerLayer(cake, layercake.ContainerID("def"), layercake.DockerImageID(""), "")
			createContainerLayer(cake, layercake.ContainerID("abc"), layercake.ContainerID("def"), "")
			createContainerLayer(cake, layercake.ContainerID("child2"), layercake.ContainerID("def"), "")
		})

		It("should return all the leaves", func() {
			leaves, err := cake.GetAllLeaves()
			Expect(err).NotTo(HaveOccurred())

			Expect(leaves).To(HaveLen(2))
			Expect(leaves).To(ContainElement(layercake.DockerImageID(layercake.ContainerID("abc").GraphID())))
			Expect(leaves).To(ContainElement(layercake.DockerImageID(layercake.ContainerID("child2").GraphID())))
		})
	})

	Describe("QuotaedPath", func() {
		Context("when not using aufs quotaed driver", func() {
			It("should return an error", func() {
				id := layercake.ContainerID("aubergine-layer")

				registerImageLayer(cake, &image.Image{
					ID: id.GraphID(),
				})

				_, err := cake.QuotaedPath(id, 10*1024*1024)
				Expect(err).To(HaveOccurred())
			})
		})

		Context("when using aufs quotaed driver", func() {
			var (
				backingStoreRoot string
				id               layercake.ContainerID
			)

			BeforeEach(func() {
				var err error

				backingStoreRoot, err = ioutil.TempDir("", "backingstore")
				Expect(err).NotTo(HaveOccurred())

				Expect(syscall.Mount("tmpfs", root, "tmpfs", 0, "")).To(Succeed())

				driver, err = graphdriver.GetDriver("aufs", root, nil)
				Expect(err).NotTo(HaveOccurred())

				driver = &quotaedaufs.QuotaedDriver{
					GraphDriver: driver,
					Unmount:     quotaedaufs.Unmount,
					BackingStoreMgr: &quotaedaufs.BackingStore{
						RootPath: backingStoreRoot,
						Logger:   lagertest.NewTestLogger("test"),
					},
					LoopMounter: &quotaedaufs.Loop{
						Retrier: retrier.New(retrier.ConstantBackoff(200, 500*time.Millisecond), nil),
						Logger:  lagertest.NewTestLogger("test"),
					},
					Retrier:  retrier.New(retrier.ConstantBackoff(200, 500*time.Millisecond), nil),
					RootPath: root,
					Logger:   lagertest.NewTestLogger("test"),
				}

				graph, err := graph.NewGraph(root, driver)
				Expect(err).NotTo(HaveOccurred())

				cake = &layercake.Docker{
					Graph:  graph,
					Driver: driver,
				}
			})

			JustBeforeEach(func() {
				parent := layercake.DockerImageID("70d8f0edf5c9008eb61c7c52c458e7e0a831649dbb238b93dde0854faae314a8")
				registerImageLayer(cake, &image.Image{
					ID:     parent.GraphID(),
					Parent: "",
				})

				id = layercake.ContainerID("abc")
				createContainerLayer(cake, id, parent, "")
			})

			AfterEach(func() {
				Expect(cake.Remove(id)).To(Succeed())
				Expect(driver.Cleanup()).To(Succeed())
				Expect(syscall.Unmount(root, 0)).To(Succeed())
				Expect(os.RemoveAll(backingStoreRoot)).To(Succeed())
			})

			It("should allow the user to write a file that does not exceed the quota", func() {
				path, err := cake.QuotaedPath(id, 10*1024*1024)
				Expect(err).NotTo(HaveOccurred())

				cmd := exec.Command("dd", "if=/dev/zero", fmt.Sprintf("of=%s", filepath.Join(path, "foo")), "bs=1M", "count=8")
				session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
				Expect(err).To(Succeed())

				Eventually(session).Should(gexec.Exit(0))
			})

			It("should not allow the user to write that exceeds the quota", func() {
				path, err := cake.QuotaedPath(id, 10*1024*1024)
				Expect(err).NotTo(HaveOccurred())

				cmd := exec.Command("dd", "if=/dev/zero", fmt.Sprintf("of=%s", filepath.Join(path, "foo")), "bs=1M", "count=11")
				session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
				Expect(err).To(Succeed())

				Eventually(session).ShouldNot(gexec.Exit(0))
			})
		})
	})
})

func createContainerLayer(cake *layercake.Docker, id, parent layercake.ID, containerId string) {
	Expect(cake.Create(id, parent, containerId)).To(Succeed())
}

func registerImageLayer(cake *layercake.Docker, img *image.Image) {
	tmp, err := ioutil.TempDir("", "my-img")
	Expect(err).NotTo(HaveOccurred())
	defer os.RemoveAll(tmp)

	Expect(ioutil.WriteFile(path.Join(tmp, img.ID), []byte("Hello"), 0700)).To(Succeed())
	archiver, _ := archive.Tar(tmp, archive.Uncompressed)

	Expect(cake.Register(img, archiver)).To(Succeed())
}
<commit_msg>Remove duplicate mount/umount in tests<commit_after>package layercake_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"syscall"
	"time"

	quotaedaufs "code.cloudfoundry.org/garden-shed/docker_drivers/aufs"
	"code.cloudfoundry.org/garden-shed/layercake"
	"code.cloudfoundry.org/lager/lagertest"
	"github.com/docker/docker/daemon/graphdriver"
	"github.com/docker/docker/graph"
	"github.com/docker/docker/image"
	"github.com/docker/docker/pkg/archive"
	"github.com/eapache/go-resiliency/retrier"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"

	_ "github.com/docker/docker/daemon/graphdriver/aufs"
	_ "github.com/docker/docker/daemon/graphdriver/vfs"
	_ "github.com/docker/docker/pkg/chrootarchive" // allow reexec of docker-applyLayer
	"github.com/docker/docker/pkg/reexec"
)

func init() {
	reexec.Init()
}

var _ = Describe("Docker", func() {
	var (
		root   string
		driver graphdriver.Driver
		cake   *layercake.Docker
	)

	BeforeEach(func() {
		var err error

		root, err = ioutil.TempDir("", "cakeroot")
		Expect(err).NotTo(HaveOccurred())

		Expect(syscall.Mount("tmpfs", root, "tmpfs", 0, "")).To(Succeed())
		driver, err := graphdriver.GetDriver("vfs", root, nil)
		Expect(err).NotTo(HaveOccurred())

		graph, err := graph.NewGraph(root, driver)
		Expect(err).NotTo(HaveOccurred())

		cake = &layercake.Docker{
			Graph:  graph,
			Driver: driver,
		}
	})

	AfterEach(func() {
		Expect(syscall.Unmount(root, 0)).To(Succeed())
		Expect(os.RemoveAll(root)).To(Succeed())
	})

	Describe("Register", func() {
		Context("after registering a layer", func() {
			var id layercake.ID
			var parent layercake.ID

			BeforeEach(func() {
				id = layercake.ContainerID("")
				parent = layercake.ContainerID("")
			})

			ItCanReadWriteTheLayer := func() {
				It("can read and write files", func() {
					p, err := cake.Path(id)
					Expect(err).NotTo(HaveOccurred())
					Expect(ioutil.WriteFile(path.Join(p, "foo"), []byte("hi"), 0700)).To(Succeed())

					p, err = cake.Path(id)
					Expect(err).NotTo(HaveOccurred())
					Expect(path.Join(p, "foo")).To(BeAnExistingFile())
				})

				It("can get back the image", func() {
					img, err := cake.Get(id)
					Expect(err).NotTo(HaveOccurred())
					Expect(img.ID).To(Equal(id.GraphID()))
					Expect(img.Parent).To(Equal(parent.GraphID()))
				})
			}

			Context("when the new layer is a docker image", func() {
				JustBeforeEach(func() {
					id = layercake.DockerImageID("70d8f0edf5c9008eb61c7c52c458e7e0a831649dbb238b93dde0854faae314a8")
					registerImageLayer(cake, &image.Image{
						ID:     id.GraphID(),
						Parent: parent.GraphID(),
					})
				})

				Context("without a parent", func() {
					ItCanReadWriteTheLayer()

					It("can read the files in the image", func() {
						p, err := cake.Path(id)
						Expect(err).NotTo(HaveOccurred())

						Expect(path.Join(p, id.GraphID())).To(BeAnExistingFile())
					})

					It("can be deleted", func() {
						cake.Remove(id)

						filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
							Expect(path).To(BeADirectory())
							return nil
						})
					})
				})

				Context("with a parent", func() {
					BeforeEach(func() {
						parent = layercake.DockerImageID("07d8fe0df5c9008eb16c7c52c548e7e0a831649dbb238b93dde0854faae3148a")
						registerImageLayer(cake, &image.Image{
							ID:     parent.GraphID(),
							Parent: "",
						})
					})

					ItCanReadWriteTheLayer()

					It("inherits files from the parent layer", func() {
						p, err := cake.Path(id)
						Expect(err).NotTo(HaveOccurred())

						Expect(path.Join(p, parent.GraphID())).To(BeAnExistingFile())
					})

					It("can read the files in the image", func() {
						p, err := cake.Path(id)
						Expect(err).NotTo(HaveOccurred())

						Expect(path.Join(p, id.GraphID())).To(BeAnExistingFile())
					})
				})
			})

			Context("when the new layer is a container", func() {
				Context("with a parent", func() {
					BeforeEach(func() {
						parent = layercake.DockerImageID("70d8f0edf5c9008eb61c7c52c458e7e0a831649dbb238b93dde0854faae314a8")
						registerImageLayer(cake, &image.Image{
							ID:     parent.GraphID(),
							Parent: "",
						})

						id = layercake.ContainerID("abc")
						createContainerLayer(cake, id, parent, "potato")
					})

					ItCanReadWriteTheLayer()

					It("inherits files from the parent layer", func() {
						p, err := cake.Path(id)
						Expect(err).NotTo(HaveOccurred())

						Expect(path.Join(p, parent.GraphID())).To(BeAnExistingFile())
					})

					It("saves the container ID in the graph", func() {
						p, err := cake.Get(id)
						Expect(err).NotTo(HaveOccurred())

						Expect(p.Container).To(Equal("potato"))
					})
				})
			})
		})
	})

	Describe("All", func() {
		BeforeEach(func() {
			createContainerLayer(cake, layercake.ContainerID("def"), layercake.DockerImageID(""), "")
			createContainerLayer(cake, layercake.ContainerID("abc"), layercake.ContainerID("def"), "")
			createContainerLayer(cake, layercake.ContainerID("child2"), layercake.ContainerID("def"), "")
		})

		It("returns all the layers in the graph", func() {
			Expect(cake.All()).To(HaveLen(3))

			var ids []string
			for _, layer := range cake.All() {
				ids = append(ids, layer.ID)
			}

			Expect(ids).To(ContainElement(layercake.ContainerID("def").GraphID()))
			Expect(ids).To(ContainElement(layercake.ContainerID("abc").GraphID()))
			Expect(ids).To(ContainElement(layercake.ContainerID("child2").GraphID()))
		})
	})

	Describe("IsLeaf", func() {
		BeforeEach(func() {
			createContainerLayer(cake, layercake.ContainerID("def"), layercake.DockerImageID(""), "")
			createContainerLayer(cake, layercake.ContainerID("abc"), layercake.ContainerID("def"), "")
			createContainerLayer(cake, layercake.ContainerID("child2"), layercake.ContainerID("def"), "")
		})

		Context("when an image has no children", func() {
			It("is a leaf", func() {
				Expect(cake.IsLeaf(layercake.ContainerID("abc"))).To(Equal(true))
			})
		})

		Context("when an image has children", func() {
			It("is not a leaf", func() {
				Expect(cake.IsLeaf(layercake.ContainerID("def"))).To(Equal(false))
			})
		})

		Context("when an image's final child is removed", func() {
			It("is becomes a leaf", func() {
				Expect(cake.IsLeaf(layercake.ContainerID("def"))).To(Equal(false))

				Expect(cake.Remove(layercake.ContainerID("abc"))).To(Succeed())
				Expect(cake.IsLeaf(layercake.ContainerID("def"))).To(Equal(false))

				Expect(cake.Remove(layercake.ContainerID("child2"))).To(Succeed())
				Expect(cake.IsLeaf(layercake.ContainerID("def"))).To(Equal(true))
			})
		})
	})

	Describe("GetAllLeaves", func() {
		BeforeEach(func() {
			createContainerLayer(cake, layercake.ContainerID("def"), layercake.DockerImageID(""), "")
			createContainerLayer(cake, layercake.ContainerID("abc"), layercake.ContainerID("def"), "")
			createContainerLayer(cake, layercake.ContainerID("child2"), layercake.ContainerID("def"), "")
		})

		It("should return all the leaves", func() {
			leaves, err := cake.GetAllLeaves()
			Expect(err).NotTo(HaveOccurred())

			Expect(leaves).To(HaveLen(2))
			Expect(leaves).To(ContainElement(layercake.DockerImageID(layercake.ContainerID("abc").GraphID())))
			Expect(leaves).To(ContainElement(layercake.DockerImageID(layercake.ContainerID("child2").GraphID())))
		})
	})

	Describe("QuotaedPath", func() {
		Context("when not using aufs quotaed driver", func() {
			It("should return an error", func() {
				id := layercake.ContainerID("aubergine-layer")

				registerImageLayer(cake, &image.Image{
					ID: id.GraphID(),
				})

				_, err := cake.QuotaedPath(id, 10*1024*1024)
				Expect(err).To(HaveOccurred())
			})
		})

		Context("when using aufs quotaed driver", func() {
			var (
				backingStoreRoot string
				id               layercake.ContainerID
			)

			BeforeEach(func() {
				var err error

				backingStoreRoot, err = ioutil.TempDir("", "backingstore")
				Expect(err).NotTo(HaveOccurred())

				driver, err = graphdriver.GetDriver("aufs", root, nil)
				Expect(err).NotTo(HaveOccurred())

				driver = &quotaedaufs.QuotaedDriver{
					GraphDriver: driver,
					Unmount:     quotaedaufs.Unmount,
					BackingStoreMgr: &quotaedaufs.BackingStore{
						RootPath: backingStoreRoot,
						Logger:   lagertest.NewTestLogger("test"),
					},
					LoopMounter: &quotaedaufs.Loop{
						Retrier: retrier.New(retrier.ConstantBackoff(200, 500*time.Millisecond), nil),
						Logger:  lagertest.NewTestLogger("test"),
					},
					Retrier:  retrier.New(retrier.ConstantBackoff(200, 500*time.Millisecond), nil),
					RootPath: root,
					Logger:   lagertest.NewTestLogger("test"),
				}

				graph, err := graph.NewGraph(root, driver)
				Expect(err).NotTo(HaveOccurred())

				cake = &layercake.Docker{
					Graph:  graph,
					Driver: driver,
				}
			})

			JustBeforeEach(func() {
				parent := layercake.DockerImageID("70d8f0edf5c9008eb61c7c52c458e7e0a831649dbb238b93dde0854faae314a8")
				registerImageLayer(cake, &image.Image{
					ID:     parent.GraphID(),
					Parent: "",
				})

				id = layercake.ContainerID("abc")
				createContainerLayer(cake, id, parent, "")
			})

			AfterEach(func() {
				Expect(cake.Remove(id)).To(Succeed())
				Expect(driver.Cleanup()).To(Succeed())
				Expect(os.RemoveAll(backingStoreRoot)).To(Succeed())
			})

			It("should allow the user to write a file that does not exceed the quota", func() {
				path, err := cake.QuotaedPath(id, 10*1024*1024)
				Expect(err).NotTo(HaveOccurred())

				cmd := exec.Command("dd", "if=/dev/zero", fmt.Sprintf("of=%s", filepath.Join(path, "foo")), "bs=1M", "count=8")
				session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
				Expect(err).To(Succeed())

				Eventually(session).Should(gexec.Exit(0))
			})

			It("should not allow the user to write that exceeds the quota", func() {
				path, err := cake.QuotaedPath(id, 10*1024*1024)
				Expect(err).NotTo(HaveOccurred())

				cmd := exec.Command("dd", "if=/dev/zero", fmt.Sprintf("of=%s", filepath.Join(path, "foo")), "bs=1M", "count=11")
				session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
				Expect(err).To(Succeed())

				Eventually(session).ShouldNot(gexec.Exit(0))
			})
		})
	})
})

func createContainerLayer(cake *layercake.Docker, id, parent layercake.ID, containerId string) {
	Expect(cake.Create(id, parent, containerId)).To(Succeed())
}

func registerImageLayer(cake *layercake.Docker, img *image.Image) {
	tmp, err := ioutil.TempDir("", "my-img")
	Expect(err).NotTo(HaveOccurred())
	defer os.RemoveAll(tmp)

	Expect(ioutil.WriteFile(path.Join(tmp, img.ID), []byte("Hello"), 0700)).To(Succeed())
	archiver, _ := archive.Tar(tmp, archive.Uncompressed)

	Expect(cake.Register(img, archiver)).To(Succeed())
}
<|endoftext|>


<commit_before>// Copyright 2020 The LUCI 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 ledcmd

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/golang/protobuf/jsonpb"
	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"

	"go.chromium.org/luci/led/job"
)

// EditRecipeBundleOpts are user-provided options for the recipe bundling
// process.
type EditRecipeBundleOpts struct {
	// Path on disk to the repo to extract the recipes from. May be a subdirectory
	// of the repo, as long as `git rev-parse --show-toplevel` can find the root
	// of the repository.
	//
	// If empty, uses the current working directory.
	RepoDir string

	// Overrides is a mapping of recipe project id (e.g. "recipe_engine") to
	// a local path to a checkout of that repo (e.g. "/path/to/recipes-py.git").
	//
	// When the bundle is created, this local repo will be used instead of the
	// pinned version of this recipe project id. This is helpful for preparing
	// bundles which have code changes in multiple recipe repos.
	Overrides map[string]string

	// DebugSleep is the amount of time to wait after the recipe completes
	// execution (either success or failure). This is injected into the generated
	// recipe bundle as a 'sleep X' command after the invocation of the recipe
	// itself.
	DebugSleep time.Duration

	// PropertyOnly determines whether to pass the recipe bundle's CAS reference
	// as a property and preserve the executable and payload of the input job
	// rather than overwriting it.
	PropertyOnly bool
}

const (
	// RecipeDirectory is a very unfortunate constant which is here for
	// a combination of reasons:
	//   1) swarming doesn't allow you to 'checkout' an isolate relative to any
	//      path in the task (other than the task root). This means that
	//      whatever value we pick for EditRecipeBundle must be used EVERYWHERE
	//      the isolated hash is used.
	//   2) Currently the 'recipe_engine/led' module will blindly take the
	//      isolated input and 'inject' it into further uses of led. This module
	//      currently doesn't specify the checkout dir, relying on kitchen's
	//      default value of (you guessed it) "kitchen-checkout".
	//
	// In order to fix this (and it will need to be fixed for bbagent support):
	//   * The 'recipe_engine/led' module needs to accept 'checkout-dir' as
	//     a parameter in its input properties.
	//   * led needs to start passing the checkout dir to the led module's input
	//     properties.
	//   * `led edit` needs a way to manipulate the checkout directory in a job
	//   * The 'recipe_engine/led' module needs to set this in the job
	//     alongside the isolate hash when it's doing the injection.
	//
	// For now, we just hard-code it.
	//
	// TODO(crbug.com/1072117): Fix this, it's weird.
	RecipeDirectory = "kitchen-checkout"

	// A property that should be set to a boolean value. If true,
	// edit-recipe-bundle will set the "led_cas_recipe_bundle" property
	// instead of overwriting the build's payload.
	LEDBuilderIsBootstrappedProperty = "led_builder_is_bootstrapped"

	// In PropertyOnly mode or if the "led_builder_is_bootstrapped" property
	// of the build is true, this property will be set with the CAS digest
	// of the executable of the recipe bundle.
	CASRecipeBundleProperty = "led_cas_recipe_bundle"
)

// EditRecipeBundle overrides the recipe bundle in the given job with one
// located on disk.
//
// It isolates the recipes from the repository in the given working directory
// into the UserPayload under the directory "kitchen-checkout/". If there's an
// existing directory in the UserPayload at that location, it will be removed.
func EditRecipeBundle(ctx context.Context, authOpts auth.Options, jd *job.Definition, opts *EditRecipeBundleOpts) error {
	if jd.GetBuildbucket() == nil {
		return errors.New("ledcmd.EditRecipeBundle is only available for Buildbucket tasks")
	}

	if opts == nil {
		opts = &EditRecipeBundleOpts{}
	}

	recipesPy, err := findRecipesPy(ctx, opts.RepoDir)
	if err != nil {
		return err
	}
	logging.Debugf(ctx, "using recipes.py: %q", recipesPy)

	extraProperties := make(map[string]string)
	setRecipeBundleProperty := opts.PropertyOnly || jd.GetBuildbucket().GetBbagentArgs().GetBuild().GetInput().GetProperties().GetFields()[LEDBuilderIsBootstrappedProperty].GetBoolValue()
	if setRecipeBundleProperty {
		// In property-only mode, we want to leave the original payload as is
		// and just upload the recipe bundle as a brand new independent CAS
		// archive for the job's executable to download.
		bundlePath, err := ioutil.TempDir("", "led-recipe-bundle")
		if err != nil {
			return errors.Annotate(err, "creating temporary recipe bundle directory").Err()
		}
		opts.prepBundle(ctx, opts.RepoDir, recipesPy, bundlePath)
		logging.Infof(ctx, "isolating recipes")
		casClient, err := newCASClient(ctx, authOpts, jd)
		if err != nil {
			return err
		}
		casRef, err := uploadToCas(ctx, casClient, bundlePath)
		if err != nil {
			return err
		}
		m := &jsonpb.Marshaler{OrigName: true}
		jsonCASRef, err := m.MarshalToString(casRef)
		if err != nil {
			return errors.Annotate(err, "encoding CAS user payload").Err()
		}
		extraProperties[CASRecipeBundleProperty] = jsonCASRef
	} else {
		if err := EditIsolated(ctx, authOpts, jd, func(ctx context.Context, dir string) error {
			bundlePath := filepath.Join(dir, RecipeDirectory)
			// Remove existing bundled recipes, if any. Ignore the error.
			os.RemoveAll(bundlePath)
			if err := opts.prepBundle(ctx, opts.RepoDir, recipesPy, bundlePath); err != nil {
				return err
			}
			logging.Infof(ctx, "isolating recipes")
			return nil
		}); err != nil {
			return err
		}
	}

	return jd.HighLevelEdit(func(je job.HighLevelEditor) {
		if setRecipeBundleProperty {
			je.Properties(extraProperties, false)
		} else {
			je.TaskPayloadSource("", "")
			je.TaskPayloadPath(RecipeDirectory)
		}
		if opts.DebugSleep != 0 {
			je.Env(map[string]string{
				"RECIPES_DEBUG_SLEEP": fmt.Sprintf("%f", opts.DebugSleep.Seconds()),
			})
		}
	})
}

func logCmd(ctx context.Context, inDir string, arg0 string, args ...string) *exec.Cmd {
	ret := exec.CommandContext(ctx, arg0, args...)
	ret.Dir = inDir
	logging.Debugf(ctx, "Running (from %q) - %s %v", inDir, arg0, args)
	return ret
}

func cmdErr(cmd *exec.Cmd, err error, reason string) error {
	if err != nil {
		ee, _ := err.(*exec.ExitError)
		outErr := ""
		if ee != nil {
			outErr = strings.TrimSpace(string(ee.Stderr))
			if len(outErr) > 128 {
				outErr = outErr[:128] + "..."
			}
		}
		err = errors.Annotate(err, "running %q: %s: %s", strings.Join(cmd.Args, " "), reason, outErr).Err()
	}
	return err
}

func appendText(path, fmtStr string, items ...interface{}) error {
	file, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0666)
	if err != nil {
		return err
	}
	defer file.Close()
	_, err = fmt.Fprintf(file, fmtStr, items...)
	return err
}

func (opts *EditRecipeBundleOpts) prepBundle(ctx context.Context, inDir, recipesPy, toDirectory string) error {
	logging.Infof(ctx, "bundling recipes")
	args := []string{
		recipesPy,
	}
	if logging.GetLevel(ctx) < logging.Info {
		args = append(args, "-v")
	}
	for projID, path := range opts.Overrides {
		args = append(args, "-O", fmt.Sprintf("%s=%s", projID, path))
	}
	args = append(args, "bundle", "--destination", filepath.Join(toDirectory))
	cmd := logCmd(ctx, inDir, "python", args...)
	if logging.GetLevel(ctx) < logging.Info {
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
	}
	return cmdErr(cmd, cmd.Run(), "creating bundle")
}

// findRecipesPy locates the current repo's `recipes.py`. It does this by:
//   * invoking git to find the repo root
//   * loading the recipes.cfg at infra/config/recipes.cfg
//   * stat'ing the recipes.py implied by the recipes_path in that cfg file.
//
// Failure will return an error.
//
// On success, the absolute path to recipes.py is returned.
func findRecipesPy(ctx context.Context, inDir string) (string, error) {
	cmd := logCmd(ctx, inDir, "git", "rev-parse", "--show-toplevel")
	out, err := cmd.Output()
	if err = cmdErr(cmd, err, "finding git repo"); err != nil {
		return "", err
	}

	repoRoot := strings.TrimSpace(string(out))

	pth := filepath.Join(repoRoot, "infra", "config", "recipes.cfg")
	switch st, err := os.Stat(pth); {
	case err != nil:
		return "", errors.Annotate(err, "reading recipes.cfg").Err()

	case !st.Mode().IsRegular():
		return "", errors.Reason("%q is not a regular file", pth).Err()
	}

	type recipesJSON struct {
		RecipesPath string `json:"recipes_path"`
	}
	rj := &recipesJSON{}

	f, err := os.Open(pth)
	if err != nil {
		return "", errors.Reason("reading recipes.cfg: %q", pth).Err()
	}
	defer f.Close()

	if err := json.NewDecoder(f).Decode(rj); err != nil {
		return "", errors.Reason("parsing recipes.cfg: %q", pth).Err()
	}

	return filepath.Join(
		repoRoot, filepath.FromSlash(rj.RecipesPath), "recipes.py"), nil
}
<commit_msg>[led] Propagate error if `python` is not in $PATH<commit_after>// Copyright 2020 The LUCI 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 ledcmd

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/golang/protobuf/jsonpb"
	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"

	"go.chromium.org/luci/led/job"
)

// EditRecipeBundleOpts are user-provided options for the recipe bundling
// process.
type EditRecipeBundleOpts struct {
	// Path on disk to the repo to extract the recipes from. May be a subdirectory
	// of the repo, as long as `git rev-parse --show-toplevel` can find the root
	// of the repository.
	//
	// If empty, uses the current working directory.
	RepoDir string

	// Overrides is a mapping of recipe project id (e.g. "recipe_engine") to
	// a local path to a checkout of that repo (e.g. "/path/to/recipes-py.git").
	//
	// When the bundle is created, this local repo will be used instead of the
	// pinned version of this recipe project id. This is helpful for preparing
	// bundles which have code changes in multiple recipe repos.
	Overrides map[string]string

	// DebugSleep is the amount of time to wait after the recipe completes
	// execution (either success or failure). This is injected into the generated
	// recipe bundle as a 'sleep X' command after the invocation of the recipe
	// itself.
	DebugSleep time.Duration

	// PropertyOnly determines whether to pass the recipe bundle's CAS reference
	// as a property and preserve the executable and payload of the input job
	// rather than overwriting it.
	PropertyOnly bool
}

const (
	// RecipeDirectory is a very unfortunate constant which is here for
	// a combination of reasons:
	//   1) swarming doesn't allow you to 'checkout' an isolate relative to any
	//      path in the task (other than the task root). This means that
	//      whatever value we pick for EditRecipeBundle must be used EVERYWHERE
	//      the isolated hash is used.
	//   2) Currently the 'recipe_engine/led' module will blindly take the
	//      isolated input and 'inject' it into further uses of led. This module
	//      currently doesn't specify the checkout dir, relying on kitchen's
	//      default value of (you guessed it) "kitchen-checkout".
	//
	// In order to fix this (and it will need to be fixed for bbagent support):
	//   * The 'recipe_engine/led' module needs to accept 'checkout-dir' as
	//     a parameter in its input properties.
	//   * led needs to start passing the checkout dir to the led module's input
	//     properties.
	//   * `led edit` needs a way to manipulate the checkout directory in a job
	//   * The 'recipe_engine/led' module needs to set this in the job
	//     alongside the isolate hash when it's doing the injection.
	//
	// For now, we just hard-code it.
	//
	// TODO(crbug.com/1072117): Fix this, it's weird.
	RecipeDirectory = "kitchen-checkout"

	// A property that should be set to a boolean value. If true,
	// edit-recipe-bundle will set the "led_cas_recipe_bundle" property
	// instead of overwriting the build's payload.
	LEDBuilderIsBootstrappedProperty = "led_builder_is_bootstrapped"

	// In PropertyOnly mode or if the "led_builder_is_bootstrapped" property
	// of the build is true, this property will be set with the CAS digest
	// of the executable of the recipe bundle.
	CASRecipeBundleProperty = "led_cas_recipe_bundle"
)

// EditRecipeBundle overrides the recipe bundle in the given job with one
// located on disk.
//
// It isolates the recipes from the repository in the given working directory
// into the UserPayload under the directory "kitchen-checkout/". If there's an
// existing directory in the UserPayload at that location, it will be removed.
func EditRecipeBundle(ctx context.Context, authOpts auth.Options, jd *job.Definition, opts *EditRecipeBundleOpts) error {
	if jd.GetBuildbucket() == nil {
		return errors.New("ledcmd.EditRecipeBundle is only available for Buildbucket tasks")
	}

	if opts == nil {
		opts = &EditRecipeBundleOpts{}
	}

	recipesPy, err := findRecipesPy(ctx, opts.RepoDir)
	if err != nil {
		return err
	}
	logging.Debugf(ctx, "using recipes.py: %q", recipesPy)

	extraProperties := make(map[string]string)
	setRecipeBundleProperty := opts.PropertyOnly || jd.GetBuildbucket().GetBbagentArgs().GetBuild().GetInput().GetProperties().GetFields()[LEDBuilderIsBootstrappedProperty].GetBoolValue()
	if setRecipeBundleProperty {
		// In property-only mode, we want to leave the original payload as is
		// and just upload the recipe bundle as a brand new independent CAS
		// archive for the job's executable to download.
		bundlePath, err := ioutil.TempDir("", "led-recipe-bundle")
		if err != nil {
			return errors.Annotate(err, "creating temporary recipe bundle directory").Err()
		}
		if err := opts.prepBundle(ctx, opts.RepoDir, recipesPy, bundlePath); err != nil {
			return err
		}
		logging.Infof(ctx, "isolating recipes")
		casClient, err := newCASClient(ctx, authOpts, jd)
		if err != nil {
			return err
		}
		casRef, err := uploadToCas(ctx, casClient, bundlePath)
		if err != nil {
			return err
		}
		m := &jsonpb.Marshaler{OrigName: true}
		jsonCASRef, err := m.MarshalToString(casRef)
		if err != nil {
			return errors.Annotate(err, "encoding CAS user payload").Err()
		}
		extraProperties[CASRecipeBundleProperty] = jsonCASRef
	} else {
		if err := EditIsolated(ctx, authOpts, jd, func(ctx context.Context, dir string) error {
			bundlePath := filepath.Join(dir, RecipeDirectory)
			// Remove existing bundled recipes, if any. Ignore the error.
			os.RemoveAll(bundlePath)
			if err := opts.prepBundle(ctx, opts.RepoDir, recipesPy, bundlePath); err != nil {
				return err
			}
			logging.Infof(ctx, "isolating recipes")
			return nil
		}); err != nil {
			return err
		}
	}

	return jd.HighLevelEdit(func(je job.HighLevelEditor) {
		if setRecipeBundleProperty {
			je.Properties(extraProperties, false)
		} else {
			je.TaskPayloadSource("", "")
			je.TaskPayloadPath(RecipeDirectory)
		}
		if opts.DebugSleep != 0 {
			je.Env(map[string]string{
				"RECIPES_DEBUG_SLEEP": fmt.Sprintf("%f", opts.DebugSleep.Seconds()),
			})
		}
	})
}

func logCmd(ctx context.Context, inDir string, arg0 string, args ...string) *exec.Cmd {
	ret := exec.CommandContext(ctx, arg0, args...)
	ret.Dir = inDir
	logging.Debugf(ctx, "Running (from %q) - %s %v", inDir, arg0, args)
	return ret
}

func cmdErr(cmd *exec.Cmd, err error, reason string) error {
	if err == nil {
		return nil
	}
	var outErr string
	if ee, ok := err.(*exec.ExitError); ok {
		outErr = strings.TrimSpace(string(ee.Stderr))
		if len(outErr) > 128 {
			outErr = outErr[:128] + "..."
		}
	} else {
		outErr = err.Error()
	}
	return errors.Annotate(err, "running %q: %s: %s", strings.Join(cmd.Args, " "), reason, outErr).Err()
}

func appendText(path, fmtStr string, items ...interface{}) error {
	file, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0666)
	if err != nil {
		return err
	}
	defer file.Close()
	_, err = fmt.Fprintf(file, fmtStr, items...)
	return err
}

func (opts *EditRecipeBundleOpts) prepBundle(ctx context.Context, inDir, recipesPy, toDirectory string) error {
	logging.Infof(ctx, "bundling recipes")
	args := []string{
		recipesPy,
	}
	if logging.GetLevel(ctx) < logging.Info {
		args = append(args, "-v")
	}
	for projID, path := range opts.Overrides {
		args = append(args, "-O", fmt.Sprintf("%s=%s", projID, path))
	}
	args = append(args, "bundle", "--destination", filepath.Join(toDirectory))
	cmd := logCmd(ctx, inDir, "python", args...)
	if logging.GetLevel(ctx) < logging.Info {
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
	}
	return cmdErr(cmd, cmd.Run(), "creating bundle")
}

// findRecipesPy locates the current repo's `recipes.py`. It does this by:
//   * invoking git to find the repo root
//   * loading the recipes.cfg at infra/config/recipes.cfg
//   * stat'ing the recipes.py implied by the recipes_path in that cfg file.
//
// Failure will return an error.
//
// On success, the absolute path to recipes.py is returned.
func findRecipesPy(ctx context.Context, inDir string) (string, error) {
	cmd := logCmd(ctx, inDir, "git", "rev-parse", "--show-toplevel")
	out, err := cmd.Output()
	if err = cmdErr(cmd, err, "finding git repo"); err != nil {
		return "", err
	}

	repoRoot := strings.TrimSpace(string(out))

	pth := filepath.Join(repoRoot, "infra", "config", "recipes.cfg")
	switch st, err := os.Stat(pth); {
	case err != nil:
		return "", errors.Annotate(err, "reading recipes.cfg").Err()

	case !st.Mode().IsRegular():
		return "", errors.Reason("%q is not a regular file", pth).Err()
	}

	type recipesJSON struct {
		RecipesPath string `json:"recipes_path"`
	}
	rj := &recipesJSON{}

	f, err := os.Open(pth)
	if err != nil {
		return "", errors.Reason("reading recipes.cfg: %q", pth).Err()
	}
	defer f.Close()

	if err := json.NewDecoder(f).Decode(rj); err != nil {
		return "", errors.Reason("parsing recipes.cfg: %q", pth).Err()
	}

	return filepath.Join(
		repoRoot, filepath.FromSlash(rj.RecipesPath), "recipes.py"), nil
}
<|endoftext|>


<commit_before>package dependency

import (
	"errors"
	"fmt"
	"log"
	"regexp"

	"github.com/hashicorp/consul/api"
)

// Node is a node entry in Consul
type Node struct {
	Node    string
	Address string
}

type CatalogNodes struct {
	rawKey     string
	DataCenter string
}

// Fetch queries the Consul API defined by the given client and returns a slice
// of Node objects
func (d *CatalogNodes) Fetch(client *api.Client, options *api.QueryOptions) (interface{}, *api.QueryMeta, error) {
	if d.DataCenter != "" {
		options.Datacenter = d.DataCenter
	}

	log.Printf("[DEBUG] (%s) querying Consul with %+v", d.Display(), options)

	catalog := client.Catalog()
	n, qm, err := catalog.Nodes(options)
	if err != nil {
		return n, qm, err
	}

	log.Printf("[DEBUG] (%s) Consul returned %d nodes", d.Display(), len(n))

	nodes := make([]*Node, 0, len(n))
	for _, node := range n {
		nodes = append(nodes, &Node{
			Node:    node.Node,
			Address: node.Address,
		})
	}

	return nodes, qm, nil
}

func (d *CatalogNodes) HashCode() string {
	return fmt.Sprintf("CatalogNodes|%s", d.rawKey)
}

func (d *CatalogNodes) Display() string {
	if d.rawKey == "" {
		return fmt.Sprintf("nodes")
	}

	return fmt.Sprintf(`nodes "%s"`, d.rawKey)
}

// ParseCatalogNodes parses a string of the format @dc.
func ParseCatalogNodes(s ...string) (*CatalogNodes, error) {
	switch len(s) {
	case 0:
		return &CatalogNodes{rawKey: ""}, nil
	case 1:
		dc := s[0]

		re := regexp.MustCompile(`\A` +
			`(@(?P<datacenter>[[:word:]\.\-]+))?` +
			`\z`)
		names := re.SubexpNames()
		match := re.FindAllStringSubmatch(dc, -1)

		if len(match) == 0 {
			return nil, errors.New("invalid node dependency format")
		}

		r := match[0]

		m := map[string]string{}
		for i, n := range r {
			if names[i] != "" {
				m[names[i]] = n
			}
		}

		nd := &CatalogNodes{
			rawKey:     dc,
			DataCenter: m["datacenter"],
		}

		return nd, nil
	default:
		return nil, fmt.Errorf("expected 0 or 1 arguments, got %d", len(s))
	}
}
<commit_msg>Return nil instead<commit_after>package dependency

import (
	"errors"
	"fmt"
	"log"
	"regexp"

	"github.com/hashicorp/consul/api"
)

// Node is a node entry in Consul
type Node struct {
	Node    string
	Address string
}

type CatalogNodes struct {
	rawKey     string
	DataCenter string
}

// Fetch queries the Consul API defined by the given client and returns a slice
// of Node objects
func (d *CatalogNodes) Fetch(client *api.Client, options *api.QueryOptions) (interface{}, *api.QueryMeta, error) {
	if d.DataCenter != "" {
		options.Datacenter = d.DataCenter
	}

	log.Printf("[DEBUG] (%s) querying Consul with %+v", d.Display(), options)

	catalog := client.Catalog()
	n, qm, err := catalog.Nodes(options)
	if err != nil {
		return nil, qm, err
	}

	log.Printf("[DEBUG] (%s) Consul returned %d nodes", d.Display(), len(n))

	nodes := make([]*Node, 0, len(n))
	for _, node := range n {
		nodes = append(nodes, &Node{
			Node:    node.Node,
			Address: node.Address,
		})
	}

	return nodes, qm, nil
}

func (d *CatalogNodes) HashCode() string {
	return fmt.Sprintf("CatalogNodes|%s", d.rawKey)
}

func (d *CatalogNodes) Display() string {
	if d.rawKey == "" {
		return fmt.Sprintf("nodes")
	}

	return fmt.Sprintf(`nodes "%s"`, d.rawKey)
}

// ParseCatalogNodes parses a string of the format @dc.
func ParseCatalogNodes(s ...string) (*CatalogNodes, error) {
	switch len(s) {
	case 0:
		return &CatalogNodes{rawKey: ""}, nil
	case 1:
		dc := s[0]

		re := regexp.MustCompile(`\A` +
			`(@(?P<datacenter>[[:word:]\.\-]+))?` +
			`\z`)
		names := re.SubexpNames()
		match := re.FindAllStringSubmatch(dc, -1)

		if len(match) == 0 {
			return nil, errors.New("invalid node dependency format")
		}

		r := match[0]

		m := map[string]string{}
		for i, n := range r {
			if names[i] != "" {
				m[names[i]] = n
			}
		}

		nd := &CatalogNodes{
			rawKey:     dc,
			DataCenter: m["datacenter"],
		}

		return nd, nil
	default:
		return nil, fmt.Errorf("expected 0 or 1 arguments, got %d", len(s))
	}
}
<|endoftext|>