<commit_before>package selenium

import (
	"flag"
	"fmt"
	"net/http"
	"strings"
	"testing"
)

var grid = flag.Bool("test.grid", false, "skip tests that fail on Selenium Grid")
var executor = flag.String("test.executor", defaultExecutor, "executor URL")
var browserName = flag.String("test.browserName", "firefox", "browser to run tests on")

func init() {
	flag.BoolVar(&Trace, "trace", false, "trace HTTP requests and responses")
	flag.Parse()

	caps["browserName"] = *browserName
}

var caps Capabilities = make(Capabilities)

var runOnSauce *bool = flag.Bool("saucelabs", false, "run on sauce")

func newRemote(testName string, t *testing.T) (wd WebDriver) {
	var err error
	if wd, err = NewRemote(caps, *executor); err != nil {
		t.Fatalf("can't start session for test %s: %s", testName, err)
	}
	return wd
}

func TestStatus(t *testing.T) {
	if *grid {
		t.Skip()
	}
	t.Parallel()
	wd := newRemote("TestStatus", t)
	defer wd.Quit()

	status, err := wd.Status()
	if err != nil {
		t.Fatal(err)
	}

	if status.OS.Name == "" {
		t.Fatal("No OS")
	}
}

func TestNewSession(t *testing.T) {
	t.Parallel()
	if *runOnSauce {
		return
	}
	wd := &remoteWebDriver{capabilities: caps, executor: defaultExecutor}
	sid, err := wd.NewSession()
	defer wd.Quit()

	if err != nil {
		t.Fatalf("error in new session - %s", err)
	}

	if sid == "" {
		t.Fatal("Empty session id")
	}

	if wd.id != sid {
		t.Fatal("Session id mismatch")
	}
}

func TestCapabilities(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestCapabilities", t)
	defer wd.Quit()

	c, err := wd.Capabilities()
	if err != nil {
		t.Fatal(err)
	}

	if c["browserName"] != caps["browserName"] {
		t.Fatalf("bad browser name - %s", c["browserName"])
	}
}

func TestSetAsyncScriptTimeout(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestSetAsyncScriptTimeout", t).T(t)
	defer wd.Quit()

	wd.SetAsyncScriptTimeout(200)
}

func TestSetImplicitWaitTimeout(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestSetImplicitWaitTimeout", t).T(t)
	defer wd.Quit()

	wd.SetImplicitWaitTimeout(200)
}

func TestCurrentWindowHandle(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestCurrentWindowHandle", t).T(t)
	defer wd.Quit()

	handle := wd.CurrentWindowHandle()

	if handle == "" {
		t.Fatal("Empty handle")
	}
}

func TestWindowHandles(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestWindowHandles", t).T(t)
	defer wd.Quit()

	handles := wd.CurrentWindowHandle()

	if handles == "" {
		t.Fatal("No handles")
	}
}

func TestGet(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestGet", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)

	newURL := wd.CurrentURL()

	if newURL != serverURL {
		t.Fatalf("%s != %s", newURL, serverURL)
	}
}

func TestNavigation(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestNavigation", t).T(t)
	defer wd.Quit()

	url1 := serverURL
	wd.Get(url1)

	url2 := serverURL + "other"
	wd.Get(url2)

	wd.Back()
	url := wd.CurrentURL()
	if url != url1 {
		t.Fatalf("back got me to %s (expected %s)", url, url1)
	}
	wd.Forward()
	url = wd.CurrentURL()
	if url != url2 {
		t.Fatalf("forward got me to %s (expected %s)", url, url2)
	}

	wd.Refresh()
	url = wd.CurrentURL()
	if url != url2 {
		t.Fatalf("refresh got me to %s (expected %s)", url, url2)
	}
}

func TestTitle(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestTitle", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	title := wd.Title()
	expectedTitle := "Go Selenium Test Suite"
	if title != expectedTitle {
		t.Fatal("Bad title %s, should be %s", title, expectedTitle)
	}
}

func TestPageSource(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestPageSource", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	source := wd.PageSource()
	if !strings.Contains(source, "The home page.") {
		t.Fatalf("Bad source\n%s", source)
	}
}

func TestFindElement(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestFindElement", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	elem := wd.FindElement(ByName, "q")

	we, ok := elem.(*webElementT).e.(*remoteWE)
	if !ok {
		t.Fatal("Can't convert to *remoteWE")
	}
	if we.id == "" {
		t.Fatal("Empty element")
	}
	if we.parent != wd.(*webDriverT).d.(*remoteWebDriver) {
		t.Fatal("Bad parent")
	}
}

func TestFindElements(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestFindElements", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	elems := wd.FindElements(ByName, "q")
	if len(elems) != 1 {
		t.Fatal("Wrong number of elements %d (should be 1)", len(elems))
	}

	we, ok := elems[0].(*webElementT).e.(*remoteWE)
	if !ok {
		t.Fatal("Can't convert to *remoteWE")
	}

	if we.id == "" {
		t.Fatal("Empty element")
	}
	if we.parent != wd.(*webDriverT).d.(*remoteWebDriver) {
		t.Fatal("Bad parent")
	}
}

func TestSendKeys(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestSendKeys", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	input := wd.FindElement(ByName, "q")
	input.SendKeys("golang\n")

	source := wd.PageSource()
	if !strings.Contains(source, "The Go Programming Language") {
		t.Fatal("Can't find Go")
	}
	if !strings.Contains(source, "golang") {
		t.Fatal("Can't find search query in source")
	}
}

func TestClick(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestClick", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	input := wd.FindElement(ByName, "q")
	input.SendKeys("golang")

	button := wd.FindElement(ById, "submit")
	button.Click()

	if !strings.Contains(wd.PageSource(), "The Go Programming Language") {
		t.Fatal("Can't find Go")
	}
}

func TestGetCookies(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestGetCookies", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	cookies := wd.GetCookies()

	if len(cookies) == 0 {
		t.Fatal("No cookies")
	}

	if cookies[0].Name == "" {
		t.Fatal("Empty cookie")
	}
}

func TestAddCookie(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestAddCookie", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	cookie := &Cookie{Name: "the nameless cookie", Value: "I have nothing"}
	wd.AddCookie(cookie)

	cookies := wd.GetCookies()
	for _, c := range cookies {
		if (c.Name == cookie.Name) && (c.Value == cookie.Value) {
			return
		}
	}

	t.Fatal("Can't find new cookie")
}

func TestDeleteCookie(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestDeleteCookie", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	cookies := wd.GetCookies()
	if len(cookies) == 0 {
		t.Fatal("No cookies")
	}
	wd.DeleteCookie(cookies[0].Name)
	newCookies := wd.GetCookies()
	if len(newCookies) != len(cookies)-1 {
		t.Fatal("Cookie not deleted")
	}

	for _, c := range newCookies {
		if c.Name == cookies[0].Name {
			t.Fatal("Deleted cookie found")
		}
	}
}

