<commit_before>package main

import (
	"testing"
)

func TestNewFile(t *testing.T) {
	startApplication()

	// Register on volume
	b := datastore.NewBlock()

	// Get shard
	shard := b.DataShards[0]

	// == FIRST FILE ==
	// Fake file
	fileBytes := []byte("Hello File")
	fileMeta := newFileMeta("/images/robin/hello.txt")

	// Add file
	updatedFileMeta, err := shard.AddFile(fileMeta, fileBytes)
	// No errors
	if err != nil {
		t.Error("Failed to add file")
	}

	// Validate size
	if updatedFileMeta.Size != uint32(len(fileBytes)) {
		t.Error("Failed size validation")
	}

	// Validate offset (first file in test)
	if updatedFileMeta.StartOffset != 0 {
		t.Error("Start offset must be 0 for first file")
	}

	// == SECOND FILE ==
	// Fake another file
	fileBytes2 := []byte("Hello number two")
	fileMeta2 := newFileMeta("/images/robin/hello2.txt")

	// Add another file
	updatedFileMeta2, err2 := shard.AddFile(fileMeta2, fileBytes2)
	if err2 != nil {
		t.Error(err2)
	}

	// Offset of file 2 should be after file 1
	if updatedFileMeta2.StartOffset != uint32(len(fileBytes)) {
		t.Error("Should start after file 1")
	}

	// Validate shard meta
	if shard.FileCount() != 2 {
		t.Error("Shard should contain 2 files")
	}

	// Validate shard index
	if shard.shardIndex.Test("/non-existing") == true {
		t.Error("Shard index false positive")
	}
	if shard.shardIndex.Test(fileMeta.FullName) == false {
		t.Error("Shard index does not contain file 1")
	}
	if shard.shardIndex.Test(fileMeta2.FullName) == false {
		t.Error("Shard index does not contain file 2")
	}

	// Validate shard meta (de)serialization
	shardFileMetaBytes := shard.shardFileMeta.Bytes()
	if len(shardFileMetaBytes) < 1 {
		panic("Shard file meta bytes must be non-zero")
	}

	// Load these bytes into new to confirm it's still intact
	newShardFileMetaInstance := newShardFileMeta()
	if len(newShardFileMetaInstance.fileMeta) != 0 {
		panic("Empty shard file meta must be empty")
	}
	newShardFileMetaInstance.FromBytes(shardFileMetaBytes)
	if len(newShardFileMetaInstance.fileMeta) != 2 {
		panic("Loaded shard file meta must be 2")
	}

	// Persist shard to disk
	shard.Persist()

	// Reset pointers in this shard
	shard.contents = nil
	shard.shardIndex = nil
	shard.shardMeta = nil
	shard.shardFileMeta = nil

	// Read from disk
	shard.Load()
}
<commit_msg>Write tests for the binary loading<commit_after>package main

import (
	"testing"
)

func TestNewFile(t *testing.T) {
	startApplication()

	// Register on volume
	b := datastore.NewBlock()

	// Get shard
	shard := b.DataShards[0]

	// == FIRST FILE ==
	// Fake file
	fileBytes := []byte("Hello File")
	fileMeta := newFileMeta("/images/robin/hello.txt")

	// Add file
	updatedFileMeta, err := shard.AddFile(fileMeta, fileBytes)
	// No errors
	if err != nil {
		t.Error("Failed to add file")
	}

	// Validate size
	if updatedFileMeta.Size != uint32(len(fileBytes)) {
		t.Error("Failed size validation")
	}

	// Validate offset (first file in test)
	if updatedFileMeta.StartOffset != 0 {
		t.Error("Start offset must be 0 for first file")
	}

	// == SECOND FILE ==
	// Fake another file
	fileBytes2 := []byte("Hello number two")
	fileMeta2 := newFileMeta("/images/robin/hello2.txt")

	// Add another file
	updatedFileMeta2, err2 := shard.AddFile(fileMeta2, fileBytes2)
	if err2 != nil {
		t.Error(err2)
	}

	// Offset of file 2 should be after file 1
	if updatedFileMeta2.StartOffset != uint32(len(fileBytes)) {
		t.Error("Should start after file 1")
	}

	// Validate shard meta
	if shard.FileCount() != 2 {
		t.Error("Shard should contain 2 files")
	}

	// Validate shard index
	if shard.shardIndex.Test("/non-existing") == true {
		t.Error("Shard index false positive")
	}
	if shard.shardIndex.Test(fileMeta.FullName) == false {
		t.Error("Shard index does not contain file 1")
	}
	if shard.shardIndex.Test(fileMeta2.FullName) == false {
		t.Error("Shard index does not contain file 2")
	}

	// Validate shard meta (de)serialization
	shardFileMetaBytes := shard.shardFileMeta.Bytes()
	if len(shardFileMetaBytes) < 1 {
		t.Error("Shard file meta bytes must be non-zero")
	}

	// Load these bytes into new to confirm it's still intact
	newShardFileMetaInstance := newShardFileMeta()
	if len(newShardFileMetaInstance.fileMeta) != 0 {
		t.Error("Empty shard file meta must be empty")
	}
	newShardFileMetaInstance.FromBytes(shardFileMetaBytes)
	if len(newShardFileMetaInstance.fileMeta) != 2 {
		t.Error("Loaded shard file meta must be 2")
	}

	// Persist shard to disk
	shard.Persist()

	// Reset pointers in this shard
	shard.contents = nil
	shard.shardIndex = nil
	shard.shardMeta = nil
	shard.shardFileMeta = nil

	// Read from disk
	shard.Load()

	// Validate (no need to check contents as those are not used directly)
	if shard.shardIndex == nil {
		t.Error("Failed to load shard index")
	}
	if shard.shardMeta == nil {
		t.Error("Failed to load shard meta")
	}
	if shard.shardFileMeta == nil {
		t.Error("Failed to load shard file meta")
	}

	// Functional validaton on the index
	if shard.shardIndex.Test(fileMeta.FullName) == false {
		t.Error("Shard index (after persist, nil, load) does not contain file 1")
	}

	// Function validation on shard meta
	if shard.FileCount() != 2 {
		t.Error("Shard (after persist, nil, load) should contain 2 files")
	}
}
<|endoftext|>


<commit_before>package scipipe

import (
	"fmt"
	"github.com/stretchr/testify/assert"
	"os"
	t "testing"
	"time"
)

func initTestLogs() {
	InitLogWarn()
}

func TestShellHasInOutPorts(t *t.T) {
	initTestLogs()

	tt := Sh("echo {i:in1} > {o:out1}")
	tt.OutPathFuncs["out1"] = func() string {
		return fmt.Sprint(tt.InPaths["in1"], ".bar")
	}
	tt.InPorts["in1"] = make(chan *FileTarget, BUFSIZE)

	go tt.Run()
	go func() {
		defer close(tt.InPorts["in1"])
		tt.InPorts["in1"] <- NewFileTarget("foo.txt")
	}()
	<-tt.OutPorts["out1"]

	assert.NotNil(t, tt.InPorts["in1"], "InPorts are nil!")
	assert.NotNil(t, tt.OutPorts["out1"], "OutPorts are nil!")

	cleanFiles("foo.txt", "foo.txt.bar")
}

func TestShellCloseOutPortOnInPortClose(t *t.T) {
	initTestLogs()

	fooTask := Sh("echo foo > {o:out1}")
	fooTask.OutPathFuncs["out1"] = func() string {
		return "foo.txt"
	}

	barReplacer := Sh("sed 's/foo/bar/g' {i:foo} > {o:bar}")
	barReplacer.OutPathFuncs["bar"] = func() string {
		return barReplacer.GetInPath("foo") + ".bar"
	}

	barReplacer.InPorts["foo"] = fooTask.OutPorts["out1"]

	go fooTask.Run()
	go barReplacer.Run()
	<-barReplacer.OutPorts["bar"]

	// Assert no more content coming on channels
	assert.Nil(t, <-fooTask.OutPorts["out1"])
	assert.Nil(t, <-barReplacer.OutPorts["bar"])

	_, fooErr := os.Stat("foo.txt")
	assert.Nil(t, fooErr)
	_, barErr := os.Stat("foo.txt.bar")
	assert.Nil(t, barErr)

	cleanFiles("foo.txt", "foo.txt.bar")
}

func TestReplacePlaceholdersInCmd(t *t.T) {
	initTestLogs()

	rawCmd := "echo {i:in1} > {o:out1}"
	tt := Sh(rawCmd)
	tt.OutPathFuncs["out1"] = func() string {
		return fmt.Sprint(tt.InPaths["in1"], ".bar")
	}

	tt.InPorts["in1"] = make(chan *FileTarget, BUFSIZE)
	ift := NewFileTarget("foo.txt")
	go func() {
		defer close(tt.InPorts["in1"])
		tt.InPorts["in1"] <- ift
	}()

	// Assert inport is still open after first read
	inPortsClosed := tt.receiveInputs()
	assert.Equal(t, false, inPortsClosed)

	// Assert inport is closed after second read
	inPortsClosed = tt.receiveInputs()
	assert.Equal(t, true, inPortsClosed)

	// Assert InPath is correct
	assert.Equal(t, "foo.txt", tt.InPaths["in1"], "foo.txt")

	// Assert placeholders are correctly replaced in command
	outTargets := tt.createOutTargets()
	cmd := tt.formatCommand(rawCmd, outTargets)
	assert.EqualValues(t, "echo foo.txt > foo.txt.bar.tmp", cmd, "Command not properly formatted!")

	// Test prepend
	tt.Prepend = "dash"
	cmd = tt.formatCommand(rawCmd, outTargets)
	assert.EqualValues(t, "dash echo foo.txt > foo.txt.bar.tmp", cmd, "Prepend not working!")
}

func TestParameterCommand(t *t.T) {
	initTestLogs()

	cmb := NewCombinatoricsTask()

	// An abc file printer
	abc := Sh("echo {p:a} {p:b} {p:c} > {o:out}")
	abc.OutPathFuncs["out"] = func() string {
		return fmt.Sprintf(
			"%s_%s_%s.txt",
			abc.Params["a"],
			abc.Params["b"],
			abc.Params["c"],
		)
	}

	// A printer task
	prt := Sh("cat {i:in} >> /tmp/log.txt; rm {i:in}")

	// Connection info
	abc.ParamPorts["a"] = cmb.A
	abc.ParamPorts["b"] = cmb.B
	abc.ParamPorts["c"] = cmb.C
	prt.InPorts["in"] = abc.OutPorts["out"]

	pl := NewPipeline()
	pl.AddTasks(cmb, abc, prt)
	pl.Run()

	// Run tests
	_, err := os.Stat("/tmp/log.txt")
	assert.Nil(t, err)

	cleanFiles("/tmp/log.txt")
}

func TestTaskWithoutInputsOutputs(t *t.T) {
	initTestLogs()
	Debug.Println("Starting test TestTaskWithoutInputsOutputs")

	f := "/tmp/hej.txt"
	tsk := Sh("echo hej > " + f)
	tsk.Run()
	_, err := os.Stat(f)
	assert.Nil(t, err)
	cleanFiles(f)
}

func TestDontOverWriteExistingOutputs(t *t.T) {
	InitLogError()
	Debug.Println("Starting test TestDontOverWriteExistingOutputs")

	f := "/tmp/hej.txt"

	// Assert file does not exist before running
	_, e1 := os.Stat(f)
	assert.NotNil(t, e1)

	// Run pipeline a first time
	tsk := Sh("echo hej > {o:hej}")
	tsk.OutPathFuncs["hej"] = func() string { return f }
	prt := Sh("echo {i:in} Done!")
	prt.InPorts["in"] = tsk.OutPorts["hej"]
	pl := NewPipeline()
	pl.AddTasks(tsk, prt)
	pl.Run()

	// Assert file DO exist after running
	fiBef, e2 := os.Stat(f)
	assert.Nil(t, e2)

	// Get modified time before
	mtBef := fiBef.ModTime()

	// Make sure some time has passed before the second write
	time.Sleep(1 * time.Millisecond)

	Debug.Println("Try running the same workflow again ...")
	// Run again with different output
	tsk = Sh("echo hej > {o:hej}")
	tsk.OutPathFuncs["hej"] = func() string { return f }
	prt.InPorts["in"] = tsk.OutPorts["hej"]
	pl = NewPipeline()
	pl.AddTasks(tsk, prt)
	pl.Run()

	// Assert exists
	fiAft, e3 := os.Stat(f)
	assert.Nil(t, e3)

	// Get modified time AFTER second run
	mtAft := fiAft.ModTime()

	// Assert file is not modified!
	assert.EqualValues(t, mtBef, mtAft)

	cleanFiles(f)
}

