<commit_before>package rpc2

import (
	"bufio"
	"bytes"
	"github.com/ugorji/go/codec"
	"io/ioutil"
	"net"
	"sync"
)

type WrapErrorFunc func(error) interface{}
type UnwrapErrorFunc func(nxt DecodeNext) (error, error)

type Transporter interface {
	RawWrite([]byte) error
	ReadByte() (byte, error)
	Decode(interface{}) error
	Encode(interface{}) error
	GetDispatcher() (Dispatcher, error)
	ReadLock()
	ReadUnlock()
}

type ConPackage struct {
	con        net.Conn
	remoteAddr net.Addr
	br         *bufio.Reader
	dec        *codec.Decoder
}

func (c *ConPackage) ReadByte() (b byte, e error) {
	return c.br.ReadByte()
}

func (c *ConPackage) Write(b []byte) (err error) {
	_, err = c.con.Write(b)
	return
}

func (c *ConPackage) Close() error {
	return c.con.Close()
}

func (c *ConPackage) GetRemoteAddr() net.Addr {
	return c.remoteAddr
}

type Transport struct {
	mh         *codec.MsgpackHandle
	cpkg       *ConPackage
	buf        *bytes.Buffer
	enc        *codec.Encoder
	mutex      *sync.Mutex
	rdlck      *sync.Mutex
	dispatcher Dispatcher
	packetizer *Packetizer
	log        LogInterface
	running    bool
	wrapError  WrapErrorFunc
}

func NewConPackage(c net.Conn, mh *codec.MsgpackHandle) *ConPackage {
	br := bufio.NewReader(c)

	return &ConPackage{
		con:        c,
		remoteAddr: c.RemoteAddr(),
		br:         br,
		dec:        codec.NewDecoder(br, mh),
	}
}

func (t *Transport) IsConnected() bool {
	t.mutex.Lock()
	ret := (t.cpkg != nil)
	t.mutex.Unlock()
	return ret
}

func (t *Transport) GetRemoteAddr() (ret net.Addr) {
	if t.cpkg != nil {
		ret = t.cpkg.GetRemoteAddr()
	}
	return
}

func NewTransport(c net.Conn, l LogFactory, wef WrapErrorFunc) *Transport {
	mh := codec.MsgpackHandle{WriteExt : true}

	buf := new(bytes.Buffer)
	ret := &Transport{
		mh:        &mh,
		cpkg:      NewConPackage(c, &mh),
		buf:       buf,
		enc:       codec.NewEncoder(buf, &mh),
		mutex:     new(sync.Mutex),
		rdlck:     new(sync.Mutex),
		wrapError: wef,
	}
	if l == nil {
		l = NewSimpleLogFactory(nil, nil)
	}
	log := l.NewLog(ret.cpkg.GetRemoteAddr())
	ret.log = log
	ret.dispatcher = NewDispatch(ret, log, wef)
	ret.packetizer = NewPacketizer(ret.dispatcher, ret)
	return ret
}

func (t *Transport) ReadLock()   { t.rdlck.Lock() }
func (t *Transport) ReadUnlock() { t.rdlck.Unlock() }

func (t *Transport) encodeToBytes(i interface{}) (v []byte, err error) {
	if err = t.enc.Encode(i); err != nil {
		return
	}
	v, _ = ioutil.ReadAll(t.buf)
	return
}

func (t *Transport) run2() (err error) {
	err = t.packetizer.Packetize()
	t.handlePacketizerFailure(err)
	return
}

func (t *Transport) handlePacketizerFailure(err error) {
	// For now, just throw everything away.  Eventually we might
	// want to make a plan for reconnecting.
	t.mutex.Lock()
	t.log.TransportError(err)
	t.running = false
	t.dispatcher.Reset()
	t.dispatcher = nil
	t.packetizer.Clear()
	t.packetizer = nil
	t.cpkg.Close()
	t.cpkg = nil
	t.mutex.Unlock()
	return
}

func (t *Transport) run(bg bool) (err error) {
	dostart := false
	t.mutex.Lock()
	if t.cpkg == nil {
		err = DisconnectedError{}
	} else if !t.running {
		dostart = true
		t.running = true
	}
	t.mutex.Unlock()
	if dostart {
		if bg {
			go t.run2()
		} else {
			err = t.run2()
		}
	}
	return
}

func (t *Transport) Encode(i interface{}) (err error) {
	var v1, v2 []byte
	if v2, err = t.encodeToBytes(i); err != nil {
		return
	}
	l := len(v2)
	if v1, err = t.encodeToBytes(l); err != nil {
		return
	}
	if err = t.RawWrite(v1); err != nil {
		return
	}
	return t.RawWrite(v2)
}

func (t *Transport) getConPackage() (ret *ConPackage, err error) {
	t.mutex.Lock()
	ret = t.cpkg
	t.mutex.Unlock()
	if ret == nil {
		err = DisconnectedError{}
	}
	return
}

func (t *Transport) ReadByte() (b byte, err error) {
	var cp *ConPackage
	if cp, err = t.getConPackage(); err == nil {
		b, err = cp.ReadByte()
	}
	return
}

func (t *Transport) Decode(i interface{}) (err error) {
	var cp *ConPackage
	if cp, err = t.getConPackage(); err == nil {
		err = cp.dec.Decode(i)
	}
	return
}

func (t *Transport) RawWrite(b []byte) (err error) {
	var cp *ConPackage
	if cp, err = t.getConPackage(); err == nil {
		err = cp.Write(b)
	}
	return err
}

func (t *Transport) GetDispatcher() (d Dispatcher, err error) {
	t.run(true)
	if !t.IsConnected() {
		err = DisconnectedError{}
	} else {
		d = t.dispatcher
	}
	return
}
<commit_msg>careful about interleaving writers<commit_after>package rpc2

import (
	"bufio"
	"bytes"
	"github.com/ugorji/go/codec"
	"io/ioutil"
	"net"
	"sync"
)

type WrapErrorFunc func(error) interface{}
type UnwrapErrorFunc func(nxt DecodeNext) (error, error)

type Transporter interface {
	RawWrite([]byte) error
	ReadByte() (byte, error)
	Decode(interface{}) error
	Encode(interface{}) error
	GetDispatcher() (Dispatcher, error)
	ReadLock()
	ReadUnlock()
}

type ConPackage struct {
	con        net.Conn
	remoteAddr net.Addr
	br         *bufio.Reader
	dec        *codec.Decoder
}

func (c *ConPackage) ReadByte() (b byte, e error) {
	return c.br.ReadByte()
}

func (c *ConPackage) Write(b []byte) (err error) {
	_, err = c.con.Write(b)
	return
}

func (c *ConPackage) Close() error {
	return c.con.Close()
}

func (c *ConPackage) GetRemoteAddr() net.Addr {
	return c.remoteAddr
}

type Transport struct {
	mh         *codec.MsgpackHandle
	cpkg       *ConPackage
	buf        *bytes.Buffer
	enc        *codec.Encoder
	mutex      *sync.Mutex
	rdlck      *sync.Mutex
	wrlck      *sync.Mutex
	dispatcher Dispatcher
	packetizer *Packetizer
	log        LogInterface
	running    bool
	wrapError  WrapErrorFunc
}

func NewConPackage(c net.Conn, mh *codec.MsgpackHandle) *ConPackage {
	br := bufio.NewReader(c)

	return &ConPackage{
		con:        c,
		remoteAddr: c.RemoteAddr(),
		br:         br,
		dec:        codec.NewDecoder(br, mh),
	}
}

func (t *Transport) IsConnected() bool {
	t.mutex.Lock()
	ret := (t.cpkg != nil)
	t.mutex.Unlock()
	return ret
}

func (t *Transport) GetRemoteAddr() (ret net.Addr) {
	if t.cpkg != nil {
		ret = t.cpkg.GetRemoteAddr()
	}
	return
}

func NewTransport(c net.Conn, l LogFactory, wef WrapErrorFunc) *Transport {
	mh := codec.MsgpackHandle{WriteExt : true}

	buf := new(bytes.Buffer)
	ret := &Transport{
		mh:        &mh,
		cpkg:      NewConPackage(c, &mh),
		buf:       buf,
		enc:       codec.NewEncoder(buf, &mh),
		mutex:     new(sync.Mutex),
		rdlck:     new(sync.Mutex),
		wrlck:     new(sync.Mutex),
		wrapError: wef,
	}
	if l == nil {
		l = NewSimpleLogFactory(nil, nil)
	}
	log := l.NewLog(ret.cpkg.GetRemoteAddr())
	ret.log = log
	ret.dispatcher = NewDispatch(ret, log, wef)
	ret.packetizer = NewPacketizer(ret.dispatcher, ret)
	return ret
}

func (t *Transport) ReadLock()   { t.rdlck.Lock() }
func (t *Transport) ReadUnlock() { t.rdlck.Unlock() }

func (t *Transport) encodeToBytes(i interface{}) (v []byte, err error) {
	if err = t.enc.Encode(i); err != nil {
		return
	}
	v, _ = ioutil.ReadAll(t.buf)
	return
}

func (t *Transport) run2() (err error) {
	err = t.packetizer.Packetize()
	t.handlePacketizerFailure(err)
	return
}

func (t *Transport) handlePacketizerFailure(err error) {
	// For now, just throw everything away.  Eventually we might
	// want to make a plan for reconnecting.
	t.mutex.Lock()
	t.log.TransportError(err)
	t.running = false
	t.dispatcher.Reset()
	t.dispatcher = nil
	t.packetizer.Clear()
	t.packetizer = nil
	t.cpkg.Close()
	t.cpkg = nil
	t.mutex.Unlock()
	return
}

func (t *Transport) run(bg bool) (err error) {
	dostart := false
	t.mutex.Lock()
	if t.cpkg == nil {
		err = DisconnectedError{}
	} else if !t.running {
		dostart = true
		t.running = true
	}
	t.mutex.Unlock()
	if dostart {
		if bg {
			go t.run2()
		} else {
			err = t.run2()
		}
	}
	return
}

func (t *Transport) Encode(i interface{}) (err error) {
	t.wrlck.Lock()
	defer t.wrlck.Unlock()

	var v1, v2 []byte
	if v2, err = t.encodeToBytes(i); err != nil {
		return
	}
	l := len(v2)
	if v1, err = t.encodeToBytes(l); err != nil {
		return
	}
	if err = t.RawWrite(v1); err != nil {
		return
	}
	return t.RawWrite(v2)
}

func (t *Transport) getConPackage() (ret *ConPackage, err error) {
	t.mutex.Lock()
	ret = t.cpkg
	t.mutex.Unlock()
	if ret == nil {
		err = DisconnectedError{}
	}
	return
}

func (t *Transport) ReadByte() (b byte, err error) {
	var cp *ConPackage
	if cp, err = t.getConPackage(); err == nil {
		b, err = cp.ReadByte()
	}
	return
}

func (t *Transport) Decode(i interface{}) (err error) {
	var cp *ConPackage
	if cp, err = t.getConPackage(); err == nil {
		err = cp.dec.Decode(i)
	}
	return
}

func (t *Transport) RawWrite(b []byte) (err error) {
	var cp *ConPackage
	if cp, err = t.getConPackage(); err == nil {
		err = cp.Write(b)
	}
	return err
}

func (t *Transport) GetDispatcher() (d Dispatcher, err error) {
	t.run(true)
	if !t.IsConnected() {
		err = DisconnectedError{}
	} else {
		d = t.dispatcher
	}
	return
}
<|endoftext|>


<commit_before>// Package mux multiplexes packets on a single socket (RFC7983)
package mux

import (
	"errors"
	"io"
	"net"
	"sync"
	"time"

	"github.com/pion/ice/v2"
	"github.com/pion/logging"
	"github.com/pion/transport/packetio"
)

// The maximum amount of data that can be buffered before returning errors.
const maxBufferSize = 1000 * 1000 // 1MB

// Config collects the arguments to mux.Mux construction into
// a single structure
type Config struct {
	Conn          net.Conn
	BufferSize    int
	LoggerFactory logging.LoggerFactory
}

// Mux allows multiplexing
type Mux struct {
	lock       sync.RWMutex
	nextConn   net.Conn
	endpoints  map[*Endpoint]MatchFunc
	bufferSize int
	closedCh   chan struct{}

	log logging.LeveledLogger
}

// NewMux creates a new Mux
func NewMux(config Config) *Mux {
	m := &Mux{
		nextConn:   config.Conn,
		endpoints:  make(map[*Endpoint]MatchFunc),
		bufferSize: config.BufferSize,
		closedCh:   make(chan struct{}),
		log:        config.LoggerFactory.NewLogger("mux"),
	}

	go m.readLoop()

	return m
}

