<commit_before>package user

import (
	"fmt"
	jwt "github.com/dgrijalva/jwt-go"
	"github.com/gin-gonic/gin"

	e "github.com/eirka/eirka-libs/errors"
)

// holds the hmac secret, is set from main
var Secret string

// checks for session cookie and handles permissions
func Auth(authenticated bool) gin.HandlerFunc {
	return func(c *gin.Context) {

		// error if theres no secret set
		if Secret == "" {
			c.JSON(e.ErrorMessage(e.ErrInternalError))
			c.Error(e.ErrNoSecret).SetMeta("auth.Auth")
			c.Abort()
			return
		}

		// set default anonymous user
		user := DefaultUser()

		// parse jwt token if its there
		token, err := jwt.ParseFromRequest(c.Request, func(token *jwt.Token) (interface{}, error) {

			// check alg to make sure its hmac
			_, ok := token.Method.(*jwt.SigningMethodHMAC)
			if !ok {
				return nil, fmt.Errorf("Unexpected signing method: %v", token.Header["alg"])
			}

			// get the issuer from claims
			token_issuer, ok := token.Claims[jwt_claim_issuer].(string)
			if !ok {
				return nil, fmt.Errorf("Couldnt find issuer")
			}

			// check the issuer
			if token_issuer != jwt_issuer {
				return nil, fmt.Errorf("Incorrect issuer")
			}

			// get uid from token
			token_uid, ok := token.Claims[jwt_claim_user_id].(float64)
			if !ok {
				return nil, fmt.Errorf("Couldnt find user id")
			}

			// set the user id
			user.SetId(uint(token_uid))
			// set authenticated
			user.SetAuthenticated()

			// check that the generated user is valid
			if !user.IsValid() {
				return nil, fmt.Errorf("Generated invalid user")
			}

			// compare with secret from settings
			return []byte(Secret), nil

		})
		// if theres some jwt error other than no token in request or the token is invalid then return unauth
		if err != nil && err != jwt.ErrNoTokenInRequest || token != nil && !token.Valid {
			c.JSON(e.ErrorMessage(e.ErrUnauthorized))
			c.Error(err).SetMeta("user.Auth")
			c.Abort()
			return
		}

		// check if user needed to be authenticated
		if authenticated && !user.IsAuthenticated {
			c.JSON(e.ErrorMessage(e.ErrUnauthorized))
			c.Error(e.ErrUnauthorized).SetMeta("user.Auth")
			c.Abort()
			return
		}

		// set user data
		c.Set("userdata", user)

		c.Next()

	}

}
<commit_msg>simplify logic in auth<commit_after>package user

import (
	"fmt"
	jwt "github.com/dgrijalva/jwt-go"
	"github.com/gin-gonic/gin"

	e "github.com/eirka/eirka-libs/errors"
)

// holds the hmac secret, is set from main
var Secret string

// checks for session cookie and handles permissions
func Auth(authenticated bool) gin.HandlerFunc {
	return func(c *gin.Context) {

		// error if theres no secret set
		if Secret == "" {
			c.JSON(e.ErrorMessage(e.ErrInternalError))
			c.Error(e.ErrNoSecret).SetMeta("auth.Auth")
			c.Abort()
			return
		}

		// set default anonymous user
		user := DefaultUser()

		// parse jwt token if its there
		token, err := jwt.ParseFromRequest(c.Request, func(token *jwt.Token) (interface{}, error) {

			// check alg to make sure its hmac
			_, ok := token.Method.(*jwt.SigningMethodHMAC)
			if !ok {
				return nil, fmt.Errorf("Unexpected signing method: %v", token.Header["alg"])
			}

			// get the issuer from claims
			token_issuer, ok := token.Claims[jwt_claim_issuer].(string)
			if !ok {
				return nil, fmt.Errorf("Couldnt find issuer")
			}

			// check the issuer
			if token_issuer != jwt_issuer {
				return nil, fmt.Errorf("Incorrect issuer")
			}

			// get uid from token
			token_uid, ok := token.Claims[jwt_claim_user_id].(float64)
			if !ok {
				return nil, fmt.Errorf("Couldnt find user id")
			}

			// set the user id
			user.SetId(uint(token_uid))
			// set authenticated
			user.SetAuthenticated()

			// check that the generated user is valid
			if !user.IsValid() {
				return nil, fmt.Errorf("Generated invalid user")
			}

			// compare with secret from settings
			return []byte(Secret), nil

		})
		// if theres some jwt error other than no token in request or the token is invalid then return unauth
		if err != nil && err != jwt.ErrNoTokenInRequest || token != nil && !token.Valid {
			c.JSON(e.ErrorMessage(e.ErrUnauthorized))
			c.Error(err).SetMeta("user.Auth")
			c.Abort()
			return
		}

		// check if user needed to be authenticated
		if user.IsAuthenticated != authenticated {
			c.JSON(e.ErrorMessage(e.ErrUnauthorized))
			c.Error(e.ErrUnauthorized).SetMeta("user.Auth")
			c.Abort()
			return
		}

		// set user data
		c.Set("userdata", user)

		c.Next()

	}

}
<|endoftext|>


<commit_before>package util

import (
	"fmt"
	"github.com/qjpcpu/sesh/batch"
	"github.com/qjpcpu/sesh/cowsay"
	"github.com/qjpcpu/sesh/dircat"
	"github.com/qjpcpu/sesh/golang.org/x/crypto/ssh/terminal"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"sync"
	"time"
)

type TaskArgs struct {
	User     string
	Password string
	Keyfile  string
	Port     string
	Cmd      string
	CmdArgs  string
	Timeout  int
	// parallel run
	Output    io.Writer
	ErrOutput io.Writer
	// scp
	Parallel bool
	Source   string
	Destdir  string
}

func GirlSay(content ...interface{}) string {
	return cowsay.Format(fmt.Sprint(content))
}

// Hook for per task state changed
func report(output io.Writer, prefix, host string, color bool) {
	if color {
		output.Write([]byte(fmt.Sprintf("\033[33m%s========== %s ==========\033[0m\n", prefix, host)))
	} else {
		output.Write([]byte(fmt.Sprintf("%s========== %s ==========\n", prefix, host)))
	}
}

func SerialRun(config TaskArgs, raw_host_arr []string, start, end int) error {
	host_arr := raw_host_arr[start:end]
	user := config.User
	pwd := config.Password
	keyfile := config.Keyfile
	cmd := config.Cmd
	args := config.CmdArgs
	timeout := config.Timeout
	port := config.Port
	// Format command
	cmd = format_cmd(cmd, args)
	printer := config.Output
	err_printer := config.ErrOutput

	BuildKnownHosts(host_arr)

	wg := new(sync.WaitGroup)

	for index, h := range host_arr {
		task := batch.NewTask(h, user, pwd, keyfile, cmd, port, printer, err_printer, wg)
		task.Timeout = timeout
		report(err_printer, fmt.Sprintf("%d/%d ", index+1+start, len(raw_host_arr)), task.Host, true)
		task.Work()
	}
	wg.Wait()
	return nil
}

func ParallelRun(config TaskArgs, raw_host_arr []string, start, end int, tmpdir string) error {
	host_arr := raw_host_arr[start:end]
	user := config.User
	pwd := config.Password
	keyfile := config.Keyfile
	cmd := config.Cmd
	args := config.CmdArgs
	timeout := config.Timeout
	port := config.Port
	cmd = format_cmd(cmd, args)
	printer := config.Output
	err_printer := config.ErrOutput

	// Create master, the master is used to manage go routines
	wg := new(sync.WaitGroup)
	// Setup tmp directory for tmp files
	dir := fmt.Sprintf("%s/.task.%d", tmpdir, time.Now().Nanosecond())
	if err := os.Mkdir(dir, os.ModeDir|os.ModePerm); err != nil {
		return err
	}

	// Listen interrupt and kill signal, clear tmp files before exit.
	intqueue := make(chan os.Signal, 1)
	signal.Notify(intqueue, os.Interrupt, os.Kill)
	// If got interrupt or kill signal, delete tmp directory first, then exit with 1
	go func() {
		<-intqueue
		os.RemoveAll(dir)
		os.Exit(1)
	}()
	// If the complete all the tasks normlly, stop listenning signals and remove tmp directory
	defer func() {
		signal.Stop(intqueue)
		os.RemoveAll(dir)
	}()

	BuildKnownHosts(host_arr)

	// Create tmp file for every host, then executes.
	var tmpfiles []*os.File
	for _, h := range host_arr {
		file, _ := os.Create(fmt.Sprintf("%s/%s", dir, h))
		err_file, _ := os.Create(fmt.Sprintf("%s/%s.err", dir, h))
		tmpfiles = append(tmpfiles, file, err_file)
		task := batch.NewTask(h, user, pwd, keyfile, cmd, port, file, err_file, wg)
		task.Timeout = timeout
		go task.Work()
	}

	// show realtime view for each host
	var dc *dircat.DirCat
	if terminal.IsTerminal(1) {
		wlist := []string{}
		for _, h := range host_arr {
			wlist = append(wlist, fmt.Sprintf("%s/%s", dir, h))
		}
		dc, _ = dircat.Init(wlist...)
		go dc.Start()
	}
	// When a host is ready and request for continue, the master would echo CONTINUE for response to allow host to run
	wg.Wait()

	if terminal.IsTerminal(1) {
		dc.Stop()
	}
	// close tmp files
	for _, f := range tmpfiles {
		f.Close()
	}
	// Merge all the hosts' output to the output file
	for _, h := range host_arr {
		report(os.Stderr, "", h, true)
		// copy err output first
		err_fn := fmt.Sprintf("%s/%s.err", dir, h)
		err_src, _ := os.Open(err_fn)
		io.Copy(err_printer, err_src)
		err_src.Close()
		// copy output then
		fn := fmt.Sprintf("%s/%s", dir, h)
		src, _ := os.Open(fn)
		io.Copy(printer, src)
		src.Close()
		// remove tmp file
		os.Remove(err_fn)
		os.Remove(fn)
	}
	return nil
}
func Interact(config TaskArgs, host string) {
	wg := new(sync.WaitGroup)
	task := batch.NewTask(
		host,
		config.User,
		config.Password,
		config.Keyfile,
		config.Cmd,
		config.Port,
		config.Output,
		config.ErrOutput,
		wg,
	)
	task.SysLogin()
}

func ScpRun(config TaskArgs, host_arr []string) error {
	user := config.User
	parallel := config.Parallel
	src := config.Source
	dest := config.Destdir
	if !strings.HasSuffix(dest, "/") {
		dest += "/"
	}
	wg := new(sync.WaitGroup)
	BuildKnownHosts(host_arr)
	for _, h := range host_arr {
		wg.Add(1)
		taskF := func(host string) {
			defer wg.Done()
			var remote_shell string = ""
			if config.Port != "" && config.Port != "22" {
				remote_shell = fmt.Sprintf(` -e 'ssh -p %s' `, config.Port)
			}
			cmdstr := fmt.Sprintf(`rsync -azh %s %s %s@%s:%s`, remote_shell, src, user, host, dest)
			cmd := exec.Command("/bin/bash", "-c", cmdstr)
			fmt.Fprintf(os.Stderr, "Start sync to  %s......\n", host)
			if err := cmd.Run(); err != nil {
				fmt.Fprintf(os.Stderr, "Sync to %s:%s fail:%v\n", host, dest, err)
			} else {
				fmt.Fprintf(os.Stderr, "Sync to %s:%s OK\n", host, dest)
			}
		}
		if parallel {
			go taskF(h)
		} else {
			taskF(h)
		}
	}

	wg.Wait()
	return nil
}

func BuildKnownHosts(host_arr []string) {
	filename := os.Getenv("HOME") + "/.ssh/known_hosts"
	dirname := os.Getenv("HOME") + "/.ssh"
	content_bytes, _ := ioutil.ReadFile(filename)
	content := string(content_bytes)
	var not_exists []string
	for _, host := range host_arr {
		if !strings.Contains(content, host) {
			not_exists = append(not_exists, host)
		}
	}
	if len(not_exists) == 0 {
		return
	}
	cmdstr := fmt.Sprintf(`
[ ! -e %s ] && mkdir -p %s && chmod 700 %s
[ ! -e %s ] && touch %s && chmod 644 %s
ssh-keyscan %s >> %s
`, dirname, dirname, dirname,
		filename, filename, filename,
		strings.Join(not_exists, " "), filename)
	cmd := exec.Command("/bin/sh", "-c", cmdstr)
	cmd.Run()
}
<commit_msg>err output<commit_after>package util

import (
	"fmt"
	"github.com/qjpcpu/sesh/batch"
	"github.com/qjpcpu/sesh/cowsay"
	"github.com/qjpcpu/sesh/dircat"
	"github.com/qjpcpu/sesh/golang.org/x/crypto/ssh/terminal"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"sync"
	"time"
)

type TaskArgs struct {
	User     string
	Password string
	Keyfile  string
	Port     string
	Cmd      string
	CmdArgs  string
	Timeout  int
	// parallel run
	Output    io.Writer
	ErrOutput io.Writer
	// scp
	Parallel bool
	Source   string
	Destdir  string
}

