<commit_before>package main

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/wayf-dk/pkcs11"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Hsm struct {
	session pkcs11.SessionHandle
	used    int
	started time.Time
	sessno int
}

var currentsessions int
var hsm1 Hsm
var sem chan Hsm
var pguard sync.Mutex
var p *pkcs11.Ctx
var config = map[string]string{
	"GOELEVEN_HSMLIB":        "",
	"GOELEVEN_INTERFACE":     "localhost:8080",
	"GOELEVEN_ALLOWEDIP":     "127.0.0.1",
	"GOELEVEN_SLOT":          "",
	"GOELEVEN_SLOT_PASSWORD": "",
	"GOELEVEN_KEY_LABEL":     "",
	"GOELEVEN_SHAREDSECRET":  "",
	"GOELEVEN_MINSESSIONS":   "1",
	"GOELEVEN_MAXSESSIONS":   "1",
	"GOELEVEN_MAXSESSIONAGE": "1000000",
	"GOELEVEN_MECH":          "CKM_RSA_PKCS",
	"GOELEVEN_DEBUG":         "false",
	"SOFTHSM_CONF":           "softhsm.conf",
	"GOELEVEN_HTTPS_KEY":     "false",
	"GOELEVEN_HTTPS_CERT":    "false",
}

var keymap map[string]pkcs11.ObjectHandle

var xauthlen = map[string]int{
	"min": 12,
	"max": 32,
}

func main() {
	currentsessions = 0
	keymap = make(map[string]pkcs11.ObjectHandle)
	//wd, _ := os.Getwd()
	initConfig()
	p = pkcs11.New(config["GOELEVEN_HSMLIB"])
	p.Initialize()
	handlesessions()
	http.HandleFunc("/", handler)
	var err error
	if config["GOELEVEN_HTTPS_CERT"] == "false" {
	    err = http.ListenAndServe(config["GOELEVEN_INTERFACE"], nil)
	} else {
        err = http.ListenAndServeTLS(config["GOELEVEN_INTERFACE"], config["GOELEVEN_HTTPS_CERT"], config["GOELEVEN_HTTPS_KEY"], nil)
    }
    if err != nil {
        fmt.Printf("main(): %s\n", err)
    }
}

// initConfig read several Environment variables and based on them initialise the configuration
func initConfig() {
	envFiles := []string{"GOELEVEN_HSMLIB"}

	// Load all Environments variables
	for k, _ := range config {
		if os.Getenv(k) != "" {
			config[k] = os.Getenv(k)
		}
	}
	// All variable MUST have a value but we can not verify the variable content
	for k, _ := range config {
		if isdebug() {
			// Don't write PASSWORD to debug
			if k == "GOELEVEN_SLOT_PASSWORD" {
				debug(fmt.Sprintf("%v: xxxxxx\n", k))
			} else {
				debug(fmt.Sprintf("%v: %v\n", k, config[k]))
			}
		}
		if config[k] == "" {
			exit(fmt.Sprintf("Problem with %s", k), 2)
		}
	}

	// Check file exists
	for _, v := range envFiles {
		_, err := os.Stat(config[v])
		if err != nil {
			exit(fmt.Sprintf("%s %s", v, err.Error()), 2)
		}
	}
	// Test XAUTH enviroment
	_, err := sanitizeXAuth(config["GOELEVEN_SHAREDSECRET"])
	if err != nil {
		exit(fmt.Sprintf("GOELEVEN_SHAREDSECRET: %v", err.Error()), 2)
	}
}

func handlesessions() {
	// String->int64->int convert
	max, _ := strconv.ParseInt(config["GOELEVEN_MAXSESSIONS"], 10, 0)
	var maxsessions int = int(max)
	sem = make(chan Hsm, maxsessions)
	for currentsessions < maxsessions {
		currentsessions++
		sem <- inithsm(currentsessions)
	}

	s := <-sem

    keys := strings.Split(config["GOELEVEN_KEY_LABEL"], ",")

    for _, v := range keys {

        template := []*pkcs11.Attribute{pkcs11.NewAttribute(pkcs11.CKA_LABEL, v), pkcs11.NewAttribute(pkcs11.CKA_CLASS, pkcs11.CKO_PRIVATE_KEY)}
        if e := p.FindObjectsInit(s.session, template); e != nil {
            panic(fmt.Sprintf("Failed to init: %s\n", e.Error()))
        }
        obj, b, e := p.FindObjects(s.session, 2)

        debug(fmt.Sprintf("Obj %v\n", obj))
        if e != nil {
            exit(fmt.Sprintf("Failed to find: %s %v\n", e.Error(), b), 2)
        }
        if e := p.FindObjectsFinal(s.session); e != nil {
            exit(fmt.Sprintf("Failed to finalize: %s\n", e.Error()), 2)
        }
        debug(fmt.Sprintf("found keys: %v\n", len(obj)))
        if len(obj) == 0 {
            exit(fmt.Sprintf("did not find a key with label '%s'", v), 2)
        }
        keymap[v] = obj[0]
    }

	fmt.Printf("hsm initialized new: %#v\n", keymap)

	sem <- s

	debug(fmt.Sprintf("sem: %v\n", len(sem)))
}

// Check if the X-Auth string are safe to use
func sanitizeXAuth(insecureXAuth string) (string, error) {
	if len(insecureXAuth) >= xauthlen["min"] && len(insecureXAuth) <= xauthlen["max"] {
		return insecureXAuth, nil
	}
	return "", errors.New("X-AUTH do not complies with the defined rules")
}

// Client authenticate/authorization
func authClient(sharedkey string, slot string, keylabel string, mech string) error {
	//  Check sharedkey
	if sharedkey != config["GOELEVEN_SHAREDSECRET"] {
		return errors.New("Shared secret does not match")
	}
	//  Check slot nummer
	if slot != config["GOELEVEN_SLOT"] {
		return errors.New("Slot number does not match")
	}
	//  Check key aliases/label
	if keymap[keylabel] == 0 {
		return errors.New(fmt.Sprintf("Key label does not match %s", keylabel))
	}
	//  Check key mech
	if mech != config["GOELEVEN_MECH"] {
		return errors.New("Mech does not match")
	}
	// client ok
	return nil
}

// TODO: Cleanup
// TODO: Documentation
// TODO: Error handling
/*
 * If error then send HTTP 500 to client and keep the server running
 *
 */
func handler(w http.ResponseWriter, r *http.Request) {

     fmt.Println("access attempt from:", r.RemoteAddr)
     ips := strings.Split(config["GOELEVEN_ALLOWEDIP"], ",")
     ip := strings.Split(r.RemoteAddr, ":")
     var allowed bool
     for _, v := range ips {
        allowed = allowed || ip[0] == v
     }

     if (!allowed) {
         fmt.Println("unauthorised access attempt from:", r.RemoteAddr)
         http.Error(w,"Unauthorized",http.StatusUnauthorized)
         return;
     }

	var err error
	var validPath = regexp.MustCompile("^/(\\d+)/([a-zA-Z0-9\\.]+)/sign$")
	mSlot := validPath.FindStringSubmatch(r.URL.Path)[1]
	mKeyAlias := validPath.FindStringSubmatch(r.URL.Path)[2]

	defer r.Body.Close()
	body, _ := ioutil.ReadAll(r.Body)

	// Parse JSON
	//var b struct { Data,Mech string }
	var b map[string]interface{}
	err = json.Unmarshal(body, &b)
	if err != nil {
		http.Error(w, "Invalid input", 500)
		fmt.Printf("json.unmarshall: %v\n", err.Error())
		return
	}
	data, err := base64.StdEncoding.DecodeString(b["data"].(string))
	if err != nil {
		http.Error(w, "Invalid input", 500)
		fmt.Printf("DecodeString: %v\n", err.Error())
		return
	}

	// Client auth
	err = authClient(b["sharedkey"].(string), mSlot, mKeyAlias, b["mech"].(string))
	if err != nil {
		http.Error(w, "Invalid input", 500)
		fmt.Printf("authClient: %v\n", err.Error())
		return
	}

	sig, err, sessno:= signing(data, keymap[mKeyAlias])
	if err != nil {
		http.Error(w, "Invalid output", 500)
		fmt.Printf("signing: %v %v\n", err.Error(), sessno)
		return
	}
	sigs := base64.StdEncoding.EncodeToString(sig)
	type Res struct {
		Slot   string `json:"slot"`
		Mech   string `json:"mech"`
		Signed string `json:"signed"`
	}
	res := Res{mSlot, "mech", sigs}
	json, err := json.Marshal(res)
	if err != nil {
		http.Error(w, "Invalid output", 500)
		fmt.Printf("json.marshall: %v\n", err.Error())
		return
	}
	fmt.Fprintf(w, "%s\n\n", json)
}

// TODO: Cleanup
// TODO: Documentation
func inithsm(sessno int) Hsm {
	pguard.Lock()
	defer pguard.Unlock()
    slot, _ := strconv.ParseUint(config["GOELEVEN_SLOT"], 10, 32)

    fmt.Printf("slot: %v\n", slot)
    session, e := p.OpenSession(uint(slot), pkcs11.CKF_SERIAL_SESSION )

    if e != nil {
        panic(fmt.Sprintf("Failed to open session: %s\n", e.Error()))
    }

	p.Login(session, pkcs11.CKU_USER, config["GOELEVEN_SLOT_PASSWORD"])

	return Hsm{session, 0, time.Now(), sessno}
}

// TODO: Cleanup
// TODO: Documentation
func signing(data []byte, key pkcs11.ObjectHandle) ([]byte, error, int) {
	// Pop HSM struct from queue
	s := <-sem
	s.used++
	if s.used > 10000 || time.Now().Sub(s.started) > 1000*time.Second {
		p.Logout(s.session)
		p.CloseSession(s.session)
		//p.Finalize()
		//p.Destroy()
		s = inithsm(s.sessno)
	}
	fmt.Printf("hsm: %v %v\n", s, key)
	//p.SignInit(s.session, []*pkcs11.Mechanism{pkcs11.NewMechanism(pkcs11.CKM_SHA256_RSA_PKCS, nil)}, key)
	p.SignInit(s.session, []*pkcs11.Mechanism{pkcs11.NewMechanism(pkcs11.CKM_RSA_PKCS, nil)}, key)
	sig, err := p.Sign(s.session, data)
	fmt.Printf("err: %v\n", err)

	// Push HSM struct back on queue
	sem <- s
	return sig, nil, s.sessno
}

// Utils

func debug(messages string) {
	if config["GOELEVEN_DEBUG"] == "true" {
		fmt.Print(messages)
	}
}

// Standard function to test for debug mode
func isdebug() bool {
	if config["GOELEVEN_DEBUG"] == "true" {
		return true
	} else {
		return false
	}
}

func exit(messages string, errorCode int) {
	// Exit code and messages based on Nagios plugin return codes (https://nagios-plugins.org/doc/guidelines.html#AEN78)
	var prefix = map[int]string{0: "OK", 1: "Warning", 2: "Critical", 3: "Unknown"}

	// Catch all unknown errorCode and convert them to Unknown
	if errorCode < 0 || errorCode > 3 {
		errorCode = 3
	}

	fmt.Printf("%s %s\n", prefix[errorCode], messages)
	os.Exit(errorCode)
}
<commit_msg>Added support for per label sharedsecret.<commit_after>package main

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/wayf-dk/pkcs11"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Hsm struct {
	session pkcs11.SessionHandle
	used    int
	started time.Time
	sessno  int
}

type aclmap struct {
	handle       pkcs11.ObjectHandle
	sharedsecret string
	label        string
}

var currentsessions int
var hsm1 Hsm
var sem chan Hsm
var pguard sync.Mutex
var p *pkcs11.Ctx
var config = map[string]string{
	"GOELEVEN_HSMLIB":        "",
	"GOELEVEN_INTERFACE":     "localhost:8080",
	"GOELEVEN_ALLOWEDIP":     "127.0.0.1",
	"GOELEVEN_SLOT":          "",
	"GOELEVEN_SLOT_PASSWORD": "",
	"GOELEVEN_KEY_LABEL":     "",
	"GOELEVEN_MINSESSIONS":   "1",
	"GOELEVEN_MAXSESSIONS":   "1",
	"GOELEVEN_MAXSESSIONAGE": "1000000",
	"GOELEVEN_MECH":          "CKM_RSA_PKCS",
	"GOELEVEN_DEBUG":         "false",
	"SOFTHSM_CONF":           "softhsm.conf",
	"GOELEVEN_HTTPS_KEY":     "false",
	"GOELEVEN_HTTPS_CERT":    "false",
}

var keymap map[string]aclmap

var sharedsecretlen = map[string]int{
	"min": 12,
	"max": 32,
}

func main() {
	currentsessions = 0
	keymap = make(map[string]aclmap)
	//wd, _ := os.Getwd()
	initConfig()
	p = pkcs11.New(config["GOELEVEN_HSMLIB"])
	p.Initialize()
	handlesessions()
	http.HandleFunc("/", handler)
	var err error
	if config["GOELEVEN_HTTPS_CERT"] == "false" {
		err = http.ListenAndServe(config["GOELEVEN_INTERFACE"], nil)
	} else {
		err = http.ListenAndServeTLS(config["GOELEVEN_INTERFACE"], config["GOELEVEN_HTTPS_CERT"], config["GOELEVEN_HTTPS_KEY"], nil)
	}
	if err != nil {
		fmt.Printf("main(): %s\n", err)
	}
}

// initConfig read several Environment variables and based on them initialise the configuration
func initConfig() {
	envFiles := []string{"GOELEVEN_HSMLIB"}

	// Load all Environments variables
	for k, _ := range config {
		if os.Getenv(k) != "" {
			config[k] = os.Getenv(k)
		}
	}
	// All variable MUST have a value but we can not verify the variable content
	for k, _ := range config {
		if isdebug() {
			// Don't write PASSWORD to debug
			if k == "GOELEVEN_SLOT_PASSWORD" {
				debug(fmt.Sprintf("%v: xxxxxx\n", k))
			} else {
				debug(fmt.Sprintf("%v: %v\n", k, config[k]))
			}
		}
		if config[k] == "" {
			exit(fmt.Sprintf("Problem with %s", k), 2)
		}
	}

	// Check file exists
	for _, v := range envFiles {
		_, err := os.Stat(config[v])
		if err != nil {
			exit(fmt.Sprintf("%s %s", v, err.Error()), 2)
		}
	}
}

