<commit_before>package integration_test

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"runtime"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
	"github.com/onsi/gomega/ghttp"

	"github.com/concourse/atc"
)

var _ = Describe("login -k Command", func() {
	var (
		atcServer *ghttp.Server

		homeDir string
	)

	BeforeEach(func() {
		var err error

		homeDir, err = ioutil.TempDir("", "fly-test")
		Expect(err).NotTo(HaveOccurred())

		if runtime.GOOS == "windows" {
			os.Setenv("USERPROFILE", homeDir)
		} else {
			os.Setenv("HOME", homeDir)
		}
	})

	AfterEach(func() {
		os.RemoveAll(homeDir)
	})

	Describe("login", func() {
		var (
			flyCmd *exec.Cmd
			stdin io.WriteCloser
		)
		BeforeEach(func() {
			l := log.New(GinkgoWriter, "TLSServer", 0)
			atcServer = ghttp.NewTLSServer()
			atcServer.HTTPTestServer.Config.ErrorLog = l
		})

		Context("to new target with invalid SSL with -k", func() {
			BeforeEach(func() {
				atcServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("GET", "/api/v1/auth/methods"),
						ghttp.RespondWithJSONEncoded(200, []atc.AuthMethod{
							{
								Type:        atc.AuthTypeBasic,
								DisplayName: "Basic",
								AuthURL:     "https://example.com/login/basic",
							},
						}),
					),
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("GET", "/api/v1/auth/token"),
						ghttp.VerifyBasicAuth("some username", "some password"),
						ghttp.RespondWithJSONEncoded(200, atc.AuthToken{
							Type:  "Bearer",
							Value: "some-token",
						}),
					),
				)

				flyCmd = exec.Command(flyPath, "-t", "some-target", "login", "-c", atcServer.URL(), "-k")

				var err error
				stdin, err = flyCmd.StdinPipe()
				Expect(err).NotTo(HaveOccurred())
			})

			It("succeeds", func() {
				sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
				Expect(err).NotTo(HaveOccurred())

				Eventually(sess.Out).Should(gbytes.Say("username: "))

				_, err = fmt.Fprintf(stdin, "some username\n")
				Expect(err).NotTo(HaveOccurred())

				Eventually(sess.Out).Should(gbytes.Say("password: "))

				_, err = fmt.Fprintf(stdin, "some password\n")
				Expect(err).NotTo(HaveOccurred())

				Consistently(sess.Out.Contents).ShouldNot(ContainSubstring("some password"))

				Eventually(sess.Out).Should(gbytes.Say("token saved"))

				err = stdin.Close()
				Expect(err).NotTo(HaveOccurred())

				<-sess.Exited
				Expect(sess.ExitCode()).To(Equal(0))
			})

			Context("login to existing target", func() {
				var otherCmd *exec.Cmd
				BeforeEach(func() {
					atcServer.AppendHandlers(
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("GET", "/api/v1/auth/methods"),
							ghttp.RespondWithJSONEncoded(200, []atc.AuthMethod{
								{
									Type:        atc.AuthTypeBasic,
									DisplayName: "Basic",
									AuthURL:     "https://example.com/login/basic",
								},
							}),
						),
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("GET", "/api/v1/auth/token"),
							ghttp.VerifyBasicAuth("some username", "some password"),
							ghttp.RespondWithJSONEncoded(200, atc.AuthToken{
								Type:  "Bearer",
								Value: "some-token",
							}),
						),
					)

					sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
					Expect(err).NotTo(HaveOccurred())
					Eventually(sess.Out).Should(gbytes.Say("username: "))
					_, err = fmt.Fprintf(stdin, "some username\n")
					Expect(err).NotTo(HaveOccurred())
					Eventually(sess.Out).Should(gbytes.Say("password: "))
					_, err = fmt.Fprintf(stdin, "some password\n")
					Expect(err).NotTo(HaveOccurred())
					Eventually(sess.Out).Should(gbytes.Say("token saved"))

					err = stdin.Close()
					Expect(err).NotTo(HaveOccurred())

					<-sess.Exited
					Expect(sess.ExitCode()).To(Equal(0))
				})

				Context("with -k", func() {
					BeforeEach(func() {
						otherCmd = exec.Command(flyPath, "-t", "some-target", "login", "-k")

						var err error
						stdin, err = otherCmd.StdinPipe()
						Expect(err).NotTo(HaveOccurred())
					})

					It("succeeds", func() {
						sess, err := gexec.Start(otherCmd, GinkgoWriter, GinkgoWriter)
						Expect(err).NotTo(HaveOccurred())

						Eventually(sess.Out).Should(gbytes.Say("username: "))

						_, err = fmt.Fprintf(stdin, "some username\n")
						Expect(err).NotTo(HaveOccurred())

						Eventually(sess.Out).Should(gbytes.Say("password: "))

						_, err = fmt.Fprintf(stdin, "some password\n")
						Expect(err).NotTo(HaveOccurred())

						Eventually(sess.Out).Should(gbytes.Say("token saved"))

						err = stdin.Close()
						Expect(err).NotTo(HaveOccurred())

						err = stdin.Close()
						Expect(err).NotTo(HaveOccurred())

						<-sess.Exited
						Expect(sess.ExitCode()).To(Equal(0))
					})
				})

				Context("without -k", func() {
					BeforeEach(func() {
						otherCmd = exec.Command(flyPath, "-t", "some-target", "login")
					})

					It("errors", func() {
						sess, err := gexec.Start(otherCmd, GinkgoWriter, GinkgoWriter)
						Expect(err).NotTo(HaveOccurred())

						<-sess.Exited
						Expect(sess.ExitCode()).To(Equal(1))
						Eventually(sess.Err).Should(gbytes.Say("x509: certificate signed by unknown authority"))
					})
				})
			})

		})

		Context("to new target with invalid SSL without -k", func() {
			BeforeEach(func() {
				flyCmd = exec.Command(flyPath, "-t", "some-target", "login", "-c", atcServer.URL())

				var err error
				stdin, err = flyCmd.StdinPipe()
				Expect(err).NotTo(HaveOccurred())
			})

			It("errors", func() {
				//TODO suppress stderr (from http library)
				sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
				Expect(err).NotTo(HaveOccurred())

				err = stdin.Close()
				Expect(err).NotTo(HaveOccurred())

				<-sess.Exited
				Expect(sess.ExitCode()).To(Equal(1))
				Eventually(sess.Err).Should(gbytes.Say("x509: certificate signed by unknown authority"))
			})
		})

		Context("to existing target with invalid SSL certificate", func() {
			Context("when 'insecure' is not set", func() {
				BeforeEach(func() {
					flyrcContents := `targets:
  some-target:
    api: ` + atcServer.URL() + `
    token:
      type: Bearer
      value: some-token`
					ioutil.WriteFile(homeDir+"/.flyrc", []byte(flyrcContents), 0777)
				})

				Context("with -k", func() {
					BeforeEach(func() {
						atcServer.AppendHandlers(
							ghttp.CombineHandlers(
								ghttp.VerifyRequest("GET", "/api/v1/auth/methods"),
								ghttp.RespondWithJSONEncoded(200, []atc.AuthMethod{
									{
										Type:        atc.AuthTypeBasic,
										DisplayName: "Basic",
										AuthURL:     "https://example.com/login/basic",
									},
								}),
							),
							ghttp.CombineHandlers(
								ghttp.VerifyRequest("GET", "/api/v1/auth/token"),
								ghttp.VerifyBasicAuth("some username", "some password"),
								ghttp.RespondWithJSONEncoded(200, atc.AuthToken{
									Type:  "Bearer",
									Value: "some-token",
								}),
							),
						)

						flyCmd = exec.Command(flyPath, "-t", "some-target", "login", "-k")

						var err error
						stdin, err = flyCmd.StdinPipe()
						Expect(err).NotTo(HaveOccurred())
					})

					It("succeeds", func() {
						sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
						Expect(err).NotTo(HaveOccurred())

						Eventually(sess.Out).Should(gbytes.Say("username: "))

						_, err = fmt.Fprintf(stdin, "some username\n")
						Expect(err).NotTo(HaveOccurred())

						Eventually(sess.Out).Should(gbytes.Say("password: "))

						_, err = fmt.Fprintf(stdin, "some password\n")
						Expect(err).NotTo(HaveOccurred())

						Consistently(sess.Out.Contents).ShouldNot(ContainSubstring("some password"))

						Eventually(sess.Out).Should(gbytes.Say("token saved"))

						err = stdin.Close()
						Expect(err).NotTo(HaveOccurred())

						<-sess.Exited
						Expect(sess.ExitCode()).To(Equal(0))
					})
				})

				Context("without -k", func() {
					BeforeEach(func() {
						flyCmd = exec.Command(flyPath, "-t", "some-target", "login")
					})

					It("errors", func() {
						sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
						Expect(err).NotTo(HaveOccurred())

						<-sess.Exited
						Expect(sess.ExitCode()).To(Equal(1))
						Eventually(sess.Err).Should(gbytes.Say("x509: certificate signed by unknown authority"))
					})
				})
			})
		})
	})
})
<commit_msg>remove the racism<commit_after>package integration_test

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"runtime"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
	"github.com/onsi/gomega/ghttp"

	"github.com/concourse/atc"
)

