<commit_before>package gpg2

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"

	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	"go.mozilla.org/autograph/signer"
)

const (
	// Type of this signer is "gpg2" represents a signer that
	// shells out to gpg2 to sign artifacts since the golang "pgp"
	// signer doesn't support signing with subkeys
	// https://godoc.org/golang.org/x/crypto/openpgp#ArmoredDetachSign
	// or loading keys exported with gnu-dummy s2k encrypted
	// passphrases https://github.com/golang/go/issues/13605
	Type = "gpg2"

	keyRingFilename = "autograph_gpg2_keyring.gpg"
	secRingFilename = "autograph_gpg2_secring.gpg"
)

// GPG2Signer holds the configuration of the signer
type GPG2Signer struct {
	signer.Configuration

	// KeyID is the fingerprint of the gpg key or subkey to use
	// e.g. 0xA2B637F535A86009
	KeyID string

	// Passphrase is the optional passphrase to use decrypt the
	// gpg secret key
	Passphrase string

	// tmpDir is the signer's temporary working directory. It
	// holds the gpg sec and keyrings
	tmpDir string
}

// New initializes a pgp signer using a configuration
func New(conf signer.Configuration) (s *GPG2Signer, err error) {
	s = new(GPG2Signer)

	if conf.Type != Type {
		return nil, errors.Errorf("gpg2: invalid type %q, must be %q", conf.Type, Type)
	}
	s.Type = conf.Type

	if conf.ID == "" {
		return nil, errors.New("gpg2: missing signer ID in signer configuration")
	}
	s.ID = conf.ID

	if conf.PrivateKey == "" {
		return nil, errors.New("gpg2: missing private key in signer configuration")
	}
	s.PrivateKey = conf.PrivateKey

	if conf.PublicKey == "" {
		return nil, errors.New("gpg2: missing public key in signer configuration")
	}
	s.PublicKey = conf.PublicKey

	if conf.KeyID == "" {
		return nil, errors.New("gpg2: missing gpg key ID in signer configuration")
	}
	s.KeyID = conf.KeyID

	s.Passphrase = conf.Passphrase

	s.tmpDir, err = createKeyRing(s)
	if err != nil {
		return nil, errors.Wrap(err, "gpg2: error creating keyring")
	}
	return
}

// createKeyRing creates a temporary gpg sec and keyrings, loads the
// private and public keys for the signer, and returns the temporary
// director holding the rings
func createKeyRing(s *GPG2Signer) (string, error) {
	// reuse keyring in tempdir
	prefix := fmt.Sprintf("autograph_%s_%s", s.Type, s.KeyID)

	dir, err := ioutil.TempDir("", prefix)
	if err != nil {
		return "", errors.Wrap(err, "gpg2: error creating tempdir for keyring")
	}

	// write the public key to a temp file in our signer's temp dir
	tmpPublicKeyFile, err := ioutil.TempFile(dir, "gpg2_publickey")
	if err != nil {
		return "", errors.Wrap(err, "gpg2: error creating tempfile for public key")
	}
	defer os.Remove(tmpPublicKeyFile.Name())
	err = ioutil.WriteFile(tmpPublicKeyFile.Name(), []byte(s.PublicKey), 0755)
	if err != nil {
		return "", errors.Wrap(err, "gpg2: error writing public key to tempfile")
	}

	// write the private key to a temp file in our signer's temp dir
	tmpPrivateKeyFile, err := ioutil.TempFile(dir, "gpg2_privatekey")
	if err != nil {
		return "", errors.Wrap(err, "gpg2: error creating tempfile for private key")
	}
	defer os.Remove(tmpPrivateKeyFile.Name())
	err = ioutil.WriteFile(tmpPrivateKeyFile.Name(), []byte(s.PrivateKey), 0755)
	if err != nil {
		return "", errors.Wrap(err, "gpg2: error writing private key to tempfile")
	}

	keyRingPath := filepath.Join(dir, keyRingFilename)
	secRingPath := filepath.Join(dir, secRingFilename)

	// call gpg to create a new keyring and load the public key in it
	gpgLoadPublicKey := exec.Command("gpg",
		"--no-default-keyring",
		"--keyring", keyRingPath,
		"--secret-keyring", secRingPath,
		"--no-tty",
		"--batch",
		"--yes",
		"--import", tmpPublicKeyFile.Name(),
	)
	out, err := gpgLoadPublicKey.CombinedOutput()
	if err != nil {
		return "", errors.Wrapf(err, "gpg2: failed to load public key into keyring: %s\n%s", err, out)
	} else {
		log.Debugf(fmt.Sprintf("gpg2: loaded public key %s", string(out)))
	}

	// call gpg to load the private key in it
	gpgLoadPrivateKey := exec.Command("gpg", "--no-default-keyring",
		"--keyring", keyRingPath,
		"--secret-keyring", secRingPath,
		"--no-tty",
		"--batch",
		"--yes",
		"--import", tmpPrivateKeyFile.Name())
	out, err = gpgLoadPrivateKey.CombinedOutput()
	if err != nil {
		return "", errors.Wrapf(err, "gpg2: failed to load private key into keyring: %s\n%s", err, out)
	} else {
		log.Debugf(fmt.Sprintf("gpg2: loaded private key %s", string(out)))
	}

	return dir, nil

}

// Config returns the configuration of the current signer
func (s *GPG2Signer) Config() signer.Configuration {
	return signer.Configuration{
		ID:         s.ID,
		Type:       s.Type,
		PrivateKey: s.PrivateKey,
		PublicKey:  s.PublicKey,
	}
}

// SignData takes data and returns an armored signature with pgp header and footer
func (s *GPG2Signer) SignData(data []byte, options interface{}) (signer.Signature, error) {
	keyRingPath := filepath.Join(s.tmpDir, keyRingFilename)
	secRingPath := filepath.Join(s.tmpDir, secRingFilename)

	// write the input to a temp file
	tmpContentFile, err := ioutil.TempFile(s.tmpDir, fmt.Sprintf("gpg2_%s_input", s.ID))
	if err != nil {
		return nil, errors.Wrap(err, "gpg2: failed to create tempfile for input to sign")
	}
	defer os.Remove(tmpContentFile.Name())
	ioutil.WriteFile(tmpContentFile.Name(), data, 0755)

	gpgVerifySig := exec.Command("gpg",
		"--no-default-keyring",
		"--keyring", keyRingPath,
		"--secret-keyring", secRingPath,
		"--armor",
		"--no-tty",
		"--batch",
		"--yes",
		"--sign-with", s.KeyID,
		"--output", "-",
		"--pinentry-mode", "loopback",
		"--passphrase-fd", "0",
		"--detach-sign", tmpContentFile.Name(),
	)
	stdin, err := gpgVerifySig.StdinPipe()
	if err != nil {
		return nil, errors.Wrap(err, "gpg2: failed to create stdin pipe for sign cmd")
	}
	go func() {
		defer stdin.Close()
		io.WriteString(stdin, s.Passphrase)
	}()
	out, err := gpgVerifySig.CombinedOutput()
	if err != nil {
		return nil, errors.Wrapf(err, "gpg2: failed to sign input %s\n%s", err, out)
	} else {
		log.Debugf("signed as:\n%s\n", string(out))
	}

	sig := new(Signature)
	sig.Data = out
	return sig, nil
}

// Signature is a PGP signature
type Signature struct {
	Data []byte
}

// Marshal doesn't do much for this signer. sig.Data already contains
// an armored signature, so we simply convert it to a string and return it
func (sig *Signature) Marshal() (string, error) {
	return string(sig.Data), nil
}

// Unmarshal also does very little. It simply converts the armored signature
// from a string to an []byte, but doesn't attempt to parse it, and returns it
// as a Signature
func Unmarshal(sigstr string) (signer.Signature, error) {
	sig := new(Signature)
	sig.Data = []byte(sigstr)
	return sig, nil
}

// Options are not implemented for this signer
type Options struct {
}

// GetDefaultOptions returns default options of the signer
func (s *GPG2Signer) GetDefaultOptions() interface{} {
	return Options{}
}
<commit_msg>signer: register cleanup handler for gpg2 signers<commit_after>package gpg2

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"syscall"

	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	"go.mozilla.org/autograph/signer"
)

const (
	// Type of this signer is "gpg2" represents a signer that
	// shells out to gpg2 to sign artifacts since the golang "pgp"
	// signer doesn't support signing with subkeys
	// https://godoc.org/golang.org/x/crypto/openpgp#ArmoredDetachSign
	// or loading keys exported with gnu-dummy s2k encrypted
	// passphrases https://github.com/golang/go/issues/13605
	Type = "gpg2"

	keyRingFilename = "autograph_gpg2_keyring.gpg"
	secRingFilename = "autograph_gpg2_secring.gpg"
)

// GPG2Signer holds the configuration of the signer
type GPG2Signer struct {
	signer.Configuration

	// KeyID is the fingerprint of the gpg key or subkey to use
	// e.g. 0xA2B637F535A86009
	KeyID string

	// Passphrase is the optional passphrase to use decrypt the
	// gpg secret key
	Passphrase string

	// tmpDir is the signer's temporary working directory. It
	// holds the gpg sec and keyrings
	tmpDir string
}

// New initializes a pgp signer using a configuration
func New(conf signer.Configuration) (s *GPG2Signer, err error) {
	s = new(GPG2Signer)

	if conf.Type != Type {
		return nil, errors.Errorf("gpg2: invalid type %q, must be %q", conf.Type, Type)
	}
	s.Type = conf.Type

	if conf.ID == "" {
		return nil, errors.New("gpg2: missing signer ID in signer configuration")
	}
	s.ID = conf.ID

	if conf.PrivateKey == "" {
		return nil, errors.New("gpg2: missing private key in signer configuration")
	}
	s.PrivateKey = conf.PrivateKey

	if conf.PublicKey == "" {
		return nil, errors.New("gpg2: missing public key in signer configuration")
	}
	s.PublicKey = conf.PublicKey

	if conf.KeyID == "" {
		return nil, errors.New("gpg2: missing gpg key ID in signer configuration")
	}
	s.KeyID = conf.KeyID

	s.Passphrase = conf.Passphrase

	s.tmpDir, err = createKeyRing(s)
	if err != nil {
		return nil, errors.Wrap(err, "gpg2: error creating keyring")
	}
	s.RegisterCleanupAndExit()
	return
}

// createKeyRing creates a temporary gpg sec and keyrings, loads the
// private and public keys for the signer, and returns the temporary
// director holding the rings
func createKeyRing(s *GPG2Signer) (string, error) {
	// reuse keyring in tempdir
	prefix := fmt.Sprintf("autograph_%s_%s", s.Type, s.KeyID)

	dir, err := ioutil.TempDir("", prefix)
	if err != nil {
		return "", errors.Wrap(err, "gpg2: error creating tempdir for keyring")
	}

	// write the public key to a temp file in our signer's temp dir
	tmpPublicKeyFile, err := ioutil.TempFile(dir, "gpg2_publickey")
	if err != nil {
		return "", errors.Wrap(err, "gpg2: error creating tempfile for public key")
	}
	defer os.Remove(tmpPublicKeyFile.Name())
	err = ioutil.WriteFile(tmpPublicKeyFile.Name(), []byte(s.PublicKey), 0755)
	if err != nil {
		return "", errors.Wrap(err, "gpg2: error writing public key to tempfile")
	}

	// write the private key to a temp file in our signer's temp dir
	tmpPrivateKeyFile, err := ioutil.TempFile(dir, "gpg2_privatekey")
	if err != nil {
		return "", errors.Wrap(err, "gpg2: error creating tempfile for private key")
	}
	defer os.Remove(tmpPrivateKeyFile.Name())
	err = ioutil.WriteFile(tmpPrivateKeyFile.Name(), []byte(s.PrivateKey), 0755)
	if err != nil {
		return "", errors.Wrap(err, "gpg2: error writing private key to tempfile")
	}

	keyRingPath := filepath.Join(dir, keyRingFilename)
	secRingPath := filepath.Join(dir, secRingFilename)

	// call gpg to create a new keyring and load the public key in it
	gpgLoadPublicKey := exec.Command("gpg",
		"--no-default-keyring",
		"--keyring", keyRingPath,
		"--secret-keyring", secRingPath,
		"--no-tty",
		"--batch",
		"--yes",
		"--import", tmpPublicKeyFile.Name(),
	)
	out, err := gpgLoadPublicKey.CombinedOutput()
	if err != nil {
		return "", errors.Wrapf(err, "gpg2: failed to load public key into keyring: %s\n%s", err, out)
	} else {
		log.Debugf(fmt.Sprintf("gpg2: loaded public key %s", string(out)))
	}

	// call gpg to load the private key in it
	gpgLoadPrivateKey := exec.Command("gpg", "--no-default-keyring",
		"--keyring", keyRingPath,
		"--secret-keyring", secRingPath,
		"--no-tty",
		"--batch",
		"--yes",
		"--import", tmpPrivateKeyFile.Name())
	out, err = gpgLoadPrivateKey.CombinedOutput()
	if err != nil {
		return "", errors.Wrapf(err, "gpg2: failed to load private key into keyring: %s\n%s", err, out)
	} else {
		log.Debugf(fmt.Sprintf("gpg2: loaded private key %s", string(out)))
	}

	return dir, nil

}

