<commit_before>package zlib

import (
	"errors"
	"fmt"
	"net"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/zmap/zgrab/ztools/x509"
	"github.com/zmap/zgrab/ztools/ztls"
)

var smtpEndRegex = regexp.MustCompile(`(?:\r\n)|^[0-9]{3} .+\r\n$`)
var pop3EndRegex = regexp.MustCompile(`(?:\r\n\.\r\n$)|(?:\r\n$)`)
var imapStatusEndRegex = regexp.MustCompile(`\r\n$`)
var ftpEndRegex = regexp.MustCompile(`^.*[0-9]{3}( [^\r\n]*)?\r?\n$`)

const (
	SMTP_COMMAND = "STARTTLS\r\n"
	POP3_COMMAND = "STLS\r\n"
	IMAP_COMMAND = "a001 STARTTLS\r\n"
)

// Implements the net.Conn interface
type Conn struct {
	// Underlying network connection
	conn    net.Conn
	tlsConn *ztls.Conn
	isTls   bool

	// Max TLS version
	maxTlsVersion uint16

	// Keep track of state / network operations
	operations []ConnectionEvent

	// Cache the deadlines so we can reapply after TLS handshake
	readDeadline  time.Time
	writeDeadline time.Time

	caPool *x509.CertPool

	onlyCBC       bool
	onlySchannel  bool
	onlyDHE       bool
	onlyExports   bool
	onlyExportsDH bool

	domain string

	// Encoding type
	ReadEncoding string
}

func (c *Conn) getUnderlyingConn() net.Conn {
	if c.isTls {
		return c.tlsConn
	}
	return c.conn
}

func (c *Conn) SetCBCOnly() {
	c.onlyCBC = true
}

func (c *Conn) SetSChannelOnly() {
	c.onlySchannel = true
}

func (c *Conn) SetDHEOnly() {
	c.onlyDHE = true
}

func (c *Conn) SetExportsOnly() {
	c.onlyExports = true
}

func (c *Conn) SetExportsDHOnly() {
	c.onlyExportsDH = true
}

func (c *Conn) SetCAPool(pool *x509.CertPool) {
	c.caPool = pool
}

func (c *Conn) SetDomain(domain string) {
	c.domain = domain
}

// Layer in the regular conn methods
func (c *Conn) LocalAddr() net.Addr {
	return c.getUnderlyingConn().LocalAddr()
}

func (c *Conn) RemoteAddr() net.Addr {
	return c.getUnderlyingConn().RemoteAddr()
}

func (c *Conn) SetDeadline(t time.Time) error {
	c.readDeadline = t
	c.writeDeadline = t
	return c.getUnderlyingConn().SetDeadline(t)
}

func (c *Conn) SetReadDeadline(t time.Time) error {
	c.readDeadline = t
	return c.getUnderlyingConn().SetReadDeadline(t)
}

func (c *Conn) SetWriteDeadline(t time.Time) error {
	c.writeDeadline = t
	return c.getUnderlyingConn().SetWriteDeadline(t)
}

// Delegate here, but record all the things
func (c *Conn) Write(b []byte) (int, error) {
	n, err := c.getUnderlyingConn().Write(b)
	w := WriteEvent{Sent: b}
	event := ConnectionEvent{
		Data:  &w,
		Error: err,
	}
	c.operations = append(c.operations, event)
	return n, err
}

func (c *Conn) Read(b []byte) (int, error) {
	n, err := c.getUnderlyingConn().Read(b)
	r := ReadEvent{
		Response: b[0:n],
		encoding: c.ReadEncoding,
	}
	c.appendEvent(&r, err)
	return n, err
}

func (c *Conn) Close() error {
	return c.getUnderlyingConn().Close()
}

// Extra method - Do a TLS Handshake and record progress
func (c *Conn) TLSHandshake() error {
	if c.isTls {
		return fmt.Errorf(
			"Attempted repeat handshake with remote host %s",
			c.RemoteAddr().String())
	}
	tlsConfig := new(ztls.Config)
	tlsConfig.InsecureSkipVerify = true
	tlsConfig.MinVersion = ztls.VersionSSL30
	tlsConfig.MaxVersion = c.maxTlsVersion
	tlsConfig.RootCAs = c.caPool
	if c.domain != "" {
		tlsConfig.ServerName = c.domain
	}
	if c.onlyCBC {
		tlsConfig.CipherSuites = ztls.CBCSuiteIDList
	}
	if c.onlySchannel {
		tlsConfig.CipherSuites = ztls.SChannelSuites
	}
	if c.onlyDHE {
		tlsConfig.CipherSuites = ztls.DHECiphers
		tlsConfig.ForceSuites = true
	}
	if c.onlyExports {
		tlsConfig.CipherSuites = ztls.ExportCiphers
		tlsConfig.ForceSuites = true
	}
	if c.onlyExportsDH {
		tlsConfig.CipherSuites = ztls.DHEExportCiphers
		tlsConfig.ForceSuites = true
	}
	c.tlsConn = ztls.Client(c.conn, tlsConfig)
	c.tlsConn.SetReadDeadline(c.readDeadline)
	c.tlsConn.SetWriteDeadline(c.writeDeadline)
	c.isTls = true
	err := c.tlsConn.Handshake()
	hl := c.tlsConn.GetHandshakeLog()
	ts := TLSHandshakeEvent{handshakeLog: hl}
	event := ConnectionEvent{
		Data:  &ts,
		Error: err,
	}
	c.operations = append(c.operations, event)
	return err
}

func (c *Conn) sendStartTLSCommand(command string) error {
	// Don't doublehandshake
	if c.isTls {
		return fmt.Errorf(
			"Attempt STARTTLS after TLS handshake with remote host %s",
			c.RemoteAddr().String())
	}
	// Send the STARTTLS message
	starttls := []byte(command)
	_, err := c.conn.Write(starttls)
	return err
}

// Do a STARTTLS handshake
func (c *Conn) SMTPStartTLSHandshake() error {
	// Make the state
	ss := StartTLSEvent{Command: SMTP_COMMAND}

	// Send the command
	if err := c.sendStartTLSCommand(SMTP_COMMAND); err != nil {
		c.appendEvent(&ss, err)
		return err
	}
	// Read the response on a successful send
	buf := make([]byte, 256)
	n, err := c.readSmtpResponse(buf)
	ss.Response = string(buf[0:n])

	// Actually check return code
	if n < 5 {
		err = errors.New("Server did not indicate support for STARTTLS")
	}
	if err == nil {
		var ret int
		ret, err = strconv.Atoi(ss.Response[0:3])
		if err != nil || ret < 200 || ret >= 300 {
			err = errors.New("Bad return code for STARTTLS")
		}
	}

	// Record everything no matter the result
	c.appendEvent(&ss, err)

	// Stop if we failed already
	if err != nil {
		return err
	}

	// Successful so far, attempt to do the actual handshake
	return c.TLSHandshake()
}

func (c *Conn) POP3StartTLSHandshake() error {
	ss := StartTLSEvent{Command: POP3_COMMAND}
	if err := c.sendStartTLSCommand(POP3_COMMAND); err != nil {
		c.appendEvent(&ss, err)
		return err
	}

	buf := make([]byte, 512)
	n, err := c.readPop3Response(buf)
	ss.Response = string(buf[0:n])
	if err == nil {
		if !strings.HasPrefix(ss.Response, "+") {
			err = errors.New("Server did not indicate support for STARTTLS")
		}
	}
	c.appendEvent(&ss, err)

	if err != nil {
		return err
	}
	return c.TLSHandshake()
}

func (c *Conn) IMAPStartTLSHandshake() error {
	ss := StartTLSEvent{Command: IMAP_COMMAND}
	if err := c.sendStartTLSCommand(IMAP_COMMAND); err != nil {
		c.appendEvent(&ss, err)
		return err
	}

	buf := make([]byte, 512)
	n, err := c.readImapStatusResponse(buf)
	ss.Response = string(buf[0:n])
	if err == nil {
		if !strings.HasPrefix(ss.Response, "a001 OK") {
			err = errors.New("Server did not indicate support for STARTTLS")
		}
	}
	c.appendEvent(&ss, err)

	if err != nil {
		return err
	}
	return c.TLSHandshake()
}

func (c *Conn) readUntilRegex(res []byte, expr *regexp.Regexp) (int, error) {
	buf := res[0:]
	length := 0
	for finished := false; !finished; {
		n, err := c.getUnderlyingConn().Read(buf)
		length += n
		if err != nil {
			return length, err
		}
		if expr.Match(res[0:length]) {
			finished = true
		}
		if length == len(res) {
			return length, errors.New("Not enough buffer space")
		}
		buf = res[length:]
	}
	return length, nil
}

func (c *Conn) readSmtpResponse(res []byte) (int, error) {
	return c.readUntilRegex(res, smtpEndRegex)
}

func (c *Conn) SMTPBanner(b []byte) (int, error) {
	n, err := c.readSmtpResponse(b)
	mb := MailBannerEvent{}
	mb.Banner = string(b[0:n])
	c.appendEvent(&mb, err)
	return n, err
}

func (c *Conn) EHLO(domain string) error {
	cmd := []byte("EHLO " + domain + "\r\n")
	ee := EHLOEvent{}
	if _, err := c.getUnderlyingConn().Write(cmd); err != nil {
		c.appendEvent(&ee, err)
		return err
	}

	buf := make([]byte, 512)
	n, err := c.readSmtpResponse(buf)
	ee.Response = string(buf[0:n])
	c.appendEvent(&ee, err)
	return err
}

func (c *Conn) SMTPHelp() error {
	cmd := []byte("HELP\r\n")
	h := new(SMTPHelpEvent)
	if _, err := c.getUnderlyingConn().Write(cmd); err != nil {
		c.appendEvent(h, err)
		return err
	}
	buf := make([]byte, 512)
	n, err := c.readSmtpResponse(buf)
	h.Response = string(buf[0:n])
	c.appendEvent(h, err)
	return err
}

func (c *Conn) readPop3Response(res []byte) (int, error) {
	return c.readUntilRegex(res, pop3EndRegex)
}

func (c *Conn) POP3Banner(b []byte) (int, error) {
	n, err := c.readPop3Response(b)
	mb := MailBannerEvent{
		Banner: string(b[0:n]),
	}
	c.appendEvent(&mb, err)
	return n, err
}

func (c *Conn) readImapStatusResponse(res []byte) (int, error) {
	return c.readUntilRegex(res, imapStatusEndRegex)
}

func (c *Conn) IMAPBanner(b []byte) (int, error) {
	n, err := c.readImapStatusResponse(b)
	mb := MailBannerEvent{
		Banner: string(b[0:n]),
	}
	c.appendEvent(&mb, err)
	return n, err
}

func (c *Conn) CheckHeartbleed(b []byte) (int, error) {
	if !c.isTls {
		return 0, fmt.Errorf(
			"Must perform TLS handshake before sending Heartbleed probe to %s",
			c.RemoteAddr().String())
	}
	n, err := c.tlsConn.CheckHeartbleed(b)
	hb := c.tlsConn.GetHeartbleedLog()
	if err == ztls.HeartbleedError {
		err = nil
	}
	event := ConnectionEvent{
		Data:  &HeartbleedEvent{heartbleedLog: hb},
		Error: err,
	}
	c.operations = append(c.operations, event)
	return n, err
}

func (c *Conn) SendModbusEcho() (int, error) {
	req := ModbusRequest{
		Function: ModbusFunctionEncapsulatedInterface,
		Data: []byte{
			0x0E, // read device info
			0x01, // product code
			0x00, // object id, should always be 0 in initial request
		},
	}

	event := new(ModbusEvent)
	data, err := req.MarshalBinary()
	w := 0
	for w < len(data) {
		written, err := c.getUnderlyingConn().Write(data[w:]) // TODO verify write
		w += written
		if err != nil {
			c.appendEvent(event, err)
			return w, errors.New("Could not write modbus request")
		}
	}

	res, err := c.GetModbusResponse()
	event.Length = res.Length
	event.UnitID = res.UnitID
	event.Function = res.Function
	event.Response = res.Data
	event.ParseSelf()
	// make sure the whole thing gets appended to the operation log
	c.appendEvent(event, err)
	return w, err
}

func (c *Conn) GetFTPBanner() error {
	event := new(FTPBannerEvent)
	res := make([]byte, 1024)
	n, err := c.readUntilRegex(res, ftpEndRegex)
	event.Banner = string(res[0:n])
	c.appendEvent(event, err)
	return err
}

func (c *Conn) States() []ConnectionEvent {
	return c.operations
}

func (c *Conn) appendEvent(data EventData, err error) {
	event := ConnectionEvent{
		Data:  data,
		Error: err,
	}
	c.operations = append(c.operations, event)
}
<commit_msg>export ciphers now means RSA export ciphers<commit_after>package zlib

