<commit_before>// Copyright 2018 The Fractal Team Authors
// This file is part of the fractal project.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

package protoadaptor

import (
	"fmt"
	"reflect"
	"sync"
	"time"

	"github.com/ethereum/go-ethereum/log"
	router "github.com/fractalplatform/fractal/event"
	"github.com/fractalplatform/fractal/p2p"
	"github.com/fractalplatform/fractal/utils/rlp"
)

type pack struct {
	From     string
	To       string
	Typecode uint32
	Payload  []byte
}

type remotePeer struct {
	peer *p2p.Peer
	ws   p2p.MsgReadWriter
}

// ProtoAdaptor is subprotocol on p2p
type ProtoAdaptor struct {
	p2p.Server
	peerMangaer
	event  chan *router.Event
	loopWG sync.WaitGroup
	quit   chan struct{}
	subs   []router.Subscription
}

// NewProtoAdaptor return new ProtoAdaptor
func NewProtoAdaptor(config *p2p.Config) *ProtoAdaptor {
	adaptor := &ProtoAdaptor{
		Server: p2p.Server{
			Config: config,
		},
		peerMangaer: peerMangaer{
			activePeers: make(map[[8]byte]*remotePeer),
			station:     nil,
		},
		event: make(chan *router.Event),
		quit:  make(chan struct{}),
		subs:  make([]router.Subscription, 0, 2),
	}
	adaptor.peerMangaer.station = router.NewBroadcastStation("broadcast", &adaptor.peerMangaer)
	adaptor.Server.Config.Protocols = adaptor.Protocols()
	return adaptor
}

// Start start p2p protocol adaptor
func (adaptor *ProtoAdaptor) Start() error {
	router.StationRegister(adaptor.peerMangaer.station)
	router.AdaptorRegister(adaptor)
	sub1 := router.Subscribe(nil, adaptor.event, router.DisconectCtrl, nil)
	sub2 := router.Subscribe(nil, adaptor.event, router.OneMinuteLimited, nil)
	adaptor.subs = append(adaptor.subs, sub1, sub2)
	adaptor.loopWG.Add(1)
	go func() {
		adaptor.adaptorEvent()
		adaptor.loopWG.Done()
	}()
	return adaptor.Server.Start()
}

func (adaptor *ProtoAdaptor) adaptorEvent() {
	timer := time.NewTimer(time.Second)
	if adaptor.PeerPeriod == 0 || adaptor.MaxPeers == 0 {
		timer.Stop()
	}
	for {
		select {
		case <-adaptor.quit:
			return
		case <-timer.C:
			if adaptor.PeerCount() == adaptor.MaxPeers {
				peer := router.WorstStation().Data().(*remotePeer)
				endtime := time.Now().Add(time.Minute)
				adaptor.Server.AddBadNode(peer.peer.Node(), &endtime) // AddBadNode also disconnect the peer
			}
			timer.Reset(time.Duration(adaptor.PeerPeriod) * time.Millisecond)
		case e := <-adaptor.event:
			switch e.Typecode {
			case router.DisconectCtrl:
				peer := e.Data.(router.Station).Data().(*remotePeer)
				peer.peer.Disconnect(p2p.DiscSubprotocolError)
			case router.OneMinuteLimited:
				peer := e.Data.(router.Station).Data().(*remotePeer)
				endtime := time.Now().Add(time.Minute)
				adaptor.Server.AddBadNode(peer.peer.Node(), &endtime) // AddBadNode also disconnect the peer
			}
		}
	}
}

func GetFnode(station router.Station) string {
	remote := station.Data().(*remotePeer)
	return remote.peer.Node().String()
}

func (adaptor *ProtoAdaptor) adaptorLoop(peer *p2p.Peer, ws p2p.MsgReadWriter) error {
	remote := remotePeer{ws: ws, peer: peer}
	log.Info("New remote station", "detail", remote.peer.String())
	station := router.NewRemoteStation(string(remote.peer.ID().Bytes()[:8]), &remote)
	adaptor.peerMangaer.addActivePeer(&remote)
	router.StationRegister(station)
	url := remote.peer.Node().String()
	router.SendTo(station, nil, router.NewPeerNotify, &url)
	defer func() {
		adaptor.peerMangaer.delActivePeer(&remote)
		router.StationUnregister(station)
		url := remote.peer.Node().String()
		router.SendTo(station, nil, router.DelPeerNotify, &url)
	}()

	monitor := make(map[int][]int64)
	for {
		msg, err := ws.ReadMsg()
		if err != nil {
			return err
		}
		pack := pack{}
		if err := msg.Decode(&pack); err != nil {
			return err
		}
		e, err := pack2event(&pack, station)
		if err != nil {
			return err
		}
		router.AddNetIn(station, 1)
		if checkDDOS(monitor, e) {
			//router.SendTo(nil, nil, router.OneMinuteLimited, e.From)
			log.Warn("DDos detection", "peer", remote.peer.String(), "typecode", e.Typecode, "count", monitor[e.Typecode][1])
			//return p2p.DiscDDOS
		}
		if e.To == nil && len(pack.To) != 0 {
			continue
		}
		router.SendEvent(e)
	}
}

func checkDDOS(m map[int][]int64, e *router.Event) bool {
	t := e.Typecode
	limit := int64(router.GetDDosLimit(t)) * 10
	if limit == 0 {
		return false
	}

	if len(m[t]) == 0 {
		m[t] = make([]int64, 2)
	}
	//m[t][0] time
	//m[t][1] request per second
	if m[t][0] == time.Now().Unix() {
		m[t][1]++
	} else {
		if m[t][1] > limit {
			return true
		}
		m[t][0] = time.Now().Unix()
		m[t][1] = 1
	}
	return false
}

// Protocols .
func (adaptor *ProtoAdaptor) Protocols() []p2p.Protocol {
	return []p2p.Protocol{
		p2p.Protocol{
			Name:    "FractalTest",
			Version: 1,
			Length:  1,
			Run:     adaptor.adaptorLoop,
		},
	}
}

// Stop .
func (adaptor *ProtoAdaptor) Stop() {
	close(adaptor.quit)
	for _, sub := range adaptor.subs {
		sub.Unsubscribe()
	}
	adaptor.Server.Stop()
	adaptor.loopWG.Wait()
	log.Info("P2P networking stopped")
}

// SendOut .
func (adaptor *ProtoAdaptor) SendOut(e *router.Event) error {
	if e.To.IsBroadcast() {
		adaptor.msgBroadcast(e)
		return nil
	}
	return adaptor.msgSend(e)
}

func (adaptor *ProtoAdaptor) msgSend(e *router.Event) error {
	pack, err := event2pack(e)
	if err != nil {
		return err
	}
	router.AddNetOut(e.To, 1)
	return p2p.Send(e.To.Data().(*remotePeer).ws, 0, pack)
}

func (adaptor *ProtoAdaptor) msgBroadcast(e *router.Event) {
	te := *e
	te.To = nil
	te.From = nil
	pack, err := event2pack(&te)
	if err != nil {
		return
	}

	send := func(peer *remotePeer) {
		//router.AddNetOut(x,1)
		p2p.Send(peer.ws, 0, pack)
	}
	if e.To.Data() != nil {
		e.To.Data().(*peerMangaer).mapActivePeer(send)
		return
	}
	adaptor.peerMangaer.mapActivePeer(send)
}

func event2pack(e *router.Event) (*pack, error) {
	buf, err := rlp.EncodeToBytes(e.Data)
	if err != nil {
		return nil, err
	}
	from := ""
	if e.From != nil {
		from = e.From.Name()
	}
	to := ""
	if e.To != nil {
		to = e.To.Name()[8:]
	}
	return &pack{
		From:     from,
		To:       to,
		Typecode: uint32(e.Typecode),
		Payload:  buf,
	}, nil
}

