<commit_before>package metacmd

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

	"github.com/knq/dburl"
	"github.com/knq/usql/drivers"
	"github.com/knq/usql/env"
	"github.com/knq/usql/text"
)

// Cmd is a command implementation.
type Cmd struct {
	Section Section
	Name    string
	Desc    string
	Min     int
	Aliases map[string]string
	Process func(Handler, string, []string) (Res, error)
}

// cmds is the set of commands.
var cmds []Cmd

// cmdMap is the map of commands and their aliases.
var cmdMap map[string]Metacmd

// sectMap is the map of sections to its respective commands.
var sectMap map[Section][]Metacmd

func init() {
	cmds = []Cmd{
		Question: {
			Section: SectionHelp,
			Name:    "?",
			Desc:    "show help on backslash commands,[commands]",
			Aliases: map[string]string{
				"?":  "show help on " + text.CommandName + " command-line options,options",
				"? ": "show help on special variables,variables",
			},
			Process: func(h Handler, _ string, _ []string) (Res, error) {
				Listing(h.IO().Stdout())
				return Res{}, nil
			},
		},

		Quit: {
			Section: SectionGeneral,
			Name:    "q",
			Desc:    "quit " + text.CommandName,
			Aliases: map[string]string{"quit": ""},
			Process: func(Handler, string, []string) (Res, error) {
				return Res{Quit: true}, nil
			},
		},

		Copyright: {
			Section: SectionGeneral,
			Name:    "copyright",
			Desc:    "show " + text.CommandName + " usage and distribution terms",
			Process: func(h Handler, _ string, _ []string) (Res, error) {
				fmt.Fprintln(h.IO().Stdout(), text.Copyright)
				return Res{}, nil
			},
		},

		ConnInfo: {
			Section: SectionConnection,
			Name:    "conninfo",
			Desc:    "display information about the current database connection",
			Process: func(h Handler, _ string, _ []string) (Res, error) {
				if u := h.URL(); u != nil {
					out := h.IO().Stdout()
					fmt.Fprintf(out, text.ConnInfo, u.Driver, u.DSN)
					fmt.Fprintln(out)
				}
				return Res{}, nil
			},
		},

		Drivers: {
			Section: SectionGeneral,
			Name:    "drivers",
			Desc:    "display information about available database drivers",
			Process: func(h Handler, _ string, _ []string) (Res, error) {
				out := h.IO().Stdout()

				names := make([]string, len(drivers.Drivers))
				var z int
				for k := range drivers.Drivers {
					names[z] = k
					z++
				}
				sort.Strings(names)

				fmt.Fprintln(out, text.AvailableDrivers)
				for _, n := range names {
					s := "  " + n

					driver, aliases := dburl.SchemeDriverAndAliases(n)
					if driver != n {
						s += " (" + driver + ")"
					}
					if len(aliases) > 0 {
						if len(aliases) > 0 {
							s += " [" + strings.Join(aliases, ", ") + "]"
						}
					}
					fmt.Fprintln(out, s)
				}

				return Res{}, nil
			},
		},

		Connect: {
			Section: SectionConnection,
			Name:    "c",
			Desc:    "connect to database with url,URL",
			Aliases: map[string]string{
				"c":       "connect to database with SQL driver and parameters,DRIVER [PARAMS]",
				"connect": "",
			},
			Min: 1,
			Process: func(h Handler, _ string, params []string) (Res, error) {
				return Res{Processed: len(params)}, h.Open(params...)
			},
		},

		Disconnect: {
			Section: SectionConnection,
			Name:    "Z",
			Desc:    "close database connection",
			Aliases: map[string]string{"disconnect": ""},
			Process: func(h Handler, _ string, _ []string) (Res, error) {
				return Res{}, h.Close()
			},
		},

		Exec: {
			Section: SectionGeneral,
			Name:    "g",
			Desc:    "execute query (and send results to file or |pipe),[FILE] or ;",
			Aliases: map[string]string{
				"gexec": "execute query and execute each value of the result",
				"gset":  "execute query and store results in " + text.CommandName + " variables,[PREFIX]",
			},
			Process: func(h Handler, cmd string, params []string) (Res, error) {
				res := Res{
					Exec: ExecOnly,
				}

				switch cmd {
				case "g":
					if len(params) > 0 {
						res.ExecParam = params[0]
						res.Processed++
					}

				case "gexec":
					res.Exec = ExecExec

				case "gset":
					res.Exec = ExecSet
					if len(params) > 0 {
						res.ExecParam = params[0]
						res.Processed++
					}
				}

				return res, nil
			},
		},

		Edit: {
			Section: SectionQueryBuffer,
			Name:    "e",
			Desc:    "edit the query buffer (or file) with external editor,[FILE] [LINE]",
			Aliases: map[string]string{"edit": ""},
			Process: func(h Handler, _ string, params []string) (Res, error) {
				var res Res
				var path, line string

				// get path, line params
				if len(params) > 0 {
					path = env.Expand(params[0], h.User().HomeDir)
					res.Processed++
				}
				if len(params) > 1 {
					line = params[1]
					res.Processed++
				}

				// get last statement
				s, buf := h.Last(), h.Buf()
				if buf.Len != 0 {
					s = buf.String()
				}

				n, err := env.EditFile(path, line, s)

				// reset if no error
				if err == nil {
					buf.Reset()
					buf.Feed(n)
				}

				return res, err
			},
		},

		Print: {
			Section: SectionQueryBuffer,
			Name:    "p",
			Desc:    "show the contents of the query buffer",
			Aliases: map[string]string{"print": ""},
			Process: func(h Handler, _ string, _ []string) (Res, error) {
				// get last statement
				s, buf := h.Last(), h.Buf()
				if buf.Len != 0 {
					s = buf.String()
				}

				if s == "" {
					s = text.QueryBufferEmpty
				}

				fmt.Fprintln(h.IO().Stdout(), s)
				return Res{}, nil
			},
		},

		Reset: {
			Section: SectionQueryBuffer,
			Name:    "r",
			Desc:    "reset (clear) the query buffer",
			Aliases: map[string]string{"reset": ""},
			Process: func(h Handler, _ string, _ []string) (Res, error) {
				h.Buf().Reset()
				fmt.Fprintln(h.IO().Stdout(), text.QueryBufferReset)
				return Res{}, nil
			},
		},

		Echo: {
			Section: SectionInputOutput,
			Name:    "echo",
			Desc:    "write string to standard output,[STRING]",
			Process: func(h Handler, _ string, params []string) (Res, error) {
				fmt.Fprintln(h.IO().Stdout(), strings.Join(params, " "))
				return Res{Processed: len(params)}, nil
			},
		},

		Write: {
			Section: SectionQueryBuffer,
			Name:    "w",
			Min:     1,
			Desc:    "write query buffer to file,FILE",
			Aliases: map[string]string{"write": ""},
			Process: func(h Handler, _ string, params []string) (Res, error) {
				// get last statement
				s, buf := h.Last(), h.Buf()
				if buf.Len != 0 {
					s = buf.String()
				}

				return Res{Processed: 1}, ioutil.WriteFile(
					params[0],
					[]byte(strings.TrimSuffix(s, "\n")+"\n"),
					0644,
				)
			},
		},

		ChangeDir: {
			Section: SectionOperatingSystem,
			Name:    "cd",
			Desc:    "change the current working directory,[DIR]",
			Process: func(h Handler, _ string, params []string) (Res, error) {
				var res Res

				home, path := h.User().HomeDir, ""
				if len(params) > 0 {
					path = env.Expand(params[0], home)
					res.Processed++
				}

				return res, os.Chdir(path)
			},
		},

		SetEnv: {
			Section: SectionOperatingSystem,
			Name:    "setenv",
			Min:     1,
			Desc:    "set or unset environment variable,NAME [VALUE]",
			Process: func(h Handler, _ string, params []string) (Res, error) {
				var err error

				n := params[0]
				if len(params) == 1 {
					err = os.Unsetenv(n)
				} else {
					err = os.Setenv(n, strings.Join(params, " "))
				}

				return Res{Processed: len(params)}, err
			},
		},

		Include: {
			Section: SectionInputOutput,
			Name:    "i",
			Min:     1,
			Desc:    "execute commands from file,FILE",
			Aliases: map[string]string{
				"ir":               `as \i, but relative to location of current script,FILE`,
				"include":          "",
				"include_relative": "",
			},
			Process: func(h Handler, cmd string, params []string) (Res, error) {
				err := h.Include(
					env.Expand(params[0], h.User().HomeDir),
					cmd == "ir" || cmd == "include_relative",
				)
				if err != nil {
					err = fmt.Errorf("%s: %v", params[0], err)
				}
				return Res{Processed: 1}, err
			},
		},
	}

	// set up map
	cmdMap = make(map[string]Metacmd, len(cmds))
	sectMap = make(map[Section][]Metacmd, len(SectionOrder))
	for i, c := range cmds {
		mc := Metacmd(i)
		if mc == None {
			continue
		}

		cmdMap[c.Name] = mc
		for alias, _ := range c.Aliases {
			cmdMap[alias] = mc
		}

		sectMap[c.Section] = append(sectMap[c.Section], mc)
	}
}
<commit_msg>Minor change to \copyright command output<commit_after>package metacmd

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

	"github.com/knq/dburl"
	"github.com/knq/usql/drivers"
	"github.com/knq/usql/env"
	"github.com/knq/usql/text"
)

// Cmd is a command implementation.
type Cmd struct {
	Section Section
	Name    string
	Desc    string
	Min     int
	Aliases map[string]string
	Process func(Handler, string, []string) (Res, error)
}

// cmds is the set of commands.
var cmds []Cmd

// cmdMap is the map of commands and their aliases.
var cmdMap map[string]Metacmd

// sectMap is the map of sections to its respective commands.
var sectMap map[Section][]Metacmd

