<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/printer"
	"go/token"
	"go/types"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"golang.org/x/tools/go/ast/astutil"
	"golang.org/x/tools/go/buildutil"
	"golang.org/x/tools/go/loader"
)

func parseOctothorpDecimal(s string) int {
	if s != "" && s[0] == '#' {
		if s, err := strconv.ParseInt(s[1:], 10, 32); err == nil {
			return int(s)
		}
	}
	return -1
}

func parsePos(pos string) (filename string, startOffset, endOffset int, err error) {
	if pos == "" {
		err = fmt.Errorf("no source position specified")
		return
	}

	colon := strings.LastIndex(pos, ":")
	if colon < 0 {
		err = fmt.Errorf("bad position syntax %q", pos)
		return
	}
	filename, offset := pos[:colon], pos[colon+1:]
	startOffset = -1
	endOffset = -1
	if hyphen := strings.Index(offset, ","); hyphen < 0 {
		// e.g. "foo.go:#123"
		startOffset = parseOctothorpDecimal(offset)
		endOffset = startOffset
	} else {
		// e.g. "foo.go:#123,#456"
		startOffset = parseOctothorpDecimal(offset[:hyphen])
		endOffset = parseOctothorpDecimal(offset[hyphen+1:])
	}
	if startOffset < 0 || endOffset < 0 {
		err = fmt.Errorf("invalid offset %q in query position", offset)
		return
	}
	return
}

func fileOffsetToPos(file *token.File, startOffset, endOffset int) (start, end token.Pos, err error) {
	// Range check [start..end], inclusive of both end-points.

	if 0 <= startOffset && startOffset <= file.Size() {
		start = file.Pos(int(startOffset))
	} else {
		err = fmt.Errorf("start position is beyond end of file")
		return
	}

	if 0 <= endOffset && endOffset <= file.Size() {
		end = file.Pos(int(endOffset))
	} else {
		err = fmt.Errorf("end position is beyond end of file")
		return
	}

	return
}

var (
	fRecursive bool
	fOneLine   bool
	fJSON      bool
)

func init() {
	flag.BoolVar(&fRecursive, "r", false, "keyify struct initializers recursively")
	flag.BoolVar(&fOneLine, "o", false, "print new struct initializer on a single line")
	flag.BoolVar(&fJSON, "json", false, "print new struct initializer as JSON")
}

func usage() {
	fmt.Printf("Usage: %s [flags] <position>\n\n", os.Args[0])
	flag.PrintDefaults()
}

func main() {
	log.SetFlags(0)
	flag.Usage = usage
	flag.Parse()
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(2)
	}
	pos := flag.Args()[0]
	name, start, _, err := parsePos(pos)
	if err != nil {
		log.Fatal(err)
	}
	eval, err := filepath.EvalSymlinks(name)
	if err != nil {
		log.Fatal(err)
	}
	name, err = filepath.Abs(eval)
	if err != nil {
		log.Fatal(err)
	}
	cwd, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	bpkg, err := buildutil.ContainingPackage(&build.Default, cwd, name)
	if err != nil {
		log.Fatal(err)
	}
	conf := &loader.Config{}
	conf.TypeCheckFuncBodies = func(s string) bool {
		return s == bpkg.ImportPath || s == bpkg.ImportPath+"_test"
	}
	conf.ImportWithTests(bpkg.ImportPath)
	lprog, err := conf.Load()
	if err != nil {
		log.Fatal(err)
	}
	var tf *token.File
	var af *ast.File
	pkg := lprog.InitialPackages()[0]
	for _, ff := range pkg.Files {
		file := lprog.Fset.File(ff.Pos())
		if file.Name() == name {
			af = ff
			tf = file
			break
		}
	}
	tstart, tend, err := fileOffsetToPos(tf, start, start)
	if err != nil {
		log.Fatal(err)
	}
	path, _ := astutil.PathEnclosingInterval(af, tstart, tend)
	var complit *ast.CompositeLit
	for _, p := range path {
		if p, ok := p.(*ast.CompositeLit); ok {
			complit = p
			break
		}
	}
	if complit == nil {
		log.Fatal("no composite literal found near point")
	}
	if len(complit.Elts) == 0 {
		printComplit(complit, complit, lprog.Fset, lprog.Fset)
		return
	}
	if _, ok := complit.Elts[0].(*ast.KeyValueExpr); ok {
		printComplit(complit, complit, lprog.Fset, lprog.Fset)
		return
	}
	st, ok := pkg.TypeOf(complit.Type).Underlying().(*types.Struct)
	if !ok {
		log.Fatal("not a struct initialiser")
		return
	}

	newFset := token.NewFileSet()
	newFile := newFset.AddFile("", -1, st.NumFields()+2)
	newComplit := &ast.CompositeLit{
		Type:   complit.Type,
		Lbrace: 1,
		Rbrace: token.Pos(st.NumFields() + 2),
	}
	newFile.AddLine(1)
	newFile.AddLine(st.NumFields() + 2)
	for i := 0; i < st.NumFields(); i++ {
		newFile.AddLine(2 + i)
		field := st.Field(i)
		elt := &ast.KeyValueExpr{
			Key:   &ast.Ident{NamePos: token.Pos(2 + i), Name: field.Name()},
			Value: copyExpr(complit.Elts[i], token.Pos(2+i)),
		}
		newComplit.Elts = append(newComplit.Elts, elt)
	}
	printComplit(complit, newComplit, lprog.Fset, newFset)
}

func printComplit(oldlit, newlit *ast.CompositeLit, oldfset, newfset *token.FileSet) {
	buf := &bytes.Buffer{}
	cfg := printer.Config{Mode: printer.UseSpaces | printer.TabIndent, Tabwidth: 8}
	_ = cfg.Fprint(buf, newfset, newlit)
	if fJSON {
		output := struct {
			Start       int    `json:"start"`
			End         int    `json:"end"`
			Replacement string `json:"replacement"`
		}{
			oldfset.Position(oldlit.Pos()).Offset,
			oldfset.Position(oldlit.End()).Offset,
			buf.String(),
		}
		_ = json.NewEncoder(os.Stdout).Encode(output)
	} else {
		fmt.Println(buf.String())
	}
}

func copyExpr(expr ast.Expr, line token.Pos) ast.Expr {
	switch expr := expr.(type) {
	case *ast.BasicLit:
		cp := *expr
		cp.ValuePos = 0
		return &cp
	case *ast.BinaryExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.OpPos = 0
		cp.Y = copyExpr(cp.Y, line)
		return &cp
	case *ast.CallExpr:
		cp := *expr
		cp.Fun = copyExpr(cp.Fun, line)
		cp.Lparen = 0
		for i, v := range cp.Args {
			cp.Args[i] = copyExpr(v, line)
		}
		cp.Ellipsis = line
		cp.Rparen = 0
		return &cp
	case *ast.CompositeLit:
		cp := *expr
		cp.Type = copyExpr(cp.Type, line)
		cp.Lbrace = 0
		for i, v := range cp.Elts {
			cp.Elts[i] = copyExpr(v, line)
		}
		cp.Rbrace = 0
		return &cp
	case *ast.Ident:
		cp := *expr
		cp.NamePos = 0
		return &cp
	case *ast.IndexExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Lbrack = 0
		cp.Index = copyExpr(cp.Index, line)
		cp.Rbrack = 0
		return &cp
	case *ast.ParenExpr:
		cp := *expr
		cp.Lparen = 0
		cp.X = copyExpr(cp.X, line)
		cp.Rparen = 0
		return &cp
	case *ast.SelectorExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Sel = copyExpr(cp.Sel, line).(*ast.Ident)
		return &cp
	case *ast.SliceExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Lbrack = 0
		cp.Low = copyExpr(cp.Low, line)
		cp.High = copyExpr(cp.High, line)
		cp.Max = copyExpr(cp.Max, line)
		cp.Rbrack = 0
		return &cp
	case *ast.StarExpr:
		cp := *expr
		cp.Star = 0
		cp.X = copyExpr(cp.X, line)
		return &cp
	case *ast.TypeAssertExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Lparen = 0
		cp.Type = copyExpr(cp.Type, line)
		cp.Rparen = 0
		return &cp
	case *ast.UnaryExpr:
		cp := *expr
		cp.OpPos = 0
		cp.X = copyExpr(cp.X, line)
		return &cp
	case nil:
		return nil
	default:
		panic(fmt.Sprintf("shouldn't happen: unknown ast.Expr of type %T", expr))
	}
	return nil
}
<commit_msg>Implement one-line printing of initializers<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/printer"
	"go/token"
	"go/types"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"golang.org/x/tools/go/ast/astutil"
	"golang.org/x/tools/go/buildutil"
	"golang.org/x/tools/go/loader"
)

func parseOctothorpDecimal(s string) int {
	if s != "" && s[0] == '#' {
		if s, err := strconv.ParseInt(s[1:], 10, 32); err == nil {
			return int(s)
		}
	}
	return -1
}

func parsePos(pos string) (filename string, startOffset, endOffset int, err error) {
	if pos == "" {
		err = fmt.Errorf("no source position specified")
		return
	}

	colon := strings.LastIndex(pos, ":")
	if colon < 0 {
		err = fmt.Errorf("bad position syntax %q", pos)
		return
	}
	filename, offset := pos[:colon], pos[colon+1:]
	startOffset = -1
	endOffset = -1
	if hyphen := strings.Index(offset, ","); hyphen < 0 {
		// e.g. "foo.go:#123"
		startOffset = parseOctothorpDecimal(offset)
		endOffset = startOffset
	} else {
		// e.g. "foo.go:#123,#456"
		startOffset = parseOctothorpDecimal(offset[:hyphen])
		endOffset = parseOctothorpDecimal(offset[hyphen+1:])
	}
	if startOffset < 0 || endOffset < 0 {
		err = fmt.Errorf("invalid offset %q in query position", offset)
		return
	}
	return
}

func fileOffsetToPos(file *token.File, startOffset, endOffset int) (start, end token.Pos, err error) {
	// Range check [start..end], inclusive of both end-points.

	if 0 <= startOffset && startOffset <= file.Size() {
		start = file.Pos(int(startOffset))
	} else {
		err = fmt.Errorf("start position is beyond end of file")
		return
	}

	if 0 <= endOffset && endOffset <= file.Size() {
		end = file.Pos(int(endOffset))
	} else {
		err = fmt.Errorf("end position is beyond end of file")
		return
	}

	return
}

var (
	fRecursive bool
	fOneLine   bool
	fJSON      bool
)

func init() {
	flag.BoolVar(&fRecursive, "r", false, "keyify struct initializers recursively")
	flag.BoolVar(&fOneLine, "o", false, "print new struct initializer on a single line")
	flag.BoolVar(&fJSON, "json", false, "print new struct initializer as JSON")
}

func usage() {
	fmt.Printf("Usage: %s [flags] <position>\n\n", os.Args[0])
	flag.PrintDefaults()
}

func main() {
	log.SetFlags(0)
	flag.Usage = usage
	flag.Parse()
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(2)
	}
	pos := flag.Args()[0]
	name, start, _, err := parsePos(pos)
	if err != nil {
		log.Fatal(err)
	}
	eval, err := filepath.EvalSymlinks(name)
	if err != nil {
		log.Fatal(err)
	}
	name, err = filepath.Abs(eval)
	if err != nil {
		log.Fatal(err)
	}
	cwd, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	bpkg, err := buildutil.ContainingPackage(&build.Default, cwd, name)
	if err != nil {
		log.Fatal(err)
	}
	conf := &loader.Config{}
	conf.TypeCheckFuncBodies = func(s string) bool {
		return s == bpkg.ImportPath || s == bpkg.ImportPath+"_test"
	}
	conf.ImportWithTests(bpkg.ImportPath)
	lprog, err := conf.Load()
	if err != nil {
		log.Fatal(err)
	}
	var tf *token.File
	var af *ast.File
	pkg := lprog.InitialPackages()[0]
	for _, ff := range pkg.Files {
		file := lprog.Fset.File(ff.Pos())
		if file.Name() == name {
			af = ff
			tf = file
			break
		}
	}
	tstart, tend, err := fileOffsetToPos(tf, start, start)
	if err != nil {
		log.Fatal(err)
	}
	path, _ := astutil.PathEnclosingInterval(af, tstart, tend)
	var complit *ast.CompositeLit
	for _, p := range path {
		if p, ok := p.(*ast.CompositeLit); ok {
			complit = p
			break
		}
	}
	if complit == nil {
		log.Fatal("no composite literal found near point")
	}
	if len(complit.Elts) == 0 {
		printComplit(complit, complit, lprog.Fset, lprog.Fset)
		return
	}
	if _, ok := complit.Elts[0].(*ast.KeyValueExpr); ok {
		printComplit(complit, complit, lprog.Fset, lprog.Fset)
		return
	}
	st, ok := pkg.TypeOf(complit.Type).Underlying().(*types.Struct)
	if !ok {
		log.Fatal("not a struct initialiser")
		return
	}

	var calcPos func(int) token.Pos
	if fOneLine {
		calcPos = func(int) token.Pos { return token.Pos(1) }
	} else {
		calcPos = func(i int) token.Pos { return token.Pos(2 + i) }
	}

	newFset := token.NewFileSet()
	newFile := newFset.AddFile("", -1, st.NumFields()+2)
	newComplit := &ast.CompositeLit{
		Type:   complit.Type,
		Lbrace: 1,
		Rbrace: token.Pos(st.NumFields() + 2),
	}
	if fOneLine {
		newComplit.Rbrace = 1
	}
	newFile.AddLine(1)
	newFile.AddLine(st.NumFields() + 2)
	for i := 0; i < st.NumFields(); i++ {
		newFile.AddLine(2 + i)
		field := st.Field(i)
		elt := &ast.KeyValueExpr{
			Key:   &ast.Ident{NamePos: calcPos(i), Name: field.Name()},
			Value: copyExpr(complit.Elts[i], calcPos(i)),
		}
		newComplit.Elts = append(newComplit.Elts, elt)
	}
	printComplit(complit, newComplit, lprog.Fset, newFset)
}

