<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"
	"github.com/demizer/go-elog"
	"strings"
	"unicode"
	"unicode/utf8"
)

type Id int

func (i Id) IdNumber() Id { return i }

// The line number of an item in the input string
type Line int

func (l Line) LineNumber() Line { return l }

// The begining location of an item in the input.
type StartPosition int

func (s StartPosition) Position() StartPosition { return s }

// itemElement are the types that are emitted by the lexer.
type itemElement int

const (
	itemEOF itemElement = iota
	itemError
	itemTitle
	itemSectionAdornment
	itemParagraph
	itemBlockquote
	itemLiteralBlock
	itemSystemMessage
	itemSpace
	itemBlankLine
)

var elements = [...]string{
	"itemEOF",
	"itemError",
	"itemTitle",
	"itemSectionAdornment",
	"itemParagraph",
	"itemBlockquote",
	"itemLiteralBlock",
	"itemSystemMessage",
	"itemSpace",
	"itemBlankLine",
}

// String implements the Stringer interface for printing itemElement types.
func (t itemElement) String() string { return elements[t] }

func (t *itemElement) UnmarshalJSON(data []byte) error {
	for num, elm := range elements {
		if elm == string(data[1:len(data)-1]) {
			*t = itemElement(num)
		}
	}
	return nil
}

// Valid section adornment runes
var sectionAdornments = []rune{'!', '"', '#', '$', '\'', '%', '&', '(', ')', '*',
	'+', ',', '-', '.', '/', ':', ';', '<', '=', '>', '?', '@', '[', '\\',
	']', '^', '_', '`', '{', '|', '}', '~'}

// Emitted by the lexer on End of File
const eof rune = -1

// Function prototype for scanner functions
type stateFn func(*lexer) stateFn

// Struct for tokens emitted by the scanning process
type item struct {
	Id            `json:"id"`
	Type          itemElement `json:"type"`
	Text          interface{} `json:"text"`
	Line          `json:"line"`
	StartPosition `json:"startPosition"`
	Length        int `json:"length"`
}

// The lexer struct tracks the state of the lexer
type lexer struct {
	name             string    // The name of the current lexer
	input            string    // The input text
	state            stateFn   // The current state of the lexer
	index            int       // Position in input
	start            int       // The start of the current token
	width            int       // The width of the current position
	items            chan item // The channel items are emitted to
	lastItem         *item     // The last item emitted to the channel
	lastItemPosition StartPosition
	id               int // Unique id for each item emitted
}

// lex is the entry point of the lexer
func lex(name, input string) *lexer {
	log.Debugln("Start")
	if !norm.NFC.IsNormalString(input) {
		input = norm.NFC.String(input)
	}
	l := &lexer{
		name:  name,
		input: input,
		items: make(chan item),
	}
	go l.run()
	log.Debugln("End")
	return l
}

// run is the engine of the lexing process.
func (l *lexer) run() {
	log.Debugln("Start")
	for l.state = lexStart; l.state != nil; {
		l.state = l.state(l)
	}
	log.Debugln("End")
}

// emit passes an item back to the client.
func (l *lexer) emit(t itemElement) {
	if l.start == l.index && int(l.index) < len(l.input) {
		l.index += 1
	}
	log.Infof("#### %s: %q start: %d pos: %d line: %d\n", t,
		l.input[l.start:l.index], l.start, l.index, l.lineNumber())
	l.id++
	length := utf8.RuneCountInString(l.input[l.start:l.index])
	nItem := item{
		Id:            Id(l.id),
		Type:          t,
		Text:          l.input[l.start:l.index],
		Line:          Line(l.lineNumber()),
		StartPosition: StartPosition(l.start + 1),
		Length:        length,
	}
	l.items <- nItem
	l.lastItem = &nItem
	l.start = l.index
}

// backup backs up the lexer by one position using the width of the last rune retrieved from the
// input.
func (l *lexer) backup() {
	l.index -= l.width
}

// current returns the rune at the current position in the input.
func (l *lexer) current() rune {
	r, _ := utf8.DecodeRuneInString(l.input[l.index:])
	return r
}

// peek looks ahead in the input by one position and returns the rune.
func (l *lexer) peek() rune {
	r := l.next()
	l.backup()
	return r
}

// advance fast forwards the lexer to the next rune in the input specified by "to".
func (l *lexer) advance(to rune) {
	for {
		if l.next() == eof || to == l.current() {
			break
		}
	}
}

// next advances the position of the lexer by one rune and returns that rune.
func (l *lexer) next() rune {
	if int(l.index) >= len(l.input) {
		log.Debugln("Reached eof!")
		l.width = 0
		return eof
	}
	r, w := utf8.DecodeRuneInString(l.input[l.index:])
	l.width = w
	l.index += l.width
	if unicode.In(r, unicode.Diacritic) {
		log.Debugf("DIACRITIC char found: %#U width: %d\n", r, w)
	}
	log.Debugf("cur: %q start: %d pos: %d\n", r, l.start, l.index)
	return r
}

// nextItem returns the next item from the input.
func (l *lexer) nextItem() *item {
	item := <-l.items
	l.lastItemPosition = item.StartPosition
	return &item

}

func (l *lexer) lineNumber() int {
	return strings.Count(l.input[:l.index-1], "\n") + 1
}

func (l *lexer) isStartOfLine() bool {
	return (l.index - l.start - l.width) == 1
}

// isSpace reports whether r is a space character.
func isSpace(r rune) bool {
	return r == ' ' || r == '\t' || r == '\n' || r == '\r'
}

// isEndOfLine reports whether r is an end-of-line character.
func isEndOfLine(r rune) bool {
	return r == '\r' || r == '\n'
}

// isSection compares a number of positions (skipping whitespace) to
// determine if the runes are sectionAdornments and returns a true if the
// positions match each other. Rune comparison begins at the current lexer
// position. isSection returns false if there is a blank line between the
// positions or if there is a rune mismatch between positions.
func isSection(l *lexer) bool {
	log.Debugln("Start")
	var lookPositions = 2
	var newLineNum int
	var runePositions []rune
	var matchCount int
	cPos := l.index

	exit := func(value bool) bool {
		l.index = cPos
		log.Debugln("Returning", value)
		return value
	}

	log.Debugln("Looking ahead", lookPositions, "position(s) for sectionAdornments...")

	// Check two runes to see if they are section adornments, if not, we will check the next
	// line (after whitespace).
	if isSectionAdornment(l.current()) && isSectionAdornment(l.peek()) &&
		l.current() == l.peek() {
		return true
	}

	// Advance to the end of the line
	l.advance('\n')

	for j := 0; j < lookPositions; j++ {
		for isSpace(l.current()) {
			if isEndOfLine(l.current()) {
				if newLineNum == 1 {
					log.Debugln("Too many newlines!")
					return exit(false)
				}
				newLineNum += 1
				log.Debugln("newLineNum:", newLineNum)
			}
			l.next()
		}
		if isSectionAdornment(l.current()) {
			log.Debugf("Found adornment: \"%s\" pos: %d\n", string(l.current()), l.index)
			runePositions = append(runePositions, l.current())
			matchCount += 1
		}
		if l.next() == eof {
			return exit(false)
		}
	}

	if len(runePositions) == 0 || matchCount != lookPositions {
		return exit(false)
	}

	return exit(true)
}

// isSectionAdornment returns true if r matches a section adornment.
func isSectionAdornment(r rune) bool {
	for _, a := range sectionAdornments {
		if a == r {
			return true
		}
	}
	return false
}

// lexStart is the first stateFn called by run(). From here other stateFn's are called depending on
// the input. When this function returns nil, the lexing is finished and run() will exit.
func lexStart(l *lexer) stateFn {
	log.Debugln("Start")
	for {
		var tokenLength = l.index - l.start
		r := l.current()
		// log.Debugln("tokenLength:", tokenLength, l.width)
		if tokenLength <= l.width && l.width > 0 && !isEndOfLine(r) {
			log.Debugln("Start of new token")
			if isSection(l) {
				return lexSection
			} else if isSpace(r) {
				return lexSpace
			}
		} else if isEndOfLine(r) {
			log.Debugln("isEndOfLine == true")
			if l.index > l.start {
				l.emit(itemParagraph)
				l.start += 1 // Skip the new line
			} else if l.start == l.index {
				l.emit(itemBlankLine)
			}
		}
		if l.next() == eof {
			break
		}
	}

	// Correctly reached eof.
	l.emit(itemEOF)
	log.Debugln("End")
	return nil
}

// lexSpace consumes space characters (space and tab) in the input and emits a itemSpace token.
func lexSpace(l *lexer) stateFn {
	log.Debugln("Start")
	for isSpace(l.current()) {
		l.next()
	}
	if l.start < l.index {
		l.emit(itemSpace)
		l.next()
	}
	log.Debugln("End")
	return lexStart
}

// lexSection is used after isSection() has determined that the next runes of input are section.
// From here, the lexTitle() and lexSectionAdornment() are called based on the input.
func lexSection(l *lexer) stateFn {
	log.Debugln("Start")
	// The order of the case statements matter here
	switch r := l.next(); {
	case isSectionAdornment(r):
		if l.lastItem != nil && l.lastItem.Type != itemTitle {
			return lexSectionAdornment
		}
		lexSectionAdornment(l)
	case unicode.IsPrint(r):
		return lexTitle
	}
	log.Debugln("Exit")
	return lexStart
}

// lexTitle consumes input until newline and emits an itemTitle token. Once the token is emitted,
// control is returned to lexSection().
func lexTitle(l *lexer) stateFn {
	log.Debugln("Start")
	for {
		if isSpace(l.current()) && l.isStartOfLine() {
			return lexSpace
		}
		l.next()
		if l.peek() == '\n' {
			l.emit(itemTitle)
			l.start += 1
			l.index += 1
			break
		}
	}
	log.Debugln("End")
	return lexSection
}

// lexSectionAdornment advances the lexer until a newline is encountered and emits a
// itemSectionAdornment token. Control is returned to lexSection() on completion.
func lexSectionAdornment(l *lexer) stateFn {
	log.Debugln("Start")
	for {
		//TODO: Add adornment rune check
		l.next()
		if l.peek() == '\n' {
			l.emit(itemSectionAdornment)
			l.start += 1
			l.index += 1
			break
		}
	}
	log.Debugln("End")
	return lexSection
}
<commit_msg>lex.go: Add pos argument to lex.backup()<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"
	"github.com/demizer/go-elog"
	"strings"
	"unicode"
	"unicode/utf8"
)

type Id int

func (i Id) IdNumber() Id { return i }

// The line number of an item in the input string
type Line int

func (l Line) LineNumber() Line { return l }

// The begining location of an item in the input.
type StartPosition int

func (s StartPosition) Position() StartPosition { return s }

// itemElement are the types that are emitted by the lexer.
type itemElement int

const (
	itemEOF itemElement = iota
	itemError
	itemTitle
	itemSectionAdornment
	itemParagraph
	itemBlockquote
	itemLiteralBlock
	itemSystemMessage
	itemSpace
	itemBlankLine
)

var elements = [...]string{
	"itemEOF",
	"itemError",
	"itemTitle",
	"itemSectionAdornment",
	"itemParagraph",
	"itemBlockquote",
	"itemLiteralBlock",
	"itemSystemMessage",
	"itemSpace",
	"itemBlankLine",
}

// String implements the Stringer interface for printing itemElement types.
func (t itemElement) String() string { return elements[t] }

func (t *itemElement) UnmarshalJSON(data []byte) error {
	for num, elm := range elements {
		if elm == string(data[1:len(data)-1]) {
			*t = itemElement(num)
		}
	}
	return nil
}

// Valid section adornment runes
var sectionAdornments = []rune{'!', '"', '#', '$', '\'', '%', '&', '(', ')', '*',
	'+', ',', '-', '.', '/', ':', ';', '<', '=', '>', '?', '@', '[', '\\',
	']', '^', '_', '`', '{', '|', '}', '~'}

// Emitted by the lexer on End of File
const eof rune = -1

// Function prototype for scanner functions
type stateFn func(*lexer) stateFn

// Struct for tokens emitted by the scanning process
type item struct {
	Id            `json:"id"`
	Type          itemElement `json:"type"`
	Text          interface{} `json:"text"`
	Line          `json:"line"`
	StartPosition `json:"startPosition"`
	Length        int `json:"length"`
}

// The lexer struct tracks the state of the lexer
type lexer struct {
	name             string    // The name of the current lexer
	input            string    // The input text
	state            stateFn   // The current state of the lexer
	index            int       // Position in input
	start            int       // The start of the current token
	width            int       // The width of the current position
	items            chan item // The channel items are emitted to
	lastItem         *item     // The last item emitted to the channel
	lastItemPosition StartPosition
	id               int // Unique id for each item emitted
}

// lex is the entry point of the lexer
func lex(name, input string) *lexer {
	log.Debugln("Start")
	if !norm.NFC.IsNormalString(input) {
		input = norm.NFC.String(input)
	}
	l := &lexer{
		name:  name,
		input: input,
		items: make(chan item),
	}
	go l.run()
	log.Debugln("End")
	return l
}

// run is the engine of the lexing process.
func (l *lexer) run() {
	log.Debugln("Start")
	for l.state = lexStart; l.state != nil; {
		l.state = l.state(l)
	}
	log.Debugln("End")
}

