<commit_before>package patcher_test

import (
	"bytes"
	"context"
	"encoding/gob"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"testing"

	"github.com/itchio/headway/united"
	"github.com/itchio/screw"
	"github.com/itchio/wharf/wsync"
	"github.com/pkg/errors"
	"github.com/stretchr/testify/assert"

	"github.com/itchio/wharf/pwr"
	"github.com/itchio/wharf/pwr/bowl"
	"github.com/itchio/wharf/pwr/patcher"
	"github.com/itchio/wharf/pwr/rediff"
	"github.com/itchio/wharf/wtest"

	"github.com/itchio/headway/state"
	"github.com/itchio/lake/pools/fspool"
	"github.com/itchio/lake/tlc"
	"github.com/itchio/savior/seeksource"

	_ "github.com/itchio/wharf/compressors/cbrotli"
	_ "github.com/itchio/wharf/decompressors/cbrotli"
)

func Test_Naive(t *testing.T) {
	dir, err := ioutil.TempDir("", "patcher-noop")
	wtest.Must(t, err)
	defer screw.RemoveAll(dir)

	v1 := filepath.Join(dir, "v1")
	wtest.MakeTestDir(t, v1, wtest.TestDirSettings{
		Entries: []wtest.TestDirEntry{
			{Path: "subdir/file-1", Seed: 0x1, Size: wtest.BlockSize*220 + 14},
			{Path: "file-1", Seed: 0x2},
			{Path: "dir2/file-2", Seed: 0x3},
			{Path: "dir3/gone", Seed: 0x4},
		},
	})

	v2 := filepath.Join(dir, "v2")
	wtest.MakeTestDir(t, v2, wtest.TestDirSettings{
		Entries: []wtest.TestDirEntry{
			{Path: "subdir/file-1", Seed: 0x1, Size: wtest.BlockSize*260 + 14, Bsmods: []wtest.Bsmod{
				{Interval: wtest.BlockSize/2 + 3, Delta: 0x4},
				{Interval: wtest.BlockSize/3 + 7, Delta: 0x18},
			}, Swaperoos: []wtest.Swaperoo{
				{OldStart: 0, NewStart: wtest.BlockSize * 210, Size: wtest.BlockSize * 10},
				{OldStart: 40, NewStart: wtest.BlockSize*10 + 8, Size: wtest.BlockSize * 40},
			}},
			{Path: "file-1", Seed: 0x2},
			{Path: "dir2/file-2", Seed: 0x3},
		},
	})

	patchBuffer := new(bytes.Buffer)
	optimizedPatchBuffer := new(bytes.Buffer)
	var sourceHashes []wsync.BlockHash
	consumer := &state.Consumer{
		OnMessage: func(level string, message string) {
			t.Logf("[%s] %s", level, message)
		},
	}

	{
		compression := &pwr.CompressionSettings{}
		compression.Algorithm = pwr.CompressionAlgorithm_BROTLI
		compression.Quality = 1

		targetContainer, err := tlc.WalkAny(v1, &tlc.WalkOpts{})
		wtest.Must(t, err)

		sourceContainer, err := tlc.WalkAny(v2, &tlc.WalkOpts{})
		wtest.Must(t, err)

		// Sign!
		t.Logf("Signing %s", sourceContainer.Stats())
		sourceHashes, err = pwr.ComputeSignature(context.Background(), sourceContainer, fspool.New(sourceContainer, v2), consumer)
		wtest.Must(t, err)

		targetPool := fspool.New(targetContainer, v1)
		targetSignature, err := pwr.ComputeSignature(context.Background(), targetContainer, targetPool, consumer)
		wtest.Must(t, err)

		pool := fspool.New(sourceContainer, v2)

		// Diff!
		t.Logf("Diffing (%s)...", compression)
		dctx := pwr.DiffContext{
			Compression: compression,
			Consumer:    consumer,

			SourceContainer: sourceContainer,
			Pool:            pool,

			TargetContainer: targetContainer,
			TargetSignature: targetSignature,
		}

		wtest.Must(t, dctx.WritePatch(context.Background(), patchBuffer, ioutil.Discard))

		// Rediff!
		t.Logf("Rediffing...")
		rc, err := rediff.NewContext(rediff.Params{
			Consumer:    consumer,
			PatchReader: seeksource.FromBytes(patchBuffer.Bytes()),
		})
		wtest.Must(t, err)

		wtest.Must(t, rc.Optimize(rediff.OptimizeParams{
			TargetPool:  targetPool,
			SourcePool:  pool,
			PatchWriter: optimizedPatchBuffer,
		}))
	}

	// Patch!
	tryPatchNoSaves := func(t *testing.T, patchBytes []byte) {
		consumer := &state.Consumer{
			OnMessage: func(level string, message string) {
				t.Logf("[%s] %s", level, message)
			},
		}

		out := filepath.Join(dir, "out")
		defer screw.RemoveAll(out)

		patchReader := seeksource.FromBytes(patchBytes)

		p, err := patcher.New(patchReader, consumer)
		wtest.Must(t, err)

		targetPool := fspool.New(p.GetTargetContainer(), v1)

		b, err := bowl.NewFreshBowl(bowl.FreshBowlParams{
			SourceContainer: p.GetSourceContainer(),
			TargetContainer: p.GetTargetContainer(),
			TargetPool:      targetPool,
			OutputFolder:    out,
		})
		wtest.Must(t, err)

		err = p.Resume(nil, targetPool, b)
		wtest.Must(t, err)

		// Validate!
		sigInfo := &pwr.SignatureInfo{
			Container: p.GetSourceContainer(),
			Hashes:    sourceHashes,
		}
		wtest.Must(t, pwr.AssertValid(out, sigInfo))
		wtest.Must(t, pwr.AssertNoGhosts(out, sigInfo))

		t.Logf("Patch applies cleanly!")
	}

	tryPatchWithSaves := func(t *testing.T, patchBytes []byte) {
		consumer := &state.Consumer{
			OnMessage: func(level string, message string) {
				t.Logf("[%s] %s", level, message)
			},
		}

		out := filepath.Join(dir, "out")
		defer screw.RemoveAll(out)

		patchReader := seeksource.FromBytes(patchBytes)

		p, err := patcher.New(patchReader, consumer)
		wtest.Must(t, err)

		var checkpoint *patcher.Checkpoint
		p.SetSaveConsumer(&patcherSaveConsumer{
			shouldSave: func() bool {
				return true
			},
			save: func(c *patcher.Checkpoint) (patcher.AfterSaveAction, error) {
				checkpoint = c
				return patcher.AfterSaveStop, nil
			},
		})

		targetPool := fspool.New(p.GetTargetContainer(), v1)

		b, err := bowl.NewFreshBowl(bowl.FreshBowlParams{
			SourceContainer: p.GetSourceContainer(),
			TargetContainer: p.GetTargetContainer(),
			TargetPool:      targetPool,
			OutputFolder:    out,
		})
		wtest.Must(t, err)

		numCheckpoints := 0
		for {
			c := checkpoint
			checkpoint = nil
			t.Logf("Resuming patcher - has checkpoint: %v", c != nil)
			err = p.Resume(c, targetPool, b)
			if errors.Cause(err) == patcher.ErrStop {
				t.Logf("Patcher returned ErrStop")

				if checkpoint == nil {
					wtest.Must(t, errors.New("patcher stopped but nil checkpoint"))
				}
				numCheckpoints++

				checkpointBuf := new(bytes.Buffer)
				enc := gob.NewEncoder(checkpointBuf)
				wtest.Must(t, enc.Encode(checkpoint))

				t.Logf("Got %s checkpoint @ %.2f%% of the patch", united.FormatBytes(int64(checkpointBuf.Len())), p.Progress()*100.0)

				checkpoint = &patcher.Checkpoint{}
				dec := gob.NewDecoder(bytes.NewReader(checkpointBuf.Bytes()))
				wtest.Must(t, dec.Decode(checkpoint))

				continue
			}

			wtest.Must(t, err)
			break
		}

		// Validate!
		wtest.Must(t, pwr.AssertValid(out, &pwr.SignatureInfo{
			Container: p.GetSourceContainer(),
			Hashes:    sourceHashes,
		}))

		t.Logf("Patch applies cleanly!")

		t.Logf("Had %d checkpoints total", numCheckpoints)
		assert.True(t, numCheckpoints > 0, "had at least one checkpoint")
	}

	tryPatch := func(kind string, patchBytes []byte) {
		t.Run(fmt.Sprintf("%s-no-saves", kind), func(t *testing.T) {
			t.Logf("Applying %s %s patch (%d bytes), no saves", united.FormatBytes(int64(len(patchBytes))), kind, len(patchBytes))
			tryPatchNoSaves(t, patchBytes)
		})

		t.Run(fmt.Sprintf("%s-with-saves", kind), func(t *testing.T) {
			t.Logf("Applying %s %s patch (%d bytes) with saves", united.FormatBytes(int64(len(patchBytes))), kind, len(patchBytes))
			tryPatchWithSaves(t, patchBytes)
		})
	}

	tryPatch("simple", patchBuffer.Bytes())
	tryPatch("optimized", optimizedPatchBuffer.Bytes())
}

//

type patcherSaveConsumer struct {
	shouldSave func() bool
	save       func(checkpoint *patcher.Checkpoint) (patcher.AfterSaveAction, error)
}

var _ patcher.SaveConsumer = (*patcherSaveConsumer)(nil)

func (psc *patcherSaveConsumer) ShouldSave() bool {
	return psc.shouldSave()
}

func (psc *patcherSaveConsumer) Save(checkpoint *patcher.Checkpoint) (patcher.AfterSaveAction, error) {
	return psc.save(checkpoint)
}
<commit_msg>Add tests for sourceIndexWhitelist<commit_after>package patcher_test

import (
	"bytes"
	"context"
	"encoding/gob"
	"fmt"
	"io"
	"io/ioutil"
	"path/filepath"
	"testing"

	"github.com/itchio/headway/united"
	"github.com/itchio/lake"
	"github.com/itchio/screw"
	"github.com/itchio/wharf/wsync"
	"github.com/pkg/errors"
	"github.com/stretchr/testify/assert"

	"github.com/itchio/wharf/pwr"
	"github.com/itchio/wharf/pwr/bowl"
	"github.com/itchio/wharf/pwr/patcher"
	"github.com/itchio/wharf/pwr/rediff"
	"github.com/itchio/wharf/wtest"

	"github.com/itchio/headway/state"
	"github.com/itchio/lake/pools/fspool"
	"github.com/itchio/lake/tlc"
	"github.com/itchio/savior/seeksource"

	_ "github.com/itchio/wharf/compressors/cbrotli"
	_ "github.com/itchio/wharf/decompressors/cbrotli"
)

