<commit_before>// A golang Skip List Implementation.
// https://github.com/huandu/skiplist/
//
// Copyright 2011, Huan Du
// Licensed under the MIT license
// https://github.com/huandu/skiplist/blob/master/LICENSE

// Package skiplist provides a go implementation for skip list.
// About skip list: http://en.wikipedia.org/wiki/Skip_list
//
// Skip list is basically an ordered set.
// Following code creates a skip list with int key and adds some values.
//     list := skiplist.New(skiplist.Int)
//
//     // adds some elements
//     list.Set(20, "Hello")
//     list.Set(10, "World")
//     list.Set(40, true)         // value can be any type
//     list.Set(40, 1000)         // replace last element with new value
//
//     // try to find one
//     e := list.Get(10)           // value is the Element with key 10
//     _ = e.Value.(string)        // it's the "World". remember to do type cast
//     v, ok := list.GetValue(20)  // directly get value. ok is false if not exists
//     v2 := list.MustGetValue(10) // directly get value. panic if key doesn't exist
//     notFound := list.Get(15)    // returns nil if key is not found
//
//     // remove element
//     old := list.Remove(40)     // remove found element and returns its pointer
//                                // returns nil if key is not found
//
//     // re-init list. it will make the list empty.
//     list.Init()
//
// Skip list elements have random number of next pointers. The max number (say
// "max level") is configurable.
//
// The variable skiplist.DefaultMaxLevel is controlling global default.
// Changing it will not affect created lists.
//     skiplist.DefaultMaxLevel = 24  // change default to 24
// Max level of a created list can also be changed even if it's not empty.
//     list.SetMaxLevel(10)
// Remember the side effect when changing this max level value.
// See its wikipedia page for more details.
//
// Most comparable built-in types are pre-defined in skiplist, including
//     byte []byte float32 float64 int int16 int32 int64 int8
//     rune string uint uint16 uint32 uint64 uint8 uintptr
// Pre-defined compare function name is similar to the type name, e.g.
// skiplist.Float32 is for float32 key. A special case is skiplist.Bytes is for []byte.
// These functions order key from small to big (say "ascending order").
// There are also reserved order functions with name like skiplist.IntDesc.
//
// For key types out of the pre-defined list, one can write a custom compare function.
//     type GreaterThanFunc func (lhs, rhs interface{}) bool
// Such compare function returns true if lhs > rhs. Note that, if lhs == rhs, compare
// function (let the name is "foo") must act as following.
//     // if lhs == rhs, following expression must be true
//     foo(lhs, rhs) == false && foo(rhs, lhs) == false
// There is another func type named LessThanFunc. It works similar to GreaterThanFunc,
// except the order is big to small.
// Here is a sample to write a compare func.
//     type Foo struct {
//         value int
//     }
//
//     // it generates a score on a given key.
//     // if key1 > key2, then there must be key1.Score() >= key2.Score().
//     // it's optional. it's worth implementing if call compare func is quite expensive.
//     func (f *Foo) Score() float64 {
//         return float64(f.value)
//     }
//
//     var greater skiplist.GreaterThanFunc = func(lhs, rhs interface{}) {
//         return lhs.(Foo).value > rhs.(Foo).value
//     }
//     list := skiplist.New(greater)
//
//     // descending version is a bit different. mind the func type.
//     var less skiplist.LessThanFunc = func(lhs, rhs interface{}) {
//         return lhs.(Foo).value < rhs.(Foo).value
//     })
//     list := skiplist.New(less)
//
// Skiplist uses global rand source in math/rand by default. This rand source acquires a
// lock when generating random number. Replacing it with a lock-free rand source can provide
// slightly better performance. Use SkipList.SetRandSource to change rand source.
package skiplist

import (
    "math/rand"
)

// Creates a new skiplist.
// keyFunc is a func checking the "greater than" logic.
// If k1 equals k2, keyFunc(k1, k2) and keyFunc(k2, k1) must both be false.
// For built-in types, keyFunc can be found in skiplist package.
// For instance, skiplist.Int is for the list with int keys.
// By default, the list with built-in type key is in ascending order.
// The keyFunc named as skiplist.IntDesc is for descending key order list.
func New(keyFunc Comparable) *SkipList {
    if DefaultMaxLevel <= 0 {
        panic("skiplist default level must not be zero or negative")
    }

    return &SkipList{
        elementNode:    elementNode{next: make([]*Element, DefaultMaxLevel)},
        prevNodesCache: make([]*elementNode, DefaultMaxLevel),
        level:          DefaultMaxLevel,
        keyFunc:        keyFunc,
        randSource:     defaultSource,
        reversed:       keyFunc.Descending(),
    }
}

// Resets a skiplist and discards all exists elements.
func (list *SkipList) Init() *SkipList {
    list.next = make([]*Element, list.level)
    list.length = 0
    return list
}

// Sets a new rand source.
//
// Skiplist uses global rand defined in math/rand by default.
// The default rand acquires a global mutex before generating any number.
// It's not necessary if the skiplist is well protected by caller.
func (list *SkipList) SetRandSource(source rand.Source) {
    list.randSource = source
}

// Gets the first element.
func (list *SkipList) Front() *Element {
    return list.next[0]
}

// Gets list length.
func (list *SkipList) Len() int {
    return list.length
}

// Sets a value in the list with key.
// If the key exists, change element value to the new one.
// Returns new element pointer.
func (list *SkipList) Set(key, value interface{}) *Element {
    var element *Element

    score := getScore(key, list.reversed)
    prevs := list.getPrevElementNodes(key, score)

    // found an element with the same key, replace its value
    if element = prevs[0].next[0]; element != nil && !list.keyFunc.Compare(element.key, key) {
        element.Value = value
        return element
    }

    element = &Element{
        elementNode: elementNode{
            next: make([]*Element, list.randLevel()),
        },
        key:   key,
        score: score,
        Value: value,
    }

    for i := range element.next {
        element.next[i] = prevs[i].next[i]
        prevs[i].next[i] = element
    }

    list.length++
    return element
}

// Gets an element.
// Returns element pointer if found, nil if not found.
func (list *SkipList) Get(key interface{}) *Element {
    var prev *elementNode = &list.elementNode
    var next *Element
    score := getScore(key, list.reversed)

    for i := list.level - 1; i >= 0; i-- {
        next = prev.next[i]

        for next != nil &&
            (score > next.score || (score == next.score && list.keyFunc.Compare(key, next.key))) {
            prev = &next.elementNode
            next = next.next[i]
        }
    }

    if next != nil && score == next.score && !list.keyFunc.Compare(next.key, key) {
        return next
    }

    return nil
}

// Gets a value. It's a short hand for Get().Value.
// Returns value and its existence status.
func (list *SkipList) GetValue(key interface{}) (interface{}, bool) {
    element := list.Get(key)

    if element == nil {
        return nil, false
    }

    return element.Value, true
}

// Gets a value. It will panic if key doesn't exist.
// Returns value.
func (list *SkipList) MustGetValue(key interface{}) interface{} {
    element := list.Get(key)

    if element == nil {
        panic("cannot find key in skiplist")
    }

    return element.Value
}

// Removes an element.
// Returns removed element pointer if found, nil if not found.
func (list *SkipList) Remove(key interface{}) *Element {
    score := getScore(key, list.reversed)
    prevs := list.getPrevElementNodes(key, score)

    // found the element, remove it
    if element := prevs[0].next[0]; element != nil && !list.keyFunc.Compare(element.key, key) {
        for k, v := range element.next {
            prevs[k].next[k] = v
        }

        list.length--
        return element
    }

    return nil
}

func (list *SkipList) getPrevElementNodes(key interface{}, score float64) []*elementNode {
    var prev *elementNode = &list.elementNode
    var next *Element

    prevs := list.prevNodesCache

    for i := list.level - 1; i >= 0; i-- {
        next = prev.next[i]

        for next != nil &&
            (score > next.score || (score == next.score && list.keyFunc.Compare(key, next.key))) {
            prev = &next.elementNode
            next = next.next[i]
        }

        prevs[i] = prev
    }

    return prevs
}

// Gets current max level value.
func (list *SkipList) MaxLevel() int {
    return list.level
}

// Changes skip list max level.
// If level is not greater than 0, just panic.
func (list *SkipList) SetMaxLevel(level int) (old int) {
    if level <= 0 {
        panic("invalid argument to SetLevel")
    }

    old, list.level = list.level, level

    if old == level {
        return
    }

    if old > level {
        list.next = list.next[:level]
        list.prevNodesCache = list.prevNodesCache[:level]
        return
    }

    next := make([]*Element, level)
    copy(next, list.next)
    list.next = next
    list.prevNodesCache = make([]*elementNode, level)

    return
}

func (list *SkipList) randLevel() int {
    l := 1

    for ((list.randSource.Int63() >> 32) & 0xFFFF) < PROPABILITY {
        l++
    }

    if l > list.level {
        l = list.level
    }

    return l
}

func getScore(key interface{}, reversed bool) (score float64) {
    switch t := key.(type) {
    case []byte:
        data := []byte(t)
        l := len(data)

        // only use first 8 bytes
        if l > 7 {
            data = data[:7]
            l = 7
        }

        var result uint64

        for i := 0; i < l; i++ {
            result |= uint64(data[i])
            result <<= 8
        }

        result <<= uint(7-l) * 8
        score = float64(result)

    case float32:
        score = float64(t)

    case float64:
        score = t

    case int:
        score = float64(t)

    case int16:
        score = float64(t)

    case int32:
        score = float64(t)

    case int64:
        score = float64(t)

    case int8:
        score = float64(t)

    case string:
        data := []byte(t)
        l := len(data)

        // only use first 8 bytes
        if l > 7 {
            data = data[:7]
            l = 7
        }

        var result uint64

        for i := 0; i < l; i++ {
            result |= uint64(data[i])
            result <<= 8
        }

        result <<= uint(7-l) * 8
        score = float64(result)

    case uint:
        score = float64(t)

    case uint16:
        score = float64(t)

    case uint32:
        score = float64(t)

    case uint64:
        score = float64(t)

    case uint8:
        score = float64(t)

    case uintptr:
        score = float64(t)

    case Scorable:
        score = t.Score()
    }

    if reversed {
        score = -score
    }

    return
}
<commit_msg>fix scoring bug. thanks @glenn-brown<commit_after>// A golang Skip List Implementation.
// https://github.com/huandu/skiplist/
//
// Copyright 2011, Huan Du
// Licensed under the MIT license
// https://github.com/huandu/skiplist/blob/master/LICENSE

// Package skiplist provides a go implementation for skip list.
// About skip list: http://en.wikipedia.org/wiki/Skip_list
//
// Skip list is basically an ordered set.
// Following code creates a skip list with int key and adds some values.
//     list := skiplist.New(skiplist.Int)
//
//     // adds some elements
//     list.Set(20, "Hello")
//     list.Set(10, "World")
//     list.Set(40, true)         // value can be any type
//     list.Set(40, 1000)         // replace last element with new value
//
//     // try to find one
//     e := list.Get(10)           // value is the Element with key 10
//     _ = e.Value.(string)        // it's the "World". remember to do type cast
//     v, ok := list.GetValue(20)  // directly get value. ok is false if not exists
//     v2 := list.MustGetValue(10) // directly get value. panic if key doesn't exist
//     notFound := list.Get(15)    // returns nil if key is not found
//
//     // remove element
//     old := list.Remove(40)     // remove found element and returns its pointer
//                                // returns nil if key is not found
//
//     // re-init list. it will make the list empty.
//     list.Init()
//
// Skip list elements have random number of next pointers. The max number (say
// "max level") is configurable.
//
// The variable skiplist.DefaultMaxLevel is controlling global default.
// Changing it will not affect created lists.
//     skiplist.DefaultMaxLevel = 24  // change default to 24
// Max level of a created list can also be changed even if it's not empty.
//     list.SetMaxLevel(10)
// Remember the side effect when changing this max level value.
// See its wikipedia page for more details.
//
// Most comparable built-in types are pre-defined in skiplist, including
//     byte []byte float32 float64 int int16 int32 int64 int8
//     rune string uint uint16 uint32 uint64 uint8 uintptr
// Pre-defined compare function name is similar to the type name, e.g.
// skiplist.Float32 is for float32 key. A special case is skiplist.Bytes is for []byte.
// These functions order key from small to big (say "ascending order").
// There are also reserved order functions with name like skiplist.IntDesc.
//
// For key types out of the pre-defined list, one can write a custom compare function.
//     type GreaterThanFunc func (lhs, rhs interface{}) bool
// Such compare function returns true if lhs > rhs. Note that, if lhs == rhs, compare
// function (let the name is "foo") must act as following.
//     // if lhs == rhs, following expression must be true
//     foo(lhs, rhs) == false && foo(rhs, lhs) == false
// There is another func type named LessThanFunc. It works similar to GreaterThanFunc,
// except the order is big to small.
// Here is a sample to write a compare func.
//     type Foo struct {
//         value int
//     }
//
//     // it generates a score on a given key.
//     // if key1 > key2, then there must be key1.Score() >= key2.Score().
//     // it's optional. it's worth implementing if call compare func is quite expensive.
//     func (f *Foo) Score() float64 {
//         return float64(f.value)
//     }
//
//     var greater skiplist.GreaterThanFunc = func(lhs, rhs interface{}) {
//         return lhs.(Foo).value > rhs.(Foo).value
//     }
//     list := skiplist.New(greater)
//
//     // descending version is a bit different. mind the func type.
//     var less skiplist.LessThanFunc = func(lhs, rhs interface{}) {
//         return lhs.(Foo).value < rhs.(Foo).value
//     })
//     list := skiplist.New(less)
//
// Skiplist uses global rand source in math/rand by default. This rand source acquires a
// lock when generating random number. Replacing it with a lock-free rand source can provide
// slightly better performance. Use SkipList.SetRandSource to change rand source.
package skiplist