// emit passes an item back to the client.
func (l *lexer) emit(t itemElement) {
	if l.start == l.index && int(l.index) < len(l.input) {
		l.index += 1
	}
	log.Infof("#### %s: %q start: %d pos: %d line: %d\n", t,
		l.input[l.start:l.index], l.start, l.index, l.lineNumber())
	l.id++
	length := utf8.RuneCountInString(l.input[l.start:l.index])
	nItem := item{
		Id:            Id(l.id),
		Type:          t,
		Text:          l.input[l.start:l.index],
		Line:          Line(l.lineNumber()),
		StartPosition: StartPosition(l.start + 1),
		Length:        length,
	}
	l.items <- nItem
	l.lastItem = &nItem
	l.start = l.index
}

// backup backs up the lexer by one position using the width of the last rune retrieved from the
// input.
func (l *lexer) backup(pos int) {
	for i := 0; i < pos; i++ {
		l.index -= l.width
	}
}

// current returns the rune at the current position in the input.
func (l *lexer) current() rune {
	r, _ := utf8.DecodeRuneInString(l.input[l.index:])
	return r
}

// peek looks ahead in the input by one position and returns the rune.
func (l *lexer) peek() rune {
	r := l.next()
	l.backup(1)
	return r
}

// advance fast forwards the lexer to the next rune in the input specified by "to".
func (l *lexer) advance(to rune) {
	for {
		if l.next() == eof || to == l.current() {
			break
		}
	}
}

// next advances the position of the lexer by one rune and returns that rune.
func (l *lexer) next() rune {
	if int(l.index) >= len(l.input) {
		log.Debugln("Reached eof!")
		l.width = 0
		return eof
	}
	r, w := utf8.DecodeRuneInString(l.input[l.index:])
	l.width = w
	l.index += l.width
	if unicode.In(r, unicode.Diacritic) {
		log.Debugf("DIACRITIC char found: %#U width: %d\n", r, w)
	}
	log.Debugf("cur: %q start: %d pos: %d\n", r, l.start, l.index)
	return r
}

// nextItem returns the next item from the input.
func (l *lexer) nextItem() *item {
	item := <-l.items
	l.lastItemPosition = item.StartPosition
	return &item

}

func (l *lexer) lineNumber() int {
	return strings.Count(l.input[:l.index-1], "\n") + 1
}

func (l *lexer) isStartOfLine() bool {
	return (l.index - l.start - l.width) == 1
}

// isSpace reports whether r is a space character.
func isSpace(r rune) bool {
	return r == ' ' || r == '\t' || r == '\n' || r == '\r'
}

// isEndOfLine reports whether r is an end-of-line character.
func isEndOfLine(r rune) bool {
	return r == '\r' || r == '\n'
}

// isSection compares a number of positions (skipping whitespace) to
// determine if the runes are sectionAdornments and returns a true if the
// positions match each other. Rune comparison begins at the current lexer
// position. isSection returns false if there is a blank line between the
// positions or if there is a rune mismatch between positions.
func isSection(l *lexer) bool {
	log.Debugln("Start")
	var lookPositions = 2
	var newLineNum int
	var runePositions []rune
	var matchCount int
	cPos := l.index

	exit := func(value bool) bool {
		l.index = cPos
		log.Debugln("Returning", value)
		return value
	}

	log.Debugln("Looking ahead", lookPositions, "position(s) for sectionAdornments...")

	// Check two runes to see if they are section adornments, if not, we will check the next
	// line (after whitespace).
	if isSectionAdornment(l.current()) && isSectionAdornment(l.peek()) &&
		l.current() == l.peek() {
		return true
	}

	// Advance to the end of the line
	l.advance('\n')

	for j := 0; j < lookPositions; j++ {
		for isSpace(l.current()) {
			if isEndOfLine(l.current()) {
				if newLineNum == 1 {
					log.Debugln("Too many newlines!")
					return exit(false)
				}
				newLineNum += 1
				log.Debugln("newLineNum:", newLineNum)
			}
			l.next()
		}
		if isSectionAdornment(l.current()) {
			log.Debugf("Found adornment: \"%s\" pos: %d\n", string(l.current()), l.index)
			runePositions = append(runePositions, l.current())
			matchCount += 1
		}
		if l.next() == eof {
			return exit(false)
		}
	}

	if len(runePositions) == 0 || matchCount != lookPositions {
		return exit(false)
	}

	return exit(true)
}

// isSectionAdornment returns true if r matches a section adornment.
func isSectionAdornment(r rune) bool {
	for _, a := range sectionAdornments {
		if a == r {
			return true
		}
	}
	return false
}

// lexStart is the first stateFn called by run(). From here other stateFn's are called depending on
// the input. When this function returns nil, the lexing is finished and run() will exit.
func lexStart(l *lexer) stateFn {
	log.Debugln("Start")
	for {
		var tokenLength = l.index - l.start
		r := l.current()
		// log.Debugln("tokenLength:", tokenLength, l.width)
		if tokenLength <= l.width && l.width > 0 && !isEndOfLine(r) {
			log.Debugln("Start of new token")
			if isSection(l) {
				return lexSection
			} else if isSpace(r) {
				return lexSpace
			}
		} else if isEndOfLine(r) {
			log.Debugln("isEndOfLine == true")
			if l.index > l.start {
				l.emit(itemParagraph)
				l.start += 1 // Skip the new line
			} else if l.start == l.index {
				l.emit(itemBlankLine)
			}
		}
		if l.next() == eof {
			break
		}
	}

	// Correctly reached eof.
	l.emit(itemEOF)
	log.Debugln("End")
	return nil
}

// lexSpace consumes space characters (space and tab) in the input and emits a itemSpace token.
func lexSpace(l *lexer) stateFn {
	log.Debugln("Start")
	for isSpace(l.current()) {
		l.next()
	}
	if l.start < l.index {
		l.emit(itemSpace)
		l.next()
	}
	log.Debugln("End")
	return lexStart
}

// lexSection is used after isSection() has determined that the next runes of input are section.
// From here, the lexTitle() and lexSectionAdornment() are called based on the input.
func lexSection(l *lexer) stateFn {
	log.Debugln("Start")
	// The order of the case statements matter here
	switch r := l.next(); {
	case isSectionAdornment(r):
		if l.lastItem != nil && l.lastItem.Type != itemTitle {
			return lexSectionAdornment
		}
		lexSectionAdornment(l)
	case unicode.IsPrint(r):
		return lexTitle
	}
	log.Debugln("Exit")
	return lexStart
}

// lexTitle consumes input until newline and emits an itemTitle token. Once the token is emitted,
// control is returned to lexSection().
func lexTitle(l *lexer) stateFn {
	log.Debugln("Start")
	l.backup(2) // Start from the newline of the previous line
	for {
		if isSpace(l.current()) && l.isStartOfLine() {
			return lexSpace
		}
		l.next()
		if l.peek() == '\n' {
			l.emit(itemTitle)
			l.start += 1
			l.index += 1
			break
		}
	}
	log.Debugln("End")
	return lexSection
}

// lexSectionAdornment advances the lexer until a newline is encountered and emits a
// itemSectionAdornment token. Control is returned to lexSection() on completion.
func lexSectionAdornment(l *lexer) stateFn {
	log.Debugln("Start")
	for {
		//TODO: Add adornment rune check
		l.next()
		if l.peek() == '\n' {
			l.emit(itemSectionAdornment)
			l.start += 1
			l.index += 1
			break
		}
	}
	log.Debugln("End")
	return lexSection
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/sha1"
	"encoding/base64"
	"encoding/json"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path"
	"regexp"
	"time"

	"github.com/cespare/blackfriday"
	"github.com/cespare/go-apachelog"
	"github.com/gorilla/pat"
)

const (
	listenAddr                 = "localhost:8389"
	staticDir                  = "public"
	pastieDir                  = "files"
	mainPastie                 = "about.markdown"
	markdownReferencePastie    = "reference.markdown"
	pygmentize                 = "./vendor/pygments/pygmentize"
	viewFile                   = "view.html"
	expirationTimeHours        = 7 * 24
	expirationCheckPeriodHours = 1
)

var (
	validLanguages     = make(map[string]struct{})
	markdownRenderer   *blackfriday.Html
	markdownExtensions int
	viewHtml           []byte
	filenameRegex      = regexp.MustCompile(`^[\w\-]{27}\.\w+$`)
)

func init() {
	var err error

	// Get the list of valid lexers from pygments.
	rawLexerList, err := exec.Command(pygmentize, "-L", "lexers").Output()
	if err != nil {
		log.Fatalln(err)
	}
	for _, line := range bytes.Split(rawLexerList, []byte("\n")) {
		if len(line) == 0 || line[0] != '*' {
			continue
		}
		for _, l := range bytes.Split(bytes.Trim(line, "* :"), []byte(",")) {
			lexer := string(bytes.TrimSpace(l))
			if len(lexer) != 0 {
				validLanguages[lexer] = struct{}{}
			}
		}
	}

	// Set up the renderer.
	flags := 0
	flags |= blackfriday.HTML_GITHUB_BLOCKCODE
	markdownRenderer = blackfriday.HtmlRenderer(flags, "", "")
	markdownRenderer.SetBlockCodeProcessor(syntaxHighlight)

	markdownExtensions = 0
	markdownExtensions |= blackfriday.EXTENSION_FENCED_CODE
	markdownExtensions |= blackfriday.EXTENSION_TABLES
	markdownExtensions |= blackfriday.EXTENSION_NO_INTRA_EMPHASIS
	markdownExtensions |= blackfriday.EXTENSION_SPACE_HEADERS
	markdownExtensions |= blackfriday.EXTENSION_AUTOLINK

	// Check that the main info file exists.
	_, err = os.Stat(pastieDir + "/" + mainPastie)
	if err != nil {
		log.Fatalln("Error with main info file: " + err.Error())
	}

	// Load in the main view template
	viewTemplate, err := template.ParseFiles(viewFile)
	if err != nil {
		log.Fatalln(err)
	}
	b := new(bytes.Buffer)
	templateData := struct{ MainId, MarkdownReferenceId string }{mainPastie, markdownReferencePastie}
	err = viewTemplate.Execute(b, templateData)
	if err != nil {
		log.Fatalln(err)
	}
	viewHtml = b.Bytes()
}

func syntaxHighlight(out io.Writer, in io.Reader, language string) {
	_, ok := validLanguages[language]
	if !ok || language == "" {
		language = "text"
	}
	pygmentsCommand := exec.Command(pygmentize, "-l", language, "-f", "html")
	pygmentsCommand.Stdin = in
	pygmentsCommand.Stdout = out
	pygmentsCommand.Run()
}

type Pastie struct {
	Text   string `json:"text"`
	Format string `json:"format"`
}

func render(text []byte, format string) []byte {
	var rendered []byte
	switch format {
	case "text":
		rendered = text
	case "markdown":
		rendered = blackfriday.Markdown(text, markdownRenderer, markdownExtensions)
	default:
		var highlighted bytes.Buffer
		in := bytes.NewBuffer(text)
		syntaxHighlight(&highlighted, in, format)
		rendered = highlighted.Bytes()
	}
	return rendered
}

// Update the mtime
func touch(path string) { os.Chtimes(path, time.Now(), time.Now()) }

func pastieHandler(w http.ResponseWriter, r *http.Request) {
	time.Sleep(0 * time.Second)
	id := r.URL.Query().Get(":id")
	if len(id) == 0 {
		http.Error(w, "No such file.", http.StatusInternalServerError)
		return
	}
	var filename string
	// If the filename is one made by Pastedown, then look in the directory structure we expect; otherwise, just
	// try to find such a file directly.
	if filenameRegex.MatchString(id) {
		filename = path.Join(pastieDir, id[:2], id[2:])
	} else {
		filename = path.Join(pastieDir, id)
	}
	touch(filename)
	contents, err := ioutil.ReadFile(filename)
	if err != nil {
		http.Error(w, "No such file.", http.StatusNotFound)
		return
	}

	// Just return the raw contents if ?rendered=true wasn't set.
	if r.URL.Query().Get("rendered") != "true" {
		w.Write(contents)
		return
	}

	// Otherwise, render the proper format according to the extension.
	extension := path.Ext(id)
	if extension == "" {
		extension = "text"
	} else {
		extension = extension[1:]
	}

	w.Write(render(contents, extension))
}

func decodePastie(r *http.Request) (*Pastie, error) {
	text, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return nil, err
	}
	pastie := &Pastie{}
	err = json.Unmarshal(text, pastie)
	if err != nil {
		return nil, err
	}
	return pastie, nil
}

func previewHandler(w http.ResponseWriter, r *http.Request) {
	preview, err := decodePastie(r)
	if err != nil {
		http.Error(w, "Could not render preview text.", http.StatusInternalServerError)
		return
	}
	w.Write(render([]byte(preview.Text), preview.Format))
}