func Test_Naive(t *testing.T) {
	dir, err := ioutil.TempDir("", "patcher-noop")
	wtest.Must(t, err)
	defer screw.RemoveAll(dir)

	v1 := filepath.Join(dir, "v1")
	wtest.MakeTestDir(t, v1, wtest.TestDirSettings{
		Entries: []wtest.TestDirEntry{
			{Path: "subdir/file-1", Seed: 0x1, Size: wtest.BlockSize*220 + 14},
			{Path: "file-1", Seed: 0x2},
			{Path: "dir2/file-2", Seed: 0x3},
			{Path: "dir3/gone", Seed: 0x4},
		},
	})

	v2 := filepath.Join(dir, "v2")
	wtest.MakeTestDir(t, v2, wtest.TestDirSettings{
		Entries: []wtest.TestDirEntry{
			{Path: "subdir/file-1", Seed: 0x1, Size: wtest.BlockSize*260 + 14, Bsmods: []wtest.Bsmod{
				{Interval: wtest.BlockSize/2 + 3, Delta: 0x4},
				{Interval: wtest.BlockSize/3 + 7, Delta: 0x18},
			}, Swaperoos: []wtest.Swaperoo{
				{OldStart: 0, NewStart: wtest.BlockSize * 210, Size: wtest.BlockSize * 10},
				{OldStart: 40, NewStart: wtest.BlockSize*10 + 8, Size: wtest.BlockSize * 40},
			}},
			{Path: "file-1", Seed: 0x2},
			{Path: "dir2/file-2", Seed: 0x3},
		},
	})

	patchBuffer := new(bytes.Buffer)
	optimizedPatchBuffer := new(bytes.Buffer)
	var sourceHashes []wsync.BlockHash
	consumer := &state.Consumer{
		OnMessage: func(level string, message string) {
			t.Logf("[%s] %s", level, message)
		},
	}

	{
		compression := &pwr.CompressionSettings{}
		compression.Algorithm = pwr.CompressionAlgorithm_BROTLI
		compression.Quality = 1

		targetContainer, err := tlc.WalkAny(v1, &tlc.WalkOpts{})
		wtest.Must(t, err)

		sourceContainer, err := tlc.WalkAny(v2, &tlc.WalkOpts{})
		wtest.Must(t, err)

		// Sign!
		t.Logf("Signing %s", sourceContainer.Stats())
		sourceHashes, err = pwr.ComputeSignature(context.Background(), sourceContainer, fspool.New(sourceContainer, v2), consumer)
		wtest.Must(t, err)

		targetPool := fspool.New(targetContainer, v1)
		targetSignature, err := pwr.ComputeSignature(context.Background(), targetContainer, targetPool, consumer)
		wtest.Must(t, err)

		pool := fspool.New(sourceContainer, v2)

		// Diff!
		t.Logf("Diffing (%s)...", compression)
		dctx := pwr.DiffContext{
			Compression: compression,
			Consumer:    consumer,

			SourceContainer: sourceContainer,
			Pool:            pool,

			TargetContainer: targetContainer,
			TargetSignature: targetSignature,
		}

		wtest.Must(t, dctx.WritePatch(context.Background(), patchBuffer, ioutil.Discard))

		// Rediff!
		t.Logf("Rediffing...")
		rc, err := rediff.NewContext(rediff.Params{
			Consumer:    consumer,
			PatchReader: seeksource.FromBytes(patchBuffer.Bytes()),
		})
		wtest.Must(t, err)

		wtest.Must(t, rc.Optimize(rediff.OptimizeParams{
			TargetPool:  targetPool,
			SourcePool:  pool,
			PatchWriter: optimizedPatchBuffer,
		}))
	}

	// Patch!
	tryPatchNoSaves := func(t *testing.T, patchBytes []byte) {
		consumer := &state.Consumer{
			OnMessage: func(level string, message string) {
				t.Logf("[%s] %s", level, message)
			},
		}

		out := filepath.Join(dir, "out")
		defer screw.RemoveAll(out)

		patchReader := seeksource.FromBytes(patchBytes)

		p, err := patcher.New(patchReader, consumer)
		wtest.Must(t, err)

		targetPool := fspool.New(p.GetTargetContainer(), v1)

		b, err := bowl.NewFreshBowl(bowl.FreshBowlParams{
			SourceContainer: p.GetSourceContainer(),
			TargetContainer: p.GetTargetContainer(),
			TargetPool:      targetPool,
			OutputFolder:    out,
		})
		wtest.Must(t, err)

		err = p.Resume(nil, targetPool, b)
		wtest.Must(t, err)

		// Validate!
		sigInfo := &pwr.SignatureInfo{
			Container: p.GetSourceContainer(),
			Hashes:    sourceHashes,
		}
		wtest.Must(t, pwr.AssertValid(out, sigInfo))
		wtest.Must(t, pwr.AssertNoGhosts(out, sigInfo))

		t.Logf("Patch applies cleanly!")
	}

	tryPatchSkip := func(t *testing.T, patchBytes []byte, all bool) {
		consumer := &state.Consumer{
			OnMessage: func(level string, message string) {
				t.Logf("[%s] %s", level, message)
			},
		}

		out := filepath.Join(dir, "out")
		defer screw.RemoveAll(out)

		patchReader := seeksource.FromBytes(patchBytes)

		p, err := patcher.New(patchReader, consumer)
		wtest.Must(t, err)

		var targetPool lake.Pool = &explodingPool{}

		sourceIndexWhitelist := make(map[int64]bool)
		if !all {
			for i := int64(0); i < int64(len(p.GetSourceContainer().Files)); i += 2 {
				sourceIndexWhitelist[i] = true
			}
			targetPool = fspool.New(p.GetTargetContainer(), v1)
		}
		p.SetSourceIndexWhitelist(sourceIndexWhitelist)

		b, err := bowl.NewFreshBowl(bowl.FreshBowlParams{
			SourceContainer: p.GetSourceContainer(),
			TargetContainer: p.GetTargetContainer(),
			TargetPool:      targetPool,
			OutputFolder:    out,
		})
		wtest.Must(t, err)

		err = p.Resume(nil, targetPool, b)
		wtest.Must(t, err)

		// Validate!
		sigInfo := &pwr.SignatureInfo{
			Container: p.GetSourceContainer(),
			Hashes:    sourceHashes,
		}
		err = pwr.AssertValid(out, sigInfo)
		assert.Error(t, err)

		t.Logf("Partially applied!")
	}

	tryPatchWithSaves := func(t *testing.T, patchBytes []byte) {
		consumer := &state.Consumer{
			OnMessage: func(level string, message string) {
				t.Logf("[%s] %s", level, message)
			},
		}

		out := filepath.Join(dir, "out")
		defer screw.RemoveAll(out)

		patchReader := seeksource.FromBytes(patchBytes)

		p, err := patcher.New(patchReader, consumer)
		wtest.Must(t, err)

		var checkpoint *patcher.Checkpoint
		p.SetSaveConsumer(&patcherSaveConsumer{
			shouldSave: func() bool {
				return true
			},
			save: func(c *patcher.Checkpoint) (patcher.AfterSaveAction, error) {
				checkpoint = c
				return patcher.AfterSaveStop, nil
			},
		})

		targetPool := fspool.New(p.GetTargetContainer(), v1)

		b, err := bowl.NewFreshBowl(bowl.FreshBowlParams{
			SourceContainer: p.GetSourceContainer(),
			TargetContainer: p.GetTargetContainer(),
			TargetPool:      targetPool,
			OutputFolder:    out,
		})
		wtest.Must(t, err)

		numCheckpoints := 0
		for {
			c := checkpoint
			checkpoint = nil
			t.Logf("Resuming patcher - has checkpoint: %v", c != nil)
			err = p.Resume(c, targetPool, b)
			if errors.Cause(err) == patcher.ErrStop {
				t.Logf("Patcher returned ErrStop")

				if checkpoint == nil {
					wtest.Must(t, errors.New("patcher stopped but nil checkpoint"))
				}
				numCheckpoints++

				checkpointBuf := new(bytes.Buffer)
				enc := gob.NewEncoder(checkpointBuf)
				wtest.Must(t, enc.Encode(checkpoint))

				t.Logf("Got %s checkpoint @ %.2f%% of the patch", united.FormatBytes(int64(checkpointBuf.Len())), p.Progress()*100.0)

				checkpoint = &patcher.Checkpoint{}
				dec := gob.NewDecoder(bytes.NewReader(checkpointBuf.Bytes()))
				wtest.Must(t, dec.Decode(checkpoint))

				continue
			}

			wtest.Must(t, err)
			break
		}

		// Validate!
		wtest.Must(t, pwr.AssertValid(out, &pwr.SignatureInfo{
			Container: p.GetSourceContainer(),
			Hashes:    sourceHashes,
		}))

		t.Logf("Patch applies cleanly!")

		t.Logf("Had %d checkpoints total", numCheckpoints)
		assert.True(t, numCheckpoints > 0, "had at least one checkpoint")
	}

	tryPatch := func(kind string, patchBytes []byte) {
		t.Run(fmt.Sprintf("%s-no-saves", kind), func(t *testing.T) {
			t.Logf("Applying %s %s patch (%d bytes), no saves", united.FormatBytes(int64(len(patchBytes))), kind, len(patchBytes))
			tryPatchNoSaves(t, patchBytes)
		})

		t.Run(fmt.Sprintf("%s-with-saves", kind), func(t *testing.T) {
			t.Logf("Applying %s %s patch (%d bytes) with saves", united.FormatBytes(int64(len(patchBytes))), kind, len(patchBytes))
			tryPatchWithSaves(t, patchBytes)
		})

		t.Run(fmt.Sprintf("%s-skip-all", kind), func(t *testing.T) {
			t.Logf("Applying %s %s patch (%d bytes) by skipping all entries", united.FormatBytes(int64(len(patchBytes))), kind, len(patchBytes))
			tryPatchSkip(t, patchBytes, true)
		})

		t.Run(fmt.Sprintf("%s-skip-some", kind), func(t *testing.T) {
			t.Logf("Applying %s %s patch (%d bytes) by skipping some entries", united.FormatBytes(int64(len(patchBytes))), kind, len(patchBytes))
			tryPatchSkip(t, patchBytes, false)
		})
	}

	tryPatch("simple", patchBuffer.Bytes())
	tryPatch("optimized", optimizedPatchBuffer.Bytes())
}

//

type patcherSaveConsumer struct {
	shouldSave func() bool
	save       func(checkpoint *patcher.Checkpoint) (patcher.AfterSaveAction, error)
}

var _ patcher.SaveConsumer = (*patcherSaveConsumer)(nil)

func (psc *patcherSaveConsumer) ShouldSave() bool {
	return psc.shouldSave()
}

func (psc *patcherSaveConsumer) Save(checkpoint *patcher.Checkpoint) (patcher.AfterSaveAction, error) {
	return psc.save(checkpoint)
}

//

type explodingPool struct{}

var _ lake.Pool = (*explodingPool)(nil)

func (ep *explodingPool) GetSize(fileIndex int64) int64 {
	panic("pool exploded")
}
func (ep *explodingPool) GetReader(fileIndex int64) (io.Reader, error) {
	panic("pool exploded")
}
func (ep *explodingPool) GetReadSeeker(fileIndex int64) (io.ReadSeeker, error) {
	panic("pool exploded")
}
func (ep *explodingPool) Close() error {
	return nil
}
<|endoftext|>


<commit_before>package workflow

import (
	"github.com/go-gorp/gorp"

	"github.com/ovh/cds/engine/api/event"
	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/log"
)

// GetWorkflowRunEventData read channel to get elements to push
func GetWorkflowRunEventData(cError <-chan error, cEvent <-chan interface{}) ([]sdk.WorkflowRun, []sdk.WorkflowNodeRun, []sdk.WorkflowNodeJobRun, error) {
	wrs := []sdk.WorkflowRun{}
	wnrs := []sdk.WorkflowNodeRun{}
	wnjrs := []sdk.WorkflowNodeJobRun{}

	for {
		select {
		case e, has := <-cError:
			log.Info("GetWorkflowRunEventData> cError has: %t err:%v", has, e)
			if !has {
				return wrs, wnrs, wnjrs, e
			}
			if e != nil {
				return nil, nil, nil, e
			}
		case w, has := <-cEvent:
			log.Info("GetWorkflowRunEventData> cEvent has: %t", has)
			if !has {
				return wrs, wnrs, wnjrs, nil
			}
			switch x := w.(type) {
			case sdk.WorkflowNodeJobRun:
				wnjrs = append(wnjrs, x)
			case sdk.WorkflowNodeRun:
				wnrs = append(wnrs, x)
			case sdk.WorkflowRun:
				wrs = append(wrs, x)
			default:
				log.Warning("GetWorkflowRunEventData> unknown type %T", w)
			}
		}
	}
}