func GirlSay(content ...interface{}) string {
	return cowsay.Format(fmt.Sprint(content))
}

// Hook for per task state changed
func report(output io.Writer, prefix, host string, color bool) {
	if color {
		output.Write([]byte(fmt.Sprintf("\033[33m%s========== %s ==========\033[0m\n", prefix, host)))
	} else {
		output.Write([]byte(fmt.Sprintf("%s========== %s ==========\n", prefix, host)))
	}
}

func SerialRun(config TaskArgs, raw_host_arr []string, start, end int) error {
	host_arr := raw_host_arr[start:end]
	user := config.User
	pwd := config.Password
	keyfile := config.Keyfile
	cmd := config.Cmd
	args := config.CmdArgs
	timeout := config.Timeout
	port := config.Port
	// Format command
	cmd = format_cmd(cmd, args)
	printer := config.Output
	err_printer := config.ErrOutput

	BuildKnownHosts(host_arr)

	wg := new(sync.WaitGroup)

	for index, h := range host_arr {
		task := batch.NewTask(h, user, pwd, keyfile, cmd, port, printer, err_printer, wg)
		task.Timeout = timeout
		report(err_printer, fmt.Sprintf("%d/%d ", index+1+start, len(raw_host_arr)), task.Host, true)
		task.Work()
	}
	wg.Wait()
	return nil
}

func ParallelRun(config TaskArgs, raw_host_arr []string, start, end int, tmpdir string) error {
	host_arr := raw_host_arr[start:end]
	user := config.User
	pwd := config.Password
	keyfile := config.Keyfile
	cmd := config.Cmd
	args := config.CmdArgs
	timeout := config.Timeout
	port := config.Port
	cmd = format_cmd(cmd, args)
	printer := config.Output
	err_printer := config.ErrOutput

	// Create master, the master is used to manage go routines
	wg := new(sync.WaitGroup)
	// Setup tmp directory for tmp files
	dir := fmt.Sprintf("%s/.task.%d", tmpdir, time.Now().Nanosecond())
	if err := os.Mkdir(dir, os.ModeDir|os.ModePerm); err != nil {
		return err
	}

	// Listen interrupt and kill signal, clear tmp files before exit.
	intqueue := make(chan os.Signal, 1)
	signal.Notify(intqueue, os.Interrupt, os.Kill)
	// If got interrupt or kill signal, delete tmp directory first, then exit with 1
	go func() {
		<-intqueue
		os.RemoveAll(dir)
		os.Exit(1)
	}()
	// If the complete all the tasks normlly, stop listenning signals and remove tmp directory
	defer func() {
		signal.Stop(intqueue)
		os.RemoveAll(dir)
	}()

	BuildKnownHosts(host_arr)

	// Create tmp file for every host, then executes.
	var tmpfiles []*os.File
	for _, h := range host_arr {
		file, _ := os.Create(fmt.Sprintf("%s/%s", dir, h))
		err_file, _ := os.Create(fmt.Sprintf("%s/%s.err", dir, h))
		tmpfiles = append(tmpfiles, file, err_file)
		task := batch.NewTask(h, user, pwd, keyfile, cmd, port, file, err_file, wg)
		task.Timeout = timeout
		go task.Work()
	}

	// show realtime view for each host
	var dc *dircat.DirCat
	if terminal.IsTerminal(1) {
		wlist := []string{}
		for _, h := range host_arr {
			wlist = append(wlist, fmt.Sprintf("%s/%s", dir, h))
		}
		dc, _ = dircat.Init(wlist...)
		go dc.Start()
	}
	// When a host is ready and request for continue, the master would echo CONTINUE for response to allow host to run
	wg.Wait()

	if terminal.IsTerminal(1) {
		dc.Stop()
	}
	// close tmp files
	for _, f := range tmpfiles {
		f.Close()
	}
	// Merge all the hosts' output to the output file
	for _, h := range host_arr {
		report(os.Stderr, "", h, true)
		// copy err output first
		err_fn := fmt.Sprintf("%s/%s.err", dir, h)
		err_src, _ := os.Open(err_fn)
		io.Copy(err_printer, err_src)
		err_src.Close()
		// copy output then
		fn := fmt.Sprintf("%s/%s", dir, h)
		src, _ := os.Open(fn)
		io.Copy(printer, src)
		src.Close()
		// remove tmp file
		os.Remove(err_fn)
		os.Remove(fn)
	}
	return nil
}
func Interact(config TaskArgs, host string) {
	wg := new(sync.WaitGroup)
	task := batch.NewTask(
		host,
		config.User,
		config.Password,
		config.Keyfile,
		config.Cmd,
		config.Port,
		config.Output,
		config.ErrOutput,
		wg,
	)
	task.SysLogin()
}

func ScpRun(config TaskArgs, host_arr []string) error {
	user := config.User
	parallel := config.Parallel
	src := config.Source
	dest := config.Destdir
	if !strings.HasSuffix(dest, "/") {
		dest += "/"
	}
	wg := new(sync.WaitGroup)
	BuildKnownHosts(host_arr)
	for _, h := range host_arr {
		wg.Add(1)
		taskF := func(host string) {
			defer wg.Done()
			var remote_shell string = ""
			if config.Port != "" && config.Port != "22" {
				remote_shell = fmt.Sprintf(` -e 'ssh -p %s' `, config.Port)
			}
			cmdstr := fmt.Sprintf(`rsync -azh %s %s %s@%s:%s`, remote_shell, src, user, host, dest)
			cmd := exec.Command("/bin/bash", "-c", cmdstr)
			fmt.Fprintf(os.Stderr, "Start sync to  %s......\n", host)
			if res, err := cmd.CombinedOutput(); err != nil {
				fmt.Fprintf(os.Stderr, "Sync to %s:%s fail:%v\n%s\n", host, dest, err, string(res))
			} else {
				fmt.Fprintf(os.Stderr, "Sync to %s:%s OK\n", host, dest)
			}
		}
		if parallel {
			go taskF(h)
		} else {
			taskF(h)
		}
	}

	wg.Wait()
	return nil
}

func BuildKnownHosts(host_arr []string) {
	filename := os.Getenv("HOME") + "/.ssh/known_hosts"
	dirname := os.Getenv("HOME") + "/.ssh"
	content_bytes, _ := ioutil.ReadFile(filename)
	content := string(content_bytes)
	var not_exists []string
	for _, host := range host_arr {
		if !strings.Contains(content, host) {
			not_exists = append(not_exists, host)
		}
	}
	if len(not_exists) == 0 {
		return
	}
	cmdstr := fmt.Sprintf(`
[ ! -e %s ] && mkdir -p %s && chmod 700 %s
[ ! -e %s ] && touch %s && chmod 644 %s
ssh-keyscan %s >> %s
`, dirname, dirname, dirname,
		filename, filename, filename,
		strings.Join(not_exists, " "), filename)
	cmd := exec.Command("/bin/sh", "-c", cmdstr)
	cmd.Run()
}
<|endoftext|>


<commit_before>package utils

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"strings"
)

// Signer interface describes facility implementing signing of files
type Signer interface {
	Init() error
	SetKey(keyRef string)
	SetKeyRing(keyring, secretKeyring string)
	DetachedSign(source string, destination string) error
	ClearSign(source string, destination string) error
}

// Verifier interface describes signature verification factility
type Verifier interface {
	InitKeyring() error
	AddKeyring(keyring string)
	VerifyDetachedSignature(signature, cleartext io.Reader) error
	VerifyClearsigned(clearsigned io.Reader) (text *os.File, err error)
}

// Test interface
var (
	_ Signer   = &GpgSigner{}
	_ Verifier = &GpgVerifier{}
)

// GpgSigner is implementation of Signer interface using gpg
type GpgSigner struct {
	keyRef                 string
	keyring, secretKeyring string
}

// SetKey sets key ID to use when signing files
func (g *GpgSigner) SetKey(keyRef string) {
	g.keyRef = keyRef
}

// SetKeyring allows to set custom keyring and secretkeyring
func (g *GpgSigner) SetKeyRing(keyring, secretKeyring string) {
	g.keyring, g.secretKeyring = keyring, secretKeyring
}

func (g *GpgSigner) gpgArgs() []string {
	args := []string{}
	if g.keyring != "" {
		args = append(args, "--no-default-keyring", "--keyring", g.keyring)
	}
	if g.secretKeyring != "" {
		args = append(args, "--secret-keyring", g.secretKeyring)
	}

	if g.keyRef != "" {
		args = append(args, "-u", g.keyRef)
	}

	return args
}

// Init verifies availability of gpg & presence of keys
func (g *GpgSigner) Init() error {
	output, err := exec.Command("gpg", "--list-keys").Output()
	if err != nil {
		return fmt.Errorf("unable to execute gpg: %s (is gpg installed?)", err)
	}

	if len(output) == 0 {
		return fmt.Errorf("looks like there are no keys in gpg, please create one (official manual: http://www.gnupg.org/gph/en/manual.html)")
	}

	return err
}

// DetachedSign signs file with detached signature in ASCII format
func (g *GpgSigner) DetachedSign(source string, destination string) error {
	fmt.Printf("Signing file '%s' with gpg, please enter your passphrase when prompted:\n", source)

	args := []string{"-o", destination, "--armor", "--yes"}
	args = append(args, g.gpgArgs()...)
	args = append(args, "--detach-sign", source)
	cmd := exec.Command("gpg", args...)
	return cmd.Run()
}

// ClearSign clear-signs the file
func (g *GpgSigner) ClearSign(source string, destination string) error {
	fmt.Printf("Clearsigning file '%s' with gpg, please enter your passphrase when prompted:\n", source)
	args := []string{"-o", destination, "--yes"}
	args = append(args, g.gpgArgs()...)
	args = append(args, "--clearsign", source)
	cmd := exec.Command("gpg", args...)
	return cmd.Run()
}

// GpgVerifier is implementation of Verifier interface using gpgv
type GpgVerifier struct {
	keyRings []string
}

// InitKeyring verifies that gpg is installed and some keys are trusted
func (g *GpgVerifier) InitKeyring() error {
	err := exec.Command("gpgv", "--version").Run()
	if err != nil {
		return fmt.Errorf("unable to execute gpgv: %s (is gpg installed?)", err)
	}

	if len(g.keyRings) == 0 {
		// using default keyring
		output, err := exec.Command("gpg", "--no-default-keyring", "--keyring", "trustedkeys.gpg", "--list-keys").Output()
		if err == nil && len(output) == 0 {
			fmt.Printf("\nLooks like your keyring with trusted keys is empty. You might consider importing some keys.\n")
			fmt.Printf("If you're running Debian or Ubuntu, it's a good idea to import current archive keys by running:\n\n")
			fmt.Printf("  gpg --keyring /usr/share/keyrings/debian-archive-keyring.gpg --export | gpg --no-default-keyring --keyring trustedkeys.gpg --import\n")
			fmt.Printf("\n(for Ubuntu, use /usr/share/keyrings/ubuntu-archive-keyring.gpg)\n\n")
		}
	}

	return nil
}

// AddKeyring adds custom keyring to GPG parameters
func (g *GpgVerifier) AddKeyring(keyring string) {
	g.keyRings = append(g.keyRings, keyring)
}

func (g *GpgVerifier) argsKeyrings() (args []string) {
	if len(g.keyRings) > 0 {
		args = make([]string, 0, 2*len(g.keyRings))
		for _, keyring := range g.keyRings {
			args = append(args, "--keyring", keyring)
		}
	} else {
		args = []string{"--keyring", "trustedkeys.gpg"}
	}
	return
}

func (g *GpgVerifier) runGpgv(args []string, context string) error {
	cmd := exec.Command("gpgv", args...)

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}
	defer stderr.Close()

	err = cmd.Start()
	if err != nil {
		return err
	}

	buffer := &bytes.Buffer{}

	_, err = io.Copy(io.MultiWriter(os.Stderr, buffer), stderr)
	if err != nil {
		return err
	}

	matches := regexp.MustCompile("ID ([0-9A-F]{8})").FindAllStringSubmatch(buffer.String(), -1)

	err = cmd.Wait()
	if err != nil {
		if len(g.keyRings) == 0 && len(matches) > 0 {
			fmt.Printf("\nLooks like some keys are missing in your trusted keyring, you may consider importing them from keyserver:\n\n")

			keyIDs := []string{}
			for _, match := range matches {
				keyIDs = append(keyIDs, match[1])
			}
			fmt.Printf("gpg --no-default-keyring --keyring trustedkeys.gpg --keyserver keys.gnupg.net --recv-keys %s\n\n",
				strings.Join(keyIDs, " "))

			fmt.Printf("Sometimes keys are stored in repository root in file named Release.key, to import such key:\n\n")
			fmt.Printf("wget -O - http://some.repo/repository/Release.key | gpg --no-default-keyring --keyring trustedkeys.gpg --import\n\n")
		}
		return fmt.Errorf("verification of %s failed: %s", context, err)
	}
	return nil
}