func TestShellExpand(t *t.T) {
	initTestLogs()

	cmdPat := "echo {p:txt} > {i:in}; cat {i:in} > {o:out}"
	expectedCmd := "echo hej > in.txt; cat in.txt > out.txt"

	params := make(map[string]string)
	params["txt"] = "hej"
	ipaths := make(map[string]string)
	ipaths["in"] = "in.txt"
	opaths := make(map[string]string)
	opaths["out"] = "out.txt"

	cmd := expandCommandParamsAndPaths(cmdPat, params, ipaths, opaths)
	assert.EqualValues(t, expectedCmd, cmd, "Command not properly expanded!")

	st := ShExp(cmdPat, ipaths, opaths, params)

	// Assert that no ports are created, since all place holders
	// are replaced by the provided maps.
	assert.EqualValues(t, 0, len(st.InPorts), "Inports created where it should not!")
	assert.EqualValues(t, 0, len(st.OutPorts), "OutPorts created where it should not!")
	assert.EqualValues(t, 0, len(st.ParamPorts), "ParamPorts created where it should not!")

	st.Run()
	assert.True(t, NewFileTarget("out.txt").Exists())

	cleanFiles("in.txt", "out.txt")
}

func TestSendsOrderedOutputs(t *t.T) {
	initTestLogs()

	fnames := []string{}
	for i := 1; i <= 10; i++ {
		fnames = append(fnames, fmt.Sprintf("/tmp/f%d.txt", i))
	}

	fq := NewFileQueue(fnames...)

	fc := Sh("sleep 1; echo {i:in} > {o:out}")
	fc.Spawn = true

	sl := Sh("sleep 1; cat {i:in} > {o:out}")
	sl.Spawn = true

	fc.OutPathFuncs["out"] = func() string { return fc.GetInPath("in") }
	sl.OutPathFuncs["out"] = func() string { return sl.GetInPath("in") + ".copy.txt" }

	go fq.Run()
	go fc.Run()
	go sl.Run()

	fc.InPorts["in"] = fq.Out
	sl.InPorts["in"] = fc.OutPorts["out"]

	assert.NotEmpty(t, sl.OutPorts)

	var expFname string
	i := 1
	for ft := range sl.OutPorts["out"] {
		expFname = fmt.Sprintf("/tmp/f%d.txt.copy.txt", i)
		assert.EqualValues(t, expFname, ft.GetPath())
		i++
	}
	expFnames := []string{}
	for i := 1; i <= 10; i++ {
		expFnames = append(expFnames, fmt.Sprintf("/tmp/f%d.txt.copy.txt", i))
	}
	cleanFiles(fnames...)
	cleanFiles(expFnames...)
}

// Helper functions

func cleanFiles(fileNames ...string) {
	Debug.Println("Starting to remove files:", fileNames)
	for _, fileName := range fileNames {
		if _, err := os.Stat(fileName); err == nil {
			os.Remove(fileName)
			Debug.Println("Successfully removed file", fileName)
		}
	}
}

// Helper tasks

type CombinatoricsTask struct {
	BaseTask
	A chan string
	B chan string
	C chan string
}

func NewCombinatoricsTask() *CombinatoricsTask {
	return &CombinatoricsTask{
		A: make(chan string, BUFSIZE),
		B: make(chan string, BUFSIZE),
		C: make(chan string, BUFSIZE),
	}
}

func (proc *CombinatoricsTask) Run() {
	defer close(proc.A)
	defer close(proc.B)
	defer close(proc.C)

	for _, a := range SS("a1", "a2", "a3") {
		for _, b := range SS("b1", "b2", "b3") {
			for _, c := range SS("c1", "c2", "c3") {
				proc.A <- a
				proc.B <- b
				proc.C <- c
			}
		}
	}
}
<commit_msg>Remove unnecessary sleeps<commit_after>package scipipe

import (
	"fmt"
	"github.com/stretchr/testify/assert"
	"os"
	t "testing"
	"time"
)

func initTestLogs() {
	InitLogWarn()
}

func TestShellHasInOutPorts(t *t.T) {
	initTestLogs()

	tt := Sh("echo {i:in1} > {o:out1}")
	tt.OutPathFuncs["out1"] = func() string {
		return fmt.Sprint(tt.InPaths["in1"], ".bar")
	}
	tt.InPorts["in1"] = make(chan *FileTarget, BUFSIZE)

	go tt.Run()
	go func() {
		defer close(tt.InPorts["in1"])
		tt.InPorts["in1"] <- NewFileTarget("foo.txt")
	}()
	<-tt.OutPorts["out1"]

	assert.NotNil(t, tt.InPorts["in1"], "InPorts are nil!")
	assert.NotNil(t, tt.OutPorts["out1"], "OutPorts are nil!")

	cleanFiles("foo.txt", "foo.txt.bar")
}

func TestShellCloseOutPortOnInPortClose(t *t.T) {
	initTestLogs()

	fooTask := Sh("echo foo > {o:out1}")
	fooTask.OutPathFuncs["out1"] = func() string {
		return "foo.txt"
	}

	barReplacer := Sh("sed 's/foo/bar/g' {i:foo} > {o:bar}")
	barReplacer.OutPathFuncs["bar"] = func() string {
		return barReplacer.GetInPath("foo") + ".bar"
	}

	barReplacer.InPorts["foo"] = fooTask.OutPorts["out1"]

	go fooTask.Run()
	go barReplacer.Run()
	<-barReplacer.OutPorts["bar"]

	// Assert no more content coming on channels
	assert.Nil(t, <-fooTask.OutPorts["out1"])
	assert.Nil(t, <-barReplacer.OutPorts["bar"])

	_, fooErr := os.Stat("foo.txt")
	assert.Nil(t, fooErr)
	_, barErr := os.Stat("foo.txt.bar")
	assert.Nil(t, barErr)

	cleanFiles("foo.txt", "foo.txt.bar")
}

func TestReplacePlaceholdersInCmd(t *t.T) {
	initTestLogs()

	rawCmd := "echo {i:in1} > {o:out1}"
	tt := Sh(rawCmd)
	tt.OutPathFuncs["out1"] = func() string {
		return fmt.Sprint(tt.InPaths["in1"], ".bar")
	}

	tt.InPorts["in1"] = make(chan *FileTarget, BUFSIZE)
	ift := NewFileTarget("foo.txt")
	go func() {
		defer close(tt.InPorts["in1"])
		tt.InPorts["in1"] <- ift
	}()

	// Assert inport is still open after first read
	inPortsClosed := tt.receiveInputs()
	assert.Equal(t, false, inPortsClosed)

	// Assert inport is closed after second read
	inPortsClosed = tt.receiveInputs()
	assert.Equal(t, true, inPortsClosed)

	// Assert InPath is correct
	assert.Equal(t, "foo.txt", tt.InPaths["in1"], "foo.txt")

	// Assert placeholders are correctly replaced in command
	outTargets := tt.createOutTargets()
	cmd := tt.formatCommand(rawCmd, outTargets)
	assert.EqualValues(t, "echo foo.txt > foo.txt.bar.tmp", cmd, "Command not properly formatted!")

	// Test prepend
	tt.Prepend = "dash"
	cmd = tt.formatCommand(rawCmd, outTargets)
	assert.EqualValues(t, "dash echo foo.txt > foo.txt.bar.tmp", cmd, "Prepend not working!")
}

func TestParameterCommand(t *t.T) {
	initTestLogs()

	cmb := NewCombinatoricsTask()

	// An abc file printer
	abc := Sh("echo {p:a} {p:b} {p:c} > {o:out}")
	abc.OutPathFuncs["out"] = func() string {
		return fmt.Sprintf(
			"%s_%s_%s.txt",
			abc.Params["a"],
			abc.Params["b"],
			abc.Params["c"],
		)
	}

	// A printer task
	prt := Sh("cat {i:in} >> /tmp/log.txt; rm {i:in}")

	// Connection info
	abc.ParamPorts["a"] = cmb.A
	abc.ParamPorts["b"] = cmb.B
	abc.ParamPorts["c"] = cmb.C
	prt.InPorts["in"] = abc.OutPorts["out"]

	pl := NewPipeline()
	pl.AddTasks(cmb, abc, prt)
	pl.Run()

	// Run tests
	_, err := os.Stat("/tmp/log.txt")
	assert.Nil(t, err)

	cleanFiles("/tmp/log.txt")
}

func TestTaskWithoutInputsOutputs(t *t.T) {
	initTestLogs()
	Debug.Println("Starting test TestTaskWithoutInputsOutputs")

	f := "/tmp/hej.txt"
	tsk := Sh("echo hej > " + f)
	tsk.Run()
	_, err := os.Stat(f)
	assert.Nil(t, err)
	cleanFiles(f)
}

func TestDontOverWriteExistingOutputs(t *t.T) {
	InitLogError()
	Debug.Println("Starting test TestDontOverWriteExistingOutputs")

	f := "/tmp/hej.txt"

	// Assert file does not exist before running
	_, e1 := os.Stat(f)
	assert.NotNil(t, e1)

	// Run pipeline a first time
	tsk := Sh("echo hej > {o:hej}")
	tsk.OutPathFuncs["hej"] = func() string { return f }
	prt := Sh("echo {i:in} Done!")
	prt.InPorts["in"] = tsk.OutPorts["hej"]
	pl := NewPipeline()
	pl.AddTasks(tsk, prt)
	pl.Run()

	// Assert file DO exist after running
	fiBef, e2 := os.Stat(f)
	assert.Nil(t, e2)

	// Get modified time before
	mtBef := fiBef.ModTime()

	// Make sure some time has passed before the second write
	time.Sleep(1 * time.Millisecond)

	Debug.Println("Try running the same workflow again ...")
	// Run again with different output
	tsk = Sh("echo hej > {o:hej}")
	tsk.OutPathFuncs["hej"] = func() string { return f }
	prt.InPorts["in"] = tsk.OutPorts["hej"]
	pl = NewPipeline()
	pl.AddTasks(tsk, prt)
	pl.Run()

	// Assert exists
	fiAft, e3 := os.Stat(f)
	assert.Nil(t, e3)

	// Get modified time AFTER second run
	mtAft := fiAft.ModTime()

	// Assert file is not modified!
	assert.EqualValues(t, mtBef, mtAft)

	cleanFiles(f)
}

func TestShellExpand(t *t.T) {
	initTestLogs()

	cmdPat := "echo {p:txt} > {i:in}; cat {i:in} > {o:out}"
	expectedCmd := "echo hej > in.txt; cat in.txt > out.txt"

	params := make(map[string]string)
	params["txt"] = "hej"
	ipaths := make(map[string]string)
	ipaths["in"] = "in.txt"
	opaths := make(map[string]string)
	opaths["out"] = "out.txt"

	cmd := expandCommandParamsAndPaths(cmdPat, params, ipaths, opaths)
	assert.EqualValues(t, expectedCmd, cmd, "Command not properly expanded!")

	st := ShExp(cmdPat, ipaths, opaths, params)

	// Assert that no ports are created, since all place holders
	// are replaced by the provided maps.
	assert.EqualValues(t, 0, len(st.InPorts), "Inports created where it should not!")
	assert.EqualValues(t, 0, len(st.OutPorts), "OutPorts created where it should not!")
	assert.EqualValues(t, 0, len(st.ParamPorts), "ParamPorts created where it should not!")

	st.Run()
	assert.True(t, NewFileTarget("out.txt").Exists())

	cleanFiles("in.txt", "out.txt")
}

func TestSendsOrderedOutputs(t *t.T) {
	initTestLogs()

	fnames := []string{}
	for i := 1; i <= 10; i++ {
		fnames = append(fnames, fmt.Sprintf("/tmp/f%d.txt", i))
	}

	fq := NewFileQueue(fnames...)

	fc := Sh("echo {i:in} > {o:out}")
	fc.Spawn = true

	sl := Sh("cat {i:in} > {o:out}")
	sl.Spawn = true

	fc.OutPathFuncs["out"] = func() string { return fc.GetInPath("in") }
	sl.OutPathFuncs["out"] = func() string { return sl.GetInPath("in") + ".copy.txt" }

	go fq.Run()
	go fc.Run()
	go sl.Run()

	fc.InPorts["in"] = fq.Out
	sl.InPorts["in"] = fc.OutPorts["out"]

	assert.NotEmpty(t, sl.OutPorts)

	var expFname string
	i := 1
	for ft := range sl.OutPorts["out"] {
		expFname = fmt.Sprintf("/tmp/f%d.txt.copy.txt", i)
		assert.EqualValues(t, expFname, ft.GetPath())
		i++
	}
	expFnames := []string{}
	for i := 1; i <= 10; i++ {
		expFnames = append(expFnames, fmt.Sprintf("/tmp/f%d.txt.copy.txt", i))
	}
	cleanFiles(fnames...)
	cleanFiles(expFnames...)
}

// Helper functions

func cleanFiles(fileNames ...string) {
	Debug.Println("Starting to remove files:", fileNames)
	for _, fileName := range fileNames {
		if _, err := os.Stat(fileName); err == nil {
			os.Remove(fileName)
			Debug.Println("Successfully removed file", fileName)
		}
	}
}

// Helper tasks

type CombinatoricsTask struct {
	BaseTask
	A chan string
	B chan string
	C chan string
}

func NewCombinatoricsTask() *CombinatoricsTask {
	return &CombinatoricsTask{
		A: make(chan string, BUFSIZE),
		B: make(chan string, BUFSIZE),
		C: make(chan string, BUFSIZE),
	}
}

func (proc *CombinatoricsTask) Run() {
	defer close(proc.A)
	defer close(proc.B)
	defer close(proc.C)

	for _, a := range SS("a1", "a2", "a3") {
		for _, b := range SS("b1", "b2", "b3") {
			for _, c := range SS("c1", "c2", "c3") {
				proc.A <- a
				proc.B <- b
				proc.C <- c
			}
		}
	}
}
<|endoftext|>