// SendEvent Send event on workflow run
func SendEvent(db gorp.SqlExecutor, wrs []sdk.WorkflowRun, wnrs []sdk.WorkflowNodeRun, wnjrs []sdk.WorkflowNodeJobRun, key string) {
	for _, wr := range wrs {
		event.PublishWorkflowRun(wr, key)
	}
	for _, wnr := range wnrs {
		wr, errWR := LoadRunByID(db, wnr.WorkflowRunID, false)
		if errWR != nil {
			log.Warning("SendEvent> Cannot load workflow run %d: %s", wnr.WorkflowRunID, errWR)
			continue
		}
		event.PublishWorkflowNodeRun(wnr, *wr, key)
	}
	for _, wnjr := range wnjrs {
		event.PublishWorkflowNodeJobRun(wnjr)
	}
}
<commit_msg>fix(api): do not return if channels are not closed (#1645)<commit_after>package workflow

import (
	"github.com/go-gorp/gorp"

	"github.com/ovh/cds/engine/api/event"
	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/log"
)

// GetWorkflowRunEventData read channel to get elements to push
func GetWorkflowRunEventData(cError <-chan error, cEvent <-chan interface{}) ([]sdk.WorkflowRun, []sdk.WorkflowNodeRun, []sdk.WorkflowNodeJobRun, error) {
	wrs := []sdk.WorkflowRun{}
	wnrs := []sdk.WorkflowNodeRun{}
	wnjrs := []sdk.WorkflowNodeJobRun{}
	var err error

	for {
		select {
		case e, has := <-cError:
			log.Info("GetWorkflowRunEventData> cError has: %t err:%v", has, e)
			err = sdk.WrapError(e, "GetWorkflowRunEventData> Error received")
			if !has {
				return wrs, wnrs, wnjrs, err
			}
		case w, has := <-cEvent:
			log.Info("GetWorkflowRunEventData> cEvent has: %t", has)
			if !has {
				return wrs, wnrs, wnjrs, err
			}
			switch x := w.(type) {
			case sdk.WorkflowNodeJobRun:
				wnjrs = append(wnjrs, x)
			case sdk.WorkflowNodeRun:
				wnrs = append(wnrs, x)
			case sdk.WorkflowRun:
				wrs = append(wrs, x)
			default:
				log.Warning("GetWorkflowRunEventData> unknown type %T", w)
			}
		}
	}
}

// SendEvent Send event on workflow run
func SendEvent(db gorp.SqlExecutor, wrs []sdk.WorkflowRun, wnrs []sdk.WorkflowNodeRun, wnjrs []sdk.WorkflowNodeJobRun, key string) {
	for _, wr := range wrs {
		event.PublishWorkflowRun(wr, key)
	}
	for _, wnr := range wnrs {
		wr, errWR := LoadRunByID(db, wnr.WorkflowRunID, false)
		if errWR != nil {
			log.Warning("SendEvent> Cannot load workflow run %d: %s", wnr.WorkflowRunID, errWR)
			continue
		}
		event.PublishWorkflowNodeRun(wnr, *wr, key)
	}
	for _, wnjr := range wnjrs {
		event.PublishWorkflowNodeJobRun(wnjr)
	}
}
<|endoftext|>


<commit_before><commit_msg>Auto commit 02/10/2012 21:00:03<commit_after><|endoftext|>


<commit_before>package model

import (
	"strings"
	"testing"
	"time"

	"github.com/oinume/lekcije/server/config"
	"github.com/oinume/lekcije/server/util"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestUpdateLessons(t *testing.T) {
	a := assert.New(t)
	r := require.New(t)

	teacherID := uint32(util.RandomInt(999999))
	datetime := time.Date(2016, 10, 1, 14, 30, 0, 0, config.LocalTimezone())
	lessons := createLessons(teacherID, datetime, "Reserved", 5)

	affected, err := lessonService.UpdateLessons(lessons)
	r.NoError(err)
	a.Equal(int64(5), affected)
	for _, l := range lessons {
		a.NotEqual(uint64(0), l.ID)
		logs, err := lessonStatusLogService.FindAllByLessonID(l.ID)
		r.NoError(err)
		a.Equal(1, len(logs))
	}

	foundLessons, err := lessonService.FindLessons(teacherID, datetime, datetime)
	r.NoError(err)
	a.Equal(len(lessons), len(foundLessons))
	for i := range lessons {
		// TODO: custom enum type
		a.Equal(strings.ToLower(lessons[i].Status), strings.ToLower(foundLessons[i].Status))
	}
}

func TestUpdateLessonsOverwrite(t *testing.T) {
	a := assert.New(t)
	r := require.New(t)

	teacherID := uint32(util.RandomInt(999999))
	datetime := time.Date(2016, 10, 1, 14, 30, 0, 0, config.LocalTimezone())
	lessons := createLessons(teacherID, datetime, "Available", 5)
	affected, err := lessonService.UpdateLessons(lessons)
	r.NoError(err)
	a.EqualValues(len(lessons), affected)

	time.Sleep(1 * time.Second)
	lessons[0].Status = "Reserved"
	affected, err = lessonService.UpdateLessons(lessons)
	r.NoError(err)
	a.EqualValues(1, affected)

	foundLessons, err := lessonService.FindLessons(teacherID, datetime, datetime)
	r.NoError(err)
	a.Equal(strings.ToLower(foundLessons[0].Status), "reserved")

	logs, err := lessonStatusLogService.FindAllByLessonID(foundLessons[0].ID)
	r.NoError(err)
	a.Equal(2, len(logs))
}

func TestGetNewAvailableLessons1(t *testing.T) {
	a := assert.New(t)

	datetime := time.Date(2016, 10, 1, 14, 30, 0, 0, config.LocalTimezone())
	lessons1 := createLessons(1, datetime, "Reserved", 3)
	lessons2 := createLessons(1, datetime, "Reserved", 3)
	lessons2[1].Status = "Available"
	// Test GetNewAvailableLessons returns a lesson when new lesson is "Available"
	availableLessons := lessonService.GetNewAvailableLessons(lessons1, lessons2)
	a.Equal(1, len(availableLessons))
	a.Equal(datetime.Add(1*time.Hour), availableLessons[0].Datetime)
}

func TestGetNewAvailableLessons2(t *testing.T) {
	a := assert.New(t)

	datetime := time.Date(2016, 10, 1, 14, 30, 0, 0, config.LocalTimezone())
	lessons1 := createLessons(1, datetime, "Reserved", 3)
	lessons2 := createLessons(1, datetime, "Reserved", 3)
	lessons1[0].Status = "Available"
	lessons2[0].Status = "Available"
	// Test GetNewAvailableLessons returns nothing when both lessons are "Available"
	availableLessons := lessonService.GetNewAvailableLessons(lessons1, lessons2)
	a.Equal(0, len(availableLessons))
}

func createLessons(teacherID uint32, baseDatetime time.Time, status string, length int) []*Lesson {
	lessons := make([]*Lesson, length)
	now := time.Now().UTC()
	for i := range lessons {
		lessons[i] = &Lesson{
			TeacherID: teacherID,
			Datetime:  baseDatetime.Add(time.Duration(i) * time.Hour),
			Status:    status,
			CreatedAt: now,
			UpdatedAt: now,
		}
	}
	return lessons
}
<commit_msg>Add test case<commit_after>package model

import (
	"strings"
	"testing"
	"time"

	"github.com/oinume/lekcije/server/config"
	"github.com/oinume/lekcije/server/util"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestUpdateLessons(t *testing.T) {
	a := assert.New(t)
	r := require.New(t)

	teacherID := uint32(util.RandomInt(999999))
	datetime := time.Date(2016, 10, 1, 14, 30, 0, 0, config.LocalTimezone())
	lessons := createLessons(teacherID, datetime, "Reserved", 5)

	affected, err := lessonService.UpdateLessons(lessons)
	r.NoError(err)
	a.Equal(int64(5), affected)
	for _, l := range lessons {
		a.NotEqual(uint64(0), l.ID)
		logs, err := lessonStatusLogService.FindAllByLessonID(l.ID)
		r.NoError(err)
		a.Equal(1, len(logs))
	}

	foundLessons, err := lessonService.FindLessons(teacherID, datetime, datetime)
	r.NoError(err)
	a.Equal(len(lessons), len(foundLessons))
	for i := range lessons {
		// TODO: custom enum type
		a.Equal(strings.ToLower(lessons[i].Status), strings.ToLower(foundLessons[i].Status))
	}
}

func TestUpdateLessonsOverwrite(t *testing.T) {
	a := assert.New(t)
	r := require.New(t)

	teacherID := uint32(util.RandomInt(999999))
	datetime := time.Date(2016, 10, 1, 14, 30, 0, 0, config.LocalTimezone())
	lessons := createLessons(teacherID, datetime, "Available", 5)
	affected, err := lessonService.UpdateLessons(lessons)
	r.NoError(err)
	a.EqualValues(len(lessons), affected)

	time.Sleep(1 * time.Second)
	lessons[0].Status = "Reserved"
	affected, err = lessonService.UpdateLessons(lessons)
	r.NoError(err)
	a.EqualValues(1, affected)

	foundLessons, err := lessonService.FindLessons(teacherID, datetime, datetime)
	r.NoError(err)
	a.Equal(strings.ToLower(foundLessons[0].Status), "reserved")

	logs, err := lessonStatusLogService.FindAllByLessonID(foundLessons[0].ID)
	r.NoError(err)
	a.Equal(2, len(logs))
}

func TestUpdateLessonsNoChange(t *testing.T) {
	a := assert.New(t)
	r := require.New(t)

	teacherID := uint32(util.RandomInt(999999))
	datetime := time.Date(2016, 10, 1, 14, 30, 0, 0, config.LocalTimezone())
	lessons := createLessons(teacherID, datetime, "Available", 5)
	affected, err := lessonService.UpdateLessons(lessons)
	r.NoError(err)
	a.EqualValues(len(lessons), affected)

	affected, err = lessonService.UpdateLessons(lessons)
	r.NoError(err)
	a.EqualValues(0, affected)

	foundLessons, err := lessonService.FindLessons(teacherID, datetime, datetime)
	r.NoError(err)
	a.Equal(strings.ToLower(foundLessons[0].Status), "available")

	logs, err := lessonStatusLogService.FindAllByLessonID(foundLessons[0].ID)
	r.NoError(err)
	a.Equal(1, len(logs))
}

func TestGetNewAvailableLessons1(t *testing.T) {
	a := assert.New(t)

	datetime := time.Date(2016, 10, 1, 14, 30, 0, 0, config.LocalTimezone())
	lessons1 := createLessons(1, datetime, "Reserved", 3)
	lessons2 := createLessons(1, datetime, "Reserved", 3)
	lessons2[1].Status = "Available"
	// Test GetNewAvailableLessons returns a lesson when new lesson is "Available"
	availableLessons := lessonService.GetNewAvailableLessons(lessons1, lessons2)
	a.Equal(1, len(availableLessons))
	a.Equal(datetime.Add(1*time.Hour), availableLessons[0].Datetime)
}

func TestGetNewAvailableLessons2(t *testing.T) {
	a := assert.New(t)

	datetime := time.Date(2016, 10, 1, 14, 30, 0, 0, config.LocalTimezone())
	lessons1 := createLessons(1, datetime, "Reserved", 3)
	lessons2 := createLessons(1, datetime, "Reserved", 3)
	lessons1[0].Status = "Available"
	lessons2[0].Status = "Available"
	// Test GetNewAvailableLessons returns nothing when both lessons are "Available"
	availableLessons := lessonService.GetNewAvailableLessons(lessons1, lessons2)
	a.Equal(0, len(availableLessons))
}

func createLessons(teacherID uint32, baseDatetime time.Time, status string, length int) []*Lesson {
	lessons := make([]*Lesson, length)
	now := time.Now().UTC()
	for i := range lessons {
		lessons[i] = &Lesson{
			TeacherID: teacherID,
			Datetime:  baseDatetime.Add(time.Duration(i) * time.Hour),
			Status:    status,
			CreatedAt: now,
			UpdatedAt: now,
		}
	}
	return lessons
}
<|endoftext|>


<commit_before><commit_msg>Fix error code and argument order (#36)<commit_after><|endoftext|>


<commit_before>package activitypub

import (
	"reflect"
	"testing"
	"time"
)

func TestCreateActivityNew(t *testing.T) {
	var testValue = ObjectID("test")
	var now time.Time

	c := CreateActivityNew(testValue, nil, nil)
	now = time.Now()
	if c.Activity.ID != testValue {
		t.Errorf("Activity Id '%v' different than expected '%v'", c.Activity.ID, testValue)
	}
	if c.Activity.Type != CreateType {
		t.Errorf("Activity Type '%v' different than expected '%v'", c.Activity.Type, CreateType)
	}
	if now.Sub(c.Published).Round(time.Millisecond) != 0 {
		t.Errorf("Published time '%v' different than expected '%v'", c.Published, now)
	}

	testValue = ObjectID("my:note")
	n := ObjectNew("my:note", NoteType)
	b := PersonNew("bob")

	c1 := CreateActivityNew(testValue, *b, n)
	now = time.Now()
	if c1.Activity.ID != testValue {
		t.Errorf("Activity Id '%v' different than expected '%v'", c1.Activity.ID, testValue)
	}
	if c1.Activity.Type != CreateType {
		t.Errorf("Activity Type '%v' different than expected '%v'", c1.Activity.Type, CreateType)
	}
	if now.Sub(c.Published).Round(time.Millisecond) != 0 {
		t.Errorf("Published time '%v' different than expected '%v'", c1.Published, now)
	}
	if !reflect.DeepEqual(c1.Activity.Actor.Object().ID, b.Object().ID) {
		t.Errorf("Actor \n'%#v'\n\n different than expected \n\n'%#v'", c1.Activity.Actor.Object().ID, b.Object().ID)
	}
	if !reflect.DeepEqual(c1.Activity.Object.(apObject), n) {
		t.Errorf("Object \n'%#v'\n\n different than expected \n\n'%#v'\n", c1.Activity.Object, n)
	}
}
<commit_msg>Adding more assertions for CreateActivity tests<commit_after>package activitypub

import (
	"reflect"
	"testing"
	"time"
)

func TestCreateActivityNew(t *testing.T) {
	var testValue = ObjectID("test")
	var now time.Time

	c := CreateActivityNew(testValue, nil, nil)
	now = time.Now()
	if c.Activity.ID != testValue {
		t.Errorf("Activity Id '%v' different than expected '%v'", c.Activity.ID, testValue)
	}
	if c.Activity.Type != CreateType {
		t.Errorf("Activity Type '%v' different than expected '%v'", c.Activity.Type, CreateType)
	}
	if now.Sub(c.Published).Round(time.Millisecond) != 0 {
		t.Errorf("Published time '%v' different than expected '%v'", c.Published, now)
	}

	testValue = ObjectID("my:note")
	n := ObjectNew("my:note", NoteType)
	b := PersonNew("bob")

	c1 := CreateActivityNew(testValue, *b, n)
	now = time.Now()
	if c1.Activity.ID != testValue {
		t.Errorf("Activity Id '%v' different than expected '%v'", c1.Activity.ID, testValue)
	}
	if c1.Activity.Type != CreateType {
		t.Errorf("Activity Type '%v' different than expected '%v'", c1.Activity.Type, CreateType)
	}
	if now.Sub(c.Published).Round(time.Millisecond) != 0 {
		t.Errorf("Published time '%v' different than expected '%v'", c1.Published, now)
	}
	if c1.Activity.Actor.Object().ID != b.ID {
		t.Errorf("Actor ID %q different than expected %q", c1.Activity.Actor.Object().ID, b.ID)
	}
	if !reflect.DeepEqual(c1.Activity.Actor.Object(), b.Object()) {
		t.Errorf("Actor %#v different than expected %#v", c1.Activity.Actor.Object(), b.Object())
	}
	if !reflect.DeepEqual(c1.Activity.Actor, *b) {
		t.Errorf("Actor %#v\n\n different than expected\n\n %#v", c1.Activity.Actor, *b)
	}
	if c1.Activity.Object.Object().ID != n.ID {
		t.Errorf("Object %q different than expected %q", c1.Activity.Object.Object().ID, n.ID)
	}
	if !reflect.DeepEqual(c1.Activity.Object.Object(), n.Object()) {
		t.Errorf("Object %#v different than expected %#v", c1.Activity.Object.Object(), n.Object())
	}
	if !reflect.DeepEqual(c1.Activity.Object, n) {
		t.Errorf("Object %#v different than expected %#v", c1.Activity.Object, n)
	}
}
<|endoftext|>


<commit_before>package server_test

import (
	"context"

	"github.com/cri-o/cri-o/internal/oci"
	"github.com/cri-o/cri-o/server/cri/types"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	specs "github.com/opencontainers/runtime-spec/specs-go"
)

// The actual test suite
var _ = t.Describe("ListPodSandbox", func() {
	// Prepare the sut
	BeforeEach(func() {
		beforeEach()
		setupSUT()
	})

	AfterEach(afterEach)

	t.Describe("ListPodSandbox", func() {
		It("should succeed", func() {
			// Given
			Expect(sut.AddSandbox(testSandbox)).To(BeNil())
			testContainer.SetState(&oci.ContainerState{
				State: specs.State{Status: oci.ContainerStateRunning},
			})
			testSandbox.SetCreated()
			Expect(testSandbox.SetInfraContainer(testContainer)).To(BeNil())

			// When
			response, err := sut.ListPodSandbox(context.Background(),
				&types.ListPodSandboxRequest{})

			// Then
			Expect(err).To(BeNil())
			Expect(response).NotTo(BeNil())
			Expect(len(response.Items)).To(BeEquivalentTo(1))
		})

		It("should succeed without infra container", func() {
			// Given
			Expect(sut.AddSandbox(testSandbox)).To(BeNil())
			testSandbox.SetCreated()

			// When
			response, err := sut.ListPodSandbox(context.Background(),
				&types.ListPodSandboxRequest{})

			// Then
			Expect(err).To(BeNil())
			Expect(response).NotTo(BeNil())
			// the sandbox is created, and even though it has no infra container, it should be displayed
			Expect(len(response.Items)).To(Equal(1))
		})

		It("should skip not created sandboxes", func() {
			// Given
			Expect(sut.AddSandbox(testSandbox)).To(BeNil())
			Expect(testSandbox.SetInfraContainer(testContainer)).To(BeNil())

			// When
			response, err := sut.ListPodSandbox(context.Background(),
				&types.ListPodSandboxRequest{})

			// Then
			Expect(err).To(BeNil())
			Expect(response).NotTo(BeNil())
			Expect(len(response.Items)).To(BeZero())
		})

		It("should succeed with filter", func() {
			// Given
			mockDirs(testManifest)
			createDummyState()
			Expect(sut.LoadSandbox(context.Background(), sandboxID)).To(BeNil())

			// When
			response, err := sut.ListPodSandbox(context.Background(),
				&types.ListPodSandboxRequest{Filter: &types.PodSandboxFilter{
					ID: sandboxID,
				}})

			// Then
			Expect(err).To(BeNil())
			Expect(response).NotTo(BeNil())
			Expect(len(response.Items)).To(BeEquivalentTo(1))
		})

		It("should succeed with filter for state", func() {
			// Given
			mockDirs(testManifest)
			createDummyState()
			Expect(sut.LoadSandbox(context.Background(), sandboxID)).To(BeNil())

			// When
			response, err := sut.ListPodSandbox(context.Background(),
				&types.ListPodSandboxRequest{Filter: &types.PodSandboxFilter{
					ID: sandboxID,
					State: &types.PodSandboxStateValue{
						State: types.PodSandboxStateSandboxReady,
					},
				}})

			// Then
			Expect(err).To(BeNil())
			Expect(response).NotTo(BeNil())
			Expect(len(response.Items)).To(BeZero())
		})

		It("should succeed with filter for label", func() {
			// Given
			mockDirs(testManifest)
			createDummyState()
			Expect(sut.LoadSandbox(context.Background(), sandboxID)).To(BeNil())

			// When
			response, err := sut.ListPodSandbox(context.Background(),
				&types.ListPodSandboxRequest{Filter: &types.PodSandboxFilter{
					ID:            sandboxID,
					LabelSelector: map[string]string{"label": "value"},
				}})

			// Then
			Expect(err).To(BeNil())
			Expect(response).NotTo(BeNil())
			Expect(len(response.Items)).To(BeZero())
		})

		It("should succeed with filter but when not finding id", func() {
			// Given
			Expect(sut.AddSandbox(testSandbox)).To(BeNil())

			// When
			response, err := sut.ListPodSandbox(context.Background(),
				&types.ListPodSandboxRequest{Filter: &types.PodSandboxFilter{
					ID: sandboxID,
				}})

			// Then
			Expect(err).To(BeNil())
			Expect(response).NotTo(BeNil())
			Expect(len(response.Items)).To(BeZero())
		})
	})
})
<commit_msg>Fix unit tests<commit_after>package server_test

import (
	"context"

	"github.com/cri-o/cri-o/internal/oci"
	"github.com/cri-o/cri-o/server/cri/types"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	specs "github.com/opencontainers/runtime-spec/specs-go"
)

// The actual test suite
var _ = t.Describe("ListPodSandbox", func() {
	// Prepare the sut
	BeforeEach(func() {
		beforeEach()
		setupSUT()
	})

	AfterEach(afterEach)

	t.Describe("ListPodSandbox", func() {
		It("should succeed", func() {
			// Given
			Expect(sut.AddSandbox(testSandbox)).To(BeNil())
			testContainer.SetState(&oci.ContainerState{
				State: specs.State{Status: oci.ContainerStateRunning},
			})
			testSandbox.SetCreated()
			Expect(testSandbox.SetInfraContainer(testContainer)).To(BeNil())

			// When
			response, err := sut.ListPodSandbox(context.Background(),
				&types.ListPodSandboxRequest{})

			// Then
			Expect(err).To(BeNil())
			Expect(response).NotTo(BeNil())
			Expect(len(response.Items)).To(BeEquivalentTo(1))
		})

		It("should succeed without infra container", func() {
			// Given
			Expect(sut.AddSandbox(testSandbox)).To(BeNil())
			testSandbox.SetCreated()

			// When
			response, err := sut.ListPodSandbox(context.Background(),
				&types.ListPodSandboxRequest{})

			// Then
			Expect(err).To(BeNil())
			Expect(response).NotTo(BeNil())
			// the sandbox is created, and even though it has no infra container, it should be displayed
			Expect(len(response.Items)).To(Equal(1))
		})

		It("should skip not created sandboxes", func() {
			// Given
			Expect(sut.AddSandbox(testSandbox)).To(BeNil())
			Expect(testSandbox.SetInfraContainer(testContainer)).To(BeNil())

			// When
			response, err := sut.ListPodSandbox(context.Background(),
				&types.ListPodSandboxRequest{})

			// Then
			Expect(err).To(BeNil())
			Expect(response).NotTo(BeNil())
			Expect(len(response.Items)).To(BeZero())
		})

		It("should succeed with filter", func() {
			// Given
			mockDirs(testManifest)
			createDummyState()
			_, err := sut.LoadSandbox(context.Background(), sandboxID)
			Expect(err).To(BeNil())

			// When
			response, err := sut.ListPodSandbox(context.Background(),
				&types.ListPodSandboxRequest{Filter: &types.PodSandboxFilter{
					ID: sandboxID,
				}})

			// Then
			Expect(err).To(BeNil())
			Expect(response).NotTo(BeNil())
			Expect(len(response.Items)).To(BeEquivalentTo(1))
		})

		It("should succeed with filter for state", func() {
			// Given
			mockDirs(testManifest)
			createDummyState()
			_, err := sut.LoadSandbox(context.Background(), sandboxID)
			Expect(err).To(BeNil())

			// When
			response, err := sut.ListPodSandbox(context.Background(),
				&types.ListPodSandboxRequest{Filter: &types.PodSandboxFilter{
					ID: sandboxID,
					State: &types.PodSandboxStateValue{
						State: types.PodSandboxStateSandboxReady,
					},
				}})

			// Then
			Expect(err).To(BeNil())
			Expect(response).NotTo(BeNil())
			Expect(len(response.Items)).To(BeZero())
		})

		It("should succeed with filter for label", func() {
			// Given
			mockDirs(testManifest)
			createDummyState()
			_, err := sut.LoadSandbox(context.Background(), sandboxID)
			Expect(err).To(BeNil())

			// When
			response, err := sut.ListPodSandbox(context.Background(),
				&types.ListPodSandboxRequest{Filter: &types.PodSandboxFilter{
					ID:            sandboxID,
					LabelSelector: map[string]string{"label": "value"},
				}})

			// Then
			Expect(err).To(BeNil())
			Expect(response).NotTo(BeNil())
			Expect(len(response.Items)).To(BeZero())
		})

		It("should succeed with filter but when not finding id", func() {
			// Given
			Expect(sut.AddSandbox(testSandbox)).To(BeNil())

			// When
			response, err := sut.ListPodSandbox(context.Background(),
				&types.ListPodSandboxRequest{Filter: &types.PodSandboxFilter{
					ID: sandboxID,
				}})

			// Then
			Expect(err).To(BeNil())
			Expect(response).NotTo(BeNil())
			Expect(len(response.Items)).To(BeZero())
		})
	})
})
<|endoftext|>


<commit_before>package detector

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"strconv"
	"sync"
	"time"

	log "github.com/golang/glog"
	mesos "github.com/mesos/mesos-go/mesosproto"
	"github.com/mesos/mesos-go/upid"
	"golang.org/x/net/context"
)

const (
	mesosHttpClientTimeout  = 10 * time.Second //TODO(jdef) configurable via fiag?
	mesosLeaderSyncInterval = 30 * time.Second //TODO(jdef) configurable via fiag?
	defaultMesosMasterPort  = 5050
)

type Standalone struct {
	ch         chan *mesos.MasterInfo
	client     *http.Client
	tr         *http.Transport
	pollOnce   sync.Once
	initial    *mesos.MasterInfo
	done       chan struct{}
	cancelOnce sync.Once
}

// Create a new stand alone master detector.
func NewStandalone(mi *mesos.MasterInfo) *Standalone {
	log.V(2).Infof("creating new standalone detector for %+v", mi)
	ch := make(chan *mesos.MasterInfo)
	tr := &http.Transport{}
	return &Standalone{
		ch: ch,
		client: &http.Client{
			Transport: tr,
			Timeout:   mesosHttpClientTimeout,
		},
		tr:      tr,
		initial: mi,
		done:    make(chan struct{}),
	}
}

func (s *Standalone) String() string {
	return fmt.Sprintf("{initial: %+v}", s.initial)
}

// Detecting the new master.
func (s *Standalone) Detect(o MasterChanged) error {
	log.V(2).Info("Detect()")
	s.pollOnce.Do(func() {
		log.V(1).Info("spinning up asyc master detector poller")
		go s.poller()
	})
	if o != nil {
		log.V(1).Info("spawning asyc master detector listener")
		go func() {
			log.V(2).Infof("waiting for polled to send updates")
		pollWaiter:
			for {
				select {
				case mi, ok := <-s.ch:
					if !ok {
						break pollWaiter
					}
					log.V(1).Infof("detected master change: %+v", mi)
					o.Notify(mi)
				case <-s.done:
					return
				}
			}
			o.Notify(nil)
		}()
	} else {
		log.Warningf("detect called with a nil master change listener")
	}
	return nil
}

func (s *Standalone) Done() <-chan struct{} {
	return s.done
}

func (s *Standalone) Cancel() {
	s.cancelOnce.Do(func() { close(s.done) })
}

// poll for changes to master leadership
func (s *Standalone) poller() {
	if s.initial == nil {
		log.Warningf("aborting master poller since initial master info is nil")
		return
	}
	//TODO(jdef) we could attempt to unpack IP addres if Host=""
	if s.initial.Hostname != nil && len(*s.initial.Hostname) == 0 {
		log.Warningf("aborted mater poller since initial master info has no host")
		return
	}
	addr := *s.initial.Hostname
	port := uint32(defaultMesosMasterPort)
	if s.initial.Port != nil && *s.initial.Port != 0 {
		port = *s.initial.Port
	}
	addr = net.JoinHostPort(addr, strconv.Itoa(int(port)))
	log.V(1).Infof("polling for master leadership at '%v'", addr)
	var lastpid *upid.UPID
	for {
		startedAt := time.Now()
		ctx, cancel := context.WithTimeout(context.Background(), mesosLeaderSyncInterval)
		if pid, err := s.fetchPid(ctx, addr); err == nil {
			if !pid.Equal(lastpid) {
				log.V(2).Infof("detected leadership change from '%v' to '%v'", lastpid, pid)
				lastpid = pid
				elapsed := time.Now().Sub(startedAt)
				mi := CreateMasterInfo(pid)
				select {
				case s.ch <- mi:
					// noop
				case <-time.After(mesosLeaderSyncInterval - elapsed):
					// no one heard the master change, oh well - poll again
					cancel()
					continue
				case <-s.done:
					cancel()
					return
				}
			} else {
				log.V(2).Infof("no change to master leadership: '%v'", lastpid)
			}
		} else {
			log.Warning(err)
		}
		remaining := mesosLeaderSyncInterval - time.Now().Sub(startedAt)
		log.V(3).Infof("master leader poller sleeping for %v", remaining)
		time.Sleep(remaining)
		cancel()
	}
}

func (s *Standalone) fetchPid(ctx context.Context, address string) (*upid.UPID, error) {
	//TODO(jdef) need better address parsing, for now assume host, or host:port format
	//TODO(jdef) need SSL support
	uri := fmt.Sprintf("http://%s/state.json", address)
	req, err := http.NewRequest("GET", uri, nil)
	if err != nil {
		return nil, err
	}
	var pid *upid.UPID
	err = s.httpDo(ctx, req, func(res *http.Response, err error) error {
		if err != nil {
			return err
		}
		defer res.Body.Close()
		if res.StatusCode != 200 {
			return fmt.Errorf("HTTP request failed with code %d: %v", res.StatusCode, res.Status)
		}
		blob, err1 := ioutil.ReadAll(res.Body)
		if err1 != nil {
			return err1
		}
		log.V(3).Infof("Got mesos state, content length %v", len(blob))
		type State struct {
			Leader string `json:"leader"` // ex: master(1)@10.22.211.18:5050
		}
		state := &State{}
		err = json.Unmarshal(blob, state)
		if err != nil {
			return err
		}
		pid, err = upid.Parse(state.Leader)
		return err
	})
	return pid, err
}

type responseHandler func(*http.Response, error) error

// hacked from https://blog.golang.org/context
func (s *Standalone) httpDo(ctx context.Context, req *http.Request, f responseHandler) error {
	// Run the HTTP request in a goroutine and pass the response to f.
	ch := make(chan error, 1)
	go func() { ch <- f(s.client.Do(req)) }()
	select {
	case <-ctx.Done():
		s.tr.CancelRequest(req)
		<-ch // Wait for f to return.
		return ctx.Err()
	case err := <-ch:
		return err
	}
}
<commit_msg>added docs for standalone poller<commit_after>package detector

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"strconv"
	"sync"
	"time"

	log "github.com/golang/glog"
	mesos "github.com/mesos/mesos-go/mesosproto"
	"github.com/mesos/mesos-go/upid"
	"golang.org/x/net/context"
)

const (
	mesosHttpClientTimeout  = 10 * time.Second //TODO(jdef) configurable via fiag?
	mesosLeaderSyncInterval = 30 * time.Second //TODO(jdef) configurable via fiag?
	defaultMesosMasterPort  = 5050
)

type Standalone struct {
	ch         chan *mesos.MasterInfo
	client     *http.Client
	tr         *http.Transport
	pollOnce   sync.Once
	initial    *mesos.MasterInfo
	done       chan struct{}
	cancelOnce sync.Once
}

// Create a new stand alone master detector.
func NewStandalone(mi *mesos.MasterInfo) *Standalone {
	log.V(2).Infof("creating new standalone detector for %+v", mi)
	ch := make(chan *mesos.MasterInfo)
	tr := &http.Transport{}
	return &Standalone{
		ch: ch,
		client: &http.Client{
			Transport: tr,
			Timeout:   mesosHttpClientTimeout,
		},
		tr:      tr,
		initial: mi,
		done:    make(chan struct{}),
	}
}

func (s *Standalone) String() string {
	return fmt.Sprintf("{initial: %+v}", s.initial)
}

// Detecting the new master.
func (s *Standalone) Detect(o MasterChanged) error {
	log.V(2).Info("Detect()")
	s.pollOnce.Do(func() {
		log.V(1).Info("spinning up asyc master detector poller")
		go s.poller()
	})
	if o != nil {
		log.V(1).Info("spawning asyc master detector listener")
		go func() {
			log.V(2).Infof("waiting for polled to send updates")
		pollWaiter:
			for {
				select {
				case mi, ok := <-s.ch:
					if !ok {
						break pollWaiter
					}
					log.V(1).Infof("detected master change: %+v", mi)
					o.Notify(mi)
				case <-s.done:
					return
				}
			}
			o.Notify(nil)
		}()
	} else {
		log.Warningf("detect called with a nil master change listener")
	}
	return nil
}

func (s *Standalone) Done() <-chan struct{} {
	return s.done
}

func (s *Standalone) Cancel() {
	s.cancelOnce.Do(func() { close(s.done) })
}

// poll for changes to master leadership via current leader's /state.json endpoint.
// we start with the `initial` leader, aborting if none was specified. thereafter,
// the `leader` property of the state.json is used to identify the next leader that
// should be polled.
//
// TODO(jdef) somehow determine all masters in cluster from the state.json?
//
func (s *Standalone) poller() {
	if s.initial == nil {
		log.Warningf("aborting master poller since initial master info is nil")
		return
	}
	//TODO(jdef) we could attempt to unpack IP addres if Host=""
	if s.initial.Hostname != nil && len(*s.initial.Hostname) == 0 {
		log.Warningf("aborted mater poller since initial master info has no host")
		return
	}
	addr := *s.initial.Hostname
	port := uint32(defaultMesosMasterPort)
	if s.initial.Port != nil && *s.initial.Port != 0 {
		port = *s.initial.Port
	}
	addr = net.JoinHostPort(addr, strconv.Itoa(int(port)))
	log.V(1).Infof("polling for master leadership at '%v'", addr)
	var lastpid *upid.UPID
	for {
		startedAt := time.Now()
		ctx, cancel := context.WithTimeout(context.Background(), mesosLeaderSyncInterval)
		if pid, err := s.fetchPid(ctx, addr); err == nil {
			if !pid.Equal(lastpid) {
				log.V(2).Infof("detected leadership change from '%v' to '%v'", lastpid, pid)
				lastpid = pid
				elapsed := time.Now().Sub(startedAt)
				mi := CreateMasterInfo(pid)
				select {
				case s.ch <- mi:
					// noop
				case <-time.After(mesosLeaderSyncInterval - elapsed):
					// no one heard the master change, oh well - poll again
					cancel()
					continue
				case <-s.done:
					cancel()
					return
				}
			} else {
				log.V(2).Infof("no change to master leadership: '%v'", lastpid)
			}
		} else {
			log.Warning(err)
		}
		remaining := mesosLeaderSyncInterval - time.Now().Sub(startedAt)
		log.V(3).Infof("master leader poller sleeping for %v", remaining)
		time.Sleep(remaining)
		cancel()
	}
}

func (s *Standalone) fetchPid(ctx context.Context, address string) (*upid.UPID, error) {
	//TODO(jdef) need better address parsing, for now assume host, or host:port format
	//TODO(jdef) need SSL support
	uri := fmt.Sprintf("http://%s/state.json", address)
	req, err := http.NewRequest("GET", uri, nil)
	if err != nil {
		return nil, err
	}
	var pid *upid.UPID
	err = s.httpDo(ctx, req, func(res *http.Response, err error) error {
		if err != nil {
			return err
		}
		defer res.Body.Close()
		if res.StatusCode != 200 {
			return fmt.Errorf("HTTP request failed with code %d: %v", res.StatusCode, res.Status)
		}
		blob, err1 := ioutil.ReadAll(res.Body)
		if err1 != nil {
			return err1
		}
		log.V(3).Infof("Got mesos state, content length %v", len(blob))
		type State struct {
			Leader string `json:"leader"` // ex: master(1)@10.22.211.18:5050
		}
		state := &State{}
		err = json.Unmarshal(blob, state)
		if err != nil {
			return err
		}
		pid, err = upid.Parse(state.Leader)
		return err
	})
	return pid, err
}

type responseHandler func(*http.Response, error) error

// hacked from https://blog.golang.org/context
func (s *Standalone) httpDo(ctx context.Context, req *http.Request, f responseHandler) error {
	// Run the HTTP request in a goroutine and pass the response to f.
	ch := make(chan error, 1)
	go func() { ch <- f(s.client.Do(req)) }()
	select {
	case <-ctx.Done():
		s.tr.CancelRequest(req)
		<-ch // Wait for f to return.
		return ctx.Err()
	case err := <-ch:
		return err
	}
}
<|endoftext|>


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

// Package client provides client support for interacting with a serverless
// log.
//
// See the /cmd/client package in this repo for an example of using this.
package client

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strconv"

	"github.com/golang/glog"
	"github.com/google/trillian-examples/formats/log"
	"github.com/google/trillian-examples/serverless/api"
	"github.com/google/trillian-examples/serverless/api/layout"
	"github.com/google/trillian/merkle"
	"github.com/google/trillian/merkle/compact"
	"github.com/google/trillian/merkle/hashers"
	"github.com/google/trillian/merkle/logverifier"
	"golang.org/x/mod/sumdb/note"
)

// Fetcher is the signature of a function which can retrieve arbitrary files from
// a log's data storage, via whatever appropriate mechanism.
// The path parameter is relative to the root of the log storage.
type Fetcher func(path string) ([]byte, error)

func fetchCheckpointAndParse(f Fetcher, v note.Verifier) (*log.Checkpoint, []byte, error) {
	cpRaw, err := f(layout.CheckpointPath)
	if err != nil {
		return nil, nil, err
	}
	n, err := note.Open(cpRaw, note.VerifierList(v))
	if err != nil {
		glog.Exitf("failed to open Checkpoint: %q", err)
	}
	cp := log.Checkpoint{}
	if _, err := cp.Unmarshal([]byte(n.Text)); err != nil {
		glog.V(1).Infof("Bad checkpoint: %q", cpRaw)
		return nil, nil, fmt.Errorf("failed to unmarshal checkpoint: %w", err)
	}
	return &cp, cpRaw, nil
}

// ProofBuilder knows how to build inclusion and consistency proofs from tiles.
// Since the tiles commit only to immutable nodes, the job of building proofs is slightly
// more complex as proofs can touch "ephemeral" nodes, so these need to be synthesized.
type ProofBuilder struct {
	cp        log.Checkpoint
	nodeCache nodeCache
	h         compact.HashFn
}

// NewProofBuilder creates a new ProofBuilder object for a given tree size.
// The returned ProofBuilder can be re-used for proofs related to a given tree size, but
// it is not thread-safe and should not be accessed concurrently.
func NewProofBuilder(cp log.Checkpoint, h compact.HashFn, f Fetcher) (*ProofBuilder, error) {
	tf := newTileFetcher(f, cp.Size)
	pb := &ProofBuilder{
		cp:        cp,
		nodeCache: newNodeCache(tf, cp.Size),
		h:         h,
	}

	hashes, err := FetchRangeNodes(cp.Size, tf)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch range nodes: %w", err)
	}
	// Create a compact range which represents the state of the log.
	r, err := (&compact.RangeFactory{Hash: h}).NewRange(0, cp.Size, hashes)
	if err != nil {
		return nil, err
	}

	// Recreate the root hash so that:
	// a) we validate the self-integrity of the log state, and
	// b) we calculate (and cache) and ephemeral nodes present in the tree,
	//    this is important since they could be required by proofs.
	sr, err := r.GetRootHash(pb.nodeCache.SetEphemeralNode)
	if err != nil {
		return nil, err
	}
	if !bytes.Equal(cp.Hash, sr) {
		return nil, fmt.Errorf("invalid checkpoint hash %x, expected %x", cp.Hash, sr)
	}
	return pb, nil
}

