<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package syslog

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"
	"text/template"
)

// tagOffset represents the substring start value for the tag to return
// the logfileName value from the syslogtag.  Substrings in syslog are
// indexed from 1, hence the + 1.
const tagOffset = len("juju-") + 1

// The rsyslog conf for state server nodes.
// Messages are gathered from other nodes and accumulated in an all-machines.log file.
//
// The apparmor profile is quite strict about where rsyslog can write files.
// Instead of poking with the profile, the local provider now logs to
// {{logDir}}-{{user}}-{{env name}}/all-machines.log, and a symlink is made
// in the local provider log dir to point to that file. The file is also
// created with 0644 so the user can read it without poking permissions. By
// default rsyslog creates files with 0644, but in the ubuntu package, the
// setting is changed to 0640, which means normal users can't read the log
// file. Using a new action directive (new as in not-legacy), we can specify
// the file create mode so it doesn't use the default.
//
// I would dearly love to write the filtering action as follows to avoid setting
// and resetting the global $FileCreateMode, but alas, precise doesn't support it
//
// if $syslogtag startswith "juju{{namespace}}-" then
//   action(type="omfile"
//          File="{{logDir}}{{namespace}}/all-machines.log"
//          Template="JujuLogFormat{{namespace}}"
//          FileCreateMode="0644")
// & stop
//
// Instead we need to mess with the global FileCreateMode.  We set it back
// to the ubuntu default after defining our rule.
const stateServerRsyslogTemplate = `
$ModLoad imfile

$InputFilePersistStateInterval 50
$InputFilePollInterval 5
$InputFileName {{logfilePath}}
$InputFileTag juju{{namespace}}-{{logfileName}}:
$InputFileStateFile {{logfileName}}{{namespace}}
$InputRunFileMonitor

$ModLoad imtcp
$DefaultNetstreamDriver gtls
$DefaultNetstreamDriverCAFile {{tlsCACertPath}}
$DefaultNetstreamDriverCertFile {{tlsCertPath}}
$DefaultNetstreamDriverKeyFile {{tlsKeyPath}}
$InputTCPServerStreamDriverAuthMode anon
$InputTCPServerStreamDriverMode 1 # run driver in TLS-only mode
$InputTCPServerRun {{portNumber}}

# Messages received from remote rsyslog machines have messages prefixed with a space,
# so add one in for local messages too if needed.
$template JujuLogFormat{{namespace}},"%syslogtag:{{tagStart}}:$%%msg:::sp-if-no-1st-sp%%msg:::drop-last-lf%\n"

$FileCreateMode 0644
:syslogtag, startswith, "juju{{namespace}}-" {{logDir}}/all-machines.log;JujuLogFormat{{namespace}}
& ~
$FileCreateMode 0640
`

// The rsyslog conf for non-state server nodes.
// Messages are forwarded to the state server node.
const nodeRsyslogTemplate = `
$ModLoad imfile
{{range $i, $bootstrapIP := bootstrapHosts}}{{if $i}}\n{{end}}
# Enable reliable forwarding.
$ActionQueueType LinkedList
$ActionQueueFileName {{logfileName}}{{namespace}}
$ActionResumeRetryCount -1
$ActionQueueSaveOnShutdown on

$InputFilePersistStateInterval 50
$InputFilePollInterval 5
$InputFileName {{logfilePath}}
$InputFileTag juju{{namespace}}-{{logfileName}}:
$InputFileStateFile {{logfileName}}{{namespace}}
$InputRunFileMonitor

$DefaultNetstreamDriver gtls
$DefaultNetstreamDriverCAFile {{tlsCACertPath}}
$ActionSendStreamDriverAuthMode anon
$ActionSendStreamDriverMode 1 # run driver in TLS-only mode

$template LongTagForwardFormat,"<%PRI%>%TIMESTAMP:::date-rfc3339% %HOSTNAME% %syslogtag%%msg:::sp-if-no-1st-sp%%msg%"

:syslogtag, startswith, "juju{{namespace}}-" @@{{$bootstrapIP}}:{{portNumber}};LongTagForwardFormat{{end}}
& ~
`

// nodeRsyslogTemplateTLSHeader is prepended to
// nodeRsyslogTemplate if TLS is to be used.
const nodeRsyslogTemplateTLSHeader = `
`

const (
	defaultConfigDir          = "/etc/rsyslog.d"
	defaultCACertFileName     = "ca-cert.pem"
	defaultServerCertFileName = "rsyslog-cert.pem"
	defaultServerKeyFileName  = "rsyslog-key.pem"
)

// SyslogConfigRenderer instances are used to generate a rsyslog conf file.
type SyslogConfigRenderer interface {
	Render() ([]byte, error)
}

// SyslogConfig provides a means to configure and generate rsyslog conf files for
// the state server nodes and unit nodes.
// rsyslog is configured to tail the specified log file.
type SyslogConfig struct {
	// the template representing the config file contents.
	configTemplate string
	// the directory where the config file is written.
	ConfigDir string
	// the config file name.
	ConfigFileName string
	// the name of the log file to tail.
	LogFileName string
	// the addresses of the state server to which messages should be forwarded.
	StateServerAddresses []string
	// CA certificate file name.
	CACertFileName string
	// Server certificate file name.
	ServerCertFileName string
	// Server private key file name.
	ServerKeyFileName string
	// the port number for the listener
	Port int
	// the directory for the logfiles
	LogDir string
	// namespace is used when there are multiple environments on one machine
	Namespace string
}

// NewForwardConfig creates a SyslogConfig instance used on unit nodes to forward log entries
// to the state server nodes.
func NewForwardConfig(logFile, logDir string, port int, namespace string, stateServerAddresses []string) *SyslogConfig {
	conf := &SyslogConfig{
		configTemplate:       nodeRsyslogTemplate,
		StateServerAddresses: stateServerAddresses,
		LogFileName:          logFile,
		Port:                 port,
		LogDir:               logDir,
	}
	if namespace != "" {
		conf.Namespace = "-" + namespace
	}
	return conf
}

// NewAccumulateConfig creates a SyslogConfig instance used to accumulate log entries from the
// various unit nodes.
func NewAccumulateConfig(logFile, logDir string, port int, namespace string) *SyslogConfig {
	conf := &SyslogConfig{
		configTemplate: stateServerRsyslogTemplate,
		LogFileName:    logFile,
		Port:           port,
		LogDir:         logDir,
	}
	if namespace != "" {
		conf.Namespace = "-" + namespace
	}
	return conf
}

func either(a, b string) string {
	if a != "" {
		return a
	}
	return b
}

func (slConfig *SyslogConfig) ConfigFilePath() string {
	dir := either(slConfig.ConfigDir, defaultConfigDir)
	return filepath.Join(dir, slConfig.ConfigFileName)
}

func (slConfig *SyslogConfig) CACertPath() string {
	filename := either(slConfig.CACertFileName, defaultCACertFileName)
	return filepath.Join(slConfig.LogDir, filename)
}

func (slConfig *SyslogConfig) ServerCertPath() string {
	filename := either(slConfig.ServerCertFileName, defaultServerCertFileName)
	return filepath.Join(slConfig.LogDir, filename)
}

func (slConfig *SyslogConfig) ServerKeyPath() string {
	filename := either(slConfig.ServerCertFileName, defaultServerKeyFileName)
	return filepath.Join(slConfig.LogDir, filename)
}

// Render generates the rsyslog config.
func (slConfig *SyslogConfig) Render() ([]byte, error) {
	var bootstrapHosts = func() []string {
		var hosts []string
		for _, addr := range slConfig.StateServerAddresses {
			parts := strings.Split(addr, ":")
			hosts = append(hosts, parts[0])
		}
		return hosts
	}

	var logFilePath = func() string {
		return fmt.Sprintf("%s/%s.log", slConfig.LogDir, slConfig.LogFileName)
	}

	t := template.New("")
	t.Funcs(template.FuncMap{
		"logfileName":    func() string { return slConfig.LogFileName },
		"bootstrapHosts": bootstrapHosts,
		"logfilePath":    logFilePath,
		"portNumber":     func() int { return slConfig.Port },
		"logDir":         func() string { return slConfig.LogDir },
		"namespace":      func() string { return slConfig.Namespace },
		"tagStart":       func() int { return tagOffset + len(slConfig.Namespace) },
		"tlsCACertPath":  slConfig.CACertPath,
		"tlsCertPath":    slConfig.ServerCertPath,
		"tlsKeyPath":     slConfig.ServerKeyPath,
	})

	// Process the rsyslog config template and echo to the conf file.
	p, err := t.Parse(slConfig.configTemplate)
	if err != nil {
		return nil, err
	}
	var confBuf bytes.Buffer
	if err := p.Execute(&confBuf, nil); err != nil {
		return nil, err
	}
	return confBuf.Bytes(), nil
}

// Write generates and writes the rsyslog config.
func (slConfig *SyslogConfig) Write() error {
	data, err := slConfig.Render()
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(slConfig.ConfigFilePath(), data, 0644)
	return err
}
<commit_msg>utils/syslog: rsyslog config for nodes<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package syslog

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"
	"text/template"
)

// tagOffset represents the substring start value for the tag to return
// the logfileName value from the syslogtag.  Substrings in syslog are
// indexed from 1, hence the + 1.
const tagOffset = len("juju-") + 1

// The rsyslog conf for state server nodes.
// Messages are gathered from other nodes and accumulated in an all-machines.log file.
//
// The apparmor profile is quite strict about where rsyslog can write files.
// Instead of poking with the profile, the local provider now logs to
// {{logDir}}-{{user}}-{{env name}}/all-machines.log, and a symlink is made
// in the local provider log dir to point to that file. The file is also
// created with 0644 so the user can read it without poking permissions. By
// default rsyslog creates files with 0644, but in the ubuntu package, the
// setting is changed to 0640, which means normal users can't read the log
// file. Using a new action directive (new as in not-legacy), we can specify
// the file create mode so it doesn't use the default.
//
// I would dearly love to write the filtering action as follows to avoid setting
// and resetting the global $FileCreateMode, but alas, precise doesn't support it
//
// if $syslogtag startswith "juju{{namespace}}-" then
//   action(type="omfile"
//          File="{{logDir}}{{namespace}}/all-machines.log"
//          Template="JujuLogFormat{{namespace}}"
//          FileCreateMode="0644")
// & stop
//
// Instead we need to mess with the global FileCreateMode.  We set it back
// to the ubuntu default after defining our rule.
const stateServerRsyslogTemplate = `
$ModLoad imfile

$InputFilePersistStateInterval 50
$InputFilePollInterval 5
$InputFileName {{logfilePath}}
$InputFileTag juju{{namespace}}-{{logfileName}}:
$InputFileStateFile {{logfileName}}{{namespace}}
$InputRunFileMonitor

$ModLoad imtcp
$DefaultNetstreamDriver gtls
$DefaultNetstreamDriverCAFile {{tlsCACertPath}}
$DefaultNetstreamDriverCertFile {{tlsCertPath}}
$DefaultNetstreamDriverKeyFile {{tlsKeyPath}}
$InputTCPServerStreamDriverAuthMode anon
$InputTCPServerStreamDriverMode 1 # run driver in TLS-only mode
$InputTCPServerRun {{portNumber}}

# Messages received from remote rsyslog machines have messages prefixed with a space,
# so add one in for local messages too if needed.
$template JujuLogFormat{{namespace}},"%syslogtag:{{tagStart}}:$%%msg:::sp-if-no-1st-sp%%msg:::drop-last-lf%\n"

$FileCreateMode 0644
:syslogtag, startswith, "juju{{namespace}}-" {{logDir}}/all-machines.log;JujuLogFormat{{namespace}}
& ~
$FileCreateMode 0640
`

// The rsyslog conf for non-state server nodes.
// Messages are forwarded to the state server node.
const nodeRsyslogTemplate = `
$ModLoad imfile
{{range $i, $bootstrapIP := bootstrapHosts}}{{if $i}}
\n
{{end}}
# Enable reliable forwarding.
$ActionQueueType LinkedList
$ActionQueueFileName {{logfileName}}{{namespace}}
$ActionResumeRetryCount -1
$ActionQueueSaveOnShutdown on

$InputFilePersistStateInterval 50
$InputFilePollInterval 5
$InputFileName {{logfilePath}}
$InputFileTag juju{{namespace}}-{{logfileName}}:
$InputFileStateFile {{logfileName}}{{namespace}}
$InputRunFileMonitor

$DefaultNetstreamDriver gtls
$DefaultNetstreamDriverCAFile {{tlsCACertPath}}
$ActionSendStreamDriverAuthMode anon
$ActionSendStreamDriverMode 1 # run driver in TLS-only mode

$template LongTagForwardFormat,"<%PRI%>%TIMESTAMP:::date-rfc3339% %HOSTNAME% %syslogtag%%msg:::sp-if-no-1st-sp%%msg%"

:syslogtag, startswith, "juju{{namespace}}-" @@{{$bootstrapIP}}:{{portNumber}};LongTagForwardFormat{{end}}
& ~
`