// NewEndpoint creates a new Endpoint
func (m *Mux) NewEndpoint(f MatchFunc) *Endpoint {
	e := &Endpoint{
		mux:    m,
		buffer: packetio.NewBuffer(),
	}

	// Set a maximum size of the buffer in bytes.
	e.buffer.SetLimitSize(maxBufferSize)

	m.lock.Lock()
	m.endpoints[e] = f
	m.lock.Unlock()

	return e
}

// RemoveEndpoint removes an endpoint from the Mux
func (m *Mux) RemoveEndpoint(e *Endpoint) {
	m.lock.Lock()
	defer m.lock.Unlock()
	delete(m.endpoints, e)
}

// Close closes the Mux and all associated Endpoints.
func (m *Mux) Close() error {
	m.lock.Lock()
	for e := range m.endpoints {
		err := e.close()
		if err != nil {
			return err
		}

		delete(m.endpoints, e)
	}
	m.lock.Unlock()

	err := m.nextConn.Close()
	if err != nil {
		return err
	}

	// Wait for readLoop to end
	<-m.closedCh

	return nil
}

func (m *Mux) readLoop() {
	defer func() {
		close(m.closedCh)
	}()

	buf := make([]byte, m.bufferSize)
	for {
		n, err := m.nextConn.Read(buf)
		switch {
		case errors.Is(err, io.EOF), errors.Is(err, ice.ErrClosed):
			return
		case errors.Is(err, io.ErrShortBuffer), errors.Is(err, packetio.ErrTimeout):
			m.log.Errorf("mux: failed to read from packetio.Buffer %s\n", err.Error())
			continue
		case err != nil:
			m.log.Errorf("mux: ending readLoop packetio.Buffer error %s\n", err.Error())
			return
		}

		if err = m.dispatch(buf[:n]); err != nil {
			m.log.Errorf("mux: ending readLoop dispatch error %s\n", err.Error())
			return
		}
	}
}

func (m *Mux) dispatch(buf []byte) error {
	var endpoint *Endpoint

	m.lock.Lock()
	for e, f := range m.endpoints {
		if f(buf) {
			endpoint = e
			break
		}
	}
	m.lock.Unlock()

	if endpoint == nil {
		if len(buf) > 0 {
			m.log.Warnf("Warning: mux: no endpoint for packet starting with %d\n", buf[0])
		} else {
			m.log.Warnf("Warning: mux: no endpoint for zero length packet")
		}
		return nil
	}

	// Attempt again later when encountering a temporarily full buffer (fix for #2152)
	var err error
	for {
		_, err = endpoint.buffer.Write(buf)
		if !errors.Is(err, packetio.ErrFull) {
			break
		}
		time.Sleep(20 * time.Millisecond)
	}

	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Unlock on error in internal/mux<commit_after>// Package mux multiplexes packets on a single socket (RFC7983)
package mux

import (
	"errors"
	"io"
	"net"
	"sync"
	"time"

	"github.com/pion/ice/v2"
	"github.com/pion/logging"
	"github.com/pion/transport/packetio"
)

// The maximum amount of data that can be buffered before returning errors.
const maxBufferSize = 1000 * 1000 // 1MB

// Config collects the arguments to mux.Mux construction into
// a single structure
type Config struct {
	Conn          net.Conn
	BufferSize    int
	LoggerFactory logging.LoggerFactory
}

// Mux allows multiplexing
type Mux struct {
	lock       sync.RWMutex
	nextConn   net.Conn
	endpoints  map[*Endpoint]MatchFunc
	bufferSize int
	closedCh   chan struct{}

	log logging.LeveledLogger
}

// NewMux creates a new Mux
func NewMux(config Config) *Mux {
	m := &Mux{
		nextConn:   config.Conn,
		endpoints:  make(map[*Endpoint]MatchFunc),
		bufferSize: config.BufferSize,
		closedCh:   make(chan struct{}),
		log:        config.LoggerFactory.NewLogger("mux"),
	}

	go m.readLoop()

	return m
}

// NewEndpoint creates a new Endpoint
func (m *Mux) NewEndpoint(f MatchFunc) *Endpoint {
	e := &Endpoint{
		mux:    m,
		buffer: packetio.NewBuffer(),
	}

	// Set a maximum size of the buffer in bytes.
	e.buffer.SetLimitSize(maxBufferSize)

	m.lock.Lock()
	m.endpoints[e] = f
	m.lock.Unlock()

	return e
}

// RemoveEndpoint removes an endpoint from the Mux
func (m *Mux) RemoveEndpoint(e *Endpoint) {
	m.lock.Lock()
	defer m.lock.Unlock()
	delete(m.endpoints, e)
}

// Close closes the Mux and all associated Endpoints.
func (m *Mux) Close() error {
	m.lock.Lock()
	for e := range m.endpoints {
		if err := e.close(); err != nil {
			m.lock.Unlock()
			return err
		}

		delete(m.endpoints, e)
	}
	m.lock.Unlock()

	err := m.nextConn.Close()
	if err != nil {
		return err
	}

	// Wait for readLoop to end
	<-m.closedCh

	return nil
}

func (m *Mux) readLoop() {
	defer func() {
		close(m.closedCh)
	}()

	buf := make([]byte, m.bufferSize)
	for {
		n, err := m.nextConn.Read(buf)
		switch {
		case errors.Is(err, io.EOF), errors.Is(err, ice.ErrClosed):
			return
		case errors.Is(err, io.ErrShortBuffer), errors.Is(err, packetio.ErrTimeout):
			m.log.Errorf("mux: failed to read from packetio.Buffer %s\n", err.Error())
			continue
		case err != nil:
			m.log.Errorf("mux: ending readLoop packetio.Buffer error %s\n", err.Error())
			return
		}

		if err = m.dispatch(buf[:n]); err != nil {
			m.log.Errorf("mux: ending readLoop dispatch error %s\n", err.Error())
			return
		}
	}
}

func (m *Mux) dispatch(buf []byte) error {
	var endpoint *Endpoint

	m.lock.Lock()
	for e, f := range m.endpoints {
		if f(buf) {
			endpoint = e
			break
		}
	}
	m.lock.Unlock()

	if endpoint == nil {
		if len(buf) > 0 {
			m.log.Warnf("Warning: mux: no endpoint for packet starting with %d\n", buf[0])
		} else {
			m.log.Warnf("Warning: mux: no endpoint for zero length packet")
		}
		return nil
	}

	// Attempt again later when encountering a temporarily full buffer (fix for #2152)
	var err error
	for {
		_, err = endpoint.buffer.Write(buf)
		if !errors.Is(err, packetio.ErrFull) {
			break
		}
		time.Sleep(20 * time.Millisecond)
	}

	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// Convert Wikipedia XML dump to JSON or extract categories
// Example inputs:
// wikidata: http://dumps.wikimedia.org/wikidatawiki/20140612/wikidatawiki-20140612-pages-articles.xml.bz2
// wikipedia:  http://dumps.wikimedia.org/huwiki/latest/huwiki-latest-pages-articles.xml.bz2
package main

import (
	"bufio"
	"encoding/json"
	"encoding/xml"
	"flag"
	"fmt"
	"io"
	"log"
	"net/url"
	"os"
	"regexp"
	"runtime"
	"strings"
)

const AppVersion = "1.1.1"

// Here is an example article from the Wikipedia XML dump
//
// <page>
// 	<title>Apollo 11</title>
//      <redirect title="Foo bar" />
// 	...
// 	<revision>
// 	...
// 	  <text xml:space="preserve">
// 	  {{Infobox Space mission
// 	  |mission_name=<!--See above-->
// 	  |insignia=Apollo_11_insignia.png
// 	...
// 	  </text>
// 	</revision>
// </page>
//
// Note how the tags on the fields of Page and Redirect below
// describe the XML schema structure.

type Redirect struct {
	Title string `xml:"title,attr" json:"title"`
}

// A page as it occurs on Wikipedia
type Page struct {
	Title          string   `xml:"title" json:"title"`
	CanonicalTitle string   `xml:"ctitle" json:"ctitle"`
	Redir          Redirect `xml:"redirect" json:"redirect"`
	Text           string   `xml:"revision>text" json:"text"`
}

// A page as it occurs on Wikidata, content will be turned from a string
// into a substructure with -d switch
type WikidataPage struct {
	Title          string      `xml:"title" json:"title"`
	CanonicalTitle string      `xml:"ctitle" json:"ctitle"`
	Redir          Redirect    `xml:"redirect" json:"redirect"`
	Content        interface{} `json:"content"`
}

func CanonicalizeTitle(title string) string {
	can := strings.ToLower(title)
	can = strings.Replace(can, " ", "_", -1)
	can = url.QueryEscape(can)
	return can
}

func categoryExtractor(in chan *Page,
	out chan *string,
	filter *regexp.Regexp,
	categoryPattern *regexp.Regexp) {
	var pp *Page
	for {
		// get the page pointer
		pp = <-in
		if pp == nil {
			break
		}
		// get the page
		p := *pp

		// do some stuff with the page
		p.CanonicalTitle = CanonicalizeTitle(p.Title)
		m := filter.MatchString(p.CanonicalTitle)
		if !m && p.Redir.Title == "" {

			// specific to category extraction
			result := categoryPattern.FindAllStringSubmatch(p.Text, -1)
			for _, value := range result {
				// replace anything after a |
				category := strings.TrimSpace(value[1])
				firstIndex := strings.Index(category, "|")
				if firstIndex != -1 {
					category = category[0:firstIndex]
				}

				line := fmt.Sprintf("%s\t%s", p.Title, category)
				out <- &line
			}
		}
	}
}

func authorityDataExtractor(in chan *Page,
	out chan *string,
	filter *regexp.Regexp,
	authorityDataPattern *regexp.Regexp) {
	var pp *Page
	for {
		// get the page pointer
		pp = <-in
		if pp == nil {
			break
		}
		// get the page
		p := *pp

		// do some stuff with the page
		p.CanonicalTitle = CanonicalizeTitle(p.Title)
		m := filter.MatchString(p.CanonicalTitle)
		if !m && p.Redir.Title == "" {

			// specific to category extraction
			result := authorityDataPattern.FindString(p.Text)
			if result != "" {
				// https://cdn.mediacru.sh/JsdjtGoLZBcR.png
				result = strings.Replace(result, "\t", "", -1)
				// fmt.Printf("%s\t%s\n", p.Title, result)
				line := fmt.Sprintf("%s\t%s", p.Title, result)
				out <- &line
			}
		}
	}
}

func wikidataEncoder(in chan *Page,
	out chan *string,
	filter *regexp.Regexp) {

	var container interface{}
	var pp *Page

	for {
		// get the page pointer
		pp = <-in
		if pp == nil {
			break
		}
		// get the page
		p := *pp

		// do some stuff with the page
		p.CanonicalTitle = CanonicalizeTitle(p.Title)
		m := filter.MatchString(p.CanonicalTitle)
		if !m && p.Redir.Title == "" {
			dec := json.NewDecoder(strings.NewReader(p.Text))
			dec.UseNumber()

			if err := dec.Decode(&container); err == io.EOF {
				break
			} else if err != nil {
				fmt.Fprintf(os.Stderr, "%s\n", err)
				continue
			}

			parsed := WikidataPage{Title: p.Title,
				CanonicalTitle: p.CanonicalTitle,
				Content:        container,
				Redir:          p.Redir}

			b, err := json.Marshal(parsed)
			if err != nil {
				fmt.Fprintf(os.Stderr, "%s\n", err)
				os.Exit(2)
			}
			// fmt.Println(string(b))
			line := string(b)
			out <- &line
		}
	}
}

func vanillaConverter(in chan *Page,
	out chan *string,
	filter *regexp.Regexp) {
	var pp *Page
	for {
		// get the page pointer
		pp = <-in
		if pp == nil {
			break
		}
		// get the page
		p := *pp

		// do some stuff with the page
		p.CanonicalTitle = CanonicalizeTitle(p.Title)
		m := filter.MatchString(p.CanonicalTitle)
		if !m && p.Redir.Title == "" {

			b, err := json.Marshal(p)
			if err != nil {
				os.Exit(2)
			}
			line := string(b)
			out <- &line
		}
	}
}

func collect(lines chan *string) {
	for line := range lines {
		fmt.Println(*line)
	}
}

// Collect output and write to file
func FileCollector(lines chan *string, filename string) {
	output, err := os.Create(filename)
	if err != nil {
		panic(err)
	}
	// close fo on exit and check for its returned error
	defer func() {
		if err := output.Close(); err != nil {
			panic(err)
		}
	}()
	// 4M buffer size
	w := bufio.NewWriter(output)
	for line := range lines {
		_, err = w.WriteString(*line + "\n")
		if err != nil {
			panic(err)
		}
	}
	w.Flush()
}

func main() {

	version := flag.Bool("v", false, "prints current version and exits")
	extractCategories := flag.String("c", "", "only extract categories TSV(page, category), argument is the prefix, e.g. Kategorie or Category, ... ")
	extractAuthorityData := flag.String("a", "", "only extract authority data (Normdaten, Authority control, ...)")
	decodeWikiData := flag.Bool("d", false, "decode the text key value")
	numWorkers := flag.Int("w", runtime.NumCPU(), "number of workers")

	filter, _ := regexp.Compile("^file:.*|^talk:.*|^special:.*|^wikipedia:.*|^wiktionary:.*|^user:.*|^user_talk:.*")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "\nExtract and convert things from wikipedia/wikidata XML dumps.\n")
		fmt.Fprintf(os.Stderr, "\nVersion: %s\n\n", AppVersion)
		flag.PrintDefaults()
	}

	flag.Parse()

	if *extractCategories != "" && *extractAuthorityData != "" {
		fmt.Fprintln(os.Stderr, "it's either -a or -c")
		os.Exit(1)
	}

	if *version {
		fmt.Println(AppVersion)
		os.Exit(0)
	}

	if flag.NArg() < 1 {
		flag.Usage()
		os.Exit(1)
	}

	runtime.GOMAXPROCS(*numWorkers)

	inputFile := flag.Args()[0]

	xmlFile, err := os.Open(inputFile)
	if err != nil {
		fmt.Println("Error opening file:", err)
		return
	}
	defer xmlFile.Close()

	// xml decoder
	decoder := xml.NewDecoder(xmlFile)
	var inElement string

	// category pattern depends on the language, e.g. Kategorie or Category, ...
	categoryPattern := regexp.MustCompile(`\[\[` + *extractCategories + `:([^\[]+)\]\]`)
	// Authority data (German only for now)
	authorityDataPattern := regexp.MustCompile(`(?mi){{` + *extractAuthorityData + `[^}]*}}`)

	// the parsed XML pages channel
	pages := make(chan *Page)
	// the strings output channel
	lines := make(chan *string)

	// start the collector
	go collect(lines)

	// start some appropriate workers
	for i := 0; i < *numWorkers; i++ {
		if *extractCategories != "" {
			go categoryExtractor(pages, lines, filter, categoryPattern)
		} else if *extractAuthorityData != "" {
			go authorityDataExtractor(pages, lines, filter, authorityDataPattern)
		} else if *decodeWikiData {
			go wikidataEncoder(pages, lines, filter)
		} else {
			go vanillaConverter(pages, lines, filter)
		}
	}

	for {
		// Read tokens from the XML document in a stream.
		t, _ := decoder.Token()
		if t == nil {
			break
		}
		// Inspect the type of the token just read.
		switch se := t.(type) {
		case xml.StartElement:
			// If we just read a StartElement token
			inElement = se.Name.Local
			// ...and its name is "page"
			if inElement == "page" {
				var p Page
				// decode a whole chunk of following XML into the
				// variable p which is a Page (se above)
				decoder.DecodeElement(&p, &se)
				pages <- &p
			}
		default:
		}
	}

	// kill workers
	for n := 0; n < *numWorkers; n++ {
		pages <- nil
	}
	close(lines)

}
<commit_msg>added support for pprof and output files<commit_after>// Convert Wikipedia XML dump to JSON or extract categories
// Example inputs:
// wikidata: http://dumps.wikimedia.org/wikidatawiki/20140612/wikidatawiki-20140612-pages-articles.xml.bz2
// wikipedia:  http://dumps.wikimedia.org/huwiki/latest/huwiki-latest-pages-articles.xml.bz2
package main