func pack2event(pack *pack, station router.Station) (*router.Event, error) {
	var elem interface{}

	isPtr := false
	typ := router.GetTypeByCode(int(pack.Typecode))
	if typ == nil {
		return nil, fmt.Errorf("unknow typecode: %d", pack.Typecode)
	}

	//for typ.Kind() == reflect.Ptr {
	if typ.Kind() == reflect.Ptr {
		isPtr = true
		typ = typ.Elem()
	}
	obj := reflect.New(typ)
	if err := rlp.DecodeBytes(pack.Payload, obj.Interface()); err != nil {
		return nil, err
	}
	if isPtr {
		elem = obj.Interface()
	} else {
		elem = obj.Elem().Interface()
	}

	if pack.From != "" {
		station = router.NewRemoteStation(station.Name()+pack.From, station.Data())
	}
	return &router.Event{
		From:     station,
		To:       router.GetStationByName(pack.To),
		Typecode: int(pack.Typecode),
		Data:     elem,
	}, nil
}
<commit_msg>p2p/protoadaptor: fixed ddos check error (#519)<commit_after>// Copyright 2018 The Fractal Team Authors
// This file is part of the fractal project.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

package protoadaptor

import (
	"fmt"
	"reflect"
	"sync"
	"time"

	"github.com/ethereum/go-ethereum/log"
	router "github.com/fractalplatform/fractal/event"
	"github.com/fractalplatform/fractal/p2p"
	"github.com/fractalplatform/fractal/utils/rlp"
)

type pack struct {
	From     string
	To       string
	Typecode uint32
	Payload  []byte
}

type remotePeer struct {
	peer *p2p.Peer
	ws   p2p.MsgReadWriter
}

// ProtoAdaptor is subprotocol on p2p
type ProtoAdaptor struct {
	p2p.Server
	peerMangaer
	event  chan *router.Event
	loopWG sync.WaitGroup
	quit   chan struct{}
	subs   []router.Subscription
}

// NewProtoAdaptor return new ProtoAdaptor
func NewProtoAdaptor(config *p2p.Config) *ProtoAdaptor {
	adaptor := &ProtoAdaptor{
		Server: p2p.Server{
			Config: config,
		},
		peerMangaer: peerMangaer{
			activePeers: make(map[[8]byte]*remotePeer),
			station:     nil,
		},
		event: make(chan *router.Event),
		quit:  make(chan struct{}),
		subs:  make([]router.Subscription, 0, 2),
	}
	adaptor.peerMangaer.station = router.NewBroadcastStation("broadcast", &adaptor.peerMangaer)
	adaptor.Server.Config.Protocols = adaptor.Protocols()
	return adaptor
}

// Start start p2p protocol adaptor
func (adaptor *ProtoAdaptor) Start() error {
	router.StationRegister(adaptor.peerMangaer.station)
	router.AdaptorRegister(adaptor)
	sub1 := router.Subscribe(nil, adaptor.event, router.DisconectCtrl, nil)
	sub2 := router.Subscribe(nil, adaptor.event, router.OneMinuteLimited, nil)
	adaptor.subs = append(adaptor.subs, sub1, sub2)
	adaptor.loopWG.Add(1)
	go func() {
		adaptor.adaptorEvent()
		adaptor.loopWG.Done()
	}()
	return adaptor.Server.Start()
}

func (adaptor *ProtoAdaptor) adaptorEvent() {
	timer := time.NewTimer(time.Second)
	if adaptor.PeerPeriod == 0 || adaptor.MaxPeers == 0 {
		timer.Stop()
	}
	for {
		select {
		case <-adaptor.quit:
			return
		case <-timer.C:
			if adaptor.PeerCount() == adaptor.MaxPeers {
				peer := router.WorstStation().Data().(*remotePeer)
				endtime := time.Now().Add(time.Minute)
				adaptor.Server.AddBadNode(peer.peer.Node(), &endtime) // AddBadNode also disconnect the peer
			}
			timer.Reset(time.Duration(adaptor.PeerPeriod) * time.Millisecond)
		case e := <-adaptor.event:
			switch e.Typecode {
			case router.DisconectCtrl:
				peer := e.Data.(router.Station).Data().(*remotePeer)
				peer.peer.Disconnect(p2p.DiscSubprotocolError)
			case router.OneMinuteLimited:
				peer := e.Data.(router.Station).Data().(*remotePeer)
				endtime := time.Now().Add(time.Minute)
				adaptor.Server.AddBadNode(peer.peer.Node(), &endtime) // AddBadNode also disconnect the peer
			}
		}
	}
}

func GetFnode(station router.Station) string {
	remote := station.Data().(*remotePeer)
	return remote.peer.Node().String()
}

func (adaptor *ProtoAdaptor) adaptorLoop(peer *p2p.Peer, ws p2p.MsgReadWriter) error {
	remote := remotePeer{ws: ws, peer: peer}
	log.Info("New remote station", "detail", remote.peer.String())
	station := router.NewRemoteStation(string(remote.peer.ID().Bytes()[:8]), &remote)
	adaptor.peerMangaer.addActivePeer(&remote)
	router.StationRegister(station)
	url := remote.peer.Node().String()
	router.SendTo(station, nil, router.NewPeerNotify, &url)
	defer func() {
		adaptor.peerMangaer.delActivePeer(&remote)
		router.StationUnregister(station)
		url := remote.peer.Node().String()
		router.SendTo(station, nil, router.DelPeerNotify, &url)
	}()

	monitor := make(map[int][]int64)
	for {
		msg, err := ws.ReadMsg()
		if err != nil {
			return err
		}
		pack := pack{}
		if err := msg.Decode(&pack); err != nil {
			return err
		}
		e, err := pack2event(&pack, station)
		if err != nil {
			return err
		}
		router.AddNetIn(station, 1)
		if checkDDOS(monitor, e) {
			//router.SendTo(nil, nil, router.OneMinuteLimited, e.From)
			log.Warn("DDos detection", "peer", remote.peer.String(), "typecode", e.Typecode, "count", monitor[e.Typecode][1])
			//return p2p.DiscDDOS
		}
		if e.To == nil && len(pack.To) != 0 {
			continue
		}
		router.SendEvent(e)
	}
}

func checkDDOS(m map[int][]int64, e *router.Event) bool {
	t := e.Typecode
	limit := int64(router.GetDDosLimit(t)) * 10
	if limit == 0 {
		return false
	}

	if len(m[t]) == 0 {
		m[t] = make([]int64, 2)
	}
	//m[t][0] time
	//m[t][1] request per second
	if m[t][0] != time.Now().Unix() {
		m[t][0] = time.Now().Unix()
		m[t][1] = 0
	}
	m[t][1]++
	return m[t][1] > limit
}

// Protocols .
func (adaptor *ProtoAdaptor) Protocols() []p2p.Protocol {
	return []p2p.Protocol{
		p2p.Protocol{
			Name:    "FractalTest",
			Version: 1,
			Length:  1,
			Run:     adaptor.adaptorLoop,
		},
	}
}

// Stop .
func (adaptor *ProtoAdaptor) Stop() {
	close(adaptor.quit)
	for _, sub := range adaptor.subs {
		sub.Unsubscribe()
	}
	adaptor.Server.Stop()
	adaptor.loopWG.Wait()
	log.Info("P2P networking stopped")
}

// SendOut .
func (adaptor *ProtoAdaptor) SendOut(e *router.Event) error {
	if e.To.IsBroadcast() {
		adaptor.msgBroadcast(e)
		return nil
	}
	return adaptor.msgSend(e)
}

func (adaptor *ProtoAdaptor) msgSend(e *router.Event) error {
	pack, err := event2pack(e)
	if err != nil {
		return err
	}
	router.AddNetOut(e.To, 1)
	return p2p.Send(e.To.Data().(*remotePeer).ws, 0, pack)
}

func (adaptor *ProtoAdaptor) msgBroadcast(e *router.Event) {
	te := *e
	te.To = nil
	te.From = nil
	pack, err := event2pack(&te)
	if err != nil {
		return
	}

	send := func(peer *remotePeer) {
		//router.AddNetOut(x,1)
		p2p.Send(peer.ws, 0, pack)
	}
	if e.To.Data() != nil {
		e.To.Data().(*peerMangaer).mapActivePeer(send)
		return
	}
	adaptor.peerMangaer.mapActivePeer(send)
}

func event2pack(e *router.Event) (*pack, error) {
	buf, err := rlp.EncodeToBytes(e.Data)
	if err != nil {
		return nil, err
	}
	from := ""
	if e.From != nil {
		from = e.From.Name()
	}
	to := ""
	if e.To != nil {
		to = e.To.Name()[8:]
	}
	return &pack{
		From:     from,
		To:       to,
		Typecode: uint32(e.Typecode),
		Payload:  buf,
	}, nil
}

func pack2event(pack *pack, station router.Station) (*router.Event, error) {
	var elem interface{}

	isPtr := false
	typ := router.GetTypeByCode(int(pack.Typecode))
	if typ == nil {
		return nil, fmt.Errorf("unknow typecode: %d", pack.Typecode)
	}

	//for typ.Kind() == reflect.Ptr {
	if typ.Kind() == reflect.Ptr {
		isPtr = true
		typ = typ.Elem()
	}
	obj := reflect.New(typ)
	if err := rlp.DecodeBytes(pack.Payload, obj.Interface()); err != nil {
		return nil, err
	}
	if isPtr {
		elem = obj.Interface()
	} else {
		elem = obj.Elem().Interface()
	}

	if pack.From != "" {
		station = router.NewRemoteStation(station.Name()+pack.From, station.Data())
	}
	return &router.Event{
		From:     station,
		To:       router.GetStationByName(pack.To),
		Typecode: int(pack.Typecode),
		Data:     elem,
	}, nil
}
<|endoftext|>


<commit_before>// Copyright (C) 2015 Thomas de Zeeuw.
//
// Licensed under the MIT license that can be found in the LICENSE file.

package logger

import (
	"errors"
	"fmt"
	"math"
	"testing"
	"time"
)

type stringer struct{}

func (s *stringer) String() string {
	return "data"
}

func TestEvent(t *testing.T) {
	defer resetEventTypes()

	now := time.Now()
	tStr := now.UTC().Format(TimeFormat)
	tStrNano := now.UTC().Format(time.RFC3339Nano)

	var msgTests = []struct {
		event        Event
		expected     string
		expectedJSON string
	}{
		{Event{DebugEvent, now, Tags{"tag1", "tag2", "tag3"}, "Message6", 0},
			tStr + " [Debug] tag1, tag2, tag3: Message6, 0",
			`{"type": "Debug", "timestamp": "` + tStrNano + `", "tags": ["tag1", "tag2", "tag3"], ` +
				`"message": "Message6", "data": "0"}`},
		{Event{InfoEvent, now, Tags{"tag1", "tag2"}, "Message4", []byte("data")},
			tStr + " [Info] tag1, tag2: Message4, data",
			`{"type": "Info", "timestamp": "` + tStrNano + `", "tags": ["tag1", "tag2"], ` +
				`"message": "Message4", "data": "data"}`},
		{Event{WarnEvent, now, Tags{"tag1"}, "Message3", &stringer{}},
			tStr + " [Warn] tag1: Message3, data",
			`{"type": "Warn", "timestamp": "` + tStrNano + `", "tags": ["tag1"], ` +
				`"message": "Message3", "data": "data"}`},
		{Event{ErrorEvent, now, Tags{"tag1"}, "Message2", "data"},
			tStr + " [Error] tag1: Message2, data",
			`{"type": "Error", "timestamp": "` + tStrNano + `", "tags": ["tag1"], ` +
				`"message": "Message2", "data": "data"}`},
		{Event{FatalEvent, now, Tags{}, "Message1", nil},
			tStr + " [Fatal] : Message1",
			`{"type": "Fatal", "timestamp": "` + tStrNano + `", "tags": [], ` +
				`"message": "Message1"}`},
		{Event{ThumbEvent, now, Tags{"tag1", "tag2", "tag3"}, "Message5", errors.New("error data")},
			tStr + " [Thumb] tag1, tag2, tag3: Message5, error data",
			`{"type": "Thumb", "timestamp": "` + tStrNano + `", "tags": ["tag1", "tag2", "tag3"], ` +
				`"message": "Message5", "data": "error data"}`},
		{Event{NewEventType("My-event-type"), now, Tags{"tag1"}, "Message7", nil},
			tStr + " [My-event-type] tag1: Message7",
			`{"type": "My-event-type", "timestamp": "` + tStrNano + `", "tags": ["tag1"], ` +
				`"message": "Message7"}`},
		{Event{NewEventType(`my-"event"-type`), now, Tags{`tag"1"`}, "Message7", `"`},
			tStr + " [my-\"event\"-type] tag\"1\": Message7, \"",
			`{"type": "my-\"event\"-type", "timestamp": "` + tStrNano + `", "tags": ["tag\"1\""], ` +
				`"message": "Message7", "data": "\""}`},
	}

	for _, test := range msgTests {
		got, gotBytes := test.event.String(), string(test.event.Bytes())
		if gotBytes != test.expected {
			t.Errorf("Expected Event(%v).Bytes() to return %q, but got %q",
				test.event, test.expected, gotBytes)
		} else if got != test.expected {
			t.Errorf("Expected Event(%v).String() to return %q, but got %q",
				test.event, test.expected, got)
		}

		if json, err := test.event.MarshalJSON(); err != nil {
			t.Errorf("Unexpected error marshaling %v into json: %s", test.event, err.Error())
		} else if got := string(json); got != test.expectedJSON {
			t.Errorf("Expected Event(%v).MarshalJSON() to return %q, but got %q",
				test.event, test.expectedJSON, got)
		}
	}
}

func TestEventType(t *testing.T) {
	defer resetEventTypes()

	tests := []struct {
		eventType    EventType
		expected     string
		expectedJSON string
	}{
		{DebugEvent, "Debug", `"Debug"`},
		{ThumbEvent, "Thumb", `"Thumb"`},
		{InfoEvent, "Info", `"Info"`},
		{WarnEvent, "Warn", `"Warn"`},
		{ErrorEvent, "Error", `"Error"`},
		{FatalEvent, "Fatal", `"Fatal"`},
		{EventType(255), "EventType(255)", `"EventType(255)"`},
		{NewEventType("my-event-type"), "my-event-type", `"my-event-type"`},
		{NewEventType("my-\"event\"-type"), "my-\"event\"-type", `"my-\"event\"-type"`},
	}

	for _, test := range tests {
		got, gotBytes := test.eventType.String(), string(test.eventType.Bytes())
		if gotBytes != test.expected {
			t.Errorf("Expected EventType(%v).Bytes() to return %q, but got %q",
				test.eventType, test.expected, gotBytes)
		} else if got != test.expected {
			t.Errorf("Expected EventType(%v).String() to return %q, but got %q",
				test.eventType, test.expected, got)
		}

		if json, err := test.eventType.MarshalJSON(); err != nil {
			t.Errorf("Unexpected error marshaling %v into json: %s", test.eventType, err.Error())
		} else if got := string(json); got != test.expectedJSON {
			t.Errorf("Expected EventType(%v).MarshalJSON() to return %q, but got %q",
				test.eventType, test.expectedJSON, got)
		}
	}
}

func TestNewEventTypeWithEmptyName(t *testing.T) {
	defer func() {
		recv := recover()
		if recv == nil {
			t.Fatal(`Expceted NewEventType("") to panic, but it didn't`)
		}

		expected := "logger: EventType name can't be empty"
		got := recv.(string)
		if got != expected {
			t.Fatalf("Expceted panic value to be %s, but got %s", expected, got)
		}
	}()

	NewEventType("")
}

var (
	// Minus builtin event types.
	maxCostumEventTypes = math.MaxUint16 - len(eventTypeIndices)

	oldEventTypeNames   = eventTypeNames
	oldEventTypeIndices = eventTypeIndices
)

func resetEventTypes() {
	eventTypeNames = oldEventTypeNames
	eventTypeIndices = oldEventTypeIndices
}

func TestNewEventTypeLimit(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping TestNewEventTypeLimit in short mode")
	}
	defer resetEventTypes()

	for i := 1; i <= maxCostumEventTypes; i++ {
		expected := fmt.Sprintf("EventType-%d", i)
		eventType := NewEventType(expected)
		if got := eventType.String(); got != expected {
			t.Fatalf("Expected NewEventType(%q).String() to return %q, but got %q",
				expected, expected, got)
		}
	}

	defer func() {
		recv := recover()
		if recv == nil {
			t.Fatal("Expected a panic after creating 65528 log levels, but didn't get one")
		}

		got, ok := recv.(string)
		if !ok {
			t.Fatalf("Expected the recoverd panic to be a string, but it's %v", recv)
		}

		expected := "logger: can't have more then 65535 EventTypes"
		if got != expected {
			t.Fatalf("Expected the recoverd panic to be %s, but got %s", expected, got)
		}
	}()

	NewEventType(fmt.Sprintf("EventType-%d", maxCostumEventTypes+1))
}
<commit_msg>Fix type-o in TestNewEventTypeWithEmptyName<commit_after>// Copyright (C) 2015 Thomas de Zeeuw.
//
// Licensed under the MIT license that can be found in the LICENSE file.

package logger

import (
	"errors"
	"fmt"
	"math"
	"testing"
	"time"
)

type stringer struct{}

func (s *stringer) String() string {
	return "data"
}

