<commit_before>package pb

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

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 uding int64 as total
func New64(total int64) *ProgressBar {
	pb := &ProgressBar{
		Total:         total,
		RefreshRate:   DEFAULT_REFRESH_RATE,
		ShowPercent:   true,
		ShowCounters:  true,
		ShowBar:       true,
		ShowTimeLeft:  true,
		ShowFinalTime: true,
		Units:         U_NO,
		ManualUpdate:  false,
		currentValue:  -1,
	}
	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)

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

	isFinish     int32
	startTime    time.Time
	currentValue int64

	prefix, postfix string

	BarStart string
	BarEnd   string
	Empty    string
	Current  string
	CurrentN string
}

// Start print
func (pb *ProgressBar) Start() *ProgressBar {
	pb.startTime = time.Now()
	if pb.Total == 0 {
		pb.ShowBar = false
		pb.ShowTimeLeft = false
		pb.ShowPercent = false
	}
	if !pb.ManualUpdate {
		go pb.writer()
	}
	return pb
}

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

// Set current value
func (pb *ProgressBar) Set(current int) {
	atomic.StoreInt64(&pb.current, int64(current))
}

// 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.prefix = prefix
	return pb
}

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

// Set custom format for bar
// Example: bar.Format("[=>_]")
func (pb *ProgressBar) Format(format string) *ProgressBar {
	formatEntries := strings.Split(format, "")
	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() {
	atomic.StoreInt32(&pb.isFinish, 1)
	pb.write(atomic.LoadInt64(&pb.current))
	if !pb.NotPrint {
		fmt.Println()
	}
}

// End print and write string 'str'
func (pb *ProgressBar) FinishPrint(str string) {
	pb.Finish()
	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
func (pb *ProgressBar) NewProxyReader(r io.Reader) *Reader {
	return &Reader{r, pb}
}

func (pb *ProgressBar) write(current int64) {
	width := pb.getWidth()

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

	// percents
	if pb.ShowPercent {
		percent := float64(current) / (float64(pb.Total) / float64(100))
		percentBox = fmt.Sprintf(" %#.02f %% ", percent)
	}

	// counters
	if pb.ShowCounters {
		if pb.Total > 0 {
			countersBox = fmt.Sprintf("%s / %s ", Format(current, pb.Units), Format(pb.Total, pb.Units))
		} else {
			countersBox = Format(current, pb.Units) + " "
		}
	}

	// time left
	fromStart := time.Now().Sub(pb.startTime)
	if atomic.LoadInt32(&pb.isFinish) != 0 {
		if pb.ShowFinalTime {
			left := (fromStart / time.Second) * time.Second
			timeLeftBox = left.String()
		}
	} else if pb.ShowTimeLeft && current > 0 {
		perEntry := fromStart / time.Duration(current)
		left := time.Duration(pb.Total-current) * perEntry
		left = (left / time.Second) * time.Second
		timeLeftBox = left.String()
	}

	// speed
	if pb.ShowSpeed && current > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		speed := float64(current) / (float64(fromStart) / float64(time.Second))
		speedBox = Format(int64(speed), pb.Units) + "/s "
	}

	// bar
	if pb.ShowBar {
		size := width - len(countersBox+pb.BarStart+pb.BarEnd+percentBox+timeLeftBox+speedBox+pb.prefix+pb.postfix)
		if size > 0 {
			curCount := int(math.Ceil((float64(current) / float64(pb.Total)) * float64(size)))
			emptCount := size - curCount
			barBox = pb.BarStart
			if emptCount < 0 {
				emptCount = 0
			}
			if curCount > size {
				curCount = size
			}
			if emptCount <= 0 {
				barBox += strings.Repeat(pb.Current, curCount)
			} else if curCount > 0 {
				barBox += strings.Repeat(pb.Current, curCount-1) + pb.CurrentN
			}

			barBox += strings.Repeat(pb.Empty, emptCount) + pb.BarEnd
		}
	}

	// check len
	out = pb.prefix + countersBox + barBox + percentBox + speedBox + timeLeftBox + pb.postfix
	if len(out) < width {
		end = strings.Repeat(" ", width-len(out))
	}

	// and print!
	switch {
	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)
	}
}

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)
	if c != pb.currentValue {
		pb.write(c)
		pb.currentValue = c
	}
}

// Internal loop for writing progressbar
func (pb *ProgressBar) writer() {
	for {
		if atomic.LoadInt32(&pb.isFinish) != 0 {
			break
		}
		pb.Update()
		time.Sleep(pb.RefreshRate)
	}
}

type window struct {
	Row    uint16
	Col    uint16
	Xpixel uint16
	Ypixel uint16
}
<commit_msg>use channel to trigger isFinished<commit_after>package pb

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

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 uding int64 as total
func New64(total int64) *ProgressBar {
	pb := &ProgressBar{
		Total:         total,
		RefreshRate:   DEFAULT_REFRESH_RATE,
		ShowPercent:   true,
		ShowCounters:  true,
		ShowBar:       true,
		ShowTimeLeft:  true,
		ShowFinalTime: true,
		Units:         U_NO,
		ManualUpdate:  false,
		isFinish:      make(chan struct{}),
		currentValue:  -1,
	}
	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)

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

	isFinish     chan struct{}
	startTime    time.Time
	currentValue int64

	prefix, postfix string

	BarStart string
	BarEnd   string
	Empty    string
	Current  string
	CurrentN string
}

// Start print
func (pb *ProgressBar) Start() *ProgressBar {
	pb.startTime = time.Now()
	if pb.Total == 0 {
		pb.ShowBar = false
		pb.ShowTimeLeft = false
		pb.ShowPercent = false
	}
	if !pb.ManualUpdate {
		go pb.writer()
	}
	return pb
}

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

// Set current value
func (pb *ProgressBar) Set(current int) {
	atomic.StoreInt64(&pb.current, int64(current))
}

// 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.prefix = prefix
	return pb
}

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

// Set custom format for bar
// Example: bar.Format("[=>_]")
func (pb *ProgressBar) Format(format string) *ProgressBar {
	formatEntries := strings.Split(format, "")
	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() {
	close(pb.isFinish)
	pb.write(atomic.LoadInt64(&pb.current))
	if !pb.NotPrint {
		fmt.Println()
	}
}

// End print and write string 'str'
func (pb *ProgressBar) FinishPrint(str string) {
	pb.Finish()
	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
func (pb *ProgressBar) NewProxyReader(r io.Reader) *Reader {
	return &Reader{r, pb}
}

func (pb *ProgressBar) write(current int64) {
	width := pb.getWidth()

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

	// percents
	if pb.ShowPercent {
		percent := float64(current) / (float64(pb.Total) / float64(100))
		percentBox = fmt.Sprintf(" %#.02f %% ", percent)
	}

	// counters
	if pb.ShowCounters {
		if pb.Total > 0 {
			countersBox = fmt.Sprintf("%s / %s ", Format(current, pb.Units), Format(pb.Total, pb.Units))
		} else {
			countersBox = Format(current, pb.Units) + " "
		}
	}

	// time left
	fromStart := time.Now().Sub(pb.startTime)
	select {
	case <-pb.isFinish:
		if pb.ShowFinalTime {
			left := (fromStart / time.Second) * time.Second
			timeLeftBox = left.String()
		}
	default:
		if pb.ShowTimeLeft && current > 0 {
			perEntry := fromStart / time.Duration(current)
			left := time.Duration(pb.Total-current) * perEntry
			left = (left / time.Second) * time.Second
			timeLeftBox = left.String()
		}
	}

	// speed
	if pb.ShowSpeed && current > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		speed := float64(current) / (float64(fromStart) / float64(time.Second))
		speedBox = Format(int64(speed), pb.Units) + "/s "
	}

	// bar
	if pb.ShowBar {
		size := width - len(countersBox+pb.BarStart+pb.BarEnd+percentBox+timeLeftBox+speedBox+pb.prefix+pb.postfix)
		if size > 0 {
			curCount := int(math.Ceil((float64(current) / float64(pb.Total)) * float64(size)))
			emptCount := size - curCount
			barBox = pb.BarStart
			if emptCount < 0 {
				emptCount = 0
			}
			if curCount > size {
				curCount = size
			}
			if emptCount <= 0 {
				barBox += strings.Repeat(pb.Current, curCount)
			} else if curCount > 0 {
				barBox += strings.Repeat(pb.Current, curCount-1) + pb.CurrentN
			}

			barBox += strings.Repeat(pb.Empty, emptCount) + pb.BarEnd
		}
	}

	// check len
	out = pb.prefix + countersBox + barBox + percentBox + speedBox + timeLeftBox + pb.postfix
	if len(out) < width {
		end = strings.Repeat(" ", width-len(out))
	}

	// and print!
	switch {
	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)
	}
}

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)
	if c != pb.currentValue {
		pb.write(c)
		pb.currentValue = c
	}
}

// Internal loop for writing progressbar
func (pb *ProgressBar) writer() {
	pb.Update()
	for {
		select {
		case <-pb.isFinish:
			return
		case <-time.After(pb.RefreshRate):
			pb.Update()
		}
	}
}

type window struct {
	Row    uint16
	Col    uint16
	Xpixel uint16
	Ypixel uint16
}
<|endoftext|>


<commit_before>/*
	Reed-Solomon Codes over GF(2^8)
	Primitive Polynomial:  x^8+x^4+x^3+x^2+1
	Galois Filed arithmetic using Intel SIMD instructions (AVX2 or SSSE3)
*/

package reedsolomon

import (
	"errors"
)

// EncodeReconster implements for Reed-Solomon Encoding/Reconstructing
type EncodeReconster interface {
	Encode(vects [][]byte) error
	Reconstruct(vects [][]byte) error
	ReconstructData(vects [][]byte) error
}

func checkCfg(data, parity int) error {
	if (data <= 0) || (parity <= 0) {
		return errors.New("rs.New: data or parity <= 0")
	}
	if data+parity > 256 {
		return errors.New("rs.New: data+parity > 256")
	}
	return nil
}

// New create an EncodeReconster (vandermonde matrix as Encoding matrix)
func New(data, parity int) (enc EncodeReconster, err error) {
	err = checkCfg(data, parity)
	if err != nil {
		return
	}
	e, err := genEncMatrixVand(data, parity)
	if err != nil {
		return
	}
	return newRS(data, parity, e), nil
}