import (
	"bufio"
	"encoding/json"
	"encoding/xml"
	"flag"
	"fmt"
	"io"
	"log"
	"net/url"
	"os"
	"regexp"
	"runtime"
	"runtime/pprof"
	"strings"
)

const AppVersion = "1.1.1"

// Here is an example article from the Wikipedia XML dump
//
// <page>
// 	<title>Apollo 11</title>
//      <redirect title="Foo bar" />
// 	...
// 	<revision>
// 	...
// 	  <text xml:space="preserve">
// 	  {{Infobox Space mission
// 	  |mission_name=<!--See above-->
// 	  |insignia=Apollo_11_insignia.png
// 	...
// 	  </text>
// 	</revision>
// </page>
//
// Note how the tags on the fields of Page and Redirect below
// describe the XML schema structure.

type Redirect struct {
	Title string `xml:"title,attr" json:"title"`
}

// A page as it occurs on Wikipedia
type Page struct {
	Title          string   `xml:"title" json:"title"`
	CanonicalTitle string   `xml:"ctitle" json:"ctitle"`
	Redir          Redirect `xml:"redirect" json:"redirect"`
	Text           string   `xml:"revision>text" json:"text"`
}

// A page as it occurs on Wikidata, content will be turned from a string
// into a substructure with -d switch
type WikidataPage struct {
	Title          string      `xml:"title" json:"title"`
	CanonicalTitle string      `xml:"ctitle" json:"ctitle"`
	Redir          Redirect    `xml:"redirect" json:"redirect"`
	Content        interface{} `json:"content"`
}

func CanonicalizeTitle(title string) string {
	can := strings.ToLower(title)
	can = strings.Replace(can, " ", "_", -1)
	can = url.QueryEscape(can)
	return can
}

// category extraction worker
func CategoryExtractor(in chan *Page,
	out chan *string,
	filter *regexp.Regexp,
	categoryPattern *regexp.Regexp) {
	var pp *Page
	for {
		// get the page pointer
		pp = <-in
		if pp == nil {
			break
		}
		// get the page
		p := *pp

		// do some stuff with the page
		p.CanonicalTitle = CanonicalizeTitle(p.Title)
		m := filter.MatchString(p.CanonicalTitle)
		if !m && p.Redir.Title == "" {

			// specific to category extraction
			result := categoryPattern.FindAllStringSubmatch(p.Text, -1)
			for _, value := range result {
				// replace anything after a |
				category := strings.TrimSpace(value[1])
				firstIndex := strings.Index(category, "|")
				if firstIndex != -1 {
					category = category[0:firstIndex]
				}

				line := fmt.Sprintf("%s\t%s", p.Title, category)
				out <- &line
			}
		}
	}
}

// authority data extraction worker
func AuthorityDataExtractor(in chan *Page,
	out chan *string,
	filter *regexp.Regexp,
	authorityDataPattern *regexp.Regexp) {
	var pp *Page
	for {
		// get the page pointer
		pp = <-in
		if pp == nil {
			break
		}
		// get the page
		p := *pp

		// do some stuff with the page
		p.CanonicalTitle = CanonicalizeTitle(p.Title)
		m := filter.MatchString(p.CanonicalTitle)
		if !m && p.Redir.Title == "" {

			// specific to category extraction
			result := authorityDataPattern.FindString(p.Text)
			if result != "" {
				// https://cdn.mediacru.sh/JsdjtGoLZBcR.png
				result = strings.Replace(result, "\t", "", -1)
				// fmt.Printf("%s\t%s\n", p.Title, result)
				line := fmt.Sprintf("%s\t%s", p.Title, result)
				out <- &line
			}
		}
	}
}

// wikidata to json worker
func WikidataEncoder(in chan *Page,
	out chan *string,
	filter *regexp.Regexp) {

	var container interface{}
	var pp *Page

	for {
		// get the page pointer
		pp = <-in
		if pp == nil {
			break
		}
		// get the page
		p := *pp

		// do some stuff with the page
		p.CanonicalTitle = CanonicalizeTitle(p.Title)
		m := filter.MatchString(p.CanonicalTitle)
		if !m && p.Redir.Title == "" {
			dec := json.NewDecoder(strings.NewReader(p.Text))
			dec.UseNumber()

			if err := dec.Decode(&container); err == io.EOF {
				break
			} else if err != nil {
				fmt.Fprintf(os.Stderr, "%s\n", err)
				continue
			}

			parsed := WikidataPage{Title: p.Title,
				CanonicalTitle: p.CanonicalTitle,
				Content:        container,
				Redir:          p.Redir}

			b, err := json.Marshal(parsed)
			if err != nil {
				fmt.Fprintf(os.Stderr, "%s\n", err)
				os.Exit(2)
			}
			// fmt.Println(string(b))
			line := string(b)
			out <- &line
		}
	}
}

// just XML to json
func VanillaConverter(in chan *Page,
	out chan *string,
	filter *regexp.Regexp) {
	var pp *Page
	for {
		// get the page pointer
		pp = <-in
		if pp == nil {
			break
		}
		// get the page
		p := *pp

		// do some stuff with the page
		p.CanonicalTitle = CanonicalizeTitle(p.Title)
		m := filter.MatchString(p.CanonicalTitle)
		if !m && p.Redir.Title == "" {

			b, err := json.Marshal(p)
			if err != nil {
				os.Exit(2)
			}
			line := string(b)
			out <- &line
		}
	}
}

// Collect output and write to Stdout
func StdoutCollector(lines chan *string) {
	for line := range lines {
		fmt.Println(*line)
	}
}

// Collect output and write to file
func FileCollector(lines chan *string, filename string) {
	output, err := os.Create(filename)
	if err != nil {
		panic(err)
	}
	// close fo on exit and check for its returned error
	defer func() {
		if err := output.Close(); err != nil {
			panic(err)
		}
	}()
	// 4M buffer size
	w := bufio.NewWriter(output)
	for line := range lines {
		_, err = w.WriteString(*line + "\n")
		if err != nil {
			panic(err)
		}
	}
	w.Flush()
}

func main() {

	version := flag.Bool("v", false, "prints current version and exits")
	extractCategories := flag.String("c", "", "only extract categories TSV(page, category), argument is the prefix, e.g. Kategorie or Category, ... ")
	extractAuthorityData := flag.String("a", "", "only extract authority data (Normdaten, Authority control, ...)")
	decodeWikiData := flag.Bool("d", false, "decode the text key value")
	numWorkers := flag.Int("w", runtime.NumCPU(), "number of workers")
	outputFilename := flag.String("o", "", "write output to file (or stdout, if empty)")
	cpuprofile := flag.String("cpuprofile", "", "write cpu profile to file")

	filter, _ := regexp.Compile("^file:.*|^talk:.*|^special:.*|^wikipedia:.*|^wiktionary:.*|^user:.*|^user_talk:.*")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "\nExtract and convert things from wikipedia/wikidata XML dumps.\n")
		fmt.Fprintf(os.Stderr, "\nVersion: %s\n\n", AppVersion)
		flag.PrintDefaults()
	}

	flag.Parse()

	if *extractCategories != "" && *extractAuthorityData != "" {
		fmt.Fprintln(os.Stderr, "it's either -a or -c")
		os.Exit(1)
	}

	if *version {
		fmt.Println(AppVersion)
		os.Exit(0)
	}

	if flag.NArg() < 1 {
		flag.Usage()
		os.Exit(1)
	}

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	runtime.GOMAXPROCS(*numWorkers)

	inputFile := flag.Args()[0]

	xmlFile, err := os.Open(inputFile)
	if err != nil {
		fmt.Println("Error opening file:", err)
		return
	}
	defer xmlFile.Close()

	// xml decoder
	decoder := xml.NewDecoder(xmlFile)
	var inElement string

	// category pattern depends on the language, e.g. Kategorie or Category, ...
	categoryPattern := regexp.MustCompile(`\[\[` + *extractCategories + `:([^\[]+)\]\]`)
	// Authority data (German only for now)
	authorityDataPattern := regexp.MustCompile(`(?mi){{` + *extractAuthorityData + `[^}]*}}`)

	// the parsed XML pages channel
	pages := make(chan *Page)
	// the strings output channel
	lines := make(chan *string)

	// start the collector
	if *outputFilename != "" {
		go FileCollector(lines, *outputFilename)
	} else {
		go StdoutCollector(lines)
	}

	// start some appropriate workers
	for i := 0; i < *numWorkers; i++ {
		if *extractCategories != "" {
			go CategoryExtractor(pages, lines, filter, categoryPattern)
		} else if *extractAuthorityData != "" {
			go AuthorityDataExtractor(pages, lines, filter, authorityDataPattern)
		} else if *decodeWikiData {
			go WikidataEncoder(pages, lines, filter)
		} else {
			go VanillaConverter(pages, lines, filter)
		}
	}

	for {
		// Read tokens from the XML document in a stream.
		t, _ := decoder.Token()
		if t == nil {
			break
		}
		// Inspect the type of the token just read.
		switch se := t.(type) {
		case xml.StartElement:
			// If we just read a StartElement token
			inElement = se.Name.Local
			// ...and its name is "page"
			if inElement == "page" {
				var p Page
				// decode a whole chunk of following XML into the
				// variable p which is a Page (se above)
				decoder.DecodeElement(&p, &se)
				pages <- &p
			}
		default:
		}
	}

	// kill workers
	for n := 0; n < *numWorkers; n++ {
		pages <- nil
	}
	// close the output channel
	close(lines)
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestGetStateFile(t *testing.T) {
	sPath := getStateFile("/var/lib", "C:/Windows/hoge")
	assert.Equal(t, sPath, "/var/lib/C/Windows/hoge", "drive letter should be cared")

	sPath = getStateFile("/var/lib", "/linux/hoge")
	assert.Equal(t, sPath, "/var/lib/linux/hoge", "drive letter should be cared")
}