func saveHandler(w http.ResponseWriter, r *http.Request) {
	preview, err := decodePastie(r)
	bytes := []byte(preview.Text)
	if err != nil {
		log.Println("Error decoding pastie for saving: " + err.Error())
		http.Error(w, "Could not save text.", http.StatusInternalServerError)
		return
	}
	sha := sha1.New()
	sha.Write(bytes)
	hash := base64.URLEncoding.EncodeToString(sha.Sum(nil))

	// The filename is constructed in the following manner:
	//
	// - Chop off the last '=' (padding character) of the hash -- all the shas are the same length anyway so
	//   we might as well get rid of the character that they all have in common.
	// - Chop the first two characters off the front of the hash and use this as the directory to limit the
	//	 number files in a single directory (git uses this trick for its object store).
	// - The full file format name is used as the extension.
	//
	// So for example:
	// { sha: jBEtyBOnX_M2rp7DNp3mQskWqwg=, filetype: markdown } => jB/EtyBOnX_M2rp7DNp3mQskWqwg.markdown
	directory := path.Join(pastieDir, hash[0:2])
	logicalName := hash[:len(hash)-1] + "." + preview.Format
	filename := path.Join(directory, hash[2:len(hash)-1]+"."+preview.Format)
	err = os.MkdirAll(directory, 0771)
	if err != nil {
		log.Println("Error creating new directory: " + err.Error())
		http.Error(w, "Could not save text.", http.StatusInternalServerError)
		return
	}
	_, err = os.Stat(filename)
	if err == nil {
		log.Println("Request for existing pastie: " + filename)
		touch(filename)
	} else {
		log.Println("Saving new pastie: " + filename)
		err = ioutil.WriteFile(filename, bytes, 0666)
		if err != nil {
			log.Println("Error writing pastie file: " + err.Error())
			http.Error(w, "Could not save text.", http.StatusInternalServerError)
			return
		}
	}
	// Otherwise file already exists
	w.Write([]byte(logicalName))
}

func viewHandler(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.Redirect(w, r, "/", http.StatusFound)
		return
	}
	w.Write(viewHtml)
}

func expire() {
	log.Println("Expiring files...")
	dirs, err := ioutil.ReadDir(pastieDir)
	if err != nil {
		log.Println("Error reading file directory: " + err.Error())
		return
	}
	expiredFiles := 0
	expiredDirs := 0
	for _, d := range dirs {
		if !d.IsDir() {
			continue
		}
		dirPath := path.Join(pastieDir, d.Name())
		files, err := ioutil.ReadDir(dirPath)
		if err != nil {
			log.Println("Error reading file directory: " + err.Error())
			continue
		}
		unexpired := false
		for _, f := range files {
			if time.Now().Sub(f.ModTime()).Hours() <= expirationTimeHours {
				unexpired = true
				continue
			}
			filepath := path.Join(dirPath, f.Name())
			if err := os.Remove(filepath); err != nil {
				log.Println("Error deleting file: " + err.Error())
				continue
			}
			expiredFiles++
		}
		if !unexpired {
			if err := os.Remove(dirPath); err != nil {
				log.Println("Error deleting empty directory: " + err.Error())
				continue
			}
			expiredDirs++
		}
	}
	log.Printf("Removed %d expired files and %d empty directories.\n", expiredFiles, expiredDirs)
}

func main() {
	// Start the background file deleter going
	go func() {
		ticker := time.NewTicker(expirationCheckPeriodHours * time.Hour)
		for _ = range ticker.C {
			expire()
		}
	}()

	// Set up the server
	mux := pat.New()

	mux.Add("GET", "/favicon.ico", http.FileServer(http.Dir(staticDir)))
	staticPath := "/" + staticDir + "/"
	mux.Add("GET", staticPath, http.StripPrefix(staticPath, http.FileServer(http.Dir("./"+staticDir))))

	mux.Get(`/files/{id:[\w\.\-]+}`, pastieHandler)
	mux.Post("/preview", previewHandler)
	mux.Put("/file", saveHandler)
	mux.Get("/", viewHandler)

	handler := apachelog.NewHandler(mux, os.Stderr)
	server := &http.Server{
		Addr:    listenAddr,
		Handler: handler,
	}
	log.Println("Now listening on", listenAddr)
	log.Fatalf(server.ListenAndServe().Error())
}
<commit_msg>Add flags for configurable options.<commit_after>package main

import (
	"bytes"
	"crypto/sha1"
	"encoding/base64"
	"encoding/json"
	"flag"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path"
	"regexp"
	"time"

	"github.com/cespare/blackfriday"
	"github.com/cespare/go-apachelog"
	"github.com/gorilla/pat"
)

const (
	staticDir                  = "public"
	pygmentize                 = "./vendor/pygments/pygmentize"
	viewFile                   = "view.html"
	expirationCheckPeriodHours = 1
)

// User-configurable values
var (
	listenAddr              string
	pastieDir               string
	mainPastie              string
	markdownReferencePastie string
	expirationTimeHours     int
)

var (
	validLanguages     = make(map[string]struct{})
	markdownRenderer   *blackfriday.Html
	markdownExtensions int
	viewHtml           []byte
	filenameRegex      = regexp.MustCompile(`^[\w\-]{27}\.\w+$`)
)

func init() {
	var err error

	// Set up flags
	flag.StringVar(&listenAddr, "listenaddr", "localhost:8389", "The server address on which to listen")
	flag.StringVar(&pastieDir, "storagedir", "files", "The directory in which to store documents")
	flag.StringVar(&mainPastie, "mainpage", "about.markdown", "The document to display on the front page")
	flag.StringVar(&markdownReferencePastie, "referencepage", "reference.markdown",
		"The document to display at the 'markdown reference' link")
	flag.IntVar(&expirationTimeHours, "expirationhours", 7*24,
		"How long to keep documents before deleting them")

	// Get the list of valid lexers from pygments.
	rawLexerList, err := exec.Command(pygmentize, "-L", "lexers").Output()
	if err != nil {
		log.Fatalln(err)
	}
	for _, line := range bytes.Split(rawLexerList, []byte("\n")) {
		if len(line) == 0 || line[0] != '*' {
			continue
		}
		for _, l := range bytes.Split(bytes.Trim(line, "* :"), []byte(",")) {
			lexer := string(bytes.TrimSpace(l))
			if len(lexer) != 0 {
				validLanguages[lexer] = struct{}{}
			}
		}
	}

	// Set up the renderer.
	flags := 0
	flags |= blackfriday.HTML_GITHUB_BLOCKCODE
	markdownRenderer = blackfriday.HtmlRenderer(flags, "", "")
	markdownRenderer.SetBlockCodeProcessor(syntaxHighlight)

	markdownExtensions = 0
	markdownExtensions |= blackfriday.EXTENSION_FENCED_CODE
	markdownExtensions |= blackfriday.EXTENSION_TABLES
	markdownExtensions |= blackfriday.EXTENSION_NO_INTRA_EMPHASIS
	markdownExtensions |= blackfriday.EXTENSION_SPACE_HEADERS
	markdownExtensions |= blackfriday.EXTENSION_AUTOLINK

	// Check that the main info file exists.
	_, err = os.Stat(pastieDir + "/" + mainPastie)
	if err != nil {
		log.Fatalln("Error with main info file: " + err.Error())
	}

	// Load in the main view template
	viewTemplate, err := template.ParseFiles(viewFile)
	if err != nil {
		log.Fatalln(err)
	}
	b := new(bytes.Buffer)
	templateData := struct{ MainId, MarkdownReferenceId string }{mainPastie, markdownReferencePastie}
	err = viewTemplate.Execute(b, templateData)
	if err != nil {
		log.Fatalln(err)
	}
	viewHtml = b.Bytes()
}

func syntaxHighlight(out io.Writer, in io.Reader, language string) {
	_, ok := validLanguages[language]
	if !ok || language == "" {
		language = "text"
	}
	pygmentsCommand := exec.Command(pygmentize, "-l", language, "-f", "html")
	pygmentsCommand.Stdin = in
	pygmentsCommand.Stdout = out
	pygmentsCommand.Run()
}

type Pastie struct {
	Text   string `json:"text"`
	Format string `json:"format"`
}

func render(text []byte, format string) []byte {
	var rendered []byte
	switch format {
	case "text":
		rendered = text
	case "markdown":
		rendered = blackfriday.Markdown(text, markdownRenderer, markdownExtensions)
	default:
		var highlighted bytes.Buffer
		in := bytes.NewBuffer(text)
		syntaxHighlight(&highlighted, in, format)
		rendered = highlighted.Bytes()
	}
	return rendered
}

// Update the mtime
func touch(path string) { os.Chtimes(path, time.Now(), time.Now()) }

func pastieHandler(w http.ResponseWriter, r *http.Request) {
	time.Sleep(0 * time.Second)
	id := r.URL.Query().Get(":id")
	if len(id) == 0 {
		http.Error(w, "No such file.", http.StatusInternalServerError)
		return
	}
	var filename string
	// If the filename is one made by Pastedown, then look in the directory structure we expect; otherwise, just
	// try to find such a file directly.
	if filenameRegex.MatchString(id) {
		filename = path.Join(pastieDir, id[:2], id[2:])
	} else {
		filename = path.Join(pastieDir, id)
	}
	touch(filename)
	contents, err := ioutil.ReadFile(filename)
	if err != nil {
		http.Error(w, "No such file.", http.StatusNotFound)
		return
	}

	// Just return the raw contents if ?rendered=true wasn't set.
	if r.URL.Query().Get("rendered") != "true" {
		w.Write(contents)
		return
	}

	// Otherwise, render the proper format according to the extension.
	extension := path.Ext(id)
	if extension == "" {
		extension = "text"
	} else {
		extension = extension[1:]
	}

	w.Write(render(contents, extension))
}

func decodePastie(r *http.Request) (*Pastie, error) {
	text, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return nil, err
	}
	pastie := &Pastie{}
	err = json.Unmarshal(text, pastie)
	if err != nil {
		return nil, err
	}
	return pastie, nil
}

func previewHandler(w http.ResponseWriter, r *http.Request) {
	preview, err := decodePastie(r)
	if err != nil {
		http.Error(w, "Could not render preview text.", http.StatusInternalServerError)
		return
	}
	w.Write(render([]byte(preview.Text), preview.Format))
}

func saveHandler(w http.ResponseWriter, r *http.Request) {
	preview, err := decodePastie(r)
	bytes := []byte(preview.Text)
	if err != nil {
		log.Println("Error decoding pastie for saving: " + err.Error())
		http.Error(w, "Could not save text.", http.StatusInternalServerError)
		return
	}
	sha := sha1.New()
	sha.Write(bytes)
	hash := base64.URLEncoding.EncodeToString(sha.Sum(nil))

	// The filename is constructed in the following manner:
	//
	// - Chop off the last '=' (padding character) of the hash -- all the shas are the same length anyway so
	//   we might as well get rid of the character that they all have in common.
	// - Chop the first two characters off the front of the hash and use this as the directory to limit the
	//	 number files in a single directory (git uses this trick for its object store).
	// - The full file format name is used as the extension.
	//
	// So for example:
	// { sha: jBEtyBOnX_M2rp7DNp3mQskWqwg=, filetype: markdown } => jB/EtyBOnX_M2rp7DNp3mQskWqwg.markdown
	directory := path.Join(pastieDir, hash[0:2])
	logicalName := hash[:len(hash)-1] + "." + preview.Format
	filename := path.Join(directory, hash[2:len(hash)-1]+"."+preview.Format)
	err = os.MkdirAll(directory, 0771)
	if err != nil {
		log.Println("Error creating new directory: " + err.Error())
		http.Error(w, "Could not save text.", http.StatusInternalServerError)
		return
	}
	_, err = os.Stat(filename)
	if err == nil {
		log.Println("Request for existing pastie: " + filename)
		touch(filename)
	} else {
		log.Println("Saving new pastie: " + filename)
		err = ioutil.WriteFile(filename, bytes, 0666)
		if err != nil {
			log.Println("Error writing pastie file: " + err.Error())
			http.Error(w, "Could not save text.", http.StatusInternalServerError)
			return
		}
	}
	// Otherwise file already exists
	w.Write([]byte(logicalName))
}

func viewHandler(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.Redirect(w, r, "/", http.StatusFound)
		return
	}
	w.Write(viewHtml)
}

func expire() {
	log.Println("Expiring files...")
	dirs, err := ioutil.ReadDir(pastieDir)
	if err != nil {
		log.Println("Error reading file directory: " + err.Error())
		return
	}
	expiredFiles := 0
	expiredDirs := 0
	for _, d := range dirs {
		if !d.IsDir() {
			continue
		}
		dirPath := path.Join(pastieDir, d.Name())
		files, err := ioutil.ReadDir(dirPath)
		if err != nil {
			log.Println("Error reading file directory: " + err.Error())
			continue
		}
		unexpired := false
		for _, f := range files {
			if time.Now().Sub(f.ModTime()).Hours() <= float64(expirationTimeHours) {
				unexpired = true
				continue
			}
			filepath := path.Join(dirPath, f.Name())
			if err := os.Remove(filepath); err != nil {
				log.Println("Error deleting file: " + err.Error())
				continue
			}
			expiredFiles++
		}
		if !unexpired {
			if err := os.Remove(dirPath); err != nil {
				log.Println("Error deleting empty directory: " + err.Error())
				continue
			}
			expiredDirs++
		}
	}
	log.Printf("Removed %d expired files and %d empty directories.\n", expiredFiles, expiredDirs)
}