func handlesessions() {
	// String->int64->int convert
	max, _ := strconv.ParseInt(config["GOELEVEN_MAXSESSIONS"], 10, 0)
	var maxsessions int = int(max)
	sem = make(chan Hsm, maxsessions)
	for currentsessions < maxsessions {
		currentsessions++
		sem <- inithsm(currentsessions)
	}

	s := <-sem

	keys := strings.Split(config["GOELEVEN_KEY_LABEL"], ",")

	for _, v := range keys {

		parts := strings.Split(v, ":")
		label := parts[0]
		sharedsecret := parts[1]
        // Test validity of key specific sharedsecret
        if len(sharedsecret) < sharedsecretlen["min"] || len(sharedsecret) > sharedsecretlen["max"] {
            exit(fmt.Sprintf("problem with sharedsecret: '%s' for label: '%s'", sharedsecret, label), 2)
        }

		template := []*pkcs11.Attribute{pkcs11.NewAttribute(pkcs11.CKA_LABEL, label), pkcs11.NewAttribute(pkcs11.CKA_CLASS, pkcs11.CKO_PRIVATE_KEY)}
		if e := p.FindObjectsInit(s.session, template); e != nil {
			panic(fmt.Sprintf("Failed to init: %s\n", e.Error()))
		}
		obj, b, e := p.FindObjects(s.session, 2)

		debug(fmt.Sprintf("Obj %v\n", obj))
		if e != nil {
			exit(fmt.Sprintf("Failed to find: %s %v\n", e.Error(), b), 2)
		}
		if e := p.FindObjectsFinal(s.session); e != nil {
			exit(fmt.Sprintf("Failed to finalize: %s\n", e.Error()), 2)
		}
		debug(fmt.Sprintf("found keys: %v\n", len(obj)))
		if len(obj) == 0 {
			exit(fmt.Sprintf("did not find a key with label '%s'", label), 2)
		}
		keymap[label] = aclmap{obj[0], sharedsecret, label}
	}

	fmt.Printf("hsm initialized new: %#v\n", keymap)

	sem <- s

	debug(fmt.Sprintf("sem: %v\n", len(sem)))
}

// Client authenticate/authorization
func authClient(sharedkey string, slot string, keylabel string, mech string) error {
	//  Check sharedkey
	//  Check slot nummer
	if slot != config["GOELEVEN_SLOT"] {
		return errors.New("Slot number does not match")
	}
	//  Check key aliases/label
	if _, present := keymap[keylabel]; !present {
		return errors.New(fmt.Sprintf("Key label does not match %s", keylabel))
	}

	if sharedkey != keymap[keylabel].sharedsecret {
		return errors.New(fmt.Sprintf("Client secret for label: '%s' does not match", keymap[keylabel].label))
	}

	//  Check key mech
	if mech != config["GOELEVEN_MECH"] {
		return errors.New("Mech does not match")
	}
	// client ok
	return nil
}

// TODO: Cleanup
// TODO: Documentation
// TODO: Error handling
/*
 * If error then send HTTP 500 to client and keep the server running
 *
 */
