<commit_before>package integrationtests

import (
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path"
	"strconv"
	"time"

	"github.com/lucas-clemente/quic-go/h2quic"
	"github.com/lucas-clemente/quic-go/protocol"
	"github.com/lucas-clemente/quic-go/testdata"
	"github.com/tebeka/selenium"

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

	"testing"
)

const (
	dataLen     = 500 * 1024       // 500 KB
	dataLongLen = 50 * 1024 * 1024 // 50 MB
)

var (
	server    *h2quic.Server
	dataMan   dataManager
	port      string
	uploadDir string

	docker *gexec.Session
)

func TestIntegration(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Integration Tests Suite")
}

var _ = BeforeSuite(func() {
	setupHTTPHandlers()
	setupQuicServer()
	setupSelenium()
})

var _ = AfterSuite(func() {
	err := server.Close()
	Expect(err).NotTo(HaveOccurred())

	stopSelenium()
}, 10)

var _ = BeforeEach(func() {
	// create a new uploadDir for every test
	var err error
	uploadDir, err = ioutil.TempDir("", "quic-upload-dest")
	Expect(err).ToNot(HaveOccurred())
	err = os.MkdirAll(uploadDir, os.ModeDir|0777)
	Expect(err).ToNot(HaveOccurred())
})

var _ = AfterEach(func() {
	// remove uploadDir
	if len(uploadDir) < 20 {
		panic("uploadDir too short")
	}
	os.RemoveAll(uploadDir)

	// remove downloaded file in docker container
	removeDownload("data")
})

func setupHTTPHandlers() {
	defer GinkgoRecover()

	http.HandleFunc("/hello", func(w http.ResponseWriter, r *http.Request) {
		defer GinkgoRecover()
		_, err := io.WriteString(w, "Hello, World!\n")
		Expect(err).NotTo(HaveOccurred())
	})

	http.HandleFunc("/data", func(w http.ResponseWriter, r *http.Request) {
		defer GinkgoRecover()
		data := dataMan.GetData()
		Expect(data).ToNot(HaveLen(0))
		_, err := w.Write(data)
		Expect(err).NotTo(HaveOccurred())
	})

	http.HandleFunc("/echo", func(w http.ResponseWriter, r *http.Request) {
		defer GinkgoRecover()
		_, err := io.Copy(w, r.Body)
		Expect(err).NotTo(HaveOccurred())
	})

	http.HandleFunc("/uploadform", func(w http.ResponseWriter, r *http.Request) {
		defer GinkgoRecover()
		_, err := io.WriteString(w, `
			<form id="form" action="https://quic.clemente.io/uploadhandler" method="post" enctype="multipart/form-data">
		    <input type="file" id="upload" name="uploadfile" />
		  </form>
			`)
		Expect(err).NotTo(HaveOccurred())
	})

	http.HandleFunc("/uploadhandler", func(w http.ResponseWriter, r *http.Request) {
		defer GinkgoRecover()

		err := r.ParseMultipartForm(100 * (1 << 20)) // max. 100 MB
		Expect(err).ToNot(HaveOccurred())
		file, handler, err := r.FormFile("uploadfile")
		Expect(err).ToNot(HaveOccurred())
		defer file.Close()
		f, err := os.OpenFile(path.Join(uploadDir, handler.Filename), os.O_WRONLY|os.O_CREATE, 0666)
		Expect(err).ToNot(HaveOccurred())
		defer f.Close()
		io.Copy(f, file)

		_, err = io.WriteString(w, "")
		Expect(err).NotTo(HaveOccurred())
	})
}

func setupQuicServer() {
	server = &h2quic.Server{
		Server: &http.Server{
			TLSConfig: testdata.GetTLSConfig(),
		},
	}

	addr, err := net.ResolveUDPAddr("udp", "0.0.0.0:0")
	Expect(err).NotTo(HaveOccurred())
	conn, err := net.ListenUDP("udp", addr)
	Expect(err).NotTo(HaveOccurred())
	port = strconv.Itoa(conn.LocalAddr().(*net.UDPAddr).Port)

	go func() {
		defer GinkgoRecover()
		server.Serve(conn)
	}()
}