import (
	"errors"
	"fmt"
	"net"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/zmap/zgrab/ztools/x509"
	"github.com/zmap/zgrab/ztools/ztls"
)

var smtpEndRegex = regexp.MustCompile(`(?:\r\n)|^[0-9]{3} .+\r\n$`)
var pop3EndRegex = regexp.MustCompile(`(?:\r\n\.\r\n$)|(?:\r\n$)`)
var imapStatusEndRegex = regexp.MustCompile(`\r\n$`)
var ftpEndRegex = regexp.MustCompile(`^.*[0-9]{3}( [^\r\n]*)?\r?\n$`)

const (
	SMTP_COMMAND = "STARTTLS\r\n"
	POP3_COMMAND = "STLS\r\n"
	IMAP_COMMAND = "a001 STARTTLS\r\n"
)

// Implements the net.Conn interface
type Conn struct {
	// Underlying network connection
	conn    net.Conn
	tlsConn *ztls.Conn
	isTls   bool

	// Max TLS version
	maxTlsVersion uint16

	// Keep track of state / network operations
	operations []ConnectionEvent

	// Cache the deadlines so we can reapply after TLS handshake
	readDeadline  time.Time
	writeDeadline time.Time

	caPool *x509.CertPool

	onlyCBC       bool
	onlySchannel  bool
	onlyDHE       bool
	onlyExports   bool
	onlyExportsDH bool

	domain string

	// Encoding type
	ReadEncoding string
}

func (c *Conn) getUnderlyingConn() net.Conn {
	if c.isTls {
		return c.tlsConn
	}
	return c.conn
}

func (c *Conn) SetCBCOnly() {
	c.onlyCBC = true
}

func (c *Conn) SetSChannelOnly() {
	c.onlySchannel = true
}

func (c *Conn) SetDHEOnly() {
	c.onlyDHE = true
}

func (c *Conn) SetExportsOnly() {
	c.onlyExports = true
}

func (c *Conn) SetExportsDHOnly() {
	c.onlyExportsDH = true
}

func (c *Conn) SetCAPool(pool *x509.CertPool) {
	c.caPool = pool
}

func (c *Conn) SetDomain(domain string) {
	c.domain = domain
}

// Layer in the regular conn methods
func (c *Conn) LocalAddr() net.Addr {
	return c.getUnderlyingConn().LocalAddr()
}

func (c *Conn) RemoteAddr() net.Addr {
	return c.getUnderlyingConn().RemoteAddr()
}

func (c *Conn) SetDeadline(t time.Time) error {
	c.readDeadline = t
	c.writeDeadline = t
	return c.getUnderlyingConn().SetDeadline(t)
}

func (c *Conn) SetReadDeadline(t time.Time) error {
	c.readDeadline = t
	return c.getUnderlyingConn().SetReadDeadline(t)
}

func (c *Conn) SetWriteDeadline(t time.Time) error {
	c.writeDeadline = t
	return c.getUnderlyingConn().SetWriteDeadline(t)
}

// Delegate here, but record all the things
func (c *Conn) Write(b []byte) (int, error) {
	n, err := c.getUnderlyingConn().Write(b)
	w := WriteEvent{Sent: b}
	event := ConnectionEvent{
		Data:  &w,
		Error: err,
	}
	c.operations = append(c.operations, event)
	return n, err
}

func (c *Conn) Read(b []byte) (int, error) {
	n, err := c.getUnderlyingConn().Read(b)
	r := ReadEvent{
		Response: b[0:n],
		encoding: c.ReadEncoding,
	}
	c.appendEvent(&r, err)
	return n, err
}

func (c *Conn) Close() error {
	return c.getUnderlyingConn().Close()
}

// Extra method - Do a TLS Handshake and record progress
func (c *Conn) TLSHandshake() error {
	if c.isTls {
		return fmt.Errorf(
			"Attempted repeat handshake with remote host %s",
			c.RemoteAddr().String())
	}
	tlsConfig := new(ztls.Config)
	tlsConfig.InsecureSkipVerify = true
	tlsConfig.MinVersion = ztls.VersionSSL30
	tlsConfig.MaxVersion = c.maxTlsVersion
	tlsConfig.RootCAs = c.caPool
	if c.domain != "" {
		tlsConfig.ServerName = c.domain
	}
	if c.onlyCBC {
		tlsConfig.CipherSuites = ztls.CBCSuiteIDList
	}
	if c.onlySchannel {
		tlsConfig.CipherSuites = ztls.SChannelSuites
	}
	if c.onlyDHE {
		tlsConfig.CipherSuites = ztls.DHECiphers
		tlsConfig.ForceSuites = true
	}
	if c.onlyExports {
		tlsConfig.CipherSuites = ztls.RSA512ExportCiphers
		tlsConfig.ForceSuites = true
	}
	if c.onlyExportsDH {
		tlsConfig.CipherSuites = ztls.DHEExportCiphers
		tlsConfig.ForceSuites = true
	}
	c.tlsConn = ztls.Client(c.conn, tlsConfig)
	c.tlsConn.SetReadDeadline(c.readDeadline)
	c.tlsConn.SetWriteDeadline(c.writeDeadline)
	c.isTls = true
	err := c.tlsConn.Handshake()
	hl := c.tlsConn.GetHandshakeLog()
	ts := TLSHandshakeEvent{handshakeLog: hl}
	event := ConnectionEvent{
		Data:  &ts,
		Error: err,
	}
	c.operations = append(c.operations, event)
	return err
}

func (c *Conn) sendStartTLSCommand(command string) error {
	// Don't doublehandshake
	if c.isTls {
		return fmt.Errorf(
			"Attempt STARTTLS after TLS handshake with remote host %s",
			c.RemoteAddr().String())
	}
	// Send the STARTTLS message
	starttls := []byte(command)
	_, err := c.conn.Write(starttls)
	return err
}

// Do a STARTTLS handshake
func (c *Conn) SMTPStartTLSHandshake() error {
	// Make the state
	ss := StartTLSEvent{Command: SMTP_COMMAND}

	// Send the command
	if err := c.sendStartTLSCommand(SMTP_COMMAND); err != nil {
		c.appendEvent(&ss, err)
		return err
	}
	// Read the response on a successful send
	buf := make([]byte, 256)
	n, err := c.readSmtpResponse(buf)
	ss.Response = string(buf[0:n])

	// Actually check return code
	if n < 5 {
		err = errors.New("Server did not indicate support for STARTTLS")
	}
	if err == nil {
		var ret int
		ret, err = strconv.Atoi(ss.Response[0:3])
		if err != nil || ret < 200 || ret >= 300 {
			err = errors.New("Bad return code for STARTTLS")
		}
	}

	// Record everything no matter the result
	c.appendEvent(&ss, err)

	// Stop if we failed already
	if err != nil {
		return err
	}

	// Successful so far, attempt to do the actual handshake
	return c.TLSHandshake()
}

func (c *Conn) POP3StartTLSHandshake() error {
	ss := StartTLSEvent{Command: POP3_COMMAND}
	if err := c.sendStartTLSCommand(POP3_COMMAND); err != nil {
		c.appendEvent(&ss, err)
		return err
	}

	buf := make([]byte, 512)
	n, err := c.readPop3Response(buf)
	ss.Response = string(buf[0:n])
	if err == nil {
		if !strings.HasPrefix(ss.Response, "+") {
			err = errors.New("Server did not indicate support for STARTTLS")
		}
	}
	c.appendEvent(&ss, err)

	if err != nil {
		return err
	}
	return c.TLSHandshake()
}

func (c *Conn) IMAPStartTLSHandshake() error {
	ss := StartTLSEvent{Command: IMAP_COMMAND}
	if err := c.sendStartTLSCommand(IMAP_COMMAND); err != nil {
		c.appendEvent(&ss, err)
		return err
	}

	buf := make([]byte, 512)
	n, err := c.readImapStatusResponse(buf)
	ss.Response = string(buf[0:n])
	if err == nil {
		if !strings.HasPrefix(ss.Response, "a001 OK") {
			err = errors.New("Server did not indicate support for STARTTLS")
		}
	}
	c.appendEvent(&ss, err)

	if err != nil {
		return err
	}
	return c.TLSHandshake()
}

func (c *Conn) readUntilRegex(res []byte, expr *regexp.Regexp) (int, error) {
	buf := res[0:]
	length := 0
	for finished := false; !finished; {
		n, err := c.getUnderlyingConn().Read(buf)
		length += n
		if err != nil {
			return length, err
		}
		if expr.Match(res[0:length]) {
			finished = true
		}
		if length == len(res) {
			return length, errors.New("Not enough buffer space")
		}
		buf = res[length:]
	}
	return length, nil
}

func (c *Conn) readSmtpResponse(res []byte) (int, error) {
	return c.readUntilRegex(res, smtpEndRegex)
}

func (c *Conn) SMTPBanner(b []byte) (int, error) {
	n, err := c.readSmtpResponse(b)
	mb := MailBannerEvent{}
	mb.Banner = string(b[0:n])
	c.appendEvent(&mb, err)
	return n, err
}

func (c *Conn) EHLO(domain string) error {
	cmd := []byte("EHLO " + domain + "\r\n")
	ee := EHLOEvent{}
	if _, err := c.getUnderlyingConn().Write(cmd); err != nil {
		c.appendEvent(&ee, err)
		return err
	}

	buf := make([]byte, 512)
	n, err := c.readSmtpResponse(buf)
	ee.Response = string(buf[0:n])
	c.appendEvent(&ee, err)
	return err
}

func (c *Conn) SMTPHelp() error {
	cmd := []byte("HELP\r\n")
	h := new(SMTPHelpEvent)
	if _, err := c.getUnderlyingConn().Write(cmd); err != nil {
		c.appendEvent(h, err)
		return err
	}
	buf := make([]byte, 512)
	n, err := c.readSmtpResponse(buf)
	h.Response = string(buf[0:n])
	c.appendEvent(h, err)
	return err
}

func (c *Conn) readPop3Response(res []byte) (int, error) {
	return c.readUntilRegex(res, pop3EndRegex)
}

func (c *Conn) POP3Banner(b []byte) (int, error) {
	n, err := c.readPop3Response(b)
	mb := MailBannerEvent{
		Banner: string(b[0:n]),
	}
	c.appendEvent(&mb, err)
	return n, err
}

func (c *Conn) readImapStatusResponse(res []byte) (int, error) {
	return c.readUntilRegex(res, imapStatusEndRegex)
}

func (c *Conn) IMAPBanner(b []byte) (int, error) {
	n, err := c.readImapStatusResponse(b)
	mb := MailBannerEvent{
		Banner: string(b[0:n]),
	}
	c.appendEvent(&mb, err)
	return n, err
}

func (c *Conn) CheckHeartbleed(b []byte) (int, error) {
	if !c.isTls {
		return 0, fmt.Errorf(
			"Must perform TLS handshake before sending Heartbleed probe to %s",
			c.RemoteAddr().String())
	}
	n, err := c.tlsConn.CheckHeartbleed(b)
	hb := c.tlsConn.GetHeartbleedLog()
	if err == ztls.HeartbleedError {
		err = nil
	}
	event := ConnectionEvent{
		Data:  &HeartbleedEvent{heartbleedLog: hb},
		Error: err,
	}
	c.operations = append(c.operations, event)
	return n, err
}

func (c *Conn) SendModbusEcho() (int, error) {
	req := ModbusRequest{
		Function: ModbusFunctionEncapsulatedInterface,
		Data: []byte{
			0x0E, // read device info
			0x01, // product code
			0x00, // object id, should always be 0 in initial request
		},
	}

	event := new(ModbusEvent)
	data, err := req.MarshalBinary()
	w := 0
	for w < len(data) {
		written, err := c.getUnderlyingConn().Write(data[w:]) // TODO verify write
		w += written
		if err != nil {
			c.appendEvent(event, err)
			return w, errors.New("Could not write modbus request")
		}
	}

	res, err := c.GetModbusResponse()
	event.Length = res.Length
	event.UnitID = res.UnitID
	event.Function = res.Function
	event.Response = res.Data
	event.ParseSelf()
	// make sure the whole thing gets appended to the operation log
	c.appendEvent(event, err)
	return w, err
}

func (c *Conn) GetFTPBanner() error {
	event := new(FTPBannerEvent)
	res := make([]byte, 1024)
	n, err := c.readUntilRegex(res, ftpEndRegex)
	event.Banner = string(res[0:n])
	c.appendEvent(event, err)
	return err
}

func (c *Conn) States() []ConnectionEvent {
	return c.operations
}

func (c *Conn) appendEvent(data EventData, err error) {
	event := ConnectionEvent{
		Data:  data,
		Error: err,
	}
	c.operations = append(c.operations, event)
}
<|endoftext|>


<commit_before>package runstats

import "os"
import "fmt"
import "flag"
import "time"
import "strconv"
import "runtime"

import "github.com/bmhatfield/g2s"