func TestEvent(t *testing.T) {
	defer resetEventTypes()

	now := time.Now()
	tStr := now.UTC().Format(TimeFormat)
	tStrNano := now.UTC().Format(time.RFC3339Nano)

	var msgTests = []struct {
		event        Event
		expected     string
		expectedJSON string
	}{
		{Event{DebugEvent, now, Tags{"tag1", "tag2", "tag3"}, "Message6", 0},
			tStr + " [Debug] tag1, tag2, tag3: Message6, 0",
			`{"type": "Debug", "timestamp": "` + tStrNano + `", "tags": ["tag1", "tag2", "tag3"], ` +
				`"message": "Message6", "data": "0"}`},
		{Event{InfoEvent, now, Tags{"tag1", "tag2"}, "Message4", []byte("data")},
			tStr + " [Info] tag1, tag2: Message4, data",
			`{"type": "Info", "timestamp": "` + tStrNano + `", "tags": ["tag1", "tag2"], ` +
				`"message": "Message4", "data": "data"}`},
		{Event{WarnEvent, now, Tags{"tag1"}, "Message3", &stringer{}},
			tStr + " [Warn] tag1: Message3, data",
			`{"type": "Warn", "timestamp": "` + tStrNano + `", "tags": ["tag1"], ` +
				`"message": "Message3", "data": "data"}`},
		{Event{ErrorEvent, now, Tags{"tag1"}, "Message2", "data"},
			tStr + " [Error] tag1: Message2, data",
			`{"type": "Error", "timestamp": "` + tStrNano + `", "tags": ["tag1"], ` +
				`"message": "Message2", "data": "data"}`},
		{Event{FatalEvent, now, Tags{}, "Message1", nil},
			tStr + " [Fatal] : Message1",
			`{"type": "Fatal", "timestamp": "` + tStrNano + `", "tags": [], ` +
				`"message": "Message1"}`},
		{Event{ThumbEvent, now, Tags{"tag1", "tag2", "tag3"}, "Message5", errors.New("error data")},
			tStr + " [Thumb] tag1, tag2, tag3: Message5, error data",
			`{"type": "Thumb", "timestamp": "` + tStrNano + `", "tags": ["tag1", "tag2", "tag3"], ` +
				`"message": "Message5", "data": "error data"}`},
		{Event{NewEventType("My-event-type"), now, Tags{"tag1"}, "Message7", nil},
			tStr + " [My-event-type] tag1: Message7",
			`{"type": "My-event-type", "timestamp": "` + tStrNano + `", "tags": ["tag1"], ` +
				`"message": "Message7"}`},
		{Event{NewEventType(`my-"event"-type`), now, Tags{`tag"1"`}, "Message7", `"`},
			tStr + " [my-\"event\"-type] tag\"1\": Message7, \"",
			`{"type": "my-\"event\"-type", "timestamp": "` + tStrNano + `", "tags": ["tag\"1\""], ` +
				`"message": "Message7", "data": "\""}`},
	}

	for _, test := range msgTests {
		got, gotBytes := test.event.String(), string(test.event.Bytes())
		if gotBytes != test.expected {
			t.Errorf("Expected Event(%v).Bytes() to return %q, but got %q",
				test.event, test.expected, gotBytes)
		} else if got != test.expected {
			t.Errorf("Expected Event(%v).String() to return %q, but got %q",
				test.event, test.expected, got)
		}

		if json, err := test.event.MarshalJSON(); err != nil {
			t.Errorf("Unexpected error marshaling %v into json: %s", test.event, err.Error())
		} else if got := string(json); got != test.expectedJSON {
			t.Errorf("Expected Event(%v).MarshalJSON() to return %q, but got %q",
				test.event, test.expectedJSON, got)
		}
	}
}

func TestEventType(t *testing.T) {
	defer resetEventTypes()

	tests := []struct {
		eventType    EventType
		expected     string
		expectedJSON string
	}{
		{DebugEvent, "Debug", `"Debug"`},
		{ThumbEvent, "Thumb", `"Thumb"`},
		{InfoEvent, "Info", `"Info"`},
		{WarnEvent, "Warn", `"Warn"`},
		{ErrorEvent, "Error", `"Error"`},
		{FatalEvent, "Fatal", `"Fatal"`},
		{EventType(255), "EventType(255)", `"EventType(255)"`},
		{NewEventType("my-event-type"), "my-event-type", `"my-event-type"`},
		{NewEventType("my-\"event\"-type"), "my-\"event\"-type", `"my-\"event\"-type"`},
	}

	for _, test := range tests {
		got, gotBytes := test.eventType.String(), string(test.eventType.Bytes())
		if gotBytes != test.expected {
			t.Errorf("Expected EventType(%v).Bytes() to return %q, but got %q",
				test.eventType, test.expected, gotBytes)
		} else if got != test.expected {
			t.Errorf("Expected EventType(%v).String() to return %q, but got %q",
				test.eventType, test.expected, got)
		}

		if json, err := test.eventType.MarshalJSON(); err != nil {
			t.Errorf("Unexpected error marshaling %v into json: %s", test.eventType, err.Error())
		} else if got := string(json); got != test.expectedJSON {
			t.Errorf("Expected EventType(%v).MarshalJSON() to return %q, but got %q",
				test.eventType, test.expectedJSON, got)
		}
	}
}

func TestNewEventTypeWithEmptyName(t *testing.T) {
	defer func() {
		recv := recover()
		if recv == nil {
			t.Fatal(`Expected NewEventType("") to panic, but it didn't`)
		}

		expected := "logger: EventType name can't be empty"
		got := recv.(string)
		if got != expected {
			t.Fatalf("Expected panic value to be %s, but got %s", expected, got)
		}
	}()

	NewEventType("")
}

var (
	// Minus builtin event types.
	maxCostumEventTypes = math.MaxUint16 - len(eventTypeIndices)

	oldEventTypeNames   = eventTypeNames
	oldEventTypeIndices = eventTypeIndices
)

func resetEventTypes() {
	eventTypeNames = oldEventTypeNames
	eventTypeIndices = oldEventTypeIndices
}

func TestNewEventTypeLimit(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping TestNewEventTypeLimit in short mode")
	}
	defer resetEventTypes()

	for i := 1; i <= maxCostumEventTypes; i++ {
		expected := fmt.Sprintf("EventType-%d", i)
		eventType := NewEventType(expected)
		if got := eventType.String(); got != expected {
			t.Fatalf("Expected NewEventType(%q).String() to return %q, but got %q",
				expected, expected, got)
		}
	}

	defer func() {
		recv := recover()
		if recv == nil {
			t.Fatal("Expected a panic after creating 65528 log levels, but didn't get one")
		}

		got, ok := recv.(string)
		if !ok {
			t.Fatalf("Expected the recoverd panic to be a string, but it's %v", recv)
		}

		expected := "logger: can't have more then 65535 EventTypes"
		if got != expected {
			t.Fatalf("Expected the recoverd panic to be %s, but got %s", expected, got)
		}
	}()

	NewEventType(fmt.Sprintf("EventType-%d", maxCostumEventTypes+1))
}
<|endoftext|>


<commit_before>package hyperv

import (
	"archive/tar"
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"time"

	"github.com/docker/machine/libmachine/drivers"
	"github.com/docker/machine/libmachine/log"
	"github.com/docker/machine/libmachine/mcnflag"
	"github.com/docker/machine/libmachine/mcnutils"
	"github.com/docker/machine/libmachine/ssh"
	"github.com/docker/machine/libmachine/state"
)

type Driver struct {
	*drivers.BaseDriver
	Boot2DockerURL string
	vSwitch        string
	diskImage      string
	DiskSize       int
	MemSize        int
}

const (
	defaultDiskSize = 20000
	defaultMemory   = 1024
)

func NewDriver(hostName, storePath string) drivers.Driver {
	return &Driver{
		DiskSize: defaultDiskSize,
		MemSize:  defaultMemory,
		BaseDriver: &drivers.BaseDriver{
			MachineName: hostName,
			StorePath:   storePath,
		},
	}
}

// GetCreateFlags registers the flags this driver adds to
// "docker hosts create"
func (d *Driver) GetCreateFlags() []mcnflag.Flag {
	return []mcnflag.Flag{
		mcnflag.StringFlag{
			Name:  "hyperv-boot2docker-url",
			Usage: "URL of the boot2docker ISO. Defaults to the latest available version.",
		},
		mcnflag.StringFlag{
			Name:  "hyperv-virtual-switch",
			Usage: "Virtual switch name. Defaults to first found.",
		},
		mcnflag.IntFlag{
			Name:  "hyperv-disk-size",
			Usage: "Maximum size of dynamically expanding disk in MB.",
			Value: defaultDiskSize,
		},
		mcnflag.IntFlag{
			Name:  "hyperv-memory",
			Usage: "Memory size for host in MB.",
			Value: defaultMemory,
		},
	}
}

func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) error {
	d.Boot2DockerURL = flags.String("hyperv-boot2docker-url")
	d.vSwitch = flags.String("hyperv-virtual-switch")
	d.DiskSize = flags.Int("hyperv-disk-size")
	d.MemSize = flags.Int("hyperv-memory")
	d.SwarmMaster = flags.Bool("swarm-master")
	d.SwarmHost = flags.String("swarm-host")
	d.SwarmDiscovery = flags.String("swarm-discovery")
	d.SSHUser = "docker"
	return nil
}

func (d *Driver) GetSSHHostname() (string, error) {
	return d.GetIP()
}

// DriverName returns the name of the driver
func (d *Driver) DriverName() string {
	return "hyperv"
}

func (d *Driver) GetURL() (string, error) {
	ip, err := d.GetIP()
	if err != nil {
		return "", err
	}

	if ip == "" {
		return "", nil
	}

	return fmt.Sprintf("tcp://%s", net.JoinHostPort(ip, "2376")), nil
}

func (d *Driver) GetState() (state.State, error) {
	stdout, err := cmdOut("(", "Get-VM", "-Name", d.MachineName, ").state")
	if err != nil {
		return state.None, fmt.Errorf("Failed to find the VM status")
	}

	resp := parseLines(stdout)
	if len(resp) < 1 {
		return state.None, nil
	}

	switch resp[0] {
	case "Running":
		return state.Running, nil
	case "Off":
		return state.Stopped, nil
	default:
		return state.None, nil
	}
}

// PreCreateCheck checks that the machine creation process can be started safely.
func (d *Driver) PreCreateCheck() error {
	if err := hypervAvailable(); err != nil {
		return err
	}

	// Downloading boot2docker to cache should be done here to make sure
	// that a download failure will not leave a machine half created.
	b2dutils := mcnutils.NewB2dUtils(d.StorePath)
	if err := b2dutils.UpdateISOCache(d.Boot2DockerURL); err != nil {
		return err
	}

	return nil
}

func (d *Driver) Create() error {
	b2dutils := mcnutils.NewB2dUtils(d.StorePath)
	if err := b2dutils.CopyIsoToMachineDir(d.Boot2DockerURL, d.MachineName); err != nil {
		return err
	}

	log.Infof("Creating SSH key...")
	if err := ssh.GenerateSSHKey(d.GetSSHKeyPath()); err != nil {
		return err
	}

	log.Infof("Creating VM...")
	virtualSwitch, err := d.chooseVirtualSwitch()
	if err != nil {
		return err
	}

	err = d.generateDiskImage()
	if err != nil {
		return err
	}

	if err := cmd("New-VM", "-Name", d.MachineName, "-Path", fmt.Sprintf("'%s'", d.ResolveStorePath(".")), "-MemoryStartupBytes", fmt.Sprintf("%dMB", d.MemSize)); err != nil {
		return err
	}

	if err := cmd("Set-VMDvdDrive", "-VMName", d.MachineName, "-Path", fmt.Sprintf("'%s'", d.ResolveStorePath("boot2docker.iso"))); err != nil {
		return err
	}

	if err := cmd("Add-VMHardDiskDrive", "-VMName", d.MachineName, "-Path", fmt.Sprintf("'%s'", d.diskImage)); err != nil {
		return err
	}

	if err := cmd("Connect-VMNetworkAdapter", "-VMName", d.MachineName, "-SwitchName", fmt.Sprintf("'%s'", virtualSwitch)); err != nil {
		return err
	}

	log.Infof("Starting VM...")
	return d.Start()
}

func (d *Driver) chooseVirtualSwitch() (string, error) {
	if d.vSwitch != "" {
		return d.vSwitch, nil
	}

	stdout, err := cmdOut("@(Get-VMSwitch).Name")
	if err != nil {
		return "", err
	}

	switches := parseLines(stdout)
	if len(switches) < 1 {
		return "", fmt.Errorf("no vswitch found")
	}

	log.Infof("Using switch %s", switches[0])

	return switches[0], nil
}

func (d *Driver) wait() error {
	log.Infof("Waiting for host to start...")

	for {
		ip, _ := d.GetIP()
		if ip != "" {
			break
		}

		time.Sleep(1 * time.Second)
	}

	return nil
}

func (d *Driver) Start() error {
	if err := cmd("Start-VM", "-Name", d.MachineName); err != nil {
		return err
	}

	if err := d.wait(); err != nil {
		return err
	}

	var err error
	d.IPAddress, err = d.GetIP()

	return err
}

func (d *Driver) Stop() error {
	if err := cmd("Stop-VM", "-Name", d.MachineName); err != nil {
		return err
	}

	for {
		s, err := d.GetState()
		if err != nil {
			return err
		}

		if s != state.Running {
			break
		}

		time.Sleep(1 * time.Second)
	}

	d.IPAddress = ""

	return nil
}

func (d *Driver) Remove() error {
	s, err := d.GetState()
	if err != nil {
		return err
	}

	if s == state.Running {
		if err := d.Kill(); err != nil {
			return err
		}
	}

	return cmd("Remove-VM", "-Name", d.MachineName, "-Force")
}

func (d *Driver) Restart() error {
	err := d.Stop()
	if err != nil {
		return err
	}

	return d.Start()
}

func (d *Driver) Kill() error {
	if err := cmd("Stop-VM", "-Name", d.MachineName, "-TurnOff"); err != nil {
		return err
	}

	for {
		s, err := d.GetState()
		if err != nil {
			return err
		}

		if s != state.Running {
			break
		}

		time.Sleep(1 * time.Second)
	}

	d.IPAddress = ""

	return nil
}