func main() {
	flag.Parse()

	// Start the background file deleter going
	go func() {
		ticker := time.NewTicker(expirationCheckPeriodHours * time.Hour)
		for _ = range ticker.C {
			expire()
		}
	}()

	// Set up the server
	mux := pat.New()

	mux.Add("GET", "/favicon.ico", http.FileServer(http.Dir(staticDir)))
	staticPath := "/" + staticDir + "/"
	mux.Add("GET", staticPath, http.StripPrefix(staticPath, http.FileServer(http.Dir("./"+staticDir))))

	mux.Get(`/files/{id:[\w\.\-]+}`, pastieHandler)
	mux.Post("/preview", previewHandler)
	mux.Put("/file", saveHandler)
	mux.Get("/", viewHandler)

	handler := apachelog.NewHandler(mux, os.Stderr)
	server := &http.Server{
		Addr:    listenAddr,
		Handler: handler,
	}
	log.Println("Now listening on", listenAddr)
	log.Fatalf(server.ListenAndServe().Error())
}
<|endoftext|>


<commit_before>package restarts

import (
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/plugins/drivers"
)

func testPolicy(success bool, mode string) *structs.RestartPolicy {
	return &structs.RestartPolicy{
		Interval: 2 * time.Minute,
		Delay:    1 * time.Second,
		Attempts: 3,
		Mode:     mode,
	}
}

// withinJitter is a helper that returns whether the returned delay is within
// the jitter.
func withinJitter(expected, actual time.Duration) bool {
	return float64((actual.Nanoseconds()-expected.Nanoseconds())/
		expected.Nanoseconds()) <= jitter
}

func testExitResult(exit int) *drivers.ExitResult {
	return &drivers.ExitResult{
		ExitCode: exit,
	}
}

func TestClient_RestartTracker_ModeDelay(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeDelay)
	rt := NewRestartTracker(p, structs.JobTypeService, nil)
	for i := 0; i < p.Attempts; i++ {
		state, when := rt.SetExitResult(testExitResult(127)).GetState()
		if state != structs.TaskRestarting {
			t.Fatalf("NextRestart() returned %v, want %v", state, structs.TaskRestarting)
		}
		if !withinJitter(p.Delay, when) {
			t.Fatalf("NextRestart() returned %v; want %v+jitter", when, p.Delay)
		}
	}

	// Follow up restarts should cause delay.
	for i := 0; i < 3; i++ {
		state, when := rt.SetExitResult(testExitResult(127)).GetState()
		if state != structs.TaskRestarting {
			t.Fail()
		}
		if !(when > p.Delay && when <= p.Interval) {
			t.Fatalf("NextRestart() returned %v; want > %v and <= %v", when, p.Delay, p.Interval)
		}
	}
}

func TestClient_RestartTracker_ModeFail(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeFail)
	rt := NewRestartTracker(p, structs.JobTypeSystem, nil)
	for i := 0; i < p.Attempts; i++ {
		state, when := rt.SetExitResult(testExitResult(127)).GetState()
		if state != structs.TaskRestarting {
			t.Fatalf("NextRestart() returned %v, want %v", state, structs.TaskRestarting)
		}
		if !withinJitter(p.Delay, when) {
			t.Fatalf("NextRestart() returned %v; want %v+jitter", when, p.Delay)
		}
	}

	// Next restart should cause fail
	if state, _ := rt.SetExitResult(testExitResult(127)).GetState(); state != structs.TaskNotRestarting {
		t.Fatalf("NextRestart() returned %v; want %v", state, structs.TaskNotRestarting)
	}
}

func TestClient_RestartTracker_NoRestartOnSuccess(t *testing.T) {
	t.Parallel()
	p := testPolicy(false, structs.RestartPolicyModeDelay)
	rt := NewRestartTracker(p, structs.JobTypeBatch, nil)
	if state, _ := rt.SetExitResult(testExitResult(0)).GetState(); state != structs.TaskTerminated {
		t.Fatalf("NextRestart() returned %v, expected: %v", state, structs.TaskTerminated)
	}
}

func TestClient_RestartTracker_ZeroAttempts(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeFail)
	p.Attempts = 0

	// Test with a non-zero exit code
	rt := NewRestartTracker(p, structs.JobTypeService, nil)
	if state, when := rt.SetExitResult(testExitResult(1)).GetState(); state != structs.TaskNotRestarting {
		t.Fatalf("expect no restart, got restart/delay: %v/%v", state, when)
	}

	// Even with a zero (successful) exit code non-batch jobs should exit
	// with TaskNotRestarting
	rt = NewRestartTracker(p, structs.JobTypeService, nil)
	if state, when := rt.SetExitResult(testExitResult(0)).GetState(); state != structs.TaskNotRestarting {
		t.Fatalf("expect no restart, got restart/delay: %v/%v", state, when)
	}

	// Batch jobs with a zero exit code and 0 attempts *do* exit cleanly
	// with Terminated
	rt = NewRestartTracker(p, structs.JobTypeBatch, nil)
	if state, when := rt.SetExitResult(testExitResult(0)).GetState(); state != structs.TaskTerminated {
		t.Fatalf("expect terminated, got restart/delay: %v/%v", state, when)
	}

	// Batch jobs with a non-zero exit code and 0 attempts exit with
	// TaskNotRestarting
	rt = NewRestartTracker(p, structs.JobTypeBatch, nil)
	if state, when := rt.SetExitResult(testExitResult(1)).GetState(); state != structs.TaskNotRestarting {
		t.Fatalf("expect no restart, got restart/delay: %v/%v", state, when)
	}
}

func TestClient_RestartTracker_TaskKilled(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeFail)
	p.Attempts = 0
	rt := NewRestartTracker(p, structs.JobTypeService, nil)
	if state, when := rt.SetKilled().GetState(); state != structs.TaskKilled && when != 0 {
		t.Fatalf("expect no restart; got %v %v", state, when)
	}
}

func TestClient_RestartTracker_RestartTriggered(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeFail)
	p.Attempts = 0
	rt := NewRestartTracker(p, structs.JobTypeService, nil)
	if state, when := rt.SetRestartTriggered(false).GetState(); state != structs.TaskRestarting && when != 0 {
		t.Fatalf("expect restart immediately, got %v %v", state, when)
	}
}

func TestClient_RestartTracker_RestartTriggered_Failure(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeFail)
	p.Attempts = 1
	rt := NewRestartTracker(p, structs.JobTypeService, nil)
	if state, when := rt.SetRestartTriggered(true).GetState(); state != structs.TaskRestarting || when == 0 {
		t.Fatalf("expect restart got %v %v", state, when)
	}
	if state, when := rt.SetRestartTriggered(true).GetState(); state != structs.TaskNotRestarting || when != 0 {
		t.Fatalf("expect failed got %v %v", state, when)
	}
}

func TestClient_RestartTracker_StartError_Recoverable_Fail(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeFail)
	rt := NewRestartTracker(p, structs.JobTypeSystem, nil)
	recErr := structs.NewRecoverableError(fmt.Errorf("foo"), true)
	for i := 0; i < p.Attempts; i++ {
		state, when := rt.SetStartError(recErr).GetState()
		if state != structs.TaskRestarting {
			t.Fatalf("NextRestart() returned %v, want %v", state, structs.TaskRestarting)
		}
		if !withinJitter(p.Delay, when) {
			t.Fatalf("NextRestart() returned %v; want %v+jitter", when, p.Delay)
		}
	}

	// Next restart should cause fail
	if state, _ := rt.SetStartError(recErr).GetState(); state != structs.TaskNotRestarting {
		t.Fatalf("NextRestart() returned %v; want %v", state, structs.TaskNotRestarting)
	}
}

func TestClient_RestartTracker_StartError_Recoverable_Delay(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeDelay)
	rt := NewRestartTracker(p, structs.JobTypeSystem, nil)
	recErr := structs.NewRecoverableError(fmt.Errorf("foo"), true)
	for i := 0; i < p.Attempts; i++ {
		state, when := rt.SetStartError(recErr).GetState()
		if state != structs.TaskRestarting {
			t.Fatalf("NextRestart() returned %v, want %v", state, structs.TaskRestarting)
		}
		if !withinJitter(p.Delay, when) {
			t.Fatalf("NextRestart() returned %v; want %v+jitter", when, p.Delay)
		}
	}

	// Next restart should cause delay
	state, when := rt.SetStartError(recErr).GetState()
	if state != structs.TaskRestarting {
		t.Fatalf("NextRestart() returned %v; want %v", state, structs.TaskRestarting)
	}
	if !(when > p.Delay && when <= p.Interval) {
		t.Fatalf("NextRestart() returned %v; want > %v and <= %v", when, p.Delay, p.Interval)
	}
}

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

	testCase := []struct {
		name                   string
		taskLifecycleConfig    *structs.TaskLifecycleConfig
		jobType                string
		shouldRestartOnSuccess bool
		shouldRestartOnFailure bool
	}{
		{
			name:                   "system job no lifecycle",
			taskLifecycleConfig:    nil,
			jobType:                structs.JobTypeSystem,
			shouldRestartOnSuccess: true,
			shouldRestartOnFailure: true,
		},
		{
			name:                   "service job no lifecycle",
			taskLifecycleConfig:    nil,
			jobType:                structs.JobTypeService,
			shouldRestartOnSuccess: true,
			shouldRestartOnFailure: true,
		},
		{
			name:                   "batch job no lifecycle",
			taskLifecycleConfig:    nil,
			jobType:                structs.JobTypeBatch,
			shouldRestartOnSuccess: false,
			shouldRestartOnFailure: true,
		},
		{
			name: "system job w/ ephemeral prestart hook",
			taskLifecycleConfig: &structs.TaskLifecycleConfig{
				Hook:    structs.TaskLifecycleHookPrestart,
				Sidecar: false,
			},
			jobType:                structs.JobTypeSystem,
			shouldRestartOnSuccess: false,
			shouldRestartOnFailure: true,
		},
		{
			name: "system job w/ sidecar prestart hook",
			taskLifecycleConfig: &structs.TaskLifecycleConfig{
				Hook:    structs.TaskLifecycleHookPrestart,
				Sidecar: true,
			},
			jobType:                structs.JobTypeSystem,
			shouldRestartOnSuccess: true,
			shouldRestartOnFailure: true,
		},
		{
			name: "service job w/ ephemeral prestart hook",
			taskLifecycleConfig: &structs.TaskLifecycleConfig{
				Hook:    structs.TaskLifecycleHookPrestart,
				Sidecar: false,
			},
			jobType:                structs.JobTypeService,
			shouldRestartOnSuccess: false,
			shouldRestartOnFailure: true,
		},
		{
			name: "service job w/ sidecar prestart hook",
			taskLifecycleConfig: &structs.TaskLifecycleConfig{
				Hook:    structs.TaskLifecycleHookPrestart,
				Sidecar: true,
			},
			jobType:                structs.JobTypeService,
			shouldRestartOnSuccess: true,
			shouldRestartOnFailure: true,
		},
		{
			name: "batch job w/ ephemeral prestart hook",
			taskLifecycleConfig: &structs.TaskLifecycleConfig{
				Hook:    structs.TaskLifecycleHookPrestart,
				Sidecar: false,
			},
			jobType:                structs.JobTypeService,
			shouldRestartOnSuccess: false,
			shouldRestartOnFailure: true,
		},
		{
			name: "batch job w/ sidecar prestart hook",
			taskLifecycleConfig: &structs.TaskLifecycleConfig{
				Hook:    structs.TaskLifecycleHookPrestart,
				Sidecar: false,
			},
			jobType:                structs.JobTypeBatch,
			shouldRestartOnSuccess: true,
			shouldRestartOnFailure: true,
		},
	}

	for _, testCase := range testCase {
		t.Run(testCase.name, func(t *testing.T) {
			restartPolicy := testPolicy(true, testCase.jobType)
			restartTracker := NewRestartTracker(restartPolicy, testCase.jobType, testCase.taskLifecycleConfig)

			state, _ := restartTracker.SetExitResult(testExitResult(0)).GetState()
			if !testCase.shouldRestartOnSuccess {
				require.Equal(t, structs.TaskTerminated, state)
			} else {
				require.Equal(t, structs.TaskRestarting, state)
			}

			state, _ = restartTracker.SetExitResult(testExitResult(127)).GetState()
			if !testCase.shouldRestartOnFailure {
				require.Equal(t, structs.TaskTerminated, state)
			} else {
				require.Equal(t, structs.TaskRestarting, state)
			}
		})
	}
}
<commit_msg>fix test<commit_after>package restarts

import (
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/plugins/drivers"
)

func testPolicy(success bool, mode string) *structs.RestartPolicy {
	return &structs.RestartPolicy{
		Interval: 2 * time.Minute,
		Delay:    1 * time.Second,
		Attempts: 3,
		Mode:     mode,
	}
}

// withinJitter is a helper that returns whether the returned delay is within
// the jitter.
func withinJitter(expected, actual time.Duration) bool {
	return float64((actual.Nanoseconds()-expected.Nanoseconds())/
		expected.Nanoseconds()) <= jitter
}

func testExitResult(exit int) *drivers.ExitResult {
	return &drivers.ExitResult{
		ExitCode: exit,
	}
}

