<commit_before>package sirius

import "testing"

type ExtensionConfigAccessorTest struct {
	t   *testing.T
	f   func(string, ExtensionConfig) interface{}
	exp valueMatches
	def interface{}
}

type valueMatches map[string]interface{}

func testExtensionConfigAccessor(test ExtensionConfigAccessorTest) {
	cfg := testingExtensionConfig()

	for field := range cfg {
		a := test.f(field, cfg)

		e, match := test.exp[field]

		if match {
			if a != e {
				test.t.Fatalf("Expected (%s) to resolve into %v, got %v", field, e, a)
			}
			return
		}

		if a != test.def {
			test.t.Fatalf("Expected default value '%v' for (%s), got %v", test.def, field, a)
		}
	}
}

func testingExtensionConfig() ExtensionConfig {
	return ExtensionConfig{
		"int_0":        0,
		"int_1":        1,
		"float_0.0":    0.0,
		"float_1.1":    1.1,
		"bool_true":    true,
		"bool_false":   false,
		"string":       "Hello World!",
		"string_empty": "",
		"list":         []string{"Hit", "Me", "Up"},
	}

}

func TestExtensionConfig_Boolean(t *testing.T) {
	testExtensionConfigAccessor(ExtensionConfigAccessorTest{
		t: t,
		f: func(k string, cfg ExtensionConfig) interface{} {
			return cfg.Boolean(k)
		},
		exp: valueMatches{
			"int_0":      false,
			"int_1":      true,
			"bool_true":  true,
			"bool_false": false,
		},
		def: false,
	})
}

func TestExtensionConfig_Integer(t *testing.T) {
	testExtensionConfigAccessor(ExtensionConfigAccessorTest{
		t: t,
		f: func(k string, cfg ExtensionConfig) interface{} {
			return cfg.Integer(k, 999)
		},
		exp: valueMatches{
			"int_0": 0,
			"int_1": 1,
		},
		def: 999,
	})
}

func TestExtensionConfig_Float(t *testing.T) {
	testExtensionConfigAccessor(ExtensionConfigAccessorTest{
		t: t,
		f: func(k string, cfg ExtensionConfig) interface{} {
			return cfg.Float(k, 999.99)
		},
		exp: valueMatches{
			"float_0.0": 0.0,
			"float_1.1": 1.1,
		},
		def: 999.99,
	})
}

func TestExtensionConfig_String(t *testing.T) {
	testExtensionConfigAccessor(ExtensionConfigAccessorTest{
		t: t,
		f: func(k string, cfg ExtensionConfig) interface{} {
			return cfg.String(k, "Darth Vader")
		},
		exp: valueMatches{
			"string":       "Hello World!",
			"string_empty": "",
		},
		def: "Darth Vader",
	})
}

func TestExtensionConfig_Read(t *testing.T) {
	testExtensionConfigAccessor(ExtensionConfigAccessorTest{
		t: t,
		f: func(k string, cfg ExtensionConfig) interface{} {
			return cfg.Read(k, nil)
		},
		exp: valueMatches{
			"int_0":        0,
			"int_1":        1,
			"float_0.0":    0.0,
			"float_1.1":    1.1,
			"bool_true":    true,
			"bool_false":   false,
			"string":       "Hello World!",
			"string_empty": "",
		},
		def: nil,
	})
}

func TestExtensionConfig_List(t *testing.T) {
	cfg := testingExtensionConfig()

	exp := map[string]bool{
		"list": true,
	}

	for field := range cfg {
		a := cfg.List(field, nil)

		e, match := exp[field]

		if match {
			if a[0] != "Hit" || a[1] != "Me" || a[2] != "Up" {
				t.Fatalf("Expected (%s) to resolve into %v, got %v", field, e, a)
			}
			return
		}

		if a != nil {
			t.Fatalf("Expected default value '%v' for (%s), got %v", nil, field, a)
		}
	}
}
<commit_msg>Improve variable name<commit_after>package sirius

import "testing"

type ExtensionConfigAccessorTest struct {
	t   *testing.T
	f   func(string, ExtensionConfig) interface{}
	exp valueMatches
	def interface{}
}

type valueMatches map[string]interface{}

func testExtensionConfigAccessor(test ExtensionConfigAccessorTest) {
	cfg := testingExtensionConfig()

	for field := range cfg {
		a := test.f(field, cfg)

		e, match := test.exp[field]

		if match {
			if a != e {
				test.t.Fatalf("Expected (%s) to resolve into %v, got %v", field, e, a)
			}
			return
		}

		if a != test.def {
			test.t.Fatalf("Expected default value '%v' for (%s), got %v", test.def, field, a)
		}
	}
}

func testingExtensionConfig() ExtensionConfig {
	return ExtensionConfig{
		"int_0":        0,
		"int_1":        1,
		"float_0.0":    0.0,
		"float_1.1":    1.1,
		"bool_true":    true,
		"bool_false":   false,
		"string_hello": "hello",
		"string_empty": "",
		"list":         []string{"Hit", "Me", "Up"},
	}

}

func TestExtensionConfig_Boolean(t *testing.T) {
	testExtensionConfigAccessor(ExtensionConfigAccessorTest{
		t: t,
		f: func(k string, cfg ExtensionConfig) interface{} {
			return cfg.Boolean(k)
		},
		exp: valueMatches{
			"int_0":      false,
			"int_1":      true,
			"bool_true":  true,
			"bool_false": false,
		},
		def: false,
	})
}

func TestExtensionConfig_Integer(t *testing.T) {
	testExtensionConfigAccessor(ExtensionConfigAccessorTest{
		t: t,
		f: func(k string, cfg ExtensionConfig) interface{} {
			return cfg.Integer(k, 999)
		},
		exp: valueMatches{
			"int_0": 0,
			"int_1": 1,
		},
		def: 999,
	})
}

func TestExtensionConfig_Float(t *testing.T) {
	testExtensionConfigAccessor(ExtensionConfigAccessorTest{
		t: t,
		f: func(k string, cfg ExtensionConfig) interface{} {
			return cfg.Float(k, 999.99)
		},
		exp: valueMatches{
			"float_0.0": 0.0,
			"float_1.1": 1.1,
		},
		def: 999.99,
	})
}

func TestExtensionConfig_String(t *testing.T) {
	testExtensionConfigAccessor(ExtensionConfigAccessorTest{
		t: t,
		f: func(k string, cfg ExtensionConfig) interface{} {
			return cfg.String(k, "Darth Vader")
		},
		exp: valueMatches{
			"string_hello": "hello",
			"string_empty": "",
		},
		def: "Darth Vader",
	})
}

func TestExtensionConfig_Read(t *testing.T) {
	testExtensionConfigAccessor(ExtensionConfigAccessorTest{
		t: t,
		f: func(k string, cfg ExtensionConfig) interface{} {
			return cfg.Read(k, nil)
		},
		exp: valueMatches{
			"int_0":        0,
			"int_1":        1,
			"float_0.0":    0.0,
			"float_1.1":    1.1,
			"bool_true":    true,
			"bool_false":   false,
			"string_hello": "hello",
			"string_empty": "",
		},
		def: nil,
	})
}

func TestExtensionConfig_List(t *testing.T) {
	cfg := testingExtensionConfig()

	exp := map[string]bool{
		"list": true,
	}

	for field := range cfg {
		a := cfg.List(field, nil)

		e, match := exp[field]

		if match {
			if a[0] != "Hit" || a[1] != "Me" || a[2] != "Up" {
				t.Fatalf("Expected (%s) to resolve into %v, got %v", field, e, a)
			}
			return
		}

		if a != nil {
			t.Fatalf("Expected default value '%v' for (%s), got %v", nil, field, a)
		}
	}
}
<|endoftext|>


<commit_before>// vim:set ts=2 sw=2 et ai ft=go:
// Copyright (c) 2012 Toby DiPasquale. See accompanying LICENSE file for
// detailed licensing information.
package main

import (
  "crypto/rand"
  "math/big"
)

const DECK_SIZE = 52

type Deck struct {
  pos   int
  cards [DECK_SIZE]Card
}

// ----- PUBLIC API ---------------------------------------------------------

func NewDeck() *Deck {
  deck := new(Deck)
  deck.pos = 0

  i := 0
  for suit := range suits {
    for rank := range ranks {
      deck.cards[i] = newCard(suits[suit], ranks[rank])
      i++
    }
  }
  return deck
}

//
// Shuffle() uses a Fisher-Yates shuffle:
//
//   http://en.wikipedia.org/wiki/Fisher-Yates_shuffle
//
func (deck *Deck) Shuffle() {
  for i := len(deck.cards) - 1; i > 0; i-- {
    if j := randInt(i + 1); j >= 0 && i != j {
      deck.swap(i, j)
    }
  }
  deck.pos = 0
}

func (deck *Deck) Empty() bool {
  return deck.pos >= len(deck.cards) - 1
}

func (deck *Deck) Deal() (card Card) {
  if deck.Empty() {
    return NoCard
  }
  card = deck.cards[deck.pos]
  deck.pos++
  return card
}

func (deck *Deck) Burn() {
  if !deck.Empty() {
    deck.Deal()
  }
}

// ----- INTERNAL FUNCTIONS -------------------------------------------------

func (deck *Deck) swap(i int, j int) {
  deck.cards[i], deck.cards[j] = deck.cards[j], deck.cards[i]
}

// randInt generates a cryptographically-secure pseudo-random number in the
// range [0,max).
// It returns the generated number on success or -1 if a number could not
// be generated or max was less than 0.
func randInt(max int) int {
  if max < 0 {
    return -1
  }
  m := big.NewInt(int64(max))
  r, e := rand.Int(rand.Reader, m)
  if e != nil {
    return -1
  }
  return int(r.Int64() % int64(max))
}

<commit_msg>should return -1 when max == 0, too<commit_after>// vim:set ts=2 sw=2 et ai ft=go:
// Copyright (c) 2012 Toby DiPasquale. See accompanying LICENSE file for
// detailed licensing information.
package main

import (
  "crypto/rand"
  "math/big"
)

const DECK_SIZE = 52

type Deck struct {
  pos   int
  cards [DECK_SIZE]Card
}

// ----- PUBLIC API ---------------------------------------------------------

func NewDeck() *Deck {
  deck := new(Deck)
  deck.pos = 0

  i := 0
  for suit := range suits {
    for rank := range ranks {
      deck.cards[i] = newCard(suits[suit], ranks[rank])
      i++
    }
  }
  return deck
}

//
// Shuffle() uses a Fisher-Yates shuffle:
//
//   http://en.wikipedia.org/wiki/Fisher-Yates_shuffle
//
func (deck *Deck) Shuffle() {
  for i := len(deck.cards) - 1; i > 0; i-- {
    if j := randInt(i + 1); j >= 0 && i != j {
      deck.swap(i, j)
    }
  }
  deck.pos = 0
}

func (deck *Deck) Empty() bool {
  return deck.pos >= len(deck.cards) - 1
}

func (deck *Deck) Deal() (card Card) {
  if deck.Empty() {
    return NoCard
  }
  card = deck.cards[deck.pos]
  deck.pos++
  return card
}

func (deck *Deck) Burn() {
  if !deck.Empty() {
    deck.Deal()
  }
}

// ----- INTERNAL FUNCTIONS -------------------------------------------------

func (deck *Deck) swap(i int, j int) {
  deck.cards[i], deck.cards[j] = deck.cards[j], deck.cards[i]
}

// randInt generates a cryptographically-secure pseudo-random number in the
// range [0,max).
// It returns the generated number on success or -1 if a number could not
// be generated or max was less than or equal to 0.
func randInt(max int) int {
  if max <= 0 {
    return -1
  }
  m := big.NewInt(int64(max))
  r, e := rand.Int(rand.Reader, m)
  if e != nil {
    return -1
  }
  return int(r.Int64() % int64(max))
}

<|endoftext|>


<commit_before>package boardgame

import (
	"errors"
	"strconv"
)

//TODO: consider making Deck be an interface again (in some cases it
//might be nice to be able to cast the Deck directly to its underlying type to
//minimize later casts)

//A Deck represents an immutable collection of a certain type of components.
//Every component lives in one deck. 1 or more Stacks index into every Deck,
//and cover every item in the deck, with no items in more than one deck. The
//zero-value of Deck is useful. The Deck will not return items until it has
//been added to a ComponentChest, which helps enforce that Decks' values never
//change. Create a new Deck with NewDeck()
type Deck struct {
	chest *ComponentChest
	//Name is only set when it's added to the component chest.
	name string
	//Components should only ever be added at initalization time. After
	//initalization, Components should be read-only.
	components             []*Component
	shadowValues           SubState
	vendedShadowComponents map[int]*Component
	//TODO: protect shadowComponents cache with mutex to make threadsafe.
}

const genericComponentSentinel = -2

func NewDeck() *Deck {
	return &Deck{
		vendedShadowComponents: make(map[int]*Component),
	}
}

//AddComponent adds a new component with the given values to the next spot in
//the deck. If the deck has already been added to a componentchest, this will
//do nothing.
func (d *Deck) AddComponent(v SubState) {
	if d.chest != nil {
		return
	}

	c := &Component{
		Deck:      d,
		DeckIndex: len(d.components),
		Values:    v,
	}

	d.components = append(d.components, c)
}

