<commit_before>package main

import (
	"fmt"
	"strings"
	"time"

	"github.com/bcicen/ctop/config"
	"github.com/bcicen/ctop/container"
	"github.com/bcicen/ctop/widgets"
	"github.com/bcicen/ctop/widgets/menu"
	ui "github.com/gizak/termui"
)

// MenuFn executes a menu window, returning the next menu or nil
type MenuFn func() MenuFn

var helpDialog = []menu.Item{
	{"<enter> - open container menu", ""},
	{"", ""},
	{"[a] - toggle display of all containers", ""},
	{"[f] - filter displayed containers", ""},
	{"[h] - open this help dialog", ""},
	{"[H] - toggle ctop header", ""},
	{"[s] - select container sort field", ""},
	{"[r] - reverse container sort order", ""},
	{"[o] - open single view", ""},
	{"[l] - view container logs ([t] to toggle timestamp when open)", ""},
	{"[e] - exec shell", ""},
	{"[c] - configure columns", ""},
	{"[S] - save current configuration to file", ""},
	{"[q] - exit ctop", ""},
}

func HelpMenu() MenuFn {
	ui.Clear()
	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	m := menu.NewMenu()
	m.BorderLabel = "Help"
	m.AddItems(helpDialog...)
	ui.Handle("/sys/wnd/resize", func(e ui.Event) {
		ui.Clear()
		ui.Render(m)
	})
	ui.Handle("/sys/kbd/", func(ui.Event) {
		ui.StopLoop()
	})
	ui.Loop()
	return nil
}

func FilterMenu() MenuFn {
	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	i := widgets.NewInput()
	i.BorderLabel = "Filter"
	i.SetY(ui.TermHeight() - i.Height)
	i.Data = config.GetVal("filterStr")
	ui.Render(i)

	// refresh container rows on input
	stream := i.Stream()
	go func() {
		for s := range stream {
			config.Update("filterStr", s)
			RefreshDisplay()
			ui.Render(i)
		}
	}()

	i.InputHandlers()
	ui.Handle("/sys/kbd/<escape>", func(ui.Event) {
		config.Update("filterStr", "")
		ui.StopLoop()
	})
	ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
		config.Update("filterStr", i.Data)
		ui.StopLoop()
	})
	ui.Loop()
	return nil
}

func SortMenu() MenuFn {
	ui.Clear()
	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	m := menu.NewMenu()
	m.Selectable = true
	m.SortItems = true
	m.BorderLabel = "Sort Field"

	for _, field := range container.SortFields() {
		m.AddItems(menu.Item{field, ""})
	}

	// set cursor position to current sort field
	m.SetCursor(config.GetVal("sortField"))

	HandleKeys("up", m.Up)
	HandleKeys("down", m.Down)
	HandleKeys("exit", ui.StopLoop)

	ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
		config.Update("sortField", m.SelectedValue())
		ui.StopLoop()
	})

	ui.Render(m)
	ui.Loop()
	return nil
}

func ColumnsMenu() MenuFn {
	const (
		enabledStr  = "[X]"
		disabledStr = "[ ]"
		padding     = 2
	)

	ui.Clear()
	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	m := menu.NewMenu()
	m.Selectable = true
	m.SortItems = false
	m.BorderLabel = "Columns"
	//m.SubText = "Enabled Columns"

	rebuild := func() {
		// get padding for right alignment of enabled status
		var maxLen int
		for _, col := range config.GlobalColumns {
			if len(col.Label) > maxLen {
				maxLen = len(col.Label)
			}
		}
		maxLen += padding

		// rebuild menu items
		m.ClearItems()
		for _, col := range config.GlobalColumns {
			txt := col.Label + strings.Repeat(" ", maxLen-len(col.Label))
			if col.Enabled {
				txt += enabledStr
			} else {
				txt += disabledStr
			}
			m.AddItems(menu.Item{col.Name, txt})
		}
	}

	upFn := func() {
		config.ColumnLeft(m.SelectedValue())
		m.Up()
		rebuild()
	}

	downFn := func() {
		config.ColumnRight(m.SelectedValue())
		m.Down()
		rebuild()
	}

	toggleFn := func() {
		config.ColumnToggle(m.SelectedValue())
		rebuild()
	}

	rebuild()

	HandleKeys("up", m.Up)
	HandleKeys("down", m.Down)
	HandleKeys("enter", toggleFn)
	HandleKeys("pgup", upFn)
	HandleKeys("pgdown", downFn)

	ui.Handle("/sys/kbd/x", func(ui.Event) { toggleFn() })
	ui.Handle("/sys/kbd/<enter>", func(ui.Event) { toggleFn() })

	HandleKeys("exit", func() {
		cSource, err := cursor.cSuper.Get()
		if err == nil {
			for _, c := range cSource.All() {
				c.RecreateWidgets()
			}
		}
		ui.StopLoop()
	})

	ui.Render(m)
	ui.Loop()
	return nil
}

func ContainerMenu() MenuFn {
	c := cursor.Selected()
	if c == nil {
		return nil
	}

	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	m := menu.NewMenu()
	m.Selectable = true
	m.BorderLabel = "Menu"

	items := []menu.Item{
		menu.Item{Val: "single", Label: "[o] single view"},
		menu.Item{Val: "logs", Label: "[l] log view"},
	}

	if c.Meta["state"] == "running" {
		items = append(items, menu.Item{Val: "stop", Label: "[s] stop"})
		items = append(items, menu.Item{Val: "pause", Label: "[p] pause"})
		items = append(items, menu.Item{Val: "restart", Label: "[r] restart"})
		items = append(items, menu.Item{Val: "exec", Label: "[e] exec shell"})
	}
	if c.Meta["state"] == "exited" || c.Meta["state"] == "created" {
		items = append(items, menu.Item{Val: "start", Label: "[s] start"})
		items = append(items, menu.Item{Val: "remove", Label: "[R] remove"})
	}
	if c.Meta["state"] == "paused" {
		items = append(items, menu.Item{Val: "unpause", Label: "[p] unpause"})
	}
	items = append(items, menu.Item{Val: "cancel", Label: "[c] cancel"})

	m.AddItems(items...)
	ui.Render(m)

	HandleKeys("up", m.Up)
	HandleKeys("down", m.Down)

	var selected string

	// shortcuts
	ui.Handle("/sys/kbd/o", func(ui.Event) {
		selected = "single"
		ui.StopLoop()
	})
	ui.Handle("/sys/kbd/l", func(ui.Event) {
		selected = "logs"
		ui.StopLoop()
	})
	if c.Meta["state"] != "paused" {
		ui.Handle("/sys/kbd/s", func(ui.Event) {
			if c.Meta["state"] == "running" {
				selected = "stop"
			} else {
				selected = "start"
			}
			ui.StopLoop()
		})
	}
	if c.Meta["state"] != "exited" || c.Meta["state"] != "created" {
		ui.Handle("/sys/kbd/p", func(ui.Event) {
			if c.Meta["state"] == "paused" {
				selected = "unpause"
			} else {
				selected = "pause"
			}
			ui.StopLoop()
		})
	}
	if c.Meta["state"] == "running" {
		ui.Handle("/sys/kbd/e", func(ui.Event) {
			selected = "exec"
			ui.StopLoop()
		})
		ui.Handle("/sys/kbd/r", func(ui.Event) {
			selected = "restart"
			ui.StopLoop()
		})
	}
	ui.Handle("/sys/kbd/R", func(ui.Event) {
		selected = "remove"
		ui.StopLoop()
	})
	ui.Handle("/sys/kbd/c", func(ui.Event) {
		ui.StopLoop()
	})

	ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
		selected = m.SelectedValue()
		ui.StopLoop()
	})
	ui.Handle("/sys/kbd/", func(ui.Event) {
		ui.StopLoop()
	})
	ui.Loop()

	var nextMenu MenuFn
	switch selected {
	case "single":
		nextMenu = SingleView
	case "logs":
		nextMenu = LogMenu
	case "exec":
		nextMenu = ExecShell
	case "start":
		nextMenu = Confirm(confirmTxt("start", c.GetMeta("name")), c.Start)
	case "stop":
		nextMenu = Confirm(confirmTxt("stop", c.GetMeta("name")), c.Stop)
	case "remove":
		nextMenu = Confirm(confirmTxt("remove", c.GetMeta("name")), c.Remove)
	case "pause":
		nextMenu = Confirm(confirmTxt("pause", c.GetMeta("name")), c.Pause)
	case "unpause":
		nextMenu = Confirm(confirmTxt("unpause", c.GetMeta("name")), c.Unpause)
	case "restart":
		nextMenu = Confirm(confirmTxt("restart", c.GetMeta("name")), c.Restart)
	}

	return nextMenu
}

func LogMenu() MenuFn {

	c := cursor.Selected()
	if c == nil {
		return nil
	}

	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	logs, quit := logReader(c)
	m := widgets.NewTextView(logs)
	m.BorderLabel = fmt.Sprintf("Logs [%s]", c.GetMeta("name"))
	ui.Render(m)

	ui.Handle("/sys/wnd/resize", func(e ui.Event) {
		m.Resize()
	})
	ui.Handle("/sys/kbd/t", func(ui.Event) {
		m.Toggle()
	})
	ui.Handle("/sys/kbd/", func(ui.Event) {
		quit <- true
		ui.StopLoop()
	})
	ui.Loop()
	return nil
}

func ExecShell() MenuFn {
	c := cursor.Selected()

	if c == nil {
		return nil
	}

	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()
	// Detect and execute default shell in container.
	// Execute Ash shell command: /bin/sh -c
	// Reset colors: printf '\e[0m\e[?25h'
	// Clear screen
	// Run default shell for the user. It's configured in /etc/passwd and looks like root:x:0:0:root:/root:/bin/bash:
	//  1. Get current user id: id -un
	//  2. Find user's line in /etc/passwd by grep
	//  3. Extract default user's shell by cutting seven's column separated by :
	//  4. Execute the shell path with eval
	if err := c.Exec([]string{"/bin/sh", "-c", "printf '\\e[0m\\e[?25h' && clear && eval `grep ^$(id -un): /etc/passwd | cut -d : -f 7-`"}); err != nil {
		log.StatusErr(err)
	}

	return nil
}

// Create a confirmation dialog with a given description string and
// func to perform if confirmed
func Confirm(txt string, fn func()) MenuFn {
	menu := func() MenuFn {
		ui.DefaultEvtStream.ResetHandlers()
		defer ui.DefaultEvtStream.ResetHandlers()

		m := menu.NewMenu()
		m.Selectable = true
		m.BorderLabel = "Confirm"
		m.SubText = txt

		items := []menu.Item{
			menu.Item{Val: "cancel", Label: "[c]ancel"},
			menu.Item{Val: "yes", Label: "[y]es"},
		}

		var response bool

		m.AddItems(items...)
		ui.Render(m)

		yes := func() {
			response = true
			ui.StopLoop()
		}

		no := func() {
			response = false
			ui.StopLoop()
		}

		HandleKeys("up", m.Up)
		HandleKeys("down", m.Down)
		HandleKeys("exit", no)
		ui.Handle("/sys/kbd/c", func(ui.Event) { no() })
		ui.Handle("/sys/kbd/y", func(ui.Event) { yes() })

		ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
			switch m.SelectedValue() {
			case "cancel":
				no()
			case "yes":
				yes()
			}
		})

		ui.Loop()
		if response {
			fn()
		}
		return nil
	}
	return menu
}

type toggleLog struct {
	timestamp time.Time
	message   string
}

func (t *toggleLog) Toggle(on bool) string {
	if on {
		return fmt.Sprintf("%s %s", t.timestamp.Format("2006-01-02T15:04:05.999Z07:00"), t.message)
	}
	return t.message
}

func logReader(container *container.Container) (logs chan widgets.ToggleText, quit chan bool) {

	logCollector := container.Logs()
	stream := logCollector.Stream()
	logs = make(chan widgets.ToggleText)
	quit = make(chan bool)

	go func() {
		for {
			select {
			case log := <-stream:
				logs <- &toggleLog{timestamp: log.Timestamp, message: log.Message}
			case <-quit:
				logCollector.Stop()
				close(logs)
				return
			}
		}
	}()
	return
}

func confirmTxt(a, n string) string { return fmt.Sprintf("%s container %s?", a, n) }
<commit_msg>Fix: enable pause [p] only for running or paused container<commit_after>package main

import (
	"fmt"
	"strings"
	"time"

	"github.com/bcicen/ctop/config"
	"github.com/bcicen/ctop/container"
	"github.com/bcicen/ctop/widgets"
	"github.com/bcicen/ctop/widgets/menu"
	ui "github.com/gizak/termui"
)

// MenuFn executes a menu window, returning the next menu or nil
type MenuFn func() MenuFn

var helpDialog = []menu.Item{
	{"<enter> - open container menu", ""},
	{"", ""},
	{"[a] - toggle display of all containers", ""},
	{"[f] - filter displayed containers", ""},
	{"[h] - open this help dialog", ""},
	{"[H] - toggle ctop header", ""},
	{"[s] - select container sort field", ""},
	{"[r] - reverse container sort order", ""},
	{"[o] - open single view", ""},
	{"[l] - view container logs ([t] to toggle timestamp when open)", ""},
	{"[e] - exec shell", ""},
	{"[c] - configure columns", ""},
	{"[S] - save current configuration to file", ""},
	{"[q] - exit ctop", ""},
}

func HelpMenu() MenuFn {
	ui.Clear()
	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	m := menu.NewMenu()
	m.BorderLabel = "Help"
	m.AddItems(helpDialog...)
	ui.Handle("/sys/wnd/resize", func(e ui.Event) {
		ui.Clear()
		ui.Render(m)
	})
	ui.Handle("/sys/kbd/", func(ui.Event) {
		ui.StopLoop()
	})
	ui.Loop()
	return nil
}

