<commit_before>package test

import (
	"fmt"
	"os"
	"os/exec"
	"path"
	"testing"
	"time"

	"go.pedge.io/dlog"

	"github.com/libopenstorage/openstorage/api"
	"github.com/libopenstorage/openstorage/volume"
	"github.com/portworx/kvdb"
	"github.com/portworx/kvdb/mem"
	"github.com/stretchr/testify/require"
)

func init() {
	kv, err := kvdb.New(mem.Name, "driver_test", []string{}, nil, dlog.Panicf)
	if err != nil {
		dlog.Panicf("Failed to intialize KVDB")
	}
	err = kvdb.SetInstance(kv)
	if err != nil {
		dlog.Panicf("Failed to set KVDB instance")
	}
}

// Context maintains current device state. It gets passed into tests
// so that tests can build on other tests' work
type Context struct {
	volume.VolumeDriver
	volID      string
	snapID     string
	mountPath  string
	devicePath string
	Filesystem api.FSType
	testPath   string
	testFile   string
	Encrypted  bool
}

func NewContext(d volume.VolumeDriver) *Context {
	return &Context{
		VolumeDriver: d,
		volID:        "",
		snapID:       "",
		Filesystem:   api.FSType_FS_TYPE_NONE,
		testPath:     path.Join("/tmp/openstorage/mount/", d.Name()),
		testFile:     path.Join("/tmp/", d.Name()),
	}
}

func RunShort(t *testing.T, ctx *Context) {
	create(t, ctx)
	inspect(t, ctx)
	set(t, ctx)
	enumerate(t, ctx)
	attach(t, ctx)
	mount(t, ctx)
	io(t, ctx)
	unmount(t, ctx)
	detach(t, ctx)
	delete(t, ctx)
	runEnd(t, ctx)
}

func Run(t *testing.T, ctx *Context) {
	RunShort(t, ctx)
	RunSnap(t, ctx)
	runEnd(t, ctx)
}

func runEnd(t *testing.T, ctx *Context) {
	time.Sleep(time.Second * 2)
	os.RemoveAll(ctx.testPath)
	os.Remove(ctx.testFile)
	shutdown(t, ctx)
}

func RunSnap(t *testing.T, ctx *Context) {
	snap(t, ctx)
	snapInspect(t, ctx)
	snapEnumerate(t, ctx)
	snapDiff(t, ctx)
	snapDelete(t, ctx)
	detach(t, ctx)
	delete(t, ctx)
}

func create(t *testing.T, ctx *Context) {
	fmt.Println("create")

	volID, err := ctx.Create(
		&api.VolumeLocator{Name: "foo", VolumeLabels: map[string]string{"oh": "create"}},
		nil,
		&api.VolumeSpec{
			Size:      1 * 1024 * 1024 * 1024,
			HaLevel:   1,
			Format:    ctx.Filesystem,
			Encrypted: ctx.Encrypted,
		})

	require.NoError(t, err, "Failed in Create")
	ctx.volID = volID
}

func inspect(t *testing.T, ctx *Context) {
	fmt.Println("inspect")

	vols, err := ctx.Inspect([]string{ctx.volID})
	require.NoError(t, err, "Failed in Inspect")
	require.NotNil(t, vols, "Nil vols")
	require.Equal(t, len(vols), 1, "Expect 1 volume actual %v volumes", len(vols))
	require.Equal(t, vols[0].Id, ctx.volID, "Expect volID %v actual %v", ctx.volID, vols[0].Id)

	vols, err = ctx.Inspect([]string{string("shouldNotExist")})
	require.Equal(t, 0, len(vols), "Expect 0 volume actual %v volumes", len(vols))
}

func set(t *testing.T, ctx *Context) {
	fmt.Println("update")

	vols, err := ctx.Inspect([]string{ctx.volID})
	require.NoError(t, err, "Failed in Inspect")
	require.NotNil(t, vols, "Nil vols")
	require.Equal(t, len(vols), 1, "Expect 1 volume actual %v volumes", len(vols))
	require.Equal(t, vols[0].Id, ctx.volID, "Expect volID %v actual %v", ctx.volID, vols[0].Id)

	vols[0].Locator.VolumeLabels["UpdateTest"] = "Success"
	err = ctx.Set(ctx.volID, vols[0].Locator, nil)
	require.NoError(t, err, "Failed in Update")

	vols, err = ctx.Inspect([]string{ctx.volID})
	require.NoError(t, err, "Failed in Inspect")
	require.NotNil(t, vols, "Nil vols")
	require.Equal(t, len(vols), 1, "Expect 1 volume actual %v volumes", len(vols))
	require.Equal(t, vols[0].Locator.VolumeLabels["UpdateTest"], "Success",
		"Expect Label %v actual %v", "UpdateTest", vols[0].Locator.VolumeLabels)
}

func enumerate(t *testing.T, ctx *Context) {
	fmt.Println("enumerate")

	vols, err := ctx.Enumerate(&api.VolumeLocator{}, nil)
	require.NoError(t, err, "Failed in Enumerate")
	require.NotNil(t, vols, "Nil vols")
	require.Equal(t, 1, len(vols), "Expect 1 volume actual %v volumes", len(vols))
	require.Equal(t, vols[0].Id, ctx.volID, "Expect volID %v actual %v", ctx.volID, vols[0].Id)

	vols, err = ctx.Enumerate(&api.VolumeLocator{Name: "foo"}, nil)
	require.NoError(t, err, "Failed in Enumerate")
	require.NotNil(t, vols, "Nil vols")
	require.Equal(t, len(vols), 1, "Expect 1 volume actual %v volumes", len(vols))
	require.Equal(t, vols[0].Id, ctx.volID, "Expect volID %v actual %v", ctx.volID, vols[0].Id)

	vols, err = ctx.Enumerate(&api.VolumeLocator{Name: "shouldNotExist"}, nil)
	require.Equal(t, len(vols), 0, "Expect 0 volume actual %v volumes", len(vols))
}

func waitReady(t *testing.T, ctx *Context) error {
	total := time.Minute * 5
	inc := time.Second * 2
	elapsed := time.Second * 0
	vols, err := ctx.Inspect([]string{ctx.volID})
	for err == nil && len(vols) == 1 && vols[0].Status != api.VolumeStatus_VOLUME_STATUS_UP && elapsed < total {
		time.Sleep(inc)
		elapsed += inc
		vols, err = ctx.Inspect([]string{ctx.volID})
	}
	if err != nil {
		return err
	}
	if len(vols) != 1 {
		return fmt.Errorf("Expect one volume from inspect got %v", len(vols))
	}
	if vols[0].Status != api.VolumeStatus_VOLUME_STATUS_UP {
		return fmt.Errorf("Timed out waiting for volume status %v", vols)
	}
	return err
}

func attach(t *testing.T, ctx *Context) {
	fmt.Println("attach")
	err := waitReady(t, ctx)
	require.NoError(t, err, "Volume status is not up")
	p, err := ctx.Attach(ctx.volID)
	if err != nil {
		require.Equal(t, err, volume.ErrNotSupported, "Error on attach %v", err)
	}
	ctx.devicePath = p

	p, err = ctx.Attach(ctx.volID)
	if err == nil {
		require.Equal(t, p, ctx.devicePath, "Multiple calls to attach if not errored should return the same path")
	}
}

func detach(t *testing.T, ctx *Context) {
	fmt.Println("detach")
	err := ctx.Detach(ctx.volID)
	if err != nil {
		require.Equal(t, ctx.devicePath, "", "Error on detach %s: %v", ctx.devicePath, err)
	}
	ctx.devicePath = ""
}

func mount(t *testing.T, ctx *Context) {
	fmt.Println("mount")

	err := os.MkdirAll(ctx.testPath, 0755)

	err = ctx.Mount(ctx.volID, ctx.testPath)
	require.NoError(t, err, "Failed in mount %v", ctx.testPath)

	ctx.mountPath = ctx.testPath
}

func multiMount(t *testing.T, ctx *Context) {
	ctx2 := *ctx
	create(t, ctx)
	attach(t, ctx)
	mount(t, ctx)

	create(t, &ctx2)
	attach(t, &ctx2)

	err := ctx2.Mount(ctx2.volID, ctx2.testPath)
	require.Error(t, err, "Mount of different devices to same path must fail")

	unmount(t, ctx)
	detach(t, ctx)
	delete(t, ctx)

	detach(t, &ctx2)
	delete(t, &ctx2)
}

func unmount(t *testing.T, ctx *Context) {
	fmt.Println("unmount")

	require.NotEqual(t, ctx.mountPath, "", "Device is not mounted")

	err := ctx.Unmount(ctx.volID, ctx.mountPath)
	require.NoError(t, err, "Failed in unmount %v", ctx.mountPath)

	ctx.mountPath = ""
}

func shutdown(t *testing.T, ctx *Context) {
	fmt.Println("shutdown")
	ctx.Shutdown()
}

func io(t *testing.T, ctx *Context) {
	fmt.Println("io")
	require.NotEqual(t, ctx.mountPath, "", "Device is not mounted")

	cmd := exec.Command("dd", "if=/dev/urandom", fmt.Sprintf("of=%s", ctx.testFile), "bs=1M", "count=10")
	o, err := cmd.CombinedOutput()
	require.NoError(t, err, "Failed to run dd %s", string(o))

	cmd = exec.Command("dd", fmt.Sprintf("if=%s", ctx.testFile), fmt.Sprintf("of=%s/xx", ctx.mountPath))
	o, err = cmd.CombinedOutput()
	require.NoError(t, err, "Failed to run dd on mountpoint %s/xx : %s",
		ctx.mountPath, string(o))

	cmd = exec.Command("diff", ctx.testFile, fmt.Sprintf("%s/xx", ctx.mountPath))
	o, err = cmd.CombinedOutput()
	require.NoError(t, err, "data mismatch %s", string(o))
}

func detachBad(t *testing.T, ctx *Context) {
	err := ctx.Detach(ctx.volID)
	require.True(t, (err == nil || err == volume.ErrNotSupported),
		"Detach on mounted device should fail")
}

func deleteBad(t *testing.T, ctx *Context) {
	fmt.Println("deleteBad")
	require.NotEqual(t, ctx.mountPath, "", "Device is not mounted")

	err := ctx.Delete(ctx.volID)
	require.Error(t, err, "Delete on mounted device must fail")
}

func delete(t *testing.T, ctx *Context) {
	fmt.Println("delete")
	err := ctx.Delete(ctx.volID)
	require.NoError(t, err, "Delete failed")
	ctx.volID = ""
}

func snap(t *testing.T, ctx *Context) {
	fmt.Println("snap")
	if ctx.volID == "" {
		create(t, ctx)
	}
	attach(t, ctx)
	labels := map[string]string{"oh": "snap"}
	require.NotEqual(t, ctx.volID, "", "invalid volume ID")
	id, err := ctx.Snapshot(ctx.volID, false,
		&api.VolumeLocator{Name: "snappy", VolumeLabels: labels})
	require.NoError(t, err, "Failed in creating a snapshot")
	ctx.snapID = id
}

func snapInspect(t *testing.T, ctx *Context) {
	fmt.Println("snapInspect")

	snaps, err := ctx.Inspect([]string{ctx.snapID})
	require.NoError(t, err, "Failed in Inspect")
	require.NotNil(t, snaps, "Nil snaps")
	require.Equal(t, len(snaps), 1, "Expect 1 snaps actual %v snaps", len(snaps))
	require.Equal(t, snaps[0].Id, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].Id)

	snaps, err = ctx.Inspect([]string{string("shouldNotExist")})
	require.Equal(t, 0, len(snaps), "Expect 0 snaps actual %v snaps", len(snaps))
}

func snapEnumerate(t *testing.T, ctx *Context) {
	fmt.Println("snapEnumerate")

	snaps, err := ctx.SnapEnumerate(nil, nil)
	require.NoError(t, err, "Failed in snapEnumerate")
	require.NotNil(t, snaps, "Nil snaps")
	require.Equal(t, 1, len(snaps), "Expect 1 snaps actual %v snaps", len(snaps))
	require.Equal(t, snaps[0].Id, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].Id)
	labels := snaps[0].Locator.VolumeLabels

	snaps, err = ctx.SnapEnumerate([]string{ctx.volID}, nil)
	require.NoError(t, err, "Failed in snapEnumerate")
	require.NotNil(t, snaps, "Nil snaps")
	require.Equal(t, len(snaps), 1, "Expect 1 snap actual %v snaps", len(snaps))
	require.Equal(t, snaps[0].Id, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].Id)

	snaps, err = ctx.SnapEnumerate([]string{string("shouldNotExist")}, nil)
	require.Equal(t, len(snaps), 0, "Expect 0 snap actual %v snaps", len(snaps))

	snaps, err = ctx.SnapEnumerate(nil, labels)
	require.NoError(t, err, "Failed in snapEnumerate")
	require.NotNil(t, snaps, "Nil snaps")
	require.Equal(t, len(snaps), 1, "Expect 1 snap actual %v snaps", len(snaps))
	require.Equal(t, snaps[0].Id, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].Id)
}

