<commit_before>package commands

import (
	"context"
	"errors"
	"os"
	"regexp"
	"strconv"
	"strings"

	"github.com/zetamatta/nyagos/shell"
)

var rxElse = regexp.MustCompile(`(?i)^\s*else`)

func cmd_if(ctx context.Context, cmd *shell.Cmd) (int, error) {
	// if "xxx" == "yyy"
	args := cmd.Args
	not := false
	start := 1

	option := map[string]struct{}{}

	for len(args) >= 2 && strings.HasPrefix(args[1], "/") {
		option[strings.ToLower(args[1])] = struct{}{}
		args = args[1:]
		start++
	}

	if len(args) >= 2 && strings.EqualFold(args[1], "not") {
		not = true
		args = args[1:]
		start++
	}
	status := false
	if len(args) >= 4 && args[2] == "==" {
		if _, ok := option["/i"]; ok {
			status = strings.EqualFold(args[1], args[3])
		} else {
			status = (args[1] == args[3])
		}
		args = args[4:]
		start += 3
	} else if len(args) >= 3 && strings.EqualFold(args[1], "exist") {
		_, err := os.Stat(args[2])
		status = (err == nil)
		args = args[3:]
		start += 2
	} else if len(args) >= 3 && strings.EqualFold(args[1], "errorlevel") {
		num, num_err := strconv.Atoi(args[2])
		if num_err == nil {
			status = (shell.LastErrorLevel >= num)
		}
		args = args[2:]
		start += 2
	}

	if not {
		status = !status
	}

	if len(args) > 0 {
		if status {
			subCmd, err := cmd.Clone()
			if err != nil {
				return 0, err
			}
			subCmd.Args = cmd.Args[start:]
			subCmd.RawArgs = cmd.RawArgs[start:]
			return subCmd.SpawnvpContext(ctx)
		}
	} else {
		stream, ok := ctx.Value("stream").(shell.Stream)
		if !ok {
			return 1, errors.New("not found stream")
		}
		thenBuffer := BufStream{}
		elseBuffer := BufStream{}
		elsePart := false

		save_prompt := os.Getenv("PROMPT")
		os.Setenv("PROMPT", "if>")
		nest := 1
		for {
			_, line, err := cmd.ReadCommand(ctx, stream)
			if err != nil {
				break
			}
			args := shell.SplitQ(line)
			name := strings.ToLower(args[0])
			if _, ok := start_list[name]; ok {
				nest++
			} else if name == "end" || name == "endif" {
				nest--
				if nest == 0 {
					break
				}
			} else if name == "else" {
				if nest == 1 {
					elsePart = true
					os.Setenv("PROMPT", "else>")
					line = rxElse.ReplaceAllString(line, "")
				}
			}
			if elsePart {
				elseBuffer.Add(line)
			} else {
				thenBuffer.Add(line)
			}
		}
		os.Setenv("PROMPT", save_prompt)

		if status {
			cmd.Loop(&thenBuffer)
		} else {
			cmd.Loop(&elseBuffer)
		}
	}
	return 0, nil
}
<commit_msg>Support `then` keyword<commit_after>package commands

import (
	"context"
	"errors"
	"os"
	"regexp"
	"strconv"
	"strings"

	"github.com/zetamatta/nyagos/shell"
)

var rxElse = regexp.MustCompile(`(?i)^\s*else`)

func cmd_if(ctx context.Context, cmd *shell.Cmd) (int, error) {
	// if "xxx" == "yyy"
	args := cmd.Args
	rawargs := cmd.RawArgs
	not := false
	start := 1

	option := map[string]struct{}{}

	for len(args) >= 2 && strings.HasPrefix(args[1], "/") {
		option[strings.ToLower(args[1])] = struct{}{}
		args = args[1:]
		rawargs = rawargs[1:]
		start++
	}

	if len(args) >= 2 && strings.EqualFold(args[1], "not") {
		not = true
		args = args[1:]
		rawargs = rawargs[1:]
		start++
	}
	status := false
	if len(args) >= 4 && args[2] == "==" {
		if _, ok := option["/i"]; ok {
			status = strings.EqualFold(args[1], args[3])
		} else {
			status = (args[1] == args[3])
		}
		args = args[4:]
		rawargs = rawargs[4:]
		start += 3
	} else if len(args) >= 3 && strings.EqualFold(args[1], "exist") {
		_, err := os.Stat(args[2])
		status = (err == nil)
		args = args[3:]
		rawargs = rawargs[3:]
		start += 2
	} else if len(args) >= 3 && strings.EqualFold(args[1], "errorlevel") {
		num, num_err := strconv.Atoi(args[2])
		if num_err == nil {
			status = (shell.LastErrorLevel >= num)
		}
		args = args[2:]
		rawargs = rawargs[2:]
		start += 2
	}

	if not {
		status = !status
	}

	thenBuffer := BufStream{}

	if len(args) > 0 {
		if args[0] == "then" {
			// inline and block `then`
			if len(args) > 1 {
				thenBuffer.Add(strings.Join(rawargs[1:], " "))
			}
			// continue
		} else {
			// inline `then`
			if status {
				subCmd, err := cmd.Clone()
				if err != nil {
					return 0, err
				}
				subCmd.Args = cmd.Args[start:]
				subCmd.RawArgs = cmd.RawArgs[start:]
				return subCmd.SpawnvpContext(ctx)
			} else {
				return 0, nil
			}
		}
	}

	// block `then` / `else`

	stream, ok := ctx.Value("stream").(shell.Stream)
	if !ok {
		return 1, errors.New("not found stream")
	}

	elseBuffer := BufStream{}
	elsePart := false

	save_prompt := os.Getenv("PROMPT")
	os.Setenv("PROMPT", "if>")
	nest := 1
	for {
		_, line, err := cmd.ReadCommand(ctx, stream)
		if err != nil {
			break
		}
		args := shell.SplitQ(line)
		name := strings.ToLower(args[0])
		if _, ok := start_list[name]; ok {
			nest++
		} else if name == "end" || name == "endif" {
			nest--
			if nest == 0 {
				break
			}
		} else if name == "else" {
			if nest == 1 {
				elsePart = true
				os.Setenv("PROMPT", "else>")
				line = rxElse.ReplaceAllString(line, "")
			}
		}
		if elsePart {
			elseBuffer.Add(line)
		} else {
			thenBuffer.Add(line)
		}
	}
	os.Setenv("PROMPT", save_prompt)

	if status {
		cmd.Loop(&thenBuffer)
	} else {
		cmd.Loop(&elseBuffer)
	}
	return 0, nil
}
<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package mqtt

import (
	"sync"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/TheThingsNetwork/ttn/core"
	"github.com/TheThingsNetwork/ttn/core/mocks"
	ttnMQTT "github.com/TheThingsNetwork/ttn/mqtt"
	. "github.com/TheThingsNetwork/ttn/utils/testing"
	. "github.com/smartystreets/assertions"
)

func TestNewAdapter(t *testing.T) {
	a := New(t)
	ctx := GetLogger(t, "TestNewAdapter")
	client := ttnMQTT.NewClient(ctx, "test", "", "", "tcp://localhost:1883")
	adapter := NewAdapter(ctx, client)

	a.So(adapter.(*defaultAdapter).client, ShouldEqual, client)
}

func TestHandleData(t *testing.T) {
	a := New(t)
	ctx := GetLogger(t, "TestHandleData")
	client := ttnMQTT.NewClient(ctx, "test", "", "", "tcp://localhost:1883")
	client.Connect()

	adapter := NewAdapter(ctx, client)

	eui := []byte{0x0a, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08}

	req := core.DataAppReq{
		Payload: []byte{0x01, 0x02},
		Metadata: []*core.Metadata{
			&core.Metadata{DataRate: "SF7BW125"},
		},
		AppEUI: eui,
		DevEUI: eui,
	}

	var wg sync.WaitGroup
	wg.Add(1)

	client.SubscribeDeviceUplink(eui, eui, func(client ttnMQTT.Client, appEUI []byte, devEUI []byte, dataUp core.DataUpAppReq) {
		a.So(appEUI, ShouldResemble, eui)
		a.So(devEUI, ShouldResemble, eui)
		a.So(dataUp.Payload, ShouldResemble, []byte{0x01, 0x02})
		a.So(dataUp.Metadata[0].DataRate, ShouldEqual, "SF7BW125")
		wg.Done()
	}).Wait()

	res, err := adapter.HandleData(context.Background(), &req)
	a.So(err, ShouldBeNil)
	a.So(res, ShouldBeNil)

	wg.Wait()

}