// NewCauchy create an EncodeReconster (cauchy matrix as Generator Matrix)
func NewCauchy(data, parity int) (enc EncodeReconster, err error) {
	err = checkCfg(data, parity)
	if err != nil {
		return
	}
	e := genEncMatrixCauchy(data, parity)
	return newRS(data, parity, e), nil
}

func checkER(d, p int, vs [][]byte, okNil bool) (size int, err error) {
	if d+p != len(vs) {
		err = errors.New("rs.checkER: vects not match rs args")
		return
	}
	for _, v := range vs {
		if len(v) != 0 {
			size = len(v)
		}
	}
	if size == 0 {
		err = errors.New("rs.checkER: vects size = 0")
		return
	}
	for _, v := range vs {
		if len(v) != size {
			if v == nil && okNil {
				continue
			}
			err = errors.New("rs.checkER: vects size mismatch")
			return
		}
	}
	return
}

func checkEnc(d, p int, vs [][]byte) (size int, err error) {
	return checkER(d, p, vs, false)
}

type encBase struct {
	data         int
	parity       int
	total        int // data+parity
	encodeMatrix matrix
	genMatrix    matrix
}

// Encode : multiply generator-matrix with data
func (e *encBase) Encode(vects [][]byte) (err error) {
	d := e.data
	p := e.parity
	_, err = checkEnc(d, p, vects)
	if err != nil {
		return
	}
	dv := vects[:d]
	pv := vects[d:]
	g := e.genMatrix
	for i := 0; i < d; i++ {
		for j := 0; j < p; j++ {
			if i != 0 {
				mulVectAdd(g[j*d+i], dv[i], pv[j])
			} else {
				mulVect(g[j*d], dv[0], pv[j])
			}
		}
	}
	return
}

func mulVect(c byte, a, b []byte) {
	t := mulTbl[c]
	for i := 0; i < len(a); i++ {
		b[i] = t[a[i]]
	}
}

func mulVectAdd(c byte, a, b []byte) {
	t := mulTbl[c]
	for i := 0; i < len(a); i++ {
		b[i] ^= t[a[i]]
	}
}

// Reconstruct : reconstruct lost data & parity
// set shard nil if lost
func (e *encBase) Reconstruct(vects [][]byte) (err error) {
	return e.reconst(vects, false)
}

// ReconstructData  : reconstruct lost data
func (e *encBase) ReconstructData(vects [][]byte) (err error) {
	return e.reconst(vects, true)
}

func checkReconst(d, p int, vs [][]byte) (size int, err error) {
	return checkER(d, p, vs, true)
}

func makeInverse(em matrix, has []int, data int) (matrix, error) {
	m := newMatrix(data, data)
	for i, p := range has {
		copy(m[i*data:i*data+data], em[p*data:p*data+data])
	}
	im, err := m.invert(data)
	if err != nil {
		return nil, err
	}
	return im, nil
}

func (e *encBase) reconst(vects [][]byte, dataOnly bool) (err error) {
	d := e.data
	p := e.parity
	size, err := checkReconst(d, p, vects)
	if err != nil {
		return
	}
	hasCnt := 0
	total := e.total
	k := 0
	dPos := make([]int, d) // for invert matrix
	dLost := make([]int, 0)
	pLost := make([]int, 0)
	for i := 0; i < total; i++ {
		if vects[i] != nil {
			hasCnt++
			if k < d {
				dPos[k] = i
				k++
			}
		} else {
			if i < d {
				dLost = append(dLost, i)
			} else {
				pLost = append(pLost, i)
			}
		}
	}
	if hasCnt == total {
		return nil
	}
	if hasCnt < d {
		return errors.New("rs.Reconst: not enough vects")
	}

	em := e.encodeMatrix
	dLCnt := len(dLost)
	if dLCnt != 0 {
		im, err2 := makeInverse(em, dPos, d)
		if err2 != nil {
			return err2
		}
		g := make([]byte, dLCnt*d)
		for i, p := range dLost {
			copy(g[i*d:i*d+d], im[p*d:p*d+d])
		}
		vtmp := make([][]byte, d+dLCnt)
		j := 0
		for i, v := range vects {
			if v != nil {
				if j < d {
					vtmp[j] = vects[i]
					j++
				}
			}
		}
		for _, i := range dLost {
			vects[i] = make([]byte, size)
			vtmp[j] = vects[i]
			j++
		}

		etmp := &encBase{data: d, parity: dLCnt, genMatrix: g}
		etmp.Encode(vtmp)
	}
	pLCnt := len(pLost)
	if pLCnt != 0 && !dataOnly {
		g := make([]byte, pLCnt*d)
		for i, p := range pLost {
			copy(g[i*d:i*d+d], em[p*d:p*d+d])
		}
		vtmp := make([][]byte, d+pLCnt)
		for i := 0; i < d; i++ {
			vtmp[i] = vects[i]
		}
		for i, p := range pLost {
			vects[p] = make([]byte, size)
			vtmp[d+i] = vects[p]
		}
		etmp := &encBase{data: d, parity: pLCnt, genMatrix: g}
		etmp.Encode(vtmp)
	}
	return nil
}
<commit_msg>encBase almost done<commit_after>/*
	Reed-Solomon Codes over GF(2^8)
	Primitive Polynomial:  x^8+x^4+x^3+x^2+1
	Galois Filed arithmetic using Intel SIMD instructions (AVX2 or SSSE3)
*/

package reedsolomon

import "errors"

// EncodeReconster implements for Reed-Solomon Encoding/Reconstructing
type EncodeReconster interface {
	Encode(vects [][]byte) error
	Reconstruct(vects [][]byte) error
	ReconstructData(vects [][]byte) error
}

func checkCfg(d, p int) error {
	if (d <= 0) || (p <= 0) {
		return errors.New("rs.New: data or parity <= 0")
	}
	if d+p > 256 {
		return errors.New("rs.New: data+parity > 256")
	}
	return nil
}

// New create an EncodeReconster (vandermonde matrix as Encoding matrix)
func New(data, parity int) (enc EncodeReconster, err error) {
	err = checkCfg(data, parity)
	if err != nil {
		return
	}
	e, err := genEncMatrixVand(data, parity)
	if err != nil {
		return
	}
	return newRS(data, parity, e), nil
}

// NewCauchy create an EncodeReconster (cauchy matrix as Generator Matrix)
func NewCauchy(data, parity int) (enc EncodeReconster, err error) {
	err = checkCfg(data, parity)
	if err != nil {
		return
	}
	e := genEncMatrixCauchy(data, parity)
	return newRS(data, parity, e), nil
}

func checkER(d, p int, vs [][]byte, okNil bool) (size int, err error) {
	if d+p != len(vs) {
		err = errors.New("rs.checkER: vects not match rs args")
		return
	}
	for _, v := range vs {
		if len(v) != 0 {
			size = len(v)
		}
	}
	if size == 0 {
		err = errors.New("rs.checkER: vects size = 0")
		return
	}
	for _, v := range vs {
		if len(v) != size {
			if v == nil && okNil {
				continue
			}
			err = errors.New("rs.checkER: vects size mismatch")
			return
		}
	}
	return
}

func checkEnc(d, p int, vs [][]byte) (size int, err error) {
	return checkER(d, p, vs, false)
}

type encBase struct {
	data         int
	parity       int
	total        int // data+parity
	encodeMatrix matrix
	genMatrix    matrix
}

// Encode : multiply generator-matrix with data
func (e *encBase) Encode(vects [][]byte) (err error) {
	d := e.data
	p := e.parity
	_, err = checkEnc(d, p, vects)
	if err != nil {
		return
	}
	dv := vects[:d]
	pv := vects[d:]
	g := e.genMatrix
	for i := 0; i < d; i++ {
		for j := 0; j < p; j++ {
			if i != 0 {
				mulVectAdd(g[j*d+i], dv[i], pv[j])
			} else {
				mulVect(g[j*d], dv[0], pv[j])
			}
		}
	}
	return
}

func mulVect(c byte, a, b []byte) {
	t := mulTbl[c]
	for i := 0; i < len(a); i++ {
		b[i] = t[a[i]]
	}
}

func mulVectAdd(c byte, a, b []byte) {
	t := mulTbl[c]
	for i := 0; i < len(a); i++ {
		b[i] ^= t[a[i]]
	}
}

// Reconstruct : reconstruct lost data & parity
// set shard nil if lost
func (e *encBase) Reconstruct(vects [][]byte) (err error) {
	return e.reconst(vects, false)
}

// ReconstructData  : reconstruct lost data
func (e *encBase) ReconstructData(vects [][]byte) (err error) {
	return e.reconst(vects, true)
}

func checkReconst(d, p int, vs [][]byte) (size int, err error) {
	return checkER(d, p, vs, true)
}

func makeInverse(em matrix, has []int, data int) (matrix, error) {
	m := newMatrix(data, data)
	for i, p := range has {
		copy(m[i*data:i*data+data], em[p*data:p*data+data])
	}
	im, err := m.invert(data)
	if err != nil {
		return nil, err
	}
	return im, nil
}