var pause *int = flag.Int("pause", 10, "Collection pause interval")
var CPU *bool = flag.Bool("cpu", true, "Collect CPU Statistics")
var MEM *bool = flag.Bool("mem", true, "Collect Memory Statistics")
var GC *bool = flag.Bool("gc", true, "Collect GC Statistics (requires Memory be enabled)")

var prefix string
var s g2s.Statter

var err error

func init() {
	s, err = g2s.Dial("udp", "localhost:8125")

	if err != nil {
		panic("Unable to connect to Statsd")
	}

	prefix, err = os.Hostname()

	if err != nil {
		prefix = "unknown_host.go"
	} else {
		prefix = fmt.Sprintf("%s.go", prefix)
	}

	go collector()
}

func collector() {
	for {
		if *CPU {
			// Goroutines
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "cpu.goroutines"), strconv.Itoa(runtime.NumGoroutine()))

			// CGo calls
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "cpu.cgo_calls"), strconv.FormatUint(uint64(runtime.NumCgoCall()), 10))
		}

		if *MEM {
			m := new(runtime.MemStats)
			runtime.ReadMemStats(m)

			// General
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.alloc"), strconv.FormatUint(m.Alloc, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.total"), strconv.FormatUint(m.TotalAlloc, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.sys"), strconv.FormatUint(m.Sys, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.lookups"), strconv.FormatUint(m.Lookups, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.malloc"), strconv.FormatUint(m.Mallocs, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.frees"), strconv.FormatUint(m.Frees, 10))

			// Heap
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.heap.alloc"), strconv.FormatUint(m.HeapAlloc, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.heap.sys"), strconv.FormatUint(m.HeapSys, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.heap.idle"), strconv.FormatUint(m.HeapIdle, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.heap.inuse"), strconv.FormatUint(m.HeapInuse, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.heap.released"), strconv.FormatUint(m.HeapReleased, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.heap.objects"), strconv.FormatUint(m.HeapObjects, 10))

			// Stack
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.stack.inuse"), strconv.FormatUint(m.StackInuse, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.stack.sys"), strconv.FormatUint(m.StackSys, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.stack.mspan_inuse"), strconv.FormatUint(m.MSpanInuse, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.stack.mspan_sys"), strconv.FormatUint(m.MSpanSys, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.stack.mcache_inuse"), strconv.FormatUint(m.MCacheInuse, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.stack.mcache_sys"), strconv.FormatUint(m.MCacheSys, 10))

			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.othersys"), strconv.FormatUint(m.OtherSys, 10))

			if *GC {
				s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.gc.sys"), strconv.FormatUint(m.GCSys, 10))
				s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.gc.next"), strconv.FormatUint(m.NextGC, 10))
				s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.gc.last"), strconv.FormatUint(m.LastGC, 10))
				s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.gc.pause_total"), strconv.FormatUint(m.PauseTotalNs, 10))
				s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.gc.pause"), strconv.FormatUint(m.PauseNs[(m.NumGC+255)%256], 10))
				s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.gc.count"), strconv.FormatUint(uint64(m.NumGC), 10))
			}
		}

		// Don't collect for a time interval
		time.Sleep(time.Duration(*pause) * time.Second)
	}
}
<commit_msg>Defer actual work to goroutine<commit_after>package runstats

import "os"
import "fmt"
import "flag"
import "time"
import "strconv"
import "runtime"

import "github.com/bmhatfield/g2s"

var statsd *string = flag.String("statsd", "localhost:8125", "Statsd host:port pair")
var user_prefix *string = flag.String("metric-prefix", "default", "Metric prefix path; detects the local hostname by default")
var pause *int = flag.Int("pause", 10, "Collection pause interval")

var CPU *bool = flag.Bool("cpu", true, "Collect CPU Statistics")
var MEM *bool = flag.Bool("mem", true, "Collect Memory Statistics")
var GC *bool = flag.Bool("gc", true, "Collect GC Statistics (requires Memory be enabled)")

var prefix string
var s g2s.Statter

var err error

func init() {
	go collector()
}

func collector() {
	if !flag.Parsed() {
		flag.Parse()
	}

	s, err = g2s.Dial("udp", *statsd)

	if err != nil {
		panic(fmt.Sprintf("Unable to connect to Statsd on %s - %s", *statsd, err))
	}

	if *user_prefix == "default" {
		prefix, err = os.Hostname()

		if err != nil {
			prefix = "unknown_host.go"
		} else {
			prefix = fmt.Sprintf("%s.go", prefix)
		}
	} else {
		prefix = *user_prefix
	}

	for {
		if *CPU {
			// Goroutines
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "cpu.goroutines"), strconv.Itoa(runtime.NumGoroutine()))

			// CGo calls
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "cpu.cgo_calls"), strconv.FormatUint(uint64(runtime.NumCgoCall()), 10))
		}

		if *MEM {
			m := new(runtime.MemStats)
			runtime.ReadMemStats(m)

			// General
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.alloc"), strconv.FormatUint(m.Alloc, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.total"), strconv.FormatUint(m.TotalAlloc, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.sys"), strconv.FormatUint(m.Sys, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.lookups"), strconv.FormatUint(m.Lookups, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.malloc"), strconv.FormatUint(m.Mallocs, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.frees"), strconv.FormatUint(m.Frees, 10))

			// Heap
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.heap.alloc"), strconv.FormatUint(m.HeapAlloc, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.heap.sys"), strconv.FormatUint(m.HeapSys, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.heap.idle"), strconv.FormatUint(m.HeapIdle, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.heap.inuse"), strconv.FormatUint(m.HeapInuse, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.heap.released"), strconv.FormatUint(m.HeapReleased, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.heap.objects"), strconv.FormatUint(m.HeapObjects, 10))

			// Stack
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.stack.inuse"), strconv.FormatUint(m.StackInuse, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.stack.sys"), strconv.FormatUint(m.StackSys, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.stack.mspan_inuse"), strconv.FormatUint(m.MSpanInuse, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.stack.mspan_sys"), strconv.FormatUint(m.MSpanSys, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.stack.mcache_inuse"), strconv.FormatUint(m.MCacheInuse, 10))
			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.stack.mcache_sys"), strconv.FormatUint(m.MCacheSys, 10))

			s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.othersys"), strconv.FormatUint(m.OtherSys, 10))

			if *GC {
				// GC
				s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.gc.sys"), strconv.FormatUint(m.GCSys, 10))
				s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.gc.next"), strconv.FormatUint(m.NextGC, 10))
				s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.gc.last"), strconv.FormatUint(m.LastGC, 10))
				s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.gc.pause_total"), strconv.FormatUint(m.PauseTotalNs, 10))
				s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.gc.pause"), strconv.FormatUint(m.PauseNs[(m.NumGC+255)%256], 10))
				s.Gauge(1.0, fmt.Sprintf("%s.%s", prefix, "mem.gc.count"), strconv.FormatUint(uint64(m.NumGC), 10))
			}
		}

		// Gauges are a 'snapshot' rather than a histogram. Pausing for some interval
		// aims to get a 'recent' snapshot out before statsd flushes metrics.
		time.Sleep(time.Duration(*pause) * time.Second)
	}
}
<|endoftext|>


<commit_before>package main

import (
  "github.com/hayesgm/fiddler/config"
  "github.com/hayesgm/fiddler/tracker"
  "github.com/hayesgm/fiddler/installer"
  "github.com/hayesgm/fiddler/launcher"
  "github.com/hayesgm/fiddler/spawner"
  "github.com/coreos/go-etcd/etcd"
  "github.com/go-contrib/uuid"
  "log"
  "flag"
  "os/exec"
  "fmt"
)

/*
  Fiddler is a daemon that will check stats amoungst a
  set of nodes in an etcd cluster.  If certain thresholds
  are exceeded across the stack, fiddle will issue
  auto-scale commands to expand the cluster.
*/

/*
  Fiddler knows how to spawn itself into CoreOS
  and start docker containers.  Thus, Fiddler is
  quickly able to coordinate replicating itself
  as load dictates.  It's also able to detect drops
  in load, where it will kill excess servers.
*/

/*
  Fiddler uses a Fidderfile for all configuration settings.
  ./fiddler -c fiddler.conf
*/

var cli = etcd.NewClient()
var myid = uuid.NewV1().String()

func printUsage() {
  fmt.Println("Fiddler is an app to auto-scale containers.")
  fmt.Println("")
  fmt.Println("Usage:")
  fmt.Println("fiddler --config=<conf> install")
  fmt.Println("\tInstalls Fiddler into systemd")
  fmt.Println("fiddler --config=<conf> launch")
  fmt.Println("\tLaunches a new container within Fiddler")
  fmt.Println("fiddler --config=<conf> spawn")
  fmt.Println("\tSpawns a new fiddler environment in the cloud")
  fmt.Println("fiddler --config=<conf> daemon")
  fmt.Println("\tRuns Fiddler like a daemon to monitor machine")
}

func main() {
  // First, let's make sure we have a configuration file
  var c = flag.String("config", "", "location of configuration file (http okay)")
  flag.Parse()

  if len(flag.Args()) < 1 {
    printUsage()
    return
  }

  if *c == "" {
    fmt.Println("Missing config option")
    printUsage()
    return
  }

  // Next, we'll load the config
  var conf *config.FiddlerConf
  conf, err := config.LoadFiddlerConfig(*c)

  if err != nil {
    log.Fatal("Unable to load config file:", err)
  }

  var cmd = flag.Args()[0]

  switch cmd {
  case "i", "install":
    // We should install Fiddler
    err := installer.InstallFiddler(*c, *conf)
    if err != nil {
      log.Fatal("Error installing Fiddler:", err)
    }
  case "s", "spawn":
    // We'll grab spawn pool and grow it
    pool, err := spawner.GetSpawnPool()
    if err != nil {
      log.Fatal("Error getting spawn pool:", err)
    }

    err = pool.Grow()
    if err != nil {
      log.Fatal("Error growing spawn pool:", err)
    }
  case "l", "launch":
    // We should launch our container
    var cmd *exec.Cmd
    cmd, err = launcher.Launch(conf.Docker)
    if err != nil {
      log.Fatal("Error launching container:",conf.Docker,err)
    }

    err = cmd.Wait() // We'll stay open so long as the docker is open.  We can ensure the docker stays open, etc.
    if err != nil {
      log.Fatal("Failed to run container:",conf.Docker,err)
    }
  case "d", "daemon":
    // We should be like a daemon, tracking stats

    // Now, we're going to make sure we're monitoring our stats
    go tracker.TrackMyStats(cli, myid, []string{"cpu"})
    go tracker.WatchStats(cli) // TODO: This should only be for leader node

    ch := make(chan int)
    <- ch // Hold forever
  default:
    printUsage()
    return
  }
}<commit_msg>Fix command line syntax and start to get spawner running<commit_after>package main

import (
  "github.com/hayesgm/fiddler/config"
  "github.com/hayesgm/fiddler/tracker"
  "github.com/hayesgm/fiddler/installer"
  "github.com/hayesgm/fiddler/launcher"
  "github.com/hayesgm/fiddler/spawner"
  "github.com/coreos/go-etcd/etcd"
  "github.com/go-contrib/uuid"
  "log"
  "flag"
  "os/exec"
  "fmt"
)

/*
  Fiddler is a daemon that will check stats amoungst a
  set of nodes in an etcd cluster.  If certain thresholds
  are exceeded across the stack, fiddle will issue
  auto-scale commands to expand the cluster.
*/

/*
  Fiddler knows how to spawn itself into CoreOS
  and start docker containers.  Thus, Fiddler is
  quickly able to coordinate replicating itself
  as load dictates.  It's also able to detect drops
  in load, where it will kill excess servers.
*/

/*
  Fiddler uses a Fidderfile for all configuration settings.
  ./fiddler -c fiddler.conf
*/

var cli = etcd.NewClient()
var myid = uuid.NewV1().String()

func printUsage() {
  fmt.Println("Fiddler is an app to auto-scale containers.")
  fmt.Println("")
  fmt.Println("Usage:")
  fmt.Println("fiddler --config=<conf> install")
  fmt.Println("\tInstalls Fiddler into systemd")
  fmt.Println("fiddler --config=<conf> launch")
  fmt.Println("\tLaunches a new container within Fiddler")
  fmt.Println("fiddler --config=<conf> spawn <name>")
  fmt.Println("\tSpawns a new fiddler environment in the cloud")
  fmt.Println("fiddler --config=<conf> daemon")
  fmt.Println("\tRuns Fiddler like a daemon to monitor machine")
}

