<commit_before>
package main

import (
	"encoding/json"
	"os"
	"time"
)

// The direction of the tracks indicates where a rabbit went
// from here.
type TrackDirection uint

const (
	// Default value.
	MinRabbits	= 1
	// Default value. The numbers of rabbits that exist
	// at any given time.
	MaxRabbits	= 15
	// Spawn chance for rabbits.
	SpawnChance	= 0.20

	// Chance to ascend deeper (closer to /). The weight
	// has to be fair, because ascending is very limited
	// and you only have one option.
	AscendChance	= 0.30

	// Chance to move twice instead of once.
	TwoStepChance	= 0.50

	// How long it takes for tracks to fade. Right
	// now, it's as long as it takes for a rabbit to
	// move.
	TrackFadeTime	= IdleTime
)

const (
	// No tracks.
	TrackNone TrackDirection = iota
	// Ascending is going "up" a directory, like cd ..
	TrackAscending
	// Descending is going "down" a directory, like cd ./data
	TrackDescending
)

type track struct {
	Timestamp	time.Time
	Direction	TrackDirection
}

type directoryForest struct {
	// List of rabbits and their locations. Only one
	// rabbit per location.
	rabbits		map[string]*Rabbit
	// Tracks at a given location. Cleared and updated after every move.
	tracks		map[string]track
	// Number of rabbits seen.
	spottedCount	uint
	// Number of rabbits caught.
	caughtCount	uint
	// Number of rabbits killed. :(
	killedCount	uint
}

func newDirectoryForest() directoryForest {
	return directoryForest{
		map[string]*Rabbit{}, map[string]track{}, 0, 0, 0,
	}
}

// For a location to exist, the directory must exist.
func (f *directoryForest) LocationExists(loc string) bool {
	fi, err := os.Stat(loc)
	if err != nil {
		return false
	}
	return fi.IsDir()
}

// Returns a location near the passed location. Nearby is
// found by a small number of random directory changes. Will
// not be the same directory, unless it has to (can't ascend
// or descend).
//
// XXX: Currently this does not check if a rabbit already
// exists at the new location. So we just lose rabbits
// if one encounters another.
func (f *directoryForest) NearbyLocation(loc string) string {
	newloc := loc

	steps := 1
	if chance(TwoStepChance) {
		steps = 2
	}

tryagain:
	added := []string{}
	for i := 0; i < steps; i++ {
		// Can't move.
		if !canAscend(newloc) && !canDescend(newloc) {
			return newloc
		} else if chance(AscendChance) {
			if canAscend(newloc) {
				newloc = ascend(newloc)
			} else {
				newloc = randDescension(newloc)
			}
		} else {
			if canDescend(newloc) {
				newloc = randDescension(newloc)
			} else {
				newloc = ascend(newloc)
			}
		}

		added = append(added, newloc)
	}

	if newloc == loc {
		// Guaranteed to not be the same because you must
		// step twice to get to the same destination.
		steps = 1
		goto tryagain
	}

	pastLoc := loc
	for _, aloc := range added {
		if isAscension(aloc, pastLoc) {
			f.tracks[pastLoc] = track{time.Now(), TrackAscending}
		} else if isDescension(aloc, pastLoc) {
			f.tracks[pastLoc] = track{time.Now(), TrackDescending}
		} else {
			panic("Rabbit didn't move to nearby location.")
		}
		pastLoc = aloc
	}

	return newloc

}

// A random faraway location. Rabbits typically start here
// and run here when they're fleeing. Faraway locations don't
// add tracks.
func (f *directoryForest) FarawayLocation(loc string) string {
	newloc := baseLocation()
	triedagain := false

	steps := 1
	if chance(TwoStepChance) {
		steps = 2
	}

tryagain:
	for i := 0; i < steps; i++ {
		if canDescend(newloc) {
			newloc = randDescension(newloc)
		}
	}

	if newloc == loc && !triedagain {
		// Invert the steps. We can't get to the same
		// location with different steps.
		if steps == 1 {
			steps = 2
		} else if steps == 2 {
			steps = 1
		}
		triedagain = true
		goto tryagain
	}

	return newloc
}

// Returns true if a rabbit is here. Only useful for checking
// before performing an action.
func (f *directoryForest) IsRabbitHere() bool {
	loc, _ := os.Getwd()
	_, ok := f.rabbits[loc]
	return ok
}

// Returns whether tracks
func (f *directoryForest) GetTracksHere() (bool, TrackDirection) {
	loc, _ := os.Getwd()
	t, ok := f.tracks[loc]
	if ok {
		return true, t.Direction
	} else {
		return false, TrackNone
	}
}

// Anytime a location is entered, a check is performed. This
// function updates every rabbit and returns a rabbit if one
// is spotted.
func (f *directoryForest) PerformCheck() (spotted *Rabbit) {
	spotted = nil

	// We always check our current directory.
	loc, _ := os.Getwd()

	newrabbits := map[string]*Rabbit{}

	for _, r := range f.rabbits {
		r.DisturbanceAt(loc)

		if (r.IsPlaying()) {
			if r.JustSpotted() {
				// It's possible for two rabbits to "wakeup"
				// to the same location in the same update.
				// Right now the most recent on will not be
				// overridden so we spotted that one.
				//
				// XXX: Fix rabbits running into each other?
				spotted = r
				f.spottedCount++
			}
			newrabbits[r.Location()] = r
		} else {
			if r.State() == Dead {
				f.killedCount++
			} else if r.State() == Caught {
				// Update in PerformCatch, otherwise
				// catching a rabbit score won't
				// be update until the next call to this
				// function.
			}
		}
	}

	f.rabbits = newrabbits

	// See if we should repopulate.
	f.repopulate()

	f.fadeTracks()

	return
}

// Attempts to catch a rabbit if it's still where we are.
func (f *directoryForest) PerformCatch() bool {
	loc, _ := os.Getwd()

	f.fadeTracks()

	rab, ok := f.rabbits[loc]
	if ok {
		delete(f.rabbits, rab.Location())
		succ := rab.TryCatch(loc)
		// We must update the table, else we can run into two rabbits.
		f.rabbits[rab.Location()] = rab
		if succ {
			f.caughtCount++
		}
		return succ
	}

	return false
}

// Attempts to tag a rabbit if it's still where we are.
func (f *directoryForest) PerformTag(tag string) bool {
	loc, _ := os.Getwd()

	f.fadeTracks()

	rab, ok := f.rabbits[loc]
	if ok {
		delete(f.rabbits, rab.Location())
		succ := rab.TryTag(loc, tag)
		// We must update the table, else we can run into two rabbits.
		f.rabbits[rab.Location()] = rab
		return succ
	}

	return false
}

// Repopulated the forest if under the minimum number of rabbits
// we want. Otherwise, chance a rabbit will spawn.
func (f *directoryForest) repopulate() {
	for len(f.rabbits) < MinRabbits {
		r := NewRabbit(f)
		f.rabbits[r.Location()] = &r
	}

	if chance(SpawnChance) {
		r := NewRabbit(f)
		f.rabbits[r.Location()] = &r
	}
}

// Fades the tracks depending on how old they are. Faded
// tracks are removed.
func (f *directoryForest) fadeTracks() {
	list := []string{}
	for loc, track := range f.tracks {
		age := time.Now().Sub(track.Timestamp)
		if age >= TrackFadeTime {
			list = append(list, loc)
		}
	}

	for _, loc := range list {
		delete(f.tracks, loc)
	}
}

// Used for marshalling/unmarshalling.
type forest struct {
	Rabbits		map[string]*Rabbit
	Tracks		map[string]track
	SpottedCount	uint
	CaughtCount	uint
	KilledCount	uint
}

// These are implemented because we can't encode private fields.
func (f *directoryForest) UnmarshalJSON(b []byte) error {
	data := forest{}
	err := json.Unmarshal(b, &data)
	if err != nil {
		return err
	}
	f.rabbits = data.Rabbits
	f.tracks = data.Tracks
	f.spottedCount = data.SpottedCount
	f.caughtCount = data.CaughtCount
	f.killedCount = data.KilledCount

	// Circular reference. Couldn't marshal their home so
	// we do it here.
	for _, r := range f.rabbits {
		(&r).ChangeHome(f)
	}
	return nil
}

func (f *directoryForest) MarshalJSON() ([]byte, error) {
	return json.Marshal(&forest{
		Rabbits:	f.rabbits,
		Tracks:		f.tracks,
		SpottedCount:	f.spottedCount,
		CaughtCount:	f.caughtCount,
		KilledCount:	f.killedCount,
	})
}
<commit_msg>Changed the lifetime of rabbit tracks.<commit_after>
package main

import (
	"encoding/json"
	"os"
	"time"
)

// The direction of the tracks indicates where a rabbit went
// from here.
type TrackDirection uint

const (
	// Default value.
	MinRabbits	= 1
	// Default value. The numbers of rabbits that exist
	// at any given time.
	MaxRabbits	= 15
	// Spawn chance for rabbits.
	SpawnChance	= 0.20

	// Chance to ascend deeper (closer to /). The weight
	// has to be fair, because ascending is very limited
	// and you only have one option.
	AscendChance	= 0.30

	// Chance to move twice instead of once.
	TwoStepChance	= 0.50

	// How long it takes for tracks to fade. Right
	// now, it's a 1/5 of the time it takes a rabbit
	// to move.
	TrackFadeTime	= IdleTime / 5
)

const (
	// No tracks.
	TrackNone TrackDirection = iota
	// Ascending is going "up" a directory, like cd ..
	TrackAscending
	// Descending is going "down" a directory, like cd ./data
	TrackDescending
)

type track struct {
	Timestamp	time.Time
	Direction	TrackDirection
}

type directoryForest struct {
	// List of rabbits and their locations. Only one
	// rabbit per location.
	rabbits		map[string]*Rabbit
	// Tracks at a given location. Cleared and updated after every move.
	tracks		map[string]track
	// Number of rabbits seen.
	spottedCount	uint
	// Number of rabbits caught.
	caughtCount	uint
	// Number of rabbits killed. :(
	killedCount	uint
}

func newDirectoryForest() directoryForest {
	return directoryForest{
		map[string]*Rabbit{}, map[string]track{}, 0, 0, 0,
	}
}

// For a location to exist, the directory must exist.
func (f *directoryForest) LocationExists(loc string) bool {
	fi, err := os.Stat(loc)
	if err != nil {
		return false
	}
	return fi.IsDir()
}

// Returns a location near the passed location. Nearby is
// found by a small number of random directory changes. Will
// not be the same directory, unless it has to (can't ascend
// or descend).
//
// XXX: Currently this does not check if a rabbit already
// exists at the new location. So we just lose rabbits
// if one encounters another.
func (f *directoryForest) NearbyLocation(loc string) string {
	newloc := loc

	steps := 1
	if chance(TwoStepChance) {
		steps = 2
	}

tryagain:
	added := []string{}
	for i := 0; i < steps; i++ {
		// Can't move.
		if !canAscend(newloc) && !canDescend(newloc) {
			return newloc
		} else if chance(AscendChance) {
			if canAscend(newloc) {
				newloc = ascend(newloc)
			} else {
				newloc = randDescension(newloc)
			}
		} else {
			if canDescend(newloc) {
				newloc = randDescension(newloc)
			} else {
				newloc = ascend(newloc)
			}
		}

		added = append(added, newloc)
	}

	if newloc == loc {
		// Guaranteed to not be the same because you must
		// step twice to get to the same destination.
		steps = 1
		goto tryagain
	}

	pastLoc := loc
	for _, aloc := range added {
		if isAscension(aloc, pastLoc) {
			f.tracks[pastLoc] = track{time.Now(), TrackAscending}
		} else if isDescension(aloc, pastLoc) {
			f.tracks[pastLoc] = track{time.Now(), TrackDescending}
		} else {
			panic("Rabbit didn't move to nearby location.")
		}
		pastLoc = aloc
	}

	return newloc

}

// A random faraway location. Rabbits typically start here
// and run here when they're fleeing. Faraway locations don't
// add tracks.
func (f *directoryForest) FarawayLocation(loc string) string {
	newloc := baseLocation()
	triedagain := false

	steps := 1
	if chance(TwoStepChance) {
		steps = 2
	}

tryagain:
	for i := 0; i < steps; i++ {
		if canDescend(newloc) {
			newloc = randDescension(newloc)
		}
	}

	if newloc == loc && !triedagain {
		// Invert the steps. We can't get to the same
		// location with different steps.
		if steps == 1 {
			steps = 2
		} else if steps == 2 {
			steps = 1
		}
		triedagain = true
		goto tryagain
	}

	return newloc
}