// InclusionProof constructs an inclusion proof for the leaf at index in a tree of
// the given size.
// This function uses the passed-in function to retrieve tiles containing any log tree
// nodes necessary to build the proof.
func (pb *ProofBuilder) InclusionProof(index uint64) ([][]byte, error) {
	nodes, err := merkle.CalcInclusionProofNodeAddresses(int64(pb.cp.Size), int64(index), int64(pb.cp.Size))
	if err != nil {
		return nil, fmt.Errorf("failed to calculate inclusion proof node list: %w", err)
	}

	ret := make([][]byte, 0)
	// TODO(al) parallelise this.
	for _, n := range nodes {
		h, err := pb.nodeCache.GetNode(n.ID)
		if err != nil {
			return nil, fmt.Errorf("failed to get node (%v): %w", n.ID, err)
		}
		ret = append(ret, h)
	}
	return ret, nil
}

// ConsistencyProof constructs a consistency proof between the two passed in tree sizes.
// This function uses the passed-in function to retrieve tiles containing any log tree
// nodes necessary to build the proof.
func (pb *ProofBuilder) ConsistencyProof(smaller, larger uint64) ([][]byte, error) {
	nodes, err := merkle.CalcConsistencyProofNodeAddresses(int64(smaller), int64(larger), int64(pb.cp.Size))
	if err != nil {
		return nil, fmt.Errorf("failed to calculate consistency proof node list: %w", err)
	}

	hashes := make([][]byte, 0)
	// TODO(al) parallelise this.
	for _, n := range nodes {
		h, err := pb.nodeCache.GetNode(n.ID)
		if err != nil {
			return nil, fmt.Errorf("failed to get node (%v): %w", n.ID, err)
		}
		hashes = append(hashes, h)
	}
	return hashes, nil
}

