<commit_before>// Package hashmap provides a lock-free and thread-safe HashMap.
package hashmap

import (
	"bytes"
	"fmt"
	"reflect"
	"strconv"
	"sync/atomic"
	"unsafe"
)

// DefaultSize is the default size for a zero allocated map.
const DefaultSize = 8

// MaxFillRate is the maximum fill rate for the slice before a resize  will happen.
const MaxFillRate = 50

// HashMap implements a read optimized hash map.
type HashMap[Key comparable, Value any] struct {
	hasher     func(Key) uintptr
	store      atomic.Pointer[store[Key, Value]] // pointer to a map instance that gets replaced if the map resizes
	linkedList *List[Key, Value]                 // key sorted linked list of elements
	// resizing marks a resizing operation in progress.
	// this is using uintptr instead of atomic.Bool to avoid using 32 bit int on 64 bit systems
	resizing atomic.Uintptr
}

// KeyValue represents a key/value that is returned by the iterator.
type KeyValue[Key comparable, Value any] struct {
	Key   Key
	Value Value
}

// New returns a new HashMap instance.
func New[Key comparable, Value any]() *HashMap[Key, Value] {
	return NewSized[Key, Value](DefaultSize)
}

// NewSized returns a new HashMap instance with a specific initialization size.
func NewSized[Key comparable, Value any](size uintptr) *HashMap[Key, Value] {
	m := &HashMap[Key, Value]{}
	m.allocate(size)
	m.setDefaultHasher()
	return m
}

// SetHasher sets a custom hasher.
func (m *HashMap[Key, Value]) SetHasher(hasher func(Key) uintptr) {
	m.hasher = hasher
}

// Len returns the number of elements within the map.
func (m *HashMap[Key, Value]) Len() int {
	return m.linkedList.Len()
}

// Get retrieves an element from the map under given hash key.
func (m *HashMap[Key, Value]) Get(key Key) (Value, bool) {
	hash := m.hasher(key)
	store := m.store.Load()
	element := store.item(hash)

	// inline HashMap.searchItem()
	for element != nil {
		if element.keyHash == hash && element.key == key {
			return element.Value(), true
		}

		if element.keyHash > hash {
			return *new(Value), false
		}

		element = element.Next()
	}
	return *new(Value), false
}

// GetOrInsert returns the existing value for the key if present.
// Otherwise, it stores and returns the given value.
// The returned bool is true if the value was loaded, false if stored.
func (m *HashMap[Key, Value]) GetOrInsert(key Key, value Value) (Value, bool) {
	hash := m.hasher(key)
	var newElement *ListElement[Key, Value]

	for {
		store := m.store.Load()
		element := store.item(hash)

		for element != nil {
			if element.keyHash == hash && element.key == key {
				actual := element.Value()
				return actual, true
			}

			if element.keyHash > hash {
				break
			}

			element = element.Next()
		}

		if newElement == nil { // allocate only once
			newElement = &ListElement[Key, Value]{
				key:     key,
				keyHash: hash,
			}
			newElement.value.Store(&value)
		}

		if m.insertElement(newElement, false) {
			return value, false
		}
	}
}

// FillRate returns the fill rate of the map as a percentage integer.
func (m *HashMap[Key, Value]) FillRate() int {
	store := m.store.Load()
	count := int(store.count.Load())
	l := len(store.index)
	return (count * 100) / l
}

// Del deletes the key from the map and returns whether the key was deleted.
func (m *HashMap[Key, Value]) Del(key Key) bool {
	hash := m.hasher(key)
	store := m.store.Load()
	element := store.item(hash)

	for ; element != nil; element = element.Next() {
		if element.keyHash == hash && element.key == key {
			m.deleteElement(element)
			m.linkedList.Delete(element)
			return true
		}

		if element.keyHash > hash {
			return false
		}
	}
	return false
}

// Insert sets the value under the specified key to the map if it does not exist yet.
// If a resizing operation is happening concurrently while calling Insert, the item might show up in the map
// after the resize operation is finished.
// Returns true if the item was inserted or false if it existed.
func (m *HashMap[Key, Value]) Insert(key Key, value Value) bool {
	hash := m.hasher(key)
	element := &ListElement[Key, Value]{
		key:     key,
		keyHash: hash,
	}
	element.value.Store(&value)
	return m.insertElement(element, false)
}

// Set sets the value under the specified key to the map. An existing item for this key will be overwritten.
// If a resizing operation is happening concurrently while calling Set, the item might show up in the map
// after the resize operation is finished.
func (m *HashMap[Key, Value]) Set(key Key, value Value) {
	hash := m.hasher(key)
	element := &ListElement[Key, Value]{
		key:     key,
		keyHash: hash,
	}
	element.value.Store(&value)
	m.insertElement(element, true)
}

// Grow resizes the hashmap to a new size, the size gets rounded up to next power of 2.
// To double the size of the hashmap use newSize 0.
// This function returns immediately, the resize operation is done in a goroutine.
// No resizing is done in case of another resize operation already being in progress.
func (m *HashMap[Key, Value]) Grow(newSize uintptr) {
	if m.resizing.CompareAndSwap(0, 1) {
		go m.grow(newSize, true)
	}
}

// String returns the map as a string, only hashed keys are printed.
func (m *HashMap[Key, Value]) String() string {
	buffer := bytes.NewBufferString("")
	buffer.WriteRune('[')

	first := m.linkedList.First()
	item := first

	for item != nil {
		if item != first {
			buffer.WriteRune(',')
		}
		fmt.Fprint(buffer, item.keyHash)
		item = item.Next()
	}
	buffer.WriteRune(']')
	return buffer.String()
}

// Range calls f sequentially for each key and value present in the map.
// If f returns false, range stops the iteration.
func (m *HashMap[Key, Value]) Range(f func(Key, Value) bool) {
	item := m.linkedList.First()

	for item != nil {
		value := item.Value()
		if !f(item.key, value) {
			return
		}
		item = item.Next()
	}
}

func (m *HashMap[Key, Value]) allocate(newSize uintptr) {
	m.linkedList = NewList[Key, Value]()
	if m.resizing.CompareAndSwap(0, 1) {
		m.grow(newSize, false)
	}
}

// setDefaultHasher sets the default hasher depending on the key type.
func (m *HashMap[Key, Value]) setDefaultHasher() {
	var key Key
	switch any(key).(type) {
	case string:
		m.hasher = m.stringHasher
	case int, uint, uintptr:
		m.hasher = m.uintptrHasher
	case int8, uint8:
		m.hasher = m.byteHasher
	case int16, uint16:
		m.hasher = m.wordHasher
	case int32, uint32, float32:
		m.hasher = m.dwordHasher
	case int64, uint64, float64, complex64:
		m.hasher = m.qwordHasher
	case complex128:
		m.hasher = m.owordHasher
	default:
		panic(fmt.Errorf("unsupported key type %T", key))
	}
}

func (m *HashMap[Key, Value]) isResizeNeeded(store *store[Key, Value], count uintptr) bool {
	l := uintptr(len(store.index)) // l can't be 0 as it gets initialized in New()
	fillRate := (count * 100) / l
	return fillRate > MaxFillRate
}

/* The Golang 1.10.1 compiler does not inline this function well
func (m *HashMap[Key, Value]) searchItem(item *ListElement[Key, Value], key Key, keyHash uintptr) (Value, bool) {
	for item != nil {
		if item.keyHash == keyHash && item.key == key {
			return item.Value(), true
		}

		if item.keyHash > keyHash {
			return *new(Value), false
		}

		item = item.Next()
	}
	return *new(Value), false
}
*/

func (m *HashMap[Key, Value]) insertElement(element *ListElement[Key, Value], update bool) bool {
	var existed, inserted bool

	for {
		store := m.store.Load()
		searchStart := store.item(element.keyHash)

		if update {
			inserted = m.linkedList.AddOrUpdate(element, searchStart)
		} else if !inserted { // if retrying after insert during grow, do not add to list again
			existed, inserted = m.linkedList.Add(element, searchStart)
			if existed {
				return false
			}
		}
		if !inserted {
			continue // a concurrent add did interfere, try again
		}

		count := store.addItem(element)
		currentStore := m.store.Load()
		if store != currentStore { // retry insert in case of insert during grow
			continue
		}

		if m.isResizeNeeded(store, count) && m.resizing.CompareAndSwap(0, 1) {
			go m.grow(0, true)
		}
		return true
	}
}

// deleteElement deletes an element from index.
func (m *HashMap[Key, Value]) deleteElement(element *ListElement[Key, Value]) {
	for {
		store := m.store.Load()
		index := element.keyHash >> store.keyShifts
		ptr := (*unsafe.Pointer)(unsafe.Pointer(uintptr(store.array) + index*intSizeBytes))

		next := element.Next()
		if next != nil && element.keyHash>>store.keyShifts != index {
			next = nil // do not set index to next item if it's not the same slice index
		}
		atomic.CompareAndSwapPointer(ptr, unsafe.Pointer(element), unsafe.Pointer(next))

		currentStore := m.store.Load()
		if store == currentStore { // check that no resize happened
			break
		}
	}
}

func (m *HashMap[Key, Value]) grow(newSize uintptr, loop bool) {
	defer m.resizing.CompareAndSwap(1, 0)

	for {
		currentStore := m.store.Load()
		if newSize == 0 {
			newSize = uintptr(len(currentStore.index)) << 1
		} else {
			newSize = roundUpPower2(newSize)
		}

		index := make([]*ListElement[Key, Value], newSize)
		header := (*reflect.SliceHeader)(unsafe.Pointer(&index))

		newStore := &store[Key, Value]{
			keyShifts: strconv.IntSize - log2(newSize),
			array:     unsafe.Pointer(header.Data), // use address of slice data storage
			index:     index,
		}

		m.fillIndexItems(newStore) // initialize new index slice with longer keys

		m.store.Store(newStore)

		m.fillIndexItems(newStore) // make sure that the new index is up-to-date with the current state of the linked list

		if !loop {
			return
		}

		// check if a new resize needs to be done already
		count := uintptr(m.Len())
		if !m.isResizeNeeded(newStore, count) {
			return
		}
		newSize = 0 // 0 means double the current size
	}
}

func (m *HashMap[Key, Value]) fillIndexItems(store *store[Key, Value]) {
	first := m.linkedList.First()
	item := first
	lastIndex := uintptr(0)

	for item != nil {
		index := item.keyHash >> store.keyShifts
		if item == first || index != lastIndex { // store item with smallest hash key for every index
			store.addItem(item)
			lastIndex = index
		}
		item = item.Next()
	}
}
<commit_msg>reduce variable allocations in getter<commit_after>// Package hashmap provides a lock-free and thread-safe HashMap.
package hashmap

import (
	"bytes"
	"fmt"
	"reflect"
	"strconv"
	"sync/atomic"
	"unsafe"
)

// DefaultSize is the default size for a zero allocated map.
const DefaultSize = 8

// MaxFillRate is the maximum fill rate for the slice before a resize  will happen.
const MaxFillRate = 50

// HashMap implements a read optimized hash map.
type HashMap[Key comparable, Value any] struct {
	hasher     func(Key) uintptr
	store      atomic.Pointer[store[Key, Value]] // pointer to a map instance that gets replaced if the map resizes
	linkedList *List[Key, Value]                 // key sorted linked list of elements
	// resizing marks a resizing operation in progress.
	// this is using uintptr instead of atomic.Bool to avoid using 32 bit int on 64 bit systems
	resizing atomic.Uintptr
}

// KeyValue represents a key/value that is returned by the iterator.
type KeyValue[Key comparable, Value any] struct {
	Key   Key
	Value Value
}

// New returns a new HashMap instance.
func New[Key comparable, Value any]() *HashMap[Key, Value] {
	return NewSized[Key, Value](DefaultSize)
}

// NewSized returns a new HashMap instance with a specific initialization size.
func NewSized[Key comparable, Value any](size uintptr) *HashMap[Key, Value] {
	m := &HashMap[Key, Value]{}
	m.allocate(size)
	m.setDefaultHasher()
	return m
}

// SetHasher sets a custom hasher.
func (m *HashMap[Key, Value]) SetHasher(hasher func(Key) uintptr) {
	m.hasher = hasher
}

// Len returns the number of elements within the map.
func (m *HashMap[Key, Value]) Len() int {
	return m.linkedList.Len()
}

// Get retrieves an element from the map under given hash key.
func (m *HashMap[Key, Value]) Get(key Key) (Value, bool) {
	hash := m.hasher(key)

	// inline HashMap.searchItem()
	for element := m.store.Load().item(hash); element != nil; element = element.Next() {
		if element.keyHash == hash && element.key == key {
			return element.Value(), true
		}

		if element.keyHash > hash {
			return *new(Value), false
		}
	}
	return *new(Value), false
}

