<commit_before>// go-rst - A reStructuredText parser for Go
// 2014 (c) The go-rst Authors
// MIT Licensed. See LICENSE for details.

package parse

import (
	"code.google.com/p/go.text/unicode/norm"
	"fmt"
	"github.com/demizer/go-elog"
	"github.com/demizer/go-spew/spew"
	"reflect"
)

var spd = spew.ConfigState{Indent: "\t", DisableMethods: true}

type systemMessageLevel int

const (
	levelInfo systemMessageLevel = iota
	levelWarning
	levelError
	levelSevere
)

var systemMessageLevels = [...]string{
	"INFO",
	"WARNING",
	"ERROR",
	"SEVERE",
}

func (s systemMessageLevel) String() string {
	return systemMessageLevels[s]
}

type parserMessage int

const (
	warningShortUnderline parserMessage = iota
	errorUnexpectedSectionTitle
	errorUnexpectedSectionTitleOrTransition
)

var parserErrors = [...]string{
	"warningShortUnderline",
	"errorUnexpectedSectionTitle",
	"errorUnexpectedSectionTitleOrTransition",
}

func (p parserMessage) String() string {
	return parserErrors[p]
}

func (p parserMessage) Message() (s string) {
	switch p {
	case warningShortUnderline:
		s = "Title underline too short."
	case errorUnexpectedSectionTitle:
		s = "Unexpected section title."
	case errorUnexpectedSectionTitleOrTransition:
		s = "Unexpected section title or transition."
	}
	return
}

func (p parserMessage) Level() (s systemMessageLevel) {
	switch p {
	case warningShortUnderline:
		s = levelWarning
	case errorUnexpectedSectionTitle:
		s = levelSevere
	case errorUnexpectedSectionTitleOrTransition:
		s = levelSevere
	}
	return
}

type sectionLevels []*SectionNode

func (s *sectionLevels) String() string {
	var out string
	for _, sec := range *s {
		out += fmt.Sprintf("level: %d, rune: %q, overline: %t, length: %d\n",
			sec.Level, sec.UnderLine.Rune, sec.OverLine != nil, sec.Length)
	}
	return out
}

// Returns nil if not found
func (s *sectionLevels) FindByRune(adornChar rune) *SectionNode {
	for _, sec := range *s {
		if sec.UnderLine.Rune == adornChar {
			return sec
		}
	}
	return nil
}

// If exists == true, a section node with the same text and underline has been found in
// sectionLevels, sec is the matching SectionNode. If exists == false, then the sec return value is
// the similarly leveled SectionNode. If exists == false and sec == nil, then the SectionNode added
// to sectionLevels is a new Node.
func (s *sectionLevels) Add(section *SectionNode) (exists bool, sec *SectionNode) {
	sec = s.FindByRune(section.UnderLine.Rune)
	if sec != nil {
		if sec.Text == section.Text {
			return true, sec
		} else if sec.Text != section.Text {
			section.Level = sec.Level
		}
	} else {
		section.Level = len(*s) + 1
	}
	exists = false
	*s = append(*s, section)
	return
}

func (s *sectionLevels) Level() int {
	return len(*s)
}

// Parse is the entry point for the reStructuredText parser.
func Parse(name, text string) (t *Tree, errors []error) {
	t = New(name)
	if !norm.NFC.IsNormalString(text) {
		text = norm.NFC.String(text)
	}
	t.text = text
	_, errors = t.Parse(text, t)
	return
}

func New(name string) *Tree {
	return &Tree{
		Name:          name,
		Nodes:         newList(),
		nodeTarget:    newList(),
		sectionLevels: new(sectionLevels),
		indentWidth:   indentWidth,
	}
}

const (
	tokenZero   = 3
	indentWidth = 4 // Default indent width
)

type Tree struct {
	Name             string
	Nodes            *NodeList // The root node list
	nodeTarget       *NodeList // Used by the parser to add nodes to a target NodeList
	Errors           []error
	text             string
	lex              *lexer
	tokenBackupCount int
	tokenPeekCount   int
	token            [7]*item
	sectionLevels    *sectionLevels // Encountered section levels
	id               int            // The unique id of the node in the tree
	indentWidth      int
	indentLevel      int
}

func (t *Tree) errorf(format string, args ...interface{}) {
	format = fmt.Sprintf("go-rst: %s:%d: %s\n", t.Name, t.lex.lineNumber(), format)
	t.Errors = append(t.Errors, fmt.Errorf(format, args...))
}

func (t *Tree) error(err error) {
	t.errorf("%s\n", err)
}

// startParse initializes the parser, using the lexer.
func (t *Tree) startParse(lex *lexer) {
	t.lex = lex
}

// stopParse terminates parsing.
func (t *Tree) stopParse() {
	t.Nodes = nil
	t.nodeTarget = nil
	t.lex = nil
}

func (t *Tree) Parse(text string, treeSet *Tree) (tree *Tree, errors []error) {
	log.Debugln("Start")
	t.startParse(lex(t.Name, text))
	t.text = text
	t.parse(treeSet)
	log.Debugln("End")
	return t, t.Errors
}

func (t *Tree) parse(tree *Tree) {
	log.Debugln("Start")

	t.nodeTarget = t.Nodes

	for t.peek(1).Type != itemEOF {
		var n Node

		token := t.next()
		log.Infof("\nParser got token: %#+v\n\n", token)

		switch token.Type {
		case itemSectionAdornment:
			n = t.section(token)
		case itemParagraph:
			n = newParagraph(token, &t.id)
		case itemSpace:
			n = t.indent(token)
			if n == nil {
				continue
			}
		case itemTitle, itemBlankLine:
			// itemTitle is consumed when evaluating itemSectionAdornment
			continue
		case itemEOF:
			goto exit
		default:
			t.errorf("%q Not implemented!", token.Type)
			continue
		}

		t.nodeTarget.append(n)
		switch n.NodeType() {
		case NodeSection, NodeBlockQuote:
			// Set the loop to append items to the NodeList of the new section
			t.nodeTarget = reflect.ValueOf(n).Elem().FieldByName("NodeList").Addr().Interface().(*NodeList)
		}
	}

	exit:
	log.Debugln("End")
}

func (t *Tree) backup() *item {
	t.tokenBackupCount++
	// log.Debugln("t.tokenBackupCount:", t.tokenPeekCount)
	for i := len(t.token) - 1; i > 0; i-- {
		t.token[i] = t.token[i-1]
		t.token[i-1] = nil
	}
	// log.Debugf("\n##### backup() aftermath #####\n\n")
	// spd.Dump(t.token)
	return t.token[tokenZero-t.tokenBackupCount]
}

func (t *Tree) peekBack(pos int) *item {
	return t.token[tokenZero-pos]
}

func (t *Tree) peek(pos int) *item {
	// log.Debugln("t.tokenPeekCount:", t.tokenPeekCount, "Pos:", pos)
	if pos < 1 {
		panic("pos cannot be < 1")
	}
	var nItem *item
	for i := 0; i < pos; i++ {
		// log.Debugln("i:", i, "peekCount:", t.tokenPeekCount, "pos:", pos)
		if t.tokenPeekCount > i {
			nItem = t.token[tokenZero+i]
			log.Debugf("Using %#+v\n", nItem)
			continue
		}
		log.Debugln(tokenZero + t.tokenPeekCount + i)
		if t.token[tokenZero + t.tokenPeekCount + i + 1] == nil {
			t.tokenPeekCount++
			// log.Debugln("Getting next item")
			t.token[tokenZero+t.tokenPeekCount+i] = t.lex.nextItem()
			nItem = t.token[tokenZero+t.tokenPeekCount+i]
		} else {
			nItem = t.token[tokenZero+t.tokenPeekCount+i]
		}
	}
	// log.Debugf("\n##### peek() aftermath #####\n\n")
	// spd.Dump(t.token)
	// log.Debugf("Returning: %#+v\n", nItem)
	return nItem
}

func (t *Tree) next() *item {
	// log.Debugln("t.tokenPeekCount:", t.tokenPeekCount)
	// skip shifts the pointers left in t.token, pos is the amount to shift
	skip := func(num int) {
		for i := num; i > 0; i-- {
			for x := 0; x < len(t.token)-1; x++ {
				t.token[x] = t.token[x+1]
				t.token[x+1] = nil
			}
		}
	}
	if t.tokenPeekCount > 0 {
		skip(t.tokenPeekCount)
	} else {
		skip(1)
		t.token[tokenZero] = t.lex.nextItem()
	}
	t.tokenBackupCount, t.tokenPeekCount = 0, 0
	// log.Debugf("\n##### next() aftermath #####\n\n")
	// spd.Dump(t.token)
	return t.token[tokenZero]
}

func (t *Tree) section(i *item) Node {
	log.Debugln("Start")
	var overAdorn, title, underAdorn *item
	var overline bool
	var sysMessage Node

	peekBack := t.peekBack(1)
	if peekBack != nil {
		if peekBack.Type == itemSpace {
			// Looking back past the white space
			if t.peekBack(2).Type == itemTitle {
				return t.systemMessage(errorUnexpectedSectionTitle)
			}
			return t.systemMessage(errorUnexpectedSectionTitleOrTransition)
		} else if peekBack.Type == itemTitle {
			if t.peekBack(2) != nil && t.peekBack(2).Type == itemSectionAdornment {
				// The overline of the section
				overline = true
				overAdorn = peekBack
			}
		}
	}

	title = t.peekBack(1)
	underAdorn = i

	// TODO: Change these into proper error messages!
	// Check adornment for proper syntax
	if underAdorn.Type == itemSpace {
		t.backup() // Put the parser back on the title
		return t.systemMessage(errorUnexpectedSectionTitle)
	} else if overline && title.Length != overAdorn.Length {
		t.errorf("Section over line not equal to title length!")
	} else if overline && overAdorn.Text != underAdorn.Text {
		t.errorf("Section title over line does not match section title under line.")
	}

	sec := newSection(title, overAdorn, underAdorn, &t.id)
	exists, eSec := t.sectionLevels.Add(sec)
	if exists && eSec != nil {
		t.errorf("SectionNode using Text \"%s\" and Rune '%s' was previously parsed!",
			sec.Text, string(sec.UnderLine.Rune))
	} else if !exists && eSec != nil {
		// There is a matching level in sectionLevels
		t.nodeTarget = &(*t.sectionLevels)[sec.Level-2].NodeList
	}

	// System messages have to be applied after the section is created in order to preserve
	// a consecutive id number.
	if title.Length != underAdorn.Length {
		sysMessage = t.systemMessage(warningShortUnderline)
		sec.NodeList = append(sec.NodeList, sysMessage)
	}

	log.Debugln("End")
	return sec
}

func (t *Tree) systemMessage(err parserMessage) Node {
	var lbText string
	var lbTextLen int
	var backToken int

	s := newSystemMessage(&item{
		Type: itemSystemMessage,
		Line: t.token[tokenZero].Line,
	},
		err.Level(), &t.id)

	msg := newParagraph(&item{
		Text:   err.Message(),
		Length: len(err.Message()),
	}, &t.id)

	switch err {
	case warningShortUnderline, errorUnexpectedSectionTitle:
		log.Debugln("FOUND", err)
		backToken = tokenZero - 1
		if t.peekBack(1).Type == itemSpace {
			backToken = tokenZero - 2
		}
		lbText = t.token[backToken].Text.(string) + "\n" + t.token[tokenZero].Text.(string)
		lbTextLen = len(lbText) + 1
	case errorUnexpectedSectionTitleOrTransition:
		log.Debugln("FOUND errorUnexpectedSectionTitleOrTransition")
		lbText = t.token[tokenZero].Text.(string)
		lbTextLen = len(lbText)
	}

	lb := newLiteralBlock(&item{
		Type:   itemLiteralBlock,
		Text:   lbText,
		Length: lbTextLen, // Add one to account for the backslash
	}, &t.id)

	s.NodeList = append(s.NodeList, msg, lb)
	return s
}

func (t *Tree) indent(i *item) Node {
	level := i.Length / t.indentWidth
	if t.peekBack(1).Type == itemBlankLine {
		if t.indentLevel == level {
			// Append to the current blockquote NodeList
			return nil
		}
		t.indentLevel = level
		return newBlockQuote(&item{Type: itemBlockquote, Line: i.Line}, level, &t.id)
	}
	return nil
}
<commit_msg>parse.go: Add peekSkip()<commit_after>// go-rst - A reStructuredText parser for Go
// 2014 (c) The go-rst Authors
// MIT Licensed. See LICENSE for details.

package parse

import (
	"code.google.com/p/go.text/unicode/norm"
	"fmt"
	"github.com/demizer/go-elog"
	"github.com/demizer/go-spew/spew"
	"reflect"
)

var spd = spew.ConfigState{Indent: "\t", DisableMethods: true}

type systemMessageLevel int

const (
	levelInfo systemMessageLevel = iota
	levelWarning
	levelError
	levelSevere
)

var systemMessageLevels = [...]string{
	"INFO",
	"WARNING",
	"ERROR",
	"SEVERE",
}

func (s systemMessageLevel) String() string {
	return systemMessageLevels[s]
}

type parserMessage int

const (
	warningShortUnderline parserMessage = iota
	errorUnexpectedSectionTitle
	errorUnexpectedSectionTitleOrTransition
)

var parserErrors = [...]string{
	"warningShortUnderline",
	"errorUnexpectedSectionTitle",
	"errorUnexpectedSectionTitleOrTransition",
}

func (p parserMessage) String() string {
	return parserErrors[p]
}

func (p parserMessage) Message() (s string) {
	switch p {
	case warningShortUnderline:
		s = "Title underline too short."
	case errorUnexpectedSectionTitle:
		s = "Unexpected section title."
	case errorUnexpectedSectionTitleOrTransition:
		s = "Unexpected section title or transition."
	}
	return
}