// nodeRsyslogTemplateTLSHeader is prepended to
// nodeRsyslogTemplate if TLS is to be used.
const nodeRsyslogTemplateTLSHeader = `
`

const (
	defaultConfigDir          = "/etc/rsyslog.d"
	defaultCACertFileName     = "ca-cert.pem"
	defaultServerCertFileName = "rsyslog-cert.pem"
	defaultServerKeyFileName  = "rsyslog-key.pem"
)

// SyslogConfigRenderer instances are used to generate a rsyslog conf file.
type SyslogConfigRenderer interface {
	Render() ([]byte, error)
}

// SyslogConfig provides a means to configure and generate rsyslog conf files for
// the state server nodes and unit nodes.
// rsyslog is configured to tail the specified log file.
type SyslogConfig struct {
	// the template representing the config file contents.
	configTemplate string
	// the directory where the config file is written.
	ConfigDir string
	// the config file name.
	ConfigFileName string
	// the name of the log file to tail.
	LogFileName string
	// the addresses of the state server to which messages should be forwarded.
	StateServerAddresses []string
	// CA certificate file name.
	CACertFileName string
	// Server certificate file name.
	ServerCertFileName string
	// Server private key file name.
	ServerKeyFileName string
	// the port number for the listener
	Port int
	// the directory for the logfiles
	LogDir string
	// namespace is used when there are multiple environments on one machine
	Namespace string
}

// NewForwardConfig creates a SyslogConfig instance used on unit nodes to forward log entries
// to the state server nodes.
func NewForwardConfig(logFile, logDir string, port int, namespace string, stateServerAddresses []string) *SyslogConfig {
	conf := &SyslogConfig{
		configTemplate:       nodeRsyslogTemplate,
		StateServerAddresses: stateServerAddresses,
		LogFileName:          logFile,
		Port:                 port,
		LogDir:               logDir,
	}
	if namespace != "" {
		conf.Namespace = "-" + namespace
	}
	return conf
}

// NewAccumulateConfig creates a SyslogConfig instance used to accumulate log entries from the
// various unit nodes.
func NewAccumulateConfig(logFile, logDir string, port int, namespace string) *SyslogConfig {
	conf := &SyslogConfig{
		configTemplate: stateServerRsyslogTemplate,
		LogFileName:    logFile,
		Port:           port,
		LogDir:         logDir,
	}
	if namespace != "" {
		conf.Namespace = "-" + namespace
	}
	return conf
}

func either(a, b string) string {
	if a != "" {
		return a
	}
	return b
}

func (slConfig *SyslogConfig) ConfigFilePath() string {
	dir := either(slConfig.ConfigDir, defaultConfigDir)
	return filepath.Join(dir, slConfig.ConfigFileName)
}

func (slConfig *SyslogConfig) CACertPath() string {
	filename := either(slConfig.CACertFileName, defaultCACertFileName)
	return filepath.Join(slConfig.LogDir, filename)
}

func (slConfig *SyslogConfig) ServerCertPath() string {
	filename := either(slConfig.ServerCertFileName, defaultServerCertFileName)
	return filepath.Join(slConfig.LogDir, filename)
}

func (slConfig *SyslogConfig) ServerKeyPath() string {
	filename := either(slConfig.ServerCertFileName, defaultServerKeyFileName)
	return filepath.Join(slConfig.LogDir, filename)
}

// Render generates the rsyslog config.
func (slConfig *SyslogConfig) Render() ([]byte, error) {
	var bootstrapHosts = func() []string {
		var hosts []string
		for _, addr := range slConfig.StateServerAddresses {
			parts := strings.Split(addr, ":")
			hosts = append(hosts, parts[0])
		}
		return hosts
	}

	var logFilePath = func() string {
		return fmt.Sprintf("%s/%s.log", slConfig.LogDir, slConfig.LogFileName)
	}

	t := template.New("")
	t.Funcs(template.FuncMap{
		"logfileName":    func() string { return slConfig.LogFileName },
		"bootstrapHosts": bootstrapHosts,
		"logfilePath":    logFilePath,
		"portNumber":     func() int { return slConfig.Port },
		"logDir":         func() string { return slConfig.LogDir },
		"namespace":      func() string { return slConfig.Namespace },
		"tagStart":       func() int { return tagOffset + len(slConfig.Namespace) },
		"tlsCACertPath":  slConfig.CACertPath,
		"tlsCertPath":    slConfig.ServerCertPath,
		"tlsKeyPath":     slConfig.ServerKeyPath,
	})

	// Process the rsyslog config template and echo to the conf file.
	p, err := t.Parse(slConfig.configTemplate)
	if err != nil {
		return nil, err
	}
	var confBuf bytes.Buffer
	if err := p.Execute(&confBuf, nil); err != nil {
		return nil, err
	}
	return confBuf.Bytes(), nil
}

// Write generates and writes the rsyslog config.
func (slConfig *SyslogConfig) Write() error {
	data, err := slConfig.Render()
	if err != nil {
		return err
	}
	err = ioutil.WriteFile(slConfig.ConfigFilePath(), data, 0644)
	return err
}
<|endoftext|>


<commit_before>package unison

import (
	"errors"
	"fmt"
	"sync"

	"github.com/Sirupsen/logrus"
	arg "github.com/alexflint/go-arg"
	"github.com/andersfylling/unison/constant"
	"github.com/bwmarrin/discordgo"
	shellquote "github.com/kballard/go-shellquote"
)

// CommandAction command logic to be executed
type CommandAction func(ctx *Context, msg *discordgo.Message, request string) error

// Command struct holds all the command details
type Command struct {
	// Name is the command title used to trigger the CommandAction
	Name string

	// Aliases are alternative command names. Usually shorter versions of the main.
	//Aliases []string. Should be stored on a per server basis, not on a per command basis.

	// Usage is a description of the command goal
	Usage string

	// Flags are optionals and bools that can be used in triggering the command
	Flags interface{} // used to generate a *arg.Parser

	// Action is the func run to execute the command
	Action CommandAction

	// Sub commands
	SubCommands []*Command

	// Deactivated true if the command should be ignored and viewed as "dead"
	Deactivated bool

	// Set the minimum required permissions for this command
	//	This level is inherited into each subcommand and must be overwritten if else is desired
	//	https://discordapp.com/developers/docs/topics/permissions
	Permissions uint32

	// Private
	//

	// go-arg parser for user input
	flagParser *arg.Parser // might be nil

	// mutex
	sync.RWMutex
}

// buildCommand configurates the command from the public fields before it gets stored.
func (cmd *Command) buildCommand() *Command {
	cmd.Lock()
	defer cmd.Unlock()

	errArr := []error{}

	// Build the command in layers..

	// Build a go-arg.Parser to handle the flags
	if cmd.Flags != nil {
		errArr = append(errArr, cmd.createParser(cmd.Flags))
	}

	// make sure the depth of sub commands is acceptable
	errArr = append(errArr, cmd.insistSubCommandDepth(constant.SubCommandDepthLimit))

	// Not sure if this matters, but make sure the User can write in the same channel that the bot
	// gets triggered from. Some one that cannot send a message should never be able to trigger a command.
	if cmd.Permissions == 0 {
		cmd.Permissions = 0x00000800
	}

	// check for issues
	for _, err := range errArr {
		if err != nil {
			logrus.Fatal(err)
		}
	}

	// completed and usable command instance
	return cmd
}

// createParser creates a go-arg parser that can be used to parse user input for flags/optionals
func (cmd *Command) createParser(dests ...interface{}) error {
	p, err := arg.NewParser(arg.Config{}, dests...)
	if err != nil {
		return err
	}

	cmd.flagParser = p

	return nil
}

// insistSubCommandDepth Make sure there aren't infinite depth of sub commands
func (cmd *Command) insistSubCommandDepth(depth int) error {

	if depth <= 0 && len(cmd.SubCommands) > 0 {
		errMsg := fmt.Sprintf("Too many recursive sub commands. Max depth is %d", constant.SubCommandDepthLimit)
		return errors.New(errMsg)
	}

	for _, c := range cmd.SubCommands {
		err := c.insistSubCommandDepth(depth - 1)
		if err != nil {
			return err
		}
	}

	return nil
}

// invokableWithPermissions checks if the permissions given has the minimum access level
func (cmd *Command) invokableWithPermissions(permissions uint32) bool {
	return (cmd.Permissions & permissions) == permissions
}

func (cmd *Command) parseInput(input string) error {
	if cmd.Flags == nil {
		return errors.New("No flags have been added")
	}

	args, err := shellquote.Split(input)
	if err != nil {
		return err
	}

	return cmd.flagParser.Parse(args)
}
<commit_msg>fixed typo<commit_after>package unison

import (
	"errors"
	"fmt"
	"sync"

	"github.com/Sirupsen/logrus"
	arg "github.com/alexflint/go-arg"
	"github.com/andersfylling/unison/constant"
	"github.com/bwmarrin/discordgo"
	shellquote "github.com/kballard/go-shellquote"
)

// CommandAction command logic to be executed
type CommandAction func(ctx *Context, msg *discordgo.Message, request string) error

// Command struct holds all the command details
type Command struct {
	// Name is the command title used to trigger the CommandAction
	Name string

	// Aliases are alternative command names. Usually shorter versions of the main.
	//Aliases []string. Should be stored on a per server basis, not on a per command basis.

	// Usage is a description of the command goal
	Usage string

	// Flags are optionals and bools that can be used in triggering the command
	Flags interface{} // used to generate a *arg.Parser

	// Action is the func run to execute the command
	Action CommandAction

	// Sub commands
	SubCommands []*Command

	// Deactivated true if the command should be ignored and viewed as "dead"
	Deactivated bool

	// Set the minimum required permissions for this command
	//	This level is inherited into each subcommand and must be overwritten if else is desired
	//	https://discordapp.com/developers/docs/topics/permissions
	Permissions uint32

	// Private
	//

	// go-arg parser for user input
	flagParser *arg.Parser // might be nil

	// mutex
	sync.RWMutex
}

// buildCommand builds the parser, checks command requirements, and sets default values
func (cmd *Command) buildCommand() *Command {
	cmd.Lock()
	defer cmd.Unlock()

	errArr := []error{}

	// Build the command in layers..

	// Build a go-arg.Parser to handle the flags
	if cmd.Flags != nil {
		errArr = append(errArr, cmd.createParser(cmd.Flags))
	}

	// make sure the depth of sub commands is acceptable
	errArr = append(errArr, cmd.insistSubCommandDepth(constant.SubCommandDepthLimit))

	// Not sure if this matters, but make sure the User can write in the same channel that the bot
	// gets triggered from. Some one that cannot send a message should never be able to trigger a command.
	if cmd.Permissions == 0 {
		cmd.Permissions = 0x00000800
	}

	// check for issues
	for _, err := range errArr {
		if err != nil {
			logrus.Fatal(err)
		}
	}

	// completed and usable command instance
	return cmd
}

// createParser creates a go-arg parser that can be used to parse user input for flags/optionals
func (cmd *Command) createParser(dests ...interface{}) error {
	p, err := arg.NewParser(arg.Config{}, dests...)
	if err != nil {
		return err
	}

	cmd.flagParser = p

	return nil
}

// insistSubCommandDepth Make sure there aren't infinite depth of sub commands
func (cmd *Command) insistSubCommandDepth(depth int) error {

	if depth <= 0 && len(cmd.SubCommands) > 0 {
		errMsg := fmt.Sprintf("Too many recursive sub commands. Max depth is %d", constant.SubCommandDepthLimit)
		return errors.New(errMsg)
	}

	for _, c := range cmd.SubCommands {
		err := c.insistSubCommandDepth(depth - 1)
		if err != nil {
			return err
		}
	}

	return nil
}

// invokableWithPermissions checks if the permissions given has the minimum access level
func (cmd *Command) invokableWithPermissions(permissions uint32) bool {
	return (cmd.Permissions & permissions) == permissions
}

func (cmd *Command) parseInput(input string) error {
	if cmd.Flags == nil {
		return errors.New("No flags have been added")
	}

	args, err := shellquote.Split(input)
	if err != nil {
		return err
	}

	return cmd.flagParser.Parse(args)
}
<|endoftext|>


<commit_before>package aero

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"mime"
	"net/http"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/OneOfOne/xxhash"
	"github.com/aerogo/session"
	"github.com/fatih/color"
	"github.com/julienschmidt/httprouter"
	cache "github.com/patrickmn/go-cache"
	"github.com/tomasen/realip"
	"github.com/valyala/fasthttp"
)

// This should be close to the MTU size of a TCP packet.
// Regarding performance it makes no sense to compress smaller files.
// Bandwidth can be saved however the savings are minimal for small files
// and the overhead of compressing can lead up to a 75% reduction
// in server speed under high load. Therefore in this case
// we're trying to optimize for performance, not bandwidth.
const gzipThreshold = 1450