// VerifyDetachedSignature verifies combination of signature and cleartext using gpgv
func (g *GpgVerifier) VerifyDetachedSignature(signature, cleartext io.Reader) error {
	args := g.argsKeyrings()

	sigf, err := ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return err
	}
	defer os.Remove(sigf.Name())
	defer sigf.Close()

	_, err = io.Copy(sigf, signature)
	if err != nil {
		return err
	}

	clearf, err := ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return err
	}
	defer os.Remove(clearf.Name())
	defer clearf.Close()

	_, err = io.Copy(clearf, cleartext)
	if err != nil {
		return err
	}

	args = append(args, sigf.Name(), clearf.Name())

	return g.runGpgv(args, "detached signature")
}

// VerifyClearsigned verifies clearsigned file using gpgv and extracts cleartext version
func (g *GpgVerifier) VerifyClearsigned(clearsigned io.Reader) (text *os.File, err error) {
	args := g.argsKeyrings()

	clearf, err := ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return
	}
	defer os.Remove(clearf.Name())
	defer clearf.Close()

	_, err = io.Copy(clearf, clearsigned)
	if err != nil {
		return
	}

	args = append(args, clearf.Name())
	err = g.runGpgv(args, "clearsigned file")
	if err != nil {
		return nil, err
	}

	text, err = ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return
	}
	defer os.Remove(text.Name())

	args = []string{"--no-default-keyring"}
	args = append(args, g.argsKeyrings()...)
	args = append(args, "--decrypt", "--batch", "--trust-model", "always", "--output", "-", clearf.Name())

	cmd := exec.Command("gpg", args...)
	cmd.Stderr = os.Stderr
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	defer stdout.Close()

	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	_, err = io.Copy(text, stdout)
	if err != nil {
		return nil, err
	}

	err = cmd.Wait()

	if err != nil {
		return nil, fmt.Errorf("extraction of clearsigned file failed: %s", err)
	}

	_, err = text.Seek(0, 0)
	if err != nil {
		return nil, err
	}

	return
}
<commit_msg>Don't compain about missing keys when using custom keyring.<commit_after>package utils

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"strings"
)

// Signer interface describes facility implementing signing of files
type Signer interface {
	Init() error
	SetKey(keyRef string)
	SetKeyRing(keyring, secretKeyring string)
	DetachedSign(source string, destination string) error
	ClearSign(source string, destination string) error
}

// Verifier interface describes signature verification factility
type Verifier interface {
	InitKeyring() error
	AddKeyring(keyring string)
	VerifyDetachedSignature(signature, cleartext io.Reader) error
	VerifyClearsigned(clearsigned io.Reader) (text *os.File, err error)
}

// Test interface
var (
	_ Signer   = &GpgSigner{}
	_ Verifier = &GpgVerifier{}
)

// GpgSigner is implementation of Signer interface using gpg
type GpgSigner struct {
	keyRef                 string
	keyring, secretKeyring string
}

// SetKey sets key ID to use when signing files
func (g *GpgSigner) SetKey(keyRef string) {
	g.keyRef = keyRef
}

// SetKeyring allows to set custom keyring and secretkeyring
func (g *GpgSigner) SetKeyRing(keyring, secretKeyring string) {
	g.keyring, g.secretKeyring = keyring, secretKeyring
}

func (g *GpgSigner) gpgArgs() []string {
	args := []string{}
	if g.keyring != "" {
		args = append(args, "--no-default-keyring", "--keyring", g.keyring)
	}
	if g.secretKeyring != "" {
		args = append(args, "--secret-keyring", g.secretKeyring)
	}

	if g.keyRef != "" {
		args = append(args, "-u", g.keyRef)
	}

	return args
}

// Init verifies availability of gpg & presence of keys
func (g *GpgSigner) Init() error {
	output, err := exec.Command("gpg", "--list-keys").Output()
	if err != nil {
		return fmt.Errorf("unable to execute gpg: %s (is gpg installed?)", err)
	}

	if g.keyring == "" && g.secretKeyring == "" && len(output) == 0 {
		return fmt.Errorf("looks like there are no keys in gpg, please create one (official manual: http://www.gnupg.org/gph/en/manual.html)")
	}

	return err
}

// DetachedSign signs file with detached signature in ASCII format
func (g *GpgSigner) DetachedSign(source string, destination string) error {
	fmt.Printf("Signing file '%s' with gpg, please enter your passphrase when prompted:\n", source)

	args := []string{"-o", destination, "--armor", "--yes"}
	args = append(args, g.gpgArgs()...)
	args = append(args, "--detach-sign", source)
	cmd := exec.Command("gpg", args...)
	return cmd.Run()
}

// ClearSign clear-signs the file
func (g *GpgSigner) ClearSign(source string, destination string) error {
	fmt.Printf("Clearsigning file '%s' with gpg, please enter your passphrase when prompted:\n", source)
	args := []string{"-o", destination, "--yes"}
	args = append(args, g.gpgArgs()...)
	args = append(args, "--clearsign", source)
	cmd := exec.Command("gpg", args...)
	return cmd.Run()
}

// GpgVerifier is implementation of Verifier interface using gpgv
type GpgVerifier struct {
	keyRings []string
}

// InitKeyring verifies that gpg is installed and some keys are trusted
func (g *GpgVerifier) InitKeyring() error {
	err := exec.Command("gpgv", "--version").Run()
	if err != nil {
		return fmt.Errorf("unable to execute gpgv: %s (is gpg installed?)", err)
	}

	if len(g.keyRings) == 0 {
		// using default keyring
		output, err := exec.Command("gpg", "--no-default-keyring", "--keyring", "trustedkeys.gpg", "--list-keys").Output()
		if err == nil && len(output) == 0 {
			fmt.Printf("\nLooks like your keyring with trusted keys is empty. You might consider importing some keys.\n")
			fmt.Printf("If you're running Debian or Ubuntu, it's a good idea to import current archive keys by running:\n\n")
			fmt.Printf("  gpg --keyring /usr/share/keyrings/debian-archive-keyring.gpg --export | gpg --no-default-keyring --keyring trustedkeys.gpg --import\n")
			fmt.Printf("\n(for Ubuntu, use /usr/share/keyrings/ubuntu-archive-keyring.gpg)\n\n")
		}
	}

	return nil
}

// AddKeyring adds custom keyring to GPG parameters
func (g *GpgVerifier) AddKeyring(keyring string) {
	g.keyRings = append(g.keyRings, keyring)
}

func (g *GpgVerifier) argsKeyrings() (args []string) {
	if len(g.keyRings) > 0 {
		args = make([]string, 0, 2*len(g.keyRings))
		for _, keyring := range g.keyRings {
			args = append(args, "--keyring", keyring)
		}
	} else {
		args = []string{"--keyring", "trustedkeys.gpg"}
	}
	return
}

func (g *GpgVerifier) runGpgv(args []string, context string) error {
	cmd := exec.Command("gpgv", args...)

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}
	defer stderr.Close()

	err = cmd.Start()
	if err != nil {
		return err
	}

	buffer := &bytes.Buffer{}

	_, err = io.Copy(io.MultiWriter(os.Stderr, buffer), stderr)
	if err != nil {
		return err
	}

	matches := regexp.MustCompile("ID ([0-9A-F]{8})").FindAllStringSubmatch(buffer.String(), -1)

	err = cmd.Wait()
	if err != nil {
		if len(g.keyRings) == 0 && len(matches) > 0 {
			fmt.Printf("\nLooks like some keys are missing in your trusted keyring, you may consider importing them from keyserver:\n\n")

			keyIDs := []string{}
			for _, match := range matches {
				keyIDs = append(keyIDs, match[1])
			}
			fmt.Printf("gpg --no-default-keyring --keyring trustedkeys.gpg --keyserver keys.gnupg.net --recv-keys %s\n\n",
				strings.Join(keyIDs, " "))

			fmt.Printf("Sometimes keys are stored in repository root in file named Release.key, to import such key:\n\n")
			fmt.Printf("wget -O - http://some.repo/repository/Release.key | gpg --no-default-keyring --keyring trustedkeys.gpg --import\n\n")
		}
		return fmt.Errorf("verification of %s failed: %s", context, err)
	}
	return nil
}

// VerifyDetachedSignature verifies combination of signature and cleartext using gpgv
func (g *GpgVerifier) VerifyDetachedSignature(signature, cleartext io.Reader) error {
	args := g.argsKeyrings()

	sigf, err := ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return err
	}
	defer os.Remove(sigf.Name())
	defer sigf.Close()

	_, err = io.Copy(sigf, signature)
	if err != nil {
		return err
	}

	clearf, err := ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return err
	}
	defer os.Remove(clearf.Name())
	defer clearf.Close()

	_, err = io.Copy(clearf, cleartext)
	if err != nil {
		return err
	}

	args = append(args, sigf.Name(), clearf.Name())

	return g.runGpgv(args, "detached signature")
}

// VerifyClearsigned verifies clearsigned file using gpgv and extracts cleartext version
func (g *GpgVerifier) VerifyClearsigned(clearsigned io.Reader) (text *os.File, err error) {
	args := g.argsKeyrings()

	clearf, err := ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return
	}
	defer os.Remove(clearf.Name())
	defer clearf.Close()

	_, err = io.Copy(clearf, clearsigned)
	if err != nil {
		return
	}

	args = append(args, clearf.Name())
	err = g.runGpgv(args, "clearsigned file")
	if err != nil {
		return nil, err
	}

	text, err = ioutil.TempFile("", "aptly-gpg")
	if err != nil {
		return
	}
	defer os.Remove(text.Name())

	args = []string{"--no-default-keyring"}
	args = append(args, g.argsKeyrings()...)
	args = append(args, "--decrypt", "--batch", "--trust-model", "always", "--output", "-", clearf.Name())

	cmd := exec.Command("gpg", args...)
	cmd.Stderr = os.Stderr
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	defer stdout.Close()

	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	_, err = io.Copy(text, stdout)
	if err != nil {
		return nil, err
	}

	err = cmd.Wait()

	if err != nil {
		return nil, fmt.Errorf("extraction of clearsigned file failed: %s", err)
	}

	_, err = text.Seek(0, 0)
	if err != nil {
		return nil, err
	}

	return
}
<|endoftext|>


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

package slogger

import (
	"errors"
	"fmt"
	"github.com/tolsen/slogger/v2/queued_set"
	"runtime"
	"strings"
	"time"
)

type Log struct {
	Prefix     string
	Level      Level
	Filename   string
	Line       int
	Timestamp  time.Time
	MessageFmt string
	Args       []interface{}
	Context    *Context
}

func SimpleLog(prefix string, level Level, callerSkip int, messageFmt string, args []interface{}) *Log {
	_, file, line, ok := runtime.Caller(callerSkip)
	if !ok {
		file = "UNKNOWN_FILE"
		line = -1
	}

	return &Log{
		Prefix:     prefix,
		Level:      level,
		Filename:   file,
		Line:       line,
		Timestamp:  time.Now(),
		MessageFmt: messageFmt,
		Args:       args,
	}
}

func (self *Log) Message() string {
	messageFmt := self.MessageFmt
	if self.Context != nil {
		messageFmt = self.Context.interpolateString(self.MessageFmt)
	}

	return fmt.Sprintf(messageFmt, self.Args...)
}

// for use as a cache key
func (self *Log) stringWithoutTime() string {
	return fmt.Sprintf(
		"%s %v %s %d %s",
		self.Prefix,
		self.Level.Type(),
		self.Filename,
		self.Line,
		self.Message(),
	)
}

type Logger struct {
	Prefix             string
	Appenders          []Appender
	StripDirs          int
	cache              *queued_set.QueuedSet
	suppressionEnabled bool
}

// Log a message and a level to a logger instance. This returns a
// pointer to a Log and a slice of errors that were gathered from every
// Appender (nil errors included).
func (self *Logger) Logf(level Level, messageFmt string, args ...interface{}) (*Log, []error) {
	return self.logf(level, messageFmt, nil, args...)
}

func (self *Logger) LogfWithContext(level Level, messageFmt string, context *Context, args ...interface{}) (*Log, []error) {
	return self.logf(level, messageFmt, context, args...)
}

func (self *Logger) DisableLogSuppression() {
	self.cache = nil
	self.suppressionEnabled = false
	return
}

func (self *Logger) EnableLogSuppression(historyCapacity int) {
	self.cache = queued_set.New(historyCapacity)
	self.suppressionEnabled = true
	return
}

// Log and return a formatted error string.
// Example:
//
// if whatIsExpected != whatIsReturned {
//     return slogger.Errorf(slogger.WARN, "Unexpected return value. Expected: %v Received: %v",
//         whatIsExpected, whatIsReturned)
// }5
//
func (self *Logger) Errorf(level Level, messageFmt string, args ...interface{}) error {
	return self.ErrorfWithContext(level, messageFmt, nil, args...)
}

func (self *Logger) ErrorfWithContext(level Level, messageFmt string, context *Context, args ...interface{}) error {
	log, _ := self.logf(level, messageFmt, context, args...)
	return errors.New(log.Message())
}

func (self *Logger) Flush() (errors []error) {
	for _, appender := range self.Appenders {
		if err := appender.Flush(); err != nil {
			errors = append(errors, err)
		}
	}
	return
}

func (self *Logger) IsSuppressionEnabled() bool {
	return self.suppressionEnabled
}