func TestLocation(t *testing.T) {
	wd := newRemote("TestLocation", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	button := wd.FindElement(ById, "submit")

	loc := button.Location()

	if (loc.X == 0) || (loc.Y == 0) {
		t.Fatalf("Bad location: %v\n", loc)
	}
}

func TestLocationInView(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestLocationInView", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	button := wd.FindElement(ById, "submit")

	loc := button.LocationInView()

	if (loc.X == 0) || (loc.Y == 0) {
		t.Fatalf("Bad location: %v\n", loc)
	}
}

func TestSize(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestSize", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	button := wd.FindElement(ById, "submit")

	size := button.Size()

	if (size.Width == 0) || (size.Height == 0) {
		t.Fatalf("Bad size: %v\n", size)
	}
}

func TestExecuteScript(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestExecuteScript", t).T(t)
	defer wd.Quit()

	script := "return arguments[0] + arguments[1]"
	args := []interface{}{1, 2}
	reply := wd.ExecuteScript(script, args)

	result, ok := reply.(float64)
	if !ok {
		t.Fatal("Not an int reply")
	}

	if result != 3 {
		t.Fatal("Bad result %d (expected 3)", result)
	}
}

func TestScreenshot(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestScreenshot", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	data := wd.Screenshot()

	if len(data) == 0 {
		t.Fatal("Empty reply")
	}
}

func TestIsSelected(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestIsSelected", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	elem := wd.FindElement(ById, "chuk")

	selected := elem.IsSelected()
	if selected {
		t.Fatal("Already selected")
	}

	elem.Click()
	selected = elem.IsSelected()
	if !selected {
		t.Fatal("Not selected")
	}
}

// Test server

var homePage = `

	<title>Go Selenium Test Suite</title>

	The home page.

	<form action="/search">
		<input name="q" /> <input type="submit" id="submit"/>

		<input id="chuk" type="checkbox" /> A checkbox.
	</form>

`

var otherPage = `

	<title>Go Selenium Test Suite - Other Page</title>

	The other page.

`

var searchPage = `

	<title>Go Selenium Test Suite - Search Page</title>

	You searched for "%s". I'll pretend I've found:

	"The Go Programming Language"

`

var pages = map[string]string{
	"/":       homePage,
	"/other":  otherPage,
	"/search": searchPage,
}

func handler(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path
	page, ok := pages[path]
	if !ok {
		http.NotFound(w, r)
		return
	}

	if path == "/search" {
		r.ParseForm()
		page = fmt.Sprintf(page, r.Form["q"][0])
	}
	// Some cookies for the tests
	for i := 0; i < 3; i++ {
		name := fmt.Sprintf("cookie-%d", i)
		value := fmt.Sprintf("value-%d", i)
		http.SetCookie(w, &http.Cookie{Name: name, Value: value})
	}
	fmt.Fprintf(w, page)
}

var serverPort = ":4793"
var serverURL = "http://localhost" + serverPort + "/"

func init() {
	go func() {
		http.HandleFunc("/", handler)
		http.ListenAndServe(serverPort, nil)
	}()
}
<commit_msg>Add tests for finding sub-element(s) of a WebElementT<commit_after>package selenium

import (
	"flag"
	"fmt"
	"net/http"
	"strings"
	"testing"
)

var grid = flag.Bool("test.grid", false, "skip tests that fail on Selenium Grid")
var executor = flag.String("test.executor", defaultExecutor, "executor URL")
var browserName = flag.String("test.browserName", "firefox", "browser to run tests on")

func init() {
	flag.BoolVar(&Trace, "trace", false, "trace HTTP requests and responses")
	flag.Parse()

	caps["browserName"] = *browserName
}

var caps Capabilities = make(Capabilities)

var runOnSauce *bool = flag.Bool("saucelabs", false, "run on sauce")

func newRemote(testName string, t *testing.T) (wd WebDriver) {
	var err error
	if wd, err = NewRemote(caps, *executor); err != nil {
		t.Fatalf("can't start session for test %s: %s", testName, err)
	}
	return wd
}

func TestStatus(t *testing.T) {
	if *grid {
		t.Skip()
	}
	t.Parallel()
	wd := newRemote("TestStatus", t)
	defer wd.Quit()

	status, err := wd.Status()
	if err != nil {
		t.Fatal(err)
	}

	if status.OS.Name == "" {
		t.Fatal("No OS")
	}
}

func TestNewSession(t *testing.T) {
	t.Parallel()
	if *runOnSauce {
		return
	}
	wd := &remoteWebDriver{capabilities: caps, executor: defaultExecutor}
	sid, err := wd.NewSession()
	defer wd.Quit()

	if err != nil {
		t.Fatalf("error in new session - %s", err)
	}

	if sid == "" {
		t.Fatal("Empty session id")
	}

	if wd.id != sid {
		t.Fatal("Session id mismatch")
	}
}

func TestCapabilities(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestCapabilities", t)
	defer wd.Quit()

	c, err := wd.Capabilities()
	if err != nil {
		t.Fatal(err)
	}

	if c["browserName"] != caps["browserName"] {
		t.Fatalf("bad browser name - %s", c["browserName"])
	}
}

func TestSetAsyncScriptTimeout(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestSetAsyncScriptTimeout", t).T(t)
	defer wd.Quit()

	wd.SetAsyncScriptTimeout(200)
}

func TestSetImplicitWaitTimeout(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestSetImplicitWaitTimeout", t).T(t)
	defer wd.Quit()

	wd.SetImplicitWaitTimeout(200)
}

func TestCurrentWindowHandle(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestCurrentWindowHandle", t).T(t)
	defer wd.Quit()

	handle := wd.CurrentWindowHandle()

	if handle == "" {
		t.Fatal("Empty handle")
	}
}

func TestWindowHandles(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestWindowHandles", t).T(t)
	defer wd.Quit()

	handles := wd.CurrentWindowHandle()

	if handles == "" {
		t.Fatal("No handles")
	}
}

func TestGet(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestGet", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)

	newURL := wd.CurrentURL()

	if newURL != serverURL {
		t.Fatalf("%s != %s", newURL, serverURL)
	}
}

func TestNavigation(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestNavigation", t).T(t)
	defer wd.Quit()

	url1 := serverURL
	wd.Get(url1)

	url2 := serverURL + "other"
	wd.Get(url2)

	wd.Back()
	url := wd.CurrentURL()
	if url != url1 {
		t.Fatalf("back got me to %s (expected %s)", url, url1)
	}
	wd.Forward()
	url = wd.CurrentURL()
	if url != url2 {
		t.Fatalf("forward got me to %s (expected %s)", url, url2)
	}

	wd.Refresh()
	url = wd.CurrentURL()
	if url != url2 {
		t.Fatalf("refresh got me to %s (expected %s)", url, url2)
	}
}

func TestTitle(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestTitle", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	title := wd.Title()
	expectedTitle := "Go Selenium Test Suite"
	if title != expectedTitle {
		t.Fatal("Bad title %s, should be %s", title, expectedTitle)
	}
}

func TestPageSource(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestPageSource", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	source := wd.PageSource()
	if !strings.Contains(source, "The home page.") {
		t.Fatalf("Bad source\n%s", source)
	}
}

type elementFinder interface {
	FindElement(by, value string) WebElementT
	FindElements(by, value string) []WebElementT
}

func TestFindElement(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestFindElement", t).T(t)
	defer wd.Quit()
	wd.Get(serverURL)
	testFindElement(t, wd, ByCSSSelector, "ol.list li", "foo")
}

func TestFindChildElement(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestFindChildElement", t).T(t)
	defer wd.Quit()
	wd.Get(serverURL)
	testFindElement(t, wd.FindElement(ByTagName, "body"), ByCSSSelector, "ol.list li", "foo")
}

func testFindElement(t *testing.T, ef elementFinder, by, value string, txt string) {
	elem := ef.FindElement(by, value)
	if want, got := txt, elem.Text(); want != got {
		t.Errorf("Elem for %q %q: want text %q, got %q", by, value, want, got)
	}
}

func TestFindElements(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestFindElements", t).T(t)
	defer wd.Quit()
	wd.Get(serverURL)
	testFindElements(t, wd, ByCSSSelector, "ol.list li", []string{"foo", "bar"})
}

func TestFindChildElements(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestFindChildElements", t).T(t)
	defer wd.Quit()
	wd.Get(serverURL)
	testFindElements(t, wd.FindElement(ByCSSSelector, "ol.list"), ByCSSSelector, "li", []string{"foo", "bar"})
}

func testFindElements(t *testing.T, ef elementFinder, by, value string, elemsTxt []string) {
	elems := ef.FindElements(by, value)
	if len(elems) != len(elemsTxt) {
		t.Fatal("Wrong number of elements %d (should be %d)", len(elems), len(elemsTxt))
	}
	t.Logf("Found %d elements for %q %q", len(elems), by, value)
	for i, txt := range elemsTxt {
		elem := elems[i]
		if want, got := txt, elem.Text(); want != got {
			t.Errorf("Elem %d for %q %q: want text %q, got %q", i, by, value, want, got)
		}
	}
}

func TestSendKeys(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestSendKeys", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	input := wd.FindElement(ByName, "q")
	input.SendKeys("golang\n")

	source := wd.PageSource()
	if !strings.Contains(source, "The Go Programming Language") {
		t.Fatal("Can't find Go")
	}
	if !strings.Contains(source, "golang") {
		t.Fatal("Can't find search query in source")
	}
}

func TestClick(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestClick", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	input := wd.FindElement(ByName, "q")
	input.SendKeys("golang")

	button := wd.FindElement(ById, "submit")
	button.Click()

	if !strings.Contains(wd.PageSource(), "The Go Programming Language") {
		t.Fatal("Can't find Go")
	}
}

func TestGetCookies(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestGetCookies", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	cookies := wd.GetCookies()

	if len(cookies) == 0 {
		t.Fatal("No cookies")
	}

	if cookies[0].Name == "" {
		t.Fatal("Empty cookie")
	}
}

func TestAddCookie(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestAddCookie", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	cookie := &Cookie{Name: "the nameless cookie", Value: "I have nothing"}
	wd.AddCookie(cookie)

	cookies := wd.GetCookies()
	for _, c := range cookies {
		if (c.Name == cookie.Name) && (c.Value == cookie.Value) {
			return
		}
	}

	t.Fatal("Can't find new cookie")
}

func TestDeleteCookie(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestDeleteCookie", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	cookies := wd.GetCookies()
	if len(cookies) == 0 {
		t.Fatal("No cookies")
	}
	wd.DeleteCookie(cookies[0].Name)
	newCookies := wd.GetCookies()
	if len(newCookies) != len(cookies)-1 {
		t.Fatal("Cookie not deleted")
	}

	for _, c := range newCookies {
		if c.Name == cookies[0].Name {
			t.Fatal("Deleted cookie found")
		}
	}
}

func TestLocation(t *testing.T) {
	wd := newRemote("TestLocation", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	button := wd.FindElement(ById, "submit")

	loc := button.Location()

	if (loc.X == 0) || (loc.Y == 0) {
		t.Fatalf("Bad location: %v\n", loc)
	}
}

func TestLocationInView(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestLocationInView", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	button := wd.FindElement(ById, "submit")

	loc := button.LocationInView()

	if (loc.X == 0) || (loc.Y == 0) {
		t.Fatalf("Bad location: %v\n", loc)
	}
}

func TestSize(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestSize", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	button := wd.FindElement(ById, "submit")

	size := button.Size()

	if (size.Width == 0) || (size.Height == 0) {
		t.Fatalf("Bad size: %v\n", size)
	}
}

func TestExecuteScript(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestExecuteScript", t).T(t)
	defer wd.Quit()

	script := "return arguments[0] + arguments[1]"
	args := []interface{}{1, 2}
	reply := wd.ExecuteScript(script, args)

	result, ok := reply.(float64)
	if !ok {
		t.Fatal("Not an int reply")
	}

	if result != 3 {
		t.Fatal("Bad result %d (expected 3)", result)
	}
}

func TestScreenshot(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestScreenshot", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	data := wd.Screenshot()

	if len(data) == 0 {
		t.Fatal("Empty reply")
	}
}

func TestIsSelected(t *testing.T) {
	t.Parallel()
	wd := newRemote("TestIsSelected", t).T(t)
	defer wd.Quit()

	wd.Get(serverURL)
	elem := wd.FindElement(ById, "chuk")

	selected := elem.IsSelected()
	if selected {
		t.Fatal("Already selected")
	}

	elem.Click()
	selected = elem.IsSelected()
	if !selected {
		t.Fatal("Not selected")
	}
}

// Test server

var homePage = `

	<title>Go Selenium Test Suite</title>

	The home page.

	<form action="/search">
		<input name="q" /> <input type="submit" id="submit"/>

		<input id="chuk" type="checkbox" /> A checkbox.
	</form>

foo

bar

baz

qux

`

var otherPage = `

	<title>Go Selenium Test Suite - Other Page</title>

	The other page.

`

var searchPage = `

	<title>Go Selenium Test Suite - Search Page</title>

	You searched for "%s". I'll pretend I've found:

	"The Go Programming Language"

`

var pages = map[string]string{
	"/":       homePage,
	"/other":  otherPage,
	"/search": searchPage,
}

func handler(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path
	page, ok := pages[path]
	if !ok {
		http.NotFound(w, r)
		return
	}

	if path == "/search" {
		r.ParseForm()
		page = fmt.Sprintf(page, r.Form["q"][0])
	}
	// Some cookies for the tests
	for i := 0; i < 3; i++ {
		name := fmt.Sprintf("cookie-%d", i)
		value := fmt.Sprintf("value-%d", i)
		http.SetCookie(w, &http.Cookie{Name: name, Value: value})
	}
	fmt.Fprintf(w, page)
}

var serverPort = ":4793"
var serverURL = "http://localhost" + serverPort + "/"

func init() {
	go func() {
		http.HandleFunc("/", handler)
		http.ListenAndServe(serverPort, nil)
	}()
}
<|endoftext|>


<commit_before>package approvals

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"os"
	"reflect"
	"regexp"
	"strings"

	"github.com/approvals/go-approval-tests/reporters"
	"github.com/approvals/go-approval-tests/utils"
)

var (
	defaultReporter            = reporters.NewDiffReporter()
	defaultFrontLoadedReporter = reporters.NewFrontLoadedReporter()
	defaultFolder              = ""
)

// Failable is an interface wrapper around testing.T
type Failable interface {
	Fail()
	Fatal(args ...interface{})
	Fatalf(format string, args ...interface{})
	Name() string
	Log(args ...interface{})
	Logf(format string, args ...interface{})
	Helper()
}

// VerifyWithExtension Example:
//   VerifyWithExtension(t, strings.NewReader("Hello"), ".json")
// Deprecated: Please use Verify with the Options() fluent syntax.
func VerifyWithExtension(t Failable, reader io.Reader, extWithDot string, opts ...verifyOptions) {
	t.Helper()

	var v verifyOptions
	if len(opts) == 0 {
		v = Options()
	} else {
		v = opts[0]
	}

	v = v.WithExtension(extWithDot)
	Verify(t, reader, v)
}

// Verify Example:
//   Verify(t, strings.NewReader("Hello"))
func Verify(t Failable, reader io.Reader, opts ...verifyOptions) {
	t.Helper()

	if len(opts) > 1 {
		panic("Please use fluent syntax for options, see documentation for more information")
	}

	var extWithDot string
	if len(opts) == 0 || opts[0].extWithDot == "" {
		extWithDot = ".txt"
	} else {
		extWithDot = opts[0].extWithDot
	}

	namer := getApprovalName(t)

	if len(opts) > 0 {
		b, err := io.ReadAll(reader)
		if err != nil {
			// TODO: do something
		}

		result := string(b)
		for _, o := range opts {
			for _, sb := range o.scrubbers {
				result = sb(result)
			}
		}

		reader = strings.NewReader(result)
	}

	reporter := getReporter()
	err := namer.compare(namer.getApprovalFile(extWithDot), namer.getReceivedFile(extWithDot), reader)
	if err != nil {
		reporter.Report(namer.getApprovalFile(extWithDot), namer.getReceivedFile(extWithDot))
		t.Log("Failed Approval: received does not match approved.")
		t.Fail()
	} else {
		_ = os.Remove(namer.getReceivedFile(extWithDot))
	}
}

// VerifyString stores the passed string into the received file and confirms
// that it matches the approved local file. On failure, it will launch a reporter.
func VerifyString(t Failable, s string, opts ...verifyOptions) {
	t.Helper()
	for _, o := range opts {
		for _, sb := range o.scrubbers {
			s = sb(s)
		}
	}
	reader := strings.NewReader(s)
	Verify(t, reader)
}

// VerifyXMLStruct Example:
//   VerifyXMLStruct(t, xml)
func VerifyXMLStruct(t Failable, obj interface{}, opts ...verifyOptions) {
	t.Helper()
	xmlContent, err := xml.MarshalIndent(obj, "", "  ")
	if err != nil {
		tip := ""
		if reflect.TypeOf(obj).Name() == "" {
			tip = "when using anonymous types be sure to include\n  XMLName xml.Name `xml:\"Your_Name_Here\"`\n"
		}
		message := fmt.Sprintf("error while pretty printing XML\n%verror:\n  %v\nXML:\n  %v\n", tip, err, obj)
		VerifyWithExtension(t, strings.NewReader(message), ".xml", opts...)
	} else {
		VerifyWithExtension(t, bytes.NewReader(xmlContent), ".xml", opts...)
	}
}

// VerifyXMLBytes Example:
//   VerifyXMLBytes(t, []byte("<Test/>"))
func VerifyXMLBytes(t Failable, bs []byte, opts ...verifyOptions) {
	t.Helper()
	type node struct {
		Attr     []xml.Attr
		XMLName  xml.Name
		Children []node `xml:",any"`
		Text     string `xml:",chardata"`
	}
	x := node{}

	err := xml.Unmarshal(bs, &x)
	if err != nil {
		message := fmt.Sprintf("error while parsing XML\nerror:\n  %s\nXML:\n  %s\n", err, string(bs))
		VerifyWithExtension(t, strings.NewReader(message), ".xml", opts...)
	} else {
		VerifyXMLStruct(t, x, opts...)
	}
}

// VerifyJSONStruct Example:
//   VerifyJSONStruct(t, json)
func VerifyJSONStruct(t Failable, obj interface{}, opts ...verifyOptions) {
	t.Helper()

	jsonb, err := json.MarshalIndent(obj, "", "  ")
	if err != nil {
		message := fmt.Sprintf("error while pretty printing JSON\nerror:\n  %s\nJSON:\n  %s\n", err, obj)
		VerifyWithExtension(t, strings.NewReader(message), ".json", opts...)
	} else {
		VerifyWithExtension(t, bytes.NewReader(jsonb), ".json", opts...)
	}
}

// VerifyJSONBytes Example:
//   VerifyJSONBytes(t, []byte("{ \"Greeting\": \"Hello\" }"))
func VerifyJSONBytes(t Failable, bs []byte, opts ...verifyOptions) {
	t.Helper()
	var obj map[string]interface{}
	err := json.Unmarshal(bs, &obj)
	if err != nil {
		message := fmt.Sprintf("error while parsing JSON\nerror:\n  %s\nJSON:\n  %s\n", err, string(bs))
		VerifyWithExtension(t, strings.NewReader(message), ".json", opts...)
	} else {
		VerifyJSONStruct(t, obj, opts...)
	}
}

// VerifyMap Example:
//   VerifyMap(t, map[string][string] { "dog": "bark" })
func VerifyMap(t Failable, m interface{}, opts ...verifyOptions) {
	t.Helper()
	outputText := utils.PrintMap(m)
	VerifyString(t, outputText, opts...)
}

// VerifyArray Example:
//   VerifyArray(t, []string{"dog", "cat"})
func VerifyArray(t Failable, array interface{}, opts ...verifyOptions) {
	t.Helper()
	outputText := utils.PrintArray(array)
	VerifyString(t, outputText, opts...)
}

// VerifyAll Example:
//   VerifyAll(t, "uppercase", []string("dog", "cat"}, func(x interface{}) string { return strings.ToUpper(x.(string)) })
func VerifyAll(t Failable, header string, collection interface{}, transform func(interface{}) string, opts ...verifyOptions) {
	t.Helper()
	if len(header) != 0 {
		header = fmt.Sprintf("%s\n\n\n", header)
	}

	outputText := header + strings.Join(utils.MapToString(collection, transform), "\n")
	VerifyString(t, outputText, opts...)
}

type reporterCloser struct {
	reporter reporters.Reporter
}

func (s *reporterCloser) Close() error {
	defaultReporter = s.reporter
	return nil
}

type frontLoadedReporterCloser struct {
	reporter reporters.Reporter
}

func (s *frontLoadedReporterCloser) Close() error {
	defaultFrontLoadedReporter = s.reporter
	return nil
}

// UseReporter configures which reporter to use on failure.
// Add at the test or method level to configure your reporter.
//
// The following examples shows how to use a reporter for all of your test cases
// in a package directory through go's setup feature.
//
//
//	func TestMain(m *testing.M) {
//		r := approvals.UseReporter(reporters.NewBeyondCompareReporter())
//		defer r.Close()
//
//		os.Exit(m.Run())
//	}
//
func UseReporter(reporter reporters.Reporter) io.Closer {
	closer := &reporterCloser{
		reporter: defaultReporter,
	}

	defaultReporter = reporter
	return closer
}

// UseFrontLoadedReporter configures reporters ahead of all other reporters to
// handle situations like CI. These reporters usually prevent reporting in
// scenarios that are headless.
func UseFrontLoadedReporter(reporter reporters.Reporter) io.Closer {
	closer := &frontLoadedReporterCloser{
		reporter: defaultFrontLoadedReporter,
	}

	defaultFrontLoadedReporter = reporter
	return closer
}

func getReporter() reporters.Reporter {
	return reporters.NewFirstWorkingReporter(
		defaultFrontLoadedReporter,
		defaultReporter,
	)
}

// UseFolder configures which folder to use to store approval files.
// By default, the approval files will be stored at the same level as the code.
//
// The following examples shows how to use the idiomatic 'testdata' folder
// for all of your test cases in a package directory.
//
//
//	func TestMain(m *testing.M) {
//		approvals.UseFolder("testdata")
//
//		os.Exit(m.Run())
//	}
//
func UseFolder(f string) {
	defaultFolder = f
}

type scrubber func(s string) string

// verifyOptions can be accessed via the approvals.Options() API enabling configuration of scrubbers
type verifyOptions struct {
	scrubbers  []scrubber
	extWithDot string
}

// Options enables providing individual Verify functions with customisations such as scrubbers
func Options() verifyOptions {
	return verifyOptions{}
}

// WithRegexScrubber allows you to 'scrub' dynamic data such as timestamps within your test input
// and replace it with a static placeholder
func (v verifyOptions) WithRegexScrubber(scrubber *regexp.Regexp, replacer string) verifyOptions {
	v.scrubbers = append(v.scrubbers, func(s string) string {
		return scrubber.ReplaceAllString(s, replacer)
	})
	return v
}

// WithExtension overrides the default file extension (.txt) for approval files.
func (v verifyOptions) WithExtension(extension string) verifyOptions {
	v.extWithDot = extension
	return v
}
<commit_msg>Clean up initial implementation of scrubbers on VerifyString()<commit_after>package approvals

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"os"
	"reflect"
	"regexp"
	"strings"

	"github.com/approvals/go-approval-tests/reporters"
	"github.com/approvals/go-approval-tests/utils"
)