func TestClient_RestartTracker_ModeDelay(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeDelay)
	rt := NewRestartTracker(p, structs.JobTypeService, nil)
	for i := 0; i < p.Attempts; i++ {
		state, when := rt.SetExitResult(testExitResult(127)).GetState()
		if state != structs.TaskRestarting {
			t.Fatalf("NextRestart() returned %v, want %v", state, structs.TaskRestarting)
		}
		if !withinJitter(p.Delay, when) {
			t.Fatalf("NextRestart() returned %v; want %v+jitter", when, p.Delay)
		}
	}

	// Follow up restarts should cause delay.
	for i := 0; i < 3; i++ {
		state, when := rt.SetExitResult(testExitResult(127)).GetState()
		if state != structs.TaskRestarting {
			t.Fail()
		}
		if !(when > p.Delay && when <= p.Interval) {
			t.Fatalf("NextRestart() returned %v; want > %v and <= %v", when, p.Delay, p.Interval)
		}
	}
}

func TestClient_RestartTracker_ModeFail(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeFail)
	rt := NewRestartTracker(p, structs.JobTypeSystem, nil)
	for i := 0; i < p.Attempts; i++ {
		state, when := rt.SetExitResult(testExitResult(127)).GetState()
		if state != structs.TaskRestarting {
			t.Fatalf("NextRestart() returned %v, want %v", state, structs.TaskRestarting)
		}
		if !withinJitter(p.Delay, when) {
			t.Fatalf("NextRestart() returned %v; want %v+jitter", when, p.Delay)
		}
	}

	// Next restart should cause fail
	if state, _ := rt.SetExitResult(testExitResult(127)).GetState(); state != structs.TaskNotRestarting {
		t.Fatalf("NextRestart() returned %v; want %v", state, structs.TaskNotRestarting)
	}
}

func TestClient_RestartTracker_NoRestartOnSuccess(t *testing.T) {
	t.Parallel()
	p := testPolicy(false, structs.RestartPolicyModeDelay)
	rt := NewRestartTracker(p, structs.JobTypeBatch, nil)
	if state, _ := rt.SetExitResult(testExitResult(0)).GetState(); state != structs.TaskTerminated {
		t.Fatalf("NextRestart() returned %v, expected: %v", state, structs.TaskTerminated)
	}
}

func TestClient_RestartTracker_ZeroAttempts(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeFail)
	p.Attempts = 0

	// Test with a non-zero exit code
	rt := NewRestartTracker(p, structs.JobTypeService, nil)
	if state, when := rt.SetExitResult(testExitResult(1)).GetState(); state != structs.TaskNotRestarting {
		t.Fatalf("expect no restart, got restart/delay: %v/%v", state, when)
	}

	// Even with a zero (successful) exit code non-batch jobs should exit
	// with TaskNotRestarting
	rt = NewRestartTracker(p, structs.JobTypeService, nil)
	if state, when := rt.SetExitResult(testExitResult(0)).GetState(); state != structs.TaskNotRestarting {
		t.Fatalf("expect no restart, got restart/delay: %v/%v", state, when)
	}

	// Batch jobs with a zero exit code and 0 attempts *do* exit cleanly
	// with Terminated
	rt = NewRestartTracker(p, structs.JobTypeBatch, nil)
	if state, when := rt.SetExitResult(testExitResult(0)).GetState(); state != structs.TaskTerminated {
		t.Fatalf("expect terminated, got restart/delay: %v/%v", state, when)
	}

	// Batch jobs with a non-zero exit code and 0 attempts exit with
	// TaskNotRestarting
	rt = NewRestartTracker(p, structs.JobTypeBatch, nil)
	if state, when := rt.SetExitResult(testExitResult(1)).GetState(); state != structs.TaskNotRestarting {
		t.Fatalf("expect no restart, got restart/delay: %v/%v", state, when)
	}
}

func TestClient_RestartTracker_TaskKilled(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeFail)
	p.Attempts = 0
	rt := NewRestartTracker(p, structs.JobTypeService, nil)
	if state, when := rt.SetKilled().GetState(); state != structs.TaskKilled && when != 0 {
		t.Fatalf("expect no restart; got %v %v", state, when)
	}
}

func TestClient_RestartTracker_RestartTriggered(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeFail)
	p.Attempts = 0
	rt := NewRestartTracker(p, structs.JobTypeService, nil)
	if state, when := rt.SetRestartTriggered(false).GetState(); state != structs.TaskRestarting && when != 0 {
		t.Fatalf("expect restart immediately, got %v %v", state, when)
	}
}

func TestClient_RestartTracker_RestartTriggered_Failure(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeFail)
	p.Attempts = 1
	rt := NewRestartTracker(p, structs.JobTypeService, nil)
	if state, when := rt.SetRestartTriggered(true).GetState(); state != structs.TaskRestarting || when == 0 {
		t.Fatalf("expect restart got %v %v", state, when)
	}
	if state, when := rt.SetRestartTriggered(true).GetState(); state != structs.TaskNotRestarting || when != 0 {
		t.Fatalf("expect failed got %v %v", state, when)
	}
}

func TestClient_RestartTracker_StartError_Recoverable_Fail(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeFail)
	rt := NewRestartTracker(p, structs.JobTypeSystem, nil)
	recErr := structs.NewRecoverableError(fmt.Errorf("foo"), true)
	for i := 0; i < p.Attempts; i++ {
		state, when := rt.SetStartError(recErr).GetState()
		if state != structs.TaskRestarting {
			t.Fatalf("NextRestart() returned %v, want %v", state, structs.TaskRestarting)
		}
		if !withinJitter(p.Delay, when) {
			t.Fatalf("NextRestart() returned %v; want %v+jitter", when, p.Delay)
		}
	}

	// Next restart should cause fail
	if state, _ := rt.SetStartError(recErr).GetState(); state != structs.TaskNotRestarting {
		t.Fatalf("NextRestart() returned %v; want %v", state, structs.TaskNotRestarting)
	}
}

func TestClient_RestartTracker_StartError_Recoverable_Delay(t *testing.T) {
	t.Parallel()
	p := testPolicy(true, structs.RestartPolicyModeDelay)
	rt := NewRestartTracker(p, structs.JobTypeSystem, nil)
	recErr := structs.NewRecoverableError(fmt.Errorf("foo"), true)
	for i := 0; i < p.Attempts; i++ {
		state, when := rt.SetStartError(recErr).GetState()
		if state != structs.TaskRestarting {
			t.Fatalf("NextRestart() returned %v, want %v", state, structs.TaskRestarting)
		}
		if !withinJitter(p.Delay, when) {
			t.Fatalf("NextRestart() returned %v; want %v+jitter", when, p.Delay)
		}
	}

	// Next restart should cause delay
	state, when := rt.SetStartError(recErr).GetState()
	if state != structs.TaskRestarting {
		t.Fatalf("NextRestart() returned %v; want %v", state, structs.TaskRestarting)
	}
	if !(when > p.Delay && when <= p.Interval) {
		t.Fatalf("NextRestart() returned %v; want > %v and <= %v", when, p.Delay, p.Interval)
	}
}

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

	testCase := []struct {
		name                   string
		taskLifecycleConfig    *structs.TaskLifecycleConfig
		jobType                string
		shouldRestartOnSuccess bool
		shouldRestartOnFailure bool
	}{
		{
			name:                   "system job no lifecycle",
			taskLifecycleConfig:    nil,
			jobType:                structs.JobTypeSystem,
			shouldRestartOnSuccess: true,
			shouldRestartOnFailure: true,
		},
		{
			name:                   "service job no lifecycle",
			taskLifecycleConfig:    nil,
			jobType:                structs.JobTypeService,
			shouldRestartOnSuccess: true,
			shouldRestartOnFailure: true,
		},
		{
			name:                   "batch job no lifecycle",
			taskLifecycleConfig:    nil,
			jobType:                structs.JobTypeBatch,
			shouldRestartOnSuccess: false,
			shouldRestartOnFailure: true,
		},
		{
			name: "system job w/ ephemeral prestart hook",
			taskLifecycleConfig: &structs.TaskLifecycleConfig{
				Hook:    structs.TaskLifecycleHookPrestart,
				Sidecar: false,
			},
			jobType:                structs.JobTypeSystem,
			shouldRestartOnSuccess: false,
			shouldRestartOnFailure: true,
		},
		{
			name: "system job w/ sidecar prestart hook",
			taskLifecycleConfig: &structs.TaskLifecycleConfig{
				Hook:    structs.TaskLifecycleHookPrestart,
				Sidecar: true,
			},
			jobType:                structs.JobTypeSystem,
			shouldRestartOnSuccess: true,
			shouldRestartOnFailure: true,
		},
		{
			name: "service job w/ ephemeral prestart hook",
			taskLifecycleConfig: &structs.TaskLifecycleConfig{
				Hook:    structs.TaskLifecycleHookPrestart,
				Sidecar: false,
			},
			jobType:                structs.JobTypeService,
			shouldRestartOnSuccess: false,
			shouldRestartOnFailure: true,
		},
		{
			name: "service job w/ sidecar prestart hook",
			taskLifecycleConfig: &structs.TaskLifecycleConfig{
				Hook:    structs.TaskLifecycleHookPrestart,
				Sidecar: true,
			},
			jobType:                structs.JobTypeService,
			shouldRestartOnSuccess: true,
			shouldRestartOnFailure: true,
		},
		{
			name: "batch job w/ ephemeral prestart hook",
			taskLifecycleConfig: &structs.TaskLifecycleConfig{
				Hook:    structs.TaskLifecycleHookPrestart,
				Sidecar: false,
			},
			jobType:                structs.JobTypeService,
			shouldRestartOnSuccess: false,
			shouldRestartOnFailure: true,
		},
		{
			name: "batch job w/ sidecar prestart hook",
			taskLifecycleConfig: &structs.TaskLifecycleConfig{
				Hook:    structs.TaskLifecycleHookPrestart,
				Sidecar: true,
			},
			jobType:                structs.JobTypeBatch,
			shouldRestartOnSuccess: true,
			shouldRestartOnFailure: true,
		},
	}

	for _, testCase := range testCase {
		t.Run(testCase.name, func(t *testing.T) {
			restartPolicy := testPolicy(true, testCase.jobType)
			restartTracker := NewRestartTracker(restartPolicy, testCase.jobType, testCase.taskLifecycleConfig)

			state, _ := restartTracker.SetExitResult(testExitResult(0)).GetState()
			if !testCase.shouldRestartOnSuccess {
				require.Equal(t, structs.TaskTerminated, state)
			} else {
				require.Equal(t, structs.TaskRestarting, state)
			}

			state, _ = restartTracker.SetExitResult(testExitResult(127)).GetState()
			if !testCase.shouldRestartOnFailure {
				require.Equal(t, structs.TaskTerminated, state)
			} else {
				require.Equal(t, structs.TaskRestarting, state)
			}
		})
	}
}
<|endoftext|>


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

package codec

import (
	"fmt"

	"github.com/henrylee2cn/goutil"
)

// protobuf codec id
const (
	NAME_STRING = "string"
	ID_STRING   = 's'
)

func init() {
	Reg(new(StringCodec))
}

// StringCodec string codec
type StringCodec struct{}

// Name returns codec string
func (StringCodec) Name() string {
	return NAME_STRING
}

// Id returns codec id
func (StringCodec) Id() byte {
	return ID_STRING
}

// Marshal returns the string encoding of v.
func (StringCodec) Marshal(v interface{}) ([]byte, error) {
	var b []byte
	switch s := v.(type) {
	case nil:
	case string:
		b = goutil.StringToBytes(s)
	case *string:
		b = goutil.StringToBytes(*s)
	case []byte:
		b = s
	case *[]byte:
		b = *s
	default:
		return nil, fmt.Errorf("%T can not be directly converted to []byte type", v)
	}
	return b, nil
}

// Unmarshal parses the string-encoded data and stores the result
// in the value pointed to by v.
func (StringCodec) Unmarshal(data []byte, v interface{}) error {
	switch s := v.(type) {
	case nil:
		return nil
	case *string:
		*s = string(data)
	case []byte:
		copy(s, data)
	case *[]byte:
		*s = make([]byte, len(data))
		copy(*s, data)
	default:
		return fmt.Errorf("[]byte can not be directly converted to %T type", v)
	}
	return nil
}
<commit_msg>codec: improve string codec error info<commit_after>// Copyright 2015-2017 HenryLee. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package codec

import (
	"fmt"

	"github.com/henrylee2cn/goutil"
)

// protobuf codec id
const (
	NAME_STRING = "string"
	ID_STRING   = 's'
)

func init() {
	Reg(new(StringCodec))
}

// StringCodec string codec
type StringCodec struct{}

// Name returns codec string
func (StringCodec) Name() string {
	return NAME_STRING
}

// Id returns codec id
func (StringCodec) Id() byte {
	return ID_STRING
}

// Marshal returns the string encoding of v.
func (StringCodec) Marshal(v interface{}) ([]byte, error) {
	var b []byte
	switch s := v.(type) {
	case nil:
	case string:
		b = goutil.StringToBytes(s)
	case *string:
		b = goutil.StringToBytes(*s)
	case []byte:
		b = s
	case *[]byte:
		b = *s
	default:
		return nil, fmt.Errorf("string codec: %T can not be directly converted to []byte type", v)
	}
	return b, nil
}