const (
	serverHeader                  = "Server"
	server                        = "Aero"
	cacheControlHeader            = "Cache-Control"
	cacheControlAlwaysValidate    = "must-revalidate"
	cacheControlMedia             = "public, max-age=864000"
	contentTypeOptionsHeader      = "X-Content-Type-Options"
	contentTypeOptions            = "nosniff"
	xssProtectionHeader           = "X-XSS-Protection"
	xssProtection                 = "1; mode=block"
	etagHeader                    = "ETag"
	contentTypeHeader             = "Content-Type"
	contentTypeHTML               = "text/html; charset=utf-8"
	contentTypeJSON               = "application/json; charset=utf-8"
	contentTypePlainText          = "text/plain; charset=utf-8"
	contentEncodingHeader         = "Content-Encoding"
	contentEncodingGzip           = "gzip"
	contentLengthHeader           = "Content-Length"
	responseTimeHeader            = "X-Response-Time"
	ifNoneMatchHeader             = "If-None-Match"
	xFrameOptionsHeader           = "X-Frame-Options"
	xFrameOptions                 = "SAMEORIGIN"
	referrerPolicyHeader          = "Referrer-Policy"
	referrerPolicySameOrigin      = "no-referrer"
	strictTransportSecurityHeader = "Strict-Transport-Security"
	strictTransportSecurity       = "max-age=31536000; includeSubDomains; preload"
	contentSecurityPolicyHeader   = "Content-Security-Policy"
)

// Context ...
type Context struct {
	// net/http
	request  *http.Request
	response http.ResponseWriter
	params   httprouter.Params

	// A pointer to the application this request occured on.
	App *Application

	// Status code
	StatusCode int

	// Start time
	start time.Time

	// User session
	session *session.Session
}

// Handle ...
type Handle func(*Context) string

// Session returns the session of the context or creates and caches a new session.
func (ctx *Context) Session() *session.Session {
	// Return cached session if available.
	if ctx.session != nil {
		return ctx.session
	}

	// Check if the client has a session cookie already.
	cookie, err := ctx.request.Cookie("sid")

	if err == nil {
		sid := cookie.Value

		if session.IsValidID(sid) {
			ctx.session, err = ctx.App.Sessions.Store.Get(sid)

			if err != nil {
				color.Red(err.Error())
			}

			if ctx.session != nil {
				return ctx.session
			}
		}
	}

	// Create a new session
	ctx.session = ctx.App.Sessions.New()

	// Create a session cookie in the client
	ctx.createSessionCookie()

	return ctx.session
}

// createSessionCookie creates a session cookie in the client.
func (ctx *Context) createSessionCookie() {
	sessionCookie := http.Cookie{
		Name:     "sid",
		Value:    ctx.session.ID(),
		HttpOnly: true,
		Secure:   true,
		MaxAge:   ctx.App.Sessions.Duration,
		Path:     "/",
	}

	http.SetCookie(ctx.response, &sessionCookie)

	// HACK: Add SameSite attribute
	// Remove this once it's available inside http.Cookie
	cookieData := ctx.response.Header().Get("Set-Cookie")
	cookieData += "; SameSite=lax"
	ctx.response.Header().Set("Set-Cookie", cookieData)
}

// HasSession indicates whether the client has a valid session or not.
func (ctx *Context) HasSession() bool {
	if ctx.session != nil {
		return true
	}

	cookie, err := ctx.request.Cookie("sid")

	if err != nil || !session.IsValidID(cookie.Value) {
		return false
	}

	ctx.session, _ = ctx.App.Sessions.Store.Get(cookie.Value)

	return ctx.session != nil
}

// JSON encodes the object to a JSON string and responds.
func (ctx *Context) JSON(value interface{}) string {
	bytes, _ := json.Marshal(value)

	ctx.SetResponseHeader(contentTypeHeader, contentTypeJSON)
	return string(bytes)
}

// HTML sends a HTML string.
func (ctx *Context) HTML(html string) string {
	ctx.SetResponseHeader(contentTypeHeader, contentTypeHTML)
	ctx.SetResponseHeader(contentTypeOptionsHeader, contentTypeOptions)
	ctx.SetResponseHeader(xssProtectionHeader, xssProtection)
	ctx.SetResponseHeader(xFrameOptionsHeader, xFrameOptions)
	ctx.SetResponseHeader(referrerPolicyHeader, referrerPolicySameOrigin)

	if ctx.App.Security.Certificate != "" {
		ctx.SetResponseHeader(strictTransportSecurityHeader, strictTransportSecurity)
		ctx.SetResponseHeader(contentSecurityPolicyHeader, ctx.App.contentSecurityPolicy)
	}

	return html
}

// Text sends a plain text string.
func (ctx *Context) Text(text string) string {
	ctx.SetResponseHeader(contentTypeHeader, contentTypePlainText)
	return text
}

// File sends the contents of a local file and determines its mime type by extension.
func (ctx *Context) File(file string) string {
	extension := filepath.Ext(file)
	mimeType := mime.TypeByExtension(extension)
	data, _ := ioutil.ReadFile(file)

	if mimeType == "" {
		mimeType = http.DetectContentType(data)
	}

	ctx.SetResponseHeader(contentTypeHeader, mimeType)
	return string(data)
}

// TryWebP tries to serve a WebP image but will fall back to the specified extension if needed.
func (ctx *Context) TryWebP(path string, extension string) string {
	if ctx.CanUseWebP() {
		extension = ".webp"
	}

	return ctx.File(path + extension)
}

// Error should be used for sending error messages to the user.
func (ctx *Context) Error(statusCode int, explanation string, err error) string {
	ctx.StatusCode = statusCode
	ctx.SetResponseHeader(contentTypeHeader, contentTypeHTML)
	// ctx.App.Logger.Error(
	// 	color.RedString(explanation),
	// 	zap.String("error", err.Error()),
	// 	zap.String("url", ctx.request.RequestURI),
	// )
	if err != nil {
		detailed := err.Error()
		color.Red(detailed)
		return fmt.Sprintf("%s (%s)", explanation, detailed)
	}

	return explanation
}

// GetRequestHeader retrieves the value for the request header.
func (ctx *Context) GetRequestHeader(header string) string {
	return ctx.request.Header.Get(header)
}

// SetResponseHeader sets response header to value.
func (ctx *Context) SetResponseHeader(header string, value string) {
	ctx.response.Header().Set(header, value)
}

// URI returns the relative path, e.g. /blog/post/123.
func (ctx *Context) URI() string {
	return ctx.request.URL.Path
}

// SetURI sets the relative path, e.g. /blog/post/123.
func (ctx *Context) SetURI(b string) {
	ctx.request.URL.Path = b
}

// Get retrieves an URL parameter.
func (ctx *Context) Get(param string) string {
	return ctx.params.ByName(param)
}

// GetInt retrieves an URL parameter as an integer.
func (ctx *Context) GetInt(param string) (int, error) {
	return strconv.Atoi(ctx.Get(param))
}

// RequestBody returns the request body as a string.
func (ctx *Context) RequestBody() []byte {
	body, err := ioutil.ReadAll(ctx.request.Body)

	if err != nil {
		panic(err)
	}

	return body
}

// RealIP tries to determine the real IP address of the request.
func (ctx *Context) RealIP() string {
	return realip.RealIP(ctx.request)
}

// UserAgent retrieves the user agent for the given request.
func (ctx *Context) UserAgent() string {
	ctx.request.URL.Query()
	return ctx.request.UserAgent()
}

// Query retrieves the value for the given URL query parameter.
func (ctx *Context) Query(param string) string {
	return ctx.request.URL.Query().Get(param)
}

// Redirect redirects to the given URL using status code 302.
func (ctx *Context) Redirect(url string) string {
	ctx.StatusCode = http.StatusFound
	ctx.SetResponseHeader("Location", url)
	return ""
}

// RedirectPermanently redirects to the given URL and indicates that this is a permanent change using status code 301.
func (ctx *Context) RedirectPermanently(url string) string {
	ctx.StatusCode = http.StatusPermanentRedirect
	ctx.SetResponseHeader("Location", url)
	return ""
}

// CanUseWebP checks the Accept header to find out if WebP is supported by the client's browser.
func (ctx *Context) CanUseWebP() bool {
	accept := ctx.GetRequestHeader("Accept")

	if strings.Index(accept, "image/webp") != -1 {
		return true
	}

	return false
}

// Respond responds either with raw code or gzipped if the
// code length is greater than the gzip threshold.
func (ctx *Context) Respond(code string) {
	ctx.RespondBytes(StringToBytesUnsafe(code))
}