// GetOrInsert returns the existing value for the key if present.
// Otherwise, it stores and returns the given value.
// The returned bool is true if the value was loaded, false if stored.
func (m *HashMap[Key, Value]) GetOrInsert(key Key, value Value) (Value, bool) {
	hash := m.hasher(key)
	var newElement *ListElement[Key, Value]

	for {
		for element := m.store.Load().item(hash); element != nil; element = element.Next() {
			if element.keyHash == hash && element.key == key {
				actual := element.Value()
				return actual, true
			}

			if element.keyHash > hash {
				break
			}
		}

		if newElement == nil { // allocate only once
			newElement = &ListElement[Key, Value]{
				key:     key,
				keyHash: hash,
			}
			newElement.value.Store(&value)
		}

		if m.insertElement(newElement, false) {
			return value, false
		}
	}
}

// FillRate returns the fill rate of the map as a percentage integer.
func (m *HashMap[Key, Value]) FillRate() int {
	store := m.store.Load()
	count := int(store.count.Load())
	l := len(store.index)
	return (count * 100) / l
}

// Del deletes the key from the map and returns whether the key was deleted.
func (m *HashMap[Key, Value]) Del(key Key) bool {
	hash := m.hasher(key)
	store := m.store.Load()
	element := store.item(hash)

	for ; element != nil; element = element.Next() {
		if element.keyHash == hash && element.key == key {
			m.deleteElement(element)
			m.linkedList.Delete(element)
			return true
		}

		if element.keyHash > hash {
			return false
		}
	}
	return false
}

// Insert sets the value under the specified key to the map if it does not exist yet.
// If a resizing operation is happening concurrently while calling Insert, the item might show up in the map
// after the resize operation is finished.
// Returns true if the item was inserted or false if it existed.
func (m *HashMap[Key, Value]) Insert(key Key, value Value) bool {
	hash := m.hasher(key)
	element := &ListElement[Key, Value]{
		key:     key,
		keyHash: hash,
	}
	element.value.Store(&value)
	return m.insertElement(element, false)
}

// Set sets the value under the specified key to the map. An existing item for this key will be overwritten.
// If a resizing operation is happening concurrently while calling Set, the item might show up in the map
// after the resize operation is finished.
func (m *HashMap[Key, Value]) Set(key Key, value Value) {
	hash := m.hasher(key)
	element := &ListElement[Key, Value]{
		key:     key,
		keyHash: hash,
	}
	element.value.Store(&value)
	m.insertElement(element, true)
}

// Grow resizes the hashmap to a new size, the size gets rounded up to next power of 2.
// To double the size of the hashmap use newSize 0.
// This function returns immediately, the resize operation is done in a goroutine.
// No resizing is done in case of another resize operation already being in progress.
func (m *HashMap[Key, Value]) Grow(newSize uintptr) {
	if m.resizing.CompareAndSwap(0, 1) {
		go m.grow(newSize, true)
	}
}

// String returns the map as a string, only hashed keys are printed.
func (m *HashMap[Key, Value]) String() string {
	buffer := bytes.NewBufferString("")
	buffer.WriteRune('[')

	first := m.linkedList.First()
	item := first

	for item != nil {
		if item != first {
			buffer.WriteRune(',')
		}
		fmt.Fprint(buffer, item.keyHash)
		item = item.Next()
	}
	buffer.WriteRune(']')
	return buffer.String()
}

// Range calls f sequentially for each key and value present in the map.
// If f returns false, range stops the iteration.
func (m *HashMap[Key, Value]) Range(f func(Key, Value) bool) {
	item := m.linkedList.First()

	for item != nil {
		value := item.Value()
		if !f(item.key, value) {
			return
		}
		item = item.Next()
	}
}

func (m *HashMap[Key, Value]) allocate(newSize uintptr) {
	m.linkedList = NewList[Key, Value]()
	if m.resizing.CompareAndSwap(0, 1) {
		m.grow(newSize, false)
	}
}

// setDefaultHasher sets the default hasher depending on the key type.
func (m *HashMap[Key, Value]) setDefaultHasher() {
	var key Key
	switch any(key).(type) {
	case string:
		m.hasher = m.stringHasher
	case int, uint, uintptr:
		m.hasher = m.uintptrHasher
	case int8, uint8:
		m.hasher = m.byteHasher
	case int16, uint16:
		m.hasher = m.wordHasher
	case int32, uint32, float32:
		m.hasher = m.dwordHasher
	case int64, uint64, float64, complex64:
		m.hasher = m.qwordHasher
	case complex128:
		m.hasher = m.owordHasher
	default:
		panic(fmt.Errorf("unsupported key type %T", key))
	}
}

func (m *HashMap[Key, Value]) isResizeNeeded(store *store[Key, Value], count uintptr) bool {
	l := uintptr(len(store.index)) // l can't be 0 as it gets initialized in New()
	fillRate := (count * 100) / l
	return fillRate > MaxFillRate
}

/* The Golang 1.10.1 compiler does not inline this function well
func (m *HashMap[Key, Value]) searchItem(item *ListElement[Key, Value], key Key, keyHash uintptr) (Value, bool) {
	for item != nil {
		if item.keyHash == keyHash && item.key == key {
			return item.Value(), true
		}

		if item.keyHash > keyHash {
			return *new(Value), false
		}

		item = item.Next()
	}
	return *new(Value), false
}
*/

func (m *HashMap[Key, Value]) insertElement(element *ListElement[Key, Value], update bool) bool {
	var existed, inserted bool

	for {
		store := m.store.Load()
		searchStart := store.item(element.keyHash)

		if update {
			inserted = m.linkedList.AddOrUpdate(element, searchStart)
		} else if !inserted { // if retrying after insert during grow, do not add to list again
			existed, inserted = m.linkedList.Add(element, searchStart)
			if existed {
				return false
			}
		}
		if !inserted {
			continue // a concurrent add did interfere, try again
		}

		count := store.addItem(element)
		currentStore := m.store.Load()
		if store != currentStore { // retry insert in case of insert during grow
			continue
		}

		if m.isResizeNeeded(store, count) && m.resizing.CompareAndSwap(0, 1) {
			go m.grow(0, true)
		}
		return true
	}
}

// deleteElement deletes an element from index.
func (m *HashMap[Key, Value]) deleteElement(element *ListElement[Key, Value]) {
	for {
		store := m.store.Load()
		index := element.keyHash >> store.keyShifts
		ptr := (*unsafe.Pointer)(unsafe.Pointer(uintptr(store.array) + index*intSizeBytes))

		next := element.Next()
		if next != nil && element.keyHash>>store.keyShifts != index {
			next = nil // do not set index to next item if it's not the same slice index
		}
		atomic.CompareAndSwapPointer(ptr, unsafe.Pointer(element), unsafe.Pointer(next))

		currentStore := m.store.Load()
		if store == currentStore { // check that no resize happened
			break
		}
	}
}

func (m *HashMap[Key, Value]) grow(newSize uintptr, loop bool) {
	defer m.resizing.CompareAndSwap(1, 0)

	for {
		currentStore := m.store.Load()
		if newSize == 0 {
			newSize = uintptr(len(currentStore.index)) << 1
		} else {
			newSize = roundUpPower2(newSize)
		}

		index := make([]*ListElement[Key, Value], newSize)
		header := (*reflect.SliceHeader)(unsafe.Pointer(&index))

		newStore := &store[Key, Value]{
			keyShifts: strconv.IntSize - log2(newSize),
			array:     unsafe.Pointer(header.Data), // use address of slice data storage
			index:     index,
		}

		m.fillIndexItems(newStore) // initialize new index slice with longer keys

		m.store.Store(newStore)

		m.fillIndexItems(newStore) // make sure that the new index is up-to-date with the current state of the linked list

		if !loop {
			return
		}

		// check if a new resize needs to be done already
		count := uintptr(m.Len())
		if !m.isResizeNeeded(newStore, count) {
			return
		}
		newSize = 0 // 0 means double the current size
	}
}

func (m *HashMap[Key, Value]) fillIndexItems(store *store[Key, Value]) {
	first := m.linkedList.First()
	item := first
	lastIndex := uintptr(0)

	for item != nil {
		index := item.keyHash >> store.keyShifts
		if item == first || index != lastIndex { // store item with smallest hash key for every index
			store.addItem(item)
			lastIndex = index
		}
		item = item.Next()
	}
}
<|endoftext|>


<commit_before>package mesosutil

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"regexp"

	log "github.com/golang/glog"
	mesos "github.com/mesos/mesos-go/mesosproto"
)

func FetchLog(logUrl string) ([]byte, error) {
	resp, err := defaultClient.Get(logUrl)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 200 {
		return nil, ProxyError{resp.Status, resp.StatusCode}
	}

	defer resp.Body.Close()
	return ioutil.ReadAll(resp.Body)
}

type ProxyError struct {
	Status     string
	StatusCode int
}

func (p ProxyError) Error() string {
	return p.Status
}

type MesosState struct {
	CompletedFrameworks []Framework `json:"completed_frameworks"`
	Frameworks          []Framework
}

func (m MesosState) Directory(frameworkId string) string {
	for _, f := range append(m.CompletedFrameworks, m.Frameworks...) {
		// should we check for the framework?
		for _, e := range append(f.CompletedExecutors, f.Executors...) {
			for _, t := range append(e.CompletedTasks, e.Tasks...) {
				if t.FrameworkId == frameworkId {
					log.V(2).Infoln("Matching task id " + frameworkId)
					return e.Directory
				}
			}
		}
	}
	return ""
}

type Tasks struct {
	FrameworkId string `json:"framework_id"`
}

type Framework struct {
	Name               string
	CompletedExecutors []Executor `json:"completed_executors"`
	Executors          []Executor
}

type Executor struct {
	CompletedTasks []Tasks `json:"completed_tasks"`
	Tasks          []Tasks
	Source         string
	Directory      string
}

type MesosTaskStatusMounts []struct {
	Source string `json:"Source"`
}

type MesosTaskStatusConfig struct {
	Hostname   string `json:"Hostname"`
	Domainname string `json:"Domainname"`
}

type MesosTaskStatusData []struct {
	Mounts MesosTaskStatusMounts `json:"Mounts"`
	Config MesosTaskStatusConfig `json:"Config"`
}

type LogData struct {
	Data   string `json:"data"`
	Offset int    `json:"offset"`
}

func FetchLogs(status *mesos.TaskStatus, offset int, file string, frameworkId string) ([]byte, error) {
	var (
		dir      string
		hostname string
		err      error
	)
	switch status.GetState() {
	case mesos.TaskState_TASK_FAILED,
		mesos.TaskState_TASK_KILLED:
		hostname = *status.ContainerStatus.NetworkInfos[0].IpAddress
		url := "http://" + hostname + ":5051/state.json"
		resp, err := defaultClient.Get(url)
		if err != nil {
			return nil, err
		}
		if resp.StatusCode != 200 {
			return nil, ProxyError{resp.Status, resp.StatusCode}
		}
		defer resp.Body.Close()
		var ms MesosState
		err = json.NewDecoder(resp.Body).Decode(&ms)
		if err != nil {
			return nil, err
		}
		dir = ms.Directory(frameworkId)
	default:
		var mtsd MesosTaskStatusData
		err = json.Unmarshal(status.Data, &mtsd)
		if err != nil {
			return nil, err
		}
		// status.Data is an array of one value :( Maybe there is a better way to marshal it?
		firstMtsd := mtsd[0]
		log.V(2).Infof("firstMtsd: %#v", firstMtsd)
		for _, mount := range firstMtsd.Mounts {
			source := mount.Source
			log.V(2).Infoln("mount: ", source)
			matched, _ := regexp.MatchString("slaves.*frameworks.*executors", source)
			if matched {
				dir = source
			}
		}
		domainName := ""
		if firstMtsd.Config.Domainname != "" {
			domainName = "." + firstMtsd.Config.Domainname
		}
		hostname = firstMtsd.Config.Hostname + domainName
	}
	logUrl := fmt.Sprintf("http://%s:5051/files/read.json?path=%s/%s&offset=%d",
		hostname, dir, file, offset)
	bodyData, _ := FetchLog(logUrl)

	var logData LogData
	err = json.Unmarshal(bodyData, &logData)
	if err != nil {
		return nil, err
	}
	return []byte(logData.Data), nil
}
<commit_msg>Slight cleanup of dup http fetching code.<commit_after>package mesosutil

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"regexp"

	log "github.com/golang/glog"
	mesos "github.com/mesos/mesos-go/mesosproto"
)

func FetchUrl(url string) ([]byte, error) {
	resp, err := defaultClient.Get(url)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 200 {
		return nil, ProxyError{resp.Status, resp.StatusCode}
	}

	defer resp.Body.Close()
	return ioutil.ReadAll(resp.Body)
}

type ProxyError struct {
	Status     string
	StatusCode int
}

func (p ProxyError) Error() string {
	return p.Status
}

type MesosState struct {
	CompletedFrameworks []Framework `json:"completed_frameworks"`
	Frameworks          []Framework
}

func (m MesosState) Directory(frameworkId string) string {
	for _, f := range append(m.CompletedFrameworks, m.Frameworks...) {
		// should we check for the framework?
		for _, e := range append(f.CompletedExecutors, f.Executors...) {
			for _, t := range append(e.CompletedTasks, e.Tasks...) {
				if t.FrameworkId == frameworkId {
					log.V(2).Infoln("Matching task id " + frameworkId)
					return e.Directory
				}
			}
		}
	}
	return ""
}