// Unmarshal parses the string-encoded data and stores the result
// in the value pointed to by v.
func (StringCodec) Unmarshal(data []byte, v interface{}) error {
	switch s := v.(type) {
	case nil:
		return nil
	case *string:
		*s = string(data)
	case []byte:
		copy(s, data)
	case *[]byte:
		*s = make([]byte, len(data))
		copy(*s, data)
	default:
		return fmt.Errorf("string codec: []byte can not be directly converted to %T type", v)
	}
	return nil
}
<|endoftext|>


<commit_before>package flow

import (
	"fmt"
	"reflect"
	"testing"
)

func TestMapSingleParameter(t *testing.T) {

	New().Slice(
		[]int{1, 2, 3, 4, 5},
	).Map(func(t int) (int, int) {
		return t, t * 7
	}).Map(func(x, y int) {
		fmt.Println("x=", x, "7*x=", y)
	}).Run()

	t.Logf("single parameter mapping runs well")

}

func TestGroupByKeyMap(t *testing.T) {

	fmt.Println("group by result mapping testing start...")
	New().Slice(
		[]int{1, 1, 1, 1, 1, 2, 2, 2, 2,
			3, 3, 3, 4, 4, 5,
			100, 234, 43, 100, 43, 43, 43},
	).Map(func(t int) (int, int) {
		return t, t * 2
	}).Map(func(key, value int) (int, int) {
		return key, value * 3
	}).GroupByKey().Map(func(key int, values []int) {
		fmt.Printf("key: %d values: %v\n", key, values)
	}).Run()

	fmt.Println("group by result mapping runs well")

}

func TestCoGroupMap(t *testing.T) {
	f := New()
	left := f.Slice(
		[]int{1, 1, 2},
	).Map(func(t int) (int, int) {
		return t, t * 2
	})

	right := f.Slice(
		[]int{1, 2, 2},
	).Map(func(t int) (int, int) {
		return t, t * 5
	})

	type result struct {
		key   int
		left  []int
		right []int
	}
	cogroup_result := left.CoGroup(right).Map(func(key int, lefts, rights []int) result {
		return result{
			key:   key,
			left:  lefts,
			right: rights,
		}
	})

	outChan := make(chan result, 0)
	cogroup_result.AddOutput(outChan)
	go cogroup_result.Run()

	got := make([]result, 0, 2)
	for item := range outChan {
		got = append(got, item)
	}

	want := []result{
		{
			key:   1,
			left:  []int{2, 2},
			right: []int{5},
		}, {
			key:   2,
			left:  []int{4},
			right: []int{10, 10},
		},
	}

	if !reflect.DeepEqual(got, want) {
		t.Fatalf("Got %v want %v", got, want)
	}
}
<commit_msg>Rewrite tests.<commit_after>package flow

import (
	"reflect"
	"testing"
)

func TestMapSingleParameter(t *testing.T) {
	dataset := New().Slice(
		[]int{1, 2, 3},
	).Map(func(t int) (int, int) {
		return t, t
	}).Map(func(x, y int) int {
		return x + y
	})
	got := collectOutput(dataset, reflect.TypeOf(int(1)))

	if want := []int{2, 4, 6}; !reflect.DeepEqual(got, want) {
		t.Errorf("Got: %v want: %v", got, want)
	}
}

func TestGroupByKeyMap(t *testing.T) {
	dataset := New().Slice(
		[]int{1, 1, 2, 2, 3, 3},
	).Map(func(t int) (int, int) {
		return t, t * 2
	}).GroupByKey().Map(func(key int, values []int) []int {
		return append([]int{key}, values...)
	})

	got := collectOutput(dataset, reflect.TypeOf([]int{}))

	if want := [][]int{{1, 2, 2}, {2, 4, 4}, {3, 6, 6}}; !reflect.DeepEqual(got, want) {
		t.Errorf("Got %v want %v", got, want)
	}
}

func TestCoGroupMap(t *testing.T) {
	f := New()
	left := f.Slice(
		[]int{1, 1, 2},
	).Map(func(t int) (int, int) {
		return t, t * 2
	})

	right := f.Slice(
		[]int{1, 2, 2},
	).Map(func(t int) (int, int) {
		return t, t * 5
	})

	type result struct {
		key   int
		left  []int
		right []int
	}
	cogroupResult := left.CoGroup(right).Map(func(key int, lefts, rights []int) result {
		return result{
			key:   key,
			left:  lefts,
			right: rights,
		}
	})

	got := collectOutput(cogroupResult, reflect.TypeOf(result{}))
	want := []result{
		{
			key:   1,
			left:  []int{2, 2},
			right: []int{5},
		}, {
			key:   2,
			left:  []int{4},
			right: []int{10, 10},
		},
	}

	if !reflect.DeepEqual(got, want) {
		t.Fatalf("Got %v want %v", got, want)
	}
}
<|endoftext|>


<commit_before>package main

import (
  "os"
  "os/exec"
  "strings"
  "strconv"
  "encoding/json"
  "io/ioutil"
  "github.com/jawher/mow.cli"
  log "github.com/Sirupsen/logrus"
)

var version string
var commit string

type configuration struct {
  Periscope periscope `json:"periscope"`
  Video video `json:"video"`
}

type periscope struct {
  Key string `json:"key"`
}

type video struct {
  Width int `json:"width"`
  Height int `json:"height"`
  Sharpness int `json:"sharpness"`
  Quality int `json:"quality"`
  Bitrate int `json:"bitrate"`
  VFlip bool `json:"vflip"`
  HFlip bool `json:"hflip"`
}

func joinProps(props map[string]string, kvSeparator string, fieldSeparator string) string {
  parts := []string{}
  for key, value := range props {
    parts = append(parts, key + kvSeparator + value)
  }
  return strings.Join(parts, fieldSeparator)
}

func showCommand(cmd *exec.Cmd) {
  log.WithFields(log.Fields{
    "cmd": cmd.Path,
    "args": strings.Join(cmd.Args[1:], " "),
  }).Debug("Running command")
}

func mergeString(l string, r string) string {
  if r == "" {
    return l
  } else {
    return r
  }
}

func mergeInt(l int, r int) int {
  if r == 0 {
    return l
  } else {
    return r
  }
}

func mergeBool(l bool, r bool, def bool) bool {
  if r == def {
    return l
  } else {
    return r
  }
}

func mergePeriscope(l *periscope, r *periscope) *periscope {
  return &periscope{
    Key: mergeString(l.Key, r.Key),
  }
}

func mergeVideo(l *video, r *video) *video {
  return &video{
    Width: mergeInt(l.Width, r.Width),
    Height: mergeInt(l.Height, r.Height),
    Sharpness: mergeInt(l.Sharpness, r.Sharpness),
    Quality: mergeInt(l.Quality, r.Quality),
    Bitrate: mergeInt(l.Bitrate, r.Bitrate),
    VFlip: mergeBool(l.VFlip, r.VFlip, false),
    HFlip: mergeBool(l.HFlip, r.HFlip, false),
  }
}

func mergeConfig(l *configuration, r *configuration) *configuration {
  return &configuration{
    Periscope: *mergePeriscope(&l.Periscope, &r.Periscope),
    Video: *mergeVideo(&l.Video, &r.Video),
  }
}

func main() {
  app := cli.App("piriscope", "Piriscope - https://github.com/schmich/piriscope")
  key := app.StringOpt("k key", "", "Periscope stream key")
  conf := app.StringOpt("c conf", "", "Configuration file")
  verbose := app.BoolOpt("v verbose", false, "Verbose output")

  app.Version("version", "piriscope " + version + " " + commit)

  app.Action = func () {
    if *verbose {
      log.SetLevel(log.DebugLevel)
    }

    var fileConfig configuration
    if *conf != "" {
      content, err := ioutil.ReadFile(*conf)
      if err != nil {
        log.Fatal(err)
      }

      log.WithFields(log.Fields{ "file": *conf }).Info("Using configuration file")

      err = json.Unmarshal(content, &fileConfig)
      if err != nil {
        log.Fatal(err)
      }
    }

    defaultConfig := &configuration{
      Periscope: periscope{
        Key: "",
      },
      Video: video{
        Width: 960,
        Height: 540,
        Sharpness: 30,
        Quality: 80,
        Bitrate: 800000,
      },
    }

    cliConfig := &configuration{
      Periscope: periscope{
        Key: *key,
      },
    }

    config := mergeConfig(defaultConfig, mergeConfig(&fileConfig, cliConfig))

    if config.Periscope.Key == "" {
      log.Fatal("Periscope stream key (-k, --key) is required.")
    }

    streamUrl := "rtmp://va.pscp.tv:80/x/" + config.Periscope.Key

    width := strconv.Itoa(config.Video.Width)
    height := strconv.Itoa(config.Video.Height)
    sharpness := strconv.Itoa(config.Video.Sharpness)
    quality := strconv.Itoa(config.Video.Quality)
    bitrate := strconv.Itoa(config.Video.Bitrate)

    videoProps := map[string]string{
      "width": width,
      "height": height,
      "pixelformat": "4",
    }

    controlProps := map[string]string{
      "sharpness": sharpness,
      "compression_quality": quality,
      "video_bitrate_mode": "1",
      "video_bitrate": bitrate,
    }

    v4l2 := exec.Command("v4l2-ctl", "--set-fmt-video=" + joinProps(videoProps, "=", ","))
    v4l2.Stdout = os.Stdout
    v4l2.Stderr = os.Stderr

    showCommand(v4l2)
    err := v4l2.Run()
    if err != nil {
      log.Fatal(err)
    }

    v4l2 = exec.Command("v4l2-ctl", "--set-ctrl=" + joinProps(controlProps, "=", ","))
    v4l2.Stdout = os.Stdout
    v4l2.Stderr = os.Stderr

    showCommand(v4l2)
    err = v4l2.Run()
    if err != nil {
      log.Fatal(err)
    }

    raspividArgs := []string{
      "-o", "-",                        // Write video to stdout.
      "-t", "0",                        // Continuous capturing (no timeout).
      "-w", width,                      // Output video width.
      "-h", height,                     // Output video height.
      "-fps", "30",                     // Capture video at 30 frames per second.
      "-b", bitrate,                    // Capture video bitrate.
    }

    if config.Video.VFlip {
      raspividArgs = append(raspividArgs, "-vf")
    }

    if config.Video.HFlip {
      raspividArgs = append(raspividArgs, "-hf")
    }

    raspivid := exec.Command("raspivid", raspividArgs...)

    ffmpegArgs := []string{
      "-re",                            // Read from input at its native framerate. Best for real-time/streaming output.
      "-f", "lavfi", "-i", "anullsrc",  // No input audio.
      "-i", "-",                        // Use stdin for video (from raspivid).
      "-acodec", "aac",                 // Use AAC codec for audio (Periscope requirement).
      "-b:a", "0",                      // Zero audio bitrate since we have no input audio.
      "-map", "0:a",                    // Use stream 0 for audio (anullsrc).
      "-map", "1:v",                    // Use stream 1 for video (stdin).
      "-f", "h264",                     // Use H.264 codec for video (Periscope requirement).
      "-vcodec", "copy",                // Copy video data directly from input.
      "-g", "60",                       // Keyframe interval: one keyframe every 60 frames (2 seconds for 30 fps video; Periscope requirement).
      "-f", "flv",                      // Package output in a Flash Video container (Periscope requirement).
      streamUrl,                        // RTMP streaming destination.
    }

    ffmpeg := exec.Command("ffmpeg", ffmpegArgs...)

    ffmpegStdin, err := ffmpeg.StdinPipe()
    if err != nil {
      log.Fatal(err)
    }

    raspivid.Stderr = os.Stderr
    raspivid.Stdout = ffmpegStdin
    ffmpeg.Stdout = os.Stdout
    ffmpeg.Stderr = os.Stderr

    defer func () {
      raspivid.Wait()
      ffmpegStdin.Close()
      ffmpeg.Wait()
    }()

    showCommand(raspivid)
    showCommand(ffmpeg)

    ffmpeg.Start()
    raspivid.Start()
  }

  app.Run(os.Args)
}
<commit_msg>Refactor stream execution code into its own method.<commit_after>package main

import (
  "os"
  "os/exec"
  "strings"
  "strconv"
  "encoding/json"
  "io/ioutil"
  "errors"
  "github.com/jawher/mow.cli"
  log "github.com/Sirupsen/logrus"
)

var version string
var commit string

type configuration struct {
  Periscope periscope `json:"periscope"`
  Video video `json:"video"`
}

type periscope struct {
  Key string `json:"key"`
}

type video struct {
  Width int `json:"width"`
  Height int `json:"height"`
  Sharpness int `json:"sharpness"`
  Quality int `json:"quality"`
  Bitrate int `json:"bitrate"`
  VFlip bool `json:"vflip"`
  HFlip bool `json:"hflip"`
}

func joinProps(props map[string]string, kvSeparator string, fieldSeparator string) string {
  parts := []string{}
  for key, value := range props {
    parts = append(parts, key + kvSeparator + value)
  }
  return strings.Join(parts, fieldSeparator)
}

func showCommand(cmd *exec.Cmd) {
  log.WithFields(log.Fields{
    "cmd": cmd.Path,
    "args": strings.Join(cmd.Args[1:], " "),
  }).Debug("Running command")
}