func (s *GPG2Signer) RegisterCleanupAndExit() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, os.Kill, syscall.SIGTERM)

	go func() {
		sig := <-c
		os.RemoveAll(s.tmpDir)
		log.Infof("gpg2: received signal %s; cleaned up %s and exiting", sig, s.tmpDir)
		os.Exit(0)
	}()
}

// Config returns the configuration of the current signer
func (s *GPG2Signer) Config() signer.Configuration {
	return signer.Configuration{
		ID:         s.ID,
		Type:       s.Type,
		PrivateKey: s.PrivateKey,
		PublicKey:  s.PublicKey,
	}
}

// SignData takes data and returns an armored signature with pgp header and footer
func (s *GPG2Signer) SignData(data []byte, options interface{}) (signer.Signature, error) {
	keyRingPath := filepath.Join(s.tmpDir, keyRingFilename)
	secRingPath := filepath.Join(s.tmpDir, secRingFilename)

	// write the input to a temp file
	tmpContentFile, err := ioutil.TempFile(s.tmpDir, fmt.Sprintf("gpg2_%s_input", s.ID))
	if err != nil {
		return nil, errors.Wrap(err, "gpg2: failed to create tempfile for input to sign")
	}
	defer os.Remove(tmpContentFile.Name())
	ioutil.WriteFile(tmpContentFile.Name(), data, 0755)

	gpgVerifySig := exec.Command("gpg",
		"--no-default-keyring",
		"--keyring", keyRingPath,
		"--secret-keyring", secRingPath,
		"--armor",
		"--no-tty",
		"--batch",
		"--yes",
		"--sign-with", s.KeyID,
		"--output", "-",
		"--pinentry-mode", "loopback",
		"--passphrase-fd", "0",
		"--detach-sign", tmpContentFile.Name(),
	)
	stdin, err := gpgVerifySig.StdinPipe()
	if err != nil {
		return nil, errors.Wrap(err, "gpg2: failed to create stdin pipe for sign cmd")
	}
	go func() {
		defer stdin.Close()
		io.WriteString(stdin, s.Passphrase)
	}()
	out, err := gpgVerifySig.CombinedOutput()
	if err != nil {
		return nil, errors.Wrapf(err, "gpg2: failed to sign input %s\n%s", err, out)
	} else {
		log.Debugf("signed as:\n%s\n", string(out))
	}

	sig := new(Signature)
	sig.Data = out
	return sig, nil
}

// Signature is a PGP signature
type Signature struct {
	Data []byte
}

// Marshal doesn't do much for this signer. sig.Data already contains
// an armored signature, so we simply convert it to a string and return it
func (sig *Signature) Marshal() (string, error) {
	return string(sig.Data), nil
}

// Unmarshal also does very little. It simply converts the armored signature
// from a string to an []byte, but doesn't attempt to parse it, and returns it
// as a Signature
func Unmarshal(sigstr string) (signer.Signature, error) {
	sig := new(Signature)
	sig.Data = []byte(sigstr)
	return sig, nil
}

// Options are not implemented for this signer
type Options struct {
}

// GetDefaultOptions returns default options of the signer
func (s *GPG2Signer) GetDefaultOptions() interface{} {
	return Options{}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strconv"
	"sync"
	"time"

	"github.com/jbreitbart/coBench/bit"
	"github.com/montanaflynn/stats"
)

// global command line parameters
var runs *int
var cpus [2]string
var threads *string
var hermitcore *bool
var resctrlPath *string

var cat *bool
var catBitChunk *int

func main() {
	commandFile := parseArgs()

	commands, err := readCommands(*commandFile)
	if err != nil {
		log.Fatalf("Error reading command file %v: %v", *commandFile, err)
	}
	if len(commands) < 2 {
		log.Fatal("You must provide at least 2 commands")
	}

	var minBits int
	var numBits int

	if *cat {
		var err error
		minBits, numBits, err = setupCAT()
		if err != nil {
			log.Fatalf("%v\n", err)
		}
	} else {
		minBits = 0
		numBits = 0
	}

	commandPairs := generateCommandPairs(commands)

	fmt.Println("Executing the following command pairs:")
	for _, c := range commandPairs {
		fmt.Println(c)
	}

	for i, c := range commandPairs {
		fmt.Printf("Running pair %v\n", i)
		fmt.Println(c)
		for bits := minBits; bits <= numBits-minBits; bits += *catBitChunk {
			bitsets := []int64{0, 0}
			bitsets[0] = bit.SetFirstN(bitsets[0], bits)
			bitsets[1] = bit.SetLastN(bitsets[1], bits, numBits-1)

			err := runPair(c, i, bitsets)
			if err != nil {
				log.Fatalf("Error while running pair %v (%v): %v", i, c, err)
			}
		}
	}
}

func runCmdMinTimes(cmd *exec.Cmd, min int, catMask int64, wg *sync.WaitGroup, measurement *string, done chan int, errs chan error) {
	var runtime []float64

	defer wg.Done()

	defer func() {
		mean, _ := stats.Mean(runtime)
		stddev, _ := stats.StandardDeviation(runtime)
		vari, _ := stats.Variance(runtime)

		s := fmt.Sprintf("%v \t %9.2fs avg. runtime \t %1.6f std. dev. \t %1.6f variance \t %v runs", cmd.Args, mean, stddev, vari, len(runtime))
		if *cat {
			s += fmt.Sprintf("\t %x", (uint)(catMask))
		}
		s += "\n"
		fmt.Println(s)
	}()

	for i := 1; ; i++ {
		// create a copy of the command
		cmd := *cmd

		start := time.Now()
		err := cmd.Run()
		elapsed := time.Since(start)

		if err != nil {
			errs <- fmt.Errorf("Error running %v: %v", cmd.Args, err)
			return
		}

		// did the other cmd result in an error?
		if len(errs) != 0 {
			return
		}

		d := <-done

		// check if the other application was running the whole time
		if d == len(cpus) {
			// no
			*measurement += "# "
		} else {
			runtime = append(runtime, elapsed.Seconds())
		}
		*measurement += strconv.FormatInt(elapsed.Nanoseconds(), 10)
		*measurement += "\n"

		// did we run min times?
		if i == min {
			d++
		}
		done <- d

		// both applications are done
		if d == len(cpus) {
			return
		}
	}
}

func runPair(cPair [2]string, id int, catConfig []int64) error {
	env := os.Environ()

	if *cat {
		if err := writeCATConfig(catConfig); err != nil {
			return fmt.Errorf("Error while writting CAT config: %v", err)
		}
	}

	var cmds [len(cpus)]*exec.Cmd
	// setup commands
	for i, _ := range cmds {
		if *hermitcore {
			cmds[i] = exec.Command("numactl", "--physcpubind", cpus[i], "/bin/sh", "-c", cPair[i])
			cmds[i].Env = append(env, "HERMIT_CPUS="+*threads, "HERMIT_MEM=4G", "HERMIT_ISLE=uhyve")
		} else {
			cmds[i] = exec.Command("/bin/sh", "-c", cPair[i])
			cmds[i].Env = append(env, "GOMP_CPU_AFFINITY="+cpus[i], "OMP_NUM_THREADS="+*threads)
		}

		outfile, err := os.Create(fmt.Sprintf("%v-%v-%x.log", id, i, catConfig[i]))
		if err != nil {
			return fmt.Errorf("Error while creating file: %v", err)
		}
		defer outfile.Close()
		cmds[i].Stdout = outfile
	}

	var measurements [len(cmds)]string
	// used to count how many apps have reached their min limit
	done := make(chan int, 1)
	done <- 0

	// used to return an error from the go-routines
	errs := make(chan error, len(cmds))

	// used to wait for the following 2 goroutines
	var wg sync.WaitGroup
	wg.Add(len(cmds))

	for i, c := range cmds {
		go runCmdMinTimes(c, *runs, catConfig[i], &wg, &measurements[i], done, errs)
	}

	wg.Wait()

	for i, s := range measurements {
		measurementsFile, err := os.Create(fmt.Sprintf("%v-%v.time", id, i))
		if err != nil {
			return fmt.Errorf("Error while creating file: %v", err)
		}
		defer measurementsFile.Close()

		out := "# runtime in nanoseconds of \"" + cPair[i] + "\" on CPUs " + cpus[i] + "while \"" + cPair[(i+1)%2] + "\" was running on cores " + cpus[(i+1)%len(cpus)]
		if *cat {
			out += fmt.Sprintf(" with CAT %x ", catConfig[i])
		}
		out += "\n"
		out += s

		_, err = measurementsFile.WriteString(out)
		if err != nil {
			return fmt.Errorf("Error while writing measurements file: %v", err)
		}
	}

	if len(errs) != 0 {
		return <-errs
	}

	return nil
}

func generateCommandPairs(commands []string) [][2]string {
	var pairs [][2]string
	for i, c0 := range commands {
		for j, c1 := range commands {
			if i >= j {
				continue
			}
			pairs = append(pairs, [2]string{c0, c1})
		}
	}
	return pairs
}

func readCommands(filename string) ([]string, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, errors.New("Error opening file " + filename + ": " + err.Error())
	}
	defer file.Close()

	var commands []string

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		commands = append(commands, scanner.Text())
	}

	if err := scanner.Err(); err != nil {
		return nil, errors.New("Error scanning commands: " + err.Error())
	}

	return commands, nil
}

func parseArgs() *string {
	runs = flag.Int("runs", 2, "Number of times the applications are executed")
	commandFile := flag.String("cmd", "cmd.txt", "Text file containing the commands to execute")

	cpus0 := flag.String("cpus0", "0-4", "List of CPUs to be used for the 1st command")
	cpus1 := flag.String("cpus1", "5-9", "List of CPUs to be used for the 2nd command")
	threads = flag.String("threads", "5", "Number of threads to be used")

	cat = flag.Bool("cat", false, "Measure with all CAT settings")
	catBitChunk = flag.Int("catChunk", 2, "Bits changed from one run to the next")
	resctrlPath = flag.String("resctrl", "/sys/fs/resctrl/", "Root path of the resctrl file system")

	hermitcore = flag.Bool("hermitcore", false, "Use if you are executing hermitcore binaries")

	flag.Parse()

	if *runs < 1 {
		fmt.Println("runs must be > 0")
		os.Exit(0)
	}
	if *catBitChunk < 1 {
		fmt.Println("catChunk must be > 0")
		os.Exit(0)
	}

	cpus[0] = *cpus0
	cpus[1] = *cpus1

	return commandFile
}
<commit_msg>Output normalized performance.<commit_after>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strconv"
	"sync"
	"time"

	"github.com/jbreitbart/coBench/bit"
	"github.com/montanaflynn/stats"
)

// global command line parameters
var runs *int
var cpus [2]string
var threads *string
var hermitcore *bool
var resctrlPath *string

var cat *bool
var catBitChunk *int

func main() {
	commandFile := parseArgs()

	commands, err := readCommands(*commandFile)
	if err != nil {
		log.Fatalf("Error reading command file %v: %v", *commandFile, err)
	}
	if len(commands) < 2 {
		log.Fatal("You must provide at least 2 commands")
	}

	var minBits int
	var numBits int

	if *cat {
		var err error
		minBits, numBits, err = setupCAT()
		if err != nil {
			log.Fatalf("%v\n", err)
		}
	} else {
		minBits = 0
		numBits = 0
	}

	commandPairs := generateCommandPairs(commands)

	fmt.Println("Executing the following command pairs:")
	for _, c := range commandPairs {
		fmt.Println(c)
	}

	for i, c := range commandPairs {
		fmt.Printf("Running pair %v\n", i)
		fmt.Println(c)
		for bits := minBits; bits <= numBits-minBits; bits += *catBitChunk {
			bitsets := []int64{0, 0}
			bitsets[0] = bit.SetFirstN(bitsets[0], bits)
			bitsets[1] = bit.SetLastN(bitsets[1], bits, numBits-1)

			err := runPair(c, i, bitsets)
			if err != nil {
				log.Fatalf("Error while running pair %v (%v): %v", i, c, err)
			}
		}
	}
}