func (d *Driver) GetIP() (string, error) {
	stdout, err := cmdOut("((", "Get-VM", "-Name", d.MachineName, ").networkadapters[0]).ipaddresses[0]")
	if err != nil {
		return "", err
	}

	resp := parseLines(stdout)
	if len(resp) < 1 {
		return "", fmt.Errorf("IP not found")
	}

	return resp[0], nil
}

func (d *Driver) publicSSHKeyPath() string {
	return d.GetSSHKeyPath() + ".pub"
}

// generateDiskImage creates a small fixed vhd, put the tar in, convert to dynamic, then resize
func (d *Driver) generateDiskImage() error {
	d.diskImage = d.ResolveStorePath("disk.vhd")
	fixed := d.ResolveStorePath("fixed.vhd")

	log.Infof("Creating VHD")
	if err := cmd("New-VHD", "-Path", fmt.Sprintf("'%s'", fixed), "-SizeBytes", "10MB", "-Fixed"); err != nil {
		return err
	}

	tarBuf, err := d.generateTar()
	if err != nil {
		return err
	}

	file, err := os.OpenFile(fixed, os.O_WRONLY, 0644)
	if err != nil {
		return err
	}
	defer file.Close()

	file.Seek(0, os.SEEK_SET)
	_, err = file.Write(tarBuf.Bytes())
	if err != nil {
		return err
	}
	file.Close()

	if err := cmd("Convert-VHD", "-Path", fmt.Sprintf("'%s'", fixed), "-DestinationPath", fmt.Sprintf("'%s'", d.diskImage), "-VHDType", "Dynamic"); err != nil {
		return err
	}

	return cmd("Resize-VHD", "-Path", fmt.Sprintf("'%s'", d.diskImage), "-SizeBytes", fmt.Sprintf("%dMB", d.DiskSize))
}

// Make a boot2docker VM disk image.
// See https://github.com/boot2docker/boot2docker/blob/master/rootfs/rootfs/etc/rc.d/automount
func (d *Driver) generateTar() (*bytes.Buffer, error) {
	magicString := "boot2docker, please format-me"

	buf := new(bytes.Buffer)
	tw := tar.NewWriter(buf)

	// magicString first so the automount script knows to format the disk
	file := &tar.Header{Name: magicString, Size: int64(len(magicString))}
	if err := tw.WriteHeader(file); err != nil {
		return nil, err
	}

	if _, err := tw.Write([]byte(magicString)); err != nil {
		return nil, err
	}

	// .ssh/key.pub => authorized_keys
	file = &tar.Header{Name: ".ssh", Typeflag: tar.TypeDir, Mode: 0700}
	if err := tw.WriteHeader(file); err != nil {
		return nil, err
	}

	pubKey, err := ioutil.ReadFile(d.publicSSHKeyPath())
	if err != nil {
		return nil, err
	}

	file = &tar.Header{Name: ".ssh/authorized_keys", Size: int64(len(pubKey)), Mode: 0644}
	if err := tw.WriteHeader(file); err != nil {
		return nil, err
	}

	if _, err := tw.Write([]byte(pubKey)); err != nil {
		return nil, err
	}

	file = &tar.Header{Name: ".ssh/authorized_keys2", Size: int64(len(pubKey)), Mode: 0644}
	if err := tw.WriteHeader(file); err != nil {
		return nil, err
	}

	if _, err := tw.Write([]byte(pubKey)); err != nil {
		return nil, err
	}

	if err := tw.Close(); err != nil {
		return nil, err
	}

	return buf, nil
}
<commit_msg>FIX #512 Check the virtual switch in PreCreateCheck()<commit_after>package hyperv

import (
	"archive/tar"
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"time"

	"github.com/docker/machine/libmachine/drivers"
	"github.com/docker/machine/libmachine/log"
	"github.com/docker/machine/libmachine/mcnflag"
	"github.com/docker/machine/libmachine/mcnutils"
	"github.com/docker/machine/libmachine/ssh"
	"github.com/docker/machine/libmachine/state"
)

type Driver struct {
	*drivers.BaseDriver
	Boot2DockerURL string
	vSwitch        string
	diskImage      string
	DiskSize       int
	MemSize        int
}

const (
	defaultDiskSize = 20000
	defaultMemory   = 1024
)

func NewDriver(hostName, storePath string) drivers.Driver {
	return &Driver{
		DiskSize: defaultDiskSize,
		MemSize:  defaultMemory,
		BaseDriver: &drivers.BaseDriver{
			MachineName: hostName,
			StorePath:   storePath,
		},
	}
}

// GetCreateFlags registers the flags this driver adds to
// "docker hosts create"
func (d *Driver) GetCreateFlags() []mcnflag.Flag {
	return []mcnflag.Flag{
		mcnflag.StringFlag{
			Name:  "hyperv-boot2docker-url",
			Usage: "URL of the boot2docker ISO. Defaults to the latest available version.",
		},
		mcnflag.StringFlag{
			Name:  "hyperv-virtual-switch",
			Usage: "Virtual switch name. Defaults to first found.",
		},
		mcnflag.IntFlag{
			Name:  "hyperv-disk-size",
			Usage: "Maximum size of dynamically expanding disk in MB.",
			Value: defaultDiskSize,
		},
		mcnflag.IntFlag{
			Name:  "hyperv-memory",
			Usage: "Memory size for host in MB.",
			Value: defaultMemory,
		},
	}
}

func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) error {
	d.Boot2DockerURL = flags.String("hyperv-boot2docker-url")
	d.vSwitch = flags.String("hyperv-virtual-switch")
	d.DiskSize = flags.Int("hyperv-disk-size")
	d.MemSize = flags.Int("hyperv-memory")
	d.SwarmMaster = flags.Bool("swarm-master")
	d.SwarmHost = flags.String("swarm-host")
	d.SwarmDiscovery = flags.String("swarm-discovery")
	d.SSHUser = "docker"
	return nil
}

func (d *Driver) GetSSHHostname() (string, error) {
	return d.GetIP()
}

// DriverName returns the name of the driver
func (d *Driver) DriverName() string {
	return "hyperv"
}

func (d *Driver) GetURL() (string, error) {
	ip, err := d.GetIP()
	if err != nil {
		return "", err
	}

	if ip == "" {
		return "", nil
	}

	return fmt.Sprintf("tcp://%s", net.JoinHostPort(ip, "2376")), nil
}

func (d *Driver) GetState() (state.State, error) {
	stdout, err := cmdOut("(", "Get-VM", "-Name", d.MachineName, ").state")
	if err != nil {
		return state.None, fmt.Errorf("Failed to find the VM status")
	}

	resp := parseLines(stdout)
	if len(resp) < 1 {
		return state.None, nil
	}

	switch resp[0] {
	case "Running":
		return state.Running, nil
	case "Off":
		return state.Stopped, nil
	default:
		return state.None, nil
	}
}

// PreCreateCheck checks that the machine creation process can be started safely.
func (d *Driver) PreCreateCheck() error {
	if err := hypervAvailable(); err != nil {
		return err
	}

	// Check that there is a virtual switch already configured
	_, err := d.chooseVirtualSwitch()
	if err != nil {
		return err
	}

	// Downloading boot2docker to cache should be done here to make sure
	// that a download failure will not leave a machine half created.
	b2dutils := mcnutils.NewB2dUtils(d.StorePath)
	if err := b2dutils.UpdateISOCache(d.Boot2DockerURL); err != nil {
		return err
	}

	return nil
}

func (d *Driver) Create() error {
	b2dutils := mcnutils.NewB2dUtils(d.StorePath)
	if err := b2dutils.CopyIsoToMachineDir(d.Boot2DockerURL, d.MachineName); err != nil {
		return err
	}

	log.Infof("Creating SSH key...")
	if err := ssh.GenerateSSHKey(d.GetSSHKeyPath()); err != nil {
		return err
	}

	log.Infof("Creating VM...")
	virtualSwitch, err := d.chooseVirtualSwitch()
	if err != nil {
		return err
	}

	log.Infof("Using switch %q", virtualSwitch)

	err = d.generateDiskImage()
	if err != nil {
		return err
	}

	if err := cmd("New-VM", "-Name", d.MachineName, "-Path", fmt.Sprintf("'%s'", d.ResolveStorePath(".")), "-MemoryStartupBytes", fmt.Sprintf("%dMB", d.MemSize)); err != nil {
		return err
	}

	if err := cmd("Set-VMDvdDrive", "-VMName", d.MachineName, "-Path", fmt.Sprintf("'%s'", d.ResolveStorePath("boot2docker.iso"))); err != nil {
		return err
	}

	if err := cmd("Add-VMHardDiskDrive", "-VMName", d.MachineName, "-Path", fmt.Sprintf("'%s'", d.diskImage)); err != nil {
		return err
	}

	if err := cmd("Connect-VMNetworkAdapter", "-VMName", d.MachineName, "-SwitchName", fmt.Sprintf("'%s'", virtualSwitch)); err != nil {
		return err
	}

	log.Infof("Starting VM...")
	return d.Start()
}

func (d *Driver) chooseVirtualSwitch() (string, error) {
	stdout, err := cmdOut("@(Get-VMSwitch).Name")
	if err != nil {
		return "", err
	}

	switches := parseLines(stdout)

	if d.vSwitch == "" {
		if len(switches) < 1 {
			return "", fmt.Errorf("no vswitch found")
		}

		return switches[0], nil
	}

	found := false
	for _, name := range switches {
		if name == d.vSwitch {
			found = true
			break
		}
	}

	if !found {
		return "", fmt.Errorf("vswitch %q not found", d.vSwitch)
	}

	return d.vSwitch, nil
}

func (d *Driver) wait() error {
	log.Infof("Waiting for host to start...")

	for {
		ip, _ := d.GetIP()
		if ip != "" {
			break
		}

		time.Sleep(1 * time.Second)
	}

	return nil
}

func (d *Driver) Start() error {
	if err := cmd("Start-VM", "-Name", d.MachineName); err != nil {
		return err
	}

	if err := d.wait(); err != nil {
		return err
	}

	var err error
	d.IPAddress, err = d.GetIP()

	return err
}

func (d *Driver) Stop() error {
	if err := cmd("Stop-VM", "-Name", d.MachineName); err != nil {
		return err
	}

	for {
		s, err := d.GetState()
		if err != nil {
			return err
		}

		if s != state.Running {
			break
		}

		time.Sleep(1 * time.Second)
	}

	d.IPAddress = ""

	return nil
}

func (d *Driver) Remove() error {
	s, err := d.GetState()
	if err != nil {
		return err
	}

	if s == state.Running {
		if err := d.Kill(); err != nil {
			return err
		}
	}

	return cmd("Remove-VM", "-Name", d.MachineName, "-Force")
}

func (d *Driver) Restart() error {
	err := d.Stop()
	if err != nil {
		return err
	}

	return d.Start()
}

func (d *Driver) Kill() error {
	if err := cmd("Stop-VM", "-Name", d.MachineName, "-TurnOff"); err != nil {
		return err
	}

	for {
		s, err := d.GetState()
		if err != nil {
			return err
		}

		if s != state.Running {
			break
		}

		time.Sleep(1 * time.Second)
	}

	d.IPAddress = ""

	return nil
}

func (d *Driver) GetIP() (string, error) {
	stdout, err := cmdOut("((", "Get-VM", "-Name", d.MachineName, ").networkadapters[0]).ipaddresses[0]")
	if err != nil {
		return "", err
	}

	resp := parseLines(stdout)
	if len(resp) < 1 {
		return "", fmt.Errorf("IP not found")
	}

	return resp[0], nil
}

func (d *Driver) publicSSHKeyPath() string {
	return d.GetSSHKeyPath() + ".pub"
}

// generateDiskImage creates a small fixed vhd, put the tar in, convert to dynamic, then resize
func (d *Driver) generateDiskImage() error {
	d.diskImage = d.ResolveStorePath("disk.vhd")
	fixed := d.ResolveStorePath("fixed.vhd")

	log.Infof("Creating VHD")
	if err := cmd("New-VHD", "-Path", fmt.Sprintf("'%s'", fixed), "-SizeBytes", "10MB", "-Fixed"); err != nil {
		return err
	}

	tarBuf, err := d.generateTar()
	if err != nil {
		return err
	}

	file, err := os.OpenFile(fixed, os.O_WRONLY, 0644)
	if err != nil {
		return err
	}
	defer file.Close()

	file.Seek(0, os.SEEK_SET)
	_, err = file.Write(tarBuf.Bytes())
	if err != nil {
		return err
	}
	file.Close()

	if err := cmd("Convert-VHD", "-Path", fmt.Sprintf("'%s'", fixed), "-DestinationPath", fmt.Sprintf("'%s'", d.diskImage), "-VHDType", "Dynamic"); err != nil {
		return err
	}

	return cmd("Resize-VHD", "-Path", fmt.Sprintf("'%s'", d.diskImage), "-SizeBytes", fmt.Sprintf("%dMB", d.DiskSize))
}

