<commit_before>// +build linux solaris darwin freebsd

package storage

import (
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/containers/storage/pkg/stringid"
	"github.com/pkg/errors"
	"golang.org/x/sys/unix"
)

func getLockFile(path string, ro bool) (Locker, error) {
	var fd int
	var err error
	if ro {
		fd, err = unix.Open(path, os.O_RDONLY, 0)
	} else {
		fd, err = unix.Open(path, os.O_RDWR|os.O_CREATE, unix.S_IRUSR|unix.S_IWUSR)
	}
	if err != nil {
		return nil, errors.Wrapf(err, "error opening %q", path)
	}
	unix.CloseOnExec(fd)

	locktype := unix.F_WRLCK
	if ro {
		locktype = unix.F_RDLCK
	}
	return &lockfile{
		stateMutex: &sync.Mutex{},
		writeMutex: &sync.Mutex{},
		file:       path,
		fd:         uintptr(fd),
		lw:         stringid.GenerateRandomID(),
		locktype:   int16(locktype),
		locked:     false,
		ro:         ro}, nil
}

type lockfile struct {
	// stateMutex is used to synchronize concurrent accesses
	stateMutex *sync.Mutex
	// writeMutex is used to serialize and avoid recursive writer locks
	writeMutex *sync.Mutex
	counter    int64
	file       string
	fd         uintptr
	lw         string
	locktype   int16
	locked     bool
	ro         bool
}

// lock locks the lockfile via FCTNL(2) based on the specified type and
// command.
func (l *lockfile) lock(l_type int16) {
	lk := unix.Flock_t{
		Type:   l_type,
		Whence: int16(os.SEEK_SET),
		Start:  0,
		Len:    0,
		Pid:    int32(os.Getpid()),
	}
	if l_type == unix.F_WRLCK {
		// If we try to lock as a writer, lock the writerMutex first to
		// avoid multiple writer acquisitions of the same process.
		// Note: it's important to lock it prior to the stateMutex to
		// avoid a deadlock.
		l.writeMutex.Lock()
	}
	l.stateMutex.Lock()
	l.locktype = l_type
	if l.counter == 0 {
		// Optimization: only use the (expensive) fcntl syscall when
		// the counter is 0.  If it's greater than that, we're owning
		// the lock already and can only be a reader.
		for unix.FcntlFlock(l.fd, unix.F_SETLKW, &lk) != nil {
			time.Sleep(10 * time.Millisecond)
		}
	}
	l.locked = true
	l.counter++
	l.stateMutex.Unlock()
}

// Lock locks the lockfile as a writer.  Note that RLock() will be called if
// the lock is a read-only one.
func (l *lockfile) Lock() {
	if l.ro {
		l.RLock()
	} else {
		l.lock(unix.F_WRLCK)
	}
}

// LockRead locks the lockfile as a reader.
func (l *lockfile) RLock() {
	l.lock(unix.F_RDLCK)
}

// Unlock unlocks the lockfile.
func (l *lockfile) Unlock() {
	lk := unix.Flock_t{
		Type:   unix.F_UNLCK,
		Whence: int16(os.SEEK_SET),
		Start:  0,
		Len:    0,
		Pid:    int32(os.Getpid()),
	}
	l.stateMutex.Lock()
	if l.locked == false {
		// Panic when unlocking an unlocked lock.  That's a violation
		// of the lock semantics and will reveal such.
		panic("calling Unlock on unlocked lock")
	}
	l.counter--
	if l.counter < 0 {
		// Panic when the counter is negative.  There is no way we can
		// recover from a corrupted lock and we need to protect the
		// storage from corruption.
		panic(fmt.Sprintf("lock %q has been unlocked too often", l.file))
	}
	if l.counter == 0 {
		// We should only release the lock when the counter is 0 to
		// avoid releasing read-locks too early; a given process may
		// acquire a read lock multiple times.
		l.locked = false
		for unix.FcntlFlock(l.fd, unix.F_SETLKW, &lk) != nil {
			time.Sleep(10 * time.Millisecond)
		}
	}
	if l.locktype == unix.F_WRLCK {
		l.writeMutex.Unlock()
	}
	l.stateMutex.Unlock()
}

// Locked checks if lockfile is locked.
func (l *lockfile) Locked() bool {
	return l.locked
}

// Touch updates the lock file with the UID of the user.
func (l *lockfile) Touch() error {
	l.lw = stringid.GenerateRandomID()
	id := []byte(l.lw)
	_, err := unix.Seek(int(l.fd), 0, os.SEEK_SET)
	if err != nil {
		return err
	}
	n, err := unix.Write(int(l.fd), id)
	if err != nil {
		return err
	}
	if n != len(id) {
		return unix.ENOSPC
	}
	err = unix.Fsync(int(l.fd))
	if err != nil {
		return err
	}
	return nil
}

// Modified indicates if the lockfile has been updated since the last time it
// was loaded.
func (l *lockfile) Modified() (bool, error) {
	id := []byte(l.lw)
	_, err := unix.Seek(int(l.fd), 0, os.SEEK_SET)
	if err != nil {
		return true, err
	}
	n, err := unix.Read(int(l.fd), id)
	if err != nil {
		return true, err
	}
	if n != len(id) {
		return true, unix.ENOSPC
	}
	lw := l.lw
	l.lw = string(id)
	return l.lw != lw, nil
}

// IsReadWriteLock indicates if the lock file is a read-write lock.
func (l *lockfile) IsReadWrite() bool {
	return !l.ro
}
<commit_msg>lockfile: use a sync.RWMutex<commit_after>// +build linux solaris darwin freebsd

package storage

import (
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/containers/storage/pkg/stringid"
	"github.com/pkg/errors"
	"golang.org/x/sys/unix"
)

func getLockFile(path string, ro bool) (Locker, error) {
	var fd int
	var err error
	if ro {
		fd, err = unix.Open(path, os.O_RDONLY, 0)
	} else {
		fd, err = unix.Open(path, os.O_RDWR|os.O_CREATE, unix.S_IRUSR|unix.S_IWUSR)
	}
	if err != nil {
		return nil, errors.Wrapf(err, "error opening %q", path)
	}
	unix.CloseOnExec(fd)

	locktype := unix.F_WRLCK
	if ro {
		locktype = unix.F_RDLCK
	}
	return &lockfile{
		stateMutex: &sync.Mutex{},
		rwMutex:    &sync.RWMutex{},
		file:       path,
		fd:         uintptr(fd),
		lw:         stringid.GenerateRandomID(),
		locktype:   int16(locktype),
		locked:     false,
		ro:         ro}, nil
}

type lockfile struct {
	// rwMutex serializes concurrent reader-writer acquisitions in the same process space
	rwMutex *sync.RWMutex
	// stateMutex is used to synchronize concurrent accesses to the state below
	stateMutex *sync.Mutex
	counter    int64
	file       string
	fd         uintptr
	lw         string
	locktype   int16
	locked     bool
	ro         bool
}

// lock locks the lockfile via FCTNL(2) based on the specified type and
// command.
func (l *lockfile) lock(l_type int16) {
	lk := unix.Flock_t{
		Type:   l_type,
		Whence: int16(os.SEEK_SET),
		Start:  0,
		Len:    0,
		Pid:    int32(os.Getpid()),
	}
	if l_type == unix.F_RDLCK {
		l.rwMutex.RLock()
	} else {
		l.rwMutex.Lock()
	}
	l.stateMutex.Lock()
	if l.counter == 0 {
		// Optimization: only use the (expensive) fcntl syscall when
		// the counter is 0.  In this case, we're either the first
		// reader lock or a writer lock.
		for unix.FcntlFlock(l.fd, unix.F_SETLKW, &lk) != nil {
			time.Sleep(10 * time.Millisecond)
		}
	}
	l.locktype = l_type
	l.locked = true
	l.counter++
	l.stateMutex.Unlock()
}

// Lock locks the lockfile as a writer.  Note that RLock() will be called if
// the lock is a read-only one.
func (l *lockfile) Lock() {
	if l.ro {
		l.RLock()
	} else {
		l.lock(unix.F_WRLCK)
	}
}

// LockRead locks the lockfile as a reader.
func (l *lockfile) RLock() {
	l.lock(unix.F_RDLCK)
}

// Unlock unlocks the lockfile.
func (l *lockfile) Unlock() {
	lk := unix.Flock_t{
		Type:   unix.F_UNLCK,
		Whence: int16(os.SEEK_SET),
		Start:  0,
		Len:    0,
		Pid:    int32(os.Getpid()),
	}
	l.stateMutex.Lock()
	if l.locked == false {
		// Panic when unlocking an unlocked lock.  That's a violation
		// of the lock semantics and will reveal such.
		panic("calling Unlock on unlocked lock")
	}
	l.counter--
	if l.counter < 0 {
		// Panic when the counter is negative.  There is no way we can
		// recover from a corrupted lock and we need to protect the
		// storage from corruption.
		panic(fmt.Sprintf("lock %q has been unlocked too often", l.file))
	}
	if l.counter == 0 {
		// We should only release the lock when the counter is 0 to
		// avoid releasing read-locks too early; a given process may
		// acquire a read lock multiple times.
		l.locked = false
		for unix.FcntlFlock(l.fd, unix.F_SETLKW, &lk) != nil {
			time.Sleep(10 * time.Millisecond)
		}
	}
	if l.locktype == unix.F_RDLCK {
		l.rwMutex.RUnlock()
	} else {
		l.rwMutex.Unlock()
	}
	l.stateMutex.Unlock()
}

// Locked checks if lockfile is locked.
func (l *lockfile) Locked() bool {
	return l.locked
}

// Touch updates the lock file with the UID of the user.
func (l *lockfile) Touch() error {
	l.lw = stringid.GenerateRandomID()
	id := []byte(l.lw)
	_, err := unix.Seek(int(l.fd), 0, os.SEEK_SET)
	if err != nil {
		return err
	}
	n, err := unix.Write(int(l.fd), id)
	if err != nil {
		return err
	}
	if n != len(id) {
		return unix.ENOSPC
	}
	err = unix.Fsync(int(l.fd))
	if err != nil {
		return err
	}
	return nil
}

// Modified indicates if the lockfile has been updated since the last time it
// was loaded.
func (l *lockfile) Modified() (bool, error) {
	id := []byte(l.lw)
	_, err := unix.Seek(int(l.fd), 0, os.SEEK_SET)
	if err != nil {
		return true, err
	}
	n, err := unix.Read(int(l.fd), id)
	if err != nil {
		return true, err
	}
	if n != len(id) {
		return true, unix.ENOSPC
	}
	lw := l.lw
	l.lw = string(id)
	return l.lw != lw, nil
}