func TestHandleInvalidData(t *testing.T) {
	a := New(t)
	client := ttnMQTT.NewClient(nil, "test", "", "", "tcp://localhost:1883")
	adapter := NewAdapter(nil, client)

	// nil Request
	_, err := adapter.HandleData(context.Background(), nil)
	a.So(err, ShouldNotBeNil)

	// Invalid Payload
	_, err = adapter.HandleData(context.Background(), &core.DataAppReq{
		Payload: []byte{},
	})
	a.So(err, ShouldNotBeNil)

	// Invalid DevEUI
	_, err = adapter.HandleData(context.Background(), &core.DataAppReq{
		Payload: []byte{0x00},
		DevEUI:  []byte{},
	})
	a.So(err, ShouldNotBeNil)

	// Invalid AppEUI
	_, err = adapter.HandleData(context.Background(), &core.DataAppReq{
		Payload: []byte{0x00},
		DevEUI:  []byte{0x0b, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		AppEUI:  []byte{},
	})
	a.So(err, ShouldNotBeNil)

	// Missing Metadata
	_, err = adapter.HandleData(context.Background(), &core.DataAppReq{
		Payload: []byte{0x00},
		DevEUI:  []byte{0x0b, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		AppEUI:  []byte{0x0b, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
	})
	a.So(err, ShouldNotBeNil)

	// Not Connected
	_, err = adapter.HandleData(context.Background(), &core.DataAppReq{
		Payload:  []byte{0x00},
		DevEUI:   []byte{0x0b, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		AppEUI:   []byte{0x0b, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		Metadata: []*core.Metadata{},
	})
	a.So(err, ShouldNotBeNil)
}

func TestHandleJoin(t *testing.T) {
	a := New(t)
	ctx := GetLogger(t, "TestHandleJoin")
	client := ttnMQTT.NewClient(ctx, "test", "", "", "tcp://localhost:1883")
	client.Connect()

	adapter := NewAdapter(ctx, client)

	eui := []byte{0x0a, 0x03, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08}

	req := core.JoinAppReq{
		AppEUI: eui,
		DevEUI: eui,
		Metadata: []*core.Metadata{
			&core.Metadata{DataRate: "SF7BW125"},
		},
	}

	var wg sync.WaitGroup
	wg.Add(1)

	client.SubscribeDeviceActivations(eui, eui, func(client ttnMQTT.Client, appEUI []byte, devEUI []byte, activation core.OTAAAppReq) {
		a.So(appEUI, ShouldResemble, eui)
		a.So(devEUI, ShouldResemble, eui)
		a.So(activation.Metadata[0].DataRate, ShouldEqual, "SF7BW125")
		wg.Done()
	}).Wait()

	res, err := adapter.HandleJoin(context.Background(), &req)
	a.So(err, ShouldBeNil)
	a.So(res, ShouldBeNil)

	wg.Wait()
}

func TestHandleInvalidJoin(t *testing.T) {
	a := New(t)
	client := ttnMQTT.NewClient(nil, "test", "", "", "tcp://localhost:1883")
	adapter := NewAdapter(nil, client)

	// nil Request
	_, err := adapter.HandleJoin(context.Background(), nil)
	a.So(err, ShouldNotBeNil)

	// Invalid DevEUI
	_, err = adapter.HandleJoin(context.Background(), &core.JoinAppReq{
		DevEUI: []byte{},
	})
	a.So(err, ShouldNotBeNil)

	// Invalid AppEUI
	_, err = adapter.HandleJoin(context.Background(), &core.JoinAppReq{
		DevEUI: []byte{0x0c, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		AppEUI: []byte{},
	})
	a.So(err, ShouldNotBeNil)

	// Missing Metadata
	_, err = adapter.HandleJoin(context.Background(), &core.JoinAppReq{
		DevEUI: []byte{0x0c, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		AppEUI: []byte{0x0c, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
	})
	a.So(err, ShouldNotBeNil)

	// Not Connected
	_, err = adapter.HandleJoin(context.Background(), &core.JoinAppReq{
		DevEUI:   []byte{0x0c, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		AppEUI:   []byte{0x0c, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		Metadata: []*core.Metadata{},
	})
	a.So(err, ShouldNotBeNil)
}

func TestSubscribeDownlink(t *testing.T) {
	a := New(t)
	ctx := GetLogger(t, "TestHandleJoin")
	client := ttnMQTT.NewClient(ctx, "test", "", "", "tcp://localhost:1883")
	client.Connect()

	adapter := NewAdapter(ctx, client)
	handler := mocks.NewHandlerServer()

	adapter.SubscribeDownlink(handler)

	appEUI := []byte{0x04, 0x03, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08}
	devEUI := []byte{0x08, 0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01}
	client.PublishDownlink(appEUI, devEUI, core.DataDownAppReq{Payload: []byte{0x01, 0x02, 0x03, 0x04}}).Wait()

	<-time.After(25 * time.Millisecond)

	expected := &core.DataDownHandlerReq{
		AppEUI:  appEUI,
		DevEUI:  devEUI,
		Payload: []byte{0x01, 0x02, 0x03, 0x04},
	}

	a.So(handler.InHandleDataDown.Req, ShouldResemble, expected)
}

func TestSubscribeInvalidDownlink(t *testing.T) {
	a := New(t)
	ctx := GetLogger(t, "TestHandleJoin")
	client := ttnMQTT.NewClient(ctx, "test", "", "", "tcp://localhost:1883")
	client.Connect()

	adapter := NewAdapter(ctx, client)
	handler := mocks.NewHandlerServer()

	adapter.SubscribeDownlink(handler)

	appEUI := []byte{0x04, 0x03, 0x03, 0x09, 0x05, 0x06, 0x07, 0x08}
	devEUI := []byte{0x08, 0x07, 0x06, 0x09, 0x04, 0x03, 0x02, 0x01}
	client.PublishDownlink(appEUI, devEUI, core.DataDownAppReq{Payload: []byte{}}).Wait()

	<-time.After(25 * time.Millisecond)

	a.So(handler.InHandleDataDown.Req, ShouldBeNil)
}
<commit_msg>[mqtt-adapter] copy-paste<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package mqtt

import (
	"sync"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/TheThingsNetwork/ttn/core"
	"github.com/TheThingsNetwork/ttn/core/mocks"
	ttnMQTT "github.com/TheThingsNetwork/ttn/mqtt"
	. "github.com/TheThingsNetwork/ttn/utils/testing"
	. "github.com/smartystreets/assertions"
)

func TestNewAdapter(t *testing.T) {
	a := New(t)
	ctx := GetLogger(t, "TestNewAdapter")
	client := ttnMQTT.NewClient(ctx, "test", "", "", "tcp://localhost:1883")
	adapter := NewAdapter(ctx, client)

	a.So(adapter.(*defaultAdapter).client, ShouldEqual, client)
}

func TestHandleData(t *testing.T) {
	a := New(t)
	ctx := GetLogger(t, "TestHandleData")
	client := ttnMQTT.NewClient(ctx, "test", "", "", "tcp://localhost:1883")
	client.Connect()

	adapter := NewAdapter(ctx, client)

	eui := []byte{0x0a, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08}

	req := core.DataAppReq{
		Payload: []byte{0x01, 0x02},
		Metadata: []*core.Metadata{
			&core.Metadata{DataRate: "SF7BW125"},
		},
		AppEUI: eui,
		DevEUI: eui,
	}

	var wg sync.WaitGroup
	wg.Add(1)

	client.SubscribeDeviceUplink(eui, eui, func(client ttnMQTT.Client, appEUI []byte, devEUI []byte, dataUp core.DataUpAppReq) {
		a.So(appEUI, ShouldResemble, eui)
		a.So(devEUI, ShouldResemble, eui)
		a.So(dataUp.Payload, ShouldResemble, []byte{0x01, 0x02})
		a.So(dataUp.Metadata[0].DataRate, ShouldEqual, "SF7BW125")
		wg.Done()
	}).Wait()

	res, err := adapter.HandleData(context.Background(), &req)
	a.So(err, ShouldBeNil)
	a.So(res, ShouldBeNil)

	wg.Wait()

}

func TestHandleInvalidData(t *testing.T) {
	a := New(t)
	client := ttnMQTT.NewClient(nil, "test", "", "", "tcp://localhost:1883")
	adapter := NewAdapter(nil, client)

	// nil Request
	_, err := adapter.HandleData(context.Background(), nil)
	a.So(err, ShouldNotBeNil)

	// Invalid Payload
	_, err = adapter.HandleData(context.Background(), &core.DataAppReq{
		Payload: []byte{},
	})
	a.So(err, ShouldNotBeNil)

	// Invalid DevEUI
	_, err = adapter.HandleData(context.Background(), &core.DataAppReq{
		Payload: []byte{0x00},
		DevEUI:  []byte{},
	})
	a.So(err, ShouldNotBeNil)

	// Invalid AppEUI
	_, err = adapter.HandleData(context.Background(), &core.DataAppReq{
		Payload: []byte{0x00},
		DevEUI:  []byte{0x0b, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		AppEUI:  []byte{},
	})
	a.So(err, ShouldNotBeNil)

	// Missing Metadata
	_, err = adapter.HandleData(context.Background(), &core.DataAppReq{
		Payload: []byte{0x00},
		DevEUI:  []byte{0x0b, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		AppEUI:  []byte{0x0b, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
	})
	a.So(err, ShouldNotBeNil)

	// Not Connected
	_, err = adapter.HandleData(context.Background(), &core.DataAppReq{
		Payload:  []byte{0x00},
		DevEUI:   []byte{0x0b, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		AppEUI:   []byte{0x0b, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		Metadata: []*core.Metadata{},
	})
	a.So(err, ShouldNotBeNil)
}

func TestHandleJoin(t *testing.T) {
	a := New(t)
	ctx := GetLogger(t, "TestHandleJoin")
	client := ttnMQTT.NewClient(ctx, "test", "", "", "tcp://localhost:1883")
	client.Connect()

	adapter := NewAdapter(ctx, client)

	eui := []byte{0x0a, 0x03, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08}

	req := core.JoinAppReq{
		AppEUI: eui,
		DevEUI: eui,
		Metadata: []*core.Metadata{
			&core.Metadata{DataRate: "SF7BW125"},
		},
	}

	var wg sync.WaitGroup
	wg.Add(1)

	client.SubscribeDeviceActivations(eui, eui, func(client ttnMQTT.Client, appEUI []byte, devEUI []byte, activation core.OTAAAppReq) {
		a.So(appEUI, ShouldResemble, eui)
		a.So(devEUI, ShouldResemble, eui)
		a.So(activation.Metadata[0].DataRate, ShouldEqual, "SF7BW125")
		wg.Done()
	}).Wait()

	res, err := adapter.HandleJoin(context.Background(), &req)
	a.So(err, ShouldBeNil)
	a.So(res, ShouldBeNil)

	wg.Wait()
}

func TestHandleInvalidJoin(t *testing.T) {
	a := New(t)
	client := ttnMQTT.NewClient(nil, "test", "", "", "tcp://localhost:1883")
	adapter := NewAdapter(nil, client)

	// nil Request
	_, err := adapter.HandleJoin(context.Background(), nil)
	a.So(err, ShouldNotBeNil)

	// Invalid DevEUI
	_, err = adapter.HandleJoin(context.Background(), &core.JoinAppReq{
		DevEUI: []byte{},
	})
	a.So(err, ShouldNotBeNil)

	// Invalid AppEUI
	_, err = adapter.HandleJoin(context.Background(), &core.JoinAppReq{
		DevEUI: []byte{0x0c, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		AppEUI: []byte{},
	})
	a.So(err, ShouldNotBeNil)

	// Missing Metadata
	_, err = adapter.HandleJoin(context.Background(), &core.JoinAppReq{
		DevEUI: []byte{0x0c, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		AppEUI: []byte{0x0c, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
	})
	a.So(err, ShouldNotBeNil)

	// Not Connected
	_, err = adapter.HandleJoin(context.Background(), &core.JoinAppReq{
		DevEUI:   []byte{0x0c, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		AppEUI:   []byte{0x0c, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08},
		Metadata: []*core.Metadata{},
	})
	a.So(err, ShouldNotBeNil)
}

func TestSubscribeDownlink(t *testing.T) {
	a := New(t)
	ctx := GetLogger(t, "TestSubscribeDownlink")
	client := ttnMQTT.NewClient(ctx, "test", "", "", "tcp://localhost:1883")
	client.Connect()

	adapter := NewAdapter(ctx, client)
	handler := mocks.NewHandlerServer()

	adapter.SubscribeDownlink(handler)

	appEUI := []byte{0x04, 0x03, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08}
	devEUI := []byte{0x08, 0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01}
	client.PublishDownlink(appEUI, devEUI, core.DataDownAppReq{Payload: []byte{0x01, 0x02, 0x03, 0x04}}).Wait()

	<-time.After(25 * time.Millisecond)

	expected := &core.DataDownHandlerReq{
		AppEUI:  appEUI,
		DevEUI:  devEUI,
		Payload: []byte{0x01, 0x02, 0x03, 0x04},
	}

	a.So(handler.InHandleDataDown.Req, ShouldResemble, expected)
}

func TestSubscribeInvalidDownlink(t *testing.T) {
	a := New(t)
	ctx := GetLogger(t, "TestSubscribeInvalidDownlink")
	client := ttnMQTT.NewClient(ctx, "test", "", "", "tcp://localhost:1883")
	client.Connect()

	adapter := NewAdapter(ctx, client)
	handler := mocks.NewHandlerServer()

	adapter.SubscribeDownlink(handler)

	appEUI := []byte{0x04, 0x03, 0x03, 0x09, 0x05, 0x06, 0x07, 0x08}
	devEUI := []byte{0x08, 0x07, 0x06, 0x09, 0x04, 0x03, 0x02, 0x01}
	client.PublishDownlink(appEUI, devEUI, core.DataDownAppReq{Payload: []byte{}}).Wait()

	<-time.After(25 * time.Millisecond)

	a.So(handler.InHandleDataDown.Req, ShouldBeNil)
}
<|endoftext|>


<commit_before>package db

import (
	"database/sql"
	"time"

	"github.com/stellar/horizon/log"

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

const (
	MaxHistoryLedger = "SELECT MAX(sequence) FROM history_ledgers"
)

// NewLedgerClosePump starts a background proc that continually watches the
// history database provided.  The watch is stopped after the provided context
// is cancelled.
//
// Every second, the proc spawned by calling this func will check to see
// if a new ledger has been imported (by ruby-horizon as of 2015-04-30, but
// should eventually end up being in this project).  If a new ledger is seen
// the the channel returned by this function emits
func NewLedgerClosePump(ctx context.Context, db *sql.DB) <-chan struct{} {
	result := make(chan struct{})

	go func() {
		var lastSeenLedger int32
		for {
			select {
			case <-time.After(1 * time.Second):
				var latestLedger int32
				row := db.QueryRow(MaxHistoryLedger)
				err := row.Scan(&latestLedger)

				if err != nil {
					log.Warn(ctx, "Failed to check latest ledger", err)
					break
				}

				if latestLedger > lastSeenLedger {
					log.Debugf(ctx, "saw new ledger: %d, prev: %d", latestLedger, lastSeenLedger)

					select {
					case result <- struct{}{}:
						lastSeenLedger = latestLedger
					default:
						log.Debug(ctx, "ledger pump channel is blocked.  waiting...")
					}
				}

			case <-ctx.Done():
				log.Info(ctx, "canceling ledger pump")
				close(result)
				return
			}
		}
	}()

	return result
}
<commit_msg>Add ledger pump reset trigger<commit_after>package db

import (
	"database/sql"
	"time"

	"github.com/stellar/horizon/log"

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

const (
	MaxHistoryLedger = "SELECT MAX(sequence) FROM history_ledgers"
)

// NewLedgerClosePump starts a background proc that continually watches the
// history database provided.  The watch is stopped after the provided context
// is cancelled.
//
// Every second, the proc spawned by calling this func will check to see
// if a new ledger has been imported (by ruby-horizon as of 2015-04-30, but
// should eventually end up being in this project).  If a new ledger is seen
// the the channel returned by this function emits
func NewLedgerClosePump(ctx context.Context, db *sql.DB) <-chan struct{} {
	result := make(chan struct{})

	go func() {
		var lastSeenLedger int32
		for {
			select {
			case <-time.After(1 * time.Second):
				var latestLedger int32
				row := db.QueryRow(MaxHistoryLedger)
				err := row.Scan(&latestLedger)

				if err != nil {
					log.Warn(ctx, "Failed to check latest ledger", err)
					break
				}

				if latestLedger > lastSeenLedger {
					log.Debugf(ctx, "saw new ledger: %d, prev: %d", latestLedger, lastSeenLedger)

					select {
					case result <- struct{}{}:
						lastSeenLedger = latestLedger
					default:
						log.Debug(ctx, "ledger pump channel is blocked.  waiting...")
					}
				} else if latestLedger < lastSeenLedger {
					log.Warn(ctx, "latest ledger went backwards! reseting ledger pump")
					lastSeenLedger = 0
				}

			case <-ctx.Done():
				log.Info(ctx, "canceling ledger pump")
				close(result)
				return
			}
		}
	}()

	return result
}
<|endoftext|>


<commit_before>/*
Copyright IBM Corp. 2016 All Rights Reserved.

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

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

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

package chaincode

import (
	"fmt"
	"sync"
	"testing"

	"github.com/hyperledger/fabric/common/util"
	pb "github.com/hyperledger/fabric/protos/peer"

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

//TestExecuteConcurrentInvokes deploys newkeyperinvoke and runs 100 concurrent invokes
//followed by concurrent 100 queries to validate
func TestExecuteConcurrentInvokes(t *testing.T) {
	chainID := util.GetTestChainID()

	lis, err := initPeer(chainID)
	if err != nil {
		t.Fail()
		t.Logf("Error creating peer: %s", err)
	}

	defer finitPeer(lis, chainID)

	var ctxt = context.Background()

	url := "github.com/hyperledger/fabric/examples/ccchecker/chaincodes/newkeyperinvoke"

	chaincodeID := &pb.ChaincodeID{Name: "nkpi", Path: url}

	args := util.ToChaincodeArgs("init", "")

	spec := &pb.ChaincodeSpec{Type: 1, ChaincodeID: chaincodeID, CtorMsg: &pb.ChaincodeInput{Args: args}}

	cccid := NewCCContext(chainID, "nkpi", "0", "", false, nil)

	defer theChaincodeSupport.Stop(ctxt, cccid, &pb.ChaincodeDeploymentSpec{ChaincodeSpec: spec})

	_, err = deploy(ctxt, cccid, spec)
	if err != nil {
		t.Fail()
		t.Logf("Error initializing chaincode %s(%s)", chaincodeID, err)
		return
	}

	var wg sync.WaitGroup

	//run 100 invokes in parallel
	numTrans := 100

	results := make([][]byte, numTrans)
	errs := make([]error, numTrans)

	e := func(inv bool, qnum int) {
		defer wg.Done()

		newkey := fmt.Sprintf("%d", qnum)

		var args [][]byte
		if inv {
			args = util.ToChaincodeArgs("put", newkey, newkey)
		} else {
			args = util.ToChaincodeArgs("get", newkey)
		}

		spec = &pb.ChaincodeSpec{Type: 1, ChaincodeID: chaincodeID, CtorMsg: &pb.ChaincodeInput{Args: args}}

		//start with a new background
		_, _, results[qnum], err = invoke(context.Background(), chainID, spec)

		if err != nil {
			errs[qnum] = fmt.Errorf("Error executing <%s>: %s", chaincodeID.Name, err)
			return
		}
	}

	wg.Add(numTrans)

	//execute transactions concurrently.
	for i := 0; i < numTrans; i++ {
		go e(true, i)
	}

	wg.Wait()

	for i := 0; i < numTrans; i++ {
		if errs[i] != nil {
			t.Fail()
			t.Logf("Error invoking chaincode iter %d %s(%s)", i, chaincodeID.Name, errs[i])
		}
		if results[i] == nil || string(results[i]) != "OK" {
			t.Fail()
			t.Logf("Error concurrent invoke %d %s", i, chaincodeID.Name)
			return
		}
	}

	wg.Add(numTrans)

	//execute queries concurrently.
	for i := 0; i < numTrans; i++ {
		go e(false, i)
	}

	wg.Wait()

	for i := 0; i < numTrans; i++ {
		if errs[i] != nil {
			t.Fail()
			t.Logf("Error querying chaincode iter %d %s(%s)", i, chaincodeID.Name, errs[i])
			return
		}
		if results[i] == nil || string(results[i]) != fmt.Sprintf("%d", i) {
			t.Fail()
			if results[i] == nil {
				t.Logf("Error concurrent query %d(%s)", i, chaincodeID.Name)
			} else {
				t.Logf("Error concurrent query %d(%s, %s, %v)", i, chaincodeID.Name, string(results[i]), results[i])
			}
			return
		}
	}
}
<commit_msg>FAB-1600 skip TestExecuteConcurrentInvokes<commit_after>/*
Copyright IBM Corp. 2016 All Rights Reserved.

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

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

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

package chaincode

import (
	"fmt"
	"sync"
	"testing"

	"github.com/hyperledger/fabric/common/util"
	pb "github.com/hyperledger/fabric/protos/peer"

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

//TestExecuteConcurrentInvokes deploys newkeyperinvoke and runs 100 concurrent invokes
//followed by concurrent 100 queries to validate
func TestExecuteConcurrentInvokes(t *testing.T) {
	//this test fails occasionally. FAB-1600 is opened to track this issue
	//skip meanwhile so as to not block CI
	t.Skip()
	chainID := util.GetTestChainID()

	lis, err := initPeer(chainID)
	if err != nil {
		t.Fail()
		t.Logf("Error creating peer: %s", err)
	}

	defer finitPeer(lis, chainID)

	var ctxt = context.Background()

	url := "github.com/hyperledger/fabric/examples/ccchecker/chaincodes/newkeyperinvoke"

	chaincodeID := &pb.ChaincodeID{Name: "nkpi", Path: url}

	args := util.ToChaincodeArgs("init", "")

	spec := &pb.ChaincodeSpec{Type: 1, ChaincodeID: chaincodeID, CtorMsg: &pb.ChaincodeInput{Args: args}}

	cccid := NewCCContext(chainID, "nkpi", "0", "", false, nil)

	defer theChaincodeSupport.Stop(ctxt, cccid, &pb.ChaincodeDeploymentSpec{ChaincodeSpec: spec})

	_, err = deploy(ctxt, cccid, spec)
	if err != nil {
		t.Fail()
		t.Logf("Error initializing chaincode %s(%s)", chaincodeID, err)
		return
	}

	var wg sync.WaitGroup

	//run 100 invokes in parallel
	numTrans := 100

	results := make([][]byte, numTrans)
	errs := make([]error, numTrans)

	e := func(inv bool, qnum int) {
		defer wg.Done()

		newkey := fmt.Sprintf("%d", qnum)

		var args [][]byte
		if inv {
			args = util.ToChaincodeArgs("put", newkey, newkey)
		} else {
			args = util.ToChaincodeArgs("get", newkey)
		}

		spec = &pb.ChaincodeSpec{Type: 1, ChaincodeID: chaincodeID, CtorMsg: &pb.ChaincodeInput{Args: args}}

		//start with a new background
		_, _, results[qnum], err = invoke(context.Background(), chainID, spec)

		if err != nil {
			errs[qnum] = fmt.Errorf("Error executing <%s>: %s", chaincodeID.Name, err)
			return
		}
	}

	wg.Add(numTrans)

	//execute transactions concurrently.
	for i := 0; i < numTrans; i++ {
		go e(true, i)
	}

	wg.Wait()

	for i := 0; i < numTrans; i++ {
		if errs[i] != nil {
			t.Fail()
			t.Logf("Error invoking chaincode iter %d %s(%s)", i, chaincodeID.Name, errs[i])
		}
		if results[i] == nil || string(results[i]) != "OK" {
			t.Fail()
			t.Logf("Error concurrent invoke %d %s", i, chaincodeID.Name)
			return
		}
	}

	wg.Add(numTrans)

	//execute queries concurrently.
	for i := 0; i < numTrans; i++ {
		go e(false, i)
	}

	wg.Wait()

	for i := 0; i < numTrans; i++ {
		if errs[i] != nil {
			t.Fail()
			t.Logf("Error querying chaincode iter %d %s(%s)", i, chaincodeID.Name, errs[i])
			return
		}
		if results[i] == nil || string(results[i]) != fmt.Sprintf("%d", i) {
			t.Fail()
			if results[i] == nil {
				t.Logf("Error concurrent query %d(%s)", i, chaincodeID.Name)
			} else {
				t.Logf("Error concurrent query %d(%s, %s, %v)", i, chaincodeID.Name, string(results[i]), results[i])
			}
			return
		}
	}
}
<|endoftext|>


<commit_before>// +build linux

package overlay2 // import "github.com/docker/docker/daemon/graphdriver/overlay2"

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"syscall"

	"github.com/docker/docker/pkg/system"
	"github.com/pkg/errors"
	"golang.org/x/sys/unix"
)

// doesSupportNativeDiff checks whether the filesystem has a bug
// which copies up the opaque flag when copying up an opaque
// directory or the kernel enable CONFIG_OVERLAY_FS_REDIRECT_DIR.
// When these exist naive diff should be used.
func doesSupportNativeDiff(d string) error {
	td, err := ioutil.TempDir(d, "opaque-bug-check")
	if err != nil {
		return err
	}
	defer func() {
		if err := os.RemoveAll(td); err != nil {
			logger.Warnf("Failed to remove check directory %v: %v", td, err)
		}
	}()

	// Make directories l1/d, l1/d1, l2/d, l3, work, merged
	if err := os.MkdirAll(filepath.Join(td, "l1", "d"), 0755); err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Join(td, "l1", "d1"), 0755); err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Join(td, "l2", "d"), 0755); err != nil {
		return err
	}
	if err := os.Mkdir(filepath.Join(td, "l3"), 0755); err != nil {
		return err
	}
	if err := os.Mkdir(filepath.Join(td, workDirName), 0755); err != nil {
		return err
	}
	if err := os.Mkdir(filepath.Join(td, mergedDirName), 0755); err != nil {
		return err
	}

	// Mark l2/d as opaque
	if err := system.Lsetxattr(filepath.Join(td, "l2", "d"), "trusted.overlay.opaque", []byte("y"), 0); err != nil {
		return errors.Wrap(err, "failed to set opaque flag on middle layer")
	}

	opts := fmt.Sprintf("lowerdir=%s:%s,upperdir=%s,workdir=%s", path.Join(td, "l2"), path.Join(td, "l1"), path.Join(td, "l3"), path.Join(td, workDirName))
	if err := unix.Mount("overlay", filepath.Join(td, mergedDirName), "overlay", 0, opts); err != nil {
		return errors.Wrap(err, "failed to mount overlay")
	}
	defer func() {
		if err := unix.Unmount(filepath.Join(td, mergedDirName), 0); err != nil {
			logger.Warnf("Failed to unmount check directory %v: %v", filepath.Join(td, mergedDirName), err)
		}
	}()

	// Touch file in d to force copy up of opaque directory "d" from "l2" to "l3"
	if err := ioutil.WriteFile(filepath.Join(td, mergedDirName, "d", "f"), []byte{}, 0644); err != nil {
		return errors.Wrap(err, "failed to write to merged directory")
	}

	// Check l3/d does not have opaque flag
	xattrOpaque, err := system.Lgetxattr(filepath.Join(td, "l3", "d"), "trusted.overlay.opaque")
	if err != nil {
		return errors.Wrap(err, "failed to read opaque flag on upper layer")
	}
	if string(xattrOpaque) == "y" {
		return errors.New("opaque flag erroneously copied up, consider update to kernel 4.8 or later to fix")
	}

	// rename "d1" to "d2"
	if err := os.Rename(filepath.Join(td, mergedDirName, "d1"), filepath.Join(td, mergedDirName, "d2")); err != nil {
		// if rename failed with syscall.EXDEV, the kernel doesn't have CONFIG_OVERLAY_FS_REDIRECT_DIR enabled
		if err.(*os.LinkError).Err == syscall.EXDEV {
			return nil
		}
		return errors.Wrap(err, "failed to rename dir in merged directory")
	}
	// get the xattr of "d2"
	xattrRedirect, err := system.Lgetxattr(filepath.Join(td, "l3", "d2"), "trusted.overlay.redirect")
	if err != nil {
		return errors.Wrap(err, "failed to read redirect flag on upper layer")
	}

	if string(xattrRedirect) == "d1" {
		return errors.New("kernel has CONFIG_OVERLAY_FS_REDIRECT_DIR enabled")
	}

	return nil
}
<commit_msg>overlay2: doesSupportNativeDiff: add fast path for userns<commit_after>// +build linux

package overlay2 // import "github.com/docker/docker/daemon/graphdriver/overlay2"

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"syscall"

	"github.com/containerd/containerd/sys"
	"github.com/docker/docker/pkg/system"
	"github.com/pkg/errors"
	"golang.org/x/sys/unix"
)

// doesSupportNativeDiff checks whether the filesystem has a bug
// which copies up the opaque flag when copying up an opaque
// directory or the kernel enable CONFIG_OVERLAY_FS_REDIRECT_DIR.
// When these exist naive diff should be used.
//
// When running in a user namespace, returns errRunningInUserNS
// immediately.
func doesSupportNativeDiff(d string) error {
	if sys.RunningInUserNS() {
		return errors.New("running in a user namespace")
	}

	td, err := ioutil.TempDir(d, "opaque-bug-check")
	if err != nil {
		return err
	}
	defer func() {
		if err := os.RemoveAll(td); err != nil {
			logger.Warnf("Failed to remove check directory %v: %v", td, err)
		}
	}()

	// Make directories l1/d, l1/d1, l2/d, l3, work, merged
	if err := os.MkdirAll(filepath.Join(td, "l1", "d"), 0755); err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Join(td, "l1", "d1"), 0755); err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Join(td, "l2", "d"), 0755); err != nil {
		return err
	}
	if err := os.Mkdir(filepath.Join(td, "l3"), 0755); err != nil {
		return err
	}
	if err := os.Mkdir(filepath.Join(td, workDirName), 0755); err != nil {
		return err
	}
	if err := os.Mkdir(filepath.Join(td, mergedDirName), 0755); err != nil {
		return err
	}

	// Mark l2/d as opaque
	if err := system.Lsetxattr(filepath.Join(td, "l2", "d"), "trusted.overlay.opaque", []byte("y"), 0); err != nil {
		return errors.Wrap(err, "failed to set opaque flag on middle layer")
	}

	opts := fmt.Sprintf("lowerdir=%s:%s,upperdir=%s,workdir=%s", path.Join(td, "l2"), path.Join(td, "l1"), path.Join(td, "l3"), path.Join(td, workDirName))
	if err := unix.Mount("overlay", filepath.Join(td, mergedDirName), "overlay", 0, opts); err != nil {
		return errors.Wrap(err, "failed to mount overlay")
	}
	defer func() {
		if err := unix.Unmount(filepath.Join(td, mergedDirName), 0); err != nil {
			logger.Warnf("Failed to unmount check directory %v: %v", filepath.Join(td, mergedDirName), err)
		}
	}()

	// Touch file in d to force copy up of opaque directory "d" from "l2" to "l3"
	if err := ioutil.WriteFile(filepath.Join(td, mergedDirName, "d", "f"), []byte{}, 0644); err != nil {
		return errors.Wrap(err, "failed to write to merged directory")
	}

	// Check l3/d does not have opaque flag
	xattrOpaque, err := system.Lgetxattr(filepath.Join(td, "l3", "d"), "trusted.overlay.opaque")
	if err != nil {
		return errors.Wrap(err, "failed to read opaque flag on upper layer")
	}
	if string(xattrOpaque) == "y" {
		return errors.New("opaque flag erroneously copied up, consider update to kernel 4.8 or later to fix")
	}

	// rename "d1" to "d2"
	if err := os.Rename(filepath.Join(td, mergedDirName, "d1"), filepath.Join(td, mergedDirName, "d2")); err != nil {
		// if rename failed with syscall.EXDEV, the kernel doesn't have CONFIG_OVERLAY_FS_REDIRECT_DIR enabled
		if err.(*os.LinkError).Err == syscall.EXDEV {
			return nil
		}
		return errors.Wrap(err, "failed to rename dir in merged directory")
	}
	// get the xattr of "d2"
	xattrRedirect, err := system.Lgetxattr(filepath.Join(td, "l3", "d2"), "trusted.overlay.redirect")
	if err != nil {
		return errors.Wrap(err, "failed to read redirect flag on upper layer")
	}

	if string(xattrRedirect) == "d1" {
		return errors.New("kernel has CONFIG_OVERLAY_FS_REDIRECT_DIR enabled")
	}

	return nil
}
<|endoftext|>


<commit_before>package container

import (
	"fmt"
	"github.com/labstack/gommon/log"
	"reflect"
)

type singletonContainer struct {
	values map[string]interface{}
}

var instance *singletonContainer

func init() {
	instance = &singletonContainer{
		make(map[string]interface{}),
	}
}

func Set(val interface{}) {
	var key string
	if reflect.TypeOf(val).Kind() == reflect.Ptr {
		key = reflect.Indirect(reflect.ValueOf(val)).Type().String()
	} else {
		key = reflect.ValueOf(val).Type().String()
	}

	log.Info(fmt.Sprintf("added %s to components container.", key))
	instance.values[key] = val
}

func Get(ptr interface{}) {
	val := reflect.ValueOf(ptr)
	key := reflect.Indirect(val).Type().String()
	component := instance.values[key]
	if component == nil {
		if reflect.Indirect(val).Type().Kind() != reflect.Interface {
			log.Warn(fmt.Sprintf("component not found. such type of %s.", key))
			return
		}
		for _, component := range instance.values {
			value := reflect.ValueOf(component)
			elm := reflect.ValueOf(ptr).Elem()
			if value.Type().Implements(elm.Type()) {
				log.Info(fmt.Sprintf("found component of %s .", key))
				elm.Set(value)
				return
			}
		}
		return
	}
	log.Debug(fmt.Sprintf("found component of %s .", key))

	elm := reflect.ValueOf(ptr).Elem()
	if reflect.TypeOf(component).Kind() == reflect.Ptr {
		elm.Set(reflect.Indirect(reflect.ValueOf(component)))
	} else {
		elm.Set(reflect.ValueOf(component))
	}
}
<commit_msg>Change log level<commit_after>package container

import (
	"fmt"
	"github.com/labstack/gommon/log"
	"reflect"
)

type singletonContainer struct {
	values map[string]interface{}
}

var instance *singletonContainer

func init() {
	instance = &singletonContainer{
		make(map[string]interface{}),
	}
}

func Set(val interface{}) {
	var key string
	if reflect.TypeOf(val).Kind() == reflect.Ptr {
		key = reflect.Indirect(reflect.ValueOf(val)).Type().String()
	} else {
		key = reflect.ValueOf(val).Type().String()
	}

	log.Debug(fmt.Sprintf("added %s to components container.", key))
	instance.values[key] = val
}

func Get(ptr interface{}) {
	val := reflect.ValueOf(ptr)
	key := reflect.Indirect(val).Type().String()
	component := instance.values[key]
	if component == nil {
		if reflect.Indirect(val).Type().Kind() != reflect.Interface {
			log.Warn(fmt.Sprintf("component not found. such type of %s.", key))
			return
		}
		for _, component := range instance.values {
			value := reflect.ValueOf(component)
			elm := reflect.ValueOf(ptr).Elem()
			if value.Type().Implements(elm.Type()) {
				log.Debug(fmt.Sprintf("found component of %s .", key))
				elm.Set(value)
				return
			}
		}
		return
	}
	log.Debug(fmt.Sprintf("found component of %s .", key))

	elm := reflect.ValueOf(ptr).Elem()
	if reflect.TypeOf(component).Kind() == reflect.Ptr {
		elm.Set(reflect.Indirect(reflect.ValueOf(component)))
	} else {
		elm.Set(reflect.ValueOf(component))
	}
}
<|endoftext|>


<commit_before>package usecases

import (
	"time"

	"github.com/webitel/cdr/src/conf"
	"github.com/webitel/cdr/src/entity"
	"github.com/webitel/cdr/src/logger"
)

type Legs struct {
	LegsB []entity.ElasticCdr "json:legs_b"
}

type CheckCalls func(bulkCount uint32, state uint8)

func (interactor *CdrInteractor) RunElastic() {
	if interactor.ElasticCdrBRepository == nil || interactor.ElasticCdrARepository == nil || interactor.SqlCdrBRepository == nil || interactor.SqlCdrARepository == nil {
		return
	}
	elasticConfig := conf.GetElastic()
	if !elasticConfig.Enabled {
		return
	}
	go LegListener(interactor.CheckLegsAFromSql, elasticConfig.RequestTimeout, elasticConfig.BulkCount)
	go LegListener(interactor.CheckLegsBFromSql, elasticConfig.RequestTimeout, elasticConfig.BulkCount)
	logger.Notice("Elastic module: start listening...")
}

func LegListener(checkCalls CheckCalls, timeout uint32, bulkCount uint32) {
	promise := time.Millisecond * time.Duration(timeout)
	ticker := time.NewTicker(promise)
	errorTicker := time.NewTicker(promise * 10)
	for {
		select {
		case <-ticker.C:
			{
				go checkCalls(bulkCount, 0)
			}
		case <-errorTicker.C:
			{
				go checkCalls(bulkCount, 4)
			}
		}

	}
}

func (interactor *CdrInteractor) CheckLegsAFromSql(bulkCount uint32, state uint8) {
	var calls []entity.ElasticCdr
	cdr, err := interactor.SqlCdrARepository.SelectPackByState(bulkCount, state, "stored")
	if err != nil {
		logger.Error(err.Error())
		return
	}
	if len(cdr) == 0 {
		//log.Println("Elastic module: listening Leg A from pg...")
		return
	}
	if err := interactor.SqlCdrARepository.UpdateState(cdr, 1, 0, "stored"); err != nil {
		logger.Error(err.Error())
		return
	}
	var (
		eCall entity.ElasticCdr
		iCall interface{}
	)
	for _, item := range cdr {
		iCall, err = readBytes(item.Event)
		if err != nil {
			interactor.SqlCdrARepository.UpdateState(cdr, 4, 0, "stored")
			logger.Error(err.Error())
			return
		}
		eCall, err = ParseToCdr(iCall)
		if err != nil {
			interactor.SqlCdrARepository.UpdateState(cdr, 4, 0, "stored")
			logger.Error(err.Error())
			return
		}
		calls = append(calls, eCall)
	}
	if err, errCalls, succCalls := interactor.ElasticCdrARepository.InsertDocs(calls); err != nil {
		if errCalls != nil && len(errCalls) > 0 {
			interactor.SqlCdrARepository.UpdateState(errCalls, 4, 0, "stored")
			if succCalls != nil && len(succCalls) > 0 {
				interactor.SqlCdrARepository.UpdateState(succCalls, 2, uint64(time.Now().UnixNano()/1000000), "stored")
			}
		} else {
			interactor.SqlCdrARepository.UpdateState(cdr, 4, 0, "stored")
		}
		logger.Error(err.Error())
	} else {
		logger.Notice("Elastic: items stored [%s, %v]", "Leg A", len(calls))
		interactor.SqlCdrARepository.UpdateState(cdr, 2, uint64(time.Now().UnixNano()/1000000), "stored")
	}
	//log.Println("Elastic module: listening Leg A from pg...")
}

func (interactor *CdrInteractor) CheckLegsBFromSql(bulkCount uint32, state uint8) {
	var calls []entity.ElasticCdr
	cdr, err := interactor.SqlCdrBRepository.SelectPackByState(bulkCount, state, "stored")
	if err != nil {
		logger.Error(err.Error())
		return
	}
	if len(cdr) == 0 {
		//log.Println("Elastic module: listening Leg B from pg...")
		return
	}
	if err := interactor.SqlCdrBRepository.UpdateState(cdr, 1, 0, "stored"); err != nil {
		logger.Error(err.Error())
		return
	}
	var (
		eCall entity.ElasticCdr
		iCall interface{}
	)
	for _, item := range cdr {
		iCall, err = readBytes(item.Event)
		if err != nil {
			interactor.SqlCdrBRepository.UpdateState(cdr, 4, 0, "stored")
			logger.Error(err.Error())
			return
		}
		eCall, err = ParseToCdr(iCall)
		if err != nil {
			interactor.SqlCdrBRepository.UpdateState(cdr, 4, 0, "stored")
			logger.Error(err.Error())
			return
		}
		calls = append(calls, eCall)
	}
	if err, errCalls, succCalls := interactor.ElasticCdrBRepository.InsertDocs(calls); err != nil {
		if errCalls != nil && len(errCalls) > 0 {
			interactor.SqlCdrBRepository.UpdateState(errCalls, 4, 0, "stored")
			if succCalls != nil && len(succCalls) > 0 {
				interactor.SqlCdrBRepository.UpdateState(succCalls, 2, uint64(time.Now().UnixNano()/1000000), "stored")
			}
		} else {
			interactor.SqlCdrBRepository.UpdateState(cdr, 4, 0, "stored")
		}
		logger.Error(err.Error())
	} else {
		logger.Notice("Elastic: items stored [%s, %v]", "Leg B", len(calls))
		interactor.SqlCdrBRepository.UpdateState(cdr, 2, uint64(time.Now().UnixNano()/1000000), "stored")
	}
	//log.Println("Elastic module: listening Leg B from pg...")
}
<commit_msg>method get calls<commit_after>package usecases

import (
	"time"

	"github.com/webitel/cdr/src/conf"
	"github.com/webitel/cdr/src/entity"
	"github.com/webitel/cdr/src/logger"
)

type CheckCalls func(bulkCount uint32, state uint8)

func (interactor *CdrInteractor) RunElastic() {
	if interactor.ElasticCdrBRepository == nil || interactor.ElasticCdrARepository == nil || interactor.SqlCdrBRepository == nil || interactor.SqlCdrARepository == nil {
		return
	}
	elasticConfig := conf.GetElastic()
	if !elasticConfig.Enabled {
		return
	}
	go LegListener(interactor.CheckLegsAFromSql, elasticConfig.RequestTimeout, elasticConfig.BulkCount)
	go LegListener(interactor.CheckLegsBFromSql, elasticConfig.RequestTimeout, elasticConfig.BulkCount)
	logger.Notice("Elastic module: start listening...")
}

func LegListener(checkCalls CheckCalls, timeout uint32, bulkCount uint32) {
	promise := time.Millisecond * time.Duration(timeout)
	ticker := time.NewTicker(promise)
	errorTicker := time.NewTicker(promise * 10)
	for {
		select {
		case <-ticker.C:
			{
				go checkCalls(bulkCount, 0)
			}
		case <-errorTicker.C:
			{
				go checkCalls(bulkCount, 4)
			}
		}

	}
}

func (interactor *CdrInteractor) CheckLegsAFromSql(bulkCount uint32, state uint8) {

	cdr, err := interactor.SqlCdrARepository.SelectPackByState(bulkCount, state, "stored")
	if err != nil {
		logger.Error(err.Error())
		return
	}
	if len(cdr) == 0 {
		//log.Println("Elastic module: listening Leg A from pg...")
		return
	}
	if err := interactor.SqlCdrARepository.UpdateState(cdr, 1, 0, "stored"); err != nil {
		logger.Error(err.Error())
		return
	}
	calls, err := getCalls(interactor.SqlCdrARepository, cdr)
	if err != nil {
		return
	}
	if err, errCalls, succCalls := interactor.ElasticCdrARepository.InsertDocs(calls); err != nil {
		if errCalls != nil && len(errCalls) > 0 {
			interactor.SqlCdrARepository.UpdateState(errCalls, 4, 0, "stored")
			if succCalls != nil && len(succCalls) > 0 {
				interactor.SqlCdrARepository.UpdateState(succCalls, 2, uint64(time.Now().UnixNano()/1000000), "stored")
				logger.Notice("Elastic: items stored [%s, %v]", "Leg A", len(succCalls))
			}
			logger.Error("Elastic: failed to store items [%s, %v]", "Leg A", len(errCalls))
		} else {
			interactor.SqlCdrARepository.UpdateState(cdr, 4, 0, "stored")
		}
		logger.Error(err.Error())
	} else {
		logger.Notice("Elastic: items stored [%s, %v]", "Leg A", len(calls))
		interactor.SqlCdrARepository.UpdateState(cdr, 2, uint64(time.Now().UnixNano()/1000000), "stored")
	}
	//log.Println("Elastic module: listening Leg A from pg...")
}

func (interactor *CdrInteractor) CheckLegsBFromSql(bulkCount uint32, state uint8) {

	cdr, err := interactor.SqlCdrBRepository.SelectPackByState(bulkCount, state, "stored")
	if err != nil {
		logger.Error(err.Error())
		return
	}
	if len(cdr) == 0 {
		//log.Println("Elastic module: listening Leg B from pg...")
		return
	}
	if err := interactor.SqlCdrBRepository.UpdateState(cdr, 1, 0, "stored"); err != nil {
		logger.Error(err.Error())
		return
	}
	calls, err := getCalls(interactor.SqlCdrBRepository, cdr)
	if err != nil {
		return
	}
	if err, errCalls, succCalls := interactor.ElasticCdrBRepository.InsertDocs(calls); err != nil {
		if errCalls != nil && len(errCalls) > 0 {
			interactor.SqlCdrBRepository.UpdateState(errCalls, 4, 0, "stored")
			if succCalls != nil && len(succCalls) > 0 {
				interactor.SqlCdrBRepository.UpdateState(succCalls, 2, uint64(time.Now().UnixNano()/1000000), "stored")
				logger.Notice("Elastic: items stored [%s, %v]", "Leg B", len(succCalls))
			}
			logger.Error("Elastic: failed to store items [%s, %v]", "Leg B", len(errCalls))
		} else {
			interactor.SqlCdrBRepository.UpdateState(cdr, 4, 0, "stored")
		}
		logger.Error(err.Error())
	} else {
		logger.Notice("Elastic: items stored [%s, %v]", "Leg B", len(calls))
		interactor.SqlCdrBRepository.UpdateState(cdr, 2, uint64(time.Now().UnixNano()/1000000), "stored")
	}
	//log.Println("Elastic module: listening Leg B from pg...")
}

func getCalls(repo entity.SqlCdrRepository, cdr []entity.SqlCdr) ([]entity.ElasticCdr, error) {
	var calls []entity.ElasticCdr
	var (
		eCall entity.ElasticCdr
		iCall interface{}
		err   error
	)
	for _, item := range cdr {
		iCall, err = readBytes(item.Event)
		if err != nil {
			repo.UpdateState(cdr, 4, 0, "stored")
			logger.Error(err.Error())
			return nil, err
		}
		eCall, err = ParseToCdr(iCall)
		if err != nil {
			repo.UpdateState(cdr, 4, 0, "stored")
			logger.Error(err.Error())
			return nil, err
		}
		calls = append(calls, eCall)
	}
	return calls, nil
}
<|endoftext|>


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

package machiner_test

import (
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api"
	"launchpad.net/juju-core/state/api/params"
	statetesting "launchpad.net/juju-core/state/testing"
	coretesting "launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/testing/checkers"
	stdtesting "testing"
)

func TestAll(t *stdtesting.T) {
	coretesting.MgoTestPackage(t)
}

type machinerSuite struct {
	testing.JujuConnSuite
	st      *api.State
	machine *state.Machine
}

var _ = Suite(&machinerSuite{})

func (s *machinerSuite) SetUpTest(c *C) {
	s.JujuConnSuite.SetUpTest(c)

	// Create a machine so we can log in as its agent.
	var err error
	s.machine, err = s.State.AddMachine("series", state.JobHostUnits)
	c.Assert(err, IsNil)
	err = s.machine.SetProvisioned("foo", "fake_nonce", nil)
	c.Assert(err, IsNil)
	err = s.machine.SetPassword("password")
	c.Assert(err, IsNil)
	s.st = s.OpenAPIAsMachine(c, s.machine.Tag(), "password", "fake_nonce")
}

func (s *machinerSuite) TearDownTest(c *C) {
	err := s.st.Close()
	c.Assert(err, IsNil)
	s.JujuConnSuite.TearDownTest(c)
}

func (s *machinerSuite) TestMachineAndMachineId(c *C) {
	machine, err := s.st.Machiner().Machine("machine-42")
	c.Assert(err, ErrorMatches, "permission denied")
	c.Assert(params.ErrCode(err), Equals, params.CodeUnauthorized)
	c.Assert(machine, IsNil)

	machine, err = s.st.Machiner().Machine("machine-0")
	c.Assert(err, IsNil)
	c.Assert(machine.Tag(), Equals, "machine-0")
}

func (s *machinerSuite) TestSetStatus(c *C) {
	machine, err := s.st.Machiner().Machine("machine-0")
	c.Assert(err, IsNil)

	status, info, err := s.machine.Status()
	c.Assert(err, IsNil)
	c.Assert(status, Equals, params.StatusPending)
	c.Assert(info, Equals, "")

	err = machine.SetStatus(params.StatusStarted, "blah")
	c.Assert(err, IsNil)

	status, info, err = s.machine.Status()
	c.Assert(err, IsNil)
	c.Assert(status, Equals, params.StatusStarted)
	c.Assert(info, Equals, "blah")
}

func (s *machinerSuite) TestEnsureDead(c *C) {
	c.Assert(s.machine.Life(), Equals, state.Alive)

	machine, err := s.st.Machiner().Machine("machine-0")
	c.Assert(err, IsNil)

	err = machine.EnsureDead()
	c.Assert(err, IsNil)

	err = s.machine.Refresh()
	c.Assert(err, IsNil)
	c.Assert(s.machine.Life(), Equals, state.Dead)

	err = machine.EnsureDead()
	c.Assert(err, IsNil)
	err = s.machine.Refresh()
	c.Assert(err, IsNil)
	c.Assert(s.machine.Life(), Equals, state.Dead)

	err = s.machine.Remove()
	c.Assert(err, IsNil)
	err = s.machine.Refresh()
	c.Assert(err, checkers.Satisfies, errors.IsNotFoundError)

	err = machine.EnsureDead()
	c.Assert(err, ErrorMatches, "machine 0 not found")
	c.Assert(params.ErrCode(err), Equals, params.CodeNotFound)
}

func (s *machinerSuite) TestRefresh(c *C) {
	machine, err := s.st.Machiner().Machine("machine-0")
	c.Assert(err, IsNil)
	c.Assert(machine.Life(), Equals, params.Alive)

	err = machine.EnsureDead()
	c.Assert(err, IsNil)
	c.Assert(machine.Life(), Equals, params.Alive)

	err = machine.Refresh()
	c.Assert(err, IsNil)
	c.Assert(machine.Life(), Equals, params.Dead)
}

func (s *machinerSuite) TestWatch(c *C) {
	machine, err := s.st.Machiner().Machine("machine-0")
	c.Assert(err, IsNil)
	c.Assert(machine.Life(), Equals, params.Alive)

	w, err := machine.Watch()
	c.Assert(err, IsNil)
	defer statetesting.AssertStop(c, w)
	wc := statetesting.NewNotifyWatcherC(c, s.BackingState, w)

	// Initial event.
	wc.AssertOneChange()

	// Change something other than the lifecycle and make sure it's
	// not detected.
	err = machine.SetStatus(params.StatusStarted, "not really")
	c.Assert(err, IsNil)
	wc.AssertNoChange()

	// Make the machine dying and check it's detected.
	err = machine.EnsureDead()
	c.Assert(err, IsNil)
	wc.AssertOneChange()

	statetesting.AssertStop(c, w)
	wc.AssertClosed()
}
<commit_msg>Changes after review<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package machiner_test

import (
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api"
	"launchpad.net/juju-core/state/api/params"
	statetesting "launchpad.net/juju-core/state/testing"
	coretesting "launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/testing/checkers"
	stdtesting "testing"
)

func TestAll(t *stdtesting.T) {
	coretesting.MgoTestPackage(t)
}

type machinerSuite struct {
	testing.JujuConnSuite
	st      *api.State
	machine *state.Machine
}

var _ = Suite(&machinerSuite{})

func (s *machinerSuite) SetUpTest(c *C) {
	s.JujuConnSuite.SetUpTest(c)

	// Create a machine so we can log in as its agent.
	var err error
	s.machine, err = s.State.AddMachine("series", state.JobHostUnits)
	c.Assert(err, IsNil)
	err = s.machine.SetProvisioned("foo", "fake_nonce", nil)
	c.Assert(err, IsNil)
	err = s.machine.SetPassword("password")
	c.Assert(err, IsNil)
	s.st = s.OpenAPIAsMachine(c, s.machine.Tag(), "password", "fake_nonce")
}

func (s *machinerSuite) TearDownTest(c *C) {
	err := s.st.Close()
	c.Assert(err, IsNil)
	s.JujuConnSuite.TearDownTest(c)
}

func (s *machinerSuite) TestMachineAndMachineId(c *C) {
	machine, err := s.st.Machiner().Machine("machine-42")
	c.Assert(err, ErrorMatches, "permission denied")
	c.Assert(params.ErrCode(err), Equals, params.CodeUnauthorized)
	c.Assert(machine, IsNil)

	machine, err = s.st.Machiner().Machine("machine-0")
	c.Assert(err, IsNil)
	c.Assert(machine.Tag(), Equals, "machine-0")
}

func (s *machinerSuite) TestSetStatus(c *C) {
	machine, err := s.st.Machiner().Machine("machine-0")
	c.Assert(err, IsNil)

	status, info, err := s.machine.Status()
	c.Assert(err, IsNil)
	c.Assert(status, Equals, params.StatusPending)
	c.Assert(info, Equals, "")

	err = machine.SetStatus(params.StatusStarted, "blah")
	c.Assert(err, IsNil)

	status, info, err = s.machine.Status()
	c.Assert(err, IsNil)
	c.Assert(status, Equals, params.StatusStarted)
	c.Assert(info, Equals, "blah")
}

func (s *machinerSuite) TestEnsureDead(c *C) {
	c.Assert(s.machine.Life(), Equals, state.Alive)

	machine, err := s.st.Machiner().Machine("machine-0")
	c.Assert(err, IsNil)

	err = machine.EnsureDead()
	c.Assert(err, IsNil)

	err = s.machine.Refresh()
	c.Assert(err, IsNil)
	c.Assert(s.machine.Life(), Equals, state.Dead)

	err = machine.EnsureDead()
	c.Assert(err, IsNil)
	err = s.machine.Refresh()
	c.Assert(err, IsNil)
	c.Assert(s.machine.Life(), Equals, state.Dead)

	err = s.machine.Remove()
	c.Assert(err, IsNil)
	err = s.machine.Refresh()
	c.Assert(err, checkers.Satisfies, errors.IsNotFoundError)

	err = machine.EnsureDead()
	c.Assert(err, ErrorMatches, "machine 0 not found")
	c.Assert(params.ErrCode(err), Equals, params.CodeNotFound)
}

func (s *machinerSuite) TestRefresh(c *C) {
	machine, err := s.st.Machiner().Machine("machine-0")
	c.Assert(err, IsNil)
	c.Assert(machine.Life(), Equals, params.Alive)

	err = machine.EnsureDead()
	c.Assert(err, IsNil)
	c.Assert(machine.Life(), Equals, params.Alive)

	err = machine.Refresh()
	c.Assert(err, IsNil)
	c.Assert(machine.Life(), Equals, params.Dead)
}

func (s *machinerSuite) TestWatch(c *C) {
	machine, err := s.st.Machiner().Machine("machine-0")
	c.Assert(err, IsNil)
	c.Assert(machine.Life(), Equals, params.Alive)

	w, err := machine.Watch()
	c.Assert(err, IsNil)
	defer statetesting.AssertStop(c, w)
	wc := statetesting.NewNotifyWatcherC(c, s.BackingState, w)

	// Initial event.
	wc.AssertOneChange()

	// Change something other than the lifecycle and make sure it's
	// not detected.
	err = machine.SetStatus(params.StatusStarted, "not really")
	c.Assert(err, IsNil)
	wc.AssertNoChange()

	// Make the machine dead and check it's detected.
	err = machine.EnsureDead()
	c.Assert(err, IsNil)
	wc.AssertOneChange()

	statetesting.AssertStop(c, w)
	wc.AssertClosed()
}
<|endoftext|>


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

package backoff

import (
	"context"
	"errors"
	"testing"
	"time"
)

func TestBackoff(t *testing.T) {
	b := Backoff{
		Min:    time.Duration(1),
		Max:    time.Duration(100),
		Factor: 2,
	}
	for _, test := range []struct {
		b     Backoff
		times int
		want  time.Duration
	}{
		{b, 1, time.Duration(1)},
		{b, 2, time.Duration(2)},
		{b, 3, time.Duration(4)},
		{b, 4, time.Duration(8)},
		{b, 8, time.Duration(100)},
	} {
		test.b.Reset()
		var got time.Duration
		for i := 0; i < test.times; i++ {
			got = test.b.Duration()
		}
		if got != test.want {
			t.Errorf("Duration() %v times: %v, want %v", test.times, got, test.want)
		}
	}
}

func TestJitter(t *testing.T) {
	b := Backoff{
		Min:    time.Duration(1) * time.Second,
		Max:    time.Duration(100) * time.Second,
		Factor: 2,
		Jitter: true,
	}
	for _, test := range []struct {
		b     Backoff
		times int
		min   time.Duration
		max   time.Duration
	}{
		{b, 1, time.Duration(1) * time.Second, time.Duration(2) * time.Second},
		{b, 2, time.Duration(2) * time.Second, time.Duration(4) * time.Second},
		{b, 3, time.Duration(4) * time.Second, time.Duration(8) * time.Second},
		{b, 4, time.Duration(8) * time.Second, time.Duration(16) * time.Second},
		{b, 8, time.Duration(100) * time.Second, time.Duration(200) * time.Second},
	} {
		test.b.Reset()
		var got1 time.Duration
		for i := 0; i < test.times; i++ {
			got1 = test.b.Duration()
		}
		if got1 < test.min || got1 > test.max {
			t.Errorf("Duration() %v times, want  %v < %v < %v", test.times, test.min, got1, test.max)
		}

		// Ensure a random value is being produced.
		test.b.Reset()
		var got2 time.Duration
		for i := 0; i < test.times; i++ {
			got2 = test.b.Duration()
		}
		if got1 == got2 {
			t.Errorf("Duration() %v times == Duration() %v times, want  %v != %v",
				test.times, test.times, got1, got2)
		}
	}
}

func TestRetry(t *testing.T) {
	b := Backoff{
		Min:    time.Duration(50) * time.Millisecond,
		Max:    time.Duration(200) * time.Millisecond,
		Factor: 2,
	}

	// callCount is used by some test funcs to count how many times they've been called.
	var callCount int
	// ctx used by Retry(), declared here to that test.ctxFunc can set it.
	var ctx context.Context
	var cancel context.CancelFunc

	for _, test := range []struct {
		name    string
		f       func() error
		ctxFunc func()
		wantErr bool
	}{
		{
			name: "func that immediately succeeds",
			f:    func() error { return nil },
		},
		{
			name: "func that succeeds on second attempt",
			f: func() error {
				callCount++
				if callCount == 1 {
					return errors.New("error")
				}
				return nil
			},
		},
		{
			name: "func that takes too long to succeed",
			f: func() error {
				// Cancel the context and return an error. This func will succeed on
				// any future calls, but it should not be retried due to the context
				// being cancelled.
				if ctx.Err() == nil {
					cancel()
					return errors.New("error")
				}
				return nil
			},
			wantErr: true,
		},
		{
			name: "context done before Retry() called",
			f: func() error {
				return nil
			},
			ctxFunc: func() {
				ctx, cancel = context.WithCancel(context.Background())
				cancel()
			},
			wantErr: true,
		},
	} {
		if test.ctxFunc != nil {
			test.ctxFunc()
		} else {
			ctx, cancel = context.WithCancel(context.Background())
		}

		callCount = 0
		err := b.Retry(ctx, test.f)
		cancel()
		if gotErr := err != nil; gotErr != test.wantErr {
			t.Errorf("%v: Retry() = %v, want err? %v", test.name, err, test.wantErr)
		}
	}
}
<commit_msg>Remove unnecessary casts to time.Duration in backoff_test.go<commit_after>// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package backoff

import (
	"context"
	"errors"
	"testing"
	"time"
)

func TestBackoff(t *testing.T) {
	b := Backoff{
		Min:    time.Duration(1),
		Max:    time.Duration(100),
		Factor: 2,
	}
	for _, test := range []struct {
		b     Backoff
		times int
		want  time.Duration
	}{
		{b, 1, time.Duration(1)},
		{b, 2, time.Duration(2)},
		{b, 3, time.Duration(4)},
		{b, 4, time.Duration(8)},
		{b, 8, time.Duration(100)},
	} {
		test.b.Reset()
		var got time.Duration
		for i := 0; i < test.times; i++ {
			got = test.b.Duration()
		}
		if got != test.want {
			t.Errorf("Duration() %v times: %v, want %v", test.times, got, test.want)
		}
	}
}

func TestJitter(t *testing.T) {
	b := Backoff{
		Min:    1 * time.Second,
		Max:    100 * time.Second,
		Factor: 2,
		Jitter: true,
	}
	for _, test := range []struct {
		b     Backoff
		times int
		min   time.Duration
		max   time.Duration
	}{
		{b, 1, 1 * time.Second, 2 * time.Second},
		{b, 2, 2 * time.Second, 4 * time.Second},
		{b, 3, 4 * time.Second, 8 * time.Second},
		{b, 4, 8 * time.Second, 16 * time.Second},
		{b, 8, 100 * time.Second, 200 * time.Second},
	} {
		test.b.Reset()
		var got1 time.Duration
		for i := 0; i < test.times; i++ {
			got1 = test.b.Duration()
		}
		if got1 < test.min || got1 > test.max {
			t.Errorf("Duration() %v times, want  %v < %v < %v", test.times, test.min, got1, test.max)
		}

		// Ensure a random value is being produced.
		test.b.Reset()
		var got2 time.Duration
		for i := 0; i < test.times; i++ {
			got2 = test.b.Duration()
		}
		if got1 == got2 {
			t.Errorf("Duration() %v times == Duration() %v times, want  %v != %v",
				test.times, test.times, got1, got2)
		}
	}
}

func TestRetry(t *testing.T) {
	b := Backoff{
		Min:    50 * time.Millisecond,
		Max:    200 * time.Millisecond,
		Factor: 2,
	}

	// callCount is used by some test funcs to count how many times they've been called.
	var callCount int
	// ctx used by Retry(), declared here to that test.ctxFunc can set it.
	var ctx context.Context
	var cancel context.CancelFunc

	for _, test := range []struct {
		name    string
		f       func() error
		ctxFunc func()
		wantErr bool
	}{
		{
			name: "func that immediately succeeds",
			f:    func() error { return nil },
		},
		{
			name: "func that succeeds on second attempt",
			f: func() error {
				callCount++
				if callCount == 1 {
					return errors.New("error")
				}
				return nil
			},
		},
		{
			name: "func that takes too long to succeed",
			f: func() error {
				// Cancel the context and return an error. This func will succeed on
				// any future calls, but it should not be retried due to the context
				// being cancelled.
				if ctx.Err() == nil {
					cancel()
					return errors.New("error")
				}
				return nil
			},
			wantErr: true,
		},
		{
			name: "context done before Retry() called",
			f: func() error {
				return nil
			},
			ctxFunc: func() {
				ctx, cancel = context.WithCancel(context.Background())
				cancel()
			},
			wantErr: true,
		},
	} {
		if test.ctxFunc != nil {
			test.ctxFunc()
		} else {
			ctx, cancel = context.WithCancel(context.Background())
		}

		callCount = 0
		err := b.Retry(ctx, test.f)
		cancel()
		if gotErr := err != nil; gotErr != test.wantErr {
			t.Errorf("%v: Retry() = %v, want err? %v", test.name, err, test.wantErr)
		}
	}
}
<|endoftext|>


<commit_before>// gotify is a client library for the Spotify API
package gotify

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"
)

const OauthAuthorizeURL = "https://accounts.spotify.com/authorize"
const OauthTokenUrl = "https://accounts.spotify.com/api/token"

type SpotifyOauth struct {
	ClientId     string
	ClientSecret string
	RedirectUri  string
	State        string
	Scope        string
	CachePath    string
}

type Token struct {
	AccessToken  string `json:"access_token"`
	TokenType    string `json:"token_type"`
	ExpiresAt    time.Time
	RefreshToken string `json:"refresh_token"`
	TTL          int    `json:"expires_in"`
}

func (t Token) isExpired() bool {
	if !(int(time.Since(t.ExpiresAt)) >= 0) {
		return false
	}
	return true
}

// Reads a cached token
func GetCachedToken(oauth SpotifyOauth) (Token, error) {
	var token Token
	var err error
	if oauth.CachePath != "" {
		cachedData, err := ioutil.ReadFile(oauth.CachePath)
		if err == nil {
			err = json.Unmarshal(cachedData, &token)
			if err == nil {
				if token.isExpired() {
					token, err = RefreshAccessToken(token.RefreshToken, oauth)
				}
				if err == nil {
					return token, nil
				}
			}
		}
	}
	return Token{}, err
}

// Takes a SpotifyOauth struct and returns the appropriate AuthorizeUrl for
// your request
func GetAuthorizeURL(oauth SpotifyOauth) (string, error) {
	var Url *url.URL
	Url, err := url.Parse(OauthAuthorizeURL)
	if err == nil {
		parameters := url.Values{}
		parameters.Add("client_id", oauth.ClientId)
		parameters.Add("response_type", "code")
		parameters.Add("redirect_uri", oauth.RedirectUri)
		if oauth.Scope != "" {
			parameters.Add("scope", oauth.Scope)
		}
		if oauth.State != "" {
			parameters.Add("state", oauth.State)
		}
		Url.RawQuery = parameters.Encode()
		return Url.String(), nil
	}
	return "", err
}

// Takes the authorization code and a SpotifyOauth
// Returns an access token
func GetAccessToken(code string, oauth SpotifyOauth) (Token, error) {
	var err error
	parameters := url.Values{}
	parameters.Add("redirect_uri", oauth.RedirectUri)
	parameters.Add("code", code)
	parameters.Add("grant_type", "authorization_code")

	token, err := sendAccessTokenRequest(parameters, oauth)
	if err == nil {
		err = saveTokenInfo(token, oauth)
		if err == nil {
			return token, nil
		}
	}
	return Token{}, err
}

// Refreshes an expired AccessToken
func RefreshAccessToken(refreshToken string, oauth SpotifyOauth) (Token, error) {
	parameters := url.Values{}
	parameters.Add("refresh_token", refreshToken)
	parameters.Add("grant_type", "refresh_token")

	token, err := sendAccessTokenRequest(parameters, oauth)
	if err == nil {
		if token.RefreshToken == "" {
			token.RefreshToken = refreshToken
		}
		err = saveTokenInfo(token, oauth)
		if err == nil {
			return token, nil
		}
	}
	return Token{}, err
}

// Parses the response code from from the query string when user is redirected
// back to the application
func ParseResponseCode(response string) (string, error) {
	u, err := url.Parse(response)
	if err != nil {
		return "", err
	}
	q, _ := url.ParseQuery(u.RawQuery)
	code := q["code"][0]
	return code, nil
}

func sendAccessTokenRequest(parameters url.Values, oauth SpotifyOauth) (Token, error) {
	client := &http.Client{}
	req, err := http.NewRequest("POST", OauthTokenUrl, strings.NewReader(parameters.Encode()))
	if err == nil {
		req.SetBasicAuth(oauth.ClientId, oauth.ClientSecret)
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
		resp, err := client.Do(req)
		if err == nil {
			if resp.StatusCode == http.StatusOK {
				var token Token
				err = json.NewDecoder(resp.Body).Decode(&token)
				if err == nil {
					token.ExpiresAt = time.Now().Add(time.Duration(token.TTL) * time.Second)
					return token, nil
				}
			} else {
				err = errors.New(resp.Status)
			}
		}
	}
	return Token{}, err
}

func saveTokenInfo(token Token, oauth SpotifyOauth) error {
	var err error
	if oauth.CachePath != "" {
		marshaledToken, err := json.Marshal(token)
		if err == nil {
			err = ioutil.WriteFile(oauth.CachePath, marshaledToken, 0x777)
			return nil
		}
	}
	return err
}
<commit_msg>Added defer resp.Body.Close()<commit_after>// gotify is a client library for the Spotify API
package gotify

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"
)

const OauthAuthorizeURL = "https://accounts.spotify.com/authorize"
const OauthTokenUrl = "https://accounts.spotify.com/api/token"

type SpotifyOauth struct {
	ClientId     string
	ClientSecret string
	RedirectUri  string
	State        string
	Scope        string
	CachePath    string
}

type Token struct {
	AccessToken  string `json:"access_token"`
	TokenType    string `json:"token_type"`
	ExpiresAt    time.Time
	RefreshToken string `json:"refresh_token"`
	TTL          int    `json:"expires_in"`
}

func (t Token) isExpired() bool {
	if !(int(time.Since(t.ExpiresAt)) >= 0) {
		return false
	}
	return true
}

// Reads a cached token
func GetCachedToken(oauth SpotifyOauth) (Token, error) {
	var token Token
	var err error
	if oauth.CachePath != "" {
		cachedData, err := ioutil.ReadFile(oauth.CachePath)
		if err == nil {
			err = json.Unmarshal(cachedData, &token)
			if token.isExpired() {
				token, err = RefreshAccessToken(token.RefreshToken, oauth)
			}
			if err == nil {
				return token, nil
			}
		}
	}
	return Token{}, err
}

// Takes a SpotifyOauth struct and returns the appropriate AuthorizeUrl for
// your request
func GetAuthorizeURL(oauth SpotifyOauth) (string, error) {
	var Url *url.URL
	Url, err := url.Parse(OauthAuthorizeURL)
	if err == nil {
		parameters := url.Values{}
		parameters.Add("client_id", oauth.ClientId)
		parameters.Add("response_type", "code")
		parameters.Add("redirect_uri", oauth.RedirectUri)
		if oauth.Scope != "" {
			parameters.Add("scope", oauth.Scope)
		}
		if oauth.State != "" {
			parameters.Add("state", oauth.State)
		}
		Url.RawQuery = parameters.Encode()
		return Url.String(), nil
	}
	return "", err
}

// Takes the authorization code and a SpotifyOauth
// Returns an access token
func GetAccessToken(code string, oauth SpotifyOauth) (Token, error) {
	var err error
	parameters := url.Values{}
	parameters.Add("redirect_uri", oauth.RedirectUri)
	parameters.Add("code", code)
	parameters.Add("grant_type", "authorization_code")

	token, err := sendAccessTokenRequest(parameters, oauth)
	if err == nil {
		err = saveTokenInfo(token, oauth)
		if err == nil {
			return token, nil
		}
	}
	return Token{}, err
}

// Refreshes an expired AccessToken
func RefreshAccessToken(refreshToken string, oauth SpotifyOauth) (Token, error) {
	parameters := url.Values{}
	parameters.Add("refresh_token", refreshToken)
	parameters.Add("grant_type", "refresh_token")

	token, err := sendAccessTokenRequest(parameters, oauth)
	if err == nil {
		if token.RefreshToken == "" {
			token.RefreshToken = refreshToken
		}
		err = saveTokenInfo(token, oauth)
		if err == nil {
			return token, nil
		}
	}
	return Token{}, err
}

// Parses the response code from from the query string when user is redirected
// back to the application
func ParseResponseCode(response string) (string, error) {
	u, err := url.Parse(response)
	if err != nil {
		return "", err
	}
	q, _ := url.ParseQuery(u.RawQuery)
	code := q["code"][0]
	return code, nil
}

func sendAccessTokenRequest(parameters url.Values, oauth SpotifyOauth) (Token, error) {
	client := &http.Client{}
	req, err := http.NewRequest("POST", OauthTokenUrl, strings.NewReader(parameters.Encode()))
	if err == nil {
		req.SetBasicAuth(oauth.ClientId, oauth.ClientSecret)
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
		resp, err := client.Do(req)
		defer resp.Body.Close()
		if err == nil {
			if resp.StatusCode == http.StatusOK {
				var token Token
				err = json.NewDecoder(resp.Body).Decode(&token)
				if err == nil {
					token.ExpiresAt = time.Now().Add(time.Duration(token.TTL) * time.Second)
					return token, nil
				}
			} else {
				err = errors.New(resp.Status)
			}
		}
	}
	return Token{}, err
}

func saveTokenInfo(token Token, oauth SpotifyOauth) error {
	var err error
	if oauth.CachePath != "" {
		marshaledToken, err := json.Marshal(token)
		if err == nil {
			err = ioutil.WriteFile(oauth.CachePath, marshaledToken, 0x777)
			return nil
		}
	}
	return err
}
<|endoftext|>


<commit_before>// Discordgo - Discord bindings for Go
// Available at https://github.com/bwmarrin/discordgo

// Copyright 2015-2016 Bruce Marriner <bruce@sqls.net>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file contains functions related to Discord OAuth2 endpoints

package discordgo

// ------------------------------------------------------------------------------------------------
// Code specific to Discord OAuth2 Applications
// ------------------------------------------------------------------------------------------------

// An Application struct stores values for a Discord OAuth2 Application
type Application struct {
	ID                  string    `json:"id,omitempty"`
	Name                string    `json:"name"`
	Description         string    `json:"description,omitempty"`
	Icon                string    `json:"icon,omitempty"`
	Secret              string    `json:"secret,omitempty"`
	RedirectURIs        *[]string `json:"redirect_uris,omitempty"`
	BotRequireCodeGrant bool      `json:"bot_require_code_grant,omitempty"`
	BotPublic           bool      `json:"bot_public,omitempty"`
	RPCApplicationState int       `json:"rpc_application_state,omitempty"`
	Flags               int       `json:"flags,omitempty"`
	Owner               *User     `json:"owner"`
	Bot                 *User     `json:"bot"`
}

// Application returns an Application structure of a specific Application
//   appID : The ID of an Application
func (s *Session) Application(appID string) (st *Application, err error) {

	body, err := s.RequestWithBucketID("GET", EndpointApplication(appID), nil, EndpointApplication(""))
	if err != nil {
		return
	}

	err = unmarshal(body, &st)
	return
}

// Applications returns all applications for the authenticated user
func (s *Session) Applications() (st []*Application, err error) {

	body, err := s.RequestWithBucketID("GET", EndpointApplications, nil, EndpointApplications)
	if err != nil {
		return
	}

	err = unmarshal(body, &st)
	return
}

// ApplicationCreate creates a new Application
//    name : Name of Application / Bot
//    uris : Redirect URIs (Not required)
func (s *Session) ApplicationCreate(ap *Application) (st *Application, err error) {

	data := struct {
		Name         string    `json:"name"`
		Description  string    `json:"description"`
		RedirectURIs *[]string `json:"redirect_uris,omitempty"`
	}{ap.Name, ap.Description, ap.RedirectURIs}

	body, err := s.RequestWithBucketID("POST", EndpointApplications, data, EndpointApplications)
	if err != nil {
		return
	}

	err = unmarshal(body, &st)
	return
}

// ApplicationUpdate updates an existing Application
//   var : desc
func (s *Session) ApplicationUpdate(appID string, ap *Application) (st *Application, err error) {

	data := struct {
		Name         string    `json:"name"`
		Description  string    `json:"description"`
		RedirectURIs *[]string `json:"redirect_uris,omitempty"`
	}{ap.Name, ap.Description, ap.RedirectURIs}

	body, err := s.RequestWithBucketID("PUT", EndpointApplication(appID), data, EndpointApplication(""))
	if err != nil {
		return
	}

	err = unmarshal(body, &st)
	return
}

// ApplicationDelete deletes an existing Application
//   appID : The ID of an Application
func (s *Session) ApplicationDelete(appID string) (err error) {

	_, err = s.RequestWithBucketID("DELETE", EndpointApplication(appID), nil, EndpointApplication(""))
	if err != nil {
		return
	}

	return
}

// Asset struct stores values for an asset of an application
type Asset struct {
	Type int    `json:"type"`
	ID   string `json:"id"`
	Name string `json:"name"`
}

// ApplicationAssets returns an application's assets
func (s *Session) ApplicationAssets(appID string) (ass []*Asset, err error) {

	body, err := s.RequestWithBucketID("GET", EndpointApplicationAssets(appID), nil, EndpointApplicationAssets(""))
	if err != nil {
		return
	}

	err = unmarshal(body, &ass)
	return
}

// ------------------------------------------------------------------------------------------------
// Code specific to Discord OAuth2 Application Bots
// ------------------------------------------------------------------------------------------------

// ApplicationBotCreate creates an Application Bot Account
//
//   appID : The ID of an Application
//
// NOTE: func name may change, if I can think up something better.
func (s *Session) ApplicationBotCreate(appID string) (st *User, err error) {

	body, err := s.RequestWithBucketID("POST", EndpointApplicationsBot(appID), nil, EndpointApplicationsBot(""))
	if err != nil {
		return
	}

	err = unmarshal(body, &st)
	return
}
<commit_msg>Add team data to applications (#787)<commit_after>// Discordgo - Discord bindings for Go
// Available at https://github.com/bwmarrin/discordgo

// Copyright 2015-2016 Bruce Marriner <bruce@sqls.net>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file contains functions related to Discord OAuth2 endpoints

package discordgo

// ------------------------------------------------------------------------------------------------
// Code specific to Discord OAuth2 Applications
// ------------------------------------------------------------------------------------------------

// The MembershipState represents whether the user is in the team or has been invited into it
type MembershipState int

// Constants for the different stages of the MembershipState
const (
	MembershipStateInvited MembershipState = iota + 1
	MembershipStateAccepted
)

// A TeamMember struct stores values for a single Team Member, extending the normal User data - note that the user field is partial
type TeamMember struct {
	User            *User           `json:"user"`
	TeamID          string          `json:"team_id"`
	MembershipState MembershipState `json:"membership_state"`
	Permissions     []string        `json:"permissions"`
}

// A Team struct stores the members of a Discord Developer Team as well as some metadata about it
type Team struct {
	ID          string        `json:"id"`
	Name        string        `json:"name"`
	Description string        `json:"description"`
	Icon        string        `json:"icon"`
	OwnerID     string        `json:"owner_user_id"`
	Members     []*TeamMember `json:"members"`
}

// An Application struct stores values for a Discord OAuth2 Application
type Application struct {
	ID                  string    `json:"id,omitempty"`
	Name                string    `json:"name"`
	Description         string    `json:"description,omitempty"`
	Icon                string    `json:"icon,omitempty"`
	Secret              string    `json:"secret,omitempty"`
	RedirectURIs        *[]string `json:"redirect_uris,omitempty"`
	BotRequireCodeGrant bool      `json:"bot_require_code_grant,omitempty"`
	BotPublic           bool      `json:"bot_public,omitempty"`
	RPCApplicationState int       `json:"rpc_application_state,omitempty"`
	Flags               int       `json:"flags,omitempty"`
	Owner               *User     `json:"owner"`
	Bot                 *User     `json:"bot"`
	Team                *Team     `json:"team"`
}

// Application returns an Application structure of a specific Application
//   appID : The ID of an Application
func (s *Session) Application(appID string) (st *Application, err error) {

	body, err := s.RequestWithBucketID("GET", EndpointApplication(appID), nil, EndpointApplication(""))
	if err != nil {
		return
	}

	err = unmarshal(body, &st)
	return
}

// Applications returns all applications for the authenticated user
func (s *Session) Applications() (st []*Application, err error) {

	body, err := s.RequestWithBucketID("GET", EndpointApplications, nil, EndpointApplications)
	if err != nil {
		return
	}

	err = unmarshal(body, &st)
	return
}

// ApplicationCreate creates a new Application
//    name : Name of Application / Bot
//    uris : Redirect URIs (Not required)
func (s *Session) ApplicationCreate(ap *Application) (st *Application, err error) {

	data := struct {
		Name         string    `json:"name"`
		Description  string    `json:"description"`
		RedirectURIs *[]string `json:"redirect_uris,omitempty"`
	}{ap.Name, ap.Description, ap.RedirectURIs}

	body, err := s.RequestWithBucketID("POST", EndpointApplications, data, EndpointApplications)
	if err != nil {
		return
	}

	err = unmarshal(body, &st)
	return
}

// ApplicationUpdate updates an existing Application
//   var : desc
func (s *Session) ApplicationUpdate(appID string, ap *Application) (st *Application, err error) {

	data := struct {
		Name         string    `json:"name"`
		Description  string    `json:"description"`
		RedirectURIs *[]string `json:"redirect_uris,omitempty"`
	}{ap.Name, ap.Description, ap.RedirectURIs}

	body, err := s.RequestWithBucketID("PUT", EndpointApplication(appID), data, EndpointApplication(""))
	if err != nil {
		return
	}

	err = unmarshal(body, &st)
	return
}

// ApplicationDelete deletes an existing Application
//   appID : The ID of an Application
func (s *Session) ApplicationDelete(appID string) (err error) {

	_, err = s.RequestWithBucketID("DELETE", EndpointApplication(appID), nil, EndpointApplication(""))
	if err != nil {
		return
	}

	return
}

// Asset struct stores values for an asset of an application
type Asset struct {
	Type int    `json:"type"`
	ID   string `json:"id"`
	Name string `json:"name"`
}

// ApplicationAssets returns an application's assets
func (s *Session) ApplicationAssets(appID string) (ass []*Asset, err error) {

	body, err := s.RequestWithBucketID("GET", EndpointApplicationAssets(appID), nil, EndpointApplicationAssets(""))
	if err != nil {
		return
	}

	err = unmarshal(body, &ass)
	return
}

// ------------------------------------------------------------------------------------------------
// Code specific to Discord OAuth2 Application Bots
// ------------------------------------------------------------------------------------------------

// ApplicationBotCreate creates an Application Bot Account
//
//   appID : The ID of an Application
//
// NOTE: func name may change, if I can think up something better.
func (s *Session) ApplicationBotCreate(appID string) (st *User, err error) {

	body, err := s.RequestWithBucketID("POST", EndpointApplicationsBot(appID), nil, EndpointApplicationsBot(""))
	if err != nil {
		return
	}

	err = unmarshal(body, &st)
	return
}
<|endoftext|>


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

// Package oauth2 contains Negroni middleware to provide
// user login via an OAuth 2.0 backend.

package oauth2

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"time"

	"github.com/codegangsta/negroni"
	sessions "github.com/goincremental/negroni-sessions"
	"golang.org/x/oauth2"
)

const (
	codeRedirect = 302
	keyToken     = "oauth2_token"
	keyNextPage  = "next"
	keyState     = "state"
	keyProvider  = "provider"

	// PathLogin sets the path to handle OAuth 2.0 logins.
	pathLogin = "/login"
	// PathLogout sets to handle OAuth 2.0 logouts.
	pathLogout = "/logout"
	// PathCallback sets the path to handle callback from OAuth 2.0 backend
	// to exchange credentials.
	pathCallback = "/oauth2callback"
	// PathError sets the path to handle error cases.
	pathError = "/oauth2error"
	// the provider
	provider = "provider"
)

type Oauth2Handler struct {
	Provider     string
	PathLogin    string
	PathLogout   string
	PathCallback string
	PathError    string
	Config       *oauth2.Config
}

func (h *Oauth2Handler) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	s := sessions.GetSession(r)

	if r.Method == "GET" {
		switch r.URL.Path {
		case h.PathLogin:
			h.login(s, w, r)
		case h.PathLogout:
			h.logout(s, w, r)
		case h.PathCallback:
			h.handleOAuth2Callback(s, w, r)
		default:
			next(w, r)
		}
	} else {
		next(w, r)
	}
}

func (h *Oauth2Handler) login(s sessions.Session, w http.ResponseWriter, r *http.Request) {
	next := r.URL.Query().Get(keyNextPage)

	if s.Get(keyToken) == nil {
		// User is not logged in.
		if next == "" {
			next = "/"
		}

		state := newState()
		// store the next url and state token in the session
		s.Set(keyState, state)
		s.Set(keyNextPage, next)
		s.Set(keyProvider, h.Provider)
		http.Redirect(w, r, h.Config.AuthCodeURL(state, oauth2.AccessTypeOffline), http.StatusFound)
		return
	}
	// No need to login, redirect to the next page.
	http.Redirect(w, r, next, http.StatusFound)
}

func (h *Oauth2Handler) logout(s sessions.Session, w http.ResponseWriter, r *http.Request) {
	next := r.URL.Query().Get(keyNextPage)
	s.Delete(keyToken)
	s.Delete("email")
	http.Redirect(w, r, next, http.StatusFound)
}

func (h *Oauth2Handler) handleOAuth2Callback(s sessions.Session, w http.ResponseWriter, r *http.Request) {
	providedState := r.URL.Query().Get("state")
	fmt.Printf("Got state from request %s\n", providedState)

	//verify that the provided state is the state we generated
	//if it is not, then redirect to the error page
	originalState := s.Get(keyState)
	fmt.Printf("Got state from session %s\n", originalState)
	if providedState != originalState {
		http.Redirect(w, r, h.PathError, http.StatusFound)
		return
	}

	next := s.Get(keyNextPage).(string)
	fmt.Printf("Got a next page from the session: %s\n", next)
	code := r.URL.Query().Get("code")
	t, err := h.Config.Exchange(oauth2.NoContext, code)
	if err != nil {
		// Pass the error message, or allow dev to provide its own
		// error handler.
		http.Redirect(w, r, h.PathError, http.StatusFound)
		return
	}

	// Store the credentials in the session.
	val, _ := json.Marshal(t)
	s.Set(keyToken, val)
	http.Redirect(w, r, next, http.StatusFound)
}

// Handler that redirects user to the login page
// if user is not logged in.
func (h *Oauth2Handler) LoginRequired() negroni.HandlerFunc {
	return func(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
		token := GetToken(r)
		if token == nil || !token.Valid() {
			// Set token to null to avoid redirection loop
			SetToken(r, nil)
			next := url.QueryEscape(r.URL.RequestURI())
			http.Redirect(rw, r, h.PathLogin+"?"+keyNextPage+"="+next, http.StatusFound)
		} else {
			next(rw, r)
		}
	}
}

type Config oauth2.Config

// Tokens Represents a container that contains
// user's OAuth 2.0 access and refresh tokens.
type Tokens interface {
	Access() string
	Refresh() string
	Valid() bool
	ExpiryTime() time.Time
	ExtraData(string) interface{}
	Get() oauth2.Token
}

type token struct {
	oauth2.Token
}

func (t *token) ExtraData(key string) interface{} {
	return t.Extra(key)
}

// Returns the access token.
func (t *token) Access() string {
	return t.AccessToken
}

// Returns the refresh token.
func (t *token) Refresh() string {
	return t.RefreshToken
}

// Returns whether the access token is
// expired or not.
func (t *token) Valid() bool {
	if t == nil {
		return true
	}
	return t.Token.Valid()
}

// Returns the expiry time of the user's
// access token.
func (t *token) ExpiryTime() time.Time {
	return t.Expiry
}

// String returns the string representation of the token.
func (t *token) String() string {
	return fmt.Sprintf("tokens: %v", t)
}

// Returns oauth2.Token.
func (t *token) Get() oauth2.Token {
	return t.Token
}

// Returns a new Google OAuth 2.0 backend endpoint.
func Google(config *Config) negroni.Handler {
	authUrl := "https://accounts.google.com/o/oauth2/auth"
	tokenUrl := "https://accounts.google.com/o/oauth2/token"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

// Returns a new Github OAuth 2.0 backend endpoint.
func Github(config *Config) negroni.Handler {
	authUrl := "https://github.com/login/oauth/authorize"
	tokenUrl := "https://github.com/login/oauth/access_token"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

func Facebook(config *Config) negroni.Handler {
	authUrl := "https://www.facebook.com/dialog/oauth"
	tokenUrl := "https://graph.facebook.com/oauth/access_token"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

func LinkedIn(config *Config) negroni.Handler {
	authUrl := "https://www.linkedin.com/uas/oauth2/authorization"
	tokenUrl := "https://www.linkedin.com/uas/oauth2/accessToken"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

// Returns a generic OAuth 2.0 backend endpoint.
func NewOAuth2Provider(config *Config, authUrl, tokenUrl string) negroni.Handler {
	c := &oauth2.Config{
		ClientID:     config.ClientID,
		ClientSecret: config.ClientSecret,
		Scopes:       config.Scopes,
		RedirectURL:  config.RedirectURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  authUrl,
			TokenURL: tokenUrl,
		},
	}

	h := &Oauth2Handler{
		Provider:     provider,
		PathLogin:    pathLogin,
		PathLogout:   pathLogout,
		PathCallback: pathCallback,
		PathError:    pathError,
		Config:       c,
	}

	return h
}

func GetToken(r *http.Request) Tokens {
	s := sessions.GetSession(r)
	t := unmarshallToken(s)

	//not doing this doesn't pass through the
	//nil return, causing a test to fail - not sure why??
	if t == nil {
		return nil
	} else {
		return t
	}
}

func SetToken(r *http.Request, t interface{}) {
	s := sessions.GetSession(r)
	val, _ := json.Marshal(t)
	s.Set(keyToken, val)
	//Check immediately to see if the token is expired
	tk := unmarshallToken(s)
	if tk != nil {
		// check if the access token is expired
		if !tk.Valid() && tk.Refresh() == "" {
			s.Delete(keyToken)
			tk = nil
		}
	}
}

func newState() string {
	var p [16]byte
	_, err := rand.Read(p[:])
	if err != nil {
		panic(err)
	}
	return hex.EncodeToString(p[:])
}

func unmarshallToken(s sessions.Session) *token {

	if s.Get(keyToken) == nil {
		return nil
	}

	data := s.Get(keyToken).([]byte)
	var tk oauth2.Token
	json.Unmarshal(data, &tk)
	return &token{tk}

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

// Package oauth2 contains Negroni middleware to provide
// user login via an OAuth 2.0 backend.

package oauth2

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"time"

	"github.com/codegangsta/negroni"
	sessions "github.com/goincremental/negroni-sessions"
	"golang.org/x/oauth2"
)

const (
	codeRedirect = 302
	keyToken     = "oauth2_token"
	keyNextPage  = "next"
	keyState     = "state"
	keyProvider  = "provider"

	// PathLogin sets the path to handle OAuth 2.0 logins.
	pathLogin = "/login"
	// PathLogout sets to handle OAuth 2.0 logouts.
	pathLogout = "/logout"
	// PathCallback sets the path to handle callback from OAuth 2.0 backend
	// to exchange credentials.
	pathCallback = "/oauth2callback"
	// PathError sets the path to handle error cases.
	pathError = "/oauth2error"
	// the provider
	provider = "provider"
)

type Oauth2Handler struct {
	Provider     string
	PathLogin    string
	PathLogout   string
	PathCallback string
	PathError    string
	Config       *oauth2.Config
}

func (h *Oauth2Handler) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	s := sessions.GetSession(r)

	if r.Method == "GET" {
		switch r.URL.Path {
		case h.PathLogin:
			h.login(s, w, r)
		case h.PathLogout:
			h.logout(s, w, r)
		case h.PathCallback:
			h.handleOAuth2Callback(s, w, r)
		default:
			next(w, r)
		}
	} else {
		next(w, r)
	}
}

func (h *Oauth2Handler) login(s sessions.Session, w http.ResponseWriter, r *http.Request) {
	next := r.URL.Query().Get(keyNextPage)

	if s.Get(keyToken) == nil {
		// User is not logged in.
		if next == "" {
			next = "/"
		}

		state := newState()
		// store the next url and state token in the session
		s.Set(keyState, state)
		s.Set(keyNextPage, next)
		s.Set(keyProvider, h.Provider)
		http.Redirect(w, r, h.Config.AuthCodeURL(state, oauth2.AccessTypeOffline), http.StatusFound)
		return
	}
	// No need to login, redirect to the next page.
	http.Redirect(w, r, next, http.StatusFound)
}

func (h *Oauth2Handler) logout(s sessions.Session, w http.ResponseWriter, r *http.Request) {
	next := r.URL.Query().Get(keyNextPage)
	s.Delete(keyToken)
	s.Delete("email")
	s.Delete("provider")
	http.Redirect(w, r, next, http.StatusFound)
}

func (h *Oauth2Handler) handleOAuth2Callback(s sessions.Session, w http.ResponseWriter, r *http.Request) {
	providedState := r.URL.Query().Get("state")
	fmt.Printf("Got state from request %s\n", providedState)

	//verify that the provided state is the state we generated
	//if it is not, then redirect to the error page
	originalState := s.Get(keyState)
	fmt.Printf("Got state from session %s\n", originalState)
	if providedState != originalState {
		http.Redirect(w, r, h.PathError, http.StatusFound)
		return
	}

	next := s.Get(keyNextPage).(string)
	fmt.Printf("Got a next page from the session: %s\n", next)
	code := r.URL.Query().Get("code")
	t, err := h.Config.Exchange(oauth2.NoContext, code)
	if err != nil {
		// Pass the error message, or allow dev to provide its own
		// error handler.
		http.Redirect(w, r, h.PathError, http.StatusFound)
		return
	}

	// Store the credentials in the session.
	val, _ := json.Marshal(t)
	s.Set(keyToken, val)
	http.Redirect(w, r, next, http.StatusFound)
}

// Handler that redirects user to the login page
// if user is not logged in.
func (h *Oauth2Handler) LoginRequired() negroni.HandlerFunc {
	return func(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
		token := GetToken(r)
		if token == nil || !token.Valid() {
			// Set token to null to avoid redirection loop
			SetToken(r, nil)
			next := url.QueryEscape(r.URL.RequestURI())
			http.Redirect(rw, r, h.PathLogin+"?"+keyNextPage+"="+next, http.StatusFound)
		} else {
			next(rw, r)
		}
	}
}

type Config oauth2.Config

// Tokens Represents a container that contains
// user's OAuth 2.0 access and refresh tokens.
type Tokens interface {
	Access() string
	Refresh() string
	Valid() bool
	ExpiryTime() time.Time
	ExtraData(string) interface{}
	Get() oauth2.Token
}

type token struct {
	oauth2.Token
}

func (t *token) ExtraData(key string) interface{} {
	return t.Extra(key)
}

// Returns the access token.
func (t *token) Access() string {
	return t.AccessToken
}

// Returns the refresh token.
func (t *token) Refresh() string {
	return t.RefreshToken
}

// Returns whether the access token is
// expired or not.
func (t *token) Valid() bool {
	if t == nil {
		return true
	}
	return t.Token.Valid()
}

// Returns the expiry time of the user's
// access token.
func (t *token) ExpiryTime() time.Time {
	return t.Expiry
}

// String returns the string representation of the token.
func (t *token) String() string {
	return fmt.Sprintf("tokens: %v", t)
}

// Returns oauth2.Token.
func (t *token) Get() oauth2.Token {
	return t.Token
}

// Returns a new Google OAuth 2.0 backend endpoint.
func Google(config *Config) negroni.Handler {
	authUrl := "https://accounts.google.com/o/oauth2/auth"
	tokenUrl := "https://accounts.google.com/o/oauth2/token"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

// Returns a new Github OAuth 2.0 backend endpoint.
func Github(config *Config) negroni.Handler {
	authUrl := "https://github.com/login/oauth/authorize"
	tokenUrl := "https://github.com/login/oauth/access_token"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

func Facebook(config *Config) negroni.Handler {
	authUrl := "https://www.facebook.com/dialog/oauth"
	tokenUrl := "https://graph.facebook.com/oauth/access_token"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

func LinkedIn(config *Config) negroni.Handler {
	authUrl := "https://www.linkedin.com/uas/oauth2/authorization"
	tokenUrl := "https://www.linkedin.com/uas/oauth2/accessToken"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

// Returns a generic OAuth 2.0 backend endpoint.
func NewOAuth2Provider(config *Config, authUrl, tokenUrl string) negroni.Handler {
	c := &oauth2.Config{
		ClientID:     config.ClientID,
		ClientSecret: config.ClientSecret,
		Scopes:       config.Scopes,
		RedirectURL:  config.RedirectURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  authUrl,
			TokenURL: tokenUrl,
		},
	}

	h := &Oauth2Handler{
		Provider:     provider,
		PathLogin:    pathLogin,
		PathLogout:   pathLogout,
		PathCallback: pathCallback,
		PathError:    pathError,
		Config:       c,
	}

	return h
}

func GetToken(r *http.Request) Tokens {
	s := sessions.GetSession(r)
	t := unmarshallToken(s)

	//not doing this doesn't pass through the
	//nil return, causing a test to fail - not sure why??
	if t == nil {
		return nil
	} else {
		return t
	}
}

func SetToken(r *http.Request, t interface{}) {
	s := sessions.GetSession(r)
	val, _ := json.Marshal(t)
	s.Set(keyToken, val)
	//Check immediately to see if the token is expired
	tk := unmarshallToken(s)
	if tk != nil {
		// check if the access token is expired
		if !tk.Valid() && tk.Refresh() == "" {
			s.Delete(keyToken)
			tk = nil
		}
	}
}

func newState() string {
	var p [16]byte
	_, err := rand.Read(p[:])
	if err != nil {
		panic(err)
	}
	return hex.EncodeToString(p[:])
}

func unmarshallToken(s sessions.Session) *token {

	if s.Get(keyToken) == nil {
		return nil
	}

	data := s.Get(keyToken).([]byte)
	var tk oauth2.Token
	json.Unmarshal(data, &tk)
	return &token{tk}

}
<|endoftext|>