type Tasks struct {
	FrameworkId string `json:"framework_id"`
}

type Framework struct {
	Name               string
	CompletedExecutors []Executor `json:"completed_executors"`
	Executors          []Executor
}

type Executor struct {
	CompletedTasks []Tasks `json:"completed_tasks"`
	Tasks          []Tasks
	Source         string
	Directory      string
}

type MesosTaskStatusMounts []struct {
	Source string `json:"Source"`
}

type MesosTaskStatusConfig struct {
	Hostname   string `json:"Hostname"`
	Domainname string `json:"Domainname"`
}

type MesosTaskStatusData []struct {
	Mounts MesosTaskStatusMounts `json:"Mounts"`
	Config MesosTaskStatusConfig `json:"Config"`
}

type LogData struct {
	Data   string `json:"data"`
	Offset int    `json:"offset"`
}

func FetchLogs(status *mesos.TaskStatus, offset int, file string, frameworkId string) ([]byte, error) {
	var (
		dir      string
		hostname string
		err      error
	)
	switch status.GetState() {
	case mesos.TaskState_TASK_FAILED,
		mesos.TaskState_TASK_KILLED:
		hostname = *status.ContainerStatus.NetworkInfos[0].IpAddress
		url := "http://" + hostname + ":5051/state.json"
		bodyData, _ := FetchUrl(url)
		var ms MesosState
		err = json.Unmarshal(bodyData, &ms)
		if err != nil {
			return nil, err
		}
		dir = ms.Directory(frameworkId)
	default:
		var mtsd MesosTaskStatusData
		err = json.Unmarshal(status.Data, &mtsd)
		if err != nil {
			return nil, err
		}
		// status.Data is an array of one value :( Maybe there is a better way to marshal it?
		firstMtsd := mtsd[0]
		log.V(2).Infof("firstMtsd: %#v", firstMtsd)
		for _, mount := range firstMtsd.Mounts {
			source := mount.Source
			log.V(2).Infoln("mount: ", source)
			matched, _ := regexp.MatchString("slaves.*frameworks.*executors", source)
			if matched {
				dir = source
			}
		}
		domainName := ""
		if firstMtsd.Config.Domainname != "" {
			domainName = "." + firstMtsd.Config.Domainname
		}
		hostname = firstMtsd.Config.Hostname + domainName
	}
	url := fmt.Sprintf("http://%s:5051/files/read.json?path=%s/%s&offset=%d",
		hostname, dir, file, offset)
	bodyData, _ := FetchUrl(url)

	var logData LogData
	err = json.Unmarshal(bodyData, &logData)
	if err != nil {
		return nil, err
	}
	return []byte(logData.Data), nil
}
<|endoftext|>


<commit_before>package golevel7

import (
	"bufio"
	"golevel7/commons"
	"io"
)

type MessageScanner struct {
	r       io.Reader
	b       *bufio.Scanner
	thisMsg *Message
	err     error
}

func NewMessageScanner(r io.Reader) {
	ms := &MessageScanner{
		r: r,
		b: bufio.NewScanner(r),
	}
	ms.b.Split(commons.CrLfSplit)
}

func (ms *MessageScanner) Scan() (gotOne bool) {
	if ms.b.Scan() {
		if ms.err = ms.b.Err(); ms.err != nil {
			if ms.b.Bytes() != nil {
				gotOne = true
			}
		}
		if gotOne {
			ms.thisMsg = NewMessage(ms.b.Bytes())
		} else {
			ms.thisMsg = nil
		}
		return gotOne
	}
	return gotOne
}

func (ms *MessageScanner) Message() *Message {
	return ms.thisMsg
}

func (ms *MessageScanner) Err() error {
	return ms.err
}
<commit_msg>feat: implemented a scanner for messages from an io.Reader, better for large sets of messages, instead of returning a slice of all of them like Decoder<commit_after>package golevel7

import (
	"bufio"
	"golevel7/commons"
	"io"
)

type MessageScanner struct {
	r       io.Reader
	b       *bufio.Scanner
	thisMsg *Message
	err     error
}

// NewMessageScanner returns a new scanner that returns
// hl7 messages from an io.Reader
func NewMessageScanner(r io.Reader) *MessageScanner {
	ms := &MessageScanner{
		r: r,
		b: bufio.NewScanner(r),
	}
	ms.b.Split(commons.CrLfSplit)
	return ms
}

func (ms *MessageScanner) Scan() (gotOne bool) {
	if scan := ms.b.Scan(); scan {
		if ms.err = ms.b.Err(); ms.err != nil {
			if ms.b.Bytes() != nil {
				gotOne = true
			}
		} else {
			gotOne = true
		}
		if gotOne {
			ms.thisMsg = NewMessage(ms.b.Bytes())
		} else {
			ms.thisMsg = nil
		}
		return gotOne
	}
	return gotOne
}

func (ms *MessageScanner) Message() *Message {
	return ms.thisMsg
}

func (ms *MessageScanner) Err() error {
	return ms.err
}
<|endoftext|>


<commit_before>package trackello

import (
	"fmt"

	"github.com/VojtechVitek/go-trello"
	"github.com/fatih/color"
)

// statistics provides a way to show statistical information about a list, card or whatnot by aggregating the updates,
// comments, checklists, and other actions under a specific item.
type statistics struct {
	comments, // represented by a horizontal ellepsis ⋯ 0x22EF
	updates, // represented by a keyboard 0x2328
	creates, // represented by plus +
	checkListItemUpdates int // represented by check mark ✓ 0x2713
}

// Statistics represents the statistics for all the actions generated for a list, card, etc.
type Statistics interface {
	AddCalculation(trello.Action)
}

func (c *Card) AddCalculation(a trello.Action) {
	switch a.Type {
	case "commentCard":
		c.stats.comments++
	case "updateCard":
		c.stats.updates++
	case "createCard", "addChecklistToCard", "addAttachmentToCard":
		c.stats.creates++
	case "updateCheckItemStateOnCard":
		c.stats.checkListItemUpdates++
	default:
		fmt.Printf("Unmapped action type: %s.  Defaulting to update\n", a.Type)
		c.stats.updates++
	}
}

// PrintStatistics will print the statistics information out.
// Example format: [ 3 +  2 ≡  0 ✓  1 … ]
func (s *statistics) PrintStatistics() string {
	if s == nil {
		s = &statistics{}
	}
	stats := "[" + color.CyanString("%d +", s.updates)
	stats = stats + color.RedString(" %d ≡", s.comments)
	stats = stats + color.GreenString(" %d ✓", s.checkListItemUpdates)
	stats = stats + color.MagentaString(" %d …", s.creates)
	stats = stats + "]"
	return stats
}

func (l *List) AddCalculation(a trello.Action) {}
<commit_msg>simplify unused interfaces<commit_after>package trackello

import (
	"fmt"

	"github.com/VojtechVitek/go-trello"
	"github.com/fatih/color"
)

// statistics provides a way to show statistical information about a list, card or whatnot by aggregating the updates,
// comments, checklists, and other actions under a specific item.
type statistics struct {
	comments, // represented by a horizontal ellepsis ⋯ 0x22EF
	updates, // represented by a keyboard 0x2328
	creates, // represented by plus +
	checkListItemUpdates int // represented by check mark ✓ 0x2713
}

// AddCalculation will the new action to the Card's statistics.
func (c *Card) AddCalculation(a trello.Action) {
	switch a.Type {
	case "commentCard":
		c.stats.comments++
	case "updateCard":
		c.stats.updates++
	case "createCard", "addChecklistToCard", "addAttachmentToCard":
		c.stats.creates++
	case "updateCheckItemStateOnCard":
		c.stats.checkListItemUpdates++
	default:
		fmt.Printf("Unmapped action type: %s.  Defaulting to update\n", a.Type)
		c.stats.updates++
	}
}

// PrintStatistics will print the statistics information out.
// Example format: [ 3 +  2 ≡  0 ✓  1 … ]
func (s *statistics) PrintStatistics() string {
	if s == nil {
		s = &statistics{}
	}
	stats := "[" + color.CyanString("%d +", s.updates)
	stats = stats + color.RedString(" %d ≡", s.comments)
	stats = stats + color.GreenString(" %d ✓", s.checkListItemUpdates)
	stats = stats + color.MagentaString(" %d …", s.creates)
	stats = stats + "]"
	return stats
}
<|endoftext|>


<commit_before>package omap

type any interface{}

type Map interface {
	Set(key, value any)
	Get(key any) (any, bool)
	Delete(key any)
}

type omap struct {
	keys []any
	body map[any]any
}

var _ Map = &omap{}

func New() Map {
	return &omap{
		keys: make([]any, 0),
		body: make(map[any]any),
	}
}

func (m *omap) Set(key, value any) {
	if _, exist := m.body[key]; !exist {
		m.keys = append(m.keys, key)
	}
	m.body[key] = value
}

func (m *omap) Get(key any) (any, bool) {
	v, ok := m.body[key]
	return v, ok
}

func (m *omap) Delete(key any) {
	i := index(m.keys, key)
	if i == -1 {
		return
	}
	delete(m.body, key)
	m.keys = append(m.keys[:i], m.keys[i+1:]...)
}

func index(keys []any, key any) int {
	for i, k := range keys {
		if k == key {
			return i
		}
	}
	return -1
}
<commit_msg>key is a string<commit_after>package omap

type any interface{}

type Map interface {
	Set(key string, value any)
	Get(key string) (any, bool)
	Delete(key string)

	Each(func(any))
}

type omap struct {
	keys []string
	body map[string]any
}

var _ Map = &omap{}

func New() Map {
	return &omap{
		keys: make([]string, 0),
		body: make(map[string]any),
	}
}

func (m *omap) Set(key string, value any) {
	if _, exist := m.body[key]; !exist {
		m.keys = append(m.keys, key)
	}
	m.body[key] = value
}

func (m *omap) Get(key string) (any, bool) {
	v, ok := m.body[key]
	return v, ok
}

func (m *omap) Delete(key string) {
	i := index(m.keys, key)
	if i == -1 {
		return
	}
	delete(m.body, key)
	m.keys = append(m.keys[:i], m.keys[i+1:]...)
}

func index(keys []string, key string) int {
	for i, k := range keys {
		if k == key {
			return i
		}
	}
	return -1
}

func (m *omap) Each(f func(any)) {
	for _, v := range m.keys {
		f(m.body[v])
	}
}

func MarshalJSON(b []byte) error {
	return nil
}
<|endoftext|>


<commit_before>package coal

import (
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo/options"
)

// Tester provides facilities to work with coal models in tests.
type Tester struct {
	// The store to use for cleaning the database.
	Store *Store

	// The registered models.
	Models []Model
}

// NewTester returns a new tester.
func NewTester(store *Store, models ...Model) *Tester {
	return &Tester{
		Store:  store,
		Models: models,
	}
}

// Clean will remove the collections of models that have been registered and
// reset the header map.
func (t *Tester) Clean() {
	for _, model := range t.Models {
		// remove all is faster than dropping the collection
		_, err := t.Store.C(model).DeleteMany(nil, bson.M{})
		if err != nil {
			panic(err)
		}
	}
}

// Save will save the specified model.
func (t *Tester) Save(model Model) Model {
	// initialize model
	model = Init(model)

	// insert to collection
	_, err := t.Store.C(model).InsertOne(nil, model)
	if err != nil {
		panic(err)
	}

	return model
}

// FindAll will return all saved models.
func (t *Tester) FindAll(model Model) interface{} {
	// initialize model
	model = Init(model)

	// find all documents
	list := model.Meta().MakeSlice()
	cursor, err := t.Store.C(model).Find(nil, bson.M{}, options.Find().SetSort(Sort("_id")))
	if err != nil {
		panic(err)
	}

	// get all results
	err = cursor.All(nil, list)
	if err != nil {
		panic(err)
	}

	// close cursor
	err = cursor.Close(nil)
	if err != nil {
		panic(err)
	}

	// initialize list
	InitSlice(list)

	return list
}

// FindLast will return the last saved model.
func (t *Tester) FindLast(model Model) Model {
	// find last document
	err := t.Store.C(model).FindOne(nil, bson.M{}, options.FindOne().SetSort(Sort("-_id"))).Decode(model)
	if err != nil {
		panic(err)
	}

	// initialize model
	Init(model)

	return model
}

// Count will count all saved models.
func (t *Tester) Count(model Model) int {
	// count all documents
	n, err := t.Store.C(model).CountDocuments(nil, bson.M{})
	if err != nil {
		panic(err)
	}

	return int(n)
}

// Fetch will return the saved model.
func (t *Tester) Fetch(model Model, id ID) Model {
	// find specific document
	err := t.Store.C(model).FindOne(nil, bson.M{
		"_id": id,
	}).Decode(model)
	if err != nil {
		panic(err)
	}

	// initialize model
	Init(model)

	return model
}