<commit_before>package somatree

import "github.com/satori/go.uuid"

type SomaTreeElemNode struct {
	Id              uuid.UUID
	Name            string
	Parent          SomaTreeNodeReceiver
	PropertyOncall  map[string]*SomaTreePropertyOncall
	PropertyService map[string]*SomaTreePropertyService
	PropertySystem  map[string]*SomaTreePropertySystem
	PropertyCustom  map[string]*SomaTreePropertyCustom
	Checks          map[string]*SomaTreeCheck
}

func NewNode() *SomaTreeElemNode {
	return new(SomaTreeElemNode)
}

func (ten *SomaTreeElemNode) GetID() string {
	return ten.Id.String()
}

func (ten *SomaTreeElemNode) GetName() string {
	return ten.Name
}

func (ten *SomaTreeElemNode) Attach(a AttachRequest) {
}

func (ten *SomaTreeElemNode) ReAttach(a AttachRequest) {
}

func (ten *SomaTreeElemNode) setParent(p SomaTreeReceiver) {
}

func (ten *SomaTreeElemNode) attachToBucket(a AttachRequest) {
}

func (ten *SomaTreeElemNode) attachToGroup(a AttachRequest) {
}

func (ten *SomaTreeElemNode) attachToCluster(a AttachRequest) {
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>Implement empty interface methods<commit_after>package somatree

import (
	"fmt"
	"log"
	"reflect"

	"github.com/satori/go.uuid"
)

type SomaTreeElemNode struct {
	Id       uuid.UUID
	Name     string
	AssetId  uuid.UUID
	Team     uuid.UUID
	ServerId uuid.UUID
	State    string
	Online   bool
	Deleted  bool
	Type     string
	Parent   SomaTreeNodeReceiver `json:"-"`
	//PropertyOncall  map[string]*SomaTreePropertyOncall
	//PropertyService map[string]*SomaTreePropertyService
	//PropertySystem  map[string]*SomaTreePropertySystem
	//PropertyCustom  map[string]*SomaTreePropertyCustom
	//Checks          map[string]*SomaTreeCheck
}

type NodeSpec struct {
	Id       uuid.UUID
	Name     string
	AssetId  uuid.UUID
	Team     uuid.UUID
	ServerId uuid.UUID
	State    string
}

//
// NEW
func NewNode(name string) *SomaTreeElemNode {
	ten := new(SomaTreeElemNode)
	ten.Id = uuid.NewV4()
	ten.Name = name
	ten.Type = "node"
	//ten.PropertyOncall = make(map[string]*SomaTreePropertyOncall)
	//ten.PropertyService = make(map[string]*SomaTreePropertyService)
	//ten.PropertySystem = make(map[string]*SomaTreePropertySystem)
	//ten.PropertyCustom = make(map[string]*SomaTreePropertyCustom)
	//ten.Checks = make(map[string]*SomaTreeCheck)

	return ten
}

//
// Interface:
func (ten *SomaTreeElemNode) GetID() string {
	return ten.Id.String()
}

func (ten *SomaTreeElemNode) GetName() string {
	return ten.Name
}

func (ten *SomaTreeElemNode) GetType() string {
	return ten.Type
}

//
// Interface: SomaTreeAttacher
func (ten *SomaTreeElemNode) Attach(a AttachRequest) {
	switch {
	case a.ParentType == "bucket":
		ten.attachToBucket(a)
	case a.ParentType == "group":
		ten.attachToGroup(a)
	case a.ParentType == "cluster":
		ten.attachToCluster(a)
	default:
		panic(`SomaTreeElemNode.Attach`)
	}
}

func (ten *SomaTreeElemNode) ReAttach(a AttachRequest) {
	log.Fatal("Not implemented")
}

func (ten *SomaTreeElemNode) setParent(p SomaTreeReceiver) {
	switch p.(type) {
	case *SomaTreeElemBucket:
		ten.setNodeParent(p.(SomaTreeNodeReceiver))
	case *SomaTreeElemGroup:
		ten.setNodeParent(p.(SomaTreeNodeReceiver))
	case *SomaTreeElemCluster:
		ten.setNodeParent(p.(SomaTreeNodeReceiver))
	default:
		fmt.Printf("Type: %s\n", reflect.TypeOf(p))
		panic(`SomaTreeElemNode.setParent`)
	}
}

func (ten *SomaTreeElemNode) setNodeParent(p SomaTreeNodeReceiver) {
	ten.Parent = p
}

//
// Interface: SomaTreeBucketAttacher
func (ten *SomaTreeElemNode) attachToBucket(a AttachRequest) {
	a.Root.Receive(ReceiveRequest{
		ParentType: a.ParentType,
		ParentId:   a.ParentId,
		ParentName: a.ParentName,
		ChildType:  ten.Type,
		Node:       ten,
	})
}

//
// Interface: SomaTreeGroupAttacher
func (ten *SomaTreeElemNode) attachToGroup(a AttachRequest) {
	a.Root.Receive(ReceiveRequest{
		ParentType: a.ParentType,
		ParentId:   a.ParentId,
		ParentName: a.ParentName,
		ChildType:  ten.Type,
		Node:       ten,
	})
}

//
// Interface: SomaTreeClusterAttacher
func (ten *SomaTreeElemNode) attachToCluster(a AttachRequest) {
	a.Root.Receive(ReceiveRequest{
		ParentType: a.ParentType,
		ParentId:   a.ParentId,
		ParentName: a.ParentName,
		ChildType:  ten.Type,
		Node:       ten,
	})
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<|endoftext|>


<commit_before>package main

import "runtime"
import "sync"
import "fmt"

const TOTAL int = 5

func Dine(id int, l_fork, r_fork chan int) {
	//Get the left fork
	left := <-l_fork
	runtime.Gosched()
	right := <-r_fork

	fmt.Println(id, "is eating with forks", left, "and", right)
	runtime.Gosched()
	fmt.Println(id, "is finished eating")

	//Give back the forks
	l_fork <- left
	r_fork <- right
	runtime.Gosched()
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	var wg sync.WaitGroup
	forks := make([]chan int, TOTAL)
	for i := 0; i < TOTAL; i++ {
		wg.Add(1)
		go func(k int, wg *sync.WaitGroup) {
			defer wg.Done()

			fork := make(chan int, 1)
			forks[k] = fork
			fork <- k
		}(i, &wg)
	}
	wg.Wait()
	var wg2 sync.WaitGroup
	for i := 0; i < TOTAL; i++ {
		wg2.Add(1)
		go func(id int, wg *sync.WaitGroup) {
			defer wg.Done()

			Dine(id, forks[id], forks[(id+1)%TOTAL])
		}(i, &wg2)
	}
	wg2.Wait()
	for i := 0; i < TOTAL; i++ {
		close(forks[i])
		<-forks[i]
	}
}
<commit_msg>Remove scheduling condition<commit_after>package main

import "runtime"
import "sync"
import "fmt"

const TOTAL int = 5

func Dine(id int, l_fork, r_fork chan int) {
	//Get the left fork
	left := <-l_fork
	right := <-r_fork

	fmt.Println(id, "is eating with forks", left, "and", right)
	runtime.Gosched()
	fmt.Println(id, "is finished eating")

	//Give back the forks
	l_fork <- left
	r_fork <- right
	runtime.Gosched()
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	var wg sync.WaitGroup
	forks := make([]chan int, TOTAL)
	for i := 0; i < TOTAL; i++ {
		wg.Add(1)
		go func(k int, wg *sync.WaitGroup) {
			defer wg.Done()

			fork := make(chan int, 1)
			forks[k] = fork
			fork <- k
		}(i, &wg)
	}
	wg.Wait()
	var wg2 sync.WaitGroup
	for i := 0; i < TOTAL; i++ {
		wg2.Add(1)
		go func(id int, wg *sync.WaitGroup) {
			defer wg.Done()

			Dine(id, forks[id], forks[(id+1)%TOTAL])
		}(i, &wg2)
	}
	wg2.Wait()
	for i := 0; i < TOTAL; i++ {
		close(forks[i])
		<-forks[i]
	}
}
<|endoftext|>


<commit_before>package ioextras

import (
	"errors"
	"io"
	"io/ioutil"
	"os"
)

type RandomAccessStore interface {
	io.ReaderAt
	io.WriterAt
	io.Closer
}

type SizedRandomAccessStore interface {
	RandomAccessStore
	Size() (int64, error)
}

type NamedRandomAccessStore interface {
	RandomAccessStore
	Name() string
}

type RandomAccessStoreFactory interface {
	RandomAccessStore() (RandomAccessStore, error)
}

type SeekerWrapper struct {
	s  RandomAccessStore
	sk io.Seeker
	ns NamedRandomAccessStore
}

func (s *SeekerWrapper) ReadAt(p []byte, offset int64) (int, error) { return s.s.ReadAt(p, offset) }

func (s *SeekerWrapper) WriteAt(p []byte, offset int64) (int, error) { return s.s.WriteAt(p, offset) }

func (s *SeekerWrapper) Close() error { return s.s.Close() }

func (s *SeekerWrapper) Name() string {
	if s.ns != nil {
		return s.ns.Name()
	} else {
		return ""
	}
}

func (s *SeekerWrapper) Size() (int64, error) {
	return s.sk.Seek(0, os.SEEK_END)
}

func NewSeekerWrapper(s RandomAccessStore) *SeekerWrapper {
	ns, _ := s.(NamedRandomAccessStore)
	return &SeekerWrapper{s, s.(io.Seeker), ns}
}

type StoreReadWriter struct {
	s    RandomAccessStore
	pos  int64
	size int64
}

func (rw *StoreReadWriter) Write(p []byte) (int, error) {
	n, err := rw.s.WriteAt(p, rw.pos)
	rw.pos += int64(n)
	return n, err
}

func (rw *StoreReadWriter) Read(p []byte) (int, error) {
	n, err := rw.s.ReadAt(p, rw.pos)
	if err == io.EOF {
		rw.size = rw.pos + int64(n)
	}
	rw.pos += int64(n)
	return n, err
}

func (rw *StoreReadWriter) Close() error { return nil }

func (rw *StoreReadWriter) Seek(pos int64, whence int) (int64, error) {
	switch whence {
	case os.SEEK_SET:
		rw.pos = pos
	case os.SEEK_CUR:
		rw.pos += pos
	case os.SEEK_END:
		if rw.size < 0 {
			return -1, errors.New("trying to seek to EOF while the store size is not known")
		}
		rw.pos = rw.size + pos
	}
	return rw.pos, nil
}

type MemoryRandomAccessStore struct {
	buf []byte
}

func (s *MemoryRandomAccessStore) WriteAt(p []byte, offset int64) (int, error) {
	err := (error)(nil)
	o := int(offset)
	e := o + len(p)
	if e > len(s.buf) {
		if e <= cap(s.buf) {
			s.buf = s.buf[0:e]
		} else {
			newBuf := make([]byte, e, cap(s.buf)*2)
			copy(newBuf, s.buf)
			s.buf = newBuf
		}
	}
	n := e - o
	copy(s.buf[o:e], p)
	return n, err
}

func (s *MemoryRandomAccessStore) ReadAt(p []byte, offset int64) (int, error) {
	err := (error)(nil)
	o := int(offset)
	e := o + len(p)
	if e > len(s.buf) {
		e = len(s.buf)
		err = io.EOF
	}
	n := e - o
	copy(p, s.buf[o:e])
	return n, err
}

func (s *MemoryRandomAccessStore) Size() (int64, error) {
	return int64(len(s.buf)), nil
}

func (s *MemoryRandomAccessStore) Close() error { return nil }

func NewMemoryRandomAccessStore() *MemoryRandomAccessStore {
	return &MemoryRandomAccessStore{
		buf: make([]byte, 0, 16),
	}
}

type MemoryRandomAccessStoreFactory struct{}

func (ras *MemoryRandomAccessStoreFactory) RandomAccessStore() (RandomAccessStore, error) {
	return NewMemoryRandomAccessStore(), nil
}

type TempFileRandomAccessStoreFactory struct {
	Dir    string
	Prefix string
	GCChan chan *os.File
}

func (ras *TempFileRandomAccessStoreFactory) RandomAccessStore() (RandomAccessStore, error) {
	f, err := ioutil.TempFile(ras.Dir, ras.Prefix)
	if err != nil {
		return nil, err
	}
	f_ := (RandomAccessStore)(f)
	if ras.GCChan != nil {
		c := ras.GCChan
		f_ = NewCloseHook(
			f_,
			func(s io.Closer) {
				c <- s.(*os.File)
			},
		)
	}
	return NewSeekerWrapper(f_), nil
}
<commit_msg>Expose fields of StoreReadWriter<commit_after>package ioextras

import (
	"errors"
	"io"
	"io/ioutil"
	"os"
)

type RandomAccessStore interface {
	io.ReaderAt
	io.WriterAt
	io.Closer
}

type SizedRandomAccessStore interface {
	RandomAccessStore
	Size() (int64, error)
}

type NamedRandomAccessStore interface {
	RandomAccessStore
	Name() string
}

type RandomAccessStoreFactory interface {
	RandomAccessStore() (RandomAccessStore, error)
}

type SeekerWrapper struct {
	s  RandomAccessStore
	sk io.Seeker
	ns NamedRandomAccessStore
}

func (s *SeekerWrapper) ReadAt(p []byte, offset int64) (int, error) { return s.s.ReadAt(p, offset) }

func (s *SeekerWrapper) WriteAt(p []byte, offset int64) (int, error) { return s.s.WriteAt(p, offset) }

func (s *SeekerWrapper) Close() error { return s.s.Close() }

func (s *SeekerWrapper) Name() string {
	if s.ns != nil {
		return s.ns.Name()
	} else {
		return ""
	}
}

func (s *SeekerWrapper) Size() (int64, error) {
	return s.sk.Seek(0, os.SEEK_END)
}

func NewSeekerWrapper(s RandomAccessStore) *SeekerWrapper {
	ns, _ := s.(NamedRandomAccessStore)
	return &SeekerWrapper{s, s.(io.Seeker), ns}
}

type StoreReadWriter struct {
	Store    RandomAccessStore
	Position int64
	Size     int64
}

func (rw *StoreReadWriter) Write(p []byte) (int, error) {
	n, err := rw.Store.WriteAt(p, rw.Position)
	rw.Position += int64(n)
	return n, err
}

func (rw *StoreReadWriter) Read(p []byte) (int, error) {
	n, err := rw.Store.ReadAt(p, rw.Position)
	if err == io.EOF {
		rw.Size = rw.Position + int64(n)
	}
	rw.Position += int64(n)
	return n, err
}

func (rw *StoreReadWriter) Close() error { return nil }

func (rw *StoreReadWriter) Seek(pos int64, whence int) (int64, error) {
	switch whence {
	case os.SEEK_SET:
		rw.Position = pos
	case os.SEEK_CUR:
		rw.Position += pos
	case os.SEEK_END:
		if rw.Size < 0 {
			return -1, errors.New("trying to seek to EOF while the store size is not known")
		}
		rw.Position = rw.Size + pos
	}
	return rw.Position, nil
}

type MemoryRandomAccessStore struct {
	buf []byte
}

func (s *MemoryRandomAccessStore) WriteAt(p []byte, offset int64) (int, error) {
	err := (error)(nil)
	o := int(offset)
	e := o + len(p)
	if e > len(s.buf) {
		if e <= cap(s.buf) {
			s.buf = s.buf[0:e]
		} else {
			newBuf := make([]byte, e, cap(s.buf)*2)
			copy(newBuf, s.buf)
			s.buf = newBuf
		}
	}
	n := e - o
	copy(s.buf[o:e], p)
	return n, err
}

func (s *MemoryRandomAccessStore) ReadAt(p []byte, offset int64) (int, error) {
	err := (error)(nil)
	o := int(offset)
	e := o + len(p)
	if e > len(s.buf) {
		e = len(s.buf)
		err = io.EOF
	}
	n := e - o
	copy(p, s.buf[o:e])
	return n, err
}

func (s *MemoryRandomAccessStore) Size() (int64, error) {
	return int64(len(s.buf)), nil
}

func (s *MemoryRandomAccessStore) Close() error { return nil }

func NewMemoryRandomAccessStore() *MemoryRandomAccessStore {
	return &MemoryRandomAccessStore{
		buf: make([]byte, 0, 16),
	}
}

type MemoryRandomAccessStoreFactory struct{}

func (ras *MemoryRandomAccessStoreFactory) RandomAccessStore() (RandomAccessStore, error) {
	return NewMemoryRandomAccessStore(), nil
}

type TempFileRandomAccessStoreFactory struct {
	Dir    string
	Prefix string
	GCChan chan *os.File
}

func (ras *TempFileRandomAccessStoreFactory) RandomAccessStore() (RandomAccessStore, error) {
	f, err := ioutil.TempFile(ras.Dir, ras.Prefix)
	if err != nil {
		return nil, err
	}
	f_ := (RandomAccessStore)(f)
	if ras.GCChan != nil {
		c := ras.GCChan
		f_ = NewCloseHook(
			f_,
			func(s io.Closer) {
				c <- s.(*os.File)
			},
		)
	}
	return NewSeekerWrapper(f_), nil
}
<|endoftext|>


<commit_before>package internal

import (
	"io"
	"net"
	"sync"
	"time"

	v2net "v2ray.com/core/common/net"
)

// ConnectionID is the ID of a connection.
type ConnectionID struct {
	Local      v2net.Address
	Remote     v2net.Address
	RemotePort v2net.Port
}

// NewConnectionID creates a new ConnectionId.
func NewConnectionID(source v2net.Address, dest v2net.Destination) ConnectionID {
	return ConnectionID{
		Local:      source,
		Remote:     dest.Address,
		RemotePort: dest.Port,
	}
}

type Reuser struct {
	userEnabled bool
	appEnable   bool
}

func ReuseConnection(reuse bool) *Reuser {
	return &Reuser{
		userEnabled: reuse,
		appEnable:   reuse,
	}
}

// Connection is an implementation of net.Conn with re-usability.
type Connection struct {
	sync.RWMutex
	id       ConnectionID
	conn     net.Conn
	listener ConnectionRecyler
	reuser   *Reuser
}

func NewConnection(id ConnectionID, conn net.Conn, manager ConnectionRecyler, reuser *Reuser) *Connection {
	return &Connection{
		id:       id,
		conn:     conn,
		listener: manager,
		reuser:   reuser,
	}
}

func (v *Connection) Read(b []byte) (int, error) {
	conn := v.underlyingConn()
	if conn == nil {
		return 0, io.EOF
	}

	return conn.Read(b)
}

func (v *Connection) Write(b []byte) (int, error) {
	conn := v.underlyingConn()
	if conn == nil {
		return 0, io.ErrClosedPipe
	}
	return conn.Write(b)
}

// Close implements net.Conn.Close(). If the connection is reusable, the underlying connection will be recycled.
func (v *Connection) Close() error {
	if v == nil {
		return io.ErrClosedPipe
	}

	v.Lock()
	defer v.Unlock()
	if v.conn == nil {
		return io.ErrClosedPipe
	}
	if v.Reusable() {
		v.listener.Put(v.id, v.conn)
		return nil
	}
	err := v.conn.Close()
	v.conn = nil
	return err
}

func (v *Connection) LocalAddr() net.Addr {
	conn := v.underlyingConn()
	if conn == nil {
		return nil
	}
	return conn.LocalAddr()
}

func (v *Connection) RemoteAddr() net.Addr {
	conn := v.underlyingConn()
	if conn == nil {
		return nil
	}
	return conn.RemoteAddr()
}

func (v *Connection) SetDeadline(t time.Time) error {
	conn := v.underlyingConn()
	if conn == nil {
		return nil
	}
	return conn.SetDeadline(t)
}

func (v *Connection) SetReadDeadline(t time.Time) error {
	conn := v.underlyingConn()
	if conn == nil {
		return nil
	}
	return conn.SetReadDeadline(t)
}

func (v *Connection) SetWriteDeadline(t time.Time) error {
	conn := v.underlyingConn()
	if conn == nil {
		return nil
	}
	return conn.SetWriteDeadline(t)
}

func (v *Connection) SetReusable(reusable bool) {
	if v == nil {
		return
	}
	v.reuser.appEnable = reusable
}

func (v *Connection) Reusable() bool {
	if v == nil {
		return false
	}
	return v.reuser.userEnabled && v.reuser.appEnable
}

func (v *Connection) SysFd() (int, error) {
	conn := v.underlyingConn()
	if conn == nil {
		return 0, io.ErrClosedPipe
	}
	return GetSysFd(conn)
}

func (v *Connection) underlyingConn() net.Conn {
	if v == nil {
		return nil
	}

	v.RLock()
	defer v.RUnlock()

	return v.conn
}
<commit_msg>comments<commit_after>package internal

import (
	"io"
	"net"
	"sync"
	"time"

	v2net "v2ray.com/core/common/net"
)

// ConnectionID is the ID of a connection.
type ConnectionID struct {
	Local      v2net.Address
	Remote     v2net.Address
	RemotePort v2net.Port
}

// NewConnectionID creates a new ConnectionId.
func NewConnectionID(source v2net.Address, dest v2net.Destination) ConnectionID {
	return ConnectionID{
		Local:      source,
		Remote:     dest.Address,
		RemotePort: dest.Port,
	}
}

// Reuser determines whether a connection can be reused or not.
type Reuser struct {
	// userEnabled indicates connection-reuse enabled by user.
	userEnabled bool
	// appEnabled indicates connection-reuse enabled by app.
	appEnabled bool
}

// ReuseConnection returns a tracker for tracking connection reusability.
func ReuseConnection(reuse bool) *Reuser {
	return &Reuser{
		userEnabled: reuse,
		appEnabled:  reuse,
	}
}

// Connection is an implementation of net.Conn with re-usability.
type Connection struct {
	sync.RWMutex
	id       ConnectionID
	conn     net.Conn
	listener ConnectionRecyler
	reuser   *Reuser
}

// NewConnection creates a new connection.
func NewConnection(id ConnectionID, conn net.Conn, manager ConnectionRecyler, reuser *Reuser) *Connection {
	return &Connection{
		id:       id,
		conn:     conn,
		listener: manager,
		reuser:   reuser,
	}
}

// Read implements net.Conn.Read().
func (v *Connection) Read(b []byte) (int, error) {
	conn := v.underlyingConn()
	if conn == nil {
		return 0, io.EOF
	}

	return conn.Read(b)
}

// Write implement net.Conn.Write().
func (v *Connection) Write(b []byte) (int, error) {
	conn := v.underlyingConn()
	if conn == nil {
		return 0, io.ErrClosedPipe
	}
	return conn.Write(b)
}

// Close implements net.Conn.Close(). If the connection is reusable, the underlying connection will be recycled.
func (v *Connection) Close() error {
	if v == nil {
		return io.ErrClosedPipe
	}

	v.Lock()
	defer v.Unlock()
	if v.conn == nil {
		return io.ErrClosedPipe
	}
	if v.Reusable() {
		v.listener.Put(v.id, v.conn)
		return nil
	}
	err := v.conn.Close()
	v.conn = nil
	return err
}

// LocalAddr implements net.Conn.LocalAddr().
func (v *Connection) LocalAddr() net.Addr {
	conn := v.underlyingConn()
	if conn == nil {
		return nil
	}
	return conn.LocalAddr()
}

// RemoteAddr implements net.Conn.RemoteAddr().
func (v *Connection) RemoteAddr() net.Addr {
	conn := v.underlyingConn()
	if conn == nil {
		return nil
	}
	return conn.RemoteAddr()
}

// SetDeadline implements net.Conn.SetDeadline().
func (v *Connection) SetDeadline(t time.Time) error {
	conn := v.underlyingConn()
	if conn == nil {
		return nil
	}
	return conn.SetDeadline(t)
}

// SetReadDeadline implements net.Conn.SetReadDeadline().
func (v *Connection) SetReadDeadline(t time.Time) error {
	conn := v.underlyingConn()
	if conn == nil {
		return nil
	}
	return conn.SetReadDeadline(t)
}

// SetWriteDeadline implements net.Conn.SetWriteDeadline().
func (v *Connection) SetWriteDeadline(t time.Time) error {
	conn := v.underlyingConn()
	if conn == nil {
		return nil
	}
	return conn.SetWriteDeadline(t)
}

// SetReusable implements internet.Reusable.SetReusable().
func (v *Connection) SetReusable(reusable bool) {
	if v == nil {
		return
	}
	v.reuser.appEnabled = reusable
}

// Reusable implements internet.Reusable.Reusable().
func (v *Connection) Reusable() bool {
	if v == nil {
		return false
	}
	return v.reuser.userEnabled && v.reuser.appEnabled
}

// SysFd implement internet.SysFd.SysFd().
func (v *Connection) SysFd() (int, error) {
	conn := v.underlyingConn()
	if conn == nil {
		return 0, io.ErrClosedPipe
	}
	return GetSysFd(conn)
}

func (v *Connection) underlyingConn() net.Conn {
	if v == nil {
		return nil
	}

	v.RLock()
	defer v.RUnlock()

	return v.conn
}
<|endoftext|>


<commit_before>// Copyright 2013 @atotto. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build freebsd linux netbsd openbsd solaris dragonfly

package clipboard

import (
	"errors"
	"os/exec"
)

const (
	xsel  = "xsel"
	xclip = "xclip"
)

var (
	Primary bool

	pasteCmdArgs []string
	copyCmdArgs  []string

	xselPasteArgs = []string{xsel, "--output", "--clipboard"}
	xselCopyArgs  = []string{xsel, "--input", "--clipboard"}

	xclipPasteArgs = []string{xclip, "-out", "-selection", "clipboard"}
	xclipCopyArgs  = []string{xclip, "-in", "-selection", "clipboard"}

	missingCommands = errors.New("No clipboard utilities available. Please install xsel or xclip.")
)

func init() {
	pasteCmdArgs = xclipPasteArgs
	copyCmdArgs = xclipCopyArgs

	if _, err := exec.LookPath(xclip); err == nil {
		return
	}

	pasteCmdArgs = xselPasteArgs
	copyCmdArgs = xselCopyArgs

	if _, err := exec.LookPath(xsel); err == nil {
		return
	}

	Unsupported = true
}

func getPasteCommand() *exec.Cmd {
	if Primary {
		pasteCmdArgs = pasteCmdArgs[:1]
	}
	return exec.Command(pasteCmdArgs[0], pasteCmdArgs[1:]...)
}

func getCopyCommand() *exec.Cmd {
	if Primary {
		copyCmdArgs = copyCmdArgs[:1]
	}
	return exec.Command(copyCmdArgs[0], copyCmdArgs[1:]...)
}

func readAll() (string, error) {
	if Unsupported {
		return "", missingCommands
	}
	pasteCmd := getPasteCommand()
	out, err := pasteCmd.Output()
	if err != nil {
		return "", err
	}
	return string(out), nil
}

func writeAll(text string) error {
	if Unsupported {
		return missingCommands
	}
	copyCmd := getCopyCommand()
	in, err := copyCmd.StdinPipe()
	if err != nil {
		return err
	}

	if err := copyCmd.Start(); err != nil {
		return err
	}
	if _, err := in.Write([]byte(text)); err != nil {
		return err
	}
	if err := in.Close(); err != nil {
		return err
	}
	return copyCmd.Wait()
}
<commit_msg>Add Termux support<commit_after>// Copyright 2013 @atotto. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build freebsd linux netbsd openbsd solaris dragonfly

package clipboard

import (
	"errors"
	"os/exec"
)

const (
	xsel               = "xsel"
	xclip              = "xclip"
	termuxClipboardGet = "termux-clipboard-get"
	termuxClipboardSet = "termux-clipboard-set"
)

var (
	Primary bool

	pasteCmdArgs []string
	copyCmdArgs  []string

	xselPasteArgs = []string{xsel, "--output", "--clipboard"}
	xselCopyArgs  = []string{xsel, "--input", "--clipboard"}

	xclipPasteArgs = []string{xclip, "-out", "-selection", "clipboard"}
	xclipCopyArgs  = []string{xclip, "-in", "-selection", "clipboard"}

	termuxPasteArgs = []string{termuxClipboardGet}
	termuxCopyArgs  = []string{termuxClipboardSet}

	missingCommands = errors.New("No clipboard utilities available. Please install xsel, xclip, or Termux:API add-on for termux-clipboard-get/set.")
)

func init() {
	pasteCmdArgs = xclipPasteArgs
	copyCmdArgs = xclipCopyArgs

	if _, err := exec.LookPath(xclip); err == nil {
		return
	}

	pasteCmdArgs = xselPasteArgs
	copyCmdArgs = xselCopyArgs

	if _, err := exec.LookPath(xsel); err == nil {
		return
	}

	pasteCmdArgs = termuxPasteArgs
	copyCmdArgs = termuxCopyArgs

	if _, err := exec.LookPath(termuxClipboardSet); err == nil {
		if _, err := exec.LookPath(termuxClipboardGet); err == nil {
			return
		}
	}

	Unsupported = true
}

func getPasteCommand() *exec.Cmd {
	if Primary {
		pasteCmdArgs = pasteCmdArgs[:1]
	}
	return exec.Command(pasteCmdArgs[0], pasteCmdArgs[1:]...)
}

func getCopyCommand() *exec.Cmd {
	if Primary {
		copyCmdArgs = copyCmdArgs[:1]
	}
	return exec.Command(copyCmdArgs[0], copyCmdArgs[1:]...)
}

func readAll() (string, error) {
	if Unsupported {
		return "", missingCommands
	}
	pasteCmd := getPasteCommand()
	out, err := pasteCmd.Output()
	if err != nil {
		return "", err
	}
	return string(out), nil
}

func writeAll(text string) error {
	if Unsupported {
		return missingCommands
	}
	copyCmd := getCopyCommand()
	in, err := copyCmd.StdinPipe()
	if err != nil {
		return err
	}

	if err := copyCmd.Start(); err != nil {
		return err
	}
	if _, err := in.Write([]byte(text)); err != nil {
		return err
	}
	if err := in.Close(); err != nil {
		return err
	}
	return copyCmd.Wait()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"strings"
	"text/template"

	"github.com/snikch/goose/lib/goose"
)

// global options. available to any subcommands.
var flagPath = flag.String("path", "db", "folder containing db info")
var flagEnv = flag.String("env", "development", "which DB environment to use")
var flagPgSchema = flag.String("pgschema", "", "which postgres-schema to migrate (default = none)")
var flagMigrationsPath = flag.String("migrations", "", "folder containing db migrations (default = none)")

// helper to create a DBConf from the given flags
func dbConfFromFlags() (dbconf *goose.DBConf, err error) {
	conf, err := goose.NewDBConf(*flagPath, *flagEnv, *flagPgSchema)
	if err != nil {
		return nil, err
	}
	if *flagMigrationsPath != "" {
		conf.MigrationsDir = *flagMigrationsPath
	}
	return conf, nil
}

var commands = []*Command{
	upCmd,
	downCmd,
	redoCmd,
	statusCmd,
	createCmd,
	dbVersionCmd,
}

func main() {

	flag.Usage = usage
	flag.Parse()

	args := flag.Args()
	if len(args) == 0 || args[0] == "-h" {
		flag.Usage()
		return
	}

	var cmd *Command
	name := args[0]
	for _, c := range commands {
		if strings.HasPrefix(c.Name, name) {
			cmd = c
			break
		}
	}

	if cmd == nil {
		fmt.Printf("error: unknown command %q\n", name)
		flag.Usage()
		os.Exit(1)
	}

	cmd.Exec(args[1:])
}

func usage() {
	fmt.Print(usagePrefix)
	flag.PrintDefaults()
	usageTmpl.Execute(os.Stdout, commands)
}

var usagePrefix = `
goose is a database migration management system for Go projects.

Usage:
    goose [options] <subcommand> [subcommand options]

Options:
`
var usageTmpl = template.Must(template.New("usage").Parse(
	`
Commands:{{range .}}
    {{.Name | printf "%-10s"}} {{.Summary}}{{end}}
`))
<commit_msg>Make obvious what the default migrations path is<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"strings"
	"text/template"

	"github.com/snikch/goose/lib/goose"
)

// global options. available to any subcommands.
var flagPath = flag.String("path", "db", "folder containing db info")
var flagEnv = flag.String("env", "development", "which DB environment to use")
var flagPgSchema = flag.String("pgschema", "", "which postgres-schema to migrate (default = none)")
var flagMigrationsPath = flag.String("migrations", "", "folder containing db migrations (default = [path]/migrations)")

// helper to create a DBConf from the given flags
func dbConfFromFlags() (dbconf *goose.DBConf, err error) {
	conf, err := goose.NewDBConf(*flagPath, *flagEnv, *flagPgSchema)
	if err != nil {
		return nil, err
	}
	if *flagMigrationsPath != "" {
		conf.MigrationsDir = *flagMigrationsPath
	}
	return conf, nil
}

var commands = []*Command{
	upCmd,
	downCmd,
	redoCmd,
	statusCmd,
	createCmd,
	dbVersionCmd,
}

func main() {

	flag.Usage = usage
	flag.Parse()

	args := flag.Args()
	if len(args) == 0 || args[0] == "-h" {
		flag.Usage()
		return
	}

	var cmd *Command
	name := args[0]
	for _, c := range commands {
		if strings.HasPrefix(c.Name, name) {
			cmd = c
			break
		}
	}

	if cmd == nil {
		fmt.Printf("error: unknown command %q\n", name)
		flag.Usage()
		os.Exit(1)
	}

	cmd.Exec(args[1:])
}

func usage() {
	fmt.Print(usagePrefix)
	flag.PrintDefaults()
	usageTmpl.Execute(os.Stdout, commands)
}

var usagePrefix = `
goose is a database migration management system for Go projects.

Usage:
    goose [options] <subcommand> [subcommand options]

Options:
`
var usageTmpl = template.Must(template.New("usage").Parse(
	`
Commands:{{range .}}
    {{.Name | printf "%-10s"}} {{.Summary}}{{end}}
`))
<|endoftext|>


<commit_before>package cmd

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"time"

	"github.com/spf13/cobra"
)

