<commit_before>package main

import (
	"container/list"
	"fmt"
	"math"
	"os"
)

type Screen interface {
	Map() int // draw the screen in the views map
	View() *View
	SetView(*View)
}

// XXX editbuffers are editable text buffers

// todo: lockable
type EditBuffer struct {
	st                             *os.FileInfo // file info
	fromDisk, fromArchive, fromNet bool

	lco   uint   // line count
	title string // buffer title
	dirty bool   // dirty

	// XXX text
	lines *list.List    // list of lines
	line  *list.Element // current line

	anchor *list.Element // first line to draw
	view   *View         // view this buffer draws to

	hist *list.List

	prev, next *EditBuffer // roll ourselves because type assertions are pobyteless in this case.

	OptLineNo bool
}

func NewEditBuffer(title string) *EditBuffer {

	b := new(EditBuffer)

	// XXX it would be better to init this to nil and add the line from another
	// place so that the file loader can use this
	b.lines = list.New()
	// b.InsertLine(NewLine([]byte("")))
	b.anchor = b.line
	b.st = nil
	b.title = title
	b.next = nil
	b.prev = nil
	b.dirty = false

	b.view = Vw
	b.OptLineNo = true

	return b
}

func (b *EditBuffer) InsertChar(ch byte) {
	if l, ok := b.line.Value.(*Line); ok {
		l.insertCharacter(ch)
	}
	b.dirty = true
}

func (b *EditBuffer) BackSpace() {
	if b.line == nil {
		Beep()
		return
	}

	if l, ok := b.line.Value.(*Line); ok {
		if l.cursor == 0 {
			if l.size != 0 && b.line.Prev() != nil {
				// combine this line and the previous
			} else {

				if b.line.Prev() != nil {
					b.DeleteCurrLine()
					b.lco--
				} else {
					Beep()
				}
			}
		} else {
			l.backspace()
		}
	}
}

func (b *EditBuffer) MoveLeft() {
	b.MoveCursor(-1)
}

func (b *EditBuffer) MoveRight() {
	b.MoveCursor(1)
}

func (b *EditBuffer) MoveCursor(d int) {
	if l, ok := b.line.Value.(*Line); ok && l.moveCursor(l.cursor+d) < 0 {
		Beep()
	}
}

func (b *EditBuffer) MoveDown() {
	if l, ok := b.line.Value.(*Line); ok {
		if n := b.line.Next(); n != nil {
			ln := n.Value.(*Line)
			if ln.moveCursor(l.cursor) < 0 {
				ln.moveCursor(ln.cursorMax())
			}
			b.line = n
		} else {
			Beep()
		}
	}
}

func (b *EditBuffer) MoveUp() {
	if l, ok := b.line.Value.(*Line); ok {
		if p := b.line.Prev(); p != nil {
			lp := p.Value.(*Line)
			if lp.moveCursor(l.cursor) < 0 {
				lp.moveCursor(lp.cursorMax())
			}
			b.line = p
		} else {
			Beep()
		}
	}
}

func (b *EditBuffer) DeleteSpan(p, l int) {
	if ln, ok := b.line.Value.(*Line); ok {
		ln.delete(p, l)
		b.dirty = true
	}
}

func (b *EditBuffer) FirstLine() {
	b.line = b.lines.Front()
}

// Insert a new line after line
// XXX this doesn't really work as desired (can't insert at 0, for instance)
func (b *EditBuffer) InsertLine(line *Line) {
	if b.line == nil {
		b.line = b.lines.PushFront(line)
		l := b.line.Value.(*Line)
		l.lno = 1
	} else {
		b.line = b.lines.InsertAfter(line, b.line)
		l := b.line.Value.(*Line)
		l.lno = b.line.Prev().Value.(*Line).lno + 1
		for p := b.line.Next(); p != nil; p = p.Next() {
			p.Value.(*Line).lno++
		}
	}
	b.lco++
	b.dirty = true
}

func (b *EditBuffer) AppendLine() {
	b.InsertLine(NewLine([]byte("")))
}

func (b *EditBuffer) NewLine(nlchar byte) {
	if l, ok := b.line.Value.(*Line); ok {
		newbuf := l.raw()[l.cursor:]
		l.insertCharacter(nlchar)
		l.ClearAfterCursor()
		l.size -= len(newbuf)
		b.InsertLine(NewLine(newbuf))
	}
}

func (b *EditBuffer) DeleteCurrLine() {
	if b.line.Prev() != nil {
		p := b.line.Prev()
		b.lines.Remove(b.line)
		b.line = p
	} else if b.line.Next() != nil {
		n := b.line.Next()
		b.lines.Remove(b.line)
		b.line = n
	} else {
		return
	}
	b.dirty = true
}

// Move to line p
func (b *EditBuffer) MoveLine(p int) {
	i := 0
	for l := b.lines.Front(); l != nil; l = l.Next() {
		if i == p {
			b.line = l
			return
		}
		i++
	}
}

func (b *EditBuffer) MoveLineNext() {
	if n := b.line.Next(); n != nil {
		b.line = n
	}
}

func (b *EditBuffer) MoveLinePrev() {
	if p := b.line.Prev(); p != nil {
		b.line = p
	}
}

func (b *EditBuffer) LnoOffset() int {
	// if we show line numbers, we reserve at least 3 lines.
	if b.OptLineNo {
		if dg := math.Log10(float64(b.lco)) + 1; dg > 2 {
			return int(dg) + 1
		} else {
			return 3
		}
	}
	return 0
}

func (b *EditBuffer) ScreenLines(ln *Line) int {
	offset := b.LnoOffset()
	actual := b.view.Cols - offset
	return int(math.Ceil(float64(len(ln.raw())) / float64(actual)))
}

func (b *EditBuffer) Map() int {
	offset := b.LnoOffset()
	i := 0
	for l := b.anchor; l != nil && i < b.view.Rows; l = l.Next() {
		ln := l.Value.(*Line)
		cnt := b.ScreenLines(ln)
		if cnt == 0 {
			cnt = 1
		}
		wrap := 0
		for lmt := i + cnt; i < lmt; i++ {
			str := make([]byte, b.view.Cols)
			// XXX this is the first part of the line, the optional
			// line number.  if this isn't the first screen line of
			// a line (for a wrapped line), then just draw empty
			// space if line numbers are on
			if wrap == 0 {
				lno := []byte(fmt.Sprintf("%*d ", offset-1, ln.lno))
				copy(str[0:offset], lno)
			} else {
				for j := 0; j < offset; j++ {
					str[j] = ' '
				}
			}

			// XXX the second part of the line, which shows actual
			// text that the user is viewing
			actual := b.view.Cols - offset
			start := actual * wrap
			if start + actual - 1 < len(ln.raw()) {
				copy(str[offset:], ln.raw()[start:start+actual - 1])
			} else {
				copy(str[offset:], ln.raw()[start:len(ln.raw())])
			}
			b.view.Lines[i] = string(str)
			wrap++
		}
	}
	for ; i < b.view.Rows-1; i++ {
		b.view.Lines[i] = NaL
	}
	return 0
}

func (b *EditBuffer) CursorCoord() (int, int) {
	l := b.line.Value.(*Line)
	return int(l.cursor) + b.LnoOffset(), int(l.lno)
}

func (b *EditBuffer) Lines() *list.List {
	return b.lines
}

func (b *EditBuffer) Title() string {
	return b.title
}
<commit_msg>i maked a comment!<commit_after>package main

import (
	"container/list"
	"fmt"
	"math"
	"os"
)

type Screen interface {
	Map() int // draw the screen in the views map
	View() *View
	SetView(*View)
}

// XXX editbuffers are editable text buffers

// todo: lockable
type EditBuffer struct {
	st                             *os.FileInfo // file info
	fromDisk, fromArchive, fromNet bool

	lco   uint   // line count
	title string // buffer title
	dirty bool   // dirty

	// XXX text
	lines *list.List    // list of lines
	line  *list.Element // current line

	anchor *list.Element // first line to draw
	view   *View         // view this buffer draws to

	hist *list.List

	prev, next *EditBuffer // roll ourselves because type assertions are pobyteless in this case.

	OptLineNo bool
}

func NewEditBuffer(title string) *EditBuffer {

	b := new(EditBuffer)

	// XXX it would be better to init this to nil and add the line from another
	// place so that the file loader can use this
	b.lines = list.New()
	// b.InsertLine(NewLine([]byte("")))
	b.anchor = b.line
	b.st = nil
	b.title = title
	b.next = nil
	b.prev = nil
	b.dirty = false

	b.view = Vw
	b.OptLineNo = true

	return b
}

func (b *EditBuffer) InsertChar(ch byte) {
	if l, ok := b.line.Value.(*Line); ok {
		l.insertCharacter(ch)
	}
	b.dirty = true
}

func (b *EditBuffer) BackSpace() {
	if b.line == nil {
		Beep()
		return
	}

	if l, ok := b.line.Value.(*Line); ok {
		if l.cursor == 0 {
			if l.size != 0 && b.line.Prev() != nil {
				// combine this line and the previous
			} else {

				if b.line.Prev() != nil {
					b.DeleteCurrLine()
					b.lco--
				} else {
					Beep()
				}
			}
		} else {
			l.backspace()
		}
	}
}

func (b *EditBuffer) MoveLeft() {
	b.MoveCursor(-1)
}

func (b *EditBuffer) MoveRight() {
	b.MoveCursor(1)
}

func (b *EditBuffer) MoveCursor(d int) {
	if l, ok := b.line.Value.(*Line); ok && l.moveCursor(l.cursor+d) < 0 {
		Beep()
	}
}

func (b *EditBuffer) MoveDown() {
	if l, ok := b.line.Value.(*Line); ok {
		if n := b.line.Next(); n != nil {
			ln := n.Value.(*Line)
			if ln.moveCursor(l.cursor) < 0 {
				ln.moveCursor(ln.cursorMax())
			}
			b.line = n
		} else {
			Beep()
		}
	}
}

func (b *EditBuffer) MoveUp() {
	if l, ok := b.line.Value.(*Line); ok {
		if p := b.line.Prev(); p != nil {
			lp := p.Value.(*Line)
			if lp.moveCursor(l.cursor) < 0 {
				lp.moveCursor(lp.cursorMax())
			}
			b.line = p
		} else {
			Beep()
		}
	}
}

func (b *EditBuffer) DeleteSpan(p, l int) {
	if ln, ok := b.line.Value.(*Line); ok {
		ln.delete(p, l)
		b.dirty = true
	}
}

func (b *EditBuffer) FirstLine() {
	b.line = b.lines.Front()
}

// Insert a new line after line
// XXX this doesn't really work as desired (can't insert at 0, for instance)
func (b *EditBuffer) InsertLine(line *Line) {
	if b.line == nil {
		b.line = b.lines.PushFront(line)
		l := b.line.Value.(*Line)
		l.lno = 1
	} else {
		b.line = b.lines.InsertAfter(line, b.line)
		l := b.line.Value.(*Line)
		l.lno = b.line.Prev().Value.(*Line).lno + 1
		for p := b.line.Next(); p != nil; p = p.Next() {
			p.Value.(*Line).lno++
		}
	}
	b.lco++
	b.dirty = true
}

func (b *EditBuffer) AppendLine() {
	b.InsertLine(NewLine([]byte("")))
}

func (b *EditBuffer) NewLine(nlchar byte) {
	if l, ok := b.line.Value.(*Line); ok {
		newbuf := l.raw()[l.cursor:]
		l.insertCharacter(nlchar)
		l.ClearAfterCursor()
		l.size -= len(newbuf)
		b.InsertLine(NewLine(newbuf))
	}
}

func (b *EditBuffer) DeleteCurrLine() {
	if b.line.Prev() != nil {
		p := b.line.Prev()
		b.lines.Remove(b.line)
		b.line = p
	} else if b.line.Next() != nil {
		n := b.line.Next()
		b.lines.Remove(b.line)
		b.line = n
	} else {
		return
	}
	b.dirty = true
}

// Move to line p
func (b *EditBuffer) MoveLine(p int) {
	i := 0
	for l := b.lines.Front(); l != nil; l = l.Next() {
		if i == p {
			b.line = l
			return
		}
		i++
	}
}

func (b *EditBuffer) MoveLineNext() {
	if n := b.line.Next(); n != nil {
		b.line = n
	}
}

func (b *EditBuffer) MoveLinePrev() {
	if p := b.line.Prev(); p != nil {
		b.line = p
	}
}

func (b *EditBuffer) LnoOffset() int {
	// if we show line numbers, we reserve at least 3 lines.
	if b.OptLineNo {
		if dg := math.Log10(float64(b.lco)) + 1; dg > 2 {
			return int(dg) + 1
		} else {
			return 3
		}
	}
	return 0
}

func (b *EditBuffer) ScreenLines(ln *Line) int {
	offset := b.LnoOffset()
	actual := b.view.Cols - offset
	return int(math.Ceil(float64(len(ln.raw())) / float64(actual)))
}

// Maps every visible line to a position on the screen.  This is a super-slow complete refresh.
func (b *EditBuffer) Map() int {
	offset := b.LnoOffset()
	i := 0
	for l := b.anchor; l != nil && i < b.view.Rows; l = l.Next() {
		ln := l.Value.(*Line)
		cnt := b.ScreenLines(ln)
		if cnt == 0 {
			cnt = 1
		}
		wrap := 0
		for lmt := i + cnt; i < lmt; i++ {
			str := make([]byte, b.view.Cols)
			// XXX this is the first part of the line, the optional
			// line number.  if this isn't the first screen line of
			// a line (for a wrapped line), then just draw empty
			// space if line numbers are on
			if wrap == 0 {
				lno := []byte(fmt.Sprintf("%*d ", offset-1, ln.lno))
				copy(str[0:offset], lno)
			} else {
				for j := 0; j < offset; j++ {
					str[j] = ' '
				}
			}

			// XXX the second part of the line, which shows actual
			// text that the user is viewing
			actual := b.view.Cols - offset
			start := actual * wrap
			if start + actual - 1 < len(ln.raw()) {
				copy(str[offset:], ln.raw()[start:start+actual - 1])
			} else {
				copy(str[offset:], ln.raw()[start:len(ln.raw())])
			}
			b.view.Lines[i] = string(str)
			wrap++
		}
	}
	for ; i < b.view.Rows-1; i++ {
		b.view.Lines[i] = NaL
	}
	return 0
}

