<commit_before>package goahocorasick

type AhoCorasick struct {
	root *node
}

type node struct {
	children map[rune]*node
	depth    int
	fail     *node
	hit      bool
}

func New(keywords []string) *AhoCorasick {
	a := AhoCorasick{root: newNode(0)}
	a.createTrie(keywords)
	a.createFail()
	return &a
}

func newNode(depth int) *node {
	return &node{
		children: map[rune]*node{},
		depth:    depth}
}

func (a *AhoCorasick) createTrie(keywords []string) {
	for _, keyword := range keywords {
		n := a.root
		for _, r := range keyword {
			v, ok := n.children[r]
			if !ok {
				v = newNode(n.depth + 1)
				n.children[r] = v
			}
			n = v
		}
		n.hit = true
	}
}

func (a *AhoCorasick) createFail() {
	for k, v := range a.root.children {
		a.walkCreateFail(v, []rune{k})
	}
}

func (a *AhoCorasick) walkCreateFail(n *node, text []rune) {
	n.fail = a.backwardMatchNode(text)
	for k, v := range n.children {
		a.walkCreateFail(v, append(text, k))
	}
}

func (a *AhoCorasick) backwardMatchNode(text []rune) *node {
	for t := text[1:]; len(t) > 0; t = t[1:] {
		n, ok := a.matchNode(t)
		if ok {
			return n
		}
	}
	return a.root
}

func (a *AhoCorasick) matchNode(text []rune) (*node, bool) {
	n := a.root
	for _, r := range text {
		v, ok := n.children[r]
		if ok {
			n = v
		} else {
			return nil, false
		}
	}
	return n, true
}

func (a *AhoCorasick) Match(text string) [][]int {
	runes := []rune(text)
	result := make([][]int, 0)
	n := a.root
	i := 0

	check := func(node *node) {
		if node.hit {
			result = append(result, []int{i - node.depth, node.depth})
		}
	}

	for i < len(runes) {
		check(n)
		child, ok := n.children[runes[i]]
		if ok {
			for n != a.root {
				n = n.fail
				check(n)
			}
			n = child
			i++
		} else if n == a.root {
			i++
		} else {
			n = n.fail
		}
	}

	return result
}
<commit_msg>Prevent mem copy<commit_after>package goahocorasick

type AhoCorasick struct {
	root *node
}

type node struct {
	children map[rune]*node
	depth    int
	fail     *node
	hit      bool
}

func New(keywords []string) *AhoCorasick {
	a := AhoCorasick{root: newNode(0)}
	a.createTrie(keywords)
	a.createFail()
	return &a
}

func newNode(depth int) *node {
	return &node{
		children: map[rune]*node{},
		depth:    depth}
}

func (a *AhoCorasick) createTrie(keywords []string) {
	for _, keyword := range keywords {
		n := a.root
		for _, r := range keyword {
			v, ok := n.children[r]
			if !ok {
				v = newNode(n.depth + 1)
				n.children[r] = v
			}
			n = v
		}
		n.hit = true
	}
}

func (a *AhoCorasick) createFail() {
	for k, v := range a.root.children {
		a.walkCreateFail(v, []rune{k})
	}
}

func (a *AhoCorasick) walkCreateFail(n *node, text []rune) {
	n.fail = a.backwardMatchNode(text)
	for k, v := range n.children {
		a.walkCreateFail(v, append(text, k))
	}
}

func (a *AhoCorasick) backwardMatchNode(text []rune) *node {
	for t := text[1:]; len(t) > 0; t = t[1:] {
		n, ok := a.matchNode(t)
		if ok {
			return n
		}
	}
	return a.root
}

func (a *AhoCorasick) matchNode(text []rune) (*node, bool) {
	n := a.root
	for _, r := range text {
		v, ok := n.children[r]
		if ok {
			n = v
		} else {
			return nil, false
		}
	}
	return n, true
}

func (a *AhoCorasick) Match(text string) [][]int {
	result := make([][]int, 0)
	n := a.root
	i := 0

	for _, r := range text {
	L:
		if n.hit {
			result = append(result, []int{i - n.depth, n.depth})
		}
		child, ok := n.children[r]
		if ok {
			for n != a.root {
				n = n.fail
				if n.hit {
					result = append(result, []int{i - n.depth, n.depth})
				}
			}
			n = child
			i++
		} else if n == a.root {
			i++
		} else {
			n = n.fail
			goto L
		}
	}

	return result
}
<|endoftext|>


<commit_before>// Copyright (C) 2015 Nippon Telegraph and Telephone Corporation.
//
// 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 cmd

import (
	"fmt"
	"io"
	"os"
	"strconv"
	"time"

	"github.com/osrg/gobgp/packet/mrt"
	"github.com/osrg/gobgp/table"
	"github.com/spf13/cobra"
)

func injectMrt(filename string, count int, skip int, onlyBest bool) error {

	file, err := os.Open(filename)
	if err != nil {
		return fmt.Errorf("failed to open file: %s", err)
	}

	idx := 0

	ch := make(chan []*table.Path, 1<<20)

	go func() {

		var peers []*mrt.Peer

		for {
			buf := make([]byte, mrt.MRT_COMMON_HEADER_LEN)
			_, err := file.Read(buf)
			if err == io.EOF {
				break
			} else if err != nil {
				exitWithError(fmt.Errorf("failed to read: %s", err))
			}

			h := &mrt.MRTHeader{}
			err = h.DecodeFromBytes(buf)
			if err != nil {
				exitWithError(fmt.Errorf("failed to parse"))
			}

			buf = make([]byte, h.Len)
			_, err = file.Read(buf)
			if err != nil {
				exitWithError(fmt.Errorf("failed to read"))
			}

			msg, err := mrt.ParseMRTBody(h, buf)
			if err != nil {
				printError(fmt.Errorf("failed to parse: %s", err))
				continue
			}

			if globalOpts.Debug {
				fmt.Println(msg)
			}

			if msg.Header.Type == mrt.TABLE_DUMPv2 {
				subType := mrt.MRTSubTypeTableDumpv2(msg.Header.SubType)
				switch subType {
				case mrt.PEER_INDEX_TABLE:
					peers = msg.Body.(*mrt.PeerIndexTable).Peers
					continue
				case mrt.RIB_IPV4_UNICAST, mrt.RIB_IPV6_UNICAST:
				case mrt.GEO_PEER_TABLE:
					fmt.Printf("WARNING: Skipping GEO_PEER_TABLE: %s", msg.Body.(*mrt.GeoPeerTable))
				default:
					exitWithError(fmt.Errorf("unsupported subType: %v", subType))
				}

				if peers == nil {
					exitWithError(fmt.Errorf("not found PEER_INDEX_TABLE"))
				}

				rib := msg.Body.(*mrt.Rib)
				nlri := rib.Prefix

				paths := make([]*table.Path, 0, len(rib.Entries))

				for _, e := range rib.Entries {
					if len(peers) < int(e.PeerIndex) {
						exitWithError(fmt.Errorf("invalid peer index: %d (PEER_INDEX_TABLE has only %d peers)\n", e.PeerIndex, len(peers)))
					}
					source := &table.PeerInfo{
						AS: peers[e.PeerIndex].AS,
						ID: peers[e.PeerIndex].BgpId,
					}
					t := time.Unix(int64(e.OriginatedTime), 0)
					paths = append(paths, table.NewPath(source, nlri, false, e.PathAttributes, t, false))
				}

				if onlyBest {
					dst := table.NewDestination(nlri, 0)
					for _, p := range paths {
						dst.AddNewPath(p)
					}
					best, _, _ := dst.Calculate().GetChanges(table.GLOBAL_RIB_NAME, false)
					if best == nil {
						exitWithError(fmt.Errorf("Can't find the best %v", nlri))
					}
					paths = []*table.Path{best}
				}

				if idx >= skip {
					ch <- paths
				}

				idx += 1
				if idx == count+skip {
					break
				}
			}
		}

		close(ch)
	}()

	stream, err := client.AddPathByStream()
	if err != nil {
		return fmt.Errorf("failed to modpath: %s", err)
	}

	for paths := range ch {
		err = stream.Send(paths...)
		if err != nil {
			return fmt.Errorf("failed to send: %s", err)
		}
	}

	if err := stream.Close(); err != nil {
		return fmt.Errorf("failed to send: %s", err)
	}
	return nil
}

func NewMrtCmd() *cobra.Command {
	globalInjectCmd := &cobra.Command{
		Use: CMD_GLOBAL,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) < 1 {
				exitWithError(fmt.Errorf("usage: gobgp mrt inject global <filename> [<count> [<skip>]]"))
			}
			filename := args[0]
			count := -1
			skip := 0
			if len(args) > 1 {
				var err error
				count, err = strconv.Atoi(args[1])
				if err != nil {
					exitWithError(fmt.Errorf("invalid count value: %s", args[1]))
				}
				if len(args) > 2 {
					skip, err = strconv.Atoi(args[2])
					if err != nil {
						exitWithError(fmt.Errorf("invalid skip value: %s", args[2]))
					}
				}
			}
			err := injectMrt(filename, count, skip, mrtOpts.Best)
			if err != nil {
				exitWithError(err)
			}
		},
	}

	injectCmd := &cobra.Command{
		Use: CMD_INJECT,
	}
	injectCmd.AddCommand(globalInjectCmd)

	mrtCmd := &cobra.Command{
		Use: CMD_MRT,
	}
	mrtCmd.AddCommand(injectCmd)

	mrtCmd.PersistentFlags().BoolVarP(&mrtOpts.Best, "only-best", "", false, "inject only best paths")
	return mrtCmd
}
<commit_msg>cli: Enable to inject MRT with Path Identifier<commit_after>// Copyright (C) 2015 Nippon Telegraph and Telephone Corporation.
//
// 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 cmd

import (
	"fmt"
	"io"
	"os"
	"strconv"
	"time"

	"github.com/osrg/gobgp/packet/mrt"
	"github.com/osrg/gobgp/table"
	"github.com/spf13/cobra"
)

func injectMrt(filename string, count int, skip int, onlyBest bool) error {

	file, err := os.Open(filename)
	if err != nil {
		return fmt.Errorf("failed to open file: %s", err)
	}

	idx := 0

	ch := make(chan []*table.Path, 1<<20)

	go func() {

		var peers []*mrt.Peer

		for {
			buf := make([]byte, mrt.MRT_COMMON_HEADER_LEN)
			_, err := file.Read(buf)
			if err == io.EOF {
				break
			} else if err != nil {
				exitWithError(fmt.Errorf("failed to read: %s", err))
			}

			h := &mrt.MRTHeader{}
			err = h.DecodeFromBytes(buf)
			if err != nil {
				exitWithError(fmt.Errorf("failed to parse"))
			}

			buf = make([]byte, h.Len)
			_, err = file.Read(buf)
			if err != nil {
				exitWithError(fmt.Errorf("failed to read"))
			}

			msg, err := mrt.ParseMRTBody(h, buf)
			if err != nil {
				printError(fmt.Errorf("failed to parse: %s", err))
				continue
			}

			if globalOpts.Debug {
				fmt.Println(msg)
			}

			if msg.Header.Type == mrt.TABLE_DUMPv2 {
				subType := mrt.MRTSubTypeTableDumpv2(msg.Header.SubType)
				switch subType {
				case mrt.PEER_INDEX_TABLE:
					peers = msg.Body.(*mrt.PeerIndexTable).Peers
					continue
				case mrt.RIB_IPV4_UNICAST, mrt.RIB_IPV6_UNICAST, mrt.RIB_IPV4_UNICAST_ADDPATH, mrt.RIB_IPV6_UNICAST_ADDPATH:
				case mrt.GEO_PEER_TABLE:
					fmt.Printf("WARNING: Skipping GEO_PEER_TABLE: %s", msg.Body.(*mrt.GeoPeerTable))
				default:
					exitWithError(fmt.Errorf("unsupported subType: %v", subType))
				}

				if peers == nil {
					exitWithError(fmt.Errorf("not found PEER_INDEX_TABLE"))
				}

				rib := msg.Body.(*mrt.Rib)
				nlri := rib.Prefix

				paths := make([]*table.Path, 0, len(rib.Entries))

				for _, e := range rib.Entries {
					if len(peers) < int(e.PeerIndex) {
						exitWithError(fmt.Errorf("invalid peer index: %d (PEER_INDEX_TABLE has only %d peers)\n", e.PeerIndex, len(peers)))
					}
					source := &table.PeerInfo{
						AS: peers[e.PeerIndex].AS,
						ID: peers[e.PeerIndex].BgpId,
					}
					t := time.Unix(int64(e.OriginatedTime), 0)
					paths = append(paths, table.NewPath(source, nlri, false, e.PathAttributes, t, false))
				}

				if onlyBest {
					dst := table.NewDestination(nlri, 0)
					for _, p := range paths {
						dst.AddNewPath(p)
					}
					best, _, _ := dst.Calculate().GetChanges(table.GLOBAL_RIB_NAME, false)
					if best == nil {
						exitWithError(fmt.Errorf("Can't find the best %v", nlri))
					}
					paths = []*table.Path{best}
				}

				if idx >= skip {
					ch <- paths
				}

				idx += 1
				if idx == count+skip {
					break
				}
			}
		}

		close(ch)
	}()

	stream, err := client.AddPathByStream()
	if err != nil {
		return fmt.Errorf("failed to modpath: %s", err)
	}

	for paths := range ch {
		err = stream.Send(paths...)
		if err != nil {
			return fmt.Errorf("failed to send: %s", err)
		}
	}

	if err := stream.Close(); err != nil {
		return fmt.Errorf("failed to send: %s", err)
	}
	return nil
}