func snapDiff(t *testing.T, ctx *Context) {
	fmt.Println("snapDiff")
}

func snapDelete(t *testing.T, ctx *Context) {
	fmt.Println("snapDelete")
}
<commit_msg>Add passphrase to test<commit_after>package test

import (
	"fmt"
	"os"
	"os/exec"
	"path"
	"testing"
	"time"

	"go.pedge.io/dlog"

	"github.com/libopenstorage/openstorage/api"
	"github.com/libopenstorage/openstorage/volume"
	"github.com/portworx/kvdb"
	"github.com/portworx/kvdb/mem"
	"github.com/stretchr/testify/require"
)

func init() {
	kv, err := kvdb.New(mem.Name, "driver_test", []string{}, nil, dlog.Panicf)
	if err != nil {
		dlog.Panicf("Failed to intialize KVDB")
	}
	err = kvdb.SetInstance(kv)
	if err != nil {
		dlog.Panicf("Failed to set KVDB instance")
	}
}

// Context maintains current device state. It gets passed into tests
// so that tests can build on other tests' work
type Context struct {
	volume.VolumeDriver
	volID      string
	snapID     string
	mountPath  string
	devicePath string
	Filesystem api.FSType
	testPath   string
	testFile   string
	Encrypted  bool
	Passphrase string
}

func NewContext(d volume.VolumeDriver) *Context {
	return &Context{
		VolumeDriver: d,
		volID:        "",
		snapID:       "",
		Filesystem:   api.FSType_FS_TYPE_NONE,
		testPath:     path.Join("/tmp/openstorage/mount/", d.Name()),
		testFile:     path.Join("/tmp/", d.Name()),
	}
}

func RunShort(t *testing.T, ctx *Context) {
	create(t, ctx)
	inspect(t, ctx)
	set(t, ctx)
	enumerate(t, ctx)
	attach(t, ctx)
	mount(t, ctx)
	io(t, ctx)
	unmount(t, ctx)
	detach(t, ctx)
	delete(t, ctx)
	runEnd(t, ctx)
}

func Run(t *testing.T, ctx *Context) {
	RunShort(t, ctx)
	RunSnap(t, ctx)
	runEnd(t, ctx)
}

func runEnd(t *testing.T, ctx *Context) {
	time.Sleep(time.Second * 2)
	os.RemoveAll(ctx.testPath)
	os.Remove(ctx.testFile)
	shutdown(t, ctx)
}

func RunSnap(t *testing.T, ctx *Context) {
	snap(t, ctx)
	snapInspect(t, ctx)
	snapEnumerate(t, ctx)
	snapDiff(t, ctx)
	snapDelete(t, ctx)
	detach(t, ctx)
	delete(t, ctx)
}

func create(t *testing.T, ctx *Context) {
	fmt.Println("create")

	volID, err := ctx.Create(
		&api.VolumeLocator{Name: "foo", VolumeLabels: map[string]string{"oh": "create"}},
		nil,
		&api.VolumeSpec{
			Size:       1 * 1024 * 1024 * 1024,
			HaLevel:    1,
			Format:     ctx.Filesystem,
			Encrypted:  ctx.Encrypted,
			Passphrase: ctx.Passphrase,
		})

	require.NoError(t, err, "Failed in Create")
	ctx.volID = volID
}

func inspect(t *testing.T, ctx *Context) {
	fmt.Println("inspect")

	vols, err := ctx.Inspect([]string{ctx.volID})
	require.NoError(t, err, "Failed in Inspect")
	require.NotNil(t, vols, "Nil vols")
	require.Equal(t, len(vols), 1, "Expect 1 volume actual %v volumes", len(vols))
	require.Equal(t, vols[0].Id, ctx.volID, "Expect volID %v actual %v", ctx.volID, vols[0].Id)

	vols, err = ctx.Inspect([]string{string("shouldNotExist")})
	require.Equal(t, 0, len(vols), "Expect 0 volume actual %v volumes", len(vols))
}

func set(t *testing.T, ctx *Context) {
	fmt.Println("update")

	vols, err := ctx.Inspect([]string{ctx.volID})
	require.NoError(t, err, "Failed in Inspect")
	require.NotNil(t, vols, "Nil vols")
	require.Equal(t, len(vols), 1, "Expect 1 volume actual %v volumes", len(vols))
	require.Equal(t, vols[0].Id, ctx.volID, "Expect volID %v actual %v", ctx.volID, vols[0].Id)

	vols[0].Locator.VolumeLabels["UpdateTest"] = "Success"
	err = ctx.Set(ctx.volID, vols[0].Locator, nil)
	require.NoError(t, err, "Failed in Update")

	vols, err = ctx.Inspect([]string{ctx.volID})
	require.NoError(t, err, "Failed in Inspect")
	require.NotNil(t, vols, "Nil vols")
	require.Equal(t, len(vols), 1, "Expect 1 volume actual %v volumes", len(vols))
	require.Equal(t, vols[0].Locator.VolumeLabels["UpdateTest"], "Success",
		"Expect Label %v actual %v", "UpdateTest", vols[0].Locator.VolumeLabels)
}

func enumerate(t *testing.T, ctx *Context) {
	fmt.Println("enumerate")

	vols, err := ctx.Enumerate(&api.VolumeLocator{}, nil)
	require.NoError(t, err, "Failed in Enumerate")
	require.NotNil(t, vols, "Nil vols")
	require.Equal(t, 1, len(vols), "Expect 1 volume actual %v volumes", len(vols))
	require.Equal(t, vols[0].Id, ctx.volID, "Expect volID %v actual %v", ctx.volID, vols[0].Id)

	vols, err = ctx.Enumerate(&api.VolumeLocator{Name: "foo"}, nil)
	require.NoError(t, err, "Failed in Enumerate")
	require.NotNil(t, vols, "Nil vols")
	require.Equal(t, len(vols), 1, "Expect 1 volume actual %v volumes", len(vols))
	require.Equal(t, vols[0].Id, ctx.volID, "Expect volID %v actual %v", ctx.volID, vols[0].Id)

	vols, err = ctx.Enumerate(&api.VolumeLocator{Name: "shouldNotExist"}, nil)
	require.Equal(t, len(vols), 0, "Expect 0 volume actual %v volumes", len(vols))
}

func waitReady(t *testing.T, ctx *Context) error {
	total := time.Minute * 5
	inc := time.Second * 2
	elapsed := time.Second * 0
	vols, err := ctx.Inspect([]string{ctx.volID})
	for err == nil && len(vols) == 1 && vols[0].Status != api.VolumeStatus_VOLUME_STATUS_UP && elapsed < total {
		time.Sleep(inc)
		elapsed += inc
		vols, err = ctx.Inspect([]string{ctx.volID})
	}
	if err != nil {
		return err
	}
	if len(vols) != 1 {
		return fmt.Errorf("Expect one volume from inspect got %v", len(vols))
	}
	if vols[0].Status != api.VolumeStatus_VOLUME_STATUS_UP {
		return fmt.Errorf("Timed out waiting for volume status %v", vols)
	}
	return err
}

func attach(t *testing.T, ctx *Context) {
	fmt.Println("attach")
	err := waitReady(t, ctx)
	require.NoError(t, err, "Volume status is not up")
	p, err := ctx.Attach(ctx.volID)
	if err != nil {
		require.Equal(t, err, volume.ErrNotSupported, "Error on attach %v", err)
	}
	ctx.devicePath = p

	p, err = ctx.Attach(ctx.volID)
	if err == nil {
		require.Equal(t, p, ctx.devicePath, "Multiple calls to attach if not errored should return the same path")
	}
}

func detach(t *testing.T, ctx *Context) {
	fmt.Println("detach")
	err := ctx.Detach(ctx.volID)
	if err != nil {
		require.Equal(t, ctx.devicePath, "", "Error on detach %s: %v", ctx.devicePath, err)
	}
	ctx.devicePath = ""
}

func mount(t *testing.T, ctx *Context) {
	fmt.Println("mount")

	err := os.MkdirAll(ctx.testPath, 0755)

	err = ctx.Mount(ctx.volID, ctx.testPath)
	require.NoError(t, err, "Failed in mount %v", ctx.testPath)

	ctx.mountPath = ctx.testPath
}

func multiMount(t *testing.T, ctx *Context) {
	ctx2 := *ctx
	create(t, ctx)
	attach(t, ctx)
	mount(t, ctx)

	create(t, &ctx2)
	attach(t, &ctx2)

	err := ctx2.Mount(ctx2.volID, ctx2.testPath)
	require.Error(t, err, "Mount of different devices to same path must fail")

	unmount(t, ctx)
	detach(t, ctx)
	delete(t, ctx)

	detach(t, &ctx2)
	delete(t, &ctx2)
}

func unmount(t *testing.T, ctx *Context) {
	fmt.Println("unmount")

	require.NotEqual(t, ctx.mountPath, "", "Device is not mounted")

	err := ctx.Unmount(ctx.volID, ctx.mountPath)
	require.NoError(t, err, "Failed in unmount %v", ctx.mountPath)

	ctx.mountPath = ""
}

func shutdown(t *testing.T, ctx *Context) {
	fmt.Println("shutdown")
	ctx.Shutdown()
}

func io(t *testing.T, ctx *Context) {
	fmt.Println("io")
	require.NotEqual(t, ctx.mountPath, "", "Device is not mounted")

	cmd := exec.Command("dd", "if=/dev/urandom", fmt.Sprintf("of=%s", ctx.testFile), "bs=1M", "count=10")
	o, err := cmd.CombinedOutput()
	require.NoError(t, err, "Failed to run dd %s", string(o))

	cmd = exec.Command("dd", fmt.Sprintf("if=%s", ctx.testFile), fmt.Sprintf("of=%s/xx", ctx.mountPath))
	o, err = cmd.CombinedOutput()
	require.NoError(t, err, "Failed to run dd on mountpoint %s/xx : %s",
		ctx.mountPath, string(o))

	cmd = exec.Command("diff", ctx.testFile, fmt.Sprintf("%s/xx", ctx.mountPath))
	o, err = cmd.CombinedOutput()
	require.NoError(t, err, "data mismatch %s", string(o))
}

func detachBad(t *testing.T, ctx *Context) {
	err := ctx.Detach(ctx.volID)
	require.True(t, (err == nil || err == volume.ErrNotSupported),
		"Detach on mounted device should fail")
}

func deleteBad(t *testing.T, ctx *Context) {
	fmt.Println("deleteBad")
	require.NotEqual(t, ctx.mountPath, "", "Device is not mounted")

	err := ctx.Delete(ctx.volID)
	require.Error(t, err, "Delete on mounted device must fail")
}

func delete(t *testing.T, ctx *Context) {
	fmt.Println("delete")
	err := ctx.Delete(ctx.volID)
	require.NoError(t, err, "Delete failed")
	ctx.volID = ""
}

func snap(t *testing.T, ctx *Context) {
	fmt.Println("snap")
	if ctx.volID == "" {
		create(t, ctx)
	}
	attach(t, ctx)
	labels := map[string]string{"oh": "snap"}
	require.NotEqual(t, ctx.volID, "", "invalid volume ID")
	id, err := ctx.Snapshot(ctx.volID, false,
		&api.VolumeLocator{Name: "snappy", VolumeLabels: labels})
	require.NoError(t, err, "Failed in creating a snapshot")
	ctx.snapID = id
}

func snapInspect(t *testing.T, ctx *Context) {
	fmt.Println("snapInspect")

	snaps, err := ctx.Inspect([]string{ctx.snapID})
	require.NoError(t, err, "Failed in Inspect")
	require.NotNil(t, snaps, "Nil snaps")
	require.Equal(t, len(snaps), 1, "Expect 1 snaps actual %v snaps", len(snaps))
	require.Equal(t, snaps[0].Id, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].Id)

	snaps, err = ctx.Inspect([]string{string("shouldNotExist")})
	require.Equal(t, 0, len(snaps), "Expect 0 snaps actual %v snaps", len(snaps))
}