// RespondBytes responds either with raw code or gzipped if the
// code length is greater than the gzip threshold. Requires a byte slice.
func (ctx *Context) RespondBytes(b []byte) {
	response := ctx.response
	header := response.Header()
	contentType := ctx.response.Header().Get(contentTypeHeader)
	isMedia := strings.HasPrefix(contentType, "image/") || strings.HasPrefix(contentType, "video/")

	// Headers
	if isMedia {
		header.Set(cacheControlHeader, cacheControlMedia)
	} else {
		header.Set(cacheControlHeader, cacheControlAlwaysValidate)
		header.Set(serverHeader, server)
		header.Set(responseTimeHeader, strconv.FormatInt(time.Since(ctx.start).Nanoseconds()/1000, 10)+" us")
	}

	// Small response
	if len(b) < gzipThreshold {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// ETag generation
	h := xxhash.NewS64(0)
	h.Write(b)
	etag := strconv.FormatUint(h.Sum64(), 16)

	// If client cache is up to date, send 304 with no response body.
	clientETag := ctx.request.Header.Get(ifNoneMatchHeader)

	if etag == clientETag {
		response.WriteHeader(304)
		return
	}

	// Set ETag
	header.Set(etagHeader, etag)

	// No GZip?
	if !ctx.App.Config.GZip || isMedia {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// GZip
	header.Set(contentEncodingHeader, contentEncodingGzip)

	if ctx.App.Config.GZipCache {
		cachedResponse, found := ctx.App.gzipCache.Get(etag)

		if found {
			cachedResponseBytes := cachedResponse.([]byte)
			header.Set(contentLengthHeader, strconv.Itoa(len(cachedResponseBytes)))
			response.WriteHeader(ctx.StatusCode)
			response.Write(cachedResponseBytes)
			return
		}
	}

	var buffer bytes.Buffer
	writer := bufio.NewWriter(&buffer)
	fasthttp.WriteGzipLevel(writer, b, 9)
	writer.Flush()
	gzippedBytes := buffer.Bytes()

	header.Set(contentLengthHeader, strconv.Itoa(len(gzippedBytes)))
	response.WriteHeader(ctx.StatusCode)
	response.Write(gzippedBytes)

	if ctx.App.Config.GZipCache {
		ctx.App.gzipCache.Set(etag, gzippedBytes, cache.DefaultExpiration)
	}
}
<commit_msg>Added IsMediaResponse()<commit_after>package aero

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"mime"
	"net/http"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/OneOfOne/xxhash"
	"github.com/aerogo/session"
	"github.com/fatih/color"
	"github.com/julienschmidt/httprouter"
	cache "github.com/patrickmn/go-cache"
	"github.com/tomasen/realip"
	"github.com/valyala/fasthttp"
)

// This should be close to the MTU size of a TCP packet.
// Regarding performance it makes no sense to compress smaller files.
// Bandwidth can be saved however the savings are minimal for small files
// and the overhead of compressing can lead up to a 75% reduction
// in server speed under high load. Therefore in this case
// we're trying to optimize for performance, not bandwidth.
const gzipThreshold = 1450

const (
	serverHeader                  = "Server"
	server                        = "Aero"
	cacheControlHeader            = "Cache-Control"
	cacheControlAlwaysValidate    = "must-revalidate"
	cacheControlMedia             = "public, max-age=864000"
	contentTypeOptionsHeader      = "X-Content-Type-Options"
	contentTypeOptions            = "nosniff"
	xssProtectionHeader           = "X-XSS-Protection"
	xssProtection                 = "1; mode=block"
	etagHeader                    = "ETag"
	contentTypeHeader             = "Content-Type"
	contentTypeHTML               = "text/html; charset=utf-8"
	contentTypeJSON               = "application/json; charset=utf-8"
	contentTypePlainText          = "text/plain; charset=utf-8"
	contentEncodingHeader         = "Content-Encoding"
	contentEncodingGzip           = "gzip"
	contentLengthHeader           = "Content-Length"
	responseTimeHeader            = "X-Response-Time"
	ifNoneMatchHeader             = "If-None-Match"
	xFrameOptionsHeader           = "X-Frame-Options"
	xFrameOptions                 = "SAMEORIGIN"
	referrerPolicyHeader          = "Referrer-Policy"
	referrerPolicySameOrigin      = "no-referrer"
	strictTransportSecurityHeader = "Strict-Transport-Security"
	strictTransportSecurity       = "max-age=31536000; includeSubDomains; preload"
	contentSecurityPolicyHeader   = "Content-Security-Policy"
)

// Context ...
type Context struct {
	// net/http
	request  *http.Request
	response http.ResponseWriter
	params   httprouter.Params

	// A pointer to the application this request occured on.
	App *Application

	// Status code
	StatusCode int

	// Start time
	start time.Time

	// User session
	session *session.Session
}

// Handle ...
type Handle func(*Context) string

// Session returns the session of the context or creates and caches a new session.
func (ctx *Context) Session() *session.Session {
	// Return cached session if available.
	if ctx.session != nil {
		return ctx.session
	}

	// Check if the client has a session cookie already.
	cookie, err := ctx.request.Cookie("sid")

	if err == nil {
		sid := cookie.Value

		if session.IsValidID(sid) {
			ctx.session, err = ctx.App.Sessions.Store.Get(sid)

			if err != nil {
				color.Red(err.Error())
			}

			if ctx.session != nil {
				return ctx.session
			}
		}
	}

	// Create a new session
	ctx.session = ctx.App.Sessions.New()

	// Create a session cookie in the client
	ctx.createSessionCookie()

	return ctx.session
}

// createSessionCookie creates a session cookie in the client.
func (ctx *Context) createSessionCookie() {
	sessionCookie := http.Cookie{
		Name:     "sid",
		Value:    ctx.session.ID(),
		HttpOnly: true,
		Secure:   true,
		MaxAge:   ctx.App.Sessions.Duration,
		Path:     "/",
	}

	http.SetCookie(ctx.response, &sessionCookie)

	// HACK: Add SameSite attribute
	// Remove this once it's available inside http.Cookie
	cookieData := ctx.response.Header().Get("Set-Cookie")
	cookieData += "; SameSite=lax"
	ctx.response.Header().Set("Set-Cookie", cookieData)
}

// HasSession indicates whether the client has a valid session or not.
func (ctx *Context) HasSession() bool {
	if ctx.session != nil {
		return true
	}

	cookie, err := ctx.request.Cookie("sid")

	if err != nil || !session.IsValidID(cookie.Value) {
		return false
	}

	ctx.session, _ = ctx.App.Sessions.Store.Get(cookie.Value)

	return ctx.session != nil
}

// JSON encodes the object to a JSON string and responds.
func (ctx *Context) JSON(value interface{}) string {
	bytes, _ := json.Marshal(value)

	ctx.SetResponseHeader(contentTypeHeader, contentTypeJSON)
	return string(bytes)
}

// HTML sends a HTML string.
func (ctx *Context) HTML(html string) string {
	ctx.SetResponseHeader(contentTypeHeader, contentTypeHTML)
	ctx.SetResponseHeader(contentTypeOptionsHeader, contentTypeOptions)
	ctx.SetResponseHeader(xssProtectionHeader, xssProtection)
	ctx.SetResponseHeader(xFrameOptionsHeader, xFrameOptions)
	ctx.SetResponseHeader(referrerPolicyHeader, referrerPolicySameOrigin)

	if ctx.App.Security.Certificate != "" {
		ctx.SetResponseHeader(strictTransportSecurityHeader, strictTransportSecurity)
		ctx.SetResponseHeader(contentSecurityPolicyHeader, ctx.App.contentSecurityPolicy)
	}

	return html
}

// Text sends a plain text string.
func (ctx *Context) Text(text string) string {
	ctx.SetResponseHeader(contentTypeHeader, contentTypePlainText)
	return text
}

// File sends the contents of a local file and determines its mime type by extension.
func (ctx *Context) File(file string) string {
	extension := filepath.Ext(file)
	mimeType := mime.TypeByExtension(extension)
	data, _ := ioutil.ReadFile(file)

	if mimeType == "" {
		mimeType = http.DetectContentType(data)
	}

	ctx.SetResponseHeader(contentTypeHeader, mimeType)
	return string(data)
}

// TryWebP tries to serve a WebP image but will fall back to the specified extension if needed.
func (ctx *Context) TryWebP(path string, extension string) string {
	if ctx.CanUseWebP() {
		extension = ".webp"
	}

	return ctx.File(path + extension)
}

// Error should be used for sending error messages to the user.
func (ctx *Context) Error(statusCode int, explanation string, err error) string {
	ctx.StatusCode = statusCode
	ctx.SetResponseHeader(contentTypeHeader, contentTypeHTML)
	// ctx.App.Logger.Error(
	// 	color.RedString(explanation),
	// 	zap.String("error", err.Error()),
	// 	zap.String("url", ctx.request.RequestURI),
	// )
	if err != nil {
		detailed := err.Error()
		color.Red(detailed)
		return fmt.Sprintf("%s (%s)", explanation, detailed)
	}

	return explanation
}

// GetRequestHeader retrieves the value for the request header.
func (ctx *Context) GetRequestHeader(header string) string {
	return ctx.request.Header.Get(header)
}

// SetResponseHeader sets response header to value.
func (ctx *Context) SetResponseHeader(header string, value string) {
	ctx.response.Header().Set(header, value)
}

// URI returns the relative path, e.g. /blog/post/123.
func (ctx *Context) URI() string {
	return ctx.request.URL.Path
}

// SetURI sets the relative path, e.g. /blog/post/123.
func (ctx *Context) SetURI(b string) {
	ctx.request.URL.Path = b
}

// Get retrieves an URL parameter.
func (ctx *Context) Get(param string) string {
	return ctx.params.ByName(param)
}

// GetInt retrieves an URL parameter as an integer.
func (ctx *Context) GetInt(param string) (int, error) {
	return strconv.Atoi(ctx.Get(param))
}

// RequestBody returns the request body as a string.
func (ctx *Context) RequestBody() []byte {
	body, err := ioutil.ReadAll(ctx.request.Body)

	if err != nil {
		panic(err)
	}

	return body
}

// RealIP tries to determine the real IP address of the request.
func (ctx *Context) RealIP() string {
	return realip.RealIP(ctx.request)
}

// UserAgent retrieves the user agent for the given request.
func (ctx *Context) UserAgent() string {
	ctx.request.URL.Query()
	return ctx.request.UserAgent()
}

// Query retrieves the value for the given URL query parameter.
func (ctx *Context) Query(param string) string {
	return ctx.request.URL.Query().Get(param)
}

// Redirect redirects to the given URL using status code 302.
func (ctx *Context) Redirect(url string) string {
	ctx.StatusCode = http.StatusFound
	ctx.SetResponseHeader("Location", url)
	return ""
}

// RedirectPermanently redirects to the given URL and indicates that this is a permanent change using status code 301.
func (ctx *Context) RedirectPermanently(url string) string {
	ctx.StatusCode = http.StatusPermanentRedirect
	ctx.SetResponseHeader("Location", url)
	return ""
}

// CanUseWebP checks the Accept header to find out if WebP is supported by the client's browser.
func (ctx *Context) CanUseWebP() bool {
	accept := ctx.GetRequestHeader("Accept")

	if strings.Index(accept, "image/webp") != -1 {
		return true
	}

	return false
}

// IsMediaResponse returns whether the given context has already set its content type to a media type.
func (ctx *Context) IsMediaResponse() bool {
	contentType := ctx.response.Header().Get(contentTypeHeader)
	return strings.HasPrefix(contentType, "image/") || strings.HasPrefix(contentType, "video/")
}

// Respond responds either with raw code or gzipped if the
// code length is greater than the gzip threshold.
func (ctx *Context) Respond(code string) {
	ctx.RespondBytes(StringToBytesUnsafe(code))
}

// RespondBytes responds either with raw code or gzipped if the
// code length is greater than the gzip threshold. Requires a byte slice.
func (ctx *Context) RespondBytes(b []byte) {
	response := ctx.response
	header := response.Header()
	isMedia := ctx.IsMediaResponse()

	// Headers
	if isMedia {
		header.Set(cacheControlHeader, cacheControlMedia)
	} else {
		header.Set(cacheControlHeader, cacheControlAlwaysValidate)
		header.Set(serverHeader, server)
		header.Set(responseTimeHeader, strconv.FormatInt(time.Since(ctx.start).Nanoseconds()/1000, 10)+" us")
	}

	// Small response
	if len(b) < gzipThreshold {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// ETag generation
	h := xxhash.NewS64(0)
	h.Write(b)
	etag := strconv.FormatUint(h.Sum64(), 16)

	// If client cache is up to date, send 304 with no response body.
	clientETag := ctx.request.Header.Get(ifNoneMatchHeader)

	if etag == clientETag {
		response.WriteHeader(304)
		return
	}

	// Set ETag
	header.Set(etagHeader, etag)

	// No GZip?
	if !ctx.App.Config.GZip || isMedia {
		header.Set(contentLengthHeader, strconv.Itoa(len(b)))
		response.WriteHeader(ctx.StatusCode)
		response.Write(b)
		return
	}

	// GZip
	header.Set(contentEncodingHeader, contentEncodingGzip)

	if ctx.App.Config.GZipCache {
		cachedResponse, found := ctx.App.gzipCache.Get(etag)

		if found {
			cachedResponseBytes := cachedResponse.([]byte)
			header.Set(contentLengthHeader, strconv.Itoa(len(cachedResponseBytes)))
			response.WriteHeader(ctx.StatusCode)
			response.Write(cachedResponseBytes)
			return
		}
	}

	var buffer bytes.Buffer
	writer := bufio.NewWriter(&buffer)
	fasthttp.WriteGzipLevel(writer, b, 9)
	writer.Flush()
	gzippedBytes := buffer.Bytes()

	header.Set(contentLengthHeader, strconv.Itoa(len(gzippedBytes)))
	response.WriteHeader(ctx.StatusCode)
	response.Write(gzippedBytes)

	if ctx.App.Config.GZipCache {
		ctx.App.gzipCache.Set(etag, gzippedBytes, cache.DefaultExpiration)
	}
}
<|endoftext|>


<commit_before>package redis

import (
	"context"
	"fmt"
	"time"

	libredis "github.com/go-redis/redis"
	"github.com/pkg/errors"

	"github.com/ulule/limiter/v3"
	"github.com/ulule/limiter/v3/drivers/store/common"
)

// Client is an interface thats allows to use a redis cluster or a redis single client seamlessly.
type Client interface {
	Ping() *libredis.StatusCmd
	Get(key string) *libredis.StringCmd
	Set(key string, value interface{}, expiration time.Duration) *libredis.StatusCmd
	Watch(handler func(*libredis.Tx) error, keys ...string) error
	Del(keys ...string) *libredis.IntCmd
	SetNX(key string, value interface{}, expiration time.Duration) *libredis.BoolCmd
	Eval(script string, keys []string, args ...interface{}) *libredis.Cmd
}

// Store is the redis store.
type Store struct {
	// Prefix used for the key.
	Prefix string
	// MaxRetry is the maximum number of retry under race conditions.
	MaxRetry int
	// client used to communicate with redis server.
	client Client
}

// NewStore returns an instance of redis store with defaults.
func NewStore(client Client) (limiter.Store, error) {
	return NewStoreWithOptions(client, limiter.StoreOptions{
		Prefix:          limiter.DefaultPrefix,
		CleanUpInterval: limiter.DefaultCleanUpInterval,
		MaxRetry:        limiter.DefaultMaxRetry,
	})
}

// NewStoreWithOptions returns an instance of redis store with options.
func NewStoreWithOptions(client Client, options limiter.StoreOptions) (limiter.Store, error) {
	store := &Store{
		client:   client,
		Prefix:   options.Prefix,
		MaxRetry: options.MaxRetry,
	}

	if store.MaxRetry <= 0 {
		store.MaxRetry = 1
	}

	_, err := store.ping()
	if err != nil {
		return nil, err
	}

	return store, nil
}

// Get returns the limit for given identifier.
func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
	key = fmt.Sprintf("%s:%s", store.Prefix, key)
	now := time.Now()

	lctx := limiter.Context{}
	onWatch := func(rtx *libredis.Tx) error {

		created, err := store.doSetValue(rtx, key, rate.Period)
		if err != nil {
			return err
		}

		if created {
			expiration := now.Add(rate.Period)
			lctx = common.GetContextFromState(now, rate, expiration, 1)
			return nil
		}

		count, ttl, err := store.doUpdateValue(rtx, key, rate.Period)
		if err != nil {
			return err
		}

		expiration := now.Add(rate.Period)
		if ttl > 0 {
			expiration = now.Add(ttl)
		}

		lctx = common.GetContextFromState(now, rate, expiration, count)
		return nil
	}

	err := store.client.Watch(onWatch, key)
	if err != nil {
		err = errors.Wrapf(err, "limiter: cannot get value for %s", key)
		return limiter.Context{}, err
	}

	return lctx, nil
}

// Peek returns the limit for given identifier, without modification on current values.
func (store *Store) Peek(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
	key = fmt.Sprintf("%s:%s", store.Prefix, key)
	now := time.Now()

	lctx := limiter.Context{}
	onWatch := func(rtx *libredis.Tx) error {
		count, ttl, err := store.doPeekValue(rtx, key)
		if err != nil {
			return err
		}

		expiration := now.Add(rate.Period)
		if ttl > 0 {
			expiration = now.Add(ttl)
		}

		lctx = common.GetContextFromState(now, rate, expiration, count)
		return nil
	}

	err := store.client.Watch(onWatch, key)
	if err != nil {
		err = errors.Wrapf(err, "limiter: cannot peek value for %s", key)
		return limiter.Context{}, err
	}

	return lctx, nil
}

// doPeekValue will execute peekValue with a retry mecanism (optimistic locking) until store.MaxRetry is reached.
func (store *Store) doPeekValue(rtx *libredis.Tx, key string) (int64, time.Duration, error) {
	for i := 0; i < store.MaxRetry; i++ {
		count, ttl, err := peekValue(rtx, key)
		if err == nil {
			return count, ttl, nil
		}
	}
	return 0, 0, errors.New("retry limit exceeded")
}

// peekValue will retrieve the counter and its expiration for given key.
func peekValue(rtx *libredis.Tx, key string) (int64, time.Duration, error) {
	pipe := rtx.Pipeline()
	value := pipe.Get(key)
	expire := pipe.PTTL(key)

	_, err := pipe.Exec()
	if err != nil && err != libredis.Nil {
		return 0, 0, err
	}

	count, err := value.Int64()
	if err != nil && err != libredis.Nil {
		return 0, 0, err
	}

	ttl, err := expire.Result()
	if err != nil {
		return 0, 0, err
	}

	return count, ttl, nil
}

// doSetValue will execute setValue with a retry mecanism (optimistic locking) until store.MaxRetry is reached.
func (store *Store) doSetValue(rtx *libredis.Tx, key string, expiration time.Duration) (bool, error) {
	for i := 0; i < store.MaxRetry; i++ {
		created, err := setValue(rtx, key, expiration)
		if err == nil {
			return created, nil
		}
	}
	return false, errors.New("retry limit exceeded")
}

// setValue will try to initialize a new counter if given key doesn't exists.
func setValue(rtx *libredis.Tx, key string, expiration time.Duration) (bool, error) {
	value := rtx.SetNX(key, 1, expiration)

	created, err := value.Result()
	if err != nil {
		return false, err
	}

	return created, nil
}

// doUpdateValue will execute setValue with a retry mecanism (optimistic locking) until store.MaxRetry is reached.
func (store *Store) doUpdateValue(rtx *libredis.Tx, key string,
	expiration time.Duration) (int64, time.Duration, error) {
	for i := 0; i < store.MaxRetry; i++ {
		count, ttl, err := updateValue(rtx, key, expiration)
		if err == nil {
			return count, ttl, nil
		}

		// If ttl is negative and there is an error, do not retry an update.
		if ttl < 0 {
			return 0, 0, err
		}
	}
	return 0, 0, errors.New("retry limit exceeded")
}

// updateValue will try to increment the counter identified by given key.
func updateValue(rtx *libredis.Tx, key string, expiration time.Duration) (int64, time.Duration, error) {
	pipe := rtx.Pipeline()
	value := pipe.Incr(key)
	expire := pipe.PTTL(key)

	_, err := pipe.Exec()
	if err != nil {
		return 0, 0, err
	}

	count, err := value.Result()
	if err != nil {
		return 0, 0, err
	}

	ttl, err := expire.Result()
	if err != nil {
		return 0, 0, err
	}

	// If ttl is negative, we have to define key expiration.
	if ttl < 0 {
		expire := rtx.Expire(key, expiration)

		ok, err := expire.Result()
		if err != nil {
			return count, ttl, err
		}

		if !ok {
			return count, ttl, errors.New("cannot configure timeout on key")
		}
	}

	return count, ttl, nil

}

// ping checks if redis is alive.
func (store *Store) ping() (bool, error) {
	cmd := store.client.Ping()

	pong, err := cmd.Result()
	if err != nil {
		return false, errors.Wrap(err, "limiter: cannot ping redis server")
	}

	return (pong == "PONG"), nil
}
<commit_msg>chore: merge branch 'master'<commit_after>package redis

import (
	"context"
	"fmt"
	"time"

	libredis "github.com/go-redis/redis"
	"github.com/pkg/errors"

	"github.com/ulule/limiter/v3"
	"github.com/ulule/limiter/v3/drivers/store/common"
)

// Client is an interface thats allows to use a redis cluster or a redis single client seamlessly.
type Client interface {
	Ping() *libredis.StatusCmd
	Get(key string) *libredis.StringCmd
	Set(key string, value interface{}, expiration time.Duration) *libredis.StatusCmd
	Watch(handler func(*libredis.Tx) error, keys ...string) error
	Del(keys ...string) *libredis.IntCmd
	SetNX(key string, value interface{}, expiration time.Duration) *libredis.BoolCmd
	Eval(script string, keys []string, args ...interface{}) *libredis.Cmd
}

// Store is the redis store.
type Store struct {
	// Prefix used for the key.
	Prefix string
	// MaxRetry is the maximum number of retry under race conditions.
	MaxRetry int
	// client used to communicate with redis server.
	client Client
}

// NewStore returns an instance of redis store with defaults.
func NewStore(client Client) (limiter.Store, error) {
	return NewStoreWithOptions(client, limiter.StoreOptions{
		Prefix:          limiter.DefaultPrefix,
		CleanUpInterval: limiter.DefaultCleanUpInterval,
		MaxRetry:        limiter.DefaultMaxRetry,
	})
}

// NewStoreWithOptions returns an instance of redis store with options.
func NewStoreWithOptions(client Client, options limiter.StoreOptions) (limiter.Store, error) {
	store := &Store{
		client:   client,
		Prefix:   options.Prefix,
		MaxRetry: options.MaxRetry,
	}

	if store.MaxRetry <= 0 {
		store.MaxRetry = 1
	}

	_, err := store.ping()
	if err != nil {
		return nil, err
	}

	return store, nil
}

// Get returns the limit for given identifier.
func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
	key = fmt.Sprintf("%s:%s", store.Prefix, key)
	now := time.Now()

	lctx := limiter.Context{}
	onWatch := func(rtx *libredis.Tx) error {

		created, err := store.doSetValue(rtx, key, rate.Period)
		if err != nil {
			return err
		}

		if created {
			expiration := now.Add(rate.Period)
			lctx = common.GetContextFromState(now, rate, expiration, 1)
			return nil
		}

		count, ttl, err := store.doUpdateValue(rtx, key, rate.Period)
		if err != nil {
			return err
		}

		expiration := now.Add(rate.Period)
		if ttl > 0 {
			expiration = now.Add(ttl)
		}

		lctx = common.GetContextFromState(now, rate, expiration, count)
		return nil
	}

	err := store.client.Watch(onWatch, key)
	if err != nil {
		err = errors.Wrapf(err, "limiter: cannot get value for %s", key)
		return limiter.Context{}, err
	}

	return lctx, nil
}