// Stackf is designed to work in tandem with `NewStackError`. This
// function is similar to `Logf`, but takes a `stackErr`
// parameter. `stackErr` is expected to be of type StackError, but does
// not have to be.
func (self *Logger) Stackf(level Level, stackErr error, messageFmt string, args ...interface{}) (*Log, []error) {
	return self.StackfWithContext(level, stackErr, messageFmt, nil, args...)
}

func (self *Logger) StackfWithContext(level Level, stackErr error, messageFmt string, context *Context, args ...interface{}) (*Log, []error) {
	messageFmt = fmt.Sprintf("%v\n%v", messageFmt, stackErr.Error())
	return self.logf(level, messageFmt, context, args...)
}

var ignoredFileNames = []string{"logger.go"}

// Add a file to the list of file names that slogger will skip when it identifies the source
// of a message.  This is useful if you have a logging library built on top of slogger.
// If you IgnoreThisFilenameToo(...) on the files of that library, logging messages
// will be marked as coming from your code that calls your library, rather than from your library.
func IgnoreThisFilenameToo(fn string) {
	ignoredFileNames = append(ignoredFileNames, fn)
}

func containsAnyIgnoredFilename(s string) bool {
	for _, ign := range ignoredFileNames {
		if strings.Contains(s, ign) {
			return true
		}
	}
	return false
}

func nonSloggerCaller() (pc uintptr, file string, line int, ok bool) {
	for skip := 0; skip < 100; skip++ {
		pc, file, line, ok := runtime.Caller(skip)
		if !ok || !containsAnyIgnoredFilename(file) {
			return pc, file, line, ok
		}
	}
	return 0, "", 0, false
}

// accepts a Context or *Context as the first element of args
func (self *Logger) logf(level Level, messageFmt string, context *Context, args ...interface{}) (*Log, []error) {
	var errors []error

	_, file, line, ok := nonSloggerCaller()
	//	_, file, line, ok := runtime.Caller(2+offset)
	if ok == false {
		return nil, []error{fmt.Errorf("Failed to find the calling method.")}
	}

	file = stripDirectories(file, self.StripDirs)
	log := &Log{
		Prefix:     self.Prefix,
		Level:      level,
		Filename:   file,
		Line:       line,
		Timestamp:  time.Now(),
		MessageFmt: messageFmt,
		Args:       args,
		Context:    context,
	}

	if !self.suppressionEnabled || self.cache.Add(log.stringWithoutTime()) {
		for _, appender := range self.Appenders {
			if err := appender.Append(log); err != nil {
				error := fmt.Errorf("Error appending. Appender: %T Error: %v", appender, err)
				errors = append(errors, error)
			}
		}
	}

	return log, errors
}

type Level uint8

// The level is in an order such that the expressions
// `level < WARN`, `level >= INFO` have intuitive meaning.
const (
	OFF Level = iota
	DEBUG
	ROUTINE
	INFO
	WARN
	ERROR
	DOOM
	topLevel
)

var strToLevel map[string]Level

var levelToStr []string

func init() {
	strToLevel = map[string]Level{
		"off":     OFF,
		"debug":   DEBUG,
		"routine": ROUTINE,
		"info":    INFO,
		"warn":    WARN,
		"error":   ERROR,
		"doom":    DOOM,
	}

	levelToStr = make([]string, len(strToLevel))
	for str, level := range strToLevel {
		levelToStr[uint8(level)] = str
	}
}

func NewLevel(levelStr string) (Level, error) {
	level, ok := strToLevel[strings.ToLower(levelStr)]

	if !ok {
		err := UnknownLevelError{levelStr}
		return OFF, err
	}

	return level, nil
}

func (self Level) Type() string {
	if self >= topLevel {
		return "off?"
	}

	return levelToStr[uint8(self)]
}

func stacktrace() []string {
	ret := make([]string, 0, 2)
	for skip := 2; true; skip++ {
		_, file, line, ok := runtime.Caller(skip)
		if ok == false {
			break
		}

		ret = append(ret, fmt.Sprintf("at %s:%d", stripDirectories(file, 2), line))
	}

	return ret
}

type StackError struct {
	Message    string
	Stacktrace []string
}

func NewStackError(messageFmt string, args ...interface{}) *StackError {
	return &StackError{
		Message:    fmt.Sprintf(messageFmt, args...),
		Stacktrace: stacktrace(),
	}
}

func (self *StackError) Error() string {
	return fmt.Sprintf("%s\n\t%s", self.Message, strings.Join(self.Stacktrace, "\n\t"))
}

func stripDirectories(filepath string, toKeep int) string {
	var idxCutoff int
	if idxCutoff = strings.LastIndex(filepath, "/"); idxCutoff == -1 {
		return filepath
	}

	for dirToKeep := 0; dirToKeep < toKeep; dirToKeep++ {
		switch idx := strings.LastIndex(filepath[:idxCutoff], "/"); idx {
		case -1:
			break
		default:
			idxCutoff = idx
		}
	}

	return filepath[idxCutoff+1:]
}

type UnknownLevelError struct {
	levelStr string
}

func (self UnknownLevelError) Error() string {
	return fmt.Sprintf("Unknown level: %s", self.levelStr)
}
<commit_msg>change order of setting cache to nil and setting suppressionEnabled to false in DisableLogSuppression() in order to avoid period of time where suppressionEnabled is true but cache is nil<commit_after>// Copyright 2013 MongoDB, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package slogger

import (
	"errors"
	"fmt"
	"github.com/tolsen/slogger/v2/queued_set"
	"runtime"
	"strings"
	"time"
)

type Log struct {
	Prefix     string
	Level      Level
	Filename   string
	Line       int
	Timestamp  time.Time
	MessageFmt string
	Args       []interface{}
	Context    *Context
}

func SimpleLog(prefix string, level Level, callerSkip int, messageFmt string, args []interface{}) *Log {
	_, file, line, ok := runtime.Caller(callerSkip)
	if !ok {
		file = "UNKNOWN_FILE"
		line = -1
	}

	return &Log{
		Prefix:     prefix,
		Level:      level,
		Filename:   file,
		Line:       line,
		Timestamp:  time.Now(),
		MessageFmt: messageFmt,
		Args:       args,
	}
}

func (self *Log) Message() string {
	messageFmt := self.MessageFmt
	if self.Context != nil {
		messageFmt = self.Context.interpolateString(self.MessageFmt)
	}

	return fmt.Sprintf(messageFmt, self.Args...)
}

// for use as a cache key
func (self *Log) stringWithoutTime() string {
	return fmt.Sprintf(
		"%s %v %s %d %s",
		self.Prefix,
		self.Level.Type(),
		self.Filename,
		self.Line,
		self.Message(),
	)
}

type Logger struct {
	Prefix             string
	Appenders          []Appender
	StripDirs          int
	cache              *queued_set.QueuedSet
	suppressionEnabled bool
}

// Log a message and a level to a logger instance. This returns a
// pointer to a Log and a slice of errors that were gathered from every
// Appender (nil errors included).
func (self *Logger) Logf(level Level, messageFmt string, args ...interface{}) (*Log, []error) {
	return self.logf(level, messageFmt, nil, args...)
}

func (self *Logger) LogfWithContext(level Level, messageFmt string, context *Context, args ...interface{}) (*Log, []error) {
	return self.logf(level, messageFmt, context, args...)
}

func (self *Logger) DisableLogSuppression() {
	self.suppressionEnabled = false
	self.cache = nil
	return
}

func (self *Logger) EnableLogSuppression(historyCapacity int) {
	self.cache = queued_set.New(historyCapacity)
	self.suppressionEnabled = true
	return
}

// Log and return a formatted error string.
// Example:
//
// if whatIsExpected != whatIsReturned {
//     return slogger.Errorf(slogger.WARN, "Unexpected return value. Expected: %v Received: %v",
//         whatIsExpected, whatIsReturned)
// }5
//
func (self *Logger) Errorf(level Level, messageFmt string, args ...interface{}) error {
	return self.ErrorfWithContext(level, messageFmt, nil, args...)
}

func (self *Logger) ErrorfWithContext(level Level, messageFmt string, context *Context, args ...interface{}) error {
	log, _ := self.logf(level, messageFmt, context, args...)
	return errors.New(log.Message())
}

func (self *Logger) Flush() (errors []error) {
	for _, appender := range self.Appenders {
		if err := appender.Flush(); err != nil {
			errors = append(errors, err)
		}
	}
	return
}

func (self *Logger) IsSuppressionEnabled() bool {
	return self.suppressionEnabled
}

// Stackf is designed to work in tandem with `NewStackError`. This
// function is similar to `Logf`, but takes a `stackErr`
// parameter. `stackErr` is expected to be of type StackError, but does
// not have to be.
func (self *Logger) Stackf(level Level, stackErr error, messageFmt string, args ...interface{}) (*Log, []error) {
	return self.StackfWithContext(level, stackErr, messageFmt, nil, args...)
}

func (self *Logger) StackfWithContext(level Level, stackErr error, messageFmt string, context *Context, args ...interface{}) (*Log, []error) {
	messageFmt = fmt.Sprintf("%v\n%v", messageFmt, stackErr.Error())
	return self.logf(level, messageFmt, context, args...)
}

var ignoredFileNames = []string{"logger.go"}

// Add a file to the list of file names that slogger will skip when it identifies the source
// of a message.  This is useful if you have a logging library built on top of slogger.
// If you IgnoreThisFilenameToo(...) on the files of that library, logging messages
// will be marked as coming from your code that calls your library, rather than from your library.
func IgnoreThisFilenameToo(fn string) {
	ignoredFileNames = append(ignoredFileNames, fn)
}

func containsAnyIgnoredFilename(s string) bool {
	for _, ign := range ignoredFileNames {
		if strings.Contains(s, ign) {
			return true
		}
	}
	return false
}

func nonSloggerCaller() (pc uintptr, file string, line int, ok bool) {
	for skip := 0; skip < 100; skip++ {
		pc, file, line, ok := runtime.Caller(skip)
		if !ok || !containsAnyIgnoredFilename(file) {
			return pc, file, line, ok
		}
	}
	return 0, "", 0, false
}

// accepts a Context or *Context as the first element of args
func (self *Logger) logf(level Level, messageFmt string, context *Context, args ...interface{}) (*Log, []error) {
	var errors []error

	_, file, line, ok := nonSloggerCaller()
	//	_, file, line, ok := runtime.Caller(2+offset)
	if ok == false {
		return nil, []error{fmt.Errorf("Failed to find the calling method.")}
	}

	file = stripDirectories(file, self.StripDirs)
	log := &Log{
		Prefix:     self.Prefix,
		Level:      level,
		Filename:   file,
		Line:       line,
		Timestamp:  time.Now(),
		MessageFmt: messageFmt,
		Args:       args,
		Context:    context,
	}

	if !self.suppressionEnabled || self.cache.Add(log.stringWithoutTime()) {
		for _, appender := range self.Appenders {
			if err := appender.Append(log); err != nil {
				error := fmt.Errorf("Error appending. Appender: %T Error: %v", appender, err)
				errors = append(errors, error)
			}
		}
	}

	return log, errors
}

type Level uint8

// The level is in an order such that the expressions
// `level < WARN`, `level >= INFO` have intuitive meaning.
const (
	OFF Level = iota
	DEBUG
	ROUTINE
	INFO
	WARN
	ERROR
	DOOM
	topLevel
)

var strToLevel map[string]Level

var levelToStr []string

func init() {
	strToLevel = map[string]Level{
		"off":     OFF,
		"debug":   DEBUG,
		"routine": ROUTINE,
		"info":    INFO,
		"warn":    WARN,
		"error":   ERROR,
		"doom":    DOOM,
	}

	levelToStr = make([]string, len(strToLevel))
	for str, level := range strToLevel {
		levelToStr[uint8(level)] = str
	}
}

func NewLevel(levelStr string) (Level, error) {
	level, ok := strToLevel[strings.ToLower(levelStr)]

	if !ok {
		err := UnknownLevelError{levelStr}
		return OFF, err
	}

	return level, nil
}

func (self Level) Type() string {
	if self >= topLevel {
		return "off?"
	}

	return levelToStr[uint8(self)]
}

func stacktrace() []string {
	ret := make([]string, 0, 2)
	for skip := 2; true; skip++ {
		_, file, line, ok := runtime.Caller(skip)
		if ok == false {
			break
		}

		ret = append(ret, fmt.Sprintf("at %s:%d", stripDirectories(file, 2), line))
	}

	return ret
}

type StackError struct {
	Message    string
	Stacktrace []string
}

func NewStackError(messageFmt string, args ...interface{}) *StackError {
	return &StackError{
		Message:    fmt.Sprintf(messageFmt, args...),
		Stacktrace: stacktrace(),
	}
}

func (self *StackError) Error() string {
	return fmt.Sprintf("%s\n\t%s", self.Message, strings.Join(self.Stacktrace, "\n\t"))
}

