<commit_before>package shell

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/zetamatta/go-mbcs"

	"github.com/zetamatta/nyagos/dos"
)

func readEnv(scan *bufio.Scanner, verbose io.Writer) (int, error) {
	errorlevel := -1
	for scan.Scan() {
		line, err := mbcs.ConsoleCpToUtf8(scan.Bytes())
		if err != nil {
			continue
		}
		line = strings.TrimSpace(line)
		eqlPos := strings.Index(line, "=")
		if eqlPos > 0 {
			left := line[:eqlPos]
			right := line[eqlPos+1:]
			if left == "ERRORLEVEL_" {
				value, err := strconv.ParseUint(right, 10, 32)
				if err != nil {
					fmt.Fprintf(verbose, "Could not read ERRORLEVEL(%s)\n", right)
				} else {
					errorlevel = int(value)
				}
			} else {
				orig := os.Getenv(left)
				if verbose != nil {
					fmt.Fprintf(verbose, "%s=%s\n", left, right)
				}
				if orig != right {
					// fmt.Fprintf(os.Stderr, "%s:=%s\n", left, right)
					os.Setenv(left, right)
				}
			}
		}
	}
	return errorlevel, scan.Err()
}

func readPwd(scan *bufio.Scanner, verbose io.Writer) error {
	if !scan.Scan() {
		return errors.New("Could not load the new current directory")
	}
	if err := scan.Err(); err != nil {
		return err
	}
	line, err := mbcs.ConsoleCpToUtf8(scan.Bytes())
	if err != nil {
		return err
	}
	line = strings.TrimSpace(line)
	if verbose != nil {
		fmt.Fprintf(verbose, "cd \"%s\"\n", line)
	}
	os.Chdir(line)
	return nil
}

// loadTmpFile - read update the current-directory and environment-variables from tmp-file.
func loadTmpFile(fname string, verbose io.Writer) (int, error) {
	fp, err := os.Open(fname)
	if err != nil {
		return -1, err
	}
	defer fp.Close()

	scan := bufio.NewScanner(fp)
	if err := readPwd(scan, verbose); err != nil {
		return -1, err
	}
	return readEnv(scan, verbose)
}

func callBatch(batch string,
	args []string,
	tmpfile string,
	verbose io.Writer,
	stdin io.Reader,
	stdout io.Writer,
	stderr io.Writer) (int, error) {
	params := []string{
		os.Getenv("COMSPEC"),
		"/C",
		batch,
	}
	fd, err := os.Create(batch)
	if err != nil {
		return 1, err
	}
	var writer *bufio.Writer
	if verbose != nil && verbose != ioutil.Discard {
		writer = bufio.NewWriter(io.MultiWriter(fd, verbose))
	} else {
		writer = bufio.NewWriter(fd)
	}
	io.WriteString(writer, "@call")
	for _, arg1 := range args {
		// UTF8 parameter to ANSI
		ansi, err := mbcs.Utf8ToConsoleCp(arg1)
		if err != nil {
			// println("utoa: " + err.Error())
			fd.Close()
			return -1, err
		}
		ansi = bytes.TrimSuffix(ansi, []byte{0})
		fmt.Fprintf(writer, " %s", ansi)
	}
	fmt.Fprintf(writer, "\r\n@set \"ERRORLEVEL_=%%ERRORLEVEL%%\"\r\n")

	// Sometimes %TEMP% has not ASCII letters.
	ansi, err := mbcs.Utf8ToConsoleCp(tmpfile)
	if err != nil {
		fd.Close()
		return -1, err
	}
	ansi = bytes.TrimSuffix(ansi, []byte{0})
	fmt.Fprintf(writer, "@(cd & set) > \"%s\"\r\n", ansi)
	fmt.Fprintf(writer, "@exit /b \"%%ERRORLEVEL_%%\"\r\n")
	writer.Flush()
	if err := fd.Close(); err != nil {
		return 1, err
	}
	cmd := exec.Cmd{
		Path:   params[0],
		Args:   params,
		Stdin:  stdin,
		Stdout: stdout,
		Stderr: stderr,
	}
	if err := cmd.Run(); err != nil {
		return 1, err
	}
	errorlevel, errorlevelOk := dos.GetErrorLevel(&cmd)
	if !errorlevelOk {
		errorlevel = 255
	}
	return errorlevel, nil
}

func RawSource(args []string, verbose io.Writer, debug bool, stdin io.Reader, stdout io.Writer, stderr io.Writer) (int, error) {
	tempDir := os.TempDir()
	pid := os.Getpid()
	batch := filepath.Join(tempDir, fmt.Sprintf("nyagos-%d.cmd", pid))
	tmpfile := filepath.Join(tempDir, fmt.Sprintf("nyagos-%d.tmp", pid))

	errorlevel, err := callBatch(
		batch,
		args,
		tmpfile,
		verbose,
		stdin,
		stdout,
		stderr)

	if err != nil {
		return errorlevel, err
	}

	if !debug {
		defer os.Remove(tmpfile)
		defer os.Remove(batch)
	}

	if errorlevel, err = loadTmpFile(tmpfile, verbose); err != nil {
		if os.IsNotExist(err) {
			return 1, fmt.Errorf("%s: the batch file may use `exit` without `/b` option. Could not find the change of the environment variables", args[0])
		}
		return 1, err
	}
	// println("ERRORLEVEL=", errorlevel)
	if err != nil {
		return errorlevel, err
	}
	if errorlevel != 0 {
		return errorlevel, fmt.Errorf("exit status %d", errorlevel)
	}
	return 0, nil
}
<commit_msg>Fix source.go for golint<commit_after>package shell

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/zetamatta/go-mbcs"

	"github.com/zetamatta/nyagos/dos"
)

func readEnv(scan *bufio.Scanner, verbose io.Writer) (int, error) {
	errorlevel := -1
	for scan.Scan() {
		line, err := mbcs.ConsoleCpToUtf8(scan.Bytes())
		if err != nil {
			continue
		}
		line = strings.TrimSpace(line)
		eqlPos := strings.Index(line, "=")
		if eqlPos > 0 {
			left := line[:eqlPos]
			right := line[eqlPos+1:]
			if left == "ERRORLEVEL_" {
				value, err := strconv.ParseUint(right, 10, 32)
				if err != nil {
					fmt.Fprintf(verbose, "Could not read ERRORLEVEL(%s)\n", right)
				} else {
					errorlevel = int(value)
				}
			} else {
				orig := os.Getenv(left)
				if verbose != nil {
					fmt.Fprintf(verbose, "%s=%s\n", left, right)
				}
				if orig != right {
					// fmt.Fprintf(os.Stderr, "%s:=%s\n", left, right)
					os.Setenv(left, right)
				}
			}
		}
	}
	return errorlevel, scan.Err()
}

func readPwd(scan *bufio.Scanner, verbose io.Writer) error {
	if !scan.Scan() {
		return errors.New("Could not load the new current directory")
	}
	if err := scan.Err(); err != nil {
		return err
	}
	line, err := mbcs.ConsoleCpToUtf8(scan.Bytes())
	if err != nil {
		return err
	}
	line = strings.TrimSpace(line)
	if verbose != nil {
		fmt.Fprintf(verbose, "cd \"%s\"\n", line)
	}
	os.Chdir(line)
	return nil
}

// loadTmpFile - read update the current-directory and environment-variables from tmp-file.
func loadTmpFile(fname string, verbose io.Writer) (int, error) {
	fp, err := os.Open(fname)
	if err != nil {
		return -1, err
	}
	defer fp.Close()

	scan := bufio.NewScanner(fp)
	if err := readPwd(scan, verbose); err != nil {
		return -1, err
	}
	return readEnv(scan, verbose)
}

func callBatch(batch string,
	args []string,
	tmpfile string,
	verbose io.Writer,
	stdin io.Reader,
	stdout io.Writer,
	stderr io.Writer) (int, error) {
	params := []string{
		os.Getenv("COMSPEC"),
		"/C",
		batch,
	}
	fd, err := os.Create(batch)
	if err != nil {
		return 1, err
	}
	var writer *bufio.Writer
	if verbose != nil && verbose != ioutil.Discard {
		writer = bufio.NewWriter(io.MultiWriter(fd, verbose))
	} else {
		writer = bufio.NewWriter(fd)
	}
	io.WriteString(writer, "@call")
	for _, arg1 := range args {
		// UTF8 parameter to ANSI
		ansi, err := mbcs.Utf8ToConsoleCp(arg1)
		if err != nil {
			// println("utoa: " + err.Error())
			fd.Close()
			return -1, err
		}
		ansi = bytes.TrimSuffix(ansi, []byte{0})
		fmt.Fprintf(writer, " %s", ansi)
	}
	fmt.Fprintf(writer, "\r\n@set \"ERRORLEVEL_=%%ERRORLEVEL%%\"\r\n")

	// Sometimes %TEMP% has not ASCII letters.
	ansi, err := mbcs.Utf8ToConsoleCp(tmpfile)
	if err != nil {
		fd.Close()
		return -1, err
	}
	ansi = bytes.TrimSuffix(ansi, []byte{0})
	fmt.Fprintf(writer, "@(cd & set) > \"%s\"\r\n", ansi)
	fmt.Fprintf(writer, "@exit /b \"%%ERRORLEVEL_%%\"\r\n")
	writer.Flush()
	if err := fd.Close(); err != nil {
		return 1, err
	}
	cmd := exec.Cmd{
		Path:   params[0],
		Args:   params,
		Stdin:  stdin,
		Stdout: stdout,
		Stderr: stderr,
	}
	if err := cmd.Run(); err != nil {
		return 1, err
	}
	errorlevel, errorlevelOk := dos.GetErrorLevel(&cmd)
	if !errorlevelOk {
		errorlevel = 255
	}
	return errorlevel, nil
}

// RawSource calls the batchfiles and load the changed variable the batchfile has done.
func RawSource(args []string, verbose io.Writer, debug bool, stdin io.Reader, stdout io.Writer, stderr io.Writer) (int, error) {
	tempDir := os.TempDir()
	pid := os.Getpid()
	batch := filepath.Join(tempDir, fmt.Sprintf("nyagos-%d.cmd", pid))
	tmpfile := filepath.Join(tempDir, fmt.Sprintf("nyagos-%d.tmp", pid))

	errorlevel, err := callBatch(
		batch,
		args,
		tmpfile,
		verbose,
		stdin,
		stdout,
		stderr)

	if err != nil {
		return errorlevel, err
	}

	if !debug {
		defer os.Remove(tmpfile)
		defer os.Remove(batch)
	}

	if errorlevel, err = loadTmpFile(tmpfile, verbose); err != nil {
		if os.IsNotExist(err) {
			return 1, fmt.Errorf("%s: the batch file may use `exit` without `/b` option. Could not find the change of the environment variables", args[0])
		}
		return 1, err
	}
	// println("ERRORLEVEL=", errorlevel)
	if err != nil {
		return errorlevel, err
	}
	if errorlevel != 0 {
		return errorlevel, fmt.Errorf("exit status %d", errorlevel)
	}
	return 0, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"sync"
	"testing"
	"time"

	. "launchpad.net/gocheck"
)

func init() {
	//log.SetFlags(log.LstdFlags | log.Lshortfile)
	log.SetOutput(ioutil.Discard)
}

// something that can wrap a gocheck.C testing.T or testing.B
// Just add more methods as we need them.
type Tester interface {
	Fatal(args ...interface{})
	Fatalf(format string, args ...interface{})
	Log(args ...interface{})
	Logf(format string, args ...interface{})
}

func Test(t *testing.T) { TestingT(t) }

type BasicSuite struct {
	servers []*testServer
	service *Service
}

var _ = Suite(&BasicSuite{})

// Make Setup and TearDown more generic, so we can bypass the gocheck Suite if
// needed.
func mySetup(s *BasicSuite, t Tester) {
	// start 4 possible backend servers
	ports := []string{"2001", "2002", "2003", "2004"}
	for _, p := range ports {
		server, err := NewTestServer("127.0.0.1:"+p, t)
		if err != nil {
			t.Fatal(err)
		}
		s.servers = append(s.servers, server)
	}

	svcCfg := ServiceConfig{
		Name: "testService",
		Addr: "127.0.0.1:2222",
	}

	if err := Registry.AddService(svcCfg); err != nil {
		t.Fatal(err)
	}

	s.service = Registry.GetService(svcCfg.Name)
}

// shutdown our backend servers
func myTearDown(s *BasicSuite, t Tester) {
	for _, s := range s.servers {
		s.Stop()
	}

	// get rid of the servers refs too!
	s.servers = nil

	err := Registry.RemoveService(s.service.Name)
	if err != nil {
		t.Fatalf("could not remove service '%s': %s", s.service.Name, err)
	}
}

func (s *BasicSuite) SetUpTest(c *C) {
	mySetup(s, c)
}

func (s *BasicSuite) TearDownTest(c *C) {
	myTearDown(s, c)
}

// Add a default backend for the next server we have running
func (s *BasicSuite) AddBackend(c Tester) {
	// get the backends via Config to use the Service's locking.
	svcCfg := s.service.Config()
	next := len(svcCfg.Backends)
	if next >= len(s.servers) {
		c.Fatal("no more servers")
	}

	name := fmt.Sprintf("backend_%d", next)
	cfg := BackendConfig{
		Name:      name,
		Addr:      s.servers[next].addr,
		CheckAddr: s.servers[next].addr,
	}

	s.service.add(NewBackend(cfg))
}

