<commit_before>package main

import (
	"log"
	"sort"
	"time"
)

type game struct {
	players     map[*player]bool
	projectiles map[*projectile]bool
	splosions   map[*splosion]bool
	register    chan *player
	unregister  chan *player
	id          chan int
}

type handshake struct {
	ID string `json:"id"`
}

var g = game{
	register:    make(chan *player),
	unregister:  make(chan *player),
	projectiles: make(map[*projectile]bool),
	splosions:   make(map[*splosion]bool),
	players:     make(map[*player]bool),
}

type robotSorter struct {
	robots []robot
}

func (s robotSorter) Len() int {
	return len(s.robots)
}

func (s robotSorter) Swap(i, j int) {
	s.robots[i], s.robots[j] = s.robots[j], s.robots[i]
}

func (s robotSorter) Less(i, j int) bool {
	return s.robots[i].Id < s.robots[j].Id
}

func (g *game) run() {
	g.id = make(chan int)
	go func() {
		for i := 0; ; i++ {
			g.id <- i
		}
	}()

	for {
		select {
		case p := <-g.register:
			log.Printf("adding player: %+v", p.Robot.Id)
			g.players[p] = true
		case p := <-g.unregister:
			delete(g.players, p)
			close(p.send)
		case <-time.Tick(time.Duration(*tick) * time.Millisecond):
			payload := payload{}
			payload.Robots = []robot{}
			payload.Projectiles = []projectile{}

			robots_remaining := 0

			for p := range g.players {
				if p.Robot.Health > 0 {
					robots_remaining++
					p.scan()
					p.nudge()
					if p.Robot.FireAt.X != 0 && p.Robot.FireAt.Y != 0 {
						p.fire()
					}
				}
				payload.Robots = append(payload.Robots, p.Robot)
			}
			sort.Sort(robotSorter{payload.Robots})

			for p := range g.projectiles {
				p.nudge()
				payload.Projectiles = append(payload.Projectiles, *p)
			}

			for s := range g.splosions {
				s.tick()
				payload.Splosions = append(payload.Splosions, *s)
			}

			if robots_remaining <= 1 && len(g.players) > 1 {
				for p := range g.players {
					if p.Robot.Health > 0 {
						log.Printf("Robot %v Wins", p.Robot.Id)
					}
					p.reset()
				}
				payload.Reset = true
			} else {
				payload.Reset = false
			}

			for p := range g.players {
				p.send <- &payload
			}

		}
	}
}
<commit_msg>adding turn counter and some perf stats to the server<commit_after>package main

import (
	"log"
	"sort"
	"time"
)

type game struct {
	players     map[*player]bool
	projectiles map[*projectile]bool
	splosions   map[*splosion]bool
	register    chan *player
	unregister  chan *player
	id          chan int
	turn        int
}

type handshake struct {
	ID string `json:"id"`
}

var g = game{
	register:    make(chan *player),
	unregister:  make(chan *player),
	projectiles: make(map[*projectile]bool),
	splosions:   make(map[*splosion]bool),
	players:     make(map[*player]bool),
	turn:        0,
}

type robotSorter struct {
	robots []robot
}

func (s robotSorter) Len() int {
	return len(s.robots)
}

func (s robotSorter) Swap(i, j int) {
	s.robots[i], s.robots[j] = s.robots[j], s.robots[i]
}

func (s robotSorter) Less(i, j int) bool {
	return s.robots[i].Id < s.robots[j].Id
}

func (g *game) run() {
	g.id = make(chan int)
	go func() {
		for i := 0; ; i++ {
			g.id <- i
		}
	}()

	for {
		select {
		case p := <-g.register:
			log.Printf("adding player: %+v", p.Robot.Id)
			g.players[p] = true
		case p := <-g.unregister:
			delete(g.players, p)
			close(p.send)
		case <-time.Tick(time.Duration(*tick) * time.Millisecond):
			g.turn++
			t0 := time.Now()
			log.Printf("\033[2JTurn: %v", g.turn)
			log.Printf("Players: %v", len(g.players))
			log.Printf("Projectiles: %v", len(g.projectiles))
			log.Printf("Explosions: %v", len(g.splosions))
			payload := payload{}
			payload.Robots = []robot{}
			payload.Projectiles = []projectile{}

			robots_remaining := 0

			for p := range g.players {
				if p.Robot.Health > 0 {
					robots_remaining++
					p.scan()
					p.nudge()
					if p.Robot.FireAt.X != 0 && p.Robot.FireAt.Y != 0 {
						p.fire()
					}
				}
				payload.Robots = append(payload.Robots, p.Robot)
			}
			sort.Sort(robotSorter{payload.Robots})

			for p := range g.projectiles {
				p.nudge()
				payload.Projectiles = append(payload.Projectiles, *p)
			}

			for s := range g.splosions {
				s.tick()
				payload.Splosions = append(payload.Splosions, *s)
			}

			if robots_remaining <= 1 && len(g.players) > 1 {
				for p := range g.players {
					if p.Robot.Health > 0 {
						log.Printf("Robot %v Wins", p.Robot.Id)
					}
					p.reset()
				}
				payload.Reset = true
			} else {
				payload.Reset = false
			}

			t1 := time.Now()
			log.Printf("Turn Processes %v\n", t1.Sub(t0))

			for p := range g.players {
				p.send <- &payload
			}

			t1 = time.Now()
			log.Printf("Sent Payload %v\n", t1.Sub(t0))

		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"runtime"
	"time"
)

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	server := WysteriaServer{
		// Allocate 3 seconds for the server connections to be severed when shutting down
		GracefulShutdownTime: time.Second * 3,

		// Pass in all config settings
		settings: Config,
	}

	// Connect to all required endpoints, die in the event of any failures
	err := server.Run()
	if err != nil {
		panic(err)
	}
	defer server.Shutdown()
}
<commit_msg>Remove superfluous GOMAXPROCS call<commit_after>package main

import (
	"time"
)

func main() {
	server := WysteriaServer{
		// Allocate 3 seconds for the server connections to be severed when shutting down
		GracefulShutdownTime: time.Second * 3,

		// Pass in all config settings
		settings: Config,
	}

	// Connect to all required endpoints, die in the event of any failures
	err := server.Run()
	if err != nil {
		panic(err)
	}
	defer server.Shutdown()
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Copyright 2012 The Gorilla Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rpc

import (
	"net/http"
	"testing"
)

type Service1Request struct {
	A int
	B int
}

type Service1Response struct {
	Result int
}

type Service1 struct {
}

func (t *Service1) Multiply(r *http.Request, req *Service1Request, res *Service1Response) error {
	res.Result = req.A * req.B
	return nil
}

func (t *Service1) Add(req *Service1Request, res *Service1Response) error {
	res.Result = req.A + req.B
	return nil
}

type Service2 struct {
}

func TestRegisterService(t *testing.T) {
	var err error
	s := NewServer()
	service1 := new(Service1)
	service2 := new(Service2)

	// Inferred name.
	err = s.RegisterService(service1, "")
	if err != nil || !s.HasMethod("Service1.Multiply") {
		t.Errorf("Expected to be registered: Service1.Multiply")
	}
	// Provided name.
	err = s.RegisterService(service1, "Foo")
	if err != nil || !s.HasMethod("Foo.Multiply") {
		t.Errorf("Expected to be registered: Foo.Multiply")
	}
	// No methods.
	err = s.RegisterService(service2, "")
	if err == nil {
		t.Errorf("Expected error on service2")
	}
}

func TestRegisterTCPService(t *testing.T) {
	var err error
	s := NewServer()
	service1 := new(Service1)
	service2 := new(Service2)

	// Inferred name.
	err = s.RegisterTCPService(service1, "")
	if err != nil || !s.HasMethod("Service1.Add") {
		t.Errorf("Expected to be registered: Service1.Add")
	}
	// Provided name.
	err = s.RegisterTCPService(service1, "Foo")
	if err != nil || !s.HasMethod("Foo.Add") {
		t.Errorf("Expected to be registered: Foo.Add")
	}
	// No methods.
	err = s.RegisterTCPService(service2, "")
	if err == nil {
		t.Errorf("Expected error on service2")
	}
}
<commit_msg>Test Register is OK<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Copyright 2012 The Gorilla Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rpc

import (
	"net/http"
	"testing"
)

type Service1Request struct {
	A int
	B int
}

type Service1Response struct {
	Result int
}

type Service1 struct {
}

func (t *Service1) Multiply(r *http.Request, req *Service1Request, res *Service1Response) error {
	res.Result = req.A * req.B
	return nil
}

func (t *Service1) Add(req *Service1Request, res *Service1Response) error {
	res.Result = req.A + req.B
	return nil
}

type Service2 struct {
}

func TestRegisterService(t *testing.T) {
	var err error
	s := NewServer()
	service1 := new(Service1)
	service2 := new(Service2)

	// Inferred name.
	err = s.RegisterService(service1, "")
	if err != nil || !s.HasMethod("Service1.Multiply") {
		t.Errorf("Expected to be registered: Service1.Multiply")
	}
	// Provided name.
	err = s.RegisterService(service1, "Foo")
	if err != nil || !s.HasMethod("Foo.Multiply") {
		t.Errorf("Expected to be registered: Foo.Multiply")
	}
	// No methods.
	err = s.RegisterService(service2, "")
	if err == nil {
		t.Errorf("Expected error on service2")
	}
}

func TestRegisterTCPService(t *testing.T) {
	var err error
	s := NewServer()
	service1 := new(Service1)
	service2 := new(Service2)

	// Inferred name.
	err = s.RegisterTCPService(service1, "")
	if err != nil || !s.HasMethod("Service1.Add") {
		t.Errorf("Expected to be registered: Service1.Add")
	}
	// Provided name.
	err = s.RegisterTCPService(service1, "Foo")
	if err != nil || !s.HasMethod("Foo.Add") {
		t.Errorf("Expected to be registered: Foo.Add")
	}
	// No methods.
	err = s.RegisterTCPService(service2, "")
	if err == nil {
		t.Errorf("Expected error on service2")
	}
}

func TestRegisterBeforeAfterFunc(t *testing.T) {
	var err error
	s := NewServer()

	err = s.RegisterBeforeFunc(func(i *RequestInfo) {})
	if err != nil {
		t.Errorf("Expected to be nil")
	}

	err = s.RegisterAfterFunc(func(i *RequestInfo) {})
	if err != nil {
		t.Errorf("Expected to be nil")
	}
}
<|endoftext|>


<commit_before>package service

import (
	"encoding/json"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/google/shlex"
	"github.com/spf13/cobra"

	"github.com/runabove/sail/internal"
)

var cmdAddLink []string
var cmdAddNetworkAllow string
var addPublish []string
var cmdAddGateway []string
var cmdAddVolume []string
var addBatch bool
var cmdAddRedeploy bool
var cmdAddBody Add
var cmdAddNetwork []string
var cmdAddCommand string
var cmdAddEntrypoint string

const cmdAddUsage = "Invalid usage. sail service add [<application>/]<repository>[:tag] [<service>]. Please see sail service add --help"

func addCmd() *cobra.Command {

	cmd := &cobra.Command{
		Use:   "add",
		Short: "Add a new docker service",
		Long: `add [<namespace>/]<repository>[:tag] [namespace/]<service-name>
		--model         Container model
		--number        Number of container to run
		[--link         name:alias]
		[--network      {public|private|<namespace name>}]
		[--network-allow [network:]ip[/mask] Use IPs whitelist]
		[--publish, -p  Publish a container's port to the host]
		[                 format: network:publishedPort:containerPort, network::containerPort, publishedPort:containerPort, containerPort]
		[--gateway      network-input:network-output
		[--restart {no|always[:<max>]|on-failure[:<max>]}]
		[--volume       /path:size] (Size in GB)
		[--batch        do not attach console on start]
		[--pool         deploy on dedicated host pool <name>]
		[--redeploy     if the service already exists, redeploy instead]

		override docker options:
			--user
			--entrypoint
			--command
			--workdir
			--environment KEY=val
		other options:
		`,
		Run: cmdAdd,
	}
	cmd.Flags().StringVarP(&cmdAddBody.ContainerModel, "model", "", "x1", "Container model")
	cmd.Flags().IntVarP(&cmdAddBody.ContainerNumber, "number", "", 1, "Number of container to run")
	cmd.Flags().StringSliceVarP(&cmdAddLink, "link", "", nil, "name:alias")
	cmd.Flags().StringSliceVar(&cmdAddNetwork, "network", []string{}, "public|private|<namespace name>")
	cmd.Flags().StringVarP(&cmdAddNetworkAllow, "network-allow", "", "", "[network:]ip[/mask] Use IPs whitelist")
	cmd.Flags().StringSliceVarP(&addPublish, "publish", "p", nil, "Publish a container's port to the host")
	cmd.Flags().StringSliceVar(&cmdAddGateway, "gateway", nil, "network-input:network-output")
	cmd.Flags().StringVarP(&cmdAddBody.RestartPolicy, "restart", "", "no", "{no|always[:<max>]|on-failure[:<max>]}")
	cmd.Flags().StringVarP(&cmdAddCommand, "command", "", "", "override docker run command")
	cmd.Flags().StringVarP(&cmdAddBody.RepositoryTag, "tag", "", "", "deploy from new image version")
	cmd.Flags().StringVarP(&cmdAddBody.ContainerWorkdir, "workdir", "", "", "override docker workdir")
	cmd.Flags().StringVarP(&cmdAddEntrypoint, "entrypoint", "", "", "override docker entrypoint")
	cmd.Flags().StringVarP(&cmdAddBody.ContainerUser, "user", "", "", "override docker user")
	cmd.Flags().StringSliceVar(&cmdAddVolume, "volume", nil, "/path:size] (Size in GB)")
	cmd.Flags().BoolVarP(&addBatch, "batch", "", false, "do not attach console on start")
	cmd.Flags().BoolVarP(&cmdAddRedeploy, "redeploy", "", false, "if the service already exists, redeploy instead")
	cmd.Flags().StringSliceVarP(&cmdAddBody.ContainerEnvironment, "env", "e", nil, "override docker environment")
	cmd.Flags().StringVarP(&cmdAddBody.Pool, "pool", "", "", "Dedicated host pool")
	return cmd
}

// PortConfig is a parameter of Add to modify exposed container ports
type PortConfig struct {
	PublishedPort string `json:"published_port"`
	Network       string `json:"network,omitempty"`
}

// VolumeConfig is a parameter of Add to modify mounted volumes
type VolumeConfig struct {
	Size string `json:"size"`
}

// Add struct holds all parameters sent to /applications/%s/services/%s?stream
type Add struct {
	Service              string                         `json:"-"`
	Volumes              map[string]VolumeConfig        `json:"volumes,omitempty"`
	Repository           string                         `json:"repository"`
	ContainerUser        string                         `json:"container_user,omitempty"`
	RestartPolicy        string                         `json:"restart_policy"`
	ContainerCommand     []string                       `json:"container_command,omitempty"`
	ContainerNetwork     map[string]map[string][]string `json:"container_network"`
	ContainerEntrypoint  []string                       `json:"container_entrypoint,omitempty"`
	ContainerNumber      int                            `json:"container_number"`
	RepositoryTag        string                         `json:"repository_tag"`
	Links                map[string]string              `json:"links"`
	Application          string                         `json:"namespace"`
	ContainerWorkdir     string                         `json:"container_workdir,omitempty"`
	ContainerEnvironment []string                       `json:"container_environment"`
	ContainerModel       string                         `json:"container_model"`
	ContainerPorts       map[string][]PortConfig        `json:"container_ports"`
	Pool                 string                         `json:"pool,omitempty"`
}

func cmdAdd(cmd *cobra.Command, args []string) {
	cmdAddBody.ContainerNetwork = make(map[string]map[string][]string)
	cmdAddBody.Links = make(map[string]string)
	cmdAddBody.ContainerPorts = make(map[string][]PortConfig)

	if len(args) < 2 {
		fmt.Fprintln(os.Stderr, cmdAddUsage)
		os.Exit(1)
	}

	// Split namespace and repository
	host, app, repo, tag, err := internal.ParseResourceName(args[0])
	internal.Check(err)
	cmdAddBody.Application = app
	cmdAddBody.Repository = repo
	cmdAddBody.RepositoryTag = tag

	if !internal.CheckHostConsistent(host) {
		fmt.Fprintf(os.Stderr, "Error: Invalid Host %s for endpoint %s\n", host, internal.Host)
		os.Exit(1)
	}

	err = internal.CheckName(cmdAddBody.Application)
	internal.Check(err)
	err = internal.CheckName(cmdAddBody.Repository)
	internal.Check(err)
	err = internal.CheckName(cmdAddBody.Service)
	internal.Check(err)

	// Service name
	if len(args) >= 2 {
		cmdAddBody.Service = args[1]
	} else {
		cmdAddBody.Service = cmdAddBody.Repository
	}

	serviceAdd(cmdAddBody)
}

func serviceAdd(args Add) {

	if args.ContainerEnvironment == nil {
		args.ContainerEnvironment = make([]string, 0)
	}

	// Parse command
	if cmdAddCommand != "" {
		command, err := shlex.Split(cmdAddCommand)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Fatal, cannot split command %s\n", err)
			return
		}
		args.ContainerCommand = command
	}

	// Parse Entrypoint
	if cmdAddEntrypoint != "" {
		entrypoint, err := shlex.Split(cmdAddEntrypoint)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Fatal, cannot split command %s\n", err)
			return
		}
		args.ContainerEntrypoint = entrypoint
	}

	// Parse volumes
	if len(cmdAddVolume) > 0 {
		args.Volumes = make(map[string]VolumeConfig)
	}
	for _, vol := range cmdAddVolume {
		t := strings.Split(vol, ":")
		if len(t) == 2 {
			args.Volumes[t[0]] = VolumeConfig{Size: t[1]}
		} else if len(t) == 1 {
			args.Volumes[t[0]] = VolumeConfig{Size: "10"}
		} else {
			fmt.Fprintf(os.Stderr, "Error: Volume parameter '%s' not formated correctly\n", vol)
			os.Exit(1)
		}
	}

	// Parse links
	if len(redeployLink) > 0 {
		args.Links = make(map[string]string)
	}

	for _, link := range cmdAddLink {
		t := strings.Split(link, ":")
		if len(t) == 1 {
			args.Links[t[0]] = t[0]
		} else {
			args.Links[t[0]] = t[1]
		}
	}

	// Parse ContainerNetworks arguments
	for _, network := range cmdAddNetwork {
		args.ContainerNetwork[network] = make(map[string][]string)
	}

	for _, gat := range cmdAddGateway {
		t := strings.Split(gat, ":")
		if len(t) != 2 {
			fmt.Fprintf(os.Stderr, "Invalid gateway parameter, should be \"input:output\". Typically, output will be one of 'predictor', 'public'")
			os.Exit(1)
		}
		if _, ok := args.ContainerNetwork[t[0]]; !ok {
			fmt.Fprintf(os.Stderr, "Automatically adding %s to network list\n", t[0])
			args.ContainerNetwork[t[0]] = make(map[string][]string)
		}
		if _, ok := args.ContainerNetwork[t[1]]; !ok {
			fmt.Fprintf(os.Stderr, "Automatically adding %s to network list\n", t[0])
			args.ContainerNetwork[t[1]] = make(map[string][]string)
		}
		args.ContainerNetwork[t[0]]["gateway_to"] = append(args.ContainerNetwork[t[0]]["gateway_to"], t[1])
	}

	// Parse ContainerPorts
	args.ContainerPorts = parsePublishedPort(addPublish)

	path := fmt.Sprintf("/applications/%s/services/%s", args.Application, args.Service)
	body, err := json.MarshalIndent(args, " ", " ")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Fatal: %s\n", err)
		return
	}

	if addBatch {
		ret, code, err := internal.Request("POST", path, body)

		// http.Request failed for some reason
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %s\n", err)
			return
		}

		//  If we are in ensure mode, fallback to redeploy
		if code == 409 && cmdAddRedeploy {
			ensureMode(args)
			return
		}

		// If API returned a json error
		e := internal.DecodeError(ret)
		if e != nil {
			fmt.Fprintf(os.Stderr, "%s\n", e)
			return
		}

		// Just print data
		internal.FormatOutputDef(ret)

		// Always start service
		if internal.Format == "pretty" {
			fmt.Fprintf(os.Stderr, "Starting service %s/%s...\n", args.Application, args.Service)
		}
		serviceStart(args.Application, args.Service, true)

		return
	}

	buffer, code, err := internal.Stream("POST", path+"?stream", body)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %s\n", err)
		return
	}

	if code == 409 && cmdAddRedeploy {
		ensureMode(args)
		return
	}

	line, err := internal.DisplayStream(buffer)
	internal.Check(err)
	if line != nil {
		var data map[string]interface{}
		err = json.Unmarshal(line, &data)
		internal.Check(err)

		fmt.Printf("Hostname: %v\n", data["hostname"])
		fmt.Printf("Running containers: %v/%v\n", data["container_number"], data["container_target"])
	}
}