func printComplit(oldlit, newlit *ast.CompositeLit, oldfset, newfset *token.FileSet) {
	buf := &bytes.Buffer{}
	cfg := printer.Config{Mode: printer.UseSpaces | printer.TabIndent, Tabwidth: 8}
	_ = cfg.Fprint(buf, newfset, newlit)
	if fJSON {
		output := struct {
			Start       int    `json:"start"`
			End         int    `json:"end"`
			Replacement string `json:"replacement"`
		}{
			oldfset.Position(oldlit.Pos()).Offset,
			oldfset.Position(oldlit.End()).Offset,
			buf.String(),
		}
		_ = json.NewEncoder(os.Stdout).Encode(output)
	} else {
		fmt.Println(buf.String())
	}
}

func copyExpr(expr ast.Expr, line token.Pos) ast.Expr {
	switch expr := expr.(type) {
	case *ast.BasicLit:
		cp := *expr
		cp.ValuePos = 0
		return &cp
	case *ast.BinaryExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.OpPos = 0
		cp.Y = copyExpr(cp.Y, line)
		return &cp
	case *ast.CallExpr:
		cp := *expr
		cp.Fun = copyExpr(cp.Fun, line)
		cp.Lparen = 0
		for i, v := range cp.Args {
			cp.Args[i] = copyExpr(v, line)
		}
		cp.Ellipsis = line
		cp.Rparen = 0
		return &cp
	case *ast.CompositeLit:
		cp := *expr
		cp.Type = copyExpr(cp.Type, line)
		cp.Lbrace = 0
		for i, v := range cp.Elts {
			cp.Elts[i] = copyExpr(v, line)
		}
		cp.Rbrace = 0
		return &cp
	case *ast.Ident:
		cp := *expr
		cp.NamePos = 0
		return &cp
	case *ast.IndexExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Lbrack = 0
		cp.Index = copyExpr(cp.Index, line)
		cp.Rbrack = 0
		return &cp
	case *ast.ParenExpr:
		cp := *expr
		cp.Lparen = 0
		cp.X = copyExpr(cp.X, line)
		cp.Rparen = 0
		return &cp
	case *ast.SelectorExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Sel = copyExpr(cp.Sel, line).(*ast.Ident)
		return &cp
	case *ast.SliceExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Lbrack = 0
		cp.Low = copyExpr(cp.Low, line)
		cp.High = copyExpr(cp.High, line)
		cp.Max = copyExpr(cp.Max, line)
		cp.Rbrack = 0
		return &cp
	case *ast.StarExpr:
		cp := *expr
		cp.Star = 0
		cp.X = copyExpr(cp.X, line)
		return &cp
	case *ast.TypeAssertExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Lparen = 0
		cp.Type = copyExpr(cp.Type, line)
		cp.Rparen = 0
		return &cp
	case *ast.UnaryExpr:
		cp := *expr
		cp.OpPos = 0
		cp.X = copyExpr(cp.X, line)
		return &cp
	case nil:
		return nil
	default:
		panic(fmt.Sprintf("shouldn't happen: unknown ast.Expr of type %T", expr))
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"

	"gopkg.in/check.v1"

	"github.com/containers/image/v5/manifest"
)

// This image is known to be x86_64 only right now
const knownNotManifestListedImage_x8664 = "docker://quay.io/coreos/11bot"

const expectedProxySemverMajor = "0.2"

// request is copied from proxy.go
// We intentionally copy to ensure that we catch any unexpected "API" changes
// in the JSON.
type request struct {
	// Method is the name of the function
	Method string `json:"method"`
	// Args is the arguments (parsed inside the fuction)
	Args []interface{} `json:"args"`
}

// reply is copied from proxy.go
type reply struct {
	// Success is true if and only if the call succeeded.
	Success bool `json:"success"`
	// Value is an arbitrary value (or values, as array/map) returned from the call.
	Value interface{} `json:"value"`
	// PipeID is an index into open pipes, and should be passed to FinishPipe
	PipeID uint32 `json:"pipeid"`
	// Error should be non-empty if Success == false
	Error string `json:"error"`
}

// maxMsgSize is also copied from proxy.go
const maxMsgSize = 32 * 1024

type proxy struct {
	c *net.UnixConn
}

type pipefd struct {
	// id is the remote identifier "pipeid"
	id uint
	fd *os.File
}

func (self *proxy) call(method string, args []interface{}) (rval interface{}, fd *pipefd, err error) {
	req := request{
		Method: method,
		Args:   args,
	}
	reqbuf, err := json.Marshal(&req)
	if err != nil {
		return
	}
	n, err := self.c.Write(reqbuf)
	if err != nil {
		return
	}
	if n != len(reqbuf) {
		err = fmt.Errorf("short write during call of %d bytes", n)
		return
	}
	oob := make([]byte, syscall.CmsgSpace(1))
	replybuf := make([]byte, maxMsgSize)
	n, oobn, _, _, err := self.c.ReadMsgUnix(replybuf, oob)
	if err != nil {
		err = fmt.Errorf("reading reply: %v", err)
		return
	}
	var reply reply
	err = json.Unmarshal(replybuf[0:n], &reply)
	if err != nil {
		err = fmt.Errorf("Failed to parse reply: %w", err)
		return
	}
	if !reply.Success {
		err = fmt.Errorf("remote error: %s", reply.Error)
		return
	}

	if reply.PipeID > 0 {
		var scms []syscall.SocketControlMessage
		scms, err = syscall.ParseSocketControlMessage(oob[:oobn])
		if err != nil {
			err = fmt.Errorf("failed to parse control message: %v", err)
			return
		}
		if len(scms) != 1 {
			err = fmt.Errorf("Expected 1 received fd, found %d", len(scms))
			return
		}
		var fds []int
		fds, err = syscall.ParseUnixRights(&scms[0])
		if err != nil {
			err = fmt.Errorf("failed to parse unix rights: %v", err)
			return
		}
		fd = &pipefd{
			fd: os.NewFile(uintptr(fds[0]), "replyfd"),
			id: uint(reply.PipeID),
		}
	}

	rval = reply.Value
	return
}

func (self *proxy) callNoFd(method string, args []interface{}) (rval interface{}, err error) {
	var fd *pipefd
	rval, fd, err = self.call(method, args)
	if err != nil {
		return
	}
	if fd != nil {
		err = fmt.Errorf("Unexpected fd from method %s", method)
		return
	}
	return rval, nil
}

func newProxy() (*proxy, error) {
	fds, err := syscall.Socketpair(syscall.AF_LOCAL, syscall.SOCK_SEQPACKET, 0)
	if err != nil {
		return nil, err
	}
	myfd := os.NewFile(uintptr(fds[0]), "myfd")
	defer myfd.Close()
	theirfd := os.NewFile(uintptr(fds[1]), "theirfd")
	defer theirfd.Close()

	mysock, err := net.FileConn(myfd)
	if err != nil {
		return nil, err
	}

	// Note ExtraFiles starts at 3
	proc := exec.Command("skopeo", "experimental-image-proxy", "--sockfd", "3")
	proc.Stderr = os.Stderr
	proc.SysProcAttr = &syscall.SysProcAttr{
		Pdeathsig: syscall.SIGTERM,
	}
	proc.ExtraFiles = append(proc.ExtraFiles, theirfd)

	if err = proc.Start(); err != nil {
		return nil, err
	}

	p := &proxy{
		c: mysock.(*net.UnixConn),
	}

	v, err := p.callNoFd("Initialize", nil)
	if err != nil {
		return nil, err
	}
	semver, ok := v.(string)
	if !ok {
		return nil, fmt.Errorf("proxy Initialize: Unexpected value %T", v)
	}
	if !strings.HasPrefix(semver, expectedProxySemverMajor) {
		return nil, fmt.Errorf("Unexpected semver %s", semver)
	}
	return p, nil
}

func init() {
	check.Suite(&ProxySuite{})
}

type ProxySuite struct {
}

func (s *ProxySuite) SetUpSuite(c *check.C) {
}

func (s *ProxySuite) TearDownSuite(c *check.C) {
}

type byteFetch struct {
	content []byte
	err     error
}

func runTestGetManifest(p *proxy, img string) error {
	v, err := p.callNoFd("OpenImage", []interface{}{knownNotManifestListedImage_x8664})
	if err != nil {
		return err
	}

	imgidv, ok := v.(float64)
	if !ok {
		return fmt.Errorf("OpenImage return value is %T", v)
	}
	imgid := uint32(imgidv)

	v, fd, err := p.call("GetManifest", []interface{}{imgid})
	if err != nil {
		return err
	}
	if fd == nil {
		return fmt.Errorf("expected GetManifest fd")
	}
	fetchchan := make(chan byteFetch)
	go func() {
		manifestBytes, err := ioutil.ReadAll(fd.fd)
		fetchchan <- byteFetch{
			content: manifestBytes,
			err:     err,
		}
	}()
	_, _, err = p.call("FinishPipe", []interface{}{fd.id})
	if err != nil {
		return err
	}
	fetchRes := <-fetchchan
	if fetchRes.err != nil {
		return err
	}
	_, err = manifest.OCI1FromManifest(fetchRes.content)
	if err != nil {
		return err
	}

	_, _, err = p.call("CloseImage", []interface{}{imgid})

	return nil
}

func (s *ProxySuite) TestProxy(c *check.C) {
	p, err := newProxy()
	c.Assert(err, check.IsNil)

	err = runTestGetManifest(p, knownNotManifestListedImage_x8664)
	if err != nil {
		err = fmt.Errorf("Testing image %s: %v", knownNotManifestListedImage_x8664, err)
	}
	c.Assert(err, check.IsNil)

	err = runTestGetManifest(p, knownListImage)
	if err != nil {
		err = fmt.Errorf("Testing image %s: %v", knownListImage, err)
	}
	c.Assert(err, check.IsNil)
}
<commit_msg>proxy_test: Add helper to read all from a reply<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"

	"gopkg.in/check.v1"

	"github.com/containers/image/v5/manifest"
)

// This image is known to be x86_64 only right now
const knownNotManifestListedImage_x8664 = "docker://quay.io/coreos/11bot"

const expectedProxySemverMajor = "0.2"

// request is copied from proxy.go
// We intentionally copy to ensure that we catch any unexpected "API" changes
// in the JSON.
type request struct {
	// Method is the name of the function
	Method string `json:"method"`
	// Args is the arguments (parsed inside the fuction)
	Args []interface{} `json:"args"`
}

// reply is copied from proxy.go
type reply struct {
	// Success is true if and only if the call succeeded.
	Success bool `json:"success"`
	// Value is an arbitrary value (or values, as array/map) returned from the call.
	Value interface{} `json:"value"`
	// PipeID is an index into open pipes, and should be passed to FinishPipe
	PipeID uint32 `json:"pipeid"`
	// Error should be non-empty if Success == false
	Error string `json:"error"`
}

// maxMsgSize is also copied from proxy.go
const maxMsgSize = 32 * 1024

type proxy struct {
	c *net.UnixConn
}

type pipefd struct {
	// id is the remote identifier "pipeid"
	id uint
	fd *os.File
}

func (self *proxy) call(method string, args []interface{}) (rval interface{}, fd *pipefd, err error) {
	req := request{
		Method: method,
		Args:   args,
	}
	reqbuf, err := json.Marshal(&req)
	if err != nil {
		return
	}
	n, err := self.c.Write(reqbuf)
	if err != nil {
		return
	}
	if n != len(reqbuf) {
		err = fmt.Errorf("short write during call of %d bytes", n)
		return
	}
	oob := make([]byte, syscall.CmsgSpace(1))
	replybuf := make([]byte, maxMsgSize)
	n, oobn, _, _, err := self.c.ReadMsgUnix(replybuf, oob)
	if err != nil {
		err = fmt.Errorf("reading reply: %v", err)
		return
	}
	var reply reply
	err = json.Unmarshal(replybuf[0:n], &reply)
	if err != nil {
		err = fmt.Errorf("Failed to parse reply: %w", err)
		return
	}
	if !reply.Success {
		err = fmt.Errorf("remote error: %s", reply.Error)
		return
	}

	if reply.PipeID > 0 {
		var scms []syscall.SocketControlMessage
		scms, err = syscall.ParseSocketControlMessage(oob[:oobn])
		if err != nil {
			err = fmt.Errorf("failed to parse control message: %v", err)
			return
		}
		if len(scms) != 1 {
			err = fmt.Errorf("Expected 1 received fd, found %d", len(scms))
			return
		}
		var fds []int
		fds, err = syscall.ParseUnixRights(&scms[0])
		if err != nil {
			err = fmt.Errorf("failed to parse unix rights: %v", err)
			return
		}
		fd = &pipefd{
			fd: os.NewFile(uintptr(fds[0]), "replyfd"),
			id: uint(reply.PipeID),
		}
	}

	rval = reply.Value
	return
}

func (self *proxy) callNoFd(method string, args []interface{}) (rval interface{}, err error) {
	var fd *pipefd
	rval, fd, err = self.call(method, args)
	if err != nil {
		return
	}
	if fd != nil {
		err = fmt.Errorf("Unexpected fd from method %s", method)
		return
	}
	return rval, nil
}

