<commit_before>package server

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/dsnet/compress/brotli"
	"github.com/khlieng/dispatch/assets"
	"github.com/spf13/viper"
)

const longCacheControl = "public, max-age=31536000, immutable"
const disabledCacheControl = "no-cache, no-store, must-revalidate"

type File struct {
	Path         string
	Asset        string
	GzipAsset    []byte
	Hash         string
	ContentType  string
	CacheControl string
	Compressed   bool
}

type h2PushAsset struct {
	path string
	hash string
}

func newH2PushAsset(name string) h2PushAsset {
	return h2PushAsset{
		path: name,
		hash: strings.Split(name, ".")[1],
	}
}

var (
	files []*File

	indexStylesheet      string
	indexScripts         []string
	inlineScript         string
	inlineScriptSha256   string
	inlineScriptSW       string
	inlineScriptSWSha256 string
	serviceWorker        []byte

	h2PushAssets      []h2PushAsset
	h2PushCookieValue string

	contentTypes = map[string]string{
		".js":    "text/javascript",
		".css":   "text/css",
		".woff2": "font/woff2",
		".woff":  "application/font-woff",
		".ttf":   "application/x-font-ttf",
	}

	hstsHeader string
	cspEnabled bool
)

func (d *Dispatch) initFileServer() {
	if viper.GetBool("dev") {
		indexScripts = []string{"bundle.js"}
	} else {
		data, err := assets.Asset("asset-manifest.json")
		if err != nil {
			log.Fatal(err)
		}

		manifest := map[string]string{}
		err = json.Unmarshal(data, &manifest)
		if err != nil {
			log.Fatal(err)
		}

		bootloader := decompressedAsset(manifest["boot.js"])
		runtime := decompressedAsset(manifest["runtime.js"])

		inlineScript = string(runtime)
		inlineScriptSW = string(bootloader) + string(runtime)

		hash := sha256.New()
		hash.Write(runtime)
		inlineScriptSha256 = base64.StdEncoding.EncodeToString(hash.Sum(nil))

		hash.Reset()
		hash.Write(bootloader)
		hash.Write(runtime)
		inlineScriptSWSha256 = base64.StdEncoding.EncodeToString(hash.Sum(nil))

		indexStylesheet = manifest["main.css"]
		indexScripts = []string{
			manifest["vendors~main.js"],
			manifest["main.js"],
		}

		h2PushAssets = []h2PushAsset{
			newH2PushAsset(indexStylesheet),
			newH2PushAsset(indexScripts[0]),
			newH2PushAsset(indexScripts[1]),
		}

		for _, asset := range h2PushAssets {
			h2PushCookieValue += asset.hash
		}

		ignoreAssets := []string{
			manifest["runtime.js"],
			manifest["boot.js"],
			manifest["sw.js"],
		}

	outer:
		for _, assetPath := range manifest {
			for _, ignored := range ignoreAssets {
				if assetPath == ignored {
					continue outer
				}
			}

			file := &File{
				Path:         assetPath,
				Asset:        strings.TrimLeft(assetPath, "/") + ".br",
				ContentType:  contentTypes[filepath.Ext(assetPath)],
				CacheControl: longCacheControl,
				Compressed:   true,
			}

			files = append(files, file)
		}

		fonts, err := assets.AssetDir("font")
		if err != nil {
			log.Fatal(err)
		}

		for _, font := range fonts {
			p := strings.TrimSuffix(font, ".br")

			file := &File{
				Path:         path.Join("font", p),
				Asset:        path.Join("font", font),
				ContentType:  contentTypes[filepath.Ext(p)],
				CacheControl: longCacheControl,
				Compressed:   strings.HasSuffix(font, ".br"),
			}

			files = append(files, file)
		}

		for _, file := range files {
			if file.Compressed {
				data, err := assets.Asset(file.Asset)
				if err != nil {
					log.Fatal(err)
				}

				file.GzipAsset = gzipAsset(data)
			}
		}

		serviceWorker = decompressedAsset("sw.js")
		hash.Reset()
		IndexTemplate(hash, nil, indexStylesheet, inlineScriptSW, indexScripts)
		indexHash := base64.StdEncoding.EncodeToString(hash.Sum(nil))

		serviceWorker = append(serviceWorker, []byte(`
workbox.precaching.precacheAndRoute([{
	revision: '`+indexHash+`',
	url: '/?sw'
}]);
workbox.routing.registerNavigationRoute('/?sw');`)...)

		if viper.GetBool("https.hsts.enabled") && viper.GetBool("https.enabled") {
			hstsHeader = "max-age=" + viper.GetString("https.hsts.max_age")

			if viper.GetBool("https.hsts.include_subdomains") {
				hstsHeader += "; includeSubDomains"
			}
			if viper.GetBool("https.hsts.preload") {
				hstsHeader += "; preload"
			}
		}

		cspEnabled = true
	}
}

func decompressAsset(data []byte) []byte {
	br, err := brotli.NewReader(bytes.NewReader(data), nil)
	if err != nil {
		log.Fatal(err)
	}

	buf := &bytes.Buffer{}
	io.Copy(buf, br)
	return buf.Bytes()
}

func decompressedAsset(name string) []byte {
	asset, err := assets.Asset(strings.TrimLeft(name, "/") + ".br")
	if err != nil {
		log.Fatal(err)
	}
	return decompressAsset(asset)
}

func gzipAsset(data []byte) []byte {
	br, err := brotli.NewReader(bytes.NewReader(data), nil)
	if err != nil {
		log.Fatal(err)
	}

	buf := &bytes.Buffer{}
	gzw, err := gzip.NewWriterLevel(buf, gzip.BestCompression)
	if err != nil {
		log.Fatal(err)
	}

	io.Copy(gzw, br)
	gzw.Close()
	return buf.Bytes()
}

func (d *Dispatch) serveFiles(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/" {
		d.serveIndex(w, r)
		return
	}

	if r.URL.Path == "/sw.js" {
		w.Header().Set("Cache-Control", disabledCacheControl)
		w.Header().Set("Content-Type", "text/javascript")
		w.Header().Set("Content-Length", strconv.Itoa(len(serviceWorker)))
		w.Write(serviceWorker)
		return
	}

	for _, file := range files {
		if strings.HasSuffix(r.URL.Path, file.Path) {
			d.serveFile(w, r, file)
			return
		}
	}

	d.serveIndex(w, r)
}

func (d *Dispatch) serveIndex(w http.ResponseWriter, r *http.Request) {
	state := d.handleAuth(w, r, false)

	_, sw := r.URL.Query()["sw"]

	if cspEnabled {
		var wsSrc string
		if r.TLS != nil {
			wsSrc = "wss://" + r.Host
		} else {
			wsSrc = "ws://" + r.Host
		}

		inlineSha := inlineScriptSha256
		if sw {
			inlineSha = inlineScriptSWSha256
		}

		w.Header().Set("Content-Security-Policy", "default-src 'none'; script-src 'self' 'sha256-"+inlineSha+"'; style-src 'self' 'unsafe-inline'; font-src 'self'; img-src data:; connect-src 'self' "+wsSrc)
	}

	w.Header().Set("Content-Type", "text/html")
	w.Header().Set("Cache-Control", disabledCacheControl)
	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.Header().Set("X-Frame-Options", "deny")
	w.Header().Set("X-XSS-Protection", "1; mode=block")

	if hstsHeader != "" {
		w.Header().Set("Strict-Transport-Security", hstsHeader)
	}

	if pusher, ok := w.(http.Pusher); ok {
		options := &http.PushOptions{
			Header: http.Header{
				"Accept-Encoding": r.Header["Accept-Encoding"],
			},
		}
		cookie, err := r.Cookie("push")
		if err != nil {
			for _, asset := range h2PushAssets {
				pusher.Push(asset.path, options)
			}

			setPushCookie(w, r)
		} else {
			pushed := false

			for i, asset := range h2PushAssets {
				if len(cookie.Value) >= (i+1)*8 &&
					asset.hash != cookie.Value[i*8:(i+1)*8] {
					pusher.Push(asset.path, options)
					pushed = true
				}
			}

			if pushed {
				setPushCookie(w, r)
			}
		}
	}

	var data *indexData
	if !sw {
		data = getIndexData(r, r.URL.EscapedPath(), state)
	}

	inline := inlineScript
	if sw {
		inline = inlineScriptSW
	}

	if strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		w.Header().Set("Content-Encoding", "gzip")

		gzw := gzip.NewWriter(w)
		IndexTemplate(gzw, data, indexStylesheet, inline, indexScripts)
		gzw.Close()
	} else {
		IndexTemplate(w, data, indexStylesheet, inline, indexScripts)
	}
}

func setPushCookie(w http.ResponseWriter, r *http.Request) {
	http.SetCookie(w, &http.Cookie{
		Name:     "push",
		Value:    h2PushCookieValue,
		Path:     "/",
		Expires:  time.Now().AddDate(1, 0, 0),
		HttpOnly: true,
		Secure:   r.TLS != nil,
	})
}

func (d *Dispatch) serveFile(w http.ResponseWriter, r *http.Request, file *File) {
	data, err := assets.Asset(file.Asset)
	if err != nil {
		http.Error(w, "", http.StatusInternalServerError)
		return
	}

	if file.CacheControl != "" {
		w.Header().Set("Cache-Control", file.CacheControl)
	}

	w.Header().Set("Content-Type", file.ContentType)

	if file.Compressed && strings.Contains(r.Header.Get("Accept-Encoding"), "br") {
		w.Header().Set("Content-Encoding", "br")
		w.Header().Set("Content-Length", strconv.Itoa(len(data)))
		w.Write(data)
	} else if file.Compressed && strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		w.Header().Set("Content-Encoding", "gzip")
		w.Header().Set("Content-Length", strconv.Itoa(len(file.GzipAsset)))
		w.Write(file.GzipAsset)
	} else if !file.Compressed {
		w.Header().Set("Content-Length", strconv.Itoa(len(data)))
		w.Write(data)
	} else {
		gzr, err := gzip.NewReader(bytes.NewReader(file.GzipAsset))
		buf, err := ioutil.ReadAll(gzr)
		if err != nil {
			http.Error(w, "", http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Length", strconv.Itoa(len(buf)))
		w.Write(buf)
	}
}
<commit_msg>Dont do auth on service worker index page requests<commit_after>package server

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/dsnet/compress/brotli"
	"github.com/khlieng/dispatch/assets"
	"github.com/spf13/viper"
)

const longCacheControl = "public, max-age=31536000, immutable"
const disabledCacheControl = "no-cache, no-store, must-revalidate"

type File struct {
	Path         string
	Asset        string
	GzipAsset    []byte
	Hash         string
	ContentType  string
	CacheControl string
	Compressed   bool
}

type h2PushAsset struct {
	path string
	hash string
}

func newH2PushAsset(name string) h2PushAsset {
	return h2PushAsset{
		path: name,
		hash: strings.Split(name, ".")[1],
	}
}

