<commit_before>// Simple console progress bars
package pb

import (
	"fmt"
	"io"
	"math"
	"strings"
	"sync"
	"sync/atomic"
	"time"
	"unicode/utf8"
)

// Current version
const Version = "1.0.25"

const (
	// Default refresh rate - 200ms
	DEFAULT_REFRESH_RATE = time.Millisecond * 200
	FORMAT               = "[=>-]"
)

// DEPRECATED
// variables for backward compatibility, from now do not work
// use pb.Format and pb.SetRefreshRate
var (
	DefaultRefreshRate                         = DEFAULT_REFRESH_RATE
	BarStart, BarEnd, Empty, Current, CurrentN string
)

// Create new progress bar object
func New(total int) *ProgressBar {
	return New64(int64(total))
}

// Create new progress bar object using int64 as total
func New64(total int64) *ProgressBar {
	pb := &ProgressBar{
		Total:           total,
		RefreshRate:     DEFAULT_REFRESH_RATE,
		ShowPercent:     true,
		ShowCounters:    true,
		ShowBar:         true,
		ShowTimeLeft:    true,
		ShowElapsedTime: false,
		ShowFinalTime:   true,
		Units:           U_NO,
		ManualUpdate:    false,
		finish:          make(chan struct{}),
	}
	return pb.Format(FORMAT)
}

// Create new object and start
func StartNew(total int) *ProgressBar {
	return New(total).Start()
}

// Callback for custom output
// For example:
// bar.Callback = func(s string) {
//     mySuperPrint(s)
// }
//
type Callback func(out string)

type ProgressBar struct {
	current  int64 // current must be first member of struct (https://code.google.com/p/go/issues/detail?id=5278)
	previous int64

	Total                            int64
	RefreshRate                      time.Duration
	ShowPercent, ShowCounters        bool
	ShowSpeed, ShowTimeLeft, ShowBar bool
	ShowFinalTime, ShowElapsedTime   bool
	Output                           io.Writer
	Callback                         Callback
	NotPrint                         bool
	Units                            Units
	Width                            int
	ForceWidth                       bool
	ManualUpdate                     bool
	AutoStat                         bool

	// Default width for the time box.
	UnitsWidth   int
	TimeBoxWidth int

	finishOnce sync.Once //Guards isFinish
	finish     chan struct{}
	isFinish   bool

	startTime  time.Time
	startValue int64

	changeTime time.Time

	prefix, postfix string

	mu        sync.Mutex
	lastPrint string

	BarStart string
	BarEnd   string
	Empty    string
	Current  string
	CurrentN string

	AlwaysUpdate bool
}

// Start print
func (pb *ProgressBar) Start() *ProgressBar {
	pb.startTime = time.Now()
	pb.startValue = atomic.LoadInt64(&pb.current)
	if atomic.LoadInt64(&pb.Total) == 0 {
		pb.ShowTimeLeft = false
		pb.ShowPercent = false
		pb.AutoStat = false
	}
	if !pb.ManualUpdate {
		pb.Update() // Initial printing of the bar before running the bar refresher.
		go pb.refresher()
	}
	return pb
}

// Increment current value
func (pb *ProgressBar) Increment() int {
	return pb.Add(1)
}

// Get current value
func (pb *ProgressBar) Get() int64 {
	c := atomic.LoadInt64(&pb.current)
	return c
}

// Set current value
func (pb *ProgressBar) Set(current int) *ProgressBar {
	return pb.Set64(int64(current))
}

// Set64 sets the current value as int64
func (pb *ProgressBar) Set64(current int64) *ProgressBar {
	atomic.StoreInt64(&pb.current, current)
	return pb
}

// Add to current value
func (pb *ProgressBar) Add(add int) int {
	return int(pb.Add64(int64(add)))
}

func (pb *ProgressBar) Add64(add int64) int64 {
	return atomic.AddInt64(&pb.current, add)
}

// Set prefix string
func (pb *ProgressBar) Prefix(prefix string) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	pb.prefix = prefix
	return pb
}

// Set postfix string
func (pb *ProgressBar) Postfix(postfix string) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	pb.postfix = postfix
	return pb
}

// Set custom format for bar
// Example: bar.Format("[=>_]")
// Example: bar.Format("[\x00=\x00>\x00-\x00]") // \x00 is the delimiter
func (pb *ProgressBar) Format(format string) *ProgressBar {
	var formatEntries []string
	if utf8.RuneCountInString(format) == 5 {
		formatEntries = strings.Split(format, "")
	} else {
		formatEntries = strings.Split(format, "\x00")
	}
	if len(formatEntries) == 5 {
		pb.BarStart = formatEntries[0]
		pb.BarEnd = formatEntries[4]
		pb.Empty = formatEntries[3]
		pb.Current = formatEntries[1]
		pb.CurrentN = formatEntries[2]
	}
	return pb
}

// Set bar refresh rate
func (pb *ProgressBar) SetRefreshRate(rate time.Duration) *ProgressBar {
	pb.RefreshRate = rate
	return pb
}

// Set units
// bar.SetUnits(U_NO) - by default
// bar.SetUnits(U_BYTES) - for Mb, Kb, etc
func (pb *ProgressBar) SetUnits(units Units) *ProgressBar {
	pb.Units = units
	return pb
}

// Set max width, if width is bigger than terminal width, will be ignored
func (pb *ProgressBar) SetMaxWidth(width int) *ProgressBar {
	pb.Width = width
	pb.ForceWidth = false
	return pb
}

// Set bar width
func (pb *ProgressBar) SetWidth(width int) *ProgressBar {
	pb.Width = width
	pb.ForceWidth = true
	return pb
}

// End print
func (pb *ProgressBar) Finish() {
	//Protect multiple calls
	pb.finishOnce.Do(func() {
		close(pb.finish)
		pb.write(atomic.LoadInt64(&pb.Total), atomic.LoadInt64(&pb.current))
		pb.mu.Lock()
		defer pb.mu.Unlock()
		switch {
		case pb.Output != nil:
			fmt.Fprintln(pb.Output)
		case !pb.NotPrint:
			fmt.Println()
		}
		pb.isFinish = true
	})
}

// IsFinished return boolean
func (pb *ProgressBar) IsFinished() bool {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	return pb.isFinish
}

// End print and write string 'str'
func (pb *ProgressBar) FinishPrint(str string) {
	pb.Finish()
	if pb.Output != nil {
		fmt.Fprintln(pb.Output, str)
	} else {
		fmt.Println(str)
	}
}

// implement io.Writer
func (pb *ProgressBar) Write(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

// implement io.Reader
func (pb *ProgressBar) Read(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

// Create new proxy reader over bar
// Takes io.Reader or io.ReadCloser
func (pb *ProgressBar) NewProxyReader(r io.Reader) *Reader {
	return &Reader{r, pb}
}

func (pb *ProgressBar) write(total, current int64) {
	width := pb.GetWidth()

	var percentBox, countersBox, timeLeftBox, timeSpentBox, speedBox, barBox, end, out string

	// percents
	if pb.ShowPercent {
		var percent float64
		if total > 0 {
			percent = float64(current) / (float64(total) / float64(100))
		} else {
			percent = float64(current) / float64(100)
		}
		percentBox = fmt.Sprintf(" %6.02f%%", percent)
	}

	// counters
	if pb.ShowCounters {
		current := Format(current).To(pb.Units).Width(pb.UnitsWidth)
		if total > 0 {
			totalS := Format(total).To(pb.Units).Width(pb.UnitsWidth)
			countersBox = fmt.Sprintf(" %s / %s ", current, totalS)
		} else {
			countersBox = fmt.Sprintf(" %s / ? ", current)
		}
	}

	// time left
	pb.mu.Lock()
	currentFromStart := current - pb.startValue
	fromStart := time.Now().Sub(pb.startTime)
	lastChangeTime := pb.changeTime
	fromChange := lastChangeTime.Sub(pb.startTime)
	pb.mu.Unlock()

	if pb.ShowElapsedTime {
		timeSpentBox = fmt.Sprintf(" %s ", (fromStart/time.Second)*time.Second)
	}

	select {
	case <-pb.finish:
		if pb.ShowFinalTime {
			var left time.Duration
			left = (fromStart / time.Second) * time.Second
			timeLeftBox = fmt.Sprintf(" %s", left.String())
		}
	default:
		if pb.ShowTimeLeft && currentFromStart > 0 {
			perEntry := fromChange / time.Duration(currentFromStart)
			var left time.Duration
			if total > 0 {
				left = time.Duration(total-currentFromStart) * perEntry
				left -= time.Since(lastChangeTime)
				left = (left / time.Second) * time.Second
			} else {
				left = time.Duration(currentFromStart) * perEntry
				left = (left / time.Second) * time.Second
			}
			if left > 0 {
				timeLeft := Format(int64(left)).To(U_DURATION).String()
				timeLeftBox = fmt.Sprintf(" %s", timeLeft)
			}
		}
	}

	if len(timeLeftBox) < pb.TimeBoxWidth {
		timeLeftBox = fmt.Sprintf("%s%s", strings.Repeat(" ", pb.TimeBoxWidth-len(timeLeftBox)), timeLeftBox)
	}

	// speed
	if pb.ShowSpeed && currentFromStart > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		speed := float64(currentFromStart) / (float64(fromStart) / float64(time.Second))
		speedBox = " " + Format(int64(speed)).To(pb.Units).Width(pb.UnitsWidth).PerSec().String()
	}

	pb.mu.Lock()
	barWidth := escapeAwareRuneCountInString(countersBox + pb.BarStart + pb.BarEnd + percentBox + timeSpentBox + timeLeftBox + speedBox + pb.prefix + pb.postfix)
	pb.mu.Unlock()
	// bar
	if pb.ShowBar {
		size := width - barWidth
		if size > 0 {
			if total > 0 {
				curSize := int(math.Ceil((float64(current) / float64(total)) * float64(size)))
				emptySize := size - curSize
				barBox = pb.BarStart
				if emptySize < 0 {
					emptySize = 0
				}
				if curSize > size {
					curSize = size
				}

				cursorLen := escapeAwareRuneCountInString(pb.Current)
				if emptySize <= 0 {
					barBox += strings.Repeat(pb.Current, curSize/cursorLen)
				} else if curSize > 0 {
					cursorEndLen := escapeAwareRuneCountInString(pb.CurrentN)
					cursorRepetitions := (curSize - cursorEndLen) / cursorLen
					barBox += strings.Repeat(pb.Current, cursorRepetitions)
					barBox += pb.CurrentN
				}

				emptyLen := escapeAwareRuneCountInString(pb.Empty)
				barBox += strings.Repeat(pb.Empty, emptySize/emptyLen)
				barBox += pb.BarEnd
			} else {
				pos := size - int(current)%int(size)
				barBox = pb.BarStart
				if pos-1 > 0 {
					barBox += strings.Repeat(pb.Empty, pos-1)
				}
				barBox += pb.Current
				if size-pos-1 > 0 {
					barBox += strings.Repeat(pb.Empty, size-pos-1)
				}
				barBox += pb.BarEnd
			}
		}
	}

	// check len
	pb.mu.Lock()
	defer pb.mu.Unlock()
	out = pb.prefix + timeSpentBox + countersBox + barBox + percentBox + speedBox + timeLeftBox + pb.postfix

	if cl := escapeAwareRuneCountInString(out); cl < width {
		end = strings.Repeat(" ", width-cl)
	}

	// and print!
	pb.lastPrint = out + end
	isFinish := pb.isFinish

	switch {
	case isFinish:
		return
	case pb.Output != nil:
		fmt.Fprint(pb.Output, "\r"+out+end)
	case pb.Callback != nil:
		pb.Callback(out + end)
	case !pb.NotPrint:
		fmt.Print("\r" + out + end)
	}
}

// GetTerminalWidth - returns terminal width for all platforms.
func GetTerminalWidth() (int, error) {
	return terminalWidth()
}

func (pb *ProgressBar) GetWidth() int {
	if pb.ForceWidth {
		return pb.Width
	}

	width := pb.Width
	termWidth, _ := terminalWidth()
	if width == 0 || termWidth <= width {
		width = termWidth
	}

	return width
}

// Write the current state of the progressbar
func (pb *ProgressBar) Update() {
	c := atomic.LoadInt64(&pb.current)
	p := atomic.LoadInt64(&pb.previous)
	t := atomic.LoadInt64(&pb.Total)
	if p != c {
		pb.mu.Lock()
		pb.changeTime = time.Now()
		pb.mu.Unlock()
		atomic.StoreInt64(&pb.previous, c)
	}
	pb.write(t, c)
	if pb.AutoStat {
		if c == 0 {
			pb.startTime = time.Now()
			pb.startValue = 0
		} else if c >= t && pb.isFinish != true {
			pb.Finish()
		}
	}
}

// String return the last bar print
func (pb *ProgressBar) String() string {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	return pb.lastPrint
}

// SetTotal atomically sets new total count
func (pb *ProgressBar) SetTotal(total int) *ProgressBar {
	return pb.SetTotal64(int64(total))
}

// SetTotal64 atomically sets new total count
func (pb *ProgressBar) SetTotal64(total int64) *ProgressBar {
	atomic.StoreInt64(&pb.Total, total)
	return pb
}

// Reset bar and set new total count
// Does effect only on finished bar
func (pb *ProgressBar) Reset(total int) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	if pb.isFinish {
		pb.SetTotal(total).Set(0)
		atomic.StoreInt64(&pb.previous, 0)
	}
	return pb
}

// Internal loop for refreshing the progressbar
func (pb *ProgressBar) refresher() {
	for {
		select {
		case <-pb.finish:
			return
		case <-time.After(pb.RefreshRate):
			pb.Update()
		}
	}
}
<commit_msg>protect write with a lock<commit_after>// Simple console progress bars
package pb