func stripDirectories(filepath string, toKeep int) string {
	var idxCutoff int
	if idxCutoff = strings.LastIndex(filepath, "/"); idxCutoff == -1 {
		return filepath
	}

	for dirToKeep := 0; dirToKeep < toKeep; dirToKeep++ {
		switch idx := strings.LastIndex(filepath[:idxCutoff], "/"); idx {
		case -1:
			break
		default:
			idxCutoff = idx
		}
	}

	return filepath[idxCutoff+1:]
}

type UnknownLevelError struct {
	levelStr string
}

func (self UnknownLevelError) Error() string {
	return fmt.Sprintf("Unknown level: %s", self.levelStr)
}
<|endoftext|>


<commit_before>package popart

var (
	errInvalidSyntax   = NewReportableError("invalid syntax")
	errUnexpectedState = NewReportableError("unexpected state transition")
)

var (
	validators = map[string]*validator{
		"APOP": validates(state(stateAuthorization), arity(2)),
		"CAPA": validates(state(stateAuthorization, stateTransaction), arity(0),
		"DELE": validates(state(stateTransaction), arity(1)),
		"LIST": validates(state(stateTransaction), arity(0, 1)),
		"NOOP": validates(state(stateTransaction), arity(0)),
		"PASS": validates(state(stateAuthorization), arity(1)),
		"QUIT": validates(state(stateAuthorization, stateTransaction), arity(0)),
		"RETR": validates(state(stateTransaction), arity(1)),
		"RSET": validates(state(stateTransaction), arity(0)),
		"STAT": validates(state(stateTransaction), arity(0)),
		"TOP":  validates(state(stateTransaction), arity(2)),
		"UIDL": validates(state(stateTransaction), arity(0, 1)),
		"USER": validates(state(stateAuthorization), arity(1)),
	}
)

type validator struct {
	allowedStates  []int
	allowedArities []int
}

type option func(*validator)

func validates(opts ...option) *validator {
	ret := &validator{}
	for _, opt := range opts {
		opt(ret)
	}
	return ret
}

func (v *validator) validate(s *session, args []string) error {
	if err := v.allowedState(s); err != nil {
		return err
	}
	return v.allowedArity(args)
}

func (v *validator) allowedState(s *session) error {
	for _, state := range v.allowedStates {
		if state == s.state {
			return nil
		}
	}
	return errUnexpectedState
}

func (v *validator) allowedArity(args []string) error {
	for _, ar := range v.allowedArities {
		if ar == len(args) {
			return nil
		}
	}
	return errInvalidSyntax
}

func state(states ...int) option {
	return func(v *validator) {
		v.allowedStates = states
	}
}

func arity(arities ...int) option {
	return func(v *validator) {
		v.allowedArities = arities
	}
}
<commit_msg>fix braces<commit_after>package popart

var (
	errInvalidSyntax   = NewReportableError("invalid syntax")
	errUnexpectedState = NewReportableError("unexpected state transition")
)

var (
	validators = map[string]*validator{
		"APOP": validates(state(stateAuthorization), arity(2)),
		"CAPA": validates(state(stateAuthorization, stateTransaction), arity(0)),
		"DELE": validates(state(stateTransaction), arity(1)),
		"LIST": validates(state(stateTransaction), arity(0, 1)),
		"NOOP": validates(state(stateTransaction), arity(0)),
		"PASS": validates(state(stateAuthorization), arity(1)),
		"QUIT": validates(state(stateAuthorization, stateTransaction), arity(0)),
		"RETR": validates(state(stateTransaction), arity(1)),
		"RSET": validates(state(stateTransaction), arity(0)),
		"STAT": validates(state(stateTransaction), arity(0)),
		"TOP":  validates(state(stateTransaction), arity(2)),
		"UIDL": validates(state(stateTransaction), arity(0, 1)),
		"USER": validates(state(stateAuthorization), arity(1)),
	}
)

type validator struct {
	allowedStates  []int
	allowedArities []int
}

type option func(*validator)

func validates(opts ...option) *validator {
	ret := &validator{}
	for _, opt := range opts {
		opt(ret)
	}
	return ret
}

func (v *validator) validate(s *session, args []string) error {
	if err := v.allowedState(s); err != nil {
		return err
	}
	return v.allowedArity(args)
}

func (v *validator) allowedState(s *session) error {
	for _, state := range v.allowedStates {
		if state == s.state {
			return nil
		}
	}
	return errUnexpectedState
}

func (v *validator) allowedArity(args []string) error {
	for _, ar := range v.allowedArities {
		if ar == len(args) {
			return nil
		}
	}
	return errInvalidSyntax
}

func state(states ...int) option {
	return func(v *validator) {
		v.allowedStates = states
	}
}

func arity(arities ...int) option {
	return func(v *validator) {
		v.allowedArities = arities
	}
}
<|endoftext|>


<commit_before><commit_msg>simplify create theme in kit<commit_after><|endoftext|>


<commit_before>package ps

import (
	"encoding/csv"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/gyuho/psn/table"
	"github.com/olekukonko/tablewriter"
)

// ProcessTableColumns is columns for CSV file.
var ProcessTableColumns = []string{
	"NAME",
	"STATE",

	"PID",
	"PPID",

	"CPU",
	"VM_RSS",
	"VM_SIZE",

	"FD",
	"THREADS",

	"CpuUsageFloat64",
	"VmRSSBytes",
	"VmSizeBytes",
}

// List finds the status by specifying the filter.
func List(filter *Process) ([]Process, error) {
	if filter != nil {
		if filter.Stat.Pid != 0 && filter.Status.Pid == 0 {
			filter.Status.Pid = filter.Stat.Pid
		}
		if filter.Stat.Pid == 0 && filter.Status.Pid != 0 {
			filter.Stat.Pid = filter.Status.Pid
		}
	}
	if filter != nil && filter.Stat.Pid != 0 && filter.Status.Pid != 0 {
		stat, err := GetStat(filter.Stat.Pid)
		if err != nil {
			return nil, err
		}
		status, err := GetStatus(filter.Status.Pid)
		if err != nil {
			return nil, err
		}
		return []Process{Process{Stat: stat, Status: status}}, err
	}

	ds, err := ioutil.ReadDir("/proc")
	if err != nil {
		return nil, err
	}
	pids := []int64{}
	for _, f := range ds {
		if f.IsDir() && isInt(f.Name()) {
			i, err := strconv.ParseInt(f.Name(), 10, 64)
			if err != nil {
				return nil, err
			}
			pids = append(pids, i)
		}
	}

	donec, errc := make(chan struct{}), make(chan error)
	rmap := make(map[int64]Process)
	var mu sync.Mutex

	for _, pid := range pids {
		go func(pid int64, filter *Process) {
			stat, err := GetStat(pid)
			if err != nil {
				errc <- err
				return
			}
			if stat.Match(filter) {
				mu.Lock()
				if v, ok := rmap[pid]; !ok {
					rmap[pid] = Process{Stat: stat}
				} else {
					v.Stat = stat
					rmap[pid] = v
				}
				mu.Unlock()
			}
			donec <- struct{}{}
		}(pid, filter)
		go func(pid int64, filter *Process) {
			status, err := GetStatus(pid)
			if err != nil {
				errc <- err
				return
			}
			if status.Match(filter) {
				mu.Lock()
				if v, ok := rmap[pid]; !ok {
					rmap[pid] = Process{Status: status}
				} else {
					v.Status = status
					rmap[pid] = v
				}
				mu.Unlock()
			}
			donec <- struct{}{}
		}(pid, filter)
	}

	cnt := 0
	for cnt != 2*len(pids) {
		select {
		case <-donec:
			cnt++
		case e := <-errc:
			return nil, e
		}
	}

	rs := []Process{}
	for _, proc := range rmap {
		if proc.Stat.Comm != "" && proc.Status.Name == "" {
			proc.Stat.Comm = proc.Status.Name
			// return nil, fmt.Errorf("%s != %s\n", proc.Stat.Comm, proc.Status.Name)
		}
		if proc.Stat.Comm == "" && proc.Status.Name != "" {
			proc.Status.Name = proc.Stat.Comm
		}
		if proc.Stat.Pid != proc.Status.Pid {
			return nil, fmt.Errorf("%s : %d != %d\n", proc.Stat.Comm, proc.Stat.Pid, proc.Status.Pid)
		}
		rs = append(rs, proc)
	}
	return rs, nil
}

// Match matches Stat to filter. Only supports name and PID.
func (s *Stat) Match(filter *Process) bool {
	if s == nil {
		return false
	}
	if filter == nil {
		return true // no need to compare
	}
	name := filter.Stat.Comm
	if name != "" {
		if name != s.Comm {
			return false
		}
	}
	pid := filter.Stat.Pid
	if pid != 0 {
		if pid != s.Pid {
			return false
		}
	}
	return true
}

// Match matches Status to filter. Only supports name and PID.
func (s *Status) Match(filter *Process) bool {
	if s == nil {
		return false
	}
	if filter == nil {
		return true // no need to compare
	}
	name := filter.Status.Name
	if name != "" {
		if name != s.Name {
			return false
		}
	}
	pid := filter.Status.Pid
	if pid != 0 {
		if pid != s.Pid {
			return false
		}
	}
	state := filter.Status.State
	if state != "" {
		if state != s.State {
			return false
		}
	}
	return true
}

// WriteToTable writes slice of Process to ASCII table.
func WriteToTable(w io.Writer, top int, pss ...Process) {
	tw := tablewriter.NewWriter(w)
	tw.SetHeader(ProcessTableColumns[:9:9])

	rows := make([][]string, len(pss))
	for i, s := range pss {
		sl := make([]string, len(ProcessTableColumns))
		sl[0] = s.Status.Name
		sl[1] = s.Status.State
		sl[2] = fmt.Sprintf("%d", s.Status.Pid)
		sl[3] = fmt.Sprintf("%d", s.Status.PPid)
		sl[4] = fmt.Sprintf("%3.2f %%", s.Stat.CpuUsage)
		sl[5] = s.Status.VmRSS
		sl[6] = s.Status.VmSize
		sl[7] = fmt.Sprintf("%d", s.Status.FDSize)
		sl[8] = fmt.Sprintf("%d", s.Status.Threads)

		sl[9] = fmt.Sprintf("%3.2f", s.Stat.CpuUsage)
		sl[10] = fmt.Sprintf("%d", s.Status.VmRSSBytes)
		sl[11] = fmt.Sprintf("%d", s.Status.VmSizeBytes)
		rows[i] = sl
	}
	table.By(
		rows,
		table.MakeDescendingIntFunc(10),    // VM_RSS
		table.MakeDescendingFloat64Func(9), // CPU
		table.MakeDescendingIntFunc(11),    // VM_SIZE
	).Sort(rows)

	if top != 0 && len(rows) > top {
		rows = rows[:top:top]
	}
	for _, row := range rows {
		tw.Append(row[:9:9])
	}

	tw.Render()
}

var once sync.Once

// WriteToTable writes slice of Process to a csv file.
func WriteToCSV(f *os.File, pss ...Process) error {
	wr := csv.NewWriter(f)

	var werr error
	writeCSVHeader := func() {
		if err := wr.Write(append([]string{"unix_ts"}, ProcessTableColumns...)); err != nil {
			werr = err
		}
	}
	once.Do(writeCSVHeader)
	if werr != nil {
		return werr
	}

	rows := make([][]string, len(pss))
	for i, s := range pss {
		sl := make([]string, len(ProcessTableColumns))
		sl[0] = s.Status.Name
		sl[1] = s.Status.State
		sl[2] = fmt.Sprintf("%d", s.Status.Pid)
		sl[3] = fmt.Sprintf("%d", s.Status.PPid)
		sl[4] = fmt.Sprintf("%3.2f %%", s.Stat.CpuUsage)
		sl[5] = s.Status.VmRSS
		sl[6] = s.Status.VmSize
		sl[7] = fmt.Sprintf("%d", s.Status.FDSize)
		sl[8] = fmt.Sprintf("%d", s.Status.Threads)

		sl[9] = fmt.Sprintf("%3.2f", s.Stat.CpuUsage)
		sl[10] = fmt.Sprintf("%d", s.Status.VmRSSBytes)
		sl[11] = fmt.Sprintf("%d", s.Status.VmSizeBytes)
		rows[i] = sl
	}
	table.By(
		rows,
		table.MakeDescendingIntFunc(10),    // VM_RSS
		table.MakeDescendingFloat64Func(9), // CPU
		table.MakeDescendingIntFunc(11),    // VM_SIZE
	).Sort(rows)

	ts := fmt.Sprintf("%d", time.Now().Unix())
	nrows := make([][]string, len(rows))
	for i, row := range rows {
		nrows[i] = append([]string{ts}, row...)
	}
	if err := wr.WriteAll(nrows); err != nil {
		return err
	}

	wr.Flush()
	return wr.Error()
}

type int64Slice []int64

func (s int64Slice) Len() int           { return len(s) }
func (s int64Slice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
func (s int64Slice) Less(i, j int) bool { return s[i] < s[j] }

// Kill kills all processes in arguments.
func Kill(w io.Writer, parent bool, pss ...Process) {
	defer func() {
		if err := recover(); err != nil {
			fmt.Fprintln(w, "Kill:", err)
		}
	}()

	pidToKill := make(map[int64]string)
	for _, s := range pss {
		pidToKill[s.Status.Pid] = s.Status.Name
		if parent && s.Status.PPid != 0 {
			pidToKill[s.Status.PPid] = s.Status.Name
		}
	}
	if len(pidToKill) == 0 {
		fmt.Fprintln(w, "no PID to kill...")
		return
	}

	pids := []int64{}
	for pid := range pidToKill {
		pids = append(pids, pid)
	}
	sort.Sort(int64Slice(pids))

	for _, pid := range pids {
		fmt.Fprintf(w, "\nsyscall.Kill: %s [PID: %d]\n", pidToKill[pid], pid)
		if err := syscall.Kill(int(pid), syscall.SIGTERM); err != nil {
			fmt.Fprintf(w, "syscall.SIGTERM error (%v)\n", err)

			shell := os.Getenv("SHELL")
			if len(shell) == 0 {
				shell = "sh"
			}
			args := []string{shell, "-c", fmt.Sprintf("sudo kill -9 %d", pid)}
			cmd := exec.Command(args[0], args[1:]...)
			cmd.Stdout = w
			cmd.Stderr = w
			fmt.Fprintf(w, "Starting: %q\n", strings.Join(cmd.Args, ""))
			if err := cmd.Start(); err != nil {
				fmt.Fprintf(w, "error when 'sudo kill -9' (%v)\n", err)
			}
			if err := cmd.Wait(); err != nil {
				fmt.Fprintf(w, "Start(%s) cmd.Wait returned %v\n", cmd.Path, err)
			}
			fmt.Fprintf(w, "    Done: %q\n", strings.Join(cmd.Args, ""))
		}
		if err := syscall.Kill(int(pid), syscall.SIGKILL); err != nil {
			fmt.Fprintf(w, "syscall.SIGKILL error (%v)\n", err)

			shell := os.Getenv("SHELL")
			if len(shell) == 0 {
				shell = "sh"
			}
			args := []string{shell, "-c", fmt.Sprintf("sudo kill -9 %d", pid)}
			cmd := exec.Command(args[0], args[1:]...)
			cmd.Stdout = w
			cmd.Stderr = w
			fmt.Fprintf(w, "Starting: %q\n", strings.Join(cmd.Args, ""))
			if err := cmd.Start(); err != nil {
				fmt.Fprintf(w, "error when 'sudo kill -9' (%v)\n", err)
			}
			if err := cmd.Wait(); err != nil {
				fmt.Fprintf(w, "Start(%s) cmd.Wait returned %v\n", cmd.Path, err)
			}
			fmt.Fprintf(w, "    Done: %q\n", strings.Join(cmd.Args, ""))
		}
	}
	fmt.Fprintln(w)
}
<commit_msg>update ps<commit_after>package ps

import (
	"encoding/csv"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/gyuho/psn/table"
	"github.com/olekukonko/tablewriter"
)

// ProcessTableColumns is columns for CSV file.
var ProcessTableColumns = []string{
	"NAME",
	"STATE",

	"PID",
	"PPID",

	"CPU",
	"VM_RSS",
	"VM_SIZE",

	"FD",
	"THREADS",

	"CpuUsageFloat64",
	"VmRSSBytes",
	"VmSizeBytes",
}

// List finds the status by specifying the filter.
func List(filter *Process) ([]Process, error) {
	if filter != nil {
		if filter.Stat.Pid != 0 && filter.Status.Pid == 0 {
			filter.Status.Pid = filter.Stat.Pid
		}
		if filter.Stat.Pid == 0 && filter.Status.Pid != 0 {
			filter.Stat.Pid = filter.Status.Pid
		}
	}
	if filter != nil && filter.Stat.Pid != 0 && filter.Status.Pid != 0 {
		stat, err := GetStat(filter.Stat.Pid)
		if err != nil {
			return nil, err
		}
		status, err := GetStatus(filter.Status.Pid)
		if err != nil {
			return nil, err
		}
		return []Process{Process{Stat: stat, Status: status}}, err
	}

	ds, err := ioutil.ReadDir("/proc")
	if err != nil {
		return nil, err
	}
	pids := []int64{}
	for _, f := range ds {
		if f.IsDir() && isInt(f.Name()) {
			i, err := strconv.ParseInt(f.Name(), 10, 64)
			if err != nil {
				return nil, err
			}
			pids = append(pids, i)
		}
	}

	donec, errc := make(chan struct{}), make(chan error)
	rmap := make(map[int64]Process)
	var mu sync.Mutex

	for _, pid := range pids {
		go func(pid int64, filter *Process) {
			stat, err := GetStat(pid)
			if err != nil {
				errc <- err
				return
			}
			if stat.Match(filter) {
				mu.Lock()
				if v, ok := rmap[pid]; !ok {
					rmap[pid] = Process{Stat: stat}
				} else {
					v.Stat = stat
					rmap[pid] = v
				}
				mu.Unlock()
			}
			donec <- struct{}{}
		}(pid, filter)
		go func(pid int64, filter *Process) {
			status, err := GetStatus(pid)
			if err != nil {
				errc <- err
				return
			}
			if status.Match(filter) {
				mu.Lock()
				if v, ok := rmap[pid]; !ok {
					rmap[pid] = Process{Status: status}
				} else {
					v.Status = status
					rmap[pid] = v
				}
				mu.Unlock()
			}
			donec <- struct{}{}
		}(pid, filter)
	}

	cnt := 0
	for cnt != 2*len(pids) {
		select {
		case <-donec:
			cnt++
		case e := <-errc:
			return nil, e
		}
	}

	rs := []Process{}
	for _, proc := range rmap {
		if proc.Stat.Comm == "" && proc.Status.Name != "" {
			proc.Stat.Comm = proc.Status.Name
			// return nil, fmt.Errorf("%s != %s\n", proc.Stat.Comm, proc.Status.Name)
		}
		if proc.Stat.Comm != "" && proc.Status.Name == "" {
			proc.Status.Name = proc.Stat.Comm
		}
		if proc.Stat.Pid == 0 && proc.Status.Pid != 0 {
			proc.Stat.Pid = proc.Status.Pid
		}
		if proc.Stat.Pid != 0 && proc.Status.Pid == 0 {
			proc.Status.Pid = proc.Stat.Pid
		}
		rs = append(rs, proc)
	}
	return rs, nil
}

// Match matches Stat to filter. Only supports name and PID.
func (s *Stat) Match(filter *Process) bool {
	if s == nil {
		return false
	}
	if filter == nil {
		return true // no need to compare
	}
	name := filter.Stat.Comm
	if name != "" {
		if name != s.Comm {
			return false
		}
	}
	pid := filter.Stat.Pid
	if pid != 0 {
		if pid != s.Pid {
			return false
		}
	}
	return true
}

// Match matches Status to filter. Only supports name and PID.
func (s *Status) Match(filter *Process) bool {
	if s == nil {
		return false
	}
	if filter == nil {
		return true // no need to compare
	}
	name := filter.Status.Name
	if name != "" {
		if name != s.Name {
			return false
		}
	}
	pid := filter.Status.Pid
	if pid != 0 {
		if pid != s.Pid {
			return false
		}
	}
	state := filter.Status.State
	if state != "" {
		if state != s.State {
			return false
		}
	}
	return true
}

// WriteToTable writes slice of Process to ASCII table.
func WriteToTable(w io.Writer, top int, pss ...Process) {
	tw := tablewriter.NewWriter(w)
	tw.SetHeader(ProcessTableColumns[:9:9])

	rows := make([][]string, len(pss))
	for i, s := range pss {
		sl := make([]string, len(ProcessTableColumns))
		sl[0] = s.Status.Name
		sl[1] = s.Status.State
		sl[2] = fmt.Sprintf("%d", s.Status.Pid)
		sl[3] = fmt.Sprintf("%d", s.Status.PPid)
		sl[4] = fmt.Sprintf("%3.2f %%", s.Stat.CpuUsage)
		sl[5] = s.Status.VmRSS
		sl[6] = s.Status.VmSize
		sl[7] = fmt.Sprintf("%d", s.Status.FDSize)
		sl[8] = fmt.Sprintf("%d", s.Status.Threads)

		sl[9] = fmt.Sprintf("%3.2f", s.Stat.CpuUsage)
		sl[10] = fmt.Sprintf("%d", s.Status.VmRSSBytes)
		sl[11] = fmt.Sprintf("%d", s.Status.VmSizeBytes)
		rows[i] = sl
	}
	table.By(
		rows,
		table.MakeDescendingIntFunc(10),    // VM_RSS
		table.MakeDescendingFloat64Func(9), // CPU
		table.MakeDescendingIntFunc(11),    // VM_SIZE
	).Sort(rows)

	if top != 0 && len(rows) > top {
		rows = rows[:top:top]
	}
	for _, row := range rows {
		tw.Append(row[:9:9])
	}

	tw.Render()
}

var once sync.Once

// WriteToTable writes slice of Process to a csv file.
func WriteToCSV(f *os.File, pss ...Process) error {
	wr := csv.NewWriter(f)

	var werr error
	writeCSVHeader := func() {
		if err := wr.Write(append([]string{"unix_ts"}, ProcessTableColumns...)); err != nil {
			werr = err
		}
	}
	once.Do(writeCSVHeader)
	if werr != nil {
		return werr
	}

	rows := make([][]string, len(pss))
	for i, s := range pss {
		sl := make([]string, len(ProcessTableColumns))
		sl[0] = s.Status.Name
		sl[1] = s.Status.State
		sl[2] = fmt.Sprintf("%d", s.Status.Pid)
		sl[3] = fmt.Sprintf("%d", s.Status.PPid)
		sl[4] = fmt.Sprintf("%3.2f %%", s.Stat.CpuUsage)
		sl[5] = s.Status.VmRSS
		sl[6] = s.Status.VmSize
		sl[7] = fmt.Sprintf("%d", s.Status.FDSize)
		sl[8] = fmt.Sprintf("%d", s.Status.Threads)

		sl[9] = fmt.Sprintf("%3.2f", s.Stat.CpuUsage)
		sl[10] = fmt.Sprintf("%d", s.Status.VmRSSBytes)
		sl[11] = fmt.Sprintf("%d", s.Status.VmSizeBytes)
		rows[i] = sl
	}
	table.By(
		rows,
		table.MakeDescendingIntFunc(10),    // VM_RSS
		table.MakeDescendingFloat64Func(9), // CPU
		table.MakeDescendingIntFunc(11),    // VM_SIZE
	).Sort(rows)

	ts := fmt.Sprintf("%d", time.Now().Unix())
	nrows := make([][]string, len(rows))
	for i, row := range rows {
		nrows[i] = append([]string{ts}, row...)
	}
	if err := wr.WriteAll(nrows); err != nil {
		return err
	}

	wr.Flush()
	return wr.Error()
}

type int64Slice []int64

func (s int64Slice) Len() int           { return len(s) }
func (s int64Slice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
func (s int64Slice) Less(i, j int) bool { return s[i] < s[j] }

// Kill kills all processes in arguments.
func Kill(w io.Writer, parent bool, pss ...Process) {
	defer func() {
		if err := recover(); err != nil {
			fmt.Fprintln(w, "Kill:", err)
		}
	}()

	pidToKill := make(map[int64]string)
	for _, s := range pss {
		pidToKill[s.Status.Pid] = s.Status.Name
		if parent && s.Status.PPid != 0 {
			pidToKill[s.Status.PPid] = s.Status.Name
		}
	}
	if len(pidToKill) == 0 {
		fmt.Fprintln(w, "no PID to kill...")
		return
	}

	pids := []int64{}
	for pid := range pidToKill {
		pids = append(pids, pid)
	}
	sort.Sort(int64Slice(pids))

	for _, pid := range pids {
		fmt.Fprintf(w, "\nsyscall.Kill: %s [PID: %d]\n", pidToKill[pid], pid)
		if err := syscall.Kill(int(pid), syscall.SIGTERM); err != nil {
			fmt.Fprintf(w, "syscall.SIGTERM error (%v)\n", err)

			shell := os.Getenv("SHELL")
			if len(shell) == 0 {
				shell = "sh"
			}
			args := []string{shell, "-c", fmt.Sprintf("sudo kill -9 %d", pid)}
			cmd := exec.Command(args[0], args[1:]...)
			cmd.Stdout = w
			cmd.Stderr = w
			fmt.Fprintf(w, "Starting: %q\n", strings.Join(cmd.Args, ""))
			if err := cmd.Start(); err != nil {
				fmt.Fprintf(w, "error when 'sudo kill -9' (%v)\n", err)
			}
			if err := cmd.Wait(); err != nil {
				fmt.Fprintf(w, "Start(%s) cmd.Wait returned %v\n", cmd.Path, err)
			}
			fmt.Fprintf(w, "    Done: %q\n", strings.Join(cmd.Args, ""))
		}
		if err := syscall.Kill(int(pid), syscall.SIGKILL); err != nil {
			fmt.Fprintf(w, "syscall.SIGKILL error (%v)\n", err)

			shell := os.Getenv("SHELL")
			if len(shell) == 0 {
				shell = "sh"
			}
			args := []string{shell, "-c", fmt.Sprintf("sudo kill -9 %d", pid)}
			cmd := exec.Command(args[0], args[1:]...)
			cmd.Stdout = w
			cmd.Stderr = w
			fmt.Fprintf(w, "Starting: %q\n", strings.Join(cmd.Args, ""))
			if err := cmd.Start(); err != nil {
				fmt.Fprintf(w, "error when 'sudo kill -9' (%v)\n", err)
			}
			if err := cmd.Wait(); err != nil {
				fmt.Fprintf(w, "Start(%s) cmd.Wait returned %v\n", cmd.Path, err)
			}
			fmt.Fprintf(w, "    Done: %q\n", strings.Join(cmd.Args, ""))
		}
	}
	fmt.Fprintln(w)
}
<|endoftext|>


<commit_before>package render

import (
	"context"
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/lomik/graphite-clickhouse/config"
	"github.com/lomik/graphite-clickhouse/helper/clickhouse"
	"github.com/lomik/graphite-clickhouse/helper/rollup"
	"github.com/lomik/graphite-clickhouse/pkg/alias"
	"github.com/lomik/graphite-clickhouse/pkg/dry"
	"github.com/lomik/graphite-clickhouse/pkg/reverse"
	"github.com/lomik/graphite-clickhouse/pkg/scope"
	"github.com/lomik/graphite-clickhouse/pkg/where"
	"go.uber.org/zap"
)

// from, until, step, function, table, prewhere, where
// arrayFilter(x->isNotNull(x)) - do not pass nulls to client
// -Resample - group time and values by time intervals and apply aggregation function
// -OrNull - if there aren't points in an interval, null will be returned
// intDiv(Time, x)*x - round Time down to step multiplier
const queryAggregated = `SELECT Path,
	arrayFilter(x->isNotNull(x), anyOrNullResample(%[1]d, %[2]d, %[3]d)(toUInt32(intDiv(Time, %[3]d)*%[3]d), Time)),
	arrayFilter(x->isNotNull(x), %[4]sOrNullResample(%[1]d, %[2]d, %[3]d)(Value, Time))
FROM %[5]s
%[6]s
%[7]s
GROUP BY Path
FORMAT RowBinary`

// table, prewhere, where
const queryUnaggregated = `SELECT Path, groupArray(Time), groupArray(Value), groupArray(Timestamp) FROM %s %s %s GROUP BY Path FORMAT RowBinary`

// TimeFrame contains information about fetch request time conditions
type TimeFrame struct {
	From          int64
	Until         int64
	MaxDataPoints int64
}

type Targets struct {
	List        []string
	AM          *alias.Map
	pointsTable string
	isReverse   bool
	rollupObj   *rollup.Rules
}

// MultiFetchRequest is a map of TimeFrame keys and targets slice of strings values
type MultiFetchRequest map[TimeFrame]*Targets

type CHResponse struct {
	Data  *Data
	From  int64
	Until int64
}

type Reply struct {
	CHResponses []CHResponse
	lock        sync.RWMutex
	cStep       *commonStep
}

func (r *Reply) Append(chr CHResponse) {
	r.lock.Lock()
	r.CHResponses = append(r.CHResponses, chr)
	r.lock.Unlock()
}

// EmptyResponse is an empty []CHResponse
var EmptyResponse []CHResponse = []CHResponse{{EmptyData, 0, 0}}

// EmptyReply is used when no metrics/data points are read
var EmptyReply *Reply = &Reply{EmptyResponse, sync.RWMutex{}, nil}

// FetchDataPoints fetches the data from ClickHouse and parses it into []CHResponse
func FetchDataPoints(ctx context.Context, cfg *config.Config, fetchRequests MultiFetchRequest, chContext string) (*Reply, error) {
	var lock sync.RWMutex
	var wg sync.WaitGroup
	logger := scope.Logger(ctx)
	cStep := &commonStep{
		result: 0,
		wg:     sync.WaitGroup{},
		lock:   sync.RWMutex{},
	}

	if cfg.ClickHouse.InternalAggregation {
		cStep.addTargets(len(fetchRequests))
	}

	reply := &Reply{
		CHResponses: make([]CHResponse, 0, len(fetchRequests)),
		lock:        sync.RWMutex{},
		cStep:       cStep,
	}
	errors := make([]error, 0, len(fetchRequests))

	for tf, targets := range fetchRequests {
		if tf.MaxDataPoints <= 0 {
			tf.MaxDataPoints = int64(cfg.ClickHouse.MaxDataPoints)
		}
		targets.pointsTable, targets.isReverse, targets.rollupObj = SelectDataTable(cfg, tf.From, tf.Until, targets.List, chContext)
		if targets.pointsTable == "" {
			err := fmt.Errorf("data tables is not specified for %v", targets.List[0])
			lock.Lock()
			errors = append(errors, err)
			lock.Unlock()
			logger.Error("data tables is not specified", zap.Error(err))
			return EmptyReply, err
		}
		wg.Add(1)
		go func(tf TimeFrame, targets *Targets) {
			defer wg.Done()
			err := reply.getDataPoints(ctx, cfg, tf, targets)
			if err != nil {
				lock.Lock()
				errors = append(errors, err)
				lock.Unlock()
				return
			}
		}(tf, targets)
	}
	wg.Wait()
	for len(errors) != 0 {
		return EmptyReply, errors[0]
	}

	return reply, nil
}

func (r *Reply) getDataPoints(ctx context.Context, cfg *config.Config, tf TimeFrame, targets *Targets) error {
	logger := scope.Logger(ctx)
	var data *Data
	var err error
	var rollupTime time.Duration
	defer func() {
		if rollupTime > 0 {
			logger.Debug(
				"rollup",
				zap.String("runtime", rollupTime.String()),
				zap.Duration("runtime_ns", rollupTime),
			)
		}
	}()

	if cfg.ClickHouse.InternalAggregation {
		data, err = r.getDataAggregated(ctx, cfg, tf, targets)
	} else {
		data, err = r.getDataUnaggregated(ctx, cfg, tf, targets)
	}

	if err != nil {
		return nil
	}

	logger.Info("data", zap.Int("read_bytes", data.length), zap.Int("read_points", data.Points.Len()))

	// ClickHouse returns sorted and uniq values, when internal aggregation is used
	// But if carbonlink is used, we still need to sort, filter and rollup points
	if !cfg.ClickHouse.InternalAggregation || carbonlinkClient(cfg) != nil {
		sortStart := time.Now()
		data.Points.Sort()
		d := time.Since(sortStart)
		logger.Debug("sort", zap.String("runtime", d.String()), zap.Duration("runtime_ns", d))

		data.Points.Uniq()
		rollupStart := time.Now()
		err = data.rollupObj.RollupPoints(data.Points, tf.From, data.commonStep)
		if err != nil {
			logger.Error("rollup failed", zap.Error(err))
			return err
		}
		rollupTime += time.Since(rollupStart)
	}

	data.Aliases = targets.AM

	r.Append(CHResponse{
		Data:  data,
		From:  tf.From,
		Until: tf.Until,
	})
	return nil
}

func (r *Reply) getDataAggregated(ctx context.Context, cfg *config.Config, tf TimeFrame, targets *Targets) (data *Data, err error) {
	// Generic prepare part
	logger := scope.Logger(ctx)
	data = EmptyData

	metricList := targets.AM.Series(targets.isReverse)

	if len(metricList) == 0 {
		return
	}

	// from carbonlink request
	carbonlinkResponseRead := queryCarbonlink(ctx, cfg, metricList)

	now := time.Now().Unix()
	age := dry.Max(0, now-tf.From)

	// end of generic prepare

	maxDataPoints := dry.Min(tf.MaxDataPoints, int64(cfg.ClickHouse.MaxDataPoints))
	metricsAggregation := make(map[string][]string)

	// Grouping metrics by aggregation steps and functions
	var step int64
	for _, m := range metricList {
		newStep, agg := targets.rollupObj.Lookup(m, uint32(age))
		step = r.cStep.calculateUnsafe(step, int64(newStep))
		if mm, ok := metricsAggregation[agg.Name()]; ok {
			metricsAggregation[agg.Name()] = append(mm, m)
		} else {
			metricsAggregation[agg.Name()] = []string{m}
		}
	}
	r.cStep.calculate(step)
	step = dry.Max(r.cStep.getResult(), (tf.Until-tf.From)/maxDataPoints)
	step = dry.CeilToMultiplier(step, r.cStep.getResult())

	b := make(chan io.ReadCloser, len(metricsAggregation))
	e := make(chan error)
	queryWg := sync.WaitGroup{}
	queryContext, cancel := context.WithCancel(ctx)

	defer func() {
		cancel()
		queryWg.Wait()
		close(e)
		close(b)
	}()

	for agg, metricList := range metricsAggregation {
		from := dry.CeilToMultiplier(tf.From, step)
		until := dry.CeilToMultiplier(tf.Until, step) - 1
		pw := where.New()
		pw.And(where.DateBetween("Date", time.Unix(from, 0), time.Unix(until, 0)))

		wr := where.New()
		wr.And(where.In("Path", metricList))

		wr.And(where.TimestampBetween("Time", from, until))
		query := fmt.Sprintf(
			queryAggregated,
			from, until, step, agg,
			targets.pointsTable, pw.PreWhereSQL(), wr.SQL(),
		)
		queryWg.Add(1)
		go func(query string) {
			defer queryWg.Done()
			body, err := clickhouse.Reader(
				scope.WithTable(ctx, targets.pointsTable),
				cfg.ClickHouse.Url,
				query,
				clickhouse.Options{Timeout: cfg.ClickHouse.DataTimeout.Value(), ConnectTimeout: cfg.ClickHouse.ConnectTimeout.Value()},
			)
			if err != nil {
				logger.Error("reader", zap.Error(err))
				select {
				case <-queryContext.Done():
					return
				case e <- err:
					return
				}
			}
			select {
			case <-queryContext.Done():
				return
			case b <- body:
				return
			}
		}(query)
	}

	carbonlinkData := carbonlinkResponseRead()
	data, err = parseAggregatedResponse(b, e, carbonlinkData, targets.isReverse)
	if err != nil {
		logger.Error("data", zap.Error(err), zap.Int("read_bytes", data.length))
		return nil, err
	}
	data.commonStep = step
	return
}

func (r *Reply) getDataUnaggregated(ctx context.Context, cfg *config.Config, tf TimeFrame, targets *Targets) (data *Data, err error) {
	// Generic
	logger := scope.Logger(ctx)
	data = EmptyData

	metricList := targets.AM.Series(targets.isReverse)

	if len(metricList) == 0 {
		return
	}

	// from carbonlink request
	carbonlinkResponseRead := queryCarbonlink(ctx, cfg, metricList)

	now := time.Now().Unix()
	age := dry.Max(0, now-tf.From)

	// end of generic prepare

	// calculate max step

	var maxStep int64
	steps := make(map[string]uint32)
	for _, m := range metricList {
		step, _ := targets.rollupObj.Lookup(m, uint32(age))
		steps[m] = step
		if int64(step) > maxStep {
			maxStep = int64(step)
		}

		if targets.isReverse {
			m = reverse.String(m)
		}
	}
	until := dry.CeilToMultiplier(tf.Until, maxStep) - 1

	pw := where.New()
	pw.And(where.DateBetween("Date", time.Unix(tf.From, 0), time.Unix(until, 0)))

	wr := where.New()
	wr.And(where.In("Path", metricList))

	wr.And(where.TimestampBetween("Time", tf.From, until))

	query := fmt.Sprintf(
		queryUnaggregated,
		targets.pointsTable, pw.PreWhereSQL(), wr.SQL(),
	)

	body, err := clickhouse.Reader(
		scope.WithTable(ctx, targets.pointsTable),
		cfg.ClickHouse.Url,
		query,
		clickhouse.Options{Timeout: cfg.ClickHouse.DataTimeout.Value(), ConnectTimeout: cfg.ClickHouse.ConnectTimeout.Value()},
	)

	if err != nil {
		logger.Error("reader", zap.Error(err))
		return
	}

	// fetch carbonlink response
	carbonlinkData := carbonlinkResponseRead()

	parseStart := time.Now()

	// pass carbonlinkData to data parser
	data, err = parseUnaggregatedResponse(body, carbonlinkData, targets.isReverse)
	if err != nil {
		logger.Error("data", zap.Error(err), zap.Int("read_bytes", data.length))
		return nil, err
	}
	d := time.Since(parseStart)

	data.rollupObj = targets.rollupObj
	data.Points.SetSteps(steps)

	logger.Debug("parse", zap.String("runtime", d.String()), zap.Duration("runtime_ns", d))

	return
}
<commit_msg>fix typo<commit_after>package render

import (
	"context"
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/lomik/graphite-clickhouse/config"
	"github.com/lomik/graphite-clickhouse/helper/clickhouse"
	"github.com/lomik/graphite-clickhouse/helper/rollup"
	"github.com/lomik/graphite-clickhouse/pkg/alias"
	"github.com/lomik/graphite-clickhouse/pkg/dry"
	"github.com/lomik/graphite-clickhouse/pkg/reverse"
	"github.com/lomik/graphite-clickhouse/pkg/scope"
	"github.com/lomik/graphite-clickhouse/pkg/where"
	"go.uber.org/zap"
)

// from, until, step, function, table, prewhere, where
// arrayFilter(x->isNotNull(x)) - do not pass nulls to client
// -Resample - group time and values by time intervals and apply aggregation function
// -OrNull - if there aren't points in an interval, null will be returned
// intDiv(Time, x)*x - round Time down to step multiplier
const queryAggregated = `SELECT Path,
	arrayFilter(x->isNotNull(x), anyOrNullResample(%[1]d, %[2]d, %[3]d)(toUInt32(intDiv(Time, %[3]d)*%[3]d), Time)),
	arrayFilter(x->isNotNull(x), %[4]sOrNullResample(%[1]d, %[2]d, %[3]d)(Value, Time))
FROM %[5]s
%[6]s
%[7]s
GROUP BY Path
FORMAT RowBinary`

// table, prewhere, where
const queryUnaggregated = `SELECT Path, groupArray(Time), groupArray(Value), groupArray(Timestamp) FROM %s %s %s GROUP BY Path FORMAT RowBinary`

// TimeFrame contains information about fetch request time conditions
type TimeFrame struct {
	From          int64
	Until         int64
	MaxDataPoints int64
}

type Targets struct {
	List        []string
	AM          *alias.Map
	pointsTable string
	isReverse   bool
	rollupObj   *rollup.Rules
}

// MultiFetchRequest is a map of TimeFrame keys and targets slice of strings values
type MultiFetchRequest map[TimeFrame]*Targets

type CHResponse struct {
	Data  *Data
	From  int64
	Until int64
}

type Reply struct {
	CHResponses []CHResponse
	lock        sync.RWMutex
	cStep       *commonStep
}

func (r *Reply) Append(chr CHResponse) {
	r.lock.Lock()
	r.CHResponses = append(r.CHResponses, chr)
	r.lock.Unlock()
}

// EmptyResponse is an empty []CHResponse
var EmptyResponse []CHResponse = []CHResponse{{EmptyData, 0, 0}}

// EmptyReply is used when no metrics/data points are read
var EmptyReply *Reply = &Reply{EmptyResponse, sync.RWMutex{}, nil}

// FetchDataPoints fetches the data from ClickHouse and parses it into []CHResponse
func FetchDataPoints(ctx context.Context, cfg *config.Config, fetchRequests MultiFetchRequest, chContext string) (*Reply, error) {
	var lock sync.RWMutex
	var wg sync.WaitGroup
	logger := scope.Logger(ctx)
	cStep := &commonStep{
		result: 0,
		wg:     sync.WaitGroup{},
		lock:   sync.RWMutex{},
	}

	if cfg.ClickHouse.InternalAggregation {
		cStep.addTargets(len(fetchRequests))
	}

	reply := &Reply{
		CHResponses: make([]CHResponse, 0, len(fetchRequests)),
		lock:        sync.RWMutex{},
		cStep:       cStep,
	}
	errors := make([]error, 0, len(fetchRequests))

	for tf, targets := range fetchRequests {
		if tf.MaxDataPoints <= 0 {
			tf.MaxDataPoints = int64(cfg.ClickHouse.MaxDataPoints)
		}
		targets.pointsTable, targets.isReverse, targets.rollupObj = SelectDataTable(cfg, tf.From, tf.Until, targets.List, chContext)
		if targets.pointsTable == "" {
			err := fmt.Errorf("data tables is not specified for %v", targets.List[0])
			lock.Lock()
			errors = append(errors, err)
			lock.Unlock()
			logger.Error("data tables is not specified", zap.Error(err))
			return EmptyReply, err
		}
		wg.Add(1)
		go func(tf TimeFrame, targets *Targets) {
			defer wg.Done()
			err := reply.getDataPoints(ctx, cfg, tf, targets)
			if err != nil {
				lock.Lock()
				errors = append(errors, err)
				lock.Unlock()
				return
			}
		}(tf, targets)
	}
	wg.Wait()
	for len(errors) != 0 {
		return EmptyReply, errors[0]
	}

	return reply, nil
}

func (r *Reply) getDataPoints(ctx context.Context, cfg *config.Config, tf TimeFrame, targets *Targets) error {
	logger := scope.Logger(ctx)
	var data *Data
	var err error
	var rollupTime time.Duration
	defer func() {
		if rollupTime > 0 {
			logger.Debug(
				"rollup",
				zap.String("runtime", rollupTime.String()),
				zap.Duration("runtime_ns", rollupTime),
			)
		}
	}()

	if cfg.ClickHouse.InternalAggregation {
		data, err = r.getDataAggregated(ctx, cfg, tf, targets)
	} else {
		data, err = r.getDataUnaggregated(ctx, cfg, tf, targets)
	}

	if err != nil {
		return nil
	}

	logger.Info("data", zap.Int("read_bytes", data.length), zap.Int("read_points", data.Points.Len()))

	// ClickHouse returns sorted and uniq values, when internal aggregation is used
	// But if carbonlink is used, we still need to sort, filter and rollup points
	if !cfg.ClickHouse.InternalAggregation || carbonlinkClient(cfg) != nil {
		sortStart := time.Now()
		data.Points.Sort()
		d := time.Since(sortStart)
		logger.Debug("sort", zap.String("runtime", d.String()), zap.Duration("runtime_ns", d))

		data.Points.Uniq()
		rollupStart := time.Now()
		err = data.rollupObj.RollupPoints(data.Points, tf.From, data.commonStep)
		if err != nil {
			logger.Error("rollup failed", zap.Error(err))
			return err
		}
		rollupTime += time.Since(rollupStart)
	}

	data.Aliases = targets.AM

	r.Append(CHResponse{
		Data:  data,
		From:  tf.From,
		Until: tf.Until,
	})
	return nil
}

func (r *Reply) getDataAggregated(ctx context.Context, cfg *config.Config, tf TimeFrame, targets *Targets) (data *Data, err error) {
	// Generic prepare part
	logger := scope.Logger(ctx)
	data = EmptyData

	metricList := targets.AM.Series(targets.isReverse)

	if len(metricList) == 0 {
		return
	}

	// from carbonlink request
	carbonlinkResponseRead := queryCarbonlink(ctx, cfg, metricList)

	now := time.Now().Unix()
	age := dry.Max(0, now-tf.From)

	// end of generic prepare

	maxDataPoints := dry.Min(tf.MaxDataPoints, int64(cfg.ClickHouse.MaxDataPoints))
	metricsAggregation := make(map[string][]string)

	// Grouping metrics by aggregation steps and functions
	var step int64
	for _, m := range metricList {
		newStep, agg := targets.rollupObj.Lookup(m, uint32(age))
		step = r.cStep.calculateUnsafe(step, int64(newStep))
		if mm, ok := metricsAggregation[agg.Name()]; ok {
			metricsAggregation[agg.Name()] = append(mm, m)
		} else {
			metricsAggregation[agg.Name()] = []string{m}
		}
	}
	r.cStep.calculate(step)
	step = dry.Max(r.cStep.getResult(), (tf.Until-tf.From)/maxDataPoints)
	step = dry.CeilToMultiplier(step, r.cStep.getResult())

	b := make(chan io.ReadCloser, len(metricsAggregation))
	e := make(chan error)
	queryWg := sync.WaitGroup{}
	queryContext, cancel := context.WithCancel(ctx)

	defer func() {
		cancel()
		queryWg.Wait()
		close(e)
		close(b)
	}()

	for agg, metricList := range metricsAggregation {
		from := dry.CeilToMultiplier(tf.From, step)
		until := dry.CeilToMultiplier(tf.Until, step) - 1
		pw := where.New()
		pw.And(where.DateBetween("Date", time.Unix(from, 0), time.Unix(until, 0)))

		wr := where.New()
		wr.And(where.In("Path", metricList))

		wr.And(where.TimestampBetween("Time", from, until))
		query := fmt.Sprintf(
			queryAggregated,
			from, until, step, agg,
			targets.pointsTable, pw.PreWhereSQL(), wr.SQL(),
		)
		queryWg.Add(1)
		go func(query string) {
			defer queryWg.Done()
			body, err := clickhouse.Reader(
				scope.WithTable(ctx, targets.pointsTable),
				cfg.ClickHouse.Url,
				query,
				clickhouse.Options{Timeout: cfg.ClickHouse.DataTimeout.Value(), ConnectTimeout: cfg.ClickHouse.ConnectTimeout.Value()},
			)
			if err != nil {
				logger.Error("reader", zap.Error(err))
				select {
				case <-queryContext.Done():
					return
				case e <- err:
					return
				}
			}
			select {
			case <-queryContext.Done():
				return
			case b <- body:
				return
			}
		}(query)
	}

	carbonlinkData := carbonlinkResponseRead()
	data, err = parseAggregatedResponse(b, e, carbonlinkData, targets.isReverse)
	if err != nil {
		logger.Error("data", zap.Error(err), zap.Int("read_bytes", data.length))
		return nil, err
	}
	data.commonStep = step
	return
}

func (r *Reply) getDataUnaggregated(ctx context.Context, cfg *config.Config, tf TimeFrame, targets *Targets) (data *Data, err error) {
	// Generic
	logger := scope.Logger(ctx)
	data = EmptyData

	metricList := targets.AM.Series(targets.isReverse)

	if len(metricList) == 0 {
		return
	}

	// from carbonlink request
	carbonlinkResponseRead := queryCarbonlink(ctx, cfg, metricList)

	now := time.Now().Unix()
	age := dry.Max(0, now-tf.From)

	// end of generic prepare

	// calculate max step

	var maxStep int64
	steps := make(map[string]uint32)
	for _, m := range metricList {
		step, _ := targets.rollupObj.Lookup(m, uint32(age))
		if int64(step) > maxStep {
			maxStep = int64(step)
		}

		if targets.isReverse {
			m = reverse.String(m)
		}

		steps[m] = step
	}
	until := dry.CeilToMultiplier(tf.Until, maxStep) - 1

	pw := where.New()
	pw.And(where.DateBetween("Date", time.Unix(tf.From, 0), time.Unix(until, 0)))

	wr := where.New()
	wr.And(where.In("Path", metricList))

	wr.And(where.TimestampBetween("Time", tf.From, until))

	query := fmt.Sprintf(
		queryUnaggregated,
		targets.pointsTable, pw.PreWhereSQL(), wr.SQL(),
	)

	body, err := clickhouse.Reader(
		scope.WithTable(ctx, targets.pointsTable),
		cfg.ClickHouse.Url,
		query,
		clickhouse.Options{Timeout: cfg.ClickHouse.DataTimeout.Value(), ConnectTimeout: cfg.ClickHouse.ConnectTimeout.Value()},
	)

	if err != nil {
		logger.Error("reader", zap.Error(err))
		return
	}

	// fetch carbonlink response
	carbonlinkData := carbonlinkResponseRead()

	parseStart := time.Now()

	// pass carbonlinkData to data parser
	data, err = parseUnaggregatedResponse(body, carbonlinkData, targets.isReverse)
	if err != nil {
		logger.Error("data", zap.Error(err), zap.Int("read_bytes", data.length))
		return nil, err
	}
	d := time.Since(parseStart)

	data.rollupObj = targets.rollupObj
	data.Points.SetSteps(steps)

	logger.Debug("parse", zap.String("runtime", d.String()), zap.Duration("runtime_ns", d))

	return
}
<|endoftext|>


<commit_before><commit_msg>In & out client registration change<commit_after><|endoftext|>


<commit_before>package reqctx

import (
	"context"
	"net/http"

	"github.com/anacrolix/missinggo/assert"
)

func NewValue() *contextValue {
	return &contextValue{new(byte)}
}

type contextValue struct {
	key interface{}
}

func (me contextValue) Get(ctx context.Context) interface{} {
	return ctx.Value(me.key)
}

func (me contextValue) SetRequestOnce(r *http.Request, val interface{}) *http.Request {
	assert.Nil(me.Get(r.Context()))
	return r.WithContext(context.WithValue(r.Context(), me.key, val))
}

func (me contextValue) SetMiddleware(val interface{}) func(http.Handler) http.Handler {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			r = me.SetRequestOnce(r, val)
			h.ServeHTTP(w, r)
		})
	}
}
<commit_msg>reqctx: Some comments<commit_after>package reqctx

import (
	"context"
	"net/http"

	"github.com/anacrolix/missinggo/assert"
)

func NewValue() *contextValue {
	return &contextValue{new(byte)}
}

type contextValue struct {
	key interface{}
}

func (me contextValue) Get(ctx context.Context) interface{} {
	return ctx.Value(me.key)
}

// Sets the value on the Request. It must not have been already set.
func (me contextValue) SetRequestOnce(r *http.Request, val interface{}) *http.Request {
	assert.Nil(me.Get(r.Context()))
	return r.WithContext(context.WithValue(r.Context(), me.key, val))
}

// Returns a middleware that sets the value in the Request's Context.
func (me contextValue) SetMiddleware(val interface{}) func(http.Handler) http.Handler {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			r = me.SetRequestOnce(r, val)
			h.ServeHTTP(w, r)
		})
	}
}
<|endoftext|>