//AddComponentMulti is like AddComponent, but creates multiple versions of the
//same component. The exact same ComponentValues will be re-used, which is
//reasonable becasue components are read-only anyway.
func (d *Deck) AddComponentMulti(v SubState, count int) {
	for i := 0; i < count; i++ {
		d.AddComponent(v)
	}
}

//Components returns a list of Components in order in this deck, but only if
//this Deck has already been added to its ComponentChest.
func (d *Deck) Components() []*Component {
	if d.chest == nil {
		return nil
	}
	return d.components
}

//Chest points back to the chest we're part of.
func (d *Deck) Chest() *ComponentChest {
	return d.chest
}

func (d *Deck) Name() string {
	return d.name
}

//ComponentAt returns the component at a given index. It handles empty indexes
//and shadow indexes correctly.
func (d *Deck) ComponentAt(index int) *Component {
	if d.chest == nil {
		return nil
	}
	if index >= len(d.components) {
		return nil
	}
	if index >= 0 {
		return d.components[index]
	}

	//d.ShadowComponent handles all negative indexes correctly, which is what
	//we have.
	return d.ShadowComponent(index)

}

//SetShadowValues sets the SubState to return for every shadow
//component that is returned. May only be set before added to a chest. Should
//generally be the same shape of componentValues as used for other components
//in the deck.
func (d *Deck) SetShadowValues(v SubState) {
	if d.chest != nil {
		return
	}
	d.shadowValues = v
}

//ShadowComponent takes an index that is negative and returns a component that
//is empty but when compared to the result of previous calls to
//ShadowComponent with that index will have equality. This is important for
//sanitized states, where depending on the policy for that property, the stack
//might have its order revealed but not its contents, which requires throwaway
//but stable indexes.
func (d *Deck) ShadowComponent(index int) *Component {
	if index >= 0 {
		return nil
	}
	if index == emptyIndexSentinel {
		return nil
	}

	shadow, ok := d.vendedShadowComponents[index]

	if !ok {
		shadow = &Component{
			Deck:      d,
			DeckIndex: index,
			Values:    d.shadowValues,
		}
		d.vendedShadowComponents[index] = shadow
	}

	return shadow

}

//GenericComponent returns the component that is considereed fully generic for
//this deck. This is the component that every component will be if a Stack is
//sanitized with PolicyLen, for example. If you want to figure out if a Stack
//was sanitized according to that policy, you can compare the component to
//this.
func (d *Deck) GenericComponent() *Component {
	return d.ShadowComponent(genericComponentSentinel)
}

//finish is called when the deck is added to a component chest. It signifies that no more items may be added.
func (d *Deck) finish(chest *ComponentChest, name string) error {

	//TODO: this should use the generic reader tester infrastructure we're building out in #464
	for i, c := range d.components {
		if c.Values == nil {
			continue
		}
		for propName, propType := range c.Values.Reader().Props() {
			if propType == TypeEnumVar || propType == TypeGrowableStack || propType == TypeSizedStack || propType == TypeTimer {
				return errors.New("Component " + strconv.Itoa(i) + " has an illegal property type for property " + propName)
			}
			if propType == TypeEnumConst {
				enumConst, err := c.Values.Reader().EnumConstProp(propName)
				if err != nil {
					return errors.New("Unexpected error for component " + strconv.Itoa(i) + ": " + err.Error())
				}
				if enumConst == nil {
					return errors.New("Component " + strconv.Itoa(i) + " had a nil enum.Const for property " + propName)
				}
			}
		}
	}

	d.chest = chest
	//If a deck has a name, it cannot receive any more items.
	d.name = name
	return nil
}
<commit_msg>deck.finish() now uses the readerValidator logic. Part of #464.<commit_after>package boardgame

import (
	"errors"
	"strconv"
)

//TODO: consider making Deck be an interface again (in some cases it
//might be nice to be able to cast the Deck directly to its underlying type to
//minimize later casts)

//A Deck represents an immutable collection of a certain type of components.
//Every component lives in one deck. 1 or more Stacks index into every Deck,
//and cover every item in the deck, with no items in more than one deck. The
//zero-value of Deck is useful. The Deck will not return items until it has
//been added to a ComponentChest, which helps enforce that Decks' values never
//change. Create a new Deck with NewDeck()
type Deck struct {
	chest *ComponentChest
	//Name is only set when it's added to the component chest.
	name string
	//Components should only ever be added at initalization time. After
	//initalization, Components should be read-only.
	components             []*Component
	shadowValues           SubState
	vendedShadowComponents map[int]*Component
	//TODO: protect shadowComponents cache with mutex to make threadsafe.
}

const genericComponentSentinel = -2

func NewDeck() *Deck {
	return &Deck{
		vendedShadowComponents: make(map[int]*Component),
	}
}

//AddComponent adds a new component with the given values to the next spot in
//the deck. If the deck has already been added to a componentchest, this will
//do nothing.
func (d *Deck) AddComponent(v SubState) {
	if d.chest != nil {
		return
	}

	c := &Component{
		Deck:      d,
		DeckIndex: len(d.components),
		Values:    v,
	}

	d.components = append(d.components, c)
}

//AddComponentMulti is like AddComponent, but creates multiple versions of the
//same component. The exact same ComponentValues will be re-used, which is
//reasonable becasue components are read-only anyway.
func (d *Deck) AddComponentMulti(v SubState, count int) {
	for i := 0; i < count; i++ {
		d.AddComponent(v)
	}
}

//Components returns a list of Components in order in this deck, but only if
//this Deck has already been added to its ComponentChest.
func (d *Deck) Components() []*Component {
	if d.chest == nil {
		return nil
	}
	return d.components
}

//Chest points back to the chest we're part of.
func (d *Deck) Chest() *ComponentChest {
	return d.chest
}

func (d *Deck) Name() string {
	return d.name
}

//ComponentAt returns the component at a given index. It handles empty indexes
//and shadow indexes correctly.
func (d *Deck) ComponentAt(index int) *Component {
	if d.chest == nil {
		return nil
	}
	if index >= len(d.components) {
		return nil
	}
	if index >= 0 {
		return d.components[index]
	}

	//d.ShadowComponent handles all negative indexes correctly, which is what
	//we have.
	return d.ShadowComponent(index)

}

//SetShadowValues sets the SubState to return for every shadow
//component that is returned. May only be set before added to a chest. Should
//generally be the same shape of componentValues as used for other components
//in the deck.
func (d *Deck) SetShadowValues(v SubState) {
	if d.chest != nil {
		return
	}
	d.shadowValues = v
}

//ShadowComponent takes an index that is negative and returns a component that
//is empty but when compared to the result of previous calls to
//ShadowComponent with that index will have equality. This is important for
//sanitized states, where depending on the policy for that property, the stack
//might have its order revealed but not its contents, which requires throwaway
//but stable indexes.
func (d *Deck) ShadowComponent(index int) *Component {
	if index >= 0 {
		return nil
	}
	if index == emptyIndexSentinel {
		return nil
	}

	shadow, ok := d.vendedShadowComponents[index]

	if !ok {
		shadow = &Component{
			Deck:      d,
			DeckIndex: index,
			Values:    d.shadowValues,
		}
		d.vendedShadowComponents[index] = shadow
	}

	return shadow

}

//GenericComponent returns the component that is considereed fully generic for
//this deck. This is the component that every component will be if a Stack is
//sanitized with PolicyLen, for example. If you want to figure out if a Stack
//was sanitized according to that policy, you can compare the component to
//this.
func (d *Deck) GenericComponent() *Component {
	return d.ShadowComponent(genericComponentSentinel)
}

var illegalComponentValuesProps = map[PropertyType]bool{
	TypeEnumVar:       true,
	TypeGrowableStack: true,
	TypeSizedStack:    true,
	TypeTimer:         true,
}

//finish is called when the deck is added to a component chest. It signifies that no more items may be added.
func (d *Deck) finish(chest *ComponentChest, name string) error {

	for i, c := range d.components {
		if c.Values == nil {
			continue
		}
		validator, err := newReaderValidator(c.Values.Reader(), c.Values, illegalComponentValuesProps, chest)
		if err != nil {
			return errors.New("Component " + strconv.Itoa(i) + "failed to validate: " + err.Error())
		}
		if err := validator.Valid(c.Values.Reader()); err != nil {
			return errors.New("Component " + strconv.Itoa(i) + " failed to validate: " + err.Error())
		}
	}

	d.chest = chest
	//If a deck has a name, it cannot receive any more items.
	d.name = name
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	oss "github.com/kyf/oss_sdk/lib"
	"net/http"
)

func init() {
	var group string = "demo"
	myHandlers[fmt.Sprintf("/%s", group)] = func(w http.ResponseWriter, r *http.Request) {
		oss.Init(OSS_ACCESS_ID, OSS_ACCESS_KEY, logger)
		/*
			oss.RmDir(BUCKET, "/sight")
			oss.Remove(BUCKET, "/sight/2015_10_14_14447985047668171981298498081.jpg")
			oss.RemoveBucket(BUCKET)
		*/
		//oss.CreateBucket(BUCKET)
		w.Write([]byte("done"))
	}
}
<commit_msg>debug<commit_after>package main

import (
	"fmt"
	oss "github.com/kyf/oss_sdk/lib"
	"net/http"
)