import (
	"fmt"
	"io"
	"math"
	"strings"
	"sync"
	"sync/atomic"
	"time"
	"unicode/utf8"
)

// Current version
const Version = "1.0.25"

const (
	// Default refresh rate - 200ms
	DEFAULT_REFRESH_RATE = time.Millisecond * 200
	FORMAT               = "[=>-]"
)

// DEPRECATED
// variables for backward compatibility, from now do not work
// use pb.Format and pb.SetRefreshRate
var (
	DefaultRefreshRate                         = DEFAULT_REFRESH_RATE
	BarStart, BarEnd, Empty, Current, CurrentN string
)

// Create new progress bar object
func New(total int) *ProgressBar {
	return New64(int64(total))
}

// Create new progress bar object using int64 as total
func New64(total int64) *ProgressBar {
	pb := &ProgressBar{
		Total:           total,
		RefreshRate:     DEFAULT_REFRESH_RATE,
		ShowPercent:     true,
		ShowCounters:    true,
		ShowBar:         true,
		ShowTimeLeft:    true,
		ShowElapsedTime: false,
		ShowFinalTime:   true,
		Units:           U_NO,
		ManualUpdate:    false,
		finish:          make(chan struct{}),
	}
	return pb.Format(FORMAT)
}

// Create new object and start
func StartNew(total int) *ProgressBar {
	return New(total).Start()
}

// Callback for custom output
// For example:
// bar.Callback = func(s string) {
//     mySuperPrint(s)
// }
//
type Callback func(out string)

type ProgressBar struct {
	current  int64 // current must be first member of struct (https://code.google.com/p/go/issues/detail?id=5278)
	previous int64

	Total                            int64
	RefreshRate                      time.Duration
	ShowPercent, ShowCounters        bool
	ShowSpeed, ShowTimeLeft, ShowBar bool
	ShowFinalTime, ShowElapsedTime   bool
	Output                           io.Writer
	Callback                         Callback
	NotPrint                         bool
	Units                            Units
	Width                            int
	ForceWidth                       bool
	ManualUpdate                     bool
	AutoStat                         bool

	// Default width for the time box.
	UnitsWidth   int
	TimeBoxWidth int

	finishOnce sync.Once //Guards isFinish
	finish     chan struct{}
	isFinish   bool

	startTime  time.Time
	startValue int64

	changeTime time.Time

	prefix, postfix string

	mu        sync.Mutex
	lastPrint string

	BarStart string
	BarEnd   string
	Empty    string
	Current  string
	CurrentN string

	AlwaysUpdate bool
}

// Start print
func (pb *ProgressBar) Start() *ProgressBar {
	pb.startTime = time.Now()
	pb.startValue = atomic.LoadInt64(&pb.current)
	if atomic.LoadInt64(&pb.Total) == 0 {
		pb.ShowTimeLeft = false
		pb.ShowPercent = false
		pb.AutoStat = false
	}
	if !pb.ManualUpdate {
		pb.Update() // Initial printing of the bar before running the bar refresher.
		go pb.refresher()
	}
	return pb
}

// Increment current value
func (pb *ProgressBar) Increment() int {
	return pb.Add(1)
}

// Get current value
func (pb *ProgressBar) Get() int64 {
	c := atomic.LoadInt64(&pb.current)
	return c
}

// Set current value
func (pb *ProgressBar) Set(current int) *ProgressBar {
	return pb.Set64(int64(current))
}

// Set64 sets the current value as int64
func (pb *ProgressBar) Set64(current int64) *ProgressBar {
	atomic.StoreInt64(&pb.current, current)
	return pb
}

// Add to current value
func (pb *ProgressBar) Add(add int) int {
	return int(pb.Add64(int64(add)))
}

func (pb *ProgressBar) Add64(add int64) int64 {
	return atomic.AddInt64(&pb.current, add)
}

// Set prefix string
func (pb *ProgressBar) Prefix(prefix string) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	pb.prefix = prefix
	return pb
}

// Set postfix string
func (pb *ProgressBar) Postfix(postfix string) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	pb.postfix = postfix
	return pb
}

// Set custom format for bar
// Example: bar.Format("[=>_]")
// Example: bar.Format("[\x00=\x00>\x00-\x00]") // \x00 is the delimiter
func (pb *ProgressBar) Format(format string) *ProgressBar {
	var formatEntries []string
	if utf8.RuneCountInString(format) == 5 {
		formatEntries = strings.Split(format, "")
	} else {
		formatEntries = strings.Split(format, "\x00")
	}
	if len(formatEntries) == 5 {
		pb.BarStart = formatEntries[0]
		pb.BarEnd = formatEntries[4]
		pb.Empty = formatEntries[3]
		pb.Current = formatEntries[1]
		pb.CurrentN = formatEntries[2]
	}
	return pb
}

// Set bar refresh rate
func (pb *ProgressBar) SetRefreshRate(rate time.Duration) *ProgressBar {
	pb.RefreshRate = rate
	return pb
}

// Set units
// bar.SetUnits(U_NO) - by default
// bar.SetUnits(U_BYTES) - for Mb, Kb, etc
func (pb *ProgressBar) SetUnits(units Units) *ProgressBar {
	pb.Units = units
	return pb
}

// Set max width, if width is bigger than terminal width, will be ignored
func (pb *ProgressBar) SetMaxWidth(width int) *ProgressBar {
	pb.Width = width
	pb.ForceWidth = false
	return pb
}

// Set bar width
func (pb *ProgressBar) SetWidth(width int) *ProgressBar {
	pb.Width = width
	pb.ForceWidth = true
	return pb
}

// End print
func (pb *ProgressBar) Finish() {
	//Protect multiple calls
	pb.finishOnce.Do(func() {
		close(pb.finish)
		pb.write(atomic.LoadInt64(&pb.Total), atomic.LoadInt64(&pb.current))
		pb.mu.Lock()
		defer pb.mu.Unlock()
		switch {
		case pb.Output != nil:
			fmt.Fprintln(pb.Output)
		case !pb.NotPrint:
			fmt.Println()
		}
		pb.isFinish = true
	})
}

// IsFinished return boolean
func (pb *ProgressBar) IsFinished() bool {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	return pb.isFinish
}

// End print and write string 'str'
func (pb *ProgressBar) FinishPrint(str string) {
	pb.Finish()
	if pb.Output != nil {
		fmt.Fprintln(pb.Output, str)
	} else {
		fmt.Println(str)
	}
}