// Peek returns the limit for given identifier, without modification on current values.
func (store *Store) Peek(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
	key = fmt.Sprintf("%s:%s", store.Prefix, key)
	now := time.Now()

	lctx := limiter.Context{}
	onWatch := func(rtx *libredis.Tx) error {
		count, ttl, err := store.doPeekValue(rtx, key)
		if err != nil {
			return err
		}

		expiration := now.Add(rate.Period)
		if ttl > 0 {
			expiration = now.Add(ttl)
		}

		lctx = common.GetContextFromState(now, rate, expiration, count)
		return nil
	}

	err := store.client.Watch(onWatch, key)
	if err != nil {
		err = errors.Wrapf(err, "limiter: cannot peek value for %s", key)
		return limiter.Context{}, err
	}

	return lctx, nil
}

// doPeekValue will execute peekValue with a retry mecanism (optimistic locking) until store.MaxRetry is reached.
func (store *Store) doPeekValue(rtx *libredis.Tx, key string) (int64, time.Duration, error) {
	for i := 0; i < store.MaxRetry; i++ {
		count, ttl, err := peekValue(rtx, key)
		if err == nil {
			return count, ttl, nil
		}
	}
	return 0, 0, errors.New("retry limit exceeded")
}

// peekValue will retrieve the counter and its expiration for given key.
func peekValue(rtx *libredis.Tx, key string) (int64, time.Duration, error) {
	pipe := rtx.Pipeline()
	value := pipe.Get(key)
	expire := pipe.PTTL(key)

	_, err := pipe.Exec()
	if err != nil && err != libredis.Nil {
		return 0, 0, err
	}

	count, err := value.Int64()
	if err != nil && err != libredis.Nil {
		return 0, 0, err
	}

	ttl, err := expire.Result()
	if err != nil {
		return 0, 0, err
	}

	return count, ttl, nil
}

// doSetValue will execute setValue with a retry mecanism (optimistic locking) until store.MaxRetry is reached.
func (store *Store) doSetValue(rtx *libredis.Tx, key string, expiration time.Duration) (bool, error) {
	for i := 0; i < store.MaxRetry; i++ {
		created, err := setValue(rtx, key, expiration)
		if err == nil {
			return created, nil
		}
	}
	return false, errors.New("retry limit exceeded")
}

// setValue will try to initialize a new counter if given key doesn't exists.
func setValue(rtx *libredis.Tx, key string, expiration time.Duration) (bool, error) {
	value := rtx.SetNX(key, 1, expiration)

	created, err := value.Result()
	if err != nil {
		return false, err
	}

	return created, nil
}

// doUpdateValue will execute setValue with a retry mecanism (optimistic locking) until store.MaxRetry is reached.
func (store *Store) doUpdateValue(rtx *libredis.Tx, key string,
	expiration time.Duration) (int64, time.Duration, error) {
	for i := 0; i < store.MaxRetry; i++ {
		count, ttl, err := updateValue(rtx, key, expiration)
		if err == nil {
			return count, ttl, nil
		}

		// If ttl is negative and there is an error, do not retry an update.
		if ttl < 0 {
			return 0, 0, err
		}
	}
	return 0, 0, errors.New("retry limit exceeded")
}

// updateValue will try to increment the counter identified by given key.
func updateValue(rtx *libredis.Tx, key string, expiration time.Duration) (int64, time.Duration, error) {
	pipe := rtx.Pipeline()
	value := pipe.Incr(key)
	expire := pipe.PTTL(key)

	_, err := pipe.Exec()
	if err != nil {
		return 0, 0, err
	}

	count, err := value.Result()
	if err != nil {
		return 0, 0, err
	}

	ttl, err := expire.Result()
	if err != nil {
		return 0, 0, err
	}

	// If ttl is -1, we have to define key expiration.
	// PTTL return values changed as of Redis 2.8
	// Now the command returns -2 if the key does not exist, and -1 if the key exists, but there is no expiry set
	// We shouldn't try to set an expiry on a key that doesn't exist
	if ttl == -1 {
		expire := rtx.Expire(key, expiration)

		ok, err := expire.Result()
		if err != nil {
			return count, ttl, err
		}

		if !ok {
			return count, ttl, errors.New("cannot configure timeout on key")
		}
	}

	return count, ttl, nil

}

// ping checks if redis is alive.
func (store *Store) ping() (bool, error) {
	cmd := store.client.Ping()

	pong, err := cmd.Result()
	if err != nil {
		return false, errors.Wrap(err, "limiter: cannot ping redis server")
	}

	return (pong == "PONG"), nil
}
<|endoftext|>


<commit_before>package gobot

import (
    . "github.com/onsi/ginkgo"
    . "github.com/onsi/gomega"
)

var _ = Describe("Robot", func() {

var (
        someRobot Robot
    )

    BeforeEach(func() {
        someRobot = Robot{
          Work: func() {
          },
        }
    })

    Describe("Robot", func() {
        Context("when valid", func() {
            PIt("should Start", func() {
                Expect(true)
            })
            PIt("should initConnections", func() {
                Expect(true)
            })
            PIt("should initDevices", func() {
                Expect(true)
            })
            PIt("should startConnections", func() {
                Expect(true)
            })
            PIt("should startDevices", func() {
                Expect(true)
            })
            PIt("should GetDevices", func() {
                Expect(true)
            })
            PIt("should GetDevice", func() {
                Expect(true)
            })
        })
    })
})
<commit_msg>Nicer format<commit_after>package gobot

import (
  . "github.com/onsi/ginkgo"
  . "github.com/onsi/gomega"
)

var _ = Describe("Robot", func() {

  var (
    someRobot Robot
  )

  BeforeEach(func() {
    someRobot = Robot{
      Work: func() {
      },
    }
  })

  Describe("Robot", func() {
    Context("when valid", func() {
      PIt("should Start", func() {
        Expect(true)
      })
      PIt("should initConnections", func() {
        Expect(true)
      })
      PIt("should initDevices", func() {
        Expect(true)
      })
      PIt("should startConnections", func() {
        Expect(true)
      })
      PIt("should startDevices", func() {
        Expect(true)
      })
      PIt("should GetDevices", func() {
        Expect(true)
      })
      PIt("should GetDevice", func() {
        Expect(true)
      })
    })
  })
})
<|endoftext|>