var _ = Describe("login -k Command", func() {
	var (
		atcServer *ghttp.Server

		homeDir string
	)

	BeforeEach(func() {
		var err error

		homeDir, err = ioutil.TempDir("", "fly-test")
		Expect(err).NotTo(HaveOccurred())

		if runtime.GOOS == "windows" {
			os.Setenv("USERPROFILE", homeDir)
		} else {
			os.Setenv("HOME", homeDir)
		}
	})

	AfterEach(func() {
		os.RemoveAll(homeDir)
	})

	Describe("login", func() {
		var (
			flyCmd *exec.Cmd
			stdin io.WriteCloser
		)
		BeforeEach(func() {
			l := log.New(GinkgoWriter, "TLSServer", 0)
			atcServer = ghttp.NewUnstartedServer()
			atcServer.HTTPTestServer.Config.ErrorLog = l
			atcServer.HTTPTestServer.StartTLS()
		})

		AfterEach(func() {
			atcServer.Close()
		})

		Context("to new target with invalid SSL with -k", func() {
			BeforeEach(func() {
				atcServer.AppendHandlers(
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("GET", "/api/v1/auth/methods"),
						ghttp.RespondWithJSONEncoded(200, []atc.AuthMethod{
							{
								Type:        atc.AuthTypeBasic,
								DisplayName: "Basic",
								AuthURL:     "https://example.com/login/basic",
							},
						}),
					),
					ghttp.CombineHandlers(
						ghttp.VerifyRequest("GET", "/api/v1/auth/token"),
						ghttp.VerifyBasicAuth("some username", "some password"),
						ghttp.RespondWithJSONEncoded(200, atc.AuthToken{
							Type:  "Bearer",
							Value: "some-token",
						}),
					),
				)

				flyCmd = exec.Command(flyPath, "-t", "some-target", "login", "-c", atcServer.URL(), "-k")

				var err error
				stdin, err = flyCmd.StdinPipe()
				Expect(err).NotTo(HaveOccurred())
			})

			It("succeeds", func() {
				sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
				Expect(err).NotTo(HaveOccurred())

				Eventually(sess.Out).Should(gbytes.Say("username: "))

				_, err = fmt.Fprintf(stdin, "some username\n")
				Expect(err).NotTo(HaveOccurred())

				Eventually(sess.Out).Should(gbytes.Say("password: "))

				_, err = fmt.Fprintf(stdin, "some password\n")
				Expect(err).NotTo(HaveOccurred())

				Consistently(sess.Out.Contents).ShouldNot(ContainSubstring("some password"))

				Eventually(sess.Out).Should(gbytes.Say("token saved"))

				err = stdin.Close()
				Expect(err).NotTo(HaveOccurred())

				<-sess.Exited
				Expect(sess.ExitCode()).To(Equal(0))
			})

			Context("login to existing target", func() {
				var otherCmd *exec.Cmd
				BeforeEach(func() {
					atcServer.AppendHandlers(
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("GET", "/api/v1/auth/methods"),
							ghttp.RespondWithJSONEncoded(200, []atc.AuthMethod{
								{
									Type:        atc.AuthTypeBasic,
									DisplayName: "Basic",
									AuthURL:     "https://example.com/login/basic",
								},
							}),
						),
						ghttp.CombineHandlers(
							ghttp.VerifyRequest("GET", "/api/v1/auth/token"),
							ghttp.VerifyBasicAuth("some username", "some password"),
							ghttp.RespondWithJSONEncoded(200, atc.AuthToken{
								Type:  "Bearer",
								Value: "some-token",
							}),
						),
					)

					sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
					Expect(err).NotTo(HaveOccurred())
					Eventually(sess.Out).Should(gbytes.Say("username: "))
					_, err = fmt.Fprintf(stdin, "some username\n")
					Expect(err).NotTo(HaveOccurred())
					Eventually(sess.Out).Should(gbytes.Say("password: "))
					_, err = fmt.Fprintf(stdin, "some password\n")
					Expect(err).NotTo(HaveOccurred())
					Eventually(sess.Out).Should(gbytes.Say("token saved"))

					err = stdin.Close()
					Expect(err).NotTo(HaveOccurred())

					<-sess.Exited
					Expect(sess.ExitCode()).To(Equal(0))
				})

				Context("with -k", func() {
					BeforeEach(func() {
						otherCmd = exec.Command(flyPath, "-t", "some-target", "login", "-k")

						var err error
						stdin, err = otherCmd.StdinPipe()
						Expect(err).NotTo(HaveOccurred())
					})

					It("succeeds", func() {
						sess, err := gexec.Start(otherCmd, GinkgoWriter, GinkgoWriter)
						Expect(err).NotTo(HaveOccurred())

						Eventually(sess.Out).Should(gbytes.Say("username: "))

						_, err = fmt.Fprintf(stdin, "some username\n")
						Expect(err).NotTo(HaveOccurred())

						Eventually(sess.Out).Should(gbytes.Say("password: "))

						_, err = fmt.Fprintf(stdin, "some password\n")
						Expect(err).NotTo(HaveOccurred())

						Eventually(sess.Out).Should(gbytes.Say("token saved"))

						err = stdin.Close()
						Expect(err).NotTo(HaveOccurred())

						err = stdin.Close()
						Expect(err).NotTo(HaveOccurred())

						<-sess.Exited
						Expect(sess.ExitCode()).To(Equal(0))
					})
				})

				Context("without -k", func() {
					BeforeEach(func() {
						otherCmd = exec.Command(flyPath, "-t", "some-target", "login")
					})

					It("errors", func() {
						sess, err := gexec.Start(otherCmd, GinkgoWriter, GinkgoWriter)
						Expect(err).NotTo(HaveOccurred())

						<-sess.Exited
						Expect(sess.ExitCode()).To(Equal(1))
						Eventually(sess.Err).Should(gbytes.Say("x509: certificate signed by unknown authority"))
					})
				})
			})

		})

		Context("to new target with invalid SSL without -k", func() {
			BeforeEach(func() {
				flyCmd = exec.Command(flyPath, "-t", "some-target", "login", "-c", atcServer.URL())

				var err error
				stdin, err = flyCmd.StdinPipe()
				Expect(err).NotTo(HaveOccurred())
			})

			It("errors", func() {
				sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
				Expect(err).NotTo(HaveOccurred())

				err = stdin.Close()
				Expect(err).NotTo(HaveOccurred())

				<-sess.Exited
				Expect(sess.ExitCode()).To(Equal(1))
				Eventually(sess.Err).Should(gbytes.Say("x509: certificate signed by unknown authority"))
			})
		})

		Context("to existing target with invalid SSL certificate", func() {
			Context("when 'insecure' is not set", func() {
				BeforeEach(func() {
					flyrcContents := `targets:
  some-target:
    api: ` + atcServer.URL() + `
    token:
      type: Bearer
      value: some-token`
					ioutil.WriteFile(homeDir+"/.flyrc", []byte(flyrcContents), 0777)
				})

				Context("with -k", func() {
					BeforeEach(func() {
						atcServer.AppendHandlers(
							ghttp.CombineHandlers(
								ghttp.VerifyRequest("GET", "/api/v1/auth/methods"),
								ghttp.RespondWithJSONEncoded(200, []atc.AuthMethod{
									{
										Type:        atc.AuthTypeBasic,
										DisplayName: "Basic",
										AuthURL:     "https://example.com/login/basic",
									},
								}),
							),
							ghttp.CombineHandlers(
								ghttp.VerifyRequest("GET", "/api/v1/auth/token"),
								ghttp.VerifyBasicAuth("some username", "some password"),
								ghttp.RespondWithJSONEncoded(200, atc.AuthToken{
									Type:  "Bearer",
									Value: "some-token",
								}),
							),
						)

						flyCmd = exec.Command(flyPath, "-t", "some-target", "login", "-k")

						var err error
						stdin, err = flyCmd.StdinPipe()
						Expect(err).NotTo(HaveOccurred())
					})

					It("succeeds", func() {
						sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
						Expect(err).NotTo(HaveOccurred())

						Eventually(sess.Out).Should(gbytes.Say("username: "))

						_, err = fmt.Fprintf(stdin, "some username\n")
						Expect(err).NotTo(HaveOccurred())

						Eventually(sess.Out).Should(gbytes.Say("password: "))

						_, err = fmt.Fprintf(stdin, "some password\n")
						Expect(err).NotTo(HaveOccurred())

						Consistently(sess.Out.Contents).ShouldNot(ContainSubstring("some password"))

						Eventually(sess.Out).Should(gbytes.Say("token saved"))

						err = stdin.Close()
						Expect(err).NotTo(HaveOccurred())

						<-sess.Exited
						Expect(sess.ExitCode()).To(Equal(0))
					})
				})

				Context("without -k", func() {
					BeforeEach(func() {
						flyCmd = exec.Command(flyPath, "-t", "some-target", "login")
					})

					It("errors", func() {
						sess, err := gexec.Start(flyCmd, GinkgoWriter, GinkgoWriter)
						Expect(err).NotTo(HaveOccurred())

						<-sess.Exited
						Expect(sess.ExitCode()).To(Equal(1))
						Eventually(sess.Err).Should(gbytes.Say("x509: certificate signed by unknown authority"))
					})
				})
			})
		})
	})
})
<|endoftext|>


<commit_before>package service

import (
	"context"
	"io"
	"io/ioutil"
	"os"
	"path"
	"testing"

	"github.com/docker/docker/api/types"
	swarmtypes "github.com/docker/docker/api/types/swarm"
	"github.com/docker/docker/api/types/swarm/runtime"
	"github.com/docker/docker/integration/internal/swarm"
	"github.com/docker/docker/internal/test/daemon"
	"github.com/docker/docker/internal/test/fixtures/plugin"
	"github.com/docker/docker/internal/test/registry"
	"gotest.tools/assert"
	"gotest.tools/poll"
	"gotest.tools/skip"
)

func TestServicePlugin(t *testing.T) {
	skip.If(t, testEnv.IsRemoteDaemon, "cannot run daemon when remote daemon")
	skip.If(t, testEnv.DaemonInfo.OSType == "windows")
	skip.If(t, os.Getenv("DOCKER_ENGINE_GOARCH") != "amd64")
	defer setupTest(t)()

	reg := registry.NewV2(t)
	defer reg.Close()

	repo := path.Join(registry.DefaultURL, "swarm", "test:v1")
	repo2 := path.Join(registry.DefaultURL, "swarm", "test:v2")
	name := "test"

	d := daemon.New(t)
	d.StartWithBusybox(t)
	apiclient := d.NewClientT(t)
	err := plugin.Create(context.Background(), apiclient, repo)
	assert.NilError(t, err)
	r, err := apiclient.PluginPush(context.Background(), repo, "")
	assert.NilError(t, err)
	_, err = io.Copy(ioutil.Discard, r)
	assert.NilError(t, err)
	err = apiclient.PluginRemove(context.Background(), repo, types.PluginRemoveOptions{})
	assert.NilError(t, err)
	err = plugin.Create(context.Background(), apiclient, repo2)
	assert.NilError(t, err)
	r, err = apiclient.PluginPush(context.Background(), repo2, "")
	assert.NilError(t, err)
	_, err = io.Copy(ioutil.Discard, r)
	assert.NilError(t, err)
	err = apiclient.PluginRemove(context.Background(), repo2, types.PluginRemoveOptions{})
	assert.NilError(t, err)
	d.Stop(t)

	d1 := swarm.NewSwarm(t, testEnv, daemon.WithExperimental)
	defer d1.Stop(t)
	d2 := daemon.New(t, daemon.WithExperimental, daemon.WithSwarmPort(daemon.DefaultSwarmPort+1))
	d2.StartAndSwarmJoin(t, d1, true)
	defer d2.Stop(t)
	d3 := daemon.New(t, daemon.WithExperimental, daemon.WithSwarmPort(daemon.DefaultSwarmPort+2))
	d3.StartAndSwarmJoin(t, d1, false)
	defer d3.Stop(t)

	id := d1.CreateService(t, makePlugin(repo, name, nil))
	poll.WaitOn(t, d1.PluginIsRunning(name), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginIsRunning(name), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginIsRunning(name), swarm.ServicePoll)

	service := d1.GetService(t, id)
	d1.UpdateService(t, service, makePlugin(repo2, name, nil))
	poll.WaitOn(t, d1.PluginReferenceIs(name, repo2), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginReferenceIs(name, repo2), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginReferenceIs(name, repo2), swarm.ServicePoll)
	poll.WaitOn(t, d1.PluginIsRunning(name), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginIsRunning(name), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginIsRunning(name), swarm.ServicePoll)

	d1.RemoveService(t, id)
	poll.WaitOn(t, d1.PluginIsNotPresent(name), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginIsNotPresent(name), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginIsNotPresent(name), swarm.ServicePoll)

	// constrain to managers only
	id = d1.CreateService(t, makePlugin(repo, name, []string{"node.role==manager"}))
	poll.WaitOn(t, d1.PluginIsRunning(name), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginIsRunning(name), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginIsNotPresent(name), swarm.ServicePoll)

	d1.RemoveService(t, id)
	poll.WaitOn(t, d1.PluginIsNotPresent(name), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginIsNotPresent(name), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginIsNotPresent(name), swarm.ServicePoll)

	// with no name
	id = d1.CreateService(t, makePlugin(repo, "", nil))
	poll.WaitOn(t, d1.PluginIsRunning(repo), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginIsRunning(repo), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginIsRunning(repo), swarm.ServicePoll)

	d1.RemoveService(t, id)
	poll.WaitOn(t, d1.PluginIsNotPresent(repo), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginIsNotPresent(repo), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginIsNotPresent(repo), swarm.ServicePoll)
}

func makePlugin(repo, name string, constraints []string) func(*swarmtypes.Service) {
	return func(s *swarmtypes.Service) {
		s.Spec.TaskTemplate.Runtime = "plugin"
		s.Spec.TaskTemplate.PluginSpec = &runtime.PluginSpec{
			Name:   name,
			Remote: repo,
		}
		if constraints != nil {
			s.Spec.TaskTemplate.Placement = &swarmtypes.Placement{
				Constraints: constraints,
			}
		}
	}
}
<commit_msg>Use constant for task runtime value<commit_after>package service

import (
	"context"
	"io"
	"io/ioutil"
	"os"
	"path"
	"testing"

	"github.com/docker/docker/api/types"
	swarmtypes "github.com/docker/docker/api/types/swarm"
	"github.com/docker/docker/api/types/swarm/runtime"
	"github.com/docker/docker/integration/internal/swarm"
	"github.com/docker/docker/internal/test/daemon"
	"github.com/docker/docker/internal/test/fixtures/plugin"
	"github.com/docker/docker/internal/test/registry"
	"gotest.tools/assert"
	"gotest.tools/poll"
	"gotest.tools/skip"
)

func TestServicePlugin(t *testing.T) {
	skip.If(t, testEnv.IsRemoteDaemon, "cannot run daemon when remote daemon")
	skip.If(t, testEnv.DaemonInfo.OSType == "windows")
	skip.If(t, os.Getenv("DOCKER_ENGINE_GOARCH") != "amd64")
	defer setupTest(t)()

	reg := registry.NewV2(t)
	defer reg.Close()

	repo := path.Join(registry.DefaultURL, "swarm", "test:v1")
	repo2 := path.Join(registry.DefaultURL, "swarm", "test:v2")
	name := "test"

	d := daemon.New(t)
	d.StartWithBusybox(t)
	apiclient := d.NewClientT(t)
	err := plugin.Create(context.Background(), apiclient, repo)
	assert.NilError(t, err)
	r, err := apiclient.PluginPush(context.Background(), repo, "")
	assert.NilError(t, err)
	_, err = io.Copy(ioutil.Discard, r)
	assert.NilError(t, err)
	err = apiclient.PluginRemove(context.Background(), repo, types.PluginRemoveOptions{})
	assert.NilError(t, err)
	err = plugin.Create(context.Background(), apiclient, repo2)
	assert.NilError(t, err)
	r, err = apiclient.PluginPush(context.Background(), repo2, "")
	assert.NilError(t, err)
	_, err = io.Copy(ioutil.Discard, r)
	assert.NilError(t, err)
	err = apiclient.PluginRemove(context.Background(), repo2, types.PluginRemoveOptions{})
	assert.NilError(t, err)
	d.Stop(t)

	d1 := swarm.NewSwarm(t, testEnv, daemon.WithExperimental)
	defer d1.Stop(t)
	d2 := daemon.New(t, daemon.WithExperimental, daemon.WithSwarmPort(daemon.DefaultSwarmPort+1))
	d2.StartAndSwarmJoin(t, d1, true)
	defer d2.Stop(t)
	d3 := daemon.New(t, daemon.WithExperimental, daemon.WithSwarmPort(daemon.DefaultSwarmPort+2))
	d3.StartAndSwarmJoin(t, d1, false)
	defer d3.Stop(t)

	id := d1.CreateService(t, makePlugin(repo, name, nil))
	poll.WaitOn(t, d1.PluginIsRunning(name), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginIsRunning(name), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginIsRunning(name), swarm.ServicePoll)

	service := d1.GetService(t, id)
	d1.UpdateService(t, service, makePlugin(repo2, name, nil))
	poll.WaitOn(t, d1.PluginReferenceIs(name, repo2), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginReferenceIs(name, repo2), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginReferenceIs(name, repo2), swarm.ServicePoll)
	poll.WaitOn(t, d1.PluginIsRunning(name), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginIsRunning(name), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginIsRunning(name), swarm.ServicePoll)

	d1.RemoveService(t, id)
	poll.WaitOn(t, d1.PluginIsNotPresent(name), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginIsNotPresent(name), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginIsNotPresent(name), swarm.ServicePoll)

	// constrain to managers only
	id = d1.CreateService(t, makePlugin(repo, name, []string{"node.role==manager"}))
	poll.WaitOn(t, d1.PluginIsRunning(name), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginIsRunning(name), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginIsNotPresent(name), swarm.ServicePoll)

	d1.RemoveService(t, id)
	poll.WaitOn(t, d1.PluginIsNotPresent(name), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginIsNotPresent(name), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginIsNotPresent(name), swarm.ServicePoll)

	// with no name
	id = d1.CreateService(t, makePlugin(repo, "", nil))
	poll.WaitOn(t, d1.PluginIsRunning(repo), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginIsRunning(repo), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginIsRunning(repo), swarm.ServicePoll)

	d1.RemoveService(t, id)
	poll.WaitOn(t, d1.PluginIsNotPresent(repo), swarm.ServicePoll)
	poll.WaitOn(t, d2.PluginIsNotPresent(repo), swarm.ServicePoll)
	poll.WaitOn(t, d3.PluginIsNotPresent(repo), swarm.ServicePoll)
}