// Make a boot2docker VM disk image.
// See https://github.com/boot2docker/boot2docker/blob/master/rootfs/rootfs/etc/rc.d/automount
func (d *Driver) generateTar() (*bytes.Buffer, error) {
	magicString := "boot2docker, please format-me"

	buf := new(bytes.Buffer)
	tw := tar.NewWriter(buf)

	// magicString first so the automount script knows to format the disk
	file := &tar.Header{Name: magicString, Size: int64(len(magicString))}
	if err := tw.WriteHeader(file); err != nil {
		return nil, err
	}

	if _, err := tw.Write([]byte(magicString)); err != nil {
		return nil, err
	}

	// .ssh/key.pub => authorized_keys
	file = &tar.Header{Name: ".ssh", Typeflag: tar.TypeDir, Mode: 0700}
	if err := tw.WriteHeader(file); err != nil {
		return nil, err
	}

	pubKey, err := ioutil.ReadFile(d.publicSSHKeyPath())
	if err != nil {
		return nil, err
	}

	file = &tar.Header{Name: ".ssh/authorized_keys", Size: int64(len(pubKey)), Mode: 0644}
	if err := tw.WriteHeader(file); err != nil {
		return nil, err
	}

	if _, err := tw.Write([]byte(pubKey)); err != nil {
		return nil, err
	}

	file = &tar.Header{Name: ".ssh/authorized_keys2", Size: int64(len(pubKey)), Mode: 0644}
	if err := tw.WriteHeader(file); err != nil {
		return nil, err
	}

	if _, err := tw.Write([]byte(pubKey)); err != nil {
		return nil, err
	}

	if err := tw.Close(); err != nil {
		return nil, err
	}

	return buf, nil
}
<|endoftext|>


<commit_before>package web

import (
    "fmt"
    "net/http"
    "net/url"
    "strings"

    "github.com/gin-gonic/gin"
    "github.com/gin-gonic/contrib/renders/multitemplate"
    "github.com/gin-gonic/contrib/sessions"

    "github.com/earaujoassis/space/config"
    "github.com/earaujoassis/space/models"
    "github.com/earaujoassis/space/oauth"
    "github.com/earaujoassis/space/services"
    "github.com/earaujoassis/space/feature"
    "github.com/earaujoassis/space/utils"
)

const (
    errorURI string = "%s?error=%s&state=%s"
)

func createCustomRender() multitemplate.Render {
    render := multitemplate.New()
    render.AddFromFiles("satellite", "web/templates/default.html", "web/templates/satellite.html")
    return render
}

func ExposeRoutes(router *gin.Engine) {
    router.LoadHTMLGlob(config.GetConfig("template_folder").(string))
    router.HTMLRender = createCustomRender()
    router.Static("/public", config.GetConfig("public_folder").(string))
    store := sessions.NewCookieStore([]byte(config.GetConfig("http.session_secret").(string)))
    router.Use(sessions.Sessions("jupiter", store))
    views := router.Group("/")
    {
        views.GET("/", jupiterHandler)
        views.GET("/profile", jupiterHandler)

        views.GET("/signup", func(c *gin.Context) {
            c.HTML(http.StatusOK, "satellite", utils.H{
                "Title": " - Sign up",
                "Satellite": "io",
                "Data": utils.H{
                    "feature.gates": utils.H{
                        "user.create": feature.Active("user.create"),
                    },
                },
            })
        })

        views.GET("/signin", func(c *gin.Context) {
            c.HTML(http.StatusOK, "satellite", utils.H{
                "Title": " - Sign in",
                "Satellite": "ganymede",
            })
        })

        views.GET("/signout", func(c *gin.Context) {
            session := sessions.Default(c)

            userPublicId := session.Get("userPublicId")
            if userPublicId != nil {
                session.Delete("userPublicId")
                session.Save()
            }

            c.Redirect(http.StatusFound, "/signin")
        })

        views.GET("/session", func(c *gin.Context) {
            session := sessions.Default(c)

            userPublicId := session.Get("userPublicId")
            if userPublicId != nil {
                c.Redirect(http.StatusFound, "/")
                return
            }

            var nextPath string = "/"
            var scope string = c.Query("scope")
            var grantType string = c.Query("grant_type")
            var code string = c.Query("code")
            var clientId string = c.Query("client_id")
            var _nextPath string = c.Query("_")
            //var state string = c.Query("state")

            if scope == "" || grantType == "" || code == "" || clientId == "" {
                c.String(http.StatusMethodNotAllowed, "Missing required parameters")
                return
            }
            if _nextPath != "" {
                if _nextPath, err := url.QueryUnescape(_nextPath); err == nil {
                    nextPath = _nextPath
                }
            }

            client := services.FindOrCreateClient("Jupiter")
            if client.Key == clientId && grantType == oauth.AuthorizationCode && scope == models.PublicScope {
                grantToken := services.FindSessionByToken(code, models.GrantToken)
                if grantToken.ID != 0 {
                    session.Set("userPublicId", grantToken.User.PublicId)
                    session.Save()
                    services.InvalidateSession(grantToken)
                    c.Redirect(http.StatusFound, nextPath)
                    return
                }
            }

            c.Redirect(http.StatusFound, "/signin")
        })

        views.GET("/authorize", authorizeHandler)
        views.POST("/authorize", authorizeHandler)

        views.GET("/error", func(c *gin.Context) {
            c.String(http.StatusMethodNotAllowed, "Not implemented")
        })

        views.POST("/token", func(c *gin.Context) {
            var grantType string = c.PostForm("grant_type")

            authorizationBasic := strings.Replace(c.Request.Header["Authorization"][0], "Basic ", "", 1)
            client := oauth.ClientAuthentication(authorizationBasic)
            if client.ID == 0 {
                c.Header("WWW-Authenticate", fmt.Sprintf("Basic realm=\"%s\"", c.Request.RequestURI))
                c.JSON(http.StatusUnauthorized, utils.H{
                    "error": oauth.AccessDenied,
                })
                return
            }

            switch grantType {
            // Authorization Code Grant
            case oauth.AuthorizationCode:
                result, err := oauth.AccessTokenRequest(utils.H{
                    "grant_type": grantType,
                    "code": c.PostForm("code"),
                    "redirect_uri": c.PostForm("redirect_uri"),
                    "client": client,
                })
                if err != nil {
                    c.JSON(http.StatusMethodNotAllowed, utils.H{
                        "error": result["error"],
                    })
                    return
                } else {
                    c.JSON(http.StatusOK, utils.H{
                        "user_id": result["user_id"],
                        "access_token": result["access_token"],
                        "token_type": result["token_type"],
                        "expires_in": result["expires_in"],
                        "refresh_token": result["refresh_token"],
                        "scope": result["scope"],
                    })
                    return
                }
                return
            // Refreshing an Access Token
            case oauth.RefreshToken:
                result, err := oauth.RefreshTokenRequest(utils.H{
                    "grant_type": grantType,
                    "refresh_token": c.PostForm("refresh_token"),
                    "scope": c.PostForm("scope"),
                    "client": client,
                })
                if err != nil {
                    c.JSON(http.StatusMethodNotAllowed, utils.H{
                        "error": result["error"],
                    })
                    return
                } else {
                    c.JSON(http.StatusOK, utils.H{
                        "user_id": result["user_id"],
                        "access_token": result["access_token"],
                        "token_type": result["token_type"],
                        "expires_in": result["expires_in"],
                        "refresh_token": result["refresh_token"],
                        "scope": result["scope"],
                    })
                    return
                }
                return
            // Resource Owner Password Credentials Grant
            // Client Credentials Grant
            case oauth.Password, oauth.ClientCredentials:
                c.JSON(http.StatusMethodNotAllowed, utils.H{
                    "error": oauth.UnsupportedGrantType,
                })
                return
            default:
                c.JSON(http.StatusBadRequest, utils.H{
                    "error": oauth.InvalidRequest,
                })
                return
            }
        })
    }
}

func jupiterHandler(c *gin.Context) {
    session := sessions.Default(c)
    userPublicId := session.Get("userPublicId")
    if userPublicId == nil {
        c.Redirect(http.StatusFound, "/signin")
        return
    }
    client := services.FindOrCreateClient("Jupiter")
    user := services.FindUserByPublicId(userPublicId.(string))
    actionToken := services.CreateAction(user, client,
        c.Request.RemoteAddr,
        c.Request.UserAgent(),
        models.ReadWriteScope)
    c.HTML(http.StatusOK, "satellite", utils.H{
        "Title": " - Mission control",
        "Satellite": "europa",
        "Data": utils.H {
            "action_token": actionToken.Token,
            "user_id": user.UUID,
        },
    })
}

func authorizeHandler(c *gin.Context) {
    var location string
    var responseType string
    var clientId string
    var redirectURI string
    var scope string
    var state string

    session := sessions.Default(c)
    userPublicId := session.Get("userPublicId")
    nextPath := url.QueryEscape(fmt.Sprintf("%s?%s", c.Request.URL.Path, c.Request.URL.RawQuery))
    if userPublicId == nil {
        location = fmt.Sprintf("/signin?_=%s", nextPath)
        c.Redirect(http.StatusFound, location)
        return
    }
    user := services.FindUserByPublicId(userPublicId.(string))
    if user.ID == 0 {
        session.Delete("userPublicId")
        session.Save()
        location = fmt.Sprintf("/signin?_=%s", nextPath)
        c.Redirect(http.StatusFound, location)
        return
    }

    responseType = c.Query("response_type")
    clientId = c.Query("client_id")
    redirectURI = c.Query("redirect_uri")
    scope = c.Query("scope")
    state = c.Query("state")

    if redirectURI == "" {
        redirectURI = "/error"
    }

    client := services.FindClientByKey(clientId)
    if client.ID == 0 {
        redirectURI = "/error"
        location = fmt.Sprintf("%s?error=%s&state=%s",
            redirectURI, oauth.UnauthorizedClient, state)
        c.Redirect(http.StatusFound, location)
        return
    }

    if scope != models.PublicScope && scope != models.ReadScope && scope != models.ReadWriteScope {
        scope = "public"
    }

    switch responseType {
    // Authorization Code Grant
    case oauth.Code:
        activeSessions := services.ActiveSessionsForClient(client.ID, user.ID)
        if c.Request.Method == "GET" && activeSessions == 0 {
            c.HTML(http.StatusOK, "satellite", utils.H{
                "Title": " - Authorize",
                "Satellite": "callisto",
                "Data": utils.H{
                    "first_name": user.FirstName,
                    "last_name": user.LastName,
                    "client_name": client.Name,
                    "client_uri": client.DefaultRedirectURI(),
                    "requested_scope": scope,
                },
            })
            return
        } else if c.Request.Method == "POST" || (activeSessions > 0 && c.Request.Method == "GET") {
            if c.PostForm("access_denied") == "true" {
                location = fmt.Sprintf(errorURI, redirectURI, oauth.AccessDenied, state)
                c.Redirect(http.StatusFound, location)
                return
            }
            result, err := oauth.AuthorizationCodeGrant(utils.H{
                "response_type": responseType,
                "client": client,
                "user": user,
                "ip": c.Request.RemoteAddr,
                "userAgent": c.Request.UserAgent(),
                "redirect_uri": redirectURI,
                "scope": scope,
                "state": state,
            })
            if err != nil {
                location = fmt.Sprintf(errorURI, redirectURI, result["error"], result["state"])
                c.Redirect(http.StatusFound, location)
            } else {
                location = fmt.Sprintf("%s?code=%s&scope=%s&state=%s",
                    redirectURI, result["code"], result["scope"], result["state"])
                c.Redirect(http.StatusFound, location)
            }
        } else {
            c.String(http.StatusNotFound, "404 Not Found")
        }
    // Implicit Grant
    case oauth.Token:
        location = fmt.Sprintf(errorURI,
            redirectURI, oauth.UnsupportedResponseType, state)
        c.Redirect(http.StatusFound, location)
        return
    default:
        location = fmt.Sprintf(errorURI,
            redirectURI, oauth.InvalidRequest, state)
        c.Redirect(http.StatusFound, location)
        return
    }
}
<commit_msg>Hotfix for getting Authorization header attribute<commit_after>package web

import (
    "fmt"
    "net/http"
    "net/url"
    "strings"

    "github.com/gin-gonic/gin"
    "github.com/gin-gonic/contrib/renders/multitemplate"
    "github.com/gin-gonic/contrib/sessions"

    "github.com/earaujoassis/space/config"
    "github.com/earaujoassis/space/models"
    "github.com/earaujoassis/space/oauth"
    "github.com/earaujoassis/space/services"
    "github.com/earaujoassis/space/feature"
    "github.com/earaujoassis/space/utils"
)

const (
    errorURI string = "%s?error=%s&state=%s"
)

func createCustomRender() multitemplate.Render {
    render := multitemplate.New()
    render.AddFromFiles("satellite", "web/templates/default.html", "web/templates/satellite.html")
    return render
}