func main() {
  // First, let's make sure we have a configuration file
  var c = flag.String("config", "", "location of configuration file (http okay)")
  flag.Parse()

  if len(flag.Args()) < 1 {
    printUsage()
    return
  }

  if *c == "" {
    fmt.Println("Missing config option")
    printUsage()
    return
  }

  // Next, we'll load the config
  var conf *config.FiddlerConf
  conf, err := config.LoadFiddlerConfig(*c)

  if err != nil {
    log.Fatal("Unable to load config file:", err)
  }

  var cmd = flag.Args()[0]

  switch cmd {
  case "i", "install":
    // We should install Fiddler
    err := installer.InstallFiddler(*c, *conf)
    if err != nil {
      log.Fatal("Error installing Fiddler:", err)
    }
  case "s", "spawn":
    if len(flag.Args()) < 2 {
      fmt.Println("Must provide env name")
      printUsage()
      return
    }

    name := flag.Args()[1]

    // We'll grab spawn pool and grow it
    pool, err := spawner.GetSpawnPool(name)
    if err != nil {
      log.Fatal("Error getting spawn pool:", err)
    }

    err = pool.Grow(*c)
    if err != nil {
      log.Fatal("Error growing spawn pool:", err)
    }
  case "l", "launch":
    // We should launch our container
    var cmd *exec.Cmd
    cmd, err = launcher.Launch(conf.Docker)
    if err != nil {
      log.Fatal("Error launching container:",conf.Docker,err)
    }

    err = cmd.Wait() // We'll stay open so long as the docker is open.  We can ensure the docker stays open, etc.
    if err != nil {
      log.Fatal("Failed to run container:",conf.Docker,err)
    }
  case "d", "daemon":
    // We should be like a daemon, tracking stats

    // Now, we're going to make sure we're monitoring our stats
    go tracker.TrackMyStats(cli, myid, []string{"cpu"})
    go tracker.WatchStats(cli) // TODO: This should only be for leader node

    ch := make(chan int)
    <- ch // Hold forever
  default:
    printUsage()
    return
  }
}<|endoftext|>


<commit_before>package role

import (
	"strings"

	"github.com/admpub/log"
	"github.com/admpub/nging/v4/application/library/perm"
	"github.com/admpub/nging/v4/application/registry/navigate"
	"github.com/webx-top/echo"
)

func NewRolePermission() *RolePermission {
	r := &RolePermission{}
	r.CommonPermission = NewCommonPermission(UserRolePermissionType, r)
	return r
}

type RolePermission struct {
	*CommonPermission
	Roles []*UserRoleWithPermissions
}

func (r *RolePermission) Init(roleList []*UserRoleWithPermissions) *RolePermission {
	r.Roles = roleList
	gts := make([]PermissionsGetter, len(r.Roles))
	for k, v := range r.Roles {
		gts[k] = v
	}
	r.CommonPermission.Init(gts)
	return r
}

func NewCommonPermission(d *echo.KVData, c navigate.Checker) *CommonPermission {
	return &CommonPermission{
		DefinedType: d,
		filter:      navigate.NewFilter(c),
		Combined:    map[string]string{},
		parsed:      map[string]interface{}{},
	}
}

type CommonPermission struct {
	DefinedType *echo.KVData
	Combined    map[string]string
	parsed      map[string]interface{}
	filter      *navigate.Filter
}

type PermissionConfiger interface {
	GetType() string
	GetPermission() string
}

type PermissionsGetter interface {
	GetPermissions() []PermissionConfiger
}

func (r *CommonPermission) Init(roleList []PermissionsGetter) *CommonPermission {
	checkeds := map[string]map[string]struct{}{}
	seperators := map[string]string{}
	for _, role := range roleList {
		for _, rolePerm := range role.GetPermissions() {
			if _, ok := checkeds[rolePerm.GetType()]; !ok {
				checkeds[rolePerm.GetType()] = map[string]struct{}{}
			}
			if _, ok := checkeds[rolePerm.GetType()][`*`]; ok {
				continue
			}
			for _, pa := range strings.Split(rolePerm.GetPermission(), `,`) {
				if _, ok := checkeds[rolePerm.GetType()][pa]; !ok {
					checkeds[rolePerm.GetType()][pa] = struct{}{}
					r.Combined[rolePerm.GetType()] += seperators[rolePerm.GetType()] + pa
					seperators[rolePerm.GetType()] = `,`
				}
			}
		}
	}
	return r
}

func (r *CommonPermission) onceParse(ctx echo.Context, typ string) bool {
	if _, ok := r.parsed[typ]; !ok {
		item := r.DefinedType.GetItem(typ)
		if item == nil {
			return false
		}
		var err error
		r.parsed[typ], err = item.X.(*perm.Handle).Parse(ctx, r.Combined[typ])
		//echo.Dump(echo.H{`combined`: r.Combined, `parsed`: r.parsed[typ]})
		if err != nil {
			log.Error(err)
			return false
		}
	}
	return true
}

func (r *CommonPermission) CheckByType(ctx echo.Context, typ string, permPath string) interface{} {
	if !r.onceParse(ctx, typ) {
		return nil
	}
	rs, err := perm.HandleCheck(ctx, r.DefinedType, permPath, typ, r.Combined[typ], r.parsed[typ])
	if err != nil {
		log.Error(err)
	}
	return rs
}

func (r *CommonPermission) Check(ctx echo.Context, permPath string) bool {
	rs := r.CheckByType(ctx, RolePermissionTypePage, permPath)
	if rs == nil {
		return false
	}
	if bl, ok := rs.(bool); ok {
		return bl
	}
	return false
}

func (r *CommonPermission) CheckCmd(ctx echo.Context, permPath string) bool {
	rs := r.CheckByType(ctx, RolePermissionTypeCommand, permPath)
	if rs == nil {
		return false
	}
	if bl, ok := rs.(bool); ok {
		return bl
	}
	return false
}

func (r *CommonPermission) CheckBehavior(ctx echo.Context, permPath string) *perm.CheckedBehavior {
	rs := r.CheckByType(ctx, RolePermissionTypeBehavior, permPath)
	if rs == nil {
		return &perm.CheckedBehavior{}
	}
	if bv, ok := rs.(*perm.CheckedBehavior); ok {
		return bv
	}
	return &perm.CheckedBehavior{}
}

//FilterNavigate 过滤导航菜单，只显示有权限的菜单
func (r *CommonPermission) FilterNavigate(ctx echo.Context, navList *navigate.List) navigate.List {
	return r.filter.FilterNavigate(ctx, navList)
}
<commit_msg>update<commit_after>package role

import (
	"strings"

	"github.com/admpub/log"
	"github.com/admpub/nging/v4/application/library/perm"
	"github.com/admpub/nging/v4/application/registry/navigate"
	"github.com/webx-top/echo"
)

func NewRolePermission() *RolePermission {
	r := &RolePermission{}
	r.CommonPermission = NewCommonPermission(UserRolePermissionType, r)
	return r
}

type RolePermission struct {
	*CommonPermission
	Roles []*UserRoleWithPermissions
}

func (r *RolePermission) Init(roleList []*UserRoleWithPermissions) *RolePermission {
	r.Roles = roleList
	gts := make([]PermissionsGetter, len(r.Roles))
	for k, v := range r.Roles {
		gts[k] = v
	}
	r.CommonPermission.Init(gts)
	return r
}

func NewCommonPermission(d *echo.KVData, c navigate.Checker) *CommonPermission {
	return &CommonPermission{
		DefinedType: d,
		filter:      navigate.NewFilter(c),
		Combined:    map[string]string{},
		parsed:      map[string]interface{}{},
	}
}

type CommonPermission struct {
	DefinedType *echo.KVData
	Combined    map[string]string
	parsed      map[string]interface{}
	filter      *navigate.Filter
}

type PermissionConfiger interface {
	GetType() string
	GetPermission() string
}

type PermissionsGetter interface {
	GetPermissions() []PermissionConfiger
}

func (r *CommonPermission) Init(roleList []PermissionsGetter) *CommonPermission {
	checkeds := map[string]map[string]struct{}{}
	seperators := map[string]string{}
	for _, role := range roleList {
		for _, rolePerm := range role.GetPermissions() {
			if _, ok := checkeds[rolePerm.GetType()]; !ok {
				checkeds[rolePerm.GetType()] = map[string]struct{}{}
			}
			if _, ok := checkeds[rolePerm.GetType()][`*`]; ok {
				continue
			}
			for _, pa := range strings.Split(rolePerm.GetPermission(), `,`) {
				if _, ok := checkeds[rolePerm.GetType()][pa]; !ok {
					checkeds[rolePerm.GetType()][pa] = struct{}{}
					r.Combined[rolePerm.GetType()] += seperators[rolePerm.GetType()] + pa
					seperators[rolePerm.GetType()] = `,`
				}
			}
		}
	}
	return r
}

func (r *CommonPermission) onceParse(ctx echo.Context, typ string) bool {
	if _, ok := r.parsed[typ]; !ok {
		item := r.DefinedType.GetItem(typ)
		if item == nil {
			return false
		}
		var err error
		r.parsed[typ], err = item.X.(*perm.Handle).Parse(ctx, r.Combined[typ])
		//echo.Dump(echo.H{`combined`: r.Combined, `parsed`: r.parsed[typ]})
		if err != nil {
			log.Error(err)
			return false
		}
	}
	return true
}

func (r *CommonPermission) Get(ctx echo.Context, typ string) interface{} {
	if !r.onceParse(ctx, typ) {
		return nil
	}
	return r.parsed[typ]
}

func (r *CommonPermission) CheckByType(ctx echo.Context, typ string, permPath string) interface{} {
	if !r.onceParse(ctx, typ) {
		return nil
	}
	rs, err := perm.HandleCheck(ctx, r.DefinedType, permPath, typ, r.Combined[typ], r.parsed[typ])
	if err != nil {
		log.Error(err)
	}
	return rs
}

func (r *CommonPermission) Check(ctx echo.Context, permPath string) bool {
	rs := r.CheckByType(ctx, RolePermissionTypePage, permPath)
	if rs == nil {
		return false
	}
	if bl, ok := rs.(bool); ok {
		return bl
	}
	return false
}

func (r *CommonPermission) CheckCmd(ctx echo.Context, permPath string) bool {
	rs := r.CheckByType(ctx, RolePermissionTypeCommand, permPath)
	if rs == nil {
		return false
	}
	if bl, ok := rs.(bool); ok {
		return bl
	}
	return false
}

func (r *CommonPermission) CheckBehavior(ctx echo.Context, permPath string) *perm.CheckedBehavior {
	rs := r.CheckByType(ctx, RolePermissionTypeBehavior, permPath)
	if rs == nil {
		return &perm.CheckedBehavior{}
	}
	if bv, ok := rs.(*perm.CheckedBehavior); ok {
		return bv
	}
	return &perm.CheckedBehavior{}
}

//FilterNavigate 过滤导航菜单，只显示有权限的菜单
func (r *CommonPermission) FilterNavigate(ctx echo.Context, navList *navigate.List) navigate.List {
	return r.filter.FilterNavigate(ctx, navList)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/md5"
	_ "expvar"
	"fmt"
	"github.com/quii/mockingjay-server/mockingjay"
	"github.com/quii/mockingjay-server/monkey"
	"io/ioutil"
	"log"
	"net/http"
	"path/filepath"
	"strings"
	"time"
)

var (
	// ErrCDCFail describes when a fake server is not compatible with a given URL
	ErrCDCFail = fmt.Errorf("At least one endpoint was incompatible with the real URL supplied")
)

type configLoader func(string) ([][]byte, []string, error)
type mockingjayLoader func([]byte) ([]mockingjay.FakeEndpoint, error)

type compatabilityChecker interface {
	CheckCompatibility(endpoints []mockingjay.FakeEndpoint, realURL string) bool
}

type serverMaker func([]mockingjay.FakeEndpoint, bool) *mockingjay.Server
type monkeyServerMaker func(http.Handler, string) (http.Handler, error)

type application struct {
	configLoader          configLoader
	mockingjayLoader      mockingjayLoader
	compatabilityChecker  compatabilityChecker
	mockingjayServerMaker serverMaker
	monkeyServerMaker     monkeyServerMaker
	logger                *log.Logger
	mjServer              *mockingjay.Server
	configPath            string
	monkeyConfigPath      string
	yamlMD5               [md5.Size]byte
}

func defaultApplication(logger *log.Logger, httpTimeout time.Duration) (app *application) {
	app = new(application)
	app.configLoader = globFileLoader
	app.mockingjayLoader = mockingjay.NewFakeEndpoints
	app.compatabilityChecker = NewCompatabilityChecker(logger, httpTimeout)
	app.mockingjayServerMaker = mockingjay.NewServer
	app.monkeyServerMaker = monkey.NewServer
	app.logger = logger

	return
}

func (a *application) PollConfig() {
	for range time.Tick(time.Millisecond * 500) {
		endpoints, err := a.loadConfig()
		if err != nil {
			log.Println(err)
		} else if len(endpoints) > 0 {
			a.logger.Println("Reloaded config")
			a.mjServer.Endpoints = endpoints
		}
	}
}