func NewMrtCmd() *cobra.Command {
	globalInjectCmd := &cobra.Command{
		Use: CMD_GLOBAL,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) < 1 {
				exitWithError(fmt.Errorf("usage: gobgp mrt inject global <filename> [<count> [<skip>]]"))
			}
			filename := args[0]
			count := -1
			skip := 0
			if len(args) > 1 {
				var err error
				count, err = strconv.Atoi(args[1])
				if err != nil {
					exitWithError(fmt.Errorf("invalid count value: %s", args[1]))
				}
				if len(args) > 2 {
					skip, err = strconv.Atoi(args[2])
					if err != nil {
						exitWithError(fmt.Errorf("invalid skip value: %s", args[2]))
					}
				}
			}
			err := injectMrt(filename, count, skip, mrtOpts.Best)
			if err != nil {
				exitWithError(err)
			}
		},
	}

	injectCmd := &cobra.Command{
		Use: CMD_INJECT,
	}
	injectCmd.AddCommand(globalInjectCmd)

	mrtCmd := &cobra.Command{
		Use: CMD_MRT,
	}
	mrtCmd.AddCommand(injectCmd)

	mrtCmd.PersistentFlags().BoolVarP(&mrtOpts.Best, "only-best", "", false, "inject only best paths")
	return mrtCmd
}
<|endoftext|>


<commit_before>package gorush

import (
	"sync"
)

// InitWorkers for initialize all workers.
func InitWorkers(workerNum int64, queueNum int64) {
	LogAccess.Debug("worker number is ", workerNum, ", queue number is ", queueNum)
	QueueNotification = make(chan PushNotification, queueNum)
	for i := int64(0); i < workerNum; i++ {
		go startWorker()
	}
}

// SendNotification is send message to iOS or Android
func SendNotification(msg PushNotification) {
	switch msg.Platform {
	case PlatFormIos:
		PushToIOS(msg)
	case PlatFormAndroid:
		PushToAndroid(msg)
	}
}

func startWorker() {
	for {
		notification := <-QueueNotification
		SendNotification(notification)
	}
}

// queueNotification add notification to queue list.
func queueNotification(req RequestPush) (int, []LogPushEntry) {
	var count int
	wg := sync.WaitGroup{}
	newNotification := []*PushNotification{}
	for i := range req.Notifications {
		notification := &req.Notifications[i]
		switch notification.Platform {
		case PlatFormIos:
			if !PushConf.Ios.Enabled {
				continue
			}
		case PlatFormAndroid:
			if !PushConf.Android.Enabled {
				continue
			}
		}
		newNotification = append(newNotification, notification)
	}

	log := make([]LogPushEntry, 0, count)
	for _, notification := range newNotification {
		if PushConf.Core.Sync {
			notification.wg = &wg
			notification.log = &log
			notification.AddWaitCount()
		}
		if !tryEnqueue(*notification, QueueNotification) {
			LogError.Error("max capacity reached")
		}
		count += len(notification.Tokens)
		// Count topic message
		if notification.To != "" {
			count++
		}
	}

	if PushConf.Core.Sync {
		wg.Wait()
	}

	StatStorage.AddTotalCount(int64(count))

	return count, log
}

// tryEnqueue tries to enqueue a job to the given job channel. Returns true if
// the operation was successful, and false if enqueuing would not have been
// possible without blocking. Job is not enqueued in the latter case.
func tryEnqueue(job PushNotification, jobChan chan<- PushNotification) bool {
	select {
	case jobChan <- job:
		return true
	default:
		return false
	}
}
<commit_msg>When queue is full, requests hang because Wait.Done() is not cal… (#430)<commit_after>package gorush

import (
	"errors"
	"sync"
)

// InitWorkers for initialize all workers.
func InitWorkers(workerNum int64, queueNum int64) {
	LogAccess.Debug("worker number is ", workerNum, ", queue number is ", queueNum)
	QueueNotification = make(chan PushNotification, queueNum)
	for i := int64(0); i < workerNum; i++ {
		go startWorker()
	}
}

// SendNotification is send message to iOS or Android
func SendNotification(msg PushNotification) {
	switch msg.Platform {
	case PlatFormIos:
		PushToIOS(msg)
	case PlatFormAndroid:
		PushToAndroid(msg)
	}
}

func startWorker() {
	for {
		notification := <-QueueNotification
		SendNotification(notification)
	}
}

// markFailedNotification adds failure logs for all tokens in push notification
func markFailedNotification(notification *PushNotification, reason string) {
	LogError.Error(reason)
	for _, token := range notification.Tokens {
		notification.AddLog(getLogPushEntry(FailedPush, token, *notification, errors.New(reason)))
	}
	notification.WaitDone()
}

// queueNotification add notification to queue list.
func queueNotification(req RequestPush) (int, []LogPushEntry) {
	var count int
	wg := sync.WaitGroup{}
	newNotification := []*PushNotification{}
	for i := range req.Notifications {
		notification := &req.Notifications[i]
		switch notification.Platform {
		case PlatFormIos:
			if !PushConf.Ios.Enabled {
				continue
			}
		case PlatFormAndroid:
			if !PushConf.Android.Enabled {
				continue
			}
		}
		newNotification = append(newNotification, notification)
	}

	log := make([]LogPushEntry, 0, count)
	for _, notification := range newNotification {
		if PushConf.Core.Sync {
			notification.wg = &wg
			notification.log = &log
			notification.AddWaitCount()
		}
		if !tryEnqueue(*notification, QueueNotification) {
			markFailedNotification(notification, "max capacity reached")
		}
		count += len(notification.Tokens)
		// Count topic message
		if notification.To != "" {
			count++
		}
	}

	if PushConf.Core.Sync {
		wg.Wait()
	}

	StatStorage.AddTotalCount(int64(count))

	return count, log
}

// tryEnqueue tries to enqueue a job to the given job channel. Returns true if
// the operation was successful, and false if enqueuing would not have been
// possible without blocking. Job is not enqueued in the latter case.
func tryEnqueue(job PushNotification, jobChan chan<- PushNotification) bool {
	select {
	case jobChan <- job:
		return true
	default:
		return false
	}
}
<|endoftext|>


<commit_before>package langdetect

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"sort"
	"strings"
	"testing"
)

func TestWikipedia(t *testing.T) {
	f, e := os.Open("test_data/wikipedia")
	if e != nil {
		panic(e)
	}
	ns, e := f.Readdirnames(-1)
	if e != nil {
		panic(e)
	}
	for _, n := range ns {
		bs, e := ioutil.ReadFile("test_data/wikipedia/" + n)
		l := n[5:8]
		if l[2] == '.' {
			l = l[:2]
		}
		if e != nil {
			panic(e)
		}
		lang := DetectLanguage(bs, "")
		var alert string
		if l != lang.String() {
			alert = "WRONG ("
			if len(l) == 2 {
				alert += Language([2]byte{l[0], l[1]}).EnglishName()
			}
			alert += ")"
		}
		if lang == Unknown {
			alert = "UNKNOWN ("
			if len(l) == 2 {
				alert += Language([2]byte{l[0], l[1]}).EnglishName()
			}
			alert += ")"
		}
		fmt.Printf("%20s => %v %s %s\n", n, lang, lang.EnglishName(), alert)
	}
}

func XTestCalcTrigrams(t *testing.T) {
	f, e := os.Open("test_data/calc_trigrams")
	if e != nil {
		panic(e)
	}
	ns, e := f.Readdirnames(-1)
	if e != nil {
		panic(e)
	}
	for _, n := range ns {
		bs, e := ioutil.ReadFile("test_data/calc_trigrams/" + n)
		fmt.Println("BUILDING TRIGRAM TABLE: " + n)
		if e != nil {
			panic(e)
		}
		vs := calcLatinByteTrigram(bs)
		lim := 500
		if lim > len(vs) {
			lim = len(vs)
		}
		vs = vs[0:lim]
		x := trigramfreqs(vs)
		sort.Sort(x)
		buf := new(bytes.Buffer)
		buf.WriteString("package langdetect\n\nvar trigram_" + n + " = trigramfreqs{\n")
		for i := 0; i < lim; i++ {
			buf.WriteString(strings.Replace(fmt.Sprintf("%#v,\n", x[i]), "langdetect.", "", -1))
		}
		buf.WriteString("}\n\n")
		ioutil.WriteFile("autogenerated_trigrams_"+n+".go", buf.Bytes(), 0600)
	}
}

func BenchmarkDiff(b *testing.B) {
	for i := 0; i < b.N; i++ {
		trigram_en.Diff(trigram_fr)
	}
}

//func BenchmarkDiffOld(b *testing.B) {
//	for i := 0; i < b.N; i++ {
//		trigram_en.DiffOld(trigram_fr)
//	}
//}

func BenchmarkDetectEn(b *testing.B) {
	bs, e := ioutil.ReadFile("test_data/wikipedia/auto_en.txt")
	if e != nil {
		panic(e)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		DetectLanguage(bs, "")
	}
}

func BenchmarkDetectZh(b *testing.B) {
	bs, e := ioutil.ReadFile("test_data/wikipedia/auto_zh.txt")
	if e != nil {
		panic(e)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		DetectLanguage(bs, "")
	}
}
<commit_msg>Make tests work without test_data<commit_after>package langdetect

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"sort"
	"strings"
	"testing"
)

func TestWikipedia(t *testing.T) {
	f, e := os.Open("test_data/wikipedia")
	if e != nil {
		fmt.Println("NO test_data PRESENT EXITING")
		return
	}
	ns, e := f.Readdirnames(-1)
	if e != nil {
		panic(e)
	}
	for _, n := range ns {
		bs, e := ioutil.ReadFile("test_data/wikipedia/" + n)
		l := n[5:8]
		if l[2] == '.' {
			l = l[:2]
		}
		if e != nil {
			panic(e)
		}
		lang := DetectLanguage(bs, "")
		var alert string
		if l != lang.String() {
			alert = "WRONG ("
			if len(l) == 2 {
				alert += Language([2]byte{l[0], l[1]}).EnglishName()
			}
			alert += ")"
		}
		if lang == Unknown {
			alert = "UNKNOWN ("
			if len(l) == 2 {
				alert += Language([2]byte{l[0], l[1]}).EnglishName()
			}
			alert += ")"
		}
		fmt.Printf("%20s => %v %s %s\n", n, lang, lang.EnglishName(), alert)
	}
}

func XTestCalcTrigrams(t *testing.T) {
	f, e := os.Open("test_data/calc_trigrams")
	if e != nil {
		panic(e)
	}
	ns, e := f.Readdirnames(-1)
	if e != nil {
		panic(e)
	}
	for _, n := range ns {
		bs, e := ioutil.ReadFile("test_data/calc_trigrams/" + n)
		fmt.Println("BUILDING TRIGRAM TABLE: " + n)
		if e != nil {
			panic(e)
		}
		vs := calcLatinByteTrigram(bs)
		lim := 500
		if lim > len(vs) {
			lim = len(vs)
		}
		vs = vs[0:lim]
		x := trigramfreqs(vs)
		sort.Sort(x)
		buf := new(bytes.Buffer)
		buf.WriteString("package langdetect\n\nvar trigram_" + n + " = trigramfreqs{\n")
		for i := 0; i < lim; i++ {
			buf.WriteString(strings.Replace(fmt.Sprintf("%#v,\n", x[i]), "langdetect.", "", -1))
		}
		buf.WriteString("}\n\n")
		ioutil.WriteFile("autogenerated_trigrams_"+n+".go", buf.Bytes(), 0600)
	}
}

func BenchmarkDiff(b *testing.B) {
	for i := 0; i < b.N; i++ {
		trigram_en.Diff(trigram_fr)
	}
}