func (p parserMessage) Level() (s systemMessageLevel) {
	switch p {
	case warningShortUnderline:
		s = levelWarning
	case errorUnexpectedSectionTitle:
		s = levelSevere
	case errorUnexpectedSectionTitleOrTransition:
		s = levelSevere
	}
	return
}

type sectionLevels []*SectionNode

func (s *sectionLevels) String() string {
	var out string
	for _, sec := range *s {
		out += fmt.Sprintf("level: %d, rune: %q, overline: %t, length: %d\n",
			sec.Level, sec.UnderLine.Rune, sec.OverLine != nil, sec.Length)
	}
	return out
}

// Returns nil if not found
func (s *sectionLevels) FindByRune(adornChar rune) *SectionNode {
	for _, sec := range *s {
		if sec.UnderLine.Rune == adornChar {
			return sec
		}
	}
	return nil
}

// If exists == true, a section node with the same text and underline has been found in
// sectionLevels, sec is the matching SectionNode. If exists == false, then the sec return value is
// the similarly leveled SectionNode. If exists == false and sec == nil, then the SectionNode added
// to sectionLevels is a new Node.
func (s *sectionLevels) Add(section *SectionNode) (exists bool, sec *SectionNode) {
	sec = s.FindByRune(section.UnderLine.Rune)
	if sec != nil {
		if sec.Text == section.Text {
			return true, sec
		} else if sec.Text != section.Text {
			section.Level = sec.Level
		}
	} else {
		section.Level = len(*s) + 1
	}
	exists = false
	*s = append(*s, section)
	return
}

func (s *sectionLevels) Level() int {
	return len(*s)
}

// Parse is the entry point for the reStructuredText parser.
func Parse(name, text string) (t *Tree, errors []error) {
	t = New(name)
	if !norm.NFC.IsNormalString(text) {
		text = norm.NFC.String(text)
	}
	t.text = text
	_, errors = t.Parse(text, t)
	return
}

func New(name string) *Tree {
	return &Tree{
		Name:          name,
		Nodes:         newList(),
		nodeTarget:    newList(),
		sectionLevels: new(sectionLevels),
		indentWidth:   indentWidth,
	}
}

const (
	tokenZero   = 3
	indentWidth = 4 // Default indent width
)

type Tree struct {
	Name             string
	Nodes            *NodeList // The root node list
	nodeTarget       *NodeList // Used by the parser to add nodes to a target NodeList
	Errors           []error
	text             string
	lex              *lexer
	tokenBackupCount int
	tokenPeekCount   int
	token            [7]*item
	sectionLevels    *sectionLevels // Encountered section levels
	id               int            // The unique id of the node in the tree
	indentWidth      int
	indentLevel      int
}

func (t *Tree) errorf(format string, args ...interface{}) {
	format = fmt.Sprintf("go-rst: %s:%d: %s\n", t.Name, t.lex.lineNumber(), format)
	t.Errors = append(t.Errors, fmt.Errorf(format, args...))
}

func (t *Tree) error(err error) {
	t.errorf("%s\n", err)
}

// startParse initializes the parser, using the lexer.
func (t *Tree) startParse(lex *lexer) {
	t.lex = lex
}

// stopParse terminates parsing.
func (t *Tree) stopParse() {
	t.Nodes = nil
	t.nodeTarget = nil
	t.lex = nil
}

func (t *Tree) Parse(text string, treeSet *Tree) (tree *Tree, errors []error) {
	log.Debugln("Start")
	t.startParse(lex(t.Name, text))
	t.text = text
	t.parse(treeSet)
	log.Debugln("End")
	return t, t.Errors
}

func (t *Tree) parse(tree *Tree) {
	log.Debugln("Start")

	t.nodeTarget = t.Nodes

	for t.peek(1).Type != itemEOF {
		var n Node

		token := t.next()
		log.Infof("\nParser got token: %#+v\n\n", token)

		switch token.Type {
		case itemSectionAdornment:
			n = t.section(token)
		case itemParagraph:
			n = newParagraph(token, &t.id)
		case itemSpace:
			n = t.indent(token)
			if n == nil {
				continue
			}
		case itemTitle, itemBlankLine:
			// itemTitle is consumed when evaluating itemSectionAdornment
			continue
		case itemEOF:
			goto exit
		default:
			t.errorf("%q Not implemented!", token.Type)
			continue
		}

		t.nodeTarget.append(n)
		switch n.NodeType() {
		case NodeSection, NodeBlockQuote:
			// Set the loop to append items to the NodeList of the new section
			t.nodeTarget = reflect.ValueOf(n).Elem().FieldByName("NodeList").Addr().Interface().(*NodeList)
		}
	}

	exit:
	log.Debugln("End")
}

func (t *Tree) backup() *item {
	t.tokenBackupCount++
	// log.Debugln("t.tokenBackupCount:", t.tokenPeekCount)
	for i := len(t.token) - 1; i > 0; i-- {
		t.token[i] = t.token[i-1]
		t.token[i-1] = nil
	}
	// log.Debugf("\n##### backup() aftermath #####\n\n")
	// spd.Dump(t.token)
	return t.token[tokenZero-t.tokenBackupCount]
}

func (t *Tree) peekBack(pos int) *item {
	return t.token[tokenZero-pos]
}

func (t *Tree) peek(pos int) *item {
	// log.Debugln("t.tokenPeekCount:", t.tokenPeekCount, "Pos:", pos)
	if pos < 1 {
		panic("pos cannot be < 1")
	}
	var nItem *item
	for i := 0; i < pos; i++ {
		// log.Debugln("i:", i, "peekCount:", t.tokenPeekCount, "pos:", pos)
		if t.tokenPeekCount > i {
			nItem = t.token[tokenZero+i]
			log.Debugf("Using %#+v\n", nItem)
			continue
		}
		log.Debugln(tokenZero + t.tokenPeekCount + i)
		if t.token[tokenZero + t.tokenPeekCount + i + 1] == nil {
			t.tokenPeekCount++
			// log.Debugln("Getting next item")
			t.token[tokenZero+t.tokenPeekCount+i] = t.lex.nextItem()
			nItem = t.token[tokenZero+t.tokenPeekCount+i]
		} else {
			nItem = t.token[tokenZero+t.tokenPeekCount+i]
		}
	}
	// log.Debugf("\n##### peek() aftermath #####\n\n")
	// spd.Dump(t.token)
	// log.Debugf("Returning: %#+v\n", nItem)
	return nItem
}

func (t *Tree) peekSkip(pos int, iSkip itemElement) *item {
	var nItem *item
	outer:
	for i := 1; i <= pos; i++ {
		for {
			nItem = t.peek(i)
			if nItem.Type == iSkip {
				continue
			} else {
				break outer
			}
		}
	}
	return nItem
}

func (t *Tree) next() *item {
	// log.Debugln("t.tokenPeekCount:", t.tokenPeekCount)
	// skip shifts the pointers left in t.token, pos is the amount to shift
	skip := func(num int) {
		for i := num; i > 0; i-- {
			for x := 0; x < len(t.token)-1; x++ {
				t.token[x] = t.token[x+1]
				t.token[x+1] = nil
			}
		}
	}
	if t.tokenPeekCount > 0 {
		skip(t.tokenPeekCount)
	} else {
		skip(1)
		t.token[tokenZero] = t.lex.nextItem()
	}
	t.tokenBackupCount, t.tokenPeekCount = 0, 0
	// log.Debugf("\n##### next() aftermath #####\n\n")
	// spd.Dump(t.token)
	return t.token[tokenZero]
}

func (t *Tree) section(i *item) Node {
	log.Debugln("Start")
	var overAdorn, title, underAdorn *item
	var overline bool
	var sysMessage Node

	peekBack := t.peekBack(1)
	if peekBack != nil {
		if peekBack.Type == itemSpace {
			// Looking back past the white space
			if t.peekBack(2).Type == itemTitle {
				return t.systemMessage(errorUnexpectedSectionTitle)
			}
			return t.systemMessage(errorUnexpectedSectionTitleOrTransition)
		} else if peekBack.Type == itemTitle {
			if t.peekBack(2) != nil && t.peekBack(2).Type == itemSectionAdornment {
				// The overline of the section
				overline = true
				overAdorn = peekBack
			}
		}
	}

	title = t.peekBack(1)
	underAdorn = i

	// TODO: Change these into proper error messages!
	// Check adornment for proper syntax
	if underAdorn.Type == itemSpace {
		t.backup() // Put the parser back on the title
		return t.systemMessage(errorUnexpectedSectionTitle)
	} else if overline && title.Length != overAdorn.Length {
		t.errorf("Section over line not equal to title length!")
	} else if overline && overAdorn.Text != underAdorn.Text {
		t.errorf("Section title over line does not match section title under line.")
	}

	sec := newSection(title, overAdorn, underAdorn, &t.id)
	exists, eSec := t.sectionLevels.Add(sec)
	if exists && eSec != nil {
		t.errorf("SectionNode using Text \"%s\" and Rune '%s' was previously parsed!",
			sec.Text, string(sec.UnderLine.Rune))
	} else if !exists && eSec != nil {
		// There is a matching level in sectionLevels
		t.nodeTarget = &(*t.sectionLevels)[sec.Level-2].NodeList
	}

	// System messages have to be applied after the section is created in order to preserve
	// a consecutive id number.
	if title.Length != underAdorn.Length {
		sysMessage = t.systemMessage(warningShortUnderline)
		sec.NodeList = append(sec.NodeList, sysMessage)
	}

	log.Debugln("End")
	return sec
}

func (t *Tree) systemMessage(err parserMessage) Node {
	var lbText string
	var lbTextLen int
	var backToken int

	s := newSystemMessage(&item{
		Type: itemSystemMessage,
		Line: t.token[tokenZero].Line,
	},
		err.Level(), &t.id)

	msg := newParagraph(&item{
		Text:   err.Message(),
		Length: len(err.Message()),
	}, &t.id)

	switch err {
	case warningShortUnderline, errorUnexpectedSectionTitle:
		log.Debugln("FOUND", err)
		backToken = tokenZero - 1
		if t.peekBack(1).Type == itemSpace {
			backToken = tokenZero - 2
		}
		lbText = t.token[backToken].Text.(string) + "\n" + t.token[tokenZero].Text.(string)
		lbTextLen = len(lbText) + 1
	case errorUnexpectedSectionTitleOrTransition:
		log.Debugln("FOUND errorUnexpectedSectionTitleOrTransition")
		lbText = t.token[tokenZero].Text.(string)
		lbTextLen = len(lbText)
	}

	lb := newLiteralBlock(&item{
		Type:   itemLiteralBlock,
		Text:   lbText,
		Length: lbTextLen, // Add one to account for the backslash
	}, &t.id)

	s.NodeList = append(s.NodeList, msg, lb)
	return s
}

func (t *Tree) indent(i *item) Node {
	level := i.Length / t.indentWidth
	if t.peekBack(1).Type == itemBlankLine {
		if t.indentLevel == level {
			// Append to the current blockquote NodeList
			return nil
		}
		t.indentLevel = level
		return newBlockQuote(&item{Type: itemBlockquote, Line: i.Line}, level, &t.id)
	}
	return nil
}
<|endoftext|>


<commit_before>package parser

import (
	"fmt"

	"github.com/twtiger/gosecco/tree"
)

// ParseError represents error parsing a policy file. It will report the filename and the line number as well as the actual error.
type ParseError struct {
	originalError error
	file          string
	line          int
}

func (e *ParseError) Error() string {
	return fmt.Sprintf("%s:%d: %s", e.file, e.line, e.originalError)
}

func parseLines(path string, lines []string) (tree.RawPolicy, error) {
	result := []interface{}{}

	for ix, l := range lines {
		switch lineType(l) {
		case commentLine: //ignore
		case emptyLine: //ignore
		case ruleLine:
			parsedRule, err := parseRule(l)
			if err != nil {
				return tree.RawPolicy{}, &ParseError{err, path, ix}
			}
			result = append(result, parsedRule)
		case assignmentLine, defaultAssignmentLine:
			parsedBinding, err := parseBinding(l)
			if err != nil {
				return tree.RawPolicy{}, &ParseError{err, path, ix}
			}
			result = append(result, parsedBinding)

		case unknownLine:
			return tree.RawPolicy{}, &ParseError{fmt.Errorf("Couldn't parse line: '%s' - it doesn't match any kind of valid syntax", l), path, ix}
		}
	}

	return tree.RawPolicy{RuleOrMacros: result}, nil
}

// ParseFile will parse the given file and return a raw parse tree or the error generated
// This function is deprecated and shouldn't be used in new code
func ParseFile(path string) (tree.RawPolicy, error) {
	return Parse(&FileSource{path})
}

// ParseString will parse the given string and return a raw parse tree or the error generated
// This function is deprecated and shouldn't be used in new code
func ParseString(str string) (tree.RawPolicy, error) {
	return Parse(&StringSource{"<string>", str})
}

// Parse will parse the given Source and return a raw parse tree or the error generated
func Parse(s Source) (tree.RawPolicy, error) {
	return s.parse()
}
<commit_msg>Expose the Parse function so it's possible to implement the Source interface<commit_after>package parser

import (
	"fmt"

	"github.com/twtiger/gosecco/tree"
)

// ParseError represents error parsing a policy file. It will report the filename and the line number as well as the actual error.
type ParseError struct {
	originalError error
	file          string
	line          int
}

func (e *ParseError) Error() string {
	return fmt.Sprintf("%s:%d: %s", e.file, e.line, e.originalError)
}