func init() {
	cmds = []Cmd{
		Question: {
			Section: SectionHelp,
			Name:    "?",
			Desc:    "show help on backslash commands,[commands]",
			Aliases: map[string]string{
				"?":  "show help on " + text.CommandName + " command-line options,options",
				"? ": "show help on special variables,variables",
			},
			Process: func(h Handler, _ string, _ []string) (Res, error) {
				Listing(h.IO().Stdout())
				return Res{}, nil
			},
		},

		Quit: {
			Section: SectionGeneral,
			Name:    "q",
			Desc:    "quit " + text.CommandName,
			Aliases: map[string]string{"quit": ""},
			Process: func(Handler, string, []string) (Res, error) {
				return Res{Quit: true}, nil
			},
		},

		Copyright: {
			Section: SectionGeneral,
			Name:    "copyright",
			Desc:    "show " + text.CommandName + " usage and distribution terms",
			Process: func(h Handler, _ string, _ []string) (Res, error) {
				out := h.IO().Stdout()
				fmt.Fprintln(out, text.Copyright)
				fmt.Fprintln(out)
				return Res{}, nil
			},
		},

		ConnInfo: {
			Section: SectionConnection,
			Name:    "conninfo",
			Desc:    "display information about the current database connection",
			Process: func(h Handler, _ string, _ []string) (Res, error) {
				if u := h.URL(); u != nil {
					out := h.IO().Stdout()
					fmt.Fprintf(out, text.ConnInfo, u.Driver, u.DSN)
					fmt.Fprintln(out)
				}
				return Res{}, nil
			},
		},

		Drivers: {
			Section: SectionGeneral,
			Name:    "drivers",
			Desc:    "display information about available database drivers",
			Process: func(h Handler, _ string, _ []string) (Res, error) {
				out := h.IO().Stdout()

				names := make([]string, len(drivers.Drivers))
				var z int
				for k := range drivers.Drivers {
					names[z] = k
					z++
				}
				sort.Strings(names)

				fmt.Fprintln(out, text.AvailableDrivers)
				for _, n := range names {
					s := "  " + n

					driver, aliases := dburl.SchemeDriverAndAliases(n)
					if driver != n {
						s += " (" + driver + ")"
					}
					if len(aliases) > 0 {
						if len(aliases) > 0 {
							s += " [" + strings.Join(aliases, ", ") + "]"
						}
					}
					fmt.Fprintln(out, s)
				}

				return Res{}, nil
			},
		},

		Connect: {
			Section: SectionConnection,
			Name:    "c",
			Desc:    "connect to database with url,URL",
			Aliases: map[string]string{
				"c":       "connect to database with SQL driver and parameters,DRIVER [PARAMS]",
				"connect": "",
			},
			Min: 1,
			Process: func(h Handler, _ string, params []string) (Res, error) {
				return Res{Processed: len(params)}, h.Open(params...)
			},
		},

		Disconnect: {
			Section: SectionConnection,
			Name:    "Z",
			Desc:    "close database connection",
			Aliases: map[string]string{"disconnect": ""},
			Process: func(h Handler, _ string, _ []string) (Res, error) {
				return Res{}, h.Close()
			},
		},

		Exec: {
			Section: SectionGeneral,
			Name:    "g",
			Desc:    "execute query (and send results to file or |pipe),[FILE] or ;",
			Aliases: map[string]string{
				"gexec": "execute query and execute each value of the result",
				"gset":  "execute query and store results in " + text.CommandName + " variables,[PREFIX]",
			},
			Process: func(h Handler, cmd string, params []string) (Res, error) {
				res := Res{
					Exec: ExecOnly,
				}

				switch cmd {
				case "g":
					if len(params) > 0 {
						res.ExecParam = params[0]
						res.Processed++
					}

				case "gexec":
					res.Exec = ExecExec

				case "gset":
					res.Exec = ExecSet
					if len(params) > 0 {
						res.ExecParam = params[0]
						res.Processed++
					}
				}

				return res, nil
			},
		},

		Edit: {
			Section: SectionQueryBuffer,
			Name:    "e",
			Desc:    "edit the query buffer (or file) with external editor,[FILE] [LINE]",
			Aliases: map[string]string{"edit": ""},
			Process: func(h Handler, _ string, params []string) (Res, error) {
				var res Res
				var path, line string

				// get path, line params
				if len(params) > 0 {
					path = env.Expand(params[0], h.User().HomeDir)
					res.Processed++
				}
				if len(params) > 1 {
					line = params[1]
					res.Processed++
				}

				// get last statement
				s, buf := h.Last(), h.Buf()
				if buf.Len != 0 {
					s = buf.String()
				}

				n, err := env.EditFile(path, line, s)

				// reset if no error
				if err == nil {
					buf.Reset()
					buf.Feed(n)
				}

				return res, err
			},
		},

		Print: {
			Section: SectionQueryBuffer,
			Name:    "p",
			Desc:    "show the contents of the query buffer",
			Aliases: map[string]string{"print": ""},
			Process: func(h Handler, _ string, _ []string) (Res, error) {
				// get last statement
				s, buf := h.Last(), h.Buf()
				if buf.Len != 0 {
					s = buf.String()
				}

				if s == "" {
					s = text.QueryBufferEmpty
				}

				fmt.Fprintln(h.IO().Stdout(), s)
				return Res{}, nil
			},
		},

		Reset: {
			Section: SectionQueryBuffer,
			Name:    "r",
			Desc:    "reset (clear) the query buffer",
			Aliases: map[string]string{"reset": ""},
			Process: func(h Handler, _ string, _ []string) (Res, error) {
				h.Buf().Reset()
				fmt.Fprintln(h.IO().Stdout(), text.QueryBufferReset)
				return Res{}, nil
			},
		},

		Echo: {
			Section: SectionInputOutput,
			Name:    "echo",
			Desc:    "write string to standard output,[STRING]",
			Process: func(h Handler, _ string, params []string) (Res, error) {
				fmt.Fprintln(h.IO().Stdout(), strings.Join(params, " "))
				return Res{Processed: len(params)}, nil
			},
		},

		Write: {
			Section: SectionQueryBuffer,
			Name:    "w",
			Min:     1,
			Desc:    "write query buffer to file,FILE",
			Aliases: map[string]string{"write": ""},
			Process: func(h Handler, _ string, params []string) (Res, error) {
				// get last statement
				s, buf := h.Last(), h.Buf()
				if buf.Len != 0 {
					s = buf.String()
				}

				return Res{Processed: 1}, ioutil.WriteFile(
					params[0],
					[]byte(strings.TrimSuffix(s, "\n")+"\n"),
					0644,
				)
			},
		},

		ChangeDir: {
			Section: SectionOperatingSystem,
			Name:    "cd",
			Desc:    "change the current working directory,[DIR]",
			Process: func(h Handler, _ string, params []string) (Res, error) {
				var res Res

				home, path := h.User().HomeDir, ""
				if len(params) > 0 {
					path = env.Expand(params[0], home)
					res.Processed++
				}

				return res, os.Chdir(path)
			},
		},

		SetEnv: {
			Section: SectionOperatingSystem,
			Name:    "setenv",
			Min:     1,
			Desc:    "set or unset environment variable,NAME [VALUE]",
			Process: func(h Handler, _ string, params []string) (Res, error) {
				var err error

				n := params[0]
				if len(params) == 1 {
					err = os.Unsetenv(n)
				} else {
					err = os.Setenv(n, strings.Join(params, " "))
				}

				return Res{Processed: len(params)}, err
			},
		},

		Include: {
			Section: SectionInputOutput,
			Name:    "i",
			Min:     1,
			Desc:    "execute commands from file,FILE",
			Aliases: map[string]string{
				"ir":               `as \i, but relative to location of current script,FILE`,
				"include":          "",
				"include_relative": "",
			},
			Process: func(h Handler, cmd string, params []string) (Res, error) {
				err := h.Include(
					env.Expand(params[0], h.User().HomeDir),
					cmd == "ir" || cmd == "include_relative",
				)
				if err != nil {
					err = fmt.Errorf("%s: %v", params[0], err)
				}
				return Res{Processed: 1}, err
			},
		},
	}

	// set up map
	cmdMap = make(map[string]Metacmd, len(cmds))
	sectMap = make(map[Section][]Metacmd, len(SectionOrder))
	for i, c := range cmds {
		mc := Metacmd(i)
		if mc == None {
			continue
		}

		cmdMap[c.Name] = mc
		for alias, _ := range c.Aliases {
			cmdMap[alias] = mc
		}

		sectMap[c.Section] = append(sectMap[c.Section], mc)
	}
}
<|endoftext|>


<commit_before>package http_util

import (
	// "crypto/md5"
	// "errors"
	"curl_cmd"
	// "fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"sync"
	"time"
)

const ChromeUserAgent string = "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.17 (KHTML, like Gecko) Chrome/24.0.1312.56 Safari/537.17"

type ResourceInfo struct {
	length   int64
	filename string
}
type RedirectError struct {
	s string
}

func (re RedirectError) Error() string {
	return re.s
}

func CheckRedirect(*http.Request, []*http.Request) error {
	return RedirectError{"don't follow redirect"}
}

type myUrlEror url.Error

func (e myUrlEror) Error() string {
	return (&e).Error()
}

var open_file_func openFileFunc = NewFile

const (
	BlockSize                 int64 = 1024 * 1024
	NBlocksPerRequest               = 300
	DownloaderChanBufferSize        = 100
	TimeoutOfGetResourceInfo        = 60 * time.Second
	TimeoutOfPerBlockDownload       = 1024 / 42 * time.Second
)

type DownloadRange struct {
	Start  int64
	Length int64
}

type File interface {
	Size() int64
	Name() string
	io.ReadWriteSeeker
	Truncate(size int64) error
	WriteAt([]byte, int64) (int, error)
	Sync() error
	Close() error
}

type openFileFunc func(string) (File, error)

type FileS struct {
	os.File
}

func (f *FileS) Size() int64 {
	stat, err := f.Stat()
	if err != nil {
		return 0
	}
	return stat.Size()
}

func NewFile(name string) (File, error) {
	f, err := os.OpenFile(name, os.O_RDWR|os.O_CREATE, 0666)
	fs := FileS{*f}
	return &fs, err
}

func Truncate(name string, size int64) error {
	f, err := os.Create(name)
	if err != nil {
		return err
	}
	f.Close()
	err = os.Truncate(name, size)
	return err
}

type Request struct {
	url    string
	header http.Header
}

type DownloadTaskInfo struct {
	DownloadRange
	Request
	Name        string
	FailedTimes int
	LastWorkerN int
}

type DownloadChunk struct {
	Data  []byte
	Start int64
	Name  string
}

func Downloader(task_info_c <-chan DownloadTaskInfo, finished_c chan<- DownloadChunk, failed_task_info_c chan<- DownloadTaskInfo, wg *sync.WaitGroup, worker_n int) {
	for task_info := range task_info_c {
		length := task_info.Length
		start := task_info.Start
		name := task_info.Name
		log.Printf("Worker[%v] %v, %v", name, worker_n, task_info.DownloadRange)
		var downloaded int64 = 0
		task_start_time := GetNowEpochInMilli()
		for try_times := 0; try_times < 3; try_times++ {
			chunk_datas := make(chan []byte, 1)
			go RangeGet(task_info.Request, start, length, chunk_datas)
			for chunk_data := range chunk_datas {
				downloaded += int64(len(chunk_data))
				log.Printf("Wroker[%v] %v %v k/s, %v%%", name, worker_n, downloaded*1000/1024/(GetNowEpochInMilli()-task_start_time), 100*downloaded/length)
				finished_c <- DownloadChunk{Data: chunk_data, Name: name, Start: start}
				start += int64(len(chunk_data))
			}
			if downloaded == length {
				wg.Done()
				break
			}
		}
		if downloaded < length {
			task_info.Start = start
			task_info.Length = length - downloaded
			task_info.LastWorkerN = worker_n
			task_info.FailedTimes++
			failed_task_info_c <- task_info
		}
	}
}

func DownloadChunkWaitGroupAutoCloser(wg *sync.WaitGroup, c chan<- DownloadChunk) {
	wg.Wait()
	close(c)
}

func Receiver(fileDownloadInfoC <-chan FileDownloadInfo, chunks <-chan DownloadChunk, finished chan<- int) {
	defer func() {
		finished <- 0
	}()
	fileDownloadInfoMap := make(map[string]FileDownloadInfo)
	fileFdMap := make(map[string]File)
	for {
		select {
		case info, ok := <-fileDownloadInfoC:
			if !ok {
				if len(fileDownloadInfoMap) == 0 {
					return
				}
				fileDownloadInfoC = nil
				break
			}
			if info.Finished() {
				break
			}
			fileDownloadInfoMap[info.Name] = info
			fd, err := open_file_func(info.Name)
			if err != nil {
				log.Fatalf("can't open %v", info.Name)
			}
			fileFdMap[info.Name] = fd
		case chunk, ok := <-chunks:
			if !ok {
				return
			}
			info, ok := fileDownloadInfoMap[chunk.Name]
			info.Update(chunk.Start, int64(len(chunk.Data)))
			info.Sync()
			if !ok {
				log.Fatalf("can't find chunk.Name in info_map, %v", chunk.Name)
			}
			fd, ok := fileFdMap[chunk.Name]
			if !ok {
				log.Fatalf("can't find chunk.Name in file_fd_map, %v", chunk.Name)
			}
			_, err := fd.WriteAt(chunk.Data, chunk.Start)
			if err != nil {
				log.Fatalf("can't write content to file : %v", chunk.Name)
			}
			err = fd.Sync()
			if err != nil {
				log.Fatalf("can't write content to file : %v", chunk.Name)
			}
			if info.Finished() {
				delete(fileDownloadInfoMap, chunk.Name)
				fd.Close()
				delete(fileFdMap, chunk.Name)
				log.Print(info, info.Finished(), len(fileDownloadInfoMap))
				if len(fileDownloadInfoMap) == 0 {
					return
				}
			}
		}
	}
}