func TestWriteBytesToSkip(t *testing.T) {
	f := ".tmp/fuga/piyo"
	err := writeBytesToSkip(f, 15)
	assert.Equal(t, err, nil, "err should be nil")

	skipBytes, err := getBytesToSkip(f)
	assert.Equal(t, err, nil, "err should be nil")
	assert.Equal(t, skipBytes, int64(15))
}

func TestSearchReader(t *testing.T) {
	dir, err := ioutil.TempDir("", "check-event-log-test")
	if err != nil {
		t.Fatalf("TempDir failed: %s", err)
	}
	defer os.RemoveAll(dir)

	opts := &logOpts{
		StateDir: dir,
		LogFile:  filepath.Join(dir, "dummy"),
		Pattern:  `FATAL`,
	}
	opts.prepare()

	content := `FATAL 11
OK
FATAL 22
Fatal
`
	r := strings.NewReader(content)
	warnNum, critNum, readBytes, errLines, err := opts.searchReader(r)

	assert.Equal(t, int64(2), warnNum, "warnNum should be 2")
	assert.Equal(t, int64(2), critNum, "critNum should be 2")
	assert.Equal(t, "FATAL 11\nFATAL 22\n", errLines, "invalid errLines")
	assert.Equal(t, int64(len(content)), readBytes, "readBytes should be 26")
}

func TestRun(t *testing.T) {
	dir, err := ioutil.TempDir("", "check-event-log-test")
	if err != nil {
		t.Errorf("something went wrong")
	}
	defer os.RemoveAll(dir)

	logf := filepath.Join(dir, "dummy")
	fh, _ := os.Create(logf)
	defer fh.Close()

	ptn := `FATAL`
	opts, _ := parseArgs([]string{"-s", dir, "-f", logf, "-p", ptn})
	opts.prepare()

	stateFile := getStateFile(opts.StateDir, logf)

	bytes, _ := getBytesToSkip(stateFile)
	assert.Equal(t, int64(0), bytes, "something went wrong")

	testEmpty := func() {
		w, c, errLines, err := opts.searchLog(logf)
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(0), w, "something went wrong")
		assert.Equal(t, int64(0), c, "something went wrong")
		assert.Equal(t, "", errLines, "something went wrong")

		bytes, _ = getBytesToSkip(stateFile)
		assert.Equal(t, int64(0), bytes, "something went wrong")
	}
	testEmpty()

	l1 := "FATAL\nFATAL\n"
	test2Line := func() {
		fh.WriteString(l1)
		w, c, errLines, err := opts.searchLog(logf)
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(2), w, "something went wrong")
		assert.Equal(t, int64(2), c, "something went wrong")
		assert.Equal(t, l1, errLines, "something went wrong")

		bytes, _ = getBytesToSkip(stateFile)
		assert.Equal(t, int64(len(l1)), bytes, "something went wrong")
	}
	test2Line()

	testReadAgain := func() {
		w, c, errLines, err := opts.searchLog(logf)
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(0), w, "something went wrong")
		assert.Equal(t, int64(0), c, "something went wrong")
		assert.Equal(t, "", errLines, "something went wrong")

		bytes, _ = getBytesToSkip(stateFile)
		assert.Equal(t, int64(len(l1)), bytes, "something went wrong")
	}
	testReadAgain()

	l2 := "SUCCESS\n"
	testRecover := func() {
		fh.WriteString(l2)
		w, c, errLines, err := opts.searchLog(logf)
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(0), w, "something went wrong")
		assert.Equal(t, int64(0), c, "something went wrong")
		assert.Equal(t, "", errLines, "something went wrong")

		bytes, _ = getBytesToSkip(stateFile)
		assert.Equal(t, int64(len(l1)+len(l2)), bytes, "something went wrong")
	}
	testRecover()

	testSuccessAgain := func() {
		w, c, errLines, err := opts.searchLog(logf)
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(0), w, "something went wrong")
		assert.Equal(t, int64(0), c, "something went wrong")
		assert.Equal(t, "", errLines, "something went wrong")

		bytes, _ = getBytesToSkip(stateFile)
		assert.Equal(t, int64(len(l1)+len(l2)), bytes, "something went wrong")
	}
	testSuccessAgain()

	testErrorAgain := func() {
		fh.WriteString(l1)
		w, c, errLines, err := opts.searchLog(logf)
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(2), w, "something went wrong")
		assert.Equal(t, int64(2), c, "something went wrong")
		assert.Equal(t, l1, errLines, "something went wrong")

		bytes, _ = getBytesToSkip(stateFile)
		assert.Equal(t, int64(len(l1)*2+len(l2)), bytes, "something went wrong")
	}
	testErrorAgain()

	testRecoverAgain := func() {
		fh.WriteString(l2)
		w, c, errLines, err := opts.searchLog(logf)
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(0), w, "something went wrong")
		assert.Equal(t, int64(0), c, "something went wrong")
		assert.Equal(t, "", errLines, "something went wrong")

		bytes, _ = getBytesToSkip(stateFile)
		assert.Equal(t, int64(len(l1)*2+len(l2)*2), bytes, "something went wrong")
	}
	testRecoverAgain()

	testRotate := func() {
		fh.Close()
		os.Remove(logf)
		fh, _ = os.Create(logf)

		fh.WriteString(l2)
		w, c, errLines, err := opts.searchLog(logf)
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(0), w, "something went wrong")
		assert.Equal(t, int64(0), c, "something went wrong")
		assert.Equal(t, "", errLines, "something went wrong")

		bytes, _ = getBytesToSkip(stateFile)
		assert.Equal(t, int64(len(l2)), bytes, "something went wrong")
	}
	testRotate()
}

func TestRunWithMiddleOfLine(t *testing.T) {
	dir, err := ioutil.TempDir("", "check-event-log-test")
	if err != nil {
		t.Errorf("something went wrong")
	}
	defer os.RemoveAll(dir)

	logf := filepath.Join(dir, "dummy")
	fh, _ := os.Create(logf)
	defer fh.Close()

	ptn := `FATAL`
	opts, _ := parseArgs([]string{"-s", dir, "-f", logf, "-p", ptn})
	opts.prepare()

	stateFile := getStateFile(opts.StateDir, logf)

	bytes, _ := getBytesToSkip(stateFile)
	assert.Equal(t, int64(0), bytes, "something went wrong")

	testMiddleOfLine := func() {
		fh.WriteString("FATA")
		w, c, errLines, err := opts.searchLog(logf)
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(0), w, "something went wrong")
		assert.Equal(t, int64(0), c, "something went wrong")
		assert.Equal(t, "", errLines, "something went wrong")

		bytes, _ = getBytesToSkip(stateFile)
		assert.Equal(t, int64(0), bytes, "something went wrong")
	}
	testMiddleOfLine()

	testFail := func() {
		fh.WriteString("L\nSUCC")
		w, c, errLines, err := opts.searchLog(logf)
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(1), w, "something went wrong")
		assert.Equal(t, int64(1), c, "something went wrong")
		assert.Equal(t, "FATAL\n", errLines, "something went wrong")

		bytes, _ = getBytesToSkip(stateFile)
		assert.Equal(t, int64(len("FATAL\n")), bytes, "something went wrong")
	}
	testFail()
}

func TestRunWithNoState(t *testing.T) {
	dir, err := ioutil.TempDir("", "check-event-log-test")
	if err != nil {
		t.Errorf("something went wrong")
	}
	defer os.RemoveAll(dir)

	logf := filepath.Join(dir, "dummy")
	fh, _ := os.Create(logf)
	defer fh.Close()

	ptn := `FATAL`
	opts, _ := parseArgs([]string{"-s", dir, "-f", logf, "-p", ptn, "--no-state"})
	opts.prepare()

	fatal := "FATAL\n"
	test2Line := func() {
		fh.WriteString(fatal)
		fh.WriteString(fatal)
		w, c, errLines, err := opts.searchLog(logf)
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(2), w, "something went wrong")
		assert.Equal(t, int64(2), c, "something went wrong")
		assert.Equal(t, strings.Repeat(fatal, 2), errLines, "something went wrong")
	}
	test2Line()

	test1LineAgain := func() {
		fh.WriteString(fatal)
		w, c, errLines, err := opts.searchLog(logf)
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(3), w, "something went wrong")
		assert.Equal(t, int64(3), c, "something went wrong")
		assert.Equal(t, strings.Repeat(fatal, 3), errLines, "something went wrong")
	}
	test1LineAgain()
}

func TestSearchReaderWithLevel(t *testing.T) {
	dir, err := ioutil.TempDir("", "check-event-log-test")
	if err != nil {
		t.Errorf("something went wrong")
	}
	defer os.RemoveAll(dir)

	logf := filepath.Join(dir, "dummy")
	ptn := `FATAL level:([0-9]+)`
	opts, _ := parseArgs([]string{"-s", dir, "-f", logf, "-i", "-p", ptn, "--critical-level=17", "--warning-level=11"})
	if !reflect.DeepEqual(&logOpts{
		StateDir:        dir,
		LogFile:         filepath.Join(dir, "dummy"),
		CaseInsensitive: true,
		Pattern:         `FATAL level:([0-9]+)`,
		WarnLevel:       11,
		CritLevel:       17,
	}, opts) {
		t.Errorf("something went wrong")
	}
	opts.prepare()

	content := `FATAL level:11
OK
FATAL level:22
Fatal level:17
`
	r := strings.NewReader(content)
	warnNum, critNum, readBytes, errLines, err := opts.searchReader(r)

	assert.Equal(t, int64(2), warnNum, "warnNum should be 2")
	assert.Equal(t, int64(1), critNum, "critNum should be 1")
	assert.Equal(t, "FATAL level:22\nFatal level:17\n", errLines, "invalid errLines")
	assert.Equal(t, int64(len(content)), readBytes, "readBytes should be 26")
}
<commit_msg>add tests<commit_after>package main

import (
	"io/ioutil"
	"os"
	"testing"

	"github.com/go-ole/go-ole"
	"github.com/go-ole/go-ole/oleutil"
	"github.com/stretchr/testify/assert"
)

func TestGetStateFile(t *testing.T) {
	sPath := getStateFile("/var/lib", "C:/Windows/hoge")
	assert.Equal(t, sPath, "/var/lib/C/Windows/hoge", "drive letter should be cared")

	sPath = getStateFile("/var/lib", "/linux/hoge")
	assert.Equal(t, sPath, "/var/lib/linux/hoge", "drive letter should be cared")
}

func TestWriteLastOffset(t *testing.T) {
	f := ".tmp/fuga/piyo"
	err := writeLastOffset(f, 15)
	assert.Equal(t, err, nil, "err should be nil")

	recordNumber, err := getLastOffset(f)
	assert.Equal(t, err, nil, "err should be nil")
	assert.Equal(t, recordNumber, int64(15))
}

func raiseEvent(t *testing.T, typ int, msg string) {
	ole.CoInitialize(0)
	defer ole.CoUninitialize()

	unk, err := oleutil.CreateObject("Wscript.Shell")
	if err != nil {
		t.Fatal(err)
	}
	disp, err := unk.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		t.Fatal(err)
	}
	oleutil.MustCallMethod(disp, "LogEvent", typ, msg)
}