func setupSelenium() {
	var err error
	pullCmd := exec.Command("docker", "pull", "lclemente/standalone-chrome:beta")
	pull, err := gexec.Start(pullCmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	// Assuming a download at 10 Mbit/s
	Eventually(pull, 10*time.Minute).Should(gexec.Exit(0))

	dockerCmd := exec.Command(
		"docker",
		"run",
		"-i",
		"--rm",
		"-p=4444:4444",
		"--name", "quic-test-selenium",
		"lclemente/standalone-chrome:beta",
	)
	docker, err = gexec.Start(dockerCmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(docker.Out, 10).Should(gbytes.Say("Selenium Server is up and running"))
}

func stopSelenium() {
	docker.Interrupt().Wait(10)
}

func getWebdriverForVersion(version protocol.VersionNumber) selenium.WebDriver {
	caps := selenium.Capabilities{
		"browserName": "chrome",
		"chromeOptions": map[string]interface{}{
			"args": []string{
				"--enable-quic",
				"--no-proxy-server",
				"--origin-to-force-quic-on=quic.clemente.io:443",
				fmt.Sprintf(`--host-resolver-rules=MAP quic.clemente.io:443 %s:%s`, GetLocalIP(), port),
				fmt.Sprintf(`--quic-version=QUIC_VERSION_%d`, version),
			},
		},
	}
	wd, err := selenium.NewRemote(caps, "http://localhost:4444/wd/hub")
	Expect(err).NotTo(HaveOccurred())
	return wd
}

func GetLocalIP() string {
	// First, try finding interface docker0
	i, err := net.InterfaceByName("docker0")
	if err == nil {
		var addrs []net.Addr
		addrs, err = i.Addrs()
		Expect(err).NotTo(HaveOccurred())
		return addrs[0].(*net.IPNet).IP.String()
	}

	addrs, err := net.InterfaceAddrs()
	Expect(err).NotTo(HaveOccurred())
	for _, address := range addrs {
		if ipnet, ok := address.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
			if ipnet.IP.To4() != nil {
				return ipnet.IP.String()
			}
		}
	}
	panic("no addr")
}

func removeDownload(filename string) {
	cmd := exec.Command("docker", "exec", "-i", "quic-test-selenium", "rm", "-f", "/home/seluser/Downloads/"+filename)
	session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(session).Should(gexec.Exit(0))
}

// getDownloadSize gets the file size of a file in the /home/seluser/Downloads folder in the docker container
func getDownloadSize(filename string) int {
	var out bytes.Buffer
	cmd := exec.Command("docker", "exec", "-i", "quic-test-selenium", "stat", "--printf=%s", "/home/seluser/Downloads/"+filename)
	session, err := gexec.Start(cmd, &out, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(session).Should(gexec.Exit())
	if session.ExitCode() != 0 {
		return 0
	}
	Expect(out.Bytes()).ToNot(BeEmpty())
	size, err := strconv.Atoi(string(out.Bytes()))
	Expect(err).NotTo(HaveOccurred())
	return size
}

// getFileSize gets the file size of a file on the local file system
func getFileSize(filename string) int {
	file, err := os.Open(filename)
	Expect(err).ToNot(HaveOccurred())
	fi, err := file.Stat()
	Expect(err).ToNot(HaveOccurred())
	return int(fi.Size())
}

// getDownloadMD5 gets the md5 sum file of a file in the /home/seluser/Downloads folder in the docker container
func getDownloadMD5(filename string) []byte {
	var out bytes.Buffer
	cmd := exec.Command("docker", "exec", "-i", "quic-test-selenium", "md5sum", "/home/seluser/Downloads/"+filename)
	session, err := gexec.Start(cmd, &out, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(session).Should(gexec.Exit())
	if session.ExitCode() != 0 {
		return nil
	}
	Expect(out.Bytes()).ToNot(BeEmpty())
	res, err := hex.DecodeString(string(out.Bytes()[0:32]))
	Expect(err).NotTo(HaveOccurred())
	return res
}

// getFileMD5 gets the md5 sum of a file on the local file system
func getFileMD5(filepath string) []byte {
	var result []byte
	file, err := os.Open(filepath)
	Expect(err).ToNot(HaveOccurred())
	defer file.Close()

	hash := md5.New()
	_, err = io.Copy(hash, file)
	Expect(err).ToNot(HaveOccurred())
	return hash.Sum(result)
}

// copyFileToDocker copies a file from the local file system into the /home/seluser/ directory in the docker container
func copyFileToDocker(filepath string) {
	cmd := exec.Command("docker", "cp", filepath, "quic-test-selenium:/home/seluser/")
	session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(session).Should(gexec.Exit(0))
}
<commit_msg>use chrome dev in integration tests for testing v36<commit_after>package integrationtests

import (
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path"
	"strconv"
	"time"

	"github.com/lucas-clemente/quic-go/h2quic"
	"github.com/lucas-clemente/quic-go/protocol"
	"github.com/lucas-clemente/quic-go/testdata"
	"github.com/tebeka/selenium"

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

	"testing"
)

const (
	dataLen     = 500 * 1024       // 500 KB
	dataLongLen = 50 * 1024 * 1024 // 50 MB
)

var (
	server    *h2quic.Server
	dataMan   dataManager
	port      string
	uploadDir string

	docker *gexec.Session
)

func TestIntegration(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Integration Tests Suite")
}

var _ = BeforeSuite(func() {
	setupHTTPHandlers()
	setupQuicServer()
	setupSelenium()
})

var _ = AfterSuite(func() {
	err := server.Close()
	Expect(err).NotTo(HaveOccurred())

	stopSelenium()
}, 10)

var _ = BeforeEach(func() {
	// create a new uploadDir for every test
	var err error
	uploadDir, err = ioutil.TempDir("", "quic-upload-dest")
	Expect(err).ToNot(HaveOccurred())
	err = os.MkdirAll(uploadDir, os.ModeDir|0777)
	Expect(err).ToNot(HaveOccurred())
})

var _ = AfterEach(func() {
	// remove uploadDir
	if len(uploadDir) < 20 {
		panic("uploadDir too short")
	}
	os.RemoveAll(uploadDir)

	// remove downloaded file in docker container
	removeDownload("data")
})

func setupHTTPHandlers() {
	defer GinkgoRecover()

	http.HandleFunc("/hello", func(w http.ResponseWriter, r *http.Request) {
		defer GinkgoRecover()
		_, err := io.WriteString(w, "Hello, World!\n")
		Expect(err).NotTo(HaveOccurred())
	})

	http.HandleFunc("/data", func(w http.ResponseWriter, r *http.Request) {
		defer GinkgoRecover()
		data := dataMan.GetData()
		Expect(data).ToNot(HaveLen(0))
		_, err := w.Write(data)
		Expect(err).NotTo(HaveOccurred())
	})

	http.HandleFunc("/echo", func(w http.ResponseWriter, r *http.Request) {
		defer GinkgoRecover()
		_, err := io.Copy(w, r.Body)
		Expect(err).NotTo(HaveOccurred())
	})

	http.HandleFunc("/uploadform", func(w http.ResponseWriter, r *http.Request) {
		defer GinkgoRecover()
		_, err := io.WriteString(w, `
			<form id="form" action="https://quic.clemente.io/uploadhandler" method="post" enctype="multipart/form-data">
		    <input type="file" id="upload" name="uploadfile" />
		  </form>
			`)
		Expect(err).NotTo(HaveOccurred())
	})

	http.HandleFunc("/uploadhandler", func(w http.ResponseWriter, r *http.Request) {
		defer GinkgoRecover()

		err := r.ParseMultipartForm(100 * (1 << 20)) // max. 100 MB
		Expect(err).ToNot(HaveOccurred())
		file, handler, err := r.FormFile("uploadfile")
		Expect(err).ToNot(HaveOccurred())
		defer file.Close()
		f, err := os.OpenFile(path.Join(uploadDir, handler.Filename), os.O_WRONLY|os.O_CREATE, 0666)
		Expect(err).ToNot(HaveOccurred())
		defer f.Close()
		io.Copy(f, file)

		_, err = io.WriteString(w, "")
		Expect(err).NotTo(HaveOccurred())
	})
}

func setupQuicServer() {
	server = &h2quic.Server{
		Server: &http.Server{
			TLSConfig: testdata.GetTLSConfig(),
		},
	}

	addr, err := net.ResolveUDPAddr("udp", "0.0.0.0:0")
	Expect(err).NotTo(HaveOccurred())
	conn, err := net.ListenUDP("udp", addr)
	Expect(err).NotTo(HaveOccurred())
	port = strconv.Itoa(conn.LocalAddr().(*net.UDPAddr).Port)

	go func() {
		defer GinkgoRecover()
		server.Serve(conn)
	}()
}

func setupSelenium() {
	var err error
	pullCmd := exec.Command("docker", "pull", "lclemente/standalone-chrome:dev")
	pull, err := gexec.Start(pullCmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	// Assuming a download at 10 Mbit/s
	Eventually(pull, 10*time.Minute).Should(gexec.Exit(0))

	dockerCmd := exec.Command(
		"docker",
		"run",
		"-i",
		"--rm",
		"-p=4444:4444",
		"--name", "quic-test-selenium",
		"lclemente/standalone-chrome:dev",
	)
	docker, err = gexec.Start(dockerCmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(docker.Out, 10).Should(gbytes.Say("Selenium Server is up and running"))
}

func stopSelenium() {
	docker.Interrupt().Wait(10)
}

func getWebdriverForVersion(version protocol.VersionNumber) selenium.WebDriver {
	caps := selenium.Capabilities{
		"browserName": "chrome",
		"chromeOptions": map[string]interface{}{
			"args": []string{
				"--enable-quic",
				"--no-proxy-server",
				"--origin-to-force-quic-on=quic.clemente.io:443",
				fmt.Sprintf(`--host-resolver-rules=MAP quic.clemente.io:443 %s:%s`, GetLocalIP(), port),
				fmt.Sprintf(`--quic-version=QUIC_VERSION_%d`, version),
			},
		},
	}
	wd, err := selenium.NewRemote(caps, "http://localhost:4444/wd/hub")
	Expect(err).NotTo(HaveOccurred())
	return wd
}

func GetLocalIP() string {
	// First, try finding interface docker0
	i, err := net.InterfaceByName("docker0")
	if err == nil {
		var addrs []net.Addr
		addrs, err = i.Addrs()
		Expect(err).NotTo(HaveOccurred())
		return addrs[0].(*net.IPNet).IP.String()
	}

	addrs, err := net.InterfaceAddrs()
	Expect(err).NotTo(HaveOccurred())
	for _, address := range addrs {
		if ipnet, ok := address.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
			if ipnet.IP.To4() != nil {
				return ipnet.IP.String()
			}
		}
	}
	panic("no addr")
}

func removeDownload(filename string) {
	cmd := exec.Command("docker", "exec", "-i", "quic-test-selenium", "rm", "-f", "/home/seluser/Downloads/"+filename)
	session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(session).Should(gexec.Exit(0))
}

// getDownloadSize gets the file size of a file in the /home/seluser/Downloads folder in the docker container
func getDownloadSize(filename string) int {
	var out bytes.Buffer
	cmd := exec.Command("docker", "exec", "-i", "quic-test-selenium", "stat", "--printf=%s", "/home/seluser/Downloads/"+filename)
	session, err := gexec.Start(cmd, &out, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(session).Should(gexec.Exit())
	if session.ExitCode() != 0 {
		return 0
	}
	Expect(out.Bytes()).ToNot(BeEmpty())
	size, err := strconv.Atoi(string(out.Bytes()))
	Expect(err).NotTo(HaveOccurred())
	return size
}

// getFileSize gets the file size of a file on the local file system
func getFileSize(filename string) int {
	file, err := os.Open(filename)
	Expect(err).ToNot(HaveOccurred())
	fi, err := file.Stat()
	Expect(err).ToNot(HaveOccurred())
	return int(fi.Size())
}

// getDownloadMD5 gets the md5 sum file of a file in the /home/seluser/Downloads folder in the docker container
func getDownloadMD5(filename string) []byte {
	var out bytes.Buffer
	cmd := exec.Command("docker", "exec", "-i", "quic-test-selenium", "md5sum", "/home/seluser/Downloads/"+filename)
	session, err := gexec.Start(cmd, &out, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(session).Should(gexec.Exit())
	if session.ExitCode() != 0 {
		return nil
	}
	Expect(out.Bytes()).ToNot(BeEmpty())
	res, err := hex.DecodeString(string(out.Bytes()[0:32]))
	Expect(err).NotTo(HaveOccurred())
	return res
}

// getFileMD5 gets the md5 sum of a file on the local file system
func getFileMD5(filepath string) []byte {
	var result []byte
	file, err := os.Open(filepath)
	Expect(err).ToNot(HaveOccurred())
	defer file.Close()

	hash := md5.New()
	_, err = io.Copy(hash, file)
	Expect(err).ToNot(HaveOccurred())
	return hash.Sum(result)
}

// copyFileToDocker copies a file from the local file system into the /home/seluser/ directory in the docker container
func copyFileToDocker(filepath string) {
	cmd := exec.Command("docker", "cp", filepath, "quic-test-selenium:/home/seluser/")
	session, err := gexec.Start(cmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())
	Eventually(session).Should(gexec.Exit(0))
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Nodetemple <hostmaster@nodetemple.com>

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 main

import (
	"strings"

	"github.com/codegangsta/cli"
	"github.com/nodetemple/nodetemple/common"
	"github.com/nodetemple/nodetemple/nodectl/util"
	"github.com/nodetemple/nodetemple/nodectl/command"
)

func main() {
	app := cli.NewApp()
	app.Name = "nodectl"
	app.Usage = "CLI for an orchestration of CoreOS and Kubernetes cluster"
	app.Version = common.Version
	app.Flags = []cli.Flag{
		cli.StringFlag{Name: "providers, p", Usage: "A comma-separated list of IaaS providers ("+strings.Join(common.AvailableProviders, ",")+") and API keys, format: 'provider:api-key,...'", EnvVar: util.EnvVarConv(app.Name, "providers"),},
		cli.BoolFlag{Name: "debug", Usage: "Print out more debug information to stderr"},
	}
	app.Commands = []cli.Command{
		command.DemoCmd,
	}
	app.CommandNotFound = func(c *cli.Context, command string) {
		util.Err("unknown command '%v'\nRun '%v help [command]' for usage information", command, c.App.Name)
	}

	if err := app.Run(os.Args); err != nil {
		util.Err(err)
	}
}
<commit_msg>Remove debug flag<commit_after>/*
Copyright 2015 Nodetemple <hostmaster@nodetemple.com>

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 main

import (
	"strings"

	"github.com/codegangsta/cli"
	"github.com/nodetemple/nodetemple/common"
	"github.com/nodetemple/nodetemple/nodectl/util"
	"github.com/nodetemple/nodetemple/nodectl/command"
)

func main() {
	app := cli.NewApp()
	app.Name = "nodectl"
	app.Usage = "CLI for an orchestration of CoreOS and Kubernetes cluster"
	app.Version = common.Version
	app.Flags = []cli.Flag{
		cli.StringFlag{Name: "providers, p", Usage: "A comma-separated list of IaaS providers ("+strings.Join(common.AvailableProviders, ",")+") and API keys, format: 'provider:api-key,...'", EnvVar: util.EnvVarConv(app.Name, "providers"),},
	}
	app.Commands = []cli.Command{
		command.DemoCmd,
	}
	app.CommandNotFound = func(c *cli.Context, command string) {
		util.Err("unknown command '%v'\nRun '%v help [command]' for usage information", command, c.App.Name)
	}

	if err := app.Run(os.Args); err != nil {
		util.Err(err)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Nodetemple <hostmaster@nodetemple.com>

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 main

import (
	"os"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/nodetemple/nodetemple/common"
	"github.com/nodetemple/nodetemple/nodectl/util"
	"github.com/nodetemple/nodetemple/nodectl/command"
)

func main() {
	app := cli.NewApp()
	app.Name = "nodectl"
	app.Usage = "CLI for an orchestration of CoreOS and Kubernetes cluster"
	app.Version = common.Version
	app.Flags = []cli.Flag{
		cli.StringFlag{Name: "providers, p", Usage: "A comma-separated list of IaaS providers ("+strings.Join(common.AvailableProviders, ",")+") and API keys, format: 'provider:api-key,...'", EnvVar: util.EnvVarConv(app.Name, "providers"),},
		cli.BoolFlag{Name: "debug", Usage: "Print out more debug information to stderr"},
	}
	app.Before = func(c *cli.Context) {
		if c.String("providers") == "" && !c.Bool("help") && !c.Bool("version") {
			util.Err("must provide API Key via NODECTL_PROVIDERS environment variable or via CLI argument.")
		}
	}
	app.Commands = []cli.Command{
		command.DemoCmd(),
	}
	app.CommandNotFound = func(c *cli.Context, command string) {
		util.Err("unknown command '%v'\nRun '%v help [command]' for usage information", command, c.App.Name)
	}

	app.RunAndExitOnError()
}
<commit_msg>Err: no providers<commit_after>/*
Copyright 2015 Nodetemple <hostmaster@nodetemple.com>

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 main

import (
	"os"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/nodetemple/nodetemple/common"
	"github.com/nodetemple/nodetemple/nodectl/util"
	"github.com/nodetemple/nodetemple/nodectl/command"
)

func main() {
	app := cli.NewApp()
	app.Name = "nodectl"
	app.Usage = "CLI for an orchestration of CoreOS and Kubernetes cluster"
	app.Version = common.Version
	app.Flags = []cli.Flag{
		cli.StringFlag{Name: "providers, p", Usage: "A comma-separated list of IaaS providers ("+strings.Join(common.AvailableProviders, ",")+") and API keys, format: 'provider:api-key,...'", EnvVar: util.EnvVarConv(app.Name, "providers"),},
		cli.BoolFlag{Name: "debug", Usage: "Print out more debug information to stderr"},
	}
	app.Before = func(c *cli.Context) {
		if c.String("providers") == "" && !c.Bool("help") && !c.Bool("version") {
			util.Err("set at least one provider with a valid API key")
		}
	}
	app.Commands = []cli.Command{
		command.DemoCmd(),
	}
	app.CommandNotFound = func(c *cli.Context, command string) {
		util.Err("unknown command '%v'\nRun '%v help [command]' for usage information", command, c.App.Name)
	}

	app.RunAndExitOnError()
}
<|endoftext|>


<commit_before>package client

import (
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"

	"golang.org/x/oauth2"
)

const (
	// rateLimit is the wait time between queries.
	rateLimit = 3 * time.Second
)

type client struct {
	// agent is the client's User-Agent in http requests.
	agent string
	// id is the bot's OAuth2 client id.
	id string
	// secret is the bot's OAuth2 client secret.
	secret string
	// user is the bot's username on reddit.
	user string
	// pass is the bot's password on reddit.
	pass string

	// authMu protects authentication fields.
	authMu sync.Mutex
	// cli is the authenticated client to execute requests with.
	cli *http.Client
	// token is the OAuth2 token cli uses to authenticate.
	token *oauth2.Token

	// rateMu protects rate limiting fields.
	rateMu sync.Mutex
	// nextReq is the time at which it is ok to make the next request.
	nextReq time.Time
}

// Do wraps the execution of http requests. It updates authentications and rate
// limits requests to Reddit to comply with the API rules. It returns the
// response body.
func (c *client) Do(r *http.Request) (io.ReadCloser, error) {
	c.rateRequest()
	if !c.token.Valid() {
		var err error
		c.cli, c.token, err = build(c.id, c.secret, c.user, c.pass)
		if err != nil {
			return nil, err
		}
	}
	return c.exec(r)
}

// exec executes an http request and returns the response body.
func (c *client) exec(r *http.Request) (io.ReadCloser, error) {
	resp, err := c.doRaw(r)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("bad response code: %d\n"+
			"request was: %v\n",
			resp.StatusCode,
			r)
	}

	if resp.Body == nil {
		return nil, fmt.Errorf("no body in response")
	}

	return resp.Body, nil
}

// doRaw executes an http Request using an authenticated client, and the configured
// user agent.
func (c *client) doRaw(r *http.Request) (*http.Response, error) {
	if r.Header == nil {
		r.Header = make(http.Header)
	}
	r.Header.Add("User-Agent", c.agent)
	return c.cli.Do(r)
}

// rateRequest blocks until the rate limits have been abided by.
func (c *client) rateRequest() {
	c.rateMu.Lock()
	defer c.rateMu.Unlock()

	if time.Now().After(c.nextReq) {
		c.nextReq = time.Now().Add(rateLimit)
		return
	}

	currentReq := c.nextReq
	c.nextReq = currentReq.Add(rateLimit)
	<-time.After(currentReq.Sub(time.Now()))
}
<commit_msg>Reduce wait period between requests.<commit_after>package client

import (
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"

	"golang.org/x/oauth2"
)

const (
	// rateLimit is the wait time between requests to Reddit.
	rateLimit = time.Second
)

type client struct {
	// agent is the client's User-Agent in http requests.
	agent string
	// id is the bot's OAuth2 client id.
	id string
	// secret is the bot's OAuth2 client secret.
	secret string
	// user is the bot's username on reddit.
	user string
	// pass is the bot's password on reddit.
	pass string

	// authMu protects authentication fields.
	authMu sync.Mutex
	// cli is the authenticated client to execute requests with.
	cli *http.Client
	// token is the OAuth2 token cli uses to authenticate.
	token *oauth2.Token

	// rateMu protects rate limiting fields.
	rateMu sync.Mutex
	// nextReq is the time at which it is ok to make the next request.
	nextReq time.Time
}

// Do wraps the execution of http requests. It updates authentications and rate
// limits requests to Reddit to comply with the API rules. It returns the
// response body.
func (c *client) Do(r *http.Request) (io.ReadCloser, error) {
	c.rateRequest()
	if !c.token.Valid() {
		var err error
		c.cli, c.token, err = build(c.id, c.secret, c.user, c.pass)
		if err != nil {
			return nil, err
		}
	}
	return c.exec(r)
}

// exec executes an http request and returns the response body.
func (c *client) exec(r *http.Request) (io.ReadCloser, error) {
	resp, err := c.doRaw(r)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("bad response code: %d\n"+
			"request was: %v\n",
			resp.StatusCode,
			r)
	}

	if resp.Body == nil {
		return nil, fmt.Errorf("no body in response")
	}

	return resp.Body, nil
}

// doRaw executes an http Request using an authenticated client, and the configured
// user agent.
func (c *client) doRaw(r *http.Request) (*http.Response, error) {
	if r.Header == nil {
		r.Header = make(http.Header)
	}
	r.Header.Add("User-Agent", c.agent)
	return c.cli.Do(r)
}

// rateRequest blocks until the rate limits have been abided by.
func (c *client) rateRequest() {
	c.rateMu.Lock()
	defer c.rateMu.Unlock()

	if time.Now().After(c.nextReq) {
		c.nextReq = time.Now().Add(rateLimit)
		return
	}

	currentReq := c.nextReq
	c.nextReq = currentReq.Add(rateLimit)
	<-time.After(currentReq.Sub(time.Now()))
}
<|endoftext|>


<commit_before>package nomad

import (
	"fmt"
	"io"
	"net"
	"os"
	"runtime"
	"time"

	"github.com/hashicorp/memberlist"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/scheduler"
	"github.com/hashicorp/raft"
	"github.com/hashicorp/serf/serf"
)

const (
	DefaultRegion   = "region1"
	DefaultDC       = "dc1"
	DefaultSerfPort = 4648
)

// These are the protocol versions that Nomad can understand
const (
	ProtocolVersionMin uint8 = 1
	ProtocolVersionMax       = 1
)

// ProtocolVersionMap is the mapping of Nomad protocol versions
// to Serf protocol versions. We mask the Serf protocols using
// our own protocol version.
var protocolVersionMap map[uint8]uint8

func init() {
	protocolVersionMap = map[uint8]uint8{
		1: 5,
	}
}

var (
	DefaultRPCAddr = &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 4647}
)

// Config is used to parameterize the server
type Config struct {
	// Bootstrap mode is used to bring up the first Consul server.
	// It is required so that it can elect a leader without any
	// other nodes being present
	Bootstrap bool

	// BootstrapExpect mode is used to automatically bring up a collection of
	// Consul servers. This can be used to automatically bring up a collection
	// of nodes.
	BootstrapExpect int

	// DataDir is the directory to store our state in
	DataDir string

	// DevMode is used for development purposes only and limits the
	// use of persistence or state.
	DevMode bool

	// DevDisableBootstrap is used to disable bootstrap mode while
	// in DevMode. This is largely used for testing.
	DevDisableBootstrap bool

	// LogOutput is the location to write logs to. If this is not set,
	// logs will go to stderr.
	LogOutput io.Writer

	// ProtocolVersion is the protocol version to speak. This must be between
	// ProtocolVersionMin and ProtocolVersionMax.
	ProtocolVersion uint8

	// RPCAddr is the RPC address used by Nomad. This should be reachable
	// by the other servers and clients
	RPCAddr *net.TCPAddr

	// RPCAdvertise is the address that is advertised to other nodes for
	// the RPC endpoint. This can differ from the RPC address, if for example
	// the RPCAddr is unspecified "0.0.0.0:4646", but this address must be
	// reachable
	RPCAdvertise *net.TCPAddr

	// RaftConfig is the configuration used for Raft in the local DC
	RaftConfig *raft.Config

	// RaftTimeout is applied to any network traffic for raft. Defaults to 10s.
	RaftTimeout time.Duration

	// RequireTLS ensures that all RPC traffic is protected with TLS
	RequireTLS bool

	// SerfConfig is the configuration for the serf cluster
	SerfConfig *serf.Config

	// Node name is the name we use to advertise. Defaults to hostname.
	NodeName string

	// Region is the region this Nomad server belongs to.
	Region string

	// Datacenter is the datacenter this Nomad server belongs to.
	Datacenter string

	// Build is a string that is gossiped around, and can be used to help
	// operators track which versions are actively deployed
	Build string

	// NumSchedulers is the number of scheduler thread that are run.
	// This can be as many as one per core, or zero to disable this server
	// from doing any scheduling work.
	NumSchedulers int

	// EnabledSchedulers controls the set of sub-schedulers that are
	// enabled for this server to handle. This will restrict the evaluations
	// that the workers dequeue for processing.
	EnabledSchedulers []string

	// ReconcileInterval controls how often we reconcile the strongly
	// consistent store with the Serf info. This is used to handle nodes
	// that are force removed, as well as intermittent unavailability during
	// leader election.
	ReconcileInterval time.Duration

	// EvalGCInterval is how often we dispatch a job to GC evaluations
	EvalGCInterval time.Duration

	// EvalGCThreshold is how "old" an evaluation must be to be eligible
	// for GC. This gives users some time to debug a failed evaluation.
	EvalGCThreshold time.Duration

	// NodeGCInterval is how often we dispatch a job to GC failed nodes.
	NodeGCInterval time.Duration

	// NodeGCThreshold is how "old" a nodemust be to be eligible
	// for GC. This gives users some time to view and debug a failed nodes.
	NodeGCThreshold time.Duration

	// EvalNackTimeout controls how long we allow a sub-scheduler to
	// work on an evaluation before we consider it failed and Nack it.
	// This allows that evaluation to be handed to another sub-scheduler
	// to work on. Defaults to 60 seconds. This should be long enough that
	// no evaluation hits it unless the sub-scheduler has failed.
	EvalNackTimeout time.Duration

	// EvalDeliveryLimit is the limit of attempts we make to deliver and
	// process an evaluation. This is used so that an eval that will never
	// complete eventually fails out of the system.
	EvalDeliveryLimit int

	// MinHeartbeatTTL is the minimum time between heartbeats.
	// This is used as a floor to prevent excessive updates.
	MinHeartbeatTTL time.Duration

	// MaxHeartbeatsPerSecond is the maximum target rate of heartbeats
	// being processed per second. This allows the TTL to be increased
	// to meet the target rate.
	MaxHeartbeatsPerSecond float64

	// HeartbeatGrace is the additional time given as a grace period
	// beyond the TTL to account for network and processing delays
	// as well as clock skew.
	HeartbeatGrace time.Duration

	// FailoverHeartbeatTTL is the TTL applied to heartbeats after
	// a new leader is elected, since we no longer know the status
	// of all the heartbeats.
	FailoverHeartbeatTTL time.Duration
}

// CheckVersion is used to check if the ProtocolVersion is valid
func (c *Config) CheckVersion() error {
	if c.ProtocolVersion < ProtocolVersionMin {
		return fmt.Errorf("Protocol version '%d' too low. Must be in range: [%d, %d]",
			c.ProtocolVersion, ProtocolVersionMin, ProtocolVersionMax)
	} else if c.ProtocolVersion > ProtocolVersionMax {
		return fmt.Errorf("Protocol version '%d' too high. Must be in range: [%d, %d]",
			c.ProtocolVersion, ProtocolVersionMin, ProtocolVersionMax)
	}
	return nil
}

// DefaultConfig returns the default configuration
func DefaultConfig() *Config {
	hostname, err := os.Hostname()
	if err != nil {
		panic(err)
	}

	c := &Config{
		Region:                 DefaultRegion,
		Datacenter:             DefaultDC,
		NodeName:               hostname,
		ProtocolVersion:        ProtocolVersionMax,
		RaftConfig:             raft.DefaultConfig(),
		RaftTimeout:            10 * time.Second,
		RPCAddr:                DefaultRPCAddr,
		SerfConfig:             serf.DefaultConfig(),
		NumSchedulers:          1,
		ReconcileInterval:      60 * time.Second,
		EvalGCInterval:         5 * time.Minute,
		EvalGCThreshold:        1 * time.Hour,
		NodeGCInterval:         5 * time.Minute,
		NodeGCThreshold:        24 * time.Hour,
		EvalNackTimeout:        60 * time.Second,
		EvalDeliveryLimit:      3,
		MinHeartbeatTTL:        10 * time.Second,
		MaxHeartbeatsPerSecond: 50.0,
		HeartbeatGrace:         10 * time.Second,
		FailoverHeartbeatTTL:   300 * time.Second,
	}

	// Enable all known schedulers by default
	c.EnabledSchedulers = make([]string, 0, len(scheduler.BuiltinSchedulers))
	for name := range scheduler.BuiltinSchedulers {
		c.EnabledSchedulers = append(c.EnabledSchedulers, name)
	}
	c.EnabledSchedulers = append(c.EnabledSchedulers, structs.JobTypeCore)

	// Default the number of schedulers to match the coores
	c.NumSchedulers = runtime.NumCPU()

	// Increase our reap interval to 3 days instead of 24h.
	c.SerfConfig.ReconnectTimeout = 3 * 24 * time.Hour

	// Serf should use the WAN timing, since we are using it
	// to communicate between DC's
	c.SerfConfig.MemberlistConfig = memberlist.DefaultWANConfig()
	c.SerfConfig.MemberlistConfig.BindPort = DefaultSerfPort

	// Disable shutdown on removal
	c.RaftConfig.ShutdownOnRemove = false
	return c
}
<commit_msg>Change serf version to 4 so the entire test suite doesn't fail<commit_after>package nomad

import (
	"fmt"
	"io"
	"net"
	"os"
	"runtime"
	"time"

	"github.com/hashicorp/memberlist"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/scheduler"
	"github.com/hashicorp/raft"
	"github.com/hashicorp/serf/serf"
)

const (
	DefaultRegion   = "region1"
	DefaultDC       = "dc1"
	DefaultSerfPort = 4648
)

// These are the protocol versions that Nomad can understand
const (
	ProtocolVersionMin uint8 = 1
	ProtocolVersionMax       = 1
)

// ProtocolVersionMap is the mapping of Nomad protocol versions
// to Serf protocol versions. We mask the Serf protocols using
// our own protocol version.
var protocolVersionMap map[uint8]uint8

func init() {
	protocolVersionMap = map[uint8]uint8{
		1: 4, // TODO change this back to 5 when serf is ready
	}
}

var (
	DefaultRPCAddr = &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 4647}
)

// Config is used to parameterize the server
type Config struct {
	// Bootstrap mode is used to bring up the first Consul server.
	// It is required so that it can elect a leader without any
	// other nodes being present
	Bootstrap bool

	// BootstrapExpect mode is used to automatically bring up a collection of
	// Consul servers. This can be used to automatically bring up a collection
	// of nodes.
	BootstrapExpect int

	// DataDir is the directory to store our state in
	DataDir string

	// DevMode is used for development purposes only and limits the
	// use of persistence or state.
	DevMode bool

	// DevDisableBootstrap is used to disable bootstrap mode while
	// in DevMode. This is largely used for testing.
	DevDisableBootstrap bool

	// LogOutput is the location to write logs to. If this is not set,
	// logs will go to stderr.
	LogOutput io.Writer

	// ProtocolVersion is the protocol version to speak. This must be between
	// ProtocolVersionMin and ProtocolVersionMax.
	ProtocolVersion uint8

	// RPCAddr is the RPC address used by Nomad. This should be reachable
	// by the other servers and clients
	RPCAddr *net.TCPAddr

	// RPCAdvertise is the address that is advertised to other nodes for
	// the RPC endpoint. This can differ from the RPC address, if for example
	// the RPCAddr is unspecified "0.0.0.0:4646", but this address must be
	// reachable
	RPCAdvertise *net.TCPAddr

	// RaftConfig is the configuration used for Raft in the local DC
	RaftConfig *raft.Config

	// RaftTimeout is applied to any network traffic for raft. Defaults to 10s.
	RaftTimeout time.Duration

	// RequireTLS ensures that all RPC traffic is protected with TLS
	RequireTLS bool

	// SerfConfig is the configuration for the serf cluster
	SerfConfig *serf.Config

	// Node name is the name we use to advertise. Defaults to hostname.
	NodeName string

	// Region is the region this Nomad server belongs to.
	Region string

	// Datacenter is the datacenter this Nomad server belongs to.
	Datacenter string

	// Build is a string that is gossiped around, and can be used to help
	// operators track which versions are actively deployed
	Build string

	// NumSchedulers is the number of scheduler thread that are run.
	// This can be as many as one per core, or zero to disable this server
	// from doing any scheduling work.
	NumSchedulers int

	// EnabledSchedulers controls the set of sub-schedulers that are
	// enabled for this server to handle. This will restrict the evaluations
	// that the workers dequeue for processing.
	EnabledSchedulers []string

	// ReconcileInterval controls how often we reconcile the strongly
	// consistent store with the Serf info. This is used to handle nodes
	// that are force removed, as well as intermittent unavailability during
	// leader election.
	ReconcileInterval time.Duration

	// EvalGCInterval is how often we dispatch a job to GC evaluations
	EvalGCInterval time.Duration

	// EvalGCThreshold is how "old" an evaluation must be to be eligible
	// for GC. This gives users some time to debug a failed evaluation.
	EvalGCThreshold time.Duration

	// NodeGCInterval is how often we dispatch a job to GC failed nodes.
	NodeGCInterval time.Duration

	// NodeGCThreshold is how "old" a nodemust be to be eligible
	// for GC. This gives users some time to view and debug a failed nodes.
	NodeGCThreshold time.Duration

	// EvalNackTimeout controls how long we allow a sub-scheduler to
	// work on an evaluation before we consider it failed and Nack it.
	// This allows that evaluation to be handed to another sub-scheduler
	// to work on. Defaults to 60 seconds. This should be long enough that
	// no evaluation hits it unless the sub-scheduler has failed.
	EvalNackTimeout time.Duration

	// EvalDeliveryLimit is the limit of attempts we make to deliver and
	// process an evaluation. This is used so that an eval that will never
	// complete eventually fails out of the system.
	EvalDeliveryLimit int

	// MinHeartbeatTTL is the minimum time between heartbeats.
	// This is used as a floor to prevent excessive updates.
	MinHeartbeatTTL time.Duration

	// MaxHeartbeatsPerSecond is the maximum target rate of heartbeats
	// being processed per second. This allows the TTL to be increased
	// to meet the target rate.
	MaxHeartbeatsPerSecond float64

	// HeartbeatGrace is the additional time given as a grace period
	// beyond the TTL to account for network and processing delays
	// as well as clock skew.
	HeartbeatGrace time.Duration

	// FailoverHeartbeatTTL is the TTL applied to heartbeats after
	// a new leader is elected, since we no longer know the status
	// of all the heartbeats.
	FailoverHeartbeatTTL time.Duration
}

// CheckVersion is used to check if the ProtocolVersion is valid
func (c *Config) CheckVersion() error {
	if c.ProtocolVersion < ProtocolVersionMin {
		return fmt.Errorf("Protocol version '%d' too low. Must be in range: [%d, %d]",
			c.ProtocolVersion, ProtocolVersionMin, ProtocolVersionMax)
	} else if c.ProtocolVersion > ProtocolVersionMax {
		return fmt.Errorf("Protocol version '%d' too high. Must be in range: [%d, %d]",
			c.ProtocolVersion, ProtocolVersionMin, ProtocolVersionMax)
	}
	return nil
}

// DefaultConfig returns the default configuration
func DefaultConfig() *Config {
	hostname, err := os.Hostname()
	if err != nil {
		panic(err)
	}

	c := &Config{
		Region:                 DefaultRegion,
		Datacenter:             DefaultDC,
		NodeName:               hostname,
		ProtocolVersion:        ProtocolVersionMax,
		RaftConfig:             raft.DefaultConfig(),
		RaftTimeout:            10 * time.Second,
		RPCAddr:                DefaultRPCAddr,
		SerfConfig:             serf.DefaultConfig(),
		NumSchedulers:          1,
		ReconcileInterval:      60 * time.Second,
		EvalGCInterval:         5 * time.Minute,
		EvalGCThreshold:        1 * time.Hour,
		NodeGCInterval:         5 * time.Minute,
		NodeGCThreshold:        24 * time.Hour,
		EvalNackTimeout:        60 * time.Second,
		EvalDeliveryLimit:      3,
		MinHeartbeatTTL:        10 * time.Second,
		MaxHeartbeatsPerSecond: 50.0,
		HeartbeatGrace:         10 * time.Second,
		FailoverHeartbeatTTL:   300 * time.Second,
	}

	// Enable all known schedulers by default
	c.EnabledSchedulers = make([]string, 0, len(scheduler.BuiltinSchedulers))
	for name := range scheduler.BuiltinSchedulers {
		c.EnabledSchedulers = append(c.EnabledSchedulers, name)
	}
	c.EnabledSchedulers = append(c.EnabledSchedulers, structs.JobTypeCore)

	// Default the number of schedulers to match the coores
	c.NumSchedulers = runtime.NumCPU()

	// Increase our reap interval to 3 days instead of 24h.
	c.SerfConfig.ReconnectTimeout = 3 * 24 * time.Hour

	// Serf should use the WAN timing, since we are using it
	// to communicate between DC's
	c.SerfConfig.MemberlistConfig = memberlist.DefaultWANConfig()
	c.SerfConfig.MemberlistConfig.BindPort = DefaultSerfPort

	// Disable shutdown on removal
	c.RaftConfig.ShutdownOnRemove = false
	return c
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"

	pubsub "google.golang.org/api/pubsub/v1"

	log "github.com/Sirupsen/logrus"
)

type Message struct {
	Data string
	Attributes map[string]string
}

type Worker struct {
	service *pubsub.Service
	topic string
	lines chan []byte
	done    bool
	error   error
}

func (w *Worker) run() {
	for {
		flds := log.Fields{}
		log.Debugln("Getting a target")
		var line []byte
		select {
		case line = <-w.lines:
		default: // Do nothing to break
		}
		if line == nil {
			log.Debugln("No target found any more")
			w.done = true
			w.error = nil
			break
		}

		flds["line"] = string(line)
		log.WithFields(flds).Debugln("Job Start")

		err := w.process(line)
		flds["error"] = err
		if err != nil {
			w.done = true
			w.error = err
			break
		}
		log.WithFields(flds).Debugln("Job Finished")
	}
}

func (w *Worker) process(line []byte) error {
	var msg Message
	err := json.Unmarshal(line, &msg)
	if err != nil {
		flds := log.Fields{"error": err, "line": string(line)}
		log.WithFields(flds).Errorln("JSON parse error")
		return err
	}

	topic := w.service.Projects.Topics
	call := topic.Publish(w.topic, &pubsub.PublishRequest{
		Messages: []*pubsub.PubsubMessage{
			&pubsub.PubsubMessage{
				Attributes: msg.Attributes,
				Data: msg.Data,
			},
		},
	})
	res, err := call.Do()
	if err != nil {
		flds := log.Fields{"attributes": msg.Attributes, "data": msg.Data}
		log.WithFields(flds).Errorln("Publish error")
		return err
	}

	flds := log.Fields{"MessageIds": res.MessageIds}
	log.WithFields(flds).Infoln("Publish successfully")

	return nil
}
<commit_msg>:+1: Encode message.Data with Base64 encoding<commit_after>package main

import (
	"encoding/base64"
	"encoding/json"

	pubsub "google.golang.org/api/pubsub/v1"

	log "github.com/Sirupsen/logrus"
)

type Message struct {
	Data string `json:"data,omitempty"`
	Attributes map[string]string `json:"attributes,omitempty"`
}

type Worker struct {
	service *pubsub.Service
	topic string
	lines chan []byte
	done    bool
	error   error
}

func (w *Worker) run() {
	for {
		flds := log.Fields{}
		log.Debugln("Getting a target")
		var line []byte
		select {
		case line = <-w.lines:
		default: // Do nothing to break
		}
		if line == nil {
			log.Debugln("No target found any more")
			w.done = true
			w.error = nil
			break
		}

		flds["line"] = string(line)
		log.WithFields(flds).Debugln("Job Start")

		err := w.process(line)
		flds["error"] = err
		if err != nil {
			w.done = true
			w.error = err
			break
		}
		log.WithFields(flds).Debugln("Job Finished")
	}
}

func (w *Worker) process(line []byte) error {
	var msg Message
	err := json.Unmarshal(line, &msg)
	if err != nil {
		flds := log.Fields{"error": err, "line": string(line)}
		log.WithFields(flds).Errorln("JSON parse error")
		return err
	}

	topic := w.service.Projects.Topics
	call := topic.Publish(w.topic, &pubsub.PublishRequest{
		Messages: []*pubsub.PubsubMessage{
			&pubsub.PubsubMessage{
				Attributes: msg.Attributes,
				Data: base64.StdEncoding.EncodeToString([]byte(msg.Data)),
			},
		},
	})
	res, err := call.Do()
	if err != nil {
		flds := log.Fields{"attributes": msg.Attributes, "data": msg.Data, "error": err}
		log.WithFields(flds).Errorln("Publish error")
		return err
	}

	flds := log.Fields{"MessageIds": res.MessageIds}
	log.WithFields(flds).Infoln("Publish successfully")

	return nil
}
<|endoftext|>


<commit_before>// Copyright ©2011 Dan Kortschak <dan.kortschak@adelaide.edu.au>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package main

import (
	"flag"
	"fmt"
	"github.com/kortschak/BioGo/interval"
	"github.com/kortschak/BioGo/io/featio/bed"
	"math"
	"os"
)

func main() {
	var (
		region *bed.Reader
		motif  *bed.Reader
		err    error
	)

	motifName := flag.String("motif", "", "Filename for motif file.")
	regionName := flag.String("region", "", "Filename for region file.")
	verbose := flag.Bool("verbose", false, "Print details of identified motifs to stderr.")
	headerLine := flag.Bool("header", false, "Print a header line.")
	help := flag.Bool("help", false, "Print this usage message.")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s -motif <motif file> -region <region file>\n", os.Args[0])
		flag.PrintDefaults()
	}

	flag.Parse()

	if *help || *regionName == "" || *motifName == "" {
		flag.Usage()
		os.Exit(1)
	}

	// Open files
	if motif, err = bed.NewReaderName(*motifName, 3); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v.", err)
		os.Exit(0)
	} else {
		fmt.Fprintf(os.Stderr, "Reading motif features from `%s'.\n", *motifName)
	}
	defer motif.Close()

	if region, err = bed.NewReaderName(*regionName, 3); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v.", err)
		os.Exit(0)
	} else {
		fmt.Fprintf(os.Stderr, "Reading region features from `%s'.\n", *regionName)
	}
	defer region.Close()

	// Read in motif features and build interval tree to search
	intervalTree := interval.NewTree()

	for line := 1; ; line++ {
		if motifLine, err := motif.Read(); err != nil {
			break
		} else {
			if motifInterval, err := interval.New(string(motifLine.Location), motifLine.Start, motifLine.End, 0, nil); err == nil {
				intervalTree.Insert(motifInterval)
			} else {
				fmt.Fprintf(os.Stderr, "Line: %d: Feature has end < start: %v\n", line, motifLine)
			}
		}
	}

	// Read in region features and search for motifs within region
	// Calculate median motif location, sample standard deviation of locations
	// and mean distance of motif from midpoint of region for motifs contained
	// within region. Report these and n of motifs within region.
	if *headerLine {
		fmt.Println("Chromosome\tStart\tEnd\tn-hits\tMeanHitPos\tStddevHitPos\tMeanMidDistance")
	}
	for line := 1; ; line++ {
		if regionLine, err := region.Read(); err != nil {
			break
		} else {
			regionMidPoint := float64(regionLine.Start+regionLine.End) / 2
			if regionInterval, err := interval.New(string(regionLine.Location), regionLine.Start, regionLine.End, 0, regionMidPoint); err == nil {
				if *verbose {
					fmt.Fprintf(os.Stderr, "%s\t%d\t%d\n", regionLine.Location, regionLine.Start, regionLine.End)
				}
				sumOfDiffs, sumOfSquares, mean, oldmean, n := 0., 0., 0., 0., 0.
				for intersector := range intervalTree.Within(regionInterval, 0) {
					motifMidPoint := float64(intersector.Start()+intersector.End()) / 2
					if *verbose {
						fmt.Fprintf(os.Stderr, "\t%s\t%d\t%d\n", intersector.Chromosome(), intersector.Start(), intersector.End())
					}

					// The Method of Provisional Means
					n++
					mean = oldmean + (motifMidPoint-oldmean)/n
					sumOfSquares += (motifMidPoint - oldmean) * (motifMidPoint - mean)
					oldmean = mean

					sumOfDiffs += math.Abs(motifMidPoint - regionMidPoint)
				}
				fmt.Printf("%s\t%d\t%d\t%0.f\t%0.f\t%f\t%f\n",
					regionLine.Location, regionLine.Start, regionLine.End,
					n, mean, math.Sqrt(sumOfSquares)/(n-1), sumOfDiffs/n)
			} else {
				fmt.Fprintf(os.Stderr, "Line: %d: Feature has end < start: %v\n", line, regionLine)
			}
		}
	}

}
<commit_msg>Typo - alters order of operations.<commit_after>// Copyright ©2011 Dan Kortschak <dan.kortschak@adelaide.edu.au>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package main

import (
	"flag"
	"fmt"
	"github.com/kortschak/BioGo/interval"
	"github.com/kortschak/BioGo/io/featio/bed"
	"math"
	"os"
)

func main() {
	var (
		region *bed.Reader
		motif  *bed.Reader
		err    error
	)

	motifName := flag.String("motif", "", "Filename for motif file.")
	regionName := flag.String("region", "", "Filename for region file.")
	verbose := flag.Bool("verbose", false, "Print details of identified motifs to stderr.")
	headerLine := flag.Bool("header", false, "Print a header line.")
	help := flag.Bool("help", false, "Print this usage message.")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s -motif <motif file> -region <region file>\n", os.Args[0])
		flag.PrintDefaults()
	}

	flag.Parse()

	if *help || *regionName == "" || *motifName == "" {
		flag.Usage()
		os.Exit(1)
	}

	// Open files
	if motif, err = bed.NewReaderName(*motifName, 3); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v.", err)
		os.Exit(0)
	} else {
		fmt.Fprintf(os.Stderr, "Reading motif features from `%s'.\n", *motifName)
	}
	defer motif.Close()

	if region, err = bed.NewReaderName(*regionName, 3); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v.", err)
		os.Exit(0)
	} else {
		fmt.Fprintf(os.Stderr, "Reading region features from `%s'.\n", *regionName)
	}
	defer region.Close()

	// Read in motif features and build interval tree to search
	intervalTree := interval.NewTree()

	for line := 1; ; line++ {
		if motifLine, err := motif.Read(); err != nil {
			break
		} else {
			if motifInterval, err := interval.New(string(motifLine.Location), motifLine.Start, motifLine.End, 0, nil); err == nil {
				intervalTree.Insert(motifInterval)
			} else {
				fmt.Fprintf(os.Stderr, "Line: %d: Feature has end < start: %v\n", line, motifLine)
			}
		}
	}

	// Read in region features and search for motifs within region
	// Calculate median motif location, sample standard deviation of locations
	// and mean distance of motif from midpoint of region for motifs contained
	// within region. Report these and n of motifs within region.
	if *headerLine {
		fmt.Println("Chromosome\tStart\tEnd\tn-hits\tMeanHitPos\tStddevHitPos\tMeanMidDistance")
	}
	for line := 1; ; line++ {
		if regionLine, err := region.Read(); err != nil {
			break
		} else {
			regionMidPoint := float64(regionLine.Start+regionLine.End) / 2
			if regionInterval, err := interval.New(string(regionLine.Location), regionLine.Start, regionLine.End, 0, regionMidPoint); err == nil {
				if *verbose {
					fmt.Fprintf(os.Stderr, "%s\t%d\t%d\n", regionLine.Location, regionLine.Start, regionLine.End)
				}
				sumOfDiffs, sumOfSquares, mean, oldmean, n := 0., 0., 0., 0., 0.
				for intersector := range intervalTree.Within(regionInterval, 0) {
					motifMidPoint := float64(intersector.Start()+intersector.End()) / 2
					if *verbose {
						fmt.Fprintf(os.Stderr, "\t%s\t%d\t%d\n", intersector.Chromosome(), intersector.Start(), intersector.End())
					}

					// The Method of Provisional Means
					n++
					mean = oldmean + (motifMidPoint-oldmean)/n
					sumOfSquares += (motifMidPoint - oldmean) * (motifMidPoint - mean)
					oldmean = mean

					sumOfDiffs += math.Abs(motifMidPoint - regionMidPoint)
				}
				fmt.Printf("%s\t%d\t%d\t%0.f\t%0.f\t%f\t%f\n",
					regionLine.Location, regionLine.Start, regionLine.End,
					n, mean, math.Sqrt(sumOfSquares/(n-1)), sumOfDiffs/n)
			} else {
				fmt.Fprintf(os.Stderr, "Line: %d: Feature has end < start: %v\n", line, regionLine)
			}
		}
	}

}
<|endoftext|>


<commit_before>package main

import (
    "fmt"
    "net/http"
    "flag"
    "strings"
    "strconv"
    "encoding/base64"
    "github.com/mssola/user_agent"
    "github.com/dchest/uniuri"
)

var Size int

func prefix(r *http.Request) string {
    ua := new(user_agent.UserAgent)
    ua.Parse(r.UserAgent())

    os := ua.OS()
    if strings.Contains(os, "Linux") {
        return "lin"
    } else if strings.Contains(os, "Windows") {
        return "win"
    } else {
        return "srv"
    }
}

func randomName(length int) string {
    return uniuri.NewLen(length)
}

func hashName(id string) string {
//    h := md5.New()
//    h.Write([]byte(id))
//    return base64.URLEncoding.EncodeToString(h.Sum(nil))
    return base64.URLEncoding.EncodeToString([]byte(id))
}

func handler(w http.ResponseWriter, r *http.Request, size int) {
    id := r.FormValue("id")

    tmpsize := r.FormValue("size")

    if len(tmpsize) != 0 {
      s, err := strconv.Atoi(tmpsize)
      size = s
      if err != nil {
        fmt.Println(err)
        return
      }
    }

    prefix := prefix(r)
    var suffix string
    if len(id) != 0 {
      suffix = hashName(id)[0:size-3]
    } else {
      suffix = randomName(size-3)
    }
    name := strings.Join([]string{ prefix, suffix},"")
    fmt.Fprintf(w, "%s", name)
}

func main() {
    port := flag.Int("port", 8080, "Port to use")
    address := flag.String("address", "", "Address to bind")
    size := flag.Int("size", 10, "Default final hostname size")

    flag.Parse()

    socket := fmt.Sprint(*address, ":", *port)
    fmt.Printf("Bind to %s", socket)
    http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
              handler(w, r, *size)
       })
    http.ListenAndServe(socket , nil)
}
<commit_msg>Manage generated base64 < size<commit_after>package main

import (
    "fmt"
    "net/http"
    "flag"
    "strings"
    "strconv"
    "encoding/base64"
    "github.com/mssola/user_agent"
    "github.com/dchest/uniuri"
)

var Size int

func prefix(r *http.Request) string {
    ua := new(user_agent.UserAgent)
    ua.Parse(r.UserAgent())

    os := ua.OS()
    if strings.Contains(os, "Linux") {
        return "lin"
    } else if strings.Contains(os, "Windows") {
        return "win"
    } else {
        return "srv"
    }
}

func randomName(length int) string {
    return uniuri.NewLen(length)
}

func hashName(id string) string {
//    h := md5.New()
//    h.Write([]byte(id))
//    return base64.URLEncoding.EncodeToString(h.Sum(nil))
    return base64.URLEncoding.EncodeToString([]byte(id))
}

func handler(w http.ResponseWriter, r *http.Request, size int) {
    id := r.FormValue("id")

    tmpsize := r.FormValue("size")

    if len(tmpsize) != 0 {
      s, err := strconv.Atoi(tmpsize)
      size = s
      if err != nil {
        fmt.Println(err)
        return
      }
    }

    prefix := prefix(r)
    var suffix string
    if len(id) != 0 {
      ts := hashName(id)
      if len(ts) >= size-3 {
        suffix = ts[0:size-3]
      } else {
        suffix = ts
      }
    } else {
      suffix = randomName(size-3)
    }
    name := strings.Join([]string{ prefix, suffix},"")
    fmt.Fprintf(w, "%s", name)
}

func main() {
    port := flag.Int("port", 8080, "Port to use")
    address := flag.String("address", "", "Address to bind")
    size := flag.Int("size", 10, "Default final hostname size")

    flag.Parse()

    socket := fmt.Sprint(*address, ":", *port)
    fmt.Printf("Bind to %s", socket)
    http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
              handler(w, r, *size)
       })
    http.ListenAndServe(socket , nil)
}
<|endoftext|>


<commit_before>// Copyright 2012-2015, Rolf Veen and contributors.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ogdl

import (
	"errors"
	"reflect"
)

// factory[] is a map that stores type constructors.
var factory map[string]func() interface{}

// functions[] is a map for storing functions with a suitable signature so that
// they can be called from within templates.
var functions map[string]func(*Graph, []interface{}) []byte // interface{}

// FunctionAddConstructor adds a factory kind of function to the context.
func FunctionAddConstructor(s string, f func() interface{}) {
	factory[s] = f
}

// FunctionAdd adds a function to the context.
func FunctionAdd(s string, f func(*Graph, []interface{}) []byte) {
	functions[s] = f
}

// Function enables calling Go functions from templates.
//
// INPUT FORMAT
//
// g is the Function's context. g.This contains the presumed class name.
// The _type subnode of g, if present, contains the function type (a Go
// interface name or 'rfunction'
//
// p is the input path, where i points to the current position to be processed.
// The arguments of the function are 1 level higher than the function name.
// p[ix] points to the class name.
//
// Example 1
//
// !p
//   T
//   !g
//     'some text'
//
// Example 2
// !p
//   math
//   Sin
//   !g
//    !e
//      1.0
//
// Functions calls are limited to whole paths.
//
func (g *Graph) function(path *Graph, ix int, typ interface{}) (interface{}, error) {

	v := reflect.ValueOf(typ)

	// Build arguments in the form []reflect.Value
	var vargs []reflect.Value

	switch v.Kind() {

	case reflect.Func:

		// Pre-evaluate
		var args []interface{}
		for _, arg := range path.Out[ix].Out {
			args = append(args, g.evalExpression(arg))

		}

		for i, arg := range args {
			if arg == nil {
				// No untyped nil support :-(
				vargs = append(vargs, reflect.Zero(v.Type().In(i)))
			} else {
				vargs = append(vargs, reflect.ValueOf(arg))
			}
		}

		if v.Type().NumIn() != len(args) {
			s := "Invalid number of arguments in function"
			return nil, errors.New(s)
		}

		// TODO: return 0..2 values
		vv := v.Call(vargs)
		if len(vv) > 0 {
			return vv[0].Interface(), nil
		}
		return nil, nil

	case reflect.Ptr:

		fn := path.GetAt(ix)
		if fn == nil {
			return nil, errors.New("No method")
		}
		fname := fn.ThisString()

		// Check if it is a method
		me := v.MethodByName(fname)

		if !me.IsValid() {
			// Try field
			if v.Kind() == reflect.Struct {
				v = v.FieldByName(fname)
				if v.IsValid() {
					return v.Interface(), nil
				}
			}

			return nil, errors.New("No method: " + fname)
		}

		// Pre-evaluate
		var args []interface{}
		for _, arg := range path.Out[ix+1].Out {
			args = append(args, g.evalExpression(arg))

		}

		for i, arg := range args {
			if arg == nil {
				// No untyped nil support :-(
				vargs = append(vargs, reflect.Zero(me.Type().In(i)))
			} else {
				vargs = append(vargs, reflect.ValueOf(arg))
			}
		}

		if me.Type().NumIn() != len(args) {
			return nil, errors.New("Invalid number of arguments in method " + fname)
		}

		for i, arg := range args {
			v := reflect.TypeOf(arg)
			if v == nil || me.Type().In(i).String() != v.String() {
				return nil, errors.New("Invalid argument for method " + fname)
			}
		}

		// TODO: return 0..2 values
		vv := me.Call(vargs)
		if len(vv) > 0 {
			return vv[0].Interface(), nil
		}
		return nil, nil

	default:
		return nil, nil
	}

}
<commit_msg>Remove function tables<commit_after>// Copyright 2012-2015, Rolf Veen and contributors.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ogdl

import (
	"errors"
	"reflect"
)

// Function enables calling Go functions from templates.
//
// INPUT FORMAT
//
// g is the Function's context. g.This contains the presumed class name.
// The _type subnode of g, if present, contains the function type (a Go
// interface name or 'rfunction'
//
// p is the input path, where i points to the current position to be processed.
// The arguments of the function are 1 level higher than the function name.
// p[ix] points to the class name.
//
// Example 1
//
// !p
//   T
//   !g
//     'some text'
//
// Example 2
// !p
//   math
//   Sin
//   !g
//    !e
//      1.0
//
// Functions calls are limited to whole paths.
//
func (g *Graph) function(path *Graph, ix int, typ interface{}) (interface{}, error) {

	v := reflect.ValueOf(typ)

	// Build arguments in the form []reflect.Value
	var vargs []reflect.Value

	switch v.Kind() {

	case reflect.Func:

		// Pre-evaluate
		var args []interface{}
		for _, arg := range path.Out[ix].Out {
			args = append(args, g.evalExpression(arg))

		}

		for i, arg := range args {
			if arg == nil {
				// No untyped nil support :-(
				vargs = append(vargs, reflect.Zero(v.Type().In(i)))
			} else {
				vargs = append(vargs, reflect.ValueOf(arg))
			}
		}

		if v.Type().NumIn() != len(args) {
			s := "Invalid number of arguments in function"
			return nil, errors.New(s)
		}

		// TODO: return 0..2 values
		vv := v.Call(vargs)
		if len(vv) > 0 {
			return vv[0].Interface(), nil
		}
		return nil, nil

	case reflect.Ptr:

		fn := path.GetAt(ix)
		if fn == nil {
			return nil, errors.New("No method")
		}
		fname := fn.ThisString()

		// Check if it is a method
		me := v.MethodByName(fname)

		if !me.IsValid() {
			// Try field
			if v.Kind() == reflect.Struct {
				v = v.FieldByName(fname)
				if v.IsValid() {
					return v.Interface(), nil
				}
			}

			return nil, errors.New("No method: " + fname)
		}

		// Pre-evaluate
		var args []interface{}
		for _, arg := range path.Out[ix+1].Out {
			args = append(args, g.evalExpression(arg))

		}

		for i, arg := range args {
			if arg == nil {
				// No untyped nil support :-(
				vargs = append(vargs, reflect.Zero(me.Type().In(i)))
			} else {
				vargs = append(vargs, reflect.ValueOf(arg))
			}
		}

		if me.Type().NumIn() != len(args) {
			return nil, errors.New("Invalid number of arguments in method " + fname)
		}

		for i, arg := range args {
			v := reflect.TypeOf(arg)
			if v == nil || me.Type().In(i).String() != v.String() {
				return nil, errors.New("Invalid argument for method " + fname)
			}
		}

		// TODO: return 0..2 values
		vv := me.Call(vargs)
		if len(vv) > 0 {
			return vv[0].Interface(), nil
		}
		return nil, nil

	default:
		return nil, nil
	}

}
<|endoftext|>


<commit_before>package big

import (
	"fmt"
	"math"
	"math/big"
	"math/rand"
)

// A wrapper around math/big.Int which makes operations easier to express.
// The big difference is that this big.Int is immutable.  Operations on
// these big.Ints are easier to write code with, but require more
// allocations under the hood.  Totally worth it.

type Int struct {
	v *big.Int
}

// Constructors

func Int64(x int64) Int {
	return Int{big.NewInt(x)}
}

func ParseInt(s string) (Int, bool) {
	y, ok := new(big.Int).SetString(s, 10)
	return Int{y}, ok
}

// Arithmetic

func (x Int) Add(y Int) Int {
	return Int{new(big.Int).Add(x.v, y.v)}
}

func (x Int) Sub(y Int) Int {
	return Int{new(big.Int).Sub(x.v, y.v)}
}

func (x Int) Mul(y Int) Int {
	return Int{new(big.Int).Mul(x.v, y.v)}
}

func (x Int) Div(y Int) Int {
	return Int{new(big.Int).Div(x.v, y.v)}
}

func (x Int) Mod(y Int) Int {
	return Int{new(big.Int).Mod(x.v, y.v)}
}

func (x Int) Add64(y int64) Int {
	z := big.NewInt(y)
	return Int{z.Add(x.v, z)}
}

func (x Int) Sub64(y int64) Int {
	z := big.NewInt(y)
	return Int{z.Sub(x.v, z)}
}

func (x Int) Mul64(y int64) Int {
	z := big.NewInt(y)
	return Int{z.Mul(x.v, z)}
}

func (x Int) Div64(y int64) Int {
	z := big.NewInt(y)
	return Int{z.Div(x.v, z)}
}

func (x Int) Mod64(y int64) int64 {
	z := big.NewInt(y)
	return z.Mod(x.v, z).Int64()
}

func (x Int) Neg() Int {
	return Int{new(big.Int).Neg(x.v)}
}

func (x Int) Lsh(n uint) Int {
	return Int{new(big.Int).Lsh(x.v, n)}
}

func (x Int) Rsh(n uint) Int {
	return Int{new(big.Int).Rsh(x.v, n)}
}

// Info extraction

func (x Int) Int64() int64 {
	return x.v.Int64()
}

func (x Int) IsZero() bool {
	return x.v.Sign() == 0
}

func (x Int) Sign() int {
	return x.v.Sign()
}

func (x Int) Cmp(y Int) int {
	return x.v.Cmp(y.v)
}

func (x Int) Cmp64(y int64) int {
	if x.BitLen() >= 64 {
		if x.Sign() > 0 {
			return 1
		}
		if x.Cmp(minInt64) == 0 && y == math.MinInt64 {
			return 0
		}
		return -1
	}
	z := x.Int64()
	if z > y {
		return 1
	}
	if z < y {
		return -1
	}
	return 0
}

func (x Int) BitLen() int {
	return x.v.BitLen()
}

func (x Int) Bit(i int) uint {
	return x.v.Bit(i)
}

func (x Int) ProbablyPrime(n int) bool {
	return x.v.ProbablyPrime(n)
}

// Other math

func (x Int) Square() Int {
	return x.Mul(x)
}

func (x Int) Exp(k int64) Int {
	b := big.NewInt(k)
	return Int{b.Exp(x.v, b, nil)}
}

func (x Int) SqrtFloor() Int {
	if x.IsZero() {
		return x
	}
	b := uint(x.BitLen())

	// invariant lo <= sqrt(x) < hi
	lo := One.Lsh((b - 1) / 2)
	hi := lo.Lsh(1)
	for {
		m := lo.Add(hi).Rsh(1)
		if m.Cmp(lo) == 0 {
			return lo
		}
		if m.Square().Cmp(x) <= 0 {
			lo = m
		} else {
			hi = m
		}
	}
}

func (x Int) SqrtCeil() Int {
	y := x.SqrtFloor()
	if y.Square().Cmp(x) != 0 {
		y = y.Add(One)
	}
	return y
}

// Discrete math stuff

func (x Int) ExpMod(k, m Int) Int {
	return Int{new(big.Int).Exp(x.v, k.v, m.v)}
}

func (x Int) ModInv(n Int) Int {
	// TODO: check gcd(x,n)==1?
	return Int{new(big.Int).ModInverse(x.v, n.v)}
}

func (x Int) GCD(y Int) Int {
	return Int{new(big.Int).GCD(nil, nil, x.v, y.v)}
}

// For printing

func (x Int) Format(s fmt.State, ch rune) {
	x.v.Format(s, ch)
}

// Rand returns a random number in [0,x)
func (x Int) Rand(rnd *rand.Rand) Int {
	return Int{new(big.Int).Rand(rnd, x.v)}
}

// Optimized routines

// Scratch space for use by Mod64s.  Mod64s is the same
// as Mod64 except it uses the scratch space to avoid allocation.
type Scratch [2]big.Int

func (x Int) Mod64s(y int64, s *Scratch) int64 {
	return s[0].Mod(x.v, s[1].SetInt64(y)).Int64()
}

// helpful constants
var Zero = Int64(0)
var One = Int64(1)
var Two = Int64(2)
var Ten = Int64(10)

var minInt64 = Int64(math.MinInt64)
<commit_msg>Keep Mod64s from allocating by using DivMod instead of Mod.<commit_after>package big

import (
	"fmt"
	"math"
	"math/big"
	"math/rand"
)

// A wrapper around math/big.Int which makes operations easier to express.
// The big difference is that this big.Int is immutable.  Operations on
// these big.Ints are easier to write code with, but require more
// allocations under the hood.  Totally worth it.

type Int struct {
	v *big.Int
}

// Constructors

func Int64(x int64) Int {
	return Int{big.NewInt(x)}
}

func ParseInt(s string) (Int, bool) {
	y, ok := new(big.Int).SetString(s, 10)
	return Int{y}, ok
}

// Arithmetic

func (x Int) Add(y Int) Int {
	return Int{new(big.Int).Add(x.v, y.v)}
}

func (x Int) Sub(y Int) Int {
	return Int{new(big.Int).Sub(x.v, y.v)}
}

func (x Int) Mul(y Int) Int {
	return Int{new(big.Int).Mul(x.v, y.v)}
}

func (x Int) Div(y Int) Int {
	return Int{new(big.Int).Div(x.v, y.v)}
}

func (x Int) Mod(y Int) Int {
	return Int{new(big.Int).Mod(x.v, y.v)}
}

func (x Int) Add64(y int64) Int {
	z := big.NewInt(y)
	return Int{z.Add(x.v, z)}
}

func (x Int) Sub64(y int64) Int {
	z := big.NewInt(y)
	return Int{z.Sub(x.v, z)}
}

func (x Int) Mul64(y int64) Int {
	z := big.NewInt(y)
	return Int{z.Mul(x.v, z)}
}

func (x Int) Div64(y int64) Int {
	z := big.NewInt(y)
	return Int{z.Div(x.v, z)}
}

func (x Int) Mod64(y int64) int64 {
	z := big.NewInt(y)
	return z.Mod(x.v, z).Int64()
}

func (x Int) Neg() Int {
	return Int{new(big.Int).Neg(x.v)}
}

func (x Int) Lsh(n uint) Int {
	return Int{new(big.Int).Lsh(x.v, n)}
}

func (x Int) Rsh(n uint) Int {
	return Int{new(big.Int).Rsh(x.v, n)}
}

// Info extraction

func (x Int) Int64() int64 {
	return x.v.Int64()
}

func (x Int) IsZero() bool {
	return x.v.Sign() == 0
}

func (x Int) Sign() int {
	return x.v.Sign()
}

func (x Int) Cmp(y Int) int {
	return x.v.Cmp(y.v)
}

func (x Int) Cmp64(y int64) int {
	if x.BitLen() >= 64 {
		if x.Sign() > 0 {
			return 1
		}
		if x.Cmp(minInt64) == 0 && y == math.MinInt64 {
			return 0
		}
		return -1
	}
	z := x.Int64()
	if z > y {
		return 1
	}
	if z < y {
		return -1
	}
	return 0
}

func (x Int) BitLen() int {
	return x.v.BitLen()
}

func (x Int) Bit(i int) uint {
	return x.v.Bit(i)
}

func (x Int) ProbablyPrime(n int) bool {
	return x.v.ProbablyPrime(n)
}

// Other math

func (x Int) Square() Int {
	return x.Mul(x)
}

func (x Int) Exp(k int64) Int {
	b := big.NewInt(k)
	return Int{b.Exp(x.v, b, nil)}
}

func (x Int) SqrtFloor() Int {
	if x.IsZero() {
		return x
	}
	b := uint(x.BitLen())

	// invariant lo <= sqrt(x) < hi
	lo := One.Lsh((b - 1) / 2)
	hi := lo.Lsh(1)
	for {
		m := lo.Add(hi).Rsh(1)
		if m.Cmp(lo) == 0 {
			return lo
		}
		if m.Square().Cmp(x) <= 0 {
			lo = m
		} else {
			hi = m
		}
	}
}

func (x Int) SqrtCeil() Int {
	y := x.SqrtFloor()
	if y.Square().Cmp(x) != 0 {
		y = y.Add(One)
	}
	return y
}

// Discrete math stuff

func (x Int) ExpMod(k, m Int) Int {
	return Int{new(big.Int).Exp(x.v, k.v, m.v)}
}

func (x Int) ModInv(n Int) Int {
	// TODO: check gcd(x,n)==1?
	return Int{new(big.Int).ModInverse(x.v, n.v)}
}

func (x Int) GCD(y Int) Int {
	return Int{new(big.Int).GCD(nil, nil, x.v, y.v)}
}

// For printing

func (x Int) Format(s fmt.State, ch rune) {
	x.v.Format(s, ch)
}

// Rand returns a random number in [0,x)
func (x Int) Rand(rnd *rand.Rand) Int {
	return Int{new(big.Int).Rand(rnd, x.v)}
}

// Optimized routines

// Scratch space for use by Mod64s.  Mod64s is the same
// as Mod64 except it uses the scratch space to avoid allocation.
type Scratch [3]big.Int

func (x Int) Mod64s(y int64, s *Scratch) int64 {
	// Note: use DivMod here instead of Mod so we can reuse
	// storage for the dividend.  Mod allocates storage for
	// the (thrown away) dividend on each call.
	s[0].DivMod(x.v, s[1].SetInt64(y), &s[2])
	return s[2].Int64()
}

// helpful constants
var Zero = Int64(0)
var One = Int64(1)
var Two = Int64(2)
var Ten = Int64(10)

var minInt64 = Int64(math.MinInt64)
<|endoftext|>


<commit_before>package wallet

import (
	"bytes"
	"os"
	"path/filepath"
	"testing"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/modules/miner"
	"github.com/NebulousLabs/Sia/types"
	"github.com/NebulousLabs/fastrand"
)

// postEncryptionTesting runs a series of checks on the wallet after it has
// been encrypted, to make sure that locking, unlocking, and spending after
// unlocking are all happening in the correct order and returning the correct
// errors.
func postEncryptionTesting(m modules.TestMiner, w *Wallet, masterKey crypto.TwofishKey) {
	if !w.Encrypted() {
		panic("wallet is not encrypted when starting postEncryptionTesting")
	}
	if w.Unlocked() {
		panic("wallet is unlocked when starting postEncryptionTesting")
	}
	if len(w.seeds) != 0 {
		panic("wallet has seeds in it when startin postEncryptionTesting")
	}

	// Try unlocking and using the wallet.
	err := w.Unlock(masterKey)
	if err != nil {
		panic(err)
	}
	err = w.Unlock(masterKey)
	if err != errAlreadyUnlocked {
		panic(err)
	}
	// Mine enough coins so that a balance appears (and some buffer for the
	// send later).
	for i := types.BlockHeight(0); i <= types.MaturityDelay+1; i++ {
		_, err := m.AddBlock()
		if err != nil {
			panic(err)
		}
	}
	siacoinBal, _, _ := w.ConfirmedBalance()
	if siacoinBal.IsZero() {
		panic("wallet balance reported as 0 after maturing some mined blocks")
	}
	err = w.Unlock(masterKey)
	if err != errAlreadyUnlocked {
		panic(err)
	}

	// Lock, unlock, and trying using the wallet some more.
	err = w.Lock()
	if err != nil {
		panic(err)
	}
	err = w.Lock()
	if err != modules.ErrLockedWallet {
		panic(err)
	}
	err = w.Unlock(crypto.TwofishKey{})
	if err != modules.ErrBadEncryptionKey {
		panic(err)
	}
	err = w.Unlock(masterKey)
	if err != nil {
		panic(err)
	}
	// Verify that the secret keys have been restored by sending coins to the
	// void. Send more coins than are received by mining a block.
	_, err = w.SendSiacoins(types.CalculateCoinbase(0), types.UnlockHash{})
	if err != nil {
		panic(err)
	}
	_, err = m.AddBlock()
	if err != nil {
		panic(err)
	}
	siacoinBal2, _, _ := w.ConfirmedBalance()
	if siacoinBal2.Cmp(siacoinBal) >= 0 {
		panic("balance did not increase")
	}
}

// TestIntegrationPreEncryption checks that the wallet operates as expected
// prior to encryption.
func TestIntegrationPreEncryption(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	wt, err := createBlankWalletTester(t.Name())
	if err != nil {
		t.Fatal(err)
	}

	// Check that the wallet knows it's not encrypted.
	if wt.wallet.Encrypted() {
		t.Error("wallet is reporting that it has been encrypted")
	}
	err = wt.wallet.Lock()
	if err != modules.ErrLockedWallet {
		t.Fatal(err)
	}
	err = wt.wallet.Unlock(crypto.TwofishKey{})
	if err != errUnencryptedWallet {
		t.Fatal(err)
	}
	wt.closeWt()

	// Create a second wallet using the same directory - make sure that if any
	// files have been created, the wallet is still being treated as new.
	w1, err := New(wt.cs, wt.tpool, filepath.Join(wt.persistDir, modules.WalletDir))
	if err != nil {
		t.Fatal(err)
	}
	if w1.Encrypted() {
		t.Error("wallet is reporting that it has been encrypted when no such action has occurred")
	}
	if w1.Unlocked() {
		t.Error("new wallet is not being treated as locked")
	}
	w1.Close()
}

// TestIntegrationUserSuppliedEncryption probes the encryption process when the
// user manually supplies an encryption key.
func TestIntegrationUserSuppliedEncryption(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	// Create and wallet and user-specified key, then encrypt the wallet and
	// run post-encryption tests on it.
	wt, err := createBlankWalletTester(t.Name())
	if err != nil {
		t.Fatal(err)
	}
	defer wt.closeWt()
	var masterKey crypto.TwofishKey
	fastrand.Read(masterKey[:])
	_, err = wt.wallet.Encrypt(masterKey)
	if err != nil {
		t.Error(err)
	}
	postEncryptionTesting(wt.miner, wt.wallet, masterKey)
}

// TestIntegrationBlankEncryption probes the encryption process when the user
// supplies a blank encryption key during the encryption process.
func TestIntegrationBlankEncryption(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	// Create the wallet.
	wt, err := createBlankWalletTester(t.Name())
	if err != nil {
		t.Fatal(err)
	}
	defer wt.closeWt()
	// Encrypt the wallet using a blank key.
	seed, err := wt.wallet.Encrypt(crypto.TwofishKey{})
	if err != nil {
		t.Error(err)
	}

	// Try unlocking the wallet using a blank key.
	err = wt.wallet.Unlock(crypto.TwofishKey{})
	if err != modules.ErrBadEncryptionKey {
		t.Fatal(err)
	}
	// Try unlocking the wallet using the correct key.
	err = wt.wallet.Unlock(crypto.TwofishKey(crypto.HashObject(seed)))
	if err != nil {
		t.Fatal(err)
	}
	err = wt.wallet.Lock()
	if err != nil {
		t.Fatal(err)
	}
	postEncryptionTesting(wt.miner, wt.wallet, crypto.TwofishKey(crypto.HashObject(seed)))
}

// TestLock checks that lock correctly wipes keys when locking the wallet,
// while still being able to track the balance of the wallet.
func TestLock(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	wt, err := createWalletTester(t.Name())
	if err != nil {
		t.Fatal(err)
	}
	defer wt.closeWt()

	// Grab a block for work - miner will not supply blocks after the wallet
	// has been locked, and the test needs to mine a block after locking the
	// wallet to verify  that the balance reporting of a locked wallet is
	// correct.
	block, target, err := wt.miner.BlockForWork()
	if err != nil {
		t.Fatal(err)
	}

	// Lock the wallet.
	siacoinBalance, _, _ := wt.wallet.ConfirmedBalance()
	err = wt.wallet.Lock()
	if err != nil {
		t.Error(err)
	}
	// Compare to the original balance.
	siacoinBalance2, _, _ := wt.wallet.ConfirmedBalance()
	if !siacoinBalance2.Equals(siacoinBalance) {
		t.Error("siacoin balance reporting changed upon closing the wallet")
	}
	// Check that the keys and seeds were wiped.
	wipedKey := make([]byte, crypto.SecretKeySize)
	for _, key := range wt.wallet.keys {
		for i := range key.SecretKeys {
			if !bytes.Equal(wipedKey, key.SecretKeys[i][:]) {
				t.Error("Key was not wiped after closing the wallet")
			}
		}
	}
	if len(wt.wallet.seeds) != 0 {
		t.Error("seeds not wiped from wallet")
	}
	if !bytes.Equal(wipedKey[:crypto.EntropySize], wt.wallet.primarySeed[:]) {
		t.Error("primary seed not wiped from memory")
	}

	// Solve the block generated earlier and add it to the consensus set, this
	// should boost the balance of the wallet.
	solvedBlock, _ := wt.miner.SolveBlock(block, target)
	err = wt.cs.AcceptBlock(solvedBlock)
	if err != nil {
		t.Fatal(err)
	}
	siacoinBalance3, _, _ := wt.wallet.ConfirmedBalance()
	if siacoinBalance3.Cmp(siacoinBalance2) <= 0 {
		t.Error("balance should increase after a block was mined")
	}
}

// TestInitFromSeed tests creating a wallet from a preexisting seed.
func TestInitFromSeed(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	// create a wallet with some money
	wt, err := createWalletTester("TestInitFromSeed0")
	if err != nil {
		t.Fatal(err)
	}
	defer wt.closeWt()
	seed, _, err := wt.wallet.PrimarySeed()
	if err != nil {
		t.Fatal(err)
	}
	origBal, _, _ := wt.wallet.ConfirmedBalance()

	// create a blank wallet
	dir := filepath.Join(build.TempDir(modules.WalletDir, "TestInitFromSeed1"), modules.WalletDir)
	w, err := New(wt.cs, wt.tpool, dir)
	if err != nil {
		t.Fatal(err)
	}
	err = w.InitFromSeed(crypto.TwofishKey{}, seed)
	if err != nil {
		t.Fatal(err)
	}
	err = w.Unlock(crypto.TwofishKey(crypto.HashObject(seed)))
	if err != nil {
		t.Fatal(err)
	}
	// starting balance should match the original wallet
	newBal, _, _ := w.ConfirmedBalance()
	if newBal.Cmp(origBal) != 0 {
		t.Log(w.UnconfirmedBalance())
		t.Fatalf("wallet should have correct balance after loading seed: wanted %v, got %v", origBal, newBal)
	}
}

// TestReset tests that Reset resets a wallet correctly.
func TestReset(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	wt, err := createBlankWalletTester(t.Name())
	if err != nil {
		t.Fatal(err)
	}
	defer wt.closeWt()

	var originalKey crypto.TwofishKey
	fastrand.Read(originalKey[:])
	_, err = wt.wallet.Encrypt(originalKey)
	if err != nil {
		t.Fatal(err)
	}
	postEncryptionTesting(wt.miner, wt.wallet, originalKey)

	err = wt.wallet.Reset()
	if err != nil {
		t.Fatal(err)
	}

	// reinitialize the miner so it mines into the new seed
	minerData := filepath.Join(wt.persistDir, modules.MinerDir)
	err = os.RemoveAll(minerData)
	if err != nil {
		t.Fatal(err)
	}
	newminer, err := miner.New(wt.cs, wt.tpool, wt.wallet, filepath.Join(wt.persistDir, modules.MinerDir))
	if err != nil {
		t.Fatal(err)
	}
	wt.miner = newminer

	var newKey crypto.TwofishKey
	fastrand.Read(newKey[:])
	_, err = wt.wallet.Encrypt(newKey)
	if err != nil {
		t.Fatal(err)
	}
	postEncryptionTesting(wt.miner, wt.wallet, newKey)
}
<commit_msg>close previous miner before reinitializing a new one<commit_after>package wallet

import (
	"bytes"
	"os"
	"path/filepath"
	"testing"

	"github.com/NebulousLabs/Sia/build"
	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/modules/miner"
	"github.com/NebulousLabs/Sia/types"
	"github.com/NebulousLabs/fastrand"
)

// postEncryptionTesting runs a series of checks on the wallet after it has
// been encrypted, to make sure that locking, unlocking, and spending after
// unlocking are all happening in the correct order and returning the correct
// errors.
func postEncryptionTesting(m modules.TestMiner, w *Wallet, masterKey crypto.TwofishKey) {
	if !w.Encrypted() {
		panic("wallet is not encrypted when starting postEncryptionTesting")
	}
	if w.Unlocked() {
		panic("wallet is unlocked when starting postEncryptionTesting")
	}
	if len(w.seeds) != 0 {
		panic("wallet has seeds in it when startin postEncryptionTesting")
	}

	// Try unlocking and using the wallet.
	err := w.Unlock(masterKey)
	if err != nil {
		panic(err)
	}
	err = w.Unlock(masterKey)
	if err != errAlreadyUnlocked {
		panic(err)
	}
	// Mine enough coins so that a balance appears (and some buffer for the
	// send later).
	for i := types.BlockHeight(0); i <= types.MaturityDelay+1; i++ {
		_, err := m.AddBlock()
		if err != nil {
			panic(err)
		}
	}
	siacoinBal, _, _ := w.ConfirmedBalance()
	if siacoinBal.IsZero() {
		panic("wallet balance reported as 0 after maturing some mined blocks")
	}
	err = w.Unlock(masterKey)
	if err != errAlreadyUnlocked {
		panic(err)
	}

	// Lock, unlock, and trying using the wallet some more.
	err = w.Lock()
	if err != nil {
		panic(err)
	}
	err = w.Lock()
	if err != modules.ErrLockedWallet {
		panic(err)
	}
	err = w.Unlock(crypto.TwofishKey{})
	if err != modules.ErrBadEncryptionKey {
		panic(err)
	}
	err = w.Unlock(masterKey)
	if err != nil {
		panic(err)
	}
	// Verify that the secret keys have been restored by sending coins to the
	// void. Send more coins than are received by mining a block.
	_, err = w.SendSiacoins(types.CalculateCoinbase(0), types.UnlockHash{})
	if err != nil {
		panic(err)
	}
	_, err = m.AddBlock()
	if err != nil {
		panic(err)
	}
	siacoinBal2, _, _ := w.ConfirmedBalance()
	if siacoinBal2.Cmp(siacoinBal) >= 0 {
		panic("balance did not increase")
	}
}

// TestIntegrationPreEncryption checks that the wallet operates as expected
// prior to encryption.
func TestIntegrationPreEncryption(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	wt, err := createBlankWalletTester(t.Name())
	if err != nil {
		t.Fatal(err)
	}

	// Check that the wallet knows it's not encrypted.
	if wt.wallet.Encrypted() {
		t.Error("wallet is reporting that it has been encrypted")
	}
	err = wt.wallet.Lock()
	if err != modules.ErrLockedWallet {
		t.Fatal(err)
	}
	err = wt.wallet.Unlock(crypto.TwofishKey{})
	if err != errUnencryptedWallet {
		t.Fatal(err)
	}
	wt.closeWt()

	// Create a second wallet using the same directory - make sure that if any
	// files have been created, the wallet is still being treated as new.
	w1, err := New(wt.cs, wt.tpool, filepath.Join(wt.persistDir, modules.WalletDir))
	if err != nil {
		t.Fatal(err)
	}
	if w1.Encrypted() {
		t.Error("wallet is reporting that it has been encrypted when no such action has occurred")
	}
	if w1.Unlocked() {
		t.Error("new wallet is not being treated as locked")
	}
	w1.Close()
}

// TestIntegrationUserSuppliedEncryption probes the encryption process when the
// user manually supplies an encryption key.
func TestIntegrationUserSuppliedEncryption(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	// Create and wallet and user-specified key, then encrypt the wallet and
	// run post-encryption tests on it.
	wt, err := createBlankWalletTester(t.Name())
	if err != nil {
		t.Fatal(err)
	}
	defer wt.closeWt()
	var masterKey crypto.TwofishKey
	fastrand.Read(masterKey[:])
	_, err = wt.wallet.Encrypt(masterKey)
	if err != nil {
		t.Error(err)
	}
	postEncryptionTesting(wt.miner, wt.wallet, masterKey)
}

// TestIntegrationBlankEncryption probes the encryption process when the user
// supplies a blank encryption key during the encryption process.
func TestIntegrationBlankEncryption(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	// Create the wallet.
	wt, err := createBlankWalletTester(t.Name())
	if err != nil {
		t.Fatal(err)
	}
	defer wt.closeWt()
	// Encrypt the wallet using a blank key.
	seed, err := wt.wallet.Encrypt(crypto.TwofishKey{})
	if err != nil {
		t.Error(err)
	}

	// Try unlocking the wallet using a blank key.
	err = wt.wallet.Unlock(crypto.TwofishKey{})
	if err != modules.ErrBadEncryptionKey {
		t.Fatal(err)
	}
	// Try unlocking the wallet using the correct key.
	err = wt.wallet.Unlock(crypto.TwofishKey(crypto.HashObject(seed)))
	if err != nil {
		t.Fatal(err)
	}
	err = wt.wallet.Lock()
	if err != nil {
		t.Fatal(err)
	}
	postEncryptionTesting(wt.miner, wt.wallet, crypto.TwofishKey(crypto.HashObject(seed)))
}

// TestLock checks that lock correctly wipes keys when locking the wallet,
// while still being able to track the balance of the wallet.
func TestLock(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	wt, err := createWalletTester(t.Name())
	if err != nil {
		t.Fatal(err)
	}
	defer wt.closeWt()

	// Grab a block for work - miner will not supply blocks after the wallet
	// has been locked, and the test needs to mine a block after locking the
	// wallet to verify  that the balance reporting of a locked wallet is
	// correct.
	block, target, err := wt.miner.BlockForWork()
	if err != nil {
		t.Fatal(err)
	}

	// Lock the wallet.
	siacoinBalance, _, _ := wt.wallet.ConfirmedBalance()
	err = wt.wallet.Lock()
	if err != nil {
		t.Error(err)
	}
	// Compare to the original balance.
	siacoinBalance2, _, _ := wt.wallet.ConfirmedBalance()
	if !siacoinBalance2.Equals(siacoinBalance) {
		t.Error("siacoin balance reporting changed upon closing the wallet")
	}
	// Check that the keys and seeds were wiped.
	wipedKey := make([]byte, crypto.SecretKeySize)
	for _, key := range wt.wallet.keys {
		for i := range key.SecretKeys {
			if !bytes.Equal(wipedKey, key.SecretKeys[i][:]) {
				t.Error("Key was not wiped after closing the wallet")
			}
		}
	}
	if len(wt.wallet.seeds) != 0 {
		t.Error("seeds not wiped from wallet")
	}
	if !bytes.Equal(wipedKey[:crypto.EntropySize], wt.wallet.primarySeed[:]) {
		t.Error("primary seed not wiped from memory")
	}

	// Solve the block generated earlier and add it to the consensus set, this
	// should boost the balance of the wallet.
	solvedBlock, _ := wt.miner.SolveBlock(block, target)
	err = wt.cs.AcceptBlock(solvedBlock)
	if err != nil {
		t.Fatal(err)
	}
	siacoinBalance3, _, _ := wt.wallet.ConfirmedBalance()
	if siacoinBalance3.Cmp(siacoinBalance2) <= 0 {
		t.Error("balance should increase after a block was mined")
	}
}

// TestInitFromSeed tests creating a wallet from a preexisting seed.
func TestInitFromSeed(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	// create a wallet with some money
	wt, err := createWalletTester("TestInitFromSeed0")
	if err != nil {
		t.Fatal(err)
	}
	defer wt.closeWt()
	seed, _, err := wt.wallet.PrimarySeed()
	if err != nil {
		t.Fatal(err)
	}
	origBal, _, _ := wt.wallet.ConfirmedBalance()

	// create a blank wallet
	dir := filepath.Join(build.TempDir(modules.WalletDir, "TestInitFromSeed1"), modules.WalletDir)
	w, err := New(wt.cs, wt.tpool, dir)
	if err != nil {
		t.Fatal(err)
	}
	err = w.InitFromSeed(crypto.TwofishKey{}, seed)
	if err != nil {
		t.Fatal(err)
	}
	err = w.Unlock(crypto.TwofishKey(crypto.HashObject(seed)))
	if err != nil {
		t.Fatal(err)
	}
	// starting balance should match the original wallet
	newBal, _, _ := w.ConfirmedBalance()
	if newBal.Cmp(origBal) != 0 {
		t.Log(w.UnconfirmedBalance())
		t.Fatalf("wallet should have correct balance after loading seed: wanted %v, got %v", origBal, newBal)
	}
}

// TestReset tests that Reset resets a wallet correctly.
func TestReset(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	wt, err := createBlankWalletTester(t.Name())
	if err != nil {
		t.Fatal(err)
	}
	defer wt.closeWt()

	var originalKey crypto.TwofishKey
	fastrand.Read(originalKey[:])
	_, err = wt.wallet.Encrypt(originalKey)
	if err != nil {
		t.Fatal(err)
	}
	postEncryptionTesting(wt.miner, wt.wallet, originalKey)

	err = wt.wallet.Reset()
	if err != nil {
		t.Fatal(err)
	}

	// reinitialize the miner so it mines into the new seed
	err = wt.miner.Close()
	if err != nil {
		t.Fatal(err)
	}
	minerData := filepath.Join(wt.persistDir, modules.MinerDir)
	err = os.RemoveAll(minerData)
	if err != nil {
		t.Fatal(err)
	}
	newminer, err := miner.New(wt.cs, wt.tpool, wt.wallet, filepath.Join(wt.persistDir, modules.MinerDir))
	if err != nil {
		t.Fatal(err)
	}
	wt.miner = newminer

	var newKey crypto.TwofishKey
	fastrand.Read(newKey[:])
	_, err = wt.wallet.Encrypt(newKey)
	if err != nil {
		t.Fatal(err)
	}
	postEncryptionTesting(wt.miner, wt.wallet, newKey)
}
<|endoftext|>


<commit_before>package dbus

import (
	"time"
)

// SubscriptionSet returns a subscription set which is like conn.Subscribe but
// can filter to only return events for a set of units.
type SubscriptionSet struct {
	*set
	conn *Conn
}

func (s *SubscriptionSet) filter(unit string) bool {
	return !s.Contains(unit)
}

// Subscribe starts listening for dbus events for all of the units in the set.
// Returns channels identical to conn.SubscribeUnits.
func (s *SubscriptionSet) Subscribe() (<-chan map[string]*UnitStatus, <-chan error) {
	// TODO: Make fully evented by using systemd 209 with properties changed values
	return s.conn.SubscribeUnitsCustom(time.Second, 0,
		mismatchUnitStatus,
		func(unit string) bool { return s.filter(unit) },
	)
}

// NewSubscriptionSet returns a new subscription set.
func (conn *Conn) NewSubscriptionSet() (*SubscriptionSet) {
	return &SubscriptionSet{newSet(), conn}
}

// mismatchUnitStatus returns true if the provided UnitStatus objects
// are not equivalent. false is returned if the objects are equivalent.
// Only the Name, Description and state-related fields are used in
// the comparison.
func mismatchUnitStatus(u1, u2 *UnitStatus) bool {
	return u1.Name != u2.Name ||
		u1.Description != u2.Description ||
		u1.LoadState != u2.LoadState ||
		u1.ActiveState != u2.ActiveState ||
		u1.SubState != u2.SubState
}
<commit_msg>subscription: gofmt<commit_after>package dbus

import (
	"time"
)

// SubscriptionSet returns a subscription set which is like conn.Subscribe but
// can filter to only return events for a set of units.
type SubscriptionSet struct {
	*set
	conn *Conn
}

func (s *SubscriptionSet) filter(unit string) bool {
	return !s.Contains(unit)
}

// Subscribe starts listening for dbus events for all of the units in the set.
// Returns channels identical to conn.SubscribeUnits.
func (s *SubscriptionSet) Subscribe() (<-chan map[string]*UnitStatus, <-chan error) {
	// TODO: Make fully evented by using systemd 209 with properties changed values
	return s.conn.SubscribeUnitsCustom(time.Second, 0,
		mismatchUnitStatus,
		func(unit string) bool { return s.filter(unit) },
	)
}

// NewSubscriptionSet returns a new subscription set.
func (conn *Conn) NewSubscriptionSet() *SubscriptionSet {
	return &SubscriptionSet{newSet(), conn}
}

// mismatchUnitStatus returns true if the provided UnitStatus objects
// are not equivalent. false is returned if the objects are equivalent.
// Only the Name, Description and state-related fields are used in
// the comparison.
func mismatchUnitStatus(u1, u2 *UnitStatus) bool {
	return u1.Name != u2.Name ||
		u1.Description != u2.Description ||
		u1.LoadState != u2.LoadState ||
		u1.ActiveState != u2.ActiveState ||
		u1.SubState != u2.SubState
}
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// 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.

// forkfd listens for tcp and spawns a waterfall server receiving on the opened connection.
package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"syscall"
)

const (
	// Path to logwrapper binary on device. This pipes stdout/stderr to logcat.
	logwrapperBin = "/system/bin/logwrapper"
)

var (
	binaryToLaunch = flag.String("binary_to_launch", "", "Path to the binary to launch.")
	fork = flag.Bool("fork", false, "Weather to listen on -addr and dup the conn fd to pass it to the forked process.")
	addr           = flag.String("addr", "localhost:8080", "Address to listen on for connections to pass to forked process.")
)

// launchAndDupConn launches a binary dupping the conn fd and making it accessible to the forked process.
func launchAndDupConn(binaryPath string, rc syscall.RawConn, passthroughArgs []string) error {
	var cmd *exec.Cmd
	var cmdErr error

	args := append([]string{binaryPath, "-addr", fmt.Sprintf("mux:fd:%d", 3)}, passthroughArgs...)
	err := rc.Control(func(fd uintptr) {
		log.Printf("Forking %s with fd: %d\n", binaryPath, fd)
		f := os.NewFile(fd, "")
		cmd = exec.Command(logwrapperBin, args...)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.ExtraFiles = append(cmd.ExtraFiles, f)
		cmdErr = cmd.Start()
	})

	if err != nil {
		return err
	}

	if cmdErr != nil {
		return err
	}

	return cmd.Wait()
}

func main() {
	flag.Parse()

	if *addr == "" {
		log.Fatalf("Need to provide -addr.")
	}

	if *binaryToLaunch == "" {
		log.Fatalf("Need to provide -binary_to_launch.")
	}

	if *fork {
		log.Printf("Launching %s\n", *binaryToLaunch)
		cmd := exec.Command(os.Args[0], "-addr", *addr, "-binary_to_launch", *binaryToLaunch)
		if err := cmd.Start(); err != nil {
			log.Fatal(err)
		}
		return
	}

	lis, err := net.Listen("tcp", *addr)
	if err != nil {
		log.Fatalf("Failed to listen: %v", err)
	}
	defer lis.Close()

	for {
		log.Printf("Listening on %s\n", *addr)
		conn, err := lis.Accept()
		if err != nil {
			log.Fatalf("Failed to accept: %v", err)
		}
		defer conn.Close()

		log.Println("Accepted connection. Starting process ...")

		rawConn, err := conn.(*net.TCPConn).SyscallConn()
		if err != nil {
			log.Fatalf("Failed to get raw conn: %v", err)
		}

		if err := launchAndDupConn(*binaryToLaunch, rawConn, flag.Args()); err != nil {
			// just wait for the next connection on error
			log.Printf("Error running process: %v\n", err)
		}
	}
}
<commit_msg>Ran gofmt<commit_after>// Copyright 2019 Google LLC
//
// 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.

// forkfd listens for tcp and spawns a waterfall server receiving on the opened connection.
package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"syscall"
)

const (
	// Path to logwrapper binary on device. This pipes stdout/stderr to logcat.
	logwrapperBin = "/system/bin/logwrapper"
)

var (
	binaryToLaunch = flag.String("binary_to_launch", "", "Path to the binary to launch.")
	fork           = flag.Bool("fork", false, "Weather to listen on -addr and dup the conn fd to pass it to the forked process.")
	addr           = flag.String("addr", "localhost:8080", "Address to listen on for connections to pass to forked process.")
)

// launchAndDupConn launches a binary dupping the conn fd and making it accessible to the forked process.
func launchAndDupConn(binaryPath string, rc syscall.RawConn, passthroughArgs []string) error {
	var cmd *exec.Cmd
	var cmdErr error

	args := append([]string{binaryPath, "-addr", fmt.Sprintf("mux:fd:%d", 3)}, passthroughArgs...)
	err := rc.Control(func(fd uintptr) {
		log.Printf("Forking %s with fd: %d\n", binaryPath, fd)
		f := os.NewFile(fd, "")
		cmd = exec.Command(logwrapperBin, args...)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.ExtraFiles = append(cmd.ExtraFiles, f)
		cmdErr = cmd.Start()
	})

	if err != nil {
		return err
	}

	if cmdErr != nil {
		return err
	}

	return cmd.Wait()
}

func main() {
	flag.Parse()

	if *addr == "" {
		log.Fatalf("Need to provide -addr.")
	}

	if *binaryToLaunch == "" {
		log.Fatalf("Need to provide -binary_to_launch.")
	}

	if *fork {
		log.Printf("Launching %s\n", *binaryToLaunch)
		cmd := exec.Command(os.Args[0], "-addr", *addr, "-binary_to_launch", *binaryToLaunch)
		if err := cmd.Start(); err != nil {
			log.Fatal(err)
		}
		return
	}

	lis, err := net.Listen("tcp", *addr)
	if err != nil {
		log.Fatalf("Failed to listen: %v", err)
	}
	defer lis.Close()

	for {
		log.Printf("Listening on %s\n", *addr)
		conn, err := lis.Accept()
		if err != nil {
			log.Fatalf("Failed to accept: %v", err)
		}
		defer conn.Close()

		log.Println("Accepted connection. Starting process ...")

		rawConn, err := conn.(*net.TCPConn).SyscallConn()
		if err != nil {
			log.Fatalf("Failed to get raw conn: %v", err)
		}

		if err := launchAndDupConn(*binaryToLaunch, rawConn, flag.Args()); err != nil {
			// just wait for the next connection on error
			log.Printf("Error running process: %v\n", err)
		}
	}
}
<|endoftext|>