func init() {
	var group string = "demo"
	myHandlers[fmt.Sprintf("/%s", group)] = func(w http.ResponseWriter, r *http.Request) {
		oss.Init(OSS_ACCESS_ID, OSS_ACCESS_KEY, logger)
		/*
			oss.RmDir(BUCKET, "/sight")
			oss.Remove(BUCKET, "/sight/2015_10_14_14447985047668171981298498081.jpg")
			oss.RemoveBucket(BUCKET)
		*/
		//oss.CreateBucket(BUCKET)
		oss.Remove(BUCKET, "/sight/2015_10_14_14447987323049666151298498081.jpg")
		oss.Remove(BUCKET, "/sight/2015_10_14_14447988769221438841298498081.jpg")
		oss.Remove(BUCKET, "/sight/2015_10_14_14448006768876860831298498081.jpg")
		oss.Remove(BUCKET, "/sight/2015_10_14_14448011257944723231298498081.jpg")
		w.Write([]byte("done"))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io"
	"net"
	"runtime"
	"strings"

	"time"

	"strconv"

	"github.com/sybrexsys/RapidKV/datamodel"
)

var telnetstop chan struct{}

type clientConnection struct {
	answers         []datamodel.CustomDataType
	answersSize     int
	needQuit        bool
	currentDatabase *Database
	authorized      bool
	inMulti         bool
	queue           datamodel.DataArray
}

func (cc *clientConnection) setCapacity(newCapacity int) {
	tmp := make([]datamodel.CustomDataType, newCapacity)
	copy(tmp, cc.answers)
	cc.answers = tmp
}

func (cc *clientConnection) grow() {
	var Delta int
	Cap := len(cc.answers)
	if Cap > 64 {
		Delta = Cap / 4
	} else {
		if Cap > 8 {
			Delta = 16
		} else {
			Delta = 4
		}
	}
	cc.setCapacity(Cap + Delta)
}

func (cc *clientConnection) pushAnswer(answer datamodel.CustomDataType) {
	if cc.answersSize == len(cc.answers) {
		cc.grow()
	}
	cc.answers[cc.answersSize] = answer
	cc.answersSize++
}
func (cc *clientConnection) popAnswers(writer *bufio.Writer) error {
	for i := 0; i < cc.answersSize; i++ {
		_, err := writer.Write(datamodel.ConvertToRASP(cc.answers[i]))
		if err != nil {
			return err
		}
		cc.answers[i] = nil
	}
	writer.Flush()
	if cc.answersSize > 1024 {
		cc.answers = make([]datamodel.CustomDataType, 100)
	}
	cc.answersSize = 0
	return nil
}

func (cc *clientConnection) processOneRESPCommandWithoutLock(command datamodel.DataArray) datamodel.CustomDataType {
	comName := command.Get(0).(datamodel.DataString).Get()
	f, ok := commandList[strings.ToLower(comName)]
	if !ok {
		return datamodel.CreateError("ERR Unknown command")
	}
	if needAuth && !cc.authorized {
		return datamodel.CreateError("ERR Not authorized")
	}

	if f == nil {
		return datamodel.CreateError("ERR Command not implemented")
	}
	if command.Count() < 2 {
		return datamodel.CreateError("ERR Unknown parameter")
	}
	key, err := getKey(command, 1)
	if err != nil {
		return datamodel.CreateError("ERR Unknown parameter")
	}
	command.Remove(0)
	command.Remove(0)
	return f(cc.currentDatabase, key, command)
}

func (cc *clientConnection) processTransaction() datamodel.CustomDataType {
	cnt := cc.queue.Count()
	answers := datamodel.CreateArray(cnt)
	cc.currentDatabase.Lock()
	defer func() {
		cc.inMulti = false
		cc.queue = datamodel.CreateArray(10)
		cc.currentDatabase.Unlock()
	}()
	for i := 0; i < cnt; i++ {
		answers.Add(cc.processOneRESPCommandWithoutLock(cc.queue.Get(i).(datamodel.DataArray)))
	}
	return answers
}

func (cc *clientConnection) processOneRESPCommand(command datamodel.CustomDataType) datamodel.CustomDataType {
	arr, ok := command.(datamodel.DataArray)
	if !ok {
		return datamodel.CreateError("ERR Invalid command")
	}
	if arr.Count() < 1 {
		return datamodel.CreateError("ERR Invalid command")
	}
	comdat := arr.Get(0)
	str, okstr := comdat.(datamodel.DataString)
	if !okstr {
		return datamodel.CreateError("ERR Invalid command")
	}
	commandName := strings.ToLower(str.Get())

	concom, isConnectionCommand := connectionCommands[commandName]
	if isConnectionCommand {
		if concom.needAuth && !cc.authorized {
			return datamodel.CreateError("ERR Not authorized")
		}
		return concom.function(cc, arr)
	}
	if commandName == "multi" {
		cc.queue = datamodel.CreateArray(10)
		cc.inMulti = true
		return datamodel.CreateSimpleString("OK")
	}
	if commandName == "discard" {
		cc.inMulti = false
		cc.queue = datamodel.CreateArray(10)
		return datamodel.CreateSimpleString("OK")
	}
	if commandName == "exec" {
		return cc.processTransaction()
	}
	if cc.inMulti {
		cc.queue.Add(command)
		return datamodel.CreateSimpleString("QUEUED")
	}
	cc.currentDatabase.RLock()
	defer cc.currentDatabase.RUnlock()
	return cc.processOneRESPCommandWithoutLock(arr)
}

func checkEOF(reader *bufio.Reader) (bool, error) {
	_, err := reader.ReadByte()
	if err == nil {
		reader.UnreadByte()
		return false, nil
	}
	if err == io.EOF {
		return true, nil
	}
	return false, err
}

func saveOneDataToOut(answer datamodel.CustomDataType, writer *bufio.Writer) error {
	_, err := writer.Write(datamodel.ConvertToRASP(answer))
	return err
}

func processRESPConnection(c net.Conn) {
	var request datamodel.CustomDataType
	defer func() {
		if e := recover(); e != nil {
			buf := make([]byte, 4096)
			n := runtime.Stack(buf, false)
			buf = buf[0:n]
			fmt.Printf("client run panic %s:%v\rLast command to server was:%s", buf, e, datamodel.DataObjectToString(request))
		}
		c.Close()
		notifier.Done()
	}()
	notifier.Add(1)
	reader := bufio.NewReader(c)
	writer := bufio.NewWriter(c)
	cc := &clientConnection{
		answers:         make([]datamodel.CustomDataType, 100),
		answersSize:     0,
		currentDatabase: firstDatabase,
		authorized:      !needAuth,
	}
	go func() {
		<-quit
		c.SetReadDeadline(time.Now().Add(0))
	}()
	for {
		c.SetReadDeadline(time.Now().Add(time.Millisecond * 200))
		ch, err := reader.ReadByte()
		select {
		case <-quit:
			return
		default:
		}
		if err != nil {
			netErr, ok := err.(net.Error)
			if ok && netErr.Timeout() && netErr.Temporary() {
				if cc.answersSize != 0 {
					if cc.popAnswers(writer) != nil {
						return
					}
				}
				continue
			}
		} else {
			if ch == '\r' {
				reader.ReadBytes(10)
				continue
			} else {
				reader.UnreadByte()
				c.SetReadDeadline(time.Now().Add(time.Second * 50))
			}
		}
		request, err = datamodel.LoadRespFromIO(reader, true)
		if err != nil {
			parseError, ok := err.(datamodel.ParseError)
			if !ok {
				return
			}
			answer := datamodel.CreateError(parseError.Error())
			cc.pushAnswer(answer)
		}
		answer := cc.processOneRESPCommand(request)
		cc.pushAnswer(answer)
		if cc.needQuit {
			cc.popAnswers(writer)
			break
		}
	}
}

func startRESPListener() {
	defer func() {
		fmt.Println("stopped RESP server...")
		notifier.Done()
	}()
	notifier.Add(1)

mainloop:
	for {
		conn, err := tcplistener.Accept()
		select {
		case <-quit:
			fmt.Println("listener stop signal was received")
			break mainloop
		default:
		}
		if err != nil {
			fmt.Println(err.Error())
			return
		}
		go processRESPConnection(conn)
	}
}

func authCommand(cc *clientConnection, command datamodel.DataArray) datamodel.CustomDataType {
	if needAuth {
		item := command.Get(1)
		str, ok := item.(datamodel.DataString)
		if !ok {
			datamodel.CreateError("ERR Invalid parameter")
		}
		cc.authorized = str.Get() == cfg.AuthPass
		if !cc.authorized {
			return datamodel.CreateError("ERR Invalid password")
		}
	}
	return datamodel.CreateSimpleString("OK")
}

func quitCommand(cc *clientConnection, command datamodel.DataArray) datamodel.CustomDataType {
	cc.needQuit = true
	return datamodel.CreateSimpleString("OK")
}

func selectCommand(cc *clientConnection, command datamodel.DataArray) datamodel.CustomDataType {
	idx, err := getInt(command, 1)
	if err != nil {
		return datamodel.CreateError("ERR Invalid parameter")
	}
	cc.currentDatabase = getDataBase(idx)
	return datamodel.CreateSimpleString("OK")
}

func echoCommand(cc *clientConnection, command datamodel.DataArray) datamodel.CustomDataType {
	key, err := getKey(command, 1)
	if err != nil {
		return datamodel.CreateError("ERR Invalid parameter")
	}
	return datamodel.CreateString(key)

}

func pingCommand(cc *clientConnection, command datamodel.DataArray) datamodel.CustomDataType {
	item := command.Get(1)
	switch value := item.(type) {
	case datamodel.DataNull:
		return datamodel.CreateSimpleString("PONG")
	case datamodel.DataString:
		return datamodel.CreateString(value.Get())
	case datamodel.DataInt:
		return datamodel.CreateString(strconv.Itoa(value.Get()))
	default:
		return datamodel.CreateError("ERR Invalid parameter")
	}
}

type connectionCommand struct {
	needAuth bool
	function func(cc *clientConnection, command datamodel.DataArray) datamodel.CustomDataType
}

var connectionCommands = map[string]connectionCommand{
	"auth": {
		needAuth: false,
		function: authCommand,
	},
	"select": {
		needAuth: true,
		function: selectCommand,
	},
	"echo": {
		needAuth: false,
		function: echoCommand,
	},
	"ping": {
		needAuth: false,
		function: pingCommand,
	},
	"quit": {
		needAuth: false,
		function: quitCommand,
	},
}
<commit_msg>error notification appended<commit_after>package main

import (
	"bufio"
	"fmt"
	"io"
	"net"
	"runtime"
	"strings"

	"time"

	"strconv"

	"github.com/sybrexsys/RapidKV/datamodel"
)

var telnetstop chan struct{}

type clientConnection struct {
	answers         []datamodel.CustomDataType
	answersSize     int
	needQuit        bool
	currentDatabase *Database
	authorized      bool
	inMulti         bool
	queue           datamodel.DataArray
}

func (cc *clientConnection) setCapacity(newCapacity int) {
	tmp := make([]datamodel.CustomDataType, newCapacity)
	copy(tmp, cc.answers)
	cc.answers = tmp
}

func (cc *clientConnection) grow() {
	var Delta int
	Cap := len(cc.answers)
	if Cap > 64 {
		Delta = Cap / 4
	} else {
		if Cap > 8 {
			Delta = 16
		} else {
			Delta = 4
		}
	}
	cc.setCapacity(Cap + Delta)
}

func (cc *clientConnection) pushAnswer(answer datamodel.CustomDataType) {
	if cc.answersSize == len(cc.answers) {
		cc.grow()
	}
	cc.answers[cc.answersSize] = answer
	cc.answersSize++
}
func (cc *clientConnection) popAnswers(writer *bufio.Writer) error {
	for i := 0; i < cc.answersSize; i++ {
		_, err := writer.Write(datamodel.ConvertToRASP(cc.answers[i]))
		if err != nil {
			return err
		}
		cc.answers[i] = nil
	}
	writer.Flush()
	if cc.answersSize > 1024 {
		cc.answers = make([]datamodel.CustomDataType, 100)
	}
	cc.answersSize = 0
	return nil
}

func (cc *clientConnection) processOneRESPCommandWithoutLock(command datamodel.DataArray) datamodel.CustomDataType {
	comName := command.Get(0).(datamodel.DataString).Get()
	f, ok := commandList[strings.ToLower(comName)]
	if !ok {
		return datamodel.CreateError("ERR Unknown command")
	}
	if needAuth && !cc.authorized {
		return datamodel.CreateError("ERR Not authorized")
	}

	if f == nil {
		return datamodel.CreateError("ERR Command not implemented")
	}
	if command.Count() < 2 {
		return datamodel.CreateError("ERR Unknown parameter")
	}
	key, err := getKey(command, 1)
	if err != nil {
		return datamodel.CreateError("ERR Unknown parameter")
	}
	command.Remove(0)
	command.Remove(0)
	return f(cc.currentDatabase, key, command)
}

func (cc *clientConnection) processTransaction() datamodel.CustomDataType {
	cnt := cc.queue.Count()
	answers := datamodel.CreateArray(cnt)
	cc.currentDatabase.Lock()
	defer func() {
		cc.inMulti = false
		cc.queue = datamodel.CreateArray(10)
		cc.currentDatabase.Unlock()
	}()
	for i := 0; i < cnt; i++ {
		answers.Add(cc.processOneRESPCommandWithoutLock(cc.queue.Get(i).(datamodel.DataArray)))
	}
	return answers
}

func (cc *clientConnection) processOneRESPCommand(command datamodel.CustomDataType) datamodel.CustomDataType {
	arr, ok := command.(datamodel.DataArray)
	if !ok {
		return datamodel.CreateError("ERR Invalid command")
	}
	if arr.Count() < 1 {
		return datamodel.CreateError("ERR Invalid command")
	}
	comdat := arr.Get(0)
	str, okstr := comdat.(datamodel.DataString)
	if !okstr {
		return datamodel.CreateError("ERR Invalid command")
	}
	commandName := strings.ToLower(str.Get())

	concom, isConnectionCommand := connectionCommands[commandName]
	if isConnectionCommand {
		if concom.needAuth && !cc.authorized {
			return datamodel.CreateError("ERR Not authorized")
		}
		return concom.function(cc, arr)
	}
	if commandName == "multi" {
		cc.queue = datamodel.CreateArray(10)
		cc.inMulti = true
		return datamodel.CreateSimpleString("OK")
	}
	if commandName == "discard" {
		cc.inMulti = false
		cc.queue = datamodel.CreateArray(10)
		return datamodel.CreateSimpleString("OK")
	}
	if commandName == "exec" {
		return cc.processTransaction()
	}
	if cc.inMulti {
		cc.queue.Add(command)
		return datamodel.CreateSimpleString("QUEUED")
	}
	cc.currentDatabase.RLock()
	defer cc.currentDatabase.RUnlock()
	return cc.processOneRESPCommandWithoutLock(arr)
}

func checkEOF(reader *bufio.Reader) (bool, error) {
	_, err := reader.ReadByte()
	if err == nil {
		reader.UnreadByte()
		return false, nil
	}
	if err == io.EOF {
		return true, nil
	}
	return false, err
}

func saveOneDataToOut(answer datamodel.CustomDataType, writer *bufio.Writer) error {
	_, err := writer.Write(datamodel.ConvertToRASP(answer))
	return err
}

func processRESPConnection(c net.Conn) {
	var request datamodel.CustomDataType
	defer func() {
		if e := recover(); e != nil {
			buf := make([]byte, 4096)
			n := runtime.Stack(buf, false)
			buf = buf[0:n]
			fmt.Printf("client run panic %s:%v\rLast command to server was:%s", buf, e, datamodel.DataObjectToString(request))
		}
		c.Close()
		notifier.Done()
	}()
	notifier.Add(1)
	reader := bufio.NewReader(c)
	writer := bufio.NewWriter(c)
	cc := &clientConnection{
		answers:         make([]datamodel.CustomDataType, 100),
		answersSize:     0,
		currentDatabase: firstDatabase,
		authorized:      !needAuth,
	}
	go func() {
		<-quit
		c.SetReadDeadline(time.Now().Add(0))
	}()
	for {
		c.SetReadDeadline(time.Now().Add(time.Millisecond * 200))
		ch, err := reader.ReadByte()
		select {
		case <-quit:
			return
		default:
		}
		if err != nil {
			netErr, ok := err.(net.Error)
			if ok && netErr.Timeout() && netErr.Temporary() {
				if cc.answersSize != 0 {
					if err := cc.popAnswers(writer); err != nil {
						fmt.Printf("Client connection lost/ Error:%s", err.Error())
						return
					}
				}
				continue
			}
		} else {
			if ch == '\r' {
				reader.ReadBytes(10)
				continue
			} else {
				reader.UnreadByte()
				c.SetReadDeadline(time.Now().Add(time.Second * 50))
			}
		}
		request, err = datamodel.LoadRespFromIO(reader, true)
		if err != nil {
			parseError, ok := err.(datamodel.ParseError)
			if !ok {
				fmt.Printf("Client connection lost/ Error:%s", err.Error())
				return
			}
			answer := datamodel.CreateError(parseError.Error())
			cc.pushAnswer(answer)
		}
		answer := cc.processOneRESPCommand(request)
		cc.pushAnswer(answer)
		if cc.needQuit {
			cc.popAnswers(writer)
			break
		}
	}
}

func startRESPListener() {
	defer func() {
		fmt.Println("stopped RESP server...")
		notifier.Done()
	}()
	notifier.Add(1)

mainloop:
	for {
		conn, err := tcplistener.Accept()
		select {
		case <-quit:
			fmt.Println("listener stop signal was received")
			break mainloop
		default:
		}
		if err != nil {
			fmt.Println(err.Error())
			return
		}
		go processRESPConnection(conn)
	}
}

func authCommand(cc *clientConnection, command datamodel.DataArray) datamodel.CustomDataType {
	if needAuth {
		item := command.Get(1)
		str, ok := item.(datamodel.DataString)
		if !ok {
			datamodel.CreateError("ERR Invalid parameter")
		}
		cc.authorized = str.Get() == cfg.AuthPass
		if !cc.authorized {
			return datamodel.CreateError("ERR Invalid password")
		}
	}
	return datamodel.CreateSimpleString("OK")
}

func quitCommand(cc *clientConnection, command datamodel.DataArray) datamodel.CustomDataType {
	cc.needQuit = true
	return datamodel.CreateSimpleString("OK")
}

func selectCommand(cc *clientConnection, command datamodel.DataArray) datamodel.CustomDataType {
	idx, err := getInt(command, 1)
	if err != nil {
		return datamodel.CreateError("ERR Invalid parameter")
	}
	cc.currentDatabase = getDataBase(idx)
	return datamodel.CreateSimpleString("OK")
}

func echoCommand(cc *clientConnection, command datamodel.DataArray) datamodel.CustomDataType {
	key, err := getKey(command, 1)
	if err != nil {
		return datamodel.CreateError("ERR Invalid parameter")
	}
	return datamodel.CreateString(key)

}

func pingCommand(cc *clientConnection, command datamodel.DataArray) datamodel.CustomDataType {
	item := command.Get(1)
	switch value := item.(type) {
	case datamodel.DataNull:
		return datamodel.CreateSimpleString("PONG")
	case datamodel.DataString:
		return datamodel.CreateString(value.Get())
	case datamodel.DataInt:
		return datamodel.CreateString(strconv.Itoa(value.Get()))
	default:
		return datamodel.CreateError("ERR Invalid parameter")
	}
}

type connectionCommand struct {
	needAuth bool
	function func(cc *clientConnection, command datamodel.DataArray) datamodel.CustomDataType
}

var connectionCommands = map[string]connectionCommand{
	"auth": {
		needAuth: false,
		function: authCommand,
	},
	"select": {
		needAuth: true,
		function: selectCommand,
	},
	"echo": {
		needAuth: false,
		function: echoCommand,
	},
	"ping": {
		needAuth: false,
		function: pingCommand,
	},
	"quit": {
		needAuth: false,
		function: quitCommand,
	},
}
<|endoftext|>


<commit_before>package chroot

import (
	"context"
	"io/ioutil"
	"net/http"
	"reflect"
	"regexp"
	"strings"
	"testing"

	"github.com/hashicorp/packer/builder/azure/common/client"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"

	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
	"github.com/Azure/go-autorest/autorest"
)

func TestStepCreateNewDisk_Run(t *testing.T) {
	tests := []struct {
		name                  string
		fields                StepCreateNewDiskset
		expectedPutDiskBodies []string
		want                  multistep.StepAction
		verifyDiskset         *Diskset
	}{
		{
			name: "from disk",
			fields: StepCreateNewDiskset{
				OSDiskID:                 "/subscriptions/SubscriptionID/resourcegroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName",
				OSDiskSizeGB:             42,
				OSDiskStorageAccountType: string(compute.PremiumLRS),
				HyperVGeneration:         string(compute.V1),
				Location:                 "westus",
				SourceOSDiskResourceID:   "SourceDisk",
			},
			expectedPutDiskBodies: []string{`
				{
					"location": "westus",
					"properties": {
						"osType": "Linux",
						"hyperVGeneration": "V1",
						"creationData": {
							"createOption": "Copy",
							"sourceResourceId": "SourceDisk"
						},
						"diskSizeGB": 42
					},
					"sku": {
						"name": "Premium_LRS"
					}
				}`},
			want:          multistep.ActionContinue,
			verifyDiskset: &Diskset{-1: resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName")},
		},
		{
			name: "from platform image",
			fields: StepCreateNewDiskset{
				OSDiskID:                 "/subscriptions/SubscriptionID/resourcegroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName",
				OSDiskStorageAccountType: string(compute.StandardLRS),
				HyperVGeneration:         string(compute.V1),
				Location:                 "westus",
				SourcePlatformImage: &client.PlatformImage{
					Publisher: "Microsoft",
					Offer:     "Windows",
					Sku:       "2016-DataCenter",
					Version:   "2016.1.4",
				},
			},
			expectedPutDiskBodies: []string{`
				{
					"location": "westus",
					"properties": {
						"osType": "Linux",
						"hyperVGeneration": "V1",
						"creationData": {
							"createOption":"FromImage",
							"imageReference": {
								"id":"/subscriptions/SubscriptionID/providers/Microsoft.Compute/locations/westus/publishers/Microsoft/artifacttypes/vmimage/offers/Windows/skus/2016-DataCenter/versions/2016.1.4"
							}
						}
					},
					"sku": {
						"name": "Standard_LRS"
					}
				}`},
			want:          multistep.ActionContinue,
			verifyDiskset: &Diskset{-1: resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName")},
		},
		{
			name: "from shared image",
			fields: StepCreateNewDiskset{
				OSDiskID:                   "/subscriptions/SubscriptionID/resourcegroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName",
				OSDiskStorageAccountType:   string(compute.StandardLRS),
				DataDiskStorageAccountType: string(compute.PremiumLRS),
				DataDiskIDPrefix:           "/subscriptions/SubscriptionID/resourcegroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryDataDisk-",
				HyperVGeneration:           string(compute.V1),
				Location:                   "westus",
				SourceImageResourceID:      "/subscriptions/SubscriptionID/resourcegroups/imagegroup/providers/Microsoft.Compute/galleries/MyGallery/images/MyImage/versions/1.2.3",
			},

			expectedPutDiskBodies: []string{`
				{
					"location": "westus",
					"properties": {
						"osType": "Linux",
						"hyperVGeneration": "V1",
						"creationData": {
							"createOption":"FromImage",
							"galleryImageReference": {
								"id":"/subscriptions/SubscriptionID/resourcegroups/imagegroup/providers/Microsoft.Compute/galleries/MyGallery/images/MyImage/versions/1.2.3"
							}
						}
					},
					"sku": {
						"name": "Standard_LRS"
					}
				}`, `
				{
					"location": "westus",
					"properties": {
						"creationData": {
							"createOption":"FromImage",
							"galleryImageReference": {
								"id": "/subscriptions/SubscriptionID/resourcegroups/imagegroup/providers/Microsoft.Compute/galleries/MyGallery/images/MyImage/versions/1.2.3",
								"lun": 5
							}
						}
					},
					"sku": {
						"name": "Premium_LRS"
					}
				}`, `
				{
					"location": "westus",
					"properties": {
						"creationData": {
							"createOption":"FromImage",
							"galleryImageReference": {
								"id": "/subscriptions/SubscriptionID/resourcegroups/imagegroup/providers/Microsoft.Compute/galleries/MyGallery/images/MyImage/versions/1.2.3",
								"lun": 9
							}
						}
					},
					"sku": {
						"name": "Premium_LRS"
					}
				}`, `
				{
					"location": "westus",
					"properties": {
						"creationData": {
							"createOption":"FromImage",
							"galleryImageReference": {
								"id": "/subscriptions/SubscriptionID/resourcegroups/imagegroup/providers/Microsoft.Compute/galleries/MyGallery/images/MyImage/versions/1.2.3",
								"lun": 3
							}
						}
					},
					"sku": {
						"name": "Premium_LRS"
					}
				}`},
			want: multistep.ActionContinue,
			verifyDiskset: &Diskset{
				-1: resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName"),
				3:  resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryDataDisk-3"),
				5:  resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryDataDisk-5"),
				9:  resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryDataDisk-9"),
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := tt.fields

			bodyCount := 0
			m := compute.NewDisksClient("SubscriptionID")
			m.Sender = autorest.SenderFunc(func(r *http.Request) (*http.Response, error) {
				if r.Method != "PUT" {
					t.Fatal("Expected only a PUT disk call")
				}
				b, _ := ioutil.ReadAll(r.Body)
				expectedPutDiskBody := regexp.MustCompile(`[\s\n]`).ReplaceAllString(tt.expectedPutDiskBodies[bodyCount], "")
				bodyCount++
				if string(b) != expectedPutDiskBody {
					t.Fatalf("expected body #%d to be %q, but got %q", bodyCount, expectedPutDiskBody, string(b))
				}
				return &http.Response{
					Request:    r,
					StatusCode: 200,
				}, nil
			})

			giv := compute.NewGalleryImageVersionsClient("SubscriptionID")
			giv.Sender = autorest.SenderFunc(func(r *http.Request) (*http.Response, error) {
				if r.Method == "GET" &&
					regexp.MustCompile(`(?i)/versions/1\.2\.3$`).MatchString(r.URL.Path) {
					return &http.Response{
						Request: r,
						Body: ioutil.NopCloser(strings.NewReader(`{
							"properties": { "storageProfile": {
								"dataDiskImages":[
									{ "lun": 5 },
									{ "lun": 9 },
									{ "lun": 3 }
								]
							} }
						}`)),
						StatusCode: 200,
					}, nil
				}
				return &http.Response{
					Request:    r,
					Status:     "Unexpected request",
					StatusCode: 500,
				}, nil
			})

			state := new(multistep.BasicStateBag)
			state.Put("azureclient", &client.AzureClientSetMock{
				SubscriptionIDMock:             "SubscriptionID",
				DisksClientMock:                m,
				GalleryImageVersionsClientMock: giv,
			})
			state.Put("ui", packer.TestUi(t))

			if got := s.Run(context.TODO(), state); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("StepCreateNewDisk.Run() = %v, want %v", got, tt.want)
			}

			ds := state.Get(stateBagKey_Diskset)
			if tt.verifyDiskset != nil && !reflect.DeepEqual(*tt.verifyDiskset, ds) {
				t.Errorf("Error verifying diskset after Run(), got %v, want %v", ds, *&tt.verifyDiskset)
			}
		})
	}
}

func resource(id string) client.Resource {
	v, err := client.ParseResourceID(id)
	if err != nil {
		panic(err)
	}
	return v
}
<commit_msg>[bug] Fix test error message<commit_after>package chroot

import (
	"context"
	"io/ioutil"
	"net/http"
	"reflect"
	"regexp"
	"strings"
	"testing"

	"github.com/hashicorp/packer/builder/azure/common/client"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"

	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
	"github.com/Azure/go-autorest/autorest"
)

func TestStepCreateNewDisk_Run(t *testing.T) {
	tests := []struct {
		name                  string
		fields                StepCreateNewDiskset
		expectedPutDiskBodies []string
		want                  multistep.StepAction
		verifyDiskset         *Diskset
	}{
		{
			name: "from disk",
			fields: StepCreateNewDiskset{
				OSDiskID:                 "/subscriptions/SubscriptionID/resourcegroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName",
				OSDiskSizeGB:             42,
				OSDiskStorageAccountType: string(compute.PremiumLRS),
				HyperVGeneration:         string(compute.V1),
				Location:                 "westus",
				SourceOSDiskResourceID:   "SourceDisk",
			},
			expectedPutDiskBodies: []string{`
				{
					"location": "westus",
					"properties": {
						"osType": "Linux",
						"hyperVGeneration": "V1",
						"creationData": {
							"createOption": "Copy",
							"sourceResourceId": "SourceDisk"
						},
						"diskSizeGB": 42
					},
					"sku": {
						"name": "Premium_LRS"
					}
				}`},
			want:          multistep.ActionContinue,
			verifyDiskset: &Diskset{-1: resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName")},
		},
		{
			name: "from platform image",
			fields: StepCreateNewDiskset{
				OSDiskID:                 "/subscriptions/SubscriptionID/resourcegroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName",
				OSDiskStorageAccountType: string(compute.StandardLRS),
				HyperVGeneration:         string(compute.V1),
				Location:                 "westus",
				SourcePlatformImage: &client.PlatformImage{
					Publisher: "Microsoft",
					Offer:     "Windows",
					Sku:       "2016-DataCenter",
					Version:   "2016.1.4",
				},
			},
			expectedPutDiskBodies: []string{`
				{
					"location": "westus",
					"properties": {
						"osType": "Linux",
						"hyperVGeneration": "V1",
						"creationData": {
							"createOption":"FromImage",
							"imageReference": {
								"id":"/subscriptions/SubscriptionID/providers/Microsoft.Compute/locations/westus/publishers/Microsoft/artifacttypes/vmimage/offers/Windows/skus/2016-DataCenter/versions/2016.1.4"
							}
						}
					},
					"sku": {
						"name": "Standard_LRS"
					}
				}`},
			want:          multistep.ActionContinue,
			verifyDiskset: &Diskset{-1: resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName")},
		},
		{
			name: "from shared image",
			fields: StepCreateNewDiskset{
				OSDiskID:                   "/subscriptions/SubscriptionID/resourcegroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName",
				OSDiskStorageAccountType:   string(compute.StandardLRS),
				DataDiskStorageAccountType: string(compute.PremiumLRS),
				DataDiskIDPrefix:           "/subscriptions/SubscriptionID/resourcegroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryDataDisk-",
				HyperVGeneration:           string(compute.V1),
				Location:                   "westus",
				SourceImageResourceID:      "/subscriptions/SubscriptionID/resourcegroups/imagegroup/providers/Microsoft.Compute/galleries/MyGallery/images/MyImage/versions/1.2.3",
			},

			expectedPutDiskBodies: []string{`
				{
					"location": "westus",
					"properties": {
						"osType": "Linux",
						"hyperVGeneration": "V1",
						"creationData": {
							"createOption":"FromImage",
							"galleryImageReference": {
								"id":"/subscriptions/SubscriptionID/resourcegroups/imagegroup/providers/Microsoft.Compute/galleries/MyGallery/images/MyImage/versions/1.2.3"
							}
						}
					},
					"sku": {
						"name": "Standard_LRS"
					}
				}`, `
				{
					"location": "westus",
					"properties": {
						"creationData": {
							"createOption":"FromImage",
							"galleryImageReference": {
								"id": "/subscriptions/SubscriptionID/resourcegroups/imagegroup/providers/Microsoft.Compute/galleries/MyGallery/images/MyImage/versions/1.2.3",
								"lun": 5
							}
						}
					},
					"sku": {
						"name": "Premium_LRS"
					}
				}`, `
				{
					"location": "westus",
					"properties": {
						"creationData": {
							"createOption":"FromImage",
							"galleryImageReference": {
								"id": "/subscriptions/SubscriptionID/resourcegroups/imagegroup/providers/Microsoft.Compute/galleries/MyGallery/images/MyImage/versions/1.2.3",
								"lun": 9
							}
						}
					},
					"sku": {
						"name": "Premium_LRS"
					}
				}`, `
				{
					"location": "westus",
					"properties": {
						"creationData": {
							"createOption":"FromImage",
							"galleryImageReference": {
								"id": "/subscriptions/SubscriptionID/resourcegroups/imagegroup/providers/Microsoft.Compute/galleries/MyGallery/images/MyImage/versions/1.2.3",
								"lun": 3
							}
						}
					},
					"sku": {
						"name": "Premium_LRS"
					}
				}`},
			want: multistep.ActionContinue,
			verifyDiskset: &Diskset{
				-1: resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryOSDiskName"),
				3:  resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryDataDisk-3"),
				5:  resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryDataDisk-5"),
				9:  resource("/subscriptions/SubscriptionID/resourceGroups/ResourceGroupName/providers/Microsoft.Compute/disks/TemporaryDataDisk-9"),
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := tt.fields

			bodyCount := 0
			m := compute.NewDisksClient("SubscriptionID")
			m.Sender = autorest.SenderFunc(func(r *http.Request) (*http.Response, error) {
				if r.Method != "PUT" {
					t.Fatal("Expected only a PUT disk call")
				}
				b, _ := ioutil.ReadAll(r.Body)
				expectedPutDiskBody := regexp.MustCompile(`[\s\n]`).ReplaceAllString(tt.expectedPutDiskBodies[bodyCount], "")
				bodyCount++
				if string(b) != expectedPutDiskBody {
					t.Fatalf("expected body #%d to be %q, but got %q", bodyCount, expectedPutDiskBody, string(b))
				}
				return &http.Response{
					Request:    r,
					StatusCode: 200,
				}, nil
			})

			giv := compute.NewGalleryImageVersionsClient("SubscriptionID")
			giv.Sender = autorest.SenderFunc(func(r *http.Request) (*http.Response, error) {
				if r.Method == "GET" &&
					regexp.MustCompile(`(?i)/versions/1\.2\.3$`).MatchString(r.URL.Path) {
					return &http.Response{
						Request: r,
						Body: ioutil.NopCloser(strings.NewReader(`{
							"properties": { "storageProfile": {
								"dataDiskImages":[
									{ "lun": 5 },
									{ "lun": 9 },
									{ "lun": 3 }
								]
							} }
						}`)),
						StatusCode: 200,
					}, nil
				}
				return &http.Response{
					Request:    r,
					Status:     "Unexpected request",
					StatusCode: 500,
				}, nil
			})

			state := new(multistep.BasicStateBag)
			state.Put("azureclient", &client.AzureClientSetMock{
				SubscriptionIDMock:             "SubscriptionID",
				DisksClientMock:                m,
				GalleryImageVersionsClientMock: giv,
			})
			state.Put("ui", packer.TestUi(t))

			if got := s.Run(context.TODO(), state); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("StepCreateNewDisk.Run() = %v, want %v", got, tt.want)
			}

			ds := state.Get(stateBagKey_Diskset)
			if tt.verifyDiskset != nil && !reflect.DeepEqual(*tt.verifyDiskset, ds) {
				t.Errorf("Error verifying diskset after Run(), got %v, want %v", ds, *tt.verifyDiskset)
			}
		})
	}
}

func resource(id string) client.Resource {
	v, err := client.ParseResourceID(id)
	if err != nil {
		panic(err)
	}
	return v
}
<|endoftext|>


<commit_before>package router

import (
	"bufio"
	"errors"
	"io"
	"log"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/fsouza/go-dockerclient"
)

func init() {
	pump := &LogsPump{
		pumps:  make(map[string]*containerPump),
		routes: make(map[chan *update]struct{}),
	}
	LogRouters.Register(pump, "pump")
	Jobs.Register(pump, "pump")
}

func getopt(name, dfault string) string {
	value := os.Getenv(name)
	if value == "" {
		value = dfault
	}
	return value
}

func debug(v ...interface{}) {
	if os.Getenv("DEBUG") != "" {
		log.Println(v...)
	}
}

func assert(err error, context string) {
	if err != nil {
		log.Fatal(context+": ", err)
	}
}

func normalName(name string) string {
	return name[1:]
}

func normalID(id string) string {
	if len(id) > 12 {
		return id[:12]
	}
	return id
}

func ignoreContainer(container *docker.Container) bool {
	for _, kv := range container.Config.Env {
		kvp := strings.SplitN(kv, "=", 2)
		if len(kvp) == 2 && kvp[0] == "LOGSPOUT" && strings.ToLower(kvp[1]) == "ignore" {
			return true
		}
	}
	return false
}

type update struct {
	*docker.APIEvents
	pump *containerPump
}

type LogsPump struct {
	mu     sync.Mutex
	pumps  map[string]*containerPump
	routes map[chan *update]struct{}
	client *docker.Client
}

func (p *LogsPump) Name() string {
	return "pump"
}

func (p *LogsPump) Setup() error {
	client, err := docker.NewClient(
		getopt("DOCKER_HOST", "unix:///var/run/docker.sock"))
	if err != nil {
		return err
	}
	p.client = client
	return nil
}

func (p *LogsPump) rename(event *docker.APIEvents) {
    p.mu.Lock()
    defer p.mu.Unlock()
    container, err := p.client.InspectContainer(event.ID)
    assert(err, "pump")
    pump, _ := p.pumps[normalID(event.ID)]
    pump.container.Name = container.Name
}

func (p *LogsPump) Run() error {
	containers, err := p.client.ListContainers(docker.ListContainersOptions{})
	if err != nil {
		return err
	}
	for _, listing := range containers {
		p.pumpLogs(&docker.APIEvents{
			ID:     normalID(listing.ID),
			Status: "start",
		}, false)
	}
	events := make(chan *docker.APIEvents)
	err = p.client.AddEventListener(events)
	if err != nil {
		return err
	}
	for event := range events {
		debug("pump.Run() event:", normalID(event.ID), event.Status)
		switch event.Status {
		case "start", "restart":
			go p.pumpLogs(event, true)
		case "rename":
			go p.rename(event)
		case "die":
			go p.update(event)
		}
	}
	return errors.New("docker event stream closed")
}

func (p *LogsPump) pumpLogs(event *docker.APIEvents, backlog bool) {
	id := normalID(event.ID)
	container, err := p.client.InspectContainer(id)
	assert(err, "pump")
	if container.Config.Tty {
		debug("pump.pumpLogs():", id, "ignored: tty enabled")
		return
	}
	if ignoreContainer(container) {
		debug("pump.pumpLogs():", id, "ignored: environ ignore")
		return
	}
	var tail string
	if backlog {
		tail = "all"
	} else {
		tail = "0"
	}
	outrd, outwr := io.Pipe()
	errrd, errwr := io.Pipe()
	p.mu.Lock()
	p.pumps[id] = newContainerPump(container, outrd, errrd)
	p.mu.Unlock()
	p.update(event)
	debug("pump.pumpLogs():", id, "started")
	go func() {
		err := p.client.Logs(docker.LogsOptions{
			Container:    id,
			OutputStream: outwr,
			ErrorStream:  errwr,
			Stdout:       true,
			Stderr:       true,
			Follow:       true,
			Tail:         tail,
		})
		if err != nil {
			debug("pump.pumpLogs():", id, "stopped:", err)
		}
		outwr.Close()
		errwr.Close()
		p.mu.Lock()
		delete(p.pumps, id)
		p.mu.Unlock()
	}()
}

func (p *LogsPump) update(event *docker.APIEvents) {
	p.mu.Lock()
	defer p.mu.Unlock()
	pump, pumping := p.pumps[normalID(event.ID)]
	if pumping {
		for r, _ := range p.routes {
			select {
			case r <- &update{event, pump}:
			case <-time.After(time.Second * 1):
				debug("pump.update(): route timeout, dropping")
				defer delete(p.routes, r)
			}
		}
	}
}

func (p *LogsPump) RoutingFrom(id string) bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	_, monitoring := p.pumps[normalID(id)]
	return monitoring
}

func (p *LogsPump) Route(route *Route, logstream chan *Message) {
	p.mu.Lock()
	for _, pump := range p.pumps {
		if route.MatchContainer(
			normalID(pump.container.ID),
			normalName(pump.container.Name)) {

			pump.add(logstream, route)
			defer pump.remove(logstream)
		}
	}
	updates := make(chan *update)
	p.routes[updates] = struct{}{}
	p.mu.Unlock()
	defer func() {
		p.mu.Lock()
		delete(p.routes, updates)
		p.mu.Unlock()
	}()
	for {
		select {
		case event := <-updates:
			switch event.Status {
			case "start", "restart":
				if route.MatchContainer(
					normalID(event.pump.container.ID),
					normalName(event.pump.container.Name)) {

					event.pump.add(logstream, route)
					defer event.pump.remove(logstream)
				}
			case "die":
				if strings.HasPrefix(route.FilterID, event.ID) {
					// If the route is just about a single container,
					// we can stop routing when it dies.
					return
				}
			}
		case <-route.Closer():
			return
		}
	}
}

type containerPump struct {
	sync.Mutex
	container  *docker.Container
	logstreams map[chan *Message]*Route
}

func newContainerPump(container *docker.Container, stdout, stderr io.Reader) *containerPump {
	cp := &containerPump{
		container:  container,
		logstreams: make(map[chan *Message]*Route),
	}
	pump := func(source string, input io.Reader) {
		buf := bufio.NewReader(input)
		for {
			line, err := buf.ReadString('\n')
			if err != nil {
				if err != io.EOF {
					debug("pump.newContainerPump():", normalID(container.ID), source+":", err)
				}
				return
			}
			cp.send(&Message{
				Data:      strings.TrimSuffix(line, "\n"),
				Container: container,
				Time:      time.Now(),
				Source:    source,
			})
		}
	}
	go pump("stdout", stdout)
	go pump("stderr", stderr)
	return cp
}

func (cp *containerPump) send(msg *Message) {
	cp.Lock()
	defer cp.Unlock()
	for logstream, route := range cp.logstreams {
		if !route.MatchMessage(msg) {
			continue
		}
		select {
		case logstream <- msg:
		case <-time.After(time.Second * 1):
			debug("pump.send(): send timeout, closing")
			// normal call to remove() triggered by
			// route.Closer() may not be able to grab
			// lock under heavy load, so we delete here
			defer delete(cp.logstreams, logstream)
		}
	}
}

func (cp *containerPump) add(logstream chan *Message, route *Route) {
	cp.Lock()
	defer cp.Unlock()
	cp.logstreams[logstream] = route
}

func (cp *containerPump) remove(logstream chan *Message) {
	cp.Lock()
	defer cp.Unlock()
	delete(cp.logstreams, logstream)
}
<commit_msg>cleaned up setup method<commit_after>package router

import (
	"bufio"
	"errors"
	"io"
	"log"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/fsouza/go-dockerclient"
)

func init() {
	pump := &LogsPump{
		pumps:  make(map[string]*containerPump),
		routes: make(map[chan *update]struct{}),
	}
	LogRouters.Register(pump, "pump")
	Jobs.Register(pump, "pump")
}

func getopt(name, dfault string) string {
	value := os.Getenv(name)
	if value == "" {
		value = dfault
	}
	return value
}

func debug(v ...interface{}) {
	if os.Getenv("DEBUG") != "" {
		log.Println(v...)
	}
}

func assert(err error, context string) {
	if err != nil {
		log.Fatal(context+": ", err)
	}
}

func normalName(name string) string {
	return name[1:]
}

func normalID(id string) string {
	if len(id) > 12 {
		return id[:12]
	}
	return id
}

func ignoreContainer(container *docker.Container) bool {
	for _, kv := range container.Config.Env {
		kvp := strings.SplitN(kv, "=", 2)
		if len(kvp) == 2 && kvp[0] == "LOGSPOUT" && strings.ToLower(kvp[1]) == "ignore" {
			return true
		}
	}
	return false
}

type update struct {
	*docker.APIEvents
	pump *containerPump
}

type LogsPump struct {
	mu     sync.Mutex
	pumps  map[string]*containerPump
	routes map[chan *update]struct{}
	client *docker.Client
}

func (p *LogsPump) Name() string {
	return "pump"
}

func (p *LogsPump) Setup() error {
	var err error
	p.client, err = docker.NewClientFromEnv()
	return err
}

func (p *LogsPump) rename(event *docker.APIEvents) {
	p.mu.Lock()
	defer p.mu.Unlock()
	container, err := p.client.InspectContainer(event.ID)
	assert(err, "pump")
	pump, _ := p.pumps[normalID(event.ID)]
	pump.container.Name = container.Name
}

func (p *LogsPump) Run() error {
	containers, err := p.client.ListContainers(docker.ListContainersOptions{})
	if err != nil {
		return err
	}
	for _, listing := range containers {
		p.pumpLogs(&docker.APIEvents{
			ID:     normalID(listing.ID),
			Status: "start",
		}, false)
	}
	events := make(chan *docker.APIEvents)
	err = p.client.AddEventListener(events)
	if err != nil {
		return err
	}
	for event := range events {
		debug("pump.Run() event:", normalID(event.ID), event.Status)
		switch event.Status {
		case "start", "restart":
			go p.pumpLogs(event, true)
		case "rename":
			go p.rename(event)
		case "die":
			go p.update(event)
		}
	}
	return errors.New("docker event stream closed")
}

func (p *LogsPump) pumpLogs(event *docker.APIEvents, backlog bool) {
	id := normalID(event.ID)
	container, err := p.client.InspectContainer(id)
	assert(err, "pump")
	if container.Config.Tty {
		debug("pump.pumpLogs():", id, "ignored: tty enabled")
		return
	}
	if ignoreContainer(container) {
		debug("pump.pumpLogs():", id, "ignored: environ ignore")
		return
	}
	var tail string
	if backlog {
		tail = "all"
	} else {
		tail = "0"
	}
	outrd, outwr := io.Pipe()
	errrd, errwr := io.Pipe()
	p.mu.Lock()
	p.pumps[id] = newContainerPump(container, outrd, errrd)
	p.mu.Unlock()
	p.update(event)
	debug("pump.pumpLogs():", id, "started")
	go func() {
		err := p.client.Logs(docker.LogsOptions{
			Container:    id,
			OutputStream: outwr,
			ErrorStream:  errwr,
			Stdout:       true,
			Stderr:       true,
			Follow:       true,
			Tail:         tail,
		})
		if err != nil {
			debug("pump.pumpLogs():", id, "stopped:", err)
		}
		outwr.Close()
		errwr.Close()
		p.mu.Lock()
		delete(p.pumps, id)
		p.mu.Unlock()
	}()
}

func (p *LogsPump) update(event *docker.APIEvents) {
	p.mu.Lock()
	defer p.mu.Unlock()
	pump, pumping := p.pumps[normalID(event.ID)]
	if pumping {
		for r := range p.routes {
			select {
			case r <- &update{event, pump}:
			case <-time.After(time.Second * 1):
				debug("pump.update(): route timeout, dropping")
				defer delete(p.routes, r)
			}
		}
	}
}

func (p *LogsPump) RoutingFrom(id string) bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	_, monitoring := p.pumps[normalID(id)]
	return monitoring
}

func (p *LogsPump) Route(route *Route, logstream chan *Message) {
	p.mu.Lock()
	for _, pump := range p.pumps {
		if route.MatchContainer(
			normalID(pump.container.ID),
			normalName(pump.container.Name)) {

			pump.add(logstream, route)
			defer pump.remove(logstream)
		}
	}
	updates := make(chan *update)
	p.routes[updates] = struct{}{}
	p.mu.Unlock()
	defer func() {
		p.mu.Lock()
		delete(p.routes, updates)
		p.mu.Unlock()
	}()
	for {
		select {
		case event := <-updates:
			switch event.Status {
			case "start", "restart":
				if route.MatchContainer(
					normalID(event.pump.container.ID),
					normalName(event.pump.container.Name)) {

					event.pump.add(logstream, route)
					defer event.pump.remove(logstream)
				}
			case "die":
				if strings.HasPrefix(route.FilterID, event.ID) {
					// If the route is just about a single container,
					// we can stop routing when it dies.
					return
				}
			}
		case <-route.Closer():
			return
		}
	}
}

type containerPump struct {
	sync.Mutex
	container  *docker.Container
	logstreams map[chan *Message]*Route
}

func newContainerPump(container *docker.Container, stdout, stderr io.Reader) *containerPump {
	cp := &containerPump{
		container:  container,
		logstreams: make(map[chan *Message]*Route),
	}
	pump := func(source string, input io.Reader) {
		buf := bufio.NewReader(input)
		for {
			line, err := buf.ReadString('\n')
			if err != nil {
				if err != io.EOF {
					debug("pump.newContainerPump():", normalID(container.ID), source+":", err)
				}
				return
			}
			cp.send(&Message{
				Data:      strings.TrimSuffix(line, "\n"),
				Container: container,
				Time:      time.Now(),
				Source:    source,
			})
		}
	}
	go pump("stdout", stdout)
	go pump("stderr", stderr)
	return cp
}

func (cp *containerPump) send(msg *Message) {
	cp.Lock()
	defer cp.Unlock()
	for logstream, route := range cp.logstreams {
		if !route.MatchMessage(msg) {
			continue
		}
		select {
		case logstream <- msg:
		case <-time.After(time.Second * 1):
			debug("pump.send(): send timeout, closing")
			// normal call to remove() triggered by
			// route.Closer() may not be able to grab
			// lock under heavy load, so we delete here
			defer delete(cp.logstreams, logstream)
		}
	}
}

func (cp *containerPump) add(logstream chan *Message, route *Route) {
	cp.Lock()
	defer cp.Unlock()
	cp.logstreams[logstream] = route
}

func (cp *containerPump) remove(logstream chan *Message) {
	cp.Lock()
	defer cp.Unlock()
	delete(cp.logstreams, logstream)
}
<|endoftext|>


<commit_before>package grunway

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

func TestNothing(t *testing.T) {
	// just a placeholder
}

type AuthorController struct {
}
type BookController struct {
}

type AuthorPayload struct {
	Entity
	Name string
}
type BookPayload struct {
	Entity
	Name     string
	AuthorId uint64
}

func (self *AuthorController) GetHandlerV1(ctx *Context) {

}

func (self *BookController) GetHandlerV1(ctx *Context) {

	var book BookPayload
	book.SetPrimaryKey(1)
	book.Name = "The Greatest Works of All Time"
	book.AuthorId = 1

	ctx.WrapAndSendPayload(book)
}
func (self *BookController) GetHandlerV2(ctx *Context) {

}
func (self *BookController) GetHandlerV003(ctx *Context) {

}
func (self *BookController) GetHandlerV018(ctx *Context) {

}
func (self *BookController) PostHandlerV1(ctx *Context) {

}
func (self *BookController) PutHandlerV1(ctx *Context) {

}
func (self *BookController) DeleteHandlerV1(ctx *Context) {

}
func (self *BookController) GetHandlerV1All(ctx *Context) {

}
func (self *BookController) GetHandlerV1Popular(ctx *Context) {

}

func makeLibrary(t *testing.T) *Router {
	routerPtr := NewRouter()
	routerPtr.BasePath = "/api/"
	routerPtr.RegisterEntity("author", &AuthorController{}, &AuthorPayload{})
	routerPtr.RegisterEntity("book", &BookController{}, &BookPayload{})

	// routerPtr.LogRoutes()
	t.Log("All Routes:\n", routerPtr.AllRoutesSummary())

	return routerPtr
}

func TestRouterSetup(t *testing.T) {
	makeLibrary(t)
}

func TestAPIRoutes(t *testing.T) {
	router := makeLibrary(t)
	ts := httptest.NewServer(router)
	defer ts.Close()

	getURLAndStatusCodes := map[string]int{
		// CRUD
		"/api/v1/book/1":   http.StatusOK,
		"/api/v2/book/1":   http.StatusOK,
		"/api/v3/book/1":   http.StatusOK,
		"/api/v4/book/1":   http.StatusNotFound,
		"/api/v18/book/1":  http.StatusOK,
		"/api/v1/book/":    http.StatusBadRequest,
		"/api/v1/book":     http.StatusBadRequest,
		"/api/v1/author/1": http.StatusOK,
		"/api/v1/bogus/1":  http.StatusNotFound,

		// Custom
		"/api/v1/book/Popular": http.StatusOK,
		"/api/v1/book/All":     http.StatusOK,
		"/api/v1/book/Bogus":   http.StatusNotFound,
	}

	for urlsuffix, expectedStatusCode := range getURLAndStatusCodes {
		response, err := http.Get(ts.URL + urlsuffix)
		if err != nil {
			t.Fatal(err)
		}

		if response.StatusCode != expectedStatusCode {
			t.Error("GET", urlsuffix, "expected ", expectedStatusCode, ", got", response.StatusCode)
		}
	}

	postURLAndStatusCodes := map[string]int{
		"/api/v1/book/": http.StatusOK,

		"/api/v2/book/":   http.StatusNotFound,
		"/api/v3/book/":   http.StatusNotFound,
		"/api/v1/author/": http.StatusNotFound,
		"/api/v1/bogus/":  http.StatusNotFound,

		"/api/v1/book/1": http.StatusBadRequest, // Create (POST) should never have a pk
	}

	for urlsuffix, expectedStatusCode := range postURLAndStatusCodes {
		buf := bytes.NewBufferString("Hello World!")
		response, err := http.Post(ts.URL+urlsuffix, "text/plain", buf)
		if err != nil {
			t.Fatal(err)
		}

		if response.StatusCode != expectedStatusCode {
			t.Error("POST", urlsuffix, "expected ", expectedStatusCode, ", got", response.StatusCode)
		}
	}

	putURLAndStatusCodes := map[string]int{
		"/api/v1/book/1": http.StatusOK,

		"/api/v2/book/1":   http.StatusNotFound,
		"/api/v3/book/1":   http.StatusNotFound,
		"/api/v1/author/1": http.StatusNotFound,
		"/api/v1/bogus/1":  http.StatusNotFound,

		"/api/v1/book/": http.StatusBadRequest, // Update (PUT) should always have a pk
	}

	for urlsuffix, expectedStatusCode := range putURLAndStatusCodes {
		buf := bytes.NewBufferString("Hello World!")

		client := new(http.Client)
		req, err := http.NewRequest("PUT", ts.URL+urlsuffix, buf)
		response, err := client.Do(req)
		if err != nil {
			t.Fatal(err)
		}

		if response.StatusCode != expectedStatusCode {
			t.Error("PUT", urlsuffix, "expected ", expectedStatusCode, ", got", response.StatusCode)
		}
	}
}

func TestPayload(t *testing.T) {
	router := makeLibrary(t)
	ts := httptest.NewServer(router)
	defer ts.Close()

	response, err := http.Get(ts.URL + "/api/v1/book/1")

	bodyBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		t.Error("9244741117", err)
	}
	defer response.Body.Close()
	bodyString := string(bodyBytes)
	t.Log(bodyString)

	// decode the json,
	var pw PayloadWrapper
	json.Unmarshal(bodyBytes, &pw)

	//check errNo == 0
	if pw.ErrNo != 0 {
		t.Errorf("918188683 expected non-zero ErrStr, got %d\npayload:%+v", pw.ErrNo, pw)
	}
	//check pk matches expected,

	if len(pw.PayloadList) != 1 {
		t.Errorf("918188684 expected 1 entity, got %v\npayload:%+v", pw.ErrStr, pw)
	}

	for i, untypedEntity := range pw.PayloadList {

		jsonBytes, err := json.Marshal(untypedEntity)
		if err != nil {
			t.Fatalf("918188685 failure to marshal entity %+v", untypedEntity)
		}

		var bookPayload BookPayload
		err = json.Unmarshal(jsonBytes, &bookPayload)
		if err != nil {
			t.Fatalf("918188686 failure to Unmarshal entity %+v", string(jsonBytes))
		}

		if bookPayload.GetPrimaryKey() != 1 {
			t.Errorf("918188687 i:%d Expected pk == 1, got %d\nentity:%+v", i, bookPayload.GetPrimaryKey(), bookPayload)
		}
	}
}

