<commit_before>package main

import (
	"io/ioutil"
	"os"
	"os/exec"
	"syscall"
	"testing"

	"github.com/stretchr/testify/suite"
)

var (
	eexist       = syscall.EEXIST.Error()
	einval       = syscall.EINVAL.Error()
	enametoolong = syscall.ENAMETOOLONG.Error()
	enoent       = syscall.ENOENT.Error()
)

const (
	// 257 * "z"
	longName = "zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz"
)

type internal struct {
	pool  string
	files []string
	suite.Suite
}

func TestSuiteInternal(t *testing.T) {
	suite.Run(t, &internal{})
}

func command(name string, arg ...string) *exec.Cmd {
	cmd := exec.Command(name, arg...)
	cmd.Stderr = os.Stderr
	return cmd
}

func (s *internal) create(pool string) {
	s.pool = pool
	files := make([]string, 5)
	for i := range files {
		f, err := ioutil.TempFile("", "gozfs-test-temp")
		if err != nil {
			panic(err)
		}
		files[i] = f.Name()
		f.Close()
	}
	s.files = files

	script := []byte(`
	set -e
	pool=$1
	shift
	zpool list $pool &>/dev/null && zpool destroy $pool
	files=($@)
	for f in ${files[*]}; do
		truncate -s1G $f
	done
	zpool create $pool ${files[*]}

	zfs create $pool/a
	zfs create $pool/a/1
	zfs create $pool/a/2
	zfs create $pool/a/4
	zfs snapshot $pool/a/1@snap1
	zfs snapshot $pool/a/2@snap1
	zfs snapshot $pool/a/2@snap2
	zfs snapshot $pool/a/2@snap3
	zfs clone $pool/a/1@snap1 $pool/a/3
	zfs hold hold1 $pool/a/2@snap1
	zfs hold hold2 $pool/a/2@snap2
	zfs unmount $pool/a/4

	zfs create $pool/b
	zfs create -V 8192 $pool/b/1
	zfs create -b 1024 -V 2048 $pool/b/2
	zfs create -V 8192 $pool/b/4
	zfs snapshot $pool/b/1@snap1
	zfs snapshot $pool/b/2@snap1
	zfs snapshot $pool/b/2@snap2
	zfs snapshot $pool/b/2@snap3
	zfs clone $pool/b/1@snap1 $pool/b/3
	zfs hold hold1 $pool/b/2@snap1
	zfs hold hold2 $pool/b/2@snap2

	zfs create $pool/c
	zfs create $pool/c/one
	zfs create $pool/c/two
	zfs create $pool/c/three
	zfs snapshot -r $pool/c@snap1

	exit 0
	`)

	args := make([]string, 3, 3+len(files))
	args[0] = "bash"
	args[1] = "/dev/stdin"
	args[2] = s.pool
	args = append(args, files...)
	cmd := command("sudo", args...)

	stdin, err := cmd.StdinPipe()
	s.Require().NoError(err)
	go func() {
		_, err := stdin.Write([]byte(script))
		s.Require().NoError(err)
	}()

	s.Require().NoError(cmd.Run())
}

func (s *internal) destroy() {
	err := command("sudo", "zpool", "destroy", s.pool).Run()
	for i := range s.files {
		os.Remove(s.files[i])
	}
	s.Require().NoError(err)
}

func (s *internal) SetupTest() {
	s.create("gozfs-test")
}

func (s *internal) TearDownTest() {
	s.destroy()
}

func (s *internal) TestClone() {
	s.EqualError(clone(s.pool+"/a/2", s.pool+"/a/1", nil), eexist)
	s.EqualError(clone(s.pool+"/a 3", s.pool+"/a/1", nil), einval)
	s.EqualError(clone(s.pool+"/a/"+longName, s.pool+"/a/1", nil), einval) // WANTE(ENAMETOOLONG)
	s.EqualError(clone(s.pool+"/a/z", s.pool+"/a/"+longName, nil), enametoolong)
	s.NoError(clone(s.pool+"/a/z", s.pool+"/a/1@snap1", nil))
}

func (s *internal) TestCreateFS() {
	s.EqualError(create("gozfs-test-no-exists/1", dmuZFS, nil), enoent)
	s.EqualError(create(s.pool+"/~1", dmuZFS, nil), einval)
	s.EqualError(create(s.pool+"/1", dmuNumtypes+1, nil), einval)
	s.EqualError(create(s.pool+"/1", dmuZFS, map[string]interface{}{"bad-prop": true}), einval)
	s.EqualError(create(s.pool+"/"+longName, dmuZFS, nil), einval) // WANTE(ENAMETOOLONG)
	s.NoError(create(s.pool+"/1", dmuZFS, nil))
	s.EqualError(create(s.pool+"/1", dmuZFS, nil), eexist)
	s.EqualError(create(s.pool+"/2/2", dmuZFS, nil), enoent)
}

func (s *internal) TestCreateVOL() {
	type p map[string]interface{}
	u1024 := uint64(1024)
	u8192 := u1024 * 8

	tests := []struct {
		name  string
		props p
		err   string
	}{
		{"gozfs-test-no-exists/1", nil, enoent},
		{s.pool + "/" + longName, nil, einval}, // WANTE(ENAMETOOLONG)
		{s.pool + "/~1", nil, einval},
		{s.pool + "/1", p{"bad-prop": true}, einval},
		{s.pool + "/2", p{"volsize": 0}, einval},
		{s.pool + "/3", p{"volsize": 1}, einval},
		{s.pool + "/4", p{"volsize": 8*1024 + 1}, einval},
		{s.pool + "/5", p{"volsize": 8 * 1024}, einval},
		{s.pool + "/6", p{"volsize": uint64(0)}, einval},
		{s.pool + "/7", p{"volsize": uint64(1)}, einval},
		{s.pool + "/8", p{"volsize": u8192 + 1}, einval},
		{s.pool + "/9", p{"volsize": u8192}, ""},
		{s.pool + "/9", p{"volsize": u8192}, eexist},
		{s.pool + "/10", p{"volsize": 0, "volblocksize": 1024}, einval},
		{s.pool + "/11", p{"volsize": 1, "volblocksize": 1024}, einval},
		{s.pool + "/12", p{"volsize": 1024 + 1, "volblocksize": 1024}, einval},
		{s.pool + "/13", p{"volsize": 1024, "volblocksize": 1024}, einval},
		{s.pool + "/14", p{"volsize": uint64(0), "volblocksize": 1024}, einval},
		{s.pool + "/15", p{"volsize": uint64(1), "volblocksize": 1024}, einval},
		{s.pool + "/16", p{"volsize": u1024 + 1, "volblocksize": 1024}, einval},
		{s.pool + "/17", p{"volsize": u1024, "volblocksize": 1024}, einval},
		{s.pool + "/18", p{"volsize": 0, "volblocksize": u1024}, einval},
		{s.pool + "/19", p{"volsize": 1, "volblocksize": u1024}, einval},
		{s.pool + "/20", p{"volsize": 1024 + 1, "volblocksize": u1024}, einval},
		{s.pool + "/21", p{"volsize": 1024, "volblocksize": u1024}, einval},
		{s.pool + "/22", p{"volsize": uint64(0), "volblocksize": u1024}, einval},
		{s.pool + "/23", p{"volsize": uint64(1), "volblocksize": u1024}, einval},
		{s.pool + "/24", p{"volsize": u1024 + 1, "volblocksize": u1024}, einval},
		{s.pool + "/25", p{"volsize": u1024, "volblocksize": u1024}, ""},
		{s.pool + "/25", p{"volsize": u1024, "volblocksize": u1024}, eexist},
		{s.pool + "/26/1", p{"volsize": u1024, "volblocksize": u1024}, enoent},
	}
	for _, test := range tests {
		err := create(test.name, dmuZVOL, test.props)
		if test.err == "" {
			s.NoError(err, "name: %v, props: %v",
				test.name, test.props)
		} else {
			s.EqualError(err, test.err, "name: %v, props: %v",
				test.name, test.props)
		}
	}

}

func (s *internal) TestListEmpty() {
	s.destroy()

	m, err := list("", nil, true, 0)
	s.NoError(err, "m: %v", m)
	s.Assert().Len(m, 0)

	s.create(s.pool)
}

func (s *internal) TestList() {
	type t map[string]bool
	tests := []struct {
		name    string
		types   t
		recurse bool
		depth   uint64
		expNum  int
		err     string
	}{
		{name: "blah", err: enoent},
		{name: s.pool + "/" + longName, err: einval}, // WANTE(ENAMETOOLONG)
		{name: s.pool, expNum: 1},
		{name: s.pool, recurse: true, expNum: 11},
		{name: s.pool, recurse: true, depth: 1, expNum: 4},
		{name: s.pool + "/a", recurse: true, expNum: 5},
		{name: s.pool + "/a", recurse: true, depth: 1, expNum: 5},
		{name: s.pool, types: t{"volume": true}, recurse: true, expNum: 4},
		{name: s.pool, types: t{"volume": true}, recurse: true, depth: 1, expNum: 0},
		{name: s.pool + "/a", types: t{"volume": true}, recurse: true, expNum: 0},
		{name: s.pool + "/b", types: t{"volume": true}, recurse: true, expNum: 4},
		{name: s.pool + "/b", types: t{"volume": true}, recurse: true, depth: 1, expNum: 4},
		{name: s.pool, types: t{"snapshot": true}, recurse: true, expNum: 12},
		{name: s.pool, types: t{"snapshot": true}, recurse: true, depth: 1, expNum: 0},
		{name: s.pool + "/a", types: t{"snapshot": true}, recurse: true, expNum: 4},
		{name: s.pool + "/a/1", types: t{"snapshot": true}, recurse: true, expNum: 1},
		{name: s.pool + "/b", types: t{"snapshot": true}, recurse: true, expNum: 4},
		{name: s.pool + "/b/1", types: t{"snapshot": true}, recurse: true, expNum: 1},
	}
	for i, test := range tests {
		m, err := list(test.name, test.types, test.recurse, test.depth)
		if test.err != "" {
			s.EqualError(err, test.err, "test num:%d", i)
			s.Nil(m, "test:%d", i)
		} else {
			s.NoError(err, "test:%d", i)
			s.Len(m, test.expNum, "test num:%d", i)
		}
	}
}
<commit_msg>gozfs: add tests for zfs_destroy<commit_after>package main

import (
	"io/ioutil"
	"os"
	"os/exec"
	"syscall"
	"testing"

	"github.com/stretchr/testify/suite"
)

var (
	ebusy        = syscall.EBUSY.Error()
	eexist       = syscall.EEXIST.Error()
	einval       = syscall.EINVAL.Error()
	enametoolong = syscall.ENAMETOOLONG.Error()
	enoent       = syscall.ENOENT.Error()
)

const (
	// 257 * "z"
	longName = "zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz"
)

type internal struct {
	pool  string
	files []string
	suite.Suite
}

func TestSuiteInternal(t *testing.T) {
	suite.Run(t, &internal{})
}

func command(name string, arg ...string) *exec.Cmd {
	cmd := exec.Command(name, arg...)
	cmd.Stderr = os.Stderr
	return cmd
}

func (s *internal) create(pool string) {
	s.pool = pool
	files := make([]string, 5)
	for i := range files {
		f, err := ioutil.TempFile("", "gozfs-test-temp")
		if err != nil {
			panic(err)
		}
		files[i] = f.Name()
		f.Close()
	}
	s.files = files

	script := []byte(`
	set -e
	pool=$1
	shift
	zpool list $pool &>/dev/null && zpool destroy $pool
	files=($@)
	for f in ${files[*]}; do
		truncate -s1G $f
	done
	zpool create $pool ${files[*]}

	zfs create $pool/a
	zfs create $pool/a/1
	zfs create $pool/a/2
	zfs create $pool/a/4
	zfs snapshot $pool/a/1@snap1
	zfs snapshot $pool/a/2@snap1
	zfs snapshot $pool/a/2@snap2
	zfs snapshot $pool/a/2@snap3
	zfs clone $pool/a/1@snap1 $pool/a/3
	zfs hold hold1 $pool/a/2@snap1
	zfs hold hold2 $pool/a/2@snap2
	zfs unmount $pool/a/4

	zfs create $pool/b
	zfs create -V 8192 $pool/b/1
	zfs create -b 1024 -V 2048 $pool/b/2
	zfs create -V 8192 $pool/b/4
	zfs snapshot $pool/b/1@snap1
	zfs snapshot $pool/b/2@snap1
	zfs snapshot $pool/b/2@snap2
	zfs snapshot $pool/b/2@snap3
	zfs clone $pool/b/1@snap1 $pool/b/3
	zfs hold hold1 $pool/b/2@snap1
	zfs hold hold2 $pool/b/2@snap2

	zfs create $pool/c
	zfs create $pool/c/one
	zfs create $pool/c/two
	zfs create $pool/c/three
	zfs snapshot -r $pool/c@snap1

	exit 0
	`)

	args := make([]string, 3, 3+len(files))
	args[0] = "bash"
	args[1] = "/dev/stdin"
	args[2] = s.pool
	args = append(args, files...)
	cmd := command("sudo", args...)

	stdin, err := cmd.StdinPipe()
	s.Require().NoError(err)
	go func() {
		_, err := stdin.Write([]byte(script))
		s.Require().NoError(err)
	}()

	s.Require().NoError(cmd.Run())
}