// Returns true if a rabbit is here. Only useful for checking
// before performing an action.
func (f *directoryForest) IsRabbitHere() bool {
	loc, _ := os.Getwd()
	_, ok := f.rabbits[loc]
	return ok
}

// Returns whether tracks
func (f *directoryForest) GetTracksHere() (bool, TrackDirection) {
	loc, _ := os.Getwd()
	t, ok := f.tracks[loc]
	if ok {
		return true, t.Direction
	} else {
		return false, TrackNone
	}
}

// Anytime a location is entered, a check is performed. This
// function updates every rabbit and returns a rabbit if one
// is spotted.
func (f *directoryForest) PerformCheck() (spotted *Rabbit) {
	spotted = nil

	// We always check our current directory.
	loc, _ := os.Getwd()

	newrabbits := map[string]*Rabbit{}

	for _, r := range f.rabbits {
		r.DisturbanceAt(loc)

		if (r.IsPlaying()) {
			if r.JustSpotted() {
				// It's possible for two rabbits to "wakeup"
				// to the same location in the same update.
				// Right now the most recent on will not be
				// overridden so we spotted that one.
				//
				// XXX: Fix rabbits running into each other?
				spotted = r
				f.spottedCount++
			}
			newrabbits[r.Location()] = r
		} else {
			if r.State() == Dead {
				f.killedCount++
			} else if r.State() == Caught {
				// Update in PerformCatch, otherwise
				// catching a rabbit score won't
				// be update until the next call to this
				// function.
			}
		}
	}

	f.rabbits = newrabbits

	// See if we should repopulate.
	f.repopulate()

	f.fadeTracks()

	return
}

// Attempts to catch a rabbit if it's still where we are.
func (f *directoryForest) PerformCatch() bool {
	loc, _ := os.Getwd()

	f.fadeTracks()

	rab, ok := f.rabbits[loc]
	if ok {
		delete(f.rabbits, rab.Location())
		succ := rab.TryCatch(loc)
		// We must update the table, else we can run into two rabbits.
		f.rabbits[rab.Location()] = rab
		if succ {
			f.caughtCount++
		}
		return succ
	}

	return false
}

// Attempts to tag a rabbit if it's still where we are.
func (f *directoryForest) PerformTag(tag string) bool {
	loc, _ := os.Getwd()

	f.fadeTracks()

	rab, ok := f.rabbits[loc]
	if ok {
		delete(f.rabbits, rab.Location())
		succ := rab.TryTag(loc, tag)
		// We must update the table, else we can run into two rabbits.
		f.rabbits[rab.Location()] = rab
		return succ
	}

	return false
}

// Repopulated the forest if under the minimum number of rabbits
// we want. Otherwise, chance a rabbit will spawn.
func (f *directoryForest) repopulate() {
	for len(f.rabbits) < MinRabbits {
		r := NewRabbit(f)
		f.rabbits[r.Location()] = &r
	}

	if chance(SpawnChance) {
		r := NewRabbit(f)
		f.rabbits[r.Location()] = &r
	}
}

// Fades the tracks depending on how old they are. Faded
// tracks are removed.
func (f *directoryForest) fadeTracks() {
	list := []string{}
	for loc, track := range f.tracks {
		age := time.Now().Sub(track.Timestamp)
		if age >= TrackFadeTime {
			list = append(list, loc)
		}
	}

	for _, loc := range list {
		delete(f.tracks, loc)
	}
}

// Used for marshalling/unmarshalling.
type forest struct {
	Rabbits		map[string]*Rabbit
	Tracks		map[string]track
	SpottedCount	uint
	CaughtCount	uint
	KilledCount	uint
}

// These are implemented because we can't encode private fields.
func (f *directoryForest) UnmarshalJSON(b []byte) error {
	data := forest{}
	err := json.Unmarshal(b, &data)
	if err != nil {
		return err
	}
	f.rabbits = data.Rabbits
	f.tracks = data.Tracks
	f.spottedCount = data.SpottedCount
	f.caughtCount = data.CaughtCount
	f.killedCount = data.KilledCount

	// Circular reference. Couldn't marshal their home so
	// we do it here.
	for _, r := range f.rabbits {
		(&r).ChangeHome(f)
	}
	return nil
}

func (f *directoryForest) MarshalJSON() ([]byte, error) {
	return json.Marshal(&forest{
		Rabbits:	f.rabbits,
		Tracks:		f.tracks,
		SpottedCount:	f.spottedCount,
		CaughtCount:	f.caughtCount,
		KilledCount:	f.killedCount,
	})
}
<|endoftext|>


<commit_before>package mint

import "testing"

// Mint (mint.Mint) is wrapper for *testing.T
// blending testing type to omit repeated `t`.
type Mint struct {
	t *testing.T
}

var (
	failToBe = 0
	failType = 1
	scolds   = map[int]string{
		failToBe: "%s at line %d\nExpected %sto be\t`%+v`\nBut actual\t`%+v`",
		failType: "%s at line %d\nExpected %stype\t`%+v`\nBut actual\t`%T`",
	}
)
var (
	redB     = "\033[1;31m"
	reset    = "\033[0m"
	colorize = map[string]func(string) string{
		"red": func(v string) string {
			return redB + v + reset
		},
	}
)

// Blend provides (blended) *mint.Mint.
// You can save writing "t" repeatedly.
func Blend(t *testing.T) *Mint {
	return &Mint{
		t,
	}
}

// Expect provides "*Testee".
// The blended mint is merely a proxy to instantiate testee.
func (m *Mint) Expect(actual interface{}) *Testee {
	return newTestee(m.t, actual)
}

// Expect provides "*mint.Testee".
// It has assertion methods such as "ToBe".
func Expect(t *testing.T, actual interface{}) *Testee {
	return newTestee(t, actual)
}

func newTestee(t *testing.T, actual interface{}) *Testee {
	return &Testee{t: t, actual: actual, result: Result{ok: true}}
}
func judge(a, b interface{}, not, deeply bool) bool {
	comparer := getComparer(a, b, deeply)
	if not {
		return !comparer.Compare(a, b)
	}
	return comparer.Compare(a, b)
}
<commit_msg>Fix presentation<commit_after>package mint

import "testing"

// Mint (mint.Mint) is wrapper for *testing.T
// blending testing type to omit repeated `t`.
type Mint struct {
	t *testing.T
}

var (
	failToBe = 0
	failType = 1
	scolds   = map[int]string{
		failToBe: "%s at line %d\n\tExpected %sto be\t`%+v`\n\tBut actual\t`%+v`",
		failType: "%s at line %d\n\tExpected %stype\t`%+v`\n\tBut actual\t`%T`",
	}
)
var (
	redB     = "\033[1;31m"
	reset    = "\033[0m"
	colorize = map[string]func(string) string{
		"red": func(v string) string {
			return redB + v + reset
		},
	}
)

// Blend provides (blended) *mint.Mint.
// You can save writing "t" repeatedly.
func Blend(t *testing.T) *Mint {
	return &Mint{
		t,
	}
}

// Expect provides "*Testee".
// The blended mint is merely a proxy to instantiate testee.
func (m *Mint) Expect(actual interface{}) *Testee {
	return newTestee(m.t, actual)
}

// Expect provides "*mint.Testee".
// It has assertion methods such as "ToBe".
func Expect(t *testing.T, actual interface{}) *Testee {
	return newTestee(t, actual)
}

func newTestee(t *testing.T, actual interface{}) *Testee {
	return &Testee{t: t, actual: actual, result: Result{ok: true}}
}
func judge(a, b interface{}, not, deeply bool) bool {
	comparer := getComparer(a, b, deeply)
	if not {
		return !comparer.Compare(a, b)
	}
	return comparer.Compare(a, b)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	cryptorand "crypto/rand"
	"crypto/sha512"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"os"
	"runtime"
	"time"
)

const (
	BSIZE    = 65536
	BSIZE7   = BSIZE - BSIZEMOD
	BSIZEMOD = BSIZE % 7
	MULTI    = 4
	SOURCES  = 8
	POOL     = SOURCES * (MULTI + 1)
)

var (
	narg   = int(0)
	phrase = ""
	pool   = make(chan []byte, POOL)
	loop   = make(chan []byte, POOL)
	data   = make([]chan []byte, SOURCES)
)

// randomSeed produces a int64 seed based on crypto/rand and time.
func randomSeed() int64 {
	var seed int64

	urandom := make([]byte, 8)
	cryptorand.Reader.Read(urandom)

	for key, value := range urandom {
		seed ^= (int64(value) ^ time.Now().UTC().UnixNano()) << (uint(key) * 8)
	}

	return seed
}

// randomBytes fills byte buffers with random data
func randomBytes(src rand.Source, out chan<- []byte) {
	var (
		r int64
		i = BSIZE
	)

	for buf, ok := <-pool; ok; buf, ok = <-pool {
		r = src.Int63()
		switch { // Go seems to eliminate impossible cases
		case BSIZEMOD == 6:
			buf[BSIZE-6] = byte(r >> 48)
			fallthrough
		case BSIZEMOD == 5:
			buf[BSIZE-5] = byte(r >> 32)
			fallthrough
		case BSIZEMOD == 4:
			buf[BSIZE-4] = byte(r >> 24)
			fallthrough
		case BSIZEMOD == 3:
			buf[BSIZE-3] = byte(r >> 16)
			fallthrough
		case BSIZEMOD == 2:
			buf[BSIZE-2] = byte(r >> 8)
			fallthrough
		case BSIZEMOD == 1:
			buf[BSIZE-1] = byte(r)
		}

		for i = 0; i < BSIZE7; i += 7 {
			r = src.Int63()
			buf[i] = byte(r)
			buf[i+1] = byte(r >> 8)
			buf[i+2] = byte(r >> 16)
			buf[i+3] = byte(r >> 24)
			buf[i+4] = byte(r >> 32)
			buf[i+5] = byte(r >> 40)
			buf[i+6] = byte(r >> 48)
		}

		out <- buf
	}
}

func garble(f *os.File, in <-chan []byte, out chan<- bool) {
	var (
		data []byte // data we read
		err  error  // I/O error
		h    int    // data index
		n    int    // data size
		pos  int64  // file pos
	)

	data = make([]byte, BSIZE)
	err = nil
	h = 0
	n = 0
	pos = 0

	for buf, ok := <-in; ok; buf, ok = <-in {
		for _, randombyte := range buf {
			if h == n {
				if h > 0 {
					f.WriteAt(data[0:h], pos)
					pos += int64(h)
				}

				h = 0
				n, err = f.Read(data)

				if err != nil && err != io.EOF {
					panic(err)
				}
				if n == 0 {
					close(out)
					for {
						<-in // sleep of no return
					}
				}
			}

			data[h] ^= randombyte
			h++
		}

		out <- true
		// close(out)
	}
}

func init() {
	flag.StringVar(&phrase, "phrase", "", "the Garble phrase, by default random")
	flag.Parse()

	narg = flag.NArg()

	if narg <= 0 {
		flag.Usage()
		os.Exit(1)
	}

	if phrase == "" {
		phrase = fmt.Sprintf("%016x", uint64(randomSeed()))
	}
}

func main() {
	// Use available CPUs:
	if runtime.GOMAXPROCS(0) == 1 &&
		runtime.NumCPU() > 1 &&
		os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	// Open files:
	files := make([]*os.File, narg)
	writers := make([]chan []byte, narg)
	signals := make([]chan bool, narg)

	for i, arg := range flag.Args() {
		f, err := os.OpenFile(arg, os.O_RDWR, 0666)
		if err != nil {
			panic(err)
		}
		defer f.Close()
		files[i] = f
		writers[i] = make(chan []byte, MULTI)
		signals[i] = make(chan bool, MULTI)
		go garble(files[i], writers[i], signals[i])
	}

	// Allocate byte buffer pool:
	buffer := make([]byte, BSIZE*POOL)

	for i := 0; i < BSIZE*POOL; i += BSIZE {
		pool <- buffer[i : i+BSIZE]
	}

	// Initialize random sources:
	hash := sha512.New()
	sum := make([]byte, hash.Size())

	for i := 0; i < SOURCES; i++ {
		var seed, s int64
		var err error

		hash.Write([]byte(":garble:" + phrase))
		hash.Sum(sum[:0])

		buf := bytes.NewReader(sum)
		s = 0

		for err == nil {
			err = binary.Read(buf, binary.LittleEndian, &s)
			seed ^= s
		}

		src := rand.NewSource(seed)
		data[i] = make(chan []byte, MULTI)
		go randomBytes(src, data[i])
	}

	// Route data channels:
	go func(data []chan []byte, writers []chan []byte, signals []chan bool) {
		var buf []byte
		for {
			for _, r := range data {
				buf = <-r

				for i, w := range writers {
					if signals[i] != nil {
						w <- buf
					}
				}

				loop <- buf
			}
		}
	}(data, writers, signals)

	var buf []byte

	for narg > 0 {
		buf = <-loop

		for i, s := range signals {
			if signals[i] != nil {
				_, ok := <-s

				if !ok {
					signals[i] = nil
					narg--
				}
			}
		}

		pool <- buf
	}

	fmt.Println("All done!")
}
<commit_msg>print phrase<commit_after>package main