func mergeString(l string, r string) string {
  if r == "" {
    return l
  } else {
    return r
  }
}

func mergeInt(l int, r int) int {
  if r == 0 {
    return l
  } else {
    return r
  }
}

func mergeBool(l bool, r bool, def bool) bool {
  if r == def {
    return l
  } else {
    return r
  }
}

func mergePeriscope(l *periscope, r *periscope) *periscope {
  return &periscope{
    Key: mergeString(l.Key, r.Key),
  }
}

func mergeVideo(l *video, r *video) *video {
  return &video{
    Width: mergeInt(l.Width, r.Width),
    Height: mergeInt(l.Height, r.Height),
    Sharpness: mergeInt(l.Sharpness, r.Sharpness),
    Quality: mergeInt(l.Quality, r.Quality),
    Bitrate: mergeInt(l.Bitrate, r.Bitrate),
    VFlip: mergeBool(l.VFlip, r.VFlip, false),
    HFlip: mergeBool(l.HFlip, r.HFlip, false),
  }
}

func mergeConfig(l *configuration, r *configuration) *configuration {
  return &configuration{
    Periscope: *mergePeriscope(&l.Periscope, &r.Periscope),
    Video: *mergeVideo(&l.Video, &r.Video),
  }
}

func runStream(config *configuration) error {
  if config.Periscope.Key == "" {
    return errors.New("Periscope stream key (-k, --key) is required.")
  }

  streamUrl := "rtmp://va.pscp.tv:80/x/" + config.Periscope.Key

  width := strconv.Itoa(config.Video.Width)
  height := strconv.Itoa(config.Video.Height)
  sharpness := strconv.Itoa(config.Video.Sharpness)
  quality := strconv.Itoa(config.Video.Quality)
  bitrate := strconv.Itoa(config.Video.Bitrate)

  videoProps := map[string]string{
    "width": width,
    "height": height,
    "pixelformat": "4",
  }

  v4l2 := exec.Command("v4l2-ctl", "--set-fmt-video=" + joinProps(videoProps, "=", ","))
  v4l2.Stdout = os.Stdout
  v4l2.Stderr = os.Stderr

  showCommand(v4l2)
  if err := v4l2.Run(); err != nil {
    return err
  }

  controlProps := map[string]string{
    "sharpness": sharpness,
    "compression_quality": quality,
    "video_bitrate_mode": "1",
    "video_bitrate": bitrate,
  }

  v4l2 = exec.Command("v4l2-ctl", "--set-ctrl=" + joinProps(controlProps, "=", ","))
  v4l2.Stdout = os.Stdout
  v4l2.Stderr = os.Stderr

  showCommand(v4l2)
  if err := v4l2.Run(); err != nil {
    return err
  }

  raspividArgs := []string{
    "-o", "-",                        // Write video to stdout.
    "-t", "0",                        // Continuous capturing (no timeout).
    "-w", width,                      // Output video width.
    "-h", height,                     // Output video height.
    "-fps", "30",                     // Capture video at 30 frames per second.
    "-b", bitrate,                    // Capture video bitrate.
  }

  if config.Video.VFlip {
    raspividArgs = append(raspividArgs, "-vf")
  }

  if config.Video.HFlip {
    raspividArgs = append(raspividArgs, "-hf")
  }

  raspivid := exec.Command("raspivid", raspividArgs...)

  ffmpegArgs := []string{
    "-re",                            // Read from input at its native framerate. Best for real-time/streaming output.
    "-f", "lavfi", "-i", "anullsrc",  // No input audio.
    "-i", "-",                        // Use stdin for video (from raspivid).
    "-acodec", "aac",                 // Use AAC codec for audio (Periscope requirement).
    "-b:a", "0",                      // Zero audio bitrate since we have no input audio.
    "-map", "0:a",                    // Use stream 0 for audio (anullsrc).
    "-map", "1:v",                    // Use stream 1 for video (stdin).
    "-f", "h264",                     // Use H.264 codec for video (Periscope requirement).
    "-vcodec", "copy",                // Copy video data directly from input.
    "-g", "60",                       // Keyframe interval: one keyframe every 60 frames (2 seconds for 30 fps video; Periscope requirement).
    "-f", "flv",                      // Package output in a Flash Video container (Periscope requirement).
    streamUrl,                        // RTMP streaming destination.
  }

  ffmpeg := exec.Command("ffmpeg", ffmpegArgs...)

  ffmpegStdin, err := ffmpeg.StdinPipe()
  if err != nil {
    return err
  }

  raspivid.Stderr = os.Stderr
  raspivid.Stdout = ffmpegStdin
  ffmpeg.Stdout = os.Stdout
  ffmpeg.Stderr = os.Stderr

  defer func () {
    raspivid.Wait()
    ffmpegStdin.Close()
    ffmpeg.Wait()
  }()

  showCommand(ffmpeg)
  if err := ffmpeg.Start(); err != nil {
    return err
  }

  showCommand(raspivid)
  if err := raspivid.Start(); err != nil {
    return err
  }

  return nil
}

func main() {
  app := cli.App("piriscope", "Piriscope - https://github.com/schmich/piriscope")
  key := app.StringOpt("k key", "", "Periscope stream key")
  conf := app.StringOpt("c conf", "", "Configuration file")
  verbose := app.BoolOpt("v verbose", false, "Verbose output")

  app.Version("version", "piriscope " + version + " " + commit)

  app.Action = func () {
    if *verbose {
      log.SetLevel(log.DebugLevel)
    }

    var fileConfig configuration
    if *conf != "" {
      content, err := ioutil.ReadFile(*conf)
      if err != nil {
        log.Fatal(err)
      }

      log.WithFields(log.Fields{ "file": *conf }).Info("Using configuration file")

      err = json.Unmarshal(content, &fileConfig)
      if err != nil {
        log.Fatal(err)
      }
    }

    defaultConfig := &configuration{
      Periscope: periscope{
        Key: "",
      },
      Video: video{
        Width: 960,
        Height: 540,
        Sharpness: 30,
        Quality: 80,
        Bitrate: 800000,
      },
    }

    cliConfig := &configuration{
      Periscope: periscope{
        Key: *key,
      },
    }

    config := mergeConfig(defaultConfig, mergeConfig(&fileConfig, cliConfig))
    if err := runStream(config); err != nil {
      log.Fatal(err)
    }
  }

  app.Run(os.Args)
}
<|endoftext|>


<commit_before>package ca

import (
	"crypto"
	"crypto/rand"
	"crypto/sha1"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/hex"
	"encoding/json"
	"errors"
	"math/big"
	"time"

	"github.com/cloudflare/cfssl/certdb"
	"github.com/cloudflare/cfssl/config"
	"github.com/cloudflare/cfssl/csr"
	"github.com/cloudflare/cfssl/helpers"
	"github.com/cloudflare/cfssl/signer"
	"github.com/cloudflare/cfssl/signer/local"
)

type StorageProvider interface {
	GetMetadata(key []byte) ([]byte, error)
	SetMetadata(key, value []byte) error
	Accessor() certdb.Accessor
}

type KeyProvider interface {
	GenerateKeyPair(label string, algo string, size int) (crypto.Signer, error)
	FindKeyPair(key crypto.PublicKey) (crypto.Signer, error)
}

// CertificationAuthority represents a certification authority.
type CertificationAuthority struct {
	sp StorageProvider
	kp KeyProvider

	signer *local.Signer
}

// Init creates a CA with given config.
func Init(cfg *Config, caFile string, kp KeyProvider) (*CertificationAuthority, error) {
	db, err := openDB(caFile)
	if err != nil {
		return nil, err
	}
	ca := &CertificationAuthority{db, kp, nil}

	req := cfg.CertificateRequest()
	policy, err := cfg.Signing()
	if err != nil {
		return nil, err
	}
	err = ca.init(req, policy)
	if err != nil {
		return nil, err
	}

	if cfg.SelfSign {
		err = ca.selfSign()
		if err != nil {
			return nil, err
		}
	}

	return ca, nil
}

// Open opens an existing CA.
func Open(caFile string, kp KeyProvider) (*CertificationAuthority, error) {
	db, err := openDB(caFile)
	if err != nil {
		return nil, err
	}
	ca := &CertificationAuthority{db, kp, nil}

	err = ca.initSigner()
	if err != nil {
		return nil, err
	}

	return ca, nil
}

// init generates a key pair and creates a certificate signing request
// for the CA.
func (ca *CertificationAuthority) init(req *csr.CertificateRequest, policy *config.Signing) error {
	if csrPEM, _ := ca.CertificateRequestPEM(); csrPEM != nil {
		return errors.New("ca csr exists")
	}

	key, err := ca.kp.GenerateKeyPair(req.CN, req.KeyRequest.Algo(), req.KeyRequest.Size())
	if err != nil {
		return err
	}

	csrPEM, err := csr.Generate(key, req)
	if err != nil {
		return err
	}
	err = ca.sp.SetMetadata([]byte("csr"), csrPEM)
	if err != nil {
		return err
	}

	err = ca.SetPolicy(policy)
	if err != nil {
		return err
	}

	err = ca.initSigner()
	if err != nil {
		return err
	}

	return nil
}

// initSigner initializes a new signer for the CA.
func (ca *CertificationAuthority) initSigner() error {
	key, err := ca.privateKey()
	if err != nil {
		return err
	}

	cert, _ := ca.Certificate()

	policy, err := ca.Policy()
	if err != nil {
		return err
	}

	signer, err := local.NewSigner(key, cert, signer.DefaultSigAlgo(key), policy)
	if err != nil {
		return err
	}

	signer.SetDBAccessor(ca.sp.Accessor())

	ca.signer = signer
	return nil
}

// Certificate returns the certificate of the CA.
func (ca *CertificationAuthority) Certificate() (*x509.Certificate, error) {
	certPEM, err := ca.CertificatePEM()
	if err != nil {
		return nil, err
	}

	return helpers.ParseCertificatePEM(certPEM)
}

// Certificate returns the certificate of the CA in PEM encoding.
func (ca *CertificationAuthority) CertificatePEM() ([]byte, error) {
	return ca.sp.GetMetadata([]byte("cert"))
}

// selfSign creates a self-signed certificate for the CA.
func (ca *CertificationAuthority) selfSign() error {
	csrPEM, err := ca.CertificateRequestPEM()
	if err != nil {
		return err
	}

	certPEM, err := ca.Issue(csrPEM)
	if err != nil {
		return err
	}

	return ca.ImportCertificate(certPEM)
}

// ImportCertificate imports the given certificate if the CA does not
// have one.
func (ca *CertificationAuthority) ImportCertificate(certPEM []byte) error {
	if cert, _ := ca.Certificate(); cert != nil {
		return errors.New("ca cert exists")
	}

	_, err := helpers.ParseCertificatePEM(certPEM)
	if err != nil {
		return err
	}

	// TODO: Check signature and compare with original CSR.

	err = ca.sp.SetMetadata([]byte("cert"), certPEM)
	if err != nil {
		return err
	}

	return ca.initSigner()
}

// CertificateRequest returns the certificate signing request of the CA.
func (ca *CertificationAuthority) CertificateRequest() (*x509.CertificateRequest, error) {
	csrPEM, err := ca.CertificateRequestPEM()
	if err != nil {
		return nil, err
	}

	return helpers.ParseCSRPEM(csrPEM)
}

// CertificateRequestPEM returns the certificate signing request of the
// CA in PEM encoding.
func (ca *CertificationAuthority) CertificateRequestPEM() ([]byte, error) {
	return ca.sp.GetMetadata([]byte("csr"))
}

// PublicKey returns the public key from the CA certificate or CSR.
func (ca *CertificationAuthority) PublicKey() (crypto.PublicKey, error) {
	if cert, _ := ca.Certificate(); cert != nil {
		return cert.PublicKey, nil
	}

	if csr, _ := ca.CertificateRequest(); csr != nil {
		return csr.PublicKey, nil
	}

	return nil, errors.New("no valid csr in db")
}

// privateKey returns the private key as a crypto.Signer.
func (ca *CertificationAuthority) privateKey() (crypto.Signer, error) {
	pub, err := ca.PublicKey()
	if err != nil {
		return nil, err
	}

	return ca.kp.FindKeyPair(pub)
}

// KeyID returns the identifier of the signing key, which will also be
// the Authority Key Identifier (AKI) for issued certificates.
func (ca *CertificationAuthority) KeyID() ([]byte, error) {
	if cert, _ := ca.Certificate(); cert != nil {
		return cert.SubjectKeyId, nil
	}

	pub, err := ca.PublicKey()
	if err != nil {
		return nil, err
	}

	pkixPub, err := x509.MarshalPKIXPublicKey(pub)
	if err != nil {
		return nil, err
	}

	var pubKeyInfo struct {
		Algo      pkix.AlgorithmIdentifier
		BitString asn1.BitString
	}
	_, err = asn1.Unmarshal(pkixPub, &pubKeyInfo)
	if err != nil {
		return nil, err
	}

	hash := sha1.New()
	hash.Write(pubKeyInfo.BitString.Bytes)
	return hash.Sum(nil), nil
}