//func BenchmarkDiffOld(b *testing.B) {
//	for i := 0; i < b.N; i++ {
//		trigram_en.DiffOld(trigram_fr)
//	}
//}

func BenchmarkDetectEn(b *testing.B) {
	bs, e := ioutil.ReadFile("test_data/wikipedia/auto_en.txt")
	b.ResetTimer()
	if e != nil {
		fmt.Println("NO test_data PRESENT EXITING")
		return
	}
	for i := 0; i < b.N; i++ {
		DetectLanguage(bs, "")
	}
}

func BenchmarkDetectZh(b *testing.B) {
	bs, e := ioutil.ReadFile("test_data/wikipedia/auto_zh.txt")
	if e != nil {
		fmt.Println("NO test_data PRESENT EXITING")
		return
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		DetectLanguage(bs, "")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"

	"io/ioutil"

	"fmt"

	"bytes"
)

func mapKeyPairs(input [][]byte, buffer []string) []map[string][]byte {
	tokenMapS := []map[string][]byte{}
	for _, pv := range input {
		tokenMap := map[string]*json.RawMessage{}

		err := json.Unmarshal(pv, &tokenMap)
		if err != nil {
			// TODO: find a better solution for this block.
			// spew.Dump(pv, tokenMap)
			// checkError(err)
		}

		tempMap := map[string][]byte{}

		for k, v := range tokenMap {
			j, err := json.Marshal(&v)
			checkError(err)
			tempMap[buffer[0]+k+buffer[1]] = j
		}
		tokenMapS = append(tokenMapS, tempMap)
	}
	return tokenMapS
}

func detokenize(input []byte, tokenMap []map[string][]byte) []byte {
	overrideCompiled := map[string][]byte{}
	for _, v := range tokenMap {
		for mk, mv := range v {
			overrideCompiled[mk] = mv
		}
	}
	for k, v := range overrideCompiled {
		input = bytes.Replace(input, []byte(k), v, -1)
	}
	return input
}

func loadFile(path string) []byte {
	file, err := ioutil.ReadFile(path)
	checkError(err)
	return file
}

func outputToFile(path string, data []byte) {
	err := ioutil.WriteFile(path, data, 0644)
	checkError(err)
}

func outputToStdout(data []byte) {
	fmt.Print(string(data))
}
<commit_msg>better detoken<commit_after>package main

import (
	"encoding/json"

	"io/ioutil"

	"fmt"

	"bytes"
)

func mapKeyPairs(input [][]byte, buffer []string) []map[string][]byte {
	tokenMapS := []map[string][]byte{}
	for _, pv := range input {
		tokenMap := map[string]*json.RawMessage{}

		err := json.Unmarshal(pv, &tokenMap)
		if err != nil {
			// TODO: find a better solution for this block.
			// spew.Dump(pv, tokenMap)
			// checkError(err)
		}

		tempMap := map[string][]byte{}

		for k, v := range tokenMap {
			j, err := json.Marshal(&v)
			checkError(err)
			tempMap[buffer[0]+k+buffer[1]] = j
		}
		tokenMapS = append(tokenMapS, tempMap)
	}
	return tokenMapS
}

func detokenize(input []byte, tokenMap []map[string][]byte) []byte {
	overrideCompiled := map[string][]byte{}
	for _, v := range tokenMap {
		for mk, mv := range v {
			overrideCompiled[mk] = mv
		}
	}
	for k, v := range overrideCompiled {
		input = bytes.Replace(input, []byte(k), v, -1)
        input = bytes.Replace(input, []byte("\\u003c"), []byte("<"), -1)
        input = bytes.Replace(input, []byte("\\u003e"), []byte(">"), -1)
        input = bytes.Replace(input, []byte("\\u0026"), []byte("&"), -1)

    }
	}
	return input
}

func loadFile(path string) []byte {
	file, err := ioutil.ReadFile(path)
	checkError(err)
	return file
}

func outputToFile(path string, data []byte) {
	err := ioutil.WriteFile(path, data, 0644)
	checkError(err)
}

func outputToStdout(data []byte) {
	fmt.Print(string(data))
}
<|endoftext|>


<commit_before>/*
Package gothic wraps common behaviour when using Goth. This makes it quick, and easy, to get up
and running with Goth. Of course, if you want complete control over how things flow, in regards
to the authentication process, feel free and use Goth directly.

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
package gothic

import (
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"os"

	"github.com/gorilla/mux"
	"github.com/gorilla/sessions"
	"github.com/markbates/goth"
)

// SessionName is the key used to access the session store.
const SessionName = "_gothic_session"

// Store can/should be set by applications using gothic. The default is a cookie store.
var Store sessions.Store
var defaultStore sessions.Store

var keySet = false

func init() {
	key := []byte(os.Getenv("SESSION_SECRET"))
	keySet = len(key) != 0

	cookieStore := sessions.NewCookieStore([]byte(key))
	cookieStore.Options.HttpOnly = true
	Store = cookieStore
	defaultStore = Store
}

/*
BeginAuthHandler is a convienence handler for starting the authentication process.
It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

BeginAuthHandler will redirect the user to the appropriate authentication end-point
for the requested provider.

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
func BeginAuthHandler(res http.ResponseWriter, req *http.Request) {
	url, err := GetAuthURL(res, req)
	if err != nil {
		res.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(res, err)
		return
	}

	http.Redirect(res, req, url, http.StatusTemporaryRedirect)
}

// SetState sets the state string associated with the given request.
// If no state string is associated with the request, one will be generated.
// This state is sent to the provider and can be retrieved during the
// callback.
var SetState = func(req *http.Request) string {
	state := req.URL.Query().Get("state")
	if len(state) > 0 {
		return state
	}

	return "state"

}

// GetState gets the state returned by the provider during the callback.
// This is used to prevent CSRF attacks, see
// http://tools.ietf.org/html/rfc6749#section-10.12
var GetState = func(req *http.Request) string {
	return req.URL.Query().Get("state")
}

/*
GetAuthURL starts the authentication process with the requested provided.
It will return a URL that should be used to send users to.

It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

I would recommend using the BeginAuthHandler instead of doing all of these steps
yourself, but that's entirely up to you.
*/
func GetAuthURL(res http.ResponseWriter, req *http.Request) (string, error) {
	if !keySet && defaultStore == Store {
		fmt.Println("goth/gothic: no SESSION_SECRET environment variable is set. The default cookie store is not available and any calls will fail. Ignore this warning if you are using a different store.")
	}

	providerName, err := GetProviderName(req)
	if err != nil {
		return "", err
	}

	provider, err := goth.GetProvider(providerName)
	if err != nil {
		return "", err
	}
	sess, err := provider.BeginAuth(SetState(req))
	if err != nil {
		return "", err
	}

	url, err := sess.GetAuthURL()
	if err != nil {
		return "", err
	}

	err = storeInSession(providerName, sess.Marshal(), req, res)

	if err != nil {
		return "", err
	}

	return url, err
}

/*
CompleteUserAuth does what it says on the tin. It completes the authentication
process and fetches all of the basic information about the user from the provider.

It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
var CompleteUserAuth = func(res http.ResponseWriter, req *http.Request) (goth.User, error) {
	if !keySet && defaultStore == Store {
		fmt.Println("goth/gothic: no SESSION_SECRET environment variable is set. The default cookie store is not available and any calls will fail. Ignore this warning if you are using a different store.")
	}

	providerName, err := GetProviderName(req)
	if err != nil {
		return goth.User{}, err
	}

	provider, err := goth.GetProvider(providerName)
	if err != nil {
		return goth.User{}, err
	}

	value, err := getFromSession(providerName, req)
	if err != nil {
		return goth.User{}, err
	}

	sess, err := provider.UnmarshalSession(value)
	if err != nil {
		return goth.User{}, err
	}

	err = validateState(req, sess)
	if err != nil {
		return goth.User{}, err
	}

	user, err := provider.FetchUser(sess)
	if err == nil {
		// user can be found with existing session data
		return user, err
	}

	// get new token and retry fetch
	_, err = sess.Authorize(provider, req.URL.Query())
	if err != nil {
		return goth.User{}, err
	}

	err = storeInSession(providerName, sess.Marshal(), req, res)

	if err != nil {
		return goth.User{}, err
	}

	return provider.FetchUser(sess)
}

// validateState ensures that the state token param from the original
// AuthURL matches the one included in the current (callback) request.
func validateState(req *http.Request, sess goth.Session) error {
	rawAuthURL, err := sess.GetAuthURL()
	if err != nil {
		return err
	}

	authURL, err := url.Parse(rawAuthURL)
	if err != nil {
		return err
	}

	originalState := authURL.Query().Get("state")
	if originalState != "" && (originalState != req.URL.Query().Get("state")) {
		return errors.New("state token mismatch")
	}
	return nil
}

// Logout invalidates a user session.
func Logout(res http.ResponseWriter, req *http.Request) error {
	providerName, err := GetProviderName(req)
	if err != nil {
		return err
	}

	session, err := Store.Get(req, providerName+SessionName)
	if err != nil {
		return err
	}
	session.Options.MaxAge = -1
	session.Values = make(map[interface{}]interface{})
	err = session.Save(req, res)
	if err != nil {
		return errors.New("Could not delete user session ")
	}
	return nil
}

// GetProviderName is a function used to get the name of a provider
// for a given request. By default, this provider is fetched from
// the URL query string. If you provide it in a different way,
// assign your own function to this variable that returns the provider
// name for your request.
var GetProviderName = getProviderName

func getProviderName(req *http.Request) (string, error) {

	// get all the used providers
	providers := goth.GetProviders()

	// loop over the used providers, if we already have a valid session for any provider (ie. user is already logged-in with a provider), then return that provider name
	for _, provider := range providers {
		p := provider.Name()
		session, _ := Store.Get(req, p+SessionName)
		value := session.Values[p]
		if _, ok := value.(string); ok {
			return p, nil
		}
	}

	// try to get it from the url param "provider"
	if p := req.URL.Query().Get("provider"); p != "" {
		return p, nil
	}

	// try to get it from the url param ":provider"
	if p := req.URL.Query().Get(":provider"); p != "" {
		return p, nil
	}

	// try to get it from the context's value of "provider" key
	if p, ok := mux.Vars(req)["provider"]; ok {
		return p, nil
	}

	//  try to get it from the go-context's value of "provider" key
	if p, ok := req.Context().Value("provider").(string); ok {
		return p, nil
	}

	// if not found then return an empty string with the corresponding error
	return "", errors.New("you must select a provider")
}

func storeInSession(key string, value string, req *http.Request, res http.ResponseWriter) error {
	session, _ := Store.Get(req, key+SessionName)

	session.Values[key] = value

	return session.Save(req, res)
}

func getFromSession(key string, req *http.Request) (string, error) {
	session, _ := Store.Get(req, key+SessionName)
	value := session.Values[key]
	if value == nil {
		return "", errors.New("could not find a matching session for this request")
	}

	return value.(string), nil
}
<commit_msg>fixed issue with gothic not handling multiple logins correctly<commit_after>/*
Package gothic wraps common behaviour when using Goth. This makes it quick, and easy, to get up
and running with Goth. Of course, if you want complete control over how things flow, in regards
to the authentication process, feel free and use Goth directly.

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
package gothic

import (
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"os"

	"github.com/gorilla/mux"
	"github.com/gorilla/sessions"
	"github.com/markbates/goth"
)

// SessionName is the key used to access the session store.
const SessionName = "_gothic_session"

// Store can/should be set by applications using gothic. The default is a cookie store.
var Store sessions.Store
var defaultStore sessions.Store

var keySet = false

func init() {
	key := []byte(os.Getenv("SESSION_SECRET"))
	keySet = len(key) != 0

	cookieStore := sessions.NewCookieStore([]byte(key))
	cookieStore.Options.HttpOnly = true
	Store = cookieStore
	defaultStore = Store
}

/*
BeginAuthHandler is a convienence handler for starting the authentication process.
It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

BeginAuthHandler will redirect the user to the appropriate authentication end-point
for the requested provider.

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
func BeginAuthHandler(res http.ResponseWriter, req *http.Request) {
	url, err := GetAuthURL(res, req)
	if err != nil {
		res.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(res, err)
		return
	}

	http.Redirect(res, req, url, http.StatusTemporaryRedirect)
}

// SetState sets the state string associated with the given request.
// If no state string is associated with the request, one will be generated.
// This state is sent to the provider and can be retrieved during the
// callback.
var SetState = func(req *http.Request) string {
	state := req.URL.Query().Get("state")
	if len(state) > 0 {
		return state
	}

	return "state"

}

// GetState gets the state returned by the provider during the callback.
// This is used to prevent CSRF attacks, see
// http://tools.ietf.org/html/rfc6749#section-10.12
var GetState = func(req *http.Request) string {
	return req.URL.Query().Get("state")
}

/*
GetAuthURL starts the authentication process with the requested provided.
It will return a URL that should be used to send users to.

It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

I would recommend using the BeginAuthHandler instead of doing all of these steps
yourself, but that's entirely up to you.
*/
func GetAuthURL(res http.ResponseWriter, req *http.Request) (string, error) {
	if !keySet && defaultStore == Store {
		fmt.Println("goth/gothic: no SESSION_SECRET environment variable is set. The default cookie store is not available and any calls will fail. Ignore this warning if you are using a different store.")
	}

	providerName, err := GetProviderName(req)
	if err != nil {
		return "", err
	}

	provider, err := goth.GetProvider(providerName)
	if err != nil {
		return "", err
	}
	sess, err := provider.BeginAuth(SetState(req))
	if err != nil {
		return "", err
	}

	url, err := sess.GetAuthURL()
	if err != nil {
		return "", err
	}

	err = storeInSession(providerName, sess.Marshal(), req, res)

	if err != nil {
		return "", err
	}

	return url, err
}