func snapEnumerate(t *testing.T, ctx *Context) {
	fmt.Println("snapEnumerate")

	snaps, err := ctx.SnapEnumerate(nil, nil)
	require.NoError(t, err, "Failed in snapEnumerate")
	require.NotNil(t, snaps, "Nil snaps")
	require.Equal(t, 1, len(snaps), "Expect 1 snaps actual %v snaps", len(snaps))
	require.Equal(t, snaps[0].Id, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].Id)
	labels := snaps[0].Locator.VolumeLabels

	snaps, err = ctx.SnapEnumerate([]string{ctx.volID}, nil)
	require.NoError(t, err, "Failed in snapEnumerate")
	require.NotNil(t, snaps, "Nil snaps")
	require.Equal(t, len(snaps), 1, "Expect 1 snap actual %v snaps", len(snaps))
	require.Equal(t, snaps[0].Id, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].Id)

	snaps, err = ctx.SnapEnumerate([]string{string("shouldNotExist")}, nil)
	require.Equal(t, len(snaps), 0, "Expect 0 snap actual %v snaps", len(snaps))

	snaps, err = ctx.SnapEnumerate(nil, labels)
	require.NoError(t, err, "Failed in snapEnumerate")
	require.NotNil(t, snaps, "Nil snaps")
	require.Equal(t, len(snaps), 1, "Expect 1 snap actual %v snaps", len(snaps))
	require.Equal(t, snaps[0].Id, ctx.snapID, "Expect snapID %v actual %v", ctx.snapID, snaps[0].Id)
}

func snapDiff(t *testing.T, ctx *Context) {
	fmt.Println("snapDiff")
}

func snapDelete(t *testing.T, ctx *Context) {
	fmt.Println("snapDelete")
}
<|endoftext|>


<commit_before>package module

import (
	"fmt"
	"github.com/davidscholberg/irkbot/lib/configure"
	"github.com/davidscholberg/irkbot/lib/message"
	"golang.org/x/net/html"
	"io"
	"mvdan.cc/xurls"
	"net"
	"net/http"
	urllib "net/url"
	"os"
	"strings"
	"time"
)

// parseUrls attempts to fetch the title of the HTML document returned by a URL
func parseUrls(cfg *configure.Config, in *message.InboundMsg, actions *actions) bool {
	urls := xurls.Strict().FindAllString(in.Msg, -1)

	for _, urlStr := range urls {
		url, err := urllib.Parse(urlStr)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			continue
		}
		if v, err := validateUrl(url); !v {
			fmt.Fprintln(os.Stderr, err)
			continue
		}
		host := url.Host
		title, err := getHtmlTitle(urlStr, cfg.Http.ResponseSizeLimit, cfg.Http.Timeout)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			continue
		}
		cleanedTitle := cleanTitleWhiteSpace(title)
		actions.say(fmt.Sprintf("^ %s - [%s]", cleanedTitle, host))
	}

	// don't consume the message, in case there are commands in it
	return false
}

// cleanTitleWhiteSpace takes the contents of an html title and makes it fit
// to be printed on a single line.
func cleanTitleWhiteSpace(title string) string {
	// split string by newline characters
	titleFields := strings.FieldsFunc(title, func(c rune) bool {
		return c == '\n' || c == '\r'
	})
	// trim each field
	trimmedFields := []string{}
	for _, field := range titleFields {
		field = strings.TrimSpace(field)
		if field != "" {
			trimmedFields = append(trimmedFields, field)
		}
	}
	// join by newline separator
	return strings.Join(trimmedFields, " / ")
}

// validateUrl ensures that the given URL is safe to GET.
func validateUrl(url *urllib.URL) (bool, error) {
	privateSubnets := []string{
		"127.0.0.0/8",
		"10.0.0.0/8",
		"172.16.0.0/12",
		"192.168.0.0/16",
		"::1/128",
		"fc00::/7"}

	ips, err := net.LookupIP(url.Host)
	if err != nil {
		return false, err
	}
	if len(ips) == 0 {
		return false, fmt.Errorf("no IPs found for %s", url.Host)
	}

	hostNotAllowed, err := isCidrMatch(&ips[0], privateSubnets)
	if err != nil {
		return false, err
	}
	if hostNotAllowed {
		return false, fmt.Errorf("host not allowed")
	}

	return true, nil
}

// isCidrMatch tests if a given IP is in any of a given set of CIDR subnets.
func isCidrMatch(ip *net.IP, subnets []string) (bool, error) {
	var cidr *net.IPNet
	var err error

	for _, subnet := range subnets {
		_, cidr, err = net.ParseCIDR(subnet)
		if err != nil {
			return false, err
		}
		if cidr.Contains(*ip) {
			return true, nil
		}
	}

	return false, nil
}

// getHtmlTitle returns the HTML title found at the given URL.
func getHtmlTitle(url string, responseSizeLimit int64, timeout int64) (string, error) {
	c := &http.Client{Timeout: time.Duration(timeout) * time.Second}
	response, err := c.Get(url)
	if err != nil {
		return "", err
	}
	defer response.Body.Close()

	// ignore response codes 400 and above
	if response.StatusCode >= 400 {
		return "", fmt.Errorf("received status %d", response.StatusCode)
	}

	doctree, err := html.Parse(io.LimitReader(response.Body, responseSizeLimit))
	if err != nil {
		return "", err
	}

	title, err := searchForHtmlTitle(doctree)
	if err != nil {
		return "", err
	}
	if len(title) == 0 {
		return "", fmt.Errorf("title not found")
	}

	return title, nil
}

// searchForHtmlTitle searches the parsed html document for the title.
func searchForHtmlTitle(n *html.Node) (string, error) {
	if n.Type == html.ElementNode && n.Data == "title" {
		if n.FirstChild.Type != html.TextNode {
			err := fmt.Errorf("child of title not TextNode type")
			return "", err
		}
		return n.FirstChild.Data, nil
	}
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		title, err := searchForHtmlTitle(c)
		if len(title) > 0 || err != nil {
			return title, err
		}
	}
	return "", nil
}
<commit_msg>url: updated to display final hostname if there were redirects<commit_after>package module

import (
	"fmt"
	"github.com/davidscholberg/irkbot/lib/configure"
	"github.com/davidscholberg/irkbot/lib/message"
	"golang.org/x/net/html"
	"io"
	"mvdan.cc/xurls"
	"net"
	"net/http"
	urllib "net/url"
	"os"
	"strings"
	"time"
)

// parseUrls attempts to fetch the title of the HTML document returned by a URL
func parseUrls(cfg *configure.Config, in *message.InboundMsg, actions *actions) bool {
	urls := xurls.Strict().FindAllString(in.Msg, -1)

	for _, urlStr := range urls {
		url, err := urllib.Parse(urlStr)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			continue
		}
		if v, err := validateUrl(url); !v {
			fmt.Fprintln(os.Stderr, err)
			continue
		}
		title, host, err := getHtmlTitle(urlStr, cfg.Http.ResponseSizeLimit, cfg.Http.Timeout)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			continue
		}
		cleanedTitle := cleanTitleWhiteSpace(title)
		actions.say(fmt.Sprintf("^ %s - [%s]", cleanedTitle, host))
	}

	// don't consume the message, in case there are commands in it
	return false
}

// cleanTitleWhiteSpace takes the contents of an html title and makes it fit
// to be printed on a single line.
func cleanTitleWhiteSpace(title string) string {
	// split string by newline characters
	titleFields := strings.FieldsFunc(title, func(c rune) bool {
		return c == '\n' || c == '\r'
	})
	// trim each field
	trimmedFields := []string{}
	for _, field := range titleFields {
		field = strings.TrimSpace(field)
		if field != "" {
			trimmedFields = append(trimmedFields, field)
		}
	}
	// join by newline separator
	return strings.Join(trimmedFields, " / ")
}

// validateUrl ensures that the given URL is safe to GET.
func validateUrl(url *urllib.URL) (bool, error) {
	privateSubnets := []string{
		"127.0.0.0/8",
		"10.0.0.0/8",
		"172.16.0.0/12",
		"192.168.0.0/16",
		"::1/128",
		"fc00::/7"}

	ips, err := net.LookupIP(url.Host)
	if err != nil {
		return false, err
	}
	if len(ips) == 0 {
		return false, fmt.Errorf("no IPs found for %s", url.Host)
	}

	hostNotAllowed, err := isCidrMatch(&ips[0], privateSubnets)
	if err != nil {
		return false, err
	}
	if hostNotAllowed {
		return false, fmt.Errorf("host not allowed")
	}

	return true, nil
}

// isCidrMatch tests if a given IP is in any of a given set of CIDR subnets.
func isCidrMatch(ip *net.IP, subnets []string) (bool, error) {
	var cidr *net.IPNet
	var err error

	for _, subnet := range subnets {
		_, cidr, err = net.ParseCIDR(subnet)
		if err != nil {
			return false, err
		}
		if cidr.Contains(*ip) {
			return true, nil
		}
	}

	return false, nil
}

// getHtmlTitle returns the HTML title found at the given URL as well as the
// final hostname (after any redirects).
func getHtmlTitle(url string, responseSizeLimit int64, timeout int64) (string, string, error) {
	c := &http.Client{Timeout: time.Duration(timeout) * time.Second}
	response, err := c.Get(url)
	if err != nil {
		return "", "", err
	}
	defer response.Body.Close()

	// ignore response codes 400 and above
	if response.StatusCode >= 400 {
		return "", "", fmt.Errorf("received status %d", response.StatusCode)
	}

	doctree, err := html.Parse(io.LimitReader(response.Body, responseSizeLimit))
	if err != nil {
		return "", "", err
	}

	title, err := searchForHtmlTitle(doctree)
	if err != nil {
		return "", "", err
	}
	if len(title) == 0 {
		return "", "", fmt.Errorf("title not found")
	}

	return title, response.Request.URL.Host, nil
}

// searchForHtmlTitle searches the parsed html document for the title.
func searchForHtmlTitle(n *html.Node) (string, error) {
	if n.Type == html.ElementNode && n.Data == "title" {
		if n.FirstChild.Type != html.TextNode {
			err := fmt.Errorf("child of title not TextNode type")
			return "", err
		}
		return n.FirstChild.Data, nil
	}
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		title, err := searchForHtmlTitle(c)
		if len(title) > 0 || err != nil {
			return title, err
		}
	}
	return "", nil
}
<|endoftext|>


<commit_before>package workers

import (
	"fmt"
	"github.com/APTrust/exchange/context"
	"github.com/APTrust/exchange/dpn/models"
	"github.com/APTrust/exchange/dpn/network"
	apt_models "github.com/APTrust/exchange/models"
	"github.com/nsqio/go-nsq"
//	"os"
	"os/exec"
//	"path/filepath"
//	"time"
)

// dpn_copier copies tarred bags from other nodes via rsync.
// This is used when replicating content from other nodes.
// For putting together DPN bags from APTrust files, see fetcher.go.

type Copier struct {
	CopyChannel         chan *CopyManifest
	ChecksumChannel     chan *CopyManifest
	Context             *context.Context
	LocalClient         *network.DPNRestClient
	RemoteClients       map[string]*network.DPNRestClient
}

type CopyManifest struct {
	NsqMessage          *nsq.Message `json:"-"`
	DPNWorkItem         *apt_models.DPNWorkItem
	ReplicationTransfer *models.ReplicationTransfer
	DPNBag              *models.DPNBag
	WorkSummary         *apt_models.WorkSummary
	LocalPath           string
	RsyncStdout         string
	RsyncStderr         string
}

func NewCopier(_context *context.Context) (*Copier, error) {
	localClient, err := network.NewDPNRestClient(
		_context.Config.DPN.RestClient.LocalServiceURL,
		_context.Config.DPN.RestClient.LocalAPIRoot,
		_context.Config.DPN.RestClient.LocalAuthToken,
		_context.Config.DPN.LocalNode,
		_context.Config.DPN)
	if err != nil {
		return nil, fmt.Errorf("Error creating local DPN REST client: %v", err)
	}
	remoteClients, err := localClient.GetRemoteClients()
	if err != nil {
		return nil, err
	}
	copier := &Copier {
		Context: _context,
		LocalClient: localClient,
		RemoteClients: remoteClients,
	}
	workerBufferSize := _context.Config.DPN.DPNCopyWorker.Workers * 4
	copier.CopyChannel = make(chan *CopyManifest, workerBufferSize)
	copier.ChecksumChannel = make(chan *CopyManifest, workerBufferSize)
	for i := 0; i < _context.Config.DPN.DPNCopyWorker.Workers; i++ {
		go copier.doCopy()
		go copier.verifyChecksum()
	}
	return copier, nil
}

func (copier *Copier) HandleMessage(message *nsq.Message) error {
	message.DisableAutoResponse()

	// Get the DPNWorkItem, the ReplicationTransfer, and the DPNBag
	copyManifest := copier.buildCopyManifest(message)

	// Start processing.
	copier.CopyChannel <- copyManifest
	copier.Context.MessageLog.Info("Put xfer request %s (bag %s) from %s " +
		" into the copy channel", copyManifest.ReplicationTransfer.ReplicationId,
		copyManifest.ReplicationTransfer, copyManifest.ReplicationTransfer.FromNode)
	return nil
}

