<commit_before>package main

import (
	"encoding/json"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/jeffwilliams/wwwmp3/play"
	"github.com/jeffwilliams/wwwmp3/scan"
)

// queryVal returns the first value for the GET query variable with name `key`.
func queryVal(r *http.Request, key string) (rs string) {
	v, ok := r.URL.Query()[key]
	if ok {
		rs = v[0]
	}
	return
}

// prependPrefix prepents the configured prefix to the MP3 path
func prependPrefix(fields map[string]string) {
	if v, ok := fields["path"]; ok {
		fields["path"] = prefix.apply(v)
	}
}

// Respond to requests for mp3 metadata: lists of artists, titles, song paths, etc.
func serveMeta(w http.ResponseWriter, r *http.Request) {
	timer := time.Now()

	// Parse a query value which is a comma separated list. If the value doesn't appear, return the empty list
	queryList := func(key string) (l []string) {
		if s := queryVal(r, key); len(s) > 0 {
			l = strings.Split(s, ",")
		}
		return
	}

	if r.Method == "GET" {
		// Query string Format:
		//  filter:
		//    artist=blah&album=blah&title=blah
		//  paging:
		//    page=0&pagesize=10
		//  fields (not present or the following:)
		//    fields=artist,album,...
		//  order (not present or the following:)
		//    order=artist,album,...

		// Calculate the filter
		filt := make(map[string]string)
		if v := queryVal(r, "artist"); len(v) > 0 {
			filt["artist"] = v
		}
		if v := queryVal(r, "album"); len(v) > 0 {
			filt["album"] = v
		}
		if v := queryVal(r, "title"); len(v) > 0 {
			filt["title"] = v
		}

		page, err := strconv.Atoi(queryVal(r, "page"))
		if err != nil {
			w.WriteHeader(400)
			w.Write([]byte("400 Bad Request: The 'page' parameter is missing or invalid."))
		}

		pageSize, err := strconv.Atoi(queryVal(r, "pagesize"))
		if err != nil {
			w.WriteHeader(400)
			w.Write([]byte("400 Bad Request: The 'pagesize' parameter is missing or invalid."))
		}

		fields := queryList("fields")

		order := queryList("order")

		ch := make(chan map[string]string)

		go scan.FindMp3sInDb(
			db,
			fields,
			filt,
			order,
			ch,
			&scan.Paging{PageSize: pageSize, Page: page},
			nil)

		enc := json.NewEncoder(w)
		w.Write([]byte("[\n"))
		j := 0
		for m := range ch {
			if j > 0 {
				w.Write([]byte(",\n"))
			}

			prependPrefix(m)

			enc.Encode(m)
			j++
		}
		w.Write([]byte("\n]\n"))
	}

	d := time.Now().Sub(timer)
	log.Info("serveMeta completed in %v", d)
}

// Perform functions on the mp3 player like play and pause.
func servePlayer(w http.ResponseWriter, r *http.Request) {
	logPrefix := "servePlayer: " + r.Method + " " + r.URL.Path + " - "
	log.Notice("%s requested", logPrefix)

	if r.Method == "GET" {
		if r.URL.Path == "/player/play" {
			log.Notice("%s play", logPrefix)
			err := player.Play()
			if err != nil {
				log.Error("%s player.Play() returned error: %v", logPrefix, err)
				w.WriteHeader(500)
				w.Write([]byte(err.Error()))
			}
		} else if r.URL.Path == "/player/pause" {
			player.Pause()
		} else if r.URL.Path == "/player/stop" {
			player.Stop()
		} else if r.URL.Path == "/player/volume" {
			v, err := play.GetVolume()
			if err != nil {
				log.Error("%s play.GetVolume() returned error: %v", logPrefix, err)
				w.WriteHeader(500)
				w.Write([]byte(err.Error()))
				return
			}
			w.Write([]byte("{\"volume\": "))
			w.Write([]byte(strconv.Itoa(int(v))))
			w.Write([]byte("}"))
			log.Notice("%s returning %d", logPrefix, int(v))
		}

	} else if r.Method == "POST" {
		decoder := (*json.Decoder)(nil)
		if r.Body != nil {
			decoder = json.NewDecoder(r.Body)
		} else {
			log.Error("%s posted with nil body.", logPrefix)
			w.WriteHeader(400)
			w.Write([]byte("400 Bad Request: no body"))
			return
		}

		decodeReq := func(i interface{}) bool {
			err := decoder.Decode(i)
			if err != nil {
				log.Error("%s decoding request failed: %v", logPrefix, err)
				w.WriteHeader(400)
				w.Write([]byte("400 Bad Request: invalid JSON"))
				return false
			}
			return true
		}

		if r.URL.Path == "/player/queue.enqueue" {
			req := struct {
				File string
			}{}

			if !decodeReq(&req) {
				return
			}

			queue.Enqueue(req.File)
		} else if r.URL.Path == "/player/queue.move" {
			req := struct {
				Indexes []int
				Delta   int
			}{}

			if !decodeReq(&req) {
				return
			}

			queue.Move(req.Indexes, req.Delta)
		} else if r.URL.Path == "/player/queue.move_to_top" {
			req := struct {
				Indexes []int
			}{}

			if !decodeReq(&req) {
				return
			}

			queue.MoveToTop(req.Indexes)
		} else if r.URL.Path == "/player/queue.remove" {
			req := struct {
				Indexes []int
			}{}

			if !decodeReq(&req) {
				return
			}

			queue.Remove(req.Indexes)
		} else if r.URL.Path == "/player/volume" {
			req := struct {
				Volume int
			}{}

			if !decodeReq(&req) {
				return
			}

			player.SetVolume(byte(req.Volume))
		} else if r.URL.Path == "/player/seek" {
			req := struct {
				Seek int
			}{}

			if !decodeReq(&req) {
				return
			}

			log.Notice("%s seek %v", logPrefix, req.Seek)
			player.Seek(req.Seek)
		} else if r.URL.Path == "/player/repeat_mode" {
			req := struct {
				Mode string
			}{}

			if !decodeReq(&req) {
				return
			}

			r, err := ParseRepeatMode(req.Mode)
			if err != nil {
				log.Warning("%s error parsing repeat mode: %v", logPrefix, err)
				w.WriteHeader(400)
				w.Write([]byte("400 Bad Request: Bad repeat mode"))
				return
			}

			repeatMode = r

			if repeatMode == RepeatOne {
				player.SetRepeat(true)
			} else {
				player.SetRepeat(false)
			}
			repeatModeTee.In <- struct{}{}
		}
	}
}

func enterScanning() bool {
	scanMutex.Lock()
	defer scanMutex.Unlock()

	if scanning {
		return false
	} else {
		scanning = true
		return true
	}
}

// Perform functions related to mp3 scanning
func serveScan(w http.ResponseWriter, r *http.Request) {
	logPrefix := "serveScan: " + r.Method + " " + r.URL.Path + " - "
	log.Notice("%s requested", logPrefix)

	if r.Method == "GET" {
		if r.URL.Path == "/scan/all" {
			if mp3Dirs != nil {
				scanMutex.Lock()
				defer scanMutex.Unlock()
				if !scanning {
					scanning = true
					go scanDirs(mp3Dirs)
				}
			} else {
				log.Notice("serveScan: scan requested but no directories to scan")
				w.WriteHeader(404)
				w.Write([]byte("404 Not Found: no directories to scan"))
			}
		}
	} else if r.Method == "POST" {
		decoder := (*json.Decoder)(nil)
		if r.Body != nil {
			decoder = json.NewDecoder(r.Body)
		} else {
			log.Error("%s posted with nil body.", logPrefix)
			w.WriteHeader(400)
			w.Write([]byte("400 Bad Request: no body"))
			return
		}

		decodeReq := func(i interface{}) bool {
			err := decoder.Decode(i)
			if err != nil {
				log.Error("%s decoding request failed: %v", logPrefix, err)
				w.WriteHeader(400)
				w.Write([]byte("400 Bad Request: invalid JSON"))
				return false
			}
			return true
		}

		if r.URL.Path == "/scan/path" {
			req := struct {
				Path string
			}{}

			if !decodeReq(&req) {
				return
			}

			if enterScanning() {
				go scanDirs([]string{req.Path})
			}
		}
	}
}

/*
serveWebsock implements a websocket connection with the browser, over which events
generated by the player are sent.

The messages sent to the browser are JSON and have the form:

{
  Offset: 0,
  Size: 0,
  State: 0,
  Volume: 0,
  Meta: {
    artist: "s",
    album: "s",
    title: "s",
    path: "s"
    }
}

State is an enumeration with the values:
  0 = empty
  1 = playing
  2 = paused

Meta may be null if there is no loaded mp3.
*/
func serveWebsock(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", 405)
		return
	}

	userAgent := r.Header.Get("User-Agent")

	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Error("Error upgrading connection to websock: %v", err)
		return
	}

	log.Notice("Websocket connection from %s %v", ws.RemoteAddr(), userAgent)

	defer log.Notice("Websocket handler for %s exiting", ws.RemoteAddr())

	// Send the full status to the browser
	d, err := jsonFullStatus(player.GetStatus(), meta, listQueue(queue), pathsToMetadatas(recent.Slice()), repeatMode)
	if err != nil {
		log.Error("Error encoding Player event as JSON: %v", err)
		return
	}
	log.Notice("Writing complete status to websocket %s", ws.RemoteAddr())
	err = websockWrite(ws, d)
	if err != nil {
		log.Error("Error writing to websocket %s: %v. Handler for socket is terminating.", ws.RemoteAddr(), err)
		// Client is probably gone. Close our channel and exit.
		return
	}

	// Add a new channel to the tee that we will use to detect changes
	c := make(chan interface{})
	eventTee.Add(c)
	defer eventTee.Del(c)

	scanEvents := make(chan interface{})
	scanTee.Add(scanEvents)
	defer scanTee.Del(scanEvents)

	repeatModeChanged := make(chan interface{})
	repeatModeTee.Add(repeatModeChanged)
	defer repeatModeTee.Del(repeatModeChanged)

	//ws.SetReadDeadline(time.Now().Add(10 * time.Millisecond))

loop:
	for {
		select {
		case e, ok := <-c:
			if !ok {
				// Player died...
				log.Error("Player closed event channel. Closing websocket %v", ws.RemoteAddr())
				ws.Close()
				break loop
			}

			if !websockHandlePlayerEvent(ws, e.(play.Event)) {
				break loop
			}

		case e, ok := <-scanEvents:
			if !ok {
				// Something closed the scanner channel.
				log.Error("Scanner channel was closed. Closing websocket %v", ws.RemoteAddr())
				ws.Close()
				break loop
			}

			d, err = jsonScan(e.(*scan.Metadata))
			if err != nil {
				log.Error("Error encoding metadata as JSON: %v", err)
				// Hopefully the next one works...
				continue loop
			}
			err = websockWrite(ws, d)
			if err != nil {
				log.Error("Error writing to websocket %v: %v", ws.RemoteAddr(), err)
				// Client is probably gone. Close our channel and exit.
				break loop
			}
		case _ = <-repeatModeChanged:
			m, err := jsonRepeat(repeatMode)
			if err != nil {
				log.Error("Error encoding metadata as JSON: %v", err)
				// Hopefully the next one works...
				continue loop
			}
			err = websockWrite(ws, m)
			if err != nil {
				log.Error("Error writing to websocket %v: %v", ws.RemoteAddr(), err)
				// Client is probably gone. Close our channel and exit.
				break loop
			}
		}
	}
}
<commit_msg>Fix cleanup when websock peer disappears<commit_after>package main

import (
	"encoding/json"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/jeffwilliams/wwwmp3/play"
	"github.com/jeffwilliams/wwwmp3/scan"
)

// queryVal returns the first value for the GET query variable with name `key`.
func queryVal(r *http.Request, key string) (rs string) {
	v, ok := r.URL.Query()[key]
	if ok {
		rs = v[0]
	}
	return
}

// prependPrefix prepents the configured prefix to the MP3 path
func prependPrefix(fields map[string]string) {
	if v, ok := fields["path"]; ok {
		fields["path"] = prefix.apply(v)
	}
}

// Respond to requests for mp3 metadata: lists of artists, titles, song paths, etc.
func serveMeta(w http.ResponseWriter, r *http.Request) {
	timer := time.Now()

	// Parse a query value which is a comma separated list. If the value doesn't appear, return the empty list
	queryList := func(key string) (l []string) {
		if s := queryVal(r, key); len(s) > 0 {
			l = strings.Split(s, ",")
		}
		return
	}

	if r.Method == "GET" {
		// Query string Format:
		//  filter:
		//    artist=blah&album=blah&title=blah
		//  paging:
		//    page=0&pagesize=10
		//  fields (not present or the following:)
		//    fields=artist,album,...
		//  order (not present or the following:)
		//    order=artist,album,...

		// Calculate the filter
		filt := make(map[string]string)
		if v := queryVal(r, "artist"); len(v) > 0 {
			filt["artist"] = v
		}
		if v := queryVal(r, "album"); len(v) > 0 {
			filt["album"] = v
		}
		if v := queryVal(r, "title"); len(v) > 0 {
			filt["title"] = v
		}

		page, err := strconv.Atoi(queryVal(r, "page"))
		if err != nil {
			w.WriteHeader(400)
			w.Write([]byte("400 Bad Request: The 'page' parameter is missing or invalid."))
		}

		pageSize, err := strconv.Atoi(queryVal(r, "pagesize"))
		if err != nil {
			w.WriteHeader(400)
			w.Write([]byte("400 Bad Request: The 'pagesize' parameter is missing or invalid."))
		}

		fields := queryList("fields")

		order := queryList("order")

		ch := make(chan map[string]string)

		go scan.FindMp3sInDb(
			db,
			fields,
			filt,
			order,
			ch,
			&scan.Paging{PageSize: pageSize, Page: page},
			nil)

		enc := json.NewEncoder(w)
		w.Write([]byte("[\n"))
		j := 0
		for m := range ch {
			if j > 0 {
				w.Write([]byte(",\n"))
			}

			prependPrefix(m)

			enc.Encode(m)
			j++
		}
		w.Write([]byte("\n]\n"))
	}

	d := time.Now().Sub(timer)
	log.Info("serveMeta completed in %v", d)
}