func (b *EditBuffer) CursorCoord() (int, int) {
	l := b.line.Value.(*Line)
	return int(l.cursor) + b.LnoOffset(), int(l.lno)
}

func (b *EditBuffer) Lines() *list.List {
	return b.lines
}

func (b *EditBuffer) Title() string {
	return b.title
}
<|endoftext|>


<commit_before>package peco

import (
	"bufio"
	"io"
	"runtime"
	"sync"
	"time"

	"golang.org/x/net/context"

	"github.com/peco/peco/pipeline"
	"github.com/pkg/errors"
)

// ErrBufferOutOfRange is returned when the index within the buffer that
// was queried was out of the containing buffer's range
var ErrBufferOutOfRange = errors.New("error: Specified index is out of range")

func (sp simplePipeline) Cancel()                 { close(sp.cancelCh) }
func (sp simplePipeline) CancelCh() chan struct{} { return sp.cancelCh }
func (sp simplePipeline) OutputCh() chan Line     { return sp.outputCh }
func (sp simplePipeline) Pipeline() (chan struct{}, chan Line) {
	return sp.cancelCh, sp.outputCh
}

func acceptPipeline(cancel chan struct{}, in chan Line, out chan Line, pc *pipelineCtx) {
	trace("acceptPipeline: START")
	defer trace("acceptPipeline: END")
	defer close(out)
	for {
		select {
		case <-cancel:
			trace("acceptPipeline: detected cancel request. Bailing out")
			return
		case l, ok := <-in:
			if l == nil && !ok {
				trace("acceptPipeline: detected end of input. Bailing out")
				if pc.onEnd != nil {
					pc.onEnd()
				}
				return
			}
			trace("acceptPipeline: forwarding to callback")
			if ll, err := pc.onIncomingLine(l); err == nil {
				trace("acceptPipeline: forwarding to out channel")
				out <- ll
			}
		}
	}
}

func (buffers *dependentBuffers) Register(lb LineBuffer) {
	*buffers = append(*buffers, lb)
}

func (buffers *dependentBuffers) Unregister(lb LineBuffer) {
	for i, x := range *buffers {
		if x == lb {
			switch i {
			case 0:
				*buffers = append([]LineBuffer(nil), (*buffers)[1:]...)
			case len(*buffers) - 1:
				*buffers = append([]LineBuffer(nil), (*buffers)[0:i-1]...)
			default:
				*buffers = append(append([]LineBuffer(nil), (*buffers)[0:i-1]...), (*buffers)[i+1:]...)
			}
			return
		}
	}
}

func (buffers dependentBuffers) InvalidateUpTo(i int) {
	for _, b := range buffers {
		b.InvalidateUpTo(i)
	}
}

func NewRawLineBuffer() *RawLineBuffer {
	return &RawLineBuffer{
		simplePipeline: simplePipeline{},
		lines:          []Line{},
		capacity:       0,
	}
}

func (rlb *RawLineBuffer) Replay() error {
	rlb.outputCh = make(chan Line)
	go func() {
		replayed := 0
		trace("RawLineBuffer.Replay (goroutine): START")
		defer func() { trace("RawLineBuffer.Replay (goroutine): END (Replayed %d lines)", replayed) }()

		defer func() { recover() }() // It's okay if we fail to replay
		defer close(rlb.outputCh)
		for _, l := range rlb.lines {
			select {
			case rlb.outputCh <- l:
				replayed++
			case <-rlb.cancelCh:
				return
			}
		}
	}()
	return nil
}

func (rlb *RawLineBuffer) Accept(p Pipeliner) {
	cancelCh, incomingCh := p.Pipeline()
	rlb.cancelCh = cancelCh
	rlb.outputCh = make(chan Line)
	go acceptPipeline(cancelCh, incomingCh, rlb.outputCh,
		&pipelineCtx{rlb.Append, rlb.onEnd})
}

func (rlb *RawLineBuffer) Append(l Line) (Line, error) {
	trace("RawLineBuffer.Append: %s", l.DisplayString())
	if rlb.capacity > 0 && len(rlb.lines) > rlb.capacity {
		diff := len(rlb.lines) - rlb.capacity

		// Golang's version of array realloc
		rlb.lines = rlb.lines[diff:rlb.capacity:rlb.capacity]
	} else {
		rlb.lines = append(rlb.lines, l)
	}

	return l, nil
}

func (rlb *RawLineBuffer) Register(lb LineBuffer) {
	rlb.buffers.Register(lb)
}

func (rlb *RawLineBuffer) Unregister(lb LineBuffer) {
	rlb.buffers.Unregister(lb)
}

// LineAt returns the line at index `i`
func (rlb RawLineBuffer) LineAt(i int) (Line, error) {
	if i < 0 || len(rlb.lines) <= i {
		return nil, ErrBufferOutOfRange
	}
	return rlb.lines[i], nil
}

// Size returns the number of lines in the buffer
func (rlb RawLineBuffer) Size() int {
	return len(rlb.lines)
}

func (rlb *RawLineBuffer) SetCapacity(capacity int) {
	if capacity < 0 {
		capacity = 0
	}
	rlb.capacity = capacity
}

func (rlb RawLineBuffer) InvalidateUpTo(_ int) {
	// no op
}

func (rlb *RawLineBuffer) AppendLine(l Line) (Line, error) {
	return rlb.Append(l)
}

func NewFilteredLineBuffer(src LineBuffer) *FilteredLineBuffer {
	flb := &FilteredLineBuffer{
		simplePipeline: simplePipeline{},
		src:            src,
		selection:      []int{},
	}
	src.Register(flb)

	runtime.SetFinalizer(flb, func(x *FilteredLineBuffer) {
		x.src.Unregister(x)
	})

	return flb
}

func (flb *FilteredLineBuffer) Accept(p Pipeliner) {
	cancelCh, incomingCh := p.Pipeline()
	flb.cancelCh = cancelCh
	flb.outputCh = make(chan Line)
	go acceptPipeline(cancelCh, incomingCh, flb.outputCh,
		&pipelineCtx{flb.Append, nil})
}

func (flb *FilteredLineBuffer) Append(l Line) (Line, error) {
	return l, nil
}

func (flb *FilteredLineBuffer) InvalidateUpTo(x int) {
	p := -1
	for i := 0; i < len(flb.selection); i++ {
		if flb.selection[i] > x {
			break
		}
		p = i
	}

	if p >= 0 {
		flb.selection = append([]int(nil), flb.selection[p:]...)
	}

	for _, b := range flb.buffers {
		b.InvalidateUpTo(p)
	}
}

// LineAt returns the line at index `i`. Note that the i-th element
// in this filtered buffer may actually correspond to a totally
// different line number in the source buffer.
func (flb FilteredLineBuffer) LineAt(i int) (Line, error) {
	if i < 0 || i >= len(flb.selection) {
		return nil, ErrBufferOutOfRange
	}
	return flb.src.LineAt(flb.selection[i])
}

// Size returns the number of lines in the buffer
func (flb FilteredLineBuffer) Size() int {
	return len(flb.selection)
}

func (flb *FilteredLineBuffer) SelectSourceLineAt(i int) {
	flb.selection = append(flb.selection, i)
}

func (flb *FilteredLineBuffer) Register(lb LineBuffer) {
	flb.buffers.Register(lb)
}

func (flb *FilteredLineBuffer) Unregister(lb LineBuffer) {
	flb.buffers.Unregister(lb)
}

// Buffer interface is used for containers for lines to be
// processed by peco.
type Buffer interface {
	LineAt(int) (Line, error)
	Size() int
}

// MemoryBuffer is an implementation of Buffer
type MemoryBuffer struct {
	lines []Line
	mutex sync.Mutex
}

// XXX go through an accessor that returns a reference so that
// we are sure we are accessing/modifying the same mutex
func (mb MemoryBuffer) locker() *sync.Mutex {
	return &mb.mutex
}

func (mb MemoryBuffer) Size() int {
	l := mb.locker()
	l.Lock()
	defer l.Unlock()

	return len(mb.lines)
}

func (mb MemoryBuffer) LineAt(n int) (Line, error) {
	l := mb.locker()
	l.Lock()
	defer l.Unlock()

	if n > mb.Size() || n < 0 {
		return nil, errors.New("empty buffer")
	}

	return mb.lines[n], nil
}

// Source implements pipline.Source, and is the buffer for the input
type Source struct {
	pipeline.OutputChannel
	MemoryBuffer

	in        io.Reader
	enableSep bool
	ready     chan struct{}
	setupOnce sync.Once
}

// Creates a new Source. Does not start processing the input until you
// call Setup()
func NewSource(in io.Reader, enableSep bool) *Source {
	return &Source{
		in:            in, // Note that this may be closed, so do not rely on it
		enableSep:     enableSep,
		ready:         make(chan struct{}),
		setupOnce:     sync.Once{},
		OutputChannel: pipeline.OutputChannel(make(chan interface{})),
	}
}

// Setup reads from the input os.File.
func (s *Source) Setup(state *Peco) {
	s.setupOnce.Do(func() {
		l := s.locker()
		l.Lock()
		defer l.Unlock()

		done := make(chan struct{})
		refresh := make(chan struct{})
		defer close(done)
		defer close(refresh)

		go func() {
			ticker := time.NewTicker(100 * time.Millisecond)
			defer ticker.Stop()

			for {
				select {
				case <-ticker.C:
					run := false
					for loop := true; loop; {
						select {
						case <-refresh:
							run = true
						default:
							loop = false
						}
					}
					if run {
						// Not a great thing to do, allowing nil to be passed
						// as state, but for testing I couldn't come up with anything
						// better for the moment
						if state != nil && !state.ExecQuery() {
							state.Hub().SendDraw(false)
						}
					}
				case <-done:
					return
				}
			}
		}()

		// This sync.Once var is used to receive the notification
		// that there was at least 1 line read from the source
		var notify sync.Once
		notifycb := func() {
			// close the ready channel so others can be notified
			// that there's at least 1 line in the buffer
			close(s.ready)
		}
		scanner := bufio.NewScanner(s.in)
		for scanner.Scan() {
			txt := scanner.Text()
			s.lines = append(s.lines, NewRawLine(txt, s.enableSep))
			notify.Do(notifycb)

			go func() {
				defer func() { recover() }()
				refresh <- struct{}{}
			}()
		}
		trace("Read all from source")
	})
}

// Start starts
func (s *Source) Start(ctx context.Context) {
	go func() {
		defer s.OutputChannel.SendEndMark("end of input")

		for i := 0; i < len(s.lines); i++ {
			select {
			case <-ctx.Done():
				return
			case s.OutputChannel <- s.lines[i]:
				// no op
			}
		}
	}()
}

// Ready returns the "input ready" channel. It will be closed as soon as
// the first line of input is processed via Setup()
func (s *Source) Ready() <-chan struct{} {
	return s.ready
}
<commit_msg>fix updater loop<commit_after>package peco

import (
	"bufio"
	"io"
	"runtime"
	"sync"
	"time"

	"golang.org/x/net/context"

	"github.com/peco/peco/pipeline"
	"github.com/pkg/errors"
)

// ErrBufferOutOfRange is returned when the index within the buffer that
// was queried was out of the containing buffer's range
var ErrBufferOutOfRange = errors.New("error: Specified index is out of range")

func (sp simplePipeline) Cancel()                 { close(sp.cancelCh) }
func (sp simplePipeline) CancelCh() chan struct{} { return sp.cancelCh }
func (sp simplePipeline) OutputCh() chan Line     { return sp.outputCh }
func (sp simplePipeline) Pipeline() (chan struct{}, chan Line) {
	return sp.cancelCh, sp.outputCh
}

func acceptPipeline(cancel chan struct{}, in chan Line, out chan Line, pc *pipelineCtx) {
	trace("acceptPipeline: START")
	defer trace("acceptPipeline: END")
	defer close(out)
	for {
		select {
		case <-cancel:
			trace("acceptPipeline: detected cancel request. Bailing out")
			return
		case l, ok := <-in:
			if l == nil && !ok {
				trace("acceptPipeline: detected end of input. Bailing out")
				if pc.onEnd != nil {
					pc.onEnd()
				}
				return
			}
			trace("acceptPipeline: forwarding to callback")
			if ll, err := pc.onIncomingLine(l); err == nil {
				trace("acceptPipeline: forwarding to out channel")
				out <- ll
			}
		}
	}
}

func (buffers *dependentBuffers) Register(lb LineBuffer) {
	*buffers = append(*buffers, lb)
}

func (buffers *dependentBuffers) Unregister(lb LineBuffer) {
	for i, x := range *buffers {
		if x == lb {
			switch i {
			case 0:
				*buffers = append([]LineBuffer(nil), (*buffers)[1:]...)
			case len(*buffers) - 1:
				*buffers = append([]LineBuffer(nil), (*buffers)[0:i-1]...)
			default:
				*buffers = append(append([]LineBuffer(nil), (*buffers)[0:i-1]...), (*buffers)[i+1:]...)
			}
			return
		}
	}
}

func (buffers dependentBuffers) InvalidateUpTo(i int) {
	for _, b := range buffers {
		b.InvalidateUpTo(i)
	}
}

func NewRawLineBuffer() *RawLineBuffer {
	return &RawLineBuffer{
		simplePipeline: simplePipeline{},
		lines:          []Line{},
		capacity:       0,
	}
}

func (rlb *RawLineBuffer) Replay() error {
	rlb.outputCh = make(chan Line)
	go func() {
		replayed := 0
		trace("RawLineBuffer.Replay (goroutine): START")
		defer func() { trace("RawLineBuffer.Replay (goroutine): END (Replayed %d lines)", replayed) }()

		defer func() { recover() }() // It's okay if we fail to replay
		defer close(rlb.outputCh)
		for _, l := range rlb.lines {
			select {
			case rlb.outputCh <- l:
				replayed++
			case <-rlb.cancelCh:
				return
			}
		}
	}()
	return nil
}