var (
	defaultReporter            = reporters.NewDiffReporter()
	defaultFrontLoadedReporter = reporters.NewFrontLoadedReporter()
	defaultFolder              = ""
)

// Failable is an interface wrapper around testing.T
type Failable interface {
	Fail()
	Fatal(args ...interface{})
	Fatalf(format string, args ...interface{})
	Name() string
	Log(args ...interface{})
	Logf(format string, args ...interface{})
	Helper()
}

// VerifyWithExtension Example:
//   VerifyWithExtension(t, strings.NewReader("Hello"), ".json")
// Deprecated: Please use Verify with the Options() fluent syntax.
func VerifyWithExtension(t Failable, reader io.Reader, extWithDot string, opts ...verifyOptions) {
	t.Helper()

	var v verifyOptions
	if len(opts) == 0 {
		v = Options()
	} else {
		v = opts[0]
	}

	v = v.WithExtension(extWithDot)
	Verify(t, reader, v)
}

// Verify Example:
//   Verify(t, strings.NewReader("Hello"))
func Verify(t Failable, reader io.Reader, opts ...verifyOptions) {
	t.Helper()

	if len(opts) > 1 {
		panic("Please use fluent syntax for options, see documentation for more information")
	}

	var extWithDot string
	if len(opts) == 0 || opts[0].extWithDot == "" {
		extWithDot = ".txt"
	} else {
		extWithDot = opts[0].extWithDot
	}

	namer := getApprovalName(t)

	if len(opts) > 0 {
		b, err := io.ReadAll(reader)
		if err != nil {
			// TODO: do something
		}

		result := string(b)
		for _, o := range opts {
			for _, sb := range o.scrubbers {
				result = sb(result)
			}
		}

		reader = strings.NewReader(result)
	}

	reporter := getReporter()
	err := namer.compare(namer.getApprovalFile(extWithDot), namer.getReceivedFile(extWithDot), reader)
	if err != nil {
		reporter.Report(namer.getApprovalFile(extWithDot), namer.getReceivedFile(extWithDot))
		t.Log("Failed Approval: received does not match approved.")
		t.Fail()
	} else {
		_ = os.Remove(namer.getReceivedFile(extWithDot))
	}
}

// VerifyString stores the passed string into the received file and confirms
// that it matches the approved local file. On failure, it will launch a reporter.
func VerifyString(t Failable, s string, opts ...verifyOptions) {
	t.Helper()
	reader := strings.NewReader(s)
	Verify(t, reader, opts...)
}

// VerifyXMLStruct Example:
//   VerifyXMLStruct(t, xml)
func VerifyXMLStruct(t Failable, obj interface{}, opts ...verifyOptions) {
	t.Helper()
	xmlContent, err := xml.MarshalIndent(obj, "", "  ")
	if err != nil {
		tip := ""
		if reflect.TypeOf(obj).Name() == "" {
			tip = "when using anonymous types be sure to include\n  XMLName xml.Name `xml:\"Your_Name_Here\"`\n"
		}
		message := fmt.Sprintf("error while pretty printing XML\n%verror:\n  %v\nXML:\n  %v\n", tip, err, obj)
		VerifyWithExtension(t, strings.NewReader(message), ".xml", opts...)
	} else {
		VerifyWithExtension(t, bytes.NewReader(xmlContent), ".xml", opts...)
	}
}

// VerifyXMLBytes Example:
//   VerifyXMLBytes(t, []byte("<Test/>"))
func VerifyXMLBytes(t Failable, bs []byte, opts ...verifyOptions) {
	t.Helper()
	type node struct {
		Attr     []xml.Attr
		XMLName  xml.Name
		Children []node `xml:",any"`
		Text     string `xml:",chardata"`
	}
	x := node{}

	err := xml.Unmarshal(bs, &x)
	if err != nil {
		message := fmt.Sprintf("error while parsing XML\nerror:\n  %s\nXML:\n  %s\n", err, string(bs))
		VerifyWithExtension(t, strings.NewReader(message), ".xml", opts...)
	} else {
		VerifyXMLStruct(t, x, opts...)
	}
}

// VerifyJSONStruct Example:
//   VerifyJSONStruct(t, json)
func VerifyJSONStruct(t Failable, obj interface{}, opts ...verifyOptions) {
	t.Helper()

	jsonb, err := json.MarshalIndent(obj, "", "  ")
	if err != nil {
		message := fmt.Sprintf("error while pretty printing JSON\nerror:\n  %s\nJSON:\n  %s\n", err, obj)
		VerifyWithExtension(t, strings.NewReader(message), ".json", opts...)
	} else {
		VerifyWithExtension(t, bytes.NewReader(jsonb), ".json", opts...)
	}
}

// VerifyJSONBytes Example:
//   VerifyJSONBytes(t, []byte("{ \"Greeting\": \"Hello\" }"))
func VerifyJSONBytes(t Failable, bs []byte, opts ...verifyOptions) {
	t.Helper()
	var obj map[string]interface{}
	err := json.Unmarshal(bs, &obj)
	if err != nil {
		message := fmt.Sprintf("error while parsing JSON\nerror:\n  %s\nJSON:\n  %s\n", err, string(bs))
		VerifyWithExtension(t, strings.NewReader(message), ".json", opts...)
	} else {
		VerifyJSONStruct(t, obj, opts...)
	}
}

// VerifyMap Example:
//   VerifyMap(t, map[string][string] { "dog": "bark" })
func VerifyMap(t Failable, m interface{}, opts ...verifyOptions) {
	t.Helper()
	outputText := utils.PrintMap(m)
	VerifyString(t, outputText, opts...)
}

// VerifyArray Example:
//   VerifyArray(t, []string{"dog", "cat"})
func VerifyArray(t Failable, array interface{}, opts ...verifyOptions) {
	t.Helper()
	outputText := utils.PrintArray(array)
	VerifyString(t, outputText, opts...)
}

// VerifyAll Example:
//   VerifyAll(t, "uppercase", []string("dog", "cat"}, func(x interface{}) string { return strings.ToUpper(x.(string)) })
func VerifyAll(t Failable, header string, collection interface{}, transform func(interface{}) string, opts ...verifyOptions) {
	t.Helper()
	if len(header) != 0 {
		header = fmt.Sprintf("%s\n\n\n", header)
	}

	outputText := header + strings.Join(utils.MapToString(collection, transform), "\n")
	VerifyString(t, outputText, opts...)
}

type reporterCloser struct {
	reporter reporters.Reporter
}

func (s *reporterCloser) Close() error {
	defaultReporter = s.reporter
	return nil
}

type frontLoadedReporterCloser struct {
	reporter reporters.Reporter
}

func (s *frontLoadedReporterCloser) Close() error {
	defaultFrontLoadedReporter = s.reporter
	return nil
}

// UseReporter configures which reporter to use on failure.
// Add at the test or method level to configure your reporter.
//
// The following examples shows how to use a reporter for all of your test cases
// in a package directory through go's setup feature.
//
//
//	func TestMain(m *testing.M) {
//		r := approvals.UseReporter(reporters.NewBeyondCompareReporter())
//		defer r.Close()
//
//		os.Exit(m.Run())
//	}
//
func UseReporter(reporter reporters.Reporter) io.Closer {
	closer := &reporterCloser{
		reporter: defaultReporter,
	}

	defaultReporter = reporter
	return closer
}

// UseFrontLoadedReporter configures reporters ahead of all other reporters to
// handle situations like CI. These reporters usually prevent reporting in
// scenarios that are headless.
func UseFrontLoadedReporter(reporter reporters.Reporter) io.Closer {
	closer := &frontLoadedReporterCloser{
		reporter: defaultFrontLoadedReporter,
	}

	defaultFrontLoadedReporter = reporter
	return closer
}

func getReporter() reporters.Reporter {
	return reporters.NewFirstWorkingReporter(
		defaultFrontLoadedReporter,
		defaultReporter,
	)
}

// UseFolder configures which folder to use to store approval files.
// By default, the approval files will be stored at the same level as the code.
//
// The following examples shows how to use the idiomatic 'testdata' folder
// for all of your test cases in a package directory.
//
//
//	func TestMain(m *testing.M) {
//		approvals.UseFolder("testdata")
//
//		os.Exit(m.Run())
//	}
//
func UseFolder(f string) {
	defaultFolder = f
}