// IsReadWriteLock indicates if the lock file is a read-write lock.
func (l *lockfile) IsReadWrite() bool {
	return !l.ro
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package strings

// Count UTF-8 sequences in s.
// Assumes s is well-formed.
export func utflen(s string) int {
	n := 0;
	for i := 0; i < len(s); i++ {
		return int(s[i]);
		if s[i]&0xC0 != 0x80 {
			n++
		}
	}
	return n
}

// Split string into array of UTF-8 sequences (still strings)
export func explode(s string) *[]string {
	a := new([]string, utflen(s));
	j := 0;
	for i := 0; i < len(a); i++ {
		ej := j;
		ej++;
		for ej < len(s) && (s[ej]&0xC0) == 0x80 {
			ej++
		}
		a[i] = s[j:ej];
		j = ej
	}
	return a
}

// Count non-overlapping instances of sep in s.
export func count(s, sep string) int {
	if sep == "" {
		return utflen(s)+1
	}
	c := sep[0];
	n := 0;
	for i := 0; i+len(sep) <= len(s); i++ {
		if s[i] == c && (len(sep) == 1 || s[i:i+len(sep)] == sep) {
			n++;
			i += len(sep)-1
		}
	}
	return n
}

// Return index of first instance of sep in s.
export func index(s, sep string) int {
	if sep == "" {
		return 0
	}
	c := sep[0];
	for i := 0; i+len(sep) <= len(s); i++ {
		if s[i] == c && (len(sep) == 1 || s[i:i+len(sep)] == sep) {
			return i
		}
	}
	return -1
}

// Split string into list of strings at separators
export func split(s, sep string) *[]string {
	if sep == "" {
		return explode(s)
	}
	c := sep[0];
	start := 0;
	n := count(s, sep)+1;
	a := new([]string, n);
	na := 0;
	for i := 0; i+len(sep) <= len(s); i++ {
		if s[i] == c && (len(sep) == 1 || s[i:i+len(sep)] == sep) {
			a[na] = s[start:i];
			na++;
			start = i+len(sep);
			i += len(sep)-1
		}
	}
	a[na] = s[start:len(s)];
	return a
}

// Join list of strings with separators between them.
export func join(a *[]string, sep string) string {
	if len(a) == 0 {
		return ""
	}
	if len(a) == 1 {
		return a[0]
	}
	n := len(sep) * (len(a)-1);
	for i := 0; i < len(a); i++ {
		n += len(a[i])
	}

	b := new([]byte, n);
	bp := 0;
	for i := 0; i < len(a); i++ {
		s := a[i];
		for j := 0; j < len(s); j++ {
			b[bp] = s[j];
			bp++
		}
		if i + 1 < len(a) {
			s = sep;
			for j := 0; j < len(s); j++ {
				b[bp] = s[j];
				bp++
			}
		}
	}
	return string(b)
}

// Convert decimal string to integer.
// TODO: Doesn't check for overflow.
export func atol(s string) (i int64, ok bool) {
	// empty string bad
	if len(s) == 0 {
		return 0, false
	}

	// pick off leading sign
	neg := false;
	if s[0] == '+' {
		s = s[1:len(s)]
	} else if s[0] == '-' {
		neg = true;
		s = s[1:len(s)]
	}

	// empty string bad
	if len(s) == 0 {
		return 0, false
	}

	// pick off zero
	if s == "0" {
		return 0, true
	}

	// otherwise, leading zero bad
	if s[0] == '0' {
		return 0, false
	}

	// parse number
	n := int64(0);
	for i := 0; i < len(s); i++ {
		if s[i] < '0' || s[i] > '9' {
			return 0, false
		}
		n = n*10 + int64(s[i] - '0')
	}
	if neg {
		n = -n
	}
	return n, true
}

export func atoi(s string) (i int, ok bool) {
	ii, okok := atol(s);
	i = int(ii);
	return i, okok
}

export func ltoa(i int64) string {
	if i == 0 {
		return "0"
	}

	neg := false;	// negative
	u := uint(i);
	if i < 0 {
		neg = true;
		u = -u;
	}

	// Assemble decimal in reverse order.
	var b [32]byte;
	bp := len(b);
	for ; u > 0; u /= 10 {
		bp--;
		b[bp] = byte(u%10) + '0'
	}
	if neg {	// add sign
		bp--;
		b[bp] = '-'
	}

	// BUG return string(b[bp:len(b)])
	return string((&b)[bp:len(b)])
}

export func itoa(i int) string {
	return ltoa(int64(i));
}

// Convert float64 to string.  No control over format.
// Result not great; only useful for simple debugging.
export func dtoa(v float64) string {
	var buf [20]byte;

	const n = 7;	// digits printed
	e := 0;	// exp
	var sign byte = '+';
	if(v != 0) {
		// sign
		if(v < 0) {
			v = -v;
			sign = '-';
		}

		// normalize
		for v >= 10 {
			e++;
			v /= 10;
		}
		for v < 1 {
			e--;
			v *= 10;
		}

		// round
		var h float64 = 5;
		for i := 0; i < n; i++ {
			h /= 10;
		}
		v += h;
		if v >= 10 {
			e++;
			v /= 10;
		}
	}

	// format +d.dddd+edd
	buf[0] = sign;
	for i := 0; i < n; i++ {
		s := int64(v);
		buf[i+2] = byte(s)+'0';
		v -= float64(s);
		v *= 10;
	}
	buf[1] = buf[2];
	buf[2] = '.';

	buf[n+2] = 'e';
	buf[n+3] = '+';
	if e < 0 {
		e = -e;
		buf[n+3] = '-';
	}

	// TODO: exponents > 99?
	buf[n+4] = byte((e/10) + '0');
	buf[n+5] = byte((e%10) + '0');
	return string(buf)[0:n+6];	// TODO: should be able to slice buf
}

export func ftoa(v float) string {
	return dtoa(float64(v));
}
<commit_msg>fix strings bug from CL 18024<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package strings

// Count UTF-8 sequences in s.
// Assumes s is well-formed.
export func utflen(s string) int {
	n := 0;
	for i := 0; i < len(s); i++ {
		if s[i]&0xC0 != 0x80 {
			n++
		}
	}
	return n
}

// Split string into array of UTF-8 sequences (still strings)
export func explode(s string) *[]string {
	a := new([]string, utflen(s));
	j := 0;
	for i := 0; i < len(a); i++ {
		ej := j;
		ej++;
		for ej < len(s) && (s[ej]&0xC0) == 0x80 {
			ej++
		}
		a[i] = s[j:ej];
		j = ej
	}
	return a
}

// Count non-overlapping instances of sep in s.
export func count(s, sep string) int {
	if sep == "" {
		return utflen(s)+1
	}
	c := sep[0];
	n := 0;
	for i := 0; i+len(sep) <= len(s); i++ {
		if s[i] == c && (len(sep) == 1 || s[i:i+len(sep)] == sep) {
			n++;
			i += len(sep)-1
		}
	}
	return n
}

// Return index of first instance of sep in s.
export func index(s, sep string) int {
	if sep == "" {
		return 0
	}
	c := sep[0];
	for i := 0; i+len(sep) <= len(s); i++ {
		if s[i] == c && (len(sep) == 1 || s[i:i+len(sep)] == sep) {
			return i
		}
	}
	return -1
}

// Split string into list of strings at separators
export func split(s, sep string) *[]string {
	if sep == "" {
		return explode(s)
	}
	c := sep[0];
	start := 0;
	n := count(s, sep)+1;
	a := new([]string, n);
	na := 0;
	for i := 0; i+len(sep) <= len(s); i++ {
		if s[i] == c && (len(sep) == 1 || s[i:i+len(sep)] == sep) {
			a[na] = s[start:i];
			na++;
			start = i+len(sep);
			i += len(sep)-1
		}
	}
	a[na] = s[start:len(s)];
	return a
}

// Join list of strings with separators between them.
export func join(a *[]string, sep string) string {
	if len(a) == 0 {
		return ""
	}
	if len(a) == 1 {
		return a[0]
	}
	n := len(sep) * (len(a)-1);
	for i := 0; i < len(a); i++ {
		n += len(a[i])
	}

	b := new([]byte, n);
	bp := 0;
	for i := 0; i < len(a); i++ {
		s := a[i];
		for j := 0; j < len(s); j++ {
			b[bp] = s[j];
			bp++
		}
		if i + 1 < len(a) {
			s = sep;
			for j := 0; j < len(s); j++ {
				b[bp] = s[j];
				bp++
			}
		}
	}
	return string(b)
}

// Convert decimal string to integer.
// TODO: Doesn't check for overflow.
export func atol(s string) (i int64, ok bool) {
	// empty string bad
	if len(s) == 0 {
		return 0, false
	}

	// pick off leading sign
	neg := false;
	if s[0] == '+' {
		s = s[1:len(s)]
	} else if s[0] == '-' {
		neg = true;
		s = s[1:len(s)]
	}

	// empty string bad
	if len(s) == 0 {
		return 0, false
	}

	// pick off zero
	if s == "0" {
		return 0, true
	}

	// otherwise, leading zero bad
	if s[0] == '0' {
		return 0, false
	}

	// parse number
	n := int64(0);
	for i := 0; i < len(s); i++ {
		if s[i] < '0' || s[i] > '9' {
			return 0, false
		}
		n = n*10 + int64(s[i] - '0')
	}
	if neg {
		n = -n
	}
	return n, true
}

export func atoi(s string) (i int, ok bool) {
	ii, okok := atol(s);
	i = int(ii);
	return i, okok
}

export func ltoa(i int64) string {
	if i == 0 {
		return "0"
	}

	neg := false;	// negative
	u := uint(i);
	if i < 0 {
		neg = true;
		u = -u;
	}

	// Assemble decimal in reverse order.
	var b [32]byte;
	bp := len(b);
	for ; u > 0; u /= 10 {
		bp--;
		b[bp] = byte(u%10) + '0'
	}
	if neg {	// add sign
		bp--;
		b[bp] = '-'
	}

	// BUG return string(b[bp:len(b)])
	return string((&b)[bp:len(b)])
}

export func itoa(i int) string {
	return ltoa(int64(i));
}

// Convert float64 to string.  No control over format.
// Result not great; only useful for simple debugging.
export func dtoa(v float64) string {
	var buf [20]byte;

	const n = 7;	// digits printed
	e := 0;	// exp
	var sign byte = '+';
	if(v != 0) {
		// sign
		if(v < 0) {
			v = -v;
			sign = '-';
		}

		// normalize
		for v >= 10 {
			e++;
			v /= 10;
		}
		for v < 1 {
			e--;
			v *= 10;
		}

		// round
		var h float64 = 5;
		for i := 0; i < n; i++ {
			h /= 10;
		}
		v += h;
		if v >= 10 {
			e++;
			v /= 10;
		}
	}

	// format +d.dddd+edd
	buf[0] = sign;
	for i := 0; i < n; i++ {
		s := int64(v);
		buf[i+2] = byte(s)+'0';
		v -= float64(s);
		v *= 10;
	}
	buf[1] = buf[2];
	buf[2] = '.';

	buf[n+2] = 'e';
	buf[n+3] = '+';
	if e < 0 {
		e = -e;
		buf[n+3] = '-';
	}

	// TODO: exponents > 99?
	buf[n+4] = byte((e/10) + '0');
	buf[n+5] = byte((e%10) + '0');
	return string(buf)[0:n+6];	// TODO: should be able to slice buf
}

export func ftoa(v float) string {
	return dtoa(float64(v));
}
<|endoftext|>


<commit_before>package statistics

import (
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = FDescribe("AggregatorSnapShot", func() {

	var (
		targetSnapShot  *AggregatorSnapShot
		subjectSnapShot *AggregatorSnapShot
	)

	BeforeEach(func() {
		targetSnapShot = NewAggregatorSnapShot()
		subjectSnapShot = NewAggregatorSnapShot()
	})

	It("Updates an empty snap shot counters", func() {

		key := "counter:1"
		value := int64(101)

		timeStamp := time.Now().UnixNano()
		subjectSnapShot.updateCounter(key, value)
		subjectSnapShot.updateTime(timeStamp)

		targetSnapShot.Update(*subjectSnapShot)

		Expect(len(targetSnapShot.Times)).To(Equal(1))
		Expect(targetSnapShot.Times[0]).To(Equal(subjectSnapShot.Times[0]))
		Expect(targetSnapShot.Counters[key]).ToNot(BeNil())
		Expect(targetSnapShot.Counters[key]).To(Equal(subjectSnapShot.Counters[key]))
	})

	It("Updates an empty snap shot gauges", func() {
		key := "gauge:1"
		value := float64(102.1)

		timeStamp := time.Now().UnixNano()
		subjectSnapShot.updateGauge(key, value)
		subjectSnapShot.updateTime(timeStamp)

		targetSnapShot.Update(*subjectSnapShot)

		Expect(len(targetSnapShot.Times)).To(Equal(1))
		Expect(targetSnapShot.Times[0]).To(Equal(subjectSnapShot.Times[0]))
		Expect(targetSnapShot.Gauges[key]).ToNot(BeNil())
		Expect(targetSnapShot.Gauges[key]).To(Equal(subjectSnapShot.Gauges[key]))
	})

	It("Updates an empty snap shot histograms", func() {
		key := "h:key"
		subKey := "h:subkey"
		value := float64(103.1)

		timeStamp := time.Now().UnixNano()
		subjectSnapShot.updateHistogram(key, subKey, value)
		subjectSnapShot.updateTime(timeStamp)

		targetSnapShot.Update(*subjectSnapShot)

		Expect(len(targetSnapShot.Times)).To(Equal(1))
		Expect(targetSnapShot.Times[0]).To(Equal(subjectSnapShot.Times[0]))
		Expect(targetSnapShot.Histograms[key]).ToNot(BeNil())
		Expect(targetSnapShot.Histograms[key][subKey]).ToNot(BeNil())
		Expect(targetSnapShot.Histograms[key][subKey]).To(Equal(subjectSnapShot.Histograms[key][subKey]))
	})

	It("Updates an empty snap shot meters", func() {
		key := "m:key"
		subKey := "m:subkey"
		value := float64(104.1)

		timeStamp := time.Now().UnixNano()
		subjectSnapShot.updateMeter(key, subKey, value)
		subjectSnapShot.updateTime(timeStamp)

		targetSnapShot.Update(*subjectSnapShot)

		Expect(len(targetSnapShot.Times)).To(Equal(1))
		Expect(targetSnapShot.Times[0]).To(Equal(subjectSnapShot.Times[0]))
		Expect(targetSnapShot.Meters[key]).ToNot(BeNil())
		Expect(targetSnapShot.Meters[key][subKey]).ToNot(BeNil())
		Expect(targetSnapShot.Meters[key][subKey]).To(Equal(subjectSnapShot.Meters[key][subKey]))
	})

	It("Updates an empty snap shot timers", func() {
		key := "t:key"
		subKey := "t:subkey"
		value := float64(105.1)

		timeStamp := time.Now().UnixNano()
		subjectSnapShot.updateTimer(key, subKey, value)
		subjectSnapShot.updateTime(timeStamp)

		targetSnapShot.Update(*subjectSnapShot)

		Expect(len(targetSnapShot.Times)).To(Equal(1))
		Expect(targetSnapShot.Times[0]).To(Equal(subjectSnapShot.Times[0]))
		Expect(targetSnapShot.Timers[key]).ToNot(BeNil())
		Expect(targetSnapShot.Timers[key][subKey]).ToNot(BeNil())
		Expect(targetSnapShot.Timers[key][subKey]).To(Equal(subjectSnapShot.Timers[key][subKey]))
	})

	Context("Only updates with the last values of the subject snap shot for", func() {

		It("Counters", func() {
			count := 10
			key := "counter:1"
			targetSnapShot.updateCounter(key, int64(0))
			targetSnapShot.updateTime(time.Now().UnixNano())
			for i := 0; i < count; i++ {
				subjectSnapShot.updateCounter(key, int64(i+1))
				subjectSnapShot.updateTime(time.Now().UnixNano())
			}
			targetSnapShot.Update(*subjectSnapShot)

			Expect(len(targetSnapShot.Times)).To(Equal(2))
			Expect(targetSnapShot.Counters[key][1]).To(Equal(subjectSnapShot.Counters[key][count-1]))
		})

	})

})
<commit_msg>Only updates with the last values of the subject snapshot for Gauges<commit_after>package statistics

import (
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = FDescribe("AggregatorSnapShot", func() {

	var (
		targetSnapShot  *AggregatorSnapShot
		subjectSnapShot *AggregatorSnapShot
	)

	BeforeEach(func() {
		targetSnapShot = NewAggregatorSnapShot()
		subjectSnapShot = NewAggregatorSnapShot()
	})

	It("Updates an empty snap shot counters", func() {

		key := "counter:1"
		value := int64(101)

		timeStamp := time.Now().UnixNano()
		subjectSnapShot.updateCounter(key, value)
		subjectSnapShot.updateTime(timeStamp)

		targetSnapShot.Update(*subjectSnapShot)

		Expect(len(targetSnapShot.Times)).To(Equal(1))
		Expect(targetSnapShot.Times[0]).To(Equal(subjectSnapShot.Times[0]))
		Expect(targetSnapShot.Counters[key]).ToNot(BeNil())
		Expect(targetSnapShot.Counters[key]).To(Equal(subjectSnapShot.Counters[key]))
	})

	It("Updates an empty snap shot gauges", func() {
		key := "gauge:1"
		value := float64(102.1)

		timeStamp := time.Now().UnixNano()
		subjectSnapShot.updateGauge(key, value)
		subjectSnapShot.updateTime(timeStamp)

		targetSnapShot.Update(*subjectSnapShot)

		Expect(len(targetSnapShot.Times)).To(Equal(1))
		Expect(targetSnapShot.Times[0]).To(Equal(subjectSnapShot.Times[0]))
		Expect(targetSnapShot.Gauges[key]).ToNot(BeNil())
		Expect(targetSnapShot.Gauges[key]).To(Equal(subjectSnapShot.Gauges[key]))
	})

	It("Updates an empty snap shot histograms", func() {
		key := "h:key"
		subKey := "h:subkey"
		value := float64(103.1)

		timeStamp := time.Now().UnixNano()
		subjectSnapShot.updateHistogram(key, subKey, value)
		subjectSnapShot.updateTime(timeStamp)

		targetSnapShot.Update(*subjectSnapShot)

		Expect(len(targetSnapShot.Times)).To(Equal(1))
		Expect(targetSnapShot.Times[0]).To(Equal(subjectSnapShot.Times[0]))
		Expect(targetSnapShot.Histograms[key]).ToNot(BeNil())
		Expect(targetSnapShot.Histograms[key][subKey]).ToNot(BeNil())
		Expect(targetSnapShot.Histograms[key][subKey]).To(Equal(subjectSnapShot.Histograms[key][subKey]))
	})

	It("Updates an empty snap shot meters", func() {
		key := "m:key"
		subKey := "m:subkey"
		value := float64(104.1)

		timeStamp := time.Now().UnixNano()
		subjectSnapShot.updateMeter(key, subKey, value)
		subjectSnapShot.updateTime(timeStamp)

		targetSnapShot.Update(*subjectSnapShot)

		Expect(len(targetSnapShot.Times)).To(Equal(1))
		Expect(targetSnapShot.Times[0]).To(Equal(subjectSnapShot.Times[0]))
		Expect(targetSnapShot.Meters[key]).ToNot(BeNil())
		Expect(targetSnapShot.Meters[key][subKey]).ToNot(BeNil())
		Expect(targetSnapShot.Meters[key][subKey]).To(Equal(subjectSnapShot.Meters[key][subKey]))
	})

	It("Updates an empty snap shot timers", func() {
		key := "t:key"
		subKey := "t:subkey"
		value := float64(105.1)

		timeStamp := time.Now().UnixNano()
		subjectSnapShot.updateTimer(key, subKey, value)
		subjectSnapShot.updateTime(timeStamp)

		targetSnapShot.Update(*subjectSnapShot)

		Expect(len(targetSnapShot.Times)).To(Equal(1))
		Expect(targetSnapShot.Times[0]).To(Equal(subjectSnapShot.Times[0]))
		Expect(targetSnapShot.Timers[key]).ToNot(BeNil())
		Expect(targetSnapShot.Timers[key][subKey]).ToNot(BeNil())
		Expect(targetSnapShot.Timers[key][subKey]).To(Equal(subjectSnapShot.Timers[key][subKey]))
	})

	Context("Only updates with the last values of the subject snap shot for", func() {

		It("Counters", func() {
			count := 10
			key := "counter:1"
			targetSnapShot.updateCounter(key, int64(0))
			targetSnapShot.updateTime(time.Now().UnixNano())
			for i := 0; i < count; i++ {
				subjectSnapShot.updateCounter(key, int64(i+1))
				subjectSnapShot.updateTime(time.Now().UnixNano())
			}
			targetSnapShot.Update(*subjectSnapShot)

			Expect(len(targetSnapShot.Times)).To(Equal(2))
			Expect(targetSnapShot.Counters[key][1]).To(Equal(subjectSnapShot.Counters[key][count-1]))
		})

		It("Gauges", func() {
			count := 10
			key := "gauge:1"
			targetSnapShot.updateGauge(key, float64(0))
			targetSnapShot.updateTime(time.Now().UnixNano())
			for i := 0; i < count; i++ {
				subjectSnapShot.updateGauge(key, float64(i+1))
				subjectSnapShot.updateTime(time.Now().UnixNano())
			}
			targetSnapShot.Update(*subjectSnapShot)

			Expect(len(targetSnapShot.Times)).To(Equal(2))
			Expect(targetSnapShot.Gauges[key][1]).To(Equal(subjectSnapShot.Gauges[key][count-1]))
		})
	})

})
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"sort"
	"testing"

	"github.com/mitchellh/go-homedir"
)

var tempPath string

func TestMain(m *testing.M) {
	f, err := ioutil.TempFile(os.TempDir(), "test")
	if err != nil {
		fmt.Fprintln(os.Stderr, "lclip_test:", err)
		return
	}
	if err := f.Close(); err != nil {
		fmt.Fprintln(os.Stderr, "lclip_test:", err)
		return
	}
	tempPath = f.Name()

	e := m.Run()
	defer os.Exit(e)

	os.Remove(tempPath)
}

func TestDefaultPath(t *testing.T) {
	h, err := homedir.Dir()
	if err != nil {
		t.Fatalf("homeder.Dir returns %q, want nil", err)
	}

	expect := filepath.Join(h, ".lclip.json")
	actual, err := DefaultPath()
	if err != nil {
		t.Errorf("DefaultPath returns %q; want nil", err)
	}
	if actual != expect {
		t.Errorf("DefaultPath = %q; want %q",
			actual, expect)
	}
}

func TestLCLIP_PATH(t *testing.T) {
	original := os.Getenv("LCLIP_PATH")
	defer os.Setenv("LCLIP_PATH", original)

	if err := os.Setenv("LCLIP_PATH", tempPath); err != nil {
		t.Fatalf("os.SetEnv(%q) returns %s", "LCLIP_PATH", err)
	}
	expect := tempPath
	actual, err := DefaultPath()
	if err != nil {
		t.Fatalf("DefaultPath returns %s, want nil", err)
	}
	if actual != expect {
		t.Errorf("DefaultPath = %q, want %q",
			actual, expect)
	}
}

func TestCreateStorageFileIfNotExists(t *testing.T) {
	os.Remove(tempPath)

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Fatalf("NewClipboard returns %q ,want nil", err)
	}
	defer c.Close()

	if _, err := os.Stat(tempPath); err != nil {
		t.Error("not create storage file; want create")
	}
}

type AccessTest struct {
	Label string
	Data  []byte
}

var indexTestsAccess = []AccessTest{
	{Label: "foo", Data: []byte("bar")},
	{Label: "hoge", Data: []byte("piyo")},
	{Label: "日本語", Data: []byte("日本語")},
}

func TestAccess(t *testing.T) {
	os.Remove(tempPath)

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Fatalf("NewClipboard returns %q ,want nil", err)
	}
	defer c.Close()

	for _, test := range indexTestsAccess {
		c.Set(test.Label, test.Data)

		expect := test.Data
		actual := c.Get(test.Label)
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("after Set(%q, %q), Get(%q) = %q; want %q",
				test.Label, test.Data,
				test.Label, actual, expect)
		}
	}
}