func ensureMode(args Add) {
	redeployBatch = addBatch
	redeployBody := Redeploy{
		Service:              args.Service,
		Volumes:              args.Volumes,
		Repository:           args.Repository,
		ContainerUser:        args.ContainerUser,
		RestartPolicy:        args.RestartPolicy,
		ContainerCommand:     args.ContainerCommand,
		ContainerNetwork:     args.ContainerNetwork,
		ContainerEntrypoint:  args.ContainerEntrypoint,
		ContainerNumber:      args.ContainerNumber,
		RepositoryTag:        args.RepositoryTag,
		Links:                args.Links,
		Application:          args.Application,
		ContainerWorkdir:     args.ContainerWorkdir,
		ContainerEnvironment: args.ContainerEnvironment,
		ContainerModel:       args.ContainerModel,
		ContainerPorts:       args.ContainerPorts,
	}
	serviceRedeploy(redeployBody)
}

func parsePublishedPort(args []string) map[string][]PortConfig {
	v := make(map[string][]PortConfig)

	for _, pub := range args {
		split := strings.Split(pub, ":")
		if len(split) == 1 { // containerPort
			v[split[0]+"/tcp"] = []PortConfig{PortConfig{PublishedPort: split[0]}}
		} else if len(split) == 2 { // network:containerPort, publishedPort:containerPort
			_, err := strconv.Atoi(split[0])
			if err != nil { // network:containerPort
				key := split[1] + "/tcp"
				v[key] = append(v[key], PortConfig{PublishedPort: split[1], Network: split[0]})
			} else { // publishedPort:containerPort
				key := split[1] + "/tcp"
				v[key] = append(v[key], PortConfig{PublishedPort: split[0]})
			}
		} else if len(split) == 3 { // network:publishedPort:containerPort, network::containerPort
			if split[1] == "" {
				split[1] = split[2]
			}

			key := split[2] + "/tcp"
			v[key] = append(v[key], PortConfig{PublishedPort: split[0], Network: split[1]})
		} else {
			fmt.Fprintf(os.Stderr, "Error: Invalid port expose rule %s.", pub)
			os.Exit(1)
		}
	}

	return v
}
<commit_msg>fix: use repo name as default service name on service add<commit_after>package service

import (
	"encoding/json"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/google/shlex"
	"github.com/spf13/cobra"

	"github.com/runabove/sail/internal"
)

var cmdAddLink []string
var cmdAddNetworkAllow string
var addPublish []string
var cmdAddGateway []string
var cmdAddVolume []string
var addBatch bool
var cmdAddRedeploy bool
var cmdAddBody Add
var cmdAddNetwork []string
var cmdAddCommand string
var cmdAddEntrypoint string

const cmdAddUsage = "Invalid usage. sail service add [<application>/]<repository>[:tag] [<service>]. Please see sail service add --help"