func TestRun(t *testing.T) {
	dir, err := ioutil.TempDir("", "check-event-log-test")
	if err != nil {
		t.Errorf("something went wrong")
	}
	defer os.RemoveAll(dir)

	opts, _ := parseArgs([]string{"-s", dir, "--log", "Application"})
	opts.prepare()

	stateFile := getStateFile(opts.StateDir, "Application")

	recordNumber, _ := getLastOffset(stateFile)
	lastNumber := recordNumber
	assert.Equal(t, int64(0), recordNumber, "something went wrong")

	testEmpty := func() {
		w, c, errLines, err := opts.searchLog("Application")
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(0), w, "something went wrong")
		assert.Equal(t, int64(0), c, "something went wrong")
		assert.Equal(t, "", errLines, "something went wrong")

		recordNumber, _ = getLastOffset(stateFile)
		assert.NotEqual(t, 0, recordNumber, "something went wrong")
	}
	testEmpty()

	lastNumber = recordNumber

	testInfo := func() {
		raiseEvent(t, 0, "check-event-log: something info occured")
		w, c, errLines, err := opts.searchLog("Application")
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(0), w, "something went wrong")
		assert.Equal(t, int64(0), c, "something went wrong")
		assert.Equal(t, "", errLines, "something went wrong")

		recordNumber, _ = getLastOffset(stateFile)
		assert.NotEqual(t, lastNumber, recordNumber, "something went wrong")
	}
	testInfo()

	lastNumber = recordNumber

	testWarning := func() {
		raiseEvent(t, 2, "check-event-log: something warning occured")
		w, c, errLines, err := opts.searchLog("Application")
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(1), w, "something went wrong")
		assert.Equal(t, int64(0), c, "something went wrong")
		assert.Equal(t, "", errLines, "something went wrong")

		recordNumber, _ = getLastOffset(stateFile)
		assert.NotEqual(t, lastNumber, recordNumber, "something went wrong")
	}
	testWarning()

	lastNumber = recordNumber

	testError := func() {
		raiseEvent(t, 1, "check-event-log: something error occured")
		w, c, errLines, err := opts.searchLog("Application")
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(0), w, "something went wrong")
		assert.Equal(t, int64(1), c, "something went wrong")
		assert.Equal(t, "", errLines, "something went wrong")

		recordNumber, _ = getLastOffset(stateFile)
		assert.NotEqual(t, lastNumber, recordNumber, "something went wrong")
	}
	testError()

	lastNumber = recordNumber

	opts, _ = parseArgs([]string{"-s", dir, "--log", "Application", "-r"})
	opts.prepare()

	testReturn := func() {
		raiseEvent(t, 1, "check-event-log: something error occured")
		raiseEvent(t, 2, "check-event-log: something warning occured")
		w, c, errLines, err := opts.searchLog("Application")
		assert.Equal(t, err, nil, "err should be nil")
		assert.Equal(t, int64(1), w, "something went wrong")
		assert.Equal(t, int64(1), c, "something went wrong")
		assert.Equal(t, "WSH:check-event-log: something error occured\nWSH:check-event-log: something warning occured\n", errLines, "something went wrong")

		recordNumber, _ = getLastOffset(stateFile)
		assert.NotEqual(t, lastNumber, recordNumber, "something went wrong")
	}
	testReturn()
}
<|endoftext|>


<commit_before>package cmd

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/dgrijalva/jwt-go"
	"github.com/go-errors/errors"
	"github.com/mdp/qrterminal"
	"github.com/privacybydesign/irmago"
	"github.com/privacybydesign/irmago/internal/fs"
	"github.com/privacybydesign/irmago/server"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

// requestCmd represents the request command
var requestCmd = &cobra.Command{
	Use:   "request",
	Short: "Generate an IRMA session request",
	Run: func(cmd *cobra.Command, args []string) {
		request, _, err := configureRequest(cmd)
		if err != nil {
			die("", err)
		}

		flags := cmd.Flags()
		authmethod, _ := flags.GetString("authmethod")
		var output string
		if authmethod == "none" || authmethod == "token" {
			output = prettyprint(request)
		} else {
			key, _ := flags.GetString("key")
			name, _ := flags.GetString("name")
			if output, err = signRequest(request, name, authmethod, key); err != nil {
				die("Failed to sign request", err)
			}
		}

		fmt.Println(output)
	},
}

func signRequest(request irma.RequestorRequest, name, authmethod, key string) (string, error) {
	var (
		err    error
		sk     interface{}
		jwtalg jwt.SigningMethod
		bts    []byte
	)
	// If the key refers to an existing file, use contents of the file as key
	if bts, err = fs.ReadKey("", key); err != nil {
		bts = []byte(key)
	}
	switch authmethod {
	case "hmac":
		jwtalg = jwt.SigningMethodHS256
		if sk, err = fs.Base64Decode(bts); err != nil {
			return "", err
		}
	case "rsa":
		jwtalg = jwt.SigningMethodRS256
		if sk, err = jwt.ParseRSAPrivateKeyFromPEM(bts); err != nil {
			return "", err
		}
	default:
		return "", errors.Errorf("Unsupported signing algorithm: '%s'", authmethod)
	}

	return irma.SignRequestorRequest(request, jwtalg, sk, name)
}

func configureRequest(cmd *cobra.Command) (irma.RequestorRequest, *irma.Configuration, error) {
	irmaconfigPath, err := cmd.Flags().GetString("schemes-path")
	if err != nil {
		return nil, nil, err
	}
	irmaconfig, err := irma.NewConfiguration(irmaconfigPath)
	if err != nil {
		return nil, nil, err
	}
	if err = irmaconfig.ParseFolder(); err != nil {
		return nil, nil, err
	}
	if len(irmaconfig.SchemeManagers) == 0 {
		if err = irmaconfig.DownloadDefaultSchemes(); err != nil {
			return nil, nil, err
		}
	}

	request, err := constructSessionRequest(cmd, irmaconfig)
	if err != nil {
		return nil, nil, err
	}

	return request, irmaconfig, nil
}

// Helper functions

// poll recursively polls the session status until a status different from initialStatus is received.
func poll(initialStatus server.Status, transport *irma.HTTPTransport, statuschan chan server.Status) {
	// First we wait
	<-time.NewTimer(pollInterval).C

	// Get session status
	var status string
	if err := transport.Get("status", &status); err != nil {
		_ = server.LogFatal(err)
	}
	status = strings.Trim(status, `"`)

	// If the status has not yet changed, schedule another poll
	if server.Status(status) == initialStatus {
		go poll(initialStatus, transport, statuschan)
	} else {
		logger.Trace("Stopped polling, new status ", status)
		statuschan <- server.Status(status)
	}
}

func constructSessionRequest(cmd *cobra.Command, conf *irma.Configuration) (irma.RequestorRequest, error) {
	disclose, _ := cmd.Flags().GetStringArray("disclose")
	issue, _ := cmd.Flags().GetStringArray("issue")
	sign, _ := cmd.Flags().GetStringArray("sign")
	message, _ := cmd.Flags().GetString("message")
	jsonrequest, _ := cmd.Flags().GetString("request")

	if len(disclose) == 0 && len(issue) == 0 && len(sign) == 0 && message == "" {
		if jsonrequest == "" {
			return nil, errors.New("Provide either a complete session request using --request or construct one using the other flags")
		}
		request, err := server.ParseSessionRequest(jsonrequest)
		if err != nil {
			return nil, err
		}
		return request, nil
	}

	if jsonrequest != "" {
		return nil, errors.New("Provide either a complete session request using --request or construct one using the other flags")
	}

	if len(sign) != 0 {
		if len(disclose) != 0 {
			return nil, errors.New("cannot combine disclosure and signature sessions, use either --disclose or --sign")
		}
		if len(issue) != 0 {
			return nil, errors.New("cannot combine issuance and signature sessions, use either --issue or --sign")
		}
		if message == "" {
			return nil, errors.New("signature sessions require a message to be signed using --message")
		}
	}

	var request irma.RequestorRequest
	if len(disclose) != 0 {
		disclose, err := parseAttrs(disclose, conf)
		if err != nil {
			return nil, err
		}
		request = &irma.ServiceProviderRequest{
			Request: irma.NewDisclosureRequest(),
		}
		request.SessionRequest().(*irma.DisclosureRequest).Disclose = disclose
	}
	if len(sign) != 0 {
		disclose, err := parseAttrs(sign, conf)
		if err != nil {
			return nil, err
		}
		request = &irma.SignatureRequestorRequest{
			Request: irma.NewSignatureRequest(message),
		}
		request.SessionRequest().(*irma.SignatureRequest).Disclose = disclose
	}
	if len(issue) != 0 {
		creds, err := parseCredentials(issue, conf)
		if err != nil {
			return nil, err
		}
		disclose, err := parseAttrs(disclose, conf)
		if err != nil {
			return nil, err
		}
		request = &irma.IdentityProviderRequest{
			Request: irma.NewIssuanceRequest(creds),
		}
		request.SessionRequest().(*irma.IssuanceRequest).Disclose = disclose
	}

	return request, nil
}

func parseCredentials(credentialsStr []string, conf *irma.Configuration) ([]*irma.CredentialRequest, error) {
	list := make([]*irma.CredentialRequest, 0, len(credentialsStr))

	for _, credStr := range credentialsStr {
		parts := strings.Split(credStr, "=")
		if len(parts) != 2 {
			return nil, errors.New("--issue argument must contain exactly 1 = sign")
		}
		credIdStr, attrsStr := parts[0], parts[1]
		credtype := conf.CredentialTypes[irma.NewCredentialTypeIdentifier(credIdStr)]
		if credtype == nil {
			return nil, errors.New("unknown credential type: " + credIdStr)
		}

		attrsSlice := strings.Split(attrsStr, ",")
		if len(attrsSlice) != len(credtype.AttributeTypes) {
			return nil, errors.Errorf("%d attributes required but %d provided for %s", len(credtype.AttributeTypes), len(attrsSlice), credIdStr)
		}

		attrs := make(map[string]string, len(attrsSlice))
		for i, typ := range credtype.AttributeTypes {
			attrs[typ.ID] = attrsSlice[i]
		}
		list = append(list, &irma.CredentialRequest{
			CredentialTypeID: irma.NewCredentialTypeIdentifier(credIdStr),
			Attributes:       attrs,
		})
	}

	return list, nil
}

func parseAttrs(attrsStr []string, conf *irma.Configuration) (irma.AttributeConDisCon, error) {
	list := make(irma.AttributeConDisCon, 0, len(attrsStr))
	for _, disjunctionStr := range attrsStr {
		disjunction := irma.AttributeDisCon{}
		attrids := strings.Split(disjunctionStr, ",")
		for _, attridStr := range attrids {
			attrid := irma.NewAttributeTypeIdentifier(attridStr)
			if conf.AttributeTypes[attrid] == nil {
				return nil, errors.New("unknown attribute: " + attridStr)
			}
			disjunction = append(disjunction, irma.AttributeCon{irma.AttributeRequest{Type: attrid}})
		}
		list = append(list, disjunction)
	}
	return list, nil
}

func startServer(port int) {
	mux := http.NewServeMux()
	mux.HandleFunc("/", irmaServer.HandlerFunc())
	httpServer = &http.Server{Addr: ":" + strconv.Itoa(port), Handler: mux}
	go func() {
		err := httpServer.ListenAndServe()
		if err != nil && err != http.ErrServerClosed {
			die("Failed to start server", err)
		}
	}()
}

func printQr(qr *irma.Qr, noqr bool) error {
	qrBts, err := json.Marshal(qr)
	if err != nil {
		return err
	}
	if noqr {
		fmt.Println(string(qrBts))
	} else {
		qrterminal.GenerateWithConfig(string(qrBts), qrterminal.Config{
			Level:     qrterminal.L,
			Writer:    os.Stdout,
			BlackChar: qrterminal.BLACK,
			WhiteChar: qrterminal.WHITE,
		})
	}
	return nil
}

func printSessionResult(result *server.SessionResult) {
	fmt.Println("Session result:")
	fmt.Println(prettyprint(result))
}

func init() {
	RootCmd.AddCommand(requestCmd)

	flags := requestCmd.Flags()
	flags.SortFlags = false

	addRequestFlags(flags)
}

func addRequestFlags(flags *pflag.FlagSet) {
	flags.StringP("schemes-path", "s", server.DefaultSchemesPath(), "path to irma_configuration")
	flags.StringP("authmethod", "a", "none", "Authentication method to server (none, token, rsa, hmac)")
	flags.String("key", "", "Key to sign request with")
	flags.String("name", "", "Requestor name")
	flags.StringArray("disclose", nil, "Add an attribute disjunction (comma-separated)")
	flags.StringArray("issue", nil, "Add a credential to issue")
	flags.StringArray("sign", nil, "Add an attribute disjunction to signature session")
	flags.String("message", "", "Message to sign in signature session")
}
<commit_msg>feat: rename --authmethod to --auth-method flag to irma session for consistency with Configuration struct<commit_after>package cmd

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/dgrijalva/jwt-go"
	"github.com/go-errors/errors"
	"github.com/mdp/qrterminal"
	"github.com/privacybydesign/irmago"
	"github.com/privacybydesign/irmago/internal/fs"
	"github.com/privacybydesign/irmago/server"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

// requestCmd represents the request command
var requestCmd = &cobra.Command{
	Use:   "request",
	Short: "Generate an IRMA session request",
	Run: func(cmd *cobra.Command, args []string) {
		request, _, err := configureRequest(cmd)
		if err != nil {
			die("", err)
		}

		flags := cmd.Flags()
		authmethod, _ := flags.GetString("authmethod")
		var output string
		if authmethod == "none" || authmethod == "token" {
			output = prettyprint(request)
		} else {
			key, _ := flags.GetString("key")
			name, _ := flags.GetString("name")
			if output, err = signRequest(request, name, authmethod, key); err != nil {
				die("Failed to sign request", err)
			}
		}

		fmt.Println(output)
	},
}