// Connect to address, and check response after write.
func checkResp(addr, expected string, c Tester) {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		c.Fatal(err)
	}
	defer conn.Close()

	if _, err := io.WriteString(conn, "testing\n"); err != nil {
		c.Fatal(err)
	}

	buff := make([]byte, 1024)
	n, err := conn.Read(buff)
	if err != nil {
		c.Fatal(err)
	}

	resp := string(buff[:n])
	if resp == "" {
		c.Fatal("No response")
	}

	if expected != "" && resp != expected {
		c.Fatal("Expected", expected, ", got", resp)
	}
}

func (s *BasicSuite) TestSingleBackend(c *C) {
	s.AddBackend(c)

	checkResp(s.service.Addr, s.servers[0].addr, c)
}

func (s *BasicSuite) TestRoundRobin(c *C) {
	s.AddBackend(c)
	s.AddBackend(c)

	checkResp(s.service.Addr, s.servers[0].addr, c)
	checkResp(s.service.Addr, s.servers[1].addr, c)
	checkResp(s.service.Addr, s.servers[0].addr, c)
	checkResp(s.service.Addr, s.servers[1].addr, c)
}

func (s *BasicSuite) TestWeightedRoundRobin(c *C) {
	s.AddBackend(c)
	s.AddBackend(c)
	s.AddBackend(c)

	s.service.Backends[0].Weight = 1
	s.service.Backends[1].Weight = 2
	s.service.Backends[2].Weight = 3

	// we already checked that we connect to the correct backends,
	// so skip the tcp connection this time.

	// one from the first server
	c.Assert(s.service.next().Name, Equals, "backend_0")
	// A weight of 2 should return twice
	c.Assert(s.service.next().Name, Equals, "backend_1")
	c.Assert(s.service.next().Name, Equals, "backend_1")
	// And a weight of 3 should return thrice
	c.Assert(s.service.next().Name, Equals, "backend_2")
	c.Assert(s.service.next().Name, Equals, "backend_2")
	c.Assert(s.service.next().Name, Equals, "backend_2")
	// and once around or good measure
	c.Assert(s.service.next().Name, Equals, "backend_0")
}

func (s *BasicSuite) TestLeastConn(c *C) {
	// replace out default service with one using LeastConn balancing
	Registry.RemoveService("testService")
	svcCfg := ServiceConfig{
		Name:    "testService",
		Addr:    "127.0.0.1:2223",
		Balance: "LC",
	}

	if err := Registry.AddService(svcCfg); err != nil {
		c.Fatal(err)
	}
	s.service = Registry.GetService("testService")

	s.AddBackend(c)
	s.AddBackend(c)

	// tie up 4 connections to the backends
	for i := 0; i < 4; i++ {
		conn, e := net.Dial("tcp", s.service.Addr)
		if e != nil {
			c.Fatal(e)
		}
		defer conn.Close()
	}

	s.AddBackend(c)

	checkResp(s.service.Addr, s.servers[2].addr, c)
	checkResp(s.service.Addr, s.servers[2].addr, c)
}

// Test health check by taking down a server from a configured backend
func (s *BasicSuite) TestFailedCheck(c *C) {
	s.service.CheckInterval = 500
	s.service.Fall = 1
	s.AddBackend(c)

	stats := s.service.Stats()
	c.Assert(stats.Backends[0].Up, Equals, true)

	// Stop the server, and see if the backend shows Down after our check
	// interval.
	s.servers[0].Stop()
	time.Sleep(800 * time.Millisecond)

	stats = s.service.Stats()
	c.Assert(stats.Backends[0].Up, Equals, false)
	c.Assert(stats.Backends[0].CheckFail, Equals, 1)

	// now try and connect to the service
	conn, err := net.Dial("tcp", s.service.Addr)
	if err != nil {
		// we should still get an initial connection
		c.Fatal(err)
	}

	b := make([]byte, 1024)
	n, err := conn.Read(b)
	if n != 0 || err != io.EOF {
		c.Fatal("connection should have been closed")
	}
}

// Update a backend in place
func (s *BasicSuite) TestUpdateBackend(c *C) {
	s.service.CheckInterval = 500
	s.service.Fall = 1
	s.AddBackend(c)

	cfg := s.service.Config()
	backendCfg := cfg.Backends[0]

	c.Assert(backendCfg.CheckAddr, Equals, backendCfg.Addr)

	backendCfg.CheckAddr = ""
	s.service.add(NewBackend(backendCfg))

	// see if the config reflects the new backend
	cfg = s.service.Config()
	c.Assert(len(cfg.Backends), Equals, 1)
	c.Assert(cfg.Backends[0].CheckAddr, Equals, "")

	// Stopping the server should not take down the backend
	// since there is no longer a Check address.
	s.servers[0].Stop()
	time.Sleep(800 * time.Millisecond)

	stats := s.service.Stats()
	c.Assert(stats.Backends[0].Up, Equals, true)
	// should have been no check failures
	c.Assert(stats.Backends[0].CheckFail, Equals, 0)
}

// Test removal of a single Backend from a service with multiple.
func (s *BasicSuite) TestRemoveBackend(c *C) {
	s.AddBackend(c)
	s.AddBackend(c)

	stats, err := Registry.ServiceStats("testService")
	if err != nil {
		c.Fatal(err)
	}

	c.Assert(len(stats.Backends), Equals, 2)

	backend1 := stats.Backends[0].Name

	err = Registry.RemoveBackend("testService", backend1)
	if err != nil {
		c.Fatal(err)
	}

	stats, err = Registry.ServiceStats("testService")
	if err != nil {
		c.Fatal(err)
	}

	c.Assert(len(stats.Backends), Equals, 1)

	_, err = Registry.BackendStats("testService", backend1)
	c.Assert(err, Equals, ErrNoBackend)
}

func (s *BasicSuite) TestUpdateService(c *C) {
	svcCfg := ServiceConfig{
		Name: "Update",
		Addr: "127.0.0.1:9324",
	}

	if err := Registry.AddService(svcCfg); err != nil {
		c.Fatal(err)
	}

	svc := Registry.GetService("Update")
	if svc == nil {
		c.Fatal(ErrNoService)
	}

	svcCfg = ServiceConfig{
		Name: "Update",
		Addr: "127.0.0.1:9425",
	}

	// Make sure we can't add it through AddService
	if err := Registry.AddService(svcCfg); err == nil {
		c.Fatal(err)
	}

	// Now update the service
	if err := Registry.UpdateService(svcCfg); err != nil {
		c.Fatal(err)
	}

	svc = Registry.GetService("Update")
	if svc == nil {
		c.Fatal(ErrNoService)
	}
	c.Assert(svc.Addr, Equals, "127.0.0.1:9425")

	if err := Registry.RemoveService("Update"); err != nil {
		c.Fatal(err)
	}
}

// Add backends and run response tests in parallel
func (s *BasicSuite) TestParallel(c *C) {
	var wg sync.WaitGroup

	client := func(i int) {
		s.AddBackend(c)
		// do a bunch of new connections in unison
		for i := 0; i < 100; i++ {
			checkResp(s.service.Addr, "", c)
		}

		conn, err := net.Dial("tcp", s.service.Addr)
		if err != nil {
			// we should still get an initial connection
			c.Fatal(err)
		}
		defer conn.Close()

		// now do some more continuous ping-pongs with the server
		buff := make([]byte, 1024)

		for i := 0; i < 1000; i++ {
			n, err := io.WriteString(conn, "Testing testing\n")
			if err != nil || n == 0 {
				c.Fatal("couldn't write:", err)
			}

			n, err = conn.Read(buff)
			if err != nil || n == 0 {
				c.Fatal("no response:", err)
			}
		}
		wg.Done()
	}

	for i := 0; i < 4; i++ {
		wg.Add(1)
		go client(i)
	}

	wg.Wait()
}
<commit_msg>Services weren't coming back up<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"sync"
	"testing"
	"time"

	. "launchpad.net/gocheck"
)

func init() {
	//log.SetFlags(log.LstdFlags | log.Lshortfile)
	log.SetOutput(ioutil.Discard)
}

// something that can wrap a gocheck.C testing.T or testing.B
// Just add more methods as we need them.
type Tester interface {
	Fatal(args ...interface{})
	Fatalf(format string, args ...interface{})
	Log(args ...interface{})
	Logf(format string, args ...interface{})
}

func Test(t *testing.T) { TestingT(t) }

type BasicSuite struct {
	servers []*testServer
	service *Service
}

var _ = Suite(&BasicSuite{})

// Make Setup and TearDown more generic, so we can bypass the gocheck Suite if
// needed.
func mySetup(s *BasicSuite, t Tester) {
	// start 4 possible backend servers
	ports := []string{"2001", "2002", "2003", "2004"}
	for _, p := range ports {
		server, err := NewTestServer("127.0.0.1:"+p, t)
		if err != nil {
			t.Fatal(err)
		}
		s.servers = append(s.servers, server)
	}

	svcCfg := ServiceConfig{
		Name: "testService",
		Addr: "127.0.0.1:2222",
	}

	if err := Registry.AddService(svcCfg); err != nil {
		t.Fatal(err)
	}

	s.service = Registry.GetService(svcCfg.Name)
}

// shutdown our backend servers
func myTearDown(s *BasicSuite, t Tester) {
	for _, s := range s.servers {
		s.Stop()
	}

	// get rid of the servers refs too!
	s.servers = nil

	err := Registry.RemoveService(s.service.Name)
	if err != nil {
		t.Fatalf("could not remove service '%s': %s", s.service.Name, err)
	}
}

func (s *BasicSuite) SetUpTest(c *C) {
	mySetup(s, c)
}

func (s *BasicSuite) TearDownTest(c *C) {
	myTearDown(s, c)
}

// Add a default backend for the next server we have running
func (s *BasicSuite) AddBackend(c Tester) {
	// get the backends via Config to use the Service's locking.
	svcCfg := s.service.Config()
	next := len(svcCfg.Backends)
	if next >= len(s.servers) {
		c.Fatal("no more servers")
	}

	name := fmt.Sprintf("backend_%d", next)
	cfg := BackendConfig{
		Name:      name,
		Addr:      s.servers[next].addr,
		CheckAddr: s.servers[next].addr,
	}

	s.service.add(NewBackend(cfg))
}

// Connect to address, and check response after write.
func checkResp(addr, expected string, c Tester) {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		c.Fatal(err)
	}
	defer conn.Close()

	if _, err := io.WriteString(conn, "testing\n"); err != nil {
		c.Fatal(err)
	}

	buff := make([]byte, 1024)
	n, err := conn.Read(buff)
	if err != nil {
		c.Fatal(err)
	}

	resp := string(buff[:n])
	if resp == "" {
		c.Fatal("No response")
	}

	if expected != "" && resp != expected {
		c.Fatal("Expected", expected, ", got", resp)
	}
}

func (s *BasicSuite) TestSingleBackend(c *C) {
	s.AddBackend(c)

	checkResp(s.service.Addr, s.servers[0].addr, c)
}

func (s *BasicSuite) TestRoundRobin(c *C) {
	s.AddBackend(c)
	s.AddBackend(c)

	checkResp(s.service.Addr, s.servers[0].addr, c)
	checkResp(s.service.Addr, s.servers[1].addr, c)
	checkResp(s.service.Addr, s.servers[0].addr, c)
	checkResp(s.service.Addr, s.servers[1].addr, c)
}

func (s *BasicSuite) TestWeightedRoundRobin(c *C) {
	s.AddBackend(c)
	s.AddBackend(c)
	s.AddBackend(c)

	s.service.Backends[0].Weight = 1
	s.service.Backends[1].Weight = 2
	s.service.Backends[2].Weight = 3

	// we already checked that we connect to the correct backends,
	// so skip the tcp connection this time.

	// one from the first server
	c.Assert(s.service.next().Name, Equals, "backend_0")
	// A weight of 2 should return twice
	c.Assert(s.service.next().Name, Equals, "backend_1")
	c.Assert(s.service.next().Name, Equals, "backend_1")
	// And a weight of 3 should return thrice
	c.Assert(s.service.next().Name, Equals, "backend_2")
	c.Assert(s.service.next().Name, Equals, "backend_2")
	c.Assert(s.service.next().Name, Equals, "backend_2")
	// and once around or good measure
	c.Assert(s.service.next().Name, Equals, "backend_0")
}

func (s *BasicSuite) TestLeastConn(c *C) {
	// replace out default service with one using LeastConn balancing
	Registry.RemoveService("testService")
	svcCfg := ServiceConfig{
		Name:    "testService",
		Addr:    "127.0.0.1:2223",
		Balance: "LC",
	}

	if err := Registry.AddService(svcCfg); err != nil {
		c.Fatal(err)
	}
	s.service = Registry.GetService("testService")

	s.AddBackend(c)
	s.AddBackend(c)

	// tie up 4 connections to the backends
	for i := 0; i < 4; i++ {
		conn, e := net.Dial("tcp", s.service.Addr)
		if e != nil {
			c.Fatal(e)
		}
		defer conn.Close()
	}

	s.AddBackend(c)

	checkResp(s.service.Addr, s.servers[2].addr, c)
	checkResp(s.service.Addr, s.servers[2].addr, c)
}