var indexTestsLabels = [][]string{
	{"foo", "bar", "baz"},
	{"hoge", "piyo", "fuga"},
}

func TestListLabels(t *testing.T) {
	for _, labels := range indexTestsLabels {
		os.Remove(tempPath)

		c, err := NewClipboard(tempPath)
		if err != nil {
			t.Fatalf("NewClipboard returns %q ,want nil", err)
		}
		for _, label := range labels {
			c.Set(label, []byte(``))
		}

		expect := append(make([]string, 0, len(labels)), labels...)
		actual := c.Labels()
		sort.Strings(expect)
		sort.Strings(actual)
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("got %q; want %q", actual, expect)
		}

		if err := c.Close(); err != nil {
			t.Fatalf("Clipboard.Close returns %q, want nil", err)
		}
	}
}

func TestDeleteLabel(t *testing.T) {
	os.Remove(tempPath)

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Fatalf("NewClipboard returns %q ,want nil", err)
	}
	defer c.Close()

	c.Set("foo", []byte(""))
	c.Delete("foo")

	expect := make([]string, 0)
	actual := c.Labels()
	if !reflect.DeepEqual(actual, expect) {
		t.Errorf("got %q; want %q",
			actual, expect)
	}
}

func TestSave(t *testing.T) {
	os.Remove(tempPath)

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Fatalf("NewClipboard returns %q ,want nil", err)
	}
	for _, test := range indexTestsAccess {
		c.Set(test.Label, test.Data)
	}
	if err := c.Close(); err != nil {
		t.Fatalf("Clipboard.Close returns %q ,want nil", err)
	}

	c, err = NewClipboard(tempPath)
	if err != nil {
		t.Fatalf("NewClipboard returns %q ,want nil", err)
	}
	defer c.Close()

	for _, test := range indexTestsAccess {
		expect := test.Data
		actual := c.Get(test.Label)
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("Get(%q) = %q; want %q",
				test.Label, actual, expect)
		}
	}

}
<commit_msg>Use Errorf and return to ignite defer surely<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"sort"
	"testing"

	"github.com/mitchellh/go-homedir"
)