import (
	"bytes"
	cryptorand "crypto/rand"
	"crypto/sha512"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"os"
	"runtime"
	"time"
)

const (
	BSIZE    = 65536
	BSIZE7   = BSIZE - BSIZEMOD
	BSIZEMOD = BSIZE % 7
	MULTI    = 4
	SOURCES  = 8
	POOL     = SOURCES * (MULTI + 1)
)

var (
	narg   = int(0)
	phrase = ""
	pool   = make(chan []byte, POOL)
	loop   = make(chan []byte, POOL)
	data   = make([]chan []byte, SOURCES)
)

// randomSeed produces a int64 seed based on crypto/rand and time.
func randomSeed() int64 {
	var seed int64

	urandom := make([]byte, 8)
	cryptorand.Reader.Read(urandom)

	for key, value := range urandom {
		seed ^= (int64(value) ^ time.Now().UTC().UnixNano()) << (uint(key) * 8)
	}

	return seed
}

// randomBytes fills byte buffers with random data
func randomBytes(src rand.Source, out chan<- []byte) {
	var (
		r int64
		i = BSIZE
	)

	for buf, ok := <-pool; ok; buf, ok = <-pool {
		r = src.Int63()
		switch { // Go seems to eliminate impossible cases
		case BSIZEMOD == 6:
			buf[BSIZE-6] = byte(r >> 48)
			fallthrough
		case BSIZEMOD == 5:
			buf[BSIZE-5] = byte(r >> 32)
			fallthrough
		case BSIZEMOD == 4:
			buf[BSIZE-4] = byte(r >> 24)
			fallthrough
		case BSIZEMOD == 3:
			buf[BSIZE-3] = byte(r >> 16)
			fallthrough
		case BSIZEMOD == 2:
			buf[BSIZE-2] = byte(r >> 8)
			fallthrough
		case BSIZEMOD == 1:
			buf[BSIZE-1] = byte(r)
		}

		for i = 0; i < BSIZE7; i += 7 {
			r = src.Int63()
			buf[i] = byte(r)
			buf[i+1] = byte(r >> 8)
			buf[i+2] = byte(r >> 16)
			buf[i+3] = byte(r >> 24)
			buf[i+4] = byte(r >> 32)
			buf[i+5] = byte(r >> 40)
			buf[i+6] = byte(r >> 48)
		}

		out <- buf
	}
}

func garble(f *os.File, in <-chan []byte, out chan<- bool) {
	var (
		data []byte // data we read
		err  error  // I/O error
		h    int    // data index
		n    int    // data size
		pos  int64  // file pos
	)

	data = make([]byte, BSIZE)
	err = nil
	h = 0
	n = 0
	pos = 0

	for buf, ok := <-in; ok; buf, ok = <-in {
		for _, randombyte := range buf {
			if h == n {
				if h > 0 {
					f.WriteAt(data[0:h], pos)
					pos += int64(h)
				}

				h = 0
				n, err = f.Read(data)

				if err != nil && err != io.EOF {
					panic(err)
				}
				if n == 0 {
					close(out)
					for {
						<-in // sleep of no return
					}
				}
			}

			data[h] ^= randombyte
			h++
		}

		out <- true
		// close(out)
	}
}

func init() {
	flag.StringVar(&phrase, "phrase", "", "the Garble phrase, by default random")
	flag.Parse()

	narg = flag.NArg()

	if narg <= 0 {
		flag.Usage()
		os.Exit(1)
	}

	if phrase == "" {
		phrase = fmt.Sprintf("%016x", uint64(randomSeed()))
	}

	fmt.Println("Using phrase:", phrase)
}

func main() {
	// Use available CPUs:
	if runtime.GOMAXPROCS(0) == 1 &&
		runtime.NumCPU() > 1 &&
		os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	// Open files:
	files := make([]*os.File, narg)
	writers := make([]chan []byte, narg)
	signals := make([]chan bool, narg)

	for i, arg := range flag.Args() {
		f, err := os.OpenFile(arg, os.O_RDWR, 0666)
		if err != nil {
			panic(err)
		}
		defer f.Close()
		files[i] = f
		writers[i] = make(chan []byte, MULTI)
		signals[i] = make(chan bool, MULTI)
		go garble(files[i], writers[i], signals[i])
	}

	// Allocate byte buffer pool:
	buffer := make([]byte, BSIZE*POOL)

	for i := 0; i < BSIZE*POOL; i += BSIZE {
		pool <- buffer[i : i+BSIZE]
	}

	// Initialize random sources:
	hash := sha512.New()
	sum := make([]byte, hash.Size())

	for i := 0; i < SOURCES; i++ {
		var seed, s int64
		var err error

		hash.Write([]byte(":garble:" + phrase))
		hash.Sum(sum[:0])

		buf := bytes.NewReader(sum)
		s = 0

		for err == nil {
			err = binary.Read(buf, binary.LittleEndian, &s)
			seed ^= s
		}

		src := rand.NewSource(seed)
		data[i] = make(chan []byte, MULTI)
		go randomBytes(src, data[i])
	}

	// Route data channels:
	go func(data []chan []byte, writers []chan []byte, signals []chan bool) {
		var buf []byte
		for {
			for _, r := range data {
				buf = <-r

				for i, w := range writers {
					if signals[i] != nil {
						w <- buf
					}
				}

				loop <- buf
			}
		}
	}(data, writers, signals)

	var buf []byte

	for narg > 0 {
		buf = <-loop

		for i, s := range signals {
			if signals[i] != nil {
				_, ok := <-s

				if !ok {
					signals[i] = nil
					narg--
				}
			}
		}

		pool <- buf
	}

	fmt.Println("All done!")
}
<|endoftext|>


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

// +build linux unix posix ansi

package ct

import (
	"fmt"
	"io"
)

var modeTable = [...]int{
	reset:      0,
	bold:       1,
	faint:      2,
	italic:     3,
	underline:  4,
	blinkSlow:  5,
	blinkRapid: 6,
	negative:   7,
	conceal:    8,
	crossedOut: 9,
}

const (
	csi      = "\x1b["
	resetSgr = csi + "0m"
)

func (m Mode) set(w io.Writer) {
	w.Write([]byte(csi))
	if m&(colorSet|colorMask)>>(colorWidth-2) == 0x3 {
		m |= Bold
	}
	var printed bool
	for _, d := range []Mode{30, 40} {
		if m&colorSet != 0 {
			if printed {
				w.Write([]byte{';'})
			}
			fmt.Fprintf(w, "%d", m&(colorMask>>1)+d)
			printed = true
		}
		m >>= colorWidth
	}
	for _, d := range []int{38, 48} {
		if m&xTermColorSet != 0 {
			if printed {
				w.Write([]byte{';'})
			}
			fmt.Fprintf(w, "%d;5;%d", d, m&xTermColorMask)
			printed = true
		}
		m >>= xTermColorWidth
	}
	for _, v := range modeTable {
		if m&1 != 0 {
			if printed {
				w.Write([]byte{';'})
			}
			fmt.Fprintf(w, "%d", v)
			printed = true
		}
		m >>= 1
	}
	w.Write([]byte{'m'})
}

func (m Mode) reset(w io.Writer) { w.Write([]byte(resetSgr)) }
<commit_msg>Reduce naive allocations/improve readability<commit_after>// Copyright ©2014 Dan Kortschak. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build linux unix posix ansi

package ct

import (
	"fmt"
	"io"
)

var modeTable = [...]int{
	reset:      0,
	bold:       1,
	faint:      2,
	italic:     3,
	underline:  4,
	blinkSlow:  5,
	blinkRapid: 6,
	negative:   7,
	conceal:    8,
	crossedOut: 9,
}

const (
	csi      = "\x1b["
	resetSgr = csi + "0m"
)

var semicolon = []byte{';'}

func (m Mode) set(w io.Writer) {
	w.Write([]byte(csi))
	if m&(colorSet|colorMask)>>(colorWidth-2) == 0x3 {
		m |= Bold
	}
	var printed bool
	for _, d := range []Mode{30, 40} {
		if m&colorSet != 0 {
			if printed {
				w.Write(semicolon)
			}
			fmt.Fprintf(w, "%d", m&(colorMask>>1)+d)
			printed = true
		}
		m >>= colorWidth
	}
	for _, d := range []int{38, 48} {
		if m&xTermColorSet != 0 {
			if printed {
				w.Write(semicolon)
			}
			fmt.Fprintf(w, "%d;5;%d", d, m&xTermColorMask)
			printed = true
		}
		m >>= xTermColorWidth
	}
	for _, v := range modeTable {
		if m&1 != 0 {
			if printed {
				w.Write(semicolon)
			}
			fmt.Fprintf(w, "%d", v)
			printed = true
		}
		m >>= 1
	}
	w.Write([]byte{'m'})
}

func (m Mode) reset(w io.Writer) { w.Write([]byte(resetSgr)) }
<|endoftext|>


<commit_before>package dnssd

import (
	"context"
	"fmt"
	"net"

	"github.com/brutella/dnssd/log"
	"github.com/miekg/dns"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
)

var (
	IPv4LinkLocalMulticast = net.ParseIP("224.0.0.251")
	IPv6LinkLocalMulticast = net.ParseIP("ff02::fb")

	AddrIPv4LinkLocalMulticast = &net.UDPAddr{
		IP:   IPv4LinkLocalMulticast,
		Port: 5353,
	}

	AddrIPv6LinkLocalMulticast = &net.UDPAddr{
		IP:   IPv6LinkLocalMulticast,
		Port: 5353,
	}

	TTLDefault  uint32 = 75 * 60 // Default ttl for mDNS resource records
	TTLHostname uint32 = 120     // TTL for mDNS resource records containing the host name
)

// Query is a mDNS query
type Query struct {
	msg   *dns.Msg       // The query message
	iface *net.Interface // The network interface to which the message is sent
}

// Response is a mDNS response
type Response struct {
	msg   *dns.Msg       // The response message
	addr  *net.UDPAddr   // Is nil for multicast response
	iface *net.Interface // The network interface to which the message is sent
}

// Request represents an incoming mDNS message
type Request struct {
	msg   *dns.Msg       // The message
	from  *net.UDPAddr   // The source addr of the message
	iface *net.Interface // The network interface from which the message was received
}

func (r Request) String() string {
	return fmt.Sprintf("%s@%s\n%v", r.from.IP, r.iface.Name, r.msg)
}

// MDNSConn represents a mDNS connection. It encapsulates an IPv4 and IPv6 UDP connection.
type MDNSConn interface {
	// SendQuery sends a mDNS query.
	SendQuery(q *Query) error

	// SendResponse sends a mDNS response
	SendResponse(resp *Response) error

	// Read returns a channel which receives mDNS messages
	Read(ctx context.Context) <-chan *Request

	// Clears the connection buffer
	Drain(ctx context.Context)

	// Close closes the connection
	Close()
}

type mdnsConn struct {
	ipv4 *ipv4.PacketConn
	ipv6 *ipv6.PacketConn
	ch   chan *Request
}

func NewMDNSConn() (MDNSConn, error) {
	return newMDNSConn()
}

func (c *mdnsConn) SendQuery(q *Query) error {
	return c.sendQuery(q.msg, q.iface)
}

func (c *mdnsConn) SendResponse(resp *Response) error {
	if resp.addr != nil {
		return c.sendResponseTo(resp.msg, resp.iface, resp.addr)
	}

	return c.sendResponse(resp.msg, resp.iface)
}

func (c *mdnsConn) Read(ctx context.Context) <-chan *Request {
	return c.read(ctx)
}