func runCmdMinTimes(cmd *exec.Cmd, min int, catMask int64, wg *sync.WaitGroup, measurement *string, done chan int, errs chan error) {
	var runtime []float64

	defer wg.Done()

	defer func() {
		mean, _ := stats.Mean(runtime)
		stddev, _ := stats.StandardDeviation(runtime)
		vari, _ := stats.Variance(runtime)
		runtimeSum, _ := stats.Sum(runtime)

		s := fmt.Sprintf("%v \t %9.2fs avg. runtime \t %1.6f std. dev. \t %1.6f variance \t %v runs", cmd.Args, mean, stddev, vari, len(runtime))
		if *cat {
			s += fmt.Sprintf("\t %x", (uint)(catMask))
		}
		s += fmt.Sprintf("\t %1.6f", (float64)(len(runtime))/runtimeSum)
		s += "\n"
		fmt.Println(s)
	}()

	for i := 1; ; i++ {
		// create a copy of the command
		cmd := *cmd

		start := time.Now()
		err := cmd.Run()
		elapsed := time.Since(start)

		if err != nil {
			errs <- fmt.Errorf("Error running %v: %v", cmd.Args, err)
			return
		}

		// did the other cmd result in an error?
		if len(errs) != 0 {
			return
		}

		d := <-done

		// check if the other application was running the whole time
		if d == len(cpus) {
			// no
			*measurement += "# "
		} else {
			runtime = append(runtime, elapsed.Seconds())
		}
		*measurement += strconv.FormatInt(elapsed.Nanoseconds(), 10)
		*measurement += "\n"

		// did we run min times?
		if i == min {
			d++
		}
		done <- d

		// both applications are done
		if d == len(cpus) {
			return
		}
	}
}

func runPair(cPair [2]string, id int, catConfig []int64) error {
	env := os.Environ()

	if *cat {
		if err := writeCATConfig(catConfig); err != nil {
			return fmt.Errorf("Error while writting CAT config: %v", err)
		}
	}

	var cmds [len(cpus)]*exec.Cmd
	// setup commands
	for i, _ := range cmds {
		if *hermitcore {
			cmds[i] = exec.Command("numactl", "--physcpubind", cpus[i], "/bin/sh", "-c", cPair[i])
			cmds[i].Env = append(env, "HERMIT_CPUS="+*threads, "HERMIT_MEM=4G", "HERMIT_ISLE=uhyve")
		} else {
			cmds[i] = exec.Command("/bin/sh", "-c", cPair[i])
			cmds[i].Env = append(env, "GOMP_CPU_AFFINITY="+cpus[i], "OMP_NUM_THREADS="+*threads)
		}

		outfile, err := os.Create(fmt.Sprintf("%v-%v-%x.log", id, i, catConfig[i]))
		if err != nil {
			return fmt.Errorf("Error while creating file: %v", err)
		}
		defer outfile.Close()
		cmds[i].Stdout = outfile
	}

	var measurements [len(cmds)]string
	// used to count how many apps have reached their min limit
	done := make(chan int, 1)
	done <- 0

	// used to return an error from the go-routines
	errs := make(chan error, len(cmds))

	// used to wait for the following 2 goroutines
	var wg sync.WaitGroup
	wg.Add(len(cmds))

	for i, c := range cmds {
		go runCmdMinTimes(c, *runs, catConfig[i], &wg, &measurements[i], done, errs)
	}

	wg.Wait()

	for i, s := range measurements {
		measurementsFile, err := os.Create(fmt.Sprintf("%v-%v.time", id, i))
		if err != nil {
			return fmt.Errorf("Error while creating file: %v", err)
		}
		defer measurementsFile.Close()

		out := "# runtime in nanoseconds of \"" + cPair[i] + "\" on CPUs " + cpus[i] + "while \"" + cPair[(i+1)%2] + "\" was running on cores " + cpus[(i+1)%len(cpus)]
		if *cat {
			out += fmt.Sprintf(" with CAT %x ", catConfig[i])
		}
		out += "\n"
		out += s

		_, err = measurementsFile.WriteString(out)
		if err != nil {
			return fmt.Errorf("Error while writing measurements file: %v", err)
		}
	}

	if len(errs) != 0 {
		return <-errs
	}

	return nil
}

func generateCommandPairs(commands []string) [][2]string {
	var pairs [][2]string
	for i, c0 := range commands {
		for j, c1 := range commands {
			if i >= j {
				continue
			}
			pairs = append(pairs, [2]string{c0, c1})
		}
	}
	return pairs
}

func readCommands(filename string) ([]string, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, errors.New("Error opening file " + filename + ": " + err.Error())
	}
	defer file.Close()

	var commands []string

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		commands = append(commands, scanner.Text())
	}

	if err := scanner.Err(); err != nil {
		return nil, errors.New("Error scanning commands: " + err.Error())
	}

	return commands, nil
}

func parseArgs() *string {
	runs = flag.Int("runs", 2, "Number of times the applications are executed")
	commandFile := flag.String("cmd", "cmd.txt", "Text file containing the commands to execute")

	cpus0 := flag.String("cpus0", "0-4", "List of CPUs to be used for the 1st command")
	cpus1 := flag.String("cpus1", "5-9", "List of CPUs to be used for the 2nd command")
	threads = flag.String("threads", "5", "Number of threads to be used")

	cat = flag.Bool("cat", false, "Measure with all CAT settings")
	catBitChunk = flag.Int("catChunk", 2, "Bits changed from one run to the next")
	resctrlPath = flag.String("resctrl", "/sys/fs/resctrl/", "Root path of the resctrl file system")

	hermitcore = flag.Bool("hermitcore", false, "Use if you are executing hermitcore binaries")

	flag.Parse()

	if *runs < 1 {
		fmt.Println("runs must be > 0")
		os.Exit(0)
	}
	if *catBitChunk < 1 {
		fmt.Println("catChunk must be > 0")
		os.Exit(0)
	}

	cpus[0] = *cpus0
	cpus[1] = *cpus1

	return commandFile
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Opsidian Ltd.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package text_test

import (
	"io/ioutil"
	"os"
	"path/filepath"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/opsidian/parsley/parsley"
	"github.com/opsidian/parsley/text"
)

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

	var (
		fs   *parsley.FileSet
		f    *text.File
		data []byte
	)

	JustBeforeEach(func() {
		f = text.NewFile("testfile", data)
		fs = parsley.NewFileSet(f)
	})

	BeforeEach(func() {
		data = []byte("ab\nc")
	})

	It("should implement the parsley.File interface", func() {
		var _ parsley.File = &text.File{}
	})

	Describe("Position()", func() {
		It("should return with the position for an offset", func() {
			Expect(f.Position(0)).To(Equal(text.NewPosition("testfile", 1, 1)))
			Expect(f.Position(1)).To(Equal(text.NewPosition("testfile", 1, 2)))
			Expect(f.Position(2)).To(Equal(text.NewPosition("testfile", 1, 3)))
			Expect(f.Position(3)).To(Equal(text.NewPosition("testfile", 2, 1)))
			Expect(f.Position(4)).To(Equal(text.NewPosition("testfile", 2, 2)))
		})

		It("should return nil position for an invalid offset", func() {
			Expect(f.Position(5)).To(Equal(parsley.NilPosition))
		})
	})

	Describe("Len()", func() {
		Context("when data is empty", func() {
			BeforeEach(func() {
				data = []byte{}
			})

			It("should return zero", func() {
				Expect(f.Len()).To(Equal(0))
			})
		})

		It("should return the data length", func() {
			Expect(f.Len()).To(Equal(4))
		})
	})

	Context("when data contains Windows-style line endings", func() {
		BeforeEach(func() {
			data = []byte("a\r\nb\r\nc\n")
		})

		It("should remove the \r characters", func() {
			Expect(f.Len()).To(Equal(6))
			Expect(f.Position(1).(*text.Position).Line).To(Equal(1))
			Expect(f.Position(2).(*text.Position).Line).To(Equal(2))
			Expect(f.Position(3).(*text.Position).Line).To(Equal(2))
			Expect(f.Position(4).(*text.Position).Line).To(Equal(3))
		})
	})

	Context("ReadFile", func() {
		var (
			f           *text.File
			readFileErr error
			filename    string
		)

		JustBeforeEach(func() {
			f, readFileErr = text.ReadFile(filename)
		})

		Context("reading an existing file", func() {
			var tmpDir string

			BeforeEach(func() {
				var err error
				tmpDir, err = ioutil.TempDir("", "parsley-test-")
				Expect(err).ToNot(HaveOccurred())
				filename = filepath.Join(tmpDir, "testfile")
				ioutil.WriteFile(filename, []byte("ab\nc"), 0600)
			})

			AfterEach(func() {
				if tmpDir != "" {
					os.RemoveAll(tmpDir)
				}
			})

			It("should succeed", func() {
				Expect(readFileErr).ToNot(HaveOccurred())
			})

			Describe("Position", func() {
				It("should return the position for an offset", func() {
					Expect(f.Position(0)).To(Equal(text.NewPosition(filename, 1, 1)))
					Expect(f.Position(1)).To(Equal(text.NewPosition(filename, 1, 2)))
					Expect(f.Position(2)).To(Equal(text.NewPosition(filename, 1, 3)))
					Expect(f.Position(3)).To(Equal(text.NewPosition(filename, 2, 1)))
					Expect(f.Position(4)).To(Equal(text.NewPosition(filename, 2, 2)))
				})

				It("should return nil position for an invalid offset", func() {
					Expect(f.Position(5)).To(Equal(parsley.NilPosition))
				})

				It("should contain a readable position string", func() {
					Expect(f.Position(2).String()).To(Equal(filename + ":1:3"))
				})
			})

			Describe("Len", func() {
				It("should return with the filesize", func() {
					Expect(f.Len()).To(Equal(4))
				})
			})
		})

		Context("reading a non-existing file", func() {
			BeforeEach(func() {
				filename = "/tmp/non-existing-filename"
			})

			It("should throw an error", func() {
				Expect(readFileErr).To(MatchError("can not read /tmp/non-existing-filename"))
			})
		})
	})

	Describe("Pos()", func() {
		It("should return with a global position (pos 0)", func() {
			pos := f.Pos(0)
			Expect(pos).To(Equal(parsley.Pos(1)))
			Expect(fs.Position(pos)).To(Equal(text.NewPosition("testfile", 1, 1)))
		})

		It("should return with a global position (at the end)", func() {
			pos := f.Pos(3)
			Expect(pos).To(Equal(parsley.Pos(4)))
			Expect(fs.Position(pos)).To(Equal(text.NewPosition("testfile", 2, 1)))
		})
	})
})
<commit_msg>Remove unnecessary variable from text.File test<commit_after>// Copyright (c) 2017 Opsidian Ltd.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package text_test

import (
	"io/ioutil"
	"os"
	"path/filepath"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/opsidian/parsley/parsley"
	"github.com/opsidian/parsley/text"
)

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

	var (
		fs   *parsley.FileSet
		f    *text.File
		data []byte
	)

	JustBeforeEach(func() {
		f = text.NewFile("testfile", data)
		fs = parsley.NewFileSet(f)
	})

	BeforeEach(func() {
		data = []byte("ab\nc")
	})

	It("should implement the parsley.File interface", func() {
		var _ parsley.File = &text.File{}
	})

	Describe("Position()", func() {
		It("should return with the position for an offset", func() {
			Expect(f.Position(0)).To(Equal(text.NewPosition("testfile", 1, 1)))
			Expect(f.Position(1)).To(Equal(text.NewPosition("testfile", 1, 2)))
			Expect(f.Position(2)).To(Equal(text.NewPosition("testfile", 1, 3)))
			Expect(f.Position(3)).To(Equal(text.NewPosition("testfile", 2, 1)))
			Expect(f.Position(4)).To(Equal(text.NewPosition("testfile", 2, 2)))
		})

		It("should return nil position for an invalid offset", func() {
			Expect(f.Position(5)).To(Equal(parsley.NilPosition))
		})
	})

	Describe("Len()", func() {
		Context("when data is empty", func() {
			BeforeEach(func() {
				data = []byte{}
			})

			It("should return zero", func() {
				Expect(f.Len()).To(Equal(0))
			})
		})

		It("should return the data length", func() {
			Expect(f.Len()).To(Equal(4))
		})
	})

	Context("when data contains Windows-style line endings", func() {
		BeforeEach(func() {
			data = []byte("a\r\nb\r\nc\n")
		})

		It("should remove the \r characters", func() {
			Expect(f.Len()).To(Equal(6))
			Expect(f.Position(1).(*text.Position).Line).To(Equal(1))
			Expect(f.Position(2).(*text.Position).Line).To(Equal(2))
			Expect(f.Position(3).(*text.Position).Line).To(Equal(2))
			Expect(f.Position(4).(*text.Position).Line).To(Equal(3))
		})
	})

	Context("ReadFile", func() {
		var (
			readFileErr error
			filename    string
		)

		JustBeforeEach(func() {
			f, readFileErr = text.ReadFile(filename)
		})

		Context("reading an existing file", func() {
			var tmpDir string

			BeforeEach(func() {
				var err error
				tmpDir, err = ioutil.TempDir("", "parsley-test-")
				Expect(err).ToNot(HaveOccurred())
				filename = filepath.Join(tmpDir, "testfile")
				ioutil.WriteFile(filename, []byte("ab\nc"), 0600)
			})

			AfterEach(func() {
				if tmpDir != "" {
					os.RemoveAll(tmpDir)
				}
			})

			It("should succeed", func() {
				Expect(readFileErr).ToNot(HaveOccurred())
			})

			Describe("Position", func() {
				It("should return the position for an offset", func() {
					Expect(f.Position(0)).To(Equal(text.NewPosition(filename, 1, 1)))
					Expect(f.Position(1)).To(Equal(text.NewPosition(filename, 1, 2)))
					Expect(f.Position(2)).To(Equal(text.NewPosition(filename, 1, 3)))
					Expect(f.Position(3)).To(Equal(text.NewPosition(filename, 2, 1)))
					Expect(f.Position(4)).To(Equal(text.NewPosition(filename, 2, 2)))
				})

				It("should return nil position for an invalid offset", func() {
					Expect(f.Position(5)).To(Equal(parsley.NilPosition))
				})

				It("should contain a readable position string", func() {
					Expect(f.Position(2).String()).To(Equal(filename + ":1:3"))
				})
			})

			Describe("Len", func() {
				It("should return with the filesize", func() {
					Expect(f.Len()).To(Equal(4))
				})
			})
		})

		Context("reading a non-existing file", func() {
			BeforeEach(func() {
				filename = "/tmp/non-existing-filename"
			})

			It("should throw an error", func() {
				Expect(readFileErr).To(MatchError("can not read /tmp/non-existing-filename"))
			})
		})
	})

	Describe("Pos()", func() {
		It("should return with a global position (pos 0)", func() {
			pos := f.Pos(0)
			Expect(pos).To(Equal(parsley.Pos(1)))
			Expect(fs.Position(pos)).To(Equal(text.NewPosition("testfile", 1, 1)))
		})

		It("should return with a global position (at the end)", func() {
			pos := f.Pos(3)
			Expect(pos).To(Equal(parsley.Pos(4)))
			Expect(fs.Position(pos)).To(Equal(text.NewPosition("testfile", 2, 1)))
		})
	})
})
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Julien Vehent jvehent@mozilla.com [:ulfr]
package main