func (e *encBase) reconst(vects [][]byte, dataOnly bool) (err error) {
	d := e.data
	p := e.parity
	total := e.total
	size, err := checkReconst(d, p, vects)
	if err != nil {
		return
	}
	has := 0
	k := 0
	dPos := make([]int, d) // for invert matrix
	dLost := make([]int, 0)
	pLost := make([]int, 0)
	for i := 0; i < total; i++ {
		if vects[i] != nil {
			has++
			if k < d {
				dPos[k] = i
				k++
			}
		} else {
			if i < d {
				dLost = append(dLost, i)
			} else {
				pLost = append(pLost, i)
			}
		}
	}
	if has == total {
		return nil
	}
	if has < d {
		return errors.New("rs.Reconst: not enough vects")
	}

	em := e.encodeMatrix
	dLCnt := len(dLost)
	if dLCnt != 0 {
		im, err2 := makeInverse(em, dPos, d)
		if err2 != nil {
			return err2
		}
		g := make([]byte, dLCnt*d)
		for i, p := range dLost {
			copy(g[i*d:i*d+d], im[p*d:p*d+d])
		}
		vtmp := make([][]byte, d+dLCnt)
		j := 0
		for i, v := range vects {
			if v != nil {
				if j < d {
					vtmp[j] = vects[i]
					j++
				}
			}
		}
		for _, i := range dLost {
			vects[i] = make([]byte, size)
			vtmp[j] = vects[i]
			j++
		}

		etmp := &encBase{data: d, parity: dLCnt, genMatrix: g}
		etmp.Encode(vtmp)
	}
	pLCnt := len(pLost)
	if pLCnt != 0 && !dataOnly {
		g := make([]byte, pLCnt*d)
		for i, p := range pLost {
			copy(g[i*d:i*d+d], em[p*d:p*d+d])
		}
		vtmp := make([][]byte, d+pLCnt)
		for i := 0; i < d; i++ {
			vtmp[i] = vects[i]
		}
		for i, p := range pLost {
			vects[p] = make([]byte, size)
			vtmp[d+i] = vects[p]
		}
		etmp := &encBase{data: d, parity: pLCnt, genMatrix: g}
		etmp.Encode(vtmp)
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"regexp"
	"strconv"
)

var hexRegex = regexp.MustCompile(`^[a-f\d]{6}$`)

var colors = map[string]string{
	"bg-none":    "0",
	"fg-black":   "30",
	"fg-red":     "31",
	"fg-green":   "32",
	"fg-yellow":  "33",
	"fg-blue":    "34",
	"fg-magenta": "35",
	"fg-cyan":    "36",
	"fg-white":   "37",
	"bg-black":   "40",
	"bg-red":     "41",
	"bg-green":   "42",
	"bg-yellow":  "43",
	"bg-blue":    "44",
	"bg-magenta": "45",
	"bg-cyan":    "46",
	"bg-white":   "47",
}

func escapeBackground(color string) string {
	switch shell {
	case "zsh":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "%{\x1b[48;2;" + escapeHex(color) + "m%}"
		}

		// 16 and 256 colors
		return "%K{" + color + "}"
	case "bash":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\\[\x1b[48;2;" + escapeHex(color) + "m\\]"
		}

		// 16 colors
		code, ok := colors["bg-"+color]
		if !ok {
			dief("invalid background color: %q", color)
		}
		return "\\[\x1b[" + code + "m\\]"
	default:
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\x1b[48;2;" + escapeHex(color) + "m"
		}

		// 256 colors
		_, err := strconv.Atoi(color)
		if err == nil {
			return "\x1b[48;5;" + color + "m"
		}

		// 16 colors
		code, ok := colors["bg-"+color]
		if !ok {
			dief("invalid background color: %q", color)
		}
		return "\x1b[" + code + "m"
	}
}

func escapeForeground(color string) string {
	switch shell {
	case "zsh":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "%{\x1b[38;2;" + escapeHex(color) + "m%}"
		}

		// 16 and 256 colors
		return "%F{" + color + "}"
	case "bash":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\\[\x1b[38;2;" + escapeHex(color) + "m\\]"
		}

		// 16 colors
		code, ok := colors["fg-"+color]
		if !ok {
			dief("invalid foreground color: %q", color)
		}
		return "\\[\x1b[" + code + "m\\]"
	default:
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\x1b[38;2;" + escapeHex(color) + "m"
		}

		// 256 colors
		_, err := strconv.Atoi(color)
		if err == nil {
			return "\x1b[38;5;" + color + "m"
		}

		// 16 colors
		code, ok := colors["fg-"+color]
		if !ok {
			dief("invalid foreground color: %q", color)
		}
		return "\x1b[" + code + "m"
	}
}

func resetColors() {
	switch shell {
	case "zsh":
		fmt.Print("%{%f%}")
	case "bash":
		fmt.Print("\\[\x1b[0m\\]")
	default:
		fmt.Print("\x1b[0m")
	}
}

func escapeHex(color string) string {
	var r, g, b byte
	fmt.Sscanf(color, "%02x%02x%02x", &r, &g, &b)
	return fmt.Sprintf("%d;%d;%d", r, g, b)
}
<commit_msg>added 256 color support for bash<commit_after>package main

import (
	"fmt"
	"regexp"
	"strconv"
)

var hexRegex = regexp.MustCompile(`^[a-f\d]{6}$`)

var colors = map[string]string{
	"bg-none":    "0",
	"fg-black":   "30",
	"fg-red":     "31",
	"fg-green":   "32",
	"fg-yellow":  "33",
	"fg-blue":    "34",
	"fg-magenta": "35",
	"fg-cyan":    "36",
	"fg-white":   "37",
	"bg-black":   "40",
	"bg-red":     "41",
	"bg-green":   "42",
	"bg-yellow":  "43",
	"bg-blue":    "44",
	"bg-magenta": "45",
	"bg-cyan":    "46",
	"bg-white":   "47",
}

func escapeBackground(color string) string {
	switch shell {
	case "zsh":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "%{\x1b[48;2;" + escapeHex(color) + "m%}"
		}

		// 16 and 256 colors
		return "%K{" + color + "}"
	case "bash":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\\[\x1b[48;2;" + escapeHex(color) + "m\\]"
		}

		// 256 colors
		_, err := strconv.Atoi(color)
		if err == nil {
			return "\\[\x1b[48;5;" + color + "m\\]"
		}

		// 16 colors
		code, ok := colors["bg-"+color]
		if !ok {
			dief("invalid background color: %q", color)
		}
		return "\\[\x1b[" + code + "m\\]"
	default:
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\x1b[48;2;" + escapeHex(color) + "m"
		}

		// 256 colors
		_, err := strconv.Atoi(color)
		if err == nil {
			return "\x1b[48;5;" + color + "m"
		}

		// 16 colors
		code, ok := colors["bg-"+color]
		if !ok {
			dief("invalid background color: %q", color)
		}
		return "\x1b[" + code + "m"
	}
}

func escapeForeground(color string) string {
	switch shell {
	case "zsh":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "%{\x1b[38;2;" + escapeHex(color) + "m%}"
		}

		// 16 and 256 colors
		return "%F{" + color + "}"
	case "bash":
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\\[\x1b[38;2;" + escapeHex(color) + "m\\]"
		}

		// 256 colors
		_, err := strconv.Atoi(color)
		if err == nil {
			return "\x1b[38;5;" + color + "m"
		}

		// 16 colors
		code, ok := colors["fg-"+color]
		if !ok {
			dief("invalid foreground color: %q", color)
		}
		return "\\[\x1b[" + code + "m\\]"
	default:
		// 24-bit color
		if hexRegex.MatchString(color) {
			return "\x1b[38;2;" + escapeHex(color) + "m"
		}

		// 256 colors
		_, err := strconv.Atoi(color)
		if err == nil {
			return "\x1b[38;5;" + color + "m"
		}

		// 16 colors
		code, ok := colors["fg-"+color]
		if !ok {
			dief("invalid foreground color: %q", color)
		}
		return "\x1b[" + code + "m"
	}
}

func resetColors() {
	switch shell {
	case "zsh":
		fmt.Print("%{%f%}")
	case "bash":
		fmt.Print("\\[\x1b[0m\\]")
	default:
		fmt.Print("\x1b[0m")
	}
}

func escapeHex(color string) string {
	var r, g, b byte
	fmt.Sscanf(color, "%02x%02x%02x", &r, &g, &b)
	return fmt.Sprintf("%d;%d;%d", r, g, b)
}
<|endoftext|>


<commit_before>// MIT License
//
// Copyright (c) 2017 Pressly Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package saml

import (
	"encoding/base64"
	"encoding/pem"
	"encoding/xml"
	"errors"
	"io/ioutil"
	"net/http"
	"os"
	"sync/atomic"
)

// ServiceProvider represents a service provider.
type ServiceProvider struct {
	IdPMetadataURL string
	IdPMetadataXML []byte
	IdPMetadata    *Metadata

	KeyFile  string
	CertFile string

	PrivkeyPEM string
	PubkeyPEM  string

	MetadataURL string
	AcsURL      string

	AllowIdpInitiated bool

	SecurityOpts

	pemCert atomic.Value
}

// PrivkeyFile returns a physical path where the SP's key can be accessed.
func (sp *ServiceProvider) PrivkeyFile() (string, error) {
	if sp.KeyFile != "" {
		return sp.KeyFile, nil
	}
	if sp.PrivkeyPEM != "" {
		return writeFile([]byte(sp.PrivkeyPEM))
	}
	return "", errors.New("No private key given.")
}

// PubkeyFile returns a physical path where the SP's public certificate can be
// accessed.
func (sp *ServiceProvider) PubkeyFile() (string, error) {
	if sp.CertFile != "" {
		return validateKeyFile(sp.CertFile, nil)
	}
	if sp.PubkeyPEM != "" {
		return validateKeyFile(writeFile([]byte(sp.PubkeyPEM)))
	}
	return "", errors.New("No public key given.")
}

// GetIdPAuthResource returns the authentication URL for the SP.
func (sp *ServiceProvider) GetIdPAuthResource() (string, error) {
	meta, err := sp.GetIdPMetadata()
	if err != nil {
		return "", err
	}

	for _, endpoint := range meta.IDPSSODescriptor.SingleSignOnService {
		if endpoint.Binding == "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect" {
			return endpoint.Location, nil
		}
		if endpoint.Binding == "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST" {
			return endpoint.Location, nil
		}
	}

	return "", errors.New("Could not find resource")
}

// GetIdPCertFile returns a physical path where the IdP certificate can be
// accessed.
func (sp *ServiceProvider) GetIdPCertFile() (string, error) {
	meta, err := sp.GetIdPMetadata()
	if err != nil {
		return "", err
	}

	cert := ""
	for _, keyDescriptor := range meta.IDPSSODescriptor.KeyDescriptor {
		if keyDescriptor.Use == "encryption" {
			cert = keyDescriptor.KeyInfo.Certificate
			break
		}
	}

	if cert == "" {
		for _, keyDescriptor := range meta.IDPSSODescriptor.KeyDescriptor {
			if keyDescriptor.KeyInfo.Certificate != "" {
				cert = keyDescriptor.KeyInfo.Certificate
				break
			}
		}
	}

	if cert == "" {
		return "", errors.New("Missing certificate data.")
	}

	certBytes, _ := base64.StdEncoding.DecodeString(cert)

	certBytes = pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certBytes,
	})

	return writeFile(certBytes)
}