// Update will update the specified model.
func (t *Tester) Update(model Model) Model {
	// initialize model
	model = Init(model)

	// insert to collection
	_, err := t.Store.C(model).ReplaceOne(nil, bson.M{
		"_id": model.ID(),
	}, model)
	if err != nil {
		panic(err)
	}

	return model
}

// Delete will delete the specified model.
func (t *Tester) Delete(model Model) {
	// initialize model
	model = Init(model)

	// insert to collection
	_, err := t.Store.C(model).DeleteOne(nil, bson.M{
		"_id": model.ID(),
	})
	if err != nil {
		panic(err)
	}
}
<commit_msg>added optional query parameters<commit_after>package coal

import (
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo/options"
)

// Tester provides facilities to work with coal models in tests.
type Tester struct {
	// The store to use for cleaning the database.
	Store *Store

	// The registered models.
	Models []Model
}

// NewTester returns a new tester.
func NewTester(store *Store, models ...Model) *Tester {
	return &Tester{
		Store:  store,
		Models: models,
	}
}

// Clean will remove the collections of models that have been registered and
// reset the header map.
func (t *Tester) Clean() {
	for _, model := range t.Models {
		// remove all is faster than dropping the collection
		_, err := t.Store.C(model).DeleteMany(nil, bson.M{})
		if err != nil {
			panic(err)
		}
	}
}

// Save will save the specified model.
func (t *Tester) Save(model Model) Model {
	// initialize model
	model = Init(model)

	// insert to collection
	_, err := t.Store.C(model).InsertOne(nil, model)
	if err != nil {
		panic(err)
	}

	return model
}

// FindAll will return all saved models.
func (t *Tester) FindAll(model Model, query ...bson.M) interface{} {
	// initialize model
	model = Init(model)

	// prepare query
	qry := bson.M{}
	if len(query) > 0 {
		qry = query[0]
	}

	// find all documents
	list := model.Meta().MakeSlice()
	cursor, err := t.Store.C(model).Find(nil, qry, options.Find().SetSort(Sort("_id")))
	if err != nil {
		panic(err)
	}

	// get all results
	err = cursor.All(nil, list)
	if err != nil {
		panic(err)
	}

	// close cursor
	err = cursor.Close(nil)
	if err != nil {
		panic(err)
	}

	// initialize list
	InitSlice(list)

	return list
}

// FindLast will return the last saved model.
func (t *Tester) FindLast(model Model, query ...bson.M) Model {
	// prepare query
	qry := bson.M{}
	if len(query) > 0 {
		qry = query[0]
	}

	// find last document
	err := t.Store.C(model).FindOne(nil, qry, options.FindOne().SetSort(Sort("-_id"))).Decode(model)
	if err != nil {
		panic(err)
	}

	// initialize model
	Init(model)

	return model
}

// Count will count all saved models.
func (t *Tester) Count(model Model, query ...bson.M) int {
	// prepare query
	qry := bson.M{}
	if len(query) > 0 {
		qry = query[0]
	}

	// count all documents
	n, err := t.Store.C(model).CountDocuments(nil, qry)
	if err != nil {
		panic(err)
	}

	return int(n)
}

// Fetch will return the saved model.
func (t *Tester) Fetch(model Model, id ID) Model {
	// find specific document
	err := t.Store.C(model).FindOne(nil, bson.M{
		"_id": id,
	}).Decode(model)
	if err != nil {
		panic(err)
	}

	// initialize model
	Init(model)

	return model
}

// Update will update the specified model.
func (t *Tester) Update(model Model) Model {
	// initialize model
	model = Init(model)

	// insert to collection
	_, err := t.Store.C(model).ReplaceOne(nil, bson.M{
		"_id": model.ID(),
	}, model)
	if err != nil {
		panic(err)
	}

	return model
}

// Delete will delete the specified model.
func (t *Tester) Delete(model Model) {
	// initialize model
	model = Init(model)

	// insert to collection
	_, err := t.Store.C(model).DeleteOne(nil, bson.M{
		"_id": model.ID(),
	})
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


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

package brk_hdl_http

import (
	"encoding/hex"
	"encoding/json"
	"fmt"
	"github.com/thethingsnetwork/core"
	"io"
	"net/http"
	"regexp"
	"strings"
)

var ErrInvalidPort = fmt.Errorf("The given port is invalid")

type Adapter struct {
	regs chan regMsg
}

type regMsg struct {
	config core.Registration
	resp   chan httpResponse
}

type httpResponse struct {
	statusCode int
	content    []byte
}

func NewAdapter(port uint) (*Adapter, error) {
	if port == 0 {
		return nil, ErrInvalidPort
	}

	a := Adapter{
		regs: make(chan regMsg),
	}
	return &a, nil
}

// Send implements the core.Adapter interface
func (a *Adapter) Send(p core.Packet, an core.AckNacker) error {
	return nil
}

// Next implements the core.Adapter inerface
func (a *Adapter) Next() (core.Packet, core.AckNacker, error) {
	return core.Packet{}, nil, nil
}

// NextRegistration implements the core.BrkHdlAdapter interface
func (a *Adapter) NextRegistration() (core.Registration, core.AckNacker, error) {
	return core.Registration{}, nil, nil
}

func (a *Adapter) listen(port uint) {
	serveMux := http.NewServeMux()
	serveMux.HandleFunc("/end-device/", func(w http.ResponseWriter, req *http.Request) {
		// Check Content-type
		contentType := req.Header.Get("Content-Type")
		if contentType != "application/json" {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Bad content type"))
			return
		}

		// Check the query parameter
		reg := regexp.MustCompile("end-device/([a-fA-F0-9]{10})$")
		query := reg.FindStringSubmatch(req.RequestURI)
		if len(query) < 2 {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Incorrect end-device address format"))
			return
		}
		devAddr := query[1]

		// Check configuration in body
		body := make([]byte, 256)
		n, err := req.Body.Read(body)
		if err != nil && err != io.EOF {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Incorrect request body"))
			return
		}
		body = body[:n]
		params := &struct {
			Id     string `json:"app_id"`
			Url    string `json:"app_url"`
			NwsKey string `json:"nws_key"`
		}{}
		if err := json.Unmarshal(body, params); err != nil {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Incorrect body payload"))
			return
		}

		nwsKey, err := hex.DecodeString(params.NwsKey)
		if err != nil || len(nwsKey) != 16 {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Incorrect network session nws_key"))
			return
		}

		params.Id = strings.Trim(params.Id, " ")
		params.Url = strings.Trim(params.Url, " ")
		if len(params.Id) <= 0 {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Incorrect config app_id"))
			return
		}
		if len(params.Url) <= 0 {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Incorrect config app_url"))
			return
		}

		// Create registration
		config := core.Registration{
			Handler: core.Recipient{Id: params.Id, Address: params.Url},
		}
		copy(config.NwsKey[:], nwsKey)
		copy(config.DevAddr[:], devAddr)

		// Send the registration and wait for ack / nack
		resp := make(chan httpResponse)
		a.regs <- regMsg{config: config, resp: resp}
		r, ok := <-resp
		if !ok {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte("Core server not responding"))
			return
		}
		w.WriteHeader(r.statusCode)
		w.Write(r.content)
	})

	go func() {
		server := http.Server{
			Addr:    fmt.Sprintf("localhost:%d", port),
			Handler: serveMux,
		}
		server.ListenAndServe()
	}()
}
<commit_msg>[broker] Add logger to adapter and start listening<commit_after>// Copyright © 2015 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package brk_hdl_http

import (
	"encoding/hex"
	"encoding/json"
	"fmt"
	"github.com/thethingsnetwork/core"
	"github.com/thethingsnetwork/core/utils/log"
	"io"
	"net/http"
	"regexp"
	"strings"
)

var ErrInvalidPort = fmt.Errorf("The given port is invalid")

type Adapter struct {
	logger log.Logger
	regs   chan regMsg
}

type regMsg struct {
	config core.Registration
	resp   chan httpResponse
}

type httpResponse struct {
	statusCode int
	content    []byte
}

func NewAdapter(port uint) (*Adapter, error) {
	if port == 0 {
		return nil, ErrInvalidPort
	}

	a := Adapter{
		regs:   make(chan regMsg),
		logger: log.DebugLogger{Tag: "BRK_HDL_ADAPTER"},
	}

	a.listen(port)

	return &a, nil
}

// Send implements the core.Adapter interface
func (a *Adapter) Send(p core.Packet, an core.AckNacker) error {
	return nil
}

// Next implements the core.Adapter inerface
func (a *Adapter) Next() (core.Packet, core.AckNacker, error) {
	return core.Packet{}, nil, nil
}

// NextRegistration implements the core.BrkHdlAdapter interface
func (a *Adapter) NextRegistration() (core.Registration, core.AckNacker, error) {
	return core.Registration{}, nil, nil
}

func (a *Adapter) listen(port uint) {
	serveMux := http.NewServeMux()
	serveMux.HandleFunc("/end-device/", func(w http.ResponseWriter, req *http.Request) {
		a.log("Receive new registration request")
		// Check Content-type
		contentType := req.Header.Get("Content-Type")
		if contentType != "application/json" {
			a.log("registration request rejected: not json")
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Bad content type"))
			return
		}

		// Check the query parameter
		reg := regexp.MustCompile("end-device/([a-fA-F0-9]{10})$")
		query := reg.FindStringSubmatch(req.RequestURI)
		if len(query) < 2 {
			a.log("registration request rejected: devAddr invalid")
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Incorrect end-device address format"))
			return
		}
		devAddr := query[1]

		// Check configuration in body
		body := make([]byte, 256)
		n, err := req.Body.Read(body)
		if err != nil && err != io.EOF {
			a.log("registration request rejected: body unreadable")
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Incorrect request body"))
			return
		}
		body = body[:n]
		params := &struct {
			Id     string `json:"app_id"`
			Url    string `json:"app_url"`
			NwsKey string `json:"nws_key"`
		}{}
		if err := json.Unmarshal(body, params); err != nil {
			a.log("registration request rejected: payload invalid or incomplete")
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Incorrect body payload"))
			return
		}

		nwsKey, err := hex.DecodeString(params.NwsKey)
		if err != nil || len(nwsKey) != 16 {
			a.log("registration request rejected: nwsKey invalid")
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Incorrect network session nws_key"))
			return
		}

		params.Id = strings.Trim(params.Id, " ")
		params.Url = strings.Trim(params.Url, " ")
		if len(params.Id) <= 0 {
			a.log("registration request rejected: appId invalid")
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Incorrect config app_id"))
			return
		}
		if len(params.Url) <= 0 {
			a.log("registration request rejected: appUrl invalid")
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte("Incorrect config app_url"))
			return
		}

		// Create registration
		config := core.Registration{
			Handler: core.Recipient{Id: params.Id, Address: params.Url},
		}
		copy(config.NwsKey[:], nwsKey)
		copy(config.DevAddr[:], devAddr)

		// Send the registration and wait for ack / nack
		resp := make(chan httpResponse)
		a.regs <- regMsg{config: config, resp: resp}
		r, ok := <-resp
		if !ok {
			a.log("Unexpected channel closure")
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte("Core server not responding"))
			return
		}
		a.log("sending normal response to registration request")
		w.WriteHeader(r.statusCode)
		w.Write(r.content)
	})

	go func() {
		server := http.Server{
			Addr:    fmt.Sprintf("localhost:%d", port),
			Handler: serveMux,
		}
		err := server.ListenAndServe()
		a.log("HTTP connection lost: %v", err)
	}()
}

func (a *Adapter) log(format string, i ...interface{}) {
	if a == nil || a.logger == nil {
		return
	}
	a.logger.Log(format, i...)
}
<|endoftext|>


<commit_before>package process_test

import (
	"fmt"
	"os"
	"os/signal"
	"reflect"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"testing"
	"time"

	"github.com/buildkite/agent/process"
)

func TestProcessRunsAndSignalsStartedAndStopped(t *testing.T) {
	var started int32
	var done int32

	p := process.Process{
		Script: []string{os.Args[0]},
		Env:    []string{"TEST_MAIN=tester"},
	}

	go func() {
		<-p.Started()
		atomic.AddInt32(&started, 1)
		<-p.Done()
		atomic.AddInt32(&done, 1)
	}()

	if err := p.Start(); err != nil {
		t.Fatal(err)
	}

	if startedVal := atomic.LoadInt32(&started); startedVal != 1 {
		t.Fatalf("Expected started to be 1, got %d", startedVal)
	}

	if doneVal := atomic.LoadInt32(&done); doneVal != 1 {
		t.Fatalf("Expected done to be 1, got %d", doneVal)
	}

	if exitStatus := p.ExitStatus; exitStatus != "0" {
		t.Fatalf("Expected ExitStatus of 0, got %v", exitStatus)
	}
}

func TestProcessCapturesOutputLineByLine(t *testing.T) {
	var lines []string
	var mu sync.Mutex

	p := process.Process{
		Script: []string{os.Args[0]},
		Env:    []string{"TEST_MAIN=tester"},
		Handler: func(line string) {
			mu.Lock()
			defer mu.Unlock()
			lines = append(lines, line)
		},
	}

	if err := p.Start(); err != nil {
		t.Error(err)
	}

	expected := []string{
		"+++ My header",
		"llamas",
		"and more llamas",
		"a very long line a very long line a very long line a very long line a very long line a very long line a very long line a very long line a very long line a very long line a very long line a very long line a very long line a very long line",
		"and some alpacas",
	}

	if !reflect.DeepEqual(expected, lines) {
		t.Fatalf("Unexpected lines: %v", lines)
	}
}

