<commit_before>package router

import (
	"testing"

	"time"

	"github.com/rodaine/statstee/bucket"
	"github.com/rodaine/statstee/datagram"
	"github.com/stretchr/testify/assert"
)

func TestRouter_Listen(t *testing.T) {
	t.Parallel()

	c := make(chan datagram.Metric, 3)
	c <- datagram.Metric{
		Name:       "foo.bar",
		Type:       datagram.Counter,
		Value:      1,
		SampleRate: 1,
	}
	c <- datagram.Metric{
		Name:       "fizz.buzz",
		Type:       datagram.Histogram,
		Value:      1,
		SampleRate: 1,
	}
	c <- datagram.Metric{
		Name:       "fizz.buzz",
		Type:       datagram.Histogram,
		Value:      1,
		SampleRate: 1,
	}
	close(c)

	r := New(c)
	r.Listen()
	<-time.NewTimer(100 * time.Millisecond).C

	m := r.Metrics()
	assert.Len(t, m, 2)

	assert.Equal(t, "foo.bar", r.Selected(), "first metric added should be selected")
	assert.Equal(t, "fizz.buzz", m[0].Name, "alphabetized metrics")
}

func TestRouter_SelectedMetric(t *testing.T) {
	t.Parallel()

	c := make(chan datagram.Metric)
	defer close(c)

	r := New(c)
	go r.Listen()

	assert.True(t, bucket.DummyWindow == r.SelectedMetric())
	r.selected = "foo"
	assert.True(t, bucket.DummyWindow == r.SelectedMetric())
	r.selected = ""

	c <- datagram.Metric{
		Name:       "foo.bar",
		Type:       datagram.Counter,
		Value:      1,
		SampleRate: 1,
	}

	assert.False(t, bucket.DummyWindow == r.SelectedMetric())
	assert.NotNil(t, r.SelectedMetric())
}

func TestRouter_PreviousNext(t *testing.T) {
	t.Parallel()

	c := make(chan datagram.Metric)
	defer close(c)

	r := New(c)
	go r.Listen()

	r.Previous()
	assert.Empty(t, r.Selected())
	r.Next()
	assert.Empty(t, r.Selected())

	c <- datagram.Metric{
		Name:       "foo.bar",
		Type:       datagram.Counter,
		Value:      1,
		SampleRate: 1,
	}
	c <- datagram.Metric{
		Name:       "fizz.buzz",
		Type:       datagram.Set,
		Value:      1,
		SampleRate: 1,
	}

	assert.Equal(t, "foo.bar", r.Selected(), "first metric added should be selected")
	r.selected = "not.found"
	r.Previous()
	assert.Equal(t, "fizz.buzz", r.Selected(), "not found metric should default to first")
	r.selected = ""
	r.Next()
	assert.Equal(t, "fizz.buzz", r.Selected(), "not found metric should default to first")

	r.Previous()
	assert.Equal(t, "fizz.buzz", r.Selected(), "if on first item, don't change on previous")
	r.Next()
	assert.Equal(t, "foo.bar", r.Selected())
	r.Next()
	assert.Equal(t, "foo.bar", r.Selected(), "if on last item, don't change on next")
	r.Previous()
	assert.Equal(t, "fizz.buzz", r.Selected())
}
<commit_msg>checkpoint<commit_after>package router

import (
	"testing"

	"time"

	"github.com/rodaine/statstee/bucket"
	"github.com/rodaine/statstee/datagram"
	"github.com/stretchr/testify/assert"
)

func TestRouter_Listen(t *testing.T) {
	t.Parallel()

	c := make(chan datagram.Metric, 3)
	c <- datagram.Metric{
		Name:       "foo.bar",
		Type:       datagram.Counter,
		Value:      1,
		SampleRate: 1,
	}
	c <- datagram.Metric{
		Name:       "fizz.buzz",
		Type:       datagram.Histogram,
		Value:      1,
		SampleRate: 1,
	}
	c <- datagram.Metric{
		Name:       "fizz.buzz",
		Type:       datagram.Histogram,
		Value:      1,
		SampleRate: 1,
	}
	close(c)

	r := New(c)
	r.Listen()
	<-time.NewTimer(100 * time.Millisecond).C

	m := r.Metrics()
	assert.Len(t, m, 2)

	assert.Equal(t, "foo.bar", r.Selected(), "first metric added should be selected")
	assert.Equal(t, "fizz.buzz", m[0].Name, "alphabetized metrics")
}

func TestRouter_SelectedMetric(t *testing.T) {
	t.Parallel()

	c := make(chan datagram.Metric)
	defer close(c)

	r := New(c)
	go r.Listen()

	assert.True(t, bucket.DummyWindow == r.SelectedMetric())
	r.selected = "foo"
	assert.True(t, bucket.DummyWindow == r.SelectedMetric())
	r.selected = ""

	c <- datagram.Metric{
		Name:       "foo.bar",
		Type:       datagram.Counter,
		Value:      1,
		SampleRate: 1,
	}

	assert.False(t, bucket.DummyWindow == r.SelectedMetric())
	assert.NotNil(t, r.SelectedMetric())
}

func TestRouter_PreviousNext(t *testing.T) {
	t.Parallel()

	c := make(chan datagram.Metric, 2)
	r := New(c)

	r.Previous()
	assert.Empty(t, r.Selected())
	r.Next()
	assert.Empty(t, r.Selected())

	c <- datagram.Metric{
		Name:       "foo.bar",
		Type:       datagram.Counter,
		Value:      1,
		SampleRate: 1,
	}
	c <- datagram.Metric{
		Name:       "fizz.buzz",
		Type:       datagram.Set,
		Value:      1,
		SampleRate: 1,
	}
	close(c)
	r.Listen()

	assert.Equal(t, "foo.bar", r.Selected(), "first metric added should be selected")
	r.selected = "not.found"
	r.Previous()
	assert.Equal(t, "fizz.buzz", r.Selected(), "not found metric should default to first")
	r.selected = ""
	r.Next()
	assert.Equal(t, "fizz.buzz", r.Selected(), "not found metric should default to first")

	r.Previous()
	assert.Equal(t, "fizz.buzz", r.Selected(), "if on first item, don't change on previous")
	r.Next()
	assert.Equal(t, "foo.bar", r.Selected())
	r.Next()
	assert.Equal(t, "foo.bar", r.Selected(), "if on last item, don't change on next")
	r.Previous()
	assert.Equal(t, "fizz.buzz", r.Selected())
}
<|endoftext|>


<commit_before>// Copyright 2016 Marc-Antoine Ruel. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

package anim1d

import (
	"encoding/json"
	"log"
	"sync"
	"time"

	"github.com/maruel/dlibox/go/donotuse/devices"
	"github.com/maruel/interrupt"
)

// Pattern is a interface to draw an animated line.
type Pattern interface {
	// NextFrame fills the buffer with the image at this time frame.
	//
	// The image should be derived from timeMS, which is the time since this
	// pattern was started.
	//
	// Calling NextFrame() with a nil pattern is valid. Patterns should be
	// callable without crashing with an object initialized with default values.
	//
	// timeMS will cycle after 49.7 days. The reason it's not using time.Duration
	// is that int64 calculation on ARM is very slow and abysmal on xtensa, which
	// this code is transpiled to.
	NextFrame(pixels Frame, timeMS uint32)
}

// Painter handles the "draw frame, write" loop.
type Painter struct {
	d             devices.Display
	c             chan newPattern
	wg            sync.WaitGroup
	frameDuration time.Duration
}

// SetPattern changes the current pattern to a new one.
//
// The pattern is in JSON encoded format. The function will return an error if
// the encoding is bad. The function is synchronous, it returns only after the
// pattern was effectively set.
func (p *Painter) SetPattern(s string, transition time.Duration) error {
	var pat SPattern
	if err := json.Unmarshal([]byte(s), &pat); err != nil {
		return err
	}
	p.c <- newPattern{pat.Pattern, transition}
	return nil
}

func (p *Painter) Close() error {
	select {
	case p.c <- newPattern{}:
	default:
	}
	close(p.c)
	p.wg.Wait()
	return nil
}

// NewPainter returns a Painter that manages updating the Patterns to the
// strip.
//
// It Assumes the display uses native RGB packed pixels.
func NewPainter(d devices.Display, fps int) *Painter {
	p := &Painter{
		d:             d,
		c:             make(chan newPattern),
		frameDuration: time.Second / time.Duration(fps),
	}
	numLights := d.Bounds().Dx()
	// Tripple buffering.
	cGen := make(chan Frame, 3)
	cWrite := make(chan Frame, cap(cGen))
	for i := 0; i < cap(cGen); i++ {
		cGen <- make(Frame, numLights)
	}
	p.wg.Add(2)
	go p.runPattern(cGen, cWrite)
	go p.runWrite(cGen, cWrite, numLights)
	return p
}

// Private stuff.

var black = &Color{}

type newPattern struct {
	p Pattern
	d time.Duration
}

func (p *Painter) runPattern(cGen <-chan Frame, cWrite chan<- Frame) {
	defer func() {
		// Tell runWrite() to quit.
		for loop := true; loop; {
			select {
			case _, loop = <-cGen:
			default:
				loop = false
			}
		}
		select {
		case cWrite <- nil:
		default:
		}
		close(cWrite)
		p.wg.Done()
	}()

	var root Pattern = black
	var since time.Duration
	for {
		select {
		case newPat, ok := <-p.c:
			if newPat.p == nil || !ok {
				// Request to terminate.
				return
			}

			// New pattern.
			if newPat.d == 0 {
				root = newPat.p
			} else {
				root = &Transition{
					Before:       SPattern{root},
					After:        SPattern{newPat.p},
					OffsetMS:     uint32(since / time.Millisecond),
					TransitionMS: uint32(newPat.d / time.Millisecond),
					Curve:        EaseOut,
				}
			}

		case pixels, ok := <-cGen:
			if !ok {
				return
			}
			for i := range pixels {
				pixels[i] = Color{}
			}
			timeMS := uint32(since / time.Millisecond)
			root.NextFrame(pixels, timeMS)
			since += p.frameDuration
			cWrite <- pixels
			if t, ok := root.(*Transition); ok {
				if t.OffsetMS+t.TransitionMS < timeMS {
					root = t.After.Pattern
				}
			}

		case <-interrupt.Channel:
			return
		}
	}
}

func (p *Painter) runWrite(cGen chan<- Frame, cWrite <-chan Frame, numLights int) {
	defer func() {
		// Tell runPattern() to quit.
		for loop := true; loop; {
			select {
			case _, loop = <-cWrite:
			default:
				loop = false
			}
		}
		select {
		case cGen <- nil:
		default:
		}
		close(cGen)
		p.wg.Done()
	}()

	tick := time.NewTicker(p.frameDuration)
	defer tick.Stop()
	var err error
	buf := make([]byte, numLights*3)
	for {
		pixels, ok := <-cWrite
		if pixels == nil || !ok {
			return
		}
		if err == nil {
			pixels.ToRGB(buf)
			if _, err = p.d.Write(buf); err != nil {
				log.Printf("Writing failed: %s", err)
			}
		}
		cGen <- pixels

		select {
		case <-tick.C:
		case <-interrupt.Channel:
			return
		}
	}
}
<commit_msg>anim1d: Fix bug in offset handling of Transition.<commit_after>// Copyright 2016 Marc-Antoine Ruel. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

package anim1d

import (
	"encoding/json"
	"log"
	"sync"
	"time"

	"github.com/maruel/dlibox/go/donotuse/devices"
	"github.com/maruel/interrupt"
)

// Pattern is a interface to draw an animated line.
type Pattern interface {
	// NextFrame fills the buffer with the image at this time frame.
	//
	// The image should be derived from timeMS, which is the time since this
	// pattern was started.
	//
	// Calling NextFrame() with a nil pattern is valid. Patterns should be
	// callable without crashing with an object initialized with default values.
	//
	// timeMS will cycle after 49.7 days. The reason it's not using time.Duration
	// is that int64 calculation on ARM is very slow and abysmal on xtensa, which
	// this code is transpiled to.
	NextFrame(pixels Frame, timeMS uint32)
}

// Painter handles the "draw frame, write" loop.
type Painter struct {
	d             devices.Display
	c             chan newPattern
	wg            sync.WaitGroup
	frameDuration time.Duration
}

// SetPattern changes the current pattern to a new one.
//
// The pattern is in JSON encoded format. The function will return an error if
// the encoding is bad. The function is synchronous, it returns only after the
// pattern was effectively set.
func (p *Painter) SetPattern(s string, transition time.Duration) error {
	var pat SPattern
	if err := json.Unmarshal([]byte(s), &pat); err != nil {
		return err
	}
	p.c <- newPattern{pat.Pattern, transition}
	return nil
}

func (p *Painter) Close() error {
	select {
	case p.c <- newPattern{}:
	default:
	}
	close(p.c)
	p.wg.Wait()
	return nil
}

// NewPainter returns a Painter that manages updating the Patterns to the
// strip.
//
// It Assumes the display uses native RGB packed pixels.
func NewPainter(d devices.Display, fps int) *Painter {
	p := &Painter{
		d:             d,
		c:             make(chan newPattern),
		frameDuration: time.Second / time.Duration(fps),
	}
	numLights := d.Bounds().Dx()
	// Tripple buffering.
	cGen := make(chan Frame, 3)
	cWrite := make(chan Frame, cap(cGen))
	for i := 0; i < cap(cGen); i++ {
		cGen <- make(Frame, numLights)
	}
	p.wg.Add(2)
	go p.runPattern(cGen, cWrite)
	go p.runWrite(cGen, cWrite, numLights)
	return p
}