/*
CompleteUserAuth does what it says on the tin. It completes the authentication
process and fetches all of the basic information about the user from the provider.

It expects to be able to get the name of the provider from the query parameters
as either "provider" or ":provider".

See https://github.com/markbates/goth/examples/main.go to see this in action.
*/
var CompleteUserAuth = func(res http.ResponseWriter, req *http.Request) (goth.User, error) {
	defer Logout(res, req)

	if !keySet && defaultStore == Store {
		fmt.Println("goth/gothic: no SESSION_SECRET environment variable is set. The default cookie store is not available and any calls will fail. Ignore this warning if you are using a different store.")
	}

	providerName, err := GetProviderName(req)
	if err != nil {
		return goth.User{}, err
	}

	provider, err := goth.GetProvider(providerName)
	if err != nil {
		return goth.User{}, err
	}

	value, err := getFromSession(providerName, req)
	if err != nil {
		return goth.User{}, err
	}

	sess, err := provider.UnmarshalSession(value)
	if err != nil {
		return goth.User{}, err
	}

	err = validateState(req, sess)
	if err != nil {
		return goth.User{}, err
	}

	user, err := provider.FetchUser(sess)
	if err == nil {
		// user can be found with existing session data
		return user, err
	}

	// get new token and retry fetch
	_, err = sess.Authorize(provider, req.URL.Query())
	if err != nil {
		return goth.User{}, err
	}

	err = storeInSession(providerName, sess.Marshal(), req, res)

	if err != nil {
		return goth.User{}, err
	}

	gu, err := provider.FetchUser(sess)
	return gu, err
}

// validateState ensures that the state token param from the original
// AuthURL matches the one included in the current (callback) request.
func validateState(req *http.Request, sess goth.Session) error {
	rawAuthURL, err := sess.GetAuthURL()
	if err != nil {
		return err
	}

	authURL, err := url.Parse(rawAuthURL)
	if err != nil {
		return err
	}

	originalState := authURL.Query().Get("state")
	if originalState != "" && (originalState != req.URL.Query().Get("state")) {
		return errors.New("state token mismatch")
	}
	return nil
}

// Logout invalidates a user session.
func Logout(res http.ResponseWriter, req *http.Request) error {
	providerName, err := GetProviderName(req)
	if err != nil {
		return err
	}

	session, err := Store.Get(req, providerName+SessionName)
	if err != nil {
		return err
	}
	session.Options.MaxAge = -1
	session.Values = make(map[interface{}]interface{})
	err = session.Save(req, res)
	if err != nil {
		return errors.New("Could not delete user session ")
	}
	return nil
}

// GetProviderName is a function used to get the name of a provider
// for a given request. By default, this provider is fetched from
// the URL query string. If you provide it in a different way,
// assign your own function to this variable that returns the provider
// name for your request.
var GetProviderName = getProviderName

func getProviderName(req *http.Request) (string, error) {

	// get all the used providers
	providers := goth.GetProviders()

	// loop over the used providers, if we already have a valid session for any provider (ie. user is already logged-in with a provider), then return that provider name
	for _, provider := range providers {
		p := provider.Name()
		session, _ := Store.Get(req, p+SessionName)
		value := session.Values[p]
		if _, ok := value.(string); ok {
			return p, nil
		}
	}

	// try to get it from the url param "provider"
	if p := req.URL.Query().Get("provider"); p != "" {
		return p, nil
	}

	// try to get it from the url param ":provider"
	if p := req.URL.Query().Get(":provider"); p != "" {
		return p, nil
	}

	// try to get it from the context's value of "provider" key
	if p, ok := mux.Vars(req)["provider"]; ok {
		return p, nil
	}

	//  try to get it from the go-context's value of "provider" key
	if p, ok := req.Context().Value("provider").(string); ok {
		return p, nil
	}

	// if not found then return an empty string with the corresponding error
	return "", errors.New("you must select a provider")
}

func storeInSession(key string, value string, req *http.Request, res http.ResponseWriter) error {
	session, _ := Store.Get(req, key+SessionName)

	session.Values[key] = value

	return session.Save(req, res)
}

func getFromSession(key string, req *http.Request) (string, error) {
	session, _ := Store.Get(req, key+SessionName)
	value := session.Values[key]
	if value == nil {
		return "", errors.New("could not find a matching session for this request")
	}

	return value.(string), nil
}
<|endoftext|>


<commit_before>package mount

import (
	"encoding/json"
	"reflect"
	"sort"
	"testing"
)

func TestMakeIDString(t *testing.T) {
	tests := map[string]struct {
		Tok     string
		ID      ID
		IsValid bool
	}{
		"valid UUIDv4": {
			Tok:     "e2d92c39-bcc6-45c8-86f6-7ca114ec7e06",
			ID:      ID("e2d92c39-bcc6-45c8-86f6-7ca114ec7e06"),
			IsValid: true,
		},
		"valid with braces": {
			Tok:     "52560e13-ad72-4cef-a3d7-f52ca197168d",
			ID:      ID("52560e13-ad72-4cef-a3d7-f52ca197168d"),
			IsValid: true,
		},
		"invalid UUIDv1": {
			Tok:     "123e4567-e89b-12d3-a456-426655440000",
			ID:      ID(""),
			IsValid: false,
		},
		"invalid": {
			Tok:     "invalid",
			ID:      ID(""),
			IsValid: false,
		},
		"empty token": {
			Tok:     "",
			ID:      ID(""),
			IsValid: false,
		},
	}

	for name, test := range tests {
		t.Run(name, func(t *testing.T) {
			mountID, err := IDFromString(test.Tok)
			if (err == nil) != test.IsValid {
				t.Fatalf("want err == nil to be %t; got %v", test.IsValid, err)
			}

			if mountID != test.ID {
				t.Errorf("want mount ID = %s; got %s", test.ID, mountID)
			}
		})
	}
}

func TestMountBook(t *testing.T) {
	var (
		idA = MakeID()
		idB = MakeID()
		idC = MakeID()
	)

	mounts := map[ID]Mount{
		idA: {
			Path:       "/home/koding/a",
			RemotePath: "/home/koding/remote/a",
		},
		idB: {
			Path:       "/home/koding/b",
			RemotePath: "/home/koding/remote/shared",
		},
		idC: {
			Path:       "/home/koding/c",
			RemotePath: "/home/koding/remote/shared",
		},
	}

	tests := map[string]struct {
		Path          string
		PathID        ID
		RemotePath    string
		RemotePathIDs IDSlice
	}{
		"paths from mount A": {
			Path:          "/home/koding/a",
			PathID:        idA,
			RemotePath:    "/home/koding/remote/a",
			RemotePathIDs: IDSlice{idA},
		},
		"shared remote path": {
			Path:          "/home/koding/b",
			PathID:        idB,
			RemotePath:    "/home/koding/remote/shared",
			RemotePathIDs: IDSlice{idB, idC},
		},
		"unknown paths": {
			Path:          "/home/koding/unknown",
			PathID:        "",
			RemotePath:    "/home/koding/remote/unknown",
			RemotePathIDs: nil,
		},
	}

	mb := NewMountBook()
	for i, id := range []ID{idA, idB, idC} {
		if err := mb.Add(id, mounts[id]); err != nil {
			t.Fatalf("want err = nil; got %v (i:%d)", err, i)
		}
	}

	for name, test := range tests {
		t.Run(name, func(t *testing.T) {
			sort.Sort(test.RemotePathIDs)
			id, err := mb.Path(test.Path)
			if test.PathID == "" && err == nil {
				t.Errorf("want err != nil, got nil")
			}
			if test.PathID != "" && err != nil {
				t.Errorf("want ID = %s, got err = %v", test.PathID, err)
			}
			if test.PathID != id {
				t.Errorf("want ID = %s, got %s", test.PathID, id)
			}

			ids, err := mb.RemotePath(test.RemotePath)
			sort.Sort(ids)
			if test.RemotePathIDs == nil && err == nil {
				t.Errorf("want err != nil, got nil")
			}
			if test.RemotePathIDs != nil && err != nil {
				t.Errorf("want IDs = %v, got err = %v", test.RemotePathIDs, err)
			}
			if !reflect.DeepEqual(test.RemotePathIDs, ids) {
				t.Errorf("want IDs = %v, got %v", test.RemotePathIDs, ids)
			}
		})
	}

}

func TestMountBookDuplicatedID(t *testing.T) {
	mb := NewMountBook()
	if err := mb.Add("duplicated", Mount{}); err != nil {
		t.Fatalf("want err = nil; got %v", err)
	}

	if err := mb.Add("duplicated", Mount{}); err == nil {
		t.Fatalf("want err != nil; got nil")
	}
}

func TestMountBookJSON(t *testing.T) {
	mounts := map[ID]Mount{
		MakeID(): {
			Path:       "/home/koding/a",
			RemotePath: "/home/koding/remote/a",
		},
		MakeID(): {
			Path:       "/home/koding/b",
			RemotePath: "/home/koding/remote/b",
		},
	}

	mb := NewMountBook()
	for id, mount := range mounts {
		if err := mb.Add(id, mount); err != nil {
			t.Fatalf("want err = nil; got %v", err)
		}
	}

	data, err := json.Marshal(mb)
	if err != nil {
		t.Fatalf("want err = nil; got %v", err)
	}

	mb = NewMountBook()
	if err = json.Unmarshal(data, mb); err != nil {
		t.Fatalf("want err = nil; got %v", err)
	}

	all := mb.All()
	if !reflect.DeepEqual(all, mounts) {
		t.Fatalf("want mount map = %#v; got %#v", mounts, all)
	}
}
<commit_msg>klient/machine: parallel MakeIDString function tests<commit_after>package mount

import (
	"encoding/json"
	"reflect"
	"sort"
	"testing"
)

func TestMakeIDString(t *testing.T) {
	tests := map[string]struct {
		Tok     string
		ID      ID
		IsValid bool
	}{
		"valid UUIDv4": {
			Tok:     "e2d92c39-bcc6-45c8-86f6-7ca114ec7e06",
			ID:      ID("e2d92c39-bcc6-45c8-86f6-7ca114ec7e06"),
			IsValid: true,
		},
		"valid with braces": {
			Tok:     "52560e13-ad72-4cef-a3d7-f52ca197168d",
			ID:      ID("52560e13-ad72-4cef-a3d7-f52ca197168d"),
			IsValid: true,
		},
		"invalid UUIDv1": {
			Tok:     "123e4567-e89b-12d3-a456-426655440000",
			ID:      ID(""),
			IsValid: false,
		},
		"invalid": {
			Tok:     "invalid",
			ID:      ID(""),
			IsValid: false,
		},
		"empty token": {
			Tok:     "",
			ID:      ID(""),
			IsValid: false,
		},
	}

	for name, test := range tests {
		test := test // Capture range variable.
		t.Run(name, func(t *testing.T) {
			t.Parallel()

			mountID, err := IDFromString(test.Tok)
			if (err == nil) != test.IsValid {
				t.Fatalf("want err == nil to be %t; got %v", test.IsValid, err)
			}

			if mountID != test.ID {
				t.Errorf("want mount ID = %s; got %s", test.ID, mountID)
			}
		})
	}
}