func Run(curl_cmd_strs []string, num_of_workers int) {
	file_download_info_c := make(chan FileDownloadInfo, 1)
	downloader_task_info_cs := make([]chan DownloadTaskInfo, num_of_workers)
	chunk_c := make(chan DownloadChunk, 1)
	url_chan_map := make(map[string]int)
	file_name_reqs_map := make(map[string]*[]Request)
	receiver_finish_channel := make(chan int, 1)
	task_info_wait_group := sync.WaitGroup{}
	failed_task_info_c := make(chan DownloadTaskInfo, 1)

	go Receiver(file_download_info_c, chunk_c, receiver_finish_channel)
	for i := 0; i < num_of_workers; i++ {
		tmp := make(chan DownloadTaskInfo, DownloaderChanBufferSize)
		downloader_task_info_cs[i] = tmp
		go Downloader(tmp, chunk_c, failed_task_info_c, &task_info_wait_group, i)
	}

	file_download_infos := []FileDownloadInfo{}
	// resource_infos := []ResourceInfo{}
	worker_n := -1
	task_infos := []DownloadTaskInfo{}
	for _, curl_cmd_str := range curl_cmd_strs {
		worker_n++
		worker_n = worker_n % num_of_workers
		url := curl_cmd.ParseCmdStr(curl_cmd_str)[1]
		// log.Printf("%v %v", url, worker_n)
		url_chan_map[url] = worker_n
		header := curl_cmd.GetHeadersFromCurlCmd(curl_cmd_str)
		req := Request{url, header}
		// reqs = append(reqs, req)
		// reqs := []Request{req}
		resource_info, err := GetResourceInfo(url, header)
		if err != nil {
			log.Printf("Can't get resource_info of url(%v), err(%v)", url, err)
			continue
		}
		file_name_reqs, ok := file_name_reqs_map[resource_info.filename]
		if !ok {
			tmp := make([]Request, 0, 1)
			file_name_reqs = &tmp
			file_name_reqs_map[resource_info.filename] = file_name_reqs
			*file_name_reqs = append(*file_name_reqs, req)
		} else {
			*file_name_reqs = append(*file_name_reqs, req)
			continue
		}

		log.Printf("Get Resource %v %v", resource_info.filename, resource_info.length)
		file_download_info, err := NewFileDownloadInfo(resource_info.filename, resource_info.length)
		if err != nil {
			log.Printf("Can't create file downloaded info of %v, %v", resource_info.filename, resource_info.length)
			continue
		}
		err = file_download_info.Sync()
		if err != nil {
			log.Printf("Can't sync file downloaded info of %v %v", resource_info.filename, resource_info.length)
			continue
		}
		file_download_info_c <- *file_download_info
		file_download_infos = append(file_download_infos, *file_download_info)
		for _, range_ := range file_download_info.UndownloadedRanges() {
			task := DownloadTaskInfo{range_, Request{}, resource_info.filename, 0, worker_n}
			task_info_wait_group.Add(1)
			task_infos = append(task_infos, task)
		}
	}
	all_task_finish_c := make(chan bool)
	go ConvertWaitGroupToBoolChan(&task_info_wait_group, all_task_finish_c)
	close(file_download_info_c)
	for file_name, reqs := range file_name_reqs_map {
		log.Printf("%v %v", file_name, len(*reqs))
	}
	worker_n = -1
	for _, task_info := range task_infos {
		worker_n++
		reqs := file_name_reqs_map[task_info.Name]
		req := (*reqs)[worker_n%len(*reqs)]
		task_info.Request = req
		// log.Printf("%v %v %v", task_info.Name, task_info.DownloadRange, url_chan_map[task_info.url])
		select {
		case downloader_task_info_cs[url_chan_map[task_info.url]] <- task_info:
		case failed_task_info := <-failed_task_info_c:
			if failed_task_info.FailedTimes < 3 {
				reqs := file_name_reqs_map[failed_task_info.Name]
				last_worker_n := failed_task_info.LastWorkerN
				new_req := (*reqs)[(last_worker_n+1)%len(*reqs)]
				failed_task_info.Request = new_req
				downloader_task_info_cs[url_chan_map[new_req.url]] <- failed_task_info
			} else {
				task_info_wait_group.Done()
			}
		}
	}
ForLoop:
	for {
		select {
		case failed_task_info := <-failed_task_info_c:
			if failed_task_info.FailedTimes < 3 {
				reqs := file_name_reqs_map[failed_task_info.Name]
				last_worker_n := failed_task_info.LastWorkerN
				new_req := (*reqs)[(last_worker_n+1)%len(*reqs)]
				failed_task_info.Request = new_req
				downloader_task_info_cs[url_chan_map[new_req.url]] <- failed_task_info
			} else {
				task_info_wait_group.Done()
			}
		case <-all_task_finish_c:
			for _, chan_ := range downloader_task_info_cs {
				close(chan_)
			}
			close(chunk_c)
			close(failed_task_info_c)
			break ForLoop
		}
	}
}
<commit_msg>minor changes<commit_after>package http_util

import (
	// "crypto/md5"
	// "errors"
	"curl_cmd"
	// "fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"sync"
	"time"
)

const ChromeUserAgent string = "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.17 (KHTML, like Gecko) Chrome/24.0.1312.56 Safari/537.17"

type ResourceInfo struct {
	length   int64
	filename string
}
type RedirectError struct {
	s string
}

func (re RedirectError) Error() string {
	return re.s
}

func CheckRedirect(*http.Request, []*http.Request) error {
	return RedirectError{"don't follow redirect"}
}

type myUrlEror url.Error

func (e myUrlEror) Error() string {
	return (&e).Error()
}

var open_file_func openFileFunc = NewFile

const (
	BlockSize                 int64 = 1024 * 1024
	NBlocksPerRequest               = 300
	DownloaderChanBufferSize        = 100
	TimeoutOfGetResourceInfo        = 60 * time.Second
	TimeoutOfPerBlockDownload       = 1024 / 42 * time.Second
)

type DownloadRange struct {
	Start  int64
	Length int64
}

type File interface {
	Size() int64
	Name() string
	io.ReadWriteSeeker
	Truncate(size int64) error
	WriteAt([]byte, int64) (int, error)
	Sync() error
	Close() error
}

type openFileFunc func(string) (File, error)

type FileS struct {
	os.File
}

func (f *FileS) Size() int64 {
	stat, err := f.Stat()
	if err != nil {
		return 0
	}
	return stat.Size()
}

func NewFile(name string) (File, error) {
	f, err := os.OpenFile(name, os.O_RDWR|os.O_CREATE, 0666)
	fs := FileS{*f}
	return &fs, err
}

func Truncate(name string, size int64) error {
	f, err := os.Create(name)
	if err != nil {
		return err
	}
	f.Close()
	err = os.Truncate(name, size)
	return err
}

type Request struct {
	url    string
	header http.Header
}

type DownloadTaskInfo struct {
	DownloadRange
	Request
	Name        string
	FailedTimes int
	LastWorkerN int
}

type DownloadChunk struct {
	Data  []byte
	Start int64
	Name  string
}

func Downloader(task_info_c <-chan DownloadTaskInfo, finished_c chan<- DownloadChunk, failed_task_info_c chan<- DownloadTaskInfo, wg *sync.WaitGroup, worker_n int) {
	for task_info := range task_info_c {
		length := task_info.Length
		start := task_info.Start
		name := task_info.Name
		log.Printf("Worker[%v] %v %v", name, worker_n, task_info.DownloadRange)
		var downloaded int64 = 0
		task_start_time := GetNowEpochInMilli()
		for try_times := 0; try_times < 3; try_times++ {
			chunk_datas := make(chan []byte, 1)
			go RangeGet(task_info.Request, start, length, chunk_datas)
			for chunk_data := range chunk_datas {
				downloaded += int64(len(chunk_data))
				log.Printf("Wroker[%v] %v %v k/s %v%%", name, worker_n, downloaded*1000/1024/(GetNowEpochInMilli()-task_start_time), 100*downloaded/length)
				finished_c <- DownloadChunk{Data: chunk_data, Name: name, Start: start}
				start += int64(len(chunk_data))
			}
			if downloaded == length {
				wg.Done()
				break
			}
		}
		if downloaded < length {
			task_info.Start = start
			task_info.Length = length - downloaded
			task_info.LastWorkerN = worker_n
			task_info.FailedTimes++
			failed_task_info_c <- task_info
		}
	}
}

func DownloadChunkWaitGroupAutoCloser(wg *sync.WaitGroup, c chan<- DownloadChunk) {
	wg.Wait()
	close(c)
}

func Receiver(fileDownloadInfoC <-chan FileDownloadInfo, chunks <-chan DownloadChunk, finished chan<- int) {
	defer func() {
		finished <- 0
	}()
	fileDownloadInfoMap := make(map[string]FileDownloadInfo)
	fileFdMap := make(map[string]File)
	for {
		select {
		case info, ok := <-fileDownloadInfoC:
			if !ok {
				if len(fileDownloadInfoMap) == 0 {
					return
				}
				fileDownloadInfoC = nil
				break
			}
			if info.Finished() {
				break
			}
			fileDownloadInfoMap[info.Name] = info
			fd, err := open_file_func(info.Name)
			if err != nil {
				log.Fatalf("can't open %v", info.Name)
			}
			fileFdMap[info.Name] = fd
		case chunk, ok := <-chunks:
			if !ok {
				return
			}
			info, ok := fileDownloadInfoMap[chunk.Name]
			info.Update(chunk.Start, int64(len(chunk.Data)))
			info.Sync()
			if !ok {
				log.Fatalf("can't find chunk.Name in info_map, %v", chunk.Name)
			}
			fd, ok := fileFdMap[chunk.Name]
			if !ok {
				log.Fatalf("can't find chunk.Name in file_fd_map, %v", chunk.Name)
			}
			_, err := fd.WriteAt(chunk.Data, chunk.Start)
			if err != nil {
				log.Fatalf("can't write content to file : %v", chunk.Name)
			}
			err = fd.Sync()
			if err != nil {
				log.Fatalf("can't write content to file : %v", chunk.Name)
			}
			if info.Finished() {
				delete(fileDownloadInfoMap, chunk.Name)
				fd.Close()
				delete(fileFdMap, chunk.Name)
				log.Print(info, info.Finished(), len(fileDownloadInfoMap))
				if len(fileDownloadInfoMap) == 0 {
					return
				}
			}
		}
	}
}