func parseLines(path string, lines []string) (tree.RawPolicy, error) {
	result := []interface{}{}

	for ix, l := range lines {
		switch lineType(l) {
		case commentLine: //ignore
		case emptyLine: //ignore
		case ruleLine:
			parsedRule, err := parseRule(l)
			if err != nil {
				return tree.RawPolicy{}, &ParseError{err, path, ix}
			}
			result = append(result, parsedRule)
		case assignmentLine, defaultAssignmentLine:
			parsedBinding, err := parseBinding(l)
			if err != nil {
				return tree.RawPolicy{}, &ParseError{err, path, ix}
			}
			result = append(result, parsedBinding)

		case unknownLine:
			return tree.RawPolicy{}, &ParseError{fmt.Errorf("Couldn't parse line: '%s' - it doesn't match any kind of valid syntax", l), path, ix}
		}
	}

	return tree.RawPolicy{RuleOrMacros: result}, nil
}

// ParseFile will parse the given file and return a raw parse tree or the error generated
// This function is deprecated and shouldn't be used in new code
func ParseFile(path string) (tree.RawPolicy, error) {
	return Parse(&FileSource{path})
}

// ParseString will parse the given string and return a raw parse tree or the error generated
// This function is deprecated and shouldn't be used in new code
func ParseString(str string) (tree.RawPolicy, error) {
	return Parse(&StringSource{"<string>", str})
}

// Parse will parse the given Source and return a raw parse tree or the error generated
func Parse(s Source) (tree.RawPolicy, error) {
	return s.Parse()
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Dongsu Park <dpark@posteo.net>
//
// Permission to use, copy, modify, and distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
// ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
// ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
// OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

package parser

import (
	"bytes"
	"fmt"
	"io"
	"regexp"
	"strings"
	"time"

	"golang.org/x/net/html"

	"github.com/dongsupark/go-subtitle/pkg"
	"github.com/dongsupark/go-subtitle/subtitle"
)

const (
	SamiStateInit      = 0
	SamiStateSyncStart = 1
	SamiStateText      = 2
	SamiStateSyncEnd   = 3
	SamiStateForceQuit = 99
)

type SamiStateType int

type SamiFormat struct {
	TypeName string
}

func (sr *SamiFormat) Read(inputData string) (subtitle.Subtitle, error) {
	var st subtitle.Subtitle
	se := new(subtitle.SubtitleEntry)
	ss := SamiStateType(SamiStateInit)

	inputData = strings.TrimSpace(inputData)

	doc, err := html.Parse(strings.NewReader(inputData))
	if err != nil {
		return subtitle.Subtitle{}, err
	}

	renl := regexp.MustCompile("\\n")

	var parseNode func(*html.Node, SamiStateType)
	parseNode = func(n *html.Node, samiState SamiStateType) {
		if n.Type == html.ElementNode && n.Data == "sync" {
			for _, a := range n.Attr {
				if a.Key == "start" {
					if samiState == SamiStateSyncStart || samiState == SamiStateText {
						se.EndValue = pkg.ComposeTimeDuration(0, 0, 0, pkg.StringToInt(a.Val))
						samiState = SamiStateSyncEnd

						st.Subtitles = append(st.Subtitles, *se)
						se = new(subtitle.SubtitleEntry)
					} else {
						se.StartValue = pkg.ComposeTimeDuration(0, 0, 0, pkg.StringToInt(a.Val))
						samiState = SamiStateSyncStart
					}
					break
				}
			}
		}
		if n.Type == html.TextNode {
			n.Data = stripComments(n.Data)

			inText := strings.TrimSpace(renl.ReplaceAllString(n.Data, " "))
			if len(inText) == 0 {
				if samiState == SamiStateSyncEnd {
					samiState = SamiStateInit
				}
			} else {
				se.Text = n.Data
				samiState = SamiStateText
			}
		}
		if n.Type == html.CommentNode {
			return
		}
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			parseNode(c, samiState)
		}
	}
	parseNode(doc, ss)

	return st, nil
}

func (sr *SamiFormat) Write(insub subtitle.Subtitle) (string, error) {
	doc := &html.Node{
		Type: html.DocumentNode,
	}
	for _, v := range insub.Subtitles {
		htmlText := strings.Replace(v.Text, "\n", "<br>", -1)

		sStartNode := &html.Node{
			Type: html.ElementNode,
			Data: fmt.Sprintf("SYNC Start=%s", timeToSami(v.StartValue)),
		}
		sPNode := &html.Node{
			Type: html.ElementNode,
			Data: "P Class=ENCC",
		}
		sPNode.AppendChild(&html.Node{
			Type: html.TextNode,
			Data: fmt.Sprintf("%s\n", htmlText),
		})
		sStartNode.AppendChild(sPNode)
		doc.AppendChild(sStartNode)

		sEndNode := &html.Node{
			Type: html.ElementNode,
			Data: fmt.Sprintf("SYNC Start=%s", timeToSami(v.EndValue)),
		}
		sPNode = &html.Node{
			Type: html.ElementNode,
			Data: "P Class=ENCC",
		}
		sPNode.AppendChild(&html.Node{
			Type: html.TextNode,
			Data: " \n",
		})
		sEndNode.AppendChild(sPNode)
		doc.AppendChild(sEndNode)
	}

	b := new(bytes.Buffer)
	if err := html.Render(b, doc); err != nil {
		return "", err
	}

	return strings.TrimSpace(b.String()), nil
}

// strip comments in every text node
func stripComments(inStr string) string {
	z := html.NewTokenizer(bytes.NewBufferString(inStr))
	for {
		tt := z.Next()
		if tt == html.ErrorToken {
			if err := z.Err(); err != nil && err != io.EOF {
				return inStr
			}
			break
		}
		if tt == html.CommentToken {
			return strings.Replace(inStr, string(z.Raw()), "", -1)
		}
	}

	return inStr
}

func timeToSami(inTime time.Duration) string {
	totalSec := inTime.Seconds()
	return fmt.Sprintf("%04d%03d", int(totalSec), int(inTime.Nanoseconds()/1000/1000%1000))
}
<commit_msg>parser: reimplement sami parser using html tokenizer<commit_after>// Copyright (c) 2017 Dongsu Park <dpark@posteo.net>
//
// Permission to use, copy, modify, and distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
// ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
// ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
// OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

package parser

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"regexp"
	"strings"
	"time"

	"golang.org/x/net/html"

	"github.com/dongsupark/go-subtitle/pkg"
	"github.com/dongsupark/go-subtitle/subtitle"
)

const (
	SamiStateInit      = 0
	SamiStateSyncStart = 1
	SamiStateText      = 2
	SamiStateSyncEnd   = 3
	SamiStateForceQuit = 99
)

var textElemTags = []string{
	"i",
	"b",
}

func matchTextElemTag(input string) bool {
	for _, e := range textElemTags {
		if input == e {
			return true
		}
	}
	return false
}

var legitElemTags = []string{
	"&nbsp",
	"i>",
	"b>",
}

func hasLegitElemTag(input string) bool {
	for _, e := range legitElemTags {
		if strings.Contains(input, e) {
			return true
		}
	}
	return false
}

type SamiStateType int

type SamiFormat struct {
	TypeName string
}

func (sr *SamiFormat) Read(inputData string) (subtitle.Subtitle, error) {
	var st subtitle.Subtitle
	se := new(subtitle.SubtitleEntry)
	samiState := SamiStateType(SamiStateInit)

	inputData = strings.TrimSpace(inputData)

	renl := regexp.MustCompile("\\n")

	z := html.NewTokenizer(strings.NewReader(inputData))
	prevStartValue := time.Duration(0)
	for {
		tok := z.Next()
		switch tok {
		case html.ErrorToken:
			if z.Err() == io.EOF {
				break
			}
			return subtitle.Subtitle{}, fmt.Errorf("got error token")
		case html.StartTagToken:
			tn, hasAttr := z.TagName()
			tnStr := string(tn)

			if hasAttr && strings.ToLower(tnStr) == "sync" {
				key, value, _ := z.TagAttr()
				if strings.ToLower(string(key)) == "start" {
					if samiState == SamiStateSyncEnd {
						se.EndValue = prevStartValue
						samiState = SamiStateInit

						st.Subtitles = append(st.Subtitles, *se)
						se = new(subtitle.SubtitleEntry)
					}

					if samiState == SamiStateInit {
						se.StartValue = pkg.ComposeTimeDuration(0, 0, 0, pkg.StringToInt(string(value)))
						prevStartValue = se.StartValue
						samiState = SamiStateSyncStart
					} else if samiState == SamiStateSyncStart || samiState == SamiStateText {
						se.EndValue = pkg.ComposeTimeDuration(0, 0, 0, pkg.StringToInt(string(value)))
						samiState = SamiStateInit

						st.Subtitles = append(st.Subtitles, *se)
						se = new(subtitle.SubtitleEntry)
					}
					break
				}
			}

			// consider this node as a text node with an in-text tag
			if matchTextElemTag(tnStr) {
				se.Text += fmt.Sprintf("<%s>", tnStr)
			}
		case html.TextToken:
			toSyncEnd := false
			parsed := ""

			if strings.Contains(string(z.Raw()), "&nbsp") {
				parsed = string(z.Raw())
				toSyncEnd = true
			} else {
				parsed = string(z.Text())
			}

			if samiState == SamiStateSyncStart || samiState == SamiStateInit {
				textStr := stripComments(parsed)

				inText := strings.TrimSpace(renl.ReplaceAllString(textStr, " "))
				if len(inText) > 0 {
					se.Text += parsed

					if toSyncEnd {
						samiState = SamiStateSyncEnd
					} else {
						samiState = SamiStateText
					}
				}
			}
		case html.EndTagToken:
			tn, _ := z.TagName()
			tnStr := string(tn)

			if matchTextElemTag(tnStr) {
				se.Text += fmt.Sprintf("</%s>", tnStr)
			}
		case html.SelfClosingTagToken, html.CommentToken, html.DoctypeToken:
			// do nothing
		}

		if z.Err() == io.EOF {
			break
		}
	}

	return st, nil
}

func (sr *SamiFormat) Write(insub subtitle.Subtitle) (string, error) {
	doc := &html.Node{
		Type: html.DocumentNode,
	}
	for _, v := range insub.Subtitles {
		htmlText := strings.TrimSpace(html.UnescapeString(v.Text))

		sStartNode := &html.Node{
			Type: html.ElementNode,
			Data: fmt.Sprintf("SYNC Start=%s", timeToSami(v.StartValue)),
		}
		sPNode := &html.Node{
			Type: html.ElementNode,
			Data: "P Class=ENCC",
		}
		sPNode.AppendChild(&html.Node{
			Type: html.TextNode,
			Data: fmt.Sprintf("%s\n", htmlText),
		})
		sStartNode.AppendChild(sPNode)
		doc.AppendChild(sStartNode)

		sEndNode := &html.Node{
			Type: html.ElementNode,
			Data: fmt.Sprintf("SYNC Start=%s", timeToSami(v.EndValue)),
		}
		sPNode = &html.Node{
			Type: html.ElementNode,
			Data: "P Class=ENCC",
		}
		sPNode.AppendChild(&html.Node{
			Type: html.TextNode,
			Data: " \n",
		})
		sEndNode.AppendChild(sPNode)
		doc.AppendChild(sEndNode)
	}

	b := new(bytes.Buffer)
	if err := samiRender(b, doc); err != nil {
		return "", err
	}

	return strings.TrimSpace(b.String()), nil
}

// strip comments in every text node
func stripComments(inStr string) string {
	z := html.NewTokenizer(bytes.NewBufferString(inStr))
	for {
		tt := z.Next()
		if tt == html.ErrorToken {
			if err := z.Err(); err != nil && err != io.EOF {
				return inStr
			}
			break
		}
		if tt == html.CommentToken {
			return strings.Replace(inStr, string(z.Raw()), "", -1)
		}
	}

	return inStr
}

func timeToSami(inTime time.Duration) string {
	totalSec := inTime.Seconds()
	totalMsec := (int(totalSec) * 1000) + int(inTime.Nanoseconds()/1000/1000%1000)
	return fmt.Sprintf("%d", totalMsec)
}

type writer interface {
	io.Writer
	io.ByteWriter
	WriteString(string) (int, error)
}

func samiRender(w io.Writer, n *html.Node) error {
	if x, ok := w.(writer); ok {
		return doSamiRender(x, n)
	}
	buf := bufio.NewWriter(w)
	if err := doSamiRender(buf, n); err != nil {
		return err
	}
	return buf.Flush()
}