func handler(w http.ResponseWriter, r *http.Request) {

	fmt.Println("access attempt from:", r.RemoteAddr)
	ips := strings.Split(config["GOELEVEN_ALLOWEDIP"], ",")
	ip := strings.Split(r.RemoteAddr, ":")
	var allowed bool
	for _, v := range ips {
		allowed = allowed || ip[0] == v
	}

	if !allowed {
		fmt.Println("unauthorised access attempt from:", r.RemoteAddr)
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	var err error
	var validPath = regexp.MustCompile("^/(\\d+)/([a-zA-Z0-9\\.]+)/sign$")
	mSlot := validPath.FindStringSubmatch(r.URL.Path)[1]
	mKeyAlias := validPath.FindStringSubmatch(r.URL.Path)[2]

	defer r.Body.Close()
	body, _ := ioutil.ReadAll(r.Body)

	// Parse JSON
	//var b struct { Data,Mech string }
	var b map[string]interface{}
	err = json.Unmarshal(body, &b)
	if err != nil {
		http.Error(w, "Invalid input", 500)
		fmt.Printf("json.unmarshall: %v\n", err.Error())
		return
	}
	data, err := base64.StdEncoding.DecodeString(b["data"].(string))
	if err != nil {
		http.Error(w, "Invalid input", 500)
		fmt.Printf("DecodeString: %v\n", err.Error())
		return
	}

	// Client auth
	err = authClient(b["sharedkey"].(string), mSlot, mKeyAlias, b["mech"].(string))
	if err != nil {
		http.Error(w, "Invalid input", 500)
		fmt.Printf("authClient: %v\n", err.Error())
		return
	}

	sig, err, sessno := signing(data, keymap[mKeyAlias].handle)
	if err != nil {
		http.Error(w, "Invalid output", 500)
		fmt.Printf("signing: %v %v\n", err.Error(), sessno)
		return
	}
	sigs := base64.StdEncoding.EncodeToString(sig)
	type Res struct {
		Slot   string `json:"slot"`
		Mech   string `json:"mech"`
		Signed string `json:"signed"`
	}
	res := Res{mSlot, "mech", sigs}
	json, err := json.Marshal(res)
	if err != nil {
		http.Error(w, "Invalid output", 500)
		fmt.Printf("json.marshall: %v\n", err.Error())
		return
	}
	fmt.Fprintf(w, "%s\n\n", json)
}

// TODO: Cleanup
// TODO: Documentation
func inithsm(sessno int) Hsm {
	pguard.Lock()
	defer pguard.Unlock()
	slot, _ := strconv.ParseUint(config["GOELEVEN_SLOT"], 10, 32)

	fmt.Printf("slot: %v\n", slot)
	session, e := p.OpenSession(uint(slot), pkcs11.CKF_SERIAL_SESSION)

	if e != nil {
		panic(fmt.Sprintf("Failed to open session: %s\n", e.Error()))
	}

	p.Login(session, pkcs11.CKU_USER, config["GOELEVEN_SLOT_PASSWORD"])

	return Hsm{session, 0, time.Now(), sessno}
}

// TODO: Cleanup
// TODO: Documentation
func signing(data []byte, key pkcs11.ObjectHandle) ([]byte, error, int) {
	// Pop HSM struct from queue
	s := <-sem
	s.used++
	if s.used > 10000 || time.Now().Sub(s.started) > 1000*time.Second {
		p.Logout(s.session)
		p.CloseSession(s.session)
		//p.Finalize()
		//p.Destroy()
		s = inithsm(s.sessno)
	}
	fmt.Printf("hsm: %v %v\n", s, key)
	//p.SignInit(s.session, []*pkcs11.Mechanism{pkcs11.NewMechanism(pkcs11.CKM_SHA256_RSA_PKCS, nil)}, key)
	p.SignInit(s.session, []*pkcs11.Mechanism{pkcs11.NewMechanism(pkcs11.CKM_RSA_PKCS, nil)}, key)
	sig, err := p.Sign(s.session, data)
	fmt.Printf("err: %v\n", err)

	// Push HSM struct back on queue
	sem <- s
	return sig, nil, s.sessno
}

// Utils

func debug(messages string) {
	if config["GOELEVEN_DEBUG"] == "true" {
		fmt.Print(messages)
	}
}

// Standard function to test for debug mode
func isdebug() bool {
	if config["GOELEVEN_DEBUG"] == "true" {
		return true
	} else {
		return false
	}
}

func exit(messages string, errorCode int) {
	// Exit code and messages based on Nagios plugin return codes (https://nagios-plugins.org/doc/guidelines.html#AEN78)
	var prefix = map[int]string{0: "OK", 1: "Warning", 2: "Critical", 3: "Unknown"}

	// Catch all unknown errorCode and convert them to Unknown
	if errorCode < 0 || errorCode > 3 {
		errorCode = 3
	}

	fmt.Printf("%s %s\n", prefix[errorCode], messages)
	os.Exit(errorCode)
}
<|endoftext|>


<commit_before>// sokoban solver, work in progress
package main

import (
	"fmt"
	"github.com/bertbaron/solve"
	"math"
	"sort"
	"strings"
	"os"
	"log"
	"runtime/pprof"
	"time"
)

const (
	floor  byte = 0
	wall   byte = 1
	box    byte = 2
	goal   byte = 4
	player byte = 8
)

var chars = map[rune]byte{
	' ': floor,
	'#': wall,
	'$': box,
	'.': goal,
	'@': player,
	'+': player | goal,
	'*': box | goal}

var reverse = map[byte]rune{
	floor:         ' ',
	wall:          '#',
	box:           '$',
	goal:          '.',
	player:        '@',
	player | goal: '+',
	box | goal:    '*'}

// -------- main problem. We only expose the states in which a block is pushed though to limit the search space
//          for the main search.
type sokoban struct {
	// the static world, without player and boxes
	world []byte
	// sorted list of goal positions
	goals  []uint16
	width  int
	height int
}

type mainstate struct {
	// sorted list of box positions
	boxes    []uint16
	position int
	cost     int
}

// returns the index of position in the sorted list of positions. Returns -1 if the position is not found
func binarySearch(positions []uint16, position int) int {
	idx := sort.Search(len(positions), func(i int) bool { return positions[i] >= uint16(position) })
	if idx < len(positions) && positions[idx] == uint16(position) {
		return idx
	}
	return -1
}

func valueOf(s sokoban, m mainstate, position int) byte {
	boxidx := binarySearch(m.boxes, position)
	var additional byte = 0
	if m.position == position {
		additional |= player
	}
	if boxidx >= 0 {
		additional |= box
	}
	return s.world[position] | additional
}

func isEmpty(value byte) bool {
	return value&(wall|box) == 0
}

func isBox(value byte) bool {
	return value&box != 0
}

func isWall(value byte) bool {
	return value&wall != 0
}

func print(s sokoban, m mainstate) {
	for position := range s.world {
		fmt.Print(string(reverse[valueOf(s, m, position)]))
		if position%s.width == s.width-1 {
			fmt.Println()
		}
	}
}

func (s mainstate) Cost(ctx solve.Context) float64 {
	return float64(s.cost)
}

func abs(value int) int {
	if value < 0 {
		return -value
	}
	return value
}

// total manhattan distance for all boxes towards their nearest target - O(boxcount^2)
func minimalManhattan(world sokoban, s mainstate) int {
	total := 0
	for _, box := range s.boxes {
		min := math.MaxInt32
		bx, by := int(box)%world.width, int(box)/world.width
		for _, goal := range world.goals {
			gx, gy := int(goal)%world.width, int(goal)/world.width
			md := abs(gx-bx) + abs(gy-by)
			if md < min {
				min = md
			}
		}
		total += min
	}
	return total
}

// number of boxes not on a goal position - O(boxcount)
func displaced(world sokoban, s mainstate) int {
	displaced := len(s.boxes)
	for _, box := range s.boxes {
		displaced -= int(world.world[box] & goal) >> 2
	}
	return displaced
}

func (s mainstate) Heuristic(ctx solve.Context) float64 {
	world := ctx.Custom.(sokoban)
	//h := displaced(world, s)
	h := minimalManhattan(world, s)
	return float64(h)
}

func (s mainstate) IsGoal(ctx solve.Context) bool {
	for i, value := range ctx.Custom.(sokoban).goals {
		if s.boxes[i] != value {
			return false
		}
	}
	return true
}

func (s mainstate) Expand(ctx solve.Context) []solve.State {
	world := ctx.Custom.(sokoban)
	targets := make([]int, 0)
	for _, box := range s.boxes {
		left := isEmpty(valueOf(world, s, int(box)-1))
		right := isEmpty(valueOf(world, s, int(box)+1))
		up := isEmpty(valueOf(world, s, int(box)-world.width))
		down := isEmpty(valueOf(world, s, int(box)+world.width))
		if left && right {
			targets = append(targets, int(box)-1)
			targets = append(targets, int(box)+1)
		}
		if up && down {
			targets = append(targets, int(box)-world.width)
			targets = append(targets, int(box)+world.width)
		}
	}
	paths := getWalkMoves(world, s, targets)

	var children []solve.State
	for _, path := range paths {
		p := path.position
		for _, dir := range [...]int{-1, 1, -world.width, world.width} {
			if isBox(valueOf(world, s, p+dir)) && isEmpty(valueOf(world, s, p+2*dir)) {
				child := push(world, s, p, dir, path.cost)
				if child != nil {
					//print(world, child)
					children = append(children, *child)
				}
			}
		}
	}
	return children
}

func push(world sokoban, s mainstate, position int, direction int, cost int) *mainstate {
	newposition := position + direction
	newbox := uint16(position + direction*2)
	newboxes := make([]uint16, len(s.boxes))
	copy(newboxes, s.boxes)
	idx := binarySearch(newboxes, newposition)
	newboxes[idx] = newbox

	n := len(newboxes)
	// insertion sort to keep boxes sorted, only needed when moving up or down
	if direction < -1 {
		for idx > 0 && newboxes[idx-1] > newbox {
			newboxes[idx-1], newboxes[idx] = newboxes[idx], newboxes[idx-1]
			idx--
		}
	}
	if direction > 1 {
		for idx < n-1 && newboxes[idx+1] < newbox {
			newboxes[idx+1], newboxes[idx] = newboxes[idx], newboxes[idx+1]
			idx++
		}
	}
	newState := mainstate{newboxes, newposition, s.cost + cost + 1}
	if deadEnd(world, newState, int(newbox)) {
		return nil
	}
	return &newState
}

// looks in a 3x3 pattern around the box position if this is a dead end
func deadEnd(world sokoban, s mainstate, position int) bool {
	if world.world[position]&goal != 0 {
		return false // box is on a goal position
	}

	// corner walls
	lu := world.world[position-1-world.width]&wall != 0
	ru := world.world[position+1-world.width]&wall != 0
	ld := world.world[position-1+world.width]&wall != 0
	rd := world.world[position+1+world.width]&wall != 0

	// orthogonal walls or boxes
	uvalue := valueOf(world, s, position-world.width)
	dvalue := valueOf(world, s, position+world.width)
	lvalue := valueOf(world, s, position-1)
	rvalue := valueOf(world, s, position+1)

	// direction is blocked if it is a wall or a block that is sideways blocked by a wall
	u := isWall(uvalue) || (isBox(uvalue) && (lu || ru))
	d := isWall(dvalue) || (isBox(dvalue) && (ld || rd))
	l := isWall(lvalue) || (isBox(lvalue) && (lu || ld))
	r := isWall(rvalue) || (isBox(rvalue) && (ru || rd))

	return u && r || r && d || d && l || l && u
}

func assertOrdered(positions []uint16) {
	for i, value := range positions[1:] {
		if value <= positions[i] {
			panic(fmt.Sprintf("Ordering invariant violated: %v", positions))
		}
	}
}

// -------------- Sub problem for moving the player to all positions in which a box can be moved -----------

type walkcontext struct {
	// the static world, without player but with boxes because we don't move them here
	world   []byte
	targets []int
	width   int
}

type walkstate struct {
	position int
	cost     int
}

func (s walkstate) Cost(ctx solve.Context) float64 {
	return float64(s.cost)
}

func (s walkstate) Heuristic(ctx solve.Context) float64 {
	return 0
}

func (s walkstate) IsGoal(ctx solve.Context) bool {
	wc := ctx.Custom.(walkcontext)
	for _, goal := range wc.targets {
		if s.position == goal {
			return true
		}
	}
	return false
}

func (s walkstate) Expand(ctx solve.Context) []solve.State {
	var children []solve.State
	wc := ctx.Custom.(walkcontext)
	children = s.addIfValid(children, s.position-1, wc)
	children = s.addIfValid(children, s.position+1, wc)
	children = s.addIfValid(children, s.position-wc.width, wc)
	children = s.addIfValid(children, s.position+wc.width, wc)
	return children
}

func (s walkstate) addIfValid(children []solve.State, newPosition int, wc walkcontext) []solve.State {
	if wc.world[newPosition]&(wall|box) == 0 {
		return append(children, walkstate{newPosition, s.cost + 1})
	}
	return children
}

type walkstateMap map[int]float64

func (c walkstateMap) Get(state solve.State) (float64, bool) {
	value, ok := c[state.(walkstate).position]
	return value, ok
}

func (c walkstateMap) Put(state solve.State, value float64) {
	c[state.(walkstate).position] = value
}

func (c *walkstateMap) Clear() {
	*c = make(walkstateMap)
}

func getWalkMoves(wc sokoban, s mainstate, targets []int) []walkstate {
	context := walkcontext{wc.world, targets, wc.width}
	context.world = make([]byte, len(wc.world))
	copy(context.world, wc.world)
	for _, boxposition := range s.boxes {
		context.world[boxposition] |= box
	}
	rootstate := walkstate{s.position, 0}
	var wsMap walkstateMap
	solver := solve.NewSolver(rootstate).
		Context(context).
		Constraint(solve.CheapestPathConstraint(&wsMap)).
		Algorithm(solve.BreadthFirst)
	solutions := make([]walkstate, 0)
	for solution := range solver.SolveAll() {
		solutions = append(solutions, solution.GoalState().(walkstate))
	}
	return solutions
}

func parse(level string) (sokoban, mainstate) {
	width := 0
	lines := strings.Split(level, "\n")
	height := len(lines)
	for _, line := range lines {
		if len(line) > width {
			width = len(line)
		}
	}
	var c sokoban
	var s mainstate
	c.width = width
	c.height = height

	c.world = make([]byte, width*height)
	c.goals = make([]uint16, 0)
	s.boxes = make([]uint16, 0)
	for y, row := range lines {
		for x, raw := range row {
			position := y*width + x
			if value, ok := chars[raw]; ok {
				c.world[position] = value &^ player &^ box
				if value&player != 0 {
					s.position = position
				}
				if value&goal != 0 {
					c.goals = append(c.goals, uint16(position))
				}
				if value&box != 0 {
					s.boxes = append(s.boxes, uint16(position))
				}
			} else {
				panic(fmt.Sprintf("Invalid level format, character %v is not valid", value))
			}
		}
	}
	return c, s
}

// For cheapest path constraint
type cpMap map[string]float64

func key(state solve.State) string {
	// nasty hack, but string seems to be the only variable-size type
	// supported as map key. Would love to be able to use slices directly
	s := state.(mainstate)
	runes := make([]rune, len(s.boxes) + 1)
	runes[0] = rune(s.position)
	for i, box := range s.boxes {
		runes[i+1] = rune(box)
	}
	return string(runes)
}

func (c cpMap) Get(state solve.State) (value float64, ok bool) {
	value, ok = c[key(state)]
	return
}

func (c cpMap) Put(state solve.State, value float64) {
	c[key(state)] = value
}

func (c *cpMap) Clear() {
	*c = make(cpMap)
}

func cheapestPathConstraint() solve.Constraint {
	var m cpMap
	return solve.CheapestPathConstraint(&m)
}

var simpleLevel = `
########
#     @#
#      #
# $ # .#
#  # # #
########`

var level = `
   ####
####  ##
#   $  #
#  *** #
#  . . ##
## * *  #
 ##***  #
  # $ ###
  # @ #
  #####`

func main() {
	f, err := os.Create("cpu.prof")
	if err != nil {
		log.Fatal(err)
	}
	pprof.StartCPUProfile(f)
	defer pprof.StopCPUProfile()

	world, root := parse(level)
	print(world, root)
	start := time.Now()
	result := solve.NewSolver(root).
		Context(world).
		Algorithm(solve.Astar).
		Constraint(cheapestPathConstraint()).
		//Limit(40).
		Solve()
	fmt.Printf("Time: %.1f seconds\n", time.Since(start).Seconds())
	if result.Solved() {
		fmt.Printf("Result:\n ")
		for _, state := range result.Solution {
			print(world, state.(mainstate))
		}
		fmt.Printf("Solved in %d moves\n", int(result.GoalState().(mainstate).cost))
	}
	fmt.Printf("visited %v main nodes\n", result.Visited)
}
<commit_msg>Tuning sokoban example<commit_after>// sokoban solver, work in progress
package main

import (
	"fmt"
	"github.com/bertbaron/solve"
	"math"
	"sort"
	"strings"
	"os"
	"log"
	"runtime/pprof"
	"time"
)

const (
	floor  byte = 0
	wall   byte = 1
	box    byte = 2
	goal   byte = 4
	player byte = 8
)

var chars = map[rune]byte{
	' ': floor,
	'#': wall,
	'$': box,
	'.': goal,
	'@': player,
	'+': player | goal,
	'*': box | goal}

var reverse = map[byte]rune{
	floor:         ' ',
	wall:          '#',
	box:           '$',
	goal:          '.',
	player:        '@',
	player | goal: '+',
	box | goal:    '*'}

// -------- main problem. We only expose the states in which a block is pushed though to limit the search space
//          for the main search.
type sokoban struct {
	// the static world, without player and boxes
	world []byte
	// sorted list of goal positions
	goals  []uint16
	width  int
	height int
}

type mainstate struct {
	// sorted list of box positions
	boxes    []uint16
	position int
	cost     int
}

// returns the index of position in the sorted list of positions. Returns -1 if the position is not found
func binarySearch(positions []uint16, position int) int {
	idx := sort.Search(len(positions), func(i int) bool { return positions[i] >= uint16(position) })
	if idx < len(positions) && positions[idx] == uint16(position) {
		return idx
	}
	return -1
}

func valueOf(s sokoban, m mainstate, position int) byte {
	boxidx := binarySearch(m.boxes, position)
	var additional byte = 0
	if m.position == position {
		additional |= player
	}
	if boxidx >= 0 {
		additional |= box
	}
	return s.world[position] | additional
}

func isEmpty(value byte) bool {
	return value&(wall|box) == 0
}

func isBox(value byte) bool {
	return value&box != 0
}

func isWall(value byte) bool {
	return value&wall != 0
}

func print(s sokoban, m mainstate) {
	for position := range s.world {
		fmt.Print(string(reverse[valueOf(s, m, position)]))
		if position%s.width == s.width-1 {
			fmt.Println()
		}
	}
}

func (s mainstate) Cost(ctx solve.Context) float64 {
	return float64(s.cost)
}

func abs(value int) int {
	if value < 0 {
		return -value
	}
	return value
}

// total manhattan distance for all boxes towards their nearest target - O(boxcount^2)
func minimalManhattan(world sokoban, s mainstate) int {
	total := 0
	for _, box := range s.boxes {
		min := math.MaxInt32
		bx, by := int(box)%world.width, int(box)/world.width
		for _, goal := range world.goals {
			gx, gy := int(goal)%world.width, int(goal)/world.width
			md := abs(gx-bx) + abs(gy-by)
			if md < min {
				min = md
			}
		}
		total += min
	}
	return total
}

// number of boxes not on a goal position - O(boxcount)
func displaced(world sokoban, s mainstate) int {
	displaced := len(s.boxes)
	for _, box := range s.boxes {
		displaced -= int(world.world[box] & goal) >> 2
	}
	return displaced
}

func (s mainstate) Heuristic(ctx solve.Context) float64 {
	world := ctx.Custom.(sokoban)
	//h := displaced(world, s)
	h := minimalManhattan(world, s)
	return float64(h)
}

func (s mainstate) IsGoal(ctx solve.Context) bool {
	for i, value := range ctx.Custom.(sokoban).goals {
		if s.boxes[i] != value {
			return false
		}
	}
	return true
}

func (s mainstate) Expand(ctx solve.Context) []solve.State {
	world := ctx.Custom.(sokoban)
	targets := make([]int, 0)
	for _, box := range s.boxes {
		left := isEmpty(valueOf(world, s, int(box)-1))
		right := isEmpty(valueOf(world, s, int(box)+1))
		up := isEmpty(valueOf(world, s, int(box)-world.width))
		down := isEmpty(valueOf(world, s, int(box)+world.width))
		if left && right {
			targets = append(targets, int(box)-1)
			targets = append(targets, int(box)+1)
		}
		if up && down {
			targets = append(targets, int(box)-world.width)
			targets = append(targets, int(box)+world.width)
		}
	}
	paths := getWalkMoves(world, s, targets)

	var children []solve.State
	for _, path := range paths {
		p := path.position
		for _, dir := range [...]int{-1, 1, -world.width, world.width} {
			if isBox(valueOf(world, s, p+dir)) && isEmpty(valueOf(world, s, p+2*dir)) {
				child := push(world, s, p, dir, path.cost)
				if child != nil {
					//print(world, child)
					children = append(children, *child)
				}
			}
		}
	}
	return children
}

func push(world sokoban, s mainstate, position int, direction int, cost int) *mainstate {
	newposition := position + direction
	newbox := uint16(position + direction*2)
	newboxes := make([]uint16, len(s.boxes))
	copy(newboxes, s.boxes)
	idx := binarySearch(newboxes, newposition)
	newboxes[idx] = newbox

	n := len(newboxes)
	// insertion sort to keep boxes sorted, only needed when moving up or down
	if direction < -1 {
		for idx > 0 && newboxes[idx-1] > newbox {
			newboxes[idx-1], newboxes[idx] = newboxes[idx], newboxes[idx-1]
			idx--
		}
	}
	if direction > 1 {
		for idx < n-1 && newboxes[idx+1] < newbox {
			newboxes[idx+1], newboxes[idx] = newboxes[idx], newboxes[idx+1]
			idx++
		}
	}
	newState := mainstate{newboxes, newposition, s.cost + cost + 1}
	if deadEnd(world, newState, int(newbox)) {
		return nil
	}
	return &newState
}

// looks in a 3x3 pattern around the box position if this is a dead end
func deadEnd(world sokoban, s mainstate, position int) bool {
	if world.world[position]&goal != 0 {
		return false // box is on a goal position
	}

	// corner walls
	lu := world.world[position-1-world.width]&wall != 0
	ru := world.world[position+1-world.width]&wall != 0
	ld := world.world[position-1+world.width]&wall != 0
	rd := world.world[position+1+world.width]&wall != 0

	// orthogonal walls or boxes
	uvalue := valueOf(world, s, position-world.width)
	dvalue := valueOf(world, s, position+world.width)
	lvalue := valueOf(world, s, position-1)
	rvalue := valueOf(world, s, position+1)

	// direction is blocked if it is a wall or a block that is sideways blocked by a wall
	u := isWall(uvalue) || (isBox(uvalue) && (lu || ru))
	d := isWall(dvalue) || (isBox(dvalue) && (ld || rd))
	l := isWall(lvalue) || (isBox(lvalue) && (lu || ld))
	r := isWall(rvalue) || (isBox(rvalue) && (ru || rd))

	return u && r || r && d || d && l || l && u
}

func assertOrdered(positions []uint16) {
	for i, value := range positions[1:] {
		if value <= positions[i] {
			panic(fmt.Sprintf("Ordering invariant violated: %v", positions))
		}
	}
}

// -------------- Sub problem for moving the player to all positions in which a box can be moved -----------

type walkcontext struct {
	// the static world, without player but with boxes because we don't move them here
	world   []byte
	targets []int
	width   int
}

type walkstate struct {
	position int
	cost     int
}

func (s walkstate) Cost(ctx solve.Context) float64 {
	return float64(s.cost)
}

func (s walkstate) Heuristic(ctx solve.Context) float64 {
	return 0
}

func (s walkstate) IsGoal(ctx solve.Context) bool {
	wc := ctx.Custom.(walkcontext)
	for _, goal := range wc.targets {
		if s.position == goal {
			return true
		}
	}
	return false
}

func (s walkstate) Expand(ctx solve.Context) []solve.State {
	var children []solve.State
	wc := ctx.Custom.(walkcontext)
	children = s.addIfValid(children, s.position-1, wc)
	children = s.addIfValid(children, s.position+1, wc)
	children = s.addIfValid(children, s.position-wc.width, wc)
	children = s.addIfValid(children, s.position+wc.width, wc)
	return children
}

func (s walkstate) addIfValid(children []solve.State, newPosition int, wc walkcontext) []solve.State {
	if wc.world[newPosition]&(wall|box) == 0 {
		return append(children, walkstate{newPosition, s.cost + 1})
	}
	return children
}

// Example of a CPMap implementation based on a slice
type walkstateMap []float64

func (c walkstateMap) Get(state solve.State) (float64, bool) {
	value := c[state.(walkstate).position];
	return value, value >= 0
}

func (c walkstateMap) Put(state solve.State, value float64) {
	c[state.(walkstate).position] = value
}

func (c walkstateMap) Clear() {
	for i := range c {
		c[i] = -1
	}
}

func getWalkMoves(wc sokoban, s mainstate, targets []int) []walkstate {
	context := walkcontext{wc.world, targets, wc.width}
	context.world = make([]byte, len(wc.world))
	copy(context.world, wc.world)
	for _, boxposition := range s.boxes {
		context.world[boxposition] |= box
	}
	rootstate := walkstate{s.position, 0}
	wsMap := make(walkstateMap, len(wc.world))
	solver := solve.NewSolver(rootstate).
		Context(context).
		Constraint(solve.CheapestPathConstraint(wsMap)).
		Algorithm(solve.BreadthFirst)
	solutions := make([]walkstate, 0)
	for solution := solver.Solve(); solution.Solved(); solution = solver.Solve() {
		solutions = append(solutions, solution.GoalState().(walkstate))
	}
	return solutions
}

func parse(level string) (sokoban, mainstate) {
	width := 0
	lines := strings.Split(level, "\n")
	height := len(lines)
	for _, line := range lines {
		if len(line) > width {
			width = len(line)
		}
	}
	var c sokoban
	var s mainstate
	c.width = width
	c.height = height

	c.world = make([]byte, width*height)
	c.goals = make([]uint16, 0)
	s.boxes = make([]uint16, 0)
	for y, row := range lines {
		for x, raw := range row {
			position := y*width + x
			if value, ok := chars[raw]; ok {
				c.world[position] = value &^ player &^ box
				if value&player != 0 {
					s.position = position
				}
				if value&goal != 0 {
					c.goals = append(c.goals, uint16(position))
				}
				if value&box != 0 {
					s.boxes = append(s.boxes, uint16(position))
				}
			} else {
				panic(fmt.Sprintf("Invalid level format, character %v is not valid", value))
			}
		}
	}
	return c, s
}

// For cheapest path constraint
type cpMap map[string]float64

func key(state solve.State) string {
	// nasty hack, but string seems to be the only variable-size type
	// supported as map key. Would love to be able to use slices directly
	s := state.(mainstate)
	runes := make([]rune, len(s.boxes) + 1)
	runes[0] = rune(s.position)
	for i, box := range s.boxes {
		runes[i+1] = rune(box)
	}
	return string(runes)
}

func (c cpMap) Get(state solve.State) (value float64, ok bool) {
	value, ok = c[key(state)]
	return
}

func (c cpMap) Put(state solve.State, value float64) {
	c[key(state)] = value
}

func (c *cpMap) Clear() {
	*c = make(cpMap)
}

func cheapestPathConstraint() solve.Constraint {
	var m cpMap
	return solve.CheapestPathConstraint(&m)
}

var simpleLevel = `
########
#     @#
#      #
# $ # .#
#  # # #
########`

var level = `
   ####
####  ##
#   $  #
#  *** #
#  . . ##
## * *  #
 ##***  #
  # $ ###
  # @ #
  #####`

func main() {
	f, err := os.Create("cpu.prof")
	if err != nil {
		log.Fatal(err)
	}
	pprof.StartCPUProfile(f)
	defer pprof.StopCPUProfile()

	world, root := parse(level)
	print(world, root)
	start := time.Now()
	result := solve.NewSolver(root).
		Context(world).
		Algorithm(solve.IDAstar).
		Constraint(cheapestPathConstraint()).
		Limit(38).
		Solve()
	fmt.Printf("Time: %.1f seconds\n", time.Since(start).Seconds())
	if result.Solved() {
		fmt.Printf("Result:\n ")
		for _, state := range result.Solution {
			print(world, state.(mainstate))
		}
		fmt.Printf("Solved in %d moves\n", int(result.GoalState().(mainstate).cost))
	}
	fmt.Printf("visited %v main nodes\n", result.Visited)
}
<|endoftext|>


<commit_before>// learn
package learn

import (
	"fmt"
	"runtime"
	//"strings"
)
/**
*全局变量定义 但不一定使用
* 局部变量 定义必须使用
*
*/
var (
	a  int
	b  int
	d  string

)
/*
*多变量赋值： 首先计算相关之的值，从左到右依次赋值
* 常量 必须是编译期 确定的值  数字  字符串  布尔值
*  len  cap  unsafe.Sizeof
*/
var (
	q int32
	m string
	r  []int
)
func Test(){

}
func  TestVar(){
    fmt.Println("------------------------------2----------------------------!")
	d:="fdfdf"
	//使用语言自带的内置方法
	fmt.Println(len(d))
	fmt.Println(d)
	v:=make([]int, 2,8)
	for index , v := range v{
		fmt.Println(index , v )
	}
	//基本数据类型之间相互转换
	fmt.Println(string([]byte(d)))
}

func TestRuntime() {
	fmt.Println("------------------------------1----------------------------!")
	fmt.Println("Hello learnner!")
	//runtime.Gosched
	fmt.Println(runtime.GOMAXPROCS(runtime.NumCPU()))
	runtime.LockOSThread()
	fmt.Println(runtime.MemProfileRate)
	fmt.Println(runtime.NumCgoCall())
	fmt.Println(runtime.NumGoroutine())
	fmt.Println(runtime.Version())
}

<commit_msg>20150607<commit_after>// learn
package learn

import (
	"fmt"
	"runtime"
	//"strings"
)

/**
*全局变量【包含常量】定义 但不一定使用  符合golang语法
* 局部变量 定义必须使用
*
 */
var (
	a int
	b int
	d string
)

/*
*多变量赋值： 首先计算相关之的值，从左到右依次赋值
* 常量: 必须是编译期 确定的值  数字  字符串  布尔值 len  cap  unsafe.Sizeof
 */
var (
	q int32
	m string
	r []int
)
/*
  字符串的操作说明：
     1.子字符串指向原来的字节数组
     2.修改了指针和长度属性
  字符串的转换需要重新分配内存地址以及复制字节数组
**/
func Test() {
	fmt.Println("------------------------------3----------------------------!")
	var m string = "张伯雨"
	fmt.Println(m[:])
	fmt.Println(m[0:3])
	fmt.Println(m[3:6])
	//输出变量m的数据类型
	fmt.Printf("%T\r\n", m)
	//输出字符串内存地址
	fmt.Println("\r\n", &m)
/*
golang 中的for 循环分为两种情况：
   1.基于字节遍历 使用字符串的len方法 使用的是字节遍历
   2.基于字符遍历 使用字符串的range 使用的是字符遍历
*/
	for i:=0;i<len(m);i++{
		fmt.Printf("%d , %c \r\n",i,m[i])
	 }
	fmt.Println("------------------------------33----------------------------!")
	for  index , v :=range m{
		fmt.Printf("%d , %c \r\n",index,v)
	}
}
func TestVar() {
	fmt.Println("------------------------------2----------------------------!")
	d := "fdfdf"
	//使用语言自带的内置方法
	fmt.Println(len(d))
	fmt.Println(d)
	v := make([]int, 2, 8)
	for index, v := range v {
		fmt.Println(index, v)
	}
	//基本数据类型之间相互转换
	fmt.Println(string([]byte(d)))
}

func TestRuntime() {
	fmt.Println("------------------------------1----------------------------!")
	fmt.Println("Hello learnner!")
	//runtime.Gosched
	fmt.Println(runtime.GOMAXPROCS(runtime.NumCPU()))
	runtime.LockOSThread()
	fmt.Println(runtime.MemProfileRate)
	fmt.Println(runtime.NumCgoCall())
	fmt.Println(runtime.NumGoroutine())
	fmt.Println(runtime.Version())
}
<|endoftext|>


<commit_before>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build windows

package mgr_test

import (
	"os"
	"path/filepath"
	"sort"
	"strings"
	"syscall"
	"testing"
	"time"

	"golang.org/x/sys/windows/svc/mgr"
)

func TestOpenLanManServer(t *testing.T) {
	m, err := mgr.Connect()
	if err != nil {
		if errno, ok := err.(syscall.Errno); ok && errno == syscall.ERROR_ACCESS_DENIED {
			t.Skip("Skipping test: we don't have rights to manage services.")
		}
		t.Fatalf("SCM connection failed: %s", err)
	}
	defer m.Disconnect()

	s, err := m.OpenService("LanmanServer")
	if err != nil {
		t.Fatalf("OpenService(lanmanserver) failed: %s", err)
	}
	defer s.Close()

	_, err = s.Config()
	if err != nil {
		t.Fatalf("Config failed: %s", err)
	}
}

func install(t *testing.T, m *mgr.Mgr, name, exepath string, c mgr.Config) {
	// Sometimes it takes a while for the service to get
	// removed after previous test run.
	for i := 0; ; i++ {
		s, err := m.OpenService(name)
		if err != nil {
			break
		}
		s.Close()

		if i > 10 {
			t.Fatalf("service %s already exists", name)
		}
		time.Sleep(300 * time.Millisecond)
	}

	s, err := m.CreateService(name, exepath, c)
	if err != nil {
		t.Fatalf("CreateService(%s) failed: %v", name, err)
	}
	defer s.Close()
}

func depString(d []string) string {
	if len(d) == 0 {
		return ""
	}
	for i := range d {
		d[i] = strings.ToLower(d[i])
	}
	ss := sort.StringSlice(d)
	ss.Sort()
	return strings.Join([]string(ss), " ")
}

func testConfig(t *testing.T, s *mgr.Service, should mgr.Config) mgr.Config {
	is, err := s.Config()
	if err != nil {
		t.Fatalf("Config failed: %s", err)
	}
	if should.DelayedAutoStart != is.DelayedAutoStart {
		t.Fatalf("config mismatch: DelayedAutoStart is %v, but should have %v", is.DelayedAutoStart, should.DelayedAutoStart)
	}
	if should.DisplayName != is.DisplayName {
		t.Fatalf("config mismatch: DisplayName is %q, but should have %q", is.DisplayName, should.DisplayName)
	}
	if should.StartType != is.StartType {
		t.Fatalf("config mismatch: StartType is %v, but should have %v", is.StartType, should.StartType)
	}
	if should.Description != is.Description {
		t.Fatalf("config mismatch: Description is %q, but should have %q", is.Description, should.Description)
	}
	if depString(should.Dependencies) != depString(is.Dependencies) {
		t.Fatalf("config mismatch: Dependencies is %v, but should have %v", is.Dependencies, should.Dependencies)
	}
	return is
}

func testRecoveryActions(t *testing.T, s *mgr.Service, should []mgr.RecoveryAction) {
	is, err := s.RecoveryActions()
	if err != nil {
		t.Fatalf("RecoveryActions failed: %s", err)
	}
	if len(should) != len(is) {
		t.Errorf("recovery action mismatch: contains %v actions, but should have %v", len(is), len(should))
	}
	for i, _ := range is {
		if should[i].Type != is[i].Type {
			t.Errorf("recovery action mismatch: Type is %v, but should have %v", is[i].Type, should[i].Type)
		}
		if should[i].Delay != is[i].Delay {
			t.Errorf("recovery action mismatch: Delay is %v, but should have %v", is[i].Delay, should[i].Delay)
		}
	}
}

func testResetPeriod(t *testing.T, s *mgr.Service, should uint32) {
	is, err := s.ResetPeriod()
	if err != nil {
		t.Fatalf("ResetPeriod failed: %s", err)
	}
	if should != is {
		t.Errorf("reset period mismatch: reset period is %v, but should have %v", is, should)
	}
}

func testSetRecoveryActions(t *testing.T, s *mgr.Service) {
	r := []mgr.RecoveryAction{
		mgr.RecoveryAction{
			Type:  mgr.NoAction,
			Delay: 60000 * time.Millisecond,
		},
		mgr.RecoveryAction{
			Type:  mgr.ServiceRestart,
			Delay: 4 * time.Minute,
		},
		mgr.RecoveryAction{
			Type:  mgr.ServiceRestart,
			Delay: time.Minute,
		},
		mgr.RecoveryAction{
			Type:  mgr.RunCommand,
			Delay: 4000 * time.Millisecond,
		},
	}

	// 4 recovery actions with reset period
	err := s.SetRecoveryActions(r, uint32(10000))
	if err != nil {
		t.Fatalf("SetRecoveryActions failed: %v", err)
	}
	testRecoveryActions(t, s, r)
	testResetPeriod(t, s, uint32(10000))

	// Infinite reset period
	err = s.SetRecoveryActions(r, syscall.INFINITE)
	if err != nil {
		t.Fatalf("SetRecoveryActions failed: %v", err)
	}
	testRecoveryActions(t, s, r)
	testResetPeriod(t, s, syscall.INFINITE)

	// nil recovery actions
	err = s.SetRecoveryActions(nil, 0)
	if err.Error() != "recoveryActions cannot be nil" {
		t.Fatalf("SetRecoveryActions failed with unexpected error message of %q", err)
	}

	// Delete all recovery actions and reset period
	err = s.ResetRecoveryActions()
	if err != nil {
		t.Fatalf("ResetRecoveryActions failed: %v", err)
	}
	testRecoveryActions(t, s, nil)
	testResetPeriod(t, s, 0)
}

func testRebootMessage(t *testing.T, s *mgr.Service, should string) {
	err := s.SetRebootMessage(should)
	if err != nil {
		t.Fatalf("SetRebootMessage failed: %v", err)
	}
	is, err := s.RebootMessage()
	if err != nil {
		t.Fatalf("RebootMessage failed: %v", err)
	}
	if should != is {
		t.Errorf("reboot message mismatch: message is %q, but should have %q", is, should)
	}
}

func testRecoveryCommand(t *testing.T, s *mgr.Service, should string) {
	err := s.SetRecoveryCommand(should)
	if err != nil {
		t.Fatalf("SetRecoveryCommand failed: %v", err)
	}
	is, err := s.RecoveryCommand()
	if err != nil {
		t.Fatalf("RecoveryCommand failed: %v", err)
	}
	if should != is {
		t.Errorf("recovery command mismatch: command is %q, but should have %q", is, should)
	}
}

func remove(t *testing.T, s *mgr.Service) {
	err := s.Delete()
	if err != nil {
		t.Fatalf("Delete failed: %s", err)
	}
}

func TestMyService(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode - it modifies system services")
	}

	const name = "myservice"

	m, err := mgr.Connect()
	if err != nil {
		if errno, ok := err.(syscall.Errno); ok && errno == syscall.ERROR_ACCESS_DENIED {
			t.Skip("Skipping test: we don't have rights to manage services.")
		}
		t.Fatalf("SCM connection failed: %s", err)
	}
	defer m.Disconnect()

	c := mgr.Config{
		StartType:    mgr.StartDisabled,
		DisplayName:  "my service",
		Description:  "my service is just a test",
		Dependencies: []string{"LanmanServer", "W32Time"},
	}

	exename := os.Args[0]
	exepath, err := filepath.Abs(exename)
	if err != nil {
		t.Fatalf("filepath.Abs(%s) failed: %s", exename, err)
	}

	install(t, m, name, exepath, c)

	s, err := m.OpenService(name)
	if err != nil {
		t.Fatalf("service %s is not installed", name)
	}
	defer s.Close()

	c.BinaryPathName = exepath
	c = testConfig(t, s, c)

	c.StartType = mgr.StartManual
	err = s.UpdateConfig(c)
	if err != nil {
		t.Fatalf("UpdateConfig failed: %v", err)
	}

	testConfig(t, s, c)

	c.StartType = mgr.StartAutomatic
	c.DelayedAutoStart = true
	err = s.UpdateConfig(c)
	if err != nil {
		t.Fatalf("UpdateConfig failed: %v", err)
	}

	testConfig(t, s, c)

	svcnames, err := m.ListServices()
	if err != nil {
		t.Fatalf("ListServices failed: %v", err)
	}
	var myserviceIsInstalled bool
	for _, sn := range svcnames {
		if sn == name {
			myserviceIsInstalled = true
			break
		}
	}
	if !myserviceIsInstalled {
		t.Errorf("ListServices failed to find %q service", name)
	}

	testSetRecoveryActions(t, s)
	testRebootMessage(t, s, "myservice failed")
	testRebootMessage(t, s, "") // delete reboot message
	testRecoveryCommand(t, s, "sc query myservice")
	testRecoveryCommand(t, s, "") // delete recovery command

	remove(t, s)
}
<commit_msg>windows/svc/mgr: rename service in TestMyService<commit_after>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build windows