// Private stuff.

var black = &Color{}

type newPattern struct {
	p Pattern
	d time.Duration
}

func (p *Painter) runPattern(cGen <-chan Frame, cWrite chan<- Frame) {
	defer func() {
		// Tell runWrite() to quit.
		for loop := true; loop; {
			select {
			case _, loop = <-cGen:
			default:
				loop = false
			}
		}
		select {
		case cWrite <- nil:
		default:
		}
		close(cWrite)
		p.wg.Done()
	}()

	var root Pattern = black
	var since time.Duration
	for {
		select {
		case newPat, ok := <-p.c:
			if newPat.p == nil || !ok {
				// Request to terminate.
				return
			}

			// New pattern.
			if newPat.d == 0 {
				root = newPat.p
			} else {
				root = &Transition{
					Before:       SPattern{root},
					After:        SPattern{newPat.p},
					OffsetMS:     uint32(since / time.Millisecond),
					TransitionMS: uint32(newPat.d / time.Millisecond),
					Curve:        EaseOut,
				}
			}

		case pixels, ok := <-cGen:
			if !ok {
				return
			}
			for i := range pixels {
				pixels[i] = Color{}
			}
			timeMS := uint32(since / time.Millisecond)
			root.NextFrame(pixels, timeMS)
			since += p.frameDuration
			cWrite <- pixels
			if t, ok := root.(*Transition); ok {
				if t.OffsetMS+t.TransitionMS < timeMS {
					root = t.After.Pattern
					since -= time.Duration(t.OffsetMS) * time.Millisecond
				}
			}

		case <-interrupt.Channel:
			return
		}
	}
}

func (p *Painter) runWrite(cGen chan<- Frame, cWrite <-chan Frame, numLights int) {
	defer func() {
		// Tell runPattern() to quit.
		for loop := true; loop; {
			select {
			case _, loop = <-cWrite:
			default:
				loop = false
			}
		}
		select {
		case cGen <- nil:
		default:
		}
		close(cGen)
		p.wg.Done()
	}()

	tick := time.NewTicker(p.frameDuration)
	defer tick.Stop()
	var err error
	buf := make([]byte, numLights*3)
	for {
		pixels, ok := <-cWrite
		if pixels == nil || !ok {
			return
		}
		if err == nil {
			pixels.ToRGB(buf)
			if _, err = p.d.Write(buf); err != nil {
				log.Printf("Writing failed: %s", err)
			}
		}
		cGen <- pixels

		select {
		case <-tick.C:
		case <-interrupt.Channel:
			return
		}
	}
}
<|endoftext|>


<commit_before>package mips

import (
	cs "github.com/bnagy/gapstone"
	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"

	"../../models"
)

var Arch = &models.Arch{
	Bits:    32,
	Radare:  "mips",
	CS_ARCH: cs.CS_ARCH_MIPS,
	CS_MODE: cs.CS_MODE_MIPS32 + cs.CS_MODE_LITTLE_ENDIAN,
	UC_ARCH: uc.UC_ARCH_MIPS,
	UC_MODE: uc.UC_MODE_MIPS32 + uc.UC_MODE_LITTLE_ENDIAN,
	SP:      uc.UC_MIPS_REG_SP,
}
<commit_msg>add mips registers<commit_after>package mips

import (
	cs "github.com/bnagy/gapstone"
	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"

	"../../models"
)

var Arch = &models.Arch{
	Bits:    32,
	Radare:  "mips",
	CS_ARCH: cs.CS_ARCH_MIPS,
	CS_MODE: cs.CS_MODE_MIPS32 + cs.CS_MODE_LITTLE_ENDIAN,
	UC_ARCH: uc.UC_ARCH_MIPS,
	UC_MODE: uc.UC_MODE_MIPS32 + uc.UC_MODE_LITTLE_ENDIAN,
	SP:      uc.UC_MIPS_REG_SP,
	Regs: map[int]string{
		uc.UC_MIPS_REG_AT: "at",
		uc.UC_MIPS_REG_V0: "v0",
		uc.UC_MIPS_REG_V1: "v1",
		uc.UC_MIPS_REG_A0: "a0",
		uc.UC_MIPS_REG_A1: "a1",
		uc.UC_MIPS_REG_A2: "a2",
		uc.UC_MIPS_REG_A3: "a3",
		uc.UC_MIPS_REG_T0: "t0",
		uc.UC_MIPS_REG_T1: "t1",
		uc.UC_MIPS_REG_T2: "t2",
		uc.UC_MIPS_REG_T3: "t3",
		uc.UC_MIPS_REG_T4: "t4",
		uc.UC_MIPS_REG_T5: "t5",
		uc.UC_MIPS_REG_T6: "t6",
		uc.UC_MIPS_REG_T7: "t7",
		uc.UC_MIPS_REG_T8: "t8",
		uc.UC_MIPS_REG_T9: "t9",
		uc.UC_MIPS_REG_S0: "s0",
		uc.UC_MIPS_REG_S1: "s1",
		uc.UC_MIPS_REG_S2: "s2",
		uc.UC_MIPS_REG_S3: "s3",
		uc.UC_MIPS_REG_S4: "s4",
		uc.UC_MIPS_REG_S5: "s5",
		uc.UC_MIPS_REG_S6: "s6",
		uc.UC_MIPS_REG_S7: "s7",
		uc.UC_MIPS_REG_S8: "s8",
		uc.UC_MIPS_REG_K0: "k0",
		uc.UC_MIPS_REG_K1: "k1",
		uc.UC_MIPS_REG_GP: "gp",
		uc.UC_MIPS_REG_SP: "sp",
		uc.UC_MIPS_REG_RA: "ra",
	},
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	jc "github.com/juju/testing/checkers"
	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/testing/testbase"
)

type ImageMetadataSuite struct {
	testbase.LoggingSuite
	environ []string
	home    *testing.FakeHome
	dir     string
}

var _ = gc.Suite(&ImageMetadataSuite{})

func (s *ImageMetadataSuite) SetUpSuite(c *gc.C) {
	s.LoggingSuite.SetUpSuite(c)
	s.environ = os.Environ()
}

func (s *ImageMetadataSuite) SetUpTest(c *gc.C) {
	s.LoggingSuite.SetUpTest(c)
	os.Clearenv()
	s.dir = c.MkDir()
	// Create a fake certificate so azure test environment can be opened.
	certfile, err := ioutil.TempFile(s.dir, "")
	c.Assert(err, gc.IsNil)
	filename := certfile.Name()
	err = ioutil.WriteFile(filename, []byte("test certificate"), 0644)
	c.Assert(err, gc.IsNil)
	envConfig := strings.Replace(metadataTestEnvConfig, "/home/me/azure.pem", filename, -1)
	s.home = testing.MakeFakeHome(c, envConfig)
	s.PatchEnvironment("AWS_ACCESS_KEY_ID", "access")
	s.PatchEnvironment("AWS_SECRET_ACCESS_KEY", "secret")
}

func (s *ImageMetadataSuite) TearDownTest(c *gc.C) {
	for _, envstring := range s.environ {
		kv := strings.SplitN(envstring, "=", 2)
		os.Setenv(kv[0], kv[1])
	}
	s.home.Restore()
	s.LoggingSuite.TearDownTest(c)
}

var seriesVersions map[string]string = map[string]string{
	"precise": "12.04",
	"raring":  "13.04",
	"trusty":  "14.04",
}

type expectedMetadata struct {
	series   string
	arch     string
	region   string
	endpoint string
}

func (s *ImageMetadataSuite) assertCommandOutput(c *gc.C, expected expectedMetadata, errOut, indexFileName, imageFileName string) {
	if expected.region == "" {
		expected.region = "region"
	}
	if expected.endpoint == "" {
		expected.endpoint = "endpoint"
	}
	strippedOut := strings.Replace(errOut, "\n", "", -1)
	c.Check(strippedOut, gc.Matches, `image metadata files have been written to.*`)
	indexpath := filepath.Join(s.dir, "images", "streams", "v1", indexFileName)
	data, err := ioutil.ReadFile(indexpath)
	c.Assert(err, gc.IsNil)
	content := string(data)
	var indices interface{}
	err = json.Unmarshal(data, &indices)
	c.Assert(err, gc.IsNil)
	c.Assert(indices.(map[string]interface{})["format"], gc.Equals, "index:1.0")
	prodId := fmt.Sprintf("com.ubuntu.cloud:server:%s:%s", seriesVersions[expected.series], expected.arch)
	c.Assert(content, jc.Contains, prodId)
	c.Assert(content, jc.Contains, fmt.Sprintf(`"region": %q`, expected.region))
	c.Assert(content, jc.Contains, fmt.Sprintf(`"endpoint": %q`, expected.endpoint))
	c.Assert(content, jc.Contains, fmt.Sprintf(`"path": "streams/v1/%s"`, imageFileName))

	imagepath := filepath.Join(s.dir, "images", "streams", "v1", imageFileName)
	data, err = ioutil.ReadFile(imagepath)
	c.Assert(err, gc.IsNil)
	content = string(data)
	var images interface{}
	err = json.Unmarshal(data, &images)
	c.Assert(err, gc.IsNil)
	c.Assert(images.(map[string]interface{})["format"], gc.Equals, "products:1.0")
	c.Assert(content, jc.Contains, prodId)
	c.Assert(content, jc.Contains, `"id": "1234"`)
}

const (
	defaultIndexFileName = "index.json"
	defaultImageFileName = "com.ubuntu.cloud:released:imagemetadata.json"
)

func (s *ImageMetadataSuite) TestImageMetadataFilesNoEnv(c *gc.C) {
	ctx := testing.Context(c)
	code := cmd.Main(
		&ImageMetadataCommand{}, ctx, []string{
			"-d", s.dir, "-i", "1234", "-r", "region", "-a", "arch", "-u", "endpoint", "-s", "raring"})
	c.Assert(code, gc.Equals, 0)
	out := testing.Stdout(ctx)
	expected := expectedMetadata{
		series: "raring",
		arch:   "arch",
	}
	s.assertCommandOutput(c, expected, out, defaultIndexFileName, defaultImageFileName)
}

func (s *ImageMetadataSuite) TestImageMetadataFilesDefaultArch(c *gc.C) {
	ctx := testing.Context(c)
	code := cmd.Main(
		&ImageMetadataCommand{}, ctx, []string{
			"-d", s.dir, "-i", "1234", "-r", "region", "-u", "endpoint", "-s", "raring"})
	c.Assert(code, gc.Equals, 0)
	out := testing.Stdout(ctx)
	expected := expectedMetadata{
		series: "raring",
		arch:   "amd64",
	}
	s.assertCommandOutput(c, expected, out, defaultIndexFileName, defaultImageFileName)
}

func (s *ImageMetadataSuite) TestImageMetadataFilesDefaultSeries(c *gc.C) {
	ctx := testing.Context(c)
	code := cmd.Main(
		&ImageMetadataCommand{}, ctx, []string{
			"-d", s.dir, "-i", "1234", "-r", "region", "-a", "arch", "-u", "endpoint"})
	c.Assert(code, gc.Equals, 0)
	out := testing.Stdout(ctx)
	expected := expectedMetadata{
		series: config.LatestLtsSeries(),
		arch:   "arch",
	}
	s.assertCommandOutput(c, expected, out, defaultIndexFileName, defaultImageFileName)
}

func (s *ImageMetadataSuite) TestImageMetadataFilesUsingEnv(c *gc.C) {
	ctx := testing.Context(c)
	code := cmd.Main(
		&ImageMetadataCommand{}, ctx, []string{"-d", s.dir, "-e", "ec2", "-i", "1234"})
	c.Assert(code, gc.Equals, 0)
	out := testing.Stdout(ctx)
	expected := expectedMetadata{
		series:   "precise",
		arch:     "amd64",
		region:   "us-east-1",
		endpoint: "https://ec2.us-east-1.amazonaws.com",
	}
	s.assertCommandOutput(c, expected, out, defaultIndexFileName, defaultImageFileName)
}

func (s *ImageMetadataSuite) TestImageMetadataFilesUsingEnvWithRegionOverride(c *gc.C) {
	ctx := testing.Context(c)
	code := cmd.Main(
		&ImageMetadataCommand{}, ctx, []string{
			"-d", s.dir, "-e", "ec2", "-r", "us-west-1", "-u", "https://ec2.us-west-1.amazonaws.com", "-i", "1234"})
	c.Assert(code, gc.Equals, 0)
	out := testing.Stdout(ctx)
	expected := expectedMetadata{
		series:   "precise",
		arch:     "amd64",
		region:   "us-west-1",
		endpoint: "https://ec2.us-west-1.amazonaws.com",
	}
	s.assertCommandOutput(c, expected, out, defaultIndexFileName, defaultImageFileName)
}

func (s *ImageMetadataSuite) TestImageMetadataFilesUsingEnvWithNoHasRegion(c *gc.C) {
	ctx := testing.Context(c)
	code := cmd.Main(
		&ImageMetadataCommand{}, ctx, []string{
			"-d", s.dir, "-e", "azure", "-r", "region", "-u", "endpoint", "-i", "1234"})
	c.Assert(code, gc.Equals, 0)
	out := testing.Stdout(ctx)
	expected := expectedMetadata{
		series:   "raring",
		arch:     "amd64",
		region:   "region",
		endpoint: "endpoint",
	}
	s.assertCommandOutput(c, expected, out, defaultIndexFileName, defaultImageFileName)
}