func addCmd() *cobra.Command {

	cmd := &cobra.Command{
		Use:   "add",
		Short: "Add a new docker service",
		Long: `add [<namespace>/]<repository>[:tag] [namespace/]<service-name>
		--model         Container model
		--number        Number of container to run
		[--link         name:alias]
		[--network      {public|private|<namespace name>}]
		[--network-allow [network:]ip[/mask] Use IPs whitelist]
		[--publish, -p  Publish a container's port to the host]
		[                 format: network:publishedPort:containerPort, network::containerPort, publishedPort:containerPort, containerPort]
		[--gateway      network-input:network-output
		[--restart {no|always[:<max>]|on-failure[:<max>]}]
		[--volume       /path:size] (Size in GB)
		[--batch        do not attach console on start]
		[--pool         deploy on dedicated host pool <name>]
		[--redeploy     if the service already exists, redeploy instead]

		override docker options:
			--user
			--entrypoint
			--command
			--workdir
			--environment KEY=val
		other options:
		`,
		Run: cmdAdd,
	}
	cmd.Flags().StringVarP(&cmdAddBody.ContainerModel, "model", "", "x1", "Container model")
	cmd.Flags().IntVarP(&cmdAddBody.ContainerNumber, "number", "", 1, "Number of container to run")
	cmd.Flags().StringSliceVarP(&cmdAddLink, "link", "", nil, "name:alias")
	cmd.Flags().StringSliceVar(&cmdAddNetwork, "network", []string{}, "public|private|<namespace name>")
	cmd.Flags().StringVarP(&cmdAddNetworkAllow, "network-allow", "", "", "[network:]ip[/mask] Use IPs whitelist")
	cmd.Flags().StringSliceVarP(&addPublish, "publish", "p", nil, "Publish a container's port to the host")
	cmd.Flags().StringSliceVar(&cmdAddGateway, "gateway", nil, "network-input:network-output")
	cmd.Flags().StringVarP(&cmdAddBody.RestartPolicy, "restart", "", "no", "{no|always[:<max>]|on-failure[:<max>]}")
	cmd.Flags().StringVarP(&cmdAddCommand, "command", "", "", "override docker run command")
	cmd.Flags().StringVarP(&cmdAddBody.RepositoryTag, "tag", "", "", "deploy from new image version")
	cmd.Flags().StringVarP(&cmdAddBody.ContainerWorkdir, "workdir", "", "", "override docker workdir")
	cmd.Flags().StringVarP(&cmdAddEntrypoint, "entrypoint", "", "", "override docker entrypoint")
	cmd.Flags().StringVarP(&cmdAddBody.ContainerUser, "user", "", "", "override docker user")
	cmd.Flags().StringSliceVar(&cmdAddVolume, "volume", nil, "/path:size] (Size in GB)")
	cmd.Flags().BoolVarP(&addBatch, "batch", "", false, "do not attach console on start")
	cmd.Flags().BoolVarP(&cmdAddRedeploy, "redeploy", "", false, "if the service already exists, redeploy instead")
	cmd.Flags().StringSliceVarP(&cmdAddBody.ContainerEnvironment, "env", "e", nil, "override docker environment")
	cmd.Flags().StringVarP(&cmdAddBody.Pool, "pool", "", "", "Dedicated host pool")
	return cmd
}

// PortConfig is a parameter of Add to modify exposed container ports
type PortConfig struct {
	PublishedPort string `json:"published_port"`
	Network       string `json:"network,omitempty"`
}

// VolumeConfig is a parameter of Add to modify mounted volumes
type VolumeConfig struct {
	Size string `json:"size"`
}

// Add struct holds all parameters sent to /applications/%s/services/%s?stream
type Add struct {
	Service              string                         `json:"-"`
	Volumes              map[string]VolumeConfig        `json:"volumes,omitempty"`
	Repository           string                         `json:"repository"`
	ContainerUser        string                         `json:"container_user,omitempty"`
	RestartPolicy        string                         `json:"restart_policy"`
	ContainerCommand     []string                       `json:"container_command,omitempty"`
	ContainerNetwork     map[string]map[string][]string `json:"container_network"`
	ContainerEntrypoint  []string                       `json:"container_entrypoint,omitempty"`
	ContainerNumber      int                            `json:"container_number"`
	RepositoryTag        string                         `json:"repository_tag"`
	Links                map[string]string              `json:"links"`
	Application          string                         `json:"namespace"`
	ContainerWorkdir     string                         `json:"container_workdir,omitempty"`
	ContainerEnvironment []string                       `json:"container_environment"`
	ContainerModel       string                         `json:"container_model"`
	ContainerPorts       map[string][]PortConfig        `json:"container_ports"`
	Pool                 string                         `json:"pool,omitempty"`
}

func cmdAdd(cmd *cobra.Command, args []string) {
	cmdAddBody.ContainerNetwork = make(map[string]map[string][]string)
	cmdAddBody.Links = make(map[string]string)
	cmdAddBody.ContainerPorts = make(map[string][]PortConfig)

	if len(args) >= 2 || len(args) < 1 {
		fmt.Fprintln(os.Stderr, cmdAddUsage)
		os.Exit(1)
	}

	// Split namespace and repository
	host, app, repo, tag, err := internal.ParseResourceName(args[0])
	internal.Check(err)
	cmdAddBody.Application = app
	cmdAddBody.Repository = repo
	cmdAddBody.RepositoryTag = tag

	if !internal.CheckHostConsistent(host) {
		fmt.Fprintf(os.Stderr, "Error: Invalid Host %s for endpoint %s\n", host, internal.Host)
		os.Exit(1)
	}

	err = internal.CheckName(cmdAddBody.Application)
	internal.Check(err)
	err = internal.CheckName(cmdAddBody.Repository)
	internal.Check(err)
	err = internal.CheckName(cmdAddBody.Service)
	internal.Check(err)

	// Service name
	if len(args) > 1 {
		cmdAddBody.Service = args[1]
	} else {
		cmdAddBody.Service = cmdAddBody.Repository
	}

	serviceAdd(cmdAddBody)
}

func serviceAdd(args Add) {

	if args.ContainerEnvironment == nil {
		args.ContainerEnvironment = make([]string, 0)
	}

	// Parse command
	if cmdAddCommand != "" {
		command, err := shlex.Split(cmdAddCommand)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Fatal, cannot split command %s\n", err)
			return
		}
		args.ContainerCommand = command
	}

	// Parse Entrypoint
	if cmdAddEntrypoint != "" {
		entrypoint, err := shlex.Split(cmdAddEntrypoint)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Fatal, cannot split command %s\n", err)
			return
		}
		args.ContainerEntrypoint = entrypoint
	}

	// Parse volumes
	if len(cmdAddVolume) > 0 {
		args.Volumes = make(map[string]VolumeConfig)
	}
	for _, vol := range cmdAddVolume {
		t := strings.Split(vol, ":")
		if len(t) == 2 {
			args.Volumes[t[0]] = VolumeConfig{Size: t[1]}
		} else if len(t) == 1 {
			args.Volumes[t[0]] = VolumeConfig{Size: "10"}
		} else {
			fmt.Fprintf(os.Stderr, "Error: Volume parameter '%s' not formated correctly\n", vol)
			os.Exit(1)
		}
	}

	// Parse links
	if len(redeployLink) > 0 {
		args.Links = make(map[string]string)
	}

	for _, link := range cmdAddLink {
		t := strings.Split(link, ":")
		if len(t) == 1 {
			args.Links[t[0]] = t[0]
		} else {
			args.Links[t[0]] = t[1]
		}
	}

	// Parse ContainerNetworks arguments
	for _, network := range cmdAddNetwork {
		args.ContainerNetwork[network] = make(map[string][]string)
	}

	for _, gat := range cmdAddGateway {
		t := strings.Split(gat, ":")
		if len(t) != 2 {
			fmt.Fprintf(os.Stderr, "Invalid gateway parameter, should be \"input:output\". Typically, output will be one of 'predictor', 'public'")
			os.Exit(1)
		}
		if _, ok := args.ContainerNetwork[t[0]]; !ok {
			fmt.Fprintf(os.Stderr, "Automatically adding %s to network list\n", t[0])
			args.ContainerNetwork[t[0]] = make(map[string][]string)
		}
		if _, ok := args.ContainerNetwork[t[1]]; !ok {
			fmt.Fprintf(os.Stderr, "Automatically adding %s to network list\n", t[0])
			args.ContainerNetwork[t[1]] = make(map[string][]string)
		}
		args.ContainerNetwork[t[0]]["gateway_to"] = append(args.ContainerNetwork[t[0]]["gateway_to"], t[1])
	}

	// Parse ContainerPorts
	args.ContainerPorts = parsePublishedPort(addPublish)

	path := fmt.Sprintf("/applications/%s/services/%s", args.Application, args.Service)
	body, err := json.MarshalIndent(args, " ", " ")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Fatal: %s\n", err)
		return
	}

	if addBatch {
		ret, code, err := internal.Request("POST", path, body)

		// http.Request failed for some reason
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error: %s\n", err)
			return
		}

		//  If we are in ensure mode, fallback to redeploy
		if code == 409 && cmdAddRedeploy {
			ensureMode(args)
			return
		}

		// If API returned a json error
		e := internal.DecodeError(ret)
		if e != nil {
			fmt.Fprintf(os.Stderr, "%s\n", e)
			return
		}

		// Just print data
		internal.FormatOutputDef(ret)

		// Always start service
		if internal.Format == "pretty" {
			fmt.Fprintf(os.Stderr, "Starting service %s/%s...\n", args.Application, args.Service)
		}
		serviceStart(args.Application, args.Service, true)

		return
	}

	buffer, code, err := internal.Stream("POST", path+"?stream", body)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %s\n", err)
		return
	}

	if code == 409 && cmdAddRedeploy {
		ensureMode(args)
		return
	}

	line, err := internal.DisplayStream(buffer)
	internal.Check(err)
	if line != nil {
		var data map[string]interface{}
		err = json.Unmarshal(line, &data)
		internal.Check(err)

		fmt.Printf("Hostname: %v\n", data["hostname"])
		fmt.Printf("Running containers: %v/%v\n", data["container_number"], data["container_target"])
	}
}

func ensureMode(args Add) {
	redeployBatch = addBatch
	redeployBody := Redeploy{
		Service:              args.Service,
		Volumes:              args.Volumes,
		Repository:           args.Repository,
		ContainerUser:        args.ContainerUser,
		RestartPolicy:        args.RestartPolicy,
		ContainerCommand:     args.ContainerCommand,
		ContainerNetwork:     args.ContainerNetwork,
		ContainerEntrypoint:  args.ContainerEntrypoint,
		ContainerNumber:      args.ContainerNumber,
		RepositoryTag:        args.RepositoryTag,
		Links:                args.Links,
		Application:          args.Application,
		ContainerWorkdir:     args.ContainerWorkdir,
		ContainerEnvironment: args.ContainerEnvironment,
		ContainerModel:       args.ContainerModel,
		ContainerPorts:       args.ContainerPorts,
	}
	serviceRedeploy(redeployBody)
}

func parsePublishedPort(args []string) map[string][]PortConfig {
	v := make(map[string][]PortConfig)

	for _, pub := range args {
		split := strings.Split(pub, ":")
		if len(split) == 1 { // containerPort
			v[split[0]+"/tcp"] = []PortConfig{PortConfig{PublishedPort: split[0]}}
		} else if len(split) == 2 { // network:containerPort, publishedPort:containerPort
			_, err := strconv.Atoi(split[0])
			if err != nil { // network:containerPort
				key := split[1] + "/tcp"
				v[key] = append(v[key], PortConfig{PublishedPort: split[1], Network: split[0]})
			} else { // publishedPort:containerPort
				key := split[1] + "/tcp"
				v[key] = append(v[key], PortConfig{PublishedPort: split[0]})
			}
		} else if len(split) == 3 { // network:publishedPort:containerPort, network::containerPort
			if split[1] == "" {
				split[1] = split[2]
			}

			key := split[2] + "/tcp"
			v[key] = append(v[key], PortConfig{PublishedPort: split[0], Network: split[1]})
		} else {
			fmt.Fprintf(os.Stderr, "Error: Invalid port expose rule %s.", pub)
			os.Exit(1)
		}
	}

	return v
}
<|endoftext|>


<commit_before>package networking

import (
	"bytes"
	"context"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	appsv1 "k8s.io/api/apps/v1"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes/scheme"
	coreclientset "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/remotecommand"
	"k8s.io/kubernetes/test/e2e/framework"
	e2enetwork "k8s.io/kubernetes/test/e2e/framework/network"
	e2eskipper "k8s.io/kubernetes/test/e2e/framework/skipper"

	"github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"
)

const (
	nodeTCPPort = 9000
	nodeUDPPort = 9999
)