var (
	files []*File

	indexStylesheet      string
	indexScripts         []string
	inlineScript         string
	inlineScriptSha256   string
	inlineScriptSW       string
	inlineScriptSWSha256 string
	serviceWorker        []byte

	h2PushAssets      []h2PushAsset
	h2PushCookieValue string

	contentTypes = map[string]string{
		".js":    "text/javascript",
		".css":   "text/css",
		".woff2": "font/woff2",
		".woff":  "application/font-woff",
		".ttf":   "application/x-font-ttf",
	}

	hstsHeader string
	cspEnabled bool
)

func (d *Dispatch) initFileServer() {
	if viper.GetBool("dev") {
		indexScripts = []string{"bundle.js"}
	} else {
		data, err := assets.Asset("asset-manifest.json")
		if err != nil {
			log.Fatal(err)
		}

		manifest := map[string]string{}
		err = json.Unmarshal(data, &manifest)
		if err != nil {
			log.Fatal(err)
		}

		bootloader := decompressedAsset(manifest["boot.js"])
		runtime := decompressedAsset(manifest["runtime.js"])

		inlineScript = string(runtime)
		inlineScriptSW = string(bootloader) + string(runtime)

		hash := sha256.New()
		hash.Write(runtime)
		inlineScriptSha256 = base64.StdEncoding.EncodeToString(hash.Sum(nil))

		hash.Reset()
		hash.Write(bootloader)
		hash.Write(runtime)
		inlineScriptSWSha256 = base64.StdEncoding.EncodeToString(hash.Sum(nil))

		indexStylesheet = manifest["main.css"]
		indexScripts = []string{
			manifest["vendors~main.js"],
			manifest["main.js"],
		}

		h2PushAssets = []h2PushAsset{
			newH2PushAsset(indexStylesheet),
			newH2PushAsset(indexScripts[0]),
			newH2PushAsset(indexScripts[1]),
		}

		for _, asset := range h2PushAssets {
			h2PushCookieValue += asset.hash
		}

		ignoreAssets := []string{
			manifest["runtime.js"],
			manifest["boot.js"],
			manifest["sw.js"],
		}

	outer:
		for _, assetPath := range manifest {
			for _, ignored := range ignoreAssets {
				if assetPath == ignored {
					continue outer
				}
			}

			file := &File{
				Path:         assetPath,
				Asset:        strings.TrimLeft(assetPath, "/") + ".br",
				ContentType:  contentTypes[filepath.Ext(assetPath)],
				CacheControl: longCacheControl,
				Compressed:   true,
			}

			files = append(files, file)
		}

		fonts, err := assets.AssetDir("font")
		if err != nil {
			log.Fatal(err)
		}

		for _, font := range fonts {
			p := strings.TrimSuffix(font, ".br")

			file := &File{
				Path:         path.Join("font", p),
				Asset:        path.Join("font", font),
				ContentType:  contentTypes[filepath.Ext(p)],
				CacheControl: longCacheControl,
				Compressed:   strings.HasSuffix(font, ".br"),
			}

			files = append(files, file)
		}

		for _, file := range files {
			if file.Compressed {
				data, err := assets.Asset(file.Asset)
				if err != nil {
					log.Fatal(err)
				}

				file.GzipAsset = gzipAsset(data)
			}
		}

		serviceWorker = decompressedAsset("sw.js")
		hash.Reset()
		IndexTemplate(hash, nil, indexStylesheet, inlineScriptSW, indexScripts)
		indexHash := base64.StdEncoding.EncodeToString(hash.Sum(nil))

		serviceWorker = append(serviceWorker, []byte(`
workbox.precaching.precacheAndRoute([{
	revision: '`+indexHash+`',
	url: '/?sw'
}]);
workbox.routing.registerNavigationRoute('/?sw');`)...)

		if viper.GetBool("https.hsts.enabled") && viper.GetBool("https.enabled") {
			hstsHeader = "max-age=" + viper.GetString("https.hsts.max_age")

			if viper.GetBool("https.hsts.include_subdomains") {
				hstsHeader += "; includeSubDomains"
			}
			if viper.GetBool("https.hsts.preload") {
				hstsHeader += "; preload"
			}
		}

		cspEnabled = true
	}
}

func decompressAsset(data []byte) []byte {
	br, err := brotli.NewReader(bytes.NewReader(data), nil)
	if err != nil {
		log.Fatal(err)
	}

	buf := &bytes.Buffer{}
	io.Copy(buf, br)
	return buf.Bytes()
}

func decompressedAsset(name string) []byte {
	asset, err := assets.Asset(strings.TrimLeft(name, "/") + ".br")
	if err != nil {
		log.Fatal(err)
	}
	return decompressAsset(asset)
}

func gzipAsset(data []byte) []byte {
	br, err := brotli.NewReader(bytes.NewReader(data), nil)
	if err != nil {
		log.Fatal(err)
	}

	buf := &bytes.Buffer{}
	gzw, err := gzip.NewWriterLevel(buf, gzip.BestCompression)
	if err != nil {
		log.Fatal(err)
	}

	io.Copy(gzw, br)
	gzw.Close()
	return buf.Bytes()
}

func (d *Dispatch) serveFiles(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/" {
		d.serveIndex(w, r)
		return
	}

	if r.URL.Path == "/sw.js" {
		w.Header().Set("Cache-Control", disabledCacheControl)
		w.Header().Set("Content-Type", "text/javascript")
		w.Header().Set("Content-Length", strconv.Itoa(len(serviceWorker)))
		w.Write(serviceWorker)
		return
	}

	for _, file := range files {
		if strings.HasSuffix(r.URL.Path, file.Path) {
			d.serveFile(w, r, file)
			return
		}
	}

	d.serveIndex(w, r)
}

func (d *Dispatch) serveIndex(w http.ResponseWriter, r *http.Request) {
	if pusher, ok := w.(http.Pusher); ok {
		options := &http.PushOptions{
			Header: http.Header{
				"Accept-Encoding": r.Header["Accept-Encoding"],
			},
		}

		cookie, err := r.Cookie("push")
		if err != nil {
			for _, asset := range h2PushAssets {
				pusher.Push(asset.path, options)
			}

			setPushCookie(w, r)
		} else {
			pushed := false

			for i, asset := range h2PushAssets {
				if len(cookie.Value) >= (i+1)*8 &&
					asset.hash != cookie.Value[i*8:(i+1)*8] {
					pusher.Push(asset.path, options)
					pushed = true
				}
			}

			if pushed {
				setPushCookie(w, r)
			}
		}
	}

	_, sw := r.URL.Query()["sw"]

	if cspEnabled {
		var wsSrc string
		if r.TLS != nil {
			wsSrc = "wss://" + r.Host
		} else {
			wsSrc = "ws://" + r.Host
		}

		inlineSha := inlineScriptSha256
		if sw {
			inlineSha = inlineScriptSWSha256
		}

		w.Header().Set("Content-Security-Policy", "default-src 'none'; script-src 'self' 'sha256-"+inlineSha+"'; style-src 'self' 'unsafe-inline'; font-src 'self'; img-src data:; connect-src 'self' "+wsSrc)
	}

	w.Header().Set("Content-Type", "text/html")
	w.Header().Set("Cache-Control", disabledCacheControl)
	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.Header().Set("X-Frame-Options", "deny")
	w.Header().Set("X-XSS-Protection", "1; mode=block")

	if hstsHeader != "" {
		w.Header().Set("Strict-Transport-Security", hstsHeader)
	}

	var data *indexData
	inline := inlineScriptSW
	if !sw {
		data = getIndexData(r, r.URL.EscapedPath(), d.handleAuth(w, r, false))
		inline = inlineScript
	}

	if strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		w.Header().Set("Content-Encoding", "gzip")

		gzw := gzip.NewWriter(w)
		IndexTemplate(gzw, data, indexStylesheet, inline, indexScripts)
		gzw.Close()
	} else {
		IndexTemplate(w, data, indexStylesheet, inline, indexScripts)
	}
}

func setPushCookie(w http.ResponseWriter, r *http.Request) {
	http.SetCookie(w, &http.Cookie{
		Name:     "push",
		Value:    h2PushCookieValue,
		Path:     "/",
		Expires:  time.Now().AddDate(1, 0, 0),
		HttpOnly: true,
		Secure:   r.TLS != nil,
	})
}