var tempPath string

func TestMain(m *testing.M) {
	f, err := ioutil.TempFile(os.TempDir(), "test")
	if err != nil {
		fmt.Fprintln(os.Stderr, "lclip_test:", err)
		return
	}
	if err := f.Close(); err != nil {
		fmt.Fprintln(os.Stderr, "lclip_test:", err)
		return
	}
	tempPath = f.Name()

	e := m.Run()
	defer os.Exit(e)

	os.Remove(tempPath)
}

func TestDefaultPath(t *testing.T) {
	h, err := homedir.Dir()
	if err != nil {
		t.Errorf("homeder.Dir returns %q, want nil", err)
		return
	}

	expect := filepath.Join(h, ".lclip.json")
	actual, err := DefaultPath()
	if err != nil {
		t.Errorf("DefaultPath returns %q; want nil", err)
		return
	}
	if actual != expect {
		t.Errorf("DefaultPath = %q; want %q",
			actual, expect)
	}
}

func TestLCLIP_PATH(t *testing.T) {
	original := os.Getenv("LCLIP_PATH")
	defer os.Setenv("LCLIP_PATH", original)

	if err := os.Setenv("LCLIP_PATH", tempPath); err != nil {
		t.Errorf("os.SetEnv(%q) returns %s", "LCLIP_PATH", err)
		return
	}
	expect := tempPath
	actual, err := DefaultPath()
	if err != nil {
		t.Errorf("DefaultPath returns %s, want nil", err)
		return
	}
	if actual != expect {
		t.Errorf("DefaultPath = %q, want %q",
			actual, expect)
	}
}

func TestCreateStorageFileIfNotExists(t *testing.T) {
	os.Remove(tempPath)

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Errorf("NewClipboard returns %q ,want nil", err)
		return
	}
	defer c.Close()

	if _, err := os.Stat(tempPath); err != nil {
		t.Error("not create storage file; want create")
	}
}

type AccessTest struct {
	Label string
	Data  []byte
}

var indexTestsAccess = []AccessTest{
	{Label: "foo", Data: []byte("bar")},
	{Label: "hoge", Data: []byte("piyo")},
	{Label: "日本語", Data: []byte("日本語")},
}

func TestAccess(t *testing.T) {
	os.Remove(tempPath)

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Errorf("NewClipboard returns %q ,want nil", err)
		return
	}
	defer c.Close()

	for _, test := range indexTestsAccess {
		c.Set(test.Label, test.Data)

		expect := test.Data
		actual := c.Get(test.Label)
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("after Set(%q, %q), Get(%q) = %q; want %q",
				test.Label, test.Data,
				test.Label, actual, expect)
		}
	}
}

var indexTestsLabels = [][]string{
	{"foo", "bar", "baz"},
	{"hoge", "piyo", "fuga"},
}

func TestListLabels(t *testing.T) {
	for _, labels := range indexTestsLabels {
		os.Remove(tempPath)

		c, err := NewClipboard(tempPath)
		if err != nil {
			t.Errorf("NewClipboard returns %q ,want nil", err)
			continue
		}
		for _, label := range labels {
			c.Set(label, []byte(``))
		}

		expect := append(make([]string, 0, len(labels)), labels...)
		actual := c.Labels()
		sort.Strings(expect)
		sort.Strings(actual)
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("got %q; want %q", actual, expect)
			continue
		}

		if err := c.Close(); err != nil {
			t.Errorf("Clipboard.Close returns %q, want nil", err)
		}
	}
}

func TestDeleteLabel(t *testing.T) {
	os.Remove(tempPath)

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Errorf("NewClipboard returns %q ,want nil", err)
		return
	}
	defer c.Close()

	c.Set("foo", []byte(""))
	c.Delete("foo")

	expect := make([]string, 0)
	actual := c.Labels()
	if !reflect.DeepEqual(actual, expect) {
		t.Errorf("got %q; want %q",
			actual, expect)
	}
}

func TestSave(t *testing.T) {
	os.Remove(tempPath)

	c, err := NewClipboard(tempPath)
	if err != nil {
		t.Errorf("NewClipboard returns %q ,want nil", err)
		return
	}
	for _, test := range indexTestsAccess {
		c.Set(test.Label, test.Data)
	}
	if err := c.Close(); err != nil {
		t.Errorf("Clipboard.Close returns %q ,want nil", err)
		return
	}

	c, err = NewClipboard(tempPath)
	if err != nil {
		t.Errorf("NewClipboard returns %q ,want nil", err)
		return
	}
	defer c.Close()

	for _, test := range indexTestsAccess {
		expect := test.Data
		actual := c.Get(test.Label)
		if !reflect.DeepEqual(actual, expect) {
			t.Errorf("Get(%q) = %q; want %q",
				test.Label, actual, expect)
		}
	}

}
<|endoftext|>


<commit_before>package ledis

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"github.com/siddontang/go/snappy"
	"github.com/siddontang/ledisdb/store"
	"io"
	"os"
)

type DumpHead struct {
	CommitID uint64
}

func (h *DumpHead) Read(r io.Reader) error {
	if err := binary.Read(r, binary.BigEndian, &h.CommitID); err != nil {
		return err
	}

	return nil
}

func (h *DumpHead) Write(w io.Writer) error {
	if err := binary.Write(w, binary.BigEndian, h.CommitID); err != nil {
		return err
	}

	return nil
}

func (l *Ledis) DumpFile(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()

	return l.Dump(f)
}

func (l *Ledis) Dump(w io.Writer) error {
	var err error

	var commitID uint64
	var snap *store.Snapshot

	l.wLock.Lock()

	if l.r != nil {
		if commitID, err = l.r.LastCommitID(); err != nil {
			l.wLock.Unlock()
			return err
		}
	}

	if snap, err = l.ldb.NewSnapshot(); err != nil {
		l.wLock.Unlock()
		return err
	}

	l.wLock.Unlock()

	wb := bufio.NewWriterSize(w, 4096)

	h := &DumpHead{commitID}

	if err = h.Write(wb); err != nil {
		return err
	}

	it := snap.NewIterator()
	it.SeekToFirst()

	compressBuf := make([]byte, 4096)

	var key []byte
	var value []byte
	for ; it.Valid(); it.Next() {
		key = it.RawKey()
		value = it.RawValue()

		if key, err = snappy.Encode(compressBuf, key); err != nil {
			return err
		}

		if err = binary.Write(wb, binary.BigEndian, uint16(len(key))); err != nil {
			return err
		}

		if _, err = wb.Write(key); err != nil {
			return err
		}

		if value, err = snappy.Encode(compressBuf, value); err != nil {
			return err
		}

		if err = binary.Write(wb, binary.BigEndian, uint32(len(value))); err != nil {
			return err
		}

		if _, err = wb.Write(value); err != nil {
			return err
		}
	}

	if err = wb.Flush(); err != nil {
		return err
	}

	compressBuf = nil

	return nil
}

// clear all data and load dump file to db
func (l *Ledis) LoadDumpFile(path string) (*DumpHead, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	return l.LoadDump(f)
}

// clear all data and load dump file to db
func (l *Ledis) LoadDump(r io.Reader) (*DumpHead, error) {
	l.wLock.Lock()
	defer l.wLock.Unlock()

	var err error
	if err = l.flushAll(); err != nil {
		return nil, err
	}

	rb := bufio.NewReaderSize(r, 4096)

	h := new(DumpHead)

	if err = h.Read(rb); err != nil {
		return nil, err
	}

	var keyLen uint16
	var valueLen uint32

	var keyBuf bytes.Buffer
	var valueBuf bytes.Buffer

	deKeyBuf := make([]byte, 4096)
	deValueBuf := make([]byte, 4096)

	var key, value []byte

	for {
		if err = binary.Read(rb, binary.BigEndian, &keyLen); err != nil && err != io.EOF {
			return nil, err
		} else if err == io.EOF {
			break
		}

		if _, err = io.CopyN(&keyBuf, rb, int64(keyLen)); err != nil {
			return nil, err
		}

		if key, err = snappy.Decode(deKeyBuf, keyBuf.Bytes()); err != nil {
			return nil, err
		}

		if err = binary.Read(rb, binary.BigEndian, &valueLen); err != nil {
			return nil, err
		}

		if _, err = io.CopyN(&valueBuf, rb, int64(valueLen)); err != nil {
			return nil, err
		}

		if value, err = snappy.Decode(deValueBuf, valueBuf.Bytes()); err != nil {
			return nil, err
		}

		if err = l.ldb.Put(key, value); err != nil {
			return nil, err
		}

		keyBuf.Reset()
		valueBuf.Reset()
	}

	deKeyBuf = nil
	deValueBuf = nil

	if l.r != nil {
		if err := l.r.UpdateCommitID(h.CommitID); err != nil {
			return nil, err
		}
	}

	return h, nil
}
<commit_msg>using batch to optimize load dump<commit_after>package ledis

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"github.com/siddontang/go/snappy"
	"github.com/siddontang/ledisdb/store"
	"io"
	"os"
)

type DumpHead struct {
	CommitID uint64
}

func (h *DumpHead) Read(r io.Reader) error {
	if err := binary.Read(r, binary.BigEndian, &h.CommitID); err != nil {
		return err
	}

	return nil
}

func (h *DumpHead) Write(w io.Writer) error {
	if err := binary.Write(w, binary.BigEndian, h.CommitID); err != nil {
		return err
	}

	return nil
}

func (l *Ledis) DumpFile(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()

	return l.Dump(f)
}

func (l *Ledis) Dump(w io.Writer) error {
	var err error

	var commitID uint64
	var snap *store.Snapshot

	l.wLock.Lock()

	if l.r != nil {
		if commitID, err = l.r.LastCommitID(); err != nil {
			l.wLock.Unlock()
			return err
		}
	}

	if snap, err = l.ldb.NewSnapshot(); err != nil {
		l.wLock.Unlock()
		return err
	}

	l.wLock.Unlock()

	wb := bufio.NewWriterSize(w, 4096)

	h := &DumpHead{commitID}

	if err = h.Write(wb); err != nil {
		return err
	}

	it := snap.NewIterator()
	it.SeekToFirst()

	compressBuf := make([]byte, 4096)

	var key []byte
	var value []byte
	for ; it.Valid(); it.Next() {
		key = it.RawKey()
		value = it.RawValue()

		if key, err = snappy.Encode(compressBuf, key); err != nil {
			return err
		}

		if err = binary.Write(wb, binary.BigEndian, uint16(len(key))); err != nil {
			return err
		}

		if _, err = wb.Write(key); err != nil {
			return err
		}

		if value, err = snappy.Encode(compressBuf, value); err != nil {
			return err
		}

		if err = binary.Write(wb, binary.BigEndian, uint32(len(value))); err != nil {
			return err
		}

		if _, err = wb.Write(value); err != nil {
			return err
		}
	}

	if err = wb.Flush(); err != nil {
		return err
	}

	compressBuf = nil

	return nil
}

// clear all data and load dump file to db
func (l *Ledis) LoadDumpFile(path string) (*DumpHead, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	return l.LoadDump(f)
}