func FilterMenu() MenuFn {
	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	i := widgets.NewInput()
	i.BorderLabel = "Filter"
	i.SetY(ui.TermHeight() - i.Height)
	i.Data = config.GetVal("filterStr")
	ui.Render(i)

	// refresh container rows on input
	stream := i.Stream()
	go func() {
		for s := range stream {
			config.Update("filterStr", s)
			RefreshDisplay()
			ui.Render(i)
		}
	}()

	i.InputHandlers()
	ui.Handle("/sys/kbd/<escape>", func(ui.Event) {
		config.Update("filterStr", "")
		ui.StopLoop()
	})
	ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
		config.Update("filterStr", i.Data)
		ui.StopLoop()
	})
	ui.Loop()
	return nil
}

func SortMenu() MenuFn {
	ui.Clear()
	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	m := menu.NewMenu()
	m.Selectable = true
	m.SortItems = true
	m.BorderLabel = "Sort Field"

	for _, field := range container.SortFields() {
		m.AddItems(menu.Item{field, ""})
	}

	// set cursor position to current sort field
	m.SetCursor(config.GetVal("sortField"))

	HandleKeys("up", m.Up)
	HandleKeys("down", m.Down)
	HandleKeys("exit", ui.StopLoop)

	ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
		config.Update("sortField", m.SelectedValue())
		ui.StopLoop()
	})

	ui.Render(m)
	ui.Loop()
	return nil
}

func ColumnsMenu() MenuFn {
	const (
		enabledStr  = "[X]"
		disabledStr = "[ ]"
		padding     = 2
	)

	ui.Clear()
	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	m := menu.NewMenu()
	m.Selectable = true
	m.SortItems = false
	m.BorderLabel = "Columns"
	//m.SubText = "Enabled Columns"

	rebuild := func() {
		// get padding for right alignment of enabled status
		var maxLen int
		for _, col := range config.GlobalColumns {
			if len(col.Label) > maxLen {
				maxLen = len(col.Label)
			}
		}
		maxLen += padding

		// rebuild menu items
		m.ClearItems()
		for _, col := range config.GlobalColumns {
			txt := col.Label + strings.Repeat(" ", maxLen-len(col.Label))
			if col.Enabled {
				txt += enabledStr
			} else {
				txt += disabledStr
			}
			m.AddItems(menu.Item{col.Name, txt})
		}
	}

	upFn := func() {
		config.ColumnLeft(m.SelectedValue())
		m.Up()
		rebuild()
	}

	downFn := func() {
		config.ColumnRight(m.SelectedValue())
		m.Down()
		rebuild()
	}

	toggleFn := func() {
		config.ColumnToggle(m.SelectedValue())
		rebuild()
	}

	rebuild()

	HandleKeys("up", m.Up)
	HandleKeys("down", m.Down)
	HandleKeys("enter", toggleFn)
	HandleKeys("pgup", upFn)
	HandleKeys("pgdown", downFn)

	ui.Handle("/sys/kbd/x", func(ui.Event) { toggleFn() })
	ui.Handle("/sys/kbd/<enter>", func(ui.Event) { toggleFn() })

	HandleKeys("exit", func() {
		cSource, err := cursor.cSuper.Get()
		if err == nil {
			for _, c := range cSource.All() {
				c.RecreateWidgets()
			}
		}
		ui.StopLoop()
	})

	ui.Render(m)
	ui.Loop()
	return nil
}

func ContainerMenu() MenuFn {
	c := cursor.Selected()
	if c == nil {
		return nil
	}

	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	m := menu.NewMenu()
	m.Selectable = true
	m.BorderLabel = "Menu"

	items := []menu.Item{
		menu.Item{Val: "single", Label: "[o] single view"},
		menu.Item{Val: "logs", Label: "[l] log view"},
	}

	if c.Meta["state"] == "running" {
		items = append(items, menu.Item{Val: "stop", Label: "[s] stop"})
		items = append(items, menu.Item{Val: "pause", Label: "[p] pause"})
		items = append(items, menu.Item{Val: "restart", Label: "[r] restart"})
		items = append(items, menu.Item{Val: "exec", Label: "[e] exec shell"})
	}
	if c.Meta["state"] == "exited" || c.Meta["state"] == "created" {
		items = append(items, menu.Item{Val: "start", Label: "[s] start"})
		items = append(items, menu.Item{Val: "remove", Label: "[R] remove"})
	}
	if c.Meta["state"] == "paused" {
		items = append(items, menu.Item{Val: "unpause", Label: "[p] unpause"})
	}
	items = append(items, menu.Item{Val: "cancel", Label: "[c] cancel"})

	m.AddItems(items...)
	ui.Render(m)

	HandleKeys("up", m.Up)
	HandleKeys("down", m.Down)

	var selected string

	// shortcuts
	ui.Handle("/sys/kbd/o", func(ui.Event) {
		selected = "single"
		ui.StopLoop()
	})
	ui.Handle("/sys/kbd/l", func(ui.Event) {
		selected = "logs"
		ui.StopLoop()
	})
	if c.Meta["state"] != "paused" {
		ui.Handle("/sys/kbd/s", func(ui.Event) {
			if c.Meta["state"] == "running" {
				selected = "stop"
			} else {
				selected = "start"
			}
			ui.StopLoop()
		})
	}
	if c.Meta["state"] != "exited" && c.Meta["state"] != "created" {
		ui.Handle("/sys/kbd/p", func(ui.Event) {
			if c.Meta["state"] == "paused" {
				selected = "unpause"
			} else {
				selected = "pause"
			}
			ui.StopLoop()
		})
	}
	if c.Meta["state"] == "running" {
		ui.Handle("/sys/kbd/e", func(ui.Event) {
			selected = "exec"
			ui.StopLoop()
		})
		ui.Handle("/sys/kbd/r", func(ui.Event) {
			selected = "restart"
			ui.StopLoop()
		})
	}
	ui.Handle("/sys/kbd/R", func(ui.Event) {
		selected = "remove"
		ui.StopLoop()
	})
	ui.Handle("/sys/kbd/c", func(ui.Event) {
		ui.StopLoop()
	})

	ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
		selected = m.SelectedValue()
		ui.StopLoop()
	})
	ui.Handle("/sys/kbd/", func(ui.Event) {
		ui.StopLoop()
	})
	ui.Loop()

	var nextMenu MenuFn
	switch selected {
	case "single":
		nextMenu = SingleView
	case "logs":
		nextMenu = LogMenu
	case "exec":
		nextMenu = ExecShell
	case "start":
		nextMenu = Confirm(confirmTxt("start", c.GetMeta("name")), c.Start)
	case "stop":
		nextMenu = Confirm(confirmTxt("stop", c.GetMeta("name")), c.Stop)
	case "remove":
		nextMenu = Confirm(confirmTxt("remove", c.GetMeta("name")), c.Remove)
	case "pause":
		nextMenu = Confirm(confirmTxt("pause", c.GetMeta("name")), c.Pause)
	case "unpause":
		nextMenu = Confirm(confirmTxt("unpause", c.GetMeta("name")), c.Unpause)
	case "restart":
		nextMenu = Confirm(confirmTxt("restart", c.GetMeta("name")), c.Restart)
	}

	return nextMenu
}

func LogMenu() MenuFn {

	c := cursor.Selected()
	if c == nil {
		return nil
	}

	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	logs, quit := logReader(c)
	m := widgets.NewTextView(logs)
	m.BorderLabel = fmt.Sprintf("Logs [%s]", c.GetMeta("name"))
	ui.Render(m)

	ui.Handle("/sys/wnd/resize", func(e ui.Event) {
		m.Resize()
	})
	ui.Handle("/sys/kbd/t", func(ui.Event) {
		m.Toggle()
	})
	ui.Handle("/sys/kbd/", func(ui.Event) {
		quit <- true
		ui.StopLoop()
	})
	ui.Loop()
	return nil
}

func ExecShell() MenuFn {
	c := cursor.Selected()

	if c == nil {
		return nil
	}

	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()
	// Detect and execute default shell in container.
	// Execute Ash shell command: /bin/sh -c
	// Reset colors: printf '\e[0m\e[?25h'
	// Clear screen
	// Run default shell for the user. It's configured in /etc/passwd and looks like root:x:0:0:root:/root:/bin/bash:
	//  1. Get current user id: id -un
	//  2. Find user's line in /etc/passwd by grep
	//  3. Extract default user's shell by cutting seven's column separated by :
	//  4. Execute the shell path with eval
	if err := c.Exec([]string{"/bin/sh", "-c", "printf '\\e[0m\\e[?25h' && clear && eval `grep ^$(id -un): /etc/passwd | cut -d : -f 7-`"}); err != nil {
		log.StatusErr(err)
	}

	return nil
}

// Create a confirmation dialog with a given description string and
// func to perform if confirmed
func Confirm(txt string, fn func()) MenuFn {
	menu := func() MenuFn {
		ui.DefaultEvtStream.ResetHandlers()
		defer ui.DefaultEvtStream.ResetHandlers()

		m := menu.NewMenu()
		m.Selectable = true
		m.BorderLabel = "Confirm"
		m.SubText = txt

		items := []menu.Item{
			menu.Item{Val: "cancel", Label: "[c]ancel"},
			menu.Item{Val: "yes", Label: "[y]es"},
		}

		var response bool

		m.AddItems(items...)
		ui.Render(m)

		yes := func() {
			response = true
			ui.StopLoop()
		}

		no := func() {
			response = false
			ui.StopLoop()
		}

		HandleKeys("up", m.Up)
		HandleKeys("down", m.Down)
		HandleKeys("exit", no)
		ui.Handle("/sys/kbd/c", func(ui.Event) { no() })
		ui.Handle("/sys/kbd/y", func(ui.Event) { yes() })

		ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
			switch m.SelectedValue() {
			case "cancel":
				no()
			case "yes":
				yes()
			}
		})

		ui.Loop()
		if response {
			fn()
		}
		return nil
	}
	return menu
}

type toggleLog struct {
	timestamp time.Time
	message   string
}

func (t *toggleLog) Toggle(on bool) string {
	if on {
		return fmt.Sprintf("%s %s", t.timestamp.Format("2006-01-02T15:04:05.999Z07:00"), t.message)
	}
	return t.message
}

func logReader(container *container.Container) (logs chan widgets.ToggleText, quit chan bool) {

	logCollector := container.Logs()
	stream := logCollector.Stream()
	logs = make(chan widgets.ToggleText)
	quit = make(chan bool)

	go func() {
		for {
			select {
			case log := <-stream:
				logs <- &toggleLog{timestamp: log.Timestamp, message: log.Message}
			case <-quit:
				logCollector.Stop()
				close(logs)
				return
			}
		}
	}()
	return
}

func confirmTxt(a, n string) string { return fmt.Sprintf("%s container %s?", a, n) }
<|endoftext|>


<commit_before>package jsonpatch

import (
	"bytes"
	"encoding/json"
	"fmt"
	"reflect"
)

func merge(cur, patch *lazyNode, mergeMerge bool) *lazyNode {
	curDoc, err := cur.intoDoc()

	if err != nil {
		pruneNulls(patch)
		return patch
	}

	patchDoc, err := patch.intoDoc()

	if err != nil {
		return patch
	}

	mergeDocs(curDoc, patchDoc, mergeMerge)

	return cur
}

func mergeDocs(doc, patch *partialDoc, mergeMerge bool) {
	for k, v := range *patch {
		if v == nil {
			if mergeMerge {
				(*doc)[k] = nil
			} else {
				delete(*doc, k)
			}
		} else {
			cur, ok := (*doc)[k]

			if !ok || cur == nil {
				pruneNulls(v)
				(*doc)[k] = v
			} else {
				(*doc)[k] = merge(cur, v, mergeMerge)
			}
		}
	}
}

func pruneNulls(n *lazyNode) {
	sub, err := n.intoDoc()

	if err == nil {
		pruneDocNulls(sub)
	} else {
		ary, err := n.intoAry()

		if err == nil {
			pruneAryNulls(ary)
		}
	}
}

func pruneDocNulls(doc *partialDoc) *partialDoc {
	for k, v := range *doc {
		if v == nil {
			delete(*doc, k)
		} else {
			pruneNulls(v)
		}
	}

	return doc
}

func pruneAryNulls(ary *partialArray) *partialArray {
	newAry := []*lazyNode{}

	for _, v := range *ary {
		if v != nil {
			pruneNulls(v)
			newAry = append(newAry, v)
		}
	}

	*ary = newAry

	return ary
}

var errBadJSONDoc = fmt.Errorf("Invalid JSON Document")
var errBadJSONPatch = fmt.Errorf("Invalid JSON Patch")
var errBadMergeTypes = fmt.Errorf("Mismatched JSON Documents")

// MergeMergePatches merges two merge patches together, such that
// applying this resulting merged merge patch to a document yields the same
// as merging each merge patch to the document in succession.
func MergeMergePatches(patch1Data, patch2Data []byte) ([]byte, error) {
	return doMergePatch(patch1Data, patch2Data, true)
}

// MergePatch merges the patchData into the docData.
func MergePatch(docData, patchData []byte) ([]byte, error) {
	return doMergePatch(docData, patchData, false)
}

func doMergePatch(docData, patchData []byte, mergeMerge bool) ([]byte, error) {
	doc := &partialDoc{}

	docErr := json.Unmarshal(docData, doc)

	patch := &partialDoc{}

	patchErr := json.Unmarshal(patchData, patch)

	if _, ok := docErr.(*json.SyntaxError); ok {
		return nil, errBadJSONDoc
	}

	if _, ok := patchErr.(*json.SyntaxError); ok {
		return nil, errBadJSONPatch
	}

	if docErr == nil && *doc == nil {
		return nil, errBadJSONDoc
	}

	if patchErr == nil && *patch == nil {
		return nil, errBadJSONPatch
	}

	if docErr != nil || patchErr != nil {
		// Not an error, just not a doc, so we turn straight into the patch
		if patchErr == nil {
			if mergeMerge {
				doc = patch
			} else {
				doc = pruneDocNulls(patch)
			}
		} else {
			patchAry := &partialArray{}
			patchErr = json.Unmarshal(patchData, patchAry)

			if patchErr != nil {
				return nil, errBadJSONPatch
			}

			pruneAryNulls(patchAry)

			out, patchErr := json.Marshal(patchAry)

			if patchErr != nil {
				return nil, errBadJSONPatch
			}

			return out, nil
		}
	} else {
		mergeDocs(doc, patch, mergeMerge)
	}

	return json.Marshal(doc)
}