// holds value of "name" flag
var jobCreateName string

// holds value of "aws-account-id" flag
var jobCreateAwsAccountID string

// holds value of "rule-type" flag
var jobCreateRuleType string

// holds value of "rule-value" flag
var jobCreateRuleValue string

// holds value of "action-type" flag
var jobCreateActionType string

// holds value of "action-value" flag
var jobCreateActionValue string

var jobCreateCmd = &cobra.Command{
	Use:   "create",
	Short: "Register new CA job",
	Long:  `Use the "/jobs" API to register new job with Cloud Automator.`,
	RunE:  execJobCreate,
}

func init() {
	jobCmd.AddCommand(jobCreateCmd)

	jobCreateCmd.Flags().StringVar(&jobCreateName, "name", "", "Specify CA job name. [required]")
	jobCreateCmd.Flags().StringVar(&jobCreateAwsAccountID, "aws-account-id", "", "Specify CA AWS account ID. [required]")
	jobCreateCmd.Flags().StringVar(&jobCreateRuleType, "rule-type", "", "Specify CA job trigger type. [required]")
	jobCreateCmd.Flags().StringVar(&jobCreateRuleValue, "rule-value", "", "Specify CA job trigger setting value. [required only \"--rule-type=cron/sqs\"]")
	jobCreateCmd.Flags().StringVar(&jobCreateActionType, "action-type", "", "Specify CA job action type. [required]")
	jobCreateCmd.Flags().StringVar(&jobCreateActionValue, "action-value", "", "Specify CA job action setting value. [required]")
}