import (
	"encoding/json"
	"fmt"
	"github.com/bobappleyard/readline"
	"github.com/jvehent/cljs"
	"io"
	"mig"
	"mig/client"
	"strconv"
	"strings"
)

// actionReader retrieves an action from the API using its numerical ID
// and enters prompt mode to analyze it
func actionReader(input string, cli client.Client) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("actionReader() -> %v", e)
		}
	}()
	inputArr := strings.Split(input, " ")
	if len(inputArr) < 2 {
		panic("wrong order format. must be 'action <actionid>'")
	}
	aid, err := strconv.ParseFloat(inputArr[1], 64)
	if err != nil {
		panic(err)
	}
	a, links, err := cli.GetAction(aid)
	if err != nil {
		panic(err)
	}
	investigators := investigatorsStringFromAction(a.Investigators, 80)

	fmt.Println("Entering action reader mode. Type \x1b[32;1mexit\x1b[0m or press \x1b[32;1mctrl+d\x1b[0m to leave. \x1b[32;1mhelp\x1b[0m may help.")
	fmt.Printf("Action: '%s'.\nLaunched by '%s' on '%s'.\nStatus '%s'.\n",
		a.Name, investigators, a.StartTime, a.Status)
	prompt := fmt.Sprintf("\x1b[31;1maction %d>\x1b[0m ", uint64(aid)%1000)
	for {
		// completion
		var symbols = []string{"command", "copy", "counters", "details", "exit", "foundanything", "foundnothing",
			"grep", "help", "investigators", "json", "ls", "found", "pretty", "r", "results", "times"}
		readline.Completer = func(query, ctx string) []string {
			var res []string
			for _, sym := range symbols {
				if strings.HasPrefix(sym, query) {
					res = append(res, sym)
				}
			}
			return res
		}

		input, err := readline.String(prompt)
		if err == io.EOF {
			break
		}
		if err != nil {
			fmt.Println("error: ", err)
			break
		}
		orders := strings.Split(strings.TrimSpace(input), " ")
		switch orders[0] {
		case "command":
			err = commandReader(input, cli)
			if err != nil {
				panic(err)
			}
		case "copy":
			err = actionLauncher(a, cli)
			if err != nil {
				panic(err)
			}
			goto exit
		case "counters":
			fmt.Printf("Sent:\t\t%d\nDone:\t\t%d\nIn Flight:\t%d\n"+
				"Cancelled:\t%d\nFailed:\t\t%d\nTimeout:\t%d\n",
				a.Counters.Sent, a.Counters.Done, a.Counters.InFlight,
				a.Counters.Cancelled, a.Counters.Failed, a.Counters.TimeOut)
		case "details":
			actionPrintDetails(a)
		case "exit":
			fmt.Printf("exit\n")
			goto exit
		case "foundanything":
			err = searchFoundAnything(a, true, cli)
			if err != nil {
				panic(err)
			}
		case "foundnothing":
			err = searchFoundAnything(a, false, cli)
			if err != nil {
				panic(err)
			}
		case "help":
			fmt.Printf(`The following orders are available:
command <id>	jump to command reader mode for command <id>
copy		enter action launcher mode using current action as template
counters	display the counters of the action
details		display the details of the action, including status & times
exit		exit this mode (also works with ctrl+d)
foundanything	list commands and agents that have found something
foundnothing	list commands and agents that have found nothing
help		show this help
investigators   print the list of investigators that signed the action
json         	show the json of the action
ls <filter>	returns the list of commands with their status
		'filter' is a pipe separated string of filter:
		ex: ls | grep server1.(dom1|dom2) | grep -v example.net
r		refresh the action (get latest version from upstream)
results <found> display results of all commands, limit to results that have
                found something is <found> is declared
times		show the various timestamps of the action
`)
		case "investigators":
			for _, i := range a.Investigators {
				fmt.Println(i.Name, "- Key ID:", i.PGPFingerprint)
			}
		case "json":
			var ajson []byte
			ajson, err = json.MarshalIndent(a, "", "  ")
			if err != nil {
				panic(err)
			}
			fmt.Printf("%s\n", ajson)
		case "ls":
			err = actionPrintLinks(links, orders)
			if err != nil {
				panic(err)
			}
		case "r":
			a, links, err = cli.GetAction(aid)
			if err != nil {
				panic(err)
			}
			fmt.Println("Reload succeeded")
		case "results":
			err = actionPrintResults(a, links, orders, cli)
			if err != nil {
				panic(err)
			}
		case "times":
			fmt.Printf("Valid from   '%s' until '%s'\nStarted on   '%s'\n"+
				"Last updated '%s'\nFinished on  '%s'\n",
				a.ValidFrom, a.ExpireAfter, a.StartTime, a.LastUpdateTime, a.FinishTime)
		case "":
			break
		default:
			fmt.Printf("Unknown order '%s'. You are in action reader mode. Try `help`.\n", orders[0])
		}
		readline.AddHistory(input)
	}
exit:
	fmt.Printf("\n")
	return
}

func searchFoundAnything(a mig.Action, wantFound bool, cli client.Client) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("searchFoundAnything() -> %v", e)
		}
	}()
	target := "search?type=command&limit=1000000&actionid=" + fmt.Sprintf("%.0f", a.ID)
	if wantFound {
		target += "&foundanything=true"
	} else {
		target += "&foundanything=false"
	}
	resource, err := cli.GetAPIResource(target)
	if err != nil {
		panic(err)
	}
	agents := make(map[float64]mig.Command)
	for _, item := range resource.Collection.Items {
		for _, data := range item.Data {
			if data.Name != "command" {
				continue
			}
			cmd, err := client.ValueToCommand(data.Value)
			if err != nil {
				panic(err)
			}
			agents[cmd.Agent.ID] = cmd
		}
	}
	if wantFound {
		fmt.Printf("%d agents have found things\n", len(agents))
	} else {
		fmt.Printf("%d agents have not found anything\n", len(agents))
	}
	if len(agents) > 0 {
		fmt.Println("---- Command ID ----    ---- Agent Name & ID----")
		for agtid, cmd := range agents {
			fmt.Printf("%20.0f    %s [%.0f]\n", cmd.ID, cmd.Agent.Name, agtid)
		}
	}
	return
}

func actionPrintShort(data interface{}) (idstr, name, datestr, invs string, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("actionPrintShort() -> %v", e)
		}
	}()
	a, err := client.ValueToAction(data)
	if err != nil {
		panic(err)
	}
	invs = investigatorsStringFromAction(a.Investigators, 23)

	idstr = fmt.Sprintf("%.0f", a.ID)
	if len(idstr) < 20 {
		for i := len(idstr); i < 20; i++ {
			idstr += " "
		}
	}

	name = a.Name
	if len(name) < 30 {
		for i := len(name); i < 30; i++ {
			name += " "
		}
	}
	if len(name) > 30 {
		name = name[0:27] + "..."
	}

	datestr = a.LastUpdateTime.Format("Mon Jan 2 3:04pm MST")
	if len(datestr) > 20 {
		datestr = datestr[0:20]
	}
	if len(datestr) < 20 {
		for i := len(datestr); i < 20; i++ {
			datestr += " "
		}
	}
	return
}

func investigatorsStringFromAction(invlist []mig.Investigator, strlen int) (investigators string) {
	for ctr, i := range invlist {
		if ctr > 0 {
			investigators += "; "
		}
		investigators += i.Name
	}
	if len(investigators) > strlen {
		investigators = investigators[0:(strlen-3)] + "..."
	}
	return
}

func actionPrintDetails(a mig.Action) {
	fmt.Printf(`
ID             %.0f
Name           %s
Target         %s
Desc           author '%s <%s>'; revision '%.0f';
               url '%s'
Threat         type '%s'; level '%s'; family '%s'; reference '%s'
Status         %s
Times          valid from %s until %s
               started %s; last updated %s; finished %s
               duration: %s
`, a.ID, a.Name, a.Target, a.Description.Author, a.Description.Email, a.Description.Revision,
		a.Description.URL, a.Threat.Type, a.Threat.Level, a.Threat.Family, a.Threat.Ref, a.Status,
		a.ValidFrom, a.ExpireAfter, a.StartTime, a.LastUpdateTime, a.FinishTime, a.LastUpdateTime.Sub(a.StartTime).String())
	fmt.Printf("Investigators  ")
	for _, i := range a.Investigators {
		fmt.Println(i.Name, "- keyid:", i.PGPFingerprint)
	}
	fmt.Printf("Operations     count=%d => ", len(a.Operations))
	for _, op := range a.Operations {
		fmt.Printf("%s; ", op.Module)
	}
	fmt.Printf("\n")
	fmt.Printf("Counters       sent=%d; done=%d; in flight=%d\n"+
		"               cancelled=%d; expired=%d; failed=%d; timeout=%d\n",
		a.Counters.Sent, a.Counters.Done, a.Counters.InFlight,
		a.Counters.Cancelled, a.Counters.Expired, a.Counters.Failed, a.Counters.TimeOut)
	return
}

func actionPrintLinks(links []cljs.Link, orders []string) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("actionPrintLinks() -> %v", e)
		}
	}()
	has_filter := false
	var filter []string
	if len(orders) > 1 {
		has_filter = true
		filter = orders[1:len(orders)]
	}
	ctr := 0
	for _, link := range links {
		if useShortNames {
			link.Rel = link.Rel[0:40] + shorten(link.Rel[40:len(link.Rel)])
		}
		if has_filter {
			str, err := filterString(link.Rel, filter)
			if err != nil {
				fmt.Printf("Invalid filter '%s': '%v'\n", filter, err)
				break
			}
			if str != "" {
				fmt.Println(str)
				ctr++
			}
		} else {
			fmt.Println(link.Rel)
			ctr++
		}
	}
	fmt.Printf("%d command", ctr)
	if ctr > 1 {
		fmt.Printf("s")
	}
	fmt.Printf(" found\n")
	return
}

func actionPrintResults(a mig.Action, links []cljs.Link, orders []string, cli client.Client) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("actionPrintResuls() -> %v", e)
		}
	}()
	found := false
	if len(orders) > 1 {
		if orders[1] == "found" {
			found = true
		} else {
			fmt.Printf("Unknown option '%s'\n", orders[1])
		}
	}
	if found {
		// if we want foundes, use the search api, it's faster than
		// iterating through each link when we have thousands of them
		target := "search?type=command&limit=1000000&foundanything=true"
		target += "&actionid=" + fmt.Sprintf("%.0f", a.ID)
		resource, err := cli.GetAPIResource(target)
		if err != nil {
			panic(err)
		}
		for _, item := range resource.Collection.Items {
			for _, data := range item.Data {
				if data.Name != "command" {
					continue
				}
				cmd, err := client.ValueToCommand(data.Value)
				if err != nil {
					panic(err)
				}
				err = commandPrintResults(cmd, true, true)
				if err != nil {
					panic(err)
				}
			}
		}
	} else {
		for _, link := range links {
			// TODO: replace the url parsing hack with proper link creation in API response
			id := strings.Split(link.Href, "=")[1]
			cmdid, err := strconv.ParseFloat(id, 64)
			if err != nil {
				panic(err)
			}
			cmd, err := cli.GetCommand(cmdid)
			if err != nil {
				panic(err)
			}
			err = commandPrintResults(cmd, found, true)
			if err != nil {
				panic(err)
			}
		}
	}
	return
}
<commit_msg>[minor] fix counters display in console<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Julien Vehent jvehent@mozilla.com [:ulfr]
package main