// Test health check by taking down a server from a configured backend
func (s *BasicSuite) TestFailedCheck(c *C) {
	s.service.CheckInterval = 500
	s.service.Fall = 1
	s.AddBackend(c)

	stats := s.service.Stats()
	c.Assert(stats.Backends[0].Up, Equals, true)

	// Stop the server, and see if the backend shows Down after our check
	// interval.
	s.servers[0].Stop()
	time.Sleep(800 * time.Millisecond)

	stats = s.service.Stats()
	c.Assert(stats.Backends[0].Up, Equals, false)
	c.Assert(stats.Backends[0].CheckFail, Equals, 1)

	// now try and connect to the service
	conn, err := net.Dial("tcp", s.service.Addr)
	if err != nil {
		// we should still get an initial connection
		c.Fatal(err)
	}

	b := make([]byte, 1024)
	n, err := conn.Read(b)
	if n != 0 || err != io.EOF {
		c.Fatal("connection should have been closed")
	}

	// now bring that server back up
	server, err := NewTestServer(s.servers[0].addr, c)
	if err != nil {
		c.Fatal(err)
	}
	s.servers[0] = server

	time.Sleep(800 * time.Millisecond)
	stats = s.service.Stats()
	c.Assert(stats.Backends[0].Up, Equals, true)
}

// Update a backend in place
func (s *BasicSuite) TestUpdateBackend(c *C) {
	s.service.CheckInterval = 500
	s.service.Fall = 1
	s.AddBackend(c)

	cfg := s.service.Config()
	backendCfg := cfg.Backends[0]

	c.Assert(backendCfg.CheckAddr, Equals, backendCfg.Addr)

	backendCfg.CheckAddr = ""
	s.service.add(NewBackend(backendCfg))

	// see if the config reflects the new backend
	cfg = s.service.Config()
	c.Assert(len(cfg.Backends), Equals, 1)
	c.Assert(cfg.Backends[0].CheckAddr, Equals, "")

	// Stopping the server should not take down the backend
	// since there is no longer a Check address.
	s.servers[0].Stop()
	time.Sleep(800 * time.Millisecond)

	stats := s.service.Stats()
	c.Assert(stats.Backends[0].Up, Equals, true)
	// should have been no check failures
	c.Assert(stats.Backends[0].CheckFail, Equals, 0)
}

// Test removal of a single Backend from a service with multiple.
func (s *BasicSuite) TestRemoveBackend(c *C) {
	s.AddBackend(c)
	s.AddBackend(c)

	stats, err := Registry.ServiceStats("testService")
	if err != nil {
		c.Fatal(err)
	}

	c.Assert(len(stats.Backends), Equals, 2)

	backend1 := stats.Backends[0].Name

	err = Registry.RemoveBackend("testService", backend1)
	if err != nil {
		c.Fatal(err)
	}

	stats, err = Registry.ServiceStats("testService")
	if err != nil {
		c.Fatal(err)
	}

	c.Assert(len(stats.Backends), Equals, 1)

	_, err = Registry.BackendStats("testService", backend1)
	c.Assert(err, Equals, ErrNoBackend)
}

func (s *BasicSuite) TestUpdateService(c *C) {
	svcCfg := ServiceConfig{
		Name: "Update",
		Addr: "127.0.0.1:9324",
	}

	if err := Registry.AddService(svcCfg); err != nil {
		c.Fatal(err)
	}

	svc := Registry.GetService("Update")
	if svc == nil {
		c.Fatal(ErrNoService)
	}

	svcCfg = ServiceConfig{
		Name: "Update",
		Addr: "127.0.0.1:9425",
	}

	// Make sure we can't add it through AddService
	if err := Registry.AddService(svcCfg); err == nil {
		c.Fatal(err)
	}

	// Now update the service
	if err := Registry.UpdateService(svcCfg); err != nil {
		c.Fatal(err)
	}

	svc = Registry.GetService("Update")
	if svc == nil {
		c.Fatal(ErrNoService)
	}
	c.Assert(svc.Addr, Equals, "127.0.0.1:9425")

	if err := Registry.RemoveService("Update"); err != nil {
		c.Fatal(err)
	}
}

// Add backends and run response tests in parallel
func (s *BasicSuite) TestParallel(c *C) {
	var wg sync.WaitGroup

	client := func(i int) {
		s.AddBackend(c)
		// do a bunch of new connections in unison
		for i := 0; i < 100; i++ {
			checkResp(s.service.Addr, "", c)
		}

		conn, err := net.Dial("tcp", s.service.Addr)
		if err != nil {
			// we should still get an initial connection
			c.Fatal(err)
		}
		defer conn.Close()

		// now do some more continuous ping-pongs with the server
		buff := make([]byte, 1024)

		for i := 0; i < 1000; i++ {
			n, err := io.WriteString(conn, "Testing testing\n")
			if err != nil || n == 0 {
				c.Fatal("couldn't write:", err)
			}

			n, err = conn.Read(buff)
			if err != nil || n == 0 {
				c.Fatal("no response:", err)
			}
		}
		wg.Done()
	}

	for i := 0; i < 4; i++ {
		wg.Add(1)
		go client(i)
	}

	wg.Wait()
}
<|endoftext|>


<commit_before>package protocol

//TODO: Handle errors

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
)

// Represents server list ping response.
type ServerListPing struct {
	Ver  Version `json:"version"`
	Pl   Players `json:"players"`
	Desc Chat    `json:"description"`
	Fav  string  `json:"favicon,omitempty"`
}

type Version struct {
	Name     string `json:"name"`
	Protocol uint32 `json:"protocol"`
}

type Players struct {
	Max    uint `json:"max"`
	Online uint `json:"online"`
}

type Chat struct {
	Text string `json:"text"`
}

type Response struct {
	data *bytes.Buffer
}

// Creates a new response.
func NewResponse() *Response {
	return &Response{data: new(bytes.Buffer)}
}

// Writes a boolean.
func (r *Response) WriteBoolean(b bool) *Response {
	if b {
		return r.WriteByte(1)
	} else {
		return r.WriteByte(0)
	}
}

// Writes a Chat JSON Object.
func (r *Response) WriteChat(obj string) *Response {
	return r.WriteJSON(Chat{obj})
}

// Writes the given object as a JSON string.
func (r *Response) WriteJSON(obj interface{}) *Response {
	j, _ := json.Marshal(obj)
	return r.WriteByteArray(j)
}

// Writes the given byte.
func (r *Response) WriteByte(b byte) *Response {
	r.data.Write([]byte{b})
	return r
}

func (r *Response) WriteUnsignedByte(b uint8) *Response {
	binary.Write(r.data, ByteOrder, b)
	return r
}

// Writes a varint.
func (r *Response) WriteVarint(i uint32) *Response {
	_, err := r.data.Write(Uvarint(i))
	if err != nil {
		panic(err)
	}
	return r
}

// Writes an integer.
func (r *Response) WriteInt(i int32) *Response {
	binary.Write(r.data, ByteOrder, i)
	return r
}

// Writes a long.
func (r *Response) WriteLong(l int64) *Response {
	binary.Write(r.data, ByteOrder, l)
	return r
}

// Writes a byte array.
func (r *Response) WriteByteArray(b []byte) *Response {
	r.WriteVarint(uint32(len(b)))
	r.data.Write(b)
	return r
}

// Writes a string.
func (r *Response) WriteString(str string) *Response {
	return r.WriteByteArray([]byte(str))
}

// Returns the raw packet created from the written bytes and the provided id.
func (r *Response) ToRawPacket(id uint64) *RawPacket {
	return NewRawPacket(id, r.data.Bytes(), nil)
}

func (r *Response) Clear() {
	r.data = new(bytes.Buffer)
}
<commit_msg>:memo: Documented the functions.<commit_after>package protocol

//TODO: Handle errors

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
)

// Represents server list ping response.
type ServerListPing struct {
	Ver  Version `json:"version"`
	Pl   Players `json:"players"`
	Desc Chat    `json:"description"`
	Fav  string  `json:"favicon,omitempty"`
}

type Version struct {
	Name     string `json:"name"`
	Protocol uint32 `json:"protocol"`
}

type Players struct {
	Max    uint `json:"max"`
	Online uint `json:"online"`
}

type Chat struct {
	Text string `json:"text"`
}

type Response struct {
	data *bytes.Buffer
}

// NewResponse creates a new response.
func NewResponse() *Response {
	return &Response{data: new(bytes.Buffer)}
}

// WriteBoolean writes the given boolean to the current response.
func (r *Response) WriteBoolean(b bool) *Response {
	if b {
		return r.WriteByte(1)
	} else {
		return r.WriteByte(0)
	}
}

// WriteChat writes a Chat JSON Object to the current response.
func (r *Response) WriteChat(obj string) *Response {
	return r.WriteJSON(Chat{obj})
}

// WriteJSON writes the given interface as a JSON string to the current response.
// (Currently ignores failure.)
func (r *Response) WriteJSON(obj interface{}) *Response {
	j, _ := json.Marshal(obj)
	return r.WriteByteArray(j)
}

// WriteByte writes the given byte to the current response.
func (r *Response) WriteByte(b byte) *Response {
	r.data.Write([]byte{b})
	return r
}

// WriteUnsignedByte writes the given unsigned byte to the current response.
func (r *Response) WriteUnsignedByte(b uint8) *Response {
	binary.Write(r.data, ByteOrder, b)
	return r
}

// WriteVarInt writes the given VarInt to the current response.
func (r *Response) WriteVarint(i uint32) *Response {
	_, err := r.data.Write(Uvarint(i))
	if err != nil {
		panic(err)
	}
	return r
}

// WriteInt writes the given integer to the current response.
func (r *Response) WriteInt(i int32) *Response {
	binary.Write(r.data, ByteOrder, i)
	return r
}

// WriteLong writes the given long to the current response.
func (r *Response) WriteLong(l int64) *Response {
	binary.Write(r.data, ByteOrder, l)
	return r
}

// WriteByteArray writes the given byte array to the current response.
func (r *Response) WriteByteArray(b []byte) *Response {
	r.WriteVarint(uint32(len(b)))
	r.data.Write(b)
	return r
}

// WriteString writes the given string to the current response.
func (r *Response) WriteString(str string) *Response {
	return r.WriteByteArray([]byte(str))
}

// Returns the raw packet created from the written bytes and the provided id.
func (r *Response) ToRawPacket(id uint64) *RawPacket {
	return NewRawPacket(id, r.data.Bytes(), nil)
}

// Clear clears the data from the response's buffer.
func (r *Response) Clear() {
	r.data = new(bytes.Buffer)
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

// Package gerrit implements a gerrit-fetcher using https://github.com/andygrunwald/go-gerrit
package gerrit

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/andygrunwald/go-gerrit"
	"github.com/sirupsen/logrus"

	"k8s.io/test-infra/prow/config"
	"k8s.io/test-infra/prow/kube"
	"k8s.io/test-infra/prow/pjutil"
)

type kubeClient interface {
	CreateProwJob(kube.ProwJob) (kube.ProwJob, error)
}

type gerritAuthentication interface {
	SetCookieAuth(name, value string)
}

type gerritAccount interface {
	GetAccount(name string) (*gerrit.AccountInfo, *gerrit.Response, error)
	SetUsername(accountID string, input *gerrit.UsernameInput) (*string, *gerrit.Response, error)
}

type gerritChange interface {
	QueryChanges(opt *gerrit.QueryChangeOptions) (*[]gerrit.ChangeInfo, *gerrit.Response, error)
	SetReview(changeID, revisionID string, input *gerrit.ReviewInput) (*gerrit.ReviewResult, *gerrit.Response, error)
}

type configAgent interface {
	Config() *config.Config
}

// Controller manages gerrit changes.
type Controller struct {
	ca configAgent

	// go-gerrit change endpoint client
	auth     gerritAuthentication
	account  gerritAccount
	gc       gerritChange
	instance string
	storage  string
	projects []string

	kc kubeClient

	lastUpdate time.Time
}

// NewController returns a new gerrit controller client
func NewController(instance, storage string, projects []string, kc *kube.Client, ca *config.Agent) (*Controller, error) {
	lastUpdate := time.Now()
	if storage != "" {
		buf, err := ioutil.ReadFile(storage)
		if err == nil {
			unix, err := strconv.ParseInt(string(buf), 10, 64)
			if err != nil {
				return nil, err
			} else {
				lastUpdate = time.Unix(unix, 0)
			}
		} else if !os.IsNotExist(err) {
			return nil, err
		}
		// fallback to time.Now() if file does not exist yet
	}

	c, err := gerrit.NewClient(instance, nil)
	if err != nil {
		return nil, err
	}

	return &Controller{
		instance:   instance,
		projects:   projects,
		kc:         kc,
		ca:         ca,
		auth:       c.Authentication,
		account:    c.Accounts,
		gc:         c.Changes,
		lastUpdate: lastUpdate,
		storage:    storage,
	}, nil
}

// Auth authenticates to gerrit server
// Token will expire, so we need to regenerate it once so often
func (c *Controller) Auth() error {
	cmd := exec.Command("python", "./git-cookie-authdaemon")
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("Fail to authenticate to gerrit using git-cookie-authdaemon : %v", err)
	}

	raw, err := ioutil.ReadFile(filepath.Join(os.Getenv("HOME"), ".git-credential-cache/cookie"))
	if err != nil {
		return err
	}
	fields := strings.Fields(string(raw))
	token := fields[len(fields)-1]

	c.auth.SetCookieAuth("o", token)

	self, _, err := c.account.GetAccount("self")
	if err != nil {
		logrus.WithError(err).Errorf("Fail to auth with token: %s", token)
		return err
	}

	logrus.Infof("Authentication successful, Username: %s", self.Name)

	return nil
}