<commit_before>// +build proto

// RPC-based remote cache. Similar to HTTP but likely higher performance.
package cache

import (
	"bytes"
	"core"
	"encoding/base64"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"

	pb "cache/proto/rpc_cache"
)

type rpcCache struct {
	connection *grpc.ClientConn
	client     pb.RpcCacheClient
	Writeable  bool
	Connected  bool
	Connecting bool
	OSName     string
}

func (cache *rpcCache) Store(target *core.BuildTarget, key []byte) {
	if cache.isConnected() && cache.Writeable {
		log.Debug("Storing %s in RPC cache...", target.Label)
		artifacts := []*pb.Artifact{}
		for out := range cacheArtifacts(target) {
			artifacts2, err := cache.loadArtifacts(target, out)
			if err != nil {
				log.Warning("RPC cache failed to load artifact %s: %s", out, err)
				return
			}
			artifacts = append(artifacts, artifacts2...)
		}
		cache.sendArtifacts(target, key, artifacts)
	}
}

func (cache *rpcCache) StoreExtra(target *core.BuildTarget, key []byte, file string) {
	if cache.isConnected() && cache.Writeable {
		log.Debug("Storing %s : %s in RPC cache...", target.Label, file)
		artifacts, err := cache.loadArtifacts(target, file)
		if err != nil {
			log.Warning("RPC cache failed to load artifact %s: %s", file, err)
			return
		}
		cache.sendArtifacts(target, key, artifacts)
	}
}

func (cache *rpcCache) loadArtifacts(target *core.BuildTarget, file string) ([]*pb.Artifact, error) {
	artifacts := []*pb.Artifact{}
	outDir := target.OutDir()
	root := path.Join(outDir, file)
	err := filepath.Walk(root, func(name string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		} else if !info.IsDir() {
			content, err := ioutil.ReadFile(name)
			if err != nil {
				return err
			}
			artifacts = append(artifacts, &pb.Artifact{
				Package: target.Label.PackageName,
				Target:  target.Label.Name,
				File:    name[len(outDir)+1 : len(name)],
				Body:    content,
			})
		}
		return nil
	})
	return artifacts, err
}

func (cache *rpcCache) sendArtifacts(target *core.BuildTarget, key []byte, artifacts []*pb.Artifact) {
	req := pb.StoreRequest{Artifacts: artifacts, Hash: key, Os: runtime.GOOS, Arch: runtime.GOARCH}
	resp, err := cache.client.Store(context.Background(), &req)
	if err != nil {
		log.Warning("Error communicating with RPC cache server: %s", err)
	} else if !resp.Success {
		log.Warning("Failed to store artifacts in RPC cache for %s", target.Label)
	}
}

func (cache *rpcCache) Retrieve(target *core.BuildTarget, key []byte) bool {
	if !cache.isConnected() {
		return false
	}
	req := pb.RetrieveRequest{Hash: key, Os: runtime.GOOS, Arch: runtime.GOARCH}
	for out := range cacheArtifacts(target) {
		artifact := pb.Artifact{Package: target.Label.PackageName, Target: target.Label.Name, File: out}
		req.Artifacts = append(req.Artifacts, &artifact)
	}
	// We can't tell from here if retrieval has been successful for a target with no outputs.
	// This is kind of weird but not actually disallowed, and we already have a test case for it,
	// so might as well try to get it right here.
	if len(req.Artifacts) == 0 {
		return false
	}
	return cache.retrieveArtifacts(target, &req)
}

func (cache *rpcCache) RetrieveExtra(target *core.BuildTarget, key []byte, file string) bool {
	if !cache.isConnected() {
		return false
	}
	artifact := pb.Artifact{Package: target.Label.PackageName, Target: target.Label.Name, File: file}
	artifacts := []*pb.Artifact{&artifact}
	req := pb.RetrieveRequest{Hash: key, Os: runtime.GOOS, Arch: runtime.GOARCH, Artifacts: artifacts}
	return cache.retrieveArtifacts(target, &req)
}

func (cache *rpcCache) retrieveArtifacts(target *core.BuildTarget, req *pb.RetrieveRequest) bool {
	response, err := cache.client.Retrieve(context.Background(), req)
	if err != nil {
		log.Warning("Failed to retrieve artifacts for %s", target.Label)
		return false
	} else if !response.Success {
		// Quiet, this is almost certainly just a 'not found'
		log.Debug("Couldn't retrieve artifacts for %s [key %s] from RPC cache", target.Label, base64.RawURLEncoding.EncodeToString(req.Hash))
		return false
	}
	for _, artifact := range response.Artifacts {
		if !cache.writeFile(target, artifact.File, artifact.Body) {
			return false
		}
	}
	return true
}

func (cache *rpcCache) writeFile(target *core.BuildTarget, file string, body []byte) bool {
	out := path.Join(target.OutDir(), file)
	if err := core.WriteFile(bytes.NewReader(body), out, fileMode(target)); err != nil {
		log.Warning("RPC cache failed to write file %s", err)
		return false
	}
	log.Debug("Retrieved %s - %s from RPC cache", target.Label, file)
	return true
}

func (cache *rpcCache) Clean(target *core.BuildTarget) {
	if cache.isConnected() && cache.Writeable {
		req := pb.DeleteRequest{Os: runtime.GOOS, Arch: runtime.GOARCH}
		artifact := pb.Artifact{Package: target.Label.PackageName, Target: target.Label.Name}
		req.Artifacts = []*pb.Artifact{&artifact}
		response, err := cache.client.Delete(context.Background(), &req)
		if err != nil || !response.Success {
			log.Errorf("Failed to remove %s from RPC cache", target.Label)
		}
	}
}

func (cache *rpcCache) connect(config *core.Configuration) {
	// Change grpc to log using our implementation
	grpclog.SetLogger(&grpcLogMabob{})
	log.Info("Connecting to RPC cache at %s", config.Cache.RpcUrl)
	// TODO(pebers): Add support for communicating over https.
	connection, err := grpc.Dial(config.Cache.RpcUrl, grpc.WithInsecure(),
		grpc.WithTimeout(time.Duration(config.Cache.RpcTimeout)*time.Second))
	if err != nil {
		cache.Connecting = false
		log.Warning("Failed to connect to RPC cache: %s", err)
		return
	}
	// Note that we have to actually send it a message here to validate the connection;
	// Dial() only seems to return errors for superficial failures like syntactically invalid addresses,
	// it will return essentially immediately even if the server doesn't exist.
	healthclient := healthpb.NewHealthClient(connection)
	resp, err := healthclient.Check(context.Background(), &healthpb.HealthCheckRequest{Service: "plz-rpc-cache"})
	if err != nil {
		cache.Connecting = false
		log.Warning("Failed to contact RPC cache: %s", err)
	} else if resp.Status != healthpb.HealthCheckResponse_SERVING {
		cache.Connecting = false
		log.Warning("RPC cache says it is not serving (%d)", resp.Status)
	} else {
		cache.connection = connection
		cache.client = pb.NewRpcCacheClient(connection)
		cache.Connected = true
		cache.Connecting = false
		log.Info("RPC cache connected")
	}
}

// isConnected checks if the cache is connected. If it's still trying to connect it allows a
// very brief wait to give it a chance to come online.
func (cache *rpcCache) isConnected() bool {
	if cache.Connected {
		return true
	} else if !cache.Connecting {
		return false
	}
	ticker := time.NewTicker(10 * time.Millisecond)
	for i := 0; i < 5 && cache.Connecting; i++ {
		<-ticker.C
	}
	ticker.Stop()
	return cache.Connected
}

func newRpcCache(config *core.Configuration) (*rpcCache, error) {
	cache := &rpcCache{Writeable: config.Cache.RpcWriteable, Connecting: true}
	go cache.connect(config)
	return cache, nil
}

// grpcLogMabob is an implementation of grpc's logging interface using our backend.
type grpcLogMabob struct{}

func (g *grpcLogMabob) Fatal(args ...interface{})                 { log.Fatal(args...) }
func (g *grpcLogMabob) Fatalf(format string, args ...interface{}) { log.Fatalf(format, args...) }
func (g *grpcLogMabob) Fatalln(args ...interface{})               { log.Fatal(args...) }
func (g *grpcLogMabob) Print(args ...interface{})                 { log.Warning("%s", args) }
func (g *grpcLogMabob) Printf(format string, args ...interface{}) { log.Warning(format, args...) }
func (g *grpcLogMabob) Println(args ...interface{})               { log.Warning("%s", args) }
<commit_msg>Small sanity check on cache retrieval<commit_after>// +build proto

// RPC-based remote cache. Similar to HTTP but likely higher performance.
package cache

import (
	"bytes"
	"core"
	"encoding/base64"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"

	pb "cache/proto/rpc_cache"
)

type rpcCache struct {
	connection *grpc.ClientConn
	client     pb.RpcCacheClient
	Writeable  bool
	Connected  bool
	Connecting bool
	OSName     string
}

func (cache *rpcCache) Store(target *core.BuildTarget, key []byte) {
	if cache.isConnected() && cache.Writeable {
		log.Debug("Storing %s in RPC cache...", target.Label)
		artifacts := []*pb.Artifact{}
		for out := range cacheArtifacts(target) {
			artifacts2, err := cache.loadArtifacts(target, out)
			if err != nil {
				log.Warning("RPC cache failed to load artifact %s: %s", out, err)
				return
			}
			artifacts = append(artifacts, artifacts2...)
		}
		cache.sendArtifacts(target, key, artifacts)
	}
}

func (cache *rpcCache) StoreExtra(target *core.BuildTarget, key []byte, file string) {
	if cache.isConnected() && cache.Writeable {
		log.Debug("Storing %s : %s in RPC cache...", target.Label, file)
		artifacts, err := cache.loadArtifacts(target, file)
		if err != nil {
			log.Warning("RPC cache failed to load artifact %s: %s", file, err)
			return
		}
		cache.sendArtifacts(target, key, artifacts)
	}
}

func (cache *rpcCache) loadArtifacts(target *core.BuildTarget, file string) ([]*pb.Artifact, error) {
	artifacts := []*pb.Artifact{}
	outDir := target.OutDir()
	root := path.Join(outDir, file)
	err := filepath.Walk(root, func(name string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		} else if !info.IsDir() {
			content, err := ioutil.ReadFile(name)
			if err != nil {
				return err
			}
			artifacts = append(artifacts, &pb.Artifact{
				Package: target.Label.PackageName,
				Target:  target.Label.Name,
				File:    name[len(outDir)+1 : len(name)],
				Body:    content,
			})
		}
		return nil
	})
	return artifacts, err
}

func (cache *rpcCache) sendArtifacts(target *core.BuildTarget, key []byte, artifacts []*pb.Artifact) {
	req := pb.StoreRequest{Artifacts: artifacts, Hash: key, Os: runtime.GOOS, Arch: runtime.GOARCH}
	resp, err := cache.client.Store(context.Background(), &req)
	if err != nil {
		log.Warning("Error communicating with RPC cache server: %s", err)
	} else if !resp.Success {
		log.Warning("Failed to store artifacts in RPC cache for %s", target.Label)
	}
}

func (cache *rpcCache) Retrieve(target *core.BuildTarget, key []byte) bool {
	if !cache.isConnected() {
		return false
	}
	req := pb.RetrieveRequest{Hash: key, Os: runtime.GOOS, Arch: runtime.GOARCH}
	for out := range cacheArtifacts(target) {
		artifact := pb.Artifact{Package: target.Label.PackageName, Target: target.Label.Name, File: out}
		req.Artifacts = append(req.Artifacts, &artifact)
	}
	// We can't tell from here if retrieval has been successful for a target with no outputs.
	// This is kind of weird but not actually disallowed, and we already have a test case for it,
	// so might as well try to get it right here.
	if len(req.Artifacts) == 0 {
		return false
	}
	return cache.retrieveArtifacts(target, &req)
}

func (cache *rpcCache) RetrieveExtra(target *core.BuildTarget, key []byte, file string) bool {
	if !cache.isConnected() {
		return false
	}
	artifact := pb.Artifact{Package: target.Label.PackageName, Target: target.Label.Name, File: file}
	artifacts := []*pb.Artifact{&artifact}
	req := pb.RetrieveRequest{Hash: key, Os: runtime.GOOS, Arch: runtime.GOARCH, Artifacts: artifacts}
	return cache.retrieveArtifacts(target, &req)
}

func (cache *rpcCache) retrieveArtifacts(target *core.BuildTarget, req *pb.RetrieveRequest) bool {
	response, err := cache.client.Retrieve(context.Background(), req)
	if err != nil {
		log.Warning("Failed to retrieve artifacts for %s", target.Label)
		return false
	} else if !response.Success {
		// Quiet, this is almost certainly just a 'not found'
		log.Debug("Couldn't retrieve artifacts for %s [key %s] from RPC cache", target.Label, base64.RawURLEncoding.EncodeToString(req.Hash))
		return false
	}
	for _, artifact := range response.Artifacts {
		if !cache.writeFile(target, artifact.File, artifact.Body) {
			return false
		}
	}
	// Sanity check: if we don't get anything back, assume it probably wasn't really a success.
	return len(response.Artifacts) > 0
}

