<commit_before>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel_test

import (
	"strings"
	"testing"
	"time"

	. "github.com/uber/tchannel-go"

	"github.com/uber/tchannel-go/raw"
	"github.com/uber/tchannel-go/testutils"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context"
)

func TestActiveCallReq(t *testing.T) {
	t.Skip("Test skipped due to unreliable way to test for protocol errors")

	ctx, cancel := NewContext(time.Second)
	defer cancel()

	// Note: This test cannot use log verification as the duplicate ID causes a log.
	// It does not use a verified server, as it leaks a message exchange due to the
	// modification of IDs in the relay.
	opts := testutils.NewOpts().DisableLogVerification()
	testutils.WithServer(t, opts, func(ch *Channel, hostPort string) {
		gotCall := make(chan struct{})
		unblock := make(chan struct{})

		testutils.RegisterFunc(ch, "blocked", func(ctx context.Context, args *raw.Args) (*raw.Res, error) {
			gotCall <- struct{}{}
			<-unblock
			return &raw.Res{}, nil
		})

		relayFunc := func(outgoing bool, frame *Frame) *Frame {
			if outgoing && frame.Header.ID == 3 {
				frame.Header.ID = 2
			}
			return frame
		}

		relayHostPort, closeRelay := testutils.FrameRelay(t, hostPort, relayFunc)
		defer closeRelay()

		firstComplete := make(chan struct{})
		go func() {
			// This call will block until we close unblock.
			raw.Call(ctx, ch, relayHostPort, ch.PeerInfo().ServiceName, "blocked", nil, nil)
			close(firstComplete)
		}()

		// Wait for the first call to be received by the server
		<-gotCall

		// Make a new call, which should fail
		_, _, _, err := raw.Call(ctx, ch, relayHostPort, ch.PeerInfo().ServiceName, "blocked", nil, nil)
		assert.Error(t, err, "Expect error")
		assert.True(t, strings.Contains(err.Error(), "already active"),
			"expected already active error, got %v", err)

		close(unblock)
		<-firstComplete
	})
}

func TestInboundConnection(t *testing.T) {
	ctx, cancel := NewContext(time.Second)
	defer cancel()

	WithVerifiedServer(t, nil, func(ch *Channel, hostPort string) {
		testutils.RegisterFunc(ch, "test", func(ctx context.Context, args *raw.Args) (*raw.Res, error) {
			c, _ := InboundConnection(CurrentCall(ctx))
			assert.Equal(t, hostPort, c.RemotePeerInfo().HostPort, "Unexpected host port")
			return &raw.Res{}, nil
		})

		_, _, _, err := raw.Call(ctx, ch, hostPort, ch.PeerInfo().ServiceName, "test", nil, nil)
		require.NoError(t, err, "Call failed")
	})
}

func TestInboundConnection_CallOptions(t *testing.T) {
	ctx, cancel := NewContext(time.Second)
	defer cancel()

	WithVerifiedServer(t, nil, func(serverCh *Channel, serverHostPort string) {
		testutils.RegisterFunc(serverCh, "test", func(ctx context.Context, args *raw.Args) (*raw.Res, error) {
			assert.Equal(t, "client", CurrentCall(ctx).CallerName(), "Expected caller name to be passed through")
			return &raw.Res{}, nil
		})

		backendName := serverCh.PeerInfo().ServiceName

		proxyCh, err := testutils.NewServerChannel(&testutils.ChannelOpts{ServiceName: "proxy"})
		require.NoError(t, err, "Proxy channel creation failed")
		defer proxyCh.Close()

		subCh := proxyCh.GetSubChannel(backendName)
		subCh.SetHandler(HandlerFunc(func(ctx context.Context, inbound *InboundCall) {
			outbound, err := proxyCh.BeginCall(ctx, serverHostPort, backendName, inbound.MethodString(), inbound.CallOptions())
			require.NoError(t, err, "Create outbound call failed")

			require.NoError(t, NewArgWriter(outbound.Arg2Writer()).Write([]byte("hi")), "Write arg2 failed")
			require.NoError(t, NewArgWriter(outbound.Arg3Writer()).Write([]byte("body")), "Write arg3 failed")
			var arg2, arg3 []byte
			require.NoError(t, NewArgReader(outbound.Response().Arg2Reader()).Read(&arg2), "Read arg2 failed")
			require.NoError(t, NewArgReader(outbound.Response().Arg3Reader()).Read(&arg3), "Read arg3 failed")
			require.NoError(t, NewArgWriter(inbound.Response().Arg2Writer()).Write(arg2), "Write arg2 failed")
			require.NoError(t, NewArgWriter(inbound.Response().Arg3Writer()).Write(arg3), "Write arg3 failed")
		}))

		clientCh, err := NewChannel("client", nil)
		require.NoError(t, err, "Create client channel failed")
		defer clientCh.Close()

		clientCh.Peers().Add(proxyCh.PeerInfo().HostPort)

		_, _, _, err = raw.Call(ctx, clientCh, proxyCh.PeerInfo().HostPort, backendName, "test", nil, nil)
		require.NoError(t, err, "Call through proxy failed")
	})
}
<commit_msg>Clean up tests<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel_test

import (
	"strings"
	"testing"
	"time"

	. "github.com/uber/tchannel-go"

	"github.com/uber/tchannel-go/raw"
	"github.com/uber/tchannel-go/testutils"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context"
)

func TestActiveCallReq(t *testing.T) {
	t.Skip("Test skipped due to unreliable way to test for protocol errors")

	ctx, cancel := NewContext(time.Second)
	defer cancel()

	// Note: This test cannot use log verification as the duplicate ID causes a log.
	// It does not use a verified server, as it leaks a message exchange due to the
	// modification of IDs in the relay.
	opts := testutils.NewOpts().DisableLogVerification()
	testutils.WithServer(t, opts, func(ch *Channel, hostPort string) {
		gotCall := make(chan struct{})
		unblock := make(chan struct{})

		testutils.RegisterFunc(ch, "blocked", func(ctx context.Context, args *raw.Args) (*raw.Res, error) {
			gotCall <- struct{}{}
			<-unblock
			return &raw.Res{}, nil
		})

		relayFunc := func(outgoing bool, frame *Frame) *Frame {
			if outgoing && frame.Header.ID == 3 {
				frame.Header.ID = 2
			}
			return frame
		}

		relayHostPort, closeRelay := testutils.FrameRelay(t, hostPort, relayFunc)
		defer closeRelay()

		firstComplete := make(chan struct{})
		go func() {
			// This call will block until we close unblock.
			raw.Call(ctx, ch, relayHostPort, ch.PeerInfo().ServiceName, "blocked", nil, nil)
			close(firstComplete)
		}()

		// Wait for the first call to be received by the server
		<-gotCall

		// Make a new call, which should fail
		_, _, _, err := raw.Call(ctx, ch, relayHostPort, ch.PeerInfo().ServiceName, "blocked", nil, nil)
		assert.Error(t, err, "Expect error")
		assert.True(t, strings.Contains(err.Error(), "already active"),
			"expected already active error, got %v", err)

		close(unblock)
		<-firstComplete
	})
}

func TestInboundConnection(t *testing.T) {
	ctx, cancel := NewContext(time.Second)
	defer cancel()

	WithVerifiedServer(t, nil, func(ch *Channel, hostPort string) {
		testutils.RegisterFunc(ch, "test", func(ctx context.Context, args *raw.Args) (*raw.Res, error) {
			c, _ := InboundConnection(CurrentCall(ctx))
			assert.Equal(t, hostPort, c.RemotePeerInfo().HostPort, "Unexpected host port")
			return &raw.Res{}, nil
		})

		_, _, _, err := raw.Call(ctx, ch, hostPort, ch.PeerInfo().ServiceName, "test", nil, nil)
		require.NoError(t, err, "Call failed")
	})
}

func TestInboundConnection_CallOptions(t *testing.T) {
	ctx, cancel := NewContext(time.Second)
	defer cancel()

	testutils.WithTestServer(t, nil, func(server *testutils.TestServer) {
		server.RegisterFunc("test", func(ctx context.Context, args *raw.Args) (*raw.Res, error) {
			assert.Equal(t, "client", CurrentCall(ctx).CallerName(), "Expected caller name to be passed through")
			return &raw.Res{}, nil
		})

		backendName := server.ServiceName()

		proxyCh := server.NewServer(&testutils.ChannelOpts{ServiceName: "proxy"})
		defer proxyCh.Close()

		subCh := proxyCh.GetSubChannel(backendName)
		subCh.SetHandler(HandlerFunc(func(ctx context.Context, inbound *InboundCall) {
			outbound, err := proxyCh.BeginCall(ctx, server.HostPort(), backendName, inbound.MethodString(), inbound.CallOptions())
			require.NoError(t, err, "Create outbound call failed")
			arg2, arg3, _, err := raw.WriteArgs(outbound, []byte("hello"), []byte("world"))
			require.NoError(t, err, "Write outbound call failed")
			require.NoError(t, raw.WriteResponse(inbound.Response(), &raw.Res{
				Arg2: arg2,
				Arg3: arg3,
			}), "Write response failed")
		}))

		clientCh := server.NewClient(&testutils.ChannelOpts{
			ServiceName: "client",
		})
		defer clientCh.Close()

		_, _, _, err := raw.Call(ctx, clientCh, proxyCh.PeerInfo().HostPort, backendName, "test", nil, nil)
		require.NoError(t, err, "Call through proxy failed")
	})
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"text/template"
	"time"
)

type ChannelType string

const (
	ChannelStable   ChannelType = "stable"
	ChannelUnstable ChannelType = "unstable"
	ChannelNightly  ChannelType = "nightly"
)

type work struct {
	src, dst string
	t        *template.Template
	info     os.FileInfo
}

type build struct {
	Package  string
	Distros  []string
	Versions []version
}

type version struct {
	Version, Revision, DownloadLinkBase string
	Channel                             ChannelType
	GetVersion                          func() (string, error)
	GetDownloadLinkBase                 func(v version) (string, error)
}

type cfg struct {
	version
	DistroName, Arch, DebArch, Package string
}

var (
	architectures = []string{"amd64", "arm", "arm64"}
	serverDistros = []string{"xenial"}
	allDistros    = []string{"xenial", "jessie", "precise", "sid", "stretch", "trusty", "utopic", "vivid", "wheezy", "wily", "yakkety"}

	builtins = map[string]interface{}{
		"date": func() string {
			return time.Now().Format(time.RFC1123Z)
		},
	}

	keepTmp = flag.Bool("keep_tmp", false, "keep tmp dir after build")
)

func runCommand(pwd string, command string, cmdArgs ...string) error {
	cmd := exec.Command(command, cmdArgs...)
	if len(pwd) != 0 {
		cmd.Dir = pwd
	}
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return err
	}
	return nil
}

func (c cfg) run() error {
	log.Printf("!!!!!!!!! doing: %#v", c)
	var w []work

	srcdir := filepath.Join(c.DistroName, c.Package)
	dstdir, err := ioutil.TempDir(os.TempDir(), "debs")
	if err != nil {
		return err
	}
	if !*keepTmp {
		defer os.RemoveAll(dstdir)
	}

	// allow base package dir to by a symlink so we can reuse packages
	// that don't change between distros
	realSrcdir, err := filepath.EvalSymlinks(srcdir)
	if err != nil {
		return err
	}

	if err := filepath.Walk(realSrcdir, func(srcfile string, f os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		dstfile := filepath.Join(dstdir, srcfile[len(realSrcdir):])
		if dstfile == dstdir {
			return nil
		}
		if f.IsDir() {
			log.Printf(dstfile)
			return os.Mkdir(dstfile, f.Mode())
		}
		t, err := template.
			New("").
			Funcs(builtins).
			Option("missingkey=error").
			ParseFiles(srcfile)
		if err != nil {
			return err
		}
		w = append(w, work{
			src:  srcfile,
			dst:  dstfile,
			t:    t.Templates()[0],
			info: f,
		})

		return nil
	}); err != nil {
		return err
	}

	for _, w := range w {
		log.Printf("w: %#v", w)
		if err := func() error {
			f, err := os.OpenFile(w.dst, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0)
			if err != nil {
				return err
			}
			defer f.Close()

			if err := w.t.Execute(f, c); err != nil {
				return err
			}
			if err := os.Chmod(w.dst, w.info.Mode()); err != nil {
				return err
			}
			return nil
		}(); err != nil {
			return err
		}
	}

	err = runCommand(dstdir, "dpkg-buildpackage", "-us", "-uc", "-b", "-a"+c.DebArch)
	if err != nil {
		return err
	}

	dstParts := []string{"bin", string(c.Channel), c.DistroName}

	dstPath := filepath.Join(dstParts...)
	os.MkdirAll(dstPath, 0777)

	fileName := fmt.Sprintf("%s_%s-%s_%s.deb", c.Package, c.Version, c.Revision, c.DebArch)
	err = runCommand("", "mv", filepath.Join("/tmp", fileName), dstPath)
	if err != nil {
		return err
	}

	return nil
}