// Perform functions on the mp3 player like play and pause.
func servePlayer(w http.ResponseWriter, r *http.Request) {
	logPrefix := "servePlayer: " + r.Method + " " + r.URL.Path + " - "
	log.Notice("%s requested", logPrefix)

	if r.Method == "GET" {
		if r.URL.Path == "/player/play" {
			log.Notice("%s play", logPrefix)
			err := player.Play()
			if err != nil {
				log.Error("%s player.Play() returned error: %v", logPrefix, err)
				w.WriteHeader(500)
				w.Write([]byte(err.Error()))
			}
		} else if r.URL.Path == "/player/pause" {
			player.Pause()
		} else if r.URL.Path == "/player/stop" {
			player.Stop()
		} else if r.URL.Path == "/player/volume" {
			v, err := play.GetVolume()
			if err != nil {
				log.Error("%s play.GetVolume() returned error: %v", logPrefix, err)
				w.WriteHeader(500)
				w.Write([]byte(err.Error()))
				return
			}
			w.Write([]byte("{\"volume\": "))
			w.Write([]byte(strconv.Itoa(int(v))))
			w.Write([]byte("}"))
			log.Notice("%s returning %d", logPrefix, int(v))
		}

	} else if r.Method == "POST" {
		decoder := (*json.Decoder)(nil)
		if r.Body != nil {
			decoder = json.NewDecoder(r.Body)
		} else {
			log.Error("%s posted with nil body.", logPrefix)
			w.WriteHeader(400)
			w.Write([]byte("400 Bad Request: no body"))
			return
		}

		decodeReq := func(i interface{}) bool {
			err := decoder.Decode(i)
			if err != nil {
				log.Error("%s decoding request failed: %v", logPrefix, err)
				w.WriteHeader(400)
				w.Write([]byte("400 Bad Request: invalid JSON"))
				return false
			}
			return true
		}

		if r.URL.Path == "/player/queue.enqueue" {
			req := struct {
				File string
			}{}

			if !decodeReq(&req) {
				return
			}

			queue.Enqueue(req.File)
		} else if r.URL.Path == "/player/queue.move" {
			req := struct {
				Indexes []int
				Delta   int
			}{}

			if !decodeReq(&req) {
				return
			}

			queue.Move(req.Indexes, req.Delta)
		} else if r.URL.Path == "/player/queue.move_to_top" {
			req := struct {
				Indexes []int
			}{}

			if !decodeReq(&req) {
				return
			}

			queue.MoveToTop(req.Indexes)
		} else if r.URL.Path == "/player/queue.remove" {
			req := struct {
				Indexes []int
			}{}

			if !decodeReq(&req) {
				return
			}

			queue.Remove(req.Indexes)
		} else if r.URL.Path == "/player/volume" {
			req := struct {
				Volume int
			}{}

			if !decodeReq(&req) {
				return
			}

			player.SetVolume(byte(req.Volume))
		} else if r.URL.Path == "/player/seek" {
			req := struct {
				Seek int
			}{}

			if !decodeReq(&req) {
				return
			}

			log.Notice("%s seek %v", logPrefix, req.Seek)
			player.Seek(req.Seek)
		} else if r.URL.Path == "/player/repeat_mode" {
			req := struct {
				Mode string
			}{}

			if !decodeReq(&req) {
				return
			}

			r, err := ParseRepeatMode(req.Mode)
			if err != nil {
				log.Warning("%s error parsing repeat mode: %v", logPrefix, err)
				w.WriteHeader(400)
				w.Write([]byte("400 Bad Request: Bad repeat mode"))
				return
			}

			repeatMode = r

			if repeatMode == RepeatOne {
				player.SetRepeat(true)
			} else {
				player.SetRepeat(false)
			}
			repeatModeTee.In <- struct{}{}
		}
	}
}

func enterScanning() bool {
	scanMutex.Lock()
	defer scanMutex.Unlock()

	if scanning {
		return false
	} else {
		scanning = true
		return true
	}
}

// Perform functions related to mp3 scanning
func serveScan(w http.ResponseWriter, r *http.Request) {
	logPrefix := "serveScan: " + r.Method + " " + r.URL.Path + " - "
	log.Notice("%s requested", logPrefix)

	if r.Method == "GET" {
		if r.URL.Path == "/scan/all" {
			if mp3Dirs != nil {
				scanMutex.Lock()
				defer scanMutex.Unlock()
				if !scanning {
					scanning = true
					go scanDirs(mp3Dirs)
				}
			} else {
				log.Notice("serveScan: scan requested but no directories to scan")
				w.WriteHeader(404)
				w.Write([]byte("404 Not Found: no directories to scan"))
			}
		}
	} else if r.Method == "POST" {
		decoder := (*json.Decoder)(nil)
		if r.Body != nil {
			decoder = json.NewDecoder(r.Body)
		} else {
			log.Error("%s posted with nil body.", logPrefix)
			w.WriteHeader(400)
			w.Write([]byte("400 Bad Request: no body"))
			return
		}

		decodeReq := func(i interface{}) bool {
			err := decoder.Decode(i)
			if err != nil {
				log.Error("%s decoding request failed: %v", logPrefix, err)
				w.WriteHeader(400)
				w.Write([]byte("400 Bad Request: invalid JSON"))
				return false
			}
			return true
		}

		if r.URL.Path == "/scan/path" {
			req := struct {
				Path string
			}{}

			if !decodeReq(&req) {
				return
			}

			if enterScanning() {
				go scanDirs([]string{req.Path})
			}
		}
	}
}

/*
serveWebsock implements a websocket connection with the browser, over which events
generated by the player are sent.

The messages sent to the browser are JSON and have the form:

{
  Offset: 0,
  Size: 0,
  State: 0,
  Volume: 0,
  Meta: {
    artist: "s",
    album: "s",
    title: "s",
    path: "s"
    }
}

State is an enumeration with the values:
  0 = empty
  1 = playing
  2 = paused

Meta may be null if there is no loaded mp3.
*/
func serveWebsock(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", 405)
		return
	}

	userAgent := r.Header.Get("User-Agent")

	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Error("Error upgrading connection to websock: %v", err)
		return
	}

	log.Notice("Websocket connection from %s %v", ws.RemoteAddr(), userAgent)

	defer log.Notice("Websocket handler for %s exiting", ws.RemoteAddr())
	defer ws.Close()

	// Send the full status to the browser
	d, err := jsonFullStatus(player.GetStatus(), meta, listQueue(queue), pathsToMetadatas(recent.Slice()), repeatMode)
	if err != nil {
		log.Error("Websock %v: Error encoding Player event as JSON: %v", ws.RemoteAddr(), err)
		return
	}
	log.Notice("Websock %v: Writing complete status to websocket", ws.RemoteAddr())
	err = websockWrite(ws, d)
	if err != nil {
		log.Error("Websock %v: Error writing to websocket: %v. Handler for socket is terminating.", ws.RemoteAddr(), err)
		// Client is probably gone. Close our channel and exit.
		return
	}

	// Add a new channel to the tee that we will use to detect changes
	c := make(chan interface{})
	eventTee.Add(c)
	defer eventTee.Del(c)

	scanEvents := make(chan interface{})
	scanTee.Add(scanEvents)
	defer scanTee.Del(scanEvents)

	repeatModeChanged := make(chan interface{})
	repeatModeTee.Add(repeatModeChanged)
	defer repeatModeTee.Del(repeatModeChanged)

	//ws.SetReadDeadline(time.Now().Add(10 * time.Millisecond))

loop:
	for {
		select {
		case e, ok := <-c:
			if !ok {
				// Player died...
				log.Error("Websock %v: Player closed event channel.", ws.RemoteAddr())
				break loop
			}

			if !websockHandlePlayerEvent(ws, e.(play.Event)) {
				break loop
			}

		case e, ok := <-scanEvents:
			if !ok {
				// Something closed the scanner channel.
				log.Error("Websock %v: Scanner channel was closed.", ws.RemoteAddr())
				ws.Close()
				break loop
			}

			d, err = jsonScan(e.(*scan.Metadata))
			if err != nil {
				log.Error("Websock %v: Error encoding metadata as JSON: %v", ws.RemoteAddr(), err)
				// Hopefully the next one works...
				continue loop
			}
			err = websockWrite(ws, d)
			if err != nil {
				log.Error("Websock %v: Error writing to websocket: %v", ws.RemoteAddr(), err)
				// Client is probably gone. Close our channel and exit.
				break loop
			}
		case _ = <-repeatModeChanged:
			m, err := jsonRepeat(repeatMode)
			if err != nil {
				log.Error("Websock %v: Error encoding metadata as JSON: %v", ws.RemoteAddr(), err)
				// Hopefully the next one works...
				continue loop
			}
			err = websockWrite(ws, m)
			if err != nil {
				log.Error("Websock %v: Error writing to websocket: %v", ws.RemoteAddr(), err)
				// Client is probably gone. Close our channel and exit.
				break loop
			}
		}
	}

	eventTee.Del(c)
	scanTee.Del(scanEvents)
	repeatModeTee.Del(repeatModeChanged)
}
<|endoftext|>


<commit_before>package command

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"runtime"
	"testing"

	"github.com/hashicorp/packer-plugin-amazon/builder/ebs"
	packersdk "github.com/hashicorp/packer-plugin-sdk/packer"
	"github.com/hashicorp/packer/builder/file"
	"github.com/hashicorp/packer/builder/null"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/post-processor/manifest"
	shell_local_pp "github.com/hashicorp/packer/post-processor/shell-local"
	filep "github.com/hashicorp/packer/provisioner/file"
	"github.com/hashicorp/packer/provisioner/shell"
	shell_local "github.com/hashicorp/packer/provisioner/shell-local"
	"github.com/hashicorp/packer/version"
)

// HasExec reports whether the current system can start new processes
// using os.StartProcess or (more commonly) exec.Command.
func HasExec() bool {
	switch runtime.GOOS {
	case "js":
		return false
	case "darwin":
		if runtime.GOARCH == "arm64" {
			return false
		}
	}
	return true
}

// MustHaveExec checks that the current system can start new processes
// using os.StartProcess or (more commonly) exec.Command.
// If not, MustHaveExec calls t.Skip with an explanation.
func MustHaveExec(t testing.TB) {
	if !HasExec() {
		t.Skipf("skipping test: cannot exec subprocess on %s/%s", runtime.GOOS, runtime.GOARCH)
	}
}

func helperCommandContext(t *testing.T, ctx context.Context, s ...string) (cmd *exec.Cmd) {
	MustHaveExec(t)

	cs := []string{"-test.run=TestHelperProcess", "--"}
	cs = append(cs, s...)
	if ctx != nil {
		cmd = exec.CommandContext(ctx, os.Args[0], cs...)
	} else {
		cmd = exec.Command(os.Args[0], cs...)
	}
	cmd.Env = append(os.Environ(), "GO_WANT_HELPER_PROCESS=1")
	return cmd
}

func helperCommand(t *testing.T, s ...string) *exec.Cmd {
	return helperCommandContext(t, nil, s...)
}

// TestHelperProcess isn't a real test. It's used as a helper process
// for TestParameterRun.
func TestHelperProcess(*testing.T) {
	if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
		return
	}
	defer os.Exit(0)

	args := os.Args
	for len(args) > 0 {
		if args[0] == "--" {
			args = args[1:]
			break
		}
		args = args[1:]
	}
	if len(args) == 0 {
		fmt.Fprintf(os.Stderr, "No command\n")
		os.Exit(2)
	}

	cmd, args := args[0], args[1:]
	switch cmd {
	case "console":
		os.Exit((&ConsoleCommand{Meta: commandMeta()}).Run(args))
	case "fmt":
		os.Exit((&FormatCommand{Meta: commandMeta()}).Run(args))
	case "inspect":
		os.Exit((&InspectCommand{Meta: commandMeta()}).Run(args))
	case "build":
		os.Exit((&BuildCommand{Meta: commandMeta()}).Run(args))
	case "hcl2_upgrade":
		os.Exit((&HCL2UpgradeCommand{Meta: commandMeta()}).Run(args))
	default:
		fmt.Fprintf(os.Stderr, "Unknown command %q\n", cmd)
		os.Exit(2)
	}
}

func commandMeta() Meta {
	basicUi := &packersdk.BasicUi{
		Reader:      os.Stdin,
		Writer:      os.Stdout,
		ErrorWriter: os.Stdout,
	}

	CommandMeta := Meta{
		CoreConfig: &packer.CoreConfig{
			Components: getBareComponentFinder(),
			Version:    version.Version,
		},
		Ui: basicUi,
	}
	return CommandMeta
}