// implement io.Writer
func (pb *ProgressBar) Write(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

// implement io.Reader
func (pb *ProgressBar) Read(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

// Create new proxy reader over bar
// Takes io.Reader or io.ReadCloser
func (pb *ProgressBar) NewProxyReader(r io.Reader) *Reader {
	return &Reader{r, pb}
}

func (pb *ProgressBar) write(total, current int64) {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	width := pb.GetWidth()

	var percentBox, countersBox, timeLeftBox, timeSpentBox, speedBox, barBox, end, out string

	// percents
	if pb.ShowPercent {
		var percent float64
		if total > 0 {
			percent = float64(current) / (float64(total) / float64(100))
		} else {
			percent = float64(current) / float64(100)
		}
		percentBox = fmt.Sprintf(" %6.02f%%", percent)
	}

	// counters
	if pb.ShowCounters {
		current := Format(current).To(pb.Units).Width(pb.UnitsWidth)
		if total > 0 {
			totalS := Format(total).To(pb.Units).Width(pb.UnitsWidth)
			countersBox = fmt.Sprintf(" %s / %s ", current, totalS)
		} else {
			countersBox = fmt.Sprintf(" %s / ? ", current)
		}
	}

	// time left
	currentFromStart := current - pb.startValue
	fromStart := time.Now().Sub(pb.startTime)
	lastChangeTime := pb.changeTime
	fromChange := lastChangeTime.Sub(pb.startTime)

	if pb.ShowElapsedTime {
		timeSpentBox = fmt.Sprintf(" %s ", (fromStart/time.Second)*time.Second)
	}

	select {
	case <-pb.finish:
		if pb.ShowFinalTime {
			var left time.Duration
			left = (fromStart / time.Second) * time.Second
			timeLeftBox = fmt.Sprintf(" %s", left.String())
		}
	default:
		if pb.ShowTimeLeft && currentFromStart > 0 {
			perEntry := fromChange / time.Duration(currentFromStart)
			var left time.Duration
			if total > 0 {
				left = time.Duration(total-currentFromStart) * perEntry
				left -= time.Since(lastChangeTime)
				left = (left / time.Second) * time.Second
			} else {
				left = time.Duration(currentFromStart) * perEntry
				left = (left / time.Second) * time.Second
			}
			if left > 0 {
				timeLeft := Format(int64(left)).To(U_DURATION).String()
				timeLeftBox = fmt.Sprintf(" %s", timeLeft)
			}
		}
	}

	if len(timeLeftBox) < pb.TimeBoxWidth {
		timeLeftBox = fmt.Sprintf("%s%s", strings.Repeat(" ", pb.TimeBoxWidth-len(timeLeftBox)), timeLeftBox)
	}

	// speed
	if pb.ShowSpeed && currentFromStart > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		speed := float64(currentFromStart) / (float64(fromStart) / float64(time.Second))
		speedBox = " " + Format(int64(speed)).To(pb.Units).Width(pb.UnitsWidth).PerSec().String()
	}

	barWidth := escapeAwareRuneCountInString(countersBox + pb.BarStart + pb.BarEnd + percentBox + timeSpentBox + timeLeftBox + speedBox + pb.prefix + pb.postfix)
	// bar
	if pb.ShowBar {
		size := width - barWidth
		if size > 0 {
			if total > 0 {
				curSize := int(math.Ceil((float64(current) / float64(total)) * float64(size)))
				emptySize := size - curSize
				barBox = pb.BarStart
				if emptySize < 0 {
					emptySize = 0
				}
				if curSize > size {
					curSize = size
				}

				cursorLen := escapeAwareRuneCountInString(pb.Current)
				if emptySize <= 0 {
					barBox += strings.Repeat(pb.Current, curSize/cursorLen)
				} else if curSize > 0 {
					cursorEndLen := escapeAwareRuneCountInString(pb.CurrentN)
					cursorRepetitions := (curSize - cursorEndLen) / cursorLen
					barBox += strings.Repeat(pb.Current, cursorRepetitions)
					barBox += pb.CurrentN
				}

				emptyLen := escapeAwareRuneCountInString(pb.Empty)
				barBox += strings.Repeat(pb.Empty, emptySize/emptyLen)
				barBox += pb.BarEnd
			} else {
				pos := size - int(current)%int(size)
				barBox = pb.BarStart
				if pos-1 > 0 {
					barBox += strings.Repeat(pb.Empty, pos-1)
				}
				barBox += pb.Current
				if size-pos-1 > 0 {
					barBox += strings.Repeat(pb.Empty, size-pos-1)
				}
				barBox += pb.BarEnd
			}
		}
	}

	// check len
	out = pb.prefix + timeSpentBox + countersBox + barBox + percentBox + speedBox + timeLeftBox + pb.postfix

	if cl := escapeAwareRuneCountInString(out); cl < width {
		end = strings.Repeat(" ", width-cl)
	}

	// and print!
	pb.lastPrint = out + end
	isFinish := pb.isFinish

	switch {
	case isFinish:
		return
	case pb.Output != nil:
		fmt.Fprint(pb.Output, "\r"+out+end)
	case pb.Callback != nil:
		pb.Callback(out + end)
	case !pb.NotPrint:
		fmt.Print("\r" + out + end)
	}
}

// GetTerminalWidth - returns terminal width for all platforms.
func GetTerminalWidth() (int, error) {
	return terminalWidth()
}

func (pb *ProgressBar) GetWidth() int {
	if pb.ForceWidth {
		return pb.Width
	}

	width := pb.Width
	termWidth, _ := terminalWidth()
	if width == 0 || termWidth <= width {
		width = termWidth
	}

	return width
}

// Write the current state of the progressbar
func (pb *ProgressBar) Update() {
	c := atomic.LoadInt64(&pb.current)
	p := atomic.LoadInt64(&pb.previous)
	t := atomic.LoadInt64(&pb.Total)
	if p != c {
		pb.mu.Lock()
		pb.changeTime = time.Now()
		pb.mu.Unlock()
		atomic.StoreInt64(&pb.previous, c)
	}
	pb.write(t, c)
	if pb.AutoStat {
		if c == 0 {
			pb.startTime = time.Now()
			pb.startValue = 0
		} else if c >= t && pb.isFinish != true {
			pb.Finish()
		}
	}
}

// String return the last bar print
func (pb *ProgressBar) String() string {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	return pb.lastPrint
}

// SetTotal atomically sets new total count
func (pb *ProgressBar) SetTotal(total int) *ProgressBar {
	return pb.SetTotal64(int64(total))
}

// SetTotal64 atomically sets new total count
func (pb *ProgressBar) SetTotal64(total int64) *ProgressBar {
	atomic.StoreInt64(&pb.Total, total)
	return pb
}

// Reset bar and set new total count
// Does effect only on finished bar
func (pb *ProgressBar) Reset(total int) *ProgressBar {
	pb.mu.Lock()
	defer pb.mu.Unlock()
	if pb.isFinish {
		pb.SetTotal(total).Set(0)
		atomic.StoreInt64(&pb.previous, 0)
	}
	return pb
}

// Internal loop for refreshing the progressbar
func (pb *ProgressBar) refresher() {
	for {
		select {
		case <-pb.finish:
			return
		case <-time.After(pb.RefreshRate):
			pb.Update()
		}
	}
}
<|endoftext|>


<commit_before>package sa

import (
	"testing"

	"github.com/letsencrypt/boulder/grpc"
	"github.com/letsencrypt/boulder/probs"

	"github.com/letsencrypt/boulder/core"
	corepb "github.com/letsencrypt/boulder/core/proto"
	"github.com/letsencrypt/boulder/test"
)

func TestModelToRegistrationNilContact(t *testing.T) {
	reg, err := modelToRegistration(&regModel{
		Key:     []byte(`{"kty":"RSA","n":"AQAB","e":"AQAB"}`),
		Contact: nil,
	})
	if err != nil {
		t.Errorf("Got error from modelToRegistration: %s", err)
	}
	if reg.Contact == nil {
		t.Errorf("Expected non-nil Contact field, got %#v", reg.Contact)
	}
	if len(*reg.Contact) != 0 {
		t.Errorf("Expected empty Contact field, got %#v", reg.Contact)
	}
}

// TestModelToRegistrationBadJSON tests that converting a model with an invalid
// JWK JSON produces the expected bad JSON error.
func TestModelToRegistrationBadJSON(t *testing.T) {
	badJSON := []byte(`{`)
	_, err := modelToRegistration(&regModel{
		Key: badJSON,
	})
	test.AssertError(t, err, "expected error from truncated reg model key")
	var badJSONErr errBadJSON
	test.AssertErrorWraps(t, err, &badJSONErr)
	test.AssertEquals(t, string(badJSONErr.json), string(badJSON))
}

func TestModelToRegistrationNonNilContact(t *testing.T) {
	reg, err := modelToRegistration(&regModel{
		Key:     []byte(`{"kty":"RSA","n":"AQAB","e":"AQAB"}`),
		Contact: []string{},
	})
	if err != nil {
		t.Errorf("Got error from modelToRegistration: %s", err)
	}
	if reg.Contact == nil {
		t.Errorf("Expected non-nil Contact field, got %#v", reg.Contact)
	}
	if len(*reg.Contact) != 0 {
		t.Errorf("Expected empty Contact field, got %#v", reg.Contact)
	}
}

func TestAuthzModel(t *testing.T) {
	authzPB := &corepb.Authorization{
		Id:             "1",
		Identifier:     "example.com",
		RegistrationID: 1,
		Status:         string(core.StatusValid),
		Expires:        1234,
		Challenges: []*corepb.Challenge{
			{
				Type:      string(core.ChallengeTypeHTTP01),
				Status:    string(core.StatusValid),
				Token:     "MTIz",
				Validated: 1234,
				Validationrecords: []*corepb.ValidationRecord{
					{
						Hostname:          "hostname",
						Port:              "port",
						AddressUsed:       []byte("1.2.3.4"),
						Url:               "url",
						AddressesResolved: [][]byte{{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4}},
						AddressesTried:    [][]byte{{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4}},
					},
				},
			},
		},
	}

	model, err := authzPBToModel(authzPB)
	test.AssertNotError(t, err, "authzPBToModel failed")

	authzPBOut, err := modelToAuthzPB(*model)
	test.AssertNotError(t, err, "modelToAuthzPB failed")
	test.AssertDeepEquals(t, authzPB.Challenges, authzPBOut.Challenges)

	validationErr := probs.ConnectionFailure("weewoo")
	authzPB.Challenges[0].Status = string(core.StatusInvalid)
	authzPB.Challenges[0].Error, err = grpc.ProblemDetailsToPB(validationErr)
	test.AssertNotError(t, err, "grpc.ProblemDetailsToPB failed")
	model, err = authzPBToModel(authzPB)
	test.AssertNotError(t, err, "authzPBToModel failed")

	authzPBOut, err = modelToAuthzPB(*model)
	test.AssertNotError(t, err, "modelToAuthzPB failed")
	test.AssertDeepEquals(t, authzPB.Challenges, authzPBOut.Challenges)

	authzPB = &corepb.Authorization{
		Id:             "1",
		Identifier:     "example.com",
		RegistrationID: 1,
		Status:         string(core.StatusInvalid),
		Expires:        1234,
		Challenges: []*corepb.Challenge{
			{
				Type:   string(core.ChallengeTypeHTTP01),
				Status: string(core.StatusInvalid),
				Token:  "MTIz",
				Validationrecords: []*corepb.ValidationRecord{
					{
						Hostname:          "hostname",
						Port:              "port",
						AddressUsed:       []byte("1.2.3.4"),
						Url:               "url",
						AddressesResolved: [][]byte{{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4}},
						AddressesTried:    [][]byte{{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4}},
					},
				},
			},
			{
				Type:   string(core.ChallengeTypeDNS01),
				Status: string(core.StatusInvalid),
				Token:  "MTIz",
				Validationrecords: []*corepb.ValidationRecord{
					{
						Hostname:          "hostname",
						Port:              "port",
						AddressUsed:       []byte("1.2.3.4"),
						Url:               "url",
						AddressesResolved: [][]byte{{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4}},
						AddressesTried:    [][]byte{{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4}},
					},
				},
			},
		},
	}
	_, err = authzPBToModel(authzPB)
	test.AssertError(t, err, "authzPBToModel didn't fail with multiple non-pending challenges")
}

// TestModelToChallengeBadJSON tests that converting a challenge model with an
// invalid validation error field or validation record field produces the
// expected bad JSON error.
func TestModelToChallengeBadJSON(t *testing.T) {
	badJSON := []byte(`{`)

	testCases := []struct {
		Name  string
		Model *challModel
	}{
		{
			Name: "Bad error field",
			Model: &challModel{
				Error: badJSON,
			},
		},
		{
			Name: "Bad validation record field",
			Model: &challModel{
				ValidationRecord: badJSON,
			},
		},
	}
	for _, tc := range testCases {
		t.Run(tc.Name, func(t *testing.T) {
			_, err := modelToChallenge(tc.Model)
			test.AssertError(t, err, "expected error from modelToChallenge")
			var badJSONErr errBadJSON
			test.AssertErrorWraps(t, err, &badJSONErr)
			test.AssertEquals(t, string(badJSONErr.json), string(badJSON))
		})
	}
}

// TestModelToOrderBADJSON tests that converting an order model with an invalid
// validation error JSON field to an Order produces the expected bad JSON error.
func TestModelToOrderBadJSON(t *testing.T) {
	badJSON := []byte(`{`)
	_, err := modelToOrder(&orderModel{
		Error: badJSON,
	})
	test.AssertError(t, err, "expected error from modelToOrder")
	var badJSONErr errBadJSON
	test.AssertErrorWraps(t, err, &badJSONErr)
	test.AssertEquals(t, string(badJSONErr.json), string(badJSON))
}

// TestPopulateAttemptedFieldsBadJSON tests that populating a challenge from an
// authz2 model with an invalid validation error or an invalid validation record
// produces the expected bad JSON error.
func TestPopulateAttemptedFieldsBadJSON(t *testing.T) {
	badJSON := []byte(`{`)

	testCases := []struct {
		Name  string
		Model *authzModel
	}{
		{
			Name: "Bad validation error field",
			Model: &authzModel{
				ValidationError: badJSON,
			},
		},
		{
			Name: "Bad validation record field",
			Model: &authzModel{
				ValidationRecord: badJSON,
			},
		},
	}
	for _, tc := range testCases {
		t.Run(tc.Name, func(t *testing.T) {
			err := populateAttemptedFields(*tc.Model, &corepb.Challenge{})
			test.AssertError(t, err, "expected error from populateAttemptedFields")
			var badJSONErr errBadJSON
			test.AssertErrorWraps(t, err, &badJSONErr)
			test.AssertEquals(t, string(badJSONErr.json), string(badJSON))
		})
	}
}
<commit_msg>SA: Add unit tests for SelectCertificate (#5503)<commit_after>package sa

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"math/big"
	"testing"
	"time"

	"github.com/jmhodges/clock"
	"github.com/letsencrypt/boulder/db"
	"github.com/letsencrypt/boulder/grpc"
	"github.com/letsencrypt/boulder/probs"

	"github.com/letsencrypt/boulder/core"
	corepb "github.com/letsencrypt/boulder/core/proto"
	"github.com/letsencrypt/boulder/test"
)

func TestModelToRegistrationNilContact(t *testing.T) {
	reg, err := modelToRegistration(&regModel{
		Key:     []byte(`{"kty":"RSA","n":"AQAB","e":"AQAB"}`),
		Contact: nil,
	})
	if err != nil {
		t.Errorf("Got error from modelToRegistration: %s", err)
	}
	if reg.Contact == nil {
		t.Errorf("Expected non-nil Contact field, got %#v", reg.Contact)
	}
	if len(*reg.Contact) != 0 {
		t.Errorf("Expected empty Contact field, got %#v", reg.Contact)
	}
}

// TestModelToRegistrationBadJSON tests that converting a model with an invalid
// JWK JSON produces the expected bad JSON error.
func TestModelToRegistrationBadJSON(t *testing.T) {
	badJSON := []byte(`{`)
	_, err := modelToRegistration(&regModel{
		Key: badJSON,
	})
	test.AssertError(t, err, "expected error from truncated reg model key")
	var badJSONErr errBadJSON
	test.AssertErrorWraps(t, err, &badJSONErr)
	test.AssertEquals(t, string(badJSONErr.json), string(badJSON))
}

func TestModelToRegistrationNonNilContact(t *testing.T) {
	reg, err := modelToRegistration(&regModel{
		Key:     []byte(`{"kty":"RSA","n":"AQAB","e":"AQAB"}`),
		Contact: []string{},
	})
	if err != nil {
		t.Errorf("Got error from modelToRegistration: %s", err)
	}
	if reg.Contact == nil {
		t.Errorf("Expected non-nil Contact field, got %#v", reg.Contact)
	}
	if len(*reg.Contact) != 0 {
		t.Errorf("Expected empty Contact field, got %#v", reg.Contact)
	}
}

func TestAuthzModel(t *testing.T) {
	authzPB := &corepb.Authorization{
		Id:             "1",
		Identifier:     "example.com",
		RegistrationID: 1,
		Status:         string(core.StatusValid),
		Expires:        1234,
		Challenges: []*corepb.Challenge{
			{
				Type:      string(core.ChallengeTypeHTTP01),
				Status:    string(core.StatusValid),
				Token:     "MTIz",
				Validated: 1234,
				Validationrecords: []*corepb.ValidationRecord{
					{
						Hostname:          "hostname",
						Port:              "port",
						AddressUsed:       []byte("1.2.3.4"),
						Url:               "url",
						AddressesResolved: [][]byte{{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4}},
						AddressesTried:    [][]byte{{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4}},
					},
				},
			},
		},
	}

	model, err := authzPBToModel(authzPB)
	test.AssertNotError(t, err, "authzPBToModel failed")

	authzPBOut, err := modelToAuthzPB(*model)
	test.AssertNotError(t, err, "modelToAuthzPB failed")
	test.AssertDeepEquals(t, authzPB.Challenges, authzPBOut.Challenges)

	validationErr := probs.ConnectionFailure("weewoo")
	authzPB.Challenges[0].Status = string(core.StatusInvalid)
	authzPB.Challenges[0].Error, err = grpc.ProblemDetailsToPB(validationErr)
	test.AssertNotError(t, err, "grpc.ProblemDetailsToPB failed")
	model, err = authzPBToModel(authzPB)
	test.AssertNotError(t, err, "authzPBToModel failed")

	authzPBOut, err = modelToAuthzPB(*model)
	test.AssertNotError(t, err, "modelToAuthzPB failed")
	test.AssertDeepEquals(t, authzPB.Challenges, authzPBOut.Challenges)

	authzPB = &corepb.Authorization{
		Id:             "1",
		Identifier:     "example.com",
		RegistrationID: 1,
		Status:         string(core.StatusInvalid),
		Expires:        1234,
		Challenges: []*corepb.Challenge{
			{
				Type:   string(core.ChallengeTypeHTTP01),
				Status: string(core.StatusInvalid),
				Token:  "MTIz",
				Validationrecords: []*corepb.ValidationRecord{
					{
						Hostname:          "hostname",
						Port:              "port",
						AddressUsed:       []byte("1.2.3.4"),
						Url:               "url",
						AddressesResolved: [][]byte{{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4}},
						AddressesTried:    [][]byte{{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4}},
					},
				},
			},
			{
				Type:   string(core.ChallengeTypeDNS01),
				Status: string(core.StatusInvalid),
				Token:  "MTIz",
				Validationrecords: []*corepb.ValidationRecord{
					{
						Hostname:          "hostname",
						Port:              "port",
						AddressUsed:       []byte("1.2.3.4"),
						Url:               "url",
						AddressesResolved: [][]byte{{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4}},
						AddressesTried:    [][]byte{{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4}},
					},
				},
			},
		},
	}
	_, err = authzPBToModel(authzPB)
	test.AssertError(t, err, "authzPBToModel didn't fail with multiple non-pending challenges")
}

// TestModelToChallengeBadJSON tests that converting a challenge model with an
// invalid validation error field or validation record field produces the
// expected bad JSON error.
func TestModelToChallengeBadJSON(t *testing.T) {
	badJSON := []byte(`{`)

	testCases := []struct {
		Name  string
		Model *challModel
	}{
		{
			Name: "Bad error field",
			Model: &challModel{
				Error: badJSON,
			},
		},
		{
			Name: "Bad validation record field",
			Model: &challModel{
				ValidationRecord: badJSON,
			},
		},
	}
	for _, tc := range testCases {
		t.Run(tc.Name, func(t *testing.T) {
			_, err := modelToChallenge(tc.Model)
			test.AssertError(t, err, "expected error from modelToChallenge")
			var badJSONErr errBadJSON
			test.AssertErrorWraps(t, err, &badJSONErr)
			test.AssertEquals(t, string(badJSONErr.json), string(badJSON))
		})
	}
}

// TestModelToOrderBADJSON tests that converting an order model with an invalid
// validation error JSON field to an Order produces the expected bad JSON error.
func TestModelToOrderBadJSON(t *testing.T) {
	badJSON := []byte(`{`)
	_, err := modelToOrder(&orderModel{
		Error: badJSON,
	})
	test.AssertError(t, err, "expected error from modelToOrder")
	var badJSONErr errBadJSON
	test.AssertErrorWraps(t, err, &badJSONErr)
	test.AssertEquals(t, string(badJSONErr.json), string(badJSON))
}

// TestPopulateAttemptedFieldsBadJSON tests that populating a challenge from an
// authz2 model with an invalid validation error or an invalid validation record
// produces the expected bad JSON error.
func TestPopulateAttemptedFieldsBadJSON(t *testing.T) {
	badJSON := []byte(`{`)

	testCases := []struct {
		Name  string
		Model *authzModel
	}{
		{
			Name: "Bad validation error field",
			Model: &authzModel{
				ValidationError: badJSON,
			},
		},
		{
			Name: "Bad validation record field",
			Model: &authzModel{
				ValidationRecord: badJSON,
			},
		},
	}
	for _, tc := range testCases {
		t.Run(tc.Name, func(t *testing.T) {
			err := populateAttemptedFields(*tc.Model, &corepb.Challenge{})
			test.AssertError(t, err, "expected error from populateAttemptedFields")
			var badJSONErr errBadJSON
			test.AssertErrorWraps(t, err, &badJSONErr)
			test.AssertEquals(t, string(badJSONErr.json), string(badJSON))
		})
	}
}

func TestCerficatesTableContainsDuplicateSerials(t *testing.T) {
	sa, fc, cleanUp := initSA(t)
	defer cleanUp()

	serialString := core.SerialToString(big.NewInt(1337))

	// Insert a certificate with a serial of `1337`.
	err := insertCertificate(sa.dbMap, fc, "1337.com", "leet", 1337, 1)
	test.AssertNotError(t, err, "couldn't insert valid certificate")

	// This should return the certificate that we just inserted.
	_, err = SelectCertificate(sa.dbMap, serialString)
	test.AssertNotError(t, err, "received an error for a valid query")

	// Insert a certificate with a serial of `1337` but for a different
	// hostname.
	err = insertCertificate(sa.dbMap, fc, "1337.net", "leet", 1337, 1)
	test.AssertNotError(t, err, "couldn't insert valid certificate")

	// With a duplicate being present, this should error.
	_, err = SelectCertificate(sa.dbMap, serialString)
	test.AssertError(t, err, "should've received an error for multiple rows")
}

func insertCertificate(dbMap *db.WrappedMap, fc clock.FakeClock, hostname, cn string, serial, regID int64) error {
	serialBigInt := big.NewInt(serial)
	serialString := core.SerialToString(serialBigInt)

	template := x509.Certificate{
		Subject: pkix.Name{
			CommonName: cn,
		},
		NotAfter:     fc.Now().Add(30 * 24 * time.Hour),
		DNSNames:     []string{hostname},
		SerialNumber: serialBigInt,
	}

	testKey := makeKey()
	certDer, _ := x509.CreateCertificate(rand.Reader, &template, &template, &testKey.PublicKey, &testKey)
	cert := &core.Certificate{
		RegistrationID: regID,
		Serial:         serialString,
		Expires:        template.NotAfter,
		DER:            certDer,
	}
	err := dbMap.Insert(cert)
	if err != nil {
		return err
	}
	return nil
}

func bigIntFromB64(b64 string) *big.Int {
	bytes, _ := base64.URLEncoding.DecodeString(b64)
	x := big.NewInt(0)
	x.SetBytes(bytes)
	return x
}

func makeKey() rsa.PrivateKey {
	n := bigIntFromB64("n4EPtAOCc9AlkeQHPzHStgAbgs7bTZLwUBZdR8_KuKPEHLd4rHVTeT-O-XV2jRojdNhxJWTDvNd7nqQ0VEiZQHz_AJmSCpMaJMRBSFKrKb2wqVwGU_NsYOYL-QtiWN2lbzcEe6XC0dApr5ydQLrHqkHHig3RBordaZ6Aj-oBHqFEHYpPe7Tpe-OfVfHd1E6cS6M1FZcD1NNLYD5lFHpPI9bTwJlsde3uhGqC0ZCuEHg8lhzwOHrtIQbS0FVbb9k3-tVTU4fg_3L_vniUFAKwuCLqKnS2BYwdq_mzSnbLY7h_qixoR7jig3__kRhuaxwUkRz5iaiQkqgc5gHdrNP5zw==")
	e := int(bigIntFromB64("AQAB").Int64())
	d := bigIntFromB64("bWUC9B-EFRIo8kpGfh0ZuyGPvMNKvYWNtB_ikiH9k20eT-O1q_I78eiZkpXxXQ0UTEs2LsNRS-8uJbvQ-A1irkwMSMkK1J3XTGgdrhCku9gRldY7sNA_AKZGh-Q661_42rINLRCe8W-nZ34ui_qOfkLnK9QWDDqpaIsA-bMwWWSDFu2MUBYwkHTMEzLYGqOe04noqeq1hExBTHBOBdkMXiuFhUq1BU6l-DqEiWxqg82sXt2h-LMnT3046AOYJoRioz75tSUQfGCshWTBnP5uDjd18kKhyv07lhfSJdrPdM5Plyl21hsFf4L_mHCuoFau7gdsPfHPxxjVOcOpBrQzwQ==")
	p := bigIntFromB64("uKE2dh-cTf6ERF4k4e_jy78GfPYUIaUyoSSJuBzp3Cubk3OCqs6grT8bR_cu0Dm1MZwWmtdqDyI95HrUeq3MP15vMMON8lHTeZu2lmKvwqW7anV5UzhM1iZ7z4yMkuUwFWoBvyY898EXvRD-hdqRxHlSqAZ192zB3pVFJ0s7pFc=")
	q := bigIntFromB64("uKE2dh-cTf6ERF4k4e_jy78GfPYUIaUyoSSJuBzp3Cubk3OCqs6grT8bR_cu0Dm1MZwWmtdqDyI95HrUeq3MP15vMMON8lHTeZu2lmKvwqW7anV5UzhM1iZ7z4yMkuUwFWoBvyY898EXvRD-hdqRxHlSqAZ192zB3pVFJ0s7pFc=")
	return rsa.PrivateKey{PublicKey: rsa.PublicKey{N: n, E: e}, D: d, Primes: []*big.Int{p, q}}
}
<|endoftext|>


<commit_before>package s3

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"time"
)