package mgr_test

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"syscall"
	"testing"
	"time"

	"golang.org/x/sys/windows/svc/mgr"
)

func TestOpenLanManServer(t *testing.T) {
	m, err := mgr.Connect()
	if err != nil {
		if errno, ok := err.(syscall.Errno); ok && errno == syscall.ERROR_ACCESS_DENIED {
			t.Skip("Skipping test: we don't have rights to manage services.")
		}
		t.Fatalf("SCM connection failed: %s", err)
	}
	defer m.Disconnect()

	s, err := m.OpenService("LanmanServer")
	if err != nil {
		t.Fatalf("OpenService(lanmanserver) failed: %s", err)
	}
	defer s.Close()

	_, err = s.Config()
	if err != nil {
		t.Fatalf("Config failed: %s", err)
	}
}

func install(t *testing.T, m *mgr.Mgr, name, exepath string, c mgr.Config) {
	// Sometimes it takes a while for the service to get
	// removed after previous test run.
	for i := 0; ; i++ {
		s, err := m.OpenService(name)
		if err != nil {
			break
		}
		s.Close()

		if i > 10 {
			t.Fatalf("service %s already exists", name)
		}
		time.Sleep(300 * time.Millisecond)
	}

	s, err := m.CreateService(name, exepath, c)
	if err != nil {
		t.Fatalf("CreateService(%s) failed: %v", name, err)
	}
	defer s.Close()
}