import (
    "math/rand"
)

// Creates a new skiplist.
// keyFunc is a func checking the "greater than" logic.
// If k1 equals k2, keyFunc(k1, k2) and keyFunc(k2, k1) must both be false.
// For built-in types, keyFunc can be found in skiplist package.
// For instance, skiplist.Int is for the list with int keys.
// By default, the list with built-in type key is in ascending order.
// The keyFunc named as skiplist.IntDesc is for descending key order list.
func New(keyFunc Comparable) *SkipList {
    if DefaultMaxLevel <= 0 {
        panic("skiplist default level must not be zero or negative")
    }

    return &SkipList{
        elementNode:    elementNode{next: make([]*Element, DefaultMaxLevel)},
        prevNodesCache: make([]*elementNode, DefaultMaxLevel),
        level:          DefaultMaxLevel,
        keyFunc:        keyFunc,
        randSource:     defaultSource,
        reversed:       keyFunc.Descending(),
    }
}

// Resets a skiplist and discards all exists elements.
func (list *SkipList) Init() *SkipList {
    list.next = make([]*Element, list.level)
    list.length = 0
    return list
}

// Sets a new rand source.
//
// Skiplist uses global rand defined in math/rand by default.
// The default rand acquires a global mutex before generating any number.
// It's not necessary if the skiplist is well protected by caller.
func (list *SkipList) SetRandSource(source rand.Source) {
    list.randSource = source
}

// Gets the first element.
func (list *SkipList) Front() *Element {
    return list.next[0]
}

// Gets list length.
func (list *SkipList) Len() int {
    return list.length
}

// Sets a value in the list with key.
// If the key exists, change element value to the new one.
// Returns new element pointer.
func (list *SkipList) Set(key, value interface{}) *Element {
    var element *Element

    score := getScore(key, list.reversed)
    prevs := list.getPrevElementNodes(key, score)

    // found an element with the same key, replace its value
    if element = prevs[0].next[0]; element != nil && !list.keyFunc.Compare(element.key, key) {
        element.Value = value
        return element
    }

    element = &Element{
        elementNode: elementNode{
            next: make([]*Element, list.randLevel()),
        },
        key:   key,
        score: score,
        Value: value,
    }

    for i := range element.next {
        element.next[i] = prevs[i].next[i]
        prevs[i].next[i] = element
    }

    list.length++
    return element
}

// Gets an element.
// Returns element pointer if found, nil if not found.
func (list *SkipList) Get(key interface{}) *Element {
    var prev *elementNode = &list.elementNode
    var next *Element
    score := getScore(key, list.reversed)

    for i := list.level - 1; i >= 0; i-- {
        next = prev.next[i]

        for next != nil &&
            (score > next.score || (score == next.score && list.keyFunc.Compare(key, next.key))) {
            prev = &next.elementNode
            next = next.next[i]
        }
    }

    if next != nil && score == next.score && !list.keyFunc.Compare(next.key, key) {
        return next
    }

    return nil
}

// Gets a value. It's a short hand for Get().Value.
// Returns value and its existence status.
func (list *SkipList) GetValue(key interface{}) (interface{}, bool) {
    element := list.Get(key)

    if element == nil {
        return nil, false
    }

    return element.Value, true
}

// Gets a value. It will panic if key doesn't exist.
// Returns value.
func (list *SkipList) MustGetValue(key interface{}) interface{} {
    element := list.Get(key)

    if element == nil {
        panic("cannot find key in skiplist")
    }

    return element.Value
}

// Removes an element.
// Returns removed element pointer if found, nil if not found.
func (list *SkipList) Remove(key interface{}) *Element {
    score := getScore(key, list.reversed)
    prevs := list.getPrevElementNodes(key, score)

    // found the element, remove it
    if element := prevs[0].next[0]; element != nil && !list.keyFunc.Compare(element.key, key) {
        for k, v := range element.next {
            prevs[k].next[k] = v
        }

        list.length--
        return element
    }

    return nil
}

func (list *SkipList) getPrevElementNodes(key interface{}, score float64) []*elementNode {
    var prev *elementNode = &list.elementNode
    var next *Element

    prevs := list.prevNodesCache

    for i := list.level - 1; i >= 0; i-- {
        next = prev.next[i]

        for next != nil &&
            (score > next.score || (score == next.score && list.keyFunc.Compare(key, next.key))) {
            prev = &next.elementNode
            next = next.next[i]
        }

        prevs[i] = prev
    }

    return prevs
}

// Gets current max level value.
func (list *SkipList) MaxLevel() int {
    return list.level
}

// Changes skip list max level.
// If level is not greater than 0, just panic.
func (list *SkipList) SetMaxLevel(level int) (old int) {
    if level <= 0 {
        panic("invalid argument to SetLevel")
    }

    old, list.level = list.level, level

    if old == level {
        return
    }

    if old > level {
        list.next = list.next[:level]
        list.prevNodesCache = list.prevNodesCache[:level]
        return
    }

    next := make([]*Element, level)
    copy(next, list.next)
    list.next = next
    list.prevNodesCache = make([]*elementNode, level)

    return
}

func (list *SkipList) randLevel() int {
    l := 1

    for ((list.randSource.Int63() >> 32) & 0xFFFF) < PROPABILITY {
        l++
    }

    if l > list.level {
        l = list.level
    }

    return l
}

func getScore(key interface{}, reversed bool) (score float64) {
    switch t := key.(type) {
    case []byte:
        var result uint64
        data := []byte(t)
        l := len(data)

        // only use first 8 bytes
        if l > 8 {
            l = 8
        }

        for i := 0; i < l; i++ {
            result |= uint64(data[i]) << uint(8 * (7 - i))
        }

        score = float64(result)

    case float32:
        score = float64(t)

    case float64:
        score = t

    case int:
        score = float64(t)

    case int16:
        score = float64(t)

    case int32:
        score = float64(t)

    case int64:
        score = float64(t)

    case int8:
        score = float64(t)

    case string:
        var result uint64
        data := []byte(t)
        l := len(data)

        // only use first 8 bytes
        if l > 8 {
            l = 8
        }

        for i := 0; i < l; i++ {
            result |= uint64(data[i]) << uint(8 * (7 - i))
        }

        score = float64(result)

    case uint:
        score = float64(t)

    case uint16:
        score = float64(t)

    case uint32:
        score = float64(t)

    case uint64:
        score = float64(t)

    case uint8:
        score = float64(t)

    case uintptr:
        score = float64(t)

    case Scorable:
        score = t.Score()
    }

    if reversed {
        score = -score
    }

    return
}
<|endoftext|>


<commit_before>package morphemekor

import (
	"bytes"
	"strings"
)

// bs := []byte{}
// b := New().Add("morphemekor")
// for {
// 	n, err := b.ReadByte()
// 	if err != nil {
// 		log.Println(err)
// 		break
// 	}
// 	bs = append(bs, n)
// }
// fmt.Println(bs)
// fmt.Println(string(bs))

// Stream is a stream of bytes.
type Stream struct {
	*bytes.Buffer
}

// New returns *morphemekor.Stream.
func New() *Stream {
	out := new(Stream)
	out.Buffer = new(bytes.Buffer)
	return out
}

// Init initializes Stream.
func (s *Stream) Init() *Stream {
	// *s = *New()
	s.Buffer.Reset()
	return s
}

// Add adds a string to the Stream.
func (s *Stream) Add(str string) *Stream {
	s.Buffer.WriteString(str)
	return s
}

// Get returns the string from the String.
func (s *Stream) Get() string {
	return s.Buffer.String()
}

// Segment segments Korean with morphemic approach.
func Segment(str string) string {
	if len(str)/3 < 4 {
		return str
	}
	isInflection1 := make(map[string]bool)
	for _, elem := range strings.Split("이히리기었렀", "") {
		isInflection1[elem] = true
	}
	isTargetEnding1 := make(map[string]bool)
	for _, elem := range strings.Split("다,고,지,마,지", ",") {
		isTargetEnding1[elem] = true
	}
	isTargetEnding2 := make(map[string]bool)
	for _, elem := range strings.Split("지만", ",") {
		isTargetEnding2[elem] = true
	}
	isPostposition1 := make(map[string]bool)
	for _, elem := range strings.Split("은,는,이,가,도,을,를,께,로,의,와,과", ",") {
		isPostposition1[elem] = true
	}
	isPostposition2 := make(map[string]bool)
	for _, elem := range strings.Split("에,한,으", ",") {
		isPostposition2[elem] = true
	}
	isPostposition3 := make(map[string]bool)
	for _, elem := range strings.Split("에게,에서,한테,으로", ",") {
		isPostposition3[elem] = true
	}
	buffer := New()
	characters := strings.Split(str, "")
	marker := 0 // another index in word phrase
	for idx := 0; idx < len(characters); idx++ {
		elem := characters[idx]
		if idx == 0 || marker == 0 {
			buffer.Add(elem)
			marker++
			continue
		}
		if elem == " " {
			buffer.Add(elem)
			marker++
			continue
		}
		// Check if the character is inflection
		if _, ok := isInflection1[elem]; ok {
			buffer.Add(elem)
			// Check the next chracter
			if len(characters) > idx+1 {
				nelem1 := characters[idx+1]
				if _, ok := isTargetEnding1[nelem1]; ok {
					// Check the second next chracter
					if len(characters) > idx+2 {
						nelem2 := characters[idx+2]
						if _, ok := isTargetEnding2[nelem1+nelem2]; ok {
							buffer.Add(nelem1 + nelem2)
							buffer.Add(" ")
							idx = idx + 2
							marker = 0
							continue
						}
					}
					buffer.Add(nelem1)
					buffer.Add(" ")
					idx++
					marker = 0
					continue
				}
			}
			marker++
			continue
		}
		// Check if the character is postposition
		if _, ok := isPostposition1[elem]; ok {
			buffer.Add(elem)
			buffer.Add(" ")
			marker = 0
			continue
		} else if _, ok := isPostposition2[elem]; ok {
			if len(characters) > idx+1 {
				nelem2 := characters[idx+1]
				if _, ok := isPostposition3[elem+nelem2]; ok {
					buffer.Add(elem + nelem2)
					buffer.Add(" ")
					idx++
					marker = 0
					continue
				}
			}
		}
		buffer.Add(elem)
		marker++
	}
	return strings.TrimSpace(buffer.Get())
}
<commit_msg>Update paper reference<commit_after>package morphemekor

import (
	"bytes"
	"strings"
)

// bs := []byte{}
// b := New().Add("morphemekor")
// for {
// 	n, err := b.ReadByte()
// 	if err != nil {
// 		log.Println(err)
// 		break
// 	}
// 	bs = append(bs, n)
// }
// fmt.Println(bs)
// fmt.Println(string(bs))

// Stream is a stream of bytes.
type Stream struct {
	*bytes.Buffer
}

// New returns *morphemekor.Stream.
func New() *Stream {
	out := new(Stream)
	out.Buffer = new(bytes.Buffer)
	return out
}

// Init initializes Stream.
func (s *Stream) Init() *Stream {
	// *s = *New()
	s.Buffer.Reset()
	return s
}

// Add adds a string to the Stream.
func (s *Stream) Add(str string) *Stream {
	s.Buffer.WriteString(str)
	return s
}

// Get returns the string from the String.
func (s *Stream) Get() string {
	return s.Buffer.String()
}

