<commit_before>package integration_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
	"golang.org/x/crypto/ssh"
	"srcd.works/go-git.v4"
	gitssh "srcd.works/go-git.v4/plumbing/transport/ssh"
)

const (
	slackRateLimitDelay  = 200 * time.Millisecond
	slackRequestAttempts = 5
)

var _ = Describe("Claimer", func() {
	var (
		claimer   string
		gitDir    string
		apiToken  string
		channelId string
		repoUrl   string
		deployKey string
	)

	BeforeSuite(func() {
		var err error

		gitDir, err = ioutil.TempDir("", "claimer-integration-tests")
		Expect(err).NotTo(HaveOccurred())

		claimer, err = gexec.Build(filepath.Join("github.com", "mdelillo", "claimer"))
		Expect(err).NotTo(HaveOccurred())

		apiToken = getEnv("CLAIMER_TEST_API_TOKEN")
		channelId = getEnv("CLAIMER_TEST_CHANNEL_ID")
		repoUrl = getEnv("CLAIMER_TEST_REPO_URL")
		deployKey = getEnv("CLAIMER_TEST_DEPLOY_KEY")
	})

	AfterEach(func() {
		gexec.KillAndWait()
	})

	AfterSuite(func() {
		gexec.CleanupBuildArtifacts()
		os.RemoveAll(gitDir)
	})

	It("claims and releases locks", func() {
		botId := getEnv("CLAIMER_TEST_BOT_ID")
		userApiToken := getEnv("CLAIMER_TEST_USER_API_TOKEN")
		username := getEnv("CLAIMER_TEST_USERNAME")
		userId := getEnv("CLAIMER_TEST_USER_ID")
		otherChannelId := getEnv("CLAIMER_TEST_OTHER_CHANNEL_ID")

		signer, err := ssh.ParsePrivateKey([]byte(deployKey))
		Expect(err).NotTo(HaveOccurred())

		_, err = git.PlainClone(gitDir, false, &git.CloneOptions{
			URL: repoUrl,
			Auth: &gitssh.PublicKeys{
				User:   "git",
				Signer: signer,
			},
		})
		Expect(err).NotTo(HaveOccurred())
		defer resetClaimerTestPool(gitDir, deployKey)

		resetClaimerTestPool(gitDir, deployKey)

		claimerCommand := exec.Command(
			claimer,
			"-apiToken", apiToken,
			"-channelId", channelId,
			"-repoUrl", repoUrl,
			"-deployKey", deployKey,
		)
		session, err := gexec.Start(claimerCommand, GinkgoWriter, GinkgoWriter)
		Expect(err).NotTo(HaveOccurred())

		Eventually(session, "20s").Should(gbytes.Say("Listening for messages"))

		By("Displaying the help message")
		postSlackMessage(fmt.Sprintf("<@%s> help", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(ContainSubstring("Available commands:"))

		By("Checking the status")
		postSlackMessage(fmt.Sprintf("<@%s> status", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("*Claimed by you:* \n*Claimed by others:* pool-3\n*Unclaimed:* pool-1"))

		By("Claiming pool-1")
		postSlackMessage(fmt.Sprintf("<@%s> claim pool-1", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("Claimed pool-1"))
		updateGitRepo(gitDir, deployKey)
		Expect(filepath.Join(gitDir, "pool-1", "claimed", "lock-a")).To(BeAnExistingFile())
		Expect(filepath.Join(gitDir, "pool-1", "unclaimed", "lock-a")).NotTo(BeAnExistingFile())

		By("Checking the status")
		postSlackMessage(fmt.Sprintf("<@%s> status", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(ContainSubstring("*Claimed by you:* pool-1\n"))
		Expect(latestSlackMessage(channelId, apiToken)).NotTo(MatchRegexp(`\*Unclaimed:\*.*pool-1`))

		By("Checking the owner of pool-1")
		postSlackMessage(fmt.Sprintf("<@%s> owner pool-1", botId), channelId, userApiToken)
		ownerMessage := fmt.Sprintf("pool-1 was claimed by %s on ", username)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").Should(ContainSubstring(ownerMessage))
		date := strings.TrimPrefix(latestSlackMessage(channelId, apiToken), ownerMessage)
		parsedDate, err := time.Parse("Mon Jan 2 15:04:05 2006 -0700", date)
		Expect(err).NotTo(HaveOccurred())
		Expect(parsedDate).To(BeTemporally("~", time.Now(), 10*time.Second))

		By("Notifying owners of their claimed locks when characters preceding @claimer exist")
		postSlackMessage(fmt.Sprintf("Reminder: <@%s> notify", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(ContainSubstring("Currently claimed locks, please release if not in use:\n"))
		Expect(latestSlackMessage(channelId, apiToken)).To(ContainSubstring(fmt.Sprintf("<@%s>: pool-1", userId)))

		By("Trying to claim pool-1 again")
		postSlackMessage(fmt.Sprintf("<@%s> claim pool-1", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("pool-1 is already claimed"))

		By("Releasing pool-1")
		postSlackMessage(fmt.Sprintf("<@%s> release pool-1", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("Released pool-1"))
		updateGitRepo(gitDir, deployKey)
		Expect(filepath.Join(gitDir, "pool-1", "unclaimed", "lock-a")).To(BeAnExistingFile())
		Expect(filepath.Join(gitDir, "pool-1", "claimed", "lock-a")).NotTo(BeAnExistingFile())

		By("Checking the status")
		postSlackMessage(fmt.Sprintf("<@%s> status", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(ContainSubstring("*Claimed by you:* \n"))
		Expect(latestSlackMessage(channelId, apiToken)).To(MatchRegexp(`\*Unclaimed:\*.*pool-1`))

		By("Checking the status of pool-1")
		postSlackMessage(fmt.Sprintf("<@%s> owner pool-1", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("pool-1 is not claimed"))

		By("Trying to release pool-1 again")
		postSlackMessage(fmt.Sprintf("<@%s> release pool-1", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("pool-1 is not claimed"))

		By("Claiming pool-1 with a message")
		postSlackMessage(fmt.Sprintf("<@%s> claim pool-1 some message", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("Claimed pool-1"))
		postSlackMessage(fmt.Sprintf("<@%s> owner pool-1", botId), channelId, userApiToken)

		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").Should(HaveSuffix(" (some message)"))

		postSlackMessage(fmt.Sprintf("<@%s> release pool-1", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("Released pool-1"))

		By("Trying to claim non-existent pool")
		postSlackMessage(fmt.Sprintf("<@%s> claim non-existent-pool", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("non-existent-pool does not exist"))

		By("Trying to claim without a pool name")
		postSlackMessage(fmt.Sprintf("<@%s> claim", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("must specify pool to claim"))

		By("Trying to release non-existent-pool")
		postSlackMessage(fmt.Sprintf("<@%s> release non-existent-pool", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("non-existent-pool does not exist"))

		By("Trying to run an unknown command")
		postSlackMessage(fmt.Sprintf("<@%s> unknown-command", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("Unknown command. Try `@claimer help` to see usage."))

		By("Mentioning claimer in a different channel")
		postSlackMessage(fmt.Sprintf("<@%s> help", botId), otherChannelId, userApiToken)
		Consistently(func() string { return latestSlackMessage(otherChannelId, apiToken) }, "10s").
			Should(Equal(fmt.Sprintf("<@%s> help", botId)))

		By("Creating a pool")
		postSlackMessage(fmt.Sprintf("<@%s> create new-pool", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("Created new-pool"))

		By("Trying to create a pool that already exists")
		postSlackMessage(fmt.Sprintf("<@%s> create new-pool", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("new-pool already exists"))

		By("Destroying a pool")
		postSlackMessage(fmt.Sprintf("<@%s> destroy new-pool", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("Destroyed new-pool"))

		By("Trying to destroy a pool that does not exist")
		postSlackMessage(fmt.Sprintf("<@%s> destroy new-pool", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(Equal("new-pool does not exist"))
	})

	It("responds with a message from a given translation file", func() {
		botId := getEnv("CLAIMER_TEST_BOT_ID")
		userApiToken := getEnv("CLAIMER_TEST_USER_API_TOKEN")

		translationFile, err := ioutil.TempFile("", "claimer-integration-tests")
		Expect(err).NotTo(HaveOccurred())
		translations := "help: {header: foo}"
		Expect(ioutil.WriteFile(translationFile.Name(), []byte(translations), 0644)).To(Succeed())
		defer os.Remove(translationFile.Name())

		claimerCommand := exec.Command(
			claimer,
			"-apiToken", apiToken,
			"-channelId", channelId,
			"-repoUrl", repoUrl,
			"-deployKey", deployKey,
			"-translationFile", translationFile.Name(),
		)
		session, err := gexec.Start(claimerCommand, GinkgoWriter, GinkgoWriter)
		Expect(err).NotTo(HaveOccurred())

		Eventually(session, "20s").Should(gbytes.Say("Listening for messages"))

		By("Displaying the custom help message")
		postSlackMessage(fmt.Sprintf("<@%s> help", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(HavePrefix("foo"))

		By("Displaying the default status message")
		postSlackMessage(fmt.Sprintf("<@%s> status", botId), channelId, userApiToken)
		Eventually(func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
			Should(ContainSubstring("Claimed by you:"))
	})
})

func getEnv(name string) string {
	value, ok := os.LookupEnv(name)
	if !ok {
		Fail(fmt.Sprintf("%s must be set", name))
	}
	return value
}

func postSlackMessage(text, channelId, apiToken string) {
	_, err := slackPostForm("https://slack.com/api/chat.postMessage", url.Values{
		"token":   {apiToken},
		"channel": {channelId},
		"text":    {text},
		"as_user": {"true"},
	})
	Expect(err).NotTo(HaveOccurred())
}

func latestSlackMessage(channelId, apiToken string) string {
	body, err := slackPostForm("https://slack.com/api/channels.history", url.Values{
		"token":   {apiToken},
		"channel": {channelId},
		"count":   {"1"},
	})
	Expect(err).NotTo(HaveOccurred())
	var slackResponse struct {
		Messages []struct {
			Text string
		}
	}
	Expect(json.Unmarshal(body, &slackResponse)).To(Succeed())

	return slackResponse.Messages[0].Text
}

func slackPostForm(url string, values url.Values) ([]byte, error) {
	delay := slackRateLimitDelay
	for i := 0; i < slackRequestAttempts; i++ {
		time.Sleep(delay)

		body, err := postForm(url, values)
		if err != nil {
			return nil, err
		}

		var slackResponse struct {
			Ok    bool
			Error string
		}
		if err := json.Unmarshal(body, &slackResponse); err != nil {
			return nil, err
		}

		if slackResponse.Ok {
			return body, nil
		} else if slackResponse.Error != "ratelimited" {
			return nil, fmt.Errorf("Slack request failed: %s", slackResponse.Error)
		}

		delay *= 2
	}
	return nil, fmt.Errorf("Slack request failed %d times", slackRequestAttempts)
}

func postForm(url string, values url.Values) ([]byte, error) {
	response, err := http.PostForm(url, values)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()

	return ioutil.ReadAll(response.Body)
}

func updateGitRepo(gitDir, deployKey string) {
	runGitCommand(gitDir, deployKey, "fetch")
	runGitCommand(gitDir, deployKey, "reset", "--hard", "origin/master")
}

func resetClaimerTestPool(gitDir, deployKey string) {
	runGitCommand(gitDir, deployKey, "checkout", "master")
	runGitCommand(gitDir, deployKey, "reset", "--hard", "initial-state")
	runGitCommand(gitDir, deployKey, "push", "--force", "origin", "master")
}

func runGitCommand(dir, deployKey string, args ...string) {
	deployKeyDir, err := ioutil.TempDir("", "claimer-integration-test-deploy-key")
	Expect(err).NotTo(HaveOccurred())
	defer os.RemoveAll(deployKeyDir)

	deployKeyPath := filepath.Join(deployKeyDir, "key.pem")
	Expect(ioutil.WriteFile(deployKeyPath, []byte(deployKey), 0600)).To(Succeed())

	cmd := exec.Command("git", args...)
	cmd.Dir = dir
	cmd.Env = append(os.Environ(), fmt.Sprintf(`GIT_SSH_COMMAND=/usr/bin/ssh -i %s`, deployKeyPath))
	output, err := cmd.CombinedOutput()
	ExpectWithOffset(1, err).NotTo(HaveOccurred(), fmt.Sprintf("Error running git command: %s", string(output)))
}
<commit_msg>Clean up integration tests<commit_after>package integration_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
	"golang.org/x/crypto/ssh"
	"srcd.works/go-git.v4"
	gitssh "srcd.works/go-git.v4/plumbing/transport/ssh"
)

const (
	slackRateLimitDelay  = 200 * time.Millisecond
	slackRequestAttempts = 5
)

var _ = Describe("Claimer", func() {
	var (
		apiToken       string
		channelId      string
		repoUrl        string
		deployKey      string
		botId          string
		userApiToken   string
		username       string
		userId         string
		otherChannelId string
		runCommand     func(string) string
		startClaimer   func(string) *gexec.Session
		gitDir         string
	)

	BeforeSuite(func() {
		claimer, err := gexec.Build(filepath.Join("github.com", "mdelillo", "claimer"))
		Expect(err).NotTo(HaveOccurred())

		apiToken = getEnv("CLAIMER_TEST_API_TOKEN")
		channelId = getEnv("CLAIMER_TEST_CHANNEL_ID")
		repoUrl = getEnv("CLAIMER_TEST_REPO_URL")
		deployKey = getEnv("CLAIMER_TEST_DEPLOY_KEY")
		botId = getEnv("CLAIMER_TEST_BOT_ID")
		userApiToken = getEnv("CLAIMER_TEST_USER_API_TOKEN")
		username = getEnv("CLAIMER_TEST_USERNAME")
		userId = getEnv("CLAIMER_TEST_USER_ID")
		otherChannelId = getEnv("CLAIMER_TEST_OTHER_CHANNEL_ID")

		runCommand = func(command string) string {
			message := fmt.Sprintf("<@%s> %s", botId, command)
			postSlackMessage(message, channelId, userApiToken)
			EventuallyWithOffset(1, func() string { return latestSlackMessage(channelId, apiToken) }, "10s").
				ShouldNot(Equal(message), fmt.Sprintf(`Did not get response from command "%s"`, command))
			return latestSlackMessage(channelId, apiToken)
		}
		startClaimer = func(translationFile string) *gexec.Session {
			args := []string{
				"-apiToken", apiToken,
				"-channelId", channelId,
				"-repoUrl", repoUrl,
				"-deployKey", deployKey,
			}
			if translationFile != "" {
				args = append(args, "-translationFile", translationFile)
			}
			claimerCommand := exec.Command(claimer, args...)
			session, err := gexec.Start(claimerCommand, GinkgoWriter, GinkgoWriter)
			ExpectWithOffset(1, err).NotTo(HaveOccurred())
			return session
		}
	})

	BeforeEach(func() {
		var err error
		gitDir, err = ioutil.TempDir("", "claimer-integration-tests")
		Expect(err).NotTo(HaveOccurred())

		signer, err := ssh.ParsePrivateKey([]byte(deployKey))
		Expect(err).NotTo(HaveOccurred())

		_, err = git.PlainClone(gitDir, false, &git.CloneOptions{
			URL: repoUrl,
			Auth: &gitssh.PublicKeys{
				User:   "git",
				Signer: signer,
			},
		})
		Expect(err).NotTo(HaveOccurred())

		resetClaimerTestPool(gitDir, deployKey)
	})

	AfterEach(func() {
		gexec.KillAndWait()
		resetClaimerTestPool(gitDir, deployKey)
		Expect(os.RemoveAll(gitDir)).To(Succeed())
	})

	AfterSuite(func() {
		gexec.CleanupBuildArtifacts()
	})

	It("claims, releases, and shows status of locks", func() {
		session := startClaimer("")
		Eventually(session, "20s").Should(gbytes.Say("Listening for messages"))

		Expect(runCommand("help")).To(ContainSubstring("Available commands:"))

		Expect(runCommand("status")).To(Equal("*Claimed by you:* \n*Claimed by others:* pool-3\n*Unclaimed:* pool-1"))

		Expect(runCommand("claim pool-1")).To(Equal("Claimed pool-1"))
		updateGitRepo(gitDir, deployKey)
		Expect(filepath.Join(gitDir, "pool-1", "claimed", "lock-a")).To(BeAnExistingFile())
		Expect(filepath.Join(gitDir, "pool-1", "unclaimed", "lock-a")).NotTo(BeAnExistingFile())

		response := runCommand("status")
		Expect(response).To(ContainSubstring("*Claimed by you:* pool-1\n"))
		Expect(response).NotTo(MatchRegexp(`\*Unclaimed:\*.*pool-1`))

		Expect(runCommand("claim pool-1")).To(Equal("pool-1 is already claimed"))

		Expect(runCommand("release pool-1")).To(Equal("Released pool-1"))
		updateGitRepo(gitDir, deployKey)
		Expect(filepath.Join(gitDir, "pool-1", "unclaimed", "lock-a")).To(BeAnExistingFile())
		Expect(filepath.Join(gitDir, "pool-1", "claimed", "lock-a")).NotTo(BeAnExistingFile())

		response = runCommand("status")
		Expect(response).To(ContainSubstring("*Claimed by you:* \n"))
		Expect(response).To(MatchRegexp(`\*Unclaimed:\*.*pool-1`))

		Expect(runCommand("release pool-1")).To(Equal("pool-1 is not claimed"))

		Expect(runCommand("claim non-existent-pool")).To(Equal("non-existent-pool does not exist"))

		Expect(runCommand("claim")).To(Equal("must specify pool to claim"))

		Expect(runCommand("release non-existent-pool")).To(Equal("non-existent-pool does not exist"))

		Expect(runCommand("unknown-command")).To(Equal("Unknown command. Try `@claimer help` to see usage."))
	})

	It("shows the owner of a lock", func() {
		session := startClaimer("")
		Eventually(session, "20s").Should(gbytes.Say("Listening for messages"))

		Expect(runCommand("owner pool-1")).To(Equal("pool-1 is not claimed"))

		claimTime := time.Now()
		Expect(runCommand("claim pool-1")).To(Equal("Claimed pool-1"))

		response := runCommand("owner pool-1")
		ownerMessage := fmt.Sprintf("pool-1 was claimed by %s on ", username)
		Expect(response).To(ContainSubstring(ownerMessage))

		date := strings.TrimPrefix(response, ownerMessage)
		parsedDate, err := time.Parse("Mon Jan 2 15:04:05 2006 -0700", date)
		Expect(err).NotTo(HaveOccurred())
		Expect(parsedDate).To(BeTemporally("~", claimTime, 10*time.Second))

		Expect(runCommand("release pool-1")).To(Equal("Released pool-1"))

		Expect(runCommand("claim pool-1 some message")).To(Equal("Claimed pool-1"))

		Expect(runCommand("owner pool-1")).To(HaveSuffix(" (some message)"))
	})

	It("notifies users who have claimed locks", func() {
		session := startClaimer("")
		Eventually(session, "20s").Should(gbytes.Say("Listening for messages"))

		Expect(runCommand("claim pool-1")).To(Equal("Claimed pool-1"))

		response := runCommand("notify")
		Expect(response).To(ContainSubstring("Currently claimed locks, please release if not in use:\n"))
		Expect(response).To(ContainSubstring(fmt.Sprintf("<@%s>: pool-1", userId)))
	})

	It("creates and destroys locks", func() {
		session := startClaimer("")
		Eventually(session, "20s").Should(gbytes.Say("Listening for messages"))

		Expect(runCommand("create new-pool")).To(Equal("Created new-pool"))

		updateGitRepo(gitDir, deployKey)
		Expect(filepath.Join(gitDir, "new-pool", "unclaimed", "new-pool")).To(BeAnExistingFile())

		Expect(runCommand("status")).To(MatchRegexp(`\*Unclaimed:\*.*new-pool`))

		Expect(runCommand("create new-pool")).To(Equal("new-pool already exists"))

		Expect(runCommand("destroy new-pool")).To(Equal("Destroyed new-pool"))

		updateGitRepo(gitDir, deployKey)
		Expect(filepath.Join(gitDir, "new-pool")).NotTo(BeADirectory())

		Expect(runCommand("destroy new-pool")).To(Equal("new-pool does not exist"))

		Expect(runCommand("status")).NotTo(MatchRegexp(`\*Unclaimed:\*.*new-pool`))
	})

	It("does not respond in other channels", func() {
		session := startClaimer("")
		Eventually(session, "20s").Should(gbytes.Say("Listening for messages"))

		postSlackMessage(fmt.Sprintf("<@%s> help", botId), otherChannelId, userApiToken)

		Consistently(func() string { return latestSlackMessage(otherChannelId, apiToken) }, "10s").
			Should(Equal(fmt.Sprintf("<@%s> help", botId)))
	})

	Context("when a translation file is provided", func() {
		var translationFilePath string

		BeforeEach(func() {
			translationFile, err := ioutil.TempFile("", "claimer-integration-tests")
			Expect(err).NotTo(HaveOccurred())
			translationFilePath = translationFile.Name()
		})

		AfterEach(func() {
			Expect(os.Remove(translationFilePath)).To(Succeed())
		})

		It("responds with a message from the given translation file", func() {
			translations := "help: {header: foo}"
			Expect(ioutil.WriteFile(translationFilePath, []byte(translations), 0644)).To(Succeed())

			session := startClaimer(translationFilePath)
			Eventually(session, "20s").Should(gbytes.Say("Listening for messages"))

			Expect(runCommand("help")).To(HavePrefix("foo"))

			Expect(runCommand("status")).To(ContainSubstring("Claimed by you:"))
		})
	})
})

func getEnv(name string) string {
	value, ok := os.LookupEnv(name)
	if !ok {
		Fail(fmt.Sprintf("%s must be set", name))
	}
	return value
}

func postSlackMessage(text, channelId, apiToken string) {
	_, err := slackPostForm("https://slack.com/api/chat.postMessage", url.Values{
		"token":   {apiToken},
		"channel": {channelId},
		"text":    {text},
		"as_user": {"true"},
	})
	Expect(err).NotTo(HaveOccurred())
}

func latestSlackMessage(channelId, apiToken string) string {
	body, err := slackPostForm("https://slack.com/api/channels.history", url.Values{
		"token":   {apiToken},
		"channel": {channelId},
		"count":   {"1"},
	})
	Expect(err).NotTo(HaveOccurred())
	var slackResponse struct {
		Messages []struct {
			Text string
		}
	}
	Expect(json.Unmarshal(body, &slackResponse)).To(Succeed())

	return slackResponse.Messages[0].Text
}

func slackPostForm(url string, values url.Values) ([]byte, error) {
	delay := slackRateLimitDelay
	for i := 0; i < slackRequestAttempts; i++ {
		time.Sleep(delay)

		body, err := postForm(url, values)
		if err != nil {
			return nil, err
		}

		var slackResponse struct {
			Ok    bool
			Error string
		}
		if err := json.Unmarshal(body, &slackResponse); err != nil {
			return nil, err
		}

		if slackResponse.Ok {
			return body, nil
		} else if slackResponse.Error != "ratelimited" {
			return nil, fmt.Errorf("Slack request failed: %s", slackResponse.Error)
		}

		delay *= 2
	}
	return nil, fmt.Errorf("Slack request failed %d times", slackRequestAttempts)
}

func postForm(url string, values url.Values) ([]byte, error) {
	response, err := http.PostForm(url, values)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()

	return ioutil.ReadAll(response.Body)
}

func updateGitRepo(gitDir, deployKey string) {
	runGitCommand(gitDir, deployKey, "fetch")
	runGitCommand(gitDir, deployKey, "reset", "--hard", "origin/master")
}

func resetClaimerTestPool(gitDir, deployKey string) {
	runGitCommand(gitDir, deployKey, "checkout", "master")
	runGitCommand(gitDir, deployKey, "reset", "--hard", "initial-state")
	runGitCommand(gitDir, deployKey, "push", "--force", "origin", "master")
}

func runGitCommand(dir, deployKey string, args ...string) {
	deployKeyDir, err := ioutil.TempDir("", "claimer-integration-test-deploy-key")
	Expect(err).NotTo(HaveOccurred())
	defer os.RemoveAll(deployKeyDir)

	deployKeyPath := filepath.Join(deployKeyDir, "key.pem")
	Expect(ioutil.WriteFile(deployKeyPath, []byte(deployKey), 0600)).To(Succeed())

	cmd := exec.Command("git", args...)
	cmd.Dir = dir
	cmd.Env = append(os.Environ(), fmt.Sprintf(`GIT_SSH_COMMAND=/usr/bin/ssh -i %s`, deployKeyPath))
	output, err := cmd.CombinedOutput()
	ExpectWithOffset(1, err).NotTo(HaveOccurred(), fmt.Sprintf("Error running git command: %s", string(output)))
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"

	"github.com/gorilla/mux"
	"github.com/nytlabs/st-core/core"
)

type ConnectionNode struct {
	Id    int `json:"id"`
	Route int `json:"route"`
}

type ConnectionLedger struct {
	Source ConnectionNode `json:"source"`
	Target ConnectionNode `json:"target"`
	Id     int            `json:"id"`
}

type ProtoConnection struct {
	Source ConnectionNode `json:"source"`
	Target ConnectionNode `json:"target"`
}

func (s *Server) ListConnections() []ConnectionLedger {
	connections := []ConnectionLedger{}
	for _, c := range s.connections {
		connections = append(connections, *c)
	}
	return connections
}

func (s *Server) ConnectionIndexHandler(w http.ResponseWriter, r *http.Request) {
	s.Lock()
	defer s.Unlock()

	c := s.ListConnections()

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	if err := json.NewEncoder(w).Encode(c); err != nil {
		panic(err)
	}
}

// CreateConnectionHandler responds to a POST request to instantiate a new connection
func (s *Server) ConnectionCreateHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	var newConn ProtoConnection
	json.Unmarshal(body, &newConn)

	s.Lock()
	defer s.Unlock()

	nc, err := s.CreateConnection(newConn)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusOK)
	writeJSON(w, nc)
}

func (s *Server) CreateConnection(newConn ProtoConnection) (*ConnectionLedger, error) {
	source, ok := s.blocks[newConn.Source.Id]
	if !ok {
		return nil, errors.New("source block does not exist")
	}

	target, ok := s.blocks[newConn.Target.Id]
	if !ok {
		return nil, errors.New("target block does not exist")
	}

	sourceRoute := core.RouteIndex(newConn.Source.Route)
	targetRoute, err := target.Block.GetInput(core.RouteIndex(newConn.Target.Route))
	if err != nil {
		return nil, err
	}

	err = source.Block.Connect(sourceRoute, targetRoute.C)
	if err != nil {
		return nil, err
	}

	conn := &ConnectionLedger{
		Source: newConn.Source,
		Target: newConn.Target,
		Id:     s.GetNextID(),
	}

	s.connections[conn.Id] = conn

	s.websocketBroadcast(Update{Action: CREATE, Type: CONNECTION, Data: conn})
	return conn, nil
}

// returns a description of the connection
func (s *Server) ConnectionHandler(w http.ResponseWriter, r *http.Request) {

	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	conn, ok := s.connections[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not find connection" + string(id)})
		return
	}

	w.WriteHeader(http.StatusOK)
	writeJSON(w, conn)
}

func (s *Server) ConnectionModifyCoordinates(w http.ResponseWriter, r *http.Request) {
}

func (s *Server) DeleteConnection(id int) error {
	c, ok := s.connections[id]
	if !ok {
		return errors.New("could not find connection")
	}

	source, ok := s.blocks[c.Source.Id]
	if !ok {
		return errors.New("could not find source block")
	}

	target, ok := s.blocks[c.Target.Id]
	if !ok {
		return errors.New("could not find target block")
	}

	route, err := target.Block.GetInput(core.RouteIndex(c.Target.Route))
	if err != nil {
		return err
	}

	err = source.Block.Disconnect(core.RouteIndex(c.Source.Route), route.C)
	if err != nil {
		return err
	}

	delete(s.connections, id)

	s.websocketBroadcast(Update{Action: DELETE, Type: CONNECTION, Data: c})
	return nil
}

func (s *Server) ConnectionDeleteHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.DeleteConnection(id)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}
<commit_msg>fixing server api<commit_after>package server

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"

	"github.com/gorilla/mux"
	"github.com/nytlabs/st-core/core"
)

type ConnectionNode struct {
	Id    int `json:"id"`
	Route int `json:"route"`
}

type ConnectionLedger struct {
	Source ConnectionNode `json:"from"`
	Target ConnectionNode `json:"to"`
	Id     int            `json:"id"`
}

type ProtoConnection struct {
	Source ConnectionNode `json:"from"`
	Target ConnectionNode `json:"to"`
}

func (s *Server) ListConnections() []ConnectionLedger {
	connections := []ConnectionLedger{}
	for _, c := range s.connections {
		connections = append(connections, *c)
	}
	return connections
}

func (s *Server) ConnectionIndexHandler(w http.ResponseWriter, r *http.Request) {
	s.Lock()
	defer s.Unlock()

	c := s.ListConnections()

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.WriteHeader(http.StatusOK)
	if err := json.NewEncoder(w).Encode(c); err != nil {
		panic(err)
	}
}

// CreateConnectionHandler responds to a POST request to instantiate a new connection
func (s *Server) ConnectionCreateHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	var newConn ProtoConnection
	json.Unmarshal(body, &newConn)

	s.Lock()
	defer s.Unlock()

	nc, err := s.CreateConnection(newConn)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusOK)
	writeJSON(w, nc)
}

func (s *Server) CreateConnection(newConn ProtoConnection) (*ConnectionLedger, error) {
	source, ok := s.blocks[newConn.Source.Id]
	if !ok {
		return nil, errors.New("source block does not exist")
	}

	target, ok := s.blocks[newConn.Target.Id]
	if !ok {
		return nil, errors.New("target block does not exist")
	}

	sourceRoute := core.RouteIndex(newConn.Source.Route)
	targetRoute, err := target.Block.GetInput(core.RouteIndex(newConn.Target.Route))
	if err != nil {
		return nil, err
	}

	err = source.Block.Connect(sourceRoute, targetRoute.C)
	if err != nil {
		return nil, err
	}

	conn := &ConnectionLedger{
		Source: newConn.Source,
		Target: newConn.Target,
		Id:     s.GetNextID(),
	}

	s.connections[conn.Id] = conn

	s.websocketBroadcast(Update{Action: CREATE, Type: CONNECTION, Data: wsConnection{*conn}})
	return conn, nil
}

// returns a description of the connection
func (s *Server) ConnectionHandler(w http.ResponseWriter, r *http.Request) {

	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	conn, ok := s.connections[id]
	if !ok {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{"could not find connection" + string(id)})
		return
	}

	w.WriteHeader(http.StatusOK)
	writeJSON(w, conn)
}

func (s *Server) ConnectionModifyCoordinates(w http.ResponseWriter, r *http.Request) {
}

func (s *Server) DeleteConnection(id int) error {
	c, ok := s.connections[id]
	if !ok {
		return errors.New("could not find connection")
	}

	source, ok := s.blocks[c.Source.Id]
	if !ok {
		return errors.New("could not find source block")
	}

	target, ok := s.blocks[c.Target.Id]
	if !ok {
		return errors.New("could not find target block")
	}

	route, err := target.Block.GetInput(core.RouteIndex(c.Target.Route))
	if err != nil {
		return err
	}

	err = source.Block.Disconnect(core.RouteIndex(c.Source.Route), route.C)
	if err != nil {
		return err
	}

	delete(s.connections, id)

	s.websocketBroadcast(Update{Action: DELETE, Type: CONNECTION, Data: wsConnection{wsId{id}}})
	return nil
}

func (s *Server) ConnectionDeleteHandler(w http.ResponseWriter, r *http.Request) {
	id, err := getIDFromMux(mux.Vars(r))
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, err)
		return
	}

	s.Lock()
	defer s.Unlock()

	err = s.DeleteConnection(id)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		writeJSON(w, Error{err.Error()})
		return
	}

	w.WriteHeader(http.StatusNoContent)
}
<|endoftext|>


<commit_before>package main

/*
 * Copyright 2014 Albert P. Tobey <atobey@datastax.com> @AlTobey
 *
 * 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.
 *
 * This program is a Docker entrypoint for Apache Cassandra.
 */

import (
	"bytes"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/tobert/sprok"
	"gopkg.in/yaml.v2"
)

const ugid = 1337

type CassandraDockerConfig struct {
	VolDir         string   // read/write data, should be a volume
	SrcConfDir     string   // root path for assets to copy to the volume
	ConfDir        string   // conf directory
	DataDir        string   // data directory
	CommitLogDir   string   // cl directory
	LogDir         string   // log directory
	SavedCachesDir string   // saved_caches directory
	CqlshDotDir    string   // ~/.cassandra
	CassandraYaml  string   // conf/cassandra.yaml
	SprokDir       string   // conf/sproks directory
	ExtraArgs      []string // args to be passed to child commands
	// Cassandra configuration items
	ClusterName      string // cluster_name in cassandra.yaml
	Seeds            string // seeds value for cassandra.yaml
	CassandraLogfile string // system.log
	DefaultIP        string // IP of the default route interface
	JmxPort          string // JMX port for nodetool
}

func main() {
	cdc := CassandraDockerConfig{
		VolDir:           "/data",
		SrcConfDir:       "/src/conf",
		ConfDir:          "/data/conf",
		DataDir:          "/data/data",
		CommitLogDir:     "/data/commitlog",
		LogDir:           "/data/log",
		SavedCachesDir:   "/data/saved_caches",
		CqlshDotDir:      "/data/.cassandra",
		CassandraYaml:    "/data/conf/cassandra.yaml",
		SprokDir:         "/data/conf/sproks",
		ClusterName:      "Docker Cluster",
		Seeds:            "127.0.0.1",
		CassandraLogfile: "/data/log/system.log",
		DefaultIP:        "127.0.0.1",
		JmxPort:          "7199",
	}

	var command, sprokFile string
	var args []string

	// extract the command, e.g. 'cassandra', 'nodetool' from os.Args
	// when not present it's assumed to be 'cassandra' even when arguments
	// are provided.
	if path.Base(os.Args[0]) != "cassandra-docker" {
		// handle symlink commands, e.g. ln -s /bin/cassandra-docker /bin/cqlsh
		command = path.Base(os.Args[0])
		args = os.Args[1:]
	} else if len(os.Args) == 1 {
		// no arguments: run cassandra
		command = "cassandra"
		args = []string{}
	} else if len(os.Args) > 1 {
		// when no command is provided, assume cassandra + flags
		// otherwise take the first argument as the command and check it below
		if strings.HasPrefix(os.Args[1], "-") {
			command = "cassandra"
			args = os.Args[1:]
		} else {
			command = os.Args[1]
			if len(os.Args) > 2 {
				args = os.Args[2:]
			} else {
				args = []string{}
			}
		}
	}

	// parse the subcommand and arguments to it
	switch command {
	case "cassandra":
		args, _, cdc.Seeds = extractArg(args, "seeds", "127.0.0.1")
		args, _, cdc.ClusterName = extractArg(args, "name", "Docker Cluster")
		sprokFile = path.Join(cdc.SprokDir, "cassandra.yaml")
	case "cqlsh":
		sprokFile = path.Join(cdc.SprokDir, "cqlsh.yaml")
	case "nodetool":
		args, _, cdc.JmxPort = extractArg(args, "p", "7199")
		sprokFile = path.Join(cdc.SprokDir, "nodetool.yaml")
	case "cassandra-stress":
		sprokFile = path.Join(cdc.SprokDir, "cassandra-stress.yaml")
	default:
		log.Fatalf("invalid command '%s'", command)
	}

	// copy the remaining command-line args to cdc so templates can render
	cdc.ExtraArgs = args

	// bootstrap - find the default IP, make directories, copy files
	cdc.setDefaultIP()
	if strings.EqualFold(cdc.Seeds, "127.0.0.1") {
		cdc.Seeds = cdc.DefaultIP
	}

	cdc.mkdirs()
	// copies files from src to data, running them through as templates
	// in the process. existing files are not overwritten
	cdc.tmplCopy()

	// load the sprok config
	fd, err := os.Open(sprokFile)
	if err != nil {
		log.Fatalf("error opening '%s' for read: %s\n", sprokFile, err)
	}

	// render the config template before unmarshaling
	// this allows sprok files to work across upgrades with smart use
	// of glob() to work around files with version numbers in them
	var data bytes.Buffer
	cdc.render(fd, &data)

	// configure the process from the yaml
	proc := sprok.NewProcess()
	err = yaml.Unmarshal(data.Bytes(), &proc)
	if err != nil {
		log.Fatalf("could not parse YAML in file '%s': %s\n", sprokFile, err)
	}

	// this is an actual execve(3p), this process is replaced with the new one
	proc.Exec()
}

func (cdc *CassandraDockerConfig) mkdirs() {
	mkdirAll(cdc.ConfDir)
	mkdirAll(cdc.DataDir)
	mkdirAll(cdc.CommitLogDir)
	mkdirAll(cdc.LogDir)
	mkdirAll(cdc.SavedCachesDir)
	mkdirAll(cdc.CqlshDotDir)
	mkdirAll(cdc.SprokDir)

	chownAll(cdc.DataDir)
	chownAll(cdc.CommitLogDir)
	chownAll(cdc.LogDir)
	chownAll(cdc.SavedCachesDir)
	chownAll(cdc.CqlshDotDir)
}

// tmplCopy reads all the files in cdc.SrcConfDir, treating them as text
// templates, then writes them to cdc.ConfDir. If a file exists in ConfDir
// it is not overwritten.
func (cdc *CassandraDockerConfig) tmplCopy() {
	walk := func(fromName string, fromFi os.FileInfo, err error) error {
		if err != nil {
			log.Fatalf("failed to find files in '%s': %s\n", cdc.SrcConfDir, err)
		}

		// only safe for same filesystem with no relative paths or symlinks
		toName := strings.Replace(fromName, cdc.SrcConfDir, cdc.ConfDir, 1)

		if exists(toName) {
			return nil // try not to overwrite any existing files
		} else if strings.HasPrefix(path.Base(fromName), ".") {
			return nil
		} else if fromFi.IsDir() {
			if exists(toName) {
				return nil
			} else {
				mkdirAll(toName)
			}
		} else if fromFi.Mode().IsRegular() {
			// don't render sprok files, only copy them
			// they will get rendered at run time
			if strings.HasSuffix(path.Dir(fromName), "sproks") {
				cp(fromName, toName)
			} else {
				cdc.renderFile(fromName, toName)
			}
		} else {
			log.Fatalf("unsupported file mode on file '%s'\n", fromName)
		}

		return nil
	}

	err := filepath.Walk(cdc.SrcConfDir, walk)
	if err != nil {
		log.Fatalf("tmplCopy() failed: %s\n", err)
	}
}

// renderFile renders one file to another using text/template
func (cdc *CassandraDockerConfig) renderFile(src, dest string) {
	in, err := os.Open(src)
	if err != nil {
		log.Fatalf("could not open '%s' for reading: %s\n", src, err)
	}
	defer in.Close()

	out, err := os.OpenFile(dest, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		log.Fatalf("could not open '%s' for write: %s\n", dest, err)
	}
	defer out.Close()

	cdc.render(in, out)
}

// render renders an io.Reader to an io.Writer using text/template
func (cdc *CassandraDockerConfig) render(in io.Reader, out io.Writer) {
	funcMap := template.FuncMap{
		"glob": Globber,
	}

	tdata, err := ioutil.ReadAll(in)
	if err != nil {
		log.Fatalf("template read failed: %s\n", err)
	}

	tmpl, err := template.New("whatever").Funcs(funcMap).Parse(string(tdata))
	if err != nil {
		log.Fatalf("template parsing failed: %s", err)
	}

	err = tmpl.Execute(out, cdc)
	if err != nil {
		log.Fatalf("template rendering failed: %s\n", err)
	}
}

// setDefaultIP finds the first configured interface that is not a loopback
// and sets the cdc.DefaultIP value
func (cdc *CassandraDockerConfig) setDefaultIP() {
	ifaces, err := net.Interfaces()
	if err != nil {
		log.Fatalf("error while listing network interfaces: %s\n", err)
	}

	for _, iface := range ifaces {
		if iface.Flags&net.FlagUp == 0 {
			continue // interface is down
		}
		if iface.Flags&net.FlagLoopback != 0 {
			continue // ignore loopback interface
		}

		addrs, err := iface.Addrs()
		if err != nil {
			log.Fatalf("error while examining network interface: %s\n", err)
		}

		// for now, just go with the first interface that is up
		// and is not a loopback, which should cover most Docker setups
		for _, addr := range addrs {
			switch v := addr.(type) {
			case *net.IPNet:
				cdc.DefaultIP = v.IP.String()
				return
			}
		}
	}
}

// Globber takes paths, performs a glob match, then returns
// all the results joined with the specified seperator.
func Globber(sep string, globs ...string) string {
	if len(globs) == 0 {
		log.Fatalf("Globber() requires at least one path.\n")
	}

	out := []string{}

	for _, glob := range globs {
		filenames, err := filepath.Glob(glob)
		if err != nil {
			log.Fatalf("file glob failed: %s\n", err)
		}

		for _, filename := range filenames {
			out = append(out, filename)
		}
	}

	return strings.Join(out, sep)
}

// exists returns boolean whether a path exists or not
func exists(name string) bool {
	_, err := os.Stat(name)
	if os.IsNotExist(err) {
		return false
	} else if err != nil {
		log.Fatalf("could not stat file '%s': %s\n", name, err)
	}

	return true
}

// mkdirAll creates a directory recursively, crashes the program on error.
func mkdirAll(name string) {
	err := os.MkdirAll(name, 0755)
	if err != nil {
		log.Fatalf("os.MkdirAll('%s') failed: %s\n", name, err)
	}
}

func chownAll(name string) {
	walk := func(fname string, fi os.FileInfo, err error) error {
		if err != nil {
			log.Fatalf("error during fs walk of '%s': %s\n", fname, err)

		}

		return os.Chown(fname, ugid, ugid)
	}

	err := filepath.Walk(name, walk)
	if err != nil {
		log.Fatalf("chownAll('%s') failed: %s\n", name, err)
	}
}

// cp copies a file, crashing the program on any errors
// It does not attempt to use rename.
func cp(from, to string) {
	in, err := os.Open(from)
	if err != nil {
		log.Fatalf("could not open '%s' for reading: %s\n", from, err)
	}
	defer in.Close()

	out, err := os.OpenFile(to, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		log.Fatalf("could not open '%s' for writing: %s\n", to, err)
	}
	defer out.Close()

	_, err = io.Copy(out, in)
	if err != nil {
		log.Fatalf("data copy failed for file '%s': %s\n", to, err)
	}
}
<commit_msg>Add LibDir for sprok templates<commit_after>package main

/*
 * Copyright 2014 Albert P. Tobey <atobey@datastax.com> @AlTobey
 *
 * 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.
 *
 * This program is a Docker entrypoint for Apache Cassandra.
 */

import (
	"bytes"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/tobert/sprok"
	"gopkg.in/yaml.v2"
)

const ugid = 1337

type CassandraDockerConfig struct {
	VolDir         string   // read/write data, should be a volume
	SrcConfDir     string   // root path for assets to copy to the volume
	ConfDir        string   // conf directory
	DataDir        string   // data directory
	CommitLogDir   string   // cl directory
	LogDir         string   // log directory
	LibDir         string   // custom classpath directory
	SavedCachesDir string   // saved_caches directory
	CqlshDotDir    string   // ~/.cassandra
	CassandraYaml  string   // conf/cassandra.yaml
	SprokDir       string   // conf/sproks directory
	ExtraArgs      []string // args to be passed to child commands
	// Cassandra configuration items
	ClusterName      string // cluster_name in cassandra.yaml
	Seeds            string // seeds value for cassandra.yaml
	CassandraLogfile string // system.log
	DefaultIP        string // IP of the default route interface
	JmxPort          string // JMX port for nodetool
}

func main() {
	cdc := CassandraDockerConfig{
		VolDir:           "/data",
		SrcConfDir:       "/src/conf",
		ConfDir:          "/data/conf",
		DataDir:          "/data/data",
		CommitLogDir:     "/data/commitlog",
		LogDir:           "/data/log",
		LibDir:           "/data/lib",
		SavedCachesDir:   "/data/saved_caches",
		CqlshDotDir:      "/data/.cassandra",
		CassandraYaml:    "/data/conf/cassandra.yaml",
		SprokDir:         "/data/conf/sproks",
		ClusterName:      "Docker Cluster",
		Seeds:            "127.0.0.1",
		CassandraLogfile: "/data/log/system.log",
		DefaultIP:        "127.0.0.1",
		JmxPort:          "7199",
	}

	var command, sprokFile string
	var args []string

	// extract the command, e.g. 'cassandra', 'nodetool' from os.Args
	// when not present it's assumed to be 'cassandra' even when arguments
	// are provided.
	if path.Base(os.Args[0]) != "cassandra-docker" {
		// handle symlink commands, e.g. ln -s /bin/cassandra-docker /bin/cqlsh
		command = path.Base(os.Args[0])
		args = os.Args[1:]
	} else if len(os.Args) == 1 {
		// no arguments: run cassandra
		command = "cassandra"
		args = []string{}
	} else if len(os.Args) > 1 {
		// when no command is provided, assume cassandra + flags
		// otherwise take the first argument as the command and check it below
		if strings.HasPrefix(os.Args[1], "-") {
			command = "cassandra"
			args = os.Args[1:]
		} else {
			command = os.Args[1]
			if len(os.Args) > 2 {
				args = os.Args[2:]
			} else {
				args = []string{}
			}
		}
	}

	// parse the subcommand and arguments to it
	switch command {
	case "cassandra":
		args, _, cdc.Seeds = extractArg(args, "seeds", "127.0.0.1")
		args, _, cdc.ClusterName = extractArg(args, "name", "Docker Cluster")
		sprokFile = path.Join(cdc.SprokDir, "cassandra.yaml")
	case "cqlsh":
		sprokFile = path.Join(cdc.SprokDir, "cqlsh.yaml")
	case "nodetool":
		args, _, cdc.JmxPort = extractArg(args, "p", "7199")
		sprokFile = path.Join(cdc.SprokDir, "nodetool.yaml")
	case "cassandra-stress":
		sprokFile = path.Join(cdc.SprokDir, "cassandra-stress.yaml")
	default:
		log.Fatalf("invalid command '%s'", command)
	}

	// copy the remaining command-line args to cdc so templates can render
	cdc.ExtraArgs = args

	// bootstrap - find the default IP, make directories, copy files
	cdc.setDefaultIP()
	if strings.EqualFold(cdc.Seeds, "127.0.0.1") {
		cdc.Seeds = cdc.DefaultIP
	}

	cdc.mkdirs()
	// copies files from src to data, running them through as templates
	// in the process. existing files are not overwritten
	cdc.tmplCopy()

	// load the sprok config
	fd, err := os.Open(sprokFile)
	if err != nil {
		log.Fatalf("error opening '%s' for read: %s\n", sprokFile, err)
	}

	// render the config template before unmarshaling
	// this allows sprok files to work across upgrades with smart use
	// of glob() to work around files with version numbers in them
	var data bytes.Buffer
	cdc.render(fd, &data)

	// configure the process from the yaml
	proc := sprok.NewProcess()
	err = yaml.Unmarshal(data.Bytes(), &proc)
	if err != nil {
		log.Fatalf("could not parse YAML in file '%s': %s\n", sprokFile, err)
	}

	// this is an actual execve(3p), this process is replaced with the new one
	proc.Exec()
}

func (cdc *CassandraDockerConfig) mkdirs() {
	mkdirAll(cdc.ConfDir)
	mkdirAll(cdc.DataDir)
	mkdirAll(cdc.CommitLogDir)
	mkdirAll(cdc.LogDir)
	mkdirAll(cdc.SavedCachesDir)
	mkdirAll(cdc.CqlshDotDir)
	mkdirAll(cdc.SprokDir)

	chownAll(cdc.DataDir)
	chownAll(cdc.CommitLogDir)
	chownAll(cdc.LogDir)
	chownAll(cdc.SavedCachesDir)
	chownAll(cdc.CqlshDotDir)
}

// tmplCopy reads all the files in cdc.SrcConfDir, treating them as text
// templates, then writes them to cdc.ConfDir. If a file exists in ConfDir
// it is not overwritten.
func (cdc *CassandraDockerConfig) tmplCopy() {
	walk := func(fromName string, fromFi os.FileInfo, err error) error {
		if err != nil {
			log.Fatalf("failed to find files in '%s': %s\n", cdc.SrcConfDir, err)
		}

		// only safe for same filesystem with no relative paths or symlinks
		toName := strings.Replace(fromName, cdc.SrcConfDir, cdc.ConfDir, 1)

		if exists(toName) {
			return nil // try not to overwrite any existing files
		} else if strings.HasPrefix(path.Base(fromName), ".") {
			return nil
		} else if fromFi.IsDir() {
			if exists(toName) {
				return nil
			} else {
				mkdirAll(toName)
			}
		} else if fromFi.Mode().IsRegular() {
			// don't render sprok files, only copy them
			// they will get rendered at run time
			if strings.HasSuffix(path.Dir(fromName), "sproks") {
				cp(fromName, toName)
			} else {
				cdc.renderFile(fromName, toName)
			}
		} else {
			log.Fatalf("unsupported file mode on file '%s'\n", fromName)
		}

		return nil
	}

	err := filepath.Walk(cdc.SrcConfDir, walk)
	if err != nil {
		log.Fatalf("tmplCopy() failed: %s\n", err)
	}
}

// renderFile renders one file to another using text/template
func (cdc *CassandraDockerConfig) renderFile(src, dest string) {
	in, err := os.Open(src)
	if err != nil {
		log.Fatalf("could not open '%s' for reading: %s\n", src, err)
	}
	defer in.Close()

	out, err := os.OpenFile(dest, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		log.Fatalf("could not open '%s' for write: %s\n", dest, err)
	}
	defer out.Close()

	cdc.render(in, out)
}

// render renders an io.Reader to an io.Writer using text/template
func (cdc *CassandraDockerConfig) render(in io.Reader, out io.Writer) {
	funcMap := template.FuncMap{
		"glob": Globber,
	}

	tdata, err := ioutil.ReadAll(in)
	if err != nil {
		log.Fatalf("template read failed: %s\n", err)
	}

	tmpl, err := template.New("whatever").Funcs(funcMap).Parse(string(tdata))
	if err != nil {
		log.Fatalf("template parsing failed: %s", err)
	}

	err = tmpl.Execute(out, cdc)
	if err != nil {
		log.Fatalf("template rendering failed: %s\n", err)
	}
}

// setDefaultIP finds the first configured interface that is not a loopback
// and sets the cdc.DefaultIP value
func (cdc *CassandraDockerConfig) setDefaultIP() {
	ifaces, err := net.Interfaces()
	if err != nil {
		log.Fatalf("error while listing network interfaces: %s\n", err)
	}

	for _, iface := range ifaces {
		if iface.Flags&net.FlagUp == 0 {
			continue // interface is down
		}
		if iface.Flags&net.FlagLoopback != 0 {
			continue // ignore loopback interface
		}

		addrs, err := iface.Addrs()
		if err != nil {
			log.Fatalf("error while examining network interface: %s\n", err)
		}

		// for now, just go with the first interface that is up
		// and is not a loopback, which should cover most Docker setups
		for _, addr := range addrs {
			switch v := addr.(type) {
			case *net.IPNet:
				cdc.DefaultIP = v.IP.String()
				return
			}
		}
	}
}

// Globber takes paths, performs a glob match, then returns
// all the results joined with the specified seperator.
func Globber(sep string, globs ...string) string {
	if len(globs) == 0 {
		log.Fatalf("Globber() requires at least one path.\n")
	}

	out := []string{}

	for _, glob := range globs {
		filenames, err := filepath.Glob(glob)
		if err != nil {
			log.Fatalf("file glob failed: %s\n", err)
		}

		for _, filename := range filenames {
			out = append(out, filename)
		}
	}

	return strings.Join(out, sep)
}

// exists returns boolean whether a path exists or not
func exists(name string) bool {
	_, err := os.Stat(name)
	if os.IsNotExist(err) {
		return false
	} else if err != nil {
		log.Fatalf("could not stat file '%s': %s\n", name, err)
	}

	return true
}

// mkdirAll creates a directory recursively, crashes the program on error.
func mkdirAll(name string) {
	err := os.MkdirAll(name, 0755)
	if err != nil {
		log.Fatalf("os.MkdirAll('%s') failed: %s\n", name, err)
	}
}

func chownAll(name string) {
	walk := func(fname string, fi os.FileInfo, err error) error {
		if err != nil {
			log.Fatalf("error during fs walk of '%s': %s\n", fname, err)

		}

		return os.Chown(fname, ugid, ugid)
	}

	err := filepath.Walk(name, walk)
	if err != nil {
		log.Fatalf("chownAll('%s') failed: %s\n", name, err)
	}
}

// cp copies a file, crashing the program on any errors
// It does not attempt to use rename.
func cp(from, to string) {
	in, err := os.Open(from)
	if err != nil {
		log.Fatalf("could not open '%s' for reading: %s\n", from, err)
	}
	defer in.Close()

	out, err := os.OpenFile(to, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		log.Fatalf("could not open '%s' for writing: %s\n", to, err)
	}
	defer out.Close()

	_, err = io.Copy(out, in)
	if err != nil {
		log.Fatalf("data copy failed for file '%s': %s\n", to, err)
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"bytes"
	"context"

	"github.com/mitchellh/cli"
	"github.com/nerdalize/nerd/svc"
	"github.com/pkg/errors"
)

//JobLogs command
type JobLogs struct {
	KubeOpts
	Tail int64 `long:"tail" short:"t" description:"only return the oldest N lines of the process logs"`

	*command
}

//JobLogsFactory creates the command
func JobLogsFactory(ui cli.Ui) cli.CommandFactory {
	cmd := &JobLogs{}
	cmd.command = createCommand(ui, cmd.Execute, cmd.Description, cmd.Usage, cmd)
	return func() (cli.Command, error) {
		return cmd, nil
	}
}

//Execute runs the command
func (cmd *JobLogs) Execute(args []string) (err error) {
	if len(args) < 1 {
		return errors.New(MessageNotEnoughArguments)
	}

	kopts := cmd.KubeOpts
	deps, err := NewDeps(cmd.Logger(), kopts)
	if err != nil {
		return errors.Wrap(err, "failed to configure")
	}

	ctx := context.Background()
	ctx, cancel := context.WithTimeout(ctx, cmd.Timeout)
	defer cancel()

	in := &svc.FetchJobLogsInput{
		Name: args[0],
		Tail: cmd.Tail,
	}

	kube := svc.NewKube(deps, kopts.Namespace)
	out, err := kube.FetchJobLogs(ctx, in)
	if err != nil {
		return renderServiceError(err, "failed to fetch job logs")
	}

	lines := string(bytes.TrimSpace(out.Data))
	if len(lines) < 1 {
		cmd.out.Errorf("No logs visible (anymore) for job '%s'. Maybe the process didn't output any logs or it was created a long time ago: old logs may be discarded", in.Name)
		return nil
	}

	cmd.out.Output(string(out.Data))
	return nil
}

// Description returns long-form help text
func (cmd *JobLogs) Description() string { return cmd.Synopsis() }

// Synopsis returns a one-line
func (cmd *JobLogs) Synopsis() string { return "Return logs for a running job" }

// Usage shows usage
func (cmd *JobLogs) Usage() string { return "nerd job logs [NAME]" }
<commit_msg>add a message when logs are trimmed<commit_after>package cmd

import (
	"bytes"
	"context"

	"github.com/mitchellh/cli"
	"github.com/nerdalize/nerd/pkg/kubevisor"
	"github.com/nerdalize/nerd/svc"
	"github.com/pkg/errors"
)

//JobLogs command
type JobLogs struct {
	KubeOpts
	Tail int64 `long:"tail" short:"t" description:"only return the oldest N lines of the process logs"`

	*command
}

//JobLogsFactory creates the command
func JobLogsFactory(ui cli.Ui) cli.CommandFactory {
	cmd := &JobLogs{}
	cmd.command = createCommand(ui, cmd.Execute, cmd.Description, cmd.Usage, cmd)
	return func() (cli.Command, error) {
		return cmd, nil
	}
}

//Execute runs the command
func (cmd *JobLogs) Execute(args []string) (err error) {
	if len(args) < 1 {
		return errors.New(MessageNotEnoughArguments)
	}

	kopts := cmd.KubeOpts
	deps, err := NewDeps(cmd.Logger(), kopts)
	if err != nil {
		return errors.Wrap(err, "failed to configure")
	}

	ctx := context.Background()
	ctx, cancel := context.WithTimeout(ctx, cmd.Timeout)
	defer cancel()

	in := &svc.FetchJobLogsInput{
		Name: args[0],
		Tail: cmd.Tail,
	}

	kube := svc.NewKube(deps, kopts.Namespace)
	out, err := kube.FetchJobLogs(ctx, in)
	if err != nil {
		return renderServiceError(err, "failed to fetch job logs")
	}

	lines := string(bytes.TrimSpace(out.Data))
	if len(lines) < 1 {
		cmd.out.Info("-- no visible logs returned --")
		return nil
	}

	cmd.out.Output(string(out.Data))
	if int64(len(out.Data)) == kubevisor.MaxLogBytes {
		cmd.out.Info("-- logs are trimmed after this point --")
	}

	return nil
}

// Description returns long-form help text
func (cmd *JobLogs) Description() string { return cmd.Synopsis() }

// Synopsis returns a one-line
func (cmd *JobLogs) Synopsis() string { return "Return logs for a running job" }

// Usage shows usage
func (cmd *JobLogs) Usage() string { return "nerd job logs [NAME]" }
<|endoftext|>


<commit_before>package horizon

import (
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"github.com/PuerkitoBio/throttled"
	"github.com/PuerkitoBio/throttled/store"
	"github.com/rcrowley/go-metrics"
	"github.com/rs/cors"
	"github.com/sebest/xff"
	"github.com/stellar/go-horizon/db"
	"github.com/stellar/go-horizon/render/problem"
	"github.com/zenazn/goji/web"
	"github.com/zenazn/goji/web/middleware"
)

// Web contains the http server related fields for go-horizon: the router,
// rate limiter, etc.
type Web struct {
	router      *web.Mux
	rateLimiter *throttled.Throttler

	requestTimer metrics.Timer
	failureMeter metrics.Meter
	successMeter metrics.Meter
}

// initWeb installed a new Web instance onto the provided app object.
func initWeb(app *App) {
	app.web = &Web{
		router:       web.New(),
		requestTimer: metrics.NewTimer(),
		failureMeter: metrics.NewMeter(),
		successMeter: metrics.NewMeter(),
	}

	// register problems
	problem.RegisterError(db.ErrNoResults, problem.NotFound)
}

// initWebMiddleware installs the middleware stack used for go-horizon onto the
// provided app.
func initWebMiddleware(app *App) {

	r := app.web.router
	r.Use(stripTrailingSlashMiddleware())
	r.Use(middleware.EnvInit)
	r.Use(app.Middleware)
	r.Use(middleware.RequestID)
	r.Use(contextMiddleware(app.ctx))
	r.Use(xff.Handler)
	r.Use(LoggerMiddleware)
	r.Use(requestMetricsMiddleware)
	r.Use(RecoverMiddleware)
	r.Use(middleware.AutomaticOptions)

	c := cors.New(cors.Options{
		AllowedOrigins: []string{"*"},
	})
	r.Use(c.Handler)

	r.Use(app.web.RateLimitMiddleware)
}

// initWebActions installs the routing configuration of go-horizon onto the
// provided app.  All route registration should be implemented here.
func initWebActions(app *App) {
	r := app.web.router
	r.Get("/", rootAction)
	r.Get("/metrics", &MetricsAction{})

	// ledger actions
	r.Get("/ledgers", &LedgerIndexAction{})
	r.Get("/ledgers/:id", &LedgerShowAction{})
	r.Get("/ledgers/:ledger_id/transactions", &TransactionIndexAction{})
	r.Get("/ledgers/:ledger_id/operations", &OperationIndexAction{})
	r.Get("/ledgers/:ledger_id/payments", &PaymentsIndexAction{})
	r.Get("/ledgers/:ledger_id/effects", &NotImplementedAction{})

	// account actions
	r.Get("/accounts", &AccountIndexAction{})
	r.Get("/accounts/:id", &AccountShowAction{})
	r.Get("/accounts/:account_id/transactions", &TransactionIndexAction{})
	r.Get("/accounts/:account_id/operations", &OperationIndexAction{})
	r.Get("/accounts/:account_id/payments", &PaymentsIndexAction{})
	r.Get("/accounts/:account_id/effects", &NotImplementedAction{})
	r.Get("/accounts/:account_id/offers", &OffersByAccountAction{})

	// transaction actions
	r.Get("/transactions", &TransactionIndexAction{})
	r.Get("/transactions/:id", &TransactionShowAction{})
	r.Get("/transactions/:tx_id/operations", &OperationIndexAction{})
	r.Get("/transactions/:tx_id/payments", &PaymentsIndexAction{})
	r.Get("/transactions/:tx_id/effects", &NotImplementedAction{})

	// operation actions
	r.Get("/operations", &OperationIndexAction{})
	r.Get("/operations/:id", &OperationShowAction{})
	r.Get("/operations/:op_id/effects", &NotImplementedAction{})

	r.Get("/payments", &PaymentsIndexAction{})

	r.Get("/offers/:id", &NotImplementedAction{})

	// go-horizon doesn't implement everything horizon did,
	// so we reverse proxy if we can
	if app.config.RubyHorizonUrl != "" {

		u, err := url.Parse(app.config.RubyHorizonUrl)
		if err != nil {
			panic("cannot parse ruby-horizon-url")
		}

		rp := httputil.NewSingleHostReverseProxy(u)
		r.Post("/transactions", rp)
		r.Post("/friendbot", rp)
		r.Get("/friendbot", rp)
	} else {
		r.Post("/transactions", &NotImplementedAction{})
		r.Post("/friendbot", &NotImplementedAction{})
		r.Get("/friendbot", &NotImplementedAction{})
	}

	r.NotFound(&NotFoundAction{})
}

func initWebRateLimiter(app *App) {
	rateLimitStore := store.NewMemStore(1000)

	if app.redis != nil {
		rateLimitStore = store.NewRedisStore(app.redis, "throttle:", 0)
	}

	rateLimiter := throttled.RateLimit(
		app.config.RateLimit,
		&throttled.VaryBy{Custom: remoteAddrIP},
		rateLimitStore,
	)

	rateLimiter.DeniedHandler = &RateLimitExceededAction{App: app, Action: Action{}}
	app.web.rateLimiter = rateLimiter
}

func remoteAddrIP(r *http.Request) string {
	ip := strings.SplitN(r.RemoteAddr, ":", 2)[0]
	return ip
}

func init() {
	appInit.Add(
		"web.init",
		initWeb,

		"app-context",
	)

	appInit.Add(
		"web.rate-limiter",
		initWebRateLimiter,

		"web.init",
	)
	appInit.Add(
		"web.middleware",
		initWebMiddleware,

		"web.init",
		"web.rate-limiter",
		"web.metrics",
	)
	appInit.Add(
		"web.actions",
		initWebActions,

		"web.init",
	)
}
<commit_msg>add AllowedHeaders to cors<commit_after>package horizon

import (
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"github.com/PuerkitoBio/throttled"
	"github.com/PuerkitoBio/throttled/store"
	"github.com/rcrowley/go-metrics"
	"github.com/rs/cors"
	"github.com/sebest/xff"
	"github.com/stellar/go-horizon/db"
	"github.com/stellar/go-horizon/render/problem"
	"github.com/zenazn/goji/web"
	"github.com/zenazn/goji/web/middleware"
)

// Web contains the http server related fields for go-horizon: the router,
// rate limiter, etc.
type Web struct {
	router      *web.Mux
	rateLimiter *throttled.Throttler

	requestTimer metrics.Timer
	failureMeter metrics.Meter
	successMeter metrics.Meter
}

// initWeb installed a new Web instance onto the provided app object.
func initWeb(app *App) {
	app.web = &Web{
		router:       web.New(),
		requestTimer: metrics.NewTimer(),
		failureMeter: metrics.NewMeter(),
		successMeter: metrics.NewMeter(),
	}

	// register problems
	problem.RegisterError(db.ErrNoResults, problem.NotFound)
}

// initWebMiddleware installs the middleware stack used for go-horizon onto the
// provided app.
func initWebMiddleware(app *App) {

	r := app.web.router
	r.Use(stripTrailingSlashMiddleware())
	r.Use(middleware.EnvInit)
	r.Use(app.Middleware)
	r.Use(middleware.RequestID)
	r.Use(contextMiddleware(app.ctx))
	r.Use(xff.Handler)
	r.Use(LoggerMiddleware)
	r.Use(requestMetricsMiddleware)
	r.Use(RecoverMiddleware)
	r.Use(middleware.AutomaticOptions)

	c := cors.New(cors.Options{
		AllowedOrigins: []string{"*"},
		AllowedHeaders: []string{"*"},
	})
	r.Use(c.Handler)

	r.Use(app.web.RateLimitMiddleware)
}

// initWebActions installs the routing configuration of go-horizon onto the
// provided app.  All route registration should be implemented here.
func initWebActions(app *App) {
	r := app.web.router
	r.Get("/", rootAction)
	r.Get("/metrics", &MetricsAction{})

	// ledger actions
	r.Get("/ledgers", &LedgerIndexAction{})
	r.Get("/ledgers/:id", &LedgerShowAction{})
	r.Get("/ledgers/:ledger_id/transactions", &TransactionIndexAction{})
	r.Get("/ledgers/:ledger_id/operations", &OperationIndexAction{})
	r.Get("/ledgers/:ledger_id/payments", &PaymentsIndexAction{})
	r.Get("/ledgers/:ledger_id/effects", &NotImplementedAction{})

	// account actions
	r.Get("/accounts", &AccountIndexAction{})
	r.Get("/accounts/:id", &AccountShowAction{})
	r.Get("/accounts/:account_id/transactions", &TransactionIndexAction{})
	r.Get("/accounts/:account_id/operations", &OperationIndexAction{})
	r.Get("/accounts/:account_id/payments", &PaymentsIndexAction{})
	r.Get("/accounts/:account_id/effects", &NotImplementedAction{})
	r.Get("/accounts/:account_id/offers", &OffersByAccountAction{})

	// transaction actions
	r.Get("/transactions", &TransactionIndexAction{})
	r.Get("/transactions/:id", &TransactionShowAction{})
	r.Get("/transactions/:tx_id/operations", &OperationIndexAction{})
	r.Get("/transactions/:tx_id/payments", &PaymentsIndexAction{})
	r.Get("/transactions/:tx_id/effects", &NotImplementedAction{})

	// operation actions
	r.Get("/operations", &OperationIndexAction{})
	r.Get("/operations/:id", &OperationShowAction{})
	r.Get("/operations/:op_id/effects", &NotImplementedAction{})

	r.Get("/payments", &PaymentsIndexAction{})

	r.Get("/offers/:id", &NotImplementedAction{})

	// go-horizon doesn't implement everything horizon did,
	// so we reverse proxy if we can
	if app.config.RubyHorizonUrl != "" {

		u, err := url.Parse(app.config.RubyHorizonUrl)
		if err != nil {
			panic("cannot parse ruby-horizon-url")
		}

		rp := httputil.NewSingleHostReverseProxy(u)
		r.Post("/transactions", rp)
		r.Post("/friendbot", rp)
		r.Get("/friendbot", rp)
	} else {
		r.Post("/transactions", &NotImplementedAction{})
		r.Post("/friendbot", &NotImplementedAction{})
		r.Get("/friendbot", &NotImplementedAction{})
	}

	r.NotFound(&NotFoundAction{})
}

func initWebRateLimiter(app *App) {
	rateLimitStore := store.NewMemStore(1000)

	if app.redis != nil {
		rateLimitStore = store.NewRedisStore(app.redis, "throttle:", 0)
	}

	rateLimiter := throttled.RateLimit(
		app.config.RateLimit,
		&throttled.VaryBy{Custom: remoteAddrIP},
		rateLimitStore,
	)

	rateLimiter.DeniedHandler = &RateLimitExceededAction{App: app, Action: Action{}}
	app.web.rateLimiter = rateLimiter
}

func remoteAddrIP(r *http.Request) string {
	ip := strings.SplitN(r.RemoteAddr, ":", 2)[0]
	return ip
}

func init() {
	appInit.Add(
		"web.init",
		initWeb,

		"app-context",
	)

	appInit.Add(
		"web.rate-limiter",
		initWebRateLimiter,

		"web.init",
	)
	appInit.Add(
		"web.middleware",
		initWebMiddleware,

		"web.init",
		"web.rate-limiter",
		"web.metrics",
	)
	appInit.Add(
		"web.actions",
		initWebActions,

		"web.init",
	)
}
<|endoftext|>


<commit_before>package goscp

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/cheggaaa/pb"
	"golang.org/x/crypto/ssh"
)

var (
	// SCP messages
	fileCopyRx  = regexp.MustCompile(`C(?P<mode>\d{4}) (?P<length>\d+) (?P<filename>.+)`)
	dirCopyRx   = regexp.MustCompile(`D(?P<mode>\d{4}) (?P<length>\d+) (?P<dirname>.+)`)
	timestampRx = regexp.MustCompile(`T(?P<mtime>\d+) 0 (?P<atime>\d+) 0`)
	endDir      = "E"
)

type Client struct {
	SSHClient        *ssh.Client
	ProgressCallback func(out string)
	DestinationPath  []string

	// Errors that have occurred while communicating with host
	errors []error

	// Verbose output when communicating with host
	Verbose bool

	// Stop transfer on OS error - occurs during filepath.Walk
	StopOnOSError bool

	// Stdin for SSH session
	scpStdinPipe io.WriteCloser

	// Stdout for SSH session
	scpStdoutPipe *Reader
}

// Returns a ssh.Client wrapper.
// DestinationPath is set to the current directory by default.
func NewClient(c *ssh.Client) *Client {
	return &Client{
		SSHClient:       c,
		DestinationPath: []string{"."},
	}
}

// Set where content will be sent
func (c *Client) SetDestinationPath(path string) {
	c.DestinationPath = []string{path}
}

func (c *Client) addError(err error) {
	c.errors = append(c.errors, err)
}

// GetLastError should be queried after a call to Download() or Upload().
func (c *Client) GetLastError() error {
	if len(c.errors) > 0 {
		return c.errors[len(c.errors)-1]
	}
	return nil
}

// GetErrorStack returns all errors that have occurred so far
func (c *Client) GetErrorStack() []error {
	return c.errors
}

// Cancel an ongoing operation
func (c *Client) Cancel() {
	if c.scpStdoutPipe != nil {
		c.scpStdoutPipe.cancel <- struct{}{}
	}
}

// Download remotePath to c.DestinationPath
func (c *Client) Download(remotePath string) {
	session, err := c.SSHClient.NewSession()
	if err != nil {
		c.addError(err)
		return
	}
	defer session.Close()

	go func() {
		c.scpStdinPipe, err = session.StdinPipe()
		if err != nil {
			c.addError(err)
			return
		}
		defer c.scpStdinPipe.Close()

		r, err := session.StdoutPipe()
		if err != nil {
			c.addError(err)
			return
		}

		// Initialise transfer
		c.sendAck()

		// Wrapper to support cancellation
		c.scpStdoutPipe = &Reader{
			Reader: bufio.NewReader(r),
			cancel: make(chan struct{}, 1),
		}

		for {
			c.outputInfo("Reading message from source")
			msg, err := c.scpStdoutPipe.ReadString('\n')
			if err != nil {
				if err != io.EOF {
					c.addError(err)
				}
				return
			}

			// Strip nulls and new lines
			msg = strings.TrimSpace(strings.Trim(msg, "\x00"))
			c.outputInfo(fmt.Sprintf("Received: %s", msg))

			// Confirm message
			c.sendAck()

			switch {
			case c.isFileCopyMsg(msg):
				// Handle incoming file
				err := c.file(msg)
				if err != nil {
					c.addError(err)
					return
				}
			case c.isDirCopyMsg(msg):
				// Handling incoming directory
				err := c.directory(msg)
				if err != nil {
					c.addError(err)
					return
				}
			case msg == endDir:
				// Directory finished, go up a directory
				c.upDirectory()
			case c.isWarningMsg(msg):
				c.addError(fmt.Errorf("Warning message: [%q]\n", msg))
				return
			case c.isErrorMsg(msg):
				c.addError(fmt.Errorf("Error message: [%q]\n", msg))
				return
			default:
				c.addError(fmt.Errorf("Unhandled message: [%q]\n", msg))
				return
			}

			// Confirm message
			c.sendAck()
		}
	}()

	cmd := fmt.Sprintf("scp -rf %s", remotePath)
	if err := session.Run(cmd); err != nil {
		c.addError(err)
		return
	}

	return
}

// Upload localPath to c.DestinationPath
func (c *Client) Upload(localPath string) {
	session, err := c.SSHClient.NewSession()
	if err != nil {
		c.addError(err)
		return
	}
	defer session.Close()

	go func() {
		c.scpStdinPipe, err = session.StdinPipe()
		if err != nil {
			c.addError(err)
			return
		}
		defer c.scpStdinPipe.Close()

		r, err := session.StdoutPipe()
		if err != nil {
			c.addError(err)
			return
		}

		// Wrapper to support cancellation
		c.scpStdoutPipe = &Reader{
			Reader: bufio.NewReader(r),
			cancel: make(chan struct{}, 1),
		}

		// This has already been used in the cmd call below
		// so it can be reused for 'end of directory' message handling
		c.DestinationPath = []string{}

		err = filepath.Walk(localPath, c.handleItem)
		if err != nil {
			c.addError(err)
			return
		}

		// End transfer
		paths := strings.Split(c.DestinationPath[0], "/")
		for range paths {
			c.sendEndOfDirectoryMessage()
		}
	}()

	cmd := fmt.Sprintf("scp -rt %s", filepath.Join(c.DestinationPath...))
	if err := session.Run(cmd); err != nil {
		c.addError(err)
		return
	}

	return
}

// Send an acknowledgement message
func (c *Client) sendAck() {
	fmt.Fprint(c.scpStdinPipe, "\x00")
}

// Send an error message
func (c *Client) sendErr() {
	fmt.Fprint(c.scpStdinPipe, "\x02")
}

// Check if an incoming message is a file copy message
func (c *Client) isFileCopyMsg(s string) bool {
	return strings.HasPrefix(s, "C")
}

// Check if an incoming message is a directory copy message
func (c *Client) isDirCopyMsg(s string) bool {
	return strings.HasPrefix(s, "D")
}

// Check if an incoming message is a warning
func (c *Client) isWarningMsg(s string) bool {
	return strings.HasPrefix(s, "\x01")
}

// Check if an incoming message is an error
func (c *Client) isErrorMsg(s string) bool {
	return strings.HasPrefix(s, "\x02")
}

// Send a directory message while in source mode
func (c *Client) sendDirectoryMessage(mode os.FileMode, dirname string) {
	msg := fmt.Sprintf("D0%o 0 %s", mode, dirname)
	fmt.Fprintln(c.scpStdinPipe, msg)
	c.outputInfo(fmt.Sprintf("Sent: %s", msg))
}

// Send a end of directory message while in source mode
func (c *Client) sendEndOfDirectoryMessage() {
	msg := endDir
	fmt.Fprintln(c.scpStdinPipe, msg)
	c.outputInfo(fmt.Sprintf("Sent: %s", msg))
}

// Send a file message while in source mode
func (c *Client) sendFileMessage(mode os.FileMode, size int64, filename string) {
	msg := fmt.Sprintf("C0%o %d %s", mode, size, filename)
	fmt.Fprintln(c.scpStdinPipe, msg)
	c.outputInfo(fmt.Sprintf("Sent: %s", msg))
}

// Handle directory copy message in sink mode
func (c *Client) directory(msg string) error {
	parts, err := c.parseMessage(msg, dirCopyRx)
	if err != nil {
		return err
	}

	err = os.Mkdir(filepath.Join(c.DestinationPath...)+string(filepath.Separator)+parts["dirname"], 0755)
	if err != nil {
		return err
	}

	// Traverse into directory
	c.DestinationPath = append(c.DestinationPath, parts["dirname"])

	return nil
}

// Handle file copy message in sink mode
func (c *Client) file(msg string) error {
	parts, err := c.parseMessage(msg, fileCopyRx)
	if err != nil {
		return err
	}

	fileLen, _ := strconv.Atoi(parts["length"])

	// Create local file
	localFile, err := os.Create(filepath.Join(c.DestinationPath...) + string(filepath.Separator) + parts["filename"])
	if err != nil {
		return err
	}
	defer localFile.Close()

	bar := c.NewProgressBar(fileLen)
	bar.Start()
	defer bar.Finish()

	mw := io.MultiWriter(localFile, bar)
	if n, err := io.CopyN(mw, c.scpStdoutPipe, int64(fileLen)); err != nil || n < int64(fileLen) {
		c.sendErr()
		return err
	}

	return nil
}

// Break down incoming protocol messages
func (c *Client) parseMessage(msg string, rx *regexp.Regexp) (map[string]string, error) {
	parts := make(map[string]string)
	matches := rx.FindStringSubmatch(msg)
	if len(matches) == 0 {
		return parts, errors.New("Could not parse protocol message: " + msg)
	}

	for i, name := range rx.SubexpNames() {
		parts[name] = matches[i]
	}
	return parts, nil
}

// Go back up one directory
func (c *Client) upDirectory() {
	c.DestinationPath = c.DestinationPath[:len(c.DestinationPath)-1]
}

// Handle each item coming through filepath.Walk
func (c *Client) handleItem(path string, info os.FileInfo, err error) error {
	if err != nil {
		// OS error
		c.outputInfo(fmt.Sprintf("Item error: %s", err))

		if c.StopOnOSError {
			return err
		}
		return nil
	}

	if info.IsDir() {
		// Handle directories
		if len(c.DestinationPath) != 0 {
			// If not first directory
			currentPath := strings.Split(c.DestinationPath[0], "/")
			newPath := strings.Split(path, "/")

			// <= slashes = going back up
			if len(newPath) <= len(currentPath) {
				// Send EOD messages for the amount of directories we go up
				for i := len(newPath) - 1; i < len(currentPath); i++ {
					c.sendEndOfDirectoryMessage()
				}
			}
		}
		c.DestinationPath = []string{path}
		c.sendDirectoryMessage(0644, filepath.Base(path))
	} else {
		// Handle regular files
		targetItem, err := os.Open(path)
		if err != nil {
			return err
		}

		c.sendFileMessage(0644, info.Size(), filepath.Base(path))

		if info.Size() > 0 {
			bar := c.NewProgressBar(int(info.Size()))
			bar.Start()
			defer bar.Finish()

			mw := io.MultiWriter(c.scpStdinPipe, bar)

			c.outputInfo(fmt.Sprintf("Sending file: %s", path))
			if _, err := io.Copy(mw, targetItem); err != nil {
				c.sendErr()
				return err
			}

			c.sendAck()
		} else {
			c.outputInfo(fmt.Sprintf("Sending empty file: %s", path))
			c.sendAck()
		}
	}

	return nil
}

func (c *Client) outputInfo(s ...string) {
	if c.Verbose {
		log.Println(s)
	}
}

// Create progress bar
func (c *Client) NewProgressBar(fileLength int) *pb.ProgressBar {
	bar := pb.New(fileLength)
	bar.Callback = c.ProgressCallback
	bar.ShowSpeed = true
	bar.ShowTimeLeft = true
	bar.ShowCounters = true
	bar.Units = pb.U_BYTES
	bar.SetRefreshRate(time.Second)
	bar.SetWidth(80)
	bar.SetMaxWidth(80)

	return bar
}

// Wrapper to support cancellation
type Reader struct {
	*bufio.Reader

	// Cancel an ongoing transfer
	cancel chan struct{}
}

// Additional cancellation check
func (r *Reader) Read(p []byte) (n int, err error) {
	select {
	case <-r.cancel:
		log.Println("CANCELLED!!")
		return 0, errors.New("Transfer cancelled")
	default:
		return r.Reader.Read(p)
	}
}
<commit_msg>Removed debug message.<commit_after>package goscp

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/cheggaaa/pb"
	"golang.org/x/crypto/ssh"
)

var (
	// SCP messages
	fileCopyRx  = regexp.MustCompile(`C(?P<mode>\d{4}) (?P<length>\d+) (?P<filename>.+)`)
	dirCopyRx   = regexp.MustCompile(`D(?P<mode>\d{4}) (?P<length>\d+) (?P<dirname>.+)`)
	timestampRx = regexp.MustCompile(`T(?P<mtime>\d+) 0 (?P<atime>\d+) 0`)
	endDir      = "E"
)

type Client struct {
	SSHClient        *ssh.Client
	ProgressCallback func(out string)
	DestinationPath  []string

	// Errors that have occurred while communicating with host
	errors []error

	// Verbose output when communicating with host
	Verbose bool

	// Stop transfer on OS error - occurs during filepath.Walk
	StopOnOSError bool

	// Stdin for SSH session
	scpStdinPipe io.WriteCloser

	// Stdout for SSH session
	scpStdoutPipe *Reader
}

// Returns a ssh.Client wrapper.
// DestinationPath is set to the current directory by default.
func NewClient(c *ssh.Client) *Client {
	return &Client{
		SSHClient:       c,
		DestinationPath: []string{"."},
	}
}

// Set where content will be sent
func (c *Client) SetDestinationPath(path string) {
	c.DestinationPath = []string{path}
}

func (c *Client) addError(err error) {
	c.errors = append(c.errors, err)
}

// GetLastError should be queried after a call to Download() or Upload().
func (c *Client) GetLastError() error {
	if len(c.errors) > 0 {
		return c.errors[len(c.errors)-1]
	}
	return nil
}

// GetErrorStack returns all errors that have occurred so far
func (c *Client) GetErrorStack() []error {
	return c.errors
}

// Cancel an ongoing operation
func (c *Client) Cancel() {
	if c.scpStdoutPipe != nil {
		c.scpStdoutPipe.cancel <- struct{}{}
	}
}

// Download remotePath to c.DestinationPath
func (c *Client) Download(remotePath string) {
	session, err := c.SSHClient.NewSession()
	if err != nil {
		c.addError(err)
		return
	}
	defer session.Close()

	go func() {
		c.scpStdinPipe, err = session.StdinPipe()
		if err != nil {
			c.addError(err)
			return
		}
		defer c.scpStdinPipe.Close()

		r, err := session.StdoutPipe()
		if err != nil {
			c.addError(err)
			return
		}

		// Initialise transfer
		c.sendAck()

		// Wrapper to support cancellation
		c.scpStdoutPipe = &Reader{
			Reader: bufio.NewReader(r),
			cancel: make(chan struct{}, 1),
		}

		for {
			c.outputInfo("Reading message from source")
			msg, err := c.scpStdoutPipe.ReadString('\n')
			if err != nil {
				if err != io.EOF {
					c.addError(err)
				}
				return
			}

			// Strip nulls and new lines
			msg = strings.TrimSpace(strings.Trim(msg, "\x00"))
			c.outputInfo(fmt.Sprintf("Received: %s", msg))

			// Confirm message
			c.sendAck()

			switch {
			case c.isFileCopyMsg(msg):
				// Handle incoming file
				err := c.file(msg)
				if err != nil {
					c.addError(err)
					return
				}
			case c.isDirCopyMsg(msg):
				// Handling incoming directory
				err := c.directory(msg)
				if err != nil {
					c.addError(err)
					return
				}
			case msg == endDir:
				// Directory finished, go up a directory
				c.upDirectory()
			case c.isWarningMsg(msg):
				c.addError(fmt.Errorf("Warning message: [%q]\n", msg))
				return
			case c.isErrorMsg(msg):
				c.addError(fmt.Errorf("Error message: [%q]\n", msg))
				return
			default:
				c.addError(fmt.Errorf("Unhandled message: [%q]\n", msg))
				return
			}

			// Confirm message
			c.sendAck()
		}
	}()

	cmd := fmt.Sprintf("scp -rf %s", remotePath)
	if err := session.Run(cmd); err != nil {
		c.addError(err)
		return
	}

	return
}

// Upload localPath to c.DestinationPath
func (c *Client) Upload(localPath string) {
	session, err := c.SSHClient.NewSession()
	if err != nil {
		c.addError(err)
		return
	}
	defer session.Close()

	go func() {
		c.scpStdinPipe, err = session.StdinPipe()
		if err != nil {
			c.addError(err)
			return
		}
		defer c.scpStdinPipe.Close()

		r, err := session.StdoutPipe()
		if err != nil {
			c.addError(err)
			return
		}

		// Wrapper to support cancellation
		c.scpStdoutPipe = &Reader{
			Reader: bufio.NewReader(r),
			cancel: make(chan struct{}, 1),
		}

		// This has already been used in the cmd call below
		// so it can be reused for 'end of directory' message handling
		c.DestinationPath = []string{}

		err = filepath.Walk(localPath, c.handleItem)
		if err != nil {
			c.addError(err)
			return
		}

		// End transfer
		paths := strings.Split(c.DestinationPath[0], "/")
		for range paths {
			c.sendEndOfDirectoryMessage()
		}
	}()

	cmd := fmt.Sprintf("scp -rt %s", filepath.Join(c.DestinationPath...))
	if err := session.Run(cmd); err != nil {
		c.addError(err)
		return
	}

	return
}

// Send an acknowledgement message
func (c *Client) sendAck() {
	fmt.Fprint(c.scpStdinPipe, "\x00")
}

// Send an error message
func (c *Client) sendErr() {
	fmt.Fprint(c.scpStdinPipe, "\x02")
}

// Check if an incoming message is a file copy message
func (c *Client) isFileCopyMsg(s string) bool {
	return strings.HasPrefix(s, "C")
}

// Check if an incoming message is a directory copy message
func (c *Client) isDirCopyMsg(s string) bool {
	return strings.HasPrefix(s, "D")
}

// Check if an incoming message is a warning
func (c *Client) isWarningMsg(s string) bool {
	return strings.HasPrefix(s, "\x01")
}

// Check if an incoming message is an error
func (c *Client) isErrorMsg(s string) bool {
	return strings.HasPrefix(s, "\x02")
}

// Send a directory message while in source mode
func (c *Client) sendDirectoryMessage(mode os.FileMode, dirname string) {
	msg := fmt.Sprintf("D0%o 0 %s", mode, dirname)
	fmt.Fprintln(c.scpStdinPipe, msg)
	c.outputInfo(fmt.Sprintf("Sent: %s", msg))
}

// Send a end of directory message while in source mode
func (c *Client) sendEndOfDirectoryMessage() {
	msg := endDir
	fmt.Fprintln(c.scpStdinPipe, msg)
	c.outputInfo(fmt.Sprintf("Sent: %s", msg))
}

// Send a file message while in source mode
func (c *Client) sendFileMessage(mode os.FileMode, size int64, filename string) {
	msg := fmt.Sprintf("C0%o %d %s", mode, size, filename)
	fmt.Fprintln(c.scpStdinPipe, msg)
	c.outputInfo(fmt.Sprintf("Sent: %s", msg))
}

// Handle directory copy message in sink mode
func (c *Client) directory(msg string) error {
	parts, err := c.parseMessage(msg, dirCopyRx)
	if err != nil {
		return err
	}

	err = os.Mkdir(filepath.Join(c.DestinationPath...)+string(filepath.Separator)+parts["dirname"], 0755)
	if err != nil {
		return err
	}

	// Traverse into directory
	c.DestinationPath = append(c.DestinationPath, parts["dirname"])

	return nil
}

// Handle file copy message in sink mode
func (c *Client) file(msg string) error {
	parts, err := c.parseMessage(msg, fileCopyRx)
	if err != nil {
		return err
	}

	fileLen, _ := strconv.Atoi(parts["length"])

	// Create local file
	localFile, err := os.Create(filepath.Join(c.DestinationPath...) + string(filepath.Separator) + parts["filename"])
	if err != nil {
		return err
	}
	defer localFile.Close()

	bar := c.NewProgressBar(fileLen)
	bar.Start()
	defer bar.Finish()

	mw := io.MultiWriter(localFile, bar)
	if n, err := io.CopyN(mw, c.scpStdoutPipe, int64(fileLen)); err != nil || n < int64(fileLen) {
		c.sendErr()
		return err
	}

	return nil
}

// Break down incoming protocol messages
func (c *Client) parseMessage(msg string, rx *regexp.Regexp) (map[string]string, error) {
	parts := make(map[string]string)
	matches := rx.FindStringSubmatch(msg)
	if len(matches) == 0 {
		return parts, errors.New("Could not parse protocol message: " + msg)
	}

	for i, name := range rx.SubexpNames() {
		parts[name] = matches[i]
	}
	return parts, nil
}

// Go back up one directory
func (c *Client) upDirectory() {
	c.DestinationPath = c.DestinationPath[:len(c.DestinationPath)-1]
}

// Handle each item coming through filepath.Walk
func (c *Client) handleItem(path string, info os.FileInfo, err error) error {
	if err != nil {
		// OS error
		c.outputInfo(fmt.Sprintf("Item error: %s", err))

		if c.StopOnOSError {
			return err
		}
		return nil
	}

	if info.IsDir() {
		// Handle directories
		if len(c.DestinationPath) != 0 {
			// If not first directory
			currentPath := strings.Split(c.DestinationPath[0], "/")
			newPath := strings.Split(path, "/")

			// <= slashes = going back up
			if len(newPath) <= len(currentPath) {
				// Send EOD messages for the amount of directories we go up
				for i := len(newPath) - 1; i < len(currentPath); i++ {
					c.sendEndOfDirectoryMessage()
				}
			}
		}
		c.DestinationPath = []string{path}
		c.sendDirectoryMessage(0644, filepath.Base(path))
	} else {
		// Handle regular files
		targetItem, err := os.Open(path)
		if err != nil {
			return err
		}

		c.sendFileMessage(0644, info.Size(), filepath.Base(path))

		if info.Size() > 0 {
			bar := c.NewProgressBar(int(info.Size()))
			bar.Start()
			defer bar.Finish()

			mw := io.MultiWriter(c.scpStdinPipe, bar)

			c.outputInfo(fmt.Sprintf("Sending file: %s", path))
			if _, err := io.Copy(mw, targetItem); err != nil {
				c.sendErr()
				return err
			}

			c.sendAck()
		} else {
			c.outputInfo(fmt.Sprintf("Sending empty file: %s", path))
			c.sendAck()
		}
	}

	return nil
}

func (c *Client) outputInfo(s ...string) {
	if c.Verbose {
		log.Println(s)
	}
}

// Create progress bar
func (c *Client) NewProgressBar(fileLength int) *pb.ProgressBar {
	bar := pb.New(fileLength)
	bar.Callback = c.ProgressCallback
	bar.ShowSpeed = true
	bar.ShowTimeLeft = true
	bar.ShowCounters = true
	bar.Units = pb.U_BYTES
	bar.SetRefreshRate(time.Second)
	bar.SetWidth(80)
	bar.SetMaxWidth(80)

	return bar
}

// Wrapper to support cancellation
type Reader struct {
	*bufio.Reader

	// Cancel an ongoing transfer
	cancel chan struct{}
}

// Additional cancellation check
func (r *Reader) Read(p []byte) (n int, err error) {
	select {
	case <-r.cancel:
		return 0, errors.New("Transfer cancelled")
	default:
		return r.Reader.Read(p)
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 Alex Browne.  All rights reserved.
// Use of this source code is governed by the MIT
// license, which can be found in the LICENSE file.

// File model.go contains code strictly related to DefaultData and Model.
// The Register() method and associated methods are also included here.

package zoom

import (
	"errors"
	"fmt"
	"github.com/stephenalexbrowne/zoom/util"
	"reflect"
)

// DefaultData should be embedded in any struct you wish to save.
// It includes all the required fields.
type DefaultData struct {
	Id string `redis:"-"`
	// TODO: add other default fields?
}

// Model is an interface encapsulating anything that can be saved.
// Any struct which includes an embedded DefaultData field satisfies
// the Model interface.
type Model interface {
	GetId() string
	SetId(string)
	// TODO: add getters and setters for other default fields?
}

type modelSpec struct {
	fieldNames []string
	sets       []*externalSet
	lists      []*externalList
	relations  map[string]relation
}

type externalSet struct {
	redisName string
	fieldName string
}

type externalList struct {
	redisName string
	fieldName string
}

type relation struct {
	redisName string
	fieldName string
	typ       relationType
}

type relationType int

const (
	ONE_TO_ONE = iota
	ONE_TO_MANY
)

// maps a type to a string identifier. The string is used
// as a key in the redis database.
var typeToName map[reflect.Type]string = make(map[reflect.Type]string)

// maps a string identifier to a type. This is so you can
// pass in a string for the *ById methods
var nameToType map[string]reflect.Type = make(map[string]reflect.Type)

// maps a string identifier to a modelSpec
var modelSpecs map[string]*modelSpec = make(map[string]*modelSpec)

// methods so that DefaultData (and any struct with DefaultData embedded)
// satisifies Model interface
func (d *DefaultData) GetId() string {
	return d.Id
}

func (d *DefaultData) SetId(id string) {
	d.Id = id
}

// Register adds a type to the list of registered types. Any struct
// you wish to save must be registered first. Both name and type of in
// must be unique, i.e. not already registered.
func Register(in interface{}, name string) error {
	typ := reflect.TypeOf(in)

	// make sure the interface is the correct type
	if typ.Kind() != reflect.Ptr {
		return errors.New("zoom: schema must be a pointer to a struct")
	} else if typ.Elem().Kind() != reflect.Struct {
		return errors.New("zoom: schema must be a pointer to a struct")
	}

	// make sure the name and type have not been previously registered
	if alreadyRegisteredType(typ) {
		return NewTypeAlreadyRegisteredError(typ)
	}
	if alreadyRegisteredName(name) {
		return NewNameAlreadyRegisteredError(name)
	}

	// create a new model spec and register its lists and sets
	ms := &modelSpec{relations: make(map[string]relation)}
	if err := compileModelSpec(typ, ms); err != nil {
		return err
	}

	typeToName[typ] = name
	nameToType[name] = typ
	modelSpecs[name] = ms

	return nil
}

func compileModelSpec(typ reflect.Type, ms *modelSpec) error {
	// iterate through fields to find slices and arrays
	elem := typ.Elem()
	numFields := elem.NumField()
	for i := 0; i < numFields; i++ {
		field := elem.Field(i)
		if field.Name != "DefaultData" {
			ms.fieldNames = append(ms.fieldNames, field.Name)
		}
		if util.TypeIsPointerToStruct(field.Type) {
			// assume we're dealing with a one-to-one relation
			// get the redisName
			tag := field.Tag
			redisName := tag.Get("redis")
			if redisName == "-" {
				continue // skip field
			} else if redisName == "" {
				redisName = field.Name
			}
			ms.relations[field.Name] = relation{
				redisName: redisName,
				fieldName: field.Name,
				typ:       ONE_TO_ONE,
			}
		} else if util.TypeIsSliceOrArray(field.Type) {
			// we're dealing with a slice or an array, which should be converted to a list, set, or one-to-many relation
			tag := field.Tag
			redisName := tag.Get("redis")
			if redisName == "-" {
				continue // skip field
			} else if redisName == "" {
				redisName = field.Name
			}
			if util.TypeIsPointerToStruct(field.Type.Elem()) {
				// assume we're dealing with a one-to-many relation
				ms.relations[field.Name] = relation{
					redisName: redisName,
					fieldName: field.Name,
					typ:       ONE_TO_MANY,
				}
				continue
			}
			redisType := tag.Get("redisType")
			if redisType == "" || redisType == "list" {
				ms.lists = append(ms.lists, &externalList{redisName: redisName, fieldName: field.Name})
			} else if redisType == "set" {
				ms.sets = append(ms.sets, &externalSet{redisName: redisName, fieldName: field.Name})
			} else {
				msg := fmt.Sprintf("zoom: invalid struct tag for redisType: %s. must be either 'set' or 'list'\n", redisType)
				return errors.New(msg)
			}
		}
	}
	return nil
}

// UnregisterName removes a type (identified by name) from the list of
// registered types. You only need to call UnregisterName or UnregisterType,
// not both.
func UnregisterName(name string) error {
	typ, ok := nameToType[name]
	if !ok {
		return NewModelNameNotRegisteredError(name)
	}
	delete(nameToType, name)
	delete(typeToName, typ)
	return nil
}

// UnregisterName removes a type from the list of registered types.
// You only need to call UnregisterName or UnregisterType, not both.
func UnregisterType(typ reflect.Type) error {
	name, ok := typeToName[typ]
	if !ok {
		return NewModelTypeNotRegisteredError(typ)
	}
	delete(nameToType, name)
	delete(typeToName, typ)
	return nil
}

// alreadyRegisteredName returns true iff the model name has already been registered
func alreadyRegisteredName(n string) bool {
	_, ok := nameToType[n]
	return ok
}

// alreadyRegisteredType returns true iff the model type has already been registered
func alreadyRegisteredType(t reflect.Type) bool {
	_, ok := typeToName[t]
	return ok
}

// getRegisteredNameFromInterface gets the registered name of the model we're
// trying to save based on the interfaces type. If the interface's name/type
// has not been registered, returns a ModelTypeNotRegisteredError
func getRegisteredNameFromInterface(in interface{}) (string, error) {
	typ := reflect.TypeOf(in)
	name, ok := typeToName[typ]
	if !ok {
		return "", NewModelTypeNotRegisteredError(typ)
	}
	return name, nil
}

// getRegisteredTypeFromName gets the registered type of the model we're trying
// to save based on the model name. If the interface's name/type has not been registered,
// returns a ModelNameNotRegisteredError
func getRegisteredTypeFromName(name string) (reflect.Type, error) {
	typ, ok := nameToType[name]
	if !ok {
		return nil, NewModelNameNotRegisteredError(name)
	}
	return typ, nil
}
<commit_msg>GetId now uses a non-pointer DefaultData receiver.<commit_after>// Copyright 2013 Alex Browne.  All rights reserved.
// Use of this source code is governed by the MIT
// license, which can be found in the LICENSE file.

// File model.go contains code strictly related to DefaultData and Model.
// The Register() method and associated methods are also included here.

package zoom

import (
	"errors"
	"fmt"
	"github.com/stephenalexbrowne/zoom/util"
	"reflect"
)

// DefaultData should be embedded in any struct you wish to save.
// It includes all the required fields.
type DefaultData struct {
	Id string `redis:"-"`
	// TODO: add other default fields?
}

// Model is an interface encapsulating anything that can be saved.
// Any struct which includes an embedded DefaultData field satisfies
// the Model interface.
type Model interface {
	GetId() string
	SetId(string)
	// TODO: add getters and setters for other default fields?
}

type modelSpec struct {
	fieldNames []string
	sets       []*externalSet
	lists      []*externalList
	relations  map[string]relation
}

type externalSet struct {
	redisName string
	fieldName string
}

type externalList struct {
	redisName string
	fieldName string
}

type relation struct {
	redisName string
	fieldName string
	typ       relationType
}

type relationType int

const (
	ONE_TO_ONE = iota
	ONE_TO_MANY
)

// maps a type to a string identifier. The string is used
// as a key in the redis database.
var typeToName map[reflect.Type]string = make(map[reflect.Type]string)

// maps a string identifier to a type. This is so you can
// pass in a string for the *ById methods
var nameToType map[string]reflect.Type = make(map[string]reflect.Type)

// maps a string identifier to a modelSpec
var modelSpecs map[string]*modelSpec = make(map[string]*modelSpec)

// methods so that DefaultData (and any struct with DefaultData embedded)
// satisifies Model interface
func (d DefaultData) GetId() string {
	return d.Id
}

func (d *DefaultData) SetId(id string) {
	d.Id = id
}

// Register adds a type to the list of registered types. Any struct
// you wish to save must be registered first. Both name and type of in
// must be unique, i.e. not already registered.
func Register(in interface{}, name string) error {
	typ := reflect.TypeOf(in)

	// make sure the interface is the correct type
	if typ.Kind() != reflect.Ptr {
		return errors.New("zoom: schema must be a pointer to a struct")
	} else if typ.Elem().Kind() != reflect.Struct {
		return errors.New("zoom: schema must be a pointer to a struct")
	}

	// make sure the name and type have not been previously registered
	if alreadyRegisteredType(typ) {
		return NewTypeAlreadyRegisteredError(typ)
	}
	if alreadyRegisteredName(name) {
		return NewNameAlreadyRegisteredError(name)
	}

	// create a new model spec and register its lists and sets
	ms := &modelSpec{relations: make(map[string]relation)}
	if err := compileModelSpec(typ, ms); err != nil {
		return err
	}

	typeToName[typ] = name
	nameToType[name] = typ
	modelSpecs[name] = ms

	return nil
}

func compileModelSpec(typ reflect.Type, ms *modelSpec) error {
	// iterate through fields to find slices and arrays
	elem := typ.Elem()
	numFields := elem.NumField()
	for i := 0; i < numFields; i++ {
		field := elem.Field(i)
		if field.Name != "DefaultData" {
			ms.fieldNames = append(ms.fieldNames, field.Name)
		}
		if util.TypeIsPointerToStruct(field.Type) {
			// assume we're dealing with a one-to-one relation
			// get the redisName
			tag := field.Tag
			redisName := tag.Get("redis")
			if redisName == "-" {
				continue // skip field
			} else if redisName == "" {
				redisName = field.Name
			}
			ms.relations[field.Name] = relation{
				redisName: redisName,
				fieldName: field.Name,
				typ:       ONE_TO_ONE,
			}
		} else if util.TypeIsSliceOrArray(field.Type) {
			// we're dealing with a slice or an array, which should be converted to a list, set, or one-to-many relation
			tag := field.Tag
			redisName := tag.Get("redis")
			if redisName == "-" {
				continue // skip field
			} else if redisName == "" {
				redisName = field.Name
			}
			if util.TypeIsPointerToStruct(field.Type.Elem()) {
				// assume we're dealing with a one-to-many relation
				ms.relations[field.Name] = relation{
					redisName: redisName,
					fieldName: field.Name,
					typ:       ONE_TO_MANY,
				}
				continue
			}
			redisType := tag.Get("redisType")
			if redisType == "" || redisType == "list" {
				ms.lists = append(ms.lists, &externalList{redisName: redisName, fieldName: field.Name})
			} else if redisType == "set" {
				ms.sets = append(ms.sets, &externalSet{redisName: redisName, fieldName: field.Name})
			} else {
				msg := fmt.Sprintf("zoom: invalid struct tag for redisType: %s. must be either 'set' or 'list'\n", redisType)
				return errors.New(msg)
			}
		}
	}
	return nil
}

// UnregisterName removes a type (identified by name) from the list of
// registered types. You only need to call UnregisterName or UnregisterType,
// not both.
func UnregisterName(name string) error {
	typ, ok := nameToType[name]
	if !ok {
		return NewModelNameNotRegisteredError(name)
	}
	delete(nameToType, name)
	delete(typeToName, typ)
	return nil
}

// UnregisterName removes a type from the list of registered types.
// You only need to call UnregisterName or UnregisterType, not both.
func UnregisterType(typ reflect.Type) error {
	name, ok := typeToName[typ]
	if !ok {
		return NewModelTypeNotRegisteredError(typ)
	}
	delete(nameToType, name)
	delete(typeToName, typ)
	return nil
}

// alreadyRegisteredName returns true iff the model name has already been registered
func alreadyRegisteredName(n string) bool {
	_, ok := nameToType[n]
	return ok
}

// alreadyRegisteredType returns true iff the model type has already been registered
func alreadyRegisteredType(t reflect.Type) bool {
	_, ok := typeToName[t]
	return ok
}

// getRegisteredNameFromInterface gets the registered name of the model we're
// trying to save based on the interfaces type. If the interface's name/type
// has not been registered, returns a ModelTypeNotRegisteredError
func getRegisteredNameFromInterface(in interface{}) (string, error) {
	typ := reflect.TypeOf(in)
	name, ok := typeToName[typ]
	if !ok {
		return "", NewModelTypeNotRegisteredError(typ)
	}
	return name, nil
}

// getRegisteredTypeFromName gets the registered type of the model we're trying
// to save based on the model name. If the interface's name/type has not been registered,
// returns a ModelNameNotRegisteredError
func getRegisteredTypeFromName(name string) (reflect.Type, error) {
	typ, ok := nameToType[name]
	if !ok {
		return nil, NewModelNameNotRegisteredError(name)
	}
	return typ, nil
}
<|endoftext|>