func TestMountBook(t *testing.T) {
	var (
		idA = MakeID()
		idB = MakeID()
		idC = MakeID()
	)

	mounts := map[ID]Mount{
		idA: {
			Path:       "/home/koding/a",
			RemotePath: "/home/koding/remote/a",
		},
		idB: {
			Path:       "/home/koding/b",
			RemotePath: "/home/koding/remote/shared",
		},
		idC: {
			Path:       "/home/koding/c",
			RemotePath: "/home/koding/remote/shared",
		},
	}

	tests := map[string]struct {
		Path          string
		PathID        ID
		RemotePath    string
		RemotePathIDs IDSlice
	}{
		"paths from mount A": {
			Path:          "/home/koding/a",
			PathID:        idA,
			RemotePath:    "/home/koding/remote/a",
			RemotePathIDs: IDSlice{idA},
		},
		"shared remote path": {
			Path:          "/home/koding/b",
			PathID:        idB,
			RemotePath:    "/home/koding/remote/shared",
			RemotePathIDs: IDSlice{idB, idC},
		},
		"unknown paths": {
			Path:          "/home/koding/unknown",
			PathID:        "",
			RemotePath:    "/home/koding/remote/unknown",
			RemotePathIDs: nil,
		},
	}

	mb := NewMountBook()
	for i, id := range []ID{idA, idB, idC} {
		if err := mb.Add(id, mounts[id]); err != nil {
			t.Fatalf("want err = nil; got %v (i:%d)", err, i)
		}
	}

	for name, test := range tests {
		t.Run(name, func(t *testing.T) {
			sort.Sort(test.RemotePathIDs)
			id, err := mb.Path(test.Path)
			if test.PathID == "" && err == nil {
				t.Errorf("want err != nil, got nil")
			}
			if test.PathID != "" && err != nil {
				t.Errorf("want ID = %s, got err = %v", test.PathID, err)
			}
			if test.PathID != id {
				t.Errorf("want ID = %s, got %s", test.PathID, id)
			}

			ids, err := mb.RemotePath(test.RemotePath)
			sort.Sort(ids)
			if test.RemotePathIDs == nil && err == nil {
				t.Errorf("want err != nil, got nil")
			}
			if test.RemotePathIDs != nil && err != nil {
				t.Errorf("want IDs = %v, got err = %v", test.RemotePathIDs, err)
			}
			if !reflect.DeepEqual(test.RemotePathIDs, ids) {
				t.Errorf("want IDs = %v, got %v", test.RemotePathIDs, ids)
			}
		})
	}

}

func TestMountBookDuplicatedID(t *testing.T) {
	mb := NewMountBook()
	if err := mb.Add("duplicated", Mount{}); err != nil {
		t.Fatalf("want err = nil; got %v", err)
	}

	if err := mb.Add("duplicated", Mount{}); err == nil {
		t.Fatalf("want err != nil; got nil")
	}
}

func TestMountBookJSON(t *testing.T) {
	mounts := map[ID]Mount{
		MakeID(): {
			Path:       "/home/koding/a",
			RemotePath: "/home/koding/remote/a",
		},
		MakeID(): {
			Path:       "/home/koding/b",
			RemotePath: "/home/koding/remote/b",
		},
	}

	mb := NewMountBook()
	for id, mount := range mounts {
		if err := mb.Add(id, mount); err != nil {
			t.Fatalf("want err = nil; got %v", err)
		}
	}

	data, err := json.Marshal(mb)
	if err != nil {
		t.Fatalf("want err = nil; got %v", err)
	}

	mb = NewMountBook()
	if err = json.Unmarshal(data, mb); err != nil {
		t.Fatalf("want err = nil; got %v", err)
	}

	all := mb.All()
	if !reflect.DeepEqual(all, mounts) {
		t.Fatalf("want mount map = %#v; got %#v", mounts, all)
	}
}
<|endoftext|>


<commit_before>package runner

import (
	"errors"
	"flag"
	"fmt"
	"net/url"
	"os"
	"os/signal"
	"socialapi/config"
	"socialapi/workers/helper"
	"strconv"
	"syscall"
	"github.com/koding/kite"
	kiteconfig "github.com/koding/kite/config"
	"github.com/koding/kite/protocol"

	"github.com/koding/bongo"
	"github.com/koding/logging"
	"github.com/koding/worker"
)

var (
	flagConfFile = flag.String("c", "", "Configuration profile from file")
	flagRegion   = flag.String("r", "", "Region name")
	flagDebug    = flag.Bool("d", false, "Debug mode")
	flagVersion  = flag.Int("v", 0, "Worker Version")

	flagKiteInit       = flag.Bool("kite-init", false, "Init kite system with the worker.")
	flagKiteLocal      = flag.Bool("kite-local", false, "Start kite system in local mode.")
	flagKiteProxy      = flag.Bool("kite-proxy", false, "Start kite system behind a proxy")
	flagKiteKontrolURL = flag.String("kite-kontrol-url", "", "Change kite's register URL to kontrol")
)

type Runner struct {
	Log             logging.Logger
	Conf            *config.Config
	Bongo           *bongo.Bongo
	Listener        *worker.Listener
	Name            string
	ShutdownHandler func()
	Done            chan error
	Kite            *kite.Kite
}

func New(name string) *Runner {
	return &Runner{Name: name}
}

func WrapWithVersion(name string, version *int) string {
	return fmt.Sprintf("%s:%d", name, *version)
}

func (r *Runner) Init() error {
	flag.Parse()
	if *flagConfFile == "" {
		return fmt.Errorf("Please define config file with -c Exiting...")
	}

	return r.InitWithConfigFile(*flagConfFile)
}

// InitWithConfigFile used for externally setting config file.
// This is used for testing purposes, and usage of Init method is encouraged
func (r *Runner) InitWithConfigFile(flagConfFile string) error {
	r.Conf = config.MustRead(flagConfFile)

	// create logger for our package
	r.Log = helper.CreateLogger(
		WrapWithVersion(r.Name, flagVersion),
		*flagDebug,
	)

	// panics if not successful
	r.Bongo = helper.MustInitBongo(
		WrapWithVersion(r.Name, flagVersion),
		WrapWithVersion(r.Conf.EventExchangeName, flagVersion),
		r.Conf,
		r.Log,
		*flagDebug,
	)

	r.ShutdownHandler = func() {}
	r.Done = make(chan error, 1)
	r.RegisterSignalHandler()

	if *flagKiteInit {
		if err := r.initKite(); err != nil {
			return err
		}
	}

	return nil
}

func (r *Runner) initKite() error {
	// init kite here
	k := kite.New(r.Name, "0.0."+strconv.Itoa(*flagVersion))

	// TODO use get
	k.Config = kiteconfig.MustGet()
	// no need to set, will be set randomly.
	// k.Config.Port = 9876
	k.Config.Environment = r.Conf.Environment
	region := *flagRegion
	// if region is not given, get it from config
	if region == "" {
		region = k.Config.Region
	}

	k.Config.Region = region
	// set kite
	r.Kite = k

	if err := r.RegisterToKontrol(); err != nil {
		return err
	}

	return nil
}

func (r *Runner) RegisterToKontrol() error {
	if r.Kite == nil {
		return errors.New("kite is not initialized yet")
	}

	registerURL := r.Kite.RegisterURL(*flagKiteLocal)
	if *flagKiteKontrolURL != "" {
		u, err := url.Parse(*flagKiteKontrolURL)
		if err != nil {
			r.Log.Fatal("Couldn't parse register url: %s", err)
		}

		registerURL = u
	}

	r.Log.Info("Going to register to kontrol with URL: %s", registerURL)
	if *flagKiteProxy {
		// Koding proxies in production only
		proxyQuery := &protocol.KontrolQuery{
			Username:    "koding",
			Environment: "production",
			Name:        "proxy",
		}

		r.Log.Info("Seaching proxy: %#v", proxyQuery)
		go r.Kite.RegisterToProxy(registerURL, proxyQuery)
		return nil
	}

	return r.Kite.RegisterForever(registerURL)
}

func (r *Runner) Listen(handler worker.Handler) {
	r.Listener = worker.NewListener(
		WrapWithVersion(r.Name, flagVersion),
		WrapWithVersion(r.Conf.EventExchangeName, flagVersion),
		r.Log,
	)

	// blocking
	// listen for events
	r.Listener.Listen(helper.NewRabbitMQ(r.Conf, r.Log), handler)
}

func (r *Runner) Close() {
	r.ShutdownHandler()
	if r.Listener != nil {
		r.Listener.Close()
	}
	r.Bongo.Close()
	if *flagKiteInit {
		if r.Kite == nil {
			return
		}

		r.Kite.Close()
	}
}

func (r *Runner) RegisterSignalHandler() {
	go func() {
		signals := make(chan os.Signal, 1)
		signal.Notify(signals)
		for {
			signal := <-signals
			switch signal {
			case syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGSTOP:
				r.Close()
				r.Done <- nil
			}
		}
	}()
}

func (r *Runner) Wait() error {
	err := <-r.Done

	return err
}
<commit_msg>Social/kite: simplify error handling<commit_after>package runner

import (
	"errors"
	"flag"
	"fmt"
	"net/url"
	"os"
	"os/signal"
	"socialapi/config"
	"socialapi/workers/helper"
	"strconv"
	"syscall"
	"github.com/koding/kite"
	kiteconfig "github.com/koding/kite/config"
	"github.com/koding/kite/protocol"

	"github.com/koding/bongo"
	"github.com/koding/logging"
	"github.com/koding/worker"
)

var (
	flagConfFile = flag.String("c", "", "Configuration profile from file")
	flagRegion   = flag.String("r", "", "Region name")
	flagDebug    = flag.Bool("d", false, "Debug mode")
	flagVersion  = flag.Int("v", 0, "Worker Version")

	flagKiteInit       = flag.Bool("kite-init", false, "Init kite system with the worker.")
	flagKiteLocal      = flag.Bool("kite-local", false, "Start kite system in local mode.")
	flagKiteProxy      = flag.Bool("kite-proxy", false, "Start kite system behind a proxy")
	flagKiteKontrolURL = flag.String("kite-kontrol-url", "", "Change kite's register URL to kontrol")
)

type Runner struct {
	Log             logging.Logger
	Conf            *config.Config
	Bongo           *bongo.Bongo
	Listener        *worker.Listener
	Name            string
	ShutdownHandler func()
	Done            chan error
	Kite            *kite.Kite
}

func New(name string) *Runner {
	return &Runner{Name: name}
}

func WrapWithVersion(name string, version *int) string {
	return fmt.Sprintf("%s:%d", name, *version)
}

func (r *Runner) Init() error {
	flag.Parse()
	if *flagConfFile == "" {
		return fmt.Errorf("Please define config file with -c Exiting...")
	}

	return r.InitWithConfigFile(*flagConfFile)
}

// InitWithConfigFile used for externally setting config file.
// This is used for testing purposes, and usage of Init method is encouraged
func (r *Runner) InitWithConfigFile(flagConfFile string) error {
	r.Conf = config.MustRead(flagConfFile)

	// create logger for our package
	r.Log = helper.CreateLogger(
		WrapWithVersion(r.Name, flagVersion),
		*flagDebug,
	)

	// panics if not successful
	r.Bongo = helper.MustInitBongo(
		WrapWithVersion(r.Name, flagVersion),
		WrapWithVersion(r.Conf.EventExchangeName, flagVersion),
		r.Conf,
		r.Log,
		*flagDebug,
	)

	r.ShutdownHandler = func() {}
	r.Done = make(chan error, 1)
	r.RegisterSignalHandler()

	if *flagKiteInit {
		if err := r.initKite(); err != nil {
			return err
		}
	}

	return nil
}

func (r *Runner) initKite() error {
	// init kite here
	k := kite.New(r.Name, "0.0."+strconv.Itoa(*flagVersion))

	// TODO use get
	k.Config = kiteconfig.MustGet()
	// no need to set, will be set randomly.
	// k.Config.Port = 9876
	k.Config.Environment = r.Conf.Environment
	region := *flagRegion
	// if region is not given, get it from config
	if region == "" {
		region = k.Config.Region
	}

	k.Config.Region = region
	// set kite
	r.Kite = k

	return r.RegisterToKontrol()
}

func (r *Runner) RegisterToKontrol() error {
	if r.Kite == nil {
		return errors.New("kite is not initialized yet")
	}

	registerURL := r.Kite.RegisterURL(*flagKiteLocal)
	if *flagKiteKontrolURL != "" {
		u, err := url.Parse(*flagKiteKontrolURL)
		if err != nil {
			r.Log.Fatal("Couldn't parse register url: %s", err)
		}

		registerURL = u
	}

	r.Log.Info("Going to register to kontrol with URL: %s", registerURL)
	if *flagKiteProxy {
		// Koding proxies in production only
		proxyQuery := &protocol.KontrolQuery{
			Username:    "koding",
			Environment: "production",
			Name:        "proxy",
		}

		r.Log.Info("Seaching proxy: %#v", proxyQuery)
		go r.Kite.RegisterToProxy(registerURL, proxyQuery)
		return nil
	}

	return r.Kite.RegisterForever(registerURL)
}