func (c *mdnsConn) Drain(ctx context.Context) {
	log.Debug.Println("Draining connection")
	for {
		select {
		case req := <-c.Read(ctx):
			log.Debug.Println("Ignoring msg from", req.from.IP)
			break
		default:
			return
		}
	}
}

func (c *mdnsConn) Close() {
	c.close()
}

func newMDNSConn() (*mdnsConn, error) {
	var errs []error
	var connIPv4 *ipv4.PacketConn
	var connIPv6 *ipv6.PacketConn

	if conn, err := net.ListenUDP("udp4", AddrIPv4LinkLocalMulticast); err != nil {
		errs = append(errs, err)
	} else {
		connIPv4 = ipv4.NewPacketConn(conn)
		connIPv4.SetControlMessage(ipv4.FlagInterface, true)
		connIPv4.SetMulticastLoopback(true)

		for _, iface := range multicastInterfaces() {
			if err := connIPv4.JoinGroup(&iface, &net.UDPAddr{IP: IPv4LinkLocalMulticast}); err != nil {
				log.Debug.Printf("Failed joining IPv4 %v: %v", iface.Name, err)
			} else {
				log.Debug.Printf("Joined IPv4 %v", iface.Name)
			}
		}
	}

	if conn, err := net.ListenUDP("udp6", AddrIPv6LinkLocalMulticast); err != nil {
		errs = append(errs, err)
	} else {
		connIPv6 = ipv6.NewPacketConn(conn)
		connIPv6.SetControlMessage(ipv6.FlagInterface, true)
		connIPv6.SetMulticastLoopback(true)

		for _, iface := range multicastInterfaces() {
			if err := connIPv6.JoinGroup(&iface, &net.UDPAddr{IP: IPv6LinkLocalMulticast}); err != nil {
				log.Debug.Printf("Failed joining IPv6 %v: %v", iface.Name, err)
			} else {
				log.Debug.Printf("Joined IPv6 %v", iface.Name)
			}
		}
	}

	if err := first(errs...); connIPv4 == nil && connIPv6 == nil {
		return nil, fmt.Errorf("Failed setting up UDP server: %v", err)
	}

	return &mdnsConn{
		ipv4: connIPv4,
		ipv6: connIPv6,
		ch:   make(chan *Request),
	}, nil
}

func (c *mdnsConn) enableMulticastLoopback(enable bool) {
	if c.ipv4 != nil {
		c.ipv4.SetMulticastLoopback(enable)
	}

	if c.ipv6 != nil {
		c.ipv6.SetMulticastLoopback(enable)
	}
}

func (c *mdnsConn) close() {
	if c.ipv4 != nil {
		c.ipv4.Close()
	}

	if c.ipv6 != nil {
		c.ipv6.Close()
	}
}

func (c *mdnsConn) read(ctx context.Context) <-chan *Request {
	c.readInto(ctx, c.ch)
	return c.ch
}

func (c *mdnsConn) readInto(ctx context.Context, ch chan *Request) {

	isDone := func(ctx context.Context) bool {
		return ctx.Err() == context.Canceled
	}

	if c.ipv4 != nil {
		go func() {
			buf := make([]byte, 65536)
			for {
				if isDone(ctx) {
					return
				}

				n, cm, from, err := c.ipv4.ReadFrom(buf)
				if err != nil {
					continue
				}

				udpAddr, ok := from.(*net.UDPAddr)
				if !ok {
					log.Info.Println("invalid source address")
					continue
				}

				var iface *net.Interface
				if cm != nil {
					iface, err = net.InterfaceByIndex(cm.IfIndex)
					if err != nil {
						continue
					}
				}

				if n > 0 {
					m := new(dns.Msg)
					if err := m.Unpack(buf); err == nil && !shouldIgnore(m) {
						ch <- &Request{m, udpAddr, iface}
					}
				}
			}
		}()
	}

	if c.ipv6 != nil {
		go func() {
			buf := make([]byte, 65536)
			for {
				if isDone(ctx) {
					return
				}

				n, cm, from, err := c.ipv6.ReadFrom(buf)
				if err != nil {
					continue
				}

				udpAddr, ok := from.(*net.UDPAddr)
				if !ok {
					log.Info.Println("invalid source address")
					continue
				}

				var iface *net.Interface
				if cm != nil {
					iface, err = net.InterfaceByIndex(cm.IfIndex)
					if err != nil {
						continue
					}
				}

				if n > 0 {
					m := new(dns.Msg)
					if err := m.Unpack(buf); err == nil && !shouldIgnore(m) {
						ch <- &Request{m, udpAddr, iface}
					}
				}
			}
		}()
	}
}

func (c *mdnsConn) sendQuery(m *dns.Msg, iface *net.Interface) error {
	sanitizeQuery(m)

	return c.writeMsg(m, iface)
}

func (c *mdnsConn) sendResponse(m *dns.Msg, iface *net.Interface) error {
	sanitizeResponse(m)

	return c.writeMsg(m, iface)
}

func (c *mdnsConn) sendResponseTo(m *dns.Msg, iface *net.Interface, addr *net.UDPAddr) error {
	sanitizeResponse(m)

	return c.writeMsgTo(m, iface, addr)
}

func (c *mdnsConn) writeMsg(m *dns.Msg, iface *net.Interface) error {
	var err error
	if c.ipv4 != nil {
		err = c.writeMsgTo(m, iface, AddrIPv4LinkLocalMulticast)
	}

	if c.ipv6 != nil {
		err = c.writeMsgTo(m, iface, AddrIPv6LinkLocalMulticast)
	}

	return err
}

func (c *mdnsConn) writeMsgTo(m *dns.Msg, iface *net.Interface, addr *net.UDPAddr) error {
	sanitizeMsg(m)

	var err error
	if c.ipv4 != nil && addr.IP.To4() != nil {
		if out, err := m.Pack(); err == nil {
			var ctrl *ipv4.ControlMessage
			if iface != nil {
				ctrl = &ipv4.ControlMessage{
					IfIndex: iface.Index,
				}
			}
			_, err = c.ipv4.WriteTo(out, ctrl, addr)
		}
	}

	if c.ipv6 != nil && addr.IP.To4() == nil {
		if out, err := m.Pack(); err == nil {
			var ctrl *ipv6.ControlMessage
			if iface != nil {
				ctrl = &ipv6.ControlMessage{
					IfIndex: iface.Index,
				}
			}
			_, err = c.ipv6.WriteTo(out, ctrl, addr)
		}
	}

	return err
}

func shouldIgnore(m *dns.Msg) bool {
	if m.Opcode != 0 {
		return true
	}

	if m.Rcode != 0 {
		return true
	}

	return false
}

func sanitizeResponse(m *dns.Msg) {
	if m.Question != nil && len(m.Question) > 0 {
		log.Info.Println("Multicast DNS responses MUST NOT contain any questions in the Question Section.  (RFC6762 6)")
		m.Question = nil
	}

	if !m.Response {
		log.Info.Println("In response messages the QR bit MUST be one (RFC6762 18.2)")
		m.Response = true
	}

	if !m.Authoritative {
		log.Info.Println("AA Bit bit MUST be set to one in response messages (RFC6762 18.4)")
		m.Authoritative = true
	}

	if m.Truncated {
		log.Info.Println("In multicast response messages, the TC bit MUST be zero on transmission. (RFC6762 18.5)")
		m.Truncated = false
	}
}

func sanitizeQuery(m *dns.Msg) {
	if m.Response {
		log.Info.Println("In query messages the QR bit MUST be zero (RFC6762 18.2)")
		m.Response = false
	}

	if m.Authoritative {
		log.Info.Println("AA Bit MUST be zero in query messages (RFC6762 18.4)")
		m.Authoritative = false
	}
}

func sanitizeMsg(m *dns.Msg) {
	if m.Opcode != 0 {
		log.Info.Println("In both multicast query and multicast response messages, the OPCODE MUST be zero on transmission (RFC6762 18.3)")
		m.Opcode = 0
	}

	if m.RecursionDesired {
		log.Info.Println("In both multicast query and multicast response messages, the Recursion Available bit MUST be zero on transmission. (RFC6762 18.7)")
		m.RecursionDesired = false
	}

	if m.Zero {
		log.Info.Println("In both query and response messages, the Zero bit MUST be zero on transmission (RFC6762 18.8)")
		m.Zero = false
	}

	if m.AuthenticatedData {
		log.Info.Println("In both multicast query and multicast response messages, the Authentic Data bit MUST be zero on transmission (RFC6762 18.9)")
		m.AuthenticatedData = false
	}

	if m.CheckingDisabled {
		log.Info.Println("In both multicast query and multicast response messages, the Checking Disabled bit MUST be zero on transmission (RFC6762 18.10)")
		m.CheckingDisabled = false
	}

	if m.Rcode != 0 {
		log.Info.Println("In both multicast query and multicast response messages, the Response Code MUST be zero on transmission. (RFC6762 18.11)")
		m.Rcode = 0
	}
}

func first(errs ...error) error {
	for _, err := range errs {
		if err != nil {
			return err
		}
	}

	return nil
}

// Sets the Top Bit of rrclass for all answer records (except PTR) to trigger a cache flush in the receivers.
func setAnswerCacheFlushBit(msg *dns.Msg) {
	// From RFC6762
	//    The most significant bit of the rrclass for a record in the Answer
	//    Section of a response message is the Multicast DNS cache-flush bit
	//    and is discussed in more detail below in Section 10.2, "Announcements
	//    to Flush Outdated Cache Entries".
	for _, a := range msg.Answer {
		switch a.(type) {
		case *dns.PTR:
			continue
		default:
			a.Header().Class |= (1 << 15)
		}
	}
}

// Sets the Top Bit of class to indicate the unicast responses are preferred for this question.
func setQuestionUnicast(q *dns.Question) {
	q.Qclass |= (1 << 15)
}

// Returns true if q requires unicast responses.
func isUnicastQuestion(q dns.Question) bool {
	// From RFC6762
	// 18.12.  Repurposing of Top Bit of qclass in Question Section
	//
	//    In the Question Section of a Multicast DNS query, the top bit of the
	//    qclass field is used to indicate that unicast responses are preferred
	//    for this particular question.  (See Section 5.4.)
	return q.Qclass&(1<<15) != 0
}
<commit_msg>Return error when sending fails<commit_after>package dnssd

import (
	"context"
	"fmt"
	"net"

	"github.com/brutella/dnssd/log"
	"github.com/miekg/dns"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
)

var (
	IPv4LinkLocalMulticast = net.ParseIP("224.0.0.251")
	IPv6LinkLocalMulticast = net.ParseIP("ff02::fb")

	AddrIPv4LinkLocalMulticast = &net.UDPAddr{
		IP:   IPv4LinkLocalMulticast,
		Port: 5353,
	}

	AddrIPv6LinkLocalMulticast = &net.UDPAddr{
		IP:   IPv6LinkLocalMulticast,
		Port: 5353,
	}

	TTLDefault  uint32 = 75 * 60 // Default ttl for mDNS resource records
	TTLHostname uint32 = 120     // TTL for mDNS resource records containing the host name
)

// Query is a mDNS query
type Query struct {
	msg   *dns.Msg       // The query message
	iface *net.Interface // The network interface to which the message is sent
}

// Response is a mDNS response
type Response struct {
	msg   *dns.Msg       // The response message
	addr  *net.UDPAddr   // Is nil for multicast response
	iface *net.Interface // The network interface to which the message is sent
}

// Request represents an incoming mDNS message
type Request struct {
	msg   *dns.Msg       // The message
	from  *net.UDPAddr   // The source addr of the message
	iface *net.Interface // The network interface from which the message was received
}

func (r Request) String() string {
	return fmt.Sprintf("%s@%s\n%v", r.from.IP, r.iface.Name, r.msg)
}

// MDNSConn represents a mDNS connection. It encapsulates an IPv4 and IPv6 UDP connection.
type MDNSConn interface {
	// SendQuery sends a mDNS query.
	SendQuery(q *Query) error

	// SendResponse sends a mDNS response
	SendResponse(resp *Response) error

	// Read returns a channel which receives mDNS messages
	Read(ctx context.Context) <-chan *Request

	// Clears the connection buffer
	Drain(ctx context.Context)

	// Close closes the connection
	Close()
}

type mdnsConn struct {
	ipv4 *ipv4.PacketConn
	ipv6 *ipv6.PacketConn
	ch   chan *Request
}

func NewMDNSConn() (MDNSConn, error) {
	return newMDNSConn()
}

func (c *mdnsConn) SendQuery(q *Query) error {
	return c.sendQuery(q.msg, q.iface)
}