func makePlugin(repo, name string, constraints []string) func(*swarmtypes.Service) {
	return func(s *swarmtypes.Service) {
		s.Spec.TaskTemplate.Runtime = swarmtypes.RuntimePlugin
		s.Spec.TaskTemplate.PluginSpec = &runtime.PluginSpec{
			Name:   name,
			Remote: repo,
		}
		if constraints != nil {
			s.Spec.TaskTemplate.Placement = &swarmtypes.Placement{
				Constraints: constraints,
			}
		}
	}
}
<|endoftext|>


<commit_before>// Package configfile reads and writes gocryptfs.conf does the key
// wrapping.
package configfile

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"

	"github.com/rfjakob/gocryptfs/internal/contentenc"
	"github.com/rfjakob/gocryptfs/internal/cryptocore"
	"github.com/rfjakob/gocryptfs/internal/exitcodes"
	"github.com/rfjakob/gocryptfs/internal/tlog"
)
import "os"

const (
	// ConfDefaultName is the default configuration file name.
	// The dot "." is not used in base64url (RFC4648), hence
	// we can never clash with an encrypted file.
	ConfDefaultName = "gocryptfs.conf"
	// ConfReverseName is the default configuration file name in reverse mode,
	// the config file gets stored next to the plain-text files. Make it hidden
	// (start with dot) to not annoy the user.
	ConfReverseName = ".gocryptfs.reverse.conf"
)

// ConfFile is the content of a config file.
type ConfFile struct {
	// Creator is the gocryptfs version string.
	// This only documents the config file for humans who look at it. The actual
	// technical info is contained in FeatureFlags.
	Creator string
	// EncryptedKey holds an encrypted AES key, unlocked using a password
	// hashed with scrypt
	EncryptedKey []byte
	// ScryptObject stores parameters for scrypt hashing (key derivation)
	ScryptObject ScryptKDF
	// Version is the On-Disk-Format version this filesystem uses
	Version uint16
	// FeatureFlags is a list of feature flags this filesystem has enabled.
	// If gocryptfs encounters a feature flag it does not support, it will refuse
	// mounting. This mechanism is analogous to the ext4 feature flags that are
	// stored in the superblock.
	FeatureFlags []string
	// TrezorPayload stores 32 random bytes used for unlocking the master key using
	// a Trezor security module. The randomness makes sure that a unique unlock
	// value is used for each gocryptfs filesystem.
	TrezorPayload []byte `json:",omitempty"`
	// Filename is the name of the config file. Not exported to JSON.
	filename string
}

// randBytesDevRandom gets "n" random bytes from /dev/random or panics
func randBytesDevRandom(n int) []byte {
	f, err := os.Open("/dev/random")
	if err != nil {
		log.Panic("Failed to open /dev/random: " + err.Error())
	}
	defer f.Close()
	b := make([]byte, n)
	_, err = io.ReadFull(f, b)
	if err != nil {
		log.Panic("Failed to read random bytes: " + err.Error())
	}
	return b
}

// Create - create a new config with a random key encrypted with
// "password" and write it to "filename".
// Uses scrypt with cost parameter logN.
func Create(filename string, password []byte, plaintextNames bool,
	logN int, creator string, aessiv bool, devrandom bool, trezorPayload []byte) error {
	var cf ConfFile
	cf.filename = filename
	cf.Creator = creator
	cf.Version = contentenc.CurrentVersion

	// Set feature flags
	cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagGCMIV128])
	cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagHKDF])
	if plaintextNames {
		cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagPlaintextNames])
	} else {
		cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagDirIV])
		cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagEMENames])
		cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagLongNames])
		cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagRaw64])
	}
	if aessiv {
		cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagAESSIV])
	}
	if len(trezorPayload) > 0 {
		cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagTrezor])
		cf.TrezorPayload = trezorPayload
	}
	{
		// Generate new random master key
		var key []byte
		if devrandom {
			key = randBytesDevRandom(cryptocore.KeyLen)
		} else {
			key = cryptocore.RandBytes(cryptocore.KeyLen)
		}
		tlog.PrintMasterkeyReminder(key)
		// Encrypt it using the password
		// This sets ScryptObject and EncryptedKey
		// Note: this looks at the FeatureFlags, so call it AFTER setting them.
		cf.EncryptKey(key, password, logN)
		for i := range key {
			key[i] = 0
		}
		// key runs out of scope here
	}
	// Write file to disk
	return cf.WriteFile()
}

// Load - read config file from disk and decrypt the
// contained key using "password".
// Returns the decrypted key and the ConfFile object
//
// If "password" is empty, the config file is read
// but the key is not decrypted (returns nil in its place).
func Load(filename string, password []byte) ([]byte, *ConfFile, error) {
	var cf ConfFile
	cf.filename = filename

	// Read from disk
	js, err := ioutil.ReadFile(filename)
	if err != nil {
		fmt.Printf("LoadConfFile: ReadFile: %#v\n", err)
		return nil, nil, err
	}

	// Unmarshal
	err = json.Unmarshal(js, &cf)
	if err != nil {
		tlog.Warn.Printf("Failed to unmarshal config file")
		return nil, nil, err
	}

	if cf.Version != contentenc.CurrentVersion {
		return nil, nil, fmt.Errorf("Unsupported on-disk format %d", cf.Version)
	}

	// Check that all set feature flags are known
	for _, flag := range cf.FeatureFlags {
		if !cf.isFeatureFlagKnown(flag) {
			return nil, nil, fmt.Errorf("Unsupported feature flag %q", flag)
		}
	}

	// Check that all required feature flags are set
	var requiredFlags []flagIota
	if cf.IsFeatureFlagSet(FlagPlaintextNames) {
		requiredFlags = requiredFlagsPlaintextNames
	} else {
		requiredFlags = requiredFlagsNormal
	}
	deprecatedFs := false
	for _, i := range requiredFlags {
		if !cf.IsFeatureFlagSet(i) {
			fmt.Fprintf(os.Stderr, "Required feature flag %q is missing\n", knownFlags[i])
			deprecatedFs = true
		}
	}
	if deprecatedFs {
		fmt.Fprintf(os.Stderr, tlog.ColorYellow+`
    The filesystem was created by gocryptfs v0.6 or earlier. This version of
    gocryptfs can no longer mount the filesystem.
    Please download gocryptfs v0.11 and upgrade your filesystem,
    see https://github.com/rfjakob/gocryptfs/wiki/Upgrading for instructions.

    If you have trouble upgrading, join the discussion at
    https://github.com/rfjakob/gocryptfs/issues/29 .

`+tlog.ColorReset)

		return nil, nil, exitcodes.NewErr("Deprecated filesystem", exitcodes.DeprecatedFS)
	}
	if len(password) == 0 {
		// We have validated the config file, but without a password we cannot
		// decrypt the master key. Return only the parsed config.
		return nil, &cf, nil
	}

	// Generate derived key from password
	scryptHash := cf.ScryptObject.DeriveKey(password)

	// Unlock master key using password-based key
	useHKDF := cf.IsFeatureFlagSet(FlagHKDF)
	ce := getKeyEncrypter(scryptHash, useHKDF)

	tlog.Warn.Enabled = false // Silence DecryptBlock() error messages on incorrect password
	key, err := ce.DecryptBlock(cf.EncryptedKey, 0, nil)
	tlog.Warn.Enabled = true
	if err != nil {
		tlog.Warn.Printf("failed to unlock master key: %s", err.Error())
		return nil, nil, exitcodes.NewErr("Password incorrect.", exitcodes.PasswordIncorrect)
	}

	return key, &cf, err
}

// EncryptKey - encrypt "key" using an scrypt hash generated from "password"
// and store it in cf.EncryptedKey.
// Uses scrypt with cost parameter logN and stores the scrypt parameters in
// cf.ScryptObject.
func (cf *ConfFile) EncryptKey(key []byte, password []byte, logN int) {
	// Generate scrypt-derived key from password
	cf.ScryptObject = NewScryptKDF(logN)
	scryptHash := cf.ScryptObject.DeriveKey(password)
	// Lock master key using password-based key
	useHKDF := cf.IsFeatureFlagSet(FlagHKDF)
	ce := getKeyEncrypter(scryptHash, useHKDF)
	cf.EncryptedKey = ce.EncryptBlock(key, 0, nil)
	// Purge scrypt-derived key
	for i := range scryptHash {
		scryptHash[i] = 0
	}
}

// WriteFile - write out config in JSON format to file "filename.tmp"
// then rename over "filename".
// This way a password change atomically replaces the file.
func (cf *ConfFile) WriteFile() error {
	tmp := cf.filename + ".tmp"
	// 0400 permissions: gocryptfs.conf should be kept secret and never be written to.
	fd, err := os.OpenFile(tmp, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0400)
	if err != nil {
		return err
	}
	js, err := json.MarshalIndent(cf, "", "\t")
	if err != nil {
		return err
	}
	// For convenience for the user, add a newline at the end.
	js = append(js, '\n')
	_, err = fd.Write(js)
	if err != nil {
		return err
	}
	err = fd.Sync()
	if err != nil {
		return err
	}
	err = fd.Close()
	if err != nil {
		return err
	}
	err = os.Rename(tmp, cf.filename)
	return err
}

// getKeyEncrypter is a helper function that returns the right ContentEnc
// instance for the "useHKDF" setting.
func getKeyEncrypter(scryptHash []byte, useHKDF bool) *contentenc.ContentEnc {
	IVLen := 96
	// gocryptfs v1.2 and older used 96-bit IVs for master key encryption.
	// v1.3 adds the "HKDF" feature flag, which also enables 128-bit nonces.
	if useHKDF {
		IVLen = contentenc.DefaultIVBits
	}
	cc := cryptocore.New(scryptHash, cryptocore.BackendGoGCM, IVLen, useHKDF, false)
	ce := contentenc.New(cc, 4096, false)
	return ce
}
<commit_msg>configfile: drop superflous Printf<commit_after>// Package configfile reads and writes gocryptfs.conf does the key
// wrapping.
package configfile

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"

	"github.com/rfjakob/gocryptfs/internal/contentenc"
	"github.com/rfjakob/gocryptfs/internal/cryptocore"
	"github.com/rfjakob/gocryptfs/internal/exitcodes"
	"github.com/rfjakob/gocryptfs/internal/tlog"
)
import "os"

const (
	// ConfDefaultName is the default configuration file name.
	// The dot "." is not used in base64url (RFC4648), hence
	// we can never clash with an encrypted file.
	ConfDefaultName = "gocryptfs.conf"
	// ConfReverseName is the default configuration file name in reverse mode,
	// the config file gets stored next to the plain-text files. Make it hidden
	// (start with dot) to not annoy the user.
	ConfReverseName = ".gocryptfs.reverse.conf"
)

// ConfFile is the content of a config file.
type ConfFile struct {
	// Creator is the gocryptfs version string.
	// This only documents the config file for humans who look at it. The actual
	// technical info is contained in FeatureFlags.
	Creator string
	// EncryptedKey holds an encrypted AES key, unlocked using a password
	// hashed with scrypt
	EncryptedKey []byte
	// ScryptObject stores parameters for scrypt hashing (key derivation)
	ScryptObject ScryptKDF
	// Version is the On-Disk-Format version this filesystem uses
	Version uint16
	// FeatureFlags is a list of feature flags this filesystem has enabled.
	// If gocryptfs encounters a feature flag it does not support, it will refuse
	// mounting. This mechanism is analogous to the ext4 feature flags that are
	// stored in the superblock.
	FeatureFlags []string
	// TrezorPayload stores 32 random bytes used for unlocking the master key using
	// a Trezor security module. The randomness makes sure that a unique unlock
	// value is used for each gocryptfs filesystem.
	TrezorPayload []byte `json:",omitempty"`
	// Filename is the name of the config file. Not exported to JSON.
	filename string
}