// resemblesJSONArray indicates whether the byte-slice "appears" to be
// a JSON array or not.
// False-positives are possible, as this function does not check the internal
// structure of the array. It only checks that the outer syntax is present and
// correct.
func resemblesJSONArray(input []byte) bool {
	input = bytes.TrimSpace(input)

	hasPrefix := bytes.HasPrefix(input, []byte("["))
	hasSuffix := bytes.HasSuffix(input, []byte("]"))

	return hasPrefix && hasSuffix
}

// CreateMergePatch will return a merge patch document capable of converting
// the original document(s) to the modified document(s).
// The parameters can be bytes of either two JSON Documents, or two arrays of
// JSON documents.
// The merge patch returned follows the specification defined at http://tools.ietf.org/html/draft-ietf-appsawg-json-merge-patch-07
func CreateMergePatch(originalJSON, modifiedJSON []byte) ([]byte, error) {
	originalResemblesArray := resemblesJSONArray(originalJSON)
	modifiedResemblesArray := resemblesJSONArray(modifiedJSON)

	if originalResemblesArray && modifiedResemblesArray {
		return createArrayMergePatch(originalJSON, modifiedJSON)
	}

	if !originalResemblesArray && !modifiedResemblesArray {
		return createObjectMergePatch(originalJSON, modifiedJSON)
	}

	return nil, errBadMergeTypes
}

// createObjectMergePatch will return a merge-patch document capable of
// converting the original document to the modified document.
func createObjectMergePatch(originalJSON, modifiedJSON []byte) ([]byte, error) {
	originalDoc := map[string]interface{}{}
	modifiedDoc := map[string]interface{}{}

	err := json.Unmarshal(originalJSON, &originalDoc)
	if err != nil {
		return nil, errBadJSONDoc
	}

	err = json.Unmarshal(modifiedJSON, &modifiedDoc)
	if err != nil {
		return nil, errBadJSONDoc
	}

	dest, err := getDiff(originalDoc, modifiedDoc)
	if err != nil {
		return nil, err
	}

	return json.Marshal(dest)
}

// createArrayMergePatch will return an array of merge-patch documents capable
// of converting the original document to the modified document for each
// pair of JSON documents provided in the arrays.
// Arrays of mismatched sizes will result in an error.
func createArrayMergePatch(originalJSON, modifiedJSON []byte) ([]byte, error) {
	originalDocs := []json.RawMessage{}
	modifiedDocs := []json.RawMessage{}

	err := json.Unmarshal(originalJSON, &originalDocs)
	if err != nil {
		return nil, errBadJSONDoc
	}

	err = json.Unmarshal(modifiedJSON, &modifiedDocs)
	if err != nil {
		return nil, errBadJSONDoc
	}

	total := len(originalDocs)
	if len(modifiedDocs) != total {
		return nil, errBadJSONDoc
	}

	result := []json.RawMessage{}
	for i := 0; i < len(originalDocs); i++ {
		original := originalDocs[i]
		modified := modifiedDocs[i]

		patch, err := createObjectMergePatch(original, modified)
		if err != nil {
			return nil, err
		}

		result = append(result, json.RawMessage(patch))
	}

	return json.Marshal(result)
}

// Returns true if the array matches (must be json types).
// As is idiomatic for go, an empty array is not the same as a nil array.
func matchesArray(a, b []interface{}) bool {
	if len(a) != len(b) {
		return false
	}
	if (a == nil && b != nil) || (a != nil && b == nil) {
		return false
	}
	for i := range a {
		if !matchesValue(a[i], b[i]) {
			return false
		}
	}
	return true
}

// Returns true if the values matches (must be json types)
// The types of the values must match, otherwise it will always return false
// If two map[string]interface{} are given, all elements must match.
func matchesValue(av, bv interface{}) bool {
	if reflect.TypeOf(av) != reflect.TypeOf(bv) {
		return false
	}
	switch at := av.(type) {
	case string:
		bt := bv.(string)
		if bt == at {
			return true
		}
	case float64:
		bt := bv.(float64)
		if bt == at {
			return true
		}
	case bool:
		bt := bv.(bool)
		if bt == at {
			return true
		}
	case nil:
		// Both nil, fine.
		return true
	case map[string]interface{}:
		bt := bv.(map[string]interface{})
		for key := range at {
			if !matchesValue(at[key], bt[key]) {
				return false
			}
		}
		for key := range bt {
			if !matchesValue(at[key], bt[key]) {
				return false
			}
		}
		return true
	case []interface{}:
		bt := bv.([]interface{})
		return matchesArray(at, bt)
	}
	return false
}

// getDiff returns the (recursive) difference between a and b as a map[string]interface{}.
func getDiff(a, b map[string]interface{}) (map[string]interface{}, error) {
	into := map[string]interface{}{}
	for key, bv := range b {
		av, ok := a[key]
		// value was added
		if !ok {
			into[key] = bv
			continue
		}
		// If types have changed, replace completely
		if reflect.TypeOf(av) != reflect.TypeOf(bv) {
			into[key] = bv
			continue
		}
		// Types are the same, compare values
		switch at := av.(type) {
		case map[string]interface{}:
			bt := bv.(map[string]interface{})
			dst := make(map[string]interface{}, len(bt))
			dst, err := getDiff(at, bt)
			if err != nil {
				return nil, err
			}
			if len(dst) > 0 {
				into[key] = dst
			}
		case string, float64, bool:
			if !matchesValue(av, bv) {
				into[key] = bv
			}
		case []interface{}:
			bt := bv.([]interface{})
			if !matchesArray(at, bt) {
				into[key] = bv
			}
		case nil:
			switch bv.(type) {
			case nil:
				// Both nil, fine.
			default:
				into[key] = bv
			}
		default:
			panic(fmt.Sprintf("Unknown type:%T in key %s", av, key))
		}
	}
	// Now add all deleted values as nil
	for key := range a {
		_, found := b[key]
		if !found {
			into[key] = nil
		}
	}
	return into, nil
}
<commit_msg>add comments to CreateMergePatch<commit_after>package jsonpatch

import (
	"bytes"
	"encoding/json"
	"fmt"
	"reflect"
)

func merge(cur, patch *lazyNode, mergeMerge bool) *lazyNode {
	curDoc, err := cur.intoDoc()

	if err != nil {
		pruneNulls(patch)
		return patch
	}

	patchDoc, err := patch.intoDoc()

	if err != nil {
		return patch
	}

	mergeDocs(curDoc, patchDoc, mergeMerge)

	return cur
}

func mergeDocs(doc, patch *partialDoc, mergeMerge bool) {
	for k, v := range *patch {
		if v == nil {
			if mergeMerge {
				(*doc)[k] = nil
			} else {
				delete(*doc, k)
			}
		} else {
			cur, ok := (*doc)[k]

			if !ok || cur == nil {
				pruneNulls(v)
				(*doc)[k] = v
			} else {
				(*doc)[k] = merge(cur, v, mergeMerge)
			}
		}
	}
}

func pruneNulls(n *lazyNode) {
	sub, err := n.intoDoc()

	if err == nil {
		pruneDocNulls(sub)
	} else {
		ary, err := n.intoAry()

		if err == nil {
			pruneAryNulls(ary)
		}
	}
}

func pruneDocNulls(doc *partialDoc) *partialDoc {
	for k, v := range *doc {
		if v == nil {
			delete(*doc, k)
		} else {
			pruneNulls(v)
		}
	}

	return doc
}

func pruneAryNulls(ary *partialArray) *partialArray {
	newAry := []*lazyNode{}

	for _, v := range *ary {
		if v != nil {
			pruneNulls(v)
			newAry = append(newAry, v)
		}
	}

	*ary = newAry

	return ary
}

var errBadJSONDoc = fmt.Errorf("Invalid JSON Document")
var errBadJSONPatch = fmt.Errorf("Invalid JSON Patch")
var errBadMergeTypes = fmt.Errorf("Mismatched JSON Documents")

// MergeMergePatches merges two merge patches together, such that
// applying this resulting merged merge patch to a document yields the same
// as merging each merge patch to the document in succession.
func MergeMergePatches(patch1Data, patch2Data []byte) ([]byte, error) {
	return doMergePatch(patch1Data, patch2Data, true)
}

// MergePatch merges the patchData into the docData.
func MergePatch(docData, patchData []byte) ([]byte, error) {
	return doMergePatch(docData, patchData, false)
}

func doMergePatch(docData, patchData []byte, mergeMerge bool) ([]byte, error) {
	doc := &partialDoc{}

	docErr := json.Unmarshal(docData, doc)

	patch := &partialDoc{}

	patchErr := json.Unmarshal(patchData, patch)

	if _, ok := docErr.(*json.SyntaxError); ok {
		return nil, errBadJSONDoc
	}

	if _, ok := patchErr.(*json.SyntaxError); ok {
		return nil, errBadJSONPatch
	}

	if docErr == nil && *doc == nil {
		return nil, errBadJSONDoc
	}

	if patchErr == nil && *patch == nil {
		return nil, errBadJSONPatch
	}

	if docErr != nil || patchErr != nil {
		// Not an error, just not a doc, so we turn straight into the patch
		if patchErr == nil {
			if mergeMerge {
				doc = patch
			} else {
				doc = pruneDocNulls(patch)
			}
		} else {
			patchAry := &partialArray{}
			patchErr = json.Unmarshal(patchData, patchAry)

			if patchErr != nil {
				return nil, errBadJSONPatch
			}

			pruneAryNulls(patchAry)

			out, patchErr := json.Marshal(patchAry)

			if patchErr != nil {
				return nil, errBadJSONPatch
			}

			return out, nil
		}
	} else {
		mergeDocs(doc, patch, mergeMerge)
	}

	return json.Marshal(doc)
}

// resemblesJSONArray indicates whether the byte-slice "appears" to be
// a JSON array or not.
// False-positives are possible, as this function does not check the internal
// structure of the array. It only checks that the outer syntax is present and
// correct.
func resemblesJSONArray(input []byte) bool {
	input = bytes.TrimSpace(input)

	hasPrefix := bytes.HasPrefix(input, []byte("["))
	hasSuffix := bytes.HasSuffix(input, []byte("]"))

	return hasPrefix && hasSuffix
}

// CreateMergePatch will return a merge patch document capable of converting
// the original document(s) to the modified document(s).
// The parameters can be bytes of either two JSON Documents, or two arrays of
// JSON documents.
// The merge patch returned follows the specification defined at http://tools.ietf.org/html/draft-ietf-appsawg-json-merge-patch-07
func CreateMergePatch(originalJSON, modifiedJSON []byte) ([]byte, error) {
	originalResemblesArray := resemblesJSONArray(originalJSON)
	modifiedResemblesArray := resemblesJSONArray(modifiedJSON)

	// Do both byte-slices seem like JSON arrays?
	if originalResemblesArray && modifiedResemblesArray {
		return createArrayMergePatch(originalJSON, modifiedJSON)
	}

	// Are both byte-slices are not arrays? Then they are likely JSON objects...
	if !originalResemblesArray && !modifiedResemblesArray {
		return createObjectMergePatch(originalJSON, modifiedJSON)
	}

	// None of the above? Then return an error because of mismatched types.
	return nil, errBadMergeTypes
}

// createObjectMergePatch will return a merge-patch document capable of
// converting the original document to the modified document.
func createObjectMergePatch(originalJSON, modifiedJSON []byte) ([]byte, error) {
	originalDoc := map[string]interface{}{}
	modifiedDoc := map[string]interface{}{}

	err := json.Unmarshal(originalJSON, &originalDoc)
	if err != nil {
		return nil, errBadJSONDoc
	}

	err = json.Unmarshal(modifiedJSON, &modifiedDoc)
	if err != nil {
		return nil, errBadJSONDoc
	}

	dest, err := getDiff(originalDoc, modifiedDoc)
	if err != nil {
		return nil, err
	}

	return json.Marshal(dest)
}

// createArrayMergePatch will return an array of merge-patch documents capable
// of converting the original document to the modified document for each
// pair of JSON documents provided in the arrays.
// Arrays of mismatched sizes will result in an error.
func createArrayMergePatch(originalJSON, modifiedJSON []byte) ([]byte, error) {
	originalDocs := []json.RawMessage{}
	modifiedDocs := []json.RawMessage{}

	err := json.Unmarshal(originalJSON, &originalDocs)
	if err != nil {
		return nil, errBadJSONDoc
	}

	err = json.Unmarshal(modifiedJSON, &modifiedDocs)
	if err != nil {
		return nil, errBadJSONDoc
	}

	total := len(originalDocs)
	if len(modifiedDocs) != total {
		return nil, errBadJSONDoc
	}

	result := []json.RawMessage{}
	for i := 0; i < len(originalDocs); i++ {
		original := originalDocs[i]
		modified := modifiedDocs[i]

		patch, err := createObjectMergePatch(original, modified)
		if err != nil {
			return nil, err
		}

		result = append(result, json.RawMessage(patch))
	}

	return json.Marshal(result)
}

// Returns true if the array matches (must be json types).
// As is idiomatic for go, an empty array is not the same as a nil array.
func matchesArray(a, b []interface{}) bool {
	if len(a) != len(b) {
		return false
	}
	if (a == nil && b != nil) || (a != nil && b == nil) {
		return false
	}
	for i := range a {
		if !matchesValue(a[i], b[i]) {
			return false
		}
	}
	return true
}