// GetIdPMetadata returns the IdP metadata value.
func (sp *ServiceProvider) GetIdPMetadata() (*Metadata, error) {
	if sp.IdPMetadata != nil {
		m := *(sp.IdPMetadata)
		return &m, nil
	}

	if len(sp.IdPMetadataXML) == 0 {
		if sp.IdPMetadataURL == "" {
			return nil, errors.New("Missing metadata URL.")
		}

		res, err := http.Get(sp.IdPMetadataURL)
		if err != nil {
			return nil, err
		}
		defer res.Body.Close()

		buf, err := ioutil.ReadAll(res.Body)
		if err != nil {
			return nil, err
		}

		sp.IdPMetadataXML = buf
	}

	var metadata Metadata
	err := xml.Unmarshal(sp.IdPMetadataXML, &metadata)
	if err != nil {
		return nil, err
	}

	sp.IdPMetadata = &metadata
	return &metadata, nil
}

// Cert returns a *pem.Block value that corresponds to the SP's certificate.
func (sp *ServiceProvider) Cert() (*pem.Block, error) {
	if v := sp.pemCert.Load(); v != nil {
		return v.(*pem.Block), nil
	}

	certFile, err := sp.PubkeyFile()
	if err != nil {
		return nil, err
	}

	fp, err := os.Open(certFile)
	if err != nil {
		return nil, err
	}
	defer fp.Close()

	buf, err := ioutil.ReadAll(fp)
	if err != nil {
		return nil, err
	}

	cert, _ := pem.Decode(buf)
	if cert == nil {
		return nil, errors.New("Invalid certificate.")
	}

	sp.pemCert.Store(cert)

	return cert, nil
}

// Metadata returns a metadata value based on the SP's data.
func (sp *ServiceProvider) Metadata() (*Metadata, error) {
	cert, err := sp.Cert()
	if err != nil {
		return nil, err
	}
	certStr := base64.StdEncoding.EncodeToString(cert.Bytes)

	metadata := &Metadata{
		EntityID:   sp.MetadataURL,
		ValidUntil: Now().Add(defaultValidDuration),
		SPSSODescriptor: &SPSSODescriptor{
			AuthnRequestsSigned:        false,
			WantAssertionsSigned:       true,
			ProtocolSupportEnumeration: "urn:oasis:names:tc:SAML:2.0:protocol",
			KeyDescriptor: []KeyDescriptor{
				KeyDescriptor{
					Use: "signing",
					KeyInfo: KeyInfo{
						Certificate: certStr,
					},
				},
				KeyDescriptor{
					Use: "encryption",
					KeyInfo: KeyInfo{
						Certificate: certStr,
					},
					EncryptionMethods: []EncryptionMethod{
						EncryptionMethod{Algorithm: "http://www.w3.org/2001/04/xmlenc#aes128-cbc"},
						EncryptionMethod{Algorithm: "http://www.w3.org/2001/04/xmlenc#aes192-cbc"},
						EncryptionMethod{Algorithm: "http://www.w3.org/2001/04/xmlenc#aes256-cbc"},
						EncryptionMethod{Algorithm: "http://www.w3.org/2001/04/xmlenc#rsa-oaep-mgf1p"},
					},
				},
			},
			AssertionConsumerService: []IndexedEndpoint{{
				Binding:  HTTPPostBinding,
				Location: sp.AcsURL,
				Index:    1,
			}},
		},
	}

	return metadata, nil
}

// MakeAuthenticationRequest produces a new AuthnRequest object for the given
// idpURL.
func (sp *ServiceProvider) MakeAuthenticationRequest(idpURL string) (*AuthnRequest, error) {
	req := AuthnRequest{
		AssertionConsumerServiceURL: sp.AcsURL,
		Destination:                 idpURL,
		ID:                          NewID(),
		IssueInstant:                Now(),
		Version:                     "2.0",
		Issuer: Issuer{
			Format: "urn:oasis:names:tc:SAML:2.0:nameid-format:entity",
			Value:  sp.MetadataURL,
		},
		NameIDPolicy: NameIDPolicy{
			AllowCreate: true,
			// TODO(ross): figure out exactly policy we need
			// urn:mace:shibboleth:1.0:nameIdentifier
			// urn:oasis:names:tc:SAML:2.0:nameid-format:transient
			Format: "urn:oasis:names:tc:SAML:2.0:nameid-format:transient",
		},
	}
	return &req, nil
}
<commit_msg>saml: fix nil pointer dereference<commit_after>// MIT License
//
// Copyright (c) 2017 Pressly Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package saml

import (
	"encoding/base64"
	"encoding/pem"
	"encoding/xml"
	"errors"
	"io/ioutil"
	"net/http"
	"os"
	"sync/atomic"
)

// ServiceProvider represents a service provider.
type ServiceProvider struct {
	IdPMetadataURL string
	IdPMetadataXML []byte
	IdPMetadata    *Metadata

	KeyFile  string
	CertFile string

	PrivkeyPEM string
	PubkeyPEM  string

	MetadataURL string
	AcsURL      string

	AllowIdpInitiated bool

	SecurityOpts

	pemCert atomic.Value
}

// PrivkeyFile returns a physical path where the SP's key can be accessed.
func (sp *ServiceProvider) PrivkeyFile() (string, error) {
	if sp.KeyFile != "" {
		return sp.KeyFile, nil
	}
	if sp.PrivkeyPEM != "" {
		return writeFile([]byte(sp.PrivkeyPEM))
	}
	return "", errors.New("No private key given.")
}

// PubkeyFile returns a physical path where the SP's public certificate can be
// accessed.
func (sp *ServiceProvider) PubkeyFile() (string, error) {
	if sp.CertFile != "" {
		return validateKeyFile(sp.CertFile, nil)
	}
	if sp.PubkeyPEM != "" {
		return validateKeyFile(writeFile([]byte(sp.PubkeyPEM)))
	}
	return "", errors.New("No public key given.")
}

// GetIdPAuthResource returns the authentication URL for the SP.
func (sp *ServiceProvider) GetIdPAuthResource() (string, error) {
	meta, err := sp.GetIdPMetadata()
	if err != nil {
		return "", err
	}

	if meta.IDPSSODescriptor == nil {
		return "", errors.New("could not find IDPSSODescriptor")
	}

	for _, endpoint := range meta.IDPSSODescriptor.SingleSignOnService {
		if endpoint.Binding == "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect" {
			return endpoint.Location, nil
		}
		if endpoint.Binding == "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST" {
			return endpoint.Location, nil
		}
	}

	return "", errors.New("could not find SingleSignOnService")
}

// GetIdPCertFile returns a physical path where the IdP certificate can be
// accessed.
func (sp *ServiceProvider) GetIdPCertFile() (string, error) {
	meta, err := sp.GetIdPMetadata()
	if err != nil {
		return "", err
	}

	cert := ""
	for _, keyDescriptor := range meta.IDPSSODescriptor.KeyDescriptor {
		if keyDescriptor.Use == "encryption" {
			cert = keyDescriptor.KeyInfo.Certificate
			break
		}
	}

	if cert == "" {
		for _, keyDescriptor := range meta.IDPSSODescriptor.KeyDescriptor {
			if keyDescriptor.KeyInfo.Certificate != "" {
				cert = keyDescriptor.KeyInfo.Certificate
				break
			}
		}
	}

	if cert == "" {
		return "", errors.New("Missing certificate data.")
	}

	certBytes, _ := base64.StdEncoding.DecodeString(cert)

	certBytes = pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certBytes,
	})

	return writeFile(certBytes)
}

// GetIdPMetadata returns the IdP metadata value.
func (sp *ServiceProvider) GetIdPMetadata() (*Metadata, error) {
	if sp.IdPMetadata != nil {
		m := *(sp.IdPMetadata)
		return &m, nil
	}

	if len(sp.IdPMetadataXML) == 0 {
		if sp.IdPMetadataURL == "" {
			return nil, errors.New("Missing metadata URL.")
		}

		res, err := http.Get(sp.IdPMetadataURL)
		if err != nil {
			return nil, err
		}
		defer res.Body.Close()

		buf, err := ioutil.ReadAll(res.Body)
		if err != nil {
			return nil, err
		}

		sp.IdPMetadataXML = buf
	}

	var metadata Metadata
	err := xml.Unmarshal(sp.IdPMetadataXML, &metadata)
	if err != nil {
		return nil, err
	}

	sp.IdPMetadata = &metadata
	return &metadata, nil
}

// Cert returns a *pem.Block value that corresponds to the SP's certificate.
func (sp *ServiceProvider) Cert() (*pem.Block, error) {
	if v := sp.pemCert.Load(); v != nil {
		return v.(*pem.Block), nil
	}

	certFile, err := sp.PubkeyFile()
	if err != nil {
		return nil, err
	}

	fp, err := os.Open(certFile)
	if err != nil {
		return nil, err
	}
	defer fp.Close()

	buf, err := ioutil.ReadAll(fp)
	if err != nil {
		return nil, err
	}

	cert, _ := pem.Decode(buf)
	if cert == nil {
		return nil, errors.New("Invalid certificate.")
	}

	sp.pemCert.Store(cert)

	return cert, nil
}

// Metadata returns a metadata value based on the SP's data.
func (sp *ServiceProvider) Metadata() (*Metadata, error) {
	cert, err := sp.Cert()
	if err != nil {
		return nil, err
	}
	certStr := base64.StdEncoding.EncodeToString(cert.Bytes)

	metadata := &Metadata{
		EntityID:   sp.MetadataURL,
		ValidUntil: Now().Add(defaultValidDuration),
		SPSSODescriptor: &SPSSODescriptor{
			AuthnRequestsSigned:        false,
			WantAssertionsSigned:       true,
			ProtocolSupportEnumeration: "urn:oasis:names:tc:SAML:2.0:protocol",
			KeyDescriptor: []KeyDescriptor{
				KeyDescriptor{
					Use: "signing",
					KeyInfo: KeyInfo{
						Certificate: certStr,
					},
				},
				KeyDescriptor{
					Use: "encryption",
					KeyInfo: KeyInfo{
						Certificate: certStr,
					},
					EncryptionMethods: []EncryptionMethod{
						EncryptionMethod{Algorithm: "http://www.w3.org/2001/04/xmlenc#aes128-cbc"},
						EncryptionMethod{Algorithm: "http://www.w3.org/2001/04/xmlenc#aes192-cbc"},
						EncryptionMethod{Algorithm: "http://www.w3.org/2001/04/xmlenc#aes256-cbc"},
						EncryptionMethod{Algorithm: "http://www.w3.org/2001/04/xmlenc#rsa-oaep-mgf1p"},
					},
				},
			},
			AssertionConsumerService: []IndexedEndpoint{{
				Binding:  HTTPPostBinding,
				Location: sp.AcsURL,
				Index:    1,
			}},
		},
	}

	return metadata, nil
}