func getBareComponentFinder() packer.ComponentFinder {
	return packer.ComponentFinder{
		PluginConfig: &packer.PluginConfig{
			Builders: packer.MapOfBuilder{
				"file":       func() (packersdk.Builder, error) { return &file.Builder{}, nil },
				"null":       func() (packersdk.Builder, error) { return &null.Builder{}, nil },
				"amazon-ebs": func() (packersdk.Builder, error) { return &ebs.Builder{}, nil },
			},
			Provisioners: packer.MapOfProvisioner{
				"shell-local": func() (packersdk.Provisioner, error) { return &shell_local.Provisioner{}, nil },
				"shell":       func() (packersdk.Provisioner, error) { return &shell.Provisioner{}, nil },
				"file":        func() (packersdk.Provisioner, error) { return &filep.Provisioner{}, nil },
			},
			PostProcessors: packer.MapOfPostProcessor{
				"shell-local": func() (packersdk.PostProcessor, error) { return &shell_local_pp.PostProcessor{}, nil },
				"manifest":    func() (packersdk.PostProcessor, error) { return &manifest.PostProcessor{}, nil },
			},
		},
	}
}
<commit_msg>Remove a skip for darwin arm64 (m1/m2 macs) for exec tests. (#12088)<commit_after>package command

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"runtime"
	"testing"

	"github.com/hashicorp/packer-plugin-amazon/builder/ebs"
	packersdk "github.com/hashicorp/packer-plugin-sdk/packer"
	"github.com/hashicorp/packer/builder/file"
	"github.com/hashicorp/packer/builder/null"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/post-processor/manifest"
	shell_local_pp "github.com/hashicorp/packer/post-processor/shell-local"
	filep "github.com/hashicorp/packer/provisioner/file"
	"github.com/hashicorp/packer/provisioner/shell"
	shell_local "github.com/hashicorp/packer/provisioner/shell-local"
	"github.com/hashicorp/packer/version"
)

// HasExec reports whether the current system can start new processes
// using os.StartProcess or (more commonly) exec.Command.
func HasExec() bool {
	switch runtime.GOOS {
	case "js":
		return false
	}
	return true
}

// MustHaveExec checks that the current system can start new processes
// using os.StartProcess or (more commonly) exec.Command.
// If not, MustHaveExec calls t.Skip with an explanation.
func MustHaveExec(t testing.TB) {
	if !HasExec() {
		t.Skipf("skipping test: cannot exec subprocess on %s/%s", runtime.GOOS, runtime.GOARCH)
	}
}

func helperCommandContext(t *testing.T, ctx context.Context, s ...string) (cmd *exec.Cmd) {
	MustHaveExec(t)

	cs := []string{"-test.run=TestHelperProcess", "--"}
	cs = append(cs, s...)
	if ctx != nil {
		cmd = exec.CommandContext(ctx, os.Args[0], cs...)
	} else {
		cmd = exec.Command(os.Args[0], cs...)
	}
	cmd.Env = append(os.Environ(), "GO_WANT_HELPER_PROCESS=1")
	return cmd
}

func helperCommand(t *testing.T, s ...string) *exec.Cmd {
	return helperCommandContext(t, nil, s...)
}

// TestHelperProcess isn't a real test. It's used as a helper process
// for TestParameterRun.
func TestHelperProcess(*testing.T) {
	if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
		return
	}
	defer os.Exit(0)

	args := os.Args
	for len(args) > 0 {
		if args[0] == "--" {
			args = args[1:]
			break
		}
		args = args[1:]
	}
	if len(args) == 0 {
		fmt.Fprintf(os.Stderr, "No command\n")
		os.Exit(2)
	}

	cmd, args := args[0], args[1:]
	switch cmd {
	case "console":
		os.Exit((&ConsoleCommand{Meta: commandMeta()}).Run(args))
	case "fmt":
		os.Exit((&FormatCommand{Meta: commandMeta()}).Run(args))
	case "inspect":
		os.Exit((&InspectCommand{Meta: commandMeta()}).Run(args))
	case "build":
		os.Exit((&BuildCommand{Meta: commandMeta()}).Run(args))
	case "hcl2_upgrade":
		os.Exit((&HCL2UpgradeCommand{Meta: commandMeta()}).Run(args))
	default:
		fmt.Fprintf(os.Stderr, "Unknown command %q\n", cmd)
		os.Exit(2)
	}
}

func commandMeta() Meta {
	basicUi := &packersdk.BasicUi{
		Reader:      os.Stdin,
		Writer:      os.Stdout,
		ErrorWriter: os.Stdout,
	}

	CommandMeta := Meta{
		CoreConfig: &packer.CoreConfig{
			Components: getBareComponentFinder(),
			Version:    version.Version,
		},
		Ui: basicUi,
	}
	return CommandMeta
}

func getBareComponentFinder() packer.ComponentFinder {
	return packer.ComponentFinder{
		PluginConfig: &packer.PluginConfig{
			Builders: packer.MapOfBuilder{
				"file":       func() (packersdk.Builder, error) { return &file.Builder{}, nil },
				"null":       func() (packersdk.Builder, error) { return &null.Builder{}, nil },
				"amazon-ebs": func() (packersdk.Builder, error) { return &ebs.Builder{}, nil },
			},
			Provisioners: packer.MapOfProvisioner{
				"shell-local": func() (packersdk.Provisioner, error) { return &shell_local.Provisioner{}, nil },
				"shell":       func() (packersdk.Provisioner, error) { return &shell.Provisioner{}, nil },
				"file":        func() (packersdk.Provisioner, error) { return &filep.Provisioner{}, nil },
			},
			PostProcessors: packer.MapOfPostProcessor{
				"shell-local": func() (packersdk.PostProcessor, error) { return &shell_local_pp.PostProcessor{}, nil },
				"manifest":    func() (packersdk.PostProcessor, error) { return &manifest.PostProcessor{}, nil },
			},
		},
	}
}
<|endoftext|>


<commit_before>package command

import (
	"fmt"
	"github.com/hashicorp/consul/command/agent"
	"github.com/hashicorp/consul/consul"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"os"
	"sync/atomic"
	"testing"
	"time"
)

var offset uint64

func init() {
	// Seed the random number generator
	rand.Seed(time.Now().UnixNano())
}

type agentWrapper struct {
	dir      string
	config   *agent.Config
	agent    *agent.Agent
	rpc      *agent.AgentRPC
	http     *agent.HTTPServer
	addr     string
	httpAddr string
}

func (a *agentWrapper) Shutdown() {
	a.rpc.Shutdown()
	a.agent.Shutdown()
	a.http.Shutdown()
	os.RemoveAll(a.dir)
}

func testAgent(t *testing.T) *agentWrapper {
	return testAgentWithConfig(t, func(c *agent.Config) {})
}

func testAgentWithConfig(t *testing.T, cb func(c *agent.Config)) *agentWrapper {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	lw := agent.NewLogWriter(512)
	mult := io.MultiWriter(os.Stderr, lw)

	conf := nextConfig()
	cb(conf)

	dir, err := ioutil.TempDir("", "agent")
	if err != nil {
		t.Fatalf(fmt.Sprintf("err: %v", err))
	}
	conf.DataDir = dir

	a, err := agent.Create(conf, lw)
	if err != nil {
		os.RemoveAll(dir)
		t.Fatalf(fmt.Sprintf("err: %v", err))
	}

	rpc := agent.NewAgentRPC(a, l, mult, lw)

	conf.Addresses.HTTP = "127.0.0.1"
	httpAddr := fmt.Sprintf("127.0.0.1:%d", conf.Ports.HTTP)
	http, err := agent.NewHTTPServers(a, conf, os.Stderr)
	if err != nil {
		os.RemoveAll(dir)
		t.Fatalf(fmt.Sprintf("err: %v", err))
	}

	if http == nil || len(http) == 0 {
		os.RemoveAll(dir)
		t.Fatalf(fmt.Sprintf("Could not create HTTP server to listen on: %s", httpAddr))
	}

	return &agentWrapper{
		dir:      dir,
		config:   conf,
		agent:    a,
		rpc:      rpc,
		http:     http[0],
		addr:     l.Addr().String(),
		httpAddr: httpAddr,
	}
}

func nextConfig() *agent.Config {
	idx := int(atomic.AddUint64(&offset, 1))
	conf := agent.DefaultConfig()

	conf.Bootstrap = true
	conf.Datacenter = "dc1"
	conf.NodeName = fmt.Sprintf("Node %d", idx)
	conf.BindAddr = "127.0.0.1"
	conf.Server = true

	conf.Ports.HTTP = 10000 + 10*idx
	conf.Ports.HTTPS = 10401 + 10*idx
	conf.Ports.RPC = 10100 + 10*idx
	conf.Ports.SerfLan = 10201 + 10*idx
	conf.Ports.SerfWan = 10202 + 10*idx
	conf.Ports.Server = 10300 + 10*idx

	cons := consul.DefaultConfig()
	conf.ConsulConfig = cons

	cons.SerfLANConfig.MemberlistConfig.ProbeTimeout = 100 * time.Millisecond
	cons.SerfLANConfig.MemberlistConfig.ProbeInterval = 100 * time.Millisecond
	cons.SerfLANConfig.MemberlistConfig.GossipInterval = 100 * time.Millisecond

	cons.SerfWANConfig.MemberlistConfig.ProbeTimeout = 100 * time.Millisecond
	cons.SerfWANConfig.MemberlistConfig.ProbeInterval = 100 * time.Millisecond
	cons.SerfWANConfig.MemberlistConfig.GossipInterval = 100 * time.Millisecond

	cons.RaftConfig.LeaderLeaseTimeout = 20 * time.Millisecond
	cons.RaftConfig.HeartbeatTimeout = 40 * time.Millisecond
	cons.RaftConfig.ElectionTimeout = 40 * time.Millisecond

	return conf
}
<commit_msg>Add test function for returning an API client<commit_after>package command

import (
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"os"
	"sync/atomic"
	"testing"
	"time"

	"github.com/hashicorp/consul/api"
	"github.com/hashicorp/consul/command/agent"
	"github.com/hashicorp/consul/consul"
)

var offset uint64

func init() {
	// Seed the random number generator
	rand.Seed(time.Now().UnixNano())
}

type agentWrapper struct {
	dir      string
	config   *agent.Config
	agent    *agent.Agent
	rpc      *agent.AgentRPC
	http     *agent.HTTPServer
	addr     string
	httpAddr string
}

func (a *agentWrapper) Shutdown() {
	a.rpc.Shutdown()
	a.agent.Shutdown()
	a.http.Shutdown()
	os.RemoveAll(a.dir)
}

func testAgent(t *testing.T) *agentWrapper {
	return testAgentWithConfig(t, func(c *agent.Config) {})
}

func testAgentWithAPIClient(t *testing.T) (*agentWrapper, *api.Client) {
	agent := testAgentWithConfig(t, func(c *agent.Config) {})
	client, err := api.NewClient(&api.Config{Address: agent.httpAddr})
	if err != nil {
		t.Fatalf("consul client: %#v", err)
	}
	return agent, client
}

func testAgentWithConfig(t *testing.T, cb func(c *agent.Config)) *agentWrapper {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	lw := agent.NewLogWriter(512)
	mult := io.MultiWriter(os.Stderr, lw)

	conf := nextConfig()
	cb(conf)

	dir, err := ioutil.TempDir("", "agent")
	if err != nil {
		t.Fatalf(fmt.Sprintf("err: %v", err))
	}
	conf.DataDir = dir

	a, err := agent.Create(conf, lw)
	if err != nil {
		os.RemoveAll(dir)
		t.Fatalf(fmt.Sprintf("err: %v", err))
	}

	rpc := agent.NewAgentRPC(a, l, mult, lw)

	conf.Addresses.HTTP = "127.0.0.1"
	httpAddr := fmt.Sprintf("127.0.0.1:%d", conf.Ports.HTTP)
	http, err := agent.NewHTTPServers(a, conf, os.Stderr)
	if err != nil {
		os.RemoveAll(dir)
		t.Fatalf(fmt.Sprintf("err: %v", err))
	}

	if http == nil || len(http) == 0 {
		os.RemoveAll(dir)
		t.Fatalf(fmt.Sprintf("Could not create HTTP server to listen on: %s", httpAddr))
	}

	return &agentWrapper{
		dir:      dir,
		config:   conf,
		agent:    a,
		rpc:      rpc,
		http:     http[0],
		addr:     l.Addr().String(),
		httpAddr: httpAddr,
	}
}