func Run(curl_cmd_strs []string, num_of_workers int) {
	file_download_info_c := make(chan FileDownloadInfo, 1)
	downloader_task_info_cs := make([]chan DownloadTaskInfo, num_of_workers)
	chunk_c := make(chan DownloadChunk, 1)
	url_chan_map := make(map[string]int)
	file_name_reqs_map := make(map[string]*[]Request)
	receiver_finish_channel := make(chan int, 1)
	task_info_wait_group := sync.WaitGroup{}
	failed_task_info_c := make(chan DownloadTaskInfo, 1)

	go Receiver(file_download_info_c, chunk_c, receiver_finish_channel)
	for i := 0; i < num_of_workers; i++ {
		tmp := make(chan DownloadTaskInfo, DownloaderChanBufferSize)
		downloader_task_info_cs[i] = tmp
		go Downloader(tmp, chunk_c, failed_task_info_c, &task_info_wait_group, i)
	}

	file_download_infos := []FileDownloadInfo{}
	// resource_infos := []ResourceInfo{}
	worker_n := -1
	task_infos := []DownloadTaskInfo{}
	for _, curl_cmd_str := range curl_cmd_strs {
		worker_n++
		worker_n = worker_n % num_of_workers
		url := curl_cmd.ParseCmdStr(curl_cmd_str)[1]
		// log.Printf("%v %v", url, worker_n)
		url_chan_map[url] = worker_n
		header := curl_cmd.GetHeadersFromCurlCmd(curl_cmd_str)
		req := Request{url, header}
		// reqs = append(reqs, req)
		// reqs := []Request{req}
		resource_info, err := GetResourceInfo(url, header)
		if err != nil {
			log.Printf("Can't get resource_info of url(%v), err(%v)", url, err)
			continue
		}
		file_name_reqs, ok := file_name_reqs_map[resource_info.filename]
		if !ok {
			tmp := make([]Request, 0, 1)
			file_name_reqs = &tmp
			file_name_reqs_map[resource_info.filename] = file_name_reqs
			*file_name_reqs = append(*file_name_reqs, req)
		} else {
			*file_name_reqs = append(*file_name_reqs, req)
			continue
		}

		log.Printf("Get Resource %v %v", resource_info.filename, resource_info.length)
		file_download_info, err := NewFileDownloadInfo(resource_info.filename, resource_info.length)
		if err != nil {
			log.Printf("Can't create file downloaded info of %v, %v", resource_info.filename, resource_info.length)
			continue
		}
		err = file_download_info.Sync()
		if err != nil {
			log.Printf("Can't sync file downloaded info of %v %v", resource_info.filename, resource_info.length)
			continue
		}
		file_download_info_c <- *file_download_info
		file_download_infos = append(file_download_infos, *file_download_info)
		for _, range_ := range file_download_info.UndownloadedRanges() {
			task := DownloadTaskInfo{range_, Request{}, resource_info.filename, 0, worker_n}
			task_info_wait_group.Add(1)
			task_infos = append(task_infos, task)
		}
	}
	all_task_finish_c := make(chan bool)
	go ConvertWaitGroupToBoolChan(&task_info_wait_group, all_task_finish_c)
	close(file_download_info_c)
	for file_name, reqs := range file_name_reqs_map {
		log.Printf("%v %v", file_name, len(*reqs))
	}
	worker_n = -1
	for _, task_info := range task_infos {
		worker_n++
		reqs := file_name_reqs_map[task_info.Name]
		req := (*reqs)[worker_n%len(*reqs)]
		task_info.Request = req
		// log.Printf("%v %v %v", task_info.Name, task_info.DownloadRange, url_chan_map[task_info.url])
		select {
		case downloader_task_info_cs[url_chan_map[task_info.url]] <- task_info:
		case failed_task_info := <-failed_task_info_c:
			if failed_task_info.FailedTimes < 3 {
				reqs := file_name_reqs_map[failed_task_info.Name]
				last_worker_n := failed_task_info.LastWorkerN
				new_req := (*reqs)[(last_worker_n+1)%len(*reqs)]
				failed_task_info.Request = new_req
				downloader_task_info_cs[url_chan_map[new_req.url]] <- failed_task_info
			} else {
				task_info_wait_group.Done()
			}
		}
	}
ForLoop:
	for {
		select {
		case failed_task_info := <-failed_task_info_c:
			if failed_task_info.FailedTimes < 3 {
				reqs := file_name_reqs_map[failed_task_info.Name]
				last_worker_n := failed_task_info.LastWorkerN
				new_req := (*reqs)[(last_worker_n+1)%len(*reqs)]
				failed_task_info.Request = new_req
				downloader_task_info_cs[url_chan_map[new_req.url]] <- failed_task_info
			} else {
				task_info_wait_group.Done()
			}
		case <-all_task_finish_c:
			for _, chan_ := range downloader_task_info_cs {
				close(chan_)
			}
			close(chunk_c)
			close(failed_task_info_c)
			break ForLoop
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"log"
	"testing"
)

func TestRun(t *testing.T) {
	got := run(NewApp())
	if got != 0 {
		t.Error("boom")
	}
}

func TestNewApp(t *testing.T) {
	if NewApp() == nil {
		t.Error("NewApp() = <nil>")
	}
}

func loggerTest(t *testing.T, funcname, want string) {
	var buf bytes.Buffer

	a := &App{logger: log.New(&buf, "", 0)}

	a.Log("text", 12)

	got := buf.String()

	if got != want {
		t.Errorf("Log(...) = %q; want %q", got, want)
	}
}

func TestApp_Log(t *testing.T) {
	var buf bytes.Buffer

	a := &App{logger: log.New(&buf, "", 0)}

	a.Log("text", 12)

	got := buf.String()
	want := "text12\n"

	if got != want {
		t.Errorf("Log(...) = %q; want %q", got, want)
	}
}

func TestApp_Logf(t *testing.T) {
	var buf bytes.Buffer

	a := &App{logger: log.New(&buf, "", 0)}

	a.Logf("%d %s", 46, "text")

	got := buf.String()
	want := "46 text\n"

	if got != want {
		t.Errorf("Logf(...) = %q; want %q", got, want)
	}
}
<commit_msg>app: clean TestApp_Log<commit_after>package main

import (
	"bytes"
	"log"
	"testing"
)

func TestRun(t *testing.T) {
	got := run(NewApp())
	if got != 0 {
		t.Error("boom")
	}
}

func TestNewApp(t *testing.T) {
	if NewApp() == nil {
		t.Error("NewApp() = <nil>")
	}
}

func loggerTest(t *testing.T, funcname, want string) {
	var buf bytes.Buffer

	a := &App{logger: log.New(&buf, "", 0)}

	a.Log("text", 12)

	got := buf.String()

	if got != want {
		t.Errorf("Log(...) = %q; want %q", got, want)
	}
}

func TestApp_Log(t *testing.T) {
	loggerTest(t, "Log", "text 12\n")
}

func TestApp_Logf(t *testing.T) {
	var buf bytes.Buffer

	a := &App{logger: log.New(&buf, "", 0)}

	a.Logf("%d %s", 46, "text")

	got := buf.String()
	want := "46 text\n"

	if got != want {
		t.Errorf("Logf(...) = %q; want %q", got, want)
	}
}
<|endoftext|>


<commit_before>package osutils

import (
	. "launchpad.net/gocheck"
)

type SystemCommandTestSuite struct{}

var _ = Suite(&SystemCommandTestSuite{})

func (s *SystemCommandTestSuite) TestBuildEnvVars(c *C) {
	sc := &SystemCommand{
		EnvVars: map[string]string{
			"FOO": "bar",
			"BAR": "baz",
		},
	}

	obtained := sc.buildEnvVars()
	expected := []string{"FOO=bar", "BAR=baz"}
	c.Assert(obtained, DeepEquals, expected)
}

func (s *SystemCommandTestSuite) TestBuildCmd_ShellExpansionDisabled(c *C) {
	path := "/foo/bar"
	args := []string{"a", "b"}

	sc := &SystemCommand{
		Path:                 path,
		Args:                 args,
		EnableShellExpansion: false,
	}

	cmd := sc.buildCmd()
	c.Assert(cmd.Path, Equals, path)
	c.Assert(cmd.Args, DeepEquals, args)
}

func (s *SystemCommandTestSuite) TestBuildCmd_ShellExpansionEnabled(c *C) {
	path := "/foo/bar"
	args := []string{"a", "b"}

	sc := &SystemCommand{
		Path:                 path,
		Args:                 args,
		EnableShellExpansion: true,
	}

	cmd := sc.buildCmd()
	c.Assert(cmd.Path, Equals, "/bin/sh")
	c.Assert(cmd.Args, DeepEquals, []string{"sh", "-c", "/foo/bar", "a", "b"})
}

func (s *SystemCommandTestSuite) TestRun_CommandFailedWrongPath(c *C) {
	sc := &SystemCommand{
		Path:                 "/path/to/inexistant/command",
		Args:                 []string{"a", "b"},
		EnableShellExpansion: false,
	}

	err := sc.Run()
	c.Assert(err, NotNil)
	expected := `Error with command "/path/to/inexistant/command a b". Error message : "fork/exec /path/to/inexistant/command: no such file or directory".`
	c.Assert(err.Error(), Equals, expected)
}

func (s *SystemCommandTestSuite) TestRun_CommandFailed(c *C) {
	sc := &SystemCommand{
		Path:                 "/usr/bin/tr",
		Args:                 []string{"--xxx"},
		EnableShellExpansion: false,
	}

	err := sc.Run()
	// TODO : improve error msg check
	c.Assert(err, NotNil)
}
<commit_msg>Fixed broken test<commit_after>package osutils

import (
	. "launchpad.net/gocheck"
)

type SystemCommandTestSuite struct{}

var _ = Suite(&SystemCommandTestSuite{})

func (s *SystemCommandTestSuite) TestBuildEnvVars(c *C) {
	sc := &SystemCommand{
		EnvVars: map[string]string{
			"FOO": "bar",
			"BAR": "baz",
		},
	}

	obtained := sc.buildEnvVars()
	expected := []string{"FOO=bar", "BAR=baz"}
	c.Assert(obtained, DeepEquals, expected)
}

func (s *SystemCommandTestSuite) TestBuildCmd_ShellExpansionDisabled(c *C) {
	path := "/foo/bar"
	args := []string{"a", "b"}

	sc := &SystemCommand{
		Path:                 path,
		Args:                 args,
		EnableShellExpansion: false,
	}

	cmd := sc.buildCmd()
	c.Assert(cmd.Path, Equals, path)
	c.Assert(cmd.Args, DeepEquals, args)
}

func (s *SystemCommandTestSuite) TestBuildCmd_ShellExpansionEnabled(c *C) {
	path := "/foo/bar"
	args := []string{"a", "b"}

	sc := &SystemCommand{
		Path:                 path,
		Args:                 args,
		EnableShellExpansion: true,
	}

	cmd := sc.buildCmd()
	c.Assert(cmd.Path, Equals, "/bin/sh")
	c.Assert(cmd.Args, DeepEquals, []string{"sh", "-c", "/foo/bar", "a", "b"})
}

func (s *SystemCommandTestSuite) TestRun_CommandFailedWrongPath(c *C) {
	sc := &SystemCommand{
		Path:                 "/path/to/inexistant/command",
		Args:                 []string{"a", "b"},
		EnableShellExpansion: false,
	}

	err := sc.Run()
	c.Assert(err, NotNil)
	expected := `Error with command "/path/to/inexistant/command a b". Error message was "fork/exec /path/to/inexistant/command: no such file or directory".`
	c.Assert(err.Error(), Equals, expected)
}

func (s *SystemCommandTestSuite) TestRun_CommandFailed(c *C) {
	sc := &SystemCommand{
		Path:                 "/usr/bin/tr",
		Args:                 []string{"--xxx"},
		EnableShellExpansion: false,
	}

	err := sc.Run()
	// TODO : improve error msg check
	c.Assert(err, NotNil)
}
<|endoftext|>


<commit_before>package lib

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"strings"
)

func backend(c *Config, r *http.Request) (string, string, bool) {
	ps := strings.SplitN(r.URL.Path, "/", 3)
	if len(ps) != 3 {
		return tryFallback(c, r)
	}
	rules, ok := c.Versions[strings.ToLower(ps[1])]
	if !ok {
		return tryFallback(c, r)
	}
	pathToMatch := "/" + ps[2]
	for k, v := range rules {
		if strings.Index(pathToMatch, k) == 0 {
			return v, pathToMatch, true
		}
	}
	return "", "", false
}

func tryFallback(c *Config, r *http.Request) (string, string, bool){
	if c.FallbackRule != "" {
		return c.FallbackRule, r.URL.Path, true
	}
	return "", "", false
}

func bytesToResponse(statusCode int, contentType string, buffer []byte) *http.Response {
	r := ioutil.NopCloser(bytes.NewReader(buffer))
	h := http.Header{}
	h.Set("Content-type", contentType)
	return &http.Response{
		Status:           "",
		StatusCode:       statusCode,
		Header: h,
		Body:             r,
		ContentLength:    int64(len(buffer)),
	}
}

//  writeJSON marshals a JSON object to bytes and writes as body, ignoring any errors (will just be an empty response)
func (c *Config) writeJSON(req *http.Request, w http.ResponseWriter, status int, obj interface{})  {
	resp, _ := json.Marshal(obj)
	w.WriteHeader(status)
	w.Header().Set("Content-type", "application/json")
	w.Write(resp)
	c.Interceptor(req, bytesToResponse(status, "application/json", resp))
}

// New creates a new gateway.
func New(c *Config) http.HandlerFunc {
	c.setDefaults()
	return func(w http.ResponseWriter, req *http.Request) {

		//  1. Apply Filter
		allow, filterStatus, filterBody := c.Filter(req)
		if !allow {
			c.writeJSON(req, w, filterStatus, filterBody)
			return
		}

		// 2. Find backend, or return "not found" if not found
		b, url, ok := backend(c, req)
		if !ok {
			c.writeJSON(req, w, http.StatusNotFound, c.NotFoundResponse)
			return
		}

		// 3. Reverse proxy request
		(&httputil.ReverseProxy{
			Director: func(r *http.Request) {
				r.URL.Scheme = c.Scheme
				r.URL.Host = b
				r.URL.Path = url
				r.Host = b
			},
			ModifyResponse: func(resp *http.Response) error {
				c.Interceptor(req, resp)
				return nil
			},
		}).ServeHTTP(w, req)
	}
}
<commit_msg>Fix issue where the request body isn't available because it has already been read<commit_after>package lib

import (
	"bytes"
	"context"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"strings"
)

func backend(c *Config, r *http.Request) (string, string, bool) {
	ps := strings.SplitN(r.URL.Path, "/", 3)
	if len(ps) != 3 {
		return tryFallback(c, r)
	}
	rules, ok := c.Versions[strings.ToLower(ps[1])]
	if !ok {
		return tryFallback(c, r)
	}
	pathToMatch := "/" + ps[2]
	for k, v := range rules {
		if strings.Index(pathToMatch, k) == 0 {
			return v, pathToMatch, true
		}
	}
	return "", "", false
}

func tryFallback(c *Config, r *http.Request) (string, string, bool){
	if c.FallbackRule != "" {
		return c.FallbackRule, r.URL.Path, true
	}
	return "", "", false
}

func bytesToResponse(statusCode int, contentType string, buffer []byte) *http.Response {
	r := ioutil.NopCloser(bytes.NewReader(buffer))
	h := http.Header{}
	h.Set("Content-type", contentType)
	return &http.Response{
		Status:           "",
		StatusCode:       statusCode,
		Header: h,
		Body:             r,
		ContentLength:    int64(len(buffer)),
	}
}

//  writeJSON marshals a JSON object to bytes and writes as body, ignoring any errors (will just be an empty response)
func (c *Config) writeJSON(req *http.Request, w http.ResponseWriter, status int, obj interface{})  {
	resp, _ := json.Marshal(obj)
	w.WriteHeader(status)
	w.Header().Set("Content-type", "application/json")
	w.Write(resp)
	c.Interceptor(req, bytesToResponse(status, "application/json", resp))
}

func clone(r *http.Request) *http.Request {
	r2 := r.Clone(context.Background())
	bodyBytes, _ := ioutil.ReadAll(r.Body)
	r.Body.Close()  //  must close
	bodyCopy := make([]byte, len(bodyBytes))
	copy(bodyBytes, bodyCopy)
	r.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes))
	r2.Body = ioutil.NopCloser(bytes.NewBuffer(bodyCopy))
	return r2
}

// New creates a new gateway.
func New(c *Config) http.HandlerFunc {
	c.setDefaults()
	return func(w http.ResponseWriter, req *http.Request) {

		//  1. Apply Filter
		allow, filterStatus, filterBody := c.Filter(req)
		if !allow {
			c.writeJSON(req, w, filterStatus, filterBody)
			return
		}

		// 2. Find backend, or return "not found" if not found
		b, url, ok := backend(c, req)
		if !ok {
			c.writeJSON(req, w, http.StatusNotFound, c.NotFoundResponse)
			return
		}
		clonedReq := clone(req)
		// 3. Reverse proxy request
		(&httputil.ReverseProxy{
			Director: func(r *http.Request) {
				r.URL.Scheme = c.Scheme
				r.URL.Host = b
				r.URL.Path = url
				r.Host = b
			},
			ModifyResponse: func(resp *http.Response) error {
				c.Interceptor(clonedReq, resp)
				return nil
			},
		}).ServeHTTP(w, req)
	}
}
<|endoftext|>


<commit_before>package libcomfo

import (
	"bytes"
	"fmt"
	"io"
	"reflect"
	"testing"
	"time"
)

// TestConn is a structure that implements io.ReadWriter and can be used
// to mock a connection with the unit. HangStart makes initial Read and Write
// calls hang. HangEOF makes Read hang on the second call when a Scanner
// would expect an EOF error to be returned.
// Not thread-safe.
type TestConn struct {
	Receives []byte
	Emits    []byte

	HangTime  time.Duration
	HangStart bool
	HangEOF   bool

	emitted  int
	received int
}