func execJobCreate(cmd *cobra.Command, args []string) error {
	postBody, err := createJobCreatePostBody()
	if err != nil {
		cmd.Help()
		return err
	}

	// get profile data
	paramProfileName := rootPprofileName
	currentProfile, err := getProfile(paramProfileName)
	if err != nil {
		return err
	}

	// preparing to call API
	endURL := currentProfile.Endpoint
	apiPath := "jobs"
	method := "POST"

	// create http client
	httpClient := &http.Client{}
	client, err := createClient(endURL, httpClient)
	if err != nil {
		return err
	}

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	// call API [/jobs]
	httpRequest, err := client.createRequest(ctx, method, apiPath, currentProfile.APIKey, postBody)
	if err != nil {
		return err
	}

	httpResponse, err := client.HTTPClient.Do(httpRequest)
	if err != nil {
		return err
	}

	// get result string
	result, err := getBodyString(httpResponse)
	if err != nil {
		return err
	}

	// draw result
	fmt.Printf("%s\n", result)

	return nil
}

func createJobCreatePostBody() (io.Reader, error) {
	params := map[string]interface{}{}

	paramName := jobCreateName
	if paramName == "" {
		return nil, errors.New("\nPlease specify flag [--name] required CA job name.")
	}
	params["name"] = paramName

	paramAwsAccountID := jobCreateAwsAccountID
	if paramAwsAccountID == "" {
		return nil, errors.New("\nPlease specify flag [--aws-account-id] required CA AWS account ID.")
	}
	params["aws_account_id"] = paramAwsAccountID

	paramRuleType := jobCreateRuleType
	if paramRuleType == "" {
		return nil, errors.New("\nPlease specify flag [--rule-type] required CA job trigger type.")
	}
	valid, required := isValidJobRuleType(paramRuleType)
	if !valid {
		return nil, errors.New("\nPlease specify flag [--rule-type] valid CA job trigger type.")
	}
	params["rule_type"] = paramRuleType

	paramRuleValue := jobCreateRuleValue
	if required {
		if paramRuleValue == "" {
			return nil, errors.New("\nPlease specify flag [--rule-value] required CA job trigger setting value.")
		}
		parsedParam, err := createJobParseObjectParameter(paramRuleValue, "\nPlease specify flag [--rule-value] required CA job trigger setting value.")
		if err != nil {
			return nil, err
		}
		params["rule_value"] = *parsedParam
	} else {
		params["rule_value"] = paramRuleValue
	}

	paramActionType := jobCreateActionType
	if paramActionType == "" {
		return nil, errors.New("\nPlease specify flag [--action-type] required CA job action type.")
	}
	params["action_type"] = paramActionType

	paramActionValue := jobCreateActionValue
	if paramActionValue == "" {
		return nil, errors.New("\nPlease specify flag [--action-value] required CA job action setting value.")
	}
	parsedParam, err := createJobParseObjectParameter(paramActionValue, "\nPlease specify flag [--action-value] required CA job action setting value.")
	if err != nil {
		return nil, err
	}
	params["action_value"] = parsedParam

	paramBytes, err := json.Marshal(params)
	if err != nil {
		return nil, err
	}

	return bytes.NewBuffer(paramBytes), nil
}
<commit_msg>add post process flag<commit_after>package cmd

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/spf13/cobra"
)

// holds value of "name" flag
var jobCreateName string

// holds value of "aws-account-id" flag
var jobCreateAwsAccountID string

// holds value of "rule-type" flag
var jobCreateRuleType string

// holds value of "rule-value" flag
var jobCreateRuleValue string

// holds value of "action-type" flag
var jobCreateActionType string

// holds value of "action-value" flag
var jobCreateActionValue string

// holds value of "completed-post-process-id" flag
var jobCreateCompletedPostProcessID string