func nextConfig() *agent.Config {
	idx := int(atomic.AddUint64(&offset, 1))
	conf := agent.DefaultConfig()

	conf.Bootstrap = true
	conf.Datacenter = "dc1"
	conf.NodeName = fmt.Sprintf("Node %d", idx)
	conf.BindAddr = "127.0.0.1"
	conf.Server = true

	conf.Ports.HTTP = 10000 + 10*idx
	conf.Ports.HTTPS = 10401 + 10*idx
	conf.Ports.RPC = 10100 + 10*idx
	conf.Ports.SerfLan = 10201 + 10*idx
	conf.Ports.SerfWan = 10202 + 10*idx
	conf.Ports.Server = 10300 + 10*idx

	cons := consul.DefaultConfig()
	conf.ConsulConfig = cons

	cons.SerfLANConfig.MemberlistConfig.ProbeTimeout = 100 * time.Millisecond
	cons.SerfLANConfig.MemberlistConfig.ProbeInterval = 100 * time.Millisecond
	cons.SerfLANConfig.MemberlistConfig.GossipInterval = 100 * time.Millisecond

	cons.SerfWANConfig.MemberlistConfig.ProbeTimeout = 100 * time.Millisecond
	cons.SerfWANConfig.MemberlistConfig.ProbeInterval = 100 * time.Millisecond
	cons.SerfWANConfig.MemberlistConfig.GossipInterval = 100 * time.Millisecond

	cons.RaftConfig.LeaderLeaseTimeout = 20 * time.Millisecond
	cons.RaftConfig.HeartbeatTimeout = 40 * time.Millisecond
	cons.RaftConfig.ElectionTimeout = 40 * time.Millisecond

	return conf
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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

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

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

package util

import (
	"bytes"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
)

// KubernetesKubectl implements the interface for talking to Kubernetes by wrapping calls
// via kubectl.
type KubernetesKubectl struct {
	KubePath string
	// Base level arguments to kubectl. User commands/arguments get appended to this.
	Arguments []string
}

func NewKubernetesKubectl(config *KubernetesConfig) Kubernetes {
	if config.KubePath == "" {
		log.Fatalf("kubectl path cannot be empty")
	}
	// If a configuration file is specified, then it will provide the server
	// address and credentials. If not, then we check for the server address
	// and credentials as individual flags.
	var args []string
	if config.KubeConfig != "" {
		config.KubeConfig = os.ExpandEnv(config.KubeConfig)
		args = append(args, fmt.Sprintf("--kubeconfig=%s", config.KubeConfig))
	} else {
		if config.KubeServer != "" {
			args = append(args, fmt.Sprintf("--server=https://%s", config.KubeServer))
		} else if config.KubeService != "" {
			addrs, err := net.LookupHost(config.KubeService)
			if err != nil || len(addrs) < 1 {
				log.Fatalf("cannot resolve DNS name: %v", config.KubeService)
			}

			args = append(args, fmt.Sprintf("--server=https://%s", addrs[0]))
		}

		if config.KubeInsecure {
			args = append(args, fmt.Sprintf("--insecure-skip-tls-verify=%s", config.KubeInsecure))
		} else {
			if config.KubeCertAuth != "" {
				args = append(args, fmt.Sprintf("--certificate-authority=%s", config.KubeCertAuth))
				if config.KubeClientCert == "" {
					args = append(args, fmt.Sprintf("--client-certificate=%s", config.KubeClientCert))
				}

				if config.KubeClientKey == "" {
					args = append(args, fmt.Sprintf("--client-key=%s", config.KubeClientKey))
				}
			}
		}
		if config.KubeToken != "" {
			args = append(args, fmt.Sprintf("--token=%s", config.KubeToken))
		} else {
			if config.KubeUsername != "" {
				args = append(args, fmt.Sprintf("--username=%s", config.KubeUsername))
			}

			if config.KubePassword != "" {
				args = append(args, fmt.Sprintf("--password=%s", config.KubePassword))
			}
		}
	}
	return &KubernetesKubectl{config.KubePath, args}
}

func (k *KubernetesKubectl) Get(name string, resourceType string) (string, error) {
	args := []string{"get"}
	// Specify output as json rather than human readable for easier machine parsing
	args = append(args, "-o", "json")
	args = append(args, resourceType)
	args = append(args, name)
	return k.execute(args, "")
}

func (k *KubernetesKubectl) Create(resource string) (string, error) {
	args := []string{"create"}
	//	args = append(args, resourceType)
	//	args = append(args, name)
	return k.execute(args, resource)
}

func (k *KubernetesKubectl) Delete(name string, resourceType string) (string, error) {
	args := []string{"delete"}
	args = append(args, resourceType)
	args = append(args, name)
	return k.execute(args, "")
}

func (k *KubernetesKubectl) execute(args []string, input string) (string, error) {
	if len(input) > 0 {
		args = append(args, "-f", "-")
	}

	// Tack on the common arguments to the end of the command line
	args = append(args, k.Arguments...)
	cmd := exec.Command(k.KubePath, args...)
	cmd.Stdin = bytes.NewBuffer([]byte(input))

	// Combine stdout and stderr into a single dynamically resized buffer
	combined := &bytes.Buffer{}
	cmd.Stdout = combined
	cmd.Stderr = combined

	if err := cmd.Start(); err != nil {
		log.Printf("cannot start kubectl %#v", err)
		return combined.String(), err
	}

	if err := cmd.Wait(); err != nil {
		log.Printf("kubectl failed:  %#v", err)
		return combined.String(), err
	}
	return combined.String(), nil
}
<commit_msg>missed two comments from the CR<commit_after>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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

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

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

package util

import (
	"bytes"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
)

// KubernetesKubectl implements the interface for talking to Kubernetes by wrapping calls
// via kubectl.
type KubernetesKubectl struct {
	KubePath string
	// Base level arguments to kubectl. User commands/arguments get appended to this.
	Arguments []string
}

func NewKubernetesKubectl(config *KubernetesConfig) Kubernetes {
	if config.KubePath == "" {
		log.Fatalf("kubectl path cannot be empty")
	}
	// If a configuration file is specified, then it will provide the server
	// address and credentials. If not, then we check for the server address
	// and credentials as individual flags.
	var args []string
	if config.KubeConfig != "" {
		config.KubeConfig = os.ExpandEnv(config.KubeConfig)
		args = append(args, fmt.Sprintf("--kubeconfig=%s", config.KubeConfig))
	} else {
		if config.KubeServer != "" {
			args = append(args, fmt.Sprintf("--server=https://%s", config.KubeServer))
		} else if config.KubeService != "" {
			addrs, err := net.LookupHost(config.KubeService)
			if err != nil || len(addrs) < 1 {
				log.Fatalf("cannot resolve DNS name: %v", config.KubeService)
			}

			args = append(args, fmt.Sprintf("--server=https://%s", addrs[0]))
		}

		if config.KubeInsecure {
			args = append(args, fmt.Sprintf("--insecure-skip-tls-verify=%s", config.KubeInsecure))
		} else {
			if config.KubeCertAuth != "" {
				args = append(args, fmt.Sprintf("--certificate-authority=%s", config.KubeCertAuth))
				if config.KubeClientCert == "" {
					args = append(args, fmt.Sprintf("--client-certificate=%s", config.KubeClientCert))
				}

				if config.KubeClientKey == "" {
					args = append(args, fmt.Sprintf("--client-key=%s", config.KubeClientKey))
				}
			}
		}
		if config.KubeToken != "" {
			args = append(args, fmt.Sprintf("--token=%s", config.KubeToken))
		} else {
			if config.KubeUsername != "" {
				args = append(args, fmt.Sprintf("--username=%s", config.KubeUsername))
			}

			if config.KubePassword != "" {
				args = append(args, fmt.Sprintf("--password=%s", config.KubePassword))
			}
		}
	}
	return &KubernetesKubectl{config.KubePath, args}
}

func (k *KubernetesKubectl) Get(name string, resourceType string) (string, error) {
	// Specify output as json rather than human readable for easier machine parsing
	args := []string{"get",
		"-o",
		"json",
		resourceType,
		name}
	return k.execute(args, "")
}

func (k *KubernetesKubectl) Create(resource string) (string, error) {
	args := []string{"create"}
	return k.execute(args, resource)
}

func (k *KubernetesKubectl) Delete(name string, resourceType string) (string, error) {
	args := []string{"delete",
		resourceType,
		name}
	return k.execute(args, "")
}

func (k *KubernetesKubectl) execute(args []string, input string) (string, error) {
	if len(input) > 0 {
		args = append(args, "-f", "-")
	}

	// Tack on the common arguments to the end of the command line
	args = append(args, k.Arguments...)
	cmd := exec.Command(k.KubePath, args...)
	cmd.Stdin = bytes.NewBuffer([]byte(input))

	// Combine stdout and stderr into a single dynamically resized buffer
	combined := &bytes.Buffer{}
	cmd.Stdout = combined
	cmd.Stderr = combined

	if err := cmd.Start(); err != nil {
		log.Printf("cannot start kubectl %#v", err)
		return combined.String(), err
	}

	if err := cmd.Wait(); err != nil {
		log.Printf("kubectl failed:  %#v", err)
		return combined.String(), err
	}
	return combined.String(), nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"

	"code.uber.internal/infra/kraken/build-index/tagclient"
	"code.uber.internal/infra/kraken/lib/dockerregistry/transfer"
	"code.uber.internal/infra/kraken/lib/healthcheck"
	"code.uber.internal/infra/kraken/lib/store"
	"code.uber.internal/infra/kraken/lib/upstream"
	"code.uber.internal/infra/kraken/metrics"
	"code.uber.internal/infra/kraken/nginx"
	"code.uber.internal/infra/kraken/origin/blobclient"
	"code.uber.internal/infra/kraken/proxy/registryoverride"
	"code.uber.internal/infra/kraken/utils/configutil"
	"code.uber.internal/infra/kraken/utils/flagutil"
	"code.uber.internal/infra/kraken/utils/log"
)

func main() {
	configFile := flag.String("config", "", "Configuration file that has to be loaded from one of UBER_CONFIG_DIR locations")
	cluster := flag.String("cluster", "", "cluster name (e.g. prod01-sjc1)")
	var ports flagutil.Ints
	flag.Var(&ports, "port", "ports to listen on (may specify multiple)")
	flag.Parse()

	var config Config
	if err := configutil.Load(*configFile, &config); err != nil {
		panic(err)
	}

	log.ConfigureLogger(config.ZapLogging)

	if len(ports) == 0 {
		log.Fatal("Must specify at least one -port")
	}

	stats, closer, err := metrics.New(config.Metrics, *cluster)
	if err != nil {
		log.Fatalf("Failed to init metrics: %s", err)
	}
	defer closer.Close()

	go metrics.EmitVersion(stats)

	cas, err := store.NewCAStore(config.CAStore, stats)
	if err != nil {
		log.Fatalf("Failed to create store: %s", err)
	}

	tls, err := config.TLS.BuildClient()
	if err != nil {
		log.Fatalf("Error building client tls config: %s", err)
	}

	origins, err := config.Origin.Build(upstream.WithHealthCheck(healthcheck.Default(tls)))
	if err != nil {
		log.Fatalf("Error building origin host list: %s", err)
	}

	r := blobclient.NewClientResolver(blobclient.NewProvider(blobclient.WithTLS(tls)), origins)
	originCluster := blobclient.NewClusterClient(r)

	buildIndexes, err := config.BuildIndex.Build(upstream.WithHealthCheck(healthcheck.Default(tls)))
	if err != nil {
		log.Fatalf("Error building build-index host list: %s", err)
	}

	tagClient := tagclient.NewClusterClient(buildIndexes, tls)

	transferer := transfer.NewReadWriteTransferer(tagClient, originCluster, cas)

	registry, err := config.Registry.Build(config.Registry.ReadWriteParameters(transferer, cas, stats))
	if err != nil {
		log.Fatalf("Error creating registry: %s", err)
	}
	go func() {
		log.Info("Starting registry...")
		log.Fatal(registry.ListenAndServe())
	}()

	ros := registryoverride.NewServer(config.RegistryOverride, tagClient)
	go func() {
		log.Fatal(ros.ListenAndServe())
	}()

	log.Info("Starting nginx...")
	log.Fatal(nginx.Run(config.Nginx, map[string]interface{}{
		"ports": ports,
		"registry_server": nginx.GetServer(
			config.Registry.Docker.HTTP.Net, config.Registry.Docker.HTTP.Addr),
		"registry_override_server": nginx.GetServer(
			config.RegistryOverride.Listener.Net, config.RegistryOverride.Listener.Addr),
	}))
}
<commit_msg>Add tls config to proxy<commit_after>package main

import (
	"flag"

	"code.uber.internal/infra/kraken/build-index/tagclient"
	"code.uber.internal/infra/kraken/lib/dockerregistry/transfer"
	"code.uber.internal/infra/kraken/lib/healthcheck"
	"code.uber.internal/infra/kraken/lib/store"
	"code.uber.internal/infra/kraken/lib/upstream"
	"code.uber.internal/infra/kraken/metrics"
	"code.uber.internal/infra/kraken/nginx"
	"code.uber.internal/infra/kraken/origin/blobclient"
	"code.uber.internal/infra/kraken/proxy/registryoverride"
	"code.uber.internal/infra/kraken/utils/configutil"
	"code.uber.internal/infra/kraken/utils/flagutil"
	"code.uber.internal/infra/kraken/utils/log"
)

func main() {
	configFile := flag.String("config", "", "Configuration file that has to be loaded from one of UBER_CONFIG_DIR locations")
	cluster := flag.String("cluster", "", "cluster name (e.g. prod01-sjc1)")
	var ports flagutil.Ints
	flag.Var(&ports, "port", "ports to listen on (may specify multiple)")
	flag.Parse()

	var config Config
	if err := configutil.Load(*configFile, &config); err != nil {
		panic(err)
	}

	log.ConfigureLogger(config.ZapLogging)

	if len(ports) == 0 {
		log.Fatal("Must specify at least one -port")
	}

	stats, closer, err := metrics.New(config.Metrics, *cluster)
	if err != nil {
		log.Fatalf("Failed to init metrics: %s", err)
	}
	defer closer.Close()

	go metrics.EmitVersion(stats)

	cas, err := store.NewCAStore(config.CAStore, stats)
	if err != nil {
		log.Fatalf("Failed to create store: %s", err)
	}

	tls, err := config.TLS.BuildClient()
	if err != nil {
		log.Fatalf("Error building client tls config: %s", err)
	}

	origins, err := config.Origin.Build(upstream.WithHealthCheck(healthcheck.Default(tls)))
	if err != nil {
		log.Fatalf("Error building origin host list: %s", err)
	}

	r := blobclient.NewClientResolver(blobclient.NewProvider(blobclient.WithTLS(tls)), origins)
	originCluster := blobclient.NewClusterClient(r)

	buildIndexes, err := config.BuildIndex.Build(upstream.WithHealthCheck(healthcheck.Default(tls)))
	if err != nil {
		log.Fatalf("Error building build-index host list: %s", err)
	}

	tagClient := tagclient.NewClusterClient(buildIndexes, tls)

	transferer := transfer.NewReadWriteTransferer(tagClient, originCluster, cas)

	registry, err := config.Registry.Build(config.Registry.ReadWriteParameters(transferer, cas, stats))
	if err != nil {
		log.Fatalf("Error creating registry: %s", err)
	}
	go func() {
		log.Info("Starting registry...")
		log.Fatal(registry.ListenAndServe())
	}()

	ros := registryoverride.NewServer(config.RegistryOverride, tagClient)
	go func() {
		log.Fatal(ros.ListenAndServe())
	}()

	log.Info("Starting nginx...")
	log.Fatal(nginx.Run(config.Nginx, map[string]interface{}{
		"ports": ports,
		"registry_server": nginx.GetServer(
			config.Registry.Docker.HTTP.Net, config.Registry.Docker.HTTP.Addr),
		"registry_override_server": nginx.GetServer(
			config.RegistryOverride.Listener.Net, config.RegistryOverride.Listener.Addr)},
		nginx.WithTLS(config.TLS)))
}
<|endoftext|>


<commit_before>package proxy

import (
	"bufio"
	"context"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/getlantern/errors"
	"github.com/getlantern/netx"
	"github.com/getlantern/preconn"
	"github.com/getlantern/proxy/filters"
)

func (opts *Opts) applyHTTPDefaults() {
	// Apply defaults
	if opts.Filter == nil {
		opts.Filter = filters.FilterFunc(defaultFilter)
	}
	if opts.OnError == nil {
		opts.OnError = defaultOnError
	}
}

// Handle implements the interface Proxy
func (proxy *proxy) Handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn) (err error) {
	defer func() {
		p := recover()
		if p != nil {
			safeClose(downstream)
			err = errors.New("Recovered from panic handling connection: %v", p)
		}
	}()

	err = proxy.handle(ctx, downstreamIn, downstream, nil)
	return
}