// Copy the file from the remote node to our local staging area.
func (copier *Copier) doCopy() {
	for copyManifest := range copier.CopyChannel {
		localPath := "?"
		rsyncCommand := GetRsyncCommand(copyManifest.ReplicationTransfer.Link,
			localPath, copier.Context.Config.DPN.UseSSHWithRsync)

		// Touch message on both sides of rsync, so NSQ doesn't time out.
		if copyManifest.NsqMessage != nil {
			copyManifest.NsqMessage.Touch()
		}
		output, err := rsyncCommand.CombinedOutput()
		copier.Context.MessageLog.Info("Rsync Output: %s", output)
		if copyManifest.NsqMessage != nil {
			copyManifest.NsqMessage.Touch()
		}
		if err != nil {
			// Something went wrong
		} else {
			// OK
		}
	}
}

// Run a checksum on the tag manifest and send that back to the
// FromNode. If the checksum is good, the FromNode will set
// the ReplicationTransfer's StoreRequested attribute to true,
// and we should store the bag. If the checksum is bad, the remote
// node will set StoreRequested to false, and we should delete
// the tar file.
func (copier *Copier) verifyChecksum() {
	//for copyManifest := range copier.ChecksumChannel {
		// 1. Calculate the sha256 digest of the tag manifest.
		// 2. Send the result the ReplicationTransfer.FromNode.
		// 3. If the updated ReplicationTransfer.StoreRequested is true,
		//    push this item into the validation queue. Otherwise,
		//    delete the bag from the local staging area.
	//}
}

func (copier *Copier) buildCopyManifest(message *nsq.Message) (*CopyManifest) {
	// 1. Get the DPNWorkItem from Pharos.
	//    Stop if it's marked complete.
	// 2. Get the ReplicationTransfer from the remote node.
	//    Stop if it's completed or cancelled.
	// 3. Get the DPNBag record from the remote node.
	//    We need to know its size.
	// 4. Build and return the CopyManifest.
	return nil
}

// Make sure we have space to copy this item from the remote node.
// We will be validating this bag in a later step without untarring it,
// so we just have to reserve enough room for the tar file.
func (copier *Copier) reserveSpaceOnVolume(copyManifest *CopyManifest) (bool) {
	okToCopy := false
	err := copier.Context.VolumeClient.Ping(500)
	if err == nil {
		path := copyManifest.LocalPath
		ok, err := copier.Context.VolumeClient.Reserve(path, uint64(copyManifest.DPNBag.Size))
		if err != nil {
			copier.Context.MessageLog.Warning("Volume service returned an error. " +
				"Will requeue ReplicationTransfer %s bag (%s) because we may not " +
				"have enough space to copy %d bytes from %s.",
				copyManifest.ReplicationTransfer.ReplicationId,
				copyManifest.ReplicationTransfer.Bag,
				copyManifest.DPNBag.Size,
				copyManifest.ReplicationTransfer.FromNode)
		} else if ok {
			// VolumeService says we have enough space for this.
			okToCopy = ok
		}
	} else {
		copier.Context.MessageLog.Warning("Volume service is not running or returned an error. " +
			"Continuing as if we have enough space to download %d bytes.",
			copyManifest.DPNBag.Size,)
		okToCopy = true
	}
	return okToCopy
}

// GetRsyncCommand returns a command object for copying from the remote
// location to the local filesystem. The copy is done via rsync over ssh,
// and the command will capture stdout and stderr. The copyFrom param
// should be a valid scp target in this format:
//
// remoteuser@remotehost:/remote/dir/bag.tar
//
// The copyTo param should be an absolute path on a locally-accessible
// file system, such as:
//
// /mnt/dpn/data/bag.tar
//
// Using this assumes a few things:
//
// 1. You have rsync installed.
// 2. You have an ssh client installed.
// 3. You have an entry in your ~/.ssh/config file specifying
//    connection and key information for the remote host.
//
// Usage:
//
// command := GetRsyncCommand("aptrust@tdr:bag.tar", "/mnt/dpn/bag.tar")
// err := command.Run()
// if err != nil {
//    ... do something ...
// }
//
// -- OR --
//
// output, err := command.CombinedOutput()
// if err != nil {
//    fmt.Println(err.Error())
//    fmt.Println(string(output))
// }
func GetRsyncCommand(copyFrom, copyTo string, useSSH bool) (*exec.Cmd) {
	//rsync -avz -e ssh remoteuser@remotehost:/remote/dir /this/dir/
	if useSSH {
		return exec.Command("rsync", "-avzW", "-e",  "ssh", copyFrom, copyTo, "--inplace")
	}
	return exec.Command("rsync", "-avzW", "--inplace", copyFrom, copyTo)
}
<commit_msg>More work on DPN copier<commit_after>package workers

import (
	"fmt"
	"github.com/APTrust/exchange/constants"
	"github.com/APTrust/exchange/context"
	"github.com/APTrust/exchange/dpn/models"
	"github.com/APTrust/exchange/dpn/network"
	"github.com/APTrust/exchange/util"
	apt_models "github.com/APTrust/exchange/models"
	"github.com/nsqio/go-nsq"
//	"os"
	"os/exec"
//	"path/filepath"
	"strconv"
//	"time"
)

// dpn_copier copies tarred bags from other nodes via rsync.
// This is used when replicating content from other nodes.
// For putting together DPN bags from APTrust files, see fetcher.go.

type Copier struct {
	CopyChannel         chan *CopyManifest
	ChecksumChannel     chan *CopyManifest
	Context             *context.Context
	LocalClient         *network.DPNRestClient
	RemoteClients       map[string]*network.DPNRestClient
}

type CopyManifest struct {
	NsqMessage          *nsq.Message `json:"-"`
	DPNWorkItem         *apt_models.DPNWorkItem
	ReplicationTransfer *models.ReplicationTransfer
	DPNBag              *models.DPNBag
	WorkSummary         *apt_models.WorkSummary
	LocalPath           string
	RsyncStdout         string
	RsyncStderr         string
}

func NewCopyManifest() (*CopyManifest) {
	return &CopyManifest{
		WorkSummary: apt_models.NewWorkSummary(),
	}
}

func NewCopier(_context *context.Context) (*Copier, error) {
	localClient, err := network.NewDPNRestClient(
		_context.Config.DPN.RestClient.LocalServiceURL,
		_context.Config.DPN.RestClient.LocalAPIRoot,
		_context.Config.DPN.RestClient.LocalAuthToken,
		_context.Config.DPN.LocalNode,
		_context.Config.DPN)
	if err != nil {
		return nil, fmt.Errorf("Error creating local DPN REST client: %v", err)
	}
	remoteClients, err := localClient.GetRemoteClients()
	if err != nil {
		return nil, err
	}
	copier := &Copier {
		Context: _context,
		LocalClient: localClient,
		RemoteClients: remoteClients,
	}
	workerBufferSize := _context.Config.DPN.DPNCopyWorker.Workers * 4
	copier.CopyChannel = make(chan *CopyManifest, workerBufferSize)
	copier.ChecksumChannel = make(chan *CopyManifest, workerBufferSize)
	for i := 0; i < _context.Config.DPN.DPNCopyWorker.Workers; i++ {
		go copier.doCopy()
		go copier.verifyChecksum()
	}
	return copier, nil
}

func (copier *Copier) HandleMessage(message *nsq.Message) error {
	message.DisableAutoResponse()

	// Get the DPNWorkItem, the ReplicationTransfer, and the DPNBag
	copyManifest := copier.buildCopyManifest(message)

	// Start processing.
	copier.CopyChannel <- copyManifest
	copier.Context.MessageLog.Info("Put xfer request %s (bag %s) from %s " +
		" into the copy channel", copyManifest.ReplicationTransfer.ReplicationId,
		copyManifest.ReplicationTransfer, copyManifest.ReplicationTransfer.FromNode)
	return nil
}

// Copy the file from the remote node to our local staging area.
func (copier *Copier) doCopy() {
	for copyManifest := range copier.CopyChannel {
		localPath := "?"
		rsyncCommand := GetRsyncCommand(copyManifest.ReplicationTransfer.Link,
			localPath, copier.Context.Config.DPN.UseSSHWithRsync)

		// Touch message on both sides of rsync, so NSQ doesn't time out.
		if copyManifest.NsqMessage != nil {
			copyManifest.NsqMessage.Touch()
		}
		output, err := rsyncCommand.CombinedOutput()
		copier.Context.MessageLog.Info("Rsync Output: %s", output)
		if copyManifest.NsqMessage != nil {
			copyManifest.NsqMessage.Touch()
		}
		if err != nil {
			// Something went wrong
		} else {
			// OK
		}
	}
}

// Run a checksum on the tag manifest and send that back to the
// FromNode. If the checksum is good, the FromNode will set
// the ReplicationTransfer's StoreRequested attribute to true,
// and we should store the bag. If the checksum is bad, the remote
// node will set StoreRequested to false, and we should delete
// the tar file.
func (copier *Copier) verifyChecksum() {
	//for copyManifest := range copier.ChecksumChannel {
		// 1. Calculate the sha256 digest of the tag manifest.
		// 2. Send the result the ReplicationTransfer.FromNode.
		// 3. If the updated ReplicationTransfer.StoreRequested is true,
		//    push this item into the validation queue. Otherwise,
		//    delete the bag from the local staging area.
	//}
}

// buildCopyManifest creates a CopyManifest for this job.
func (copier *Copier) buildCopyManifest(message *nsq.Message) (*CopyManifest) {
	// 1. Get the DPNWorkItem from Pharos.
	//    Stop if it's marked complete.
	// 2. Get the ReplicationTransfer from the remote node.
	//    Stop if it's completed or cancelled.
	// 3. Get the DPNBag record from the remote node.
	//    We need to know its size.
	// 4. Build and return the CopyManifest.
	copyManifest := NewCopyManifest()
	copyManifest.NsqMessage = message
	copier.getDPNWorkItem(copyManifest)
	if copyManifest.WorkSummary.HasErrors() {
		return copyManifest
	}
	copier.getXferRequest(copyManifest)
	if copyManifest.WorkSummary.HasErrors() {
		return copyManifest
	}
	copier.getDPNBag(copyManifest)
	return copyManifest
}

// getDPNWorkItem returns the DPNWorkItem associated with this message,
// and a boolean indicating whether or not processing should continue.
func (copier *Copier) getDPNWorkItem(copyManifest *CopyManifest) {
	workItemId, err := strconv.Atoi(string(copyManifest.NsqMessage.Body))
	if err != nil {
		msg := fmt.Sprintf("Could not get DPNWorkItemId from" +
			"NSQ message body '%s': %v", copyManifest.NsqMessage.Body, err)
		copyManifest.WorkSummary.AddError(msg)
		copyManifest.WorkSummary.ErrorIsFatal = true
		return
	}
	resp := copier.Context.PharosClient.DPNWorkItemGet(workItemId)
	if resp.Error != nil {
		msg := fmt.Sprintf("Could not get DPNWorkItem (id %d) " +
			"from Pharos: %v", workItemId, resp.Error)
		copyManifest.WorkSummary.AddError(msg)
		copyManifest.WorkSummary.ErrorIsFatal = true
		return
	}
	dpnWorkItem := resp.DPNWorkItem()
	copyManifest.DPNWorkItem = dpnWorkItem
	if dpnWorkItem == nil {
		msg := fmt.Sprintf("Pharos returned nil for DPNWorkItem %d",
			workItemId)
		copyManifest.WorkSummary.AddError(msg)
		copyManifest.WorkSummary.ErrorIsFatal = true
		return
	}
	if dpnWorkItem.Task != constants.DPNTaskReplication {
		msg := fmt.Sprintf("DPNWorkItem %d has task type %s, " +
			"and does not belong in this queue!", workItemId, dpnWorkItem.Task)
		copyManifest.WorkSummary.AddError(msg)
		copyManifest.WorkSummary.ErrorIsFatal = true
	}
	if !util.LooksLikeUUID(dpnWorkItem.Identifier) {
		msg := fmt.Sprintf("DPNWorkItem %d has identifier '%s', " +
			"which does not look like a UUID", workItemId, dpnWorkItem.Identifier)
		copyManifest.WorkSummary.AddError(msg)
		copyManifest.WorkSummary.ErrorIsFatal = true
	}
}