type errTestParams struct {
	args []string
}

var errTests = []errTestParams{
	{
		// Missing image id
		args: []string{"-r", "region", "-a", "arch", "-u", "endpoint", "-s", "precise"},
	},
	{
		// Missing region
		args: []string{"-i", "1234", "-a", "arch", "-u", "endpoint", "-s", "precise"},
	},
	{
		// Missing endpoint
		args: []string{"-i", "1234", "-u", "endpoint", "-a", "arch", "-s", "precise"},
	},
	{
		// Missing endpoint/region for environment with no HasRegion interface
		args: []string{"-i", "1234", "-e", "azure"},
	},
}

func (s *ImageMetadataSuite) TestImageMetadataBadArgs(c *gc.C) {
	defer testing.MakeEmptyFakeHome(c).Restore()
	for i, t := range errTests {
		c.Logf("test: %d", i)
		ctx := testing.Context(c)
		code := cmd.Main(&ImageMetadataCommand{}, ctx, t.args)
		c.Check(code, gc.Equals, 1)
	}
}
<commit_msg>Change test case name to TestImageMetadataFilesLatestLts.<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	jc "github.com/juju/testing/checkers"
	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/testing/testbase"
)

type ImageMetadataSuite struct {
	testbase.LoggingSuite
	environ []string
	home    *testing.FakeHome
	dir     string
}

var _ = gc.Suite(&ImageMetadataSuite{})

func (s *ImageMetadataSuite) SetUpSuite(c *gc.C) {
	s.LoggingSuite.SetUpSuite(c)
	s.environ = os.Environ()
}

func (s *ImageMetadataSuite) SetUpTest(c *gc.C) {
	s.LoggingSuite.SetUpTest(c)
	os.Clearenv()
	s.dir = c.MkDir()
	// Create a fake certificate so azure test environment can be opened.
	certfile, err := ioutil.TempFile(s.dir, "")
	c.Assert(err, gc.IsNil)
	filename := certfile.Name()
	err = ioutil.WriteFile(filename, []byte("test certificate"), 0644)
	c.Assert(err, gc.IsNil)
	envConfig := strings.Replace(metadataTestEnvConfig, "/home/me/azure.pem", filename, -1)
	s.home = testing.MakeFakeHome(c, envConfig)
	s.PatchEnvironment("AWS_ACCESS_KEY_ID", "access")
	s.PatchEnvironment("AWS_SECRET_ACCESS_KEY", "secret")
}

func (s *ImageMetadataSuite) TearDownTest(c *gc.C) {
	for _, envstring := range s.environ {
		kv := strings.SplitN(envstring, "=", 2)
		os.Setenv(kv[0], kv[1])
	}
	s.home.Restore()
	s.LoggingSuite.TearDownTest(c)
}

var seriesVersions map[string]string = map[string]string{
	"precise": "12.04",
	"raring":  "13.04",
	"trusty":  "14.04",
}

type expectedMetadata struct {
	series   string
	arch     string
	region   string
	endpoint string
}

func (s *ImageMetadataSuite) assertCommandOutput(c *gc.C, expected expectedMetadata, errOut, indexFileName, imageFileName string) {
	if expected.region == "" {
		expected.region = "region"
	}
	if expected.endpoint == "" {
		expected.endpoint = "endpoint"
	}
	strippedOut := strings.Replace(errOut, "\n", "", -1)
	c.Check(strippedOut, gc.Matches, `image metadata files have been written to.*`)
	indexpath := filepath.Join(s.dir, "images", "streams", "v1", indexFileName)
	data, err := ioutil.ReadFile(indexpath)
	c.Assert(err, gc.IsNil)
	content := string(data)
	var indices interface{}
	err = json.Unmarshal(data, &indices)
	c.Assert(err, gc.IsNil)
	c.Assert(indices.(map[string]interface{})["format"], gc.Equals, "index:1.0")
	prodId := fmt.Sprintf("com.ubuntu.cloud:server:%s:%s", seriesVersions[expected.series], expected.arch)
	c.Assert(content, jc.Contains, prodId)
	c.Assert(content, jc.Contains, fmt.Sprintf(`"region": %q`, expected.region))
	c.Assert(content, jc.Contains, fmt.Sprintf(`"endpoint": %q`, expected.endpoint))
	c.Assert(content, jc.Contains, fmt.Sprintf(`"path": "streams/v1/%s"`, imageFileName))

	imagepath := filepath.Join(s.dir, "images", "streams", "v1", imageFileName)
	data, err = ioutil.ReadFile(imagepath)
	c.Assert(err, gc.IsNil)
	content = string(data)
	var images interface{}
	err = json.Unmarshal(data, &images)
	c.Assert(err, gc.IsNil)
	c.Assert(images.(map[string]interface{})["format"], gc.Equals, "products:1.0")
	c.Assert(content, jc.Contains, prodId)
	c.Assert(content, jc.Contains, `"id": "1234"`)
}

const (
	defaultIndexFileName = "index.json"
	defaultImageFileName = "com.ubuntu.cloud:released:imagemetadata.json"
)

func (s *ImageMetadataSuite) TestImageMetadataFilesNoEnv(c *gc.C) {
	ctx := testing.Context(c)
	code := cmd.Main(
		&ImageMetadataCommand{}, ctx, []string{
			"-d", s.dir, "-i", "1234", "-r", "region", "-a", "arch", "-u", "endpoint", "-s", "raring"})
	c.Assert(code, gc.Equals, 0)
	out := testing.Stdout(ctx)
	expected := expectedMetadata{
		series: "raring",
		arch:   "arch",
	}
	s.assertCommandOutput(c, expected, out, defaultIndexFileName, defaultImageFileName)
}

func (s *ImageMetadataSuite) TestImageMetadataFilesDefaultArch(c *gc.C) {
	ctx := testing.Context(c)
	code := cmd.Main(
		&ImageMetadataCommand{}, ctx, []string{
			"-d", s.dir, "-i", "1234", "-r", "region", "-u", "endpoint", "-s", "raring"})
	c.Assert(code, gc.Equals, 0)
	out := testing.Stdout(ctx)
	expected := expectedMetadata{
		series: "raring",
		arch:   "amd64",
	}
	s.assertCommandOutput(c, expected, out, defaultIndexFileName, defaultImageFileName)
}

func (s *ImageMetadataSuite) TestImageMetadataFilesLatestLts(c *gc.C) {
	ctx := testing.Context(c)
	code := cmd.Main(
		&ImageMetadataCommand{}, ctx, []string{
			"-d", s.dir, "-i", "1234", "-r", "region", "-a", "arch", "-u", "endpoint"})
	c.Assert(code, gc.Equals, 0)
	out := testing.Stdout(ctx)
	expected := expectedMetadata{
		series: config.LatestLtsSeries(),
		arch:   "arch",
	}
	s.assertCommandOutput(c, expected, out, defaultIndexFileName, defaultImageFileName)
}

func (s *ImageMetadataSuite) TestImageMetadataFilesUsingEnv(c *gc.C) {
	ctx := testing.Context(c)
	code := cmd.Main(
		&ImageMetadataCommand{}, ctx, []string{"-d", s.dir, "-e", "ec2", "-i", "1234"})
	c.Assert(code, gc.Equals, 0)
	out := testing.Stdout(ctx)
	expected := expectedMetadata{
		series:   "precise",
		arch:     "amd64",
		region:   "us-east-1",
		endpoint: "https://ec2.us-east-1.amazonaws.com",
	}
	s.assertCommandOutput(c, expected, out, defaultIndexFileName, defaultImageFileName)
}

func (s *ImageMetadataSuite) TestImageMetadataFilesUsingEnvWithRegionOverride(c *gc.C) {
	ctx := testing.Context(c)
	code := cmd.Main(
		&ImageMetadataCommand{}, ctx, []string{
			"-d", s.dir, "-e", "ec2", "-r", "us-west-1", "-u", "https://ec2.us-west-1.amazonaws.com", "-i", "1234"})
	c.Assert(code, gc.Equals, 0)
	out := testing.Stdout(ctx)
	expected := expectedMetadata{
		series:   "precise",
		arch:     "amd64",
		region:   "us-west-1",
		endpoint: "https://ec2.us-west-1.amazonaws.com",
	}
	s.assertCommandOutput(c, expected, out, defaultIndexFileName, defaultImageFileName)
}

func (s *ImageMetadataSuite) TestImageMetadataFilesUsingEnvWithNoHasRegion(c *gc.C) {
	ctx := testing.Context(c)
	code := cmd.Main(
		&ImageMetadataCommand{}, ctx, []string{
			"-d", s.dir, "-e", "azure", "-r", "region", "-u", "endpoint", "-i", "1234"})
	c.Assert(code, gc.Equals, 0)
	out := testing.Stdout(ctx)
	expected := expectedMetadata{
		series:   "raring",
		arch:     "amd64",
		region:   "region",
		endpoint: "endpoint",
	}
	s.assertCommandOutput(c, expected, out, defaultIndexFileName, defaultImageFileName)
}

type errTestParams struct {
	args []string
}

var errTests = []errTestParams{
	{
		// Missing image id
		args: []string{"-r", "region", "-a", "arch", "-u", "endpoint", "-s", "precise"},
	},
	{
		// Missing region
		args: []string{"-i", "1234", "-a", "arch", "-u", "endpoint", "-s", "precise"},
	},
	{
		// Missing endpoint
		args: []string{"-i", "1234", "-u", "endpoint", "-a", "arch", "-s", "precise"},
	},
	{
		// Missing endpoint/region for environment with no HasRegion interface
		args: []string{"-i", "1234", "-e", "azure"},
	},
}

func (s *ImageMetadataSuite) TestImageMetadataBadArgs(c *gc.C) {
	defer testing.MakeEmptyFakeHome(c).Restore()
	for i, t := range errTests {
		c.Logf("test: %d", i)
		ctx := testing.Context(c)
		code := cmd.Main(&ImageMetadataCommand{}, ctx, t.args)
		c.Check(code, gc.Equals, 1)
	}
}
<|endoftext|>


<commit_before>package gubled

import (
	"os"
	"testing"

	"github.com/stretchr/testify/assert"

	"fmt"
	"time"

	"io/ioutil"
	"log"

	"github.com/smancke/guble/client"
	"github.com/smancke/guble/gubled/config"
	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/testutil"
)

type testgroup struct {
	t                   *testing.T
	groupID             int
	addr                string
	done                chan bool
	messagesToSend      int
	consumer, publisher client.Client
	topic               string
}

func newTestgroup(t *testing.T, groupID int, addr string, messagesToSend int) *testgroup {
	return &testgroup{
		t:              t,
		groupID:        groupID,
		addr:           addr,
		done:           make(chan bool),
		messagesToSend: messagesToSend,
	}
}

func TestThroughput(t *testing.T) {
	// defer testutil.EnableDebugForMethod()()
	testutil.SkipIfShort(t)
	defer testutil.ResetDefaultRegistryHealthCheck()

	dir, _ := ioutil.TempDir("", "guble_benchmarking_test")

	*config.HttpListen = "localhost:0"
	*config.KVS = "memory"
	*config.MS = "file"
	*config.StoragePath = dir

	service := StartService()

	testgroupCount := 4
	messagesPerGroup := 100
	log.Printf("init the %v testgroups", testgroupCount)
	testgroups := make([]*testgroup, testgroupCount, testgroupCount)
	for i := range testgroups {
		testgroups[i] = newTestgroup(t, i, service.WebServer().GetAddr(), messagesPerGroup)
	}

	// init test
	log.Print("init the testgroups")
	for i := range testgroups {
		testgroups[i].Init()
	}

	defer func() {
		// cleanup tests
		log.Print("cleanup the testgroups")
		for i := range testgroups {
			testgroups[i].Clean()
		}

		service.Stop()

		os.RemoveAll(dir)
	}()

	// start test
	log.Print("start the testgroups")
	start := time.Now()
	for i := range testgroups {
		go testgroups[i].Start()
	}

	log.Print("wait for finishing")
	for i, test := range testgroups {
		select {
		case successFlag := <-test.done:
			if !successFlag {
				t.Logf("testgroup %v returned with error", i)
				t.FailNow()
				return
			}
		case <-time.After(time.Second * 20):
			t.Log("timeout. testgroups not ready before timeout")
			t.Fail()
			return
		}
	}

	end := time.Now()
	totalMessages := testgroupCount * messagesPerGroup
	throughput := float64(totalMessages) / end.Sub(start).Seconds()
	log.Printf("finished! Throughput: %v/sec (%v message in %v)", int(throughput), totalMessages, end.Sub(start))

	time.Sleep(time.Second * 1)
}

func (tg *testgroup) Init() {
	tg.topic = fmt.Sprintf("/%v-foo", tg.groupID)
	var err error
	location := "ws://" + tg.addr + "/stream/user/xy"
	//location := "ws://gathermon.mancke.net:8080/stream/"
	//location := "ws://127.0.0.1:8080/stream/"
	tg.consumer, err = client.Open(location, "http://localhost/", 10, false)
	if err != nil {
		panic(err)
	}
	tg.publisher, err = client.Open(location, "http://localhost/", 10, false)
	if err != nil {
		panic(err)
	}

	tg.expectStatusMessage(protocol.SUCCESS_CONNECTED, "You are connected to the server.")

	tg.consumer.Subscribe(tg.topic)
	time.Sleep(time.Millisecond * 1)
	//test.expectStatusMessage(protocol.SUCCESS_SUBSCRIBED_TO, test.topic)
}