// holds value of "failed-post-process-id" flag
var jobCreateFailedPostProcessID string

var jobCreateCmd = &cobra.Command{
	Use:   "create",
	Short: "Register new CA job",
	Long:  `Use the "/jobs" API to register new job with Cloud Automator.`,
	RunE:  execJobCreate,
}

func init() {
	jobCmd.AddCommand(jobCreateCmd)

	jobCreateCmd.Flags().StringVar(&jobCreateName, "name", "", "Specify CA job name. [required]")
	jobCreateCmd.Flags().StringVar(&jobCreateAwsAccountID, "aws-account-id", "", "Specify CA AWS account ID. [required]")
	jobCreateCmd.Flags().StringVar(&jobCreateRuleType, "rule-type", "", "Specify CA job trigger type. [required]")
	jobCreateCmd.Flags().StringVar(&jobCreateRuleValue, "rule-value", "", "Specify CA job trigger setting values. [required only \"--rule-type=cron/sqs\"]")
	jobCreateCmd.Flags().StringVar(&jobCreateActionType, "action-type", "", "Specify CA job action type. [required]")
	jobCreateCmd.Flags().StringVar(&jobCreateActionValue, "action-value", "", "Specify CA job action setting values. [required]")
	jobCreateCmd.Flags().StringVar(&jobCreateCompletedPostProcessID, "completed-post-process-id", "", "Specify array that contains post-processing IDs to be executed if the CA job succeeds.")
	jobCreateCmd.Flags().StringVar(&jobCreateFailedPostProcessID, "failed-post-process-id", "", "Specify array that contains post-processing IDs to be executed if the CA job faileds.")
}

func execJobCreate(cmd *cobra.Command, args []string) error {
	postBody, err := createJobCreatePostBody()
	if err != nil {
		cmd.Help()
		return err
	}

	// get profile data
	paramProfileName := rootPprofileName
	currentProfile, err := getProfile(paramProfileName)
	if err != nil {
		return err
	}

	// preparing to call API
	endURL := currentProfile.Endpoint
	apiPath := "jobs"
	method := "POST"

	// create http client
	httpClient := &http.Client{}
	client, err := createClient(endURL, httpClient)
	if err != nil {
		return err
	}

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	// call API [/jobs]
	httpRequest, err := client.createRequest(ctx, method, apiPath, currentProfile.APIKey, postBody)
	if err != nil {
		return err
	}

	httpResponse, err := client.HTTPClient.Do(httpRequest)
	if err != nil {
		return err
	}

	// get result string
	result, err := getBodyString(httpResponse)
	if err != nil {
		return err
	}

	// draw result
	fmt.Printf("%s\n", result)

	return nil
}

func createJobCreatePostBody() (io.Reader, error) {
	params := map[string]interface{}{}

	paramName := jobCreateName
	if paramName == "" {
		return nil, errors.New("\nPlease specify flag [--name] required CA job name.")
	}
	params["name"] = paramName

	paramAwsAccountID := jobCreateAwsAccountID
	if paramAwsAccountID == "" {
		return nil, errors.New("\nPlease specify flag [--aws-account-id] required CA AWS account ID.")
	}
	params["aws_account_id"] = paramAwsAccountID

	paramRuleType := jobCreateRuleType
	if paramRuleType == "" {
		return nil, errors.New("\nPlease specify flag [--rule-type] required CA job trigger type.")
	}
	valid, required := isValidJobRuleType(paramRuleType)
	if !valid {
		return nil, errors.New("\nPlease specify flag [--rule-type] valid CA job trigger type.")
	}
	params["rule_type"] = paramRuleType

	paramRuleValue := jobCreateRuleValue
	if required {
		if paramRuleValue == "" {
			return nil, errors.New("\nPlease specify flag [--rule-value] required CA job trigger setting values.")
		}
		parsedParam, err := createJobParseObjectParameter(paramRuleValue, "\nPlease specify flag [--rule-value] required CA job trigger setting values.")
		if err != nil {
			return nil, err
		}
		params["rule_value"] = *parsedParam
	} else {
		params["rule_value"] = paramRuleValue
	}

	paramActionType := jobCreateActionType
	if paramActionType == "" {
		return nil, errors.New("\nPlease specify flag [--action-type] required CA job action type.")
	}
	params["action_type"] = paramActionType

	paramActionValue := jobCreateActionValue
	if paramActionValue == "" {
		return nil, errors.New("\nPlease specify flag [--action-value] required CA job action setting values.")
	}
	parsedParam, err := createJobParseObjectParameter(paramActionValue, "\nPlease specify flag [--action-value] required CA job action setting values.")
	if err != nil {
		return nil, err
	}
	params["action_value"] = *parsedParam

	// optional
	paramCompletedPostProcessID := jobCreateCompletedPostProcessID
	if paramCompletedPostProcessID != "" {
		params["completed_post_process_id"] = strings.Split(paramCompletedPostProcessID, ",")
	}

	// optional
	paramFailedPostProcessID := jobCreateFailedPostProcessID
	if paramFailedPostProcessID != "" {
		params["failed_post_process_id"] = strings.Split(paramFailedPostProcessID, ",")
	}

	paramBytes, err := json.Marshal(params)
	if err != nil {
		return nil, err
	}

	return bytes.NewBuffer(paramBytes), 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.

// proxy is an HTTP/S proxy configurable via an HTTP API.
//
// It can be dynamically configured/queried at runtime by issuing requests to
// proxy specific paths using JSON.
//
// Supported configuration endpoints:
//
//   POST http://martian.proxy/configure
//
// sets the request and response modifier of the proxy; modifiers adhere to the
// following top-level JSON structure:
//
//   {
//     "package.Modifier": {
//       "scope": ["request", "response"],
//       "attribute 1": "value",
//       "attribute 2": "value"
//     }
//   }
//
// modifiers may be "stacked" to provide support for additional behaviors; for
// example, to add a "Martian-Test" header with the value "true" for requests
// with the domain "www.example.com" the JSON message would be:
//
//   {
//     "url.Filter": {
//       "scope": ["request"],
//       "host": "www.example.com",
//       "modifier": {
//         "header.Modifier": {
//           "name": "Martian-Test",
//           "value": "true"
//         }
//       }
//     }
//   }
//
// url.Filter parses the JSON object in the value of the "url.Filter" attribute;
// the "host" key tells the url.Filter to filter requests if the host explicitly
// matches "www.example.com"
//
// the "modifier" key within the "url.Filter" JSON object contains another
// modifier message of the type header.Modifier to run iff the filter passes
//
// groups may also be used to run multiple modifiers sequentially; for example to
// log requests and responses after adding the "Martian-Test" header to the
// request, but only when the host matches www.example.com:
//
//   {
//     "url.Filter": {
//       "host": "www.example.com",
//       "modifier": {
//         "fifo.Group": {
//           "modifiers": [
//             {
//               "header.Modifier": {
//                 "scope": ["request"],
//                 "name": "Martian-Test",
//                 "value": "true"
//               }
//             },
//             {
//               "log.Logger": { }
//             }
//           ]
//         }
//       }
//     }
//   }
//
// modifiers are designed to be composed together in ways that allow the user to
// write a single JSON structure to accomplish a variety of functionality
//
//   GET http://martian.proxy/verify
//
// retrieves the verifications errors as JSON with the following structure:
//
//   {
//     "errors": [
//       {
//         "message": "request(url) verification failure"
//       },
//       {
//         "message": "response(url) verification failure"
//       }
//     ]
//   }
//
// verifiers also adhere to the modifier interface and thus can be included in the
// modifier configuration request; for example, to verify that all requests to
// "www.example.com" are sent over HTTPS send the following JSON to the
// configuration endpoint:
//
//   {
//     "url.Filter": {
//       "scope": ["request"],
//       "host": "www.example.com",
//       "modifier": {
//         "url.Verifier": {
//           "scope": ["request"],
//           "scheme": "https"
//         }
//       }
//     }
//   }
//
// sending a request to "http://martian.proxy/verify" will then return errors from the url.Verifier
//
//   POST http://martian.proxy/verify/reset
//
// resets the verifiers to their initial state; note some verifiers may start in
// a failure state (e.g., pingback.Verifier is failed if no requests have been
// seen by the proxy)
//
//   GET http://martian.proxy/authority.cer
//
// prompts the user to install the CA certificate used by the proxy if MITM is enabled
//
//   GET http://martian.proxy/har/logs
//
// retrieves the HAR logs for all requests and responses seen by the proxy if
// the HAR flag is enabled
//
//   DELETE http://martian.proxy/har/logs/reset
//
// reset the in-memory HAR log; note that the log will grow unbounded unless it
// is periodically reset
//
// passing the -cors flag will enable CORS support for the endpoints so that they
// may be called via AJAX
//
// The flags are:
//   -addr=":8080"
//     host:port of the proxy
//   -api="martian.proxy"
//     hostname that can be used to reference the configuration API when
//     configuring through the proxy
//   -cert=""
//     PEM encoded X.509 CA certificate; if set, it will be set as the
//     issuer for dynamically-generated certificates during man-in-the-middle
//   -key=""
//     PEM encoded private key of cert (RSA or ECDSA); if set, the key will be used
//     to sign dynamically-generated certificates during man-in-the-middle
//   -generate-ca-cert=false
//     generates a CA certificate and private key to use for man-in-the-middle;
//     the certificate is only valid while the proxy is running and will be
//     discarded on shutdown
//   -organization="Martian Proxy"
//     organization name set on the dynamically-generated certificates during
//     man-in-the-middle
//   -validity="1h"
//     window of time around the time of request that the dynamically-generated
//     certificate is valid for; the duration is set such that the total valid
//     timeframe is double the value of validity (1h before & 1h after)
//   -cors=false
//     allow the proxy to be configured via CORS requests; such as when
//     configuring the proxy via AJAX
//   -har=false
//     enable logging endpoints for retrieving full request/response logs in
//     HAR format.
package main

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"time"

	"github.com/google/martian"
	"github.com/google/martian/cors"
	"github.com/google/martian/fifo"
	"github.com/google/martian/har"
	"github.com/google/martian/header"
	"github.com/google/martian/martianhttp"
	"github.com/google/martian/mitm"
	"github.com/google/martian/verify"

	_ "github.com/google/martian/body"
	_ "github.com/google/martian/cookie"
	_ "github.com/google/martian/martianlog"
	_ "github.com/google/martian/martianurl"
	_ "github.com/google/martian/method"
	_ "github.com/google/martian/pingback"
	_ "github.com/google/martian/priority"
	_ "github.com/google/martian/querystring"
	_ "github.com/google/martian/status"
)

var (
	addr         = flag.String("addr", ":8080", "host:port of the proxy")
	api          = flag.String("api", "martian.proxy", "hostname for the API")
	generateCA   = flag.Bool("generate-ca-cert", false, "generate CA certificate and private key for MITM")
	cert         = flag.String("cert", "", "CA certificate used to sign MITM certificates")
	key          = flag.String("key", "", "private key of the CA used to sign MITM certificates")
	organization = flag.String("organization", "Martian Proxy", "organization name for MITM certificates")
	validity     = flag.Duration("validity", time.Hour, "window of time that MITM certificates are valid")
	allowCORS    = flag.Bool("cors", false, "allow CORS requests to configure the proxy")
	harLogging   = flag.Bool("har", false, "enable HAR logging API")
)

func main() {
	flag.Parse()

	p := martian.NewProxy()

	var x509c *x509.Certificate
	var priv interface{}

	if *generateCA {
		var err error
		x509c, priv, err = mitm.NewAuthority("martian.proxy", "Martian Authority", 30*24*time.Hour)
		if err != nil {
			log.Fatal(err)
		}
	} else if *cert != "" && *key != "" {
		tlsc, err := tls.LoadX509KeyPair(*cert, *key)
		if err != nil {
			log.Fatal(err)
		}
		priv = tlsc.PrivateKey

		x509c, err = x509.ParseCertificate(tlsc.Certificate[0])
		if err != nil {
			log.Fatal(err)
		}
	}

	if x509c != nil && priv != nil {
		mc, err := mitm.NewConfig(x509c, priv)
		if err != nil {
			log.Fatal(err)
		}

		mc.SetValidity(*validity)
		mc.SetOrganization(*organization)

		p.SetMITM(mc)

		// Expose certificate authority.
		ah := martianhttp.NewAuthorityHandler(x509c)
		configure("/authority.cer", ah)
	}

	fg := fifo.NewGroup()

	hbhm := header.NewHopByHopModifier()
	fg.AddRequestModifier(hbhm)
	fg.AddRequestModifier(header.NewForwardedModifier())
	fg.AddRequestModifier(header.NewBadFramingModifier())

	vm := header.NewViaModifier("martian")
	fg.AddRequestModifier(vm)

	m := martianhttp.NewModifier()
	fg.AddRequestModifier(m)
	fg.AddResponseModifier(m)

	fg.AddResponseModifier(hbhm)
	fg.AddResponseModifier(vm)

	if *harLogging {
		hl := har.NewLogger("martian", "2.0.0")
		fg.AddRequestModifier(hl)
		fg.AddResponseModifier(hl)

		configure("/har/logs", har.NewExportHandler(hl))
		configure("/har/logs/reset", har.NewResetHandler(hl))
	}

	p.SetRequestModifier(fg)
	p.SetResponseModifier(fg)

	// Proxy specific handlers.
	// These handlers take precendence over proxy traffic and will not be
	// intercepted.

	// Configure modifiers.
	configure("/configure", m)

	// Verify assertions.
	vh := verify.NewHandler()
	vh.SetRequestVerifier(m)
	vh.SetResponseVerifier(m)
	configure("/verify", vh)

	// Reset verifications.
	rh := verify.NewResetHandler()
	rh.SetRequestVerifier(m)
	rh.SetResponseVerifier(m)
	configure("/verify/reset", rh)

	l, err := net.Listen("tcp", *addr)
	if err != nil {
		log.Fatal(err)
	}

	log.Println("martian: proxy started on:", l.Addr())

	go p.Serve(l)

	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc, os.Interrupt, os.Kill)

	<-sigc

	log.Println("martian: shutting down")
}