// Represents AWS credentials and config.
type Credentials struct {
	Region string
	Bucket string
	AccessKeyID string
	SecretAccessKey string
}

// Represents presigned POST information.
type PresignedPOST struct {
	Key string         `json:"key"`
	Policy string      `json:"policy"`
	Signature string   `json:"signature"`
	Action string      `json:"action"`
	Credential string  `json:"credential"`
	Date string        `json:"date"`
}

// Creates a new presigned POST.
func NewPresignedPOST(key string, c *Credentials) (*PresignedPOST, error) {
	p := NewPolicy(key, c)
	b64Policy := p.Base64()
	signature := createSignature(p.C, p.Date[:8], b64Policy)
	action := fmt.Sprintf("https://%s.s3.amazonaws.com/", p.Bucket)
	post := &PresignedPOST{
		Key: p.Key,
		Policy: b64Policy,
		Signature: signature,
		Action: action,
		Credential: p.Credential,
		Date: p.Date,
	}
	return post, nil
}

// Creates the signature for a string.
func createSignature(c *Credentials, formattedShortTime, stringToSign string) string {
	h1 := makeHmac([]byte("AWS4"+c.SecretAccessKey), []byte(formattedShortTime))
	h2 := makeHmac(h1, []byte(c.Region))
	h3 := makeHmac(h2, []byte("s3"))
	h4 := makeHmac(h3, []byte("aws4_request"))
	signature := makeHmac(h4, []byte(stringToSign))
	return hex.EncodeToString(signature)
}

// Helper to make the HMAC-SHA256.
func makeHmac(key []byte, data []byte) []byte {
	hash := hmac.New(sha256.New, key)
	hash.Write(data)
	return hash.Sum(nil)
}

// Policy template.
const policyDocument = `
{ "expiration": "%s",
  "conditions": [
    {"bucket": "%s"},
    ["starts-with", "$key", "%s"],
    {"acl": "public-read"},
    ["starts-with", "$x-amz-meta-tag", ""],

    {"x-amz-credential": "%s"},
    {"x-amz-algorithm": "AWS4-HMAC-SHA256"},
    {"x-amz-date": "%s" }
  ]
}
`

const (
	expirationFormat = "2006-01-02T15:04:05.000Z"
	timeFormat = "20060102T150405Z"
	shortTimeFormat = "20060102"
)

// Represents a new policy for uploading sounds.
type policy struct {
	Expiration string
	Region string
	Bucket string
	Key string
	Credential string
	Date string
	C *Credentials
}

// Creates a new policy.
func NewPolicy(key string, c *Credentials) *policy {
	// expires in 5 minutes
	t := time.Now().Add(time.Minute * 5)
	formattedShortTime := t.UTC().Format(shortTimeFormat)
	date := t.UTC().Format(timeFormat)
	cred := fmt.Sprintf("%s/%s/%s/s3/aws4_request", c.AccessKeyID, formattedShortTime, c.Region)
	return &policy{
		Expiration: t.UTC().Format(expirationFormat),
		Region: c.Region,
		Bucket: c.Bucket,
		Key: key,
		Credential: cred,
		Date: date,
		C: c,
	}
}

// Returns the policy as a string.
func (p *policy) String() string {
	return fmt.Sprintf(policyDocument,
		p.Expiration,
		p.Bucket,
		p.Key,
		p.Credential,
		p.Date,
	)
}

// Returns the policy as a base64 encoded string.
func (p *policy) Base64() string {
	return base64.StdEncoding.EncodeToString([]byte(p.String()))
}
<commit_msg>Adds a redirect URL as well as max content size to about 0.5MB.<commit_after>package s3

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"time"
)

// Represents AWS credentials and config.
type Credentials struct {
	Region string
	Bucket string
	AccessKeyID string
	SecretAccessKey string
	RedirectURL string
}

// Represents presigned POST information.
type PresignedPOST struct {
	Key string         `json:"key"`
	Policy string      `json:"policy"`
	Signature string   `json:"signature"`
	Action string      `json:"action"`
	RedirectURL string `json:"redirect"`
	Credential string  `json:"credential"`
	Date string        `json:"date"`
}

// Creates a new presigned POST.
func NewPresignedPOST(key string, c *Credentials) (*PresignedPOST, error) {
	p := NewPolicy(key, c)
	b64Policy := p.Base64()
	signature := createSignature(p.C, p.Date[:8], b64Policy)
	action := fmt.Sprintf("https://%s.s3.amazonaws.com/", p.Bucket)
	post := &PresignedPOST{
		Key: p.Key,
		Policy: b64Policy,
		Signature: signature,
		Action: action,
		RedirectURL: p.RedirectURL,
		Credential: p.Credential,
		Date: p.Date,
	}
	return post, nil
}

// Creates the signature for a string.
func createSignature(c *Credentials, formattedShortTime, stringToSign string) string {
	h1 := makeHmac([]byte("AWS4"+c.SecretAccessKey), []byte(formattedShortTime))
	h2 := makeHmac(h1, []byte(c.Region))
	h3 := makeHmac(h2, []byte("s3"))
	h4 := makeHmac(h3, []byte("aws4_request"))
	signature := makeHmac(h4, []byte(stringToSign))
	return hex.EncodeToString(signature)
}

// Helper to make the HMAC-SHA256.
func makeHmac(key []byte, data []byte) []byte {
	hash := hmac.New(sha256.New, key)
	hash.Write(data)
	return hash.Sum(nil)
}

// Policy template.
const policyDocument = `
{ "expiration": "%s",
  "conditions": [
    {"bucket": "%s"},
    ["starts-with", "$key", "%s"],
    {"acl": "public-read"},
    {"success_action_redirect": "%s"},
    ["starts-with", "$x-amz-meta-tag", ""],
    ["content-length-range", 1, 524288],

    {"x-amz-credential": "%s"},
    {"x-amz-algorithm": "AWS4-HMAC-SHA256"},
    {"x-amz-date": "%s" }
  ]
}
`

const (
	expirationFormat = "2006-01-02T15:04:05.000Z"
	timeFormat = "20060102T150405Z"
	shortTimeFormat = "20060102"
)

// Represents a new policy for uploading sounds.
type policy struct {
	Expiration string
	Region string
	Bucket string
	Key string
	RedirectURL string
	Credential string
	Date string
	C *Credentials
}

// Creates a new policy.
func NewPolicy(key string, c *Credentials) *policy {
	// expires in 5 minutes
	t := time.Now().Add(time.Minute * 5)
	url := fmt.Sprintf("%s/%s", c.RedirectURL, key)
	formattedShortTime := t.UTC().Format(shortTimeFormat)
	date := t.UTC().Format(timeFormat)
	cred := fmt.Sprintf("%s/%s/%s/s3/aws4_request", c.AccessKeyID, formattedShortTime, c.Region)
	return &policy{
		Expiration: t.UTC().Format(expirationFormat),
		Region: c.Region,
		Bucket: c.Bucket,
		Key: key,
		RedirectURL: url,
		Credential: cred,
		Date: date,
		C: c,
	}
}

// Returns the policy as a string.
func (p *policy) String() string {
	return fmt.Sprintf(policyDocument,
		p.Expiration,
		p.Bucket,
		p.Key,
		p.RedirectURL,
		p.Credential,
		p.Date,
	)
}

// Returns the policy as a base64 encoded string.
func (p *policy) Base64() string {
	return base64.StdEncoding.EncodeToString([]byte(p.String()))
}
<|endoftext|>


