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

package main

import (
	"os"
	"os/signal"
	"runtime"
	"sync"
	"time"

	"github.com/yarpc/yab/limiter"
	"github.com/yarpc/yab/statsd"
	"github.com/yarpc/yab/transport"
)

// setGoMaxProcs sets runtime.GOMAXPROCS if the option is set
// and returns the number of GOMAXPROCS configured.
func (o BenchmarkOptions) setGoMaxProcs() int {
	if o.NumCPUs > 0 {
		runtime.GOMAXPROCS(o.NumCPUs)
	}
	return runtime.GOMAXPROCS(-1)
}

func (o BenchmarkOptions) getNumConnections(goMaxProcs int) int {
	if o.Connections > 0 {
		return o.Connections
	}

	// If the user doesn't specify a number of connections, choose a sane default.
	return goMaxProcs * 2
}

func runWorker(t transport.Transport, m benchmarkMethod, s *benchmarkState, run *limiter.Run) {
	for cur := run; cur.More(); {
		latency, err := m.call(t)
		if err != nil {
			s.recordError(err)
			continue
		}

		s.recordLatency(latency)
	}
}

func runBenchmark(out output, allOpts Options, m benchmarkMethod) {
	opts := allOpts.BOpts

	// By default, benchmarks are disabled. At least MaxDuration needs to
	// be set to enable them.
	if opts.MaxDuration == 0 {
		return
	}

	goMaxProcs := opts.setGoMaxProcs()
	numConns := opts.getNumConnections(goMaxProcs)
	out.Printf("Benchmark parameters:\n")
	out.Printf("  CPUs:            %v\n", goMaxProcs)
	out.Printf("  Connections:     %v\n", numConns)
	out.Printf("  Concurrency:     %v\n", opts.Concurrency)
	out.Printf("  Max requests:    %v\n", opts.MaxRequests)
	out.Printf("  Max duration:    %v\n", opts.MaxDuration)
	out.Printf("  Max RPS:         %v\n", opts.RPS)

	// Warm up number of connections.
	connections, err := m.WarmTransports(numConns, allOpts.TOpts, opts.WarmupRequests)
	if err != nil {
		out.Fatalf("Failed to warmup connections for benchmark: %v", err)
	}

	statter, err := statsd.NewClient(opts.StatsdHostPort, allOpts.TOpts.ServiceName, allOpts.ROpts.MethodName)
	if err != nil {
		out.Fatalf("Failed to create statsd client for benchmark: %v", err)
	}

	var wg sync.WaitGroup
	states := make([]*benchmarkState, len(connections)*opts.Concurrency)
	for i := range states {
		states[i] = newBenchmarkState(statter)
	}

	run := limiter.New(opts.MaxRequests, opts.RPS, opts.MaxDuration)
	stopOnInterrupt(out, run)

	start := time.Now()
	for i, c := range connections {
		for j := 0; j < opts.Concurrency; j++ {
			state := states[i*opts.Concurrency+j]

			wg.Add(1)
			go func(c transport.Transport) {
				defer wg.Done()
				runWorker(c, m, state, run)
			}(c)
		}
	}

	// TODO: Support streaming updates.

	// Wait for all the worker goroutines to end.
	wg.Wait()
	total := time.Since(start)

	// Merge all the states into 0
	overall := states[0]
	for _, s := range states[1:] {
		overall.merge(s)
	}

	overall.printErrors(out)
	overall.printLatencies(out)

	out.Printf("Elapsed time:      %v\n", (total / time.Millisecond * time.Millisecond))
	out.Printf("Total requests:    %v\n", len(overall.latencies))
	out.Printf("RPS:               %.2f\n", float64(len(overall.latencies))/total.Seconds())
}

// stopOnInterrupt sets up a signal that will trigger the run to stop.
func stopOnInterrupt(out output, r *limiter.Run) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	go func() {
		<-c
		// Preceding newline since Ctrl-C will be printed inline.
		out.Printf("\n!!Benchmark interrupted!!\n")
		r.Stop()
	}()
}
<commit_msg>Cap MaxRequests by RPS * MaxDuration<commit_after>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package main

import (
	"os"
	"os/signal"
	"runtime"
	"sync"
	"time"

	"github.com/yarpc/yab/limiter"
	"github.com/yarpc/yab/statsd"
	"github.com/yarpc/yab/transport"
)

// setGoMaxProcs sets runtime.GOMAXPROCS if the option is set
// and returns the number of GOMAXPROCS configured.
func (o BenchmarkOptions) setGoMaxProcs() int {
	if o.NumCPUs > 0 {
		runtime.GOMAXPROCS(o.NumCPUs)
	}
	return runtime.GOMAXPROCS(-1)
}

func (o BenchmarkOptions) getNumConnections(goMaxProcs int) int {
	if o.Connections > 0 {
		return o.Connections
	}

	// If the user doesn't specify a number of connections, choose a sane default.
	return goMaxProcs * 2
}

func runWorker(t transport.Transport, m benchmarkMethod, s *benchmarkState, run *limiter.Run) {
	for cur := run; cur.More(); {
		latency, err := m.call(t)
		if err != nil {
			s.recordError(err)
			continue
		}

		s.recordLatency(latency)
	}
}

func runBenchmark(out output, allOpts Options, m benchmarkMethod) {
	opts := allOpts.BOpts

	// By default, benchmarks are disabled. At least MaxDuration needs to
	// be set to enable them.
	if opts.MaxDuration == 0 {
		return
	}

	if opts.RPS > 0 {
		// The RPS * duration in seconds may cap opts.MaxRequests.
		rpsMax := int(float64(opts.RPS) * opts.MaxDuration.Seconds())
		if rpsMax < opts.MaxRequests {
			opts.MaxRequests = rpsMax
		}
	}

	goMaxProcs := opts.setGoMaxProcs()
	numConns := opts.getNumConnections(goMaxProcs)
	out.Printf("Benchmark parameters:\n")
	out.Printf("  CPUs:            %v\n", goMaxProcs)
	out.Printf("  Connections:     %v\n", numConns)
	out.Printf("  Concurrency:     %v\n", opts.Concurrency)
	out.Printf("  Max requests:    %v\n", opts.MaxRequests)
	out.Printf("  Max duration:    %v\n", opts.MaxDuration)
	out.Printf("  Max RPS:         %v\n", opts.RPS)

	// Warm up number of connections.
	connections, err := m.WarmTransports(numConns, allOpts.TOpts, opts.WarmupRequests)
	if err != nil {
		out.Fatalf("Failed to warmup connections for benchmark: %v", err)
	}

	statter, err := statsd.NewClient(opts.StatsdHostPort, allOpts.TOpts.ServiceName, allOpts.ROpts.MethodName)
	if err != nil {
		out.Fatalf("Failed to create statsd client for benchmark: %v", err)
	}

	var wg sync.WaitGroup
	states := make([]*benchmarkState, len(connections)*opts.Concurrency)
	for i := range states {
		states[i] = newBenchmarkState(statter)
	}

	run := limiter.New(opts.MaxRequests, opts.RPS, opts.MaxDuration)
	stopOnInterrupt(out, run)

	start := time.Now()
	for i, c := range connections {
		for j := 0; j < opts.Concurrency; j++ {
			state := states[i*opts.Concurrency+j]

			wg.Add(1)
			go func(c transport.Transport) {
				defer wg.Done()
				runWorker(c, m, state, run)
			}(c)
		}
	}

	// TODO: Support streaming updates.

	// Wait for all the worker goroutines to end.
	wg.Wait()
	total := time.Since(start)

	// Merge all the states into 0
	overall := states[0]
	for _, s := range states[1:] {
		overall.merge(s)
	}

	overall.printErrors(out)
	overall.printLatencies(out)

	out.Printf("Elapsed time:      %v\n", (total / time.Millisecond * time.Millisecond))
	out.Printf("Total requests:    %v\n", len(overall.latencies))
	out.Printf("RPS:               %.2f\n", float64(len(overall.latencies))/total.Seconds())
}

// stopOnInterrupt sets up a signal that will trigger the run to stop.
func stopOnInterrupt(out output, r *limiter.Run) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	go func() {
		<-c
		// Preceding newline since Ctrl-C will be printed inline.
		out.Printf("\n!!Benchmark interrupted!!\n")
		r.Stop()
	}()
}
<|endoftext|>


<commit_before>package molecule

import (
	"sync"

	cmn "github.com/RxnWeaver/rxnweaver/common"
)

// nextMolIdHolder is a synchronised struct used to assign a
// globally-unique ID to each molecule.
type nextMolIdHolder struct {
	mu     sync.Mutex
	nextId uint32
}

// The only instance of `nextMolIdHolder`.
var nextMolId nextMolIdHolder

func nextMoleculeId() uint32 {
	nextMolId.mu.Lock()
	defer nextMolId.mu.Unlock()

	nextMolId.nextId++
	return nextMolId.nextId
}

// Molecule represents a chemical molecule.
//
// It holds information concerning its atom, bonds, rings, etc.  Note
// that a molecule is expected to be a single connected component.
type Molecule struct {
	id uint32 // The globally-unique ID of this molecule.

	atoms       []*_Atom       // List of atoms in this molecule.
	bonds       []*_Bond       // List of bonds in this molecule.
	rings       []*_Ring       // List of rings in this molecule.
	ringSystems []*_RingSystem // List of ring systems in this molecule.

	nextAtomIid      uint16 // Running number for atom input IDs.
	nextBondId       uint16 // Running number for bond IDs.
	nextRingId       uint8  // Running number for ring IDs.
	nextRingSystemId uint8  // Running number for ring system IDs.

	vendor           string // Optional string identifying the supplier.
	vendorMoleculeId string // Optional supplier-specified ID.

	attributes []Attribute // Optional list of annotations.

	dists [][]int // Matrix of pair-wise distances between atoms.
	paths [][]int // Lists of pair-wise paths between atoms.
}

// New creates and initialises a molecule.
func New() *Molecule {
	mol := new(Molecule)
	mol.id = nextMoleculeId()

	mol.atoms = make([]*_Atom, 0, cmn.ListSizeLarge)
	mol.bonds = make([]*_Bond, 0, cmn.ListSizeLarge)
	mol.rings = make([]*_Ring, 0, cmn.ListSizeSmall)
	mol.ringSystems = make([]*_RingSystem, 0, cmn.ListSizeSmall)

	mol.nextAtomIid = 1
	mol.nextBondId = 1
	mol.nextRingId = 1
	mol.nextRingSystemId = 1

	mol.attributes = make([]Attribute, 0, cmn.ListSizeTiny)

	return mol
}

// NewAtomBuilder answers a new atom builder.
func (m *Molecule) NewAtomBuilder() *AtomBuilder {
	return &AtomBuilder{m, nil}
}

// NewBondBuilder answers a new bond builder.
func (m *Molecule) NewBondBuilder() *BondBuilder {
	return &BondBuilder{m, nil}
}

// Id answers the globally-unique ID of this molecule.
func (m *Molecule) Id() uint32 {
	return m.id
}

// atomWithIid answers the atom for the given input ID, if found.
// Answers `nil` otherwise.
func (m *Molecule) atomWithIid(id uint16) *_Atom {
	for _, a := range m.atoms {
		if a.iId == id {
			return a
		}
	}

	return nil
}

// atomWithNid answers the atom for the given normalised ID, if found.
// Answers `nil` otherwise.
func (m *Molecule) atomWithNid(id uint16) *_Atom {
	for _, a := range m.atoms {
		if a.nId == id {
			return a
		}
	}

	return nil
}

// bondWithId answers the bond for the given ID, if found.  Answers
// `nil` otherwise.
func (m *Molecule) bondWithId(id uint16) *_Bond {
	for _, b := range m.bonds {
		if b.id == id {
			return b
		}
	}

	return nil
}

// ringWithId answers the ring for the given ID, if found.  Answers
// `nil` otherwise.
func (m *Molecule) ringWithId(id uint8) *_Ring {
	for _, r := range m.rings {
		if r.id == id {
			return r
		}
	}

	return nil
}

// bondBetween answers the bond between the two given atoms, if one
// such exists.  Answers `nil` otherwise.
//
// Note that the two given atoms are represented by their input IDs,
// NOT normalised IDs.
func (m *Molecule) bondBetween(a1id, a2id uint16) *_Bond {
	for _, b := range m.bonds {
		if (b.a1 == a1id && b.a2 == a2id) || (b.a2 == a1id && b.a1 == a2id) {
			return b
		}
	}

	return nil
}