// configure installs a configuration handler at path.
func configure(path string, handler http.Handler) {
	if *allowCORS {
		handler = cors.NewHandler(handler)
	}

	http.Handle(filepath.Join(*api, path), handler)
}
<commit_msg>Respond with 404 to any unhandled internal request.<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.

// proxy is an HTTP/S proxy configurable via an HTTP API.
//
// It can be dynamically configured/queried at runtime by issuing requests to
// proxy specific paths using JSON.
//
// Supported configuration endpoints:
//
//   POST http://martian.proxy/configure
//
// sets the request and response modifier of the proxy; modifiers adhere to the
// following top-level JSON structure:
//
//   {
//     "package.Modifier": {
//       "scope": ["request", "response"],
//       "attribute 1": "value",
//       "attribute 2": "value"
//     }
//   }
//
// modifiers may be "stacked" to provide support for additional behaviors; for
// example, to add a "Martian-Test" header with the value "true" for requests
// with the domain "www.example.com" the JSON message would be:
//
//   {
//     "url.Filter": {
//       "scope": ["request"],
//       "host": "www.example.com",
//       "modifier": {
//         "header.Modifier": {
//           "name": "Martian-Test",
//           "value": "true"
//         }
//       }
//     }
//   }
//
// url.Filter parses the JSON object in the value of the "url.Filter" attribute;
// the "host" key tells the url.Filter to filter requests if the host explicitly
// matches "www.example.com"
//
// the "modifier" key within the "url.Filter" JSON object contains another
// modifier message of the type header.Modifier to run iff the filter passes
//
// groups may also be used to run multiple modifiers sequentially; for example to
// log requests and responses after adding the "Martian-Test" header to the
// request, but only when the host matches www.example.com:
//
//   {
//     "url.Filter": {
//       "host": "www.example.com",
//       "modifier": {
//         "fifo.Group": {
//           "modifiers": [
//             {
//               "header.Modifier": {
//                 "scope": ["request"],
//                 "name": "Martian-Test",
//                 "value": "true"
//               }
//             },
//             {
//               "log.Logger": { }
//             }
//           ]
//         }
//       }
//     }
//   }
//
// modifiers are designed to be composed together in ways that allow the user to
// write a single JSON structure to accomplish a variety of functionality
//
//   GET http://martian.proxy/verify
//
// retrieves the verifications errors as JSON with the following structure:
//
//   {
//     "errors": [
//       {
//         "message": "request(url) verification failure"
//       },
//       {
//         "message": "response(url) verification failure"
//       }
//     ]
//   }
//
// verifiers also adhere to the modifier interface and thus can be included in the
// modifier configuration request; for example, to verify that all requests to
// "www.example.com" are sent over HTTPS send the following JSON to the
// configuration endpoint:
//
//   {
//     "url.Filter": {
//       "scope": ["request"],
//       "host": "www.example.com",
//       "modifier": {
//         "url.Verifier": {
//           "scope": ["request"],
//           "scheme": "https"
//         }
//       }
//     }
//   }
//
// sending a request to "http://martian.proxy/verify" will then return errors from the url.Verifier
//
//   POST http://martian.proxy/verify/reset
//
// resets the verifiers to their initial state; note some verifiers may start in
// a failure state (e.g., pingback.Verifier is failed if no requests have been
// seen by the proxy)
//
//   GET http://martian.proxy/authority.cer
//
// prompts the user to install the CA certificate used by the proxy if MITM is enabled
//
//   GET http://martian.proxy/har/logs
//
// retrieves the HAR logs for all requests and responses seen by the proxy if
// the HAR flag is enabled
//
//   DELETE http://martian.proxy/har/logs/reset
//
// reset the in-memory HAR log; note that the log will grow unbounded unless it
// is periodically reset
//
// passing the -cors flag will enable CORS support for the endpoints so that they
// may be called via AJAX
//
// The flags are:
//   -addr=":8080"
//     host:port of the proxy
//   -api="martian.proxy"
//     hostname that can be used to reference the configuration API when
//     configuring through the proxy
//   -cert=""
//     PEM encoded X.509 CA certificate; if set, it will be set as the
//     issuer for dynamically-generated certificates during man-in-the-middle
//   -key=""
//     PEM encoded private key of cert (RSA or ECDSA); if set, the key will be used
//     to sign dynamically-generated certificates during man-in-the-middle
//   -generate-ca-cert=false
//     generates a CA certificate and private key to use for man-in-the-middle;
//     the certificate is only valid while the proxy is running and will be
//     discarded on shutdown
//   -organization="Martian Proxy"
//     organization name set on the dynamically-generated certificates during
//     man-in-the-middle
//   -validity="1h"
//     window of time around the time of request that the dynamically-generated
//     certificate is valid for; the duration is set such that the total valid
//     timeframe is double the value of validity (1h before & 1h after)
//   -cors=false
//     allow the proxy to be configured via CORS requests; such as when
//     configuring the proxy via AJAX
//   -har=false
//     enable logging endpoints for retrieving full request/response logs in
//     HAR format.
package main

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"time"

	"github.com/google/martian"
	"github.com/google/martian/cors"
	"github.com/google/martian/fifo"
	"github.com/google/martian/har"
	"github.com/google/martian/header"
	"github.com/google/martian/martianhttp"
	"github.com/google/martian/mitm"
	"github.com/google/martian/verify"

	_ "github.com/google/martian/body"
	_ "github.com/google/martian/cookie"
	_ "github.com/google/martian/martianlog"
	_ "github.com/google/martian/martianurl"
	_ "github.com/google/martian/method"
	_ "github.com/google/martian/pingback"
	_ "github.com/google/martian/priority"
	_ "github.com/google/martian/querystring"
	_ "github.com/google/martian/status"
)

var (
	addr         = flag.String("addr", ":8080", "host:port of the proxy")
	api          = flag.String("api", "martian.proxy", "hostname for the API")
	generateCA   = flag.Bool("generate-ca-cert", false, "generate CA certificate and private key for MITM")
	cert         = flag.String("cert", "", "CA certificate used to sign MITM certificates")
	key          = flag.String("key", "", "private key of the CA used to sign MITM certificates")
	organization = flag.String("organization", "Martian Proxy", "organization name for MITM certificates")
	validity     = flag.Duration("validity", time.Hour, "window of time that MITM certificates are valid")
	allowCORS    = flag.Bool("cors", false, "allow CORS requests to configure the proxy")
	harLogging   = flag.Bool("har", false, "enable HAR logging API")
)

func main() {
	flag.Parse()

	p := martian.NewProxy()

	// Respond with 404 to any unknown proxy path.
	http.Handle(*api+"/", http.HandlerFunc(http.NotFound))

	var x509c *x509.Certificate
	var priv interface{}

	if *generateCA {
		var err error
		x509c, priv, err = mitm.NewAuthority("martian.proxy", "Martian Authority", 30*24*time.Hour)
		if err != nil {
			log.Fatal(err)
		}
	} else if *cert != "" && *key != "" {
		tlsc, err := tls.LoadX509KeyPair(*cert, *key)
		if err != nil {
			log.Fatal(err)
		}
		priv = tlsc.PrivateKey

		x509c, err = x509.ParseCertificate(tlsc.Certificate[0])
		if err != nil {
			log.Fatal(err)
		}
	}

	if x509c != nil && priv != nil {
		mc, err := mitm.NewConfig(x509c, priv)
		if err != nil {
			log.Fatal(err)
		}

		mc.SetValidity(*validity)
		mc.SetOrganization(*organization)

		p.SetMITM(mc)

		// Expose certificate authority.
		ah := martianhttp.NewAuthorityHandler(x509c)
		configure("/authority.cer", ah)
	}

	fg := fifo.NewGroup()

	hbhm := header.NewHopByHopModifier()
	fg.AddRequestModifier(hbhm)
	fg.AddRequestModifier(header.NewForwardedModifier())
	fg.AddRequestModifier(header.NewBadFramingModifier())

	vm := header.NewViaModifier("martian")
	fg.AddRequestModifier(vm)

	m := martianhttp.NewModifier()
	fg.AddRequestModifier(m)
	fg.AddResponseModifier(m)

	fg.AddResponseModifier(hbhm)
	fg.AddResponseModifier(vm)

	if *harLogging {
		hl := har.NewLogger("martian", "2.0.0")
		fg.AddRequestModifier(hl)
		fg.AddResponseModifier(hl)

		configure("/har/logs", har.NewExportHandler(hl))
		configure("/har/logs/reset", har.NewResetHandler(hl))
	}

	p.SetRequestModifier(fg)
	p.SetResponseModifier(fg)

	// Proxy specific handlers.
	// These handlers take precendence over proxy traffic and will not be
	// intercepted.

	// Configure modifiers.
	configure("/configure", m)

	// Verify assertions.
	vh := verify.NewHandler()
	vh.SetRequestVerifier(m)
	vh.SetResponseVerifier(m)
	configure("/verify", vh)

	// Reset verifications.
	rh := verify.NewResetHandler()
	rh.SetRequestVerifier(m)
	rh.SetResponseVerifier(m)
	configure("/verify/reset", rh)

	l, err := net.Listen("tcp", *addr)
	if err != nil {
		log.Fatal(err)
	}

	log.Println("martian: proxy started on:", l.Addr())

	go p.Serve(l)

	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc, os.Interrupt, os.Kill)

	<-sigc

	log.Println("martian: shutting down")
}

// configure installs a configuration handler at path.
func configure(path string, handler http.Handler) {
	if *allowCORS {
		handler = cors.NewHandler(handler)
	}

	http.Handle(filepath.Join(*api, path), handler)
}
<|endoftext|>


<commit_before>package incoming

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"time"
)

type Request struct {
	Text         string       `json:"text"`
	Notification string       `json:"notification"`
	Markdown     bool         `json:"markdown"`
	Channel      string       `json:"channel"`
	User         string       `json:"user"`
	Attachments  []Attachment `json:"attachments"`
}

type Response struct {
	Code   int         `json:"code"`
	Error  string      `json:"error"`
	Result interface{} `json:"result"`
}

func (p *Response) Err() error {
	if p.Code != 0 {
		return fmt.Errorf("code: %d, %s", p.Code, p.Error)
	}
	return nil
}

type Image struct {
	URL string `json:"url"`
}

type Attachment struct {
	Title  string  `json:"title"`
	Text   string  `json:"text"`
	Color  string  `json:"color"`
	Images []Image `json:"images"`
}

type Client struct {
	client *http.Client
}

type ClientOption func(*Client)

func NewClient(opts ...ClientOption) *Client {
	cli := &Client{
		client: &http.Client{},
	}

	cli.Options(opts...)

	return cli
}

func (p *Client) Send(robotId string, token string, req *Request) (resp *Response, err error) {

	if req == nil {
		err = errors.New("empty request")
		return
	}

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

	url := fmt.Sprintf("https://hook.bearychat.com/%s/incoming/%s", robotId, token)

	httpResp, err := http.DefaultClient.Post(url, "application/json", bytes.NewBuffer(body))
	if err != nil {
		return
	}

	defer httpResp.Body.Close()

	r := Response{}

	decoder := json.NewDecoder(httpResp.Body)
	decoder.UseNumber()

	err = decoder.Decode(&r)

	if err != nil {
		return
	}

	resp = &r

	return

}

func (p *Client) Options(opts ...ClientOption) {
	for i := 0; i < len(opts); i++ {
		opts[i](p)
	}
}

func TransportOption(transport *http.Transport) ClientOption {
	return func(c *Client) {
		c.client.Transport = transport
	}
}

func TimeoutOption(timeout time.Duration) ClientOption {
	return func(c *Client) {
		c.client.Timeout = timeout
	}
}
<commit_msg>improve http client impl<commit_after>package incoming

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"time"
)

type Request struct {
	Text         string       `json:"text"`
	Notification string       `json:"notification"`
	Markdown     bool         `json:"markdown"`
	Channel      string       `json:"channel"`
	User         string       `json:"user"`
	Attachments  []Attachment `json:"attachments"`
}

type Response struct {
	Code   int         `json:"code"`
	Error  string      `json:"error"`
	Result interface{} `json:"result"`
}

func (p *Response) Err() error {
	if p.Code != 0 {
		return fmt.Errorf("code: %d, %s", p.Code, p.Error)
	}
	return nil
}

type Image struct {
	URL string `json:"url"`
}

type Attachment struct {
	Title  string  `json:"title"`
	Text   string  `json:"text"`
	Color  string  `json:"color"`
	Images []Image `json:"images"`
}