// randBytesDevRandom gets "n" random bytes from /dev/random or panics
func randBytesDevRandom(n int) []byte {
	f, err := os.Open("/dev/random")
	if err != nil {
		log.Panic("Failed to open /dev/random: " + err.Error())
	}
	defer f.Close()
	b := make([]byte, n)
	_, err = io.ReadFull(f, b)
	if err != nil {
		log.Panic("Failed to read random bytes: " + err.Error())
	}
	return b
}

// Create - create a new config with a random key encrypted with
// "password" and write it to "filename".
// Uses scrypt with cost parameter logN.
func Create(filename string, password []byte, plaintextNames bool,
	logN int, creator string, aessiv bool, devrandom bool, trezorPayload []byte) error {
	var cf ConfFile
	cf.filename = filename
	cf.Creator = creator
	cf.Version = contentenc.CurrentVersion

	// Set feature flags
	cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagGCMIV128])
	cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagHKDF])
	if plaintextNames {
		cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagPlaintextNames])
	} else {
		cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagDirIV])
		cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagEMENames])
		cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagLongNames])
		cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagRaw64])
	}
	if aessiv {
		cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagAESSIV])
	}
	if len(trezorPayload) > 0 {
		cf.FeatureFlags = append(cf.FeatureFlags, knownFlags[FlagTrezor])
		cf.TrezorPayload = trezorPayload
	}
	{
		// Generate new random master key
		var key []byte
		if devrandom {
			key = randBytesDevRandom(cryptocore.KeyLen)
		} else {
			key = cryptocore.RandBytes(cryptocore.KeyLen)
		}
		tlog.PrintMasterkeyReminder(key)
		// Encrypt it using the password
		// This sets ScryptObject and EncryptedKey
		// Note: this looks at the FeatureFlags, so call it AFTER setting them.
		cf.EncryptKey(key, password, logN)
		for i := range key {
			key[i] = 0
		}
		// key runs out of scope here
	}
	// Write file to disk
	return cf.WriteFile()
}

// Load - read config file from disk and decrypt the
// contained key using "password".
// Returns the decrypted key and the ConfFile object
//
// If "password" is empty, the config file is read
// but the key is not decrypted (returns nil in its place).
func Load(filename string, password []byte) ([]byte, *ConfFile, error) {
	var cf ConfFile
	cf.filename = filename

	// Read from disk
	js, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, nil, err
	}

	// Unmarshal
	err = json.Unmarshal(js, &cf)
	if err != nil {
		tlog.Warn.Printf("Failed to unmarshal config file")
		return nil, nil, err
	}

	if cf.Version != contentenc.CurrentVersion {
		return nil, nil, fmt.Errorf("Unsupported on-disk format %d", cf.Version)
	}

	// Check that all set feature flags are known
	for _, flag := range cf.FeatureFlags {
		if !cf.isFeatureFlagKnown(flag) {
			return nil, nil, fmt.Errorf("Unsupported feature flag %q", flag)
		}
	}

	// Check that all required feature flags are set
	var requiredFlags []flagIota
	if cf.IsFeatureFlagSet(FlagPlaintextNames) {
		requiredFlags = requiredFlagsPlaintextNames
	} else {
		requiredFlags = requiredFlagsNormal
	}
	deprecatedFs := false
	for _, i := range requiredFlags {
		if !cf.IsFeatureFlagSet(i) {
			fmt.Fprintf(os.Stderr, "Required feature flag %q is missing\n", knownFlags[i])
			deprecatedFs = true
		}
	}
	if deprecatedFs {
		fmt.Fprintf(os.Stderr, tlog.ColorYellow+`
    The filesystem was created by gocryptfs v0.6 or earlier. This version of
    gocryptfs can no longer mount the filesystem.
    Please download gocryptfs v0.11 and upgrade your filesystem,
    see https://github.com/rfjakob/gocryptfs/wiki/Upgrading for instructions.

    If you have trouble upgrading, join the discussion at
    https://github.com/rfjakob/gocryptfs/issues/29 .

`+tlog.ColorReset)

		return nil, nil, exitcodes.NewErr("Deprecated filesystem", exitcodes.DeprecatedFS)
	}
	if len(password) == 0 {
		// We have validated the config file, but without a password we cannot
		// decrypt the master key. Return only the parsed config.
		return nil, &cf, nil
	}

	// Generate derived key from password
	scryptHash := cf.ScryptObject.DeriveKey(password)

	// Unlock master key using password-based key
	useHKDF := cf.IsFeatureFlagSet(FlagHKDF)
	ce := getKeyEncrypter(scryptHash, useHKDF)

	tlog.Warn.Enabled = false // Silence DecryptBlock() error messages on incorrect password
	key, err := ce.DecryptBlock(cf.EncryptedKey, 0, nil)
	tlog.Warn.Enabled = true
	if err != nil {
		tlog.Warn.Printf("failed to unlock master key: %s", err.Error())
		return nil, nil, exitcodes.NewErr("Password incorrect.", exitcodes.PasswordIncorrect)
	}

	return key, &cf, err
}

// EncryptKey - encrypt "key" using an scrypt hash generated from "password"
// and store it in cf.EncryptedKey.
// Uses scrypt with cost parameter logN and stores the scrypt parameters in
// cf.ScryptObject.
func (cf *ConfFile) EncryptKey(key []byte, password []byte, logN int) {
	// Generate scrypt-derived key from password
	cf.ScryptObject = NewScryptKDF(logN)
	scryptHash := cf.ScryptObject.DeriveKey(password)
	// Lock master key using password-based key
	useHKDF := cf.IsFeatureFlagSet(FlagHKDF)
	ce := getKeyEncrypter(scryptHash, useHKDF)
	cf.EncryptedKey = ce.EncryptBlock(key, 0, nil)
	// Purge scrypt-derived key
	for i := range scryptHash {
		scryptHash[i] = 0
	}
}

// WriteFile - write out config in JSON format to file "filename.tmp"
// then rename over "filename".
// This way a password change atomically replaces the file.
func (cf *ConfFile) WriteFile() error {
	tmp := cf.filename + ".tmp"
	// 0400 permissions: gocryptfs.conf should be kept secret and never be written to.
	fd, err := os.OpenFile(tmp, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0400)
	if err != nil {
		return err
	}
	js, err := json.MarshalIndent(cf, "", "\t")
	if err != nil {
		return err
	}
	// For convenience for the user, add a newline at the end.
	js = append(js, '\n')
	_, err = fd.Write(js)
	if err != nil {
		return err
	}
	err = fd.Sync()
	if err != nil {
		return err
	}
	err = fd.Close()
	if err != nil {
		return err
	}
	err = os.Rename(tmp, cf.filename)
	return err
}

// getKeyEncrypter is a helper function that returns the right ContentEnc
// instance for the "useHKDF" setting.
func getKeyEncrypter(scryptHash []byte, useHKDF bool) *contentenc.ContentEnc {
	IVLen := 96
	// gocryptfs v1.2 and older used 96-bit IVs for master key encryption.
	// v1.3 adds the "HKDF" feature flag, which also enables 128-bit nonces.
	if useHKDF {
		IVLen = contentenc.DefaultIVBits
	}
	cc := cryptocore.New(scryptHash, cryptocore.BackendGoGCM, IVLen, useHKDF, false)
	ce := contentenc.New(cc, 4096, false)
	return ce
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	. "github.com/NetSys/di/minion/proto"
	"github.com/fsouza/go-dockerclient"
)

const ETCD = "quay.io/coreos/etcd:v2.1.3"
const OVN_NORTHD = "quay.io/netsys/ovn-northd"
const OVN_CONTROLLER = "quay.io/netsys/ovn-controller"
const OVS_VSWITCHD = "quay.io/netsys/ovs-vswitchd"
const OVS_OVSDB = "quay.io/netsys/ovsdb-server"
const DOCKER_SOCK_PATH = "unix:///var/run/docker.sock"

type Supervisor struct {
	dk  *docker.Client
	cfg MinionConfig
}

func defaultHC() *docker.HostConfig {
	return &docker.HostConfig{
		NetworkMode: "host",
		Binds:       []string{"/usr/share/ca-certificates:/etc/ssl/certs"}}
}

func ovsHC() *docker.HostConfig {
	hc := defaultHC()
	hc.VolumesFrom = []string{"ovsdb-server"}
	return hc
}

func (sv Supervisor) pullImage(image string) error {
	return sv.dk.PullImage(docker.PullImageOptions{
		Repository: image},
		docker.AuthConfiguration{})
}

func NewSupervisor() (*Supervisor, error) {
	dk, err := docker.NewClient(DOCKER_SOCK_PATH)
	if err != nil {
		log.Warning("Failed to create docker client: %s", err)
		return nil, err
	}

	sv := &Supervisor{dk: dk}

	images := []string{OVS_OVSDB, OVN_NORTHD, OVN_CONTROLLER, OVS_VSWITCHD, ETCD}
	for _, image := range images {
		go sv.pullImage(image)
	}

	return sv, nil
}

func (sv Supervisor) getContainer(name string) *string {
	containers, err := sv.dk.ListContainers(docker.ListContainersOptions{All: true})
	if err != nil {
		log.Warning("Failed to list containers: %s", err)
		return nil
	}

	name = "/" + name
	for _, c := range containers {
		for _, cname := range c.Names {
			if name == cname {
				return &c.ID
			}
		}
	}

	return nil
}

func (sv Supervisor) runContainer(name, image string, hc *docker.HostConfig,
	args []string) error {
	log.Info("Booting Container: %s", name)

	if err := sv.pullImage(image); err != nil {
		return err
	}

	id := sv.getContainer(name)
	if id == nil {
		container, err := sv.dk.CreateContainer(docker.CreateContainerOptions{
			Name: name,
			Config: &docker.Config{
				Image: image,
				Cmd:   args,
			},
		})
		if err != nil {
			return err
		}
		id = &container.ID
	}

	err := sv.dk.StartContainer(*id, hc)
	if err != nil {
		if _, ok := err.(*docker.ContainerAlreadyRunning); ok {
			log.Info("Container already running: %s", name)
			return nil
		} else {
			return err
		}
	}

	log.Info("Successfully booted %s", name)
	return nil
}

func (sv Supervisor) execInContainer(name string, cmd []string) error {
	log.Info("Executing in container: %s", name)

	id := sv.getContainer(name)
	exec, err := sv.dk.CreateExec(docker.CreateExecOptions{
		Container: *id,
		Cmd:       cmd})

	if err != nil {
		return err
	}

	err = sv.dk.StartExec(exec.ID, docker.StartExecOptions{})

	if err != nil {
		return err
	}

	return nil
}

func (sv Supervisor) removeContainer(name string) error {
	log.Info("Remove Container: %s", name)

	id := sv.getContainer(name)
	if id == nil {
		log.Info("Failed to stop missing container: %s", name)
		return nil
	}

	err := sv.dk.RemoveContainer(docker.RemoveContainerOptions{ID: *id, Force: true})
	if err != nil {
		return err
	}

	return nil
}

func (sv *Supervisor) Configure(cfg MinionConfig) error {
	sv.cfg = cfg

	/* XXX: We should really being using an ID that we get from the master. */
	hostname, err := os.Hostname()
	if err != nil {
		return err
	}

	if err := sv.runContainer("ovsdb-server", OVS_OVSDB, ovsHC(), nil); err != nil {
		return err
	}

	var etcdArgs []string
	switch cfg.Role {
	case MinionConfig_MASTER:
		advertiseClient := fmt.Sprintf("http://%s:2379", cfg.PrivateIP)
		listenClient := fmt.Sprintf("http://0.0.0.0:2379")

		initialAdvertisePeer := fmt.Sprintf("http://%s:2380", cfg.PrivateIP)
		listenPeer := fmt.Sprintf("http://%s:2380", cfg.PrivateIP)

		etcdArgs = []string{"--name=" + hostname,
			"--discovery=" + cfg.EtcdToken,
			"--advertise-client-urls=" + advertiseClient,
			"--initial-advertise-peer-urls=" + initialAdvertisePeer,
			"--listen-client-urls=" + listenClient,
			"--listen-peer-urls=" + listenPeer}

	case MinionConfig_WORKER:
		etcdArgs = []string{"--discovery=" + cfg.EtcdToken, "--proxy=on"}

		hc := ovsHC()
		hc.Privileged = true
		err := sv.runContainer("ovs-vswitchd", OVS_VSWITCHD, hc, nil)
		if err != nil {
			return err
		}

		err = sv.runContainer("ovn-controller", OVN_CONTROLLER, ovsHC(), nil)
		if err != nil {
			return err
		}

	default:
		panic(fmt.Sprintf("Unknown minion role: %d", cfg.Role))
	}

	if err := sv.runContainer("etcd", ETCD, defaultHC(), etcdArgs); err != nil {
		return err
	}

	return nil
}