func ExposeRoutes(router *gin.Engine) {
    router.LoadHTMLGlob(config.GetConfig("template_folder").(string))
    router.HTMLRender = createCustomRender()
    router.Static("/public", config.GetConfig("public_folder").(string))
    store := sessions.NewCookieStore([]byte(config.GetConfig("http.session_secret").(string)))
    router.Use(sessions.Sessions("jupiter", store))
    views := router.Group("/")
    {
        views.GET("/", jupiterHandler)
        views.GET("/profile", jupiterHandler)

        views.GET("/signup", func(c *gin.Context) {
            c.HTML(http.StatusOK, "satellite", utils.H{
                "Title": " - Sign up",
                "Satellite": "io",
                "Data": utils.H{
                    "feature.gates": utils.H{
                        "user.create": feature.Active("user.create"),
                    },
                },
            })
        })

        views.GET("/signin", func(c *gin.Context) {
            c.HTML(http.StatusOK, "satellite", utils.H{
                "Title": " - Sign in",
                "Satellite": "ganymede",
            })
        })

        views.GET("/signout", func(c *gin.Context) {
            session := sessions.Default(c)

            userPublicId := session.Get("userPublicId")
            if userPublicId != nil {
                session.Delete("userPublicId")
                session.Save()
            }

            c.Redirect(http.StatusFound, "/signin")
        })

        views.GET("/session", func(c *gin.Context) {
            session := sessions.Default(c)

            userPublicId := session.Get("userPublicId")
            if userPublicId != nil {
                c.Redirect(http.StatusFound, "/")
                return
            }

            var nextPath string = "/"
            var scope string = c.Query("scope")
            var grantType string = c.Query("grant_type")
            var code string = c.Query("code")
            var clientId string = c.Query("client_id")
            var _nextPath string = c.Query("_")
            //var state string = c.Query("state")

            if scope == "" || grantType == "" || code == "" || clientId == "" {
                c.String(http.StatusMethodNotAllowed, "Missing required parameters")
                return
            }
            if _nextPath != "" {
                if _nextPath, err := url.QueryUnescape(_nextPath); err == nil {
                    nextPath = _nextPath
                }
            }

            client := services.FindOrCreateClient("Jupiter")
            if client.Key == clientId && grantType == oauth.AuthorizationCode && scope == models.PublicScope {
                grantToken := services.FindSessionByToken(code, models.GrantToken)
                if grantToken.ID != 0 {
                    session.Set("userPublicId", grantToken.User.PublicId)
                    session.Save()
                    services.InvalidateSession(grantToken)
                    c.Redirect(http.StatusFound, nextPath)
                    return
                }
            }

            c.Redirect(http.StatusFound, "/signin")
        })

        views.GET("/authorize", authorizeHandler)
        views.POST("/authorize", authorizeHandler)

        views.GET("/error", func(c *gin.Context) {
            c.String(http.StatusMethodNotAllowed, "Not implemented")
        })

        views.POST("/token", func(c *gin.Context) {
            var grantType string = c.PostForm("grant_type")

            authorizationBasic := strings.Replace(c.Request.Header.Get("Authorization"), "Basic ", "", 1)
            client := oauth.ClientAuthentication(authorizationBasic)
            if client.ID == 0 {
                c.Header("WWW-Authenticate", fmt.Sprintf("Basic realm=\"%s\"", c.Request.RequestURI))
                c.JSON(http.StatusUnauthorized, utils.H{
                    "error": oauth.AccessDenied,
                })
                return
            }

            switch grantType {
            // Authorization Code Grant
            case oauth.AuthorizationCode:
                result, err := oauth.AccessTokenRequest(utils.H{
                    "grant_type": grantType,
                    "code": c.PostForm("code"),
                    "redirect_uri": c.PostForm("redirect_uri"),
                    "client": client,
                })
                if err != nil {
                    c.JSON(http.StatusMethodNotAllowed, utils.H{
                        "error": result["error"],
                    })
                    return
                } else {
                    c.JSON(http.StatusOK, utils.H{
                        "user_id": result["user_id"],
                        "access_token": result["access_token"],
                        "token_type": result["token_type"],
                        "expires_in": result["expires_in"],
                        "refresh_token": result["refresh_token"],
                        "scope": result["scope"],
                    })
                    return
                }
                return
            // Refreshing an Access Token
            case oauth.RefreshToken:
                result, err := oauth.RefreshTokenRequest(utils.H{
                    "grant_type": grantType,
                    "refresh_token": c.PostForm("refresh_token"),
                    "scope": c.PostForm("scope"),
                    "client": client,
                })
                if err != nil {
                    c.JSON(http.StatusMethodNotAllowed, utils.H{
                        "error": result["error"],
                    })
                    return
                } else {
                    c.JSON(http.StatusOK, utils.H{
                        "user_id": result["user_id"],
                        "access_token": result["access_token"],
                        "token_type": result["token_type"],
                        "expires_in": result["expires_in"],
                        "refresh_token": result["refresh_token"],
                        "scope": result["scope"],
                    })
                    return
                }
                return
            // Resource Owner Password Credentials Grant
            // Client Credentials Grant
            case oauth.Password, oauth.ClientCredentials:
                c.JSON(http.StatusMethodNotAllowed, utils.H{
                    "error": oauth.UnsupportedGrantType,
                })
                return
            default:
                c.JSON(http.StatusBadRequest, utils.H{
                    "error": oauth.InvalidRequest,
                })
                return
            }
        })
    }
}

func jupiterHandler(c *gin.Context) {
    session := sessions.Default(c)
    userPublicId := session.Get("userPublicId")
    if userPublicId == nil {
        c.Redirect(http.StatusFound, "/signin")
        return
    }
    client := services.FindOrCreateClient("Jupiter")
    user := services.FindUserByPublicId(userPublicId.(string))
    actionToken := services.CreateAction(user, client,
        c.Request.RemoteAddr,
        c.Request.UserAgent(),
        models.ReadWriteScope)
    c.HTML(http.StatusOK, "satellite", utils.H{
        "Title": " - Mission control",
        "Satellite": "europa",
        "Data": utils.H {
            "action_token": actionToken.Token,
            "user_id": user.UUID,
        },
    })
}

func authorizeHandler(c *gin.Context) {
    var location string
    var responseType string
    var clientId string
    var redirectURI string
    var scope string
    var state string

    session := sessions.Default(c)
    userPublicId := session.Get("userPublicId")
    nextPath := url.QueryEscape(fmt.Sprintf("%s?%s", c.Request.URL.Path, c.Request.URL.RawQuery))
    if userPublicId == nil {
        location = fmt.Sprintf("/signin?_=%s", nextPath)
        c.Redirect(http.StatusFound, location)
        return
    }
    user := services.FindUserByPublicId(userPublicId.(string))
    if user.ID == 0 {
        session.Delete("userPublicId")
        session.Save()
        location = fmt.Sprintf("/signin?_=%s", nextPath)
        c.Redirect(http.StatusFound, location)
        return
    }

    responseType = c.Query("response_type")
    clientId = c.Query("client_id")
    redirectURI = c.Query("redirect_uri")
    scope = c.Query("scope")
    state = c.Query("state")

    if redirectURI == "" {
        redirectURI = "/error"
    }

    client := services.FindClientByKey(clientId)
    if client.ID == 0 {
        redirectURI = "/error"
        location = fmt.Sprintf("%s?error=%s&state=%s",
            redirectURI, oauth.UnauthorizedClient, state)
        c.Redirect(http.StatusFound, location)
        return
    }

    if scope != models.PublicScope && scope != models.ReadScope && scope != models.ReadWriteScope {
        scope = "public"
    }

    switch responseType {
    // Authorization Code Grant
    case oauth.Code:
        activeSessions := services.ActiveSessionsForClient(client.ID, user.ID)
        if c.Request.Method == "GET" && activeSessions == 0 {
            c.HTML(http.StatusOK, "satellite", utils.H{
                "Title": " - Authorize",
                "Satellite": "callisto",
                "Data": utils.H{
                    "first_name": user.FirstName,
                    "last_name": user.LastName,
                    "client_name": client.Name,
                    "client_uri": client.DefaultRedirectURI(),
                    "requested_scope": scope,
                },
            })
            return
        } else if c.Request.Method == "POST" || (activeSessions > 0 && c.Request.Method == "GET") {
            if c.PostForm("access_denied") == "true" {
                location = fmt.Sprintf(errorURI, redirectURI, oauth.AccessDenied, state)
                c.Redirect(http.StatusFound, location)
                return
            }
            result, err := oauth.AuthorizationCodeGrant(utils.H{
                "response_type": responseType,
                "client": client,
                "user": user,
                "ip": c.Request.RemoteAddr,
                "userAgent": c.Request.UserAgent(),
                "redirect_uri": redirectURI,
                "scope": scope,
                "state": state,
            })
            if err != nil {
                location = fmt.Sprintf(errorURI, redirectURI, result["error"], result["state"])
                c.Redirect(http.StatusFound, location)
            } else {
                location = fmt.Sprintf("%s?code=%s&scope=%s&state=%s",
                    redirectURI, result["code"], result["scope"], result["state"])
                c.Redirect(http.StatusFound, location)
            }
        } else {
            c.String(http.StatusNotFound, "404 Not Found")
        }
    // Implicit Grant
    case oauth.Token:
        location = fmt.Sprintf(errorURI,
            redirectURI, oauth.UnsupportedResponseType, state)
        c.Redirect(http.StatusFound, location)
        return
    default:
        location = fmt.Sprintf(errorURI,
            redirectURI, oauth.InvalidRequest, state)
        c.Redirect(http.StatusFound, location)
        return
    }
}
<|endoftext|>


<commit_before><commit_msg>* bayesian sensor code<commit_after><|endoftext|>


<commit_before><commit_msg>Blocklist<commit_after><|endoftext|>


<commit_before>package dynamodb

import (
	SDK "github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/mitchellh/mapstructure"
)

const defaultResultTag = "dynamodb"

// QueryResult is struct for result of Query operation.
type QueryResult struct {
	Items            []map[string]*SDK.AttributeValue
	LastEvaluatedKey map[string]*SDK.AttributeValue
	Count            int64
	ScannedCount     int64
	tagName          string
}

// ToSliceMap converts result to slice of map.
func (r QueryResult) ToSliceMap() []map[string]interface{} {
	m := make([]map[string]interface{}, len(r.Items))
	for i, item := range r.Items {
		m[i] = UnmarshalAttributeValue(item)
	}
	return m
}

// Unmarshal parse DynamoDB item data and mapping value to given slice pointer sturct.
//     e.g. err = Unmarshal(&[]*yourStruct)
func (r QueryResult) Unmarshal(v interface{}) error {
	m := r.ToSliceMap()

	tagName := r.tagName
	if tagName == "" {
		tagName = defaultResultTag
	}

	decoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{
		Metadata: nil,
		Result:   v,
		TagName:  tagName,
	})
	if err != nil {
		return err
	}
	return decoder.Decode(m)
}
<commit_msg>[DynamoDB] Fix DynamoDB Unmarshal to use weak type (#22)<commit_after>package dynamodb

import (
	SDK "github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/mitchellh/mapstructure"
)

const defaultResultTag = "dynamodb"

// QueryResult is struct for result of Query operation.
type QueryResult struct {
	Items            []map[string]*SDK.AttributeValue
	LastEvaluatedKey map[string]*SDK.AttributeValue
	Count            int64
	ScannedCount     int64
	tagName          string
}

// ToSliceMap converts result to slice of map.
func (r QueryResult) ToSliceMap() []map[string]interface{} {
	m := make([]map[string]interface{}, len(r.Items))
	for i, item := range r.Items {
		m[i] = UnmarshalAttributeValue(item)
	}
	return m
}

// Unmarshal parse DynamoDB item data and mapping value to given slice pointer sturct.
//     e.g. err = Unmarshal(&[]*yourStruct)
func (r QueryResult) Unmarshal(v interface{}) error {
	m := r.ToSliceMap()

	tagName := r.tagName
	if tagName == "" {
		tagName = defaultResultTag
	}

	decoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{
		Metadata:         nil,
		Result:           v,
		TagName:          tagName,
		WeaklyTypedInput: true,
	})
	if err != nil {
		return err
	}
	return decoder.Decode(m)
}
<|endoftext|>


<commit_before>package agent

import (
	"bufio"
	"fmt"
	"github.com/hashicorp/logutils"
	"github.com/hashicorp/serf/serf"
	"github.com/mitchellh/mapstructure"
	"github.com/ugorji/go/codec"
	"log"
	"net"
	"sync"
	"sync/atomic"
)

var (
	clientClosed = fmt.Errorf("client closed")
)

type rpcResponseTuple struct {
	response interface{}
	err      error
}

type seqListener struct {
	// The 'once' handler is called on the first message
	first func(resp interface{}, err error)

	// The 'stream' handler is called for each subsequent message
	stream func(resp interface{}, err error)
}

// RPCClient is the RPC client to make requests to the agent RPC.
type RPCClient struct {
	seq          int32
	conn         *net.TCPConn
	reader       *bufio.Reader
	writer       *bufio.Writer
	dec          *codec.Decoder
	enc          *codec.Encoder
	writeLock    sync.Mutex
	dispatch     map[int]*seqListener
	dispatchLock sync.Mutex
	shutdown     bool
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex
}

// send is used to send an object using the MsgPack encoding. send
// is serialized to prevent write overlaps, while properly buffering.
func (c *RPCClient) send(obj interface{}) error {
	c.writeLock.Lock()
	defer c.writeLock.Unlock()

	if c.shutdown {
		return clientClosed
	}

	if err := c.enc.Encode(obj); err != nil {
		return err
	}

	if err := c.writer.Flush(); err != nil {
		return err
	}

	return nil
}