func depString(d []string) string {
	if len(d) == 0 {
		return ""
	}
	for i := range d {
		d[i] = strings.ToLower(d[i])
	}
	ss := sort.StringSlice(d)
	ss.Sort()
	return strings.Join([]string(ss), " ")
}

func testConfig(t *testing.T, s *mgr.Service, should mgr.Config) mgr.Config {
	is, err := s.Config()
	if err != nil {
		t.Fatalf("Config failed: %s", err)
	}
	if should.DelayedAutoStart != is.DelayedAutoStart {
		t.Fatalf("config mismatch: DelayedAutoStart is %v, but should have %v", is.DelayedAutoStart, should.DelayedAutoStart)
	}
	if should.DisplayName != is.DisplayName {
		t.Fatalf("config mismatch: DisplayName is %q, but should have %q", is.DisplayName, should.DisplayName)
	}
	if should.StartType != is.StartType {
		t.Fatalf("config mismatch: StartType is %v, but should have %v", is.StartType, should.StartType)
	}
	if should.Description != is.Description {
		t.Fatalf("config mismatch: Description is %q, but should have %q", is.Description, should.Description)
	}
	if depString(should.Dependencies) != depString(is.Dependencies) {
		t.Fatalf("config mismatch: Dependencies is %v, but should have %v", is.Dependencies, should.Dependencies)
	}
	return is
}

func testRecoveryActions(t *testing.T, s *mgr.Service, should []mgr.RecoveryAction) {
	is, err := s.RecoveryActions()
	if err != nil {
		t.Fatalf("RecoveryActions failed: %s", err)
	}
	if len(should) != len(is) {
		t.Errorf("recovery action mismatch: contains %v actions, but should have %v", len(is), len(should))
	}
	for i, _ := range is {
		if should[i].Type != is[i].Type {
			t.Errorf("recovery action mismatch: Type is %v, but should have %v", is[i].Type, should[i].Type)
		}
		if should[i].Delay != is[i].Delay {
			t.Errorf("recovery action mismatch: Delay is %v, but should have %v", is[i].Delay, should[i].Delay)
		}
	}
}

func testResetPeriod(t *testing.T, s *mgr.Service, should uint32) {
	is, err := s.ResetPeriod()
	if err != nil {
		t.Fatalf("ResetPeriod failed: %s", err)
	}
	if should != is {
		t.Errorf("reset period mismatch: reset period is %v, but should have %v", is, should)
	}
}

func testSetRecoveryActions(t *testing.T, s *mgr.Service) {
	r := []mgr.RecoveryAction{
		mgr.RecoveryAction{
			Type:  mgr.NoAction,
			Delay: 60000 * time.Millisecond,
		},
		mgr.RecoveryAction{
			Type:  mgr.ServiceRestart,
			Delay: 4 * time.Minute,
		},
		mgr.RecoveryAction{
			Type:  mgr.ServiceRestart,
			Delay: time.Minute,
		},
		mgr.RecoveryAction{
			Type:  mgr.RunCommand,
			Delay: 4000 * time.Millisecond,
		},
	}

	// 4 recovery actions with reset period
	err := s.SetRecoveryActions(r, uint32(10000))
	if err != nil {
		t.Fatalf("SetRecoveryActions failed: %v", err)
	}
	testRecoveryActions(t, s, r)
	testResetPeriod(t, s, uint32(10000))

	// Infinite reset period
	err = s.SetRecoveryActions(r, syscall.INFINITE)
	if err != nil {
		t.Fatalf("SetRecoveryActions failed: %v", err)
	}
	testRecoveryActions(t, s, r)
	testResetPeriod(t, s, syscall.INFINITE)

	// nil recovery actions
	err = s.SetRecoveryActions(nil, 0)
	if err.Error() != "recoveryActions cannot be nil" {
		t.Fatalf("SetRecoveryActions failed with unexpected error message of %q", err)
	}

	// Delete all recovery actions and reset period
	err = s.ResetRecoveryActions()
	if err != nil {
		t.Fatalf("ResetRecoveryActions failed: %v", err)
	}
	testRecoveryActions(t, s, nil)
	testResetPeriod(t, s, 0)
}

func testRebootMessage(t *testing.T, s *mgr.Service, should string) {
	err := s.SetRebootMessage(should)
	if err != nil {
		t.Fatalf("SetRebootMessage failed: %v", err)
	}
	is, err := s.RebootMessage()
	if err != nil {
		t.Fatalf("RebootMessage failed: %v", err)
	}
	if should != is {
		t.Errorf("reboot message mismatch: message is %q, but should have %q", is, should)
	}
}

func testRecoveryCommand(t *testing.T, s *mgr.Service, should string) {
	err := s.SetRecoveryCommand(should)
	if err != nil {
		t.Fatalf("SetRecoveryCommand failed: %v", err)
	}
	is, err := s.RecoveryCommand()
	if err != nil {
		t.Fatalf("RecoveryCommand failed: %v", err)
	}
	if should != is {
		t.Errorf("recovery command mismatch: command is %q, but should have %q", is, should)
	}
}

func remove(t *testing.T, s *mgr.Service) {
	err := s.Delete()
	if err != nil {
		t.Fatalf("Delete failed: %s", err)
	}
}

func TestMyService(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode - it modifies system services")
	}

	const name = "mymgrservice"

	m, err := mgr.Connect()
	if err != nil {
		if errno, ok := err.(syscall.Errno); ok && errno == syscall.ERROR_ACCESS_DENIED {
			t.Skip("Skipping test: we don't have rights to manage services.")
		}
		t.Fatalf("SCM connection failed: %s", err)
	}
	defer m.Disconnect()

	c := mgr.Config{
		StartType:    mgr.StartDisabled,
		DisplayName:  "my service",
		Description:  "my service is just a test",
		Dependencies: []string{"LanmanServer", "W32Time"},
	}

	exename := os.Args[0]
	exepath, err := filepath.Abs(exename)
	if err != nil {
		t.Fatalf("filepath.Abs(%s) failed: %s", exename, err)
	}

	install(t, m, name, exepath, c)

	s, err := m.OpenService(name)
	if err != nil {
		t.Fatalf("service %s is not installed", name)
	}
	defer s.Close()

	c.BinaryPathName = exepath
	c = testConfig(t, s, c)

	c.StartType = mgr.StartManual
	err = s.UpdateConfig(c)
	if err != nil {
		t.Fatalf("UpdateConfig failed: %v", err)
	}

	testConfig(t, s, c)

	c.StartType = mgr.StartAutomatic
	c.DelayedAutoStart = true
	err = s.UpdateConfig(c)
	if err != nil {
		t.Fatalf("UpdateConfig failed: %v", err)
	}

	testConfig(t, s, c)

	svcnames, err := m.ListServices()
	if err != nil {
		t.Fatalf("ListServices failed: %v", err)
	}
	var myserviceIsInstalled bool
	for _, sn := range svcnames {
		if sn == name {
			myserviceIsInstalled = true
			break
		}
	}
	if !myserviceIsInstalled {
		t.Errorf("ListServices failed to find %q service", name)
	}

	testSetRecoveryActions(t, s)
	testRebootMessage(t, s, fmt.Sprintf("%s failed", name))
	testRebootMessage(t, s, "") // delete reboot message
	testRecoveryCommand(t, s, fmt.Sprintf("sc query %s", name))
	testRecoveryCommand(t, s, "") // delete recovery command

	remove(t, s)
}
<|endoftext|>


<commit_before>package dukedb

import (
	"bytes"
	"fmt"
	"reflect"
	"strconv"
	"strings"

	"github.com/theduke/go-apperror"
)

// StatementTranslator takes a Statement or Expression and converts them to a string.
type ExpressionTranslator struct {
	buffer    bytes.Buffer
	Arguments []interface{}
}

func (ExpressionTranslator) QuoteIdentifier(id string) string {
	return "\"" + strings.Replace(id, "\"", "\\\"", -1) + "\""
}

func (ExpressionTranslator) QuoteValue(val interface{}) string {
	return "'" + fmt.Sprintf("%v", val) + "'"
}

func (ExpressionTranslator) PlaceHolder() string {
	return "?"
}

// Arg adds an argument to the arguments list.
func (t *ExpressionTranslator) Arg(val interface{}) {
	t.Arguments = append(t.Arguments, val)
}

// W writes a string to the buffer.
func (t *ExpressionTranslator) W(strs ...string) {
	for _, str := range strs {
		if _, err := t.buffer.WriteString(str); err != nil {
			panic(fmt.Sprintf("Could not write to buffer: %v", err))
		}
	}
}

// WQ quotes an identifier and writes it to the buffer.
func (t *ExpressionTranslator) WQ(str string) {
	t.W(t.QuoteIdentifier(str))
}

// Reset the buffer.
func (t *ExpressionTranslator) Reset() {
	t.buffer.Reset()
}

func (t *ExpressionTranslator) Translate(expression Expression) apperror.Error {
	panic("translator.Panic() not overwritten.")
}

func (t *ExpressionTranslator) String() string {
	return t.buffer.String()
}

/**
 * Basic SQL translater.
 */

type SqlTranslator struct {
	ExpressionTranslator
}