func (d *Dispatch) serveFile(w http.ResponseWriter, r *http.Request, file *File) {
	data, err := assets.Asset(file.Asset)
	if err != nil {
		http.Error(w, "", http.StatusInternalServerError)
		return
	}

	if file.CacheControl != "" {
		w.Header().Set("Cache-Control", file.CacheControl)
	}

	w.Header().Set("Content-Type", file.ContentType)

	if file.Compressed && strings.Contains(r.Header.Get("Accept-Encoding"), "br") {
		w.Header().Set("Content-Encoding", "br")
		w.Header().Set("Content-Length", strconv.Itoa(len(data)))
		w.Write(data)
	} else if file.Compressed && strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		w.Header().Set("Content-Encoding", "gzip")
		w.Header().Set("Content-Length", strconv.Itoa(len(file.GzipAsset)))
		w.Write(file.GzipAsset)
	} else if !file.Compressed {
		w.Header().Set("Content-Length", strconv.Itoa(len(data)))
		w.Write(data)
	} else {
		gzr, err := gzip.NewReader(bytes.NewReader(file.GzipAsset))
		buf, err := ioutil.ReadAll(gzr)
		if err != nil {
			http.Error(w, "", http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Length", strconv.Itoa(len(buf)))
		w.Write(buf)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"net"
	"testing"

	"github.com/admiraldolphin/govhack2017/server/game"

	"github.com/kr/pretty"
)

func TestGame(t *testing.T) {
	s := server{state: game.New(testDeck)}
	r := &response{}
	if err := s.listenAndServe("localhost:0"); err != nil {
		t.Fatalf("Couldn't start game server: %v", err)
	}
	defer s.Close()

	conn0, err := net.Dial("tcp", s.Addr().String())
	if err != nil {
		t.Fatalf("Couldn't connect player 0 to game server: %v", err)
	}
	defer conn0.Close()

	send0 := json.NewEncoder(conn0)
	recv0 := json.NewDecoder(conn0)

	// Should get a response immediately
	if err := recv0.Decode(r); err != nil {
		t.Fatalf("Couldn't get state immediately: %v", err)
	}
	if got, want := r.PlayerID, 0; got != want {
		t.Errorf("response.PlayerID = %d, want %d", got, want)
	}
	if got, want := r.State.State, game.StateLobby; got != want {
		t.Errorf("response.State.State = %d, want %d", got, want)
	}

	conn1, err := net.Dial("tcp", s.Addr().String())
	if err != nil {
		t.Fatalf("Couldn't connect player 1 to game server: %v", err)
	}
	defer conn1.Close()

	send1 := json.NewEncoder(conn1)
	recv1 := json.NewDecoder(conn1)

	// Should get a response immediately
	if err := recv1.Decode(r); err != nil {
		t.Fatalf("Couldn't get state immediately: %v", err)
	}
	if got, want := r.PlayerID, 1; got != want {
		t.Errorf("response.PlayerID = %d, want %d", got, want)
	}
	if got, want := r.State.State, game.StateLobby; got != want {
		t.Errorf("response.State.State = %d, want %d", got, want)
	}

	// Play a game!
	g := []struct {
		send   *json.Encoder
		recv   *json.Decoder
		action *game.Action
		want   *game.State
	}{
		{
			// Player 1 starts the game
			send:   send1,
			recv:   recv1,
			action: &game.Action{Act: game.ActStartGame},
			want: &game.State{
				State: game.StateInGame,
				Players: map[int]*game.Player{
					0: {
						Name: "Player 0",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action0"}},
								{Card: &game.ActionCard{Name: "action1"}},
								{Card: &game.ActionCard{Name: "action2"}},
								{Card: &game.ActionCard{Name: "action3"}},
								{Card: &game.ActionCard{Name: "action4"}},
								{Card: &game.ActionCard{Name: "action5"}},
								{Card: &game.ActionCard{Name: "action6"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person0"}},
								{Card: &game.PersonCard{Name: "person1"}},
								{Card: &game.PersonCard{Name: "person2"}},
								{Card: &game.PersonCard{Name: "person3"}},
								{Card: &game.PersonCard{Name: "person4"}},
							},
						},
					},
					1: {
						Name: "Player 1",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action7"}},
								{Card: &game.ActionCard{Name: "action8"}},
								{Card: &game.ActionCard{Name: "action9"}},
								{Card: &game.ActionCard{Name: "action10"}},
								{Card: &game.ActionCard{Name: "action11"}},
								{Card: &game.ActionCard{Name: "action12"}},
								{Card: &game.ActionCard{Name: "action13"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person5"}},
								{Card: &game.PersonCard{Name: "person6"}},
								{Card: &game.PersonCard{Name: "person7"}},
								{Card: &game.PersonCard{Name: "person8"}},
								{Card: &game.PersonCard{Name: "person9"}},
							},
						},
					},
				},
				WhoseTurn: 0,
				Clock:     0,
			},
		},
		{
			// Player 0 plays a card
			send:   send0,
			recv:   recv0,
			action: &game.Action{Act: game.ActPlayCard, Card: 3},
			want: &game.State{
				State: game.StateInGame,
				Players: map[int]*game.Player{
					0: {
						Name: "Player 0",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action0"}},
								{Card: &game.ActionCard{Name: "action1"}},
								{Card: &game.ActionCard{Name: "action2"}},
								{Card: &game.ActionCard{Name: "action14"}},
								{Card: &game.ActionCard{Name: "action4"}},
								{Card: &game.ActionCard{Name: "action5"}},
								{Card: &game.ActionCard{Name: "action6"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person0"}},
								{Card: &game.PersonCard{Name: "person1"}},
								{Card: &game.PersonCard{Name: "person2"}},
								{Card: &game.PersonCard{Name: "person3"}},
								{Card: &game.PersonCard{Name: "person4"}},
							},
						},
						Played: []*game.ActionCardState{
							{Card: &game.ActionCard{Name: "action3"}, Played: true},
						},
					},
					1: {
						Name: "Player 1",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action7"}},
								{Card: &game.ActionCard{Name: "action8"}},
								{Card: &game.ActionCard{Name: "action9"}},
								{Card: &game.ActionCard{Name: "action10"}},
								{Card: &game.ActionCard{Name: "action11"}},
								{Card: &game.ActionCard{Name: "action12"}},
								{Card: &game.ActionCard{Name: "action13"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person5"}},
								{Card: &game.PersonCard{Name: "person6"}},
								{Card: &game.PersonCard{Name: "person7"}},
								{Card: &game.PersonCard{Name: "person8"}},
								{Card: &game.PersonCard{Name: "person9"}},
							},
						},
					},
				},
				WhoseTurn: 1,
				Clock:     0,
			},
		},
		{
			// Player 0 does a no-op
			send:   send0,
			recv:   recv0,
			action: &game.Action{Act: game.ActNoOp},
			want: &game.State{
				State: game.StateInGame,
				Players: map[int]*game.Player{
					0: {
						Name: "Player 0",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action0"}},
								{Card: &game.ActionCard{Name: "action1"}},
								{Card: &game.ActionCard{Name: "action2"}},
								{Card: &game.ActionCard{Name: "action14"}},
								{Card: &game.ActionCard{Name: "action4"}},
								{Card: &game.ActionCard{Name: "action5"}},
								{Card: &game.ActionCard{Name: "action6"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person0"}},
								{Card: &game.PersonCard{Name: "person1"}},
								{Card: &game.PersonCard{Name: "person2"}},
								{Card: &game.PersonCard{Name: "person3"}},
								{Card: &game.PersonCard{Name: "person4"}},
							},
						},
						Played: []*game.ActionCardState{
							{Card: &game.ActionCard{Name: "action3"}, Played: true},
						},
					},
					1: {
						Name: "Player 1",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action7"}},
								{Card: &game.ActionCard{Name: "action8"}},
								{Card: &game.ActionCard{Name: "action9"}},
								{Card: &game.ActionCard{Name: "action10"}},
								{Card: &game.ActionCard{Name: "action11"}},
								{Card: &game.ActionCard{Name: "action12"}},
								{Card: &game.ActionCard{Name: "action13"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person5"}},
								{Card: &game.PersonCard{Name: "person6"}},
								{Card: &game.PersonCard{Name: "person7"}},
								{Card: &game.PersonCard{Name: "person8"}},
								{Card: &game.PersonCard{Name: "person9"}},
							},
						},
					},
				},
				WhoseTurn: 1,
				Clock:     0,
			},
		},
		{
			// Player 1 discards a card
			send:   send1,
			recv:   recv1,
			action: &game.Action{Act: game.ActDiscard, Card: 2},
			want: &game.State{
				State: game.StateInGame,
				Players: map[int]*game.Player{
					0: {
						Name: "Player 0",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action0"}},
								{Card: &game.ActionCard{Name: "action1"}},
								{Card: &game.ActionCard{Name: "action2"}},
								{Card: &game.ActionCard{Name: "action14"}},
								{Card: &game.ActionCard{Name: "action4"}},
								{Card: &game.ActionCard{Name: "action5"}},
								{Card: &game.ActionCard{Name: "action6"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person0"}},
								{Card: &game.PersonCard{Name: "person1"}},
								{Card: &game.PersonCard{Name: "person2"}},
								{Card: &game.PersonCard{Name: "person3"}},
								{Card: &game.PersonCard{Name: "person4"}},
							},
						},
						Played: []*game.ActionCardState{
							{Card: &game.ActionCard{Name: "action3"}, Played: true},
						},
					},
					1: {
						Name: "Player 1",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action7"}},
								{Card: &game.ActionCard{Name: "action8"}},
								{Card: &game.ActionCard{Name: "action15"}},
								{Card: &game.ActionCard{Name: "action10"}},
								{Card: &game.ActionCard{Name: "action11"}},
								{Card: &game.ActionCard{Name: "action12"}},
								{Card: &game.ActionCard{Name: "action13"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person5"}},
								{Card: &game.PersonCard{Name: "person6"}},
								{Card: &game.PersonCard{Name: "person7"}},
								{Card: &game.PersonCard{Name: "person8"}},
								{Card: &game.PersonCard{Name: "person9"}},
							},
						},
						Discarded: []*game.ActionCardState{
							{Card: &game.ActionCard{Name: "action9"}, Discarded: true},
						},
					},
				},
				WhoseTurn: 0,
				Clock:     1,
			},
		},
	}

	for i, p := range g {
		if err := p.send.Encode(p.action); err != nil {
			t.Fatalf("Message %d [%v] got error %v", i, p.action, err)
		}
		if err := recv0.Decode(r); err != nil {
			t.Errorf("After message %d [%v]: recv0.Decode = error %v", i, p.action, err)
		}
		if err := recv1.Decode(r); err != nil {
			t.Errorf("After message %d [%v]: recv1.Decode = error %v", i, p.action, err)
		}

		if got, want := r.State.State, p.want.State; got != want {
			t.Errorf("After message %d [%v]: state.State = %#v, want %#v", i, p.action, got, want)
		}
		if diff := pretty.Diff(p.want.Players, r.State.Players); len(diff) > 0 {
			t.Errorf("After message %d [%v]: state.Players diff (-want +got): %v", i, p.action, diff)
		}
		if got, want := r.State.WhoseTurn, p.want.WhoseTurn; got != want {
			t.Errorf("After message %d [%v]: state.WhoseTurn = %#v, want %#v", i, p.action, got, want)
		}
		if got, want := r.State.Clock, p.want.Clock; got != want {
			t.Errorf("After message %d [%v]: state.Clock = %#v, want %#v", i, p.action, got, want)
		}
	}
}
<commit_msg>Fix test<commit_after>package main

import (
	"encoding/json"
	"net"
	"testing"

	"github.com/admiraldolphin/govhack2017/server/game"

	"github.com/kr/pretty"
)

func TestGame(t *testing.T) {
	s := server{state: game.New(testDeck)}
	r := &response{}
	if err := s.listenAndServe("localhost:0"); err != nil {
		t.Fatalf("Couldn't start game server: %v", err)
	}
	defer s.Close()

	conn0, err := net.Dial("tcp", s.Addr().String())
	if err != nil {
		t.Fatalf("Couldn't connect player 0 to game server: %v", err)
	}
	defer conn0.Close()

	send0 := json.NewEncoder(conn0)
	recv0 := json.NewDecoder(conn0)

	// Should get a response immediately
	if err := recv0.Decode(r); err != nil {
		t.Fatalf("Couldn't get state immediately: %v", err)
	}
	if got, want := r.PlayerID, 0; got != want {
		t.Errorf("response.PlayerID = %d, want %d", got, want)
	}
	if got, want := r.State.State, game.StateLobby; got != want {
		t.Errorf("response.State.State = %d, want %d", got, want)
	}

	conn1, err := net.Dial("tcp", s.Addr().String())
	if err != nil {
		t.Fatalf("Couldn't connect player 1 to game server: %v", err)
	}
	defer conn1.Close()

	send1 := json.NewEncoder(conn1)
	recv1 := json.NewDecoder(conn1)

	// Should get a response immediately
	if err := recv1.Decode(r); err != nil {
		t.Fatalf("Couldn't get state immediately: %v", err)
	}
	if got, want := r.PlayerID, 1; got != want {
		t.Errorf("response.PlayerID = %d, want %d", got, want)
	}
	if got, want := r.State.State, game.StateLobby; got != want {
		t.Errorf("response.State.State = %d, want %d", got, want)
	}

	// Play a game!
	g := []struct {
		send   *json.Encoder
		recv   *json.Decoder
		action *game.Action
		want   *game.State
	}{
		{
			// Player 1 starts the game
			send:   send1,
			recv:   recv1,
			action: &game.Action{Act: game.ActStartGame},
			want: &game.State{
				State: game.StateInGame,
				Players: map[int]*game.Player{
					0: {
						Name: "Player 0",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action0"}},
								{Card: &game.ActionCard{Name: "action1"}},
								{Card: &game.ActionCard{Name: "action2"}},
								{Card: &game.ActionCard{Name: "action3"}},
								{Card: &game.ActionCard{Name: "action4"}},
								{Card: &game.ActionCard{Name: "action5"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person0"}},
								{Card: &game.PersonCard{Name: "person1"}},
								{Card: &game.PersonCard{Name: "person2"}},
								{Card: &game.PersonCard{Name: "person3"}},
								{Card: &game.PersonCard{Name: "person4"}},
							},
						},
					},
					1: {
						Name: "Player 1",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action6"}},
								{Card: &game.ActionCard{Name: "action7"}},
								{Card: &game.ActionCard{Name: "action8"}},
								{Card: &game.ActionCard{Name: "action9"}},
								{Card: &game.ActionCard{Name: "action10"}},
								{Card: &game.ActionCard{Name: "action11"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person5"}},
								{Card: &game.PersonCard{Name: "person6"}},
								{Card: &game.PersonCard{Name: "person7"}},
								{Card: &game.PersonCard{Name: "person8"}},
								{Card: &game.PersonCard{Name: "person9"}},
							},
						},
					},
				},
				WhoseTurn: 0,
				Clock:     4,
			},
		},
		{
			// Player 0 plays a card
			send:   send0,
			recv:   recv0,
			action: &game.Action{Act: game.ActPlayCard, Card: 3},
			want: &game.State{
				State: game.StateInGame,
				Players: map[int]*game.Player{
					0: {
						Name: "Player 0",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action0"}},
								{Card: &game.ActionCard{Name: "action1"}},
								{Card: &game.ActionCard{Name: "action2"}},
								{Card: &game.ActionCard{Name: "action12"}},
								{Card: &game.ActionCard{Name: "action4"}},
								{Card: &game.ActionCard{Name: "action5"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person0"}},
								{Card: &game.PersonCard{Name: "person1"}},
								{Card: &game.PersonCard{Name: "person2"}},
								{Card: &game.PersonCard{Name: "person3"}},
								{Card: &game.PersonCard{Name: "person4"}},
							},
						},
						Played: []*game.ActionCardState{
							{Card: &game.ActionCard{Name: "action3"}, Played: true},
						},
					},
					1: {
						Name: "Player 1",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action6"}},
								{Card: &game.ActionCard{Name: "action7"}},
								{Card: &game.ActionCard{Name: "action8"}},
								{Card: &game.ActionCard{Name: "action9"}},
								{Card: &game.ActionCard{Name: "action10"}},
								{Card: &game.ActionCard{Name: "action11"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person5"}},
								{Card: &game.PersonCard{Name: "person6"}},
								{Card: &game.PersonCard{Name: "person7"}},
								{Card: &game.PersonCard{Name: "person8"}},
								{Card: &game.PersonCard{Name: "person9"}},
							},
						},
					},
				},
				WhoseTurn: 1,
				Clock:     4,
			},
		},
		{
			// Player 0 does a no-op
			send:   send0,
			recv:   recv0,
			action: &game.Action{Act: game.ActNoOp},
			want: &game.State{
				State: game.StateInGame,
				Players: map[int]*game.Player{
					0: {
						Name: "Player 0",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action0"}},
								{Card: &game.ActionCard{Name: "action1"}},
								{Card: &game.ActionCard{Name: "action2"}},
								{Card: &game.ActionCard{Name: "action12"}},
								{Card: &game.ActionCard{Name: "action4"}},
								{Card: &game.ActionCard{Name: "action5"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person0"}},
								{Card: &game.PersonCard{Name: "person1"}},
								{Card: &game.PersonCard{Name: "person2"}},
								{Card: &game.PersonCard{Name: "person3"}},
								{Card: &game.PersonCard{Name: "person4"}},
							},
						},
						Played: []*game.ActionCardState{
							{Card: &game.ActionCard{Name: "action3"}, Played: true},
						},
					},
					1: {
						Name: "Player 1",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action6"}},
								{Card: &game.ActionCard{Name: "action7"}},
								{Card: &game.ActionCard{Name: "action8"}},
								{Card: &game.ActionCard{Name: "action9"}},
								{Card: &game.ActionCard{Name: "action10"}},
								{Card: &game.ActionCard{Name: "action11"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person5"}},
								{Card: &game.PersonCard{Name: "person6"}},
								{Card: &game.PersonCard{Name: "person7"}},
								{Card: &game.PersonCard{Name: "person8"}},
								{Card: &game.PersonCard{Name: "person9"}},
							},
						},
					},
				},
				WhoseTurn: 1,
				Clock:     4,
			},
		},
		{
			// Player 1 discards a card
			send:   send1,
			recv:   recv1,
			action: &game.Action{Act: game.ActDiscard, Card: 2},
			want: &game.State{
				State: game.StateInGame,
				Players: map[int]*game.Player{
					0: {
						Name: "Player 0",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action0"}},
								{Card: &game.ActionCard{Name: "action1"}},
								{Card: &game.ActionCard{Name: "action2"}},
								{Card: &game.ActionCard{Name: "action12"}},
								{Card: &game.ActionCard{Name: "action4"}},
								{Card: &game.ActionCard{Name: "action5"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person0"}},
								{Card: &game.PersonCard{Name: "person1"}},
								{Card: &game.PersonCard{Name: "person2"}},
								{Card: &game.PersonCard{Name: "person3"}},
								{Card: &game.PersonCard{Name: "person4"}},
							},
						},
						Played: []*game.ActionCardState{
							{Card: &game.ActionCard{Name: "action3"}, Played: true},
						},
					},
					1: {
						Name: "Player 1",
						Hand: &game.HandState{
							Actions: []*game.ActionCardState{
								{Card: &game.ActionCard{Name: "action6"}},
								{Card: &game.ActionCard{Name: "action7"}},
								{Card: &game.ActionCard{Name: "action13"}},
								{Card: &game.ActionCard{Name: "action9"}},
								{Card: &game.ActionCard{Name: "action10"}},
								{Card: &game.ActionCard{Name: "action11"}},
							},
							People: []*game.PersonCardState{
								{Card: &game.PersonCard{Name: "person5"}},
								{Card: &game.PersonCard{Name: "person6"}},
								{Card: &game.PersonCard{Name: "person7"}},
								{Card: &game.PersonCard{Name: "person8"}},
								{Card: &game.PersonCard{Name: "person9"}},
							},
						},
						Discarded: []*game.ActionCardState{
							{Card: &game.ActionCard{Name: "action8"}, Discarded: true},
						},
					},
				},
				WhoseTurn: 0,
				Clock:     3,
			},
		},
	}

	for i, p := range g {
		if err := p.send.Encode(p.action); err != nil {
			t.Fatalf("Message %d [%v] got error %v", i, p.action, err)
		}
		if err := recv0.Decode(r); err != nil {
			t.Errorf("After message %d [%v]: recv0.Decode = error %v", i, p.action, err)
		}
		if err := recv1.Decode(r); err != nil {
			t.Errorf("After message %d [%v]: recv1.Decode = error %v", i, p.action, err)
		}

		if got, want := r.State.State, p.want.State; got != want {
			t.Errorf("After message %d [%v]: state.State = %#v, want %#v", i, p.action, got, want)
		}
		if diff := pretty.Diff(p.want.Players, r.State.Players); len(diff) > 0 {
			t.Errorf("After message %d [%v]: state.Players diff (-want +got): %v", i, p.action, diff)
		}
		if got, want := r.State.WhoseTurn, p.want.WhoseTurn; got != want {
			t.Errorf("After message %d [%v]: state.WhoseTurn = %#v, want %#v", i, p.action, got, want)
		}
		if got, want := r.State.Clock, p.want.Clock; got != want {
			t.Errorf("After message %d [%v]: state.Clock = %#v, want %#v", i, p.action, got, want)
		}
	}
}
<|endoftext|>


<commit_before>package services

import (
	"errors"
	"fmt"
	"github.com/stitchfix/flotilla-os/log"
	"math/rand"
	"strings"
	"time"

	"github.com/stitchfix/flotilla-os/clients/cluster"
	"github.com/stitchfix/flotilla-os/clients/registry"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/exceptions"
	"github.com/stitchfix/flotilla-os/execution/engine"
	"github.com/stitchfix/flotilla-os/state"
)

//
// ExecutionService interacts with the state manager and queue manager to queue runs, and perform
// CRUD operations on them
// * Acts as an intermediary layer between state and the execution engine
//
type ExecutionService interface {
	Create(definitionID string, clusterName string, env *state.EnvList, ownerID string, command *string, memory *int64, cpu *int64, engine *string, ephemeralStorage *int64, nodeLifecycle *string) (state.Run, error)
	CreateByAlias(alias string, clusterName string, env *state.EnvList, ownerID string, command *string, memory *int64, cpu *int64, engine *string, ephemeralStorage *int64, nodeLifecycle *string) (state.Run, error)
	List(
		limit int,
		offset int,
		sortOrder string,
		sortField string,
		filters map[string][]string,
		envFilters map[string]string) (state.RunList, error)
	Get(runID string) (state.Run, error)
	UpdateStatus(runID string, status string, exitCode *int64) error
	Terminate(runID string) error
	ReservedVariables() []string
	ListClusters() ([]string, error)
	GetEvents(run state.Run) (state.RunEventList, error)
}

type executionService struct {
	stateManager       state.Manager
	ecsClusterClient   cluster.Client
	eksClusterClient   cluster.Client
	registryClient     registry.Client
	ecsExecutionEngine engine.Engine
	eksExecutionEngine engine.Engine
	reservedEnv        map[string]func(run state.Run) string
	eksClusterOverride string
	eksOverridePercent int
}

func (es *executionService) GetEvents(run state.Run) (state.RunEventList, error) {
	return es.eksExecutionEngine.GetEvents(run)
}

//
// NewExecutionService configures and returns an ExecutionService
//
func NewExecutionService(conf config.Config,
	ecsExecutionEngine engine.Engine,
	eksExecutionEngine engine.Engine,
	sm state.Manager,
	ecsClusterClient cluster.Client,
	eksClusterClient cluster.Client,
	rc registry.Client,
	log log.Logger) (ExecutionService, error) {
	es := executionService{
		stateManager:       sm,
		ecsClusterClient:   ecsClusterClient,
		eksClusterClient:   eksClusterClient,
		registryClient:     rc,
		ecsExecutionEngine: ecsExecutionEngine,
		eksExecutionEngine: eksExecutionEngine,
	}
	//
	// Reserved environment variables dynamically generated
	// per run

	ownerKey := conf.GetString("owner_id_var")
	if len(ownerKey) == 0 {
		ownerKey = "FLOTILLA_RUN_OWNER_ID"
	}

	es.eksClusterOverride = conf.GetString("eks.cluster_override")
	es.eksOverridePercent = conf.GetInt("eks.cluster_override_percent")

	es.reservedEnv = map[string]func(run state.Run) string{
		"FLOTILLA_SERVER_MODE": func(run state.Run) string {
			return conf.GetString("flotilla_mode")
		},
		"FLOTILLA_RUN_ID": func(run state.Run) string {
			return run.RunID
		},
		ownerKey: func(run state.Run) string {
			return run.User
		},
	}
	// Warm cached cluster list
	_, _ = es.ecsClusterClient.ListClusters()
	return &es, nil
}

//
// ReservedVariables returns the list of reserved run environment variable
// names
//
func (es *executionService) ReservedVariables() []string {
	var keys []string
	for k := range es.reservedEnv {
		keys = append(keys, k)
	}
	return keys
}

//
// Create constructs and queues a new Run on the cluster specified
//
func (es *executionService) Create(definitionID string, clusterName string, env *state.EnvList, ownerID string, command *string, memory *int64, cpu *int64, engine *string, ephemeralStorage *int64, nodeLifecycle *string) (state.Run, error) {

	// Ensure definition exists
	definition, err := es.stateManager.GetDefinition(definitionID)
	if err != nil {
		return state.Run{}, err
	}

	if engine == nil {
		engine = &state.DefaultEngine
	}

	// Added to facilitate migration of ECS jobs to EKS.
	if engine != &state.EKSEngine && es.eksOverridePercent > 0 && *definition.Privileged == false {
		modulo := 100 / es.eksOverridePercent
		if rand.Int()%modulo == 0 {
			clusterName = es.eksClusterOverride
			engine = &state.EKSEngine

			if strings.Contains(strings.ToLower(definition.Alias), "ondemand") {
				nodeLifecycle = &state.OndemandLifecycle
			} else {
				nodeLifecycle = &state.SpotLifecycle
			}
		}
	}

	return es.createFromDefinition(definition, clusterName, env, ownerID, command, memory, cpu, engine, nodeLifecycle, ephemeralStorage)
}

//
// Create constructs and queues a new Run on the cluster specified, based on an alias
//
func (es *executionService) CreateByAlias(alias string, clusterName string, env *state.EnvList, ownerID string, command *string, memory *int64, cpu *int64, engine *string, ephemeralStorage *int64, nodeLifecycle *string) (state.Run, error) {

	// Ensure definition exists
	definition, err := es.stateManager.GetDefinitionByAlias(alias)
	if err != nil {
		return state.Run{}, err
	}

	if engine == nil {
		engine = &state.DefaultEngine
	}

	// Added to facilitate migration of ECS jobs to EKS.
	if engine != &state.EKSEngine && es.eksOverridePercent > 0 && *definition.Privileged == false {
		modulo := 100 / es.eksOverridePercent
		if rand.Int()%modulo == 0 {
			clusterName = es.eksClusterOverride
			engine = &state.EKSEngine

			if strings.Contains(strings.ToLower(definition.Alias), "ondemand") {
				nodeLifecycle = &state.OndemandLifecycle
			} else {
				nodeLifecycle = &state.SpotLifecycle
			}
		}
	}

	return es.createFromDefinition(definition, clusterName, env, ownerID, command, memory, cpu, engine, nodeLifecycle, ephemeralStorage)
}

func (es *executionService) createFromDefinition(definition state.Definition, clusterName string, env *state.EnvList, ownerID string, command *string, memory *int64, cpu *int64, engine *string, nodeLifecycle *string, ephemeralStorage *int64) (state.Run, error) {
	var (
		run state.Run
		err error
	)

	// Validate that definition can be run (image exists, cluster has resources)
	if err = es.canBeRun(clusterName, definition, env, *engine); err != nil {
		return run, err
	}

	// Construct run object with StatusQueued and new UUID4 run id
	run, err = es.constructRun(clusterName, definition, env, ownerID, command, memory, cpu, engine, nodeLifecycle, ephemeralStorage)
	if err != nil {
		return run, err
	}

	// Save run to source of state - it is *CRITICAL* to do this
	// -before- queuing to avoid processing unsaved runs
	if err = es.stateManager.CreateRun(run); err != nil {
		return run, err
	}

	// ECS Queue run
	if *engine == state.ECSEngine {
		err = es.ecsExecutionEngine.Enqueue(run)
	}
	if *engine == state.EKSEngine {
		err = es.eksExecutionEngine.Enqueue(run)
	}

	queuedAt := time.Now()

	if err != nil {
		return run, err
	}

	// UpdateStatus the run's QueuedAt field
	if run, err = es.stateManager.UpdateRun(run.RunID, state.Run{QueuedAt: &queuedAt}); err != nil {
		return run, err
	}

	return run, nil
}

func (es *executionService) constructRun(clusterName string, definition state.Definition, env *state.EnvList, ownerID string, command *string, memory *int64, cpu *int64, engine *string, nodeLifecycle *string, ephemeralStorage *int64) (state.Run, error) {

	var (
		run state.Run
		err error
	)

	if engine == nil {
		engine = &state.DefaultEngine
	}

	runID, err := state.NewRunID(engine)
	if err != nil {
		return run, err
	}

	run = state.Run{
		RunID:            runID,
		ClusterName:      clusterName,
		GroupName:        definition.GroupName,
		DefinitionID:     definition.DefinitionID,
		Alias:            definition.Alias,
		Image:            definition.Image,
		Status:           state.StatusQueued,
		User:             ownerID,
		Command:          command,
		Memory:           memory,
		Cpu:              cpu,
		Gpu:              definition.Gpu,
		Engine:           engine,
		NodeLifecycle:    nodeLifecycle,
		EphemeralStorage: ephemeralStorage,
	}

	runEnv := es.constructEnviron(run, env)
	run.Env = &runEnv
	return run, nil
}

func (es *executionService) constructEnviron(run state.Run, env *state.EnvList) state.EnvList {
	size := len(es.reservedEnv)
	if env != nil {
		size += len(*env)
	}
	runEnv := make([]state.EnvVar, size)
	i := 0
	for k, f := range es.reservedEnv {
		runEnv[i] = state.EnvVar{
			Name:  k,
			Value: f(run),
		}
		i++
	}
	if env != nil {
		for j, e := range *env {
			runEnv[i+j] = e
		}
	}
	return state.EnvList(runEnv)
}

func (es *executionService) canBeRun(clusterName string, definition state.Definition, env *state.EnvList, engine string) error {
	if env != nil {
		for _, e := range *env {
			_, usingRestricted := es.reservedEnv[e.Name]
			if usingRestricted {
				return exceptions.ConflictingResource{
					ErrorString: fmt.Sprintf("environment variable %s is reserved", e.Name)}
			}
		}
	}

	ok, err := es.registryClient.IsImageValid(definition.Image)
	if err != nil {
		return err
	}
	if !ok {
		return exceptions.MissingResource{
			ErrorString: fmt.Sprintf(
				"image [%s] was not found in any of the configured repositories", definition.Image)}
	}

	if engine == state.ECSEngine {
		ok, err = es.ecsClusterClient.CanBeRun(clusterName, definition)
	}
	if engine == state.EKSEngine {
		if *definition.Privileged == true {
			ok, err = false, errors.New("eks cannot run containers with privileged mode")
		} else {
			ok, err = true, nil
		}
	}

	if err != nil {
		return err
	}
	if !ok {
		return exceptions.MalformedInput{
			ErrorString: fmt.Sprintf(
				"definition [%s] cannot be run on cluster [%s]", definition.DefinitionID, clusterName)}
	}
	return nil
}

//
// List returns a list of Runs
// * validates definition_id and status filters
//
func (es *executionService) List(
	limit int,
	offset int,
	sortOrder string,
	sortField string,
	filters map[string][]string,
	envFilters map[string]string) (state.RunList, error) {

	// If definition_id is present in filters, validate its
	// existence first
	definitionID, ok := filters["definition_id"]
	if ok {
		_, err := es.stateManager.GetDefinition(definitionID[0])
		if err != nil {
			return state.RunList{}, err
		}
	}

	if statusFilters, ok := filters["status"]; ok {
		for _, status := range statusFilters {
			if !state.IsValidStatus(status) {
				// Status filter is invalid
				err := exceptions.MalformedInput{
					ErrorString: fmt.Sprintf("invalid status [%s]", status)}
				return state.RunList{}, err
			}
		}
	}
	return es.stateManager.ListRuns(limit, offset, sortField, sortOrder, filters, envFilters, []string{state.ECSEngine, state.EKSEngine})
}

//
// Get returns the run with the given runID
//
func (es *executionService) Get(runID string) (state.Run, error) {
	return es.stateManager.GetRun(runID)
}

//
// UpdateStatus is for supporting some legacy runs that still manually update their status
//
func (es *executionService) UpdateStatus(runID string, status string, exitCode *int64) error {
	if !state.IsValidStatus(status) {
		return exceptions.MalformedInput{ErrorString: fmt.Sprintf("status %s is invalid", status)}
	}
	_, err := es.stateManager.UpdateRun(runID, state.Run{Status: status, ExitCode: exitCode})
	return err
}

//
// Terminate stops the run with the given runID
//
func (es *executionService) Terminate(runID string) error {
	run, err := es.stateManager.GetRun(runID)
	if err != nil {
		return err
	}

	if run.Engine == nil {
		run.Engine = &state.ECSEngine
	}

	// If it's queued and not submitted, set status to stopped (checked by submit worker)
	if run.Status == state.StatusQueued {
		_, err = es.stateManager.UpdateRun(runID, state.Run{Status: state.StatusStopped})
		return err
	}

	if *run.Engine == state.ECSEngine {
		// If it's been submitted, let the status update workers handle setting it to stopped
		if run.Status != state.StatusStopped && len(run.TaskArn) > 0 && len(run.ClusterName) > 0 {
			return es.ecsExecutionEngine.Terminate(run)
		}
	}

	if *run.Engine == state.EKSEngine && run.Status != state.StatusStopped {
		err = es.eksExecutionEngine.Terminate(run)
		if err == nil {
			exitReason := "Task terminated by the user."
			exitCode := int64(1)
			finishedAt := time.Now()
			_, err = es.stateManager.UpdateRun(run.RunID, state.Run{
				Status:     state.StatusStopped,
				ExitReason: &exitReason,
				ExitCode:   &exitCode,
				FinishedAt: &finishedAt,
			})
			return err
		}
		return nil
	}

	return exceptions.MalformedInput{
		ErrorString: fmt.Sprintf(
			"invalid run, state: %s, arn: %s, clusterName: %s", run.Status, run.TaskArn, run.ClusterName)}
}

//
// ListClusters returns a list of all execution clusters available
//
func (es *executionService) ListClusters() ([]string, error) {
	return es.ecsClusterClient.ListClusters()
}
<commit_msg>cluster ondemand override (#197)<commit_after>package services

import (
	"errors"
	"fmt"
	"github.com/stitchfix/flotilla-os/log"
	"math/rand"
	"time"

	"github.com/stitchfix/flotilla-os/clients/cluster"
	"github.com/stitchfix/flotilla-os/clients/registry"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/exceptions"
	"github.com/stitchfix/flotilla-os/execution/engine"
	"github.com/stitchfix/flotilla-os/state"
)

//
// ExecutionService interacts with the state manager and queue manager to queue runs, and perform
// CRUD operations on them
// * Acts as an intermediary layer between state and the execution engine
//
type ExecutionService interface {
	Create(definitionID string, clusterName string, env *state.EnvList, ownerID string, command *string, memory *int64, cpu *int64, engine *string, ephemeralStorage *int64, nodeLifecycle *string) (state.Run, error)
	CreateByAlias(alias string, clusterName string, env *state.EnvList, ownerID string, command *string, memory *int64, cpu *int64, engine *string, ephemeralStorage *int64, nodeLifecycle *string) (state.Run, error)
	List(
		limit int,
		offset int,
		sortOrder string,
		sortField string,
		filters map[string][]string,
		envFilters map[string]string) (state.RunList, error)
	Get(runID string) (state.Run, error)
	UpdateStatus(runID string, status string, exitCode *int64) error
	Terminate(runID string) error
	ReservedVariables() []string
	ListClusters() ([]string, error)
	GetEvents(run state.Run) (state.RunEventList, error)
}

type executionService struct {
	stateManager             state.Manager
	ecsClusterClient         cluster.Client
	eksClusterClient         cluster.Client
	registryClient           registry.Client
	ecsExecutionEngine       engine.Engine
	eksExecutionEngine       engine.Engine
	reservedEnv              map[string]func(run state.Run) string
	eksClusterOverride       string
	eksOverridePercent       int
	clusterOndemandWhitelist []string
}

func (es *executionService) GetEvents(run state.Run) (state.RunEventList, error) {
	return es.eksExecutionEngine.GetEvents(run)
}

//
// NewExecutionService configures and returns an ExecutionService
//
func NewExecutionService(conf config.Config,
	ecsExecutionEngine engine.Engine,
	eksExecutionEngine engine.Engine,
	sm state.Manager,
	ecsClusterClient cluster.Client,
	eksClusterClient cluster.Client,
	rc registry.Client,
	log log.Logger) (ExecutionService, error) {
	es := executionService{
		stateManager:       sm,
		ecsClusterClient:   ecsClusterClient,
		eksClusterClient:   eksClusterClient,
		registryClient:     rc,
		ecsExecutionEngine: ecsExecutionEngine,
		eksExecutionEngine: eksExecutionEngine,
	}
	//
	// Reserved environment variables dynamically generated
	// per run

	ownerKey := conf.GetString("owner_id_var")
	if len(ownerKey) == 0 {
		ownerKey = "FLOTILLA_RUN_OWNER_ID"
	}

	es.eksClusterOverride = conf.GetString("eks.cluster_override")
	es.eksOverridePercent = conf.GetInt("eks.cluster_override_percent")
	es.clusterOndemandWhitelist = conf.GetStringSlice("eks.cluster_ondemand_whitelist")

	es.reservedEnv = map[string]func(run state.Run) string{
		"FLOTILLA_SERVER_MODE": func(run state.Run) string {
			return conf.GetString("flotilla_mode")
		},
		"FLOTILLA_RUN_ID": func(run state.Run) string {
			return run.RunID
		},
		ownerKey: func(run state.Run) string {
			return run.User
		},
	}
	// Warm cached cluster list
	_, _ = es.ecsClusterClient.ListClusters()
	return &es, nil
}

//
// ReservedVariables returns the list of reserved run environment variable
// names
//
func (es *executionService) ReservedVariables() []string {
	var keys []string
	for k := range es.reservedEnv {
		keys = append(keys, k)
	}
	return keys
}
func contains(s []string, e string) bool {
	for _, a := range s {
		if a == e {
			return true
		}
	}
	return false
}

//
// Create constructs and queues a new Run on the cluster specified
//
func (es *executionService) Create(definitionID string, clusterName string, env *state.EnvList, ownerID string, command *string, memory *int64, cpu *int64, engine *string, ephemeralStorage *int64, nodeLifecycle *string) (state.Run, error) {

	// Ensure definition exists
	definition, err := es.stateManager.GetDefinition(definitionID)
	if err != nil {
		return state.Run{}, err
	}

	if engine == nil {
		engine = &state.DefaultEngine
	}

	// Added to facilitate migration of ECS jobs to EKS.
	if engine != &state.EKSEngine && es.eksOverridePercent > 0 && *definition.Privileged == false {
		modulo := 100 / es.eksOverridePercent
		if rand.Int()%modulo == 0 {
			engine = &state.EKSEngine
			if contains(es.clusterOndemandWhitelist, clusterName) {
				nodeLifecycle = &state.OndemandLifecycle
			} else {
				nodeLifecycle = &state.SpotLifecycle
			}
			clusterName = es.eksClusterOverride
		}
	}

	return es.createFromDefinition(definition, clusterName, env, ownerID, command, memory, cpu, engine, nodeLifecycle, ephemeralStorage)
}

//
// Create constructs and queues a new Run on the cluster specified, based on an alias
//
func (es *executionService) CreateByAlias(alias string, clusterName string, env *state.EnvList, ownerID string, command *string, memory *int64, cpu *int64, engine *string, ephemeralStorage *int64, nodeLifecycle *string) (state.Run, error) {

	// Ensure definition exists
	definition, err := es.stateManager.GetDefinitionByAlias(alias)
	if err != nil {
		return state.Run{}, err
	}

	if engine == nil {
		engine = &state.DefaultEngine
	}

	// Added to facilitate migration of ECS jobs to EKS.
	if engine != &state.EKSEngine && es.eksOverridePercent > 0 && *definition.Privileged == false {
		modulo := 100 / es.eksOverridePercent
		if rand.Int()%modulo == 0 {
			engine = &state.EKSEngine
			if contains(es.clusterOndemandWhitelist, clusterName) {
				nodeLifecycle = &state.OndemandLifecycle
			} else {
				nodeLifecycle = &state.SpotLifecycle
			}
			clusterName = es.eksClusterOverride
		}
	}
	return es.createFromDefinition(definition, clusterName, env, ownerID, command, memory, cpu, engine, nodeLifecycle, ephemeralStorage)
}

func (es *executionService) createFromDefinition(definition state.Definition, clusterName string, env *state.EnvList, ownerID string, command *string, memory *int64, cpu *int64, engine *string, nodeLifecycle *string, ephemeralStorage *int64) (state.Run, error) {
	var (
		run state.Run
		err error
	)

	// Validate that definition can be run (image exists, cluster has resources)
	if err = es.canBeRun(clusterName, definition, env, *engine); err != nil {
		return run, err
	}

	// Construct run object with StatusQueued and new UUID4 run id
	run, err = es.constructRun(clusterName, definition, env, ownerID, command, memory, cpu, engine, nodeLifecycle, ephemeralStorage)
	if err != nil {
		return run, err
	}

	// Save run to source of state - it is *CRITICAL* to do this
	// -before- queuing to avoid processing unsaved runs
	if err = es.stateManager.CreateRun(run); err != nil {
		return run, err
	}

	// ECS Queue run
	if *engine == state.ECSEngine {
		err = es.ecsExecutionEngine.Enqueue(run)
	}
	if *engine == state.EKSEngine {
		err = es.eksExecutionEngine.Enqueue(run)
	}

	queuedAt := time.Now()

	if err != nil {
		return run, err
	}

	// UpdateStatus the run's QueuedAt field
	if run, err = es.stateManager.UpdateRun(run.RunID, state.Run{QueuedAt: &queuedAt}); err != nil {
		return run, err
	}

	return run, nil
}

func (es *executionService) constructRun(clusterName string, definition state.Definition, env *state.EnvList, ownerID string, command *string, memory *int64, cpu *int64, engine *string, nodeLifecycle *string, ephemeralStorage *int64) (state.Run, error) {

	var (
		run state.Run
		err error
	)

	if engine == nil {
		engine = &state.DefaultEngine
	}

	runID, err := state.NewRunID(engine)
	if err != nil {
		return run, err
	}

	run = state.Run{
		RunID:            runID,
		ClusterName:      clusterName,
		GroupName:        definition.GroupName,
		DefinitionID:     definition.DefinitionID,
		Alias:            definition.Alias,
		Image:            definition.Image,
		Status:           state.StatusQueued,
		User:             ownerID,
		Command:          command,
		Memory:           memory,
		Cpu:              cpu,
		Gpu:              definition.Gpu,
		Engine:           engine,
		NodeLifecycle:    nodeLifecycle,
		EphemeralStorage: ephemeralStorage,
	}

	runEnv := es.constructEnviron(run, env)
	run.Env = &runEnv
	return run, nil
}

func (es *executionService) constructEnviron(run state.Run, env *state.EnvList) state.EnvList {
	size := len(es.reservedEnv)
	if env != nil {
		size += len(*env)
	}
	runEnv := make([]state.EnvVar, size)
	i := 0
	for k, f := range es.reservedEnv {
		runEnv[i] = state.EnvVar{
			Name:  k,
			Value: f(run),
		}
		i++
	}
	if env != nil {
		for j, e := range *env {
			runEnv[i+j] = e
		}
	}
	return state.EnvList(runEnv)
}

func (es *executionService) canBeRun(clusterName string, definition state.Definition, env *state.EnvList, engine string) error {
	if env != nil {
		for _, e := range *env {
			_, usingRestricted := es.reservedEnv[e.Name]
			if usingRestricted {
				return exceptions.ConflictingResource{
					ErrorString: fmt.Sprintf("environment variable %s is reserved", e.Name)}
			}
		}
	}

	ok, err := es.registryClient.IsImageValid(definition.Image)
	if err != nil {
		return err
	}
	if !ok {
		return exceptions.MissingResource{
			ErrorString: fmt.Sprintf(
				"image [%s] was not found in any of the configured repositories", definition.Image)}
	}

	if engine == state.ECSEngine {
		ok, err = es.ecsClusterClient.CanBeRun(clusterName, definition)
	}
	if engine == state.EKSEngine {
		if *definition.Privileged == true {
			ok, err = false, errors.New("eks cannot run containers with privileged mode")
		} else {
			ok, err = true, nil
		}
	}

	if err != nil {
		return err
	}
	if !ok {
		return exceptions.MalformedInput{
			ErrorString: fmt.Sprintf(
				"definition [%s] cannot be run on cluster [%s]", definition.DefinitionID, clusterName)}
	}
	return nil
}

//
// List returns a list of Runs
// * validates definition_id and status filters
//
func (es *executionService) List(
	limit int,
	offset int,
	sortOrder string,
	sortField string,
	filters map[string][]string,
	envFilters map[string]string) (state.RunList, error) {

	// If definition_id is present in filters, validate its
	// existence first
	definitionID, ok := filters["definition_id"]
	if ok {
		_, err := es.stateManager.GetDefinition(definitionID[0])
		if err != nil {
			return state.RunList{}, err
		}
	}

	if statusFilters, ok := filters["status"]; ok {
		for _, status := range statusFilters {
			if !state.IsValidStatus(status) {
				// Status filter is invalid
				err := exceptions.MalformedInput{
					ErrorString: fmt.Sprintf("invalid status [%s]", status)}
				return state.RunList{}, err
			}
		}
	}
	return es.stateManager.ListRuns(limit, offset, sortField, sortOrder, filters, envFilters, []string{state.ECSEngine, state.EKSEngine})
}

//
// Get returns the run with the given runID
//
func (es *executionService) Get(runID string) (state.Run, error) {
	return es.stateManager.GetRun(runID)
}

//
// UpdateStatus is for supporting some legacy runs that still manually update their status
//
func (es *executionService) UpdateStatus(runID string, status string, exitCode *int64) error {
	if !state.IsValidStatus(status) {
		return exceptions.MalformedInput{ErrorString: fmt.Sprintf("status %s is invalid", status)}
	}
	_, err := es.stateManager.UpdateRun(runID, state.Run{Status: status, ExitCode: exitCode})
	return err
}

//
// Terminate stops the run with the given runID
//
func (es *executionService) Terminate(runID string) error {
	run, err := es.stateManager.GetRun(runID)
	if err != nil {
		return err
	}

	if run.Engine == nil {
		run.Engine = &state.ECSEngine
	}

	// If it's queued and not submitted, set status to stopped (checked by submit worker)
	if run.Status == state.StatusQueued {
		_, err = es.stateManager.UpdateRun(runID, state.Run{Status: state.StatusStopped})
		return err
	}

	if *run.Engine == state.ECSEngine {
		// If it's been submitted, let the status update workers handle setting it to stopped
		if run.Status != state.StatusStopped && len(run.TaskArn) > 0 && len(run.ClusterName) > 0 {
			return es.ecsExecutionEngine.Terminate(run)
		}
	}

	if *run.Engine == state.EKSEngine && run.Status != state.StatusStopped {
		err = es.eksExecutionEngine.Terminate(run)
		if err == nil {
			exitReason := "Task terminated by the user."
			exitCode := int64(1)
			finishedAt := time.Now()
			_, err = es.stateManager.UpdateRun(run.RunID, state.Run{
				Status:     state.StatusStopped,
				ExitReason: &exitReason,
				ExitCode:   &exitCode,
				FinishedAt: &finishedAt,
			})
			return err
		}
		return nil
	}

	return exceptions.MalformedInput{
		ErrorString: fmt.Sprintf(
			"invalid run, state: %s, arn: %s, clusterName: %s", run.Status, run.TaskArn, run.ClusterName)}
}

//
// ListClusters returns a list of all execution clusters available
//
func (es *executionService) ListClusters() ([]string, error) {
	return es.ecsClusterClient.ListClusters()
}
<|endoftext|>


<commit_before>package nats

import (
	"log"
	"net"
)

func Start(port int, host string) error {
	log.Println("running nats mock server", port, host)

	ln, err := net.Listen("tcp", host+":"+string(port))

	if err != nil {
		return err
	}

	for {
		conn, err := ln.Accept()

		if err != nil {
			return err
		}

		go func(c net.Conn) {
		}(conn)
	}

	return nil
}
<commit_msg>nats: convert port to string<commit_after>package nats

import (
	"log"
	"net"
	"strconv"
)

func Start(port int, host string) error {
	log.Println("running nats mock server", port, host)

	ln, err := net.Listen("tcp", host+":"+strconv.Itoa(port))

	if err != nil {
		return err
	}

	for {
		conn, err := ln.Accept()

		if err != nil {
			return err
		}

		go func(c net.Conn) {
		}(conn)
	}

	return nil
}
<|endoftext|>


<commit_before>package shape

import (
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestSquareOutine(t *testing.T) {
	out, err := Rectangle.Outline(2, 2)
	assert.Nil(t, err)
	fmt.Println(out, err)
	assert.Equal(t, len(out), 4)

	out, err = Rectangle.Outline(3, 3)
	assert.Nil(t, err)
	fmt.Println(out, err)
	assert.Equal(t, len(out), 8)

	out, err = Rectangle.Outline(4, 4)
	assert.Nil(t, err)
	fmt.Println(out, err)
	assert.Equal(t, len(out), 12)

	in := JustIn(func(x, y int, sizes ...int) bool {
		return x > 5
	})
	out, err = in.Outline(10, 3)
	assert.Nil(t, err)
	assert.Equal(t, len(out), 10)

	out, err = in.Outline(3, 10)
	assert.NotNil(t, err)
	assert.Equal(t, len(out), 0)

	in = JustIn(func(x, y int, sizes ...int) bool {
		return x == 0 && y == 0
	})
	out, err = in.Outline(2, 2)
	assert.Nil(t, err)
	assert.Equal(t, len(out), 1)
}
<commit_msg>More outline tests<commit_after>package shape

import (
	"testing"

	"github.com/stretchr/testify/require"
)

func TestRectangleOutline(t *testing.T) {
	type testCase struct {
		w, h           int
		expectedLength int
		expectedErr    error
	}
	tcs := []testCase{
		{
			w:              2,
			h:              2,
			expectedLength: 4,
		},
		{
			w:              3,
			h:              3,
			expectedLength: 8,
		},
		{
			w:              4,
			h:              4,
			expectedLength: 12,
		},
	}
	for _, tc := range tcs {
		out, err := Rectangle.Outline(tc.w, tc.h)
		if tc.expectedErr == nil {
			require.Nil(t, err)
		} else {
			require.Contains(t, tc.expectedErr.Error(), err.Error())
		}
		require.Equal(t, tc.expectedLength, len(out))
	}
}

func TestToOutline4(t *testing.T) {
	type testCase struct {
		w, h        int
		shape       Shape
		outlineLen  int
		outline4Len int
	}
	tcs := []testCase{
		{
			w:           10,
			h:           10,
			shape:       Heart,
			outlineLen:  24,
			outline4Len: 32,
		},
	}
	for _, tc := range tcs {
		out, _ := tc.shape.Outline(tc.w, tc.h)
		require.Equal(t, tc.outlineLen, len(out))

		out, _ = ToOutline4(tc.shape)(tc.w, tc.h)
		require.Equal(t, tc.outline4Len, len(out))
	}
}

func TestJustInOutline(t *testing.T) {
	type testCase struct {
		w, h        int
		in          func(int, int, ...int) bool
		expectedLen int
		shouldErr   bool
	}
	tcs := []testCase{
		{
			w: 10,
			h: 3,
			in: func(x, y int, sizes ...int) bool {
				return x > 5
			},
			expectedLen: 10,
		},
		{
			w: 3,
			h: 10,
			in: func(x, y int, sizes ...int) bool {
				return x > 5
			},
			shouldErr: true,
		},
		{
			w: 2,
			h: 2,
			in: func(x, y int, sizes ...int) bool {
				return x == 0 && y == 0
			},
			expectedLen: 1,
		},
	}
	for _, tc := range tcs {
		out, err := JustIn(tc.in).Outline(tc.w, tc.h)
		if tc.shouldErr {
			require.NotNil(t, err)
		} else {
			require.Nil(t, err)
		}
		require.Equal(t, tc.expectedLen, len(out))
	}
}
<|endoftext|>


<commit_before>/*
   Package cryptographer implements a simple library for on-the-fly
   NaCl secret key encryption and decryption.

   It is meant to be used for symmetric-key encryption schemes.
   Optionally it can (de)compress the data before (dec)encryption.

   Beyond the recommended methods (Encrypt, Decrypt) it also implements
   the io.ReadWriter interface which is slower. You may run the benchmarks
   from cryptographer_test.go to decide if it is acceptable.

   One bit of the NaCl's nonce is used to indicate whether the message was
   compressed before encrypting. Still the algorithm should remain safe
   since nonce collisions are again extremely rare.
*/
package cryptographer

import (
	"bytes"
	"compress/zlib"
	"crypto/rand"
	"errors"
	"golang.org/x/crypto/nacl/secretbox"
	"io"
	"io/ioutil"
)

// These are defined in golang.org/x/crypto/nacl/secretbox
const (
	keySize   = 32
	nonceSize = 24
)

const compressBit byte = 0x01

// A Cryptographer holds the instance's key and the compression settings.
type Cryptographer struct {
	key      *[keySize]byte
	compress bool
}

// New creates a new Cryptographer instance. key is the key used for encryption,
// pad is the padding to be used (at least 32 bytes), if the key is smaller than 32 bytes.
// compress indicates whether the data should be compessed (zlib) before encrypting.
// The pad can be a const in your code.
func New(key, pad string, compress bool) (*Cryptographer, error) {
	naclKey, err := constructKey(key, pad)
	if err != nil {
		return nil, err
	}
	return &Cryptographer{naclKey, compress}, nil
}

func constructKey(key, pad string) (naclKey *[32]byte, e error) {
	tKey := []byte(key)
	tPad := []byte(pad)
	if len(tPad) < 32 {
		return nil, errors.New("cryptographer pad should be 32 bytes or more")
	}
	tKey = append(tKey, tPad...)
	naclKey = new([keySize]byte)
	copy(naclKey[:], tKey[:keySize])
	return naclKey, nil
}

// Encrypt encrypts a message and returns the encrypted msg (nonce + ciphertext).
// If you have enabled compression, it will compress the msg before encrypting it.
func (c Cryptographer) Encrypt(msg []byte) (out []byte, e error) {
	nonce := new([nonceSize]byte)
	_, err := io.ReadFull(rand.Reader, nonce[:])
	if err != nil {
		return nil, err
	}

	// We use the last bit of the nonce as a compression indicator.
	// This should still keep you safe (extremely rare collisions).
	nonce[23] &= ^compressBit

	if c.compress {
		var b bytes.Buffer
		w := zlib.NewWriter(&b)
		w.Write(msg)
		w.Close()
		msg = b.Bytes()
		nonce[23] |= compressBit
	}

	out = make([]byte, nonceSize)
	copy(out, nonce[:])

	out = secretbox.Seal(out, msg, nonce, c.key)
	return out, nil
}

// Decrypt decrypts an encrypted message and returns it (plaintext).
// If you have enabled compression, it wil detect it and decompress
// the msg after decrypting it.
func (c Cryptographer) Decrypt(msg []byte) ([]byte, error) {
	if len(msg) < nonceSize+secretbox.Overhead {
		return nil, errors.New("encrypted message length too short")
	}

	nonce := new([nonceSize]byte)
	copy(nonce[:], msg[:nonceSize])

	out, ok := secretbox.Open(nil, msg[nonceSize:], nonce, c.key)
	if !ok {
		return nil, errors.New("could not decrypt message")
	}

	if nonce[23]&compressBit == compressBit {
		r, err := zlib.NewReader(bytes.NewReader(out))
		if err != nil {
			return nil, err
		}
		r.Close()
		out, err = ioutil.ReadAll(r)
		if err != nil {
			return nil, err
		}
	}
	return out, nil
}

// A Reader reads data from another Reader, decrypts and, if needed,
// decompress them into a []byte variable.
// A Reader may be re-used by using Reset.
type Reader struct {
	r         io.Reader
	msg       []byte
	done      bool
	firstRead bool
	c         *Cryptographer
}

func newInternal(naclKey *[32]byte, compress bool) *Cryptographer {
	return &Cryptographer{naclKey, compress}
}

// NewReader creates a new Reader. Reads from the returned Reader read,
// unencrypt and decompress, if needed, data from r. The implementation
// needs to read all data from r at once, since we do not use a stream
// cipher. Pad should be at least 32 bytes long.
func NewReader(r io.Reader, key, pad string) (*Reader, error) {
	naclKey, err := constructKey(key, pad)
	if err != nil {
		return nil, err
	}
	return &Reader{r, nil, false, true, newInternal(naclKey, false)}, nil
}

// Read reads into p an unecrypted and, if needed, uncompressed form
// of the bytes from the underlying Reader. Read needs to read all
// data from the underlying Reader before it can decrypt them.
func (r *Reader) Read(p []byte) (n int, err error) {
	if r.done {
		return 0, io.EOF
	}

	msg, err := ioutil.ReadAll(r.r)
	if err != nil {
		return 0, err
	}

	if r.firstRead {
		r.msg, err = r.c.Decrypt(msg)
		if err != nil {
			r.done = true
			return 0, errors.New("could not decrypt input: " + err.Error())
		}
		r.firstRead = false
	}

	length := len(r.msg)
	if len(p) < length {
		length = len(p)
	}
	for i := 0; i < length; i++ {
		p[i] = r.msg[i]
	}
	if length < len(r.msg) {
		r.msg = r.msg[len(p):]
	} else {
		r.msg = r.msg[:0]
		r.done = true
	}
	return length, nil
}

// Reset returns Reader to its initial state, except it now reads from r.
func (d *Reader) Reset(r io.Reader) {
	d.r = r
	d.msg = nil
	d.done = false
	d.firstRead = true
}

// A Writer takes data written to it and writes the encrypted and, if needed,
// compressed form of that data to an underlying writer.
type Writer struct {
	w           io.Writer
	unencrypted []byte
	c           *Cryptographer
}

// NewWriter creates a new writer. Writes to the returned Writer are encrypted and,
// if needed, compressed and written to w.
//
// It is the caller's responsibility to call Close() on WriteCloser when done, since
// we do not use a stream cipher, we need to have all the data before encrypting them.
func NewWriter(w io.Writer, key, pad string, compress bool) (*Writer, error) {
	naclKey, err := constructKey(key, pad)
	if err != nil {
		return nil, err
	}
	return &Writer{w, nil, newInternal(naclKey, compress)}, nil
}

// Write writes and encrypted and, if needed, compressed form of p to the underlying
// io.Writer. The compressed bytes are not written until the Writer is closed or
// explicitly flushed.
func (w *Writer) Write(p []byte) (n int, err error) {
	w.unencrypted = append(w.unencrypted, p...)
	return len(p), nil
}

// Flush encrypts and compresses, if needed, the data written to the writer.
// After a Flush, the writer has to be Reset in order to write to it again.
func (w *Writer) Flush() error {
	var err error
	write, err := w.c.Encrypt(w.unencrypted)
	if err != nil {
		return err
	}
	_, err = w.w.Write(write)
	if err != nil {
		return err
	}
	return nil
}

// Close acts as a placeholder for Flush.
func (w *Writer) Close() error {
	return w.Flush()
}

// Reset clears the sate of the Writer w such that it is equivalent to its
// initial state from NewWriter, but instead writing to w.
func (e *Writer) Reset(w io.Writer) {
	e.w = w
	e.unencrypted = nil
}
<commit_msg>golint<commit_after>/*
Package cryptographer implements a simple library for on-the-fly
NaCl secret key encryption and decryption.

It is meant to be used for symmetric-key encryption schemes.
Optionally it can (de)compress the data before (dec)encryption.

Beyond the recommended methods (Encrypt, Decrypt) it also implements
the io.ReadWriter interface which is slower. You may run the benchmarks
from cryptographer_test.go to decide if it is acceptable.

One bit of the NaCl's nonce is used to indicate whether the message was
compressed before encrypting. Still the algorithm should remain safe
since nonce collisions are again extremely rare.
*/
package cryptographer

import (
	"bytes"
	"compress/zlib"
	"crypto/rand"
	"errors"
	"golang.org/x/crypto/nacl/secretbox"
	"io"
	"io/ioutil"
)

// These are defined in golang.org/x/crypto/nacl/secretbox
const (
	keySize   = 32
	nonceSize = 24
)

const compressBit byte = 0x01

// A Cryptographer holds the instance's key and the compression settings.
type Cryptographer struct {
	key      *[keySize]byte
	compress bool
}

// New creates a new Cryptographer instance. key is the key used for encryption,
// pad is the padding to be used (at least 32 bytes), if the key is smaller than 32 bytes.
// compress indicates whether the data should be compessed (zlib) before encrypting.
// The pad can be a const in your code.
func New(key, pad string, compress bool) (*Cryptographer, error) {
	naclKey, err := constructKey(key, pad)
	if err != nil {
		return nil, err
	}
	return &Cryptographer{naclKey, compress}, nil
}

func constructKey(key, pad string) (naclKey *[32]byte, e error) {
	tKey := []byte(key)
	tPad := []byte(pad)
	if len(tPad) < 32 {
		return nil, errors.New("cryptographer pad should be 32 bytes or more")
	}
	tKey = append(tKey, tPad...)
	naclKey = new([keySize]byte)
	copy(naclKey[:], tKey[:keySize])
	return naclKey, nil
}

// Encrypt encrypts a message and returns the encrypted msg (nonce + ciphertext).
// If you have enabled compression, it will compress the msg before encrypting it.
func (c Cryptographer) Encrypt(msg []byte) (out []byte, e error) {
	nonce := new([nonceSize]byte)
	_, err := io.ReadFull(rand.Reader, nonce[:])
	if err != nil {
		return nil, err
	}

	// We use the last bit of the nonce as a compression indicator.
	// This should still keep you safe (extremely rare collisions).
	nonce[23] &= ^compressBit

	if c.compress {
		var b bytes.Buffer
		w := zlib.NewWriter(&b)
		w.Write(msg)
		w.Close()
		msg = b.Bytes()
		nonce[23] |= compressBit
	}

	out = make([]byte, nonceSize)
	copy(out, nonce[:])

	out = secretbox.Seal(out, msg, nonce, c.key)
	return out, nil
}

// Decrypt decrypts an encrypted message and returns it (plaintext).
// If you have enabled compression, it wil detect it and decompress
// the msg after decrypting it.
func (c Cryptographer) Decrypt(msg []byte) ([]byte, error) {
	if len(msg) < nonceSize+secretbox.Overhead {
		return nil, errors.New("encrypted message length too short")
	}

	nonce := new([nonceSize]byte)
	copy(nonce[:], msg[:nonceSize])

	out, ok := secretbox.Open(nil, msg[nonceSize:], nonce, c.key)
	if !ok {
		return nil, errors.New("could not decrypt message")
	}

	if nonce[23]&compressBit == compressBit {
		r, err := zlib.NewReader(bytes.NewReader(out))
		if err != nil {
			return nil, err
		}
		r.Close()
		out, err = ioutil.ReadAll(r)
		if err != nil {
			return nil, err
		}
	}
	return out, nil
}

// A Reader reads data from another Reader, decrypts and, if needed,
// decompress them into a []byte variable.
// A Reader may be re-used by using Reset.
type Reader struct {
	r         io.Reader
	msg       []byte
	done      bool
	firstRead bool
	c         *Cryptographer
}

func newInternal(naclKey *[32]byte, compress bool) *Cryptographer {
	return &Cryptographer{naclKey, compress}
}

// NewReader creates a new Reader. Reads from the returned Reader read,
// unencrypt and decompress, if needed, data from r. The implementation
// needs to read all data from r at once, since we do not use a stream
// cipher. Pad should be at least 32 bytes long.
func NewReader(r io.Reader, key, pad string) (*Reader, error) {
	naclKey, err := constructKey(key, pad)
	if err != nil {
		return nil, err
	}
	return &Reader{r, nil, false, true, newInternal(naclKey, false)}, nil
}

// Read reads into p an unecrypted and, if needed, uncompressed form
// of the bytes from the underlying Reader. Read needs to read all
// data from the underlying Reader before it can decrypt them.
func (d *Reader) Read(p []byte) (n int, err error) {
	if d.done {
		return 0, io.EOF
	}

	msg, err := ioutil.ReadAll(d.r)
	if err != nil {
		return 0, err
	}

	if d.firstRead {
		d.msg, err = d.c.Decrypt(msg)
		if err != nil {
			d.done = true
			return 0, errors.New("could not decrypt input: " + err.Error())
		}
		d.firstRead = false
	}

	length := len(d.msg)
	if len(p) < length {
		length = len(p)
	}
	for i := 0; i < length; i++ {
		p[i] = d.msg[i]
	}
	if length < len(d.msg) {
		d.msg = d.msg[len(p):]
	} else {
		d.msg = d.msg[:0]
		d.done = true
	}
	return length, nil
}

// Reset returns Reader to its initial state, except it now reads from r.
func (d *Reader) Reset(r io.Reader) {
	d.r = r
	d.msg = nil
	d.done = false
	d.firstRead = true
}

// A Writer takes data written to it and writes the encrypted and, if needed,
// compressed form of that data to an underlying writer.
type Writer struct {
	w           io.Writer
	unencrypted []byte
	c           *Cryptographer
}

// NewWriter creates a new writer. Writes to the returned Writer are encrypted and,
// if needed, compressed and written to w.
//
// It is the caller's responsibility to call Close() on WriteCloser when done, since
// we do not use a stream cipher, we need to have all the data before encrypting them.
func NewWriter(w io.Writer, key, pad string, compress bool) (*Writer, error) {
	naclKey, err := constructKey(key, pad)
	if err != nil {
		return nil, err
	}
	return &Writer{w, nil, newInternal(naclKey, compress)}, nil
}

// Write writes and encrypted and, if needed, compressed form of p to the underlying
// io.Writer. The compressed bytes are not written until the Writer is closed or
// explicitly flushed.
func (e *Writer) Write(p []byte) (n int, err error) {
	e.unencrypted = append(e.unencrypted, p...)
	return len(p), nil
}

// Flush encrypts and compresses, if needed, the data written to the writer.
// After a Flush, the writer has to be Reset in order to write to it again.
func (e *Writer) Flush() error {
	var err error
	write, err := e.c.Encrypt(e.unencrypted)
	if err != nil {
		return err
	}
	_, err = e.w.Write(write)
	if err != nil {
		return err
	}
	return nil
}

// Close acts as a placeholder for Flush.
func (e *Writer) Close() error {
	return e.Flush()
}

// Reset clears the sate of the Writer w such that it is equivalent to its
// initial state from NewWriter, but instead writing to w.
func (e *Writer) Reset(w io.Writer) {
	e.w = w
	e.unencrypted = nil
}
<|endoftext|>