// Benchmark our routeKey Algorithms.  this is called every request.

//As of 2013-09-19, Go 1.1, rMBP
//BenchmarkRouteKeyJoinString	 5000000	       483 ns/op
func BenchmarkRouteKeyJoinString(b *testing.B) {
	for i := 0; i < b.N; i++ {
		routeKeyJoinString("GET", "1", "book", "all")
	}
}

//As of 2013-09-19, Go 1.1, rMBP
//BenchmarkRouteKeyFormatString	 1000000	      1249 ns/op
func BenchmarkRouteKeyFormatString(b *testing.B) {
	for i := 0; i < b.N; i++ {
		routeKeyFormatString("GET", "1", "book", "all")
	}
}
<commit_msg>fix unit tests<commit_after>package grunway

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

func TestNothing(t *testing.T) {
	// just a placeholder
}

type AuthorController struct {
}
type BookController struct {
}

type AuthorPayload struct {
	PKey int64
	Name string
}
type BookPayload struct {
	PKey     int64
	Name     string
	AuthorId uint64
}

func (self *AuthorController) GetHandlerV1(ctx *Context) {

}

func (self *BookController) GetHandlerV1(ctx *Context) {

	var book BookPayload
	book.PKey = 1
	book.Name = "The Greatest Works of All Time"
	book.AuthorId = 1

	ctx.WrapAndSendPayload(book)
}
func (self *BookController) GetHandlerV2(ctx *Context) {

}
func (self *BookController) GetHandlerV003(ctx *Context) {

}
func (self *BookController) GetHandlerV018(ctx *Context) {

}
func (self *BookController) PostHandlerV1(ctx *Context) {

}
func (self *BookController) PutHandlerV1(ctx *Context) {

}
func (self *BookController) DeleteHandlerV1(ctx *Context) {

}
func (self *BookController) GetHandlerV1All(ctx *Context) {

}
func (self *BookController) GetHandlerV1Popular(ctx *Context) {

}