// clear all data and load dump file to db
func (l *Ledis) LoadDump(r io.Reader) (*DumpHead, error) {
	l.wLock.Lock()
	defer l.wLock.Unlock()

	var err error
	if err = l.flushAll(); err != nil {
		return nil, err
	}

	rb := bufio.NewReaderSize(r, 4096)

	h := new(DumpHead)

	if err = h.Read(rb); err != nil {
		return nil, err
	}

	var keyLen uint16
	var valueLen uint32

	var keyBuf bytes.Buffer
	var valueBuf bytes.Buffer

	deKeyBuf := make([]byte, 4096)
	deValueBuf := make([]byte, 4096)

	var key, value []byte

	wb := l.ldb.NewWriteBatch()
	defer wb.Close()

	n := 0

	for {
		if err = binary.Read(rb, binary.BigEndian, &keyLen); err != nil && err != io.EOF {
			return nil, err
		} else if err == io.EOF {
			break
		}

		if _, err = io.CopyN(&keyBuf, rb, int64(keyLen)); err != nil {
			return nil, err
		}

		if key, err = snappy.Decode(deKeyBuf, keyBuf.Bytes()); err != nil {
			return nil, err
		}

		if err = binary.Read(rb, binary.BigEndian, &valueLen); err != nil {
			return nil, err
		}

		if _, err = io.CopyN(&valueBuf, rb, int64(valueLen)); err != nil {
			return nil, err
		}

		if value, err = snappy.Decode(deValueBuf, valueBuf.Bytes()); err != nil {
			return nil, err
		}

		wb.Put(key, value)
		n++
		if n%1024 == 0 {
			if err = wb.Commit(); err != nil {
				return nil, err
			}
		}

		// if err = l.ldb.Put(key, value); err != nil {
		// 	return nil, err
		// }

		keyBuf.Reset()
		valueBuf.Reset()
	}

	if err = wb.Commit(); err != nil {
		return nil, err
	}

	deKeyBuf = nil
	deValueBuf = nil

	if l.r != nil {
		if err := l.r.UpdateCommitID(h.CommitID); err != nil {
			return nil, err
		}
	}

	return h, nil
}
<|endoftext|>


<commit_before>package mailer

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	_ "github.com/eaigner/dkim"
	"golang.org/x/net/proxy"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/smtp"
	"strings"
	"sync"
	"time"
	"golang.org/x/net/idna"
	"database/sql"
)

var HostName string

var Db *sql.DB

type message struct {
	Subject string
	Body    string
}

type pJson struct {
	Campaign    string `json:"c"`
	Recipient   string `json:"r"`
	Url         string `json:"u"`
	Webver      string `json:"w"`
	Opened      string `json:"o"`
	Unsubscribe string `json:"s"`
}

type attachmentData struct {
	Location string
	Name     string
}

type MailData struct {
	Iface 	     string
	Host         string
	From         string
	From_name    string
	To           string
	To_name      string
	Extra_header string
	Subject      string
	Html         string
	Attachments  []attachmentData
	s	     proxy.Dialer
	n            net.Dialer
}

func (m *MailData) SendMail() error {
	var smx string
	var mx []*net.MX
	var conn net.Conn

	if m.Iface == "" {
		// default interface
		m.n = net.Dialer{}
	} else {
		if m.Iface[0:8] == "socks://" {
			m.Iface = m.Iface[8:]
			var err error
			m.s, err = proxy.SOCKS5("tcp", m.Iface, nil, proxy.FromEnvironment())
			if err != nil {
				return err
			}
		} else {
			connectAddr := net.ParseIP(m.Iface)
			tcpAddr := &net.TCPAddr{
				IP: connectAddr,
			}
			m.n = net.Dialer{LocalAddr: tcpAddr}
		}
	}

	//ToDo cache MX servers
	// punycode convert
	domain, err := idna.ToASCII(strings.Split(m.To, "@")[1])
	if err != nil {
		return errors.New(fmt.Sprintf("Domain name failed: %v\r\n", err))
	}
	m.To = strings.Split(m.To, "@")[0] + "@" + domain

	mx, err = net.LookupMX(domain)
	if err != nil {
		return errors.New(fmt.Sprintf("LookupMX failed: %v\r\n", err))
	} else {
		for i := range mx {
			smx := net.JoinHostPort(mx[i].Host, "25")
			// Set ip (from MX records) and port mail server
			if m.s != nil {
				conn, err = m.s.Dial("tcp", smx)
			} else {
				conn, err = m.n.Dial("tcp", smx)
			}
			if err == nil {
				defer conn.Close()
				break
			}
		}
	}
	if err != nil {
		return err
	}

	host, _, _ := net.SplitHostPort(smx)
	c, err := smtp.NewClient(conn, host)
	if err != nil {
		return err
	}

	if err := c.Hello(m.Host); err != nil {
		return err
	}

	// Set the sender and recipient first
	if err := c.Mail(m.From); err != nil {
		return err
	}

	if err := c.Rcpt(m.To); err != nil {
		return err
	}

	msg := m.makeMail()

	//dkim.New()

	w, err := c.Data()
	if err != nil {
		return err
	}
	_, err = fmt.Fprintf(w, msg)
	if err != nil {
		return err
	}

	err = w.Close()
	if err != nil {
		return err
	}

	// Send the QUIT command and close the connection.
	return c.Quit()
}

func (m *MailData) makeMail() (msg string) {
	marker := makeMarker()

	msg = ""

	if m.From_name == "" {
		msg += `From: ` + m.From + "\r\n"
	} else {
		msg += `From: "` + encodeRFC2047(m.From_name) + `" <` + m.From + `>` +"\r\n"
	}

	if m.To_name == "" {
		msg += `To: ` + m.To + "\r\n"
	} else {
		msg += `To: "` + encodeRFC2047(m.To_name) + `" <` + m.To + `>` + "\r\n"
	}

	// -------------- head ----------------------------------------------------------

	msg += "Subject: " + encodeRFC2047(m.Subject) + "\r\n"
	msg += "MIME-Version: 1.0\r\n"
	msg += "Content-Type: multipart/mixed;\r\n	boundary=\"" + marker + "\"\r\n"
	msg += "X-Mailer: Gonder 0.2\r\n"
	msg += m.Extra_header + "\r\n"
	// ------------- /head ---------------------------------------------------------

	// ------------- body ----------------------------------------------------------
	msg += "\r\n--" + marker + "\r\nContent-Type: text/html; charset=\"utf-8\"\r\nContent-Transfer-Encoding: 8bit\r\n\r\n"
	msg += m.Html
	// ------------ /body ---------------------------------------------------------

	// ----------- attachments ----------------------------------------------------
	for _, file := range m.Attachments {

		msg += "\r\n--" + marker
		//read and encode attachment
		content, _ := ioutil.ReadFile(file.Location + file.Name)
		encoded := base64.StdEncoding.EncodeToString(content)

		//split the encoded file in lines (doesn't matter, but low enough not to hit a max limit)
		lineMaxLength := 500
		nbrLines := len(encoded) / lineMaxLength

		//append lines to buffer
		var buf bytes.Buffer
		for i := 0; i < nbrLines; i++ {
			buf.WriteString(encoded[i*lineMaxLength:(i+1)*lineMaxLength] + "\n")
		} //for

		//append last line in buffer
		buf.WriteString(encoded[nbrLines*lineMaxLength:])

		//part 3 will be the attachment
		msg += fmt.Sprintf("\r\nContent-Type: %s;\r\n	name=\"%s\"\r\nContent-Transfer-Encoding: base64\r\nContent-Disposition: attachment;\r\n	filename=\"%s\"\r\n\r\n%s\r\n", http.DetectContentType(content), file.Name, file.Name, buf.String())
	}
	// ----------- /attachments ---------------------------------------------------

	return
}

func makeMarker() string {
	b := make([]byte, 30)
	rand.Read(b)
	en := base64.StdEncoding // or URLEncoding
	d := make([]byte, en.EncodedLen(len(b)))
	en.Encode(d, b)
	return string(d)
}

func encodeRFC2047(s string) string {
	// use code from net/mail for rfc2047 encode any string
	// UTF-8 "Q" encoding
	b := bytes.NewBufferString("=?utf-8?q?")
	for i := 0; i < len(s); i++ {
		switch c := s[i]; {
		case c == ' ':
			b.WriteByte('_')
		case isVchar(c) && c != '=' && c != '?' && c != '_':
			b.WriteByte(c)
		default:
			fmt.Fprintf(b, "=%02X", c)
		}
	}
	b.WriteString("?= ")
	return b.String()
}

// isVchar returns true if c is an RFC 5322 VCHAR character.
func isVchar(c byte) bool {
	// Visible (printing) characters.
	return '!' <= c && c <= '~'
}

func Sender() {

	type campaignData struct {
		id, from, from_name, iface, host string
		attachments []attachmentData
		stream      int
		delay       int
	}

	type recipientData struct {
		id, to, to_name	string
	}

	campaign, err := Db.Prepare("SELECT t1.`id`,t1.`from`,t1.`from_name`,t2.`iface`,t2.`host`,t2.`stream`,t2.`delay` FROM `campaign` t1 INNER JOIN `interface` t2 ON t2.`id`=t1.`interface_id` WHERE NOW() BETWEEN t1.`start_time` AND t1.`end_time`")
	checkErr(err)
	defer campaign.Close()

	for {
		camp, err := campaign.Query()
		checkErr(err)
		defer camp.Close()

		var wc sync.WaitGroup
		for camp.Next() {

			c := new(campaignData)

			err = camp.Scan(&c.id, &c.from, &c.from_name, &c.iface, &c.host, &c.stream, &c.delay)
			checkErr(err)

			wc.Add(1)
			go func(cPart *campaignData) {
				attachment, err := Db.Prepare("SELECT `path`, `file` FROM attachment WHERE campaign_id=?")
				checkErr(err)
				defer attachment.Close()

				attach, err := attachment.Query(cPart.id)
				checkErr(err)
				defer attach.Close()

				cPart.attachments = nil

				var location string
				var name string
				for attach.Next() {
					err = attach.Scan(&location, &name)
					checkErr(err)
					cPart.attachments = append(cPart.attachments, attachmentData{Location: location, Name: name})
				}

				recipient, err := Db.Prepare("SELECT `id`, `email`, `name` FROM recipient WHERE campaign_id=? AND status IS NULL LIMIT ?")
				checkErr(err)
				defer recipient.Close()

				recip, err := recipient.Query(cPart.id, cPart.stream)
				checkErr(err)
				defer recip.Close()

				var wr sync.WaitGroup
				for recip.Next() {

					r := new(recipientData)

					err = recip.Scan(&r.id, &r.to, &r.to_name)
					checkErr(err)

					wr.Add(1)
					go func(cData *campaignData, rData *recipientData ) {
						data := new(MailData)
						data.Iface = cData.iface
						data.From = cData.from
						data.From_name = cData.from_name
						data.Host = cData.host
						data.Attachments = cData.attachments

						d := getMailMessage(cData.id, rData.id)
						data.Subject = d.Subject
						data.Html = d.Body
						data.Extra_header = "List-Unsubscribe: " + getUnsubscribeUrl(cData.id, rData.id) + "\r\nPrecedence: bulk\r\n"

						data.To = r.to
						data.To_name = r.to_name

						// Send mail
						res := data.SendMail()

						var rs string
						if res == nil {
							rs = "Ok"
						} else {
							rs = res.Error()
						}

						log.Printf("Send mail for recipient id %s email %s is %s", rData.id, data.To, rs)
						rows, err := Db.Query("UPDATE recipient SET status=?, date=NOW() WHERE id=?", rs, rData.id)
						checkErr(err)
						defer rows.Close()
						defer wr.Done()
					}(cPart, r)
				}
				wr.Wait()
				time.Sleep(time.Second + time.Duration(c.delay)*time.Second)
				defer wc.Done()
			}(c)
		}
		wc.Wait()
		time.Sleep(15 * time.Second) // easy with database
	}
}

func checkErr(err error) {
	if err != nil {
		log.Println(err)
	}
}
<commit_msg>correcting typos<commit_after>package mailer

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	_ "github.com/eaigner/dkim"
	"golang.org/x/net/proxy"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/smtp"
	"strings"
	"sync"
	"time"
	"golang.org/x/net/idna"
	"database/sql"
)

var HostName string

var Db *sql.DB

type message struct {
	Subject string
	Body    string
}

type pJson struct {
	Campaign    string `json:"c"`
	Recipient   string `json:"r"`
	Url         string `json:"u"`
	Webver      string `json:"w"`
	Opened      string `json:"o"`
	Unsubscribe string `json:"s"`
}

type attachmentData struct {
	Location string
	Name     string
}

type MailData struct {
	Iface 	     string
	Host         string
	From         string
	From_name    string
	To           string
	To_name      string
	Extra_header string
	Subject      string
	Html         string
	Attachments  []attachmentData
	s	     proxy.Dialer
	n            net.Dialer
}