func (t *SqlTranslator) Translate(expression Expression) apperror.Error {
	if err := expression.Validate(); err != nil {
		return apperror.Wrap(err, "invalid_"+expression.Type()+"_expression")
	}

	switch e := expression.(type) {
	case *CreateCollectionStatement:
		t.Reset()

		t.W("CREATE TABLE ")
		if e.IfNotExists {
			t.W("IF NOT EXISTS ")
		}
		t.WQ(e.Collection)
		t.WQ("(")

		lastIndex := len(e.Fields) - 1
		for i, field := range e.Fields {
			if err := t.Translate(field); err != nil {
				return err
			}
			if i < lastIndex {
				t.W(", ")
			}
		}

		t.W(")")

	case *RenameCollectionStatement:
		t.Reset()

		t.W("ALTER TABLE ")
		t.WQ(e.Collection)
		t.W(" RENAME TO ")
		t.WQ(e.NewCollection)

	case *DropCollectionStatement:
		t.Reset()

		t.W("DROP TABLE ")
		if e.IfExists {
			t.W("IF EXISTS ")
		}
		t.WQ(e.Collection)
		if e.Cascade {
			t.W(" CASCADE")
		}

	case *AddCollectionFieldStatement:
		t.Reset()

		t.W("ALTER TABLE ")
		t.WQ(e.Collection)
		t.WQ(" ADD COLUMN ")
		if err := t.Translate(e.Field); err != nil {
			return err
		}

	case *RenameCollectionFieldStatement:
		t.Reset()

		t.W("ALTER TABLE ")
		t.WQ(e.Collection)
		t.W(" RENAME COLUMN ")
		t.WQ(e.Field)
		t.W(" TO ")
		t.WQ(e.NewName)

	case *DropCollectionFieldStatement:
		t.Reset()

		t.W("ALTER TABLE ")
		t.WQ(e.Collection)
		t.W("DROP COLUMN ")
		if e.IfExists {
			t.W("IF EXISTS ")
		}
		t.WQ(e.Field)
		if e.Cascasde {
			t.W(" CASCADE")
		}

	case *CreateIndexStatement:
		t.Reset()

		t.W("CREATE ")
		if e.Unique {
			t.W("UNIQUE ")
		}
		t.W("INDEX ")
		t.WQ(e.IndexName)
		t.W(" ON ")
		t.WQ(e.Collection)
		if e.Method != "" {
			t.W("USING ", e.Method, " ")
		}

		t.W("(")
		lastIndex := len(e.Expressions) - 1
		for index, expr := range e.Expressions {
			if err := t.Translate(expr); err != nil {
				return err
			}
			if index < lastIndex {
				t.W(", ")
			}
		}
		t.W(")")

	case *DropIndexStatement:
		t.Reset()

		t.W("DROP INDEX ")
		if e.IfExists {
			t.W(" IF EXISTS ")
		}
		t.WQ(e.IndexName)
		if e.Cascade {
			t.W(" CASCADE")
		}

	case *SelectStatement:
		t.Reset()

		t.W("SELECT ")

		// Field expressions.
		lastIndex := len(e.Fields) - 1
		for i, expr := range e.Fields {
			if err := t.Translate(expr.GetExpression()); err != nil {
				return err
			}
			t.W(" AS ")
			t.WQ(e.Collection + "." + expr.GetName())
			if i < lastIndex {
				t.W(", ")
			}
		}

		// Join fields.
		for _, join := range e.Joins {
			if join.RelationType != RELATION_TYPE_HAS_ONE {
				continue
			}
			t.W(", ")
			lastIndex := len(join.Fields) - 1
			for i, field := range join.Fields {
				if err := t.Translate(field.GetExpression()); err != nil {
					return err
				}
				t.W(" AS ")
				name := join.RelationName
				if name == "" {
					name = join.Collection
				}
				t.WQ(name + "." + field.GetName())
				if i < lastIndex {
					t.W(", ")
				}
			}
		}

		t.W(" FROM ")
		t.WQ(e.Collection)
		t.W(" ")

		// Join clauses.
		for _, join := range e.Joins {
			if err := t.Translate(join); err != nil {
				return err
			}
			t.W(" ")
		}

		if e.Filter != nil {
			t.W("WHERE ")
			if err := t.Translate(e.Filter); err != nil {
				return err
			}
			t.W(" ")
		}

		if len(e.Sorts) > 0 {
			t.W("ORDER BY ")
			lastIndex := len(e.Sorts) - 1
			for i, sort := range e.Sorts {
				if err := t.Translate(sort); err != nil {
					return err
				}
				if i < lastIndex {
					t.W(", ")
				}
			}
		}

		if e.Limit > 0 {
			t.W(" LIMIT ", strconv.Itoa(e.Limit))
		}
		if e.Offset > 0 {
			t.W(" OFFSET ", strconv.Itoa(e.Offset))
		}

	case *JoinStatement:
		name := e.RelationName
		if name == "" {
			name = e.Collection
		}
		t.W(JOIN_SQL_MAP[e.JoinType])
		t.WQ(name)
		t.W(" ON ")
		if err := t.Translate(e.JoinCondition); err != nil {
			return err
		}

	case *CreateStatement:
		t.Reset()

		t.W("INSERT INTO ")
		t.WQ(e.Collection)
		t.W("(")
		lastIndex := len(e.Values) - 1
		for i, field := range e.Values {
			if err := t.Translate(field); err != nil {
				return err
			}
			if i < lastIndex {
				t.W(", ")
			}
		}
		t.W(") VALUES(")
		for i, field := range e.Values {
			if err := t.Translate(field); err != nil {
				return err
			}
			if i < lastIndex {
				t.W(",")
			}
		}
		t.W(")")

	case *UpdateStatement:
		t.Reset()

		t.W("UPDATE ")
		t.WQ(e.Collection)
		t.W(" SET ")

		lastIndex := len(e.Values) - 1
		for i, field := range e.Values {
			if err := t.Translate(field); err != nil {
				return err
			}
			if i < lastIndex {
				t.W(", ")
			}
		}

		if e.Select != nil {
			if e.Select.Filter != nil {
				t.W("WHERE ")
				if err := t.Translate(e.Select.Filter); err != nil {
					return err
				}
				t.W(" ")
			}

			if len(e.Select.Sorts) > 0 {
				t.W("ORDER BY ")
				lastIndex := len(e.Select.Sorts) - 1
				for i, sort := range e.Select.Sorts {
					if err := t.Translate(sort); err != nil {
						return nil
					}
					if i < lastIndex {
						t.W(", ")
					}
				}
			}

			if e.Select.Limit > 0 {
				t.W(" LIMIT ", strconv.Itoa(e.Select.Limit))
			}
			if e.Select.Offset > 0 {
				t.W(" OFFSET ", strconv.Itoa(e.Select.Offset))
			}
		}

	case *NamedNestedExpr:
		return t.Translate(e.Expression)

	case *TextExpression:
		t.W(e.Text)

	case *FieldTypeExpression:
		t.W(e.Typ)

	case *ValueExpression:
		t.W(t.PlaceHolder())
		t.Arg(e.Val)

	case *IdentifierExpression:
		t.WQ(e.Identifier)

	case *CollectionFieldIdentifierExpression:
		if e.Collection != "" {
			t.WQ(e.Collection)
			t.W(".")
		}
		t.WQ(e.Field)

	case *NotNullConstraint:
		t.W("NOT NULL")

	case *UniqueConstraint:
		t.W("UNIQUE")

	case *UniqueFieldsConstraint:
		t.W("UNIQUE (")
		lastIndex := len(e.Fields) - 1
		for i, field := range e.Fields {
			if err := t.Translate(field); err != nil {
				return err
			}
			if i < lastIndex {
				t.W(", ")
			}
		}
		t.W(")")

	case *PrimaryKeyConstraint:
		t.W("PRIMARY KEY")

	case *AutoIncrementConstraint:
		t.W("AUTO INCREMENT")

	case *DefaultValueConstraint:
		t.W("DEFAULT ")
		if err := t.Translate(e.Value); err != nil {
			return err
		}

	case *FieldUpdateConstraint:
		t.W("ON UPDATE ", e.Action)

	case *FieldDeleteConstraint:
		t.W("ON DELETE ", e.Action)

	case *IndexConstraint:

	case *CheckConstraint:
		t.W("CHECK (")
		if err := t.Translate(e.Check); err != nil {
			return err
		}
		t.W(")")

	case *ReferenceConstraint:
		t.W("REFERENCES ")
		t.WQ(e.ForeignKey.Collection)
		t.W(" (")
		t.WQ(e.ForeignKey.Field)
		t.W(")")

	case *FieldExpression:
		t.WQ(e.Name)
		t.W(" ")
		if err := t.Translate(e.Typ); err != nil {
			return err
		}
		if len(e.Constraints) > 0 {
			t.W(" ")
			for _, constraint := range e.Constraints {
				if err := t.Translate(constraint); err != nil {
					return err
				}
			}
		}

	case *FieldValueExpression:
		if err := t.Translate(e.Field); err != nil {
			return err
		}
		t.W("=")
		if err := t.Translate(e.Value); err != nil {
			return err
		}

	case *FunctionExpression:
		t.W(e.Function, "(")
		if err := t.Translate(e.Nested); err != nil {
			return err
		}
		t.W(")")

	case *AndExpression:
		t.W("(")
		lastIndex := len(e.Expressions) - 1
		for i, expr := range e.Expressions {
			if err := t.Translate(expr); err != nil {
				return err
			}
			if i < lastIndex {
				t.W(" AND ")
			}
		}
		t.W(")")

	case *OrExpression:
		t.W("(")
		lastIndex := len(e.Expressions) - 1
		for i, expr := range e.Expressions {
			if err := t.Translate(expr); err != nil {
				return err
			}
			if i < lastIndex {
				t.W(" OR ")
			}
		}
		t.W(")")

	case *NotExpression:
		t.W("NOT ")
		if err := t.Translate(e.Nested); err != nil {
			return err
		}

	case FilterExpression:
		if err := t.Translate(e.GetField()); err != nil {
			return err
		}
		t.W(" ", OperatorToSql(e.GetOperator()), " ")
		if err := t.Translate(e.GetClause()); err != nil {
			return err
		}

	case *SortExpression:
		if err := t.Translate(e.Field); err != nil {
			return err
		}
		if e.Ascending {
			t.W(" ASC")
		} else {
			t.W(" DESC")
		}

	default:
		panic(fmt.Sprintf("Unhandled statement type: %v", reflect.TypeOf(expression)))
	}

	return nil
}
<commit_msg>Updated expression translator to new stmts/expressions<commit_after>package dukedb

import (
	"bytes"
	"fmt"
	"reflect"
	"strconv"
	"strings"

	"github.com/theduke/go-apperror"
)

// StatementTranslator takes a Statement or Expression and converts them to a string.
type ExpressionTranslator struct {
	buffer    bytes.Buffer
	arguments []interface{}

	// Counter for the translations performed since last .Reset().
	translationCounter int
}

func (ExpressionTranslator) QuoteIdentifier(id string) string {
	return "\"" + strings.Replace(id, "\"", "\\\"", -1) + "\""
}

func (ExpressionTranslator) QuoteValue(val interface{}) string {
	return "'" + fmt.Sprintf("%v", val) + "'"
}

func (ExpressionTranslator) PlaceHolder() string {
	return "?"
}

// Arg adds an argument to the arguments list.
func (t *ExpressionTranslator) Arg(val interface{}) {
	t.arguments = append(t.arguments, val)
}

// W writes a string to the buffer.
func (t *ExpressionTranslator) W(strs ...string) {
	for _, str := range strs {
		if _, err := t.buffer.WriteString(str); err != nil {
			panic(fmt.Sprintf("Could not write to buffer: %v", err))
		}
	}
}

// WQ quotes an identifier and writes it to the buffer.
func (t *ExpressionTranslator) WQ(str string) {
	t.W(t.QuoteIdentifier(str))
}

// Reset the buffer.
func (t *ExpressionTranslator) Reset() {
	t.buffer.Reset()
	t.translationCounter = 0
}

func (t *ExpressionTranslator) Translate(expression Expression) apperror.Error {
	panic("translator.Panic() not overwritten.")
}

func (t *ExpressionTranslator) String() string {
	return t.buffer.String()
}

func (t ExpressionTranslator) Arguments() []interface{} {
	return t.arguments
}

/**
 * Basic SQL translater.
 */

type SqlTranslator struct {
	ExpressionTranslator
}