// SaveLastSync saves last sync time in Unix to a volume
func (c *Controller) SaveLastSync(lastSync time.Time) error {
	if c.storage == "" {
		return nil
	}

	lastSyncUnix := strconv.FormatInt(lastSync.Unix(), 10)
	logrus.Infof("Writing last sync: %s", lastSyncUnix)

	err := ioutil.WriteFile(c.storage+".tmp", []byte(lastSyncUnix), 0644)
	if err != nil {
		return err
	}
	return os.Rename(c.storage+".tmp", c.storage)
}

// Sync looks for newly made gerrit changes
// and creates prowjobs according to presubmit specs
func (c *Controller) Sync() error {
	syncTime := time.Now()
	changes := c.QueryChanges()

	for _, change := range changes {
		if err := c.ProcessChange(change); err != nil {
			logrus.WithError(err).Errorf("Failed process change %v", change.CurrentRevision)
		}
	}

	c.lastUpdate = syncTime
	if err := c.SaveLastSync(syncTime); err != nil {
		logrus.WithError(err).Errorf("last sync %v, cannot save to path %v", syncTime, c.storage)
	}
	logrus.Infof("Processed %d changes", len(changes))
	return nil
}

func (c *Controller) queryProjectChanges(proj string) ([]gerrit.ChangeInfo, error) {
	pending := []gerrit.ChangeInfo{}

	opt := &gerrit.QueryChangeOptions{}
	opt.Query = append(opt.Query, "project:"+proj+"+status:open")
	opt.AdditionalFields = []string{"CURRENT_REVISION", "CURRENT_COMMIT"}

	start := 0

	for {
		opt.Limit = c.ca.Config().Gerrit.RateLimit
		opt.Start = start

		// The change output is sorted by the last update time, most recently updated to oldest updated.
		// Gerrit API docs: https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#list-changes
		changes, _, err := c.gc.QueryChanges(opt)
		if err != nil {
			// should not happen? Let next sync loop catch up
			return pending, fmt.Errorf("failed to query gerrit changes: %v", err)
		}

		logrus.Infof("Find %d changes from query %v", len(*changes), opt.Query)

		if len(*changes) == 0 {
			return pending, nil
		}
		start += len(*changes)

		for _, change := range *changes {
			// if we already processed this change, then we stop the current sync loop
			const layout = "2006-01-02 15:04:05"
			updated, err := time.Parse(layout, change.Updated)
			if err != nil {
				logrus.WithError(err).Errorf("Parse time %v failed", change.Updated)
				continue
			}

			// process if updated later than last updated
			// stop if update was stale
			if updated.After(c.lastUpdate) {
				// we need to make sure the change update is from a new commit change
				rev, ok := change.Revisions[change.CurrentRevision]
				if !ok {
					logrus.WithError(err).Errorf("(should not happen?)cannot find current revision for change %v", change.ID)
					continue
				}

				created, err := time.Parse(layout, rev.Created)
				if err != nil {
					logrus.WithError(err).Errorf("Parse time %v failed", rev.Created)
					continue
				}

				if !created.After(c.lastUpdate) {
					// stale commit
					continue
				}

				pending = append(pending, change)
			} else {
				return pending, nil
			}
		}
	}
}

// QueryChanges will query all valid gerrit changes since controller's last sync loop
func (c *Controller) QueryChanges() []gerrit.ChangeInfo {
	// store a map of changeID:change
	pending := []gerrit.ChangeInfo{}

	// can only query against one project at a time :-(
	for _, proj := range c.projects {
		if res, err := c.queryProjectChanges(proj); err != nil {
			logrus.WithError(err).Errorf("fail to query changes for project %s", proj)
		} else {
			pending = append(pending, res...)
		}
	}

	return pending
}

// ProcessChange creates new presubmit prowjobs base off the gerrit changes
func (c *Controller) ProcessChange(change gerrit.ChangeInfo) error {
	rev, ok := change.Revisions[change.CurrentRevision]
	if !ok {
		return fmt.Errorf("cannot find current revision for change %v", change.ID)
	}

	parentSHA := ""
	if len(rev.Commit.Parents) > 0 {
		parentSHA = rev.Commit.Parents[0].Commit
	}

	logger := logrus.WithField("gerrit change", change.Number)

	type triggeredJob struct {
		Name, URL string
	}
	triggeredJobs := []triggeredJob{}

	for _, spec := range c.ca.Config().Presubmits[c.instance+"/"+change.Project] {
		kr := kube.Refs{
			Org:     c.instance,
			Repo:    change.Project,
			BaseRef: change.Branch,
			BaseSHA: parentSHA,
			Pulls: []kube.Pull{
				{
					Number: change.Number,
					Author: rev.Commit.Author.Name,
					SHA:    change.CurrentRevision,
					Ref:    rev.Ref,
				},
			},
		}

		// TODO(krzyzacy): Support AlwaysRun and RunIfChanged
		pj := pjutil.NewProwJob(pjutil.PresubmitSpec(spec, kr), map[string]string{})
		logger.WithFields(pjutil.ProwJobFields(&pj)).Infof("Creating a new prowjob for change %s.", change.Number)
		if _, err := c.kc.CreateProwJob(pj); err != nil {
			logger.WithError(err).Errorf("fail to create prowjob %v", pj)
		} else {
			var b bytes.Buffer
			url := ""
			template := c.ca.Config().Plank.JobURLTemplate
			if template != nil {
				if err := template.Execute(&b, &pj); err != nil {
					logger.WithFields(pjutil.ProwJobFields(&pj)).Errorf("error executing URL template: %v", err)
				}
				// TODO(krzyzacy): We doesn't have buildID here yet - do a hack to get a proper URL to the PR
				// Remove this once we have proper report interface.

				// mangle
				// https://gubernator.k8s.io/build/gob-prow/pr-logs/pull/some/repo/8940/pull-test-infra-presubmit//
				// to
				// https://gubernator.k8s.io/builds/gob-prow/pr-logs/pull/some_repo/8940/pull-test-infra-presubmit/
				url = b.String()
				url = strings.Replace(url, "build", "builds", 1)
				// TODO(krzyzacy): gerrit path can be foo.googlesource.com/bar/baz, which means we took bar/baz as the repo
				// we are mangling the path in bootstrap.py, we need to handle this better in podutils
				url = strings.Replace(url, change.Project, strings.Replace(change.Project, "/", "_", -1), 1)
				url = strings.TrimSuffix(url, "//")
			}
			triggeredJobs = append(triggeredJobs, triggeredJob{Name: spec.Name, URL: url})
		}
	}

	if len(triggeredJobs) > 0 {
		// comment back to gerrit
		message := "Triggered presubmit:"
		for _, job := range triggeredJobs {
			if job.URL != "" {
				message += fmt.Sprintf("\n  * Name: %s, URL: %s", job.Name, job.URL)
			} else {
				message += fmt.Sprintf("\n  * Name: %s", job.Name)
			}
		}

		if _, _, err := c.gc.SetReview(change.ID, change.CurrentRevision, &gerrit.ReviewInput{
			Message: message,
		}); err != nil {
			return fmt.Errorf("cannot comment to gerrit: %v", err)
		}
	}

	return nil
}
<commit_msg>add some more log entries for gerrit<commit_after>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

// Package gerrit implements a gerrit-fetcher using https://github.com/andygrunwald/go-gerrit
package gerrit

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/andygrunwald/go-gerrit"
	"github.com/sirupsen/logrus"

	"k8s.io/test-infra/prow/config"
	"k8s.io/test-infra/prow/kube"
	"k8s.io/test-infra/prow/pjutil"
)

type kubeClient interface {
	CreateProwJob(kube.ProwJob) (kube.ProwJob, error)
}

type gerritAuthentication interface {
	SetCookieAuth(name, value string)
}

type gerritAccount interface {
	GetAccount(name string) (*gerrit.AccountInfo, *gerrit.Response, error)
	SetUsername(accountID string, input *gerrit.UsernameInput) (*string, *gerrit.Response, error)
}

type gerritChange interface {
	QueryChanges(opt *gerrit.QueryChangeOptions) (*[]gerrit.ChangeInfo, *gerrit.Response, error)
	SetReview(changeID, revisionID string, input *gerrit.ReviewInput) (*gerrit.ReviewResult, *gerrit.Response, error)
}

type configAgent interface {
	Config() *config.Config
}

// Controller manages gerrit changes.
type Controller struct {
	ca configAgent

	// go-gerrit change endpoint client
	auth     gerritAuthentication
	account  gerritAccount
	gc       gerritChange
	instance string
	storage  string
	projects []string

	kc kubeClient

	lastUpdate time.Time
}

// NewController returns a new gerrit controller client
func NewController(instance, storage string, projects []string, kc *kube.Client, ca *config.Agent) (*Controller, error) {
	lastUpdate := time.Now()
	if storage != "" {
		buf, err := ioutil.ReadFile(storage)
		if err == nil {
			unix, err := strconv.ParseInt(string(buf), 10, 64)
			if err != nil {
				return nil, err
			} else {
				lastUpdate = time.Unix(unix, 0)
			}
		} else if !os.IsNotExist(err) {
			return nil, err
		}
		// fallback to time.Now() if file does not exist yet
	}

	c, err := gerrit.NewClient(instance, nil)
	if err != nil {
		return nil, err
	}

	return &Controller{
		instance:   instance,
		projects:   projects,
		kc:         kc,
		ca:         ca,
		auth:       c.Authentication,
		account:    c.Accounts,
		gc:         c.Changes,
		lastUpdate: lastUpdate,
		storage:    storage,
	}, nil
}

// Auth authenticates to gerrit server
// Token will expire, so we need to regenerate it once so often
func (c *Controller) Auth() error {
	cmd := exec.Command("python", "./git-cookie-authdaemon")
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("Fail to authenticate to gerrit using git-cookie-authdaemon : %v", err)
	}

	raw, err := ioutil.ReadFile(filepath.Join(os.Getenv("HOME"), ".git-credential-cache/cookie"))
	if err != nil {
		return err
	}
	fields := strings.Fields(string(raw))
	token := fields[len(fields)-1]

	c.auth.SetCookieAuth("o", token)

	self, _, err := c.account.GetAccount("self")
	if err != nil {
		logrus.WithError(err).Errorf("Fail to auth with token: %s", token)
		return err
	}

	logrus.Infof("Authentication successful, Username: %s", self.Name)

	return nil
}

// SaveLastSync saves last sync time in Unix to a volume
func (c *Controller) SaveLastSync(lastSync time.Time) error {
	if c.storage == "" {
		return nil
	}

	lastSyncUnix := strconv.FormatInt(lastSync.Unix(), 10)
	logrus.Infof("Writing last sync: %s", lastSyncUnix)

	err := ioutil.WriteFile(c.storage+".tmp", []byte(lastSyncUnix), 0644)
	if err != nil {
		return err
	}
	return os.Rename(c.storage+".tmp", c.storage)
}

// Sync looks for newly made gerrit changes
// and creates prowjobs according to presubmit specs
func (c *Controller) Sync() error {
	syncTime := time.Now()
	changes := c.QueryChanges()

	for _, change := range changes {
		if err := c.ProcessChange(change); err != nil {
			logrus.WithError(err).Errorf("Failed process change %v", change.CurrentRevision)
		}
	}

	c.lastUpdate = syncTime
	if err := c.SaveLastSync(syncTime); err != nil {
		logrus.WithError(err).Errorf("last sync %v, cannot save to path %v", syncTime, c.storage)
	}
	logrus.Infof("Processed %d changes", len(changes))
	return nil
}

func (c *Controller) queryProjectChanges(proj string) ([]gerrit.ChangeInfo, error) {
	pending := []gerrit.ChangeInfo{}

	opt := &gerrit.QueryChangeOptions{}
	opt.Query = append(opt.Query, "project:"+proj+"+status:open")
	opt.AdditionalFields = []string{"CURRENT_REVISION", "CURRENT_COMMIT"}

	start := 0

	for {
		opt.Limit = c.ca.Config().Gerrit.RateLimit
		opt.Start = start

		// The change output is sorted by the last update time, most recently updated to oldest updated.
		// Gerrit API docs: https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#list-changes
		changes, _, err := c.gc.QueryChanges(opt)
		if err != nil {
			// should not happen? Let next sync loop catch up
			return pending, fmt.Errorf("failed to query gerrit changes: %v", err)
		}

		logrus.Infof("Find %d changes from query %v", len(*changes), opt.Query)

		if len(*changes) == 0 {
			return pending, nil
		}
		start += len(*changes)

		for _, change := range *changes {
			// if we already processed this change, then we stop the current sync loop
			const layout = "2006-01-02 15:04:05"
			updated, err := time.Parse(layout, change.Updated)
			if err != nil {
				logrus.WithError(err).Errorf("Parse time %v failed", change.Updated)
				continue
			}

			logrus.Infof("Change %s, last updated %s", change.Number, change.Updated)

			// process if updated later than last updated
			// stop if update was stale
			if updated.After(c.lastUpdate) {
				// we need to make sure the change update is from a new commit change
				rev, ok := change.Revisions[change.CurrentRevision]
				if !ok {
					logrus.WithError(err).Errorf("(should not happen?)cannot find current revision for change %v", change.ID)
					continue
				}

				created, err := time.Parse(layout, rev.Created)
				if err != nil {
					logrus.WithError(err).Errorf("Parse time %v failed", rev.Created)
					continue
				}

				if !created.After(c.lastUpdate) {
					// stale commit
					continue
				}

				pending = append(pending, change)
			} else {
				return pending, nil
			}
		}
	}
}