func (self *proxy) callReadAllBytes(method string, args []interface{}) (rval interface{}, buf []byte, err error) {
	var fd *pipefd
	rval, fd, err = self.call(method, args)
	if err != nil {
		return
	}
	if fd == nil {
		err = fmt.Errorf("Expected fd from method %s", method)
		return
	}
	fetchchan := make(chan byteFetch)
	go func() {
		manifestBytes, err := ioutil.ReadAll(fd.fd)
		fetchchan <- byteFetch{
			content: manifestBytes,
			err:     err,
		}
	}()
	_, err = self.callNoFd("FinishPipe", []interface{}{fd.id})
	if err != nil {
		return
	}
	select {
	case fetchRes := <-fetchchan:
		err = fetchRes.err
		if err != nil {
			return
		}

		buf = fetchRes.content
	case <-time.After(5 * time.Minute):
		err = fmt.Errorf("timed out during proxy fetch")
	}
	return
}

func newProxy() (*proxy, error) {
	fds, err := syscall.Socketpair(syscall.AF_LOCAL, syscall.SOCK_SEQPACKET, 0)
	if err != nil {
		return nil, err
	}
	myfd := os.NewFile(uintptr(fds[0]), "myfd")
	defer myfd.Close()
	theirfd := os.NewFile(uintptr(fds[1]), "theirfd")
	defer theirfd.Close()

	mysock, err := net.FileConn(myfd)
	if err != nil {
		return nil, err
	}

	// Note ExtraFiles starts at 3
	proc := exec.Command("skopeo", "experimental-image-proxy", "--sockfd", "3")
	proc.Stderr = os.Stderr
	proc.SysProcAttr = &syscall.SysProcAttr{
		Pdeathsig: syscall.SIGTERM,
	}
	proc.ExtraFiles = append(proc.ExtraFiles, theirfd)

	if err = proc.Start(); err != nil {
		return nil, err
	}

	p := &proxy{
		c: mysock.(*net.UnixConn),
	}

	v, err := p.callNoFd("Initialize", nil)
	if err != nil {
		return nil, err
	}
	semver, ok := v.(string)
	if !ok {
		return nil, fmt.Errorf("proxy Initialize: Unexpected value %T", v)
	}
	if !strings.HasPrefix(semver, expectedProxySemverMajor) {
		return nil, fmt.Errorf("Unexpected semver %s", semver)
	}
	return p, nil
}

func init() {
	check.Suite(&ProxySuite{})
}

type ProxySuite struct {
}

func (s *ProxySuite) SetUpSuite(c *check.C) {
}

func (s *ProxySuite) TearDownSuite(c *check.C) {
}

type byteFetch struct {
	content []byte
	err     error
}

func runTestGetManifest(p *proxy, img string) error {
	v, err := p.callNoFd("OpenImage", []interface{}{knownNotManifestListedImage_x8664})
	if err != nil {
		return err
	}

	imgidv, ok := v.(float64)
	if !ok {
		return fmt.Errorf("OpenImage return value is %T", v)
	}
	imgid := uint32(imgidv)

	v, manifestBytes, err := p.callReadAllBytes("GetManifest", []interface{}{imgid})
	if err != nil {
		return err
	}
	_, err = manifest.OCI1FromManifest(manifestBytes)
	if err != nil {
		return err
	}

	_, err = p.callNoFd("CloseImage", []interface{}{imgid})

	return nil
}

func (s *ProxySuite) TestProxy(c *check.C) {
	p, err := newProxy()
	c.Assert(err, check.IsNil)

	err = runTestGetManifest(p, knownNotManifestListedImage_x8664)
	if err != nil {
		err = fmt.Errorf("Testing image %s: %v", knownNotManifestListedImage_x8664, err)
	}
	c.Assert(err, check.IsNil)

	err = runTestGetManifest(p, knownListImage)
	if err != nil {
		err = fmt.Errorf("Testing image %s: %v", knownListImage, err)
	}
	c.Assert(err, check.IsNil)
}
<|endoftext|>