// CreateServer will create a fake server from the configuration found in configPath with optional performance constraints from configutation found in monkeyConfigPath
func (a *application) CreateServer(configPath string, monkeyConfigPath string, debugMode bool) (server http.Handler, err error) {
	a.configPath = configPath
	a.monkeyConfigPath = monkeyConfigPath
	endpoints, err := a.loadConfig()

	if err != nil || len(endpoints) == 0 {
		return
	}

	return a.createFakeServer(endpoints, debugMode)
}

// CheckCompatibility will run a MJ config against a realURL to see if it's compatible
func (a *application) CheckCompatibility(configPath string, realURL string) error {
	a.configPath = configPath
	endpoints, err := a.loadConfig()

	if err != nil {
		return err
	}

	if a.compatabilityChecker.CheckCompatibility(endpoints, realURL) {
		a.logger.Println("All endpoints are compatible")
		return nil
	}

	return ErrCDCFail
}

func (a *application) loadConfig() (endpoints []mockingjay.FakeEndpoint, err error) {

	configs, paths, err := a.configLoader(a.configPath)

	if err != nil {
		return
	}

	if newMD5 := md5.Sum(bytes.Join(configs, []byte{})); newMD5 != a.yamlMD5 {
		a.yamlMD5 = newMD5

		for _, configData := range configs {
			mjEndpoint, err := a.mockingjayLoader(configData)

			if err != nil {
				return nil, err
			}

			endpoints = append(endpoints, mjEndpoint...)
		}
	}

	return
}

func (a *application) createFakeServer(endpoints []mockingjay.FakeEndpoint, debugMode bool) (server http.Handler, err error) {
	go a.PollConfig()
	a.mjServer = a.mockingjayServerMaker(endpoints, debugMode)
	monkeyServer, err := a.monkeyServerMaker(a.mjServer, a.monkeyConfigPath)

	if err != nil {
		return nil, err
	}

	router := http.NewServeMux()
	router.Handle("/", monkeyServer)

	return router, nil
}

func globFileLoader(path string) (data [][]byte, paths []string, err error) {
	files, err := filepath.Glob(path)

	if err != nil {
		return nil, nil, fmt.Errorf("Failed to get files from file path (glob) %s, %v", path, err)
	}

	if len(files) == 0 {
		return nil, nil, fmt.Errorf("No files found in path %s", path)
	}

	var configs [][]byte
	for _, file := range files {

		configData, err := ioutil.ReadFile(file)

		if err != nil {
			return nil, nil, err
		}

		configs = append(configs, configData)
	}

	return configs, files, nil
}
<commit_msg>too hasty<commit_after>package main

import (
	"bytes"
	"crypto/md5"
	_ "expvar"
	"fmt"
	"github.com/quii/mockingjay-server/mockingjay"
	"github.com/quii/mockingjay-server/monkey"
	"io/ioutil"
	"log"
	"net/http"
	"path/filepath"
	"time"
)

var (
	// ErrCDCFail describes when a fake server is not compatible with a given URL
	ErrCDCFail = fmt.Errorf("At least one endpoint was incompatible with the real URL supplied")
)

type configLoader func(string) ([][]byte, []string, error)
type mockingjayLoader func([]byte) ([]mockingjay.FakeEndpoint, error)

type compatabilityChecker interface {
	CheckCompatibility(endpoints []mockingjay.FakeEndpoint, realURL string) bool
}

type serverMaker func([]mockingjay.FakeEndpoint, bool) *mockingjay.Server
type monkeyServerMaker func(http.Handler, string) (http.Handler, error)

type application struct {
	configLoader          configLoader
	mockingjayLoader      mockingjayLoader
	compatabilityChecker  compatabilityChecker
	mockingjayServerMaker serverMaker
	monkeyServerMaker     monkeyServerMaker
	logger                *log.Logger
	mjServer              *mockingjay.Server
	configPath            string
	monkeyConfigPath      string
	yamlMD5               [md5.Size]byte
}

func defaultApplication(logger *log.Logger, httpTimeout time.Duration) (app *application) {
	app = new(application)
	app.configLoader = globFileLoader
	app.mockingjayLoader = mockingjay.NewFakeEndpoints
	app.compatabilityChecker = NewCompatabilityChecker(logger, httpTimeout)
	app.mockingjayServerMaker = mockingjay.NewServer
	app.monkeyServerMaker = monkey.NewServer
	app.logger = logger

	return
}

func (a *application) PollConfig() {
	for range time.Tick(time.Millisecond * 500) {
		endpoints, err := a.loadConfig()
		if err != nil {
			log.Println(err)
		} else if len(endpoints) > 0 {
			a.logger.Println("Reloaded config")
			a.mjServer.Endpoints = endpoints
		}
	}
}

// CreateServer will create a fake server from the configuration found in configPath with optional performance constraints from configutation found in monkeyConfigPath
func (a *application) CreateServer(configPath string, monkeyConfigPath string, debugMode bool) (server http.Handler, err error) {
	a.configPath = configPath
	a.monkeyConfigPath = monkeyConfigPath
	endpoints, err := a.loadConfig()

	if err != nil || len(endpoints) == 0 {
		return
	}

	return a.createFakeServer(endpoints, debugMode)
}

// CheckCompatibility will run a MJ config against a realURL to see if it's compatible
func (a *application) CheckCompatibility(configPath string, realURL string) error {
	a.configPath = configPath
	endpoints, err := a.loadConfig()

	if err != nil {
		return err
	}

	if a.compatabilityChecker.CheckCompatibility(endpoints, realURL) {
		a.logger.Println("All endpoints are compatible")
		return nil
	}

	return ErrCDCFail
}

func (a *application) loadConfig() (endpoints []mockingjay.FakeEndpoint, err error) {

	configs, _, err := a.configLoader(a.configPath)

	if err != nil {
		return
	}

	if newMD5 := md5.Sum(bytes.Join(configs, []byte{})); newMD5 != a.yamlMD5 {
		a.yamlMD5 = newMD5

		for _, configData := range configs {
			mjEndpoint, err := a.mockingjayLoader(configData)

			if err != nil {
				return nil, err
			}

			endpoints = append(endpoints, mjEndpoint...)
		}
	}

	return
}

func (a *application) createFakeServer(endpoints []mockingjay.FakeEndpoint, debugMode bool) (server http.Handler, err error) {
	go a.PollConfig()
	a.mjServer = a.mockingjayServerMaker(endpoints, debugMode)
	monkeyServer, err := a.monkeyServerMaker(a.mjServer, a.monkeyConfigPath)

	if err != nil {
		return nil, err
	}

	router := http.NewServeMux()
	router.Handle("/", monkeyServer)

	return router, nil
}

func globFileLoader(path string) (data [][]byte, paths []string, err error) {
	files, err := filepath.Glob(path)

	if err != nil {
		return nil, nil, fmt.Errorf("Failed to get files from file path (glob) %s, %v", path, err)
	}

	if len(files) == 0 {
		return nil, nil, fmt.Errorf("No files found in path %s", path)
	}

	var configs [][]byte
	for _, file := range files {

		configData, err := ioutil.ReadFile(file)

		if err != nil {
			return nil, nil, err
		}

		configs = append(configs, configData)
	}

	return configs, files, nil
}
<|endoftext|>


<commit_before>package fire

import (
	"fmt"
	"time"

	"github.com/labstack/echo"
	"github.com/labstack/echo/engine"
	"github.com/labstack/echo/engine/standard"
	"github.com/labstack/echo/middleware"
	"gopkg.in/mgo.v2"
)

// An Application provides an out-of-the-box configuration of components to
// get started with building JSON APIs.
type Application struct {
	set    *Set
	router *echo.Echo

	bodyLimit      string
	allowedOrigins []string
	allowedHeaders []string

	forceEncryption        bool
	disableCORS            bool
	disableCompression     bool
	disableRecovery        bool
	disableCommonSecurity  bool
	enableMethodOverriding bool
	enableDevMode          bool
}

// New creates and returns a new Application.
func New(mongoURI, prefix string) *Application {
	// create router
	router := echo.New()

	// connect to database
	sess, err := mgo.Dial(mongoURI)
	if err != nil {
		panic(err)
	}

	// create controller set
	set := NewSet(sess, router, prefix)

	return &Application{
		set:            set,
		router:         router,
		bodyLimit:      "4K",
		allowedOrigins: []string{"*"},
		allowedHeaders: []string{
			echo.HeaderOrigin,
			echo.HeaderContentType,
			echo.HeaderAuthorization,
		},
	}
}

// Mount will add controllers to the set and register them on the router.
//
// Note: Each controller should only be mounted once.
func (a *Application) Mount(controllers ...*Controller) {
	a.set.Mount(controllers...)
}

// Router will return the internally used echo instance.
func (a *Application) Router() *echo.Echo {
	return a.router
}

// ForceEncryption will make the application enforce and only respond to
// encrypted requests.
//
// Note: ForceEncryption will be automatically enabled when calling SecureStart.
func (a *Application) ForceEncryption() {
	a.forceEncryption = true
}

// EnableMethodOverriding will enable the usage of the X-HTTP-Method-Override
// header to set a request method when using the POST method.
//
// Note: This method must be called before calling Run or Start.
func (a *Application) EnableMethodOverriding() {
	a.enableMethodOverriding = true
}

// SetBodyLimit can be used to override the default body limit of 4K with a new
// value in the form of 4K, 2M, 1G or 1P.
//
// Note: This method must be called before calling Run or Start.
func (a *Application) SetBodyLimit(size string) {
	a.bodyLimit = size
}

// SetAllowedOrigins will replace the default allowed origin set `*`.
func (a *Application) SetAllowedOrigins(origins ...string) {
	a.allowedOrigins = origins
}

// AddAllowedHeaders will allow additional headers.
func (a *Application) AddAllowedHeaders(headers ...string) {
	a.allowedHeaders = append(a.allowedHeaders, headers...)
}

// DisableCORS will turn off CORS support.
//
// Note: This method must be called before calling Run or Start.
func (a *Application) DisableCORS(origins ...string) {
	a.disableCORS = true
}

// DisableCompression will turn of gzip compression.
//
// Note: This method must be called before calling Run or Start.
func (a *Application) DisableCompression() {
	a.disableCompression = true
}

// DisableRecovery will disable the automatic recover mechanism.
//
// Note: This method must be called before calling Run or Start.
func (a *Application) DisableRecovery() {
	a.disableRecovery = true
}

// DisableCommonSecurity will disable common security features including:
// protection against cross-site scripting attacks by setting the
// `X-XSS-Protection` header, protection against overriding Content-Type
// header by setting the `X-Content-Type-Options` header and protection against
// clickjacking by setting the `X-Frame-Options` header.
//
// Note: This method must be called before calling Run or Start.
func (a *Application) DisableCommonSecurity() {
	a.disableCommonSecurity = true
}

// EnableDevMode will enable the development mode that prints all registered
// handlers on boot and all incoming requests.
func (a *Application) EnableDevMode() {
	a.enableDevMode = true
}

// Start will run the application on the specified address.
func (a *Application) Start(addr string) {
	a.run(standard.New(addr))
}

// SecureStart will run the application on the specified address using a TLS
// certificate.
func (a *Application) SecureStart(addr, certFile, keyFile string) {
	a.forceEncryption = true

	a.run(standard.WithTLS(addr, certFile, keyFile))
}

func (a *Application) run(server engine.Server) {
	// set body limit
	a.router.Use(middleware.BodyLimit(a.bodyLimit))

	// force encryption
	if a.forceEncryption {
		// TODO: register https redirect middleware with next release
		// a.router.Pre(middleware.HTTPSRedirect())
	}

	// enable cors
	if !a.disableCORS {
		allowedHeaders := a.allowedHeaders

		// add method override header if enabled
		if a.enableMethodOverriding {
			allowedHeaders = append(allowedHeaders, echo.HeaderXHTTPMethodOverride)
		}

		// add cors middleware
		a.router.Use(middleware.CORSWithConfig(middleware.CORSConfig{
			AllowOrigins: a.allowedOrigins,
			AllowMethods: []string{echo.GET, echo.POST, echo.PATCH, echo.DELETE},
			AllowHeaders: allowedHeaders,
			MaxAge:       60,
		}))
	}

	// enable gzip compression
	if !a.disableCompression {
		a.router.Use(middleware.Gzip())
	}

	// enable automatic recovery
	if !a.disableRecovery {
		a.router.Use(middleware.Recover())
	}

	// enable common security
	if !a.disableCommonSecurity {
		config := middleware.DefaultSecureConfig

		// keep using TLS for 60 minutes on just that domain
		// TODO: Make that configurable.
		if a.forceEncryption {
			config.HSTSMaxAge = 3600
			config.HSTSExcludeSubdomains = true
		}

		a.router.Use(middleware.SecureWithConfig(config))
	}

	// enable method overriding
	if a.enableMethodOverriding {
		a.router.Pre(middleware.MethodOverride())
	}

	// enable dev mode
	if a.enableDevMode {
		a.printInfo()
		a.router.Use(a.logger)
	}

	a.router.Run(server)
}