func (tg *testgroup) expectStatusMessage(name string, arg string) {
	select {
	case notify := <-tg.consumer.StatusMessages():
		assert.Equal(tg.t, name, notify.Name)
		assert.Equal(tg.t, arg, notify.Arg)
	case <-time.After(time.Second * 1):
		tg.t.Logf("[%v] no notification of type %s until timeout", tg.groupID, name)
		tg.done <- false
		tg.t.Fail()
		return
	}
}

func (tg *testgroup) Start() {
	go func() {
		for i := 0; i < tg.messagesToSend; i++ {
			body := fmt.Sprintf("Hallo-%d", i)
			tg.publisher.Send(tg.topic, body, "")
		}
	}()

	for i := 0; i < tg.messagesToSend; i++ {
		body := fmt.Sprintf("Hallo-%d", i)

		select {
		case msg := <-tg.consumer.Messages():
			assert.Equal(tg.t, tg.topic, string(msg.Path))
			if !assert.Equal(tg.t, body, msg.BodyAsString()) {
				tg.t.FailNow()
				tg.done <- false
			}
		case msg := <-tg.consumer.Errors():
			tg.t.Logf("[%v] received error: %v", tg.groupID, msg)
			tg.done <- false
			tg.t.Fail()
			return
		case <-time.After(time.Second * 5):
			tg.t.Logf("[%v] no message received until timeout, expected message %v", tg.groupID, i)
			tg.done <- false
			tg.t.Fail()
			return
		}
	}
	tg.done <- true
}

func (tg *testgroup) Clean() {
	tg.consumer.Close()
	tg.publisher.Close()
}
<commit_msg>Add comment for disabling the TestThroughput<commit_after>package gubled

import (
	"os"
	"testing"

	"github.com/stretchr/testify/assert"

	"fmt"
	"time"

	"io/ioutil"
	"log"

	"github.com/smancke/guble/client"
	"github.com/smancke/guble/gubled/config"
	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/testutil"
)

type testgroup struct {
	t                   *testing.T
	groupID             int
	addr                string
	done                chan bool
	messagesToSend      int
	consumer, publisher client.Client
	topic               string
}

func newTestgroup(t *testing.T, groupID int, addr string, messagesToSend int) *testgroup {
	return &testgroup{
		t:              t,
		groupID:        groupID,
		addr:           addr,
		done:           make(chan bool),
		messagesToSend: messagesToSend,
	}
}

func TestThroughput(t *testing.T) {
	// TODO: We disabled this test because the receiver implementation of fetching messages
	// should be reimplemented according to the new message store
	testutil.SkipIfShort(t)
	defer testutil.ResetDefaultRegistryHealthCheck()

	dir, _ := ioutil.TempDir("", "guble_benchmarking_test")

	*config.HttpListen = "localhost:0"
	*config.KVS = "memory"
	*config.MS = "file"
	*config.StoragePath = dir

	service := StartService()

	testgroupCount := 4
	messagesPerGroup := 100
	log.Printf("init the %v testgroups", testgroupCount)
	testgroups := make([]*testgroup, testgroupCount, testgroupCount)
	for i := range testgroups {
		testgroups[i] = newTestgroup(t, i, service.WebServer().GetAddr(), messagesPerGroup)
	}

	// init test
	log.Print("init the testgroups")
	for i := range testgroups {
		testgroups[i].Init()
	}

	defer func() {
		// cleanup tests
		log.Print("cleanup the testgroups")
		for i := range testgroups {
			testgroups[i].Clean()
		}

		service.Stop()

		os.RemoveAll(dir)
	}()

	// start test
	log.Print("start the testgroups")
	start := time.Now()
	for i := range testgroups {
		go testgroups[i].Start()
	}

	log.Print("wait for finishing")
	for i, test := range testgroups {
		select {
		case successFlag := <-test.done:
			if !successFlag {
				t.Logf("testgroup %v returned with error", i)
				t.FailNow()
				return
			}
		case <-time.After(time.Second * 20):
			t.Log("timeout. testgroups not ready before timeout")
			t.Fail()
			return
		}
	}

	end := time.Now()
	totalMessages := testgroupCount * messagesPerGroup
	throughput := float64(totalMessages) / end.Sub(start).Seconds()
	log.Printf("finished! Throughput: %v/sec (%v message in %v)", int(throughput), totalMessages, end.Sub(start))

	time.Sleep(time.Second * 1)
}

func (tg *testgroup) Init() {
	tg.topic = fmt.Sprintf("/%v-foo", tg.groupID)
	var err error
	location := "ws://" + tg.addr + "/stream/user/xy"
	//location := "ws://gathermon.mancke.net:8080/stream/"
	//location := "ws://127.0.0.1:8080/stream/"
	tg.consumer, err = client.Open(location, "http://localhost/", 10, false)
	if err != nil {
		panic(err)
	}
	tg.publisher, err = client.Open(location, "http://localhost/", 10, false)
	if err != nil {
		panic(err)
	}

	tg.expectStatusMessage(protocol.SUCCESS_CONNECTED, "You are connected to the server.")

	tg.consumer.Subscribe(tg.topic)
	time.Sleep(time.Millisecond * 1)
	//test.expectStatusMessage(protocol.SUCCESS_SUBSCRIBED_TO, test.topic)
}

func (tg *testgroup) expectStatusMessage(name string, arg string) {
	select {
	case notify := <-tg.consumer.StatusMessages():
		assert.Equal(tg.t, name, notify.Name)
		assert.Equal(tg.t, arg, notify.Arg)
	case <-time.After(time.Second * 1):
		tg.t.Logf("[%v] no notification of type %s until timeout", tg.groupID, name)
		tg.done <- false
		tg.t.Fail()
		return
	}
}

func (tg *testgroup) Start() {
	go func() {
		for i := 0; i < tg.messagesToSend; i++ {
			body := fmt.Sprintf("Hallo-%d", i)
			tg.publisher.Send(tg.topic, body, "")
		}
	}()

	for i := 0; i < tg.messagesToSend; i++ {
		body := fmt.Sprintf("Hallo-%d", i)

		select {
		case msg := <-tg.consumer.Messages():
			assert.Equal(tg.t, tg.topic, string(msg.Path))
			if !assert.Equal(tg.t, body, msg.BodyAsString()) {
				tg.t.FailNow()
				tg.done <- false
			}
		case msg := <-tg.consumer.Errors():
			tg.t.Logf("[%v] received error: %v", tg.groupID, msg)
			tg.done <- false
			tg.t.Fail()
			return
		case <-time.After(time.Second * 5):
			tg.t.Logf("[%v] no message received until timeout, expected message %v", tg.groupID, i)
			tg.done <- false
			tg.t.Fail()
			return
		}
	}
	tg.done <- true
}

func (tg *testgroup) Clean() {
	tg.consumer.Close()
	tg.publisher.Close()
}
<|endoftext|>


<commit_before>package rrs

import (
	"fmt"
	"runtime"
	"sync"
	"testing"
	"time"

	"github.com/facebookgo/ensure"
	"github.com/teh-cmc/seq"
)

// -----------------------------------------------------------------------------

// NOTE: run these tests with `go test -race -cpu 1,4,8`

func TestRRSeq_New_BufSize(t *testing.T) {
	var s *RRSeq
	var err error

	name := fmt.Sprintf("TestRRSeq_New_BufSize(gomaxprocs:%d)", runtime.GOMAXPROCS(0))

	s, err = NewRRSeq(name, -42, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}
	ensure.DeepEqual(t, cap(s.ids), 0)
	_ = s.Close()

	s, err = NewRRSeq(name, 0, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}
	ensure.DeepEqual(t, cap(s.ids), 0)
	_ = s.Close()

	s, err = NewRRSeq(name, 1, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}
	ensure.DeepEqual(t, cap(s.ids), 1)
	_ = s.Close()

	s, err = NewRRSeq(name, 1e6, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}
	ensure.DeepEqual(t, cap(s.ids), int(1e6))
	_ = s.Close()
}

func TestRRSeq_FirstID(t *testing.T) {
	name := fmt.Sprintf("TestRRSeq_FirstID(gomaxprocs:%d)", runtime.GOMAXPROCS(0))
	s, err := NewRRSeq(name, 1e2, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}
	ensure.DeepEqual(t, <-s.GetStream(), seq.ID(1))
	_ = s.Close()
}

// -----------------------------------------------------------------------------

func testRRSeq_SingleClient(bufSize int, t *testing.T) {
	name := fmt.Sprintf(
		"testRRSeq_SingleClient(bufsz:%d)(gomaxprocs:%d)", bufSize, runtime.GOMAXPROCS(0),
	)
	s, err := NewRRSeq(name, bufSize, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}
	lastID := seq.ID(0)

	go func() {
		<-time.After(time.Millisecond * 500)
		_ = s.Close()
	}()

	for id := range s.GetStream() {
		ensure.DeepEqual(t, id, lastID+1)
		lastID = id
	}
}

func TestRRSeq_BufSize0_SingleClient(t *testing.T) {
	testRRSeq_SingleClient(0, t)
}

func TestRRSeq_BufSize1_SingleClient(t *testing.T) {
	testRRSeq_SingleClient(1, t)
}

func TestRRSeq_BufSize2_SingleClient(t *testing.T) {
	testRRSeq_SingleClient(2, t)
}

func TestRRSeq_BufSize1024_SingleClient(t *testing.T) {
	testRRSeq_SingleClient(1024, t)
}

// -----------------------------------------------------------------------------

func testRRSeq_MultiClient_Local(bufSize int, t *testing.T) {
	name := fmt.Sprintf(
		"testRRSeq_MultiClient_Local(bufsz:%d)(gomaxprocs:%d)", bufSize, runtime.GOMAXPROCS(0),
	)
	s, err := NewRRSeq(name, bufSize, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}
	lastID := seq.ID(0)

	go func() {
		<-time.After(time.Millisecond * 500)
		_ = s.Close()
	}()

	s1, s2, s3 := s.GetStream(), s.GetStream(), s.GetStream()
	for {
		id1 := s1.Next()
		if id1 == 0 {
			break
		}
		ensure.DeepEqual(t, id1, lastID+1)
		lastID++
		id2 := s2.Next()
		if id2 == 0 {
			break
		}
		ensure.DeepEqual(t, id2, id1+1)
		lastID++
		id3 := s3.Next()
		if id3 == 0 {
			break
		}
		ensure.DeepEqual(t, id3, id2+1)
		lastID++
	}
}

func TestRRSeq_BufSize0_MultiClient_Local(t *testing.T) {
	testRRSeq_MultiClient_Local(0, t)
}

func TestRRSeq_BufSize1_MultiClient_Local(t *testing.T) {
	testRRSeq_MultiClient_Local(1, t)
}

func TestRRSeq_BufSize2_MultiClient_Local(t *testing.T) {
	testRRSeq_MultiClient_Local(2, t)
}

func TestRRSeq_BufSize1024_MultiClient_Local(t *testing.T) {
	testRRSeq_MultiClient_Local(1024, t)
}

// -----------------------------------------------------------------------------

func testRRSeq_ConcurrentClients256_Local(bufSize int, t *testing.T) {
	name := fmt.Sprintf(
		"testRRSeq_ConcurrentClients256_Local(bufsz:%d)(gomaxprocs:%d)", bufSize, runtime.GOMAXPROCS(0),
	)
	s, err := NewRRSeq(name, bufSize, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}

	go func() {
		<-time.After(time.Millisecond * 500)
		_ = s.Close()
	}()

	wg := &sync.WaitGroup{}
	for i := 0; i < 256; i++ {
		wg.Add(1)
		go func() {
			for id := range s.GetStream() {
				_ = id
			}
			wg.Done()
		}()
	}
	wg.Wait()
}

func TestRRSeq_BufSize0_ConcurrentClients256_Local(t *testing.T) {
	testRRSeq_ConcurrentClients256_Local(0, t)
}

func TestRRSeq_BufSize1_ConcurrentClients256_Local(t *testing.T) {
	testRRSeq_ConcurrentClients256_Local(1, t)
}

func TestRRSeq_BufSize2_ConcurrentClients256_Local(t *testing.T) {
	testRRSeq_ConcurrentClients256_Local(2, t)
}

func TestRRSeq_BufSize1024_ConcurrentClients256_Local(t *testing.T) {
	testRRSeq_ConcurrentClients256_Local(1024, t)
}

// -----------------------------------------------------------------------------

// This is certainly the most important of the standard (i.e. non-mayhem) tests,
// as it checks that a cluster of `RRServer`s, being bombarded of NextID queries
// on its every nodes, still consistently deliver coherent, sequential `ID`s.
func testRRSeq_ConcurrentClients32_Distributed(bufSize int, t *testing.T) {
	name := fmt.Sprintf(
		"testRRSeq_ConcurrentClients32_Distributed(bufsz:%d)(gomaxprocs:%d)",
		bufSize, runtime.GOMAXPROCS(0),
	)

	ids := make(seq.IDSlice, 0, 32*bufSize*3)
	idsLock := &sync.Mutex{}

	wg := &sync.WaitGroup{}
	for i := 0; i < 32; i++ {
		wg.Add(1)
		go func() {
			s, err := NewRRSeq(name, bufSize, testingRRServerAddrs...)
			if err != nil {
				t.Fatal(err)
			}
			allIDs := make(seq.IDSlice, bufSize*3)
			lastID := seq.ID(0)
			j := 0
			for id := range s.GetStream() {
				if j >= len(allIDs) {
					break
				}
				ensure.True(t, id > lastID)
				lastID = id
				allIDs[j] = id
				j++
			}
			_ = s.Close()

			idsLock.Lock()
			ids = append(ids, allIDs...)
			idsLock.Unlock()

			wg.Done()
		}()
	}
	wg.Wait()

	// this checks that, within a healthy cluster, the complete set of `ID`s
	// returned is monotonically increasing in its entirety
	ids = ids.Sort()
	for i := 0; i < len(ids)-1; i++ {
		ensure.True(t, ids[i]+1 == ids[i+1])
	}
}