// QueryChanges will query all valid gerrit changes since controller's last sync loop
func (c *Controller) QueryChanges() []gerrit.ChangeInfo {
	// store a map of changeID:change
	pending := []gerrit.ChangeInfo{}

	// can only query against one project at a time :-(
	for _, proj := range c.projects {
		if res, err := c.queryProjectChanges(proj); err != nil {
			logrus.WithError(err).Errorf("fail to query changes for project %s", proj)
		} else {
			pending = append(pending, res...)
		}
	}

	return pending
}

// ProcessChange creates new presubmit prowjobs base off the gerrit changes
func (c *Controller) ProcessChange(change gerrit.ChangeInfo) error {
	rev, ok := change.Revisions[change.CurrentRevision]
	if !ok {
		return fmt.Errorf("cannot find current revision for change %v", change.ID)
	}

	parentSHA := ""
	if len(rev.Commit.Parents) > 0 {
		parentSHA = rev.Commit.Parents[0].Commit
	}

	logger := logrus.WithField("gerrit change", change.Number)

	type triggeredJob struct {
		Name, URL string
	}
	triggeredJobs := []triggeredJob{}

	for _, spec := range c.ca.Config().Presubmits[c.instance+"/"+change.Project] {
		kr := kube.Refs{
			Org:     c.instance,
			Repo:    change.Project,
			BaseRef: change.Branch,
			BaseSHA: parentSHA,
			Pulls: []kube.Pull{
				{
					Number: change.Number,
					Author: rev.Commit.Author.Name,
					SHA:    change.CurrentRevision,
					Ref:    rev.Ref,
				},
			},
		}

		// TODO(krzyzacy): Support AlwaysRun and RunIfChanged
		pj := pjutil.NewProwJob(pjutil.PresubmitSpec(spec, kr), map[string]string{})
		logger.WithFields(pjutil.ProwJobFields(&pj)).Infof("Creating a new prowjob for change %s.", change.Number)
		if _, err := c.kc.CreateProwJob(pj); err != nil {
			logger.WithError(err).Errorf("fail to create prowjob %v", pj)
		} else {
			var b bytes.Buffer
			url := ""
			template := c.ca.Config().Plank.JobURLTemplate
			if template != nil {
				if err := template.Execute(&b, &pj); err != nil {
					logger.WithFields(pjutil.ProwJobFields(&pj)).Errorf("error executing URL template: %v", err)
				}
				// TODO(krzyzacy): We doesn't have buildID here yet - do a hack to get a proper URL to the PR
				// Remove this once we have proper report interface.

				// mangle
				// https://gubernator.k8s.io/build/gob-prow/pr-logs/pull/some/repo/8940/pull-test-infra-presubmit//
				// to
				// https://gubernator.k8s.io/builds/gob-prow/pr-logs/pull/some_repo/8940/pull-test-infra-presubmit/
				url = b.String()
				url = strings.Replace(url, "build", "builds", 1)
				// TODO(krzyzacy): gerrit path can be foo.googlesource.com/bar/baz, which means we took bar/baz as the repo
				// we are mangling the path in bootstrap.py, we need to handle this better in podutils
				url = strings.Replace(url, change.Project, strings.Replace(change.Project, "/", "_", -1), 1)
				url = strings.TrimSuffix(url, "//")
			}
			triggeredJobs = append(triggeredJobs, triggeredJob{Name: spec.Name, URL: url})
		}
	}

	if len(triggeredJobs) > 0 {
		// comment back to gerrit
		message := "Triggered presubmit:"
		for _, job := range triggeredJobs {
			if job.URL != "" {
				message += fmt.Sprintf("\n  * Name: %s, URL: %s", job.Name, job.URL)
			} else {
				message += fmt.Sprintf("\n  * Name: %s", job.Name)
			}
		}

		if _, _, err := c.gc.SetReview(change.ID, change.CurrentRevision, &gerrit.ReviewInput{
			Message: message,
		}); err != nil {
			return fmt.Errorf("cannot comment to gerrit: %v", err)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	flag "github.com/ogier/pflag"
	"io/ioutil"
	"os"
	"strings"
	"text/template"
)

const helpStr = `    jgtr - JSON Go Template Renderer

USAGE:
    jgtr [OPTIONS]

    jgtr consumes a data file and a template file written in Go's text/template
    language. The values available in the data file are then used to render the
    template file and generate output.

    By default, jgtr reads the data and template from stdin, and writes output
    to stdout. Note that data and template cannot both come from stdin - at
    least one of the two must be specified via an option.

    jgtr can consume data from JSON, YAML 1.1 or TOML v0.2.0. You can specify
    which type to use via an option. If no such option is given, jgtr attempts
    to guess from the extension of the data file (if any). If the format is
    still ambiguous, jgtr uses JSON as the default.

OPTIONS:
    -d FILE, --data=FILE
        Read data data from FILE. Specify "-" (the default) to use stdin.

    -t FILE, --template=FILE
        Read template from FILE. Specify "-" (the default) to use stdin.

    -o FILE, --output=FILE
        Write rendered template to FILE. Specify "-" (the default) to use
        stdout.

    -j, --json
    	Specify the data format as JSON (default).

    -y, --yaml
    	Specify the data format as YAML.

    -T, --toml
    	Specify the data format as TOML.

    -h, --help
        Display this help.

    -V, --version
        Display jgtr version.`

const versionStr = `0.7.0`

func main() {
	help := flag.BoolP("help", "h", false, "show help")
	version := flag.BoolP("version", "V", false, "show version")

	dataPath := flag.StringP("data", "d", "-", "data file (JSON by default)")
	tmplPath := flag.StringP("template", "t", "-", "Go template file")
	outPath := flag.StringP("output", "o", "-", "output file")

	jsonFlag := flag.BoolP("json", "j", false, "interpret data as JSON")
	yamlFlag := flag.BoolP("yaml", "y", false, "interpret data as YAML")
	tomlFlag := flag.BoolP("toml", "T", false, "interpret data as TOML")

	flag.Parse()

	if *help {
		fmt.Println(helpStr)
		return
	}
	if *version {
		println(versionStr)
		return
	}

	if *dataPath == "-" && *tmplPath == "-" {
		println("Cannot use stdin for data and template simultaneously")
		os.Exit(1)
	}

	// no flag set - check extensions
	if !(*jsonFlag || *yamlFlag || *tomlFlag) {
		if strings.HasSuffix(*dataPath, ".yaml") || strings.HasSuffix(*dataPath, ".yml") {
			flag.Set("yaml", "true")
		} else if strings.HasSuffix(*dataPath, ".toml") {
			flag.Set("toml", "true")
		} else {
			flag.Set("json", "true")
		}
	}

	var data interface{} = nil
	var err error = nil
	if *yamlFlag {
		data, err = loadYAMLData(*dataPath)
	} else if *tomlFlag {
		data, err = loadTOMLData(*dataPath)
	} else {
		data, err = loadJSONData(*dataPath)
	}
	if err != nil {
		panic(err)
	}

	tmpl, err := loadGoTemplate(*tmplPath)
	if err != nil {
		panic(err)
	}

	outFile, err := createStream(*outPath)
	if err != nil {
		panic(err)
	}
	defer closeStream(outFile)

	err = tmpl.Execute(outFile, data)
	if err != nil {
		panic(err)
	}
}

// loadGoTemplate parses a Go text template from the file specified by path.
// The file contents are parsed into a top-level template with the name "root".
// If the path is "-", then the template will be parsed from os.Stdin.
func loadGoTemplate(path string) (tmpl *template.Template, err error) {
	file, err := openStream(path)
	if err != nil {
		return
	}
	defer closeStream(file)

	rawTmpl, err := ioutil.ReadAll(file)
	if err != nil {
		return
	}

	// explicitly parse the raw string rather than using ParseFiles
	// ParseFiles creates templates whose names are those of the files
	// and associates them with the parent template
	// this creates some confusing behavior in Template.Parse
	// see http://stackoverflow.com/questions/11805356/text-template-issue-parse-vs-parsefiles
	// also note: functions have to be added before the template is parsed
	tmpl, err = template.New("root").Funcs(tmplFuncs).Parse(string(rawTmpl))
	return // again, whether err==nil or not, this is finished
}

// openStream behaves like os.Open, except that if the path is "-", then it
// simply returns os.Stdin.
func openStream(path string) (file *os.File, err error) {
	if path == "-" {
		file = os.Stdin
	} else {
		file, err = os.Open(path)
	}
	return
}

// createStream behaves like os.Create, except that if the path is "-", then it
// simply returns os.Stdout.
func createStream(path string) (file *os.File, err error) {
	if path == "-" {
		file = os.Stdout
	} else {
		file, err = os.Create(path)
	}
	return
}

// closeStream behaves like file.Close, except that if the file is os.Stdin or
// os.Stdout, it does nothing.
func closeStream(file *os.File) (err error) {
	if file == os.Stdout || file == os.Stdin {
		return
	}
	return file.Close()
}
<commit_msg>Add error detection if multiple formats selected<commit_after>package main

import (
	"fmt"
	flag "github.com/ogier/pflag"
	"io/ioutil"
	"os"
	"strings"
	"text/template"
)

const helpStr = `    jgtr - JSON Go Template Renderer

USAGE:
    jgtr [OPTIONS]

    jgtr consumes a data file and a template file written in Go's text/template
    language. The values available in the data file are then used to render the
    template file and generate output.

    By default, jgtr reads the data and template from stdin, and writes output
    to stdout. Note that data and template cannot both come from stdin - at
    least one of the two must be specified via an option.

    jgtr can consume data from JSON, YAML 1.1 or TOML v0.2.0. You can specify
    which type to use via an option. If no such option is given, jgtr attempts
    to guess from the extension of the data file (if any). If the format is
    still ambiguous, jgtr uses JSON as the default.

OPTIONS:
    -d FILE, --data=FILE
        Read data data from FILE. Specify "-" (the default) to use stdin.

    -t FILE, --template=FILE
        Read template from FILE. Specify "-" (the default) to use stdin.

    -o FILE, --output=FILE
        Write rendered template to FILE. Specify "-" (the default) to use
        stdout.

    -j, --json
    	Specify the data format as JSON (default).

    -y, --yaml
    	Specify the data format as YAML.

    -T, --toml
    	Specify the data format as TOML.

    -h, --help
        Display this help.

    -V, --version
        Display jgtr version.`

const versionStr = `0.7.1`

func main() {
	help := flag.BoolP("help", "h", false, "show help")
	version := flag.BoolP("version", "V", false, "show version")

	dataPath := flag.StringP("data", "d", "-", "data file (JSON by default)")
	tmplPath := flag.StringP("template", "t", "-", "Go template file")
	outPath := flag.StringP("output", "o", "-", "output file")

	jsonFlag := flag.BoolP("json", "j", false, "interpret data as JSON")
	yamlFlag := flag.BoolP("yaml", "y", false, "interpret data as YAML")
	tomlFlag := flag.BoolP("toml", "T", false, "interpret data as TOML")

	flag.Parse()

	if *help {
		fmt.Println(helpStr)
		return
	}
	if *version {
		println(versionStr)
		return
	}

	if *dataPath == "-" && *tmplPath == "-" {
		println("Cannot use stdin for data and template simultaneously")
		os.Exit(1)
	}

	if (*jsonFlag && *yamlFlag) ||
		(*jsonFlag && *tomlFlag) ||
		(*yamlFlag && *tomlFlag) {
		println("Cannot select multiple data formats simultaneously")
		os.Exit(1)
	}

	// no flag set - check extensions
	if !(*jsonFlag || *yamlFlag || *tomlFlag) {
		if strings.HasSuffix(*dataPath, ".yaml") || strings.HasSuffix(*dataPath, ".yml") {
			flag.Set("yaml", "true")
		} else if strings.HasSuffix(*dataPath, ".toml") {
			flag.Set("toml", "true")
		} else {
			flag.Set("json", "true")
		}
	}

	var data interface{} = nil
	var err error = nil
	if *yamlFlag {
		data, err = loadYAMLData(*dataPath)
	} else if *tomlFlag {
		data, err = loadTOMLData(*dataPath)
	} else {
		data, err = loadJSONData(*dataPath)
	}
	if err != nil {
		panic(err)
	}

	tmpl, err := loadGoTemplate(*tmplPath)
	if err != nil {
		panic(err)
	}

	outFile, err := createStream(*outPath)
	if err != nil {
		panic(err)
	}
	defer closeStream(outFile)

	err = tmpl.Execute(outFile, data)
	if err != nil {
		panic(err)
	}
}

// loadGoTemplate parses a Go text template from the file specified by path.
// The file contents are parsed into a top-level template with the name "root".
// If the path is "-", then the template will be parsed from os.Stdin.
func loadGoTemplate(path string) (tmpl *template.Template, err error) {
	file, err := openStream(path)
	if err != nil {
		return
	}
	defer closeStream(file)

	rawTmpl, err := ioutil.ReadAll(file)
	if err != nil {
		return
	}

	// explicitly parse the raw string rather than using ParseFiles
	// ParseFiles creates templates whose names are those of the files
	// and associates them with the parent template
	// this creates some confusing behavior in Template.Parse
	// see http://stackoverflow.com/questions/11805356/text-template-issue-parse-vs-parsefiles
	// also note: functions have to be added before the template is parsed
	tmpl, err = template.New("root").Funcs(tmplFuncs).Parse(string(rawTmpl))
	return // again, whether err==nil or not, this is finished
}

// openStream behaves like os.Open, except that if the path is "-", then it
// simply returns os.Stdin.
func openStream(path string) (file *os.File, err error) {
	if path == "-" {
		file = os.Stdin
	} else {
		file, err = os.Open(path)
	}
	return
}

// createStream behaves like os.Create, except that if the path is "-", then it
// simply returns os.Stdout.
func createStream(path string) (file *os.File, err error) {
	if path == "-" {
		file = os.Stdout
	} else {
		file, err = os.Create(path)
	}
	return
}

// closeStream behaves like file.Close, except that if the file is os.Stdin or
// os.Stdout, it does nothing.
func closeStream(file *os.File) (err error) {
	if file == os.Stdout || file == os.Stdin {
		return
	}
	return file.Close()
}
<|endoftext|>


<commit_before>/*
Test out the AHRS code in ahrs/ahrs.go.
Define a flight path/attitude in code, and then synthesize the matching GPS, gyro, accel (and other) data
Add some noise if desired.
Then see if the AHRS code can replicate the "true" attitude given the noisy and limited input data
*/

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"
	"strings"

	"../ahrs"
)

func parseFloatArrayString(str string, a *[]float64) (err error) {
	for i, s := range strings.Split(str, ",") {
		(*a)[i], err = strconv.ParseFloat(s, 64)
		if err != nil {
			break
		}
	}
	return
}

func main() {
	// Handle some shell arguments
	var (
		pdt, udt                                            float64
		gyroBiasStr, accelBiasStr, magBiasStr               string
		gyroNoise, accelNoise, gpsNoise, asiNoise, magNoise float64
		asiBias                                             float64
		gyroBias, accelBias, magBias                        []float64
		gpsInop, magInop, asiInop                           bool
		algo                                                string
		s                                                   ahrs.AHRSProvider
		scenario                                            string
		sit                                                 Situation
		err                                                 error
		ahrsLogger                                          *ahrs.AHRSLogger
	)
	gyroBias = make([]float64, 3)
	accelBias = make([]float64, 3)
	magBias = make([]float64, 3)

	const (
		defaultPdt        = 0.05
		pdtUsage          = "Kalman filter predict period, seconds"
		defaultUdt        = 0.05
		udtUsage          = "Kalman filter update period, seconds"
		defaultGyroNoise  = 0.0
		gyroNoiseUsage    = "Amount of noise to add to gyro measurements, °/s"
		defaultGyroBias   = "0,0,0"
		gyroBiasUsage     = "Amount of bias to add to gyro measurements, \"x,y,z\" °/s"
		defaultAccelNoise = 0.0
		accelNoiseUsage   = "Amount of noise to add to accel measurements, G"
		defaultAccelBias  = "0,0,0"
		accelBiasUsage    = "Amount of bias to add to accel measurements, \"x,y,z\" G"
		defaultGPSNoise   = 0.0
		gpsNoiseUsage     = "Amount of noise to add to GPS speed measurements, kt"
		defaultASINoise   = 0.0
		asiNoiseUsage     = "Amount of noise to add to airspeed measurements, kt"
		defaultASIBias    = 0.0
		asiBiasUsage      = "Amount of bias to add to airspeed measurements, kt"
		defaultMagNoise   = 0.0
		magNoiseUsage     = "Amount of noise to add to magnetometer measurements, μT"
		defaultMagBias    = "0,0,0"
		magBiasUsage      = "Amount of bias to add to magnetometer measurements, \"x,y,z\" μT"
		defaultGPSInop    = false
		gpsInopUsage      = "Make the GPS inoperative"
		defaultASIInop    = true
		asiInopUsage      = "Make the Airspeed sensor inoperative"
		defaultMagInop    = false
		magInopUsage      = "Make the Magnetometer inoperative"
		defaultScenario   = "takeoff"
		scenarioUsage     = "Scenario to use: filename or \"takeoff\" or \"turn\""
		defaultAlgo       = "simple"
		algoUsage         = "Algo to use for AHRS: simple (default), heuristic, kalman, kalman1, kalman2"
	)

	flag.Float64Var(&pdt, "pdt", defaultPdt, pdtUsage)
	flag.Float64Var(&udt, "udt", defaultUdt, udtUsage)
	flag.Float64Var(&gyroNoise, "gyro-noise", defaultGyroNoise, gyroNoiseUsage)
	flag.Float64Var(&gyroNoise, "g", defaultGyroNoise, gyroNoiseUsage)
	flag.StringVar(&gyroBiasStr, "gyro-bias", defaultGyroBias, gyroBiasUsage)
	flag.StringVar(&gyroBiasStr, "h", defaultGyroBias, gyroBiasUsage)
	flag.Float64Var(&accelNoise, "accel-noise", defaultAccelNoise, accelNoiseUsage)
	flag.Float64Var(&accelNoise, "a", defaultAccelNoise, accelNoiseUsage)
	flag.StringVar(&accelBiasStr, "accel-bias", defaultAccelBias, accelBiasUsage)
	flag.StringVar(&accelBiasStr, "i", defaultAccelBias, accelBiasUsage)
	flag.Float64Var(&gpsNoise, "gps-noise", defaultGPSNoise, gpsNoiseUsage)
	flag.Float64Var(&gpsNoise, "n", defaultGPSNoise, gpsNoiseUsage)
	flag.Float64Var(&asiNoise, "asi-noise", defaultASINoise, asiNoiseUsage)
	flag.Float64Var(&asiNoise, "v", defaultASINoise, asiNoiseUsage)
	flag.Float64Var(&asiBias, "asi-bias", defaultASIBias, asiBiasUsage)
	flag.Float64Var(&asiBias, "j", defaultASIBias, asiBiasUsage)
	flag.Float64Var(&magNoise, "mag-noise", defaultMagNoise, magNoiseUsage)
	flag.Float64Var(&magNoise, "b", defaultMagNoise, magNoiseUsage)
	flag.StringVar(&magBiasStr, "mag-bias", defaultMagBias, magBiasUsage)
	flag.StringVar(&magBiasStr, "k", defaultMagBias, magBiasUsage)
	flag.BoolVar(&gpsInop, "w", defaultGPSInop, gpsInopUsage)
	flag.BoolVar(&asiInop, "u", defaultASIInop, asiInopUsage)
	flag.BoolVar(&magInop, "m", defaultMagInop, magInopUsage)
	flag.StringVar(&scenario, "scenario", defaultScenario, scenarioUsage)
	flag.StringVar(&scenario, "s", defaultScenario, scenarioUsage)
	flag.StringVar(&algo, "algo", defaultAlgo, algoUsage)
	flag.Parse()

	switch scenario {
	/*
	case "takeoff":
		sit = sitTakeoffDef
	case "turn":
		sit = sitTurnDef
	*/
	default:
		log.Printf("Loading data from %s\n", scenario)
		sit, err = NewSituationFromFile(scenario)
		if err != nil {
			log.Fatalln(err)
		}
	}

	s0 := new(ahrs.State)      // Actual state from simulation, for comparison
	m := ahrs.NewMeasurement() // Measurement from IMU

	fmt.Println("Simulation parameters:")
	switch strings.ToLower(algo) {
	case "simple":
		fmt.Println("Running simple AHRS")
		ioutil.WriteFile("config.json", []byte(ahrs.SimpleJSONConfig), 0644)
		s = ahrs.InitializeSimple()
	/*
	case "kalman":
		fmt.Println("Running Kalman AHRS")
		ioutil.WriteFile("config.json", []byte(ahrs.KalmanJSONConfig), 0644)
		s = ahrs.InitializeKalman(m)
	case "kalman1":
		fmt.Println("Running Kalman1 AHRS")
		ioutil.WriteFile("config.json", []byte(ahrs.Kalman1JSONConfig), 0644)
		s = ahrs.InitializeKalman1(m)
	*/
	default:
		fmt.Printf("No such AHRS implementation: %s\n", algo)
		return
	}

	if err := parseFloatArrayString(gyroBiasStr, &gyroBias); err != nil {
		fmt.Printf("Error %v parsing %s\n", err, gyroBiasStr)
		return
	}
	if err := parseFloatArrayString(accelBiasStr, &accelBias); err != nil {
		fmt.Printf("Error %v parsing %s\n", err, accelBiasStr)
		return
	}
	if err := parseFloatArrayString(magBiasStr, &magBias); err != nil {
		fmt.Printf("Error %v parsing %s\n", err, magBiasStr)
		return
	}

	fmt.Println("Timing:")
	fmt.Printf("\tPredict Freqency: %d Hz\n", int(1/pdt))
	fmt.Printf("\tUpdate  Freqency: %d Hz\n", int(1/udt))
	fmt.Println("Accelerometer:")
	fmt.Printf("\tNoise: %f G\n", accelNoise)
	fmt.Printf("\tBias: %f,%f,%f\n", accelBias[0], accelBias[1], accelBias[2])
	fmt.Println("Gyro:")
	fmt.Printf("\tNoise: %f °/s\n", gyroNoise)
	fmt.Printf("\tBias: %f,%f,%f\n", gyroBias[0], gyroBias[1], gyroBias[2])
	fmt.Println("GPS:")
	fmt.Printf("\tInop: %t\n", gpsInop)
	fmt.Printf("\tNoise: %f kt\n", gpsNoise)
	fmt.Println("ASI:")
	fmt.Printf("\tInop: %t\n", asiInop)
	fmt.Printf("\tNoise: %f kt\n", asiNoise)
	fmt.Println("Magnetometer:")
	fmt.Printf("\tInop: %t\n", magInop)
	fmt.Printf("\tNoise: %f G\n", magNoise)
	fmt.Printf("\tBias: %f,%f,%f\n", magBias[0], magBias[1], magBias[2])

	uBias := []float64{asiBias, 0, 0}

	// Set up logging
	ahrsLogger = ahrs.NewAHRSLogger("ahrs.csv", s.GetLogMap())

	// This is where it all happens
	fmt.Println("Running Simulation")
	sit.BeginTime()
	sit.UpdateMeasurement(m, !asiInop, !gpsInop, true, !magInop,
		asiNoise, gpsNoise, accelNoise, gyroNoise, magNoise,
		uBias, accelBias, gyroBias, magBias)

	for {
		// Peek behind the curtain: the "actual" state, which the algorithm doesn't know
		if err := sit.UpdateState(s0, accelBias, gyroBias, magBias); err != nil {
			log.Printf("Interpolation error at time %f: %s\n", m.T, err)
			break
		}
		//TODO westphae: log actual state

		// Take sensor measurements
		if err := sit.UpdateMeasurement(m, !asiInop, !gpsInop, true, !magInop,
			asiNoise, gpsNoise, accelNoise, gyroNoise, magNoise,
			uBias, accelBias, gyroBias, magBias); err != nil {
			log.Printf("Measurement error at time %f: %s\n", m.T, err)
			break
		}

		s.Compute(m)
		ahrsLogger.Log() // Log to csv for serving

		err = sit.NextTime()
		if err != nil {
			log.Println(err)
			break
		}

	}

	// Run analysis web server
	fmt.Println("Serving charts")
	http.Handle("/", http.FileServer(http.Dir("./")))
	http.ListenAndServe(":8080", nil)
}
<commit_msg>Allow sim to change ahrs config.<commit_after>/*
Test out the AHRS code in ahrs/ahrs.go.
Define a flight path/attitude in code, and then synthesize the matching GPS, gyro, accel (and other) data
Add some noise if desired.
Then see if the AHRS code can replicate the "true" attitude given the noisy and limited input data
*/

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"
	"strings"

	"../ahrs"
	"encoding/json"
)