// Segment segments Korean with morphemic approach.
// (Reference: https://github.com/gyuho/learn_other/blob/master/ling105/gyuho.pdf)
func Segment(str string) string {
	if len(str)/3 < 4 {
		return str
	}
	isInflection1 := make(map[string]bool)
	for _, elem := range strings.Split("이히리기었렀", "") {
		isInflection1[elem] = true
	}
	isTargetEnding1 := make(map[string]bool)
	for _, elem := range strings.Split("다,고,지,마,지", ",") {
		isTargetEnding1[elem] = true
	}
	isTargetEnding2 := make(map[string]bool)
	for _, elem := range strings.Split("지만", ",") {
		isTargetEnding2[elem] = true
	}
	isPostposition1 := make(map[string]bool)
	for _, elem := range strings.Split("은,는,이,가,도,을,를,께,로,의,와,과", ",") {
		isPostposition1[elem] = true
	}
	isPostposition2 := make(map[string]bool)
	for _, elem := range strings.Split("에,한,으", ",") {
		isPostposition2[elem] = true
	}
	isPostposition3 := make(map[string]bool)
	for _, elem := range strings.Split("에게,에서,한테,으로", ",") {
		isPostposition3[elem] = true
	}
	buffer := New()
	characters := strings.Split(str, "")
	marker := 0 // another index in word phrase
	for idx := 0; idx < len(characters); idx++ {
		elem := characters[idx]
		if idx == 0 || marker == 0 {
			buffer.Add(elem)
			marker++
			continue
		}
		if elem == " " {
			buffer.Add(elem)
			marker++
			continue
		}
		// Check if the character is inflection
		if _, ok := isInflection1[elem]; ok {
			buffer.Add(elem)
			// Check the next chracter
			if len(characters) > idx+1 {
				nelem1 := characters[idx+1]
				if _, ok := isTargetEnding1[nelem1]; ok {
					// Check the second next chracter
					if len(characters) > idx+2 {
						nelem2 := characters[idx+2]
						if _, ok := isTargetEnding2[nelem1+nelem2]; ok {
							buffer.Add(nelem1 + nelem2)
							buffer.Add(" ")
							idx = idx + 2
							marker = 0
							continue
						}
					}
					buffer.Add(nelem1)
					buffer.Add(" ")
					idx++
					marker = 0
					continue
				}
			}
			marker++
			continue
		}
		// Check if the character is postposition
		if _, ok := isPostposition1[elem]; ok {
			buffer.Add(elem)
			buffer.Add(" ")
			marker = 0
			continue
		} else if _, ok := isPostposition2[elem]; ok {
			if len(characters) > idx+1 {
				nelem2 := characters[idx+1]
				if _, ok := isPostposition3[elem+nelem2]; ok {
					buffer.Add(elem + nelem2)
					buffer.Add(" ")
					idx++
					marker = 0
					continue
				}
			}
		}
		buffer.Add(elem)
		marker++
	}
	return strings.TrimSpace(buffer.Get())
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/commands/admin"
	"github.com/urfave/cli"
)

func generateHelp([]cli.Command) {
	for idx, cmd := range commands {
		switch cmd.Name {
		case "admin":
			commands[idx].Description = cmd.Description + "" + app.GenerateCommandsHelp(admin.Commands)
		case "commands":
			commands[idx].Description = cmd.Description + "" + app.GenerateCommandsHelp(commands)
		}
	}
}
<commit_msg>Fix lint issue with generateHelp<commit_after>package main

import (
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/commands/admin"
	"github.com/urfave/cli"
)

