<commit_before>package letsencrypt

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"os"

	"github.com/mholt/caddy/middleware"
	"github.com/mholt/caddy/middleware/redirect"
	"github.com/mholt/caddy/server"
	"github.com/xenolf/lego/acme"
)

// Activate sets up TLS for each server config in configs
// as needed. It only skips the config if the cert and key
// are already provided or if plaintext http is explicitly
// specified as the port.
//
// This function may prompt the user to provide an email
// address if none is available through other means. It
// prefers the email address specified in the config, but
// if that is not available it will check the command line
// argument. If absent, it will use the most recent email
// address from last time. If there isn't one, the user
// will be prompted. If the user leaves email blank, <TODO>.
func Activate(configs []server.Config) ([]server.Config, error) {
	// Group configs by LE email address; this will help us
	// reduce round-trips when getting the certs.
	initMap, err := groupConfigsByEmail(configs)
	if err != nil {
		return configs, err
	}

	// Loop through each email address and obtain certs; we can obtain more
	// than one certificate per email address, and still save them individually.
	for leEmail, serverConfigs := range initMap {
		// make client to service this email address with CA server
		client, err := newClient(leEmail)
		if err != nil {
			return configs, err
		}

		// client is ready, so let's get free, trusted SSL certificates! yeah!
		certificates, err := obtainCertificates(client, serverConfigs)
		if err != nil {
			return configs, err
		}

		// ... that's it. save the certs, keys, and metadata files to disk
		err = saveCertsAndKeys(certificates)
		if err != nil {
			return configs, err
		}

		// it all comes down to this: filling in the file path of a valid certificate automatically
		for _, cfg := range serverConfigs {
			cfg.TLS.Certificate = storage.SiteCertFile(cfg.Host)
			cfg.TLS.Key = storage.SiteKeyFile(cfg.Host)
			cfg.TLS.Enabled = true
			cfg.Port = "https"

			// Is there a plaintext HTTP config for the same host? If not, make
			// one and have it redirect all requests to this HTTPS host.
			var plaintextHostFound bool
			for _, otherCfg := range configs {
				if cfg.Host == otherCfg.Host && otherCfg.Port == "http" {
					plaintextHostFound = true
					break
				}
			}

			if !plaintextHostFound {
				// Make one that redirects to HTTPS for all requests
				configs = append(configs, redirPlaintextHost(*cfg))
			}
		}
	}

	return configs, nil
}

// groupConfigsByEmail groups configs by the Let's Encrypt email address
// associated to them or to the default Let's Encrypt email address. If the
// default email is not available, the user will be prompted to provide one.
func groupConfigsByEmail(configs []server.Config) (map[string][]*server.Config, error) {
	initMap := make(map[string][]*server.Config)
	for i := 0; i < len(configs); i++ {
		if configs[i].TLS.Certificate == "" && configs[i].TLS.Key == "" && configs[i].Port != "http" { // TODO: && !cfg.Host.IsLoopback()
			leEmail := getEmail(configs[i])
			if leEmail == "" {
				return nil, errors.New("must have email address to serve HTTPS without existing certificate and key")
			}
			initMap[leEmail] = append(initMap[leEmail], &configs[i])
		}
	}
	return initMap, nil
}

// newClient creates a new ACME client to facilitate communication
// with the Let's Encrypt CA server on behalf of the user specified
// by leEmail. As part of this process, a user will be loaded from
// disk (if already exists) or created new and registered via ACME
// and saved to the file system for next time.
func newClient(leEmail string) (*acme.Client, error) {
	// Look up or create the LE user account
	leUser, err := getUser(leEmail)
	if err != nil {
		return nil, err
	}

	// The client facilitates our communication with the CA server.
	client := acme.NewClient(caURL, &leUser, rsaKeySize, exposePort)

	// If not registered, the user must register an account with the CA
	// and agree to terms
	if leUser.Registration == nil {
		reg, err := client.Register()
		if err != nil {
			return nil, errors.New("registration error: " + err.Error())
		}
		leUser.Registration = reg

		// TODO: we can just do the agreement once: when registering, right?
		err = client.AgreeToTos()
		if err != nil {
			saveUser(leUser) // TODO: Might as well try, right? Error check?
			return nil, errors.New("error agreeing to terms: " + err.Error())
		}

		// save user to the file system
		err = saveUser(leUser)
		if err != nil {
			return nil, errors.New("could not save user: " + err.Error())
		}
	}

	return client, nil
}

// obtainCertificates obtains certificates from the CA server for
// the configurations in serverConfigs using client.
func obtainCertificates(client *acme.Client, serverConfigs []*server.Config) ([]acme.CertificateResource, error) {
	// collect all the hostnames into one slice
	var hosts []string
	for _, cfg := range serverConfigs {
		hosts = append(hosts, cfg.Host)
	}

	certificates, err := client.ObtainCertificates(hosts)
	if err != nil {
		return nil, errors.New("error obtaining certs: " + err.Error())
	}

	return certificates, nil
}

// saveCertificates saves each certificate resource to disk. This
// includes the certificate file itself, the private key, and the
// metadata file.
func saveCertsAndKeys(certificates []acme.CertificateResource) error {
	for _, cert := range certificates {
		os.MkdirAll(storage.Site(cert.Domain), 0700)

		// Save cert
		err := saveCertificate(cert.Certificate, storage.SiteCertFile(cert.Domain))
		if err != nil {
			return err
		}

		// Save private key
		err = ioutil.WriteFile(storage.SiteKeyFile(cert.Domain), cert.PrivateKey, 0600)
		if err != nil {
			return err
		}

		// Save cert metadata
		jsonBytes, err := json.MarshalIndent(&CertificateMeta{URL: cert.CertURL, Domain: cert.Domain}, "", "\t")
		if err != nil {
			return err
		}
		err = ioutil.WriteFile(storage.SiteMetaFile(cert.Domain), jsonBytes, 0600)
		if err != nil {
			return err
		}
	}
	return nil
}

// redirPlaintextHost returns a new plaintext HTTP configuration for
// a virtualHost that simply redirects to cfg, which is assumed to
// be the HTTPS configuration. The returned configuration is set
// to listen on the "http" port (port 80).
func redirPlaintextHost(cfg server.Config) server.Config {
	redirMidware := func(next middleware.Handler) middleware.Handler {
		return redirect.Redirect{Next: next, Rules: []redirect.Rule{
			{
				FromScheme: "http",
				FromPath:   "/",
				To:         "https://" + cfg.Host + "{uri}",
				Code:       http.StatusMovedPermanently,
			},
		}}
	}

	return server.Config{
		Host: cfg.Host,
		Port: "http",
		Middleware: map[string][]middleware.Middleware{
			"/": []middleware.Middleware{redirMidware},
		},
	}
}

var (
	// Let's Encrypt account email to use if none provided
	DefaultEmail string

	// Whether user has agreed to the Let's Encrypt SA
	Agreed bool
)

// Some essential values related to the Let's Encrypt process
const (
	// Size of RSA keys in bits
	rsaKeySize = 2048

	// The base URL to the Let's Encrypt CA
	caURL = "http://192.168.99.100:4000"

	// The port to expose to the CA server for Simple HTTP Challenge
	exposePort = "5001"
)

// KeySize represents the length of a key in bits
type KeySize int

// Key sizes
const (
	ECC_224  KeySize = 224
	ECC_256          = 256
	RSA_2048         = 2048
	RSA_4096         = 4096
)

type CertificateMeta struct {
	Domain, URL string
}
<commit_msg>Little more refactoring in letsencrypt<commit_after>package letsencrypt

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"os"

	"github.com/mholt/caddy/middleware"
	"github.com/mholt/caddy/middleware/redirect"
	"github.com/mholt/caddy/server"
	"github.com/xenolf/lego/acme"
)

// Activate sets up TLS for each server config in configs
// as needed. It only skips the config if the cert and key
// are already provided or if plaintext http is explicitly
// specified as the port.
//
// This function may prompt the user to provide an email
// address if none is available through other means. It
// prefers the email address specified in the config, but
// if that is not available it will check the command line
// argument. If absent, it will use the most recent email
// address from last time. If there isn't one, the user
// will be prompted. If the user leaves email blank, <TODO>.
func Activate(configs []server.Config) ([]server.Config, error) {
	// Group configs by LE email address; this will help us
	// reduce round-trips when getting the certs.
	initMap, err := groupConfigsByEmail(configs)
	if err != nil {
		return configs, err
	}

	// Loop through each email address and obtain certs; we can obtain more
	// than one certificate per email address, and still save them individually.
	for leEmail, serverConfigs := range initMap {
		// make client to service this email address with CA server
		client, err := newClient(leEmail)
		if err != nil {
			return configs, err
		}

		// client is ready, so let's get free, trusted SSL certificates! yeah!
		certificates, err := obtainCertificates(client, serverConfigs)
		if err != nil {
			return configs, err
		}

		// ... that's it. save the certs, keys, and metadata files to disk
		err = saveCertsAndKeys(certificates)
		if err != nil {
			return configs, err
		}

		// it all comes down to this: filling in the file path of a valid certificate automatically
		configs = autoConfigure(configs, serverConfigs)
	}

	return configs, nil
}

// groupConfigsByEmail groups configs by the Let's Encrypt email address
// associated to them or to the default Let's Encrypt email address. If the
// default email is not available, the user will be prompted to provide one.
func groupConfigsByEmail(configs []server.Config) (map[string][]*server.Config, error) {
	initMap := make(map[string][]*server.Config)
	for i := 0; i < len(configs); i++ {
		if configs[i].TLS.Certificate == "" && configs[i].TLS.Key == "" && configs[i].Port != "http" { // TODO: && !cfg.Host.IsLoopback()
			leEmail := getEmail(configs[i])
			if leEmail == "" {
				return nil, errors.New("must have email address to serve HTTPS without existing certificate and key")
			}
			initMap[leEmail] = append(initMap[leEmail], &configs[i])
		}
	}
	return initMap, nil
}

// newClient creates a new ACME client to facilitate communication
// with the Let's Encrypt CA server on behalf of the user specified
// by leEmail. As part of this process, a user will be loaded from
// disk (if already exists) or created new and registered via ACME
// and saved to the file system for next time.
func newClient(leEmail string) (*acme.Client, error) {
	// Look up or create the LE user account
	leUser, err := getUser(leEmail)
	if err != nil {
		return nil, err
	}

	// The client facilitates our communication with the CA server.
	client := acme.NewClient(caURL, &leUser, rsaKeySize, exposePort)

	// If not registered, the user must register an account with the CA
	// and agree to terms
	if leUser.Registration == nil {
		reg, err := client.Register()
		if err != nil {
			return nil, errors.New("registration error: " + err.Error())
		}
		leUser.Registration = reg

		// TODO: we can just do the agreement once: when registering, right?
		err = client.AgreeToTos()
		if err != nil {
			saveUser(leUser) // TODO: Might as well try, right? Error check?
			return nil, errors.New("error agreeing to terms: " + err.Error())
		}

		// save user to the file system
		err = saveUser(leUser)
		if err != nil {
			return nil, errors.New("could not save user: " + err.Error())
		}
	}

	return client, nil
}

// obtainCertificates obtains certificates from the CA server for
// the configurations in serverConfigs using client.
func obtainCertificates(client *acme.Client, serverConfigs []*server.Config) ([]acme.CertificateResource, error) {
	// collect all the hostnames into one slice
	var hosts []string
	for _, cfg := range serverConfigs {
		hosts = append(hosts, cfg.Host)
	}

	certificates, err := client.ObtainCertificates(hosts)
	if err != nil {
		return nil, errors.New("error obtaining certs: " + err.Error())
	}

	return certificates, nil
}

// saveCertificates saves each certificate resource to disk. This
// includes the certificate file itself, the private key, and the
// metadata file.
func saveCertsAndKeys(certificates []acme.CertificateResource) error {
	for _, cert := range certificates {
		os.MkdirAll(storage.Site(cert.Domain), 0700)

		// Save cert
		err := saveCertificate(cert.Certificate, storage.SiteCertFile(cert.Domain))
		if err != nil {
			return err
		}

		// Save private key
		err = ioutil.WriteFile(storage.SiteKeyFile(cert.Domain), cert.PrivateKey, 0600)
		if err != nil {
			return err
		}

		// Save cert metadata
		jsonBytes, err := json.MarshalIndent(&CertificateMeta{URL: cert.CertURL, Domain: cert.Domain}, "", "\t")
		if err != nil {
			return err
		}
		err = ioutil.WriteFile(storage.SiteMetaFile(cert.Domain), jsonBytes, 0600)
		if err != nil {
			return err
		}
	}
	return nil
}