func TestProcessIsKilledGracefully(t *testing.T) {
	var lines []string
	var mu sync.Mutex

	p := process.Process{
		Script: []string{os.Args[0]},
		Env:    []string{"TEST_MAIN=tester-signal"},
		Handler: func(line string) {
			mu.Lock()
			defer mu.Unlock()
			lines = append(lines, line)
		},
	}

	go func() {
		<-p.Started()

		// give the signal handler some time to install
		time.Sleep(time.Millisecond * 50)

		p.Kill()
	}()

	if err := p.Start(); err != nil {
		t.Fatal(err)
	}

	mu.Lock()
	defer mu.Unlock()

	output := strings.Join(lines, "\n")
	if output != `SIG terminated` {
		t.Fatalf("Bad output: %q", output)
	}
}

// Invoked by `go test`, switch between helper and running tests based on env
func TestMain(m *testing.M) {
	switch os.Getenv("TEST_MAIN") {
	case "tester":
		for _, line := range strings.Split(strings.TrimSuffix(longTestOutput, "\n"), "\n") {
			fmt.Printf("%s\n", line)
			time.Sleep(time.Millisecond * 20)
		}
		os.Exit(0)

	case "tester-signal":
		signals := make(chan os.Signal, 1)
		signal.Notify(signals, os.Interrupt,
			syscall.SIGTERM,
			syscall.SIGINT,
		)
		sig := <-signals
		fmt.Printf("SIG %v", sig)
		os.Exit(0)

	default:
		os.Exit(m.Run())
	}
}
<commit_msg>Reflect kill signature change in tests<commit_after>package process_test

import (
	"fmt"
	"os"
	"os/signal"
	"reflect"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"testing"
	"time"

	"github.com/buildkite/agent/process"
)

func TestProcessRunsAndSignalsStartedAndStopped(t *testing.T) {
	var started int32
	var done int32

	p := process.Process{
		Script: []string{os.Args[0]},
		Env:    []string{"TEST_MAIN=tester"},
	}

	go func() {
		<-p.Started()
		atomic.AddInt32(&started, 1)
		<-p.Done()
		atomic.AddInt32(&done, 1)
	}()

	if err := p.Start(); err != nil {
		t.Fatal(err)
	}

	if startedVal := atomic.LoadInt32(&started); startedVal != 1 {
		t.Fatalf("Expected started to be 1, got %d", startedVal)
	}

	if doneVal := atomic.LoadInt32(&done); doneVal != 1 {
		t.Fatalf("Expected done to be 1, got %d", doneVal)
	}

	if exitStatus := p.ExitStatus; exitStatus != "0" {
		t.Fatalf("Expected ExitStatus of 0, got %v", exitStatus)
	}
}

func TestProcessCapturesOutputLineByLine(t *testing.T) {
	var lines []string
	var mu sync.Mutex

	p := process.Process{
		Script: []string{os.Args[0]},
		Env:    []string{"TEST_MAIN=tester"},
		Handler: func(line string) {
			mu.Lock()
			defer mu.Unlock()
			lines = append(lines, line)
		},
	}

	if err := p.Start(); err != nil {
		t.Error(err)
	}

	expected := []string{
		"+++ My header",
		"llamas",
		"and more llamas",
		"a very long line a very long line a very long line a very long line a very long line a very long line a very long line a very long line a very long line a very long line a very long line a very long line a very long line a very long line",
		"and some alpacas",
	}

	if !reflect.DeepEqual(expected, lines) {
		t.Fatalf("Unexpected lines: %v", lines)
	}
}

func TestProcessIsKilledGracefully(t *testing.T) {
	var lines []string
	var mu sync.Mutex

	p := process.Process{
		Script: []string{os.Args[0]},
		Env:    []string{"TEST_MAIN=tester-signal"},
		Handler: func(line string) {
			mu.Lock()
			defer mu.Unlock()
			lines = append(lines, line)
		},
	}

	go func() {
		<-p.Started()

		// give the signal handler some time to install
		time.Sleep(time.Millisecond * 50)

		p.Kill(time.Millisecond * 20)
	}()

	if err := p.Start(); err != nil {
		t.Fatal(err)
	}

	mu.Lock()
	defer mu.Unlock()

	output := strings.Join(lines, "\n")
	if output != `SIG terminated` {
		t.Fatalf("Bad output: %q", output)
	}
}

// Invoked by `go test`, switch between helper and running tests based on env
func TestMain(m *testing.M) {
	switch os.Getenv("TEST_MAIN") {
	case "tester":
		for _, line := range strings.Split(strings.TrimSuffix(longTestOutput, "\n"), "\n") {
			fmt.Printf("%s\n", line)
			time.Sleep(time.Millisecond * 20)
		}
		os.Exit(0)

	case "tester-signal":
		signals := make(chan os.Signal, 1)
		signal.Notify(signals, os.Interrupt,
			syscall.SIGTERM,
			syscall.SIGINT,
		)
		sig := <-signals
		fmt.Printf("SIG %v", sig)
		os.Exit(0)

	default:
		os.Exit(m.Run())
	}
}
<|endoftext|>


<commit_before>package middleware

import (
	"errors"
	"fmt"
	"net/http"
	"reflect"
	"strings"

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

type (
	// JWTConfig defines the config for JWT middleware.
	JWTConfig struct {
		// Skipper defines a function to skip middleware.
		Skipper Skipper

		// Signing key to validate token.
		// Required.
		SigningKey interface{}

		// Signing method, used to check token signing method.
		// Optional. Default value HS256.
		SigningMethod string

		// Context key to store user information from the token into context.
		// Optional. Default value "user".
		ContextKey string

		// Claims are extendable claims data defining token content.
		// Optional. Default value jwt.MapClaims
		Claims jwt.Claims

		// TokenLookup is a string in the form of "<source>:<name>" that is used
		// to extract token from the request.
		// Optional. Default value "header:Authorization".
		// Possible values:
		// - "header:<name>"
		// - "query:<name>"
		// - "cookie:<name>"
		TokenLookup string

		// AuthScheme to be used in the Authorization header.
		// Optional. Default value "Bearer".
		AuthScheme string

		keyFunc jwt.Keyfunc
	}

	jwtExtractor func(echo.Context) (string, error)
)

// Algorithms
const (
	AlgorithmHS256 = "HS256"
)

var (
	// DefaultJWTConfig is the default JWT auth middleware config.
	DefaultJWTConfig = JWTConfig{
		Skipper:       DefaultSkipper,
		SigningMethod: AlgorithmHS256,
		ContextKey:    "user",
		TokenLookup:   "header:" + echo.HeaderAuthorization,
		AuthScheme:    "Bearer",
		Claims:        jwt.MapClaims{},
	}
)

// JWT returns a JSON Web Token (JWT) auth middleware.
//
// For valid token, it sets the user in context and calls next handler.
// For invalid token, it returns "401 - Unauthorized" error.
// For missing token, it returns "400 - Bad Request" error.
//
// See: https://jwt.io/introduction
// See `JWTConfig.TokenLookup`
func JWT(key []byte) echo.MiddlewareFunc {
	c := DefaultJWTConfig
	c.SigningKey = key
	return JWTWithConfig(c)
}

// JWTWithConfig returns a JWT auth middleware with config.
// See: `JWT()`.
func JWTWithConfig(config JWTConfig) echo.MiddlewareFunc {
	// Defaults
	if config.Skipper == nil {
		config.Skipper = DefaultJWTConfig.Skipper
	}
	if config.SigningKey == nil {
		panic("echo: jwt middleware requires signing key")
	}
	if config.SigningMethod == "" {
		config.SigningMethod = DefaultJWTConfig.SigningMethod
	}
	if config.ContextKey == "" {
		config.ContextKey = DefaultJWTConfig.ContextKey
	}
	if config.Claims == nil {
		config.Claims = DefaultJWTConfig.Claims
	}
	if config.TokenLookup == "" {
		config.TokenLookup = DefaultJWTConfig.TokenLookup
	}
	if config.AuthScheme == "" {
		config.AuthScheme = DefaultJWTConfig.AuthScheme
	}
	config.keyFunc = func(t *jwt.Token) (interface{}, error) {
		// Check the signing method
		if t.Method.Alg() != config.SigningMethod {
			return nil, fmt.Errorf("Unexpected jwt signing method=%v", t.Header["alg"])
		}
		return config.SigningKey, nil
	}

	// Initialize
	parts := strings.Split(config.TokenLookup, ":")
	extractor := jwtFromHeader(parts[1], config.AuthScheme)
	switch parts[0] {
	case "query":
		extractor = jwtFromQuery(parts[1])
	case "cookie":
		extractor = jwtFromCookie(parts[1])
	}

	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			if config.Skipper(c) {
				return next(c)
			}

			auth, err := extractor(c)
			if err != nil {
				return echo.NewHTTPError(http.StatusBadRequest, err.Error())
			}
			token := new(jwt.Token)
			// Issue #647, #656
			if _, ok := config.Claims.(jwt.MapClaims); ok {
				token, err = jwt.Parse(auth, config.keyFunc)
			} else {
				claims := reflect.ValueOf(config.Claims).Interface().(jwt.Claims)
				token, err = jwt.ParseWithClaims(auth, claims, config.keyFunc)
			}
			if err == nil && token.Valid {
				// Store user information from token into context.
				c.Set(config.ContextKey, token)
				return next(c)
			}
			return echo.ErrUnauthorized
		}
	}
}

// jwtFromHeader returns a `jwtExtractor` that extracts token from the request header.
func jwtFromHeader(header string, authScheme string) jwtExtractor {
	return func(c echo.Context) (string, error) {
		auth := c.Request().Header.Get(header)
		l := len(authScheme)
		if len(auth) > l+1 && auth[:l] == authScheme {
			return auth[l+1:], nil
		}
		return "", errors.New("Missing or invalid jwt in the request header")
	}
}

// jwtFromQuery returns a `jwtExtractor` that extracts token from the query string.
func jwtFromQuery(param string) jwtExtractor {
	return func(c echo.Context) (string, error) {
		token := c.QueryParam(param)
		if token == "" {
			return "", errors.New("Missing jwt in the query string")
		}
		return token, nil
	}
}

// jwtFromCookie returns a `jwtExtractor` that extracts token from the named cookie.
func jwtFromCookie(name string) jwtExtractor {
	return func(c echo.Context) (string, error) {
		cookie, err := c.Cookie(name)
		if err != nil {
			return "", errors.New("Missing jwt in the cookie")
		}
		return cookie.Value, nil
	}
}
<commit_msg>Allow interface{} as key for JWT middleware argument (#957)<commit_after>package middleware

import (
	"errors"
	"fmt"
	"net/http"
	"reflect"
	"strings"

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

type (
	// JWTConfig defines the config for JWT middleware.
	JWTConfig struct {
		// Skipper defines a function to skip middleware.
		Skipper Skipper

		// Signing key to validate token.
		// Required.
		SigningKey interface{}

		// Signing method, used to check token signing method.
		// Optional. Default value HS256.
		SigningMethod string

		// Context key to store user information from the token into context.
		// Optional. Default value "user".
		ContextKey string

		// Claims are extendable claims data defining token content.
		// Optional. Default value jwt.MapClaims
		Claims jwt.Claims

		// TokenLookup is a string in the form of "<source>:<name>" that is used
		// to extract token from the request.
		// Optional. Default value "header:Authorization".
		// Possible values:
		// - "header:<name>"
		// - "query:<name>"
		// - "cookie:<name>"
		TokenLookup string

		// AuthScheme to be used in the Authorization header.
		// Optional. Default value "Bearer".
		AuthScheme string

		keyFunc jwt.Keyfunc
	}

	jwtExtractor func(echo.Context) (string, error)
)

// Algorithms
const (
	AlgorithmHS256 = "HS256"
)

var (
	// DefaultJWTConfig is the default JWT auth middleware config.
	DefaultJWTConfig = JWTConfig{
		Skipper:       DefaultSkipper,
		SigningMethod: AlgorithmHS256,
		ContextKey:    "user",
		TokenLookup:   "header:" + echo.HeaderAuthorization,
		AuthScheme:    "Bearer",
		Claims:        jwt.MapClaims{},
	}
)

// JWT returns a JSON Web Token (JWT) auth middleware.
//
// For valid token, it sets the user in context and calls next handler.
// For invalid token, it returns "401 - Unauthorized" error.
// For missing token, it returns "400 - Bad Request" error.
//
// See: https://jwt.io/introduction
// See `JWTConfig.TokenLookup`
func JWT(key interface{}) echo.MiddlewareFunc {
	c := DefaultJWTConfig
	c.SigningKey = key
	return JWTWithConfig(c)
}

// JWTWithConfig returns a JWT auth middleware with config.
// See: `JWT()`.
func JWTWithConfig(config JWTConfig) echo.MiddlewareFunc {
	// Defaults
	if config.Skipper == nil {
		config.Skipper = DefaultJWTConfig.Skipper
	}
	if config.SigningKey == nil {
		panic("echo: jwt middleware requires signing key")
	}
	if config.SigningMethod == "" {
		config.SigningMethod = DefaultJWTConfig.SigningMethod
	}
	if config.ContextKey == "" {
		config.ContextKey = DefaultJWTConfig.ContextKey
	}
	if config.Claims == nil {
		config.Claims = DefaultJWTConfig.Claims
	}
	if config.TokenLookup == "" {
		config.TokenLookup = DefaultJWTConfig.TokenLookup
	}
	if config.AuthScheme == "" {
		config.AuthScheme = DefaultJWTConfig.AuthScheme
	}
	config.keyFunc = func(t *jwt.Token) (interface{}, error) {
		// Check the signing method
		if t.Method.Alg() != config.SigningMethod {
			return nil, fmt.Errorf("Unexpected jwt signing method=%v", t.Header["alg"])
		}
		return config.SigningKey, nil
	}

	// Initialize
	parts := strings.Split(config.TokenLookup, ":")
	extractor := jwtFromHeader(parts[1], config.AuthScheme)
	switch parts[0] {
	case "query":
		extractor = jwtFromQuery(parts[1])
	case "cookie":
		extractor = jwtFromCookie(parts[1])
	}

	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			if config.Skipper(c) {
				return next(c)
			}

			auth, err := extractor(c)
			if err != nil {
				return echo.NewHTTPError(http.StatusBadRequest, err.Error())
			}
			token := new(jwt.Token)
			// Issue #647, #656
			if _, ok := config.Claims.(jwt.MapClaims); ok {
				token, err = jwt.Parse(auth, config.keyFunc)
			} else {
				claims := reflect.ValueOf(config.Claims).Interface().(jwt.Claims)
				token, err = jwt.ParseWithClaims(auth, claims, config.keyFunc)
			}
			if err == nil && token.Valid {
				// Store user information from token into context.
				c.Set(config.ContextKey, token)
				return next(c)
			}
			return echo.ErrUnauthorized
		}
	}
}