func (tr *TestConn) Read(p []byte) (n int, err error) {

	// Simulate read delay before returning output or EOF
	if tr.HangTime > 0 && tr.HangStart {
		time.Sleep(tr.HangTime)
	}

	// Return EOF when buffer is fully copied
	if tr.emitted == len(tr.Emits) {

		if tr.HangTime > 0 && tr.HangEOF {
			time.Sleep(tr.HangTime)
		}

		return 0, io.EOF
	}

	// Return the data in the connection's Emits field
	n = copy(p, tr.Emits)

	// Keep track of how many bytes we've been able to copy
	tr.emitted = tr.emitted + n

	return
}

func (tr *TestConn) Write(p []byte) (n int, err error) {

	// Simulate write delay
	if tr.HangTime > 0 {
		time.Sleep(tr.HangTime)
	}

	// Bounds check for Receives
	if tr.received+len(p) > len(tr.Receives) {
		return n,
			fmt.Errorf("write to TestConn exceeded expected input specified in Receives: %v",
				// This formula determines the offset of the slice
				// of surplus characters over the struct's 'Receives' field.
				p[len(p)-(tr.received+len(p)-len(tr.Receives)):])
	}

	// Expect the written data to correspond to what's in the 'Receives' field.
	// Offset the comparison against the write index (received).
	if want, got := tr.Receives[tr.received:tr.received+len(p)], p; !bytes.Equal(want, got) {
		err = fmt.Errorf(
			"unexpected write to TestConn:\n  - want: %v\n  -  got: %v", want, got)
	}

	// Keep a record of the bytes we successfully received
	// and compared against the 'Receives' buffer.
	tr.received = tr.received + len(p)

	// Return the length of the 'written' slice
	return len(p), err
}

func TestWaitTimeout(t *testing.T) {

	returnChan := make(chan bool)

	// Return timer at 2 milliseconds
	returnTimer := func() {
		time.Sleep(time.Millisecond * 2)
		returnChan <- true
	}

	// Run test 10 times back-to-back
	for i := 0; i < 10; i++ {
		t.Run(t.Name(), func(t *testing.T) {

			// Start a timeout timer with a timeout higher than the return timer
			timeOutTimer := time.NewTimer(time.Millisecond * 3)
			go returnTimer() // Start return timer

			// Expect returnChan to unblock before timeOutTimer
			_, err := WaitTimeout(returnChan, timeOutTimer)
			if err != nil {
				t.Fatal("returnChan did not unblock before timeOutTimer")
			}

			// Start a timeout timer with a lower timeout than return timer
			timeOutTimer = time.NewTimer(time.Millisecond * 1)
			go returnTimer() // Start return timer

			// Expect returnChan to unblock before timeOutTimer
			_, err = WaitTimeout(returnChan, timeOutTimer)
			if err != errTimeout {
				t.Fatal("timeOutTimer did not expire before returnChan unblock")
			}

			// Wait for the return timer to send on channel
			// This value needs to be read to prevent it from interfering other tests
			<-returnChan
		})
	}
}

func TestReadPacket(t *testing.T) {

	rt := []struct {
		name string
		tc   TestConn
		pkt  Packet
		err  error
	}{
		{
			name: "single and double escaped 0x07s in payload",
			tc: TestConn{
				Emits: []byte{ // Response
					esc, ack,
					0x07, 0xF0, // Frame Start
					0x00, uint8(getFans + 1), // response type
					0x06,       // Length
					0xAA, 0xBB, // in/out percents
					0x11, 0x07, 0x07, // in speed (one seven)
					0x07, 0x07, 0x07, 0x07, // out speed (two sevens)
					0x4A, // Checksum
					esc, end,
					esc, ack,
				},
			},
			pkt: Packet{
				Command: uint8(getFans + 1),
				Data:    []byte{0xAA, 0xBB, 0x11, 0x7, 0x7, 0x7},
			},
		},
		{
			name: "timeout reading first response ACK",
			tc: TestConn{
				HangStart: true,
				HangTime:  2 * time.Millisecond,
			},
			err: errTimeout,
		},
		{
			name: "timeout reading start sequence",
			tc: TestConn{
				Emits:    []byte{esc, ack},
				HangEOF:  true,
				HangTime: 2 * time.Millisecond,
			},
			err: errTimeout,
		},
		{
			name: "timeout reading end sequence",
			tc: TestConn{
				Emits:    []byte{esc, ack, esc, start},
				HangEOF:  true,
				HangTime: 2 * time.Millisecond,
			},
			err: errTimeout,
		},
		{
			name: "garbage before ACK",
			tc: TestConn{
				Emits: []byte{esc, ack, 0xF0, 0x0B, esc, start},
			},
			err: errScanInput,
		},
		{
			name: "impossible packet length (need at least 4 bytes)",
			tc: TestConn{
				Emits: []byte{esc, ack, esc, start, 0xF0, 0x0B, esc, end},
			},
			err: errTooShort,
		},
		{
			name: "packet unmarshal error (any)",
			tc: TestConn{
				Emits: []byte{esc, ack, esc, start, 0xF0, 0x0B, 0xBA, 0xAA, esc, end},
			},
			err: errPayloadSize,
		},
	}

	for _, tt := range rt {
		t.Run(tt.name, func(t *testing.T) {

			// Copy the TestConn structure of this test cycle into
			// the goroutine to prevent the test engine from modifying
			// the structure before it is read by the library.
			tc := tt.tc

			// Start timeout timer
			to := time.NewTimer(time.Millisecond * 1)

			// Attempt to read packet from mock connection
			pkt, err := ReadPacket(&tc, to, true)

			if want, got := tt.err, err; want != got {
				t.Fatalf("unexpected error reading packet:\n- want: %v\n-  got: %v",
					want, got)
			}

			if want, got := tt.pkt, pkt; !reflect.DeepEqual(want, got) {
				t.Fatalf("unexpected packet read from connection:\n- want: %v\n-  got: %v",
					want, got)
			}
		})
	}
}
<commit_msg>libcomfo - TestConn: i/o limit for mock connection ReadWriter<commit_after>package libcomfo

import (
	"bytes"
	"fmt"
	"io"
	"reflect"
	"testing"
	"time"
)

// TestConn is a structure that implements io.ReadWriter and can be used
// to mock a connection with the unit. HangStart makes initial Read and Write
// calls hang. HangEOF makes Read hang on the second call when a Scanner
// would expect an EOF error to be returned.
// Not thread-safe.
type TestConn struct {
	Receives []byte
	Emits    []byte

	HangTime  time.Duration
	HangStart bool
	HangEOF   bool

	Limit int

	emitted  int
	received int
}

func (tr *TestConn) Read(p []byte) (n int, err error) {

	// Simulate read delay before returning output or EOF
	if tr.HangTime > 0 && tr.HangStart {
		time.Sleep(tr.HangTime)
	}

	// Return EOF when buffer is fully copied
	if tr.emitted == len(tr.Emits) {

		if tr.HangTime > 0 && tr.HangEOF {
			time.Sleep(tr.HangTime)
		}

		return 0, io.EOF
	}

	// Return the data in the connection's Emits field
	n = copy(p, tr.Emits)

	// Keep track of how many bytes we've been able to copy
	tr.emitted = tr.emitted + n

	// Return mocked read bytes total
	if tr.Limit > 0 {
		n = tr.Limit
	}

	return
}

func (tr *TestConn) Write(p []byte) (n int, err error) {

	// Simulate write delay
	if tr.HangTime > 0 {
		time.Sleep(tr.HangTime)
	}

	// Bounds check for Receives
	if tr.received+len(p) > len(tr.Receives) {
		return n,
			fmt.Errorf("write to TestConn exceeded expected input specified in Receives: %v",
				// This formula determines the offset of the slice
				// of surplus characters over the struct's 'Receives' field.
				p[len(p)-(tr.received+len(p)-len(tr.Receives)):])
	}

	// Expect the written data to correspond to what's in the 'Receives' field.
	// Offset the comparison against the write index (received).
	if want, got := tr.Receives[tr.received:tr.received+len(p)], p; !bytes.Equal(want, got) {
		err = fmt.Errorf(
			"unexpected write to TestConn:\n  - want: %v\n  -  got: %v", want, got)
	}

	// Keep a record of the bytes we successfully received
	// and compared against the 'Receives' buffer.
	tr.received = tr.received + len(p)

	// Return mocked written bytes total
	if tr.Limit > 0 {
		n = tr.Limit
	} else {
		// Return the length of the 'written' slice
		n = len(p)
	}

	return
}

func TestWaitTimeout(t *testing.T) {

	returnChan := make(chan bool)

	// Return timer at 2 milliseconds
	returnTimer := func() {
		time.Sleep(time.Millisecond * 2)
		returnChan <- true
	}

	// Run test 10 times back-to-back
	for i := 0; i < 10; i++ {
		t.Run(t.Name(), func(t *testing.T) {

			// Start a timeout timer with a timeout higher than the return timer
			timeOutTimer := time.NewTimer(time.Millisecond * 3)
			go returnTimer() // Start return timer

			// Expect returnChan to unblock before timeOutTimer
			_, err := WaitTimeout(returnChan, timeOutTimer)
			if err != nil {
				t.Fatal("returnChan did not unblock before timeOutTimer")
			}

			// Start a timeout timer with a lower timeout than return timer
			timeOutTimer = time.NewTimer(time.Millisecond * 1)
			go returnTimer() // Start return timer

			// Expect returnChan to unblock before timeOutTimer
			_, err = WaitTimeout(returnChan, timeOutTimer)
			if err != errTimeout {
				t.Fatal("timeOutTimer did not expire before returnChan unblock")
			}

			// Wait for the return timer to send on channel
			// This value needs to be read to prevent it from interfering other tests
			<-returnChan
		})
	}
}

func TestReadPacket(t *testing.T) {

	rt := []struct {
		name string
		tc   TestConn
		pkt  Packet
		err  error
	}{
		{
			name: "single and double escaped 0x07s in payload",
			tc: TestConn{
				Emits: []byte{ // Response
					esc, ack,
					0x07, 0xF0, // Frame Start
					0x00, uint8(getFans + 1), // response type
					0x06,       // Length
					0xAA, 0xBB, // in/out percents
					0x11, 0x07, 0x07, // in speed (one seven)
					0x07, 0x07, 0x07, 0x07, // out speed (two sevens)
					0x4A, // Checksum
					esc, end,
					esc, ack,
				},
			},
			pkt: Packet{
				Command: uint8(getFans + 1),
				Data:    []byte{0xAA, 0xBB, 0x11, 0x7, 0x7, 0x7},
			},
		},
		{
			name: "timeout reading first response ACK",
			tc: TestConn{
				HangStart: true,
				HangTime:  2 * time.Millisecond,
			},
			err: errTimeout,
		},
		{
			name: "timeout reading start sequence",
			tc: TestConn{
				Emits:    []byte{esc, ack},
				HangEOF:  true,
				HangTime: 2 * time.Millisecond,
			},
			err: errTimeout,
		},
		{
			name: "timeout reading end sequence",
			tc: TestConn{
				Emits:    []byte{esc, ack, esc, start},
				HangEOF:  true,
				HangTime: 2 * time.Millisecond,
			},
			err: errTimeout,
		},
		{
			name: "garbage before ACK",
			tc: TestConn{
				Emits: []byte{esc, ack, 0xF0, 0x0B, esc, start},
			},
			err: errScanInput,
		},
		{
			name: "impossible packet length (need at least 4 bytes)",
			tc: TestConn{
				Emits: []byte{esc, ack, esc, start, 0xF0, 0x0B, esc, end},
			},
			err: errTooShort,
		},
		{
			name: "packet unmarshal error (any)",
			tc: TestConn{
				Emits: []byte{esc, ack, esc, start, 0xF0, 0x0B, 0xBA, 0xAA, esc, end},
			},
			err: errPayloadSize,
		},
	}

	for _, tt := range rt {
		t.Run(tt.name, func(t *testing.T) {

			// Copy the TestConn structure of this test cycle into
			// the goroutine to prevent the test engine from modifying
			// the structure before it is read by the library.
			tc := tt.tc

			// Start timeout timer
			to := time.NewTimer(time.Millisecond * 1)

			// Attempt to read packet from mock connection
			pkt, err := ReadPacket(&tc, to, true)

			if want, got := tt.err, err; want != got {
				t.Fatalf("unexpected error reading packet:\n- want: %v\n-  got: %v",
					want, got)
			}

			if want, got := tt.pkt, pkt; !reflect.DeepEqual(want, got) {
				t.Fatalf("unexpected packet read from connection:\n- want: %v\n-  got: %v",
					want, got)
			}
		})
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2019-2021, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package grpcutils

import (
	"fmt"
	"math"
	"net"
	"net/http"
	"time"

	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/anypb"

	spb "google.golang.org/genproto/googleapis/rpc/status"

	tspb "google.golang.org/protobuf/types/known/timestamppb"

	httppb "github.com/ava-labs/avalanchego/proto/pb/http"
)

const (
	// Server:

	// MinTime is the minimum amount of time a client should wait before sending
	// a keepalive ping. grpc-go default 5 mins
	defaultServerKeepAliveMinTime = 5 * time.Second
	// After a duration of this time if the server doesn't see any activity it
	// pings the client to see if the transport is still alive.
	// If set below 1s, a minimum value of 1s will be used instead.
	// grpc-go default 2h
	defaultServerKeepAliveInterval = 2 * time.Hour
	// After having pinged for keepalive check, the server waits for a duration
	// of Timeout and if no activity is seen even after that the connection is
	// closed. grpc-go default 20s
	defaultServerKeepAliveTimeout = 20 * time.Second

	// Client:

	// After a duration of this time if the client doesn't see any activity it
	// pings the server to see if the transport is still alive.
	// If set below 10s, a minimum value of 10s will be used instead.
	// grpc-go default infinity
	defaultClientKeepAliveTime = 30 * time.Second
	// After having pinged for keepalive check, the client waits for a duration
	// of Timeout and if no activity is seen even after that the connection is
	// closed. grpc-go default 20s
	defaultClientKeepAliveTimeOut = 10 * time.Second
	// If true, client sends keepalive pings even with no active RPCs. If false,
	// when there are no active RPCs, Time and Timeout will be ignored and no
	// keepalive pings will be sent. grpc-go default false
	defaultPermitWithoutStream = true
)

var (
	DefaultDialOptions = []grpc.DialOption{
		grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(math.MaxInt)),
		grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(math.MaxInt)),
		grpc.WithKeepaliveParams(keepalive.ClientParameters{
			Time:                defaultClientKeepAliveTime,
			Timeout:             defaultClientKeepAliveTimeOut,
			PermitWithoutStream: defaultPermitWithoutStream,
		}),
	}

	DefaultServerOptions = []grpc.ServerOption{
		grpc.MaxRecvMsgSize(math.MaxInt),
		grpc.MaxSendMsgSize(math.MaxInt),
		grpc.MaxConcurrentStreams(math.MaxUint32),
		grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{
			MinTime:             defaultServerKeepAliveMinTime,
			PermitWithoutStream: defaultPermitWithoutStream,
		}),
		grpc.KeepaliveParams(keepalive.ServerParameters{
			Time:    defaultServerKeepAliveInterval,
			Timeout: defaultServerKeepAliveTimeout,
		}),
	}
)