func safeClose(conn net.Conn) {
	defer func() {
		p := recover()
		if p != nil {
			log.Errorf("Panic on closing connection: %v", p)
		}
	}()

	conn.Close()
}

func (proxy *proxy) logInitialReadError(downstream net.Conn, err error) error {
	rem := downstream.RemoteAddr()
	r := ""
	if rem != nil {
		r = rem.String()
	}
	txt := err.Error()
	// Ignore our generated error that should have already been reported.
	if strings.HasPrefix(txt, "Client Hello has no cipher suites") {
		log.Debugf("No cipher suites in common -- old Lantern client")
		return err
	}
	// These errors should all typically be internal go errors, typically with TLS. Break them up
	// for stackdriver grouping.
	if strings.Contains(txt, "oversized") {
		return log.Errorf("Oversized record on initial read: %v from %v", err, r)
	}
	if strings.Contains(txt, "first record does not") {
		return log.Errorf("Not a TLS client connection: %v from %v", err, r)
	}
	return log.Errorf("Initial ReadRequest: %v from %v", err, r)
}

func (proxy *proxy) handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn, upstream net.Conn) error {
	defer func() {
		if closeErr := downstream.Close(); closeErr != nil {
			log.Tracef("Error closing downstream connection: %s", closeErr)
		}
	}()

	downstreamBuffered := bufio.NewReader(downstreamIn)
	fctx := filters.WrapContext(withAwareConn(ctx), downstream)

	// Read initial request
	req, err := http.ReadRequest(downstreamBuffered)
	if req != nil {
		remoteAddr := downstream.RemoteAddr()
		if remoteAddr != nil {
			req.RemoteAddr = remoteAddr.String()
		}
		if origURLScheme(ctx) == "" {
			fctx = fctx.
				WithValue(ctxKeyOrigURLScheme, req.URL.Scheme).
				WithValue(ctxKeyOrigURLHost, req.URL.Host).
				WithValue(ctxKeyOrigHost, req.Host)
		}
	}

	if err != nil {
		if isUnexpected(err) {
			errResp := proxy.OnError(fctx, req, true, err)
			if errResp != nil {
				proxy.writeResponse(downstream, req, errResp)
			}

			return proxy.logInitialReadError(downstream, err)
		}
		return nil
	}

	var next filters.Next
	if req.Method == http.MethodConnect {
		next = proxy.nextCONNECT(downstream)
	} else {
		var tr idleClosingTransport
		if upstream != nil {
			setUpstreamForAwareConn(fctx, upstream)
			tr = &addressLoggingTransport{
				Transport: &http.Transport{
					DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
						// always use the supplied upstream connection, but don't allow it to
						// be closed by the transport
						return &noCloseConn{upstream}, nil
					},
					// this transport is only used once, don't keep any idle connections,
					// however still allow the transport to close the connection after using
					// it
					MaxIdleConnsPerHost: -1,
				},
				upstream: upstream,
			}
		} else {
			tr = &http.Transport{
				DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
					conn, err := proxy.Dial(ctx, false, net, addr)
					if err == nil {
						// On first dialing conn, handle RequestAware
						setUpstreamForAwareConn(ctx, conn)
						handleRequestAware(ctx)
					}
					return conn, err
				},
				IdleConnTimeout: proxy.IdleTimeout,
				// since we have one transport per downstream connection, we don't need
				// more than this
				MaxIdleConnsPerHost: 1,
			}
		}

		defer tr.CloseIdleConnections()
		next = func(ctx filters.Context, modifiedReq *http.Request) (*http.Response, filters.Context, error) {
			modifiedReq = modifiedReq.WithContext(ctx)
			setRequestForAwareConn(ctx, modifiedReq)
			handleRequestAware(ctx)
			resp, err := tr.RoundTrip(prepareRequest(modifiedReq))
			handleResponseAware(ctx, modifiedReq, resp, err)
			if err != nil {
				err = errors.New("Unable to round-trip http request to upstream: %v", err)
			}
			return resp, ctx, err
		}
	}

	return proxy.processRequests(fctx, req.RemoteAddr, req, downstream, downstreamBuffered, next)
}

func (proxy *proxy) processRequests(ctx filters.Context, remoteAddr string, req *http.Request, downstream net.Conn, downstreamBuffered *bufio.Reader, next filters.Next) error {
	var readErr error
	var resp *http.Response
	var err error

	for {
		if req.URL.Scheme == "" {
			req.URL.Scheme = origURLScheme(ctx)
		}
		if req.URL.Host == "" {
			req.URL.Host = origURLHost(ctx)
		}
		if req.Host == "" {
			req.Host = origHost(ctx)
		}
		resp, ctx, err = proxy.Filter.Apply(ctx, req, next)
		if err != nil && resp == nil {
			resp = proxy.OnError(ctx, req, false, err)
			if resp != nil {
				log.Debugf("Closing client connection on error: %v", err)
				// On error, we will always close the connection
				resp.Close = true
			}
		}

		if resp != nil {
			writeErr := proxy.writeResponse(downstream, req, resp)
			if writeErr != nil {
				if isUnexpected(writeErr) {
					return log.Errorf("Unable to write response to downstream: %v", writeErr)
				}
				// Error is not unexpected, but we're done
				return err
			}
		}

		if err != nil {
			// We encountered an error on round-tripping, stop now
			return err
		}

		upstream := upstreamConn(ctx)
		upstreamAddr := upstreamAddr(ctx)
		isConnect := upstream != nil || upstreamAddr != ""

		buffered := downstreamBuffered.Buffered()
		if buffered > 0 {
			b, _ := downstreamBuffered.Peek(buffered)
			downstream = preconn.Wrap(downstream, b)
		}

		if isConnect {
			return proxy.proceedWithConnect(ctx, req, upstreamAddr, upstream, downstream)
		}

		if req.Close {
			// Client signaled that they would close the connection after this
			// request, finish
			return err
		}

		if err == nil && resp != nil && resp.Close {
			// Last response, finish
			return err
		}

		// read the next request
		req, readErr = http.ReadRequest(downstreamBuffered)
		if readErr != nil {
			if isUnexpected(readErr) {
				errResp := proxy.OnError(ctx, req, true, readErr)
				if errResp != nil {
					proxy.writeResponse(downstream, req, errResp)
				}
				return log.Errorf("Unable to read next request from downstream: %v", readErr)
			}
			return err
		}

		// Preserve remote address from original request
		ctx = ctx.IncrementRequestNumber()
		req.RemoteAddr = remoteAddr
		req = req.WithContext(ctx)
	}
}

func handleRequestAware(ctx context.Context) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case RequestAware:
			req := requestForAwareConn(ctx)
			t.OnRequest(req)
		}
		return true
	})
}

func handleResponseAware(ctx context.Context, req *http.Request, resp *http.Response, err error) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case ResponseAware:
			t.OnResponse(req, resp, err)
		}
		return true
	})
}

func (proxy *proxy) writeResponse(downstream io.Writer, req *http.Request, resp *http.Response) error {
	if resp.Request == nil {
		resp.Request = req
	}
	out := downstream
	if resp.ProtoMajor == 0 {
		resp.ProtoMajor = 1
		resp.ProtoMinor = 1
	}
	belowHTTP11 := !resp.ProtoAtLeast(1, 1)
	if belowHTTP11 && resp.StatusCode < 200 {
		// HTTP 1.0 doesn't define status codes below 200, discard response
		// see http://coad.measurement-factory.com/cgi-bin/coad/SpecCgi?spec_id=rfc2616#excerpt/rfc2616/859a092cb26bde76c25284196171c94d
		out = ioutil.Discard
	} else {
		resp = prepareResponse(resp, belowHTTP11)
		proxy.addIdleKeepAlive(resp.Header)
	}
	err := resp.Write(out)
	// resp.Write closes the body only if it's successfully sent. Close
	// manually when error happens.
	if err != nil && resp.Body != nil {
		resp.Body.Close()
	}
	return err
}

// prepareRequest prepares the request in line with the HTTP spec for proxies.
func prepareRequest(req *http.Request) *http.Request {
	req.Proto = "HTTP/1.1"
	req.ProtoMajor = 1
	req.ProtoMinor = 1
	// Overwrite close flag: keep persistent connection for the backend servers
	req.Close = false

	// Request Header
	newHeader := make(http.Header)
	copyHeadersForForwarding(newHeader, req.Header)
	// Ensure we have a HOST header (important for Go 1.6+ because http.Server
	// strips the HOST header from the inbound request)
	newHeader.Set("Host", req.Host)
	req.Header = newHeader

	// Request URL
	req.URL = cloneURL(req.URL)
	// If req.URL.Scheme was blank, it's http. Otherwise, it's https and we leave
	// it alone.
	if req.URL.Scheme == "" {
		req.URL.Scheme = "http"
	}
	// We need to make sure the host is defined in the URL (not the actual URI)
	req.URL.Host = req.Host

	userAgent := req.UserAgent()
	if userAgent == "" {
		req.Header.Del("User-Agent")
	} else {
		req.Header.Set("User-Agent", userAgent)
	}

	return req
}

// prepareResponse prepares the response in line with the HTTP spec
func prepareResponse(resp *http.Response, belowHTTP11 bool) *http.Response {
	origHeader := resp.Header
	resp.Header = make(http.Header)
	copyHeadersForForwarding(resp.Header, origHeader)
	// Below added due to CoAdvisor test failure
	if resp.Header.Get("Date") == "" {
		resp.Header.Set("Date", time.Now().Format(time.RFC850))
	}
	if belowHTTP11 {
		// Also, make sure we're not sending chunked transfer encoding to 1.0 clients
		resp.TransferEncoding = nil
	}
	return resp
}

// cloneURL provides update safe copy by avoiding shallow copying User field
func cloneURL(i *url.URL) *url.URL {
	out := *i
	if i.User != nil {
		out.User = &(*i.User)
	}
	return &out
}

// copyHeadersForForwarding will copy the headers but filter those that shouldn't be
// forwarded
func copyHeadersForForwarding(dst, src http.Header) {
	var extraHopByHopHeaders []string
	for k, vv := range src {
		switch k {
		// Skip hop-by-hop headers, ref section 13.5.1 of http://www.ietf.org/rfc/rfc2616.txt
		case "Connection":
			// section 14.10 of rfc2616
			// the slice is short typically, don't bother sort it to speed up lookup
			extraHopByHopHeaders = vv
		case "Keep-Alive":
		case "Proxy-Connection": // It's non-standard but every browser adds it when sending request to proxy.
		case "Proxy-Authenticate":
		case "Proxy-Authorization":
		case "TE":
		case "Trailers":
		case "Transfer-Encoding":
		case "Upgrade":
		default:
			if !contains(k, extraHopByHopHeaders) {
				for _, v := range vv {
					dst.Add(k, v)
				}
			}
		}
	}
}