// NewRPCClient is used to create a new RPC client given the address.
// This will properly dial, handshake, and start listening
func NewRPCClient(addr string) (*RPCClient, error) {
	// Try to dial to serf
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}

	// Create the client
	client := &RPCClient{
		seq:        0,
		conn:       conn.(*net.TCPConn),
		reader:     bufio.NewReader(conn),
		writer:     bufio.NewWriter(conn),
		dispatch:   make(map[int]*seqListener),
		shutdownCh: make(chan struct{}),
	}
	client.dec = codec.NewDecoder(client.reader,
		&codec.MsgpackHandle{RawToString: true, WriteExt: true})
	client.enc = codec.NewEncoder(client.writer,
		&codec.MsgpackHandle{RawToString: true, WriteExt: true})
	go client.listen()

	// Do the initial handshake
	if err := client.handshake(); err != nil {
		client.Close()
		return nil, err
	}
	return client, err
}

// StreamHandle is an opaque handle passed to stop to stop streaming
type StreamHandle int

// Close is used to free any resources associated with the client
func (c *RPCClient) Close() error {
	c.shutdownLock.Lock()
	defer c.shutdownLock.Unlock()

	if !c.shutdown {
		c.shutdown = true
		close(c.shutdownCh)
		return c.conn.Close()
	}
	return nil
}

// ForceLeave is used to ask the agent to issue a leave command for
// a given node
func (c *RPCClient) ForceLeave(node string) error {
	req := forceLeaveRequest{
		Command: forceLeaveCommand,
		Seq:     c.getSeq(),
		Node:    node,
	}
	return c.genericRPC(req.Seq, req)
}

// Join is used to instruct the agent to attempt a join
func (c *RPCClient) Join(addrs []string, ignoreOld bool) (n int, err error) {
	req := joinRequest{
		Command:  joinCommand,
		Seq:      c.getSeq(),
		Existing: addrs,
		Replay:   !ignoreOld,
	}
	ch := c.waitSeq(req.Seq)
	if err := c.send(req); err != nil {
		return 0, err
	}

	select {
	case tuple := <-ch:
		if tuple.err != nil {
			return 0, tuple.err
		}

		// Decode the response
		var resp joinResponse
		if err := mapstructure.Decode(tuple.response, &resp); err != nil {
			return 0, err
		}
		return resp.Num, strError(resp.Error)

	case <-c.shutdownCh:
		return 0, clientClosed
	}
}

// Members is used to fetch a list of known members
func (c *RPCClient) Members() ([]serf.Member, error) {
	req := membersRequest{
		Command: membersCommand,
		Seq:     c.getSeq(),
	}
	ch := c.waitSeq(req.Seq)
	if err := c.send(req); err != nil {
		return nil, err
	}

	select {
	case tuple := <-ch:
		// Check for an error
		if tuple.err != nil {
			return nil, tuple.err
		}

		// Decode the response
		var resp membersResponse
		if err := mapstructure.Decode(tuple.response, &resp); err != nil {
			return nil, err
		}
		return resp.Members, nil

	case <-c.shutdownCh:
		return nil, clientClosed
	}
}

// UserEvent is used to trigger sending an event
func (c *RPCClient) UserEvent(name string, payload []byte, coalesce bool) error {
	req := eventRequest{
		Command:  eventCommand,
		Seq:      c.getSeq(),
		Name:     name,
		Payload:  payload,
		Coalesce: coalesce,
	}
	return c.genericRPC(req.Seq, req)
}

// Monitor is used to subscribe to the logs of the agent
func (c *RPCClient) Monitor(level logutils.LogLevel, ch chan<- string) (StreamHandle, error) {
	// Setup the request
	seq := c.getSeq()
	req := monitorRequest{
		Command:  monitorCommand,
		Seq:      seq,
		LogLevel: string(level),
	}

	// Create a stream handler
	handler := func(resp interface{}, err error) {
		resp_map, ok := resp.(map[string]interface{})
		if !ok {
			return
		}
		raw, ok := getField(resp_map, "Log")
		if !ok {
			return
		}
		if log, ok := raw.(string); ok {
			ch <- log
		}
	}
	c.streamHandler(seq, handler)

	// Try to establish the
	if err := c.genericRPC(req.Seq, req); err != nil {
		c.deregisterHandler(seq)
		return 0, err
	}

	// Return a handle
	return StreamHandle(seq), nil
}

// Stream is used to subscribe to events
func (c *RPCClient) Stream(filter string, ch chan<- map[string]interface{}) (StreamHandle, error) {
	// Setup the request
	seq := c.getSeq()
	req := streamRequest{
		Command: streamCommand,
		Seq:     seq,
		Type:    filter,
	}

	// Create a stream handler
	handler := func(resp interface{}, err error) {
		resp_map, ok := resp.(map[string]interface{})
		if !ok {
			return
		}
		ch <- resp_map
	}
	c.streamHandler(seq, handler)

	if err := c.genericRPC(req.Seq, req); err != nil {
		c.deregisterHandler(seq)
		return 0, err
	}

	// Return a handle
	return StreamHandle(seq), nil
}

// Stop is used to unsubscribe from logs or event streams
func (c *RPCClient) Stop(handle StreamHandle) error {
	// Deregister locally first to stop delivery
	c.deregisterHandler(int(handle))

	req := stopRequest{
		Command: stopCommand,
		Seq:     c.getSeq(),
		Stop:    int(handle),
	}
	return c.genericRPC(req.Seq, req)
}

// handshake is used to perform the initial handshake on connect
func (c *RPCClient) handshake() error {
	req := handshakeRequest{
		Command: handshakeCommand,
		Seq:     c.getSeq(),
		Version: MaxIPCVersion,
	}
	return c.genericRPC(req.Seq, req)
}

// genericRPC is used to send a request and wait for an
// errorSequenceResponse, potentially returning an error
func (c *RPCClient) genericRPC(seq int, req interface{}) error {
	ch := c.waitSeq(seq)
	if err := c.send(req); err != nil {
		return err
	}
	select {
	case tuple := <-ch:
		if tuple.err != nil {
			return tuple.err
		}
		return c.maybeError(tuple.response)
	case <-c.shutdownCh:
		return clientClosed
	}
}

// maybeError is used for errorSeqResponse, where we may have an error
func (c *RPCClient) maybeError(resp interface{}) error {
	var errSeq errorSeqResponse
	if err := mapstructure.Decode(resp, &errSeq); err != nil {
		return err
	}
	return strError(errSeq.Error)
}

// strError converts a string to an error if not blank
func strError(s string) error {
	if s != "" {
		return fmt.Errorf(s)
	}
	return nil
}

// getSeq returns the next sequence number in a safe manner
func (c *RPCClient) getSeq() int {
	return int(atomic.AddInt32(&c.seq, 1))
}

// deregisterHandler is used to deregister a handler
func (c *RPCClient) deregisterHandler(seq int) {
	c.dispatchLock.Lock()
	defer c.dispatchLock.Unlock()
	delete(c.dispatch, seq)
}

// streamHandler is used to set a stream handler that is
// used after the first response to a sequence number is received
func (c *RPCClient) streamHandler(seq int, handler func(resp interface{}, err error)) {
	c.dispatchLock.Lock()
	defer c.dispatchLock.Unlock()

	seqL := c.dispatch[seq]
	if seqL == nil {
		seqL = new(seqListener)
	}
	seqL.stream = handler
	c.dispatch[seq] = seqL
}

// waitSeq is used to setup a channel to wait on a response for
// a given sequence number.
func (c *RPCClient) waitSeq(seq int) chan rpcResponseTuple {
	c.dispatchLock.Lock()
	defer c.dispatchLock.Unlock()

	ch := make(chan rpcResponseTuple, 1)
	handler := func(resp interface{}, err error) {
		ch <- rpcResponseTuple{resp, err}
	}

	seqL := c.dispatch[seq]
	if seqL == nil {
		seqL = new(seqListener)
	}
	seqL.first = handler
	c.dispatch[seq] = seqL
	return ch
}

// respondSeq is used to respond to a given sequence number
func (c *RPCClient) respondSeq(seq int, resp interface{}, err error) {
	c.dispatchLock.Lock()
	defer c.dispatchLock.Unlock()

	seqL, ok := c.dispatch[seq]
	if !ok {
		return
	}

	// Invoke the first handler if it exists
	if seqL.first != nil {
		seqL.first(resp, nil)
		seqL.first = nil
		if seqL.stream == nil {
			delete(c.dispatch, seq)
		}
		return
	}

	// Invoke the stream handler otherwise
	seqL.stream(resp, nil)
}

// listen is used to processes data coming over the IPC channel,
// and wrote it to the correct destination based on seq no
func (c *RPCClient) listen() {
	defer c.Close()
	var resp map[string]interface{}
	for {
		resp = nil
		if err := c.dec.Decode(&resp); err != nil {
			if !c.shutdown {
				log.Printf("[ERR] agent.client: Failed to decode response: %v", err)
			}
			break
		}

		// Look for the seq
		seq_raw, ok := getField(resp, "Seq")
		if !ok {
			log.Printf("[ERR] agent.client: Response missing Seq: %#v", resp)
			continue
		}

		// Try to convert
		var seq int
		switch typed := seq_raw.(type) {
		case int:
			seq = typed
		case int32:
			seq = int(typed)
		case int64:
			seq = int(typed)
		default:
			log.Printf("[ERR] agent.client: Seq not int: %#v", seq_raw)
			continue
		}

		// Try to dispatch
		c.respondSeq(seq, resp, nil)
	}
}
<commit_msg>agent: RPCClient closes channel on deregister<commit_after>package agent

import (
	"bufio"
	"fmt"
	"github.com/hashicorp/logutils"
	"github.com/hashicorp/serf/serf"
	"github.com/mitchellh/mapstructure"
	"github.com/ugorji/go/codec"
	"log"
	"net"
	"sync"
	"sync/atomic"
)

var (
	clientClosed = fmt.Errorf("client closed")
)

type rpcResponseTuple struct {
	response interface{}
	err      error
}

type seqListener struct {
	// The 'once' handler is called on the first message
	first func(resp interface{}, err error)

	// The 'stream' handler is called for each subsequent message
	stream func(resp interface{}, err error)

	// The cleanup func is called whent he listener is being deregistered
	cleanup func()
}

// RPCClient is the RPC client to make requests to the agent RPC.
type RPCClient struct {
	seq          int32
	conn         *net.TCPConn
	reader       *bufio.Reader
	writer       *bufio.Writer
	dec          *codec.Decoder
	enc          *codec.Encoder
	writeLock    sync.Mutex
	dispatch     map[int]*seqListener
	dispatchLock sync.Mutex
	shutdown     bool
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex
}

// send is used to send an object using the MsgPack encoding. send
// is serialized to prevent write overlaps, while properly buffering.
func (c *RPCClient) send(obj interface{}) error {
	c.writeLock.Lock()
	defer c.writeLock.Unlock()

	if c.shutdown {
		return clientClosed
	}

	if err := c.enc.Encode(obj); err != nil {
		return err
	}

	if err := c.writer.Flush(); err != nil {
		return err
	}

	return nil
}

// NewRPCClient is used to create a new RPC client given the address.
// This will properly dial, handshake, and start listening
func NewRPCClient(addr string) (*RPCClient, error) {
	// Try to dial to serf
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}

	// Create the client
	client := &RPCClient{
		seq:        0,
		conn:       conn.(*net.TCPConn),
		reader:     bufio.NewReader(conn),
		writer:     bufio.NewWriter(conn),
		dispatch:   make(map[int]*seqListener),
		shutdownCh: make(chan struct{}),
	}
	client.dec = codec.NewDecoder(client.reader,
		&codec.MsgpackHandle{RawToString: true, WriteExt: true})
	client.enc = codec.NewEncoder(client.writer,
		&codec.MsgpackHandle{RawToString: true, WriteExt: true})
	go client.listen()

	// Do the initial handshake
	if err := client.handshake(); err != nil {
		client.Close()
		return nil, err
	}
	return client, err
}

// StreamHandle is an opaque handle passed to stop to stop streaming
type StreamHandle int

// Close is used to free any resources associated with the client
func (c *RPCClient) Close() error {
	c.shutdownLock.Lock()
	defer c.shutdownLock.Unlock()

	if !c.shutdown {
		c.shutdown = true
		close(c.shutdownCh)
		c.deregisterAll()
		return c.conn.Close()
	}
	return nil
}

// ForceLeave is used to ask the agent to issue a leave command for
// a given node
func (c *RPCClient) ForceLeave(node string) error {
	req := forceLeaveRequest{
		Command: forceLeaveCommand,
		Seq:     c.getSeq(),
		Node:    node,
	}
	return c.genericRPC(req.Seq, req)
}

// Join is used to instruct the agent to attempt a join
func (c *RPCClient) Join(addrs []string, ignoreOld bool) (n int, err error) {
	req := joinRequest{
		Command:  joinCommand,
		Seq:      c.getSeq(),
		Existing: addrs,
		Replay:   !ignoreOld,
	}
	ch := c.waitSeq(req.Seq)
	if err := c.send(req); err != nil {
		return 0, err
	}

	select {
	case tuple := <-ch:
		if tuple.err != nil {
			return 0, tuple.err
		}

		// Decode the response
		var resp joinResponse
		if err := mapstructure.Decode(tuple.response, &resp); err != nil {
			return 0, err
		}
		return resp.Num, strError(resp.Error)

	case <-c.shutdownCh:
		return 0, clientClosed
	}
}