type scrubber func(s string) string

// verifyOptions can be accessed via the approvals.Options() API enabling configuration of scrubbers
type verifyOptions struct {
	scrubbers  []scrubber
	extWithDot string
}

// Options enables providing individual Verify functions with customisations such as scrubbers
func Options() verifyOptions {
	return verifyOptions{}
}

// WithRegexScrubber allows you to 'scrub' dynamic data such as timestamps within your test input
// and replace it with a static placeholder
func (v verifyOptions) WithRegexScrubber(scrubber *regexp.Regexp, replacer string) verifyOptions {
	v.scrubbers = append(v.scrubbers, func(s string) string {
		return scrubber.ReplaceAllString(s, replacer)
	})
	return v
}

// WithExtension overrides the default file extension (.txt) for approval files.
func (v verifyOptions) WithExtension(extension string) verifyOptions {
	v.extWithDot = extension
	return v
}
<|endoftext|>


<commit_before>package stop

// Done is the variable sent on the Chan to indicate
// that something has stopped.
var Done = struct{}{}

// Chan is a stop channel.
type Chan chan struct{}

// Stopper represents types that implement
// the stop channel pattern.
type Stopper interface {
	Stop() Chan
}

// Stopped returns a Chan that signals immediately. Useful for
// cases when no tear-down work is required and stopping is
// immediate.
func Stopped() Chan {
	c := MakeChan()
	c <- Done
	return c
}

// MakeChan makes a new Chan used to indicate when
// stopping has finished. Sends to Chan will not block.
func MakeChan() Chan {
	return make(chan struct{}, 1)
}

// All stops all Stopper types and returns another Chan
// on which Done will be sent once all things have
// finished stopping.
func All(stoppers ...Stopper) Chan {
	all := MakeChan()
	go func() {
		var allChans []Chan
		for _, stopper := range stoppers {
			allChans = append(allChans, stopper.Stop())
		}
		for _, ch := range allChans {
			<-ch
		}
		all <- Done
	}()
	return all
}
<commit_msg>added comment for Stop<commit_after>package stop

// Done is the variable sent on the Chan to indicate
// that something has stopped.
var Done = struct{}{}

// Chan is a stop channel.
type Chan chan struct{}

// Stopper represents types that implement
// the stop channel pattern.
type Stopper interface {
	// Stop instructs the type to halt operations and
	// returns a Chan on which a stop.Done signal is sent
	// when stopping has completed.
	Stop() Chan
}

// Stopped returns a Chan that signals immediately. Useful for
// cases when no tear-down work is required and stopping is
// immediate.
func Stopped() Chan {
	c := MakeChan()
	c <- Done
	return c
}

// MakeChan makes a new Chan used to indicate when
// stopping has finished. Sends to Chan will not block.
func MakeChan() Chan {
	return make(chan struct{}, 1)
}

// All stops all Stopper types and returns another Chan
// on which Done will be sent once all things have
// finished stopping.
func All(stoppers ...Stopper) Chan {
	all := MakeChan()
	go func() {
		var allChans []Chan
		for _, stopper := range stoppers {
			allChans = append(allChans, stopper.Stop())
		}
		for _, ch := range allChans {
			<-ch
		}
		all <- Done
	}()
	return all
}
<|endoftext|>


<commit_before>// Copyright 2014 The rkt Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//+build linux

package main

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"syscall"

	"github.com/appc/spec/schema"
	"github.com/appc/spec/schema/types"
	"github.com/coreos/go-systemd/util"
	"github.com/coreos/rkt/common"
	"github.com/coreos/rkt/networking"
	stage1commontypes "github.com/coreos/rkt/stage1/common/types"
	stage1initcommon "github.com/coreos/rkt/stage1/init/common"
	"github.com/coreos/rkt/stage1/init/kvm"
	"github.com/hashicorp/errwrap"
)

const journalDir = "/var/log/journal"

// KvmNetworkingToSystemd generates systemd unit files for a pod according to network configuration
func KvmNetworkingToSystemd(p *stage1commontypes.Pod, n *networking.Networking) error {
	podRoot := common.Stage1RootfsPath(p.Root)

	// networking
	netDescriptions := kvm.GetNetworkDescriptions(n)
	if err := kvm.GenerateNetworkInterfaceUnits(filepath.Join(podRoot, stage1initcommon.UnitsDir), netDescriptions); err != nil {
		return errwrap.Wrap(errors.New("failed to transform networking to units"), err)
	}

	return nil
}

func mountSharedVolumes(root string, p *stage1commontypes.Pod, ra *schema.RuntimeApp) error {
	app := ra.App
	appName := ra.Name
	volumes := p.Manifest.Volumes
	vols := make(map[types.ACName]types.Volume)
	for _, v := range volumes {
		vols[v.Name] = v
	}

	sharedVolPath := common.SharedVolumesPath(root)
	if err := os.MkdirAll(sharedVolPath, stage1initcommon.SharedVolPerm); err != nil {
		return errwrap.Wrap(errors.New("could not create shared volumes directory"), err)
	}
	if err := os.Chmod(sharedVolPath, stage1initcommon.SharedVolPerm); err != nil {
		return errwrap.Wrap(fmt.Errorf("could not change permissions of %q", sharedVolPath), err)
	}

	imageManifest := p.Images[appName.String()]
	mounts := stage1initcommon.GenerateMounts(ra, vols, imageManifest)
	for _, m := range mounts {
		vol := vols[m.Volume]

		absRoot, err := filepath.Abs(p.Root) // Absolute path to the pod's rootfs.
		if err != nil {
			return errwrap.Wrap(errors.New("could not get pod's root absolute path"), err)
		}

		absAppRootfs := common.AppRootfsPath(absRoot, appName)
		if err != nil {
			return fmt.Errorf(`could not evaluate absolute path for application rootfs in app: %v`, appName)
		}

		mntPath, err := stage1initcommon.EvaluateSymlinksInsideApp(absAppRootfs, m.Path)
		if err != nil {
			return errwrap.Wrap(fmt.Errorf("could not evaluate path %v", m.Path), err)
		}
		absDestination := filepath.Join(absAppRootfs, mntPath)
		shPath := filepath.Join(sharedVolPath, vol.Name.String())
		if err := stage1initcommon.PrepareMountpoints(shPath, absDestination, &vol, m.DockerImplicit); err != nil {
			return err
		}

		readOnly := stage1initcommon.IsMountReadOnly(vol, app.MountPoints)
		var source string
		switch vol.Kind {
		case "host":
			source = vol.Source
		case "empty":
			source = filepath.Join(common.SharedVolumesPath(root), vol.Name.String())
		default:
			return fmt.Errorf(`invalid volume kind %q. Must be one of "host" or "empty"`, vol.Kind)
		}
		if cleanedSource, err := filepath.EvalSymlinks(source); err != nil {
			return errwrap.Wrap(fmt.Errorf("could not resolve symlink for source: %v", source), err)
		} else if err := ensureDestinationExists(cleanedSource, absDestination); err != nil {
			return errwrap.Wrap(fmt.Errorf("could not create destination mount point: %v", absDestination), err)
		} else if err := doBindMount(cleanedSource, absDestination, readOnly); err != nil {
			return errwrap.Wrap(fmt.Errorf("could not bind mount path %v (s: %v, d: %v)", m.Path, source, absDestination), err)
		}
	}
	return nil
}

func doBindMount(source, destination string, readOnly bool) error {
	if err := syscall.Mount(source, destination, "bind", syscall.MS_BIND, ""); err != nil {
		return err
	}
	if readOnly {
		return syscall.Mount(source, destination, "bind", syscall.MS_REMOUNT|syscall.MS_RDONLY|syscall.MS_BIND, "")
	}
	return nil
}

func ensureDestinationExists(source, destination string) error {
	fileInfo, err := os.Stat(source)
	if err != nil {
		return errwrap.Wrap(fmt.Errorf("could not stat source location: %v", source), err)
	}

	targetPathParent, _ := filepath.Split(destination)
	if err := os.MkdirAll(targetPathParent, stage1initcommon.SharedVolPerm); err != nil {
		return errwrap.Wrap(fmt.Errorf("could not create parent directory: %v", targetPathParent), err)
	}

	if fileInfo.IsDir() {
		if err := os.Mkdir(destination, stage1initcommon.SharedVolPerm); !os.IsExist(err) {
			return err
		}
	} else {
		if file, err := os.OpenFile(destination, os.O_CREATE, stage1initcommon.SharedVolPerm); err != nil {
			return err
		} else {
			file.Close()
		}
	}
	return nil
}

func prepareMountsForApp(s1Root string, p *stage1commontypes.Pod, ra *schema.RuntimeApp) error {
	// bind mount all shared volumes (we don't use mechanism for bind-mounting given by nspawn)
	if err := mountSharedVolumes(s1Root, p, ra); err != nil {
		return errwrap.Wrap(errors.New("failed to prepare mount point"), err)
	}

	return nil
}

func KvmPrepareMounts(s1Root string, p *stage1commontypes.Pod) error {
	for i := range p.Manifest.Apps {
		ra := &p.Manifest.Apps[i]
		if err := prepareMountsForApp(s1Root, p, ra); err != nil {
			return errwrap.Wrap(fmt.Errorf("failed prepare mounts for app %q", ra.Name), err)
		}
	}

	return nil
}

func linkJournal(s1Root, machineID string) error {
	if !util.IsRunningSystemd() {
		return nil
	}

	absS1Root, err := filepath.Abs(s1Root)
	if err != nil {
		return err
	}

	// /var/log/journal doesn't exist on the host, don't do anything
	if _, err := os.Stat(journalDir); os.IsNotExist(err) {
		return nil
	}

	machineJournalDir := filepath.Join(journalDir, machineID)
	podJournalDir := filepath.Join(absS1Root, machineJournalDir)

	hostMachineID, err := util.GetMachineID()
	if err != nil {
		return err
	}

	// unlikely, machine ID is random (== pod UUID)
	if hostMachineID == machineID {
		return fmt.Errorf("host and pod machine IDs are equal (%s)", machineID)
	}

	fi, err := os.Lstat(machineJournalDir)
	switch {
	case os.IsNotExist(err):
		// good, we'll create the symlink
	case err != nil:
		return err
	// unlikely, machine ID is random (== pod UUID)
	default:
		if fi.IsDir() {
			if err := os.Remove(machineJournalDir); err != nil {
				return err
			}
		}

		link, err := os.Readlink(machineJournalDir)
		if err != nil {
			return err
		}

		if link == podJournalDir {
			return nil
		} else {
			if err := os.Remove(machineJournalDir); err != nil {
				return err
			}
		}
	}

	if err := os.Symlink(podJournalDir, machineJournalDir); err != nil {
		return err
	}

	return nil
}
<commit_msg>stage1/kvm: volumes: implement recursive option<commit_after>// Copyright 2014 The rkt Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//+build linux

package main

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"syscall"

	"github.com/appc/spec/schema"
	"github.com/appc/spec/schema/types"
	"github.com/coreos/go-systemd/util"
	"github.com/coreos/rkt/common"
	"github.com/coreos/rkt/networking"
	stage1commontypes "github.com/coreos/rkt/stage1/common/types"
	stage1initcommon "github.com/coreos/rkt/stage1/init/common"
	"github.com/coreos/rkt/stage1/init/kvm"
	"github.com/hashicorp/errwrap"
)

const journalDir = "/var/log/journal"

// KvmNetworkingToSystemd generates systemd unit files for a pod according to network configuration
func KvmNetworkingToSystemd(p *stage1commontypes.Pod, n *networking.Networking) error {
	podRoot := common.Stage1RootfsPath(p.Root)

	// networking
	netDescriptions := kvm.GetNetworkDescriptions(n)
	if err := kvm.GenerateNetworkInterfaceUnits(filepath.Join(podRoot, stage1initcommon.UnitsDir), netDescriptions); err != nil {
		return errwrap.Wrap(errors.New("failed to transform networking to units"), err)
	}

	return nil
}