func (cache *rpcCache) writeFile(target *core.BuildTarget, file string, body []byte) bool {
	out := path.Join(target.OutDir(), file)
	if err := core.WriteFile(bytes.NewReader(body), out, fileMode(target)); err != nil {
		log.Warning("RPC cache failed to write file %s", err)
		return false
	}
	log.Debug("Retrieved %s - %s from RPC cache", target.Label, file)
	return true
}

func (cache *rpcCache) Clean(target *core.BuildTarget) {
	if cache.isConnected() && cache.Writeable {
		req := pb.DeleteRequest{Os: runtime.GOOS, Arch: runtime.GOARCH}
		artifact := pb.Artifact{Package: target.Label.PackageName, Target: target.Label.Name}
		req.Artifacts = []*pb.Artifact{&artifact}
		response, err := cache.client.Delete(context.Background(), &req)
		if err != nil || !response.Success {
			log.Errorf("Failed to remove %s from RPC cache", target.Label)
		}
	}
}

func (cache *rpcCache) connect(config *core.Configuration) {
	// Change grpc to log using our implementation
	grpclog.SetLogger(&grpcLogMabob{})
	log.Info("Connecting to RPC cache at %s", config.Cache.RpcUrl)
	// TODO(pebers): Add support for communicating over https.
	connection, err := grpc.Dial(config.Cache.RpcUrl, grpc.WithInsecure(),
		grpc.WithTimeout(time.Duration(config.Cache.RpcTimeout)*time.Second))
	if err != nil {
		cache.Connecting = false
		log.Warning("Failed to connect to RPC cache: %s", err)
		return
	}
	// Note that we have to actually send it a message here to validate the connection;
	// Dial() only seems to return errors for superficial failures like syntactically invalid addresses,
	// it will return essentially immediately even if the server doesn't exist.
	healthclient := healthpb.NewHealthClient(connection)
	resp, err := healthclient.Check(context.Background(), &healthpb.HealthCheckRequest{Service: "plz-rpc-cache"})
	if err != nil {
		cache.Connecting = false
		log.Warning("Failed to contact RPC cache: %s", err)
	} else if resp.Status != healthpb.HealthCheckResponse_SERVING {
		cache.Connecting = false
		log.Warning("RPC cache says it is not serving (%d)", resp.Status)
	} else {
		cache.connection = connection
		cache.client = pb.NewRpcCacheClient(connection)
		cache.Connected = true
		cache.Connecting = false
		log.Info("RPC cache connected")
	}
}

// isConnected checks if the cache is connected. If it's still trying to connect it allows a
// very brief wait to give it a chance to come online.
func (cache *rpcCache) isConnected() bool {
	if cache.Connected {
		return true
	} else if !cache.Connecting {
		return false
	}
	ticker := time.NewTicker(10 * time.Millisecond)
	for i := 0; i < 5 && cache.Connecting; i++ {
		<-ticker.C
	}
	ticker.Stop()
	return cache.Connected
}

func newRpcCache(config *core.Configuration) (*rpcCache, error) {
	cache := &rpcCache{Writeable: config.Cache.RpcWriteable, Connecting: true}
	go cache.connect(config)
	return cache, nil
}

// grpcLogMabob is an implementation of grpc's logging interface using our backend.
type grpcLogMabob struct{}

func (g *grpcLogMabob) Fatal(args ...interface{})                 { log.Fatal(args...) }
func (g *grpcLogMabob) Fatalf(format string, args ...interface{}) { log.Fatalf(format, args...) }
func (g *grpcLogMabob) Fatalln(args ...interface{})               { log.Fatal(args...) }
func (g *grpcLogMabob) Print(args ...interface{})                 { log.Warning("%s", args) }
func (g *grpcLogMabob) Printf(format string, args ...interface{}) { log.Warning(format, args...) }
func (g *grpcLogMabob) Println(args ...interface{})               { log.Warning("%s", args) }
<|endoftext|>


<commit_before>// goncurses - ncurses library for Go.
// Copyright 2011 Rob Thornton. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/* ncurses library

   1. No functions which operate only on stdscr have been implemented because
   it makes little sense to do so in a Go implementation. Stdscr is treated the
   same as any other window.

   2. Whenever possible, versions of ncurses functions which could potentially
   have a buffer overflow, like the getstr() family of functions, have not been
   implemented. Instead, only the mvwgetnstr() and wgetnstr() can be used. */
package goncurses

/*
#cgo LDFLAGS: -lncurses
#include <ncurses.h>
#include "goncurses.h"
*/
import "C"

import (
	"errors"
	"fmt"
	"unsafe"
)

// BaudRate returns the speed of the terminal in bits per second
func BaudRate() int {
	return int(C.baudrate())
}

// Beep requests the terminal make an audible bell or, if not available,
// flashes the screen. Note that screen flashing doesn't work on all
// terminals
func Beep() {
	C.beep()
}

// Turn on/off buffering; raw user signals are passed to the program for
// handling. Overrides raw mode
func CBreak(on bool) {
	if on {
		C.cbreak()
		return
	}
	C.nocbreak()
}

// Test whether colour values can be changed
func CanChangeColor() bool {
	return bool(C.bool(C.can_change_color()))
}

// Get RGB values for specified colour
func ColorContent(col int) (int, int, int) {
	var r, g, b C.short
	C.color_content(C.short(col), (*C.short)(&r), (*C.short)(&g),
		(*C.short)(&b))
	return int(r), int(g), int(b)
}

// Return the value of a color pair which can be passed to functions which
// accept attributes like AddChar or AttrOn/Off.
func ColorPair(pair int) int {
	return int(C.COLOR_PAIR(C.int(pair)))
}

// CursesVersion returns the version of the ncurses library currently linked to
func CursesVersion() string {
	return C.GoString(C.curses_version())
}

// Set the cursor visibility. Options are: 0 (invisible/hidden), 1 (normal)
// and 2 (extra-visible)
func Cursor(vis byte) error {
	if C.curs_set(C.int(vis)) == C.ERR {
		return errors.New("Failed to enable ")
	}
	return nil
}

// Echo turns on/off the printing of typed characters
func Echo(on bool) {
	if on {
		C.echo()
		return
	}
	C.noecho()
}

// Must be called prior to exiting the program in order to make sure the
// terminal returns to normal operation
func End() {
	C.endwin()
}

// Flash requests the terminal flashes the screen or, if not available,
// make an audible bell. Note that screen flashing doesn't work on all
// terminals
func Flash() {
	C.flash()
}

// FlushInput flushes all input
func FlushInput() error {
	if C.flushinp() == C.ERR {
		return errors.New("Flush input failed")
	}
	return nil
}

// Returns an array of integers representing the following, in order:
// x, y and z coordinates, id of the device, and a bit masked state of
// the devices buttons
func GetMouse() ([]int, error) {
	if bool(C.ncurses_has_mouse()) != true {
		return nil, errors.New("Mouse support not enabled")
	}
	var event C.MEVENT
	if C.getmouse(&event) != C.OK {
		return nil, errors.New("Failed to get mouse event")
	}
	return []int{int(event.x), int(event.y), int(event.z), int(event.id),
		int(event.bstate)}, nil
}

// Behaves like cbreak() but also adds a timeout for input. If timeout is
// exceeded after a call to Getch() has been made then GetChar will return
// with an error.
func HalfDelay(delay int) error {
	var cerr C.int
	if delay > 0 {
		cerr = C.halfdelay(C.int(delay))
	}
	if cerr == C.ERR {
		return errors.New("Unable to set delay mode")
	}
	return nil
}

// HasColors returns true if terminal can display colors
func HasColors() bool {
	return bool(C.has_colors())
}

// HasInsertCharacter return true if the terminal has insert and delete
// character capabilities
func HasInsertCharacter() bool {
	return bool(C.has_ic())
}

// HasInsertLine returns true if the terminal has insert and delete line
// capabilities. See ncurses documentation for more details
func HasInsertLine() bool {
	return bool(C.has_il())
}

// HasKey returns true if terminal recognized the given character
func HasKey(ch Key) bool {
	if C.has_key(C.int(ch)) == 1 {
		return true
	}
	return false
}

// InitColor is used to set 'color' to the specified RGB values. Values may
// be between 0 and 1000.
func InitColor(col int, r, g, b int) error {
	if C.init_color(C.short(col), C.short(r), C.short(g),
		C.short(b)) == C.ERR {
		return errors.New("Failed to set new color definition")
	}
	return nil
}

// InitPair sets a colour pair designated by 'pair' to fg and bg colors
func InitPair(pair byte, fg, bg int) error {
	if pair == 0 || C.int(pair) > (C.COLOR_PAIRS-1) {
		return errors.New("Invalid color pair selected")
	}
	if C.init_pair(C.short(pair), C.short(fg), C.short(bg)) == C.ERR {
		return errors.New("Failed to init color pair")
	}
	return nil
}

// Initialize the ncurses library. You must run this function prior to any
// other goncurses function in order for the library to work
func Init() (stdscr Window, err error) {
	stdscr = Window{C.initscr()}
	if unsafe.Pointer(stdscr.win) == nil {
		err = errors.New("An error occurred initializing ncurses")
	}
	return
}

// IsEnd returns true if End() has been called, otherwise false
func IsEnd() bool {
	return bool(C.isendwin())
}

// IsTermResized returns true if ResizeTerm would modify any current Windows
// if called with the given parameters
func IsTermResized(nlines, ncols int) bool {
	return bool(C.is_term_resized(C.int(nlines), C.int(ncols)))
}

// Returns a string representing the value of input returned by Getch
func KeyString(k Key) string {
	key, ok := keyList[k]
	if !ok {
		key = fmt.Sprintf("%c", int(k))
	}
	return key
}

func Mouse() bool {
	return bool(C.ncurses_has_mouse())
}

func MouseInterval() {
}

// MouseMask accepts a single int of OR'd mouse events. If a mouse event
// is triggered, GetChar() will return KEY_MOUSE. To retrieve the actual
// event use GetMouse() to pop it off the queue. Pass a pointer as the
// second argument to store the prior events being monitored or nil.
func MouseMask(mask int, old *int) (m int) {
	if bool(C.ncurses_has_mouse()) {
		m = int(C.mousemask((C.mmask_t)(mask),
			(*C.mmask_t)(unsafe.Pointer(old))))
	}
	return
}

// NapMilliseconds is used to sleep for ms milliseconds
func NapMilliseconds(ms int) {
	C.napms(C.int(ms))
}

// NewWindow creates a window of size h(eight) and w(idth) at y, x
func NewWindow(h, w, y, x int) (window Window, err error) {
	window = Window{C.newwin(C.int(h), C.int(w), C.int(y), C.int(x))}
	if window.win == nil {
		err = errors.New("Failed to create a new window")
	}
	return
}

// NL turns newline translation on/off.
func NL(on bool) {
	if on {
		C.nl()
		return
	}
	C.nonl()
}

// Raw turns on input buffering; user signals are disabled and the key strokes
// are passed directly to input. Set to false if you wish to turn this mode
// off
func Raw(on bool) {
	if on {
		C.raw()
		return
	}
	C.noraw()
}

// ResizeTerm will attempt to resize the terminal. This only has an effect if
// the terminal is in an XWindows (GUI) environment.
func ResizeTerm(nlines, ncols int) error {
	if C.resizeterm(C.int(nlines), C.int(ncols)) == C.ERR {
		return errors.New("Failed to resize terminal")
	}
	return nil
}

// Enables colors to be displayed. Will return an error if terminal is not
// capable of displaying colors
func StartColor() error {
	if C.has_colors() == C.bool(false) {
		return errors.New("Terminal does not support colors")
	}
	if C.start_color() == C.ERR {
		return errors.New("Failed to enable color mode")
	}
	return nil
}

// Update the screen, refreshing all windows
func Update() error {
	if C.doupdate() == C.ERR {
		return errors.New("Failed to update")
	}
	return nil
}

// UseEnvironment specifies whether the LINES and COLUMNS environmental
// variables should be used or not
func UseEnvironment(use bool) {
	C.use_env(C.bool(use))
}
<commit_msg>Implement UnGetChar<commit_after>// goncurses - ncurses library for Go.
// Copyright 2011 Rob Thornton. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/* ncurses library

   1. No functions which operate only on stdscr have been implemented because
   it makes little sense to do so in a Go implementation. Stdscr is treated the
   same as any other window.

   2. Whenever possible, versions of ncurses functions which could potentially
   have a buffer overflow, like the getstr() family of functions, have not been
   implemented. Instead, only the mvwgetnstr() and wgetnstr() can be used. */
package goncurses