func makeLibrary(t *testing.T) *Router {
	routerPtr := NewRouter()
	routerPtr.BasePath = "/api/"
	routerPtr.RegisterEntity("author", &AuthorController{}, &AuthorPayload{})
	routerPtr.RegisterEntity("book", &BookController{}, &BookPayload{})

	// routerPtr.LogRoutes()
	t.Log("All Routes:\n", routerPtr.AllRoutesSummary())

	return routerPtr
}

func TestRouterSetup(t *testing.T) {
	makeLibrary(t)
}

func TestAPIRoutes(t *testing.T) {
	router := makeLibrary(t)
	ts := httptest.NewServer(router)
	defer ts.Close()

	getURLAndStatusCodes := map[string]int{
		// CRUD
		"/api/v1/book/1":   http.StatusOK,
		"/api/v2/book/1":   http.StatusOK,
		"/api/v3/book/1":   http.StatusOK,
		"/api/v4/book/1":   http.StatusNotFound,
		"/api/v18/book/1":  http.StatusOK,
		"/api/v1/book/":    http.StatusBadRequest,
		"/api/v1/book":     http.StatusBadRequest,
		"/api/v1/author/1": http.StatusOK,
		"/api/v1/bogus/1":  http.StatusNotFound,

		// Custom
		"/api/v1/book/Popular": http.StatusOK,
		"/api/v1/book/All":     http.StatusOK,
		"/api/v1/book/Bogus":   http.StatusNotFound,
	}

	for urlsuffix, expectedStatusCode := range getURLAndStatusCodes {
		response, err := http.Get(ts.URL + urlsuffix)
		if err != nil {
			t.Fatal(err)
		}

		if response.StatusCode != expectedStatusCode {
			t.Error("GET", urlsuffix, "expected ", expectedStatusCode, ", got", response.StatusCode)
		}
	}

	postURLAndStatusCodes := map[string]int{
		"/api/v1/book/": http.StatusOK,

		"/api/v2/book/":   http.StatusNotFound,
		"/api/v3/book/":   http.StatusNotFound,
		"/api/v1/author/": http.StatusNotFound,
		"/api/v1/bogus/":  http.StatusNotFound,

		"/api/v1/book/1": http.StatusBadRequest, // Create (POST) should never have a pk
	}

	for urlsuffix, expectedStatusCode := range postURLAndStatusCodes {
		buf := bytes.NewBufferString("Hello World!")
		response, err := http.Post(ts.URL+urlsuffix, "text/plain", buf)
		if err != nil {
			t.Fatal(err)
		}

		if response.StatusCode != expectedStatusCode {
			t.Error("POST", urlsuffix, "expected ", expectedStatusCode, ", got", response.StatusCode)
		}
	}

	putURLAndStatusCodes := map[string]int{
		"/api/v1/book/1": http.StatusOK,

		"/api/v2/book/1":   http.StatusNotFound,
		"/api/v3/book/1":   http.StatusNotFound,
		"/api/v1/author/1": http.StatusNotFound,
		"/api/v1/bogus/1":  http.StatusNotFound,

		"/api/v1/book/": http.StatusBadRequest, // Update (PUT) should always have a pk
	}

	for urlsuffix, expectedStatusCode := range putURLAndStatusCodes {
		buf := bytes.NewBufferString("Hello World!")

		client := new(http.Client)
		req, err := http.NewRequest("PUT", ts.URL+urlsuffix, buf)
		response, err := client.Do(req)
		if err != nil {
			t.Fatal(err)
		}

		if response.StatusCode != expectedStatusCode {
			t.Error("PUT", urlsuffix, "expected ", expectedStatusCode, ", got", response.StatusCode)
		}
	}
}