func signRequest(request irma.RequestorRequest, name, authmethod, key string) (string, error) {
	var (
		err    error
		sk     interface{}
		jwtalg jwt.SigningMethod
		bts    []byte
	)
	// If the key refers to an existing file, use contents of the file as key
	if bts, err = fs.ReadKey("", key); err != nil {
		bts = []byte(key)
	}
	switch authmethod {
	case "hmac":
		jwtalg = jwt.SigningMethodHS256
		if sk, err = fs.Base64Decode(bts); err != nil {
			return "", err
		}
	case "rsa":
		jwtalg = jwt.SigningMethodRS256
		if sk, err = jwt.ParseRSAPrivateKeyFromPEM(bts); err != nil {
			return "", err
		}
	default:
		return "", errors.Errorf("Unsupported signing algorithm: '%s'", authmethod)
	}

	return irma.SignRequestorRequest(request, jwtalg, sk, name)
}

func configureRequest(cmd *cobra.Command) (irma.RequestorRequest, *irma.Configuration, error) {
	irmaconfigPath, err := cmd.Flags().GetString("schemes-path")
	if err != nil {
		return nil, nil, err
	}
	irmaconfig, err := irma.NewConfiguration(irmaconfigPath)
	if err != nil {
		return nil, nil, err
	}
	if err = irmaconfig.ParseFolder(); err != nil {
		return nil, nil, err
	}
	if len(irmaconfig.SchemeManagers) == 0 {
		if err = irmaconfig.DownloadDefaultSchemes(); err != nil {
			return nil, nil, err
		}
	}

	request, err := constructSessionRequest(cmd, irmaconfig)
	if err != nil {
		return nil, nil, err
	}

	return request, irmaconfig, nil
}

// Helper functions

// poll recursively polls the session status until a status different from initialStatus is received.
func poll(initialStatus server.Status, transport *irma.HTTPTransport, statuschan chan server.Status) {
	// First we wait
	<-time.NewTimer(pollInterval).C

	// Get session status
	var status string
	if err := transport.Get("status", &status); err != nil {
		_ = server.LogFatal(err)
	}
	status = strings.Trim(status, `"`)

	// If the status has not yet changed, schedule another poll
	if server.Status(status) == initialStatus {
		go poll(initialStatus, transport, statuschan)
	} else {
		logger.Trace("Stopped polling, new status ", status)
		statuschan <- server.Status(status)
	}
}

func constructSessionRequest(cmd *cobra.Command, conf *irma.Configuration) (irma.RequestorRequest, error) {
	disclose, _ := cmd.Flags().GetStringArray("disclose")
	issue, _ := cmd.Flags().GetStringArray("issue")
	sign, _ := cmd.Flags().GetStringArray("sign")
	message, _ := cmd.Flags().GetString("message")
	jsonrequest, _ := cmd.Flags().GetString("request")

	if len(disclose) == 0 && len(issue) == 0 && len(sign) == 0 && message == "" {
		if jsonrequest == "" {
			return nil, errors.New("Provide either a complete session request using --request or construct one using the other flags")
		}
		request, err := server.ParseSessionRequest(jsonrequest)
		if err != nil {
			return nil, err
		}
		return request, nil
	}

	if jsonrequest != "" {
		return nil, errors.New("Provide either a complete session request using --request or construct one using the other flags")
	}

	if len(sign) != 0 {
		if len(disclose) != 0 {
			return nil, errors.New("cannot combine disclosure and signature sessions, use either --disclose or --sign")
		}
		if len(issue) != 0 {
			return nil, errors.New("cannot combine issuance and signature sessions, use either --issue or --sign")
		}
		if message == "" {
			return nil, errors.New("signature sessions require a message to be signed using --message")
		}
	}

	var request irma.RequestorRequest
	if len(disclose) != 0 {
		disclose, err := parseAttrs(disclose, conf)
		if err != nil {
			return nil, err
		}
		request = &irma.ServiceProviderRequest{
			Request: irma.NewDisclosureRequest(),
		}
		request.SessionRequest().(*irma.DisclosureRequest).Disclose = disclose
	}
	if len(sign) != 0 {
		disclose, err := parseAttrs(sign, conf)
		if err != nil {
			return nil, err
		}
		request = &irma.SignatureRequestorRequest{
			Request: irma.NewSignatureRequest(message),
		}
		request.SessionRequest().(*irma.SignatureRequest).Disclose = disclose
	}
	if len(issue) != 0 {
		creds, err := parseCredentials(issue, conf)
		if err != nil {
			return nil, err
		}
		disclose, err := parseAttrs(disclose, conf)
		if err != nil {
			return nil, err
		}
		request = &irma.IdentityProviderRequest{
			Request: irma.NewIssuanceRequest(creds),
		}
		request.SessionRequest().(*irma.IssuanceRequest).Disclose = disclose
	}

	return request, nil
}

func parseCredentials(credentialsStr []string, conf *irma.Configuration) ([]*irma.CredentialRequest, error) {
	list := make([]*irma.CredentialRequest, 0, len(credentialsStr))

	for _, credStr := range credentialsStr {
		parts := strings.Split(credStr, "=")
		if len(parts) != 2 {
			return nil, errors.New("--issue argument must contain exactly 1 = sign")
		}
		credIdStr, attrsStr := parts[0], parts[1]
		credtype := conf.CredentialTypes[irma.NewCredentialTypeIdentifier(credIdStr)]
		if credtype == nil {
			return nil, errors.New("unknown credential type: " + credIdStr)
		}

		attrsSlice := strings.Split(attrsStr, ",")
		if len(attrsSlice) != len(credtype.AttributeTypes) {
			return nil, errors.Errorf("%d attributes required but %d provided for %s", len(credtype.AttributeTypes), len(attrsSlice), credIdStr)
		}

		attrs := make(map[string]string, len(attrsSlice))
		for i, typ := range credtype.AttributeTypes {
			attrs[typ.ID] = attrsSlice[i]
		}
		list = append(list, &irma.CredentialRequest{
			CredentialTypeID: irma.NewCredentialTypeIdentifier(credIdStr),
			Attributes:       attrs,
		})
	}

	return list, nil
}

func parseAttrs(attrsStr []string, conf *irma.Configuration) (irma.AttributeConDisCon, error) {
	list := make(irma.AttributeConDisCon, 0, len(attrsStr))
	for _, disjunctionStr := range attrsStr {
		disjunction := irma.AttributeDisCon{}
		attrids := strings.Split(disjunctionStr, ",")
		for _, attridStr := range attrids {
			attrid := irma.NewAttributeTypeIdentifier(attridStr)
			if conf.AttributeTypes[attrid] == nil {
				return nil, errors.New("unknown attribute: " + attridStr)
			}
			disjunction = append(disjunction, irma.AttributeCon{irma.AttributeRequest{Type: attrid}})
		}
		list = append(list, disjunction)
	}
	return list, nil
}

func startServer(port int) {
	mux := http.NewServeMux()
	mux.HandleFunc("/", irmaServer.HandlerFunc())
	httpServer = &http.Server{Addr: ":" + strconv.Itoa(port), Handler: mux}
	go func() {
		err := httpServer.ListenAndServe()
		if err != nil && err != http.ErrServerClosed {
			die("Failed to start server", err)
		}
	}()
}

func printQr(qr *irma.Qr, noqr bool) error {
	qrBts, err := json.Marshal(qr)
	if err != nil {
		return err
	}
	if noqr {
		fmt.Println(string(qrBts))
	} else {
		qrterminal.GenerateWithConfig(string(qrBts), qrterminal.Config{
			Level:     qrterminal.L,
			Writer:    os.Stdout,
			BlackChar: qrterminal.BLACK,
			WhiteChar: qrterminal.WHITE,
		})
	}
	return nil
}

func printSessionResult(result *server.SessionResult) {
	fmt.Println("Session result:")
	fmt.Println(prettyprint(result))
}

func init() {
	RootCmd.AddCommand(requestCmd)

	flags := requestCmd.Flags()
	flags.SortFlags = false

	addRequestFlags(flags)
}

func authmethodAlias(f *pflag.FlagSet, name string) pflag.NormalizedName {
	switch name {
	case "authmethod":
		name = "auth-method"
		break
	}
	return pflag.NormalizedName(name)
}

func addRequestFlags(flags *pflag.FlagSet) {
	flags.StringP("schemes-path", "s", server.DefaultSchemesPath(), "path to irma_configuration")
	flags.StringP("auth-method", "a", "none", "Authentication method to server (none, token, rsa, hmac)")
	flags.SetNormalizeFunc(authmethodAlias)
	flags.String("key", "", "Key to sign request with")
	flags.String("name", "", "Requestor name")
	flags.StringArray("disclose", nil, "Add an attribute disjunction (comma-separated)")
	flags.StringArray("issue", nil, "Add a credential to issue")
	flags.StringArray("sign", nil, "Add an attribute disjunction to signature session")
	flags.String("message", "", "Message to sign in signature session")
}
<|endoftext|>


<commit_before>package lfs

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"hash"
	"io"
	"os"
	"os/exec"
	"sort"
	"strings"
)

// An Extension describes how to manipulate files during smudge and clean.
// Extensions are parsed from the Git config.
type Extension struct {
	Name     string
	Clean    string
	Smudge   string
	Priority int
}

type pipeRequest struct {
	action     string
	reader     io.Reader
	fileName   string
	extensions []Extension
}

type pipeResponse struct {
	file    *os.File
	results []*pipeExtResult
}

type pipeExtResult struct {
	name   string
	oidIn  string
	oidOut string
}

type extCommand struct {
	cmd    *exec.Cmd
	out    io.WriteCloser
	err    *bytes.Buffer
	hasher hash.Hash
	result *pipeExtResult
}

// SortExtensions sorts a map of extensions in ascending order by Priority
func SortExtensions(m map[string]Extension) ([]Extension, error) {
	pMap := make(map[int]Extension)
	for n, ext := range m {
		p := ext.Priority
		if _, exist := pMap[p]; exist {
			err := fmt.Errorf("duplicate priority %d on %s", p, n)
			return nil, err
		}
		pMap[p] = ext
	}

	var priorities []int
	for p := range pMap {
		priorities = append(priorities, p)
	}

	sort.Ints(priorities)

	var result []Extension
	for _, p := range priorities {
		result = append(result, pMap[p])
	}

	return result, nil
}

func pipeExtensions(request *pipeRequest) (response pipeResponse, err error) {
	var extcmds []*extCommand
	for _, e := range request.extensions {
		var pieces []string
		switch request.action {
		case "clean":
			pieces = strings.Split(e.Clean, " ")
		case "smudge":
			pieces = strings.Split(e.Smudge, " ")
		default:
			err = fmt.Errorf("Invalid action: " + request.action)
			return
		}
		name := strings.Trim(pieces[0], " ")
		var args []string
		for _, value := range pieces[1:] {
			arg := strings.Replace(value, "%f", request.fileName, -1)
			args = append(args, arg)
		}
		cmd := exec.Command(name, args...)
		ec := &extCommand{cmd: cmd, result: &pipeExtResult{name: e.Name}}
		extcmds = append(extcmds, ec)
	}

	hasher := sha256.New()
	pipeReader, pipeWriter := io.Pipe()
	multiWriter := io.MultiWriter(hasher, pipeWriter)

	var input io.Reader
	var output io.WriteCloser
	input = pipeReader
	extcmds[0].cmd.Stdin = input
	if response.file, err = TempFile(""); err != nil {
		return
	}
	defer response.file.Close()
	output = response.file

	last := len(extcmds) - 1
	for i, ec := range extcmds {
		ec.hasher = sha256.New()

		if i == last {
			ec.cmd.Stdout = io.MultiWriter(ec.hasher, output)
			ec.out = output
			continue
		}

		nextec := extcmds[i+1]
		var nextStdin io.WriteCloser
		var stdout io.ReadCloser
		if nextStdin, err = nextec.cmd.StdinPipe(); err != nil {
			return
		}
		if stdout, err = ec.cmd.StdoutPipe(); err != nil {
			return
		}

		ec.cmd.Stdin = input
		ec.cmd.Stdout = io.MultiWriter(ec.hasher, nextStdin)
		ec.out = nextStdin

		input = stdout

		var errBuff bytes.Buffer
		ec.err = &errBuff
		ec.cmd.Stderr = ec.err
	}

	for _, ec := range extcmds {
		if err = ec.cmd.Start(); err != nil {
			return
		}
	}

	if _, err = io.Copy(multiWriter, request.reader); err != nil {
		return
	}
	if err = pipeWriter.Close(); err != nil {
		return
	}

	for _, ec := range extcmds {
		if err = ec.cmd.Wait(); err != nil {
			if ec.err != nil {
				errStr := ec.err.String()
				err = fmt.Errorf("Extension '%s' failed with: %s", ec.result.name, errStr)
			}
			return
		}
		if err = ec.out.Close(); err != nil {
			return
		}
	}

	oid := hex.EncodeToString(hasher.Sum(nil))
	for _, ec := range extcmds {
		ec.result.oidIn = oid
		oid = hex.EncodeToString(ec.hasher.Sum(nil))
		ec.result.oidOut = oid
		response.results = append(response.results, ec.result)
	}
	return
}
<commit_msg>refactor lfs.SortExtensions<commit_after>package lfs

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"hash"
	"io"
	"os"
	"os/exec"
	"sort"
	"strings"
)