func generateHelp([]cli.Command) {
	for idx, cmd := range commands {
		switch cmd.Name {
		case "admin":
			commands[idx].Description = cmd.Description + app.GenerateCommandsHelp(admin.Commands)
		case "commands":
			commands[idx].Description = cmd.Description + app.GenerateCommandsHelp(commands)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strings"
	"time"

	"code.cloudfoundry.org/guardian/guardiancmd"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/localip"
	"github.com/concourse/atc"
	"github.com/jessevdk/go-flags"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
)

type GardenBackend struct {
	guardiancmd.ServerCommand

	DNS DNSConfig `group:"DNS Proxy Configuration" namespace:"dns-proxy"`
}

func (cmd WorkerCommand) lessenRequirements(command *flags.Command) {
	command.FindOptionByLongName("garden-bind-port").Default = []string{"7777"}

	// configured as work-dir/depot
	command.FindOptionByLongName("garden-depot").Required = false

	// un-configure graph (default /var/gdn/graph)
	command.FindOptionByLongName("garden-graph").Required = false
	command.FindOptionByLongName("garden-graph").Default = []string{}

	// these are provided as assets embedded in the 'concourse' binary
	command.FindOptionByLongName("garden-runc-bin").Required = false
	command.FindOptionByLongName("garden-dadoo-bin").Required = false
	command.FindOptionByLongName("garden-init-bin").Required = false
	command.FindOptionByLongName("garden-nstar-bin").Required = false
	command.FindOptionByLongName("garden-tar-bin").Required = false

	// configured as work-dir/volumes
	command.FindOptionByLongName("baggageclaim-volumes").Required = false
}

func (cmd *WorkerCommand) gardenRunner(logger lager.Logger, hasAssets bool) (atc.Worker, ifrit.Runner, error) {
	err := cmd.checkRoot()
	if err != nil {
		return atc.Worker{}, nil, err
	}

	depotDir := filepath.Join(cmd.WorkDir.Path(), "depot")

	// must be readable by other users so unprivileged containers can run their
	// own `initc' process
	err = os.MkdirAll(depotDir, 0755)
	if err != nil {
		return atc.Worker{}, nil, err
	}

	cmd.Garden.Server.BindIP = guardiancmd.IPFlag(cmd.BindIP)
	cmd.Garden.Containers.Dir = depotDir

	cmd.Garden.Network.AllowHostAccess = true

	worker := atc.Worker{
		Platform: "linux",
		Tags:     cmd.Tags,
		Team:     cmd.TeamName,

		HTTPProxyURL:  cmd.HTTPProxy.String(),
		HTTPSProxyURL: cmd.HTTPSProxy.String(),
		NoProxy:       strings.Join(cmd.NoProxy, ","),
		StartTime:     time.Now().Unix(),
	}

	if hasAssets {
		cmd.Garden.Runtime.Plugin = cmd.assetPath("bin", "runc")
		cmd.Garden.Bin.Dadoo = guardiancmd.FileFlag(cmd.assetPath("bin", "dadoo"))
		cmd.Garden.Bin.Init = guardiancmd.FileFlag(cmd.assetPath("bin", "init"))
		cmd.Garden.Bin.NSTar = guardiancmd.FileFlag(cmd.assetPath("bin", "nstar"))
		cmd.Garden.Bin.Tar = guardiancmd.FileFlag(cmd.assetPath("bin", "tar"))

		iptablesDir := cmd.assetPath("iptables")
		cmd.Garden.Bin.IPTables = guardiancmd.FileFlag(filepath.Join(iptablesDir, "sbin", "iptables"))

		worker.ResourceTypes, err = cmd.extractResources(logger.Session("extract-resources"))
		if err != nil {
			return atc.Worker{}, nil, err
		}
	}

	worker.Name, err = cmd.workerName()
	if err != nil {
		return atc.Worker{}, nil, err
	}

	members := grouper.Members{
		{
			Name:   "garden-runc",
			Runner: &cmd.Garden.ServerCommand,
		},
	}

	if cmd.Garden.DNS.Enable {
		dnsProxyRunner, err := cmd.dnsProxyRunner(logger.Session("dns-proxy"))
		if err != nil {
			return atc.Worker{}, nil, err
		}

		lip, err := localip.LocalIP()
		if err != nil {
			return atc.Worker{}, nil, err
		}

		cmd.Garden.Network.AdditionalDNSServers = append(
			cmd.Garden.Network.AdditionalDNSServers,
			guardiancmd.IPFlag(net.ParseIP(lip)),
		)

		members = append(members, grouper.Member{
			Name:   "dns-proxy",
			Runner: dnsProxyRunner,
		})
	}

	return worker, grouper.NewParallel(os.Interrupt, members), nil
}

func (cmd *WorkerCommand) extractResources(logger lager.Logger) ([]atc.WorkerResourceType, error) {
	var resourceTypes []atc.WorkerResourceType

	resourcesDir := cmd.assetPath("resources")

	infos, err := ioutil.ReadDir(resourcesDir)
	if err != nil {
		logger.Error("failed-to-list-resource-assets", err)
		return nil, err
	}

	for _, info := range infos {
		resourceType := info.Name()

		workerResourceType, err := cmd.extractResource(
			logger.Session("extract", lager.Data{"resource-type": resourceType}),
			resourcesDir,
			resourceType,
		)
		if err != nil {
			logger.Error("failed-to-extract-resource", err)
			return nil, err
		}

		resourceTypes = append(resourceTypes, workerResourceType)
	}

	return resourceTypes, nil
}

func (cmd *WorkerCommand) extractResource(
	logger lager.Logger,
	resourcesDir string,
	resourceType string,
) (atc.WorkerResourceType, error) {
	resourceImagesDir := cmd.assetPath("resource-images")
	tarBin := cmd.assetPath("bin", "tar")

	archive := filepath.Join(resourcesDir, resourceType, "rootfs.tar.gz")

	extractedDir := filepath.Join(resourceImagesDir, resourceType)

	rootfsDir := filepath.Join(extractedDir, "rootfs")
	okMarker := filepath.Join(extractedDir, "ok")

	var version string
	versionFile, err := os.Open(filepath.Join(resourcesDir, resourceType, "version"))
	if err != nil {
		logger.Error("failed-to-read-version", err)
		return atc.WorkerResourceType{}, err
	}

	_, err = fmt.Fscanf(versionFile, "%s", &version)
	if err != nil {
		logger.Error("failed-to-parse-version", err)
		return atc.WorkerResourceType{}, err
	}

	defer versionFile.Close()

	var privileged bool
	_, err = os.Stat(filepath.Join(resourcesDir, resourceType, "privileged"))
	if err == nil {
		privileged = true
	}

	_, err = os.Stat(okMarker)
	if err == nil {
		logger.Info("already-extracted")
	} else {
		logger.Info("extracting")

		err := os.RemoveAll(rootfsDir)
		if err != nil {
			logger.Error("failed-to-clear-out-existing-rootfs", err)
			return atc.WorkerResourceType{}, err
		}

		err = os.MkdirAll(rootfsDir, 0755)
		if err != nil {
			logger.Error("failed-to-create-rootfs-dir", err)
			return atc.WorkerResourceType{}, err
		}

		tar := exec.Command(tarBin, "-zxf", archive, "-C", rootfsDir)

		output, err := tar.CombinedOutput()
		if err != nil {
			logger.Error("failed-to-extract-resource", err, lager.Data{
				"output": string(output),
			})
			return atc.WorkerResourceType{}, err
		}

		ok, err := os.Create(okMarker)
		if err != nil {
			logger.Error("failed-to-create-ok-marker", err)
			return atc.WorkerResourceType{}, err
		}

		err = ok.Close()
		if err != nil {
			logger.Error("failed-to-close-ok-marker", err)
			return atc.WorkerResourceType{}, err
		}
	}

	return atc.WorkerResourceType{
		Type:       resourceType,
		Image:      rootfsDir,
		Version:    version,
		Privileged: privileged,
	}, nil
}

var ErrNotRoot = errors.New("worker must be run as root")

func (cmd *WorkerCommand) checkRoot() error {
	currentUser, err := user.Current()
	if err != nil {
		return err
	}

	if currentUser.Uid != "0" {
		return ErrNotRoot
	}

	return nil
}

func (cmd *WorkerCommand) dnsProxyRunner(logger lager.Logger) (ifrit.Runner, error) {
	server, err := cmd.Garden.DNS.Server()
	if err != nil {
		return nil, err
	}

	return ifrit.RunFunc(func(signals <-chan os.Signal, ready chan<- struct{}) error {
		server.NotifyStartedFunc = func() {
			close(ready)
			logger.Info("started")
		}

		serveErr := make(chan error, 1)

		go func() {
			serveErr <- server.ListenAndServe()
		}()

		for {
			select {
			case err := <-serveErr:
				return err
			case <-signals:
				server.Shutdown()
			}
		}
	}), nil
}
<commit_msg>fix flag for new garden command<commit_after>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strings"
	"time"

	"code.cloudfoundry.org/guardian/guardiancmd"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/localip"
	"github.com/concourse/atc"
	"github.com/jessevdk/go-flags"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
)

type GardenBackend struct {
	guardiancmd.ServerCommand

	DNS DNSConfig `group:"DNS Proxy Configuration" namespace:"dns-proxy"`
}

func (cmd WorkerCommand) lessenRequirements(command *flags.Command) {
	command.FindOptionByLongName("garden-bind-port").Default = []string{"7777"}

	// configured as work-dir/depot
	command.FindOptionByLongName("garden-depot").Required = false

	// un-configure graph (default /var/gdn/graph)
	command.FindOptionByLongName("garden-graph").Required = false
	command.FindOptionByLongName("garden-graph").Default = []string{}

	// these are provided as assets embedded in the 'concourse' binary
	command.FindOptionByLongName("garden-runtime-plugin").Required = false
	command.FindOptionByLongName("garden-dadoo-bin").Required = false
	command.FindOptionByLongName("garden-init-bin").Required = false
	command.FindOptionByLongName("garden-nstar-bin").Required = false
	command.FindOptionByLongName("garden-tar-bin").Required = false

	// configured as work-dir/volumes
	command.FindOptionByLongName("baggageclaim-volumes").Required = false
}

func (cmd *WorkerCommand) gardenRunner(logger lager.Logger, hasAssets bool) (atc.Worker, ifrit.Runner, error) {
	err := cmd.checkRoot()
	if err != nil {
		return atc.Worker{}, nil, err
	}

	depotDir := filepath.Join(cmd.WorkDir.Path(), "depot")

	// must be readable by other users so unprivileged containers can run their
	// own `initc' process
	err = os.MkdirAll(depotDir, 0755)
	if err != nil {
		return atc.Worker{}, nil, err
	}

	cmd.Garden.Server.BindIP = guardiancmd.IPFlag(cmd.BindIP)
	cmd.Garden.Containers.Dir = depotDir

	cmd.Garden.Network.AllowHostAccess = true

	worker := atc.Worker{
		Platform: "linux",
		Tags:     cmd.Tags,
		Team:     cmd.TeamName,

		HTTPProxyURL:  cmd.HTTPProxy.String(),
		HTTPSProxyURL: cmd.HTTPSProxy.String(),
		NoProxy:       strings.Join(cmd.NoProxy, ","),
		StartTime:     time.Now().Unix(),
	}

	if hasAssets {
		cmd.Garden.Runtime.Plugin = cmd.assetPath("bin", "runc")
		cmd.Garden.Bin.Dadoo = guardiancmd.FileFlag(cmd.assetPath("bin", "dadoo"))
		cmd.Garden.Bin.Init = guardiancmd.FileFlag(cmd.assetPath("bin", "init"))
		cmd.Garden.Bin.NSTar = guardiancmd.FileFlag(cmd.assetPath("bin", "nstar"))
		cmd.Garden.Bin.Tar = guardiancmd.FileFlag(cmd.assetPath("bin", "tar"))

		iptablesDir := cmd.assetPath("iptables")
		cmd.Garden.Bin.IPTables = guardiancmd.FileFlag(filepath.Join(iptablesDir, "sbin", "iptables"))

		worker.ResourceTypes, err = cmd.extractResources(logger.Session("extract-resources"))
		if err != nil {
			return atc.Worker{}, nil, err
		}
	}

	worker.Name, err = cmd.workerName()
	if err != nil {
		return atc.Worker{}, nil, err
	}

	members := grouper.Members{
		{
			Name:   "garden-runc",
			Runner: &cmd.Garden.ServerCommand,
		},
	}

	if cmd.Garden.DNS.Enable {
		dnsProxyRunner, err := cmd.dnsProxyRunner(logger.Session("dns-proxy"))
		if err != nil {
			return atc.Worker{}, nil, err
		}

		lip, err := localip.LocalIP()
		if err != nil {
			return atc.Worker{}, nil, err
		}

		cmd.Garden.Network.AdditionalDNSServers = append(
			cmd.Garden.Network.AdditionalDNSServers,
			guardiancmd.IPFlag(net.ParseIP(lip)),
		)

		members = append(members, grouper.Member{
			Name:   "dns-proxy",
			Runner: dnsProxyRunner,
		})
	}

	return worker, grouper.NewParallel(os.Interrupt, members), nil
}

func (cmd *WorkerCommand) extractResources(logger lager.Logger) ([]atc.WorkerResourceType, error) {
	var resourceTypes []atc.WorkerResourceType

	resourcesDir := cmd.assetPath("resources")

	infos, err := ioutil.ReadDir(resourcesDir)
	if err != nil {
		logger.Error("failed-to-list-resource-assets", err)
		return nil, err
	}

	for _, info := range infos {
		resourceType := info.Name()

		workerResourceType, err := cmd.extractResource(
			logger.Session("extract", lager.Data{"resource-type": resourceType}),
			resourcesDir,
			resourceType,
		)
		if err != nil {
			logger.Error("failed-to-extract-resource", err)
			return nil, err
		}

		resourceTypes = append(resourceTypes, workerResourceType)
	}

	return resourceTypes, nil
}

func (cmd *WorkerCommand) extractResource(
	logger lager.Logger,
	resourcesDir string,
	resourceType string,
) (atc.WorkerResourceType, error) {
	resourceImagesDir := cmd.assetPath("resource-images")
	tarBin := cmd.assetPath("bin", "tar")

	archive := filepath.Join(resourcesDir, resourceType, "rootfs.tar.gz")

	extractedDir := filepath.Join(resourceImagesDir, resourceType)

	rootfsDir := filepath.Join(extractedDir, "rootfs")
	okMarker := filepath.Join(extractedDir, "ok")

	var version string
	versionFile, err := os.Open(filepath.Join(resourcesDir, resourceType, "version"))
	if err != nil {
		logger.Error("failed-to-read-version", err)
		return atc.WorkerResourceType{}, err
	}

	_, err = fmt.Fscanf(versionFile, "%s", &version)
	if err != nil {
		logger.Error("failed-to-parse-version", err)
		return atc.WorkerResourceType{}, err
	}

	defer versionFile.Close()

	var privileged bool
	_, err = os.Stat(filepath.Join(resourcesDir, resourceType, "privileged"))
	if err == nil {
		privileged = true
	}

	_, err = os.Stat(okMarker)
	if err == nil {
		logger.Info("already-extracted")
	} else {
		logger.Info("extracting")

		err := os.RemoveAll(rootfsDir)
		if err != nil {
			logger.Error("failed-to-clear-out-existing-rootfs", err)
			return atc.WorkerResourceType{}, err
		}

		err = os.MkdirAll(rootfsDir, 0755)
		if err != nil {
			logger.Error("failed-to-create-rootfs-dir", err)
			return atc.WorkerResourceType{}, err
		}

		tar := exec.Command(tarBin, "-zxf", archive, "-C", rootfsDir)

		output, err := tar.CombinedOutput()
		if err != nil {
			logger.Error("failed-to-extract-resource", err, lager.Data{
				"output": string(output),
			})
			return atc.WorkerResourceType{}, err
		}

		ok, err := os.Create(okMarker)
		if err != nil {
			logger.Error("failed-to-create-ok-marker", err)
			return atc.WorkerResourceType{}, err
		}

		err = ok.Close()
		if err != nil {
			logger.Error("failed-to-close-ok-marker", err)
			return atc.WorkerResourceType{}, err
		}
	}

	return atc.WorkerResourceType{
		Type:       resourceType,
		Image:      rootfsDir,
		Version:    version,
		Privileged: privileged,
	}, nil
}

var ErrNotRoot = errors.New("worker must be run as root")

func (cmd *WorkerCommand) checkRoot() error {
	currentUser, err := user.Current()
	if err != nil {
		return err
	}

	if currentUser.Uid != "0" {
		return ErrNotRoot
	}

	return nil
}

func (cmd *WorkerCommand) dnsProxyRunner(logger lager.Logger) (ifrit.Runner, error) {
	server, err := cmd.Garden.DNS.Server()
	if err != nil {
		return nil, err
	}

	return ifrit.RunFunc(func(signals <-chan os.Signal, ready chan<- struct{}) error {
		server.NotifyStartedFunc = func() {
			close(ready)
			logger.Info("started")
		}

		serveErr := make(chan error, 1)

		go func() {
			serveErr <- server.ListenAndServe()
		}()

		for {
			select {
			case err := <-serveErr:
				return err
			case <-signals:
				server.Shutdown()
			}
		}
	}), nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strings"
	"time"

	"code.cloudfoundry.org/guardian/guardiancmd"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/localip"
	"github.com/concourse/atc"
	"github.com/jessevdk/go-flags"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
)

type Certs struct {
	Dir string `long:"certs-dir" description:"Directory to use when creating the resource certificates volume."`
}

type GardenBackend struct {
	guardiancmd.ServerCommand

	DNS DNSConfig `group:"DNS Proxy Configuration" namespace:"dns-proxy"`
}

func (cmd WorkerCommand) lessenRequirements(prefix string, command *flags.Command) {
	command.FindOptionByLongName(prefix + "garden-bind-port").Default = []string{"7777"}

	// configured as work-dir/depot
	command.FindOptionByLongName(prefix + "garden-depot").Required = false

	// un-configure graph (default /var/gdn/graph)
	command.FindOptionByLongName(prefix + "garden-graph").Required = false
	command.FindOptionByLongName(prefix + "garden-graph").Default = []string{}

	// these are provided as assets embedded in the 'concourse' binary
	command.FindOptionByLongName(prefix + "garden-runtime-plugin").Required = false
	command.FindOptionByLongName(prefix + "garden-dadoo-bin").Required = false
	command.FindOptionByLongName(prefix + "garden-init-bin").Required = false
	command.FindOptionByLongName(prefix + "garden-nstar-bin").Required = false
	command.FindOptionByLongName(prefix + "garden-tar-bin").Required = false

	// configured as work-dir/volumes
	command.FindOptionByLongName(prefix + "baggageclaim-volumes").Required = false
}

func (cmd *WorkerCommand) gardenRunner(logger lager.Logger, hasAssets bool) (atc.Worker, ifrit.Runner, error) {
	err := cmd.checkRoot()
	if err != nil {
		return atc.Worker{}, nil, err
	}

	depotDir := filepath.Join(cmd.WorkDir.Path(), "depot")

	// must be readable by other users so unprivileged containers can run their
	// own `initc' process
	err = os.MkdirAll(depotDir, 0755)
	if err != nil {
		return atc.Worker{}, nil, err
	}

	cmd.Garden.Server.BindIP = guardiancmd.IPFlag(cmd.BindIP.IP)
	cmd.Garden.Containers.Dir = depotDir

	cmd.Garden.Network.AllowHostAccess = true

	worker := atc.Worker{
		Platform:  "linux",
		Tags:      cmd.Tags,
		Team:      cmd.TeamName,
		CertsPath: &cmd.Certs.Dir,

		HTTPProxyURL:  cmd.HTTPProxy.String(),
		HTTPSProxyURL: cmd.HTTPSProxy.String(),
		NoProxy:       strings.Join(cmd.NoProxy, ","),
		StartTime:     time.Now().Unix(),
	}

	if hasAssets {
		cmd.Garden.Runtime.Plugin = cmd.assetPath("bin", "runc")
		cmd.Garden.Bin.Dadoo = guardiancmd.FileFlag(cmd.assetPath("bin", "dadoo"))
		cmd.Garden.Bin.Init = guardiancmd.FileFlag(cmd.assetPath("bin", "init"))
		cmd.Garden.Bin.NSTar = guardiancmd.FileFlag(cmd.assetPath("bin", "nstar"))
		cmd.Garden.Bin.Tar = guardiancmd.FileFlag(cmd.assetPath("bin", "tar"))

		iptablesDir := cmd.assetPath("iptables")
		cmd.Garden.Bin.IPTables = guardiancmd.FileFlag(filepath.Join(iptablesDir, "sbin", "iptables"))

		worker.ResourceTypes, err = cmd.extractResources(logger.Session("extract-resources"))
		if err != nil {
			return atc.Worker{}, nil, err
		}
	}

	worker.Name, err = cmd.workerName()
	if err != nil {
		return atc.Worker{}, nil, err
	}

	members := grouper.Members{
		{
			Name:   "garden-runc",
			Runner: &cmd.Garden.ServerCommand,
		},
	}

	if cmd.Garden.DNS.Enable {
		dnsProxyRunner, err := cmd.dnsProxyRunner(logger.Session("dns-proxy"))
		if err != nil {
			return atc.Worker{}, nil, err
		}

		lip, err := localip.LocalIP()
		if err != nil {
			return atc.Worker{}, nil, err
		}

		cmd.Garden.Network.AdditionalDNSServers = append(
			cmd.Garden.Network.AdditionalDNSServers,
			guardiancmd.IPFlag(net.ParseIP(lip)),
		)

		members = append(members, grouper.Member{
			Name:   "dns-proxy",
			Runner: dnsProxyRunner,
		})
	}

	return worker, grouper.NewParallel(os.Interrupt, members), nil
}

func (cmd *WorkerCommand) extractResources(logger lager.Logger) ([]atc.WorkerResourceType, error) {
	var resourceTypes []atc.WorkerResourceType

	resourcesDir := cmd.assetPath("resources")

	infos, err := ioutil.ReadDir(resourcesDir)
	if err != nil {
		logger.Error("failed-to-list-resource-assets", err)
		return nil, err
	}

	for _, info := range infos {
		resourceType := info.Name()

		workerResourceType, err := cmd.extractResource(
			logger.Session("extract", lager.Data{"resource-type": resourceType}),
			resourcesDir,
			resourceType,
		)
		if err != nil {
			logger.Error("failed-to-extract-resource", err)
			return nil, err
		}

		resourceTypes = append(resourceTypes, workerResourceType)
	}

	return resourceTypes, nil
}

func (cmd *WorkerCommand) extractResource(
	logger lager.Logger,
	resourcesDir string,
	resourceType string,
) (atc.WorkerResourceType, error) {
	resourceImagesDir := cmd.assetPath("resource-images")
	tarBin := cmd.assetPath("bin", "tar")

	archive := filepath.Join(resourcesDir, resourceType, "rootfs.tar.gz")

	extractedDir := filepath.Join(resourceImagesDir, resourceType)

	rootfsDir := filepath.Join(extractedDir, "rootfs")
	okMarker := filepath.Join(extractedDir, "ok")

	var version string
	versionFile, err := os.Open(filepath.Join(resourcesDir, resourceType, "version"))
	if err != nil {
		logger.Error("failed-to-read-version", err)
		return atc.WorkerResourceType{}, err
	}

	_, err = fmt.Fscanf(versionFile, "%s", &version)
	if err != nil {
		logger.Error("failed-to-parse-version", err)
		return atc.WorkerResourceType{}, err
	}

	defer versionFile.Close()

	var privileged bool
	_, err = os.Stat(filepath.Join(resourcesDir, resourceType, "privileged"))
	if err == nil {
		privileged = true
	}

	_, err = os.Stat(okMarker)
	if err == nil {
		logger.Info("already-extracted")
	} else {
		logger.Info("extracting")

		err := os.RemoveAll(rootfsDir)
		if err != nil {
			logger.Error("failed-to-clear-out-existing-rootfs", err)
			return atc.WorkerResourceType{}, err
		}

		err = os.MkdirAll(rootfsDir, 0755)
		if err != nil {
			logger.Error("failed-to-create-rootfs-dir", err)
			return atc.WorkerResourceType{}, err
		}

		tar := exec.Command(tarBin, "-zxf", archive, "-C", rootfsDir)

		output, err := tar.CombinedOutput()
		if err != nil {
			logger.Error("failed-to-extract-resource", err, lager.Data{
				"output": string(output),
			})
			return atc.WorkerResourceType{}, err
		}

		ok, err := os.Create(okMarker)
		if err != nil {
			logger.Error("failed-to-create-ok-marker", err)
			return atc.WorkerResourceType{}, err
		}

		err = ok.Close()
		if err != nil {
			logger.Error("failed-to-close-ok-marker", err)
			return atc.WorkerResourceType{}, err
		}
	}

	return atc.WorkerResourceType{
		Type:       resourceType,
		Image:      rootfsDir,
		Version:    version,
		Privileged: privileged,
	}, nil
}

var ErrNotRoot = errors.New("worker must be run as root")

func (cmd *WorkerCommand) checkRoot() error {
	currentUser, err := user.Current()
	if err != nil {
		return err
	}

	if currentUser.Uid != "0" {
		return ErrNotRoot
	}

	return nil
}

func (cmd *WorkerCommand) dnsProxyRunner(logger lager.Logger) (ifrit.Runner, error) {
	server, err := cmd.Garden.DNS.Server()
	if err != nil {
		return nil, err
	}

	return ifrit.RunFunc(func(signals <-chan os.Signal, ready chan<- struct{}) error {
		server.NotifyStartedFunc = func() {
			close(ready)
			logger.Info("started")
		}

		serveErr := make(chan error, 1)

		go func() {
			serveErr <- server.ListenAndServe()
		}()

		for {
			select {
			case err := <-serveErr:
				return err
			case <-signals:
				server.Shutdown()
			}
		}
	}), nil
}
<commit_msg>Fix compilation issues for linux worker<commit_after>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"time"

	"code.cloudfoundry.org/guardian/guardiancmd"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/localip"
	"github.com/concourse/atc"
	"github.com/jessevdk/go-flags"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
)