// bondCount answers the total number of bonds of the given type in
// this molecule.
func (m *Molecule) bondCount(typ cmn.BondType) int {
	c := 0
	for _, b := range m.bonds {
		if b.bType == typ {
			c++
		}
	}

	return c
}

// singleBondCount answers the total number of single bonds in this
// molecule.
func (m *Molecule) singleBondCount() int {
	return m.bondCount(cmn.BondTypeSingle)
}

// doubleBondCount answers the total number of double bonds in this
// molecule.
func (m *Molecule) doubleBondCount() int {
	return m.bondCount(cmn.BondTypeDouble)
}

// tripleBondCount answers the total number of triple bonds in this
// molecule.
func (m *Molecule) tripleBondCount() int {
	return m.bondCount(cmn.BondTypeTriple)
}

// aromaticRingCount answers the number of aromatic rings in this
// molecule.
func (m *Molecule) aromaticRingCount() int {
	c := 0
	for _, r := range m.rings {
		if r.isAro {
			c++
		}
	}

	return c
}

// aromaticRingSystemCount answers the number of aromatic ring systems
// in this molecule.
func (m *Molecule) aromaticRingSystemCount() int {
	c := 0
	for _, rs := range m.ringSystems {
		if rs.isAro {
			c++
		}
	}

	return c
}
<commit_msg>Introduce communication channel pair per molecule<commit_after>package molecule

import (
	"sync"

	cmn "github.com/RxnWeaver/rxnweaver/common"
)

// nextMolIdHolder is a synchronised struct used to assign a
// globally-unique ID to each molecule.
type nextMolIdHolder struct {
	mu     sync.Mutex
	nextId uint32
}

// The only instance of `nextMolIdHolder`.
var nextMolId nextMolIdHolder

func nextMoleculeId() uint32 {
	nextMolId.mu.Lock()
	defer nextMolId.mu.Unlock()

	nextMolId.nextId++
	return nextMolId.nextId
}

// Molecule represents a chemical molecule.
//
// It holds information concerning its atom, bonds, rings, etc.  Note
// that a molecule is expected to be a single connected component.
type Molecule struct {
	id uint32 // The globally-unique ID of this molecule.

	// Channel on which this molecule receives requests and
	// notifications.
	inChannel chan InMessage
	// Channel on which this molecule sends responses to received
	// requests.
	outChannel chan OutMessage

	atoms       []*_Atom       // List of atoms in this molecule.
	bonds       []*_Bond       // List of bonds in this molecule.
	rings       []*_Ring       // List of rings in this molecule.
	ringSystems []*_RingSystem // List of ring systems in this molecule.

	nextAtomIid      uint16 // Running number for atom input IDs.
	nextBondId       uint16 // Running number for bond IDs.
	nextRingId       uint8  // Running number for ring IDs.
	nextRingSystemId uint8  // Running number for ring system IDs.

	vendor           string // Optional string identifying the supplier.
	vendorMoleculeId string // Optional supplier-specified ID.

	attributes []Attribute // Optional list of annotations.

	dists [][]int // Matrix of pair-wise distances between atoms.
	paths [][]int // Lists of pair-wise paths between atoms.
}

// New creates and initialises a molecule.
func New() *Molecule {
	mol := new(Molecule)
	mol.id = nextMoleculeId()

	mol.atoms = make([]*_Atom, 0, cmn.ListSizeLarge)
	mol.bonds = make([]*_Bond, 0, cmn.ListSizeLarge)
	mol.rings = make([]*_Ring, 0, cmn.ListSizeSmall)
	mol.ringSystems = make([]*_RingSystem, 0, cmn.ListSizeSmall)

	mol.attributes = make([]Attribute, 0, cmn.ListSizeTiny)

	return mol
}

// NewAtomBuilder answers a new atom builder.
func (m *Molecule) NewAtomBuilder() *AtomBuilder {
	return &AtomBuilder{m, nil}
}

// NewBondBuilder answers a new bond builder.
func (m *Molecule) NewBondBuilder() *BondBuilder {
	return &BondBuilder{m, nil}
}

// Id answers the globally-unique ID of this molecule.
func (m *Molecule) Id() uint32 {
	return m.id
}

// InChannel answers the input channel of this molecule.
func (m *Molecule) InChannel() chan InMessage {
	return m.inChannel
}

// OutChannel answers the output channel of this molecule.
func (m *Molecule) OutChannel() chan OutMessage {
	return m.outChannel
}

// atomWithIid answers the atom for the given input ID, if found.
// Answers `nil` otherwise.
func (m *Molecule) atomWithIid(id uint16) *_Atom {
	for _, a := range m.atoms {
		if a.iId == id {
			return a
		}
	}

	return nil
}

// atomWithNid answers the atom for the given normalised ID, if found.
// Answers `nil` otherwise.
func (m *Molecule) atomWithNid(id uint16) *_Atom {
	for _, a := range m.atoms {
		if a.nId == id {
			return a
		}
	}

	return nil
}

// bondWithId answers the bond for the given ID, if found.  Answers
// `nil` otherwise.
func (m *Molecule) bondWithId(id uint16) *_Bond {
	for _, b := range m.bonds {
		if b.id == id {
			return b
		}
	}

	return nil
}

// ringWithId answers the ring for the given ID, if found.  Answers
// `nil` otherwise.
func (m *Molecule) ringWithId(id uint8) *_Ring {
	for _, r := range m.rings {
		if r.id == id {
			return r
		}
	}

	return nil
}

// bondBetween answers the bond between the two given atoms, if one
// such exists.  Answers `nil` otherwise.
//
// Note that the two given atoms are represented by their input IDs,
// NOT normalised IDs.
func (m *Molecule) bondBetween(a1id, a2id uint16) *_Bond {
	for _, b := range m.bonds {
		if (b.a1 == a1id && b.a2 == a2id) || (b.a2 == a1id && b.a1 == a2id) {
			return b
		}
	}

	return nil
}

// bondCount answers the total number of bonds of the given type in
// this molecule.
func (m *Molecule) bondCount(typ cmn.BondType) int {
	c := 0
	for _, b := range m.bonds {
		if b.bType == typ {
			c++
		}
	}

	return c
}

// singleBondCount answers the total number of single bonds in this
// molecule.
func (m *Molecule) singleBondCount() int {
	return m.bondCount(cmn.BondTypeSingle)
}

// doubleBondCount answers the total number of double bonds in this
// molecule.
func (m *Molecule) doubleBondCount() int {
	return m.bondCount(cmn.BondTypeDouble)
}

// tripleBondCount answers the total number of triple bonds in this
// molecule.
func (m *Molecule) tripleBondCount() int {
	return m.bondCount(cmn.BondTypeTriple)
}

// aromaticRingCount answers the number of aromatic rings in this
// molecule.
func (m *Molecule) aromaticRingCount() int {
	c := 0
	for _, r := range m.rings {
		if r.isAro {
			c++
		}
	}

	return c
}

// aromaticRingSystemCount answers the number of aromatic ring systems
// in this molecule.
func (m *Molecule) aromaticRingSystemCount() int {
	c := 0
	for _, rs := range m.ringSystems {
		if rs.isAro {
			c++
		}
	}

	return c
}
<|endoftext|>


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

package format

import (
	"github.com/trivago/gollum/core"
	"strconv"
)

// Runlength formatter plugin
// Runlength is a formatter that prepends the length of the message, followed by
// a ":". The actual message is formatted by a nested formatter.
// Configuration example
//
//  - format.Runlength
//      Separator: ":"
//      ApplyTo: "payload" # payload or <metaKey>
//
// Separator sets the separator character placed after the runlength.
// This is set to ":" by default. If no separator is set the runlength will only set.
//
// ApplyTo defines the formatter content for the data transferred
type Runlength struct {
	core.SimpleFormatter
	separator []byte
}

func init() {
	core.TypeRegistry.Register(Runlength{})
}

// Configure initializes this formatter with values from a plugin config.
func (format *Runlength) Configure(conf core.PluginConfigReader) error {
	format.SimpleFormatter.Configure(conf)

	format.separator = []byte(conf.GetString("Separator", ":"))
	return conf.Errors.OrNil()
}

// ApplyFormatter update message payload
func (format *Runlength) ApplyFormatter(msg *core.Message) error {
	content := format.GetAppliedContent(msg)
	contentLen := len(content)
	separatorLen := len(format.separator)

	lengthStr := strconv.Itoa(contentLen)
	var payload []byte

	if separatorLen > 0 {
		dataSize := len(lengthStr) + separatorLen + contentLen
		payload = core.MessageDataPool.Get(dataSize)

		offset := copy(payload, []byte(lengthStr))
		offset += copy(payload[offset:], format.separator)
		copy(payload[offset:], content)
	} else {
		payload = []byte(lengthStr)
	}

	format.SetAppliedContent(msg, payload)
	return nil
}
<commit_msg>Code style<commit_after>// Copyright 2015-2017 trivago GmbH
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package format

import (
	"github.com/trivago/gollum/core"
	"strconv"
)

// Runlength formatter plugin
// Runlength is a formatter that prepends the length of the message, followed by
// a ":". The actual message is formatted by a nested formatter.
// Configuration example
//
//  - format.Runlength
//      Separator: ":"
//      ApplyTo: "payload" # payload or <metaKey>
//
// Separator sets the separator character placed after the runlength.
// This is set to ":" by default. If no separator is set the runlength will only set.
//
// ApplyTo defines the formatter content for the data transferred
type Runlength struct {
	core.SimpleFormatter
	separator []byte
}

func init() {
	core.TypeRegistry.Register(Runlength{})
}

// Configure initializes this formatter with values from a plugin config.
func (format *Runlength) Configure(conf core.PluginConfigReader) error {
	format.SimpleFormatter.Configure(conf)

	format.separator = []byte(conf.GetString("Separator", ":"))
	return conf.Errors.OrNil()
}

// ApplyFormatter update message payload
func (format *Runlength) ApplyFormatter(msg *core.Message) error {
	content := format.GetAppliedContent(msg)
	contentLen := len(content)
	separatorLen := len(format.separator)
	lengthStr := strconv.Itoa(contentLen)

	var payload []byte
	if separatorLen > 0 {
		dataSize := len(lengthStr) + separatorLen + contentLen
		payload = core.MessageDataPool.Get(dataSize)

		offset := copy(payload, []byte(lengthStr))
		offset += copy(payload[offset:], format.separator)
		copy(payload[offset:], content)
	} else {
		payload = []byte(lengthStr)
	}

	format.SetAppliedContent(msg, payload)
	return nil
}
<|endoftext|>


<commit_before>package zoom

// File contains code strictly related to the database, including
// setting up the database with given config, generating unique,
// random ids, and creating and managing a connection pool. There
// are also convenience functions for (e.g.) checking if a key exists
// in redis.

import (
	"github.com/dchest/uniuri"
	"github.com/stephenalexbrowne/zoom/redis"
	"strconv"
	"time"
)

type Configuration struct {
	Address  string // Address to connect to. Default: "localhost:6380"
	Network  string // Network to use. Default: "tcp"
	Database int    // Database id to use (using SELECT). Default: 0
}

var pool *redis.Pool

var defaultConfiguration = Configuration{
	Address:  "localhost:6379",
	Network:  "tcp",
	Database: 0,
}

func GetConn() redis.Conn {
	return pool.Get()
}