import (
	"encoding/json"
	"fmt"
	"github.com/bobappleyard/readline"
	"github.com/jvehent/cljs"
	"io"
	"mig"
	"mig/client"
	"strconv"
	"strings"
)

// actionReader retrieves an action from the API using its numerical ID
// and enters prompt mode to analyze it
func actionReader(input string, cli client.Client) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("actionReader() -> %v", e)
		}
	}()
	inputArr := strings.Split(input, " ")
	if len(inputArr) < 2 {
		panic("wrong order format. must be 'action <actionid>'")
	}
	aid, err := strconv.ParseFloat(inputArr[1], 64)
	if err != nil {
		panic(err)
	}
	a, links, err := cli.GetAction(aid)
	if err != nil {
		panic(err)
	}
	investigators := investigatorsStringFromAction(a.Investigators, 80)

	fmt.Println("Entering action reader mode. Type \x1b[32;1mexit\x1b[0m or press \x1b[32;1mctrl+d\x1b[0m to leave. \x1b[32;1mhelp\x1b[0m may help.")
	fmt.Printf("Action: '%s'.\nLaunched by '%s' on '%s'.\nStatus '%s'.\n",
		a.Name, investigators, a.StartTime, a.Status)
	prompt := fmt.Sprintf("\x1b[31;1maction %d>\x1b[0m ", uint64(aid)%1000)
	for {
		// completion
		var symbols = []string{"command", "copy", "counters", "details", "exit", "foundanything", "foundnothing",
			"grep", "help", "investigators", "json", "ls", "found", "pretty", "r", "results", "times"}
		readline.Completer = func(query, ctx string) []string {
			var res []string
			for _, sym := range symbols {
				if strings.HasPrefix(sym, query) {
					res = append(res, sym)
				}
			}
			return res
		}

		input, err := readline.String(prompt)
		if err == io.EOF {
			break
		}
		if err != nil {
			fmt.Println("error: ", err)
			break
		}
		orders := strings.Split(strings.TrimSpace(input), " ")
		switch orders[0] {
		case "command":
			err = commandReader(input, cli)
			if err != nil {
				panic(err)
			}
		case "copy":
			err = actionLauncher(a, cli)
			if err != nil {
				panic(err)
			}
			goto exit
		case "counters":
			fmt.Printf("Sent:\t\t%d\nDone:\t\t%d\nIn Flight:\t%d\n"+
				"Success:\t%d\nCancelled:\t%d\nExpired:\t%d\nFailed:\t\t%d\nTimeout:\t%d\n",
				a.Counters.Sent, a.Counters.Done, a.Counters.InFlight,
				a.Counters.Success, a.Counters.Cancelled, a.Counters.Expired, a.Counters.Failed, a.Counters.TimeOut)
		case "details":
			actionPrintDetails(a)
		case "exit":
			fmt.Printf("exit\n")
			goto exit
		case "foundanything":
			err = searchFoundAnything(a, true, cli)
			if err != nil {
				panic(err)
			}
		case "foundnothing":
			err = searchFoundAnything(a, false, cli)
			if err != nil {
				panic(err)
			}
		case "help":
			fmt.Printf(`The following orders are available:
command <id>	jump to command reader mode for command <id>
copy		enter action launcher mode using current action as template
counters	display the counters of the action
details		display the details of the action, including status & times
exit		exit this mode (also works with ctrl+d)
foundanything	list commands and agents that have found something
foundnothing	list commands and agents that have found nothing
help		show this help
investigators   print the list of investigators that signed the action
json         	show the json of the action
ls <filter>	returns the list of commands with their status
		'filter' is a pipe separated string of filter:
		ex: ls | grep server1.(dom1|dom2) | grep -v example.net
r		refresh the action (get latest version from upstream)
results <found> display results of all commands, limit to results that have
                found something is <found> is declared
times		show the various timestamps of the action
`)
		case "investigators":
			for _, i := range a.Investigators {
				fmt.Println(i.Name, "- Key ID:", i.PGPFingerprint)
			}
		case "json":
			var ajson []byte
			ajson, err = json.MarshalIndent(a, "", "  ")
			if err != nil {
				panic(err)
			}
			fmt.Printf("%s\n", ajson)
		case "ls":
			err = actionPrintLinks(links, orders)
			if err != nil {
				panic(err)
			}
		case "r":
			a, links, err = cli.GetAction(aid)
			if err != nil {
				panic(err)
			}
			fmt.Println("Reload succeeded")
		case "results":
			err = actionPrintResults(a, links, orders, cli)
			if err != nil {
				panic(err)
			}
		case "times":
			fmt.Printf("Valid from   '%s' until '%s'\nStarted on   '%s'\n"+
				"Last updated '%s'\nFinished on  '%s'\n",
				a.ValidFrom, a.ExpireAfter, a.StartTime, a.LastUpdateTime, a.FinishTime)
		case "":
			break
		default:
			fmt.Printf("Unknown order '%s'. You are in action reader mode. Try `help`.\n", orders[0])
		}
		readline.AddHistory(input)
	}
exit:
	fmt.Printf("\n")
	return
}

func searchFoundAnything(a mig.Action, wantFound bool, cli client.Client) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("searchFoundAnything() -> %v", e)
		}
	}()
	target := "search?type=command&limit=1000000&actionid=" + fmt.Sprintf("%.0f", a.ID)
	if wantFound {
		target += "&foundanything=true"
	} else {
		target += "&foundanything=false"
	}
	resource, err := cli.GetAPIResource(target)
	if err != nil {
		panic(err)
	}
	agents := make(map[float64]mig.Command)
	for _, item := range resource.Collection.Items {
		for _, data := range item.Data {
			if data.Name != "command" {
				continue
			}
			cmd, err := client.ValueToCommand(data.Value)
			if err != nil {
				panic(err)
			}
			agents[cmd.Agent.ID] = cmd
		}
	}
	if wantFound {
		fmt.Printf("%d agents have found things\n", len(agents))
	} else {
		fmt.Printf("%d agents have not found anything\n", len(agents))
	}
	if len(agents) > 0 {
		fmt.Println("---- Command ID ----    ---- Agent Name & ID----")
		for agtid, cmd := range agents {
			fmt.Printf("%20.0f    %s [%.0f]\n", cmd.ID, cmd.Agent.Name, agtid)
		}
	}
	return
}

func actionPrintShort(data interface{}) (idstr, name, datestr, invs string, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("actionPrintShort() -> %v", e)
		}
	}()
	a, err := client.ValueToAction(data)
	if err != nil {
		panic(err)
	}
	invs = investigatorsStringFromAction(a.Investigators, 23)

	idstr = fmt.Sprintf("%.0f", a.ID)
	if len(idstr) < 20 {
		for i := len(idstr); i < 20; i++ {
			idstr += " "
		}
	}

	name = a.Name
	if len(name) < 30 {
		for i := len(name); i < 30; i++ {
			name += " "
		}
	}
	if len(name) > 30 {
		name = name[0:27] + "..."
	}

	datestr = a.LastUpdateTime.Format("Mon Jan 2 3:04pm MST")
	if len(datestr) > 20 {
		datestr = datestr[0:20]
	}
	if len(datestr) < 20 {
		for i := len(datestr); i < 20; i++ {
			datestr += " "
		}
	}
	return
}

func investigatorsStringFromAction(invlist []mig.Investigator, strlen int) (investigators string) {
	for ctr, i := range invlist {
		if ctr > 0 {
			investigators += "; "
		}
		investigators += i.Name
	}
	if len(investigators) > strlen {
		investigators = investigators[0:(strlen-3)] + "..."
	}
	return
}

func actionPrintDetails(a mig.Action) {
	fmt.Printf(`
ID             %.0f
Name           %s
Target         %s
Desc           author '%s <%s>'; revision '%.0f';
               url '%s'
Threat         type '%s'; level '%s'; family '%s'; reference '%s'
Status         %s
Times          valid from %s until %s
               started %s; last updated %s; finished %s
               duration: %s
`, a.ID, a.Name, a.Target, a.Description.Author, a.Description.Email, a.Description.Revision,
		a.Description.URL, a.Threat.Type, a.Threat.Level, a.Threat.Family, a.Threat.Ref, a.Status,
		a.ValidFrom, a.ExpireAfter, a.StartTime, a.LastUpdateTime, a.FinishTime, a.LastUpdateTime.Sub(a.StartTime).String())
	fmt.Printf("Investigators  ")
	for _, i := range a.Investigators {
		fmt.Println(i.Name, "- keyid:", i.PGPFingerprint)
	}
	fmt.Printf("Operations     count=%d => ", len(a.Operations))
	for _, op := range a.Operations {
		fmt.Printf("%s; ", op.Module)
	}
	fmt.Printf("\n")
	fmt.Printf("Counters       sent=%d; done=%d; in flight=%d\n"+
		"               sucess=%d; cancelled=%d; expired=%d; failed=%d; timeout=%d\n",
		a.Counters.Sent, a.Counters.Done, a.Counters.InFlight, a.Counters.Success,
		a.Counters.Cancelled, a.Counters.Expired, a.Counters.Failed, a.Counters.TimeOut)
	return
}

func actionPrintLinks(links []cljs.Link, orders []string) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("actionPrintLinks() -> %v", e)
		}
	}()
	has_filter := false
	var filter []string
	if len(orders) > 1 {
		has_filter = true
		filter = orders[1:len(orders)]
	}
	ctr := 0
	for _, link := range links {
		if useShortNames {
			link.Rel = link.Rel[0:40] + shorten(link.Rel[40:len(link.Rel)])
		}
		if has_filter {
			str, err := filterString(link.Rel, filter)
			if err != nil {
				fmt.Printf("Invalid filter '%s': '%v'\n", filter, err)
				break
			}
			if str != "" {
				fmt.Println(str)
				ctr++
			}
		} else {
			fmt.Println(link.Rel)
			ctr++
		}
	}
	fmt.Printf("%d command", ctr)
	if ctr > 1 {
		fmt.Printf("s")
	}
	fmt.Printf(" found\n")
	return
}

func actionPrintResults(a mig.Action, links []cljs.Link, orders []string, cli client.Client) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("actionPrintResuls() -> %v", e)
		}
	}()
	found := false
	if len(orders) > 1 {
		if orders[1] == "found" {
			found = true
		} else {
			fmt.Printf("Unknown option '%s'\n", orders[1])
		}
	}
	if found {
		// if we want foundes, use the search api, it's faster than
		// iterating through each link when we have thousands of them
		target := "search?type=command&limit=1000000&foundanything=true"
		target += "&actionid=" + fmt.Sprintf("%.0f", a.ID)
		resource, err := cli.GetAPIResource(target)
		if err != nil {
			panic(err)
		}
		for _, item := range resource.Collection.Items {
			for _, data := range item.Data {
				if data.Name != "command" {
					continue
				}
				cmd, err := client.ValueToCommand(data.Value)
				if err != nil {
					panic(err)
				}
				err = commandPrintResults(cmd, true, true)
				if err != nil {
					panic(err)
				}
			}
		}
	} else {
		for _, link := range links {
			// TODO: replace the url parsing hack with proper link creation in API response
			id := strings.Split(link.Href, "=")[1]
			cmdid, err := strconv.ParseFloat(id, 64)
			if err != nil {
				panic(err)
			}
			cmd, err := cli.GetCommand(cmdid)
			if err != nil {
				panic(err)
			}
			err = commandPrintResults(cmd, found, true)
			if err != nil {
				panic(err)
			}
		}
	}
	return
}
<|endoftext|>


<commit_before>package GoSDK

import (
//"fmt"
)

const (
	_THROTTLERS_PREAMBLE = "/admin/throttlers"
)

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