func mountSharedVolumes(root string, p *stage1commontypes.Pod, ra *schema.RuntimeApp) error {
	app := ra.App
	appName := ra.Name
	volumes := p.Manifest.Volumes
	vols := make(map[types.ACName]types.Volume)
	for _, v := range volumes {
		vols[v.Name] = v
	}

	sharedVolPath := common.SharedVolumesPath(root)
	if err := os.MkdirAll(sharedVolPath, stage1initcommon.SharedVolPerm); err != nil {
		return errwrap.Wrap(errors.New("could not create shared volumes directory"), err)
	}
	if err := os.Chmod(sharedVolPath, stage1initcommon.SharedVolPerm); err != nil {
		return errwrap.Wrap(fmt.Errorf("could not change permissions of %q", sharedVolPath), err)
	}

	imageManifest := p.Images[appName.String()]
	mounts := stage1initcommon.GenerateMounts(ra, vols, imageManifest)
	for _, m := range mounts {
		vol := vols[m.Volume]

		absRoot, err := filepath.Abs(p.Root) // Absolute path to the pod's rootfs.
		if err != nil {
			return errwrap.Wrap(errors.New("could not get pod's root absolute path"), err)
		}

		absAppRootfs := common.AppRootfsPath(absRoot, appName)
		if err != nil {
			return fmt.Errorf(`could not evaluate absolute path for application rootfs in app: %v`, appName)
		}

		mntPath, err := stage1initcommon.EvaluateSymlinksInsideApp(absAppRootfs, m.Path)
		if err != nil {
			return errwrap.Wrap(fmt.Errorf("could not evaluate path %v", m.Path), err)
		}
		absDestination := filepath.Join(absAppRootfs, mntPath)
		shPath := filepath.Join(sharedVolPath, vol.Name.String())
		if err := stage1initcommon.PrepareMountpoints(shPath, absDestination, &vol, m.DockerImplicit); err != nil {
			return err
		}

		readOnly := stage1initcommon.IsMountReadOnly(vol, app.MountPoints)
		var source string
		switch vol.Kind {
		case "host":
			source = vol.Source
		case "empty":
			source = filepath.Join(common.SharedVolumesPath(root), vol.Name.String())
		default:
			return fmt.Errorf(`invalid volume kind %q. Must be one of "host" or "empty"`, vol.Kind)
		}
		if cleanedSource, err := filepath.EvalSymlinks(source); err != nil {
			return errwrap.Wrap(fmt.Errorf("could not resolve symlink for source: %v", source), err)
		} else if err := ensureDestinationExists(cleanedSource, absDestination); err != nil {
			return errwrap.Wrap(fmt.Errorf("could not create destination mount point: %v", absDestination), err)
		} else if err := doBindMount(cleanedSource, absDestination, readOnly, vol.Recursive); err != nil {
			return errwrap.Wrap(fmt.Errorf("could not bind mount path %v (s: %v, d: %v)", m.Path, source, absDestination), err)
		}
	}
	return nil
}

func doBindMount(source, destination string, readOnly bool, recursive *bool) error {
	var flags uintptr = syscall.MS_BIND | syscall.MS_REC
	if readOnly {
		flags |= syscall.MS_RDONLY
	}

	// Enable recursive by default and remove it if explicitly requested
	if recursive != nil && *recursive == false {
		flags ^= syscall.MS_REC
	}
	return syscall.Mount(source, destination, "bind", flags, "")
}

func ensureDestinationExists(source, destination string) error {
	fileInfo, err := os.Stat(source)
	if err != nil {
		return errwrap.Wrap(fmt.Errorf("could not stat source location: %v", source), err)
	}

	targetPathParent, _ := filepath.Split(destination)
	if err := os.MkdirAll(targetPathParent, stage1initcommon.SharedVolPerm); err != nil {
		return errwrap.Wrap(fmt.Errorf("could not create parent directory: %v", targetPathParent), err)
	}

	if fileInfo.IsDir() {
		if err := os.Mkdir(destination, stage1initcommon.SharedVolPerm); !os.IsExist(err) {
			return err
		}
	} else {
		if file, err := os.OpenFile(destination, os.O_CREATE, stage1initcommon.SharedVolPerm); err != nil {
			return err
		} else {
			file.Close()
		}
	}
	return nil
}

func prepareMountsForApp(s1Root string, p *stage1commontypes.Pod, ra *schema.RuntimeApp) error {
	// bind mount all shared volumes (we don't use mechanism for bind-mounting given by nspawn)
	if err := mountSharedVolumes(s1Root, p, ra); err != nil {
		return errwrap.Wrap(errors.New("failed to prepare mount point"), err)
	}

	return nil
}

func KvmPrepareMounts(s1Root string, p *stage1commontypes.Pod) error {
	for i := range p.Manifest.Apps {
		ra := &p.Manifest.Apps[i]
		if err := prepareMountsForApp(s1Root, p, ra); err != nil {
			return errwrap.Wrap(fmt.Errorf("failed prepare mounts for app %q", ra.Name), err)
		}
	}

	return nil
}