func doSamiRender(w writer, n *html.Node) error {
	// Render non-element nodes; these are the easy cases.
	switch n.Type {
	case html.ErrorNode:
		return errors.New("html: cannot render an ErrorNode node")
	case html.TextNode:
		// NOTE: we need to prevent several strings from being escaped,
		// for example, &nbsp", a special termination tag in sami format.
		if hasLegitElemTag(n.Data) {
			_, err := w.WriteString(n.Data)
			return err
		} else {
			return escape(w, n.Data)
		}
	case html.DocumentNode:
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			if err := doSamiRender(w, c); err != nil {
				return err
			}
		}
		return nil
	case html.ElementNode:
		// No-op.
	case html.CommentNode:
		if _, err := w.WriteString("<!--"); err != nil {
			return err
		}
		if _, err := w.WriteString(n.Data); err != nil {
			return err
		}
		if _, err := w.WriteString("-->"); err != nil {
			return err
		}
		return nil
	case html.DoctypeNode:
		if _, err := w.WriteString("<!DOCTYPE "); err != nil {
			return err
		}
		if _, err := w.WriteString(n.Data); err != nil {
			return err
		}
		if n.Attr != nil {
			var p, s string
			for _, a := range n.Attr {
				switch a.Key {
				case "public":
					p = a.Val
				case "system":
					s = a.Val
				}
			}
			if p != "" {
				if _, err := w.WriteString(" PUBLIC "); err != nil {
					return err
				}
				if err := writeQuoted(w, p); err != nil {
					return err
				}
				if s != "" {
					if err := w.WriteByte(' '); err != nil {
						return err
					}
					if err := writeQuoted(w, s); err != nil {
						return err
					}
				}
			} else if s != "" {
				if _, err := w.WriteString(" SYSTEM "); err != nil {
					return err
				}
				if err := writeQuoted(w, s); err != nil {
					return err
				}
			}
		}
		return w.WriteByte('>')
	default:
		return errors.New("html: unknown node type")
	}

	// Render the <xxx> opening tag.
	if err := w.WriteByte('<'); err != nil {
		return err
	}
	if _, err := w.WriteString(n.Data); err != nil {
		return err
	}
	for _, a := range n.Attr {
		if err := w.WriteByte(' '); err != nil {
			return err
		}
		if a.Namespace != "" {
			if _, err := w.WriteString(a.Namespace); err != nil {
				return err
			}
			if err := w.WriteByte(':'); err != nil {
				return err
			}
		}
		if _, err := w.WriteString(a.Key); err != nil {
			return err
		}
		if _, err := w.WriteString(`="`); err != nil {
			return err
		}
		if err := escape(w, a.Val); err != nil {
			return err
		}
		if err := w.WriteByte('"'); err != nil {
			return err
		}
	}
	if err := w.WriteByte('>'); err != nil {
		return err
	}

	// Add initial newline where there is danger of a newline beging ignored.
	if c := n.FirstChild; c != nil && c.Type == html.TextNode && strings.HasPrefix(c.Data, "\n") {
		switch n.Data {
		case "pre", "listing", "textarea":
			if err := w.WriteByte('\n'); err != nil {
				return err
			}
		}
	}

	// Render any child nodes.
	switch n.Data {
	case "iframe", "noembed", "noframes", "noscript", "plaintext", "script", "style", "xmp":
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			if c.Type == html.TextNode {
				if _, err := w.WriteString(c.Data); err != nil {
					return err
				}
			} else {
				if err := doSamiRender(w, c); err != nil {
					return err
				}
			}
		}
	default:
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			if err := doSamiRender(w, c); err != nil {
				return err
			}
		}
	}

	// NOTE: don't render a closing tag at all, for sami files.
	return nil
}

const escapedChars = "&'<>\"\r"