func (c *mdnsConn) SendResponse(resp *Response) error {
	if resp.addr != nil {
		return c.sendResponseTo(resp.msg, resp.iface, resp.addr)
	}

	return c.sendResponse(resp.msg, resp.iface)
}

func (c *mdnsConn) Read(ctx context.Context) <-chan *Request {
	return c.read(ctx)
}

func (c *mdnsConn) Drain(ctx context.Context) {
	log.Debug.Println("Draining connection")
	for {
		select {
		case req := <-c.Read(ctx):
			log.Debug.Println("Ignoring msg from", req.from.IP)
			break
		default:
			return
		}
	}
}

func (c *mdnsConn) Close() {
	c.close()
}

func newMDNSConn() (*mdnsConn, error) {
	var errs []error
	var connIPv4 *ipv4.PacketConn
	var connIPv6 *ipv6.PacketConn

	if conn, err := net.ListenUDP("udp4", AddrIPv4LinkLocalMulticast); err != nil {
		errs = append(errs, err)
	} else {
		connIPv4 = ipv4.NewPacketConn(conn)
		connIPv4.SetControlMessage(ipv4.FlagInterface, true)
		connIPv4.SetMulticastLoopback(true)

		for _, iface := range multicastInterfaces() {
			if err := connIPv4.JoinGroup(&iface, &net.UDPAddr{IP: IPv4LinkLocalMulticast}); err != nil {
				log.Debug.Printf("Failed joining IPv4 %v: %v", iface.Name, err)
			} else {
				log.Debug.Printf("Joined IPv4 %v", iface.Name)
			}
		}
	}

	if conn, err := net.ListenUDP("udp6", AddrIPv6LinkLocalMulticast); err != nil {
		errs = append(errs, err)
	} else {
		connIPv6 = ipv6.NewPacketConn(conn)
		connIPv6.SetControlMessage(ipv6.FlagInterface, true)
		connIPv6.SetMulticastLoopback(true)

		for _, iface := range multicastInterfaces() {
			if err := connIPv6.JoinGroup(&iface, &net.UDPAddr{IP: IPv6LinkLocalMulticast}); err != nil {
				log.Debug.Printf("Failed joining IPv6 %v: %v", iface.Name, err)
			} else {
				log.Debug.Printf("Joined IPv6 %v", iface.Name)
			}
		}
	}

	if err := first(errs...); connIPv4 == nil && connIPv6 == nil {
		return nil, fmt.Errorf("Failed setting up UDP server: %v", err)
	}

	return &mdnsConn{
		ipv4: connIPv4,
		ipv6: connIPv6,
		ch:   make(chan *Request),
	}, nil
}

func (c *mdnsConn) enableMulticastLoopback(enable bool) {
	if c.ipv4 != nil {
		c.ipv4.SetMulticastLoopback(enable)
	}

	if c.ipv6 != nil {
		c.ipv6.SetMulticastLoopback(enable)
	}
}

func (c *mdnsConn) close() {
	if c.ipv4 != nil {
		c.ipv4.Close()
	}

	if c.ipv6 != nil {
		c.ipv6.Close()
	}
}

func (c *mdnsConn) read(ctx context.Context) <-chan *Request {
	c.readInto(ctx, c.ch)
	return c.ch
}

func (c *mdnsConn) readInto(ctx context.Context, ch chan *Request) {

	isDone := func(ctx context.Context) bool {
		return ctx.Err() == context.Canceled
	}

	if c.ipv4 != nil {
		go func() {
			buf := make([]byte, 65536)
			for {
				if isDone(ctx) {
					return
				}

				n, cm, from, err := c.ipv4.ReadFrom(buf)
				if err != nil {
					continue
				}

				udpAddr, ok := from.(*net.UDPAddr)
				if !ok {
					log.Info.Println("invalid source address")
					continue
				}

				var iface *net.Interface
				if cm != nil {
					iface, err = net.InterfaceByIndex(cm.IfIndex)
					if err != nil {
						continue
					}
				}

				if n > 0 {
					m := new(dns.Msg)
					if err := m.Unpack(buf); err == nil && !shouldIgnore(m) {
						ch <- &Request{m, udpAddr, iface}
					}
				}
			}
		}()
	}

	if c.ipv6 != nil {
		go func() {
			buf := make([]byte, 65536)
			for {
				if isDone(ctx) {
					return
				}

				n, cm, from, err := c.ipv6.ReadFrom(buf)
				if err != nil {
					continue
				}

				udpAddr, ok := from.(*net.UDPAddr)
				if !ok {
					log.Info.Println("invalid source address")
					continue
				}

				var iface *net.Interface
				if cm != nil {
					iface, err = net.InterfaceByIndex(cm.IfIndex)
					if err != nil {
						continue
					}
				}

				if n > 0 {
					m := new(dns.Msg)
					if err := m.Unpack(buf); err == nil && !shouldIgnore(m) {
						ch <- &Request{m, udpAddr, iface}
					}
				}
			}
		}()
	}
}

func (c *mdnsConn) sendQuery(m *dns.Msg, iface *net.Interface) error {
	sanitizeQuery(m)

	return c.writeMsg(m, iface)
}

func (c *mdnsConn) sendResponse(m *dns.Msg, iface *net.Interface) error {
	sanitizeResponse(m)

	return c.writeMsg(m, iface)
}

func (c *mdnsConn) sendResponseTo(m *dns.Msg, iface *net.Interface, addr *net.UDPAddr) error {
	sanitizeResponse(m)

	return c.writeMsgTo(m, iface, addr)
}

func (c *mdnsConn) writeMsg(m *dns.Msg, iface *net.Interface) error {
	var err error
	if c.ipv4 != nil {
		err = c.writeMsgTo(m, iface, AddrIPv4LinkLocalMulticast)
	}

	if c.ipv6 != nil {
		err = c.writeMsgTo(m, iface, AddrIPv6LinkLocalMulticast)
	}

	return err
}

func (c *mdnsConn) writeMsgTo(m *dns.Msg, iface *net.Interface, addr *net.UDPAddr) error {
	sanitizeMsg(m)

	if c.ipv4 != nil && addr.IP.To4() != nil {
		if out, err := m.Pack(); err == nil {
			var ctrl *ipv4.ControlMessage
			if iface != nil {
				ctrl = &ipv4.ControlMessage{
					IfIndex: iface.Index,
				}
			}
			if _, err = c.ipv4.WriteTo(out, ctrl, addr); err != nil {
				return err
			}
		}
	}

	if c.ipv6 != nil && addr.IP.To4() == nil {
		if out, err := m.Pack(); err == nil {
			var ctrl *ipv6.ControlMessage
			if iface != nil {
				ctrl = &ipv6.ControlMessage{
					IfIndex: iface.Index,
				}
			}
			if _, err = c.ipv6.WriteTo(out, ctrl, addr); err != nil {
				return err
			}
		}
	}

	return nil
}

func shouldIgnore(m *dns.Msg) bool {
	if m.Opcode != 0 {
		return true
	}

	if m.Rcode != 0 {
		return true
	}

	return false
}

func sanitizeResponse(m *dns.Msg) {
	if m.Question != nil && len(m.Question) > 0 {
		log.Info.Println("Multicast DNS responses MUST NOT contain any questions in the Question Section.  (RFC6762 6)")
		m.Question = nil
	}

	if !m.Response {
		log.Info.Println("In response messages the QR bit MUST be one (RFC6762 18.2)")
		m.Response = true
	}

	if !m.Authoritative {
		log.Info.Println("AA Bit bit MUST be set to one in response messages (RFC6762 18.4)")
		m.Authoritative = true
	}

	if m.Truncated {
		log.Info.Println("In multicast response messages, the TC bit MUST be zero on transmission. (RFC6762 18.5)")
		m.Truncated = false
	}
}

func sanitizeQuery(m *dns.Msg) {
	if m.Response {
		log.Info.Println("In query messages the QR bit MUST be zero (RFC6762 18.2)")
		m.Response = false
	}

	if m.Authoritative {
		log.Info.Println("AA Bit MUST be zero in query messages (RFC6762 18.4)")
		m.Authoritative = false
	}
}

func sanitizeMsg(m *dns.Msg) {
	if m.Opcode != 0 {
		log.Info.Println("In both multicast query and multicast response messages, the OPCODE MUST be zero on transmission (RFC6762 18.3)")
		m.Opcode = 0
	}

	if m.RecursionDesired {
		log.Info.Println("In both multicast query and multicast response messages, the Recursion Available bit MUST be zero on transmission. (RFC6762 18.7)")
		m.RecursionDesired = false
	}

	if m.Zero {
		log.Info.Println("In both query and response messages, the Zero bit MUST be zero on transmission (RFC6762 18.8)")
		m.Zero = false
	}

	if m.AuthenticatedData {
		log.Info.Println("In both multicast query and multicast response messages, the Authentic Data bit MUST be zero on transmission (RFC6762 18.9)")
		m.AuthenticatedData = false
	}

	if m.CheckingDisabled {
		log.Info.Println("In both multicast query and multicast response messages, the Checking Disabled bit MUST be zero on transmission (RFC6762 18.10)")
		m.CheckingDisabled = false
	}

	if m.Rcode != 0 {
		log.Info.Println("In both multicast query and multicast response messages, the Response Code MUST be zero on transmission. (RFC6762 18.11)")
		m.Rcode = 0
	}
}

func first(errs ...error) error {
	for _, err := range errs {
		if err != nil {
			return err
		}
	}

	return nil
}

// Sets the Top Bit of rrclass for all answer records (except PTR) to trigger a cache flush in the receivers.
func setAnswerCacheFlushBit(msg *dns.Msg) {
	// From RFC6762
	//    The most significant bit of the rrclass for a record in the Answer
	//    Section of a response message is the Multicast DNS cache-flush bit
	//    and is discussed in more detail below in Section 10.2, "Announcements
	//    to Flush Outdated Cache Entries".
	for _, a := range msg.Answer {
		switch a.(type) {
		case *dns.PTR:
			continue
		default:
			a.Header().Class |= (1 << 15)
		}
	}
}

// Sets the Top Bit of class to indicate the unicast responses are preferred for this question.
func setQuestionUnicast(q *dns.Question) {
	q.Qclass |= (1 << 15)
}

// Returns true if q requires unicast responses.
func isUnicastQuestion(q dns.Question) bool {
	// From RFC6762
	// 18.12.  Repurposing of Top Bit of qclass in Question Section
	//
	//    In the Question Section of a Multicast DNS query, the top bit of the
	//    qclass field is used to indicate that unicast responses are preferred
	//    for this particular question.  (See Section 5.4.)
	return q.Qclass&(1<<15) != 0
}
<|endoftext|>


<commit_before>package slack

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

// HTTPRequester defines the minimal interface needed for an http.Client to be implemented.
//
// Use it in conjunction with the SetHTTPClient function to allow for other capabilities
// like a tracing http.Client
type HTTPRequester interface {
	Do(*http.Request) (*http.Response, error)
}

var customHTTPClient HTTPRequester

// HTTPClient sets a custom http.Client
// deprecated: in favor of SetHTTPClient()
var HTTPClient = &http.Client{}

type WebResponse struct {
	Ok    bool      `json:"ok"`
	Error *WebError `json:"error"`
}

type WebError string

func (s WebError) Error() string {
	return string(s)
}

type RateLimitedError struct {
	RetryAfter time.Duration
}

func (e *RateLimitedError) Error() string {
	return fmt.Sprintf("Slack rate limit exceeded, retry after %s", e.RetryAfter)
}

func fileUploadReq(ctx context.Context, path, fieldname, filename string, values url.Values, r io.Reader) (*http.Request, error) {
	body := &bytes.Buffer{}
	wr := multipart.NewWriter(body)

	ioWriter, err := wr.CreateFormFile(fieldname, filename)
	if err != nil {
		wr.Close()
		return nil, err
	}
	_, err = io.Copy(ioWriter, r)
	if err != nil {
		wr.Close()
		return nil, err
	}
	// Close the multipart writer or the footer won't be written
	wr.Close()
	req, err := http.NewRequest("POST", path, body)
	req = req.WithContext(ctx)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Content-Type", wr.FormDataContentType())
	req.URL.RawQuery = (values).Encode()
	return req, nil
}

func parseResponseBody(body io.ReadCloser, intf *interface{}, debug bool) error {
	response, err := ioutil.ReadAll(body)
	if err != nil {
		return err
	}

	// FIXME: will be api.Debugf
	if debug {
		logger.Printf("parseResponseBody: %s\n", string(response))
	}

	return json.Unmarshal(response, &intf)
}