func (s *internal) destroy() {
	err := command("sudo", "zpool", "destroy", s.pool).Run()
	for i := range s.files {
		os.Remove(s.files[i])
	}
	s.Require().NoError(err)
}

func (s *internal) SetupTest() {
	s.create("gozfs-test")
}

func (s *internal) TearDownTest() {
	s.destroy()
}

func (s *internal) TestClone() {
	s.EqualError(clone(s.pool+"/a/2", s.pool+"/a/1", nil), eexist)
	s.EqualError(clone(s.pool+"/a 3", s.pool+"/a/1", nil), einval)
	s.EqualError(clone(s.pool+"/a/"+longName, s.pool+"/a/1", nil), einval) // WANTE(ENAMETOOLONG)
	s.EqualError(clone(s.pool+"/a/z", s.pool+"/a/"+longName, nil), enametoolong)
	s.NoError(clone(s.pool+"/a/z", s.pool+"/a/1@snap1", nil))
}

func (s *internal) TestCreateFS() {
	s.EqualError(create("gozfs-test-no-exists/1", dmuZFS, nil), enoent)
	s.EqualError(create(s.pool+"/~1", dmuZFS, nil), einval)
	s.EqualError(create(s.pool+"/1", dmuNumtypes+1, nil), einval)
	s.EqualError(create(s.pool+"/1", dmuZFS, map[string]interface{}{"bad-prop": true}), einval)
	s.EqualError(create(s.pool+"/"+longName, dmuZFS, nil), einval) // WANTE(ENAMETOOLONG)
	s.NoError(create(s.pool+"/1", dmuZFS, nil))
	s.EqualError(create(s.pool+"/1", dmuZFS, nil), eexist)
	s.EqualError(create(s.pool+"/2/2", dmuZFS, nil), enoent)
}

func (s *internal) TestCreateVOL() {
	type p map[string]interface{}
	u1024 := uint64(1024)
	u8192 := u1024 * 8

	tests := []struct {
		name  string
		props p
		err   string
	}{
		{"gozfs-test-no-exists/1", nil, enoent},
		{s.pool + "/" + longName, nil, einval}, // WANTE(ENAMETOOLONG)
		{s.pool + "/~1", nil, einval},
		{s.pool + "/1", p{"bad-prop": true}, einval},
		{s.pool + "/2", p{"volsize": 0}, einval},
		{s.pool + "/3", p{"volsize": 1}, einval},
		{s.pool + "/4", p{"volsize": 8*1024 + 1}, einval},
		{s.pool + "/5", p{"volsize": 8 * 1024}, einval},
		{s.pool + "/6", p{"volsize": uint64(0)}, einval},
		{s.pool + "/7", p{"volsize": uint64(1)}, einval},
		{s.pool + "/8", p{"volsize": u8192 + 1}, einval},
		{s.pool + "/9", p{"volsize": u8192}, ""},
		{s.pool + "/9", p{"volsize": u8192}, eexist},
		{s.pool + "/10", p{"volsize": 0, "volblocksize": 1024}, einval},
		{s.pool + "/11", p{"volsize": 1, "volblocksize": 1024}, einval},
		{s.pool + "/12", p{"volsize": 1024 + 1, "volblocksize": 1024}, einval},
		{s.pool + "/13", p{"volsize": 1024, "volblocksize": 1024}, einval},
		{s.pool + "/14", p{"volsize": uint64(0), "volblocksize": 1024}, einval},
		{s.pool + "/15", p{"volsize": uint64(1), "volblocksize": 1024}, einval},
		{s.pool + "/16", p{"volsize": u1024 + 1, "volblocksize": 1024}, einval},
		{s.pool + "/17", p{"volsize": u1024, "volblocksize": 1024}, einval},
		{s.pool + "/18", p{"volsize": 0, "volblocksize": u1024}, einval},
		{s.pool + "/19", p{"volsize": 1, "volblocksize": u1024}, einval},
		{s.pool + "/20", p{"volsize": 1024 + 1, "volblocksize": u1024}, einval},
		{s.pool + "/21", p{"volsize": 1024, "volblocksize": u1024}, einval},
		{s.pool + "/22", p{"volsize": uint64(0), "volblocksize": u1024}, einval},
		{s.pool + "/23", p{"volsize": uint64(1), "volblocksize": u1024}, einval},
		{s.pool + "/24", p{"volsize": u1024 + 1, "volblocksize": u1024}, einval},
		{s.pool + "/25", p{"volsize": u1024, "volblocksize": u1024}, ""},
		{s.pool + "/25", p{"volsize": u1024, "volblocksize": u1024}, eexist},
		{s.pool + "/26/1", p{"volsize": u1024, "volblocksize": u1024}, enoent},
	}
	for _, test := range tests {
		err := create(test.name, dmuZVOL, test.props)
		if test.err == "" {
			s.NoError(err, "name: %v, props: %v",
				test.name, test.props)
		} else {
			s.EqualError(err, test.err, "name: %v, props: %v",
				test.name, test.props)
		}
	}

}

func unmount(ds string) error {
	return command("sudo", "zfs", "unmount", ds).Run()
}

func unhold(tag, snapshot string) error {
	return command("sudo", "zfs", "release", tag, snapshot).Run()
}

func destroyKids(ds string) error {
	script := `
		ds=` + ds + `
		for fs in $(zfs list -H -t all $ds | sort -r | tail -n +2 | awk '{print $1}'); do
			zfs destroy -r $fs
		done
		zfs list -t all $ds;
		`
	return command("sudo", "bash", "-c", script).Run()
}

func (s *internal) TestDestroy() {
	s.EqualError(destroy("non-existent-pool", false), enoent)
	s.EqualError(destroy("non-existent-pool", true), enoent)
	s.EqualError(destroy(s.pool+"/"+longName, false), einval) // WANTE(ENAMETOOLONG)
	s.EqualError(destroy(s.pool+"/"+longName, true), einval)  // WANTE(ENAMETOOLONG)
	s.EqualError(destroy(s.pool+"/z", false), enoent)
	s.EqualError(destroy(s.pool+"/z", true), enoent)

	// mounted
	s.EqualError(destroy(s.pool+"/a/3", false), ebusy)
	s.NoError(unmount(s.pool + "/a/3"))
	s.NoError(destroy(s.pool+"/a/3", true))

	// has holds
	s.EqualError(destroy(s.pool+"/a/2@snap1", false), ebusy)
	s.NoError(unhold("hold1", s.pool+"/a/2@snap1"))
	s.NoError(destroy(s.pool+"/a/2@snap1", false))
	s.EqualError(destroy(s.pool+"/a/2@snap2", false), ebusy)
	s.NoError(unhold("hold2", s.pool+"/a/2@snap2"))
	s.NoError(destroy(s.pool+"/a/2@snap2", true))

	s.EqualError(destroy(s.pool+"/b/2@snap1", false), ebusy)
	s.NoError(unhold("hold1", s.pool+"/b/2@snap1"))
	s.NoError(destroy(s.pool+"/b/2@snap1", true))
	s.EqualError(destroy(s.pool+"/b/2@snap2", false), ebusy)
	s.NoError(unhold("hold2", s.pool+"/b/2@snap2"))
	s.NoError(destroy(s.pool+"/b/2@snap2", true))

	// misc
	s.NoError(destroy(s.pool+"/a/4", true))
	s.EqualError(destroy(s.pool+"/a/4", false), enoent)
	s.EqualError(destroy(s.pool+"/a/4", true), enoent)
	s.NoError(destroy(s.pool+"/b/4", false))
	s.NoError(destroy(s.pool+"/b/3", true))

	// has child datasets
	s.EqualError(destroy(s.pool+"/a", false), ebusy)
	s.NoError(unmount(s.pool + "/a"))
	s.EqualError(destroy(s.pool+"/a", false), eexist)
	s.NoError(destroyKids(s.pool + "/a"))
	s.NoError(destroy(s.pool+"/a", false))

	s.EqualError(destroy(s.pool+"/b", false), ebusy)
	s.NoError(destroyKids(s.pool + "/b"))
	s.NoError(unmount(s.pool + "/b"))
	s.NoError(destroy(s.pool+"/b", true))

}

func (s *internal) TestListEmpty() {
	s.destroy()

	m, err := list("", nil, true, 0)
	s.NoError(err, "m: %v", m)
	s.Assert().Len(m, 0)

	s.create(s.pool)
}

func (s *internal) TestList() {
	type t map[string]bool
	tests := []struct {
		name    string
		types   t
		recurse bool
		depth   uint64
		expNum  int
		err     string
	}{
		{name: "blah", err: enoent},
		{name: s.pool + "/" + longName, err: einval}, // WANTE(ENAMETOOLONG)
		{name: s.pool, expNum: 1},
		{name: s.pool, recurse: true, expNum: 11},
		{name: s.pool, recurse: true, depth: 1, expNum: 4},
		{name: s.pool + "/a", recurse: true, expNum: 5},
		{name: s.pool + "/a", recurse: true, depth: 1, expNum: 5},
		{name: s.pool, types: t{"volume": true}, recurse: true, expNum: 4},
		{name: s.pool, types: t{"volume": true}, recurse: true, depth: 1, expNum: 0},
		{name: s.pool + "/a", types: t{"volume": true}, recurse: true, expNum: 0},
		{name: s.pool + "/b", types: t{"volume": true}, recurse: true, expNum: 4},
		{name: s.pool + "/b", types: t{"volume": true}, recurse: true, depth: 1, expNum: 4},
		{name: s.pool, types: t{"snapshot": true}, recurse: true, expNum: 12},
		{name: s.pool, types: t{"snapshot": true}, recurse: true, depth: 1, expNum: 0},
		{name: s.pool + "/a", types: t{"snapshot": true}, recurse: true, expNum: 4},
		{name: s.pool + "/a/1", types: t{"snapshot": true}, recurse: true, expNum: 1},
		{name: s.pool + "/b", types: t{"snapshot": true}, recurse: true, expNum: 4},
		{name: s.pool + "/b/1", types: t{"snapshot": true}, recurse: true, expNum: 1},
	}
	for i, test := range tests {
		m, err := list(test.name, test.types, test.recurse, test.depth)
		if test.err != "" {
			s.EqualError(err, test.err, "test num:%d", i)
			s.Nil(m, "test:%d", i)
		} else {
			s.NoError(err, "test:%d", i)
			s.Len(m, test.expNum, "test num:%d", i)
		}
	}
}
<|endoftext|>


<commit_before>package scaffold

import (
	"bytes"
	"strings"
	"testing"
)

func TestSplitFilename(t *testing.T) {

	tests := []struct {
		input string
		bare  string
		ext   string
	}{
		{"test.txt", "test", ".txt"},
		{"test.a.b", "test.a", ".b"},
	}

	for _, test := range tests {
		bare, ext := splitFilename(test.input)

		if bare != test.bare || ext != test.ext {
			t.Errorf("splitFilename(%#v) = %#v, %#v; want %#v, %#v", test.input, bare, ext, test.bare, test.ext)
		}
	}

}

func TestIsLowercase(t *testing.T) {

	tests := []struct {
		input    string
		expected bool
	}{
		{"A", false},
		{"B", false},
		{"C", false},
		{"D", false},
		{"X", false},
		{"Y", false},
		{"Z", false},
		{"a", true},
		{"b", true},
		{"c", true},
		{"d", true},
		{"x", true},
		{"y", true},
		{"z", true},
	}

	for _, test := range tests {

		if got, want := isLowercase(test.input), test.expected; got != want {
			t.Errorf("isLowercase(%#v) = %v; want %v", test.input, got, want)
		}
	}

}

func TestFileName(t *testing.T) {

	/*
		([^-a-zA-Z_0-9])
	*/

	tests := []struct {
		input, expected string
	}{
		{" \\A/ want's 853 : until . # * +-0?@µ ", "A-wants-853-until-0"},
	}

	for _, test := range tests {

		if got, want := FileName(test.input), test.expected; got != want {
			t.Errorf("FileName(%#v) = %#v; want %#v", test.input, got, want)
		}
	}

}