<commit_before>package s3

// https://github.com/aws/aws-sdk-go
// https://docs.aws.amazon.com/sdk-for-go/api/service/s3.html

// https://github.com/goamz/goamz/blob/master/aws/aws.go
// https://github.com/goamz/goamz/blob/master/s3/s3.go

import (
	"bufio"
	"fmt"
	"github.com/goamz/goamz/aws"
	aws_s3 "github.com/goamz/goamz/s3"
	"github.com/jeffail/tunny"
	"github.com/whosonfirst/go-whosonfirst-crawl"
	log "github.com/whosonfirst/go-whosonfirst-log"
	pool "github.com/whosonfirst/go-whosonfirst-pool"
	utils "github.com/whosonfirst/go-whosonfirst-utils"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

type Sync struct {
	ACL           aws_s3.ACL
	Bucket        aws_s3.Bucket
	Prefix        string
	WorkPool      tunny.WorkPool
	Logger        *log.WOFLogger
	Debug         bool
	Success       int64
	Error         int64
	Skipped       int64
	Scheduled     int64
	Completed     int64
	Retried       int64
	TimeToProcess *time.Duration
	Retries       *pool.LIFOPool
	MaxRetries    float64 // max percentage of errors over scheduled
}

func NewSync(auth aws.Auth, region aws.Region, acl aws_s3.ACL, bucket string, prefix string, procs int, debug bool, logger *log.WOFLogger) *Sync {

	logger.Info("creating a new Sync thing-y with %d processes", procs)

	runtime.GOMAXPROCS(procs)

	workpool, _ := tunny.CreatePoolGeneric(procs).Open()

	retries := pool.NewLIFOPool()

	s := aws_s3.New(auth, region)
	b := s.Bucket(bucket)

	ttp := new(time.Duration)

	return &Sync{
		ACL:           acl,
		Bucket:        *b,
		Prefix:        prefix,
		WorkPool:      *workpool,
		Debug:         debug,
		Logger:        logger,
		Scheduled:     0,
		Completed:     0,
		Skipped:       0,
		Error:         0,
		Success:       0,
		Retried:       0,
		TimeToProcess: ttp,
		Retries:       retries,
		MaxRetries:    25.0, // maybe allow this to be user-defined ?
	}
}

func WOFSync(auth aws.Auth, bucket string, prefix string, procs int, debug bool, logger *log.WOFLogger) *Sync {

	return NewSync(auth, aws.USEast, aws_s3.PublicRead, bucket, prefix, procs, debug, logger)
}

func (sink *Sync) SyncDirectory(root string) error {

	defer sink.WorkPool.Close()

	t0 := time.Now()

	wg := new(sync.WaitGroup)

	callback := func(source string, info os.FileInfo) error {

		wg.Add(1)

		if info.IsDir() {
			return nil
		}

		err := sink.SyncFile(source, root, wg)

		if err != nil {
			sink.Logger.Error("failed to sync %s, because '%s'", source, err)
		}

		return nil
	}

	c := crawl.NewCrawler(root)
	_ = c.Crawl(callback)

	wg.Wait()

	sink.ProcessRetries(root)

	ttp := time.Since(t0)
	sink.TimeToProcess = &ttp

	return nil
}

func (sink *Sync) SyncFiles(files []string, root string) error {

	defer sink.WorkPool.Close()

	t0 := time.Now()

	wg := new(sync.WaitGroup)

	for _, path := range files {

		wg.Add(1)

		go func(path string, root string, wg *sync.WaitGroup) {
			sink.SyncFile(path, root, wg)
		}(path, root, wg)
	}

	wg.Wait()

	sink.ProcessRetries(root)

	ttp := time.Since(t0)
	sink.TimeToProcess = &ttp

	return nil
}

func (sink *Sync) SyncFileList(path string, root string) error {

	defer sink.WorkPool.Close()

	t0 := time.Now()

	file, err := os.Open(path)

	if err != nil {
		return err
	}

	defer file.Close()

	scanner := bufio.NewScanner(file)

	wg := new(sync.WaitGroup)

	for scanner.Scan() {

		path := scanner.Text()

		wg.Add(1)

		go func(path string, root string, wg *sync.WaitGroup) {
			sink.SyncFile(path, root, wg)
		}(path, root, wg)
	}

	wg.Wait()

	sink.ProcessRetries(root)

	ttp := time.Since(t0)
	sink.TimeToProcess = &ttp

	return nil
}

func (sink *Sync) SyncFile(source string, root string, wg *sync.WaitGroup) error {

	atomic.AddInt64(&sink.Scheduled, 1)

	_, err := sink.WorkPool.SendWork(func() {

		defer wg.Done()

		dest := source

		dest = strings.Replace(dest, root, "", -1)

		if sink.Prefix != "" {
			dest = path.Join(sink.Prefix, dest)
		}

		// Note: both HasChanged and SyncFile will ioutil.ReadFile(source)
		// which is a potential waste of time and resource. Or maybe we just
		// don't care? (20150930/thisisaaronland)

		sink.Logger.Debug("Looking for changes to %s (%s)", dest, sink.Prefix)

		change, ch_err := sink.HasChanged(source, dest)

		if ch_err != nil {
			atomic.AddInt64(&sink.Completed, 1)
			atomic.AddInt64(&sink.Error, 1)
			sink.Logger.Warning("failed to determine whether %s had changed, because '%s'", source, ch_err)

			sink.Retries.Push(&pool.PoolString{String: source})
			return
		}

		if sink.Debug == true {
			atomic.AddInt64(&sink.Completed, 1)
			atomic.AddInt64(&sink.Skipped, 1)
			sink.Logger.Debug("has %s changed? the answer is %v but does it really matter since debugging is enabled?", source, change)
			return
		}

		if !change {
			atomic.AddInt64(&sink.Completed, 1)
			atomic.AddInt64(&sink.Skipped, 1)
			sink.Logger.Debug("%s has not changed, skipping", source)
			return
		}

		err := sink.DoSyncFile(source, dest)

		if err != nil {
			sink.Retries.Push(&pool.PoolString{String: source})
			atomic.AddInt64(&sink.Error, 1)
		} else {
			atomic.AddInt64(&sink.Success, 1)
		}

		atomic.AddInt64(&sink.Completed, 1)
	})

	if err != nil {
		wg.Done()
		atomic.AddInt64(&sink.Error, 1)
		sink.Logger.Error("failed to schedule %s for processing, because %v", source, err)
		return err
	}

	sink.Logger.Debug("schedule %s for processing", source)
	return nil
}

func (sink *Sync) DoSyncFile(source string, dest string) error {

	sink.Logger.Debug("prepare %s for syncing", source)

	body, err := ioutil.ReadFile(source)

	if err != nil {
		sink.Logger.Error("Failed to read %s, because %v", source, err)
		return err
	}

	sink.Logger.Debug("PUT %s as %s", dest, sink.ACL)

	o := aws_s3.Options{}

	err = sink.Bucket.Put(dest, body, "text/plain", sink.ACL, o)

	if err != nil {
		sink.Logger.Error("failed to PUT %s, because '%s'", dest, err)
		return err
	}

	return nil
}

func (sink *Sync) HasChanged(source string, dest string) (ch bool, err error) {

	headers := make(http.Header)
	rsp, err := sink.Bucket.Head(dest, headers)

	if err != nil {

		if e, ok := err.(*aws_s3.Error); ok && e.StatusCode == 404 {
			sink.Logger.Debug("%s is 404 so assuming it has changed (WHOA)", dest)
			return true, nil
		}

		sink.Logger.Error("failed to HEAD %s because %s", dest, err)
		return false, err
	}

	local_hash, err := utils.HashFile(source)

	if err != nil {
		sink.Logger.Warning("failed to hash %s, because %v", source, err)
		return false, err
	}

	etag := rsp.Header.Get("Etag")
	remote_hash := strings.Replace(etag, "\"", "", -1)

	if local_hash == remote_hash {
		return false, nil
	}

	// Okay so we think that things have changed but let's just check
	// modification times to be extra sure (20151112/thisisaaronland)

	info, err := os.Stat(source)

	if err != nil {
		sink.Logger.Error("failed to stat %s because %s", source, err)
		return false, err
	}

	mtime_local := info.ModTime()

	last_mod := rsp.Header.Get("Last-Modified")
	mtime_remote, err := time.Parse(time.RFC1123, last_mod)

	if err != nil {
		sink.Logger.Error("failed to parse timestamp %s because %s", last_mod, err)
		return false, err
	}

	// Because who remembers this stuff anyway...
	// func (t Time) Before(u Time) bool
	// Before reports whether the time instant t is before u.

	sink.Logger.Debug("local %s %s", mtime_local, source)
	sink.Logger.Debug("remote %s %s", mtime_remote, dest)

	if mtime_local.Before(mtime_remote) {
		sink.Logger.Warning("remote copy of %s has a more recent modification date (local: %s remote: %s)", source, mtime_local, mtime_remote)
		return false, nil
	}

	return true, nil
}

func (sink *Sync) ProcessRetries(root string) bool {

	to_retry := sink.Retries.Length()

	if to_retry > 0 {

		scheduled_f := float64(sink.Scheduled)
		retry_f := float64(to_retry)

		pct := (retry_f / scheduled_f) * 100.0

		if pct > sink.MaxRetries {
			sink.Logger.Warning("E_EXCESSIVE_ERRORS, %f percent of scheduled processes failed thus undermining our faith that they will work now...", pct)
			return false
		}

		sink.Logger.Info("There are %d failed requests that will now be retried", to_retry)

		wg := new(sync.WaitGroup)

		for sink.Retries.Length() > 0 {

			r, ok := sink.Retries.Pop()

			if !ok {
				sink.Logger.Error("failed to pop retries because... computers?")
				break
			}

			source := r.StringValue()

			wg.Add(1)

			go func(source string, root string, wg *sync.WaitGroup) {

				atomic.AddInt64(&sink.Scheduled, 1)

				sink.WorkPool.SendWork(func() {

					atomic.AddInt64(&sink.Retried, 1)

					sink.Logger.Info("retry syncing %s", source)

					sink.SyncFile(source, root, wg)

					atomic.AddInt64(&sink.Completed, 1)
				})

			}(source, root, wg)
		}

		wg.Wait()
	}

	return true
}

func (sink *Sync) MonitorStatus() {

	go func() {

		t0 := time.Now()

		for {

			rpt := sink.StatusReport()
			ttp := time.Since(t0)

			sink.Logger.Info("%s Time %v", rpt, ttp)

			time.Sleep(10 * time.Second)

			if sink.Scheduled == sink.Completed {
				break
			}
		}

		sink.Logger.Info(sink.StatusReport())
		sink.Logger.Info("monitoring complete")
	}()
}

func (sink *Sync) StatusReport() string {
	return fmt.Sprintf("Scheduled %d Completed %d Success %d Error %d Skipped %d Retried %d Goroutines %d",
		sink.Scheduled, sink.Completed, sink.Success, sink.Error, sink.Skipped, sink.Retried, runtime.NumGoroutine())
}
<commit_msg>add a buffered channel to prevent sync w/ file list from spiraling out of control<commit_after>package s3

// https://github.com/aws/aws-sdk-go
// https://docs.aws.amazon.com/sdk-for-go/api/service/s3.html

// https://github.com/goamz/goamz/blob/master/aws/aws.go
// https://github.com/goamz/goamz/blob/master/s3/s3.go

import (
	"bufio"
	"fmt"
	"github.com/goamz/goamz/aws"
	aws_s3 "github.com/goamz/goamz/s3"
	"github.com/jeffail/tunny"
	"github.com/whosonfirst/go-whosonfirst-crawl"
	log "github.com/whosonfirst/go-whosonfirst-log"
	pool "github.com/whosonfirst/go-whosonfirst-pool"
	utils "github.com/whosonfirst/go-whosonfirst-utils"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

type Sync struct {
	ACL           aws_s3.ACL
	Bucket        aws_s3.Bucket
	Prefix        string
	WorkPool      tunny.WorkPool
	Logger        *log.WOFLogger
	Debug         bool
	Success       int64
	Error         int64
	Skipped       int64
	Scheduled     int64
	Completed     int64
	Retried       int64
	TimeToProcess *time.Duration
	Retries       *pool.LIFOPool
	MaxRetries    float64 // max percentage of errors over scheduled
}

func NewSync(auth aws.Auth, region aws.Region, acl aws_s3.ACL, bucket string, prefix string, procs int, debug bool, logger *log.WOFLogger) *Sync {

	logger.Info("creating a new Sync thing-y with %d processes", procs)

	runtime.GOMAXPROCS(procs)

	workpool, _ := tunny.CreatePoolGeneric(procs).Open()

	retries := pool.NewLIFOPool()

	s := aws_s3.New(auth, region)
	b := s.Bucket(bucket)

	ttp := new(time.Duration)

	return &Sync{
		ACL:           acl,
		Bucket:        *b,
		Prefix:        prefix,
		WorkPool:      *workpool,
		Debug:         debug,
		Logger:        logger,
		Scheduled:     0,
		Completed:     0,
		Skipped:       0,
		Error:         0,
		Success:       0,
		Retried:       0,
		TimeToProcess: ttp,
		Retries:       retries,
		MaxRetries:    25.0, // maybe allow this to be user-defined ?
	}
}

func WOFSync(auth aws.Auth, bucket string, prefix string, procs int, debug bool, logger *log.WOFLogger) *Sync {

	return NewSync(auth, aws.USEast, aws_s3.PublicRead, bucket, prefix, procs, debug, logger)
}

func (sink *Sync) SyncDirectory(root string) error {

	defer sink.WorkPool.Close()

	t0 := time.Now()

	wg := new(sync.WaitGroup)

	callback := func(source string, info os.FileInfo) error {

		wg.Add(1)

		if info.IsDir() {
			return nil
		}

		err := sink.SyncFile(source, root, wg)

		if err != nil {
			sink.Logger.Error("failed to sync %s, because '%s'", source, err)
		}

		return nil
	}

	c := crawl.NewCrawler(root)
	_ = c.Crawl(callback)

	wg.Wait()

	sink.ProcessRetries(root)

	ttp := time.Since(t0)
	sink.TimeToProcess = &ttp

	return nil
}

func (sink *Sync) SyncFiles(files []string, root string) error {

	defer sink.WorkPool.Close()

	t0 := time.Now()

	wg := new(sync.WaitGroup)

	for _, path := range files {

		wg.Add(1)

		go func(path string, root string, wg *sync.WaitGroup) {
			sink.SyncFile(path, root, wg)
		}(path, root, wg)
	}

	wg.Wait()

	sink.ProcessRetries(root)

	ttp := time.Since(t0)
	sink.TimeToProcess = &ttp

	return nil
}

func (sink *Sync) SyncFileList(path string, root string) error {

	defer sink.WorkPool.Close()

	t0 := time.Now()

	file, err := os.Open(path)

	if err != nil {
		return err
	}

	defer file.Close()

	scanner := bufio.NewScanner(file)

	count := 100000
	ch := make(chan bool, count)

	go func() {
		for i := 0; i < count; i++ {
			ch <- true
		}
	}()

	wg := new(sync.WaitGroup)

	for scanner.Scan() {

		<-ch

		path := scanner.Text()

		wg.Add(1)

		go func(path string, root string, wg *sync.WaitGroup, ch chan bool) {
			sink.SyncFile(path, root, wg)
			ch <- true
		}(path, root, wg, ch)
	}

	wg.Wait()

	sink.ProcessRetries(root)

	ttp := time.Since(t0)
	sink.TimeToProcess = &ttp

	return nil
}

func (sink *Sync) SyncFile(source string, root string, wg *sync.WaitGroup) error {

	atomic.AddInt64(&sink.Scheduled, 1)

	_, err := sink.WorkPool.SendWork(func() {

		defer wg.Done()

		dest := source

		dest = strings.Replace(dest, root, "", -1)

		if sink.Prefix != "" {
			dest = path.Join(sink.Prefix, dest)
		}

		// Note: both HasChanged and SyncFile will ioutil.ReadFile(source)
		// which is a potential waste of time and resource. Or maybe we just
		// don't care? (20150930/thisisaaronland)

		sink.Logger.Debug("Looking for changes to %s (%s)", dest, sink.Prefix)

		change, ch_err := sink.HasChanged(source, dest)

		if ch_err != nil {
			atomic.AddInt64(&sink.Completed, 1)
			atomic.AddInt64(&sink.Error, 1)
			sink.Logger.Warning("failed to determine whether %s had changed, because '%s'", source, ch_err)

			sink.Retries.Push(&pool.PoolString{String: source})
			return
		}

		if sink.Debug == true {
			atomic.AddInt64(&sink.Completed, 1)
			atomic.AddInt64(&sink.Skipped, 1)
			sink.Logger.Debug("has %s changed? the answer is %v but does it really matter since debugging is enabled?", source, change)
			return
		}

		if !change {
			atomic.AddInt64(&sink.Completed, 1)
			atomic.AddInt64(&sink.Skipped, 1)
			sink.Logger.Debug("%s has not changed, skipping", source)
			return
		}

		err := sink.DoSyncFile(source, dest)

		if err != nil {
			sink.Retries.Push(&pool.PoolString{String: source})
			atomic.AddInt64(&sink.Error, 1)
		} else {
			atomic.AddInt64(&sink.Success, 1)
		}

		atomic.AddInt64(&sink.Completed, 1)
	})

	if err != nil {
		wg.Done()
		atomic.AddInt64(&sink.Error, 1)
		sink.Logger.Error("failed to schedule %s for processing, because %v", source, err)
		return err
	}

	sink.Logger.Debug("schedule %s for processing", source)
	return nil
}

func (sink *Sync) DoSyncFile(source string, dest string) error {

	sink.Logger.Debug("prepare %s for syncing", source)

	body, err := ioutil.ReadFile(source)

	if err != nil {
		sink.Logger.Error("Failed to read %s, because %v", source, err)
		return err
	}

	sink.Logger.Debug("PUT %s as %s", dest, sink.ACL)

	o := aws_s3.Options{}

	err = sink.Bucket.Put(dest, body, "text/plain", sink.ACL, o)

	if err != nil {
		sink.Logger.Error("failed to PUT %s, because '%s'", dest, err)
		return err
	}

	return nil
}

func (sink *Sync) HasChanged(source string, dest string) (ch bool, err error) {

	sink.Logger.Debug("HEAD %s", dest)

	headers := make(http.Header)
	rsp, err := sink.Bucket.Head(dest, headers)

	if err != nil {

		if e, ok := err.(*aws_s3.Error); ok && e.StatusCode == 404 {
			sink.Logger.Debug("%s is 404 so assuming it has changed (WHOA)", dest)
			return true, nil
		}

		sink.Logger.Error("failed to HEAD %s because %s", dest, err)
		return false, err
	}

	local_hash, err := utils.HashFile(source)

	if err != nil {
		sink.Logger.Warning("failed to hash %s, because %v", source, err)
		return false, err
	}

	etag := rsp.Header.Get("Etag")
	remote_hash := strings.Replace(etag, "\"", "", -1)

	sink.Logger.Debug("local hash is %s remote hash is %s", local_hash, remote_hash)

	if local_hash == remote_hash {
		return false, nil
	}

	// Okay so we think that things have changed but let's just check
	// modification times to be extra sure (20151112/thisisaaronland)

	info, err := os.Stat(source)

	if err != nil {
		sink.Logger.Error("failed to stat %s because %s", source, err)
		return false, err
	}

	mtime_local := info.ModTime()

	last_mod := rsp.Header.Get("Last-Modified")
	mtime_remote, err := time.Parse(time.RFC1123, last_mod)

	if err != nil {
		sink.Logger.Error("failed to parse timestamp %s because %s", last_mod, err)
		return false, err
	}

	// Because who remembers this stuff anyway...
	// func (t Time) Before(u Time) bool
	// Before reports whether the time instant t is before u.

	sink.Logger.Debug("local %s %s", mtime_local, source)
	sink.Logger.Debug("remote %s %s", mtime_remote, dest)

	if mtime_local.Before(mtime_remote) {
		sink.Logger.Warning("remote copy of %s has a more recent modification date (local: %s remote: %s)", source, mtime_local, mtime_remote)
		return false, nil
	}

	return true, nil
}

func (sink *Sync) ProcessRetries(root string) bool {

	to_retry := sink.Retries.Length()

	if to_retry > 0 {

		scheduled_f := float64(sink.Scheduled)
		retry_f := float64(to_retry)

		pct := (retry_f / scheduled_f) * 100.0

		if pct > sink.MaxRetries {
			sink.Logger.Warning("E_EXCESSIVE_ERRORS, %f percent of scheduled processes failed thus undermining our faith that they will work now...", pct)
			return false
		}

		sink.Logger.Info("There are %d failed requests that will now be retried", to_retry)

		wg := new(sync.WaitGroup)

		for sink.Retries.Length() > 0 {

			r, ok := sink.Retries.Pop()

			if !ok {
				sink.Logger.Error("failed to pop retries because... computers?")
				break
			}

			source := r.StringValue()

			wg.Add(1)

			go func(source string, root string, wg *sync.WaitGroup) {

				atomic.AddInt64(&sink.Scheduled, 1)

				sink.WorkPool.SendWork(func() {

					atomic.AddInt64(&sink.Retried, 1)

					sink.Logger.Info("retry syncing %s", source)

					sink.SyncFile(source, root, wg)

					atomic.AddInt64(&sink.Completed, 1)
				})

			}(source, root, wg)
		}

		wg.Wait()
	}

	return true
}

func (sink *Sync) MonitorStatus() {

	go func() {

		t0 := time.Now()

		for {

			rpt := sink.StatusReport()
			ttp := time.Since(t0)

			sink.Logger.Info("%s Time %v", rpt, ttp)

			time.Sleep(10 * time.Second)

			if sink.Scheduled == sink.Completed {
				break
			}
		}

		sink.Logger.Info(sink.StatusReport())
		sink.Logger.Info("monitoring complete")
	}()
}

func (sink *Sync) StatusReport() string {
	return fmt.Sprintf("Scheduled %d Completed %d Success %d Error %d Skipped %d Retried %d Goroutines %d",
		sink.Scheduled, sink.Completed, sink.Success, sink.Error, sink.Skipped, sink.Retried, runtime.NumGoroutine())
}
<|endoftext|>


<commit_before>package server

import (
	"net/http"

	"github.com/influxdata/chronograf"
)

// AuthRoute are the routes for each type of OAuth2 provider
type AuthRoute struct {
	Name     string `json:"name"`     // Name uniquely identifies the provider
	Label    string `json:"label"`    // Label is a user-facing string to present in the UI
	Login    string `json:"login"`    // Login is the route to the login redirect path
	Logout   string `json:"logout"`   // Logout is the route to the logout redirect path
	Callback string `json:"callback"` // Callback is the route the provider calls to exchange the code/state
}

// AuthRoutes contains all OAuth2 provider routes.
type AuthRoutes []AuthRoute

// Lookup searches all the routes for a specific provider
func (r *AuthRoutes) Lookup(provider string) (AuthRoute, bool) {
	for _, route := range *r {
		if route.Name == provider {
			return route, true
		}
	}
	return AuthRoute{}, false
}

type getRoutesResponse struct {
	Layouts    string      `json:"layouts"`    // Location of the layouts endpoint
	Mappings   string      `json:"mappings"`   // Location of the application mappings endpoint
	Sources    string      `json:"sources"`    // Location of the sources endpoint
	Me         string      `json:"me"`         // Location of the me endpoint
	Dashboards string      `json:"dashboards"` // Location of the dashboards endpoint
	Auth       []AuthRoute `json:"auth"`       // Location of all auth routes.
}

// AllRoutes returns all top level routes within chronograf
func AllRoutes(authRoutes []AuthRoute, logger chronograf.Logger) http.HandlerFunc {
	routes := getRoutesResponse{
		Sources:    "/chronograf/v1/sources",
		Layouts:    "/chronograf/v1/layouts",
		Me:         "/chronograf/v1/me",
		Mappings:   "/chronograf/v1/mappings",
		Dashboards: "/chronograf/v1/dashboards",
		Auth:       make([]AuthRoute, len(authRoutes)),
	}

	for i, route := range authRoutes {
		routes.Auth[i] = route
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		encodeJSON(w, http.StatusOK, routes, logger)
		return
	})
}
<commit_msg>add kapacitors to AllRoutes<commit_after>package server

import (
	"net/http"

	"github.com/influxdata/chronograf"
)

// AuthRoute are the routes for each type of OAuth2 provider
type AuthRoute struct {
	Name     string `json:"name"`     // Name uniquely identifies the provider
	Label    string `json:"label"`    // Label is a user-facing string to present in the UI
	Login    string `json:"login"`    // Login is the route to the login redirect path
	Logout   string `json:"logout"`   // Logout is the route to the logout redirect path
	Callback string `json:"callback"` // Callback is the route the provider calls to exchange the code/state
}

// AuthRoutes contains all OAuth2 provider routes.
type AuthRoutes []AuthRoute

// Lookup searches all the routes for a specific provider
func (r *AuthRoutes) Lookup(provider string) (AuthRoute, bool) {
	for _, route := range *r {
		if route.Name == provider {
			return route, true
		}
	}
	return AuthRoute{}, false
}

type getRoutesResponse struct {
	Layouts    string      `json:"layouts"`    // Location of the layouts endpoint
	Mappings   string      `json:"mappings"`   // Location of the application mappings endpoint
	Sources    string      `json:"sources"`    // Location of the sources endpoint
	Me         string      `json:"me"`         // Location of the me endpoint
	Dashboards string      `json:"dashboards"` // Location of the dashboards endpoint
	Kapacitors string      `json:"kapacitors"` // Location of the kapacitors endpoint
	Auth       []AuthRoute `json:"auth"`       // Location of all auth routes.
}

// AllRoutes returns all top level routes within chronograf
func AllRoutes(authRoutes []AuthRoute, logger chronograf.Logger) http.HandlerFunc {
	routes := getRoutesResponse{
		Sources:    "/chronograf/v1/sources",
		Layouts:    "/chronograf/v1/layouts",
		Me:         "/chronograf/v1/me",
		Mappings:   "/chronograf/v1/mappings",
		Dashboards: "/chronograf/v1/dashboards",
		Kapacitors: "/chronograf/v1/kapacitors",
		Auth:       make([]AuthRoute, len(authRoutes)),
	}

	for i, route := range authRoutes {
		routes.Auth[i] = route
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		encodeJSON(w, http.StatusOK, routes, logger)
		return
	})
}
<|endoftext|>


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