func (copier *Copier) getXferRequest(copyManifest *CopyManifest) {
	if copyManifest == nil || copyManifest.DPNWorkItem == nil {
		msg := fmt.Sprintf("getXferRequest: CopyManifest.DPNWorkItem cannot be nil.")
		copyManifest.WorkSummary.AddError(msg)
		copyManifest.WorkSummary.ErrorIsFatal = true
		return
	}
	resp := copier.LocalClient.ReplicationTransferGet(copyManifest.DPNWorkItem.Identifier)
	if resp.Error != nil {
		msg := fmt.Sprintf("Could not get ReplicationTransfer %s " +
			"from DPN server: %v", copyManifest.DPNWorkItem.Identifier, resp.Error)
		copyManifest.WorkSummary.AddError(msg)
		copyManifest.WorkSummary.ErrorIsFatal = true
		return
	}
	xfer := resp.ReplicationTransfer()
	copyManifest.ReplicationTransfer = xfer
	if xfer == nil {
		msg := fmt.Sprintf("DPN server returned nil for ReplicationId %s",
			copyManifest.DPNWorkItem.Identifier)
		copyManifest.WorkSummary.AddError(msg)
		copyManifest.WorkSummary.ErrorIsFatal = true
		return
	}
	if xfer.Stored {
		msg := fmt.Sprintf("ReplicationId %s is already marked as Stored. Nothing left to do.",
			copyManifest.DPNWorkItem.Identifier)
		copyManifest.WorkSummary.AddError(msg)
		copyManifest.WorkSummary.ErrorIsFatal = true
		return
	}
	if xfer.Cancelled {
		msg := fmt.Sprintf("ReplicationId %s was cancelled. Nothing left to do.",
			copyManifest.DPNWorkItem.Identifier)
		copyManifest.WorkSummary.AddError(msg)
		copyManifest.WorkSummary.ErrorIsFatal = true
	}
}

func (copier *Copier) getDPNBag(copyManifest *CopyManifest) {
	if copyManifest == nil || copyManifest.ReplicationTransfer == nil {
		msg := fmt.Sprintf("getDPNBag: CopyManifest.ReplicationTransfer cannot be nil.")
		copyManifest.WorkSummary.ErrorIsFatal = true
		copyManifest.WorkSummary.AddError(msg)
		return
	}
	resp := copier.LocalClient.DPNBagGet(copyManifest.ReplicationTransfer.Bag)
	if resp.Error != nil {
		msg := fmt.Sprintf("Could not get ReplicationTransfer %s " +
			"from DPN server: %v", copyManifest.DPNWorkItem.Identifier, resp.Error)
		copyManifest.WorkSummary.AddError(msg)
		copyManifest.WorkSummary.ErrorIsFatal = true
		return
	}
	dpnBag := resp.Bag()
	copyManifest.DPNBag = dpnBag
	if dpnBag == nil {
		msg := fmt.Sprintf("DPN server returned nil for Bag %s",
			copyManifest.ReplicationTransfer.Bag)
		copyManifest.WorkSummary.AddError(msg)
		copyManifest.WorkSummary.ErrorIsFatal = true
		return
	}
}

// Make sure we have space to copy this item from the remote node.
// We will be validating this bag in a later step without untarring it,
// so we just have to reserve enough room for the tar file.
func (copier *Copier) reserveSpaceOnVolume(copyManifest *CopyManifest) (bool) {
	okToCopy := false
	err := copier.Context.VolumeClient.Ping(500)
	if err == nil {
		path := copyManifest.LocalPath
		ok, err := copier.Context.VolumeClient.Reserve(path, uint64(copyManifest.DPNBag.Size))
		if err != nil {
			copier.Context.MessageLog.Warning("Volume service returned an error. " +
				"Will requeue ReplicationTransfer %s bag (%s) because we may not " +
				"have enough space to copy %d bytes from %s.",
				copyManifest.ReplicationTransfer.ReplicationId,
				copyManifest.ReplicationTransfer.Bag,
				copyManifest.DPNBag.Size,
				copyManifest.ReplicationTransfer.FromNode)
		} else if ok {
			// VolumeService says we have enough space for this.
			okToCopy = ok
		}
	} else {
		copier.Context.MessageLog.Warning("Volume service is not running or returned an error. " +
			"Continuing as if we have enough space to download %d bytes.",
			copyManifest.DPNBag.Size,)
		okToCopy = true
	}
	return okToCopy
}

// GetRsyncCommand returns a command object for copying from the remote
// location to the local filesystem. The copy is done via rsync over ssh,
// and the command will capture stdout and stderr. The copyFrom param
// should be a valid scp target in this format:
//
// remoteuser@remotehost:/remote/dir/bag.tar
//
// The copyTo param should be an absolute path on a locally-accessible
// file system, such as:
//
// /mnt/dpn/data/bag.tar
//
// Using this assumes a few things:
//
// 1. You have rsync installed.
// 2. You have an ssh client installed.
// 3. You have an entry in your ~/.ssh/config file specifying
//    connection and key information for the remote host.
//
// Usage:
//
// command := GetRsyncCommand("aptrust@tdr:bag.tar", "/mnt/dpn/bag.tar")
// err := command.Run()
// if err != nil {
//    ... do something ...
// }
//
// -- OR --
//
// output, err := command.CombinedOutput()
// if err != nil {
//    fmt.Println(err.Error())
//    fmt.Println(string(output))
// }
func GetRsyncCommand(copyFrom, copyTo string, useSSH bool) (*exec.Cmd) {
	//rsync -avz -e ssh remoteuser@remotehost:/remote/dir /this/dir/
	if useSSH {
		return exec.Command("rsync", "-avzW", "-e",  "ssh", copyFrom, copyTo, "--inplace")
	}
	return exec.Command("rsync", "-avzW", "--inplace", copyFrom, copyTo)
}
<|endoftext|>


<commit_before><commit_msg>Fix build errors & move deleteCommit work into txn<commit_after><|endoftext|>


<commit_before>// toolbag - dynamic tools for my website
// Copyright (C) 2015 Austin Adams
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package tools

import (
	"errors"
	"html/template"
	"net/http"
	"strings"

	"code.austinjadams.com/execd"
	tb "code.austinjadams.com/toolbag"
)

type Figlet struct {
	defaultFont string
	fonts       map[string][]string
	templ       *template.Template
	net, addr   string

	args struct {
		template, unix, tcp string
	}
}

func NewFiglet() *Figlet {
	return &Figlet{}
}

func (f *Figlet) Name() string { return "figlet" }
func (f *Figlet) Desc() string { return "a web frontend to figlet" }
func (f *Figlet) Path() string { return "/" + f.Name() }
func (f *Figlet) AddArgs(toolbag *tb.ToolBag) {
	toolbag.StringVar(&f.args.template, tb.Arg(f, "template"), "", "path to template")
	toolbag.StringVar(&f.args.unix, tb.Arg(f, "unix"), "", "path to unix socket to execd")
	toolbag.StringVar(&f.args.tcp, tb.Arg(f, "tcp"), "", "tcp address to execd")
}

func (f *Figlet) makeClient() (*execd.Client, error) {
	return execd.DialClient(f.net, f.addr)
}

func (f *Figlet) fontCategory(needle string) string {
	for category, fonts := range f.fonts {
		for _, font := range fonts {
			if needle == font {
				return category
			}
		}
	}
	// no match
	return ""
}

func (f *Figlet) Init() error {
	if f.args.template == "" {
		return errors.New("missing template arg")
	}
	if (f.args.unix == "") == (f.args.tcp == "") {
		return errors.New("specify either unix or tcp, but not both")
	}

	templ, err := template.ParseFiles(f.args.template)
	if err != nil {
		return err
	}
	f.templ = templ

	if f.args.unix != "" {
		f.net = "unix"
		f.addr = f.args.unix
	} else {
		f.net = "tcp"
		f.addr = f.args.tcp
	}

	client, err := f.makeClient()
	if err != nil {
		return err
	}
	// find default font
	defaultFont, err := client.ExecString("", "fig", "default")
	if err != nil {
		return err
	}
	f.defaultFont = strings.TrimSpace(defaultFont)

	// find categories of fonts
	output, err := client.ExecString("", "fig", "ls")
	if err != nil {
		return err
	}
	f.fonts = make(map[string][]string)
	for _, v := range splitLines(output) {
		f.fonts[v] = nil
	}

	// find the fonts in each category
	for category, _ := range f.fonts {
		output, err = client.ExecString("", "fig", "ls", category)
		if err != nil {
			return err
		}
		f.fonts[category] = splitLines(output)
	}

	return nil
}

// serve
func (f *Figlet) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	font := r.PostFormValue("font")
	text := r.PostFormValue("text")
	result := ""

	if font == "" {
		font = f.defaultFont
	}

	if text != "" {
		category := f.fontCategory(font)
		if category == "" {
			tb.WhineString(f, w, "nice try, lad")
			return
		}

		client, err := f.makeClient()
		if err != nil {
			tb.Whine(f, w, err)
			return
		}
		// otherwise
		defer client.Close()

		result, err = client.ExecString(text, "fig", category, font)
		if err != nil {
			tb.Whine(f, w, err)
			return
		}
	}

	err := f.templ.Execute(w, &struct {
		Font         string
		Fonts        map[string][]string
		Text, Result string
	}{font, f.fonts, text, result})
	if err != nil {
		tb.LogWhine(f, err)
		return
	}
}
<commit_msg>figlet: limit request body to 4K, clean up Init()<commit_after>// toolbag - dynamic tools for my website
// Copyright (C) 2015 Austin Adams
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package tools

import (
	"errors"
	"html/template"
	"net/http"
	"strings"

	"code.austinjadams.com/execd"
	tb "code.austinjadams.com/toolbag"
)

// by default, limit request bodies to 4KiB (2^12 bytes)
const defaultMaxReqBody int64 = 1 << 12

type Figlet struct {
	defaultFont string
	fonts       map[string][]string
	templ       *template.Template
	net, addr   string
	args        struct {
		maxReqBody          int64
		template, unix, tcp string
	}
}

func NewFiglet() *Figlet {
	return &Figlet{}
}

func (f *Figlet) Name() string { return "figlet" }
func (f *Figlet) Desc() string { return "a web frontend to figlet" }
func (f *Figlet) Path() string { return "/" + f.Name() }
func (f *Figlet) AddArgs(toolbag *tb.ToolBag) {
	toolbag.StringVar(&f.args.template, tb.Arg(f, "template"), "", "path to template")
	toolbag.StringVar(&f.args.unix, tb.Arg(f, "unix"), "", "path to unix socket to execd")
	toolbag.StringVar(&f.args.tcp, tb.Arg(f, "tcp"), "", "tcp address to execd")
	toolbag.Int64Var(&f.args.maxReqBody, tb.Arg(f, "maxReqBody"), defaultMaxReqBody, "maximum size of a request body in bytes")
}

func (f *Figlet) makeClient() (*execd.Client, error) {
	return execd.DialClient(f.net, f.addr)
}

func (f *Figlet) fontCategory(needle string) string {
	for category, fonts := range f.fonts {
		for _, font := range fonts {
			if needle == font {
				return category
			}
		}
	}
	// no match
	return ""
}

func (f *Figlet) parseArgs() error {
	if f.args.template == "" {
		return errors.New("missing template arg")
	}
	if (f.args.unix == "") == (f.args.tcp == "") {
		return errors.New("specify either unix or tcp, but not both")
	}

	templ, err := template.ParseFiles(f.args.template)
	if err != nil {
		return err
	}
	f.templ = templ

	if f.args.unix != "" {
		f.net = "unix"
		f.addr = f.args.unix
	} else {
		f.net = "tcp"
		f.addr = f.args.tcp
	}

	return nil
}

func (f *Figlet) findDefaultFont(client *execd.Client) error {
	// find default font
	defaultFont, err := client.ExecString("", "fig", "default")
	if err != nil {
		return err
	}
	f.defaultFont = strings.TrimSpace(defaultFont)

	return nil
}

func (f *Figlet) findFonts(client *execd.Client) error {
	// find categories of fonts
	output, err := client.ExecString("", "fig", "ls")
	if err != nil {
		return err
	}
	f.fonts = make(map[string][]string)
	for _, v := range splitLines(output) {
		f.fonts[v] = nil
	}

	// find the fonts in each category
	for category, _ := range f.fonts {
		output, err = client.ExecString("", "fig", "ls", category)
		if err != nil {
			return err
		}
		f.fonts[category] = splitLines(output)
	}

	return nil
}

func (f *Figlet) Init() error {
	err := f.parseArgs()

	client, err := f.makeClient()
	if err != nil {
		return err
	}

	err = f.findDefaultFont(client)
	if err != nil {
		return err
	}

	err = f.findFonts(client)
	if err != nil {
		return err
	}

	return nil
}