func contains(k string, s []string) bool {
	for _, h := range s {
		if k == h {
			return true
		}
	}
	return false
}

func isUnexpected(err error) bool {
	if err == nil {
		return false
	}
	if err == io.EOF {
		return false
	}
	// This is okay per the HTTP spec.
	// See https://www.w3.org/Protocols/rfc2616/rfc2616-sec8.html#sec8.1.4
	if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
		return false
	}

	text := err.Error()
	return !strings.HasSuffix(text, "EOF") &&
		!strings.Contains(text, "i/o timeout") &&
		!strings.Contains(text, "Use of idled network connection") &&
		!strings.Contains(text, "use of closed network connection") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "broken pipe") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "connection reset by peer")
}

func defaultFilter(ctx filters.Context, req *http.Request, next filters.Next) (*http.Response, filters.Context, error) {
	return next(ctx, req)
}

func defaultOnError(ctx filters.Context, req *http.Request, read bool, err error) *http.Response {
	return nil
}

type idleClosingTransport interface {
	RoundTrip(req *http.Request) (*http.Response, error)
	CloseIdleConnections()
}

type addressLoggingTransport struct {
	*http.Transport
	upstream net.Conn
}

func (alt *addressLoggingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	resp, err := alt.Transport.RoundTrip(req)
	if err != nil {
		err = errors.New("Error round-tripping to %v: %v", alt.upstream.RemoteAddr(), err)
	}
	return resp, err
}
<commit_msg>do not uniformly strip Proxy-Connection<commit_after>package proxy

import (
	"bufio"
	"context"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/getlantern/errors"
	"github.com/getlantern/netx"
	"github.com/getlantern/preconn"
	"github.com/getlantern/proxy/filters"
)

func (opts *Opts) applyHTTPDefaults() {
	// Apply defaults
	if opts.Filter == nil {
		opts.Filter = filters.FilterFunc(defaultFilter)
	}
	if opts.OnError == nil {
		opts.OnError = defaultOnError
	}
}

// Handle implements the interface Proxy
func (proxy *proxy) Handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn) (err error) {
	defer func() {
		p := recover()
		if p != nil {
			safeClose(downstream)
			err = errors.New("Recovered from panic handling connection: %v", p)
		}
	}()

	err = proxy.handle(ctx, downstreamIn, downstream, nil)
	return
}

func safeClose(conn net.Conn) {
	defer func() {
		p := recover()
		if p != nil {
			log.Errorf("Panic on closing connection: %v", p)
		}
	}()

	conn.Close()
}

func (proxy *proxy) logInitialReadError(downstream net.Conn, err error) error {
	rem := downstream.RemoteAddr()
	r := ""
	if rem != nil {
		r = rem.String()
	}
	txt := err.Error()
	// Ignore our generated error that should have already been reported.
	if strings.HasPrefix(txt, "Client Hello has no cipher suites") {
		log.Debugf("No cipher suites in common -- old Lantern client")
		return err
	}
	// These errors should all typically be internal go errors, typically with TLS. Break them up
	// for stackdriver grouping.
	if strings.Contains(txt, "oversized") {
		return log.Errorf("Oversized record on initial read: %v from %v", err, r)
	}
	if strings.Contains(txt, "first record does not") {
		return log.Errorf("Not a TLS client connection: %v from %v", err, r)
	}
	return log.Errorf("Initial ReadRequest: %v from %v", err, r)
}

func (proxy *proxy) handle(ctx context.Context, downstreamIn io.Reader, downstream net.Conn, upstream net.Conn) error {
	defer func() {
		if closeErr := downstream.Close(); closeErr != nil {
			log.Tracef("Error closing downstream connection: %s", closeErr)
		}
	}()

	downstreamBuffered := bufio.NewReader(downstreamIn)
	fctx := filters.WrapContext(withAwareConn(ctx), downstream)

	// Read initial request
	req, err := http.ReadRequest(downstreamBuffered)
	if req != nil {
		remoteAddr := downstream.RemoteAddr()
		if remoteAddr != nil {
			req.RemoteAddr = remoteAddr.String()
		}
		if origURLScheme(ctx) == "" {
			fctx = fctx.
				WithValue(ctxKeyOrigURLScheme, req.URL.Scheme).
				WithValue(ctxKeyOrigURLHost, req.URL.Host).
				WithValue(ctxKeyOrigHost, req.Host)
		}
	}

	if err != nil {
		if isUnexpected(err) {
			errResp := proxy.OnError(fctx, req, true, err)
			if errResp != nil {
				proxy.writeResponse(downstream, req, errResp)
			}

			return proxy.logInitialReadError(downstream, err)
		}
		return nil
	}

	var next filters.Next
	if req.Method == http.MethodConnect {
		next = proxy.nextCONNECT(downstream)
	} else {
		var tr idleClosingTransport
		if upstream != nil {
			setUpstreamForAwareConn(fctx, upstream)
			tr = &addressLoggingTransport{
				Transport: &http.Transport{
					DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
						// always use the supplied upstream connection, but don't allow it to
						// be closed by the transport
						return &noCloseConn{upstream}, nil
					},
					// this transport is only used once, don't keep any idle connections,
					// however still allow the transport to close the connection after using
					// it
					MaxIdleConnsPerHost: -1,
				},
				upstream: upstream,
			}
		} else {
			tr = &http.Transport{
				DialContext: func(ctx context.Context, net, addr string) (net.Conn, error) {
					conn, err := proxy.Dial(ctx, false, net, addr)
					if err == nil {
						// On first dialing conn, handle RequestAware
						setUpstreamForAwareConn(ctx, conn)
						handleRequestAware(ctx)
					}
					return conn, err
				},
				IdleConnTimeout: proxy.IdleTimeout,
				// since we have one transport per downstream connection, we don't need
				// more than this
				MaxIdleConnsPerHost: 1,
			}
		}

		defer tr.CloseIdleConnections()
		next = func(ctx filters.Context, modifiedReq *http.Request) (*http.Response, filters.Context, error) {
			modifiedReq = modifiedReq.WithContext(ctx)
			setRequestForAwareConn(ctx, modifiedReq)
			handleRequestAware(ctx)
			resp, err := tr.RoundTrip(prepareRequest(modifiedReq))
			handleResponseAware(ctx, modifiedReq, resp, err)
			if err != nil {
				err = errors.New("Unable to round-trip http request to upstream: %v", err)
			}
			return resp, ctx, err
		}
	}

	return proxy.processRequests(fctx, req.RemoteAddr, req, downstream, downstreamBuffered, next)
}

func (proxy *proxy) processRequests(ctx filters.Context, remoteAddr string, req *http.Request, downstream net.Conn, downstreamBuffered *bufio.Reader, next filters.Next) error {
	var readErr error
	var resp *http.Response
	var err error

	for {
		if req.URL.Scheme == "" {
			req.URL.Scheme = origURLScheme(ctx)
		}
		if req.URL.Host == "" {
			req.URL.Host = origURLHost(ctx)
		}
		if req.Host == "" {
			req.Host = origHost(ctx)
		}
		resp, ctx, err = proxy.Filter.Apply(ctx, req, next)
		if err != nil && resp == nil {
			resp = proxy.OnError(ctx, req, false, err)
			if resp != nil {
				log.Debugf("Closing client connection on error: %v", err)
				// On error, we will always close the connection
				resp.Close = true
			}
		}

		if resp != nil {
			writeErr := proxy.writeResponse(downstream, req, resp)
			if writeErr != nil {
				if isUnexpected(writeErr) {
					return log.Errorf("Unable to write response to downstream: %v", writeErr)
				}
				// Error is not unexpected, but we're done
				return err
			}
		}

		if err != nil {
			// We encountered an error on round-tripping, stop now
			return err
		}

		upstream := upstreamConn(ctx)
		upstreamAddr := upstreamAddr(ctx)
		isConnect := upstream != nil || upstreamAddr != ""

		buffered := downstreamBuffered.Buffered()
		if buffered > 0 {
			b, _ := downstreamBuffered.Peek(buffered)
			downstream = preconn.Wrap(downstream, b)
		}

		if isConnect {
			return proxy.proceedWithConnect(ctx, req, upstreamAddr, upstream, downstream)
		}

		if req.Close {
			// Client signaled that they would close the connection after this
			// request, finish
			return err
		}

		if err == nil && resp != nil && resp.Close {
			// Last response, finish
			return err
		}

		// read the next request
		req, readErr = http.ReadRequest(downstreamBuffered)
		if readErr != nil {
			if isUnexpected(readErr) {
				errResp := proxy.OnError(ctx, req, true, readErr)
				if errResp != nil {
					proxy.writeResponse(downstream, req, errResp)
				}
				return log.Errorf("Unable to read next request from downstream: %v", readErr)
			}
			return err
		}

		// Preserve remote address from original request
		ctx = ctx.IncrementRequestNumber()
		req.RemoteAddr = remoteAddr
		req = req.WithContext(ctx)
	}
}

func handleRequestAware(ctx context.Context) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case RequestAware:
			req := requestForAwareConn(ctx)
			t.OnRequest(req)
		}
		return true
	})
}

func handleResponseAware(ctx context.Context, req *http.Request, resp *http.Response, err error) {
	upstream := upstreamForAwareConn(ctx)
	if upstream == nil {
		return
	}

	netx.WalkWrapped(upstream, func(wrapped net.Conn) bool {
		switch t := wrapped.(type) {
		case ResponseAware:
			t.OnResponse(req, resp, err)
		}
		return true
	})
}

func (proxy *proxy) writeResponse(downstream io.Writer, req *http.Request, resp *http.Response) error {
	if resp.Request == nil {
		resp.Request = req
	}
	out := downstream
	if resp.ProtoMajor == 0 {
		resp.ProtoMajor = 1
		resp.ProtoMinor = 1
	}
	belowHTTP11 := !resp.ProtoAtLeast(1, 1)
	if belowHTTP11 && resp.StatusCode < 200 {
		// HTTP 1.0 doesn't define status codes below 200, discard response
		// see http://coad.measurement-factory.com/cgi-bin/coad/SpecCgi?spec_id=rfc2616#excerpt/rfc2616/859a092cb26bde76c25284196171c94d
		out = ioutil.Discard
	} else {
		resp = prepareResponse(resp, belowHTTP11)
		proxy.addIdleKeepAlive(resp.Header)
	}
	err := resp.Write(out)
	// resp.Write closes the body only if it's successfully sent. Close
	// manually when error happens.
	if err != nil && resp.Body != nil {
		resp.Body.Close()
	}
	return err
}

// prepareRequest prepares the request in line with the HTTP spec for proxies.
func prepareRequest(req *http.Request) *http.Request {
	req.Proto = "HTTP/1.1"
	req.ProtoMajor = 1
	req.ProtoMinor = 1
	// Overwrite close flag: keep persistent connection for the backend servers
	req.Close = false

	// Request Header
	newHeader := make(http.Header)
	copyHeadersForForwarding(newHeader, req.Header)
	// Ensure we have a HOST header (important for Go 1.6+ because http.Server
	// strips the HOST header from the inbound request)
	newHeader.Set("Host", req.Host)
	req.Header = newHeader

	// Request URL
	req.URL = cloneURL(req.URL)
	// If req.URL.Scheme was blank, it's http. Otherwise, it's https and we leave
	// it alone.
	if req.URL.Scheme == "" {
		req.URL.Scheme = "http"
	}
	// We need to make sure the host is defined in the URL (not the actual URI)
	req.URL.Host = req.Host

	userAgent := req.UserAgent()
	if userAgent == "" {
		req.Header.Del("User-Agent")
	} else {
		req.Header.Set("User-Agent", userAgent)
	}

	return req
}

// prepareResponse prepares the response in line with the HTTP spec
func prepareResponse(resp *http.Response, belowHTTP11 bool) *http.Response {
	origHeader := resp.Header
	resp.Header = make(http.Header)
	copyHeadersForForwarding(resp.Header, origHeader)
	// Below added due to CoAdvisor test failure
	if resp.Header.Get("Date") == "" {
		resp.Header.Set("Date", time.Now().Format(time.RFC850))
	}
	if belowHTTP11 {
		// Also, make sure we're not sending chunked transfer encoding to 1.0 clients
		resp.TransferEncoding = nil
	}
	return resp
}

// cloneURL provides update safe copy by avoiding shallow copying User field
func cloneURL(i *url.URL) *url.URL {
	out := *i
	if i.User != nil {
		out.User = &(*i.User)
	}
	return &out
}

// copyHeadersForForwarding will copy the headers but filter those that shouldn't be
// forwarded
func copyHeadersForForwarding(dst, src http.Header) {
	var extraHopByHopHeaders []string
	for k, vv := range src {
		switch k {
		// Skip hop-by-hop headers, ref section 13.5.1 of http://www.ietf.org/rfc/rfc2616.txt
		case "Connection":
			// section 14.10 of rfc2616
			// the slice is short typically, don't bother sort it to speed up lookup
			extraHopByHopHeaders = vv
		case "Keep-Alive":
		case "Proxy-Authenticate":
		case "Proxy-Authorization":
		case "TE":
		case "Trailers":
		case "Transfer-Encoding":
		case "Upgrade":
		default:
			if !contains(k, extraHopByHopHeaders) {
				for _, v := range vv {
					dst.Add(k, v)
				}
			}
		}
	}
}