/*
#cgo LDFLAGS: -lncurses
#include <ncurses.h>
#include "goncurses.h"
*/
import "C"

import (
	"errors"
	"fmt"
	"unsafe"
)

// BaudRate returns the speed of the terminal in bits per second
func BaudRate() int {
	return int(C.baudrate())
}

// Beep requests the terminal make an audible bell or, if not available,
// flashes the screen. Note that screen flashing doesn't work on all
// terminals
func Beep() {
	C.beep()
}

// Turn on/off buffering; raw user signals are passed to the program for
// handling. Overrides raw mode
func CBreak(on bool) {
	if on {
		C.cbreak()
		return
	}
	C.nocbreak()
}

// Test whether colour values can be changed
func CanChangeColor() bool {
	return bool(C.bool(C.can_change_color()))
}

// Get RGB values for specified colour
func ColorContent(col int) (int, int, int) {
	var r, g, b C.short
	C.color_content(C.short(col), (*C.short)(&r), (*C.short)(&g),
		(*C.short)(&b))
	return int(r), int(g), int(b)
}

// Return the value of a color pair which can be passed to functions which
// accept attributes like AddChar or AttrOn/Off.
func ColorPair(pair int) int {
	return int(C.COLOR_PAIR(C.int(pair)))
}

// CursesVersion returns the version of the ncurses library currently linked to
func CursesVersion() string {
	return C.GoString(C.curses_version())
}

// Set the cursor visibility. Options are: 0 (invisible/hidden), 1 (normal)
// and 2 (extra-visible)
func Cursor(vis byte) error {
	if C.curs_set(C.int(vis)) == C.ERR {
		return errors.New("Failed to enable ")
	}
	return nil
}

// Echo turns on/off the printing of typed characters
func Echo(on bool) {
	if on {
		C.echo()
		return
	}
	C.noecho()
}

// Must be called prior to exiting the program in order to make sure the
// terminal returns to normal operation
func End() {
	C.endwin()
}

// Flash requests the terminal flashes the screen or, if not available,
// make an audible bell. Note that screen flashing doesn't work on all
// terminals
func Flash() {
	C.flash()
}

// FlushInput flushes all input
func FlushInput() error {
	if C.flushinp() == C.ERR {
		return errors.New("Flush input failed")
	}
	return nil
}

// Returns an array of integers representing the following, in order:
// x, y and z coordinates, id of the device, and a bit masked state of
// the devices buttons
func GetMouse() ([]int, error) {
	if bool(C.ncurses_has_mouse()) != true {
		return nil, errors.New("Mouse support not enabled")
	}
	var event C.MEVENT
	if C.getmouse(&event) != C.OK {
		return nil, errors.New("Failed to get mouse event")
	}
	return []int{int(event.x), int(event.y), int(event.z), int(event.id),
		int(event.bstate)}, nil
}

// Behaves like cbreak() but also adds a timeout for input. If timeout is
// exceeded after a call to Getch() has been made then GetChar will return
// with an error.
func HalfDelay(delay int) error {
	var cerr C.int
	if delay > 0 {
		cerr = C.halfdelay(C.int(delay))
	}
	if cerr == C.ERR {
		return errors.New("Unable to set delay mode")
	}
	return nil
}

// HasColors returns true if terminal can display colors
func HasColors() bool {
	return bool(C.has_colors())
}

// HasInsertCharacter return true if the terminal has insert and delete
// character capabilities
func HasInsertCharacter() bool {
	return bool(C.has_ic())
}

// HasInsertLine returns true if the terminal has insert and delete line
// capabilities. See ncurses documentation for more details
func HasInsertLine() bool {
	return bool(C.has_il())
}

// HasKey returns true if terminal recognized the given character
func HasKey(ch Key) bool {
	if C.has_key(C.int(ch)) == 1 {
		return true
	}
	return false
}

// InitColor is used to set 'color' to the specified RGB values. Values may
// be between 0 and 1000.
func InitColor(col int, r, g, b int) error {
	if C.init_color(C.short(col), C.short(r), C.short(g),
		C.short(b)) == C.ERR {
		return errors.New("Failed to set new color definition")
	}
	return nil
}

// InitPair sets a colour pair designated by 'pair' to fg and bg colors
func InitPair(pair byte, fg, bg int) error {
	if pair == 0 || C.int(pair) > (C.COLOR_PAIRS-1) {
		return errors.New("Invalid color pair selected")
	}
	if C.init_pair(C.short(pair), C.short(fg), C.short(bg)) == C.ERR {
		return errors.New("Failed to init color pair")
	}
	return nil
}

// Initialize the ncurses library. You must run this function prior to any
// other goncurses function in order for the library to work
func Init() (stdscr Window, err error) {
	stdscr = Window{C.initscr()}
	if unsafe.Pointer(stdscr.win) == nil {
		err = errors.New("An error occurred initializing ncurses")
	}
	return
}

// IsEnd returns true if End() has been called, otherwise false
func IsEnd() bool {
	return bool(C.isendwin())
}

// IsTermResized returns true if ResizeTerm would modify any current Windows
// if called with the given parameters
func IsTermResized(nlines, ncols int) bool {
	return bool(C.is_term_resized(C.int(nlines), C.int(ncols)))
}

// Returns a string representing the value of input returned by Getch
func KeyString(k Key) string {
	key, ok := keyList[k]
	if !ok {
		key = fmt.Sprintf("%c", int(k))
	}
	return key
}

func Mouse() bool {
	return bool(C.ncurses_has_mouse())
}

func MouseInterval() {
}

// MouseMask accepts a single int of OR'd mouse events. If a mouse event
// is triggered, GetChar() will return KEY_MOUSE. To retrieve the actual
// event use GetMouse() to pop it off the queue. Pass a pointer as the
// second argument to store the prior events being monitored or nil.
func MouseMask(mask int, old *int) (m int) {
	if bool(C.ncurses_has_mouse()) {
		m = int(C.mousemask((C.mmask_t)(mask),
			(*C.mmask_t)(unsafe.Pointer(old))))
	}
	return
}

// NapMilliseconds is used to sleep for ms milliseconds
func NapMilliseconds(ms int) {
	C.napms(C.int(ms))
}

// NewWindow creates a window of size h(eight) and w(idth) at y, x
func NewWindow(h, w, y, x int) (window Window, err error) {
	window = Window{C.newwin(C.int(h), C.int(w), C.int(y), C.int(x))}
	if window.win == nil {
		err = errors.New("Failed to create a new window")
	}
	return
}

// NL turns newline translation on/off.
func NL(on bool) {
	if on {
		C.nl()
		return
	}
	C.nonl()
}

// Raw turns on input buffering; user signals are disabled and the key strokes
// are passed directly to input. Set to false if you wish to turn this mode
// off
func Raw(on bool) {
	if on {
		C.raw()
		return
	}
	C.noraw()
}

// ResizeTerm will attempt to resize the terminal. This only has an effect if
// the terminal is in an XWindows (GUI) environment.
func ResizeTerm(nlines, ncols int) error {
	if C.resizeterm(C.int(nlines), C.int(ncols)) == C.ERR {
		return errors.New("Failed to resize terminal")
	}
	return nil
}

// Enables colors to be displayed. Will return an error if terminal is not
// capable of displaying colors
func StartColor() error {
	if C.has_colors() == C.bool(false) {
		return errors.New("Terminal does not support colors")
	}
	if C.start_color() == C.ERR {
		return errors.New("Failed to enable color mode")
	}
	return nil
}

// UnGetChar places the character back into the input queue
func UnGetChar(ch Character) {
	C.ungetch(C.int(ch))
}

// Update the screen, refreshing all windows
func Update() error {
	if C.doupdate() == C.ERR {
		return errors.New("Failed to update")
	}
	return nil
}

// UseEnvironment specifies whether the LINES and COLUMNS environmental
// variables should be used or not
func UseEnvironment(use bool) {
	C.use_env(C.bool(use))
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/gob"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/kr/pty"
	"github.com/pkg/term"
)

func link(socketPath string) {
	conn, err := net.Dial("unix", socketPath)
	if err != nil {
		fatal(err)
	}

	var b [2048]byte
	var oob [2048]byte

	n, oobn, _, _, err := conn.(*net.UnixConn).ReadMsgUnix(b[:], oob[:])
	if err != nil {
		log.Fatalln("failed to read unix msg:", err, n, oobn)
	}

	scms, err := syscall.ParseSocketControlMessage(oob[:oobn])
	if err != nil {
		log.Fatalln("failed to parse socket control message:", err)
	}

	if len(scms) < 1 {
		log.Fatalln("no socket control messages sent")
	}

	scm := scms[0]

	fds, err := syscall.ParseUnixRights(&scm)
	if err != nil {
		log.Fatalln("failed to parse unix rights", err)
	}

	if len(fds) != 3 {
		log.Fatalln("invalid number of fds; need 3, got", len(fds))
	}

	stdout := os.NewFile(uintptr(fds[0]), "stdout")
	stderr := os.NewFile(uintptr(fds[1]), "stderr")
	status := os.NewFile(uintptr(fds[2]), "status")

	streaming := &sync.WaitGroup{}

	inputWriter := &inputWriter{gob.NewEncoder(conn)}

	var input io.Reader

	if *tty {
		term, err := term.Open(os.Stdin.Name())
		if err == nil {
			err = term.SetRaw()
			if err == nil {
				input = term
				defer term.Restore()
			}
		}

		resized := make(chan os.Signal, 10)

		go func() {
			for {
				<-resized

				rows, cols, err := pty.Getsize(os.Stdin)
				if err == nil {
					inputWriter.SetWindowSize(cols, rows)
				}
			}
		}()

		signal.Notify(resized, syscall.SIGWINCH)
	} else {
		input = os.Stdin
	}

	// do not add stdin to the waitgroup; it appears to cause things to hang.
	// doesn't make much sense anyway; if stdout/stderr closed we probably
	// can't write any more to stdin in the first place.
	go func() {
		io.Copy(inputWriter, input)
		inputWriter.Close()
	}()

	streaming.Add(1)
	go func() {
		io.Copy(os.Stdout, stdout)
		stdout.Close()
		streaming.Done()
	}()

	streaming.Add(1)
	go func() {
		io.Copy(os.Stderr, stderr)
		stderr.Close()
		streaming.Done()
	}()

	streaming.Wait()

	var exitStatus int
	_, err = fmt.Fscanf(status, "%d\n", &exitStatus)
	if err != nil {
		log.Fatalln("failed to read status:", err)
	}

	os.Exit(exitStatus)
}
<commit_msg>remove setting tty to raw when linking<commit_after>package main

import (
	"encoding/gob"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/kr/pty"
)

func link(socketPath string) {
	conn, err := net.Dial("unix", socketPath)
	if err != nil {
		fatal(err)
	}

	var b [2048]byte
	var oob [2048]byte

	n, oobn, _, _, err := conn.(*net.UnixConn).ReadMsgUnix(b[:], oob[:])
	if err != nil {
		log.Fatalln("failed to read unix msg:", err, n, oobn)
	}

	scms, err := syscall.ParseSocketControlMessage(oob[:oobn])
	if err != nil {
		log.Fatalln("failed to parse socket control message:", err)
	}

	if len(scms) < 1 {
		log.Fatalln("no socket control messages sent")
	}

	scm := scms[0]

	fds, err := syscall.ParseUnixRights(&scm)
	if err != nil {
		log.Fatalln("failed to parse unix rights", err)
	}

	if len(fds) != 3 {
		log.Fatalln("invalid number of fds; need 3, got", len(fds))
	}

	stdout := os.NewFile(uintptr(fds[0]), "stdout")
	stderr := os.NewFile(uintptr(fds[1]), "stderr")
	status := os.NewFile(uintptr(fds[2]), "status")

	streaming := &sync.WaitGroup{}

	inputWriter := &inputWriter{gob.NewEncoder(conn)}

	resized := make(chan os.Signal, 10)

	go func() {
		for {
			<-resized

			rows, cols, err := pty.Getsize(os.Stdin)
			if err == nil {
				inputWriter.SetWindowSize(cols, rows)
			}
		}
	}()

	signal.Notify(resized, syscall.SIGWINCH)

	// do not add stdin to the waitgroup; it appears to cause things to hang.
	// doesn't make much sense anyway; if stdout/stderr closed we probably
	// can't write any more to stdin in the first place.
	go func() {
		io.Copy(inputWriter, os.Stdin)
		inputWriter.Close()
	}()

	streaming.Add(1)
	go func() {
		io.Copy(os.Stdout, stdout)
		stdout.Close()
		streaming.Done()
	}()

	streaming.Add(1)
	go func() {
		io.Copy(os.Stderr, stderr)
		stderr.Close()
		streaming.Done()
	}()

	streaming.Wait()

	var exitStatus int
	_, err = fmt.Fscanf(status, "%d\n", &exitStatus)
	if err != nil {
		log.Fatalln("failed to read status:", err)
	}

	os.Exit(exitStatus)
}
<|endoftext|>