// FetchRangeNodes returns the set of nodes representing the compact range covering
// a log of size s.
func FetchRangeNodes(s uint64, gt GetTileFunc) ([][]byte, error) {
	nc := newNodeCache(gt, s)
	nIDs := compact.RangeNodes(0, s)
	ret := make([][]byte, len(nIDs))
	for i, n := range nIDs {
		h, err := nc.GetNode(n)
		if err != nil {
			return nil, err
		}
		ret[i] = h
	}
	return ret, nil
}

// FetchLeafHashes fetches N consecutive leaf hashes starting with the leaf at index first.
func FetchLeafHashes(f Fetcher, first, N, logSize uint64) ([][]byte, error) {
	nc := newNodeCache(newTileFetcher(f, logSize), logSize)
	ret := make([][]byte, N, 0)
	for i, seq := uint64(0), first; i < N; i, seq = i+1, seq+1 {
		nID := compact.NodeID{Level: 0, Index: seq}
		h, err := nc.GetNode(nID)
		if err != nil {
			return nil, fmt.Errorf("failed to fetch node %v: %v", nID, err)
		}
		ret = append(ret, h)
	}
	return ret, nil
}

// nodeCache hides the tiles abstraction away, and improves
// performance by caching tiles it's seen.
// Not threadsafe, and intended to be only used throughout the course
// of a single request.
type nodeCache struct {
	logSize   uint64
	ephemeral map[compact.NodeID][]byte
	tiles     map[tileKey]api.Tile
	getTile   GetTileFunc
}

// GetTileFunc is the signature of a function which knows how to fetch a
// specific tile.
type GetTileFunc func(level, index uint64) (*api.Tile, error)

// tileKey is used as a key in nodeCache's tile map.
type tileKey struct {
	tileLevel uint64
	tileIndex uint64
}

// newNodeCache creates a new nodeCache instance for a given log size.
func newNodeCache(f GetTileFunc, logSize uint64) nodeCache {
	return nodeCache{
		logSize:   logSize,
		ephemeral: make(map[compact.NodeID][]byte),
		tiles:     make(map[tileKey]api.Tile),
		getTile:   f,
	}
}

// SetEphemeralNode stored a derived "ephemeral" tree node.
func (n *nodeCache) SetEphemeralNode(id compact.NodeID, h []byte) {
	n.ephemeral[id] = h
}

// GetNode returns the internal log tree node hash for the specified node ID.
// A previously set ephemeral node will be returned if id matches, otherwise
// the tile containing the requested node will be fetched and cached, and the
// node hash returned.
func (n *nodeCache) GetNode(id compact.NodeID) ([]byte, error) {
	// First check for ephemeral nodes:
	if e := n.ephemeral[id]; len(e) != 0 {
		return e, nil
	}
	// Otherwise look in fetched tiles:
	tileLevel, tileIndex, nodeLevel, nodeIndex := layout.NodeCoordsToTileAddress(uint64(id.Level), uint64(id.Index))
	tKey := tileKey{tileLevel, tileIndex}
	t, ok := n.tiles[tKey]
	if !ok {
		tile, err := n.getTile(tileLevel, tileIndex)
		if err != nil {
			return nil, fmt.Errorf("failed to fetch tile: %w", err)
		}
		t = *tile
		n.tiles[tKey] = *tile
	}
	node := t.Nodes[api.TileNodeKey(nodeLevel, nodeIndex)]
	if node == nil {
		return nil, fmt.Errorf("node %v (tile coords [%d,%d]/[%d,%d]) unknown", id, tileLevel, tileIndex, nodeLevel, nodeIndex)
	}
	return node, nil
}

// newTileFetcher returns a GetTileFunc based on the passed in Fetcher and log size.
func newTileFetcher(f Fetcher, logSize uint64) GetTileFunc {
	return func(level, index uint64) (*api.Tile, error) {
		tileSize := layout.PartialTileSize(level, index, logSize)
		p := filepath.Join(layout.TilePath("", level, index, tileSize))
		t, err := f(p)
		if err != nil {
			if !errors.Is(err, os.ErrNotExist) {
				return nil, fmt.Errorf("failed to read tile at %q: %w", p, err)
			}
			return nil, err
		}

		var tile api.Tile
		if err := tile.UnmarshalText(t); err != nil {
			return nil, fmt.Errorf("failed to parse tile: %w", err)
		}
		return &tile, nil
	}
}

// LookupIndex fetches the leafhash->seq mapping file from the log, and returns
// its parsed contents.
func LookupIndex(f Fetcher, lh []byte) (uint64, error) {
	p := filepath.Join(layout.LeafPath("", lh))
	sRaw, err := f(p)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return 0, fmt.Errorf("leafhash unknown (%w)", err)
		}
		return 0, fmt.Errorf("failed to fetch leafhash->seq file: %w", err)
	}
	return strconv.ParseUint(string(sRaw), 16, 64)
}

// LogStateTracker represents a client-side view of a target log's state.
// This tracker handles verification that updates to the tracked log state are
// consistent with previously seen states.
type LogStateTracker struct {
	Hasher   hashers.LogHasher
	Verifier logverifier.LogVerifier
	Fetcher  Fetcher

	// LatestConsistentRaw holds the raw bytes of the latest proven-consistent
	// LogState seen by this tracker.
	LatestConsistentRaw []byte

	// LatestConsistent is the deserialised form of LatestConsistentRaw
	LatestConsistent log.Checkpoint
	CpSigVerifier    note.Verifier
}

// NewLogStateTracker creates a newly initialised tracker.
// If a serialised LogState representation is provided then this is used as the
// initial tracked state, otherwise a log state is fetched from the target log.
func NewLogStateTracker(f Fetcher, h hashers.LogHasher, checkpointRaw []byte, nV note.Verifier) (LogStateTracker, error) {

	ret := LogStateTracker{
		Fetcher:          f,
		Hasher:           h,
		Verifier:         logverifier.New(h),
		LatestConsistent: log.Checkpoint{},
		CpSigVerifier:    nV,
	}
	if len(checkpointRaw) > 0 {
		ret.LatestConsistentRaw = checkpointRaw
		if _, err := ret.LatestConsistent.Unmarshal(checkpointRaw); err != nil {
			return ret, err
		}
		return ret, nil
	}
	return ret, ret.Update()
}