// Returns true if the values matches (must be json types)
// The types of the values must match, otherwise it will always return false
// If two map[string]interface{} are given, all elements must match.
func matchesValue(av, bv interface{}) bool {
	if reflect.TypeOf(av) != reflect.TypeOf(bv) {
		return false
	}
	switch at := av.(type) {
	case string:
		bt := bv.(string)
		if bt == at {
			return true
		}
	case float64:
		bt := bv.(float64)
		if bt == at {
			return true
		}
	case bool:
		bt := bv.(bool)
		if bt == at {
			return true
		}
	case nil:
		// Both nil, fine.
		return true
	case map[string]interface{}:
		bt := bv.(map[string]interface{})
		for key := range at {
			if !matchesValue(at[key], bt[key]) {
				return false
			}
		}
		for key := range bt {
			if !matchesValue(at[key], bt[key]) {
				return false
			}
		}
		return true
	case []interface{}:
		bt := bv.([]interface{})
		return matchesArray(at, bt)
	}
	return false
}

// getDiff returns the (recursive) difference between a and b as a map[string]interface{}.
func getDiff(a, b map[string]interface{}) (map[string]interface{}, error) {
	into := map[string]interface{}{}
	for key, bv := range b {
		av, ok := a[key]
		// value was added
		if !ok {
			into[key] = bv
			continue
		}
		// If types have changed, replace completely
		if reflect.TypeOf(av) != reflect.TypeOf(bv) {
			into[key] = bv
			continue
		}
		// Types are the same, compare values
		switch at := av.(type) {
		case map[string]interface{}:
			bt := bv.(map[string]interface{})
			dst := make(map[string]interface{}, len(bt))
			dst, err := getDiff(at, bt)
			if err != nil {
				return nil, err
			}
			if len(dst) > 0 {
				into[key] = dst
			}
		case string, float64, bool:
			if !matchesValue(av, bv) {
				into[key] = bv
			}
		case []interface{}:
			bt := bv.([]interface{})
			if !matchesArray(at, bt) {
				into[key] = bv
			}
		case nil:
			switch bv.(type) {
			case nil:
				// Both nil, fine.
			default:
				into[key] = bv
			}
		default:
			panic(fmt.Sprintf("Unknown type:%T in key %s", av, key))
		}
	}
	// Now add all deleted values as nil
	for key := range a {
		_, found := b[key]
		if !found {
			into[key] = nil
		}
	}
	return into, nil
}
<|endoftext|>


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

package api

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"

	"github.com/tsuru/config"
	"github.com/tsuru/tsuru/app"
	"github.com/tsuru/tsuru/db"
	"github.com/tsuru/tsuru/db/dbtest"
	"github.com/tsuru/tsuru/io"
	"github.com/tsuru/tsuru/provision/provisiontest"
	"gopkg.in/check.v1"
)

type PlatformSuite struct{}

var _ = check.Suite(&PlatformSuite{})

func (s *PlatformSuite) SetUpTest(c *check.C) {
	config.Set("database:url", "127.0.0.1:27017")
	config.Set("database:name", "tsuru_api_platform_test")
	conn, err := db.Conn()
	c.Assert(err, check.IsNil)
	defer conn.Close()
	dbtest.ClearAllCollections(conn.Apps().Database)
}

func (p *PlatformSuite) TestPlatformAdd(c *check.C) {
	provisioner := provisiontest.ExtensibleFakeProvisioner{
		FakeProvisioner: provisiontest.NewFakeProvisioner(),
	}
	oldProvisioner := app.Provisioner
	app.Provisioner = &provisioner
	defer func() {
		app.Provisioner = oldProvisioner
	}()
	dockerfile_url := "http://localhost/Dockerfile"
	body := fmt.Sprintf("name=%s&dockerfile=%s", "teste", dockerfile_url)
	request, _ := http.NewRequest("POST", "/platforms/add", strings.NewReader(body))
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	recorder := httptest.NewRecorder()
	result := platformAdd(recorder, request, nil)
	c.Assert(result, check.IsNil)
	c.Assert(recorder.Body.String(), check.Equals, "\nOK!\n")
}

func (p *PlatformSuite) TestPlatformUpdate(c *check.C) {
	provisioner := provisiontest.ExtensibleFakeProvisioner{
		FakeProvisioner: provisiontest.NewFakeProvisioner(),
	}
	oldProvisioner := app.Provisioner
	app.Provisioner = &provisioner
	defer func() {
		app.Provisioner = oldProvisioner
	}()
	err := app.PlatformAdd("wat", nil, nil)
	c.Assert(err, check.IsNil)
	dockerfile_url := "http://localhost/Dockerfile"
	body := fmt.Sprintf("dockerfile=%s", dockerfile_url)
	request, _ := http.NewRequest("PUT", "/platforms/wat?:name=wat", strings.NewReader(body))
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	recorder := httptest.NewRecorder()
	result := platformUpdate(recorder, request, nil)
	c.Assert(result, check.IsNil)
	b, err := ioutil.ReadAll(recorder.Body)
	c.Assert(err, check.IsNil)
	var msg io.SimpleJsonMessage
	json.Unmarshal(b, &msg)
	c.Assert(errors.New(msg.Error), check.ErrorMatches, "")
}

func (p *PlatformSuite) TestPlatformUpdateOnlyDisableTrue(c *check.C) {
	provisioner := provisiontest.ExtensibleFakeProvisioner{
		FakeProvisioner: provisiontest.NewFakeProvisioner(),
	}
	oldProvisioner := app.Provisioner
	app.Provisioner = &provisioner
	defer func() {
		app.Provisioner = oldProvisioner
	}()
	err := app.PlatformAdd("wat", nil, nil)
	c.Assert(err, check.IsNil)
	dockerfile_url := ""
	body := fmt.Sprintf("dockerfile=%s", dockerfile_url)
	request, _ := http.NewRequest("PUT", "/platforms/wat?:name=wat&disabled=true", strings.NewReader(body))
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	recorder := httptest.NewRecorder()
	result := platformUpdate(recorder, request, nil)
	c.Assert(result, check.IsNil)
	b, err := ioutil.ReadAll(recorder.Body)
	c.Assert(err, check.IsNil)
	var msg io.SimpleJsonMessage
	json.Unmarshal(b, &msg)
	c.Assert(errors.New(msg.Error), check.ErrorMatches, "")
}

func (p *PlatformSuite) TestPlatformUpdateDisableTrueAndDockerfile(c *check.C) {
	provisioner := provisiontest.ExtensibleFakeProvisioner{
		FakeProvisioner: provisiontest.NewFakeProvisioner(),
	}
	oldProvisioner := app.Provisioner
	app.Provisioner = &provisioner
	defer func() {
		app.Provisioner = oldProvisioner
	}()
	err := app.PlatformAdd("wat", nil, nil)
	c.Assert(err, check.IsNil)
	dockerfile_url := "http://localhost/Dockerfile"
	body := fmt.Sprintf("dockerfile=%s", dockerfile_url)
	request, _ := http.NewRequest("PUT", "/platforms/wat?:name=wat&disabled=true", strings.NewReader(body))
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	recorder := httptest.NewRecorder()
	result := platformUpdate(recorder, request, nil)
	c.Assert(result, check.IsNil)
	b, err := ioutil.ReadAll(recorder.Body)
	c.Assert(err, check.IsNil)
	var msg io.SimpleJsonMessage
	json.Unmarshal(b, &msg)
	c.Assert(errors.New(msg.Error), check.ErrorMatches, "")
}

func (p *PlatformSuite) TestPlatformUpdateOnlyDisableFalse(c *check.C) {
	provisioner := provisiontest.ExtensibleFakeProvisioner{
		FakeProvisioner: provisiontest.NewFakeProvisioner(),
	}
	oldProvisioner := app.Provisioner
	app.Provisioner = &provisioner
	defer func() {
		app.Provisioner = oldProvisioner
	}()
	err := app.PlatformAdd("wat", nil, nil)
	c.Assert(err, check.IsNil)
	dockerfile_url := ""
	body := fmt.Sprintf("dockerfile=%s", dockerfile_url)
	request, _ := http.NewRequest("PUT", "/platforms/wat?:name=wat&disabled=false", strings.NewReader(body))
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	recorder := httptest.NewRecorder()
	result := platformUpdate(recorder, request, nil)
	c.Assert(result, check.IsNil)
	b, err := ioutil.ReadAll(recorder.Body)
	c.Assert(err, check.IsNil)
	var msg io.SimpleJsonMessage
	json.Unmarshal(b, &msg)
	c.Assert(errors.New(msg.Error), check.ErrorMatches, "")
}

func (p *PlatformSuite) TestPlatformUpdateDisableFalseAndDockerfile(c *check.C) {
	provisioner := provisiontest.ExtensibleFakeProvisioner{
		FakeProvisioner: provisiontest.NewFakeProvisioner(),
	}
	oldProvisioner := app.Provisioner
	app.Provisioner = &provisioner
	defer func() {
		app.Provisioner = oldProvisioner
	}()
	err := app.PlatformAdd("wat", nil, nil)
	c.Assert(err, check.IsNil)
	dockerfile_url := "http://localhost/Dockerfile"
	body := fmt.Sprintf("dockerfile=%s", dockerfile_url)
	request, _ := http.NewRequest("PUT", "/platforms/wat?:name=wat&disabled=false", strings.NewReader(body))
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	recorder := httptest.NewRecorder()
	result := platformUpdate(recorder, request, nil)
	c.Assert(result, check.IsNil)
	b, err := ioutil.ReadAll(recorder.Body)
	c.Assert(err, check.IsNil)
	var msg io.SimpleJsonMessage
	json.Unmarshal(b, &msg)
	c.Assert(errors.New(msg.Error), check.ErrorMatches, "")
}

func (p *PlatformSuite) TestPlatformRemove(c *check.C) {
	provisioner := provisiontest.ExtensibleFakeProvisioner{
		FakeProvisioner: provisiontest.NewFakeProvisioner(),
	}
	oldProvisioner := app.Provisioner
	app.Provisioner = &provisioner
	defer func() {
		app.Provisioner = oldProvisioner
	}()
	err := app.PlatformAdd("test", nil, nil)
	c.Assert(err, check.IsNil)
	request, _ := http.NewRequest("DELETE", "/platforms/test?:name=test", nil)
	recorder := httptest.NewRecorder()
	err = platformRemove(recorder, request, nil)
	c.Assert(err, check.IsNil)
}
<commit_msg>api: fix build<commit_after>// Copyright 2015 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package api

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"

	"github.com/tsuru/config"
	"github.com/tsuru/tsuru/app"
	"github.com/tsuru/tsuru/db"
	"github.com/tsuru/tsuru/db/dbtest"
	"github.com/tsuru/tsuru/io"
	"github.com/tsuru/tsuru/provision"
	"github.com/tsuru/tsuru/provision/provisiontest"
	"gopkg.in/check.v1"
)

type PlatformSuite struct{}

var _ = check.Suite(&PlatformSuite{})

func (s *PlatformSuite) SetUpTest(c *check.C) {
	config.Set("database:url", "127.0.0.1:27017")
	config.Set("database:name", "tsuru_api_platform_test")
	conn, err := db.Conn()
	c.Assert(err, check.IsNil)
	defer conn.Close()
	dbtest.ClearAllCollections(conn.Apps().Database)
}

func (p *PlatformSuite) TestPlatformAdd(c *check.C) {
	provisioner := provisiontest.ExtensibleFakeProvisioner{
		FakeProvisioner: provisiontest.NewFakeProvisioner(),
	}
	oldProvisioner := app.Provisioner
	app.Provisioner = &provisioner
	defer func() {
		app.Provisioner = oldProvisioner
	}()
	dockerfile_url := "http://localhost/Dockerfile"
	body := fmt.Sprintf("name=%s&dockerfile=%s", "teste", dockerfile_url)
	request, _ := http.NewRequest("POST", "/platforms/add", strings.NewReader(body))
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	recorder := httptest.NewRecorder()
	result := platformAdd(recorder, request, nil)
	c.Assert(result, check.IsNil)
	c.Assert(recorder.Body.String(), check.Equals, "\nOK!\n")
}

func (p *PlatformSuite) TestPlatformUpdate(c *check.C) {
	provisioner := provisiontest.ExtensibleFakeProvisioner{
		FakeProvisioner: provisiontest.NewFakeProvisioner(),
	}
	oldProvisioner := app.Provisioner
	app.Provisioner = &provisioner
	defer func() {
		app.Provisioner = oldProvisioner
	}()
	err := app.PlatformAdd(provision.PlatformOptions{Name: "wat", Args: nil, Output: nil})
	c.Assert(err, check.IsNil)
	dockerfile_url := "http://localhost/Dockerfile"
	body := fmt.Sprintf("dockerfile=%s", dockerfile_url)
	request, _ := http.NewRequest("PUT", "/platforms/wat?:name=wat", strings.NewReader(body))
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	recorder := httptest.NewRecorder()
	result := platformUpdate(recorder, request, nil)
	c.Assert(result, check.IsNil)
	b, err := ioutil.ReadAll(recorder.Body)
	c.Assert(err, check.IsNil)
	var msg io.SimpleJsonMessage
	json.Unmarshal(b, &msg)
	c.Assert(errors.New(msg.Error), check.ErrorMatches, "")
}

func (p *PlatformSuite) TestPlatformUpdateOnlyDisableTrue(c *check.C) {
	provisioner := provisiontest.ExtensibleFakeProvisioner{
		FakeProvisioner: provisiontest.NewFakeProvisioner(),
	}
	oldProvisioner := app.Provisioner
	app.Provisioner = &provisioner
	defer func() {
		app.Provisioner = oldProvisioner
	}()
	err := app.PlatformAdd(provision.PlatformOptions{Name: "wat", Args: nil, Output: nil})
	c.Assert(err, check.IsNil)
	dockerfile_url := ""
	body := fmt.Sprintf("dockerfile=%s", dockerfile_url)
	request, _ := http.NewRequest("PUT", "/platforms/wat?:name=wat&disabled=true", strings.NewReader(body))
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	recorder := httptest.NewRecorder()
	result := platformUpdate(recorder, request, nil)
	c.Assert(result, check.IsNil)
	b, err := ioutil.ReadAll(recorder.Body)
	c.Assert(err, check.IsNil)
	var msg io.SimpleJsonMessage
	json.Unmarshal(b, &msg)
	c.Assert(errors.New(msg.Error), check.ErrorMatches, "")
}