// autoConfigure enables TLS on all the configs in serverConfigs
// and appends, if necessary, new configs to allConfigs that redirect
// plaintext HTTP to their HTTPS counterparts.
func autoConfigure(allConfigs []server.Config, serverConfigs []*server.Config) []server.Config {
	for _, cfg := range serverConfigs {
		cfg.TLS.Certificate = storage.SiteCertFile(cfg.Host)
		cfg.TLS.Key = storage.SiteKeyFile(cfg.Host)
		cfg.TLS.Enabled = true
		cfg.Port = "https"

		// Is there a plaintext HTTP config for the same host? If not, make
		// one and have it redirect all requests to this HTTPS host.
		var plaintextHostFound bool
		for _, otherCfg := range allConfigs {
			if cfg.Host == otherCfg.Host && otherCfg.Port == "http" {
				plaintextHostFound = true
				break
			}
		}

		if !plaintextHostFound {
			// Make one that redirects to HTTPS for all requests
			allConfigs = append(allConfigs, redirPlaintextHost(*cfg))
		}
	}

	return allConfigs
}

// redirPlaintextHost returns a new plaintext HTTP configuration for
// a virtualHost that simply redirects to cfg, which is assumed to
// be the HTTPS configuration. The returned configuration is set
// to listen on the "http" port (port 80).
func redirPlaintextHost(cfg server.Config) server.Config {
	redirMidware := func(next middleware.Handler) middleware.Handler {
		return redirect.Redirect{Next: next, Rules: []redirect.Rule{
			{
				FromScheme: "http",
				FromPath:   "/",
				To:         "https://" + cfg.Host + "{uri}",
				Code:       http.StatusMovedPermanently,
			},
		}}
	}

	return server.Config{
		Host: cfg.Host,
		Port: "http",
		Middleware: map[string][]middleware.Middleware{
			"/": []middleware.Middleware{redirMidware},
		},
	}
}

var (
	// Let's Encrypt account email to use if none provided
	DefaultEmail string

	// Whether user has agreed to the Let's Encrypt SA
	Agreed bool
)

// Some essential values related to the Let's Encrypt process
const (
	// Size of RSA keys in bits
	rsaKeySize = 2048

	// The base URL to the Let's Encrypt CA
	caURL = "http://192.168.99.100:4000"

	// The port to expose to the CA server for Simple HTTP Challenge
	exposePort = "5001"
)

// KeySize represents the length of a key in bits
type KeySize int

// Key sizes
const (
	ECC_224  KeySize = 224
	ECC_256          = 256
	RSA_2048         = 2048
	RSA_4096         = 4096
)

type CertificateMeta struct {
	Domain, URL string
}
<|endoftext|>