func linkJournal(s1Root, machineID string) error {
	if !util.IsRunningSystemd() {
		return nil
	}

	absS1Root, err := filepath.Abs(s1Root)
	if err != nil {
		return err
	}

	// /var/log/journal doesn't exist on the host, don't do anything
	if _, err := os.Stat(journalDir); os.IsNotExist(err) {
		return nil
	}

	machineJournalDir := filepath.Join(journalDir, machineID)
	podJournalDir := filepath.Join(absS1Root, machineJournalDir)

	hostMachineID, err := util.GetMachineID()
	if err != nil {
		return err
	}

	// unlikely, machine ID is random (== pod UUID)
	if hostMachineID == machineID {
		return fmt.Errorf("host and pod machine IDs are equal (%s)", machineID)
	}

	fi, err := os.Lstat(machineJournalDir)
	switch {
	case os.IsNotExist(err):
		// good, we'll create the symlink
	case err != nil:
		return err
	// unlikely, machine ID is random (== pod UUID)
	default:
		if fi.IsDir() {
			if err := os.Remove(machineJournalDir); err != nil {
				return err
			}
		}

		link, err := os.Readlink(machineJournalDir)
		if err != nil {
			return err
		}

		if link == podJournalDir {
			return nil
		} else {
			if err := os.Remove(machineJournalDir); err != nil {
				return err
			}
		}
	}

	if err := os.Symlink(podJournalDir, machineJournalDir); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state

import (
	"github.com/juju/errors"
	"github.com/juju/names"
	"gopkg.in/juju/charm.v5"

	"github.com/juju/juju/process"
)

// TODO(ericsnow) Track juju-level status in the status collection.

// TODO(ericsnow) Rename Register to Add and Unregister to Remove.

// TODO(ericsnow) Distinguish between marking as unregistered and
// removing a proc from state?

// UnitProcesses exposes high-level interaction with workload processes
// for a unit.
type UnitProcesses interface {
	// Register registers a workload process in state. If a process is
	// already registered for the same (unit, proc name, plugin ID)
	// then the request will fail. The unit must also be "alive".
	Register(info process.Info) error
	// SetStatus sets the raw status of a workload process. The process
	// ID is in the format provided by process.Info.ID()
	// ("<proc name>/<plugin ID>"). If the process is not in state then
	// the request will fail.
	SetStatus(id string, status process.PluginStatus) error
	// List builds the list of workload processes registered for
	// the given unit and IDs. If no IDs are provided then all
	// registered processes for the unit are returned. In the case that
	// IDs are provided, any that are not in state are ignored and only
	// the found ones are returned. It is up to the caller to
	// extrapolate the list of missing IDs.
	List(ids ...string) ([]process.Info, error)
	// Unregister marks the identified process as unregistered. If the
	// given ID is not in state then the request will fail.
	Unregister(id string) error
}

// ProcessDefiniitions provides the state functionality related to
// workload process definitions.
type ProcessDefinitions interface {
	// EnsureDefined adds the definitions to state if they aren't there
	// already. If they are there then it verfies that the existing
	// definitions match the provided ones.
	EnsureDefined(definitions ...charm.Process) error
}

// TODO(ericsnow) Use a more generic component registration mechanism?

type newUnitProcessesFunc func(persist Persistence, unit names.UnitTag, charm names.CharmTag) (UnitProcesses, error)

type newProcessDefinitionsFunc func(persist Persistence, charm names.CharmTag) (ProcessDefinitions, error)

var (
	newUnitProcesses      newUnitProcessesFunc
	newProcessDefinitions newProcessDefinitionsFunc
)

// SetProcessesComponent registers the functions that provide the state
// functionality related to workload processes.
func SetProcessesComponent(upFunc newUnitProcessesFunc, pdFunc newProcessDefinitionsFunc) {
	newUnitProcesses = upFunc
	newProcessDefinitions = pdFunc
}

type unitProcesses struct {
	UnitProcesses
	charm *Charm
	st    *State
}

// UnitProcesses exposes interaction with workload processes in state
// for a the given unit. It also ensures that all of the process
// definitions in the charm's metadata are in state.
func (st *State) UnitProcesses(unit *Unit) (UnitProcesses, error) {
	if newUnitProcesses == nil {
		return nil, errors.Errorf("unit processes not supported")
	}

	charm, err := unit.charm()
	if err != nil {
		return nil, errors.Trace(err)
	}
	charmTag := charm.Tag().(names.CharmTag)

	// TODO(ericsnow) Instead call st.defineProcesses when a charm
	// is added or its URL changes?
	if err := st.defineProcesses(charmTag, *charm.Meta()); err != nil {
		return nil, errors.Trace(err)
	}

	persist := st.newPersistence()
	unitProcs, err := newUnitProcesses(persist, unit.UnitTag(), charmTag)
	if err != nil {
		return nil, errors.Trace(err)
	}

	return unitProcs, nil
}

// TODO(ericsnow) DestroyProcess: Mark the proc as Dying.

// defineProcesses adds the workload process definitions from the provided
// charm metadata to state.
func (st *State) defineProcesses(charmTag names.CharmTag, meta charm.Meta) error {
	if newProcessDefinitions == nil {
		return errors.Errorf("process definitions not supported")
	}

	var definitions []charm.Process
	for _, definition := range meta.Processes {
		definitions = append(definitions, definition)
	}

	persist := st.newPersistence()
	pd, err := newProcessDefinitions(persist, charmTag)
	if err != nil {
		return errors.Trace(err)
	}

	if err := pd.EnsureDefined(definitions...); err != nil {
		return errors.Trace(err)
	}
	return nil
}
<commit_msg>Fix some error messages.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state

import (
	"github.com/juju/errors"
	"github.com/juju/names"
	"gopkg.in/juju/charm.v5"

	"github.com/juju/juju/process"
)

// TODO(ericsnow) Track juju-level status in the status collection.

// TODO(ericsnow) Rename Register to Add and Unregister to Remove.

// TODO(ericsnow) Distinguish between marking as unregistered and
// removing a proc from state?

// UnitProcesses exposes high-level interaction with workload processes
// for a unit.
type UnitProcesses interface {
	// Register registers a workload process in state. If a process is
	// already registered for the same (unit, proc name, plugin ID)
	// then the request will fail. The unit must also be "alive".
	Register(info process.Info) error
	// SetStatus sets the raw status of a workload process. The process
	// ID is in the format provided by process.Info.ID()
	// ("<proc name>/<plugin ID>"). If the process is not in state then
	// the request will fail.
	SetStatus(id string, status process.PluginStatus) error
	// List builds the list of workload processes registered for
	// the given unit and IDs. If no IDs are provided then all
	// registered processes for the unit are returned. In the case that
	// IDs are provided, any that are not in state are ignored and only
	// the found ones are returned. It is up to the caller to
	// extrapolate the list of missing IDs.
	List(ids ...string) ([]process.Info, error)
	// Unregister marks the identified process as unregistered. If the
	// given ID is not in state then the request will fail.
	Unregister(id string) error
}

// ProcessDefiniitions provides the state functionality related to
// workload process definitions.
type ProcessDefinitions interface {
	// EnsureDefined adds the definitions to state if they aren't there
	// already. If they are there then it verfies that the existing
	// definitions match the provided ones.
	EnsureDefined(definitions ...charm.Process) error
}

// TODO(ericsnow) Use a more generic component registration mechanism?

type newUnitProcessesFunc func(persist Persistence, unit names.UnitTag, charm names.CharmTag) (UnitProcesses, error)

type newProcessDefinitionsFunc func(persist Persistence, charm names.CharmTag) (ProcessDefinitions, error)

var (
	newUnitProcesses      newUnitProcessesFunc
	newProcessDefinitions newProcessDefinitionsFunc
)

// SetProcessesComponent registers the functions that provide the state
// functionality related to workload processes.
func SetProcessesComponent(upFunc newUnitProcessesFunc, pdFunc newProcessDefinitionsFunc) {
	newUnitProcesses = upFunc
	newProcessDefinitions = pdFunc
}

type unitProcesses struct {
	UnitProcesses
	charm *Charm
	st    *State
}

// UnitProcesses exposes interaction with workload processes in state
// for a the given unit. It also ensures that all of the process
// definitions in the charm's metadata are in state.
func (st *State) UnitProcesses(unit *Unit) (UnitProcesses, error) {
	if newUnitProcesses == nil {
		return nil, errors.Errorf("workload processes not supported")
	}

	charm, err := unit.charm()
	if err != nil {
		return nil, errors.Trace(err)
	}
	charmTag := charm.Tag().(names.CharmTag)

	// TODO(ericsnow) Instead call st.defineProcesses when a charm
	// is added or its URL changes?
	if err := st.defineProcesses(charmTag, *charm.Meta()); err != nil {
		return nil, errors.Trace(err)
	}

	persist := st.newPersistence()
	unitProcs, err := newUnitProcesses(persist, unit.UnitTag(), charmTag)
	if err != nil {
		return nil, errors.Trace(err)
	}

	return unitProcs, nil
}

// TODO(ericsnow) DestroyProcess: Mark the proc as Dying.

// defineProcesses adds the workload process definitions from the provided
// charm metadata to state.
func (st *State) defineProcesses(charmTag names.CharmTag, meta charm.Meta) error {
	if newProcessDefinitions == nil {
		return errors.Errorf("workload processes not supported")
	}

	var definitions []charm.Process
	for _, definition := range meta.Processes {
		definitions = append(definitions, definition)
	}

	persist := st.newPersistence()
	pd, err := newProcessDefinitions(persist, charmTag)
	if err != nil {
		return errors.Trace(err)
	}

	if err := pd.EnsureDefined(definitions...); err != nil {
		return errors.Trace(err)
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// A command that processes all of the blobs in a GCS bucket. For each blob:
//
//  1. The GCS object record is verified for internal consistency. In
//     particular, comeback's expected MD5 and CRC32C are checked against the
//     GCS-reported values.
//
//  2. The object contents are read and the SHA-1 computed and checked against
//     the blob's name and expected SHA-1.
//
//  3. The edges in the directory graph are printed to an output file.
//
// (In "fast mode", only #1 is performed, and the output is a simple list of
// scores, one per line.)
//
// Output is space-separated lines of the form:
//
//     <score> <child score>*
//
// where <score> is a verified score, and the list is a list of scores of
// children for directories (and empty for files).

package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"sync"
	"time"

	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/crypto"
	"github.com/jacobsa/comeback/repr"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/syncutil"
	"golang.org/x/net/context"
)

var cmdScanBlobs = &Command{
	Name: "scan_blobs",
}

var fOutputFile = cmdScanBlobs.Flags.String(
	"output_file",
	"",
	"Path to scan_blobs output file. Will be truncated.")

var fFast = cmdScanBlobs.Flags.Bool(
	"fast",
	false,
	"Only verify object metadata.")

func init() {
	cmdScanBlobs.Run = runScanBlobs // Break flag-related dependency loop.
}

type scoreAndContents struct {
	score    blob.Score
	contents []byte
}

type verifiedScore struct {
	score    blob.Score
	children []blob.Score
}

// List all blob objects in the GCS bucket, and verify their metadata. Write
// their scores into the supplied channel.
func listBlobs(
	ctx context.Context,
	bucket gcs.Bucket,
	scores chan<- blob.Score) (err error) {
	req := &gcs.ListObjectsRequest{
		Prefix: blobObjectNamePrefix,
	}

	// List until we run out.
	for {
		// Fetch the next batch.
		var listing *gcs.Listing
		listing, err = bucket.ListObjects(ctx, req)
		if err != nil {
			err = fmt.Errorf("ListObjects: %v", err)
			return
		}

		// Transform to scores, verifying in the process.
		var batch []blob.Score
		for _, o := range listing.Objects {
			// Skip an object named simply blobObjectNamePrefix, which we allow to
			// exist for gcsfuse compatibility.
			if o.Name == blobObjectNamePrefix {
				continue
			}

			// Parse and verify the record.
			var score blob.Score
			score, err = blob.ParseObjectRecord(o, blobObjectNamePrefix)
			if err != nil {
				err = fmt.Errorf("ParseObjectRecord: %v", err)
				return
			}

			batch = append(batch, score)
		}

		// Feed out each score.
		for _, score := range batch {
			select {
			case scores <- score:

				// Cancelled?
			case <-ctx.Done():
				err = ctx.Err()
				return
			}
		}

		// Are we done?
		if listing.ContinuationToken == "" {
			break
		}

		req.ContinuationToken = listing.ContinuationToken
	}

	return
}

func readAndClose(rc io.ReadCloser) (d []byte, err error) {
	// Read.
	d, err = ioutil.ReadAll(rc)
	if err != nil {
		rc.Close()
		err = fmt.Errorf("ReadAll: %v", err)
		return
	}

	// Close.
	err = rc.Close()
	if err != nil {
		err = fmt.Errorf("Close: %v", err)
		return
	}

	return
}

func readBlob(
	ctx context.Context,
	bucket gcs.Bucket,
	score blob.Score) (blob scoreAndContents, err error) {
	// Fill in the score.
	blob.score = score

	// Obtain a reader.
	req := &gcs.ReadObjectRequest{
		Name: blobObjectNamePrefix + score.Hex(),
	}

	rc, err := bucket.NewReader(ctx, req)
	if err != nil {
		err = fmt.Errorf("NewReader: %v", err)
		return
	}

	// Consume it.
	blob.contents, err = readAndClose(rc)
	if err != nil {
		err = fmt.Errorf("readAndClose: %v", err)
		return
	}

	return
}

func readBlobWithRetryLoop(
	ctx context.Context,
	bucket gcs.Bucket,
	score blob.Score) (blob scoreAndContents, err error) {
	const maxRetries = 5
	for n := 0; n < maxRetries; n++ {
		blob, err = readBlob(ctx, bucket, score)
		if err == nil {
			return
		}

		time.Sleep(500 * time.Millisecond)
	}

	return
}

// Read the contents of blobs specified on the incoming channel. Do not close
// the outgoing channel.
func readBlobs(
	ctx context.Context,
	bucket gcs.Bucket,
	scores <-chan blob.Score,
	blobs chan<- scoreAndContents) (err error) {
	for score := range scores {
		var blob scoreAndContents

		// Read the contents for this score.
		blob, err = readBlobWithRetryLoop(ctx, bucket, score)
		if err != nil {
			err = fmt.Errorf("readBlob: %v", err)
			return
		}

		select {
		case blobs <- blob:
		case <-ctx.Done():
			err = ctx.Err()
			return
		}
	}

	return
}

func scoresEqual(a, b blob.Score) bool {
	for i := 0; i < blob.ScoreLength; i++ {
		if a[i] != b[i] {
			return false
		}
	}

	return true
}

// Parse the blob ciphertext as a directory if appropriate, returning a list of
// children. If not a directory, return the empty list.
func unmarshalBlob(
	crypter crypto.Crypter,
	b []byte) (children []blob.Score, err error) {
	// Attempt to decrypt.
	b, err = crypter.Decrypt(b)
	if err != nil {
		err = fmt.Errorf("Decrypt: %v", err)
		return
	}

	// If this is a file, simply make sure it is in a legal format.
	if !repr.IsDir(b) {
		_, err = repr.UnmarshalFile(b)
		if err != nil {
			err = fmt.Errorf("UnmarshalFile: %v", err)
			return
		}

		return
	}

	// Otherwise, parse it as a directory.
	entries, err := repr.UnmarshalDir(b)
	if err != nil {
		err = fmt.Errorf("UnmarshalDir: %v", err)
		return
	}

	// Pull out the child scores.
	for _, entry := range entries {
		for _, score := range entry.Scores {
			children = append(children, score)
		}
	}

	return
}

// Verify the contents of the incoming blobs. Do not close the outgoing
// channel.
func verifyScores(
	ctx context.Context,
	crypter crypto.Crypter,
	blobs <-chan scoreAndContents,
	results chan<- verifiedScore) (err error) {
	for b := range blobs {
		result := verifiedScore{
			score: b.score,
		}

		// Make sure the score matches.
		computed := blob.ComputeScore(b.contents)
		if !scoresEqual(b.score, computed) {
			err = fmt.Errorf(
				"Score mismatch: %v vs. %v",
				b.score.Hex(),
				computed.Hex())

			return
		}

		// Parse the blob.
		result.children, err = unmarshalBlob(crypter, b.contents)
		if err != nil {
			err = fmt.Errorf("unmarshalBlob: %v", err)
			return
		}

		// Write out the result.
		select {
		case results <- result:
		case <-ctx.Done():
			err = ctx.Err()
			return
		}
	}

	return
}

// Write results to stdout.
func writeResults(
	ctx context.Context,
	output io.Writer,
	results <-chan verifiedScore) (err error) {
	// Log also to stdout.
	output = io.MultiWriter(output, os.Stdout)

	// Process each result.
	for result := range results {
		// Write the score.
		_, err = fmt.Fprintf(output, "%s", result.score.Hex())
		if err != nil {
			return
		}

		// Write its children.
		for _, child := range result.children {
			_, err = fmt.Fprintf(output, " %s", child.Hex())
			if err != nil {
				return
			}
		}

		// Add a line separator.
		_, err = fmt.Fprintf(output, "\n")
		if err != nil {
			return
		}
	}

	return
}

func runScanBlobs(args []string) {
	// Die on error.
	var err error
	defer func() {
		if err != nil {
			log.Fatalln(err)
		}
	}()

	// Open the output file.
	if *fOutputFile == "" {
		err = fmt.Errorf("You must set --output_file.")
		return
	}

	output, err := os.OpenFile(
		*fOutputFile,
		os.O_WRONLY|os.O_TRUNC|os.O_CREATE,
		0666)

	if err != nil {
		err = fmt.Errorf("OpenFile: %v", err)
		return
	}

	defer output.Close()

	// Grab dependencies.
	bucket := getBucket()
	var crypter crypto.Crypter
	if !*fFast {
		crypter = getCrypter()
	}

	// Allow parallelism.
	runtime.GOMAXPROCS(runtime.NumCPU())

	b := syncutil.NewBundle(context.Background())

	// List all of the scores in the bucket.
	scores := make(chan blob.Score, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(scores)
		err = listBlobs(ctx, bucket, scores)
		return
	})

	results := make(chan verifiedScore, 100)
	if !*fFast {
		// Read the contents of the corresponding blobs in parallel, bounding how
		// hard we hammer GCS by bounding the parallelism.
		const readWorkers = 8
		var readWaitGroup sync.WaitGroup

		blobs := make(chan scoreAndContents, 10)
		for i := 0; i < readWorkers; i++ {
			readWaitGroup.Add(1)
			b.Add(func(ctx context.Context) (err error) {
				defer readWaitGroup.Done()
				err = readBlobs(ctx, bucket, scores, blobs)
				return
			})
		}

		go func() {
			readWaitGroup.Wait()
			close(blobs)
		}()

		// Verify the blob contents and summarize their children. Use one worker per
		// CPU.
		var verifyWaitGroup sync.WaitGroup
		for i := 0; i < runtime.NumCPU(); i++ {
			verifyWaitGroup.Add(1)
			b.Add(func(ctx context.Context) (err error) {
				defer verifyWaitGroup.Done()
				err = verifyScores(ctx, crypter, blobs, results)
				return
			})
		}

		go func() {
			verifyWaitGroup.Wait()
			close(results)
		}()
	} else {
		// Fast mode. Pretend that everything is verified, and that nothing has
		// children.
		b.Add(func(ctx context.Context) (err error) {
			defer close(results)
			for score := range scores {
				select {
				case results <- verifiedScore{score: score}:

				// Cancelled?
				case <-ctx.Done():
					err = ctx.Err()
					return
				}
			}

			return
		})
	}

	// Process results.
	b.Add(func(ctx context.Context) (err error) {
		err = writeResults(ctx, output, results)
		return
	})

	// Wait for everything to complete.
	err = b.Join()
	return
}
<commit_msg>use a larger buffer for listing.<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// A command that processes all of the blobs in a GCS bucket. For each blob:
//
//  1. The GCS object record is verified for internal consistency. In
//     particular, comeback's expected MD5 and CRC32C are checked against the
//     GCS-reported values.
//
//  2. The object contents are read and the SHA-1 computed and checked against
//     the blob's name and expected SHA-1.
//
//  3. The edges in the directory graph are printed to an output file.
//
// (In "fast mode", only #1 is performed, and the output is a simple list of
// scores, one per line.)
//
// Output is space-separated lines of the form:
//
//     <score> <child score>*
//
// where <score> is a verified score, and the list is a list of scores of
// children for directories (and empty for files).

package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"sync"
	"time"

	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/crypto"
	"github.com/jacobsa/comeback/repr"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/syncutil"
	"golang.org/x/net/context"
)

var cmdScanBlobs = &Command{
	Name: "scan_blobs",
}

var fOutputFile = cmdScanBlobs.Flags.String(
	"output_file",
	"",
	"Path to scan_blobs output file. Will be truncated.")

var fFast = cmdScanBlobs.Flags.Bool(
	"fast",
	false,
	"Only verify object metadata.")

func init() {
	cmdScanBlobs.Run = runScanBlobs // Break flag-related dependency loop.
}

type scoreAndContents struct {
	score    blob.Score
	contents []byte
}

type verifiedScore struct {
	score    blob.Score
	children []blob.Score
}