func (p *PlatformSuite) TestPlatformUpdateDisableTrueAndDockerfile(c *check.C) {
	provisioner := provisiontest.ExtensibleFakeProvisioner{
		FakeProvisioner: provisiontest.NewFakeProvisioner(),
	}
	oldProvisioner := app.Provisioner
	app.Provisioner = &provisioner
	defer func() {
		app.Provisioner = oldProvisioner
	}()
	err := app.PlatformAdd(provision.PlatformOptions{Name: "wat", Args: nil, Output: nil})
	c.Assert(err, check.IsNil)
	dockerfile_url := "http://localhost/Dockerfile"
	body := fmt.Sprintf("dockerfile=%s", dockerfile_url)
	request, _ := http.NewRequest("PUT", "/platforms/wat?:name=wat&disabled=true", strings.NewReader(body))
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	recorder := httptest.NewRecorder()
	result := platformUpdate(recorder, request, nil)
	c.Assert(result, check.IsNil)
	b, err := ioutil.ReadAll(recorder.Body)
	c.Assert(err, check.IsNil)
	var msg io.SimpleJsonMessage
	json.Unmarshal(b, &msg)
	c.Assert(errors.New(msg.Error), check.ErrorMatches, "")
}

func (p *PlatformSuite) TestPlatformUpdateOnlyDisableFalse(c *check.C) {
	provisioner := provisiontest.ExtensibleFakeProvisioner{
		FakeProvisioner: provisiontest.NewFakeProvisioner(),
	}
	oldProvisioner := app.Provisioner
	app.Provisioner = &provisioner
	defer func() {
		app.Provisioner = oldProvisioner
	}()
	err := app.PlatformAdd(provision.PlatformOptions{Name: "wat", Args: nil, Output: nil})
	c.Assert(err, check.IsNil)
	dockerfile_url := ""
	body := fmt.Sprintf("dockerfile=%s", dockerfile_url)
	request, _ := http.NewRequest("PUT", "/platforms/wat?:name=wat&disabled=false", strings.NewReader(body))
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	recorder := httptest.NewRecorder()
	result := platformUpdate(recorder, request, nil)
	c.Assert(result, check.IsNil)
	b, err := ioutil.ReadAll(recorder.Body)
	c.Assert(err, check.IsNil)
	var msg io.SimpleJsonMessage
	json.Unmarshal(b, &msg)
	c.Assert(errors.New(msg.Error), check.ErrorMatches, "")
}

func (p *PlatformSuite) TestPlatformUpdateDisableFalseAndDockerfile(c *check.C) {
	provisioner := provisiontest.ExtensibleFakeProvisioner{
		FakeProvisioner: provisiontest.NewFakeProvisioner(),
	}
	oldProvisioner := app.Provisioner
	app.Provisioner = &provisioner
	defer func() {
		app.Provisioner = oldProvisioner
	}()
	err := app.PlatformAdd(provision.PlatformOptions{Name: "wat", Args: nil, Output: nil})
	c.Assert(err, check.IsNil)
	dockerfile_url := "http://localhost/Dockerfile"
	body := fmt.Sprintf("dockerfile=%s", dockerfile_url)
	request, _ := http.NewRequest("PUT", "/platforms/wat?:name=wat&disabled=false", strings.NewReader(body))
	request.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	recorder := httptest.NewRecorder()
	result := platformUpdate(recorder, request, nil)
	c.Assert(result, check.IsNil)
	b, err := ioutil.ReadAll(recorder.Body)
	c.Assert(err, check.IsNil)
	var msg io.SimpleJsonMessage
	json.Unmarshal(b, &msg)
	c.Assert(errors.New(msg.Error), check.ErrorMatches, "")
}

func (p *PlatformSuite) TestPlatformRemove(c *check.C) {
	provisioner := provisiontest.ExtensibleFakeProvisioner{
		FakeProvisioner: provisiontest.NewFakeProvisioner(),
	}
	oldProvisioner := app.Provisioner
	app.Provisioner = &provisioner
	defer func() {
		app.Provisioner = oldProvisioner
	}()
	err := app.PlatformAdd(provision.PlatformOptions{Name: "test", Args: nil, Output: nil})
	c.Assert(err, check.IsNil)
	request, _ := http.NewRequest("DELETE", "/platforms/test?:name=test", nil)
	recorder := httptest.NewRecorder()
	err = platformRemove(recorder, request, nil)
	c.Assert(err, check.IsNil)
}
<|endoftext|>


<commit_before>// Copyright © 2017 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package router

import (
	"context"
	"errors"
	"io"
	"sync"

	"github.com/TheThingsNetwork/go-utils/grpc/restartstream"
	"github.com/TheThingsNetwork/go-utils/log"
	"github.com/TheThingsNetwork/ttn/api"
	"github.com/TheThingsNetwork/ttn/api/gateway"
	"github.com/TheThingsNetwork/ttn/utils"
	"github.com/golang/protobuf/ptypes/empty"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

// GenericStream is used for sending to and receiving from the router.
type GenericStream interface {
	Uplink(*UplinkMessage)
	Status(*gateway.Status)
	Downlink() (<-chan *DownlinkMessage, error)
	Close()
}

// ClientConfig for router Client
type ClientConfig struct {
	BackgroundContext context.Context
	BufferSize        int
}

// DefaultClientConfig for router Client
var DefaultClientConfig = ClientConfig{
	BackgroundContext: context.Background(),
	BufferSize:        10,
}

// NewClient creates a new Client with the given configuration
func NewClient(config ClientConfig) *Client {
	ctx, cancel := context.WithCancel(config.BackgroundContext)

	return &Client{
		log:    log.Get(),
		ctx:    ctx,
		cancel: cancel,

		config: config,
	}
}

// Client for router
type Client struct {
	log    log.Interface
	ctx    context.Context
	cancel context.CancelFunc

	config      ClientConfig
	serverConns []*serverConn
}

// AddServer adds a router server
func (c *Client) AddServer(name string, conn *grpc.ClientConn) {
	log := c.log.WithField("Router", name)
	log.Info("Adding Router server")
	s := &serverConn{
		ctx:  log,
		name: name,
		conn: conn,
	}
	c.serverConns = append(c.serverConns, s)
}

// Close the client and all its connections
func (c *Client) Close() {
	c.cancel()
	for _, server := range c.serverConns {
		server.Close()
	}
}

type serverConn struct {
	ctx  log.Interface
	name string

	ready chan struct{}
	conn  *grpc.ClientConn
}

func (c *serverConn) Close() {
	if c.ready != nil {
		<-c.ready
	}
	if c.conn != nil {
		c.conn.Close()
	}
}

type gatewayStreams struct {
	log    log.Interface
	ctx    context.Context
	cancel context.CancelFunc

	mu     sync.RWMutex
	uplink map[string]chan *UplinkMessage
	status map[string]chan *gateway.Status

	downlink chan *DownlinkMessage
}

func (s *gatewayStreams) Uplink(msg *UplinkMessage) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	s.log.WithField("Routers", len(s.uplink)).Debug("Sending UplinkMessage to router")
	for serverName, ch := range s.uplink {
		select {
		case ch <- msg:
		default:
			s.log.WithField("Router", serverName).Warn("UplinkMessage buffer full")
		}
	}
}

func (s *gatewayStreams) Status(msg *gateway.Status) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	s.log.WithField("Routers", len(s.status)).Debug("Sending Status to router")
	for serverName, ch := range s.status {
		select {
		case ch <- msg:
		default:
			s.log.WithField("Router", serverName).Warn("GatewayStatus buffer full")
		}
	}
}

// ErrDownlinkInactive is returned by the Downlink func if downlink is inactive
var ErrDownlinkInactive = errors.New("Downlink stream not active")

func (s *gatewayStreams) Downlink() (<-chan *DownlinkMessage, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.downlink == nil {
		return nil, ErrDownlinkInactive
	}
	return s.downlink, nil
}

func (s *gatewayStreams) Close() {
	s.cancel()
}

// NewGatewayStreams returns new streams using the given gateway ID and token
func (c *Client) NewGatewayStreams(id string, token string, downlinkActive bool) GenericStream {
	log := c.log.WithField("GatewayID", id)
	ctx, cancel := context.WithCancel(c.ctx)
	ctx = api.ContextWithID(ctx, id)
	ctx = api.ContextWithToken(ctx, token)
	s := &gatewayStreams{
		log:    log,
		ctx:    ctx,
		cancel: cancel,

		uplink: make(map[string]chan *UplinkMessage),
		status: make(map[string]chan *gateway.Status),
	}

	var wg utils.WaitGroup

	var wgDown sync.WaitGroup
	if downlinkActive {
		s.downlink = make(chan *DownlinkMessage, c.config.BufferSize)
		defer func() {
			go func() {
				wgDown.Wait()
				close(s.downlink)
			}()
		}()
	}

	// Hook up the router servers
	for _, server := range c.serverConns {
		wg.Add(1)
		wgDown.Add(1)
		go func(server *serverConn) {
			if server.ready != nil {
				select {
				case <-ctx.Done():
					return
				case <-server.ready:
				}
			}
			if server.conn == nil {
				return
			}
			log := log.WithField("Router", server.name)
			cli := NewRouterClient(server.conn)

			logStreamErr := func(streamName string, err error) {
				switch {
				case err == nil:
					log.Debugf("%s stream closed", streamName)
				case err == io.EOF:
					log.WithError(err).Debugf("%s stream ended", streamName)
				case err == context.Canceled || grpc.Code(err) == codes.Canceled:
					log.WithError(err).Debugf("%s stream canceled", streamName)
				case err == context.DeadlineExceeded || grpc.Code(err) == codes.DeadlineExceeded:
					log.WithError(err).Debugf("%s stream deadline exceeded", streamName)
				case grpc.ErrorDesc(err) == grpc.ErrClientConnClosing.Error():
					log.WithError(err).Debugf("%s stream connection closed", streamName)
				default:
					log.WithError(err).Warnf("%s stream closed unexpectedly", streamName)
				}
			}

			// Stream channels
			chUplink := make(chan *UplinkMessage, c.config.BufferSize)
			chStatus := make(chan *gateway.Status, c.config.BufferSize)

			defer func() {
				s.mu.Lock()
				defer s.mu.Unlock()
				delete(s.uplink, server.name)
				delete(s.status, server.name)
				close(chUplink)
				close(chStatus)
			}()

			// Uplink stream
			uplink, err := cli.Uplink(ctx)
			if err != nil {
				log.WithError(err).Warn("Could not set up Uplink stream")
			} else {
				s.mu.Lock()
				s.uplink[server.name] = chUplink
				s.mu.Unlock()
				go func() {
					err := uplink.RecvMsg(new(empty.Empty))
					logStreamErr("Uplink", err)
					s.mu.Lock()
					defer s.mu.Unlock()
					delete(s.uplink, server.name)
				}()
			}

			// Downlink stream
			if downlinkActive {
				downlink, err := cli.Subscribe(ctx, &SubscribeRequest{})
				if err != nil {
					log.WithError(err).Warn("Could not set up Subscribe stream")
					wgDown.Done()
				} else {
					go func() {
						defer func() {
							wgDown.Done()
						}()
						for {
							msg, err := downlink.Recv()
							if err != nil {
								logStreamErr("Subscribe", err)
								return
							}
							select {
							case s.downlink <- msg:
							default:
								log.Warn("Downlink buffer full")
							}
						}
					}()
				}
			}

			// Status stream
			status, err := cli.GatewayStatus(ctx)
			if err != nil {
				log.WithError(err).Warn("Could not set up GatewayStatus stream")
			} else {
				s.mu.Lock()
				s.status[server.name] = chStatus
				s.mu.Unlock()
				go func() {
					err := status.RecvMsg(new(empty.Empty))
					logStreamErr("GatewayStatus", err)
					s.mu.Lock()
					defer s.mu.Unlock()
					delete(s.status, server.name)
				}()
			}

			wg.Done()
			log.Debug("Start handling Gateway streams")
			defer log.Debug("Done handling Gateway streams")
			for {
				select {
				case <-ctx.Done():
					return
				case msg := <-chStatus:
					if err := status.Send(msg); err != nil {
						log.WithError(err).Warn("Could not send GatewayStatus to router")
						if err == restartstream.ErrStreamClosed {
							return
						}
					}
				case msg := <-chUplink:
					if err := uplink.Send(msg); err != nil {
						log.WithError(err).Warn("Could not send UplinkMessage to router")
						if err == restartstream.ErrStreamClosed {
							return
						}
					}
				}
			}

		}(server)
	}

	if api.WaitForStreams > 0 {
		wg.WaitForMax(api.WaitForStreams)
	}

	return s
}
<commit_msg>Assume successful stream setup<commit_after>// Copyright © 2017 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package router