// MakeAuthenticationRequest produces a new AuthnRequest object for the given
// idpURL.
func (sp *ServiceProvider) MakeAuthenticationRequest(idpURL string) (*AuthnRequest, error) {
	req := AuthnRequest{
		AssertionConsumerServiceURL: sp.AcsURL,
		Destination:                 idpURL,
		ID:                          NewID(),
		IssueInstant:                Now(),
		Version:                     "2.0",
		Issuer: Issuer{
			Format: "urn:oasis:names:tc:SAML:2.0:nameid-format:entity",
			Value:  sp.MetadataURL,
		},
		NameIDPolicy: NameIDPolicy{
			AllowCreate: true,
			// TODO(ross): figure out exactly policy we need
			// urn:mace:shibboleth:1.0:nameIdentifier
			// urn:oasis:names:tc:SAML:2.0:nameid-format:transient
			Format: "urn:oasis:names:tc:SAML:2.0:nameid-format:transient",
		},
	}
	return &req, nil
}
<|endoftext|>


<commit_before>package stomp

import (
	"errors"
	"io"
)

var ErrTxDone = errors.New("stomp: transaction has already been committed or aborted")

type Tx struct {
	tid       string
	done      bool
	receipts  *receipts
	transport *Transport
}

func (t *Tx) Commit(receipt bool) error {
	if t.done {
		return ErrTxDone
	}
	defer func() {
		t.done = true
	}()

	if receipt {
		return doWithReceipt(t.receipts, func(rid string) error {
			return t.transport.TxCommit(t.tid, &rid)
		})
	}
	return t.transport.TxCommit(t.tid, nil)
}

func (t *Tx) Abort(receipt bool) error {
	if t.done {
		return nil
	}
	defer func() {
		t.done = true
	}()

	if receipt {
		return doWithReceipt(t.receipts, func(rid string) error {
			return t.transport.TxAbort(t.tid, &rid)
		})
	}
	return t.transport.TxAbort(t.tid, nil)
}

func (t *Tx) TxSend(dest string, hdrs *map[string]string, bodyType string, body io.Reader) error {
	if t.done {
		return ErrTxDone
	}
	return t.transport.TxSend(t.tid, dest, hdrs, bodyType, body)
}

func (t *Tx) Ack(id string) error {
	if t.done {
		return ErrTxDone
	}
	return t.transport.TxAck(t.tid, id)
}

func (t *Tx) Nack(id string) error {
	if t.done {
		return ErrTxDone
	}
	return t.transport.TxNack(t.tid, id)
}
<commit_msg>Renmae tx send.<commit_after>package stomp

import (
	"errors"
	"io"
)

var ErrTxDone = errors.New("stomp: transaction has already been committed or aborted")

type Tx struct {
	tid       string
	done      bool
	receipts  *receipts
	transport *Transport
}

func (t *Tx) Commit(receipt bool) error {
	if t.done {
		return ErrTxDone
	}
	defer func() {
		t.done = true
	}()

	if receipt {
		return doWithReceipt(t.receipts, func(rid string) error {
			return t.transport.TxCommit(t.tid, &rid)
		})
	}
	return t.transport.TxCommit(t.tid, nil)
}

func (t *Tx) Abort(receipt bool) error {
	if t.done {
		return nil
	}
	defer func() {
		t.done = true
	}()

	if receipt {
		return doWithReceipt(t.receipts, func(rid string) error {
			return t.transport.TxAbort(t.tid, &rid)
		})
	}
	return t.transport.TxAbort(t.tid, nil)
}

func (t *Tx) Send(dest string, hdrs *map[string]string, bodyType string, body io.Reader) error {
	if t.done {
		return ErrTxDone
	}
	return t.transport.TxSend(t.tid, dest, hdrs, bodyType, body)
}

func (t *Tx) Ack(id string) error {
	if t.done {
		return ErrTxDone
	}
	return t.transport.TxAck(t.tid, id)
}

func (t *Tx) Nack(id string) error {
	if t.done {
		return ErrTxDone
	}
	return t.transport.TxNack(t.tid, id)
}
<|endoftext|>


<commit_before>package pgx

import (
	"bytes"
	"context"
	"fmt"

	"github.com/jackc/pgconn"
	errors "golang.org/x/xerrors"
)

type TxIsoLevel string

// Transaction isolation levels
const (
	Serializable    = TxIsoLevel("serializable")
	RepeatableRead  = TxIsoLevel("repeatable read")
	ReadCommitted   = TxIsoLevel("read committed")
	ReadUncommitted = TxIsoLevel("read uncommitted")
)

type TxAccessMode string

// Transaction access modes
const (
	ReadWrite = TxAccessMode("read write")
	ReadOnly  = TxAccessMode("read only")
)

type TxDeferrableMode string

// Transaction deferrable modes
const (
	Deferrable    = TxDeferrableMode("deferrable")
	NotDeferrable = TxDeferrableMode("not deferrable")
)

const (
	TxStatusInProgress      = 0
	TxStatusCommitFailure   = -1
	TxStatusRollbackFailure = -2
	TxStatusInFailure       = -3
	TxStatusCommitSuccess   = 1
	TxStatusRollbackSuccess = 2
)

type TxOptions struct {
	IsoLevel       TxIsoLevel
	AccessMode     TxAccessMode
	DeferrableMode TxDeferrableMode
}

func (txOptions TxOptions) beginSQL() string {
	buf := &bytes.Buffer{}
	buf.WriteString("begin")
	if txOptions.IsoLevel != "" {
		fmt.Fprintf(buf, " isolation level %s", txOptions.IsoLevel)
	}
	if txOptions.AccessMode != "" {
		fmt.Fprintf(buf, " %s", txOptions.AccessMode)
	}
	if txOptions.DeferrableMode != "" {
		fmt.Fprintf(buf, " %s", txOptions.DeferrableMode)
	}

	return buf.String()
}

var ErrTxClosed = errors.New("tx is closed")
var ErrTxInFailure = errors.New("tx failed")

// ErrTxCommitRollback occurs when an error has occurred in a transaction and
// Commit() is called. PostgreSQL accepts COMMIT on aborted transactions, but
// it is treated as ROLLBACK.
var ErrTxCommitRollback = errors.New("commit unexpectedly resulted in rollback")

// Begin starts a transaction. Unlike database/sql, the context only affects the begin command. i.e. there is no
// auto-rollback on context cancelation.
func (c *Conn) Begin(ctx context.Context) (*Tx, error) {
	return c.BeginEx(ctx, TxOptions{})
}

// BeginEx starts a transaction with txOptions determining the transaction mode. Unlike database/sql, the context only
// affects the begin command. i.e. there is no auto-rollback on context cancelation.
func (c *Conn) BeginEx(ctx context.Context, txOptions TxOptions) (*Tx, error) {
	_, err := c.Exec(ctx, txOptions.beginSQL())
	if err != nil {
		// begin should never fail unless there is an underlying connection issue or
		// a context timeout. In either case, the connection is possibly broken.
		c.die(errors.New("failed to begin transaction"))
		return nil, err
	}

	return &Tx{conn: c}, nil
}

// Tx represents a database transaction.
//
// All Tx methods return ErrTxClosed if Commit or Rollback has already been
// called on the Tx.
type Tx struct {
	conn   *Conn
	err    error
	status int8
}

// Commit commits the transaction.
func (tx *Tx) Commit(ctx context.Context) error {
	if tx.status != TxStatusInProgress {
		return ErrTxClosed
	}

	commandTag, err := tx.conn.Exec(ctx, "commit")
	if err == nil && string(commandTag) == "COMMIT" {
		tx.status = TxStatusCommitSuccess
	} else if err == nil && string(commandTag) == "ROLLBACK" {
		tx.status = TxStatusCommitFailure
		tx.err = ErrTxCommitRollback
	} else {
		tx.status = TxStatusCommitFailure
		tx.err = err
		// A commit failure leaves the connection in an undefined state
		tx.conn.die(errors.New("commit failed"))
	}

	return tx.err
}

// Rollback rolls back the transaction. Rollback will return ErrTxClosed if the
// Tx is already closed, but is otherwise safe to call multiple times. Hence, a
// defer tx.Rollback() is safe even if tx.Commit() will be called first in a
// non-error condition.
func (tx *Tx) Rollback(ctx context.Context) error {
	if tx.status != TxStatusInProgress {
		return ErrTxClosed
	}

	_, tx.err = tx.conn.Exec(ctx, "rollback")
	if tx.err == nil {
		tx.status = TxStatusRollbackSuccess
	} else {
		tx.status = TxStatusRollbackFailure
		// A rollback failure leaves the connection in an undefined state
		tx.conn.die(errors.New("rollback failed"))
	}

	return tx.err
}

// Exec delegates to the underlying *Conn
func (tx *Tx) Exec(ctx context.Context, sql string, arguments ...interface{}) (commandTag pgconn.CommandTag, err error) {
	return tx.conn.Exec(ctx, sql, arguments...)
}

// Prepare delegates to the underlying *Conn
func (tx *Tx) Prepare(ctx context.Context, name, sql string) (*PreparedStatement, error) {
	if tx.status != TxStatusInProgress {
		return nil, ErrTxClosed
	}

	return tx.conn.Prepare(ctx, name, sql)
}

// Query delegates to the underlying *Conn
func (tx *Tx) Query(ctx context.Context, sql string, args ...interface{}) (Rows, error) {
	if tx.status != TxStatusInProgress {
		// Because checking for errors can be deferred to the *Rows, build one with the error
		err := ErrTxClosed
		return &connRows{closed: true, err: err}, err
	}

	return tx.conn.Query(ctx, sql, args...)
}

// QueryRow delegates to the underlying *Conn
func (tx *Tx) QueryRow(ctx context.Context, sql string, args ...interface{}) Row {
	rows, _ := tx.Query(ctx, sql, args...)
	return (*connRow)(rows.(*connRows))
}