// Members is used to fetch a list of known members
func (c *RPCClient) Members() ([]serf.Member, error) {
	req := membersRequest{
		Command: membersCommand,
		Seq:     c.getSeq(),
	}
	ch := c.waitSeq(req.Seq)
	if err := c.send(req); err != nil {
		return nil, err
	}

	select {
	case tuple := <-ch:
		// Check for an error
		if tuple.err != nil {
			return nil, tuple.err
		}

		// Decode the response
		var resp membersResponse
		if err := mapstructure.Decode(tuple.response, &resp); err != nil {
			return nil, err
		}
		return resp.Members, nil

	case <-c.shutdownCh:
		return nil, clientClosed
	}
}

// UserEvent is used to trigger sending an event
func (c *RPCClient) UserEvent(name string, payload []byte, coalesce bool) error {
	req := eventRequest{
		Command:  eventCommand,
		Seq:      c.getSeq(),
		Name:     name,
		Payload:  payload,
		Coalesce: coalesce,
	}
	return c.genericRPC(req.Seq, req)
}

// Monitor is used to subscribe to the logs of the agent
func (c *RPCClient) Monitor(level logutils.LogLevel, ch chan<- string) (StreamHandle, error) {
	// Setup the request
	seq := c.getSeq()
	req := monitorRequest{
		Command:  monitorCommand,
		Seq:      seq,
		LogLevel: string(level),
	}

	// Create a stream handler
	handler := func(resp interface{}, err error) {
		resp_map, ok := resp.(map[string]interface{})
		if !ok {
			return
		}
		raw, ok := getField(resp_map, "Log")
		if !ok {
			return
		}
		if log, ok := raw.(string); ok {
			ch <- log
		}
	}
	cleanup := func() {
		close(ch)
	}
	c.streamHandler(seq, handler, cleanup)

	// Try to establish the
	if err := c.genericRPC(req.Seq, req); err != nil {
		c.deregisterHandler(seq)
		return 0, err
	}

	// Return a handle
	return StreamHandle(seq), nil
}

// Stream is used to subscribe to events
func (c *RPCClient) Stream(filter string, ch chan<- map[string]interface{}) (StreamHandle, error) {
	// Setup the request
	seq := c.getSeq()
	req := streamRequest{
		Command: streamCommand,
		Seq:     seq,
		Type:    filter,
	}

	// Create a stream handler
	handler := func(resp interface{}, err error) {
		resp_map, ok := resp.(map[string]interface{})
		if !ok {
			return
		}
		ch <- resp_map
	}
	cleanup := func() {
		close(ch)
	}
	c.streamHandler(seq, handler, cleanup)

	if err := c.genericRPC(req.Seq, req); err != nil {
		c.deregisterHandler(seq)
		return 0, err
	}

	// Return a handle
	return StreamHandle(seq), nil
}

// Stop is used to unsubscribe from logs or event streams
func (c *RPCClient) Stop(handle StreamHandle) error {
	// Deregister locally first to stop delivery
	c.deregisterHandler(int(handle))

	req := stopRequest{
		Command: stopCommand,
		Seq:     c.getSeq(),
		Stop:    int(handle),
	}
	return c.genericRPC(req.Seq, req)
}

// handshake is used to perform the initial handshake on connect
func (c *RPCClient) handshake() error {
	req := handshakeRequest{
		Command: handshakeCommand,
		Seq:     c.getSeq(),
		Version: MaxIPCVersion,
	}
	return c.genericRPC(req.Seq, req)
}

// genericRPC is used to send a request and wait for an
// errorSequenceResponse, potentially returning an error
func (c *RPCClient) genericRPC(seq int, req interface{}) error {
	ch := c.waitSeq(seq)
	if err := c.send(req); err != nil {
		return err
	}
	select {
	case tuple := <-ch:
		if tuple.err != nil {
			return tuple.err
		}
		return c.maybeError(tuple.response)
	case <-c.shutdownCh:
		return clientClosed
	}
}

// maybeError is used for errorSeqResponse, where we may have an error
func (c *RPCClient) maybeError(resp interface{}) error {
	var errSeq errorSeqResponse
	if err := mapstructure.Decode(resp, &errSeq); err != nil {
		return err
	}
	return strError(errSeq.Error)
}

// strError converts a string to an error if not blank
func strError(s string) error {
	if s != "" {
		return fmt.Errorf(s)
	}
	return nil
}

// getSeq returns the next sequence number in a safe manner
func (c *RPCClient) getSeq() int {
	return int(atomic.AddInt32(&c.seq, 1))
}

// deregisterAll is used to deregister all handlers
func (c *RPCClient) deregisterAll() {
	c.dispatchLock.Lock()
	defer c.dispatchLock.Unlock()

	for _, seqL := range c.dispatch {
		if seqL.cleanup != nil {
			seqL.cleanup()
		}
	}
	c.dispatch = make(map[int]*seqListener)
}

// deregisterHandler is used to deregister a handler
func (c *RPCClient) deregisterHandler(seq int) {
	c.dispatchLock.Lock()
	defer c.dispatchLock.Unlock()

	seqL, ok := c.dispatch[seq]
	delete(c.dispatch, seq)

	if ok && seqL.cleanup != nil {
		seqL.cleanup()
	}
}

// streamHandler is used to set a stream handler that is
// used after the first response to a sequence number is received
func (c *RPCClient) streamHandler(seq int, handler func(resp interface{}, err error), cleanup func()) {
	c.dispatchLock.Lock()
	defer c.dispatchLock.Unlock()

	seqL := c.dispatch[seq]
	if seqL == nil {
		seqL = new(seqListener)
	}
	seqL.stream = handler
	seqL.cleanup = cleanup
	c.dispatch[seq] = seqL
}

// waitSeq is used to setup a channel to wait on a response for
// a given sequence number.
func (c *RPCClient) waitSeq(seq int) chan rpcResponseTuple {
	c.dispatchLock.Lock()
	defer c.dispatchLock.Unlock()

	ch := make(chan rpcResponseTuple, 1)
	handler := func(resp interface{}, err error) {
		ch <- rpcResponseTuple{resp, err}
	}

	seqL := c.dispatch[seq]
	if seqL == nil {
		seqL = new(seqListener)
	}
	seqL.first = handler
	c.dispatch[seq] = seqL
	return ch
}

// respondSeq is used to respond to a given sequence number
func (c *RPCClient) respondSeq(seq int, resp interface{}, err error) {
	c.dispatchLock.Lock()
	defer c.dispatchLock.Unlock()

	seqL, ok := c.dispatch[seq]
	if !ok {
		return
	}

	// Invoke the first handler if it exists
	if seqL.first != nil {
		seqL.first(resp, nil)
		seqL.first = nil
		if seqL.stream == nil {
			delete(c.dispatch, seq)
			if seqL.cleanup != nil {
				seqL.cleanup()
			}
		}
		return
	}

	// Invoke the stream handler otherwise
	seqL.stream(resp, nil)
}

// listen is used to processes data coming over the IPC channel,
// and wrote it to the correct destination based on seq no
func (c *RPCClient) listen() {
	defer c.Close()
	var resp map[string]interface{}
	for {
		resp = nil
		if err := c.dec.Decode(&resp); err != nil {
			if !c.shutdown {
				log.Printf("[ERR] agent.client: Failed to decode response: %v", err)
			}
			break
		}

		// Look for the seq
		seq_raw, ok := getField(resp, "Seq")
		if !ok {
			log.Printf("[ERR] agent.client: Response missing Seq: %#v", resp)
			continue
		}

		// Try to convert
		var seq int
		switch typed := seq_raw.(type) {
		case int:
			seq = typed
		case int32:
			seq = int(typed)
		case int64:
			seq = int(typed)
		default:
			log.Printf("[ERR] agent.client: Seq not int: %#v", seq_raw)
			continue
		}

		// Try to dispatch
		c.respondSeq(seq, resp, nil)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors.

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 capabilities

import (
	"sync"
)

// Capabilities defines the set of capabilities available within the system.
// For now these are global.  Eventually they may be per-user
type Capabilities struct {
	AllowPrivileged bool

	// Pod sources from which to allow privileged capabilities like host networking, sharing the host
	// IPC namespace, and sharing the host PID namespace.
	PrivilegedSources PrivilegedSources

	// PerConnectionBandwidthLimitBytesPerSec limits the throughput of each connection (currently only used for proxy, exec, attach)
	PerConnectionBandwidthLimitBytesPerSec int64
}

// PrivilegedSources defines the pod sources allowed to make privileged requests for certain types
// of capabilities like host networking, sharing the host IPC namespace, and sharing the host PID namespace.
type PrivilegedSources struct {
	// List of pod sources for which using host network is allowed.
	HostNetworkSources []string

	// List of pod sources for which using host pid namespace is allowed.
	HostPIDSources []string

	// List of pod sources for which using host ipc is allowed.
	HostIPCSources []string
}

// TODO: Clean these up into a singleton
var once sync.Once
var lock sync.Mutex
var capabilities *Capabilities

// Initialize the capability set.  This can only be done once per binary, subsequent calls are ignored.
func Initialize(c Capabilities) {
	// Only do this once
	once.Do(func() {
		capabilities = &c
	})
}

// Setup the capability set.  It wraps Initialize for improving usability.
func Setup(allowPrivileged bool, privilegedSources PrivilegedSources, perConnectionBytesPerSec int64) {
	Initialize(Capabilities{
		AllowPrivileged:                        allowPrivileged,
		PrivilegedSources:                      privilegedSources,
		PerConnectionBandwidthLimitBytesPerSec: perConnectionBytesPerSec,
	})
}

// SetCapabilitiesForTests.  Convenience method for testing.  This should only be called from tests.
func SetForTests(c Capabilities) {
	lock.Lock()
	defer lock.Unlock()
	capabilities = &c
}

// Returns a read-only copy of the system capabilities.
func Get() Capabilities {
	lock.Lock()
	defer lock.Unlock()
	// This check prevents clobbering of capabilities that might've been set via SetForTests
	if capabilities == nil {
		Initialize(Capabilities{
			AllowPrivileged: false,
			PrivilegedSources: PrivilegedSources{
				HostNetworkSources: []string{},
				HostPIDSources:     []string{},
				HostIPCSources:     []string{},
			},
		})
	}
	return *capabilities
}
<commit_msg>Fix comment for method SetForTests<commit_after>/*
Copyright 2014 The Kubernetes Authors.

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 capabilities

import (
	"sync"
)

// Capabilities defines the set of capabilities available within the system.
// For now these are global.  Eventually they may be per-user
type Capabilities struct {
	AllowPrivileged bool

	// Pod sources from which to allow privileged capabilities like host networking, sharing the host
	// IPC namespace, and sharing the host PID namespace.
	PrivilegedSources PrivilegedSources

	// PerConnectionBandwidthLimitBytesPerSec limits the throughput of each connection (currently only used for proxy, exec, attach)
	PerConnectionBandwidthLimitBytesPerSec int64
}

// PrivilegedSources defines the pod sources allowed to make privileged requests for certain types
// of capabilities like host networking, sharing the host IPC namespace, and sharing the host PID namespace.
type PrivilegedSources struct {
	// List of pod sources for which using host network is allowed.
	HostNetworkSources []string

	// List of pod sources for which using host pid namespace is allowed.
	HostPIDSources []string

	// List of pod sources for which using host ipc is allowed.
	HostIPCSources []string
}

// TODO: Clean these up into a singleton
var once sync.Once
var lock sync.Mutex
var capabilities *Capabilities

// Initialize the capability set.  This can only be done once per binary, subsequent calls are ignored.
func Initialize(c Capabilities) {
	// Only do this once
	once.Do(func() {
		capabilities = &c
	})
}

// Setup the capability set.  It wraps Initialize for improving usability.
func Setup(allowPrivileged bool, privilegedSources PrivilegedSources, perConnectionBytesPerSec int64) {
	Initialize(Capabilities{
		AllowPrivileged:                        allowPrivileged,
		PrivilegedSources:                      privilegedSources,
		PerConnectionBandwidthLimitBytesPerSec: perConnectionBytesPerSec,
	})
}

// SetForTests sets capabilities for tests.  Convenience method for testing.  This should only be called from tests.
func SetForTests(c Capabilities) {
	lock.Lock()
	defer lock.Unlock()
	capabilities = &c
}

// Returns a read-only copy of the system capabilities.
func Get() Capabilities {
	lock.Lock()
	defer lock.Unlock()
	// This check prevents clobbering of capabilities that might've been set via SetForTests
	if capabilities == nil {
		Initialize(Capabilities{
			AllowPrivileged: false,
			PrivilegedSources: PrivilegedSources{
				HostNetworkSources: []string{},
				HostPIDSources:     []string{},
				HostIPCSources:     []string{},
			},
		})
	}
	return *capabilities
}
<|endoftext|>