func walkBuilds(builds []build, f func(pkg, distro, arch string, v version) error) error {
	for _, a := range architectures {
		for _, b := range builds {
			for _, d := range b.Distros {
				for _, v := range b.Versions {
					// Populate the version if it doesn't exist
					if len(v.Version) == 0 && v.GetVersion != nil {
						var err error
						v.Version, err = v.GetVersion()
						if err != nil {
							return err
						}
					}

					// Populate the version if it doesn't exist
					if len(v.DownloadLinkBase) == 0 && v.GetDownloadLinkBase != nil {
						var err error
						v.DownloadLinkBase, err = v.GetDownloadLinkBase(v)
						if err != nil {
							return err
						}
					}

					if err := f(b.Package, d, a, v); err != nil {
						return err
					}
				}
			}
		}
	}
	return nil
}

func fetchVersion(url string) (string, error) {
	res, err := http.Get(url)
	if err != nil {
		return "", err
	}

	versionBytes, err := ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		return "", err
	}
	// Remove a newline and the v prefix from the string
	return strings.Replace(strings.Replace(string(versionBytes), "v", "", 1), "\n", "", 1), nil
}

func getStableKubeVersion() (string, error) {
	return fetchVersion("https://dl.k8s.io/release/stable.txt")
}

func getLatestKubeVersion() (string, error) {
	return fetchVersion("https://dl.k8s.io/release/latest.txt")
}

func getLatestCIVersion() (string, error) {
	latestVersion, err := getLatestKubeCIBuild()
	if err != nil {
		return "", err
	}

	// Replace the "+" with a "-" to make it semver-compliant
	return strings.Replace(latestVersion, "+", "-", 1), nil
}

func getLatestKubeCIBuild() (string, error) {
	return fetchVersion("https://dl.k8s.io/ci-cross/latest.txt")
}

func getCIBuildsDownloadLinkBase(_ version) (string, error) {
	latestCiVersion, err := getLatestKubeCIBuild()
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("https://dl.k8s.io/ci-cross/v%s", latestCiVersion), nil
}

func getReleaseDownloadLinkBase(v version) (string, error) {
	return fmt.Sprintf("https://dl.k8s.io/v%s", v.Version), nil
}

func main() {
	flag.Parse()

	builds := []build{
		{
			Package: "kubectl",
			Distros: allDistros,
			Versions: []version{
				{
					GetVersion:          getStableKubeVersion,
					Revision:            "00",
					Channel:             ChannelStable,
					GetDownloadLinkBase: getReleaseDownloadLinkBase,
				},
				{
					GetVersion:          getLatestKubeVersion,
					Revision:            "00",
					Channel:             ChannelUnstable,
					GetDownloadLinkBase: getReleaseDownloadLinkBase,
				},
				{
					GetVersion:          getLatestCIVersion,
					Revision:            "00",
					Channel:             ChannelNightly,
					GetDownloadLinkBase: getCIBuildsDownloadLinkBase,
				},
			},
		},
		{
			Package: "kubelet",
			Distros: serverDistros,
			Versions: []version{
				{
					GetVersion:          getStableKubeVersion,
					Revision:            "00",
					Channel:             ChannelStable,
					GetDownloadLinkBase: getReleaseDownloadLinkBase,
				},
				{
					GetVersion:          getLatestKubeVersion,
					Revision:            "00",
					Channel:             ChannelUnstable,
					GetDownloadLinkBase: getReleaseDownloadLinkBase,
				},
				{
					GetVersion:          getLatestCIVersion,
					Revision:            "00",
					Channel:             ChannelNightly,
					GetDownloadLinkBase: getCIBuildsDownloadLinkBase,
				},
			},
		},
		{
			Package: "kubernetes-cni",
			Distros: serverDistros,
			Versions: []version{
				{
					Version:  "0.3.0.1-07a8a2",
					Revision: "00",
					Channel:  ChannelStable,
				},
				{
					Version:  "0.3.0.1-07a8a2",
					Revision: "00",
					Channel:  ChannelUnstable,
				},
				{
					Version:  "0.3.0.1-07a8a2",
					Revision: "00",
					Channel:  ChannelNightly,
				},
			},
		},
		{
			Package: "kubeadm",
			Distros: serverDistros,
			Versions: []version{
				{
					// Remember to update xenial/kubeadm/debian/rules with the same version
					Version:             "1.6.0-alpha.0.2074-a092d8e0f95f52",
					Revision:            "00",
					Channel:             ChannelStable,
					GetDownloadLinkBase: getCIBuildsDownloadLinkBase,
				},
				{
					GetVersion:          getLatestCIVersion,
					Revision:            "00",
					Channel:             ChannelUnstable,
					GetDownloadLinkBase: getCIBuildsDownloadLinkBase,
				},
				{
					GetVersion:          getLatestCIVersion,
					Revision:            "00",
					Channel:             ChannelNightly,
					GetDownloadLinkBase: getCIBuildsDownloadLinkBase,
				},
			},
		},
	}

	if err := walkBuilds(builds, func(pkg, distro, arch string, v version) error {
		c := cfg{
			Package:    pkg,
			version:    v,
			DistroName: distro,
			Arch:       arch,
		}
		if c.Arch == "arm" {
			c.DebArch = "armhf"
		} else {
			c.DebArch = c.Arch
		}
		return c.run()
	}); err != nil {
		log.Fatalf("err: %v", err)
	}
}
<commit_msg>Add ppc64le and s390x builds and fix a bug<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"text/template"
	"time"
)

type ChannelType string

const (
	ChannelStable   ChannelType = "stable"
	ChannelUnstable ChannelType = "unstable"
	ChannelNightly  ChannelType = "nightly"
)

type work struct {
	src, dst string
	t        *template.Template
	info     os.FileInfo
}

type build struct {
	Package  string
	Distros  []string
	Versions []version
}

type version struct {
	Version, Revision, DownloadLinkBase string
	Channel                             ChannelType
	GetVersion                          func() (string, error)
	GetDownloadLinkBase                 func(v version) (string, error)
}

type cfg struct {
	version
	DistroName, Arch, DebArch, Package string
}

var (
	architectures = []string{"amd64", "arm", "arm64", "ppc64le", "s390x"}
	serverDistros = []string{"xenial"}
	allDistros    = []string{"xenial", "jessie", "precise", "sid", "stretch", "trusty", "utopic", "vivid", "wheezy", "wily", "yakkety"}

	builtins = map[string]interface{}{
		"date": func() string {
			return time.Now().Format(time.RFC1123Z)
		},
	}

	keepTmp = flag.Bool("keep_tmp", false, "keep tmp dir after build")
)

func runCommand(pwd string, command string, cmdArgs ...string) error {
	cmd := exec.Command(command, cmdArgs...)
	if len(pwd) != 0 {
		cmd.Dir = pwd
	}
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return err
	}
	return nil
}

func (c cfg) run() error {
	log.Printf("!!!!!!!!! doing: %#v", c)
	var w []work

	srcdir := filepath.Join(c.DistroName, c.Package)
	dstdir, err := ioutil.TempDir(os.TempDir(), "debs")
	if err != nil {
		return err
	}
	if !*keepTmp {
		defer os.RemoveAll(dstdir)
	}

	// allow base package dir to by a symlink so we can reuse packages
	// that don't change between distros
	realSrcdir, err := filepath.EvalSymlinks(srcdir)
	if err != nil {
		return err
	}

	if err := filepath.Walk(realSrcdir, func(srcfile string, f os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		dstfile := filepath.Join(dstdir, srcfile[len(realSrcdir):])
		if dstfile == dstdir {
			return nil
		}
		if f.IsDir() {
			log.Printf(dstfile)
			return os.Mkdir(dstfile, f.Mode())
		}
		t, err := template.
			New("").
			Funcs(builtins).
			Option("missingkey=error").
			ParseFiles(srcfile)
		if err != nil {
			return err
		}
		w = append(w, work{
			src:  srcfile,
			dst:  dstfile,
			t:    t.Templates()[0],
			info: f,
		})

		return nil
	}); err != nil {
		return err
	}

	for _, w := range w {
		log.Printf("w: %#v", w)
		if err := func() error {
			f, err := os.OpenFile(w.dst, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0)
			if err != nil {
				return err
			}
			defer f.Close()

			if err := w.t.Execute(f, c); err != nil {
				return err
			}
			if err := os.Chmod(w.dst, w.info.Mode()); err != nil {
				return err
			}
			return nil
		}(); err != nil {
			return err
		}
	}

	err = runCommand(dstdir, "dpkg-buildpackage", "-us", "-uc", "-b", "-a"+c.DebArch)
	if err != nil {
		return err
	}

	dstParts := []string{"bin", string(c.Channel), c.DistroName}

	dstPath := filepath.Join(dstParts...)
	os.MkdirAll(dstPath, 0777)

	fileName := fmt.Sprintf("%s_%s-%s_%s.deb", c.Package, c.Version, c.Revision, c.DebArch)
	err = runCommand("", "mv", filepath.Join("/tmp", fileName), dstPath)
	if err != nil {
		return err
	}

	return nil
}

func walkBuilds(builds []build, f func(pkg, distro, arch string, v version) error) error {
	for _, a := range architectures {
		for _, b := range builds {
			for _, d := range b.Distros {
				for _, v := range b.Versions {
					// Populate the version if it doesn't exist
					if len(v.Version) == 0 && v.GetVersion != nil {
						var err error
						v.Version, err = v.GetVersion()
						if err != nil {
							return err
						}
					}

					// Populate the version if it doesn't exist
					if len(v.DownloadLinkBase) == 0 && v.GetDownloadLinkBase != nil {
						var err error
						v.DownloadLinkBase, err = v.GetDownloadLinkBase(v)
						if err != nil {
							return err
						}
					}

					if err := f(b.Package, d, a, v); err != nil {
						return err
					}
				}
			}
		}
	}
	return nil
}

func fetchVersion(url string) (string, error) {
	res, err := http.Get(url)
	if err != nil {
		return "", err
	}

	versionBytes, err := ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		return "", err
	}
	// Remove a newline and the v prefix from the string
	return strings.Replace(strings.Replace(string(versionBytes), "v", "", 1), "\n", "", 1), nil
}

func getStableKubeVersion() (string, error) {
	return fetchVersion("https://dl.k8s.io/release/stable.txt")
}

func getLatestKubeVersion() (string, error) {
	return fetchVersion("https://dl.k8s.io/release/latest.txt")
}

func getLatestCIVersion() (string, error) {
	latestVersion, err := getLatestKubeCIBuild()
	if err != nil {
		return "", err
	}

	// Replace the "+" with a "-" to make it semver-compliant
	return strings.Replace(latestVersion, "+", "-", 1), nil
}

func getLatestKubeCIBuild() (string, error) {
	return fetchVersion("https://dl.k8s.io/ci-cross/latest.txt")
}

func getCIBuildsDownloadLinkBase(_ version) (string, error) {
	latestCiVersion, err := getLatestKubeCIBuild()
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("https://dl.k8s.io/ci-cross/v%s", latestCiVersion), nil
}

func getCIBuildsDownloadLinkBaseFromVersion(v version) (string, error) {
	return fmt.Sprintf("https://dl.k8s.io/ci-cross/v%s", v.Version), nil
}

func getReleaseDownloadLinkBase(v version) (string, error) {
	return fmt.Sprintf("https://dl.k8s.io/v%s", v.Version), nil
}