type Certs struct {
	Dir string `long:"certs-dir" description:"Directory to use when creating the resource certificates volume."`
}

type GardenBackend struct {
	guardiancmd.ServerCommand

	DNS DNSConfig `group:"DNS Proxy Configuration" namespace:"dns-proxy"`
}

func (cmd WorkerCommand) lessenRequirements(prefix string, command *flags.Command) {
	command.FindOptionByLongName(prefix + "garden-bind-port").Default = []string{"7777"}

	// configured as work-dir/depot
	command.FindOptionByLongName(prefix + "garden-depot").Required = false

	// un-configure graph (default /var/gdn/graph)
	command.FindOptionByLongName(prefix + "garden-graph").Required = false
	command.FindOptionByLongName(prefix + "garden-graph").Default = []string{}

	// these are provided as assets embedded in the 'concourse' binary
	command.FindOptionByLongName(prefix + "garden-runtime-plugin").Required = false
	command.FindOptionByLongName(prefix + "garden-dadoo-bin").Required = false
	command.FindOptionByLongName(prefix + "garden-init-bin").Required = false
	command.FindOptionByLongName(prefix + "garden-nstar-bin").Required = false
	command.FindOptionByLongName(prefix + "garden-tar-bin").Required = false

	// configured as work-dir/volumes
	command.FindOptionByLongName(prefix + "baggageclaim-volumes").Required = false
}

func (cmd *WorkerCommand) gardenRunner(logger lager.Logger, hasAssets bool) (atc.Worker, ifrit.Runner, error) {
	err := cmd.checkRoot()
	if err != nil {
		return atc.Worker{}, nil, err
	}

	depotDir := filepath.Join(cmd.WorkDir.Path(), "depot")

	// must be readable by other users so unprivileged containers can run their
	// own `initc' process
	err = os.MkdirAll(depotDir, 0755)
	if err != nil {
		return atc.Worker{}, nil, err
	}

	cmd.Garden.Server.BindIP = guardiancmd.IPFlag(cmd.BindIP.IP)
	cmd.Garden.Containers.Dir = depotDir

	cmd.Garden.Network.AllowHostAccess = true

	worker := atc.Worker{
		Platform:  "linux",
		Tags:      cmd.Worker.Tags,
		Team:      cmd.Worker.TeamName,
		CertsPath: &cmd.Certs.Dir,

		HTTPProxyURL:  cmd.Worker.HTTPProxy,
		HTTPSProxyURL: cmd.Worker.HTTPSProxy,
		NoProxy:       cmd.Worker.NoProxy,
		StartTime:     time.Now().Unix(),
	}

	if hasAssets {
		cmd.Garden.Runtime.Plugin = cmd.assetPath("bin", "runc")
		cmd.Garden.Bin.Dadoo = guardiancmd.FileFlag(cmd.assetPath("bin", "dadoo"))
		cmd.Garden.Bin.Init = guardiancmd.FileFlag(cmd.assetPath("bin", "init"))
		cmd.Garden.Bin.NSTar = guardiancmd.FileFlag(cmd.assetPath("bin", "nstar"))
		cmd.Garden.Bin.Tar = guardiancmd.FileFlag(cmd.assetPath("bin", "tar"))

		iptablesDir := cmd.assetPath("iptables")
		cmd.Garden.Bin.IPTables = guardiancmd.FileFlag(filepath.Join(iptablesDir, "sbin", "iptables"))

		worker.ResourceTypes, err = cmd.extractResources(logger.Session("extract-resources"))
		if err != nil {
			return atc.Worker{}, nil, err
		}
	}

	worker.Name, err = cmd.workerName()
	if err != nil {
		return atc.Worker{}, nil, err
	}

	members := grouper.Members{
		{
			Name:   "garden-runc",
			Runner: &cmd.Garden.ServerCommand,
		},
	}

	if cmd.Garden.DNS.Enable {
		dnsProxyRunner, err := cmd.dnsProxyRunner(logger.Session("dns-proxy"))
		if err != nil {
			return atc.Worker{}, nil, err
		}

		lip, err := localip.LocalIP()
		if err != nil {
			return atc.Worker{}, nil, err
		}

		cmd.Garden.Network.AdditionalDNSServers = append(
			cmd.Garden.Network.AdditionalDNSServers,
			guardiancmd.IPFlag(net.ParseIP(lip)),
		)

		members = append(members, grouper.Member{
			Name:   "dns-proxy",
			Runner: dnsProxyRunner,
		})
	}

	return worker, grouper.NewParallel(os.Interrupt, members), nil
}

func (cmd *WorkerCommand) extractResources(logger lager.Logger) ([]atc.WorkerResourceType, error) {
	var resourceTypes []atc.WorkerResourceType

	resourcesDir := cmd.assetPath("resources")

	infos, err := ioutil.ReadDir(resourcesDir)
	if err != nil {
		logger.Error("failed-to-list-resource-assets", err)
		return nil, err
	}

	for _, info := range infos {
		resourceType := info.Name()

		workerResourceType, err := cmd.extractResource(
			logger.Session("extract", lager.Data{"resource-type": resourceType}),
			resourcesDir,
			resourceType,
		)
		if err != nil {
			logger.Error("failed-to-extract-resource", err)
			return nil, err
		}

		resourceTypes = append(resourceTypes, workerResourceType)
	}

	return resourceTypes, nil
}

func (cmd *WorkerCommand) extractResource(
	logger lager.Logger,
	resourcesDir string,
	resourceType string,
) (atc.WorkerResourceType, error) {
	resourceImagesDir := cmd.assetPath("resource-images")
	tarBin := cmd.assetPath("bin", "tar")

	archive := filepath.Join(resourcesDir, resourceType, "rootfs.tar.gz")

	extractedDir := filepath.Join(resourceImagesDir, resourceType)

	rootfsDir := filepath.Join(extractedDir, "rootfs")
	okMarker := filepath.Join(extractedDir, "ok")

	var version string
	versionFile, err := os.Open(filepath.Join(resourcesDir, resourceType, "version"))
	if err != nil {
		logger.Error("failed-to-read-version", err)
		return atc.WorkerResourceType{}, err
	}

	_, err = fmt.Fscanf(versionFile, "%s", &version)
	if err != nil {
		logger.Error("failed-to-parse-version", err)
		return atc.WorkerResourceType{}, err
	}

	defer versionFile.Close()

	var privileged bool
	_, err = os.Stat(filepath.Join(resourcesDir, resourceType, "privileged"))
	if err == nil {
		privileged = true
	}

	_, err = os.Stat(okMarker)
	if err == nil {
		logger.Info("already-extracted")
	} else {
		logger.Info("extracting")

		err := os.RemoveAll(rootfsDir)
		if err != nil {
			logger.Error("failed-to-clear-out-existing-rootfs", err)
			return atc.WorkerResourceType{}, err
		}

		err = os.MkdirAll(rootfsDir, 0755)
		if err != nil {
			logger.Error("failed-to-create-rootfs-dir", err)
			return atc.WorkerResourceType{}, err
		}

		tar := exec.Command(tarBin, "-zxf", archive, "-C", rootfsDir)

		output, err := tar.CombinedOutput()
		if err != nil {
			logger.Error("failed-to-extract-resource", err, lager.Data{
				"output": string(output),
			})
			return atc.WorkerResourceType{}, err
		}

		ok, err := os.Create(okMarker)
		if err != nil {
			logger.Error("failed-to-create-ok-marker", err)
			return atc.WorkerResourceType{}, err
		}

		err = ok.Close()
		if err != nil {
			logger.Error("failed-to-close-ok-marker", err)
			return atc.WorkerResourceType{}, err
		}
	}

	return atc.WorkerResourceType{
		Type:       resourceType,
		Image:      rootfsDir,
		Version:    version,
		Privileged: privileged,
	}, nil
}

var ErrNotRoot = errors.New("worker must be run as root")