// DialOptsWithMetrics registers gRPC client metrics via chain interceptors.
func DialOptsWithMetrics(clientMetrics *grpc_prometheus.ClientMetrics) []grpc.DialOption {
	return append(DefaultDialOptions,
		// Use chain interceptors to ensure custom/default interceptors are
		// applied correctly.
		// ref. https://github.com/kubernetes/kubernetes/pull/105069
		grpc.WithChainStreamInterceptor(clientMetrics.StreamClientInterceptor()),
		grpc.WithChainUnaryInterceptor(clientMetrics.UnaryClientInterceptor()),
	)
}

func Errorf(code int, tmpl string, args ...interface{}) error {
	return GetGRPCErrorFromHTTPResponse(&httppb.HandleSimpleHTTPResponse{
		Code: int32(code),
		Body: []byte(fmt.Sprintf(tmpl, args...)),
	})
}

// GetGRPCErrorFromHTTPRespone takes an HandleSimpleHTTPResponse as input and returns a gRPC error.
func GetGRPCErrorFromHTTPResponse(resp *httppb.HandleSimpleHTTPResponse) error {
	a, err := anypb.New(resp)
	if err != nil {
		return err
	}

	return status.ErrorProto(&spb.Status{
		Code:    resp.Code,
		Message: string(resp.Body),
		Details: []*anypb.Any{a},
	})
}

// GetHTTPResponseFromError takes an gRPC error as input and returns a gRPC
// HandleSimpleHTTPResponse.
func GetHTTPResponseFromError(err error) (*httppb.HandleSimpleHTTPResponse, bool) {
	s, ok := status.FromError(err)
	if !ok {
		return nil, false
	}

	status := s.Proto()
	if len(status.Details) != 1 {
		return nil, false
	}

	var resp httppb.HandleSimpleHTTPResponse
	if err := anypb.UnmarshalTo(status.Details[0], &resp, proto.UnmarshalOptions{}); err != nil {
		return nil, false
	}

	return &resp, true
}

// GetHTTPHeader takes an http.Header as input and returns a slice of Header.
func GetHTTPHeader(hs http.Header) []*httppb.Element {
	result := make([]*httppb.Element, 0, len(hs))
	for k, vs := range hs {
		result = append(result, &httppb.Element{
			Key:    k,
			Values: vs,
		})
	}
	return result
}

// MergeHTTPHeader takes a slice of Header and merges with http.Header map.
func MergeHTTPHeader(hs []*httppb.Element, header http.Header) {
	for _, h := range hs {
		header[h.Key] = h.Values
	}
}

func Serve(listener net.Listener, grpcServerFunc func([]grpc.ServerOption) *grpc.Server) {
	var opts []grpc.ServerOption
	grpcServer := grpcServerFunc(opts)

	// TODO: While errors will be reported later, it could be useful to somehow
	//       log this if it is the primary error.
	//
	// There is nothing to with the error returned by serve here. Later requests
	// will propegate their error if they occur.
	_ = grpcServer.Serve(listener)

	// Similarly, there is nothing to with an error when the listener is closed.
	_ = listener.Close()
}

func createClientConn(addr string, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
	opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials()))
	return grpc.Dial(addr, opts...)
}

// NewDefaultServer ensures the plugin service is served with proper
// defaults. This should always be passed to GRPCServer field of
// plugin.ServeConfig.
func NewDefaultServer(opts []grpc.ServerOption) *grpc.Server {
	if len(opts) == 0 {
		opts = append(opts, DefaultServerOptions...)
	}
	return grpc.NewServer(opts...)
}

// TimestampAsTime validates timestamppb timestamp and returns time.Time.
func TimestampAsTime(ts *tspb.Timestamp) (time.Time, error) {
	if err := ts.CheckValid(); err != nil {
		return time.Time{}, fmt.Errorf("invalid timestamp: %w", err)
	}
	return ts.AsTime(), nil
}

// TimestampFromTime converts time.Time to a timestamppb timestamp.
func TimestampFromTime(time time.Time) *tspb.Timestamp {
	return tspb.New(time)
}

// EnsureValidResponseCode ensures that the response code is valid otherwise it returns 500.
func EnsureValidResponseCode(code int) int {
	// Response code outside of this range is invalid and could panic.
	// ref. https://www.w3.org/Protocols/rfc2616/rfc2616-sec10.html
	if code < 100 || code > 599 {
		return http.StatusInternalServerError
	}
	return code
}
<commit_msg>Add default max connection lifetime config to gRPC servers (#1673)<commit_after>// Copyright (C) 2019-2021, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package grpcutils

import (
	"fmt"
	"math"
	"net"
	"net/http"
	"time"

	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/anypb"

	spb "google.golang.org/genproto/googleapis/rpc/status"

	tspb "google.golang.org/protobuf/types/known/timestamppb"

	httppb "github.com/ava-labs/avalanchego/proto/pb/http"
)

const (
	// Server:

	// MinTime is the minimum amount of time a client should wait before sending
	// a keepalive ping. grpc-go default 5 mins
	defaultServerKeepAliveMinTime = 5 * time.Second
	// After a duration of this time if the server doesn't see any activity it
	// pings the client to see if the transport is still alive.
	// If set below 1s, a minimum value of 1s will be used instead.
	// grpc-go default 2h
	defaultServerKeepAliveInterval = 2 * time.Hour
	// After having pinged for keepalive check, the server waits for a duration
	// of Timeout and if no activity is seen even after that the connection is
	// closed. grpc-go default 20s
	defaultServerKeepAliveTimeout = 20 * time.Second
	// Duration for the maximum amount of time a http2 connection can exist
	// before sending GOAWAY. Internally in gRPC a +-10% jitter is added to
	// mitigate retry storms.
	defaultServerMaxConnectionAge = 10 * time.Minute
	// Grace period after max defaultServerMaxConnectionAge after
	// which the http2 connection is closed. 1 second is the minimum possible
	// value. Anything less will be internally overridden to 1s by grpc.
	defaultServerMaxConnectionAgeGrace = 1 * time.Second

	// Client:

	// After a duration of this time if the client doesn't see any activity it
	// pings the server to see if the transport is still alive.
	// If set below 10s, a minimum value of 10s will be used instead.
	// grpc-go default infinity
	defaultClientKeepAliveTime = 30 * time.Second
	// After having pinged for keepalive check, the client waits for a duration
	// of Timeout and if no activity is seen even after that the connection is
	// closed. grpc-go default 20s
	defaultClientKeepAliveTimeOut = 10 * time.Second
	// If true, client sends keepalive pings even with no active RPCs. If false,
	// when there are no active RPCs, Time and Timeout will be ignored and no
	// keepalive pings will be sent. grpc-go default false
	defaultPermitWithoutStream = true
)

var (
	DefaultDialOptions = []grpc.DialOption{
		grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(math.MaxInt)),
		grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(math.MaxInt)),
		grpc.WithKeepaliveParams(keepalive.ClientParameters{
			Time:                defaultClientKeepAliveTime,
			Timeout:             defaultClientKeepAliveTimeOut,
			PermitWithoutStream: defaultPermitWithoutStream,
		}),
	}

	DefaultServerOptions = []grpc.ServerOption{
		grpc.MaxRecvMsgSize(math.MaxInt),
		grpc.MaxSendMsgSize(math.MaxInt),
		grpc.MaxConcurrentStreams(math.MaxUint32),
		grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{
			MinTime:             defaultServerKeepAliveMinTime,
			PermitWithoutStream: defaultPermitWithoutStream,
		}),
		grpc.KeepaliveParams(keepalive.ServerParameters{
			Time:                  defaultServerKeepAliveInterval,
			Timeout:               defaultServerKeepAliveTimeout,
			MaxConnectionAge:      defaultServerMaxConnectionAge,
			MaxConnectionAgeGrace: defaultServerMaxConnectionAgeGrace,
		}),
	}
)

// DialOptsWithMetrics registers gRPC client metrics via chain interceptors.
func DialOptsWithMetrics(clientMetrics *grpc_prometheus.ClientMetrics) []grpc.DialOption {
	return append(DefaultDialOptions,
		// Use chain interceptors to ensure custom/default interceptors are
		// applied correctly.
		// ref. https://github.com/kubernetes/kubernetes/pull/105069
		grpc.WithChainStreamInterceptor(clientMetrics.StreamClientInterceptor()),
		grpc.WithChainUnaryInterceptor(clientMetrics.UnaryClientInterceptor()),
	)
}

func Errorf(code int, tmpl string, args ...interface{}) error {
	return GetGRPCErrorFromHTTPResponse(&httppb.HandleSimpleHTTPResponse{
		Code: int32(code),
		Body: []byte(fmt.Sprintf(tmpl, args...)),
	})
}

// GetGRPCErrorFromHTTPRespone takes an HandleSimpleHTTPResponse as input and returns a gRPC error.
func GetGRPCErrorFromHTTPResponse(resp *httppb.HandleSimpleHTTPResponse) error {
	a, err := anypb.New(resp)
	if err != nil {
		return err
	}

	return status.ErrorProto(&spb.Status{
		Code:    resp.Code,
		Message: string(resp.Body),
		Details: []*anypb.Any{a},
	})
}

// GetHTTPResponseFromError takes an gRPC error as input and returns a gRPC
// HandleSimpleHTTPResponse.
func GetHTTPResponseFromError(err error) (*httppb.HandleSimpleHTTPResponse, bool) {
	s, ok := status.FromError(err)
	if !ok {
		return nil, false
	}

	status := s.Proto()
	if len(status.Details) != 1 {
		return nil, false
	}

	var resp httppb.HandleSimpleHTTPResponse
	if err := anypb.UnmarshalTo(status.Details[0], &resp, proto.UnmarshalOptions{}); err != nil {
		return nil, false
	}

	return &resp, true
}

// GetHTTPHeader takes an http.Header as input and returns a slice of Header.
func GetHTTPHeader(hs http.Header) []*httppb.Element {
	result := make([]*httppb.Element, 0, len(hs))
	for k, vs := range hs {
		result = append(result, &httppb.Element{
			Key:    k,
			Values: vs,
		})
	}
	return result
}

// MergeHTTPHeader takes a slice of Header and merges with http.Header map.
func MergeHTTPHeader(hs []*httppb.Element, header http.Header) {
	for _, h := range hs {
		header[h.Key] = h.Values
	}
}