func (sv Supervisor) WatchLeaderChannel(leaderChan chan *string) {
	for {
		leaderIp := <-leaderChan
		if leaderIp == nil {
			log.Warning("Leader went away.")
			/* XXX: Handle this. */
			continue
		}
		/* XXX: update ovsdb, kubernetes, etc... */
	}
}

func (sv Supervisor) WatchElectionChannel(electionChan chan bool) {
	for {
		if <-electionChan {
			err := sv.runContainer("ovn-northd", OVN_NORTHD, ovsHC(), nil)
			if err != nil {
				/* XXX: If we fail to boot ovn-northd, we should give up
				 * our leadership somehow.  This ties into the general
				 * problem of monitoring health. */
				log.Warning("Failed to boot ovn-northd: %s", err)
			}
		} else {
			err := sv.removeContainer("ovn-northd")
			if err != nil {
				log.Warning("Failed to remove ovn-northd: %s", err)
			}
		}
	}
}
<commit_msg>supervisor: Boot kubernetes on leader change<commit_after>package main

import (
	"bytes"
	"fmt"
	"net/http"
	"os"
	"time"

	. "github.com/NetSys/di/minion/proto"
	"github.com/fsouza/go-dockerclient"
)

const ETCD = "quay.io/coreos/etcd:v2.1.3"
const KUBELET = "quay.io/netsys/kubelet"
const OVN_NORTHD = "quay.io/netsys/ovn-northd"
const OVN_CONTROLLER = "quay.io/netsys/ovn-controller"
const OVS_VSWITCHD = "quay.io/netsys/ovs-vswitchd"
const OVS_OVSDB = "quay.io/netsys/ovsdb-server"
const DOCKER_SOCK_PATH = "unix:///var/run/docker.sock"

type Supervisor struct {
	dk  *docker.Client
	cfg MinionConfig
}

func kubeHC() *docker.HostConfig {
	binds := []string{
		"/var/run:/var/run:rw",
		"/var/lib/docker:/var/lib/docker:rw",
		"/etc/docker:/etc/docker:rw",
		"/dev:/dev:rw",
		"/sys:/sys:rw"}
	return &docker.HostConfig{
		Binds:       binds,
		Privileged:  true,
		NetworkMode: "host",
		/* Use PID=host per the kubernetes documentation */
		PidMode: "host"}
}

func defaultHC() *docker.HostConfig {
	return &docker.HostConfig{
		NetworkMode: "host",
		Binds:       []string{"/usr/share/ca-certificates:/etc/ssl/certs"}}
}

func ovsHC() *docker.HostConfig {
	hc := defaultHC()
	hc.VolumesFrom = []string{"ovsdb-server"}
	return hc
}

func (sv Supervisor) pullImage(image string) error {
	return sv.dk.PullImage(docker.PullImageOptions{
		Repository: image},
		docker.AuthConfiguration{})
}

func NewSupervisor() (*Supervisor, error) {
	dk, err := docker.NewClient(DOCKER_SOCK_PATH)
	if err != nil {
		log.Warning("Failed to create docker client: %s", err)
		return nil, err
	}

	sv := &Supervisor{dk: dk}

	images := []string{OVS_OVSDB, OVN_NORTHD, OVN_CONTROLLER, OVS_VSWITCHD, ETCD}
	for _, image := range images {
		go sv.pullImage(image)
	}

	return sv, nil
}

func (sv Supervisor) getContainer(name string) *string {
	containers, err := sv.dk.ListContainers(docker.ListContainersOptions{All: true})
	if err != nil {
		log.Warning("Failed to list containers: %s", err)
		return nil
	}

	name = "/" + name
	for _, c := range containers {
		for _, cname := range c.Names {
			if name == cname {
				return &c.ID
			}
		}
	}

	return nil
}

func (sv Supervisor) runContainer(name, image string, hc *docker.HostConfig,
	args []string) error {
	log.Info("Booting Container: %s", name)

	if err := sv.pullImage(image); err != nil {
		return err
	}

	id := sv.getContainer(name)
	if id == nil {
		container, err := sv.dk.CreateContainer(docker.CreateContainerOptions{
			Name: name,
			Config: &docker.Config{
				Image: image,
				Cmd:   args,
			},
		})
		if err != nil {
			return err
		}
		id = &container.ID
	}

	err := sv.dk.StartContainer(*id, hc)
	if err != nil {
		if _, ok := err.(*docker.ContainerAlreadyRunning); ok {
			log.Info("Container already running: %s", name)
			return nil
		} else {
			return err
		}
	}

	log.Info("Successfully booted %s", name)
	return nil
}

func (sv Supervisor) execInContainer(name string, cmd []string) error {
	log.Info("Executing in container: %s", name)

	id := sv.getContainer(name)
	exec, err := sv.dk.CreateExec(docker.CreateExecOptions{
		Container: *id,
		Cmd:       cmd})

	if err != nil {
		return err
	}

	err = sv.dk.StartExec(exec.ID, docker.StartExecOptions{})

	if err != nil {
		return err
	}

	return nil
}

func (sv Supervisor) removeContainer(name string) error {
	log.Info("Remove Container: %s", name)

	id := sv.getContainer(name)
	if id == nil {
		log.Info("Failed to stop missing container: %s", name)
		return nil
	}

	err := sv.dk.RemoveContainer(docker.RemoveContainerOptions{ID: *id, Force: true})
	if err != nil {
		return err
	}

	return nil
}

func (sv *Supervisor) Configure(cfg MinionConfig) error {
	sv.cfg = cfg

	/* XXX: We should really being using an ID that we get from the master. */
	hostname, err := os.Hostname()
	if err != nil {
		return err
	}

	if err := sv.runContainer("ovsdb-server", OVS_OVSDB, ovsHC(), nil); err != nil {
		return err
	}

	var etcdArgs []string
	switch cfg.Role {
	case MinionConfig_MASTER:
		advertiseClient := fmt.Sprintf("http://%s:2379", cfg.PrivateIP)
		listenClient := fmt.Sprintf("http://0.0.0.0:2379")

		initialAdvertisePeer := fmt.Sprintf("http://%s:2380", cfg.PrivateIP)
		listenPeer := fmt.Sprintf("http://%s:2380", cfg.PrivateIP)

		etcdArgs = []string{"--name=" + hostname,
			"--discovery=" + cfg.EtcdToken,
			"--advertise-client-urls=" + advertiseClient,
			"--initial-advertise-peer-urls=" + initialAdvertisePeer,
			"--listen-client-urls=" + listenClient,
			"--listen-peer-urls=" + listenPeer}

	case MinionConfig_WORKER:
		etcdArgs = []string{"--discovery=" + cfg.EtcdToken, "--proxy=on"}

		hc := ovsHC()
		hc.Privileged = true
		err := sv.runContainer("ovs-vswitchd", OVS_VSWITCHD, hc, nil)
		if err != nil {
			return err
		}

		err = sv.runContainer("ovn-controller", OVN_CONTROLLER, ovsHC(), nil)
		if err != nil {
			return err
		}

	default:
		panic(fmt.Sprintf("Unknown minion role: %d", cfg.Role))
	}

	if err := sv.runContainer("etcd", ETCD, defaultHC(), etcdArgs); err != nil {
		return err
	}

	if cfg.Role == MinionConfig_MASTER {
		kubeletArgs := []string{"/usr/bin/boot-master", sv.cfg.PrivateIP}
		err := sv.runContainer("di-kubelet", KUBELET, kubeHC(), kubeletArgs)
		if err != nil {
			return err
		}

		/* Create the kube-system namespace. This is where info about all of the
		 * kubernetes pods, conatiners, replication controllers, etc... will
		 * live. */
		go func() {
			body := `{"apiVersion":"v1","kind":"Namespace","metadata":{"name":"kube-system"}}`
			url := "http://127.0.0.1:9000/api/v1/namespaces"
			ctype := "application/json"
			for {
				_, err := http.Post(url, ctype, bytes.NewBuffer([]byte(body)))
				if err == nil {
					break
				}
				time.Sleep(5 * time.Second)
			}
		}()

		/* XXX: Should we create a separate namespace for user pods, e.g.
		 * the AWS namespace we use right now? There is probably a convention
		 * for this mentioned somewhere in the kubernetes community. */
	}

	return nil
}

func (sv Supervisor) WatchLeaderChannel(leaderChan chan *string) {
	for {
		leaderIp := <-leaderChan
		if leaderIp == nil {
			log.Warning("Leader went away.")
			/* XXX: Handle this. */
			continue
		}

		kubeletArgs := []string{"/usr/bin/boot-worker", sv.cfg.PrivateIP,
			*leaderIp}
		err := sv.runContainer("di-kubelet", KUBELET, kubeHC(), kubeletArgs)
		if err != nil {
			log.Warning("Failed to boot di-kubelet: %s", err)
			if err := sv.removeContainer("di-kubelet"); err != nil {
				log.Warning("Failed to remove di-kubelet: %s", err)
			}
		}
	}
}