func postLocalWithMultipartResponse(ctx context.Context, path, fpath, fieldname string, values url.Values, intf interface{}, debug bool) error {
	fullpath, err := filepath.Abs(fpath)
	if err != nil {
		return err
	}
	file, err := os.Open(fullpath)
	if err != nil {
		return err
	}
	defer file.Close()
	return postWithMultipartResponse(ctx, path, filepath.Base(fpath), fieldname, values, file, intf, debug)
}

func postWithMultipartResponse(ctx context.Context, path, name, fieldname string, values url.Values, r io.Reader, intf interface{}, debug bool) error {
	req, err := fileUploadReq(ctx, SLACK_API+path, fieldname, name, values, r)
	if err != nil {
		return err
	}
	req = req.WithContext(ctx)
	resp, err := getHTTPClient().Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusTooManyRequests {
		retry, err := strconv.ParseInt(resp.Header.Get("Retry-After"), 10, 64)
		if err != nil {
			return err
		}
		return &RateLimitedError{time.Duration(retry) * time.Second}
	}

	// Slack seems to send an HTML body along with 5xx error codes. Don't parse it.
	if resp.StatusCode != 200 {
		logResponse(resp, debug)
		return fmt.Errorf("Slack server error: %s.", resp.Status)
	}

	return parseResponseBody(resp.Body, &intf, debug)
}

func postForm(ctx context.Context, endpoint string, values url.Values, intf interface{}, debug bool) error {
	reqBody := strings.NewReader(values.Encode())
	req, err := http.NewRequest("POST", endpoint, reqBody)
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	req = req.WithContext(ctx)
	resp, err := getHTTPClient().Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusTooManyRequests {
		retry, err := strconv.ParseInt(resp.Header.Get("Retry-After"), 10, 64)
		if err != nil {
			return err
		}
		return &RateLimitedError{time.Duration(retry) * time.Second}
	}

	// Slack seems to send an HTML body along with 5xx error codes. Don't parse it.
	if resp.StatusCode != 200 {
		logResponse(resp, debug)
		return fmt.Errorf("Slack server error: %s.", resp.Status)
	}

	return parseResponseBody(resp.Body, &intf, debug)
}

func post(ctx context.Context, path string, values url.Values, intf interface{}, debug bool) error {
	return postForm(ctx, SLACK_API+path, values, intf, debug)
}

func parseAdminResponse(ctx context.Context, method string, teamName string, values url.Values, intf interface{}, debug bool) error {
	endpoint := fmt.Sprintf(SLACK_WEB_API_FORMAT, teamName, method, time.Now().Unix())
	return postForm(ctx, endpoint, values, intf, debug)
}

func logResponse(resp *http.Response, debug bool) error {
	if debug {
		text, err := httputil.DumpResponse(resp, true)
		if err != nil {
			return err
		}

		logger.Print(string(text))
	}

	return nil
}

func getHTTPClient() HTTPRequester {
	if customHTTPClient != nil {
		return customHTTPClient
	}

	return HTTPClient
}

// SetHTTPClient allows you to specify a custom http.Client
// Use this instead of the package level HTTPClient variable if you want to use a custom client like the
// Stackdriver Trace HTTPClient https://godoc.org/cloud.google.com/go/trace#HTTPClient
func SetHTTPClient(client HTTPRequester) {
	customHTTPClient = client
}
<commit_msg>Use net/http constants for 200 status<commit_after>package slack

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

// HTTPRequester defines the minimal interface needed for an http.Client to be implemented.
//
// Use it in conjunction with the SetHTTPClient function to allow for other capabilities
// like a tracing http.Client
type HTTPRequester interface {
	Do(*http.Request) (*http.Response, error)
}

var customHTTPClient HTTPRequester

// HTTPClient sets a custom http.Client
// deprecated: in favor of SetHTTPClient()
var HTTPClient = &http.Client{}

type WebResponse struct {
	Ok    bool      `json:"ok"`
	Error *WebError `json:"error"`
}

type WebError string

func (s WebError) Error() string {
	return string(s)
}

type RateLimitedError struct {
	RetryAfter time.Duration
}

func (e *RateLimitedError) Error() string {
	return fmt.Sprintf("Slack rate limit exceeded, retry after %s", e.RetryAfter)
}

func fileUploadReq(ctx context.Context, path, fieldname, filename string, values url.Values, r io.Reader) (*http.Request, error) {
	body := &bytes.Buffer{}
	wr := multipart.NewWriter(body)

	ioWriter, err := wr.CreateFormFile(fieldname, filename)
	if err != nil {
		wr.Close()
		return nil, err
	}
	_, err = io.Copy(ioWriter, r)
	if err != nil {
		wr.Close()
		return nil, err
	}
	// Close the multipart writer or the footer won't be written
	wr.Close()
	req, err := http.NewRequest("POST", path, body)
	req = req.WithContext(ctx)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Content-Type", wr.FormDataContentType())
	req.URL.RawQuery = (values).Encode()
	return req, nil
}

func parseResponseBody(body io.ReadCloser, intf *interface{}, debug bool) error {
	response, err := ioutil.ReadAll(body)
	if err != nil {
		return err
	}

	// FIXME: will be api.Debugf
	if debug {
		logger.Printf("parseResponseBody: %s\n", string(response))
	}

	return json.Unmarshal(response, &intf)
}

func postLocalWithMultipartResponse(ctx context.Context, path, fpath, fieldname string, values url.Values, intf interface{}, debug bool) error {
	fullpath, err := filepath.Abs(fpath)
	if err != nil {
		return err
	}
	file, err := os.Open(fullpath)
	if err != nil {
		return err
	}
	defer file.Close()
	return postWithMultipartResponse(ctx, path, filepath.Base(fpath), fieldname, values, file, intf, debug)
}

func postWithMultipartResponse(ctx context.Context, path, name, fieldname string, values url.Values, r io.Reader, intf interface{}, debug bool) error {
	req, err := fileUploadReq(ctx, SLACK_API+path, fieldname, name, values, r)
	if err != nil {
		return err
	}
	req = req.WithContext(ctx)
	resp, err := getHTTPClient().Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusTooManyRequests {
		retry, err := strconv.ParseInt(resp.Header.Get("Retry-After"), 10, 64)
		if err != nil {
			return err
		}
		return &RateLimitedError{time.Duration(retry) * time.Second}
	}

	// Slack seems to send an HTML body along with 5xx error codes. Don't parse it.
	if resp.StatusCode != http.StatusOK {
		logResponse(resp, debug)
		return fmt.Errorf("Slack server error: %s.", resp.Status)
	}

	return parseResponseBody(resp.Body, &intf, debug)
}

func postForm(ctx context.Context, endpoint string, values url.Values, intf interface{}, debug bool) error {
	reqBody := strings.NewReader(values.Encode())
	req, err := http.NewRequest("POST", endpoint, reqBody)
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	req = req.WithContext(ctx)
	resp, err := getHTTPClient().Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusTooManyRequests {
		retry, err := strconv.ParseInt(resp.Header.Get("Retry-After"), 10, 64)
		if err != nil {
			return err
		}
		return &RateLimitedError{time.Duration(retry) * time.Second}
	}

	// Slack seems to send an HTML body along with 5xx error codes. Don't parse it.
	if resp.StatusCode != http.StatusOK {
		logResponse(resp, debug)
		return fmt.Errorf("Slack server error: %s.", resp.Status)
	}

	return parseResponseBody(resp.Body, &intf, debug)
}

func post(ctx context.Context, path string, values url.Values, intf interface{}, debug bool) error {
	return postForm(ctx, SLACK_API+path, values, intf, debug)
}

func parseAdminResponse(ctx context.Context, method string, teamName string, values url.Values, intf interface{}, debug bool) error {
	endpoint := fmt.Sprintf(SLACK_WEB_API_FORMAT, teamName, method, time.Now().Unix())
	return postForm(ctx, endpoint, values, intf, debug)
}

func logResponse(resp *http.Response, debug bool) error {
	if debug {
		text, err := httputil.DumpResponse(resp, true)
		if err != nil {
			return err
		}

		logger.Print(string(text))
	}

	return nil
}

func getHTTPClient() HTTPRequester {
	if customHTTPClient != nil {
		return customHTTPClient
	}

	return HTTPClient
}

// SetHTTPClient allows you to specify a custom http.Client
// Use this instead of the package level HTTPClient variable if you want to use a custom client like the
// Stackdriver Trace HTTPClient https://godoc.org/cloud.google.com/go/trace#HTTPClient
func SetHTTPClient(client HTTPRequester) {
	customHTTPClient = client
}
<|endoftext|>


<commit_before>/*******************************************************************************
The MIT License (MIT)

Copyright (c) 2015 Hajime Nakagami

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*******************************************************************************/

package toybroker

import (
	"fmt"
	"net"
)

func MqttMainLoop(conn net.Conn, topics Topics, hook Hook) {
	connect, _, _, _, _, remaining, err := readMessage(conn)
	if connect != CONNECT || err != nil {
		conn.Write([]byte{CONNACK * 16, 2, 0, CONNACK_Rejected})
		return
	}
	clientID, _, _, loginName, loginPassword, err := unpackCONNECT(remaining)
	status := hook.Login(conn, clientID, loginName, loginPassword)
	if status != CONNACK_Success {
		conn.Write([]byte{CONNACK * 16, 2, 0, CONNACK_Rejected})
		return
	}

	oldClient := hook.GetClient(clientID)
	client := NewClient(clientID, loginName, conn, oldClient)
	hook.SetClient(client)
	client.Send(packCONNACK(status))

	for {
		command, _, header_qos, retain, _, remaining, err := readMessage(conn)
		if err != nil {
			hook.Logout(clientID)
			break
		}
		switch command {
		case PUBLISH:
			topic, messageID, payload, err := unpackPUBLISH(remaining)
			debugOutput(fmt.Sprintf("PUBLISH:%s,%d,%v,%v", topic, messageID, payload, err))

			clientList, qosList := topics.List(topic)
			for i, clientID := range clientList {
				target := hook.GetClient(clientID)
				if qosList[i] > 0 && header_qos > 0 {
					qosList[i] = 1
				}
				target.Publish(false, qosList[i], topic, payload)
				if retain {
					topics.AddRetainMessage(topic, payload)
				}
			}
		case PUBACK:
			debugOutput("PUBACK")
		case PUBREL:
			debugOutput("PUBREL")
		case SUBSCRIBE:
			messageID, subscribe_topics, err := unpackSUBSCRIBE(remaining)
			debugOutput(fmt.Sprintf("SUBSCRIBE:%v,%d,%v", subscribe_topics, messageID, err))
			qos := make([]byte, len(subscribe_topics))
			for i, topicName := range subscribe_topics {
				qos[i] = byte(hook.Subscribe(topics, topicName, clientID, 0))
			}
			client.Send(packSUBACK(messageID, qos))

			// Publish retain messages
			for i, topicName := range subscribe_topics {
				payload := topics.GetRetainMessage(topicName)
				if len(payload) != 0 {
					client.Publish(false, int(qos[i]), topicName, payload)
				}
			}

		case UNSUBSCRIBE:
			messageID, unsubscribe_topics, err := unpackUNSUBSCRIBE(remaining)
			debugOutput(fmt.Sprintf("UNSUBSCRIBE:%v,%d,%v", unsubscribe_topics, messageID, err))
			for _, topicName := range unsubscribe_topics {
				hook.Unsubscribe(topics, topicName, clientID)
			}
			client.Send(packUNSUBACK(messageID))
		case PINGREQ:
			debugOutput("PINGREQ")
			client.Send(packPINGRESP())
		case DISCONNECT:
			debugOutput("DISCONNECT")
			hook.Logout(clientID)
			break
		default:
			debugOutput("Invalid Command")
			hook.Logout(clientID)
			break
		}
	}
}
<commit_msg>add outputDebug()<commit_after>/*******************************************************************************
The MIT License (MIT)

Copyright (c) 2015 Hajime Nakagami

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*******************************************************************************/

package toybroker

import (
	"fmt"
	"net"
)