// Policy returns the signing policy of the CA.
func (ca *CertificationAuthority) Policy() (*config.Signing, error) {
	if ca.signer != nil {
		return ca.signer.Policy(), nil
	}

	policyJSON, err := ca.sp.GetMetadata([]byte("policy"))
	if err != nil {
		return nil, err
	}
	var policy *config.Signing
	err = json.Unmarshal(policyJSON, &policy)
	if err != nil {
		return nil, err
	}

	return policy, nil
}

// SetPolicy sets the signing policy of the CA.
func (ca *CertificationAuthority) SetPolicy(policy *config.Signing) error {
	if !policy.Valid() {
		return errors.New("invalid policy")
	}

	policyJSON, err := json.Marshal(policy)
	if err != nil {
		return err
	}
	err = ca.sp.SetMetadata([]byte("policy"), policyJSON)
	if err != nil {
		return err
	}

	if ca.signer != nil {
		ca.signer.SetPolicy(policy)
	}

	return nil
}

// Issue signs a PEM-encoded CSR and returns the certificate in PEM.
func (ca *CertificationAuthority) Issue(csrPEM []byte) ([]byte, error) {
	var oidExtensionAuthorityKeyId = config.OID([]int{2, 5, 29, 35})

	keyID, err := ca.KeyID()
	if err != nil {
		return nil, err
	}

	type authKeyId struct {
		KeyIdentifier []byte `asn1:"tag:0"`
	}
	aki, err := asn1.Marshal(authKeyId{keyID})
	if err != nil {
		return nil, err
	}

	req := signer.SignRequest{
		Request: string(csrPEM),
		Extensions: []signer.Extension{
			signer.Extension{
				ID:       oidExtensionAuthorityKeyId,
				Critical: false,
				Value:    hex.EncodeToString(aki),
			},
		},
	}

	return ca.signer.Sign(req)
}

// Revoke marks the certificate identified by its serial number revoked.
// The reasonCode is defined in RFC 5280 5.3.1.
func (ca *CertificationAuthority) Revoke(serial string, reasonCode int) error {
	keyID, err := ca.KeyID()
	if err != nil {
		return err
	}
	aki := hex.EncodeToString(keyID)

	return ca.sp.Accessor().RevokeCertificate(serial, aki, reasonCode)
}

// CRL returns a DER-encoded Certificate Revocation List, signed by the CA.
func (ca *CertificationAuthority) CRL(ttl time.Duration) ([]byte, error) {
	certs, err := ca.sp.Accessor().GetRevokedAndUnexpiredCertificates()
	if err != nil {
		return nil, err
	}

	var revokedCerts []pkix.RevokedCertificate
	for _, certRecord := range certs {
		serialInt := new(big.Int)
		serialInt.SetString(certRecord.Serial, 10)

		revokedCert := pkix.RevokedCertificate{
			SerialNumber:   serialInt,
			RevocationTime: certRecord.RevokedAt,
		}
		revokedCerts = append(revokedCerts, revokedCert)
	}

	cert, err := ca.Certificate()
	if err != nil {
		return nil, err
	}

	key, err := ca.privateKey()
	if err != nil {
		return nil, err
	}

	return cert.CreateCRL(rand.Reader, key, revokedCerts, time.Now(), time.Now().Add(ttl))
}
<commit_msg>Add AKI extension in self-signed SignRequests only<commit_after>package ca

import (
	"crypto"
	"crypto/rand"
	"crypto/sha1"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/hex"
	"encoding/json"
	"errors"
	"math/big"
	"time"

	"github.com/cloudflare/cfssl/certdb"
	"github.com/cloudflare/cfssl/config"
	"github.com/cloudflare/cfssl/csr"
	"github.com/cloudflare/cfssl/helpers"
	"github.com/cloudflare/cfssl/signer"
	"github.com/cloudflare/cfssl/signer/local"
)

type StorageProvider interface {
	GetMetadata(key []byte) ([]byte, error)
	SetMetadata(key, value []byte) error
	Accessor() certdb.Accessor
}

type KeyProvider interface {
	GenerateKeyPair(label string, algo string, size int) (crypto.Signer, error)
	FindKeyPair(key crypto.PublicKey) (crypto.Signer, error)
}

// CertificationAuthority represents a certification authority.
type CertificationAuthority struct {
	sp StorageProvider
	kp KeyProvider

	signer *local.Signer
}

// Init creates a CA with given config.
func Init(cfg *Config, caFile string, kp KeyProvider) (*CertificationAuthority, error) {
	db, err := openDB(caFile)
	if err != nil {
		return nil, err
	}
	ca := &CertificationAuthority{db, kp, nil}

	req := cfg.CertificateRequest()
	policy, err := cfg.Signing()
	if err != nil {
		return nil, err
	}
	err = ca.init(req, policy)
	if err != nil {
		return nil, err
	}

	if cfg.SelfSign {
		err = ca.selfSign()
		if err != nil {
			return nil, err
		}
	}

	return ca, nil
}

// Open opens an existing CA.
func Open(caFile string, kp KeyProvider) (*CertificationAuthority, error) {
	db, err := openDB(caFile)
	if err != nil {
		return nil, err
	}
	ca := &CertificationAuthority{db, kp, nil}

	err = ca.initSigner()
	if err != nil {
		return nil, err
	}

	return ca, nil
}

// init generates a key pair and creates a certificate signing request
// for the CA.
func (ca *CertificationAuthority) init(req *csr.CertificateRequest, policy *config.Signing) error {
	if csrPEM, _ := ca.CertificateRequestPEM(); csrPEM != nil {
		return errors.New("ca csr exists")
	}

	key, err := ca.kp.GenerateKeyPair(req.CN, req.KeyRequest.Algo(), req.KeyRequest.Size())
	if err != nil {
		return err
	}

	csrPEM, err := csr.Generate(key, req)
	if err != nil {
		return err
	}
	err = ca.sp.SetMetadata([]byte("csr"), csrPEM)
	if err != nil {
		return err
	}

	err = ca.SetPolicy(policy)
	if err != nil {
		return err
	}

	err = ca.initSigner()
	if err != nil {
		return err
	}

	return nil
}

// initSigner initializes a new signer for the CA.
func (ca *CertificationAuthority) initSigner() error {
	key, err := ca.privateKey()
	if err != nil {
		return err
	}

	cert, _ := ca.Certificate()

	policy, err := ca.Policy()
	if err != nil {
		return err
	}

	signer, err := local.NewSigner(key, cert, signer.DefaultSigAlgo(key), policy)
	if err != nil {
		return err
	}

	signer.SetDBAccessor(ca.sp.Accessor())

	ca.signer = signer
	return nil
}

// Certificate returns the certificate of the CA.
func (ca *CertificationAuthority) Certificate() (*x509.Certificate, error) {
	certPEM, err := ca.CertificatePEM()
	if err != nil {
		return nil, err
	}

	return helpers.ParseCertificatePEM(certPEM)
}

// Certificate returns the certificate of the CA in PEM encoding.
func (ca *CertificationAuthority) CertificatePEM() ([]byte, error) {
	return ca.sp.GetMetadata([]byte("cert"))
}

// selfSign creates a self-signed certificate for the CA.
func (ca *CertificationAuthority) selfSign() error {
	csrPEM, err := ca.CertificateRequestPEM()
	if err != nil {
		return err
	}

	var oidExtensionAuthorityKeyId = config.OID([]int{2, 5, 29, 35})

	keyID, err := ca.KeyID()
	if err != nil {
		return err
	}

	type authKeyId struct {
		KeyIdentifier []byte `asn1:"tag:0"`
	}
	aki, err := asn1.Marshal(authKeyId{keyID})
	if err != nil {
		return err
	}

	akiExt := signer.Extension{
		ID:       oidExtensionAuthorityKeyId,
		Critical: false,
		Value:    hex.EncodeToString(aki),
	}

	certPEM, err := ca.Issue(csrPEM, akiExt)
	if err != nil {
		return err
	}

	return ca.ImportCertificate(certPEM)
}

// ImportCertificate imports the given certificate if the CA does not
// have one.
func (ca *CertificationAuthority) ImportCertificate(certPEM []byte) error {
	if cert, _ := ca.Certificate(); cert != nil {
		return errors.New("ca cert exists")
	}

	_, err := helpers.ParseCertificatePEM(certPEM)
	if err != nil {
		return err
	}

	// TODO: Check signature and compare with original CSR.

	err = ca.sp.SetMetadata([]byte("cert"), certPEM)
	if err != nil {
		return err
	}

	return ca.initSigner()
}

// CertificateRequest returns the certificate signing request of the CA.
func (ca *CertificationAuthority) CertificateRequest() (*x509.CertificateRequest, error) {
	csrPEM, err := ca.CertificateRequestPEM()
	if err != nil {
		return nil, err
	}

	return helpers.ParseCSRPEM(csrPEM)
}

// CertificateRequestPEM returns the certificate signing request of the
// CA in PEM encoding.
func (ca *CertificationAuthority) CertificateRequestPEM() ([]byte, error) {
	return ca.sp.GetMetadata([]byte("csr"))
}

// PublicKey returns the public key from the CA certificate or CSR.
func (ca *CertificationAuthority) PublicKey() (crypto.PublicKey, error) {
	if cert, _ := ca.Certificate(); cert != nil {
		return cert.PublicKey, nil
	}

	if csr, _ := ca.CertificateRequest(); csr != nil {
		return csr.PublicKey, nil
	}

	return nil, errors.New("no valid csr in db")
}

// privateKey returns the private key as a crypto.Signer.
func (ca *CertificationAuthority) privateKey() (crypto.Signer, error) {
	pub, err := ca.PublicKey()
	if err != nil {
		return nil, err
	}

	return ca.kp.FindKeyPair(pub)
}

// KeyID returns the identifier of the signing key, which will also be
// the Authority Key Identifier (AKI) for issued certificates.
func (ca *CertificationAuthority) KeyID() ([]byte, error) {
	if cert, _ := ca.Certificate(); cert != nil {
		return cert.SubjectKeyId, nil
	}

	pub, err := ca.PublicKey()
	if err != nil {
		return nil, err
	}

	pkixPub, err := x509.MarshalPKIXPublicKey(pub)
	if err != nil {
		return nil, err
	}

	var pubKeyInfo struct {
		Algo      pkix.AlgorithmIdentifier
		BitString asn1.BitString
	}
	_, err = asn1.Unmarshal(pkixPub, &pubKeyInfo)
	if err != nil {
		return nil, err
	}

	hash := sha1.New()
	hash.Write(pubKeyInfo.BitString.Bytes)
	return hash.Sum(nil), nil
}

// Policy returns the signing policy of the CA.
func (ca *CertificationAuthority) Policy() (*config.Signing, error) {
	if ca.signer != nil {
		return ca.signer.Policy(), nil
	}

	policyJSON, err := ca.sp.GetMetadata([]byte("policy"))
	if err != nil {
		return nil, err
	}
	var policy *config.Signing
	err = json.Unmarshal(policyJSON, &policy)
	if err != nil {
		return nil, err
	}

	return policy, nil
}

// SetPolicy sets the signing policy of the CA.
func (ca *CertificationAuthority) SetPolicy(policy *config.Signing) error {
	if !policy.Valid() {
		return errors.New("invalid policy")
	}

	policyJSON, err := json.Marshal(policy)
	if err != nil {
		return err
	}
	err = ca.sp.SetMetadata([]byte("policy"), policyJSON)
	if err != nil {
		return err
	}

	if ca.signer != nil {
		ca.signer.SetPolicy(policy)
	}

	return nil
}

// Issue signs a PEM-encoded CSR and returns the certificate in PEM.
func (ca *CertificationAuthority) Issue(csrPEM []byte, exts ...signer.Extension) ([]byte, error) {
	req := signer.SignRequest{
		Request:    string(csrPEM),
		Extensions: exts,
	}

	return ca.signer.Sign(req)
}

// Revoke marks the certificate identified by its serial number revoked.
// The reasonCode is defined in RFC 5280 5.3.1.
func (ca *CertificationAuthority) Revoke(serial string, reasonCode int) error {
	keyID, err := ca.KeyID()
	if err != nil {
		return err
	}
	aki := hex.EncodeToString(keyID)

	return ca.sp.Accessor().RevokeCertificate(serial, aki, reasonCode)
}

// CRL returns a DER-encoded Certificate Revocation List, signed by the CA.
func (ca *CertificationAuthority) CRL(ttl time.Duration) ([]byte, error) {
	certs, err := ca.sp.Accessor().GetRevokedAndUnexpiredCertificates()
	if err != nil {
		return nil, err
	}

	var revokedCerts []pkix.RevokedCertificate
	for _, certRecord := range certs {
		serialInt := new(big.Int)
		serialInt.SetString(certRecord.Serial, 10)

		revokedCert := pkix.RevokedCertificate{
			SerialNumber:   serialInt,
			RevocationTime: certRecord.RevokedAt,
		}
		revokedCerts = append(revokedCerts, revokedCert)
	}

	cert, err := ca.Certificate()
	if err != nil {
		return nil, err
	}

	key, err := ca.privateKey()
	if err != nil {
		return nil, err
	}

	return cert.CreateCRL(rand.Reader, key, revokedCerts, time.Now(), time.Now().Add(ttl))
}
<|endoftext|>