func (sv Supervisor) WatchElectionChannel(electionChan chan bool) {
	for {
		if <-electionChan {
			err := sv.runContainer("ovn-northd", OVN_NORTHD, ovsHC(), nil)
			if err != nil {
				/* XXX: If we fail to boot ovn-northd, we should give up
				 * our leadership somehow.  This ties into the general
				 * problem of monitoring health. */
				log.Warning("Failed to boot ovn-northd: %s", err)
			}
		} else {
			err := sv.removeContainer("ovn-northd")
			if err != nil {
				log.Warning("Failed to remove ovn-northd: %s", err)
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"net/url"
	"os"
	"strings"

	log "github.com/avabot/ava/Godeps/_workspace/src/github.com/Sirupsen/logrus"
	"github.com/avabot/ava/Godeps/_workspace/src/github.com/garyburd/go-oauth/oauth"
	"github.com/avabot/ava/Godeps/_workspace/src/github.com/jmoiron/sqlx"
	"github.com/avabot/ava/shared/datatypes"
	"github.com/avabot/ava/shared/knowledge"
	"github.com/avabot/ava/shared/language"
	"github.com/avabot/ava/shared/pkg"
)

type Yelp string

type client struct {
	client oauth.Client
	token  oauth.Credentials
}

type yelpResp struct {
	Businesses []struct {
		Name         string
		ImageURL     string `json:"image_url"`
		MobileURL    string `json:"mobile_url"`
		DisplayPhone string `json:"display_phone"`
		Distance     int
		Rating       float64
		Location     struct {
			City           string
			DisplayAddress []string `json:"display_address"`
		}
	}
}

var port = flag.Int("port", 0, "Port used to communicate with Ava.")
var ErrNoBusinesses = errors.New("no businesses")

var c client
var db *sqlx.DB
var p *pkg.Pkg
var l *log.Entry

func main() {
	flag.Parse()
	l = log.WithFields(log.Fields{
		"pkg": "yelp",
	})
	c.client.Credentials.Token = os.Getenv("YELP_CONSUMER_KEY")
	c.client.Credentials.Secret = os.Getenv("YELP_CONSUMER_SECRET")
	c.token.Token = os.Getenv("YELP_TOKEN")
	c.token.Secret = os.Getenv("YELP_TOKEN_SECRET")
	var err error
	db, err = pkg.ConnectDB()
	if err != nil {
		l.Fatalln(err)
	}
	trigger := &dt.StructuredInput{
		Commands: []string{
			"find",
			"where",
			"show",
			"recommend",
			"recommendation",
			"recommendations",
		},
		Objects: language.Foods(),
	}
	p, err = pkg.NewPackage("yelp", *port, trigger)
	if err != nil {
		l.Fatalln("building", err)
	}
	yelp := new(Yelp)
	if err := p.Register(yelp); err != nil {
		l.Fatalln("registering", err)
	}
}

func (t *Yelp) Run(m *dt.Msg, respMsg *dt.RespMsg) error {
	// NOTE optional: get state before this package was run, enabling
	// chaining. For instance, allow passing context from a FourSquare
	// result into this. See the opening lines of FollowUp() for an example
	resp := m.NewResponse()
	resp.State = map[string]interface{}{
		"query":      "",
		"location":   "",
		"offset":     float64(0),
		"businesses": []interface{}{},
	}
	si := m.Input.StructuredInput
	query := ""
	for _, o := range si.Objects {
		query += o + " "
	}
	for _, p := range si.Places {
		query += p + " "
	}
	resp.State["query"] = query
	if len(si.Places) == 0 {
		l.Infoln("no place entered, getting location")
		loc, question, err := knowledge.GetLocation(db, m.User)
		if err != nil {
			return err
		}
		if len(question) > 0 {
			if loc != nil && len(loc.Name) > 0 {
				resp.State["location"] = loc.Name
			}
			resp.Sentence = question
			return p.SaveResponse(respMsg, resp)
		}
		resp.State["location"] = loc.Name
	}
	// Occurs in the case of "nearby" or other contextual place terms, where
	// no previous context was available to expand it.
	if len(resp.State["location"].(string)) == 0 {
		loc, question, err := knowledge.GetLocation(db, m.User)
		if err != nil {
			return err
		}
		if len(question) > 0 {
			if loc != nil && len(loc.Name) > 0 {
				resp.State["location"] = loc.Name
			}
			resp.Sentence = question
			return p.SaveResponse(respMsg, resp)
		}
		resp.State["location"] = loc.Name
	}
	if err := t.searchYelp(resp); err != nil {
		l.WithField("fn", "searchYelp").Errorln(err)
	}
	return p.SaveResponse(respMsg, resp)
}

// FollowUp handles dialog question/answers and additional user queries
func (t *Yelp) FollowUp(m *dt.Msg,
	respMsg *dt.RespMsg) error {
	// Retrieve the conversation's context
	if err := m.GetLastResponse(db); err != nil {
		l.WithField("fn", "GetLastResponse").Errorln(err)
		return err
	}
	resp := m.NewResponse()

	// First we handle dialog. If we asked for a location, use the response
	if resp.State["location"] == "" {
		loc := m.Input.StructuredInput.All()
		resp.State["location"] = loc
		if err := t.searchYelp(resp); err != nil {
			l.WithField("fn", "searchYelp").Errorln(err)
		}
		return p.SaveResponse(respMsg, resp)
	}

	// If no businesses are returned inform the user now
	if resp.State["businesses"] != nil &&
		len(resp.State["businesses"].([]interface{})) == 0 {
		resp.Sentence = "I couldn't find anything like that"
		return p.SaveResponse(respMsg, resp)
	}

	// Responses were returned, and the user has asked this package an
	// additional query. Handle the query by keyword
	words := strings.Fields(m.Input.Sentence)
	offI := int(resp.State["offset"].(float64))
	var s string
	for _, w := range words {
		w = strings.TrimRight(w, ").,;?!:")
		switch strings.ToLower(w) {
		case "rated", "rating", "review", "recommend", "recommended":
			s = fmt.Sprintf("It has a %s star review on Yelp",
				getRating(resp, offI))
			resp.Sentence = s
		case "number", "phone":
			s = getPhone(resp, offI)
			resp.Sentence = s
		case "call":
			s = fmt.Sprintf("You can reach them here: %s",
				getPhone(resp, offI))
			resp.Sentence = s
		case "information", "info":
			s = fmt.Sprintf("Here's some more info: %s",
				getURL(resp, offI))
			resp.Sentence = s
		case "where", "location", "address", "direction", "directions",
			"addr":
			s = fmt.Sprintf("It's at %s", getAddress(resp, offI))
			resp.Sentence = s
		case "pictures", "pic", "pics":
			s = fmt.Sprintf("I found some pics here: %s",
				getURL(resp, offI))
			resp.Sentence = s
		case "menu", "have":
			s = fmt.Sprintf("Yelp might have a menu... %s",
				getURL(resp, offI))
			resp.Sentence = s
		case "not", "else", "no", "anything", "something":
			resp.State["offset"] = float64(offI + 1)
			if err := t.searchYelp(resp); err != nil {
				l.WithField("fn", "searchYelp").Errorln(err)
			}
		// TODO perhaps handle this case and "thanks" at the AVA level?
		// with bayesian classification
		case "good", "great", "yes", "perfect":
			// TODO feed into learning engine
			resp.Sentence = language.Positive()
		case "thanks", "thank":
			resp.Sentence = language.Welcome()
		}
		if len(resp.Sentence) > 0 {
			return p.SaveResponse(respMsg, resp)
		}
	}
	return p.SaveResponse(respMsg, resp)
}

func getRating(r *dt.Resp, offset int) string {
	businesses := r.State["businesses"].([]interface{})
	firstBusiness := businesses[offset].(map[string]interface{})
	return fmt.Sprintf("%.1f", firstBusiness["Rating"].(float64))
}

func getURL(r *dt.Resp, offset int) string {
	businesses := r.State["businesses"].([]interface{})
	firstBusiness := businesses[offset].(map[string]interface{})
	return firstBusiness["mobile_url"].(string)
}

func getPhone(r *dt.Resp, offset int) string {
	businesses := r.State["businesses"].([]interface{})
	firstBusiness := businesses[offset].(map[string]interface{})
	return firstBusiness["display_phone"].(string)
}

func getAddress(r *dt.Resp, offset int) string {
	businesses := r.State["businesses"].([]interface{})
	firstBusiness := businesses[offset].(map[string]interface{})
	location := firstBusiness["Location"].(map[string]interface{})
	dispAddr := location["display_address"].([]interface{})
	if len(dispAddr) > 1 {
		str1 := dispAddr[0].(string)
		str2 := dispAddr[1].(string)
		return fmt.Sprintf("%s in %s", str1, str2)
	}
	return dispAddr[0].(string)
}

func (c *client) get(urlStr string, params url.Values, v interface{}) error {
	resp, err := c.client.Get(nil, &c.token, urlStr, params)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return fmt.Errorf("yelp status %d", resp.StatusCode)
	}
	return json.NewDecoder(resp.Body).Decode(v)
}

func (t *Yelp) searchYelp(resp *dt.Resp) error {
	q := resp.State["query"].(string)
	loc := resp.State["location"].(string)
	offset := resp.State["offset"].(float64)
	l.WithFields(log.Fields{
		"q":      q,
		"loc":    loc,
		"offset": offset,
	}).Infoln("searching yelp")
	form := url.Values{
		"term":     {q},
		"location": {loc},
		"limit":    {fmt.Sprintf("%.0f", offset+1)},
	}
	var data yelpResp
	err := c.get("http://api.yelp.com/v2/search", form, &data)
	if err != nil {
		resp.Sentence = "I can't find that for you now. " +
			"Let's try again later."
		l.WithField("fn", "get").Errorln(err)
		return err
	}
	resp.State["businesses"] = data.Businesses
	if len(data.Businesses) == 0 {
		resp.Sentence = "I couldn't find any places like that nearby."
		return nil
	}
	offI := int(offset)
	if len(data.Businesses) <= offI {
		resp.Sentence = "That's all I could find."
		return nil
	}
	b := data.Businesses[offI]
	addr := ""
	if len(b.Location.DisplayAddress) > 0 {
		addr = b.Location.DisplayAddress[0]
	}
	if offI == 0 {
		resp.Sentence = "Ok. How does this place look? " + b.Name +
			" at " + addr
	} else {
		resp.Sentence = fmt.Sprintf("What about %s instead?", b.Name)
	}
	return nil
}
<commit_msg>Improve Yelp failed search handling<commit_after>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"net/url"
	"os"
	"strings"

	log "github.com/avabot/ava/Godeps/_workspace/src/github.com/Sirupsen/logrus"
	"github.com/avabot/ava/Godeps/_workspace/src/github.com/garyburd/go-oauth/oauth"
	"github.com/avabot/ava/Godeps/_workspace/src/github.com/jmoiron/sqlx"
	"github.com/avabot/ava/shared/datatypes"
	"github.com/avabot/ava/shared/knowledge"
	"github.com/avabot/ava/shared/language"
	"github.com/avabot/ava/shared/pkg"
)

type Yelp string

type client struct {
	client oauth.Client
	token  oauth.Credentials
}

type yelpResp struct {
	Businesses []struct {
		Name         string
		ImageURL     string `json:"image_url"`
		MobileURL    string `json:"mobile_url"`
		DisplayPhone string `json:"display_phone"`
		Distance     int
		Rating       float64
		Location     struct {
			City           string
			DisplayAddress []string `json:"display_address"`
		}
	}
}

var port = flag.Int("port", 0, "Port used to communicate with Ava.")
var ErrNoBusinesses = errors.New("no businesses")

var c client
var db *sqlx.DB
var p *pkg.Pkg
var l *log.Entry

func main() {
	flag.Parse()
	l = log.WithFields(log.Fields{
		"pkg": "yelp",
	})
	c.client.Credentials.Token = os.Getenv("YELP_CONSUMER_KEY")
	c.client.Credentials.Secret = os.Getenv("YELP_CONSUMER_SECRET")
	c.token.Token = os.Getenv("YELP_TOKEN")
	c.token.Secret = os.Getenv("YELP_TOKEN_SECRET")
	var err error
	db, err = pkg.ConnectDB()
	if err != nil {
		l.Fatalln(err)
	}
	trigger := &dt.StructuredInput{
		Commands: []string{
			"find",
			"where",
			"show",
			"recommend",
			"recommendation",
			"recommendations",
		},
		Objects: language.Foods(),
	}
	p, err = pkg.NewPackage("yelp", *port, trigger)
	if err != nil {
		l.Fatalln("building", err)
	}
	yelp := new(Yelp)
	if err := p.Register(yelp); err != nil {
		l.Fatalln("registering", err)
	}
}

func (t *Yelp) Run(m *dt.Msg, respMsg *dt.RespMsg) error {
	// NOTE optional: get state before this package was run, enabling
	// chaining. For instance, allow passing context from a FourSquare
	// result into this. See the opening lines of FollowUp() for an example
	resp := m.NewResponse()
	resp.State = map[string]interface{}{
		"query":      "",
		"location":   "",
		"offset":     float64(0),
		"businesses": []interface{}{},
	}
	si := m.Input.StructuredInput
	query := ""
	for _, o := range si.Objects {
		query += o + " "
	}
	for _, p := range si.Places {
		query += p + " "
	}
	resp.State["query"] = query
	if len(si.Places) == 0 {
		l.Infoln("no place entered, getting location")
		loc, question, err := knowledge.GetLocation(db, m.User)
		if err != nil {
			return err
		}
		if len(question) > 0 {
			if loc != nil && len(loc.Name) > 0 {
				resp.State["location"] = loc.Name
			}
			resp.Sentence = question
			return p.SaveResponse(respMsg, resp)
		}
		resp.State["location"] = loc.Name
	}
	// Occurs in the case of "nearby" or other contextual place terms, where
	// no previous context was available to expand it.
	if len(resp.State["location"].(string)) == 0 {
		loc, question, err := knowledge.GetLocation(db, m.User)
		if err != nil {
			return err
		}
		if len(question) > 0 {
			if loc != nil && len(loc.Name) > 0 {
				resp.State["location"] = loc.Name
			}
			resp.Sentence = question
			return p.SaveResponse(respMsg, resp)
		}
		resp.State["location"] = loc.Name
	}
	if err := t.searchYelp(resp); err != nil {
		l.WithField("fn", "searchYelp").Errorln(err)
	}
	return p.SaveResponse(respMsg, resp)
}

// FollowUp handles dialog question/answers and additional user queries
func (t *Yelp) FollowUp(m *dt.Msg,
	respMsg *dt.RespMsg) error {
	// Retrieve the conversation's context
	if err := m.GetLastResponse(db); err != nil {
		l.WithField("fn", "GetLastResponse").Errorln(err)
		return err
	}
	resp := m.NewResponse()

	// First we handle dialog. If we asked for a location, use the response
	if resp.State["location"] == "" {
		loc := m.Input.StructuredInput.All()
		resp.State["location"] = loc
		if err := t.searchYelp(resp); err != nil {
			l.WithField("fn", "searchYelp").Errorln(err)
		}
		return p.SaveResponse(respMsg, resp)
	}

	// If no businesses are returned inform the user now
	if resp.State["businesses"] != nil &&
		len(resp.State["businesses"].([]interface{})) == 0 {
		resp.Sentence = "I couldn't find anything like that"
		return p.SaveResponse(respMsg, resp)
	}

	// Responses were returned, and the user has asked this package an
	// additional query. Handle the query by keyword
	words := strings.Fields(m.Input.Sentence)
	offI := int(resp.State["offset"].(float64))
	var s string
	for _, w := range words {
		w = strings.TrimRight(w, ").,;?!:")
		switch strings.ToLower(w) {
		case "rated", "rating", "review", "recommend", "recommended":
			s = fmt.Sprintf("It has a %s star review on Yelp",
				getRating(resp, offI))
			resp.Sentence = s
		case "number", "phone":
			s = getPhone(resp, offI)
			resp.Sentence = s
		case "call":
			s = fmt.Sprintf("You can reach them here: %s",
				getPhone(resp, offI))
			resp.Sentence = s
		case "information", "info":
			s = fmt.Sprintf("Here's some more info: %s",
				getURL(resp, offI))
			resp.Sentence = s
		case "where", "location", "address", "direction", "directions",
			"addr":
			s = fmt.Sprintf("It's at %s", getAddress(resp, offI))
			resp.Sentence = s
		case "pictures", "pic", "pics":
			s = fmt.Sprintf("I found some pics here: %s",
				getURL(resp, offI))
			resp.Sentence = s
		case "menu", "have":
			s = fmt.Sprintf("Yelp might have a menu... %s",
				getURL(resp, offI))
			resp.Sentence = s
		case "not", "else", "no", "anything", "something":
			resp.State["offset"] = float64(offI + 1)
			if err := t.searchYelp(resp); err != nil {
				l.WithField("fn", "searchYelp").Errorln(err)
			}
		// TODO perhaps handle this case and "thanks" at the AVA level?
		// with bayesian classification
		case "good", "great", "yes", "perfect":
			// TODO feed into learning engine
			resp.Sentence = language.Positive()
		case "thanks", "thank":
			resp.Sentence = language.Welcome()
		}
		if len(resp.Sentence) > 0 {
			return p.SaveResponse(respMsg, resp)
		}
	}
	return p.SaveResponse(respMsg, resp)
}

func getRating(r *dt.Resp, offset int) string {
	businesses := r.State["businesses"].([]interface{})
	firstBusiness := businesses[offset].(map[string]interface{})
	return fmt.Sprintf("%.1f", firstBusiness["Rating"].(float64))
}

func getURL(r *dt.Resp, offset int) string {
	businesses := r.State["businesses"].([]interface{})
	firstBusiness := businesses[offset].(map[string]interface{})
	return firstBusiness["mobile_url"].(string)
}

func getPhone(r *dt.Resp, offset int) string {
	businesses := r.State["businesses"].([]interface{})
	firstBusiness := businesses[offset].(map[string]interface{})
	return firstBusiness["display_phone"].(string)
}

func getAddress(r *dt.Resp, offset int) string {
	businesses := r.State["businesses"].([]interface{})
	firstBusiness := businesses[offset].(map[string]interface{})
	location := firstBusiness["Location"].(map[string]interface{})
	dispAddr := location["display_address"].([]interface{})
	if len(dispAddr) > 1 {
		str1 := dispAddr[0].(string)
		str2 := dispAddr[1].(string)
		return fmt.Sprintf("%s in %s", str1, str2)
	}
	return dispAddr[0].(string)
}

func (c *client) get(urlStr string, params url.Values, v interface{}) error {
	resp, err := c.client.Get(nil, &c.token, urlStr, params)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return fmt.Errorf("yelp status %d", resp.StatusCode)
	}
	return json.NewDecoder(resp.Body).Decode(v)
}