func (rlb *RawLineBuffer) Accept(p Pipeliner) {
	cancelCh, incomingCh := p.Pipeline()
	rlb.cancelCh = cancelCh
	rlb.outputCh = make(chan Line)
	go acceptPipeline(cancelCh, incomingCh, rlb.outputCh,
		&pipelineCtx{rlb.Append, rlb.onEnd})
}

func (rlb *RawLineBuffer) Append(l Line) (Line, error) {
	trace("RawLineBuffer.Append: %s", l.DisplayString())
	if rlb.capacity > 0 && len(rlb.lines) > rlb.capacity {
		diff := len(rlb.lines) - rlb.capacity

		// Golang's version of array realloc
		rlb.lines = rlb.lines[diff:rlb.capacity:rlb.capacity]
	} else {
		rlb.lines = append(rlb.lines, l)
	}

	return l, nil
}

func (rlb *RawLineBuffer) Register(lb LineBuffer) {
	rlb.buffers.Register(lb)
}

func (rlb *RawLineBuffer) Unregister(lb LineBuffer) {
	rlb.buffers.Unregister(lb)
}

// LineAt returns the line at index `i`
func (rlb RawLineBuffer) LineAt(i int) (Line, error) {
	if i < 0 || len(rlb.lines) <= i {
		return nil, ErrBufferOutOfRange
	}
	return rlb.lines[i], nil
}

// Size returns the number of lines in the buffer
func (rlb RawLineBuffer) Size() int {
	return len(rlb.lines)
}

func (rlb *RawLineBuffer) SetCapacity(capacity int) {
	if capacity < 0 {
		capacity = 0
	}
	rlb.capacity = capacity
}

func (rlb RawLineBuffer) InvalidateUpTo(_ int) {
	// no op
}

func (rlb *RawLineBuffer) AppendLine(l Line) (Line, error) {
	return rlb.Append(l)
}

func NewFilteredLineBuffer(src LineBuffer) *FilteredLineBuffer {
	flb := &FilteredLineBuffer{
		simplePipeline: simplePipeline{},
		src:            src,
		selection:      []int{},
	}
	src.Register(flb)

	runtime.SetFinalizer(flb, func(x *FilteredLineBuffer) {
		x.src.Unregister(x)
	})

	return flb
}

func (flb *FilteredLineBuffer) Accept(p Pipeliner) {
	cancelCh, incomingCh := p.Pipeline()
	flb.cancelCh = cancelCh
	flb.outputCh = make(chan Line)
	go acceptPipeline(cancelCh, incomingCh, flb.outputCh,
		&pipelineCtx{flb.Append, nil})
}

func (flb *FilteredLineBuffer) Append(l Line) (Line, error) {
	return l, nil
}

func (flb *FilteredLineBuffer) InvalidateUpTo(x int) {
	p := -1
	for i := 0; i < len(flb.selection); i++ {
		if flb.selection[i] > x {
			break
		}
		p = i
	}

	if p >= 0 {
		flb.selection = append([]int(nil), flb.selection[p:]...)
	}

	for _, b := range flb.buffers {
		b.InvalidateUpTo(p)
	}
}

// LineAt returns the line at index `i`. Note that the i-th element
// in this filtered buffer may actually correspond to a totally
// different line number in the source buffer.
func (flb FilteredLineBuffer) LineAt(i int) (Line, error) {
	if i < 0 || i >= len(flb.selection) {
		return nil, ErrBufferOutOfRange
	}
	return flb.src.LineAt(flb.selection[i])
}

// Size returns the number of lines in the buffer
func (flb FilteredLineBuffer) Size() int {
	return len(flb.selection)
}

func (flb *FilteredLineBuffer) SelectSourceLineAt(i int) {
	flb.selection = append(flb.selection, i)
}

func (flb *FilteredLineBuffer) Register(lb LineBuffer) {
	flb.buffers.Register(lb)
}

func (flb *FilteredLineBuffer) Unregister(lb LineBuffer) {
	flb.buffers.Unregister(lb)
}

// Buffer interface is used for containers for lines to be
// processed by peco.
type Buffer interface {
	LineAt(int) (Line, error)
	Size() int
}

// MemoryBuffer is an implementation of Buffer
type MemoryBuffer struct {
	lines []Line
	mutex sync.Mutex
}

// XXX go through an accessor that returns a reference so that
// we are sure we are accessing/modifying the same mutex
func (mb MemoryBuffer) locker() *sync.Mutex {
	return &mb.mutex
}

func (mb MemoryBuffer) Size() int {
	l := mb.locker()
	l.Lock()
	defer l.Unlock()

	return len(mb.lines)
}

func (mb MemoryBuffer) LineAt(n int) (Line, error) {
	l := mb.locker()
	l.Lock()
	defer l.Unlock()

	if n > mb.Size() || n < 0 {
		return nil, errors.New("empty buffer")
	}

	return mb.lines[n], nil
}

// Source implements pipline.Source, and is the buffer for the input
type Source struct {
	pipeline.OutputChannel
	MemoryBuffer

	in        io.Reader
	enableSep bool
	ready     chan struct{}
	setupOnce sync.Once
}

// Creates a new Source. Does not start processing the input until you
// call Setup()
func NewSource(in io.Reader, enableSep bool) *Source {
	return &Source{
		in:            in, // Note that this may be closed, so do not rely on it
		enableSep:     enableSep,
		ready:         make(chan struct{}),
		setupOnce:     sync.Once{},
		OutputChannel: pipeline.OutputChannel(make(chan interface{})),
	}
}

// Setup reads from the input os.File.
func (s *Source) Setup(state *Peco) {
	s.setupOnce.Do(func() {
		l := s.locker()
		l.Lock()
		defer l.Unlock()

		done := make(chan struct{})
		refresh := make(chan struct{}, 1)
		defer close(done)
		defer close(refresh)

		draw := func(state *Peco, refresh chan struct{}) {
			run := false
			for loop := true; loop; {
				select {
				case _, ok := <-refresh:
					run = true
					loop = ok
				default:
					loop = false
				}
			}
			if !run {
				return
			}
			// Not a great thing to do, allowing nil to be passed
			// as state, but for testing I couldn't come up with anything
			// better for the moment
			if state != nil && !state.ExecQuery() {
				state.Hub().SendDraw(false)
			}
		}

		go func() {
			ticker := time.NewTicker(100 * time.Millisecond)
			defer ticker.Stop()

			for {
				select {
				case <-done:
					draw(state, refresh)
					return
				case <-ticker.C:
					draw(state, refresh)
				}
			}
		}()

		// This sync.Once var is used to receive the notification
		// that there was at least 1 line read from the source
		var notify sync.Once
		notifycb := func() {
			// close the ready channel so others can be notified
			// that there's at least 1 line in the buffer
			close(s.ready)
		}
		scanner := bufio.NewScanner(s.in)
		for scanner.Scan() {
			txt := scanner.Text()
			s.lines = append(s.lines, NewRawLine(txt, s.enableSep))
			notify.Do(notifycb)

			go func() {
				defer func() { recover() }()
				refresh <- struct{}{}
			}()
		}
		trace("Read all from source")
	})
}

// Start starts
func (s *Source) Start(ctx context.Context) {
	go func() {
		defer s.OutputChannel.SendEndMark("end of input")

		for i := 0; i < len(s.lines); i++ {
			select {
			case <-ctx.Done():
				return
			case s.OutputChannel <- s.lines[i]:
				// no op
			}
		}
	}()
}

// Ready returns the "input ready" channel. It will be closed as soon as
// the first line of input is processed via Setup()
func (s *Source) Ready() <-chan struct{} {
	return s.ready
}
<|endoftext|>


<commit_before>package buffer

import (
	"io"
	"os"
)

const MAXINT64 = 9223372036854775807

// Buffer is used to Write() data which will be Read() later.
type Buffer interface {
	Len() int64 // How much data is Buffered in bytes
	Cap() int64 // How much data can be Buffered at once in bytes.
	io.Reader   // Read() will read from the top of the buffer [io.EOF if empty]
	io.Writer   // Write() will write to the end of the buffer [io.ErrShortWrite if not enough space]
	Reset()     // Truncates the buffer, Len() == 0.
}

// BufferAt is a buffer which supports io.ReaderAt and io.WriterAt
type BufferAt interface {
	Buffer
	io.ReaderAt
	io.WriterAt
}

func len64(p []byte) int64 {
	return int64(len(p))
}

// Gap returns buf.Cap() - buf.Len()
func Gap(buf Buffer) int64 {
	return buf.Cap() - buf.Len()
}

// Full returns true iff buf.Len() == buf.Cap()
func Full(buf Buffer) bool {
	return buf.Len() == buf.Cap()
}

// Empty returns false iff buf.Len() == 0
func Empty(buf Buffer) bool {
	return buf.Len() == 0
}

// NewUnboundedBuffer returns a Buffer which buffers "mem" bytes to memory
// and then creates file's of size "file" to buffer above "mem" bytes.
func NewUnboundedBuffer(mem, file int64) Buffer {
	return NewMulti(New(mem), NewPartition(NewFilePool(file, os.TempDir())))
}
<commit_msg>added package doc<commit_after>// buffer implements a series of Buffers which can be composed to implment complicated buffering strategies
package buffer

import (
	"io"
	"os"
)

const MAXINT64 = 9223372036854775807

// Buffer is used to Write() data which will be Read() later.
type Buffer interface {
	Len() int64 // How much data is Buffered in bytes
	Cap() int64 // How much data can be Buffered at once in bytes.
	io.Reader   // Read() will read from the top of the buffer [io.EOF if empty]
	io.Writer   // Write() will write to the end of the buffer [io.ErrShortWrite if not enough space]
	Reset()     // Truncates the buffer, Len() == 0.
}

// BufferAt is a buffer which supports io.ReaderAt and io.WriterAt
type BufferAt interface {
	Buffer
	io.ReaderAt
	io.WriterAt
}

func len64(p []byte) int64 {
	return int64(len(p))
}

// Gap returns buf.Cap() - buf.Len()
func Gap(buf Buffer) int64 {
	return buf.Cap() - buf.Len()
}

// Full returns true iff buf.Len() == buf.Cap()
func Full(buf Buffer) bool {
	return buf.Len() == buf.Cap()
}

// Empty returns false iff buf.Len() == 0
func Empty(buf Buffer) bool {
	return buf.Len() == 0
}

// NewUnboundedBuffer returns a Buffer which buffers "mem" bytes to memory
// and then creates file's of size "file" to buffer above "mem" bytes.
func NewUnboundedBuffer(mem, file int64) Buffer {
	return NewMulti(New(mem), NewPartition(NewFilePool(file, os.TempDir())))
}
<|endoftext|>


<commit_before>package epub

import (
	"archive/zip"
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
)

const (
	containerFilename     = "container.xml"
	containerFileTemplate = `<?xml version="1.0" encoding="UTF-8"?>
<container version="1.0" xmlns="urn:oasis:names:tc:opendocument:xmlns:container">
  <rootfiles>
    <rootfile full-path="%s/%s" media-type="application/oebps-package+xml" />
  </rootfiles>
</container>
`
	contentFolderName = "EPUB"
	// Permissions for any new directories we create
	dirPermissions = 0755
	// Permissions for any new files we create
	filePermissions   = 0644
	mediaTypeNcx      = "application/x-dtbncx+xml"
	mediaTypeEpub     = "application/epub+zip"
	mediaTypeXhtml    = "application/xhtml+xml"
	metaInfFolderName = "META-INF"
	mimetypeFilename  = "mimetype"
	pkgFilename       = "package.opf"
	sectionFileFormat = "section%04d.xhtml"
	tempDirPrefix     = "go-epub"
	xhtmlFolderName   = "xhtml"
)

func (e *Epub) Write(destFilePath string) error {
	tempDir, err := ioutil.TempDir("", tempDirPrefix)
	defer os.Remove(tempDir)
	if err != nil {
		log.Fatalf("os.Remove error: %s", err)
	}

	err = createEpubFolders(tempDir)
	if err != nil {
		return err
	}

	err = writeMimetype(tempDir)
	if err != nil {
		return err
	}

	err = e.writeSections(tempDir)
	if err != nil {
		return err
	}

	err = e.writeToc(tempDir)
	if err != nil {
		return err
	}

	err = writeContainerFile(tempDir)
	if err != nil {
		return err
	}

	err = e.writePackageFile(tempDir)
	if err != nil {
		return err
	}

	err = e.writeEpub(tempDir, destFilePath)
	if err != nil {
		return err
	}

	// TODO

	//	output, err := xml.MarshalIndent(e.toc.navDoc.xml, "", "  ")
	//	output = append([]byte(xhtmlDoctype), output...)

	output, err := xml.MarshalIndent(e.toc.ncxXml, "", "  ")

	//	output, err := xml.MarshalIndent(e.pkg.xml, "", "  ")

	output = append([]byte(xml.Header), output...)
	fmt.Println(string(output))

	return nil
}

func createEpubFolders(tempDir string) error {
	if err := os.Mkdir(
		filepath.Join(
			tempDir,
			contentFolderName,
		),
		dirPermissions); err != nil {
		return err
	}

	if err := os.Mkdir(
		filepath.Join(
			tempDir,
			contentFolderName,
			xhtmlFolderName,
		),
		dirPermissions); err != nil {
		return err
	}

	if err := os.Mkdir(
		filepath.Join(
			tempDir,
			metaInfFolderName,
		),
		dirPermissions); err != nil {
		return err
	}

	return nil
}

func writeContainerFile(tempDir string) error {
	containerFilePath := filepath.Join(tempDir, metaInfFolderName, containerFilename)
	if err := ioutil.WriteFile(
		containerFilePath,
		[]byte(
			fmt.Sprintf(
				containerFileTemplate,
				contentFolderName,
				pkgFilename,
			),
		),
		filePermissions,
	); err != nil {
		return err
	}

	return nil
}