func escape(w writer, s string) error {
	i := strings.IndexAny(s, escapedChars)
	for i != -1 {
		if _, err := w.WriteString(s[:i]); err != nil {
			return err
		}
		var esc string
		switch s[i] {
		case '&':
			esc = "&"
		case '\'':
			// "'" is shorter than "'" and apos was not in HTML until HTML5.
			esc = "'"
		case '<':
			esc = "<"
		case '>':
			esc = ">"
		case '"':
			// """ is shorter than """.
			esc = """
		case '\r':
			esc = "
"
		default:
			panic("unrecognized escape character")
		}
		s = s[i+1:]
		if _, err := w.WriteString(esc); err != nil {
			return err
		}
		i = strings.IndexAny(s, escapedChars)
	}
	_, err := w.WriteString(s)
	return err
}

// writeQuoted writes s to w surrounded by quotes. Normally it will use double
// quotes, but if s contains a double quote, it will use single quotes.
// It is used for writing the identifiers in a doctype declaration.
// In valid HTML, they can't contain both types of quotes.
func writeQuoted(w writer, s string) error {
	var q byte = '"'
	if strings.Contains(s, `"`) {
		q = '\''
	}
	if err := w.WriteByte(q); err != nil {
		return err
	}
	if _, err := w.WriteString(s); err != nil {
		return err
	}
	if err := w.WriteByte(q); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package queue

import (
	"encoding/json"
	"github.com/streadway/amqp"
	"log"
)

var conn *amqp.Connection

type TaskChannel struct {
	ch *amqp.Channel
}

func NewChannel() (*TaskChannel, error) {
	var err error

	if conn == nil {
		conn, err = amqp.Dial("amqp://guest:guest@localhost:5672/")
		if err != nil {
			return nil, err
		}
	}

	ch, err := conn.Channel()
	if err != nil {
		return nil, err
	}

	err = ch.Qos(
		1,     // prefetch count
		0,     // prefetch size
		false, // global
	)
	if err != nil {
		return nil, err
	}

	return &TaskChannel{
		ch: ch,
	}, nil
}

func (c *TaskChannel) Close() error {
	if conn != nil {
		// Connection exists, defer close
		defer conn.Close()
	}

	if c.ch != nil {
		err := c.ch.Close()
		if err != nil {
			return err
		}
	}

	return nil
}

type TaskQueue struct {
	c *TaskChannel
	q *amqp.Queue
}

func NewTaskQueue(c *TaskChannel, queue_name string) (*TaskQueue, error) {
	q, err := c.ch.QueueDeclare(
		queue_name, // name
		true,       // durable
		false,      // delete when unused
		false,      // exclusive
		false,      // no-wait
		nil,        // arguments
	)
	if err != nil {
		return nil, err
	}

	tq := TaskQueue{
		c: c,
		q: &q,
	}

	return &tq, nil
}

func (t TaskQueue) StartConsumer(worker func(interface{}) error, params interface{}, errc chan error) error {
	msgs, err := t.c.ch.Consume(
		t.q.Name, // queue
		"",       // consumer
		false,    // auto-ack
		false,    // exclusive
		false,    // no-local
		false,    // no-wait
		nil,      // args
	)
	if err != nil {
		return err
	}

	go func() {
		for d := range msgs {
			log.Printf("Received a message: %s", d.Body)

			err := json.Unmarshal(d.Body, &params)
			if err != nil {
				// Message is fucked up, don't retry
				d.Reject(false)
				errc <- err
			}

			err = worker(params)
			if err != nil {
				// Reject, don't retry (for now)
				d.Reject(false)
				errc <- err
			}

			d.Ack(false)
		}
	}()

	log.Printf("Started worker for queue '%s'", t.q.Name)

	return nil
}

func (t TaskQueue) AddTask(params interface{}) error {
	body, err := json.Marshal(params)
	if err != nil {
		return err
	}

	err = t.c.ch.Publish(
		"",       // exchange
		t.q.Name, // routing key
		false,    // mandatory
		false,    // immediate
		amqp.Publishing{
			DeliveryMode: amqp.Persistent,
			ContentType:  "application/json",
			Body:         body,
		})
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>DO retry rejects.<commit_after>package queue

import (
	"encoding/json"
	"github.com/streadway/amqp"
	"log"
)

var conn *amqp.Connection

type TaskChannel struct {
	ch *amqp.Channel
}

func NewChannel() (*TaskChannel, error) {
	var err error

	if conn == nil {
		conn, err = amqp.Dial("amqp://guest:guest@localhost:5672/")
		if err != nil {
			return nil, err
		}
	}

	ch, err := conn.Channel()
	if err != nil {
		return nil, err
	}

	err = ch.Qos(
		1,     // prefetch count
		0,     // prefetch size
		false, // global
	)
	if err != nil {
		return nil, err
	}

	return &TaskChannel{
		ch: ch,
	}, nil
}

func (c *TaskChannel) Close() error {
	if conn != nil {
		// Connection exists, defer close
		defer conn.Close()
	}

	if c.ch != nil {
		err := c.ch.Close()
		if err != nil {
			return err
		}
	}

	return nil
}

type TaskQueue struct {
	c *TaskChannel
	q *amqp.Queue
}

func NewTaskQueue(c *TaskChannel, queue_name string) (*TaskQueue, error) {
	q, err := c.ch.QueueDeclare(
		queue_name, // name
		true,       // durable
		false,      // delete when unused
		false,      // exclusive
		false,      // no-wait
		nil,        // arguments
	)
	if err != nil {
		return nil, err
	}

	tq := TaskQueue{
		c: c,
		q: &q,
	}

	return &tq, nil
}

func (t TaskQueue) StartConsumer(worker func(interface{}) error, params interface{}, errc chan error) error {
	msgs, err := t.c.ch.Consume(
		t.q.Name, // queue
		"",       // consumer
		false,    // auto-ack
		false,    // exclusive
		false,    // no-local
		false,    // no-wait
		nil,      // args
	)
	if err != nil {
		return err
	}

	go func() {
		for d := range msgs {
			log.Printf("Received a message: %s", d.Body)

			err := json.Unmarshal(d.Body, &params)
			if err != nil {
				// Message is fucked up, don't retry
				d.Reject(false)
				errc <- err
			}

			err = worker(params)
			if err != nil {
				// Reject, retry
				d.Reject(true)
				errc <- err
			}

			d.Ack(false)
		}
	}()

	log.Printf("Started worker for queue '%s'", t.q.Name)

	return nil
}

func (t TaskQueue) AddTask(params interface{}) error {
	body, err := json.Marshal(params)
	if err != nil {
		return err
	}

	err = t.c.ch.Publish(
		"",       // exchange
		t.q.Name, // routing key
		false,    // mandatory
		false,    // immediate
		amqp.Publishing{
			DeliveryMode: amqp.Persistent,
			ContentType:  "application/json",
			Body:         body,
		})
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package queue

import "github.com/ateleshev/go-bin/cmp"

type Result interface {
	Id() int
	Err() error
	Value() interface{}

	Init(string)
	Bind(interface{}, error)
	Create() interface{}
	Metric() Metric

	cmp.ResetReleaser
}

func NewResult(id int, creator cmp.Creator) Result { // {{{
	return &result{
		id:      id,
		creator: creator,
	}
} // }}}

type result struct {
	id      int
	err     error
	value   interface{}
	metric  Metric
	creator cmp.Creator
}

func (this *result) Id() int { // {{{
	return this.id
} // }}}

func (this *result) Create() interface{} { // {{{
	return this.creator.Create()
} // }}}

func (this *result) Err() error { // {{{
	return this.err
} // }}}

func (this *result) Value() interface{} { // {{{
	return this.value
} // }}}

func (this *result) Metric() Metric { // {{{
	return this.metric
} // }}}

func (this *result) Init(name string) { // {{{
	this.metric = NewMetric(name)
	this.metric.Begin()
} // }}}

func (this *result) Bind(value interface{}, err error) { // {{{
	this.err = err
	this.value = value

	if this.metric != nil {
		this.metric.Finish()
	}
} // }}}

func (this *result) Reset() { // {{{
	this.id = 0
	this.err = nil
	this.value = nil
	this.metric = nil
	this.creator = nil
} // }}}

func (this *result) Release() { // {{{
	resultPoolPut(this)
} // }}}
<commit_msg>Bugfix: correct bind error to queue result<commit_after>package queue

import "github.com/ateleshev/go-bin/cmp"

type Result interface {
	Id() int
	Err() error
	Value() interface{}

	Init(string)
	Bind(interface{}, error)
	Create() interface{}
	Metric() Metric

	cmp.ResetReleaser
}

func NewResult(id int, creator cmp.Creator) Result { // {{{
	return &result{
		id:      id,
		creator: creator,
	}
} // }}}

type result struct {
	id      int
	err     error
	value   interface{}
	metric  Metric
	creator cmp.Creator
}

func (this *result) Id() int { // {{{
	return this.id
} // }}}

func (this *result) Create() interface{} { // {{{
	return this.creator.Create()
} // }}}

func (this *result) Err() error { // {{{
	return this.err
} // }}}

func (this *result) Value() interface{} { // {{{
	return this.value
} // }}}

func (this *result) Metric() Metric { // {{{
	return this.metric
} // }}}

func (this *result) Init(name string) { // {{{
	this.metric = NewMetric(name)
	this.metric.Begin()
} // }}}

func (this *result) Bind(value interface{}, err error) { // {{{
	if err != nil {
		this.err = err
	} else {
		this.value = value
	}

	if this.metric != nil {
		this.metric.Finish()
	}
} // }}}

func (this *result) Reset() { // {{{
	this.id = 0
	this.err = nil
	this.value = nil
	this.metric = nil
	this.creator = nil
} // }}}

func (this *result) Release() { // {{{
	resultPoolPut(this)
} // }}}
<|endoftext|>


<commit_before>package main

import (
	"reflect"
	"testing"
)

func TestNormalizePorts(t *testing.T) {
	actual := NormalizePorts([]int{1, 3, 5})
	expected := map[int]bool{
		1: true,
		3: true,
		5: true,
	}
	if !reflect.DeepEqual(actual, expected) {
		t.Errorf("got %q\nwant %q", actual, expected)
	}
}

func TestNormalizeHost(t *testing.T) {
	suites := map[string]string{
		"127.0.0.1":             "127.0.0.1",
		"www.exsample.com":      "www.exsample.com",
		"www.exsample.com:8080": "www.exsample.com",
	}

	for host, expected := range suites {
		actual := NormalizeHost(host)
		if actual != expected {
			t.Errorf("got %v\nwant %v", actual, expected)
		}
	}
}
<commit_msg>Add more tests.<commit_after>package main

import (
	"bytes"
	"fmt"
	"net/http"
	"net/url"
	"reflect"
	"testing"
)

type TestProxyOptions struct {
	OptAllowedPorts    []int
	OptDestinationHost string
	OptDestinationPort int
	OptVerbose         bool
}

func (o TestProxyOptions) AllowedPorts() []int {
	return o.OptAllowedPorts
}

func (o TestProxyOptions) DestinationHost() string {
	return o.OptDestinationHost
}

func (o TestProxyOptions) DestinationPort() int {
	return o.OptDestinationPort
}

func (o TestProxyOptions) Verbose() bool {
	return o.OptVerbose
}

type TestResponseWriter struct {
	buffer *[]byte
}

func (w TestResponseWriter) Header() http.Header {
	return http.Header{}
}

func (w TestResponseWriter) Write(bytes []byte) (int, error) {
	*w.buffer = append(*w.buffer, bytes...)
	return len(bytes), nil
}

func (w TestResponseWriter) WriteHeader(i int) {
}

func init() {
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, r.URL.Path[1:])
	})
	go (func() {
		http.ListenAndServe(":8980", handler)
	})()
}

func TestNormalizePorts(t *testing.T) {
	actual := NormalizePorts([]int{1, 3, 5})
	expected := map[int]bool{
		1: true,
		3: true,
		5: true,
	}
	if !reflect.DeepEqual(actual, expected) {
		t.Errorf("got %q\nwant %q", actual, expected)
	}
}

func TestNormalizeHost(t *testing.T) {
	suites := map[string]string{
		"127.0.0.1":             "127.0.0.1",
		"www.exsample.com":      "www.exsample.com",
		"www.exsample.com:8080": "www.exsample.com",
	}

	for host, expected := range suites {
		actual := NormalizeHost(host)
		if actual != expected {
			t.Errorf("got %v\nwant %v", actual, expected)
		}
	}
}

func NewWriterReqest() (TestResponseWriter, *http.Request) {
	b := new([]byte)
	w := TestResponseWriter{buffer: b}
	r := new(http.Request)
	r.RemoteAddr = "127.0.0.1:9999"
	r.Header = http.Header{}

	return w, r
}

func TestNewProxyHttpServerGET(t *testing.T) {
	expected := "OK"

	opts := TestProxyOptions{
		OptAllowedPorts:    []int{8980, 443},
		OptDestinationHost: "",
		OptDestinationPort: 0,
		OptVerbose:         false,
	}

	w, r := NewWriterReqest()

	r.URL, _ = url.Parse("http://www.example.com:8980/" + expected)

	server := NewProxyHttpServer(opts)
	server.ServeHTTP(w, r)

	if !bytes.Equal(*w.buffer, []byte(expected)) {
		t.Errorf("got %q\nwant %q", *w.buffer, expected)
	}
}

func TestNewProxyHttpServerDenied(t *testing.T) {
	expected := "You are not allowed to connect to this port: 80"

	opts := TestProxyOptions{
		OptAllowedPorts:    []int{443},
		OptDestinationHost: "",
		OptDestinationPort: 0,
		OptVerbose:         false,
	}

	w, r := NewWriterReqest()

	r.URL, _ = url.Parse("http://www.example.com/OK")

	server := NewProxyHttpServer(opts)
	server.ServeHTTP(w, r)

	if !bytes.Equal(*w.buffer, []byte(expected)) {
		t.Errorf("got %q\nwant %q", *w.buffer, expected)
	}
}

func TestNewProxyHttpServerFixedHost(t *testing.T) {
	expected := "OK"

	opts := TestProxyOptions{
		OptAllowedPorts:    []int{8980, 443},
		OptDestinationHost: "localhost:50",
		OptDestinationPort: 0,
		OptVerbose:         false,
	}

	w, r := NewWriterReqest()

	r.URL, _ = url.Parse("http://www.example.com:8980/" + expected)

	server := NewProxyHttpServer(opts)
	server.ServeHTTP(w, r)

	if !bytes.Equal(*w.buffer, []byte(expected)) {
		t.Errorf("got %q\nwant %q", *w.buffer, expected)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"time"

	"github.com/AdRoll/goamz/aws"
	"github.com/AdRoll/goamz/s3"
)

type mathProblem struct {
	Num1 int
	Num2 int
}

func main() {

	if len(os.Args) > 1 {
		rawData := os.Args[1]

		var mp mathProblem
		err := json.Unmarshal([]byte(rawData), &mp)
		if err != nil {
			log.Println("Error:", err)
			return
		}
		t := time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC)
		h := md5.New()
		io.WriteString(h, fmt.Sprintf("add-%v", t))
		filename := fmt.Sprintf("add-%x", h.Sum(nil))
		answer := fmt.Sprintf("%v + %v = %v", mp.Num1, mp.Num2, mp.Num1+mp.Num2)
		writeToBuck(filename, answer)
		return
	}
	log.Println("Error: os.Args was 1 length.")
}

func getSettings() (string, string, error) {
	file, err := ioutil.ReadFile("./settings.json")
	if err != nil {
		return "", "", nil
	}
	settingsMap := make(map[string]string)
	json.Unmarshal(file, &settingsMap)
	return settingsMap["Access"], settingsMap["Secret"], nil
}

func writeToBuck(f string, a string) {
	p, s, setErr := getSettings()
	if setErr != nil {
		log.Println("Error:", setErr)
		return
	}
	auth := aws.Auth{AccessKey: p, SecretKey: s}

	S3 := s3.New(auth, aws.APNortheast)
	bucket := S3.Bucket("math-results")
	err := bucket.Put(f, []byte(a), "text/plain", s3.PublicRead, s3.Options{})
	if err != nil {
		log.Println("ERROR:", err)
	}
}
<commit_msg>fix time stamp, started kinesis<commit_after>package main

import (
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"time"

	"github.com/AdRoll/goamz/aws"
	"github.com/AdRoll/goamz/s3"
)

type mathProblem struct {
	Num1 int
	Num2 int
}

func main() {

	if len(os.Args) > 1 {
		rawData := os.Args[1]

		var mp mathProblem
		err := json.Unmarshal([]byte(rawData), &mp)
		if err != nil {
			log.Println("Error:", err)
			return
		}
		t := strconv.FormatInt(time.Now().UTC().UnixNano(), 10)
		h := md5.New()
		io.WriteString(h, fmt.Sprintf("add-%v", t))
		filename := fmt.Sprintf("add-%x", h.Sum(nil))
		answer := fmt.Sprintf("%v + %v = %v", mp.Num1, mp.Num2, mp.Num1+mp.Num2)
		writeToBuck(filename, answer)
		return
	}
	log.Println("Error: os.Args was 1 length.")
}

func getSettings() (string, string, error) {
	file, err := ioutil.ReadFile("./settings.json")
	if err != nil {
		return "", "", nil
	}
	settingsMap := make(map[string]string)
	json.Unmarshal(file, &settingsMap)
	return settingsMap["Access"], settingsMap["Secret"], nil
}

func writeToBuck(f string, a string) {
	p, s, setErr := getSettings()
	if setErr != nil {
		log.Println("Error:", setErr)
		return
	}
	auth := aws.Auth{AccessKey: p, SecretKey: s}

	S3 := s3.New(auth, aws.APNortheast)
	bucket := S3.Bucket("math-results")
	err := bucket.Put(f, []byte(a), "text/plain", s3.PublicRead, s3.Options{})
	if err != nil {
		log.Println("ERROR:", err)
	}
}
<|endoftext|>


<commit_before>// +build darwin dragonfly freebsd linux nacl netbsd openbsd solaris

package godirwalk

import (
	"bytes"
	"os"
	"path/filepath"
	"reflect"
	"syscall"
	"unsafe"

	"github.com/pkg/errors"
)

const bufsize = 16 * 1024

func readdirents(osDirname string) (Dirents, error) {
	dh, err := os.Open(osDirname)
	if err != nil {
		return nil, errors.Wrap(err, "cannot Open")
	}

	var entries Dirents
	fd := int(dh.Fd())

	scratchBuffer := make([]byte, bufsize)

	var nameBytes []byte                                     // will be updated to point to syscall.Dirent.Name
	sh := (*reflect.SliceHeader)(unsafe.Pointer(&nameBytes)) // save slice header, so we can re-use each loop

	for {
		n, err := syscall.ReadDirent(fd, scratchBuffer)
		if err != nil {
			_ = dh.Close() // ignore potential error returned by Close
			return nil, errors.Wrap(err, "cannot ReadDirent")
		}
		if n <= 0 {
			break // end of directory reached
		}
		// Loop over the bytes returned by reading the directory entries.
		buf := scratchBuffer[:n]
		for len(buf) > 0 {
			// unshift left-most directory entry from the buffer
			de := (*syscall.Dirent)(unsafe.Pointer(&buf[0]))
			buf = buf[de.Reclen:]

			if de.Ino == 0 {
				continue // this item has been deleted, but not yet removed from directory
			}

			// Convert syscall.Dirent.Name, which is array of int8, to []byte,
			// by overwriting Cap, Len, and Data slice members to values from
			// syscall.Dirent.
			//
			// ??? Set the upper bound on the Cap and the Len to the size of the
			// record length of the dirent.
			sh.Cap, sh.Len, sh.Data = int(de.Reclen), int(de.Reclen), uintptr(unsafe.Pointer(&de.Name[0]))
			nameLength := bytes.IndexByte(nameBytes, 0) // look for NULL byte
			if nameLength == -1 {
				nameLength = len(de.Name)
			}
			nameBytes = nameBytes[:nameLength]

			// Skip "." and ".." entries.
			if nameLength == 1 && nameBytes[0] == '.' || nameLength == 2 && nameBytes[0] == '.' && nameBytes[1] == '.' {
				continue
			}

			nameString := string(nameBytes)

			// Convert syscall constant, which is in purview of OS, to a
			// constant defined by Go, assumed by this project to be stable.
			var mode os.FileMode
			switch de.Type {
			case syscall.DT_REG:
				// regular file
			case syscall.DT_DIR:
				mode = os.ModeDir
			case syscall.DT_LNK:
				mode = os.ModeSymlink
			case syscall.DT_BLK:
				mode = os.ModeDevice
			case syscall.DT_CHR:
				mode = os.ModeDevice | os.ModeCharDevice
			case syscall.DT_FIFO:
				mode = os.ModeNamedPipe
			case syscall.DT_SOCK:
				mode = os.ModeSocket
			default:
				// If syscall returned unknown type (e.g., DT_UNKNOWN, DT_WHT),
				// then resolve actual mode by getting stat.
				fi, err := os.Stat(filepath.Join(osDirname, nameString))
				if err != nil {
					_ = dh.Close() // ignore potential error returned by Close
					return nil, errors.Wrap(err, "cannot Stat")
				}
				mode = fi.Mode()
			}

			// We only care about the bits that identify the type of a file
			// system node, and can ignore append, exclusive, temporary, setuid,
			// setgid, permission bits, and sticky bits, which are coincident to
			// bits which declare type of the file system node.
			entries = append(entries, &Dirent{name: nameString, modeType: mode & os.ModeType})
		}
	}
	if err = dh.Close(); err != nil {
		return nil, err
	}
	return entries, nil
}

func readdirnames(osDirname string) ([]string, error) {
	des, err := readdirents(osDirname)
	if err != nil {
		return nil, err
	}
	names := make([]string, len(des))
	for i, v := range des {
		names[i] = v.name
	}
	return names, nil
}
<commit_msg>unix readdirents sets cap of rec name correctly<commit_after>// +build darwin dragonfly freebsd linux nacl netbsd openbsd solaris

package godirwalk

import (
	"bytes"
	"os"
	"path/filepath"
	"reflect"
	"syscall"
	"unsafe"

	"github.com/pkg/errors"
)

const bufsize = 16 * 1024

func readdirents(osDirname string) (Dirents, error) {
	dh, err := os.Open(osDirname)
	if err != nil {
		return nil, errors.Wrap(err, "cannot Open")
	}

	var de *syscall.Dirent

	var entries Dirents
	fd := int(dh.Fd())

	scratchBuffer := make([]byte, bufsize)

	var nameBytes []byte                                     // will be updated to point to syscall.Dirent.Name
	sh := (*reflect.SliceHeader)(unsafe.Pointer(&nameBytes)) // save slice header, so we can re-use each loop

	for {
		n, err := syscall.ReadDirent(fd, scratchBuffer)
		if err != nil {
			_ = dh.Close() // ignore potential error returned by Close
			return nil, errors.Wrap(err, "cannot ReadDirent")
		}
		if n <= 0 {
			break // end of directory reached
		}
		// Loop over the bytes returned by reading the directory entries.
		buf := scratchBuffer[:n]
		for len(buf) > 0 {
			// unshift left-most directory entry from the buffer
			de = (*syscall.Dirent)(unsafe.Pointer(&buf[0]))
			buf = buf[de.Reclen:]

			if de.Ino == 0 {
				continue // this item has been deleted, but not yet removed from directory
			}

			// Convert syscall.Dirent.Name, which is array of int8, to []byte,
			// by overwriting Cap, Len, and Data slice members to values from
			// syscall.Dirent.
			sh.Cap, sh.Len, sh.Data = len(de.Name), len(de.Name), uintptr(unsafe.Pointer(&de.Name[0]))
			nameLength := bytes.IndexByte(nameBytes, 0) // look for NULL byte
			if nameLength == -1 {
				nameLength = len(de.Name)
			}
			nameBytes = nameBytes[:nameLength]

			// Skip "." and ".." entries.
			if nameLength == 1 && nameBytes[0] == '.' || nameLength == 2 && nameBytes[0] == '.' && nameBytes[1] == '.' {
				continue
			}

			nameString := string(nameBytes)

			// Convert syscall constant, which is in purview of OS, to a
			// constant defined by Go, assumed by this project to be stable.
			var mode os.FileMode
			switch de.Type {
			case syscall.DT_REG:
				// regular file
			case syscall.DT_DIR:
				mode = os.ModeDir
			case syscall.DT_LNK:
				mode = os.ModeSymlink
			case syscall.DT_BLK:
				mode = os.ModeDevice
			case syscall.DT_CHR:
				mode = os.ModeDevice | os.ModeCharDevice
			case syscall.DT_FIFO:
				mode = os.ModeNamedPipe
			case syscall.DT_SOCK:
				mode = os.ModeSocket
			default:
				// If syscall returned unknown type (e.g., DT_UNKNOWN, DT_WHT),
				// then resolve actual mode by getting stat.
				fi, err := os.Stat(filepath.Join(osDirname, nameString))
				if err != nil {
					_ = dh.Close() // ignore potential error returned by Close
					return nil, errors.Wrap(err, "cannot Stat")
				}
				mode = fi.Mode()
			}

			// We only care about the bits that identify the type of a file
			// system node, and can ignore append, exclusive, temporary, setuid,
			// setgid, permission bits, and sticky bits, which are coincident to
			// bits which declare type of the file system node.
			entries = append(entries, &Dirent{name: nameString, modeType: mode & os.ModeType})
		}
	}
	if err = dh.Close(); err != nil {
		return nil, err
	}
	return entries, nil
}

func readdirnames(osDirname string) ([]string, error) {
	des, err := readdirents(osDirname)
	if err != nil {
		return nil, err
	}
	names := make([]string, len(des))
	for i, v := range des {
		names[i] = v.name
	}
	return names, nil
}
<|endoftext|>


<commit_before>// +build !windows

package godirwalk

import (
	"os"
	"syscall"
	"unsafe"
)

func readDirents(osDirname string, scratchBuffer []byte) ([]*Dirent, error) {
	dh, err := os.Open(osDirname)
	if err != nil {
		return nil, err
	}
	fd := int(dh.Fd())

	if len(scratchBuffer) < MinimumScratchBufferSize {
		scratchBuffer = make([]byte, MinimumScratchBufferSize)
	}

	var entries []*Dirent
	var workBuffer []byte
	var sde *syscall.Dirent

	for {
		// When the work buffer has nothing remaining to decode, we need to load
		// more data from disk.
		if len(workBuffer) == 0 {
			n, err := syscall.ReadDirent(fd, scratchBuffer)
			if err != nil {
				_ = dh.Close()
				return nil, err
			}
			if n <= 0 { // end of directory: normal exit
				if err = dh.Close(); err != nil {
					return nil, err
				}
				return entries, nil
			}
			workBuffer = scratchBuffer[:n] // trim work buffer to number of bytes read
		}

		// Loop until we have a usable file system entry, or we run out of data
		// in the work buffer.
		for len(workBuffer) > 0 {
			sde = (*syscall.Dirent)(unsafe.Pointer(&workBuffer[0])) // point entry to first syscall.Dirent in buffer
			workBuffer = workBuffer[reclen(sde):]                   // advance buffer for next iteration through loop

			if inoFromDirent(sde) == 0 {
				continue // inode set to 0 indicates an entry that was marked as deleted
			}

			nameSlice := nameFromDirent(sde)
			nameLength := len(nameSlice)

			if nameLength == 0 || (nameSlice[0] == '.' && (nameLength == 1 || (nameLength == 2 && nameSlice[1] == '.'))) {
				continue
			}

			childName := string(nameSlice)
			mt, err := modeTypeFromDirent(sde, osDirname, childName)
			if err != nil {
				_ = dh.Close()
				return nil, err
			}
			entries = append(entries, &Dirent{name: childName, modeType: mt})
		}
		// No more data in the work buffer, so loop around in the outside loop
		// to fetch more data.
	}

	panic("should not get here")
}

func readDirnames(osDirname string, scratchBuffer []byte) ([]string, error) {
	var entries []string
	scanner, err := NewScannerWithScratchBuffer(osDirname, scratchBuffer)
	if err != nil {
		return nil, err
	}
	for scanner.Scan() {
		entries = append(entries, scanner.Name())
	}
	if err = scanner.Err(); err != nil {
		return nil, err
	}
	return entries, nil
}
<commit_msg>readDirnames optimized<commit_after>// +build !windows

package godirwalk

import (
	"os"
	"syscall"
	"unsafe"
)

func readDirents(osDirname string, scratchBuffer []byte) ([]*Dirent, error) {
	dh, err := os.Open(osDirname)
	if err != nil {
		return nil, err
	}
	fd := int(dh.Fd())

	if len(scratchBuffer) < MinimumScratchBufferSize {
		scratchBuffer = make([]byte, MinimumScratchBufferSize)
	}

	var entries []*Dirent
	var workBuffer []byte
	var sde *syscall.Dirent

	for {
		// When the work buffer has nothing remaining to decode, we need to load
		// more data from disk.
		if len(workBuffer) == 0 {
			n, err := syscall.ReadDirent(fd, scratchBuffer)
			if err != nil {
				_ = dh.Close()
				return nil, err
			}
			if n <= 0 { // end of directory: normal exit
				if err = dh.Close(); err != nil {
					return nil, err
				}
				return entries, nil
			}
			workBuffer = scratchBuffer[:n] // trim work buffer to number of bytes read
		}

		// Loop until we have a usable file system entry, or we run out of data
		// in the work buffer.
		for len(workBuffer) > 0 {
			sde = (*syscall.Dirent)(unsafe.Pointer(&workBuffer[0])) // point entry to first syscall.Dirent in buffer
			workBuffer = workBuffer[reclen(sde):]                   // advance buffer for next iteration through loop

			if inoFromDirent(sde) == 0 {
				continue // inode set to 0 indicates an entry that was marked as deleted
			}

			nameSlice := nameFromDirent(sde)
			nameLength := len(nameSlice)

			if nameLength == 0 || (nameSlice[0] == '.' && (nameLength == 1 || (nameLength == 2 && nameSlice[1] == '.'))) {
				continue
			}

			childName := string(nameSlice)
			mt, err := modeTypeFromDirent(sde, osDirname, childName)
			if err != nil {
				_ = dh.Close()
				return nil, err
			}
			entries = append(entries, &Dirent{name: childName, modeType: mt})
		}
		// No more data in the work buffer, so loop around in the outside loop
		// to fetch more data.
	}

	panic("should not get here")
}

func readDirnames(osDirname string, scratchBuffer []byte) ([]string, error) {
	dh, err := os.Open(osDirname)
	if err != nil {
		return nil, err
	}
	fd := int(dh.Fd())

	if len(scratchBuffer) < MinimumScratchBufferSize {
		scratchBuffer = make([]byte, MinimumScratchBufferSize)
	}

	var entries []string
	var workBuffer []byte
	var sde *syscall.Dirent

	for {
		// When the work buffer has nothing remaining to decode, we need to load
		// more data from disk.
		if len(workBuffer) == 0 {
			n, err := syscall.ReadDirent(fd, scratchBuffer)
			if err != nil {
				_ = dh.Close()
				return nil, err
			}
			if n <= 0 { // end of directory: normal exit
				if err = dh.Close(); err != nil {
					return nil, err
				}
				return entries, nil
			}
			workBuffer = scratchBuffer[:n] // trim work buffer to number of bytes read
		}

		// Loop until we have a usable file system entry, or we run out of data
		// in the work buffer.
		for len(workBuffer) > 0 {
			sde = (*syscall.Dirent)(unsafe.Pointer(&workBuffer[0])) // point entry to first syscall.Dirent in buffer
			workBuffer = workBuffer[reclen(sde):]                   // advance buffer for next iteration through loop

			if inoFromDirent(sde) == 0 {
				continue // inode set to 0 indicates an entry that was marked as deleted
			}

			nameSlice := nameFromDirent(sde)
			nameLength := len(nameSlice)

			if nameLength == 0 || (nameSlice[0] == '.' && (nameLength == 1 || (nameLength == 2 && nameSlice[1] == '.'))) {
				continue
			}

			entries = append(entries, string(nameSlice))
		}
		// No more data in the work buffer, so loop around in the outside loop
		// to fetch more data.
	}

	panic("should not get here")
}
<|endoftext|>


<commit_before>package reader

import (
	"encoding/json"
	"errors"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"fmt"

	"github.com/Shopify/sarama"
	"github.com/qiniu/log"
	"github.com/qiniu/logkit/utils"
	"github.com/wvanbergen/kafka/consumergroup"
)

type KafkaReader struct {
	meta            *Meta
	ConsumerGroup   string
	Topics          []string
	ZookeeperPeers  []string
	ZookeeperChroot string
	Whence          string

	Consumer *consumergroup.ConsumerGroup
	in       <-chan *sarama.ConsumerMessage
	curMsg   *sarama.ConsumerMessage

	readChan chan json.RawMessage
	errs     <-chan error

	status   int32
	mux      sync.Mutex
	startMux sync.Mutex
	started  bool

	stats     utils.StatsInfo
	statsLock sync.RWMutex
}

func NewKafkaReader(meta *Meta, consumerGroup string,
	topics []string, zookeeper []string, whence string) (kr *KafkaReader, err error) {
	kr = &KafkaReader{
		meta:           meta,
		ConsumerGroup:  consumerGroup,
		ZookeeperPeers: zookeeper,
		Topics:         topics,
		Whence:         whence,
		readChan:       make(chan json.RawMessage),
		errs:           make(chan error, 1000),
		status:         StatusInit,
		mux:            sync.Mutex{},
		startMux:       sync.Mutex{},
		statsLock:      sync.RWMutex{},
		started:        false,
	}
	return kr, nil
}

func (kr *KafkaReader) Name() string {
	return fmt.Sprintf("KafkaReader:[%s],[%s]", strings.Join(kr.Topics, ","), kr.ConsumerGroup)
}

func (kr *KafkaReader) Status() utils.StatsInfo {
	kr.statsLock.RLock()
	defer kr.statsLock.RUnlock()
	return kr.stats
}

func (kr *KafkaReader) setStatsError(err string) {
	kr.statsLock.Lock()
	defer kr.statsLock.Unlock()
	kr.stats.LastError = err
}

func (kr *KafkaReader) Source() string {
	return fmt.Sprintf("[%s],[%s]", strings.Join(kr.Topics, ","), kr.ConsumerGroup)
}

func (kr *KafkaReader) ReadLine() (data string, err error) {
	if !kr.started {
		kr.Start()
	}
	timer := time.NewTimer(time.Second)
	select {
	case dat := <-kr.readChan:
		data = string(dat)
	case <-timer.C:
	}
	timer.Stop()
	return
}

func (kr *KafkaReader) Close() (err error) {
	if atomic.CompareAndSwapInt32(&kr.status, StatusRunning, StatusStoping) {
		log.Infof("Runner[%v] %v stopping", kr.meta.RunnerName, kr.Name())
	} else {
		close(kr.readChan)
	}
	return
}

func (kr *KafkaReader) SyncMeta() {
	kr.mux.Lock()
	if kr.curMsg != nil {
		kr.Consumer.CommitUpto(kr.curMsg)
	}
	kr.mux.Unlock()
	return
}

func (kr *KafkaReader) Start() {
	kr.startMux.Lock()
	defer kr.startMux.Unlock()
	if kr.started {
		return
	}
	config := consumergroup.NewConfig()
	config.Zookeeper.Chroot = ""
	switch strings.ToLower(kr.Whence) {
	case "oldest", "":
		config.Offsets.Initial = sarama.OffsetOldest
	case "newest":
		config.Offsets.Initial = sarama.OffsetNewest
	default:
		log.Warnf("Runner[%v] WARNING: Kafka consumer invalid offset '%s', using 'oldest'\n", kr.meta.RunnerName, kr.Whence)
		config.Offsets.Initial = sarama.OffsetNewest
	}
	if kr.Consumer == nil {
		var consumerErr error
		kr.Consumer, consumerErr = consumergroup.JoinConsumerGroup(
			kr.ConsumerGroup,
			kr.Topics,
			kr.ZookeeperPeers,
			config,
		)
		if consumerErr != nil {
			log.Errorf("%v", consumerErr)
			return
		}
		kr.in = kr.Consumer.Messages()
	}
	go kr.run()
	kr.started = true
	log.Infof("Runner[%v] %v pull data daemon started", kr.meta.RunnerName, kr.Name())
}

func (kr *KafkaReader) run() {
	// 防止并发run
	for {
		if atomic.LoadInt32(&kr.status) == StatusStopped {
			return
		}
		if atomic.CompareAndSwapInt32(&kr.status, StatusInit, StatusRunning) {
			break
		}
	}
	// running在退出状态改为Init
	defer func() {
		atomic.CompareAndSwapInt32(&kr.status, StatusRunning, StatusInit)
		if atomic.CompareAndSwapInt32(&kr.status, StatusStoping, StatusStopped) {
			close(kr.readChan)
		}
		log.Infof("Runner[%v] %v successfully finished", kr.meta.RunnerName, kr.Name())
	}()
	// 开始work逻辑
	for {
		if atomic.LoadInt32(&kr.status) == StatusStoping {
			log.Warnf("Runner[%v] %v stopped from running", kr.meta.RunnerName, kr.Name())
			return
		}
		select {
		case err := <-kr.errs:
			if err != nil {
				log.Errorf("Runner[%v] Consumer Error: %s\n", kr.meta.RunnerName, err)
				kr.setStatsError("Runner[" + kr.meta.RunnerName + "] Consumer Error: " + err.Error())
			}
		case msg := <-kr.in:
			kr.readChan <- msg.Value
			kr.curMsg = msg
		}
	}
}

func (kr *KafkaReader) SetMode(mode string, v interface{}) error {
	return errors.New("KafkaReader not support read mode")
}
<commit_msg>修复kafka reader close方法没有调用consumer close的bug<commit_after>package reader

import (
	"encoding/json"
	"errors"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"fmt"

	"github.com/Shopify/sarama"
	"github.com/qiniu/log"
	"github.com/qiniu/logkit/utils"
	"github.com/wvanbergen/kafka/consumergroup"
)

type KafkaReader struct {
	meta            *Meta
	ConsumerGroup   string
	Topics          []string
	ZookeeperPeers  []string
	ZookeeperChroot string
	Whence          string

	Consumer *consumergroup.ConsumerGroup
	in       <-chan *sarama.ConsumerMessage
	curMsg   *sarama.ConsumerMessage

	readChan chan json.RawMessage
	errs     <-chan error

	status   int32
	mux      sync.Mutex
	startMux sync.Mutex
	started  bool

	stats     utils.StatsInfo
	statsLock sync.RWMutex
}

func NewKafkaReader(meta *Meta, consumerGroup string,
	topics []string, zookeeper []string, whence string) (kr *KafkaReader, err error) {
	kr = &KafkaReader{
		meta:           meta,
		ConsumerGroup:  consumerGroup,
		ZookeeperPeers: zookeeper,
		Topics:         topics,
		Whence:         whence,
		readChan:       make(chan json.RawMessage),
		errs:           make(chan error, 1000),
		status:         StatusInit,
		mux:            sync.Mutex{},
		startMux:       sync.Mutex{},
		statsLock:      sync.RWMutex{},
		started:        false,
	}
	return kr, nil
}

func (kr *KafkaReader) Name() string {
	return fmt.Sprintf("KafkaReader:[%s],[%s]", strings.Join(kr.Topics, ","), kr.ConsumerGroup)
}

func (kr *KafkaReader) Status() utils.StatsInfo {
	kr.statsLock.RLock()
	defer kr.statsLock.RUnlock()
	return kr.stats
}

func (kr *KafkaReader) setStatsError(err string) {
	kr.statsLock.Lock()
	defer kr.statsLock.Unlock()
	kr.stats.LastError = err
}

func (kr *KafkaReader) Source() string {
	return fmt.Sprintf("[%s],[%s]", strings.Join(kr.Topics, ","), kr.ConsumerGroup)
}

func (kr *KafkaReader) ReadLine() (data string, err error) {
	if !kr.started {
		kr.Start()
	}
	timer := time.NewTimer(time.Second)
	select {
	case dat := <-kr.readChan:
		data = string(dat)
	case <-timer.C:
	}
	timer.Stop()
	return
}

func (kr *KafkaReader) Close() (err error) {
	if atomic.CompareAndSwapInt32(&kr.status, StatusRunning, StatusStoping) {
		log.Infof("Runner[%v] %v stopping", kr.meta.RunnerName, kr.Name())
	} else {
		close(kr.readChan)
	}
	return
}

func (kr *KafkaReader) SyncMeta() {
	kr.mux.Lock()
	if kr.curMsg != nil {
		kr.Consumer.CommitUpto(kr.curMsg)
	}
	kr.mux.Unlock()
	return
}

func (kr *KafkaReader) Start() {
	kr.startMux.Lock()
	defer kr.startMux.Unlock()
	if kr.started {
		return
	}
	config := consumergroup.NewConfig()
	config.Zookeeper.Chroot = ""
	switch strings.ToLower(kr.Whence) {
	case "oldest", "":
		config.Offsets.Initial = sarama.OffsetOldest
	case "newest":
		config.Offsets.Initial = sarama.OffsetNewest
	default:
		log.Warnf("Runner[%v] WARNING: Kafka consumer invalid offset '%s', using 'oldest'\n", kr.meta.RunnerName, kr.Whence)
		config.Offsets.Initial = sarama.OffsetNewest
	}
	if kr.Consumer == nil || kr.Consumer.Closed() {
		var consumerErr error
		kr.Consumer, consumerErr = consumergroup.JoinConsumerGroup(
			kr.ConsumerGroup,
			kr.Topics,
			kr.ZookeeperPeers,
			config,
		)
		if consumerErr != nil {
			log.Errorf("%v", consumerErr)
			return
		}
		kr.in = kr.Consumer.Messages()
	}
	go kr.run()
	kr.started = true
	log.Infof("Runner[%v] %v pull data daemon started", kr.meta.RunnerName, kr.Name())
}

func (kr *KafkaReader) run() {
	// 防止并发run
	for {
		if atomic.LoadInt32(&kr.status) == StatusStopped {
			return
		}
		if atomic.CompareAndSwapInt32(&kr.status, StatusInit, StatusRunning) {
			break
		}
	}
	// running在退出状态改为Init
	defer func() {
		atomic.CompareAndSwapInt32(&kr.status, StatusRunning, StatusInit)
		if atomic.CompareAndSwapInt32(&kr.status, StatusStoping, StatusStopped) {
			close(kr.readChan)
			go func() {
				kr.mux.Lock()
				defer kr.mux.Unlock()
				err := kr.Consumer.Close()
				if err != nil {
					log.Infof("Runner[%v] %v stop failed error: %v", kr.meta.RunnerName, kr.Name(), err)
				}
			}()
		}
		log.Infof("Runner[%v] %v successfully finished", kr.meta.RunnerName, kr.Name())
	}()
	// 开始work逻辑
	for {
		if atomic.LoadInt32(&kr.status) == StatusStoping {
			log.Warnf("Runner[%v] %v stopped from running", kr.meta.RunnerName, kr.Name())
			return
		}
		select {
		case err := <-kr.errs:
			if err != nil {
				log.Errorf("Runner[%v] Consumer Error: %s\n", kr.meta.RunnerName, err)
				kr.setStatsError("Runner[" + kr.meta.RunnerName + "] Consumer Error: " + err.Error())
			}
		case msg := <-kr.in:
			if msg != nil && msg.Value != nil && len(msg.Value) > 0 {
				kr.readChan <- msg.Value
				kr.curMsg = msg
			}
		default:
			time.Sleep(time.Second)
		}
	}
}

func (kr *KafkaReader) SetMode(mode string, v interface{}) error {
	return errors.New("KafkaReader not support read mode")
}
<|endoftext|>


<commit_before>package pes

import "bytes"
import "github.com/32bitkid/mpeg/ts"
import "github.com/32bitkid/bitreader"

type Decoder interface {
	Output() PacketChannel
	Err() error
}

type decoder struct {
	output chan *Packet
	err    error
}

func (d *decoder) Output() PacketChannel {
	return d.output
}

func (d *decoder) Err() error {
	return d.err
}

func NewTsDecoder(newFn bitreader.NewReader32Fn, input ts.PacketChannel) Decoder {
	d := decoder{
		output: make(chan *Packet),
	}

	buffer := &bytes.Buffer{}
	reader := newFn(buffer)

	go func() {
		defer close(d.output)

		for tsPacket := range input {

			if tsPacket.PayloadUnitStartIndicator && buffer.Len() > 0 {
				// Drain
				pesPacket, err := ReadPacket(reader, buffer.Len())
				if err != nil {
					d.err = err
					return
				}
				d.output <- pesPacket
			}

			// Fill
			buffer.Write(tsPacket.Payload)
		}
	}()

	return &d
}
<commit_msg>Let pes decoders have multiple outputs<commit_after>package pes

import "bytes"
import "github.com/32bitkid/mpeg/ts"
import "github.com/32bitkid/bitreader"

type Decoder interface {
	Output() PacketChannel
	Err() error
}

type decoder struct {
	outputs []chan *Packet
	err     error
}

func (d *decoder) Output() PacketChannel {
	newOutput := make(chan *Packet)
	d.outputs = append(d.outputs, newOutput)
	return newOutput
}

func (d *decoder) Err() error {
	return d.err
}

func NewTsDecoder(newFn bitreader.NewReader32Fn, input ts.PacketChannel) Decoder {
	d := decoder{}

	buffer := &bytes.Buffer{}
	reader := newFn(buffer)

	go func() {
		defer func() {
			for _, output := range d.outputs {
				close(output)
			}
		}()

		for tsPacket := range input {

			if tsPacket.PayloadUnitStartIndicator && buffer.Len() > 0 {
				// Drain
				pesPacket, err := ReadPacket(reader, buffer.Len())
				if err != nil {
					d.err = err
					return
				}
				for _, output := range d.outputs {
					output <- pesPacket
				}
			}

			// Fill
			buffer.Write(tsPacket.Payload)
		}
	}()

	return &d
}
<|endoftext|>


<commit_before>// Copyright 2011 Bobby Powers. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.
package seshcookie

import (
	"bufio"
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/base64"
	"encoding/gob"
	"errors"
	"hash"
	"io"
	"log"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
)

// we want 16 byte blocks, for AES-128
const blockSize = 16

var (
	// if you don't need multiple independent seshcookie
	// instances, you can use this RequestSessions instance to
	// manage & access your sessions.  Simply use it as the final
	// parameter in your call to seshcookie.NewSessionHandler, and
	// whenever you want to access the current session from an
	// embedded http.Handler you can simply call:
	//
	//     seshcookie.Session.Get(req)
	Session = &RequestSessions{HttpOnly: true}

	// Hash validation of the decrypted cookie failed. Most likely
	// the session was encoded with a different cookie than we're
	// using to decode it, but its possible the client (or someone
	// else) tried to modify the session.
	HashError = errors.New("Hash validation failed")

	// The cookie is too short, so we must exit decoding early.
	LenError = errors.New("Bad cookie length")
)

type sessionResponseWriter struct {
	http.ResponseWriter
	h   *SessionHandler
	req *http.Request
	// int32 so we can use the sync/atomic functions on it
	wroteHeader int32
}

type SessionHandler struct {
	http.Handler
	CookieName string // name of the cookie to store our session in
	CookiePath string // resource path the cookie is valid for
	RS         *RequestSessions
	encKey     []byte
	hmacKey    []byte
}

type RequestSessions struct {
	HttpOnly bool // don't allow javascript to access cookie
	Secure   bool // only send session over HTTPS
	lk       sync.Mutex
	m        map[*http.Request]map[string]interface{}
	// stores a hash of the serialized session (the gob) that we
	// received with the start of the request.  Before setting a
	// cookie for the reply, check to see if the session has
	// actually changed.  If it hasn't, then we don't need to send
	// a new cookie.
	hm map[*http.Request][]byte
}

func (rs *RequestSessions) Get(req *http.Request) map[string]interface{} {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	if rs.m == nil {
		log.Print("seshcookie: warning! trying to get session " +
			"data for unknown request. Perhaps your handler " +
			"isn't wrapped by a SessionHandler?")
		return nil
	}

	return rs.m[req]
}

func (rs *RequestSessions) getHash(req *http.Request) []byte {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	if rs.hm == nil {
		return nil
	}

	return rs.hm[req]
}

func (rs *RequestSessions) Set(req *http.Request, val map[string]interface{}, gobHash []byte) {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	if rs.m == nil {
		rs.m = map[*http.Request]map[string]interface{}{}
		rs.hm = map[*http.Request][]byte{}
	}

	rs.m[req] = val
	rs.hm[req] = gobHash
}

func (rs *RequestSessions) Clear(req *http.Request) {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	delete(rs.m, req)
}

func encodeGob(obj interface{}) ([]byte, error) {
	buf := bytes.NewBuffer(nil)
	enc := gob.NewEncoder(buf)
	err := enc.Encode(obj)
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func decodeGob(encoded []byte) (map[string]interface{}, error) {
	buf := bytes.NewBuffer(encoded)
	dec := gob.NewDecoder(buf)
	var out map[string]interface{}
	err := dec.Decode(&out)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// encode uses the given block cipher (in CTR mode) to encrypt the
// data, along with a hash, returning the iv and the ciphertext. What
// is returned looks like:
//
//   encrypted(salt + sessionData) + iv + hmac
//
func encode(block cipher.Block, hmac hash.Hash, data []byte) ([]byte, error) {

	buf := bytes.NewBuffer(nil)

	salt := make([]byte, block.BlockSize())
	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
		return nil, err
	}
	buf.Write(salt)
	buf.Write(data)

	session := buf.Bytes()

	iv := make([]byte, block.BlockSize())
	if _, err := rand.Read(iv); err != nil {
		return nil, err
	}

	stream := cipher.NewCTR(block, iv)
	stream.XORKeyStream(session, session)

	buf.Write(iv)
	hmac.Write(buf.Bytes())
	buf.Write(hmac.Sum())

	return buf.Bytes(), nil
}

func encodeCookie(content interface{}, encKey, hmacKey []byte) (string, []byte, error) {
	encodedGob, err := encodeGob(content)
	if err != nil {
		return "", nil, err
	}

	gobHash := sha1.New()
	gobHash.Write(encodedGob)

	aesCipher, err := aes.NewCipher(encKey)
	if err != nil {
		return "", nil, err
	}

	hmacHash := hmac.NewSHA256(hmacKey)

	sessionBytes, err := encode(aesCipher, hmacHash, encodedGob)
	if err != nil {
		return "", nil, err
	}

	return base64.StdEncoding.EncodeToString(sessionBytes), gobHash.Sum(), nil
}

// decode uses the given block cipher (in CTR mode) to decrypt the
// data, and validate the hash.  If hash validation fails, an error is
// returned.
func decode(block cipher.Block, hmac hash.Hash, ciphertext []byte) ([]byte, error) {
	if len(ciphertext) < 2*block.BlockSize()+hmac.Size() {
		return nil, LenError
	}

	receivedHmac := ciphertext[len(ciphertext)-hmac.Size():]
	ciphertext = ciphertext[:len(ciphertext)-hmac.Size()]

	hmac.Write(ciphertext)
	if subtle.ConstantTimeCompare(hmac.Sum(), receivedHmac) != 1 {
		return nil, HashError
	}

	// split the iv and session bytes
	iv := ciphertext[len(ciphertext)-block.BlockSize():]
	session := ciphertext[:len(ciphertext)-block.BlockSize()]

	stream := cipher.NewCTR(block, iv)
	stream.XORKeyStream(session, session)

	// skip past the iv
	session = session[block.BlockSize():]

	return session, nil
}

func decodeCookie(encoded string, encKey, hmacKey []byte) (map[string]interface{}, []byte, error) {
	sessionBytes, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		return nil, nil, err
	}
	aesCipher, err := aes.NewCipher(encKey)
	if err != nil {
		return nil, nil, err
	}

	hmacHash := hmac.NewSHA256(hmacKey)
	gobBytes, err := decode(aesCipher, hmacHash, sessionBytes)
	if err != nil {
		return nil, nil, err
	}

	gobHash := sha1.New()
	gobHash.Write(gobBytes)

	session, err := decodeGob(gobBytes)
	if err != nil {
		log.Printf("decodeGob: %s\n", err)
		return nil, nil, err
	}
	return session, gobHash.Sum(), nil
}

func (s sessionResponseWriter) WriteHeader(code int) {
	if atomic.AddInt32(&s.wroteHeader, 1) == 1 {
		origCookie, err := s.req.Cookie(s.h.CookieName)
		var origCookieVal string
		if err != nil {
			origCookieVal = ""
		} else {
			origCookieVal = origCookie.Value
		}

		session := s.h.RS.Get(s.req)
		if len(session) == 0 {
			// if we have an empty session, but the
			// request didn't start out that way, we
			// assume the user wants us to clear the
			// session
			if origCookieVal != "" {
				//log.Println("clearing cookie")
				var cookie http.Cookie
				cookie.Name = s.h.CookieName
				cookie.Value = ""
				cookie.Path = "/"
				// a cookie is expired by setting it
				// with an expiration time in the past
				cookie.Expires = *time.SecondsToUTC(0)
				http.SetCookie(s, &cookie)
			}
			goto write
		}
		encoded, gobHash, err := encodeCookie(session, s.h.encKey, s.h.hmacKey)
		if err != nil {
			log.Printf("createCookie: %s\n", err)
			goto write
		}

		if bytes.Equal(gobHash, s.h.RS.getHash(s.req)) {
			//log.Println("not re-setting identical cookie")
			goto write
		}

		var cookie http.Cookie
		cookie.Name = s.h.CookieName
		cookie.Value = encoded
		cookie.Path = s.h.CookiePath
		cookie.HttpOnly = s.h.RS.HttpOnly
		cookie.Secure = s.h.RS.Secure
		http.SetCookie(s, &cookie)
	}
write:
	s.ResponseWriter.WriteHeader(code)
}

func (s sessionResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hijacker, _ := s.ResponseWriter.(http.Hijacker)
	return hijacker.Hijack()
}

func (h *SessionHandler) getCookieSession(req *http.Request) (map[string]interface{}, []byte) {
	cookie, err := req.Cookie(h.CookieName)
	if err != nil {
		//log.Printf("getCookieSesh: '%#v' not found\n",
		//	h.CookieName)
		return map[string]interface{}{}, nil
	}
	session, gobHash, err := decodeCookie(cookie.Value, h.encKey, h.hmacKey)
	if err != nil {
		log.Printf("decodeCookie: %s\n", err)
		return map[string]interface{}{}, nil
	}

	return session, gobHash
}

func (h *SessionHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// get our session a little early, so that we can add our
	// authentication information to it if we get some
	session, gobHash := h.getCookieSession(req)

	h.RS.Set(req, session, gobHash)

	sessionWriter := sessionResponseWriter{rw, h, req, 0}
	h.Handler.ServeHTTP(sessionWriter, req)

	h.RS.Clear(req)
}

func NewSessionHandler(handler http.Handler, key string, rs *RequestSessions) *SessionHandler {
	// sha1 sums are 20 bytes long.  we use the first 16 bytes as
	// the aes key.
	encHash := sha1.New()
	encHash.Write([]byte(key))
	encHash.Write([]byte("-encryption"))
	hmacHash := sha1.New()
	hmacHash.Write([]byte(key))
	hmacHash.Write([]byte("-hmac"))

	// if the user hasn't specified a session handler, use the
	// package's default one
	if rs == nil {
		rs = Session
	}

	return &SessionHandler{
		Handler:    handler,
		CookieName: "session",
		CookiePath: "/",
		RS:         rs,
		encKey:     encHash.Sum()[:blockSize],
		hmacKey:    hmacHash.Sum()[:blockSize],
	}
}
<commit_msg>don't hold on to (leak) http.Requests<commit_after>// Copyright 2011 Bobby Powers. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.
package seshcookie

import (
	"bufio"
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/base64"
	"encoding/gob"
	"errors"
	"hash"
	"io"
	"log"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
)

// we want 16 byte blocks, for AES-128
const blockSize = 16

var (
	// if you don't need multiple independent seshcookie
	// instances, you can use this RequestSessions instance to
	// manage & access your sessions.  Simply use it as the final
	// parameter in your call to seshcookie.NewSessionHandler, and
	// whenever you want to access the current session from an
	// embedded http.Handler you can simply call:
	//
	//     seshcookie.Session.Get(req)
	Session = &RequestSessions{HttpOnly: true}

	// Hash validation of the decrypted cookie failed. Most likely
	// the session was encoded with a different cookie than we're
	// using to decode it, but its possible the client (or someone
	// else) tried to modify the session.
	HashError = errors.New("Hash validation failed")

	// The cookie is too short, so we must exit decoding early.
	LenError = errors.New("Bad cookie length")
)

type sessionResponseWriter struct {
	http.ResponseWriter
	h   *SessionHandler
	req *http.Request
	// int32 so we can use the sync/atomic functions on it
	wroteHeader int32
}

type SessionHandler struct {
	http.Handler
	CookieName string // name of the cookie to store our session in
	CookiePath string // resource path the cookie is valid for
	RS         *RequestSessions
	encKey     []byte
	hmacKey    []byte
}

type RequestSessions struct {
	HttpOnly bool // don't allow javascript to access cookie
	Secure   bool // only send session over HTTPS
	lk       sync.Mutex
	m        map[*http.Request]map[string]interface{}
	// stores a hash of the serialized session (the gob) that we
	// received with the start of the request.  Before setting a
	// cookie for the reply, check to see if the session has
	// actually changed.  If it hasn't, then we don't need to send
	// a new cookie.
	hm map[*http.Request][]byte
}

func (rs *RequestSessions) Get(req *http.Request) map[string]interface{} {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	if rs.m == nil {
		log.Print("seshcookie: warning! trying to get session " +
			"data for unknown request. Perhaps your handler " +
			"isn't wrapped by a SessionHandler?")
		return nil
	}

	return rs.m[req]
}

func (rs *RequestSessions) getHash(req *http.Request) []byte {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	if rs.hm == nil {
		return nil
	}

	return rs.hm[req]
}

func (rs *RequestSessions) Set(req *http.Request, val map[string]interface{}, gobHash []byte) {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	if rs.m == nil {
		rs.m = map[*http.Request]map[string]interface{}{}
		rs.hm = map[*http.Request][]byte{}
	}

	rs.m[req] = val
	rs.hm[req] = gobHash
}

func (rs *RequestSessions) Clear(req *http.Request) {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	delete(rs.m, req)
	delete(rs.hm, req)
}

func encodeGob(obj interface{}) ([]byte, error) {
	buf := bytes.NewBuffer(nil)
	enc := gob.NewEncoder(buf)
	err := enc.Encode(obj)
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func decodeGob(encoded []byte) (map[string]interface{}, error) {
	buf := bytes.NewBuffer(encoded)
	dec := gob.NewDecoder(buf)
	var out map[string]interface{}
	err := dec.Decode(&out)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// encode uses the given block cipher (in CTR mode) to encrypt the
// data, along with a hash, returning the iv and the ciphertext. What
// is returned looks like:
//
//   encrypted(salt + sessionData) + iv + hmac
//
func encode(block cipher.Block, hmac hash.Hash, data []byte) ([]byte, error) {

	buf := bytes.NewBuffer(nil)

	salt := make([]byte, block.BlockSize())
	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
		return nil, err
	}
	buf.Write(salt)
	buf.Write(data)

	session := buf.Bytes()

	iv := make([]byte, block.BlockSize())
	if _, err := rand.Read(iv); err != nil {
		return nil, err
	}

	stream := cipher.NewCTR(block, iv)
	stream.XORKeyStream(session, session)

	buf.Write(iv)
	hmac.Write(buf.Bytes())
	buf.Write(hmac.Sum())

	return buf.Bytes(), nil
}

func encodeCookie(content interface{}, encKey, hmacKey []byte) (string, []byte, error) {
	encodedGob, err := encodeGob(content)
	if err != nil {
		return "", nil, err
	}

	gobHash := sha1.New()
	gobHash.Write(encodedGob)

	aesCipher, err := aes.NewCipher(encKey)
	if err != nil {
		return "", nil, err
	}

	hmacHash := hmac.NewSHA256(hmacKey)

	sessionBytes, err := encode(aesCipher, hmacHash, encodedGob)
	if err != nil {
		return "", nil, err
	}

	return base64.StdEncoding.EncodeToString(sessionBytes), gobHash.Sum(), nil
}

// decode uses the given block cipher (in CTR mode) to decrypt the
// data, and validate the hash.  If hash validation fails, an error is
// returned.
func decode(block cipher.Block, hmac hash.Hash, ciphertext []byte) ([]byte, error) {
	if len(ciphertext) < 2*block.BlockSize()+hmac.Size() {
		return nil, LenError
	}

	receivedHmac := ciphertext[len(ciphertext)-hmac.Size():]
	ciphertext = ciphertext[:len(ciphertext)-hmac.Size()]

	hmac.Write(ciphertext)
	if subtle.ConstantTimeCompare(hmac.Sum(), receivedHmac) != 1 {
		return nil, HashError
	}

	// split the iv and session bytes
	iv := ciphertext[len(ciphertext)-block.BlockSize():]
	session := ciphertext[:len(ciphertext)-block.BlockSize()]

	stream := cipher.NewCTR(block, iv)
	stream.XORKeyStream(session, session)

	// skip past the iv
	session = session[block.BlockSize():]

	return session, nil
}

func decodeCookie(encoded string, encKey, hmacKey []byte) (map[string]interface{}, []byte, error) {
	sessionBytes, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		return nil, nil, err
	}
	aesCipher, err := aes.NewCipher(encKey)
	if err != nil {
		return nil, nil, err
	}

	hmacHash := hmac.NewSHA256(hmacKey)
	gobBytes, err := decode(aesCipher, hmacHash, sessionBytes)
	if err != nil {
		return nil, nil, err
	}

	gobHash := sha1.New()
	gobHash.Write(gobBytes)

	session, err := decodeGob(gobBytes)
	if err != nil {
		log.Printf("decodeGob: %s\n", err)
		return nil, nil, err
	}
	return session, gobHash.Sum(), nil
}

func (s sessionResponseWriter) WriteHeader(code int) {
	if atomic.AddInt32(&s.wroteHeader, 1) == 1 {
		origCookie, err := s.req.Cookie(s.h.CookieName)
		var origCookieVal string
		if err != nil {
			origCookieVal = ""
		} else {
			origCookieVal = origCookie.Value
		}

		session := s.h.RS.Get(s.req)
		if len(session) == 0 {
			// if we have an empty session, but the
			// request didn't start out that way, we
			// assume the user wants us to clear the
			// session
			if origCookieVal != "" {
				//log.Println("clearing cookie")
				var cookie http.Cookie
				cookie.Name = s.h.CookieName
				cookie.Value = ""
				cookie.Path = "/"
				// a cookie is expired by setting it
				// with an expiration time in the past
				cookie.Expires = *time.SecondsToUTC(0)
				http.SetCookie(s, &cookie)
			}
			goto write
		}
		encoded, gobHash, err := encodeCookie(session, s.h.encKey, s.h.hmacKey)
		if err != nil {
			log.Printf("createCookie: %s\n", err)
			goto write
		}

		if bytes.Equal(gobHash, s.h.RS.getHash(s.req)) {
			//log.Println("not re-setting identical cookie")
			goto write
		}

		var cookie http.Cookie
		cookie.Name = s.h.CookieName
		cookie.Value = encoded
		cookie.Path = s.h.CookiePath
		cookie.HttpOnly = s.h.RS.HttpOnly
		cookie.Secure = s.h.RS.Secure
		http.SetCookie(s, &cookie)
	}
write:
	s.ResponseWriter.WriteHeader(code)
}

func (s sessionResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hijacker, _ := s.ResponseWriter.(http.Hijacker)
	return hijacker.Hijack()
}

func (h *SessionHandler) getCookieSession(req *http.Request) (map[string]interface{}, []byte) {
	cookie, err := req.Cookie(h.CookieName)
	if err != nil {
		//log.Printf("getCookieSesh: '%#v' not found\n",
		//	h.CookieName)
		return map[string]interface{}{}, nil
	}
	session, gobHash, err := decodeCookie(cookie.Value, h.encKey, h.hmacKey)
	if err != nil {
		log.Printf("decodeCookie: %s\n", err)
		return map[string]interface{}{}, nil
	}

	return session, gobHash
}

func (h *SessionHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// get our session a little early, so that we can add our
	// authentication information to it if we get some
	session, gobHash := h.getCookieSession(req)

	h.RS.Set(req, session, gobHash)
	defer h.RS.Clear(req)

	sessionWriter := sessionResponseWriter{rw, h, req, 0}
	h.Handler.ServeHTTP(sessionWriter, req)
}

func NewSessionHandler(handler http.Handler, key string, rs *RequestSessions) *SessionHandler {
	// sha1 sums are 20 bytes long.  we use the first 16 bytes as
	// the aes key.
	encHash := sha1.New()
	encHash.Write([]byte(key))
	encHash.Write([]byte("-encryption"))
	hmacHash := sha1.New()
	hmacHash.Write([]byte(key))
	hmacHash.Write([]byte("-hmac"))

	// if the user hasn't specified a session handler, use the
	// package's default one
	if rs == nil {
		rs = Session
	}

	return &SessionHandler{
		Handler:    handler,
		CookieName: "session",
		CookiePath: "/",
		RS:         rs,
		encKey:     encHash.Sum()[:blockSize],
		hmacKey:    hmacHash.Sum()[:blockSize],
	}
}
<|endoftext|>