// initializes a connection pool to be used to conect to database
// TODO: add some config options
func Init(passedConfig *Configuration) {

	config := getConfiguration(passedConfig)

	pool = &redis.Pool{
		MaxIdle:     3,
		MaxActive:   0,
		IdleTimeout: 240 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial(config.Network, config.Address)
			if err != nil {
				return nil, err
			}
			if _, err := c.Do("select", strconv.Itoa(config.Database)); err != nil {
				c.Close()
				return nil, err
			}
			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
}

// closes the connection pool
// Should be run when application exits
func Close() {
	pool.Close()
}

// Returns true iff a given key exists in redis
// If conn is nil, a connection will be created for you
// said connection will be closed before the end of the function
func KeyExists(key string, conn redis.Conn) (bool, error) {
	if conn == nil {
		conn = pool.Get()
		defer conn.Close()
	}
	return redis.Bool(conn.Do("exists", key))
}

// Returns true iff the redis set identified by key contains member.
// If conn is nil, a connection will be created for you.
// said connection will be closed before the end of the function.
func SetContains(key, member string, conn redis.Conn) (bool, error) {
	if conn == nil {
		conn = pool.Get()
		defer conn.Close()
	}
	return redis.Bool(conn.Do("sismember", key, member))
}

// generates a random string that is more or less
// garunteed to be unique. Used as Ids for records
// where an Id is not otherwise provided.
func generateRandomId() string {
	timeInt := time.Now().Unix()
	timeString := strconv.FormatInt(timeInt, 36)
	randomString := uniuri.NewLen(16)
	return randomString + timeString
}

// adds value as a member of a redis set identified by {name}:index
// where {name} is the name of the model you want to index.
// If the conn paramater is nil, will get a connection from the
// pool and close it before returning. If conn is a redis.Conn,
// it will use the existing connection.
func addToIndex(name, value string, conn redis.Conn) error {
	if conn == nil {
		conn = pool.Get()
		defer conn.Close()
	}
	key := name + ":index"
	_, err := conn.Do("sadd", key, value)
	return err
}

// return a proper configuration struct.
// if the passed in struct is nil, return defaultConfiguration
// else, for each attribute, if the passed in struct is "", 0, etc,
// use the default value for that attribute.
func getConfiguration(passedConfig *Configuration) Configuration {

	if passedConfig == nil {
		return defaultConfiguration
	}

	// copy the passedConfig
	newConfig := *passedConfig

	if newConfig.Address == "" {
		newConfig.Address = defaultConfiguration.Address
	}
	if newConfig.Network == "" {
		newConfig.Network = defaultConfiguration.Network
	}
	// since the zero value for int is 0, we can skip config.Database

	return newConfig
}
<commit_msg>Fix typo in comment on default connection Address<commit_after>package zoom

// File contains code strictly related to the database, including
// setting up the database with given config, generating unique,
// random ids, and creating and managing a connection pool. There
// are also convenience functions for (e.g.) checking if a key exists
// in redis.

import (
	"github.com/dchest/uniuri"
	"github.com/stephenalexbrowne/zoom/redis"
	"strconv"
	"time"
)

type Configuration struct {
	Address  string // Address to connect to. Default: "localhost:6379"
	Network  string // Network to use. Default: "tcp"
	Database int    // Database id to use (using SELECT). Default: 0
}

var pool *redis.Pool

var defaultConfiguration = Configuration{
	Address:  "localhost:6379",
	Network:  "tcp",
	Database: 0,
}

func GetConn() redis.Conn {
	return pool.Get()
}

// initializes a connection pool to be used to conect to database
// TODO: add some config options
func Init(passedConfig *Configuration) {

	config := getConfiguration(passedConfig)

	pool = &redis.Pool{
		MaxIdle:     3,
		MaxActive:   0,
		IdleTimeout: 240 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial(config.Network, config.Address)
			if err != nil {
				return nil, err
			}
			if _, err := c.Do("select", strconv.Itoa(config.Database)); err != nil {
				c.Close()
				return nil, err
			}
			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
}

// closes the connection pool
// Should be run when application exits
func Close() {
	pool.Close()
}

// Returns true iff a given key exists in redis
// If conn is nil, a connection will be created for you
// said connection will be closed before the end of the function
func KeyExists(key string, conn redis.Conn) (bool, error) {
	if conn == nil {
		conn = pool.Get()
		defer conn.Close()
	}
	return redis.Bool(conn.Do("exists", key))
}

// Returns true iff the redis set identified by key contains member.
// If conn is nil, a connection will be created for you.
// said connection will be closed before the end of the function.
func SetContains(key, member string, conn redis.Conn) (bool, error) {
	if conn == nil {
		conn = pool.Get()
		defer conn.Close()
	}
	return redis.Bool(conn.Do("sismember", key, member))
}

// generates a random string that is more or less
// garunteed to be unique. Used as Ids for records
// where an Id is not otherwise provided.
func generateRandomId() string {
	timeInt := time.Now().Unix()
	timeString := strconv.FormatInt(timeInt, 36)
	randomString := uniuri.NewLen(16)
	return randomString + timeString
}

// adds value as a member of a redis set identified by {name}:index
// where {name} is the name of the model you want to index.
// If the conn paramater is nil, will get a connection from the
// pool and close it before returning. If conn is a redis.Conn,
// it will use the existing connection.
func addToIndex(name, value string, conn redis.Conn) error {
	if conn == nil {
		conn = pool.Get()
		defer conn.Close()
	}
	key := name + ":index"
	_, err := conn.Do("sadd", key, value)
	return err
}

// return a proper configuration struct.
// if the passed in struct is nil, return defaultConfiguration
// else, for each attribute, if the passed in struct is "", 0, etc,
// use the default value for that attribute.
func getConfiguration(passedConfig *Configuration) Configuration {

	if passedConfig == nil {
		return defaultConfiguration
	}

	// copy the passedConfig
	newConfig := *passedConfig

	if newConfig.Address == "" {
		newConfig.Address = defaultConfiguration.Address
	}
	if newConfig.Network == "" {
		newConfig.Network = defaultConfiguration.Network
	}
	// since the zero value for int is 0, we can skip config.Database

	return newConfig
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"strings"
	"time"

	"github.com/jinzhu/gorm"
	_ "github.com/mattn/go-sqlite3"
)

const (
	sqlSelectByFreq = "cmd, dir, '' as hostname, '' as shell_session_id, time" // TODO: do max(time) here.
	sqlGroupByFreq  = "cmd, dir"
	sqlSelectByDate = "cmd, dir, hostname, shell_session_id, time"
)

// setting holds the persisted settings.
type setting struct {
	ID            int
	SortByFreq    bool
	OnlyMySession bool
	OnlyMyCwd     bool
}

// record holds the data recorded for a single shell command.
type record struct {
	Cmd            string `sql:"size:65535"`
	Dir            string `sql:"size:65535"`
	Hostname       string `sql:"size:65535"`
	ShellSessionID string `sql:"size:65535"`
	Time           time.Time
}

// query holds the components of a database query.
type query struct {
	Cmd            *string
	Dir            *string
	Hostname       *string
	ShellSessionID *string
	SortByFreq     bool
	Limit          int
}

// database holds a database connection and provides insert and retrieval.
type database interface {
	Add(*record) error
	Close() error
	Query(query) ([]record, error)
	Setting() (setting, error)
	WriteSetting(*setting) error
}

type sqlDatabase struct {
	db gorm.DB
	database
}

func newDatabase(path string) (database, error) {
	d := &sqlDatabase{}
	// Check if the DB already exists, or if we must create the table.
	_, err := os.Stat(path)
	n := os.IsNotExist(err)
	d.db, err = gorm.Open("sqlite3", path)
	// If new, we must create the table.
	if n {
		if err := d.db.CreateTable(&setting{}).Error; err != nil {
			return nil, err
		}
		if err := d.db.CreateTable(&record{}).Error; err != nil {
			return nil, err
		}
	}
	return d, err
}

func (d *sqlDatabase) Add(r *record) error {
	return d.db.Create(r).Error
}

func (d *sqlDatabase) Close() error {
	return d.db.Close()
}

func (d *sqlDatabase) Query(q query) ([]record, error) {
	var rs []record
	var ws []string
	var ps []interface{}
	if q.Cmd != nil {
		ws = append(ws, "cmd LIKE ?")
		ps = append(ps, "%"+*q.Cmd+"%")
	}
	if q.Dir != nil {
		ws = append(ws, "dir = ?")
		ps = append(ps, q.Dir)
	}
	if q.Hostname != nil {
		ws = append(ws, "hostname = ?")
		ps = append(ps, q.Hostname)
	}
	if q.ShellSessionID != nil {
		ws = append(ws, "shell_session_id = ?")
		ps = append(ps, q.ShellSessionID)
	}
	var db *gorm.DB
	if q.SortByFreq {
		db = d.db.Table("records").Select(sqlSelectByFreq).Where(strings.Join(ws, " and "), ps...).Group(sqlGroupByFreq)
	} else {
		db = d.db.Table("records").Select(sqlSelectByDate).Where(strings.Join(ws, " and "), ps...)
	}
	err := db.Scan(&rs).Error
	return rs, err
}

func (d *sqlDatabase) Setting() (setting, error) {
	var s setting
	r := d.db.First(&s)
	if r.RecordNotFound() {
		return s, nil
	}
	return s, r.Error
}

func (d *sqlDatabase) WriteSetting(s *setting) error {
	return d.db.Save(s).Error
}
<commit_msg>For 'by frequency', group only by command.<commit_after>package main

import (
	"os"
	"strings"
	"time"

	"github.com/jinzhu/gorm"
	_ "github.com/mattn/go-sqlite3"
)

const (
	sqlSelectByFreq = "cmd, dir, '' as hostname, '' as shell_session_id, time" // TODO: do max(time) here.
	sqlGroupByFreq  = "cmd"
	sqlSelectByDate = "cmd, dir, hostname, shell_session_id, time"
)

// setting holds the persisted settings.
type setting struct {
	ID            int
	SortByFreq    bool
	OnlyMySession bool
	OnlyMyCwd     bool
}

// record holds the data recorded for a single shell command.
type record struct {
	Cmd            string `sql:"size:65535"`
	Dir            string `sql:"size:65535"`
	Hostname       string `sql:"size:65535"`
	ShellSessionID string `sql:"size:65535"`
	Time           time.Time
}

// query holds the components of a database query.
type query struct {
	Cmd            *string
	Dir            *string
	Hostname       *string
	ShellSessionID *string
	SortByFreq     bool
	Limit          int
}

// database holds a database connection and provides insert and retrieval.
type database interface {
	Add(*record) error
	Close() error
	Query(query) ([]record, error)
	Setting() (setting, error)
	WriteSetting(*setting) error
}

type sqlDatabase struct {
	db gorm.DB
	database
}

func newDatabase(path string) (database, error) {
	d := &sqlDatabase{}
	// Check if the DB already exists, or if we must create the table.
	_, err := os.Stat(path)
	n := os.IsNotExist(err)
	d.db, err = gorm.Open("sqlite3", path)
	// If new, we must create the table.
	if n {
		if err := d.db.CreateTable(&setting{}).Error; err != nil {
			return nil, err
		}
		if err := d.db.CreateTable(&record{}).Error; err != nil {
			return nil, err
		}
	}
	return d, err
}

func (d *sqlDatabase) Add(r *record) error {
	return d.db.Create(r).Error
}

func (d *sqlDatabase) Close() error {
	return d.db.Close()
}

func (d *sqlDatabase) Query(q query) ([]record, error) {
	var rs []record
	var ws []string
	var ps []interface{}
	if q.Cmd != nil {
		ws = append(ws, "cmd LIKE ?")
		ps = append(ps, "%"+*q.Cmd+"%")
	}
	if q.Dir != nil {
		ws = append(ws, "dir = ?")
		ps = append(ps, q.Dir)
	}
	if q.Hostname != nil {
		ws = append(ws, "hostname = ?")
		ps = append(ps, q.Hostname)
	}
	if q.ShellSessionID != nil {
		ws = append(ws, "shell_session_id = ?")
		ps = append(ps, q.ShellSessionID)
	}
	var db *gorm.DB
	if q.SortByFreq {
		db = d.db.Table("records").Select(sqlSelectByFreq).Where(strings.Join(ws, " and "), ps...).Group(sqlGroupByFreq)
	} else {
		db = d.db.Table("records").Select(sqlSelectByDate).Where(strings.Join(ws, " and "), ps...)
	}
	err := db.Scan(&rs).Error
	return rs, err
}

func (d *sqlDatabase) Setting() (setting, error) {
	var s setting
	r := d.db.First(&s)
	if r.RecordNotFound() {
		return s, nil
	}
	return s, r.Error
}

func (d *sqlDatabase) WriteSetting(s *setting) error {
	return d.db.Save(s).Error
}
<|endoftext|>


<commit_before>/*
Copyright 2013 CoreOS Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package dbus

import (
	"errors"
	"time"

	"github.com/godbus/dbus"
)

const (
	cleanIgnoreInterval = int64(10 * time.Second)
	ignoreInterval      = int64(30 * time.Millisecond)
)

// Subscribe sets up this connection to subscribe to all systemd dbus events.
// This is required before calling SubscribeUnits. When the connection closes
// systemd will automatically stop sending signals so there is no need to
// explicitly call Unsubscribe().
func (c *Conn) Subscribe() error {
	c.sysconn.BusObject().Call("org.freedesktop.DBus.AddMatch", 0,
		"type='signal',interface='org.freedesktop.systemd1.Manager',member='UnitNew'")
	c.sysconn.BusObject().Call("org.freedesktop.DBus.AddMatch", 0,
		"type='signal',interface='org.freedesktop.DBus.Properties',member='PropertiesChanged'")

	err := c.sysobj.Call("org.freedesktop.systemd1.Manager.Subscribe", 0).Store()
	if err != nil {
		c.sysconn.Close()
		return err
	}

	return nil
}

// Unsubscribe this connection from systemd dbus events.
func (c *Conn) Unsubscribe() error {
	err := c.sysobj.Call("org.freedesktop.systemd1.Manager.Unsubscribe", 0).Store()
	if err != nil {
		c.sysconn.Close()
		return err
	}

	return nil
}

func (c *Conn) initSubscription() {
	c.subscriber.ignore = make(map[dbus.ObjectPath]int64)
}

func (c *Conn) initDispatch() {
	ch := make(chan *dbus.Signal, signalBuffer)

	c.sysconn.Signal(ch)

	go func() {
		for {
			signal := <-ch
			switch signal.Name {
			case "org.freedesktop.systemd1.Manager.JobRemoved":
				c.jobComplete(signal)

				unitName := signal.Body[2].(string)
				var unitPath dbus.ObjectPath
				c.sysobj.Call("org.freedesktop.systemd1.Manager.GetUnit", 0, unitName).Store(&unitPath)
				if unitPath != dbus.ObjectPath("") {
					c.sendSubStateUpdate(unitPath)
				}
			case "org.freedesktop.systemd1.Manager.UnitNew":
				c.sendSubStateUpdate(signal.Body[1].(dbus.ObjectPath))
			case "org.freedesktop.DBus.Properties.PropertiesChanged":
				if signal.Body[0].(string) == "org.freedesktop.systemd1.Unit" {
					// we only care about SubState updates, which are a Unit property
					c.sendSubStateUpdate(signal.Path)
				}
			}
		}
	}()
}

// Returns two unbuffered channels which will receive all changed units every
// interval.  Deleted units are sent as nil.
func (c *Conn) SubscribeUnits(interval time.Duration) (<-chan map[string]*UnitStatus, <-chan error) {
	return c.SubscribeUnitsCustom(interval, 0, func(u1, u2 *UnitStatus) bool { return *u1 != *u2 }, nil)
}

// SubscribeUnitsCustom is like SubscribeUnits but lets you specify the buffer
// size of the channels, the comparison function for detecting changes and a filter
// function for cutting down on the noise that your channel receives.
func (c *Conn) SubscribeUnitsCustom(interval time.Duration, buffer int, isChanged func(*UnitStatus, *UnitStatus) bool, filterUnit func (string) bool) (<-chan map[string]*UnitStatus, <-chan error) {
	old := make(map[string]*UnitStatus)
	statusChan := make(chan map[string]*UnitStatus, buffer)
	errChan := make(chan error, buffer)

	go func() {
		for {
			timerChan := time.After(interval)

			units, err := c.ListUnits()
			if err == nil {
				cur := make(map[string]*UnitStatus)
				for i := range units {
					if filterUnit != nil && filterUnit(units[i].Name) {
						continue
					}
					cur[units[i].Name] = &units[i]
				}

				// add all new or changed units
				changed := make(map[string]*UnitStatus)
				for n, u := range cur {
					if oldU, ok := old[n]; !ok || isChanged(oldU, u) {
						changed[n] = u
					}
					delete(old, n)
				}

				// add all deleted units
				for oldN := range old {
					changed[oldN] = nil
				}

				old = cur

				if len(changed) != 0 {
					statusChan <- changed
				}
			} else {
				errChan <- err
			}

			<-timerChan
		}
	}()

	return statusChan, errChan
}

type SubStateUpdate struct {
	UnitName string
	SubState string
}

// SetSubStateSubscriber writes to updateCh when any unit's substate changes.
// Although this writes to updateCh on every state change, the reported state
// may be more recent than the change that generated it (due to an unavoidable
// race in the systemd dbus interface).  That is, this method provides a good
// way to keep a current view of all units' states, but is not guaranteed to
// show every state transition they go through.  Furthermore, state changes
// will only be written to the channel with non-blocking writes.  If updateCh
// is full, it attempts to write an error to errCh; if errCh is full, the error
// passes silently.
func (c *Conn) SetSubStateSubscriber(updateCh chan<- *SubStateUpdate, errCh chan<- error) {
	c.subscriber.Lock()
	defer c.subscriber.Unlock()
	c.subscriber.updateCh = updateCh
	c.subscriber.errCh = errCh
}

func (c *Conn) sendSubStateUpdate(path dbus.ObjectPath) {
	c.subscriber.Lock()
	defer c.subscriber.Unlock()
	if c.subscriber.updateCh == nil {
		return
	}

	if c.shouldIgnore(path) {
		return
	}

	info, err := c.GetUnitProperties(string(path))
	if err != nil {
		select {
		case c.subscriber.errCh <- err:
		default:
		}
	}

	name := info["Id"].(string)
	substate := info["SubState"].(string)

	update := &SubStateUpdate{name, substate}
	select {
	case c.subscriber.updateCh <- update:
	default:
		select {
		case c.subscriber.errCh <- errors.New("update channel full!"):
		default:
		}
	}

	c.updateIgnore(path, info)
}

// The ignore functions work around a wart in the systemd dbus interface.
// Requesting the properties of an unloaded unit will cause systemd to send a
// pair of UnitNew/UnitRemoved signals.  Because we need to get a unit's
// properties on UnitNew (as that's the only indication of a new unit coming up
// for the first time), we would enter an infinite loop if we did not attempt
// to detect and ignore these spurious signals.  The signal themselves are
// indistinguishable from relevant ones, so we (somewhat hackishly) ignore an
// unloaded unit's signals for a short time after requesting its properties.
// This means that we will miss e.g. a transient unit being restarted
// *immediately* upon failure and also a transient unit being started
// immediately after requesting its status (with systemctl status, for example,
// because this causes a UnitNew signal to be sent which then causes us to fetch
// the properties).

func (c *Conn) shouldIgnore(path dbus.ObjectPath) bool {
	t, ok := c.subscriber.ignore[path]
	return ok && t >= time.Now().UnixNano()
}

func (c *Conn) updateIgnore(path dbus.ObjectPath, info map[string]interface{}) {
	c.cleanIgnore()

	// unit is unloaded - it will trigger bad systemd dbus behavior
	if info["LoadState"].(string) == "not-found" {
		c.subscriber.ignore[path] = time.Now().UnixNano() + ignoreInterval
	}
}

// without this, ignore would grow unboundedly over time
func (c *Conn) cleanIgnore() {
	now := time.Now().UnixNano()
	if c.subscriber.cleanIgnore < now {
		c.subscriber.cleanIgnore = now + cleanIgnoreInterval

		for p, t := range c.subscriber.ignore {
			if t < now {
				delete(c.subscriber.ignore, p)
			}
		}
	}
}
<commit_msg>fix(subscription): Fix handling of Close()<commit_after>/*
Copyright 2013 CoreOS Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package dbus

import (
	"errors"
	"time"

	"github.com/godbus/dbus"
)

const (
	cleanIgnoreInterval = int64(10 * time.Second)
	ignoreInterval      = int64(30 * time.Millisecond)
)

// Subscribe sets up this connection to subscribe to all systemd dbus events.
// This is required before calling SubscribeUnits. When the connection closes
// systemd will automatically stop sending signals so there is no need to
// explicitly call Unsubscribe().
func (c *Conn) Subscribe() error {
	c.sysconn.BusObject().Call("org.freedesktop.DBus.AddMatch", 0,
		"type='signal',interface='org.freedesktop.systemd1.Manager',member='UnitNew'")
	c.sysconn.BusObject().Call("org.freedesktop.DBus.AddMatch", 0,
		"type='signal',interface='org.freedesktop.DBus.Properties',member='PropertiesChanged'")

	err := c.sysobj.Call("org.freedesktop.systemd1.Manager.Subscribe", 0).Store()
	if err != nil {
		return err
	}

	return nil
}

// Unsubscribe this connection from systemd dbus events.
func (c *Conn) Unsubscribe() error {
	err := c.sysobj.Call("org.freedesktop.systemd1.Manager.Unsubscribe", 0).Store()
	if err != nil {
		return err
	}

	return nil
}

func (c *Conn) initSubscription() {
	c.subscriber.ignore = make(map[dbus.ObjectPath]int64)
}

func (c *Conn) initDispatch() {
	ch := make(chan *dbus.Signal, signalBuffer)

	c.sysconn.Signal(ch)

	go func() {
		for {
			signal, ok := <-ch
			if !ok {
				return
			}

			switch signal.Name {
			case "org.freedesktop.systemd1.Manager.JobRemoved":
				c.jobComplete(signal)

				unitName := signal.Body[2].(string)
				var unitPath dbus.ObjectPath
				c.sysobj.Call("org.freedesktop.systemd1.Manager.GetUnit", 0, unitName).Store(&unitPath)
				if unitPath != dbus.ObjectPath("") {
					c.sendSubStateUpdate(unitPath)
				}
			case "org.freedesktop.systemd1.Manager.UnitNew":
				c.sendSubStateUpdate(signal.Body[1].(dbus.ObjectPath))
			case "org.freedesktop.DBus.Properties.PropertiesChanged":
				if signal.Body[0].(string) == "org.freedesktop.systemd1.Unit" {
					// we only care about SubState updates, which are a Unit property
					c.sendSubStateUpdate(signal.Path)
				}
			}
		}
	}()
}

// Returns two unbuffered channels which will receive all changed units every
// interval.  Deleted units are sent as nil.
func (c *Conn) SubscribeUnits(interval time.Duration) (<-chan map[string]*UnitStatus, <-chan error) {
	return c.SubscribeUnitsCustom(interval, 0, func(u1, u2 *UnitStatus) bool { return *u1 != *u2 }, nil)
}

// SubscribeUnitsCustom is like SubscribeUnits but lets you specify the buffer
// size of the channels, the comparison function for detecting changes and a filter
// function for cutting down on the noise that your channel receives.
func (c *Conn) SubscribeUnitsCustom(interval time.Duration, buffer int, isChanged func(*UnitStatus, *UnitStatus) bool, filterUnit func (string) bool) (<-chan map[string]*UnitStatus, <-chan error) {
	old := make(map[string]*UnitStatus)
	statusChan := make(chan map[string]*UnitStatus, buffer)
	errChan := make(chan error, buffer)

	go func() {
		for {
			timerChan := time.After(interval)

			units, err := c.ListUnits()
			if err == nil {
				cur := make(map[string]*UnitStatus)
				for i := range units {
					if filterUnit != nil && filterUnit(units[i].Name) {
						continue
					}
					cur[units[i].Name] = &units[i]
				}

				// add all new or changed units
				changed := make(map[string]*UnitStatus)
				for n, u := range cur {
					if oldU, ok := old[n]; !ok || isChanged(oldU, u) {
						changed[n] = u
					}
					delete(old, n)
				}

				// add all deleted units
				for oldN := range old {
					changed[oldN] = nil
				}

				old = cur

				if len(changed) != 0 {
					statusChan <- changed
				}
			} else {
				errChan <- err
			}

			<-timerChan
		}
	}()

	return statusChan, errChan
}

type SubStateUpdate struct {
	UnitName string
	SubState string
}

// SetSubStateSubscriber writes to updateCh when any unit's substate changes.
// Although this writes to updateCh on every state change, the reported state
// may be more recent than the change that generated it (due to an unavoidable
// race in the systemd dbus interface).  That is, this method provides a good
// way to keep a current view of all units' states, but is not guaranteed to
// show every state transition they go through.  Furthermore, state changes
// will only be written to the channel with non-blocking writes.  If updateCh
// is full, it attempts to write an error to errCh; if errCh is full, the error
// passes silently.
func (c *Conn) SetSubStateSubscriber(updateCh chan<- *SubStateUpdate, errCh chan<- error) {
	c.subscriber.Lock()
	defer c.subscriber.Unlock()
	c.subscriber.updateCh = updateCh
	c.subscriber.errCh = errCh
}

func (c *Conn) sendSubStateUpdate(path dbus.ObjectPath) {
	c.subscriber.Lock()
	defer c.subscriber.Unlock()
	if c.subscriber.updateCh == nil {
		return
	}

	if c.shouldIgnore(path) {
		return
	}

	info, err := c.GetUnitProperties(string(path))
	if err != nil {
		select {
		case c.subscriber.errCh <- err:
		default:
		}
	}

	name := info["Id"].(string)
	substate := info["SubState"].(string)

	update := &SubStateUpdate{name, substate}
	select {
	case c.subscriber.updateCh <- update:
	default:
		select {
		case c.subscriber.errCh <- errors.New("update channel full!"):
		default:
		}
	}

	c.updateIgnore(path, info)
}

// The ignore functions work around a wart in the systemd dbus interface.
// Requesting the properties of an unloaded unit will cause systemd to send a
// pair of UnitNew/UnitRemoved signals.  Because we need to get a unit's
// properties on UnitNew (as that's the only indication of a new unit coming up
// for the first time), we would enter an infinite loop if we did not attempt
// to detect and ignore these spurious signals.  The signal themselves are
// indistinguishable from relevant ones, so we (somewhat hackishly) ignore an
// unloaded unit's signals for a short time after requesting its properties.
// This means that we will miss e.g. a transient unit being restarted
// *immediately* upon failure and also a transient unit being started
// immediately after requesting its status (with systemctl status, for example,
// because this causes a UnitNew signal to be sent which then causes us to fetch
// the properties).

func (c *Conn) shouldIgnore(path dbus.ObjectPath) bool {
	t, ok := c.subscriber.ignore[path]
	return ok && t >= time.Now().UnixNano()
}

func (c *Conn) updateIgnore(path dbus.ObjectPath, info map[string]interface{}) {
	c.cleanIgnore()

	// unit is unloaded - it will trigger bad systemd dbus behavior
	if info["LoadState"].(string) == "not-found" {
		c.subscriber.ignore[path] = time.Now().UnixNano() + ignoreInterval
	}
}

// without this, ignore would grow unboundedly over time
func (c *Conn) cleanIgnore() {
	now := time.Now().UnixNano()
	if c.subscriber.cleanIgnore < now {
		c.subscriber.cleanIgnore = now + cleanIgnoreInterval

		for p, t := range c.subscriber.ignore {
			if t < now {
				delete(c.subscriber.ignore, p)
			}
		}
	}
}
<|endoftext|>


<commit_before>package oak

import (
	"bytes"
	"testing"

	"github.com/oakmound/oak/v3/collision"
	"github.com/oakmound/oak/v3/event"
	"github.com/oakmound/oak/v3/mouse"
	"github.com/oakmound/oak/v3/render"
)

type ent struct{}

func (e ent) Init() event.CID {
	return 0
}

func TestDebugConsole(t *testing.T) {
	c1 := NewController()
	c1.config.LoadBuiltinCommands = true
	triggered := false
	err := c1.AddCommand("test", func([]string) {
		triggered = true
	})
	if err != nil {
		t.Fatalf("failed to add test command")
	}

	render.UpdateDebugMap("r", render.EmptyRenderable())

	event.NextID(ent{})

	r := bytes.NewBufferString(
		"test\n" +
			"nothing\n" +
			"fade nothing\n" +
			"fade nothing 100\n" +
			"fade r\n" +
			"skip nothing\n" +
			"print nothing\n" +
			"print 2\n" +
			"print 1\n" +
			"mouse nothing\n" +
			"mouse details\n" +
			"garbage input\n" +
			"\n" +
			"skip scene\n")
	go c1.debugConsole(r)
	sleep()
	sleep()
	if !triggered {
		t.Fatalf("debug console did not trigger test command")
	}
	<-c1.skipSceneCh
}

func TestMouseDetails(t *testing.T) {
	c1 := NewController()

	c1.mouseDetails(0, mouse.NewZeroEvent(0, 0))
	s := collision.NewUnassignedSpace(-1, -1, 2, 2)
	collision.Add(s)
	c1.mouseDetails(0, mouse.NewZeroEvent(0, 0))
	collision.Remove(s)

	// This should spew this nothing entity, but it doesn't.
	id := event.NextID(ent{})
	s = collision.NewSpace(-1, -1, 2, 2, id)
	c1.mouseDetails(0, mouse.NewZeroEvent(0, 0))
	collision.Remove(s)
}
<commit_msg>oak: correct debug console tests to use pointer mouse events<commit_after>package oak

import (
	"bytes"
	"testing"

	"github.com/oakmound/oak/v3/collision"
	"github.com/oakmound/oak/v3/event"
	"github.com/oakmound/oak/v3/mouse"
	"github.com/oakmound/oak/v3/render"
)

type ent struct{}

func (e ent) Init() event.CID {
	return 0
}

func TestDebugConsole(t *testing.T) {
	c1 := NewController()
	c1.config.LoadBuiltinCommands = true
	triggered := false
	err := c1.AddCommand("test", func([]string) {
		triggered = true
	})
	if err != nil {
		t.Fatalf("failed to add test command")
	}

	render.UpdateDebugMap("r", render.EmptyRenderable())

	event.NextID(ent{})

	r := bytes.NewBufferString(
		"test\n" +
			"nothing\n" +
			"fade nothing\n" +
			"fade nothing 100\n" +
			"fade r\n" +
			"skip nothing\n" +
			"print nothing\n" +
			"print 2\n" +
			"print 1\n" +
			"mouse nothing\n" +
			"mouse details\n" +
			"garbage input\n" +
			"\n" +
			"skip scene\n")
	go c1.debugConsole(r)
	sleep()
	sleep()
	if !triggered {
		t.Fatalf("debug console did not trigger test command")
	}
	<-c1.skipSceneCh
}

func TestMouseDetails(t *testing.T) {
	c1 := NewController()

	ev := mouse.NewZeroEvent(0, 0)
	c1.mouseDetails(0, &ev)
	s := collision.NewUnassignedSpace(-1, -1, 2, 2)
	collision.Add(s)
	c1.mouseDetails(0, &ev)
	collision.Remove(s)

	// This should spew this nothing entity, but it doesn't.
	id := event.NextID(ent{})
	s = collision.NewSpace(-1, -1, 2, 2, id)
	c1.mouseDetails(0, &ev)
	collision.Remove(s)
}
<|endoftext|>


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

package diam

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"time"
	"unsafe"
)

type AVP struct {
	Code     uint32
	Flags    uint8
	Length   uint32
	VendorId uint32
	Data     interface{}
	RawData  []byte
	Padding  int
}

type Grouped []*AVP

type rfcHdr1 struct {
	Code   uint32
	Flags  uint8
	Length [3]uint8
}

type rfcHdr2 struct {
	Code     uint32
	Flags    uint8
	Length   [3]uint8
	VendorId uint32
}

// ReadAVP reads an AVP and returns the number of extra bytes read and parsed
// AVP, or an error.
// Extra bytes are read when the content of the AVP is OctetString and
// needs padding. Total bytes read is avp.Length + extra.
func ReadAVP(r io.Reader, dict *Dict) (uint32, *AVP, error) {
	var (
		err error
		raw rfcHdr1
	)
	if err = binary.Read(r, binary.BigEndian, &raw); err != nil {
		return 0, nil, err
	}
	avp := &AVP{
		Code:   raw.Code,
		Flags:  raw.Flags,
		Length: uint24to32(raw.Length),
	}
	dlen := avp.Length - uint32(unsafe.Sizeof(raw))
	// Read VendorId when necessary.
	if raw.Flags&0x20 > 0 {
		if err = binary.Read(r, binary.BigEndian, &avp.VendorId); err != nil {
			return 0, nil, err
		}
		dlen -= uint32(unsafe.Sizeof(avp.VendorId))
	}
	// Find this AVP in a pre-loaded dict so we know how to parse it,
	// pad it, or load grouped AVPs from it's data.
	var davp *DictAVP
	if davp, err = dict.AVP(avp.Code); err != nil {
		return 0, nil, fmt.Errorf("Unknown AVP code %d: missing dict?", avp.Code)
	}
	// Read grouped (embedded) AVPs.
	// Grouped AVPs are the only reason why this function returns
	// "extra" bytes, otherwise callers would have to walk through
	// these grouped AVPs and sum their padding + length to figure
	// out the total of bytes read.
	// TODO: Handle dynamically grouped AVPs in case of 260, 279 or 284.
	if davp.Data.Type == "Grouped" {
		if eb, gavp, err := ReadAVP(r, dict); err != nil {
			return 0, nil, err
		} else {
			if avp.Data == nil {
				avp.Data = Grouped{}
			}
			avp.Data = append(avp.Data.(Grouped), gavp)
			return eb, avp, nil
		}
	}
	var needPadding bool
	// Read binary AVP data.
	avp.Data = make([]byte, dlen)
	if err = binary.Read(r, binary.BigEndian, avp.Data); err != nil {
		return 0, nil, err
	}
	switch davp.Data.Type {
	case "Address":
		needPadding = true
		// TODO: Double check this.
		if len(avp.Data.([]byte)) == 6 {
			b := avp.Data.([]byte)
			if b[1] == 1 { // Family address 1 == IPv4
				avp.Data = net.IPv4(b[2], b[3], b[4], b[5])
			}
			break
		}
		// Fallback to string instead of IPv4.
		avp.Data = string(avp.Data.([]byte))
	case "Time":
		fallthrough
	case "DiameterIdentity":
		fallthrough
	case "DiameterURI":
		fallthrough
	case "OctetString":
		fallthrough
	case "UTF8String":
		needPadding = true
		avp.Data = string(avp.Data.([]byte))
	case "Enumerated":
		fallthrough
	case "Unsigned32":
		if dlen != 4 {
			return 0, nil, fmt.Errorf("Expecting Unsigned32, got %d instead", dlen*8)
		}
		avp.Data = bytes2uint32(avp.Data.([]byte))
	case "Unsigned64":
		if dlen != 8 {
			return 0, nil, fmt.Errorf("Expecting Unsigned64, got %d instead", dlen*8)
		}
		avp.Data = bytes2uint64(avp.Data.([]byte))
	}
	// Check if there's extra data to read due to padding of OctetString.
	//
	// http://tools.ietf.org/html/rfc3588#section-4
	//
	// Each AVP of type OctetString MUST be padded to align on a 32-bit
	// boundary, while other AVP types align naturally.  A number of zero-
	// valued bytes are added to the end of the AVP Data field till a word
	// boundary is reached.  The length of the padding is not reflected in
	// the AVP Length field.
	//
	// This also applies to subtypes of OctetString such as Address.
	var extrabytes uint32
	if needPadding {
		// Read and discard pad bytes.
		if avp.Padding = int(pad4(dlen) - dlen); avp.Padding > 0 {
			extrabytes += uint32(avp.Padding)
			b := make([]byte, avp.Padding)
			if _, err = io.ReadFull(r, b); err != nil {
				return 0, nil, err
			}
		}
	}
	return extrabytes, avp, nil
}

// String returns the AVP in human readable format.
func (avp *AVP) String() string {
	// TODO: Lookup the vendor id from AVP in the dictionary.
	var name string
	if davp, err := BaseDict.AVP(avp.Code); err != nil {
		name = "Unknown"
	} else {
		name = davp.Name
	}
	v := fmt.Sprintf("%s AVP{Code=%d,Flags=%#x,Length=%d,VendorId=%#x,Padding=%d,",
		name, avp.Code, avp.Flags, avp.Length, avp.VendorId, avp.Padding)
	switch avp.Data.(type) {
	case string:
		v += fmt.Sprintf("string('%s')", avp.Data.(string))
	case DictEnumItem:
		v += fmt.Sprintf("enum(%s)", avp.Data.(DictEnumItem))
	case time.Time:
		v += fmt.Sprintf("time(%s)", avp.Data.(time.Time))
	case uint32:
		v += fmt.Sprintf("uint32(%d)", avp.Data.(uint32))
	case uint64:
		v += fmt.Sprintf("uint64(%d)", avp.Data.(uint64))
	case net.IP:
		v += fmt.Sprintf("net.IP(%s)", avp.Data.(net.IP))
	case Grouped:
		v += fmt.Sprintf("Grouped(%s)", avp.Data.(Grouped))
	default:
		v += fmt.Sprintf("Unknown(%s)", avp.Data)
	}
	return v + "}"
}

// NewAVP allocates and returns a new AVP. Used for building messages.
func NewAVP(code uint32, flags uint8, vendor uint32, data interface{}) *AVP {
	avp := &AVP{
		Code:     code,
		Flags:    flags,
		VendorId: vendor,
		Data:     data,
	}
	if flags&0x20 > 0 {
		avp.Length = uint32(unsafe.Sizeof(rfcHdr2{}))
	} else {
		avp.Length = uint32(unsafe.Sizeof(rfcHdr1{}))
	}
	switch data.(type) {
	case string:
		s := data.(string)
		sl := uint32(len(s))
		avp.Length += sl
		if extra := pad4(sl) - sl; extra > 0 {
			avp.Padding = int(extra)
			avp.RawData = make([]byte, sl+extra)
			copy(avp.RawData, s)
		} else {
			avp.RawData = []byte(s)
		}
	case net.IP:
		avp.Length += 6
		ip := data.(net.IP).To4()
		// TODO: Fix this static family 1 + IPv4 address.
		s := []byte{0, 1, ip[0], ip[1], ip[2], ip[3]}
		sl := uint32(len(s))
		if extra := pad4(sl) - sl; extra > 0 {
			avp.Padding = int(extra)
			avp.RawData = make([]byte, sl+extra)
			copy(avp.RawData, s)
		} else {
			avp.RawData = []byte(s)
		}
	case uint32:
		avp.Length += 4
		avp.RawData = uint32tobytes(data.(uint32))
	case uint64:
		avp.Length += 8
		avp.RawData = uint64tobytes(data.(uint64))
	}
	return avp
}

// Marshal returns an AVP in binary form so it can be attached to a Message
// before sent to a connection.
func (avp *AVP) Marshal() []byte {
	b := bytes.NewBuffer(make([]byte, 0))
	if avp.Flags&0x20 > 0 {
		hdr := rfcHdr2{
			Code:     avp.Code,
			Flags:    avp.Flags,
			Length:   uint32to24(avp.Length),
			VendorId: avp.VendorId,
		}
		binary.Write(b, binary.BigEndian, hdr)
	} else {
		hdr := rfcHdr1{
			Code:   avp.Code,
			Flags:  avp.Flags,
			Length: uint32to24(avp.Length),
		}
		binary.Write(b, binary.BigEndian, hdr)
	}
	binary.Write(b, binary.BigEndian, avp.RawData)
	return b.Bytes()
}
<commit_msg>Optimizations<commit_after>// Copyright 2013 Alexandre Fiori
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package diam

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"time"
	"unsafe"
)

type AVP struct {
	Code     uint32
	Flags    uint8
	Length   uint32
	VendorId uint32
	Data     interface{}
	RawData  []byte
	Padding  int
}

type Grouped []*AVP

type rfcHdr1 struct {
	Code   uint32
	Flags  uint8
	Length [3]uint8
}

type rfcHdr2 struct {
	Code     uint32
	Flags    uint8
	Length   [3]uint8
	VendorId uint32
}

// ReadAVP reads an AVP and returns the number of extra bytes read and parsed
// AVP, or an error.
// Extra bytes are read when the content of the AVP is OctetString and
// needs padding. Total bytes read is avp.Length + extra.
func ReadAVP(r io.Reader, dict *Dict) (uint32, *AVP, error) {
	var (
		err error
		raw rfcHdr1
	)
	if err = binary.Read(r, binary.BigEndian, &raw); err != nil {
		return 0, nil, err
	}
	avp := &AVP{
		Code:   raw.Code,
		Flags:  raw.Flags,
		Length: uint24to32(raw.Length),
	}
	dlen := avp.Length - uint32(unsafe.Sizeof(raw))
	// Read VendorId when necessary.
	if raw.Flags&0x20 > 0 {
		if err = binary.Read(r, binary.BigEndian, &avp.VendorId); err != nil {
			return 0, nil, err
		}
		dlen -= uint32(unsafe.Sizeof(avp.VendorId))
	}
	// Find this AVP in a pre-loaded dict so we know how to parse it,
	// pad it, or load grouped AVPs from it's data.
	var davp *DictAVP
	if davp, err = dict.AVP(avp.Code); err != nil {
		return 0, nil, fmt.Errorf("Unknown AVP code %d: missing dict?", avp.Code)
	}
	// Read grouped (embedded) AVPs.
	// Grouped AVPs are the only reason why this function returns
	// "extra" bytes, otherwise callers would have to walk through
	// these grouped AVPs and sum their padding + length to figure
	// out the total of bytes read.
	// TODO: Handle dynamically grouped AVPs in case of 260, 279 or 284.
	if davp.Data.Type == "Grouped" {
		if eb, gavp, err := ReadAVP(r, dict); err != nil {
			return 0, nil, err
		} else {
			if avp.Data == nil {
				avp.Data = Grouped{}
			}
			avp.Data = append(avp.Data.(Grouped), gavp)
			return eb, avp, nil
		}
	}
	var needPadding bool
	// Read binary AVP data.
	avp.Data = make([]byte, dlen)
	if err = binary.Read(r, binary.BigEndian, avp.Data); err != nil {
		return 0, nil, err
	}
	switch davp.Data.Type {
	case "Address":
		needPadding = true
		// TODO: Double check this.
		if len(avp.Data.([]byte)) == 6 {
			b := avp.Data.([]byte)
			if b[1] == 1 { // Family address 1 == IPv4
				avp.Data = net.IPv4(b[2], b[3], b[4], b[5])
			}
			break
		}
		// Fallback to string instead of IPv4.
		avp.Data = string(avp.Data.([]byte))
	case "Time":
		fallthrough
	case "DiameterIdentity":
		fallthrough
	case "DiameterURI":
		fallthrough
	case "OctetString":
		fallthrough
	case "UTF8String":
		needPadding = true
		avp.Data = string(avp.Data.([]byte))
	case "Enumerated":
		fallthrough
	case "Unsigned32":
		if dlen != 4 {
			return 0, nil, fmt.Errorf("Expecting Unsigned32, got %d instead", dlen*8)
		}
		avp.Data = bytes2uint32(avp.Data.([]byte))
	case "Unsigned64":
		if dlen != 8 {
			return 0, nil, fmt.Errorf("Expecting Unsigned64, got %d instead", dlen*8)
		}
		avp.Data = bytes2uint64(avp.Data.([]byte))
	}
	// Check if there's extra data to read due to padding of OctetString.
	//
	// http://tools.ietf.org/html/rfc3588#section-4
	//
	// Each AVP of type OctetString MUST be padded to align on a 32-bit
	// boundary, while other AVP types align naturally.  A number of zero-
	// valued bytes are added to the end of the AVP Data field till a word
	// boundary is reached.  The length of the padding is not reflected in
	// the AVP Length field.
	//
	// This also applies to subtypes of OctetString such as Address.
	var extrabytes uint32
	if needPadding {
		// Read and discard pad bytes.
		if avp.Padding = int(pad4(dlen) - dlen); avp.Padding > 0 {
			extrabytes += uint32(avp.Padding)
			b := make([]byte, avp.Padding)
			if _, err = io.ReadFull(r, b); err != nil {
				return 0, nil, err
			}
		}
	}
	return extrabytes, avp, nil
}

// String returns the AVP in human readable format.
func (avp *AVP) String() string {
	// TODO: Lookup the vendor id from AVP in the dictionary.
	var name string
	if davp, err := BaseDict.AVP(avp.Code); err != nil {
		name = "Unknown"
	} else {
		name = davp.Name
	}
	v := fmt.Sprintf("%s AVP{Code=%d,Flags=%#x,Length=%d,VendorId=%#x,Padding=%d,",
		name, avp.Code, avp.Flags, avp.Length, avp.VendorId, avp.Padding)
	switch avp.Data.(type) {
	case string:
		v += fmt.Sprintf("string('%s')", avp.Data.(string))
	case DictEnumItem:
		v += fmt.Sprintf("enum(%s)", avp.Data.(DictEnumItem))
	case time.Time:
		v += fmt.Sprintf("time(%s)", avp.Data.(time.Time))
	case uint32:
		v += fmt.Sprintf("uint32(%d)", avp.Data.(uint32))
	case uint64:
		v += fmt.Sprintf("uint64(%d)", avp.Data.(uint64))
	case net.IP:
		v += fmt.Sprintf("net.IP(%s)", avp.Data.(net.IP))
	case Grouped:
		v += fmt.Sprintf("Grouped(%s)", avp.Data.(Grouped))
	default:
		v += fmt.Sprintf("Unknown(%s)", avp.Data)
	}
	return v + "}"
}

// NewAVP allocates and returns a new AVP. Used for building messages.
func NewAVP(code uint32, flags uint8, vendor uint32, data interface{}) *AVP {
	avp := &AVP{
		Code:     code,
		Flags:    flags,
		VendorId: vendor,
		Data:     data,
	}
	if flags&0x20 > 0 {
		avp.Length = uint32(unsafe.Sizeof(rfcHdr2{}))
	} else {
		avp.Length = uint32(unsafe.Sizeof(rfcHdr1{}))
	}
	switch data.(type) {
	case string:
		avpString(avp, []byte(data.(string)))
	case net.IP:
		ip := data.(net.IP).To4()
		// TODO: Fix this static family 1 + IPv4 address.
		avpString(avp, []byte{0, 1, ip[0], ip[1], ip[2], ip[3]})
	case uint32:
		avp.Length += 4
		avp.RawData = uint32tobytes(data.(uint32))
	case uint64:
		avp.Length += 8
		avp.RawData = uint64tobytes(data.(uint64))
	}
	return avp
}

// avpString encodes the AVP data with padding if needed, and updates the
// Length header accordingly. It also sets the Padding attribute.
func avpString(avp *AVP, s []byte) {
	length := uint32(len(s))
	avp.Length += length
	if extra := pad4(length) - length; extra > 0 {
		avp.Padding = int(extra)
		avp.RawData = make([]byte, length+extra)
		copy(avp.RawData, s)
	} else {
		avp.RawData = s
	}
}

// Marshal returns an AVP in binary form so it can be attached to a Message
// before sent to a connection.
func (avp *AVP) Marshal() []byte {
	b := bytes.NewBuffer(make([]byte, 0))
	if avp.Flags&0x20 > 0 {
		hdr := rfcHdr2{
			Code:     avp.Code,
			Flags:    avp.Flags,
			Length:   uint32to24(avp.Length),
			VendorId: avp.VendorId,
		}
		binary.Write(b, binary.BigEndian, hdr)
	} else {
		hdr := rfcHdr1{
			Code:   avp.Code,
			Flags:  avp.Flags,
			Length: uint32to24(avp.Length),
		}
		binary.Write(b, binary.BigEndian, hdr)
	}
	binary.Write(b, binary.BigEndian, avp.RawData)
	return b.Bytes()
}
<|endoftext|>


<commit_before>// This is a wrapper of standard http package, exposing out-of-the-box flags like proxy, timeout, json handling.
package goutils

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"

	"github.com/hoveychen/go-utils/flags"
	"github.com/pkg/errors"

	"golang.org/x/net/proxy"
)

var (
	proxyAddr      = flags.String("proxy", "", "Specify proxy address to fetch data")
	proxyType      = flags.String("proxyType", "sock5", "Either sock5 or http for proxy.")
	requestTimeout = flags.Int("requestTimeout", 10, "Timeout in sec when fetching a remote page.")
	downloadClient *http.Client
)

func init() {
	PkgInit(func() {
		httpTransport := &http.Transport{}
		if *proxyAddr != "" {
			switch *proxyType {
			case "http":
				proxyUrl, err := url.Parse(*proxyAddr)
				if err != nil {
					LogFatal("Failed to parse --proxyAddr", err)
				}
				httpTransport.Proxy = http.ProxyURL(proxyUrl)
			case "sock5":
				dialer, err := proxy.SOCKS5("tcp", *proxyAddr, nil, proxy.Direct)
				if err != nil {
					LogFatal("Failed to dial sock5", err)
				}
				httpTransport.Dial = dialer.Dial
			default:
				LogFatal("Unknown proxy type:", *proxyType)
			}
		}

		downloadClient = &http.Client{Transport: httpTransport}
		downloadClient.Timeout = time.Duration(*requestTimeout) * time.Second
		downloadClient.CheckRedirect = modifiedCheckRedirect
	})
}

func modifiedCheckRedirect(req *http.Request, via []*http.Request) error {
	if len(via) >= 15 {
		return errors.New("stopped after 15 redirects")
	}
	return nil
}

func Get(url string) (*http.Response, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, errors.Wrap(err, "New get request")
	}
	resp, err := downloadClient.Do(req)
	if err != nil {
		return nil, errors.Wrap(err, "Do get request")
	}
	return resp, nil
}

func PostForm(uri string, data map[string]string) (*http.Response, error) {
	values := url.Values{}
	for k, v := range data {
		values.Set(k, v)
	}
	req, err := http.NewRequest("POST", uri, bytes.NewBufferString(values.Encode()))
	if err != nil {
		return nil, errors.Wrap(err, "New post request")
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := downloadClient.Do(req)
	if err != nil {
		return nil, errors.Wrap(err, "Do post request")
	}
	return resp, nil
}

func PostJson(url string, data interface{}) (*http.Response, error) {
	encodedData, err := json.Marshal(data)
	if err != nil {
		return nil, errors.Wrap(err, "Encode json")
	}
	req, err := http.NewRequest("POST", url, bytes.NewBuffer(encodedData))
	if err != nil {
		return nil, errors.Wrap(err, "New post request")
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := downloadClient.Do(req)
	if err != nil {
		return nil, errors.Wrap(err, "Do post request")
	}
	return resp, nil
}

// FetchData is a helper function to load local/remote data in the same function.
// Local: goutils.FetchData("/absolute/path/to/file")
// Remote: goutils.FetchData("https://www.google.com")
// Also, it's integrated with proxy in flags.
// TODO(yuheng): Allow more options, while keeping easy use.
func FetchData(path string) ([]byte, error) {
	url, err := url.Parse(path)
	if err != nil {
		return nil, errors.Wrap(err, "Decode url")
	}

	switch url.Scheme {
	case "http", "https":
		resp, err := Get(path)
		if err != nil {
			return nil, errors.Wrap(err, "Http get")
		}
		defer resp.Body.Close()
		data, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, errors.Wrap(err, "Read response")
		}
		if resp.StatusCode != http.StatusOK {
			return nil, errors.New(resp.Status + ":" + string(data))
		}

		return data, nil
	case "":
		data, err := ioutil.ReadFile(url.Path)
		if err != nil {
			return nil, errors.Wrap(err, "Read file")
		}
		return data, nil
	default:
		return nil, errors.New("Unknown scheme:" + url.Scheme)
	}
}

// FetchJson is a wrapper to call FetchData() and parse results from json.
func FetchJson(path string, resp interface{}) error {
	d, err := FetchData(path)
	if err != nil {
		return errors.Wrap(err, "Fetch data")
	}
	if err := json.Unmarshal(d, resp); err != nil {
		return errors.Wrap(err, "Decode json")
	}
	return nil
}

// FetchXml is a wrapper to call FetchData() and parse results from xml.
func FetchXml(path string, resp interface{}) error {
	d, err := FetchData(path)
	if err != nil {
		return errors.Wrap(err, "Fetch data")
	}
	if err := xml.Unmarshal(d, resp); err != nil {
		return errors.Wrap(err, "Decode xml")
	}
	return nil
}
<commit_msg>[Request] Export download client for complex request operations.<commit_after>// This is a wrapper of standard http package, exposing out-of-the-box flags like proxy, timeout, json handling.
package goutils

import (
	"bytes"
	"encoding/json"
	"encoding/xml"
	"io/ioutil"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/hoveychen/go-utils/flags"
	"github.com/pkg/errors"

	"golang.org/x/net/proxy"
)

var (
	proxyAddr      = flags.String("proxy", "", "Specify proxy address to fetch data")
	proxyType      = flags.String("proxyType", "sock5", "Either sock5 or http for proxy.")
	requestTimeout = flags.Int("requestTimeout", 10, "Timeout in sec when fetching a remote page.")
	downloadClient *http.Client
	requestOnce    sync.Once
)

func modifiedCheckRedirect(req *http.Request, via []*http.Request) error {
	if len(via) >= 15 {
		return errors.New("stopped after 15 redirects")
	}
	return nil
}

func GetDownloadClient() *http.Client {
	requestOnce.Do(func() {
		httpTransport := &http.Transport{}
		if *proxyAddr != "" {
			switch *proxyType {
			case "http":
				proxyUrl, err := url.Parse(*proxyAddr)
				if err != nil {
					LogFatal("Failed to parse --proxyAddr", err)
				}
				httpTransport.Proxy = http.ProxyURL(proxyUrl)
			case "sock5":
				dialer, err := proxy.SOCKS5("tcp", *proxyAddr, nil, proxy.Direct)
				if err != nil {
					LogFatal("Failed to dial sock5", err)
				}
				httpTransport.Dial = dialer.Dial
			default:
				LogFatal("Unknown proxy type:", *proxyType)
			}
		}

		downloadClient = &http.Client{Transport: httpTransport}
		downloadClient.Timeout = time.Duration(*requestTimeout) * time.Second
		downloadClient.CheckRedirect = modifiedCheckRedirect
	})

	return downloadClient
}

func Get(url string) (*http.Response, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, errors.Wrap(err, "New get request")
	}
	resp, err := GetDownloadClient().Do(req)
	if err != nil {
		return nil, errors.Wrap(err, "Do get request")
	}
	return resp, nil
}

func PostForm(uri string, data map[string]string) (*http.Response, error) {
	values := url.Values{}
	for k, v := range data {
		values.Set(k, v)
	}
	req, err := http.NewRequest("POST", uri, bytes.NewBufferString(values.Encode()))
	if err != nil {
		return nil, errors.Wrap(err, "New post request")
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := GetDownloadClient().Do(req)
	if err != nil {
		return nil, errors.Wrap(err, "Do post request")
	}
	return resp, nil
}

func PostJson(url string, data interface{}) (*http.Response, error) {
	encodedData, err := json.Marshal(data)
	if err != nil {
		return nil, errors.Wrap(err, "Encode json")
	}
	req, err := http.NewRequest("POST", url, bytes.NewBuffer(encodedData))
	if err != nil {
		return nil, errors.Wrap(err, "New post request")
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := GetDownloadClient().Do(req)
	if err != nil {
		return nil, errors.Wrap(err, "Do post request")
	}
	return resp, nil
}

// FetchData is a helper function to load local/remote data in the same function.
// Local: goutils.FetchData("/absolute/path/to/file")
// Remote: goutils.FetchData("https://www.google.com")
// Also, it's integrated with proxy in flags.
// TODO(yuheng): Allow more options, while keeping easy use.
func FetchData(path string) ([]byte, error) {
	url, err := url.Parse(path)
	if err != nil {
		return nil, errors.Wrap(err, "Decode url")
	}

	switch url.Scheme {
	case "http", "https":
		resp, err := Get(path)
		if err != nil {
			return nil, errors.Wrap(err, "Http get")
		}
		defer resp.Body.Close()
		data, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, errors.Wrap(err, "Read response")
		}
		if resp.StatusCode != http.StatusOK {
			return nil, errors.New(resp.Status + ":" + string(data))
		}

		return data, nil
	case "":
		data, err := ioutil.ReadFile(url.Path)
		if err != nil {
			return nil, errors.Wrap(err, "Read file")
		}
		return data, nil
	default:
		return nil, errors.New("Unknown scheme:" + url.Scheme)
	}
}

// FetchJson is a wrapper to call FetchData() and parse results from json.
func FetchJson(path string, resp interface{}) error {
	d, err := FetchData(path)
	if err != nil {
		return errors.Wrap(err, "Fetch data")
	}
	if err := json.Unmarshal(d, resp); err != nil {
		return errors.Wrap(err, "Decode json")
	}
	return nil
}

// FetchXml is a wrapper to call FetchData() and parse results from xml.
func FetchXml(path string, resp interface{}) error {
	d, err := FetchData(path)
	if err != nil {
		return errors.Wrap(err, "Fetch data")
	}
	if err := xml.Unmarshal(d, resp); err != nil {
		return errors.Wrap(err, "Decode xml")
	}
	return nil
}
<|endoftext|>


<commit_before>package kite

import (
	"errors"
	"fmt"
	"kite/dnode"
	"kite/dnode/rpc"
	"kite/kitekey"
	"reflect"
	"runtime/debug"
	"time"

	"github.com/dgrijalva/jwt-go"
)

// runMethod is called when a method is received from remote Kite.
func runMethod(method string, handlerFunc reflect.Value, args dnode.Arguments, tr dnode.Transport) {
	var (
		// The request that will be constructed from incoming dnode message.
		request *Request

		// Will hold the return values from handler func.
		values []reflect.Value

		// First value to the response.
		result interface{}

		// Second value to the response.
		kiteErr *Error

		// Will send the response when called.
		callback dnode.Function
	)

	kite := tr.Properties()["localKite"].(*Kite)

	// Send result if "responseCallback" exists in the request.
	defer func() {
		if callback == nil {
			return
		}

		// Only argument to the callback.
		response := callbackArg{
			Result: result,
			Error:  errorForSending(kiteErr),
		}

		// Call response callback function.
		if err := callback(response); err != nil {
			kite.Log.Error(err.Error())
		}
	}()

	// Recover dnode argument errors. The caller can use functions like
	// MustString(), MustSlice()... without the fear of panic.
	defer kite.recoverError(&kiteErr)()

	request, callback = kite.parseRequest(method, args, tr)

	if !kite.disableAuthenticate {
		kiteErr = request.authenticate()
		if kiteErr != nil {
			return
		}
	}

	// Call the handler function.
	callArgs := []reflect.Value{reflect.ValueOf(request)}
	values = handlerFunc.Call(callArgs)

	result = values[0].Interface()

	if err := values[1].Interface(); err != nil {
		panic(err) // This will be recoverd from kite.recoverError() above.
	}
}

// Error is the type of the kite related errors returned from kite package.
type Error struct {
	Type    string `json:"type"`
	Message string `json:"message"`
}

func (e Error) Error() string {
	return fmt.Sprintf("kite error %s - %s", e.Type, e.Message)
}

// When a callback is called we always pass this as the only argument.
type callbackArg struct {
	Error  errorForSending `json:"error"`
	Result interface{}     `json:"result"`
}

// errorForSending is a for sending the error as an argument in a dnode message.
// Normally Error method of the Error struct is sent as a callback since it is
// exported. We do not want this behavior.
type errorForSending *Error

// recoverError returns a function which recovers the error and sets to the
// given argument as kite.Error.
//
// TODO: change that it doesn't use a pointer of pointer,
// a simpler and cleaner solution would work in the future,
func (k *Kite) recoverError(kiteErr **Error) func() {
	return func() {
		r := recover()
		if r == nil {
			return
		}

		switch err := r.(type) {
		case *Error:
			*kiteErr = err
		case *dnode.ArgumentError:
			*kiteErr = &Error{"argumentError", err.Error()}
		default:
			*kiteErr = &Error{"genericError", fmt.Sprint(r)}
			debug.PrintStack()
		}

		k.Log.Warning("Error in received message %s", (*kiteErr).Error())

	}
}

// HandlerFunc is the type of the handlers registered to Kite.
// The returned result must be Marshalable with json package.
type HandlerFunc func(*Request) (result interface{}, err error)

// HandleFunc registers a handler to run when a method call is received from a Kite.
func (k *Kite) HandleFunc(method string, handler HandlerFunc) {
	k.server.HandleFunc(method, handler)
}

// Request contains information about the incoming request.
type Request struct {
	Method         string
	Args           dnode.Arguments
	LocalKite      *Kite
	RemoteKite     *RemoteKite
	Username       string
	Authentication Authentication
	RemoteAddr     string
}

// Wrap your function with Callback to send it as an argument to a RemoteKite.
type Callback func(r *Request)

func (c Callback) MarshalJSON() ([]byte, error) {
	return []byte(`"[Function]"`), nil
}

// runCallback is called when a callback method call is received from remote Kite.
func runCallback(method string, handlerFunc reflect.Value, args dnode.Arguments, tr dnode.Transport) {
	kite := tr.Properties()["localKite"].(*Kite)

	kiteErr := new(Error)               // Not used. For recovering the error.
	defer kite.recoverError(&kiteErr)() // Do not panic no matter what.

	request, _ := kite.parseRequest(method, args, tr)

	// Call the callback function.
	callArgs := []reflect.Value{reflect.ValueOf(request)}
	handlerFunc.Call(callArgs) // No return value from callback function.
}

// parseRequest is used to read a dnode message.
// It is called when a method or callback is received to parse the message.
func (k *Kite) parseRequest(method string, arguments dnode.Arguments, tr dnode.Transport) (*Request, dnode.Function) {
	// Parse dnode method arguments: [options]
	var options callOptions
	arguments.One().MustUnmarshal(&options)

	// Properties about the client...
	properties := tr.Properties()

	// Create a new RemoteKite instance to pass it to the handler, so
	// the handler can call methods on the other site on the same connection.
	var remoteKite *RemoteKite
	if properties["remoteKite"] == nil {
		// Do not create a new RemoteKite on every request,
		// cache it in Transport.Properties().
		client := tr.(*rpc.Client) // We only have a dnode/rpc.Client for now.
		remoteKite = k.newRemoteKiteWithClient(options.Kite, options.Authentication, client)
		properties["remoteKite"] = remoteKite

		// Notify Kite.OnConnect handlers.
		k.notifyRemoteKiteConnected(remoteKite)
	} else {
		remoteKite = properties["remoteKite"].(*RemoteKite)

		// Update Kite info but do not change the value of username.
		username := remoteKite.Kite.Username
		remoteKite.Kite = options.Kite
		remoteKite.Kite.Username = username

		// TODO can't set Authentication. Fix it.
		// remoteKite.Authentication = options.Authentication
	}

	request := &Request{
		Method:         method,
		Args:           options.WithArgs,
		LocalKite:      k,
		RemoteKite:     remoteKite,
		RemoteAddr:     tr.RemoteAddr(),
		Authentication: options.Authentication,
	}

	return request, options.ResponseCallback
}

// authenticate tries to authenticate the user by selecting appropriate
// authenticator function.
func (r *Request) authenticate() *Error {
	// Trust the Kite if we have initiated the connection.
	// RemoteAddr() returns "" if this is an outgoing connection.
	if r.RemoteAddr == "" {
		return nil
	}

	// Select authenticator function.
	f := r.LocalKite.Authenticators[r.Authentication.Type]
	if f == nil {
		return &Error{
			Type:    "authenticationError",
			Message: fmt.Sprintf("Unknown authentication type: %s", r.Authentication.Type),
		}
	}

	// Call authenticator function. It sets the Request.Username field.
	err := f(r)
	if err != nil {
		return &Error{
			Type:    "authenticationError",
			Message: err.Error(),
		}
	}

	// Fix username of the remote Kite if it is invalid.
	// This prevents a Kite to impersonate someone else's Kite.
	r.RemoteKite.Kite.Username = r.Username

	return nil
}

// AuthenticateFromToken is the default Authenticator for Kite.
func (k *Kite) AuthenticateFromToken(r *Request) error {
	token, err := jwt.Parse(r.Authentication.Key, r.LocalKite.getRSAKey)
	if err != nil {
		return err
	}

	if audience, ok := token.Claims["aud"].(string); !ok || audience != k.ID {
		return errors.New("Invalid audience in token")
	}

	if expiration, ok := token.Claims["exp"].(float64); !ok || int64(expiration) < time.Now().UTC().Unix() {
		return errors.New("Token is expired")
	}

	if notBefore, ok := token.Claims["nbf"].(float64); !ok || int64(notBefore) > time.Now().UTC().Unix() {
		return errors.New("Token is not valid yet")
	}

	if username, ok := token.Claims["sub"].(string); !ok {
		return errors.New("Username is not present in token")
	} else {
		r.Username = username
	}

	return nil
}

// AuthenticateFromKiteKey authenticates user from Koding Key.
// Kontrol makes requests with a Koding Key.
func (k *Kite) AuthenticateFromKiteKey(r *Request) error {
	token, err := jwt.Parse(r.Authentication.Key, kitekey.GetKontrolKey)
	if err != nil {
		return err
	}

	if username, ok := token.Claims["sub"].(string); !ok {
		return errors.New("Username is not present in token")
	} else {
		r.Username = username
	}

	return nil
}
<commit_msg>check validity of the signature in jwt tokens<commit_after>package kite

import (
	"errors"
	"fmt"
	"kite/dnode"
	"kite/dnode/rpc"
	"kite/kitekey"
	"reflect"
	"runtime/debug"
	"time"

	"github.com/dgrijalva/jwt-go"
)

// runMethod is called when a method is received from remote Kite.
func runMethod(method string, handlerFunc reflect.Value, args dnode.Arguments, tr dnode.Transport) {
	var (
		// The request that will be constructed from incoming dnode message.
		request *Request

		// Will hold the return values from handler func.
		values []reflect.Value

		// First value to the response.
		result interface{}

		// Second value to the response.
		kiteErr *Error

		// Will send the response when called.
		callback dnode.Function
	)

	kite := tr.Properties()["localKite"].(*Kite)

	// Send result if "responseCallback" exists in the request.
	defer func() {
		if callback == nil {
			return
		}

		// Only argument to the callback.
		response := callbackArg{
			Result: result,
			Error:  errorForSending(kiteErr),
		}

		// Call response callback function.
		if err := callback(response); err != nil {
			kite.Log.Error(err.Error())
		}
	}()

	// Recover dnode argument errors. The caller can use functions like
	// MustString(), MustSlice()... without the fear of panic.
	defer kite.recoverError(&kiteErr)()

	request, callback = kite.parseRequest(method, args, tr)

	if !kite.disableAuthenticate {
		kiteErr = request.authenticate()
		if kiteErr != nil {
			return
		}
	}

	// Call the handler function.
	callArgs := []reflect.Value{reflect.ValueOf(request)}
	values = handlerFunc.Call(callArgs)

	result = values[0].Interface()

	if err := values[1].Interface(); err != nil {
		panic(err) // This will be recoverd from kite.recoverError() above.
	}
}

// Error is the type of the kite related errors returned from kite package.
type Error struct {
	Type    string `json:"type"`
	Message string `json:"message"`
}

func (e Error) Error() string {
	return fmt.Sprintf("kite error %s - %s", e.Type, e.Message)
}

// When a callback is called we always pass this as the only argument.
type callbackArg struct {
	Error  errorForSending `json:"error"`
	Result interface{}     `json:"result"`
}

// errorForSending is a for sending the error as an argument in a dnode message.
// Normally Error method of the Error struct is sent as a callback since it is
// exported. We do not want this behavior.
type errorForSending *Error

// recoverError returns a function which recovers the error and sets to the
// given argument as kite.Error.
//
// TODO: change that it doesn't use a pointer of pointer,
// a simpler and cleaner solution would work in the future,
func (k *Kite) recoverError(kiteErr **Error) func() {
	return func() {
		r := recover()
		if r == nil {
			return
		}

		switch err := r.(type) {
		case *Error:
			*kiteErr = err
		case *dnode.ArgumentError:
			*kiteErr = &Error{"argumentError", err.Error()}
		default:
			*kiteErr = &Error{"genericError", fmt.Sprint(r)}
			debug.PrintStack()
		}

		k.Log.Warning("Error in received message %s", (*kiteErr).Error())

	}
}

// HandlerFunc is the type of the handlers registered to Kite.
// The returned result must be Marshalable with json package.
type HandlerFunc func(*Request) (result interface{}, err error)

// HandleFunc registers a handler to run when a method call is received from a Kite.
func (k *Kite) HandleFunc(method string, handler HandlerFunc) {
	k.server.HandleFunc(method, handler)
}

// Request contains information about the incoming request.
type Request struct {
	Method         string
	Args           dnode.Arguments
	LocalKite      *Kite
	RemoteKite     *RemoteKite
	Username       string
	Authentication Authentication
	RemoteAddr     string
}

// Wrap your function with Callback to send it as an argument to a RemoteKite.
type Callback func(r *Request)

func (c Callback) MarshalJSON() ([]byte, error) {
	return []byte(`"[Function]"`), nil
}

// runCallback is called when a callback method call is received from remote Kite.
func runCallback(method string, handlerFunc reflect.Value, args dnode.Arguments, tr dnode.Transport) {
	kite := tr.Properties()["localKite"].(*Kite)

	kiteErr := new(Error)               // Not used. For recovering the error.
	defer kite.recoverError(&kiteErr)() // Do not panic no matter what.

	request, _ := kite.parseRequest(method, args, tr)

	// Call the callback function.
	callArgs := []reflect.Value{reflect.ValueOf(request)}
	handlerFunc.Call(callArgs) // No return value from callback function.
}

// parseRequest is used to read a dnode message.
// It is called when a method or callback is received to parse the message.
func (k *Kite) parseRequest(method string, arguments dnode.Arguments, tr dnode.Transport) (*Request, dnode.Function) {
	// Parse dnode method arguments: [options]
	var options callOptions
	arguments.One().MustUnmarshal(&options)

	// Properties about the client...
	properties := tr.Properties()

	// Create a new RemoteKite instance to pass it to the handler, so
	// the handler can call methods on the other site on the same connection.
	var remoteKite *RemoteKite
	if properties["remoteKite"] == nil {
		// Do not create a new RemoteKite on every request,
		// cache it in Transport.Properties().
		client := tr.(*rpc.Client) // We only have a dnode/rpc.Client for now.
		remoteKite = k.newRemoteKiteWithClient(options.Kite, options.Authentication, client)
		properties["remoteKite"] = remoteKite

		// Notify Kite.OnConnect handlers.
		k.notifyRemoteKiteConnected(remoteKite)
	} else {
		remoteKite = properties["remoteKite"].(*RemoteKite)

		// Update Kite info but do not change the value of username.
		username := remoteKite.Kite.Username
		remoteKite.Kite = options.Kite
		remoteKite.Kite.Username = username

		// TODO can't set Authentication. Fix it.
		// remoteKite.Authentication = options.Authentication
	}

	request := &Request{
		Method:         method,
		Args:           options.WithArgs,
		LocalKite:      k,
		RemoteKite:     remoteKite,
		RemoteAddr:     tr.RemoteAddr(),
		Authentication: options.Authentication,
	}

	return request, options.ResponseCallback
}

// authenticate tries to authenticate the user by selecting appropriate
// authenticator function.
func (r *Request) authenticate() *Error {
	// Trust the Kite if we have initiated the connection.
	// RemoteAddr() returns "" if this is an outgoing connection.
	if r.RemoteAddr == "" {
		return nil
	}

	// Select authenticator function.
	f := r.LocalKite.Authenticators[r.Authentication.Type]
	if f == nil {
		return &Error{
			Type:    "authenticationError",
			Message: fmt.Sprintf("Unknown authentication type: %s", r.Authentication.Type),
		}
	}

	// Call authenticator function. It sets the Request.Username field.
	err := f(r)
	if err != nil {
		return &Error{
			Type:    "authenticationError",
			Message: err.Error(),
		}
	}

	// Fix username of the remote Kite if it is invalid.
	// This prevents a Kite to impersonate someone else's Kite.
	r.RemoteKite.Kite.Username = r.Username

	return nil
}

// AuthenticateFromToken is the default Authenticator for Kite.
func (k *Kite) AuthenticateFromToken(r *Request) error {
	token, err := jwt.Parse(r.Authentication.Key, r.LocalKite.getRSAKey)
	if err != nil {
		return err
	}

	if !token.Valid {
		return errors.New("Invalid signature in token")
	}

	if audience, ok := token.Claims["aud"].(string); !ok || audience != k.ID {
		return errors.New("Invalid audience in token")
	}

	if expiration, ok := token.Claims["exp"].(float64); !ok || int64(expiration) < time.Now().UTC().Unix() {
		return errors.New("Token is expired")
	}

	if notBefore, ok := token.Claims["nbf"].(float64); !ok || int64(notBefore) > time.Now().UTC().Unix() {
		return errors.New("Token is not valid yet")
	}

	if username, ok := token.Claims["sub"].(string); !ok {
		return errors.New("Username is not present in token")
	} else {
		r.Username = username
	}

	return nil
}

// AuthenticateFromKiteKey authenticates user from Koding Key.
// Kontrol makes requests with a Koding Key.
func (k *Kite) AuthenticateFromKiteKey(r *Request) error {
	token, err := jwt.Parse(r.Authentication.Key, kitekey.GetKontrolKey)
	if err != nil {
		return err
	}

	if !token.Valid {
		return errors.New("Invalid signature in token")
	}

	if username, ok := token.Claims["sub"].(string); !ok {
		return errors.New("Username is not present in token")
	} else {
		r.Username = username
	}

	return nil
}
<|endoftext|>