/*
    type Rootfile struct {
        FullPath string `xml:"full-path,attr"`
        MediaType string `xml:"media-type,attr"`
    }

    type Rootfiles struct {
        Rootfile []Rootfile `xml:"rootfile"`
    }

    type Container struct {
        XMLName xml.Name `xml:"urn:oasis:names:tc:opendocument:xmlns:container container"`
        Version string `xml:"version,attr"`
        Rootfiles Rootfiles `xml:"rootfiles"`
    }

    v := &Container{Version: containerVersion}

    v.Rootfiles.Rootfile = append(
        v.Rootfiles.Rootfile,
        Rootfile{
            FullPath: filepath.Join(oebpsFolderName, contentFilename),
            MediaType: mediaTypeOebpsPackage,
        },
    )

    output, err := xml.MarshalIndent(v, "", `   `)
    if err != nil {
        return err
    }
    // Add the xml header to the output
    containerContent := append([]byte(xml.Header), output...)
    // It's generally nice to have files end with a newline
    containerContent = append(containerContent, "\n"...)

    containerFilePath := filepath.Join(metaInfFolderPath, containerFilename)
    if err := ioutil.WriteFile(containerFilePath, containerContent, filePermissions); err != nil {
        return err
    }

    return nil
}

/*
func writeContentFile(tempDir string) error {
    oebpsFolderPath := filepath.Join(tempDir, oebpsFolderName)
    if err := os.Mkdir(oebpsFolderPath, dirPermissions); err != nil {
        return err
    }

    type Package struct {
        XMLName xml.Name `xml:"http://www.idpf.org/2007/opf package"`
        UniqueIdentifier string `xml:"unique-identifier,attr"`
        Version string `xml:"version,attr"`
    }

    type Rootfile struct {
        FullPath string `xml:"full-path,attr"`
        MediaType string `xml:"media-type,attr"`
    }

    type Rootfiles struct {
        Rootfile []Rootfile `xml:"rootfile"`
    }

    type Container struct {
        XMLName xml.Name `xml:"urn:oasis:names:tc:opendocument:xmlns:container container"`
        Version string `xml:"version,attr"`
        Rootfiles Rootfiles `xml:"rootfiles"`
    }

    v := &Container{Version: containerVersion}

    v.Rootfiles.Rootfile = append(
        v.Rootfiles.Rootfile,
        Rootfile{
            FullPath: filepath.Join(oebpsFolderName, contentFilename),
            MediaType: mediaTypeOebpsPackage,
        },
    )

    output, err := xml.MarshalIndent(v, "", `   `)
    if err != nil {
        return err
    }
    // Add the xml header to the output
    containerContent := append([]byte(xml.Header), output...)
    // It's generally nice to have files end with a newline
    containerContent = append(containerContent, "\n"...)

    containerFilePath := filepath.Join(metaInfFolderPath, containerFilename)
    if err := ioutil.WriteFile(containerFilePath, containerContent, filePermissions); err != nil {
        return err
    }

    return nil
}
*/
func writeMimetype(tempDir string) error {
	mimetypeFilePath := filepath.Join(tempDir, mimetypeFilename)

	if err := ioutil.WriteFile(mimetypeFilePath, []byte(mediaTypeEpub), filePermissions); err != nil {
		return err
	}

	return nil
}

func (e *Epub) writeEpub(tempDir string, destFilePath string) error {
	f, err := os.Create(destFilePath)
	if err != nil {
		log.Fatalf("os.Create error: %s", err)
	}
	defer func() {
		if err := f.Close(); err != nil {
			log.Fatalf("os.File.Close error: %s", err)
		}
	}()

	z := zip.NewWriter(f)
	defer func() {
		if err := z.Close(); err != nil {
			log.Fatalf("zip.Writer.Close error: %s", err)
		}
	}()

	skipMimetypeFile := false

	var addFileToZip = func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Get the path of the file relative to the folder we're zipping
		relativePath, err := filepath.Rel(tempDir, path)
		if err != nil {
			log.Fatalf("filepath.Rel error: %s", err)
		}

		// Only include regular files, not directories
		if !info.Mode().IsRegular() {
			return nil
		}

		var w io.Writer
		if path == filepath.Join(tempDir, mimetypeFilename) {
			// Skip the mimetype file
			if skipMimetypeFile == true {
				return nil
			}
			// The mimetype file must be uncompressed according to the EPUB spec
			w, err = z.CreateHeader(&zip.FileHeader{
				Name:   relativePath,
				Method: zip.Store,
			})
		} else {
			w, err = z.Create(relativePath)
		}
		if err != nil {
			log.Fatalf("zip.Writer.Create error: %s", err)
		}

		r, err := os.Open(path)
		if err != nil {
			log.Fatalf("os.Open error: %s", err)
		}
		defer func() {
			if err := r.Close(); err != nil {
				log.Fatalf("os.File.Close error: %s", err)
			}
		}()

		_, err = io.Copy(w, r)
		if err != nil {
			log.Fatalf("io.Copy error: %s", err)
		}

		return nil
	}

	// Add the mimetype file first
	mimetypeFilePath := filepath.Join(tempDir, mimetypeFilename)
	mimetypeInfo, err := os.Lstat(mimetypeFilePath)
	if err != nil {
		log.Fatalf("os.Lstat error: %s", err)
	}
	addFileToZip(mimetypeFilePath, mimetypeInfo, nil)

	skipMimetypeFile = true

	err = filepath.Walk(tempDir, addFileToZip)
	if err != nil {
		log.Fatalf("os.Lstat error: %s", err)
	}

	return nil
}

func (e *Epub) writePackageFile(tempDir string) error {
	err := e.pkg.write(tempDir)
	if err != nil {
		return err
	}

	return nil
}

func (e *Epub) writeSections(tempDir string) error {
	for i, section := range e.sections {
		sectionIndex := i + 1
		sectionFilename := fmt.Sprintf(sectionFileFormat, sectionIndex)
		sectionFilePath := filepath.Join(tempDir, contentFolderName, xhtmlFolderName, sectionFilename)

		if err := section.write(sectionFilePath); err != nil {
			return err
		}

		relativePath := filepath.Join(xhtmlFolderName, sectionFilename)
		e.toc.addSection(sectionIndex, section.Title(), relativePath)
		e.pkg.addToManifest(sectionFilename, relativePath, mediaTypeXhtml, "")
		e.pkg.addToSpine(sectionFilename)
	}

	return nil
}

func (e *Epub) writeToc(tempDir string) error {
	e.pkg.addToManifest(tocNavItemId, tocNavFilename, mediaTypeXhtml, tocNavItemProperties)
	e.pkg.addToManifest(tocNcxItemId, tocNcxFilename, mediaTypeNcx, "")

	err := e.toc.write(tempDir)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Comment out some debugging code<commit_after>package epub

import (
	"archive/zip"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
)

const (
	containerFilename     = "container.xml"
	containerFileTemplate = `<?xml version="1.0" encoding="UTF-8"?>
<container version="1.0" xmlns="urn:oasis:names:tc:opendocument:xmlns:container">
  <rootfiles>
    <rootfile full-path="%s/%s" media-type="application/oebps-package+xml" />
  </rootfiles>
</container>
`
	contentFolderName = "EPUB"
	// Permissions for any new directories we create
	dirPermissions = 0755
	// Permissions for any new files we create
	filePermissions   = 0644
	mediaTypeNcx      = "application/x-dtbncx+xml"
	mediaTypeEpub     = "application/epub+zip"
	mediaTypeXhtml    = "application/xhtml+xml"
	metaInfFolderName = "META-INF"
	mimetypeFilename  = "mimetype"
	pkgFilename       = "package.opf"
	sectionFileFormat = "section%04d.xhtml"
	tempDirPrefix     = "go-epub"
	xhtmlFolderName   = "xhtml"
)

func (e *Epub) Write(destFilePath string) error {
	tempDir, err := ioutil.TempDir("", tempDirPrefix)
	defer os.Remove(tempDir)
	if err != nil {
		log.Fatalf("os.Remove error: %s", err)
	}

	err = createEpubFolders(tempDir)
	if err != nil {
		return err
	}

	err = writeMimetype(tempDir)
	if err != nil {
		return err
	}

	err = e.writeSections(tempDir)
	if err != nil {
		return err
	}

	err = e.writeToc(tempDir)
	if err != nil {
		return err
	}

	err = writeContainerFile(tempDir)
	if err != nil {
		return err
	}

	err = e.writePackageFile(tempDir)
	if err != nil {
		return err
	}

	err = e.writeEpub(tempDir, destFilePath)
	if err != nil {
		return err
	}

	// TODO

	//	output, err := xml.MarshalIndent(e.toc.navDoc.xml, "", "  ")
	//	output = append([]byte(xhtmlDoctype), output...)

	//	output, err := xml.MarshalIndent(e.pkg.xml, "", "  ")

	//  output, err := xml.MarshalIndent(e.toc.ncxXml, "", "  ")
	//	output = append([]byte(xml.Header), output...)
	//	fmt.Println(string(output))

	return nil
}

func createEpubFolders(tempDir string) error {
	if err := os.Mkdir(
		filepath.Join(
			tempDir,
			contentFolderName,
		),
		dirPermissions); err != nil {
		return err
	}

	if err := os.Mkdir(
		filepath.Join(
			tempDir,
			contentFolderName,
			xhtmlFolderName,
		),
		dirPermissions); err != nil {
		return err
	}

	if err := os.Mkdir(
		filepath.Join(
			tempDir,
			metaInfFolderName,
		),
		dirPermissions); err != nil {
		return err
	}

	return nil
}

func writeContainerFile(tempDir string) error {
	containerFilePath := filepath.Join(tempDir, metaInfFolderName, containerFilename)
	if err := ioutil.WriteFile(
		containerFilePath,
		[]byte(
			fmt.Sprintf(
				containerFileTemplate,
				contentFolderName,
				pkgFilename,
			),
		),
		filePermissions,
	); err != nil {
		return err
	}

	return nil
}

/*
    type Rootfile struct {
        FullPath string `xml:"full-path,attr"`
        MediaType string `xml:"media-type,attr"`
    }

    type Rootfiles struct {
        Rootfile []Rootfile `xml:"rootfile"`
    }

    type Container struct {
        XMLName xml.Name `xml:"urn:oasis:names:tc:opendocument:xmlns:container container"`
        Version string `xml:"version,attr"`
        Rootfiles Rootfiles `xml:"rootfiles"`
    }

    v := &Container{Version: containerVersion}

    v.Rootfiles.Rootfile = append(
        v.Rootfiles.Rootfile,
        Rootfile{
            FullPath: filepath.Join(oebpsFolderName, contentFilename),
            MediaType: mediaTypeOebpsPackage,
        },
    )

    output, err := xml.MarshalIndent(v, "", `   `)
    if err != nil {
        return err
    }
    // Add the xml header to the output
    containerContent := append([]byte(xml.Header), output...)
    // It's generally nice to have files end with a newline
    containerContent = append(containerContent, "\n"...)

    containerFilePath := filepath.Join(metaInfFolderPath, containerFilename)
    if err := ioutil.WriteFile(containerFilePath, containerContent, filePermissions); err != nil {
        return err
    }

    return nil
}

/*
func writeContentFile(tempDir string) error {
    oebpsFolderPath := filepath.Join(tempDir, oebpsFolderName)
    if err := os.Mkdir(oebpsFolderPath, dirPermissions); err != nil {
        return err
    }

    type Package struct {
        XMLName xml.Name `xml:"http://www.idpf.org/2007/opf package"`
        UniqueIdentifier string `xml:"unique-identifier,attr"`
        Version string `xml:"version,attr"`
    }

    type Rootfile struct {
        FullPath string `xml:"full-path,attr"`
        MediaType string `xml:"media-type,attr"`
    }

    type Rootfiles struct {
        Rootfile []Rootfile `xml:"rootfile"`
    }

    type Container struct {
        XMLName xml.Name `xml:"urn:oasis:names:tc:opendocument:xmlns:container container"`
        Version string `xml:"version,attr"`
        Rootfiles Rootfiles `xml:"rootfiles"`
    }

    v := &Container{Version: containerVersion}

    v.Rootfiles.Rootfile = append(
        v.Rootfiles.Rootfile,
        Rootfile{
            FullPath: filepath.Join(oebpsFolderName, contentFilename),
            MediaType: mediaTypeOebpsPackage,
        },
    )

    output, err := xml.MarshalIndent(v, "", `   `)
    if err != nil {
        return err
    }
    // Add the xml header to the output
    containerContent := append([]byte(xml.Header), output...)
    // It's generally nice to have files end with a newline
    containerContent = append(containerContent, "\n"...)

    containerFilePath := filepath.Join(metaInfFolderPath, containerFilename)
    if err := ioutil.WriteFile(containerFilePath, containerContent, filePermissions); err != nil {
        return err
    }

    return nil
}
*/
func writeMimetype(tempDir string) error {
	mimetypeFilePath := filepath.Join(tempDir, mimetypeFilename)

	if err := ioutil.WriteFile(mimetypeFilePath, []byte(mediaTypeEpub), filePermissions); err != nil {
		return err
	}

	return nil
}

func (e *Epub) writeEpub(tempDir string, destFilePath string) error {
	f, err := os.Create(destFilePath)
	if err != nil {
		log.Fatalf("os.Create error: %s", err)
	}
	defer func() {
		if err := f.Close(); err != nil {
			log.Fatalf("os.File.Close error: %s", err)
		}
	}()

	z := zip.NewWriter(f)
	defer func() {
		if err := z.Close(); err != nil {
			log.Fatalf("zip.Writer.Close error: %s", err)
		}
	}()

	skipMimetypeFile := false

	var addFileToZip = func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Get the path of the file relative to the folder we're zipping
		relativePath, err := filepath.Rel(tempDir, path)
		if err != nil {
			log.Fatalf("filepath.Rel error: %s", err)
		}

		// Only include regular files, not directories
		if !info.Mode().IsRegular() {
			return nil
		}

		var w io.Writer
		if path == filepath.Join(tempDir, mimetypeFilename) {
			// Skip the mimetype file
			if skipMimetypeFile == true {
				return nil
			}
			// The mimetype file must be uncompressed according to the EPUB spec
			w, err = z.CreateHeader(&zip.FileHeader{
				Name:   relativePath,
				Method: zip.Store,
			})
		} else {
			w, err = z.Create(relativePath)
		}
		if err != nil {
			log.Fatalf("zip.Writer.Create error: %s", err)
		}

		r, err := os.Open(path)
		if err != nil {
			log.Fatalf("os.Open error: %s", err)
		}
		defer func() {
			if err := r.Close(); err != nil {
				log.Fatalf("os.File.Close error: %s", err)
			}
		}()

		_, err = io.Copy(w, r)
		if err != nil {
			log.Fatalf("io.Copy error: %s", err)
		}

		return nil
	}

	// Add the mimetype file first
	mimetypeFilePath := filepath.Join(tempDir, mimetypeFilename)
	mimetypeInfo, err := os.Lstat(mimetypeFilePath)
	if err != nil {
		log.Fatalf("os.Lstat error: %s", err)
	}
	addFileToZip(mimetypeFilePath, mimetypeInfo, nil)

	skipMimetypeFile = true

	err = filepath.Walk(tempDir, addFileToZip)
	if err != nil {
		log.Fatalf("os.Lstat error: %s", err)
	}

	return nil
}