func parseFloatArrayString(str string, a *[]float64) (err error) {
	for i, s := range strings.Split(str, ",") {
		(*a)[i], err = strconv.ParseFloat(s, 64)
		if err != nil {
			break
		}
	}
	return
}

func main() {
	// Handle some shell arguments
	var (
		pdt, udt                                            float64
		gyroBiasStr, accelBiasStr, magBiasStr               string
		gyroNoise, accelNoise, gpsNoise, asiNoise, magNoise float64
		asiBias                                             float64
		gyroBias, accelBias, magBias                        []float64
		gpsInop, magInop, asiInop                           bool
		algo                                                string
		ahrsConfigStr                                       string
		ahrsConfig                                          map[string]float64
		s                                                   ahrs.AHRSProvider
		scenario                                            string
		sit                                                 Situation
		err                                                 error
		ahrsLogger                                          *ahrs.AHRSLogger
	)
	gyroBias = make([]float64, 3)
	accelBias = make([]float64, 3)
	magBias = make([]float64, 3)

	const (
		defaultPdt        = 0.05
		pdtUsage          = "Kalman filter predict period, seconds"
		defaultUdt        = 0.05
		udtUsage          = "Kalman filter update period, seconds"
		defaultGyroNoise  = 0.0
		gyroNoiseUsage    = "Amount of noise to add to gyro measurements, °/s"
		defaultGyroBias   = "0,0,0"
		gyroBiasUsage     = "Amount of bias to add to gyro measurements, \"x,y,z\" °/s"
		defaultAccelNoise = 0.0
		accelNoiseUsage   = "Amount of noise to add to accel measurements, G"
		defaultAccelBias  = "0,0,0"
		accelBiasUsage    = "Amount of bias to add to accel measurements, \"x,y,z\" G"
		defaultGPSNoise   = 0.0
		gpsNoiseUsage     = "Amount of noise to add to GPS speed measurements, kt"
		defaultASINoise   = 0.0
		asiNoiseUsage     = "Amount of noise to add to airspeed measurements, kt"
		defaultASIBias    = 0.0
		asiBiasUsage      = "Amount of bias to add to airspeed measurements, kt"
		defaultMagNoise   = 0.0
		magNoiseUsage     = "Amount of noise to add to magnetometer measurements, μT"
		defaultMagBias    = "0,0,0"
		magBiasUsage      = "Amount of bias to add to magnetometer measurements, \"x,y,z\" μT"
		defaultGPSInop    = false
		gpsInopUsage      = "Make the GPS inoperative"
		defaultASIInop    = true
		asiInopUsage      = "Make the Airspeed sensor inoperative"
		defaultMagInop    = false
		magInopUsage      = "Make the Magnetometer inoperative"
		defaultScenario   = "takeoff"
		scenarioUsage     = "Scenario to use: filename or \"takeoff\" or \"turn\""
		defaultAlgo       = "simple"
		algoUsage         = "Algo to use for AHRS: simple (default), heuristic, kalman, kalman1, kalman2"
		defaultConfig     = ""
		configUsage       = "json-formatted map for AHRS Config"
	)

	flag.Float64Var(&pdt, "pdt", defaultPdt, pdtUsage)
	flag.Float64Var(&udt, "udt", defaultUdt, udtUsage)
	flag.Float64Var(&gyroNoise, "gyro-noise", defaultGyroNoise, gyroNoiseUsage)
	flag.Float64Var(&gyroNoise, "g", defaultGyroNoise, gyroNoiseUsage)
	flag.StringVar(&gyroBiasStr, "gyro-bias", defaultGyroBias, gyroBiasUsage)
	flag.StringVar(&gyroBiasStr, "h", defaultGyroBias, gyroBiasUsage)
	flag.Float64Var(&accelNoise, "accel-noise", defaultAccelNoise, accelNoiseUsage)
	flag.Float64Var(&accelNoise, "a", defaultAccelNoise, accelNoiseUsage)
	flag.StringVar(&accelBiasStr, "accel-bias", defaultAccelBias, accelBiasUsage)
	flag.StringVar(&accelBiasStr, "i", defaultAccelBias, accelBiasUsage)
	flag.Float64Var(&gpsNoise, "gps-noise", defaultGPSNoise, gpsNoiseUsage)
	flag.Float64Var(&gpsNoise, "n", defaultGPSNoise, gpsNoiseUsage)
	flag.Float64Var(&asiNoise, "asi-noise", defaultASINoise, asiNoiseUsage)
	flag.Float64Var(&asiNoise, "v", defaultASINoise, asiNoiseUsage)
	flag.Float64Var(&asiBias, "asi-bias", defaultASIBias, asiBiasUsage)
	flag.Float64Var(&asiBias, "j", defaultASIBias, asiBiasUsage)
	flag.Float64Var(&magNoise, "mag-noise", defaultMagNoise, magNoiseUsage)
	flag.Float64Var(&magNoise, "b", defaultMagNoise, magNoiseUsage)
	flag.StringVar(&magBiasStr, "mag-bias", defaultMagBias, magBiasUsage)
	flag.StringVar(&magBiasStr, "k", defaultMagBias, magBiasUsage)
	flag.BoolVar(&gpsInop, "w", defaultGPSInop, gpsInopUsage)
	flag.BoolVar(&asiInop, "u", defaultASIInop, asiInopUsage)
	flag.BoolVar(&magInop, "m", defaultMagInop, magInopUsage)
	flag.StringVar(&scenario, "scenario", defaultScenario, scenarioUsage)
	flag.StringVar(&scenario, "s", defaultScenario, scenarioUsage)
	flag.StringVar(&algo, "algo", defaultAlgo, algoUsage)
	flag.StringVar(&ahrsConfigStr, "config", defaultConfig, configUsage)
	flag.StringVar(&ahrsConfigStr, "c", defaultConfig, configUsage)
	flag.Parse()

	switch scenario {
	/*
	case "takeoff":
		sit = sitTakeoffDef
	case "turn":
		sit = sitTurnDef
	*/
	default:
		log.Printf("Loading data from %s\n", scenario)
		sit, err = NewSituationFromFile(scenario)
		if err != nil {
			log.Fatalln(err)
		}
	}

	s0 := new(ahrs.State)      // Actual state from simulation, for comparison
	m := ahrs.NewMeasurement() // Measurement from IMU

	fmt.Println("Simulation parameters:")
	switch strings.ToLower(algo) {
	/*
	case "kalman":
		fmt.Println("Running Kalman AHRS")
		ioutil.WriteFile("config.json", []byte(ahrs.KalmanJSONConfig), 0644)
		s = ahrs.InitializeKalman(m)
	case "kalman1":
		fmt.Println("Running Kalman1 AHRS")
		ioutil.WriteFile("config.json", []byte(ahrs.Kalman1JSONConfig), 0644)
		s = ahrs.InitializeKalman1(m)
	*/
	case "simple":
		fallthrough // simple is the default.
	default:
		fmt.Println("Running simple AHRS")
		ioutil.WriteFile("config.json", []byte(ahrs.SimpleJSONConfig), 0644)
		s = ahrs.InitializeSimple()
	}

	if err := parseFloatArrayString(gyroBiasStr, &gyroBias); err != nil {
		fmt.Printf("Error %v parsing %s\n", err, gyroBiasStr)
		return
	}
	if err := parseFloatArrayString(accelBiasStr, &accelBias); err != nil {
		fmt.Printf("Error %v parsing %s\n", err, accelBiasStr)
		return
	}
	if err := parseFloatArrayString(magBiasStr, &magBias); err != nil {
		fmt.Printf("Error %v parsing %s\n", err, magBiasStr)
		return
	}

	fmt.Println("Timing:")
	fmt.Printf("\tPredict Freqency: %d Hz\n", int(1/pdt))
	fmt.Printf("\tUpdate  Freqency: %d Hz\n", int(1/udt))
	fmt.Println("Accelerometer:")
	fmt.Printf("\tNoise: %f G\n", accelNoise)
	fmt.Printf("\tBias: %f,%f,%f\n", accelBias[0], accelBias[1], accelBias[2])
	fmt.Println("Gyro:")
	fmt.Printf("\tNoise: %f °/s\n", gyroNoise)
	fmt.Printf("\tBias: %f,%f,%f\n", gyroBias[0], gyroBias[1], gyroBias[2])
	fmt.Println("GPS:")
	fmt.Printf("\tInop: %t\n", gpsInop)
	fmt.Printf("\tNoise: %f kt\n", gpsNoise)
	fmt.Println("ASI:")
	fmt.Printf("\tInop: %t\n", asiInop)
	fmt.Printf("\tNoise: %f kt\n", asiNoise)
	fmt.Println("Magnetometer:")
	fmt.Printf("\tInop: %t\n", magInop)
	fmt.Printf("\tNoise: %f G\n", magNoise)
	fmt.Printf("\tBias: %f,%f,%f\n", magBias[0], magBias[1], magBias[2])

	uBias := []float64{asiBias, 0, 0}

	if err := json.Unmarshal([]byte(ahrsConfigStr), &ahrsConfig); err != nil {
		log.Printf("Bad config: %s\n", err.Error())
	}
	log.Printf("ahrs config: %v\n", ahrsConfig)
	s.SetConfig(ahrsConfig)

	// Set up logging
	ahrsLogger = ahrs.NewAHRSLogger("ahrs.csv", s.GetLogMap())

	// This is where it all happens
	fmt.Println("Running Simulation")
	sit.BeginTime()
	sit.UpdateMeasurement(m, !asiInop, !gpsInop, true, !magInop,
		asiNoise, gpsNoise, accelNoise, gyroNoise, magNoise,
		uBias, accelBias, gyroBias, magBias)

	for {
		// Peek behind the curtain: the "actual" state, which the algorithm doesn't know
		if err := sit.UpdateState(s0, accelBias, gyroBias, magBias); err != nil {
			log.Printf("Interpolation error at time %f: %s\n", m.T, err)
			break
		}
		//TODO westphae: log actual state

		// Take sensor measurements
		if err := sit.UpdateMeasurement(m, !asiInop, !gpsInop, true, !magInop,
			asiNoise, gpsNoise, accelNoise, gyroNoise, magNoise,
			uBias, accelBias, gyroBias, magBias); err != nil {
			log.Printf("Measurement error at time %f: %s\n", m.T, err)
			break
		}

		s.Compute(m)
		ahrsLogger.Log() // Log to csv for serving

		err = sit.NextTime()
		if err != nil {
			log.Println(err)
			break
		}

	}

	// Run analysis web server
	fmt.Println("Serving charts")
	http.Handle("/", http.FileServer(http.Dir("./")))
	http.ListenAndServe(":8080", nil)
}
<|endoftext|>