func (cmd *WorkerCommand) checkRoot() error {
	currentUser, err := user.Current()
	if err != nil {
		return err
	}

	if currentUser.Uid != "0" {
		return ErrNotRoot
	}

	return nil
}

func (cmd *WorkerCommand) dnsProxyRunner(logger lager.Logger) (ifrit.Runner, error) {
	server, err := cmd.Garden.DNS.Server()
	if err != nil {
		return nil, err
	}

	return ifrit.RunFunc(func(signals <-chan os.Signal, ready chan<- struct{}) error {
		server.NotifyStartedFunc = func() {
			close(ready)
			logger.Info("started")
		}

		serveErr := make(chan error, 1)

		go func() {
			serveErr <- server.ListenAndServe()
		}()

		for {
			select {
			case err := <-serveErr:
				return err
			case <-signals:
				server.Shutdown()
			}
		}
	}), nil
}
<|endoftext|>


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

// The ePoxy boot server is the first point of contact for managed machines as
// they boot. The boot server serves all client connections over TLS. And, the
// boot server restricts all state-changing requests to administrative users
// (any machine) and managed machines (only itself).
//
// Managed machines progress through three boot stages:
//   stage1) local boot media like an iPXE ROM, or an immutable CD image
//   stage2) a minimal, linux-based network boot environment
//   stage3) the final system image.
//
// Managed machines are treated as stateless. So, the ePoxy boot server acts as
// an external state manager that mediates the transition of successive boot
// stages. Managed machines positively acknowlege every stage transition using
// session ids generated on the first request and known only to the ePoxy boot
// server and the remote machine.
//
// So, if a managed machine acknowleges the final stage successfully, then we
// know that this machine is the same one that first contacted the ePoxy boot
// server.
package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"

	"cloud.google.com/go/datastore"
	"github.com/gorilla/mux"
	"github.com/m-lab/epoxy/handler"
	"github.com/m-lab/epoxy/storage"
)

var (
	// Environment variables are preferred to flags for deployments in
	// AppEngine. And, using environment variables is encouraged for
	// twelve-factor apps -- https://12factor.net/config

	// projectID must be set using the GCLOUD_PROJECT environment variable.
	projectID = os.Getenv("GCLOUD_PROJECT")

	// publicAddr must be set if *not* running in AppEngine. When running in
	// AppEngine publicAddr is set automatically from a combination of
	// GCLOUD_PROJECT and GAE_SERVICE environment variables. When not running in
	// AppEngine, or to override the default in AppEngine, the PUBLIC_ADDRESS
	// environment variable must be set instead.
	publicAddr = os.Getenv("PUBLIC_ADDRESS")

	// bindAddress may be set using the LISTEN environment variable. By default, ePoxy
	// listens on all available interfaces.
	bindAddress = os.Getenv("LISTEN")

	// bindPort may be set using the PORT environment variable.
	bindPort = "8080"
)

// init checks the environment for configuration values.
func init() {
	// Only use the automatic public address if PUBLIC_ADDRESS is not already set.
	if service := os.Getenv("GAE_SERVICE"); service != "" && projectID != "" && publicAddr == "" {
		publicAddr = fmt.Sprintf("%s-dot-%s.appspot.com", service, projectID)
	}
	if port := os.Getenv("PORT"); port != "" {
		bindPort = port
	}
}

// addRoute adds a new handler for a pattern-based URL target to a Gorilla mux.Router.
func addRoute(router *mux.Router, method, pattern string, handler http.Handler) {
	router.Methods(method).Path(pattern).Handler(handler)
}

// newRouter creates and initializes all routes for the ePoxy boot server.
func newRouter(env *handler.Env) *mux.Router {
	router := mux.NewRouter()

	// A health checker for running in Docker or AppEngine.
	addRoute(router, "GET", "/_ah/health", http.HandlerFunc(checkHealth))

	// Stage1 scripts are always the first script fetched by a booting machine.
	// "stage1.ipxe" is the target for ROM-based iPXE clients.
	addRoute(router, "POST", "/v1/boot/{hostname}/stage1.ipxe",
		http.HandlerFunc(env.GenerateStage1IPXE))

	// TODO(soltesz): add a target for CD-based ePoxy clients.
	// addRoute(router, "POST", "/v1/boot/{hostname}/stage1.json", generateStage1Json)

	//TODO: make the names stage2 and stage3 arbitrary when we need to support
	///the case where not every machine has the same stage2 or stage3

	// Stage2, stage3, and report targets load after stage1 runs successfully.
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/stage2",
		http.HandlerFunc(env.GenerateJSONConfig))
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/stage3",
		http.HandlerFunc(env.GenerateJSONConfig))
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/report",
		http.HandlerFunc(env.ReceiveReport))
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/extension/{operation}",
		http.HandlerFunc(env.HandleExtension))

	// TODO(soltesz): add a target or retrieving all published SSH host keys.
	// addRoute(router, "GET", "/v1/boot/known_hosts", getKnownHosts)
	return router
}

// checkHealth reports whether the server is healthy. checkHealth will
// typically be registered as the http.Handler for the path "/_ah/health" when
// running in Docker or AppEngine.
func checkHealth(rw http.ResponseWriter, req *http.Request) {
	fmt.Fprint(rw, "ok")
}

func main() {
	if projectID == "" {
		log.Fatalf("Environment variable GCLOUD_PROJECT must specify a project ID for Datastore.")
	}
	if publicAddr == "" {
		log.Fatalf("Environment variable PUBLIC_ADDRESS must specify a public service name.")
	}

	// TODO(soltesz): support TLS natively for stand-alone mode. Though, this is not necessary for AppEngine.
	addr := fmt.Sprintf("%s:%s", bindAddress, bindPort)
	ctx := context.Background()
	client, err := datastore.NewClient(ctx, projectID)
	if err != nil {
		log.Fatalf("Failed to create new datastore client: %s", err)
	}
	env := &handler.Env{storage.NewDatastoreConfig(client), publicAddr}
	http.ListenAndServe(addr, newRouter(env))
}
<commit_msg>Add additional comments describing boot and extension targets<commit_after>// Copyright 2016 ePoxy Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//////////////////////////////////////////////////////////////////////////////

// The ePoxy boot server is the first point of contact for managed machines as
// they boot. The boot server serves all client connections over TLS. And, the
// boot server restricts all state-changing requests to administrative users
// (any machine) and managed machines (only itself).
//
// Managed machines progress through three boot stages:
//   stage1) local boot media like an iPXE ROM, or an immutable CD image
//   stage2) a minimal, linux-based network boot environment
//   stage3) the final system image.
//
// Managed machines are treated as stateless. So, the ePoxy boot server acts as
// an external state manager that mediates the transition of successive boot
// stages. Managed machines positively acknowlege every stage transition using
// session ids generated on the first request and known only to the ePoxy boot
// server and the remote machine.
//
// So, if a managed machine acknowleges the final stage successfully, then we
// know that this machine is the same one that first contacted the ePoxy boot
// server.
package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"

	"cloud.google.com/go/datastore"
	"github.com/gorilla/mux"
	"github.com/m-lab/epoxy/handler"
	"github.com/m-lab/epoxy/storage"
)

var (
	// Environment variables are preferred to flags for deployments in
	// AppEngine. And, using environment variables is encouraged for
	// twelve-factor apps -- https://12factor.net/config

	// projectID must be set using the GCLOUD_PROJECT environment variable.
	projectID = os.Getenv("GCLOUD_PROJECT")

	// publicAddr must be set if *not* running in AppEngine. When running in
	// AppEngine publicAddr is set automatically from a combination of
	// GCLOUD_PROJECT and GAE_SERVICE environment variables. When not running in
	// AppEngine, or to override the default in AppEngine, the PUBLIC_ADDRESS
	// environment variable must be set instead.
	publicAddr = os.Getenv("PUBLIC_ADDRESS")

	// bindAddress may be set using the LISTEN environment variable. By default, ePoxy
	// listens on all available interfaces.
	bindAddress = os.Getenv("LISTEN")

	// bindPort may be set using the PORT environment variable.
	bindPort = "8080"
)

// init checks the environment for configuration values.
func init() {
	// Only use the automatic public address if PUBLIC_ADDRESS is not already set.
	if service := os.Getenv("GAE_SERVICE"); service != "" && projectID != "" && publicAddr == "" {
		publicAddr = fmt.Sprintf("%s-dot-%s.appspot.com", service, projectID)
	}
	if port := os.Getenv("PORT"); port != "" {
		bindPort = port
	}
}

// addRoute adds a new handler for a pattern-based URL target to a Gorilla mux.Router.
func addRoute(router *mux.Router, method, pattern string, handler http.Handler) {
	router.Methods(method).Path(pattern).Handler(handler)
}

// newRouter creates and initializes all routes for the ePoxy boot server.
func newRouter(env *handler.Env) *mux.Router {
	router := mux.NewRouter()

	// A health checker for running in Docker or AppEngine.
	addRoute(router, "GET", "/_ah/health", http.HandlerFunc(checkHealth))

	///////////////////////////////////////////////////////////////////////////
	// Boot stage targets.
	//
	// Immediately after boot, a maching unconditionally requests a stage1 target.
	// After that, the machine should incrementally request the stage2, stage3, and
	// report targets in sequence. As each is requested, the session ID for the
	// previous is invalidated.

	// Stage1 scripts are always the first script fetched by a booting machine.
	// "stage1.ipxe" is the target for ROM-based iPXE clients.
	addRoute(router, "POST", "/v1/boot/{hostname}/stage1.ipxe",
		http.HandlerFunc(env.GenerateStage1IPXE))

	// TODO(soltesz): add a target for CD-based ePoxy clients.
	// addRoute(router, "POST", "/v1/boot/{hostname}/stage1.json", generateStage1Json)

	// TODO: make the names stage2 and stage3 arbitrary when we need to support
	// the case where not every machine has the same stage2 or stage3.

	// TODO: consider placing stage sequence names under their own subpath, e.g.
	//   /v1/boot/{hostname}/{sessionID}/stage/{stage}"

	// Stage2, stage3, and report targets load after stage1 runs successfully. Stage2
	// and stage3 targets return an epoxy action. The report target returns no content.
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/stage2",
		http.HandlerFunc(env.GenerateJSONConfig))
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/stage3",
		http.HandlerFunc(env.GenerateJSONConfig))
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/report",
		http.HandlerFunc(env.ReceiveReport))

	///////////////////////////////////////////////////////////////////////////
	// Extension targets.
	//
	// Extension operations may be requested at any time during boot. The session
	// is revoked after successful use. Extensions may return any content type
	// supported by the extension service.
	addRoute(router, "POST", "/v1/boot/{hostname}/{sessionID}/extension/{operation}",
		http.HandlerFunc(env.HandleExtension))

	return router
}

// checkHealth reports whether the server is healthy. checkHealth will
// typically be registered as the http.Handler for the path "/_ah/health" when
// running in Docker or AppEngine.
func checkHealth(rw http.ResponseWriter, req *http.Request) {
	fmt.Fprint(rw, "ok")
}

func main() {
	if projectID == "" {
		log.Fatalf("Environment variable GCLOUD_PROJECT must specify a project ID for Datastore.")
	}
	if publicAddr == "" {
		log.Fatalf("Environment variable PUBLIC_ADDRESS must specify a public service name.")
	}

	// TODO(soltesz): support TLS natively for stand-alone mode. Though, this is not necessary for AppEngine.
	addr := fmt.Sprintf("%s:%s", bindAddress, bindPort)
	ctx := context.Background()
	client, err := datastore.NewClient(ctx, projectID)
	if err != nil {
		log.Fatalf("Failed to create new datastore client: %s", err)
	}
	env := &handler.Env{
		Config:     storage.NewDatastoreConfig(client),
		ServerAddr: publicAddr,
	}
	http.ListenAndServe(addr, newRouter(env))
}
<|endoftext|>


<commit_before>// Copyright (c) 2018, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"testing"
	"time"

	"mvdan.cc/fdroidcl/adb"
)

// chosenApp is the app that will be installed and uninstalled on a connected
// device. This one was chosen because it's tiny, requires no permissions, and
// should be compatible with every device.
//
// It also stores no data, so it is fine to uninstall it and the user won't lose
// any data.
const chosenApp = "org.vi_server.red_screen"