var _ = ginkgo.Describe("[sig-network] Internal connectivity", func() {
	f := framework.NewDefaultFramework("k8s-nettest")

	ginkgo.It("for TCP and UDP on ports 9000-9999 is allowed", func() {
		e2eskipper.SkipUnlessNodeCountIsAtLeast(2)
		clientConfig := f.ClientConfig()

		one := int64(0)
		ds := &appsv1.DaemonSet{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "webserver",
				Namespace: f.Namespace.Name,
			},
			Spec: appsv1.DaemonSetSpec{
				Selector: &metav1.LabelSelector{
					MatchLabels: map[string]string{
						"apps": "webserver",
					},
				},
				Template: v1.PodTemplateSpec{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							"apps": "webserver",
						},
					},
					Spec: v1.PodSpec{
						Tolerations: []v1.Toleration{
							{
								Key:      "node-role.kubernetes.io/master",
								Operator: v1.TolerationOpExists,
								Effect:   v1.TaintEffectNoSchedule,
							},
						},
						HostNetwork:                   true,
						TerminationGracePeriodSeconds: &one,
						Containers: []v1.Container{
							{
								Name:    "webserver",
								Image:   e2enetwork.NetexecImageName,
								Command: []string{"/bin/bash", "-c", fmt.Sprintf("#!/bin/bash\napk add -q --update tcpdump\n./agnhost netexec --http-port=%v --udp-port=%v &\nexec tcpdump -i any port %v or port %v -n", nodeTCPPort, nodeUDPPort, nodeTCPPort, nodeUDPPort)},
								Ports: []v1.ContainerPort{
									{Name: "tcp", ContainerPort: nodeTCPPort},
									{Name: "udp", ContainerPort: nodeUDPPort},
								},
								ReadinessProbe: &v1.Probe{
									InitialDelaySeconds: 10,
									Handler: v1.Handler{
										HTTPGet: &v1.HTTPGetAction{
											Port: intstr.FromInt(nodeTCPPort),
										},
									},
								},
								SecurityContext: &v1.SecurityContext{
									Capabilities: &v1.Capabilities{
										Add: []v1.Capability{"NET_RAW"},
									},
								},
							},
						},
					},
				},
			},
		}
		name := ds.Name
		ds, err := f.ClientSet.AppsV1().DaemonSets(f.Namespace.Name).Create(context.Background(), ds, metav1.CreateOptions{})
		o.Expect(err).NotTo(o.HaveOccurred())
		err = wait.PollImmediate(5*time.Second, 5*time.Minute, func() (bool, error) {
			ds, err = f.ClientSet.AppsV1().DaemonSets(f.Namespace.Name).Get(context.Background(), name, metav1.GetOptions{})
			if err != nil {
				framework.Logf("unable to retrieve daemonset: %v", err)
				return false, nil
			}
			if ds.Status.ObservedGeneration != ds.Generation || ds.Status.NumberAvailable == 0 || ds.Status.NumberAvailable != ds.Status.DesiredNumberScheduled {
				framework.Logf("waiting for daemonset: %#v", ds.Status)
				return false, nil
			}
			return true, nil
		})
		o.Expect(err).NotTo(o.HaveOccurred())
		framework.Logf("daemonset ready: %#v", ds.Status)

		pods, err := f.ClientSet.CoreV1().Pods(f.Namespace.Name).List(context.Background(), metav1.ListOptions{LabelSelector: labels.Set(ds.Spec.Selector.MatchLabels).String()})
		o.Expect(err).NotTo(o.HaveOccurred())
		o.Expect(len(pods.Items)).To(o.Equal(int(ds.Status.NumberAvailable)), fmt.Sprintf("%#v", pods.Items))

		// verify connectivity across pairs of pods in parallel
		// TODO: on large clusters this is O(N^2), we could potentially sample or split by topology
		var testFns []func() error
		protocols := []v1.Protocol{v1.ProtocolTCP, v1.ProtocolUDP}
		ports := []int{nodeTCPPort, nodeUDPPort}
		for j := range pods.Items {
			for i := range pods.Items {
				if i == j {
					continue
				}
				for k := range protocols {
					func(i, j, k int) {
						testFns = append(testFns, func() error {
							from := pods.Items[j]
							to := pods.Items[i]
							protocol := protocols[k]
							testingMsg := fmt.Sprintf("[%s: %s -> %s:%d]", protocol, from.Spec.NodeName, to.Spec.NodeName, ports[k])
							testMsg := fmt.Sprintf("%s-from-%s-to-%s", "hello", from.Status.PodIP, to.Status.PodIP)
							command, err := testRemoteConnectivityCommand(protocol, "localhost:"+strconv.Itoa(nodeTCPPort), to.Status.HostIP, ports[k], testMsg)
							if err != nil {
								return fmt.Errorf("test of %s failed: %v", testingMsg, err)
							}
							res, err := commandResult(f.ClientSet.CoreV1(), clientConfig, from.Namespace, from.Name, "webserver", []string{"/bin/sh", "-cex", strings.Join(command, " ")})
							if err != nil {
								return fmt.Errorf("test of %s failed: %v", testingMsg, err)
							}
							if res != `{"responses":["`+testMsg+`"]}` {
								return fmt.Errorf("test of %s failed, unexpected response: %s", testingMsg, res)
							}
							return nil
						})
					}(i, j, k)
				}
			}
		}
		errs := parallelTest(6, testFns)
		o.Expect(errs).To(o.Equal([]error(nil)))
	})
})

// parallelTest runs the provided fns in parallel with at most workers and returns an array of all
// non nil errors.
func parallelTest(workers int, fns []func() error) []error {
	var wg sync.WaitGroup
	work := make(chan func() error, workers)
	results := make(chan error, workers)

	go func() {
		for _, fn := range fns {
			work <- fn
			wg.Add(1)
		}
		close(work)
		wg.Wait()
		close(results)
	}()

	for i := 0; i < workers; i++ {
		go func() {
			for fn := range work {
				results <- fn()
				wg.Done()
			}
		}()
	}

	var errs []error
	for err := range results {
		if err != nil {
			errs = append(errs, err)
		}
	}
	return errs
}

func testRemoteConnectivityCommand(protocol v1.Protocol, localHostPort, host string, port int, echoMessage string) ([]string, error) {
	var protocolType string
	var dialCommand string
	switch protocol {
	case v1.ProtocolTCP:
		protocolType = "http"
		dialCommand = fmt.Sprintf("echo?msg=%s", echoMessage)
	case v1.ProtocolUDP:
		protocolType = "udp"
		dialCommand = fmt.Sprintf("echo%%20%s", echoMessage)
	default:
		return nil, fmt.Errorf("curl does not support protocol %s", protocol)
	}

	//func (config *NetworkingTestConfig) DialFromContainer(protocol, dialCommand, containerIP, targetIP string, containerHTTPPort, targetPort, maxTries, minTries int, expectedResponses sets.String) {
	// The current versions of curl included in CentOS and RHEL distros
	// misinterpret square brackets around IPv6 as globbing, so use the -g
	// argument to disable globbing to handle the IPv6 case.
	command := []string{
		"curl", "-g", "-q", "-s",
		fmt.Sprintf("'http://%s/dial?request=%s&protocol=%s&host=%s&port=%d&tries=1'",
			localHostPort,
			dialCommand,
			protocolType,
			host,
			port),
	}
	return command, nil
}

func testConnectivityCommand(protocol v1.Protocol, host string, port, timeout int) ([]string, error) {
	command := []string{
		"nc",
		"-vz",
		"-w", strconv.Itoa(timeout),
	}
	switch protocol {
	case v1.ProtocolTCP:
		command = append(command, "-t")
	case v1.ProtocolUDP:
		command = append(command, "-u")
	default:
		return nil, fmt.Errorf("nc does not support protocol %s", protocol)
	}
	command = append(command, host, strconv.Itoa(port))
	return command, nil
}

// commandContents fetches the result of invoking a command in the provided container from stdout.
func commandResult(podClient coreclientset.CoreV1Interface, podRESTConfig *rest.Config, ns, name, containerName string, command []string) (string, error) {
	u := podClient.RESTClient().Post().Resource("pods").Namespace(ns).Name(name).SubResource("exec").VersionedParams(&v1.PodExecOptions{
		Container: containerName,
		Stdout:    true,
		Stderr:    true,
		Command:   command,
	}, scheme.ParameterCodec).URL()

	e, err := remotecommand.NewSPDYExecutor(podRESTConfig, "POST", u)
	if err != nil {
		return "", fmt.Errorf("could not initialize a new SPDY executor: %v", err)
	}
	buf := &bytes.Buffer{}
	errBuf := &bytes.Buffer{}
	if err := e.Stream(remotecommand.StreamOptions{
		Stdout: buf,
		Stdin:  nil,
		Stderr: errBuf,
	}); err != nil {
		framework.Logf("exec error: %s", errBuf.String())
		return "", err
	}
	return buf.String(), nil
}
<commit_msg>test: remove apk add tcpdump to internal ports test<commit_after>package networking

import (
	"bytes"
	"context"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	appsv1 "k8s.io/api/apps/v1"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes/scheme"
	coreclientset "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/remotecommand"
	"k8s.io/kubernetes/test/e2e/framework"
	e2enetwork "k8s.io/kubernetes/test/e2e/framework/network"
	e2eskipper "k8s.io/kubernetes/test/e2e/framework/skipper"

	"github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"
)

const (
	nodeTCPPort = 9000
	nodeUDPPort = 9999
)

var _ = ginkgo.Describe("[sig-network] Internal connectivity", func() {
	f := framework.NewDefaultFramework("k8s-nettest")

	ginkgo.It("for TCP and UDP on ports 9000-9999 is allowed", func() {
		e2eskipper.SkipUnlessNodeCountIsAtLeast(2)
		clientConfig := f.ClientConfig()

		one := int64(0)
		ds := &appsv1.DaemonSet{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "webserver",
				Namespace: f.Namespace.Name,
			},
			Spec: appsv1.DaemonSetSpec{
				Selector: &metav1.LabelSelector{
					MatchLabels: map[string]string{
						"apps": "webserver",
					},
				},
				Template: v1.PodTemplateSpec{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							"apps": "webserver",
						},
					},
					Spec: v1.PodSpec{
						Tolerations: []v1.Toleration{
							{
								Key:      "node-role.kubernetes.io/master",
								Operator: v1.TolerationOpExists,
								Effect:   v1.TaintEffectNoSchedule,
							},
						},
						HostNetwork:                   true,
						TerminationGracePeriodSeconds: &one,
						Containers: []v1.Container{
							{
								Name:    "webserver",
								Image:   e2enetwork.NetexecImageName,
								Command: []string{"/agnhost", "netexec", fmt.Sprintf("--http-port=%v", nodeTCPPort), fmt.Sprintf("--udp-port=%v", nodeUDPPort)},
								Ports: []v1.ContainerPort{
									{Name: "tcp", ContainerPort: nodeTCPPort},
									{Name: "udp", ContainerPort: nodeUDPPort},
								},
								ReadinessProbe: &v1.Probe{
									InitialDelaySeconds: 10,
									Handler: v1.Handler{
										HTTPGet: &v1.HTTPGetAction{
											Port: intstr.FromInt(nodeTCPPort),
										},
									},
								},
								SecurityContext: &v1.SecurityContext{
									Capabilities: &v1.Capabilities{
										Add: []v1.Capability{"NET_RAW"},
									},
								},
							},
						},
					},
				},
			},
		}
		name := ds.Name
		ds, err := f.ClientSet.AppsV1().DaemonSets(f.Namespace.Name).Create(context.Background(), ds, metav1.CreateOptions{})
		o.Expect(err).NotTo(o.HaveOccurred())
		err = wait.PollImmediate(5*time.Second, 5*time.Minute, func() (bool, error) {
			ds, err = f.ClientSet.AppsV1().DaemonSets(f.Namespace.Name).Get(context.Background(), name, metav1.GetOptions{})
			if err != nil {
				framework.Logf("unable to retrieve daemonset: %v", err)
				return false, nil
			}
			if ds.Status.ObservedGeneration != ds.Generation || ds.Status.NumberAvailable == 0 || ds.Status.NumberAvailable != ds.Status.DesiredNumberScheduled {
				framework.Logf("waiting for daemonset: %#v", ds.Status)
				return false, nil
			}
			return true, nil
		})
		o.Expect(err).NotTo(o.HaveOccurred())
		framework.Logf("daemonset ready: %#v", ds.Status)

		pods, err := f.ClientSet.CoreV1().Pods(f.Namespace.Name).List(context.Background(), metav1.ListOptions{LabelSelector: labels.Set(ds.Spec.Selector.MatchLabels).String()})
		o.Expect(err).NotTo(o.HaveOccurred())
		o.Expect(len(pods.Items)).To(o.Equal(int(ds.Status.NumberAvailable)), fmt.Sprintf("%#v", pods.Items))

		// verify connectivity across pairs of pods in parallel
		// TODO: on large clusters this is O(N^2), we could potentially sample or split by topology
		var testFns []func() error
		protocols := []v1.Protocol{v1.ProtocolTCP, v1.ProtocolUDP}
		ports := []int{nodeTCPPort, nodeUDPPort}
		for j := range pods.Items {
			for i := range pods.Items {
				if i == j {
					continue
				}
				for k := range protocols {
					func(i, j, k int) {
						testFns = append(testFns, func() error {
							from := pods.Items[j]
							to := pods.Items[i]
							protocol := protocols[k]
							testingMsg := fmt.Sprintf("[%s: %s -> %s:%d]", protocol, from.Spec.NodeName, to.Spec.NodeName, ports[k])
							testMsg := fmt.Sprintf("%s-from-%s-to-%s", "hello", from.Status.PodIP, to.Status.PodIP)
							command, err := testRemoteConnectivityCommand(protocol, "localhost:"+strconv.Itoa(nodeTCPPort), to.Status.HostIP, ports[k], testMsg)
							if err != nil {
								return fmt.Errorf("test of %s failed: %v", testingMsg, err)
							}
							res, err := commandResult(f.ClientSet.CoreV1(), clientConfig, from.Namespace, from.Name, "webserver", []string{"/bin/sh", "-cex", strings.Join(command, " ")})
							if err != nil {
								return fmt.Errorf("test of %s failed: %v", testingMsg, err)
							}
							if res != `{"responses":["`+testMsg+`"]}` {
								return fmt.Errorf("test of %s failed, unexpected response: %s", testingMsg, res)
							}
							return nil
						})
					}(i, j, k)
				}
			}
		}
		errs := parallelTest(6, testFns)
		o.Expect(errs).To(o.Equal([]error(nil)))
	})
})