<commit_before>/*
	Copyright (c) 2014-2015, Percona LLC and/or its affiliates. All rights reserved.

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU Affero General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU Affero General Public License for more details.

	You should have received a copy of the GNU Affero General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

// Package slow implements a MySQL slow log parser.
package slow

import (
	"bufio"
	"fmt"
	"io"
	l "log"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/percona/go-mysql/log"
)

// Regular expressions to match important lines in slow log.
var (
	timeRe    = regexp.MustCompile(`Time: (\S+\s{1,2}\S+)`)
	timeNewRe = regexp.MustCompile(`Time:\s+(\d{4}-\d{2}-\d{2}\S+)`)
	userRe    = regexp.MustCompile(`User@Host: ([^\[]+|\[[^[]+\]).*?@ (\S*) \[(.*)\]`)
	schema    = regexp.MustCompile(`Schema: +(.*?) +Last_errno:`)
	headerRe  = regexp.MustCompile(`^#\s+[A-Z]`)
	metricsRe = regexp.MustCompile(`(\w+): (\S+|\z)`)
	adminRe   = regexp.MustCompile(`command: (.+)`)
	setRe     = regexp.MustCompile(`^SET (?:last_insert_id|insert_id|timestamp)`)
	useRe     = regexp.MustCompile(`^(?i)use `)
)

// A SlowLogParser parses a MySQL slow log. It implements the LogParser interface.
type SlowLogParser struct {
	file *os.File
	opt  log.Options
	// --
	stopChan    chan bool
	eventChan   chan *log.Event
	inHeader    bool
	inQuery     bool
	headerLines uint
	queryLines  uint64
	bytesRead   uint64
	lineOffset  uint64
	endOffset   uint64
	stopped     bool
	event       *log.Event
}

// NewSlowLogParser returns a new SlowLogParser that reads from the open file.
func NewSlowLogParser(file *os.File, opt log.Options) *SlowLogParser {
	if opt.DefaultLocation == nil {
		// Old MySQL format assumes time is taken from SYSTEM.
		opt.DefaultLocation = time.Local
	}
	p := &SlowLogParser{
		file: file,
		opt:  opt,
		// --
		stopChan:    make(chan bool, 1),
		eventChan:   make(chan *log.Event),
		inHeader:    false,
		inQuery:     false,
		headerLines: 0,
		queryLines:  0,
		lineOffset:  0,
		bytesRead:   opt.StartOffset,
		event:       log.NewEvent(),
	}
	return p
}

// EventChan returns the unbuffered event channel on which the caller can
// receive events.
func (p *SlowLogParser) EventChan() <-chan *log.Event {
	return p.eventChan
}

// Stop stops the parser before parsing the next event or while blocked on
// sending the current event to the event channel.
func (p *SlowLogParser) Stop() {
	if p.opt.Debug {
		l.Println("stopping")
	}
	p.stopChan <- true
	return
}

// Start starts the parser. Events are sent to the unbuffered event channel.
// Parsing stops on EOF, error, or call to Stop. The event channel is closed
// when parsing stops. The file is not closed.
func (p *SlowLogParser) Start() error {
	if p.opt.Debug {
		l.SetFlags(l.Ltime | l.Lmicroseconds)
		fmt.Println()
		l.Println("parsing " + p.file.Name())
	}

	// Seek to the offset, if any.
	// @todo error if start off > file size
	if p.opt.StartOffset > 0 {
		if _, err := p.file.Seek(int64(p.opt.StartOffset), os.SEEK_SET); err != nil {
			return err
		}
	}

	defer close(p.eventChan)

	r := bufio.NewReader(p.file)

SCANNER_LOOP:
	for !p.stopped {
		select {
		case <-p.stopChan:
			p.stopped = true
			break SCANNER_LOOP
		default:
		}

		line, err := r.ReadString('\n')
		if err != nil {
			if err != io.EOF {
				return err
			}
			break SCANNER_LOOP
		}

		lineLen := uint64(len(line))
		p.bytesRead += lineLen
		p.lineOffset = p.bytesRead - lineLen
		if p.opt.Debug {
			fmt.Println()
			l.Printf("+%d line: %s", p.lineOffset, line)
		}

		// Filter out meta lines:
		//   /usr/local/bin/mysqld, Version: 5.6.15-62.0-tokudb-7.1.0-tokudb-log (binary). started with:
		//   Tcp port: 3306  Unix socket: /var/lib/mysql/mysql.sock
		//   Time                 Id Command    Argument
		if lineLen >= 20 && ((line[0] == '/' && line[lineLen-6:lineLen] == "with:\n") ||
			(line[0:5] == "Time ") ||
			(line[0:4] == "Tcp ") ||
			(line[0:4] == "TCP ")) {
			if p.opt.Debug {
				l.Println("meta")
			}
			continue
		}

		// PMM-1834: Filter out empty comments and MariaDB explain:
		if line == "#\n" || strings.HasPrefix(line, "# explain:") {
			continue
		}

		// Remove \n.
		line = line[0 : lineLen-1]

		if p.inHeader {
			p.parseHeader(line)
		} else if p.inQuery {
			p.parseQuery(line)
		} else if headerRe.MatchString(line) {
			p.inHeader = true
			p.inQuery = false
			p.parseHeader(line)
		}
	}

	if !p.stopped && p.queryLines > 0 {
		p.endOffset = p.bytesRead
		p.sendEvent(false, false)
	}

	if p.opt.Debug {
		l.Printf("\ndone")
	}
	return nil
}

// --------------------------------------------------------------------------

func (p *SlowLogParser) parseHeader(line string) {
	if p.opt.Debug {
		l.Println("header")
	}

	if !headerRe.MatchString(line) {
		p.inHeader = false
		p.inQuery = true
		p.parseQuery(line)
		return
	}

	if p.headerLines == 0 {
		p.event.Offset = p.lineOffset
	}
	p.headerLines++

	if strings.HasPrefix(line, "# Time") {
		if p.opt.Debug {
			l.Println("time")
		}
		m := timeRe.FindStringSubmatch(line)
		if len(m) == 2 {
			p.event.Ts, _ = time.ParseInLocation("060102 15:04:05", m[1], p.opt.DefaultLocation)
		} else {
			m = timeNewRe.FindStringSubmatch(line)
			if len(m) == 2 {
				p.event.Ts, _ = time.Parse(time.RFC3339Nano, m[1])
			} else {
				return
			}
		}
		if userRe.MatchString(line) {
			if p.opt.Debug {
				l.Println("user (bad format)")
			}
			m := userRe.FindStringSubmatch(line)
			p.event.User = m[1]
			p.event.Host = m[2]
		}
	} else if strings.HasPrefix(line, "# User") {
		if p.opt.Debug {
			l.Println("user")
		}
		m := userRe.FindStringSubmatch(line)
		if len(m) < 3 {
			return
		}
		p.event.User = m[1]
		p.event.Host = m[2]
	} else if strings.HasPrefix(line, "# admin") {
		p.parseAdmin(line)
	} else {
		if p.opt.Debug {
			l.Println("metrics")
		}
		submatch := schema.FindStringSubmatch(line)
		if len(submatch) == 2 {
			p.event.Db = submatch[1]
		}

		m := metricsRe.FindAllStringSubmatch(line, -1)
		for _, smv := range m {
			// [String, Metric, Value], e.g. ["Query_time: 2", "Query_time", "2"]
			if strings.HasSuffix(smv[1], "_time") || strings.HasSuffix(smv[1], "_wait") {
				// microsecond value
				val, _ := strconv.ParseFloat(smv[2], 64)
				p.event.TimeMetrics[smv[1]] = val
			} else if smv[2] == "Yes" || smv[2] == "No" {
				// boolean value
				if smv[2] == "Yes" {
					p.event.BoolMetrics[smv[1]] = true
				} else {
					p.event.BoolMetrics[smv[1]] = false
				}
			} else if smv[1] == "Schema" {
				p.event.Db = smv[2]
			} else if smv[1] == "Log_slow_rate_type" {
				p.event.RateType = smv[2]
			} else if smv[1] == "Log_slow_rate_limit" {
				val, _ := strconv.ParseUint(smv[2], 10, 64)
				p.event.RateLimit = uint(val)
			} else {
				// integer value
				val, _ := strconv.ParseUint(smv[2], 10, 64)
				p.event.NumberMetrics[smv[1]] = val
			}
		}
	}
}

func (p *SlowLogParser) parseQuery(line string) {
	if p.opt.Debug {
		l.Println("query")
	}

	if strings.HasPrefix(line, "# admin") {
		p.parseAdmin(line)
		return
	} else if headerRe.MatchString(line) {
		if p.opt.Debug {
			l.Println("next event")
		}
		p.inHeader = true
		p.inQuery = false
		p.endOffset = p.lineOffset
		p.sendEvent(true, false)
		p.parseHeader(line)
		return
	}

	isUse := useRe.FindString(line)
	if p.queryLines == 0 && isUse != "" {
		if p.opt.Debug {
			l.Println("use db")
		}
		db := strings.TrimPrefix(line, isUse)
		db = strings.TrimRight(db, ";")
		db = strings.Trim(db, "`")
		p.event.Db = db
		// Set the 'use' as the query itself.
		// In case we are on a group of lines like in test 23, lines 6~8, the
		// query will be replaced by the real query "select field...."
		// In case we are on a group of lines like in test23, lines 27~28, the
		// query will be "use dbnameb" since the user executed a use command
		p.event.Query = line
	} else if setRe.MatchString(line) {
		if p.opt.Debug {
			l.Println("set var")
		}
		// @todo ignore or use these lines?
	} else {
		if p.opt.Debug {
			l.Println("query")
		}
		if p.queryLines > 0 {
			p.event.Query += "\n" + line
		} else {
			p.event.Query = line
		}
		p.queryLines++
	}
}

func (p *SlowLogParser) parseAdmin(line string) {
	if p.opt.Debug {
		l.Println("admin")
	}
	p.event.Admin = true
	m := adminRe.FindStringSubmatch(line)
	p.event.Query = m[1]
	p.event.Query = strings.TrimSuffix(p.event.Query, ";") // makes FilterAdminCommand work

	// admin commands should be the last line of the event.
	if filtered := p.opt.FilterAdminCommand[p.event.Query]; !filtered {
		if p.opt.Debug {
			l.Println("not filtered")
		}
		p.endOffset = p.bytesRead
		p.sendEvent(false, false)
	} else {
		p.inHeader = false
		p.inQuery = false
	}
}

func (p *SlowLogParser) sendEvent(inHeader bool, inQuery bool) {
	if p.opt.Debug {
		l.Println("send event")
	}

	p.event.OffsetEnd = p.endOffset

	// Make a new event and reset our metadata.
	defer func() {
		p.event = log.NewEvent()
		p.headerLines = 0
		p.queryLines = 0
		p.inHeader = inHeader
		p.inQuery = inQuery
	}()

	if _, ok := p.event.TimeMetrics["Query_time"]; !ok {
		if p.headerLines == 0 {
			l.Panicf("No Query_time in event at %d: %#v", p.lineOffset, p.event)
		}
		// Started parsing in header after Query_time.  Throw away event.
		return
	}

	// Clean up the event.
	p.event.Db = strings.TrimSuffix(p.event.Db, ";\n")
	p.event.Query = strings.TrimSuffix(p.event.Query, ";")

	// Send the event.  This will block.
	select {
	case p.eventChan <- p.event:
	case <-p.stopChan:
		p.stopped = true
	}
}
<commit_msg>PMM-2859: Use default location also for parsing new log format. This should only affect case when DefaultLocation = time.Local.<commit_after>/*
	Copyright (c) 2014-2015, Percona LLC and/or its affiliates. All rights reserved.

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU Affero General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU Affero General Public License for more details.

	You should have received a copy of the GNU Affero General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

// Package slow implements a MySQL slow log parser.
package slow

import (
	"bufio"
	"fmt"
	"io"
	l "log"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/percona/go-mysql/log"
)

// Regular expressions to match important lines in slow log.
var (
	timeRe    = regexp.MustCompile(`Time: (\S+\s{1,2}\S+)`)
	timeNewRe = regexp.MustCompile(`Time:\s+(\d{4}-\d{2}-\d{2}\S+)`)
	userRe    = regexp.MustCompile(`User@Host: ([^\[]+|\[[^[]+\]).*?@ (\S*) \[(.*)\]`)
	schema    = regexp.MustCompile(`Schema: +(.*?) +Last_errno:`)
	headerRe  = regexp.MustCompile(`^#\s+[A-Z]`)
	metricsRe = regexp.MustCompile(`(\w+): (\S+|\z)`)
	adminRe   = regexp.MustCompile(`command: (.+)`)
	setRe     = regexp.MustCompile(`^SET (?:last_insert_id|insert_id|timestamp)`)
	useRe     = regexp.MustCompile(`^(?i)use `)
)

// A SlowLogParser parses a MySQL slow log. It implements the LogParser interface.
type SlowLogParser struct {
	file *os.File
	opt  log.Options
	// --
	stopChan    chan bool
	eventChan   chan *log.Event
	inHeader    bool
	inQuery     bool
	headerLines uint
	queryLines  uint64
	bytesRead   uint64
	lineOffset  uint64
	endOffset   uint64
	stopped     bool
	event       *log.Event
}

// NewSlowLogParser returns a new SlowLogParser that reads from the open file.
func NewSlowLogParser(file *os.File, opt log.Options) *SlowLogParser {
	if opt.DefaultLocation == nil {
		// Old MySQL format assumes time is taken from SYSTEM.
		opt.DefaultLocation = time.Local
	}
	p := &SlowLogParser{
		file: file,
		opt:  opt,
		// --
		stopChan:    make(chan bool, 1),
		eventChan:   make(chan *log.Event),
		inHeader:    false,
		inQuery:     false,
		headerLines: 0,
		queryLines:  0,
		lineOffset:  0,
		bytesRead:   opt.StartOffset,
		event:       log.NewEvent(),
	}
	return p
}

// EventChan returns the unbuffered event channel on which the caller can
// receive events.
func (p *SlowLogParser) EventChan() <-chan *log.Event {
	return p.eventChan
}

// Stop stops the parser before parsing the next event or while blocked on
// sending the current event to the event channel.
func (p *SlowLogParser) Stop() {
	if p.opt.Debug {
		l.Println("stopping")
	}
	p.stopChan <- true
	return
}

// Start starts the parser. Events are sent to the unbuffered event channel.
// Parsing stops on EOF, error, or call to Stop. The event channel is closed
// when parsing stops. The file is not closed.
func (p *SlowLogParser) Start() error {
	if p.opt.Debug {
		l.SetFlags(l.Ltime | l.Lmicroseconds)
		fmt.Println()
		l.Println("parsing " + p.file.Name())
	}

	// Seek to the offset, if any.
	// @todo error if start off > file size
	if p.opt.StartOffset > 0 {
		if _, err := p.file.Seek(int64(p.opt.StartOffset), os.SEEK_SET); err != nil {
			return err
		}
	}

	defer close(p.eventChan)

	r := bufio.NewReader(p.file)

SCANNER_LOOP:
	for !p.stopped {
		select {
		case <-p.stopChan:
			p.stopped = true
			break SCANNER_LOOP
		default:
		}

		line, err := r.ReadString('\n')
		if err != nil {
			if err != io.EOF {
				return err
			}
			break SCANNER_LOOP
		}

		lineLen := uint64(len(line))
		p.bytesRead += lineLen
		p.lineOffset = p.bytesRead - lineLen
		if p.opt.Debug {
			fmt.Println()
			l.Printf("+%d line: %s", p.lineOffset, line)
		}

		// Filter out meta lines:
		//   /usr/local/bin/mysqld, Version: 5.6.15-62.0-tokudb-7.1.0-tokudb-log (binary). started with:
		//   Tcp port: 3306  Unix socket: /var/lib/mysql/mysql.sock
		//   Time                 Id Command    Argument
		if lineLen >= 20 && ((line[0] == '/' && line[lineLen-6:lineLen] == "with:\n") ||
			(line[0:5] == "Time ") ||
			(line[0:4] == "Tcp ") ||
			(line[0:4] == "TCP ")) {
			if p.opt.Debug {
				l.Println("meta")
			}
			continue
		}

		// PMM-1834: Filter out empty comments and MariaDB explain:
		if line == "#\n" || strings.HasPrefix(line, "# explain:") {
			continue
		}

		// Remove \n.
		line = line[0 : lineLen-1]

		if p.inHeader {
			p.parseHeader(line)
		} else if p.inQuery {
			p.parseQuery(line)
		} else if headerRe.MatchString(line) {
			p.inHeader = true
			p.inQuery = false
			p.parseHeader(line)
		}
	}

	if !p.stopped && p.queryLines > 0 {
		p.endOffset = p.bytesRead
		p.sendEvent(false, false)
	}

	if p.opt.Debug {
		l.Printf("\ndone")
	}
	return nil
}

// --------------------------------------------------------------------------

func (p *SlowLogParser) parseHeader(line string) {
	if p.opt.Debug {
		l.Println("header")
	}

	if !headerRe.MatchString(line) {
		p.inHeader = false
		p.inQuery = true
		p.parseQuery(line)
		return
	}

	if p.headerLines == 0 {
		p.event.Offset = p.lineOffset
	}
	p.headerLines++

	if strings.HasPrefix(line, "# Time") {
		if p.opt.Debug {
			l.Println("time")
		}
		m := timeRe.FindStringSubmatch(line)
		if len(m) == 2 {
			p.event.Ts, _ = time.ParseInLocation("060102 15:04:05", m[1], p.opt.DefaultLocation)
		} else {
			m = timeNewRe.FindStringSubmatch(line)
			if len(m) == 2 {
				p.event.Ts, _ = time.ParseInLocation(time.RFC3339Nano, m[1], p.opt.DefaultLocation)
			} else {
				return
			}
		}
		if userRe.MatchString(line) {
			if p.opt.Debug {
				l.Println("user (bad format)")
			}
			m := userRe.FindStringSubmatch(line)
			p.event.User = m[1]
			p.event.Host = m[2]
		}
	} else if strings.HasPrefix(line, "# User") {
		if p.opt.Debug {
			l.Println("user")
		}
		m := userRe.FindStringSubmatch(line)
		if len(m) < 3 {
			return
		}
		p.event.User = m[1]
		p.event.Host = m[2]
	} else if strings.HasPrefix(line, "# admin") {
		p.parseAdmin(line)
	} else {
		if p.opt.Debug {
			l.Println("metrics")
		}
		submatch := schema.FindStringSubmatch(line)
		if len(submatch) == 2 {
			p.event.Db = submatch[1]
		}

		m := metricsRe.FindAllStringSubmatch(line, -1)
		for _, smv := range m {
			// [String, Metric, Value], e.g. ["Query_time: 2", "Query_time", "2"]
			if strings.HasSuffix(smv[1], "_time") || strings.HasSuffix(smv[1], "_wait") {
				// microsecond value
				val, _ := strconv.ParseFloat(smv[2], 64)
				p.event.TimeMetrics[smv[1]] = val
			} else if smv[2] == "Yes" || smv[2] == "No" {
				// boolean value
				if smv[2] == "Yes" {
					p.event.BoolMetrics[smv[1]] = true
				} else {
					p.event.BoolMetrics[smv[1]] = false
				}
			} else if smv[1] == "Schema" {
				p.event.Db = smv[2]
			} else if smv[1] == "Log_slow_rate_type" {
				p.event.RateType = smv[2]
			} else if smv[1] == "Log_slow_rate_limit" {
				val, _ := strconv.ParseUint(smv[2], 10, 64)
				p.event.RateLimit = uint(val)
			} else {
				// integer value
				val, _ := strconv.ParseUint(smv[2], 10, 64)
				p.event.NumberMetrics[smv[1]] = val
			}
		}
	}
}

func (p *SlowLogParser) parseQuery(line string) {
	if p.opt.Debug {
		l.Println("query")
	}

	if strings.HasPrefix(line, "# admin") {
		p.parseAdmin(line)
		return
	} else if headerRe.MatchString(line) {
		if p.opt.Debug {
			l.Println("next event")
		}
		p.inHeader = true
		p.inQuery = false
		p.endOffset = p.lineOffset
		p.sendEvent(true, false)
		p.parseHeader(line)
		return
	}

	isUse := useRe.FindString(line)
	if p.queryLines == 0 && isUse != "" {
		if p.opt.Debug {
			l.Println("use db")
		}
		db := strings.TrimPrefix(line, isUse)
		db = strings.TrimRight(db, ";")
		db = strings.Trim(db, "`")
		p.event.Db = db
		// Set the 'use' as the query itself.
		// In case we are on a group of lines like in test 23, lines 6~8, the
		// query will be replaced by the real query "select field...."
		// In case we are on a group of lines like in test23, lines 27~28, the
		// query will be "use dbnameb" since the user executed a use command
		p.event.Query = line
	} else if setRe.MatchString(line) {
		if p.opt.Debug {
			l.Println("set var")
		}
		// @todo ignore or use these lines?
	} else {
		if p.opt.Debug {
			l.Println("query")
		}
		if p.queryLines > 0 {
			p.event.Query += "\n" + line
		} else {
			p.event.Query = line
		}
		p.queryLines++
	}
}

func (p *SlowLogParser) parseAdmin(line string) {
	if p.opt.Debug {
		l.Println("admin")
	}
	p.event.Admin = true
	m := adminRe.FindStringSubmatch(line)
	p.event.Query = m[1]
	p.event.Query = strings.TrimSuffix(p.event.Query, ";") // makes FilterAdminCommand work

	// admin commands should be the last line of the event.
	if filtered := p.opt.FilterAdminCommand[p.event.Query]; !filtered {
		if p.opt.Debug {
			l.Println("not filtered")
		}
		p.endOffset = p.bytesRead
		p.sendEvent(false, false)
	} else {
		p.inHeader = false
		p.inQuery = false
	}
}

func (p *SlowLogParser) sendEvent(inHeader bool, inQuery bool) {
	if p.opt.Debug {
		l.Println("send event")
	}

	p.event.OffsetEnd = p.endOffset

	// Make a new event and reset our metadata.
	defer func() {
		p.event = log.NewEvent()
		p.headerLines = 0
		p.queryLines = 0
		p.inHeader = inHeader
		p.inQuery = inQuery
	}()

	if _, ok := p.event.TimeMetrics["Query_time"]; !ok {
		if p.headerLines == 0 {
			l.Panicf("No Query_time in event at %d: %#v", p.lineOffset, p.event)
		}
		// Started parsing in header after Query_time.  Throw away event.
		return
	}

	// Clean up the event.
	p.event.Db = strings.TrimSuffix(p.event.Db, ";\n")
	p.event.Query = strings.TrimSuffix(p.event.Query, ";")

	// Send the event.  This will block.
	select {
	case p.eventChan <- p.event:
	case <-p.stopChan:
		p.stopped = true
	}
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"github.com/apex/log"
	"github.com/apex/log/handlers/cli"
	color2 "github.com/fatih/color"
	"github.com/mattn/go-colorable"
	"github.com/pkg/errors"
	"io"
	"os"
	"sync"
	"time"
)

var Default = New(os.Stderr, true)

var bold = color2.New(color2.Bold)

var Strings = [...]string{
	log.DebugLevel: "DEBUG",
	log.InfoLevel:  " INFO",
	log.WarnLevel:  " WARN",
	log.ErrorLevel: "ERROR",
	log.FatalLevel: "FATAL",
}

type Handler struct {
	mu      sync.Mutex
	Writer  io.Writer
	Padding int
}

func New(w io.Writer, useColors bool) *Handler {
	if f, ok := w.(*os.File); ok {
		if useColors {
			return &Handler{Writer: colorable.NewColorable(f), Padding: 2}
		}
	}

	return &Handler{Writer: colorable.NewNonColorable(w), Padding: 2}
}

type tracer interface {
	StackTrace() errors.StackTrace
}

// HandleLog implements log.Handler.
func (h *Handler) HandleLog(e *log.Entry) error {
	color := cli.Colors[e.Level]
	level := Strings[e.Level]
	names := e.Fields.Names()

	h.mu.Lock()
	defer h.mu.Unlock()

	color.Fprintf(h.Writer, "%s: [%s] %-25s", bold.Sprintf("%*s", h.Padding+1, level), time.Now().Format(time.StampMilli), e.Message)

	for _, name := range names {
		if name == "source" {
			continue
		}

		fmt.Fprintf(h.Writer, " %s=%v", color.Sprint(name), e.Fields.Get(name))
	}

	fmt.Fprintln(h.Writer)

	for _, name := range names {
		if name != "error" {
			continue
		}

		if err, ok := e.Fields.Get("error").(error); ok {
			br := color2.New(color2.Bold, color2.FgRed)

			if e, ok := errors.Cause(err).(tracer); ok {
				st := e.StackTrace()

				l := len(st)
				if l > 10 {
					l = 10
				}

				fmt.Fprintf(h.Writer, "\n%s%+v\n\n", br.Sprintf("Stacktrace:"), st[0:l])
			} else {
				fmt.Fprintf(h.Writer, "\n%s\n%+v\n\n", br.Sprintf("Stacktrace:"), err)
			}
		} else {
			fmt.Printf("\n\nINVALID ERROR\n\n")
		}
	}

	return nil
}
<commit_msg>Simplify logic<commit_after>package cli

import (
	"fmt"
	"github.com/apex/log"
	"github.com/apex/log/handlers/cli"
	color2 "github.com/fatih/color"
	"github.com/mattn/go-colorable"
	"github.com/pkg/errors"
	"io"
	"math"
	"os"
	"sync"
	"time"
)

var Default = New(os.Stderr, true)

var bold = color2.New(color2.Bold)
var boldred = bold.Add(color2.FgRed)

var Strings = [...]string{
	log.DebugLevel: "DEBUG",
	log.InfoLevel:  " INFO",
	log.WarnLevel:  " WARN",
	log.ErrorLevel: "ERROR",
	log.FatalLevel: "FATAL",
}

type Handler struct {
	mu      sync.Mutex
	Writer  io.Writer
	Padding int
}

func New(w io.Writer, useColors bool) *Handler {
	if f, ok := w.(*os.File); ok {
		if useColors {
			return &Handler{Writer: colorable.NewColorable(f), Padding: 2}
		}
	}

	return &Handler{Writer: colorable.NewNonColorable(w), Padding: 2}
}

type tracer interface {
	StackTrace() errors.StackTrace
}

// HandleLog implements log.Handler.
func (h *Handler) HandleLog(e *log.Entry) error {
	color := cli.Colors[e.Level]
	level := Strings[e.Level]
	names := e.Fields.Names()

	h.mu.Lock()
	defer h.mu.Unlock()

	color.Fprintf(h.Writer, "%s: [%s] %-25s", bold.Sprintf("%*s", h.Padding+1, level), time.Now().Format(time.StampMilli), e.Message)

	for _, name := range names {
		if name == "source" {
			continue
		}
		fmt.Fprintf(h.Writer, " %s=%v", color.Sprint(name), e.Fields.Get(name))
	}

	fmt.Fprintln(h.Writer)

	for _, name := range names {
		if name != "error" {
			continue
		}
		if err, ok := e.Fields.Get("error").(error); ok {
			if e, ok := errors.Cause(err).(tracer); ok {
				st := e.StackTrace()
				l := math.Min(float64(len(st)), 10)
				fmt.Fprintf(h.Writer, "\n%s%+v\n\n", boldred.Sprintf("Stacktrace:"), st[0:int(l)])
			} else {
				fmt.Fprintf(h.Writer, "\n%s\n%+v\n\n", boldred.Sprintf("Stacktrace:"), err)
			}
		}
	}

	return nil
}
<|endoftext|>


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

// Package logging implements log library for other applications. It provides
// functions Debug, Info, Warning, Error, Critical, and formatting version
// Logf.
//
// Example:
//
//	logger := logging.SimpleLogger("main")
//	logger.SetLevel(logging.WARNING)
//	logger.Error("test for error")
//	logger.Warning("test for warning", "second parameter")
//	logger.Debug("test for debug")
//
package logging

import (
	"io"
	"os"
	"sync"
	"time"
)

// Pre-defined formats
const (
	defaultFileName   = "logging.log"                   // default logging filename
	configFileName    = "logging.conf"                  // default configuration filename
	defaultTimeFormat = "2006-01-02 15:04:05.999999999" // defaulttime format
	bufSize           = 1000                            // buffer size for writer
	queueSize         = 1000                            // chan queue size in async logging
)

// Logger is the logging struct.
type Logger struct {

	// Be careful of the alignment issue of the variable seqid because it
	// uses the sync/atomic.AddUint64() operation. If the alignment is
	// wrong, it will cause a panic. To solve the alignment issue in an
	// easy way, we put seqid to the beginning of the structure.
	// seqid is only visiable internally.
	seqid uint64 // last used sequence number in record

	// These variables can be configured by users.
	name       string    // logger name
	level      Level     // record level higher than this will be printed
	format     string    // format of the record
	fargs      []string  // arguments to be used in the format
	out        io.Writer // writer
	startTime  time.Time // start time of the logger
	sync       bool      // use sync or async way to record logs
	timeFormat string    // format for time

	// Internal used variables, which don't have get and set functions.
	lock  sync.Mutex  // writer lock
	queue chan string // queue used in async logging
	flush chan bool   // flush signal for the watcher to write
	quit  chan bool   // quit signal for the watcher to quit
	fd    *os.File    // file handler, used to close the file on destroy
}

// SimpleLogger creates a new logger with simple configuration.
func SimpleLogger(name string) (*Logger, error) {
	return createLogger(name, WARNING, BasicFormat, os.Stdout, false)
}

// BasicLogger creates a new logger with basic configuration.
func BasicLogger(name string) (*Logger, error) {
	return FileLogger(name, WARNING, BasicFormat, defaultFileName, false)
}

// RichLogger creates a new logger with simple configuration.
func RichLogger(name string) (*Logger, error) {
	return FileLogger(name, NOTSET, RichFormat, defaultFileName, false)
}

// FileLogger creates a new logger with file output.
func FileLogger(name string, level Level, format string, file string, sync bool) (*Logger, error) {
	out, err := os.Create(file)
	if err != nil {
		return new(Logger), err
	}
	logger, err := createLogger(name, level, format, out, sync)
	if err == nil {
		logger.fd = out
	}
	return logger, err
}

// createLogger create a new logger
func createLogger(name string, level Level, format string, out io.Writer, sync bool) (*Logger, error) {
	logger := new(Logger)

	err := logger.SetFormat(format)
	if err != nil {
		return logger, err
	}

	// asign values to logger
	logger.name = name
	logger.level = level
	logger.out = out
	logger.seqid = 0
	logger.sync = sync
	logger.queue = make(chan string, queueSize)
	logger.flush = make(chan bool)
	logger.quit = make(chan bool)
	logger.startTime = time.Now()
	logger.fd = nil
	logger.timeFormat = defaultTimeFormat

	// start watcher to write logs if it is async
	if sync == false {
		go logger.watcher()
	}

	return logger, nil
}

// Destroy sends quit signal to watcher and releases all the resources.
func (logger *Logger) Destroy() {

	logger.quitWatcher()

	// clean up
	if logger.fd != nil {
		logger.fd.Close()
	}
}

func (logger *Logger) quitWatcher() {
	if logger.sync == false {
		// quit watcher
		logger.quit <- true

		// wait for watcher quit
		<-logger.quit
	}
}

// Flush the writer
func (logger *Logger) Flush() {
	if logger.sync == false {
		logger.flush <- true
	}
}

// Get and set the configuration of the logger

func (logger *Logger) Name() string {
	return logger.name
}

func (logger *Logger) SetName(name string) {
	logger.name = name
}

func (logger *Logger) TimeFormat() string {
	return logger.timeFormat
}

func (logger *Logger) SetTimeFormat(format string) {
	logger.timeFormat = format
}

func (logger *Logger) Level() Level {
	return logger.level
}

func (logger *Logger) SetLevel(level Level) {
	logger.level = Level(level)
}

func (logger *Logger) LevelName() string {
	name, _ := levelNames[logger.level]
	return name
}

func (logger *Logger) SetLevelName(name string) {
	level, ok := levelValues[name]
	if ok {
		logger.level = level
	}
}

func (logger *Logger) Format() string {
	return logger.format
}

func (logger *Logger) Fargs() []string {
	return logger.fargs
}

func (logger *Logger) SetFormat(format string) error {
	return logger.parseFormat(format)
}

func (logger *Logger) Writer() io.Writer {
	return logger.out
}

func (logger *Logger) AddWriter(out io.Writer) {
	logger.out = io.MultiWriter(logger.out, out)
}

func (logger *Logger) AddWriters(out ...io.Writer) {
	logger.out = io.MultiWriter(append(out, logger.out)...)
}

func (logger *Logger) SetWriter(out io.Writer) {
	logger.out = out
}

func (logger *Logger) SetWriters(out ...io.Writer) {
	logger.out = io.MultiWriter(out...)
}

func (logger *Logger) Sync() bool {
	return logger.sync
}

func (logger *Logger) SetSync(sync bool) {
	if sync == logger.sync {
		return
	}
	logger.quitWatcher()
	logger.sync = sync
	if logger.sync == false {
		go logger.watcher()
	}
}
<commit_msg>add writerlogger()<commit_after>// Copyright 2013, Cong Ding. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// author: Cong Ding <dinggnu@gmail.com>

// Package logging implements log library for other applications. It provides
// functions Debug, Info, Warning, Error, Critical, and formatting version
// Logf.
//
// Example:
//
//	logger := logging.SimpleLogger("main")
//	logger.SetLevel(logging.WARNING)
//	logger.Error("test for error")
//	logger.Warning("test for warning", "second parameter")
//	logger.Debug("test for debug")
//
package logging

import (
	"io"
	"os"
	"sync"
	"time"
)

// Pre-defined formats
const (
	defaultFileName   = "logging.log"                   // default logging filename
	configFileName    = "logging.conf"                  // default configuration filename
	defaultTimeFormat = "2006-01-02 15:04:05.999999999" // defaulttime format
	bufSize           = 1000                            // buffer size for writer
	queueSize         = 1000                            // chan queue size in async logging
)

// Logger is the logging struct.
type Logger struct {

	// Be careful of the alignment issue of the variable seqid because it
	// uses the sync/atomic.AddUint64() operation. If the alignment is
	// wrong, it will cause a panic. To solve the alignment issue in an
	// easy way, we put seqid to the beginning of the structure.
	// seqid is only visiable internally.
	seqid uint64 // last used sequence number in record

	// These variables can be configured by users.
	name       string    // logger name
	level      Level     // record level higher than this will be printed
	format     string    // format of the record
	fargs      []string  // arguments to be used in the format
	out        io.Writer // writer
	startTime  time.Time // start time of the logger
	sync       bool      // use sync or async way to record logs
	timeFormat string    // format for time

	// Internal used variables, which don't have get and set functions.
	lock  sync.Mutex  // writer lock
	queue chan string // queue used in async logging
	flush chan bool   // flush signal for the watcher to write
	quit  chan bool   // quit signal for the watcher to quit
	fd    *os.File    // file handler, used to close the file on destroy
}

// SimpleLogger creates a new logger with simple configuration.
func SimpleLogger(name string) (*Logger, error) {
	return createLogger(name, WARNING, BasicFormat, os.Stdout, false)
}

// BasicLogger creates a new logger with basic configuration.
func BasicLogger(name string) (*Logger, error) {
	return FileLogger(name, WARNING, BasicFormat, defaultFileName, false)
}

// RichLogger creates a new logger with simple configuration.
func RichLogger(name string) (*Logger, error) {
	return FileLogger(name, NOTSET, RichFormat, defaultFileName, false)
}

// FileLogger creates a new logger with file output.
func FileLogger(name string, level Level, format string, file string, sync bool) (*Logger, error) {
	out, err := os.Create(file)
	if err != nil {
		return new(Logger), err
	}
	logger, err := createLogger(name, level, format, out, sync)
	if err == nil {
		logger.fd = out
	}
	return logger, err
}

func WriterLogger(name string, level Level, format string, out io.Writer, sync bool) (*Logger, error) {
	return createLogger(name, level, format, out, sync)
}

// createLogger create a new logger
func createLogger(name string, level Level, format string, out io.Writer, sync bool) (*Logger, error) {
	logger := new(Logger)

	err := logger.SetFormat(format)
	if err != nil {
		return logger, err
	}

	// asign values to logger
	logger.name = name
	logger.level = level
	logger.out = out
	logger.seqid = 0
	logger.sync = sync
	logger.queue = make(chan string, queueSize)
	logger.flush = make(chan bool)
	logger.quit = make(chan bool)
	logger.startTime = time.Now()
	logger.fd = nil
	logger.timeFormat = defaultTimeFormat

	// start watcher to write logs if it is async
	if sync == false {
		go logger.watcher()
	}

	return logger, nil
}

// Destroy sends quit signal to watcher and releases all the resources.
func (logger *Logger) Destroy() {

	logger.quitWatcher()

	// clean up
	if logger.fd != nil {
		logger.fd.Close()
	}
}

func (logger *Logger) quitWatcher() {
	if logger.sync == false {
		// quit watcher
		logger.quit <- true

		// wait for watcher quit
		<-logger.quit
	}
}

// Flush the writer
func (logger *Logger) Flush() {
	if logger.sync == false {
		logger.flush <- true
	}
}

// Get and set the configuration of the logger

func (logger *Logger) Name() string {
	return logger.name
}

func (logger *Logger) SetName(name string) {
	logger.name = name
}

func (logger *Logger) TimeFormat() string {
	return logger.timeFormat
}

func (logger *Logger) SetTimeFormat(format string) {
	logger.timeFormat = format
}

func (logger *Logger) Level() Level {
	return logger.level
}

func (logger *Logger) SetLevel(level Level) {
	logger.level = Level(level)
}

func (logger *Logger) LevelName() string {
	name, _ := levelNames[logger.level]
	return name
}

func (logger *Logger) SetLevelName(name string) {
	level, ok := levelValues[name]
	if ok {
		logger.level = level
	}
}

func (logger *Logger) Format() string {
	return logger.format
}

func (logger *Logger) Fargs() []string {
	return logger.fargs
}

func (logger *Logger) SetFormat(format string) error {
	return logger.parseFormat(format)
}

func (logger *Logger) Writer() io.Writer {
	return logger.out
}

func (logger *Logger) AddWriter(out io.Writer) {
	logger.out = io.MultiWriter(logger.out, out)
}

func (logger *Logger) AddWriters(out ...io.Writer) {
	logger.out = io.MultiWriter(append(out, logger.out)...)
}

func (logger *Logger) SetWriter(out io.Writer) {
	logger.out = out
}

func (logger *Logger) SetWriters(out ...io.Writer) {
	logger.out = io.MultiWriter(out...)
}

func (logger *Logger) Sync() bool {
	return logger.sync
}

func (logger *Logger) SetSync(sync bool) {
	if sync == logger.sync {
		return
	}
	logger.quitWatcher()
	logger.sync = sync
	if logger.sync == false {
		go logger.watcher()
	}
}
<|endoftext|>


<commit_before>package logging

import (
	"encoding/json"
	"io/ioutil"
	"log/syslog"

	"github.com/Shopify/sarama"
	"github.com/mailgun/kafka-pixy/config"
	"github.com/mailgun/logrus-hooks/kafkahook"
	"github.com/mailgun/logrus-hooks/levelfilter"
	"github.com/pkg/errors"
	"github.com/samuel/go-zookeeper/zk"
	log "github.com/sirupsen/logrus"
	"github.com/sirupsen/logrus/hooks/syslog"
)

// Init initializes sirupsen/logrus hooks from the JSON config string. It also
// sets the sirupsen/logrus as a logger for 3rd party libraries.
func Init(jsonCfg string, cfg *config.App) error {
	var loggingCfg []loggerCfg
	if err := json.Unmarshal([]byte(jsonCfg), &loggingCfg); err != nil {
		return errors.Wrap(err, "failed to parse logger config")
	}

	formatter := &textFormatter{}
	log.SetFormatter(formatter)

	stdoutEnabled := false
	nonStdoutEnabled := false
	for _, loggerCfg := range loggingCfg {
		switch loggerCfg.Name {
		case "console":
			stdoutEnabled = true
		case "syslog":
			h, err := logrus_syslog.NewSyslogHook("udp", "127.0.0.1:514", syslog.LOG_INFO|syslog.LOG_MAIL, "kafka-pixy")
			if err != nil {
				continue
			}
			log.AddHook(levelfilter.New(h, loggerCfg.level()))
			nonStdoutEnabled = true
		case "udplog":
			if cfg == nil {
				return errors.Errorf("App config must be provided")
			}
			// If a Kafka cluster is not specified in logging config or does
			// not exist in the Kafka-Pixy config, then the default cluster is
			// used.
			cluster := loggerCfg.Params["cluster"]
			proxyCfg := cfg.Proxies[cluster]
			if proxyCfg == nil {
				proxyCfg = cfg.Proxies[cfg.DefaultCluster]
			}
			// If the log topic is not specified then "udplog" is assumed.
			topic := loggerCfg.Params["topic"]
			if topic == "" {
				topic = "udplog"
			}
			h, err := kafkahook.New(kafkahook.Config{
				Endpoints: proxyCfg.Kafka.SeedPeers,
				Topic:     topic,
			})
			if err != nil {
				continue
			}
			log.AddHook(levelfilter.New(h, loggerCfg.level()))
			nonStdoutEnabled = true
		}
	}

	saramaLogger := log.New()
	saramaLogger.Formatter = &saramaFormatter{formatter}
	sarama.Logger = saramaLogger

	zk.DefaultLogger = log.StandardLogger()

	if !stdoutEnabled || nonStdoutEnabled {
		log.SetOutput(ioutil.Discard)
		saramaLogger.Out = ioutil.Discard
	}
	return nil
}

// loggerCfg represents a configuration of an individual logger.
type loggerCfg struct {
	// Name defines a logger to be used. It can be one of: console, syslog, or
	// udplog.
	Name string `json:"name"`

	// Severity indicates the minimum severity a logger will be logging messages at.
	Severity string `json:"severity"`

	// Logger parameters
	Params map[string]string `json:"params"`
}

func (lc *loggerCfg) level() log.Level {
	level, err := log.ParseLevel(lc.Severity)
	if err != nil {
		return log.WarnLevel
	}
	return level
}

// saramaFormatter is a sirupsen/logrus formatter that strips trailing new
// lines from the log lines.
type saramaFormatter struct {
	parentFormatter log.Formatter
}

func (f *saramaFormatter) Format(entry *log.Entry) ([]byte, error) {
	lastByteIdx := len(entry.Message) - 1
	if lastByteIdx >= 0 && entry.Message[lastByteIdx] == '\n' {
		entry.Message = entry.Message[:lastByteIdx]
	}
	return f.parentFormatter.Format(entry)
}
<commit_msg>Add category to 3rd-party loggers<commit_after>package logging

import (
	"encoding/json"
	"io/ioutil"
	"log/syslog"

	"github.com/Shopify/sarama"
	"github.com/mailgun/kafka-pixy/config"
	"github.com/mailgun/logrus-hooks/kafkahook"
	"github.com/mailgun/logrus-hooks/levelfilter"
	"github.com/pkg/errors"
	"github.com/samuel/go-zookeeper/zk"
	log "github.com/sirupsen/logrus"
	"github.com/sirupsen/logrus/hooks/syslog"
)

// Init initializes sirupsen/logrus hooks from the JSON config string. It also
// sets the sirupsen/logrus as a logger for 3rd party libraries.
func Init(jsonCfg string, cfg *config.App) error {
	var loggingCfg []loggerCfg
	if err := json.Unmarshal([]byte(jsonCfg), &loggingCfg); err != nil {
		return errors.Wrap(err, "failed to parse logger config")
	}

	formatter := &textFormatter{}
	log.SetFormatter(formatter)

	stdoutEnabled := false
	nonStdoutEnabled := false
	for _, loggerCfg := range loggingCfg {
		switch loggerCfg.Name {
		case "console":
			stdoutEnabled = true
		case "syslog":
			h, err := logrus_syslog.NewSyslogHook("udp", "127.0.0.1:514", syslog.LOG_INFO|syslog.LOG_MAIL, "kafka-pixy")
			if err != nil {
				continue
			}
			log.AddHook(levelfilter.New(h, loggerCfg.level()))
			nonStdoutEnabled = true
		case "udplog":
			if cfg == nil {
				return errors.Errorf("App config must be provided")
			}
			// If a Kafka cluster is not specified in logging config or does
			// not exist in the Kafka-Pixy config, then the default cluster is
			// used.
			cluster := loggerCfg.Params["cluster"]
			proxyCfg := cfg.Proxies[cluster]
			if proxyCfg == nil {
				proxyCfg = cfg.Proxies[cfg.DefaultCluster]
			}
			// If the log topic is not specified then "udplog" is assumed.
			topic := loggerCfg.Params["topic"]
			if topic == "" {
				topic = "udplog"
			}
			h, err := kafkahook.New(kafkahook.Config{
				Endpoints: proxyCfg.Kafka.SeedPeers,
				Topic:     topic,
			})
			if err != nil {
				continue
			}
			log.AddHook(levelfilter.New(h, loggerCfg.level()))
			nonStdoutEnabled = true
		}
	}

	saramaLogger := log.New()
	saramaLogger.Formatter = &saramaFormatter{formatter}
	sarama.Logger = saramaLogger.WithField("category", "sarama")

	zk.DefaultLogger = log.StandardLogger().WithField("category", "zk")

	if !stdoutEnabled || nonStdoutEnabled {
		log.SetOutput(ioutil.Discard)
		saramaLogger.Out = ioutil.Discard
	}
	return nil
}

// loggerCfg represents a configuration of an individual logger.
type loggerCfg struct {
	// Name defines a logger to be used. It can be one of: console, syslog, or
	// udplog.
	Name string `json:"name"`

	// Severity indicates the minimum severity a logger will be logging messages at.
	Severity string `json:"severity"`

	// Logger parameters
	Params map[string]string `json:"params"`
}

func (lc *loggerCfg) level() log.Level {
	level, err := log.ParseLevel(lc.Severity)
	if err != nil {
		return log.WarnLevel
	}
	return level
}

// saramaFormatter is a sirupsen/logrus formatter that strips trailing new
// lines from the log lines.
type saramaFormatter struct {
	parentFormatter log.Formatter
}

func (f *saramaFormatter) Format(entry *log.Entry) ([]byte, error) {
	lastByteIdx := len(entry.Message) - 1
	if lastByteIdx >= 0 && entry.Message[lastByteIdx] == '\n' {
		entry.Message = entry.Message[:lastByteIdx]
	}
	return f.parentFormatter.Format(entry)
}
<|endoftext|>


<commit_before>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 blob

import (
	"bytes"
	"crypto/md5"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"strings"

	"golang.org/x/net/context"

	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcsutil"
)

// A key placed in GCS object metadata by GCSStore containing the hex SHA-1
// expected for the object contents. This is of course redundant with the
// object name; we use it as a paranoid check against GCS returning the
// metadata or contents for the wrong object.
const metadataKey_SHA1Hex = "comeback_sha1_hex"

// A key placed in GCS object metadata by GCSStore containing the hex MD5 sum
// expected for the object contents. If GCS reports a different MD5 sum or
// returns contents with a different MD5 sum, we know something screwy has
// happened.
//
// See here for more info: https://github.com/jacobsa/comeback/issues/18
const metadataKey_MD5Hex = "comeback_md5_hex"

// Return a blob store that stores blobs in the supplied GCS bucket. GCS object
// names look like:
//
//     <prefix><score>
//
// where <score> is the result of calling Score.Hex.
//
// The blob store trusts that it has full ownership of this portion of the
// bucket's namespace -- if a score name exists, then it points to the correct
// data.
//
// The returned store does not support Flush or Contains; these methods must
// not be called.
func NewGCSStore(
	bucket gcs.Bucket,
	prefix string) (store *GCSStore) {
	store = &GCSStore{
		bucket:     bucket,
		namePrefix: prefix,
	}

	return
}

type GCSStore struct {
	bucket     gcs.Bucket
	namePrefix string
}

// Parse and verify the internal consistency of the supplied object record in
// the same manner that a GCSStore configured with the supplied object name
// prefix would. Return the score of the blob that the object contains.
func ParseObjectRecord(
	o *gcs.Object,
	namePrefix string) (score Score, err error) {
	// Is the name of the appropriate form?
	if !strings.HasPrefix(o.Name, namePrefix) {
		err = fmt.Errorf("Unexpected object name: %q", o.Name)
		return
	}

	// Parse the hex score.
	hexScore := strings.TrimPrefix(o.Name, namePrefix)
	score, err = ParseHexScore(hexScore)
	if err != nil {
		err = fmt.Errorf("Unexpected hex score %q: %v", hexScore, err)
		return
	}

	// We expect the hex score to match the hex SHA-1 in the metadata.
	hexSHA1, ok := o.Metadata[metadataKey_SHA1Hex]
	if !ok {
		err = fmt.Errorf(
			"Object %q is missing metadata key %q",
			o.Name,
			metadataKey_SHA1Hex)
		return
	}

	if hexSHA1 != hexScore {
		err = fmt.Errorf(
			"Score/SHA-1 metadata mismatch for object %q: %q",
			o.Name,
			hexSHA1)
		return
	}

	// We expect the hex MD5 in the object metadata to align with what GCS says
	// the object's MD5 is.
	hexMD5, ok := o.Metadata[metadataKey_MD5Hex]
	if !ok {
		err = fmt.Errorf(
			"Object %q is missing metadata key %q",
			o.Name,
			metadataKey_MD5Hex)
		return
	}

	if hex.DecodedLen(len(hexMD5)) != md5.Size {
		err = fmt.Errorf(
			"Object %q has weird hex MD5 metadata: %q",
			o.Name,
			hexMD5)
		return
	}

	var md5 [md5.Size]byte
	_, err = hex.Decode(md5[:], []byte(hexMD5))
	if err != nil {
		err = fmt.Errorf(
			"Object %q has invalid hex MD5 in metadata: %q",
			o.Name,
			hexMD5)
		return
	}

	if md5 != o.MD5 {
		err = fmt.Errorf(
			"MD5 mismatch for object %q: %s vs. %s",
			o.Name,
			hex.EncodeToString(md5[:]),
			hex.EncodeToString(o.MD5[:]))
		return
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

func (s *GCSStore) makeName(score Score) (name string) {
	name = s.namePrefix + score.Hex()
	return
}

////////////////////////////////////////////////////////////////////////
// Public interface
////////////////////////////////////////////////////////////////////////

func (s *GCSStore) Store(blob []byte) (score Score, err error) {
	// Compute a score and an object name.
	score = ComputeScore(blob)
	name := s.makeName(score)

	// Create the object.
	crc32c := *gcsutil.CRC32C(blob)
	md5 := *gcsutil.MD5(blob)
	sha1 := sha1.Sum(blob)

	req := &gcs.CreateObjectRequest{
		Name:     name,
		Contents: bytes.NewReader(blob),
		CRC32C:   &crc32c,
		MD5:      &md5,

		Metadata: map[string]string{
			metadataKey_SHA1Hex: hex.EncodeToString(sha1[:]),
			metadataKey_MD5Hex:  hex.EncodeToString(md5[:]),
		},
	}

	o, err := s.bucket.CreateObject(context.Background(), req)
	if err != nil {
		err = fmt.Errorf("CreateObject: %v", err)
		return
	}

	// Paranoid check: what we get back from GCS should match what we put in.
	if o.CRC32C != crc32c {
		panic(fmt.Sprintf(
			"CRC32C mismatch for object %q: 0x%08xv vs. 0x%08x",
			o.Name,
			o.CRC32C,
			crc32c))
	}

	if o.MD5 != md5 {
		panic(fmt.Sprintf(
			"MD5 mismatch for object %q: %s vs. %s",
			o.Name,
			hex.EncodeToString(o.MD5[:]),
			hex.EncodeToString(md5[:])))
	}

	return
}

func (s *GCSStore) Flush() (err error) {
	panic("GCSStore.Flush not supported; wiring code bug?")
}

func (s *GCSStore) Contains(score Score) (b bool) {
	panic("GCSStore.Contains not supported; wiring code bug?")
}

func (s *GCSStore) Load(score Score) (blob []byte, err error) {
	// Create a ReadCloser.
	req := &gcs.ReadObjectRequest{
		Name: s.makeName(score),
	}

	rc, err := s.bucket.NewReader(context.Background(), req)
	if err != nil {
		err = fmt.Errorf("NewReader: %v", err)
		return
	}

	// Read from it.
	blob, err = ioutil.ReadAll(rc)
	if err != nil {
		rc.Close()
		err = fmt.Errorf("ReadAll: %v", err)
		return
	}

	// Close it.
	err = rc.Close()
	if err != nil {
		err = fmt.Errorf("Close: %v", err)
		return
	}

	return
}

// List all of the blobs that are known to be durable in the bucket.
func (s *GCSStore) List() (scores []Score, err error) {
	req := &gcs.ListObjectsRequest{
		Prefix: s.namePrefix,
	}

	// List repeatedly until we're done.
	for {
		// Call the bucket.
		var listing *gcs.Listing
		listing, err = s.bucket.ListObjects(context.Background(), req)
		if err != nil {
			err = fmt.Errorf("ListObjects: %v", err)
			return
		}

		// Process results.
		for _, o := range listing.Objects {
			var score Score
			score, err = ParseObjectRecord(o, s.namePrefix)
			if err != nil {
				return
			}

			scores = append(scores, score)
		}

		// Continue?
		if listing.ContinuationToken == "" {
			break
		}

		req.ContinuationToken = listing.ContinuationToken
	}

	return
}
<commit_msg>Improved an error message.<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 blob

import (
	"bytes"
	"crypto/md5"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"strings"

	"golang.org/x/net/context"

	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcsutil"
)

// A key placed in GCS object metadata by GCSStore containing the hex SHA-1
// expected for the object contents. This is of course redundant with the
// object name; we use it as a paranoid check against GCS returning the
// metadata or contents for the wrong object.
const metadataKey_SHA1Hex = "comeback_sha1_hex"

// A key placed in GCS object metadata by GCSStore containing the hex MD5 sum
// expected for the object contents. If GCS reports a different MD5 sum or
// returns contents with a different MD5 sum, we know something screwy has
// happened.
//
// See here for more info: https://github.com/jacobsa/comeback/issues/18
const metadataKey_MD5Hex = "comeback_md5_hex"

// Return a blob store that stores blobs in the supplied GCS bucket. GCS object
// names look like:
//
//     <prefix><score>
//
// where <score> is the result of calling Score.Hex.
//
// The blob store trusts that it has full ownership of this portion of the
// bucket's namespace -- if a score name exists, then it points to the correct
// data.
//
// The returned store does not support Flush or Contains; these methods must
// not be called.
func NewGCSStore(
	bucket gcs.Bucket,
	prefix string) (store *GCSStore) {
	store = &GCSStore{
		bucket:     bucket,
		namePrefix: prefix,
	}

	return
}

type GCSStore struct {
	bucket     gcs.Bucket
	namePrefix string
}

// Parse and verify the internal consistency of the supplied object record in
// the same manner that a GCSStore configured with the supplied object name
// prefix would. Return the score of the blob that the object contains.
func ParseObjectRecord(
	o *gcs.Object,
	namePrefix string) (score Score, err error) {
	// Is the name of the appropriate form?
	if !strings.HasPrefix(o.Name, namePrefix) {
		err = fmt.Errorf("Unexpected object name: %q", o.Name)
		return
	}

	// Parse the hex score.
	hexScore := strings.TrimPrefix(o.Name, namePrefix)
	score, err = ParseHexScore(hexScore)
	if err != nil {
		err = fmt.Errorf(
			"Unexpected hex score %q for object %q: %v",
			hexScore,
			o.Name,
			err)
		return
	}

	// We expect the hex score to match the hex SHA-1 in the metadata.
	hexSHA1, ok := o.Metadata[metadataKey_SHA1Hex]
	if !ok {
		err = fmt.Errorf(
			"Object %q is missing metadata key %q",
			o.Name,
			metadataKey_SHA1Hex)
		return
	}

	if hexSHA1 != hexScore {
		err = fmt.Errorf(
			"Score/SHA-1 metadata mismatch for object %q: %q",
			o.Name,
			hexSHA1)
		return
	}

	// We expect the hex MD5 in the object metadata to align with what GCS says
	// the object's MD5 is.
	hexMD5, ok := o.Metadata[metadataKey_MD5Hex]
	if !ok {
		err = fmt.Errorf(
			"Object %q is missing metadata key %q",
			o.Name,
			metadataKey_MD5Hex)
		return
	}

	if hex.DecodedLen(len(hexMD5)) != md5.Size {
		err = fmt.Errorf(
			"Object %q has weird hex MD5 metadata: %q",
			o.Name,
			hexMD5)
		return
	}

	var md5 [md5.Size]byte
	_, err = hex.Decode(md5[:], []byte(hexMD5))
	if err != nil {
		err = fmt.Errorf(
			"Object %q has invalid hex MD5 in metadata: %q",
			o.Name,
			hexMD5)
		return
	}

	if md5 != o.MD5 {
		err = fmt.Errorf(
			"MD5 mismatch for object %q: %s vs. %s",
			o.Name,
			hex.EncodeToString(md5[:]),
			hex.EncodeToString(o.MD5[:]))
		return
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

func (s *GCSStore) makeName(score Score) (name string) {
	name = s.namePrefix + score.Hex()
	return
}

////////////////////////////////////////////////////////////////////////
// Public interface
////////////////////////////////////////////////////////////////////////

func (s *GCSStore) Store(blob []byte) (score Score, err error) {
	// Compute a score and an object name.
	score = ComputeScore(blob)
	name := s.makeName(score)

	// Create the object.
	crc32c := *gcsutil.CRC32C(blob)
	md5 := *gcsutil.MD5(blob)
	sha1 := sha1.Sum(blob)

	req := &gcs.CreateObjectRequest{
		Name:     name,
		Contents: bytes.NewReader(blob),
		CRC32C:   &crc32c,
		MD5:      &md5,

		Metadata: map[string]string{
			metadataKey_SHA1Hex: hex.EncodeToString(sha1[:]),
			metadataKey_MD5Hex:  hex.EncodeToString(md5[:]),
		},
	}

	o, err := s.bucket.CreateObject(context.Background(), req)
	if err != nil {
		err = fmt.Errorf("CreateObject: %v", err)
		return
	}

	// Paranoid check: what we get back from GCS should match what we put in.
	if o.CRC32C != crc32c {
		panic(fmt.Sprintf(
			"CRC32C mismatch for object %q: 0x%08xv vs. 0x%08x",
			o.Name,
			o.CRC32C,
			crc32c))
	}

	if o.MD5 != md5 {
		panic(fmt.Sprintf(
			"MD5 mismatch for object %q: %s vs. %s",
			o.Name,
			hex.EncodeToString(o.MD5[:]),
			hex.EncodeToString(md5[:])))
	}

	return
}

func (s *GCSStore) Flush() (err error) {
	panic("GCSStore.Flush not supported; wiring code bug?")
}

func (s *GCSStore) Contains(score Score) (b bool) {
	panic("GCSStore.Contains not supported; wiring code bug?")
}

func (s *GCSStore) Load(score Score) (blob []byte, err error) {
	// Create a ReadCloser.
	req := &gcs.ReadObjectRequest{
		Name: s.makeName(score),
	}

	rc, err := s.bucket.NewReader(context.Background(), req)
	if err != nil {
		err = fmt.Errorf("NewReader: %v", err)
		return
	}

	// Read from it.
	blob, err = ioutil.ReadAll(rc)
	if err != nil {
		rc.Close()
		err = fmt.Errorf("ReadAll: %v", err)
		return
	}

	// Close it.
	err = rc.Close()
	if err != nil {
		err = fmt.Errorf("Close: %v", err)
		return
	}

	return
}

// List all of the blobs that are known to be durable in the bucket.
func (s *GCSStore) List() (scores []Score, err error) {
	req := &gcs.ListObjectsRequest{
		Prefix: s.namePrefix,
	}

	// List repeatedly until we're done.
	for {
		// Call the bucket.
		var listing *gcs.Listing
		listing, err = s.bucket.ListObjects(context.Background(), req)
		if err != nil {
			err = fmt.Errorf("ListObjects: %v", err)
			return
		}

		// Process results.
		for _, o := range listing.Objects {
			var score Score
			score, err = ParseObjectRecord(o, s.namePrefix)
			if err != nil {
				return
			}

			scores = append(scores, score)
		}

		// Continue?
		if listing.ContinuationToken == "" {
			break
		}

		req.ContinuationToken = listing.ContinuationToken
	}

	return
}
<|endoftext|>


<commit_before>package imageregistry

import (
	"context"
	"crypto/tls"
	"net"
	"net/http"
	"time"

	"github.com/onsi/ginkgo"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes/scheme"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/events"
	"k8s.io/kubernetes/test/e2e/framework"
	"k8s.io/kubernetes/test/e2e/upgrades"

	"github.com/openshift/origin/pkg/monitor"
	"github.com/openshift/origin/test/extended/util/disruption"
)

// AvailableTest tests that the image registry is available before, during, and
// after a cluster upgrade.
type AvailableTest struct {
}

func (AvailableTest) Name() string { return "image-registry-available" }
func (AvailableTest) DisplayName() string {
	return "[sig-imageregistry] Image registry remain available"
}

// Setup does nothing.
func (t *AvailableTest) Setup(f *framework.Framework) {
}

// Test runs a connectivity check to the service.
func (t *AvailableTest) Test(f *framework.Framework, done <-chan struct{}, upgrade upgrades.UpgradeType) {
	client, err := framework.LoadClientset()
	framework.ExpectNoError(err)

	stopCh := make(chan struct{})
	defer close(stopCh)

	newBroadcaster := events.NewBroadcaster(&events.EventSinkImpl{Interface: client.EventsV1()})
	r := newBroadcaster.NewRecorder(scheme.Scheme, "openshift.io/image-registry-available-test")
	newBroadcaster.StartRecordingToSink(stopCh)

	ginkgo.By("continuously hitting image registry")

	ctx, cancel := context.WithCancel(context.Background())
	m := monitor.NewMonitorWithInterval(1 * time.Second)
	err = startEndpointMonitoring(ctx, m, r)
	framework.ExpectNoError(err, "unable to monitor route")

	start := time.Now()
	m.StartSampling(ctx)

	// wait to ensure API is still up after the test ends
	<-done
	ginkgo.By("waiting for any post disruption failures")
	time.Sleep(15 * time.Second)
	cancel()
	end := time.Now()

	disruption.ExpectNoDisruption(f, 0.20, end.Sub(start), m.Events(time.Time{}, time.Time{}), "Image registry was unreachable during disruption")
}

// Teardown cleans up any remaining resources.
func (t *AvailableTest) Teardown(f *framework.Framework) {
	// rely on the namespace deletion to clean up everything
}

func startEndpointMonitoring(ctx context.Context, m *monitor.Monitor, r events.EventRecorder) error {
	const (
		url     = "https://image-registry.openshift-image-registry.svc:5000"
		path    = "/healthz"
		locator = "image-registry"
	)

	// this client reuses connections and detects abrupt breaks
	continuousClient, err := rest.UnversionedRESTClientFor(&rest.Config{
		Host:    url,
		Timeout: 10 * time.Second,
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: 15 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 15 * time.Second,
			TLSClientConfig:     &tls.Config{InsecureSkipVerify: true},
		},
		ContentConfig: rest.ContentConfig{
			NegotiatedSerializer: runtime.NewSimpleNegotiatedSerializer(scheme.Codecs.SupportedMediaTypes()[0]),
		},
	})
	if err != nil {
		return err
	}
	m.AddSampler(
		monitor.StartSampling(ctx, m, time.Second, func(previous bool) (condition *monitor.Condition, next bool) {
			_, err := continuousClient.Get().AbsPath(path).DoRaw(ctx)
			switch {
			case err == nil && !previous:
				condition = &monitor.Condition{
					Level:   monitor.Info,
					Locator: locator,
					Message: "Service started responding to GET requests on reused connections",
				}
			case err != nil && previous:
				framework.Logf("Service image-registry is unreachable on reused connections: %v", err)
				r.Eventf(&v1.ObjectReference{Kind: "Service", Namespace: "openshift-image-registry", Name: "image-registry"}, nil, v1.EventTypeWarning, "Unreachable", "detected", "on reused connections")
				condition = &monitor.Condition{
					Level:   monitor.Error,
					Locator: locator,
					Message: "Service stopped responding to GET requests on reused connections",
				}
			case err != nil:
				framework.Logf("Service image-registry is unreachable on reused connections: %v", err)
			}
			return condition, err == nil
		}).ConditionWhenFailing(&monitor.Condition{
			Level:   monitor.Error,
			Locator: locator,
			Message: "Service is not responding to GET requests on reused connections",
		}),
	)

	// this client creates fresh connections and detects failure to establish connections
	client, err := rest.UnversionedRESTClientFor(&rest.Config{
		Host:    url,
		Timeout: 10 * time.Second,
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout:   15 * time.Second,
				KeepAlive: -1,
			}).Dial,
			TLSHandshakeTimeout: 15 * time.Second,
			IdleConnTimeout:     15 * time.Second,
			DisableKeepAlives:   true,
			TLSClientConfig:     &tls.Config{InsecureSkipVerify: true},
		},
		ContentConfig: rest.ContentConfig{
			NegotiatedSerializer: runtime.NewSimpleNegotiatedSerializer(scheme.Codecs.SupportedMediaTypes()[0]),
		},
	})
	if err != nil {
		return err
	}
	m.AddSampler(
		monitor.StartSampling(ctx, m, time.Second, func(previous bool) (condition *monitor.Condition, next bool) {
			_, err := client.Get().AbsPath(path).DoRaw(ctx)
			switch {
			case err == nil && !previous:
				condition = &monitor.Condition{
					Level:   monitor.Info,
					Locator: locator,
					Message: "Service started responding to GET requests over new connections",
				}
			case err != nil && previous:
				framework.Logf("Service image-registry is unreachable on new connections: %v", err)
				r.Eventf(&v1.ObjectReference{Kind: "Service", Namespace: "openshift-image-registry", Name: "image-registry"}, nil, v1.EventTypeWarning, "Unreachable", "detected", "on new connections")
				condition = &monitor.Condition{
					Level:   monitor.Error,
					Locator: locator,
					Message: "Service stopped responding to GET requests over new connections",
				}
			case err != nil:
				framework.Logf("Service image-registry is unreachable on new connections: %v", err)
			}
			return condition, err == nil
		}).ConditionWhenFailing(&monitor.Condition{
			Level:   monitor.Error,
			Locator: locator,
			Message: "Service is not responding to GET requests over new connections",
		}),
	)

	return nil
}
<commit_msg>Monitor availability of registry through router<commit_after>package imageregistry

import (
	"context"
	"crypto/tls"
	"net"
	"net/http"
	"time"

	"github.com/onsi/ginkgo"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes/scheme"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/events"
	"k8s.io/kubernetes/test/e2e/framework"
	"k8s.io/kubernetes/test/e2e/upgrades"

	routev1 "github.com/openshift/api/route/v1"
	routeclient "github.com/openshift/client-go/route/clientset/versioned"
	"github.com/openshift/origin/pkg/monitor"
	"github.com/openshift/origin/test/extended/util/disruption"
)

// AvailableTest tests that the image registry is available before, during, and
// after a cluster upgrade.
type AvailableTest struct {
	routeRef *corev1.ObjectReference
	host     string
}

func (AvailableTest) Name() string { return "image-registry-available" }
func (AvailableTest) DisplayName() string {
	return "[sig-imageregistry] Image registry remain available"
}

// Setup creates a route that exposes the registry to tests.
func (t *AvailableTest) Setup(f *framework.Framework) {
	ctx := context.Background()

	config, err := framework.LoadConfig()
	framework.ExpectNoError(err)
	routeClient, err := routeclient.NewForConfig(config)
	framework.ExpectNoError(err)

	route, err := routeClient.RouteV1().Routes("openshift-image-registry").Create(ctx, &routev1.Route{
		ObjectMeta: metav1.ObjectMeta{
			Name: "test-disruption",
		},
		Spec: routev1.RouteSpec{
			To: routev1.RouteTargetReference{
				Kind: "Service",
				Name: "image-registry",
			},
			Port: &routev1.RoutePort{
				TargetPort: intstr.FromInt(5000),
			},
			TLS: &routev1.TLSConfig{
				Termination:                   routev1.TLSTerminationPassthrough,
				InsecureEdgeTerminationPolicy: routev1.InsecureEdgeTerminationPolicyRedirect,
			},
		},
	}, metav1.CreateOptions{})
	framework.ExpectNoError(err)

	err = wait.PollImmediate(1*time.Second, 30*time.Second, func() (bool, error) {
		route, err = routeClient.RouteV1().Routes("openshift-image-registry").Get(ctx, "test-disruption", metav1.GetOptions{})
		if err != nil {
			return false, err
		}

		for _, ingress := range route.Status.Ingress {
			if len(ingress.Host) > 0 {
				t.host = ingress.Host
				return true, nil
			}
		}
		return false, nil
	})
	framework.ExpectNoError(err, "failed to get route host")

	t.routeRef = &corev1.ObjectReference{
		Kind:      "Route",
		Namespace: route.Namespace,
		Name:      route.Name,
	}
}

// Test runs a connectivity check to the service.
func (t *AvailableTest) Test(f *framework.Framework, done <-chan struct{}, upgrade upgrades.UpgradeType) {
	client, err := framework.LoadClientset()
	framework.ExpectNoError(err)

	stopCh := make(chan struct{})
	defer close(stopCh)

	newBroadcaster := events.NewBroadcaster(&events.EventSinkImpl{Interface: client.EventsV1()})
	r := newBroadcaster.NewRecorder(scheme.Scheme, "openshift.io/image-registry-available-test")
	newBroadcaster.StartRecordingToSink(stopCh)

	ginkgo.By("continuously hitting image registry")

	ctx, cancel := context.WithCancel(context.Background())
	m := monitor.NewMonitorWithInterval(1 * time.Second)
	err = t.startEndpointMonitoring(ctx, m, r)
	framework.ExpectNoError(err, "unable to monitor route")

	start := time.Now()
	m.StartSampling(ctx)

	// wait to ensure API is still up after the test ends
	<-done
	ginkgo.By("waiting for any post disruption failures")
	time.Sleep(15 * time.Second)
	cancel()
	end := time.Now()

	disruption.ExpectNoDisruption(f, 0.20, end.Sub(start), m.Events(time.Time{}, time.Time{}), "Image registry was unreachable during disruption")
}

// Teardown cleans up any remaining resources.
func (t *AvailableTest) Teardown(f *framework.Framework) {
	ctx := context.Background()

	config, err := framework.LoadConfig()
	framework.ExpectNoError(err)
	routeClient, err := routeclient.NewForConfig(config)
	framework.ExpectNoError(err)

	err = routeClient.RouteV1().Routes(t.routeRef.Namespace).Delete(ctx, t.routeRef.Name, metav1.DeleteOptions{})
	framework.ExpectNoError(err, "failed to delete route")
}

func (t *AvailableTest) startEndpointMonitoring(ctx context.Context, m *monitor.Monitor, r events.EventRecorder) error {
	var (
		url     = "https://" + t.host
		path    = "/healthz"
		locator = "image-registry"
	)

	// this client reuses connections and detects abrupt breaks
	continuousClient, err := rest.UnversionedRESTClientFor(&rest.Config{
		Host:    url,
		Timeout: 10 * time.Second,
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: 15 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 15 * time.Second,
			TLSClientConfig:     &tls.Config{InsecureSkipVerify: true},
		},
		ContentConfig: rest.ContentConfig{
			NegotiatedSerializer: runtime.NewSimpleNegotiatedSerializer(scheme.Codecs.SupportedMediaTypes()[0]),
		},
	})
	if err != nil {
		return err
	}
	m.AddSampler(
		monitor.StartSampling(ctx, m, time.Second, func(previous bool) (condition *monitor.Condition, next bool) {
			_, err := continuousClient.Get().AbsPath(path).DoRaw(ctx)
			switch {
			case err == nil && !previous:
				condition = &monitor.Condition{
					Level:   monitor.Info,
					Locator: locator,
					Message: "Route started responding to GET requests on reused connections",
				}
			case err != nil && previous:
				framework.Logf("Route for image-registry is unreachable on reused connections: %v", err)
				r.Eventf(t.routeRef, nil, corev1.EventTypeWarning, "Unreachable", "detected", "on reused connections")
				condition = &monitor.Condition{
					Level:   monitor.Error,
					Locator: locator,
					Message: "Route stopped responding to GET requests on reused connections",
				}
			case err != nil:
				framework.Logf("Route for image-registry is unreachable on reused connections: %v", err)
			}
			return condition, err == nil
		}).ConditionWhenFailing(&monitor.Condition{
			Level:   monitor.Error,
			Locator: locator,
			Message: "Route is not responding to GET requests on reused connections",
		}),
	)

	// this client creates fresh connections and detects failure to establish connections
	client, err := rest.UnversionedRESTClientFor(&rest.Config{
		Host:    url,
		Timeout: 10 * time.Second,
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout:   15 * time.Second,
				KeepAlive: -1,
			}).Dial,
			TLSHandshakeTimeout: 15 * time.Second,
			IdleConnTimeout:     15 * time.Second,
			DisableKeepAlives:   true,
			TLSClientConfig:     &tls.Config{InsecureSkipVerify: true},
		},
		ContentConfig: rest.ContentConfig{
			NegotiatedSerializer: runtime.NewSimpleNegotiatedSerializer(scheme.Codecs.SupportedMediaTypes()[0]),
		},
	})
	if err != nil {
		return err
	}
	m.AddSampler(
		monitor.StartSampling(ctx, m, time.Second, func(previous bool) (condition *monitor.Condition, next bool) {
			_, err := client.Get().AbsPath(path).DoRaw(ctx)
			switch {
			case err == nil && !previous:
				condition = &monitor.Condition{
					Level:   monitor.Info,
					Locator: locator,
					Message: "Route started responding to GET requests over new connections",
				}
			case err != nil && previous:
				framework.Logf("Route for image-registry is unreachable on new connections: %v", err)
				r.Eventf(t.routeRef, nil, corev1.EventTypeWarning, "Unreachable", "detected", "on new connections")
				condition = &monitor.Condition{
					Level:   monitor.Error,
					Locator: locator,
					Message: "Route stopped responding to GET requests over new connections",
				}
			case err != nil:
				framework.Logf("Route for image-registry is unreachable on new connections: %v", err)
			}
			return condition, err == nil
		}).ConditionWhenFailing(&monitor.Condition{
			Level:   monitor.Error,
			Locator: locator,
			Message: "Route is not responding to GET requests over new connections",
		}),
	)

	return nil
}
<|endoftext|>


<commit_before>//go:build linux && cgo && !agent

package state

import (
	"context"
	"net/http"
	"net/url"

	"github.com/lxc/lxd/lxd/bgp"
	clusterConfig "github.com/lxc/lxd/lxd/cluster/config"
	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/dns"
	"github.com/lxc/lxd/lxd/endpoints"
	"github.com/lxc/lxd/lxd/events"
	"github.com/lxc/lxd/lxd/firewall"
	"github.com/lxc/lxd/lxd/fsmonitor"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/maas"
	"github.com/lxc/lxd/lxd/sys"
	"github.com/lxc/lxd/shared"
)

// State is a gateway to the two main stateful components of LXD, the database
// and the operating system. It's typically used by model entities such as
// containers, volumes, etc. in order to perform changes.
type State struct {
	// Shutdown Context
	ShutdownCtx context.Context

	// Databases
	DB *db.DB

	// MAAS server
	MAAS *maas.Controller

	// BGP server
	BGP *bgp.Server

	// DNS server
	DNS *dns.Server

	// OS access
	OS    *sys.OS
	Proxy func(req *http.Request) (*url.URL, error)

	// LXD server
	Endpoints *endpoints.Endpoints

	// Event server
	DevlxdEvents *events.DevLXDServer
	Events       *events.Server

	// Firewall instance
	Firewall firewall.Firewall

	// Server certificate
	ServerCert             func() *shared.CertInfo
	UpdateCertificateCache func()

	// Available instance types based on operational drivers.
	InstanceTypes map[instancetype.Type]error

	// Filesystem monitor
	DevMonitor fsmonitor.FSMonitor

	// Global configuration
	GlobalConfig *clusterConfig.Config

	// Local server name.
	ServerName string
}
<commit_msg>lxd/state/state: Add LocalConfig to State<commit_after>//go:build linux && cgo && !agent

package state

import (
	"context"
	"net/http"
	"net/url"

	"github.com/lxc/lxd/lxd/bgp"
	clusterConfig "github.com/lxc/lxd/lxd/cluster/config"
	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/dns"
	"github.com/lxc/lxd/lxd/endpoints"
	"github.com/lxc/lxd/lxd/events"
	"github.com/lxc/lxd/lxd/firewall"
	"github.com/lxc/lxd/lxd/fsmonitor"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/maas"
	"github.com/lxc/lxd/lxd/node"
	"github.com/lxc/lxd/lxd/sys"
	"github.com/lxc/lxd/shared"
)

// State is a gateway to the two main stateful components of LXD, the database
// and the operating system. It's typically used by model entities such as
// containers, volumes, etc. in order to perform changes.
type State struct {
	// Shutdown Context
	ShutdownCtx context.Context

	// Databases
	DB *db.DB

	// MAAS server
	MAAS *maas.Controller

	// BGP server
	BGP *bgp.Server

	// DNS server
	DNS *dns.Server

	// OS access
	OS    *sys.OS
	Proxy func(req *http.Request) (*url.URL, error)

	// LXD server
	Endpoints *endpoints.Endpoints

	// Event server
	DevlxdEvents *events.DevLXDServer
	Events       *events.Server

	// Firewall instance
	Firewall firewall.Firewall

	// Server certificate
	ServerCert             func() *shared.CertInfo
	UpdateCertificateCache func()

	// Available instance types based on operational drivers.
	InstanceTypes map[instancetype.Type]error

	// Filesystem monitor
	DevMonitor fsmonitor.FSMonitor

	// Global configuration
	GlobalConfig *clusterConfig.Config

	// Local configuration
	LocalConfig *node.Config

	// Local server name.
	ServerName string
}
<|endoftext|>


<commit_before>package lzma

import (
	"errors"
	"io"
)

// bwriter implements a ByteWriter on top of a Writer
type bwriter struct {
	w   io.Writer
	buf []byte
}

// WriteByte writes a byte using the standard writer. Note that the
// function is not thread-safe.
func (bw *bwriter) WriteByte(c byte) error {
	bw.buf[0] = c
	_, err := bw.w.Write(bw.buf)
	return err
}

// ByteWriterFromWriter converts a Writer to a ByteWriter. If the Writer
// doesn't support the ByteWriter interface directly a new structure
// will be allocated and a non-thread-safe ByteWriter will be returned.
func ByteWriterFromWriter(w io.Writer) io.ByteWriter {
	if bw, ok := w.(io.ByteWriter); ok {
		return bw
	}
	return &bwriter{w, make([]byte, 1)}
}

// ErrLimit indicates that the limit has been reached.
var ErrLimit = errors.New("limit reached")

// LimitedByteWriter provides a byte writer that can be written until a
// limit is reached. The field N provides the number of remaining
// bytes.
type LimitedByteWriter struct {
	BW io.ByteWriter
	N  int64
}

// WriteByte writes a single byte to the limited byte writer. It returns
// ErrLimit if the limit has been reached. If the byte is successfully
// written the field N of the LimitedByteWriter will be decremented by
// one.
func (l *LimitedByteWriter) WriteByte(c byte) error {
	if l.N <= 0 {
		return ErrLimit
	}
	if err := l.BW.WriteByte(c); err != nil {
		return err
	}
	l.N--
	return nil
}
<commit_msg>lzma: remove ByteWriterFromWriter<commit_after>package lzma

import (
	"errors"
	"io"
)

// ErrLimit indicates that the limit has been reached.
var ErrLimit = errors.New("limit reached")

// LimitedByteWriter provides a byte writer that can be written until a
// limit is reached. The field N provides the number of remaining
// bytes.
type LimitedByteWriter struct {
	BW io.ByteWriter
	N  int64
}

// WriteByte writes a single byte to the limited byte writer. It returns
// ErrLimit if the limit has been reached. If the byte is successfully
// written the field N of the LimitedByteWriter will be decremented by
// one.
func (l *LimitedByteWriter) WriteByte(c byte) error {
	if l.N <= 0 {
		return ErrLimit
	}
	if err := l.BW.WriteByte(c); err != nil {
		return err
	}
	l.N--
	return nil
}
<|endoftext|>