func TestCamelCase1(t *testing.T) {

	tests := []struct {
		input, expected string
	}{
		{"field_name", "FieldName"},
		{"FieldName", "FieldName"},
		{"Field_name", "FieldName"},
		{"field_Name", "FieldName"},
		{"fieldname", "Fieldname"},
	}

	for _, test := range tests {

		if got, want := CamelCase1(test.input), test.expected; got != want {
			t.Errorf("CamelCase1(%v) = %v; want %v", test.input, got, want)
		}
	}

}

func TestCamelCase2(t *testing.T) {

	tests := []struct {
		input, expected string
	}{
		{"field_name", "fieldName"},
		{"FieldName", "FieldName"},
		{"Field_name", "FieldName"},
		{"field_Name", "fieldName"},
		{"fieldname", "fieldname"},
	}

	for _, test := range tests {

		if got, want := CamelCase2(test.input), test.expected; got != want {
			t.Errorf("CamelCase2(%#v) = %#v; want %#v", test.input, got, want)
		}
	}

}

func TestReplace(t *testing.T) {

	tests := []struct {
		input, src, dest, expected string
	}{
		{"field_name", "_", "*", "field*name"},
		{"field_name_x", "_", "*", "field*name*x"},
		{"_field_name_", "_", "*", "*field*name*"},
	}

	for _, test := range tests {

		if got, want := Replace(test.input, test.src, test.dest), test.expected; got != want {
			t.Errorf("Replace(%#v, %#v, %#v) = %#v; want %#v", test.input, test.src, test.dest, got, want)
		}
	}

}

var validHead = `{
	"Models": [
		{
			"Name": "",
			"Fields": [
				{"Name": "", "Type": ""}
			]
		}
	]
}`

var validBody = `{{range .Models}}
>>>models/
>>>{{toLower .Name}}/
>>>model.go
package {{replace .Name "_" "."}}

type {{camelCase1 .Name}} struct {
{{range .Fields}}
	{{camelCase1 .Name}} {{.Type}}
{{end}}
}

<<<model.go
<<<{{toLower .Name}}/
<<<models/
{{end}}`

var validJSON = `
{
	"Models": [
		{
			"Name": "person",
			"Fields": [
				{
					"Name": "first_name",
					"Type": "string"
				},
				{
					"Name": "last_name",
					"Type": "string"
				}
			]
		},
		{
			"Name": "address",
			"Fields": [
				{
					"Name": "street_no",
					"Type": "string"
				},
				{
					"Name": "city",
					"Type": "string"
				}
			]
		}
	]
}
`

var validTemplate = validHead + "\n\n" + validBody

func TestSplitTemplate(t *testing.T) {
	h, b := SplitTemplate(validTemplate)

	if h != validHead {
		t.Errorf("wrong head: %#v, expecting %#v", h, validHead)
	}

	if b != validBody {
		t.Errorf("wrong body: %#v, expecting %#v", b, validBody)
	}
}

func TestRun(t *testing.T) {

	tests := []struct {
		dir, body, json, expected string
	}{
		{
			"start",
			">>>file.txt\n<<<file.txt",
			`{}`,
			"start/file.txt\n",
		},
		{
			"start",
			">>>file1.txt\n<<<file1.txt\n>>>file2.txt\n<<<file2.txt",
			`{}`,
			"start/file1.txt\nstart/file2.txt\n",
		},
		{
			"start",
			">>>a/\n>>>b/\n>>>file1.txt\n<<<file1.txt\n>>>file2.txt\n<<<file2.txt\n<<<b/\n<<<a/\n",
			`{}`,
			"start/a/b/file1.txt\nstart/a/b/file2.txt\n",
		},
		{
			"a/dir",
			"{{range .Files}}>>>{{.Name}}.txt\n<<<{{.Name}}.txt\n{{end}}",
			`{"Files": [{"Name": "file1"},{"Name": "file2"}]}`,
			"a/dir/file1.txt\na/dir/file2.txt\n",
		},
		{
			"start/dir",
			validBody,
			validJSON,
			"start/dir/models/person/model.go\nstart/dir/models/address/model.go\n",
		},
	}

	for _, test := range tests {
		var log bytes.Buffer
		err := Run(test.dir, test.body, strings.NewReader(test.json), &log, true)
		if err != nil {
			t.Errorf("Run(%#v, %#v, %#v,...) returned error: %v", test.dir, test.body, test.json, err)
			continue
		}
		if got, want := log.String(), strings.Replace(test.expected, "\\\\", "/", -1); got != want {
			t.Errorf("Run(%#v, %#v, %#v,...) = %#v; want %#v", test.dir, test.body, test.json, got, want)
		}
	}
}

func TestRunErrors(t *testing.T) {

	tests := []struct {
		dir, body, json string
	}{
		{
			"start",
			">>>file.txt\n<<<file.txt",
			`{`,
		},
		{
			"start",
			">>file.txt\n<<<file.txt",
			`{}`,
		},
		{
			"start",
			">>>file1.txt\n<<<file2.txt",
			`{}`,
		},
		{
			"start",
			">>>a/\n>>>b\n>>>file1.txt\n<<<file1.txt\n<<<a/\n<<<b\n",
			`{}`,
		},
		{
			"start",
			">>>file1.txt\nho\n>>>file2.txt\nhu<<<file2.txt\n<<<file1.txt\n",
			`{}`,
		},
		{
			"start",
			"{{range .x}}",
			`{}`,
		},
	}

	for _, test := range tests {
		var log bytes.Buffer
		err := Run(test.dir, test.body, strings.NewReader(test.json), &log, true)
		if err == nil {
			t.Errorf("Run(%#v, %#v, %#v,...) returned no error", test.dir, test.body, test.json)
		}
	}
}
<commit_msg>another try for file paths on windows<commit_after>package scaffold

import (
	"bytes"
	"strings"
	"testing"
)

func TestSplitFilename(t *testing.T) {

	tests := []struct {
		input string
		bare  string
		ext   string
	}{
		{"test.txt", "test", ".txt"},
		{"test.a.b", "test.a", ".b"},
	}

	for _, test := range tests {
		bare, ext := splitFilename(test.input)

		if bare != test.bare || ext != test.ext {
			t.Errorf("splitFilename(%#v) = %#v, %#v; want %#v, %#v", test.input, bare, ext, test.bare, test.ext)
		}
	}

}

func TestIsLowercase(t *testing.T) {

	tests := []struct {
		input    string
		expected bool
	}{
		{"A", false},
		{"B", false},
		{"C", false},
		{"D", false},
		{"X", false},
		{"Y", false},
		{"Z", false},
		{"a", true},
		{"b", true},
		{"c", true},
		{"d", true},
		{"x", true},
		{"y", true},
		{"z", true},
	}

	for _, test := range tests {

		if got, want := isLowercase(test.input), test.expected; got != want {
			t.Errorf("isLowercase(%#v) = %v; want %v", test.input, got, want)
		}
	}

}

func TestFileName(t *testing.T) {

	/*
		([^-a-zA-Z_0-9])
	*/

	tests := []struct {
		input, expected string
	}{
		{" \\A/ want's 853 : until . # * +-0?@µ ", "A-wants-853-until-0"},
	}

	for _, test := range tests {

		if got, want := FileName(test.input), test.expected; got != want {
			t.Errorf("FileName(%#v) = %#v; want %#v", test.input, got, want)
		}
	}

}

func TestCamelCase1(t *testing.T) {

	tests := []struct {
		input, expected string
	}{
		{"field_name", "FieldName"},
		{"FieldName", "FieldName"},
		{"Field_name", "FieldName"},
		{"field_Name", "FieldName"},
		{"fieldname", "Fieldname"},
	}

	for _, test := range tests {

		if got, want := CamelCase1(test.input), test.expected; got != want {
			t.Errorf("CamelCase1(%v) = %v; want %v", test.input, got, want)
		}
	}

}

func TestCamelCase2(t *testing.T) {

	tests := []struct {
		input, expected string
	}{
		{"field_name", "fieldName"},
		{"FieldName", "FieldName"},
		{"Field_name", "FieldName"},
		{"field_Name", "fieldName"},
		{"fieldname", "fieldname"},
	}

	for _, test := range tests {

		if got, want := CamelCase2(test.input), test.expected; got != want {
			t.Errorf("CamelCase2(%#v) = %#v; want %#v", test.input, got, want)
		}
	}

}

func TestReplace(t *testing.T) {

	tests := []struct {
		input, src, dest, expected string
	}{
		{"field_name", "_", "*", "field*name"},
		{"field_name_x", "_", "*", "field*name*x"},
		{"_field_name_", "_", "*", "*field*name*"},
	}

	for _, test := range tests {

		if got, want := Replace(test.input, test.src, test.dest), test.expected; got != want {
			t.Errorf("Replace(%#v, %#v, %#v) = %#v; want %#v", test.input, test.src, test.dest, got, want)
		}
	}

}

var validHead = `{
	"Models": [
		{
			"Name": "",
			"Fields": [
				{"Name": "", "Type": ""}
			]
		}
	]
}`

var validBody = `{{range .Models}}
>>>models/
>>>{{toLower .Name}}/
>>>model.go
package {{replace .Name "_" "."}}

type {{camelCase1 .Name}} struct {
{{range .Fields}}
	{{camelCase1 .Name}} {{.Type}}
{{end}}
}

<<<model.go
<<<{{toLower .Name}}/
<<<models/
{{end}}`

var validJSON = `
{
	"Models": [
		{
			"Name": "person",
			"Fields": [
				{
					"Name": "first_name",
					"Type": "string"
				},
				{
					"Name": "last_name",
					"Type": "string"
				}
			]
		},
		{
			"Name": "address",
			"Fields": [
				{
					"Name": "street_no",
					"Type": "string"
				},
				{
					"Name": "city",
					"Type": "string"
				}
			]
		}
	]
}
`

var validTemplate = validHead + "\n\n" + validBody

func TestSplitTemplate(t *testing.T) {
	h, b := SplitTemplate(validTemplate)

	if h != validHead {
		t.Errorf("wrong head: %#v, expecting %#v", h, validHead)
	}

	if b != validBody {
		t.Errorf("wrong body: %#v, expecting %#v", b, validBody)
	}
}

func TestRun(t *testing.T) {

	tests := []struct {
		dir, body, json, expected string
	}{
		{
			"start",
			">>>file.txt\n<<<file.txt",
			`{}`,
			"start/file.txt\n",
		},
		{
			"start",
			">>>file1.txt\n<<<file1.txt\n>>>file2.txt\n<<<file2.txt",
			`{}`,
			"start/file1.txt\nstart/file2.txt\n",
		},
		{
			"start",
			">>>a/\n>>>b/\n>>>file1.txt\n<<<file1.txt\n>>>file2.txt\n<<<file2.txt\n<<<b/\n<<<a/\n",
			`{}`,
			"start/a/b/file1.txt\nstart/a/b/file2.txt\n",
		},
		{
			"a/dir",
			"{{range .Files}}>>>{{.Name}}.txt\n<<<{{.Name}}.txt\n{{end}}",
			`{"Files": [{"Name": "file1"},{"Name": "file2"}]}`,
			"a/dir/file1.txt\na/dir/file2.txt\n",
		},
		{
			"start/dir",
			validBody,
			validJSON,
			"start/dir/models/person/model.go\nstart/dir/models/address/model.go\n",
		},
	}

	for _, test := range tests {
		var log bytes.Buffer
		err := Run(test.dir, test.body, strings.NewReader(test.json), &log, true)
		if err != nil {
			t.Errorf("Run(%#v, %#v, %#v,...) returned error: %v", test.dir, test.body, test.json, err)
			continue
		}
		if got, want := log.String(), strings.Replace(test.expected, "\\", "/", -1); got != want {
			t.Errorf("Run(%#v, %#v, %#v,...) = %#v; want %#v", test.dir, test.body, test.json, got, want)
		}
	}
}