// List all blob objects in the GCS bucket, and verify their metadata. Write
// their scores into the supplied channel.
func listBlobs(
	ctx context.Context,
	bucket gcs.Bucket,
	scores chan<- blob.Score) (err error) {
	req := &gcs.ListObjectsRequest{
		Prefix: blobObjectNamePrefix,
	}

	// List until we run out.
	for {
		// Fetch the next batch.
		var listing *gcs.Listing
		listing, err = bucket.ListObjects(ctx, req)
		if err != nil {
			err = fmt.Errorf("ListObjects: %v", err)
			return
		}

		// Transform to scores, verifying in the process.
		var batch []blob.Score
		for _, o := range listing.Objects {
			// Skip an object named simply blobObjectNamePrefix, which we allow to
			// exist for gcsfuse compatibility.
			if o.Name == blobObjectNamePrefix {
				continue
			}

			// Parse and verify the record.
			var score blob.Score
			score, err = blob.ParseObjectRecord(o, blobObjectNamePrefix)
			if err != nil {
				err = fmt.Errorf("ParseObjectRecord: %v", err)
				return
			}

			batch = append(batch, score)
		}

		// Feed out each score.
		for _, score := range batch {
			select {
			case scores <- score:

				// Cancelled?
			case <-ctx.Done():
				err = ctx.Err()
				return
			}
		}

		// Are we done?
		if listing.ContinuationToken == "" {
			break
		}

		req.ContinuationToken = listing.ContinuationToken
	}

	return
}

func readAndClose(rc io.ReadCloser) (d []byte, err error) {
	// Read.
	d, err = ioutil.ReadAll(rc)
	if err != nil {
		rc.Close()
		err = fmt.Errorf("ReadAll: %v", err)
		return
	}

	// Close.
	err = rc.Close()
	if err != nil {
		err = fmt.Errorf("Close: %v", err)
		return
	}

	return
}

func readBlob(
	ctx context.Context,
	bucket gcs.Bucket,
	score blob.Score) (blob scoreAndContents, err error) {
	// Fill in the score.
	blob.score = score

	// Obtain a reader.
	req := &gcs.ReadObjectRequest{
		Name: blobObjectNamePrefix + score.Hex(),
	}

	rc, err := bucket.NewReader(ctx, req)
	if err != nil {
		err = fmt.Errorf("NewReader: %v", err)
		return
	}

	// Consume it.
	blob.contents, err = readAndClose(rc)
	if err != nil {
		err = fmt.Errorf("readAndClose: %v", err)
		return
	}

	return
}

func readBlobWithRetryLoop(
	ctx context.Context,
	bucket gcs.Bucket,
	score blob.Score) (blob scoreAndContents, err error) {
	const maxRetries = 5
	for n := 0; n < maxRetries; n++ {
		blob, err = readBlob(ctx, bucket, score)
		if err == nil {
			return
		}

		time.Sleep(500 * time.Millisecond)
	}

	return
}

// Read the contents of blobs specified on the incoming channel. Do not close
// the outgoing channel.
func readBlobs(
	ctx context.Context,
	bucket gcs.Bucket,
	scores <-chan blob.Score,
	blobs chan<- scoreAndContents) (err error) {
	for score := range scores {
		var blob scoreAndContents

		// Read the contents for this score.
		blob, err = readBlobWithRetryLoop(ctx, bucket, score)
		if err != nil {
			err = fmt.Errorf("readBlob: %v", err)
			return
		}

		select {
		case blobs <- blob:
		case <-ctx.Done():
			err = ctx.Err()
			return
		}
	}

	return
}

func scoresEqual(a, b blob.Score) bool {
	for i := 0; i < blob.ScoreLength; i++ {
		if a[i] != b[i] {
			return false
		}
	}

	return true
}

// Parse the blob ciphertext as a directory if appropriate, returning a list of
// children. If not a directory, return the empty list.
func unmarshalBlob(
	crypter crypto.Crypter,
	b []byte) (children []blob.Score, err error) {
	// Attempt to decrypt.
	b, err = crypter.Decrypt(b)
	if err != nil {
		err = fmt.Errorf("Decrypt: %v", err)
		return
	}

	// If this is a file, simply make sure it is in a legal format.
	if !repr.IsDir(b) {
		_, err = repr.UnmarshalFile(b)
		if err != nil {
			err = fmt.Errorf("UnmarshalFile: %v", err)
			return
		}

		return
	}

	// Otherwise, parse it as a directory.
	entries, err := repr.UnmarshalDir(b)
	if err != nil {
		err = fmt.Errorf("UnmarshalDir: %v", err)
		return
	}

	// Pull out the child scores.
	for _, entry := range entries {
		for _, score := range entry.Scores {
			children = append(children, score)
		}
	}

	return
}

// Verify the contents of the incoming blobs. Do not close the outgoing
// channel.
func verifyScores(
	ctx context.Context,
	crypter crypto.Crypter,
	blobs <-chan scoreAndContents,
	results chan<- verifiedScore) (err error) {
	for b := range blobs {
		result := verifiedScore{
			score: b.score,
		}

		// Make sure the score matches.
		computed := blob.ComputeScore(b.contents)
		if !scoresEqual(b.score, computed) {
			err = fmt.Errorf(
				"Score mismatch: %v vs. %v",
				b.score.Hex(),
				computed.Hex())

			return
		}

		// Parse the blob.
		result.children, err = unmarshalBlob(crypter, b.contents)
		if err != nil {
			err = fmt.Errorf("unmarshalBlob: %v", err)
			return
		}

		// Write out the result.
		select {
		case results <- result:
		case <-ctx.Done():
			err = ctx.Err()
			return
		}
	}

	return
}

// Write results to stdout.
func writeResults(
	ctx context.Context,
	output io.Writer,
	results <-chan verifiedScore) (err error) {
	// Log also to stdout.
	output = io.MultiWriter(output, os.Stdout)

	// Process each result.
	for result := range results {
		// Write the score.
		_, err = fmt.Fprintf(output, "%s", result.score.Hex())
		if err != nil {
			return
		}

		// Write its children.
		for _, child := range result.children {
			_, err = fmt.Fprintf(output, " %s", child.Hex())
			if err != nil {
				return
			}
		}

		// Add a line separator.
		_, err = fmt.Fprintf(output, "\n")
		if err != nil {
			return
		}
	}

	return
}

func runScanBlobs(args []string) {
	// Die on error.
	var err error
	defer func() {
		if err != nil {
			log.Fatalln(err)
		}
	}()

	// Open the output file.
	if *fOutputFile == "" {
		err = fmt.Errorf("You must set --output_file.")
		return
	}

	output, err := os.OpenFile(
		*fOutputFile,
		os.O_WRONLY|os.O_TRUNC|os.O_CREATE,
		0666)

	if err != nil {
		err = fmt.Errorf("OpenFile: %v", err)
		return
	}

	defer output.Close()

	// Grab dependencies.
	bucket := getBucket()
	var crypter crypto.Crypter
	if !*fFast {
		crypter = getCrypter()
	}

	// Allow parallelism.
	runtime.GOMAXPROCS(runtime.NumCPU())

	b := syncutil.NewBundle(context.Background())

	// List all of the scores in the bucket.
	scores := make(chan blob.Score, 5000)
	b.Add(func(ctx context.Context) (err error) {
		defer close(scores)
		err = listBlobs(ctx, bucket, scores)
		return
	})

	results := make(chan verifiedScore, 100)
	if !*fFast {
		// Read the contents of the corresponding blobs in parallel, bounding how
		// hard we hammer GCS by bounding the parallelism.
		const readWorkers = 8
		var readWaitGroup sync.WaitGroup

		blobs := make(chan scoreAndContents, 10)
		for i := 0; i < readWorkers; i++ {
			readWaitGroup.Add(1)
			b.Add(func(ctx context.Context) (err error) {
				defer readWaitGroup.Done()
				err = readBlobs(ctx, bucket, scores, blobs)
				return
			})
		}

		go func() {
			readWaitGroup.Wait()
			close(blobs)
		}()

		// Verify the blob contents and summarize their children. Use one worker per
		// CPU.
		var verifyWaitGroup sync.WaitGroup
		for i := 0; i < runtime.NumCPU(); i++ {
			verifyWaitGroup.Add(1)
			b.Add(func(ctx context.Context) (err error) {
				defer verifyWaitGroup.Done()
				err = verifyScores(ctx, crypter, blobs, results)
				return
			})
		}

		go func() {
			verifyWaitGroup.Wait()
			close(results)
		}()
	} else {
		// Fast mode. Pretend that everything is verified, and that nothing has
		// children.
		b.Add(func(ctx context.Context) (err error) {
			defer close(results)
			for score := range scores {
				select {
				case results <- verifiedScore{score: score}:

				// Cancelled?
				case <-ctx.Done():
					err = ctx.Err()
					return
				}
			}

			return
		})
	}

	// Process results.
	b.Add(func(ctx context.Context) (err error) {
		err = writeResults(ctx, output, results)
		return
	})

	// Wait for everything to complete.
	err = b.Join()
	return
}
<|endoftext|>


<commit_before>package backend

// MemStore is a key value store in memory.
type MemStore struct {
	Data map[string][]string
}

// Init initializes the data store
func (m *MemStore) Init() {
	m.Data = make(map[string][]string)
}

// Add inserts a key into the backend
func (m *MemStore) Add(user string, key string) (bool, error) {
	m.Data[user] = append(m.Data[user], key)
	return true, nil
}

// RM removes a key from the backend
func (m *MemStore) RM(user string, key string) (bool, error) {
	var idx int
	for i := 0; i < len(m.Data[user]); i++ {
		if m.Data[user][i] == key {
			idx = i
		}
	}

	m.Data[user] = append(m.Data[user][:idx], m.Data[user][idx+1:]...)

	return true, nil
}

// RMAll removes all of the keys for a given user
func (m *MemStore) RMAll(user string) (bool, error) {
	delete(m.Data, user)
	return true, nil
}

// Get queries the backend for 'user' and returns all the keys
func (m *MemStore) Get(user string) ([]string, error) {
	return m.Data[user], nil
}

// GetAll returns the entire datastructure
func (m *MemStore) GetAll() (map[string][]string, error) {
	return m.Data, nil
}

// GetKeyCount queries the number of keys a given user has
func (m *MemStore) GetKeyCount(user string) (int, error) {
	return len(m.Data[user]), nil
}

// GetCount queries the number of keys stored in a backend
func (m *MemStore) GetCount() (int, error) {
	var i int
	for j := range m.Data {
		for _ = range m.Data[j] {
			i++
		}
	}
	return i, nil
}
<commit_msg>add close, return []string for GetAll<commit_after>package backend

// MemStore is a key value store in memory.
type MemStore struct {
	Data map[string][]string
}

// Init initializes the data store
func (m *MemStore) Init() error {
	m.Data = make(map[string][]string)
	return nil
}

// Close "closes" the store.. really we don't do anything because RAM IS CHEAP, RIGHT?!
func (m *MemStore) Close() {
}

// Add inserts a key into the backend
func (m *MemStore) Add(user string, key string) (bool, error) {
	m.Data[user] = append(m.Data[user], key)
	return true, nil
}

// RM removes a key from the backend
func (m *MemStore) RM(user string, key string) (bool, error) {
	var idx int
	for i := 0; i < len(m.Data[user]); i++ {
		if m.Data[user][i] == key {
			idx = i
		}
	}

	m.Data[user] = append(m.Data[user][:idx], m.Data[user][idx+1:]...)

	return true, nil
}

// RMAll removes all of the keys for a given user
func (m *MemStore) RMAll(user string) (bool, error) {
	delete(m.Data, user)
	return true, nil
}

// Get queries the backend for 'user' and returns all the keys
func (m *MemStore) Get(user string) ([]string, error) {
	return m.Data[user], nil
}

// GetAll returns the entire datastructure
func (m *MemStore) GetAll() ([]string, error) {
	var s []string
	for i := range m.Data {
		for j := range m.Data[i] {
			s = append(s, m.Data[i][j])
		}
	}
	return s, nil
}

// GetKeyCount queries the number of keys a given user has
func (m *MemStore) GetKeyCount(user string) (int, error) {
	return len(m.Data[user]), nil
}

// GetCount queries the number of keys stored in a backend
func (m *MemStore) GetCount() (int, error) {
	var i int
	for j := range m.Data {
		for _ = range m.Data[j] {
			i++
		}
	}
	return i, nil
}
<|endoftext|>


<commit_before>package engine

import (
	"context"
	"fmt"
	"time"

	"code.cloudfoundry.org/clock"
	"code.cloudfoundry.org/lager/lagerctx"
	"github.com/concourse/concourse/atc"
	"github.com/concourse/concourse/atc/db"
	"github.com/concourse/concourse/atc/db/lock"
	"github.com/concourse/concourse/atc/event"
	"github.com/concourse/concourse/atc/exec"
	"github.com/concourse/concourse/atc/policy"
	"github.com/concourse/concourse/atc/worker"
)

//go:generate counterfeiter . RateLimiter

type RateLimiter interface {
	Wait(context.Context) error
}