// An Extension describes how to manipulate files during smudge and clean.
// Extensions are parsed from the Git config.
type Extension struct {
	Name     string
	Clean    string
	Smudge   string
	Priority int
}

type pipeRequest struct {
	action     string
	reader     io.Reader
	fileName   string
	extensions []Extension
}

type pipeResponse struct {
	file    *os.File
	results []*pipeExtResult
}

type pipeExtResult struct {
	name   string
	oidIn  string
	oidOut string
}

type extCommand struct {
	cmd    *exec.Cmd
	out    io.WriteCloser
	err    *bytes.Buffer
	hasher hash.Hash
	result *pipeExtResult
}

// SortExtensions sorts a map of extensions in ascending order by Priority
func SortExtensions(m map[string]Extension) ([]Extension, error) {
	pMap := make(map[int]Extension)
	priorities := make([]int, 0, len(m))
	for n, ext := range m {
		p := ext.Priority
		if _, exist := pMap[p]; exist {
			err := fmt.Errorf("duplicate priority %d on %s", p, n)
			return nil, err
		}
		pMap[p] = ext
		priorities = append(priorities, p)
	}

	sort.Ints(priorities)

	result := make([]Extension, len(priorities))
	for i, p := range priorities {
		result[i] = pMap[p]
	}

	return result, nil
}

func pipeExtensions(request *pipeRequest) (response pipeResponse, err error) {
	var extcmds []*extCommand
	for _, e := range request.extensions {
		var pieces []string
		switch request.action {
		case "clean":
			pieces = strings.Split(e.Clean, " ")
		case "smudge":
			pieces = strings.Split(e.Smudge, " ")
		default:
			err = fmt.Errorf("Invalid action: " + request.action)
			return
		}
		name := strings.Trim(pieces[0], " ")
		var args []string
		for _, value := range pieces[1:] {
			arg := strings.Replace(value, "%f", request.fileName, -1)
			args = append(args, arg)
		}
		cmd := exec.Command(name, args...)
		ec := &extCommand{cmd: cmd, result: &pipeExtResult{name: e.Name}}
		extcmds = append(extcmds, ec)
	}

	hasher := sha256.New()
	pipeReader, pipeWriter := io.Pipe()
	multiWriter := io.MultiWriter(hasher, pipeWriter)

	var input io.Reader
	var output io.WriteCloser
	input = pipeReader
	extcmds[0].cmd.Stdin = input
	if response.file, err = TempFile(""); err != nil {
		return
	}
	defer response.file.Close()
	output = response.file

	last := len(extcmds) - 1
	for i, ec := range extcmds {
		ec.hasher = sha256.New()

		if i == last {
			ec.cmd.Stdout = io.MultiWriter(ec.hasher, output)
			ec.out = output
			continue
		}

		nextec := extcmds[i+1]
		var nextStdin io.WriteCloser
		var stdout io.ReadCloser
		if nextStdin, err = nextec.cmd.StdinPipe(); err != nil {
			return
		}
		if stdout, err = ec.cmd.StdoutPipe(); err != nil {
			return
		}

		ec.cmd.Stdin = input
		ec.cmd.Stdout = io.MultiWriter(ec.hasher, nextStdin)
		ec.out = nextStdin

		input = stdout

		var errBuff bytes.Buffer
		ec.err = &errBuff
		ec.cmd.Stderr = ec.err
	}

	for _, ec := range extcmds {
		if err = ec.cmd.Start(); err != nil {
			return
		}
	}

	if _, err = io.Copy(multiWriter, request.reader); err != nil {
		return
	}
	if err = pipeWriter.Close(); err != nil {
		return
	}

	for _, ec := range extcmds {
		if err = ec.cmd.Wait(); err != nil {
			if ec.err != nil {
				errStr := ec.err.String()
				err = fmt.Errorf("Extension '%s' failed with: %s", ec.result.name, errStr)
			}
			return
		}
		if err = ec.out.Close(); err != nil {
			return
		}
	}

	oid := hex.EncodeToString(hasher.Sum(nil))
	for _, ec := range extcmds {
		ec.result.oidIn = oid
		oid = hex.EncodeToString(ec.hasher.Sum(nil))
		ec.result.oidOut = oid
		response.results = append(response.results, ec.result)
	}
	return
}
<|endoftext|>


<commit_before>package main

// An example streaming XML parser.

import (
	"encoding/json"
	"encoding/xml"
	"flag"
	"fmt"
	"net/url"
	"os"
	"regexp"
	"strings"
)

var filter, _ = regexp.Compile("^file:.*|^talk:.*|^special:.*|^wikipedia:.*|^wiktionary:.*|^user:.*|^user_talk:.*")

// Here is an example article from the Wikipedia XML dump
//
// <page>
// 	<title>Apollo 11</title>
//      <redirect title="Foo bar" />
// 	...
// 	<revision>
// 	...
// 	  <text xml:space="preserve">
// 	  {{Infobox Space mission
// 	  |mission_name=<!--See above-->
// 	  |insignia=Apollo_11_insignia.png
// 	...
// 	  </text>
// 	</revision>
// </page>
//
// Note how the tags on the fields of Page and Redirect below
// describe the XML schema structure.

type Redirect struct {
	Title string `xml:"title,attr" json:"title"`
}

type Page struct {
	Title          string   `xml:"title" json:"title"`
	CanonicalTitle string   `xml:"ctitle" json:"ctitle"`
	Redir          Redirect `xml:"redirect" json:"redirect"`
	Text           string   `xml:"revision>text" json:"text"`
}

func CanonicalizeTitle(title string) string {
	can := strings.ToLower(title)
	can = strings.Replace(can, " ", "_", -1)
	can = url.QueryEscape(can)
	return can
}

func main() {
	flag.Parse()

	if flag.NArg() < 1 {
		fmt.Println("Usage: wptojson WIKIPEDIA-XML-DUMP")
		os.Exit(1)
	}
	inputFile := flag.Args()[0]

	xmlFile, err := os.Open(inputFile)
	if err != nil {
		fmt.Println("Error opening file:", err)
		return
	}
	defer xmlFile.Close()

	decoder := xml.NewDecoder(xmlFile)
	var inElement string
	for {
		// Read tokens from the XML document in a stream.
		t, _ := decoder.Token()
		if t == nil {
			break
		}
		// Inspect the type of the token just read.
		switch se := t.(type) {
		case xml.StartElement:
			// If we just read a StartElement token
			inElement = se.Name.Local
			// ...and its name is "page"
			if inElement == "page" {
				var p Page
				// decode a whole chunk of following XML into the
				// variable p which is a Page (se above)
				decoder.DecodeElement(&p, &se)

				// Do some stuff with the page.
				p.CanonicalTitle = CanonicalizeTitle(p.Title)
				m := filter.MatchString(p.CanonicalTitle)
				if !m && p.Redir.Title == "" {
					b, err := json.Marshal(p)
					if err != nil {
						os.Exit(2)
					}
					fmt.Println(string(b))
				}
			}
		default:
		}
	}
}
<commit_msg>added title filter (no regex)<commit_after>package main

// An example streaming XML parser.

import (
	"encoding/json"
	"encoding/xml"
	"flag"
	"fmt"
	"net/url"
	"os"
	"regexp"
	"strings"
)

var titleFilter = flag.String("f", "", "only filter pages that contain the given text (no regex)")
var filter, _ = regexp.Compile("^file:.*|^talk:.*|^special:.*|^wikipedia:.*|^wiktionary:.*|^user:.*|^user_talk:.*")

// Here is an example article from the Wikipedia XML dump
//
// <page>
// 	<title>Apollo 11</title>
//      <redirect title="Foo bar" />
// 	...
// 	<revision>
// 	...
// 	  <text xml:space="preserve">
// 	  {{Infobox Space mission
// 	  |mission_name=<!--See above-->
// 	  |insignia=Apollo_11_insignia.png
// 	...
// 	  </text>
// 	</revision>
// </page>
//
// Note how the tags on the fields of Page and Redirect below
// describe the XML schema structure.

type Redirect struct {
	Title string `xml:"title,attr" json:"title"`
}

type Page struct {
	Title          string   `xml:"title" json:"title"`
	CanonicalTitle string   `xml:"ctitle" json:"ctitle"`
	Redir          Redirect `xml:"redirect" json:"redirect"`
	Text           string   `xml:"revision>text" json:"text"`
}

func CanonicalizeTitle(title string) string {
	can := strings.ToLower(title)
	can = strings.Replace(can, " ", "_", -1)
	can = url.QueryEscape(can)
	return can
}

func main() {
	flag.Parse()

	if flag.NArg() < 1 {
		fmt.Println("Usage: wptojson WIKIPEDIA-XML-DUMP")
		os.Exit(1)
	}
	inputFile := flag.Args()[0]

	xmlFile, err := os.Open(inputFile)
	if err != nil {
		fmt.Println("Error opening file:", err)
		return
	}
	defer xmlFile.Close()

	decoder := xml.NewDecoder(xmlFile)
	var inElement string
	for {
		// Read tokens from the XML document in a stream.
		t, _ := decoder.Token()
		if t == nil {
			break
		}
		// Inspect the type of the token just read.
		switch se := t.(type) {
		case xml.StartElement:
			// If we just read a StartElement token
			inElement = se.Name.Local
			// ...and its name is "page"
			if inElement == "page" {
				var p Page
				// decode a whole chunk of following XML into the
				// variable p which is a Page (se above)
				decoder.DecodeElement(&p, &se)

				// Do some stuff with the page.
				p.CanonicalTitle = CanonicalizeTitle(p.Title)
				m := filter.MatchString(p.CanonicalTitle)
				if !m && p.Redir.Title == "" {
					if strings.Contains(p.Title, *titleFilter) {
						b, err := json.Marshal(p)
						if err != nil {
							os.Exit(2)
						}
						fmt.Println(string(b))
					}
				}
			}
		default:
		}
	}
}
<|endoftext|>


<commit_before>package crawler

import (
	"testing"
	"time"

	"google.golang.org/appengine/aetest"
)

func Test_crawlChannelClients(t *testing.T) {
	ctx, done, err := aetest.NewContext()
	if err != nil {
		t.Error(err)
	}
	defer done()

	tests := []struct {
		date      string
		expectNum int
	}{
		{"2017-01-21 16:55:00 JST", 7},
		{"2017-01-21 19:00:00 JST", 7},
		{"2017-01-22 16:55:00 JST", 8},
		{"2017-01-22 19:00:00 JST", 8},
	}

	for _, test := range tests {
		if test.date != "" {
			timeNow = func() time.Time {
				t, _ := time.Parse("2006-01-02 15:04:05 MST", test.date)
				return t
			}
		}
		clients := crawlChannelClients(ctx)
		if len(clients) != test.expectNum {
			t.Errorf("Expected number of clients %d, got %d", test.expectNum, len(clients))
		}
	}
}
<commit_msg>Fix test<commit_after>package crawler

import (
	"testing"
	"time"

	"google.golang.org/appengine/aetest"
)