func TestPayload(t *testing.T) {
	router := makeLibrary(t)
	ts := httptest.NewServer(router)
	defer ts.Close()

	response, err := http.Get(ts.URL + "/api/v1/book/1")

	bodyBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		t.Error("9244741117", err)
	}
	defer response.Body.Close()
	bodyString := string(bodyBytes)
	t.Log(bodyString)

	// decode the json,
	var pw PayloadWrapper
	json.Unmarshal(bodyBytes, &pw)

	//check errNo == 0
	if pw.ErrNo != 0 {
		t.Errorf("918188683 expected non-zero ErrStr, got %d\npayload:%+v", pw.ErrNo, pw)
	}
	//check pk matches expected,

	if len(pw.PayloadList) != 1 {
		t.Errorf("918188684 expected 1 entity, got %v\npayload:%+v", pw.ErrStr, pw)
	}

	for i, untypedEntity := range pw.PayloadList {

		jsonBytes, err := json.Marshal(untypedEntity)
		if err != nil {
			t.Fatalf("918188685 failure to marshal entity %+v", untypedEntity)
		}

		var bookPayload BookPayload
		err = json.Unmarshal(jsonBytes, &bookPayload)
		if err != nil {
			t.Fatalf("918188686 failure to Unmarshal entity %+v", string(jsonBytes))
		}

		if bookPayload.PKey != 1 {
			t.Errorf("918188687 i:%d Expected pk == 1, got %d\nentity:%+v", i, bookPayload.PKey, bookPayload)
		}
	}
}