package main

import (
	"fmt"
	"net/http"

	"github.com/satori/go.uuid"
	"github.com/trackit/jsonlog"

	_ "github.com/trackit/trackit2/aws"
	"github.com/trackit/trackit2/config"
	"github.com/trackit/trackit2/routes"
	_ "github.com/trackit/trackit2/users"
)

var buildNumber string
var backendId = getBackendId()

func main() {
	logger := jsonlog.DefaultLogger
	logger.Info("Started.", struct {
		BackendId string `json:"backendId"`
	}{backendId})
	initializeHandlers()
	logger.Info(fmt.Sprintf("Listening on %s.", config.HttpAddress), nil)
	err := http.ListenAndServe(config.HttpAddress, nil)
	logger.Error("Server stopped.", err.Error())
}

// initializeHandlers sets the HTTP server up with handler functions.
func initializeHandlers() {
	globalDecorators := []routes.Decorator{
		routes.RequestId{},
		routes.RouteLog{},
		routes.BackendId{backendId},
		routes.ErrorBody{},
		routes.Cors{
			AllowCredentials: true,
			AllowHeaders:     []string{"Content-Type", "Accept", "Authorization"},
			AllowOrigin:      []string{"*"},
		},
	}
	logger := jsonlog.DefaultLogger
	routes.DocumentationHandler().Register("/docs")
	for _, rh := range routes.RegisteredHandlers {
		applyDecoratorsAndHandle(rh.Pattern, rh.Handler, globalDecorators)
		logger.Info(fmt.Sprintf("Registered route %s.", rh.Pattern), nil)
	}
}