// parallelTest runs the provided fns in parallel with at most workers and returns an array of all
// non nil errors.
func parallelTest(workers int, fns []func() error) []error {
	var wg sync.WaitGroup
	work := make(chan func() error, workers)
	results := make(chan error, workers)

	go func() {
		for _, fn := range fns {
			work <- fn
			wg.Add(1)
		}
		close(work)
		wg.Wait()
		close(results)
	}()

	for i := 0; i < workers; i++ {
		go func() {
			for fn := range work {
				results <- fn()
				wg.Done()
			}
		}()
	}

	var errs []error
	for err := range results {
		if err != nil {
			errs = append(errs, err)
		}
	}
	return errs
}

func testRemoteConnectivityCommand(protocol v1.Protocol, localHostPort, host string, port int, echoMessage string) ([]string, error) {
	var protocolType string
	var dialCommand string
	switch protocol {
	case v1.ProtocolTCP:
		protocolType = "http"
		dialCommand = fmt.Sprintf("echo?msg=%s", echoMessage)
	case v1.ProtocolUDP:
		protocolType = "udp"
		dialCommand = fmt.Sprintf("echo%%20%s", echoMessage)
	default:
		return nil, fmt.Errorf("curl does not support protocol %s", protocol)
	}

	//func (config *NetworkingTestConfig) DialFromContainer(protocol, dialCommand, containerIP, targetIP string, containerHTTPPort, targetPort, maxTries, minTries int, expectedResponses sets.String) {
	// The current versions of curl included in CentOS and RHEL distros
	// misinterpret square brackets around IPv6 as globbing, so use the -g
	// argument to disable globbing to handle the IPv6 case.
	command := []string{
		"curl", "-g", "-q", "-s",
		fmt.Sprintf("'http://%s/dial?request=%s&protocol=%s&host=%s&port=%d&tries=1'",
			localHostPort,
			dialCommand,
			protocolType,
			host,
			port),
	}
	return command, nil
}

func testConnectivityCommand(protocol v1.Protocol, host string, port, timeout int) ([]string, error) {
	command := []string{
		"nc",
		"-vz",
		"-w", strconv.Itoa(timeout),
	}
	switch protocol {
	case v1.ProtocolTCP:
		command = append(command, "-t")
	case v1.ProtocolUDP:
		command = append(command, "-u")
	default:
		return nil, fmt.Errorf("nc does not support protocol %s", protocol)
	}
	command = append(command, host, strconv.Itoa(port))
	return command, nil
}