func (a *Application) printInfo() {
	fmt.Println("==> Fire application starting...")
	fmt.Println("==> Registered routes:")

	for _, route := range a.router.Routes() {
		fmt.Printf("%6s  %-30s\n", route.Method, route.Path)
	}

	fmt.Println("==> Ready to go!")
}

func (a *Application) logger(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) (err error) {
		req := c.Request()
		res := c.Response()

		start := time.Now()
		if err = next(c); err != nil {
			c.Error(err)
		}

		duration := time.Since(start).String()

		path := req.URL().Path()
		if path == "" {
			path = "/"
		}

		fmt.Printf("%6s  %-30s  %d  %s\n", req.Method(), path, res.Status(), duration)

		return
	}
}
<commit_msg>added note<commit_after>package fire

import (
	"fmt"
	"time"

	"github.com/labstack/echo"
	"github.com/labstack/echo/engine"
	"github.com/labstack/echo/engine/standard"
	"github.com/labstack/echo/middleware"
	"gopkg.in/mgo.v2"
)

// An Application provides an out-of-the-box configuration of components to
// get started with building JSON APIs.
type Application struct {
	set    *Set
	router *echo.Echo

	bodyLimit      string
	allowedOrigins []string
	allowedHeaders []string

	forceEncryption        bool
	disableCORS            bool
	disableCompression     bool
	disableRecovery        bool
	disableCommonSecurity  bool
	enableMethodOverriding bool
	enableDevMode          bool
}

// New creates and returns a new Application.
func New(mongoURI, prefix string) *Application {
	// create router
	router := echo.New()

	// connect to database
	sess, err := mgo.Dial(mongoURI)
	if err != nil {
		panic(err)
	}

	// create controller set
	set := NewSet(sess, router, prefix)

	return &Application{
		set:            set,
		router:         router,
		bodyLimit:      "4K",
		allowedOrigins: []string{"*"},
		allowedHeaders: []string{
			echo.HeaderOrigin,
			echo.HeaderContentType,
			echo.HeaderAuthorization,
		},
	}
}

// Mount will add controllers to the set and register them on the router.
//
// Note: Each controller should only be mounted once.
func (a *Application) Mount(controllers ...*Controller) {
	a.set.Mount(controllers...)
}

// Router will return the internally used echo instance.
func (a *Application) Router() *echo.Echo {
	return a.router
}

// ForceEncryption will make the application enforce and only respond to
// encrypted requests.
//
// Note: ForceEncryption will be automatically enabled when calling SecureStart.
func (a *Application) ForceEncryption() {
	a.forceEncryption = true
}

// EnableMethodOverriding will enable the usage of the X-HTTP-Method-Override
// header to set a request method when using the POST method.
//
// Note: This method must be called before calling Run or Start.
func (a *Application) EnableMethodOverriding() {
	a.enableMethodOverriding = true
}

// SetBodyLimit can be used to override the default body limit of 4K with a new
// value in the form of 4K, 2M, 1G or 1P.
//
// Note: This method must be called before calling Run or Start.
func (a *Application) SetBodyLimit(size string) {
	a.bodyLimit = size
}

// SetAllowedOrigins will replace the default allowed origin set `*`.
func (a *Application) SetAllowedOrigins(origins ...string) {
	a.allowedOrigins = origins
}

// AddAllowedHeaders will allow additional headers.
func (a *Application) AddAllowedHeaders(headers ...string) {
	a.allowedHeaders = append(a.allowedHeaders, headers...)
}

// DisableCORS will turn off CORS support.
//
// Note: This method must be called before calling Run or Start.
func (a *Application) DisableCORS(origins ...string) {
	a.disableCORS = true
}

// DisableCompression will turn of gzip compression.
//
// Note: This method must be called before calling Run or Start.
func (a *Application) DisableCompression() {
	a.disableCompression = true
}

// DisableRecovery will disable the automatic recover mechanism.
//
// Note: This method must be called before calling Run or Start.
func (a *Application) DisableRecovery() {
	a.disableRecovery = true
}

// DisableCommonSecurity will disable common security features including:
// protection against cross-site scripting attacks by setting the
// `X-XSS-Protection` header, protection against overriding Content-Type
// header by setting the `X-Content-Type-Options` header and protection against
// clickjacking by setting the `X-Frame-Options` header.
//
// Note: This method must be called before calling Run or Start.
func (a *Application) DisableCommonSecurity() {
	a.disableCommonSecurity = true
}

// EnableDevMode will enable the development mode that prints all registered
// handlers on boot and all incoming requests.
func (a *Application) EnableDevMode() {
	a.enableDevMode = true
}

// Start will run the application on the specified address.
func (a *Application) Start(addr string) {
	a.run(standard.New(addr))
}

// SecureStart will run the application on the specified address using a TLS
// certificate.
func (a *Application) SecureStart(addr, certFile, keyFile string) {
	a.forceEncryption = true

	a.run(standard.WithTLS(addr, certFile, keyFile))
}

func (a *Application) run(server engine.Server) {
	// set body limit
	a.router.Use(middleware.BodyLimit(a.bodyLimit))

	// force encryption
	if a.forceEncryption {
		// TODO: register https redirect middleware with next release
		// a.router.Pre(middleware.HTTPSRedirect())
	}

	// enable cors
	if !a.disableCORS {
		allowedHeaders := a.allowedHeaders

		// add method override header if enabled
		if a.enableMethodOverriding {
			allowedHeaders = append(allowedHeaders, echo.HeaderXHTTPMethodOverride)
		}

		// add cors middleware
		a.router.Use(middleware.CORSWithConfig(middleware.CORSConfig{
			AllowOrigins: a.allowedOrigins,
			AllowMethods: []string{echo.GET, echo.POST, echo.PATCH, echo.DELETE},
			AllowHeaders: allowedHeaders,
			MaxAge:       60,
		}))
	}

	// enable gzip compression
	if !a.disableCompression {
		a.router.Use(middleware.Gzip())
	}

	// enable automatic recovery
	if !a.disableRecovery {
		a.router.Use(middleware.Recover())
	}

	// enable common security
	if !a.disableCommonSecurity {
		config := middleware.DefaultSecureConfig

		// keep using TLS for 60 minutes on just that domain
		// TODO: Make that configurable.
		if a.forceEncryption {
			config.HSTSMaxAge = 3600
			config.HSTSExcludeSubdomains = true
		}

		a.router.Use(middleware.SecureWithConfig(config))
	}

	// enable method overriding
	if a.enableMethodOverriding {
		a.router.Pre(middleware.MethodOverride())
	}

	// enable dev mode
	if a.enableDevMode {
		a.printInfo()
		a.router.Use(a.logger)
	}

	a.router.Run(server)
}

func (a *Application) printInfo() {
	fmt.Println("==> Fire application starting...")
	fmt.Println("==> Registered routes:")

	// TODO: Order routes.

	for _, route := range a.router.Routes() {
		fmt.Printf("%6s  %-30s\n", route.Method, route.Path)
	}

	fmt.Println("==> Ready to go!")
}

func (a *Application) logger(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) (err error) {
		req := c.Request()
		res := c.Response()

		start := time.Now()
		if err = next(c); err != nil {
			c.Error(err)
		}

		duration := time.Since(start).String()

		path := req.URL().Path()
		if path == "" {
			path = "/"
		}

		fmt.Printf("%6s  %-30s  %d  %s\n", req.Method(), path, res.Status(), duration)

		return
	}
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"github.com/pterodactyl/wings/config"
	"github.com/pterodactyl/wings/environment"
	"io"
	"io/ioutil"
	"os"
	"sync"
)

var stateMutex sync.Mutex

// Returns the state of the servers.
func CachedServerStates() (map[string]string, error) {
	// Request a lock after we check if the file exists.
	stateMutex.Lock()
	defer stateMutex.Unlock()

	// Open the states file.
	f, err := os.OpenFile(config.Get().System.GetStatesPath(), os.O_RDONLY|os.O_CREATE, 0644)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	// Convert the json object to a map.
	states := map[string]string{}
	if err := json.NewDecoder(f).Decode(&states); err != nil && err != io.EOF {
		return nil, err
	}

	return states, nil
}

// saveServerStates .
func saveServerStates() error {
	// Get the states of all servers on the daemon.
	states := map[string]string{}
	for _, s := range GetServers().All() {
		states[s.Id()] = s.Environment.State()
	}

	// Convert the map to a json object.
	data, err := json.Marshal(states)
	if err != nil {
		return err
	}

	stateMutex.Lock()
	defer stateMutex.Unlock()

	// Write the data to the file
	if err := ioutil.WriteFile(config.Get().System.GetStatesPath(), data, 0644); err != nil {
		return err
	}

	return nil
}

// Sets the state of the server internally. This function handles crash detection as
// well as reporting to event listeners for the server.
func (s *Server) OnStateChange() {
	prevState := s.Proc().State.Load()

	st := s.Environment.State()
	// Update the currently tracked state for the server.
	s.Proc().State.Store(st)

	// Emit the event to any listeners that are currently registered.
	if prevState != s.Environment.State() {
		s.Log().WithField("status", st).Debug("saw server status change event")
		s.Events().Publish(StatusEvent, st)
	}

	// Persist this change to the disk immediately so that should the Daemon be stopped or
	// crash we can immediately restore the server state.
	//
	// This really only makes a difference if all of the Docker containers are also stopped,
	// but this was a highly requested feature and isn't hard to work with, so lets do it.
	//
	// We also get the benefit of server status changes always propagating corrected configurations
	// to the disk should we forget to do it elsewhere.
	go func() {
		if err := saveServerStates(); err != nil {
			s.Log().WithField("error", err).Warn("failed to write server states to disk")
		}
	}()

	// Reset the resource usage to 0 when the process fully stops so that all of the UI
	// views in the Panel correctly display 0.
	if st == environment.ProcessOfflineState {
		s.resources.mu.Lock()
		s.resources.Empty()
		s.resources.mu.Unlock()

		s.emitProcUsage()
	}

	// If server was in an online state, and is now in an offline state we should handle
	// that as a crash event. In that scenario, check the last crash time, and the crash
	// counter.
	//
	// In the event that we have passed the thresholds, don't do anything, otherwise
	// automatically attempt to start the process back up for the user. This is done in a
	// separate thread as to not block any actions currently taking place in the flow
	// that called this function.
	if (prevState == environment.ProcessStartingState || prevState == environment.ProcessRunningState) && s.Environment.State() == environment.ProcessOfflineState {
		s.Log().Info("detected server as entering a crashed state; running crash handler")

		go func(server *Server) {
			if err := server.handleServerCrash(); err != nil {
				if IsTooFrequentCrashError(err) {
					server.Log().Info("did not restart server after crash; occurred too soon after the last")
				} else {
					server.Log().WithField("error", err).Error("failed to handle server crash")
				}
			}
		}(s)
	}
}

// Returns the current state of the server in a race-safe manner.
// Deprecated
// use Environment.State()
func (s *Server) GetState() string {
	return s.Environment.State()
}

// Determines if the server state is running or not. This is different than the
// environment state, it is simply the tracked state from this daemon instance, and
// not the response from Docker.
func (s *Server) IsRunning() bool {
	st := s.Environment.State()

	return st == environment.ProcessRunningState || st == environment.ProcessStartingState
}
<commit_msg>Publish message when error occurs during crash handling<commit_after>package server

import (
	"encoding/json"
	"github.com/pterodactyl/wings/config"
	"github.com/pterodactyl/wings/environment"
	"io"
	"io/ioutil"
	"os"
	"sync"
)

var stateMutex sync.Mutex

// Returns the state of the servers.
func CachedServerStates() (map[string]string, error) {
	// Request a lock after we check if the file exists.
	stateMutex.Lock()
	defer stateMutex.Unlock()

	// Open the states file.
	f, err := os.OpenFile(config.Get().System.GetStatesPath(), os.O_RDONLY|os.O_CREATE, 0644)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	// Convert the json object to a map.
	states := map[string]string{}
	if err := json.NewDecoder(f).Decode(&states); err != nil && err != io.EOF {
		return nil, err
	}

	return states, nil
}

// saveServerStates .
func saveServerStates() error {
	// Get the states of all servers on the daemon.
	states := map[string]string{}
	for _, s := range GetServers().All() {
		states[s.Id()] = s.Environment.State()
	}

	// Convert the map to a json object.
	data, err := json.Marshal(states)
	if err != nil {
		return err
	}

	stateMutex.Lock()
	defer stateMutex.Unlock()

	// Write the data to the file
	if err := ioutil.WriteFile(config.Get().System.GetStatesPath(), data, 0644); err != nil {
		return err
	}

	return nil
}