func (r *Runner) Listen(handler worker.Handler) {
	r.Listener = worker.NewListener(
		WrapWithVersion(r.Name, flagVersion),
		WrapWithVersion(r.Conf.EventExchangeName, flagVersion),
		r.Log,
	)

	// blocking
	// listen for events
	r.Listener.Listen(helper.NewRabbitMQ(r.Conf, r.Log), handler)
}

func (r *Runner) Close() {
	r.ShutdownHandler()
	if r.Listener != nil {
		r.Listener.Close()
	}
	r.Bongo.Close()
	if *flagKiteInit {
		if r.Kite == nil {
			return
		}

		r.Kite.Close()
	}
}

func (r *Runner) RegisterSignalHandler() {
	go func() {
		signals := make(chan os.Signal, 1)
		signal.Notify(signals)
		for {
			signal := <-signals
			switch signal {
			case syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGSTOP:
				r.Close()
				r.Done <- nil
			}
		}
	}()
}

func (r *Runner) Wait() error {
	err := <-r.Done

	return err
}
<|endoftext|>


<commit_before>// Copyright 2016 Stratumn SAS. All rights reserved.
// Use of this source code is governed by an Apache License 2.0
// that can be found in the LICENSE file.

// Package blockcypher defines primitives to work with the BlockCypher API.
package blockcypher

import (
	"encoding/hex"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/blockcypher/gobcy"
	"github.com/btcsuite/btcutil/base58"
	"github.com/stratumn/go/types"
	"github.com/stratumn/goprivate/blockchain/btc"
)

const (
	// DefaultLimiterInterval is the default BlockCypher API limiter interval.
	DefaultLimiterInterval = 2 * time.Minute

	// DefaultLimiterSize is the default BlockCypher API limiter size.
	DefaultLimiterSize = 2
)

// Config contains configuration options for the client.
type Config struct {
	// Network is the Bitcoin network.
	Network btc.Network

	// APIKey is an optional BlockCypher API key.
	APIKey string

	// LimiterInterval is the BlockCypher API limiter interval.
	LimiterInterval time.Duration

	// LimiterSize is the BlockCypher API limiter size.
	LimiterSize int
}

// Client is a BlockCypher API client.
type Client struct {
	config    *Config
	api       *gobcy.API
	limiter   chan struct{}
	closeChan chan struct{}
	timer     *time.Timer
	waitGroup sync.WaitGroup
}

// New creates a client for a Bitcoin network, using an optional BlockCypher API key.
func New(c *Config) *Client {
	parts := strings.Split(c.Network.String(), ":")
	size := c.LimiterSize
	if size == 0 {
		size = DefaultLimiterSize
	}
	limiter := make(chan struct{}, size)

	return &Client{
		config:    c,
		api:       &gobcy.API{c.APIKey, "btc", parts[1]},
		limiter:   limiter,
		closeChan: make(chan struct{}),
	}
}

// FindUnspent implements github.com/stratumn/goprivate/blockchain/btc.UnspentFinder.FindUnspent.
func (c *Client) FindUnspent(address *types.ReversedBytes20, amount int64) ([]btc.Output, int64, error) {
	if err := c.wait(); err != nil {
		return nil, 0, err
	}
	c.waitGroup.Add(1)
	defer c.waitGroup.Done()

	addr := base58.CheckEncode(address[:], c.config.Network.ID())
	addrInfo, err := c.api.GetAddr(addr, map[string]string{
		"unspentOnly":   "true",
		"includeScript": "true",
		"limit":         "50",
	})
	if err != nil {
		return nil, 0, err
	}

	var (
		outputs []btc.Output
		total   int64
	)

TX_LOOP:
	for _, TXRef := range addrInfo.TXRefs {
		output := btc.Output{Index: TXRef.TXOutputN}
		if err := output.TXHash.Unstring(TXRef.TXHash); err != nil {
			return nil, 0, err
		}

		output.PKScript, err = hex.DecodeString(TXRef.Script)
		if err != nil {
			return nil, 0, err
		}

		outputs = append(outputs, output)

		total += int64(TXRef.Value)
		if total >= amount {
			break TX_LOOP
		}
	}

	if total < amount {
		return nil, 0, fmt.Errorf("could not get amount %d got %d", amount, total)
	}

	return outputs, total, nil
}

// Broadcast implements github.com/stratumn/goprivate/blockchain/btc.Broadcaster.Broadcast.
func (c *Client) Broadcast(raw []byte) error {
	if err := c.wait(); err != nil {
		return err
	}
	c.waitGroup.Add(1)
	defer c.waitGroup.Done()

	_, err := c.api.PushTX(hex.EncodeToString(raw))
	return err
}

// Start starts the client.
func (c *Client) Start() {
	size := c.config.LimiterSize
	if size == 0 {
		size = DefaultLimiterSize
	}
	interval := c.config.LimiterInterval
	if interval == 0 {
		interval = DefaultLimiterInterval
	}
	for i := 0; i < size; i++ {
		c.limiter <- struct{}{}
	}

	c.timer = time.NewTimer(interval)

	go func() {
		for {
			for _ = range c.timer.C {
				break
			}
			if c.closeChan == nil {
				return
			}
			c.timer = time.NewTimer(interval)
			c.limiter <- struct{}{}
		}
	}()

	<-c.closeChan
}

// Stop stops the client.
func (c *Client) Stop() {
	c.closeChan <- struct{}{}
	close(c.closeChan)
	c.closeChan = nil

	if !c.timer.Stop() {
		<-c.timer.C
	}

	c.waitGroup.Wait()
	<-c.limiter
	close(c.limiter)
}

func (c *Client) wait() error {
	for _ = range c.limiter {
		break
	}
	if c.closeChan == nil {
		return errors.New("client is stopped")
	}
	return nil
}
<commit_msg>blockchain: Fix batch interval<commit_after>// Copyright 2016 Stratumn SAS. All rights reserved.
// Use of this source code is governed by an Apache License 2.0
// that can be found in the LICENSE file.

// Package blockcypher defines primitives to work with the BlockCypher API.
package blockcypher

import (
	"encoding/hex"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/blockcypher/gobcy"
	"github.com/btcsuite/btcutil/base58"
	"github.com/stratumn/go/types"
	"github.com/stratumn/goprivate/blockchain/btc"
)

const (
	// DefaultLimiterInterval is the default BlockCypher API limiter interval.
	DefaultLimiterInterval = time.Minute

	// DefaultLimiterSize is the default BlockCypher API limiter size.
	DefaultLimiterSize = 2
)

// Config contains configuration options for the client.
type Config struct {
	// Network is the Bitcoin network.
	Network btc.Network

	// APIKey is an optional BlockCypher API key.
	APIKey string

	// LimiterInterval is the BlockCypher API limiter interval.
	LimiterInterval time.Duration

	// LimiterSize is the BlockCypher API limiter size.
	LimiterSize int
}

// Client is a BlockCypher API client.
type Client struct {
	config    *Config
	api       *gobcy.API
	limiter   chan struct{}
	closeChan chan struct{}
	timer     *time.Timer
	waitGroup sync.WaitGroup
}

// New creates a client for a Bitcoin network, using an optional BlockCypher API key.
func New(c *Config) *Client {
	parts := strings.Split(c.Network.String(), ":")
	size := c.LimiterSize
	if size == 0 {
		size = DefaultLimiterSize
	}
	limiter := make(chan struct{}, size)

	return &Client{
		config:    c,
		api:       &gobcy.API{c.APIKey, "btc", parts[1]},
		limiter:   limiter,
		closeChan: make(chan struct{}),
	}
}

// FindUnspent implements github.com/stratumn/goprivate/blockchain/btc.UnspentFinder.FindUnspent.
func (c *Client) FindUnspent(address *types.ReversedBytes20, amount int64) ([]btc.Output, int64, error) {
	if err := c.wait(); err != nil {
		return nil, 0, err
	}
	c.waitGroup.Add(1)
	defer c.waitGroup.Done()

	addr := base58.CheckEncode(address[:], c.config.Network.ID())
	addrInfo, err := c.api.GetAddr(addr, map[string]string{
		"unspentOnly":   "true",
		"includeScript": "true",
		"limit":         "50",
	})
	if err != nil {
		return nil, 0, err
	}

	var (
		outputs []btc.Output
		total   int64
	)

TX_LOOP:
	for _, TXRef := range addrInfo.TXRefs {
		output := btc.Output{Index: TXRef.TXOutputN}
		if err := output.TXHash.Unstring(TXRef.TXHash); err != nil {
			return nil, 0, err
		}

		output.PKScript, err = hex.DecodeString(TXRef.Script)
		if err != nil {
			return nil, 0, err
		}

		outputs = append(outputs, output)

		total += int64(TXRef.Value)
		if total >= amount {
			break TX_LOOP
		}
	}

	if total < amount {
		return nil, 0, fmt.Errorf("could not get amount %d got %d", amount, total)
	}

	return outputs, total, nil
}

// Broadcast implements github.com/stratumn/goprivate/blockchain/btc.Broadcaster.Broadcast.
func (c *Client) Broadcast(raw []byte) error {
	if err := c.wait(); err != nil {
		return err
	}
	c.waitGroup.Add(1)
	defer c.waitGroup.Done()

	_, err := c.api.PushTX(hex.EncodeToString(raw))
	return err
}

// Start starts the client.
func (c *Client) Start() {
	size := c.config.LimiterSize
	if size == 0 {
		size = DefaultLimiterSize
	}
	interval := c.config.LimiterInterval
	if interval == 0 {
		interval = DefaultLimiterInterval
	}
	for i := 0; i < size; i++ {
		c.limiter <- struct{}{}
	}

	c.timer = time.NewTimer(interval)

	go func() {
		for {
			for _ = range c.timer.C {
				break
			}
			if c.closeChan == nil {
				return
			}
			c.timer = time.NewTimer(interval)
			c.limiter <- struct{}{}
		}
	}()

	<-c.closeChan
}

// Stop stops the client.
func (c *Client) Stop() {
	c.closeChan <- struct{}{}
	close(c.closeChan)
	c.closeChan = nil

	if !c.timer.Stop() {
		<-c.timer.C
	}

	c.waitGroup.Wait()
	<-c.limiter
	close(c.limiter)
}

func (c *Client) wait() error {
	for _ = range c.limiter {
		break
	}
	if c.closeChan == nil {
		return errors.New("client is stopped")
	}
	return nil
}
<|endoftext|>