func TestRRSeq_BufSize0_ConcurrentClients32_Distributed(t *testing.T) {
	testRRSeq_ConcurrentClients32_Distributed(0, t)
}

func TestRRSeq_BufSize1_ConcurrentClients32_Distributed(t *testing.T) {
	testRRSeq_ConcurrentClients32_Distributed(1, t)
}

func TestRRSeq_BufSize2_ConcurrentClients32_Distributed(t *testing.T) {
	testRRSeq_ConcurrentClients32_Distributed(2, t)
}

func TestRRSeq_BufSize1024_ConcurrentClients32_Distributed(t *testing.T) {
	testRRSeq_ConcurrentClients32_Distributed(1024, t)
}

// -----------------------------------------------------------------------------

// NOTE: run these benchmarks with `go test -run=none -bench=. -cpu 1,8,32`

func benchmarkRRSeq_SingleClient(bufSize int, b *testing.B) {
	name := fmt.Sprintf(
		"benchmarkRRSeq_SingleClient(bufsz:%d)(gomaxprocs:%d)", bufSize, runtime.GOMAXPROCS(0),
	)
	s, err := NewRRSeq(name, bufSize, testingRRServerAddrs...)
	if err != nil {
		b.Fatal(err)
	}
	ids := s.GetStream()

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = ids.Next()
	}
	_ = s.Close()
}

func BenchmarkRRSeq_BufSize0_SingleClient(b *testing.B) {
	benchmarkRRSeq_SingleClient(0, b)
}

func BenchmarkRRSeq_BufSize1_SingleClient(b *testing.B) {
	benchmarkRRSeq_SingleClient(1, b)
}

func BenchmarkRRSeq_BufSize2_SingleClient(b *testing.B) {
	benchmarkRRSeq_SingleClient(2, b)
}

func BenchmarkRRSeq_BufSize1024_SingleClient(b *testing.B) {
	benchmarkRRSeq_SingleClient(1024, b)
}

// -----------------------------------------------------------------------------

func benchmarkRRSeq_MultiClient_Local(bufSize int, b *testing.B) {
	name := fmt.Sprintf(
		"benchmarkRRSeq_MultiClient_Local(bufsz:%d)(gomaxprocs:%d)", bufSize, runtime.GOMAXPROCS(0),
	)
	s, err := NewRRSeq(name, bufSize, testingRRServerAddrs...)
	if err != nil {
		b.Fatal(err)
	}
	b.RunParallel(func(pb *testing.PB) {
		ids := s.GetStream()
		for pb.Next() {
			_ = ids.Next()
		}
	})
	_ = s.Close()
}

func BenchmarkRRSeq_BufSize0_MultiClient_Local(b *testing.B) {
	benchmarkRRSeq_MultiClient_Local(0, b)
}

func BenchmarkRRSeq_BufSize1_MultiClient_Local(b *testing.B) {
	benchmarkRRSeq_MultiClient_Local(1, b)
}

func BenchmarkRRSeq_BufSize2_MultiClient_Local(b *testing.B) {
	benchmarkRRSeq_MultiClient_Local(2, b)
}

func BenchmarkRRSeq_BufSize1024_MultiClient_Local(b *testing.B) {
	benchmarkRRSeq_MultiClient_Local(1024, b)
}

// -----------------------------------------------------------------------------

func ExampleRRSeq() {
	s, err := NewRRSeq("myseq", 2, testingRRServerAddrs...)
	if err != nil {
		panic(err)
	}

	ids := make([]seq.ID, 0)
	for id := range s.GetStream() {
		ids = append(ids, id)
		if id == 10 { // won't stop until 11: 11 is already buffered
			_ = s.Close()
		}
	}
	fmt.Println(ids)

	// Output: [1 2 3 4 5 6 7 8 9 10 11]
}
<commit_msg>RRSeq tests: improved testRRSeq_ConcurrentClients32_Distributed implementation<commit_after>package rrs

import (
	"fmt"
	"runtime"
	"sync"
	"testing"
	"time"

	"github.com/facebookgo/ensure"
	"github.com/teh-cmc/seq"
)

// -----------------------------------------------------------------------------

// NOTE: run these tests with `go test -race -cpu 1,4,8`

func TestRRSeq_New_BufSize(t *testing.T) {
	var s *RRSeq
	var err error

	name := fmt.Sprintf("TestRRSeq_New_BufSize(gomaxprocs:%d)", runtime.GOMAXPROCS(0))

	s, err = NewRRSeq(name, -42, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}
	ensure.DeepEqual(t, cap(s.ids), 0)
	_ = s.Close()

	s, err = NewRRSeq(name, 0, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}
	ensure.DeepEqual(t, cap(s.ids), 0)
	_ = s.Close()

	s, err = NewRRSeq(name, 1, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}
	ensure.DeepEqual(t, cap(s.ids), 1)
	_ = s.Close()

	s, err = NewRRSeq(name, 1e6, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}
	ensure.DeepEqual(t, cap(s.ids), int(1e6))
	_ = s.Close()
}

func TestRRSeq_FirstID(t *testing.T) {
	name := fmt.Sprintf("TestRRSeq_FirstID(gomaxprocs:%d)", runtime.GOMAXPROCS(0))
	s, err := NewRRSeq(name, 1e2, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}
	ensure.DeepEqual(t, <-s.GetStream(), seq.ID(1))
	_ = s.Close()
}

// -----------------------------------------------------------------------------

func testRRSeq_SingleClient(bufSize int, t *testing.T) {
	name := fmt.Sprintf(
		"testRRSeq_SingleClient(bufsz:%d)(gomaxprocs:%d)", bufSize, runtime.GOMAXPROCS(0),
	)
	s, err := NewRRSeq(name, bufSize, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}
	lastID := seq.ID(0)

	go func() {
		<-time.After(time.Millisecond * 500)
		_ = s.Close()
	}()

	for id := range s.GetStream() {
		ensure.DeepEqual(t, id, lastID+1)
		lastID = id
	}
}

func TestRRSeq_BufSize0_SingleClient(t *testing.T) {
	testRRSeq_SingleClient(0, t)
}

func TestRRSeq_BufSize1_SingleClient(t *testing.T) {
	testRRSeq_SingleClient(1, t)
}

func TestRRSeq_BufSize2_SingleClient(t *testing.T) {
	testRRSeq_SingleClient(2, t)
}

func TestRRSeq_BufSize1024_SingleClient(t *testing.T) {
	testRRSeq_SingleClient(1024, t)
}

// -----------------------------------------------------------------------------

func testRRSeq_MultiClient_Local(bufSize int, t *testing.T) {
	name := fmt.Sprintf(
		"testRRSeq_MultiClient_Local(bufsz:%d)(gomaxprocs:%d)", bufSize, runtime.GOMAXPROCS(0),
	)
	s, err := NewRRSeq(name, bufSize, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}
	lastID := seq.ID(0)

	go func() {
		<-time.After(time.Millisecond * 500)
		_ = s.Close()
	}()

	s1, s2, s3 := s.GetStream(), s.GetStream(), s.GetStream()
	for {
		id1 := s1.Next()
		if id1 == 0 {
			break
		}
		ensure.DeepEqual(t, id1, lastID+1)
		lastID++
		id2 := s2.Next()
		if id2 == 0 {
			break
		}
		ensure.DeepEqual(t, id2, id1+1)
		lastID++
		id3 := s3.Next()
		if id3 == 0 {
			break
		}
		ensure.DeepEqual(t, id3, id2+1)
		lastID++
	}
}

func TestRRSeq_BufSize0_MultiClient_Local(t *testing.T) {
	testRRSeq_MultiClient_Local(0, t)
}

func TestRRSeq_BufSize1_MultiClient_Local(t *testing.T) {
	testRRSeq_MultiClient_Local(1, t)
}

func TestRRSeq_BufSize2_MultiClient_Local(t *testing.T) {
	testRRSeq_MultiClient_Local(2, t)
}

func TestRRSeq_BufSize1024_MultiClient_Local(t *testing.T) {
	testRRSeq_MultiClient_Local(1024, t)
}

// -----------------------------------------------------------------------------

func testRRSeq_ConcurrentClients256_Local(bufSize int, t *testing.T) {
	name := fmt.Sprintf(
		"testRRSeq_ConcurrentClients256_Local(bufsz:%d)(gomaxprocs:%d)", bufSize, runtime.GOMAXPROCS(0),
	)
	s, err := NewRRSeq(name, bufSize, testingRRServerAddrs...)
	if err != nil {
		t.Fatal(err)
	}

	go func() {
		<-time.After(time.Millisecond * 500)
		_ = s.Close()
	}()

	wg := &sync.WaitGroup{}
	for i := 0; i < 256; i++ {
		wg.Add(1)
		go func() {
			for id := range s.GetStream() {
				_ = id
			}
			wg.Done()
		}()
	}
	wg.Wait()
}

func TestRRSeq_BufSize0_ConcurrentClients256_Local(t *testing.T) {
	testRRSeq_ConcurrentClients256_Local(0, t)
}

func TestRRSeq_BufSize1_ConcurrentClients256_Local(t *testing.T) {
	testRRSeq_ConcurrentClients256_Local(1, t)
}

func TestRRSeq_BufSize2_ConcurrentClients256_Local(t *testing.T) {
	testRRSeq_ConcurrentClients256_Local(2, t)
}

func TestRRSeq_BufSize1024_ConcurrentClients256_Local(t *testing.T) {
	testRRSeq_ConcurrentClients256_Local(1024, t)
}

// -----------------------------------------------------------------------------

// This is certainly the most important of the standard (i.e. non-mayhem) tests,
// as it checks that a cluster of `RRServer`s, being bombarded of NextID queries
// on its every nodes, still consistently deliver coherent, sequential `ID`s.
func testRRSeq_ConcurrentClients32_Distributed(bufSize int, t *testing.T) {
	name := fmt.Sprintf(
		"testRRSeq_ConcurrentClients32_Distributed(bufsz:%d)(gomaxprocs:%d)",
		bufSize, runtime.GOMAXPROCS(0),
	)

	ids := make(seq.IDSlice, 0, 32*(bufSize+1)*10)
	idsLock := &sync.Mutex{}

	wg := &sync.WaitGroup{}
	for i := 0; i < 32; i++ {
		wg.Add(1)
		go func() {
			s, err := NewRRSeq(name, bufSize, testingRRServerAddrs...)
			if err != nil {
				t.Fatal(err)
			}
			allIDs := make(seq.IDSlice, 0, (bufSize+1)*10)
			lastID := seq.ID(0)
			j := 0
			for id := range s.GetStream() {
				if j >= len(allIDs) {
					break
				}
				ensure.True(t, id > lastID)
				lastID = id
				allIDs = append(allIDs, id)
				j++
			}
			_ = s.Close()

			idsLock.Lock()
			ids = append(ids, allIDs...)
			idsLock.Unlock()

			wg.Done()
		}()
	}
	wg.Wait()

	// this checks that, within a healthy cluster, the complete set of `ID`s
	// returned is monotonically increasing in its entirety
	ids = ids.Sort()
	for i := 0; i < len(ids)-1; i++ {
		ensure.True(t, ids[i]+1 == ids[i+1])
	}
}

func TestRRSeq_BufSize0_ConcurrentClients32_Distributed(t *testing.T) {
	testRRSeq_ConcurrentClients32_Distributed(0, t)
}

func TestRRSeq_BufSize1_ConcurrentClients32_Distributed(t *testing.T) {
	testRRSeq_ConcurrentClients32_Distributed(1, t)
}

func TestRRSeq_BufSize2_ConcurrentClients32_Distributed(t *testing.T) {
	testRRSeq_ConcurrentClients32_Distributed(2, t)
}

func TestRRSeq_BufSize1024_ConcurrentClients32_Distributed(t *testing.T) {
	testRRSeq_ConcurrentClients32_Distributed(1024, t)
}

// -----------------------------------------------------------------------------

// NOTE: run these benchmarks with `go test -run=none -bench=. -cpu 1,8,32`

func benchmarkRRSeq_SingleClient(bufSize int, b *testing.B) {
	name := fmt.Sprintf(
		"benchmarkRRSeq_SingleClient(bufsz:%d)(gomaxprocs:%d)", bufSize, runtime.GOMAXPROCS(0),
	)
	s, err := NewRRSeq(name, bufSize, testingRRServerAddrs...)
	if err != nil {
		b.Fatal(err)
	}
	ids := s.GetStream()

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = ids.Next()
	}
	_ = s.Close()
}

func BenchmarkRRSeq_BufSize0_SingleClient(b *testing.B) {
	benchmarkRRSeq_SingleClient(0, b)
}