// ErrInconsistency should be returned when there has been an error proving consistency
// between log states.
// The raw log state representations are included as-returned by the target log, this
// ensures that evidence of inconsistent log updates are available to the caller of
// the method(s) returning this error.
type ErrInconsistency struct {
	SmallerRaw []byte
	LargerRaw  []byte
	Proof      [][]byte

	Wrapped error
}

func (e ErrInconsistency) Unwrap() error {
	return e.Wrapped
}

func (e ErrInconsistency) Error() string {
	return fmt.Sprintf("log consistency check failed: %s", e.Wrapped)
}

// Update attempts to update the local view of the target log's state.
// If a more recent logstate is found, this method will attempt to prove
// that it is consistent with the local state before updating the tracker's
// view.
func (lst *LogStateTracker) Update() error {
	c, cRaw, err := fetchCheckpointAndParse(lst.Fetcher, lst.CpSigVerifier)
	if err != nil {
		return err
	}
	if lst.LatestConsistent.Size > 0 {
		if c.Size > lst.LatestConsistent.Size {
			builder, err := NewProofBuilder(*c, lst.Hasher.HashChildren, lst.Fetcher)
			if err != nil {
				return fmt.Errorf("failed to create proof builder: %w", err)
			}
			p, err := builder.ConsistencyProof(lst.LatestConsistent.Size, c.Size)
			if err != nil {
				return err
			}
			glog.V(1).Infof("Built consistency proof %x", p)
			if err := lst.Verifier.VerifyConsistencyProof(int64(lst.LatestConsistent.Size), int64(c.Size), lst.LatestConsistent.Hash, c.Hash, p); err != nil {
				return ErrInconsistency{
					SmallerRaw: lst.LatestConsistentRaw,
					LargerRaw:  cRaw,
					Proof:      p,
					Wrapped:    err,
				}
			}
		}
	}
	lst.LatestConsistentRaw, lst.LatestConsistent = cRaw, *c
	return nil
}
<commit_msg>Don't exit on a bad checkpoint<commit_after>// Copyright 2021 Google LLC. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package client provides client support for interacting with a serverless
// log.
//
// See the /cmd/client package in this repo for an example of using this.
package client

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strconv"

	"github.com/golang/glog"
	"github.com/google/trillian-examples/formats/log"
	"github.com/google/trillian-examples/serverless/api"
	"github.com/google/trillian-examples/serverless/api/layout"
	"github.com/google/trillian/merkle"
	"github.com/google/trillian/merkle/compact"
	"github.com/google/trillian/merkle/hashers"
	"github.com/google/trillian/merkle/logverifier"
	"golang.org/x/mod/sumdb/note"
)

// Fetcher is the signature of a function which can retrieve arbitrary files from
// a log's data storage, via whatever appropriate mechanism.
// The path parameter is relative to the root of the log storage.
type Fetcher func(path string) ([]byte, error)

func fetchCheckpointAndParse(f Fetcher, v note.Verifier) (*log.Checkpoint, []byte, error) {
	cpRaw, err := f(layout.CheckpointPath)
	if err != nil {
		return nil, nil, err
	}
	n, err := note.Open(cpRaw, note.VerifierList(v))
	if err != nil {
		return nil, nil, fmt.Errorf("failed to open Checkpoint: %v", err)
	}
	cp := log.Checkpoint{}
	if _, err := cp.Unmarshal([]byte(n.Text)); err != nil {
		glog.V(1).Infof("Bad checkpoint: %q", cpRaw)
		return nil, nil, fmt.Errorf("failed to unmarshal checkpoint: %w", err)
	}
	return &cp, cpRaw, nil
}

// ProofBuilder knows how to build inclusion and consistency proofs from tiles.
// Since the tiles commit only to immutable nodes, the job of building proofs is slightly
// more complex as proofs can touch "ephemeral" nodes, so these need to be synthesized.
type ProofBuilder struct {
	cp        log.Checkpoint
	nodeCache nodeCache
	h         compact.HashFn
}

// NewProofBuilder creates a new ProofBuilder object for a given tree size.
// The returned ProofBuilder can be re-used for proofs related to a given tree size, but
// it is not thread-safe and should not be accessed concurrently.
func NewProofBuilder(cp log.Checkpoint, h compact.HashFn, f Fetcher) (*ProofBuilder, error) {
	tf := newTileFetcher(f, cp.Size)
	pb := &ProofBuilder{
		cp:        cp,
		nodeCache: newNodeCache(tf, cp.Size),
		h:         h,
	}

	hashes, err := FetchRangeNodes(cp.Size, tf)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch range nodes: %w", err)
	}
	// Create a compact range which represents the state of the log.
	r, err := (&compact.RangeFactory{Hash: h}).NewRange(0, cp.Size, hashes)
	if err != nil {
		return nil, err
	}

	// Recreate the root hash so that:
	// a) we validate the self-integrity of the log state, and
	// b) we calculate (and cache) and ephemeral nodes present in the tree,
	//    this is important since they could be required by proofs.
	sr, err := r.GetRootHash(pb.nodeCache.SetEphemeralNode)
	if err != nil {
		return nil, err
	}
	if !bytes.Equal(cp.Hash, sr) {
		return nil, fmt.Errorf("invalid checkpoint hash %x, expected %x", cp.Hash, sr)
	}
	return pb, nil
}

// InclusionProof constructs an inclusion proof for the leaf at index in a tree of
// the given size.
// This function uses the passed-in function to retrieve tiles containing any log tree
// nodes necessary to build the proof.
func (pb *ProofBuilder) InclusionProof(index uint64) ([][]byte, error) {
	nodes, err := merkle.CalcInclusionProofNodeAddresses(int64(pb.cp.Size), int64(index), int64(pb.cp.Size))
	if err != nil {
		return nil, fmt.Errorf("failed to calculate inclusion proof node list: %w", err)
	}

	ret := make([][]byte, 0)
	// TODO(al) parallelise this.
	for _, n := range nodes {
		h, err := pb.nodeCache.GetNode(n.ID)
		if err != nil {
			return nil, fmt.Errorf("failed to get node (%v): %w", n.ID, err)
		}
		ret = append(ret, h)
	}
	return ret, nil
}

// ConsistencyProof constructs a consistency proof between the two passed in tree sizes.
// This function uses the passed-in function to retrieve tiles containing any log tree
// nodes necessary to build the proof.
func (pb *ProofBuilder) ConsistencyProof(smaller, larger uint64) ([][]byte, error) {
	nodes, err := merkle.CalcConsistencyProofNodeAddresses(int64(smaller), int64(larger), int64(pb.cp.Size))
	if err != nil {
		return nil, fmt.Errorf("failed to calculate consistency proof node list: %w", err)
	}

	hashes := make([][]byte, 0)
	// TODO(al) parallelise this.
	for _, n := range nodes {
		h, err := pb.nodeCache.GetNode(n.ID)
		if err != nil {
			return nil, fmt.Errorf("failed to get node (%v): %w", n.ID, err)
		}
		hashes = append(hashes, h)
	}
	return hashes, nil
}

// FetchRangeNodes returns the set of nodes representing the compact range covering
// a log of size s.
func FetchRangeNodes(s uint64, gt GetTileFunc) ([][]byte, error) {
	nc := newNodeCache(gt, s)
	nIDs := compact.RangeNodes(0, s)
	ret := make([][]byte, len(nIDs))
	for i, n := range nIDs {
		h, err := nc.GetNode(n)
		if err != nil {
			return nil, err
		}
		ret[i] = h
	}
	return ret, nil
}

// FetchLeafHashes fetches N consecutive leaf hashes starting with the leaf at index first.
func FetchLeafHashes(f Fetcher, first, N, logSize uint64) ([][]byte, error) {
	nc := newNodeCache(newTileFetcher(f, logSize), logSize)
	ret := make([][]byte, N, 0)
	for i, seq := uint64(0), first; i < N; i, seq = i+1, seq+1 {
		nID := compact.NodeID{Level: 0, Index: seq}
		h, err := nc.GetNode(nID)
		if err != nil {
			return nil, fmt.Errorf("failed to fetch node %v: %v", nID, err)
		}
		ret = append(ret, h)
	}
	return ret, nil
}

// nodeCache hides the tiles abstraction away, and improves
// performance by caching tiles it's seen.
// Not threadsafe, and intended to be only used throughout the course
// of a single request.
type nodeCache struct {
	logSize   uint64
	ephemeral map[compact.NodeID][]byte
	tiles     map[tileKey]api.Tile
	getTile   GetTileFunc
}

// GetTileFunc is the signature of a function which knows how to fetch a
// specific tile.
type GetTileFunc func(level, index uint64) (*api.Tile, error)

// tileKey is used as a key in nodeCache's tile map.
type tileKey struct {
	tileLevel uint64
	tileIndex uint64
}

// newNodeCache creates a new nodeCache instance for a given log size.
func newNodeCache(f GetTileFunc, logSize uint64) nodeCache {
	return nodeCache{
		logSize:   logSize,
		ephemeral: make(map[compact.NodeID][]byte),
		tiles:     make(map[tileKey]api.Tile),
		getTile:   f,
	}
}

// SetEphemeralNode stored a derived "ephemeral" tree node.
func (n *nodeCache) SetEphemeralNode(id compact.NodeID, h []byte) {
	n.ephemeral[id] = h
}

// GetNode returns the internal log tree node hash for the specified node ID.
// A previously set ephemeral node will be returned if id matches, otherwise
// the tile containing the requested node will be fetched and cached, and the
// node hash returned.
func (n *nodeCache) GetNode(id compact.NodeID) ([]byte, error) {
	// First check for ephemeral nodes:
	if e := n.ephemeral[id]; len(e) != 0 {
		return e, nil
	}
	// Otherwise look in fetched tiles:
	tileLevel, tileIndex, nodeLevel, nodeIndex := layout.NodeCoordsToTileAddress(uint64(id.Level), uint64(id.Index))
	tKey := tileKey{tileLevel, tileIndex}
	t, ok := n.tiles[tKey]
	if !ok {
		tile, err := n.getTile(tileLevel, tileIndex)
		if err != nil {
			return nil, fmt.Errorf("failed to fetch tile: %w", err)
		}
		t = *tile
		n.tiles[tKey] = *tile
	}
	node := t.Nodes[api.TileNodeKey(nodeLevel, nodeIndex)]
	if node == nil {
		return nil, fmt.Errorf("node %v (tile coords [%d,%d]/[%d,%d]) unknown", id, tileLevel, tileIndex, nodeLevel, nodeIndex)
	}
	return node, nil
}

// newTileFetcher returns a GetTileFunc based on the passed in Fetcher and log size.
func newTileFetcher(f Fetcher, logSize uint64) GetTileFunc {
	return func(level, index uint64) (*api.Tile, error) {
		tileSize := layout.PartialTileSize(level, index, logSize)
		p := filepath.Join(layout.TilePath("", level, index, tileSize))
		t, err := f(p)
		if err != nil {
			if !errors.Is(err, os.ErrNotExist) {
				return nil, fmt.Errorf("failed to read tile at %q: %w", p, err)
			}
			return nil, err
		}

		var tile api.Tile
		if err := tile.UnmarshalText(t); err != nil {
			return nil, fmt.Errorf("failed to parse tile: %w", err)
		}
		return &tile, nil
	}
}

// LookupIndex fetches the leafhash->seq mapping file from the log, and returns
// its parsed contents.
func LookupIndex(f Fetcher, lh []byte) (uint64, error) {
	p := filepath.Join(layout.LeafPath("", lh))
	sRaw, err := f(p)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return 0, fmt.Errorf("leafhash unknown (%w)", err)
		}
		return 0, fmt.Errorf("failed to fetch leafhash->seq file: %w", err)
	}
	return strconv.ParseUint(string(sRaw), 16, 64)
}

// LogStateTracker represents a client-side view of a target log's state.
// This tracker handles verification that updates to the tracked log state are
// consistent with previously seen states.
type LogStateTracker struct {
	Hasher   hashers.LogHasher
	Verifier logverifier.LogVerifier
	Fetcher  Fetcher

	// LatestConsistentRaw holds the raw bytes of the latest proven-consistent
	// LogState seen by this tracker.
	LatestConsistentRaw []byte

	// LatestConsistent is the deserialised form of LatestConsistentRaw
	LatestConsistent log.Checkpoint
	CpSigVerifier    note.Verifier
}

// NewLogStateTracker creates a newly initialised tracker.
// If a serialised LogState representation is provided then this is used as the
// initial tracked state, otherwise a log state is fetched from the target log.
func NewLogStateTracker(f Fetcher, h hashers.LogHasher, checkpointRaw []byte, nV note.Verifier) (LogStateTracker, error) {

	ret := LogStateTracker{
		Fetcher:          f,
		Hasher:           h,
		Verifier:         logverifier.New(h),
		LatestConsistent: log.Checkpoint{},
		CpSigVerifier:    nV,
	}
	if len(checkpointRaw) > 0 {
		ret.LatestConsistentRaw = checkpointRaw
		if _, err := ret.LatestConsistent.Unmarshal(checkpointRaw); err != nil {
			return ret, err
		}
		return ret, nil
	}
	return ret, ret.Update()
}

// ErrInconsistency should be returned when there has been an error proving consistency
// between log states.
// The raw log state representations are included as-returned by the target log, this
// ensures that evidence of inconsistent log updates are available to the caller of
// the method(s) returning this error.
type ErrInconsistency struct {
	SmallerRaw []byte
	LargerRaw  []byte
	Proof      [][]byte

	Wrapped error
}

func (e ErrInconsistency) Unwrap() error {
	return e.Wrapped
}

func (e ErrInconsistency) Error() string {
	return fmt.Sprintf("log consistency check failed: %s", e.Wrapped)
}