// Sets the state of the server internally. This function handles crash detection as
// well as reporting to event listeners for the server.
func (s *Server) OnStateChange() {
	prevState := s.Proc().State.Load()

	st := s.Environment.State()
	// Update the currently tracked state for the server.
	s.Proc().State.Store(st)

	// Emit the event to any listeners that are currently registered.
	if prevState != s.Environment.State() {
		s.Log().WithField("status", st).Debug("saw server status change event")
		s.Events().Publish(StatusEvent, st)
	}

	// Persist this change to the disk immediately so that should the Daemon be stopped or
	// crash we can immediately restore the server state.
	//
	// This really only makes a difference if all of the Docker containers are also stopped,
	// but this was a highly requested feature and isn't hard to work with, so lets do it.
	//
	// We also get the benefit of server status changes always propagating corrected configurations
	// to the disk should we forget to do it elsewhere.
	go func() {
		if err := saveServerStates(); err != nil {
			s.Log().WithField("error", err).Warn("failed to write server states to disk")
		}
	}()

	// Reset the resource usage to 0 when the process fully stops so that all of the UI
	// views in the Panel correctly display 0.
	if st == environment.ProcessOfflineState {
		s.resources.mu.Lock()
		s.resources.Empty()
		s.resources.mu.Unlock()

		s.emitProcUsage()
	}

	// If server was in an online state, and is now in an offline state we should handle
	// that as a crash event. In that scenario, check the last crash time, and the crash
	// counter.
	//
	// In the event that we have passed the thresholds, don't do anything, otherwise
	// automatically attempt to start the process back up for the user. This is done in a
	// separate thread as to not block any actions currently taking place in the flow
	// that called this function.
	if (prevState == environment.ProcessStartingState || prevState == environment.ProcessRunningState) && s.Environment.State() == environment.ProcessOfflineState {
		s.Log().Info("detected server as entering a crashed state; running crash handler")

		go func(server *Server) {
			if err := server.handleServerCrash(); err != nil {
				if IsTooFrequentCrashError(err) {
					server.Log().Info("did not restart server after crash; occurred too soon after the last")
				} else {
					s.PublishConsoleOutputFromDaemon("Server crash was detected but an error occurred while handling it.")
					server.Log().WithField("error", err).Error("failed to handle server crash")
				}
			}
		}(s)
	}
}

// Returns the current state of the server in a race-safe manner.
// Deprecated
// use Environment.State()
func (s *Server) GetState() string {
	return s.Environment.State()
}

// Determines if the server state is running or not. This is different than the
// environment state, it is simply the tracked state from this daemon instance, and
// not the response from Docker.
func (s *Server) IsRunning() bool {
	st := s.Environment.State()

	return st == environment.ProcessRunningState || st == environment.ProcessStartingState
}
<|endoftext|>


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

package minio

import "io"

// PutObjectWithProgress - With progress.
func (c Client) PutObjectWithProgress(bucketName, objectName string, reader io.Reader, contentType string, progress io.Reader) (n int64, err error) {
	// Input validation.
	if err := isValidBucketName(bucketName); err != nil {
		return 0, err
	}
	if err := isValidObjectName(objectName); err != nil {
		return 0, err
	}
	if reader == nil {
		return 0, ErrInvalidArgument("Input reader is invalid, cannot be nil.")
	}

	// Size of the object.
	var size int64

	// Get reader size.
	size, err = getReaderSize(reader)
	if err != nil {
		return 0, err
	}

	// Check for largest object size allowed.
	if size > int64(maxMultipartPutObjectSize) {
		return 0, ErrEntityTooLarge(size, maxMultipartPutObjectSize, bucketName, objectName)
	}

	// NOTE: Google Cloud Storage does not implement Amazon S3 Compatible multipart PUT.
	// So we fall back to single PUT operation with the maximum limit of 5GiB.
	if isGoogleEndpoint(c.endpointURL) {
		if size <= -1 {
			return 0, ErrorResponse{
				Code:       "NotImplemented",
				Message:    "Content-Length cannot be negative for file uploads to Google Cloud Storage.",
				Key:        objectName,
				BucketName: bucketName,
			}
		}
		if size > maxSinglePutObjectSize {
			return 0, ErrEntityTooLarge(size, maxSinglePutObjectSize, bucketName, objectName)
		}
		// Do not compute MD5 for Google Cloud Storage. Uploads up to 5GiB in size.
		return c.putObjectNoChecksum(bucketName, objectName, reader, size, contentType, progress)
	}

	// NOTE: S3 doesn't allow anonymous multipart requests.
	if isAmazonEndpoint(c.endpointURL) && c.anonymous {
		if size <= -1 {
			return 0, ErrorResponse{
				Code:       "NotImplemented",
				Message:    "Content-Length cannot be negative for anonymous requests.",
				Key:        objectName,
				BucketName: bucketName,
			}
		}
		if size > maxSinglePutObjectSize {
			return 0, ErrEntityTooLarge(size, maxSinglePutObjectSize, bucketName, objectName)
		}
		// Do not compute MD5 for anonymous requests to Amazon
		// S3. Uploads up to 5GiB in size.
		return c.putObjectNoChecksum(bucketName, objectName, reader, size, contentType, progress)
	}

	// putSmall object.
	if size < minPartSize && size >= 0 {
		return c.putObjectSingle(bucketName, objectName, reader, size, contentType, progress)
	}
	// For all sizes greater than 5MiB do multipart.
	n, err = c.putObjectMultipart(bucketName, objectName, reader, size, contentType, progress)
	if err != nil {
		errResp := ToErrorResponse(err)
		// Verify if multipart functionality is not available, if not
		// fall back to single PutObject operation.
		if errResp.Code == "AccessDenied" && errResp.Message == "Access Denied." {
			// Verify if size of reader is greater than '5GiB'.
			if size > maxSinglePutObjectSize {
				return 0, ErrEntityTooLarge(size, maxSinglePutObjectSize, bucketName, objectName)
			}
			// Fall back to uploading as single PutObject operation.
			return c.putObjectSingle(bucketName, objectName, reader, size, contentType, progress)
		}
		return n, err
	}
	return n, nil
}
<commit_msg>Fixing a bug in multipart upload flow. (#521)<commit_after>/*
 * Minio Go Library for Amazon S3 Compatible Cloud Storage (C) 2015 Minio, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package minio

import (
	"io"
	"strings"
)

// PutObjectWithProgress - With progress.
func (c Client) PutObjectWithProgress(bucketName, objectName string, reader io.Reader, contentType string, progress io.Reader) (n int64, err error) {
	// Input validation.
	if err := isValidBucketName(bucketName); err != nil {
		return 0, err
	}
	if err := isValidObjectName(objectName); err != nil {
		return 0, err
	}
	if reader == nil {
		return 0, ErrInvalidArgument("Input reader is invalid, cannot be nil.")
	}

	// Size of the object.
	var size int64

	// Get reader size.
	size, err = getReaderSize(reader)
	if err != nil {
		return 0, err
	}

	// Check for largest object size allowed.
	if size > int64(maxMultipartPutObjectSize) {
		return 0, ErrEntityTooLarge(size, maxMultipartPutObjectSize, bucketName, objectName)
	}

	// NOTE: Google Cloud Storage does not implement Amazon S3 Compatible multipart PUT.
	// So we fall back to single PUT operation with the maximum limit of 5GiB.
	if isGoogleEndpoint(c.endpointURL) {
		if size <= -1 {
			return 0, ErrorResponse{
				Code:       "NotImplemented",
				Message:    "Content-Length cannot be negative for file uploads to Google Cloud Storage.",
				Key:        objectName,
				BucketName: bucketName,
			}
		}
		if size > maxSinglePutObjectSize {
			return 0, ErrEntityTooLarge(size, maxSinglePutObjectSize, bucketName, objectName)
		}
		// Do not compute MD5 for Google Cloud Storage. Uploads up to 5GiB in size.
		return c.putObjectNoChecksum(bucketName, objectName, reader, size, contentType, progress)
	}

	// NOTE: S3 doesn't allow anonymous multipart requests.
	if isAmazonEndpoint(c.endpointURL) && c.anonymous {
		if size <= -1 {
			return 0, ErrorResponse{
				Code:       "NotImplemented",
				Message:    "Content-Length cannot be negative for anonymous requests.",
				Key:        objectName,
				BucketName: bucketName,
			}
		}
		if size > maxSinglePutObjectSize {
			return 0, ErrEntityTooLarge(size, maxSinglePutObjectSize, bucketName, objectName)
		}
		// Do not compute MD5 for anonymous requests to Amazon
		// S3. Uploads up to 5GiB in size.
		return c.putObjectNoChecksum(bucketName, objectName, reader, size, contentType, progress)
	}

	// putSmall object.
	if size < minPartSize && size >= 0 {
		return c.putObjectSingle(bucketName, objectName, reader, size, contentType, progress)
	}
	// For all sizes greater than 5MiB do multipart.
	n, err = c.putObjectMultipart(bucketName, objectName, reader, size, contentType, progress)
	if err != nil {
		errResp := ToErrorResponse(err)
		// Verify if multipart functionality is not available, if not
		// fall back to single PutObject operation.
		if errResp.Code == "AccessDenied" && strings.Contains(errResp.Message, "Access Denied") {
			// Verify if size of reader is greater than '5GiB'.
			if size > maxSinglePutObjectSize {
				return 0, ErrEntityTooLarge(size, maxSinglePutObjectSize, bucketName, objectName)
			}
			// Fall back to uploading as single PutObject operation.
			return c.putObjectSingle(bucketName, objectName, reader, size, contentType, progress)
		}
		return n, err
	}
	return n, nil
}
<|endoftext|>


<commit_before><commit_msg>wafcontext: govet<commit_after><|endoftext|>


<commit_before>// Package appsettings provides simple key/value store functionality
package appsettings

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"os"
	"strconv"
	"sync"
)

// DataTree is the host of key and branches of values
type DataTree interface {
	GetString(key string) (string, error)
	SetString(key string, val string)
	GetInt(key string) (int, error)
	SetInt(key string, val int)
	GetInt64(key string) (int64, error)
	SetInt64(key string, val int64)
	Delete(key string)
	GetTree(key string) DataTree

	GetLeaves() map[string]string
}

type tree struct {
	Branches map[string]*tree
	Leaves   map[string]string

	sync.Mutex
}

// AppSettings is the root most DataTree
type AppSettings struct {
	filename string

	*tree
}

// NewAppSettings gets a new AppSettings struct
func NewAppSettings(dbFilename string) (*AppSettings, error) {
	a := &AppSettings{
		filename: dbFilename,
		tree: &tree{
			Branches: make(map[string]*tree),
			Leaves:   make(map[string]string),
		},
	}

	if _, err := os.Stat(dbFilename); os.IsNotExist(err) {
		d1, _ := json.Marshal(a)
		err := ioutil.WriteFile(dbFilename, d1, 0644)
		if err != nil {
			return nil, err
		}
	} else {
		d1, err := ioutil.ReadFile(dbFilename)
		if err != nil {
			return nil, err
		}

		err = json.Unmarshal(d1, a)
		if err != nil {
			return nil, err
		}
	}

	return a, nil
}

// ErrUndefinedKey is returned when the key requested from get is undefined.
var ErrUndefinedKey = errors.New("undefined key")

func (a *tree) GetString(key string) (string, error) {
	if _, ok := a.Leaves[key]; !ok {
		return "", ErrUndefinedKey
	}

	return a.Leaves[key], nil
}

func (a *tree) SetString(key string, val string) {
	a.Leaves[key] = val
}

func (a *tree) GetInt(key string) (int, error) {
	str, err := a.GetString(key)
	if err != nil {
		return 0, err
	}

	i, err := strconv.Atoi(str)
	if err != nil {
		return 0, err
	}
	return i, nil
}

func (a *tree) SetInt(key string, val int) {
	a.Leaves[key] = strconv.Itoa(val)
}

func (a *tree) GetInt64(key string) (int64, error) {
	str, err := a.GetString(key)
	if err != nil {
		return 0, err
	}

	i, err := strconv.ParseInt(str, 10, 64)
	if err != nil {
		return 0, err
	}
	return i, nil
}

func (a *tree) SetInt64(key string, val int64) {
	// y := *a
	a.Leaves[key] = strconv.FormatInt(val, 10)
}

func (a *tree) Delete(key string) {
	delete(a.Leaves, key)
}

func (a *tree) GetLeaves() map[string]string {
	return a.Leaves
}

// GetTree fetches a tree for app setting storage
func (a *tree) GetTree(key string) DataTree {
	// a.Lock()
	// defer a.Unlock()

	if _, ok := a.Branches[key]; !ok {
		a.Branches[key] = &tree{
			Branches: make(map[string]*tree),
			Leaves:   make(map[string]string),
		}
	}

	return a.Branches[key]
}