func BenchmarkRRSeq_BufSize1_SingleClient(b *testing.B) {
	benchmarkRRSeq_SingleClient(1, b)
}

func BenchmarkRRSeq_BufSize2_SingleClient(b *testing.B) {
	benchmarkRRSeq_SingleClient(2, b)
}

func BenchmarkRRSeq_BufSize1024_SingleClient(b *testing.B) {
	benchmarkRRSeq_SingleClient(1024, b)
}

// -----------------------------------------------------------------------------

func benchmarkRRSeq_MultiClient_Local(bufSize int, b *testing.B) {
	name := fmt.Sprintf(
		"benchmarkRRSeq_MultiClient_Local(bufsz:%d)(gomaxprocs:%d)", bufSize, runtime.GOMAXPROCS(0),
	)
	s, err := NewRRSeq(name, bufSize, testingRRServerAddrs...)
	if err != nil {
		b.Fatal(err)
	}
	b.RunParallel(func(pb *testing.PB) {
		ids := s.GetStream()
		for pb.Next() {
			_ = ids.Next()
		}
	})
	_ = s.Close()
}

func BenchmarkRRSeq_BufSize0_MultiClient_Local(b *testing.B) {
	benchmarkRRSeq_MultiClient_Local(0, b)
}

func BenchmarkRRSeq_BufSize1_MultiClient_Local(b *testing.B) {
	benchmarkRRSeq_MultiClient_Local(1, b)
}

func BenchmarkRRSeq_BufSize2_MultiClient_Local(b *testing.B) {
	benchmarkRRSeq_MultiClient_Local(2, b)
}

func BenchmarkRRSeq_BufSize1024_MultiClient_Local(b *testing.B) {
	benchmarkRRSeq_MultiClient_Local(1024, b)
}

// -----------------------------------------------------------------------------

func ExampleRRSeq() {
	s, err := NewRRSeq("myseq", 2, testingRRServerAddrs...)
	if err != nil {
		panic(err)
	}

	ids := make([]seq.ID, 0)
	for id := range s.GetStream() {
		ids = append(ids, id)
		if id == 10 { // won't stop until 11: 11 is already buffered
			_ = s.Close()
		}
	}
	fmt.Println(ids)

	// Output: [1 2 3 4 5 6 7 8 9 10 11]
}
<|endoftext|>


<commit_before>package torrent

import (
	"context"
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/anacrolix/log"
	"github.com/anacrolix/missinggo/v2"
)

// Accesses Torrent data via a Client. Reads block until the data is available. Seeks and readahead
// also drive Client behaviour.
type Reader interface {
	io.ReadSeekCloser
	missinggo.ReadContexter
	// Configure the number of bytes ahead of a read that should also be prioritized in preparation
	// for further reads. Overridden by non-nil readahead func, see SetReadaheadFunc.
	SetReadahead(int64)
	// If non-nil, the provided function is called when the implementation needs to know the
	// readahead for the current reader. Calls occur during Reads and Seeks, and while the Client is
	// locked.
	SetReadaheadFunc(ReadaheadFunc)
	// Don't wait for pieces to complete and be verified. Read calls return as soon as they can when
	// the underlying chunks become available.
	SetResponsive()
}

// Piece range by piece index, [begin, end).
type pieceRange struct {
	begin, end pieceIndex
}

type ReadaheadContext struct {
	ContiguousReadStartPos int64
	CurrentPos             int64
}

// Returns the desired readahead for a Reader.
type ReadaheadFunc func(ReadaheadContext) int64

type reader struct {
	t *Torrent
	// Adjust the read/seek window to handle Readers locked to File extents and the like.
	offset, length int64

	// Function to dynamically calculate readahead. If nil, readahead is static.
	readaheadFunc ReadaheadFunc

	// Required when modifying pos and readahead.
	mu sync.Locker

	readahead, pos int64
	// Position that reads have continued contiguously from.
	contiguousReadStartPos int64
	// The cached piece range this reader wants downloaded. The zero value corresponds to nothing.
	// We cache this so that changes can be detected, and bubbled up to the Torrent only as
	// required.
	pieces pieceRange

	// Reads have been initiated since the last seek. This is used to prevent readaheads occurring
	// after a seek or with a new reader at the starting position.
	reading    bool
	responsive bool
}

var _ io.ReadSeekCloser = (*reader)(nil)

func (r *reader) SetResponsive() {
	r.responsive = true
	r.t.cl.event.Broadcast()
}

// Disable responsive mode. TODO: Remove?
func (r *reader) SetNonResponsive() {
	r.responsive = false
	r.t.cl.event.Broadcast()
}

func (r *reader) SetReadahead(readahead int64) {
	r.mu.Lock()
	r.readahead = readahead
	r.readaheadFunc = nil
	r.posChanged()
	r.mu.Unlock()
}

func (r *reader) SetReadaheadFunc(f ReadaheadFunc) {
	r.mu.Lock()
	r.readaheadFunc = f
	r.posChanged()
	r.mu.Unlock()
}

// How many bytes are available to read. Max is the most we could require.
func (r *reader) available(off, max int64) (ret int64) {
	off += r.offset
	for max > 0 {
		req, ok := r.t.offsetRequest(off)
		if !ok {
			break
		}
		if !r.responsive && !r.t.pieceComplete(pieceIndex(req.Index)) {
			break
		}
		if !r.t.haveChunk(req) {
			break
		}
		len1 := int64(req.Length) - (off - r.t.requestOffset(req))
		max -= len1
		ret += len1
		off += len1
	}
	// Ensure that ret hasn't exceeded our original max.
	if max < 0 {
		ret += max
	}
	return
}

// Calculates the pieces this reader wants downloaded, ignoring the cached value at r.pieces.
func (r *reader) piecesUncached() (ret pieceRange) {
	ra := r.readahead
	if r.readaheadFunc != nil {
		ra = r.readaheadFunc(ReadaheadContext{
			ContiguousReadStartPos: r.contiguousReadStartPos,
			CurrentPos:             r.pos,
		})
	}
	if ra < 1 {
		// Needs to be at least 1, because [x, x) means we don't want
		// anything.
		ra = 1
	}
	if !r.reading {
		ra = 0
	}
	if ra > r.length-r.pos {
		ra = r.length - r.pos
	}
	ret.begin, ret.end = r.t.byteRegionPieces(r.torrentOffset(r.pos), ra)
	return
}

func (r *reader) Read(b []byte) (n int, err error) {
	return r.ReadContext(context.Background(), b)
}

func (r *reader) ReadContext(ctx context.Context, b []byte) (n int, err error) {
	if len(b) > 0 {
		r.reading = true
		// TODO: Rework reader piece priorities so we don't have to push updates in to the Client
		// and take the lock here.
		r.mu.Lock()
		r.posChanged()
		r.mu.Unlock()
	}
	n, err = r.readOnceAt(ctx, b, r.pos)
	if n == 0 {
		if err == nil && len(b) > 0 {
			panic("expected error")
		} else {
			return
		}
	}

	r.mu.Lock()
	r.pos += int64(n)
	r.posChanged()
	r.mu.Unlock()
	if r.pos >= r.length {
		err = io.EOF
	} else if err == io.EOF {
		err = io.ErrUnexpectedEOF
	}
	return
}

var closedChan = make(chan struct{})

func init() {
	close(closedChan)
}

// Wait until some data should be available to read. Tickles the client if it isn't. Returns how
// much should be readable without blocking.
func (r *reader) waitAvailable(ctx context.Context, pos, wanted int64, wait bool) (avail int64, err error) {
	t := r.t
	for {
		r.t.cl.rLock()
		avail = r.available(pos, wanted)
		readerCond := t.piece(int((r.offset + pos) / t.info.PieceLength)).readerCond.Signaled()
		r.t.cl.rUnlock()
		if avail != 0 {
			return
		}
		var dontWait <-chan struct{}
		if !wait || wanted == 0 {
			dontWait = closedChan
		}
		select {
		case <-r.t.closed.Done():
			err = errors.New("torrent closed")
			return
		case <-ctx.Done():
			err = ctx.Err()
			return
		case <-r.t.dataDownloadDisallowed.On():
			err = errors.New("torrent data downloading disabled")
		case <-r.t.networkingEnabled.Off():
			err = errors.New("torrent networking disabled")
			return
		case <-dontWait:
			return
		case <-readerCond:
		}
	}
}

// Adds the reader's torrent offset to the reader object offset (for example the reader might be
// constrainted to a particular file within the torrent).
func (r *reader) torrentOffset(readerPos int64) int64 {
	return r.offset + readerPos
}

// Performs at most one successful read to torrent storage.
func (r *reader) readOnceAt(ctx context.Context, b []byte, pos int64) (n int, err error) {
	if pos >= r.length {
		err = io.EOF
		return
	}
	for {
		var avail int64
		avail, err = r.waitAvailable(ctx, pos, int64(len(b)), n == 0)
		if avail == 0 {
			return
		}
		firstPieceIndex := pieceIndex(r.torrentOffset(pos) / r.t.info.PieceLength)
		firstPieceOffset := r.torrentOffset(pos) % r.t.info.PieceLength
		b1 := missinggo.LimitLen(b, avail)
		n, err = r.t.readAt(b1, r.torrentOffset(pos))
		if n != 0 {
			err = nil
			return
		}
		r.t.cl.lock()
		// TODO: Just reset pieces in the readahead window. This might help
		// prevent thrashing with small caches and file and piece priorities.
		r.log(log.Fstr("error reading torrent %s piece %d offset %d, %d bytes: %v",
			r.t.infoHash.HexString(), firstPieceIndex, firstPieceOffset, len(b1), err))
		if !r.t.updatePieceCompletion(firstPieceIndex) {
			r.log(log.Fstr("piece %d completion unchanged", firstPieceIndex))
		}
		// Update the rest of the piece completions in the readahead window, without alerting to
		// changes (since only the first piece, the one above, could have generated the read error
		// we're currently handling).
		if r.pieces.begin != firstPieceIndex {
			panic(fmt.Sprint(r.pieces.begin, firstPieceIndex))
		}
		for index := r.pieces.begin + 1; index < r.pieces.end; index++ {
			r.t.updatePieceCompletion(index)
		}
		r.t.cl.unlock()
	}
}

// Hodor
func (r *reader) Close() error {
	r.t.cl.lock()
	r.t.deleteReader(r)
	r.t.cl.unlock()
	return nil
}

func (r *reader) posChanged() {
	to := r.piecesUncached()
	from := r.pieces
	if to == from {
		return
	}
	r.pieces = to
	// log.Printf("reader pos changed %v->%v", from, to)
	r.t.readerPosChanged(from, to)
}

func (r *reader) Seek(off int64, whence int) (newPos int64, err error) {
	switch whence {
	case io.SeekStart:
		newPos = off
		r.mu.Lock()
	case io.SeekCurrent:
		r.mu.Lock()
		newPos = r.pos + off
	case io.SeekEnd:
		newPos = r.length + off
		r.mu.Lock()
	default:
		return 0, errors.New("bad whence")
	}
	if newPos != r.pos {
		r.reading = false
		r.pos = newPos
		r.contiguousReadStartPos = newPos
		r.posChanged()
	}
	r.mu.Unlock()
	return
}

func (r *reader) log(m log.Msg) {
	r.t.logger.LogLevel(log.Debug, m.Skip(1))
}

// Implementation inspired by https://news.ycombinator.com/item?id=27019613.
func defaultReadaheadFunc(r ReadaheadContext) int64 {
	return r.CurrentPos - r.ContiguousReadStartPos
}
<commit_msg>Ensure unlock occurs on panic in reader<commit_after>package torrent

import (
	"context"
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/anacrolix/log"
	"github.com/anacrolix/missinggo/v2"
)

// Accesses Torrent data via a Client. Reads block until the data is available. Seeks and readahead
// also drive Client behaviour.
type Reader interface {
	io.ReadSeekCloser
	missinggo.ReadContexter
	// Configure the number of bytes ahead of a read that should also be prioritized in preparation
	// for further reads. Overridden by non-nil readahead func, see SetReadaheadFunc.
	SetReadahead(int64)
	// If non-nil, the provided function is called when the implementation needs to know the
	// readahead for the current reader. Calls occur during Reads and Seeks, and while the Client is
	// locked.
	SetReadaheadFunc(ReadaheadFunc)
	// Don't wait for pieces to complete and be verified. Read calls return as soon as they can when
	// the underlying chunks become available.
	SetResponsive()
}

// Piece range by piece index, [begin, end).
type pieceRange struct {
	begin, end pieceIndex
}

type ReadaheadContext struct {
	ContiguousReadStartPos int64
	CurrentPos             int64
}

// Returns the desired readahead for a Reader.
type ReadaheadFunc func(ReadaheadContext) int64

type reader struct {
	t *Torrent
	// Adjust the read/seek window to handle Readers locked to File extents and the like.
	offset, length int64

	// Function to dynamically calculate readahead. If nil, readahead is static.
	readaheadFunc ReadaheadFunc

	// Required when modifying pos and readahead.
	mu sync.Locker

	readahead, pos int64
	// Position that reads have continued contiguously from.
	contiguousReadStartPos int64
	// The cached piece range this reader wants downloaded. The zero value corresponds to nothing.
	// We cache this so that changes can be detected, and bubbled up to the Torrent only as
	// required.
	pieces pieceRange

	// Reads have been initiated since the last seek. This is used to prevent readaheads occurring
	// after a seek or with a new reader at the starting position.
	reading    bool
	responsive bool
}