// CopyFrom delegates to the underlying *Conn
func (tx *Tx) CopyFrom(ctx context.Context, tableName Identifier, columnNames []string, rowSrc CopyFromSource) (int64, error) {
	if tx.status != TxStatusInProgress {
		return 0, ErrTxClosed
	}

	return tx.conn.CopyFrom(ctx, tableName, columnNames, rowSrc)
}

// SendBatch delegates to the underlying *Conn
func (tx *Tx) SendBatch(ctx context.Context, b *Batch) BatchResults {
	if tx.status != TxStatusInProgress {
		return &batchResults{err: ErrTxClosed}
	}

	return tx.conn.SendBatch(ctx, b)
}

// Status returns the status of the transaction from the set of
// pgx.TxStatus* constants.
func (tx *Tx) Status() int8 {
	if tx.status == TxStatusInProgress && tx.conn.pgConn.TxStatus == 'E' {
		return TxStatusInFailure
	}
	return tx.status
}

// Err returns the final error state, if any, of calling Commit or Rollback.
func (tx *Tx) Err() error {
	return tx.err
}
<commit_msg>Remove unused exported error value<commit_after>package pgx

import (
	"bytes"
	"context"
	"fmt"

	"github.com/jackc/pgconn"
	errors "golang.org/x/xerrors"
)

type TxIsoLevel string

// Transaction isolation levels
const (
	Serializable    = TxIsoLevel("serializable")
	RepeatableRead  = TxIsoLevel("repeatable read")
	ReadCommitted   = TxIsoLevel("read committed")
	ReadUncommitted = TxIsoLevel("read uncommitted")
)

type TxAccessMode string

// Transaction access modes
const (
	ReadWrite = TxAccessMode("read write")
	ReadOnly  = TxAccessMode("read only")
)

type TxDeferrableMode string

// Transaction deferrable modes
const (
	Deferrable    = TxDeferrableMode("deferrable")
	NotDeferrable = TxDeferrableMode("not deferrable")
)

const (
	TxStatusInProgress      = 0
	TxStatusCommitFailure   = -1
	TxStatusRollbackFailure = -2
	TxStatusInFailure       = -3
	TxStatusCommitSuccess   = 1
	TxStatusRollbackSuccess = 2
)

type TxOptions struct {
	IsoLevel       TxIsoLevel
	AccessMode     TxAccessMode
	DeferrableMode TxDeferrableMode
}

func (txOptions TxOptions) beginSQL() string {
	buf := &bytes.Buffer{}
	buf.WriteString("begin")
	if txOptions.IsoLevel != "" {
		fmt.Fprintf(buf, " isolation level %s", txOptions.IsoLevel)
	}
	if txOptions.AccessMode != "" {
		fmt.Fprintf(buf, " %s", txOptions.AccessMode)
	}
	if txOptions.DeferrableMode != "" {
		fmt.Fprintf(buf, " %s", txOptions.DeferrableMode)
	}

	return buf.String()
}

var ErrTxClosed = errors.New("tx is closed")

// ErrTxCommitRollback occurs when an error has occurred in a transaction and
// Commit() is called. PostgreSQL accepts COMMIT on aborted transactions, but
// it is treated as ROLLBACK.
var ErrTxCommitRollback = errors.New("commit unexpectedly resulted in rollback")

// Begin starts a transaction. Unlike database/sql, the context only affects the begin command. i.e. there is no
// auto-rollback on context cancelation.
func (c *Conn) Begin(ctx context.Context) (*Tx, error) {
	return c.BeginEx(ctx, TxOptions{})
}

// BeginEx starts a transaction with txOptions determining the transaction mode. Unlike database/sql, the context only
// affects the begin command. i.e. there is no auto-rollback on context cancelation.
func (c *Conn) BeginEx(ctx context.Context, txOptions TxOptions) (*Tx, error) {
	_, err := c.Exec(ctx, txOptions.beginSQL())
	if err != nil {
		// begin should never fail unless there is an underlying connection issue or
		// a context timeout. In either case, the connection is possibly broken.
		c.die(errors.New("failed to begin transaction"))
		return nil, err
	}

	return &Tx{conn: c}, nil
}

// Tx represents a database transaction.
//
// All Tx methods return ErrTxClosed if Commit or Rollback has already been
// called on the Tx.
type Tx struct {
	conn   *Conn
	err    error
	status int8
}

// Commit commits the transaction.
func (tx *Tx) Commit(ctx context.Context) error {
	if tx.status != TxStatusInProgress {
		return ErrTxClosed
	}

	commandTag, err := tx.conn.Exec(ctx, "commit")
	if err == nil && string(commandTag) == "COMMIT" {
		tx.status = TxStatusCommitSuccess
	} else if err == nil && string(commandTag) == "ROLLBACK" {
		tx.status = TxStatusCommitFailure
		tx.err = ErrTxCommitRollback
	} else {
		tx.status = TxStatusCommitFailure
		tx.err = err
		// A commit failure leaves the connection in an undefined state
		tx.conn.die(errors.New("commit failed"))
	}

	return tx.err
}

// Rollback rolls back the transaction. Rollback will return ErrTxClosed if the
// Tx is already closed, but is otherwise safe to call multiple times. Hence, a
// defer tx.Rollback() is safe even if tx.Commit() will be called first in a
// non-error condition.
func (tx *Tx) Rollback(ctx context.Context) error {
	if tx.status != TxStatusInProgress {
		return ErrTxClosed
	}

	_, tx.err = tx.conn.Exec(ctx, "rollback")
	if tx.err == nil {
		tx.status = TxStatusRollbackSuccess
	} else {
		tx.status = TxStatusRollbackFailure
		// A rollback failure leaves the connection in an undefined state
		tx.conn.die(errors.New("rollback failed"))
	}

	return tx.err
}

// Exec delegates to the underlying *Conn
func (tx *Tx) Exec(ctx context.Context, sql string, arguments ...interface{}) (commandTag pgconn.CommandTag, err error) {
	return tx.conn.Exec(ctx, sql, arguments...)
}

// Prepare delegates to the underlying *Conn
func (tx *Tx) Prepare(ctx context.Context, name, sql string) (*PreparedStatement, error) {
	if tx.status != TxStatusInProgress {
		return nil, ErrTxClosed
	}

	return tx.conn.Prepare(ctx, name, sql)
}

// Query delegates to the underlying *Conn
func (tx *Tx) Query(ctx context.Context, sql string, args ...interface{}) (Rows, error) {
	if tx.status != TxStatusInProgress {
		// Because checking for errors can be deferred to the *Rows, build one with the error
		err := ErrTxClosed
		return &connRows{closed: true, err: err}, err
	}

	return tx.conn.Query(ctx, sql, args...)
}

// QueryRow delegates to the underlying *Conn
func (tx *Tx) QueryRow(ctx context.Context, sql string, args ...interface{}) Row {
	rows, _ := tx.Query(ctx, sql, args...)
	return (*connRow)(rows.(*connRows))
}

// CopyFrom delegates to the underlying *Conn
func (tx *Tx) CopyFrom(ctx context.Context, tableName Identifier, columnNames []string, rowSrc CopyFromSource) (int64, error) {
	if tx.status != TxStatusInProgress {
		return 0, ErrTxClosed
	}

	return tx.conn.CopyFrom(ctx, tableName, columnNames, rowSrc)
}

// SendBatch delegates to the underlying *Conn
func (tx *Tx) SendBatch(ctx context.Context, b *Batch) BatchResults {
	if tx.status != TxStatusInProgress {
		return &batchResults{err: ErrTxClosed}
	}

	return tx.conn.SendBatch(ctx, b)
}

// Status returns the status of the transaction from the set of
// pgx.TxStatus* constants.
func (tx *Tx) Status() int8 {
	if tx.status == TxStatusInProgress && tx.conn.pgConn.TxStatus == 'E' {
		return TxStatusInFailure
	}
	return tx.status
}

// Err returns the final error state, if any, of calling Commit or Rollback.
func (tx *Tx) Err() error {
	return tx.err
}
<|endoftext|>


<commit_before>package main
import (
    "fmt"
    "os"
    "strings"
    "bufio"
    "path/filepath"
    "strconv"
)

type dataset struct {
     title string
     path_to_file string
     json_data []map[string]interface{}
}

type uistate struct {
     phase int
     active_set *dataset
     key string
     value string
     datasets []dataset
     scanner *bufio.Scanner
}

func scan_or_exit(scanner *bufio.Scanner) {
    if scanner.Scan() == false && scanner.Err() == nil {
        os.Exit(0)
    }
}

func enter_interactive_loop(directory string) {
    scanner := bufio.NewScanner(os.Stdin)

    println("JSON Search Tool")

    datasets := find_datasets(directory)

    for {
        key, value, set := request_search_fields(datasets, scanner)
        if set != nil && set.json_data != nil {
            search_json(set.json_data, key, value)
        }
    }
    if err := scanner.Err(); err != nil {
        fmt.Println(os.Stderr, "error:", err)
        os.Exit(1)
    }
}

func find_datasets(directory string) []dataset {

    // Assumes no-one will be malicious enough to create a direcory with a '.json' suffix
    path_elements := []string { directory, "*.json" }
    files, _ := filepath.Glob(strings.Join(path_elements, "/"))

    datasets := make([]dataset, 0)

    for _, path_to_file := range files {
        elements := strings.Split(path_to_file, "/")
        filename := elements[len(elements)-1]
        set := dataset{strings.Title(strings.TrimSuffix(filename, ".json")), path_to_file, nil}
        datasets = append(datasets, set)
    }
    return datasets
}

func unpack_dataset(set *dataset) bool {
    if set.json_data == nil {
        valid, json_data := parse_file(set.path_to_file)
	if valid {
            (*set).json_data = json_data
	}
    }
    if set.json_data != nil {
        return true
    }
    return false
}

func select_dataset(state uistate) uistate {
    fmt.Println("\nPlease select a dataset to search, or 'quit' to exit:")
    fmt.Printf("   ")
    for index, set := range state.datasets {
        if strings.Compare(set.title, "") != 0 {
            fmt.Printf("%v) %v ", index+1, set.title)
        }
    }
    fmt.Printf("\n# ")

    scan_or_exit(state.scanner)
    user_input := state.scanner.Text()
    if strings.Compare("quit", strings.ToLower(user_input)) == 0 {
        state.phase = -1
	return state
    }

    index, _ := strconv.ParseInt(user_input, 10, 32)
    if index > 0 && index <= int64(len(state.datasets)) {
        if unpack_dataset(&(state.datasets[index-1])) {
            state.active_set = &(state.datasets[index-1])
	    state.phase += 1

        } else {
	    badset := state.datasets[index-1]
            fmt.Printf("Data source %v is corrupted. Please choose again.\n", badset.title)
            state.datasets[index-1].title = ""
        }

    } else {
        fmt.Printf("Invalid selection: '%v'.\n", user_input)
    }

    return state
}