// Update attempts to update the local view of the target log's state.
// If a more recent logstate is found, this method will attempt to prove
// that it is consistent with the local state before updating the tracker's
// view.
func (lst *LogStateTracker) Update() error {
	c, cRaw, err := fetchCheckpointAndParse(lst.Fetcher, lst.CpSigVerifier)
	if err != nil {
		return err
	}
	if lst.LatestConsistent.Size > 0 {
		if c.Size > lst.LatestConsistent.Size {
			builder, err := NewProofBuilder(*c, lst.Hasher.HashChildren, lst.Fetcher)
			if err != nil {
				return fmt.Errorf("failed to create proof builder: %w", err)
			}
			p, err := builder.ConsistencyProof(lst.LatestConsistent.Size, c.Size)
			if err != nil {
				return err
			}
			glog.V(1).Infof("Built consistency proof %x", p)
			if err := lst.Verifier.VerifyConsistencyProof(int64(lst.LatestConsistent.Size), int64(c.Size), lst.LatestConsistent.Hash, c.Hash, p); err != nil {
				return ErrInconsistency{
					SmallerRaw: lst.LatestConsistentRaw,
					LargerRaw:  cRaw,
					Proof:      p,
					Wrapped:    err,
				}
			}
		}
	}
	lst.LatestConsistentRaw, lst.LatestConsistent = cRaw, *c
	return nil
}
<|endoftext|>


<commit_before>package vtrace_test

import (
	"strings"
	"testing"

	"veyron.io/veyron/veyron2/context"
	"veyron.io/veyron/veyron2/ipc"
	"veyron.io/veyron/veyron2/ipc/stream"
	"veyron.io/veyron/veyron2/naming"
	"veyron.io/veyron/veyron2/security"
	"veyron.io/veyron/veyron2/vlog"
	"veyron.io/veyron/veyron2/vtrace"

	_ "veyron.io/veyron/veyron/lib/tcp"
	iipc "veyron.io/veyron/veyron/runtimes/google/ipc"
	"veyron.io/veyron/veyron/runtimes/google/ipc/stream/manager"
	tnaming "veyron.io/veyron/veyron/runtimes/google/testing/mocks/naming"
	truntime "veyron.io/veyron/veyron/runtimes/google/testing/mocks/runtime"
	ivtrace "veyron.io/veyron/veyron/runtimes/google/vtrace"
)

// We need a special way to create contexts for tests.  We
// can't create a real runtime in the runtime implementation
// so we use a fake one that panics if used.  The runtime
// implementation should not ever use the Runtime from a context.
func testContext() context.T {
	return iipc.InternalNewContext(&truntime.PanicRuntime{})
}

func TestNewFromContext(t *testing.T) {
	c0 := testContext()
	c1, s1 := ivtrace.WithNewSpan(c0, "s1")
	c2, s2 := ivtrace.WithNewSpan(c1, "s2")
	c3, s3 := ivtrace.WithNewSpan(c2, "s3")
	expected := map[context.T]vtrace.Span{
		c0: nil,
		c1: s1,
		c2: s2,
		c3: s3,
	}
	for ctx, expectedSpan := range expected {
		if s := ivtrace.FromContext(ctx); s != expectedSpan {
			t.Errorf("Wrong span for ctx %v.  Got %v, want %v", c0, s, expectedSpan)
		}
	}
}

type fakeAuthorizer int

func (fakeAuthorizer) Authorize(security.Context) error {
	return nil
}

type testServer struct {
	sm           stream.Manager
	ns           naming.Namespace
	name         string
	child        string
	stop         func() error
	forceCollect bool
}

func (c *testServer) Run(ctx ipc.ServerContext) error {
	if c.forceCollect {
		ivtrace.FromContext(ctx).Trace().ForceCollect()
	}

	client, err := iipc.InternalNewClient(c.sm, c.ns)
	if err != nil {
		vlog.Error(err)
		return err
	}

	ivtrace.FromContext(ctx).Annotate(c.name + "-begin")

	if c.child != "" {
		var call ipc.Call
		if call, err = client.StartCall(ctx, c.child, "Run", []interface{}{}); err != nil {
			vlog.Error(err)
			return err
		}
		var outerr error
		if err = call.Finish(&outerr); err != nil {
			vlog.Error(err)
			return err
		}
		if outerr != nil {
			vlog.Error(outerr)
			return outerr
		}
	}
	ivtrace.FromContext(ctx).Annotate(c.name + "-end")

	return nil
}

func makeTestServer(ns naming.Namespace, name, child string, forceCollect bool) (*testServer, error) {
	sm := manager.InternalNew(naming.FixedRoutingID(0x111111111))
	ctx := testContext()
	s, err := iipc.InternalNewServer(ctx, sm, ns, nil)
	if err != nil {
		return nil, err
	}
	if _, err := s.Listen(ipc.ListenSpec{Protocol: "tcp", Address: "127.0.0.1:0"}); err != nil {
		return nil, err
	}

	c := &testServer{
		sm:           sm,
		ns:           ns,
		name:         name,
		child:        child,
		stop:         s.Stop,
		forceCollect: forceCollect,
	}

	if err := s.Serve(name, c, fakeAuthorizer(0)); err != nil {
		return nil, err
	}

	return c, nil
}

func summary(span *vtrace.SpanRecord) string {
	summary := span.Name
	if len(span.Annotations) > 0 {
		msgs := []string{}
		for _, annotation := range span.Annotations {
			msgs = append(msgs, annotation.Message)
		}
		summary += ": " + strings.Join(msgs, ", ")
	}
	return summary
}

func expectSequence(t *testing.T, trace vtrace.TraceRecord, expectedSpans []string) {
	// It's okay to have additional spans - someone may have inserted
	// additional spans for more debugging.
	if got, want := len(trace.Spans), len(expectedSpans); got < want {
		t.Errorf("Found %d spans, want %d", got, want)
	}

	spans := map[string]*vtrace.SpanRecord{}
	summaries := []string{}
	for i := range trace.Spans {
		span := &trace.Spans[i]

		// All spans should have a start.
		if span.Start == 0 {
			t.Errorf("span missing start: %#v", span)
		}
		// All spans except the root should have an end.
		if span.Name != "" && span.End == 0 {
			t.Errorf("span missing end: %#v", span)
			if span.Start >= span.End {
				t.Errorf("span end should be after start: %#v", span)
			}
		}

		summary := summary(span)
		summaries = append(summaries, summary)
		spans[summary] = span
	}

	for i := range expectedSpans {
		child, ok := spans[expectedSpans[i]]
		if !ok {
			t.Errorf("expected span %s not found in %#v", expectedSpans[i], summaries)
			continue
		}
		if i == 0 {
			continue
		}
		parent, ok := spans[expectedSpans[i-1]]
		if !ok {
			t.Errorf("expected span %s not found in %#v", expectedSpans[i-1], summaries)
			continue
		}
		if child.Parent != parent.ID {
			t.Errorf("%v should be a child of %v, but it's not.", child, parent)
		}
	}
}

func runCallChain(t *testing.T, ctx context.T, force1, force2 bool) {
	sm := manager.InternalNew(naming.FixedRoutingID(0x555555555))
	ns := tnaming.NewSimpleNamespace()

	client, err := iipc.InternalNewClient(sm, ns)
	if err != nil {
		t.Error(err)
	}

	c1, err := makeTestServer(ns, "c1", "c2", force1)
	if err != nil {
		t.Fatal("Can't start server:", err)
	}
	defer c1.stop()

	c2, err := makeTestServer(ns, "c2", "", force2)
	if err != nil {
		t.Fatal("Can't start server:", err)
	}
	defer c2.stop()

	call, err := client.StartCall(ctx, "c1", "Run", []interface{}{})
	if err != nil {
		t.Fatal("can't call: ", err)
	}
	var outerr error
	if err = call.Finish(&outerr); err != nil {
		t.Error(err)
	}
	if outerr != nil {
		t.Error(outerr)
	}
}

// TestCancellationPropagation tests that cancellation propogates along an
// RPC call chain without user intervention.
func TestTraceAcrossRPCs(t *testing.T) {
	ctx, span := ivtrace.WithNewSpan(testContext(), "")
	span.Trace().ForceCollect()
	span.Annotate("c0-begin")

	runCallChain(t, ctx, false, false)

	span.Annotate("c0-end")

	expectedSpans := []string{
		": c0-begin, c0-end",
		"<client>\"c1\".Run",
		"\"\".Run: c1-begin, c1-end",
		"<client>\"c2\".Run",
		"\"\".Run: c2-begin, c2-end",
	}
	expectSequence(t, span.Trace().Record(), expectedSpans)
}

// TestCancellationPropagationLateForce tests that cancellation propogates along an
// RPC call chain when tracing is initiated by someone deep in the call chain.
func TestTraceAcrossRPCsLateForce(t *testing.T) {
	ctx, span := ivtrace.WithNewSpan(testContext(), "")
	span.Annotate("c0-begin")

	runCallChain(t, ctx, false, true)

	span.Annotate("c0-end")

	expectedSpans := []string{
		": c0-end",
		"<client>\"c1\".Run",
		"\"\".Run: c1-end",
		"<client>\"c2\".Run",
		"\"\".Run: c2-begin, c2-end",
	}
	expectSequence(t, span.Trace().Record(), expectedSpans)
}
<commit_msg>veyron/runtimes/google/vtrace: Fix flaky vtrace test.<commit_after>package vtrace_test

import (
	"bytes"
	"strings"
	"testing"

	"veyron.io/veyron/veyron2/context"
	"veyron.io/veyron/veyron2/ipc"
	"veyron.io/veyron/veyron2/ipc/stream"
	"veyron.io/veyron/veyron2/naming"
	"veyron.io/veyron/veyron2/security"
	"veyron.io/veyron/veyron2/vlog"
	"veyron.io/veyron/veyron2/vtrace"

	_ "veyron.io/veyron/veyron/lib/tcp"
	iipc "veyron.io/veyron/veyron/runtimes/google/ipc"
	"veyron.io/veyron/veyron/runtimes/google/ipc/stream/manager"
	tnaming "veyron.io/veyron/veyron/runtimes/google/testing/mocks/naming"
	truntime "veyron.io/veyron/veyron/runtimes/google/testing/mocks/runtime"
	ivtrace "veyron.io/veyron/veyron/runtimes/google/vtrace"
)

// We need a special way to create contexts for tests.  We
// can't create a real runtime in the runtime implementation
// so we use a fake one that panics if used.  The runtime
// implementation should not ever use the Runtime from a context.
func testContext() context.T {
	return iipc.InternalNewContext(&truntime.PanicRuntime{})
}

func TestNewFromContext(t *testing.T) {
	c0 := testContext()
	c1, s1 := ivtrace.WithNewSpan(c0, "s1")
	c2, s2 := ivtrace.WithNewSpan(c1, "s2")
	c3, s3 := ivtrace.WithNewSpan(c2, "s3")
	expected := map[context.T]vtrace.Span{
		c0: nil,
		c1: s1,
		c2: s2,
		c3: s3,
	}
	for ctx, expectedSpan := range expected {
		if s := ivtrace.FromContext(ctx); s != expectedSpan {
			t.Errorf("Wrong span for ctx %v.  Got %v, want %v", c0, s, expectedSpan)
		}
	}
}

type fakeAuthorizer int

func (fakeAuthorizer) Authorize(security.Context) error {
	return nil
}

type testServer struct {
	sm           stream.Manager
	ns           naming.Namespace
	name         string
	child        string
	stop         func() error
	forceCollect bool
}

func (c *testServer) Run(ctx ipc.ServerContext) error {
	if c.forceCollect {
		ivtrace.FromContext(ctx).Trace().ForceCollect()
	}

	client, err := iipc.InternalNewClient(c.sm, c.ns)
	if err != nil {
		vlog.Error(err)
		return err
	}

	ivtrace.FromContext(ctx).Annotate(c.name + "-begin")

	if c.child != "" {
		var call ipc.Call
		if call, err = client.StartCall(ctx, c.child, "Run", []interface{}{}); err != nil {
			vlog.Error(err)
			return err
		}
		var outerr error
		if err = call.Finish(&outerr); err != nil {
			vlog.Error(err)
			return err
		}
		if outerr != nil {
			vlog.Error(outerr)
			return outerr
		}
	}
	ivtrace.FromContext(ctx).Annotate(c.name + "-end")

	return nil
}

func makeTestServer(ns naming.Namespace, name, child string, forceCollect bool) (*testServer, error) {
	sm := manager.InternalNew(naming.FixedRoutingID(0x111111111))
	ctx := testContext()
	s, err := iipc.InternalNewServer(ctx, sm, ns, nil)
	if err != nil {
		return nil, err
	}
	if _, err := s.Listen(ipc.ListenSpec{Protocol: "tcp", Address: "127.0.0.1:0"}); err != nil {
		return nil, err
	}

	c := &testServer{
		sm:           sm,
		ns:           ns,
		name:         name,
		child:        child,
		stop:         s.Stop,
		forceCollect: forceCollect,
	}

	if err := s.Serve(name, c, fakeAuthorizer(0)); err != nil {
		return nil, err
	}

	return c, nil
}

func summary(span *vtrace.SpanRecord) string {
	summary := span.Name
	if len(span.Annotations) > 0 {
		msgs := []string{}
		for _, annotation := range span.Annotations {
			msgs = append(msgs, annotation.Message)
		}
		summary += ": " + strings.Join(msgs, ", ")
	}
	return summary
}

func traceString(trace *vtrace.TraceRecord) string {
	var b bytes.Buffer
	vtrace.FormatTrace(&b, trace, nil)
	return b.String()
}

func expectSequence(t *testing.T, trace vtrace.TraceRecord, expectedSpans []string) {
	// It's okay to have additional spans - someone may have inserted
	// additional spans for more debugging.
	if got, want := len(trace.Spans), len(expectedSpans); got < want {
		t.Errorf("Found %d spans, want %d", got, want)
	}

	spans := map[string]*vtrace.SpanRecord{}
	summaries := []string{}
	for i := range trace.Spans {
		span := &trace.Spans[i]

		// All spans should have a start.
		if span.Start == 0 {
			t.Errorf("span missing start: %x, %s", span.ID[12:], traceString(&trace))
		}
		// All spans except the root should have a valid end.
		// TODO(mattr): For now I'm also skipping connectFlow and
		// vc.HandshakeDialedVC spans because the ws endpoints are
		// currently non-deterministic in terms of whether they fail
		// before the test ends or not.  In the future it will be
		// configurable whether we listen on ws or not and then we should
		// adjust the test to not listen and remove this check.
		if span.Name != "" &&
			span.Name != "<client>connectFlow" &&
			span.Name != "vc.HandshakeDialedVC" {
			if span.End == 0 {
				t.Errorf("span missing end: %x, %s", span.ID[12:], traceString(&trace))
			} else if span.Start >= span.End {
				t.Errorf("span end should be after start: %x, %s", span.ID[12:], traceString(&trace))
			}
		}

		summary := summary(span)
		summaries = append(summaries, summary)
		spans[summary] = span
	}

	for i := range expectedSpans {
		child, ok := spans[expectedSpans[i]]
		if !ok {
			t.Errorf("expected span %s not found in %#v", expectedSpans[i], summaries)
			continue
		}
		if i == 0 {
			continue
		}
		parent, ok := spans[expectedSpans[i-1]]
		if !ok {
			t.Errorf("expected span %s not found in %#v", expectedSpans[i-1], summaries)
			continue
		}
		if child.Parent != parent.ID {
			t.Errorf("%v should be a child of %v, but it's not.", child, parent)
		}
	}
}