func Test_crawlChannelClients(t *testing.T) {
	ctx, done, err := aetest.NewContext()
	if err != nil {
		t.Error(err)
	}
	defer done()

	tests := []struct {
		date      string
		expectNum int
	}{
		{"2017-01-21 16:55:00 +0900", 7},
		{"2017-01-21 19:00:00 +0900", 7},
		{"2017-01-22 16:55:00 +0900", 8},
		{"2017-01-22 19:00:00 +0900", 8},
	}

	for _, test := range tests {
		timeNow = func() time.Time {
			t, _ := time.Parse("2006-01-02 15:04:05 -0700", test.date)
			return t
		}
		clients := crawlChannelClients(ctx)
		if len(clients) != test.expectNum {
			t.Errorf("Expected number of clients %d, got %d", test.expectNum, len(clients))
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2010, Kyle Lemons <kyle@kylelemons.net>.  All rights reserved.

package log4go

import (
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

var (
	Global Logger
)

func init() {
	// auto load config from default position
	Global = NewDefaultLogger(DEBUG)
	file, _ := exec.LookPath(os.Args[0])
	dir := filepath.Dir(file)
	if _, err := os.Stat("log4go.xml"); !os.IsNotExist(err) {
		Global.LoadConfiguration("log4go.xml")
	} else if _, err := os.Stat(filepath.Join(dir, "/log4go.xml")); !os.IsNotExist(err) {
		Global.LoadConfiguration(filepath.Join(dir, "log4go.xml"))
	} else if _, err := os.Stat(filepath.Join(dir, "/conf/log4go.xml")); !os.IsNotExist(err) {
		Global.LoadConfiguration(filepath.Join(dir, "/conf/log4go.xml"))
	} else {
		fmt.Fprintf(os.Stdout, "log4go config not found, exec dir is: %s, u need to load it by yourself.\n", dir)
	}
}

// Wrapper for (*Logger).LoadConfiguration
func LoadConfiguration(filename string) {
	Global.LoadConfiguration(filename)
}

// Wrapper for (*Logger).AddFilter
func AddFilter(name string, lvl Level, writer LogWriter) {
	Global.AddFilter(name, lvl, writer)
}

// Wrapper for (*Logger).Close (closes and removes all logwriters)
func Close() {
	Global.Close()
}

func Crash(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(CRITICAL, strings.Repeat(" %v", len(args))[1:], args...)
	}
	panic(args)
}

// Logs the given message and crashes the program
func Crashf(format string, args ...interface{}) {
	Global.intLogf(CRITICAL, format, args...)
	Global.Close() // so that hopefully the messages get logged
	panic(fmt.Sprintf(format, args...))
}

// Compatibility with `log`
func Exit(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(ERROR, strings.Repeat(" %v", len(args))[1:], args...)
	}
	Global.Close() // so that hopefully the messages get logged
	os.Exit(0)
}

// Compatibility with `log`
func Exitf(format string, args ...interface{}) {
	Global.intLogf(ERROR, format, args...)
	Global.Close() // so that hopefully the messages get logged
	os.Exit(0)
}

// Compatibility with `log`
func Stderr(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(ERROR, strings.Repeat(" %v", len(args))[1:], args...)
	}
}

// Compatibility with `log`
func Stderrf(format string, args ...interface{}) {
	Global.intLogf(ERROR, format, args...)
}

// Compatibility with `log`
func Stdout(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(INFO, strings.Repeat(" %v", len(args))[1:], args...)
	}
}

// Compatibility with `log`
func Stdoutf(format string, args ...interface{}) {
	Global.intLogf(INFO, format, args...)
}

// Send a log message manually
// Wrapper for (*Logger).Log
func Log(lvl Level, source, message string) {
	Global.Log(lvl, source, message)
}

// Send a formatted log message easily
// Wrapper for (*Logger).Logf
func Logf(lvl Level, format string, args ...interface{}) {
	Global.intLogf(lvl, format, args...)
}

// Send a closure log message
// Wrapper for (*Logger).Logc
func Logc(lvl Level, closure func() string) {
	Global.intLogc(lvl, closure)
}

// Utility for finest log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Finest
func Finest(arg0 interface{}, args ...interface{}) {
	const (
		lvl = FINEST
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for fine log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Fine
func Fine(arg0 interface{}, args ...interface{}) {
	const (
		lvl = FINE
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for debug log messages
// When given a string as the first argument, this behaves like Logf but with the DEBUG log level (e.g. the first argument is interpreted as a format for the latter arguments)
// When given a closure of type func()string, this logs the string returned by the closure iff it will be logged.  The closure runs at most one time.
// When given anything else, the log message will be each of the arguments formatted with %v and separated by spaces (ala Sprint).
// Wrapper for (*Logger).Debug
func Debug(arg0 interface{}, args ...interface{}) {
	const (
		lvl = DEBUG
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for trace log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Trace
func Trace(arg0 interface{}, args ...interface{}) {
	const (
		lvl = TRACE
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for info log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Info
func Info(arg0 interface{}, args ...interface{}) {
	const (
		lvl = INFO
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for Access log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Info
func Access(arg0 interface{}, args ...interface{}) {
	const (
		lvl = ACCESS
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for warn log messages (returns an error for easy function returns) (see Debug() for parameter explanation)
// These functions will execute a closure exactly once, to build the error message for the return
// Wrapper for (*Logger).Warn
func Warn(arg0 interface{}, args ...interface{}) error {
	const (
		lvl = WARNING
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
		return errors.New(fmt.Sprintf(first, args...))
	case func() string:
		// Log the closure (no other arguments used)
		str := first()
		Global.intLogf(lvl, "%s", str)
		return errors.New(str)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(first)+strings.Repeat(" %v", len(args)), args...)
		return errors.New(fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...))
	}
	return nil
}

// Utility for error log messages (returns an error for easy function returns) (see Debug() for parameter explanation)
// These functions will execute a closure exactly once, to build the error message for the return
// Wrapper for (*Logger).Error
func Error(arg0 interface{}, args ...interface{}) error {
	const (
		lvl = ERROR
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
		return errors.New(fmt.Sprintf(first, args...))
	case func() string:
		// Log the closure (no other arguments used)
		str := first()
		Global.intLogf(lvl, "%s", str)
		return errors.New(str)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(first)+strings.Repeat(" %v", len(args)), args...)
		return errors.New(fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...))
	}
	return nil
}

// Utility for critical log messages (returns an error for easy function returns) (see Debug() for parameter explanation)
// These functions will execute a closure exactly once, to build the error message for the return
// Wrapper for (*Logger).Critical. This method will log the call stack
func Critical(arg0 interface{}, args ...interface{}) error {
	const (
		lvl = CRITICAL
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		msg := fmt.Sprintf("%s\n%s", fmt.Sprintf(first, args...), CallStack(3))
		Global.intLogf(lvl, msg)
		//Global.intLogf(lvl, "%s", CallStack(3))
		return errors.New(fmt.Sprintf(first, args...))
	case func() string:
		// Log the closure (no other arguments used)
		str := first()
		Global.intLogf(lvl, "%s\n%s", str, CallStack(3))
		//Global.intLogf(lvl, "%s", CallStack(3))
		return errors.New(str)
	default:
		// Build a format string so that it will be similar to Sprint
		msg := fmt.Sprintf("%s\n%s", fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...), CallStack(3))
		Global.intLogf(lvl, msg)
		return errors.New(fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...))
	}
	return nil
}
<commit_msg>add log4go.Recovery() to print stack log only when there is a panic<commit_after>// Copyright (C) 2010, Kyle Lemons <kyle@kylelemons.net>.  All rights reserved.

package log4go

import (
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

var (
	Global Logger
)

func init() {
	// auto load config from default position
	Global = NewDefaultLogger(DEBUG)
	file, _ := exec.LookPath(os.Args[0])
	dir := filepath.Dir(file)
	if _, err := os.Stat("log4go.xml"); !os.IsNotExist(err) {
		Global.LoadConfiguration("log4go.xml")
	} else if _, err := os.Stat(filepath.Join(dir, "/log4go.xml")); !os.IsNotExist(err) {
		Global.LoadConfiguration(filepath.Join(dir, "log4go.xml"))
	} else if _, err := os.Stat(filepath.Join(dir, "/conf/log4go.xml")); !os.IsNotExist(err) {
		Global.LoadConfiguration(filepath.Join(dir, "/conf/log4go.xml"))
	} else {
		fmt.Fprintf(os.Stdout, "log4go config not found, exec dir is: %s, u need to load it by yourself.\n", dir)
	}
}

// Wrapper for (*Logger).LoadConfiguration
func LoadConfiguration(filename string) {
	Global.LoadConfiguration(filename)
}

// Wrapper for (*Logger).AddFilter
func AddFilter(name string, lvl Level, writer LogWriter) {
	Global.AddFilter(name, lvl, writer)
}

// Wrapper for (*Logger).Close (closes and removes all logwriters)
func Close() {
	Global.Close()
}

func Crash(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(CRITICAL, strings.Repeat(" %v", len(args))[1:], args...)
	}
	panic(args)
}

// Logs the given message and crashes the program
func Crashf(format string, args ...interface{}) {
	Global.intLogf(CRITICAL, format, args...)
	Global.Close() // so that hopefully the messages get logged
	panic(fmt.Sprintf(format, args...))
}

// Compatibility with `log`
func Exit(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(ERROR, strings.Repeat(" %v", len(args))[1:], args...)
	}
	Global.Close() // so that hopefully the messages get logged
	os.Exit(0)
}

// Compatibility with `log`
func Exitf(format string, args ...interface{}) {
	Global.intLogf(ERROR, format, args...)
	Global.Close() // so that hopefully the messages get logged
	os.Exit(0)
}

// Compatibility with `log`
func Stderr(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(ERROR, strings.Repeat(" %v", len(args))[1:], args...)
	}
}

// Compatibility with `log`
func Stderrf(format string, args ...interface{}) {
	Global.intLogf(ERROR, format, args...)
}

// Compatibility with `log`
func Stdout(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(INFO, strings.Repeat(" %v", len(args))[1:], args...)
	}
}

// Compatibility with `log`
func Stdoutf(format string, args ...interface{}) {
	Global.intLogf(INFO, format, args...)
}

// Send a log message manually
// Wrapper for (*Logger).Log
func Log(lvl Level, source, message string) {
	Global.Log(lvl, source, message)
}

// Send a formatted log message easily
// Wrapper for (*Logger).Logf
func Logf(lvl Level, format string, args ...interface{}) {
	Global.intLogf(lvl, format, args...)
}

// Send a closure log message
// Wrapper for (*Logger).Logc
func Logc(lvl Level, closure func() string) {
	Global.intLogc(lvl, closure)
}

// Utility for finest log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Finest
func Finest(arg0 interface{}, args ...interface{}) {
	const (
		lvl = FINEST
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for fine log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Fine
func Fine(arg0 interface{}, args ...interface{}) {
	const (
		lvl = FINE
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for debug log messages
// When given a string as the first argument, this behaves like Logf but with the DEBUG log level (e.g. the first argument is interpreted as a format for the latter arguments)
// When given a closure of type func()string, this logs the string returned by the closure iff it will be logged.  The closure runs at most one time.
// When given anything else, the log message will be each of the arguments formatted with %v and separated by spaces (ala Sprint).
// Wrapper for (*Logger).Debug
func Debug(arg0 interface{}, args ...interface{}) {
	const (
		lvl = DEBUG
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for trace log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Trace
func Trace(arg0 interface{}, args ...interface{}) {
	const (
		lvl = TRACE
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for info log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Info
func Info(arg0 interface{}, args ...interface{}) {
	const (
		lvl = INFO
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for Access log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Info
func Access(arg0 interface{}, args ...interface{}) {
	const (
		lvl = ACCESS
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for warn log messages (returns an error for easy function returns) (see Debug() for parameter explanation)
// These functions will execute a closure exactly once, to build the error message for the return
// Wrapper for (*Logger).Warn
func Warn(arg0 interface{}, args ...interface{}) error {
	const (
		lvl = WARNING
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
		return errors.New(fmt.Sprintf(first, args...))
	case func() string:
		// Log the closure (no other arguments used)
		str := first()
		Global.intLogf(lvl, "%s", str)
		return errors.New(str)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(first)+strings.Repeat(" %v", len(args)), args...)
		return errors.New(fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...))
	}
	return nil
}

// Utility for error log messages (returns an error for easy function returns) (see Debug() for parameter explanation)
// These functions will execute a closure exactly once, to build the error message for the return
// Wrapper for (*Logger).Error
func Error(arg0 interface{}, args ...interface{}) error {
	const (
		lvl = ERROR
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
		return errors.New(fmt.Sprintf(first, args...))
	case func() string:
		// Log the closure (no other arguments used)
		str := first()
		Global.intLogf(lvl, "%s", str)
		return errors.New(str)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(first)+strings.Repeat(" %v", len(args)), args...)
		return errors.New(fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...))
	}
	return nil
}

// Utility for critical log messages (returns an error for easy function returns) (see Debug() for parameter explanation)
// These functions will execute a closure exactly once, to build the error message for the return
// Wrapper for (*Logger).Critical. This method will log the call stack
func Critical(arg0 interface{}, args ...interface{}) error {
	const (
		lvl = CRITICAL
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		msg := fmt.Sprintf("%s\n%s", fmt.Sprintf(first, args...), CallStack(3))
		Global.intLogf(lvl, msg)
		//Global.intLogf(lvl, "%s", CallStack(3))
		return errors.New(fmt.Sprintf(first, args...))
	case func() string:
		// Log the closure (no other arguments used)
		str := first()
		Global.intLogf(lvl, "%s\n%s", str, CallStack(3))
		//Global.intLogf(lvl, "%s", CallStack(3))
		return errors.New(str)
	default:
		// Build a format string so that it will be similar to Sprint
		msg := fmt.Sprintf("%s\n%s", fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...), CallStack(3))
		Global.intLogf(lvl, msg)
		return errors.New(fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...))
	}
	return nil
}

func Recover(arg0 interface{}, args ...interface{}) {
	if err := recover(); err != nil {
		Critical(arg0, args...)
	} else {
		Error(arg0, args...)
	}
}<|endoftext|>