func MqttMainLoop(conn net.Conn, topics Topics, hook Hook) {
	connect, _, _, _, _, remaining, err := readMessage(conn)
	if connect != CONNECT || err != nil {
		conn.Write([]byte{CONNACK * 16, 2, 0, CONNACK_Rejected})
		return
	}
	clientID, _, _, loginName, loginPassword, err := unpackCONNECT(remaining)
	status := hook.Login(conn, clientID, loginName, loginPassword)
	if status != CONNACK_Success {
		conn.Write([]byte{CONNACK * 16, 2, 0, CONNACK_Rejected})
		return
	}

	oldClient := hook.GetClient(clientID)
	client := NewClient(clientID, loginName, conn, oldClient)
	hook.SetClient(client)
	client.Send(packCONNACK(status))

	for {
		command, _, header_qos, retain, _, remaining, err := readMessage(conn)
		if err != nil {
			hook.Logout(clientID)
			break
		}
		switch command {
		case PUBLISH:
			topic, messageID, payload, err := unpackPUBLISH(remaining)
			debugOutput(fmt.Sprintf("PUBLISH:%s,%d,%v,%v", topic, messageID, payload, err))

			clientList, qosList := topics.List(topic)
			for i, clientID := range clientList {
				target := hook.GetClient(clientID)
				if qosList[i] > 0 && header_qos > 0 {
					qosList[i] = 1
				}
				target.Publish(false, qosList[i], topic, payload)
				if retain {
					debugOutput("AddRetainMessage()")
					topics.AddRetainMessage(topic, payload)
				}
			}
		case PUBACK:
			debugOutput("PUBACK")
		case PUBREL:
			debugOutput("PUBREL")
		case SUBSCRIBE:
			messageID, subscribe_topics, err := unpackSUBSCRIBE(remaining)
			debugOutput(fmt.Sprintf("SUBSCRIBE:%v,%d,%v", subscribe_topics, messageID, err))
			qos := make([]byte, len(subscribe_topics))
			for i, topicName := range subscribe_topics {
				qos[i] = byte(hook.Subscribe(topics, topicName, clientID, 0))
			}
			client.Send(packSUBACK(messageID, qos))

			// Publish retain messages
			for i, topicName := range subscribe_topics {
				payload := topics.GetRetainMessage(topicName)
				if len(payload) != 0 {
					debugOutput(fmt.Sprintf("SendRetainMessage:%s->%v", topicName, payload))
					client.Publish(false, int(qos[i]), topicName, payload)
				}
			}

		case UNSUBSCRIBE:
			messageID, unsubscribe_topics, err := unpackUNSUBSCRIBE(remaining)
			debugOutput(fmt.Sprintf("UNSUBSCRIBE:%v,%d,%v", unsubscribe_topics, messageID, err))
			for _, topicName := range unsubscribe_topics {
				hook.Unsubscribe(topics, topicName, clientID)
			}
			client.Send(packUNSUBACK(messageID))
		case PINGREQ:
			debugOutput("PINGREQ")
			client.Send(packPINGRESP())
		case DISCONNECT:
			debugOutput("DISCONNECT")
			hook.Logout(clientID)
			break
		default:
			debugOutput("Invalid Command")
			hook.Logout(clientID)
			break
		}
	}
}
<|endoftext|>


<commit_before>package transforms

import "github.com/go-audio/audio"

// MonoDownmix converts the buffer to a mono buffer
// by downmixing the channels together.
func MonoDownmix(buf *audio.FloatBuffer) error {
	if buf == nil || buf.Format == nil {
		return audio.ErrInvalidBuffer
	}
	if buf.Format == nil {
		return audio.ErrInvalidBuffer
	}
	nChans := buf.Format.NumChannels
	if nChans < 2 {
		return nil
	}
	nChansF := float64(nChans)

	frameCount := len(buf.Data) / buf.Format.NumChannels
	newData := make([]float64, frameCount)
	for i := 0; i < frameCount; i++ {
		newData[i] = 0
		for j := 0; j < nChans; j++ {
			newData[i] += buf.Data[i*nChans+j]
		}
		newData[i] /= nChansF
	}
	buf.Data = newData
	buf.Format.NumChannels = 1

	return nil
}
<commit_msg>properly use the buffer API<commit_after>package transforms

import "github.com/go-audio/audio"

// MonoDownmix converts the buffer to a mono buffer
// by downmixing the channels together.
func MonoDownmix(buf *audio.FloatBuffer) error {
	if buf == nil || buf.Format == nil {
		return audio.ErrInvalidBuffer
	}
	if buf.Format == nil {
		return audio.ErrInvalidBuffer
	}
	nChans := buf.Format.NumChannels
	if nChans < 2 {
		return nil
	}
	nChansF := float64(nChans)

	frameCount := buf.NumFrames()
	newData := make([]float64, frameCount)
	for i := 0; i < frameCount; i++ {
		newData[i] = 0
		for j := 0; j < nChans; j++ {
			newData[i] += buf.Data[i*nChans+j]
		}
		newData[i] /= nChansF
	}
	buf.Data = newData
	buf.Format.NumChannels = 1

	return nil
}
<|endoftext|>


<commit_before>// Implementation of MQTT V3.1.
// See http://public.dhe.ibm.com/software/dw/webservices/ws-mqtt/mqtt-v3r1.html
package mqtt

import (
	"bytes"
	"errors"
	"io"
)

var (
	badMsgTypeError        = errors.New("mqtt: message type is invalid")
	badQosError            = errors.New("mqtt: QoS is invalid")
	badWillQosError        = errors.New("mqtt: will QoS is invalid")
	badLengthEncodingError = errors.New("mqtt: remaining length field exceeded maximum of 4 bytes")
	badReturnCodeError     = errors.New("mqtt: is invalid")
	dataExceedsPacketError = errors.New("mqtt: data exceeds packet length")
	msgTooLongError        = errors.New("mqtt: message is too long")
)

const (
	QosAtMostOnce = QosLevel(iota)
	QosAtLeastOnce
	QosExactlyOnce

	qosFirstInvalid
)

type QosLevel uint8

func (qos QosLevel) IsValid() bool {
	return qos < qosFirstInvalid
}

func (qos QosLevel) HasId() bool {
	return qos == QosAtLeastOnce || qos == QosExactlyOnce
}

const (
	RetCodeAccepted = ReturnCode(iota)
	RetCodeUnacceptableProtocolVersion
	RetCodeIdentifierRejected
	RetCodeServerUnavailable
	RetCodeBadUsernameOrPassword
	RetCodeNotAuthorized

	retCodeFirstInvalid
)

type ReturnCode uint8

func (rc ReturnCode) IsValid() bool {
	return rc >= RetCodeAccepted && rc < retCodeFirstInvalid
}

func getUint8(r io.Reader, packetRemaining *int32) uint8 {
	if *packetRemaining < 1 {
		raiseError(dataExceedsPacketError)
	}

	var b [1]byte
	if _, err := io.ReadFull(r, b[:]); err != nil {
		raiseError(err)
	}
	*packetRemaining--

	return b[0]
}

func getUint16(r io.Reader, packetRemaining *int32) uint16 {
	if *packetRemaining < 2 {
		raiseError(dataExceedsPacketError)
	}

	var b [2]byte
	if _, err := io.ReadFull(r, b[:]); err != nil {
		raiseError(err)
	}
	*packetRemaining -= 2

	return uint16(b[0]<<8) + uint16(b[1])
}

func getString(r io.Reader, packetRemaining *int32) string {
	strLen := int(getUint16(r, packetRemaining))

	if int(*packetRemaining) < strLen {
		raiseError(dataExceedsPacketError)
	}

	b := make([]byte, strLen)
	if _, err := io.ReadFull(r, b); err != nil {
		raiseError(err)
	}
	*packetRemaining -= int32(strLen)

	return string(b)
}

func DecodeOneMessage(r io.Reader) (msg Message, err error) {
	var hdr Header
	var msgType MessageType
	var packetRemaining int32
	msgType, packetRemaining, err = hdr.Decode(r)

	msg, err = NewMessage(msgType)

	return msg, msg.Decode(r, hdr, packetRemaining)
}

func NewMessage(msgType MessageType) (msg Message, err error) {
	switch msgType {
	case MsgConnect:
		msg = new(Connect)
	case MsgConnAck:
		msg = new(ConnAck)
	case MsgPublish:
		msg = new(Publish)
	case MsgPubAck:
		msg = new(PubAck)
	case MsgPubRec:
		msg = new(PubRec)
	case MsgPubRel:
		msg = new(PubRel)
	case MsgPubComp:
		msg = new(PubComp)
	case MsgSubscribe:
		msg = new(Subscribe)
	case MsgUnsubAck:
		msg = new(UnsubAck)
	case MsgSubAck:
		msg = new(SubAck)
	case MsgUnsubscribe:
		msg = new(Unsubscribe)
	default:
		return nil, badMsgTypeError
	}

	return
}

func setUint8(val uint8, buf *bytes.Buffer) {
	buf.WriteByte(byte(val))
}

func setUint16(val uint16, buf *bytes.Buffer) {
	buf.WriteByte(byte(val & 0xff00 >> 8))
	buf.WriteByte(byte(val & 0x00ff))
}

func setString(val string, buf *bytes.Buffer) {
	length := uint16(len(val))
	setUint16(length, buf)
	buf.WriteString(val)
}

func boolToByte(val bool) byte {
	if val {
		return byte(1)
	}
	return byte(0)
}

func decodeLength(r io.Reader) int32 {
	var v int32
	var buf [1]byte
	var shift uint
	for i := 0; i < 4; i++ {
		if _, err := io.ReadFull(r, buf[:]); err != nil {
			raiseError(err)
		}

		b := buf[0]
		v |= int32(b&0x7f) << shift

		if b&0x80 == 0 {
			return v
		}
		shift += 7
	}

	raiseError(badLengthEncodingError)
	panic("unreachable")
}

func encodeLength(length int32, buf *bytes.Buffer) {
	if length == 0 {
		buf.WriteByte(byte(0))
		return
	}
	var lbuf bytes.Buffer
	for length > 0 {
		digit := length % 128
		length = length / 128
		if length > 0 {
			digit = digit | 0x80
		}
		lbuf.WriteByte(byte(digit))
	}
	blen := lbuf.Bytes()
	for i := 1; i <= len(blen); i += 1 {
		buf.WriteByte(blen[len(blen)-i])
	}
}

// panicErr wraps an error that caused a problem that needs to bail out of the
// API, such that errors can be recovered and returned as errors from the
// public API.
type panicErr struct {
	err error
}

func (p panicErr) Error() string {
	return p.err.Error()
}

func raiseError(err error) {
	panic(panicErr{err})
}

// recoverError recovers any panic in flight and, iff it's an error from
// raiseError, will return the error. Otherwise re-raises the panic value.
// If no panic is in flight, it returns existingErr.
//
// This must be used in combination with a defer in all public API entry
// points where raiseError could be called.
func recoverError(existingErr error) error {
	if p := recover(); p != nil {
		if pErr, ok := p.(panicErr); ok {
			return pErr.err
		} else {
			panic(p)
		}
	}
	return existingErr
}
<commit_msg>Document DecodeOneMessage and NewMessage.<commit_after>// Implementation of MQTT V3.1.
// See http://public.dhe.ibm.com/software/dw/webservices/ws-mqtt/mqtt-v3r1.html
package mqtt

import (
	"bytes"
	"errors"
	"io"
)

var (
	badMsgTypeError        = errors.New("mqtt: message type is invalid")
	badQosError            = errors.New("mqtt: QoS is invalid")
	badWillQosError        = errors.New("mqtt: will QoS is invalid")
	badLengthEncodingError = errors.New("mqtt: remaining length field exceeded maximum of 4 bytes")
	badReturnCodeError     = errors.New("mqtt: is invalid")
	dataExceedsPacketError = errors.New("mqtt: data exceeds packet length")
	msgTooLongError        = errors.New("mqtt: message is too long")
)

const (
	QosAtMostOnce = QosLevel(iota)
	QosAtLeastOnce
	QosExactlyOnce

	qosFirstInvalid
)

type QosLevel uint8

func (qos QosLevel) IsValid() bool {
	return qos < qosFirstInvalid
}

func (qos QosLevel) HasId() bool {
	return qos == QosAtLeastOnce || qos == QosExactlyOnce
}

const (
	RetCodeAccepted = ReturnCode(iota)
	RetCodeUnacceptableProtocolVersion
	RetCodeIdentifierRejected
	RetCodeServerUnavailable
	RetCodeBadUsernameOrPassword
	RetCodeNotAuthorized

	retCodeFirstInvalid
)

type ReturnCode uint8

func (rc ReturnCode) IsValid() bool {
	return rc >= RetCodeAccepted && rc < retCodeFirstInvalid
}