func (e *Epub) writePackageFile(tempDir string) error {
	err := e.pkg.write(tempDir)
	if err != nil {
		return err
	}

	return nil
}

func (e *Epub) writeSections(tempDir string) error {
	for i, section := range e.sections {
		sectionIndex := i + 1
		sectionFilename := fmt.Sprintf(sectionFileFormat, sectionIndex)
		sectionFilePath := filepath.Join(tempDir, contentFolderName, xhtmlFolderName, sectionFilename)

		if err := section.write(sectionFilePath); err != nil {
			return err
		}

		relativePath := filepath.Join(xhtmlFolderName, sectionFilename)
		e.toc.addSection(sectionIndex, section.Title(), relativePath)
		e.pkg.addToManifest(sectionFilename, relativePath, mediaTypeXhtml, "")
		e.pkg.addToSpine(sectionFilename)
	}

	return nil
}

func (e *Epub) writeToc(tempDir string) error {
	e.pkg.addToManifest(tocNavItemId, tocNavFilename, mediaTypeXhtml, tocNavItemProperties)
	e.pkg.addToManifest(tocNcxItemId, tocNcxFilename, mediaTypeNcx, "")

	err := e.toc.write(tempDir)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// (C) 2016 by Ricardo Branco
//
// MIT License
//
// v0.3.3
//
// TODO:
// + Support HMAC
// + Read filenames from file
// + Support -c option like md5sum(1)
// + Use getopt
// + Use different output formats for display

package main

import (
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"flag"
	"fmt"
	ossl_blake2b "github.com/ricardobranco777/dgst/blake2b"
	ossl_blake2s "github.com/ricardobranco777/dgst/blake2s"
	ossl_md4 "github.com/ricardobranco777/dgst/md4"
	ossl_md5 "github.com/ricardobranco777/dgst/md5"
	ossl_ripemd160 "github.com/ricardobranco777/dgst/ripemd160"
	ossl_sha1 "github.com/ricardobranco777/dgst/sha1"
	ossl_sha256 "github.com/ricardobranco777/dgst/sha256"
	ossl_sha512 "github.com/ricardobranco777/dgst/sha512"
	"golang.org/x/crypto/blake2b"
	"golang.org/x/crypto/blake2s"
	"golang.org/x/crypto/md4"
	"golang.org/x/crypto/ripemd160"
	"golang.org/x/crypto/sha3"
	"hash"
	"io"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
)

var hashes = map[string]*struct {
	sum string
	hash.Hash
	size int
}{
	"BLAKE2b256": {"", nil, 32},
	"BLAKE2b384": {"", nil, 48},
	"BLAKE2b512": {"", nil, 64},
	"BLAKE2s256": {"", nil, 32},
	"MD4":        {"", nil, 16},
	"MD5":        {"", nil, 16},
	"RIPEMD160":  {"", nil, 20},
	"SHA1":       {"", nil, 20},
	"SHA224":     {"", nil, 28},
	"SHA256":     {"", nil, 32},
	"SHA384":     {"", nil, 48},
	"SHA512":     {"", nil, 64},
	"SHA512-224": {"", nil, 28},
	"SHA512-256": {"", nil, 32},
	"SHA3-224":   {"", nil, 28},
	"SHA3-256":   {"", nil, 32},
	"SHA3-384":   {"", nil, 48},
	"SHA3-512":   {"", nil, 64},
}

// The keys to the above dictionary in sorted order
var keys []string

var progname string

var done chan error

func main() {
	progname = path.Base(os.Args[0])

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [OPTIONS] [-s STRING...]|[FILE... DIRECTORY...]\n\n", progname)
		flag.PrintDefaults()
	}

	chosen := make(map[string]*bool)
	for h := range hashes {
		chosen[h] = flag.Bool(strings.ToLower(h), false, fmt.Sprintf("%s algorithm", h))
	}

	var all, ossl, isString *bool
	all = flag.Bool("all", false, "all algorithms")
	ossl = flag.Bool("ossl", false, "use OpenSSL")
	isString = flag.Bool("s", false, "treat arguments as strings")

	var size_hashes = map[int]*struct {
		hashes []string
		set    *bool
	}{
		128: {},
		160: {},
		224: {},
		256: {},
		384: {},
		512: {},
	}

	for h := range hashes {
		size_hashes[hashes[h].size*8].hashes = append(size_hashes[hashes[h].size*8].hashes, h)
	}

	for size := range size_hashes {
		sizeStr := strconv.Itoa(size)
		if len(size_hashes[size].hashes) > 0 {
			size_hashes[size].set = flag.Bool(sizeStr, false, "all "+sizeStr+" bits algorithms")
		}
	}
	var sha3Opt, sha2Opt, blake2Opt *bool
	sha3Opt = flag.Bool("sha3", false, "all SHA-3 algorithms")
	sha2Opt = flag.Bool("sha2", false, "all SHA-2 algorithms")
	blake2Opt = flag.Bool("blake2", false, "all BLAKE2 algorithms")

	flag.Parse()

	for size := range size_hashes {
		if !*(size_hashes[size].set) {
			continue
		}
		for _, h := range size_hashes[size].hashes {
			*chosen[h] = true
		}
	}

	if *sha3Opt {
		for h := range hashes {
			if strings.HasPrefix(h, "SHA3-") {
				*chosen[h] = true
			}
		}
	}
	if *sha2Opt {
		for h := range hashes {
			if strings.HasPrefix(h, "SHA2") || strings.HasPrefix(h, "SHA384") || strings.HasPrefix(h, "SHA512") {
				*chosen[h] = true
			}
		}
	}
	if *blake2Opt {
		for h := range hashes {
			if strings.HasPrefix(h, "BLAKE2") {
				*chosen[h] = true
			}
		}
	}

	if !(*all || choices(chosen)) {
		*all = true
	}

	for h := range chosen {
		if (*all && *chosen[h]) || !(*all || *chosen[h]) {
			delete(hashes, h)
		}
	}

	for k := range hashes {
		keys = append(keys, k)
		switch k {
		case "BLAKE2b256":
			if *ossl {
				hashes[k].Hash = blake2_(ossl_blake2b.New256)
			} else {
				hashes[k].Hash = blake2_(blake2b.New256)
			}
		case "BLAKE2b384":
			if *ossl {
				hashes[k].Hash = blake2_(ossl_blake2b.New384)
			} else {
				hashes[k].Hash = blake2_(blake2b.New384)
			}
		case "BLAKE2b512":
			if *ossl {
				hashes[k].Hash = blake2_(ossl_blake2b.New512)
			} else {
				hashes[k].Hash = blake2_(blake2b.New512)
			}
		case "BLAKE2s256":
			if *ossl {
				hashes[k].Hash = blake2_(ossl_blake2s.New256)
			} else {
				hashes[k].Hash = blake2_(blake2s.New256)
			}
		case "MD4":
			if *ossl {
				hashes[k].Hash = ossl_md4.New()
			} else {
				hashes[k].Hash = md4.New()
			}
		case "MD5":
			if *ossl {
				hashes[k].Hash = ossl_md5.New()
			} else {
				hashes[k].Hash = md5.New()
			}
		case "RIPEMD160":
			if *ossl {
				hashes[k].Hash = ossl_ripemd160.New()
			} else {
				hashes[k].Hash = ripemd160.New()
			}
		case "SHA1":
			if *ossl {
				hashes[k].Hash = ossl_sha1.New()
			} else {
				hashes[k].Hash = sha1.New()
			}
		case "SHA224":
			if *ossl {
				hashes[k].Hash = ossl_sha256.New224()
			} else {
				hashes[k].Hash = sha256.New224()
			}
		case "SHA256":
			if *ossl {
				hashes[k].Hash = ossl_sha256.New()
			} else {
				hashes[k].Hash = sha256.New()
			}
		case "SHA384":
			if *ossl {
				hashes[k].Hash = ossl_sha512.New384()
			} else {
				hashes[k].Hash = sha512.New384()
			}
		case "SHA512":
			if *ossl {
				hashes[k].Hash = ossl_sha512.New()
			} else {
				hashes[k].Hash = sha512.New()
			}
		case "SHA512-224":
			hashes[k].Hash = sha512.New512_224()
		case "SHA512-256":
			hashes[k].Hash = sha512.New512_256()
		case "SHA3-224":
			hashes[k].Hash = sha3.New224()
		case "SHA3-256":
			hashes[k].Hash = sha3.New256()
		case "SHA3-384":
			hashes[k].Hash = sha3.New384()
		case "SHA3-512":
			hashes[k].Hash = sha3.New512()
		}
	}

	sort.Strings(keys)

	if len(hashes) == 0 {
		flag.Usage()
		os.Exit(1)
	}

	if flag.NArg() == 0 {
		hash_stdin()
		os.Exit(0)
	}

	var errors bool

	if *isString {
		for _, s := range flag.Args() {
			hash_string(s)
		}
	} else {
		done = make(chan error, len(hashes))
		defer close(done)

		for _, pathname := range flag.Args() {
			info, err := os.Stat(pathname)
			if err != nil {
				fmt.Fprintf(os.Stderr, "%s: %v\n", progname, err)
				errors = true
			} else if info.IsDir() {
				errors = hash_dir(pathname)
			} else {
				errors = hash_file(pathname)
			}
		}
	}

	if errors {
		os.Exit(1)
	}
}

// Wrapper for the Blake2 New() methods that needs an optional for MAC
func blake2_(f func([]byte) (hash.Hash, error)) hash.Hash {
	h, _ := f(nil)
	return h
}

// Returns true if at least some algorithm was specified on the command line
func choices(chosen map[string]*bool) bool {
	for h := range chosen {
		if *chosen[h] {
			return true
		}
	}
	return false
}

func display() {
	for _, k := range keys {
		fmt.Println(hashes[k].sum)
	}
}

func hash_string(str string) {
	var wg sync.WaitGroup
	wg.Add(len(hashes))
	for h := range hashes {
		go func(h string) {
			defer wg.Done()
			hashes[h].Write([]byte(str))
			hashes[h].sum = fmt.Sprintf("%s(\"%s\") = %x", h, str, hashes[h].Sum(nil))
			hashes[h].Reset()
		}(h)
	}
	wg.Wait()
	display()
}

func hash_file(filename string) (errors bool) {
	for h := range hashes {
		go func(h string) {
			f, err := os.Open(filename)
			if err != nil {
				done <- err
				return
			}
			defer f.Close()

			if _, err := io.Copy(hashes[h], f); err != nil {
				done <- err
				return
			}
			hashes[h].sum = fmt.Sprintf("%s(%s) = %x", h, filename, hashes[h].Sum(nil))
			hashes[h].Reset()
			done <- nil
		}(h)
	}
	for range hashes {
		err := <-done
		if err != nil {
			if (!errors) {
				fmt.Fprintf(os.Stderr, "%s: %v\n", progname, err)
			}
			errors = true
		}
	}
	if (!errors) {
		display()
	}
	return
}

func visit(path string, f os.FileInfo, err error) error {
	if err != nil && !os.IsPermission(err) {
		return err
	}
	if f.Mode().IsRegular() {
		hash_file(path)
	}
	return nil
}

func hash_dir(dir string) bool {
	err := filepath.Walk(dir, visit)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s: %v\n", progname, err)
		return false
	}
	return true
}

func hash_stdin() (errors bool) {
	for h := range hashes {
		go func(h string) {
			if _, err := io.Copy(hashes[h], os.Stdin); err != nil {
				done <- err
				return
			}
			hashes[h].sum = fmt.Sprintf("%s() = %x", h, hashes[h].Sum(nil))
			hashes[h].Reset()
			done <- nil
		}(h)
	}
	for range hashes {
		err := <-done
		if err != nil {
			if !errors {
				fmt.Fprintf(os.Stderr, "%s: %v\n", progname, err)
			}
			errors = true
		}
	}
	if (!errors) {
		display()
	}
	return
}
<commit_msg>go fmt<commit_after>// (C) 2016 by Ricardo Branco
//
// MIT License
//
// v0.3.3
//
// TODO:
// + Support HMAC
// + Read filenames from file
// + Support -c option like md5sum(1)
// + Use getopt
// + Use different output formats for display

package main

import (
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"flag"
	"fmt"
	ossl_blake2b "github.com/ricardobranco777/dgst/blake2b"
	ossl_blake2s "github.com/ricardobranco777/dgst/blake2s"
	ossl_md4 "github.com/ricardobranco777/dgst/md4"
	ossl_md5 "github.com/ricardobranco777/dgst/md5"
	ossl_ripemd160 "github.com/ricardobranco777/dgst/ripemd160"
	ossl_sha1 "github.com/ricardobranco777/dgst/sha1"
	ossl_sha256 "github.com/ricardobranco777/dgst/sha256"
	ossl_sha512 "github.com/ricardobranco777/dgst/sha512"
	"golang.org/x/crypto/blake2b"
	"golang.org/x/crypto/blake2s"
	"golang.org/x/crypto/md4"
	"golang.org/x/crypto/ripemd160"
	"golang.org/x/crypto/sha3"
	"hash"
	"io"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
)