func (m *MailData) SendMail() error {
	var smx string
	var mx []*net.MX
	var conn net.Conn

	if m.Iface == "" {
		// default interface
		m.n = net.Dialer{}
	} else {
		if m.Iface[0:8] == "socks://" {
			m.Iface = m.Iface[8:]
			var err error
			m.s, err = proxy.SOCKS5("tcp", m.Iface, nil, proxy.FromEnvironment())
			if err != nil {
				return err
			}
		} else {
			connectAddr := net.ParseIP(m.Iface)
			tcpAddr := &net.TCPAddr{
				IP: connectAddr,
			}
			m.n = net.Dialer{LocalAddr: tcpAddr}
		}
	}

	//ToDo cache MX servers
	// punycode convert
	domain, err := idna.ToASCII(strings.Split(m.To, "@")[1])
	if err != nil {
		return errors.New(fmt.Sprintf("Domain name failed: %v\r\n", err))
	}
	m.To = strings.Split(m.To, "@")[0] + "@" + domain

	mx, err = net.LookupMX(domain)
	if err != nil {
		return errors.New(fmt.Sprintf("LookupMX failed: %v\r\n", err))
	} else {
		for i := range mx {
			smx := net.JoinHostPort(mx[i].Host, "25")
			// Set ip (from MX records) and port mail server
			if m.s != nil {
				conn, err = m.s.Dial("tcp", smx)
			} else {
				conn, err = m.n.Dial("tcp", smx)
			}
			if err == nil {
				defer conn.Close()
				break
			}
		}
	}
	if err != nil {
		return err
	}

	host, _, _ := net.SplitHostPort(smx)
	c, err := smtp.NewClient(conn, host)
	if err != nil {
		return err
	}

	if err := c.Hello(m.Host); err != nil {
		return err
	}

	// Set the sender and recipient first
	if err := c.Mail(m.From); err != nil {
		return err
	}

	if err := c.Rcpt(m.To); err != nil {
		return err
	}

	msg := m.makeMail()

	//dkim.New()

	w, err := c.Data()
	if err != nil {
		return err
	}
	_, err = fmt.Fprintf(w, msg)
	if err != nil {
		return err
	}

	err = w.Close()
	if err != nil {
		return err
	}

	// Send the QUIT command and close the connection.
	return c.Quit()
}

func (m *MailData) makeMail() (msg string) {
	marker := makeMarker()

	msg = ""

	if m.From_name == "" {
		msg += `From: ` + m.From + "\r\n"
	} else {
		msg += `From: "` + encodeRFC2047(m.From_name) + `" <` + m.From + `>` +"\r\n"
	}

	if m.To_name == "" {
		msg += `To: ` + m.To + "\r\n"
	} else {
		msg += `To: "` + encodeRFC2047(m.To_name) + `" <` + m.To + `>` + "\r\n"
	}

	// -------------- head ----------------------------------------------------------

	msg += "Subject: " + encodeRFC2047(m.Subject) + "\r\n"
	msg += "MIME-Version: 1.0\r\n"
	msg += "Content-Type: multipart/mixed;\r\n	boundary=\"" + marker + "\"\r\n"
	msg += "X-Mailer: Gonder 0.2\r\n"
	msg += m.Extra_header + "\r\n"
	// ------------- /head ---------------------------------------------------------

	// ------------- body ----------------------------------------------------------
	msg += "\r\n--" + marker + "\r\nContent-Type: text/html; charset=\"utf-8\"\r\nContent-Transfer-Encoding: 8bit\r\n\r\n"
	msg += m.Html
	// ------------ /body ---------------------------------------------------------

	// ----------- attachments ----------------------------------------------------
	for _, file := range m.Attachments {

		msg += "\r\n--" + marker
		//read and encode attachment
		content, _ := ioutil.ReadFile(file.Location + file.Name)
		encoded := base64.StdEncoding.EncodeToString(content)

		//split the encoded file in lines (doesn't matter, but low enough not to hit a max limit)
		lineMaxLength := 500
		nbrLines := len(encoded) / lineMaxLength

		//append lines to buffer
		var buf bytes.Buffer
		for i := 0; i < nbrLines; i++ {
			buf.WriteString(encoded[i*lineMaxLength:(i+1)*lineMaxLength] + "\n")
		} //for

		//append last line in buffer
		buf.WriteString(encoded[nbrLines*lineMaxLength:])

		//part 3 will be the attachment
		msg += fmt.Sprintf("\r\nContent-Type: %s;\r\n	name=\"%s\"\r\nContent-Transfer-Encoding: base64\r\nContent-Disposition: attachment;\r\n	filename=\"%s\"\r\n\r\n%s\r\n", http.DetectContentType(content), file.Name, file.Name, buf.String())
	}
	// ----------- /attachments ---------------------------------------------------

	return
}

func makeMarker() string {
	b := make([]byte, 30)
	rand.Read(b)
	en := base64.StdEncoding // or URLEncoding
	d := make([]byte, en.EncodedLen(len(b)))
	en.Encode(d, b)
	return string(d)
}

func encodeRFC2047(s string) string {
	// use code from net/mail for rfc2047 encode any string
	// UTF-8 "Q" encoding
	b := bytes.NewBufferString("=?utf-8?q?")
	for i := 0; i < len(s); i++ {
		switch c := s[i]; {
		case c == ' ':
			b.WriteByte('_')
		case isVchar(c) && c != '=' && c != '?' && c != '_':
			b.WriteByte(c)
		default:
			fmt.Fprintf(b, "=%02X", c)
		}
	}
	b.WriteString("?= ")
	return b.String()
}

// isVchar returns true if c is an RFC 5322 VCHAR character.
func isVchar(c byte) bool {
	// Visible (printing) characters.
	return '!' <= c && c <= '~'
}

func Sender() {

	type campaignData struct {
		id, from, from_name, iface, host string
		attachments []attachmentData
		stream      int
		delay       int
	}

	type recipientData struct {
		id, to, to_name	string
	}

	campaign, err := Db.Prepare("SELECT t1.`id`,t1.`from`,t1.`from_name`,t2.`iface`,t2.`host`,t2.`stream`,t2.`delay` FROM `campaign` t1 INNER JOIN `interface` t2 ON t2.`id`=t1.`interface_id` WHERE NOW() BETWEEN t1.`start_time` AND t1.`end_time`")
	checkErr(err)
	defer campaign.Close()

	for {
		camp, err := campaign.Query()
		checkErr(err)
		defer camp.Close()

		var wc sync.WaitGroup
		for camp.Next() {

			c := new(campaignData)

			err = camp.Scan(&c.id, &c.from, &c.from_name, &c.iface, &c.host, &c.stream, &c.delay)
			checkErr(err)

			wc.Add(1)
			go func(cPart *campaignData) {
				attachment, err := Db.Prepare("SELECT `path`, `file` FROM attachment WHERE campaign_id=?")
				checkErr(err)
				defer attachment.Close()

				attach, err := attachment.Query(cPart.id)
				checkErr(err)
				defer attach.Close()

				cPart.attachments = nil

				var location string
				var name string
				for attach.Next() {
					err = attach.Scan(&location, &name)
					checkErr(err)
					cPart.attachments = append(cPart.attachments, attachmentData{Location: location, Name: name})
				}

				recipient, err := Db.Prepare("SELECT `id`, `email`, `name` FROM recipient WHERE campaign_id=? AND status IS NULL LIMIT ?")
				checkErr(err)
				defer recipient.Close()

				recip, err := recipient.Query(cPart.id, cPart.stream)
				checkErr(err)
				defer recip.Close()

				var wr sync.WaitGroup
				for recip.Next() {

					r := new(recipientData)

					err = recip.Scan(&r.id, &r.to, &r.to_name)
					checkErr(err)

					wr.Add(1)
					go func(cData *campaignData, rData *recipientData ) {
						data := new(MailData)
						data.Iface = cData.iface
						data.From = cData.from
						data.From_name = cData.from_name
						data.Host = cData.host
						data.Attachments = cData.attachments

						d := getMailMessage(cData.id, rData.id)
						data.Subject = d.Subject
						data.Html = d.Body
						data.Extra_header = "List-Unsubscribe: " + getUnsubscribeUrl(cData.id, rData.id) + "\r\nPrecedence: bulk\r\n"

						data.To = rData.to
						data.To_name = rData.to_name

						// Send mail
						res := data.SendMail()

						var rs string
						if res == nil {
							rs = "Ok"
						} else {
							rs = res.Error()
						}

						log.Printf("Send mail for recipient id %s email %s is %s", rData.id, data.To, rs)
						rows, err := Db.Query("UPDATE recipient SET status=?, date=NOW() WHERE id=?", rs, rData.id)
						checkErr(err)
						defer rows.Close()
						defer wr.Done()
					}(cPart, r)
				}
				wr.Wait()
				time.Sleep(time.Second + time.Duration(c.delay)*time.Second)
				defer wc.Done()
			}(c)
		}
		wc.Wait()
		time.Sleep(15 * time.Second) // easy with database
	}
}

func checkErr(err error) {
	if err != nil {
		log.Println(err)
	}
}
<|endoftext|>


<commit_before>/*Copyright (c) 2014 Alec Lofquist

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.*/

// This code was adapted from https://github.com/Aloshi/EmulationStation/pull/200

package mamedb

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"regexp"
	"strconv"
)

const (
	path = "http://www.mamedb.com/game/"
)