func getUint8(r io.Reader, packetRemaining *int32) uint8 {
	if *packetRemaining < 1 {
		raiseError(dataExceedsPacketError)
	}

	var b [1]byte
	if _, err := io.ReadFull(r, b[:]); err != nil {
		raiseError(err)
	}
	*packetRemaining--

	return b[0]
}

func getUint16(r io.Reader, packetRemaining *int32) uint16 {
	if *packetRemaining < 2 {
		raiseError(dataExceedsPacketError)
	}

	var b [2]byte
	if _, err := io.ReadFull(r, b[:]); err != nil {
		raiseError(err)
	}
	*packetRemaining -= 2

	return uint16(b[0]<<8) + uint16(b[1])
}

func getString(r io.Reader, packetRemaining *int32) string {
	strLen := int(getUint16(r, packetRemaining))

	if int(*packetRemaining) < strLen {
		raiseError(dataExceedsPacketError)
	}

	b := make([]byte, strLen)
	if _, err := io.ReadFull(r, b); err != nil {
		raiseError(err)
	}
	*packetRemaining -= int32(strLen)

	return string(b)
}

// DecodeOneMessage decodes one message from r.
func DecodeOneMessage(r io.Reader) (msg Message, err error) {
	var hdr Header
	var msgType MessageType
	var packetRemaining int32
	msgType, packetRemaining, err = hdr.Decode(r)

	msg, err = NewMessage(msgType)

	return msg, msg.Decode(r, hdr, packetRemaining)
}

// NewMessage creates an instance of a Message value for the given message
// type. An error is returned if msgType is invalid.
func NewMessage(msgType MessageType) (msg Message, err error) {
	switch msgType {
	case MsgConnect:
		msg = new(Connect)
	case MsgConnAck:
		msg = new(ConnAck)
	case MsgPublish:
		msg = new(Publish)
	case MsgPubAck:
		msg = new(PubAck)
	case MsgPubRec:
		msg = new(PubRec)
	case MsgPubRel:
		msg = new(PubRel)
	case MsgPubComp:
		msg = new(PubComp)
	case MsgSubscribe:
		msg = new(Subscribe)
	case MsgUnsubAck:
		msg = new(UnsubAck)
	case MsgSubAck:
		msg = new(SubAck)
	case MsgUnsubscribe:
		msg = new(Unsubscribe)
	default:
		return nil, badMsgTypeError
	}

	return
}

func setUint8(val uint8, buf *bytes.Buffer) {
	buf.WriteByte(byte(val))
}

func setUint16(val uint16, buf *bytes.Buffer) {
	buf.WriteByte(byte(val & 0xff00 >> 8))
	buf.WriteByte(byte(val & 0x00ff))
}

func setString(val string, buf *bytes.Buffer) {
	length := uint16(len(val))
	setUint16(length, buf)
	buf.WriteString(val)
}

func boolToByte(val bool) byte {
	if val {
		return byte(1)
	}
	return byte(0)
}

func decodeLength(r io.Reader) int32 {
	var v int32
	var buf [1]byte
	var shift uint
	for i := 0; i < 4; i++ {
		if _, err := io.ReadFull(r, buf[:]); err != nil {
			raiseError(err)
		}

		b := buf[0]
		v |= int32(b&0x7f) << shift

		if b&0x80 == 0 {
			return v
		}
		shift += 7
	}

	raiseError(badLengthEncodingError)
	panic("unreachable")
}

func encodeLength(length int32, buf *bytes.Buffer) {
	if length == 0 {
		buf.WriteByte(byte(0))
		return
	}
	var lbuf bytes.Buffer
	for length > 0 {
		digit := length % 128
		length = length / 128
		if length > 0 {
			digit = digit | 0x80
		}
		lbuf.WriteByte(byte(digit))
	}
	blen := lbuf.Bytes()
	for i := 1; i <= len(blen); i += 1 {
		buf.WriteByte(blen[len(blen)-i])
	}
}

// panicErr wraps an error that caused a problem that needs to bail out of the
// API, such that errors can be recovered and returned as errors from the
// public API.
type panicErr struct {
	err error
}

func (p panicErr) Error() string {
	return p.err.Error()
}

func raiseError(err error) {
	panic(panicErr{err})
}

// recoverError recovers any panic in flight and, iff it's an error from
// raiseError, will return the error. Otherwise re-raises the panic value.
// If no panic is in flight, it returns existingErr.
//
// This must be used in combination with a defer in all public API entry
// points where raiseError could be called.
func recoverError(existingErr error) error {
	if p := recover(); p != nil {
		if pErr, ok := p.(panicErr); ok {
			return pErr.err
		} else {
			panic(p)
		}
	}
	return existingErr
}
<|endoftext|>


<commit_before>// Package router provides a fast HTTP router.
package router

import (
	"fmt"
	"net/http"
	"sort"
	"strings"
)

type node struct {
	s        string
	params   []string // Parameter's names from the parent node to this one.
	children nodes
	handler  http.Handler
}

func (n *node) string(level int) (s string) {
	s += fmt.Sprintf("%s%q  %v  %v\n", strings.Repeat("\t", level), n.s, n.params, n.handler)
	for _, n := range n.children {
		s += n.string(level + 1)
	}
	return
}

func (n *node) isWildcard() bool {
	return n.s[len(n.s)-1] == '/' && len(n.children) == 0
}

func (n *node) countChildren() (i int) {
	for _, n := range n.children {
		i++
		i += n.countChildren()
	}
	return
}

type nodes []*node

// makeChild adds a node to the tree.
func (nn *nodes) makeChild(path string, params []string, handler http.Handler) {
NodesLoop:
	for _, n := range *nn {
		minlen := len(n.s)
		if len(path) < minlen {
			minlen = len(path)
		}
		for i := 0; i < minlen; i++ {
			if n.s[i] == path[i] {
				continue
			}
			if i == 0 { // No match from the first byte: see next same-level node.
				continue NodesLoop
			}
			// Difference in the middle of a node: split current node to make subnode and transfer handler to it.
			*n = node{
				s: n.s[:i],
				children: nodes{
					{s: n.s[i:], params: n.params, children: n.children, handler: n.handler},
					{s: path[i:], params: params, handler: handler},
				},
			}
			return
		}
		if len(path) < len(n.s) { // s fully matched first part of n.s: split node.
			*n = node{
				s:      n.s[:len(path)],
				params: params,
				children: nodes{
					{s: n.s[len(path):], params: n.params, children: n.children, handler: n.handler},
				},
				handler: handler,
			}
		} else if len(path) > len(n.s) { // n.s fully matched first part of s: see subnodes for the rest.
			n.children.makeChild(path[len(n.s):], params, handler)
		} else { // s == n.s and no rest: node has no handler or route is duplicated.
			if handler == nil { // No handler provided (must be a non-ending path parameter): don't overwrite.
				return
			}
			if n.handler != nil { // Handler provided but n.handler already set: route is duplicated.
				panic(fmt.Errorf("router: two or more routes have same path"))
			}
			n.params = params
			n.handler = handler
		}
		return
	}
	*nn = append(*nn, &node{s: path, params: params, handler: handler}) // Not a single byte match on same-level nodes: append a new one.
}

func (nn *nodes) findChild(path string, params *[]string) *node {
	for _, n := range *nn {
		if n.s == ":" { // Handle parameter node.
			paramEnd := strings.IndexByte(path, '/')
			if paramEnd == -1 { // Path ends with the parameter.
				if n.handler != nil { // Performance: append parameter only if the node has a handler (otherwise useless).
					*params = append(*params, path)
				}
				return n
			}
			*params = append(*params, path[:paramEnd])
			return n.children.findChild(path[paramEnd:], params)
		}
		if !strings.HasPrefix(path, n.s) { // Node doesn't match beginning of path.
			continue
		}
		if n.isWildcard() {
			*params = append(*params, path[len(n.s):])
			return n
		}
		if len(path) == len(n.s) { // Node matched until the end of path.
			return n
		}
		child := n.children.findChild(path[len(n.s):], params)
		if child == nil {
			continue // No match from children, maybe there is a parameter in next same-level node.
		}
		return child
	}
	return nil
}

// sort puts nodes with most subnodes on top and plain strings before parameter.
func (nn *nodes) sort() {
	sort.Slice(*nn, func(i, j int) bool {
		if (*nn)[i].s == ":" || (*nn)[i].isWildcard() {
			return false
		}
		if (*nn)[j].s == ":" || (*nn)[j].isWildcard() {
			return true
		}
		return (*nn)[i].countChildren() > (*nn)[j].countChildren()
	})
	for _, n := range *nn {
		n.children.sort()
	}
}
<commit_msg>Fix package doc<commit_after>package router

import (
	"fmt"
	"net/http"
	"sort"
	"strings"
)

type node struct {
	s        string
	params   []string // Parameter's names from the parent node to this one.
	children nodes
	handler  http.Handler
}

func (n *node) string(level int) (s string) {
	s += fmt.Sprintf("%s%q  %v  %v\n", strings.Repeat("\t", level), n.s, n.params, n.handler)
	for _, n := range n.children {
		s += n.string(level + 1)
	}
	return
}

func (n *node) isWildcard() bool {
	return n.s[len(n.s)-1] == '/' && len(n.children) == 0
}

func (n *node) countChildren() (i int) {
	for _, n := range n.children {
		i++
		i += n.countChildren()
	}
	return
}

type nodes []*node

// makeChild adds a node to the tree.
func (nn *nodes) makeChild(path string, params []string, handler http.Handler) {
NodesLoop:
	for _, n := range *nn {
		minlen := len(n.s)
		if len(path) < minlen {
			minlen = len(path)
		}
		for i := 0; i < minlen; i++ {
			if n.s[i] == path[i] {
				continue
			}
			if i == 0 { // No match from the first byte: see next same-level node.
				continue NodesLoop
			}
			// Difference in the middle of a node: split current node to make subnode and transfer handler to it.
			*n = node{
				s: n.s[:i],
				children: nodes{
					{s: n.s[i:], params: n.params, children: n.children, handler: n.handler},
					{s: path[i:], params: params, handler: handler},
				},
			}
			return
		}
		if len(path) < len(n.s) { // s fully matched first part of n.s: split node.
			*n = node{
				s:      n.s[:len(path)],
				params: params,
				children: nodes{
					{s: n.s[len(path):], params: n.params, children: n.children, handler: n.handler},
				},
				handler: handler,
			}
		} else if len(path) > len(n.s) { // n.s fully matched first part of s: see subnodes for the rest.
			n.children.makeChild(path[len(n.s):], params, handler)
		} else { // s == n.s and no rest: node has no handler or route is duplicated.
			if handler == nil { // No handler provided (must be a non-ending path parameter): don't overwrite.
				return
			}
			if n.handler != nil { // Handler provided but n.handler already set: route is duplicated.
				panic(fmt.Errorf("router: two or more routes have same path"))
			}
			n.params = params
			n.handler = handler
		}
		return
	}
	*nn = append(*nn, &node{s: path, params: params, handler: handler}) // Not a single byte match on same-level nodes: append a new one.
}

func (nn *nodes) findChild(path string, params *[]string) *node {
	for _, n := range *nn {
		if n.s == ":" { // Handle parameter node.
			paramEnd := strings.IndexByte(path, '/')
			if paramEnd == -1 { // Path ends with the parameter.
				if n.handler != nil { // Performance: append parameter only if the node has a handler (otherwise useless).
					*params = append(*params, path)
				}
				return n
			}
			*params = append(*params, path[:paramEnd])
			return n.children.findChild(path[paramEnd:], params)
		}
		if !strings.HasPrefix(path, n.s) { // Node doesn't match beginning of path.
			continue
		}
		if n.isWildcard() {
			*params = append(*params, path[len(n.s):])
			return n
		}
		if len(path) == len(n.s) { // Node matched until the end of path.
			return n
		}
		child := n.children.findChild(path[len(n.s):], params)
		if child == nil {
			continue // No match from children, maybe there is a parameter in next same-level node.
		}
		return child
	}
	return nil
}

// sort puts nodes with most subnodes on top and plain strings before parameter.
func (nn *nodes) sort() {
	sort.Slice(*nn, func(i, j int) bool {
		if (*nn)[i].s == ":" || (*nn)[i].isWildcard() {
			return false
		}
		if (*nn)[j].s == ":" || (*nn)[j].isWildcard() {
			return true
		}
		return (*nn)[i].countChildren() > (*nn)[j].countChildren()
	})
	for _, n := range *nn {
		n.children.sort()
	}
}
<|endoftext|>