func main() {
	flag.Parse()

	builds := []build{
		{
			Package: "kubectl",
			Distros: allDistros,
			Versions: []version{
				{
					GetVersion:          getStableKubeVersion,
					Revision:            "00",
					Channel:             ChannelStable,
					GetDownloadLinkBase: getReleaseDownloadLinkBase,
				},
				{
					GetVersion:          getLatestKubeVersion,
					Revision:            "00",
					Channel:             ChannelUnstable,
					GetDownloadLinkBase: getReleaseDownloadLinkBase,
				},
				{
					GetVersion:          getLatestCIVersion,
					Revision:            "00",
					Channel:             ChannelNightly,
					GetDownloadLinkBase: getCIBuildsDownloadLinkBase,
				},
			},
		},
		{
			Package: "kubelet",
			Distros: serverDistros,
			Versions: []version{
				{
					GetVersion:          getStableKubeVersion,
					Revision:            "00",
					Channel:             ChannelStable,
					GetDownloadLinkBase: getReleaseDownloadLinkBase,
				},
				{
					GetVersion:          getLatestKubeVersion,
					Revision:            "00",
					Channel:             ChannelUnstable,
					GetDownloadLinkBase: getReleaseDownloadLinkBase,
				},
				{
					GetVersion:          getLatestCIVersion,
					Revision:            "00",
					Channel:             ChannelNightly,
					GetDownloadLinkBase: getCIBuildsDownloadLinkBase,
				},
			},
		},
		{
			Package: "kubernetes-cni",
			Distros: serverDistros,
			Versions: []version{
				{
					Version:  "0.3.0.1-07a8a2",
					Revision: "00",
					Channel:  ChannelStable,
				},
				{
					Version:  "0.3.0.1-07a8a2",
					Revision: "00",
					Channel:  ChannelUnstable,
				},
				{
					Version:  "0.3.0.1-07a8a2",
					Revision: "00",
					Channel:  ChannelNightly,
				},
			},
		},
		{
			Package: "kubeadm",
			Distros: serverDistros,
			Versions: []version{
				{
					Version:             "1.6.0-alpha.0.2074-a092d8e0f95f52",
					Revision:            "00",
					Channel:             ChannelStable,
					GetDownloadLinkBase: getCIBuildsDownloadLinkBaseFromVersion,
				},
				{
					GetVersion:          getLatestCIVersion,
					Revision:            "00",
					Channel:             ChannelUnstable,
					GetDownloadLinkBase: getCIBuildsDownloadLinkBase,
				},
				{
					GetVersion:          getLatestCIVersion,
					Revision:            "00",
					Channel:             ChannelNightly,
					GetDownloadLinkBase: getCIBuildsDownloadLinkBase,
				},
			},
		},
	}

	if err := walkBuilds(builds, func(pkg, distro, arch string, v version) error {
		c := cfg{
			Package:    pkg,
			version:    v,
			DistroName: distro,
			Arch:       arch,
		}
		if c.Arch == "arm" {
			c.DebArch = "armhf"
		} else if c.Arch == "ppc64le" {
			c.DebArch = "ppc64el"
		} else {
			c.DebArch = c.Arch
		}
		// Skip platforms that do not have binaries for a channel
		if len(v.Channel) != 0 {
			if v.Channel == "stable" && (c.Arch == "s390x" || c.Arch == "ppc64le") {
				return nil
			}
		}
		return c.run()
	}); err != nil {
		log.Fatalf("err: %v", err)
	}
}
<|endoftext|>


<commit_before>package devereux

import (
	"errors"
	"io/ioutil"
	"os"
)

func promptForKeyIfNecessary(key string, prompt string, envVariable string) (string, error) {
	if key != "" {
		return key, nil
	}

	key = os.Getenv(envVariable)
	if key != "" {
		return key, nil
	}

	var err error
	for key == "" {
		key, err = promptUserForInput(prompt)
		if err != nil {
			return "", err
		}
	}

	return key, nil
}

func GetPassword(passwordName string, repoName string, key string) (string, error) {
	key, err := promptForKeyIfNecessary(key, "Enter repository key> ", "DVRX_KEY")
	if err != nil {
		return "", err
	}

	os.MkdirAll(REPO_PATH, 0777)
	config, err := loadConfig(CONFIG_PATH)
	if err != nil {
		return "", err
	}

	if repoName == "" {
		repoName = config.CurrentRepository
	}

	if repoName == "" {
		return "", errors.New("No default repository found. Specify one with the -r flag.")
	}

	repo := &Repository{Name: repoName, config: config}
	return repo.GetPassword(passwordName, key)
}

func SetPassword(passwordName string, repoName string, key string, password string) error {
	key, err := promptForKeyIfNecessary(key, "Enter repository key> ", "DVRX_KEY")
	if err != nil {
		return err
	}

	password, err = promptForKeyIfNecessary(password, "Enter password> ", "DVRX_PASS")
	if err != nil {
		return err
	}

	os.MkdirAll(REPO_PATH, 0777)
	config, err := loadConfig(CONFIG_PATH)
	if err != nil {
		return err
	}

	if repoName == "" {
		repoName = config.CurrentRepository
	}

	if repoName == "" {
		return "", errors.New("No default repository found. Specify one with the -r flag.")
	}

	repo := &Repository{Name: repoName, config: config}
	return repo.SetPassword(passwordName, password, key)
}

func CreateRepository(name string, setAsDefault bool, key string) (string, error) {
	key, err := promptForKeyIfNecessary(key, "Enter repository key> ", "DVRX_KEY")
	if err != nil {
		return "", err
	}

	os.MkdirAll(REPO_PATH, 0777)
	config, err := loadConfig(CONFIG_PATH)
	if err != nil {
		return "", err
	}

	repo := &Repository{Name: name, config: config}
	err = repo.Create(key, setAsDefault, CONFIG_PATH)
	if err != nil {
		return "", err
	}

	return repo.Name, nil
}

func GetRepositories() ([]string, error) {
	names := make([]string, 0)

	files, err := ioutil.ReadDir(REPO_PATH)
	if err != nil {
		return names, err
	}

	for _, f := range files {
		names = append(names, f.Name())
	}
	return names, nil
}

func DeleteRepository(name string) error {
	repo := &Repository{Name: name}
	err := repo.Delete()
	if err != nil {
		return err
	}

	config, err := loadConfig(CONFIG_PATH)
	if err != nil {
		return err
	}

	if config.CurrentRepository == name {
		config.CurrentRepository = ""
		config.Save(CONFIG_PATH)
	}

	return nil
}
<commit_msg>Fixed minor bug.<commit_after>package devereux

import (
	"errors"
	"io/ioutil"
	"os"
)

func promptForKeyIfNecessary(key string, prompt string, envVariable string) (string, error) {
	if key != "" {
		return key, nil
	}

	key = os.Getenv(envVariable)
	if key != "" {
		return key, nil
	}

	var err error
	for key == "" {
		key, err = promptUserForInput(prompt)
		if err != nil {
			return "", err
		}
	}

	return key, nil
}

func GetPassword(passwordName string, repoName string, key string) (string, error) {
	key, err := promptForKeyIfNecessary(key, "Enter repository key> ", "DVRX_KEY")
	if err != nil {
		return "", err
	}

	os.MkdirAll(REPO_PATH, 0777)
	config, err := loadConfig(CONFIG_PATH)
	if err != nil {
		return "", err
	}

	if repoName == "" {
		repoName = config.CurrentRepository
	}

	if repoName == "" {
		return "", errors.New("No default repository found. Specify one with the -r flag.")
	}

	repo := &Repository{Name: repoName, config: config}
	return repo.GetPassword(passwordName, key)
}

func SetPassword(passwordName string, repoName string, key string, password string) error {
	key, err := promptForKeyIfNecessary(key, "Enter repository key> ", "DVRX_KEY")
	if err != nil {
		return err
	}

	password, err = promptForKeyIfNecessary(password, "Enter password> ", "DVRX_PASS")
	if err != nil {
		return err
	}

	os.MkdirAll(REPO_PATH, 0777)
	config, err := loadConfig(CONFIG_PATH)
	if err != nil {
		return err
	}

	if repoName == "" {
		repoName = config.CurrentRepository
	}

	if repoName == "" {
		return errors.New("No default repository found. Specify one with the -r flag.")
	}

	repo := &Repository{Name: repoName, config: config}
	return repo.SetPassword(passwordName, password, key)
}

func CreateRepository(name string, setAsDefault bool, key string) (string, error) {
	key, err := promptForKeyIfNecessary(key, "Enter repository key> ", "DVRX_KEY")
	if err != nil {
		return "", err
	}

	os.MkdirAll(REPO_PATH, 0777)
	config, err := loadConfig(CONFIG_PATH)
	if err != nil {
		return "", err
	}

	repo := &Repository{Name: name, config: config}
	err = repo.Create(key, setAsDefault, CONFIG_PATH)
	if err != nil {
		return "", err
	}

	return repo.Name, nil
}

func GetRepositories() ([]string, error) {
	names := make([]string, 0)

	files, err := ioutil.ReadDir(REPO_PATH)
	if err != nil {
		return names, err
	}

	for _, f := range files {
		names = append(names, f.Name())
	}
	return names, nil
}

func DeleteRepository(name string) error {
	repo := &Repository{Name: name}
	err := repo.Delete()
	if err != nil {
		return err
	}

	config, err := loadConfig(CONFIG_PATH)
	if err != nil {
		return err
	}

	if config.CurrentRepository == name {
		config.CurrentRepository = ""
		config.Save(CONFIG_PATH)
	}

	return nil
}
<|endoftext|>