// serve
func (f *Figlet) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, f.args.maxReqBody)
	err := r.ParseForm()
	if err != nil {
		tb.Whine(f, w, err)
		return
	}

	font := r.PostFormValue("font")
	text := r.PostFormValue("text")
	result := ""

	if font == "" {
		font = f.defaultFont
	}

	if text != "" {
		category := f.fontCategory(font)
		if category == "" {
			tb.WhineString(f, w, "nice try, lad")
			return
		}

		client, err := f.makeClient()
		if err != nil {
			tb.Whine(f, w, err)
			return
		}
		// otherwise
		defer client.Close()

		result, err = client.ExecString(text, "fig", category, font)
		if err != nil {
			tb.Whine(f, w, err)
			return
		}
	}

	err = f.templ.Execute(w, &struct {
		Font         string
		Fonts        map[string][]string
		Text, Result string
	}{font, f.fonts, text, result})
	if err != nil {
		tb.LogWhine(f, err)
		return
	}
}
<|endoftext|>


<commit_before>// mksyscall_windows.pl -l32 symlink/symlink_windows.go
// MACHINE GENERATED BY THE COMMAND ABOVE; DO NOT EDIT

package symlink

import "unsafe"
import "syscall"

var (
	modkernel32 = syscall.NewLazyDLL("kernel32.dll")

	procCreateSymbolicLinkW       = modkernel32.NewProc("CreateSymbolicLinkW")
	procGetFinalPathNameByHandleW = modkernel32.NewProc("GetFinalPathNameByHandleW")
)