func (t *SqlTranslator) Translate(expression Expression) apperror.Error {
	if err := expression.Validate(); err != nil {
		return apperror.Wrap(err, "invalid_"+expression.Type()+"_expression")
	}

	switch e := expression.(type) {
	case NamedNestedExpression:
		return t.Translate(e.Expression())

	case TextExpression:
		t.W(e.Text())

	case FieldTypeExpression:
		t.W(e.FieldType())

	case ValueExpression:
		t.W(t.PlaceHolder())
		t.Arg(e.Value())

	case IdentifierExpression:
		t.WQ(e.Identifier())

	case CollectionFieldIdentifierExpression:
		if e.Collection() != "" {
			t.WQ(e.Collection())
			t.W(".")
		}
		t.WQ(e.Field())

	case ConstraintExpression:
		switch e.Constraint() {
		case CONSTRAINT_NOT_NULL:
			t.W("NOT NULL")
		case CONSTRAINT_UNIQUE:
			t.W("UNIQUE")
		case CONSTRAINT_PRIMARY_KEY:
			t.W("PRIMARY KEY")
		case CONSTRAINT_AUTO_INCREMENT:
			t.W("AUTO_INCREMENT")
		}

	case ActionConstraint:
		t.W("ON ", strings.ToUpper(e.Event()), " ", ACTIONS_MAP[e.Action()])

	/*
		case UniqueFieldsConstraint:
			t.W("UNIQUE(")
			lastIndex := len(e.Fields()) - 1
			for i, field := range e.Fields {
				if err := t.Translate(field); err != nil {
					return err
				}
				if i < lastIndex {
					t.W(", ")
				}
			}
			t.W(")")
	*/

	case DefaultValueConstraint:
		t.W("DEFAULT ")
		if err := t.Translate(e.DefaultValue()); err != nil {
			return err
		}

	case CheckConstraint:
		t.W("CHECK (")
		if err := t.Translate(e.Check()); err != nil {
			return err
		}
		t.W(")")

	case ReferenceConstraint:
		t.W("REFERENCES ")
		fk := e.ForeignKey()
		t.WQ(fk.Collection())
		t.W(" (")
		t.WQ(fk.Field())
		t.W(")")

	case FieldExpression:
		t.WQ(e.Name())
		t.W(" ")
		if err := t.Translate(e.FieldType()); err != nil {
			return err
		}

		constraints := e.Constraints()
		if len(constraints) > 0 {
			t.W(" ")
			lastIndex := len(constraints) - 1
			for index, constraint := range constraints {
				if err := t.Translate(constraint); err != nil {
					return err
				}
				if index < lastIndex {
					t.W(" ")
				}
			}
		}

	case FieldValueExpression:
		if err := t.Translate(e.Field()); err != nil {
			return err
		}
		t.W(" = ")
		if err := t.Translate(e.Value()); err != nil {
			return err
		}

	case FunctionExpression:
		t.W(e.Function(), "(")
		if err := t.Translate(e.Expression()); err != nil {
			return err
		}
		t.W(")")

	case *AndExpression:
		lastIndex := len(e.Expressions()) - 1
		if lastIndex > 1 {
			t.W("(")
		}
		for i, expr := range e.Expressions() {
			if err := t.Translate(expr); err != nil {
				return err
			}
			if i < lastIndex {
				t.W(" AND ")
			}
		}
		if lastIndex > 1 {
			t.W(")")
		}

	case *OrExpression:
		lastIndex := len(e.Expressions()) - 1

		// Wrap in parantheses if more than one filter.
		if lastIndex > 1 {
			t.W("(")
		}
		for i, expr := range e.Expressions() {
			if err := t.Translate(expr); err != nil {
				return err
			}
			if i < lastIndex {
				t.W(" OR ")
			}
		}
		if lastIndex > 1 {
			t.W(")")
		}

	case NotExpression:
		t.W("NOT ")
		if err := t.Translate(e.Not()); err != nil {
			return err
		}

	case FilterExpression:
		if err := t.Translate(e.Field()); err != nil {
			return err
		}
		t.W(" ", e.Operator(), " ")
		if err := t.Translate(e.Clause()); err != nil {
			return err
		}

	case SortExpression:
		if err := t.Translate(e.Expression()); err != nil {
			return err
		}
		if e.Ascending() {
			t.W(" ASC")
		} else {
			t.W(" DESC")
		}

	case CreateCollectionStatement:
		t.W("CREATE TABLE ")
		if e.IfNotExists() {
			t.W("IF NOT EXISTS ")
		}
		t.WQ(e.Collection())
		t.W(" (")

		// Fields.
		lastIndex := len(e.Fields()) - 1
		for i, field := range e.Fields() {
			if err := t.Translate(field); err != nil {
				return err
			}
			if i < lastIndex {
				t.W(", ")
			}
		}

		// Constraints.
		lastIndex = len(e.Constraints()) - 1
		for i, constraint := range e.Constraints() {
			t.W(", ")
			if err := t.Translate(constraint); err != nil {
				return err
			}
			if i < lastIndex {
				t.W(", ")
			}
		}

		t.W(")")

	case RenameCollectionStatement:
		t.W("ALTER TABLE ")
		t.WQ(e.Collection())
		t.W(" RENAME TO ")
		t.WQ(e.NewName())

	case DropCollectionStatement:
		t.W("DROP TABLE ")
		if e.IfExists() {
			t.W("IF EXISTS ")
		}
		t.WQ(e.Collection())
		if e.Cascade() {
			t.W(" CASCADE")
		}

	case CreateFieldStatement:
		t.W("ALTER TABLE ")
		t.WQ(e.Collection())
		t.W(" ADD COLUMN ")
		if err := t.Translate(e.Field()); err != nil {
			return err
		}

	case RenameFieldStatement:
		t.W("ALTER TABLE ")
		t.WQ(e.Collection())
		t.W(" RENAME COLUMN ")
		t.WQ(e.Field())
		t.W(" TO ")
		t.WQ(e.NewName())

	case DropFieldStatement:
		t.W("ALTER TABLE ")
		t.WQ(e.Collection())
		t.W(" DROP COLUMN ")
		if e.IfExists() {
			t.W("IF EXISTS ")
		}
		t.WQ(e.Field())
		if e.Cascade() {
			t.W(" CASCADE")
		}

	case CreateIndexStatement:
		t.W("CREATE ")
		if e.Unique() {
			t.W("UNIQUE ")
		}
		t.W("INDEX ")
		t.WQ(e.IndexName())
		t.W(" ON ")
		if err := t.Translate(e.IndexExpression()); err != nil {
			return err
		}
		if e.Method() != "" {
			t.W("USING ", e.Method(), " ")
		}

		t.W("(")
		lastIndex := len(e.Expressions()) - 1
		for index, expr := range e.Expressions() {
			if err := t.Translate(expr); err != nil {
				return err
			}
			if index < lastIndex {
				t.W(", ")
			}
		}
		t.W(")")

	case DropIndexStatement:
		t.W("DROP INDEX ")
		if e.IfExists() {
			t.W("IF EXISTS ")
		}
		t.WQ(e.IndexName())
		if e.Cascade() {
			t.W(" CASCADE")
		}

	case SelectStatement:
		// If counter is bigger than 0, this is a subquery and needs to be
		// wrapped in parantheses.
		isSubQuery := t.translationCounter > 0
		if isSubQuery {
			t.W("(")
		}

		t.W("SELECT ")

		// Field expressions.
		lastIndex := len(e.Fields()) - 1
		for i, expr := range e.Fields() {
			if err := t.Translate(expr); err != nil {
				return err
			}
			if i < lastIndex {
				t.W(", ")
			}
		}

		// Join fields.
		for _, join := range e.Joins() {
			if join.RelationType() != RELATION_TYPE_HAS_ONE {
				continue
			}
			t.W(", ")
			lastIndex := len(join.Fields()) - 1
			for i, field := range join.Fields() {
				if err := t.Translate(field); err != nil {
					return err
				}
				if i < lastIndex {
					t.W(", ")
				}
			}
		}

		t.W(" FROM ")
		t.WQ(e.Collection())

		// Join clauses.
		for _, join := range e.Joins() {
			t.W(" ")
			if err := t.Translate(join); err != nil {
				return err
			}
		}

		if e.Filter() != nil {
			t.W(" WHERE ")
			if err := t.Translate(e.Filter()); err != nil {
				return err
			}
		}

		if len(e.Sorts()) > 0 {
			t.W(" ORDER BY ")
			lastIndex := len(e.Sorts()) - 1
			for i, sort := range e.Sorts() {
				if err := t.Translate(sort); err != nil {
					return err
				}
				if i < lastIndex {
					t.W(", ")
				}
			}
		}

		if e.Limit() > 0 {
			t.W(" LIMIT ", strconv.Itoa(e.Limit()))
		}
		if e.Offset() > 0 {
			t.W(" OFFSET ", strconv.Itoa(e.Offset()))
		}

		// If counter is bigger than 0, this is a subquery and needs to be
		// wrapped in parantheses.
		if isSubQuery {
			t.W(")")
		}

	case JoinStatement:
		name := e.RelationName()
		if name == "" {
			name = e.Collection()
		}
		t.W(JOIN_MAP[e.JoinType()])
		t.WQ(name)
		t.W(" ON ")
		if err := t.Translate(e.JoinCondition()); err != nil {
			return err
		}

	case UpdateStatement:
		t.W("UPDATE ")
		t.WQ(e.Collection())
		t.W(" SET ")

		lastIndex := len(e.Values()) - 1
		for i, field := range e.Values() {
			if err := t.Translate(field); err != nil {
				return err
			}
			if i < lastIndex {
				t.W(", ")
			}
		}

		sel := e.Select()
		if sel != nil {
			if sel.Filter() != nil {
				t.W("WHERE ")
				if err := t.Translate(sel.Filter()); err != nil {
					return err
				}
				t.W(" ")
			}

			if len(sel.Sorts()) > 0 {
				t.W("ORDER BY ")
				lastIndex := len(sel.Sorts()) - 1
				for i, sort := range sel.Sorts() {
					if err := t.Translate(sort); err != nil {
						return nil
					}
					if i < lastIndex {
						t.W(", ")
					}
				}
			}

			if sel.Limit() > 0 {
				t.W(" LIMIT ", strconv.Itoa(sel.Limit()))
			}
			if sel.Offset() > 0 {
				t.W(" OFFSET ", strconv.Itoa(sel.Offset()))
			}
		}

	case CreateStatement:
		t.W("INSERT INTO ")
		t.WQ(e.Collection())
		t.W("(")
		lastIndex := len(e.Values()) - 1
		for i, field := range e.Values() {
			if err := t.Translate(field); err != nil {
				return err
			}
			if i < lastIndex {
				t.W(", ")
			}
		}
		t.W(") VALUES(")
		for i, field := range e.Values() {
			if err := t.Translate(field); err != nil {
				return err
			}
			if i < lastIndex {
				t.W(",")
			}
		}
		t.W(")")

	default:
		panic(fmt.Sprintf("Unhandled statement type: %v", reflect.TypeOf(expression)))
	}

	t.translationCounter += 1
	return nil
}
<|endoftext|>


<commit_before>package workflows

import (
	"fmt"
	"github.com/stelligent/mu/common"
	"github.com/stelligent/mu/templates"
	"strconv"
	"strings"
)

// NewServiceDeployer create a new workflow for deploying a service in an environment
func NewServiceDeployer(ctx *common.Context, environmentName string, tag string) Executor {

	workflow := new(serviceWorkflow)
	workflow.codeRevision = ctx.Config.Repo.Revision
	workflow.repoName = ctx.Config.Repo.Slug

	stackParams := make(map[string]string)

	return newPipelineExecutor(
		workflow.serviceLoader(ctx, tag, ""),
		workflow.serviceEnvironmentLoader(environmentName, ctx.StackManager),
		workflow.serviceApplyCommonParams(&ctx.Config.Service, stackParams, environmentName, ctx.StackManager, ctx.ElbManager, ctx.ParamManager),
		newConditionalExecutor(workflow.isEcsProvider(),
			newPipelineExecutor(
				workflow.serviceRepoUpserter(&ctx.Config.Service, ctx.StackManager, ctx.StackManager),
				workflow.serviceApplyEcsParams(&ctx.Config.Service, stackParams),
				workflow.serviceEcsDeployer(&ctx.Config.Service, stackParams, environmentName, ctx.StackManager, ctx.StackManager),
			),
			newPipelineExecutor(
				workflow.serviceBucketUpserter(&ctx.Config.Service, ctx.StackManager, ctx.StackManager),
				workflow.serviceAppUpserter(&ctx.Config.Service, ctx.StackManager, ctx.StackManager),
				workflow.serviceApplyEc2Params(stackParams),
				workflow.serviceEc2Deployer(&ctx.Config.Service, stackParams, environmentName, ctx.StackManager, ctx.StackManager),
			),
		),
	)
}

func getMaxPriority(elbRuleLister common.ElbRuleLister, listenerArn string) int {
	rules, err := elbRuleLister.ListRules(listenerArn)
	if err != nil {
		log.Debugf("Error finding max priority for listener '%s': %v", listenerArn, err)
		return 0
	}
	maxPriority := 0
	for _, rule := range rules {
		priority, _ := strconv.Atoi(common.StringValue(rule.Priority))
		if priority > maxPriority {
			maxPriority = priority
		}
	}
	return maxPriority
}

func (workflow *serviceWorkflow) serviceEnvironmentLoader(environmentName string, stackWaiter common.StackWaiter) Executor {
	return func() error {
		lbStackName := common.CreateStackName(common.StackTypeLoadBalancer, environmentName)
		workflow.lbStack = stackWaiter.AwaitFinalStatus(lbStackName)

		envStackName := common.CreateStackName(common.StackTypeEnv, environmentName)
		workflow.envStack = stackWaiter.AwaitFinalStatus(envStackName)

		if workflow.envStack == nil {
			return fmt.Errorf("Unable to find stack '%s' for environment '%s'", envStackName, environmentName)
		}

		if workflow.isEcsProvider()() {
			workflow.artifactProvider = common.ArtifactProviderEcr
		} else {
			workflow.artifactProvider = common.ArtifactProviderS3
		}

		return nil
	}
}

func (workflow *serviceWorkflow) serviceApplyEcsParams(service *common.Service, params map[string]string) Executor {
	return func() error {

		params["EcsCluster"] = fmt.Sprintf("%s-EcsCluster", workflow.envStack.Name)
		params["ImageUrl"] = workflow.serviceImage

		if service.CPU != 0 {
			params["ServiceCpu"] = strconv.Itoa(service.CPU)
		}
		if service.Memory != 0 {
			params["ServiceMemory"] = strconv.Itoa(service.Memory)
		}

		return nil
	}
}

func (workflow *serviceWorkflow) serviceApplyEc2Params(params map[string]string) Executor {
	return func() error {

		params["AppName"] = workflow.appName
		params["RevisionBucket"] = workflow.appRevisionBucket
		params["RevisionKey"] = workflow.appRevisionKey
		params["RevisionBundleType"] = "zip"

		for _, key := range [...]string{
			"SshAllow",
			"InstanceType",
			"ImageId",
			"MaxSize",
			"KeyName",
			"ScaleInThreshold",
			"ScaleOutThreshold",
			"HttpProxy",
			"ConsulServerAutoScalingGroup",
			"ElbSecurityGroup",
			"ConsulRpcClientSecurityGroup",
			"InstanceSubnetIds",
			"InstanceSecurityGroup",
		} {
			params[key] = workflow.envStack.Parameters[key]
		}

		return nil
	}
}

func (workflow *serviceWorkflow) serviceApplyCommonParams(service *common.Service, params map[string]string, environmentName string, stackWaiter common.StackWaiter, elbRuleLister common.ElbRuleLister, paramGetter common.ParamGetter) Executor {
	return func() error {
		params["VpcId"] = fmt.Sprintf("%s-VpcId", workflow.envStack.Name)

		nextAvailablePriority := 0
		if workflow.lbStack.Outputs["ElbHttpListenerArn"] != "" {
			params["ElbHttpListenerArn"] = fmt.Sprintf("%s-ElbHttpListenerArn", workflow.lbStack.Name)
			nextAvailablePriority = 1 + getMaxPriority(elbRuleLister, workflow.lbStack.Outputs["ElbHttpListenerArn"])
		}
		if workflow.lbStack.Outputs["ElbHttpsListenerArn"] != "" {
			params["ElbHttpsListenerArn"] = fmt.Sprintf("%s-ElbHttpsListenerArn", workflow.lbStack.Name)
			if nextAvailablePriority == 0 {
				nextAvailablePriority = 1 + getMaxPriority(elbRuleLister, workflow.lbStack.Outputs["ElbHttpsListenerArn"])
			}
		}

		dbStackName := common.CreateStackName(common.StackTypeDatabase, workflow.serviceName, environmentName)
		dbStack := stackWaiter.AwaitFinalStatus(dbStackName)
		if dbStack != nil {
			params["DatabaseName"] = dbStack.Outputs["DatabaseName"]
			params["DatabaseEndpointAddress"] = dbStack.Outputs["DatabaseEndpointAddress"]
			params["DatabaseEndpointPort"] = dbStack.Outputs["DatabaseEndpointPort"]
			params["DatabaseMasterUsername"] = dbStack.Outputs["DatabaseMasterUsername"]

			dbPass, _ := paramGetter.GetParam(fmt.Sprintf("%s-%s", dbStackName, "DatabaseMasterPassword"))
			params["DatabaseMasterPassword"] = dbPass
		}

		svcStackName := common.CreateStackName(common.StackTypeService, workflow.serviceName, environmentName)
		svcStack := stackWaiter.AwaitFinalStatus(svcStackName)

		if workflow.priority > 0 {
			params["PathListenerRulePriority"] = strconv.Itoa(workflow.priority)
			params["HostListenerRulePriority"] = strconv.Itoa(workflow.priority + 1)
		} else if svcStack != nil {
			// no value in config, and this is an update...use prior value
			params["PathListenerRulePriority"] = ""
			params["HostListenerRulePriority"] = ""
		} else {
			// no value in config, and this is a create...use next available
			params["PathListenerRulePriority"] = strconv.Itoa(nextAvailablePriority)
			params["HostListenerRulePriority"] = strconv.Itoa(nextAvailablePriority + 1)
		}

		params["ServiceName"] = workflow.serviceName
		if service.Port != 0 {
			params["ServicePort"] = strconv.Itoa(service.Port)
		}
		if service.Protocol != "" {
			params["ServiceProtocol"] = strings.ToUpper(service.Protocol)
		}
		if service.HealthEndpoint != "" {
			params["ServiceHealthEndpoint"] = service.HealthEndpoint
		}
		if service.DesiredCount != 0 {
			params["ServiceDesiredCount"] = strconv.Itoa(service.DesiredCount)
		}
		if len(service.PathPatterns) > 0 {
			params["PathPattern"] = strings.Join(service.PathPatterns, ",")
		}
		if len(service.HostPatterns) > 0 {
			params["HostPattern"] = strings.Join(service.HostPatterns, ",")
		}

		return nil
	}
}