func NewCheckDelegate(
	build db.Build,
	plan atc.Plan,
	state exec.RunState,
	clock clock.Clock,
	limiter RateLimiter,
	policyChecker policy.Checker,
	artifactSourcer worker.ArtifactSourcer,
) exec.CheckDelegate {
	return &checkDelegate{
		BuildStepDelegate: NewBuildStepDelegate(build, plan.ID, state, clock, policyChecker, artifactSourcer),

		build:       build,
		plan:        plan.Check,
		eventOrigin: event.Origin{ID: event.OriginID(plan.ID)},
		clock:       clock,

		limiter: limiter,
	}
}

type checkDelegate struct {
	exec.BuildStepDelegate

	build       db.Build
	plan        *atc.CheckPlan
	eventOrigin event.Origin
	clock       clock.Clock

	// stashed away just so we don't have to query them multiple times
	cachedPipeline     db.Pipeline
	cachedResource     db.Resource
	cachedResourceType db.ResourceType

	limiter RateLimiter
}

func (d *checkDelegate) FindOrCreateScope(config db.ResourceConfig) (db.ResourceConfigScope, error) {
	resource, _, err := d.resource()
	if err != nil {
		return nil, fmt.Errorf("get resource: %w", err)
	}

	scope, err := config.FindOrCreateScope(resource) // ignore found, nil is ok
	if err != nil {
		return nil, fmt.Errorf("find or create scope: %w", err)
	}

	return scope, nil
}

func (d *checkDelegate) WaitToRun(ctx context.Context, scope db.ResourceConfigScope) (lock.Lock, bool, error) {
	logger := lagerctx.FromContext(ctx)

	// rate limit periodic resource checks so worker load (plus load on external
	// services) isn't too spiky
	if !d.build.IsManuallyTriggered() && d.plan.Resource != "" {
		err := d.limiter.Wait(ctx)
		if err != nil {
			return nil, false, fmt.Errorf("rate limit: %w", err)
		}
	}

	var err error

	var interval time.Duration
	if d.plan.Interval != "" {
		interval, err = time.ParseDuration(d.plan.Interval)
		if err != nil {
			return nil, false, err
		}
	}

	var lock lock.Lock = lock.NoopLock{}
	if d.plan.Resource != "" {
		for {
			var acquired bool
			lock, acquired, err = scope.AcquireResourceCheckingLock(logger)
			if err != nil {
				return nil, false, fmt.Errorf("acquire lock: %w", err)
			}

			if acquired {
				break
			}

			d.clock.Sleep(time.Second)
		}
	}

	end, err := scope.LastCheckEndTime()
	if err != nil {
		if releaseErr := lock.Release(); releaseErr != nil {
			logger.Error("failed-to-release-lock", releaseErr)
		}

		return nil, false, fmt.Errorf("get last check end time: %w", err)
	}

	runAt := end.Add(interval)

	shouldRun := false
	if d.build.IsManuallyTriggered() {
		// ignore interval for manually triggered builds
		shouldRun = true
	} else if !d.clock.Now().Before(runAt) {
		// run if we're past the last check end time
		shouldRun = true
	} else {
		// XXX(check-refactor): we could potentially sleep here until runAt is
		// reached.
		//
		// then the check build queueing logic is to just make sure there's a build
		// running for every resource, without having to check if intervals have
		// elapsed.
		//
		// this could be expanded upon to short-circuit the waiting with events
		// triggered by webhooks so that webhooks are super responsive: rather than
		// queueing a build, it would just wake up a goroutine.
	}

	if !shouldRun {
		err := lock.Release()
		if err != nil {
			return nil, false, fmt.Errorf("release lock: %w", err)
		}

		return nil, false, nil
	}

	return lock, true, nil
}

func (d *checkDelegate) PointToCheckedConfig(scope db.ResourceConfigScope) error {
	resource, found, err := d.resource()
	if err != nil {
		return fmt.Errorf("get resource: %w", err)
	}

	if found {
		err := resource.SetResourceConfigScope(scope)
		if err != nil {
			return fmt.Errorf("set resource scope: %w", err)
		}
	}

	resourceType, found, err := d.resourceType()
	if err != nil {
		return fmt.Errorf("get resource type: %w", err)
	}

	if found {
		err := resourceType.SetResourceConfigScope(scope)
		if err != nil {
			return fmt.Errorf("set resource type scope: %w", err)
		}
	}

	return nil
}

func (d *checkDelegate) pipeline() (db.Pipeline, error) {
	if d.cachedPipeline != nil {
		return d.cachedPipeline, nil
	}

	pipeline, found, err := d.build.Pipeline()
	if err != nil {
		return nil, fmt.Errorf("get build pipeline: %w", err)
	}

	if !found {
		return nil, fmt.Errorf("pipeline not found")
	}

	d.cachedPipeline = pipeline

	return d.cachedPipeline, nil
}

func (d *checkDelegate) resource() (db.Resource, bool, error) {
	if d.plan.Resource == "" {
		return nil, false, nil
	}

	if d.cachedResource != nil {
		return d.cachedResource, true, nil
	}

	pipeline, err := d.pipeline()
	if err != nil {
		return nil, false, err
	}

	resource, found, err := pipeline.Resource(d.plan.Resource)
	if err != nil {
		return nil, false, fmt.Errorf("get pipeline resource: %w", err)
	}

	if !found {
		return nil, false, fmt.Errorf("resource '%s' deleted", d.plan.Resource)
	}

	d.cachedResource = resource

	return d.cachedResource, true, nil
}

func (d *checkDelegate) resourceType() (db.ResourceType, bool, error) {
	if d.plan.ResourceType == "" {
		return nil, false, nil
	}

	if d.cachedResourceType != nil {
		return d.cachedResourceType, true, nil
	}

	pipeline, err := d.pipeline()
	if err != nil {
		return nil, false, err
	}

	resourceType, found, err := pipeline.ResourceType(d.plan.ResourceType)
	if err != nil {
		return nil, false, fmt.Errorf("get pipeline resource type: %w", err)
	}

	if !found {
		return nil, false, fmt.Errorf("resource type '%s' deleted", d.plan.ResourceType)
	}

	d.cachedResourceType = resourceType

	return d.cachedResourceType, true, nil
}
<commit_msg>Remove unused else{} case<commit_after>package engine

import (
	"context"
	"fmt"
	"time"

	"code.cloudfoundry.org/clock"
	"code.cloudfoundry.org/lager/lagerctx"
	"github.com/concourse/concourse/atc"
	"github.com/concourse/concourse/atc/db"
	"github.com/concourse/concourse/atc/db/lock"
	"github.com/concourse/concourse/atc/event"
	"github.com/concourse/concourse/atc/exec"
	"github.com/concourse/concourse/atc/policy"
	"github.com/concourse/concourse/atc/worker"
)

//go:generate counterfeiter . RateLimiter

type RateLimiter interface {
	Wait(context.Context) error
}

func NewCheckDelegate(
	build db.Build,
	plan atc.Plan,
	state exec.RunState,
	clock clock.Clock,
	limiter RateLimiter,
	policyChecker policy.Checker,
	artifactSourcer worker.ArtifactSourcer,
) exec.CheckDelegate {
	return &checkDelegate{
		BuildStepDelegate: NewBuildStepDelegate(build, plan.ID, state, clock, policyChecker, artifactSourcer),

		build:       build,
		plan:        plan.Check,
		eventOrigin: event.Origin{ID: event.OriginID(plan.ID)},
		clock:       clock,

		limiter: limiter,
	}
}

type checkDelegate struct {
	exec.BuildStepDelegate

	build       db.Build
	plan        *atc.CheckPlan
	eventOrigin event.Origin
	clock       clock.Clock

	// stashed away just so we don't have to query them multiple times
	cachedPipeline     db.Pipeline
	cachedResource     db.Resource
	cachedResourceType db.ResourceType

	limiter RateLimiter
}

func (d *checkDelegate) FindOrCreateScope(config db.ResourceConfig) (db.ResourceConfigScope, error) {
	resource, _, err := d.resource()
	if err != nil {
		return nil, fmt.Errorf("get resource: %w", err)
	}

	scope, err := config.FindOrCreateScope(resource) // ignore found, nil is ok
	if err != nil {
		return nil, fmt.Errorf("find or create scope: %w", err)
	}

	return scope, nil
}

func (d *checkDelegate) WaitToRun(ctx context.Context, scope db.ResourceConfigScope) (lock.Lock, bool, error) {
	logger := lagerctx.FromContext(ctx)

	// rate limit periodic resource checks so worker load (plus load on external
	// services) isn't too spiky
	if !d.build.IsManuallyTriggered() && d.plan.Resource != "" {
		err := d.limiter.Wait(ctx)
		if err != nil {
			return nil, false, fmt.Errorf("rate limit: %w", err)
		}
	}

	var err error

	var interval time.Duration
	if d.plan.Interval != "" {
		interval, err = time.ParseDuration(d.plan.Interval)
		if err != nil {
			return nil, false, err
		}
	}

	var lock lock.Lock = lock.NoopLock{}
	if d.plan.Resource != "" {
		for {
			var acquired bool
			lock, acquired, err = scope.AcquireResourceCheckingLock(logger)
			if err != nil {
				return nil, false, fmt.Errorf("acquire lock: %w", err)
			}

			if acquired {
				break
			}

			d.clock.Sleep(time.Second)
		}
	}

	end, err := scope.LastCheckEndTime()
	if err != nil {
		if releaseErr := lock.Release(); releaseErr != nil {
			logger.Error("failed-to-release-lock", releaseErr)
		}

		return nil, false, fmt.Errorf("get last check end time: %w", err)
	}

	runAt := end.Add(interval)

	shouldRun := false
	if d.build.IsManuallyTriggered() {
		// ignore interval for manually triggered builds
		shouldRun = true
	} else if !d.clock.Now().Before(runAt) {
		// run if we're past the last check end time
		shouldRun = true
	}
	// XXX(check-refactor): we could add an else{} case and potentially sleep
	// here until runAt is reached.
	//
	// then the check build queueing logic is to just make sure there's a build
	// running for every resource, without having to check if intervals have
	// elapsed.
	//
	// this could be expanded upon to short-circuit the waiting with events
	// triggered by webhooks so that webhooks are super responsive: rather than
	// queueing a build, it would just wake up a goroutine.

	if !shouldRun {
		err := lock.Release()
		if err != nil {
			return nil, false, fmt.Errorf("release lock: %w", err)
		}

		return nil, false, nil
	}

	return lock, true, nil
}

func (d *checkDelegate) PointToCheckedConfig(scope db.ResourceConfigScope) error {
	resource, found, err := d.resource()
	if err != nil {
		return fmt.Errorf("get resource: %w", err)
	}

	if found {
		err := resource.SetResourceConfigScope(scope)
		if err != nil {
			return fmt.Errorf("set resource scope: %w", err)
		}
	}

	resourceType, found, err := d.resourceType()
	if err != nil {
		return fmt.Errorf("get resource type: %w", err)
	}

	if found {
		err := resourceType.SetResourceConfigScope(scope)
		if err != nil {
			return fmt.Errorf("set resource type scope: %w", err)
		}
	}

	return nil
}

func (d *checkDelegate) pipeline() (db.Pipeline, error) {
	if d.cachedPipeline != nil {
		return d.cachedPipeline, nil
	}

	pipeline, found, err := d.build.Pipeline()
	if err != nil {
		return nil, fmt.Errorf("get build pipeline: %w", err)
	}

	if !found {
		return nil, fmt.Errorf("pipeline not found")
	}

	d.cachedPipeline = pipeline

	return d.cachedPipeline, nil
}

func (d *checkDelegate) resource() (db.Resource, bool, error) {
	if d.plan.Resource == "" {
		return nil, false, nil
	}

	if d.cachedResource != nil {
		return d.cachedResource, true, nil
	}

	pipeline, err := d.pipeline()
	if err != nil {
		return nil, false, err
	}

	resource, found, err := pipeline.Resource(d.plan.Resource)
	if err != nil {
		return nil, false, fmt.Errorf("get pipeline resource: %w", err)
	}

	if !found {
		return nil, false, fmt.Errorf("resource '%s' deleted", d.plan.Resource)
	}

	d.cachedResource = resource

	return d.cachedResource, true, nil
}

func (d *checkDelegate) resourceType() (db.ResourceType, bool, error) {
	if d.plan.ResourceType == "" {
		return nil, false, nil
	}

	if d.cachedResourceType != nil {
		return d.cachedResourceType, true, nil
	}

	pipeline, err := d.pipeline()
	if err != nil {
		return nil, false, err
	}

	resourceType, found, err := pipeline.ResourceType(d.plan.ResourceType)
	if err != nil {
		return nil, false, fmt.Errorf("get pipeline resource type: %w", err)
	}

	if !found {
		return nil, false, fmt.Errorf("resource type '%s' deleted", d.plan.ResourceType)
	}

	d.cachedResourceType = resourceType

	return d.cachedResourceType, true, nil
}
<|endoftext|>