<commit_before>package registry

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"path"
	"regexp"
	"strings"
	"time"

	"github.com/docker/docker/utils"
)

var (
	ErrAlreadyExists         = errors.New("Image already exists")
	ErrInvalidRepositoryName = errors.New("Invalid repository name (ex: \"registry.domain.tld/myrepos\")")
	ErrDoesNotExist          = errors.New("Image does not exist")
	errLoginRequired         = errors.New("Authentication is required.")
	validHex                 = regexp.MustCompile(`^([a-f0-9]{64})$`)
	validNamespace           = regexp.MustCompile(`^([a-z0-9_]{4,30})$`)
	validRepo                = regexp.MustCompile(`^([a-z0-9-_.]+)$`)
)

type TimeoutType uint32

const (
	NoTimeout TimeoutType = iota
	ReceiveTimeout
	ConnectTimeout
)

func newClient(jar http.CookieJar, roots *x509.CertPool, cert *tls.Certificate, timeout TimeoutType) *http.Client {
	tlsConfig := tls.Config{
		RootCAs: roots,
		// Avoid fallback to SSL protocols < TLS1.0
		MinVersion: tls.VersionTLS10,
	}

	if cert != nil {
		tlsConfig.Certificates = append(tlsConfig.Certificates, *cert)
	}

	httpTransport := &http.Transport{
		DisableKeepAlives: true,
		Proxy:             http.ProxyFromEnvironment,
		TLSClientConfig:   &tlsConfig,
	}

	switch timeout {
	case ConnectTimeout:
		httpTransport.Dial = func(proto string, addr string) (net.Conn, error) {
			// Set the connect timeout to 5 seconds
			conn, err := net.DialTimeout(proto, addr, 5*time.Second)
			if err != nil {
				return nil, err
			}
			// Set the recv timeout to 10 seconds
			conn.SetDeadline(time.Now().Add(10 * time.Second))
			return conn, nil
		}
	case ReceiveTimeout:
		httpTransport.Dial = func(proto string, addr string) (net.Conn, error) {
			conn, err := net.Dial(proto, addr)
			if err != nil {
				return nil, err
			}
			conn = utils.NewTimeoutConn(conn, 1*time.Minute)
			return conn, nil
		}
	}

	return &http.Client{
		Transport:     httpTransport,
		CheckRedirect: AddRequiredHeadersToRedirectedRequests,
		Jar:           jar,
	}
}

func doRequest(req *http.Request, jar http.CookieJar, timeout TimeoutType) (*http.Response, *http.Client, error) {
	hasFile := func(files []os.FileInfo, name string) bool {
		for _, f := range files {
			if f.Name() == name {
				return true
			}
		}
		return false
	}

	hostDir := path.Join("/etc/docker/certs.d", req.URL.Host)
	fs, err := ioutil.ReadDir(hostDir)
	if err != nil && !os.IsNotExist(err) {
		return nil, nil, err
	}

	var (
		pool  *x509.CertPool
		certs []*tls.Certificate
	)

	for _, f := range fs {
		if strings.HasSuffix(f.Name(), ".crt") {
			if pool == nil {
				pool = x509.NewCertPool()
			}
			data, err := ioutil.ReadFile(path.Join(hostDir, f.Name()))
			if err != nil {
				return nil, nil, err
			}
			pool.AppendCertsFromPEM(data)
		}
		if strings.HasSuffix(f.Name(), ".cert") {
			certName := f.Name()
			keyName := certName[:len(certName)-5] + ".key"
			if !hasFile(fs, keyName) {
				return nil, nil, fmt.Errorf("Missing key %s for certificate %s", keyName, certName)
			}
			cert, err := tls.LoadX509KeyPair(path.Join(hostDir, certName), path.Join(hostDir, keyName))
			if err != nil {
				return nil, nil, err
			}
			certs = append(certs, &cert)
		}
		if strings.HasSuffix(f.Name(), ".key") {
			keyName := f.Name()
			certName := keyName[:len(keyName)-4] + ".cert"
			if !hasFile(fs, certName) {
				return nil, nil, fmt.Errorf("Missing certificate %s for key %s", certName, keyName)
			}
		}
	}

	if len(certs) == 0 {
		client := newClient(jar, pool, nil, timeout)
		res, err := client.Do(req)
		if err != nil {
			return nil, nil, err
		}
		return res, client, nil
	}
	for i, cert := range certs {
		client := newClient(jar, pool, cert, timeout)
		res, err := client.Do(req)
		// If this is the last cert, otherwise, continue to next cert if 403 or 5xx
		if i == len(certs)-1 || err == nil && res.StatusCode != 403 && res.StatusCode < 500 {
			return res, client, err
		}
	}

	return nil, nil, nil
}