// jwtFromHeader returns a `jwtExtractor` that extracts token from the request header.
func jwtFromHeader(header string, authScheme string) jwtExtractor {
	return func(c echo.Context) (string, error) {
		auth := c.Request().Header.Get(header)
		l := len(authScheme)
		if len(auth) > l+1 && auth[:l] == authScheme {
			return auth[l+1:], nil
		}
		return "", errors.New("Missing or invalid jwt in the request header")
	}
}

// jwtFromQuery returns a `jwtExtractor` that extracts token from the query string.
func jwtFromQuery(param string) jwtExtractor {
	return func(c echo.Context) (string, error) {
		token := c.QueryParam(param)
		if token == "" {
			return "", errors.New("Missing jwt in the query string")
		}
		return token, nil
	}
}

// jwtFromCookie returns a `jwtExtractor` that extracts token from the named cookie.
func jwtFromCookie(name string) jwtExtractor {
	return func(c echo.Context) (string, error) {
		cookie, err := c.Cookie(name)
		if err != nil {
			return "", errors.New("Missing jwt in the cookie")
		}
		return cookie.Value, nil
	}
}
<|endoftext|>


<commit_before>package main

import (
	"reflect"
	"regexp"
	"testing"
)

var genMatcherTests = []struct {
	src string
	dst *regexp.Regexp
}{
	// one branch
	{`abc`, regexp.MustCompile(`(abc)`)},
	{`abcdef`, regexp.MustCompile(`(abcdef)`)},

	// multiple branches
	{`a,b`, regexp.MustCompile(`(a|b)`)},
	{`a,,b,c`, regexp.MustCompile(`(a||b|c)`)},
	{`a,bc,def`, regexp.MustCompile(`(a|bc|def)`)},
	{`,a`, regexp.MustCompile(`(|a)`)},
	{`a,`, regexp.MustCompile(`(a|)`)},
	{`,a,`, regexp.MustCompile(`(|a|)`)},

	// use escape
	{`a\,b`, regexp.MustCompile(`(a,b)`)},
	{`a\,bc\,def`, regexp.MustCompile(`(a,bc,def)`)},

	// multiple branches with escape
	{`a\,b,c`, regexp.MustCompile(`(a,b|c)`)},
	{`a,bc\,def`, regexp.MustCompile(`(a|bc,def)`)},

	// regexp quote
	{`a+b`, regexp.MustCompile(`(a\+b)`)},
	{`(a|bc)*def`, regexp.MustCompile(`(\(a\|bc\)\*def)`)},

	// unquote special values
	{`a\\bc`, regexp.MustCompile("(a\\\\bc)")},
	{`a\tb\,c`, regexp.MustCompile("(a\tb,c)")},
	{`a\tbc\n\ndef`, regexp.MustCompile("(a\tbc\n\ndef)")},

	// multiple sequenses
	{`a/b`, regexp.MustCompile("(a)(b)")},
	{`a//b/c`, regexp.MustCompile(`(a)()(b)(c)`)},
	{`a/bc/def`, regexp.MustCompile("(a)(bc)(def)")},
	{`a,b/c`, regexp.MustCompile("(a|b)(c)")},
	{`/a`, regexp.MustCompile(`()(a)`)},
	{`a/`, regexp.MustCompile(`(a)()`)},
	{`/a/`, regexp.MustCompile(`()(a)()`)},

	// multiple sequenses with escape
	{`a/b\/c`, regexp.MustCompile("(a)(b/c)")},
	{`a/\/bc\//def`, regexp.MustCompile("(a)(/bc/)(def)")},
	{`a\,b,c/d,e\/f`, regexp.MustCompile("(a,b|c)(d|e/f)")},
}

func TestGenMatcher(t *testing.T) {
	for _, test := range genMatcherTests {
		expect := test.dst
		actual, err := newMatcher(test.src)
		if err != nil {
			t.Errorf("newMatcher(%q) returns %q, want nil",
				test.src, err)
		}
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("%q: got %q, want %q",
				test.src, actual, expect)
		}
	}
}

var genReplacementTests = []struct {
	srcFrom string
	srcTo   string
	dst     []map[string]string
}{
	// one branch
	{
		"abc",
		"def",
		[]map[string]string{
			map[string]string{
				"abc": "def",
			},
		},
	},
	{
		"abcdef",
		"ghijkl",
		[]map[string]string{
			map[string]string{
				"abcdef": "ghijkl",
			},
		},
	},

	// multiple branches
	{
		"a,b",
		"b,a",
		[]map[string]string{
			map[string]string{
				"a": "b",
				"b": "a",
			},
		},
	},
	{
		"a,,b,c",
		"d,e,f,g",
		[]map[string]string{
			map[string]string{
				"a": "d",
				"":  "e",
				"b": "f",
				"c": "g",
			},
		},
	},
	{
		",a",
		"a,",
		[]map[string]string{
			map[string]string{"": "a", "a": ""},
		},
	},
	{
		"a,b,c",
		",d,",
		[]map[string]string{
			map[string]string{
				"a": "",
				"b": "d",
				"c": "",
			},
		},
	},

	// multiple sequenses
	{
		"a/b",
		"c/d",
		[]map[string]string{
			map[string]string{
				"a": "c",
			},
			map[string]string{
				"b": "d",
			},
		},
	},
	{
		"a//b/c",
		"d/e/f/g",
		[]map[string]string{
			map[string]string{
				"a": "d",
			},
			map[string]string{
				"": "e",
			},
			map[string]string{
				"b": "f",
			},
			map[string]string{
				"c": "g",
			},
		},
	},
	{
		"a,b/c",
		"d,e/f",
		[]map[string]string{
			map[string]string{
				"a": "d",
				"b": "e",
			},
			map[string]string{
				"c": "f",
			},
		},
	},
	{
		"/a",
		"a/",
		[]map[string]string{
			map[string]string{
				"": "a",
			},
			map[string]string{
				"a": "",
			},
		},
	},
	{
		"/a/",
		"b/c/d",
		[]map[string]string{
			map[string]string{
				"": "b",
			},
			map[string]string{
				"a": "c",
			},
			map[string]string{
				"": "d",
			},
		},
	},
}

func TestGenReplacement(t *testing.T) {
	for _, test := range genReplacementTests {
		expect := test.dst
		actual, err := newReplacement(test.srcFrom, test.srcTo)
		if err != nil {
			t.Errorf("newReplacement(%q, %q) returns %q, want nil",
				test.srcFrom, test.srcTo, err)
		}
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("%q, %q: got %q, want %q",
				test.srcFrom, test.srcTo, actual, expect)
		}
	}
}

var replaceTests = []struct {
	srcFrom string
	srcTo   string
	srcText string
	dst     string
}{
	// one branch
	{
		"abc",
		"def",
		"foo bar",
		"foo bar",
	},
	{
		"abc",
		"def",
		"abc def",
		"def def",
	},
	{
		"a",
		"b",
		"a b c a b c",
		"b b c b b c",
	},
}

func TestReplace(t *testing.T) {
	for _, test := range replaceTests {
		r, err := NewReplacer(test.srcFrom, test.srcTo)
		if err != nil {
			t.Errorf("NewReplacer(%q, %q) returns %q, want nil",
				test.srcFrom, test.srcTo, err)
		}

		expect := test.dst
		actual := r.ReplaceAll(test.srcText)
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("Replacer{%q, %q}: %q: got %q, want %q",
				test.srcFrom, test.srcTo, test.srcText, actual, expect)
		}
	}
}
<commit_msg>Add cases of multiple branch for Replace<commit_after>package main

import (
	"reflect"
	"regexp"
	"testing"
)

var genMatcherTests = []struct {
	src string
	dst *regexp.Regexp
}{
	// one branch
	{`abc`, regexp.MustCompile(`(abc)`)},
	{`abcdef`, regexp.MustCompile(`(abcdef)`)},

	// multiple branches
	{`a,b`, regexp.MustCompile(`(a|b)`)},
	{`a,,b,c`, regexp.MustCompile(`(a||b|c)`)},
	{`a,bc,def`, regexp.MustCompile(`(a|bc|def)`)},
	{`,a`, regexp.MustCompile(`(|a)`)},
	{`a,`, regexp.MustCompile(`(a|)`)},
	{`,a,`, regexp.MustCompile(`(|a|)`)},

	// use escape
	{`a\,b`, regexp.MustCompile(`(a,b)`)},
	{`a\,bc\,def`, regexp.MustCompile(`(a,bc,def)`)},

	// multiple branches with escape
	{`a\,b,c`, regexp.MustCompile(`(a,b|c)`)},
	{`a,bc\,def`, regexp.MustCompile(`(a|bc,def)`)},

	// regexp quote
	{`a+b`, regexp.MustCompile(`(a\+b)`)},
	{`(a|bc)*def`, regexp.MustCompile(`(\(a\|bc\)\*def)`)},

	// unquote special values
	{`a\\bc`, regexp.MustCompile("(a\\\\bc)")},
	{`a\tb\,c`, regexp.MustCompile("(a\tb,c)")},
	{`a\tbc\n\ndef`, regexp.MustCompile("(a\tbc\n\ndef)")},

	// multiple sequenses
	{`a/b`, regexp.MustCompile("(a)(b)")},
	{`a//b/c`, regexp.MustCompile(`(a)()(b)(c)`)},
	{`a/bc/def`, regexp.MustCompile("(a)(bc)(def)")},
	{`a,b/c`, regexp.MustCompile("(a|b)(c)")},
	{`/a`, regexp.MustCompile(`()(a)`)},
	{`a/`, regexp.MustCompile(`(a)()`)},
	{`/a/`, regexp.MustCompile(`()(a)()`)},

	// multiple sequenses with escape
	{`a/b\/c`, regexp.MustCompile("(a)(b/c)")},
	{`a/\/bc\//def`, regexp.MustCompile("(a)(/bc/)(def)")},
	{`a\,b,c/d,e\/f`, regexp.MustCompile("(a,b|c)(d|e/f)")},
}

func TestGenMatcher(t *testing.T) {
	for _, test := range genMatcherTests {
		expect := test.dst
		actual, err := newMatcher(test.src)
		if err != nil {
			t.Errorf("newMatcher(%q) returns %q, want nil",
				test.src, err)
		}
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("%q: got %q, want %q",
				test.src, actual, expect)
		}
	}
}