func TestCommands(t *testing.T) {
	url := config.Repos[0].URL
	client := http.Client{Timeout: 2 * time.Second}
	if _, err := client.Get(url); err != nil {
		t.Skipf("skipping since %s is unreachable: %v", url, err)
	}

	dir, err := ioutil.TempDir("", "fdroidcl")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)
	testBasedir = dir

	mustSucceed := func(t *testing.T, wantRe, negRe string, cmd *Command, args ...string) {
		mustRun(t, true, wantRe, negRe, cmd, args...)
	}
	mustFail := func(t *testing.T, wantRe, negRe string, cmd *Command, args ...string) {
		mustRun(t, false, wantRe, negRe, cmd, args...)
	}

	t.Run("Version", func(t *testing.T) {
		mustSucceed(t, `^v`, ``, cmdVersion)
	})

	t.Run("SearchBeforeUpdate", func(t *testing.T) {
		mustFail(t, `could not open index`, ``, cmdSearch)
	})
	t.Run("UpdateFirst", func(t *testing.T) {
		mustSucceed(t, `done`, ``, cmdUpdate)
	})
	t.Run("UpdateCached", func(t *testing.T) {
		mustSucceed(t, `not modified`, ``, cmdUpdate)
	})

	t.Run("SearchNoArgs", func(t *testing.T) {
		mustSucceed(t, `F-Droid`, ``, cmdSearch)
	})
	t.Run("SearchWithArgs", func(t *testing.T) {
		mustSucceed(t, `F-Droid`, ``, cmdSearch, "fdroid.fdroid")
	})
	t.Run("SearchWithArgsNone", func(t *testing.T) {
		mustSucceed(t, `^$`, ``, cmdSearch, "nomatches")
	})
	t.Run("SearchOnlyPackageNames", func(t *testing.T) {
		mustSucceed(t, `^[^ ]*$`, ``, cmdSearch, "-q", "fdroid.fdroid")
	})

	t.Run("ShowOne", func(t *testing.T) {
		mustSucceed(t, `fdroid/fdroidclient`, ``, cmdShow, "org.fdroid.fdroid")
	})
	t.Run("ShowMany", func(t *testing.T) {
		mustSucceed(t, `fdroid/fdroidclient.*fdroid/privileged-extension`, ``,
			cmdShow, "org.fdroid.fdroid", "org.fdroid.fdroid.privileged")
	})

	t.Run("ListCategories", func(t *testing.T) {
		mustSucceed(t, `Development`, ``, cmdList, "categories")
	})

	if err := startAdbIfNeeded(); err != nil {
		t.Fatal(err)
	}
	devices, err := adb.Devices()
	if err != nil {
		t.Fatal(err)
	}
	switch len(devices) {
	case 0:
		t.Log("skipping the device tests as none was found via ADB")
	case 1:
		// continue below
	default:
		t.Log("skipping the device tests as too many were found via ADB")
	}

	t.Run("DevicesOne", func(t *testing.T) {
		mustSucceed(t, `\n`, ``, cmdDevices)
	})

	// try to uninstall the app first
	devices[0].Uninstall(chosenApp)
	t.Run("UninstallMissing", func(t *testing.T) {
		mustFail(t, `not installed$`, ``, cmdUninstall, chosenApp)
	})
	t.Run("SearchInstalledMissing", func(t *testing.T) {
		mustSucceed(t, ``, regexp.QuoteMeta(chosenApp), cmdSearch, "-i", "-q")
	})
	t.Run("SearchUpgradableMissing", func(t *testing.T) {
		mustSucceed(t, ``, regexp.QuoteMeta(chosenApp), cmdSearch, "-u", "-q")
	})
	t.Run("InstallVersioned", func(t *testing.T) {
		mustSucceed(t, `Installing `+regexp.QuoteMeta(chosenApp), ``,
			cmdInstall, chosenApp+":1")
	})
	t.Run("SearchInstalled", func(t *testing.T) {
		time.Sleep(3 * time.Second)
		mustSucceed(t, regexp.QuoteMeta(chosenApp), ``, cmdSearch, "-i", "-q")
	})
	t.Run("SearchUpgradable", func(t *testing.T) {
		mustSucceed(t, regexp.QuoteMeta(chosenApp), ``, cmdSearch, "-u", "-q")
	})
	t.Run("Upgrade", func(t *testing.T) {
		mustSucceed(t, `Upgrading `+regexp.QuoteMeta(chosenApp), ``,
			cmdUpgrade, chosenApp)
	})
	t.Run("SearchUpgradableUpToDate", func(t *testing.T) {
		mustSucceed(t, ``, regexp.QuoteMeta(chosenApp), cmdSearch, "-u", "-q")
	})
	t.Run("UpgradeAlreadyInstalled", func(t *testing.T) {
		mustFail(t, `is up to date$`, ``, cmdUpgrade, chosenApp)
	})
	t.Run("UninstallExisting", func(t *testing.T) {
		mustSucceed(t, `Uninstalling `+regexp.QuoteMeta(chosenApp), ``,
			cmdUninstall, chosenApp)
	})
}

func mustRun(t *testing.T, success bool, wantRe, negRe string, cmd *Command, args ...string) {
	var buf bytes.Buffer
	stdout, stderr = &buf, &buf
	err := cmd.Run(args)
	out := buf.String()
	if success && err != nil {
		t.Fatalf("unexpected error: %v\n%s", err, out)
	} else if !success && err == nil {
		t.Fatalf("expected error, got none\n%s", out)
	}
	if err != nil {
		out += err.Error()
	}
	// Let '.' match newlines, and treat the output as a single line.
	wantRe = "(?sm)" + wantRe
	if !regexp.MustCompile(wantRe).MatchString(out) {
		t.Fatalf("output does not match %#q:\n%s", wantRe, out)
	}
	if negRe != "" {
		negRe = "(?sm)" + negRe
		if regexp.MustCompile(negRe).MatchString(out) {
			t.Fatalf("output does match %#q:\n%s", negRe, out)
		}
	}
}
<commit_msg>cmd/fdroidcl: skip tests if adb is not installed<commit_after>// Copyright (c) 2018, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"testing"
	"time"

	"mvdan.cc/fdroidcl/adb"
)

// chosenApp is the app that will be installed and uninstalled on a connected
// device. This one was chosen because it's tiny, requires no permissions, and
// should be compatible with every device.
//
// It also stores no data, so it is fine to uninstall it and the user won't lose
// any data.
const chosenApp = "org.vi_server.red_screen"

func TestCommands(t *testing.T) {
	url := config.Repos[0].URL
	client := http.Client{Timeout: 2 * time.Second}
	if _, err := client.Get(url); err != nil {
		t.Skipf("skipping since %s is unreachable: %v", url, err)
	}

	dir, err := ioutil.TempDir("", "fdroidcl")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)
	testBasedir = dir

	mustSucceed := func(t *testing.T, wantRe, negRe string, cmd *Command, args ...string) {
		mustRun(t, true, wantRe, negRe, cmd, args...)
	}
	mustFail := func(t *testing.T, wantRe, negRe string, cmd *Command, args ...string) {
		mustRun(t, false, wantRe, negRe, cmd, args...)
	}

	t.Run("Version", func(t *testing.T) {
		mustSucceed(t, `^v`, ``, cmdVersion)
	})

	t.Run("SearchBeforeUpdate", func(t *testing.T) {
		mustFail(t, `could not open index`, ``, cmdSearch)
	})
	t.Run("UpdateFirst", func(t *testing.T) {
		mustSucceed(t, `done`, ``, cmdUpdate)
	})
	t.Run("UpdateCached", func(t *testing.T) {
		mustSucceed(t, `not modified`, ``, cmdUpdate)
	})

	t.Run("SearchNoArgs", func(t *testing.T) {
		mustSucceed(t, `F-Droid`, ``, cmdSearch)
	})
	t.Run("SearchWithArgs", func(t *testing.T) {
		mustSucceed(t, `F-Droid`, ``, cmdSearch, "fdroid.fdroid")
	})
	t.Run("SearchWithArgsNone", func(t *testing.T) {
		mustSucceed(t, `^$`, ``, cmdSearch, "nomatches")
	})
	t.Run("SearchOnlyPackageNames", func(t *testing.T) {
		mustSucceed(t, `^[^ ]*$`, ``, cmdSearch, "-q", "fdroid.fdroid")
	})

	t.Run("ShowOne", func(t *testing.T) {
		mustSucceed(t, `fdroid/fdroidclient`, ``, cmdShow, "org.fdroid.fdroid")
	})
	t.Run("ShowMany", func(t *testing.T) {
		mustSucceed(t, `fdroid/fdroidclient.*fdroid/privileged-extension`, ``,
			cmdShow, "org.fdroid.fdroid", "org.fdroid.fdroid.privileged")
	})

	t.Run("ListCategories", func(t *testing.T) {
		mustSucceed(t, `Development`, ``, cmdList, "categories")
	})

	if err := startAdbIfNeeded(); err != nil {
		t.Log("skipping the device tests as ADB is not installed")
		return
	}
	devices, err := adb.Devices()
	if err != nil {
		t.Fatal(err)
	}
	switch len(devices) {
	case 0:
		t.Log("skipping the device tests as none was found via ADB")
		return
	case 1:
		// continue below
	default:
		t.Log("skipping the device tests as too many were found via ADB")
		return
	}

	t.Run("DevicesOne", func(t *testing.T) {
		mustSucceed(t, `\n`, ``, cmdDevices)
	})

	// try to uninstall the app first
	devices[0].Uninstall(chosenApp)
	t.Run("UninstallMissing", func(t *testing.T) {
		mustFail(t, `not installed$`, ``, cmdUninstall, chosenApp)
	})
	t.Run("SearchInstalledMissing", func(t *testing.T) {
		mustSucceed(t, ``, regexp.QuoteMeta(chosenApp), cmdSearch, "-i", "-q")
	})
	t.Run("SearchUpgradableMissing", func(t *testing.T) {
		mustSucceed(t, ``, regexp.QuoteMeta(chosenApp), cmdSearch, "-u", "-q")
	})
	t.Run("InstallVersioned", func(t *testing.T) {
		mustSucceed(t, `Installing `+regexp.QuoteMeta(chosenApp), ``,
			cmdInstall, chosenApp+":1")
	})
	t.Run("SearchInstalled", func(t *testing.T) {
		time.Sleep(3 * time.Second)
		mustSucceed(t, regexp.QuoteMeta(chosenApp), ``, cmdSearch, "-i", "-q")
	})
	t.Run("SearchUpgradable", func(t *testing.T) {
		mustSucceed(t, regexp.QuoteMeta(chosenApp), ``, cmdSearch, "-u", "-q")
	})
	t.Run("Upgrade", func(t *testing.T) {
		mustSucceed(t, `Upgrading `+regexp.QuoteMeta(chosenApp), ``,
			cmdUpgrade, chosenApp)
	})
	t.Run("SearchUpgradableUpToDate", func(t *testing.T) {
		mustSucceed(t, ``, regexp.QuoteMeta(chosenApp), cmdSearch, "-u", "-q")
	})
	t.Run("UpgradeAlreadyInstalled", func(t *testing.T) {
		mustFail(t, `is up to date$`, ``, cmdUpgrade, chosenApp)
	})
	t.Run("UninstallExisting", func(t *testing.T) {
		mustSucceed(t, `Uninstalling `+regexp.QuoteMeta(chosenApp), ``,
			cmdUninstall, chosenApp)
	})
}

func mustRun(t *testing.T, success bool, wantRe, negRe string, cmd *Command, args ...string) {
	var buf bytes.Buffer
	stdout, stderr = &buf, &buf
	err := cmd.Run(args)
	out := buf.String()
	if success && err != nil {
		t.Fatalf("unexpected error: %v\n%s", err, out)
	} else if !success && err == nil {
		t.Fatalf("expected error, got none\n%s", out)
	}
	if err != nil {
		out += err.Error()
	}
	// Let '.' match newlines, and treat the output as a single line.
	wantRe = "(?sm)" + wantRe
	if !regexp.MustCompile(wantRe).MatchString(out) {
		t.Fatalf("output does not match %#q:\n%s", wantRe, out)
	}
	if negRe != "" {
		negRe = "(?sm)" + negRe
		if regexp.MustCompile(negRe).MatchString(out) {
			t.Fatalf("output does match %#q:\n%s", negRe, out)
		}
	}
}
<|endoftext|>


<commit_before>package loadspec

import (
	"bufio"
	"encoding/json"
	"net/url"
	"os"
	"regexp"
	"sort"
	"strings"
	"time"

	"github.com/danielfireman/esperf/loadspec"
	"github.com/spf13/cobra"
)