<commit_before>/*
 * Copyright 2019 The Kythe Authors. All rights reserved.
 *
 * 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 createcmd

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"sort"
	"strings"

	"kythe.io/kythe/go/util/kytheuri"
	"kythe.io/kythe/go/util/vnameutil"

	"bitbucket.org/creachadair/stringset"
	"github.com/golang/protobuf/jsonpb"

	anypb "github.com/golang/protobuf/ptypes/any"
	apb "kythe.io/kythe/proto/analysis_go_proto"
)

type repeatedString []string

// Set implements part of the flag.Getter interface and will append a new value to the flag.
func (f *repeatedString) Set(s string) error {
	*f = append(*f, s)
	return nil
}

// String implements part of the flag.Getter interface and returns a string-ish value for the flag.
func (f *repeatedString) String() string {
	if f == nil {
		return ""
	}
	return strings.Join(*f, ",")
}

// Get implements flag.Getter and returns a slice of string values.
func (f *repeatedString) Get() interface{} {
	if f == nil {
		return []string(nil)
	}
	return *f
}

type repeatedStringSet stringset.Set

// Set implements part of the flag.Getter interface and will append a new value to the flag.
func (f *repeatedStringSet) Set(s string) error {
	(*stringset.Set)(f).Add(s)
	return nil
}

// update adds the values from other to the contained stringset.
func (f *repeatedStringSet) update(o repeatedStringSet) bool {
	return (*stringset.Set)(f).Update(stringset.Set(o))
}

// Elements returns the set of elements as a sorted slice.
func (f *repeatedStringSet) Elements() []string {
	return (*stringset.Set)(f).Elements()
}

// Len returns the number of elements.
func (f *repeatedStringSet) Len() int {
	return (*stringset.Set)(f).Len()
}

// String implements part of the flag.Getter interface and returns a string-ish value for the flag.
func (f *repeatedStringSet) String() string {
	if f == nil {
		return ""
	}
	return strings.Join(f.Elements(), ",")
}

// Get implements flag.Getter and returns a slice of string values.
func (f *repeatedStringSet) Get() interface{} {
	if f == nil {
		return stringset.Set(nil)
	}
	return *f
}

type repeatedEnv map[string]string

// Set implements part of the flag.Getter interface and will append a new value to the flag.
func (f *repeatedEnv) Set(s string) error {
	parts := strings.SplitN(s, "=", 2)
	if len(parts) != 2 {
		return fmt.Errorf("invalid environment entry: %v", s)
	}
	if *f == nil {
		*f = make(map[string]string)
	}
	(*f)[parts[0]] = parts[1]
	return nil
}

// String implements part of the flag.Getter interface and returns a string-ish value for the flag.
func (f *repeatedEnv) String() string {
	if f == nil || *f == nil {
		return ""
	}
	var values []string
	for key, value := range *f {
		values = append(values, key+"="+value)
	}
	return strings.Join(values, "\n")
}

// Get implements flag.Getter and returns a slice of string values.
func (f *repeatedEnv) Get() interface{} {
	if f == nil {
		return map[string]string(nil)
	}
	return *f
}

// ToProto returns a []*apb.CompilationUnit_Env for the mapped environment.
func (f *repeatedEnv) ToProto() []*apb.CompilationUnit_Env {
	if f == nil || *f == nil {
		return nil
	}
	var result []*apb.CompilationUnit_Env
	for key, value := range *f {
		result = append(result, &apb.CompilationUnit_Env{
			Name:  key,
			Value: value,
		})
	}
	sort.Slice(result, func(i, j int) bool { return result[i].Name < result[j].Name })
	return result
}

type kytheURI kytheuri.URI

// Set implements part of the flag.Value interface and will append a new value to the flag.
func (f *kytheURI) Set(s string) error {
	uri, err := kytheuri.Parse(s)
	switch {
	case err != nil:
		return err
	case uri.Corpus == "":
		return errors.New("missing required URI field: corpus")
	case uri.Language == "":
		return errors.New("missing required URI field: language")

	}
	*f = *(*kytheURI)(uri)
	return err
}

// String implements part of the flag.Value interface and returns a string-ish value for the flag.
func (f *kytheURI) String() string {
	return (*kytheuri.URI)(f).String()
}

// repeatedAny is a repeated flag of JSON-encoded Any messages.
type repeatedAny []*anypb.Any

// Set implements part of the flag.Getter interface and will append a new value to the flag.
func (f *repeatedAny) Set(s string) error {
	dec := json.NewDecoder(strings.NewReader(s))
	for dec.More() {
		var detail anypb.Any
		if err := jsonpb.UnmarshalNext(dec, &detail); err != nil {
			return err
		}
		*f = append(*f, &detail)
	}
	return nil
}

var toJSON = &jsonpb.Marshaler{OrigName: true}

// String implements part of the flag.Getter interface and returns a string-ish value for the flag.
func (f *repeatedAny) String() string {
	if f == nil {
		return ""
	}
	var result []string
	for _, val := range *f {
		str, err := toJSON.MarshalToString(val)
		if err != nil {
			panic(err)
		}
		result = append(result, str)
	}
	return strings.Join(result, " ")
}

// Get implements flag.Getter and returns a slice of string values.
func (f *repeatedAny) Get() interface{} {
	if f == nil {
		return []*anypb.Any(nil)
	}
	return *f
}

// vnameRules is a path-valued flag used for loading VName mapping rules from a vnames.json file.
type vnameRules struct {
	filename string
	vnameutil.Rules
}

// Set implements part of the flag.Value interface and will append a new value to the flag.
func (f *vnameRules) Set(s string) error {
	f.filename = s
	data, err := ioutil.ReadFile(f.filename)
	if err != nil {
		return fmt.Errorf("reading vname rules: %v", err)
	}
	f.Rules, err = vnameutil.ParseRules(data)
	if err != nil {
		return fmt.Errorf("reading vname rules: %v", err)
	}
	return nil
}

// String implements part of the flag.Value interface and returns a string-ish value for the flag.
func (f *vnameRules) String() string {
	return f.filename
}
<commit_msg>fix(kzip_create): implement flag.Getter for all flat types (#3544)<commit_after>/*
 * Copyright 2019 The Kythe Authors. All rights reserved.
 *
 * 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 createcmd

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"sort"
	"strings"

	"kythe.io/kythe/go/util/kytheuri"
	"kythe.io/kythe/go/util/vnameutil"

	"bitbucket.org/creachadair/stringset"
	"github.com/golang/protobuf/jsonpb"

	anypb "github.com/golang/protobuf/ptypes/any"
	apb "kythe.io/kythe/proto/analysis_go_proto"
)

type repeatedString []string

// Set implements part of the flag.Getter interface and will append a new value to the flag.
func (f *repeatedString) Set(s string) error {
	*f = append(*f, s)
	return nil
}

// String implements part of the flag.Getter interface and returns a string-ish value for the flag.
func (f *repeatedString) String() string {
	if f == nil {
		return ""
	}
	return strings.Join(*f, ",")
}

// Get implements flag.Getter and returns a slice of string values.
func (f *repeatedString) Get() interface{} {
	if f == nil {
		return []string(nil)
	}
	return *f
}

type repeatedStringSet stringset.Set

// Set implements part of the flag.Getter interface and will append a new value to the flag.
func (f *repeatedStringSet) Set(s string) error {
	(*stringset.Set)(f).Add(s)
	return nil
}

// update adds the values from other to the contained stringset.
func (f *repeatedStringSet) update(o repeatedStringSet) bool {
	return (*stringset.Set)(f).Update(stringset.Set(o))
}

// Elements returns the set of elements as a sorted slice.
func (f *repeatedStringSet) Elements() []string {
	return (*stringset.Set)(f).Elements()
}

// Len returns the number of elements.
func (f *repeatedStringSet) Len() int {
	return (*stringset.Set)(f).Len()
}

// String implements part of the flag.Getter interface and returns a string-ish value for the flag.
func (f *repeatedStringSet) String() string {
	if f == nil {
		return ""
	}
	return strings.Join(f.Elements(), ",")
}

// Get implements flag.Getter and returns a slice of string values.
func (f *repeatedStringSet) Get() interface{} {
	if f == nil {
		return stringset.Set(nil)
	}
	return *f
}

type repeatedEnv map[string]string

// Set implements part of the flag.Getter interface and will append a new value to the flag.
func (f *repeatedEnv) Set(s string) error {
	parts := strings.SplitN(s, "=", 2)
	if len(parts) != 2 {
		return fmt.Errorf("invalid environment entry: %v", s)
	}
	if *f == nil {
		*f = make(map[string]string)
	}
	(*f)[parts[0]] = parts[1]
	return nil
}

// String implements part of the flag.Getter interface and returns a string-ish value for the flag.
func (f *repeatedEnv) String() string {
	if f == nil || *f == nil {
		return ""
	}
	var values []string
	for key, value := range *f {
		values = append(values, key+"="+value)
	}
	return strings.Join(values, "\n")
}

// Get implements flag.Getter and returns a slice of string values.
func (f *repeatedEnv) Get() interface{} {
	if f == nil {
		return map[string]string(nil)
	}
	return *f
}

// ToProto returns a []*apb.CompilationUnit_Env for the mapped environment.
func (f *repeatedEnv) ToProto() []*apb.CompilationUnit_Env {
	if f == nil || *f == nil {
		return nil
	}
	var result []*apb.CompilationUnit_Env
	for key, value := range *f {
		result = append(result, &apb.CompilationUnit_Env{
			Name:  key,
			Value: value,
		})
	}
	sort.Slice(result, func(i, j int) bool { return result[i].Name < result[j].Name })
	return result
}

type kytheURI kytheuri.URI

// Set implements part of the flag.Value interface and will append a new value to the flag.
func (f *kytheURI) Set(s string) error {
	uri, err := kytheuri.Parse(s)
	switch {
	case err != nil:
		return err
	case uri.Corpus == "":
		return errors.New("missing required URI field: corpus")
	case uri.Language == "":
		return errors.New("missing required URI field: language")

	}
	*f = *(*kytheURI)(uri)
	return err
}

// String implements part of the flag.Value interface and returns a string-ish value for the flag.
func (f *kytheURI) String() string {
	return (*kytheuri.URI)(f).String()
}

// Get implements part of the flag.Getter interface.
func (f *kytheURI) Get() interface{} {
	return f
}

// repeatedAny is a repeated flag of JSON-encoded Any messages.
type repeatedAny []*anypb.Any

// Set implements part of the flag.Getter interface and will append a new value to the flag.
func (f *repeatedAny) Set(s string) error {
	dec := json.NewDecoder(strings.NewReader(s))
	for dec.More() {
		var detail anypb.Any
		if err := jsonpb.UnmarshalNext(dec, &detail); err != nil {
			return err
		}
		*f = append(*f, &detail)
	}
	return nil
}

var toJSON = &jsonpb.Marshaler{OrigName: true}

// String implements part of the flag.Getter interface and returns a string-ish value for the flag.
func (f *repeatedAny) String() string {
	if f == nil {
		return ""
	}
	var result []string
	for _, val := range *f {
		str, err := toJSON.MarshalToString(val)
		if err != nil {
			panic(err)
		}
		result = append(result, str)
	}
	return strings.Join(result, " ")
}

// Get implements flag.Getter and returns a slice of string values.
func (f *repeatedAny) Get() interface{} {
	if f == nil {
		return []*anypb.Any(nil)
	}
	return *f
}

// vnameRules is a path-valued flag used for loading VName mapping rules from a vnames.json file.
type vnameRules struct {
	filename string
	vnameutil.Rules
}

// Set implements part of the flag.Value interface and will append a new value to the flag.
func (f *vnameRules) Set(s string) error {
	f.filename = s
	data, err := ioutil.ReadFile(f.filename)
	if err != nil {
		return fmt.Errorf("reading vname rules: %v", err)
	}
	f.Rules, err = vnameutil.ParseRules(data)
	if err != nil {
		return fmt.Errorf("reading vname rules: %v", err)
	}
	return nil
}

// String implements part of the flag.Value interface and returns a string-ish value for the flag.
func (f *vnameRules) String() string {
	return f.filename
}

// Get implements part of the flag.Getter interface.
func (f *vnameRules) Get() interface{} {
	return f
}
<|endoftext|>


<commit_before>package gsmake

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/url"
	"path/filepath"
	"strings"
	"time"

	"github.com/gsdocker/gserrors"
	"github.com/gsdocker/gslogger"

	"github.com/gsdocker/gsos/fs"
	"github.com/gsmake/gsmake/vfs"
)

// gsmake predeined environment variables
const (
	EnvHome          = "GSMAKE_HOME"
	VersionGSMake    = "release/v2.0"
	PacakgeAnonymous = "github.com/gsmake/gsmake.anonymous"
	Logfmt           = "[$tag] $content"
	LogTimefmt       = ""
)

// Errors .
var (
	ErrLoad = errors.New("load package error")
)

// Loader package loader
type Loader struct {
	gslogger.Log                                // mixin Log APIs
	packages     map[string]map[string]*Package // loaded packages
	checkerOfDCG []*Package                     // DCG check stack
	rootfs       vfs.RootFS                     // vfs object
	targetpath   string                         // the loading package path
}

func load(rootfs vfs.RootFS) (*Loader, error) {

	loader := &Loader{
		Log:        gslogger.Get("loader"),
		targetpath: rootfs.TargetPath(),
		packages:   make(map[string]map[string]*Package),
		rootfs:     rootfs,
	}

	loader.I("load package ...")

	start := time.Now()

	err := loader.load()

	if err != nil {
		return nil, err
	}

	loader.I("load package -- success %s", time.Now().Sub(start))

	return loader, nil
}

func (loader *Loader) addpackage(domain string, pkg *Package) {
	packages, ok := loader.packages[domain]

	if !ok {
		packages = make(map[string]*Package)

		loader.packages[domain] = packages
	}

	pkg.loadPath = make([]*Package, len(loader.checkerOfDCG))

	copy(pkg.loadPath, loader.checkerOfDCG)

	packages[pkg.Name] = pkg
}

func (loader *Loader) querypackage(domain string, name string) (*Package, bool) {
	if packages, ok := loader.packages[domain]; ok {

		if pkg, ok := packages[name]; ok {
			return pkg, true
		}

	}

	return nil, false
}

func (loader *Loader) load() error {

	jsonfile := filepath.Join(loader.targetpath, ".gsmake.json")

	if !fs.Exists(jsonfile) {

		return gserrors.Newf(
			ErrLoad,
			"target package not exists or is not a gsmake package\n\t%s",
			loader.targetpath,
		)
	}

	pkg, err := loadjson(jsonfile)

	if err != nil {
		return err
	}

	domains := ParseDomain(pkg.Domain, DomainDefault)

	for _, domain := range domains {

		_, _, err := loader.tryMount(domain, pkg.Name, loader.targetpath, "current", "file")

		if err != nil {
			return err
		}

		pkg, err = loader.loadpackagev2(domain, pkg.Name, loader.targetpath)

		if err != nil {
			return err
		}

		loader.addpackage(domain, pkg)

	}

	loader.D("loaded domain : [%s]", strings.Join(domains, ","))

	// try load gsmake
	if _, ok := loader.querypackage("task", "github.com/gsmake/gsmake"); !ok {

		pkg, err := loader.loadpackage(Import{
			Name:    "github.com/gsmake/gsmake",
			Version: VersionGSMake,
			SCM:     "git",
			Domain:  "task",
		})

		if err != nil {
			return gserrors.Newf(err, "load package github.com/gsmake/gsmake error")
		}

		loader.addpackage("task", pkg)
	}

	// dismount not loaded packages

	err = loader.rootfs.List(func(src, target *vfs.Entry) bool {

		if _, ok := loader.querypackage(target.Domain(), target.Name()); !ok {
			loader.I("dismount :%s", target)
			loader.rootfs.Dismount(target.String())
		}

		return true
	})

	if err != nil {
		return gserrors.Newf(err, "dismount unreference packages error")
	}

	return nil
}

func (loader *Loader) checkDCG(name string) error {

	var stream bytes.Buffer

	for _, pkg := range loader.checkerOfDCG {
		if pkg.Name == name || stream.Len() != 0 {
			stream.WriteString(fmt.Sprintf("\t%s import\n", pkg.Name))
		}
	}

	if stream.Len() != 0 {
		return gserrors.Newf(ErrLoad, "circular package import :\n%s\t%s", stream.String(), name)
	}

	return nil
}

func (loader *Loader) tryMount(domain, name, srcpath, version, scm string) (string, string, error) {
	target := fmt.Sprintf("gsmake://%s?domain=%s", name, domain)

	src := fmt.Sprintf("%s://%s?version=%s", scm, srcpath, version)

	if !loader.rootfs.Mounted(src, target) {

		if err := loader.rootfs.Mount(src, target); err != nil {
			return src, target, err
		}
	}

	return src, target, nil
}

func loadpath(path []*Package, name, version string) string {
	var buff bytes.Buffer

	for _, pkg := range path {
		buff.WriteString(fmt.Sprintf("\t\t%s %s\n", pkg.Name, pkg.version))
	}

	buff.WriteString(fmt.Sprintf("\t\t%s %s\n", name, version))

	return buff.String()
}

func (loader *Loader) loadpackage(i Import) (*Package, error) {

	if packages, ok := loader.packages[i.Domain]; ok {

		if pkg, ok := packages[i.Name]; ok {

			if pkg.version != i.Version {
				return nil, gserrors.Newf(
					ErrLoad,
					"import package with diff version\n\tthe one:\n%s\n\tthe other:\n%s",
					loadpath(pkg.loadPath, pkg.Name, pkg.version),
					loadpath(loader.checkerOfDCG, i.Name, i.Version),
				)
			}

			return pkg, nil
		}
	}

	// DCG check
	if err := loader.checkDCG(i.Name); err != nil {
		return nil, err
	}

	_, target, err := loader.tryMount(i.Domain, i.Name, i.Name, i.Version, i.SCM)

	if err != nil {
		return nil, err
	}

	_, entry, err := loader.rootfs.Open(target)

	if err != nil {
		return nil, err
	}

	importpkg, err := loader.loadpackagev2(i.Domain, i.Name, entry.Mapping)

	if err != nil {
		return nil, err
	}

	importpkg.version = i.Version

	return importpkg, nil
}

func (loader *Loader) loadpackagev2(currentDomain, name, fullpath string) (*Package, error) {

	jsonfile := filepath.Join(fullpath, ".gsmake.json")

	if !fs.Exists(jsonfile) {
		// this package is a traditional golang package
		return &Package{
			Name:   name,
			Domain: currentDomain,
		}, nil
	}

	pkg, err := loadjson(jsonfile)

	if err != nil {
		return nil, err
	}

	loader.checkerOfDCG = append(loader.checkerOfDCG, pkg)

	defer func() {
		loader.checkerOfDCG = loader.checkerOfDCG[:len(loader.checkerOfDCG)-1]
	}()

	for _, ir := range pkg.Import {

		if err := loader.importPackage(currentDomain, pkg, ir); err != nil {
			return nil, err
		}
	}

	for _, task := range pkg.Task {
		task.Package = name
	}

	return pkg, nil
}

func (loader *Loader) importPackage(currentDomain string, parent *Package, ir Import) error {

	if ir.Version == "" {
		ir.Version = "current"
	}

	// calc scm url
	if ir.SCM == "" {

		u, err := url.Parse(fmt.Sprintf("https://%s", ir.Name))

		if err != nil {
			return gserrors.Newf(err, "%s invalid import package :%s", parent.Name, ir.Name)
		}

		ir.SCM = loader.rootfs.Protocol(u.Host)
	}

	domains := ParseDomain(ir.Domain, parent.Domain)

	parentDomains := ParseDomain(parent.Domain, DomainDefault)

	for _, domain := range domains {

		valid := false

		// check domain if valid .
		for _, parentDomain := range parentDomains {
			if domain == parentDomain {
				valid = true
				break
			}
		}

		if !valid {
			return gserrors.Newf(
				ErrLoad,
				"%s %s invalid import %s %s\n\t unknown domain",
				parent.Name,
				parent.version,
				ir.Name,
				ir.Version,
			)
		}

		if domain == currentDomain {

			ir.Domain = domain

			pkg, err := loader.loadpackage(ir)

			if err != nil {
				return err
			}

			loader.addpackage(domain, pkg)

			return nil
		}
	}

	return nil
}

func loadjson(file string) (*Package, error) {

	content, err := ioutil.ReadFile(file)

	if err != nil {
		return nil, gserrors.Newf(err, "load config file err\n\t%s", file)
	}

	config := &Package{
		Domain: DomainDefault,
	}

	err = json.Unmarshal(content, &config)

	if err != nil {
		return nil, gserrors.Newf(err, "unmarshal .gsmake.json file error\n\tfile:%s", file)
	}

	return config, nil
}
<commit_msg>add load domain check<commit_after>package gsmake

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/url"
	"path/filepath"
	"strings"
	"time"

	"github.com/gsdocker/gserrors"
	"github.com/gsdocker/gslogger"

	"github.com/gsdocker/gsos/fs"
	"github.com/gsmake/gsmake/vfs"
)

// gsmake predeined environment variables
const (
	EnvHome          = "GSMAKE_HOME"
	VersionGSMake    = "release/v2.0"
	PacakgeAnonymous = "github.com/gsmake/gsmake.anonymous"
	Logfmt           = "[$tag] $content"
	LogTimefmt       = ""
)

// Errors .
var (
	ErrLoad = errors.New("load package error")
)

// Loader package loader
type Loader struct {
	gslogger.Log                                // mixin Log APIs
	packages     map[string]map[string]*Package // loaded packages
	checkerOfDCG []*Package                     // DCG check stack
	rootfs       vfs.RootFS                     // vfs object
	targetpath   string                         // the loading package path
}

func load(rootfs vfs.RootFS) (*Loader, error) {

	loader := &Loader{
		Log:        gslogger.Get("loader"),
		targetpath: rootfs.TargetPath(),
		packages:   make(map[string]map[string]*Package),
		rootfs:     rootfs,
	}

	loader.I("load package ...")

	start := time.Now()

	err := loader.load()

	if err != nil {
		return nil, err
	}

	loader.I("load package -- success %s", time.Now().Sub(start))

	return loader, nil
}

func (loader *Loader) addpackage(domain string, pkg *Package) {
	packages, ok := loader.packages[domain]

	if !ok {
		packages = make(map[string]*Package)

		loader.packages[domain] = packages
	}

	pkg.loadPath = make([]*Package, len(loader.checkerOfDCG))

	copy(pkg.loadPath, loader.checkerOfDCG)

	packages[pkg.Name] = pkg
}

func (loader *Loader) querypackage(domain string, name string) (*Package, bool) {
	if packages, ok := loader.packages[domain]; ok {

		if pkg, ok := packages[name]; ok {
			return pkg, true
		}

	}

	return nil, false
}

func (loader *Loader) load() error {

	jsonfile := filepath.Join(loader.targetpath, ".gsmake.json")

	if !fs.Exists(jsonfile) {

		return gserrors.Newf(
			ErrLoad,
			"target package not exists or is not a gsmake package\n\t%s",
			loader.targetpath,
		)
	}

	pkg, err := loadjson(jsonfile)

	if err != nil {
		return err
	}

	domains := ParseDomain(pkg.Domain, DomainDefault)

	for _, domain := range domains {

		_, _, err := loader.tryMount(domain, pkg.Name, loader.targetpath, "current", "file")

		if err != nil {
			return err
		}

		pkg, err = loader.loadpackagev2(domain, pkg.Name, loader.targetpath)

		if err != nil {
			return err
		}

		loader.addpackage(domain, pkg)

	}

	loader.D("loaded domain : [%s]", strings.Join(domains, ","))

	// try load gsmake
	if _, ok := loader.querypackage("task", "github.com/gsmake/gsmake"); !ok {

		pkg, err := loader.loadpackage(Import{
			Name:    "github.com/gsmake/gsmake",
			Version: VersionGSMake,
			SCM:     "git",
			Domain:  "task",
		})

		if err != nil {
			return gserrors.Newf(err, "load package github.com/gsmake/gsmake error")
		}

		loader.addpackage("task", pkg)
	}

	// dismount not loaded packages

	err = loader.rootfs.List(func(src, target *vfs.Entry) bool {

		if _, ok := loader.querypackage(target.Domain(), target.Name()); !ok {
			loader.I("dismount :%s", target)
			loader.rootfs.Dismount(target.String())
		}

		return true
	})

	if err != nil {
		return gserrors.Newf(err, "dismount unreference packages error")
	}

	return nil
}

func (loader *Loader) checkDCG(name string) error {

	var stream bytes.Buffer

	for _, pkg := range loader.checkerOfDCG {
		if pkg.Name == name || stream.Len() != 0 {
			stream.WriteString(fmt.Sprintf("\t%s import\n", pkg.Name))
		}
	}

	if stream.Len() != 0 {
		return gserrors.Newf(ErrLoad, "circular package import :\n%s\t%s", stream.String(), name)
	}

	return nil
}

func (loader *Loader) tryMount(domain, name, srcpath, version, scm string) (string, string, error) {
	target := fmt.Sprintf("gsmake://%s?domain=%s", name, domain)

	src := fmt.Sprintf("%s://%s?version=%s", scm, srcpath, version)

	if !loader.rootfs.Mounted(src, target) {

		if err := loader.rootfs.Mount(src, target); err != nil {
			return src, target, err
		}
	}

	return src, target, nil
}

func loadpath(path []*Package, name, version string) string {
	var buff bytes.Buffer

	for _, pkg := range path {
		buff.WriteString(fmt.Sprintf("\t\t%s %s\n", pkg.Name, pkg.version))
	}

	buff.WriteString(fmt.Sprintf("\t\t%s %s\n", name, version))

	return buff.String()
}

func (loader *Loader) loadpackage(i Import) (*Package, error) {

	if packages, ok := loader.packages[i.Domain]; ok {

		if pkg, ok := packages[i.Name]; ok {

			if pkg.version != i.Version {
				return nil, gserrors.Newf(
					ErrLoad,
					"import package with diff version\n\tthe one:\n%s\n\tthe other:\n%s",
					loadpath(pkg.loadPath, pkg.Name, pkg.version),
					loadpath(loader.checkerOfDCG, i.Name, i.Version),
				)
			}

			return pkg, nil
		}
	}

	// DCG check
	if err := loader.checkDCG(i.Name); err != nil {
		return nil, err
	}

	_, target, err := loader.tryMount(i.Domain, i.Name, i.Name, i.Version, i.SCM)

	if err != nil {
		return nil, err
	}

	_, entry, err := loader.rootfs.Open(target)

	if err != nil {
		return nil, err
	}

	importpkg, err := loader.loadpackagev2(i.Domain, i.Name, entry.Mapping)

	if err != nil {
		return nil, err
	}

	importpkg.version = i.Version

	return importpkg, nil
}

func (loader *Loader) loadpackagev2(currentDomain, name, fullpath string) (*Package, error) {

	jsonfile := filepath.Join(fullpath, ".gsmake.json")

	if !fs.Exists(jsonfile) {
		// this package is a traditional golang package
		return &Package{
			Name:   name,
			Domain: currentDomain,
		}, nil
	}

	pkg, err := loadjson(jsonfile)

	if err != nil {
		return nil, err
	}

	valid := false

	for _, domain := range ParseDomain(pkg.Domain, DomainDefault) {
		if domain == currentDomain {
			valid = true
		}
	}

	if !valid {
		return nil, gserrors.Newf(err, "%s unsupport domain :%s", pkg.Name, currentDomain)
	}

	loader.checkerOfDCG = append(loader.checkerOfDCG, pkg)

	defer func() {
		loader.checkerOfDCG = loader.checkerOfDCG[:len(loader.checkerOfDCG)-1]
	}()

	for _, ir := range pkg.Import {

		if err := loader.importPackage(currentDomain, pkg, ir); err != nil {
			return nil, err
		}
	}

	for _, task := range pkg.Task {
		task.Package = name
	}

	return pkg, nil
}

func (loader *Loader) importPackage(currentDomain string, parent *Package, ir Import) error {

	if ir.Version == "" {
		ir.Version = "current"
	}

	// calc scm url
	if ir.SCM == "" {

		u, err := url.Parse(fmt.Sprintf("https://%s", ir.Name))

		if err != nil {
			return gserrors.Newf(err, "%s invalid import package :%s", parent.Name, ir.Name)
		}

		ir.SCM = loader.rootfs.Protocol(u.Host)
	}

	domains := ParseDomain(ir.Domain, parent.Domain)

	parentDomains := ParseDomain(parent.Domain, DomainDefault)

	for _, domain := range domains {

		valid := false

		// check domain if valid .
		for _, parentDomain := range parentDomains {
			if domain == parentDomain {
				valid = true
				break
			}
		}

		if !valid {
			return gserrors.Newf(
				ErrLoad,
				"%s %s invalid import %s %s\n\t unknown domain",
				parent.Name,
				parent.version,
				ir.Name,
				ir.Version,
			)
		}

		if domain == currentDomain {

			ir.Domain = domain

			pkg, err := loader.loadpackage(ir)

			if err != nil {
				return err
			}

			loader.addpackage(domain, pkg)

			return nil
		}
	}

	return nil
}

func loadjson(file string) (*Package, error) {

	content, err := ioutil.ReadFile(file)

	if err != nil {
		return nil, gserrors.Newf(err, "load config file err\n\t%s", file)
	}

	config := &Package{
		Domain: DomainDefault,
	}

	err = json.Unmarshal(content, &config)

	if err != nil {
		return nil, gserrors.Newf(err, "unmarshal .gsmake.json file error\n\tfile:%s", file)
	}

	return config, nil
}
<|endoftext|>


<commit_before>package pg

import (
	"errors"
	"fmt"
	"reflect"
	"strconv"
)

var (
	Discard Loader = discardLoader{}

	_ Loader = &structLoader{}
	_ Loader = &valuesLoader{}
	_ Loader = &Strings{}
	_ Loader = &Ints{}
)

type Loader interface {
	Load(colIdx int, colName string, b []byte) error
}

//------------------------------------------------------------------------------

type discardLoader struct{}

func (discardLoader) Load(colIdx int, colName string, b []byte) error {
	return nil
}

//------------------------------------------------------------------------------

type structLoader struct {
	v      reflect.Value
	fields map[string][]int
}

func newStructLoader(v reflect.Value) *structLoader {
	return &structLoader{
		v:      v,
		fields: structs.Fields(v.Type()),
	}
}

func (l *structLoader) Load(colIdx int, colName string, b []byte) error {
	indx, ok := l.fields[colName]
	if !ok {
		return fmt.Errorf("pg: cannot map field %q", colName)
	}
	return decodeValue(l.v.FieldByIndex(indx), b)
}

//------------------------------------------------------------------------------

type valuesLoader struct {
	values []interface{}
}

func LoadInto(values ...interface{}) Loader {
	return &valuesLoader{values}
}

func (l *valuesLoader) Load(colIdx int, colName string, b []byte) error {
	return Decode(l.values[colIdx], b)
}

//------------------------------------------------------------------------------

type Strings []string

func (strings *Strings) New() interface{} {
	return strings
}

func (strings *Strings) Load(colIdx int, colName string, b []byte) error {
	*strings = append(*strings, string(b))
	return nil
}

func (strings Strings) Append(dst []byte) []byte {
	if len(strings) <= 0 {
		return dst
	}

	for _, s := range strings {
		dst = appendString(dst, s)
		dst = append(dst, ',')
	}
	dst = dst[:len(dst)-1]
	return dst
}

//------------------------------------------------------------------------------

type Ints []int64

func (ints *Ints) New() interface{} {
	return ints
}

func (ints *Ints) Load(colIdx int, colName string, b []byte) error {
	n, err := strconv.ParseInt(string(b), 10, 64)
	if err != nil {
		return err
	}
	*ints = append(*ints, n)
	return nil
}

func (ints Ints) Append(dst []byte) []byte {
	if len(ints) <= 0 {
		return dst
	}

	for _, v := range ints {
		dst = strconv.AppendInt(dst, v, 10)
		dst = append(dst, ',')
	}
	dst = dst[:len(dst)-1]
	return dst
}

//------------------------------------------------------------------------------

func newLoader(dst interface{}) (Loader, error) {
	v := reflect.ValueOf(dst)
	if !v.IsValid() {
		return nil, errors.New("pg: Decode(" + v.String() + ")")
	}
	if v.Kind() != reflect.Ptr {
		return nil, errors.New("pg: pointer expected")
	}
	v = v.Elem()
	switch v.Kind() {
	case reflect.Struct:
		return newStructLoader(v), nil
	}
	return nil, fmt.Errorf("pg: unsupported type %v", v.Type().String())
}
<commit_msg>Add IntsSet.<commit_after>package pg

import (
	"errors"
	"fmt"
	"reflect"
	"strconv"
)

var (
	Discard Loader = discardLoader{}

	_ Loader = &structLoader{}
	_ Loader = &valuesLoader{}
	_ Loader = &Strings{}
	_ Loader = &Ints{}
)

type Loader interface {
	Load(colIdx int, colName string, b []byte) error
}

//------------------------------------------------------------------------------

type discardLoader struct{}

func (discardLoader) Load(colIdx int, colName string, b []byte) error {
	return nil
}

//------------------------------------------------------------------------------

type structLoader struct {
	v      reflect.Value
	fields map[string][]int
}

func newStructLoader(v reflect.Value) *structLoader {
	return &structLoader{
		v:      v,
		fields: structs.Fields(v.Type()),
	}
}

func (l *structLoader) Load(colIdx int, colName string, b []byte) error {
	indx, ok := l.fields[colName]
	if !ok {
		return fmt.Errorf("pg: cannot map field %q", colName)
	}
	return decodeValue(l.v.FieldByIndex(indx), b)
}

//------------------------------------------------------------------------------

type valuesLoader struct {
	values []interface{}
}

func LoadInto(values ...interface{}) Loader {
	return &valuesLoader{values}
}

func (l *valuesLoader) Load(colIdx int, colName string, b []byte) error {
	return Decode(l.values[colIdx], b)
}

//------------------------------------------------------------------------------

type Strings []string

func (strings *Strings) New() interface{} {
	return strings
}

func (strings *Strings) Load(colIdx int, colName string, b []byte) error {
	*strings = append(*strings, string(b))
	return nil
}

func (strings Strings) Append(dst []byte) []byte {
	if len(strings) <= 0 {
		return dst
	}

	for _, s := range strings {
		dst = appendString(dst, s)
		dst = append(dst, ',')
	}
	dst = dst[:len(dst)-1]
	return dst
}

//------------------------------------------------------------------------------

type Ints []int64

func (ints *Ints) New() interface{} {
	return ints
}

func (ints *Ints) Load(colIdx int, colName string, b []byte) error {
	n, err := strconv.ParseInt(string(b), 10, 64)
	if err != nil {
		return err
	}
	*ints = append(*ints, n)
	return nil
}

func (ints Ints) Append(dst []byte) []byte {
	if len(ints) <= 0 {
		return dst
	}

	for _, v := range ints {
		dst = strconv.AppendInt(dst, v, 10)
		dst = append(dst, ',')
	}
	dst = dst[:len(dst)-1]
	return dst
}

//------------------------------------------------------------------------------

type IntsSet map[int64]struct{}

func (set IntsSet) New() interface{} {
	return set
}

func (set IntsSet) Load(colIdx int, colName string, b []byte) error {
	n, err := strconv.ParseInt(string(b), 10, 64)
	if err != nil {
		return err
	}
	set[n] = struct{}{}
	return nil
}

//------------------------------------------------------------------------------

func newLoader(dst interface{}) (Loader, error) {
	v := reflect.ValueOf(dst)
	if !v.IsValid() {
		return nil, errors.New("pg: Decode(" + v.String() + ")")
	}
	if v.Kind() != reflect.Ptr {
		return nil, errors.New("pg: pointer expected")
	}
	v = v.Elem()
	switch v.Kind() {
	case reflect.Struct:
		return newStructLoader(v), nil
	}
	return nil, fmt.Errorf("pg: unsupported type %v", v.Type().String())
}
<|endoftext|>


<commit_before>package wally

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"testing"

	rdb "github.com/dancannon/gorethink"
	"github.com/stretchr/testify/assert"
)

var (
	session *rdb.Session
)

func init() {
	Database = "testing"
	DocumentTable = "documents"
	IndexTable = "indexes"

	var err error
	session, err = rdb.Connect(rdb.ConnectOpts{
		Address:  os.Getenv("RETHINKDB_URL"),
		Database: "test",
	})

	if err != nil {
		fmt.Errorf(err.Error())
	}
}

func DbBootstrap() {
	// Reset database
	rdb.DbDrop(Database).Exec(session)
	rdb.DbCreate(Database).Exec(session)

	// Drop tables
	rdb.Db(Database).TableDrop(DocumentTable).Run(session)
	rdb.Db(Database).TableDrop(IndexTable).Run(session)

	// Create tables
	rdb.Db(Database).TableCreate(DocumentTable).Run(session)
	rdb.Db(Database).TableCreate(IndexTable).Run(session)
}

func TestIndexer_Stopper(t *testing.T) {
	tests := []struct {
		Input  string
		Output string
	}{
		{
			"computer",
			"computer",
		},

		{
			"the",
			"",
		},

		{
			"technology",
			"technology",
		},

		{
			"wasn't",
			"",
		},
	}

	for _, test := range tests {
		assert.Equal(t, test.Output, Stopper(test.Input))
	}
}

func TestIndexer_SplitTextIntoWords(t *testing.T) {
	tests := []struct {
		Input  interface{}
		Output []string
	}{
		{
			"I am a block of text",
			[]string{"I", "am", "a", "block", "of", "text"},
		},

		{
			"        superfluous    whitespace ",
			[]string{"superfluous", "whitespace"},
		},

		{
			"              ",
			[]string{},
		},

		{
			[]byte("fancy a byte?"),
			[]string{"fancy", "a", "byte?"},
		},

		{
			32,
			[]string{},
		},
	}

	for _, test := range tests {
		assert.Equal(t, test.Output, SplitTextIntoWords(test.Input))
	}
}

func TestIndexer_Indexer(t *testing.T) {
	docId := "12345-67890-ABCDE"

	tests := []struct {
		Input  interface{}
		Output []Index
	}{
		{
			"I am a block of text and I am going to be parsed",
			[]Index{
				Index{Id: "block", DocumentId: docId, Count: 0},
				Index{Id: "text", DocumentId: docId, Count: 0},
				Index{Id: "going", DocumentId: docId, Count: 0},
				Index{Id: "parsed", DocumentId: docId, Count: 0},
			},
		},

		{
			[]byte("I am another block of text but now I am in bytes"),
			[]Index{
				Index{Id: "another", DocumentId: docId, Count: 0},
				Index{Id: "block", DocumentId: docId, Count: 0},
				Index{Id: "text", DocumentId: docId, Count: 0},
				Index{Id: "now", DocumentId: docId, Count: 0},
				Index{Id: "bytes", DocumentId: docId, Count: 0},
			},
		},
	}

	for _, test := range tests {
		assert.Equal(t, test.Output, Indexer(test.Input, docId))
	}
}

func TestIndexer_IndexString(t *testing.T) {
	indexId := "world"

	index := Index{Id: indexId}

	assert.Equal(t, "Index#world", index.String())
}

func TestIndexer_IndexPut(t *testing.T) {
	DbBootstrap()

	index := Index{Id: "hello", Count: 5, DocumentId: "12345-67890-ABCDE"}

	err := index.Put(session)
	assert.Nil(t, err)

	res, err := rdb.Db(Database).Table(IndexTable).Get("hello").Run(session)
	assert.Nil(t, err)

	var i Index
	err = res.One(&i)
	assert.Nil(t, err)

	assert.Equal(t, i.Id, "hello")
	assert.Equal(t, i.Count, 5)
	assert.Equal(t, i.DocumentId, "12345-67890-ABCDE")
}

func TestIndexer_DocumentString(t *testing.T) {
	docId := "12345-67890-ABCDE"

	doc := Document{Id: docId}

	assert.Equal(t, "Document#12345-67890-ABCDE", doc.String())
}

func TestIndexer_DocumentPut(t *testing.T) {
	DbBootstrap()

	doc := Document{
		Source:  "www.google.com",
		Content: "Lorem ipsum dolor sit amet.",
	}

	err := doc.Put(session)
	assert.Nil(t, err)

	res, err := rdb.Db(Database).Table(DocumentTable).Run(session)
	assert.Nil(t, err)

	var d Document
	err = res.One(&d)
	assert.Nil(t, err)

	assert.NotEqual(t, d.Id, "")
	assert.Equal(t, d.Source, "www.google.com")
	assert.Equal(t, d.Content, "Lorem ipsum dolor sit amet.")
}

func BenchmarkSplitTextIntoWords(b *testing.B) {
	file, err := ioutil.ReadFile("test_data/test.txt") // 30654 words
	if err != nil {
		b.Error("Could not load test data")
	}

	for n := 0; n < b.N; n++ {
		SplitTextIntoWords(file)
	}
}

func BenchmarkSplitTextIntoWords_two(b *testing.B) {
	file, err := ioutil.ReadFile("test_data/test_2.txt") // 1891 words
	if err != nil {
		b.Error("Could not load test data")
	}

	for n := 0; n < b.N; n++ {
		SplitTextIntoWords(file)
	}
}

func BenchmarkStopper(b *testing.B) {
	file, err := ioutil.ReadFile("test_data/test.txt") // 30654 words
	if err != nil {
		b.Error("Could not load test data")
	}

	data := strings.Fields(string(file))

	for n := 0; n < b.N; n++ {
		for _, word := range data {
			Stopper(word)
		}
	}
}

func BenchmarkStopper_two(b *testing.B) {
	file, err := ioutil.ReadFile("test_data/test_2.txt") // 1891 words
	if err != nil {
		b.Error("Could not load test data")
	}

	data := strings.Fields(string(file))

	for n := 0; n < b.N; n++ {
		for _, word := range data {
			Stopper(word)
		}
	}
}

func BenchmarkIndexer(b *testing.B) {
	docId := "12345-67890-ABCDE"
	file, err := ioutil.ReadFile("test_data/test.txt") // 30654 words
	if err != nil {
		b.Error("Could not load test data")
	}

	for n := 0; n < b.N; n++ {
		Indexer(file, docId)
	}
}

func BenchmarkIndexer_two(b *testing.B) {
	docId := "12345-67890-ABCDE"
	file, err := ioutil.ReadFile("test_data/test_2.txt") // 30654 words
	if err != nil {
		b.Error("Could not load test data")
	}

	for n := 0; n < b.N; n++ {
		Indexer(file, docId)
	}
}
<commit_msg>fix #6 Don't recreate database each time, recreate once before any tests are run<commit_after>package wally

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"testing"

	rdb "github.com/dancannon/gorethink"
	"github.com/stretchr/testify/assert"
)

var (
	session *rdb.Session
)

func init() {
	Database = "testing"
	DocumentTable = "documents"
	IndexTable = "indexes"

	var err error
	session, err = rdb.Connect(rdb.ConnectOpts{
		Address:  os.Getenv("RETHINKDB_URL"),
		Database: "test",
	})

	if err != nil {
		fmt.Errorf(err.Error())
	}

	// Reset database
	rdb.DbDrop(Database).Exec(session)
	rdb.DbCreate(Database).Exec(session)
}

func DbBootstrap() {
	// Drop tables
	rdb.Db(Database).TableDrop(DocumentTable).Run(session)
	rdb.Db(Database).TableDrop(IndexTable).Run(session)

	// Create tables
	rdb.Db(Database).TableCreate(DocumentTable).Run(session)
	rdb.Db(Database).TableCreate(IndexTable).Run(session)
}

func TestIndexer_Stopper(t *testing.T) {
	tests := []struct {
		Input  string
		Output string
	}{
		{
			"computer",
			"computer",
		},

		{
			"the",
			"",
		},

		{
			"technology",
			"technology",
		},

		{
			"wasn't",
			"",
		},
	}

	for _, test := range tests {
		assert.Equal(t, test.Output, Stopper(test.Input))
	}
}

func TestIndexer_SplitTextIntoWords(t *testing.T) {
	tests := []struct {
		Input  interface{}
		Output []string
	}{
		{
			"I am a block of text",
			[]string{"I", "am", "a", "block", "of", "text"},
		},

		{
			"        superfluous    whitespace ",
			[]string{"superfluous", "whitespace"},
		},

		{
			"              ",
			[]string{},
		},

		{
			[]byte("fancy a byte?"),
			[]string{"fancy", "a", "byte?"},
		},

		{
			32,
			[]string{},
		},
	}

	for _, test := range tests {
		assert.Equal(t, test.Output, SplitTextIntoWords(test.Input))
	}
}

func TestIndexer_Indexer(t *testing.T) {
	docId := "12345-67890-ABCDE"

	tests := []struct {
		Input  interface{}
		Output []Index
	}{
		{
			"I am a block of text and I am going to be parsed",
			[]Index{
				Index{Id: "block", DocumentId: docId, Count: 0},
				Index{Id: "text", DocumentId: docId, Count: 0},
				Index{Id: "going", DocumentId: docId, Count: 0},
				Index{Id: "parsed", DocumentId: docId, Count: 0},
			},
		},

		{
			[]byte("I am another block of text but now I am in bytes"),
			[]Index{
				Index{Id: "another", DocumentId: docId, Count: 0},
				Index{Id: "block", DocumentId: docId, Count: 0},
				Index{Id: "text", DocumentId: docId, Count: 0},
				Index{Id: "now", DocumentId: docId, Count: 0},
				Index{Id: "bytes", DocumentId: docId, Count: 0},
			},
		},
	}

	for _, test := range tests {
		assert.Equal(t, test.Output, Indexer(test.Input, docId))
	}
}

func TestIndexer_IndexString(t *testing.T) {
	indexId := "world"

	index := Index{Id: indexId}

	assert.Equal(t, "Index#world", index.String())
}

func TestIndexer_IndexPut(t *testing.T) {
	DbBootstrap()

	index := Index{Id: "hello", Count: 5, DocumentId: "12345-67890-ABCDE"}

	err := index.Put(session)
	assert.Nil(t, err)

	res, err := rdb.Db(Database).Table(IndexTable).Get("hello").Run(session)
	assert.Nil(t, err)

	var i Index
	err = res.One(&i)
	assert.Nil(t, err)

	assert.Equal(t, i.Id, "hello")
	assert.Equal(t, i.Count, 5)
	assert.Equal(t, i.DocumentId, "12345-67890-ABCDE")
}

func TestIndexer_DocumentString(t *testing.T) {
	docId := "12345-67890-ABCDE"

	doc := Document{Id: docId}

	assert.Equal(t, "Document#12345-67890-ABCDE", doc.String())
}

func TestIndexer_DocumentPut(t *testing.T) {
	DbBootstrap()

	doc := Document{
		Source:  "www.google.com",
		Content: "Lorem ipsum dolor sit amet.",
	}

	err := doc.Put(session)
	assert.Nil(t, err)

	res, err := rdb.Db(Database).Table(DocumentTable).Run(session)
	assert.Nil(t, err)

	var d Document
	err = res.One(&d)
	assert.Nil(t, err)

	assert.NotEqual(t, d.Id, "")
	assert.Equal(t, d.Source, "www.google.com")
	assert.Equal(t, d.Content, "Lorem ipsum dolor sit amet.")
}

func BenchmarkSplitTextIntoWords(b *testing.B) {
	file, err := ioutil.ReadFile("test_data/test.txt") // 30654 words
	if err != nil {
		b.Error("Could not load test data")
	}

	for n := 0; n < b.N; n++ {
		SplitTextIntoWords(file)
	}
}

func BenchmarkSplitTextIntoWords_two(b *testing.B) {
	file, err := ioutil.ReadFile("test_data/test_2.txt") // 1891 words
	if err != nil {
		b.Error("Could not load test data")
	}

	for n := 0; n < b.N; n++ {
		SplitTextIntoWords(file)
	}
}

func BenchmarkStopper(b *testing.B) {
	file, err := ioutil.ReadFile("test_data/test.txt") // 30654 words
	if err != nil {
		b.Error("Could not load test data")
	}

	data := strings.Fields(string(file))

	for n := 0; n < b.N; n++ {
		for _, word := range data {
			Stopper(word)
		}
	}
}

func BenchmarkStopper_two(b *testing.B) {
	file, err := ioutil.ReadFile("test_data/test_2.txt") // 1891 words
	if err != nil {
		b.Error("Could not load test data")
	}

	data := strings.Fields(string(file))

	for n := 0; n < b.N; n++ {
		for _, word := range data {
			Stopper(word)
		}
	}
}

func BenchmarkIndexer(b *testing.B) {
	docId := "12345-67890-ABCDE"
	file, err := ioutil.ReadFile("test_data/test.txt") // 30654 words
	if err != nil {
		b.Error("Could not load test data")
	}

	for n := 0; n < b.N; n++ {
		Indexer(file, docId)
	}
}

func BenchmarkIndexer_two(b *testing.B) {
	docId := "12345-67890-ABCDE"
	file, err := ioutil.ReadFile("test_data/test_2.txt") // 30654 words
	if err != nil {
		b.Error("Could not load test data")
	}

	for n := 0; n < b.N; n++ {
		Indexer(file, docId)
	}
}
<|endoftext|>


<commit_before>package aws

import (
	mach "github.com/jeffjen/machine/lib/machine"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/urfave/cli"

	"fmt"
	"math/rand"
	"net"
	"os"
	"strings"
	"sync"
)

var (
	DEVICE_NAME = []string{"b", "c", "d", "e", "f", "g", "h", "i"}
)

type ec2state struct {
	*ec2.Instance
	name string
	err  error
}

func newEc2State(inst *ec2.Instance, err error) ec2state {
	return ec2state{inst, "", err}
}

func getEc2InstanceName(inst *ec2.Instance) (name string) {
	for _, t := range inst.Tags {
		if *t.Key == "Name" {
			return *t.Value
		}
	}
	return
}

func ec2Init() error {
	var resp = new(ec2.DescribeInstancesOutput)

	for more := true; more; {
		params := &ec2.DescribeInstancesInput{}
		if resp.NextToken != nil {
			params.NextToken = resp.NextToken
		}
		resp, err := svc.DescribeInstances(params)
		if err != nil {
			more = false
		} else if len(resp.Reservations) < 1 {
			more = false
		} else {
			for _, r := range resp.Reservations {
				for _, inst := range r.Instances {
					var instanceName = getEc2InstanceName(inst)
					info, ok := instList[instanceName]
					if !ok {
						info = &mach.Instance{Id: *inst.InstanceId}
					}
					if *inst.State.Name == "terminated" {
						if info.Id == *inst.InstanceId {
							delete(instList, instanceName)
						} else {
							continue // ignore
						}
					} else {
						info.Driver = "aws"
						info.State = *inst.State.Name
						info.Id = *inst.InstanceId
						func() {
							var addr *net.TCPAddr
							if inst.PublicIpAddress != nil {
								addr, _ = net.ResolveTCPAddr("tcp", *inst.PublicIpAddress+":2376")
							}
							info.DockerHost = addr
						}()
						instList[instanceName] = info
					}
				}
			}
			more = (resp.NextToken != nil)
		}
	}

	return nil
}

func ec2_getSubnet(profile *VPCProfile, public bool) (subnetId *string) {
	var collection []*string
	for _, subnet := range profile.Subnet {
		if public && *subnet.Public {
			collection = append(collection, subnet.Id)
		} else if !public && !*subnet.Public {
			collection = append(collection, subnet.Id)
		}
	}
	idx := rand.Intn(len(collection))
	return collection[idx]
}

func ec2_findSecurityGroup(profile *VPCProfile, name ...string) (sgId []*string) {
	sgId = make([]*string, 0)
	for _, grp := range name {
		for _, sgrp := range profile.SecurityGroup {
			if *sgrp.Name == grp {
				sgId = append(sgId, sgrp.Id)
			}
		}
	}
	return
}

func ec2_tagInstanceParam(tags []string) (*ec2.CreateTagsInput, error) {
	tagparam := &ec2.CreateTagsInput{
		Tags:      make([]*ec2.Tag, 0),
		Resources: make([]*string, 0),
	}
	for _, tag := range tags {
		var parts = strings.SplitN(tag, "=", 2)
		if len(parts) == 2 {
			tagparam.Tags = append(tagparam.Tags, &ec2.Tag{
				Key:   aws.String(parts[0]),
				Value: aws.String(parts[1]),
			})
		} else {
			return nil, fmt.Errorf("Skipping bad tag spec - %s\n", tag)
		}
	}
	return tagparam, nil
}

func ec2_EbsRoot(size int) (mapping *ec2.BlockDeviceMapping) {
	return &ec2.BlockDeviceMapping{
		DeviceName: aws.String("/dev/sda1"),
		Ebs: &ec2.EbsBlockDevice{
			DeleteOnTermination: aws.Bool(true),
			VolumeSize:          aws.Int64(int64(size)),
			VolumeType:          aws.String(ec2.VolumeTypeGp2),
		},
	}
}

func ec2_EbsVols(size ...int) (mapping []*ec2.BlockDeviceMapping, err error) {
	mapping = make([]*ec2.BlockDeviceMapping, 0)
	for i, volSize := range size {
		if volSize <= 0 {
			fmt.Fprintln(os.Stderr, "Skipping bad volume size", volSize)
			continue
		}
		if i >= len(DEVICE_NAME) {
			err = fmt.Errorf("You had more volumes then AWS allowed")
			return
		}
		mapping = append(mapping, &ec2.BlockDeviceMapping{
			DeviceName: aws.String("xvd" + DEVICE_NAME[i]),
			Ebs: &ec2.EbsBlockDevice{
				DeleteOnTermination: aws.Bool(true),
				VolumeSize:          aws.Int64(int64(volSize)),
				VolumeType:          aws.String(ec2.VolumeTypeGp2),
			},
		})
	}
	return
}

func ec2_WaitForReady(instId *string) <-chan ec2state {
	out := make(chan ec2state)
	go func() {
		defer close(out)
		param := &ec2.DescribeInstancesInput{InstanceIds: []*string{instId}}
		if err := svc.WaitUntilInstanceRunning(param); err != nil {
			out <- newEc2State(nil, fmt.Errorf("%s - %s", *instId, err))
			return
		}
		resp, err := svc.DescribeInstances(param)
		if err != nil {
			out <- newEc2State(nil, fmt.Errorf("%s - %s", *instId, err))
		} else {
			out <- newEc2State(resp.Reservations[0].Instances[0], nil)
		}
		// NOTE: this should end here
	}()
	return out
}

func newEC2Inst(c *cli.Context, profile *Profile, num2Launch int) (instances []*ec2.Instance, err error) {
	var (
		amiId            = c.String("ami-id")
		networkACLGroups = c.StringSlice("group")
		iamProfile       = c.String("iam-role")
		instVolRoot      = c.Int("root-size")
		keyName          = c.String("ssh-key")
		isPrivate        = c.Bool("subnet-private")
		subnetId         = c.String("subnet-id")
		instTags         = c.StringSlice("tag")
		instType         = c.String("type")
		instVols         = c.IntSlice("volume-size")

		ec2param = &ec2.RunInstancesInput{
			InstanceType:     aws.String(instType),
			MaxCount:         aws.Int64(int64(num2Launch)),
			MinCount:         aws.Int64(1),
			SecurityGroupIds: ec2_findSecurityGroup(&profile.VPC, networkACLGroups...),
		}
	)

	// Step 1: determine the Amazone Machine Image ID
	if amiId != "" {
		ec2param.ImageId = aws.String(amiId)
	} else if len(profile.Ami) != 0 {
		ec2param.ImageId = profile.Ami[0].Id
	} else {
		return nil, fmt.Errorf("Cannot proceed without an AMI")
	}

	// Step 2: determine keypair to use for remote access
	if keyName != "" {
		ec2param.KeyName = aws.String(keyName)
	} else if len(profile.KeyPair) != 0 {
		ec2param.KeyName = profile.KeyPair[0].Name
	} else {
		return nil, fmt.Errorf("Cannot proceed without SSH keypair")
	}

	// Step 3: determine EBS Volume configuration
	ec2param.BlockDeviceMappings = make([]*ec2.BlockDeviceMapping, 0)
	if instVolRoot > 0 {
		ec2param.BlockDeviceMappings = append(ec2param.BlockDeviceMappings, ec2_EbsRoot(instVolRoot))
	}
	if mapping, brr := ec2_EbsVols(instVols...); brr != nil {
		return nil, brr
	} else {
		ec2param.BlockDeviceMappings = append(ec2param.BlockDeviceMappings, mapping...)
	}

	// Step 4: assign IAM role for the EC2 machine
	if strings.HasPrefix(iamProfile, "arn:aws:iam") {
		ec2param.IamInstanceProfile = &ec2.IamInstanceProfileSpecification{
			Arn: aws.String(iamProfile),
		}
	} else if iamProfile != "" {
		ec2param.IamInstanceProfile = &ec2.IamInstanceProfileSpecification{
			Name: aws.String(iamProfile),
		}
	}

	// Step 5: assign accessibility of EC2 instance by subnet
	if subnetId != "" {
		ec2param.SubnetId = aws.String(subnetId)
	} else {
		ec2param.SubnetId = ec2_getSubnet(&profile.VPC, !isPrivate)
	}

	// Step 6: create tags from spec
	tag, err := ec2_tagInstanceParam(instTags)
	if err != nil {
		return
	}

	// Last step: launch + tag instances
	resp, err := svc.RunInstances(ec2param)
	if err != nil {
		return
	}
	if len(tag.Tags) > 0 {
		for _, inst := range resp.Instances {
			tag.Resources = append(tag.Resources, inst.InstanceId)
		}
		_, err = svc.CreateTags(tag)
		if err != nil {
			return
		}
	}

	fmt.Println("Launched instances...")
	return resp.Instances, nil
}

func deployEC2Inst(user, cert, name, org, certpath string, num2Launch int, useDocker bool, instances []*ec2.Instance) <-chan ec2state {
	var wg sync.WaitGroup
	out := make(chan ec2state)
	go func() {
		defer close(out)
		wg.Add(len(instances))
		for idx, inst := range instances {
			go func(index int, ch <-chan ec2state) {
				var state = <-ch
				if state.Instance == nil {
					state.err = fmt.Errorf("Unexpected Instance launch failure")
				} else {
					if num2Launch > 1 {
						state.name = fmt.Sprintf("%s-%03d", name, index)
					} else {
						state.name = name
					}
					tagparam := &ec2.CreateTagsInput{
						Tags: []*ec2.Tag{
							{
								Key:   aws.String("Name"),
								Value: aws.String(state.name),
							},
						},
						Resources: []*string{
							state.InstanceId,
						},
					}
					_, state.err = svc.CreateTags(tagparam)
					if useDocker {
						host := mach.NewDockerHost(org, certpath, user, cert)
						if state.err == nil {
							state.err = host.InstallDockerEngine(*state.PublicIpAddress)
						}
						if state.err == nil {
							state.err = host.InstallDockerEngineCertificate(*state.PublicIpAddress, *state.PrivateIpAddress)
						}
					}
				}
				out <- state
				wg.Done()
			}(idx+1, ec2_WaitForReady(inst.InstanceId))
		}
		wg.Wait()
	}()
	return out
}
<commit_msg>UPDATE: aws create driver use <name>-<id> for identification<commit_after>package aws

import (
	mach "github.com/jeffjen/machine/lib/machine"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/urfave/cli"

	"fmt"
	"math/rand"
	"net"
	"os"
	"strings"
	"sync"
)

var (
	DEVICE_NAME = []string{"b", "c", "d", "e", "f", "g", "h", "i"}
)

type ec2state struct {
	*ec2.Instance
	name string
	err  error
}

func newEc2State(inst *ec2.Instance, err error) ec2state {
	return ec2state{inst, "", err}
}

func getEc2InstanceName(inst *ec2.Instance) (name string) {
	for _, t := range inst.Tags {
		if *t.Key == "Name" {
			return *t.Value
		}
	}
	return
}

func ec2Init() error {
	var resp = new(ec2.DescribeInstancesOutput)

	for more := true; more; {
		params := &ec2.DescribeInstancesInput{}
		if resp.NextToken != nil {
			params.NextToken = resp.NextToken
		}
		resp, err := svc.DescribeInstances(params)
		if err != nil {
			more = false
		} else if len(resp.Reservations) < 1 {
			more = false
		} else {
			for _, r := range resp.Reservations {
				for _, inst := range r.Instances {
					var instanceName = getEc2InstanceName(inst)
					info, ok := instList[instanceName]
					if !ok {
						info = &mach.Instance{Id: *inst.InstanceId}
					}
					if *inst.State.Name == "terminated" {
						if info.Id == *inst.InstanceId {
							delete(instList, instanceName)
						} else {
							continue // ignore
						}
					} else {
						info.Driver = "aws"
						info.State = *inst.State.Name
						info.Id = *inst.InstanceId
						func() {
							var addr *net.TCPAddr
							if inst.PublicIpAddress != nil {
								addr, _ = net.ResolveTCPAddr("tcp", *inst.PublicIpAddress+":2376")
							}
							info.DockerHost = addr
						}()
						instList[instanceName] = info
					}
				}
			}
			more = (resp.NextToken != nil)
		}
	}

	return nil
}

func ec2_getSubnet(profile *VPCProfile, public bool) (subnetId *string) {
	var collection []*string
	for _, subnet := range profile.Subnet {
		if public && *subnet.Public {
			collection = append(collection, subnet.Id)
		} else if !public && !*subnet.Public {
			collection = append(collection, subnet.Id)
		}
	}
	idx := rand.Intn(len(collection))
	return collection[idx]
}

func ec2_findSecurityGroup(profile *VPCProfile, name ...string) (sgId []*string) {
	sgId = make([]*string, 0)
	for _, grp := range name {
		for _, sgrp := range profile.SecurityGroup {
			if *sgrp.Name == grp {
				sgId = append(sgId, sgrp.Id)
			}
		}
	}
	return
}

func ec2_tagInstanceParam(tags []string) (*ec2.CreateTagsInput, error) {
	tagparam := &ec2.CreateTagsInput{
		Tags:      make([]*ec2.Tag, 0),
		Resources: make([]*string, 0),
	}
	for _, tag := range tags {
		var parts = strings.SplitN(tag, "=", 2)
		if len(parts) == 2 {
			tagparam.Tags = append(tagparam.Tags, &ec2.Tag{
				Key:   aws.String(parts[0]),
				Value: aws.String(parts[1]),
			})
		} else {
			return nil, fmt.Errorf("Skipping bad tag spec - %s\n", tag)
		}
	}
	return tagparam, nil
}

func ec2_EbsRoot(size int) (mapping *ec2.BlockDeviceMapping) {
	return &ec2.BlockDeviceMapping{
		DeviceName: aws.String("/dev/sda1"),
		Ebs: &ec2.EbsBlockDevice{
			DeleteOnTermination: aws.Bool(true),
			VolumeSize:          aws.Int64(int64(size)),
			VolumeType:          aws.String(ec2.VolumeTypeGp2),
		},
	}
}

func ec2_EbsVols(size ...int) (mapping []*ec2.BlockDeviceMapping, err error) {
	mapping = make([]*ec2.BlockDeviceMapping, 0)
	for i, volSize := range size {
		if volSize <= 0 {
			fmt.Fprintln(os.Stderr, "Skipping bad volume size", volSize)
			continue
		}
		if i >= len(DEVICE_NAME) {
			err = fmt.Errorf("You had more volumes then AWS allowed")
			return
		}
		mapping = append(mapping, &ec2.BlockDeviceMapping{
			DeviceName: aws.String("xvd" + DEVICE_NAME[i]),
			Ebs: &ec2.EbsBlockDevice{
				DeleteOnTermination: aws.Bool(true),
				VolumeSize:          aws.Int64(int64(volSize)),
				VolumeType:          aws.String(ec2.VolumeTypeGp2),
			},
		})
	}
	return
}

func ec2_WaitForReady(instId *string) <-chan ec2state {
	out := make(chan ec2state)
	go func() {
		defer close(out)
		param := &ec2.DescribeInstancesInput{InstanceIds: []*string{instId}}
		if err := svc.WaitUntilInstanceRunning(param); err != nil {
			out <- newEc2State(nil, fmt.Errorf("%s - %s", *instId, err))
			return
		}
		resp, err := svc.DescribeInstances(param)
		if err != nil {
			out <- newEc2State(nil, fmt.Errorf("%s - %s", *instId, err))
		} else {
			out <- newEc2State(resp.Reservations[0].Instances[0], nil)
		}
		// NOTE: this should end here
	}()
	return out
}

func newEC2Inst(c *cli.Context, profile *Profile, num2Launch int) (instances []*ec2.Instance, err error) {
	var (
		amiId            = c.String("ami-id")
		networkACLGroups = c.StringSlice("group")
		iamProfile       = c.String("iam-role")
		instVolRoot      = c.Int("root-size")
		keyName          = c.String("ssh-key")
		isPrivate        = c.Bool("subnet-private")
		subnetId         = c.String("subnet-id")
		instTags         = c.StringSlice("tag")
		instType         = c.String("type")
		instVols         = c.IntSlice("volume-size")

		ec2param = &ec2.RunInstancesInput{
			InstanceType:     aws.String(instType),
			MaxCount:         aws.Int64(int64(num2Launch)),
			MinCount:         aws.Int64(1),
			SecurityGroupIds: ec2_findSecurityGroup(&profile.VPC, networkACLGroups...),
		}
	)

	// Step 1: determine the Amazone Machine Image ID
	if amiId != "" {
		ec2param.ImageId = aws.String(amiId)
	} else if len(profile.Ami) != 0 {
		ec2param.ImageId = profile.Ami[0].Id
	} else {
		return nil, fmt.Errorf("Cannot proceed without an AMI")
	}

	// Step 2: determine keypair to use for remote access
	if keyName != "" {
		ec2param.KeyName = aws.String(keyName)
	} else if len(profile.KeyPair) != 0 {
		ec2param.KeyName = profile.KeyPair[0].Name
	} else {
		return nil, fmt.Errorf("Cannot proceed without SSH keypair")
	}

	// Step 3: determine EBS Volume configuration
	ec2param.BlockDeviceMappings = make([]*ec2.BlockDeviceMapping, 0)
	if instVolRoot > 0 {
		ec2param.BlockDeviceMappings = append(ec2param.BlockDeviceMappings, ec2_EbsRoot(instVolRoot))
	}
	if mapping, brr := ec2_EbsVols(instVols...); brr != nil {
		return nil, brr
	} else {
		ec2param.BlockDeviceMappings = append(ec2param.BlockDeviceMappings, mapping...)
	}

	// Step 4: assign IAM role for the EC2 machine
	if strings.HasPrefix(iamProfile, "arn:aws:iam") {
		ec2param.IamInstanceProfile = &ec2.IamInstanceProfileSpecification{
			Arn: aws.String(iamProfile),
		}
	} else if iamProfile != "" {
		ec2param.IamInstanceProfile = &ec2.IamInstanceProfileSpecification{
			Name: aws.String(iamProfile),
		}
	}

	// Step 5: assign accessibility of EC2 instance by subnet
	if subnetId != "" {
		ec2param.SubnetId = aws.String(subnetId)
	} else {
		ec2param.SubnetId = ec2_getSubnet(&profile.VPC, !isPrivate)
	}

	// Step 6: create tags from spec
	tag, err := ec2_tagInstanceParam(instTags)
	if err != nil {
		return
	}

	// Last step: launch + tag instances
	resp, err := svc.RunInstances(ec2param)
	if err != nil {
		return
	}
	if len(tag.Tags) > 0 {
		for _, inst := range resp.Instances {
			tag.Resources = append(tag.Resources, inst.InstanceId)
		}
		_, err = svc.CreateTags(tag)
		if err != nil {
			return
		}
	}

	fmt.Println("Launched instances...")
	return resp.Instances, nil
}

func deployEC2Inst(user, cert, name, org, certpath string, num2Launch int, useDocker bool, instances []*ec2.Instance) <-chan ec2state {
	var wg sync.WaitGroup
	out := make(chan ec2state)
	go func() {
		defer close(out)
		wg.Add(len(instances))
		for _, inst := range instances {
			go func(ch <-chan ec2state) {
				var state = <-ch
				if state.Instance == nil {
					state.err = fmt.Errorf("Unexpected Instance launch failure")
				} else {
					state.name = fmt.Sprintf("%s-%s", name, *state.InstanceId)
					tagparam := &ec2.CreateTagsInput{
						Tags: []*ec2.Tag{
							{
								Key:   aws.String("Name"),
								Value: aws.String(state.name),
							},
						},
						Resources: []*string{
							state.InstanceId,
						},
					}
					_, state.err = svc.CreateTags(tagparam)
					if useDocker {
						host := mach.NewDockerHost(org, certpath, user, cert)
						if state.err == nil {
							state.err = host.InstallDockerEngine(*state.PublicIpAddress)
						}
						if state.err == nil {
							state.err = host.InstallDockerEngineCertificate(*state.PublicIpAddress, *state.PrivateIpAddress)
						}
					}
				}
				out <- state
				wg.Done()
			}(ec2_WaitForReady(inst.InstanceId))
		}
		wg.Wait()
	}()
	return out
}
<|endoftext|>