// applyDecoratorsAndHandle applies a list of decorators to a handler and
// registers it.
func applyDecoratorsAndHandle(p string, h routes.Handler, ds []routes.Decorator) {
	h = h.With(ds...)
	http.Handle(p, h)
}

// getBackendId returns an ID unique to the current process. It can also be set
// in the config to a determined string. It contains the build number.
func getBackendId() string {
	if config.BackendId != "" {
		return config.BackendId
	} else {
		return fmt.Sprintf("%s-%s", uuid.NewV1().String(), buildNumber)
	}
}
<commit_msg>server: enable PanicAsError global decorator<commit_after>//   Copyright 2017 MSolution.IO
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.

package main

import (
	"fmt"
	"net/http"

	"github.com/satori/go.uuid"
	"github.com/trackit/jsonlog"

	_ "github.com/trackit/trackit2/aws"
	"github.com/trackit/trackit2/config"
	"github.com/trackit/trackit2/routes"
	_ "github.com/trackit/trackit2/users"
)

var buildNumber string
var backendId = getBackendId()

func main() {
	logger := jsonlog.DefaultLogger
	logger.Info("Started.", struct {
		BackendId string `json:"backendId"`
	}{backendId})
	initializeHandlers()
	logger.Info(fmt.Sprintf("Listening on %s.", config.HttpAddress), nil)
	err := http.ListenAndServe(config.HttpAddress, nil)
	logger.Error("Server stopped.", err.Error())
}

// initializeHandlers sets the HTTP server up with handler functions.
func initializeHandlers() {
	globalDecorators := []routes.Decorator{
		routes.RequestId{},
		routes.RouteLog{},
		routes.BackendId{backendId},
		routes.ErrorBody{},
		routes.PanicAsError{},
		routes.Cors{
			AllowCredentials: true,
			AllowHeaders:     []string{"Content-Type", "Accept", "Authorization"},
			AllowOrigin:      []string{"*"},
		},
	}
	logger := jsonlog.DefaultLogger
	routes.DocumentationHandler().Register("/docs")
	for _, rh := range routes.RegisteredHandlers {
		applyDecoratorsAndHandle(rh.Pattern, rh.Handler, globalDecorators)
		logger.Info(fmt.Sprintf("Registered route %s.", rh.Pattern), nil)
	}
}

// applyDecoratorsAndHandle applies a list of decorators to a handler and
// registers it.
func applyDecoratorsAndHandle(p string, h routes.Handler, ds []routes.Decorator) {
	h = h.With(ds...)
	http.Handle(p, h)
}

// getBackendId returns an ID unique to the current process. It can also be set
// in the config to a determined string. It contains the build number.
func getBackendId() string {
	if config.BackendId != "" {
		return config.BackendId
	} else {
		return fmt.Sprintf("%s-%s", uuid.NewV1().String(), buildNumber)
	}
}
<|endoftext|>


<commit_before>package server

import (
	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/gorilla/mux"
	"github.com/meatballhat/negroni-logrus"
	"github.com/oxfeeefeee/appgo"
	"github.com/oxfeeefeee/appgo/auth"
	"github.com/phyber/negroni-gzip/gzip"
	"github.com/rs/cors"
	"github.com/unrolled/render"
	"html/template"
	"net/http"
	_ "net/http/pprof"
	"strings"
)

type Server struct {
	ts          TokenStore
	middlewares []negroni.HandlerFunc
	ver         *versioning
	*mux.Router
}
type TokenStore interface {
	Validate(token auth.Token) bool
}

func NewServer(ts TokenStore, middlewares []negroni.HandlerFunc) *Server {
	return &Server{
		ts,
		middlewares,
		newVersioning(),
		mux.NewRouter(),
	}
}

func (s *Server) AddRest(path string, rests []interface{}) {
	renderer := render.New(render.Options{
		Directory:     "N/A",
		IndentJSON:    appgo.Conf.DevMode,
		IsDevelopment: appgo.Conf.DevMode,
	})
	for _, api := range rests {
		h := newHandler(api, HandlerTypeJson, s.ts, renderer)
		s.Handle(path+h.path, h).Methods(h.supports...)
	}
}

func (s *Server) AddHtml(path, layout string, htmls []interface{}, funcs template.FuncMap) {
	// add "static" template function
	static := func(path string) string {
		return s.ver.getStatic(path)
	}
	if funcs == nil {
		funcs = template.FuncMap{}
	}
	funcs["static"] = static

	renderer := render.New(render.Options{
		Directory:     appgo.Conf.TemplatePath,
		Layout:        layout,
		Funcs:         []template.FuncMap{funcs},
		IsDevelopment: appgo.Conf.DevMode,
	})
	for _, api := range htmls {
		h := newHandler(api, HandlerTypeHtml, s.ts, renderer)
		s.Handle(path+h.path, h).Methods("GET")
	}
}

func (s *Server) AddFeed(path string, feeds []interface{}) {
	// renderer is only for rendering error
	renderer := render.New(render.Options{
		Directory:     "N/A",
		IsDevelopment: appgo.Conf.DevMode,
	})
	for _, api := range feeds {
		h := newHandler(api, HandlerTypeFeed, s.ts, renderer)
		s.Handle(path+h.path, h).Methods("GET")
	}
}

func (s *Server) AddProxy(path string, handler http.Handler) {
	s.PathPrefix(path).Handler(http.StripPrefix(path, handler))
}

func (s *Server) AddStatic(path, fileDir string) {
	s.ver.addMap(path, fileDir)
	s.AddProxy(path, http.FileServer(http.Dir(fileDir)))
}

func (s *Server) AddAppleAppSiteAsso(content []byte) {
	f := func(w http.ResponseWriter, req *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.Write(content)
	}
	s.HandleFunc("/apple-app-site-association", f)
}

func (s *Server) Serve() {
	if appgo.Conf.Pprof.Enable {
		go func() {
			log.Infoln(http.ListenAndServe(":"+appgo.Conf.Pprof.Port, nil))
		}()
	}

	n := negroni.New()
	n.Use(negroni.NewRecovery())
	n.Use(negronilogrus.NewCustomMiddleware(
		appgo.Conf.LogLevel, &log.TextFormatter{}, "appgo"))
	n.Use(cors.New(corsOptions()))
	if appgo.Conf.Negroni.GZip {
		n.Use(gzip.Gzip(gzip.BestSpeed))
	}
	for _, mw := range s.middlewares {
		n.Use(negroni.HandlerFunc(mw))
	}
	n.UseHandler(s)
	n.Run(appgo.Conf.Negroni.Port)
}

func corsOptions() cors.Options {
	origins := strings.Split(appgo.Conf.Cors.AllowedOrigins, ",")
	methods := strings.Split(appgo.Conf.Cors.AllowedMethods, ",")
	headers := strings.Split(appgo.Conf.Cors.AllowedHeaders, ",")
	return cors.Options{
		AllowedOrigins:     origins,
		AllowedMethods:     methods,
		AllowedHeaders:     headers,
		OptionsPassthrough: appgo.Conf.Cors.OptionsPassthrough,
		Debug:              appgo.Conf.Cors.Debug,
	}
}
<commit_msg>middleware fixes<commit_after>package server

import (
	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/gorilla/mux"
	"github.com/meatballhat/negroni-logrus"
	"github.com/oxfeeefeee/appgo"
	"github.com/oxfeeefeee/appgo/auth"
	"github.com/phyber/negroni-gzip/gzip"
	"github.com/rs/cors"
	"github.com/unrolled/render"
	"html/template"
	"net/http"
	_ "net/http/pprof"
	"strings"
)

type Server struct {
	ts          TokenStore
	middlewares []negroni.HandlerFunc
	ver         *versioning
	*mux.Router
}
type TokenStore interface {
	Validate(token auth.Token) bool
}

func NewServer(ts TokenStore, middlewares []negroni.HandlerFunc) *Server {
	return &Server{
		ts,
		middlewares,
		newVersioning(),
		mux.NewRouter(),
	}
}

func (s *Server) AddRest(path string, rests []interface{}) {
	renderer := render.New(render.Options{
		Directory:     "N/A",
		IndentJSON:    appgo.Conf.DevMode,
		IsDevelopment: appgo.Conf.DevMode,
	})
	for _, api := range rests {
		h := newHandler(api, HandlerTypeJson, s.ts, renderer)
		s.Handle(path+h.path, h).Methods(h.supports...)
	}
}

func (s *Server) AddHtml(path, layout string, htmls []interface{}, funcs template.FuncMap) {
	// add "static" template function
	static := func(path string) string {
		return s.ver.getStatic(path)
	}
	if funcs == nil {
		funcs = template.FuncMap{}
	}
	funcs["static"] = static

	renderer := render.New(render.Options{
		Directory:     appgo.Conf.TemplatePath,
		Layout:        layout,
		Funcs:         []template.FuncMap{funcs},
		IsDevelopment: appgo.Conf.DevMode,
	})
	for _, api := range htmls {
		h := newHandler(api, HandlerTypeHtml, s.ts, renderer)
		s.Handle(path+h.path, h).Methods("GET")
	}
}

func (s *Server) AddFeed(path string, feeds []interface{}) {
	// renderer is only for rendering error
	renderer := render.New(render.Options{
		Directory:     "N/A",
		IsDevelopment: appgo.Conf.DevMode,
	})
	for _, api := range feeds {
		h := newHandler(api, HandlerTypeFeed, s.ts, renderer)
		s.Handle(path+h.path, h).Methods("GET")
	}
}

func (s *Server) AddProxy(path string, handler http.Handler) {
	s.PathPrefix(path).Handler(http.StripPrefix(path, handler))
}

func (s *Server) AddStatic(path, fileDir string) {
	s.ver.addMap(path, fileDir)
	s.AddProxy(path, http.FileServer(http.Dir(fileDir)))
}

func (s *Server) AddAppleAppSiteAsso(content []byte) {
	f := func(w http.ResponseWriter, req *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.Write(content)
	}
	s.HandleFunc("/apple-app-site-association", f)
}

func (s *Server) Serve() {
	if appgo.Conf.Pprof.Enable {
		go func() {
			log.Infoln(http.ListenAndServe(":"+appgo.Conf.Pprof.Port, nil))
		}()
	}

	n := negroni.New()
	rec := negroni.NewRecovery()
	rec.StackAll = true
	n.Use(rec)
	llog := negronilogrus.NewCustomMiddleware(
		appgo.Conf.LogLevel, &log.TextFormatter{}, "appgo")
	llog.Logger = log.StandardLogger()
	n.Use(llog)
	n.Use(cors.New(corsOptions()))
	for _, mw := range s.middlewares {
		n.Use(negroni.HandlerFunc(mw))
	}
	if appgo.Conf.Negroni.GZip {
		n.Use(gzip.Gzip(gzip.BestSpeed))
	}
	n.UseHandler(s)
	n.Run(appgo.Conf.Negroni.Port)
}

func corsOptions() cors.Options {
	origins := strings.Split(appgo.Conf.Cors.AllowedOrigins, ",")
	methods := strings.Split(appgo.Conf.Cors.AllowedMethods, ",")
	headers := strings.Split(appgo.Conf.Cors.AllowedHeaders, ",")
	return cors.Options{
		AllowedOrigins:     origins,
		AllowedMethods:     methods,
		AllowedHeaders:     headers,
		OptionsPassthrough: appgo.Conf.Cors.OptionsPassthrough,
		Debug:              appgo.Conf.Cors.Debug,
	}
}
<|endoftext|>


<commit_before>package server

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/WhiteHatCP/seclab-listener/backend"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"time"
)

const (
	keyLength   = 32
	reqOpen     = 0xff
	reqClose    = 0x00
	reqKeygen   = 0xaa
	respAllGood = 0xff
	respNewKey  = 0x55
)

var outLog = log.New(os.Stdout, "", log.LstdFlags)
var errLog = log.New(os.Stderr, "", log.LstdFlags)

// Server is the interface that handles the network protocol
type Server interface {
	AddBackend(backend.Backend)
	CheckMessage([]byte) error
	DispatchRequest(byte) ([]byte, error)
	KeyRotate() ([]byte, error)
	Serve(net.Listener)
}

type server struct {
	keypath  string
	maxAge   int
	backends []backend.Backend
}

// New creates a new instance of a Server
func New(keypath string, maxAge int) Server {
	return &server{
		keypath:  keypath,
		maxAge:   maxAge,
		backends: nil,
	}
}

func (s *server) AddBackend(b backend.Backend) {
	s.backends = append(s.backends, b)
}

func checkHash(key []byte, payload []byte, hash []byte) bool {
	mac := hmac.New(sha256.New, key)
	mac.Write(payload)
	expected := mac.Sum(nil)
	return hmac.Equal(hash, expected)
}

// Read the status byte, validate HMAC and timestamp
func (s *server) CheckMessage(data []byte) error {
	key, err := ioutil.ReadFile(s.keypath)
	if err != nil {
		return err
	}
	if !checkHash(key, data[:9], data[9:]) {
		return errors.New("Incorrect HMAC signature")
	}
	ts := int64(binary.BigEndian.Uint64(data[1:9]))
	if time.Now().Unix()-ts > int64(s.maxAge) {
		return errors.New("Request expired")
	}
	return nil
}