func contains(k string, s []string) bool {
	for _, h := range s {
		if k == h {
			return true
		}
	}
	return false
}

func isUnexpected(err error) bool {
	if err == nil {
		return false
	}
	if err == io.EOF {
		return false
	}
	// This is okay per the HTTP spec.
	// See https://www.w3.org/Protocols/rfc2616/rfc2616-sec8.html#sec8.1.4
	if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
		return false
	}

	text := err.Error()
	return !strings.HasSuffix(text, "EOF") &&
		!strings.Contains(text, "i/o timeout") &&
		!strings.Contains(text, "Use of idled network connection") &&
		!strings.Contains(text, "use of closed network connection") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "broken pipe") &&
		// usually caused by client disconnecting
		!strings.Contains(text, "connection reset by peer")
}

func defaultFilter(ctx filters.Context, req *http.Request, next filters.Next) (*http.Response, filters.Context, error) {
	return next(ctx, req)
}

func defaultOnError(ctx filters.Context, req *http.Request, read bool, err error) *http.Response {
	return nil
}

type idleClosingTransport interface {
	RoundTrip(req *http.Request) (*http.Response, error)
	CloseIdleConnections()
}

type addressLoggingTransport struct {
	*http.Transport
	upstream net.Conn
}

func (alt *addressLoggingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	resp, err := alt.Transport.RoundTrip(req)
	if err != nil {
		err = errors.New("Error round-tripping to %v: %v", alt.upstream.RemoteAddr(), err)
	}
	return resp, err
}
<|endoftext|>


<commit_before>package pt

import (
	"math"
	"math/rand"
)

type Sampler interface {
	Sample(scene *Scene, ray Ray, rnd *rand.Rand) Color
}

/*
pick a random light
cast from camera N times, store hits
cast from light M times, store hits
cast shadow ray for all pairs
compute probability based on cos a, cos a?
combine
*/

func NewSampler(firstHitSamples, maxBounces int) Sampler {
	return &DefaultSampler{firstHitSamples, maxBounces}
}

type DefaultSampler struct {
	FirstHitSamples int
	MaxBounces      int
}

func (s *DefaultSampler) Sample(scene *Scene, ray Ray, rnd *rand.Rand) Color {
	return s.sample(scene, ray, true, s.FirstHitSamples, s.MaxBounces, rnd)
}

func (s *DefaultSampler) sample(scene *Scene, ray Ray, emission bool, samples, depth int, rnd *rand.Rand) Color {
	if depth < 0 {
		return Color{}
	}
	hit := scene.Intersect(ray)
	if !hit.Ok() {
		if scene.Texture != nil {
			d := ray.Direction
			u := math.Atan2(d.Z, d.X)
			v := math.Atan2(d.Y, Vector{d.X, 0, d.Z}.Length())
			u = (u + math.Pi) / (2 * math.Pi)
			v = (v + math.Pi/2) / math.Pi
			return scene.Texture.Sample(u, v)
		}
		return scene.Color
	}
	info := hit.Info(ray)
	result := Color{}
	emittance := info.Material.Emittance
	if emittance > 0 {
		if !emission {
			return Color{}
		}
		attenuation := info.Material.Attenuation.Compute(hit.T)
		result = result.Add(info.Color.MulScalar(emittance * attenuation * float64(samples)))
	}
	n := int(math.Sqrt(float64(samples)))
	for u := 0; u < n; u++ {
		for v := 0; v < n; v++ {
			fu := (float64(u) + rnd.Float64()) / float64(n)
			fv := (float64(v) + rnd.Float64()) / float64(n)
			newRay, reflected := ray.Bounce(&info, fu, fv, rnd)
			indirect := s.sample(scene, newRay, reflected, 1, depth-1, rnd)
			if reflected {
				tinted := indirect.Mix(info.Color.Mul(indirect), info.Material.Tint)
				result = result.Add(tinted)
			} else {
				direct := s.directLight(scene, info.Ray, rnd)
				result = result.Add(info.Color.Mul(direct.Add(indirect)))
			}
		}
	}
	return result.DivScalar(float64(n * n))
}

func (s *DefaultSampler) directLight(scene *Scene, n Ray, rnd *rand.Rand) Color {
	nLights := len(scene.Lights)
	if nLights == 0 {
		return Color{}
	}

	// pick a random light
	light := scene.Lights[rand.Intn(nLights)]

	// get bounding sphere center and radius
	var center Vector
	var radius float64
	switch t := light.(type) {
	case *Sphere:
		radius = t.Radius
		center = t.Center
	default:
		// get bounding sphere from bounding box
		box := t.BoundingBox()
		radius = box.OuterRadius()
		center = box.Center()
	}

	// get random point on sphere surface
	// TODO: use disk instead, this is biased?
	point := RandomUnitVector(rnd).MulScalar(radius).Add(center)

	// construct ray toward light point
	ray := Ray{n.Origin, point.Sub(n.Origin).Normalize()}

	// get cosine term
	diffuse := ray.Direction.Dot(n.Direction)
	if diffuse <= 0 {
		return Color{}
	}

	// check for light visibility
	hit := scene.Intersect(ray)
	if !hit.Ok() || hit.Shape != light {
		return Color{}
	}

	// get material properties from light
	material := light.MaterialAt(point)
	color := light.ColorAt(point)
	emittance := material.Emittance
	attenuation := material.Attenuation.Compute(hit.T)

	// compute solid angle (hemisphere coverage)
	hyp := center.Sub(n.Origin).Length()
	opp := radius
	// TODO: fix issue where hyp < opp (point inside sphere)
	theta := math.Asin(opp / hyp)
	adj := opp / math.Tan(theta)
	d := math.Cos(theta) * adj
	r := math.Sin(theta) * adj
	coverage := (r * r) / (d * d)

	m := diffuse * emittance * attenuation * coverage * float64(nLights)
	return color.MulScalar(m)
}
<commit_msg>DirectLighting bool flag<commit_after>package pt

import (
	"math"
	"math/rand"
)

type Sampler interface {
	Sample(scene *Scene, ray Ray, rnd *rand.Rand) Color
}

/*
pick a random light
cast from camera N times, store hits
cast from light M times, store hits
cast shadow ray for all pairs
compute probability based on cos a, cos a?
combine
*/

func NewSampler(firstHitSamples, maxBounces int) Sampler {
	return &DefaultSampler{firstHitSamples, maxBounces, true}
}

type DefaultSampler struct {
	FirstHitSamples int
	MaxBounces      int
	DirectLighting  bool
}

func (s *DefaultSampler) Sample(scene *Scene, ray Ray, rnd *rand.Rand) Color {
	return s.sample(scene, ray, true, s.FirstHitSamples, s.MaxBounces, rnd)
}

func (s *DefaultSampler) sample(scene *Scene, ray Ray, emission bool, samples, depth int, rnd *rand.Rand) Color {
	if depth < 0 {
		return Color{}
	}
	hit := scene.Intersect(ray)
	if !hit.Ok() {
		if scene.Texture != nil {
			d := ray.Direction
			u := math.Atan2(d.Z, d.X)
			v := math.Atan2(d.Y, Vector{d.X, 0, d.Z}.Length())
			u = (u + math.Pi) / (2 * math.Pi)
			v = (v + math.Pi/2) / math.Pi
			return scene.Texture.Sample(u, v)
		}
		return scene.Color
	}
	info := hit.Info(ray)
	result := Color{}
	emittance := info.Material.Emittance
	if emittance > 0 {
		if s.DirectLighting && !emission {
			return Color{}
		}
		attenuation := info.Material.Attenuation.Compute(hit.T)
		result = result.Add(info.Color.MulScalar(emittance * attenuation * float64(samples)))
	}
	n := int(math.Sqrt(float64(samples)))
	for u := 0; u < n; u++ {
		for v := 0; v < n; v++ {
			fu := (float64(u) + rnd.Float64()) / float64(n)
			fv := (float64(v) + rnd.Float64()) / float64(n)
			newRay, reflected := ray.Bounce(&info, fu, fv, rnd)
			indirect := s.sample(scene, newRay, reflected, 1, depth-1, rnd)
			if reflected {
				tinted := indirect.Mix(info.Color.Mul(indirect), info.Material.Tint)
				result = result.Add(tinted)
			} else {
				direct := Color{}
				if s.DirectLighting {
					direct = s.directLight(scene, info.Ray, rnd)
				}
				result = result.Add(info.Color.Mul(direct.Add(indirect)))
			}
		}
	}
	return result.DivScalar(float64(n * n))
}

func (s *DefaultSampler) directLight(scene *Scene, n Ray, rnd *rand.Rand) Color {
	nLights := len(scene.Lights)
	if nLights == 0 {
		return Color{}
	}

	// pick a random light
	light := scene.Lights[rand.Intn(nLights)]

	// get bounding sphere center and radius
	var center Vector
	var radius float64
	switch t := light.(type) {
	case *Sphere:
		radius = t.Radius
		center = t.Center
	default:
		// get bounding sphere from bounding box
		box := t.BoundingBox()
		radius = box.OuterRadius()
		center = box.Center()
	}

	// get random point on sphere surface
	// TODO: use disk instead, this is biased?
	point := RandomUnitVector(rnd).MulScalar(radius).Add(center)

	// construct ray toward light point
	ray := Ray{n.Origin, point.Sub(n.Origin).Normalize()}

	// get cosine term
	diffuse := ray.Direction.Dot(n.Direction)
	if diffuse <= 0 {
		return Color{}
	}

	// check for light visibility
	hit := scene.Intersect(ray)
	if !hit.Ok() || hit.Shape != light {
		return Color{}
	}

	// get material properties from light
	material := light.MaterialAt(point)
	color := light.ColorAt(point)
	emittance := material.Emittance
	attenuation := material.Attenuation.Compute(hit.T)

	// compute solid angle (hemisphere coverage)
	hyp := center.Sub(n.Origin).Length()
	opp := radius
	// TODO: fix issue where hyp < opp (point inside sphere)
	theta := math.Asin(opp / hyp)
	adj := opp / math.Tan(theta)
	d := math.Cos(theta) * adj
	r := math.Sin(theta) * adj
	coverage := (r * r) / (d * d)

	m := diffuse * emittance * attenuation * coverage * float64(nLights)
	return color.MulScalar(m)
}
<|endoftext|>


<commit_before>// +build off
package python

import (
	"fmt"
	"log"
	"regexp"
	"strings"

	"sourcegraph.com/sourcegraph/srclib/doc"
)

var rstDirectiveRegex = regexp.MustCompile(`\.\. [a-zA-Z]+:: `)

// Properly format python documentation by doing the following:
// Check if txt is rST. If so, compile into html, otherwise return html approximation to plaintext.
func formatDocs(txt string) string {
	txt = doc.StripNulls(txt)

	var format doc.Formatter
	if strings.Contains(txt, "~~~") || len(rstDirectiveRegex.FindStringIndex(txt)) > 0 {
		format = doc.ReStructuredText

		var err error
		txt, err = rstPreprocess(txt)
		if err != nil {
			log.Printf("rstPreprocess failed: %s", err)
		}
	} else if strings.Contains(txt, "====") {
		format = doc.Markdown
	} else {
		format = doc.Text
	}

	html, err := doc.ToHTML(format, []byte(txt))
	if err != nil {
		log.Printf("doc.ToHTML failed: %s", err)
	}
	return string(html)
}

var indentMatcher = regexp.MustCompile("(\\s*)[^\\s].*")
var rstRoleMatcher = regexp.MustCompile(":[A-Za-z]+:`([^`]+)`")

// Preprocess docstring before it can be fed to rst2html.  This is
// necessary, because sometimes indentations in docstrings are
// formatted to look good with code, rather than be valid rst.
func rstPreprocess(txt string) (string, error) {
	lines := strings.Split(txt, "\n")

	globalIndent := ""
	foundFirstIndent := false
	for l, line := range lines {
		if l == 0 && !strings.HasPrefix(line, "..") {
			continue
		} else if strings.TrimSpace(line) == "" {
			continue
		}

		matches := indentMatcher.FindStringSubmatch(line)
		if matches == nil || len(matches) < 2 {
			return "", fmt.Errorf("Could not process line: %s", line)
		}

		if !foundFirstIndent {
			foundFirstIndent = true
			globalIndent = matches[1]
		} else {
			if len(globalIndent) > len(matches[1]) {
				globalIndent = matches[1]
			}
		}
	}

	for l, line := range lines {
		if strings.HasPrefix(line, globalIndent) {
			lines[l] = line[len(globalIndent):]
		}
	}

	indentNormalizedTxt := strings.Join(lines, "\n")
	roleRemovedTxt := rstRoleMatcher.ReplaceAllString(indentNormalizedTxt, ":code:`$1`")

	return roleRemovedTxt, nil
}
<commit_msg>fix build flag<commit_after>// +build off

package python

import (
	"fmt"
	"log"
	"regexp"
	"strings"

	"sourcegraph.com/sourcegraph/srclib/doc"
)

var rstDirectiveRegex = regexp.MustCompile(`\.\. [a-zA-Z]+:: `)