// commandContents fetches the result of invoking a command in the provided container from stdout.
func commandResult(podClient coreclientset.CoreV1Interface, podRESTConfig *rest.Config, ns, name, containerName string, command []string) (string, error) {
	u := podClient.RESTClient().Post().Resource("pods").Namespace(ns).Name(name).SubResource("exec").VersionedParams(&v1.PodExecOptions{
		Container: containerName,
		Stdout:    true,
		Stderr:    true,
		Command:   command,
	}, scheme.ParameterCodec).URL()

	e, err := remotecommand.NewSPDYExecutor(podRESTConfig, "POST", u)
	if err != nil {
		return "", fmt.Errorf("could not initialize a new SPDY executor: %v", err)
	}
	buf := &bytes.Buffer{}
	errBuf := &bytes.Buffer{}
	if err := e.Stream(remotecommand.StreamOptions{
		Stdout: buf,
		Stdin:  nil,
		Stderr: errBuf,
	}); err != nil {
		framework.Logf("exec error: %s", errBuf.String())
		return "", err
	}
	return buf.String(), nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2018, Google LLC All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tpmutil

import (
	"bytes"
	"io"
	"io/ioutil"
	"reflect"
	"testing"
)

type invalidPacked struct {
	A []int
	B uint32
}

func testEncodingInvalidSlices(t *testing.T, f func(io.Writer, interface{}) error) {
	d := ioutil.Discard

	// The packedSize function doesn't handle slices to anything other than bytes.
	var invalid []int
	if err := f(d, invalid); err == nil {
		t.Fatal("The packing function incorrectly succeeds for a slice of integers")
	}
	if err := f(d, &invalid); err == nil {
		t.Fatal("The packing function incorrectly succeeds for a pointer to a slice of integers")
	}

	invalid2 := invalidPacked{
		A: make([]int, 10),
		B: 137,
	}
	if err := f(d, invalid2); err == nil {
		t.Fatal("The packing function incorrectly succeeds for a struct that contains an integer slice")
	}
	if err := f(d, &invalid2); err == nil {
		t.Fatal("The packing function incorrectly succeeds for a pointer to a struct that contains an integer slice")
	}

	if err := f(d, d); err == nil {
		t.Fatal("The packing function incorrectly succeeds for a non-packable value")
	}
}

func TestEncodingPackTypeInvalid(t *testing.T) {
	f := func(w io.Writer, i interface{}) error {
		return packType(w, i)
	}

	testEncodingInvalidSlices(t, f)
}

type simplePacked struct {
	A uint32
	B uint32
}

type nestedPacked struct {
	SP simplePacked
	C  uint32
}

type nestedSlice struct {
	A uint32
	S U32Bytes
}

func TestEncodingPackType(t *testing.T) {
	buf := make([]byte, 10)
	inputs := []interface{}{
		uint32(3),
		buf,
		&buf,
		simplePacked{137, 138},
		nestedPacked{simplePacked{137, 138}, 139},
		nestedSlice{137, buf},
		[]byte(nil),
		RawBytes(buf),
	}
	for _, i := range inputs {
		if err := packType(ioutil.Discard, i); err != nil {
			t.Errorf("packType(%#v): %v", i, err)
		}
	}
}

func TestEncodingPackTypeWriteFail(t *testing.T) {
	u32WithOneByte := U32Bytes([]byte{1})
	u32Empty := U32Bytes([]byte(nil))

	tests := []struct {
		limit int
		in    interface{}
	}{
		{4, &u32WithOneByte},
		{3, &u32Empty},
	}
	for _, tt := range tests {
		if err := packType(&limitedDiscard{tt.limit}, tt.in); err == nil {
			t.Errorf("packType(%#v) with write size limit %d returned nil, want error", tt.in, tt.limit)
		}
	}
}

// limitedDiscard is an implementation of io.Writer that accepts a given number
// of bytes before returning errors.
type limitedDiscard struct {
	remaining int
}

// Write writes p to the limitedDiscard instance.
func (l *limitedDiscard) Write(p []byte) (n int, err error) {
	n = len(p)
	if n > l.remaining {
		n = l.remaining
		err = io.EOF
	}

	l.remaining -= n
	return
}

func TestEncodingCommandHeaderInvalidBody(t *testing.T) {
	var invalid []int
	ch := commandHeader{1, 0, 2}
	_, err := packWithHeader(ch, invalid)
	if err == nil {
		t.Fatal("packWithHeader incorrectly packed a body that with an invalid int slice member")
	}
}

func TestEncodingInvalidPack(t *testing.T) {
	var invalid []int
	ch := commandHeader{1, 0, 2}
	_, err := packWithHeader(ch, invalid)
	if err == nil {
		t.Fatal("packWithHeader incorrectly packed a body that with an invalid int slice member")
	}

	_, err = Pack(invalid)
	if err == nil {
		t.Fatal("pack incorrectly packed a slice of int")
	}
}

func TestEncodingCommandHeaderEncoding(t *testing.T) {
	ch := commandHeader{1, 0, 2}
	var c uint32 = 137
	in := c

	b, err := packWithHeader(ch, in)
	if err != nil {
		t.Fatal("Couldn't pack the bytes:", err)
	}

	var hdr commandHeader
	var size uint32
	if _, err := Unpack(b, &hdr, &size); err != nil {
		t.Fatal("Couldn't unpack the packed bytes")
	}

	if size != 137 {
		t.Fatal("Got the wrong size back")
	}
}

func TestEncodingInvalidUnpack(t *testing.T) {
	var i *uint32
	i = nil
	// The value ui is a serialization of uint32(0).
	ui := []byte{0, 0, 0, 0}
	uiBuf := bytes.NewBuffer(ui)
	if err := UnpackBuf(uiBuf, i); err == nil {
		t.Fatal("UnpackBuf incorrectly deserialized into a nil pointer")
	}

	var ii uint32
	if err := UnpackBuf(uiBuf, ii); err == nil {
		t.Fatal("UnpackBuf incorrectly deserialized into a non pointer")
	}

	var b U32Bytes
	var empty []byte
	emptyBuf := bytes.NewBuffer(empty)
	if err := UnpackBuf(emptyBuf, &b); err == nil {
		t.Fatal("UnpackBuf incorrectly deserialized an empty byte array into U32Bytes")
	}

	// Try to deserialize a byte array that has a length but not enough bytes.
	// The slice ui represents uint32(1), which is the length of an empty byte array.
	ui2 := []byte{0, 0, 0, 1}
	uiBuf2 := bytes.NewBuffer(ui2)
	if err := UnpackBuf(uiBuf2, &b); err == nil {
		t.Fatal("UnpackBuf incorrectly deserialized a byte array that didn't have enough bytes available")
	}

	var iii []int
	ui3 := []byte{0, 0, 0, 1}
	uiBuf3 := bytes.NewBuffer(ui3)
	if err := UnpackBuf(uiBuf3, &iii); err == nil {
		t.Fatal("UnpackBuf incorrectly deserialized into a slice of ints (only byte slices are supported)")
	}

}

func TestSelfMarshaler(t *testing.T) {
	var empty16 U16Bytes
	var empty32 U32Bytes
	subTests := []struct {
		encoded []byte
		decoded interface{}
	}{
		{[]byte{0, 0}, &empty16},
		{[]byte{0, 1, 137}, &empty16},
		{[]byte{0, 0, 0, 0}, &empty32},
		{[]byte{0, 0, 0, 1, 137}, &empty32},
	}
	for _, st := range subTests {
		t.Logf("Attempting to Marshal/Unmarshal %#v into %T", st.encoded, st.decoded)
		buffer := bytes.NewBuffer(st.encoded)
		if err := UnpackBuf(buffer, st.decoded); err != nil {
			t.Fatalf("UnpackBuf failed: %v", err)
		}
		packed, err := Pack(st.decoded)
		if err != nil {
			t.Fatalf("Pack failed: %v", err)
		}
		if !bytes.Equal(packed, st.encoded) {
			t.Fatalf("Pack failed: got %#v, want: %#v", packed, st.encoded)
		}
	}
}

func TestEncodingUnpack(t *testing.T) {
	// Deserialize the empty byte array.
	var b U32Bytes
	// The slice ui represents uint32(0), which is the length of an empty byte array.
	ui := []byte{0, 0, 0, 0}
	uiBuf := bytes.NewBuffer(ui)
	if err := UnpackBuf(uiBuf, &b); err != nil {
		t.Fatal("UnpackBuf failed to unpack the empty byte array")
	}

	// A byte slice of length 1 with a single entry: b[0] == 137
	ui2 := []byte{0, 0, 0, 1, 137}
	uiBuf2 := bytes.NewBuffer(ui2)
	if err := UnpackBuf(uiBuf2, &b); err != nil {
		t.Fatal("UnpackBuf failed to unpack a byte array with a single value in it")
	}

	if !bytes.Equal([]byte(b), []byte{137}) {
		t.Fatal("UnpackBuf unpacked a small byte array incorrectly")
	}

	sp := simplePacked{137, 138}
	bsp, err := Pack(sp)
	if err != nil {
		t.Fatal("Couldn't pack a simple struct:", err)
	}
	var sp2 simplePacked
	if _, err := Unpack(bsp, &sp2); err != nil {
		t.Fatal("Couldn't unpack a simple struct:", err)
	}

	if sp.A != sp2.A || sp.B != sp2.B {
		t.Fatal("Unpacked simple struct didn't match the original")
	}

	// Try unpacking a version that's missing a byte at the end.
	if _, err := Unpack(bsp[:len(bsp)-1], &sp2); err == nil {
		t.Fatal("unpack incorrectly unpacked from a byte array that didn't have enough values")
	}

	np := nestedPacked{sp, 139}
	bnp, err := Pack(np)
	if err != nil {
		t.Fatal("Couldn't pack a nested struct")
	}
	var np2 nestedPacked
	if _, err := Unpack(bnp, &np2); err != nil {
		t.Fatal("Couldn't unpack a nested struct:", err)
	}
	if np.SP.A != np2.SP.A || np.SP.B != np2.SP.B || np.C != np2.C {
		t.Fatal("Unpacked nested struct didn't match the original")
	}

	ns := nestedSlice{137, b}
	bns, err := Pack(&ns)
	if err != nil {
		t.Fatal("Couldn't pack a struct with a nested byte slice:", err)
	}
	var ns2 nestedSlice
	if _, err := Unpack(bns, &ns2); err != nil {
		t.Fatal("Couldn't unpacked a struct with a nested slice:", err)
	}
	if ns.A != ns2.A || !bytes.Equal(ns.S, ns2.S) {
		t.Logf("orginal = %+v", ns)
		t.Logf("decoded = %+v", ns2)
		t.Fatal("Unpacked struct with nested slice didn't match the original")
	}

	var hs []Handle
	if _, err := Unpack([]byte{0, 3, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}, &hs); err != nil {
		t.Fatal("Couldn't unpack a list of Handles:", err)
	}
	if want := []Handle{0x01020304, 0x05060708, 0x090a0b0c}; !reflect.DeepEqual(want, hs) {
		t.Fatalf("Unpacking []Handle: got %v, want %v", hs, want)
	}
}

func TestPartialUnpack(t *testing.T) {
	u1, u2 := uint32(1), uint32(2)
	buf, err := Pack(u1, u2)
	if err != nil {
		t.Fatalf("packing uint32 value: %v", err)
	}

	var gu1, gu2 uint32
	read1, err := Unpack(buf, &gu1)
	if err != nil {
		t.Fatalf("unpacking first uint32 value: %v", err)
	}
	if gu1 != u1 {
		t.Errorf("first unpacked value: got %d, want %d", gu1, u1)
	}
	read2, err := Unpack(buf[read1:], &gu2)
	if err != nil {
		t.Fatalf("unpacking second uint32 value: %v", err)
	}
	if gu2 != u2 {
		t.Errorf("second unpacked value: got %d, want %d", gu2, u2)
	}

	if read1+read2 != len(buf) {
		t.Errorf("sum of bytes read doesn't ad up to total packed size: got %d+%d=%d, want %d", read1, read2, read1+read2, len(buf))
	}
}

func TestUnpackHandlesArea(t *testing.T) {
	buf := []byte{
		0, 2,
		0, 0, 0, 1,
		0, 0, 5, 57,
	}
	var out []Handle

	if _, err := Unpack(buf, &out); err != nil {
		t.Fatalf("Unpack(%v, %T) failed: %v", buf, &out, err)
	}
	if want := []Handle{1, 1337}; !reflect.DeepEqual(out, want) {
		t.Errorf("Unpack(%v, %T): %T = %v, want %v", buf, &out, out, out, want)
	}
}

func TestUnpackMalformedBytes(t *testing.T) {
	// buf is malformed because the leading size prefix is illegally large.
	buf := []byte{0xff, 0xff, 0xff, 0xff, 0x20}

	var u32 U32Bytes
	if _, err := Unpack(buf, &u32); err != bytes.ErrTooLarge {
		t.Errorf("Unpack(U32Bytes{}) returned %q, want %q", err, bytes.ErrTooLarge)
	}
	var u16 U16Bytes
	if _, err := Unpack(buf, &u16); err != io.ErrUnexpectedEOF {
		t.Errorf("Unpack(U16Bytes{}) returned %q, want %q", err, io.ErrUnexpectedEOF)
	}
}
<commit_msg>Remove Invalid Pack test for ioutil.Discard (#241)<commit_after>// Copyright (c) 2018, Google LLC All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tpmutil

import (
	"bytes"
	"io"
	"io/ioutil"
	"reflect"
	"testing"
)

type invalidPacked struct {
	A []int
	B uint32
}

func TestEncodingPackTypeInvalid(t *testing.T) {
	d := ioutil.Discard

	// The packedSize function doesn't handle slices to anything other than bytes.
	var invalid []int
	if err := packType(d, invalid); err == nil {
		t.Fatal("The packing function incorrectly succeeds for a slice of integers")
	}
	if err := packType(d, &invalid); err == nil {
		t.Fatal("The packing function incorrectly succeeds for a pointer to a slice of integers")
	}

	invalid2 := invalidPacked{
		A: make([]int, 10),
		B: 137,
	}
	if err := packType(d, invalid2); err == nil {
		t.Fatal("The packing function incorrectly succeeds for a struct that contains an integer slice")
	}
	if err := packType(d, &invalid2); err == nil {
		t.Fatal("The packing function incorrectly succeeds for a pointer to a struct that contains an integer slice")
	}
}

type simplePacked struct {
	A uint32
	B uint32
}

type nestedPacked struct {
	SP simplePacked
	C  uint32
}

type nestedSlice struct {
	A uint32
	S U32Bytes
}

func TestEncodingPackType(t *testing.T) {
	buf := make([]byte, 10)
	inputs := []interface{}{
		uint32(3),
		buf,
		&buf,
		simplePacked{137, 138},
		nestedPacked{simplePacked{137, 138}, 139},
		nestedSlice{137, buf},
		[]byte(nil),
		RawBytes(buf),
	}
	for _, i := range inputs {
		if err := packType(ioutil.Discard, i); err != nil {
			t.Errorf("packType(%#v): %v", i, err)
		}
	}
}

func TestEncodingPackTypeWriteFail(t *testing.T) {
	u32WithOneByte := U32Bytes([]byte{1})
	u32Empty := U32Bytes([]byte(nil))

	tests := []struct {
		limit int
		in    interface{}
	}{
		{4, &u32WithOneByte},
		{3, &u32Empty},
	}
	for _, tt := range tests {
		if err := packType(&limitedDiscard{tt.limit}, tt.in); err == nil {
			t.Errorf("packType(%#v) with write size limit %d returned nil, want error", tt.in, tt.limit)
		}
	}
}

// limitedDiscard is an implementation of io.Writer that accepts a given number
// of bytes before returning errors.
type limitedDiscard struct {
	remaining int
}

// Write writes p to the limitedDiscard instance.
func (l *limitedDiscard) Write(p []byte) (n int, err error) {
	n = len(p)
	if n > l.remaining {
		n = l.remaining
		err = io.EOF
	}

	l.remaining -= n
	return
}

func TestEncodingCommandHeaderInvalidBody(t *testing.T) {
	var invalid []int
	ch := commandHeader{1, 0, 2}
	_, err := packWithHeader(ch, invalid)
	if err == nil {
		t.Fatal("packWithHeader incorrectly packed a body that with an invalid int slice member")
	}
}

func TestEncodingInvalidPack(t *testing.T) {
	var invalid []int
	ch := commandHeader{1, 0, 2}
	_, err := packWithHeader(ch, invalid)
	if err == nil {
		t.Fatal("packWithHeader incorrectly packed a body that with an invalid int slice member")
	}

	_, err = Pack(invalid)
	if err == nil {
		t.Fatal("pack incorrectly packed a slice of int")
	}
}

func TestEncodingCommandHeaderEncoding(t *testing.T) {
	ch := commandHeader{1, 0, 2}
	var c uint32 = 137
	in := c

	b, err := packWithHeader(ch, in)
	if err != nil {
		t.Fatal("Couldn't pack the bytes:", err)
	}

	var hdr commandHeader
	var size uint32
	if _, err := Unpack(b, &hdr, &size); err != nil {
		t.Fatal("Couldn't unpack the packed bytes")
	}

	if size != 137 {
		t.Fatal("Got the wrong size back")
	}
}

func TestEncodingInvalidUnpack(t *testing.T) {
	var i *uint32
	i = nil
	// The value ui is a serialization of uint32(0).
	ui := []byte{0, 0, 0, 0}
	uiBuf := bytes.NewBuffer(ui)
	if err := UnpackBuf(uiBuf, i); err == nil {
		t.Fatal("UnpackBuf incorrectly deserialized into a nil pointer")
	}

	var ii uint32
	if err := UnpackBuf(uiBuf, ii); err == nil {
		t.Fatal("UnpackBuf incorrectly deserialized into a non pointer")
	}

	var b U32Bytes
	var empty []byte
	emptyBuf := bytes.NewBuffer(empty)
	if err := UnpackBuf(emptyBuf, &b); err == nil {
		t.Fatal("UnpackBuf incorrectly deserialized an empty byte array into U32Bytes")
	}

	// Try to deserialize a byte array that has a length but not enough bytes.
	// The slice ui represents uint32(1), which is the length of an empty byte array.
	ui2 := []byte{0, 0, 0, 1}
	uiBuf2 := bytes.NewBuffer(ui2)
	if err := UnpackBuf(uiBuf2, &b); err == nil {
		t.Fatal("UnpackBuf incorrectly deserialized a byte array that didn't have enough bytes available")
	}

	var iii []int
	ui3 := []byte{0, 0, 0, 1}
	uiBuf3 := bytes.NewBuffer(ui3)
	if err := UnpackBuf(uiBuf3, &iii); err == nil {
		t.Fatal("UnpackBuf incorrectly deserialized into a slice of ints (only byte slices are supported)")
	}

}

func TestSelfMarshaler(t *testing.T) {
	var empty16 U16Bytes
	var empty32 U32Bytes
	subTests := []struct {
		encoded []byte
		decoded interface{}
	}{
		{[]byte{0, 0}, &empty16},
		{[]byte{0, 1, 137}, &empty16},
		{[]byte{0, 0, 0, 0}, &empty32},
		{[]byte{0, 0, 0, 1, 137}, &empty32},
	}
	for _, st := range subTests {
		t.Logf("Attempting to Marshal/Unmarshal %#v into %T", st.encoded, st.decoded)
		buffer := bytes.NewBuffer(st.encoded)
		if err := UnpackBuf(buffer, st.decoded); err != nil {
			t.Fatalf("UnpackBuf failed: %v", err)
		}
		packed, err := Pack(st.decoded)
		if err != nil {
			t.Fatalf("Pack failed: %v", err)
		}
		if !bytes.Equal(packed, st.encoded) {
			t.Fatalf("Pack failed: got %#v, want: %#v", packed, st.encoded)
		}
	}
}

func TestEncodingUnpack(t *testing.T) {
	// Deserialize the empty byte array.
	var b U32Bytes
	// The slice ui represents uint32(0), which is the length of an empty byte array.
	ui := []byte{0, 0, 0, 0}
	uiBuf := bytes.NewBuffer(ui)
	if err := UnpackBuf(uiBuf, &b); err != nil {
		t.Fatal("UnpackBuf failed to unpack the empty byte array")
	}

	// A byte slice of length 1 with a single entry: b[0] == 137
	ui2 := []byte{0, 0, 0, 1, 137}
	uiBuf2 := bytes.NewBuffer(ui2)
	if err := UnpackBuf(uiBuf2, &b); err != nil {
		t.Fatal("UnpackBuf failed to unpack a byte array with a single value in it")
	}

	if !bytes.Equal([]byte(b), []byte{137}) {
		t.Fatal("UnpackBuf unpacked a small byte array incorrectly")
	}

	sp := simplePacked{137, 138}
	bsp, err := Pack(sp)
	if err != nil {
		t.Fatal("Couldn't pack a simple struct:", err)
	}
	var sp2 simplePacked
	if _, err := Unpack(bsp, &sp2); err != nil {
		t.Fatal("Couldn't unpack a simple struct:", err)
	}

	if sp.A != sp2.A || sp.B != sp2.B {
		t.Fatal("Unpacked simple struct didn't match the original")
	}

	// Try unpacking a version that's missing a byte at the end.
	if _, err := Unpack(bsp[:len(bsp)-1], &sp2); err == nil {
		t.Fatal("unpack incorrectly unpacked from a byte array that didn't have enough values")
	}

	np := nestedPacked{sp, 139}
	bnp, err := Pack(np)
	if err != nil {
		t.Fatal("Couldn't pack a nested struct")
	}
	var np2 nestedPacked
	if _, err := Unpack(bnp, &np2); err != nil {
		t.Fatal("Couldn't unpack a nested struct:", err)
	}
	if np.SP.A != np2.SP.A || np.SP.B != np2.SP.B || np.C != np2.C {
		t.Fatal("Unpacked nested struct didn't match the original")
	}

	ns := nestedSlice{137, b}
	bns, err := Pack(&ns)
	if err != nil {
		t.Fatal("Couldn't pack a struct with a nested byte slice:", err)
	}
	var ns2 nestedSlice
	if _, err := Unpack(bns, &ns2); err != nil {
		t.Fatal("Couldn't unpacked a struct with a nested slice:", err)
	}
	if ns.A != ns2.A || !bytes.Equal(ns.S, ns2.S) {
		t.Logf("orginal = %+v", ns)
		t.Logf("decoded = %+v", ns2)
		t.Fatal("Unpacked struct with nested slice didn't match the original")
	}

	var hs []Handle
	if _, err := Unpack([]byte{0, 3, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}, &hs); err != nil {
		t.Fatal("Couldn't unpack a list of Handles:", err)
	}
	if want := []Handle{0x01020304, 0x05060708, 0x090a0b0c}; !reflect.DeepEqual(want, hs) {
		t.Fatalf("Unpacking []Handle: got %v, want %v", hs, want)
	}
}

func TestPartialUnpack(t *testing.T) {
	u1, u2 := uint32(1), uint32(2)
	buf, err := Pack(u1, u2)
	if err != nil {
		t.Fatalf("packing uint32 value: %v", err)
	}

	var gu1, gu2 uint32
	read1, err := Unpack(buf, &gu1)
	if err != nil {
		t.Fatalf("unpacking first uint32 value: %v", err)
	}
	if gu1 != u1 {
		t.Errorf("first unpacked value: got %d, want %d", gu1, u1)
	}
	read2, err := Unpack(buf[read1:], &gu2)
	if err != nil {
		t.Fatalf("unpacking second uint32 value: %v", err)
	}
	if gu2 != u2 {
		t.Errorf("second unpacked value: got %d, want %d", gu2, u2)
	}

	if read1+read2 != len(buf) {
		t.Errorf("sum of bytes read doesn't ad up to total packed size: got %d+%d=%d, want %d", read1, read2, read1+read2, len(buf))
	}
}

func TestUnpackHandlesArea(t *testing.T) {
	buf := []byte{
		0, 2,
		0, 0, 0, 1,
		0, 0, 5, 57,
	}
	var out []Handle

	if _, err := Unpack(buf, &out); err != nil {
		t.Fatalf("Unpack(%v, %T) failed: %v", buf, &out, err)
	}
	if want := []Handle{1, 1337}; !reflect.DeepEqual(out, want) {
		t.Errorf("Unpack(%v, %T): %T = %v, want %v", buf, &out, out, out, want)
	}
}

func TestUnpackMalformedBytes(t *testing.T) {
	// buf is malformed because the leading size prefix is illegally large.
	buf := []byte{0xff, 0xff, 0xff, 0xff, 0x20}

	var u32 U32Bytes
	if _, err := Unpack(buf, &u32); err != bytes.ErrTooLarge {
		t.Errorf("Unpack(U32Bytes{}) returned %q, want %q", err, bytes.ErrTooLarge)
	}
	var u16 U16Bytes
	if _, err := Unpack(buf, &u16); err != io.ErrUnexpectedEOF {
		t.Errorf("Unpack(U16Bytes{}) returned %q, want %q", err, io.ErrUnexpectedEOF)
	}
}
<|endoftext|>


<commit_before>package etcd

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"path"
	"strings"
	"sync"
	"time"
)

// Errors introduced by handling requests
var (
	ErrRequestCancelled = errors.New("sending request is cancelled")
)

type RawRequest struct {
	Method       string
	RelativePath string
	Values       url.Values
	Cancel       <-chan bool
}

// NewRawRequest returns a new RawRequest
func NewRawRequest(method, relativePath string, values url.Values, cancel <-chan bool) *RawRequest {
	return &RawRequest{
		Method:       method,
		RelativePath: relativePath,
		Values:       values,
		Cancel:       cancel,
	}
}

// getCancelable issues a cancelable GET request
func (c *Client) getCancelable(key string, options Options,
	cancel <-chan bool) (*RawResponse, error) {
	logger.Debugf("get %s [%s]", key, c.cluster.Leader)
	p := keyToPath(key)

	// If consistency level is set to STRONG, append
	// the `consistent` query string.
	if c.config.Consistency == STRONG_CONSISTENCY {
		options["consistent"] = true
	}

	str, err := options.toParameters(VALID_GET_OPTIONS)
	if err != nil {
		return nil, err
	}
	p += str

	req := NewRawRequest("GET", p, nil, cancel)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// get issues a GET request
func (c *Client) get(key string, options Options) (*RawResponse, error) {
	return c.getCancelable(key, options, nil)
}

// put issues a PUT request
func (c *Client) put(key string, value string, ttl uint64,
	options Options) (*RawResponse, error) {

	logger.Debugf("put %s, %s, ttl: %d, [%s]", key, value, ttl, c.cluster.Leader)
	p := keyToPath(key)

	str, err := options.toParameters(VALID_PUT_OPTIONS)
	if err != nil {
		return nil, err
	}
	p += str

	req := NewRawRequest("PUT", p, buildValues(value, ttl), nil)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// post issues a POST request
func (c *Client) post(key string, value string, ttl uint64) (*RawResponse, error) {
	logger.Debugf("post %s, %s, ttl: %d, [%s]", key, value, ttl, c.cluster.Leader)
	p := keyToPath(key)

	req := NewRawRequest("POST", p, buildValues(value, ttl), nil)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// delete issues a DELETE request
func (c *Client) delete(key string, options Options) (*RawResponse, error) {
	logger.Debugf("delete %s [%s]", key, c.cluster.Leader)
	p := keyToPath(key)

	str, err := options.toParameters(VALID_DELETE_OPTIONS)
	if err != nil {
		return nil, err
	}
	p += str

	req := NewRawRequest("DELETE", p, nil, nil)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// SendRequest sends a HTTP request and returns a Response as defined by etcd
func (c *Client) SendRequest(rr *RawRequest) (*RawResponse, error) {

	var req *http.Request
	var resp *http.Response
	var httpPath string
	var err error
	var respBody []byte

	var numReqs = 1

	checkRetry := c.CheckRetry
	if checkRetry == nil {
		checkRetry = DefaultCheckRetry
	}

	cancelled := make(chan bool, 1)
	reqLock := new(sync.Mutex)

	if rr.Cancel != nil {
		cancelRoutine := make(chan bool)
		defer close(cancelRoutine)

		go func() {
			select {
			case <-rr.Cancel:
				cancelled <- true
				logger.Debug("send.request is cancelled")
			case <-cancelRoutine:
				return
			}

			// Repeat canceling request until this thread is stopped
			// because we have no idea about whether it succeeds.
			for {
				reqLock.Lock()
				c.httpClient.Transport.(*http.Transport).CancelRequest(req)
				reqLock.Unlock()

				select {
				case <-time.After(100 * time.Millisecond):
				case <-cancelRoutine:
					return
				}
			}
		}()
	}

	// If we connect to a follower and consistency is required, retry until
	// we connect to a leader
	sleep := 25 * time.Millisecond
	maxSleep := time.Second

	for attempt := 0; ; attempt++ {
		if attempt > 0 {
			select {
			case <-cancelled:
				return nil, ErrRequestCancelled
			case <-time.After(sleep):
				sleep = sleep * 2
				if sleep > maxSleep {
					sleep = maxSleep
				}
			}
		}

		logger.Debug("Connecting to etcd: attempt ", attempt+1, " for ", rr.RelativePath)

		if rr.Method == "GET" && c.config.Consistency == WEAK_CONSISTENCY {
			// If it's a GET and consistency level is set to WEAK,
			// then use a random machine.
			httpPath = c.getHttpPath(true, rr.RelativePath)
		} else {
			// Else use the leader.
			httpPath = c.getHttpPath(false, rr.RelativePath)
		}

		// Return a cURL command if curlChan is set
		if c.cURLch != nil {
			command := fmt.Sprintf("curl -X %s %s", rr.Method, httpPath)
			for key, value := range rr.Values {
				command += fmt.Sprintf(" -d %s=%s", key, value[0])
			}
			c.sendCURL(command)
		}

		logger.Debug("send.request.to ", httpPath, " | method ", rr.Method)

		req, err := func() (*http.Request, error) {
			reqLock.Lock()
			defer reqLock.Unlock()

			if rr.Values == nil {
				if req, err = http.NewRequest(rr.Method, httpPath, nil); err != nil {
					return nil, err
				}
			} else {
				body := strings.NewReader(rr.Values.Encode())
				if req, err = http.NewRequest(rr.Method, httpPath, body); err != nil {
					return nil, err
				}

				req.Header.Set("Content-Type",
					"application/x-www-form-urlencoded; param=value")
			}
			return req, nil
		}()

		if err != nil {
			return nil, err
		}

		resp, err = c.httpClient.Do(req)
		defer func() {
			if resp != nil {
				resp.Body.Close()
			}
		}()

		// If the request was cancelled, return ErrRequestCancelled directly
		select {
		case <-cancelled:
			return nil, ErrRequestCancelled
		default:
		}

		numReqs++

		// network error, change a machine!
		if err != nil {
			logger.Debug("network error: ", err.Error())
			lastResp := http.Response{}
			if checkErr := checkRetry(c.cluster, numReqs, lastResp, err); checkErr != nil {
				return nil, checkErr
			}

			c.cluster.switchLeader(attempt % len(c.cluster.Machines))
			continue
		}

		// if there is no error, it should receive response
		logger.Debug("recv.response.from ", httpPath)

		if validHttpStatusCode[resp.StatusCode] {
			// try to read byte code and break the loop
			respBody, err = ioutil.ReadAll(resp.Body)
			if err == nil {
				logger.Debug("recv.success ", httpPath)
				break
			}
			// ReadAll error may be caused due to cancel request
			select {
			case <-cancelled:
				return nil, ErrRequestCancelled
			default:
			}

			if err == io.ErrUnexpectedEOF {
				// underlying connection was closed prematurely, probably by timeout
				// TODO: empty body or unexpectedEOF can cause http.Transport to get hosed;
				// this allows the client to detect that and take evasive action. Need
				// to revisit once code.google.com/p/go/issues/detail?id=8648 gets fixed.
				respBody = []byte{}
				break
			}
		}

		// if resp is TemporaryRedirect, set the new leader and retry
		if resp.StatusCode == http.StatusTemporaryRedirect {
			u, err := resp.Location()

			if err != nil {
				logger.Warning(err)
			} else {
				// Update cluster leader based on redirect location
				// because it should point to the leader address
				c.cluster.updateLeaderFromURL(u)
				logger.Debug("recv.response.relocate ", u.String())
			}
			resp.Body.Close()
			continue
		}

		if checkErr := checkRetry(c.cluster, numReqs, *resp,
			errors.New("Unexpected HTTP status code")); checkErr != nil {
			return nil, checkErr
		}
		resp.Body.Close()
	}

	r := &RawResponse{
		StatusCode: resp.StatusCode,
		Body:       respBody,
		Header:     resp.Header,
	}

	return r, nil
}

// DefaultCheckRetry defines the retrying behaviour for bad HTTP requests
// If we have retried 2 * machine number, stop retrying.
// If status code is InternalServerError, sleep for 200ms.
func DefaultCheckRetry(cluster *Cluster, numReqs int, lastResp http.Response,
	err error) error {

	if isEmptyResponse(lastResp) {
		if !isConnectionError(err) {
			return err
		}
	} else if !shouldRetry(lastResp) {
		body := []byte("nil")
		if lastResp.Body != nil {
			if b, err := ioutil.ReadAll(lastResp.Body); err == nil {
				body = b
			}
		}
		errStr := fmt.Sprintf("unhandled http status [%s] with body [%s]", http.StatusText(lastResp.StatusCode), body)
		return newError(ErrCodeUnhandledHTTPStatus, errStr, 0)
	}

	if numReqs > 2*len(cluster.Machines) {
		errStr := fmt.Sprintf("failed to propose on members %v twice [last error: %v]", cluster.Machines, err)
		return newError(ErrCodeEtcdNotReachable, errStr, 0)
	}
	if shouldRetry(lastResp) {
		// sleep some time and expect leader election finish
		time.Sleep(time.Millisecond * 200)
	}

	logger.Warning("bad response status code", lastResp.StatusCode)
	return nil
}

func isEmptyResponse(r http.Response) bool { return r.StatusCode == 0 }

func isConnectionError(err error) bool {
	_, ok := err.(*net.OpError)
	return ok
}

// shouldRetry returns whether the reponse deserves retry.
func shouldRetry(r http.Response) bool {
	// TODO: only retry when the cluster is in leader election
	// We cannot do it exactly because etcd doesn't support it well.
	return r.StatusCode == http.StatusInternalServerError
}

func (c *Client) getHttpPath(random bool, s ...string) string {
	var machine string
	if random {
		machine = c.cluster.Machines[rand.Intn(len(c.cluster.Machines))]
	} else {
		machine = c.cluster.Leader
	}

	fullPath := machine + "/" + version
	for _, seg := range s {
		fullPath = fullPath + "/" + seg
	}

	return fullPath
}

// buildValues builds a url.Values map according to the given value and ttl
func buildValues(value string, ttl uint64) url.Values {
	v := url.Values{}

	if value != "" {
		v.Set("value", value)
	}

	if ttl > 0 {
		v.Set("ttl", fmt.Sprintf("%v", ttl))
	}

	return v
}

// convert key string to http path exclude version, including URL escaping
// for example: key[foo] -> path[keys/foo]
// key[/%z] -> path[keys/%25z]
// key[/] -> path[keys/]
func keyToPath(key string) string {
	// URL-escape our key, except for slashes
	p := strings.Replace(url.QueryEscape(path.Join("keys", key)), "%2F", "/", -1)

	// corner case: if key is "/" or "//" ect
	// path join will clear the tailing "/"
	// we need to add it back
	if p == "keys" {
		p = "keys/"
	}

	return p
}
<commit_msg>requests.go: alwasy retry when failing to get response<commit_after>package etcd

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"path"
	"strings"
	"sync"
	"time"
)

// Errors introduced by handling requests
var (
	ErrRequestCancelled = errors.New("sending request is cancelled")
)

type RawRequest struct {
	Method       string
	RelativePath string
	Values       url.Values
	Cancel       <-chan bool
}

// NewRawRequest returns a new RawRequest
func NewRawRequest(method, relativePath string, values url.Values, cancel <-chan bool) *RawRequest {
	return &RawRequest{
		Method:       method,
		RelativePath: relativePath,
		Values:       values,
		Cancel:       cancel,
	}
}

// getCancelable issues a cancelable GET request
func (c *Client) getCancelable(key string, options Options,
	cancel <-chan bool) (*RawResponse, error) {
	logger.Debugf("get %s [%s]", key, c.cluster.Leader)
	p := keyToPath(key)

	// If consistency level is set to STRONG, append
	// the `consistent` query string.
	if c.config.Consistency == STRONG_CONSISTENCY {
		options["consistent"] = true
	}

	str, err := options.toParameters(VALID_GET_OPTIONS)
	if err != nil {
		return nil, err
	}
	p += str

	req := NewRawRequest("GET", p, nil, cancel)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// get issues a GET request
func (c *Client) get(key string, options Options) (*RawResponse, error) {
	return c.getCancelable(key, options, nil)
}

// put issues a PUT request
func (c *Client) put(key string, value string, ttl uint64,
	options Options) (*RawResponse, error) {

	logger.Debugf("put %s, %s, ttl: %d, [%s]", key, value, ttl, c.cluster.Leader)
	p := keyToPath(key)

	str, err := options.toParameters(VALID_PUT_OPTIONS)
	if err != nil {
		return nil, err
	}
	p += str

	req := NewRawRequest("PUT", p, buildValues(value, ttl), nil)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// post issues a POST request
func (c *Client) post(key string, value string, ttl uint64) (*RawResponse, error) {
	logger.Debugf("post %s, %s, ttl: %d, [%s]", key, value, ttl, c.cluster.Leader)
	p := keyToPath(key)

	req := NewRawRequest("POST", p, buildValues(value, ttl), nil)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// delete issues a DELETE request
func (c *Client) delete(key string, options Options) (*RawResponse, error) {
	logger.Debugf("delete %s [%s]", key, c.cluster.Leader)
	p := keyToPath(key)

	str, err := options.toParameters(VALID_DELETE_OPTIONS)
	if err != nil {
		return nil, err
	}
	p += str

	req := NewRawRequest("DELETE", p, nil, nil)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// SendRequest sends a HTTP request and returns a Response as defined by etcd
func (c *Client) SendRequest(rr *RawRequest) (*RawResponse, error) {

	var req *http.Request
	var resp *http.Response
	var httpPath string
	var err error
	var respBody []byte

	var numReqs = 1

	checkRetry := c.CheckRetry
	if checkRetry == nil {
		checkRetry = DefaultCheckRetry
	}

	cancelled := make(chan bool, 1)
	reqLock := new(sync.Mutex)

	if rr.Cancel != nil {
		cancelRoutine := make(chan bool)
		defer close(cancelRoutine)

		go func() {
			select {
			case <-rr.Cancel:
				cancelled <- true
				logger.Debug("send.request is cancelled")
			case <-cancelRoutine:
				return
			}

			// Repeat canceling request until this thread is stopped
			// because we have no idea about whether it succeeds.
			for {
				reqLock.Lock()
				c.httpClient.Transport.(*http.Transport).CancelRequest(req)
				reqLock.Unlock()

				select {
				case <-time.After(100 * time.Millisecond):
				case <-cancelRoutine:
					return
				}
			}
		}()
	}

	// If we connect to a follower and consistency is required, retry until
	// we connect to a leader
	sleep := 25 * time.Millisecond
	maxSleep := time.Second

	for attempt := 0; ; attempt++ {
		if attempt > 0 {
			select {
			case <-cancelled:
				return nil, ErrRequestCancelled
			case <-time.After(sleep):
				sleep = sleep * 2
				if sleep > maxSleep {
					sleep = maxSleep
				}
			}
		}

		logger.Debug("Connecting to etcd: attempt ", attempt+1, " for ", rr.RelativePath)

		if rr.Method == "GET" && c.config.Consistency == WEAK_CONSISTENCY {
			// If it's a GET and consistency level is set to WEAK,
			// then use a random machine.
			httpPath = c.getHttpPath(true, rr.RelativePath)
		} else {
			// Else use the leader.
			httpPath = c.getHttpPath(false, rr.RelativePath)
		}

		// Return a cURL command if curlChan is set
		if c.cURLch != nil {
			command := fmt.Sprintf("curl -X %s %s", rr.Method, httpPath)
			for key, value := range rr.Values {
				command += fmt.Sprintf(" -d %s=%s", key, value[0])
			}
			c.sendCURL(command)
		}

		logger.Debug("send.request.to ", httpPath, " | method ", rr.Method)

		req, err := func() (*http.Request, error) {
			reqLock.Lock()
			defer reqLock.Unlock()

			if rr.Values == nil {
				if req, err = http.NewRequest(rr.Method, httpPath, nil); err != nil {
					return nil, err
				}
			} else {
				body := strings.NewReader(rr.Values.Encode())
				if req, err = http.NewRequest(rr.Method, httpPath, body); err != nil {
					return nil, err
				}

				req.Header.Set("Content-Type",
					"application/x-www-form-urlencoded; param=value")
			}
			return req, nil
		}()

		if err != nil {
			return nil, err
		}

		resp, err = c.httpClient.Do(req)
		defer func() {
			if resp != nil {
				resp.Body.Close()
			}
		}()

		// If the request was cancelled, return ErrRequestCancelled directly
		select {
		case <-cancelled:
			return nil, ErrRequestCancelled
		default:
		}

		numReqs++

		// network error, change a machine!
		if err != nil {
			logger.Debug("network error: ", err.Error())
			lastResp := http.Response{}
			if checkErr := checkRetry(c.cluster, numReqs, lastResp, err); checkErr != nil {
				return nil, checkErr
			}

			c.cluster.switchLeader(attempt % len(c.cluster.Machines))
			continue
		}

		// if there is no error, it should receive response
		logger.Debug("recv.response.from ", httpPath)

		if validHttpStatusCode[resp.StatusCode] {
			// try to read byte code and break the loop
			respBody, err = ioutil.ReadAll(resp.Body)
			if err == nil {
				logger.Debug("recv.success ", httpPath)
				break
			}
			// ReadAll error may be caused due to cancel request
			select {
			case <-cancelled:
				return nil, ErrRequestCancelled
			default:
			}

			if err == io.ErrUnexpectedEOF {
				// underlying connection was closed prematurely, probably by timeout
				// TODO: empty body or unexpectedEOF can cause http.Transport to get hosed;
				// this allows the client to detect that and take evasive action. Need
				// to revisit once code.google.com/p/go/issues/detail?id=8648 gets fixed.
				respBody = []byte{}
				break
			}
		}

		// if resp is TemporaryRedirect, set the new leader and retry
		if resp.StatusCode == http.StatusTemporaryRedirect {
			u, err := resp.Location()

			if err != nil {
				logger.Warning(err)
			} else {
				// Update cluster leader based on redirect location
				// because it should point to the leader address
				c.cluster.updateLeaderFromURL(u)
				logger.Debug("recv.response.relocate ", u.String())
			}
			resp.Body.Close()
			continue
		}

		if checkErr := checkRetry(c.cluster, numReqs, *resp,
			errors.New("Unexpected HTTP status code")); checkErr != nil {
			return nil, checkErr
		}
		resp.Body.Close()
	}

	r := &RawResponse{
		StatusCode: resp.StatusCode,
		Body:       respBody,
		Header:     resp.Header,
	}

	return r, nil
}

// DefaultCheckRetry defines the retrying behaviour for bad HTTP requests
// If we have retried 2 * machine number, stop retrying.
// If status code is InternalServerError, sleep for 200ms.
func DefaultCheckRetry(cluster *Cluster, numReqs int, lastResp http.Response,
	err error) error {

	if isEmptyResponse(lastResp) {
		// always retry if it failed to get response from one machine
		return err
	} else if !shouldRetry(lastResp) {
		body := []byte("nil")
		if lastResp.Body != nil {
			if b, err := ioutil.ReadAll(lastResp.Body); err == nil {
				body = b
			}
		}
		errStr := fmt.Sprintf("unhandled http status [%s] with body [%s]", http.StatusText(lastResp.StatusCode), body)
		return newError(ErrCodeUnhandledHTTPStatus, errStr, 0)
	}

	if numReqs > 2*len(cluster.Machines) {
		errStr := fmt.Sprintf("failed to propose on members %v twice [last error: %v]", cluster.Machines, err)
		return newError(ErrCodeEtcdNotReachable, errStr, 0)
	}
	if shouldRetry(lastResp) {
		// sleep some time and expect leader election finish
		time.Sleep(time.Millisecond * 200)
	}

	logger.Warning("bad response status code", lastResp.StatusCode)
	return nil
}

func isEmptyResponse(r http.Response) bool { return r.StatusCode == 0 }

// shouldRetry returns whether the reponse deserves retry.
func shouldRetry(r http.Response) bool {
	// TODO: only retry when the cluster is in leader election
	// We cannot do it exactly because etcd doesn't support it well.
	return r.StatusCode == http.StatusInternalServerError
}

func (c *Client) getHttpPath(random bool, s ...string) string {
	var machine string
	if random {
		machine = c.cluster.Machines[rand.Intn(len(c.cluster.Machines))]
	} else {
		machine = c.cluster.Leader
	}

	fullPath := machine + "/" + version
	for _, seg := range s {
		fullPath = fullPath + "/" + seg
	}

	return fullPath
}

// buildValues builds a url.Values map according to the given value and ttl
func buildValues(value string, ttl uint64) url.Values {
	v := url.Values{}

	if value != "" {
		v.Set("value", value)
	}

	if ttl > 0 {
		v.Set("ttl", fmt.Sprintf("%v", ttl))
	}

	return v
}

// convert key string to http path exclude version, including URL escaping
// for example: key[foo] -> path[keys/foo]
// key[/%z] -> path[keys/%25z]
// key[/] -> path[keys/]
func keyToPath(key string) string {
	// URL-escape our key, except for slashes
	p := strings.Replace(url.QueryEscape(path.Join("keys", key)), "%2F", "/", -1)

	// corner case: if key is "/" or "//" ect
	// path join will clear the tailing "/"
	// we need to add it back
	if p == "keys" {
		p = "keys/"
	}

	return p
}
<|endoftext|>