var genReplacementTests = []struct {
	srcFrom string
	srcTo   string
	dst     []map[string]string
}{
	// one branch
	{
		"abc",
		"def",
		[]map[string]string{
			map[string]string{
				"abc": "def",
			},
		},
	},
	{
		"abcdef",
		"ghijkl",
		[]map[string]string{
			map[string]string{
				"abcdef": "ghijkl",
			},
		},
	},

	// multiple branches
	{
		"a,b",
		"b,a",
		[]map[string]string{
			map[string]string{
				"a": "b",
				"b": "a",
			},
		},
	},
	{
		"a,,b,c",
		"d,e,f,g",
		[]map[string]string{
			map[string]string{
				"a": "d",
				"":  "e",
				"b": "f",
				"c": "g",
			},
		},
	},
	{
		",a",
		"a,",
		[]map[string]string{
			map[string]string{"": "a", "a": ""},
		},
	},
	{
		"a,b,c",
		",d,",
		[]map[string]string{
			map[string]string{
				"a": "",
				"b": "d",
				"c": "",
			},
		},
	},

	// multiple sequenses
	{
		"a/b",
		"c/d",
		[]map[string]string{
			map[string]string{
				"a": "c",
			},
			map[string]string{
				"b": "d",
			},
		},
	},
	{
		"a//b/c",
		"d/e/f/g",
		[]map[string]string{
			map[string]string{
				"a": "d",
			},
			map[string]string{
				"": "e",
			},
			map[string]string{
				"b": "f",
			},
			map[string]string{
				"c": "g",
			},
		},
	},
	{
		"a,b/c",
		"d,e/f",
		[]map[string]string{
			map[string]string{
				"a": "d",
				"b": "e",
			},
			map[string]string{
				"c": "f",
			},
		},
	},
	{
		"/a",
		"a/",
		[]map[string]string{
			map[string]string{
				"": "a",
			},
			map[string]string{
				"a": "",
			},
		},
	},
	{
		"/a/",
		"b/c/d",
		[]map[string]string{
			map[string]string{
				"": "b",
			},
			map[string]string{
				"a": "c",
			},
			map[string]string{
				"": "d",
			},
		},
	},
}

func TestGenReplacement(t *testing.T) {
	for _, test := range genReplacementTests {
		expect := test.dst
		actual, err := newReplacement(test.srcFrom, test.srcTo)
		if err != nil {
			t.Errorf("newReplacement(%q, %q) returns %q, want nil",
				test.srcFrom, test.srcTo, err)
		}
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("%q, %q: got %q, want %q",
				test.srcFrom, test.srcTo, actual, expect)
		}
	}
}

var replaceTests = []struct {
	srcFrom string
	srcTo   string
	srcText string
	dst     string
}{
	// one branch
	{
		"abc",
		"def",
		"foo bar",
		"foo bar",
	},
	{
		"abc",
		"def",
		"abc def",
		"def def",
	},
	{
		"a",
		"b",
		"a b c a b c",
		"b b c b b c",
	},

	// multiple branches
	{
		"abc,def",
		"def,abc",
		"abc def",
		"def abc",
	},
	{
		"a,b,c,d",
		"e,f,g,h",
		"d c b a",
		"h g f e",
	},
	{
		"a, ",
		" ,a",
		"a a a",
		" a a ",
	},
}

func TestReplace(t *testing.T) {
	for _, test := range replaceTests {
		r, err := NewReplacer(test.srcFrom, test.srcTo)
		if err != nil {
			t.Errorf("NewReplacer(%q, %q) returns %q, want nil",
				test.srcFrom, test.srcTo, err)
		}

		expect := test.dst
		actual := r.ReplaceAll(test.srcText)
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("Replacer{%q, %q}: %q: got %q, want %q",
				test.srcFrom, test.srcTo, test.srcText, actual, expect)
		}
	}
}
<|endoftext|>


<commit_before>package orm

import (
	"bytes"
	"fmt"
	"sort"
	"strconv"
	"strings"

	"github.com/go-pg/pg/internal/parser"
	"github.com/go-pg/pg/types"
)

var formatter Formatter

type FormatAppender interface {
	AppendFormat([]byte, QueryFormatter) []byte
}

type sepFormatAppender interface {
	FormatAppender
	AppendSep([]byte) []byte
}

//------------------------------------------------------------------------------

type queryParamsAppender struct {
	query  string
	params []interface{}
}

var _ FormatAppender = (*queryParamsAppender)(nil)

func Q(query string, params ...interface{}) queryParamsAppender {
	return queryParamsAppender{query, params}
}

func (q queryParamsAppender) AppendFormat(b []byte, f QueryFormatter) []byte {
	return f.FormatQuery(b, q.query, q.params...)
}

func (q queryParamsAppender) AppendValue(b []byte, quote int) ([]byte, error) {
	return q.AppendFormat(b, formatter), nil
}

//------------------------------------------------------------------------------

type whereGroupAppender struct {
	where []sepFormatAppender
}

var _ FormatAppender = (*whereAppender)(nil)
var _ sepFormatAppender = (*whereAppender)(nil)

func (q whereGroupAppender) AppendSep(b []byte) []byte {
	return append(b, "AND"...)
}

func (q whereGroupAppender) AppendFormat(b []byte, f QueryFormatter) []byte {
	b = append(b, '(')
	for i, app := range q.where {
		if i > 0 {
			b = append(b, ' ')
			b = app.AppendSep(b)
			b = append(b, ' ')
		}
		b = app.AppendFormat(b, f)
	}
	b = append(b, ')')
	return b
}

//------------------------------------------------------------------------------

type whereAppender struct {
	conj   string
	query  string
	params []interface{}
}

var _ FormatAppender = (*whereAppender)(nil)
var _ sepFormatAppender = (*whereAppender)(nil)

func (q whereAppender) AppendSep(b []byte) []byte {
	return append(b, q.conj...)
}

func (q whereAppender) AppendFormat(b []byte, f QueryFormatter) []byte {
	b = append(b, '(')
	b = f.FormatQuery(b, q.query, q.params...)
	b = append(b, ')')
	return b
}

//------------------------------------------------------------------------------

type fieldAppender struct {
	field string
}

var _ FormatAppender = (*fieldAppender)(nil)

func (a fieldAppender) AppendFormat(b []byte, f QueryFormatter) []byte {
	return types.AppendField(b, a.field, 1)
}

//------------------------------------------------------------------------------

type Formatter struct {
	namedParams map[string]interface{}
}