type Client struct {
	client *http.Client
}

type ClientOption func(*Client)

func NewClient(opts ...ClientOption) *Client {
	cli := &Client{
		client: &http.Client{},
	}

	cli.Options(opts...)

	return cli
}

func (p *Client) Send(robotId string, token string, req *Request) (resp *Response, err error) {

	if len(robotId) == 0 {
		err = errors.New("robot id is empty")
		return
	}

	if len(token) == 0 {
		err = errors.New("token is empty")
		return
	}

	if req == nil {
		err = errors.New("empty request")
		return
	}

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

	url := fmt.Sprintf("https://hook.bearychat.com/%s/incoming/%s", robotId, token)

	httpResp, err := p.client.Post(url, "application/json", bytes.NewBuffer(body))
	if err != nil {
		return
	}

	defer httpResp.Body.Close()

	r := Response{}

	decoder := json.NewDecoder(httpResp.Body)
	decoder.UseNumber()

	err = decoder.Decode(&r)

	if err != nil {
		return
	}

	resp = &r

	return

}

func (p *Client) Options(opts ...ClientOption) {
	for i := 0; i < len(opts); i++ {
		opts[i](p)
	}
}

func TransportOption(transport *http.Transport) ClientOption {
	return func(c *Client) {
		c.client.Transport = transport
	}
}

func TimeoutOption(timeout time.Duration) ClientOption {
	return func(c *Client) {
		c.client.Timeout = timeout
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	"gopkg.in/redis.v5"
)

// GCOptions are provided by the caller of RunGarbageCollectKeys.
type GCOptions struct {
	RedisMasterFile string // Path to the redis master file.
	GcThreshold     int    // Number of seconds after which a key should be considered collectible.
	GcDatabases     string // List of databases to scan.
	GcKeyFile       string // Name of file containing keys to collect.
}

// GCState holds options and collector state, most importantly the current redis
// connection and database.
type GCState struct {
	opts          GCOptions
	currentMaster string
	currentDB     int
	redis         *redis.Client // current connection
	keySuffixes   []string
}

func (s *GCState) key(msgId, suffix string) string {
	return fmt.Sprintf("%s:%s", msgId, suffix)
}

func (s *GCState) keys(msgId string) []string {
	res := make([]string, 0)
	for _, suffix := range s.keySuffixes {
		res = append(res, s.key(msgId, suffix))
	}
	return res
}

func (s *GCState) msgId(key string) string {
	re := regexp.MustCompile("^(msgid:[^:]+:[-0-9a-f]+):.*$")
	matches := re.FindStringSubmatch(key)
	if len(matches) == 0 {
		panic(fmt.Errorf("msgid could not be extracted from key '%s'", key))
	}
	return matches[1]
}

func (s *GCState) gcKey(key string, threshold uint64) (bool, error) {
	v, err := s.redis.Get(key).Result()
	if err != nil {
		if err == redis.Nil {
			logDebug("key not found: %s", key)
			return false, nil
		}
		return false, err
	}
	expires, err := strconv.ParseUint(v, 10, 64)
	if err != nil {
		return false, err
	}
	if expires >= threshold {
		t := time.Duration(expires-threshold) * time.Second
		logDebug("key %s expires in %s", key, t)
		return false, err
	}
	t := time.Duration(threshold-expires) * time.Second
	logDebug("key %s has expired %s ago", key, t)
	msgId := s.msgId(key)
	keys := s.keys(msgId)
	// logDebug("deleting keys: %s", strings.Join(keys, ", "))
	_, err = s.redis.Del(keys...).Result()
	return err == nil, err
}

func (s *GCState) garbageCollectKeys(db int) {
	var total, expired int
	var cursor uint64
	defer func() { logInfo("expired %d keys out of %d in db %d", expired, total, db) }()
	ticker := time.NewTicker(1 * time.Second)
collecting:
	for _ = range ticker.C {
		if interrupted {
			break
		}
		if s.getMaster(db) {
			if cursor == 0 {
				logInfo("starting SCAN on db %d", db)
			}
			logDebug("cursor: %d", cursor)
			var err error
			var keys []string
			keys, cursor, err = s.redis.Scan(cursor, "msgid:*:expires", 10000).Result()
			if err != nil {
				logError("starting over: %v", err)
				cursor = 0
				total = 0
				expired = 0
				continue collecting
			}
			logDebug("retrieved %d keys from db %d", len(keys), db)
			total += len(keys)
			threshold := time.Now().Unix() + int64(s.opts.GcThreshold)
			for _, key := range keys {
				collected, err := s.gcKey(key, uint64(threshold))
				if err != nil {
					logError("starting over: %v", err)
					cursor = 0
					total = 0
					expired = 0
					goto collecting
				}
				if collected {
					expired++
				}
			}
			if cursor == 0 {
				return
			}
		}
	}
}

func (s *GCState) garbageCollectKeysFromFile(db int, filePath string) {
	var total, expired int
	defer func() { logInfo("expired %d keys out of %d in db %d", expired, total, db) }()

	file, err := os.Open(filePath)
	if err != nil {
		logError("%v", err)
		return
	}
	defer file.Close()

	s.getMaster(db)

	threshold := time.Now().Unix() + int64(s.opts.GcThreshold)
	re := regexp.MustCompile("^(msgid:[^:]+:[-0-9a-f]+):expires$")
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		if interrupted {
			break
		}
		line := scanner.Text()
		if !re.MatchString(line) {
			continue
		}
		total++
		collected, err := s.gcKey(line, uint64(threshold))
		if err != nil {
			logError("could not collect %s: %v", line, err)
			continue
		}
		if collected {
			expired++
		}
	}
	if err := scanner.Err(); err != nil {
		logError("%v", err)
	}
}

func (s *GCState) getMaster(db int) bool {
	server := ReadRedisMasterFile(s.opts.RedisMasterFile)
	if s.currentMaster != server || s.currentDB != db {
		s.currentMaster = server
		s.currentDB = db
		if server == "" {
			if s.redis != nil {
				s.redis.Close()
			}
			s.redis = nil
		} else {
			s.redis = redis.NewClient(&redis.Options{Addr: server, DB: db})
		}
	}
	if s.redis == nil {
		logError("could not determine redis master: %v, db: %d", s.currentMaster, db)
	}
	return s.redis != nil
}

// RunGarbageCollectKeys runs a garbage collection on the redis master using the
// redis SCAN operation. Restarts from the beginning, should the master change
// while running the scan. Terminates as soon as a full scan has been performed
// on all databases which need GC.
func RunGarbageCollectKeys(opts GCOptions) error {
	logDebug("garbage collecting keys with options: %+v", opts)
	state := &GCState{opts: opts}
	state.keySuffixes = []string{"status", "ack_count", "timeout", "delay", "attempts", "exceptions", "mutex", "expires"}
	for _, s := range strings.Split(opts.GcDatabases, ",") {
		if interrupted {
			break
		}
		if s != "" {
			db, err := strconv.Atoi(s)
			if err != nil {
				logError("%v", err)
				continue
			}
			if opts.GcKeyFile == "" {
				state.garbageCollectKeys(db)
			} else {
				state.garbageCollectKeysFromFile(db, opts.GcKeyFile)
			}
		}
	}
	if state.redis != nil {
		return state.redis.Close()
	}
	return nil
}
<commit_msg>key GC: don't panic when a malformed key is found<commit_after>package main

import (
	"bufio"
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	"gopkg.in/redis.v5"
)

// GCOptions are provided by the caller of RunGarbageCollectKeys.
type GCOptions struct {
	RedisMasterFile string // Path to the redis master file.
	GcThreshold     int    // Number of seconds after which a key should be considered collectible.
	GcDatabases     string // List of databases to scan.
	GcKeyFile       string // Name of file containing keys to collect.
}

// GCState holds options and collector state, most importantly the current redis
// connection and database.
type GCState struct {
	opts          GCOptions
	currentMaster string
	currentDB     int
	redis         *redis.Client // current connection
	keySuffixes   []string
}

func (s *GCState) key(msgId, suffix string) string {
	return fmt.Sprintf("%s:%s", msgId, suffix)
}

func (s *GCState) keys(msgId string) []string {
	res := make([]string, 0)
	for _, suffix := range s.keySuffixes {
		res = append(res, s.key(msgId, suffix))
	}
	return res
}

func (s *GCState) msgId(key string) string {
	re := regexp.MustCompile("^(msgid:[^:]+:[-0-9a-f]+):.*$")
	matches := re.FindStringSubmatch(key)
	if len(matches) == 0 {
		logError("msgid could not be extracted from key '%s'", key)
		return ""
	}
	return matches[1]
}

func (s *GCState) gcKey(key string, threshold uint64) (bool, error) {
	v, err := s.redis.Get(key).Result()
	if err != nil {
		if err == redis.Nil {
			logDebug("key not found: %s", key)
			return false, nil
		}
		return false, err
	}
	expires, err := strconv.ParseUint(v, 10, 64)
	if err != nil {
		return false, err
	}
	if expires >= threshold {
		t := time.Duration(expires-threshold) * time.Second
		logDebug("key %s expires in %s", key, t)
		return false, err
	}
	t := time.Duration(threshold-expires) * time.Second
	logDebug("key %s has expired %s ago", key, t)
	msgID := s.msgId(key)
	if msgID == "" {
		return false, nil
	}
	keys := s.keys(msgID)
	// logDebug("deleting keys: %s", strings.Join(keys, ", "))
	_, err = s.redis.Del(keys...).Result()
	return err == nil, err
}

func (s *GCState) garbageCollectKeys(db int) {
	var total, expired int
	var cursor uint64
	defer func() { logInfo("expired %d keys out of %d in db %d", expired, total, db) }()
	ticker := time.NewTicker(1 * time.Second)
collecting:
	for _ = range ticker.C {
		if interrupted {
			break
		}
		if s.getMaster(db) {
			if cursor == 0 {
				logInfo("starting SCAN on db %d", db)
			}
			logDebug("cursor: %d", cursor)
			var err error
			var keys []string
			keys, cursor, err = s.redis.Scan(cursor, "msgid:*:expires", 10000).Result()
			if err != nil {
				logError("starting over: %v", err)
				cursor = 0
				total = 0
				expired = 0
				continue collecting
			}
			logDebug("retrieved %d keys from db %d", len(keys), db)
			total += len(keys)
			threshold := time.Now().Unix() + int64(s.opts.GcThreshold)
			for _, key := range keys {
				collected, err := s.gcKey(key, uint64(threshold))
				if err != nil {
					logError("starting over: %v", err)
					cursor = 0
					total = 0
					expired = 0
					goto collecting
				}
				if collected {
					expired++
				}
			}
			if cursor == 0 {
				return
			}
		}
	}
}

func (s *GCState) garbageCollectKeysFromFile(db int, filePath string) {
	var total, expired int
	defer func() { logInfo("expired %d keys out of %d in db %d", expired, total, db) }()

	file, err := os.Open(filePath)
	if err != nil {
		logError("%v", err)
		return
	}
	defer file.Close()

	s.getMaster(db)

	threshold := time.Now().Unix() + int64(s.opts.GcThreshold)
	re := regexp.MustCompile("^(msgid:[^:]+:[-0-9a-f]+):expires$")
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		if interrupted {
			break
		}
		line := scanner.Text()
		if !re.MatchString(line) {
			continue
		}
		total++
		collected, err := s.gcKey(line, uint64(threshold))
		if err != nil {
			logError("could not collect %s: %v", line, err)
			continue
		}
		if collected {
			expired++
		}
	}
	if err := scanner.Err(); err != nil {
		logError("%v", err)
	}
}

func (s *GCState) getMaster(db int) bool {
	server := ReadRedisMasterFile(s.opts.RedisMasterFile)
	if s.currentMaster != server || s.currentDB != db {
		s.currentMaster = server
		s.currentDB = db
		if server == "" {
			if s.redis != nil {
				s.redis.Close()
			}
			s.redis = nil
		} else {
			s.redis = redis.NewClient(&redis.Options{Addr: server, DB: db})
		}
	}
	if s.redis == nil {
		logError("could not determine redis master: %v, db: %d", s.currentMaster, db)
	}
	return s.redis != nil
}

// RunGarbageCollectKeys runs a garbage collection on the redis master using the
// redis SCAN operation. Restarts from the beginning, should the master change
// while running the scan. Terminates as soon as a full scan has been performed
// on all databases which need GC.
func RunGarbageCollectKeys(opts GCOptions) error {
	logDebug("garbage collecting keys with options: %+v", opts)
	state := &GCState{opts: opts}
	state.keySuffixes = []string{"status", "ack_count", "timeout", "delay", "attempts", "exceptions", "mutex", "expires"}
	for _, s := range strings.Split(opts.GcDatabases, ",") {
		if interrupted {
			break
		}
		if s != "" {
			db, err := strconv.Atoi(s)
			if err != nil {
				logError("%v", err)
				continue
			}
			if opts.GcKeyFile == "" {
				state.garbageCollectKeys(db)
			} else {
				state.garbageCollectKeysFromFile(db, opts.GcKeyFile)
			}
		}
	}
	if state.redis != nil {
		return state.redis.Close()
	}
	return nil
}
<|endoftext|>