var (
	infoLineRE   = regexp.MustCompile("
Game Details
(.*?)Clock Speed")
	titleRE      = regexp.MustCompile("Name:&nbsp(?P<title>.*?)
.*?Year:&nbsp *(?P<date>.*?)
Manufacturer:&nbsp *(?P<developer>.*?)
Filename: (?P<filename>.*?)
")
	cleanTitleRE = regexp.MustCompile("^(.*?)&nbsp.*$")
	scoreRE      = regexp.MustCompile("Score: (.*?) \\(.*? votes\\)
")
	genreRE      = regexp.MustCompile("Category: (.*?)
")
	playersRE    = regexp.MustCompile("Players: (.*?)
")
	snapRE       = regexp.MustCompile("<img src='/snap/(.*?)\\.png'")
	titleImgRE   = regexp.MustCompile("<img src='/title/(.*?)\\.png'")
	cabinetRE    = regexp.MustCompile("<img src='/cabinets.small/(.*?)\\.(png|jpg|jpeg)'")
	marqueeRE    = regexp.MustCompile("<img src='/marquees.small/(.*?)\\.(png|jpg|jpeg)'")
	NotFound     = errors.New("rom not found")
)

type Game struct {
	ID        string
	Name      string
	Art       string
	Developer string
	Rating    float64
	Players   int64
	Genre     string
	Date      string
	Source    string
}

func GetGame(name string, imgPriority []string) (*Game, error) {
	var g Game
	g.Source = "mamedb.com"
	g.ID = name
	resp, err := http.Get(path + name)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode == 404 {
		return nil, NotFound
	}
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("got %d status", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	resp.Body.Close()
	ilm := infoLineRE.FindSubmatch(body)
	if ilm == nil {
		return nil, fmt.Errorf("ILM Bad HTML")
	}
	tm := titleRE.FindSubmatch(ilm[1])
	if tm == nil {
		return nil, fmt.Errorf("TM Bad HTML")
	}
	for i, n := range titleRE.SubexpNames() {
		switch n {
		case "title":
			ctn := cleanTitleRE.FindSubmatch(tm[i])
			if ctn != nil {
				g.Name = string(ctn[1])
			} else {
				g.Name = string(tm[i])
			}
		case "date":
			g.Date = string(tm[i])
		case "developer":
			developer := string(tm[i])
			if developer != "<unknown></unknown>" {
				g.Developer = developer
			}
		}
	}
	gm := genreRE.FindSubmatch(ilm[1])
	if gm != nil {
		g.Genre = string(gm[1])
	}
	rm := scoreRE.FindSubmatch(body)
	if rm != nil {
		rating, err := strconv.ParseFloat(string(rm[1]), 64)
		if err == nil {
			g.Rating = rating
		}
	}
	pm := playersRE.FindSubmatch(ilm[1])
	if pm != nil {
		players, err := strconv.ParseInt(string(pm[1]), 10, 64)
		if err == nil {
			g.Players = players
		}
	}
Loop:
	for _, i := range imgPriority {
		switch i {
		case "s":
			sm := snapRE.FindSubmatch(body)
			if sm != nil {
				g.Art = fmt.Sprintf("http://www.mamedb.com/snap/%s.png", string(sm[1]))
				break Loop
			}
		case "m":
			mm := marqueeRE.FindSubmatch(body)
			if mm != nil {
				g.Art = fmt.Sprintf("http://mamedb.com/marquees/%s.png", string(mm[1]))
				break Loop
			}
		case "t":
			tim := titleImgRE.FindSubmatch(body)
			if tim != nil {
				g.Art = fmt.Sprintf("http://mamedb.com/titles/%s.png", string(tim[1]))
				break Loop
			}
		case "c":
			cm := cabinetRE.FindSubmatch(body)
			if cm != nil {
				g.Art = fmt.Sprintf("http://mamedb.com/cabinets/%s.png", string(cm[1]))
				break Loop
			}
		}
	}
	return &g, nil
}
<commit_msg>Fix for #21<commit_after>/*Copyright (c) 2014 Alec Lofquist

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.*/

// This code was adapted from https://github.com/Aloshi/EmulationStation/pull/200

package mamedb

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"regexp"
	"strconv"
)

const (
	path = "http://www.mamedb.com/game/"
)

var (
	infoLineRE   = regexp.MustCompile("
Game Details
(.*?)Clock Speed")
	titleRE      = regexp.MustCompile("Name:&nbsp(?P<title>.*?)
.*?Year:&nbsp *(?P<date>.*?)
Manufacturer:&nbsp *(?P<developer>.*?)
Filename: (?P<filename>.*?)
")
	cleanTitleRE = regexp.MustCompile("^(.*?)&nbsp.*$")
	scoreRE      = regexp.MustCompile("Score: (.*?) \\(.*? votes\\)
")
	genreRE      = regexp.MustCompile("Category: (.*?)
")
	playersRE    = regexp.MustCompile("Players: (.*?)
")
	snapRE       = regexp.MustCompile("<img src='/snap/(.*?)\\.png'")
	titleImgRE   = regexp.MustCompile("<img src='/titles/(.*?)\\.png'")
	cabinetRE    = regexp.MustCompile("<img src='/cabinets.small/(.*?)\\.(png|jpg|jpeg)'")
	marqueeRE    = regexp.MustCompile("<img src='/marquees.small/(.*?)\\.(png|jpg|jpeg)'")
	NotFound     = errors.New("rom not found")
)

type Game struct {
	ID        string
	Name      string
	Art       string
	Developer string
	Rating    float64
	Players   int64
	Genre     string
	Date      string
	Source    string
}

func GetGame(name string, imgPriority []string) (*Game, error) {
	var g Game
	g.Source = "mamedb.com"
	g.ID = name
	resp, err := http.Get(path + name)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode == 404 {
		return nil, NotFound
	}
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("got %d status", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	resp.Body.Close()
	ilm := infoLineRE.FindSubmatch(body)
	if ilm == nil {
		return nil, fmt.Errorf("ILM Bad HTML")
	}
	tm := titleRE.FindSubmatch(ilm[1])
	if tm == nil {
		return nil, fmt.Errorf("TM Bad HTML")
	}
	for i, n := range titleRE.SubexpNames() {
		switch n {
		case "title":
			ctn := cleanTitleRE.FindSubmatch(tm[i])
			if ctn != nil {
				g.Name = string(ctn[1])
			} else {
				g.Name = string(tm[i])
			}
		case "date":
			g.Date = string(tm[i])
		case "developer":
			developer := string(tm[i])
			if developer != "<unknown></unknown>" {
				g.Developer = developer
			}
		}
	}
	gm := genreRE.FindSubmatch(ilm[1])
	if gm != nil {
		g.Genre = string(gm[1])
	}
	rm := scoreRE.FindSubmatch(body)
	if rm != nil {
		rating, err := strconv.ParseFloat(string(rm[1]), 64)
		if err == nil {
			g.Rating = rating
		}
	}
	pm := playersRE.FindSubmatch(ilm[1])
	if pm != nil {
		players, err := strconv.ParseInt(string(pm[1]), 10, 64)
		if err == nil {
			g.Players = players
		}
	}
Loop:
	for _, i := range imgPriority {
		switch i {
		case "s":
			sm := snapRE.FindSubmatch(body)
			if sm != nil {
				g.Art = fmt.Sprintf("http://www.mamedb.com/snap/%s.png", string(sm[1]))
				break Loop
			}
		case "m":
			mm := marqueeRE.FindSubmatch(body)
			if mm != nil {
				g.Art = fmt.Sprintf("http://mamedb.com/marquees/%s.png", string(mm[1]))
				break Loop
			}
		case "t":
			tim := titleImgRE.FindSubmatch(body)
			if tim != nil {
				g.Art = fmt.Sprintf("http://mamedb.com/titles/%s.png", string(tim[1]))
				break Loop
			}
		case "c":
			cm := cabinetRE.FindSubmatch(body)
			if cm != nil {
				g.Art = fmt.Sprintf("http://mamedb.com/cabinets/%s.png", string(cm[1]))
				break Loop
			}
		}
	}
	return &g, nil
}
<|endoftext|>


<commit_before>package mem

import (
	"io"
	"testing"
	"time"
)

func TestFileDataNameRace(t *testing.T) {
	t.Parallel()
	const someName = "someName"
	const someOtherName = "someOtherName"
	d := FileData{
		name: someName,
	}

	if d.Name() != someName {
		t.Errorf("Failed to read correct Name, was %v", d.Name())
	}

	ChangeFileName(&d, someOtherName)
	if d.Name() != someOtherName {
		t.Errorf("Failed to set Name, was %v", d.Name())
	}

	go func() {
		ChangeFileName(&d, someName)
	}()

	if d.Name() != someName && d.Name() != someOtherName {
		t.Errorf("Failed to read either Name, was %v", d.Name())
	}
}

func TestFileDataModTimeRace(t *testing.T) {
	t.Parallel()
	someTime := time.Now()
	someOtherTime := someTime.Add(1 * time.Minute)

	d := FileData{
		modtime: someTime,
	}

	s := FileInfo{
		FileData: &d,
	}

	if s.ModTime() != someTime {
		t.Errorf("Failed to read correct value, was %v", s.ModTime())
	}

	SetModTime(&d, someOtherTime)
	if s.ModTime() != someOtherTime {
		t.Errorf("Failed to set ModTime, was %v", s.ModTime())
	}

	go func() {
		SetModTime(&d, someTime)
	}()

	if s.ModTime() != someTime && s.ModTime() != someOtherTime {
		t.Errorf("Failed to read either modtime, was %v", s.ModTime())
	}
}

func TestFileDataModeRace(t *testing.T) {
	t.Parallel()
	const someMode = 0777
	const someOtherMode = 0660

	d := FileData{
		mode: someMode,
	}

	s := FileInfo{
		FileData: &d,
	}

	if s.Mode() != someMode {
		t.Errorf("Failed to read correct value, was %v", s.Mode())
	}

	SetMode(&d, someOtherMode)
	if s.Mode() != someOtherMode {
		t.Errorf("Failed to set Mode, was %v", s.Mode())
	}

	go func() {
		SetMode(&d, someMode)
	}()

	if s.Mode() != someMode && s.Mode() != someOtherMode {
		t.Errorf("Failed to read either mode, was %v", s.Mode())
	}
}

func TestFileDataIsDirRace(t *testing.T) {
	t.Parallel()

	d := FileData{
		dir: true,
	}

	s := FileInfo{
		FileData: &d,
	}

	if s.IsDir() != true {
		t.Errorf("Failed to read correct value, was %v", s.IsDir())
	}

	go func() {
		s.Lock()
		d.dir = false
		s.Unlock()
	}()

	//just logging the value to trigger a read:
	t.Logf("Value is %v", s.IsDir())
}

func TestFileDataSizeRace(t *testing.T) {
	t.Parallel()

	const someData = "Hello"
	const someOtherDataSize = "Hello World"

	d := FileData{
		data: []byte(someData),
		dir:  false,
	}

	s := FileInfo{
		FileData: &d,
	}

	if s.Size() != int64(len(someData)) {
		t.Errorf("Failed to read correct value, was %v", s.Size())
	}

	go func() {
		s.Lock()
		d.data = []byte(someOtherDataSize)
		s.Unlock()
	}()

	//just logging the value to trigger a read:
	t.Logf("Value is %v", s.Size())

	//Testing the Dir size case
	d.dir = true
	if s.Size() != int64(42) {
		t.Errorf("Failed to read correct value for dir, was %v", s.Size())
	}
}

func TestFileReadAtSeekOffset(t *testing.T) {
	t.Parallel()

	fd := CreateFile("foo")
	f := NewFileHandle(fd)

	_, err := f.WriteString("TEST")
	if err != nil {
		t.Fatal(err)
	}
	offset, err := f.Seek(0, io.SeekStart)
	if err != nil {
		t.Fatal(err)
	}
	if offset != 0 {
		t.Fail()
	}

	offsetBeforeReadAt, err := f.Seek(0, io.SeekCurrent)
	if err != nil {
		t.Fatal(err)
	}
	if offsetBeforeReadAt != 0 {
		t.Fatal("expected 0")
	}

	b := make([]byte, 4)
	n, err := f.ReadAt(b, 0)
	if err != nil {
		t.Fatal(err)
	}
	if n != 4 {
		t.Fail()
	}
	if string(b) != "TEST" {
		t.Fail()
	}

	offsetAfterReadAt, err := f.Seek(0, io.SeekCurrent)
	if err != nil {
		t.Fatal(err)
	}
	if offsetAfterReadAt != offsetBeforeReadAt {
		t.Fatal("ReadAt should not affect offset")
	}

	err = f.Close()
	if err != nil {
		t.Fatal(err)
	}
}
<commit_msg>mem: add test for Write/Seek<commit_after>package mem

import (
	"bytes"
	"io"
	"testing"
	"time"
)

func TestFileDataNameRace(t *testing.T) {
	t.Parallel()
	const someName = "someName"
	const someOtherName = "someOtherName"
	d := FileData{
		name: someName,
	}

	if d.Name() != someName {
		t.Errorf("Failed to read correct Name, was %v", d.Name())
	}

	ChangeFileName(&d, someOtherName)
	if d.Name() != someOtherName {
		t.Errorf("Failed to set Name, was %v", d.Name())
	}

	go func() {
		ChangeFileName(&d, someName)
	}()

	if d.Name() != someName && d.Name() != someOtherName {
		t.Errorf("Failed to read either Name, was %v", d.Name())
	}
}

func TestFileDataModTimeRace(t *testing.T) {
	t.Parallel()
	someTime := time.Now()
	someOtherTime := someTime.Add(1 * time.Minute)

	d := FileData{
		modtime: someTime,
	}

	s := FileInfo{
		FileData: &d,
	}

	if s.ModTime() != someTime {
		t.Errorf("Failed to read correct value, was %v", s.ModTime())
	}

	SetModTime(&d, someOtherTime)
	if s.ModTime() != someOtherTime {
		t.Errorf("Failed to set ModTime, was %v", s.ModTime())
	}

	go func() {
		SetModTime(&d, someTime)
	}()

	if s.ModTime() != someTime && s.ModTime() != someOtherTime {
		t.Errorf("Failed to read either modtime, was %v", s.ModTime())
	}
}

func TestFileDataModeRace(t *testing.T) {
	t.Parallel()
	const someMode = 0777
	const someOtherMode = 0660

	d := FileData{
		mode: someMode,
	}

	s := FileInfo{
		FileData: &d,
	}

	if s.Mode() != someMode {
		t.Errorf("Failed to read correct value, was %v", s.Mode())
	}

	SetMode(&d, someOtherMode)
	if s.Mode() != someOtherMode {
		t.Errorf("Failed to set Mode, was %v", s.Mode())
	}

	go func() {
		SetMode(&d, someMode)
	}()

	if s.Mode() != someMode && s.Mode() != someOtherMode {
		t.Errorf("Failed to read either mode, was %v", s.Mode())
	}
}

func TestFileDataIsDirRace(t *testing.T) {
	t.Parallel()

	d := FileData{
		dir: true,
	}

	s := FileInfo{
		FileData: &d,
	}

	if s.IsDir() != true {
		t.Errorf("Failed to read correct value, was %v", s.IsDir())
	}

	go func() {
		s.Lock()
		d.dir = false
		s.Unlock()
	}()

	//just logging the value to trigger a read:
	t.Logf("Value is %v", s.IsDir())
}

func TestFileDataSizeRace(t *testing.T) {
	t.Parallel()

	const someData = "Hello"
	const someOtherDataSize = "Hello World"

	d := FileData{
		data: []byte(someData),
		dir:  false,
	}

	s := FileInfo{
		FileData: &d,
	}

	if s.Size() != int64(len(someData)) {
		t.Errorf("Failed to read correct value, was %v", s.Size())
	}

	go func() {
		s.Lock()
		d.data = []byte(someOtherDataSize)
		s.Unlock()
	}()

	//just logging the value to trigger a read:
	t.Logf("Value is %v", s.Size())

	//Testing the Dir size case
	d.dir = true
	if s.Size() != int64(42) {
		t.Errorf("Failed to read correct value for dir, was %v", s.Size())
	}
}

func TestFileReadAtSeekOffset(t *testing.T) {
	t.Parallel()

	fd := CreateFile("foo")
	f := NewFileHandle(fd)

	_, err := f.WriteString("TEST")
	if err != nil {
		t.Fatal(err)
	}
	offset, err := f.Seek(0, io.SeekStart)
	if err != nil {
		t.Fatal(err)
	}
	if offset != 0 {
		t.Fail()
	}

	offsetBeforeReadAt, err := f.Seek(0, io.SeekCurrent)
	if err != nil {
		t.Fatal(err)
	}
	if offsetBeforeReadAt != 0 {
		t.Fatal("expected 0")
	}

	b := make([]byte, 4)
	n, err := f.ReadAt(b, 0)
	if err != nil {
		t.Fatal(err)
	}
	if n != 4 {
		t.Fail()
	}
	if string(b) != "TEST" {
		t.Fail()
	}

	offsetAfterReadAt, err := f.Seek(0, io.SeekCurrent)
	if err != nil {
		t.Fatal(err)
	}
	if offsetAfterReadAt != offsetBeforeReadAt {
		t.Fatal("ReadAt should not affect offset")
	}

	err = f.Close()
	if err != nil {
		t.Fatal(err)
	}
}

func TestFileWriteAndSeek(t *testing.T) {
	fd := CreateFile("foo")
	f := NewFileHandle(fd)

	assert := func(expected bool, v ...interface{}) {
		if !expected {
			t.Helper()
			t.Fatal(v...)
		}
	}

	data4 := []byte{0, 1, 2, 3}
	data20 := bytes.Repeat(data4, 5)
	var off int64

	for i := 0; i < 100; i++ {
		// write 20 bytes
		n, err := f.Write(data20)
		assert(err == nil, err)
		off += int64(n)
		assert(n == len(data20), n)
		assert(off == int64((i+1)*len(data20)), off)

		// rewind to start and write 4 bytes there
		cur, err := f.Seek(-off, io.SeekCurrent)
		assert(err == nil, err)
		assert(cur == 0, cur)

		n, err = f.Write(data4)
		assert(err == nil, err)
		assert(n == len(data4), n)

		// back at the end
		cur, err = f.Seek(off-int64(n), io.SeekCurrent)
		assert(err == nil, err)
		assert(cur == off, cur, off)
	}
}
<|endoftext|>


<commit_before>package lib

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws/arn"
	"github.com/aws/aws-sdk-go/aws/defaults"
	"github.com/go-ini/ini"
	"github.com/mbndr/logo"
	"os"
	"strings"
)

// Prefix for role ARNs and Virtual MFA devices
// (physical MFA devices use device serial number, not ARN)
const IAM_ARN = "arn:aws:iam::"

// Profile information from known configuration attributes
// in the aws sdk configuration file
type AWSProfile struct {
	SourceProfile   string `ini:"source_profile"`
	RoleArn         string `ini:"role_arn"`
	MfaSerial       string `ini:"mfa_serial"`
	ExternalId      string `ini:"external_id"`
	RoleSessionName string `ini:"role_session_name"`
	Region          string `ini:"region"`
	name            string
}

// Interface for managing an ini-formatted configuration file
// providing the ability to retrieve a default profile, profile
// by name, and to build a configuration file based on a
// given list of Roles
type ConfigManager interface {
	DefaultProfile() (*AWSProfile, error)
	GetProfile(name *string) (*AWSProfile, error)
	BuildConfig(roles Roles, mfa *string) error
}

type ConfigManagerOptions struct {
	LogLevel logo.Level
}

// A ConfigManager specific to the aws sdk configuration.  It will use
// the environment variable AWS_CONFIG_FILE to determine the file to
// load, and if not defined fall back to the SDK default value.
func NewAwsConfigManager(opts *ConfigManagerOptions) (ConfigManager, error) {
	c := defaults.SharedConfigFilename()
	e, ok := os.LookupEnv("AWS_CONFIG_FILE")
	if ok && len(e) > 0 {
		c = e
	}

	f, err := ini.Load(c)
	if err != nil {
		return nil, err
	}
	f.BlockMode = false

	logger := logo.NewSimpleLogger(os.Stderr, opts.LogLevel, "aws-runas.ConfigManager", true)

	return &awsConfigManager{config: f, log: logger}, nil
}

type awsConfigManager struct {
	config *ini.File
	log    *logo.Logger
}

// Return an AWSProfile for the default aws configuration profile.
// The profile name set in the AWS_DEFAULT_PROFILE environment variable
// and if not set, use a value of "default"
func (c *awsConfigManager) DefaultProfile() (*AWSProfile, error) {
	s := defaultSection()
	p := &AWSProfile{SourceProfile: s, name: s}
	if err := c.profile(p); err != nil {
		return nil, err
	}
	return p, nil
}

// Retrieve an AWSProfile by name.  The default profile will be looked up
// first to provide default settings, and then the profile-specific values
// will be retrieved.  If the specified profile contains a role_arn setting
// that value will be checked to ensure it's a valid IAM arn.
func (c *awsConfigManager) GetProfile(p *string) (*AWSProfile, error) {
	if p == nil || len(*p) < 1 {
		return nil, fmt.Errorf("nil or empty profile name")
	}

	profile, err := c.DefaultProfile()
	if err != nil {
		return nil, err
	}
	profile.name = *p

	if err := c.profile(profile); err != nil {
		return nil, err
	}

	if len(profile.RoleArn) > 0 {
		// Validate that RoleArn is a correctly formatted ARN
		a, err := arn.Parse(profile.RoleArn)
		if err != nil {
			return nil, err
		}

		// Validate that RoleArn is an IAM ARN
		if !strings.HasPrefix(a.String(), IAM_ARN) {
			return nil, fmt.Errorf("role ARN format error, does not start with %s", IAM_ARN)
		}

		// TODO also check that SourceProfile is valid?
		if len(profile.SourceProfile) < 1 {
			return nil, fmt.Errorf("role_arn configured, but missing required source_profile")
		}
	}

	return profile, nil
}

// Build an AWS SDK compliant, ini-formatted, configuration file based on the location
// configured during the NewAwsConfigManager() call.  The generated config file will
// have a default section, with all configured roles to use the default section as
// the source_profile, and MFA configured, if the mfa parameter is not nil or empty.
func (c *awsConfigManager) BuildConfig(r Roles, mfa *string) error {
	// TODO build config based on provided Roles using file name in c.config
	// Do NOT overwrite file if it already exists!
	return nil
}

func (c *awsConfigManager) profile(p *AWSProfile) error {
	name := p.name
	if name != p.SourceProfile {
		name = "profile " + name
	}

	s, err := c.config.GetSection(name)
	if err != nil {
		return err
	}

	if err := s.MapTo(p); err != nil {
		return err
	}
	return nil
}

func defaultSection() string {
	s := "default"
	v, ok := os.LookupEnv("AWS_DEFAULT_PROFILE")
	if ok {
		s = v
	}
	return s
}
<commit_msg>Remove hard-coded "default" profile name<commit_after>package lib

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws/arn"
	"github.com/aws/aws-sdk-go/aws/defaults"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/go-ini/ini"
	"github.com/mbndr/logo"
	"os"
	"strings"
)

// Prefix for role ARNs and Virtual MFA devices
// (physical MFA devices use device serial number, not ARN)
const IAM_ARN = "arn:aws:iam::"

// Profile information from known configuration attributes
// in the aws sdk configuration file
type AWSProfile struct {
	SourceProfile   string `ini:"source_profile"`
	RoleArn         string `ini:"role_arn"`
	MfaSerial       string `ini:"mfa_serial"`
	ExternalId      string `ini:"external_id"`
	RoleSessionName string `ini:"role_session_name"`
	Region          string `ini:"region"`
	name            string
}

// Interface for managing an ini-formatted configuration file
// providing the ability to retrieve a default profile, profile
// by name, and to build a configuration file based on a
// given list of Roles
type ConfigManager interface {
	DefaultProfile() (*AWSProfile, error)
	GetProfile(name *string) (*AWSProfile, error)
	BuildConfig(roles Roles, mfa *string) error
}

type ConfigManagerOptions struct {
	LogLevel logo.Level
}

// A ConfigManager specific to the aws sdk configuration.  It will use
// the environment variable AWS_CONFIG_FILE to determine the file to
// load, and if not defined fall back to the SDK default value.
func NewAwsConfigManager(opts *ConfigManagerOptions) (ConfigManager, error) {
	c := defaults.SharedConfigFilename()
	e, ok := os.LookupEnv("AWS_CONFIG_FILE")
	if ok && len(e) > 0 {
		c = e
	}

	f, err := ini.Load(c)
	if err != nil {
		return nil, err
	}
	f.BlockMode = false

	logger := logo.NewSimpleLogger(os.Stderr, opts.LogLevel, "aws-runas.ConfigManager", true)

	return &awsConfigManager{config: f, log: logger}, nil
}

type awsConfigManager struct {
	config *ini.File
	log    *logo.Logger
}

// Return an AWSProfile for the default aws configuration profile.
// The profile name set in the AWS_DEFAULT_PROFILE environment variable
// and if not set, use a value of "default"
func (c *awsConfigManager) DefaultProfile() (*AWSProfile, error) {
	s := defaultSection()
	p := &AWSProfile{SourceProfile: s, name: s}
	if err := c.profile(p); err != nil {
		return nil, err
	}
	return p, nil
}

// Retrieve an AWSProfile by name.  The default profile will be looked up
// first to provide default settings, and then the profile-specific values
// will be retrieved.  If the specified profile contains a role_arn setting
// that value will be checked to ensure it's a valid IAM arn.
func (c *awsConfigManager) GetProfile(p *string) (*AWSProfile, error) {
	if p == nil || len(*p) < 1 {
		return nil, fmt.Errorf("nil or empty profile name")
	}

	profile, err := c.DefaultProfile()
	if err != nil {
		return nil, err
	}
	profile.name = *p

	if err := c.profile(profile); err != nil {
		return nil, err
	}

	if len(profile.RoleArn) > 0 {
		// Validate that RoleArn is a correctly formatted ARN
		a, err := arn.Parse(profile.RoleArn)
		if err != nil {
			return nil, err
		}

		// Validate that RoleArn is an IAM ARN
		if !strings.HasPrefix(a.String(), IAM_ARN) {
			return nil, fmt.Errorf("role ARN format error, does not start with %s", IAM_ARN)
		}

		// TODO also check that SourceProfile is valid?
		if len(profile.SourceProfile) < 1 {
			return nil, fmt.Errorf("role_arn configured, but missing required source_profile")
		}
	}

	return profile, nil
}

// Build an AWS SDK compliant, ini-formatted, configuration file based on the location
// configured during the NewAwsConfigManager() call.  The generated config file will
// have a default section, with all configured roles to use the default section as
// the source_profile, and MFA configured, if the mfa parameter is not nil or empty.
func (c *awsConfigManager) BuildConfig(r Roles, mfa *string) error {
	// TODO build config based on provided Roles using file name in c.config
	// Do NOT overwrite file if it already exists!
	return nil
}

func (c *awsConfigManager) profile(p *AWSProfile) error {
	name := p.name
	if name != p.SourceProfile {
		name = "profile " + name
	}

	s, err := c.config.GetSection(name)
	if err != nil {
		return err
	}

	if err := s.MapTo(p); err != nil {
		return err
	}
	return nil
}

func defaultSection() string {
	s := session.DefaultSharedConfigProfile
	v, ok := os.LookupEnv("AWS_DEFAULT_PROFILE")
	if ok {
		s = v
	}
	return s
}
<|endoftext|>