var parseSlowlogCmd = &cobra.Command{
	Use:   "parseslowlog",
	Short: "Outputs a replayable loadspec based on the passed-in slowlog and parameters.",
	Long:  "Outputs a replayable loadspec based on the passed-in slowlog and parameters.",
	RunE: func(cmd *cobra.Command, args []string) error {
		// These constants need to be in sync with the regular expression bellow.
		const (
			logTypeField    = "log_type"
			hostField       = "host"
			timestampField  = "ts"
			indexField      = "index"
			typesField      = "types"
			searchTypeField = "search_type"
			sourceField     = "source"
			numFields       = 6
		)

		var urlArg string
		if len(args) > 0 {
			urlArg = args[0]
		}
		// Regular expression setup.
		// The solution is based on regexp's named matches. For each entry, we build a map of
		// of fields and values. This map is encoded as json and (buffered) written to stdout.
		re, err := regexp.Compile(`\[(?P<ts>[^]]+)\].?\[.*\].?\[(?P<log_type>[^]]+)\].?\[(?P<host>[^]]+)\].?\[(?P<index>[^]]+)\].?\[.*\].*types\[(?P<types>[^]]+)\].*search_type\[(?P<search_type>[^]]+)\].*source\[(?P<source>.*)\], extra_source`)
		if err != nil {
			return err
		}
		subExpNames := re.SubexpNames()

		var entries loadspec.ByDelaySinceLastNanos
		fields := make(map[string]string, numFields)
		scanner := bufio.NewScanner(os.Stdin)
		for scanner.Scan() {
			// Building a map using named matches.
			matches := re.FindAllStringSubmatch(scanner.Text(), -1)[0]
			for i, n := range matches {
				// Removing the first match, which is the whole line.
				if i > 0 {
					fields[subExpNames[i]] = n
				}
			}
			// For now, only processing queries.
			if fields[logTypeField] != "index.search.slowlog.query" {
				continue
			}
			entry := loadspec.Entry{Source: fields[sourceField]}
			// Making timestamp relative to the previous one. Simulate inter-arrival time can be as easy
			// as a time.Sleep and trigger a goroutine.
			t, err := time.Parse(timeLayout, strings.Replace(fields[timestampField], ",", ".", 1))
			if err != nil {
				return err
			}
			// Keeping timestamp here for post-processing bellow.
			entry.DelaySinceLastNanos = t.UnixNano()
			// Host argument is treated as full URL. This keeps consistency between here and gen command.
			var u *url.URL
			if urlArg != "" {
				u, err = url.Parse(urlArg)
				if err != nil {
					return err
				}
			} else {
				path := []string{fields[hostField], fields[indexField], fields[typesField], "_search"}
				u, err = url.Parse(strings.Join(path, "/"))
				if err != nil {
					return err
				}
				q := u.Query()
				if fields[searchTypeField] != "" {
					q.Set("search_type", strings.ToLower(fields[searchTypeField]))
				}
				u.RawQuery = q.Encode()
			}
			entry.URL = u.String()
			entries = append(entries, &entry)
		}
		if err := scanner.Err(); err != nil {
			return err
		}
		// Slow log entries are not guaranteed to be timestamp ordered.
		sort.Sort(entries)

		// Writer and encoding configuration.
		writer := bufio.NewWriter(os.Stdout)
		defer writer.Flush()
		enc := json.NewEncoder(writer)
		var previousTimestamp, currTimestamp int64
		for i, e := range entries {
			currTimestamp = e.DelaySinceLastNanos
			if i == 0 {
				e.DelaySinceLastNanos = 0
			} else {
				e.DelaySinceLastNanos -= previousTimestamp
			}
			previousTimestamp = currTimestamp
			if err := enc.Encode(&e); err != nil {
				return err
			}
		}
		return nil
	},
}
<commit_msg>Introducing hability to override only indexes on parseslowlog<commit_after>package loadspec

import (
	"bufio"
	"encoding/json"
	"os"
	"regexp"
	"sort"
	"strings"
	"time"

	"fmt"

	"github.com/danielfireman/esperf/loadspec"
	"github.com/spf13/cobra"
)

var (
	indexOverrides []string
)

func init() {
	parseSlowlogCmd.Flags().StringSliceVar(&indexOverrides, "index_overrides", []string{}, "Only queries to those indexes are going to be considered in the generated loadspec.")
}

var parseSlowlogCmd = &cobra.Command{
	Use:   "parseslowlog",
	Short: "Outputs a replayable loadspec based on the passed-in slowlog and parameters.",
	Long:  "Outputs a replayable loadspec based on the passed-in slowlog and parameters.",
	RunE: func(cmd *cobra.Command, args []string) error {
		// These constants need to be in sync with the regular expression bellow.
		const (
			logTypeField    = "log_type"
			hostField       = "host"
			timestampField  = "ts"
			indexField      = "index"
			typesField      = "types"
			searchTypeField = "search_type"
			sourceField     = "source"
			numFields       = 6
		)

		var urlArg string
		if len(args) > 0 {
			// To keep in par with gen, we only consider the host or host:port part of the URL.
			urlArg = args[0]
			prefix := ""
			switch {
			case strings.HasPrefix(urlArg, "http://"):
				urlArg = strings.TrimPrefix(urlArg, "http://")
				prefix = "http://"
			case strings.HasPrefix(urlArg, "https://"):
				urlArg = strings.TrimPrefix(urlArg, "https://")
				prefix = "https://"
			}
			i := strings.Index(urlArg, "/")
			if i > 0 {
				urlArg = prefix + urlArg[:i]
			}
		}

		// Regular expression setup.
		// The solution is based on regexp's named matches. For each entry, we build a map of
		// of fields and values. This map is encoded as json and (buffered) written to stdout.
		re, err := regexp.Compile(`\[(?P<ts>[^]]+)\].?\[.*\].?\[(?P<log_type>[^]]+)\].?\[(?P<host>[^]]+)\].?\[(?P<index>[^]]+)\].?\[.*\].*types\[(?P<types>[^]]+)\].*search_type\[(?P<search_type>[^]]+)\].*source\[(?P<source>.*)\], extra_source`)
		if err != nil {
			return err
		}
		subExpNames := re.SubexpNames()

		var entries loadspec.ByDelaySinceLastNanos
		fields := make(map[string]string, numFields)
		scanner := bufio.NewScanner(os.Stdin)
		count := 0
		for scanner.Scan() {
			// Building a map using named matches.
			matches := re.FindAllStringSubmatch(scanner.Text(), -1)[0]
			for i, n := range matches {
				// Removing the first match, which is the whole line.
				if i > 0 {
					fields[subExpNames[i]] = n
				}
			}
			// For now, only processing queries.
			if fields[logTypeField] != "index.search.slowlog.query" {
				continue
			}

			entry := loadspec.Entry{Source: fields[sourceField]}
			// Making timestamp relative to the previous one. Simulate inter-arrival time can be as easy
			// as a time.Sleep and trigger a goroutine.
			t, err := time.Parse(timeLayout, strings.Replace(fields[timestampField], ",", ".", 1))
			if err != nil {
				return err
			}
			// Keeping timestamp here for post-processing bellow.
			entry.DelaySinceLastNanos = t.UnixNano()
			// Host argument is treated as full URL. This keeps consistency between here and gen command.

			host := fields[hostField]
			if urlArg != "" {
				host = urlArg
			}
			index := fields[indexField]
			if len(indexOverrides) > 0 {
				index = indexOverrides[count%len(indexOverrides)]
			}

			// I would love to use url.URL, life is hard.
			// More on that: https://github.com/golang/go/issues/18824
			// TL;DR; We would like to use http://localhost:9200, but since go1.8 it is not allowed anymore.
			path := []string{host, index, fields[typesField], "_search"}
			st := ""
			if fields[searchTypeField] != "" {
				st = fmt.Sprintf("?search_type=%s", strings.ToLower(fields[searchTypeField]))
			}
			entry.URL = fmt.Sprintf("%s%s", strings.Join(path, "/"), st)
			entries = append(entries, &entry)
			count++
		}
		if err := scanner.Err(); err != nil {
			return err
		}
		// Slow log entries are not guaranteed to be timestamp ordered.
		sort.Sort(entries)

		// Writer and encoding configuration.
		writer := bufio.NewWriter(os.Stdout)
		defer writer.Flush()
		enc := json.NewEncoder(writer)
		var previousTimestamp, currTimestamp int64
		for i, e := range entries {
			currTimestamp = e.DelaySinceLastNanos
			if i == 0 {
				e.DelaySinceLastNanos = 0
			} else {
				e.DelaySinceLastNanos -= previousTimestamp
			}
			previousTimestamp = currTimestamp
			if err := enc.Encode(&e); err != nil {
				return err
			}
		}
		return nil
	},
}
<|endoftext|>


<commit_before>package commands

import (
	"fmt"
	"os"
	"os/exec"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/pkg/errors"
	"github.com/versent/saml2aws"
)

// Exec execute the supplied command after seeding the environment
func Exec(profile string, providerName string, skipVerify bool, cmdline []string) error {

	if len(cmdline) < 1 {
		return fmt.Errorf("Command to execute required.")
	}

	ok, err := checkToken(profile)
	if err != nil {
		return errors.Wrap(err, "error validating token")
	}

	if !ok {
		err = Login(profile, providerName, skipVerify)
	}
	if err != nil {
		return errors.Wrap(err, "error logging in")
	}

	sharedCreds := saml2aws.NewSharedCredentials(profile)

	id, secret, token, err := sharedCreds.Load()
	if err != nil {
		return errors.Wrap(err, "error loading credentials")
	}

	c := strings.Join(cmdline, " ")

	cs := []string{"/bin/sh", "-c", c}
	cmd := exec.Command(cs[0], cs[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Env = append(os.Environ(), buildEnvVars(id, secret, token)...)

	return cmd.Run()
}

func checkToken(profile string) (bool, error) {
	sess, err := session.NewSessionWithOptions(session.Options{
		Profile: profile,
	})
	if err != nil {
		return false, err
	}

	svc := sts.New(sess)

	params := &sts.GetCallerIdentityInput{}

	resp, err := svc.GetCallerIdentity(params)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ExpiredToken" {
				return false, nil
			}
		}

		return false, err
	}

	fmt.Println("Running command as:", aws.StringValue(resp.Arn))
	return true, nil
}

func buildEnvVars(id, secret, token string) []string {
	return []string{
		fmt.Sprintf("AWS_ACCESS_KEY_ID=%s", id),
		fmt.Sprintf("AWS_SECRET_ACCESS_KEY=%s", secret),
		fmt.Sprintf("AWS_SESSION_TOKEN=%s", token),
		fmt.Sprintf("AWS_SECURITY_TOKEN=%s", token),
		fmt.Sprintf("EC2_SECURITY_TOKEN=%s", token),
	}
}
<commit_msg>Changed info message to stderr to enable pipes.. my bad.<commit_after>package commands

import (
	"fmt"
	"os"
	"os/exec"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/pkg/errors"
	"github.com/versent/saml2aws"
)

// Exec execute the supplied command after seeding the environment
func Exec(profile string, providerName string, skipVerify bool, cmdline []string) error {

	if len(cmdline) < 1 {
		return fmt.Errorf("Command to execute required.")
	}

	ok, err := checkToken(profile)
	if err != nil {
		return errors.Wrap(err, "error validating token")
	}

	if !ok {
		err = Login(profile, providerName, skipVerify)
	}
	if err != nil {
		return errors.Wrap(err, "error logging in")
	}

	sharedCreds := saml2aws.NewSharedCredentials(profile)

	id, secret, token, err := sharedCreds.Load()
	if err != nil {
		return errors.Wrap(err, "error loading credentials")
	}

	c := strings.Join(cmdline, " ")

	cs := []string{"/bin/sh", "-c", c}
	cmd := exec.Command(cs[0], cs[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Env = append(os.Environ(), buildEnvVars(id, secret, token)...)

	return cmd.Run()
}

func checkToken(profile string) (bool, error) {
	sess, err := session.NewSessionWithOptions(session.Options{
		Profile: profile,
	})
	if err != nil {
		return false, err
	}

	svc := sts.New(sess)

	params := &sts.GetCallerIdentityInput{}

	resp, err := svc.GetCallerIdentity(params)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ExpiredToken" {
				return false, nil
			}
		}

		return false, err
	}

	fmt.Fprintln(os.Stderr, "Running command as:", aws.StringValue(resp.Arn))
	return true, nil
}

func buildEnvVars(id, secret, token string) []string {
	return []string{
		fmt.Sprintf("AWS_ACCESS_KEY_ID=%s", id),
		fmt.Sprintf("AWS_SECRET_ACCESS_KEY=%s", secret),
		fmt.Sprintf("AWS_SESSION_TOKEN=%s", token),
		fmt.Sprintf("AWS_SECURITY_TOKEN=%s", token),
		fmt.Sprintf("EC2_SECURITY_TOKEN=%s", token),
	}
}
<|endoftext|>