<commit_before>/*
 * This file is part of remco.
 * © 2016 The Remco Authors
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

package template

import (
	"net"
	"os"

	. "gopkg.in/check.v1"
)

type FunctionTestSuite struct{}

var _ = Suite(&FunctionTestSuite{})

func (s *FunctionTestSuite) TestAddFuncs(t *C) {
	in := map[string]interface{}{
		"a": "hallo",
		"b": "hello",
	}
	out := make(map[string]interface{})

	addFuncs(out, in)

	t.Check(len(out), Equals, len(in))
}

func (s *FunctionTestSuite) TestLookupIP(t *C) {
	ips, err := lookupIP("localhost")
	if err != nil {
		t.Error(err)
	}
	if len(ips) > 0 {
		t.Check(ips[0], Equals, "127.0.0.1")
	} else {
		t.Error("lookupIP failed")
	}
}

func (s *FunctionTestSuite) TestLookupSRV(t *C) {
	expected := []*net.SRV{
		{
			Target:   "alt1.xmpp-server.l.google.com.",
			Port:     5269,
			Priority: 20,
			Weight:   0,
		},
		{
			Target:   "alt2.xmpp-server.l.google.com.",
			Port:     5269,
			Priority: 20,
			Weight:   0,
		},
		{
			Target:   "alt3.xmpp-server.l.google.com.",
			Port:     5269,
			Priority: 20,
			Weight:   0,
		},
		{
			Target:   "alt4.xmpp-server.l.google.com.",
			Port:     5269,
			Priority: 20,
			Weight:   0,
		},
		{
			Target:   "xmpp-server.l.google.com.",
			Port:     5269,
			Priority: 5,
			Weight:   0,
		},
	}

	srv, err := lookupSRV("xmpp-server", "tcp", "google.com")
	if err != nil {
		t.Error(err)
	}
	t.Check(srv, DeepEquals, expected)
}

func (s *FunctionTestSuite) TestGetEnv(t *C) {
	key := "coolEnvVar"
	expected := "mmmh lecker saure Gurken!"
	err := os.Setenv(key, expected)
	if err != nil {
		t.Error(err)
	}

	t.Check(getenv(key), Equals, expected)
}

func (s *FunctionTestSuite) TestGetEnvDefault(t *C) {
	key := "ihopethisenvvardontexists"
	expected := "default"

	t.Check(getenv(key, "default"), Equals, expected)
}

func (s *FunctionTestSuite) TestInterfaceSet(t *C) {
	set := createSet()
	set.Append("Hallo")
	set.Append("Hallo")
	set.Append(1)
	set.Append(true)
	set.Append(false)

	t.Check(len(set.S), Equals, 4)
	t.Check(set.Contains("Hallo"), Equals, true)
	set.Remove("Hallo")
	t.Check(len(set.S), Equals, 3)
	t.Check(set.Contains("Hallo"), Equals, false)
	t.Check(set.Contains(false), Equals, true)
}

func (s *FunctionTestSuite) TestTemplateMap(t *C) {
	m := createMap()
	m.Set("Hallo", "OneOneOne")
	m.Set("Test", "Snickers")
	m.Set("One", 1)

	t.Check(m.Get("Hallo"), DeepEquals, "OneOneOne")
	t.Check(m.Get("One"), DeepEquals, 1)

	m.Remove("One")
	t.Check(m.Get("One"), DeepEquals, nil)
}
<commit_msg>Fix tests for template_funcs<commit_after>/*
 * This file is part of remco.
 * © 2016 The Remco Authors
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

package template

import (
	"net"
	"os"

	. "gopkg.in/check.v1"
)

type FunctionTestSuite struct{}

var _ = Suite(&FunctionTestSuite{})

func (s *FunctionTestSuite) TestAddFuncs(t *C) {
	in := map[string]interface{}{
		"a": "hallo",
		"b": "hello",
	}
	out := make(map[string]interface{})

	addFuncs(out, in)

	t.Check(len(out), Equals, len(in))
}

func (s *FunctionTestSuite) TestLookupIP(t *C) {
	ips, err := lookupIP("localhost")
	if err != nil {
		t.Error(err)
	}
	if len(ips) > 0 {
		t.Check(ips[0], Equals, "127.0.0.1")
	} else {
		t.Error("lookupIP failed")
	}
}

func (s *FunctionTestSuite) TestLookupSRV(t *C) {
	expected := []*net.SRV{
		{
			Target:   "alt1.xmpp-server.l.google.com.",
			Port:     5269,
			Priority: 20,
			Weight:   0,
		},
		{
			Target:   "alt2.xmpp-server.l.google.com.",
			Port:     5269,
			Priority: 20,
			Weight:   0,
		},
		{
			Target:   "alt3.xmpp-server.l.google.com.",
			Port:     5269,
			Priority: 20,
			Weight:   0,
		},
		{
			Target:   "alt4.xmpp-server.l.google.com.",
			Port:     5269,
			Priority: 20,
			Weight:   0,
		},
		{
			Target:   "xmpp-server.l.google.com.",
			Port:     5269,
			Priority: 5,
			Weight:   0,
		},
	}

	srv, err := lookupSRV("xmpp-server", "tcp", "google.com")
	if err != nil {
		t.Error(err)
	}
	t.Check(srv, DeepEquals, expected)
}

func (s *FunctionTestSuite) TestGetEnv(t *C) {
	key := "coolEnvVar"
	expected := "mmmh lecker saure Gurken!"
	err := os.Setenv(key, expected)
	if err != nil {
		t.Error(err)
	}

	t.Check(getenv(key), Equals, expected)
}

func (s *FunctionTestSuite) TestGetEnvDefault(t *C) {
	key := "ihopethisenvvardontexists"
	expected := "default"

	t.Check(getenv(key, "default"), Equals, expected)
}

func (s *FunctionTestSuite) TestInterfaceSet(t *C) {
	set := createSet()
	set.Append("Hallo")
	set.Append("Hallo")
	set.Append(1)
	set.Append(true)
	set.Append(false)

	t.Check(len(set), Equals, 4)
	t.Check(set.Contains("Hallo"), Equals, true)
	set.Remove("Hallo")
	t.Check(len(set), Equals, 3)
	t.Check(set.Contains("Hallo"), Equals, false)
	t.Check(set.Contains(false), Equals, true)
}

func (s *FunctionTestSuite) TestTemplateMap(t *C) {
	m := createMap()
	m.Set("Hallo", "OneOneOne")
	m.Set("Test", "Snickers")
	m.Set("One", 1)

	t.Check(m.Get("Hallo"), DeepEquals, "OneOneOne")
	t.Check(m.Get("One"), DeepEquals, 1)

	m.Remove("One")
	t.Check(m.Get("One"), DeepEquals, nil)
}
<|endoftext|>


<commit_before>package writerlease

import (
	"sync"
	"time"

	"github.com/golang/glog"

	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/util/workqueue"
)

// Lease performs the equivalent of leader election by competing to perform work (such as
// updating a contended resource). Every successful work unit is considered a lease renewal,
// while work that is observed from others or that fails is treated as renewing another processes
// lease. When a lease expires (no work is detected within the lease term) the writer competes
// to perform work. When competing for the lease, exponential backoff is used.
type Lease interface {
	// Wait waits for the first work function to complete and then returns whether the current
	// process is the leader. This function will block forever if no work has been requested or if the
	// work retries forever.
	Wait() bool
	// WaitUntil waits at most the provided duration for the frist work function to complete.
	// If the duration expires without work completing it will return false for expired, otherwise
	// it will return whether the lease is held by this process.
	WaitUntil(t time.Duration) (leader bool, ok bool)
	// Try runs the provided function when the lease is held is the leader. It retries work until
	// the work func indicates retry is not necessary.
	Try(key string, fn WorkFunc)
	// Extend indicates that the caller has observed another writer performing work against
	// the specified key. This will clear the work remaining for the lease and extend the lease
	// interval.
	Extend(key string)
	// Remove clears any pending work for the provided key.
	Remove(key string)
}

// WorkFunc is a retriable unit of work. It should return an error if the work couldn't be
// completed successfully, or true if we can assume our lease has been extended. If the
// lease could not be extended, we drop this unit of work.
type WorkFunc func() (extendLease, retry bool)

// LimitRetries allows a work function to be retried up to retries times.
func LimitRetries(retries int, fn WorkFunc) WorkFunc {
	i := 0
	return func() (bool, bool) {
		extend, retry := fn()
		if retry {
			retry = i < retries
			i++
		}
		return extend, retry
	}
}

// State is the state of the lease.
type State int

const (
	// Election is before a work unit has been completed.
	Election State = iota
	Leader
	Follower
)

var nowFn = time.Now

type WriterLease struct {
	name          string
	backoff       wait.Backoff
	maxBackoff    time.Duration
	retryInterval time.Duration
	once          chan struct{}

	lock    sync.Mutex
	queued  map[string]WorkFunc
	queue   workqueue.DelayingInterface
	state   State
	expires time.Time
	tick    int
}

// New creates a new Lease. Specify the duration to hold leases for and the retry
// interval on requests that fail.
func New(leaseDuration, retryInterval time.Duration) *WriterLease {
	backoff := wait.Backoff{
		Duration: 20 * time.Millisecond,
		Factor:   4,
		Steps:    5,
		Jitter:   0.5,
	}
	return &WriterLease{
		backoff:       backoff,
		maxBackoff:    leaseDuration,
		retryInterval: retryInterval,

		queued: make(map[string]WorkFunc),
		queue:  workqueue.NewDelayingQueue(),
		once:   make(chan struct{}),
	}
}

// NewWithBackoff creates a new Lease. Specify the duration to hold leases for and the retry
// interval on requests that fail.
func NewWithBackoff(name string, leaseDuration, retryInterval time.Duration, backoff wait.Backoff) *WriterLease {
	return &WriterLease{
		name:          name,
		backoff:       backoff,
		maxBackoff:    leaseDuration,
		retryInterval: retryInterval,

		queued: make(map[string]WorkFunc),
		queue:  workqueue.NewNamedDelayingQueue(name),
		once:   make(chan struct{}),
	}
}

func (l *WriterLease) Run(stopCh <-chan struct{}) {
	defer utilruntime.HandleCrash()
	defer l.queue.ShutDown()

	go func() {
		defer utilruntime.HandleCrash()
		for l.work() {
		}
		glog.V(4).Infof("[%s] Worker stopped", l.name)
	}()

	<-stopCh
}

func (l *WriterLease) Expire() {
	l.lock.Lock()
	defer l.lock.Unlock()
	l.expires = time.Time{}
}

func (l *WriterLease) Wait() bool {
	<-l.once
	state, _, _ := l.leaseState()
	return state == Leader
}

func (l *WriterLease) WaitUntil(t time.Duration) (bool, bool) {
	select {
	case <-l.once:
	case <-time.After(t):
		return false, false
	}
	state, _, _ := l.leaseState()
	return state == Leader, true
}

func (l *WriterLease) Try(key string, fn WorkFunc) {
	l.lock.Lock()
	defer l.lock.Unlock()
	l.queued[key] = fn
	if l.state == Follower {
		delay := l.expires.Sub(nowFn())
		// no matter what, always wait at least some amount of time as a follower to give the nominal
		// leader a chance to win
		if delay < l.backoff.Duration*2 {
			delay = l.backoff.Duration * 2
		}
		l.queue.AddAfter(key, delay)
	} else {
		l.queue.Add(key)
	}
}

func (l *WriterLease) Extend(key string) {
	l.lock.Lock()
	defer l.lock.Unlock()
	if _, ok := l.queued[key]; ok {
		delete(l.queued, key)
		switch l.state {
		case Follower:
			l.tick++
			backoff := l.nextBackoff()
			glog.V(4).Infof("[%s] Clearing work for %s and extending lease by %s", l.name, key, backoff)
			l.expires = nowFn().Add(backoff)
		}
	}
}

func (l *WriterLease) Len() int {
	l.lock.Lock()
	defer l.lock.Unlock()
	return len(l.queued)
}

func (l *WriterLease) Remove(key string) {
	l.lock.Lock()
	defer l.lock.Unlock()
	delete(l.queued, key)
}

func (l *WriterLease) get(key string) WorkFunc {
	l.lock.Lock()
	defer l.lock.Unlock()
	return l.queued[key]
}

func (l *WriterLease) leaseState() (State, time.Time, int) {
	l.lock.Lock()
	defer l.lock.Unlock()
	return l.state, l.expires, l.tick
}

func (l *WriterLease) work() bool {
	item, shutdown := l.queue.Get()
	if shutdown {
		return false
	}
	key := item.(string)

	fn := l.get(key)
	if fn == nil {
		glog.V(4).Infof("[%s] Work item %s was cleared, done", l.name, key)
		l.queue.Done(key)
		return true
	}

	leaseState, leaseExpires, _ := l.leaseState()
	if leaseState == Follower {
		// if we are following, continue to defer work until the lease expires
		if remaining := leaseExpires.Sub(nowFn()); remaining > 0 {
			glog.V(4).Infof("[%s] Follower, %s remaining in lease", l.name, remaining)
			l.queue.AddAfter(key, remaining)
			l.queue.Done(key)
			return true
		}
		glog.V(4).Infof("[%s] Lease expired, running %s", l.name, key)
	} else {
		glog.V(4).Infof("[%s] Lease owner or electing, running %s", l.name, key)
	}

	isLeader, retry := fn()
	if retry {
		// come back in a bit
		glog.V(4).Infof("[%s] Retrying %s", l.name, key)
		l.queue.AddAfter(key, l.retryInterval)
		l.queue.Done(key)
		return true
	}

	l.finishKey(key, isLeader)
	return true
}

func (l *WriterLease) finishKey(key string, isLeader bool) {
	l.lock.Lock()
	defer l.lock.Unlock()

	resolvedElection := l.state == Election
	if isLeader {
		switch l.state {
		case Election, Follower:
			l.tick = 0
			l.state = Leader
		}
		l.expires = nowFn().Add(l.maxBackoff)
	} else {
		switch l.state {
		case Election, Leader:
			l.tick = 0
			l.state = Follower
		case Follower:
			l.tick++
		}
		l.expires = nowFn().Add(l.nextBackoff())
	}
	delete(l.queued, key)
	// close the channel before we remove the key from the queue to prevent races in Wait
	if resolvedElection {
		close(l.once)
	}
	l.queue.Done(key)
	glog.V(4).Infof("[%s] Completed work for %s in state=%d tick=%d expires=%s", l.name, key, l.state, l.tick, l.expires)
}

func (l *WriterLease) nextBackoff() time.Duration {
	step := l.tick
	b := l.backoff
	if step > b.Steps {
		return l.maxBackoff
	}
	duration := b.Duration
	for i := 0; i < step; i++ {
		adjusted := duration
		if b.Jitter > 0.0 {
			adjusted = wait.Jitter(duration, b.Jitter)
		}
		duration = time.Duration(float64(adjusted) * b.Factor)
		if duration > l.maxBackoff {
			return l.maxBackoff
		}
	}
	return duration
}
<commit_msg>Ensure that two functions added for the same work key get run<commit_after>package writerlease

import (
	"sync"
	"time"

	"github.com/golang/glog"

	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/util/workqueue"
)

// Lease performs the equivalent of leader election by competing to perform work (such as
// updating a contended resource). Every successful work unit is considered a lease renewal,
// while work that is observed from others or that fails is treated as renewing another processes
// lease. When a lease expires (no work is detected within the lease term) the writer competes
// to perform work. When competing for the lease, exponential backoff is used.
type Lease interface {
	// Wait waits for the first work function to complete and then returns whether the current
	// process is the leader. This function will block forever if no work has been requested or if the
	// work retries forever.
	Wait() bool
	// WaitUntil waits at most the provided duration for the frist work function to complete.
	// If the duration expires without work completing it will return false for expired, otherwise
	// it will return whether the lease is held by this process.
	WaitUntil(t time.Duration) (leader bool, ok bool)
	// Try runs the provided function when the lease is held is the leader. It retries work until
	// the work func indicates retry is not necessary.
	Try(key string, fn WorkFunc)
	// Extend indicates that the caller has observed another writer performing work against
	// the specified key. This will clear the work remaining for the lease and extend the lease
	// interval.
	Extend(key string)
	// Remove clears any pending work for the provided key.
	Remove(key string)
}

// WorkFunc is a retriable unit of work. It should return an error if the work couldn't be
// completed successfully, or true if we can assume our lease has been extended. If the
// lease could not be extended, we drop this unit of work.
type WorkFunc func() (extendLease, retry bool)

// LimitRetries allows a work function to be retried up to retries times.
func LimitRetries(retries int, fn WorkFunc) WorkFunc {
	i := 0
	return func() (bool, bool) {
		extend, retry := fn()
		if retry {
			retry = i < retries
			i++
		}
		return extend, retry
	}
}

// State is the state of the lease.
type State int

const (
	// Election is before a work unit has been completed.
	Election State = iota
	Leader
	Follower
)

var nowFn = time.Now

type work struct {
	id int
	fn WorkFunc
}

type WriterLease struct {
	name          string
	backoff       wait.Backoff
	maxBackoff    time.Duration
	retryInterval time.Duration
	once          chan struct{}

	lock    sync.Mutex
	id      int
	queued  map[string]*work
	queue   workqueue.DelayingInterface
	state   State
	expires time.Time
	tick    int
}

// New creates a new Lease. Specify the duration to hold leases for and the retry
// interval on requests that fail.
func New(leaseDuration, retryInterval time.Duration) *WriterLease {
	backoff := wait.Backoff{
		Duration: 20 * time.Millisecond,
		Factor:   4,
		Steps:    5,
		Jitter:   0.5,
	}
	return &WriterLease{
		backoff:       backoff,
		maxBackoff:    leaseDuration,
		retryInterval: retryInterval,

		queued: make(map[string]*work),
		queue:  workqueue.NewDelayingQueue(),
		once:   make(chan struct{}),
	}
}

// NewWithBackoff creates a new Lease. Specify the duration to hold leases for and the retry
// interval on requests that fail.
func NewWithBackoff(name string, leaseDuration, retryInterval time.Duration, backoff wait.Backoff) *WriterLease {
	return &WriterLease{
		name:          name,
		backoff:       backoff,
		maxBackoff:    leaseDuration,
		retryInterval: retryInterval,

		queued: make(map[string]*work),
		queue:  workqueue.NewNamedDelayingQueue(name),
		once:   make(chan struct{}),
	}
}

func (l *WriterLease) Run(stopCh <-chan struct{}) {
	defer utilruntime.HandleCrash()
	defer l.queue.ShutDown()

	go func() {
		defer utilruntime.HandleCrash()
		for l.work() {
		}
		glog.V(4).Infof("[%s] Worker stopped", l.name)
	}()

	<-stopCh
}

func (l *WriterLease) Expire() {
	l.lock.Lock()
	defer l.lock.Unlock()
	l.expires = time.Time{}
}

func (l *WriterLease) Wait() bool {
	<-l.once
	state, _, _ := l.leaseState()
	return state == Leader
}

func (l *WriterLease) WaitUntil(t time.Duration) (bool, bool) {
	select {
	case <-l.once:
	case <-time.After(t):
		return false, false
	}
	state, _, _ := l.leaseState()
	return state == Leader, true
}

func (l *WriterLease) Try(key string, fn WorkFunc) {
	l.lock.Lock()
	defer l.lock.Unlock()
	l.id++
	l.queued[key] = &work{fn: fn, id: l.id}
	if l.state == Follower {
		delay := l.expires.Sub(nowFn())
		// no matter what, always wait at least some amount of time as a follower to give the nominal
		// leader a chance to win
		if delay < l.backoff.Duration*2 {
			delay = l.backoff.Duration * 2
		}
		l.queue.AddAfter(key, delay)
	} else {
		l.queue.Add(key)
	}
}

func (l *WriterLease) Extend(key string) {
	l.lock.Lock()
	defer l.lock.Unlock()
	if _, ok := l.queued[key]; ok {
		delete(l.queued, key)
		switch l.state {
		case Follower:
			l.tick++
			backoff := l.nextBackoff()
			glog.V(4).Infof("[%s] Clearing work for %s and extending lease by %s", l.name, key, backoff)
			l.expires = nowFn().Add(backoff)
		}
	}
}

func (l *WriterLease) Len() int {
	l.lock.Lock()
	defer l.lock.Unlock()
	return len(l.queued)
}

func (l *WriterLease) Remove(key string) {
	l.lock.Lock()
	defer l.lock.Unlock()
	delete(l.queued, key)
}

func (l *WriterLease) get(key string) *work {
	l.lock.Lock()
	defer l.lock.Unlock()
	return l.queued[key]
}

func (l *WriterLease) leaseState() (State, time.Time, int) {
	l.lock.Lock()
	defer l.lock.Unlock()
	return l.state, l.expires, l.tick
}

func (l *WriterLease) work() bool {
	item, shutdown := l.queue.Get()
	if shutdown {
		return false
	}
	key := item.(string)

	work := l.get(key)
	if work == nil {
		glog.V(4).Infof("[%s] Work item %s was cleared, done", l.name, key)
		l.queue.Done(key)
		return true
	}

	leaseState, leaseExpires, _ := l.leaseState()
	if leaseState == Follower {
		// if we are following, continue to defer work until the lease expires
		if remaining := leaseExpires.Sub(nowFn()); remaining > 0 {
			glog.V(4).Infof("[%s] Follower, %s remaining in lease", l.name, remaining)
			l.queue.AddAfter(key, remaining)
			l.queue.Done(key)
			return true
		}
		glog.V(4).Infof("[%s] Lease expired, running %s", l.name, key)
	} else {
		glog.V(4).Infof("[%s] Lease owner or electing, running %s", l.name, key)
	}

	isLeader, retry := work.fn()
	if retry {
		// come back in a bit
		glog.V(4).Infof("[%s] Retrying %s", l.name, key)
		l.queue.AddAfter(key, l.retryInterval)
		l.queue.Done(key)
		return true
	}

	l.finishKey(key, isLeader, work.id)
	return true
}

func (l *WriterLease) finishKey(key string, isLeader bool, id int) {
	l.lock.Lock()
	defer l.lock.Unlock()

	resolvedElection := l.state == Election
	if isLeader {
		switch l.state {
		case Election, Follower:
			l.tick = 0
			l.state = Leader
		}
		l.expires = nowFn().Add(l.maxBackoff)
	} else {
		switch l.state {
		case Election, Leader:
			l.tick = 0
			l.state = Follower
		case Follower:
			l.tick++
		}
		l.expires = nowFn().Add(l.nextBackoff())
	}
	if work, ok := l.queued[key]; ok && work.id == id {
		delete(l.queued, key)
	}
	// close the channel before we remove the key from the queue to prevent races in Wait
	if resolvedElection {
		close(l.once)
	}
	l.queue.Done(key)
	glog.V(4).Infof("[%s] Completed work for %s in state=%d tick=%d expires=%s", l.name, key, l.state, l.tick, l.expires)
}

func (l *WriterLease) nextBackoff() time.Duration {
	step := l.tick
	b := l.backoff
	if step > b.Steps {
		return l.maxBackoff
	}
	duration := b.Duration
	for i := 0; i < step; i++ {
		adjusted := duration
		if b.Jitter > 0.0 {
			adjusted = wait.Jitter(duration, b.Jitter)
		}
		duration = time.Duration(float64(adjusted) * b.Factor)
		if duration > l.maxBackoff {
			return l.maxBackoff
		}
	}
	return duration
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
)

// Game map.
// Coordinates formatted as:
// 0123
// 1
// 2
// 3
type GameMap struct {
	cellMap      [][]bool
	changeEvents []ChangeEvent
}

func (gameMap *GameMap) GetSize() (int, int) {
	height := len(gameMap.cellMap)
	if height == 0 {
		panic("Height must not be 0")
	}

	width := len(gameMap.cellMap[0])
	if width == 0 {
		panic("Width must not be 0")
	}

	return width, height
}

func (gameMap *GameMap) SetValue(width, height int, value bool) {
	gameMap.cellMap[height][width] = value
}

func (gameMap *GameMap) GetValue(width, height int) bool {
	return gameMap.cellMap[height][width]
}

//simple linear algorithm
func (gameMap *GameMap) Update() {
	width, height := gameMap.GetSize()

	for i := 0; i < width; i++ {
		for j := 0; j < height; j++ {
			aliveNeighboursCount := gameMap.getAliveNeighboursCountEff(i, j)

			cellAlive := gameMap.GetValue(i, j)
			if cellAlive {
				if !cellContinueToLive(aliveNeighboursCount) {
					gameMap.addChange(i, j, false)
				}
			} else {
				if aliveNeighboursCount == 3 {
					gameMap.addChange(i, j, true)
				}
			}
		}
	}

	gameMap.performAllChanges()
}

func (gameMap *GameMap) performAllChanges() {
	for _, changeEvent := range gameMap.changeEvents {
		gameMap.SetValue(changeEvent.width, changeEvent.height, changeEvent.newValue)
	}
	gameMap.changeEvents = nil
}

func (gameMap *GameMap) addChange(width, height int, newValue bool) {
	//TODO concurrent safety to be added
	changeEvents := gameMap.changeEvents
	changeEvents = append(changeEvents, ChangeEvent{width, height, newValue})
	gameMap.changeEvents = changeEvents
}

func cellContinueToLive(aliveNeighboursCount int) bool {
	return aliveNeighboursCount == 2 || aliveNeighboursCount == 3
}

func (gameMap *GameMap) getAliveNeighboursCountEff(width, height int) int {
	aliveNeighbours := 0

	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width - 1, height - 1))
	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width - 1, height))
	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width - 1, height + 1))

	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width, height - 1))
	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width, height + 1))

	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width + 1, height - 1))
	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width + 1, height))
	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width + 1, height + 1))

	return aliveNeighbours
}

func oneIfTrue(is bool) int {
	if is {
		return 1
	} else {
		return 0
	}
}

func (gameMap *GameMap) isNeighbourAlive(width, height int) bool {
	mapWidth, mapHeight := gameMap.GetSize()

	possibleHeight := checkPossibleValue(mapHeight, height)
	possibleWidth := checkPossibleValue(mapWidth, width)

	if possibleHeight && possibleWidth {
		return gameMap.GetValue(width, height)
	}

	return false
}

//obsolete code
func (gameMap *GameMap) getAliveNeighboursCount(width, height int) int {
	neighbours := gameMap.getCellNeighbours(width, height)

	aliveNeighboursCount := 0

	for _, cellAlive := range neighbours {
		if cellAlive {
			aliveNeighboursCount++
		}
	}

	return aliveNeighboursCount
}

//obsolete code
func (gameMap *GameMap) getCellNeighbours(width, height int) []bool {
	var neighbours []bool

	neighbours = gameMap.addNeighbour(neighbours, width - 1, height - 1)
	neighbours = gameMap.addNeighbour(neighbours, width - 1, height)
	neighbours = gameMap.addNeighbour(neighbours, width - 1, height + 1)

	neighbours = gameMap.addNeighbour(neighbours, width, height - 1)
	neighbours = gameMap.addNeighbour(neighbours, width, height + 1)

	neighbours = gameMap.addNeighbour(neighbours, width + 1, height - 1)
	neighbours = gameMap.addNeighbour(neighbours, width + 1, height)
	neighbours = gameMap.addNeighbour(neighbours, width + 1, height + 1)

	return neighbours
}

//obsolete code
func (gameMap *GameMap) addNeighbour(neighbours []bool, width, height int) []bool {
	mapWidth, mapHeight := gameMap.GetSize()

	possibleHeight := checkPossibleValue(mapHeight, height)
	possibleWidth := checkPossibleValue(mapWidth, width)

	if possibleHeight && possibleWidth {
		neighbours = append(neighbours, gameMap.GetValue(width, height)) //TODO not efficient copying of array
	}

	return neighbours
}

func checkPossibleValue(basicValue int, checkValue int) bool {
	return checkValue >= 0 && checkValue < basicValue
}

func (gameMap *GameMap) DoForEveryCell(f func(width, height int, gameMap *GameMap)) {
	width, height := gameMap.GetSize()

	for i := 0; i < width; i++ {
		for j := 0; j < height; j++ {
			f(i, j, gameMap)
		}
	}
}

func (gameMap *GameMap) String() string {
	var buffer bytes.Buffer

	width, height := gameMap.GetSize()

	for i := 0; i < width; i++ {
		for j := 0; j < height; j++ {
			cellAlive := gameMap.GetValue(i, j)
			if cellAlive {
				buffer.WriteString("1")
			} else {
				buffer.WriteString("0")
			}
		}
		buffer.WriteByte(13)
		buffer.WriteByte(10)
	}

	return buffer.String()
}

type ChangeEvent struct {
	width, height int
	newValue      bool
}<commit_msg>Initial concurrent commit - doesn't work<commit_after>package main

import (
	"bytes"
	"sync"
	"runtime"
)

var WORKERS_COUNT int = 0

// Game map.
// Coordinates formatted as:
// 0123
// 1
// 2
// 3
type GameMap struct {
	cellMap      [][]bool
	changeEvents []*ChangeEvent
}

func (gameMap *GameMap) GetSize() (int, int) {
	height := len(gameMap.cellMap)
	if height == 0 {
		panic("Height must not be 0")
	}

	width := len(gameMap.cellMap[0])
	if width == 0 {
		panic("Width must not be 0")
	}

	return width, height
}

func (gameMap *GameMap) SetValue(width, height int, value bool) {
	gameMap.cellMap[height][width] = value
}

func (gameMap *GameMap) GetValue(width, height int) bool {
	return gameMap.cellMap[height][width]
}

//simple linear algorithm
func (gameMap *GameMap) Update() {
	_, height := gameMap.GetSize()

	rowsChan := make(chan int, height)
	changesChan := make(chan *ChangeEvent)
	finished := make(chan bool)

	var wg sync.WaitGroup
	wg.Add(height) //every row to be handled
	go waitForWorkFinished(finished, &wg)

	createRowWorkers(gameMap, rowsChan, changesChan, &wg)
	go addRowsWorks(height, rowsChan)

	select {
	case changeEvent := <-changesChan:
		gameMap.addChange(changeEvent)
	case <-finished:
		break
	}

	gameMap.performAllChanges()
}

func waitForWorkFinished(finished chan bool, wg *sync.WaitGroup) {
	wg.Wait()
	finished <- true
}

func addRowsWorks(height int, rowsChan chan int) {
	for rowNum := 0; rowNum < height; rowNum++ {
		rowsChan <- rowNum
	}
}

func createRowWorkers(gameMap *GameMap, rowsChan chan int, changesChan chan *ChangeEvent, wg *sync.WaitGroup) {
	for i := 0; i < getWorkersCount(); i++ {
		go rowWorker(gameMap, rowsChan, changesChan, wg)
	}
}

func getWorkersCount() int {
	if WORKERS_COUNT == 0 {
		WORKERS_COUNT = runtime.NumCPU()
	}

	return WORKERS_COUNT
}

func rowWorker(gameMap *GameMap, rowsChan chan int, changesChan chan *ChangeEvent, wg *sync.WaitGroup) {
	for rowNumber := range rowsChan {
		handleRow(gameMap, rowNumber, changesChan)
		wg.Done()
	}
}

func handleRow(gameMap *GameMap, rowNum int, changesChan chan *ChangeEvent) {
	row := gameMap.cellMap[rowNum]
	rowLength := len(row)

	for xCoordinate := 0; xCoordinate < rowLength; xCoordinate++ {
		handleCell(gameMap, xCoordinate, rowNum, changesChan)
	}
}

func handleCell(gameMap *GameMap, xCoordinate int, rowNum int, changesChan chan *ChangeEvent) {
	aliveNeighboursCount := gameMap.getAliveNeighboursCountEff(xCoordinate, rowNum)

	cellAlive := gameMap.GetValue(xCoordinate, rowNum)
	if cellAlive {
		if !cellContinueToLive(aliveNeighboursCount) {
			changesChan <- &ChangeEvent{xCoordinate, rowNum, false}
		}
	} else {
		if aliveNeighboursCount == 3 {
			changesChan <- &ChangeEvent{xCoordinate, rowNum, true}
		}
	}
}

func (gameMap *GameMap) performAllChanges() {
	for _, changeEvent := range gameMap.changeEvents {
		gameMap.SetValue(changeEvent.width, changeEvent.height, changeEvent.newValue)
	}
	gameMap.changeEvents = nil
}

func (gameMap *GameMap) addChange(changeEvent *ChangeEvent) {
	//TODO concurrent safety to be added
	changeEvents := gameMap.changeEvents
	changeEvents = append(changeEvents, changeEvent)
	gameMap.changeEvents = changeEvents
}

func cellContinueToLive(aliveNeighboursCount int) bool {
	return aliveNeighboursCount == 2 || aliveNeighboursCount == 3
}

func (gameMap *GameMap) getAliveNeighboursCountEff(width, height int) int {
	aliveNeighbours := 0

	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width - 1, height - 1))
	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width - 1, height))
	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width - 1, height + 1))

	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width, height - 1))
	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width, height + 1))

	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width + 1, height - 1))
	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width + 1, height))
	aliveNeighbours += oneIfTrue(gameMap.isNeighbourAlive(width + 1, height + 1))

	return aliveNeighbours
}

func oneIfTrue(is bool) int {
	if is {
		return 1
	} else {
		return 0
	}
}

func (gameMap *GameMap) isNeighbourAlive(width, height int) bool {
	mapWidth, mapHeight := gameMap.GetSize()

	possibleHeight := checkPossibleValue(mapHeight, height)
	possibleWidth := checkPossibleValue(mapWidth, width)

	if possibleHeight && possibleWidth {
		return gameMap.GetValue(width, height)
	}

	return false
}

//obsolete code
func (gameMap *GameMap) getAliveNeighboursCount(width, height int) int {
	neighbours := gameMap.getCellNeighbours(width, height)

	aliveNeighboursCount := 0

	for _, cellAlive := range neighbours {
		if cellAlive {
			aliveNeighboursCount++
		}
	}

	return aliveNeighboursCount
}

//obsolete code
func (gameMap *GameMap) getCellNeighbours(width, height int) []bool {
	var neighbours []bool

	neighbours = gameMap.addNeighbour(neighbours, width - 1, height - 1)
	neighbours = gameMap.addNeighbour(neighbours, width - 1, height)
	neighbours = gameMap.addNeighbour(neighbours, width - 1, height + 1)

	neighbours = gameMap.addNeighbour(neighbours, width, height - 1)
	neighbours = gameMap.addNeighbour(neighbours, width, height + 1)

	neighbours = gameMap.addNeighbour(neighbours, width + 1, height - 1)
	neighbours = gameMap.addNeighbour(neighbours, width + 1, height)
	neighbours = gameMap.addNeighbour(neighbours, width + 1, height + 1)

	return neighbours
}

//obsolete code
func (gameMap *GameMap) addNeighbour(neighbours []bool, width, height int) []bool {
	mapWidth, mapHeight := gameMap.GetSize()

	possibleHeight := checkPossibleValue(mapHeight, height)
	possibleWidth := checkPossibleValue(mapWidth, width)

	if possibleHeight && possibleWidth {
		neighbours = append(neighbours, gameMap.GetValue(width, height)) //TODO not efficient copying of array
	}

	return neighbours
}

func checkPossibleValue(basicValue int, checkValue int) bool {
	return checkValue >= 0 && checkValue < basicValue
}

func (gameMap *GameMap) DoForEveryCell(f func(width, height int, gameMap *GameMap)) {
	width, height := gameMap.GetSize()

	for i := 0; i < width; i++ {
		for j := 0; j < height; j++ {
			f(i, j, gameMap)
		}
	}
}

func (gameMap *GameMap) String() string {
	var buffer bytes.Buffer

	width, height := gameMap.GetSize()

	for i := 0; i < width; i++ {
		for j := 0; j < height; j++ {
			cellAlive := gameMap.GetValue(i, j)
			if cellAlive {
				buffer.WriteString("1")
			} else {
				buffer.WriteString("0")
			}
		}
		buffer.WriteByte(13)
		buffer.WriteByte(10)
	}

	return buffer.String()
}

type ChangeEvent struct {
	width, height int
	newValue      bool
}<|endoftext|>


<commit_before>/*
 Copyright 2018 Padduck, LLC

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

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

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

package environments

import (
	"fmt"
	"github.com/pkg/errors"
	"github.com/pufferpanel/apufferi/common"
	"github.com/pufferpanel/pufferd/cache"
	"github.com/pufferpanel/pufferd/utils"
)

var mapping map[string]EnvironmentFactory

func LoadModules() {
	mapping := make(map[string]EnvironmentFactory)

	mapping["standard"] = StandardFactory{}

	loadAdditionalModules(mapping)
}

func Create(environmentType, folder, id string, environmentSection map[string]interface{}) (Environment, error) {
	factory := mapping[environmentType]

	if factory == nil {
		return nil, errors.New(fmt.Sprintf("undefined environment: %s", environmentType))
	}

	serverRoot := common.JoinPath(folder, id)
	rootDirectory := common.GetStringOrDefault(environmentSection, "root", serverRoot)
	cache := cache.CreateCache()
	wsManager := utils.CreateWSManager()

	env := factory.Create(folder, id, environmentSection, rootDirectory, cache, wsManager)

	return env, nil
}

func GetSupportedEnvironments() []string {
	result := make([]string, len(mapping))
	i := 0
	for _, v := range mapping {
		result[i] = v.Key()
		i++;
	}
}<commit_msg>Fix compiler<commit_after>/*
 Copyright 2018 Padduck, LLC

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

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

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

package environments

import (
	"fmt"
	"github.com/pkg/errors"
	"github.com/pufferpanel/apufferi/common"
	"github.com/pufferpanel/pufferd/cache"
	"github.com/pufferpanel/pufferd/utils"
)

var mapping map[string]EnvironmentFactory

func LoadModules() {
	mapping := make(map[string]EnvironmentFactory)

	mapping["standard"] = StandardFactory{}

	loadAdditionalModules(mapping)
}

func Create(environmentType, folder, id string, environmentSection map[string]interface{}) (Environment, error) {
	factory := mapping[environmentType]

	if factory == nil {
		return nil, errors.New(fmt.Sprintf("undefined environment: %s", environmentType))
	}

	serverRoot := common.JoinPath(folder, id)
	rootDirectory := common.GetStringOrDefault(environmentSection, "root", serverRoot)
	cache := cache.CreateCache()
	wsManager := utils.CreateWSManager()

	env := factory.Create(folder, id, environmentSection, rootDirectory, cache, wsManager)

	return env, nil
}

func GetSupportedEnvironments() []string {
	result := make([]string, len(mapping))
	i := 0
	for _, v := range mapping {
		result[i] = v.Key()
		i++
	}

	return result
}<|endoftext|>


<commit_before>package create

import (
	"context"
	"fmt"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"github.com/gopasspw/gopass/internal/backend"
	"github.com/gopasspw/gopass/internal/cui"
	"github.com/gopasspw/gopass/internal/out"
	"github.com/gopasspw/gopass/internal/set"
	"github.com/gopasspw/gopass/internal/store/root"
	"github.com/gopasspw/gopass/pkg/ctxutil"
	"github.com/gopasspw/gopass/pkg/debug"
	"github.com/gopasspw/gopass/pkg/fsutil"
	"github.com/gopasspw/gopass/pkg/gopass/secrets"
	"github.com/gopasspw/gopass/pkg/pwgen"
	"github.com/gopasspw/gopass/pkg/pwgen/pwrules"
	"github.com/gopasspw/gopass/pkg/termio"
	"github.com/martinhoefling/goxkcdpwgen/xkcdpwgen"
	"github.com/urfave/cli/v2"
	"gopkg.in/yaml.v3"
)

const (
	defaultLength     = 24
	defaultXKCDLength = 4
)

// Attribute is a credential attribute that is being asked for
// when populating a template.
type Attribute struct {
	Name    string `yaml:"name"`
	Type    string `yaml:"type"`
	Prompt  string `yaml:"prompt"`
	Charset string `yaml:"charset"`
	Min     int    `yaml:"min"`
	Max     int    `yaml:"max"`
}

// Template is an action template for the create wizard.
type Template struct {
	Name       string      `yaml:"name"`
	Priority   int         `yaml:"priority"`
	Prefix     string      `yaml:"prefix"`
	NameFrom   []string    `yaml:"name_from"`
	Welcome    string      `yaml:"welcome"`
	Attributes []Attribute `yaml:"attributes"`
}

// Wizard is the templateable credential creation wizard.
type Wizard struct {
	Templates []Template
}

// New creates a new instance of the wizard. It will parse the user
// supplied templates and add the default templates.
func New(ctx context.Context, s backend.Storage) (*Wizard, error) {
	w := &Wizard{
		Templates: []Template{
			{
				Name:     "Website login",
				Priority: 0,
				Prefix:   "websites",
				NameFrom: []string{"url", "username"},
				Welcome:  "🧪 Creating Website login",
				Attributes: []Attribute{
					{
						Name:   "url",
						Type:   "hostname",
						Prompt: "Website URL",
						Min:    1,
						Max:    255,
					},
					{
						Name:   "username",
						Type:   "string",
						Prompt: "Login",
						Min:    1,
					},
					{
						Name:   "password",
						Type:   "password",
						Prompt: "Password for the Website",
					},
				},
			},
			{
				Name:     "PIN Code (numerical)",
				Priority: 1,
				Prefix:   "pins",
				NameFrom: []string{
					"authority",
					"application",
				},
				Welcome: "🔑 Creating PIN Code",
				Attributes: []Attribute{
					{
						Name:   "authority",
						Type:   "string",
						Prompt: "Authority (Issuer)",
						Min:    1,
					},
					{
						Name:   "application",
						Type:   "string",
						Prompt: "Entity (e.g. debit, credit card, etc.)",
						Min:    1,
					},
					{
						Name:    "password",
						Type:    "password",
						Prompt:  "PIN Code",
						Min:     1,
						Max:     64,
						Charset: "0123456789",
					},
					{
						Name:   "comment",
						Type:   "string",
						Prompt: "Comment",
					},
				},
			},
		},
	}
	tpls, err := s.List(ctx, ".gopass/create/")
	if err != nil {
		return nil, err
	}

	for _, f := range tpls {
		if !strings.HasSuffix(f, ".yml") && !strings.HasSuffix(f, ".yaml") {
			debug.Log("ignoring unknown file extension: %s", f)

			continue
		}
		buf, err := s.Get(ctx, f)
		if err != nil {
			debug.Log("failed to parse template %s: %s", f, err)

			continue
		}
		tpl := Template{}
		if err := yaml.Unmarshal(buf, &tpl); err != nil {
			debug.Log("failed to parse template %s: %s", f, err)
			out.Errorf(ctx, "Bad template %s: %s\n%s", f, err, string(buf))

			continue
		}

		w.Templates = append(w.Templates, tpl)
	}

	sort.Slice(w.Templates, func(i, j int) bool {
		return w.Templates[i].Priority < w.Templates[j].Priority
	})

	return w, nil
}

// ActionCallback is the callback for the creation calls to print and copy the credentials.
type ActionCallback func(context.Context, *cli.Context, string, string, bool) error

// Actions returns a list of actions that can be performed on the wizard. The actions directly
// interact with the underlying storage.
func (w *Wizard) Actions(s *root.Store, cb ActionCallback) cui.Actions {
	sort.Slice(w.Templates, func(i, j int) bool {
		return w.Templates[i].Priority < w.Templates[j].Priority
	})

	acts := make(cui.Actions, 0, len(w.Templates))
	for _, tpl := range w.Templates {
		acts = append(acts, cui.Action{
			Name: tpl.Name,
			Fn:   mkActFunc(tpl, s, cb),
		})
	}

	return acts
}

func mkActFunc(tpl Template, s *root.Store, cb ActionCallback) func(context.Context, *cli.Context) error { //nolint:cyclop
	debug.Log("creating action func for %+v, cb: %p", tpl, cb)

	return func(ctx context.Context, c *cli.Context) error {
		name := c.Args().First()
		store := c.String("store")
		force := c.Bool("force")

		sec := secrets.New()

		out.Print(ctx, tpl.Welcome)

		// genPW is needed for the callback
		var genPw bool
		// password is needed for the callback
		var password string
		// hostname is needed in later iterations (e.g. password rule lookup)
		var hostname string
		// wantForName is a list of attributes that will be used to build the name
		wantForName := set.Map(tpl.NameFrom)
		// nameParts are the components the name will be built from
		var nameParts []string
		// step is only used for printing the progress
		var step int
		for _, v := range tpl.Attributes {
			step++
			k := v.Name

			// if no prompt is set default to the key
			if v.Prompt == "" {
				v.Prompt = strings.ToTitle(k)
			}

			switch v.Type {
			case "string":
				sv, err := termio.AskForString(ctx, fmtfn(2, strconv.Itoa(step), v.Prompt), "")
				if err != nil {
					return err
				}
				if wantForName[k] {
					nameParts = append(nameParts, sv)
				}
				_ = sec.Set(k, sv)
			case "hostname":
				sv, err := termio.AskForString(ctx, fmtfn(2, strconv.Itoa(step), v.Prompt), "")
				if err != nil {
					return err
				}
				hostname = extractHostname(sv)
				if hostname == "" {
					return fmt.Errorf("can not parse URL %s", sv)
				}
				if wantForName[k] {
					nameParts = append(nameParts, hostname)
				}
				if u := pwrules.LookupChangeURL(hostname); u != "" {
					_ = sec.Set("password-change-url", u)
				}
				_ = sec.Set(k, sv)
			case "password":
				var err error
				genPw, err = termio.AskForBool(ctx, fmtfn(2, strconv.Itoa(step), "Generate Password?"), true)
				if err != nil {
					return err
				}

				if genPw {
					password, err = generatePassword(ctx, hostname, v.Charset)
					if err != nil {
						return err
					}
				} else {
					password, err = termio.AskForPassword(ctx, v.Prompt, true)
					if err != nil {
						return err
					}
				}

				sec.SetPassword(password)
			}
		}

		// select store.
		if store == "" {
			store = cui.AskForStore(ctx, s)
		}

		// now we can generate a name. If it's already take we can the user for an alternative
		// name.

		// make sure the store is properly separated from the name.
		if store != "" {
			store += "/"
		}

		// by default create will generate a name for the secret based on the user
		// input. Only when the force flag is given it will accept a secrets path
		// as the first argument.
		if name == "" || !force {
			for i, s := range nameParts {
				nameParts[i] = fsutil.CleanFilename(s)
			}
			name = fmt.Sprintf("%s%s/%s", store, tpl.Prefix, filepath.Join(nameParts...))
		}
		if force && !strings.HasPrefix(name, store) {
			out.Warningf(ctx, "User supplied secret name %q does not match requested mount %q. Ignoring store flag.", name, store)
		}

		// force will also override the check for existing entries.
		if s.Exists(ctx, name) && !force {
			step++
			var err error
			name, err = termio.AskForString(ctx, fmtfn(2, strconv.Itoa(step), "Secret already exists. Choose another path or enter to overwrite"), name)
			if err != nil {
				return err
			}
		}

		if err := s.Set(ctxutil.WithCommitMessage(ctx, "Created new entry"), name, sec); err != nil {
			return fmt.Errorf("failed to set %q: %w", name, err)
		}
		out.OKf(ctx, "Credentials saved to %q", name)

		return cb(ctx, c, name, password, genPw)
	}
}

// generatePasssword will walk through the password generation steps.
func generatePassword(ctx context.Context, hostname, charset string) (string, error) {
	if charset != "" {
		length, err := termio.AskForInt(ctx, fmtfn(4, "a", "How long?"), 4)
		if err != nil {
			return "", err
		}

		return pwgen.GeneratePasswordCharset(length, charset), nil
	}
	if _, found := pwrules.LookupRule(hostname); found {
		out.Noticef(ctx, "Using password rules for %s ...", hostname)
		length, err := termio.AskForInt(ctx, fmtfn(4, "b", "How long?"), defaultLength)
		if err != nil {
			return "", err
		}

		return pwgen.NewCrypticForDomain(length, hostname).Password(), nil
	}
	xkcd, err := termio.AskForBool(ctx, fmtfn(4, "a", "Human-pronounceable passphrase?"), false)
	if err != nil {
		return "", err
	}
	if xkcd {
		length, err := termio.AskForInt(ctx, fmtfn(4, "b", "How many words?"), defaultXKCDLength)
		if err != nil {
			return "", err
		}
		g := xkcdpwgen.NewGenerator()
		g.SetNumWords(length)
		g.SetDelimiter(" ")
		g.SetCapitalize(true)

		return string(g.GeneratePassword()), nil
	}

	length, err := termio.AskForInt(ctx, fmtfn(4, "b", "How long?"), defaultLength)
	if err != nil {
		return "", err
	}

	symbols, err := termio.AskForBool(ctx, fmtfn(4, "c", "Include symbols?"), false)
	if err != nil {
		return "", err
	}

	corp, err := termio.AskForBool(ctx, fmtfn(4, "d", "Strict rules?"), false)
	if err != nil {
		return "", err
	}
	if corp {
		return pwgen.GeneratePasswordWithAllClasses(length, symbols)
	}

	return pwgen.GeneratePassword(length, symbols), nil
}
<commit_msg>Enforce min and max length for string and password types in the wizard (#2293)<commit_after>package create

import (
	"context"
	"fmt"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"github.com/gopasspw/gopass/internal/backend"
	"github.com/gopasspw/gopass/internal/cui"
	"github.com/gopasspw/gopass/internal/out"
	"github.com/gopasspw/gopass/internal/set"
	"github.com/gopasspw/gopass/internal/store/root"
	"github.com/gopasspw/gopass/pkg/ctxutil"
	"github.com/gopasspw/gopass/pkg/debug"
	"github.com/gopasspw/gopass/pkg/fsutil"
	"github.com/gopasspw/gopass/pkg/gopass/secrets"
	"github.com/gopasspw/gopass/pkg/pwgen"
	"github.com/gopasspw/gopass/pkg/pwgen/pwrules"
	"github.com/gopasspw/gopass/pkg/termio"
	"github.com/martinhoefling/goxkcdpwgen/xkcdpwgen"
	"github.com/urfave/cli/v2"
	"gopkg.in/yaml.v3"
)

const (
	defaultLength     = 24
	defaultXKCDLength = 4
)

// Attribute is a credential attribute that is being asked for
// when populating a template.
type Attribute struct {
	Name    string `yaml:"name"`
	Type    string `yaml:"type"`
	Prompt  string `yaml:"prompt"`
	Charset string `yaml:"charset"`
	Min     int    `yaml:"min"`
	Max     int    `yaml:"max"`
}

// Template is an action template for the create wizard.
type Template struct {
	Name       string      `yaml:"name"`
	Priority   int         `yaml:"priority"`
	Prefix     string      `yaml:"prefix"`
	NameFrom   []string    `yaml:"name_from"`
	Welcome    string      `yaml:"welcome"`
	Attributes []Attribute `yaml:"attributes"`
}

// Wizard is the templateable credential creation wizard.
type Wizard struct {
	Templates []Template
}

// New creates a new instance of the wizard. It will parse the user
// supplied templates and add the default templates.
func New(ctx context.Context, s backend.Storage) (*Wizard, error) {
	w := &Wizard{
		Templates: []Template{
			{
				Name:     "Website login",
				Priority: 0,
				Prefix:   "websites",
				NameFrom: []string{"url", "username"},
				Welcome:  "🧪 Creating Website login",
				Attributes: []Attribute{
					{
						Name:   "url",
						Type:   "hostname",
						Prompt: "Website URL",
						Min:    1,
						Max:    255,
					},
					{
						Name:   "username",
						Type:   "string",
						Prompt: "Login",
						Min:    1,
					},
					{
						Name:   "password",
						Type:   "password",
						Prompt: "Password for the Website",
					},
				},
			},
			{
				Name:     "PIN Code (numerical)",
				Priority: 1,
				Prefix:   "pins",
				NameFrom: []string{
					"authority",
					"application",
				},
				Welcome: "🔑 Creating PIN Code",
				Attributes: []Attribute{
					{
						Name:   "authority",
						Type:   "string",
						Prompt: "Authority (Issuer)",
						Min:    1,
					},
					{
						Name:   "application",
						Type:   "string",
						Prompt: "Entity (e.g. debit, credit card, etc.)",
						Min:    1,
					},
					{
						Name:    "password",
						Type:    "password",
						Prompt:  "PIN Code",
						Min:     1,
						Max:     64,
						Charset: "0123456789",
					},
					{
						Name:   "comment",
						Type:   "string",
						Prompt: "Comment",
					},
				},
			},
		},
	}
	tpls, err := s.List(ctx, ".gopass/create/")
	if err != nil {
		return nil, err
	}

	for _, f := range tpls {
		if !strings.HasSuffix(f, ".yml") && !strings.HasSuffix(f, ".yaml") {
			debug.Log("ignoring unknown file extension: %s", f)

			continue
		}
		buf, err := s.Get(ctx, f)
		if err != nil {
			debug.Log("failed to parse template %s: %s", f, err)

			continue
		}
		tpl := Template{}
		if err := yaml.Unmarshal(buf, &tpl); err != nil {
			debug.Log("failed to parse template %s: %s", f, err)
			out.Errorf(ctx, "Bad template %s: %s\n%s", f, err, string(buf))

			continue
		}

		w.Templates = append(w.Templates, tpl)
	}

	sort.Slice(w.Templates, func(i, j int) bool {
		return w.Templates[i].Priority < w.Templates[j].Priority
	})

	return w, nil
}

// ActionCallback is the callback for the creation calls to print and copy the credentials.
type ActionCallback func(context.Context, *cli.Context, string, string, bool) error

// Actions returns a list of actions that can be performed on the wizard. The actions directly
// interact with the underlying storage.
func (w *Wizard) Actions(s *root.Store, cb ActionCallback) cui.Actions {
	sort.Slice(w.Templates, func(i, j int) bool {
		return w.Templates[i].Priority < w.Templates[j].Priority
	})

	acts := make(cui.Actions, 0, len(w.Templates))
	for _, tpl := range w.Templates {
		acts = append(acts, cui.Action{
			Name: tpl.Name,
			Fn:   mkActFunc(tpl, s, cb),
		})
	}

	return acts
}

func mkActFunc(tpl Template, s *root.Store, cb ActionCallback) func(context.Context, *cli.Context) error { //nolint:cyclop
	debug.Log("creating action func for %+v, cb: %p", tpl, cb)

	return func(ctx context.Context, c *cli.Context) error {
		name := c.Args().First()
		store := c.String("store")
		force := c.Bool("force")

		sec := secrets.New()

		out.Print(ctx, tpl.Welcome)

		// genPW is needed for the callback
		var genPw bool
		// password is needed for the callback
		var password string
		// hostname is needed in later iterations (e.g. password rule lookup)
		var hostname string
		// wantForName is a list of attributes that will be used to build the name
		wantForName := set.Map(tpl.NameFrom)
		// nameParts are the components the name will be built from
		var nameParts []string
		// step is only used for printing the progress
		var step int
		for _, v := range tpl.Attributes {
			step++
			k := v.Name

			// if no prompt is set default to the key
			if v.Prompt == "" {
				v.Prompt = strings.ToTitle(k)
			}

			switch v.Type {
			case "string":
				sv, err := termio.AskForString(ctx, fmtfn(2, strconv.Itoa(step), v.Prompt), "")
				if err != nil {
					return err
				}
				if v.Min > 0 && len(sv) < v.Min {
					return fmt.Errorf("%s is too short (needs %d)", v.Name, v.Min)
				}
				if v.Max > 0 && len(sv) > v.Min {
					return fmt.Errorf("%s is too long (at most %d)", v.Name, v.Max)
				}
				if wantForName[k] {
					nameParts = append(nameParts, sv)
				}
				_ = sec.Set(k, sv)
			case "hostname":
				sv, err := termio.AskForString(ctx, fmtfn(2, strconv.Itoa(step), v.Prompt), "")
				if err != nil {
					return err
				}
				hostname = extractHostname(sv)
				if hostname == "" {
					return fmt.Errorf("can not parse URL %s", sv)
				}
				if wantForName[k] {
					nameParts = append(nameParts, hostname)
				}
				if u := pwrules.LookupChangeURL(hostname); u != "" {
					_ = sec.Set("password-change-url", u)
				}
				_ = sec.Set(k, sv)
			case "password":
				var err error
				genPw, err = termio.AskForBool(ctx, fmtfn(2, strconv.Itoa(step), "Generate Password?"), true)
				if err != nil {
					return err
				}

				if genPw {
					password, err = generatePassword(ctx, hostname, v.Charset)
					if err != nil {
						return err
					}
				} else {
					password, err = termio.AskForPassword(ctx, v.Prompt, true)
					if err != nil {
						return err
					}
					if v.Min > 0 && len(password) < v.Min {
						return fmt.Errorf("%s is too short (needs %d)", v.Name, v.Min)
					}
					if v.Max > 0 && len(password) > v.Min {
						return fmt.Errorf("%s is too long (at most %d)", v.Name, v.Max)
					}
				}

				sec.SetPassword(password)
			}
		}

		// select store.
		if store == "" {
			store = cui.AskForStore(ctx, s)
		}

		// now we can generate a name. If it's already take we can the user for an alternative
		// name.

		// make sure the store is properly separated from the name.
		if store != "" {
			store += "/"
		}

		// by default create will generate a name for the secret based on the user
		// input. Only when the force flag is given it will accept a secrets path
		// as the first argument.
		if name == "" || !force {
			for i, s := range nameParts {
				nameParts[i] = fsutil.CleanFilename(s)
			}
			name = fmt.Sprintf("%s%s/%s", store, tpl.Prefix, filepath.Join(nameParts...))
		}
		if force && !strings.HasPrefix(name, store) {
			out.Warningf(ctx, "User supplied secret name %q does not match requested mount %q. Ignoring store flag.", name, store)
		}

		// force will also override the check for existing entries.
		if s.Exists(ctx, name) && !force {
			step++
			var err error
			name, err = termio.AskForString(ctx, fmtfn(2, strconv.Itoa(step), "Secret already exists. Choose another path or enter to overwrite"), name)
			if err != nil {
				return err
			}
		}

		if err := s.Set(ctxutil.WithCommitMessage(ctx, "Created new entry"), name, sec); err != nil {
			return fmt.Errorf("failed to set %q: %w", name, err)
		}
		out.OKf(ctx, "Credentials saved to %q", name)

		return cb(ctx, c, name, password, genPw)
	}
}

// generatePasssword will walk through the password generation steps.
func generatePassword(ctx context.Context, hostname, charset string) (string, error) {
	if charset != "" {
		length, err := termio.AskForInt(ctx, fmtfn(4, "a", "How long?"), 4)
		if err != nil {
			return "", err
		}

		return pwgen.GeneratePasswordCharset(length, charset), nil
	}
	if _, found := pwrules.LookupRule(hostname); found {
		out.Noticef(ctx, "Using password rules for %s ...", hostname)
		length, err := termio.AskForInt(ctx, fmtfn(4, "b", "How long?"), defaultLength)
		if err != nil {
			return "", err
		}

		return pwgen.NewCrypticForDomain(length, hostname).Password(), nil
	}
	xkcd, err := termio.AskForBool(ctx, fmtfn(4, "a", "Human-pronounceable passphrase?"), false)
	if err != nil {
		return "", err
	}
	if xkcd {
		length, err := termio.AskForInt(ctx, fmtfn(4, "b", "How many words?"), defaultXKCDLength)
		if err != nil {
			return "", err
		}
		g := xkcdpwgen.NewGenerator()
		g.SetNumWords(length)
		g.SetDelimiter(" ")
		g.SetCapitalize(true)

		return string(g.GeneratePassword()), nil
	}

	length, err := termio.AskForInt(ctx, fmtfn(4, "b", "How long?"), defaultLength)
	if err != nil {
		return "", err
	}

	symbols, err := termio.AskForBool(ctx, fmtfn(4, "c", "Include symbols?"), false)
	if err != nil {
		return "", err
	}

	corp, err := termio.AskForBool(ctx, fmtfn(4, "d", "Strict rules?"), false)
	if err != nil {
		return "", err
	}
	if corp {
		return pwgen.GeneratePasswordWithAllClasses(length, symbols)
	}

	return pwgen.GeneratePassword(length, symbols), nil
}
<|endoftext|>


<commit_before>package forwarder

import (
	"encoding/json"
	"time"

	log "github.com/cihub/seelog"
	"github.com/golang/protobuf/proto"

	pb "github.com/bankpossible/iamdev/shared/messages"
)

var (
	debug = true
)

func Start(traceChan chan []byte, numWorkers, bufferSize int) {

	log.Infof("Starting %v forwarders with buffer size of %v", numWorkers, bufferSize)

	for i := 0; i < numWorkers; i++ {
		f := &forwarder{
			id:            i,
			ch:            traceChan,
			messageBuffer: make([][]byte, 0, bufferSize),
			bufferSize:    bufferSize,
		}

		// Do something useful
		go f.work()
	}

}

type forwarder struct {
	id            int
	ch            chan []byte
	messageBuffer [][]byte
	bufferSize    int
}

func (f *forwarder) work() {

	log.Debugf("[Forwarder %v] started", f.id)

	var b []byte
	var i int
	var decoded *pb.TraceFrame
	var js []byte

	metricsTick := time.NewTicker(5 * time.Second)
	timeoutTick := time.NewTicker(2 * time.Second)

	for {
		select {
		case b = <-f.ch:
			i++

			// Log the frame if we're in debug mode
			if debug {
				decoded = &pb.TraceFrame{}
				if err := proto.Unmarshal(b, decoded); err != nil {
					log.Infof("[Forwarder %v] Couldn't decode trace frame", f.id)
					continue
				}
				js, _ = json.Marshal(decoded)
				log.Infof("[Forwarder %v] Received message: %s", f.id, string(js))
			}

			// Add message to our buffer
			f.messageBuffer = append(f.messageBuffer, b)

			// Forward on if we're at our buffer size
			if len(f.messageBuffer) >= f.bufferSize {
				f.send()
			}
		case <-timeoutTick.C:
			f.send()
		case <-metricsTick.C:
			log.Debugf("[Forwarder %v] Processed %v messages", f.id, i)
		}
	}
}

func (f *forwarder) send() {

	log.Infof("[Forwarder %v] Sent %v messages", f.id, len(f.messageBuffer))

	// Empty the buffer
	f.messageBuffer = nil
}
<commit_msg>Lower some log levels<commit_after>package forwarder

import (
	"encoding/json"
	"time"

	log "github.com/cihub/seelog"
	"github.com/golang/protobuf/proto"

	pb "github.com/bankpossible/iamdev/shared/messages"
)

func Start(traceChan chan []byte, numWorkers, bufferSize int) {

	log.Infof("Starting %v forwarders with buffer size of %v", numWorkers, bufferSize)

	for i := 0; i < numWorkers; i++ {
		f := &forwarder{
			id:            i,
			ch:            traceChan,
			messageBuffer: make([][]byte, 0, bufferSize),
			bufferSize:    bufferSize,
		}

		// Do something useful
		go f.work()
	}

}

type forwarder struct {
	id            int
	ch            chan []byte
	messageBuffer [][]byte
	bufferSize    int
}

func (f *forwarder) work() {

	log.Debugf("[Forwarder %v] started", f.id)

	var b []byte
	var i int
	var decoded *pb.TraceFrame
	var js []byte

	metricsTick := time.NewTicker(5 * time.Second)
	timeoutTick := time.NewTicker(2 * time.Second)

	for {
		select {
		case b = <-f.ch:
			i++

			// Log the frame if we're in debug mode
			decoded = &pb.TraceFrame{}
			if err := proto.Unmarshal(b, decoded); err != nil {
				log.Warnf("[Forwarder %v] Couldn't decode trace frame", f.id)
				continue
			}
			js, _ = json.Marshal(decoded)
			log.Tracef("[Forwarder %v] Received message: %s", f.id, string(js))

			// Add message to our buffer
			f.messageBuffer = append(f.messageBuffer, b)

			// Forward on if we're at our buffer size
			if len(f.messageBuffer) >= f.bufferSize {
				f.send()
			}
		case <-timeoutTick.C:
			f.send()
		case <-metricsTick.C:
			log.Tracef("[Forwarder %v] Processed %v messages", f.id, i)
		}
	}
}

func (f *forwarder) send() {

	log.Debugf("[Forwarder %v] Sent %v messages", f.id, len(f.messageBuffer))

	// Empty the buffer
	f.messageBuffer = nil
}
<|endoftext|>


<commit_before>package errors

import (
	"errors"

	"zombiezen.com/go/capnproto2"
	"zombiezen.com/go/capnproto2/rpc"
	rpc_capnp "zombiezen.com/go/capnproto2/std/capnp/rpc"
)

var (
	NotImplemented = errors.New("Not implemented")
)

// Return an 'unimplemented' capnp exception allocated in s. If an
// error occurs in generating the exception, that error will be
// returned instead.
func UnImplementedExn(s *capnp.Segment) error {
	exn, err := rpc_capnp.NewException(s)
	if err != nil {
		return err
	}
	exn.SetType(rpc_capnp.Exception_Type_unimplemented)
	return rpc.Exception{exn}
}
<commit_msg>Slight refactoring of internal errors package.<commit_after>package errors

import (
	"errors"

	"zombiezen.com/go/capnproto2"
	"zombiezen.com/go/capnproto2/rpc"
	rpc_capnp "zombiezen.com/go/capnproto2/std/capnp/rpc"
)

var (
	NotImplemented = errors.New("Not implemented")
)

func newException(s *capnp.Segment, which rpc_capnp.Exception_Type) error {
	exn, err := rpc_capnp.NewException(s)
	if err != nil {
		return err
	}
	exn.SetType(which)
	return rpc.Exception{exn}
}

// Return an 'unimplemented' capnp exception allocated in s. If an
// error occurs in generating the exception, that error will be
// returned instead.
func UnImplementedExn(s *capnp.Segment) error {
	return newException(s, rpc_capnp.Exception_Type_unimplemented)
}
<|endoftext|>


<commit_before>package runner

import (
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	pb "github.com/ii/xds-test-harness/api/adapter"

	core "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
	discovery "github.com/envoyproxy/go-control-plane/envoy/service/discovery/v3"
	"github.com/rs/zerolog/log"
	"google.golang.org/grpc"
)

var (
	opts []grpc.DialOption = []grpc.DialOption{
		grpc.WithInsecure(),
		grpc.WithBlock(),
		grpc.WithTimeout(time.Second * 10),
	}
)

type ClientConfig struct {
	Port string
	Conn *grpc.ClientConn
}

type Cache struct {
	StartState     *pb.Snapshot
	StateSnapshots []*pb.Snapshot
	FinalResponse  *discovery.DiscoveryResponse
}

type Validate struct {
	RequestCount  int
	ResponseCount int
	Resources     map[string]map[string]string
}

type Runner struct {
	Adapter          *ClientConfig
	Target           *ClientConfig
	NodeID           string
	Cache            *Cache
	Aggregated       bool
	Service          *XDSService
	SubscribeRequest *discovery.DiscoveryRequest
	Validate         *Validate
}

func FreshRunner(current ...*Runner) *Runner {
	var (
		adapter    = &ClientConfig{}
		target     = &ClientConfig{}
		nodeID     = ""
		aggregated = false
	)

	if len(current) > 0 {
		adapter = current[0].Adapter
		target = current[0].Target
		nodeID = current[0].NodeID
		aggregated = current[0].Aggregated

	}

	return &Runner{
		Adapter:    adapter,
		Target:     target,
		NodeID:     nodeID,
		Cache:      &Cache{},
		Service:    &XDSService{},
		Aggregated: aggregated,
		Validate:   &Validate{},
	}
}

func (r *Runner) ConnectClient(server, address string) error {
	var client *ClientConfig
	if server == "target" {
		client = r.Target
	}
	if server == "adapter" {
		client = r.Adapter
	}
	if strings.HasPrefix(address, ":") {
		client.Port = address
	} else {
		client.Port = ":" + address
	}
	conn, err := connectViaGRPC(client, server)
	if err != nil {
		return err
	}
	client.Conn = conn
	return nil
}

func (r *Runner) Ack(service *XDSService) {
	service.Channels.Req <- r.SubscribeRequest
	service.Cache.Requests = append(service.Cache.Requests, r.SubscribeRequest)
	for {
		select {
		case res := <-service.Channels.Res:
			service.Cache.Responses = append(service.Cache.Responses, res)
			ack := newAckFromResponse(res, r.SubscribeRequest)
			log.Debug().
				Msgf("Sending Ack: %v", ack)
			service.Channels.Req <- ack
			service.Cache.Requests = append(service.Cache.Requests, ack)
		case <-service.Channels.Done:
			log.Debug().
				Msg("Received Done signal, shutting down request channel")
			close(service.Channels.Req)
			return
		}
	}
}

func (r *Runner) Stream(service *XDSService) error {
	defer service.Context.cancel()
	defer close(service.Channels.Err)

	var wg sync.WaitGroup
	go func() {
		for {
			wg.Add(1)
			in, err := service.Stream.Recv()
			if err == io.EOF {
				log.Debug().
					Msgf("No more Discovery Responses from %v stream", r.Service.Name)
				close(service.Channels.Res)
				return
			}
			if err != nil {
				log.Err(err).Msgf("error receiving responses on %v stream", r.Service.Name)
				service.Channels.Err <- err
				return
			}
			log.Debug().
				Msgf("Received discovery response: %v", in)
			r.Validate.ResponseCount++
			service.Channels.Res <- in
		}
	}()

	for req := range service.Channels.Req {
		if err := service.Stream.Send(req); err != nil {
			log.Err(err).
				Msg("Error sending discovery request")
			service.Channels.Err <- err
		}
		r.Validate.RequestCount++
	}
	service.Stream.CloseSend()
	wg.Wait()
	return nil
}

func connectViaGRPC(client *ClientConfig, server string) (conn *grpc.ClientConn, err error) {
	conn, err = grpc.Dial(client.Port, opts...)
	if err != nil {
		err = fmt.Errorf("Cannot connect at %v: %v\n", client.Port, err)
		return nil, err
	}
	log.Debug().
		Msgf("Runner connected to %v", server)
	return conn, nil
}

func newAckFromResponse(res *discovery.DiscoveryResponse, initReq *discovery.DiscoveryRequest) *discovery.DiscoveryRequest {
	// Only the first request should need the node ID,
	// so we do not include it in the followups.  If this
	// causes an error, it's a non-conformant error.
	request := &discovery.DiscoveryRequest{
		VersionInfo:   res.VersionInfo,
		ResourceNames: initReq.ResourceNames,
		TypeUrl:       initReq.TypeUrl,
		ResponseNonce: res.Nonce,
	}
	return request
}

func newRequest(resourceNames []string, typeURL, nodeID string) *discovery.DiscoveryRequest {
	return &discovery.DiscoveryRequest{
		VersionInfo: "",
		Node: &core.Node{
			Id: nodeID,
		},
		ResourceNames: resourceNames,
		TypeUrl:       typeURL,
	}
}
<commit_msg>small linting for standard ordering of imports<commit_after>package runner

import (
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	pb "github.com/ii/xds-test-harness/api/adapter"
	"github.com/ii/xds-test-harness/internal/parser"

	core "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
	discovery "github.com/envoyproxy/go-control-plane/envoy/service/discovery/v3"
	"github.com/rs/zerolog/log"
	"google.golang.org/grpc"
)

var (
	opts []grpc.DialOption = []grpc.DialOption{
		grpc.WithInsecure(),
		grpc.WithBlock(),
		grpc.WithTimeout(time.Second * 10),
	}
)

type ClientConfig struct {
	Port string
	Conn *grpc.ClientConn
}

type Cache struct {
	StartState     *pb.Snapshot
	StateSnapshots []*pb.Snapshot
	FinalResponse  *discovery.DiscoveryResponse
}

type Validate struct {
	RequestCount  int
	ResponseCount int
	Resources     map[string]map[string]string
}

type Runner struct {
	Adapter          *ClientConfig
	Target           *ClientConfig
	NodeID           string
	Cache            *Cache
	Aggregated       bool
	Service          *XDSService
	SubscribeRequest *discovery.DiscoveryRequest
	Validate         *Validate
}

func FreshRunner(current ...*Runner) *Runner {
	var (
		adapter    = &ClientConfig{}
		target     = &ClientConfig{}
		nodeID     = ""
		aggregated = false
	)

	if len(current) > 0 {
		adapter = current[0].Adapter
		target = current[0].Target
		nodeID = current[0].NodeID
		aggregated = current[0].Aggregated

	}

	return &Runner{
		Adapter:    adapter,
		Target:     target,
		NodeID:     nodeID,
		Cache:      &Cache{},
		Service:    &XDSService{},
		Aggregated: aggregated,
		Validate:   &Validate{},
	}
}

func (r *Runner) ConnectClient(server, address string) error {
	var client *ClientConfig
	if server == "target" {
		client = r.Target
	}
	if server == "adapter" {
		client = r.Adapter
	}
	if strings.HasPrefix(address, ":") {
		client.Port = address
	} else {
		client.Port = ":" + address
	}
	conn, err := connectViaGRPC(client, server)
	if err != nil {
		return err
	}
	client.Conn = conn
	return nil
}

func (r *Runner) Ack(service *XDSService) {
	service.Channels.Req <- r.SubscribeRequest
	service.Cache.Requests = append(service.Cache.Requests, r.SubscribeRequest)
	for {
		select {
		case res := <-service.Channels.Res:
			service.Cache.Responses = append(service.Cache.Responses, res)
			ack := newAckFromResponse(res, r.SubscribeRequest)
			log.Debug().
				Msgf("Sending Ack: %v", ack)
			service.Channels.Req <- ack
			service.Cache.Requests = append(service.Cache.Requests, ack)
		case <-service.Channels.Done:
			log.Debug().
				Msg("Received Done signal, shutting down request channel")
			close(service.Channels.Req)
			return
		}
	}
}

func (r *Runner) Stream(service *XDSService) error {
	defer service.Context.cancel()
	defer close(service.Channels.Err)

	var wg sync.WaitGroup
	go func() {
		for {
			wg.Add(1)
			in, err := service.Stream.Recv()
			if err == io.EOF {
				log.Debug().
					Msgf("No more Discovery Responses from %v stream", r.Service.Name)
				close(service.Channels.Res)
				return
			}
			if err != nil {
				log.Err(err).Msgf("error receiving responses on %v stream", r.Service.Name)
				service.Channels.Err <- err
				return
			}
			log.Debug().
				Msgf("Received discovery response: %v", in)
			r.Validate.ResponseCount++
			service.Channels.Res <- in
		}
	}()

	for req := range service.Channels.Req {
		if err := service.Stream.Send(req); err != nil {
			log.Err(err).
				Msg("Error sending discovery request")
			service.Channels.Err <- err
		}
		r.Validate.RequestCount++
	}
	service.Stream.CloseSend()
	wg.Wait()
	return nil
}

func connectViaGRPC(client *ClientConfig, server string) (conn *grpc.ClientConn, err error) {
	conn, err = grpc.Dial(client.Port, opts...)
	if err != nil {
		err = fmt.Errorf("Cannot connect at %v: %v\n", client.Port, err)
		return nil, err
	}
	log.Debug().
		Msgf("Runner connected to %v", server)
	return conn, nil
}

func newAckFromResponse(res *discovery.DiscoveryResponse, initReq *discovery.DiscoveryRequest) *discovery.DiscoveryRequest {
	// Only the first request should need the node ID,
	// so we do not include it in the followups.  If this
	// causes an error, it's a non-conformant error.
	request := &discovery.DiscoveryRequest{
		VersionInfo:   res.VersionInfo,
		ResourceNames: initReq.ResourceNames,
		TypeUrl:       initReq.TypeUrl,
		ResponseNonce: res.Nonce,
	}
	return request
}

func newRequest(resourceNames []string, typeURL, nodeID string) *discovery.DiscoveryRequest {
	return &discovery.DiscoveryRequest{
		VersionInfo: "",
		Node: &core.Node{
			Id: nodeID,
		},
		ResourceNames: resourceNames,
		TypeUrl:       typeURL,
	}
}
<|endoftext|>