import (
	"context"
	"errors"
	"io"
	"sync"

	"github.com/TheThingsNetwork/go-utils/grpc/restartstream"
	"github.com/TheThingsNetwork/go-utils/log"
	"github.com/TheThingsNetwork/ttn/api"
	"github.com/TheThingsNetwork/ttn/api/gateway"
	"github.com/TheThingsNetwork/ttn/utils"
	"github.com/golang/protobuf/ptypes/empty"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

// GenericStream is used for sending to and receiving from the router.
type GenericStream interface {
	Uplink(*UplinkMessage)
	Status(*gateway.Status)
	Downlink() (<-chan *DownlinkMessage, error)
	Close()
}

// ClientConfig for router Client
type ClientConfig struct {
	BackgroundContext context.Context
	BufferSize        int
}

// DefaultClientConfig for router Client
var DefaultClientConfig = ClientConfig{
	BackgroundContext: context.Background(),
	BufferSize:        10,
}

// NewClient creates a new Client with the given configuration
func NewClient(config ClientConfig) *Client {
	ctx, cancel := context.WithCancel(config.BackgroundContext)

	return &Client{
		log:    log.Get(),
		ctx:    ctx,
		cancel: cancel,

		config: config,
	}
}

// Client for router
type Client struct {
	log    log.Interface
	ctx    context.Context
	cancel context.CancelFunc

	config      ClientConfig
	serverConns []*serverConn
}

// AddServer adds a router server
func (c *Client) AddServer(name string, conn *grpc.ClientConn) {
	log := c.log.WithField("Router", name)
	log.Info("Adding Router server")
	s := &serverConn{
		ctx:  log,
		name: name,
		conn: conn,
	}
	c.serverConns = append(c.serverConns, s)
}

// Close the client and all its connections
func (c *Client) Close() {
	c.cancel()
	for _, server := range c.serverConns {
		server.Close()
	}
}

type serverConn struct {
	ctx  log.Interface
	name string

	ready chan struct{}
	conn  *grpc.ClientConn
}

func (c *serverConn) Close() {
	if c.ready != nil {
		<-c.ready
	}
	if c.conn != nil {
		c.conn.Close()
	}
}

type gatewayStreams struct {
	log    log.Interface
	ctx    context.Context
	cancel context.CancelFunc

	mu     sync.RWMutex
	uplink map[string]chan *UplinkMessage
	status map[string]chan *gateway.Status

	downlink chan *DownlinkMessage
}

func (s *gatewayStreams) Uplink(msg *UplinkMessage) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	s.log.WithField("Routers", len(s.uplink)).Debug("Sending UplinkMessage to router")
	for serverName, ch := range s.uplink {
		select {
		case ch <- msg:
		default:
			s.log.WithField("Router", serverName).Warn("UplinkMessage buffer full")
		}
	}
}

func (s *gatewayStreams) Status(msg *gateway.Status) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	s.log.WithField("Routers", len(s.status)).Debug("Sending Status to router")
	for serverName, ch := range s.status {
		select {
		case ch <- msg:
		default:
			s.log.WithField("Router", serverName).Warn("GatewayStatus buffer full")
		}
	}
}

// ErrDownlinkInactive is returned by the Downlink func if downlink is inactive
var ErrDownlinkInactive = errors.New("Downlink stream not active")

func (s *gatewayStreams) Downlink() (<-chan *DownlinkMessage, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.downlink == nil {
		return nil, ErrDownlinkInactive
	}
	return s.downlink, nil
}

func (s *gatewayStreams) Close() {
	s.cancel()
}

// NewGatewayStreams returns new streams using the given gateway ID and token
func (c *Client) NewGatewayStreams(id string, token string, downlinkActive bool) GenericStream {
	log := c.log.WithField("GatewayID", id)
	ctx, cancel := context.WithCancel(c.ctx)
	ctx = api.ContextWithID(ctx, id)
	ctx = api.ContextWithToken(ctx, token)
	s := &gatewayStreams{
		log:    log,
		ctx:    ctx,
		cancel: cancel,

		uplink: make(map[string]chan *UplinkMessage),
		status: make(map[string]chan *gateway.Status),
	}

	var wg utils.WaitGroup

	var wgDown sync.WaitGroup
	if downlinkActive {
		s.downlink = make(chan *DownlinkMessage, c.config.BufferSize)
		defer func() {
			go func() {
				wgDown.Wait()
				close(s.downlink)
			}()
		}()
	}

	// Hook up the router servers
	for _, server := range c.serverConns {
		wg.Add(1)
		wgDown.Add(1)

		// Stream channels
		chUplink := make(chan *UplinkMessage, c.config.BufferSize)
		chStatus := make(chan *gateway.Status, c.config.BufferSize)

		s.mu.Lock()
		s.uplink[server.name] = chUplink
		s.status[server.name] = chStatus
		s.mu.Unlock()

		go func(server *serverConn) {
			defer func() {
				s.mu.Lock()
				defer s.mu.Unlock()
				delete(s.uplink, server.name)
				delete(s.status, server.name)
				close(chUplink)
				close(chStatus)
			}()

			if server.ready != nil {
				select {
				case <-ctx.Done():
					return
				case <-server.ready:
				}
			}
			if server.conn == nil {
				return
			}
			log := log.WithField("Router", server.name)
			cli := NewRouterClient(server.conn)

			logStreamErr := func(streamName string, err error) {
				switch {
				case err == nil:
					log.Debugf("%s stream closed", streamName)
				case err == io.EOF:
					log.WithError(err).Debugf("%s stream ended", streamName)
				case err == context.Canceled || grpc.Code(err) == codes.Canceled:
					log.WithError(err).Debugf("%s stream canceled", streamName)
				case err == context.DeadlineExceeded || grpc.Code(err) == codes.DeadlineExceeded:
					log.WithError(err).Debugf("%s stream deadline exceeded", streamName)
				case grpc.ErrorDesc(err) == grpc.ErrClientConnClosing.Error():
					log.WithError(err).Debugf("%s stream connection closed", streamName)
				default:
					log.WithError(err).Warnf("%s stream closed unexpectedly", streamName)
				}
			}

			// Uplink stream
			uplink, err := cli.Uplink(ctx)
			if err != nil {
				log.WithError(err).Warn("Could not set up Uplink stream")
				s.mu.Lock()
				delete(s.uplink, server.name)
				s.mu.Unlock()
			} else {
				go func() {
					err := uplink.RecvMsg(new(empty.Empty))
					logStreamErr("Uplink", err)
					s.mu.Lock()
					defer s.mu.Unlock()
					delete(s.uplink, server.name)
				}()
			}

			// Downlink stream
			if downlinkActive {
				downlink, err := cli.Subscribe(ctx, &SubscribeRequest{})
				if err != nil {
					log.WithError(err).Warn("Could not set up Subscribe stream")
					wgDown.Done()
				} else {
					go func() {
						defer func() {
							wgDown.Done()
						}()
						for {
							msg, err := downlink.Recv()
							if err != nil {
								logStreamErr("Subscribe", err)
								return
							}
							select {
							case s.downlink <- msg:
							default:
								log.Warn("Downlink buffer full")
							}
						}
					}()
				}
			}

			// Status stream
			status, err := cli.GatewayStatus(ctx)
			if err != nil {
				log.WithError(err).Warn("Could not set up GatewayStatus stream")
				s.mu.Lock()
				delete(s.status, server.name)
				s.mu.Unlock()
			} else {
				go func() {
					err := status.RecvMsg(new(empty.Empty))
					logStreamErr("GatewayStatus", err)
					s.mu.Lock()
					defer s.mu.Unlock()
					delete(s.status, server.name)
				}()
			}

			wg.Done()
			log.Debug("Start handling Gateway streams")
			defer log.Debug("Done handling Gateway streams")
			for {
				select {
				case <-ctx.Done():
					return
				case msg := <-chStatus:
					if err := status.Send(msg); err != nil {
						log.WithError(err).Warn("Could not send GatewayStatus to router")
						if err == restartstream.ErrStreamClosed {
							return
						}
					}
				case msg := <-chUplink:
					if err := uplink.Send(msg); err != nil {
						log.WithError(err).Warn("Could not send UplinkMessage to router")
						if err == restartstream.ErrStreamClosed {
							return
						}
					}
				}
			}

		}(server)
	}

	if api.WaitForStreams > 0 {
		wg.WaitForMax(api.WaitForStreams)
	}

	return s
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"path"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/fnproject/fn/api"
	"github.com/fnproject/fn/api/id"
	"github.com/fnproject/fn/api/models"
	"github.com/fnproject/fn/api/runner"
	"github.com/fnproject/fn/api/runner/common"
	"github.com/fnproject/fn/api/runner/task"
	"github.com/gin-gonic/gin"
	cache "github.com/patrickmn/go-cache"
)

type runnerResponse struct {
	RequestID string            `json:"request_id,omitempty"`
	Error     *models.ErrorBody `json:"error,omitempty"`
}

func toEnvName(envtype, name string) string {
	name = strings.ToUpper(strings.Replace(name, "-", "_", -1))
	if envtype == "" {
		return name
	}
	return fmt.Sprintf("%s_%s", envtype, name)
}

func (s *Server) handleRequest(c *gin.Context, enqueue models.Enqueue) {
	if strings.HasPrefix(c.Request.URL.Path, "/v1") {
		c.Status(http.StatusNotFound)
		return
	}

	ctx := c.Request.Context()

	reqID := id.New().String()
	ctx, log := common.LoggerWithFields(ctx, logrus.Fields{"call_id": reqID})

	var err error
	var payload io.Reader

	if c.Request.Method == "POST" {
		payload = c.Request.Body
		// Load complete body and close
		defer func() {
			io.Copy(ioutil.Discard, c.Request.Body)
			c.Request.Body.Close()
		}()
	} else if c.Request.Method == "GET" {
		reqPayload := c.Request.URL.Query().Get("payload")
		payload = strings.NewReader(reqPayload)
	}

	r, routeExists := c.Get(api.Path)
	if !routeExists {
		r = "/"
	}

	reqRoute := &models.Route{
		AppName: c.MustGet(api.AppName).(string),
		Path:    path.Clean(r.(string)),
	}

	s.FireBeforeDispatch(ctx, reqRoute)

	appName := reqRoute.AppName
	path := reqRoute.Path

	app, err := s.Datastore.GetApp(ctx, appName)
	if err != nil {
		handleErrorResponse(c, err)
		return
	} else if app == nil {
		handleErrorResponse(c, models.ErrAppsNotFound)
		return
	}

	log.WithFields(logrus.Fields{"app": appName, "path": path}).Debug("Finding route on datastore")
	route, err := s.loadroute(ctx, appName, path)
	if err != nil {
		handleErrorResponse(c, err)
		return
	}

	if route == nil {
		handleErrorResponse(c, models.ErrRoutesNotFound)
		return
	}

	log = log.WithFields(logrus.Fields{"app": appName, "path": route.Path, "image": route.Image})
	log.Debug("Got route from datastore")

	if s.serve(ctx, c, appName, route, app, path, reqID, payload, enqueue) {
		s.FireAfterDispatch(ctx, reqRoute)
		return
	}

	handleErrorResponse(c, models.ErrRoutesNotFound)
}

func (s *Server) loadroute(ctx context.Context, appName, path string) (*models.Route, error) {
	if route, ok := s.cacheget(appName, path); ok {
		return route, nil
	}
	key := routeCacheKey(appName, path)
	resp, err := s.singleflight.do(
		key,
		func() (interface{}, error) {
			return s.Datastore.GetRoute(ctx, appName, path)
		},
	)
	if err != nil {
		return nil, err
	}
	route := resp.(*models.Route)
	s.routeCache.Set(key, route, cache.DefaultExpiration)
	return route, nil
}

// TODO: Should remove *gin.Context from these functions, should use only context.Context
func (s *Server) serve(ctx context.Context, c *gin.Context, appName string, route *models.Route, app *models.App, path, reqID string, payload io.Reader, enqueue models.Enqueue) (ok bool) {
	ctx, log := common.LoggerWithFields(ctx, logrus.Fields{"app": appName, "route": route.Path, "image": route.Image})

	params, match := matchRoute(route.Path, path)
	if !match {
		return false
	}

	var stdout bytes.Buffer // TODO: should limit the size of this, error if gets too big. akin to: https://golang.org/pkg/io/#LimitReader

	if route.Format == "" {
		route.Format = "default"
	}

	// baseVars are the vars on the route & app, not on this specific request [for hot functions]
	baseVars := make(map[string]string, len(app.Config)+len(route.Config)+3)
	baseVars["FN_FORMAT"] = route.Format
	baseVars["APP_NAME"] = appName
	baseVars["ROUTE"] = route.Path
	baseVars["MEMORY_MB"] = fmt.Sprintf("%d", route.Memory)

	// app config
	for k, v := range app.Config {
		k = toEnvName("", k)
		baseVars[k] = v
	}
	for k, v := range route.Config {
		k = toEnvName("", k)
		baseVars[k] = v
	}

	// envVars contains the full set of env vars, per request + base
	envVars := make(map[string]string, len(baseVars)+len(params)+len(c.Request.Header)+3)

	for k, v := range baseVars {
		envVars[k] = v
	}

	envVars["CALL_ID"] = reqID
	envVars["METHOD"] = c.Request.Method
	envVars["REQUEST_URL"] = fmt.Sprintf("%v//%v%v", func() string {
		if c.Request.TLS == nil {
			return "http"
		}
		return "https"
	}(), c.Request.Host, c.Request.URL.String())

	// params
	for _, param := range params {
		envVars[toEnvName("PARAM", param.Key)] = param.Value
	}

	// headers
	for header, value := range c.Request.Header {
		envVars[toEnvName("HEADER", header)] = strings.Join(value, ", ")
	}

	cfg := &task.Config{
		AppName:      appName,
		Path:         route.Path,
		BaseEnv:      baseVars,
		Env:          envVars,
		Format:       route.Format,
		ID:           reqID,
		Image:        route.Image,
		Memory:       route.Memory,
		Stdin:        payload,
		Stdout:       &stdout,
		Timeout:      time.Duration(route.Timeout) * time.Second,
		IdleTimeout:  time.Duration(route.IdleTimeout) * time.Second,
		ReceivedTime: time.Now(),
		Ready:        make(chan struct{}),
	}

	// ensure valid values
	if cfg.Timeout <= 0 {
		cfg.Timeout = runner.DefaultTimeout
	}
	if cfg.IdleTimeout <= 0 {
		cfg.IdleTimeout = runner.DefaultIdleTimeout
	}

	s.Runner.Enqueue()
	newTask := task.TaskFromConfig(cfg)

	switch route.Type {
	case "async":
		// TODO we should be able to do hot input to async. plumb protocol stuff
		// TODO enqueue should unravel the payload?

		// Read payload
		pl, err := ioutil.ReadAll(cfg.Stdin)
		if err != nil {
			handleErrorResponse(c, models.ErrInvalidPayload)
			return true
		}
		// Add in payload
		newTask.Payload = string(pl)

		// Push to queue
		_, err = enqueue(c, s.MQ, newTask)
		if err != nil {
			handleErrorResponse(c, err)
			return true
		}

		log.Info("Added new task to queue")
		c.JSON(http.StatusAccepted, map[string]string{"call_id": cfg.ID})

	default:
		result, err := s.Runner.RunTrackedTask(newTask, ctx, cfg)
		if result != nil {
			waitTime := result.StartTime().Sub(cfg.ReceivedTime)
			c.Header("XXX-FXLB-WAIT", waitTime.String())
		}

		if err != nil {
			c.JSON(http.StatusInternalServerError, runnerResponse{
				RequestID: cfg.ID,
				Error: &models.ErrorBody{
					Message: err.Error(),
				},
			})
			log.WithError(err).Error("Failed to run task")
			break
		}

		for k, v := range route.Headers {
			c.Header(k, v[0])
		}

		// this will help users to track sync execution in a manner of async
		// FN_CALL_ID is an equivalent of call_id
		c.Header("FN_CALL_ID", newTask.ID)

		switch result.Status() {
		case "success":
			c.Data(http.StatusOK, "", stdout.Bytes())
		case "timeout":
			c.JSON(http.StatusGatewayTimeout, runnerResponse{
				RequestID: cfg.ID,
				Error: &models.ErrorBody{
					Message: models.ErrRunnerTimeout.Error(),
				},
			})
		default:
			c.JSON(http.StatusInternalServerError, runnerResponse{
				RequestID: cfg.ID,
				Error: &models.ErrorBody{
					Message: result.Error(),
				},
			})
		}
	}

	return true
}