var hashes = map[string]*struct {
	sum string
	hash.Hash
	size int
}{
	"BLAKE2b256": {"", nil, 32},
	"BLAKE2b384": {"", nil, 48},
	"BLAKE2b512": {"", nil, 64},
	"BLAKE2s256": {"", nil, 32},
	"MD4":        {"", nil, 16},
	"MD5":        {"", nil, 16},
	"RIPEMD160":  {"", nil, 20},
	"SHA1":       {"", nil, 20},
	"SHA224":     {"", nil, 28},
	"SHA256":     {"", nil, 32},
	"SHA384":     {"", nil, 48},
	"SHA512":     {"", nil, 64},
	"SHA512-224": {"", nil, 28},
	"SHA512-256": {"", nil, 32},
	"SHA3-224":   {"", nil, 28},
	"SHA3-256":   {"", nil, 32},
	"SHA3-384":   {"", nil, 48},
	"SHA3-512":   {"", nil, 64},
}

// The keys to the above dictionary in sorted order
var keys []string

var progname string

var done chan error

func main() {
	progname = path.Base(os.Args[0])

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [OPTIONS] [-s STRING...]|[FILE... DIRECTORY...]\n\n", progname)
		flag.PrintDefaults()
	}

	chosen := make(map[string]*bool)
	for h := range hashes {
		chosen[h] = flag.Bool(strings.ToLower(h), false, fmt.Sprintf("%s algorithm", h))
	}

	var all, ossl, isString *bool
	all = flag.Bool("all", false, "all algorithms")
	ossl = flag.Bool("ossl", false, "use OpenSSL")
	isString = flag.Bool("s", false, "treat arguments as strings")

	var size_hashes = map[int]*struct {
		hashes []string
		set    *bool
	}{
		128: {},
		160: {},
		224: {},
		256: {},
		384: {},
		512: {},
	}

	for h := range hashes {
		size_hashes[hashes[h].size*8].hashes = append(size_hashes[hashes[h].size*8].hashes, h)
	}

	for size := range size_hashes {
		sizeStr := strconv.Itoa(size)
		if len(size_hashes[size].hashes) > 0 {
			size_hashes[size].set = flag.Bool(sizeStr, false, "all "+sizeStr+" bits algorithms")
		}
	}
	var sha3Opt, sha2Opt, blake2Opt *bool
	sha3Opt = flag.Bool("sha3", false, "all SHA-3 algorithms")
	sha2Opt = flag.Bool("sha2", false, "all SHA-2 algorithms")
	blake2Opt = flag.Bool("blake2", false, "all BLAKE2 algorithms")

	flag.Parse()

	for size := range size_hashes {
		if !*(size_hashes[size].set) {
			continue
		}
		for _, h := range size_hashes[size].hashes {
			*chosen[h] = true
		}
	}

	if *sha3Opt {
		for h := range hashes {
			if strings.HasPrefix(h, "SHA3-") {
				*chosen[h] = true
			}
		}
	}
	if *sha2Opt {
		for h := range hashes {
			if strings.HasPrefix(h, "SHA2") || strings.HasPrefix(h, "SHA384") || strings.HasPrefix(h, "SHA512") {
				*chosen[h] = true
			}
		}
	}
	if *blake2Opt {
		for h := range hashes {
			if strings.HasPrefix(h, "BLAKE2") {
				*chosen[h] = true
			}
		}
	}

	if !(*all || choices(chosen)) {
		*all = true
	}

	for h := range chosen {
		if (*all && *chosen[h]) || !(*all || *chosen[h]) {
			delete(hashes, h)
		}
	}

	for k := range hashes {
		keys = append(keys, k)
		switch k {
		case "BLAKE2b256":
			if *ossl {
				hashes[k].Hash = blake2_(ossl_blake2b.New256)
			} else {
				hashes[k].Hash = blake2_(blake2b.New256)
			}
		case "BLAKE2b384":
			if *ossl {
				hashes[k].Hash = blake2_(ossl_blake2b.New384)
			} else {
				hashes[k].Hash = blake2_(blake2b.New384)
			}
		case "BLAKE2b512":
			if *ossl {
				hashes[k].Hash = blake2_(ossl_blake2b.New512)
			} else {
				hashes[k].Hash = blake2_(blake2b.New512)
			}
		case "BLAKE2s256":
			if *ossl {
				hashes[k].Hash = blake2_(ossl_blake2s.New256)
			} else {
				hashes[k].Hash = blake2_(blake2s.New256)
			}
		case "MD4":
			if *ossl {
				hashes[k].Hash = ossl_md4.New()
			} else {
				hashes[k].Hash = md4.New()
			}
		case "MD5":
			if *ossl {
				hashes[k].Hash = ossl_md5.New()
			} else {
				hashes[k].Hash = md5.New()
			}
		case "RIPEMD160":
			if *ossl {
				hashes[k].Hash = ossl_ripemd160.New()
			} else {
				hashes[k].Hash = ripemd160.New()
			}
		case "SHA1":
			if *ossl {
				hashes[k].Hash = ossl_sha1.New()
			} else {
				hashes[k].Hash = sha1.New()
			}
		case "SHA224":
			if *ossl {
				hashes[k].Hash = ossl_sha256.New224()
			} else {
				hashes[k].Hash = sha256.New224()
			}
		case "SHA256":
			if *ossl {
				hashes[k].Hash = ossl_sha256.New()
			} else {
				hashes[k].Hash = sha256.New()
			}
		case "SHA384":
			if *ossl {
				hashes[k].Hash = ossl_sha512.New384()
			} else {
				hashes[k].Hash = sha512.New384()
			}
		case "SHA512":
			if *ossl {
				hashes[k].Hash = ossl_sha512.New()
			} else {
				hashes[k].Hash = sha512.New()
			}
		case "SHA512-224":
			hashes[k].Hash = sha512.New512_224()
		case "SHA512-256":
			hashes[k].Hash = sha512.New512_256()
		case "SHA3-224":
			hashes[k].Hash = sha3.New224()
		case "SHA3-256":
			hashes[k].Hash = sha3.New256()
		case "SHA3-384":
			hashes[k].Hash = sha3.New384()
		case "SHA3-512":
			hashes[k].Hash = sha3.New512()
		}
	}

	sort.Strings(keys)

	if len(hashes) == 0 {
		flag.Usage()
		os.Exit(1)
	}

	if flag.NArg() == 0 {
		hash_stdin()
		os.Exit(0)
	}

	var errors bool

	if *isString {
		for _, s := range flag.Args() {
			hash_string(s)
		}
	} else {
		done = make(chan error, len(hashes))
		defer close(done)

		for _, pathname := range flag.Args() {
			info, err := os.Stat(pathname)
			if err != nil {
				fmt.Fprintf(os.Stderr, "%s: %v\n", progname, err)
				errors = true
			} else if info.IsDir() {
				errors = hash_dir(pathname)
			} else {
				errors = hash_file(pathname)
			}
		}
	}

	if errors {
		os.Exit(1)
	}
}

// Wrapper for the Blake2 New() methods that needs an optional for MAC
func blake2_(f func([]byte) (hash.Hash, error)) hash.Hash {
	h, _ := f(nil)
	return h
}

// Returns true if at least some algorithm was specified on the command line
func choices(chosen map[string]*bool) bool {
	for h := range chosen {
		if *chosen[h] {
			return true
		}
	}
	return false
}

func display() {
	for _, k := range keys {
		fmt.Println(hashes[k].sum)
	}
}

func hash_string(str string) {
	var wg sync.WaitGroup
	wg.Add(len(hashes))
	for h := range hashes {
		go func(h string) {
			defer wg.Done()
			hashes[h].Write([]byte(str))
			hashes[h].sum = fmt.Sprintf("%s(\"%s\") = %x", h, str, hashes[h].Sum(nil))
			hashes[h].Reset()
		}(h)
	}
	wg.Wait()
	display()
}

func hash_file(filename string) (errors bool) {
	for h := range hashes {
		go func(h string) {
			f, err := os.Open(filename)
			if err != nil {
				done <- err
				return
			}
			defer f.Close()

			if _, err := io.Copy(hashes[h], f); err != nil {
				done <- err
				return
			}
			hashes[h].sum = fmt.Sprintf("%s(%s) = %x", h, filename, hashes[h].Sum(nil))
			hashes[h].Reset()
			done <- nil
		}(h)
	}
	for range hashes {
		err := <-done
		if err != nil {
			if !errors {
				fmt.Fprintf(os.Stderr, "%s: %v\n", progname, err)
			}
			errors = true
		}
	}
	if !errors {
		display()
	}
	return
}

func visit(path string, f os.FileInfo, err error) error {
	if err != nil && !os.IsPermission(err) {
		return err
	}
	if f.Mode().IsRegular() {
		hash_file(path)
	}
	return nil
}

func hash_dir(dir string) bool {
	err := filepath.Walk(dir, visit)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s: %v\n", progname, err)
		return false
	}
	return true
}

func hash_stdin() (errors bool) {
	for h := range hashes {
		go func(h string) {
			if _, err := io.Copy(hashes[h], os.Stdin); err != nil {
				done <- err
				return
			}
			hashes[h].sum = fmt.Sprintf("%s() = %x", h, hashes[h].Sum(nil))
			hashes[h].Reset()
			done <- nil
		}(h)
	}
	for range hashes {
		err := <-done
		if err != nil {
			if !errors {
				fmt.Fprintf(os.Stderr, "%s: %v\n", progname, err)
			}
			errors = true
		}
	}
	if !errors {
		display()
	}
	return
}
<|endoftext|>


<commit_before>package sparta

import (
	"fmt"
	"strings"

	_ "github.com/aws/aws-lambda-go/lambda"        // Force dep to resolve
	_ "github.com/aws/aws-lambda-go/lambdacontext" // Force dep to resolve
)

////////////////////////////////////////////////////////////////////////////////
// Constants
////////////////////////////////////////////////////////////////////////////////

const (
	// ProperName is the DRY name definition
	ProperName = "Sparta"
)
const (
	// SpartaVersion defines the current Sparta release
	SpartaVersion = "1.8.0"
	// GoLambdaVersion is the Go version runtime used for the lambda function
	GoLambdaVersion = "go1.x"
	// LambdaBinaryTag is the build tag name used when building the binary
	LambdaBinaryTag = "lambdabinary"
)

var (
	// SpartaBinaryName is binary name that exposes the Go lambda function
	SpartaBinaryName = fmt.Sprintf("%s.lambda.amd64", ProperName)
)

const (
	// Custom Resource typename used to create new cloudFormationUserDefinedFunctionCustomResource
	cloudFormationLambda = "Custom::SpartaLambdaCustomResource"
	// divider length is the length of a divider in the text
	// based CLI output
	dividerLength = 48
)
const (
	// envVarLogLevel is the provision time debug value
	// carried into the execution environment
	envVarLogLevel = "SPARTA_LOG_LEVEL"
	// spartaEnvVarFunctionName is the name of this function in the
	// map. It's the function that will be registered to run
	// envVarFunctionName = "SPARTA_FUNC_NAME"
	// envVarDiscoveryInformation is the name of the discovery information
	// published into the environment
	envVarDiscoveryInformation = "SPARTA_DISCOVERY_INFO"
)

var (
	// internal logging header
	headerDivider = strings.Repeat("═", dividerLength)
)

// AWS Principal ARNs from http://docs.aws.amazon.com/general/latest/gr/aws-arns-and-namespaces.html
// See also
// http://docs.aws.amazon.com/general/latest/gr/rande.html
// for region specific principal names
const (
	// @enum AWSPrincipal
	APIGatewayPrincipal = "apigateway.amazonaws.com"
	// @enum AWSPrincipal
	CloudWatchEventsPrincipal = "events.amazonaws.com"
	// @enum AWSPrincipal
	SESPrincipal = "ses.amazonaws.com"
	// @enum AWSPrincipal
	SNSPrincipal = "sns.amazonaws.com"
	// @enum AWSPrincipal
	EC2Principal = "ec2.amazonaws.com"
	// @enum AWSPrincipal
	LambdaPrincipal = "lambda.amazonaws.com"
)

type contextKey int

const (
	// ContextKeyLogger is the request-independent *logrus.Logger
	// instance common to all requests
	ContextKeyLogger contextKey = iota
	// ContextKeyRequestLogger is the *logrus.Entry instance
	// that is annotated with request-identifying
	// information extracted from the AWS context object
	ContextKeyRequestLogger
	// ContextKeyLambdaContext is the *sparta.LambdaContext
	// pointer in the request
	// DEPRECATED
	ContextKeyLambdaContext
	// ContextKeyLambdaError is the possible error that was returned
	// from the lambda function
	ContextKeyLambdaError
	// ContextKeyLambdaResponse is the possible response that
	// was returned from the lambda function
	ContextKeyLambdaResponse
)
<commit_msg>Bump version<commit_after>package sparta

import (
	"fmt"
	"strings"

	_ "github.com/aws/aws-lambda-go/lambda"        // Force dep to resolve
	_ "github.com/aws/aws-lambda-go/lambdacontext" // Force dep to resolve
)

////////////////////////////////////////////////////////////////////////////////
// Constants
////////////////////////////////////////////////////////////////////////////////

const (
	// ProperName is the DRY name definition
	ProperName = "Sparta"
)
const (
	// SpartaVersion defines the current Sparta release
	SpartaVersion = "1.9.0"
	// GoLambdaVersion is the Go version runtime used for the lambda function
	GoLambdaVersion = "go1.x"
	// LambdaBinaryTag is the build tag name used when building the binary
	LambdaBinaryTag = "lambdabinary"
)

var (
	// SpartaBinaryName is binary name that exposes the Go lambda function
	SpartaBinaryName = fmt.Sprintf("%s.lambda.amd64", ProperName)
)

const (
	// Custom Resource typename used to create new cloudFormationUserDefinedFunctionCustomResource
	cloudFormationLambda = "Custom::SpartaLambdaCustomResource"
	// divider length is the length of a divider in the text
	// based CLI output
	dividerLength = 48
)
const (
	// envVarLogLevel is the provision time debug value
	// carried into the execution environment
	envVarLogLevel = "SPARTA_LOG_LEVEL"
	// spartaEnvVarFunctionName is the name of this function in the
	// map. It's the function that will be registered to run
	// envVarFunctionName = "SPARTA_FUNC_NAME"
	// envVarDiscoveryInformation is the name of the discovery information
	// published into the environment
	envVarDiscoveryInformation = "SPARTA_DISCOVERY_INFO"
)