// Persist causes the current state of the app settings to be persisted.
func (a *AppSettings) Persist() error {
	a.Lock()
	defer a.Unlock()

	d1, _ := json.Marshal(a)
	err := ioutil.WriteFile(a.filename, d1, 0644)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Adds a handful of mutex handling<commit_after>// Package appsettings provides simple key/value store functionality
package appsettings

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"os"
	"strconv"
	"sync"
)

// DataTree is the host of key and branches of values
type DataTree interface {
	GetString(key string) (string, error)
	SetString(key string, val string)
	GetInt(key string) (int, error)
	SetInt(key string, val int)
	GetInt64(key string) (int64, error)
	SetInt64(key string, val int64)
	Delete(key string)
	GetTree(key string) DataTree

	GetLeaves() map[string]string
}

type tree struct {
	Branches map[string]*tree
	Leaves   map[string]string

	sync.Mutex
}

// AppSettings is the root most DataTree
type AppSettings struct {
	filename string

	*tree
}

// NewAppSettings gets a new AppSettings struct
func NewAppSettings(dbFilename string) (*AppSettings, error) {
	a := &AppSettings{
		filename: dbFilename,
		tree: &tree{
			Branches: make(map[string]*tree),
			Leaves:   make(map[string]string),
		},
	}

	if _, err := os.Stat(dbFilename); os.IsNotExist(err) {
		d1, _ := json.Marshal(a)
		err := ioutil.WriteFile(dbFilename, d1, 0644)
		if err != nil {
			return nil, err
		}
	} else {
		d1, err := ioutil.ReadFile(dbFilename)
		if err != nil {
			return nil, err
		}

		err = json.Unmarshal(d1, a)
		if err != nil {
			return nil, err
		}
	}

	return a, nil
}

// ErrUndefinedKey is returned when the key requested from get is undefined.
var ErrUndefinedKey = errors.New("undefined key")

func (a *tree) GetString(key string) (string, error) {
	a.Lock()
	defer a.Unlock()

	if _, ok := a.Leaves[key]; !ok {
		return "", ErrUndefinedKey
	}

	return a.Leaves[key], nil
}

func (a *tree) SetString(key string, val string) {
	a.Lock()
	defer a.Unlock()

	a.Leaves[key] = val
}

func (a *tree) GetInt(key string) (int, error) {
	a.Lock()
	defer a.Unlock()

	str, err := a.GetString(key)
	if err != nil {
		return 0, err
	}

	i, err := strconv.Atoi(str)
	if err != nil {
		return 0, err
	}
	return i, nil
}

func (a *tree) SetInt(key string, val int) {
	a.Lock()
	defer a.Unlock()

	a.Leaves[key] = strconv.Itoa(val)
}

func (a *tree) GetInt64(key string) (int64, error) {
	a.Lock()
	defer a.Unlock()

	str, err := a.GetString(key)
	if err != nil {
		return 0, err
	}

	i, err := strconv.ParseInt(str, 10, 64)
	if err != nil {
		return 0, err
	}
	return i, nil
}

func (a *tree) SetInt64(key string, val int64) {
	a.Lock()
	defer a.Unlock()

	a.Leaves[key] = strconv.FormatInt(val, 10)
}

func (a *tree) Delete(key string) {
	a.Lock()
	defer a.Unlock()

	delete(a.Leaves, key)
}

func (a *tree) GetLeaves() map[string]string {
	return a.Leaves
}

// GetTree fetches a tree for app setting storage
func (a *tree) GetTree(key string) DataTree {
	a.Lock()
	defer a.Unlock()

	if _, ok := a.Branches[key]; !ok {
		a.Branches[key] = &tree{
			Branches: make(map[string]*tree),
			Leaves:   make(map[string]string),
		}
	}

	return a.Branches[key]
}

// Persist causes the current state of the app settings to be persisted.
func (a *AppSettings) Persist() error {
	a.Lock()
	defer a.Unlock()

	d1, _ := json.Marshal(a)
	err := ioutil.WriteFile(a.filename, d1, 0644)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package archive

import (
	"bufio"
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// VCS is a struct that explains how to get the file list for a given
// VCS.
type VCS struct {
	Name string

	// Detect is a list of files/folders that if they exist, signal that
	// this VCS is the VCS in use.
	Detect []string

	// Files returns the files that are under version control for the
	// given path.
	Files VCSFilesFunc

	// Metadata returns arbitrary metadata about the underlying VCS for the
	// given path.
	Metadata VCSMetadataFunc
}

// VCSList is the list of VCS we recognize.
var VCSList = []*VCS{
	&VCS{
		Name:   "git",
		Detect: []string{".git/"},
		Files:  vcsFilesCmd("git", "ls-files"),
	},
	&VCS{
		Name:   "hg",
		Detect: []string{".hg/"},
		Files:  vcsTrimCmd(vcsFilesCmd("hg", "locate", "-f", "--include", ".")),
	},
	&VCS{
		Name:   "svn",
		Detect: []string{".svn/"},
		Files:  vcsFilesCmd("svn", "ls"),
	},
}

// VCSFilesFunc is the callback invoked to return the files in the VCS.
//
// The return value should be paths relative to the given path.
type VCSFilesFunc func(string) ([]string, error)

// VCSMetadataFunc is the callback invoked to get arbitrary information about
// the current VCS.
//
// The return value should be a map of key-value pairs.
type VCSMetadataFunc func(string) (map[string]string, error)

// vcsDetect detects the VCS that is used for path.
func vcsDetect(path string) (*VCS, error) {
	for _, v := range VCSList {
		dir := path
		for {
			for _, f := range v.Detect {
				check := filepath.Join(dir, f)
				if _, err := os.Stat(check); err == nil {
					return v, nil
				}
			}

			lastDir := dir
			dir = filepath.Dir(dir)
			if dir == lastDir {
				break
			}
		}
	}

	return nil, nil
}

// vcsFiles returns the files for the VCS directory path.
func vcsFiles(path string) ([]string, error) {
	vcs, err := vcsDetect(path)
	if err != nil {
		return nil, fmt.Errorf("Error detecting VCS: %s", err)
	}
	if vcs == nil {
		return nil, fmt.Errorf("No VCS found for path: %s", path)
	}

	if vcs.Files {
		return vcs.Files(path)
	}

	return nil, nil
}

// vcsFilesCmd creates a Files-compatible function that reads the files
// by executing the command in the repository path and returning each
// line in stdout.
func vcsFilesCmd(args ...string) VCSFilesFunc {
	return func(path string) ([]string, error) {
		var stderr, stdout bytes.Buffer

		cmd := exec.Command(args[0], args[1:]...)
		cmd.Dir = path
		cmd.Stdout = &stdout
		cmd.Stderr = &stderr
		if err := cmd.Run(); err != nil {
			return nil, fmt.Errorf(
				"Error executing %s: %s",
				strings.Join(args, " "),
				err)
		}

		// Read each line of output as a path
		result := make([]string, 0, 100)
		scanner := bufio.NewScanner(&stdout)
		for scanner.Scan() {
			result = append(result, scanner.Text())
		}

		return result, nil
	}
}

// vcsTrimCmd trims the prefix from the paths returned by another VCSFilesFunc.
// This should be used to wrap another function if the return value is known
// to have full paths rather than relative paths
func vcsTrimCmd(f VCSFilesFunc) VCSFilesFunc {
	return func(path string) ([]string, error) {
		absPath, err := filepath.Abs(path)
		if err != nil {
			return nil, fmt.Errorf(
				"Error expanding VCS path: %s", err)
		}

		// Now that we have the root path, get the inner files
		fs, err := f(path)
		if err != nil {
			return nil, err
		}

		// Trim the root path from the files
		result := make([]string, 0, len(fs))
		for _, f := range fs {
			if !strings.HasPrefix(f, absPath) {
				continue
			}

			f, err = filepath.Rel(absPath, f)
			if err != nil {
				return nil, fmt.Errorf(
					"Error determining path: %s", err)
			}

			result = append(result, f)
		}

		return result, nil
	}
}

// vcsMetadata returns the metadata for the VCS directory path.
func vcsMetadata(path string) (map[string]string, error) {
	vcs, err := vcsDetect(path)
	if err != nil {
		return nil, fmt.Errorf("Error detecting VCS: %s", err)
	}
	if vcs == nil {
		return nil, fmt.Errorf("No VCS found for path: %s", path)
	}

	if vcs.Metadata {
		return vcs.Metadata(path)
	}

	return nil, nil
}
<commit_msg>Act like a Go developer<commit_after>package archive

import (
	"bufio"
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// VCS is a struct that explains how to get the file list for a given
// VCS.
type VCS struct {
	Name string

	// Detect is a list of files/folders that if they exist, signal that
	// this VCS is the VCS in use.
	Detect []string

	// Files returns the files that are under version control for the
	// given path.
	Files VCSFilesFunc

	// Metadata returns arbitrary metadata about the underlying VCS for the
	// given path.
	Metadata VCSMetadataFunc
}

// VCSList is the list of VCS we recognize.
var VCSList = []*VCS{
	&VCS{
		Name:   "git",
		Detect: []string{".git/"},
		Files:  vcsFilesCmd("git", "ls-files"),
	},
	&VCS{
		Name:   "hg",
		Detect: []string{".hg/"},
		Files:  vcsTrimCmd(vcsFilesCmd("hg", "locate", "-f", "--include", ".")),
	},
	&VCS{
		Name:   "svn",
		Detect: []string{".svn/"},
		Files:  vcsFilesCmd("svn", "ls"),
	},
}

// VCSFilesFunc is the callback invoked to return the files in the VCS.
//
// The return value should be paths relative to the given path.
type VCSFilesFunc func(string) ([]string, error)

// VCSMetadataFunc is the callback invoked to get arbitrary information about
// the current VCS.
//
// The return value should be a map of key-value pairs.
type VCSMetadataFunc func(string) (map[string]string, error)

// vcsDetect detects the VCS that is used for path.
func vcsDetect(path string) (*VCS, error) {
	for _, v := range VCSList {
		dir := path
		for {
			for _, f := range v.Detect {
				check := filepath.Join(dir, f)
				if _, err := os.Stat(check); err == nil {
					return v, nil
				}
			}

			lastDir := dir
			dir = filepath.Dir(dir)
			if dir == lastDir {
				break
			}
		}
	}

	return nil, nil
}

// vcsFiles returns the files for the VCS directory path.
func vcsFiles(path string) ([]string, error) {
	vcs, err := vcsDetect(path)
	if err != nil {
		return nil, fmt.Errorf("Error detecting VCS: %s", err)
	}
	if vcs == nil {
		return nil, fmt.Errorf("No VCS found for path: %s", path)
	}

	if vcs.Files != nil {
		return vcs.Files(path)
	}

	return nil, nil
}

// vcsFilesCmd creates a Files-compatible function that reads the files
// by executing the command in the repository path and returning each
// line in stdout.
func vcsFilesCmd(args ...string) VCSFilesFunc {
	return func(path string) ([]string, error) {
		var stderr, stdout bytes.Buffer

		cmd := exec.Command(args[0], args[1:]...)
		cmd.Dir = path
		cmd.Stdout = &stdout
		cmd.Stderr = &stderr
		if err := cmd.Run(); err != nil {
			return nil, fmt.Errorf(
				"Error executing %s: %s",
				strings.Join(args, " "),
				err)
		}

		// Read each line of output as a path
		result := make([]string, 0, 100)
		scanner := bufio.NewScanner(&stdout)
		for scanner.Scan() {
			result = append(result, scanner.Text())
		}

		return result, nil
	}
}

// vcsTrimCmd trims the prefix from the paths returned by another VCSFilesFunc.
// This should be used to wrap another function if the return value is known
// to have full paths rather than relative paths
func vcsTrimCmd(f VCSFilesFunc) VCSFilesFunc {
	return func(path string) ([]string, error) {
		absPath, err := filepath.Abs(path)
		if err != nil {
			return nil, fmt.Errorf(
				"Error expanding VCS path: %s", err)
		}

		// Now that we have the root path, get the inner files
		fs, err := f(path)
		if err != nil {
			return nil, err
		}

		// Trim the root path from the files
		result := make([]string, 0, len(fs))
		for _, f := range fs {
			if !strings.HasPrefix(f, absPath) {
				continue
			}

			f, err = filepath.Rel(absPath, f)
			if err != nil {
				return nil, fmt.Errorf(
					"Error determining path: %s", err)
			}

			result = append(result, f)
		}

		return result, nil
	}
}

// vcsMetadata returns the metadata for the VCS directory path.
func vcsMetadata(path string) (map[string]string, error) {
	vcs, err := vcsDetect(path)
	if err != nil {
		return nil, fmt.Errorf("Error detecting VCS: %s", err)
	}
	if vcs == nil {
		return nil, fmt.Errorf("No VCS found for path: %s", path)
	}

	if vcs.Metadata != nil {
		return vcs.Metadata(path)
	}

	return nil, nil
}
<|endoftext|>