func validateRepositoryName(repositoryName string) error {
	var (
		namespace string
		name      string
	)
	nameParts := strings.SplitN(repositoryName, "/", 2)
	if len(nameParts) < 2 {
		namespace = "library"
		name = nameParts[0]

		if validHex.MatchString(name) {
			return fmt.Errorf("Invalid repository name (%s), cannot specify 64-byte hexadecimal strings", name)
		}
	} else {
		namespace = nameParts[0]
		name = nameParts[1]
	}
	if !validNamespace.MatchString(namespace) {
		return fmt.Errorf("Invalid namespace name (%s), only [a-z0-9_] are allowed, size between 4 and 30", namespace)
	}
	if !validRepo.MatchString(name) {
		return fmt.Errorf("Invalid repository name (%s), only [a-z0-9-_.] are allowed", name)
	}
	return nil
}

// Resolves a repository name to a hostname + name
func ResolveRepositoryName(reposName string) (string, string, error) {
	if strings.Contains(reposName, "://") {
		// It cannot contain a scheme!
		return "", "", ErrInvalidRepositoryName
	}
	nameParts := strings.SplitN(reposName, "/", 2)
	if len(nameParts) == 1 || (!strings.Contains(nameParts[0], ".") && !strings.Contains(nameParts[0], ":") &&
		nameParts[0] != "localhost") {
		// This is a Docker Index repos (ex: samalba/hipache or ubuntu)
		err := validateRepositoryName(reposName)
		return IndexServerAddress(), reposName, err
	}
	hostname := nameParts[0]
	reposName = nameParts[1]
	if strings.Contains(hostname, "index.docker.io") {
		return "", "", fmt.Errorf("Invalid repository name, try \"%s\" instead", reposName)
	}
	if err := validateRepositoryName(reposName); err != nil {
		return "", "", err
	}

	return hostname, reposName, nil
}

func trustedLocation(req *http.Request) bool {
	var (
		trusteds = []string{"docker.com", "docker.io"}
		hostname = strings.SplitN(req.Host, ":", 2)[0]
	)
	if req.URL.Scheme != "https" {
		return false
	}

	for _, trusted := range trusteds {
		if hostname == trusted || strings.HasSuffix(hostname, "."+trusted) {
			return true
		}
	}
	return false
}

func AddRequiredHeadersToRedirectedRequests(req *http.Request, via []*http.Request) error {
	if via != nil && via[0] != nil {
		if trustedLocation(req) && trustedLocation(via[0]) {
			req.Header = via[0].Header
			return nil
		}
		for k, v := range via[0].Header {
			if k != "Authorization" {
				for _, vv := range v {
					req.Header.Add(k, vv)
				}
			}
		}
	}
	return nil
}
<commit_msg>On Red Hat Registry Servers we return 404 on certification errors.<commit_after>package registry

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"path"
	"regexp"
	"strings"
	"time"

	"github.com/docker/docker/utils"
)

var (
	ErrAlreadyExists         = errors.New("Image already exists")
	ErrInvalidRepositoryName = errors.New("Invalid repository name (ex: \"registry.domain.tld/myrepos\")")
	ErrDoesNotExist          = errors.New("Image does not exist")
	errLoginRequired         = errors.New("Authentication is required.")
	validHex                 = regexp.MustCompile(`^([a-f0-9]{64})$`)
	validNamespace           = regexp.MustCompile(`^([a-z0-9_]{4,30})$`)
	validRepo                = regexp.MustCompile(`^([a-z0-9-_.]+)$`)
)

type TimeoutType uint32

const (
	NoTimeout TimeoutType = iota
	ReceiveTimeout
	ConnectTimeout
)

func newClient(jar http.CookieJar, roots *x509.CertPool, cert *tls.Certificate, timeout TimeoutType) *http.Client {
	tlsConfig := tls.Config{
		RootCAs: roots,
		// Avoid fallback to SSL protocols < TLS1.0
		MinVersion: tls.VersionTLS10,
	}

	if cert != nil {
		tlsConfig.Certificates = append(tlsConfig.Certificates, *cert)
	}

	httpTransport := &http.Transport{
		DisableKeepAlives: true,
		Proxy:             http.ProxyFromEnvironment,
		TLSClientConfig:   &tlsConfig,
	}

	switch timeout {
	case ConnectTimeout:
		httpTransport.Dial = func(proto string, addr string) (net.Conn, error) {
			// Set the connect timeout to 5 seconds
			conn, err := net.DialTimeout(proto, addr, 5*time.Second)
			if err != nil {
				return nil, err
			}
			// Set the recv timeout to 10 seconds
			conn.SetDeadline(time.Now().Add(10 * time.Second))
			return conn, nil
		}
	case ReceiveTimeout:
		httpTransport.Dial = func(proto string, addr string) (net.Conn, error) {
			conn, err := net.Dial(proto, addr)
			if err != nil {
				return nil, err
			}
			conn = utils.NewTimeoutConn(conn, 1*time.Minute)
			return conn, nil
		}
	}

	return &http.Client{
		Transport:     httpTransport,
		CheckRedirect: AddRequiredHeadersToRedirectedRequests,
		Jar:           jar,
	}
}

func doRequest(req *http.Request, jar http.CookieJar, timeout TimeoutType) (*http.Response, *http.Client, error) {
	hasFile := func(files []os.FileInfo, name string) bool {
		for _, f := range files {
			if f.Name() == name {
				return true
			}
		}
		return false
	}

	hostDir := path.Join("/etc/docker/certs.d", req.URL.Host)
	fs, err := ioutil.ReadDir(hostDir)
	if err != nil && !os.IsNotExist(err) {
		return nil, nil, err
	}

	var (
		pool  *x509.CertPool
		certs []*tls.Certificate
	)

	for _, f := range fs {
		if strings.HasSuffix(f.Name(), ".crt") {
			if pool == nil {
				pool = x509.NewCertPool()
			}
			data, err := ioutil.ReadFile(path.Join(hostDir, f.Name()))
			if err != nil {
				return nil, nil, err
			}
			pool.AppendCertsFromPEM(data)
		}
		if strings.HasSuffix(f.Name(), ".cert") {
			certName := f.Name()
			keyName := certName[:len(certName)-5] + ".key"
			if !hasFile(fs, keyName) {
				return nil, nil, fmt.Errorf("Missing key %s for certificate %s", keyName, certName)
			}
			cert, err := tls.LoadX509KeyPair(path.Join(hostDir, certName), path.Join(hostDir, keyName))
			if err != nil {
				return nil, nil, err
			}
			certs = append(certs, &cert)
		}
		if strings.HasSuffix(f.Name(), ".key") {
			keyName := f.Name()
			certName := keyName[:len(keyName)-4] + ".cert"
			if !hasFile(fs, certName) {
				return nil, nil, fmt.Errorf("Missing certificate %s for key %s", certName, keyName)
			}
		}
	}

	if len(certs) == 0 {
		client := newClient(jar, pool, nil, timeout)
		res, err := client.Do(req)
		if err != nil {
			return nil, nil, err
		}
		return res, client, nil
	}
	for i, cert := range certs {
		client := newClient(jar, pool, cert, timeout)
		res, err := client.Do(req)
		// If this is the last cert, otherwise, continue to next cert if 403 or 5xx
		if i == len(certs)-1 || err == nil &&
			res.StatusCode != 403 &&
			res.StatusCode != 404 &&
			res.StatusCode < 500 {
			return res, client, err
		}
	}

	return nil, nil, nil
}

func validateRepositoryName(repositoryName string) error {
	var (
		namespace string
		name      string
	)
	nameParts := strings.SplitN(repositoryName, "/", 2)
	if len(nameParts) < 2 {
		namespace = "library"
		name = nameParts[0]

		if validHex.MatchString(name) {
			return fmt.Errorf("Invalid repository name (%s), cannot specify 64-byte hexadecimal strings", name)
		}
	} else {
		namespace = nameParts[0]
		name = nameParts[1]
	}
	if !validNamespace.MatchString(namespace) {
		return fmt.Errorf("Invalid namespace name (%s), only [a-z0-9_] are allowed, size between 4 and 30", namespace)
	}
	if !validRepo.MatchString(name) {
		return fmt.Errorf("Invalid repository name (%s), only [a-z0-9-_.] are allowed", name)
	}
	return nil
}

// Resolves a repository name to a hostname + name
func ResolveRepositoryName(reposName string) (string, string, error) {
	if strings.Contains(reposName, "://") {
		// It cannot contain a scheme!
		return "", "", ErrInvalidRepositoryName
	}
	nameParts := strings.SplitN(reposName, "/", 2)
	if len(nameParts) == 1 || (!strings.Contains(nameParts[0], ".") && !strings.Contains(nameParts[0], ":") &&
		nameParts[0] != "localhost") {
		// This is a Docker Index repos (ex: samalba/hipache or ubuntu)
		err := validateRepositoryName(reposName)
		return IndexServerAddress(), reposName, err
	}
	hostname := nameParts[0]
	reposName = nameParts[1]
	if strings.Contains(hostname, "index.docker.io") {
		return "", "", fmt.Errorf("Invalid repository name, try \"%s\" instead", reposName)
	}
	if err := validateRepositoryName(reposName); err != nil {
		return "", "", err
	}

	return hostname, reposName, nil
}

func trustedLocation(req *http.Request) bool {
	var (
		trusteds = []string{"docker.com", "docker.io"}
		hostname = strings.SplitN(req.Host, ":", 2)[0]
	)
	if req.URL.Scheme != "https" {
		return false
	}

	for _, trusted := range trusteds {
		if hostname == trusted || strings.HasSuffix(hostname, "."+trusted) {
			return true
		}
	}
	return false
}

func AddRequiredHeadersToRedirectedRequests(req *http.Request, via []*http.Request) error {
	if via != nil && via[0] != nil {
		if trustedLocation(req) && trustedLocation(via[0]) {
			req.Header = via[0].Header
			return nil
		}
		for k, v := range via[0].Header {
			if k != "Authorization" {
				for _, vv := range v {
					req.Header.Add(k, vv)
				}
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package sink

import (
	"fmt"
	"net"
	"time"

	"github.com/bahadley/esp/log"
	"github.com/bahadley/esp/operator"
	"github.com/bahadley/esp/system"
)

var (
	outputChan chan *operator.SensorTuple
)

func Ingress() {
	sinkAddr, err := net.ResolveUDPAddr("udp",
		system.NodeAddr()+":"+system.SinkPort())
	if err != nil {
		log.Error.Fatal(err.Error())
	}

	conn, err := net.ListenUDP("udp", sinkAddr)
	if err != nil {
		log.Error.Fatal(err.Error())
	}

	defer conn.Close()

	log.Info.Printf("Listening for aggregation tuples (%s UDP) ...",
		sinkAddr.String())

	buf := make([]byte, system.TupleBufLen(), system.TupleBufCap())
	for {
		n, caddr, err := conn.ReadFromUDP(buf)
		if err != nil {
			log.Warning.Println(err.Error())
			continue
		}

		log.Trace.Printf("Rx(%s): %s", caddr, buf[0:n])

		aggTuple, err := operator.Unmarshal(buf[0:n])
		if err != nil {
			log.Warning.Printf("Failed to unmarshal tuple: %s", buf[0:n])
			continue
		}

		aggTuple.Timestamp = time.Now().UnixNano()

		outputChan <- aggTuple
	}
}

func Output() {
	outputChan = make(chan *operator.SensorTuple, system.ChannelBufSz())

	for {
		aggTuple := <-outputChan

		fmt.Printf("%d,%s,%.2f\n", aggTuple.Timestamp,
			aggTuple.Type, aggTuple.Data)
	}
}
<commit_msg>Set arrival time immediately after udp read<commit_after>package sink

import (
	"fmt"
	"net"
	"time"

	"github.com/bahadley/esp/log"
	"github.com/bahadley/esp/operator"
	"github.com/bahadley/esp/system"
)

var (
	outputChan chan *operator.SensorTuple
)

func Ingress() {
	sinkAddr, err := net.ResolveUDPAddr("udp",
		system.NodeAddr()+":"+system.SinkPort())
	if err != nil {
		log.Error.Fatal(err.Error())
	}

	conn, err := net.ListenUDP("udp", sinkAddr)
	if err != nil {
		log.Error.Fatal(err.Error())
	}

	defer conn.Close()

	log.Info.Printf("Listening for aggregation tuples (%s UDP) ...",
		sinkAddr.String())

	buf := make([]byte, system.TupleBufLen(), system.TupleBufCap())
	for {
		n, caddr, err := conn.ReadFromUDP(buf)
		if err != nil {
			log.Warning.Println(err.Error())
			continue
		}

		arrivalTime := time.Now().UnixNano()

		log.Trace.Printf("Rx(%s): %s", caddr, buf[0:n])

		aggTuple, err := operator.Unmarshal(buf[0:n])
		if err != nil {
			log.Warning.Printf("Failed to unmarshal tuple: %s", buf[0:n])
			continue
		}

		aggTuple.Timestamp = arrivalTime

		outputChan <- aggTuple
	}
}

func Output() {
	outputChan = make(chan *operator.SensorTuple, system.ChannelBufSz())

	for {
		aggTuple := <-outputChan

		fmt.Printf("%d,%s,%.2f\n", aggTuple.Timestamp,
			aggTuple.Type, aggTuple.Data)
	}
}
<|endoftext|>