func runCallChain(t *testing.T, ctx context.T, force1, force2 bool) {
	sm := manager.InternalNew(naming.FixedRoutingID(0x555555555))
	ns := tnaming.NewSimpleNamespace()

	client, err := iipc.InternalNewClient(sm, ns)
	if err != nil {
		t.Error(err)
	}

	c1, err := makeTestServer(ns, "c1", "c2", force1)
	if err != nil {
		t.Fatal("Can't start server:", err)
	}
	defer c1.stop()

	c2, err := makeTestServer(ns, "c2", "", force2)
	if err != nil {
		t.Fatal("Can't start server:", err)
	}
	defer c2.stop()

	call, err := client.StartCall(ctx, "c1", "Run", []interface{}{})
	if err != nil {
		t.Fatal("can't call: ", err)
	}
	var outerr error
	if err = call.Finish(&outerr); err != nil {
		t.Error(err)
	}
	if outerr != nil {
		t.Error(outerr)
	}
}

// TestCancellationPropagation tests that cancellation propogates along an
// RPC call chain without user intervention.
func TestTraceAcrossRPCs(t *testing.T) {
	ctx, span := ivtrace.WithNewSpan(testContext(), "")
	span.Trace().ForceCollect()
	span.Annotate("c0-begin")

	runCallChain(t, ctx, false, false)

	span.Annotate("c0-end")

	expectedSpans := []string{
		": c0-begin, c0-end",
		"<client>\"c1\".Run",
		"\"\".Run: c1-begin, c1-end",
		"<client>\"c2\".Run",
		"\"\".Run: c2-begin, c2-end",
	}
	expectSequence(t, span.Trace().Record(), expectedSpans)
}

// TestCancellationPropagationLateForce tests that cancellation propogates along an
// RPC call chain when tracing is initiated by someone deep in the call chain.
func TestTraceAcrossRPCsLateForce(t *testing.T) {
	ctx, span := ivtrace.WithNewSpan(testContext(), "")
	span.Annotate("c0-begin")

	runCallChain(t, ctx, false, true)

	span.Annotate("c0-end")

	expectedSpans := []string{
		": c0-end",
		"<client>\"c1\".Run",
		"\"\".Run: c1-end",
		"<client>\"c2\".Run",
		"\"\".Run: c2-begin, c2-end",
	}
	expectSequence(t, span.Trace().Record(), expectedSpans)
}
<|endoftext|>


<commit_before>/*
 * Copyright © 2015-2018 Aeneas Rekkas <aeneas+oss@aeneas.io>
 *
 * 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.
 *
 * @author		Aeneas Rekkas <aeneas+oss@aeneas.io>
 * @copyright 	2015-2018 Aeneas Rekkas <aeneas+oss@aeneas.io>
 * @license 	Apache-2.0
 */

package cli

import (
	"fmt"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"

	"github.com/ory/hydra/client"
	"github.com/ory/hydra/config"
	"github.com/ory/hydra/consent"
	"github.com/ory/hydra/jwk"
	"github.com/ory/hydra/oauth2"
	"github.com/ory/hydra/pkg"
)

type MigrateHandler struct {
	c *config.Config
}

func newMigrateHandler(c *config.Config) *MigrateHandler {
	return &MigrateHandler{c: c}
}

type schemaCreator interface {
	CreateSchemas() (int, error)
}

func (h *MigrateHandler) connectToSql(dsn string) (*sqlx.DB, error) {
	var db *sqlx.DB

	u, err := url.Parse(dsn)
	if err != nil {
		return nil, errors.Errorf("Could not parse DATABASE_URL: %s", err)
	}

	if err := pkg.Retry(h.c.GetLogger(), time.Second*15, time.Minute*2, func() error {
		if u.Scheme == "mysql" {
			dsn = strings.Replace(dsn, "mysql://", "", -1)
		}

		if db, err = sqlx.Open(u.Scheme, dsn); err != nil {
			return errors.Errorf("Could not connect to SQL: %s", err)
		} else if err := db.Ping(); err != nil {
			return errors.Errorf("Could not connect to SQL: %s", err)
		}

		return nil
	}); err != nil {
		return nil, err
	}

	return db, nil
}

func getDBUrl(cmd *cobra.Command, args []string, position int) (dburl string) {
	if readFromEnv, _ := cmd.Flags().GetBool("read-from-env"); readFromEnv {
		if len(viper.GetString("DATABASE_URL")) == 0 {
			fmt.Println(cmd.UsageString())
			fmt.Println("")
			fmt.Println("When using flag -e, environment variable DATABASE_URL must be set")
			return
		}
		dburl = viper.GetString("DATABASE_URL")
	} else {
		if len(args) < position {
			fmt.Println(cmd.UsageString())
			return
		}
		dburl = args[position]
	}
	if dburl == "" {
		fmt.Println(cmd.UsageString())
		return
	}
	return
}

func (h *MigrateHandler) MigrateSecret(cmd *cobra.Command, args []string) {
	dburl := getDBUrl(cmd, args, 0)
	if dburl == "" {
		return
	}

	db, err := h.connectToSql(dburl)
	if err != nil {
		fmt.Printf("An error occurred while connecting to SQL: %s", err)
		os.Exit(1)
		return
	}

	oldSecret := viper.GetString("OLD_SYSTEM_SECRET")
	newSecret := viper.GetString("NEW_SYSTEM_SECRET")

	if len(oldSecret) == 0 {
		fmt.Println("You did not specify the old system secret, please set environment variable OLD_SYSTEM_SECRET.")
		os.Exit(1)
		return
	} else if len(newSecret) == 0 {
		fmt.Println("You did not specify the old system secret, please set environment variable NEW_SYSTEM_SECRET.")
		os.Exit(1)
		return
	}

	fmt.Println("Rotating encryption keys for JSON Web Key storage...")
	manager := jwk.NewSQLManager(db, []byte(oldSecret))
	if err := manager.RotateKeys(&jwk.AEAD{Key: []byte(newSecret)}); err != nil {
		fmt.Printf("Error \"%s\" occurred while trying to rotate the JSON Web Key storage. All changes have been reverted.", err)
	}
	fmt.Println("Rotating encryption keys for JSON Web Key storage completed successfully!")
	fmt.Printf(`You may now run ORY Hydra with the new system secret. If you wish that old OAuth 2.0 Access and Refres
tokens stay valid, please set environment variable ROTATED_SYSTEM_SECRET to the new secret:

ROTATED_SYSTEM_SECRET=%s hydra serve ...

If you wish that OAuth 2.0 Access and Refresh Tokens issued with the old secret are revoked, simply omit environment variable
ROTATED_SYSTEM_SECRET. This will NOT affect OpenID Connect ID Tokens!
`, newSecret)
}

func (h *MigrateHandler) MigrateSQL(cmd *cobra.Command, args []string) {
	dburl := getDBUrl(cmd, args, 0)
	if dburl == "" {
		return
	}

	db, err := h.connectToSql(dburl)
	if err != nil {
		fmt.Printf("An error occurred while connecting to SQL: %s", err)
		os.Exit(1)
		return
	}

	if err := h.runMigrateSQL(db); err != nil {
		fmt.Printf("An error occurred while running the migrations: %s", err)
		os.Exit(1)
		return
	}
	fmt.Println("Migration successful!")
}

func (h *MigrateHandler) runMigrateSQL(db *sqlx.DB) error {
	var total int
	for k, m := range map[string]schemaCreator{
		"client":  &client.SQLManager{DB: db},
		"oauth2":  &oauth2.FositeSQLStore{DB: db},
		"jwk":     &jwk.SQLManager{DB: db},
		"consent": consent.NewSQLManager(db, nil, nil),
	} {
		fmt.Printf("Applying `%s` SQL migrations...\n", k)
		if num, err := m.CreateSchemas(); err != nil {
			return errors.Wrapf(err, "Could not apply `%s` SQL migrations", k)
		} else {
			fmt.Printf("Applied %d `%s` SQL migrations.\n", num, k)
			total += num
		}
	}

	fmt.Printf("Migration successful! Applied a total of %d SQL migrations.\n", total)
	return nil
}
<commit_msg>cli: Improve migrate error messages<commit_after>/*
 * Copyright © 2015-2018 Aeneas Rekkas <aeneas+oss@aeneas.io>
 *
 * 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.
 *
 * @author		Aeneas Rekkas <aeneas+oss@aeneas.io>
 * @copyright 	2015-2018 Aeneas Rekkas <aeneas+oss@aeneas.io>
 * @license 	Apache-2.0
 */

package cli

import (
	"context"
	"fmt"
	"github.com/ory/x/cmdx"
	"net/url"
	"strings"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"

	"github.com/ory/hydra/client"
	"github.com/ory/hydra/config"
	"github.com/ory/hydra/consent"
	"github.com/ory/hydra/jwk"
	"github.com/ory/hydra/oauth2"
	"github.com/ory/hydra/pkg"
)

type MigrateHandler struct {
	c *config.Config
}

func newMigrateHandler(c *config.Config) *MigrateHandler {
	return &MigrateHandler{c: c}
}

type schemaCreator interface {
	CreateSchemas() (int, error)
}

func (h *MigrateHandler) connectToSql(dsn string) (*sqlx.DB, error) {
	var db *sqlx.DB

	u, err := url.Parse(dsn)
	if err != nil {
		return nil, errors.Errorf("could not parse DATABASE_URL: %s", err)
	}

	if err := pkg.Retry(h.c.GetLogger(), time.Second*15, time.Minute*2, func() error {
		if u.Scheme == "mysql" {
			dsn = strings.Replace(dsn, "mysql://", "", -1)
		}

		if db, err = sqlx.Open(u.Scheme, dsn); err != nil {
			return errors.Errorf("could not connect to SQL: %s", err)
		} else if err := db.Ping(); err != nil {
			return errors.Errorf("could not connect to SQL: %s", err)
		}

		return nil
	}); err != nil {
		return nil, err
	}

	return db, nil
}

func getDBUrl(cmd *cobra.Command, args []string, position int) (dburl string) {
	if readFromEnv, _ := cmd.Flags().GetBool("read-from-env"); readFromEnv {
		if len(viper.GetString("DATABASE_URL")) == 0 {
			fmt.Println(cmd.UsageString())
			fmt.Println("")
			fmt.Println("When using flag -e, environment variable DATABASE_URL must be set")
			return
		}
		dburl = viper.GetString("DATABASE_URL")
	} else {
		if len(args) < position {
			fmt.Println(cmd.UsageString())
			return
		}
		dburl = args[position]
	}
	if dburl == "" {
		fmt.Println(cmd.UsageString())
		return
	}
	return
}

func (h *MigrateHandler) MigrateSecret(cmd *cobra.Command, args []string) {
	dburl := getDBUrl(cmd, args, 0)
	if dburl == "" {
		return
	}

	db, err := h.connectToSql(dburl)
	cmdx.Must(err, "An error occurred while connecting to SQL: %s", err)

	oldSecret := viper.GetString("OLD_SYSTEM_SECRET")
	newSecret := viper.GetString("NEW_SYSTEM_SECRET")

	if len(oldSecret) != 32 {
		cmdx.Fatalf("Value of environment variable OLD_SYSTEM_SECRET has to be exactly 32 characters long but got: %d", len(oldSecret))
	}

	if len(newSecret) != 32 {
		cmdx.Fatalf("Value of environment variable NEW_SYSTEM_SECRET has to be exactly 32 characters long but got: %d", len(oldSecret))
	}

	fmt.Println("Rotating encryption keys for JSON Web Key storage...")

	manager := jwk.NewSQLManager(db, []byte(oldSecret))
	err = manager.RotateKeys(context.TODO(), &jwk.AEAD{Key: []byte(newSecret)})
	cmdx.Must(err, "Unable to rotate JSON Web Keys: %s\nAll changes have been rolled back.", err)

	fmt.Println("Rotating encryption keys for JSON Web Key storage completed successfully!")
	fmt.Printf(`You may now run ORY Hydra with the new system secret. If you wish that old OAuth 2.0 Access and Refres
tokens stay valid, please set environment variable ROTATED_SYSTEM_SECRET to the new secret:

ROTATED_SYSTEM_SECRET=%s hydra serve ...

If you wish that OAuth 2.0 Access and Refresh Tokens issued with the old secret are revoked, simply omit environment variable
ROTATED_SYSTEM_SECRET. This will NOT affect OpenID Connect ID Tokens!
`, newSecret)
}

func (h *MigrateHandler) MigrateSQL(cmd *cobra.Command, args []string) {
	dburl := getDBUrl(cmd, args, 0)
	if dburl == "" {
		return
	}

	db, err := h.connectToSql(dburl)
	cmdx.Must(err, "An error occurred while connecting to SQL: %s", err)

	err = h.runMigrateSQL(db)
	cmdx.Must(err, "An error occurred while running the migrations: %s", err)

	fmt.Println("Migration successful!")
}

func (h *MigrateHandler) runMigrateSQL(db *sqlx.DB) error {
	var total int
	for k, m := range map[string]schemaCreator{
		"client":  &client.SQLManager{DB: db},
		"oauth2":  &oauth2.FositeSQLStore{DB: db},
		"jwk":     &jwk.SQLManager{DB: db},
		"consent": consent.NewSQLManager(db, nil, nil),
	} {
		fmt.Printf("Applying `%s` SQL migrations...\n", k)
		if num, err := m.CreateSchemas(); err != nil {
			return errors.Wrapf(err, "could not apply %s SQL migrations", k)
		} else {
			fmt.Printf("Applied %d `%s` SQL migrations.\n", num, k)
			total += num
		}
	}

	fmt.Printf("Migration successful! Applied a total of %d SQL migrations.\n", total)
	return nil
}
<|endoftext|>