func (t *Yelp) searchYelp(resp *dt.Resp) error {
	q := resp.State["query"].(string)
	loc := resp.State["location"].(string)
	offset := resp.State["offset"].(float64)
	l.WithFields(log.Fields{
		"q":      q,
		"loc":    loc,
		"offset": offset,
	}).Infoln("searching yelp")
	form := url.Values{
		"term":     {q},
		"location": {loc},
		"limit":    {fmt.Sprintf("%.0f", offset+1)},
	}
	var data yelpResp
	err := c.get("http://api.yelp.com/v2/search", form, &data)
	if err != nil {
		/*
			resp.Sentence = "I can't find that for you now. " +
				"Let's try again later."
			l.WithField("fn", "get").Errorln(err)
			return err
		*/
		// return for confused response, given Yelp errors are rare, but
		// unintentional runs of Yelp queries are much more common
		return nil
	}
	resp.State["businesses"] = data.Businesses
	if len(data.Businesses) == 0 {
		resp.Sentence = "I couldn't find any places like that nearby."
		return nil
	}
	offI := int(offset)
	if len(data.Businesses) <= offI {
		resp.Sentence = "That's all I could find."
		return nil
	}
	b := data.Businesses[offI]
	addr := ""
	if len(b.Location.DisplayAddress) > 0 {
		addr = b.Location.DisplayAddress[0]
	}
	if offI == 0 {
		resp.Sentence = "Ok. How does this place look? " + b.Name +
			" at " + addr
	} else {
		resp.Sentence = fmt.Sprintf("What about %s instead?", b.Name)
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cmd

import (
	"fmt"
	"io"
	"net/url"
	"time"

	"github.com/golang/glog"
	"github.com/spf13/cobra"

	apiv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/runtime"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/tools/remotecommand"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/kubectl/cmd/templates"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/kubectl/genericclioptions"
	"k8s.io/kubernetes/pkg/kubectl/genericclioptions/resource"
	"k8s.io/kubernetes/pkg/kubectl/polymorphichelpers"
	"k8s.io/kubernetes/pkg/kubectl/scheme"
	"k8s.io/kubernetes/pkg/kubectl/util/i18n"
)

var (
	attachExample = templates.Examples(i18n.T(`
		# Get output from running pod 123456-7890, using the first container by default
		kubectl attach 123456-7890

		# Get output from ruby-container from pod 123456-7890
		kubectl attach 123456-7890 -c ruby-container

		# Switch to raw terminal mode, sends stdin to 'bash' in ruby-container from pod 123456-7890
		# and sends stdout/stderr from 'bash' back to the client
		kubectl attach 123456-7890 -c ruby-container -i -t

		# Get output from the first pod of a ReplicaSet named nginx
		kubectl attach rs/nginx
		`))
)

const (
	defaultPodAttachTimeout = 60 * time.Second
	defaultPodLogsTimeout   = 20 * time.Second
)

// AttachOptions declare the arguments accepted by the Exec command
type AttachOptions struct {
	StreamOptions

	// whether to disable the use of standard error when streaming output from tty
	DisableStderr bool

	CommandName       string
	SuggestedCmdUsage string

	Pod *apiv1.Pod

	AttachFunc       func(*AttachOptions, *apiv1.Container, bool, remotecommand.TerminalSizeQueue) func() error
	Resources        []string
	Builder          func() *resource.Builder
	AttachablePodFn  polymorphichelpers.AttachableLogsForObjectFunc
	restClientGetter genericclioptions.RESTClientGetter

	Attach        RemoteAttach
	GetPodTimeout time.Duration
	Config        *restclient.Config
}

func NewAttachOptions(streams genericclioptions.IOStreams) *AttachOptions {
	return &AttachOptions{
		StreamOptions: StreamOptions{
			IOStreams: streams,
		},
		Attach: &DefaultRemoteAttach{},
	}
}

func NewCmdAttach(f cmdutil.Factory, streams genericclioptions.IOStreams) *cobra.Command {
	o := NewAttachOptions(streams)
	cmd := &cobra.Command{
		Use: "attach (POD | TYPE/NAME) -c CONTAINER",
		DisableFlagsInUseLine: true,
		Short:   i18n.T("Attach to a running container"),
		Long:    "Attach to a process that is already running inside an existing container.",
		Example: attachExample,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, cmd, args))
			cmdutil.CheckErr(o.Validate())
			cmdutil.CheckErr(o.Run())
		},
	}
	cmdutil.AddPodRunningTimeoutFlag(cmd, defaultPodAttachTimeout)
	cmd.Flags().StringVarP(&o.ContainerName, "container", "c", o.ContainerName, "Container name. If omitted, the first container in the pod will be chosen")
	cmd.Flags().BoolVarP(&o.Stdin, "stdin", "i", o.Stdin, "Pass stdin to the container")
	cmd.Flags().BoolVarP(&o.TTY, "tty", "t", o.TTY, "Stdin is a TTY")
	return cmd
}

// RemoteAttach defines the interface accepted by the Attach command - provided for test stubbing
type RemoteAttach interface {
	Attach(method string, url *url.URL, config *restclient.Config, stdin io.Reader, stdout, stderr io.Writer, tty bool, terminalSizeQueue remotecommand.TerminalSizeQueue) error
}

func defaultAttachFunc(o *AttachOptions, containerToAttach *apiv1.Container, raw bool, sizeQueue remotecommand.TerminalSizeQueue) func() error {
	return func() error {
		restClient, err := restclient.RESTClientFor(o.Config)
		if err != nil {
			return err
		}
		req := restClient.Post().
			Resource("pods").
			Name(o.Pod.Name).
			Namespace(o.Pod.Namespace).
			SubResource("attach")
		req.VersionedParams(&apiv1.PodAttachOptions{
			Container: containerToAttach.Name,
			Stdin:     o.Stdin,
			Stdout:    o.Out != nil,
			Stderr:    !o.DisableStderr,
			TTY:       raw,
		}, legacyscheme.ParameterCodec)

		return o.Attach.Attach("POST", req.URL(), o.Config, o.In, o.Out, o.ErrOut, raw, sizeQueue)
	}
}

// DefaultRemoteAttach is the standard implementation of attaching
type DefaultRemoteAttach struct{}

func (*DefaultRemoteAttach) Attach(method string, url *url.URL, config *restclient.Config, stdin io.Reader, stdout, stderr io.Writer, tty bool, terminalSizeQueue remotecommand.TerminalSizeQueue) error {
	exec, err := remotecommand.NewSPDYExecutor(config, method, url)
	if err != nil {
		return err
	}
	return exec.Stream(remotecommand.StreamOptions{
		Stdin:             stdin,
		Stdout:            stdout,
		Stderr:            stderr,
		Tty:               tty,
		TerminalSizeQueue: terminalSizeQueue,
	})
}

// Complete verifies command line arguments and loads data from the command environment
func (o *AttachOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	var err error
	o.Namespace, _, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}

	o.AttachablePodFn = polymorphichelpers.AttachablePodForObjectFn

	o.GetPodTimeout, err = cmdutil.GetPodRunningTimeoutFlag(cmd)
	if err != nil {
		return cmdutil.UsageErrorf(cmd, err.Error())
	}

	o.Builder = f.NewBuilder
	o.Resources = args
	o.restClientGetter = f

	fullCmdName := ""
	cmdParent := cmd.Parent()
	if cmdParent != nil {
		fullCmdName = cmdParent.CommandPath()
	}
	if len(fullCmdName) > 0 && cmdutil.IsSiblingCommandExists(cmd, "describe") {
		o.SuggestedCmdUsage = fmt.Sprintf("Use '%s describe pod/%s -n %s' to see all of the containers in this pod.", fullCmdName, o.PodName, o.Namespace)
	}

	config, err := f.ToRESTConfig()
	if err != nil {
		return err
	}
	o.Config = config

	o.AttachFunc = defaultAttachFunc

	if o.CommandName == "" {
		o.CommandName = cmd.CommandPath()
	}

	return nil
}

// Validate checks that the provided attach options are specified.
func (o *AttachOptions) Validate() error {
	if len(o.Resources) == 0 {
		return fmt.Errorf("at least 1 argument is required for attach")
	}
	if len(o.Resources) > 2 {
		return fmt.Errorf("expected POD, TYPE/NAME, or TYPE NAME, (at most 2 arguments) saw %d: %v", len(o.Resources), o.Resources)
	}
	if o.GetPodTimeout <= 0 {
		return fmt.Errorf("--pod-running-timeout must be higher than zero")
	}

	return nil
}

// Run executes a validated remote execution against a pod.
func (o *AttachOptions) Run() error {
	if o.Pod == nil {
		b := o.Builder().
			WithScheme(scheme.Scheme, scheme.Scheme.PrioritizedVersionsAllGroups()...).
			NamespaceParam(o.Namespace).DefaultNamespace()

		switch len(o.Resources) {
		case 1:
			b.ResourceNames("pods", o.Resources[0])
		case 2:
			b.ResourceNames(o.Resources[0], o.Resources[1])
		}

		obj, err := b.Do().Object()
		if err != nil {
			return err
		}

		o.Pod, err = o.findAttachablePod(obj)
		if err != nil {
			return err
		}

		if o.Pod.Status.Phase == apiv1.PodSucceeded || o.Pod.Status.Phase == apiv1.PodFailed {
			return fmt.Errorf("cannot attach a container in a completed pod; current phase is %s", o.Pod.Status.Phase)
		}
		// TODO: convert this to a clean "wait" behavior
	}

	// check for TTY
	containerToAttach, err := o.containerToAttachTo(o.Pod)
	if err != nil {
		return fmt.Errorf("cannot attach to the container: %v", err)
	}
	if o.TTY && !containerToAttach.TTY {
		o.TTY = false
		if o.ErrOut != nil {
			fmt.Fprintf(o.ErrOut, "Unable to use a TTY - container %s did not allocate one\n", containerToAttach.Name)
		}
	} else if !o.TTY && containerToAttach.TTY {
		// the container was launched with a TTY, so we have to force a TTY here, otherwise you'll get
		// an error "Unrecognized input header"
		o.TTY = true
	}

	// ensure we can recover the terminal while attached
	t := o.setupTTY()

	var sizeQueue remotecommand.TerminalSizeQueue
	if t.Raw {
		if size := t.GetSize(); size != nil {
			// fake resizing +1 and then back to normal so that attach-detach-reattach will result in the
			// screen being redrawn
			sizePlusOne := *size
			sizePlusOne.Width++
			sizePlusOne.Height++

			// this call spawns a goroutine to monitor/update the terminal size
			sizeQueue = t.MonitorSize(&sizePlusOne, size)
		}

		o.DisableStderr = true
	}

	if !o.Quiet {
		fmt.Fprintln(o.ErrOut, "If you don't see a command prompt, try pressing enter.")
	}
	if err := t.Safe(o.AttachFunc(o, containerToAttach, t.Raw, sizeQueue)); err != nil {
		return err
	}

	if o.Stdin && t.Raw && o.Pod.Spec.RestartPolicy == apiv1.RestartPolicyAlways {
		fmt.Fprintf(o.Out, "Session ended, resume using '%s %s -c %s -i -t' command when the pod is running\n", o.CommandName, o.Pod.Name, containerToAttach.Name)
	}
	return nil
}

func (o *AttachOptions) findAttachablePod(obj runtime.Object) (*apiv1.Pod, error) {
	attachablePod, err := o.AttachablePodFn(o.restClientGetter, obj, o.GetPodTimeout)
	if err != nil {
		return nil, err
	}

	o.StreamOptions.PodName = attachablePod.Name
	return attachablePod, nil
}