func prompt_for_input(prompt string, help string) {
    fmt.Printf("\nEnter a %v to search for:\n", prompt)
    fmt.Printf("   '?' to see %v,\n", help)
    fmt.Printf("   '..' to go back\n")
    fmt.Printf("   'quit' to exit\n")
}

func select_field(state uistate) uistate {
    prompt_for_input("term", "available fields")
    fmt.Printf("%v # ", state.active_set.title)

    scan_or_exit(state.scanner)
    user_input := state.scanner.Text()

    if len(user_input) == 0 {
        // Could we ever wish to search all fields for a specific value?
        fmt.Printf("Invalid selection")

    } else if strings.Compare(user_input, "quit") == 0 {
        state.phase = -1

    } else if strings.Compare(user_input, "..") == 0 {
        state.phase -= 1

    } else if strings.Compare(user_input, "?") == 0 {
        if len(state.active_set.json_data) > 0 {
            // Assume for now that records are sufficiently uniform
            fmt.Printf("\n%v records contain the following fields\n", strings.TrimSuffix(state.active_set.title, "s"))
            for key, _:= range state.active_set.json_data[0] {
                  fmt.Printf("* %s\n", key)
            }
        } else {
            fmt.Printf("* No records found *\n")
        }

    } else {
        state.key = user_input
        state.phase += 1
    }

    return state
}

func select_value(state uistate) uistate {
    prompt_for_input("value", "example values")
    fmt.Printf("%v[%v] # ", state.active_set.title, state.key)

    scan_or_exit(state.scanner)
    user_input := state.scanner.Text()

    if len(user_input) == 0 {
        fmt.Printf("Searching for empty '%v' fields", state.key)

    } else if strings.Compare(user_input, "quit") == 0 {
        state.phase = -1

    } else if strings.Compare(user_input, "?") == 0 {
        if len(state.active_set.json_data) > 0 {
            fmt.Printf("\n%v records contain values like:\n", strings.TrimSuffix(state.active_set.title, "s"))

            // TODO: Nicer handling complex types (arrays)
            for index, record := range state.active_set.json_data {
                fmt.Printf("* %v\n", record[state.key])
                if index > 2 {
                        break
                }
            }
        } else {
            fmt.Printf("* No records found *\n")
        }

    } else if strings.Compare(user_input, "..") == 0 {
        state.phase -= 1

    } else {
        state.value = user_input
        state.phase += 1
    }
    return state
}

func request_search_fields(datasets []dataset, scanner *bufio.Scanner) (string, string, *dataset) {

    state := uistate{0, nil, "", "", datasets, scanner}

    for {
        switch state.phase {
            case 0:
	        state = select_dataset(state)
	    case 1:
	        state = select_field(state)
	    case 2:
	        state = select_value(state)
	    case 3:
	        return state.key, state.value, state.active_set
	    case -1:
                os.Exit(0)
            default:
	        state.phase = 0
	}
    }
}

<commit_msg>Massage the user facing text<commit_after>package main
import (
    "fmt"
    "os"
    "strings"
    "bufio"
    "path/filepath"
    "strconv"
)

type dataset struct {
     title string
     path_to_file string
     json_data []map[string]interface{}
}

type uistate struct {
     phase int
     active_set *dataset
     key string
     value string
     datasets []dataset
     scanner *bufio.Scanner
}

func scan_or_exit(scanner *bufio.Scanner) {
    if scanner.Scan() == false && scanner.Err() == nil {
        os.Exit(0)
    }
}

func enter_interactive_loop(directory string) {
    scanner := bufio.NewScanner(os.Stdin)

    println("JSON Search Tool")

    datasets := find_datasets(directory)

    for {
        key, value, set := request_search_fields(datasets, scanner)
        if set != nil && set.json_data != nil {
            search_json(set.json_data, key, value)
        }
    }
    if err := scanner.Err(); err != nil {
        fmt.Println(os.Stderr, "error:", err)
        os.Exit(1)
    }
}

func find_datasets(directory string) []dataset {

    // Assumes no-one will be malicious enough to create a direcory with a '.json' suffix
    path_elements := []string { directory, "*.json" }
    files, _ := filepath.Glob(strings.Join(path_elements, "/"))

    datasets := make([]dataset, 0)

    for _, path_to_file := range files {
        elements := strings.Split(path_to_file, "/")
        filename := elements[len(elements)-1]
        set := dataset{strings.Title(strings.TrimSuffix(filename, ".json")), path_to_file, nil}
        datasets = append(datasets, set)
    }
    return datasets
}

func unpack_dataset(set *dataset) bool {
    if set.json_data == nil {
        valid, json_data := parse_file(set.path_to_file)
	if valid {
            (*set).json_data = json_data
	}
    }
    if set.json_data != nil {
        return true
    }
    return false
}

func select_dataset(state uistate) uistate {
    fmt.Println("\nPlease select a dataset to search, or 'quit' to exit:")
    fmt.Printf("   ")
    for index, set := range state.datasets {
        if strings.Compare(set.title, "") != 0 {
            fmt.Printf("%v) %v ", index+1, set.title)
        }
    }
    fmt.Printf("\n# ")

    scan_or_exit(state.scanner)
    user_input := state.scanner.Text()
    if strings.Compare("quit", strings.ToLower(user_input)) == 0 {
        state.phase = -1
	return state
    }

    index, _ := strconv.ParseInt(user_input, 10, 32)
    if index > 0 && index <= int64(len(state.datasets)) {
        if unpack_dataset(&(state.datasets[index-1])) {
            state.active_set = &(state.datasets[index-1])
	    state.phase += 1

        } else {
	    badset := state.datasets[index-1]
            fmt.Printf("Data source %v is corrupted. Please choose again.\n", badset.title)
            state.datasets[index-1].title = ""
        }

    } else {
        fmt.Printf("Invalid selection: '%v'.  Enter the number of a dataset instead.\n", user_input)
    }

    return state
}

func prompt_for_input(prompt string, help string) {
    fmt.Printf("\n   '?' to see %v\n", help)
    fmt.Printf("   '..' to go back\n")
    fmt.Printf("   'quit' to exit\n")
    fmt.Printf("\nEnter a %v to search for:\n", prompt)
}

func select_field(state uistate) uistate {
    prompt_for_input("term", "available fields")
    fmt.Printf("%v # ", state.active_set.title)

    scan_or_exit(state.scanner)
    user_input := state.scanner.Text()

    if len(user_input) == 0 {
        // Could we ever wish to search all fields for a specific value?
        fmt.Printf("Invalid selection")

    } else if strings.Compare(user_input, "quit") == 0 {
        state.phase = -1

    } else if strings.Compare(user_input, "..") == 0 {
        state.phase -= 1

    } else if strings.Compare(user_input, "?") == 0 {
        if len(state.active_set.json_data) > 0 {
            // Assume for now that records are sufficiently uniform
            fmt.Printf("\n%v records contain the following fields:\n", strings.TrimSuffix(state.active_set.title, "s"))
            for key, _:= range state.active_set.json_data[0] {
                  fmt.Printf("* %s\n", key)
            }
        } else {
            fmt.Printf("* No records found *\n")
        }

    } else {
        state.key = user_input
        state.phase += 1
    }

    return state
}

func select_value(state uistate) uistate {
    prompt_for_input("value", "example values")
    fmt.Printf("%v[%v] # ", state.active_set.title, state.key)

    scan_or_exit(state.scanner)
    user_input := state.scanner.Text()

    if len(user_input) == 0 {
        fmt.Printf("Searching for empty '%v' fields", state.key)

    } else if strings.Compare(user_input, "quit") == 0 {
        state.phase = -1

    } else if strings.Compare(user_input, "?") == 0 {
        if len(state.active_set.json_data) > 0 {
            fmt.Printf("\n%v records contain values like:\n", strings.TrimSuffix(state.active_set.title, "s"))

            // TODO: Nicer handling complex types (arrays)
            for index, record := range state.active_set.json_data {
                fmt.Printf("* %v\n", record[state.key])
                if index > 2 {
                        break
                }
            }
        } else {
            fmt.Printf("* No records found *\n")
        }

    } else if strings.Compare(user_input, "..") == 0 {
        state.phase -= 1

    } else {
        state.value = user_input
        state.phase += 1
    }
    return state
}

func request_search_fields(datasets []dataset, scanner *bufio.Scanner) (string, string, *dataset) {

    state := uistate{0, nil, "", "", datasets, scanner}

    for {
        switch state.phase {
            case 0:
	        state = select_dataset(state)
	    case 1:
	        state = select_field(state)
	    case 2:
	        state = select_value(state)
	    case 3:
	        return state.key, state.value, state.active_set
	    case -1:
                os.Exit(0)
            default:
	        state.phase = 0
	}
    }
}

<|endoftext|>


<commit_before>package proxyproto

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"io"
	"strconv"
)

var (
	lengthV4Bytes = func() []byte {
		a := make([]byte, 2)
		binary.BigEndian.PutUint16(a, 12)
		return a
	}()
	lengthV6Bytes = func() []byte {
		a := make([]byte, 2)
		binary.BigEndian.PutUint16(a, 36)
		return a
	}()
	lengthUnixBytes = func() []byte {
		a := make([]byte, 2)
		binary.BigEndian.PutUint16(a, 218)
		return a
	}()
)

type _ports struct {
	SrcPort uint16
	DstPort uint16
}

type _addr4 struct {
	Src     [4]byte
	Dst     [4]byte
	SrcPort uint16
	DstPort uint16
}

type _addr6 struct {
	Src [16]byte
	Dst [16]byte
	_ports
}

type _addrUnix struct {
	Src [108]byte
	Dst [108]byte
}