// Benchmark our routeKey Algorithms.  this is called every request.

//As of 2013-09-19, Go 1.1, rMBP
//BenchmarkRouteKeyJoinString	 5000000	       483 ns/op
func BenchmarkRouteKeyJoinString(b *testing.B) {
	for i := 0; i < b.N; i++ {
		routeKeyJoinString("GET", "1", "book", "all")
	}
}

//As of 2013-09-19, Go 1.1, rMBP
//BenchmarkRouteKeyFormatString	 1000000	      1249 ns/op
func BenchmarkRouteKeyFormatString(b *testing.B) {
	for i := 0; i < b.N; i++ {
		routeKeyFormatString("GET", "1", "book", "all")
	}
}
<|endoftext|>


<commit_before>package relay

import (
	"context"
	"fmt"
	"io"
	"sync"
	"time"

	pb "github.com/libp2p/go-libp2p-circuit/pb"

	logging "github.com/ipfs/go-log"
	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
)

var log = logging.Logger("relay")

const ProtoID = "/libp2p/circuit/relay/0.1.0"

const maxMessageSize = 4096

var RelayAcceptTimeout = time.Minute
var HopConnectTimeout = 10 * time.Second

type Relay struct {
	host host.Host
	ctx  context.Context
	self peer.ID

	active bool
	hop    bool

	incoming chan *Conn

	arLk         sync.Mutex
	activeRelays []*Conn
}

type RelayOpt int

var (
	OptActive = RelayOpt(0)
	OptHop    = RelayOpt(1)
)

type RelayError struct {
	Code pb.CircuitRelay_Status
}

func (e RelayError) Error() string {
	return fmt.Sprintf("error opening relay circuit: %s (%d)", pb.CircuitRelay_Status_name[int32(e.Code)], e.Code)
}

func NewRelay(ctx context.Context, h host.Host, opts ...RelayOpt) (*Relay, error) {
	r := &Relay{
		host:     h,
		ctx:      ctx,
		self:     h.ID(),
		incoming: make(chan *Conn),
	}

	for _, opt := range opts {
		switch opt {
		case OptActive:
			r.active = true
		case OptHop:
			r.hop = true
		default:
			return nil, fmt.Errorf("unrecognized option: %d", opt)
		}
	}

	h.SetStreamHandler(ProtoID, r.handleNewStream)

	return r, nil
}

func (r *Relay) DialPeer(ctx context.Context, relay pstore.PeerInfo, dest pstore.PeerInfo) (*Conn, error) {

	if len(relay.Addrs) > 0 {
		r.host.Peerstore().AddAddrs(relay.ID, relay.Addrs, pstore.TempAddrTTL)
	}

	s, err := r.host.NewStream(ctx, relay.ID, ProtoID)
	if err != nil {
		return nil, err
	}

	rd := newDelimitedReader(s, maxMessageSize)
	wr := newDelimitedWriter(s)

	var msg pb.CircuitRelay

	msg.Type = pb.CircuitRelay_HOP.Enum()
	msg.SrcPeer = peerInfoToPeer(r.host.Peerstore().PeerInfo(r.self))
	msg.DstPeer = peerInfoToPeer(dest)

	err = wr.WriteMsg(&msg)
	if err != nil {
		s.Close()
		return nil, err
	}

	msg.Reset()

	err = rd.ReadMsg(&msg)
	if err != nil {
		s.Close()
		return nil, err
	}

	if msg.GetType() != pb.CircuitRelay_STATUS {
		s.Close()
		return nil, fmt.Errorf("unexpected relay response; not a status message (%d)", msg.GetType())
	}

	if msg.GetCode() != pb.CircuitRelay_SUCCESS {
		s.Close()
		return nil, RelayError{msg.GetCode()}
	}

	return &Conn{Stream: s, remote: dest, transport: r.Transport()}, nil
}

func (r *Relay) handleNewStream(s inet.Stream) {
	log.Infof("new relay stream from: %s", s.Conn().RemotePeer())

	rd := newDelimitedReader(s, maxMessageSize)

	var msg pb.CircuitRelay

	err := rd.ReadMsg(&msg)
	if err != nil {
		r.handleError(s, pb.CircuitRelay_MALFORMED_MESSAGE)
		return
	}

	switch msg.GetType() {
	case pb.CircuitRelay_HOP:
		r.handleHopStream(s, &msg)
	case pb.CircuitRelay_STOP:
		r.handleStopStream(s, &msg)
	case pb.CircuitRelay_CAN_HOP:
		r.handleCanHop(s, &msg)
	default:
		log.Warningf("unexpected relay handshake: %d", msg.GetType())
		r.handleError(s, pb.CircuitRelay_MALFORMED_MESSAGE)
	}
}

func (r *Relay) handleHopStream(s inet.Stream, msg *pb.CircuitRelay) {
	if !r.hop {
		r.handleError(s, pb.CircuitRelay_HOP_CANT_SPEAK_RELAY)
		return
	}

	src, err := peerToPeerInfo(msg.GetSrcPeer())
	if err != nil {
		r.handleError(s, pb.CircuitRelay_HOP_SRC_MULTIADDR_INVALID)
		return
	}

	if src.ID != s.Conn().RemotePeer() {
		r.handleError(s, pb.CircuitRelay_HOP_SRC_MULTIADDR_INVALID)
		return
	}

	dst, err := peerToPeerInfo(msg.GetDstPeer())
	if err != nil {
		r.handleError(s, pb.CircuitRelay_HOP_DST_MULTIADDR_INVALID)
		return
	}

	if dst.ID == r.self {
		r.handleError(s, pb.CircuitRelay_HOP_CANT_RELAY_TO_SELF)
		return
	}

	// open stream
	ctp := r.host.Network().ConnsToPeer(dst.ID)

	if len(ctp) == 0 && !r.active {
		r.handleError(s, pb.CircuitRelay_HOP_NO_CONN_TO_DST)
		return
	}

	if len(dst.Addrs) > 0 {
		r.host.Peerstore().AddAddrs(dst.ID, dst.Addrs, pstore.TempAddrTTL)
	}

	ctx, cancel := context.WithTimeout(r.ctx, HopConnectTimeout)
	defer cancel()

	bs, err := r.host.NewStream(ctx, dst.ID, ProtoID)
	if err != nil {
		log.Debugf("error opening relay stream to %s: %s", dst.ID.Pretty(), err.Error())
		r.handleError(s, pb.CircuitRelay_HOP_CANT_DIAL_DST)
		return
	}

	// stop handshake
	rd := newDelimitedReader(bs, maxMessageSize)
	wr := newDelimitedWriter(bs)

	msg.Type = pb.CircuitRelay_STOP.Enum()

	err = wr.WriteMsg(msg)
	if err != nil {
		log.Debugf("error writing stop handshake: %s", err.Error())
		bs.Close()
		r.handleError(s, pb.CircuitRelay_HOP_CANT_OPEN_DST_STREAM)
		return
	}

	msg.Reset()

	err = rd.ReadMsg(msg)
	if err != nil {
		log.Debugf("error reading stop response: %s", err.Error())
		bs.Close()
		r.handleError(s, pb.CircuitRelay_HOP_CANT_OPEN_DST_STREAM)
		return
	}

	if msg.GetType() != pb.CircuitRelay_STATUS {
		log.Debugf("unexpected relay stop response: not a status message (%d)", msg.GetType())
		bs.Close()
		r.handleError(s, pb.CircuitRelay_HOP_CANT_OPEN_DST_STREAM)
		return
	}

	if msg.GetCode() != pb.CircuitRelay_SUCCESS {
		log.Debugf("relay stop failure: %d", msg.GetCode())
		bs.Close()
		r.handleError(s, msg.GetCode())
		return
	}

	err = r.writeResponse(s, pb.CircuitRelay_SUCCESS)
	if err != nil {
		log.Debugf("error writing relay response: %s", err.Error())
		bs.Close()
		s.Close()
		return
	}

	// relay connection
	log.Infof("relaying connection between %s and %s", src.ID.Pretty(), dst.ID.Pretty())

	go func() {
		count, err := io.Copy(s, bs)
		if err != io.EOF && err != nil {
			log.Debugf("relay copy error: %s", err)
		}
		s.Close()
		log.Debugf("relayed %d bytes from %s to %s", count, dst.ID.Pretty(), src.ID.Pretty())
	}()

	go func() {
		count, err := io.Copy(bs, s)
		if err != io.EOF && err != nil {
			log.Debugf("relay copy error: %s", err)
		}
		bs.Close()
		log.Debugf("relayed %d bytes from %s to %s", count, src.ID.Pretty(), dst.ID.Pretty())
	}()
}