var (
	// internal logging header
	headerDivider = strings.Repeat("═", dividerLength)
)

// AWS Principal ARNs from http://docs.aws.amazon.com/general/latest/gr/aws-arns-and-namespaces.html
// See also
// http://docs.aws.amazon.com/general/latest/gr/rande.html
// for region specific principal names
const (
	// @enum AWSPrincipal
	APIGatewayPrincipal = "apigateway.amazonaws.com"
	// @enum AWSPrincipal
	CloudWatchEventsPrincipal = "events.amazonaws.com"
	// @enum AWSPrincipal
	SESPrincipal = "ses.amazonaws.com"
	// @enum AWSPrincipal
	SNSPrincipal = "sns.amazonaws.com"
	// @enum AWSPrincipal
	EC2Principal = "ec2.amazonaws.com"
	// @enum AWSPrincipal
	LambdaPrincipal = "lambda.amazonaws.com"
)

type contextKey int

const (
	// ContextKeyLogger is the request-independent *logrus.Logger
	// instance common to all requests
	ContextKeyLogger contextKey = iota
	// ContextKeyRequestLogger is the *logrus.Entry instance
	// that is annotated with request-identifying
	// information extracted from the AWS context object
	ContextKeyRequestLogger
	// ContextKeyLambdaContext is the *sparta.LambdaContext
	// pointer in the request
	// DEPRECATED
	ContextKeyLambdaContext
	// ContextKeyLambdaError is the possible error that was returned
	// from the lambda function
	ContextKeyLambdaError
	// ContextKeyLambdaResponse is the possible response that
	// was returned from the lambda function
	ContextKeyLambdaResponse
)
<|endoftext|>


<commit_before>package migrations

import (
	"log"
	"regexp"
	"sync"

	"github.com/mitchellh/goamz/s3"
	"github.com/thomaso-mirodin/go-shorten/storage"
)

// This function is designed to assist with migrating an S3 storage with
// migrating from the original implementation to the V2 design. I.e. get to this:
//	 sha256(short)/
//	               long -> url
//	               change_history/
//	                              TimeInRFC3339Nano() -> {URL: oldURL, Owner: "TODO"}
// from this:
//	short -> url
func MigrateS3FromV1ToV2(storage *storage.S3, writenew bool, deleteold bool) {
	log.Println("Begining S3 migration from v1 design to v2")

	s3b := storage.Bucket

	var (
		wg     sync.WaitGroup
		marker string
	)

	for retryCount := 0; retryCount < 10; {
		resp, err := s3b.List("", "", marker, 1000)
		if err != nil {
			log.Println("Failed to list S3 bucket because: %v", err)
			retryCount++
			continue
		}

		log.Printf("List found %v available items", len(resp.Contents))
		for _, k := range resp.Contents {
			if isVersionedKey(k.Key) {
				log.Printf("Skipping '%v' because it looks like a versioned key", k.Key)
				continue
			}

			wg.Add(1)
			go func(k s3.Key) {
				defer wg.Done()

				b, err := s3b.Get(k.Key)
				if err != nil {
					log.Printf("Failed to get url from key '%s' because: %v", k.Key, err)
					return
				}

				short := k.Key
				url := string(b)

				log.Printf("Migrating the pair '%v'->'%v' to its new home", short, url)
				if writenew {
					err = storage.SaveName(short, url)
					if err != nil {
						log.Printf("Failed to migrate short '%s' to its new home because: %s", short, err)
						return // This is pretty important :D
					}

					if deleteold {
						if err := s3b.Del(short); err != nil {
							log.Printf("Failed to clean out old short code '%s' because: %s", short, err)
						}
					}
				}
			}(k)
		}

		wg.Wait()

		if resp.IsTruncated {
			marker = resp.NextMarker
		} else {
			break
		}
	}
}

var versionedKey = regexp.MustCompile("^v\\d+\\/.+")

func isVersionedKey(key string) bool {
	if versionedKey.MatchString(key) {
		return true
	}
	return false
}
<commit_msg>Missed updating the migration script with the new goamz<commit_after>package migrations

import (
	"log"
	"regexp"
	"sync"

	"github.com/goamz/goamz/s3"
	"github.com/thomaso-mirodin/go-shorten/storage"
)

// This function is designed to assist with migrating an S3 storage with
// migrating from the original implementation to the V2 design. I.e. get to this:
//	 sha256(short)/
//	               long -> url
//	               change_history/
//	                              TimeInRFC3339Nano() -> {URL: oldURL, Owner: "TODO"}
// from this:
//	short -> url
func MigrateS3FromV1ToV2(storage *storage.S3, writenew bool, deleteold bool) {
	log.Println("Begining S3 migration from v1 design to v2")

	s3b := storage.Bucket

	var (
		wg     sync.WaitGroup
		marker string
	)

	for retryCount := 0; retryCount < 10; {
		resp, err := s3b.List("", "", marker, 1000)
		if err != nil {
			log.Println("Failed to list S3 bucket because: %v", err)
			retryCount++
			continue
		}

		log.Printf("List found %v available items", len(resp.Contents))
		for _, k := range resp.Contents {
			if isVersionedKey(k.Key) {
				log.Printf("Skipping '%v' because it looks like a versioned key", k.Key)
				continue
			}

			wg.Add(1)
			go func(k s3.Key) {
				defer wg.Done()

				b, err := s3b.Get(k.Key)
				if err != nil {
					log.Printf("Failed to get url from key '%s' because: %v", k.Key, err)
					return
				}

				short := k.Key
				url := string(b)

				log.Printf("Migrating the pair '%v'->'%v' to its new home", short, url)
				if writenew {
					err = storage.SaveName(short, url)
					if err != nil {
						log.Printf("Failed to migrate short '%s' to its new home because: %s", short, err)
						return // This is pretty important :D
					}

					if deleteold {
						if err := s3b.Del(short); err != nil {
							log.Printf("Failed to clean out old short code '%s' because: %s", short, err)
						}
					}
				}
			}(k)
		}

		wg.Wait()

		if resp.IsTruncated {
			marker = resp.NextMarker
		} else {
			break
		}
	}
}

var versionedKey = regexp.MustCompile("^v\\d+\\/.+")