func parseVersion2(reader *bufio.Reader) (header *Header, err error) {
	// Skip first 12 bytes (signature)
	for i := 0; i < 12; i++ {
		if _, err = reader.ReadByte(); err != nil {
			return nil, ErrCantReadProtocolVersionAndCommand
		}
	}

	header = new(Header)
	header.Version = 2

	// Read the 13th byte, protocol version and command
	b13, err := reader.ReadByte()
	if err != nil {
		return nil, ErrCantReadProtocolVersionAndCommand
	}
	header.Command = ProtocolVersionAndCommand(b13)
	if _, ok := supportedCommand[header.Command]; !ok {
		return nil, ErrUnsupportedProtocolVersionAndCommand
	}
	// If command is LOCAL, header ends here
	if header.Command.IsLocal() {
		return header, nil
	}

	// Read the 14th byte, address family and protocol
	b14, err := reader.ReadByte()
	if err != nil {
		return nil, ErrCantReadAddressFamilyAndProtocol
	}
	header.TransportProtocol = AddressFamilyAndProtocol(b14)
	if _, ok := supportedTransportProtocol[header.TransportProtocol]; !ok {
		return nil, ErrUnsupportedAddressFamilyAndProtocol
	}

	// Read addresses and ports
	var length uint16
	if err := binary.Read(io.LimitReader(reader, 2), binary.BigEndian, &length); err != nil {
		return nil, ErrCantReadLength
	}
	if !header.validateLength(length) {
		return nil, ErrInvalidLength
	}

	if header.TransportProtocol.IsIPv4() {
		var addr _addr4
		if err := binary.Read(io.LimitReader(reader, int64(length)), binary.BigEndian, &addr); err != nil {
			return nil, ErrInvalidAddress
		}
		header.SourceAddress = addr.Src[:]
		header.DestinationAddress = addr.Dst[:]
		header.SourcePort = addr.SrcPort
		header.DestinationPort = addr.DstPort
	} else if header.TransportProtocol.IsIPv6() {
		var addr _addr6
		if err := binary.Read(io.LimitReader(reader, int64(length)), binary.BigEndian, &addr); err != nil {
			return nil, ErrInvalidAddress
		}
		header.SourceAddress = addr.Src[:]
		header.DestinationAddress = addr.Dst[:]
		header.SourcePort = addr.SrcPort
		header.DestinationPort = addr.DstPort
	}
	// TODO fully support Unix addresses
	//	else if header.TransportProtocol.IsUnix() {
	//		var addr _addrUnix
	//		if err := binary.Read(io.LimitReader(reader, int64(length)), binary.BigEndian, &addr); err != nil {
	//			return nil, ErrInvalidAddress
	//		}
	//
	//if header.SourceAddress, err = net.ResolveUnixAddr("unix", string(addr.Src[:])); err != nil {
	//	return nil, ErrCantResolveSourceUnixAddress
	//}
	//if header.DestinationAddress, err = net.ResolveUnixAddr("unix", string(addr.Dst[:])); err != nil {
	//	return nil, ErrCantResolveDestinationUnixAddress
	//}
	//}

	// TODO add encapsulated TLV support

	return header, nil
}

func (header *Header) writeVersion2(w io.Writer) (int64, error) {
	var buf bytes.Buffer
	buf.Write(SIGV2)
	buf.WriteByte(header.Command.toByte())
	if !header.Command.IsLocal() {
		buf.WriteByte(header.TransportProtocol.toByte())
		// TODO add encapsulated TLV length
		var addrSrc, addrDst []byte
		if header.TransportProtocol.IsIPv4() {
			buf.Write(lengthV4Bytes)
			addrSrc = header.SourceAddress.To4()
			addrDst = header.DestinationAddress.To4()
		} else if header.TransportProtocol.IsIPv6() {
			buf.Write(lengthV6Bytes)
			addrSrc = header.SourceAddress.To16()
			addrDst = header.DestinationAddress.To16()
		} else if header.TransportProtocol.IsUnix() {
			buf.Write(lengthUnixBytes)
			// TODO is below right?
			addrSrc = []byte(header.SourceAddress.String())
			addrDst = []byte(header.DestinationAddress.String())
		}
		buf.Write(addrSrc)
		buf.Write(addrDst)
	}
	buf.WriteString(strconv.Itoa(int(header.SourcePort)))
	buf.WriteString(strconv.Itoa(int(header.DestinationPort)))

	return buf.WriteTo(w)
}

func (header *Header) validateLength(length uint16) bool {
	if header.TransportProtocol.IsIPv4() {
		return length == 12
	} else if header.TransportProtocol.IsIPv6() {
		return length == 36
	} else if header.TransportProtocol.IsUnix() {
		return length == 218
	}
	return false
}
<commit_msg>Fixed ports writing in protocol version 2.<commit_after>package proxyproto

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"io"
)

var (
	lengthV4Bytes = func() []byte {
		a := make([]byte, 2)
		binary.BigEndian.PutUint16(a, 12)
		return a
	}()
	lengthV6Bytes = func() []byte {
		a := make([]byte, 2)
		binary.BigEndian.PutUint16(a, 36)
		return a
	}()
	lengthUnixBytes = func() []byte {
		a := make([]byte, 2)
		binary.BigEndian.PutUint16(a, 218)
		return a
	}()
)

type _ports struct {
	SrcPort uint16
	DstPort uint16
}

type _addr4 struct {
	Src     [4]byte
	Dst     [4]byte
	SrcPort uint16
	DstPort uint16
}

type _addr6 struct {
	Src [16]byte
	Dst [16]byte
	_ports
}

type _addrUnix struct {
	Src [108]byte
	Dst [108]byte
}

func parseVersion2(reader *bufio.Reader) (header *Header, err error) {
	// Skip first 12 bytes (signature)
	for i := 0; i < 12; i++ {
		if _, err = reader.ReadByte(); err != nil {
			return nil, ErrCantReadProtocolVersionAndCommand
		}
	}

	header = new(Header)
	header.Version = 2

	// Read the 13th byte, protocol version and command
	b13, err := reader.ReadByte()
	if err != nil {
		return nil, ErrCantReadProtocolVersionAndCommand
	}
	header.Command = ProtocolVersionAndCommand(b13)
	if _, ok := supportedCommand[header.Command]; !ok {
		return nil, ErrUnsupportedProtocolVersionAndCommand
	}
	// If command is LOCAL, header ends here
	if header.Command.IsLocal() {
		return header, nil
	}

	// Read the 14th byte, address family and protocol
	b14, err := reader.ReadByte()
	if err != nil {
		return nil, ErrCantReadAddressFamilyAndProtocol
	}
	header.TransportProtocol = AddressFamilyAndProtocol(b14)
	if _, ok := supportedTransportProtocol[header.TransportProtocol]; !ok {
		return nil, ErrUnsupportedAddressFamilyAndProtocol
	}

	// Read addresses and ports
	var length uint16
	if err := binary.Read(io.LimitReader(reader, 2), binary.BigEndian, &length); err != nil {
		return nil, ErrCantReadLength
	}
	if !header.validateLength(length) {
		return nil, ErrInvalidLength
	}

	if header.TransportProtocol.IsIPv4() {
		var addr _addr4
		if err := binary.Read(io.LimitReader(reader, int64(length)), binary.BigEndian, &addr); err != nil {
			return nil, ErrInvalidAddress
		}
		header.SourceAddress = addr.Src[:]
		header.DestinationAddress = addr.Dst[:]
		header.SourcePort = addr.SrcPort
		header.DestinationPort = addr.DstPort
	} else if header.TransportProtocol.IsIPv6() {
		var addr _addr6
		if err := binary.Read(io.LimitReader(reader, int64(length)), binary.BigEndian, &addr); err != nil {
			return nil, ErrInvalidAddress
		}
		header.SourceAddress = addr.Src[:]
		header.DestinationAddress = addr.Dst[:]
		header.SourcePort = addr.SrcPort
		header.DestinationPort = addr.DstPort
	}
	// TODO fully support Unix addresses
	//	else if header.TransportProtocol.IsUnix() {
	//		var addr _addrUnix
	//		if err := binary.Read(io.LimitReader(reader, int64(length)), binary.BigEndian, &addr); err != nil {
	//			return nil, ErrInvalidAddress
	//		}
	//
	//if header.SourceAddress, err = net.ResolveUnixAddr("unix", string(addr.Src[:])); err != nil {
	//	return nil, ErrCantResolveSourceUnixAddress
	//}
	//if header.DestinationAddress, err = net.ResolveUnixAddr("unix", string(addr.Dst[:])); err != nil {
	//	return nil, ErrCantResolveDestinationUnixAddress
	//}
	//}

	// TODO add encapsulated TLV support

	return header, nil
}

func (header *Header) writeVersion2(w io.Writer) (int64, error) {
	var buf bytes.Buffer
	buf.Write(SIGV2)
	buf.WriteByte(header.Command.toByte())
	if !header.Command.IsLocal() {
		buf.WriteByte(header.TransportProtocol.toByte())
		// TODO add encapsulated TLV length
		var addrSrc, addrDst []byte
		if header.TransportProtocol.IsIPv4() {
			buf.Write(lengthV4Bytes)
			addrSrc = header.SourceAddress.To4()
			addrDst = header.DestinationAddress.To4()
		} else if header.TransportProtocol.IsIPv6() {
			buf.Write(lengthV6Bytes)
			addrSrc = header.SourceAddress.To16()
			addrDst = header.DestinationAddress.To16()
		} else if header.TransportProtocol.IsUnix() {
			buf.Write(lengthUnixBytes)
			// TODO is below right?
			addrSrc = []byte(header.SourceAddress.String())
			addrDst = []byte(header.DestinationAddress.String())
		}
		buf.Write(addrSrc)
		buf.Write(addrDst)

		portSrcBytes := func() []byte {
			a := make([]byte, 2)
			binary.BigEndian.PutUint16(a, header.SourcePort)
			return a
		}()
		buf.Write(portSrcBytes)

		portDstBytes := func() []byte {
			a := make([]byte, 2)
			binary.BigEndian.PutUint16(a, header.DestinationPort)
			return a
		}()
		buf.Write(portDstBytes)

	}

	return buf.WriteTo(w)
}

func (header *Header) validateLength(length uint16) bool {
	if header.TransportProtocol.IsIPv4() {
		return length == 12
	} else if header.TransportProtocol.IsIPv6() {
		return length == 36
	} else if header.TransportProtocol.IsUnix() {
		return length == 218
	}
	return false
}
<|endoftext|>