func Serve(listener net.Listener, grpcServerFunc func([]grpc.ServerOption) *grpc.Server) {
	var opts []grpc.ServerOption
	grpcServer := grpcServerFunc(opts)

	// TODO: While errors will be reported later, it could be useful to somehow
	//       log this if it is the primary error.
	//
	// There is nothing to with the error returned by serve here. Later requests
	// will propegate their error if they occur.
	_ = grpcServer.Serve(listener)

	// Similarly, there is nothing to with an error when the listener is closed.
	_ = listener.Close()
}

func createClientConn(addr string, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
	opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials()))
	return grpc.Dial(addr, opts...)
}

// NewDefaultServer ensures the plugin service is served with proper
// defaults. This should always be passed to GRPCServer field of
// plugin.ServeConfig.
func NewDefaultServer(opts []grpc.ServerOption) *grpc.Server {
	if len(opts) == 0 {
		opts = append(opts, DefaultServerOptions...)
	}
	return grpc.NewServer(opts...)
}

// TimestampAsTime validates timestamppb timestamp and returns time.Time.
func TimestampAsTime(ts *tspb.Timestamp) (time.Time, error) {
	if err := ts.CheckValid(); err != nil {
		return time.Time{}, fmt.Errorf("invalid timestamp: %w", err)
	}
	return ts.AsTime(), nil
}

// TimestampFromTime converts time.Time to a timestamppb timestamp.
func TimestampFromTime(time time.Time) *tspb.Timestamp {
	return tspb.New(time)
}

// EnsureValidResponseCode ensures that the response code is valid otherwise it returns 500.
func EnsureValidResponseCode(code int) int {
	// Response code outside of this range is invalid and could panic.
	// ref. https://www.w3.org/Protocols/rfc2616/rfc2616-sec10.html
	if code < 100 || code > 599 {
		return http.StatusInternalServerError
	}
	return code
}
<|endoftext|>


<commit_before>// Copyright 2014 The Cayley 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 graph

// Define the general iterator interface.

import (
	"strings"
	"sync"

	"github.com/barakmich/glog"
)

type Tagger struct {
	tags      []string
	fixedTags map[string]Value
}

// Adds a tag to the iterator.
func (t *Tagger) Add(tag string) {
	t.tags = append(t.tags, tag)
}

func (t *Tagger) AddFixed(tag string, value Value) {
	if t.fixedTags == nil {
		t.fixedTags = make(map[string]Value)
	}
	t.fixedTags[tag] = value
}

// Returns the tags. The returned value must not be mutated.
func (t *Tagger) Tags() []string {
	return t.tags
}

// Returns the fixed tags. The returned value must not be mutated.
func (t *Tagger) Fixed() map[string]Value {
	return t.fixedTags
}

func (t *Tagger) CopyFrom(src Iterator) {
	for _, tag := range src.Tagger().Tags() {
		t.Add(tag)
	}

	for k, v := range src.Tagger().Fixed() {
		t.AddFixed(k, v)
	}

}

type Iterator interface {
	Tagger() *Tagger

	// Fills a tag-to-result-value map.
	TagResults(map[string]Value)

	// Returns the current result.
	Result() Value

	// DEPRECATED -- Fills a ResultTree struct with Result().
	ResultTree() *ResultTree

	// These methods are the heart and soul of the iterator, as they constitute
	// the iteration interface.
	//
	// To get the full results of iteration, do the following:
	//
	//  for graph.Next(it) {
	//  	val := it.Result()
	//  	... do things with val.
	//  	for it.NextPath() {
	//  		... find other paths to iterate
	//  	}
	//  }
	//
	// All of them should set iterator.Last to be the last returned value, to
	// make results work.
	//
	// NextPath() advances iterators that may have more than one valid result,
	// from the bottom up.
	NextPath() bool

	// Contains returns whether the value is within the set held by the iterator.
	Contains(Value) bool

	// Start iteration from the beginning
	Reset()

	// Create a new iterator just like this one
	Clone() Iterator

	// These methods relate to choosing the right iterator, or optimizing an
	// iterator tree
	//
	// Stats() returns the relative costs of calling the iteration methods for
	// this iterator, as well as the size. Roughly, it will take NextCost * Size
	// "cost units" to get everything out of the iterator. This is a wibbly-wobbly
	// thing, and not exact, but a useful heuristic.
	Stats() IteratorStats

	// Helpful accessor for the number of things in the iterator. The first return
	// value is the size, and the second return value is whether that number is exact,
	// or a conservative estimate.
	Size() (int64, bool)

	// Returns a string relating to what the function of the iterator is. By
	// knowing the names of the iterators, we can devise optimization strategies.
	Type() Type

	// Optimizes an iterator. Can replace the iterator, or merely move things
	// around internally. if it chooses to replace it with a better iterator,
	// returns (the new iterator, true), if not, it returns (self, false).
	Optimize() (Iterator, bool)

	// Return a slice of the subiterators for this iterator.
	SubIterators() []Iterator

	// Return a string representation of the iterator, indented by the given amount.
	DebugString(int) string

	// Close the iterator and do internal cleanup.
	Close()

	// UID returns the unique identifier of the iterator.
	UID() uint64
}

type Nexter interface {
	// Next advances the iterator to the next value, which will then be available through
	// the Result method. It returns false if no further advancement is possible.
	Next() bool

	Iterator
}

// Next is a convenience function that conditionally calls the Next method
// of an Iterator if it is a Nexter. If the Iterator is not a Nexter, Next
// returns false.
func Next(it Iterator) bool {
	if n, ok := it.(Nexter); ok {
		return n.Next()
	}
	glog.Errorln("Nexting an un-nextable iterator")
	return false
}

// Height is a convienence function to measure the height of an iterator tree.
func Height(it Iterator, until Type) int {
	if it.Type() == until {
		return 1
	}
	subs := it.SubIterators()
	maxDepth := 0
	for _, sub := range subs {
		h := Height(sub, until)
		if h > maxDepth {
			maxDepth = h
		}
	}
	return maxDepth + 1
}

// FixedIterator wraps iterators that are modifiable by addition of fixed value sets.
type FixedIterator interface {
	Iterator
	Add(Value)
}

type IteratorStats struct {
	ContainsCost int64
	NextCost     int64
	Size         int64
	Next         int64
	Contains     int64
	ContainsNext int64
}

// Type enumerates the set of Iterator types.
type Type int

const (
	Invalid Type = iota
	All
	And
	Or
	HasA
	LinksTo
	Comparison
	Null
	Fixed
	Not
	Optional
	Materialize
)

var (
	// We use a sync.Mutex rather than an RWMutex since the client packages keep
	// the Type that was returned, so the only possibility for contention is at
	// initialization.
	lock sync.Mutex
	// These strings must be kept in order consistent with the Type const block above.
	types = []string{
		"invalid",
		"all",
		"and",
		"or",
		"hasa",
		"linksto",
		"comparison",
		"null",
		"fixed",
		"not",
		"optional",
		"materialize",
	}
)

// RegisterIterator adds a new iterator type to the set of acceptable types, returning
// the registered Type.
// Calls to Register are idempotent and must be made prior to use of the iterator.
// The conventional approach for use is to include a call to Register in a package
// init() function, saving the Type to a private package var.
func RegisterIterator(name string) Type {
	lock.Lock()
	defer lock.Unlock()
	for i, t := range types {
		if t == name {
			return Type(i)
		}
	}
	types = append(types, name)
	return Type(len(types) - 1)
}

// String returns a string representation of the Type.
func (t Type) String() string {
	if t < 0 || int(t) >= len(types) {
		return "illegal-type"
	}
	return types[t]
}

type StatsContainer struct {
	IteratorStats
	Kind   string
	Uid    uint64
	SubIts []StatsContainer
}

func DumpStats(it Iterator) StatsContainer {
	var out StatsContainer
	out.IteratorStats = it.Stats()
	out.Kind = it.Type().String()
	out.Uid = it.UID()
	for _, sub := range it.SubIterators() {
		out.SubIts = append(out.SubIts, DumpStats(sub))
	}
	return out
}

// Utility logging functions for when an iterator gets called Next upon, or Contains upon, as
// well as what they return. Highly useful for tracing the execution path of a query.
func ContainsLogIn(it Iterator, val Value) {
	if glog.V(4) {
		glog.V(4).Infof("%s %d CHECK CONTAINS %d", strings.ToUpper(it.Type().String()), it.UID(), val)
	}
}

func ContainsLogOut(it Iterator, val Value, good bool) bool {
	if glog.V(4) {
		if good {
			glog.V(4).Infof("%s %d CHECK CONTAINS %d GOOD", strings.ToUpper(it.Type().String()), it.UID(), val)
		} else {
			glog.V(4).Infof("%s %d CHECK CONTAINS %d BAD", strings.ToUpper(it.Type().String()), it.UID(), val)
		}
	}
	return good
}

func NextLogIn(it Iterator) {
	if glog.V(4) {
		glog.V(4).Infof("%s %d NEXT", strings.ToUpper(it.Type().String()), it.UID())
	}
}

func NextLogOut(it Iterator, val Value, ok bool) bool {
	if glog.V(4) {
		if ok {
			glog.V(4).Infof("%s %d NEXT IS %d", strings.ToUpper(it.Type().String()), it.UID(), val)
		} else {
			glog.V(4).Infof("%s %d NEXT DONE", strings.ToUpper(it.Type().String()), it.UID())
		}
	}
	return ok
}
<commit_msg>Do tagger copying with less iteration<commit_after>// Copyright 2014 The Cayley 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 graph

// Define the general iterator interface.

import (
	"strings"
	"sync"

	"github.com/barakmich/glog"
)

type Tagger struct {
	tags      []string
	fixedTags map[string]Value
}

// Adds a tag to the iterator.
func (t *Tagger) Add(tag string) {
	t.tags = append(t.tags, tag)
}

func (t *Tagger) AddFixed(tag string, value Value) {
	if t.fixedTags == nil {
		t.fixedTags = make(map[string]Value)
	}
	t.fixedTags[tag] = value
}

// Returns the tags. The returned value must not be mutated.
func (t *Tagger) Tags() []string {
	return t.tags
}

// Returns the fixed tags. The returned value must not be mutated.
func (t *Tagger) Fixed() map[string]Value {
	return t.fixedTags
}

func (t *Tagger) CopyFrom(src Iterator) {
	st := src.Tagger()

	t.tags = append(t.tags, st.tags...)

	if t.fixedTags == nil {
		t.fixedTags = make(map[string]Value, len(st.fixedTags))
	}
	for k, v := range st.fixedTags {
		t.fixedTags[k] = v
	}
}

type Iterator interface {
	Tagger() *Tagger

	// Fills a tag-to-result-value map.
	TagResults(map[string]Value)

	// Returns the current result.
	Result() Value

	// DEPRECATED -- Fills a ResultTree struct with Result().
	ResultTree() *ResultTree

	// These methods are the heart and soul of the iterator, as they constitute
	// the iteration interface.
	//
	// To get the full results of iteration, do the following:
	//
	//  for graph.Next(it) {
	//  	val := it.Result()
	//  	... do things with val.
	//  	for it.NextPath() {
	//  		... find other paths to iterate
	//  	}
	//  }
	//
	// All of them should set iterator.Last to be the last returned value, to
	// make results work.
	//
	// NextPath() advances iterators that may have more than one valid result,
	// from the bottom up.
	NextPath() bool

	// Contains returns whether the value is within the set held by the iterator.
	Contains(Value) bool

	// Start iteration from the beginning
	Reset()

	// Create a new iterator just like this one
	Clone() Iterator

	// These methods relate to choosing the right iterator, or optimizing an
	// iterator tree
	//
	// Stats() returns the relative costs of calling the iteration methods for
	// this iterator, as well as the size. Roughly, it will take NextCost * Size
	// "cost units" to get everything out of the iterator. This is a wibbly-wobbly
	// thing, and not exact, but a useful heuristic.
	Stats() IteratorStats

	// Helpful accessor for the number of things in the iterator. The first return
	// value is the size, and the second return value is whether that number is exact,
	// or a conservative estimate.
	Size() (int64, bool)

	// Returns a string relating to what the function of the iterator is. By
	// knowing the names of the iterators, we can devise optimization strategies.
	Type() Type

	// Optimizes an iterator. Can replace the iterator, or merely move things
	// around internally. if it chooses to replace it with a better iterator,
	// returns (the new iterator, true), if not, it returns (self, false).
	Optimize() (Iterator, bool)

	// Return a slice of the subiterators for this iterator.
	SubIterators() []Iterator

	// Return a string representation of the iterator, indented by the given amount.
	DebugString(int) string

	// Close the iterator and do internal cleanup.
	Close()

	// UID returns the unique identifier of the iterator.
	UID() uint64
}

type Nexter interface {
	// Next advances the iterator to the next value, which will then be available through
	// the Result method. It returns false if no further advancement is possible.
	Next() bool

	Iterator
}

// Next is a convenience function that conditionally calls the Next method
// of an Iterator if it is a Nexter. If the Iterator is not a Nexter, Next
// returns false.
func Next(it Iterator) bool {
	if n, ok := it.(Nexter); ok {
		return n.Next()
	}
	glog.Errorln("Nexting an un-nextable iterator")
	return false
}