// containerToAttach returns a reference to the container to attach to, given
// by name or the first container if name is empty.
func (o *AttachOptions) containerToAttachTo(pod *apiv1.Pod) (*apiv1.Container, error) {
	if len(o.ContainerName) > 0 {
		for i := range pod.Spec.Containers {
			if pod.Spec.Containers[i].Name == o.ContainerName {
				return &pod.Spec.Containers[i], nil
			}
		}
		for i := range pod.Spec.InitContainers {
			if pod.Spec.InitContainers[i].Name == o.ContainerName {
				return &pod.Spec.InitContainers[i], nil
			}
		}
		return nil, fmt.Errorf("container not found (%s)", o.ContainerName)
	}

	if len(o.SuggestedCmdUsage) > 0 {
		fmt.Fprintf(o.ErrOut, "Defaulting container name to %s.\n", pod.Spec.Containers[0].Name)
		fmt.Fprintf(o.ErrOut, "%s\n", o.SuggestedCmdUsage)
	}

	glog.V(4).Infof("defaulting container name to %s", pod.Spec.Containers[0].Name)
	return &pod.Spec.Containers[0], nil
}

// GetContainerName returns the name of the container to attach to, with a fallback.
func (o *AttachOptions) GetContainerName(pod *apiv1.Pod) (string, error) {
	c, err := o.containerToAttachTo(pod)
	if err != nil {
		return "", err
	}
	return c.Name, nil
}
<commit_msg>UPSTREAM: 67615: attach: Move the AttachFunc default function to the initializer<commit_after>/*
Copyright 2014 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cmd

import (
	"fmt"
	"io"
	"net/url"
	"time"

	"github.com/golang/glog"
	"github.com/spf13/cobra"

	apiv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/runtime"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/tools/remotecommand"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/kubectl/cmd/templates"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/kubectl/genericclioptions"
	"k8s.io/kubernetes/pkg/kubectl/genericclioptions/resource"
	"k8s.io/kubernetes/pkg/kubectl/polymorphichelpers"
	"k8s.io/kubernetes/pkg/kubectl/scheme"
	"k8s.io/kubernetes/pkg/kubectl/util/i18n"
)

var (
	attachExample = templates.Examples(i18n.T(`
		# Get output from running pod 123456-7890, using the first container by default
		kubectl attach 123456-7890

		# Get output from ruby-container from pod 123456-7890
		kubectl attach 123456-7890 -c ruby-container

		# Switch to raw terminal mode, sends stdin to 'bash' in ruby-container from pod 123456-7890
		# and sends stdout/stderr from 'bash' back to the client
		kubectl attach 123456-7890 -c ruby-container -i -t

		# Get output from the first pod of a ReplicaSet named nginx
		kubectl attach rs/nginx
		`))
)

const (
	defaultPodAttachTimeout = 60 * time.Second
	defaultPodLogsTimeout   = 20 * time.Second
)

// AttachOptions declare the arguments accepted by the Exec command
type AttachOptions struct {
	StreamOptions

	// whether to disable the use of standard error when streaming output from tty
	DisableStderr bool

	CommandName       string
	SuggestedCmdUsage string

	Pod *apiv1.Pod

	AttachFunc       func(*AttachOptions, *apiv1.Container, bool, remotecommand.TerminalSizeQueue) func() error
	Resources        []string
	Builder          func() *resource.Builder
	AttachablePodFn  polymorphichelpers.AttachableLogsForObjectFunc
	restClientGetter genericclioptions.RESTClientGetter

	Attach        RemoteAttach
	GetPodTimeout time.Duration
	Config        *restclient.Config
}

func NewAttachOptions(streams genericclioptions.IOStreams) *AttachOptions {
	return &AttachOptions{
		StreamOptions: StreamOptions{
			IOStreams: streams,
		},
		Attach:     &DefaultRemoteAttach{},
		AttachFunc: defaultAttachFunc,
	}
}

func NewCmdAttach(f cmdutil.Factory, streams genericclioptions.IOStreams) *cobra.Command {
	o := NewAttachOptions(streams)
	cmd := &cobra.Command{
		Use: "attach (POD | TYPE/NAME) -c CONTAINER",
		DisableFlagsInUseLine: true,
		Short:   i18n.T("Attach to a running container"),
		Long:    "Attach to a process that is already running inside an existing container.",
		Example: attachExample,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, cmd, args))
			cmdutil.CheckErr(o.Validate())
			cmdutil.CheckErr(o.Run())
		},
	}
	cmdutil.AddPodRunningTimeoutFlag(cmd, defaultPodAttachTimeout)
	cmd.Flags().StringVarP(&o.ContainerName, "container", "c", o.ContainerName, "Container name. If omitted, the first container in the pod will be chosen")
	cmd.Flags().BoolVarP(&o.Stdin, "stdin", "i", o.Stdin, "Pass stdin to the container")
	cmd.Flags().BoolVarP(&o.TTY, "tty", "t", o.TTY, "Stdin is a TTY")
	return cmd
}

// RemoteAttach defines the interface accepted by the Attach command - provided for test stubbing
type RemoteAttach interface {
	Attach(method string, url *url.URL, config *restclient.Config, stdin io.Reader, stdout, stderr io.Writer, tty bool, terminalSizeQueue remotecommand.TerminalSizeQueue) error
}

func defaultAttachFunc(o *AttachOptions, containerToAttach *apiv1.Container, raw bool, sizeQueue remotecommand.TerminalSizeQueue) func() error {
	return func() error {
		restClient, err := restclient.RESTClientFor(o.Config)
		if err != nil {
			return err
		}
		req := restClient.Post().
			Resource("pods").
			Name(o.Pod.Name).
			Namespace(o.Pod.Namespace).
			SubResource("attach")
		req.VersionedParams(&apiv1.PodAttachOptions{
			Container: containerToAttach.Name,
			Stdin:     o.Stdin,
			Stdout:    o.Out != nil,
			Stderr:    !o.DisableStderr,
			TTY:       raw,
		}, legacyscheme.ParameterCodec)

		return o.Attach.Attach("POST", req.URL(), o.Config, o.In, o.Out, o.ErrOut, raw, sizeQueue)
	}
}

// DefaultRemoteAttach is the standard implementation of attaching
type DefaultRemoteAttach struct{}

func (*DefaultRemoteAttach) Attach(method string, url *url.URL, config *restclient.Config, stdin io.Reader, stdout, stderr io.Writer, tty bool, terminalSizeQueue remotecommand.TerminalSizeQueue) error {
	exec, err := remotecommand.NewSPDYExecutor(config, method, url)
	if err != nil {
		return err
	}
	return exec.Stream(remotecommand.StreamOptions{
		Stdin:             stdin,
		Stdout:            stdout,
		Stderr:            stderr,
		Tty:               tty,
		TerminalSizeQueue: terminalSizeQueue,
	})
}

// Complete verifies command line arguments and loads data from the command environment
func (o *AttachOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	var err error
	o.Namespace, _, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}

	o.AttachablePodFn = polymorphichelpers.AttachablePodForObjectFn

	o.GetPodTimeout, err = cmdutil.GetPodRunningTimeoutFlag(cmd)
	if err != nil {
		return cmdutil.UsageErrorf(cmd, err.Error())
	}

	o.Builder = f.NewBuilder
	o.Resources = args
	o.restClientGetter = f

	fullCmdName := ""
	cmdParent := cmd.Parent()
	if cmdParent != nil {
		fullCmdName = cmdParent.CommandPath()
	}
	if len(fullCmdName) > 0 && cmdutil.IsSiblingCommandExists(cmd, "describe") {
		o.SuggestedCmdUsage = fmt.Sprintf("Use '%s describe pod/%s -n %s' to see all of the containers in this pod.", fullCmdName, o.PodName, o.Namespace)
	}

	config, err := f.ToRESTConfig()
	if err != nil {
		return err
	}
	o.Config = config

	if o.CommandName == "" {
		o.CommandName = cmd.CommandPath()
	}

	return nil
}

// Validate checks that the provided attach options are specified.
func (o *AttachOptions) Validate() error {
	if len(o.Resources) == 0 {
		return fmt.Errorf("at least 1 argument is required for attach")
	}
	if len(o.Resources) > 2 {
		return fmt.Errorf("expected POD, TYPE/NAME, or TYPE NAME, (at most 2 arguments) saw %d: %v", len(o.Resources), o.Resources)
	}
	if o.GetPodTimeout <= 0 {
		return fmt.Errorf("--pod-running-timeout must be higher than zero")
	}

	return nil
}

// Run executes a validated remote execution against a pod.
func (o *AttachOptions) Run() error {
	if o.Pod == nil {
		b := o.Builder().
			WithScheme(scheme.Scheme, scheme.Scheme.PrioritizedVersionsAllGroups()...).
			NamespaceParam(o.Namespace).DefaultNamespace()

		switch len(o.Resources) {
		case 1:
			b.ResourceNames("pods", o.Resources[0])
		case 2:
			b.ResourceNames(o.Resources[0], o.Resources[1])
		}

		obj, err := b.Do().Object()
		if err != nil {
			return err
		}

		o.Pod, err = o.findAttachablePod(obj)
		if err != nil {
			return err
		}

		if o.Pod.Status.Phase == apiv1.PodSucceeded || o.Pod.Status.Phase == apiv1.PodFailed {
			return fmt.Errorf("cannot attach a container in a completed pod; current phase is %s", o.Pod.Status.Phase)
		}
		// TODO: convert this to a clean "wait" behavior
	}

	// check for TTY
	containerToAttach, err := o.containerToAttachTo(o.Pod)
	if err != nil {
		return fmt.Errorf("cannot attach to the container: %v", err)
	}
	if o.TTY && !containerToAttach.TTY {
		o.TTY = false
		if o.ErrOut != nil {
			fmt.Fprintf(o.ErrOut, "Unable to use a TTY - container %s did not allocate one\n", containerToAttach.Name)
		}
	} else if !o.TTY && containerToAttach.TTY {
		// the container was launched with a TTY, so we have to force a TTY here, otherwise you'll get
		// an error "Unrecognized input header"
		o.TTY = true
	}

	// ensure we can recover the terminal while attached
	t := o.setupTTY()

	var sizeQueue remotecommand.TerminalSizeQueue
	if t.Raw {
		if size := t.GetSize(); size != nil {
			// fake resizing +1 and then back to normal so that attach-detach-reattach will result in the
			// screen being redrawn
			sizePlusOne := *size
			sizePlusOne.Width++
			sizePlusOne.Height++

			// this call spawns a goroutine to monitor/update the terminal size
			sizeQueue = t.MonitorSize(&sizePlusOne, size)
		}

		o.DisableStderr = true
	}

	if !o.Quiet {
		fmt.Fprintln(o.ErrOut, "If you don't see a command prompt, try pressing enter.")
	}
	if err := t.Safe(o.AttachFunc(o, containerToAttach, t.Raw, sizeQueue)); err != nil {
		return err
	}

	if o.Stdin && t.Raw && o.Pod.Spec.RestartPolicy == apiv1.RestartPolicyAlways {
		fmt.Fprintf(o.Out, "Session ended, resume using '%s %s -c %s -i -t' command when the pod is running\n", o.CommandName, o.Pod.Name, containerToAttach.Name)
	}
	return nil
}

func (o *AttachOptions) findAttachablePod(obj runtime.Object) (*apiv1.Pod, error) {
	attachablePod, err := o.AttachablePodFn(o.restClientGetter, obj, o.GetPodTimeout)
	if err != nil {
		return nil, err
	}

	o.StreamOptions.PodName = attachablePod.Name
	return attachablePod, nil
}

// containerToAttach returns a reference to the container to attach to, given
// by name or the first container if name is empty.
func (o *AttachOptions) containerToAttachTo(pod *apiv1.Pod) (*apiv1.Container, error) {
	if len(o.ContainerName) > 0 {
		for i := range pod.Spec.Containers {
			if pod.Spec.Containers[i].Name == o.ContainerName {
				return &pod.Spec.Containers[i], nil
			}
		}
		for i := range pod.Spec.InitContainers {
			if pod.Spec.InitContainers[i].Name == o.ContainerName {
				return &pod.Spec.InitContainers[i], nil
			}
		}
		return nil, fmt.Errorf("container not found (%s)", o.ContainerName)
	}

	if len(o.SuggestedCmdUsage) > 0 {
		fmt.Fprintf(o.ErrOut, "Defaulting container name to %s.\n", pod.Spec.Containers[0].Name)
		fmt.Fprintf(o.ErrOut, "%s\n", o.SuggestedCmdUsage)
	}

	glog.V(4).Infof("defaulting container name to %s", pod.Spec.Containers[0].Name)
	return &pod.Spec.Containers[0], nil
}

// GetContainerName returns the name of the container to attach to, with a fallback.
func (o *AttachOptions) GetContainerName(pod *apiv1.Pod) (string, error) {
	c, err := o.containerToAttachTo(pod)
	if err != nil {
		return "", err
	}
	return c.Name, nil
}
<|endoftext|>