func TestRunErrors(t *testing.T) {

	tests := []struct {
		dir, body, json string
	}{
		{
			"start",
			">>>file.txt\n<<<file.txt",
			`{`,
		},
		{
			"start",
			">>file.txt\n<<<file.txt",
			`{}`,
		},
		{
			"start",
			">>>file1.txt\n<<<file2.txt",
			`{}`,
		},
		{
			"start",
			">>>a/\n>>>b\n>>>file1.txt\n<<<file1.txt\n<<<a/\n<<<b\n",
			`{}`,
		},
		{
			"start",
			">>>file1.txt\nho\n>>>file2.txt\nhu<<<file2.txt\n<<<file1.txt\n",
			`{}`,
		},
		{
			"start",
			"{{range .x}}",
			`{}`,
		},
	}

	for _, test := range tests {
		var log bytes.Buffer
		err := Run(test.dir, test.body, strings.NewReader(test.json), &log, true)
		if err == nil {
			t.Errorf("Run(%#v, %#v, %#v,...) returned no error", test.dir, test.body, test.json)
		}
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017, 2018 Ankyra

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 script

import (
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"reflect"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/ankyra/escape-core/util"
)

type StdlibFunc struct {
	Id     string
	Func   Script
	Doc    string
	ActsOn string
	Args   string
}

var trackMajorVersion = ShouldParse(`$func(v) { $v.split(".")[:1].join(".").concat(".@") }`)
var trackMinorVersion = ShouldParse(`$func(v) { $v.split(".")[:2].join(".").concat(".@") }`)
var trackPatchVersion = ShouldParse(`$func(v) { $v.split(".")[:3].join(".").concat(".@") }`)
var trackVersion = ShouldParse(`$func(v) { $v.concat(".@") }`)

var Stdlib = []StdlibFunc{
	StdlibFunc{"id", LiftFunction(builtinId), "Returns its argument", "everything", "parameter :: *"},
	StdlibFunc{"equals", LiftFunction(builtinEquals), "Returns true if the arguments are of the same type and have the same value", "everything", "parameter :: *"},
	StdlibFunc{"env_lookup", LiftFunction(builtinEnvLookup), "Lookup key in environment. Usually called implicitly when using '$'", "lists", "key :: string"},
	StdlibFunc{"concat", LiftFunction(builtinConcat), "Concatate stringable arguments", "strings", "v1 :: string, v2 :: string, ..."},
	StdlibFunc{"lower", ShouldLift(strings.ToLower), "Returns a copy of the string v with all Unicode characters mapped to their lower case", "strings", "v :: string"},
	StdlibFunc{"upper", ShouldLift(strings.ToUpper), "Returns a copy of the string v with all Unicode characters mapped to their upper case", "strings", "v :: string"},
	StdlibFunc{"title", ShouldLift(strings.ToTitle), "Returns a copy of the string v with all Unicode characters mapped to their title case", "strings", "v :: string"},
	StdlibFunc{"split", ShouldLift(strings.Split), "Split slices s into all substrings separated by sep and returns a slice of the substrings between those separators. If sep is empty, Split splits after each UTF-8 sequence.", "strings", "sep :: string"},
	StdlibFunc{"path_exists", ShouldLift(util.PathExists), "Returns true if the path exists, false if not", "strings", ""},
	StdlibFunc{"file_exists", ShouldLift(util.FileExists), "Returns true if the path exists and if it's not a directory, false otherwise", "strings", ""},
	StdlibFunc{"dir_exists", ShouldLift(util.IsDir), "Returns true if the path exists and if it is a directory, false otherwise", "strings", ""},
	StdlibFunc{"join", ShouldLift(strings.Join), "Join concatenates the elements of a to create a single string. The separator string sep is placed between elements in the resulting string. ", "lists", "sep :: string"},
	StdlibFunc{"replace", ShouldLift(strings.Replace), "Replace returns a copy of the string s with the first n non-overlapping instances of old replaced by new. If old is empty, it matches at the beginning of the string and after each UTF-8 sequence, yielding up to k+1 replacements for a k-rune string. If n < 0, there is no limit on the number of replacements.", "strings", "old :: string, new :: string, n :: integer"},
	StdlibFunc{"base64_encode", ShouldLift(base64.StdEncoding.EncodeToString), "Encode string to base64", "strings", ""},
	StdlibFunc{"base64_decode", ShouldLift(base64.StdEncoding.DecodeString), "Decode string from base64", "strings", ""},
	StdlibFunc{"trim", ShouldLift(strings.TrimSpace), "Returns a slice of the string s, with all leading and trailing white space removed, as defined by Unicode. ", "strings", ""},
	StdlibFunc{"list_index", LiftFunction(builtinListIndex), "Index a list at position `n`. Usually accessed implicitly using indexing syntax (eg. `list[0]`)", "lists", "n :: integer"},
	StdlibFunc{"length", LiftFunction(builtinListLength), "Returns the length of the list", "lists", "n :: integer"},
	StdlibFunc{"list_slice", LiftFunction(builtinListSlice), "Slice a list. Usually accessed implicitly using slice syntax (eg. `list[0:5]`)", "lists", "i :: integer, j :: integer"},
	StdlibFunc{"add", ShouldLift(builtinAdd), "Add two integers", "integers", "y :: integer"},
	StdlibFunc{"timestamp", ShouldLift(builtinTimestamp), "Returns a UNIX timestamp", "", ""},
	StdlibFunc{"read_file", ShouldLift(builtinReadfile), "Read the contents of a file", "strings", ""},
	StdlibFunc{"track_major_version", trackMajorVersion, "Track major version", "strings", ""},
	StdlibFunc{"track_minor_version", trackMinorVersion, "Track minor version", "strings", ""},
	StdlibFunc{"track_patch_version", trackPatchVersion, "Track patch version", "strings", ""},
	StdlibFunc{"track_version", trackVersion, "Track version", "strings", ""},
	StdlibFunc{"not", LiftFunction(builtinNot), "Logical NOT operation", "bool", ""},
}

func LiftGoFunc(f interface{}) Script {
	name := runtime.FuncForPC(reflect.ValueOf(f).Pointer()).Name()
	typ := reflect.TypeOf(f)
	nInputs := typ.NumIn()
	nOutputs := typ.NumOut()
	scriptFunc := func(env *ScriptEnvironment, args []Script) (Script, error) {
		if err := builtinArgCheck(nInputs, name, args); err != nil {
			return nil, err
		}

		goArgs := []reflect.Value{}
		for i := 0; i < nInputs; i++ {
			argType := typ.In(i)
			arg := args[i]

			if argType.Kind() == reflect.String {
				if !IsStringAtom(arg) {
					return nil, fmt.Errorf("Expecting string argument in call to %s, but got %s", name, arg.Type().Name())
				} else {
					goArgs = append(goArgs, reflect.ValueOf(ExpectStringAtom(arg)))
				}
			} else if argType.Kind() == reflect.Int {
				if !IsIntegerAtom(arg) {
					return nil, fmt.Errorf("Expecting integer argument in call to %s, but got %s", name, arg.Type().Name())
				} else {
					goArgs = append(goArgs, reflect.ValueOf(ExpectIntegerAtom(arg)))
				}
			} else if argType.Kind() == reflect.Slice {
				if !IsListAtom(arg) {
					if argType.Elem().Kind() == reflect.Uint8 && IsStringAtom(arg) {
						goArgs = append(goArgs, reflect.ValueOf([]byte(ExpectStringAtom(arg))))
					} else {
						return nil, fmt.Errorf("Expecting list argument in call to %s, but got %s", name, arg.Type().Name())
					}
				} else {
					lst := ExpectListAtom(arg) // []Script
					if argType.Elem().Kind() == reflect.String {
						strArg := []string{}
						for k := 0; k < len(lst); k++ {
							if !IsStringAtom(lst[k]) {
								return nil, fmt.Errorf("Expecting string value in list in call to %s, but got %s", name, arg.Type().Name())
							} else {
								strArg = append(strArg, ExpectStringAtom(lst[k]))
							}
						}
						goArgs = append(goArgs, reflect.ValueOf(strArg))
					} else {
						return nil, fmt.Errorf("Unsupported slice type in function %s", name)
					}
				}
			} else {
				return nil, fmt.Errorf("Unsupported argument type '%s' in function %s", argType.Kind(), name)
			}
		}

		outputs := reflect.ValueOf(f).Call(goArgs)
		if nOutputs == 1 {
			return Lift(outputs[0].Interface())
		}
		if nOutputs == 2 {
			_, isError := outputs[1].Interface().(error)
			if isError && outputs[1].Interface() != nil {
				return nil, fmt.Errorf("Error in call to %s: %s", name, outputs[1].Interface().(error))
			}
			return Lift(outputs[0].Interface())
		}
		if nOutputs != 1 {
			return nil, fmt.Errorf("Go functions with multiple outputs are not supported at this time")
		}
		return Lift(outputs[0].Interface())
	}
	return LiftFunction(scriptFunc)
}

/*
   Builtins
*/
func builtinArgCheck(expected int, funcName string, inputValues []Script) error {
	if len(inputValues) != expected {
		return fmt.Errorf("Expecting %d argument(s) in call to '%s', got %d",
			expected, funcName, len(inputValues))
	}
	return nil
}

func builtinId(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if err := builtinArgCheck(1, "id", inputValues); err != nil {
		return nil, err
	}
	return inputValues[0], nil
}

func builtinEnvLookup(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if err := builtinArgCheck(1, "env_lookup", inputValues); err != nil {
		return nil, err
	}
	arg := inputValues[0]
	if !IsStringAtom(arg) {
		return nil, fmt.Errorf("Expecting string argument in environment lookup call, but got '%s'", arg.Type().Name())
	}
	key := ExpectStringAtom(arg)
	val, ok := (*env)[key]
	if !ok {
		return nil, fmt.Errorf("Field '%s' was not found in environment.", key)
	}
	return val, nil
}

func builtinConcat(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	result := ""
	for _, val := range inputValues {
		if IsStringAtom(val) {
			result += ExpectStringAtom(val)
		} else if IsIntegerAtom(val) {
			result += strconv.Itoa(ExpectIntegerAtom(val))
		} else {
			return nil, fmt.Errorf("Can't concatenate value of type %s", val.Type().Name())
		}
	}
	return LiftString(result), nil
}

func builtinListLength(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if err := builtinArgCheck(1, "list_index", inputValues); err != nil {
		return nil, err
	}
	lstArg := inputValues[0]
	if !IsListAtom(lstArg) {
		if IsStringAtom(lstArg) {
			str := ExpectStringAtom(lstArg)
			return LiftInteger(len(str)), nil
		}
		return nil, fmt.Errorf("Expecting list or string argument in length call, but got '%s'", lstArg.Type().Name())
	}
	lst := ExpectListAtom(inputValues[0])
	return LiftInteger(len(lst)), nil
}

func builtinListIndex(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if err := builtinArgCheck(2, "list_index", inputValues); err != nil {
		return nil, err
	}
	lstArg := inputValues[0]
	if !IsListAtom(lstArg) {
		return nil, fmt.Errorf("Expecting list argument in list index call, but got '%s'", lstArg.Type().Name())
	}
	indexArg := inputValues[1]
	if !IsIntegerAtom(indexArg) {
		return nil, fmt.Errorf("Expecting integer argument in list index call, but got '%s'", indexArg.Type().Name())
	}
	lst := ExpectListAtom(inputValues[0])
	index := ExpectIntegerAtom(inputValues[1])
	if index < 0 || index >= len(lst) {
		return nil, fmt.Errorf("Index '%d' out of range (len: %d)", index, len(lst))
	}
	return Lift(lst[index])
}

func builtinListSlice(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if len(inputValues) < 2 || len(inputValues) > 3 {
		return nil, fmt.Errorf("Expecting at least %d argument(s) (but not more than 3) in call to '%s', got %d",
			2, "list slice", len(inputValues))
	}
	lstArg := inputValues[0]
	if !IsListAtom(lstArg) {
		return nil, fmt.Errorf("Expecting list argument in list slice call, but got '%s'", lstArg.Type().Name())
	}
	indexArg := inputValues[1]
	if !IsIntegerAtom(indexArg) {
		return nil, fmt.Errorf("Expecting integer argument in list slice call, but got '%s'", indexArg.Type().Name())
	}
	lst := ExpectListAtom(inputValues[0])
	index := ExpectIntegerAtom(inputValues[1])

	if len(inputValues) == 3 {
		endSliceArg := inputValues[2]
		if !IsIntegerAtom(endSliceArg) {
			return nil, fmt.Errorf("Expecting integer argument in list slice call, but got '%s'", endSliceArg.Type().Name())
		}
		endIndex := ExpectIntegerAtom(inputValues[2])
		if endIndex < 0 {
			endIndex = len(lst) + endIndex
		}
		return Lift(lst[index:endIndex])
	}
	return Lift(lst[index:])
}

func builtinEquals(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if err := builtinArgCheck(2, "equals", inputValues); err != nil {
		return nil, err
	}
	i1 := inputValues[0]
	i2 := inputValues[1]
	return Lift(i1.Equals(i2))
}

func builtinNot(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if err := builtinArgCheck(1, "not", inputValues); err != nil {
		return nil, err
	}
	boolArg := inputValues[0]
	if !IsBoolAtom(boolArg) {
		return nil, fmt.Errorf("Expecting bool argument in not call, but got '%s'", boolArg.Type().Name())
	}
	bool := ExpectBoolAtom(boolArg)
	return Lift(!bool)
}

func builtinReadfile(arg string) (string, error) {
	bytes, err := ioutil.ReadFile(arg)
	if err != nil {
		return "", err
	}
	return string(bytes), nil
}

func builtinTimestamp() string {
	return strconv.Itoa(int(time.Now().Unix()))
}

func builtinAdd(x, y int) int {
	return x + y
}
<commit_msg>Add 'and' and 'or' function to stdlib<commit_after>/*
Copyright 2017, 2018 Ankyra

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 script

import (
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"reflect"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/ankyra/escape-core/util"
)

type StdlibFunc struct {
	Id     string
	Func   Script
	Doc    string
	ActsOn string
	Args   string
}

var trackMajorVersion = ShouldParse(`$func(v) { $v.split(".")[:1].join(".").concat(".@") }`)
var trackMinorVersion = ShouldParse(`$func(v) { $v.split(".")[:2].join(".").concat(".@") }`)
var trackPatchVersion = ShouldParse(`$func(v) { $v.split(".")[:3].join(".").concat(".@") }`)
var trackVersion = ShouldParse(`$func(v) { $v.concat(".@") }`)

var Stdlib = []StdlibFunc{
	StdlibFunc{"id", LiftFunction(builtinId), "Returns its argument", "everything", "parameter :: *"},
	StdlibFunc{"equals", LiftFunction(builtinEquals), "Returns true if the arguments are of the same type and have the same value", "everything", "parameter :: *"},
	StdlibFunc{"env_lookup", LiftFunction(builtinEnvLookup), "Lookup key in environment. Usually called implicitly when using '$'", "lists", "key :: string"},
	StdlibFunc{"concat", LiftFunction(builtinConcat), "Concatate stringable arguments", "strings", "v1 :: string, v2 :: string, ..."},
	StdlibFunc{"lower", ShouldLift(strings.ToLower), "Returns a copy of the string v with all Unicode characters mapped to their lower case", "strings", "v :: string"},
	StdlibFunc{"upper", ShouldLift(strings.ToUpper), "Returns a copy of the string v with all Unicode characters mapped to their upper case", "strings", "v :: string"},
	StdlibFunc{"title", ShouldLift(strings.ToTitle), "Returns a copy of the string v with all Unicode characters mapped to their title case", "strings", "v :: string"},
	StdlibFunc{"split", ShouldLift(strings.Split), "Split slices s into all substrings separated by sep and returns a slice of the substrings between those separators. If sep is empty, Split splits after each UTF-8 sequence.", "strings", "sep :: string"},
	StdlibFunc{"path_exists", ShouldLift(util.PathExists), "Returns true if the path exists, false if not", "strings", ""},
	StdlibFunc{"file_exists", ShouldLift(util.FileExists), "Returns true if the path exists and if it's not a directory, false otherwise", "strings", ""},
	StdlibFunc{"dir_exists", ShouldLift(util.IsDir), "Returns true if the path exists and if it is a directory, false otherwise", "strings", ""},
	StdlibFunc{"join", ShouldLift(strings.Join), "Join concatenates the elements of a to create a single string. The separator string sep is placed between elements in the resulting string. ", "lists", "sep :: string"},
	StdlibFunc{"replace", ShouldLift(strings.Replace), "Replace returns a copy of the string s with the first n non-overlapping instances of old replaced by new. If old is empty, it matches at the beginning of the string and after each UTF-8 sequence, yielding up to k+1 replacements for a k-rune string. If n < 0, there is no limit on the number of replacements.", "strings", "old :: string, new :: string, n :: integer"},
	StdlibFunc{"base64_encode", ShouldLift(base64.StdEncoding.EncodeToString), "Encode string to base64", "strings", ""},
	StdlibFunc{"base64_decode", ShouldLift(base64.StdEncoding.DecodeString), "Decode string from base64", "strings", ""},
	StdlibFunc{"trim", ShouldLift(strings.TrimSpace), "Returns a slice of the string s, with all leading and trailing white space removed, as defined by Unicode. ", "strings", ""},
	StdlibFunc{"list_index", LiftFunction(builtinListIndex), "Index a list at position `n`. Usually accessed implicitly using indexing syntax (eg. `list[0]`)", "lists", "n :: integer"},
	StdlibFunc{"length", LiftFunction(builtinListLength), "Returns the length of the list", "lists", "n :: integer"},
	StdlibFunc{"list_slice", LiftFunction(builtinListSlice), "Slice a list. Usually accessed implicitly using slice syntax (eg. `list[0:5]`)", "lists", "i :: integer, j :: integer"},
	StdlibFunc{"add", ShouldLift(builtinAdd), "Add two integers", "integers", "y :: integer"},
	StdlibFunc{"timestamp", ShouldLift(builtinTimestamp), "Returns a UNIX timestamp", "", ""},
	StdlibFunc{"read_file", ShouldLift(builtinReadfile), "Read the contents of a file", "strings", ""},
	StdlibFunc{"track_major_version", trackMajorVersion, "Track major version", "strings", ""},
	StdlibFunc{"track_minor_version", trackMinorVersion, "Track minor version", "strings", ""},
	StdlibFunc{"track_patch_version", trackPatchVersion, "Track patch version", "strings", ""},
	StdlibFunc{"track_version", trackVersion, "Track version", "strings", ""},
	StdlibFunc{"not", LiftFunction(builtinNOT), "Logical NOT operation", "bool", ""},
	StdlibFunc{"and", LiftFunction(builtinAND), "Logical AND operation", "bool", "b2 :: bool"},
	StdlibFunc{"or", LiftFunction(builtinOR), "Logical OR operation", "bool", "b2 :: bool"},
}

func LiftGoFunc(f interface{}) Script {
	name := runtime.FuncForPC(reflect.ValueOf(f).Pointer()).Name()
	typ := reflect.TypeOf(f)
	nInputs := typ.NumIn()
	nOutputs := typ.NumOut()
	scriptFunc := func(env *ScriptEnvironment, args []Script) (Script, error) {
		if err := builtinArgCheck(nInputs, name, args); err != nil {
			return nil, err
		}

		goArgs := []reflect.Value{}
		for i := 0; i < nInputs; i++ {
			argType := typ.In(i)
			arg := args[i]

			if argType.Kind() == reflect.String {
				if !IsStringAtom(arg) {
					return nil, fmt.Errorf("Expecting string argument in call to %s, but got %s", name, arg.Type().Name())
				} else {
					goArgs = append(goArgs, reflect.ValueOf(ExpectStringAtom(arg)))
				}
			} else if argType.Kind() == reflect.Int {
				if !IsIntegerAtom(arg) {
					return nil, fmt.Errorf("Expecting integer argument in call to %s, but got %s", name, arg.Type().Name())
				} else {
					goArgs = append(goArgs, reflect.ValueOf(ExpectIntegerAtom(arg)))
				}
			} else if argType.Kind() == reflect.Slice {
				if !IsListAtom(arg) {
					if argType.Elem().Kind() == reflect.Uint8 && IsStringAtom(arg) {
						goArgs = append(goArgs, reflect.ValueOf([]byte(ExpectStringAtom(arg))))
					} else {
						return nil, fmt.Errorf("Expecting list argument in call to %s, but got %s", name, arg.Type().Name())
					}
				} else {
					lst := ExpectListAtom(arg) // []Script
					if argType.Elem().Kind() == reflect.String {
						strArg := []string{}
						for k := 0; k < len(lst); k++ {
							if !IsStringAtom(lst[k]) {
								return nil, fmt.Errorf("Expecting string value in list in call to %s, but got %s", name, arg.Type().Name())
							} else {
								strArg = append(strArg, ExpectStringAtom(lst[k]))
							}
						}
						goArgs = append(goArgs, reflect.ValueOf(strArg))
					} else {
						return nil, fmt.Errorf("Unsupported slice type in function %s", name)
					}
				}
			} else {
				return nil, fmt.Errorf("Unsupported argument type '%s' in function %s", argType.Kind(), name)
			}
		}

		outputs := reflect.ValueOf(f).Call(goArgs)
		if nOutputs == 1 {
			return Lift(outputs[0].Interface())
		}
		if nOutputs == 2 {
			_, isError := outputs[1].Interface().(error)
			if isError && outputs[1].Interface() != nil {
				return nil, fmt.Errorf("Error in call to %s: %s", name, outputs[1].Interface().(error))
			}
			return Lift(outputs[0].Interface())
		}
		if nOutputs != 1 {
			return nil, fmt.Errorf("Go functions with multiple outputs are not supported at this time")
		}
		return Lift(outputs[0].Interface())
	}
	return LiftFunction(scriptFunc)
}

/*
   Builtins
*/
func builtinArgCheck(expected int, funcName string, inputValues []Script) error {
	if len(inputValues) != expected {
		return fmt.Errorf("Expecting %d argument(s) in call to '%s', got %d",
			expected, funcName, len(inputValues))
	}
	return nil
}

func builtinId(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if err := builtinArgCheck(1, "id", inputValues); err != nil {
		return nil, err
	}
	return inputValues[0], nil
}

func builtinEnvLookup(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if err := builtinArgCheck(1, "env_lookup", inputValues); err != nil {
		return nil, err
	}
	arg := inputValues[0]
	if !IsStringAtom(arg) {
		return nil, fmt.Errorf("Expecting string argument in environment lookup call, but got '%s'", arg.Type().Name())
	}
	key := ExpectStringAtom(arg)
	val, ok := (*env)[key]
	if !ok {
		return nil, fmt.Errorf("Field '%s' was not found in environment.", key)
	}
	return val, nil
}

func builtinConcat(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	result := ""
	for _, val := range inputValues {
		if IsStringAtom(val) {
			result += ExpectStringAtom(val)
		} else if IsIntegerAtom(val) {
			result += strconv.Itoa(ExpectIntegerAtom(val))
		} else {
			return nil, fmt.Errorf("Can't concatenate value of type %s", val.Type().Name())
		}
	}
	return LiftString(result), nil
}

func builtinListLength(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if err := builtinArgCheck(1, "list_index", inputValues); err != nil {
		return nil, err
	}
	lstArg := inputValues[0]
	if !IsListAtom(lstArg) {
		if IsStringAtom(lstArg) {
			str := ExpectStringAtom(lstArg)
			return LiftInteger(len(str)), nil
		}
		return nil, fmt.Errorf("Expecting list or string argument in length call, but got '%s'", lstArg.Type().Name())
	}
	lst := ExpectListAtom(inputValues[0])
	return LiftInteger(len(lst)), nil
}

func builtinListIndex(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if err := builtinArgCheck(2, "list_index", inputValues); err != nil {
		return nil, err
	}
	lstArg := inputValues[0]
	if !IsListAtom(lstArg) {
		return nil, fmt.Errorf("Expecting list argument in list index call, but got '%s'", lstArg.Type().Name())
	}
	indexArg := inputValues[1]
	if !IsIntegerAtom(indexArg) {
		return nil, fmt.Errorf("Expecting integer argument in list index call, but got '%s'", indexArg.Type().Name())
	}
	lst := ExpectListAtom(inputValues[0])
	index := ExpectIntegerAtom(inputValues[1])
	if index < 0 || index >= len(lst) {
		return nil, fmt.Errorf("Index '%d' out of range (len: %d)", index, len(lst))
	}
	return Lift(lst[index])
}

func builtinListSlice(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if len(inputValues) < 2 || len(inputValues) > 3 {
		return nil, fmt.Errorf("Expecting at least %d argument(s) (but not more than 3) in call to '%s', got %d",
			2, "list slice", len(inputValues))
	}
	lstArg := inputValues[0]
	if !IsListAtom(lstArg) {
		return nil, fmt.Errorf("Expecting list argument in list slice call, but got '%s'", lstArg.Type().Name())
	}
	indexArg := inputValues[1]
	if !IsIntegerAtom(indexArg) {
		return nil, fmt.Errorf("Expecting integer argument in list slice call, but got '%s'", indexArg.Type().Name())
	}
	lst := ExpectListAtom(inputValues[0])
	index := ExpectIntegerAtom(inputValues[1])

	if len(inputValues) == 3 {
		endSliceArg := inputValues[2]
		if !IsIntegerAtom(endSliceArg) {
			return nil, fmt.Errorf("Expecting integer argument in list slice call, but got '%s'", endSliceArg.Type().Name())
		}
		endIndex := ExpectIntegerAtom(inputValues[2])
		if endIndex < 0 {
			endIndex = len(lst) + endIndex
		}
		return Lift(lst[index:endIndex])
	}
	return Lift(lst[index:])
}

func builtinEquals(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if err := builtinArgCheck(2, "equals", inputValues); err != nil {
		return nil, err
	}
	i1 := inputValues[0]
	i2 := inputValues[1]
	return Lift(i1.Equals(i2))
}

func builtinNOT(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if err := builtinArgCheck(1, "not", inputValues); err != nil {
		return nil, err
	}
	boolArg := inputValues[0]
	if !IsBoolAtom(boolArg) {
		return nil, fmt.Errorf("Expecting bool argument in not call, but got '%s'", boolArg.Type().Name())
	}
	bool := ExpectBoolAtom(boolArg)
	return Lift(!bool)
}

func builtinAND(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if err := builtinArgCheck(2, "and", inputValues); err != nil {
		return nil, err
	}
	boolArg1 := inputValues[0]
	boolArg2 := inputValues[1]
	if !IsBoolAtom(boolArg1) && !IsBoolAtom(boolArg2) {
		return nil, fmt.Errorf("Expecting bool argument in and call, but got '%s'", boolArg.Type().Name())
	}
	bool1 := ExpectBoolAtom(boolArg1)
	bool2 := ExpectBoolAtom(boolArg2)
	return Lift(bool1 && bool2)
}

func builtinOR(env *ScriptEnvironment, inputValues []Script) (Script, error) {
	if err := builtinArgCheck(2, "and", inputValues); err != nil {
		return nil, err
	}
	boolArg1 := inputValues[0]
	boolArg2 := inputValues[1]
	if !IsBoolAtom(boolArg1) && !IsBoolAtom(boolArg2) {
		return nil, fmt.Errorf("Expecting bool argument in and call, but got '%s'", boolArg.Type().Name())
	}
	bool1 := ExpectBoolAtom(boolArg1)
	bool2 := ExpectBoolAtom(boolArg2)
	return Lift(bool1 || bool2)
}

func builtinReadfile(arg string) (string, error) {
	bytes, err := ioutil.ReadFile(arg)
	if err != nil {
		return "", err
	}
	return string(bytes), nil
}

func builtinTimestamp() string {
	return strconv.Itoa(int(time.Now().Unix()))
}

func builtinAdd(x, y int) int {
	return x + y
}
<|endoftext|>


<commit_before>package mustache

import (
	"os"
	"path"
)

type PartialProvider interface {
	Get(name string) (*Template, error)
}

type FileProvider struct {
	Paths      []string
	Extensions []string
}

func (fp *FileProvider) Get(name string) (*Template, error) {
	var filename string

	for _, p := range fp.Paths {
		for _, e := range fp.Extensions {
			name := path.Join(p, name+e)
			f, err := os.Open(name)
			if err == nil {
				filename = name
				f.Close()
				break
			}
		}
	}

	if filename == "" {
		return &Template{"", "{{", "}}", 0, 1, "", []interface{}{}, nil}, nil
	}

	return ParseFile(filename)
}

type StaticProvider map[string]string

func (sp StaticProvider) Get(name string) (*Template, error) {
	if data, ok := sp[name]; ok {
		tmpl := Template{data, "{{", "}}", 0, 1, "", []interface{}{}, sp}
		err := tmpl.parse()
		if err != nil {
			return nil, err
		}

		return &tmpl, nil
	}

	return &Template{"", "{{", "}}", 0, 1, "", []interface{}{}, nil}, nil
}
<commit_msg>Use only public members<commit_after>package mustache

import (
	"os"
	"path"
)

type PartialProvider interface {
	Get(name string) (*Template, error)
}

type FileProvider struct {
	Paths      []string
	Extensions []string
}

func (fp *FileProvider) Get(name string) (*Template, error) {
	var filename string

	for _, p := range fp.Paths {
		for _, e := range fp.Extensions {
			name := path.Join(p, name+e)
			f, err := os.Open(name)
			if err == nil {
				filename = name
				f.Close()
				break
			}
		}
	}

	if filename == "" {
		return ParseString("")
	}

	return ParseFile(filename)
}

type StaticProvider map[string]string

func (sp StaticProvider) Get(name string) (*Template, error) {
	if data, ok := sp[name]; ok {
		return ParseStringPartials(data, sp)
	}

	return ParseString("")
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Arista Networks, Inc.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package path

import (
	"bytes"
	"fmt"
	"sort"

	"github.com/aristanetworks/goarista/key"
	"github.com/aristanetworks/goarista/pathmap"
)

// TODO: Update comments.

// Map associates Paths to values. It allows wildcards. The
// primary use of Map is to be able to register handlers to paths
// that can be efficiently looked up every time a path is updated.
//
// For example:
//
// m.Set({key.New("interfaces"), key.New("*"), key.New("adminStatus")}, AdminStatusHandler)
// m.Set({key.New("interface"), key.New("Management1"), key.New("adminStatus")},
//     Management1AdminStatusHandler)
//
// m.Visit(Path{key.New("interfaces"), key.New("Ethernet3/32/1"), key.New("adminStatus")},
//     HandlerExecutor)
// >> AdminStatusHandler gets passed to HandlerExecutor
// m.Visit(Path{key.New("interfaces"), key.New("Management1"), key.New("adminStatus")},
//    HandlerExecutor)
// >> AdminStatusHandler and Management1AdminStatusHandler gets passed to HandlerExecutor
//
// Note, Visit performance is typically linearly with the length of
// the path. But, it can be as bad as O(2^len(Path)) when TreeMap
// nodes have children and a wildcard associated with it. For example,
// if these paths were registered:
//
// m.Set(Path{key.New("foo"), key.New("bar"), key.New("baz")}, 1)
// m.Set(Path{key.New("*"), key.New("bar"), key.New("baz")}, 2)
// m.Set(Path{key.New("*"), key.New("*"), key.New("baz")}, 3)
// m.Set(Path{key.New("*"), key.New("*"), key.New("*")}, 4)
// m.Set(Path{key.New("foo"), key.New("*"), key.New("*")}, 5)
// m.Set(Path{key.New("foo"), key.New("bar"), key.New("*")}, 6)
// m.Set(Path{key.New("foo"), key.New("*"), key.New("baz")}, 7)
// m.Set(Path{key.New("*"), key.New("bar"), key.New("*")}, 8)
//
// m.Visit(Path{key.New("foo"),key.New("bar"),key.New("baz")}, Foo) // 2^3 nodes traversed
//
// This shouldn't be a concern with our paths because it is likely
// that a TreeMap node will either have a wildcard or children, not
// both. A TreeMap node that corresponds to a collection will often be a
// wildcard, otherwise it will have specific children.
type Map interface {
	// Visit calls f for every registration in the Map that
	// matches path. For example,
	//
	// m.Set(Path{key.New("foo"), key.New("bar")}, 1)
	// m.Set(Path{key.New("*"), key.New("bar")}, 2)
	//
	// m.Visit(Path{key.New("foo"), key.New("bar")}, Printer)
	// >> Calls Printer(1) and Printer(2)
	Visit(p Path, f pathmap.VisitorFunc) error

	// VisitPrefix calls f for every registration in the Map that
	// is a prefix of path. For example,
	//
	// m.Set(Path{}, 0)
	// m.Set(Path{key.New("foo")}, 1)
	// m.Set(Path{key.New("foo"), key.New("bar")}, 2)
	// m.Set(Path{key.New("foo"), key.New("quux")}, 3)
	// m.Set(Path{key.New("*"), key.New("bar")}, 4)
	//
	// m.VisitPrefix(Path{key.New("foo"), key.New("bar"), key.New("baz")}, Printer)
	// >> Calls Printer on values 0, 1, 2, and 4
	VisitPrefix(p Path, f pathmap.VisitorFunc) error

	// Get returns the mapping for path. This returns the exact
	// mapping for path. For example, if you register two paths
	//
	// m.Set(Path{key.New("foo"), key.New("bar")}, 1)
	// m.Set(Path{key.New("*"), key.New("bar")}, 2)
	//
	// m.Get(Path{key.New("foo"), key.New("bar")}) => 1
	// m.Get(Path{key.New("*"), key.New("bar")}) => 2
	Get(p Path) interface{}

	// Set a mapping of path to value. Path may contain wildcards. Set
	// replaces what was there before.
	Set(p Path, v interface{})

	// Delete removes the mapping for path
	Delete(p Path) bool
}

// Wildcard is a special key representing any possible path
var Wildcard = wildcard{}

type wildcard struct{}

func (w wildcard) Key() interface{} {
	return struct{}{}
}

func (w wildcard) String() string {
	return "*"
}

func (w wildcard) Equal(other interface{}) bool {
	_, ok := other.(wildcard)
	return ok
}

type node struct {
	val      interface{}
	wildcard *node
	children map[key.Key]*node
}

// NewMap creates a new Map
func NewMap() Map {
	return &node{}
}

// Visit calls f for every matching registration in the Map
func (n *node) Visit(p Path, f pathmap.VisitorFunc) error {
	for i, element := range p {
		if n.wildcard != nil {
			if err := n.wildcard.Visit(p[i+1:], f); err != nil {
				return err
			}
		}
		next, ok := n.children[element]
		if !ok {
			return nil
		}
		n = next
	}
	if n.val == nil {
		return nil
	}
	return f(n.val)
}

// VisitPrefix calls f for every registered path that is a prefix of
// the path
func (n *node) VisitPrefix(p Path, f pathmap.VisitorFunc) error {
	for i, element := range p {
		// Call f on each node we visit
		if n.val != nil {
			if err := f(n.val); err != nil {
				return err
			}
		}
		if n.wildcard != nil {
			if err := n.wildcard.VisitPrefix(p[i+1:], f); err != nil {
				return err
			}
		}
		next, ok := n.children[element]
		if !ok {
			return nil
		}
		n = next
	}
	if n.val == nil {
		return nil
	}
	// Call f on the final node
	return f(n.val)
}

// Get returns the mapping for path
func (n *node) Get(p Path) interface{} {
	for _, element := range p {
		if element.Equal(Wildcard) {
			if n.wildcard == nil {
				return nil
			}
			n = n.wildcard
			continue
		}
		next, ok := n.children[element]
		if !ok {
			return nil
		}
		n = next
	}
	return n.val
}

// Set a mapping of path to value. Path may contain wildcards. Set
// replaces what was there before.
func (n *node) Set(p Path, v interface{}) {
	for _, element := range p {
		if element.Equal(Wildcard) {
			if n.wildcard == nil {
				n.wildcard = &node{}
			}
			n = n.wildcard
			continue
		}
		if n.children == nil {
			n.children = map[key.Key]*node{}
		}
		next, ok := n.children[element]
		if !ok {
			next = &node{}
			n.children[element] = next
		}
		n = next
	}
	n.val = v
}

// Delete removes the mapping for path
func (n *node) Delete(p Path) bool {
	nodes := make([]*node, len(p)+1)
	for i, element := range p {
		nodes[i] = n
		if element.Equal(Wildcard) {
			if n.wildcard == nil {
				return false
			}
			n = n.wildcard
			continue
		}
		next, ok := n.children[element]
		if !ok {
			return false
		}
		n = next
	}
	n.val = nil
	nodes[len(p)] = n

	// See if we can delete any node objects
	for i := len(p); i > 0; i-- {
		n = nodes[i]
		if n.val != nil || n.wildcard != nil || len(n.children) > 0 {
			break
		}
		parent := nodes[i-1]
		element := p[i-1]
		if element.Equal(Wildcard) {
			parent.wildcard = nil
		} else {
			delete(parent.children, element)
		}
	}
	return true
}

func (n *node) String() string {
	var b bytes.Buffer
	n.write(&b, "")
	return b.String()
}

func (n *node) write(b *bytes.Buffer, indent string) {
	if n.val != nil {
		b.WriteString(indent)
		fmt.Fprintf(b, "Val: %v", n.val)
		b.WriteString("\n")
	}
	if n.wildcard != nil {
		b.WriteString(indent)
		fmt.Fprintf(b, "Child %q:\n", Wildcard)
		n.wildcard.write(b, indent+"  ")
	}
	children := make([]key.Key, 0, len(n.children))
	for key := range n.children {
		children = append(children, key)
	}
	sort.Slice(children, func(i, j int) bool {
		return children[i].String() < children[j].String()
	})

	for _, key := range children {
		child := n.children[key]
		b.WriteString(indent)
		fmt.Fprintf(b, "Child %q:\n", key.String())
		child.write(b, indent+"  ")
	}
}
<commit_msg>path: Update comments on Map<commit_after>// Copyright (c) 2017 Arista Networks, Inc.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package path

import (
	"bytes"
	"fmt"
	"sort"

	"github.com/aristanetworks/goarista/key"
	"github.com/aristanetworks/goarista/pathmap"
)

// Map associates paths to values. It allows wildcards. A Map
// is primarily used to register handlers with paths that can
// be easily looked up each time a path is updated.
type Map interface {
	// Visit calls a function fn for every value in the Map
	// that is registered with a match of a path p. In the
	// general case, time complexity is linear with respect
	// to the length of p but it can be as bad as O(2^len(p))
	// if there are a lot of paths with wildcards registered.
	//
	// Example:
	//
	// a := path.New("foo", "bar", "baz")
	// b := path.New("foo", path.Wildcard, "baz")
	// c := path.New(path.Wildcard, "bar", "baz")
	// d := path.New("foo", "bar", path.Wildcard)
	// e := path.New(path.Wildcard, path.Wildcard, "baz")
	// f := path.New(path.Wildcard, "bar", path.Wildcard)
	// g := path.New("foo", path.Wildcard, path.Wildcard)
	// h := path.New(path.Wildcard, path.Wildcard, path.Wildcard)
	//
	// m.Set(a, 1)
	// m.Set(b, 2)
	// m.Set(c, 3)
	// m.Set(d, 4)
	// m.Set(e, 5)
	// m.Set(f, 6)
	// m.Set(g, 7)
	// m.Set(h, 8)
	//
	// p := path.New("foo", "bar", "baz")
	//
	// m.Visit(p, fn)
	//
	// Result: fn(1), fn(2), fn(3), fn(4), fn(5), fn(6), fn(7) and fn(8)
	Visit(p Path, fn pathmap.VisitorFunc) error

	// VisitPrefix calls a function fn for every value in the
	// Map that is registered with a prefix of a path p.
	//
	// Example:
	//
	// a := path.New()
	// b := path.New("foo")
	// c := path.New("foo", "bar")
	// d := path.New("foo", "baz")
	// e := path.New(path.Wildcard, "bar")
	//
	// m.Set(a, 1)
	// m.Set(b, 2)
	// m.Set(c, 3)
	// m.Set(d, 4)
	// m.Set(e, 5)
	//
	// p := path.New("foo", "bar", "baz")
	//
	// m.VisitPrefix(p, fn)
	//
	// Result: fn(1), fn(2), fn(3), fn(5)
	VisitPrefix(p Path, fn pathmap.VisitorFunc) error

	// Get returns the value registered with an exact match of a
	// path p. If there is no exact match for p, Get returns nil.
	//
	// Example:
	//
	// m.Set(path.New("foo", "bar"), 1)
	//
	// a := m.Get(path.New("foo", "bar"))
	// b := m.Get(path.New("foo", path.Wildcard))
	//
	// Result: a == 1 and b == nil
	Get(p Path) interface{}

	// Set registers a path p with a value. Any previous value that
	// was registered with p is overwritten.
	//
	// Example:
	//
	// p := path.New("foo", "bar")
	//
	// m.Set(p, 0)
	// m.Set(p, 1)
	//
	// v := m.Get(p)
	//
	// Result: v == 1
	Set(p Path, v interface{})

	// Delete unregisters the value registered with a path. It
	// returns true if a value was deleted and false otherwise.
	//
	// Example:
	//
	// p := path.New("foo", "bar")
	//
	// m.Set(p, 0)
	//
	// a := m.Delete(p)
	// b := m.Delete(p)
	//
	// Result: a == true and b == false
	Delete(p Path) bool
}

// Wildcard is a special key representing any possible path.
var Wildcard = wildcard{}

type wildcard struct{}

func (w wildcard) Key() interface{} {
	return struct{}{}
}

func (w wildcard) String() string {
	return "*"
}

func (w wildcard) Equal(other interface{}) bool {
	_, ok := other.(wildcard)
	return ok
}

type node struct {
	val      interface{}
	wildcard *node
	children map[key.Key]*node
}

// NewMap creates a new Map
func NewMap() Map {
	return &node{}
}

// Visit calls f for every matching registration in the Map
func (n *node) Visit(p Path, f pathmap.VisitorFunc) error {
	for i, element := range p {
		if n.wildcard != nil {
			if err := n.wildcard.Visit(p[i+1:], f); err != nil {
				return err
			}
		}
		next, ok := n.children[element]
		if !ok {
			return nil
		}
		n = next
	}
	if n.val == nil {
		return nil
	}
	return f(n.val)
}

// VisitPrefix calls f for every registered path that is a prefix of
// the path
func (n *node) VisitPrefix(p Path, f pathmap.VisitorFunc) error {
	for i, element := range p {
		if n.val != nil {
			if err := f(n.val); err != nil {
				return err
			}
		}
		if n.wildcard != nil {
			if err := n.wildcard.VisitPrefix(p[i+1:], f); err != nil {
				return err
			}
		}
		next, ok := n.children[element]
		if !ok {
			return nil
		}
		n = next
	}
	if n.val == nil {
		return nil
	}
	return f(n.val)
}

// Get returns the mapping for path
func (n *node) Get(p Path) interface{} {
	for _, element := range p {
		if element.Equal(Wildcard) {
			if n.wildcard == nil {
				return nil
			}
			n = n.wildcard
			continue
		}
		next, ok := n.children[element]
		if !ok {
			return nil
		}
		n = next
	}
	return n.val
}

// Set a mapping of path to value. Path may contain wildcards. Set
// replaces what was there before.
func (n *node) Set(p Path, v interface{}) {
	for _, element := range p {
		if element.Equal(Wildcard) {
			if n.wildcard == nil {
				n.wildcard = &node{}
			}
			n = n.wildcard
			continue
		}
		if n.children == nil {
			n.children = map[key.Key]*node{}
		}
		next, ok := n.children[element]
		if !ok {
			next = &node{}
			n.children[element] = next
		}
		n = next
	}
	n.val = v
}

// Delete removes the mapping for path
func (n *node) Delete(p Path) bool {
	nodes := make([]*node, len(p)+1)
	for i, element := range p {
		nodes[i] = n
		if element.Equal(Wildcard) {
			if n.wildcard == nil {
				return false
			}
			n = n.wildcard
			continue
		}
		next, ok := n.children[element]
		if !ok {
			return false
		}
		n = next
	}
	n.val = nil
	nodes[len(p)] = n

	// See if we can delete any node objects.
	for i := len(p); i > 0; i-- {
		n = nodes[i]
		if n.val != nil || n.wildcard != nil || len(n.children) > 0 {
			break
		}
		parent := nodes[i-1]
		element := p[i-1]
		if element.Equal(Wildcard) {
			parent.wildcard = nil
		} else {
			delete(parent.children, element)
		}
	}
	return true
}

func (n *node) String() string {
	var b bytes.Buffer
	n.write(&b, "")
	return b.String()
}

func (n *node) write(b *bytes.Buffer, indent string) {
	if n.val != nil {
		b.WriteString(indent)
		fmt.Fprintf(b, "Val: %v", n.val)
		b.WriteString("\n")
	}
	if n.wildcard != nil {
		b.WriteString(indent)
		fmt.Fprintf(b, "Child %q:\n", Wildcard)
		n.wildcard.write(b, indent+"  ")
	}
	children := make([]key.Key, 0, len(n.children))
	for key := range n.children {
		children = append(children, key)
	}
	sort.Slice(children, func(i, j int) bool {
		return children[i].String() < children[j].String()
	})

	for _, key := range children {
		child := n.children[key]
		b.WriteString(indent)
		fmt.Fprintf(b, "Child %q:\n", key.String())
		child.write(b, indent+"  ")
	}
}
<|endoftext|>


<commit_before>package scoring

import (
	"fmt"
	"math"
	"time"
)

// A point of reference Score.Update and Score.Relevance use to reference the
// current time. It is used in testing, so we always have the same current
// time. This is okay for this programs as it won't run for long.
var Now time.Time

// Represents a weight of a score and the age of it.
type Score struct {
	Weight int64
	Age    time.Time
}

// Update the weight and age of the current score.
func (s *Score) Update() {
	s.Weight++
	s.Age = Now
}

// Relevance of a score is the difference between the current time and when the
// score was last updated.
func (s *Score) Relevance() time.Duration {
	return Now.Sub(s.Age)
}

// Calculate the final score from the score weight and the age.
func (s *Score) Calculate() float64 {
	return float64(s.Weight) * math.Log(float64(s.Relevance()))
}

// Calculate the final score from the score weight and the age.
func (s *Score) String() string {
	return fmt.Sprintf("{%s %s}", s.Weight, s.Age)
}

// Create a new score object with default weight of 1 and age set to now.
func NewScore() *Score {
	return &Score{1, Now}
}

func init() {
	Now = time.Now()
}
<commit_msg>Make sure s.Relevance is normalized (0, 1)<commit_after>package scoring

import (
	"fmt"
	"math"
	"time"
)

// A point of reference Score.Update and Score.Relevance use to reference the
// current time. It is used in testing, so we always have the same current
// time. This is okay for this programs as it won't run for long.
var Now time.Time

// Represents a weight of a score and the age of it.
type Score struct {
	Weight int64
	Age    time.Time
}

// Update the weight and age of the current score.
func (s *Score) Update() {
	s.Weight++
	s.Age = Now
}

// Relevance of a score is the difference between the current time and when the
// score was last updated.
func (s *Score) Relevance() float64 {
	return float64(s.Age.Unix()) / float64(Now.Unix())
}

// Calculate the final score from the score weight and the age.
func (s *Score) Calculate() float64 {
	return float64(s.Weight) * math.Log(s.Relevance())
}

// Calculate the final score from the score weight and the age.
func (s *Score) String() string {
	return fmt.Sprintf("{%s %s}", s.Weight, s.Age)
}

// Create a new score object with default weight of 1 and age set to now.
func NewScore() *Score {
	return &Score{1, Now}
}

func init() {
	Now = time.Now()
}
<|endoftext|>


<commit_before>package quorum

import (
	"log"

	"gopkg.in/errgo.v1"
)

type memSender struct {
	mboxes map[string]memMbox
}

type memMbox chan Ballot

// NewMemSender returns a new in-memory implementation of Sender.
func NewMemSender(names ...string) Sender {
	sender := &memSender{
		mboxes: map[string]memMbox{},
	}
	return sender
}

// ValidateRecipient implements Sender by verifying that the recipient is
// registered with the sender.
func (s *memSender) ValidateRecipient(recipient string) error {
	_, ok := s.mboxes[recipient]
	if !ok {
		return ErrNotFound
	}
	return nil
}

// Send implements Sender by sending a ballot to the intended recipient.
func (s *memSender) Send(ballot Ballot) error {
	mbox, ok := s.mboxes[ballot.Recipient]
	if !ok {
		return ErrNotFound
	}
	select {
	case mbox <- ballot:
		return nil
	}
	return errgo.Newf("%q isn't receiving messages", ballot.Recipient)
}

// Close releases all resources used by the Sender, and unregisters all
// recipients.
func (s *memSender) Close() {
	for _, mbox := range s.mboxes {
		close(mbox)
	}
	s.mboxes = map[string]memMbox{}
}

// Register registers a recipient with the Sender.
func (s *memSender) Register(recipient string, handler func(Ballot) error) error {
	mbox, ok := s.mboxes[recipient]
	if ok {
		close(mbox)
	}
	mbox = make(memMbox)
	s.mboxes[recipient] = mbox
	go func() {
		for {
			select {
			case ballot, ok := <-mbox:
				if !ok {
					return
				}
				err := handler(ballot)
				if err != nil {
					log.Println("error handling ballot for recipient %q: %v", recipient, errgo.Details(err))
				}
			}
		}
	}()
	return nil
}
<commit_msg>Add copyright.<commit_after>/*
 * Copyright 2015 Casey Marshall
 *
 * 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 quorum

import (
	"log"

	"gopkg.in/errgo.v1"
)

type memSender struct {
	mboxes map[string]memMbox
}

type memMbox chan Ballot

// NewMemSender returns a new in-memory implementation of Sender.
func NewMemSender(names ...string) Sender {
	sender := &memSender{
		mboxes: map[string]memMbox{},
	}
	return sender
}

// ValidateRecipient implements Sender by verifying that the recipient is
// registered with the sender.
func (s *memSender) ValidateRecipient(recipient string) error {
	_, ok := s.mboxes[recipient]
	if !ok {
		return ErrNotFound
	}
	return nil
}

// Send implements Sender by sending a ballot to the intended recipient.
func (s *memSender) Send(ballot Ballot) error {
	mbox, ok := s.mboxes[ballot.Recipient]
	if !ok {
		return ErrNotFound
	}
	select {
	case mbox <- ballot:
		return nil
	}
	return errgo.Newf("%q isn't receiving messages", ballot.Recipient)
}

// Close releases all resources used by the Sender, and unregisters all
// recipients.
func (s *memSender) Close() {
	for _, mbox := range s.mboxes {
		close(mbox)
	}
	s.mboxes = map[string]memMbox{}
}

// Register registers a recipient with the Sender.
func (s *memSender) Register(recipient string, handler func(Ballot) error) error {
	mbox, ok := s.mboxes[recipient]
	if ok {
		close(mbox)
	}
	mbox = make(memMbox)
	s.mboxes[recipient] = mbox
	go func() {
		for {
			select {
			case ballot, ok := <-mbox:
				if !ok {
					return
				}
				err := handler(ballot)
				if err != nil {
					log.Println("error handling ballot for recipient %q: %v", recipient, errgo.Details(err))
				}
			}
		}
	}()
	return nil
}
<|endoftext|>


<commit_before>// Memstream is an expandable ReadWriteSeeker for Golang that works with an
// internally managed byte buffer. Operation is usage is intended to be seamless
// and smooth.
//
// In situations where the maximum read/write sizes are known, a fixed
// []byte/byte buffer will likely offer better performance.
package memstream

import (
	"errors"
	"fmt"
	"io"
)

// MemoryStream is a memory-based, automatically resizing stream that can
// easily fill the role of any file-based IO.
type MemoryStream struct {
	buff []byte
	loc  int
}

// DefaultCapacity is the size in bytes of a new MemoryStream's backing buffer
const DefaultCapacity = 512

// New creates a new MemoryStream instance
func New() *MemoryStream {
	return NewCapacity(DefaultCapacity)
}

// NewCapacity starts the returned MemoryStream with the given capacity
func NewCapacity(cap int) *MemoryStream {
	return &MemoryStream{buff: make([]byte, 0, DefaultCapacity), loc: 0}
}

// Seek sets the offset for the next Read or Write to offset, interpreted
// according to whence: 0 means relative to the origin of the file, 1 means
// relative to the current offset, and 2 means relative to the end. Seek
// returns the new offset and an error, if any.
//
// Seeking to a negative offset is an error. Seeking to any positive offset is
// legal. If the location is beyond the end of the current length, the position
// will be placed at length.
func (m *MemoryStream) Seek(offset int64, whence int) (int64, error) {
	newLoc := m.loc
	switch whence {
	case 0:
		newLoc = int(offset)
	case 1:
		newLoc += whence
	case 2:
		newLoc = len(m.buff) - whence
	}

	if newLoc < 0 {
		return int64(m.loc), errors.New("Unable to seek to a location <0")
	}

	if newLoc > len(m.buff) {
		newLoc = len(m.buff)
	}

	m.loc = newLoc

	return int64(m.loc), nil
}

// Read puts up to len(p) bytes into p. Will return the number of bytes read.
func (m *MemoryStream) Read(p []byte) (n int, err error) {
	n = copy(p, m.buff[m.loc:len(m.buff)])
	m.loc += n

	if m.loc == len(m.buff) {
		return n, io.EOF
	}

	return n, nil
}

// Write writes the given bytes into the memory stream. If needed, the underlying
// buffer will be expanded to fit the new bytes.
func (m *MemoryStream) Write(p []byte) (n int, err error) {
	// Do we have space?
	if available := cap(m.buff) - m.loc; available < len(p) {
		// How much should we expand by?
		addCap := cap(m.buff)
		if addCap < len(p) {
			addCap = len(p)
		}

		newBuff := make([]byte, len(m.buff), cap(m.buff)+addCap)

		copy(newBuff, m.buff)

		fmt.Printf("Expanded to %v bytes from %v\n", cap(newBuff), cap(m.buff))

		m.buff = newBuff
	}

	// Write
	n = copy(m.buff[m.loc:cap(m.buff)], p)
	m.loc += n
	if len(m.buff) < m.loc {
		m.buff = m.buff[:m.loc]
	}

	return n, nil
}
<commit_msg>Correct format for package comment<commit_after>// Package memstream is an expandable ReadWriteSeeker for Golang that works with an
// internally managed byte buffer. Operation is usage is intended to be seamless
// and smooth.
//
// In situations where the maximum read/write sizes are known, a fixed
// []byte/byte buffer will likely offer better performance.
package memstream

import (
	"errors"
	"fmt"
	"io"
)

// MemoryStream is a memory-based, automatically resizing stream that can
// easily fill the role of any file-based IO.
type MemoryStream struct {
	buff []byte
	loc  int
}

// DefaultCapacity is the size in bytes of a new MemoryStream's backing buffer
const DefaultCapacity = 512

// New creates a new MemoryStream instance
func New() *MemoryStream {
	return NewCapacity(DefaultCapacity)
}

// NewCapacity starts the returned MemoryStream with the given capacity
func NewCapacity(cap int) *MemoryStream {
	return &MemoryStream{buff: make([]byte, 0, DefaultCapacity), loc: 0}
}

// Seek sets the offset for the next Read or Write to offset, interpreted
// according to whence: 0 means relative to the origin of the file, 1 means
// relative to the current offset, and 2 means relative to the end. Seek
// returns the new offset and an error, if any.
//
// Seeking to a negative offset is an error. Seeking to any positive offset is
// legal. If the location is beyond the end of the current length, the position
// will be placed at length.
func (m *MemoryStream) Seek(offset int64, whence int) (int64, error) {
	newLoc := m.loc
	switch whence {
	case 0:
		newLoc = int(offset)
	case 1:
		newLoc += whence
	case 2:
		newLoc = len(m.buff) - whence
	}

	if newLoc < 0 {
		return int64(m.loc), errors.New("Unable to seek to a location <0")
	}

	if newLoc > len(m.buff) {
		newLoc = len(m.buff)
	}

	m.loc = newLoc

	return int64(m.loc), nil
}

// Read puts up to len(p) bytes into p. Will return the number of bytes read.
func (m *MemoryStream) Read(p []byte) (n int, err error) {
	n = copy(p, m.buff[m.loc:len(m.buff)])
	m.loc += n

	if m.loc == len(m.buff) {
		return n, io.EOF
	}

	return n, nil
}

// Write writes the given bytes into the memory stream. If needed, the underlying
// buffer will be expanded to fit the new bytes.
func (m *MemoryStream) Write(p []byte) (n int, err error) {
	// Do we have space?
	if available := cap(m.buff) - m.loc; available < len(p) {
		// How much should we expand by?
		addCap := cap(m.buff)
		if addCap < len(p) {
			addCap = len(p)
		}

		newBuff := make([]byte, len(m.buff), cap(m.buff)+addCap)

		copy(newBuff, m.buff)

		fmt.Printf("Expanded to %v bytes from %v\n", cap(newBuff), cap(m.buff))

		m.buff = newBuff
	}

	// Write
	n = copy(m.buff[m.loc:cap(m.buff)], p)
	m.loc += n
	if len(m.buff) < m.loc {
		m.buff = m.buff[:m.loc]
	}

	return n, nil
}
<|endoftext|>


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

package worker

import (
	"github.com/Matir/webborer/logging"
	"github.com/Matir/webborer/util"
	"github.com/Matir/webborer/workqueue"
	"golang.org/x/net/html"
	"io"
	"net/http"
	"net/url"
	"strings"
)

type HTMLWorker struct {
	// Function to add future work
	adder workqueue.QueueAddFunc
}

func NewHTMLWorker(adder workqueue.QueueAddFunc) *HTMLWorker {
	return &HTMLWorker{adder: adder}
}

// Work on this response
func (w *HTMLWorker) Handle(URL *url.URL, body io.Reader) {
	links := w.GetLinks(body)
	foundURLs := make([]*url.URL, 0, len(links))
	for _, l := range links {
		u, err := url.Parse(l)
		if err != nil {
			logging.Logf(logging.LogInfo, "Error parsing URL (%s): %s", l, err.Error())
			continue
		}
		resolved := URL.ResolveReference(u)
		foundURLs = append(foundURLs, resolved)
		// Include parents of the found URL.
		// Worker will remove duplicates
		foundURLs = append(foundURLs, util.GetParentPaths(resolved)...)
	}
	w.adder(foundURLs...)
}

// Check if this response can be handled by this worker
func (*HTMLWorker) Eligible(resp *http.Response) bool {
	ct := resp.Header.Get("Content-type")
	if strings.ToLower(ct) != "text/html" {
		return false
	}
	return resp.ContentLength > 0 && resp.ContentLength < 1024*1024
}

// Get the links for the body.
func (*HTMLWorker) GetLinks(body io.Reader) []string {
	tree, err := html.Parse(body)
	if err != nil {
		logging.Logf(logging.LogInfo, "Unable to parse HTML document: %s", err.Error())
		return nil
	}
	links := collectElementAttributes(tree, "a", "href")
	links = append(links, collectElementAttributes(tree, "img", "src")...)
	links = append(links, collectElementAttributes(tree, "script", "src")...)
	links = append(links, collectElementAttributes(tree, "style", "src")...)
	return util.DedupeStrings(links)
}

func getElementsByTagName(root *html.Node, name string) []*html.Node {
	results := make([]*html.Node, 0)
	var handleNode func(*html.Node)
	handleNode = func(node *html.Node) {
		if node.Type == html.ElementNode && strings.ToLower(node.Data) == name {
			results = append(results, node)
		}
		for n := node.FirstChild; n != nil; n = n.NextSibling {
			handleNode(n)
		}
	}
	handleNode(root)
	return results
}

func getElementAttribute(node *html.Node, attrName string) *string {
	for _, a := range node.Attr {
		if strings.ToLower(a.Key) == attrName {
			return &a.Val
		}
	}
	return nil
}

func collectElementAttributes(root *html.Node, tagName, attrName string) []string {
	results := make([]string, 0)
	for _, el := range getElementsByTagName(root, tagName) {
		if val := getElementAttribute(el, attrName); val != nil {
			results = append(results, *val)
		}
	}
	return results
}
<commit_msg>Better job of spidering HTML.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package worker

import (
	"github.com/Matir/webborer/logging"
	"github.com/Matir/webborer/util"
	"github.com/Matir/webborer/workqueue"
	"golang.org/x/net/html"
	"io"
	"net/http"
	"net/url"
	"strings"
)

const (
  maxHTMLWorkerSize = 10*1024*1024
)

type HTMLWorker struct {
	// Function to add future work
	adder workqueue.QueueAddFunc
}

func NewHTMLWorker(adder workqueue.QueueAddFunc) *HTMLWorker {
	return &HTMLWorker{adder: adder}
}

// Work on this response
func (w *HTMLWorker) Handle(URL *url.URL, body io.Reader) {
  limitedBody := io.LimitReader(body, maxHTMLWorkerSize)
	links := w.GetLinks(limitedBody)
	foundURLs := make([]*url.URL, 0, len(links))
	for _, l := range links {
		u, err := url.Parse(l)
		if err != nil {
			logging.Logf(logging.LogInfo, "Error parsing URL (%s): %s", l, err.Error())
			continue
		}
		resolved := URL.ResolveReference(u)
		foundURLs = append(foundURLs, resolved)
		// Include parents of the found URL.
		// Worker will remove duplicates
		foundURLs = append(foundURLs, util.GetParentPaths(resolved)...)
	}
	w.adder(foundURLs...)
}

// Check if this response can be handled by this worker
func (*HTMLWorker) Eligible(resp *http.Response) bool {
	ct := resp.Header.Get("Content-type")
	if strings.ToLower(ct) != "text/html" {
		return false
	}
	// ContentLength is often -1, indicating unknown, so we'll try to parse those
	return resp.ContentLength == -1 || (
	  resp.ContentLength > 0 && resp.ContentLength < maxHTMLWorkerSize)
}

// Get the links for the body.
func (*HTMLWorker) GetLinks(body io.Reader) []string {
	tree, err := html.Parse(body)
	if err != nil {
		logging.Logf(logging.LogInfo, "Unable to parse HTML document: %s", err.Error())
		return nil
	}
	links := collectElementAttributes(tree, "a", "href")
	links = append(links, collectElementAttributes(tree, "img", "src")...)
	links = append(links, collectElementAttributes(tree, "script", "src")...)
	links = append(links, collectElementAttributes(tree, "style", "src")...)
	return util.DedupeStrings(links)
}

func getElementsByTagName(root *html.Node, name string) []*html.Node {
	results := make([]*html.Node, 0)
	var handleNode func(*html.Node)
	handleNode = func(node *html.Node) {
		if node.Type == html.ElementNode && strings.ToLower(node.Data) == name {
			results = append(results, node)
		}
		for n := node.FirstChild; n != nil; n = n.NextSibling {
			handleNode(n)
		}
	}
	handleNode(root)
	return results
}

func getElementAttribute(node *html.Node, attrName string) *string {
	for _, a := range node.Attr {
		if strings.ToLower(a.Key) == attrName {
			return &a.Val
		}
	}
	return nil
}

func collectElementAttributes(root *html.Node, tagName, attrName string) []string {
	results := make([]string, 0)
	for _, el := range getElementsByTagName(root, tagName) {
		if val := getElementAttribute(el, attrName); val != nil {
			results = append(results, *val)
		}
	}
	return results
}
<|endoftext|>