func (workflow *serviceWorkflow) serviceEc2Deployer(service *common.Service, stackParams map[string]string, environmentName string, stackUpserter common.StackUpserter, stackWaiter common.StackWaiter) Executor {
	return func() error {

		log.Noticef("Deploying service '%s' to '%s'", workflow.serviceName, environmentName)

		svcStackName := common.CreateStackName(common.StackTypeService, workflow.serviceName, environmentName)

		resolveServiceEnvironment(service, environmentName)
		overrides := common.GetStackOverrides(svcStackName)
		template, err := templates.NewTemplate("service-ec2.yml", service, overrides)
		if err != nil {
			return err
		}

		err = stackUpserter.UpsertStack(svcStackName, template, stackParams, buildServiceTags(workflow.serviceName, environmentName, workflow.envStack.Outputs["provider"], common.StackTypeService, workflow.codeRevision, workflow.repoName))
		if err != nil {
			return err
		}
		log.Debugf("Waiting for stack '%s' to complete", svcStackName)
		stack := stackWaiter.AwaitFinalStatus(svcStackName)
		if stack == nil {
			return fmt.Errorf("Unable to create stack %s", svcStackName)
		}
		if strings.HasSuffix(stack.Status, "ROLLBACK_COMPLETE") || !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		return nil
	}
}

func (workflow *serviceWorkflow) serviceEcsDeployer(service *common.Service, stackParams map[string]string, environmentName string, stackUpserter common.StackUpserter, stackWaiter common.StackWaiter) Executor {
	return func() error {
		log.Noticef("Deploying service '%s' to '%s' from '%s'", workflow.serviceName, environmentName, workflow.serviceImage)

		svcStackName := common.CreateStackName(common.StackTypeService, workflow.serviceName, environmentName)

		resolveServiceEnvironment(service, environmentName)
		overrides := common.GetStackOverrides(svcStackName)
		template, err := templates.NewTemplate("service-ecs.yml", service, overrides)
		if err != nil {
			return err
		}

		err = stackUpserter.UpsertStack(svcStackName, template, stackParams, buildServiceTags(workflow.serviceName, environmentName, workflow.envStack.Outputs["provider"], common.StackTypeService, workflow.codeRevision, workflow.repoName))
		if err != nil {
			return err
		}
		log.Debugf("Waiting for stack '%s' to complete", svcStackName)
		stack := stackWaiter.AwaitFinalStatus(svcStackName)
		if stack == nil {
			return fmt.Errorf("Unable to create stack %s", svcStackName)
		}
		if strings.HasSuffix(stack.Status, "ROLLBACK_COMPLETE") || !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		return nil
	}

}

func resolveServiceEnvironment(service *common.Service, environment string) {
	for key, value := range service.Environment {
		switch value.(type) {
		case map[interface{}]interface{}:
			found := false
			for env, v := range value.(map[interface{}]interface{}) {
				if env.(string) == environment {
					service.Environment[key] = v.(string)
					found = true
				}
			}
			if found != true {
				service.Environment[key] = ""
			}
		case string:
			// do nothing
		default:
			log.Warningf("Unable to resolve environment '%s': %v", key, value)
		}

	}
}

func buildServiceTags(serviceName string, environmentName string, envProvider string, stackType common.StackType, codeRevision string, repoName string) map[string]string {
	return map[string]string{
		"type":        string(stackType),
		"environment": environmentName,
		"provider":    envProvider,
		"service":     serviceName,
		"revision":    codeRevision,
		"repo":        repoName,
	}
}
<commit_msg>defect with ec2 deployments and instance security groups<commit_after>package workflows

import (
	"fmt"
	"github.com/stelligent/mu/common"
	"github.com/stelligent/mu/templates"
	"strconv"
	"strings"
)

// NewServiceDeployer create a new workflow for deploying a service in an environment
func NewServiceDeployer(ctx *common.Context, environmentName string, tag string) Executor {

	workflow := new(serviceWorkflow)
	workflow.codeRevision = ctx.Config.Repo.Revision
	workflow.repoName = ctx.Config.Repo.Slug

	stackParams := make(map[string]string)

	return newPipelineExecutor(
		workflow.serviceLoader(ctx, tag, ""),
		workflow.serviceEnvironmentLoader(environmentName, ctx.StackManager),
		workflow.serviceApplyCommonParams(&ctx.Config.Service, stackParams, environmentName, ctx.StackManager, ctx.ElbManager, ctx.ParamManager),
		newConditionalExecutor(workflow.isEcsProvider(),
			newPipelineExecutor(
				workflow.serviceRepoUpserter(&ctx.Config.Service, ctx.StackManager, ctx.StackManager),
				workflow.serviceApplyEcsParams(&ctx.Config.Service, stackParams),
				workflow.serviceEcsDeployer(&ctx.Config.Service, stackParams, environmentName, ctx.StackManager, ctx.StackManager),
			),
			newPipelineExecutor(
				workflow.serviceBucketUpserter(&ctx.Config.Service, ctx.StackManager, ctx.StackManager),
				workflow.serviceAppUpserter(&ctx.Config.Service, ctx.StackManager, ctx.StackManager),
				workflow.serviceApplyEc2Params(stackParams),
				workflow.serviceEc2Deployer(&ctx.Config.Service, stackParams, environmentName, ctx.StackManager, ctx.StackManager),
			),
		),
	)
}

func getMaxPriority(elbRuleLister common.ElbRuleLister, listenerArn string) int {
	rules, err := elbRuleLister.ListRules(listenerArn)
	if err != nil {
		log.Debugf("Error finding max priority for listener '%s': %v", listenerArn, err)
		return 0
	}
	maxPriority := 0
	for _, rule := range rules {
		priority, _ := strconv.Atoi(common.StringValue(rule.Priority))
		if priority > maxPriority {
			maxPriority = priority
		}
	}
	return maxPriority
}

func (workflow *serviceWorkflow) serviceEnvironmentLoader(environmentName string, stackWaiter common.StackWaiter) Executor {
	return func() error {
		lbStackName := common.CreateStackName(common.StackTypeLoadBalancer, environmentName)
		workflow.lbStack = stackWaiter.AwaitFinalStatus(lbStackName)

		envStackName := common.CreateStackName(common.StackTypeEnv, environmentName)
		workflow.envStack = stackWaiter.AwaitFinalStatus(envStackName)

		if workflow.envStack == nil {
			return fmt.Errorf("Unable to find stack '%s' for environment '%s'", envStackName, environmentName)
		}

		if workflow.isEcsProvider()() {
			workflow.artifactProvider = common.ArtifactProviderEcr
		} else {
			workflow.artifactProvider = common.ArtifactProviderS3
		}

		return nil
	}
}

func (workflow *serviceWorkflow) serviceApplyEcsParams(service *common.Service, params map[string]string) Executor {
	return func() error {

		params["EcsCluster"] = fmt.Sprintf("%s-EcsCluster", workflow.envStack.Name)
		params["ImageUrl"] = workflow.serviceImage

		if service.CPU != 0 {
			params["ServiceCpu"] = strconv.Itoa(service.CPU)
		}
		if service.Memory != 0 {
			params["ServiceMemory"] = strconv.Itoa(service.Memory)
		}

		return nil
	}
}

func (workflow *serviceWorkflow) serviceApplyEc2Params(params map[string]string) Executor {
	return func() error {

		params["AppName"] = workflow.appName
		params["RevisionBucket"] = workflow.appRevisionBucket
		params["RevisionKey"] = workflow.appRevisionKey
		params["RevisionBundleType"] = "zip"

		for _, key := range [...]string{
			"SshAllow",
			"InstanceType",
			"ImageId",
			"MaxSize",
			"KeyName",
			"ScaleInThreshold",
			"ScaleOutThreshold",
			"HttpProxy",
			"ConsulServerAutoScalingGroup",
			"ElbSecurityGroup",
			"ConsulRpcClientSecurityGroup",
			"InstanceSubnetIds",
			"InstanceSecurityGroup",
		} {
			params[key] = workflow.envStack.Outputs[key]
		}

		return nil
	}
}

func (workflow *serviceWorkflow) serviceApplyCommonParams(service *common.Service, params map[string]string, environmentName string, stackWaiter common.StackWaiter, elbRuleLister common.ElbRuleLister, paramGetter common.ParamGetter) Executor {
	return func() error {
		params["VpcId"] = fmt.Sprintf("%s-VpcId", workflow.envStack.Name)

		nextAvailablePriority := 0
		if workflow.lbStack.Outputs["ElbHttpListenerArn"] != "" {
			params["ElbHttpListenerArn"] = fmt.Sprintf("%s-ElbHttpListenerArn", workflow.lbStack.Name)
			nextAvailablePriority = 1 + getMaxPriority(elbRuleLister, workflow.lbStack.Outputs["ElbHttpListenerArn"])
		}
		if workflow.lbStack.Outputs["ElbHttpsListenerArn"] != "" {
			params["ElbHttpsListenerArn"] = fmt.Sprintf("%s-ElbHttpsListenerArn", workflow.lbStack.Name)
			if nextAvailablePriority == 0 {
				nextAvailablePriority = 1 + getMaxPriority(elbRuleLister, workflow.lbStack.Outputs["ElbHttpsListenerArn"])
			}
		}

		dbStackName := common.CreateStackName(common.StackTypeDatabase, workflow.serviceName, environmentName)
		dbStack := stackWaiter.AwaitFinalStatus(dbStackName)
		if dbStack != nil {
			params["DatabaseName"] = dbStack.Outputs["DatabaseName"]
			params["DatabaseEndpointAddress"] = dbStack.Outputs["DatabaseEndpointAddress"]
			params["DatabaseEndpointPort"] = dbStack.Outputs["DatabaseEndpointPort"]
			params["DatabaseMasterUsername"] = dbStack.Outputs["DatabaseMasterUsername"]

			dbPass, _ := paramGetter.GetParam(fmt.Sprintf("%s-%s", dbStackName, "DatabaseMasterPassword"))
			params["DatabaseMasterPassword"] = dbPass
		}

		svcStackName := common.CreateStackName(common.StackTypeService, workflow.serviceName, environmentName)
		svcStack := stackWaiter.AwaitFinalStatus(svcStackName)

		if workflow.priority > 0 {
			params["PathListenerRulePriority"] = strconv.Itoa(workflow.priority)
			params["HostListenerRulePriority"] = strconv.Itoa(workflow.priority + 1)
		} else if svcStack != nil {
			// no value in config, and this is an update...use prior value
			params["PathListenerRulePriority"] = ""
			params["HostListenerRulePriority"] = ""
		} else {
			// no value in config, and this is a create...use next available
			params["PathListenerRulePriority"] = strconv.Itoa(nextAvailablePriority)
			params["HostListenerRulePriority"] = strconv.Itoa(nextAvailablePriority + 1)
		}

		params["ServiceName"] = workflow.serviceName
		if service.Port != 0 {
			params["ServicePort"] = strconv.Itoa(service.Port)
		}
		if service.Protocol != "" {
			params["ServiceProtocol"] = strings.ToUpper(service.Protocol)
		}
		if service.HealthEndpoint != "" {
			params["ServiceHealthEndpoint"] = service.HealthEndpoint
		}
		if service.DesiredCount != 0 {
			params["ServiceDesiredCount"] = strconv.Itoa(service.DesiredCount)
		}
		if len(service.PathPatterns) > 0 {
			params["PathPattern"] = strings.Join(service.PathPatterns, ",")
		}
		if len(service.HostPatterns) > 0 {
			params["HostPattern"] = strings.Join(service.HostPatterns, ",")
		}

		return nil
	}
}

func (workflow *serviceWorkflow) serviceEc2Deployer(service *common.Service, stackParams map[string]string, environmentName string, stackUpserter common.StackUpserter, stackWaiter common.StackWaiter) Executor {
	return func() error {

		log.Noticef("Deploying service '%s' to '%s'", workflow.serviceName, environmentName)

		svcStackName := common.CreateStackName(common.StackTypeService, workflow.serviceName, environmentName)

		resolveServiceEnvironment(service, environmentName)
		overrides := common.GetStackOverrides(svcStackName)
		template, err := templates.NewTemplate("service-ec2.yml", service, overrides)
		if err != nil {
			return err
		}

		err = stackUpserter.UpsertStack(svcStackName, template, stackParams, buildServiceTags(workflow.serviceName, environmentName, workflow.envStack.Outputs["provider"], common.StackTypeService, workflow.codeRevision, workflow.repoName))
		if err != nil {
			return err
		}
		log.Debugf("Waiting for stack '%s' to complete", svcStackName)
		stack := stackWaiter.AwaitFinalStatus(svcStackName)
		if stack == nil {
			return fmt.Errorf("Unable to create stack %s", svcStackName)
		}
		if strings.HasSuffix(stack.Status, "ROLLBACK_COMPLETE") || !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		return nil
	}
}

func (workflow *serviceWorkflow) serviceEcsDeployer(service *common.Service, stackParams map[string]string, environmentName string, stackUpserter common.StackUpserter, stackWaiter common.StackWaiter) Executor {
	return func() error {
		log.Noticef("Deploying service '%s' to '%s' from '%s'", workflow.serviceName, environmentName, workflow.serviceImage)

		svcStackName := common.CreateStackName(common.StackTypeService, workflow.serviceName, environmentName)

		resolveServiceEnvironment(service, environmentName)
		overrides := common.GetStackOverrides(svcStackName)
		template, err := templates.NewTemplate("service-ecs.yml", service, overrides)
		if err != nil {
			return err
		}

		err = stackUpserter.UpsertStack(svcStackName, template, stackParams, buildServiceTags(workflow.serviceName, environmentName, workflow.envStack.Outputs["provider"], common.StackTypeService, workflow.codeRevision, workflow.repoName))
		if err != nil {
			return err
		}
		log.Debugf("Waiting for stack '%s' to complete", svcStackName)
		stack := stackWaiter.AwaitFinalStatus(svcStackName)
		if stack == nil {
			return fmt.Errorf("Unable to create stack %s", svcStackName)
		}
		if strings.HasSuffix(stack.Status, "ROLLBACK_COMPLETE") || !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		return nil
	}

}

func resolveServiceEnvironment(service *common.Service, environment string) {
	for key, value := range service.Environment {
		switch value.(type) {
		case map[interface{}]interface{}:
			found := false
			for env, v := range value.(map[interface{}]interface{}) {
				if env.(string) == environment {
					service.Environment[key] = v.(string)
					found = true
				}
			}
			if found != true {
				service.Environment[key] = ""
			}
		case string:
			// do nothing
		default:
			log.Warningf("Unable to resolve environment '%s': %v", key, value)
		}

	}
}

func buildServiceTags(serviceName string, environmentName string, envProvider string, stackType common.StackType, codeRevision string, repoName string) map[string]string {
	return map[string]string{
		"type":        string(stackType),
		"environment": environmentName,
		"provider":    envProvider,
		"service":     serviceName,
		"revision":    codeRevision,
		"repo":        repoName,
	}
}
<|endoftext|>