func (s *server) KeyRotate() ([]byte, error) {
	resp := make([]byte, 9+keyLength)
	key := resp[9:]
	if _, err := rand.Read(key); err != nil {
		return nil, err
	}
	if err := ioutil.WriteFile(s.keypath, key, 0600); err != nil {
		return nil, err
	}
	resp[0] = respNewKey
	binary.BigEndian.PutUint64(resp[1:9], uint64(time.Now().Unix()))
	return resp, nil
}

func (s *server) open() error {
	for _, b := range s.backends {
		if err := b.Open(); err != nil {
			return err
		}
	}
	return nil
}

func (s *server) close() error {
	for _, b := range s.backends {
		if err := b.Close(); err != nil {
			return err
		}
	}
	return nil
}

func (s *server) DispatchRequest(status byte) ([]byte, error) {
	if status == reqOpen {
		outLog.Print("Received request: open")
		return []byte{respAllGood}, s.open()
	} else if status == reqClose {
		outLog.Print("Received request: close")
		return []byte{respAllGood}, s.close()
	} else if status == reqKeygen {
		return s.KeyRotate()
	}
	return nil, fmt.Errorf("Unrecognized status byte: 0x%02x", status)
}

func (s *server) handleConnection(conn net.Conn) {
	defer conn.Close()
	data := make([]byte, 9+keyLength)
	for {
		if _, err := io.ReadFull(conn, data); err != nil {
			if err != io.EOF {
				errLog.Print(err)
			}
			return
		}
		err := s.CheckMessage(data)
		if err != nil {
			errLog.Print(err)
			return
		}
		resp, err := s.DispatchRequest(data[0])
		if err != nil {
			errLog.Print(err)
			return
		}
		conn.Write(resp)
	}
}

func (s *server) Serve(ln net.Listener) {
	outLog.Print("Seclab listener started")
	for {
		conn, err := ln.Accept()
		if err != nil {
			errLog.Fatal(err)
		}
		go s.handleConnection(conn)
	}
}
<commit_msg>Use Sentry to report server errors<commit_after>package server

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/WhiteHatCP/seclab-listener/backend"
	"github.com/getsentry/raven-go"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"time"
)

const (
	keyLength   = 32
	reqOpen     = 0xff
	reqClose    = 0x00
	reqKeygen   = 0xaa
	respAllGood = 0xff
	respNewKey  = 0x55
)

var outLog = log.New(os.Stdout, "", log.LstdFlags)
var errLog = log.New(os.Stderr, "", log.LstdFlags)

// Server is the interface that handles the network protocol
type Server interface {
	AddBackend(backend.Backend)
	CheckMessage([]byte) error
	DispatchRequest(byte) ([]byte, error)
	KeyRotate() ([]byte, error)
	Serve(net.Listener)
}

type server struct {
	keypath  string
	maxAge   int
	backends []backend.Backend
}

// New creates a new instance of a Server
func New(keypath string, maxAge int) Server {
	return &server{
		keypath:  keypath,
		maxAge:   maxAge,
		backends: nil,
	}
}

func (s *server) AddBackend(b backend.Backend) {
	s.backends = append(s.backends, b)
}

func checkHash(key []byte, payload []byte, hash []byte) bool {
	mac := hmac.New(sha256.New, key)
	mac.Write(payload)
	expected := mac.Sum(nil)
	return hmac.Equal(hash, expected)
}

// Read the status byte, validate HMAC and timestamp
func (s *server) CheckMessage(data []byte) error {
	key, err := ioutil.ReadFile(s.keypath)
	if err != nil {
		return err
	}
	if !checkHash(key, data[:9], data[9:]) {
		return errors.New("Incorrect HMAC signature")
	}
	ts := int64(binary.BigEndian.Uint64(data[1:9]))
	if time.Now().Unix()-ts > int64(s.maxAge) {
		return errors.New("Request expired")
	}
	return nil
}

func (s *server) KeyRotate() ([]byte, error) {
	resp := make([]byte, 9+keyLength)
	key := resp[9:]
	if _, err := rand.Read(key); err != nil {
		return nil, err
	}
	if err := ioutil.WriteFile(s.keypath, key, 0600); err != nil {
		return nil, err
	}
	resp[0] = respNewKey
	binary.BigEndian.PutUint64(resp[1:9], uint64(time.Now().Unix()))
	return resp, nil
}

func (s *server) open() error {
	for _, b := range s.backends {
		if err := b.Open(); err != nil {
			return err
		}
	}
	return nil
}

func (s *server) close() error {
	for _, b := range s.backends {
		if err := b.Close(); err != nil {
			return err
		}
	}
	return nil
}

func (s *server) DispatchRequest(status byte) ([]byte, error) {
	if status == reqOpen {
		outLog.Print("Received request: open")
		return []byte{respAllGood}, s.open()
	} else if status == reqClose {
		outLog.Print("Received request: close")
		return []byte{respAllGood}, s.close()
	} else if status == reqKeygen {
		return s.KeyRotate()
	}
	return nil, fmt.Errorf("Unrecognized status byte: 0x%02x", status)
}

func (s *server) readAndRespond(conn net.Conn) error {
	defer conn.Close()
	data := make([]byte, 9+keyLength)
	for {
		if _, err := io.ReadFull(conn, data); err != nil {
			if err != io.EOF {
				return err
			}
			return nil
		}
		if err := s.CheckMessage(data); err != nil {
			return err
		}
		resp, err := s.DispatchRequest(data[0])
		if err != nil {
			return err
		}
		if _, err = conn.Write(resp); err != nil {
			return err
		}
	}
}

func (s *server) handleConnection(conn net.Conn) {
	if err := s.readAndRespond(conn); err != nil {
		raven.CaptureError(err, nil)
		errLog.Print(err)
	}
}

func (s *server) Serve(ln net.Listener) {
	outLog.Print("Seclab listener started")
	for {
		conn, err := ln.Accept()
		if err != nil {
			errLog.Fatal(err)
		}
		go s.handleConnection(conn)
	}
}
<|endoftext|>


<commit_before>package server

import (
	"crypto/rsa"
	"fmt"
	"log"
	"net"
	"os"
	"time"

	"github.com/lologarithm/survival/server/messages"
)

const (
	port string = ":24816"
)

type Server struct {
	conn             *net.UDPConn
	disconnectPlayer chan Client
	outToNetwork     chan OutgoingMessage
	toGameManager    chan GameMessage
	inputBuffer      []byte
	encryptionKey    *rsa.PrivateKey

	connections map[string]*Client
	gameManager *GameManager
	clientID    uint32
}

func (s *Server) handleMessage() {
	// TODO: Add timeout on read to check for stale connections and add new user connections.
	s.conn.SetReadDeadline(time.Now().Add(time.Second * 5))
	n, addr, err := s.conn.ReadFromUDP(s.inputBuffer)

	if err != nil {
		return
	}
	addrkey := addr.String()
	if n == 0 {
		s.DisconnectConn(addrkey)
	}
	if _, ok := s.connections[addrkey]; !ok {
		s.clientID++
		fmt.Printf("New Connection: %v, ID: %d\n", addrkey, s.clientID)
		s.connections[addrkey] = &Client{
			address:         addr,
			FromNetwork:     NewBytePipe(0),
			FromGameManager: make(chan InternalMessage, 10),
			toGameManager:   s.toGameManager,
			ID:              s.clientID,
		}
		go s.connections[addrkey].ProcessBytes(s.outToNetwork, s.disconnectPlayer)
	}

	s.connections[addrkey].FromNetwork.Write(s.inputBuffer[0:n])
}

func (s *Server) DisconnectConn(addrkey string) {
	// close(s.connections[addrkey].FromNetwork)
	delete(s.connections, addrkey)
}

func (s *Server) sendMessages() {
	for {
		msg := <-s.outToNetwork
		numMsg := (len(msg.msg.RawBytes) / 512) + 1
		st := 0
		b := 512
		for i := 0; i < numMsg; i++ {
			if i == numMsg-1 {
				b = len(msg.msg.RawBytes) % 512
			}
			if n, err := s.conn.WriteToUDP(msg.msg.RawBytes[st:st+b], msg.dest.address); err != nil {
				fmt.Println("Error: ", err, " Bytes Written: ", n)
			} else {
				log.Printf("Wrote message %d with %d bytes to %v.", msg.msg.Frame.MsgType, n, msg.dest.address)
			}
			st += b
		}
	}
}

func RunServer(exit chan int) {
	toGameManager := make(chan GameMessage, 1024)
	outToNetwork := make(chan OutgoingMessage, 1024)

	manager := NewGameManager(exit, toGameManager, outToNetwork)
	go manager.Run()

	udpAddr, err := net.ResolveUDPAddr("udp", port)
	if err != nil {
		log.Printf("Failed to open UDP port: %s", err)
		os.Exit(1)
	}
	fmt.Println("Now listening on port", port)

	var s Server
	s.connections = make(map[string]*Client, 512)
	s.inputBuffer = make([]byte, 8092)
	s.toGameManager = toGameManager
	s.outToNetwork = outToNetwork
	s.disconnectPlayer = make(chan Client, 512)
	s.conn, err = net.ListenUDP("udp", udpAddr)
	if err != nil {
		log.Printf("Failed to open UDP port: %s", err)
		os.Exit(1)
	}

	go s.sendMessages()
	fmt.Println("Server Started!")

	run := true
	for run {
		select {
		case <-exit:
			fmt.Println("Killing Socket Server")
			s.conn.Close()
			run = false
		case client := <-s.disconnectPlayer:
			s.DisconnectConn(client.address.String())
		default:
			s.handleMessage()
		}
	}
}

type OutgoingMessage struct {
	dest *Client
	msg  messages.Message
}
<commit_msg>Added note to actually fix frames for multipart messages<commit_after>package server

import (
	"crypto/rsa"
	"fmt"
	"log"
	"net"
	"os"
	"time"

	"github.com/lologarithm/survival/server/messages"
)

const (
	port string = ":24816"
)

type Server struct {
	conn             *net.UDPConn
	disconnectPlayer chan Client
	outToNetwork     chan OutgoingMessage
	toGameManager    chan GameMessage
	inputBuffer      []byte
	encryptionKey    *rsa.PrivateKey

	connections map[string]*Client
	gameManager *GameManager
	clientID    uint32
}

func (s *Server) handleMessage() {
	// TODO: Add timeout on read to check for stale connections and add new user connections.
	s.conn.SetReadDeadline(time.Now().Add(time.Second * 5))
	n, addr, err := s.conn.ReadFromUDP(s.inputBuffer)

	if err != nil {
		return
	}
	addrkey := addr.String()
	if n == 0 {
		s.DisconnectConn(addrkey)
	}
	if _, ok := s.connections[addrkey]; !ok {
		s.clientID++
		fmt.Printf("New Connection: %v, ID: %d\n", addrkey, s.clientID)
		s.connections[addrkey] = &Client{
			address:         addr,
			FromNetwork:     NewBytePipe(0),
			FromGameManager: make(chan InternalMessage, 10),
			toGameManager:   s.toGameManager,
			ID:              s.clientID,
		}
		go s.connections[addrkey].ProcessBytes(s.outToNetwork, s.disconnectPlayer)
	}

	s.connections[addrkey].FromNetwork.Write(s.inputBuffer[0:n])
}

func (s *Server) DisconnectConn(addrkey string) {
	// close(s.connections[addrkey].FromNetwork)
	delete(s.connections, addrkey)
}

func (s *Server) sendMessages() {
	for {
		msg := <-s.outToNetwork
		numMsg := (len(msg.msg.RawBytes) / 512) + 1
		st := 0
		b := 512
		for i := 0; i < numMsg; i++ {
			if i == numMsg-1 {
				b = len(msg.msg.RawBytes) % 512
			}
			// TODO: write frames for each piece with a message type of 'continue' maybe. (SEQ IS IMPORTANT!) 		msg.dest.Seq
			if n, err := s.conn.WriteToUDP(msg.msg.RawBytes[st:st+b], msg.dest.address); err != nil {
				fmt.Println("Error: ", err, " Bytes Written: ", n)
			} else {
				log.Printf("Wrote message %d with %d bytes to %v.", msg.msg.Frame.MsgType, n, msg.dest.address)
			}
			st += b
		}
	}
}

func RunServer(exit chan int) {
	toGameManager := make(chan GameMessage, 1024)
	outToNetwork := make(chan OutgoingMessage, 1024)

	manager := NewGameManager(exit, toGameManager, outToNetwork)
	go manager.Run()

	udpAddr, err := net.ResolveUDPAddr("udp", port)
	if err != nil {
		log.Printf("Failed to open UDP port: %s", err)
		os.Exit(1)
	}
	fmt.Println("Now listening on port", port)

	var s Server
	s.connections = make(map[string]*Client, 512)
	s.inputBuffer = make([]byte, 8092)
	s.toGameManager = toGameManager
	s.outToNetwork = outToNetwork
	s.disconnectPlayer = make(chan Client, 512)
	s.conn, err = net.ListenUDP("udp", udpAddr)
	if err != nil {
		log.Printf("Failed to open UDP port: %s", err)
		os.Exit(1)
	}

	go s.sendMessages()
	fmt.Println("Server Started!")

	run := true
	for run {
		select {
		case <-exit:
			fmt.Println("Killing Socket Server")
			s.conn.Close()
			run = false
		case client := <-s.disconnectPlayer:
			s.DisconnectConn(client.address.String())
		default:
			s.handleMessage()
		}
	}
}

type OutgoingMessage struct {
	dest *Client
	msg  messages.Message
}
<|endoftext|>