<commit_before>/*
 * Copyright 2017 Manuel Gauto (github.com/twa16)
 *
 * 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 userspaced

import (
	"context"
	"github.com/jinzhu/gorm"
	"github.com/pkg/errors"
	"math/rand"
	"strconv"
	"sync"
	"github.com/fsouza/go-dockerclient"
)

//Contains running instances of docker hosts
var DockerInstances []*DockerInstance

//initDockerHosts Pulls docker hosts from the DB, adds them to the cache and connects to them
func initDockerHosts(db *gorm.DB) {
	instances := getAllDockerInstanceConfigurations(db)
	log.Infof("Initiating Connections to %d Docker Host(s)\n", len(instances))
	for _, instance := range instances {
		_, err := addAndConnectToDockerInstance(db, &instance)
		if err != nil {
			log.Critical("Adding Docker Host %s Failed: %s\n", instance.Name, err.Error())
			continue
		}
		log.Infof("Connected to Docker Host: %s\n", instance.Name)
	}
	log.Info("Initialized All Docker Hosts!")
}

var dockerInstanceSliceLock sync.Mutex

//addAndConnectToDockerInstance Adds a new host to in-memory cache and establishes connection to it.
func addAndConnectToDockerInstance(db *gorm.DB, instance *DockerInstance) (*DockerInstance, error) {
	//It is technically possible for two hosts to be added at once, so let's lock the slice
	dockerInstanceSliceLock.Lock()
	defer dockerInstanceSliceLock.Unlock()
	_, err := startDockerClient(instance)
	if err != nil {
		log.Fatalf("Error Starting Docker Client: %s\n", err.Error())
		return nil, err
	}
	db.Save(&instance)
	DockerInstances = append(DockerInstances, instance)
	return instance, nil
}

//getAllDockerInstanceConfigurations Gets all hosts from the db
func getAllDockerInstanceConfigurations(db *gorm.DB) []DockerInstance {
	configs := []DockerInstance{}
	db.Find(&configs)
	return configs
}

//getHostByID Helper method that gets a host from the db by id
func getHostByID(hostID uint) *DockerInstance {
	for i, instance := range DockerInstances {
		if instance.ID == hostID {
			return DockerInstances[i]
		}
	}
	return nil
}

//getImageByID Helper method that gets an image by id
func getImageByID(db *gorm.DB, imageID uint) SpaceImage {
	var image SpaceImage
	db.Find(&image, imageID)
	return image
}

//checkImageExists Check if an image exists
func checkImageExists(db *gorm.DB, imageID uint) bool {
	var image SpaceImage
	return !db.Find(&image, imageID).RecordNotFound()
}

//securePortForSpace Picks an open port between 20000 and 30000. Saves new PortLink
func securePortForSpace(db *gorm.DB, space *Space, destPort uint16) *SpacePortLink {
	log.Debugf("Attempting to secure port for %d: %d\n", space.ID, destPort)
	spaceHost := getHostByID(space.HostID)
	originalPort := space.PortLinks
	for true {
		//Copy over the basics
		portMapping := SpacePortLink{}
		portMapping.ExternalAddress = spaceHost.ExternalAddress
		portMapping.DisplayAddress = spaceHost.ExternalDisplayAddress
		portMapping.SpacePort = destPort
		//Generate a new port
		var portTry = 20000 + rand.Intn(10000)
		//Set it and append the mapping
		portMapping.ExternalPort = uint16(portTry)
		log.Debugf("Trying to secure port %d for %d\n", portTry, space.ID)
		space.PortLinks = append(originalPort, portMapping)
		//db.Model(&space).Related(&space.PortLinks)
		//Try to update and see if we get an error
		//Maybe we should add an attempt cap here
		//err := db.Save(&space).Error
		if isPortOpen(db, portMapping.ExternalAddress, portMapping.ExternalPort) {
			err := db.Save(&space).Error
			if err != nil {
				log.Criticalf("Error Saving Port Mapping: %s\n", err.Error())
			} else {
				log.Infof("Secured port mapping for space %d: %d -> %d\n", space.ID, portTry, destPort)
				return &portMapping
			}
		} else {
			log.Infof("Port %d is taken\n", portTry)
		}
	}
	//This should never happen unless something went horribly wrong
	return nil
}

func isPortOpen(db *gorm.DB, externalAddress string, port uint16) bool {
	var holder SpacePortLink
	return db.Where(&SpacePortLink{ExternalAddress: externalAddress, ExternalPort: port}).First(&holder).RecordNotFound()
}

//TODO: Make this do the thing. Returns first instance since when this was written spaces weren't created yet. This will probably be done with raw sql.
//selectLeastOccupiedHost Returns the host that has the fewest number of instances.
func selectLeastOccupiedHost(db *gorm.DB) (*DockerInstance, error) {
	if len(DockerInstances) == 0 {
		return nil, errors.New("No Hosts Have Been Added!")
	}
	return DockerInstances[0], nil
}

//startSpace Creates and starts a new space
func startSpace(db *gorm.DB, space *Space, creationStatusChan chan string) (error, *Space) {
	//======Initialization Steps=====
	//Check if the requested image exists
	if !checkImageExists(db, space.ImageID) {
		creationStatusChan <- "Error: Invalid Image"
		return errors.New("Invalid Image Specified"), nil
	}
	//Pick a host
	dockerHost, err := selectLeastOccupiedHost(db)
	if err != nil {
		log.Critical("No hosts have been added.")
		return err, nil
	}
	space.HostID = dockerHost.ID
	space.SpaceState = "creation started"
	client := dockerHost.DockerClient
	//Save it
	db.Create(&space)
	creationStatusChan <- "Host Chosen"
	log.Infof("Selected Host %d for space %d\n", space.HostID, space.ID)

	//======Container Config=====
	var containerConfig docker.Config
	//Set the image
	containerConfig.Image = getImageByID(db, space.ImageID).DockerImage

	//Empty placeholder struct
	var v struct{}

	//Create empty volume set
	//containerConfig.Volumes = make(map[string]struct{})
	//Ports
	//port80, _ := nat.NewPort("tcp", "80")
	containerConfig.ExposedPorts = make(map[docker.Port]struct{})
	containerConfig.ExposedPorts["22/tcp"] = v
	containerConfig.ExposedPorts["1337/tcp"] = v
	containerConfig.ExposedPorts["1337/udp"] = v

	//=====Host Config======
	var hostConfig docker.HostConfig
	//Secure Ports in DB
	sshPortLink := securePortForSpace(db, space, 22)
	servicePortLink := securePortForSpace(db, space, 1337)
	//Setup Port Maps
	//Forward a dynamic host port to container. Listen on localhost so that nginx can proxy.
	hostConfig.PortBindings = make(map[docker.Port][]docker.PortBinding)
	hostConfig.PortBindings["22/tcp"] = append(hostConfig.PortBindings["22/tcp"], docker.PortBinding{HostIP: "127.0.0.1", HostPort: strconv.Itoa(int(sshPortLink.ExternalPort))})
	hostConfig.PortBindings["1337/tcp"] = append(hostConfig.PortBindings["1337/tcp"], docker.PortBinding{HostIP: "127.0.0.1", HostPort: strconv.Itoa(int(servicePortLink.ExternalPort))})
	hostConfig.PortBindings["1337/udp"] = append(hostConfig.PortBindings["1337/udp"], docker.PortBinding{HostIP: "127.0.0.1", HostPort: strconv.Itoa(int(servicePortLink.ExternalPort))})
	//Save PortLinks
	sshPortLink.ExternalAddress = dockerHost.ExternalAddress
	sshPortLink.DisplayAddress = dockerHost.ExternalDisplayAddress
	sshPortLink.SpacePort = 22

	servicePortLink.ExternalAddress = dockerHost.ExternalAddress
	servicePortLink.DisplayAddress = dockerHost.ExternalDisplayAddress
	servicePortLink.SpacePort = 1337

	space.PortLinks = append(space.PortLinks, *sshPortLink)
	space.PortLinks = append(space.PortLinks, *servicePortLink)
	//======Network Config=====
	var networkConfig docker.NetworkingConfig

	//======Container Creation=====
	//Wrapup config
	var config docker.CreateContainerOptions
	config.Config = &containerConfig
	config.HostConfig = &hostConfig
	config.NetworkingConfig = &networkConfig
	config.Context = context.Background()
	//Create Container
	c, err := client.CreateContainer(config)
	if err != nil {
		log.Criticalf("Failed to create container for space %d\n", space.ID)
		log.Debug(err)
		space.SpaceState = "Error Creating"
		db.Save(&space)
		creationStatusChan <- "Error: Error Creating Container"
		return err, nil
	}
	creationStatusChan <- "Container Created"
	//Set container
	space.ContainerID = c.ID
	//At first, let us keep all containers alive
	//TODO: Make this configurable and with a quota
	space.KeepAlive = true
	space.SpaceState = "created"
	db.Save(&space)
	log.Infof("Created container for space %d: %s\n", space.ID, space.ContainerID)

	err = client.StartContainer(space.ContainerID, nil)
	if err != nil {
		log.Criticalf("Error starting container for space %d: %s\n", space.ID, err.Error())
		space.SpaceState = "error starting"
		db.Save(&space)
	} else {
		log.Infof("Container for Space %d started: %s\n", space.ID, space.ContainerID)
		space.SpaceState = "running"
		db.Save(&space)
	}
	creationStatusChan <- "Creation Complete"
	//execInSpace(database, *space, []string{"touch", "/root/testblop"})
	return nil, space
}

//execInSpace Executes a command in a space
func execInSpace(db *gorm.DB, space Space, command []string) error {
	dockerHost := getHostByID(space.HostID)

	execOptions := docker.CreateExecOptions{}
	execOptions.Cmd = command
	execOptions.AttachStderr = false
	execOptions.AttachStdin = false
	execOptions.AttachStdout = false
	execOptions.Tty = false
	exec, err := dockerHost.DockerClient.CreateExec(execOptions)
	if err != nil {
		log.Warningf("Error Executing Command on Host %s: %s\n", dockerHost.Name, err.Error())
		return err
	}
	err = dockerHost.DockerClient.StartExec(exec.ID, docker.StartExecOptions{})
	if err != nil {
		log.Warningf("Error Executing Command on Host %s: %s\n", dockerHost.Name, err.Error())
		return err
	}
	return nil
}

//startDockerClient Opens a connection to a docker instance
func startDockerClient(instance *DockerInstance) (*docker.Client, error) {
	log.Infof("Connecting to Docker Host %s using connection type %s\n", instance.Name, instance.ConnectionType)
	var cli *docker.Client
	var err error
	//Check for type and init accordingly
	if instance.ConnectionType == "local" {
		if instance.Endpoint == "" {
			cli, err = docker.NewClientFromEnv()
		} else {
			cli, err = docker.NewClient(instance.Endpoint)
		}
	} else {
		cli, err = docker.NewTLSClient(instance.Endpoint, instance.ClientCertPath, instance.ClientKeyPath, instance.CaCertPath)
	}

	if err != nil {
		log.Criticalf("Failed to start docker client for %s: %s\n", instance.Name, err.Error())
		return nil, err
	}
	//Put new connection data into the struct
	instance.IsConnected = true
	instance.DockerClient = cli

	env, _ := cli.Version()
	log.Info("Connection Suceeded! API Version: " + env.Get("ApiVersion"))
	return cli, err
}

//ensureStarterImages Pulls docker images needed for spaces to all hosts
func ensureStarterImages(db *gorm.DB) {
	ubuntuImage := SpaceImage{}
	db.Where("docker_image = ? AND docker_image_tag = ?", "userspace/ubuntu", "latest").First(&ubuntuImage)
	if ubuntuImage.Active == false {
		log.Info("Downloading Start Images")
		ubuntuImage.Active = true
		ubuntuImage.Description = "Basic Ubuntu Image"
		ubuntuImage.DockerImage = "userspace/ubuntu"
		ubuntuImage.DockerImageTag = "latest"
		ubuntuImage.Name = "Ubuntu"
		db.Create(&ubuntuImage)
	}
}

//downloadDockerImages Download an image to host
func downloadDockerImages(db *gorm.DB) {
	images := []SpaceImage{}
	db.Find(&images)
	for _, image := range images {
		for _, instance := range DockerInstances {
			if instance.IsConnected {
				pullDockerImage(instance.DockerClient, image.DockerImage, image.DockerImageTag)
				log.Infof("Downloaded image to %s\n", instance.Name)
			} else {
				log.Warningf("Skipping %s as it is not connected!\n", instance.Name)
			}
		}
	}
}

//pullDockerImage Pulls a docker image from the hub
func pullDockerImage(dClient *docker.Client, image string, tag string) error {
	pullOptions := docker.PullImageOptions{Repository: image, Tag: tag}
	authOptions := docker.AuthConfiguration{}
	err := dClient.PullImage(pullOptions, authOptions)
	return err
}
<commit_msg>Move RemoveSpace() to docker.go<commit_after>/*
 * Copyright 2017 Manuel Gauto (github.com/twa16)
 *
 * 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 userspaced

import (
	"context"
	"math/rand"
	"strconv"
	"sync"

	"github.com/fsouza/go-dockerclient"
	"github.com/jinzhu/gorm"
	"github.com/pkg/errors"
)

//Contains running instances of docker hosts
var DockerInstances []*DockerInstance

//initDockerHosts Pulls docker hosts from the DB, adds them to the cache and connects to them
func initDockerHosts(db *gorm.DB) {
	instances := getAllDockerInstanceConfigurations(db)
	log.Infof("Initiating Connections to %d Docker Host(s)\n", len(instances))
	for _, instance := range instances {
		_, err := addAndConnectToDockerInstance(db, &instance)
		if err != nil {
			log.Critical("Adding Docker Host %s Failed: %s\n", instance.Name, err.Error())
			continue
		}
		log.Infof("Connected to Docker Host: %s\n", instance.Name)
	}
	log.Info("Initialized All Docker Hosts!")
}

var dockerInstanceSliceLock sync.Mutex

//addAndConnectToDockerInstance Adds a new host to in-memory cache and establishes connection to it.
func addAndConnectToDockerInstance(db *gorm.DB, instance *DockerInstance) (*DockerInstance, error) {
	//It is technically possible for two hosts to be added at once, so let's lock the slice
	dockerInstanceSliceLock.Lock()
	defer dockerInstanceSliceLock.Unlock()
	_, err := startDockerClient(instance)
	if err != nil {
		log.Fatalf("Error Starting Docker Client: %s\n", err.Error())
		return nil, err
	}
	db.Save(&instance)
	DockerInstances = append(DockerInstances, instance)
	return instance, nil
}

//getAllDockerInstanceConfigurations Gets all hosts from the db
func getAllDockerInstanceConfigurations(db *gorm.DB) []DockerInstance {
	configs := []DockerInstance{}
	db.Find(&configs)
	return configs
}

//getHostByID Helper method that gets a host from the db by id
func getHostByID(hostID uint) *DockerInstance {
	for i, instance := range DockerInstances {
		if instance.ID == hostID {
			return DockerInstances[i]
		}
	}
	return nil
}

//getImageByID Helper method that gets an image by id
func getImageByID(db *gorm.DB, imageID uint) SpaceImage {
	var image SpaceImage
	db.Find(&image, imageID)
	return image
}

//checkImageExists Check if an image exists
func checkImageExists(db *gorm.DB, imageID uint) bool {
	var image SpaceImage
	return !db.Find(&image, imageID).RecordNotFound()
}

//securePortForSpace Picks an open port between 20000 and 30000. Saves new PortLink
func securePortForSpace(db *gorm.DB, space *Space, destPort uint16) *SpacePortLink {
	log.Debugf("Attempting to secure port for %d: %d\n", space.ID, destPort)
	spaceHost := getHostByID(space.HostID)
	originalPort := space.PortLinks
	for true {
		//Copy over the basics
		portMapping := SpacePortLink{}
		portMapping.ExternalAddress = spaceHost.ExternalAddress
		portMapping.DisplayAddress = spaceHost.ExternalDisplayAddress
		portMapping.SpacePort = destPort
		//Generate a new port
		var portTry = 20000 + rand.Intn(10000)
		//Set it and append the mapping
		portMapping.ExternalPort = uint16(portTry)
		log.Debugf("Trying to secure port %d for %d\n", portTry, space.ID)
		space.PortLinks = append(originalPort, portMapping)
		//db.Model(&space).Related(&space.PortLinks)
		//Try to update and see if we get an error
		//Maybe we should add an attempt cap here
		//err := db.Save(&space).Error
		if isPortOpen(db, portMapping.ExternalAddress, portMapping.ExternalPort) {
			err := db.Save(&space).Error
			if err != nil {
				log.Criticalf("Error Saving Port Mapping: %s\n", err.Error())
			} else {
				log.Infof("Secured port mapping for space %d: %d -> %d\n", space.ID, portTry, destPort)
				return &portMapping
			}
		} else {
			log.Infof("Port %d is taken\n", portTry)
		}
	}
	//This should never happen unless something went horribly wrong
	return nil
}

func isPortOpen(db *gorm.DB, externalAddress string, port uint16) bool {
	var holder SpacePortLink
	return db.Where(&SpacePortLink{ExternalAddress: externalAddress, ExternalPort: port}).First(&holder).RecordNotFound()
}

//TODO: Make this do the thing. Returns first instance since when this was written spaces weren't created yet. This will probably be done with raw sql.
//selectLeastOccupiedHost Returns the host that has the fewest number of instances.
func selectLeastOccupiedHost(db *gorm.DB) (*DockerInstance, error) {
	if len(DockerInstances) == 0 {
		return nil, errors.New("No Hosts Have Been Added!")
	}
	return DockerInstances[0], nil
}

//startSpace Creates and starts a new space
func startSpace(db *gorm.DB, space *Space, creationStatusChan chan string) (error, *Space) {
	//======Initialization Steps=====
	//Check if the requested image exists
	if !checkImageExists(db, space.ImageID) {
		creationStatusChan <- "Error: Invalid Image"
		return errors.New("Invalid Image Specified"), nil
	}
	//Pick a host
	dockerHost, err := selectLeastOccupiedHost(db)
	if err != nil {
		log.Critical("No hosts have been added.")
		return err, nil
	}
	space.HostID = dockerHost.ID
	space.SpaceState = "creation started"
	client := dockerHost.DockerClient
	//Save it
	db.Create(&space)
	creationStatusChan <- "Host Chosen"
	log.Infof("Selected Host %d for space %d\n", space.HostID, space.ID)

	//======Container Config=====
	var containerConfig docker.Config
	//Set the image
	containerConfig.Image = getImageByID(db, space.ImageID).DockerImage

	//Empty placeholder struct
	var v struct{}

	//Create empty volume set
	//containerConfig.Volumes = make(map[string]struct{})
	//Ports
	//port80, _ := nat.NewPort("tcp", "80")
	containerConfig.ExposedPorts = make(map[docker.Port]struct{})
	containerConfig.ExposedPorts["22/tcp"] = v
	containerConfig.ExposedPorts["1337/tcp"] = v
	containerConfig.ExposedPorts["1337/udp"] = v

	//=====Host Config======
	var hostConfig docker.HostConfig
	//Secure Ports in DB
	sshPortLink := securePortForSpace(db, space, 22)
	servicePortLink := securePortForSpace(db, space, 1337)
	//Setup Port Maps
	//Forward a dynamic host port to container. Listen on localhost so that nginx can proxy.
	hostConfig.PortBindings = make(map[docker.Port][]docker.PortBinding)
	hostConfig.PortBindings["22/tcp"] = append(hostConfig.PortBindings["22/tcp"], docker.PortBinding{HostIP: "127.0.0.1", HostPort: strconv.Itoa(int(sshPortLink.ExternalPort))})
	hostConfig.PortBindings["1337/tcp"] = append(hostConfig.PortBindings["1337/tcp"], docker.PortBinding{HostIP: "127.0.0.1", HostPort: strconv.Itoa(int(servicePortLink.ExternalPort))})
	hostConfig.PortBindings["1337/udp"] = append(hostConfig.PortBindings["1337/udp"], docker.PortBinding{HostIP: "127.0.0.1", HostPort: strconv.Itoa(int(servicePortLink.ExternalPort))})
	//Save PortLinks
	sshPortLink.ExternalAddress = dockerHost.ExternalAddress
	sshPortLink.DisplayAddress = dockerHost.ExternalDisplayAddress
	sshPortLink.SpacePort = 22

	servicePortLink.ExternalAddress = dockerHost.ExternalAddress
	servicePortLink.DisplayAddress = dockerHost.ExternalDisplayAddress
	servicePortLink.SpacePort = 1337

	space.PortLinks = append(space.PortLinks, *sshPortLink)
	space.PortLinks = append(space.PortLinks, *servicePortLink)
	//======Network Config=====
	var networkConfig docker.NetworkingConfig

	//======Container Creation=====
	//Wrapup config
	var config docker.CreateContainerOptions
	config.Config = &containerConfig
	config.HostConfig = &hostConfig
	config.NetworkingConfig = &networkConfig
	config.Context = context.Background()
	//Create Container
	c, err := client.CreateContainer(config)
	if err != nil {
		log.Criticalf("Failed to create container for space %d\n", space.ID)
		log.Debug(err)
		space.SpaceState = "Error Creating"
		db.Save(&space)
		creationStatusChan <- "Error: Error Creating Container"
		return err, nil
	}
	creationStatusChan <- "Container Created"
	//Set container
	space.ContainerID = c.ID
	//At first, let us keep all containers alive
	//TODO: Make this configurable and with a quota
	space.KeepAlive = true
	space.SpaceState = "created"
	db.Save(&space)
	log.Infof("Created container for space %d: %s\n", space.ID, space.ContainerID)

	err = client.StartContainer(space.ContainerID, nil)
	if err != nil {
		log.Criticalf("Error starting container for space %d: %s\n", space.ID, err.Error())
		space.SpaceState = "error starting"
		db.Save(&space)
	} else {
		log.Infof("Container for Space %d started: %s\n", space.ID, space.ContainerID)
		space.SpaceState = "running"
		db.Save(&space)
	}
	creationStatusChan <- "Creation Complete"
	//execInSpace(database, *space, []string{"touch", "/root/testblop"})
	return nil, space
}

//execInSpace Executes a command in a space
func execInSpace(db *gorm.DB, space Space, command []string) error {
	dockerHost := getHostByID(space.HostID)

	execOptions := docker.CreateExecOptions{}
	execOptions.Cmd = command
	execOptions.AttachStderr = false
	execOptions.AttachStdin = false
	execOptions.AttachStdout = false
	execOptions.Tty = false
	exec, err := dockerHost.DockerClient.CreateExec(execOptions)
	if err != nil {
		log.Warningf("Error Executing Command on Host %s: %s\n", dockerHost.Name, err.Error())
		return err
	}
	err = dockerHost.DockerClient.StartExec(exec.ID, docker.StartExecOptions{})
	if err != nil {
		log.Warningf("Error Executing Command on Host %s: %s\n", dockerHost.Name, err.Error())
		return err
	}
	return nil
}

//startDockerClient Opens a connection to a docker instance
func startDockerClient(instance *DockerInstance) (*docker.Client, error) {
	log.Infof("Connecting to Docker Host %s using connection type %s\n", instance.Name, instance.ConnectionType)
	var cli *docker.Client
	var err error
	//Check for type and init accordingly
	if instance.ConnectionType == "local" {
		if instance.Endpoint == "" {
			cli, err = docker.NewClientFromEnv()
		} else {
			cli, err = docker.NewClient(instance.Endpoint)
		}
	} else {
		cli, err = docker.NewTLSClient(instance.Endpoint, instance.ClientCertPath, instance.ClientKeyPath, instance.CaCertPath)
	}

	if err != nil {
		log.Criticalf("Failed to start docker client for %s: %s\n", instance.Name, err.Error())
		return nil, err
	}
	//Put new connection data into the struct
	instance.IsConnected = true
	instance.DockerClient = cli

	env, _ := cli.Version()
	log.Info("Connection Suceeded! API Version: " + env.Get("ApiVersion"))
	return cli, err
}

//ensureStarterImages Pulls docker images needed for spaces to all hosts
func ensureStarterImages(db *gorm.DB) {
	ubuntuImage := SpaceImage{}
	db.Where("docker_image = ? AND docker_image_tag = ?", "userspace/ubuntu", "latest").First(&ubuntuImage)
	if ubuntuImage.Active == false {
		log.Info("Downloading Start Images")
		ubuntuImage.Active = true
		ubuntuImage.Description = "Basic Ubuntu Image"
		ubuntuImage.DockerImage = "userspace/ubuntu"
		ubuntuImage.DockerImageTag = "latest"
		ubuntuImage.Name = "Ubuntu"
		db.Create(&ubuntuImage)
	}
}

//downloadDockerImages Download an image to host
func downloadDockerImages(db *gorm.DB) {
	images := []SpaceImage{}
	db.Find(&images)
	for _, image := range images {
		for _, instance := range DockerInstances {
			if instance.IsConnected {
				pullDockerImage(instance.DockerClient, image.DockerImage, image.DockerImageTag)
				log.Infof("Downloaded image to %s\n", instance.Name)
			} else {
				log.Warningf("Skipping %s as it is not connected!\n", instance.Name)
			}
		}
	}
}

//pullDockerImage Pulls a docker image from the hub
func pullDockerImage(dClient *docker.Client, image string, tag string) error {
	pullOptions := docker.PullImageOptions{Repository: image, Tag: tag}
	authOptions := docker.AuthConfiguration{}
	err := dClient.PullImage(pullOptions, authOptions)
	return err
}

//RemoveSpace Removes a Space's container and DB entry
func RemoveSpace(db *gorm.DB, space Space) error {
	//Get Host Docker Connection
	hostID := space.HostID
	dockerHost := getHostByID(hostID)
	//Ensure the host is connected
	if !dockerHost.IsConnected {
		log.Critical("Attempted to remove container %s from disconnected host.")
		return errors.New("Attempted to remove contaienr from disconnected host.")
	}
	//Set the space state
	space.SpaceState = "deleting"
	db.Save(&space)

	//Only remove the container if the container is not already dead
	if space.SpaceState != "error" {
		//Stop the container
		dClient := dockerHost.DockerClient
		err := dClient.StopContainer(space.ContainerID, 30)
		//Catch any errors
		if err != nil {
			log.Criticalf("Error stopping container %s: %s", space.ContainerID, err.Error())
			//return err
		} else {
			//Remove the container
			removeOptions := docker.RemoveContainerOptions{
				ID:            space.ContainerID,
				RemoveVolumes: true,
				Force:         true,
				Context:       context.Background(),
			}
			err = dClient.RemoveContainer(removeOptions)
			if err != nil {
				log.Criticalf("Error removing container %s: %s", space.ContainerID, err.Error())
				//return err
			}
		}
	}
	//Remove the db object
	err := db.Delete(&space).Error
	if err != nil {
		log.Criticalf("Error removing space record for %d\n", space.ID, err.Error())
		return err
	}
	return nil
}
<|endoftext|>