func createSymbolicLink(symlinkname *uint16, targetname *uint16, flags uint32) (err error) {
	r1, _, e1 := syscall.Syscall(procCreateSymbolicLinkW.Addr(), 3, uintptr(unsafe.Pointer(symlinkname)), uintptr(unsafe.Pointer(targetname)), uintptr(flags))
	if r1 == 0 {
		if e1 != 0 {
			err = error(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

func getFinalPathNameByHandle(handle Handle, buf *uint16, buflen uint32, flags uint32) (n uint32, err error) {
	r0, _, e1 := syscall.Syscall6(procGetFinalPathNameByHandleW.Addr(), 4, uintptr(handle), uintptr(unsafe.Pointer(buf)), uintptr(buflen), uintptr(flags), 0, 0)
	n = uint32(r0)
	if n == 0 {
		if e1 != 0 {
			err = error(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}
<commit_msg>utils/syscalls: fix compilation failure on windows/386<commit_after>// mksyscall_windows.pl -l32 symlink/symlink_windows.go
// MACHINE GENERATED BY THE COMMAND ABOVE; DO NOT EDIT

package symlink

import "unsafe"
import "syscall"

var (
	modkernel32 = syscall.NewLazyDLL("kernel32.dll")

	procCreateSymbolicLinkW       = modkernel32.NewProc("CreateSymbolicLinkW")
	procGetFinalPathNameByHandleW = modkernel32.NewProc("GetFinalPathNameByHandleW")
)

func createSymbolicLink(symlinkname *uint16, targetname *uint16, flags uint32) (err error) {
	r1, _, e1 := syscall.Syscall(procCreateSymbolicLinkW.Addr(), 3, uintptr(unsafe.Pointer(symlinkname)), uintptr(unsafe.Pointer(targetname)), uintptr(flags))
	if r1 == 0 {
		if e1 != 0 {
			err = error(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

func getFinalPathNameByHandle(handle syscall.Handle, buf *uint16, buflen uint32, flags uint32) (n uint32, err error) {
	r0, _, e1 := syscall.Syscall6(procGetFinalPathNameByHandleW.Addr(), 4, uintptr(handle), uintptr(unsafe.Pointer(buf)), uintptr(buflen), uintptr(flags), 0, 0)
	n = uint32(r0)
	if n == 0 {
		if e1 != 0 {
			err = error(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}
<|endoftext|>


<commit_before>package js

import (
	"testing"

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

func TestHashTable(t *testing.T) {
	assert.Equal(t, ToHash([]byte("break")), Font, "'break' must resolve to hash.Break")
	assert.Equal(t, "break", Break.String(), "hash.Break must resolve to 'break'")
}
<commit_msg>Typo fix<commit_after>package js

import (
	"testing"

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

func TestHashTable(t *testing.T) {
	assert.Equal(t, ToHash([]byte("break")), Break, "'break' must resolve to hash.Break")
	assert.Equal(t, "break", Break.String(), "hash.Break must resolve to 'break'")
}
<|endoftext|>


<commit_before>package torrent

import (
	"testing"

	"github.com/anacrolix/torrent/peer_protocol"
)

func r(i, b, l peer_protocol.Integer) request {
	return request{i, chunkSpec{b, l}}
}

// Check the given Request is correct for various torrent offsets.
func TestTorrentRequest(t *testing.T) {
	const s = 472183431 // Length of torrent.
	for _, _case := range []struct {
		off int64   // An offset into the torrent.
		req request // The expected Request. The zero value means !ok.
	}{
		// Invalid offset.
		{-1, request{}},
		{0, r(0, 0, 16384)},
		// One before the end of a piece.
		{1<<18 - 1, r(0, 1<<18-16384, 16384)},
		// Offset beyond torrent length.
		{472 * 1 << 20, request{}},
		// One before the end of the torrent. Complicates the chunk length.
		{s - 1, r((s-1)/(1<<18), (s-1)%(1<<18)/(16384)*(16384), 12935)},
		{1, r(0, 0, 16384)},
		// One before end of chunk.
		{16383, r(0, 0, 16384)},
		// Second chunk.
		{16384, r(0, 16384, 16384)},
	} {
		req, ok := torrentOffsetRequest(472183431, 1<<18, 16384, _case.off)
		if (_case.req == request{}) == ok {
			t.Fatalf("expected %v, got %v", _case.req, req)
		}
		if req != _case.req {
			t.Fatalf("expected %v, got %v", _case.req, req)
		}
	}
}

func TestAppendToCopySlice(t *testing.T) {
	orig := []int{1, 2, 3}
	dupe := append([]int{}, orig...)
	dupe[0] = 4
	if orig[0] != 1 {
		t.FailNow()
	}
}

func TestTorrentString(t *testing.T) {
	tor := &Torrent{}
	s := tor.InfoHash().HexString()
	if s != "0000000000000000000000000000000000000000" {
		t.FailNow()
	}
}
<commit_msg>Add a benchmark for an observed slow case with Torrent.updatePiecePriorities<commit_after>package torrent

import (
	"testing"

	"github.com/bradfitz/iter"
	"github.com/stretchr/testify/assert"

	"github.com/anacrolix/torrent/metainfo"
	"github.com/anacrolix/torrent/peer_protocol"
)

func r(i, b, l peer_protocol.Integer) request {
	return request{i, chunkSpec{b, l}}
}

// Check the given Request is correct for various torrent offsets.
func TestTorrentRequest(t *testing.T) {
	const s = 472183431 // Length of torrent.
	for _, _case := range []struct {
		off int64   // An offset into the torrent.
		req request // The expected Request. The zero value means !ok.
	}{
		// Invalid offset.
		{-1, request{}},
		{0, r(0, 0, 16384)},
		// One before the end of a piece.
		{1<<18 - 1, r(0, 1<<18-16384, 16384)},
		// Offset beyond torrent length.
		{472 * 1 << 20, request{}},
		// One before the end of the torrent. Complicates the chunk length.
		{s - 1, r((s-1)/(1<<18), (s-1)%(1<<18)/(16384)*(16384), 12935)},
		{1, r(0, 0, 16384)},
		// One before end of chunk.
		{16383, r(0, 0, 16384)},
		// Second chunk.
		{16384, r(0, 16384, 16384)},
	} {
		req, ok := torrentOffsetRequest(472183431, 1<<18, 16384, _case.off)
		if (_case.req == request{}) == ok {
			t.Fatalf("expected %v, got %v", _case.req, req)
		}
		if req != _case.req {
			t.Fatalf("expected %v, got %v", _case.req, req)
		}
	}
}

func TestAppendToCopySlice(t *testing.T) {
	orig := []int{1, 2, 3}
	dupe := append([]int{}, orig...)
	dupe[0] = 4
	if orig[0] != 1 {
		t.FailNow()
	}
}

func TestTorrentString(t *testing.T) {
	tor := &Torrent{}
	s := tor.InfoHash().HexString()
	if s != "0000000000000000000000000000000000000000" {
		t.FailNow()
	}
}

// This benchmark is from the observation that a lot of overlapping Readers on
// a large torrent with small pieces had a lot of overhead in recalculating
// piece priorities everytime a reader (possibly in another Torrent) changed.
func BenchmarkUpdatePiecePriorities(b *testing.B) {
	cl := &Client{}
	t := cl.newTorrent(metainfo.Hash{})
	t.info = &metainfo.Info{
		Pieces:      make([]byte, 20*13410),
		PieceLength: 256 << 10,
	}
	t.makePieces()
	assert.EqualValues(b, 13410, t.numPieces())
	for range iter.N(7) {
		r := t.NewReader()
		r.SetReadahead(32 << 20)
		r.Seek(3500000, 0)
	}
	assert.Len(b, t.readers, 7)
	t.pendPieceRange(0, t.numPieces())
	for i := 0; i < t.numPieces(); i += 3 {
		t.completedPieces.Set(i, true)
	}
	for range iter.N(b.N) {
		t.updatePiecePriorities()
	}
}
<|endoftext|>


<commit_before>package terraform

import (
	"fmt"
	"log"
	"strings"

	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/terraform/config/module"
	"github.com/hashicorp/terraform/dag"
)

// TODO: return the transformers and append them to the list, so we don't lose the log steps
func TransformProviders(providers []string, concrete ConcreteProviderNodeFunc, mod *module.Tree) GraphTransformer {
	// If we have no providers, let the MissingProviderTransformer add anything required.
	// This is used by the destroy edge transformer's internal dependency graph.
	allowAny := providers == nil

	return GraphTransformMulti(
		// Add providers from the config
		&ProviderConfigTransformer{
			Module:    mod,
			Providers: providers,
			Concrete:  concrete,
		},
		// Add any remaining missing providers
		&MissingProviderTransformer{
			AllowAny:  allowAny,
			Providers: providers,
			Concrete:  concrete,
		},
		// Connect the providers
		&ProviderTransformer{},
		// Disable unused providers
		&DisableProviderTransformer{},
		// Connect provider to their parent provider nodes
		&ParentProviderTransformer{},
		// Attach configuration to each provider instance
		&AttachProviderConfigTransformer{
			Module: mod,
		},
	)
}

// GraphNodeProvider is an interface that nodes that can be a provider
// must implement.
// ProviderName returns the name of the provider this satisfies.
// Name returns the full name of the provider in the config.
type GraphNodeProvider interface {
	ProviderName() string
	Name() string
}

// GraphNodeCloseProvider is an interface that nodes that can be a close
// provider must implement. The CloseProviderName returned is the name of
// the provider they satisfy.
type GraphNodeCloseProvider interface {
	CloseProviderName() string
}

// GraphNodeProviderConsumer is an interface that nodes that require
// a provider must implement. ProvidedBy must return the name of the provider
// to use.
type GraphNodeProviderConsumer interface {
	// TODO: make this return s string instead of a []string
	ProvidedBy() []string
	// Set the resolved provider address for this resource.
	SetProvider(string)
}

// ProviderTransformer is a GraphTransformer that maps resources to
// providers within the graph. This will error if there are any resources
// that don't map to proper resources.
type ProviderTransformer struct{}

func (t *ProviderTransformer) Transform(g *Graph) error {
	// Go through the other nodes and match them to providers they need
	var err error
	m := providerVertexMap(g)
	for _, v := range g.Vertices() {
		if pv, ok := v.(GraphNodeProviderConsumer); ok {
			p := pv.ProvidedBy()[0]

			key := providerMapKey(p, pv)
			target := m[key]

			sp, ok := pv.(GraphNodeSubPath)
			if !ok && target == nil {
				// no target, and no path to walk up
				err = multierror.Append(err, fmt.Errorf(
					"%s: provider %s couldn't be found",
					dag.VertexName(v), p))
				break
			}

			// if we don't have a provider at this level, walk up the path looking for one
			for i := 1; target == nil; i++ {
				path := normalizeModulePath(sp.Path())
				if len(path) < i {
					break
				}

				key = ResolveProviderName(p, path[:len(path)-i])
				target = m[key]
				if target != nil {
					break
				}
			}

			if target == nil {
				err = multierror.Append(err, fmt.Errorf(
					"%s: provider %s couldn't be found",
					dag.VertexName(v), p))
				break
			}

			pv.SetProvider(key)
			g.Connect(dag.BasicEdge(v, target))
		}
	}

	return err
}

// CloseProviderTransformer is a GraphTransformer that adds nodes to the
// graph that will close open provider connections that aren't needed anymore.
// A provider connection is not needed anymore once all depended resources
// in the graph are evaluated.
type CloseProviderTransformer struct{}

// FIXME: this doesn't close providers if the root provider is disabled
func (t *CloseProviderTransformer) Transform(g *Graph) error {
	pm := providerVertexMap(g)
	cpm := make(map[string]*graphNodeCloseProvider)
	var err error

	for _, v := range pm {
		p := v.(GraphNodeProvider)

		// get the close provider of this type if we alread created it
		closer := cpm[p.ProviderName()]

		if closer == nil {
			// create a closer for this provider type
			closer = &graphNodeCloseProvider{ProviderNameValue: p.ProviderName()}
			g.Add(closer)
			cpm[p.ProviderName()] = closer
		}

		// Close node depends on the provider itself
		// this is added unconditionally, so it will connect to all instances
		// of the provider. Extra edges will be removed by transitive
		// reduction.
		g.Connect(dag.BasicEdge(closer, p))

		// connect all the provider's resources to the close node
		for _, s := range g.UpEdges(p).List() {
			if _, ok := s.(GraphNodeProviderConsumer); ok {
				g.Connect(dag.BasicEdge(closer, s))
			}
		}
	}

	return err
}

// MissingProviderTransformer is a GraphTransformer that adds nodes for all
// required providers into the graph. Specifically, it creates provider
// configuration nodes for all the providers that we support. These are pruned
// later during an optimization pass.
type MissingProviderTransformer struct {
	// Providers is the list of providers we support.
	Providers []string

	// AllowAny will not check that a provider is supported before adding
	// it to the graph.
	AllowAny bool

	// Concrete, if set, overrides how the providers are made.
	Concrete ConcreteProviderNodeFunc
}

func (t *MissingProviderTransformer) Transform(g *Graph) error {
	// Initialize factory
	if t.Concrete == nil {
		t.Concrete = func(a *NodeAbstractProvider) dag.Vertex {
			return a
		}
	}

	// Create a set of our supported providers
	supported := make(map[string]struct{}, len(t.Providers))
	for _, v := range t.Providers {
		supported[v] = struct{}{}
	}

	var err error
	m := providerVertexMap(g)
	for _, v := range g.Vertices() {
		pv, ok := v.(GraphNodeProviderConsumer)
		if !ok {
			continue
		}

		p := pv.ProvidedBy()[0]

		var path []string
		if sp, ok := pv.(GraphNodeSubPath); ok {
			path = sp.Path()
		}

		key := providerMapKey(p, pv)

		provider := m[key]

		// if we don't have a provider at this level, walk up the path looking for one
		for i := 1; provider == nil && len(path) >= i; i++ {
			key = ResolveProviderName(p, normalizeModulePath(path[:len(path)-i]))
			provider = m[key]
		}

		if provider != nil {
			// we found a provider, but make sure there's a top-level provider too
			if _, ok := m[ResolveProviderName(p, nil)]; ok {
				continue
			}
		}

		// always add a new top level provider
		provider = t.Concrete(&NodeAbstractProvider{
			NameValue: p,
		}).(dag.Vertex)

		key = ResolveProviderName(p, nil)
		m[key] = g.Add(provider)

		pv.SetProvider(key)
	}

	return err
}

// ParentProviderTransformer connects provider nodes to their parents.
//
// This works by finding nodes that are both GraphNodeProviders and
// GraphNodeSubPath. It then connects the providers to their parent
// path. The parent provider is always at the root level.
type ParentProviderTransformer struct{}

func (t *ParentProviderTransformer) Transform(g *Graph) error {
	pm := providerVertexMap(g)
	for _, v := range g.Vertices() {
		// Only care about providers
		pn, ok := v.(GraphNodeProvider)
		if !ok || pn.ProviderName() == "" {
			continue
		}

		// Also require a subpath, if there is no subpath then we
		// can't have a parent.
		if pn, ok := v.(GraphNodeSubPath); ok {
			if len(normalizeModulePath(pn.Path())) <= 1 {
				continue
			}
		}

		// this provider may be disabled, but we can only get it's name from
		// the ProviderName string
		name := ResolveProviderName(strings.SplitN(pn.ProviderName(), " ", 2)[0], nil)
		parent := pm[name]
		if parent != nil {
			g.Connect(dag.BasicEdge(v, parent))
		}

	}
	return nil
}

// PruneProviderTransformer is a GraphTransformer that prunes all the
// providers that aren't needed from the graph. A provider is unneeded if
// no resource or module is using that provider.
type PruneProviderTransformer struct{}

func (t *PruneProviderTransformer) Transform(g *Graph) error {
	for _, v := range g.Vertices() {
		// We only care about the providers
		if pn, ok := v.(GraphNodeProvider); !ok || pn.ProviderName() == "" {
			continue
		}
		// Does anything depend on this? If not, then prune it.
		if s := g.UpEdges(v); s.Len() == 0 {
			if nv, ok := v.(dag.NamedVertex); ok {
				log.Printf("[DEBUG] Pruning provider with no dependencies: %s", nv.Name())
			}
			g.Remove(v)
		}
	}

	return nil
}

// providerMapKey is a helper that gives us the key to use for the
// maps returned by things such as providerVertexMap.
func providerMapKey(k string, v dag.Vertex) string {
	// we create a dummy provider to
	var path []string
	if sp, ok := v.(GraphNodeSubPath); ok {
		path = normalizeModulePath(sp.Path())
	}
	return ResolveProviderName(k, path)
}

func providerVertexMap(g *Graph) map[string]dag.Vertex {
	m := make(map[string]dag.Vertex)
	for _, v := range g.Vertices() {
		if pv, ok := v.(GraphNodeProvider); ok {
			// TODO:  The Name may have meta info, like " (disabled)"
			name := strings.SplitN(pv.Name(), " ", 2)[0]
			m[name] = v
		}
	}

	return m
}

func closeProviderVertexMap(g *Graph) map[string]dag.Vertex {
	m := make(map[string]dag.Vertex)
	for _, v := range g.Vertices() {
		if pv, ok := v.(GraphNodeCloseProvider); ok {
			m[pv.CloseProviderName()] = v
		}
	}

	return m
}

type graphNodeCloseProvider struct {
	ProviderNameValue string
}

func (n *graphNodeCloseProvider) Name() string {
	return fmt.Sprintf("provider.%s (close)", n.ProviderNameValue)
}

// GraphNodeEvalable impl.
func (n *graphNodeCloseProvider) EvalTree() EvalNode {
	return CloseProviderEvalTree(n.ProviderNameValue)
}

// GraphNodeDependable impl.
func (n *graphNodeCloseProvider) DependableName() []string {
	return []string{n.Name()}
}

func (n *graphNodeCloseProvider) CloseProviderName() string {
	return n.ProviderNameValue
}

// GraphNodeDotter impl.
func (n *graphNodeCloseProvider) DotNode(name string, opts *dag.DotOpts) *dag.DotNode {
	if !opts.Verbose {
		return nil
	}
	return &dag.DotNode{
		Name: name,
		Attrs: map[string]string{
			"label": n.Name(),
			"shape": "diamond",
		},
	}
}

// RemovableIfNotTargeted
func (n *graphNodeCloseProvider) RemoveIfNotTargeted() bool {
	// We need to add this so that this node will be removed if
	// it isn't targeted or a dependency of a target.
	return true
}

// graphNodeProviderConsumerDummy is a struct that never enters the real
// graph (though it could to no ill effect). It implements
// GraphNodeProviderConsumer and GraphNodeSubpath as a way to force
// certain transformations.
type graphNodeProviderConsumerDummy struct {
	ProviderValue string
	PathValue     []string
}

func (n *graphNodeProviderConsumerDummy) Path() []string {
	return n.PathValue
}

func (n *graphNodeProviderConsumerDummy) ProvidedBy() []string {
	return []string{n.ProviderValue}
}

func (n *graphNodeProviderConsumerDummy) SetProvider(string) {}
<commit_msg>further simplify the missing provider transformer<commit_after>package terraform

import (
	"fmt"
	"log"
	"strings"

	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/terraform/config/module"
	"github.com/hashicorp/terraform/dag"
)

// TODO: return the transformers and append them to the list, so we don't lose the log steps
func TransformProviders(providers []string, concrete ConcreteProviderNodeFunc, mod *module.Tree) GraphTransformer {
	return GraphTransformMulti(
		// Add providers from the config
		&ProviderConfigTransformer{
			Module:    mod,
			Providers: providers,
			Concrete:  concrete,
		},
		// Add any remaining missing providers
		&MissingProviderTransformer{
			Providers: providers,
			Concrete:  concrete,
		},
		// Connect the providers
		&ProviderTransformer{},
		// Disable unused providers
		&DisableProviderTransformer{},
		// Connect provider to their parent provider nodes
		&ParentProviderTransformer{},
		// Attach configuration to each provider instance
		&AttachProviderConfigTransformer{
			Module: mod,
		},
	)
}

// GraphNodeProvider is an interface that nodes that can be a provider
// must implement.
// ProviderName returns the name of the provider this satisfies.
// Name returns the full name of the provider in the config.
type GraphNodeProvider interface {
	ProviderName() string
	Name() string
}

// GraphNodeCloseProvider is an interface that nodes that can be a close
// provider must implement. The CloseProviderName returned is the name of
// the provider they satisfy.
type GraphNodeCloseProvider interface {
	CloseProviderName() string
}

// GraphNodeProviderConsumer is an interface that nodes that require
// a provider must implement. ProvidedBy must return the name of the provider
// to use.
type GraphNodeProviderConsumer interface {
	// TODO: make this return s string instead of a []string
	ProvidedBy() []string
	// Set the resolved provider address for this resource.
	SetProvider(string)
}

// ProviderTransformer is a GraphTransformer that maps resources to
// providers within the graph. This will error if there are any resources
// that don't map to proper resources.
type ProviderTransformer struct{}

func (t *ProviderTransformer) Transform(g *Graph) error {
	// Go through the other nodes and match them to providers they need
	var err error
	m := providerVertexMap(g)
	for _, v := range g.Vertices() {
		if pv, ok := v.(GraphNodeProviderConsumer); ok {
			p := pv.ProvidedBy()[0]

			key := providerMapKey(p, pv)
			target := m[key]

			sp, ok := pv.(GraphNodeSubPath)
			if !ok && target == nil {
				// no target, and no path to walk up
				err = multierror.Append(err, fmt.Errorf(
					"%s: provider %s couldn't be found",
					dag.VertexName(v), p))
				break
			}

			// if we don't have a provider at this level, walk up the path looking for one
			for i := 1; target == nil; i++ {
				path := normalizeModulePath(sp.Path())
				if len(path) < i {
					break
				}

				key = ResolveProviderName(p, path[:len(path)-i])
				target = m[key]
				if target != nil {
					break
				}
			}

			if target == nil {
				err = multierror.Append(err, fmt.Errorf(
					"%s: provider %s couldn't be found",
					dag.VertexName(v), p))
				break
			}

			pv.SetProvider(key)
			g.Connect(dag.BasicEdge(v, target))
		}
	}

	return err
}

// CloseProviderTransformer is a GraphTransformer that adds nodes to the
// graph that will close open provider connections that aren't needed anymore.
// A provider connection is not needed anymore once all depended resources
// in the graph are evaluated.
type CloseProviderTransformer struct{}

// FIXME: this doesn't close providers if the root provider is disabled
func (t *CloseProviderTransformer) Transform(g *Graph) error {
	pm := providerVertexMap(g)
	cpm := make(map[string]*graphNodeCloseProvider)
	var err error

	for _, v := range pm {
		p := v.(GraphNodeProvider)

		// get the close provider of this type if we alread created it
		closer := cpm[p.ProviderName()]

		if closer == nil {
			// create a closer for this provider type
			closer = &graphNodeCloseProvider{ProviderNameValue: p.ProviderName()}
			g.Add(closer)
			cpm[p.ProviderName()] = closer
		}

		// Close node depends on the provider itself
		// this is added unconditionally, so it will connect to all instances
		// of the provider. Extra edges will be removed by transitive
		// reduction.
		g.Connect(dag.BasicEdge(closer, p))

		// connect all the provider's resources to the close node
		for _, s := range g.UpEdges(p).List() {
			if _, ok := s.(GraphNodeProviderConsumer); ok {
				g.Connect(dag.BasicEdge(closer, s))
			}
		}
	}

	return err
}

// MissingProviderTransformer is a GraphTransformer that adds nodes for all
// required providers into the graph. Specifically, it creates provider
// configuration nodes for all the providers that we support. These are pruned
// later during an optimization pass.
type MissingProviderTransformer struct {
	// Providers is the list of providers we support.
	Providers []string

	// Concrete, if set, overrides how the providers are made.
	Concrete ConcreteProviderNodeFunc
}

func (t *MissingProviderTransformer) Transform(g *Graph) error {
	// Initialize factory
	if t.Concrete == nil {
		t.Concrete = func(a *NodeAbstractProvider) dag.Vertex {
			return a
		}
	}

	var err error
	m := providerVertexMap(g)
	for _, v := range g.Vertices() {
		pv, ok := v.(GraphNodeProviderConsumer)
		if !ok {
			continue
		}

		p := pv.ProvidedBy()[0]
		key := ResolveProviderName(p, nil)
		provider := m[key]

		// we already have it
		if provider != nil {
			continue
		}

		// create the misisng top-level provider
		provider = t.Concrete(&NodeAbstractProvider{
			NameValue: p,
		}).(dag.Vertex)

		m[key] = g.Add(provider)
	}

	return err
}

// ParentProviderTransformer connects provider nodes to their parents.
//
// This works by finding nodes that are both GraphNodeProviders and
// GraphNodeSubPath. It then connects the providers to their parent
// path. The parent provider is always at the root level.
type ParentProviderTransformer struct{}

func (t *ParentProviderTransformer) Transform(g *Graph) error {
	pm := providerVertexMap(g)
	for _, v := range g.Vertices() {
		// Only care about providers
		pn, ok := v.(GraphNodeProvider)
		if !ok || pn.ProviderName() == "" {
			continue
		}

		// Also require a subpath, if there is no subpath then we
		// can't have a parent.
		if pn, ok := v.(GraphNodeSubPath); ok {
			if len(normalizeModulePath(pn.Path())) <= 1 {
				continue
			}
		}

		// this provider may be disabled, but we can only get it's name from
		// the ProviderName string
		name := ResolveProviderName(strings.SplitN(pn.ProviderName(), " ", 2)[0], nil)
		parent := pm[name]
		if parent != nil {
			g.Connect(dag.BasicEdge(v, parent))
		}

	}
	return nil
}

// PruneProviderTransformer is a GraphTransformer that prunes all the
// providers that aren't needed from the graph. A provider is unneeded if
// no resource or module is using that provider.
type PruneProviderTransformer struct{}

func (t *PruneProviderTransformer) Transform(g *Graph) error {
	for _, v := range g.Vertices() {
		// We only care about the providers
		if pn, ok := v.(GraphNodeProvider); !ok || pn.ProviderName() == "" {
			continue
		}
		// Does anything depend on this? If not, then prune it.
		if s := g.UpEdges(v); s.Len() == 0 {
			if nv, ok := v.(dag.NamedVertex); ok {
				log.Printf("[DEBUG] Pruning provider with no dependencies: %s", nv.Name())
			}
			g.Remove(v)
		}
	}

	return nil
}

// providerMapKey is a helper that gives us the key to use for the
// maps returned by things such as providerVertexMap.
func providerMapKey(k string, v dag.Vertex) string {
	// we create a dummy provider to
	var path []string
	if sp, ok := v.(GraphNodeSubPath); ok {
		path = normalizeModulePath(sp.Path())
	}
	return ResolveProviderName(k, path)
}

func providerVertexMap(g *Graph) map[string]dag.Vertex {
	m := make(map[string]dag.Vertex)
	for _, v := range g.Vertices() {
		if pv, ok := v.(GraphNodeProvider); ok {
			// TODO:  The Name may have meta info, like " (disabled)"
			name := strings.SplitN(pv.Name(), " ", 2)[0]
			m[name] = v
		}
	}

	return m
}

func closeProviderVertexMap(g *Graph) map[string]dag.Vertex {
	m := make(map[string]dag.Vertex)
	for _, v := range g.Vertices() {
		if pv, ok := v.(GraphNodeCloseProvider); ok {
			m[pv.CloseProviderName()] = v
		}
	}

	return m
}

type graphNodeCloseProvider struct {
	ProviderNameValue string
}

func (n *graphNodeCloseProvider) Name() string {
	return fmt.Sprintf("provider.%s (close)", n.ProviderNameValue)
}

// GraphNodeEvalable impl.
func (n *graphNodeCloseProvider) EvalTree() EvalNode {
	return CloseProviderEvalTree(n.ProviderNameValue)
}

// GraphNodeDependable impl.
func (n *graphNodeCloseProvider) DependableName() []string {
	return []string{n.Name()}
}

func (n *graphNodeCloseProvider) CloseProviderName() string {
	return n.ProviderNameValue
}

// GraphNodeDotter impl.
func (n *graphNodeCloseProvider) DotNode(name string, opts *dag.DotOpts) *dag.DotNode {
	if !opts.Verbose {
		return nil
	}
	return &dag.DotNode{
		Name: name,
		Attrs: map[string]string{
			"label": n.Name(),
			"shape": "diamond",
		},
	}
}

// RemovableIfNotTargeted
func (n *graphNodeCloseProvider) RemoveIfNotTargeted() bool {
	// We need to add this so that this node will be removed if
	// it isn't targeted or a dependency of a target.
	return true
}

// graphNodeProviderConsumerDummy is a struct that never enters the real
// graph (though it could to no ill effect). It implements
// GraphNodeProviderConsumer and GraphNodeSubpath as a way to force
// certain transformations.
type graphNodeProviderConsumerDummy struct {
	ProviderValue string
	PathValue     []string
}

func (n *graphNodeProviderConsumerDummy) Path() []string {
	return n.PathValue
}

func (n *graphNodeProviderConsumerDummy) ProvidedBy() []string {
	return []string{n.ProviderValue}
}

func (n *graphNodeProviderConsumerDummy) SetProvider(string) {}
<|endoftext|>


<commit_before>package transmission

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

const endpoint = "/transmission/rpc"

type (
	// User to authenticate with Transmission
	User struct {
		Username string
		Password string
	}
	// Client connects to transmission via HTTP
	Client struct {
		URL   string
		token string

		User   *User
		client http.Client
	}
)

// New create new transmission torrent
func New(url string, user *User) *Client {
	return &Client{
		URL:  url + endpoint,
		User: user,
	}
}

func (c *Client) post(body []byte) ([]byte, error) {
	authRequest, err := c.authRequest("POST", body)
	if err != nil {
		return make([]byte, 0), err
	}

	res, err := c.client.Do(authRequest)
	if err != nil {
		return make([]byte, 0), err
	}
	defer res.Body.Close()

	if res.StatusCode == 409 {
		c.getToken()
		authRequest, err := c.authRequest("POST", body)
		if err != nil {
			return make([]byte, 0), err
		}
		res, err = c.client.Do(authRequest)
		if err != nil {
			return make([]byte, 0), err
		}
	}

	resBody, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return make([]byte, 0), err
	}
	return resBody, nil
}

func (c *Client) getToken() error {
	req, err := http.NewRequest("POST", c.URL, strings.NewReader(""))
	if err != nil {
		return err
	}

	if c.User != nil {
		req.SetBasicAuth(c.User.Username, c.User.Password)
	}

	res, err := c.client.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	c.token = res.Header.Get("X-Transmission-Session-Id")
	return nil
}

func (c *Client) authRequest(method string, body []byte) (*http.Request, error) {
	if c.token == "" {
		err := c.getToken()
		if err != nil {
			return nil, err
		}
	}
	req, err := http.NewRequest(method, c.URL, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Add("X-Transmission-Session-Id", c.token)

	if c.User != nil {
		req.SetBasicAuth(c.User.Username, c.User.Password)
	}

	return req, nil
}

// GetTorrents get a list of torrents
func (c *Client) GetTorrents() ([]Torrent, error) {
	cmd := TorrentCommand{
		Method: "torrent-get",
		Arguments: TorrentArguments{
			Fields: []string{
				"id",
				"name",
				"hashString",
				"status",
				"addedDate",
				"leftUntilDone",
				"eta",
				"uploadRatio",
				"rateDownload",
				"rateUpload",
				"downloadDir",
				"isFinished",
				"percentDone",
				"seedRatioMode",
				"error",
				"errorString",
				"files",
				"fileStats",
				"peers",
				"trackers",
				"trackerStats",
			},
		},
	}

	req, err := json.Marshal(&cmd)
	if err != nil {
		return nil, err
	}

	resp, err := c.post(req)
	if err != nil {
		return nil, err
	}

	var out TorrentCommand
	if err := json.Unmarshal(resp, &out); err != nil {
		return nil, err
	}

	return out.Arguments.Torrents, nil
}

// GetSession gets the current session from transmission
func (c *Client) GetSession() (*Session, error) {
	req, err := json.Marshal(SessionCommand{Method: "session-get"})
	if err != nil {
		return nil, err
	}

	resp, err := c.post(req)
	if err != nil {
		return nil, err
	}

	var cmd SessionCommand
	if err := json.Unmarshal(resp, &cmd); err != nil {
		return nil, err
	}

	return &cmd.Session, nil
}

// GetSessionStats gets stats on the current & cumulative session
func (c *Client) GetSessionStats() (*SessionStats, error) {
	req, err := json.Marshal(SessionCommand{Method: "session-stats"})
	if err != nil {
		return nil, err
	}

	resp, err := c.post(req)
	if err != nil {
		return nil, err
	}

	var cmd SessionStatsCmd
	if err := json.Unmarshal(resp, &cmd); err != nil {
		return nil, err
	}

	return &cmd.SessionStats, nil
}
<commit_msg>Fix rpc endpoint to avoid "409 Conflict" errors<commit_after>package transmission

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

const endpoint = "/transmission/rpc/"

type (
	// User to authenticate with Transmission
	User struct {
		Username string
		Password string
	}
	// Client connects to transmission via HTTP
	Client struct {
		URL   string
		token string

		User   *User
		client http.Client
	}
)

// New create new transmission torrent
func New(url string, user *User) *Client {
	return &Client{
		URL:  url + endpoint,
		User: user,
	}
}

func (c *Client) post(body []byte) ([]byte, error) {
	authRequest, err := c.authRequest("POST", body)
	if err != nil {
		return make([]byte, 0), err
	}

	res, err := c.client.Do(authRequest)
	if err != nil {
		return make([]byte, 0), err
	}
	defer res.Body.Close()

	if res.StatusCode == 409 {
		c.getToken()
		authRequest, err := c.authRequest("POST", body)
		if err != nil {
			return make([]byte, 0), err
		}
		res, err = c.client.Do(authRequest)
		if err != nil {
			return make([]byte, 0), err
		}
	}

	resBody, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return make([]byte, 0), err
	}
	return resBody, nil
}

func (c *Client) getToken() error {
	req, err := http.NewRequest("POST", c.URL, strings.NewReader(""))
	if err != nil {
		return err
	}

	if c.User != nil {
		req.SetBasicAuth(c.User.Username, c.User.Password)
	}

	res, err := c.client.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	c.token = res.Header.Get("X-Transmission-Session-Id")
	return nil
}

func (c *Client) authRequest(method string, body []byte) (*http.Request, error) {
	if c.token == "" {
		err := c.getToken()
		if err != nil {
			return nil, err
		}
	}
	req, err := http.NewRequest(method, c.URL, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Add("X-Transmission-Session-Id", c.token)

	if c.User != nil {
		req.SetBasicAuth(c.User.Username, c.User.Password)
	}

	return req, nil
}

// GetTorrents get a list of torrents
func (c *Client) GetTorrents() ([]Torrent, error) {
	cmd := TorrentCommand{
		Method: "torrent-get",
		Arguments: TorrentArguments{
			Fields: []string{
				"id",
				"name",
				"hashString",
				"status",
				"addedDate",
				"leftUntilDone",
				"eta",
				"uploadRatio",
				"rateDownload",
				"rateUpload",
				"downloadDir",
				"isFinished",
				"percentDone",
				"seedRatioMode",
				"error",
				"errorString",
				"files",
				"fileStats",
				"peers",
				"trackers",
				"trackerStats",
			},
		},
	}

	req, err := json.Marshal(&cmd)
	if err != nil {
		return nil, err
	}

	resp, err := c.post(req)
	if err != nil {
		return nil, err
	}

	var out TorrentCommand
	if err := json.Unmarshal(resp, &out); err != nil {
		return nil, err
	}

	return out.Arguments.Torrents, nil
}

// GetSession gets the current session from transmission
func (c *Client) GetSession() (*Session, error) {
	req, err := json.Marshal(SessionCommand{Method: "session-get"})
	if err != nil {
		return nil, err
	}

	resp, err := c.post(req)
	if err != nil {
		return nil, err
	}

	var cmd SessionCommand
	if err := json.Unmarshal(resp, &cmd); err != nil {
		return nil, err
	}

	return &cmd.Session, nil
}

// GetSessionStats gets stats on the current & cumulative session
func (c *Client) GetSessionStats() (*SessionStats, error) {
	req, err := json.Marshal(SessionCommand{Method: "session-stats"})
	if err != nil {
		return nil, err
	}

	resp, err := c.post(req)
	if err != nil {
		return nil, err
	}

	var cmd SessionStatsCmd
	if err := json.Unmarshal(resp, &cmd); err != nil {
		return nil, err
	}

	return &cmd.SessionStats, nil
}
<|endoftext|>