var _ io.ReadSeekCloser = (*reader)(nil)

func (r *reader) SetResponsive() {
	r.responsive = true
	r.t.cl.event.Broadcast()
}

// Disable responsive mode. TODO: Remove?
func (r *reader) SetNonResponsive() {
	r.responsive = false
	r.t.cl.event.Broadcast()
}

func (r *reader) SetReadahead(readahead int64) {
	r.mu.Lock()
	r.readahead = readahead
	r.readaheadFunc = nil
	r.posChanged()
	r.mu.Unlock()
}

func (r *reader) SetReadaheadFunc(f ReadaheadFunc) {
	r.mu.Lock()
	r.readaheadFunc = f
	r.posChanged()
	r.mu.Unlock()
}

// How many bytes are available to read. Max is the most we could require.
func (r *reader) available(off, max int64) (ret int64) {
	off += r.offset
	for max > 0 {
		req, ok := r.t.offsetRequest(off)
		if !ok {
			break
		}
		if !r.responsive && !r.t.pieceComplete(pieceIndex(req.Index)) {
			break
		}
		if !r.t.haveChunk(req) {
			break
		}
		len1 := int64(req.Length) - (off - r.t.requestOffset(req))
		max -= len1
		ret += len1
		off += len1
	}
	// Ensure that ret hasn't exceeded our original max.
	if max < 0 {
		ret += max
	}
	return
}

// Calculates the pieces this reader wants downloaded, ignoring the cached value at r.pieces.
func (r *reader) piecesUncached() (ret pieceRange) {
	ra := r.readahead
	if r.readaheadFunc != nil {
		ra = r.readaheadFunc(ReadaheadContext{
			ContiguousReadStartPos: r.contiguousReadStartPos,
			CurrentPos:             r.pos,
		})
	}
	if ra < 1 {
		// Needs to be at least 1, because [x, x) means we don't want
		// anything.
		ra = 1
	}
	if !r.reading {
		ra = 0
	}
	if ra > r.length-r.pos {
		ra = r.length - r.pos
	}
	ret.begin, ret.end = r.t.byteRegionPieces(r.torrentOffset(r.pos), ra)
	return
}

func (r *reader) Read(b []byte) (n int, err error) {
	return r.ReadContext(context.Background(), b)
}

func (r *reader) ReadContext(ctx context.Context, b []byte) (n int, err error) {
	if len(b) > 0 {
		r.reading = true
		// TODO: Rework reader piece priorities so we don't have to push updates in to the Client
		// and take the lock here.
		r.mu.Lock()
		r.posChanged()
		r.mu.Unlock()
	}
	n, err = r.readOnceAt(ctx, b, r.pos)
	if n == 0 {
		if err == nil && len(b) > 0 {
			panic("expected error")
		} else {
			return
		}
	}

	r.mu.Lock()
	r.pos += int64(n)
	r.posChanged()
	r.mu.Unlock()
	if r.pos >= r.length {
		err = io.EOF
	} else if err == io.EOF {
		err = io.ErrUnexpectedEOF
	}
	return
}

var closedChan = make(chan struct{})

func init() {
	close(closedChan)
}

// Wait until some data should be available to read. Tickles the client if it isn't. Returns how
// much should be readable without blocking.
func (r *reader) waitAvailable(ctx context.Context, pos, wanted int64, wait bool) (avail int64, err error) {
	t := r.t
	for {
		r.t.cl.rLock()
		avail = r.available(pos, wanted)
		readerCond := t.piece(int((r.offset + pos) / t.info.PieceLength)).readerCond.Signaled()
		r.t.cl.rUnlock()
		if avail != 0 {
			return
		}
		var dontWait <-chan struct{}
		if !wait || wanted == 0 {
			dontWait = closedChan
		}
		select {
		case <-r.t.closed.Done():
			err = errors.New("torrent closed")
			return
		case <-ctx.Done():
			err = ctx.Err()
			return
		case <-r.t.dataDownloadDisallowed.On():
			err = errors.New("torrent data downloading disabled")
		case <-r.t.networkingEnabled.Off():
			err = errors.New("torrent networking disabled")
			return
		case <-dontWait:
			return
		case <-readerCond:
		}
	}
}

// Adds the reader's torrent offset to the reader object offset (for example the reader might be
// constrainted to a particular file within the torrent).
func (r *reader) torrentOffset(readerPos int64) int64 {
	return r.offset + readerPos
}

// Performs at most one successful read to torrent storage.
func (r *reader) readOnceAt(ctx context.Context, b []byte, pos int64) (n int, err error) {
	if pos >= r.length {
		err = io.EOF
		return
	}
	for {
		var avail int64
		avail, err = r.waitAvailable(ctx, pos, int64(len(b)), n == 0)
		if avail == 0 {
			return
		}
		firstPieceIndex := pieceIndex(r.torrentOffset(pos) / r.t.info.PieceLength)
		firstPieceOffset := r.torrentOffset(pos) % r.t.info.PieceLength
		b1 := missinggo.LimitLen(b, avail)
		n, err = r.t.readAt(b1, r.torrentOffset(pos))
		if n != 0 {
			err = nil
			return
		}
		r.t.cl.lock()
		// I think there's a panic here caused by the Client being closed before obtaining this
		// lock. TestDropTorrentWithMmapStorageWhileHashing seems to tickle occasionally in CI.
		func() {
			// Just add exceptions already.
			defer r.t.cl.unlock()
			// TODO: Just reset pieces in the readahead window. This might help
			// prevent thrashing with small caches and file and piece priorities.
			r.log(log.Fstr("error reading torrent %s piece %d offset %d, %d bytes: %v",
				r.t.infoHash.HexString(), firstPieceIndex, firstPieceOffset, len(b1), err))
			if !r.t.updatePieceCompletion(firstPieceIndex) {
				r.log(log.Fstr("piece %d completion unchanged", firstPieceIndex))
			}
			// Update the rest of the piece completions in the readahead window, without alerting to
			// changes (since only the first piece, the one above, could have generated the read error
			// we're currently handling).
			if r.pieces.begin != firstPieceIndex {
				panic(fmt.Sprint(r.pieces.begin, firstPieceIndex))
			}
			for index := r.pieces.begin + 1; index < r.pieces.end; index++ {
				r.t.updatePieceCompletion(index)
			}
		}()
	}
}

// Hodor
func (r *reader) Close() error {
	r.t.cl.lock()
	r.t.deleteReader(r)
	r.t.cl.unlock()
	return nil
}

func (r *reader) posChanged() {
	to := r.piecesUncached()
	from := r.pieces
	if to == from {
		return
	}
	r.pieces = to
	// log.Printf("reader pos changed %v->%v", from, to)
	r.t.readerPosChanged(from, to)
}

func (r *reader) Seek(off int64, whence int) (newPos int64, err error) {
	switch whence {
	case io.SeekStart:
		newPos = off
		r.mu.Lock()
	case io.SeekCurrent:
		r.mu.Lock()
		newPos = r.pos + off
	case io.SeekEnd:
		newPos = r.length + off
		r.mu.Lock()
	default:
		return 0, errors.New("bad whence")
	}
	if newPos != r.pos {
		r.reading = false
		r.pos = newPos
		r.contiguousReadStartPos = newPos
		r.posChanged()
	}
	r.mu.Unlock()
	return
}

func (r *reader) log(m log.Msg) {
	r.t.logger.LogLevel(log.Debug, m.Skip(1))
}

// Implementation inspired by https://news.ycombinator.com/item?id=27019613.
func defaultReadaheadFunc(r ReadaheadContext) int64 {
	return r.CurrentPos - r.ContiguousReadStartPos
}
<|endoftext|>


<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2015 Dylan Carney
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package main

import (
	"fmt"
	"go/ast"
	"log"
	"regexp"
	"strconv"
	"strings"
)

const (
	KWDescription     = "Description"
	KWNotes           = "Notes"
	KWSuccessResponse = "Success Response"
	KWErrorResponse   = "Error Response"
	KWExample         = "Example"
	KWParameter       = "Parameter"
	KWMethod          = "Method"
	KWNone            = "(none)"
)

var (
	apidocMarker    = `(apidoc)\(([^)]+)\):?`                           // apidoc(name), name of at least 1 char
	apidocMarkerRx  = regexp.MustCompile(`^[ \t]*` + apidocMarker)      // the marker at text start
	apidocCommentRx = regexp.MustCompile(`^/[/*][ \t]*` + apidocMarker) // the marker at comment start

	// Example: GET /some/path/:foo
	httpVerbRx = regexp.MustCompile(`(GET|PUT|POST|DELETE|HEAD|OPTIONS|TRACE|CONNECT|PATCH)\s+(/.*)`)

	// Parameter docs follow the pattern:   name [, "required"] [, type]
	// Examples:
	//    foobar
	//    foobar, string
	//    foobar, required, string
	//    foobar, required
	//    foobar, required, array of strings
	parameterRx = regexp.MustCompile(`^([\w-]+)(?:\s*,\s*(required))?(?:\s*,\s*([\w\s]+))?$`)
)

// TODO: change this to a regexp?
func startsWithKeyword(str string) string {

	switch {
	case strings.HasPrefix(str, KWDescription):
		return KWDescription
	case strings.HasPrefix(str, KWSuccessResponse):
		return KWSuccessResponse
	case strings.HasPrefix(str, KWErrorResponse):
		return KWErrorResponse
	case strings.HasPrefix(str, KWExample):
		return KWExample
	case strings.HasPrefix(str, KWParameter):
		return KWParameter
	case strings.HasPrefix(str, KWNotes):
		return KWNotes
	case httpVerbRx.MatchString(str):
		return KWMethod
	}
	return KWNone
}

// stripKeyword removes the given keyword from the first line, as well as any
// leading whitespace-only lines
func stripKeyword(kw string, lines []string) []string {
	lines[0] = strings.TrimSpace(strings.TrimPrefix(lines[0], kw))
	i := -1
	for j, line := range lines {
		if strings.TrimSpace(line) != "" {
			i = j
			break
		}
	}

	if i > 0 {
		return lines[i:]
	}
	return lines
}

// parseKeyword populates Endpoint fields, based on the content in the
// supplied comment lines.
func parseKeyword(e *Endpoint, kw string, lines []string) error {

	switch kw {
	case KWMethod:
		matches := httpVerbRx.FindStringSubmatch(lines[0])
		if len(matches) > 0 {
			e.Method = matches[1]
			e.URLTemplate = matches[2]
		}
	case KWDescription:
		lines = stripKeyword("Description", lines)
		e.Description = strings.Join(lines, "\n")
	case KWParameter:
		lines = stripKeyword("Parameter", lines)
		matches := parameterRx.FindStringSubmatch(lines[0])
		if len(matches) > 0 {
			p := Parameter{
				Name:        matches[1],
				Required:    matches[2] == "required",
				Type:        matches[3],
				Description: strings.Join(lines[1:], " "),
			}

			e.URLParams = append(e.URLParams, p)
		}

	case KWSuccessResponse:
		lines = stripKeyword("Success Response", lines)
		code, err := strconv.Atoi(lines[0])
		if err != nil {
			return err
		}
		e.SuccessResponse = Response{
			Code:    code,
			Content: strings.Join(lines[1:], "\n"),
		}
	case KWErrorResponse:
		lines = stripKeyword("Error Response", lines)
		code, err := strconv.Atoi(lines[0])
		if err != nil {
			return err
		}
		er := Response{
			Code:    code,
			Content: strings.Join(lines[1:], "\n"),
		}
		e.ErrorResponses = append(e.ErrorResponses, er)
	case KWExample:
		lines = stripKeyword("Example", lines)
		e.Examples = append(e.Examples, strings.Join(lines, "\n"))
	default:
		return fmt.Errorf("Unknown keyword: %s", kw)
	}

	return nil
}

// parseEndpoint takes an apidoc body (which consists of one or more
// newline-separated lines) and parses the various keyword sections, populating
// an Endpoint.  The body for each keyword extends until the next keyword,
// or until the end of the body.
func parseEndpoint(body string) (*Endpoint, error) {
	lines := strings.Split(body, "\n")

	e := &Endpoint{}
	i := -1
	lastKw := KWNone
	for j, line := range lines {
		kw := startsWithKeyword(line)
		if kw != KWNone {
			if i >= 0 {
				if err := parseKeyword(e, lastKw, lines[i:j]); err != nil {
					return e, err
				}
			}
			i = j
			lastKw = kw
		}
	}
	if i >= 0 {
		if err := parseKeyword(e, lastKw, lines[i:]); err != nil {
			return e, err
		}
	}

	return e, nil
}

// A reader read a series of CommentGroups, looking for, and attempting to parse
// apidoc text blocks.
type reader struct {
	strict    bool
	endpoints []*Endpoint
}

// readDocs extracts apidoc from comments.  An apidoc must start at the
// beginning of a comment with "apidoc(name):" and is followed by the lines
// that make up the body.  The apidoc ends at the end of the comment group or
// at the start of another apidoc in the same comment group, whichever comes
// first.
func (r *reader) readDocs(comments []*ast.CommentGroup) error {
	for _, group := range comments {
		i := -1 // comment index of most recent note start, valid if >= 0
		list := group.List
		for j, c := range list {
			if apidocCommentRx.MatchString(c.Text) {
				if i >= 0 {
					if err := r.readDoc(list[i:j]); err != nil {
						return err
					}
				}
				i = j
			}
		}
		if i >= 0 {
			if err := r.readDoc(list[i:]); err != nil {
				return err
			}
		}
	}
	return nil
}