var fakeHandler = func(http.ResponseWriter, *http.Request, Params) {}

func matchRoute(baseRoute, route string) (Params, bool) {
	tree := &node{}
	tree.addRoute(baseRoute, fakeHandler)
	handler, p, _ := tree.getValue(route)
	if handler == nil {
		return nil, false
	}

	return p, true
}
<commit_msg>Fixed the formatting of the passed URL<commit_after>package server

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"path"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/fnproject/fn/api"
	"github.com/fnproject/fn/api/id"
	"github.com/fnproject/fn/api/models"
	"github.com/fnproject/fn/api/runner"
	"github.com/fnproject/fn/api/runner/common"
	"github.com/fnproject/fn/api/runner/task"
	"github.com/gin-gonic/gin"
	cache "github.com/patrickmn/go-cache"
)

type runnerResponse struct {
	RequestID string            `json:"request_id,omitempty"`
	Error     *models.ErrorBody `json:"error,omitempty"`
}

func toEnvName(envtype, name string) string {
	name = strings.ToUpper(strings.Replace(name, "-", "_", -1))
	if envtype == "" {
		return name
	}
	return fmt.Sprintf("%s_%s", envtype, name)
}

func (s *Server) handleRequest(c *gin.Context, enqueue models.Enqueue) {
	if strings.HasPrefix(c.Request.URL.Path, "/v1") {
		c.Status(http.StatusNotFound)
		return
	}

	ctx := c.Request.Context()

	reqID := id.New().String()
	ctx, log := common.LoggerWithFields(ctx, logrus.Fields{"call_id": reqID})

	var err error
	var payload io.Reader

	if c.Request.Method == "POST" {
		payload = c.Request.Body
		// Load complete body and close
		defer func() {
			io.Copy(ioutil.Discard, c.Request.Body)
			c.Request.Body.Close()
		}()
	} else if c.Request.Method == "GET" {
		reqPayload := c.Request.URL.Query().Get("payload")
		payload = strings.NewReader(reqPayload)
	}

	r, routeExists := c.Get(api.Path)
	if !routeExists {
		r = "/"
	}

	reqRoute := &models.Route{
		AppName: c.MustGet(api.AppName).(string),
		Path:    path.Clean(r.(string)),
	}

	s.FireBeforeDispatch(ctx, reqRoute)

	appName := reqRoute.AppName
	path := reqRoute.Path

	app, err := s.Datastore.GetApp(ctx, appName)
	if err != nil {
		handleErrorResponse(c, err)
		return
	} else if app == nil {
		handleErrorResponse(c, models.ErrAppsNotFound)
		return
	}

	log.WithFields(logrus.Fields{"app": appName, "path": path}).Debug("Finding route on datastore")
	route, err := s.loadroute(ctx, appName, path)
	if err != nil {
		handleErrorResponse(c, err)
		return
	}

	if route == nil {
		handleErrorResponse(c, models.ErrRoutesNotFound)
		return
	}

	log = log.WithFields(logrus.Fields{"app": appName, "path": route.Path, "image": route.Image})
	log.Debug("Got route from datastore")

	if s.serve(ctx, c, appName, route, app, path, reqID, payload, enqueue) {
		s.FireAfterDispatch(ctx, reqRoute)
		return
	}

	handleErrorResponse(c, models.ErrRoutesNotFound)
}

func (s *Server) loadroute(ctx context.Context, appName, path string) (*models.Route, error) {
	if route, ok := s.cacheget(appName, path); ok {
		return route, nil
	}
	key := routeCacheKey(appName, path)
	resp, err := s.singleflight.do(
		key,
		func() (interface{}, error) {
			return s.Datastore.GetRoute(ctx, appName, path)
		},
	)
	if err != nil {
		return nil, err
	}
	route := resp.(*models.Route)
	s.routeCache.Set(key, route, cache.DefaultExpiration)
	return route, nil
}

// TODO: Should remove *gin.Context from these functions, should use only context.Context
func (s *Server) serve(ctx context.Context, c *gin.Context, appName string, route *models.Route, app *models.App, path, reqID string, payload io.Reader, enqueue models.Enqueue) (ok bool) {
	ctx, log := common.LoggerWithFields(ctx, logrus.Fields{"app": appName, "route": route.Path, "image": route.Image})

	params, match := matchRoute(route.Path, path)
	if !match {
		return false
	}

	var stdout bytes.Buffer // TODO: should limit the size of this, error if gets too big. akin to: https://golang.org/pkg/io/#LimitReader

	if route.Format == "" {
		route.Format = "default"
	}

	// baseVars are the vars on the route & app, not on this specific request [for hot functions]
	baseVars := make(map[string]string, len(app.Config)+len(route.Config)+3)
	baseVars["FN_FORMAT"] = route.Format
	baseVars["APP_NAME"] = appName
	baseVars["ROUTE"] = route.Path
	baseVars["MEMORY_MB"] = fmt.Sprintf("%d", route.Memory)

	// app config
	for k, v := range app.Config {
		k = toEnvName("", k)
		baseVars[k] = v
	}
	for k, v := range route.Config {
		k = toEnvName("", k)
		baseVars[k] = v
	}

	// envVars contains the full set of env vars, per request + base
	envVars := make(map[string]string, len(baseVars)+len(params)+len(c.Request.Header)+3)

	for k, v := range baseVars {
		envVars[k] = v
	}

	envVars["CALL_ID"] = reqID
	envVars["METHOD"] = c.Request.Method
	envVars["REQUEST_URL"] = fmt.Sprintf("%v://%v%v", func() string {
		if c.Request.TLS == nil {
			return "http"
		}
		return "https"
	}(), c.Request.Host, c.Request.URL.String())

	// params
	for _, param := range params {
		envVars[toEnvName("PARAM", param.Key)] = param.Value
	}

	// headers
	for header, value := range c.Request.Header {
		envVars[toEnvName("HEADER", header)] = strings.Join(value, ", ")
	}

	cfg := &task.Config{
		AppName:      appName,
		Path:         route.Path,
		BaseEnv:      baseVars,
		Env:          envVars,
		Format:       route.Format,
		ID:           reqID,
		Image:        route.Image,
		Memory:       route.Memory,
		Stdin:        payload,
		Stdout:       &stdout,
		Timeout:      time.Duration(route.Timeout) * time.Second,
		IdleTimeout:  time.Duration(route.IdleTimeout) * time.Second,
		ReceivedTime: time.Now(),
		Ready:        make(chan struct{}),
	}

	// ensure valid values
	if cfg.Timeout <= 0 {
		cfg.Timeout = runner.DefaultTimeout
	}
	if cfg.IdleTimeout <= 0 {
		cfg.IdleTimeout = runner.DefaultIdleTimeout
	}

	s.Runner.Enqueue()
	newTask := task.TaskFromConfig(cfg)

	switch route.Type {
	case "async":
		// TODO we should be able to do hot input to async. plumb protocol stuff
		// TODO enqueue should unravel the payload?

		// Read payload
		pl, err := ioutil.ReadAll(cfg.Stdin)
		if err != nil {
			handleErrorResponse(c, models.ErrInvalidPayload)
			return true
		}
		// Add in payload
		newTask.Payload = string(pl)

		// Push to queue
		_, err = enqueue(c, s.MQ, newTask)
		if err != nil {
			handleErrorResponse(c, err)
			return true
		}

		log.Info("Added new task to queue")
		c.JSON(http.StatusAccepted, map[string]string{"call_id": cfg.ID})

	default:
		result, err := s.Runner.RunTrackedTask(newTask, ctx, cfg)
		if result != nil {
			waitTime := result.StartTime().Sub(cfg.ReceivedTime)
			c.Header("XXX-FXLB-WAIT", waitTime.String())
		}

		if err != nil {
			c.JSON(http.StatusInternalServerError, runnerResponse{
				RequestID: cfg.ID,
				Error: &models.ErrorBody{
					Message: err.Error(),
				},
			})
			log.WithError(err).Error("Failed to run task")
			break
		}

		for k, v := range route.Headers {
			c.Header(k, v[0])
		}

		// this will help users to track sync execution in a manner of async
		// FN_CALL_ID is an equivalent of call_id
		c.Header("FN_CALL_ID", newTask.ID)

		switch result.Status() {
		case "success":
			c.Data(http.StatusOK, "", stdout.Bytes())
		case "timeout":
			c.JSON(http.StatusGatewayTimeout, runnerResponse{
				RequestID: cfg.ID,
				Error: &models.ErrorBody{
					Message: models.ErrRunnerTimeout.Error(),
				},
			})
		default:
			c.JSON(http.StatusInternalServerError, runnerResponse{
				RequestID: cfg.ID,
				Error: &models.ErrorBody{
					Message: result.Error(),
				},
			})
		}
	}

	return true
}

var fakeHandler = func(http.ResponseWriter, *http.Request, Params) {}

func matchRoute(baseRoute, route string) (Params, bool) {
	tree := &node{}
	tree.addRoute(baseRoute, fakeHandler)
	handler, p, _ := tree.getValue(route)
	if handler == nil {
		return nil, false
	}

	return p, true
}
<|endoftext|>


<commit_before>// Package models creates the models for the modules
package models

import (
	"bytes"
	"fmt"
	"strings"
	"text/template"

	"github.com/cihangir/gene/generators/constants"
	"github.com/cihangir/gene/generators/constructors"
	"github.com/cihangir/gene/generators/sql/statements"
	"github.com/cihangir/gene/generators/validators"
	"github.com/cihangir/gene/writers"
	"github.com/cihangir/schema"
	"github.com/cihangir/stringext"
)

// Generate creates the models and write them to the required paths
func Generate(rootPath string, s *schema.Schema) error {

	for _, def := range s.Definitions {
		// create models only for objects
		if def.Type != nil {
			if t, ok := def.Type.(string); ok {
				if t != "object" {
					continue
				}
			}
		}

		moduleName := strings.ToLower(def.Title)

		modelFilePath := fmt.Sprintf(
			"%smodels/%s.go",
			rootPath,
			moduleName,
		)

		f, err := GenerateModel(def)
		if err != nil {
			return err
		}

		if err := writers.WriteFormattedFile(modelFilePath, f); err != nil {
			return err
		}
	}
	return nil
}

// GenerateStatements generates the basic CRUD statements for the models
func GenerateStatements(rootPath string, s *schema.Schema) error {

	for _, def := range s.Definitions {
		// create models only for objects
		if def.Type != nil {
			if t, ok := def.Type.(string); ok {
				if t != "object" {
					continue
				}
			}
		}

		moduleName := strings.ToLower(def.Title)

		modelFilePath := fmt.Sprintf(
			"%smodels/%s_statements.go",
			rootPath,
			moduleName,
		)

		f, err := GenerateModelStatements(def)
		if err != nil {
			return err
		}

		if err := writers.WriteFormattedFile(modelFilePath, f); err != nil {
			return err
		}
	}
	return nil
}

// GenerateModel generates the model itself
func GenerateModel(s *schema.Schema) ([]byte, error) {
	var buf bytes.Buffer

	packageLine, err := GeneratePackage(s)
	if err != nil {
		return nil, err
	}

	consts, err := constants.Generate(s)
	if err != nil {
		return nil, err
	}

	schema, err := GenerateSchema(s)
	if err != nil {
		return nil, err
	}

	constructor, err := constructors.Generate(s)
	if err != nil {
		return nil, err
	}

	validators, err := validators.Generate(s)
	if err != nil {
		return nil, err
	}

	buf.Write(packageLine)
	buf.Write(consts)
	buf.Write(schema)
	buf.Write(constructor)
	if validators != nil {
		buf.Write(validators)
	}

	return writers.Clear(buf)
}