// Properly format python documentation by doing the following:
// Check if txt is rST. If so, compile into html, otherwise return html approximation to plaintext.
func formatDocs(txt string) string {
	txt = doc.StripNulls(txt)

	var format doc.Formatter
	if strings.Contains(txt, "~~~") || len(rstDirectiveRegex.FindStringIndex(txt)) > 0 {
		format = doc.ReStructuredText

		var err error
		txt, err = rstPreprocess(txt)
		if err != nil {
			log.Printf("rstPreprocess failed: %s", err)
		}
	} else if strings.Contains(txt, "====") {
		format = doc.Markdown
	} else {
		format = doc.Text
	}

	html, err := doc.ToHTML(format, []byte(txt))
	if err != nil {
		log.Printf("doc.ToHTML failed: %s", err)
	}
	return string(html)
}

var indentMatcher = regexp.MustCompile("(\\s*)[^\\s].*")
var rstRoleMatcher = regexp.MustCompile(":[A-Za-z]+:`([^`]+)`")

// Preprocess docstring before it can be fed to rst2html.  This is
// necessary, because sometimes indentations in docstrings are
// formatted to look good with code, rather than be valid rst.
func rstPreprocess(txt string) (string, error) {
	lines := strings.Split(txt, "\n")

	globalIndent := ""
	foundFirstIndent := false
	for l, line := range lines {
		if l == 0 && !strings.HasPrefix(line, "..") {
			continue
		} else if strings.TrimSpace(line) == "" {
			continue
		}

		matches := indentMatcher.FindStringSubmatch(line)
		if matches == nil || len(matches) < 2 {
			return "", fmt.Errorf("Could not process line: %s", line)
		}

		if !foundFirstIndent {
			foundFirstIndent = true
			globalIndent = matches[1]
		} else {
			if len(globalIndent) > len(matches[1]) {
				globalIndent = matches[1]
			}
		}
	}

	for l, line := range lines {
		if strings.HasPrefix(line, globalIndent) {
			lines[l] = line[len(globalIndent):]
		}
	}

	indentNormalizedTxt := strings.Join(lines, "\n")
	roleRemovedTxt := rstRoleMatcher.ReplaceAllString(indentNormalizedTxt, ":code:`$1`")

	return roleRemovedTxt, nil
}
<|endoftext|>


<commit_before>package compressor

// Merger func is the joining together of some files so that they become one.
func Merger() {
	// TODO
}
<commit_msg>Update merger compressor<commit_after>package compressor

import "errors"

// Merger func is the joining together of some files so that they become one.
func Merger(dist string, srcs ...string) error {
	// TODO
	if len(srcs) <= 0 {
		return errors.New("Source files must require")
	}
	return nil
}
<|endoftext|>


<commit_before>package raft_test

import (
	"reflect"
	"testing"

	"github.com/relab/raft"
	gorums "github.com/relab/raft/cmd/raftgorums/gorumspb"
	pb "github.com/relab/raft/raftpb"
)

var requestVoteQFTests = []struct {
	name    string
	request *pb.RequestVoteRequest
	replies []*pb.RequestVoteResponse
	quorum  bool
	reply   *pb.RequestVoteResponse
}{
	{
		"do not grant vote",
		&pb.RequestVoteRequest{Term: 2},
		[]*pb.RequestVoteResponse{
			{Term: 2, VoteGranted: false},
		},
		false,
		nil,
	},
	{
		"grant vote",
		&pb.RequestVoteRequest{Term: 3},
		[]*pb.RequestVoteResponse{
			{Term: 3, VoteGranted: true},
		},
		true,
		&pb.RequestVoteResponse{Term: 3, VoteGranted: true},
	},
	{
		"reply with higher Term",
		&pb.RequestVoteRequest{Term: 3},
		[]*pb.RequestVoteResponse{
			{Term: 4, VoteGranted: false},
		},
		true,
		&pb.RequestVoteResponse{Term: 4, VoteGranted: false},
	},
}

type AETestCase struct {
	name    string
	request *pb.AppendEntriesRequest
	replies []*pb.AppendEntriesResponse
	quorum  bool
	reply   *pb.AppendEntriesResponse
}

var appendEntriesCommonQFTests = []AETestCase{
	{
		"reply with higher Term",
		&pb.AppendEntriesRequest{Term: 5},
		[]*pb.AppendEntriesResponse{
			{
				Term: 6,
			},
		},
		true,
		&pb.AppendEntriesResponse{
			Term: 6,
		},
	},
	{
		"one unsuccessful MatchIndex",
		&pb.AppendEntriesRequest{Term: 5},
		[]*pb.AppendEntriesResponse{
			{
				Term:       5,
				MatchIndex: 50,
				Success:    false,
			},
		},
		false,
		&pb.AppendEntriesResponse{
			Term:       5,
			MatchIndex: 50,
			Success:    false,
		},
	},
	{
		"two unsuccessful same MatchIndex",
		&pb.AppendEntriesRequest{Term: 5},
		[]*pb.AppendEntriesResponse{
			{
				Term:       5,
				MatchIndex: 100,
				Success:    false,
			},
			{
				Term:       5,
				MatchIndex: 100,
				Success:    false,
			},
		},
		true,
		&pb.AppendEntriesResponse{
			Term:       5,
			MatchIndex: 100,
			Success:    false,
		},
	},
	{
		"two unsuccessful different MatchIndex",
		&pb.AppendEntriesRequest{Term: 5},
		[]*pb.AppendEntriesResponse{
			{
				Term:       5,
				MatchIndex: 50,
				Success:    false,
			},
			{
				Term:       5,
				MatchIndex: 100,
				Success:    false,
			},
		},
		true,
		&pb.AppendEntriesResponse{
			Term:       5,
			MatchIndex: 50,
			Success:    false,
		},
	},
}

var appendEntriesFastQFTests = []AETestCase{
	// Stops after successful quorum, ignoring the last response.
	{
		"quorum successful",
		&pb.AppendEntriesRequest{Term: 5},
		[]*pb.AppendEntriesResponse{
			{
				Term:       5,
				MatchIndex: 100,
				Success:    true,
			},
		},
		true,
		&pb.AppendEntriesResponse{
			Term:       5,
			MatchIndex: 100,
			Success:    true,
		},
	},
	// Stops after successful quorum, ignoring the last response.
	// This tests shows how one follower is left behind, as from this point
	// it will never receive updates as long as the first follower is faster.
	{
		"one of two successful MatchIndex",
		&pb.AppendEntriesRequest{Term: 5},
		[]*pb.AppendEntriesResponse{
			{
				Term:       5,
				MatchIndex: 100,
				Success:    true,
			},
			{
				Term:       5,
				MatchIndex: 50,
				Success:    false,
			},
		},
		true,
		&pb.AppendEntriesResponse{
			Term:       5,
			MatchIndex: 100,
			Success:    true,
		},
	},
}

var qspecs = []struct {
	name string
	spec gorums.QuorumSpec
}{
	{
		"QuorumSpec N3 Q1",
		raft.NewQuorumSpec(3),
	},
}

func TestRequestVoteQF(t *testing.T) {
	for _, qs := range qspecs {
		for _, test := range requestVoteQFTests {
			t.Run(qs.name+"-"+test.name, func(t *testing.T) {
				reply, quorum := qs.spec.RequestVoteQF(test.request, test.replies)

				if quorum != test.quorum {
					t.Errorf("got %t, want %t", quorum, test.quorum)
				}

				if !reflect.DeepEqual(reply, test.reply) {
					t.Errorf("got %+v, want %+v", reply, test.reply)
				}
			})
		}
	}
}

func TestAppendEntriesFastQF(t *testing.T) {
	for _, qs := range qspecs {
		for _, test := range append(appendEntriesCommonQFTests, appendEntriesFastQFTests...) {
			t.Run(qs.name+"-"+test.name, func(t *testing.T) {
				reply, quorum := qs.spec.AppendEntriesQF(test.request, test.replies)

				if quorum != test.quorum {
					t.Errorf("got %t, want %t", quorum, test.quorum)
				}

				if !reflect.DeepEqual(reply, test.reply) {
					t.Errorf("got %+v, want %+v", reply, test.reply)
				}
			})
		}
	}
}
<commit_msg>qspec_test.go: Refactor tests<commit_after>package raft_test

import (
	"reflect"
	"testing"

	"github.com/relab/raft"
	gorums "github.com/relab/raft/cmd/raftgorums/gorumspb"
	pb "github.com/relab/raft/raftpb"
)

var requestVoteQFTests = []struct {
	name    string
	qs      spec
	request *pb.RequestVoteRequest
	replies []*pb.RequestVoteResponse
	quorum  bool
	reply   *pb.RequestVoteResponse
}{
	{
		"do not grant vote",
		n3q1,
		&pb.RequestVoteRequest{Term: 2},
		[]*pb.RequestVoteResponse{
			{Term: 2, VoteGranted: false},
		},
		false,
		nil,
	},
	{
		"grant vote",
		n3q1,
		&pb.RequestVoteRequest{Term: 3},
		[]*pb.RequestVoteResponse{
			{Term: 3, VoteGranted: true},
		},
		true,
		&pb.RequestVoteResponse{Term: 3, VoteGranted: true},
	},
	{
		"reply with higher Term",
		n3q1,
		&pb.RequestVoteRequest{Term: 3},
		[]*pb.RequestVoteResponse{
			{Term: 4, VoteGranted: false},
		},
		true,
		&pb.RequestVoteResponse{Term: 4, VoteGranted: false},
	},
}

var appendEntriesQFTests = []struct {
	name    string
	qs      spec
	request *pb.AppendEntriesRequest
	replies []*pb.AppendEntriesResponse
	quorum  bool
	reply   *pb.AppendEntriesResponse
}{
	{
		"reply with higher Term",
		n3q1,
		&pb.AppendEntriesRequest{Term: 5},
		[]*pb.AppendEntriesResponse{
			{
				Term: 6,
			},
		},
		true,
		&pb.AppendEntriesResponse{
			Term: 6,
		},
	},
	{
		"one unsuccessful MatchIndex",
		n3q1,
		&pb.AppendEntriesRequest{Term: 5},
		[]*pb.AppendEntriesResponse{
			{
				Term:       5,
				MatchIndex: 50,
				Success:    false,
			},
		},
		false,
		&pb.AppendEntriesResponse{
			Term:       5,
			MatchIndex: 50,
			Success:    false,
		},
	},
	{
		"two unsuccessful same MatchIndex",
		n3q1,
		&pb.AppendEntriesRequest{Term: 5},
		[]*pb.AppendEntriesResponse{
			{
				Term:       5,
				MatchIndex: 100,
				Success:    false,
			},
			{
				Term:       5,
				MatchIndex: 100,
				Success:    false,
			},
		},
		true,
		&pb.AppendEntriesResponse{
			Term:       5,
			MatchIndex: 100,
			Success:    false,
		},
	},
	{
		"two unsuccessful different MatchIndex",
		n3q1,
		&pb.AppendEntriesRequest{Term: 5},
		[]*pb.AppendEntriesResponse{
			{
				Term:       5,
				MatchIndex: 50,
				Success:    false,
			},
			{
				Term:       5,
				MatchIndex: 100,
				Success:    false,
			},
		},
		true,
		&pb.AppendEntriesResponse{
			Term:       5,
			MatchIndex: 50,
			Success:    false,
		},
	},
	{
		"quorum successful",
		n3q1,
		&pb.AppendEntriesRequest{Term: 5},
		[]*pb.AppendEntriesResponse{
			{
				Term:       5,
				MatchIndex: 100,
				Success:    true,
			},
		},
		true,
		&pb.AppendEntriesResponse{
			Term:       5,
			MatchIndex: 100,
			Success:    true,
		},
	},
}

type spec struct {
	name string
	qs   gorums.QuorumSpec
}

var (
	n1q1 = spec{
		"QuorumSpec N1 Q1",
		raft.NewQuorumSpec(3),
	}
	n2q1 = spec{
		"QuorumSpec N2 Q1",
		raft.NewQuorumSpec(3),
	}
	n3q1 = spec{
		"QuorumSpec N3 Q1",
		raft.NewQuorumSpec(3),
	}
	n7q4 = spec{
		"QuorumSpec N7 Q4",
		raft.NewQuorumSpec(3),
	}
)

func TestRequestVoteQF(t *testing.T) {
	for _, test := range requestVoteQFTests {
		t.Run(test.qs.name+"-"+test.name, func(t *testing.T) {
			reply, quorum := test.qs.qs.RequestVoteQF(test.request, test.replies)

			if quorum != test.quorum {
				t.Errorf("got %t, want %t", quorum, test.quorum)
			}

			if !reflect.DeepEqual(reply, test.reply) {
				t.Errorf("got %+v, want %+v", reply, test.reply)
			}
		})
	}
}

func TestAppendEntriesFastQF(t *testing.T) {
	for _, test := range appendEntriesQFTests {
		t.Run(test.qs.name+"-"+test.name, func(t *testing.T) {
			reply, quorum := test.qs.qs.AppendEntriesQF(test.request, test.replies)

			if quorum != test.quorum {
				t.Errorf("got %t, want %t", quorum, test.quorum)
			}

			if !reflect.DeepEqual(reply, test.reply) {
				t.Errorf("got %+v, want %+v", reply, test.reply)
			}
		})
	}
}
<|endoftext|>