// Height is a convienence function to measure the height of an iterator tree.
func Height(it Iterator, until Type) int {
	if it.Type() == until {
		return 1
	}
	subs := it.SubIterators()
	maxDepth := 0
	for _, sub := range subs {
		h := Height(sub, until)
		if h > maxDepth {
			maxDepth = h
		}
	}
	return maxDepth + 1
}

// FixedIterator wraps iterators that are modifiable by addition of fixed value sets.
type FixedIterator interface {
	Iterator
	Add(Value)
}

type IteratorStats struct {
	ContainsCost int64
	NextCost     int64
	Size         int64
	Next         int64
	Contains     int64
	ContainsNext int64
}

// Type enumerates the set of Iterator types.
type Type int

const (
	Invalid Type = iota
	All
	And
	Or
	HasA
	LinksTo
	Comparison
	Null
	Fixed
	Not
	Optional
	Materialize
)

var (
	// We use a sync.Mutex rather than an RWMutex since the client packages keep
	// the Type that was returned, so the only possibility for contention is at
	// initialization.
	lock sync.Mutex
	// These strings must be kept in order consistent with the Type const block above.
	types = []string{
		"invalid",
		"all",
		"and",
		"or",
		"hasa",
		"linksto",
		"comparison",
		"null",
		"fixed",
		"not",
		"optional",
		"materialize",
	}
)

// RegisterIterator adds a new iterator type to the set of acceptable types, returning
// the registered Type.
// Calls to Register are idempotent and must be made prior to use of the iterator.
// The conventional approach for use is to include a call to Register in a package
// init() function, saving the Type to a private package var.
func RegisterIterator(name string) Type {
	lock.Lock()
	defer lock.Unlock()
	for i, t := range types {
		if t == name {
			return Type(i)
		}
	}
	types = append(types, name)
	return Type(len(types) - 1)
}

// String returns a string representation of the Type.
func (t Type) String() string {
	if t < 0 || int(t) >= len(types) {
		return "illegal-type"
	}
	return types[t]
}

type StatsContainer struct {
	IteratorStats
	Kind   string
	Uid    uint64
	SubIts []StatsContainer
}

func DumpStats(it Iterator) StatsContainer {
	var out StatsContainer
	out.IteratorStats = it.Stats()
	out.Kind = it.Type().String()
	out.Uid = it.UID()
	for _, sub := range it.SubIterators() {
		out.SubIts = append(out.SubIts, DumpStats(sub))
	}
	return out
}

// Utility logging functions for when an iterator gets called Next upon, or Contains upon, as
// well as what they return. Highly useful for tracing the execution path of a query.
func ContainsLogIn(it Iterator, val Value) {
	if glog.V(4) {
		glog.V(4).Infof("%s %d CHECK CONTAINS %d", strings.ToUpper(it.Type().String()), it.UID(), val)
	}
}

func ContainsLogOut(it Iterator, val Value, good bool) bool {
	if glog.V(4) {
		if good {
			glog.V(4).Infof("%s %d CHECK CONTAINS %d GOOD", strings.ToUpper(it.Type().String()), it.UID(), val)
		} else {
			glog.V(4).Infof("%s %d CHECK CONTAINS %d BAD", strings.ToUpper(it.Type().String()), it.UID(), val)
		}
	}
	return good
}

func NextLogIn(it Iterator) {
	if glog.V(4) {
		glog.V(4).Infof("%s %d NEXT", strings.ToUpper(it.Type().String()), it.UID())
	}
}

func NextLogOut(it Iterator, val Value, ok bool) bool {
	if glog.V(4) {
		if ok {
			glog.V(4).Infof("%s %d NEXT IS %d", strings.ToUpper(it.Type().String()), it.UID(), val)
		} else {
			glog.V(4).Infof("%s %d NEXT DONE", strings.ToUpper(it.Type().String()), it.UID())
		}
	}
	return ok
}
<|endoftext|>


<commit_before>package git_pipeline_test

import (
	"fmt"
	"net/http"
	"os/exec"
	"strings"

	"github.com/cloudfoundry-incubator/garden"
	"github.com/cloudfoundry-incubator/garden/client"
	"github.com/mgutz/ansi"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
	"github.com/pivotal-golang/lager/lagertest"

	gconn "github.com/cloudfoundry-incubator/garden/client/connection"

	"testing"
	"time"
)

// has ruby, curl
const guidServerRootfs = "/var/vcap/packages/bosh_deployment_resource"

// has git, curl
const gitServerRootfs = "/var/vcap/packages/git_resource"

var flyBin string

var (
	gardenClient garden.Client

	atcURL string

	pipelineName string
)

var _ = SynchronizedBeforeSuite(func() []byte {
	flyBinPath, err := gexec.Build("github.com/concourse/fly", "-race")
	Expect(err).NotTo(HaveOccurred())

	return []byte(flyBinPath)
}, func(flyBinPath []byte) {
	flyBin = string(flyBinPath)

	// observed jobs taking ~1m30s, so set the timeout pretty high
	SetDefaultEventuallyTimeout(5 * time.Minute)

	// poll less frequently
	SetDefaultEventuallyPollingInterval(time.Second)

	logger := lagertest.NewTestLogger("testflight")

	gardenClient = client.New(gconn.NewWithLogger("tcp", "10.244.15.2:7777", logger.Session("garden-connection")))
	Eventually(gardenClient.Ping).ShouldNot(HaveOccurred())

	atcURL = "http://10.244.15.2:8080"

	Eventually(errorPolling(atcURL)).ShouldNot(HaveOccurred())

	pipelineName = fmt.Sprintf("test-pipeline-%d", GinkgoParallelNode())
})

func TestGitPipeline(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Pipelines Suite")
}

func destroyPipeline() {
	destroyCmd := exec.Command(
		flyBin,
		"-t", atcURL,
		"destroy-pipeline",
		"-p", pipelineName,
	)

	stdin, err := destroyCmd.StdinPipe()
	Expect(err).NotTo(HaveOccurred())

	defer stdin.Close()

	destroy, err := gexec.Start(destroyCmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())

	Eventually(destroy).Should(gbytes.Say("are you sure?"))

	fmt.Fprintln(stdin, "y")

	<-destroy.Exited

	if destroy.ExitCode() == 1 {
		if strings.Contains(string(destroy.Err.Contents()), "does not exist") {
			return
		}
	}

	Expect(destroy).To(gexec.Exit(0))
}

func configurePipeline(argv ...string) {
	destroyPipeline()

	args := append([]string{
		"-t", atcURL,
		"set-config",
		"-p", pipelineName,
		"--paused", "false",
	}, argv...)

	configureCmd := exec.Command(flyBin, args...)

	stdin, err := configureCmd.StdinPipe()
	Expect(err).NotTo(HaveOccurred())

	defer stdin.Close()

	configure, err := gexec.Start(configureCmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())

	Eventually(configure).Should(gbytes.Say("apply configuration?"))

	fmt.Fprintln(stdin, "y")

	Eventually(configure).Should(gexec.Exit(0))
}

func errorPolling(url string) func() error {
	return func() error {
		resp, err := http.Get(url)
		if err == nil {
			resp.Body.Close()
		}

		return err
	}
}

func flyWatch(jobName string, buildName ...string) *gexec.Session {
	args := []string{
		"-t", atcURL,
		"watch",
		"-j", pipelineName + "/" + jobName,
	}

	if len(buildName) > 0 {
		args = append(args, "-b", buildName[0])
	}

	for {
		session := start(exec.Command(flyBin, args...))

		<-session.Exited

		if session.ExitCode() == 1 {
			output := strings.TrimSpace(string(session.Err.Contents()))
			if output == "job has no builds" || output == "build not found" {
				// build hasn't started yet; keep polling
				time.Sleep(time.Second)
				continue
			}
		}

		return session
	}
}

func start(cmd *exec.Cmd) *gexec.Session {
	session, err := gexec.Start(
		cmd,
		gexec.NewPrefixedWriter(
			fmt.Sprintf("%s%s ", ansi.Color("[o]", "green"), ansi.Color("[fly]", "blue")),
			GinkgoWriter,
		),
		gexec.NewPrefixedWriter(
			fmt.Sprintf("%s%s ", ansi.Color("[e]", "red+bright"), ansi.Color("[fly]", "blue")),
			GinkgoWriter,
		),
	)
	Expect(err).NotTo(HaveOccurred())

	return session
}
<commit_msg>flyWatch should be a little more liberal when scanning fly output<commit_after>package git_pipeline_test

import (
	"fmt"
	"net/http"
	"os/exec"
	"regexp"
	"strings"

	"github.com/cloudfoundry-incubator/garden"
	"github.com/cloudfoundry-incubator/garden/client"
	"github.com/mgutz/ansi"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
	"github.com/pivotal-golang/lager/lagertest"

	gconn "github.com/cloudfoundry-incubator/garden/client/connection"

	"testing"
	"time"
)

// has ruby, curl
const guidServerRootfs = "/var/vcap/packages/bosh_deployment_resource"

// has git, curl
const gitServerRootfs = "/var/vcap/packages/git_resource"

var flyBin string

var (
	gardenClient garden.Client

	atcURL string

	pipelineName string
)

var _ = SynchronizedBeforeSuite(func() []byte {
	flyBinPath, err := gexec.Build("github.com/concourse/fly", "-race")
	Expect(err).NotTo(HaveOccurred())

	return []byte(flyBinPath)
}, func(flyBinPath []byte) {
	flyBin = string(flyBinPath)

	// observed jobs taking ~1m30s, so set the timeout pretty high
	SetDefaultEventuallyTimeout(5 * time.Minute)

	// poll less frequently
	SetDefaultEventuallyPollingInterval(time.Second)

	logger := lagertest.NewTestLogger("testflight")

	gardenClient = client.New(gconn.NewWithLogger("tcp", "10.244.15.2:7777", logger.Session("garden-connection")))
	Eventually(gardenClient.Ping).ShouldNot(HaveOccurred())

	atcURL = "http://10.244.15.2:8080"

	Eventually(errorPolling(atcURL)).ShouldNot(HaveOccurred())

	pipelineName = fmt.Sprintf("test-pipeline-%d", GinkgoParallelNode())
})

func TestGitPipeline(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Pipelines Suite")
}

func destroyPipeline() {
	destroyCmd := exec.Command(
		flyBin,
		"-t", atcURL,
		"destroy-pipeline",
		"-p", pipelineName,
	)

	stdin, err := destroyCmd.StdinPipe()
	Expect(err).NotTo(HaveOccurred())

	defer stdin.Close()

	destroy, err := gexec.Start(destroyCmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())

	Eventually(destroy).Should(gbytes.Say("are you sure?"))

	fmt.Fprintln(stdin, "y")

	<-destroy.Exited

	if destroy.ExitCode() == 1 {
		if strings.Contains(string(destroy.Err.Contents()), "does not exist") {
			return
		}
	}

	Expect(destroy).To(gexec.Exit(0))
}

func configurePipeline(argv ...string) {
	destroyPipeline()

	args := append([]string{
		"-t", atcURL,
		"set-config",
		"-p", pipelineName,
		"--paused", "false",
	}, argv...)

	configureCmd := exec.Command(flyBin, args...)

	stdin, err := configureCmd.StdinPipe()
	Expect(err).NotTo(HaveOccurred())

	defer stdin.Close()

	configure, err := gexec.Start(configureCmd, GinkgoWriter, GinkgoWriter)
	Expect(err).NotTo(HaveOccurred())

	Eventually(configure).Should(gbytes.Say("apply configuration?"))

	fmt.Fprintln(stdin, "y")

	Eventually(configure).Should(gexec.Exit(0))
}

func errorPolling(url string) func() error {
	return func() error {
		resp, err := http.Get(url)
		if err == nil {
			resp.Body.Close()
		}

		return err
	}
}

func flyWatch(jobName string, buildName ...string) *gexec.Session {
	args := []string{
		"-t", atcURL,
		"watch",
		"-j", pipelineName + "/" + jobName,
	}

	if len(buildName) > 0 {
		args = append(args, "-b", buildName[0])
	}

	keepPollingCheck := regexp.MustCompile("job has no builds|build not found|failed to get build")
	for {
		session := start(exec.Command(flyBin, args...))

		<-session.Exited

		if session.ExitCode() == 1 {
			output := strings.TrimSpace(string(session.Err.Contents()))
			if keepPollingCheck.MatchString(output) {
				// build hasn't started yet; keep polling
				time.Sleep(time.Second)
				continue
			}
		}

		return session
	}
}

func start(cmd *exec.Cmd) *gexec.Session {
	session, err := gexec.Start(
		cmd,
		gexec.NewPrefixedWriter(
			fmt.Sprintf("%s%s ", ansi.Color("[o]", "green"), ansi.Color("[fly]", "blue")),
			GinkgoWriter,
		),
		gexec.NewPrefixedWriter(
			fmt.Sprintf("%s%s ", ansi.Color("[e]", "red+bright"), ansi.Color("[fly]", "blue")),
			GinkgoWriter,
		),
	)
	Expect(err).NotTo(HaveOccurred())

	return session
}
<|endoftext|>