// readDoc collects a single api doc from a sequence of comments.
func (r *reader) readDoc(list []*ast.Comment) error {
	text := (&ast.CommentGroup{List: list}).Text()
	if m := apidocMarkerRx.FindStringSubmatchIndex(text); m != nil {
		// The doc body starts after the marker.
		body := text[m[1]:]
		if body != "" {
			e, err := parseEndpoint(body)
			if err != nil {
				return err
			}
			if err := e.Validate(); err != nil {
				if r.strict {
					log.Fatalf("validation error: %s\n", err.Error())
				} else {
					log.Printf("validation error: %s\n", err.Error())
				}
			}
			r.endpoints = append(r.endpoints, e)
		}
	}
	return nil
}
<commit_msg>Adds parsing for "Notes" keyword<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2015 Dylan Carney
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package main

import (
	"fmt"
	"go/ast"
	"log"
	"regexp"
	"strconv"
	"strings"
)

const (
	KWDescription     = "Description"
	KWNotes           = "Notes"
	KWSuccessResponse = "Success Response"
	KWErrorResponse   = "Error Response"
	KWExample         = "Example"
	KWParameter       = "Parameter"
	KWMethod          = "Method"
	KWNone            = "(none)"
)

var (
	apidocMarker    = `(apidoc)\(([^)]+)\):?`                           // apidoc(name), name of at least 1 char
	apidocMarkerRx  = regexp.MustCompile(`^[ \t]*` + apidocMarker)      // the marker at text start
	apidocCommentRx = regexp.MustCompile(`^/[/*][ \t]*` + apidocMarker) // the marker at comment start

	// Example: GET /some/path/:foo
	httpVerbRx = regexp.MustCompile(`(GET|PUT|POST|DELETE|HEAD|OPTIONS|TRACE|CONNECT|PATCH)\s+(/.*)`)

	// Parameter docs follow the pattern:   name [, "required"] [, type]
	// Examples:
	//    foobar
	//    foobar, string
	//    foobar, required, string
	//    foobar, required
	//    foobar, required, array of strings
	parameterRx = regexp.MustCompile(`^([\w-]+)(?:\s*,\s*(required))?(?:\s*,\s*([\w\s]+))?$`)
)

// TODO: change this to a regexp?
func startsWithKeyword(str string) string {

	switch {
	case strings.HasPrefix(str, KWDescription):
		return KWDescription
	case strings.HasPrefix(str, KWSuccessResponse):
		return KWSuccessResponse
	case strings.HasPrefix(str, KWErrorResponse):
		return KWErrorResponse
	case strings.HasPrefix(str, KWExample):
		return KWExample
	case strings.HasPrefix(str, KWParameter):
		return KWParameter
	case strings.HasPrefix(str, KWNotes):
		return KWNotes
	case httpVerbRx.MatchString(str):
		return KWMethod
	}
	return KWNone
}

// stripKeyword removes the given keyword from the first line, as well as any
// leading whitespace-only lines
func stripKeyword(kw string, lines []string) []string {
	lines[0] = strings.TrimSpace(strings.TrimPrefix(lines[0], kw))
	i := -1
	for j, line := range lines {
		if strings.TrimSpace(line) != "" {
			i = j
			break
		}
	}

	if i > 0 {
		return lines[i:]
	}
	return lines
}

// parseKeyword populates Endpoint fields, based on the content in the
// supplied comment lines.
func parseKeyword(e *Endpoint, kw string, lines []string) error {

	switch kw {
	case KWMethod:
		matches := httpVerbRx.FindStringSubmatch(lines[0])
		if len(matches) > 0 {
			e.Method = matches[1]
			e.URLTemplate = matches[2]
		}
	case KWDescription:
		lines = stripKeyword(KWDescription, lines)
		e.Description = strings.Join(lines, "\n")
	case KWParameter:
		lines = stripKeyword(KWParameter, lines)
		matches := parameterRx.FindStringSubmatch(lines[0])
		if len(matches) > 0 {
			p := Parameter{
				Name:        matches[1],
				Required:    matches[2] == "required",
				Type:        matches[3],
				Description: strings.Join(lines[1:], " "),
			}

			e.URLParams = append(e.URLParams, p)
		}

	case KWSuccessResponse:
		lines = stripKeyword(KWSuccessResponse, lines)
		code, err := strconv.Atoi(lines[0])
		if err != nil {
			return err
		}
		e.SuccessResponse = Response{
			Code:    code,
			Content: strings.Join(lines[1:], "\n"),
		}
	case KWErrorResponse:
		lines = stripKeyword(KWErrorResponse, lines)
		code, err := strconv.Atoi(lines[0])
		if err != nil {
			return err
		}
		er := Response{
			Code:    code,
			Content: strings.Join(lines[1:], "\n"),
		}
		e.ErrorResponses = append(e.ErrorResponses, er)
	case KWExample:
		lines = stripKeyword(KWExample, lines)
		e.Examples = append(e.Examples, strings.Join(lines, "\n"))
	case KWNotes:
		lines = stripKeyword(KWNotes, lines)
		e.Notes = strings.Join(lines, "\n")
	default:
		return fmt.Errorf("Unknown keyword: %s", kw)
	}

	return nil
}

// parseEndpoint takes an apidoc body (which consists of one or more
// newline-separated lines) and parses the various keyword sections, populating
// an Endpoint.  The body for each keyword extends until the next keyword,
// or until the end of the body.
func parseEndpoint(body string) (*Endpoint, error) {
	lines := strings.Split(body, "\n")

	e := &Endpoint{}
	i := -1
	lastKw := KWNone
	for j, line := range lines {
		kw := startsWithKeyword(line)
		if kw != KWNone {
			if i >= 0 {
				if err := parseKeyword(e, lastKw, lines[i:j]); err != nil {
					return e, err
				}
			}
			i = j
			lastKw = kw
		}
	}
	if i >= 0 {
		if err := parseKeyword(e, lastKw, lines[i:]); err != nil {
			return e, err
		}
	}

	return e, nil
}

// A reader read a series of CommentGroups, looking for, and attempting to parse
// apidoc text blocks.
type reader struct {
	strict    bool
	endpoints []*Endpoint
}

// readDocs extracts apidoc from comments.  An apidoc must start at the
// beginning of a comment with "apidoc(name):" and is followed by the lines
// that make up the body.  The apidoc ends at the end of the comment group or
// at the start of another apidoc in the same comment group, whichever comes
// first.
func (r *reader) readDocs(comments []*ast.CommentGroup) error {
	for _, group := range comments {
		i := -1 // comment index of most recent note start, valid if >= 0
		list := group.List
		for j, c := range list {
			if apidocCommentRx.MatchString(c.Text) {
				if i >= 0 {
					if err := r.readDoc(list[i:j]); err != nil {
						return err
					}
				}
				i = j
			}
		}
		if i >= 0 {
			if err := r.readDoc(list[i:]); err != nil {
				return err
			}
		}
	}
	return nil
}

// readDoc collects a single api doc from a sequence of comments.
func (r *reader) readDoc(list []*ast.Comment) error {
	text := (&ast.CommentGroup{List: list}).Text()
	if m := apidocMarkerRx.FindStringSubmatchIndex(text); m != nil {
		// The doc body starts after the marker.
		body := text[m[1]:]
		if body != "" {
			e, err := parseEndpoint(body)
			if err != nil {
				return err
			}
			if err := e.Validate(); err != nil {
				if r.strict {
					log.Fatalf("validation error: %s\n", err.Error())
				} else {
					log.Printf("validation error: %s\n", err.Error())
				}
			}
			r.endpoints = append(r.endpoints, e)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"io"
	"log"
	"net"
	"regexp"
	"sync/atomic"

	"github.com/google/gopacket"
	"github.com/google/gopacket/tcpassembly"
	"github.com/google/gopacket/tcpassembly/tcpreader"
	"github.com/kung-foo/certgrep/tls_clone"
)

var DPDSSLServer = regexp.MustCompile(`^(\x16\x03[\x00\x01\x02\x03]..\x02...\x03[\x00\x01\x02\x03]).*`)
var DPDSSLClient = regexp.MustCompile(`^(\x16\x03[\x00\x01\x02\x03]..\x01...\x03[\x00\x01\x02\x03]).*`)

var flow_idx uint64 = 0

type FakeConn struct {
	net.Conn
	flow io.Reader
	idx  uint64
}

func (f *FakeConn) Read(b []byte) (n int, err error) {
	r, err := f.flow.Read(b)
	//log.Printf("F%02d   read  %d %d %v\n", f.idx, len(b), r, err)
	//log.Printf("%02x %02x %02x %02x\n", b[0], b[1], b[2], b[3])
	return r, err
}

func (f *FakeConn) Write(b []byte) (n int, err error) {
	//log.Printf("F%02d   write %d %d\n", f.idx, len(b), len(b))
	return len(b), nil
}

type ReaderFactory struct{}

func (t *ReaderFactory) New(netFlow gopacket.Flow, TCPflow gopacket.Flow) tcpassembly.Stream {
	r := tcpreader.NewReaderStream()
	go handleStream(&r, netFlow, TCPflow)
	return &r
}

func handleStream(r io.Reader, netflow gopacket.Flow, tcpflow gopacket.Flow) {
	defer func() {
		tcpreader.DiscardBytesToEOF(r)
	}()

	idx := atomic.AddUint64(&flow_idx, 1)

	data := bufio.NewReader(r)
	header, err := data.Peek(256)

	if err != nil {
		if err != io.EOF {
			log.Println(err)
		}
		return
	}

	//src, _ := tcpflow.Endpoints()
	if true {
		found_cert := false
		conn := tls_clone.Client(&FakeConn{flow: data, idx: idx}, &tls_clone.Config{InsecureSkipVerify: true})
		conn.Handshake()

		for _, cert := range conn.PeerCertificates() {
			if len(cert.DNSNames) > 0 {
				found_cert = true
				log.Printf("F%04d   %v %v\n", idx, netflow, tcpflow)
				log.Printf("F%04d   %v\n", idx, cert.Subject.CommonName)
				log.Printf("F%04d   %s\n\n\n    ", idx, cert.DNSNames)
			}
			//j, _ := json.Marshal(cert)
			//fmt.Println(string(j))
		}

		if found_cert && !DPDSSLServer.Match(header) {
			//log.Printf("F%04d HHMHMHMHMHMHM  %v %v\n", idx, netflow, tcpflow)
		}
	}
}
<commit_msg>make fakeconn private<commit_after>package main

import (
	"bufio"
	"io"
	"log"
	"net"
	"regexp"
	"sync/atomic"

	"github.com/google/gopacket"
	"github.com/google/gopacket/tcpassembly"
	"github.com/google/gopacket/tcpassembly/tcpreader"
	"github.com/kung-foo/certgrep/tls_clone"
)

var DPDSSLServer = regexp.MustCompile(`^(\x16\x03[\x00\x01\x02\x03]..\x02...\x03[\x00\x01\x02\x03]).*`)
var DPDSSLClient = regexp.MustCompile(`^(\x16\x03[\x00\x01\x02\x03]..\x01...\x03[\x00\x01\x02\x03]).*`)

var flow_idx uint64 = 0

type fakeConn struct {
	net.Conn
	flow io.Reader
	idx  uint64
}

func (f *fakeConn) Read(b []byte) (n int, err error) {
	r, err := f.flow.Read(b)
	//log.Printf("F%02d   read  %d %d %v\n", f.idx, len(b), r, err)
	//log.Printf("%02x %02x %02x %02x\n", b[0], b[1], b[2], b[3])
	return r, err
}

func (f *fakeConn) Write(b []byte) (n int, err error) {
	//log.Printf("F%02d   write %d %d\n", f.idx, len(b), len(b))
	return len(b), nil
}

type ReaderFactory struct{}

func (t *ReaderFactory) New(netFlow gopacket.Flow, TCPflow gopacket.Flow) tcpassembly.Stream {
	r := tcpreader.NewReaderStream()
	go handleStream(&r, netFlow, TCPflow)
	return &r
}

func handleStream(r io.Reader, netflow gopacket.Flow, tcpflow gopacket.Flow) {
	defer func() {
		tcpreader.DiscardBytesToEOF(r)
	}()

	idx := atomic.AddUint64(&flow_idx, 1)

	data := bufio.NewReader(r)
	header, err := data.Peek(256)

	if err != nil {
		if err != io.EOF {
			log.Println(err)
		}
		return
	}

	//src, _ := tcpflow.Endpoints()
	if true {
		found_cert := false
		conn := tls_clone.Client(&fakeConn{flow: data, idx: idx}, &tls_clone.Config{InsecureSkipVerify: true})
		conn.Handshake()

		for _, cert := range conn.PeerCertificates() {
			if len(cert.DNSNames) > 0 {
				found_cert = true
				log.Printf("F%04d   %v %v\n", idx, netflow, tcpflow)
				log.Printf("F%04d   %v\n", idx, cert.Subject.CommonName)
				log.Printf("F%04d   %s\n\n\n    ", idx, cert.DNSNames)
			}
			//j, _ := json.Marshal(cert)
			//fmt.Println(string(j))
		}

		if found_cert && !DPDSSLServer.Match(header) {
			//log.Printf("F%04d HHMHMHMHMHMHM  %v %v\n", idx, netflow, tcpflow)
		}
	}
}
<|endoftext|>