func (f Formatter) String() string {
	if len(f.namedParams) == 0 {
		return ""
	}

	var keys []string
	for k, _ := range f.namedParams {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	var ss []string
	for _, k := range keys {
		ss = append(ss, fmt.Sprintf("%s=%v", k, f.namedParams[k]))
	}
	return " " + strings.Join(ss, " ")
}

func (f Formatter) copy() Formatter {
	var cp Formatter
	for param, value := range f.namedParams {
		cp.SetParam(param, value)
	}
	return cp
}

func (f *Formatter) SetParam(param string, value interface{}) {
	if f.namedParams == nil {
		f.namedParams = make(map[string]interface{})
	}
	f.namedParams[param] = value
}

func (f *Formatter) WithParam(param string, value interface{}) Formatter {
	cp := f.copy()
	cp.SetParam(param, value)
	return cp
}

func (f Formatter) Append(dst []byte, src string, params ...interface{}) []byte {
	if (params == nil && f.namedParams == nil) || strings.IndexByte(src, '?') == -1 {
		return append(dst, src...)
	}
	return f.append(dst, parser.NewString(src), params)
}

func (f Formatter) AppendBytes(dst, src []byte, params ...interface{}) []byte {
	if (params == nil && f.namedParams == nil) || bytes.IndexByte(src, '?') == -1 {
		return append(dst, src...)
	}
	return f.append(dst, parser.New(src), params)
}

func (f Formatter) FormatQuery(dst []byte, query string, params ...interface{}) []byte {
	return f.Append(dst, query, params...)
}

func (f Formatter) append(dst []byte, p *parser.Parser, params []interface{}) []byte {
	var paramsIndex int
	var namedParams *tableParams
	var namedParamsInit bool
	var model tableModel

	if len(params) > 0 {
		var ok bool
		model, ok = params[len(params)-1].(tableModel)
		if ok {
			params = params[:len(params)-1]
		}
	}

	for p.Valid() {
		b, ok := p.ReadSep('?')
		if !ok {
			dst = append(dst, b...)
			continue
		}
		if len(b) > 0 && b[len(b)-1] == '\\' {
			dst = append(dst, b[:len(b)-1]...)
			dst = append(dst, '?')
			continue
		}
		dst = append(dst, b...)

		if id, numeric := p.ReadIdentifier(); id != "" {
			if numeric {
				idx, err := strconv.Atoi(id)
				if err != nil {
					goto restore_param
				}

				if idx >= len(params) {
					goto restore_param
				}

				dst = f.appendParam(dst, params[idx])
				continue
			}

			if f.namedParams != nil {
				if param, ok := f.namedParams[id]; ok {
					dst = f.appendParam(dst, param)
					continue
				}
			}

			if !namedParamsInit && len(params) > 0 {
				namedParams, ok = newTableParams(params[len(params)-1])
				if ok {
					params = params[:len(params)-1]
				}
				namedParamsInit = true
			}

			if namedParams != nil {
				dst, ok = namedParams.AppendParam(dst, id)
				if ok {
					continue
				}
			}

			if model != nil {
				dst, ok = model.AppendParam(dst, id)
				if ok {
					continue
				}
			}

		restore_param:
			dst = append(dst, '?')
			dst = append(dst, id...)
			continue
		}

		if paramsIndex >= len(params) {
			dst = append(dst, '?')
			continue
		}

		param := params[paramsIndex]
		paramsIndex++

		if fa, ok := param.(FormatAppender); ok {
			dst = fa.AppendFormat(dst, f)
		} else {
			dst = types.Append(dst, param, 1)
		}
	}

	return dst
}

func (f Formatter) appendParam(b []byte, param interface{}) []byte {
	if fa, ok := param.(FormatAppender); ok {
		return fa.AppendFormat(b, f)
	}
	return types.Append(b, param, 1)
}
<commit_msg>orm: better var names<commit_after>package orm

import (
	"bytes"
	"fmt"
	"sort"
	"strconv"
	"strings"

	"github.com/go-pg/pg/internal/parser"
	"github.com/go-pg/pg/types"
)

var formatter Formatter

type FormatAppender interface {
	AppendFormat([]byte, QueryFormatter) []byte
}

type sepFormatAppender interface {
	FormatAppender
	AppendSep([]byte) []byte
}

//------------------------------------------------------------------------------

type queryParamsAppender struct {
	query  string
	params []interface{}
}

var _ FormatAppender = (*queryParamsAppender)(nil)

func Q(query string, params ...interface{}) queryParamsAppender {
	return queryParamsAppender{query, params}
}

func (q queryParamsAppender) AppendFormat(b []byte, f QueryFormatter) []byte {
	return f.FormatQuery(b, q.query, q.params...)
}

func (q queryParamsAppender) AppendValue(b []byte, quote int) ([]byte, error) {
	return q.AppendFormat(b, formatter), nil
}

//------------------------------------------------------------------------------

type whereGroupAppender struct {
	where []sepFormatAppender
}

var _ FormatAppender = (*whereAppender)(nil)
var _ sepFormatAppender = (*whereAppender)(nil)

func (q whereGroupAppender) AppendSep(b []byte) []byte {
	return append(b, "AND"...)
}

func (q whereGroupAppender) AppendFormat(b []byte, f QueryFormatter) []byte {
	b = append(b, '(')
	for i, app := range q.where {
		if i > 0 {
			b = append(b, ' ')
			b = app.AppendSep(b)
			b = append(b, ' ')
		}
		b = app.AppendFormat(b, f)
	}
	b = append(b, ')')
	return b
}

//------------------------------------------------------------------------------

type whereAppender struct {
	conj   string
	query  string
	params []interface{}
}

var _ FormatAppender = (*whereAppender)(nil)
var _ sepFormatAppender = (*whereAppender)(nil)

func (q whereAppender) AppendSep(b []byte) []byte {
	return append(b, q.conj...)
}

func (q whereAppender) AppendFormat(b []byte, f QueryFormatter) []byte {
	b = append(b, '(')
	b = f.FormatQuery(b, q.query, q.params...)
	b = append(b, ')')
	return b
}

//------------------------------------------------------------------------------

type fieldAppender struct {
	field string
}

var _ FormatAppender = (*fieldAppender)(nil)

func (a fieldAppender) AppendFormat(b []byte, f QueryFormatter) []byte {
	return types.AppendField(b, a.field, 1)
}

//------------------------------------------------------------------------------

type Formatter struct {
	namedParams map[string]interface{}
}

func (f Formatter) String() string {
	if len(f.namedParams) == 0 {
		return ""
	}

	var keys []string
	for k, _ := range f.namedParams {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	var ss []string
	for _, k := range keys {
		ss = append(ss, fmt.Sprintf("%s=%v", k, f.namedParams[k]))
	}
	return " " + strings.Join(ss, " ")
}

func (f Formatter) copy() Formatter {
	var cp Formatter
	for param, value := range f.namedParams {
		cp.SetParam(param, value)
	}
	return cp
}

func (f *Formatter) SetParam(param string, value interface{}) {
	if f.namedParams == nil {
		f.namedParams = make(map[string]interface{})
	}
	f.namedParams[param] = value
}

func (f *Formatter) WithParam(param string, value interface{}) Formatter {
	cp := f.copy()
	cp.SetParam(param, value)
	return cp
}

func (f Formatter) Append(dst []byte, src string, params ...interface{}) []byte {
	if (params == nil && f.namedParams == nil) || strings.IndexByte(src, '?') == -1 {
		return append(dst, src...)
	}
	return f.append(dst, parser.NewString(src), params)
}

func (f Formatter) AppendBytes(dst, src []byte, params ...interface{}) []byte {
	if (params == nil && f.namedParams == nil) || bytes.IndexByte(src, '?') == -1 {
		return append(dst, src...)
	}
	return f.append(dst, parser.New(src), params)
}

func (f Formatter) FormatQuery(dst []byte, query string, params ...interface{}) []byte {
	return f.Append(dst, query, params...)
}

func (f Formatter) append(dst []byte, p *parser.Parser, params []interface{}) []byte {
	var paramsIndex int
	var tableParamsOnce bool
	var tableParams *tableParams
	var model tableModel

	if len(params) > 0 {
		var ok bool
		model, ok = params[len(params)-1].(tableModel)
		if ok {
			params = params[:len(params)-1]
		}
	}

	for p.Valid() {
		b, ok := p.ReadSep('?')
		if !ok {
			dst = append(dst, b...)
			continue
		}
		if len(b) > 0 && b[len(b)-1] == '\\' {
			dst = append(dst, b[:len(b)-1]...)
			dst = append(dst, '?')
			continue
		}
		dst = append(dst, b...)

		if id, numeric := p.ReadIdentifier(); id != "" {
			if numeric {
				idx, err := strconv.Atoi(id)
				if err != nil {
					goto restore_param
				}

				if idx >= len(params) {
					goto restore_param
				}

				dst = f.appendParam(dst, params[idx])
				continue
			}

			if f.namedParams != nil {
				if param, ok := f.namedParams[id]; ok {
					dst = f.appendParam(dst, param)
					continue
				}
			}

			if !tableParamsOnce && len(params) > 0 {
				tableParams, ok = newTableParams(params[len(params)-1])
				if ok {
					params = params[:len(params)-1]
				}
				tableParamsOnce = true
			}

			if tableParams != nil {
				dst, ok = tableParams.AppendParam(dst, id)
				if ok {
					continue
				}
			}

			if model != nil {
				dst, ok = model.AppendParam(dst, id)
				if ok {
					continue
				}
			}

		restore_param:
			dst = append(dst, '?')
			dst = append(dst, id...)
			continue
		}

		if paramsIndex >= len(params) {
			dst = append(dst, '?')
			continue
		}

		param := params[paramsIndex]
		paramsIndex++

		if fa, ok := param.(FormatAppender); ok {
			dst = fa.AppendFormat(dst, f)
		} else {
			dst = types.Append(dst, param, 1)
		}
	}

	return dst
}

func (f Formatter) appendParam(b []byte, param interface{}) []byte {
	if fa, ok := param.(FormatAppender); ok {
		return fa.AppendFormat(b, f)
	}
	return types.Append(b, param, 1)
}
<|endoftext|>


<commit_before>package test

import (
	"github.com/icrowley/fake"

	"testing"
)

func TestSetLang(t *testing.T) {
	err := fake.SetLang("ru")
	if err != nil {
		t.Error("SetLang should successfully set lang")
	}

	err = fake.SetLang("sd")
	if err == nil {
		t.Error("SetLang with nonexistent lang should return error")
	}
}

func TestFakerRuWithoutCallback(t *testing.T) {
	fake.SetLang("ru")
	fake.EnFallback(false)
	brand := fake.Brand()
	if brand != "" {
		t.Error("Fake call with no samples should return blank string")
	}
}

func TestFakerRuWithCallback(t *testing.T) {
	fake.SetLang("ru")
	fake.EnFallback(true)
	brand := fake.Brand()
	if brand == "" {
		t.Error("Fake call for name with no samples with callback should not return blank string")
	}
}

// TestConcurrentSafety runs fake methods in multiple go routines concurrently.
// This test should be run with the race detector enabled.
func TestConcurrentSafety(t *testing.T) {
	workerCount := 10
	doneChan := make(chan struct{})

	for i := 0; i < workerCount; i++ {
		go func() {
			for j := 0; j < 1000; j++ {
				fake.FirstName()
			}
			doneChan <- struct{}{}
		}()
	}

	for i := 0; i < workerCount; i++ {
		<-doneChan
	}
}
<commit_msg>Test more functions in concurrent safety test<commit_after>package test

import (
	"github.com/icrowley/fake"

	"testing"
)

func TestSetLang(t *testing.T) {
	err := fake.SetLang("ru")
	if err != nil {
		t.Error("SetLang should successfully set lang")
	}

	err = fake.SetLang("sd")
	if err == nil {
		t.Error("SetLang with nonexistent lang should return error")
	}
}

func TestFakerRuWithoutCallback(t *testing.T) {
	fake.SetLang("ru")
	fake.EnFallback(false)
	brand := fake.Brand()
	if brand != "" {
		t.Error("Fake call with no samples should return blank string")
	}
}

func TestFakerRuWithCallback(t *testing.T) {
	fake.SetLang("ru")
	fake.EnFallback(true)
	brand := fake.Brand()
	if brand == "" {
		t.Error("Fake call for name with no samples with callback should not return blank string")
	}
}

// TestConcurrentSafety runs fake methods in multiple go routines concurrently.
// This test should be run with the race detector enabled.
func TestConcurrentSafety(t *testing.T) {
	workerCount := 10
	doneChan := make(chan struct{})

	for i := 0; i < workerCount; i++ {
		go func() {
			for j := 0; j < 1000; j++ {
				fake.FirstName()
				fake.LastName()
				fake.Gender()
				fake.FullName()
				fake.Day()
				fake.Country()
				fake.Company()
				fake.Industry()
				fake.Street()
			}
			doneChan <- struct{}{}
		}()
	}

	for i := 0; i < workerCount; i++ {
		<-doneChan
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"gopkg.in/yaml.v2"
	"html/template"
	"io/ioutil"
	"os"
	"path/filepath"
)

type Page struct {
	TemplatePath string `yaml:"template"`
	TitlePath    string `yaml:"title"`
	TitleRaw     string `yaml:"title-raw"`
	HeaderPath   string `yaml:"header"`
	HeaderRaw    string `yaml:"header-raw"`
	BodyPath     string `yaml:"body"`
	BodyRaw      string `yaml:"body-raw"`
	FooterPath   string `yaml:"footer"`
	FooterRaw    string `yaml:"footer-raw"`
	Output       string `yaml:"output-html"`
}

func (p Page) String() string {
	return fmt.Sprintf("%v\n%s\n%s\n%s\n", p.Title, p.Header, p.Body, p.Footer)
}

func (page *Page) gethtml(raw *string, path *string) template.HTML {
	var html string
	if &raw == nil {
		htmlbytes, _ := ioutil.ReadFile(*path)
		html = string(htmlbytes)
	} else {
		html = *raw
	}
	return template.HTML(html)
}

func (page *Page) Title() template.HTML {
	return page.gethtml(&page.TitleRaw, &page.TitlePath)
}

func (page *Page) Header() template.HTML {
	return page.gethtml(&page.HeaderRaw, &page.HeaderPath)
}

func (page *Page) Body() template.HTML {
	return page.gethtml(&page.BodyRaw, &page.BodyPath)
}

func (page *Page) Footer() template.HTML {
	return page.gethtml(&page.FooterRaw, &page.FooterPath)
}

func readconfig(path string) (*Page, error) {
	var page Page
	var err error

	data, err := ioutil.ReadFile(path)
	if err != nil {
		return &page, err
	}
	err = yaml.Unmarshal(data, &page)
	return &page, err
}

func buildpage(templatefile string, fout *os.File, page *Page) error {
	templatebytes, err := ioutil.ReadFile(templatefile)
	if err != nil {
		return err
	}
	templatename := templatefile[:len(templatefile)-5]
	t := template.Must(template.New(templatename).Parse(string(templatebytes)))
	err = t.Execute(fout, page)
	return err
}

func main() {

	var args []string
	args = os.Args[1:]
	if len(args) == 0 {
		fmt.Println("At least one content file must be specified")
	}

	var pagePtr *Page
	var fnm_html, templatepath string
	for _, fnm := range args {
		if _, err := os.Stat(fnm); !os.IsNotExist(err) {

			fmt.Println("parsing", fnm)
			pagePtr, err = readconfig(fnm)
			if err != nil {
				fmt.Println(err)
				break
			}
			fmt.Println("using template at", pagePtr.TemplatePath)

			if &pagePtr.Output != nil {
				fnm_html = pagePtr.Output
			} else {
				fnm_html = filepath.Base(fnm[:len(fnm)-5]) + ".html"
			}

			fmt.Println("writing to", fnm_html)
			fout, err := os.Create(fnm_html)
			defer fout.Close()
			if err != nil {
				fmt.Println(err)
				break
			}

			templatepath = filepath.Join(filepath.Dir(fnm), pagePtr.TemplatePath)
			err = buildpage(templatepath, fout, pagePtr)
			if err != nil {
				fmt.Println(err)
			}
		} else {
			fmt.Println(fnm, "does not exist")
		}
	}
}
<commit_msg>compute include paths; check null strings properly<commit_after>package main

import (
	"fmt"
	"gopkg.in/yaml.v2"
	"html/template"
	"io/ioutil"
	"os"
	"path/filepath"
)

type Page struct {
	TemplatePath string `yaml:"template"`
	TitlePath    string `yaml:"title"`
	TitleRaw     string `yaml:"title-raw"`
	HeaderPath   string `yaml:"header"`
	HeaderRaw    string `yaml:"header-raw"`
	BodyPath     string `yaml:"body"`
	BodyRaw      string `yaml:"body-raw"`
	FooterPath   string `yaml:"footer"`
	FooterRaw    string `yaml:"footer-raw"`
	Output       string `yaml:"output-html,omitempty"`
}

func (p Page) String() string {
	return fmt.Sprintf("%v\n%s\n%s\n%s\n", p.Title, p.Header, p.Body, p.Footer)
}

func (page *Page) gethtml(raw *string, path *string) template.HTML {
	var html string
	if len(*raw) == 0 {
		htmlbytes, err := ioutil.ReadFile(*path)
		if err != nil {
			fmt.Println(err)
		}
		html = string(htmlbytes)
	} else {
		html = *raw
	}
	return template.HTML(html)
}

func (page *Page) Title() template.HTML {
	return page.gethtml(&page.TitleRaw, &page.TitlePath)
}

func (page *Page) Header() template.HTML {
	return page.gethtml(&page.HeaderRaw, &page.HeaderPath)
}

func (page *Page) Body() template.HTML {
	return page.gethtml(&page.BodyRaw, &page.BodyPath)
}

func (page *Page) Footer() template.HTML {
	return page.gethtml(&page.FooterRaw, &page.FooterPath)
}

func readconfig(path string) (*Page, error) {
	var page Page
	var err error

	data, err := ioutil.ReadFile(path)
	if err != nil {
		return &page, err
	}
	err = yaml.Unmarshal(data, &page)

	basepath := filepath.Dir(path)
	if len(page.TitlePath) != 0 {
		page.TitlePath = filepath.Join(basepath, page.TitlePath)
	}
	if len(page.HeaderPath) != 0 {
		page.HeaderPath = filepath.Join(basepath, page.HeaderPath)
	}
	if len(page.BodyPath) != 0 {
		page.BodyPath = filepath.Join(basepath, page.BodyPath)
	}
	if len(page.FooterPath) != 0 {
		page.FooterPath = filepath.Join(basepath, page.FooterPath)
	}

	return &page, err
}

func buildpage(templatefile string, fout *os.File, page *Page) error {
	templatebytes, err := ioutil.ReadFile(templatefile)
	if err != nil {
		return err
	}
	templatename := templatefile[:len(templatefile)-5]
	t := template.Must(template.New(templatename).Parse(string(templatebytes)))
	err = t.Execute(fout, page)
	return err
}

func main() {

	var args []string
	args = os.Args[1:]
	if len(args) == 0 {
		fmt.Println("At least one content file must be specified")
	}

	var pagePtr *Page
	var fnm_html, templatepath string
	for _, fnm := range args {
		if _, err := os.Stat(fnm); !os.IsNotExist(err) {

			fmt.Println("parsing", fnm)
			pagePtr, err = readconfig(fnm)
			if err != nil {
				fmt.Println(err)
				break
			}
			fmt.Println("using template at", pagePtr.TemplatePath)

			if len(pagePtr.Output) == 0 {
				fnm_html = filepath.Base(fnm[:len(fnm)-5]) + ".html"
			} else {
				fnm_html = pagePtr.Output
			}

			fmt.Println("writing to", fnm_html)
			fout, err := os.Create(fnm_html)
			if err != nil {
				fmt.Println(err)
				break
			}

			templatepath = filepath.Join(filepath.Dir(fnm), pagePtr.TemplatePath)
			err = buildpage(templatepath, fout, pagePtr)
			if err != nil {
				fmt.Println(err)
			}
		} else {
			fmt.Println(fnm, "does not exist")
		}
	}
}
<|endoftext|>