func isVersionedKey(key string) bool {
	if versionedKey.MatchString(key) {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package preconditions

import (
	"fmt"

	"github.com/Everlane/evan/common"

	"github.com/google/go-github/github"
)

type GithubCombinedStatusPrecondition struct{}

func (gh *GithubCombinedStatusPrecondition) Status(deployment common.Deployment) common.PreconditionResult {
	repo := deployment.Application().Repository()
	ref := deployment.SHA1()
	if ref == "" {
		ref = deployment.Ref()
	}

	client, err := common.GithubClient(deployment)
	if err != nil {
		return createResult(gh, err)
	}

	status, _, err := client.Repositories.GetCombinedStatus(repo.Owner(), repo.Name(), ref, nil)
	fmt.Printf("status: %+v\n", status)
	if err != nil {
		return createResult(gh, err)
	}

	var result error = nil
	if *status.State != "success" {
		result = fmt.Errorf("Non-success status for ref: %v", *status.State)
	}
	return createResult(gh, result)
}

// Require that the branch for deployment not be behind the default branch
// on GitHub. If `AutoMerge` is true then it will try to create a merge via
// the GitHub API if the deployed ref is behind.
type GithubRequireAheadPrecondition struct {
	AutoMerge bool
}

type GithubRequireAheadContext struct {
	RepoClient *common.GithubRepository
	RepoDetails *github.Repository
}

// Compares the ref being deployed against the default branch on GitHub to
// determine whether or not a merge needs to happen. Returns `false` if it's
// a force deployment.
func (gh *GithubRequireAheadPrecondition) NeedsMerge(deployment common.Deployment, ctx *GithubRequireAheadContext) (bool, error) {
	if deployment.IsForce() {
		return false, nil
	}

	base := *ctx.RepoDetails.DefaultBranch
	head := deployment.Ref()

	comparison, err := ctx.RepoClient.CompareCommits(base, head)
	if err != nil {
		return false, err
	}

	return (*comparison.BehindBy > 0), nil
}

// Creates merge commit to get the target branch (deployment's ref) up-to-date
// with the default branch of the repository.
func (gh *GithubRequireAheadPrecondition) Merge(deployment common.Deployment, ctx *GithubRequireAheadContext) (string, error) {
	base := deployment.Ref()
	head := *ctx.RepoDetails.DefaultBranch
	commitMessage := fmt.Sprintf("Merge '%v' into '%v'", head, base)
	commit, err := ctx.RepoClient.Merge(base, head, commitMessage)
	if err != nil {
		return "", err
	} else {
		return *commit.SHA, nil
	}
}

func (gh *GithubRequireAheadPrecondition) Status(deployment common.Deployment) common.PreconditionResult {
	githubClient, err := common.GithubClient(deployment)
	if err != nil {
		return createResult(gh, err)
	}

	repoClient := &common.GithubRepository{
		Repository:   deployment.Application().Repository(),
		GithubClient: githubClient,
	}
	repoDetails, err := repoClient.Get()
	if err != nil {
		return createResult(gh, err)
	}

	ctx := &GithubRequireAheadContext{
		RepoClient: repoClient,
		RepoDetails: repoDetails,
	}

	needsMerge, err := gh.NeedsMerge(deployment, ctx)
	if err != nil {
		return createResult(gh, err)
	}

	// Halt if we don't need to merge!
	if !needsMerge {
		return createResult(gh, nil)
	}

	if !gh.AutoMerge {
		return createResult(gh, fmt.Errorf("Merge needed for ref '%v'", deployment.Ref()))
	}

	sha1, err := gh.Merge(deployment, ctx)
	if err != nil {
		return createResult(gh, err)
	}
	// Update the SHA1 to point to the new merge commit
	deployment.SetSHA1(sha1)

	return createResult(gh, nil)
}
<commit_msg>Ensure there are status checks defined before checking the combined status<commit_after>package preconditions

import (
	"fmt"

	"github.com/Everlane/evan/common"

	"github.com/google/go-github/github"
)

type GithubCombinedStatusPrecondition struct{}

func (gh *GithubCombinedStatusPrecondition) Status(deployment common.Deployment) common.PreconditionResult {
	repo := deployment.Application().Repository()
	ref := deployment.SHA1()
	if ref == "" {
		ref = deployment.Ref()
	}

	client, err := common.GithubClient(deployment)
	if err != nil {
		return createResult(gh, err)
	}

	status, _, err := client.Repositories.GetCombinedStatus(repo.Owner(), repo.Name(), ref, nil)
	if err != nil {
		return createResult(gh, err)
	}

	state := *status.State
	totalCount := *status.TotalCount

	var result error = nil
	// GitHub will report a "pending" status if there are zero status checks,
	// therefore we want to only check the status if there are one or more
	// status checks.
	if totalCount > 0 && state != "success" {
		result = fmt.Errorf("Non-success status for ref: %v", *status.State)
	}
	return createResult(gh, result)
}

// Require that the branch for deployment not be behind the default branch
// on GitHub. If `AutoMerge` is true then it will try to create a merge via
// the GitHub API if the deployed ref is behind.
type GithubRequireAheadPrecondition struct {
	AutoMerge bool
}

type GithubRequireAheadContext struct {
	RepoClient *common.GithubRepository
	RepoDetails *github.Repository
}

// Compares the ref being deployed against the default branch on GitHub to
// determine whether or not a merge needs to happen. Returns `false` if it's
// a force deployment.
func (gh *GithubRequireAheadPrecondition) NeedsMerge(deployment common.Deployment, ctx *GithubRequireAheadContext) (bool, error) {
	if deployment.IsForce() {
		return false, nil
	}

	base := *ctx.RepoDetails.DefaultBranch
	head := deployment.Ref()

	comparison, err := ctx.RepoClient.CompareCommits(base, head)
	if err != nil {
		return false, err
	}

	return (*comparison.BehindBy > 0), nil
}

// Creates merge commit to get the target branch (deployment's ref) up-to-date
// with the default branch of the repository.
func (gh *GithubRequireAheadPrecondition) Merge(deployment common.Deployment, ctx *GithubRequireAheadContext) (string, error) {
	base := deployment.Ref()
	head := *ctx.RepoDetails.DefaultBranch
	commitMessage := fmt.Sprintf("Merge '%v' into '%v'", head, base)
	commit, err := ctx.RepoClient.Merge(base, head, commitMessage)
	if err != nil {
		return "", err
	} else {
		return *commit.SHA, nil
	}
}

func (gh *GithubRequireAheadPrecondition) Status(deployment common.Deployment) common.PreconditionResult {
	githubClient, err := common.GithubClient(deployment)
	if err != nil {
		return createResult(gh, err)
	}

	repoClient := &common.GithubRepository{
		Repository:   deployment.Application().Repository(),
		GithubClient: githubClient,
	}
	repoDetails, err := repoClient.Get()
	if err != nil {
		return createResult(gh, err)
	}

	ctx := &GithubRequireAheadContext{
		RepoClient: repoClient,
		RepoDetails: repoDetails,
	}

	needsMerge, err := gh.NeedsMerge(deployment, ctx)
	if err != nil {
		return createResult(gh, err)
	}

	// Halt if we don't need to merge!
	if !needsMerge {
		return createResult(gh, nil)
	}

	if !gh.AutoMerge {
		return createResult(gh, fmt.Errorf("Merge needed for ref '%v'", deployment.Ref()))
	}

	sha1, err := gh.Merge(deployment, ctx)
	if err != nil {
		return createResult(gh, err)
	}
	// Update the SHA1 to point to the new merge commit
	deployment.SetSHA1(sha1)

	return createResult(gh, nil)
}
<|endoftext|>


<commit_before>package prifisocks

import (
	"encoding/binary"
)

// Possible message/packet types.
const (
	DummyData = iota
	SocksConnect
	StallCommunication
	ResumeCommunication
	SocksData
	SocksClosed
	SocksError
)

// Size of the data wrapper header.
const (
	SocksPacketHeaderSize = uint16(10)
)

//SocksPacket represents the packet communicated across the network. It contains the header components and the data.
type SocksPacket() struct {
	// Header
	Type          uint16
	ID            uint32 // SOCKS5 Connection ID
	MessageLength uint16 // The length of useful data
	PacketLength  uint16 // The length of the packet including the header

	// Data
	Data []byte // The data segment of the packet (Always of size PacketLength-HeaderLength)
}

//ToBytes converts the SocksPacket to a byte array
func (d *SocksPacket) ToBytes() []byte {

	// Make sure the data and messagelength are of appropriate length
	d.Data, d.MessageLength = socksTrimAndPadPayload(d.Data, d.MessageLength, d.PacketLength)

	buffer := make([]byte, int(SocksPacketHeaderSize)) // Temporary byte buffer to store the header

	// Fill up the header buffer
	binary.BigEndian.PutUint16(buffer[0:2], d.Type)
	binary.BigEndian.PutUint32(buffer[2:6], d.ID)
	binary.BigEndian.PutUint16(buffer[6:8], d.MessageLength)
	binary.BigEndian.PutUint16(buffer[8:10], d.PacketLength)

	return append(buffer, d.Data...) // Append the data to the header
}

//NewSocksPacket creates a new socks packet
func NewSocksPacket(Type uint16, ID uint32, MessageLength uint16, PacketLength uint16, Data []byte) SocksPacket {

	// Make sure the received data and messagelength are of appropriate length
	Data, MessageLength = socksTrimAndPadPayload(Data, MessageLength, PacketLength)

	return SocksPacket{Type, ID, MessageLength, PacketLength, Data} //Create the new packet
}

/*
socksTrimAndPadPayload checks for consistency withing the data in the packet and fixes any inconsistencies

Properties:
	- Actual message length cannot exceed the maximum possible length (which is PacketLength - HeaderLength)
	- Data should always be at maximum possible length, padding is added if needed
*/
func socksTrimAndPadPayload(data []byte, messageLength uint16, packetLength uint16) ([]byte, uint16) {

	// Get the maximum possible length of the data
	maxMessageLength := packetLength - SocksPacketHeaderSize

	// Check if the message length exceeds the maximum possible length, and truncate the length
	if messageLength > maxMessageLength {
		messageLength = maxMessageLength
	}

	// Add the necessary padding to the data and truncate it to maximum length
	padding := make([]byte, maxMessageLength)
	data = append(data, padding...)
	data = data[:maxMessageLength]

	// Return the modified data and message length
	return data, messageLength
}

//ParseSocksPacketFromBytes extracts the SocksPacket packet from an array of bytes.
func ParseSocksPacketFromBytes(buffer []byte) SocksPacket {

	if len(buffer) < int(SocksPacketHeaderSize) {
		return NewSocksPacket(0, 0, 0, 0, make([]byte, 0))
	}
	//Extract the content of the header
	packetType, connID, messageLength, packetLength := ParseSocksHeaderFromBytes(buffer)

	//Construct and return a new packet
	if int(messageLength) <= len(buffer)-int(SocksPacketHeaderSize) {
		return NewSocksPacket(packetType, connID, messageLength, packetLength, buffer[SocksPacketHeaderSize:SocksPacketHeaderSize+messageLength])
	}
	return NewSocksPacket(packetType, connID, messageLength, packetLength, buffer[SocksPacketHeaderSize:])

}

//ParseSocksHeaderFromBytes extracts the SocksPacket header from an array of bytes.
func ParseSocksHeaderFromBytes(buffer []byte) (uint16, uint32, uint16, uint16) {

	//Extract the content of the header from the buffer
	packetType := binary.BigEndian.Uint16(buffer[0:2])
	connID := binary.BigEndian.Uint32(buffer[2:6])
	messageLength := binary.BigEndian.Uint16(buffer[6:8])
	packetLength := binary.BigEndian.Uint16(buffer[8:10])

	//Return the header components
	return packetType, connID, messageLength, packetLength
}
<commit_msg>Fixed the error<commit_after>package prifisocks

import (
	"encoding/binary"
)

// Possible message/packet types.
const (
	DummyData = iota
	SocksConnect
	StallCommunication
	ResumeCommunication
	SocksData
	SocksClosed
	SocksError
)

// Size of the data wrapper header.
const (
	SocksPacketHeaderSize = uint16(10)
)

//SocksPacket represents the packet communicated across the network. It contains the header components and the data.
type SocksPacket struct {
	// Header
	Type          uint16
	ID            uint32 // SOCKS5 Connection ID
	MessageLength uint16 // The length of useful data
	PacketLength  uint16 // The length of the packet including the header

	// Data
	Data []byte // The data segment of the packet (Always of size PacketLength-HeaderLength)
}

//ToBytes converts the SocksPacket to a byte array
func (d *SocksPacket) ToBytes() []byte {

	// Make sure the data and messagelength are of appropriate length
	d.Data, d.MessageLength = socksTrimAndPadPayload(d.Data, d.MessageLength, d.PacketLength)

	buffer := make([]byte, int(SocksPacketHeaderSize)) // Temporary byte buffer to store the header

	// Fill up the header buffer
	binary.BigEndian.PutUint16(buffer[0:2], d.Type)
	binary.BigEndian.PutUint32(buffer[2:6], d.ID)
	binary.BigEndian.PutUint16(buffer[6:8], d.MessageLength)
	binary.BigEndian.PutUint16(buffer[8:10], d.PacketLength)

	return append(buffer, d.Data...) // Append the data to the header
}

//NewSocksPacket creates a new socks packet
func NewSocksPacket(Type uint16, ID uint32, MessageLength uint16, PacketLength uint16, Data []byte) SocksPacket {

	// Make sure the received data and messagelength are of appropriate length
	Data, MessageLength = socksTrimAndPadPayload(Data, MessageLength, PacketLength)

	return SocksPacket{Type, ID, MessageLength, PacketLength, Data} //Create the new packet
}

/*
socksTrimAndPadPayload checks for consistency withing the data in the packet and fixes any inconsistencies

Properties:
	- Actual message length cannot exceed the maximum possible length (which is PacketLength - HeaderLength)
	- Data should always be at maximum possible length, padding is added if needed
*/
func socksTrimAndPadPayload(data []byte, messageLength uint16, packetLength uint16) ([]byte, uint16) {

	// Get the maximum possible length of the data
	maxMessageLength := packetLength - SocksPacketHeaderSize

	// Check if the message length exceeds the maximum possible length, and truncate the length
	if messageLength > maxMessageLength {
		messageLength = maxMessageLength
	}

	// Add the necessary padding to the data and truncate it to maximum length
	padding := make([]byte, maxMessageLength)
	data = append(data, padding...)
	data = data[:maxMessageLength]

	// Return the modified data and message length
	return data, messageLength
}

//ParseSocksPacketFromBytes extracts the SocksPacket packet from an array of bytes.
func ParseSocksPacketFromBytes(buffer []byte) SocksPacket {

	if len(buffer) < int(SocksPacketHeaderSize) {
		return NewSocksPacket(0, 0, 0, 0, make([]byte, 0))
	}
	//Extract the content of the header
	packetType, connID, messageLength, packetLength := ParseSocksHeaderFromBytes(buffer)

	//Construct and return a new packet
	if int(messageLength) <= len(buffer)-int(SocksPacketHeaderSize) {
		return NewSocksPacket(packetType, connID, messageLength, packetLength, buffer[SocksPacketHeaderSize:SocksPacketHeaderSize+messageLength])
	}
	return NewSocksPacket(packetType, connID, messageLength, packetLength, buffer[SocksPacketHeaderSize:])

}

//ParseSocksHeaderFromBytes extracts the SocksPacket header from an array of bytes.
func ParseSocksHeaderFromBytes(buffer []byte) (uint16, uint32, uint16, uint16) {

	//Extract the content of the header from the buffer
	packetType := binary.BigEndian.Uint16(buffer[0:2])
	connID := binary.BigEndian.Uint32(buffer[2:6])
	messageLength := binary.BigEndian.Uint16(buffer[6:8])
	packetLength := binary.BigEndian.Uint16(buffer[8:10])

	//Return the header components
	return packetType, connID, messageLength, packetLength
}
<|endoftext|>


<commit_before>package printer_test

import (
	"testing"

	"github.com/goccy/go-yaml/lexer"
	"github.com/goccy/go-yaml/printer"
)

func Test_Printer(t *testing.T) {
	yml := `---
text: aaaa
text2: aaaa
 bbbb
 cccc
 dddd
 eeee
text3: ffff
 gggg
 hhhh
 iiii
 jjjj
bool: true
number: 10
anchor: &x 1
alias: *x
`
	t.Run("print starting from tokens[3]", func(t *testing.T) {
		tokens := lexer.Tokenize(yml)
		var p printer.Printer
		actual := "\n" + p.PrintErrorToken(tokens[3], false)
		expect := `
   1 | ---
>  2 | text: aaaa
            ^
   3 | text2: aaaa
   4 |  bbbb
   5 |  cccc
   6 |  dddd
   7 |  eeee
   8 | `
		if actual != expect {
			t.Fatalf("unexpected output: expect:[%s]\n actual:[%s]", expect, actual)
		}
	})
	t.Run("print stargin from tokens[4]", func(t *testing.T) {
		tokens := lexer.Tokenize(yml)
		var p printer.Printer
		actual := "\n" + p.PrintErrorToken(tokens[4], false)
		expect := `
   1 | ---
   2 | text: aaaa
>  3 | text2: aaaa
   4 |  bbbb
   5 |  cccc
   6 |  dddd
   7 |  eeee
      ^
`
		if actual != expect {
			t.Fatalf("unexpected output: expect:[%s]\n actual:[%s]", expect, actual)
		}
	})
	t.Run("print starting from tokens[6]", func(t *testing.T) {
		tokens := lexer.Tokenize(yml)
		var p printer.Printer
		actual := "\n" + p.PrintErrorToken(tokens[6], false)
		expect := `
   1 | ---
   2 | text: aaaa
>  3 | text2: aaaa
   4 |  bbbb
   5 |  cccc
   6 |  dddd
   7 |  eeee
             ^
   8 | text3: ffff
   9 |  gggg
  10 |  hhhh
  11 |  iiii
  12 |  jjjj
  13 | `
		if actual != expect {
			t.Fatalf("unexpected output: expect:[%s]\n actual:[%s]", expect, actual)
		}
	})
	t.Run("output with color", func(t *testing.T) {
		t.Run("token6", func(t *testing.T) {
			tokens := lexer.Tokenize(yml)
			var p printer.Printer
			t.Logf("%s", p.PrintErrorToken(tokens[6], true))
		})
		t.Run("token9", func(t *testing.T) {
			tokens := lexer.Tokenize(yml)
			var p printer.Printer
			t.Logf("%s", p.PrintErrorToken(tokens[9], true))
		})
		t.Run("token12", func(t *testing.T) {
			tokens := lexer.Tokenize(yml)
			var p printer.Printer
			t.Logf("%s", p.PrintErrorToken(tokens[12], true))
		})
	})
}
<commit_msg>Increase coverage for printer<commit_after>package printer_test

import (
	"testing"

	"github.com/goccy/go-yaml/lexer"
	"github.com/goccy/go-yaml/printer"
)

func Test_Printer(t *testing.T) {
	yml := `---
text: aaaa
text2: aaaa
 bbbb
 cccc
 dddd
 eeee
text3: ffff
 gggg
 hhhh
 iiii
 jjjj
bool: true
number: 10
anchor: &x 1
alias: *x
`
	t.Run("print starting from tokens[3]", func(t *testing.T) {
		tokens := lexer.Tokenize(yml)
		var p printer.Printer
		actual := "\n" + p.PrintErrorToken(tokens[3], false)
		expect := `
   1 | ---
>  2 | text: aaaa
            ^
   3 | text2: aaaa
   4 |  bbbb
   5 |  cccc
   6 |  dddd
   7 |  eeee
   8 | `
		if actual != expect {
			t.Fatalf("unexpected output: expect:[%s]\n actual:[%s]", expect, actual)
		}
	})
	t.Run("print stargin from tokens[4]", func(t *testing.T) {
		tokens := lexer.Tokenize(yml)
		var p printer.Printer
		actual := "\n" + p.PrintErrorToken(tokens[4], false)
		expect := `
   1 | ---
   2 | text: aaaa
>  3 | text2: aaaa
   4 |  bbbb
   5 |  cccc
   6 |  dddd
   7 |  eeee
      ^
`
		if actual != expect {
			t.Fatalf("unexpected output: expect:[%s]\n actual:[%s]", expect, actual)
		}
	})
	t.Run("print starting from tokens[6]", func(t *testing.T) {
		tokens := lexer.Tokenize(yml)
		var p printer.Printer
		actual := "\n" + p.PrintErrorToken(tokens[6], false)
		expect := `
   1 | ---
   2 | text: aaaa
>  3 | text2: aaaa
   4 |  bbbb
   5 |  cccc
   6 |  dddd
   7 |  eeee
             ^
   8 | text3: ffff
   9 |  gggg
  10 |  hhhh
  11 |  iiii
  12 |  jjjj
  13 | `
		if actual != expect {
			t.Fatalf("unexpected output: expect:[%s]\n actual:[%s]", expect, actual)
		}
	})
	t.Run("output with color", func(t *testing.T) {
		t.Run("token6", func(t *testing.T) {
			tokens := lexer.Tokenize(yml)
			var p printer.Printer
			t.Logf("%s", p.PrintErrorToken(tokens[6], true))
		})
		t.Run("token9", func(t *testing.T) {
			tokens := lexer.Tokenize(yml)
			var p printer.Printer
			t.Logf("%s", p.PrintErrorToken(tokens[9], true))
		})
		t.Run("token12", func(t *testing.T) {
			tokens := lexer.Tokenize(yml)
			var p printer.Printer
			t.Logf("%s", p.PrintErrorToken(tokens[12], true))
		})
	})
	t.Run("print error message", func(t *testing.T) {
		var p printer.Printer
		src := "message"
		msg := p.PrintErrorMessage(src, false)
		if msg != src {
			t.Fatal("unexpected result")
		}
		p.PrintErrorMessage(src, true)
	})
}
<|endoftext|>