// GenerateModelStatements generates the CRUD statements for the model struct
func GenerateModelStatements(s *schema.Schema) ([]byte, error) {
	var buf bytes.Buffer

	packageLine, err := GeneratePackage(s)
	if err != nil {
		return nil, err
	}

	createStatements, err := statements.GenerateCreate(s)
	if err != nil {
		return nil, err
	}

	updateStatements, err := statements.GenerateUpdate(s)
	if err != nil {
		return nil, err
	}

	deleteStatements, err := statements.GenerateDelete(s)
	if err != nil {
		return nil, err
	}

	selectStatements, err := statements.GenerateSelect(s)
	if err != nil {
		return nil, err
	}

	tableName, err := statements.GenerateTableName(s)
	if err != nil {
		return nil, err
	}

	buf.Write(packageLine)
	buf.Write(createStatements)
	buf.Write(updateStatements)
	buf.Write(deleteStatements)
	buf.Write(selectStatements)
	buf.Write(tableName)

	return writers.Clear(buf)
}

// GeneratePackage generates the imports according to the schema.
// TODO remove this function
func GeneratePackage(s *schema.Schema) ([]byte, error) {
	temp := template.New("package.tmpl")
	_, err := temp.Parse(PackageTemplate)
	if err != nil {
		return nil, err
	}

	var buf bytes.Buffer

	// name := strings.ToLower(strings.Split(s.Title, " ")[0])
	name := "models"
	err = temp.ExecuteTemplate(&buf, "package.tmpl", name)
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

// GenerateSchema generates the schema.
func GenerateSchema(s *schema.Schema) ([]byte, error) {

	temp := template.New("schema.tmpl")
	temp.Funcs(schema.Helpers)

	_, err := temp.Parse(StructTemplate)
	if err != nil {
		return nil, err
	}

	var buf bytes.Buffer

	context := struct {
		Name       string
		Definition *schema.Schema
	}{
		Name:       s.Title,
		Definition: s,
	}

	err = temp.ExecuteTemplate(&buf, "schema.tmpl", context)
	if err != nil {
		return nil, err
	}

	return writers.Clear(buf)
}

// GenerateFunctions generates the functions according to the schema.
func GenerateFunctions(s *schema.Schema) ([]byte, error) {

	temp := template.New("functions.tmpl")
	temp.Funcs(template.FuncMap{
		"Pointerize": stringext.Pointerize,
	})

	_, err := temp.Parse(FunctionsTemplate)
	if err != nil {
		return nil, err
	}

	var buf bytes.Buffer

	context := struct {
		Name  string
		Funcs []string
	}{
		Name: s.Title,
		Funcs: []string{
			"Create", "Update", "Delete", "ById",
			"Some", "One",
		},
	}

	if err := temp.ExecuteTemplate(&buf, "functions.tmpl", context); err != nil {
		return nil, err
	}

	return writers.Clear(buf)
}
<commit_msg>generator: instead of only object items, do limit for configs<commit_after>// Package models creates the models for the modules
package models

import (
	"bytes"
	"fmt"
	"strings"
	"text/template"

	"github.com/cihangir/gene/generators/constants"
	"github.com/cihangir/gene/generators/constructors"
	"github.com/cihangir/gene/generators/sql/statements"
	"github.com/cihangir/gene/generators/validators"
	"github.com/cihangir/gene/writers"
	"github.com/cihangir/schema"
	"github.com/cihangir/stringext"
)

// Generate creates the models and write them to the required paths
func Generate(rootPath string, s *schema.Schema) error {

	for _, def := range s.Definitions {
		// create models only for objects
		if def.Type != nil {
			if t, ok := def.Type.(string); ok {
				if t == "config" {
					continue
				}
			}
		}

		moduleName := strings.ToLower(def.Title)

		modelFilePath := fmt.Sprintf(
			"%smodels/%s.go",
			rootPath,
			moduleName,
		)

		f, err := GenerateModel(def)
		if err != nil {
			return err
		}

		if err := writers.WriteFormattedFile(modelFilePath, f); err != nil {
			return err
		}
	}
	return nil
}

// GenerateStatements generates the basic CRUD statements for the models
func GenerateStatements(rootPath string, s *schema.Schema) error {

	for _, def := range s.Definitions {
		// create models only for objects
		if def.Type != nil {
			if t, ok := def.Type.(string); ok {
				if t != "object" {
					continue
				}
			}
		}

		moduleName := strings.ToLower(def.Title)

		modelFilePath := fmt.Sprintf(
			"%smodels/%s_statements.go",
			rootPath,
			moduleName,
		)

		f, err := GenerateModelStatements(def)
		if err != nil {
			return err
		}

		if err := writers.WriteFormattedFile(modelFilePath, f); err != nil {
			return err
		}
	}
	return nil
}

// GenerateModel generates the model itself
func GenerateModel(s *schema.Schema) ([]byte, error) {
	var buf bytes.Buffer

	packageLine, err := GeneratePackage(s)
	if err != nil {
		return nil, err
	}

	consts, err := constants.Generate(s)
	if err != nil {
		return nil, err
	}

	schema, err := GenerateSchema(s)
	if err != nil {
		return nil, err
	}

	constructor, err := constructors.Generate(s)
	if err != nil {
		return nil, err
	}

	validators, err := validators.Generate(s)
	if err != nil {
		return nil, err
	}

	buf.Write(packageLine)
	buf.Write(consts)
	buf.Write(schema)
	buf.Write(constructor)
	if validators != nil {
		buf.Write(validators)
	}

	return writers.Clear(buf)
}

// GenerateModelStatements generates the CRUD statements for the model struct
func GenerateModelStatements(s *schema.Schema) ([]byte, error) {
	var buf bytes.Buffer

	packageLine, err := GeneratePackage(s)
	if err != nil {
		return nil, err
	}

	createStatements, err := statements.GenerateCreate(s)
	if err != nil {
		return nil, err
	}

	updateStatements, err := statements.GenerateUpdate(s)
	if err != nil {
		return nil, err
	}

	deleteStatements, err := statements.GenerateDelete(s)
	if err != nil {
		return nil, err
	}

	selectStatements, err := statements.GenerateSelect(s)
	if err != nil {
		return nil, err
	}

	tableName, err := statements.GenerateTableName(s)
	if err != nil {
		return nil, err
	}

	buf.Write(packageLine)
	buf.Write(createStatements)
	buf.Write(updateStatements)
	buf.Write(deleteStatements)
	buf.Write(selectStatements)
	buf.Write(tableName)

	return writers.Clear(buf)
}

// GeneratePackage generates the imports according to the schema.
// TODO remove this function
func GeneratePackage(s *schema.Schema) ([]byte, error) {
	temp := template.New("package.tmpl")
	_, err := temp.Parse(PackageTemplate)
	if err != nil {
		return nil, err
	}

	var buf bytes.Buffer

	// name := strings.ToLower(strings.Split(s.Title, " ")[0])
	name := "models"
	err = temp.ExecuteTemplate(&buf, "package.tmpl", name)
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

// GenerateSchema generates the schema.
func GenerateSchema(s *schema.Schema) ([]byte, error) {

	temp := template.New("schema.tmpl")
	temp.Funcs(schema.Helpers)

	_, err := temp.Parse(StructTemplate)
	if err != nil {
		return nil, err
	}

	var buf bytes.Buffer

	context := struct {
		Name       string
		Definition *schema.Schema
	}{
		Name:       s.Title,
		Definition: s,
	}

	err = temp.ExecuteTemplate(&buf, "schema.tmpl", context)
	if err != nil {
		return nil, err
	}

	return writers.Clear(buf)
}

// GenerateFunctions generates the functions according to the schema.
func GenerateFunctions(s *schema.Schema) ([]byte, error) {

	temp := template.New("functions.tmpl")
	temp.Funcs(template.FuncMap{
		"Pointerize": stringext.Pointerize,
	})

	_, err := temp.Parse(FunctionsTemplate)
	if err != nil {
		return nil, err
	}

	var buf bytes.Buffer

	context := struct {
		Name  string
		Funcs []string
	}{
		Name: s.Title,
		Funcs: []string{
			"Create", "Update", "Delete", "ById",
			"Some", "One",
		},
	}

	if err := temp.ExecuteTemplate(&buf, "functions.tmpl", context); err != nil {
		return nil, err
	}

	return writers.Clear(buf)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"log"
	"os"
	"os/exec"
	"strconv"

	"github.com/fjukstad/kvik/utils"
	zmq "github.com/pebbe/zmq4"

	"flag"
	"fmt"
)

// workerPort holds the next available port to start up a worker.
var workerPort int

func worker(b []byte, filename string) error {
	err := storeScript(b, filename)
	if err != nil {
		return err
	}
	err = startWorker(filename)
	if err != nil {
		return err
	}
	return nil
}

func startWorker(filename string) error {
	cmd := exec.Command("python", "worker/worker.py", strconv.Itoa(workerPort),
		filename)
	err := cmd.Start()
	if err != nil {
		return err
	}
	return nil
}

// storeScript stores the r script from *b* into filename
func storeScript(b []byte, filename string) error {
	err := utils.CreateDirectories(filename)
	if err != nil {
		log.Panic(err)
	}

	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()
	_, err = file.Write(b)
	if err != nil {
		return err
	}
	return nil
}

func killAllWorkers() error {
	cmd := exec.Command("killall", "Python")
	return cmd.Start()
}

func errorResponse(msg string) utils.ComputeResponse {
	return utils.ComputeResponse{-1, msg}
}
func successResponse(msg string) utils.ComputeResponse {
	return utils.ComputeResponse{1, msg}
}

// Compute is the component of Kvik responsible for starting/stopping workers
// that perform the statistical analyses. It exposes a http rest interface to
// the outside world.
func main() {

	var ip = flag.String("ip", "*", "ip to run on")
	var port = flag.String("port", ":8888", "port to run on")

	flag.Parse()

	responder, _ := zmq.NewSocket(zmq.REP)
	defer responder.Close()

	addr := "tcp://" + *ip + *port

	err := responder.Bind(addr)
	if err != nil {
		log.Panic(err)
	}

	// ID to identify client
	id := 0

	workerPort = 5000

	for {
		msg, err := responder.Recv(0)

		if err != nil {
			log.Println(err)
			continue
		}
		fmt.Println("Received ", msg)

		cmd := new(utils.Command)
		err = json.Unmarshal([]byte(msg), cmd)
		if err != nil {
			log.Panic(err)
		}

		var Response utils.ComputeResponse

		switch {
		case cmd.Command == "stop":
			break

		case cmd.Command == "startWorker":
			path := "scripts/" + strconv.Itoa(id) + "/script.r"
			err = worker(cmd.File, path)
			if err != nil {
				log.Println(err)
				Response = errorResponse("Could not start worker")
			} else {
				Response = successResponse(strconv.Itoa(workerPort))
			}

		case cmd.Command == "killall":
			err = killAllWorkers()
			if err != nil {
				Response = errorResponse("Could not kill all workers")
			} else {
				Response = successResponse("Killed off all workers")
			}
		}

		resp, _ := json.Marshal(Response)
		responder.Send(string(resp), 0)
		id += 1
		workerPort += 1
	}
}
<commit_msg>now starts a webserver to host images<commit_after>package main

import (
	"encoding/json"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path"
	"strconv"

	"github.com/fjukstad/kvik/utils"
	zmq "github.com/pebbe/zmq4"

	"flag"
	"fmt"
)

// workerPort holds the next available port to start up a worker.
var workerPort int

func worker(b []byte, filename string) error {
	err := storeScript(b, filename)
	if err != nil {
		return err
	}
	err = startWorker(filename)
	if err != nil {
		return err
	}
	workerPort += 1

	path := path.Dir(filename) + "/images"
	err = startWebServer(path)

	return err
}

func startWebServer(path string) error {
	p := strconv.Itoa(workerPort)
	port := ":" + p
	workerPort += 1
	fmt.Println("Starting web server at", port, path)
	err := http.ListenAndServe(port, http.FileServer(http.Dir(path)))
	return err
}

func startWorker(filename string) error {
	cmd := exec.Command("python", "worker/worker.py", strconv.Itoa(workerPort),
		filename)
	err := cmd.Start()
	if err != nil {
		return err
	}
	return nil
}

// storeScript stores the r script from *b* into filename
func storeScript(b []byte, filename string) error {
	err := utils.CreateDirectories(filename)
	if err != nil {
		log.Panic(err)
	}

	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()
	_, err = file.Write(b)
	if err != nil {
		return err
	}
	return nil
}

func killAllWorkers() error {
	cmd := exec.Command("killall", "Python")
	return cmd.Start()
}

func errorResponse(msg string) utils.ComputeResponse {
	return utils.ComputeResponse{-1, msg}
}
func successResponse(msg string) utils.ComputeResponse {
	return utils.ComputeResponse{1, msg}
}

// Compute is the component of Kvik responsible for starting/stopping workers
// that perform the statistical analyses. It exposes a http rest interface to
// the outside world.
func main() {

	var ip = flag.String("ip", "*", "ip to run on")
	var port = flag.String("port", ":8888", "port to run on")

	flag.Parse()

	responder, _ := zmq.NewSocket(zmq.REP)
	defer responder.Close()

	addr := "tcp://" + *ip + *port

	err := responder.Bind(addr)
	if err != nil {
		log.Panic(err)
	}

	// ID to identify client
	id := 0

	workerPort = 5000

	for {
		msg, err := responder.Recv(0)

		if err != nil {
			log.Println(err)
			continue
		}
		fmt.Println("Received ", msg)

		cmd := new(utils.Command)
		err = json.Unmarshal([]byte(msg), cmd)
		if err != nil {
			log.Panic(err)
		}

		var Response utils.ComputeResponse

		switch {
		case cmd.Command == "stop":
			break

		case cmd.Command == "startWorker":
			path := "scripts/" + strconv.Itoa(id) + "/script.r"
			err = worker(cmd.File, path)
			if err != nil {
				log.Println(err)
				Response = errorResponse("Could not start worker")
			} else {
				Response = successResponse(strconv.Itoa(workerPort))
			}

		case cmd.Command == "killall":
			err = killAllWorkers()
			if err != nil {
				Response = errorResponse("Could not kill all workers")
			} else {
				Response = successResponse("Killed off all workers")
			}
		}

		resp, _ := json.Marshal(Response)
		responder.Send(string(resp), 0)
		id += 1
	}
}
<|endoftext|>