func (d *DevClient) GetThrottlerDefinitions() (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(get(d, _THROTTLERS_PREAMBLE+"/definitions", nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

func (d *DevClient) GetAllThrottlers() (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(get(d, _THROTTLERS_PREAMBLE, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

func (d *DevClient) GetThrottler(throttlerName string) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(get(d, _THROTTLERS_PREAMBLE+"/"+throttlerName, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

func (d *DevClient) GetThrottlerCasesAndExceptions(throttlerName string) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(get(d, _THROTTLERS_PREAMBLE+"/"+throttlerName, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

func (d *DevClient) DeleteThrottlerCasesAndExceptions(throttlerName string) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(delete(d, _THROTTLERS_PREAMBLE+"/"+throttlerName, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

func (d *DevClient) GetThrottlerCases(throttlerName string) ([]interface{}, error) {
	return d.getThrottlerCases(throttlerName, "/cases")
}

func (d *DevClient) GetThrottlerExceptions(throttlerName string) ([]interface{}, error) {
	return d.getThrottlerCases(throttlerName, "/exceptions")
}

func (d *DevClient) getThrottlerCases(throttlerName, pathTail string) ([]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(get(d, _THROTTLERS_PREAMBLE+"/"+throttlerName+pathTail, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.([]interface{}), nil
}

func (d *DevClient) CreateThrottlerCase(throttlerName string, caseInfo map[string]interface{}) (map[string]interface{}, error) {
	return d.createThrottlerCase(throttlerName, "/cases", caseInfo)
}

func (d *DevClient) CreateThrottlerException(throttlerName string, caseInfo map[string]interface{}) (map[string]interface{}, error) {
	return d.createThrottlerCase(throttlerName, "/exceptions", caseInfo)
}

func (d *DevClient) createThrottlerCase(throttlerName, pathTail string, caseInfo map[string]interface{}) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(post(d, _THROTTLERS_PREAMBLE+"/"+throttlerName+pathTail, caseInfo, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

func (d *DevClient) DeleteAllThrottlerCases(throttlerName string) ([]interface{}, error) {
	return d.deleteAllThrottlerCases(throttlerName, "/cases")
}

func (d *DevClient) DeleteAllThrottlerExceptions(throttlerName string) ([]interface{}, error) {
	return d.deleteAllThrottlerCases(throttlerName, "/exceptions")
}

func (d *DevClient) deleteAllThrottlerCases(throttlerName, pathTail string) ([]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(delete(d, _THROTTLERS_PREAMBLE+"/"+throttlerName+pathTail, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.([]interface{}), nil
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

func (d *DevClient) GetThrottlerCase(throttlerName, caseName string) (map[string]interface{}, error) {
	return d.getThrottlerCase(throttlerName, caseName, "/cases")
}

func (d *DevClient) GetThrottlerException(throttlerName, caseName string) (map[string]interface{}, error) {
	return d.getThrottlerCase(throttlerName, caseName, "/exceptions")
}

func (d *DevClient) getThrottlerCase(throttlerName, caseName, pathTail string) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(get(d, _THROTTLERS_PREAMBLE+"/"+throttlerName+pathTail+"/"+caseName, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

func (d *DevClient) UpdateThrottlerCase(throttlerName, caseName string, updates map[string]interface{}) (map[string]interface{}, error) {
	return d.updateThrottlerCase(throttlerName, caseName, "/cases", updates)
}

func (d *DevClient) UpdateThrottlerException(throttlerName, caseName string, updates map[string]interface{}) (map[string]interface{}, error) {
	return d.updateThrottlerCase(throttlerName, caseName, "/exceptions", updates)
}

func (d *DevClient) updateThrottlerCase(throttlerName, caseName, pathTail string, updates map[string]interface{}) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(put(d, _THROTTLERS_PREAMBLE+"/"+throttlerName+pathTail+"/"+caseName, updates, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

func (d *DevClient) DeleteThrottlerCase(throttlerName, caseName string) (map[string]interface{}, error) {
	return d.deleteThrottlerCase(throttlerName, caseName, "/cases")
}

func (d *DevClient) DeleteThrottlerException(throttlerName, caseName string) (map[string]interface{}, error) {
	return d.deleteThrottlerCase(throttlerName, caseName, "/exceptions")
}

func (d *DevClient) deleteThrottlerCase(throttlerName, caseName, pathTail string) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(delete(d, _THROTTLERS_PREAMBLE+"/"+throttlerName+pathTail+"/"+caseName, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}
<commit_msg>got rid of definitions stuff<commit_after>package GoSDK

import (
//"fmt"
)

const (
	_THROTTLERS_PREAMBLE = "/admin/throttlers"
)

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

func (d *DevClient) GetAllThrottlers() (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(get(d, _THROTTLERS_PREAMBLE, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

func (d *DevClient) GetThrottler(throttlerName string) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(get(d, _THROTTLERS_PREAMBLE+"/"+throttlerName, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

func (d *DevClient) GetThrottlerCasesAndExceptions(throttlerName string) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(get(d, _THROTTLERS_PREAMBLE+"/"+throttlerName, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

func (d *DevClient) DeleteThrottlerCasesAndExceptions(throttlerName string) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(delete(d, _THROTTLERS_PREAMBLE+"/"+throttlerName, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

func (d *DevClient) GetThrottlerCases(throttlerName string) ([]interface{}, error) {
	return d.getThrottlerCases(throttlerName, "/cases")
}

func (d *DevClient) GetThrottlerExceptions(throttlerName string) ([]interface{}, error) {
	return d.getThrottlerCases(throttlerName, "/exceptions")
}

func (d *DevClient) getThrottlerCases(throttlerName, pathTail string) ([]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(get(d, _THROTTLERS_PREAMBLE+"/"+throttlerName+pathTail, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.([]interface{}), nil
}

func (d *DevClient) CreateThrottlerCase(throttlerName string, caseInfo map[string]interface{}) (map[string]interface{}, error) {
	return d.createThrottlerCase(throttlerName, "/cases", caseInfo)
}

func (d *DevClient) CreateThrottlerException(throttlerName string, caseInfo map[string]interface{}) (map[string]interface{}, error) {
	return d.createThrottlerCase(throttlerName, "/exceptions", caseInfo)
}

func (d *DevClient) createThrottlerCase(throttlerName, pathTail string, caseInfo map[string]interface{}) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(post(d, _THROTTLERS_PREAMBLE+"/"+throttlerName+pathTail, caseInfo, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

func (d *DevClient) DeleteAllThrottlerCases(throttlerName string) ([]interface{}, error) {
	return d.deleteAllThrottlerCases(throttlerName, "/cases")
}

func (d *DevClient) DeleteAllThrottlerExceptions(throttlerName string) ([]interface{}, error) {
	return d.deleteAllThrottlerCases(throttlerName, "/exceptions")
}

func (d *DevClient) deleteAllThrottlerCases(throttlerName, pathTail string) ([]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(delete(d, _THROTTLERS_PREAMBLE+"/"+throttlerName+pathTail, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.([]interface{}), nil
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

func (d *DevClient) GetThrottlerCase(throttlerName, caseName string) (map[string]interface{}, error) {
	return d.getThrottlerCase(throttlerName, caseName, "/cases")
}

func (d *DevClient) GetThrottlerException(throttlerName, caseName string) (map[string]interface{}, error) {
	return d.getThrottlerCase(throttlerName, caseName, "/exceptions")
}

func (d *DevClient) getThrottlerCase(throttlerName, caseName, pathTail string) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(get(d, _THROTTLERS_PREAMBLE+"/"+throttlerName+pathTail+"/"+caseName, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

func (d *DevClient) UpdateThrottlerCase(throttlerName, caseName string, updates map[string]interface{}) (map[string]interface{}, error) {
	return d.updateThrottlerCase(throttlerName, caseName, "/cases", updates)
}

func (d *DevClient) UpdateThrottlerException(throttlerName, caseName string, updates map[string]interface{}) (map[string]interface{}, error) {
	return d.updateThrottlerCase(throttlerName, caseName, "/exceptions", updates)
}

func (d *DevClient) updateThrottlerCase(throttlerName, caseName, pathTail string, updates map[string]interface{}) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(put(d, _THROTTLERS_PREAMBLE+"/"+throttlerName+pathTail+"/"+caseName, updates, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}

func (d *DevClient) DeleteThrottlerCase(throttlerName, caseName string) (map[string]interface{}, error) {
	return d.deleteThrottlerCase(throttlerName, caseName, "/cases")
}

func (d *DevClient) DeleteThrottlerException(throttlerName, caseName string) (map[string]interface{}, error) {
	return d.deleteThrottlerCase(throttlerName, caseName, "/exceptions")
}

func (d *DevClient) deleteThrottlerCase(throttlerName, caseName, pathTail string) (map[string]interface{}, error) {
	creds, err := d.credentials()
	if err != nil {
		return nil, err
	}
	resp, err := mapResponse(delete(d, _THROTTLERS_PREAMBLE+"/"+throttlerName+pathTail+"/"+caseName, nil, creds, nil))
	if err != nil {
		return nil, err
	}
	return resp.Body.(map[string]interface{}), nil
}
<|endoftext|>


<commit_before>// Copyright 2016 Martin Hebnes Pedersen (LA5NTA). All rights reserved.
// Use of this source code is governed by the MIT-license that can be
// found in the LICENSE file.

package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"io"
	"log"
	"os"
	"path"
	"runtime"
	"sync"
	"time"

	"github.com/fsnotify/fsnotify"
	"github.com/gorilla/websocket"

	"github.com/la5nta/pat/internal/osutil"
	"github.com/la5nta/wl2k-go/mailbox"
)

const KeepaliveInterval = 4 * time.Minute

// WSConn represent one connection in the WSHub pool
type WSConn struct {
	conn *websocket.Conn
	out  chan interface{}
}

// WSHub is a hub for broadcasting data to several websocket connections
type WSHub struct {
	mu   sync.Mutex
	pool map[*WSConn]struct{}
}

func NewWSHub() *WSHub {
	w := &WSHub{pool: map[*WSConn]struct{}{}}
	go w.watchMBox()
	return w
}

func (w *WSHub) UpdateStatus()            { w.WriteJSON(struct{ Status Status }{getStatus()}) }
func (w *WSHub) WriteProgress(p Progress) { w.WriteJSON(struct{ Progress Progress }{p}) }
func (w *WSHub) WriteNotification(n Notification) {
	w.WriteJSON(struct{ Notification Notification }{n})
}

func (w *WSHub) Prompt(p Prompt) {
	w.WriteJSON(struct{ Prompt Prompt }{p})
	go func() { <-p.cancel; w.WriteJSON(struct{ PromptAbort Prompt }{p}) }()
}

func (w *WSHub) WriteJSON(v interface{}) {
	if w == nil {
		return
	}

	w.mu.Lock()
	for c := range w.pool {
		select {
		case c.out <- v:
		case <-time.After(3 * time.Second):
			log.Println("Closing one unresponsive web socket")
			c.conn.Close()
			delete(w.pool, c)
		}
	}
	w.mu.Unlock()
}

func (w *WSHub) ClientAddrs() []string {
	if w == nil {
		return nil
	}
	w.mu.Lock()
	defer w.mu.Unlock()
	addrs := make([]string, 0, len(w.pool))
	for c := range w.pool {
		addrs = append(addrs, c.conn.RemoteAddr().String())
	}
	return addrs
}

func (w *WSHub) watchMBox() {
	// Maximise ulimit -n:
	//   fsnotify opens a file descriptor for every file in the directories it watches, which
	//   may more files than the current soft limit. The is especially a problem on macOS which
	//   has a default soft limit of only 256 files. Windows does not have a such a limit.
	if runtime.GOOS != "windows" {
		if err := osutil.RaiseOpenFileLimit(4096); err != nil {
			log.Printf("Unable to raise open file limit: %v", err)
		}
	}

	fsWatcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Println("Unable to start fs watcher: ", err)
	} else {
		p := path.Join(mbox.MBoxPath, mailbox.DIR_INBOX)
		if err := fsWatcher.Add(p); err != nil {
			log.Printf("Unable to add path '%s' to fs watcher: %s", p, err)
		}

		// These will probably fail if the first failed, but it's not important to log all.
		fsWatcher.Add(path.Join(mbox.MBoxPath, mailbox.DIR_OUTBOX))
		fsWatcher.Add(path.Join(mbox.MBoxPath, mailbox.DIR_SENT))
		fsWatcher.Add(path.Join(mbox.MBoxPath, mailbox.DIR_ARCHIVE))
		defer fsWatcher.Close()
	}

	for {
		select {
		// Filesystem events
		case e := <-fsWatcher.Events:
			if e.Op == fsnotify.Chmod {
				continue
			}
			// Make sure we don't send many of these events over a short period.
			drainUntilSilence(fsWatcher, 100*time.Millisecond)
			websocketHub.WriteJSON(struct {
				UpdateMailbox bool
			}{true})
		case err := <-fsWatcher.Errors:
			log.Println(err)
		}
	}
}

// Handle adds a new websocket to the hub
//
// It will block until the client either stops responding or closes the connection.
func (w *WSHub) Handle(conn *websocket.Conn) {
	c := &WSConn{
		conn: conn,
		out:  make(chan interface{}, 1),
	}

	w.mu.Lock()
	w.pool[c] = struct{}{}
	w.mu.Unlock()

	// Initial status update
	// (broadcasted as it includes info to other clients about this new one)
	w.UpdateStatus()

	quit := wsReadLoop(conn)

	lines, done, err := tailFile(fOptions.LogPath)
	if err != nil {
		log.Println(err)
		return
	}
	defer close(done)

	for {
		select {
		case <-time.After(KeepaliveInterval):
			c.conn.WriteJSON(struct {
				Ping bool
			}{true})
		case line := <-lines:
			c.conn.WriteJSON(struct {
				LogLine string
			}{string(line)})
		case v := <-c.out:
			err := c.conn.WriteJSON(v)
			if err != nil {
				log.Println(err)
			}
		case <-quit:
			// The read loop failed/disconnected. Remove from hub.
			c.conn.Close()
			w.mu.Lock()
			delete(w.pool, c)
			w.mu.Unlock()
			w.UpdateStatus()
			return
		}
	}
}

// drainEvents reads from w.Events and blocks until the channel has been silent for at least 50 ms.
func drainUntilSilence(w *fsnotify.Watcher, silenceDur time.Duration) {
	timer := time.NewTimer(silenceDur)
	defer timer.Stop()
	for {
		select {
		case <-w.Events:
			if !timer.Stop() {
				<-timer.C
			}
			timer.Reset(silenceDur)
		case <-timer.C:
			return
		}
	}
}

// Expects the file to never get renamed/truncated or deleted
func tailFile(path string) (<-chan []byte, chan<- struct{}, error) {
	lines := make(chan []byte)
	done := make(chan struct{})
	file, err := os.Open(path)
	if err != nil {
		return nil, nil, err
	}

	go func() {
		rd := bufio.NewReader(file)
		for {
			data, _, err := rd.ReadLine()
			if errors.Is(err, io.EOF) {
				time.Sleep(time.Millisecond * 100)
				continue
			}

			select {
			case <-done:
				file.Close()
				return
			case lines <- data:
			}
		}
	}()

	return lines, done, nil
}

func handleWSMessage(v map[string]json.RawMessage) {
	raw, ok := v["prompt_response"]
	if !ok {
		return
	}
	var resp PromptResponse
	json.Unmarshal(raw, &resp)
	promptHub.Respond(resp.ID, resp.Value, resp.Err)
}

func wsReadLoop(c *websocket.Conn) <-chan struct{} {
	quit := make(chan struct{})
	go func() {
		for {
			v := map[string]json.RawMessage{}
			err := c.ReadJSON(&v)
			if err != nil {
				close(quit)
				return
			}
			if _, ok := v["Pong"]; ok {
				// For Keepalive only. Could use a delayed ping response to detect dead connections in the future.
				continue
			}
			go handleWSMessage(v)
		}
	}()
	return quit
}
<commit_msg>Improve websocket handling (read/write deadlines)<commit_after>// Copyright 2016 Martin Hebnes Pedersen (LA5NTA). All rights reserved.
// Use of this source code is governed by the MIT-license that can be
// found in the LICENSE file.

package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"io"
	"log"
	"os"
	"path"
	"runtime"
	"sync"
	"time"

	"github.com/fsnotify/fsnotify"
	"github.com/gorilla/websocket"

	"github.com/la5nta/pat/internal/debug"
	"github.com/la5nta/pat/internal/osutil"
	"github.com/la5nta/wl2k-go/mailbox"
)

const KeepaliveInterval = 4 * time.Minute

// WSConn represent one connection in the WSHub pool
type WSConn struct {
	conn *websocket.Conn
	out  chan interface{}
}

// WSHub is a hub for broadcasting data to several websocket connections
type WSHub struct {
	mu   sync.Mutex
	pool map[*WSConn]struct{}
}

func NewWSHub() *WSHub {
	w := &WSHub{pool: map[*WSConn]struct{}{}}
	go w.watchMBox()
	return w
}

func (w *WSHub) UpdateStatus()            { w.WriteJSON(struct{ Status Status }{getStatus()}) }
func (w *WSHub) WriteProgress(p Progress) { w.WriteJSON(struct{ Progress Progress }{p}) }
func (w *WSHub) WriteNotification(n Notification) {
	w.WriteJSON(struct{ Notification Notification }{n})
}

func (w *WSHub) Prompt(p Prompt) {
	w.WriteJSON(struct{ Prompt Prompt }{p})
	go func() { <-p.cancel; w.WriteJSON(struct{ PromptAbort Prompt }{p}) }()
}

func (w *WSHub) WriteJSON(v interface{}) {
	if w == nil {
		return
	}

	w.mu.Lock()
	for c := range w.pool {
		select {
		case c.out <- v:
		case <-time.After(3 * time.Second):
			log.Println("Closing one unresponsive web socket")
			c.conn.Close()
			delete(w.pool, c)
		}
	}
	w.mu.Unlock()
}

func (w *WSHub) ClientAddrs() []string {
	if w == nil {
		return nil
	}
	w.mu.Lock()
	defer w.mu.Unlock()
	addrs := make([]string, 0, len(w.pool))
	for c := range w.pool {
		addrs = append(addrs, c.conn.RemoteAddr().String())
	}
	return addrs
}

func (w *WSHub) watchMBox() {
	// Maximise ulimit -n:
	//   fsnotify opens a file descriptor for every file in the directories it watches, which
	//   may more files than the current soft limit. The is especially a problem on macOS which
	//   has a default soft limit of only 256 files. Windows does not have a such a limit.
	if runtime.GOOS != "windows" {
		if err := osutil.RaiseOpenFileLimit(4096); err != nil {
			log.Printf("Unable to raise open file limit: %v", err)
		}
	}

	fsWatcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Println("Unable to start fs watcher: ", err)
	} else {
		p := path.Join(mbox.MBoxPath, mailbox.DIR_INBOX)
		if err := fsWatcher.Add(p); err != nil {
			log.Printf("Unable to add path '%s' to fs watcher: %s", p, err)
		}

		// These will probably fail if the first failed, but it's not important to log all.
		fsWatcher.Add(path.Join(mbox.MBoxPath, mailbox.DIR_OUTBOX))
		fsWatcher.Add(path.Join(mbox.MBoxPath, mailbox.DIR_SENT))
		fsWatcher.Add(path.Join(mbox.MBoxPath, mailbox.DIR_ARCHIVE))
		defer fsWatcher.Close()
	}

	for {
		select {
		// Filesystem events
		case e := <-fsWatcher.Events:
			if e.Op == fsnotify.Chmod {
				continue
			}
			// Make sure we don't send many of these events over a short period.
			drainUntilSilence(fsWatcher, 100*time.Millisecond)
			websocketHub.WriteJSON(struct {
				UpdateMailbox bool
			}{true})
		case err := <-fsWatcher.Errors:
			log.Println(err)
		}
	}
}

// Handle adds a new websocket to the hub
//
// It will block until the client either stops responding or closes the connection.
func (w *WSHub) Handle(conn *websocket.Conn) {
	debug.Printf("ws[%s] subscribed", conn.RemoteAddr())
	c := &WSConn{
		conn: conn,
		out:  make(chan interface{}, 1),
	}

	w.mu.Lock()
	w.pool[c] = struct{}{}
	w.mu.Unlock()

	// Initial status update
	// (broadcasted as it includes info to other clients about this new one)
	w.UpdateStatus()

	quit := wsReadLoop(conn)

	// Disconnect and remove client when this handler returns.
	defer func() {
		debug.Printf("ws[%s] unsubscribing...", conn.RemoteAddr())
		c.conn.Close()
		w.mu.Lock()
		delete(w.pool, c)
		w.mu.Unlock()
		w.UpdateStatus()
		debug.Printf("ws[%s] unsubscribed", conn.RemoteAddr())
	}()

	lines, done, err := tailFile(fOptions.LogPath)
	if err != nil {
		log.Println(err)
		return
	}
	defer close(done)
	for {
		var err error
		c.conn.SetWriteDeadline(time.Time{})
		select {
		case <-time.After(KeepaliveInterval):
			debug.Printf("ws[%s] ping", conn.RemoteAddr())
			c.conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
			err = c.conn.WriteJSON(struct {
				Ping bool
			}{true})
		case line := <-lines:
			c.conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
			err = c.conn.WriteJSON(struct {
				LogLine string
			}{string(line)})
		case v := <-c.out:
			c.conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
			err = c.conn.WriteJSON(v)
		case <-quit:
			// The read loop failed/disconnected. Abort.
			return
		}
		if err != nil {
			debug.Printf("ws[%s] write error: %v", conn.RemoteAddr(), err)
			return
		}
	}
}

// drainEvents reads from w.Events and blocks until the channel has been silent for at least 50 ms.
func drainUntilSilence(w *fsnotify.Watcher, silenceDur time.Duration) {
	timer := time.NewTimer(silenceDur)
	defer timer.Stop()
	for {
		select {
		case <-w.Events:
			if !timer.Stop() {
				<-timer.C
			}
			timer.Reset(silenceDur)
		case <-timer.C:
			return
		}
	}
}

// Expects the file to never get renamed/truncated or deleted
func tailFile(path string) (<-chan []byte, chan<- struct{}, error) {
	lines := make(chan []byte)
	done := make(chan struct{})
	file, err := os.Open(path)
	if err != nil {
		return nil, nil, err
	}

	go func() {
		rd := bufio.NewReader(file)
		for {
			data, _, err := rd.ReadLine()
			if errors.Is(err, io.EOF) {
				time.Sleep(time.Millisecond * 100)
				continue
			}

			select {
			case <-done:
				file.Close()
				return
			case lines <- data:
			}
		}
	}()

	return lines, done, nil
}

func handleWSMessage(v map[string]json.RawMessage) {
	raw, ok := v["prompt_response"]
	if !ok {
		return
	}
	var resp PromptResponse
	json.Unmarshal(raw, &resp)
	promptHub.Respond(resp.ID, resp.Value, resp.Err)
}

func wsReadLoop(c *websocket.Conn) <-chan struct{} {
	quit := make(chan struct{})
	go func() {
		for {
			v := map[string]json.RawMessage{}
			// We should at least get a ping response once per KeepaliveInterval.
			c.SetReadDeadline(time.Now().Add(KeepaliveInterval + 10*time.Second))
			err := c.ReadJSON(&v)
			if err != nil {
				debug.Printf("ws[%s] read error: %v", c.RemoteAddr(), err)
				close(quit)
				return
			}
			if _, ok := v["Pong"]; ok {
				// That's the Ping response.
				debug.Printf("ws[%s] pong", c.RemoteAddr())
				continue
			}
			go handleWSMessage(v)
		}
	}()
	return quit
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"io"
	"os"
	"registry"
	"sync"

	log "github.com/GameGophers/libs/nsq-logger"
	"github.com/GameGophers/libs/services"
	. "github.com/GameGophers/libs/services/proto"
	"golang.org/x/net/context"
)

import (
	. "types"
)

var (
	ERROR_NOT_AUTHORIZED         = errors.New("User not authorized")
	ERROR_CANNOT_FIND_GAMESERVER = errors.New("Cannot find Game Server")
)

/*
 * 长连接 双向流 RPC 到game servers, 转发agent <==> game 消息
 */

var (
	_default_forward forwarder
)

type forwarder struct {
	gs map[string]GameService_PacketClient
	sync.Mutex
}

func init() {
	_default_forward = forwarder{}
	//TODO connect all game server and forward packet
	clients, err := services.GetAllService(services.SERVICE_GAME)
	if err != nil {
		log.Critical(err)
		os.Exit(-1)
	}
	for k, cli := range clients {
		service, _ := cli.(GameServiceClient)
		stream, err := service.Packet(context.Background())
		if err != nil {
			log.Critical(err)
			os.Exit(-1)
		}
		_default_forward.gs[k] = stream

	}
	for k, c := range _default_forward.gs {
		go _default_forward.recv(k, c)
	}
}

func (f *forwarder) get(id string) GameService_PacketClient {
	f.Lock()
	defer f.Unlock()
	return f.gs[string(services.SERVICE_GAME)+"/"+id]

}

func (f *forwarder) recv(key string, cli GameService_PacketClient) {
	for {
		in, err := cli.Recv()
		if err == io.EOF {
			log.Infof("stream recv EOF err : %v", err)
			return
		}
		if err != nil {
			log.Critical("Failed to receive a note : %v", err)
			continue
		}
		registry.Deliver(in.Uid, in.Content)
	}
}

// forwarding messages to game server
func forward(sess *Session, p []byte) error {
	pkt := &Game_Packet{
		Uid:     sess.UserId,
		Content: p,
	}

	if sess.Flag&SESS_AUTHORIZED != 0 {
		// get connection
		c := _default_forward.get(sess.GSID)
		if c == nil {
			log.Criticalf("gsid %v cannot find it's game service:", sess.GSID)
			return ERROR_CANNOT_FIND_GAMESERVER
		}

		// send the packet
		if err := c.Send(pkt); err != nil {
			log.Criticalf("Failed to send pkt %v", err)
			return err
		}
		return nil
	}
	return ERROR_NOT_AUTHORIZED
}
<commit_msg>update<commit_after>package main

import (
	"errors"
	"os"
	"registry"
	"sync"

	log "github.com/GameGophers/libs/nsq-logger"
	"github.com/GameGophers/libs/services"
	. "github.com/GameGophers/libs/services/proto"
	"golang.org/x/net/context"
)

import (
	. "types"
)

var (
	ERROR_NOT_AUTHORIZED         = errors.New("User not authorized")
	ERROR_CANNOT_FIND_GAMESERVER = errors.New("Cannot find Game Server")
)

/*
 * 长连接 双向流 RPC 到game servers, 转发agent <==> game 消息
 */

var (
	_default_forward forwarder
)

type forwarder struct {
	gs map[string]GameService_PacketClient
	sync.Mutex
}

func init() {
	_default_forward = forwarder{}
	//TODO connect all game server and forward packet
	clients, err := services.GetAllService(services.SERVICE_GAME)
	if err != nil {
		log.Critical(err)
		os.Exit(-1)
	}
	for k, cli := range clients {
		service, _ := cli.(GameServiceClient)
		stream, err := service.Packet(context.Background())
		if err != nil {
			log.Critical(err)
			os.Exit(-1)
		}
		_default_forward.gs[k] = stream

	}
	for k, c := range _default_forward.gs {
		go _default_forward.recv(k, c)
	}
}

func (f *forwarder) get(id string) GameService_PacketClient {
	f.Lock()
	defer f.Unlock()
	return f.gs[string(services.SERVICE_GAME)+"/"+id]

}

func (f *forwarder) recv(key string, cli GameService_PacketClient) {
	for {
		in, err := cli.Recv()
		if err != nil {
			log.Critical(err)
			return
		}
		registry.Deliver(in.Uid, in.Content)
	}
}

// forwarding messages to game server
func forward(sess *Session, p []byte) error {
	pkt := &Game_Packet{
		Uid:     sess.UserId,
		Content: p,
	}

	if sess.Flag&SESS_AUTHORIZED != 0 {
		// get connection
		c := _default_forward.get(sess.GSID)
		if c == nil {
			log.Criticalf("gsid %v cannot find it's game service:", sess.GSID)
			return ERROR_CANNOT_FIND_GAMESERVER
		}

		// send the packet
		if err := c.Send(pkt); err != nil {
			log.Critical(err)
			return err
		}
		return nil
	}
	return ERROR_NOT_AUTHORIZED
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC
//
// 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 pubsub

import (
	"io"
	"sync"
	"time"

	gax "github.com/googleapis/gax-go"
	"golang.org/x/net/context"
	pb "google.golang.org/genproto/googleapis/pubsub/v1"
	"google.golang.org/grpc"
)

// A pullStream supports the methods of a StreamingPullClient, but re-opens
// the stream on a retryable error.
type pullStream struct {
	ctx  context.Context
	open func() (pb.Subscriber_StreamingPullClient, error)

	mu  sync.Mutex
	spc *pb.Subscriber_StreamingPullClient
	err error // permanent error
}

// for testing
type streamingPullFunc func(context.Context, ...gax.CallOption) (pb.Subscriber_StreamingPullClient, error)

func newPullStream(ctx context.Context, streamingPull streamingPullFunc, subName string) *pullStream {
	ctx = withSubscriptionKey(ctx, subName)
	return &pullStream{
		ctx: ctx,
		open: func() (pb.Subscriber_StreamingPullClient, error) {
			spc, err := streamingPull(ctx, gax.WithGRPCOptions(grpc.MaxCallRecvMsgSize(maxSendRecvBytes)))
			if err == nil {
				recordStat(ctx, StreamRequestCount, 1)
				err = spc.Send(&pb.StreamingPullRequest{
					Subscription: subName,
					// We modack messages when we receive them, so this value doesn't matter too much.
					StreamAckDeadlineSeconds: 60,
				})
			}
			if err != nil {
				return nil, err
			}
			return spc, nil
		},
	}
}

// get returns either a valid *StreamingPullClient (SPC), or a permanent error.
// If the argument is nil, this is the first call for an RPC, and the current
// SPC will be returned (or a new one will be opened). Otherwise, this call is a
// request to re-open the stream because of a retryable error, and the argument
// is a pointer to the SPC that returned the error.
func (s *pullStream) get(spc *pb.Subscriber_StreamingPullClient) (*pb.Subscriber_StreamingPullClient, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	// A stored error is permanent.
	if s.err != nil {
		return nil, s.err
	}
	// If the context is done, so are we.
	select {
	case <-s.ctx.Done():
		s.err = s.ctx.Err()
		return nil, s.err
	default:
	}
	// TODO(jba): We can use the following instead of the above after we drop support for 1.8:
	// s.err = s.ctx.Err()
	// if s.err != nil {
	// 	return nil, s.err
	// }

	// If the current and argument SPCs differ, return the current one. This subsumes two cases:
	// 1. We have an SPC and the caller is getting the stream for the first time.
	// 2. The caller wants to retry, but they have an older SPC; we've already retried.
	if spc != s.spc {
		return s.spc, nil
	}
	// Either this is the very first call on this stream (s.spc == nil), or we have a valid
	// retry request. Either way, open a new stream.
	// The lock is held here for a long time, but it doesn't matter because no callers could get
	// anything done anyway.
	s.spc = new(pb.Subscriber_StreamingPullClient)
	*s.spc, s.err = s.openWithRetry() // Any error from openWithRetry is permanent.
	return s.spc, s.err
}

func (s *pullStream) openWithRetry() (pb.Subscriber_StreamingPullClient, error) {
	var bo gax.Backoff
	for {
		recordStat(s.ctx, StreamOpenCount, 1)
		spc, err := s.open()
		if err != nil && isRetryable(err) {
			recordStat(s.ctx, StreamRetryCount, 1)
			if err := gax.Sleep(s.ctx, bo.Pause()); err != nil {
				return nil, err
			}
			continue
		}
		return spc, err
	}
}

func (s *pullStream) call(f func(pb.Subscriber_StreamingPullClient) error) error {
	var (
		spc *pb.Subscriber_StreamingPullClient
		err error
		bo  gax.Backoff
	)
	for {
		spc, err = s.get(spc)
		if err != nil {
			return err
		}
		start := time.Now()
		err = f(*spc)
		if err != nil {
			if isRetryable(err) {
				recordStat(s.ctx, StreamRetryCount, 1)
				if time.Since(start) < 30*time.Second { // don't sleep if we've been blocked for a while
					if err := gax.Sleep(s.ctx, bo.Pause()); err != nil {
						return err
					}
				}
				continue
			}
			s.mu.Lock()
			s.err = err
			s.mu.Unlock()
		}
		return err
	}
}

func (s *pullStream) Send(req *pb.StreamingPullRequest) error {
	return s.call(func(spc pb.Subscriber_StreamingPullClient) error {
		recordStat(s.ctx, AckCount, int64(len(req.AckIds)))
		zeroes := 0
		for _, mds := range req.ModifyDeadlineSeconds {
			if mds == 0 {
				zeroes++
			}
		}
		recordStat(s.ctx, NackCount, int64(zeroes))
		recordStat(s.ctx, ModAckCount, int64(len(req.ModifyDeadlineSeconds)-zeroes))
		recordStat(s.ctx, StreamRequestCount, 1)
		return spc.Send(req)
	})
}

func (s *pullStream) Recv() (*pb.StreamingPullResponse, error) {
	var res *pb.StreamingPullResponse
	err := s.call(func(spc pb.Subscriber_StreamingPullClient) error {
		var err error
		recordStat(s.ctx, StreamResponseCount, 1)
		res, err = spc.Recv()
		if err == nil {
			recordStat(s.ctx, PullCount, int64(len(res.ReceivedMessages)))
		}
		return err
	})
	return res, err
}

func (s *pullStream) CloseSend() error {
	err := s.call(func(spc pb.Subscriber_StreamingPullClient) error {
		return spc.CloseSend()
	})
	s.mu.Lock()
	s.err = io.EOF // should not be retried
	s.mu.Unlock()
	return err
}
<commit_msg>pubsub: simplify ctx.Err check<commit_after>// Copyright 2018 Google LLC
//
// 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 pubsub

import (
	"io"
	"sync"
	"time"

	gax "github.com/googleapis/gax-go"
	"golang.org/x/net/context"
	pb "google.golang.org/genproto/googleapis/pubsub/v1"
	"google.golang.org/grpc"
)

// A pullStream supports the methods of a StreamingPullClient, but re-opens
// the stream on a retryable error.
type pullStream struct {
	ctx  context.Context
	open func() (pb.Subscriber_StreamingPullClient, error)

	mu  sync.Mutex
	spc *pb.Subscriber_StreamingPullClient
	err error // permanent error
}

// for testing
type streamingPullFunc func(context.Context, ...gax.CallOption) (pb.Subscriber_StreamingPullClient, error)

func newPullStream(ctx context.Context, streamingPull streamingPullFunc, subName string) *pullStream {
	ctx = withSubscriptionKey(ctx, subName)
	return &pullStream{
		ctx: ctx,
		open: func() (pb.Subscriber_StreamingPullClient, error) {
			spc, err := streamingPull(ctx, gax.WithGRPCOptions(grpc.MaxCallRecvMsgSize(maxSendRecvBytes)))
			if err == nil {
				recordStat(ctx, StreamRequestCount, 1)
				err = spc.Send(&pb.StreamingPullRequest{
					Subscription: subName,
					// We modack messages when we receive them, so this value doesn't matter too much.
					StreamAckDeadlineSeconds: 60,
				})
			}
			if err != nil {
				return nil, err
			}
			return spc, nil
		},
	}
}

// get returns either a valid *StreamingPullClient (SPC), or a permanent error.
// If the argument is nil, this is the first call for an RPC, and the current
// SPC will be returned (or a new one will be opened). Otherwise, this call is a
// request to re-open the stream because of a retryable error, and the argument
// is a pointer to the SPC that returned the error.
func (s *pullStream) get(spc *pb.Subscriber_StreamingPullClient) (*pb.Subscriber_StreamingPullClient, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	// A stored error is permanent.
	if s.err != nil {
		return nil, s.err
	}
	// If the context is done, so are we.
	s.err = s.ctx.Err()
	if s.err != nil {
		return nil, s.err
	}

	// If the current and argument SPCs differ, return the current one. This subsumes two cases:
	// 1. We have an SPC and the caller is getting the stream for the first time.
	// 2. The caller wants to retry, but they have an older SPC; we've already retried.
	if spc != s.spc {
		return s.spc, nil
	}
	// Either this is the very first call on this stream (s.spc == nil), or we have a valid
	// retry request. Either way, open a new stream.
	// The lock is held here for a long time, but it doesn't matter because no callers could get
	// anything done anyway.
	s.spc = new(pb.Subscriber_StreamingPullClient)
	*s.spc, s.err = s.openWithRetry() // Any error from openWithRetry is permanent.
	return s.spc, s.err
}

func (s *pullStream) openWithRetry() (pb.Subscriber_StreamingPullClient, error) {
	var bo gax.Backoff
	for {
		recordStat(s.ctx, StreamOpenCount, 1)
		spc, err := s.open()
		if err != nil && isRetryable(err) {
			recordStat(s.ctx, StreamRetryCount, 1)
			if err := gax.Sleep(s.ctx, bo.Pause()); err != nil {
				return nil, err
			}
			continue
		}
		return spc, err
	}
}

func (s *pullStream) call(f func(pb.Subscriber_StreamingPullClient) error) error {
	var (
		spc *pb.Subscriber_StreamingPullClient
		err error
		bo  gax.Backoff
	)
	for {
		spc, err = s.get(spc)
		if err != nil {
			return err
		}
		start := time.Now()
		err = f(*spc)
		if err != nil {
			if isRetryable(err) {
				recordStat(s.ctx, StreamRetryCount, 1)
				if time.Since(start) < 30*time.Second { // don't sleep if we've been blocked for a while
					if err := gax.Sleep(s.ctx, bo.Pause()); err != nil {
						return err
					}
				}
				continue
			}
			s.mu.Lock()
			s.err = err
			s.mu.Unlock()
		}
		return err
	}
}

func (s *pullStream) Send(req *pb.StreamingPullRequest) error {
	return s.call(func(spc pb.Subscriber_StreamingPullClient) error {
		recordStat(s.ctx, AckCount, int64(len(req.AckIds)))
		zeroes := 0
		for _, mds := range req.ModifyDeadlineSeconds {
			if mds == 0 {
				zeroes++
			}
		}
		recordStat(s.ctx, NackCount, int64(zeroes))
		recordStat(s.ctx, ModAckCount, int64(len(req.ModifyDeadlineSeconds)-zeroes))
		recordStat(s.ctx, StreamRequestCount, 1)
		return spc.Send(req)
	})
}

func (s *pullStream) Recv() (*pb.StreamingPullResponse, error) {
	var res *pb.StreamingPullResponse
	err := s.call(func(spc pb.Subscriber_StreamingPullClient) error {
		var err error
		recordStat(s.ctx, StreamResponseCount, 1)
		res, err = spc.Recv()
		if err == nil {
			recordStat(s.ctx, PullCount, int64(len(res.ReceivedMessages)))
		}
		return err
	})
	return res, err
}

func (s *pullStream) CloseSend() error {
	err := s.call(func(spc pb.Subscriber_StreamingPullClient) error {
		return spc.CloseSend()
	})
	s.mu.Lock()
	s.err = io.EOF // should not be retried
	s.mu.Unlock()
	return err
}
<|endoftext|>