func (r *Relay) handleStopStream(s inet.Stream, msg *pb.CircuitRelay) {
	src, err := peerToPeerInfo(msg.GetSrcPeer())
	if err != nil || len(src.Addrs) == 0 {
		r.handleError(s, pb.CircuitRelay_STOP_SRC_MULTIADDR_INVALID)
		return
	}

	dst, err := peerToPeerInfo(msg.GetDstPeer())
	if err != nil || dst.ID != r.self {
		r.handleError(s, pb.CircuitRelay_STOP_DST_MULTIADDR_INVALID)
		return
	}

	log.Infof("relay connection from: %s", src.ID)

	r.host.Peerstore().AddAddrs(src.ID, src.Addrs, pstore.TempAddrTTL)

	select {
	case r.incoming <- &Conn{Stream: s, remote: src, transport: r.Transport()}:
	case <-time.After(RelayAcceptTimeout):
		r.handleError(s, pb.CircuitRelay_STOP_RELAY_REFUSED)
	}
}

func (r *Relay) handleCanHop(s inet.Stream, msg *pb.CircuitRelay) {
	var err error

	if r.hop {
		err = r.writeResponse(s, pb.CircuitRelay_SUCCESS)
	} else {
		err = r.writeResponse(s, pb.CircuitRelay_HOP_CANT_SPEAK_RELAY)
	}

	if err != nil {
		log.Debugf("error writing relay response: %s", err.Error())
	}

	s.Close()
}

func (r *Relay) handleError(s inet.Stream, code pb.CircuitRelay_Status) {
	log.Warningf("relay error: %s (%d)", pb.CircuitRelay_Status_name[int32(code)], code)
	err := r.writeResponse(s, code)
	if err != nil {
		log.Debugf("error writing relay response: %s", err.Error())
	}
	s.Close()
}

func (r *Relay) writeResponse(s inet.Stream, code pb.CircuitRelay_Status) error {
	wr := newDelimitedWriter(s)

	var msg pb.CircuitRelay
	msg.Type = pb.CircuitRelay_STATUS.Enum()
	msg.Code = code.Enum()

	return wr.WriteMsg(&msg)
}
<commit_msg>relay: remove unused fields from Relay<commit_after>package relay

import (
	"context"
	"fmt"
	"io"
	"time"

	pb "github.com/libp2p/go-libp2p-circuit/pb"

	logging "github.com/ipfs/go-log"
	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
)

var log = logging.Logger("relay")

const ProtoID = "/libp2p/circuit/relay/0.1.0"

const maxMessageSize = 4096

var RelayAcceptTimeout = time.Minute
var HopConnectTimeout = 10 * time.Second

type Relay struct {
	host host.Host
	ctx  context.Context
	self peer.ID

	active bool
	hop    bool

	incoming chan *Conn
}

type RelayOpt int

var (
	OptActive = RelayOpt(0)
	OptHop    = RelayOpt(1)
)

type RelayError struct {
	Code pb.CircuitRelay_Status
}

func (e RelayError) Error() string {
	return fmt.Sprintf("error opening relay circuit: %s (%d)", pb.CircuitRelay_Status_name[int32(e.Code)], e.Code)
}

func NewRelay(ctx context.Context, h host.Host, opts ...RelayOpt) (*Relay, error) {
	r := &Relay{
		host:     h,
		ctx:      ctx,
		self:     h.ID(),
		incoming: make(chan *Conn),
	}

	for _, opt := range opts {
		switch opt {
		case OptActive:
			r.active = true
		case OptHop:
			r.hop = true
		default:
			return nil, fmt.Errorf("unrecognized option: %d", opt)
		}
	}

	h.SetStreamHandler(ProtoID, r.handleNewStream)

	return r, nil
}

func (r *Relay) DialPeer(ctx context.Context, relay pstore.PeerInfo, dest pstore.PeerInfo) (*Conn, error) {

	if len(relay.Addrs) > 0 {
		r.host.Peerstore().AddAddrs(relay.ID, relay.Addrs, pstore.TempAddrTTL)
	}

	s, err := r.host.NewStream(ctx, relay.ID, ProtoID)
	if err != nil {
		return nil, err
	}

	rd := newDelimitedReader(s, maxMessageSize)
	wr := newDelimitedWriter(s)

	var msg pb.CircuitRelay

	msg.Type = pb.CircuitRelay_HOP.Enum()
	msg.SrcPeer = peerInfoToPeer(r.host.Peerstore().PeerInfo(r.self))
	msg.DstPeer = peerInfoToPeer(dest)

	err = wr.WriteMsg(&msg)
	if err != nil {
		s.Close()
		return nil, err
	}

	msg.Reset()

	err = rd.ReadMsg(&msg)
	if err != nil {
		s.Close()
		return nil, err
	}

	if msg.GetType() != pb.CircuitRelay_STATUS {
		s.Close()
		return nil, fmt.Errorf("unexpected relay response; not a status message (%d)", msg.GetType())
	}

	if msg.GetCode() != pb.CircuitRelay_SUCCESS {
		s.Close()
		return nil, RelayError{msg.GetCode()}
	}

	return &Conn{Stream: s, remote: dest, transport: r.Transport()}, nil
}

func (r *Relay) handleNewStream(s inet.Stream) {
	log.Infof("new relay stream from: %s", s.Conn().RemotePeer())

	rd := newDelimitedReader(s, maxMessageSize)

	var msg pb.CircuitRelay

	err := rd.ReadMsg(&msg)
	if err != nil {
		r.handleError(s, pb.CircuitRelay_MALFORMED_MESSAGE)
		return
	}

	switch msg.GetType() {
	case pb.CircuitRelay_HOP:
		r.handleHopStream(s, &msg)
	case pb.CircuitRelay_STOP:
		r.handleStopStream(s, &msg)
	case pb.CircuitRelay_CAN_HOP:
		r.handleCanHop(s, &msg)
	default:
		log.Warningf("unexpected relay handshake: %d", msg.GetType())
		r.handleError(s, pb.CircuitRelay_MALFORMED_MESSAGE)
	}
}

func (r *Relay) handleHopStream(s inet.Stream, msg *pb.CircuitRelay) {
	if !r.hop {
		r.handleError(s, pb.CircuitRelay_HOP_CANT_SPEAK_RELAY)
		return
	}

	src, err := peerToPeerInfo(msg.GetSrcPeer())
	if err != nil {
		r.handleError(s, pb.CircuitRelay_HOP_SRC_MULTIADDR_INVALID)
		return
	}

	if src.ID != s.Conn().RemotePeer() {
		r.handleError(s, pb.CircuitRelay_HOP_SRC_MULTIADDR_INVALID)
		return
	}

	dst, err := peerToPeerInfo(msg.GetDstPeer())
	if err != nil {
		r.handleError(s, pb.CircuitRelay_HOP_DST_MULTIADDR_INVALID)
		return
	}

	if dst.ID == r.self {
		r.handleError(s, pb.CircuitRelay_HOP_CANT_RELAY_TO_SELF)
		return
	}

	// open stream
	ctp := r.host.Network().ConnsToPeer(dst.ID)

	if len(ctp) == 0 && !r.active {
		r.handleError(s, pb.CircuitRelay_HOP_NO_CONN_TO_DST)
		return
	}

	if len(dst.Addrs) > 0 {
		r.host.Peerstore().AddAddrs(dst.ID, dst.Addrs, pstore.TempAddrTTL)
	}

	ctx, cancel := context.WithTimeout(r.ctx, HopConnectTimeout)
	defer cancel()

	bs, err := r.host.NewStream(ctx, dst.ID, ProtoID)
	if err != nil {
		log.Debugf("error opening relay stream to %s: %s", dst.ID.Pretty(), err.Error())
		r.handleError(s, pb.CircuitRelay_HOP_CANT_DIAL_DST)
		return
	}

	// stop handshake
	rd := newDelimitedReader(bs, maxMessageSize)
	wr := newDelimitedWriter(bs)

	msg.Type = pb.CircuitRelay_STOP.Enum()

	err = wr.WriteMsg(msg)
	if err != nil {
		log.Debugf("error writing stop handshake: %s", err.Error())
		bs.Close()
		r.handleError(s, pb.CircuitRelay_HOP_CANT_OPEN_DST_STREAM)
		return
	}

	msg.Reset()

	err = rd.ReadMsg(msg)
	if err != nil {
		log.Debugf("error reading stop response: %s", err.Error())
		bs.Close()
		r.handleError(s, pb.CircuitRelay_HOP_CANT_OPEN_DST_STREAM)
		return
	}

	if msg.GetType() != pb.CircuitRelay_STATUS {
		log.Debugf("unexpected relay stop response: not a status message (%d)", msg.GetType())
		bs.Close()
		r.handleError(s, pb.CircuitRelay_HOP_CANT_OPEN_DST_STREAM)
		return
	}

	if msg.GetCode() != pb.CircuitRelay_SUCCESS {
		log.Debugf("relay stop failure: %d", msg.GetCode())
		bs.Close()
		r.handleError(s, msg.GetCode())
		return
	}

	err = r.writeResponse(s, pb.CircuitRelay_SUCCESS)
	if err != nil {
		log.Debugf("error writing relay response: %s", err.Error())
		bs.Close()
		s.Close()
		return
	}

	// relay connection
	log.Infof("relaying connection between %s and %s", src.ID.Pretty(), dst.ID.Pretty())

	go func() {
		count, err := io.Copy(s, bs)
		if err != io.EOF && err != nil {
			log.Debugf("relay copy error: %s", err)
		}
		s.Close()
		log.Debugf("relayed %d bytes from %s to %s", count, dst.ID.Pretty(), src.ID.Pretty())
	}()

	go func() {
		count, err := io.Copy(bs, s)
		if err != io.EOF && err != nil {
			log.Debugf("relay copy error: %s", err)
		}
		bs.Close()
		log.Debugf("relayed %d bytes from %s to %s", count, src.ID.Pretty(), dst.ID.Pretty())
	}()
}

func (r *Relay) handleStopStream(s inet.Stream, msg *pb.CircuitRelay) {
	src, err := peerToPeerInfo(msg.GetSrcPeer())
	if err != nil || len(src.Addrs) == 0 {
		r.handleError(s, pb.CircuitRelay_STOP_SRC_MULTIADDR_INVALID)
		return
	}

	dst, err := peerToPeerInfo(msg.GetDstPeer())
	if err != nil || dst.ID != r.self {
		r.handleError(s, pb.CircuitRelay_STOP_DST_MULTIADDR_INVALID)
		return
	}

	log.Infof("relay connection from: %s", src.ID)

	r.host.Peerstore().AddAddrs(src.ID, src.Addrs, pstore.TempAddrTTL)

	select {
	case r.incoming <- &Conn{Stream: s, remote: src, transport: r.Transport()}:
	case <-time.After(RelayAcceptTimeout):
		r.handleError(s, pb.CircuitRelay_STOP_RELAY_REFUSED)
	}
}

func (r *Relay) handleCanHop(s inet.Stream, msg *pb.CircuitRelay) {
	var err error

	if r.hop {
		err = r.writeResponse(s, pb.CircuitRelay_SUCCESS)
	} else {
		err = r.writeResponse(s, pb.CircuitRelay_HOP_CANT_SPEAK_RELAY)
	}

	if err != nil {
		log.Debugf("error writing relay response: %s", err.Error())
	}

	s.Close()
}

func (r *Relay) handleError(s inet.Stream, code pb.CircuitRelay_Status) {
	log.Warningf("relay error: %s (%d)", pb.CircuitRelay_Status_name[int32(code)], code)
	err := r.writeResponse(s, code)
	if err != nil {
		log.Debugf("error writing relay response: %s", err.Error())
	}
	s.Close()
}

func (r *Relay) writeResponse(s inet.Stream, code pb.CircuitRelay_Status) error {
	wr := newDelimitedWriter(s)

	var msg pb.CircuitRelay
	msg.Type = pb.CircuitRelay_STATUS.Enum()
	msg.Code = code.Enum()

	return wr.WriteMsg(&msg)
}
<|endoftext|>