<commit_before>package main

import (
	"crypto/tls"
	"fmt"
	"io"
	"log"
	"net"
	"net/http/httputil"
	"net/url"
	"os"
	"strconv"
	"strings"

	"github.com/docker/docker/pkg/term"
	"github.com/remind101/empire/pkg/heroku"
)

var (
	detachedRun bool
	dynoSize    string
)

var cmdRun = &Command{
	Run:             maybeMessage(runRun),
	Usage:           "run [-s <size>] [-d] <command> [<argument>...]",
	NeedsApp:        true,
	OptionalMessage: true,
	Category:        "dyno",
	Short:           "run a process in a dyno",
	Long: `
Run a process on Heroku. Flags such as` + " `-a` " + `may be parsed out of
the command unless the command is quoted or provided after a
double-dash (--).

You may need to press enter after this command to continue.

Options:

    -s <size>  set the size for this dyno (e.g. 2X)
    -d         run in detached mode instead of attached to terminal

Examples:

    $ emp run echo "hello"
    Running ` + "`echo \"hello\"`" + ` on myapp as run.1234:
    "hello"

    $ emp run console
    Running ` + "`console`" + ` on myapp as run.5678:
    Loading production environment (Rails 3.2.14)
    irb(main):001:0> ...

    $ emp run -d -s 2X bin/my_worker
    Ran ` + "`bin/my_worker`" + ` on myapp as run.4321, detached.

    $ emp run -a myapp -- ls -a /
    Running ` + "`ls -a bin /`" + ` on myapp as run.8650:
    /:
    .  ..  app  bin  dev  etc  home  lib  lib64  lost+found  proc  sbin  tmp  usr  var
`,
}

func init() {
	cmdRun.Flag.BoolVarP(&detachedRun, "detached", "d", false, "detached")
	cmdRun.Flag.StringVarP(&dynoSize, "size", "s", "", "dyno size")
}

func runRun(cmd *Command, args []string) {
	if len(args) == 0 {
		cmd.PrintUsage()
		os.Exit(2)
	}
	appname := mustApp()
	message := getMessage()

	w, err := term.GetWinsize(inFd)
	if err != nil {
		// If syscall.TIOCGWINSZ is not supported by the device, we're
		// probably trying to run tests. Set w to some sensible default.
		if err.Error() == "operation not supported by device" {
			w = &term.Winsize{
				Height: 20,
				Width:  80,
			}
		} else {
			printFatal(err.Error())
		}
	}

	attached := !detachedRun
	opts := heroku.DynoCreateOpts{Attach: &attached, Message: message}
	if attached {
		env := map[string]string{
			"COLUMNS": strconv.Itoa(int(w.Width)),
			"LINES":   strconv.Itoa(int(w.Height)),
			"TERM":    os.Getenv("TERM"),
		}
		opts.Env = &env
	}
	if dynoSize != "" {
		opts.Size = &dynoSize
	}

	command := strings.Join(args, " ")
	if detachedRun {
		dyno, err := client.DynoCreate(appname, command, &opts)
		must(err)

		log.Printf("Ran `%s` on %s as %s, detached.", dyno.Command, appname, dyno.Name)
		return
	}

	params := struct {
		Command string             `json:"command"`
		Attach  *bool              `json:"attach,omitempty"`
		Env     *map[string]string `json:"env,omitempty"`
		Size    *string            `json:"size,omitempty"`
	}{
		Command: command,
		Attach:  opts.Attach,
		Env:     opts.Env,
		Size:    opts.Size,
	}

	rh := heroku.RequestHeaders{CommitMessage: message}
	req, err := client.NewRequest("POST", "/apps/"+appname+"/dynos", params, rh.Headers())
	must(err)

	u, err := url.Parse(apiURL)
	must(err)

	proto, address := dialParams(u)

	var dial net.Conn
	if proto == "tls" {
		dial, err = tlsDial("tcp", address, &tls.Config{})
		if err != nil {
			printFatal(err.Error())
		}
	} else {
		dial, err = net.Dial(proto, address)
		if err != nil {
			printFatal(err.Error())
		}
	}

	clientconn := httputil.NewClientConn(dial, nil)
	defer clientconn.Close()
	res, err := clientconn.Do(req)
	defer res.Body.Close()
	if err = heroku.CheckResp(res); err != nil {
		printFatal(err.Error())
	}
	if err != nil && err != httputil.ErrPersistEOF {
		printFatal(err.Error())
	}
	rwc, br := clientconn.Hijack()
	defer rwc.Close()

	if isTerminalIn && isTerminalOut {
		state, err := term.SetRawTerminal(inFd)
		if err != nil {
			printFatal(err.Error())
		}
		defer term.RestoreTerminal(inFd, state)
	}

	errChanOut := make(chan error, 1)
	errChanIn := make(chan error, 1)
	exit := make(chan bool)
	go func() {
		defer close(exit)
		defer close(errChanOut)
		var err error
		_, err = io.Copy(os.Stdout, br)
		errChanOut <- err
	}()
	go func() {
		_, err := io.Copy(rwc, os.Stdin)
		errChanIn <- err
		rwc.(interface {
			CloseWrite() error
		}).CloseWrite()
	}()
	<-exit
	select {
	case err = <-errChanIn:
		must(err)
	case err = <-errChanOut:
		must(err)
	}
}

func dialParams(u *url.URL) (proto, address string) {
	// u.Host can be either host or host:port
	host, port := splitHost(u.Host)

	switch u.Scheme {
	case "https":
		proto = "tls"
		if port == "" {
			port = "443"
		}
	case "http":
		proto = "tcp"
		if port == "" {
			port = "80"
		}
	default:
		panic(fmt.Sprintf("I don't know what to do for the %v protocol", u.Scheme))
	}

	address = fmt.Sprintf("%s:%s", host, port)

	return
}

func splitHost(hostPort string) (host, port string) {
	parts := strings.SplitN(hostPort, ":", 2)
	host = parts[0]
	if len(parts) == 2 {
		port = parts[1]
	}
	return
}
<commit_msg>Update run.go<commit_after>package main

import (
	"crypto/tls"
	"fmt"
	"io"
	"log"
	"net"
	"net/http/httputil"
	"net/url"
	"os"
	"strconv"
	"strings"

	"github.com/docker/docker/pkg/term"
	"github.com/remind101/empire/pkg/heroku"
)

var (
	detachedRun bool
	dynoSize    string
)

var cmdRun = &Command{
	Run:             maybeMessage(runRun),
	Usage:           "run [-s <size>] [-d] <command> [<argument>...]",
	NeedsApp:        true,
	OptionalMessage: true,
	Category:        "dyno",
	Short:           "run a process in a dyno",
	Long: `
Run a process on Heroku. Flags such as` + " `-a` " + `may be parsed out of
the command unless the command is quoted or provided after a
double-dash (--).

When running an attached process that reads from stdin (like bash) you may experience a "hang".
Usually, pressing a key like "enter" will flush the output to your terminal.
See https://github.com/remind101/empire/issues/609

Options:

    -s <size>  set the size for this dyno (e.g. 2X)
    -d         run in detached mode instead of attached to terminal

Examples:

    $ emp run echo "hello"
    Running ` + "`echo \"hello\"`" + ` on myapp as run.1234:
    "hello"

    $ emp run console
    Running ` + "`console`" + ` on myapp as run.5678:
    Loading production environment (Rails 3.2.14)
    irb(main):001:0> ...

    $ emp run -d -s 2X bin/my_worker
    Ran ` + "`bin/my_worker`" + ` on myapp as run.4321, detached.

    $ emp run -a myapp -- ls -a /
    Running ` + "`ls -a bin /`" + ` on myapp as run.8650:
    /:
    .  ..  app  bin  dev  etc  home  lib  lib64  lost+found  proc  sbin  tmp  usr  var
`,
}

func init() {
	cmdRun.Flag.BoolVarP(&detachedRun, "detached", "d", false, "detached")
	cmdRun.Flag.StringVarP(&dynoSize, "size", "s", "", "dyno size")
}

func runRun(cmd *Command, args []string) {
	if len(args) == 0 {
		cmd.PrintUsage()
		os.Exit(2)
	}
	appname := mustApp()
	message := getMessage()

	w, err := term.GetWinsize(inFd)
	if err != nil {
		// If syscall.TIOCGWINSZ is not supported by the device, we're
		// probably trying to run tests. Set w to some sensible default.
		if err.Error() == "operation not supported by device" {
			w = &term.Winsize{
				Height: 20,
				Width:  80,
			}
		} else {
			printFatal(err.Error())
		}
	}

	attached := !detachedRun
	opts := heroku.DynoCreateOpts{Attach: &attached, Message: message}
	if attached {
		env := map[string]string{
			"COLUMNS": strconv.Itoa(int(w.Width)),
			"LINES":   strconv.Itoa(int(w.Height)),
			"TERM":    os.Getenv("TERM"),
		}
		opts.Env = &env
	}
	if dynoSize != "" {
		opts.Size = &dynoSize
	}

	command := strings.Join(args, " ")
	if detachedRun {
		dyno, err := client.DynoCreate(appname, command, &opts)
		must(err)

		log.Printf("Ran `%s` on %s as %s, detached.", dyno.Command, appname, dyno.Name)
		return
	}

	params := struct {
		Command string             `json:"command"`
		Attach  *bool              `json:"attach,omitempty"`
		Env     *map[string]string `json:"env,omitempty"`
		Size    *string            `json:"size,omitempty"`
	}{
		Command: command,
		Attach:  opts.Attach,
		Env:     opts.Env,
		Size:    opts.Size,
	}

	rh := heroku.RequestHeaders{CommitMessage: message}
	req, err := client.NewRequest("POST", "/apps/"+appname+"/dynos", params, rh.Headers())
	must(err)

	u, err := url.Parse(apiURL)
	must(err)

	proto, address := dialParams(u)

	var dial net.Conn
	if proto == "tls" {
		dial, err = tlsDial("tcp", address, &tls.Config{})
		if err != nil {
			printFatal(err.Error())
		}
	} else {
		dial, err = net.Dial(proto, address)
		if err != nil {
			printFatal(err.Error())
		}
	}

	clientconn := httputil.NewClientConn(dial, nil)
	defer clientconn.Close()
	res, err := clientconn.Do(req)
	defer res.Body.Close()
	if err = heroku.CheckResp(res); err != nil {
		printFatal(err.Error())
	}
	if err != nil && err != httputil.ErrPersistEOF {
		printFatal(err.Error())
	}
	rwc, br := clientconn.Hijack()
	defer rwc.Close()

	if isTerminalIn && isTerminalOut {
		state, err := term.SetRawTerminal(inFd)
		if err != nil {
			printFatal(err.Error())
		}
		defer term.RestoreTerminal(inFd, state)
	}

	errChanOut := make(chan error, 1)
	errChanIn := make(chan error, 1)
	exit := make(chan bool)
	go func() {
		defer close(exit)
		defer close(errChanOut)
		var err error
		_, err = io.Copy(os.Stdout, br)
		errChanOut <- err
	}()
	go func() {
		_, err := io.Copy(rwc, os.Stdin)
		errChanIn <- err
		rwc.(interface {
			CloseWrite() error
		}).CloseWrite()
	}()
	<-exit
	select {
	case err = <-errChanIn:
		must(err)
	case err = <-errChanOut:
		must(err)
	}
}

func dialParams(u *url.URL) (proto, address string) {
	// u.Host can be either host or host:port
	host, port := splitHost(u.Host)

	switch u.Scheme {
	case "https":
		proto = "tls"
		if port == "" {
			port = "443"
		}
	case "http":
		proto = "tcp"
		if port == "" {
			port = "80"
		}
	default:
		panic(fmt.Sprintf("I don't know what to do for the %v protocol", u.Scheme))
	}

	address = fmt.Sprintf("%s:%s", host, port)

	return
}

func splitHost(hostPort string) (host, port string) {
	parts := strings.SplitN(hostPort, ":", 2)
	host = parts[0]
	if len(parts) == 2 {
		port = parts[1]
	}
	return
}
<|endoftext|>


<commit_before>package env

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
)

type Environment struct {
	Catalogs
	Halos
	MemoDir string

}

//////////////
// Catalogs //
//////////////

type Catalogs struct {
	snapMin int
	names [][]string
}

func (cat *Catalogs) Blocks() int {
	return len(cat.names[0])
}

func (cat *Catalogs) ParticleCatalog(snap, block int) string {
	return cat.names[snap - cat.snapMin][block]
}

func (cat *Catalogs) InitGotetra(
	format string, snapMin, snapMax int64, blockMins, blockMaxes []int64,
	validate bool,
) error {
	if len(blockMins) != 3 {
		return fmt.Errorf(
			"'BlockMins' had %d elements, but 3 are required for " +
			"gotetra catalogs.", len(blockMins),
		)
	}

	cat.names = make([][]string, snapMax - snapMin + 1)
	cat.snapMin = int(snapMin)

	for snap := snapMin; snap <= snapMax; snap++ {
		for x := blockMins[0]; x <= blockMaxes[0]; x++ {
			for y := blockMins[1]; y <= blockMaxes[1]; y++ {
				for z := blockMins[2]; z <= blockMaxes[2]; z++ {

					fname := fmt.Sprintf(format, snap, x, y, z)
					if validate {
						_, err := os.Stat(fname)
						if err != nil { return err }
					}

					cat.names[snap - snapMin] = append(
						cat.names[snap - snapMin], fname,
					)
				}
			}
		}
	}

	return nil
}

func (cat *Catalogs) InitLGadget2(
	format string, snapMin, snapMax int64, blockMins, blockMaxes []int,
) error {
	panic("NYI")
}

///////////
// Halos //
///////////

type Halos struct {
	snapMin int
	snapOffset int
	names []string
}

func (h *Halos) HaloCatalog(snap int) string {
	return h.names[snap - h.snapMin]
}

func (h *Halos) SnapOffset() int {
	return h.snapOffset
}

func (h *Halos) InitRockstar(dir string, snapMin, snapMax int64) error {
	infos, err := ioutil.ReadDir(dir)
	if err != nil { return err }

	h.snapOffset = int(snapMax) - len(infos)

	h.snapMin = int(snapMin)
	h.names = []string{}
	for i := range infos {
		h.names = append(h.names, path.Join(dir, infos[i].Name()))
	}

	if len(h.names) < int(snapMax - snapMin) + 1 {
		return fmt.Errorf(
			"There are %d files in the 'HaloDir' directory, %s, but " +
			"'SnapMin' = %d and 'SnapMax' = %d.",
			len(h.names), dir, snapMin, snapMax,
		)
	}
	h.names = h.names[len(h.names) - int(snapMax - snapMin + 1):]

	return nil
}
<commit_msg>Added env.CatalogType, env.HaloType, and env.TreeType.<commit_after>package env

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
)

type (
	CatalogType int
	HaloType int
	TreeType int
)

const (
	Gotetra CatalogType = iota

	Rockstar HaloType = iota
	NilHalo

	ConsistentTrees TreeType = iota
	NilTree
)

type Environment struct {
	Catalogs
	Halos
	MemoDir string

}

//////////////
// Catalogs //
//////////////

type Catalogs struct {
	CatalogType
	snapMin int
	names [][]string
}

func (cat *Catalogs) Blocks() int {
	return len(cat.names[0])
}

func (cat *Catalogs) ParticleCatalog(snap, block int) string {
	return cat.names[snap - cat.snapMin][block]
}

func (cat *Catalogs) InitGotetra(
	format string, snapMin, snapMax int64, blockMins, blockMaxes []int64,
	validate bool,
) error {
	cat.CatalogType = Gotetra

	if len(blockMins) != 3 {
		return fmt.Errorf(
			"'BlockMins' had %d elements, but 3 are required for " +
			"gotetra catalogs.", len(blockMins),
		)
	}

	cat.names = make([][]string, snapMax - snapMin + 1)
	cat.snapMin = int(snapMin)

	for snap := snapMin; snap <= snapMax; snap++ {
		for x := blockMins[0]; x <= blockMaxes[0]; x++ {
			for y := blockMins[1]; y <= blockMaxes[1]; y++ {
				for z := blockMins[2]; z <= blockMaxes[2]; z++ {

					fname := fmt.Sprintf(format, snap, x, y, z)
					if validate {
						_, err := os.Stat(fname)
						if err != nil { return err }
					}

					cat.names[snap - snapMin] = append(
						cat.names[snap - snapMin], fname,
					)
				}
			}
		}
	}

	return nil
}

func (cat *Catalogs) InitLGadget2(
	format string, snapMin, snapMax int64, blockMins, blockMaxes []int,
) error {
	panic("NYI")
}

///////////
// Halos //
///////////

type Halos struct {
	HaloType
	TreeType
	snapMin int
	snapOffset int
	names []string
}

func (h *Halos) HaloCatalog(snap int) string {
	return h.names[snap - h.snapMin]
}

func (h *Halos) SnapOffset() int {
	return h.snapOffset
}

func (h *Halos) InitRockstar(dir string, snapMin, snapMax int64) error {
	h.HaloType = Rockstar
	h.TreeType = ConsistentTrees

	infos, err := ioutil.ReadDir(dir)
	if err != nil { return err }

	h.snapOffset = int(snapMax) - len(infos)

	h.snapMin = int(snapMin)
	h.names = []string{}
	for i := range infos {
		h.names = append(h.names, path.Join(dir, infos[i].Name()))
	}

	if len(h.names) < int(snapMax - snapMin) + 1 {
		return fmt.Errorf(
			"There are %d files in the 'HaloDir' directory, %s, but " +
			"'SnapMin' = %d and 'SnapMax' = %d.",
			len(h.names), dir, snapMin, snapMax,
		)
	}
	h.names = h.names[len(h.names) - int(snapMax - snapMin + 1):]

	return nil
}

func (h *Halos) InitNilHalo(dir string, snapMin, snapMax int64) error {
	h.HaloType = NilHalo
	h.TreeType = NilTree

	return nil
}<|endoftext|>


<commit_before>package cmd

import (
	"github.com/perriea/tfversion/terraform"
	"github.com/spf13/cobra"
)

// installCmd represents the install command
var installCmd = &cobra.Command{
	Use:   "install",
	Short: "Install new versions or switch",
	Long:  `Install new versions or switch.`,
	PreRun: func(cmd *cobra.Command, args []string) {
		terraform.Init()
	},
	Run: func(cmd *cobra.Command, args []string) {
		terraform.Install(args[0])
	},
}

func init() {
	rootCmd.AddCommand(installCmd)
}
<commit_msg>Fixing no specific version in command install<commit_after>package cmd

import (
	"fmt"

	"github.com/perriea/tfversion/terraform"
	"github.com/spf13/cobra"
)

// installCmd represents the install command
var installCmd = &cobra.Command{
	Use:   "install",
	Short: "Install new versions or switch",
	Long:  `Install new versions or switch.`,
	PreRun: func(cmd *cobra.Command, args []string) {
		terraform.Init()
	},
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) > 0 {
			terraform.Install(args[0])
		} else {
			fmt.Printf("\033[1;31mNone version specified\n")
		}
	},
}

func init() {
	rootCmd.AddCommand(installCmd)
}
<|endoftext|>


<commit_before>// A useful example app.  You can use this to debug your tokens on the command line.
// This is also a great place to look at how you might use this library.
//
// Example usage:
// The following will create and sign a token, then verify it and output the original claims.
//     echo {\"foo\":\"bar\"} | bin/jwt -key test/sample_key -alg RS256 -sign - | bin/jwt -key test/sample_key.pub -verify -
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"regexp"

	"github.com/dgrijalva/jwt-go"
)

var (
	// Options
	flagAlg     = flag.String("alg", "", "signing algorithm identifier")
	flagKey     = flag.String("key", "", "path to key file or '-' to read from stdin")
	flagCompact = flag.Bool("compact", false, "output compact JSON")
	flagDebug   = flag.Bool("debug", false, "print out all kinds of debug data")

	// Modes - exactly one of these is required
	flagSign   = flag.String("sign", "", "path to claims object to sign or '-' to read from stdin")
	flagVerify = flag.String("verify", "", "path to JWT token to verify or '-' to read from stdin")
)

func main() {
	// Usage message if you ask for -help or if you mess up inputs.
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  One of the following flags is required: sign, verify\n")
		flag.PrintDefaults()
	}

	// Parse command line options
	flag.Parse()

	// Do the thing.  If something goes wrong, print error to stderr
	// and exit with a non-zero status code
	if err := start(); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}
}

// Figure out which thing to do and then do that
func start() error {
	if *flagSign != "" {
		return signToken()
	} else if *flagVerify != "" {
		return verifyToken()
	} else {
		flag.Usage()
		return fmt.Errorf("None of the required flags are present.  What do you want me to do?")
	}
}

// Helper func:  Read input from specified file or stdin
func loadData(p string) ([]byte, error) {
	if p == "" {
		return nil, fmt.Errorf("No path specified")
	}

	var rdr io.Reader
	if p == "-" {
		rdr = os.Stdin
	} else {
		if f, err := os.Open(p); err == nil {
			rdr = f
			defer f.Close()
		} else {
			return nil, err
		}
	}
	return ioutil.ReadAll(rdr)
}

// Print a json object in accordance with the prophecy (or the command line options)
func printJSON(j interface{}) error {
	var out []byte
	var err error

	if *flagCompact == false {
		out, err = json.MarshalIndent(j, "", "    ")
	} else {
		out, err = json.Marshal(j)
	}

	if err == nil {
		fmt.Println(string(out))
	}

	return err
}

// Verify a token and output the claims.  This is a great example
// of how to verify and view a token.
func verifyToken() error {
	// get the token
	tokData, err := loadData(*flagVerify)
	if err != nil {
		return fmt.Errorf("Couldn't read token: %v", err)
	}

	// trim possible whitespace from token
	tokData = regexp.MustCompile(`\s*$`).ReplaceAll(tokData, []byte{})
	if *flagDebug {
		fmt.Fprintf(os.Stderr, "Token len: %v bytes\n", len(tokData))
	}

	// Parse the token.  Load the key from command line option
	token, err := jwt.Parse(string(tokData), func(t *jwt.Token) (interface{}, error) {
		data, err := loadData(*flagKey)
		if err != nil {
			return nil, err
		}
		if isEs() {
			return jwt.ParseECPublicKeyFromPEM(data)
		}
		return data, nil
	})

	// Print some debug data
	if *flagDebug && token != nil {
		fmt.Fprintf(os.Stderr, "Header:\n%v\n", token.Header)
		fmt.Fprintf(os.Stderr, "Claims:\n%v\n", token.Claims)
	}

	// Print an error if we can't parse for some reason
	if err != nil {
		return fmt.Errorf("Couldn't parse token: %v", err)
	}

	// Is token invalid?
	if !token.Valid {
		return fmt.Errorf("Token is invalid")
	}

	// Print the token details
	if err := printJSON(token.Claims); err != nil {
		return fmt.Errorf("Failed to output claims: %v", err)
	}

	return nil
}

// Create, sign, and output a token.  This is a great, simple example of
// how to use this library to create and sign a token.
func signToken() error {
	// get the token data from command line arguments
	tokData, err := loadData(*flagSign)
	if err != nil {
		return fmt.Errorf("Couldn't read token: %v", err)
	} else if *flagDebug {
		fmt.Fprintf(os.Stderr, "Token: %v bytes", len(tokData))
	}

	// parse the JSON of the claims
	var claims map[string]interface{}
	if err := json.Unmarshal(tokData, &claims); err != nil {
		return fmt.Errorf("Couldn't parse claims JSON: %v", err)
	}

	// get the key
	var key interface{}
	key, err = loadData(*flagKey)
	if err != nil {
		return fmt.Errorf("Couldn't read key: %v", err)
	}

	// get the signing alg
	alg := jwt.GetSigningMethod(*flagAlg)
	if alg == nil {
		return fmt.Errorf("Couldn't find signing method: %v", *flagAlg)
	}

	// create a new token
	token := jwt.New(alg)
	token.Claims = claims

	if isEs() {
		if k, ok := key.([]byte); !ok {
			return fmt.Errorf("Couldn't convert key data to key")
		} else {
			key, err = jwt.ParseECPrivateKeyFromPEM(k)
			if err != nil {
				return err
			}
		}
	}

	if out, err := token.SignedString(key); err == nil {
		fmt.Println(out)
	} else {
		return fmt.Errorf("Error signing token: %v", err)
	}

	return nil
}

func isEs() bool {
	if alg := *flagAlg; len(alg) >= 2 && alg[0:2] == "ES" {
		return true
	}
	return false
}
<commit_msg>use cleaner version of prefix checking (thanks shurcooL)<commit_after>// A useful example app.  You can use this to debug your tokens on the command line.
// This is also a great place to look at how you might use this library.
//
// Example usage:
// The following will create and sign a token, then verify it and output the original claims.
//     echo {\"foo\":\"bar\"} | bin/jwt -key test/sample_key -alg RS256 -sign - | bin/jwt -key test/sample_key.pub -verify -
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"regexp"
	"strings"

	"github.com/dgrijalva/jwt-go"
)

var (
	// Options
	flagAlg     = flag.String("alg", "", "signing algorithm identifier")
	flagKey     = flag.String("key", "", "path to key file or '-' to read from stdin")
	flagCompact = flag.Bool("compact", false, "output compact JSON")
	flagDebug   = flag.Bool("debug", false, "print out all kinds of debug data")

	// Modes - exactly one of these is required
	flagSign   = flag.String("sign", "", "path to claims object to sign or '-' to read from stdin")
	flagVerify = flag.String("verify", "", "path to JWT token to verify or '-' to read from stdin")
)

func main() {
	// Usage message if you ask for -help or if you mess up inputs.
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  One of the following flags is required: sign, verify\n")
		flag.PrintDefaults()
	}

	// Parse command line options
	flag.Parse()

	// Do the thing.  If something goes wrong, print error to stderr
	// and exit with a non-zero status code
	if err := start(); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}
}

// Figure out which thing to do and then do that
func start() error {
	if *flagSign != "" {
		return signToken()
	} else if *flagVerify != "" {
		return verifyToken()
	} else {
		flag.Usage()
		return fmt.Errorf("None of the required flags are present.  What do you want me to do?")
	}
}

// Helper func:  Read input from specified file or stdin
func loadData(p string) ([]byte, error) {
	if p == "" {
		return nil, fmt.Errorf("No path specified")
	}

	var rdr io.Reader
	if p == "-" {
		rdr = os.Stdin
	} else {
		if f, err := os.Open(p); err == nil {
			rdr = f
			defer f.Close()
		} else {
			return nil, err
		}
	}
	return ioutil.ReadAll(rdr)
}

// Print a json object in accordance with the prophecy (or the command line options)
func printJSON(j interface{}) error {
	var out []byte
	var err error

	if *flagCompact == false {
		out, err = json.MarshalIndent(j, "", "    ")
	} else {
		out, err = json.Marshal(j)
	}

	if err == nil {
		fmt.Println(string(out))
	}

	return err
}

// Verify a token and output the claims.  This is a great example
// of how to verify and view a token.
func verifyToken() error {
	// get the token
	tokData, err := loadData(*flagVerify)
	if err != nil {
		return fmt.Errorf("Couldn't read token: %v", err)
	}

	// trim possible whitespace from token
	tokData = regexp.MustCompile(`\s*$`).ReplaceAll(tokData, []byte{})
	if *flagDebug {
		fmt.Fprintf(os.Stderr, "Token len: %v bytes\n", len(tokData))
	}

	// Parse the token.  Load the key from command line option
	token, err := jwt.Parse(string(tokData), func(t *jwt.Token) (interface{}, error) {
		data, err := loadData(*flagKey)
		if err != nil {
			return nil, err
		}
		if isEs() {
			return jwt.ParseECPublicKeyFromPEM(data)
		}
		return data, nil
	})

	// Print some debug data
	if *flagDebug && token != nil {
		fmt.Fprintf(os.Stderr, "Header:\n%v\n", token.Header)
		fmt.Fprintf(os.Stderr, "Claims:\n%v\n", token.Claims)
	}

	// Print an error if we can't parse for some reason
	if err != nil {
		return fmt.Errorf("Couldn't parse token: %v", err)
	}

	// Is token invalid?
	if !token.Valid {
		return fmt.Errorf("Token is invalid")
	}

	// Print the token details
	if err := printJSON(token.Claims); err != nil {
		return fmt.Errorf("Failed to output claims: %v", err)
	}

	return nil
}

// Create, sign, and output a token.  This is a great, simple example of
// how to use this library to create and sign a token.
func signToken() error {
	// get the token data from command line arguments
	tokData, err := loadData(*flagSign)
	if err != nil {
		return fmt.Errorf("Couldn't read token: %v", err)
	} else if *flagDebug {
		fmt.Fprintf(os.Stderr, "Token: %v bytes", len(tokData))
	}

	// parse the JSON of the claims
	var claims map[string]interface{}
	if err := json.Unmarshal(tokData, &claims); err != nil {
		return fmt.Errorf("Couldn't parse claims JSON: %v", err)
	}

	// get the key
	var key interface{}
	key, err = loadData(*flagKey)
	if err != nil {
		return fmt.Errorf("Couldn't read key: %v", err)
	}

	// get the signing alg
	alg := jwt.GetSigningMethod(*flagAlg)
	if alg == nil {
		return fmt.Errorf("Couldn't find signing method: %v", *flagAlg)
	}

	// create a new token
	token := jwt.New(alg)
	token.Claims = claims

	if isEs() {
		if k, ok := key.([]byte); !ok {
			return fmt.Errorf("Couldn't convert key data to key")
		} else {
			key, err = jwt.ParseECPrivateKeyFromPEM(k)
			if err != nil {
				return err
			}
		}
	}

	if out, err := token.SignedString(key); err == nil {
		fmt.Println(out)
	} else {
		return fmt.Errorf("Error signing token: %v", err)
	}

	return nil
}

func isEs() bool {
	return strings.HasPrefix(*flagAlg, "ES")
}
<|endoftext|>


<commit_before>package coal

import (
	"sync"

	"github.com/globalsign/mgo"
	"github.com/globalsign/mgo/bson"
)

// Event defines the event type.
type Event string

const (
	// Created is emitted when a document has been created.
	Created Event = "created"

	// Updated is emitted when a document has been updated.
	Updated Event = "updated"

	// Deleted is emitted when a document has been deleted.
	Deleted Event = "deleted"
)

// Receiver is a callback that receives stream events.
type Receiver func(Event, bson.ObjectId, Model)

// Stream simplifies the handling of change streams to receives changes to
// documents.
type Stream struct {
	store *Store
	model Model

	// Reporter is called with errors.
	Reporter func(error)

	mutex   sync.Mutex
	current *mgo.ChangeStream
	token   *bson.Raw
	closed  bool
}

// NewStream creates and returns a new stream.
func NewStream(store *Store, model Model) *Stream {
	return &Stream{
		store: store,
		model: model,
	}
}

// Open will open the stream and continuously forward events to the specified
// receiver until the stream is closed. The provided open function is called
// when the stream has been opened the first time.
func (s *Stream) Open(rec Receiver, open func()) {
	go s.open(rec, open)
}

// Close will close the stream.
func (s *Stream) Close() {
	// get mutex
	s.mutex.Lock()
	defer s.mutex.Unlock()

	// set flag
	s.closed = true

	// close active change stream
	if s.current != nil {
		_ = s.current.Close()
	}
}

func (s *Stream) open(rec Receiver, open func()) {
	// prepare once
	var once sync.Once

	// prepare opener
	opener := func() {
		once.Do(func() {
			if open != nil {
				open()
			}
		})
	}

	// run forever and call reporter with eventual errors
	for {
		// check status
		s.mutex.Lock()
		closed := s.closed
		s.mutex.Unlock()

		// return if closed
		if closed {
			return
		}

		// tail stream
		err := s.tail(rec, opener)
		if err != nil {
			if s.Reporter != nil {
				s.Reporter(err)
			}
		}
	}
}

func (s *Stream) tail(rec Receiver, open func()) error {
	// copy store
	store := s.store.Copy()
	defer store.Close()

	// open change stream
	cs, err := store.C(s.model).Watch([]bson.M{}, mgo.ChangeStreamOptions{
		FullDocument: mgo.UpdateLookup,
		ResumeAfter:  s.token,
	})
	if err != nil {
		return err
	}

	// ensure stream is closed
	defer cs.Close()

	// save reference and get status
	s.mutex.Lock()
	closed := s.closed
	if !closed {
		s.current = cs
	}
	s.mutex.Unlock()

	// return if closed
	if closed {
		return nil
	}

	// signal open
	open()

	// iterate on elements forever
	var ch change
	for cs.Next(&ch) {
		// prepare type
		var typ Event

		// parse operation type
		if ch.OperationType == "insert" {
			typ = Created
		} else if ch.OperationType == "replace" || ch.OperationType == "update" {
			typ = Updated
		} else if ch.OperationType == "delete" {
			typ = Deleted
		} else {
			continue
		}

		// prepare record
		var record Model

		// unmarshal document for created and updated events
		if typ != Deleted {
			// unmarshal record
			record = s.model.Meta().Make()
			err = ch.FullDocument.Unmarshal(record)
			if err != nil {
				return err
			}

			// init record
			Init(record)
		}

		// call receiver
		rec(typ, ch.DocumentKey.ID, record)
	}

	// close stream and check error
	err = cs.Close()
	if err != nil {
		return err
	}

	// unset reference
	s.mutex.Lock()
	s.current = nil
	s.mutex.Unlock()

	// save token
	s.token = cs.ResumeToken()

	return nil
}

type change struct {
	OperationType string `bson:"operationType"`
	DocumentKey   struct {
		ID bson.ObjectId `bson:"_id"`
	} `bson:"documentKey"`
	FullDocument bson.Raw `bson:"fullDocument"`
}
<commit_msg>set resume token after yielding to receiver<commit_after>package coal

import (
	"sync"

	"github.com/globalsign/mgo"
	"github.com/globalsign/mgo/bson"
)

// Event defines the event type.
type Event string

const (
	// Created is emitted when a document has been created.
	Created Event = "created"

	// Updated is emitted when a document has been updated.
	Updated Event = "updated"

	// Deleted is emitted when a document has been deleted.
	Deleted Event = "deleted"
)

// Receiver is a callback that receives stream events.
type Receiver func(Event, bson.ObjectId, Model)

// Stream simplifies the handling of change streams to receives changes to
// documents.
type Stream struct {
	store *Store
	model Model

	// Reporter is called with errors.
	Reporter func(error)

	mutex   sync.Mutex
	current *mgo.ChangeStream
	token   *bson.Raw
	closed  bool
}

// NewStream creates and returns a new stream.
func NewStream(store *Store, model Model) *Stream {
	return &Stream{
		store: store,
		model: model,
	}
}

// Open will open the stream and continuously forward events to the specified
// receiver until the stream is closed. The provided open function is called
// when the stream has been opened the first time.
func (s *Stream) Open(rec Receiver, open func()) {
	go s.open(rec, open)
}

// Close will close the stream.
func (s *Stream) Close() {
	// get mutex
	s.mutex.Lock()
	defer s.mutex.Unlock()

	// set flag
	s.closed = true

	// close active change stream
	if s.current != nil {
		_ = s.current.Close()
	}
}

func (s *Stream) open(rec Receiver, open func()) {
	// prepare once
	var once sync.Once

	// prepare opener
	opener := func() {
		once.Do(func() {
			if open != nil {
				open()
			}
		})
	}

	// run forever and call reporter with eventual errors
	for {
		// check status
		s.mutex.Lock()
		closed := s.closed
		s.mutex.Unlock()

		// return if closed
		if closed {
			return
		}

		// tail stream
		err := s.tail(rec, opener)
		if err != nil {
			if s.Reporter != nil {
				s.Reporter(err)
			}
		}
	}
}

func (s *Stream) tail(rec Receiver, open func()) error {
	// copy store
	store := s.store.Copy()
	defer store.Close()

	// open change stream
	cs, err := store.C(s.model).Watch([]bson.M{}, mgo.ChangeStreamOptions{
		FullDocument: mgo.UpdateLookup,
		ResumeAfter:  s.token,
	})
	if err != nil {
		return err
	}

	// ensure stream is closed
	defer cs.Close()

	// save reference and get status
	s.mutex.Lock()
	closed := s.closed
	if !closed {
		s.current = cs
	}
	s.mutex.Unlock()

	// return if closed
	if closed {
		return nil
	}

	// signal open
	open()

	// iterate on elements forever
	var ch change
	for cs.Next(&ch) {
		// prepare type
		var typ Event

		// parse operation type
		if ch.OperationType == "insert" {
			typ = Created
		} else if ch.OperationType == "replace" || ch.OperationType == "update" {
			typ = Updated
		} else if ch.OperationType == "delete" {
			typ = Deleted
		} else {
			continue
		}

		// prepare record
		var record Model

		// unmarshal document for created and updated events
		if typ != Deleted {
			// unmarshal record
			record = s.model.Meta().Make()
			err = ch.FullDocument.Unmarshal(record)
			if err != nil {
				return err
			}

			// init record
			Init(record)
		}

		// call receiver
		rec(typ, ch.DocumentKey.ID, record)

		// save token
		s.token = &ch.ResumeToken
	}

	// close stream and check error
	err = cs.Close()
	if err != nil {
		return err
	}

	// unset reference
	s.mutex.Lock()
	s.current = nil
	s.mutex.Unlock()

	return nil
}

type change struct {
	ResumeToken bson.Raw `bson:"_id"`
	OperationType string `bson:"operationType"`
	DocumentKey   struct {
		ID bson.ObjectId `bson:"_id"`
	} `bson:"documentKey"`
	FullDocument bson.Raw `bson:"fullDocument"`
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/hex"
	"fmt"
	"log"

	"github.com/dt/thile/gen"
	"github.com/foursquare/gohfile"
)

type Collection struct {
	cfg    *hfile.CollectionConfig
	reader *hfile.Reader
}

type CollectionSet struct {
	settings    *Settings
	collections map[string]Collection
}

func LoadCollections(settings *Settings, collections []hfile.CollectionConfig) (*CollectionSet, error) {
	cs := new(CollectionSet)
	cs.settings = settings
	cs.collections = make(map[string]Collection)

	for _, cfg := range collections {
		reader, err := hfile.NewReaderFromConfig(&cfg, settings.debug)
		if err != nil {
			return nil, err
		}

		cs.collections[cfg.Name] = Collection{&cfg, reader}
	}

	return cs, nil
}

func (cs *CollectionSet) readerFor(name string) (*hfile.Reader, error) {
	c, ok := cs.collections[name]
	if !ok {
		return nil, fmt.Errorf("not configured with reader for collection %s", name)
	}
	return c.reader, nil
}

func (cs *CollectionSet) scannerFor(c string) (*hfile.Scanner, error) {
	reader, err := cs.readerFor(c)
	if err != nil {
		return nil, err
	}
	s := hfile.NewScanner(reader)
	return s, nil
}

func (cs *CollectionSet) GetValuesSingle(req *gen.SingleHFileKeyRequest) (r *gen.SingleHFileKeyResponse, err error) {
	if cs.settings.debug {
		log.Printf("[GetValuesSingle] %s (%d keys)\n", *req.HfileName, len(req.SortedKeys))
	}
	reader, err := cs.scannerFor(*req.HfileName)
	if err != nil {
		return nil, err
	}

	if req.PerKeyValueLimit != nil {
		log.Println("[GetValuesSingle] PerKeyValueLimit. oh well.")
	}

	if req.CountOnly != nil {
		log.Println("[GetValuesSingle] CountOnly. oh well.")
	}

	res := new(gen.SingleHFileKeyResponse)
	res.Values = make(map[int32][]byte)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		if cs.settings.debug {
			log.Printf("[GetValuesSingle] key: %s\n", hex.EncodeToString(key))
		}
		value, err, ok := reader.GetFirst(key)
		if err != nil {
			return nil, err
		}
		if ok {
			found++
			res.Values[int32(idx)] = value
		}
	}

	if cs.settings.debug {
		log.Printf("[GetValuesSingle] %s found %d of %d.\n", *req.HfileName, found, len(req.SortedKeys))
	}
	res.KeyCount = &found
	return res, nil
}

func (cs *CollectionSet) GetValuesMulti(req *gen.SingleHFileKeyRequest) (r *gen.MultiHFileKeyResponse, err error) {
	log.Println("[GetValuesMulti]", len(req.SortedKeys))
	reader, err := cs.scannerFor(*req.HfileName)
	if err != nil {
		return nil, err
	}

	//out := make(map[int32][]byte)
	res := new(gen.MultiHFileKeyResponse)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		values, err := reader.GetAll(key)
		if err != nil {
			return nil, err
		}
		if len(values) > 0 {
			found += int32(len(values))
			res.Values[int32(idx)] = values
		}
	}

	res.KeyCount = &found
	return res, nil

}

func (cs *CollectionSet) GetValuesForPrefixes(req *gen.PrefixRequest) (r *gen.PrefixResponse, err error) {
	res := new(gen.PrefixResponse)
	if reader, err := cs.readerFor(*req.HfileName); err != nil {
		return nil, err
	} else {
		i := reader.NewIterator()
		if res.Values, err = i.AllForPrfixes(req.SortedKeys); err != nil {
			return nil, err
		} else {
			return res, nil
		}
	}
}

func (cs *CollectionSet) GetValuesMultiSplitKeys(req *gen.MultiHFileSplitKeyRequest) (r *gen.KeyToValuesResponse, err error) {
	return nil, nil
}

func (cs *CollectionSet) GetIterator(req *gen.IteratorRequest) (r *gen.IteratorResponse, err error) {
	return nil, nil
}

func (cs *CollectionSet) GetInfo(req *gen.InfoRequest) (r []*gen.HFileInfo, err error) {
	return nil, nil
}

func (cs *CollectionSet) TestTimeout(waitInMillis int32) (r int32, err error) {
	return 0, nil
}
<commit_msg>add split-key support<commit_after>package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"log"

	"github.com/dt/thile/gen"
	"github.com/foursquare/gohfile"
)

type Collection struct {
	cfg    *hfile.CollectionConfig
	reader *hfile.Reader
}

type CollectionSet struct {
	settings    *Settings
	collections map[string]Collection
}

func LoadCollections(settings *Settings, collections []hfile.CollectionConfig) (*CollectionSet, error) {
	cs := new(CollectionSet)
	cs.settings = settings
	cs.collections = make(map[string]Collection)

	for _, cfg := range collections {
		reader, err := hfile.NewReaderFromConfig(&cfg, settings.debug)
		if err != nil {
			return nil, err
		}

		cs.collections[cfg.Name] = Collection{&cfg, reader}
	}

	return cs, nil
}

func (cs *CollectionSet) readerFor(name string) (*hfile.Reader, error) {
	c, ok := cs.collections[name]
	if !ok {
		return nil, fmt.Errorf("not configured with reader for collection %s", name)
	}
	return c.reader, nil
}

func (cs *CollectionSet) scannerFor(c string) (*hfile.Scanner, error) {
	reader, err := cs.readerFor(c)
	if err != nil {
		return nil, err
	}
	s := hfile.NewScanner(reader)
	return s, nil
}

func (cs *CollectionSet) GetValuesSingle(req *gen.SingleHFileKeyRequest) (r *gen.SingleHFileKeyResponse, err error) {
	if cs.settings.debug {
		log.Printf("[GetValuesSingle] %s (%d keys)\n", *req.HfileName, len(req.SortedKeys))
	}
	reader, err := cs.scannerFor(*req.HfileName)
	if err != nil {
		return nil, err
	}

	if req.PerKeyValueLimit != nil {
		log.Println("[GetValuesSingle] PerKeyValueLimit. oh well.")
	}

	if req.CountOnly != nil {
		log.Println("[GetValuesSingle] CountOnly. oh well.")
	}

	res := new(gen.SingleHFileKeyResponse)
	res.Values = make(map[int32][]byte)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		if cs.settings.debug {
			log.Printf("[GetValuesSingle] key: %s\n", hex.EncodeToString(key))
		}
		value, err, ok := reader.GetFirst(key)
		if err != nil {
			return nil, err
		}
		if ok {
			found++
			res.Values[int32(idx)] = value
		}
	}

	if cs.settings.debug {
		log.Printf("[GetValuesSingle] %s found %d of %d.\n", *req.HfileName, found, len(req.SortedKeys))
	}
	res.KeyCount = &found
	return res, nil
}

func (cs *CollectionSet) GetValuesMulti(req *gen.SingleHFileKeyRequest) (r *gen.MultiHFileKeyResponse, err error) {
	log.Println("[GetValuesMulti]", len(req.SortedKeys))
	reader, err := cs.scannerFor(*req.HfileName)
	if err != nil {
		return nil, err
	}

	//out := make(map[int32][]byte)
	res := new(gen.MultiHFileKeyResponse)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		values, err := reader.GetAll(key)
		if err != nil {
			return nil, err
		}
		if len(values) > 0 {
			found += int32(len(values))
			res.Values[int32(idx)] = values
		}
	}

	res.KeyCount = &found
	return res, nil

}

func (cs *CollectionSet) GetValuesForPrefixes(req *gen.PrefixRequest) (r *gen.PrefixResponse, err error) {
	res := new(gen.PrefixResponse)
	if reader, err := cs.readerFor(*req.HfileName); err != nil {
		return nil, err
	} else {
		i := reader.NewIterator()
		if res.Values, err = i.AllForPrfixes(req.SortedKeys); err != nil {
			return nil, err
		} else {
			return res, nil
		}
	}
}

func (cs *CollectionSet) GetValuesMultiSplitKeys(req *gen.MultiHFileSplitKeyRequest) (r *gen.KeyToValuesResponse, err error) {
	res := make(map[string][][]byte)
	scanner, err := cs.scannerFor(*req.HfileName)
	if err != nil {
		return nil, err
	}

	for _, parts := range RevProduct(req.SplitKey) {
		// TODO(davidt): avoid allocing concated key by adding split-key search lower down.
		key := bytes.Join(parts, nil)

		if values, err := scanner.GetAll(key); err != nil {
			return nil, err
		} else if len(values) > 0 {
			res[string(key)] = values
		}
	}
	return &gen.KeyToValuesResponse{res}, nil
}

func (cs *CollectionSet) GetIterator(req *gen.IteratorRequest) (r *gen.IteratorResponse, err error) {
	return nil, nil
}

func (cs *CollectionSet) GetInfo(req *gen.InfoRequest) (r []*gen.HFileInfo, err error) {
	return nil, nil
}

func (cs *CollectionSet) TestTimeout(waitInMillis int32) (r int32, err error) {
	return 0, nil
}
<|endoftext|>


<commit_before>package main

import (
	log "code.google.com/p/log4go"
	"sync"
	"time"
)

const (
	timerFormat      = "2006-01-02 15:04:05"
	infiniteDuration = time.Duration(1<<63 - 1)
)

var (
	timerLazyDelay = 300 * time.Millisecond
)

type TimerData struct {
	Key    string
	expire time.Time
	fn     func()
	index  int
	next   *TimerData
}

func (td *TimerData) Delay() time.Duration {
	return td.expire.Sub(time.Now())
}

func (td *TimerData) ExpireString() string {
	return td.expire.Format(timerFormat)
}

type Timer struct {
	lock   sync.Mutex
	free   *TimerData
	timers []*TimerData
	signal *time.Timer
}

// A heap must be initialized before any of the heap operations
// can be used. Init is idempotent with respect to the heap invariants
// and may be called whenever the heap invariants may have been invalidated.
// Its complexity is O(n) where n = h.Len().
//
func NewTimer(num int) (t *Timer) {
	t = new(Timer)
	t.init(num)
	return t
}

// Init init the timer.
func (t *Timer) Init(num int) {
	t.init(num)
}

func (t *Timer) init(num int) {
	t.signal = time.NewTimer(infiniteDuration)
	t.timers = make([]*TimerData, 0, num)
	tds := make([]TimerData, num)
	t.free = &(tds[0])
	td := t.free
	for i := 1; i < num; i++ {
		td.next = &(tds[i])
		td = td.next
	}
	go t.start()
}

// get get a free timer data.
func (t *Timer) get() (td *TimerData) {
	var (
		i   int
		num = len(t.timers)
		tds []TimerData
	)
	if td = t.free; td == nil {
		tds = make([]TimerData, num)
		t.free = &(tds[0])
		td = t.free
		for i = 1; i < num; i++ {
			td.next = &(tds[i])
			td = td.next
		}
		td = t.free
	}
	t.free = td.next
	return
}

// put put back a timer data.
func (t *Timer) put(td *TimerData) {
	td.next = t.free
	t.free = td
}

// Push pushes the element x onto the heap. The complexity is
// O(log(n)) where n = h.Len().
func (t *Timer) Add(expire time.Duration, fn func()) (td *TimerData) {
	t.lock.Lock()
	td = t.get()
	td.expire = time.Now().Add(expire)
	td.fn = fn
	t.add(td)
	t.lock.Unlock()
	return
}

// Push pushes the element x onto the heap. The complexity is
// O(log(n)) where n = h.Len().
func (t *Timer) add(td *TimerData) {
	var d time.Duration
	td.index = len(t.timers)
	// add to the minheap last node
	t.timers = append(t.timers, td)
	t.up(td.index)
	if td.index == 0 {
		// if first node, signal start goroutine
		d = td.Delay()
		t.signal.Reset(d)
		if Debug {
			log.Debug("timer: add reset delay %d ms", int64(d)/int64(time.Millisecond))
		}
	}
	if Debug {
		log.Debug("timer: push item key: %s, expire: %s, index: %d", td.Key, td.ExpireString(), td.index)
	}
	return
}

// Del removes the element at index i from the heap.
// The complexity is O(log(n)) where n = h.Len().
func (t *Timer) Del(td *TimerData) {
	t.lock.Lock()
	if t.del(td) {
		t.put(td)
	}
	t.lock.Unlock()
	return
}

func (t *Timer) del(td *TimerData) bool {
	var (
		i    = td.index
		last = len(t.timers) - 1
	)
	if i < 0 || i > last || t.timers[i] != td {
		// already remove, usually by expire
		if Debug {
			log.Debug("timer del i: %d, last: %d, %p,%p", i, last, t.timers[i], td)
		}
		return false
	}
	if i != last {
		t.swap(i, last)
		t.down(i, last)
		t.up(i)
	}
	// remove item is the last node
	t.timers[last].index = -1 // for safety
	t.timers = t.timers[:last]
	if Debug {
		log.Debug("timer: remove item key: %s, expire: %s, index: %d", td.Key, td.ExpireString(), td.index)
	}
	return true
}

// Set update timer data.
func (t *Timer) Set(td *TimerData, expire time.Duration) {
	t.lock.Lock()
	t.del(td)
	td.expire = time.Now().Add(expire)
	t.add(td)
	t.lock.Unlock()
	return
}

// start start the timer.
func (t *Timer) start() {
	for {
		t.expire()
		<-t.signal.C
	}
}

// expire removes the minimum element (according to Less) from the heap.
// The complexity is O(log(n)) where n = max.
// It is equivalent to Del(0).
func (t *Timer) expire() {
	var (
		td *TimerData
		d  time.Duration
	)
	t.lock.Lock()
	for {
		if len(t.timers) == 0 {
			d = infiniteDuration
			if Debug {
				log.Debug("timer: no other instance")
			}
			break
		}
		td = t.timers[0]
		if d = td.Delay(); d > 0 {
			break
		}
		if td.fn == nil {
			log.Warn("expire timer no fn")
		} else {
			if Debug {
				log.Debug("timer key: %s, expire: %s, index: %d expired, call fn", td.Key, td.ExpireString(), td.index)
			}
			td.fn()
		}
		// let caller put back
		t.del(td)
	}
	t.signal.Reset(d)
	if Debug {
		log.Debug("timer: expier reset delay %d ms", int64(d)/int64(time.Millisecond))
	}
	t.lock.Unlock()
	return
}

func (t *Timer) up(j int) {
	for {
		i := (j - 1) / 2 // parent
		if i == j || !t.less(j, i) {
			break
		}
		t.swap(i, j)
		j = i
	}
}

func (t *Timer) down(i, n int) {
	for {
		j1 := 2*i + 1
		if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
			break
		}
		j := j1 // left child
		if j2 := j1 + 1; j2 < n && !t.less(j1, j2) {
			j = j2 // = 2*i + 2  // right child
		}
		if !t.less(j, i) {
			break
		}
		t.swap(i, j)
		i = j
	}
}

func (t *Timer) less(i, j int) bool {
	return t.timers[i].expire.Before(t.timers[j].expire)
}

func (t *Timer) swap(i, j int) {
	//log.Debug("swap(%d, %d)", i, j)
	t.timers[i], t.timers[j] = t.timers[j], t.timers[i]
	t.timers[i].index = i
	t.timers[j].index = j
}
<commit_msg>fix timer bug<commit_after>package main

import (
	log "code.google.com/p/log4go"
	"sync"
	"time"
)

const (
	timerFormat      = "2006-01-02 15:04:05"
	infiniteDuration = time.Duration(1<<63 - 1)
)

var (
	timerLazyDelay = 300 * time.Millisecond
)

type TimerData struct {
	Key    string
	expire time.Time
	fn     func()
	index  int
	next   *TimerData
}

func (td *TimerData) Delay() time.Duration {
	return td.expire.Sub(time.Now())
}

func (td *TimerData) ExpireString() string {
	return td.expire.Format(timerFormat)
}

type Timer struct {
	lock   sync.Mutex
	free   *TimerData
	timers []*TimerData
	signal *time.Timer
	num    int
}

// A heap must be initialized before any of the heap operations
// can be used. Init is idempotent with respect to the heap invariants
// and may be called whenever the heap invariants may have been invalidated.
// Its complexity is O(n) where n = h.Len().
//
func NewTimer(num int) (t *Timer) {
	t = new(Timer)
	t.init(num)
	return t
}

// Init init the timer.
func (t *Timer) Init(num int) {
	t.init(num)
}

func (t *Timer) init(num int) {
	t.signal = time.NewTimer(infiniteDuration)
	t.timers = make([]*TimerData, 0, num)
	tds := make([]TimerData, num)
	t.free = &(tds[0])
	td := t.free
	for i := 1; i < num; i++ {
		td.next = &(tds[i])
		td = td.next
	}
	t.num = num
	go t.start()
}

// get get a free timer data.
func (t *Timer) get() (td *TimerData) {
	var (
		i   int
		tds []TimerData
	)
	if td = t.free; td == nil {
		tds = make([]TimerData, t.num)
		t.free = &(tds[0])
		td = t.free
		for i = 1; i < t.num; i++ {
			td.next = &(tds[i])
			td = td.next
		}
		td = t.free
	}
	t.free = td.next
	return
}

// put put back a timer data.
func (t *Timer) put(td *TimerData) {
	td.next = t.free
	t.free = td
}

// Push pushes the element x onto the heap. The complexity is
// O(log(n)) where n = h.Len().
func (t *Timer) Add(expire time.Duration, fn func()) (td *TimerData) {
	t.lock.Lock()
	td = t.get()
	td.expire = time.Now().Add(expire)
	td.fn = fn
	t.add(td)
	t.lock.Unlock()
	return
}

// Del removes the element at index i from the heap.
// The complexity is O(log(n)) where n = h.Len().
func (t *Timer) Del(td *TimerData) {
	t.lock.Lock()
	t.del(td)
	t.put(td)
	t.lock.Unlock()
	return
}

// Push pushes the element x onto the heap. The complexity is
// O(log(n)) where n = h.Len().
func (t *Timer) add(td *TimerData) {
	var d time.Duration
	td.index = len(t.timers)
	// add to the minheap last node
	t.timers = append(t.timers, td)
	t.up(td.index)
	if td.index == 0 {
		// if first node, signal start goroutine
		d = td.Delay()
		t.signal.Reset(d)
		if Debug {
			log.Debug("timer: add reset delay %d ms", int64(d)/int64(time.Millisecond))
		}
	}
	if Debug {
		log.Debug("timer: push item key: %s, expire: %s, index: %d", td.Key, td.ExpireString(), td.index)
	}
	return
}

func (t *Timer) del(td *TimerData) {
	var (
		i    = td.index
		last = len(t.timers) - 1
	)
	if i < 0 || i > last || t.timers[i] != td {
		// already remove, usually by expire
		if Debug {
			log.Debug("timer del i: %d, last: %d, %p,%p", i, last, t.timers[i], td)
		}
		return
	}
	if i != last {
		t.swap(i, last)
		t.down(i, last)
		t.up(i)
	}
	// remove item is the last node
	t.timers[last].index = -1 // for safety
	t.timers = t.timers[:last]
	if Debug {
		log.Debug("timer: remove item key: %s, expire: %s, index: %d", td.Key, td.ExpireString(), td.index)
	}
	return
}

// Set update timer data.
func (t *Timer) Set(td *TimerData, expire time.Duration) {
	t.lock.Lock()
	t.del(td)
	td.expire = time.Now().Add(expire)
	t.add(td)
	t.lock.Unlock()
	return
}

// start start the timer.
func (t *Timer) start() {
	for {
		t.expire()
		<-t.signal.C
	}
}

// expire removes the minimum element (according to Less) from the heap.
// The complexity is O(log(n)) where n = max.
// It is equivalent to Del(0).
func (t *Timer) expire() {
	var (
		td *TimerData
		d  time.Duration
	)
	t.lock.Lock()
	for {
		if len(t.timers) == 0 {
			d = infiniteDuration
			if Debug {
				log.Debug("timer: no other instance")
			}
			break
		}
		td = t.timers[0]
		if d = td.Delay(); d > 0 {
			break
		}
		if td.fn == nil {
			log.Warn("expire timer no fn")
		} else {
			if Debug {
				log.Debug("timer key: %s, expire: %s, index: %d expired, call fn", td.Key, td.ExpireString(), td.index)
			}
			td.fn()
		}
		// let caller put back
		t.del(td)
	}
	t.signal.Reset(d)
	if Debug {
		log.Debug("timer: expier reset delay %d ms", int64(d)/int64(time.Millisecond))
	}
	t.lock.Unlock()
	return
}

func (t *Timer) up(j int) {
	for {
		i := (j - 1) / 2 // parent
		if i == j || !t.less(j, i) {
			break
		}
		t.swap(i, j)
		j = i
	}
}

func (t *Timer) down(i, n int) {
	for {
		j1 := 2*i + 1
		if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
			break
		}
		j := j1 // left child
		if j2 := j1 + 1; j2 < n && !t.less(j1, j2) {
			j = j2 // = 2*i + 2  // right child
		}
		if !t.less(j, i) {
			break
		}
		t.swap(i, j)
		i = j
	}
}

func (t *Timer) less(i, j int) bool {
	return t.timers[i].expire.Before(t.timers[j].expire)
}

func (t *Timer) swap(i, j int) {
	t.timers[i], t.timers[j] = t.timers[j], t.timers[i]
	t.timers[i].index = i
	t.timers[j].index = j
}
<|endoftext|>


<commit_before>package services

import (
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/matrix-org/go-neb/database"
	"github.com/matrix-org/go-neb/matrix"
	"github.com/matrix-org/go-neb/polling"
	"github.com/matrix-org/go-neb/types"
	"github.com/mmcdole/gofeed"
	"html"
	"time"
)

const minPollingIntervalSeconds = 60 // 1 min (News feeds can be genuinely spammy)

type feedPoller struct{}

func (p *feedPoller) IntervalSecs() int64 { return 10 }
func (p *feedPoller) OnPoll(s types.Service, cli *matrix.Client) {
	logger := log.WithFields(log.Fields{
		"service_id":   s.ServiceID(),
		"service_type": s.ServiceType(),
	})

	frService, ok := s.(*rssBotService)
	if !ok {
		logger.Error("RSSBot: OnPoll called without a Feed Service instance")
		return
	}
	now := time.Now().Unix() // Second resolution

	// Work out which feeds should be polled
	var pollFeeds []string
	for u, feedInfo := range frService.Feeds {
		if feedInfo.NextPollTimestampSecs == 0 || now >= feedInfo.NextPollTimestampSecs {
			// re-query this feed
			pollFeeds = append(pollFeeds, u)
		}
	}

	if len(pollFeeds) == 0 {
		return
	}

	// Query each feed and send new items to subscribed rooms
	for _, u := range pollFeeds {
		feed, items, err := p.queryFeed(frService, u)
		if err != nil {
			logger.WithField("feed_url", u).WithError(err).Error("Failed to query feed")
			continue
		}
		// Loop backwards since [0] is the most recent and we want to send in chronological order
		for i := len(items) - 1; i >= 0; i-- {
			item := items[i]
			if err := p.sendToRooms(frService, cli, u, feed, item); err != nil {
				logger.WithFields(log.Fields{
					"feed_url":   u,
					log.ErrorKey: err,
					"item":       item,
				}).Error("Failed to send item to room")
			}
		}
	}

	// Persist the service to save the next poll times
	if _, err := database.GetServiceDB().StoreService(frService); err != nil {
		logger.WithError(err).Error("Failed to persist next poll times for service")
	}
}

// Query the given feed, update relevant timestamps and return NEW items
func (p *feedPoller) queryFeed(s *rssBotService, feedURL string) (*gofeed.Feed, []gofeed.Item, error) {
	log.WithField("feed_url", feedURL).Info("Querying feed")
	var items []gofeed.Item
	fp := gofeed.NewParser()
	feed, err := fp.ParseURL(feedURL)
	if err != nil {
		return nil, items, err
	}

	// Work out which items are new, if any (based on the last updated TS we have)
	// If the TS is 0 then this is the first ever poll, so let's not send 10s of events
	// into the room and just do new ones from this point onwards.
	if s.Feeds[feedURL].FeedUpdatedTimestampSecs != 0 {
		for _, i := range feed.Items {
			if i == nil || i.PublishedParsed == nil {
				continue
			}
			if i.PublishedParsed.Unix() > s.Feeds[feedURL].FeedUpdatedTimestampSecs {
				items = append(items, *i)
			}
		}
	}

	now := time.Now().Unix() // Second resolution

	// Work out when this feed was last updated
	var feedLastUpdatedTs int64
	if feed.UpdatedParsed != nil {
		feedLastUpdatedTs = feed.UpdatedParsed.Unix()
	} else if len(feed.Items) > 0 {
		i := feed.Items[0]
		if i != nil && i.PublishedParsed != nil {
			feedLastUpdatedTs = i.PublishedParsed.Unix()
		}
	}

	// Work out when to next poll this feed
	nextPollTsSec := now + minPollingIntervalSeconds
	if s.Feeds[feedURL].PollIntervalMins > int(minPollingIntervalSeconds/60) {
		nextPollTsSec = now + int64(s.Feeds[feedURL].PollIntervalMins*60)
	}
	// TODO: Handle the 'sy' Syndication extension to control update interval.
	// See http://www.feedforall.com/syndication.htm and http://web.resource.org/rss/1.0/modules/syndication/

	p.updateFeedInfo(s, feedURL, nextPollTsSec, feedLastUpdatedTs)
	return feed, items, nil
}

func (p *feedPoller) updateFeedInfo(s *rssBotService, feedURL string, nextPollTs, feedUpdatedTs int64) {
	for u := range s.Feeds {
		if u != feedURL {
			continue
		}
		f := s.Feeds[u]
		f.NextPollTimestampSecs = nextPollTs
		f.FeedUpdatedTimestampSecs = feedUpdatedTs
		s.Feeds[u] = f
	}
}

func (p *feedPoller) sendToRooms(s *rssBotService, cli *matrix.Client, feedURL string, feed *gofeed.Feed, item gofeed.Item) error {
	logger := log.WithField("feed_url", feedURL).WithField("title", item.Title)
	logger.Info("New feed item")
	for _, roomID := range s.Feeds[feedURL].Rooms {
		if _, err := cli.SendMessageEvent(roomID, "m.room.message", itemToHTML(feed, item)); err != nil {
			logger.WithError(err).WithField("room_id", roomID).Error("Failed to send to room")
		}
	}
	return nil
}

// SomeOne posted a new article: Title Of The Entry ( https://someurl.com/blag )
func itemToHTML(feed *gofeed.Feed, item gofeed.Item) matrix.HTMLMessage {
	return matrix.GetHTMLMessage("m.notice", fmt.Sprintf(
		"%s posted a new article: %s ( %s )",
		html.EscapeString(feed.Title), html.EscapeString(item.Title), html.EscapeString(item.Link),
	))
}

type rssBotService struct {
	types.DefaultService
	id            string
	serviceUserID string
	Feeds         map[string]struct { // feed_url => { }
		PollIntervalMins         int      `json:"poll_interval_mins"`
		Rooms                    []string `json:"rooms"`
		NextPollTimestampSecs    int64    // Internal: When we should poll again
		FeedUpdatedTimestampSecs int64    // Internal: The last time the feed was updated
	} `json:"feeds"`
}

func (s *rssBotService) ServiceUserID() string { return s.serviceUserID }
func (s *rssBotService) ServiceID() string     { return s.id }
func (s *rssBotService) ServiceType() string   { return "rssbot" }
func (s *rssBotService) Poller() types.Poller  { return &feedPoller{} }

// Register will check the liveness of each RSS feed given. If all feeds check out okay, no error is returned.
func (s *rssBotService) Register(oldService types.Service, client *matrix.Client) error {
	if len(s.Feeds) == 0 {
		// this is an error UNLESS the old service had some feeds in which case they are deleting us :(
		var numOldFeeds int
		oldFeedService, ok := oldService.(*rssBotService)
		if !ok {
			log.WithField("service_id", oldService.ServiceID()).Error("Old service isn't a rssBotService")
		} else {
			numOldFeeds = len(oldFeedService.Feeds)
		}
		if numOldFeeds == 0 {
			return errors.New("An RSS feed must be specified.")
		}
		return nil
	}
	// Make sure we can parse the feed
	for feedURL, feedInfo := range s.Feeds {
		fp := gofeed.NewParser()
		if _, err := fp.ParseURL(feedURL); err != nil {
			return fmt.Errorf("Failed to read URL %s: %s", feedURL, err.Error())
		}
		if len(feedInfo.Rooms) == 0 {
			return fmt.Errorf("Feed %s has no rooms to send updates to", feedURL)
		}
	}
	return nil
}

func (s *rssBotService) PostRegister(oldService types.Service) {
	if len(s.Feeds) == 0 { // bye-bye :(
		logger := log.WithFields(log.Fields{
			"service_id":   s.ServiceID(),
			"service_type": s.ServiceType(),
		})
		logger.Info("Deleting service: No feeds remaining.")
		polling.StopPolling(s)
		if err := database.GetServiceDB().DeleteService(s.ServiceID()); err != nil {
			logger.WithError(err).Error("Failed to delete service")
		}
	}
}

func init() {
	types.RegisterService(func(serviceID, serviceUserID, webhookEndpointURL string) types.Service {
		r := &rssBotService{
			id:            serviceID,
			serviceUserID: serviceUserID,
		}
		return r
	})
}
<commit_msg>Fix NPE if oldService is nil<commit_after>package services

import (
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/matrix-org/go-neb/database"
	"github.com/matrix-org/go-neb/matrix"
	"github.com/matrix-org/go-neb/polling"
	"github.com/matrix-org/go-neb/types"
	"github.com/mmcdole/gofeed"
	"html"
	"time"
)

const minPollingIntervalSeconds = 60 // 1 min (News feeds can be genuinely spammy)

type feedPoller struct{}

func (p *feedPoller) IntervalSecs() int64 { return 10 }
func (p *feedPoller) OnPoll(s types.Service, cli *matrix.Client) {
	logger := log.WithFields(log.Fields{
		"service_id":   s.ServiceID(),
		"service_type": s.ServiceType(),
	})

	frService, ok := s.(*rssBotService)
	if !ok {
		logger.Error("RSSBot: OnPoll called without a Feed Service instance")
		return
	}
	now := time.Now().Unix() // Second resolution

	// Work out which feeds should be polled
	var pollFeeds []string
	for u, feedInfo := range frService.Feeds {
		if feedInfo.NextPollTimestampSecs == 0 || now >= feedInfo.NextPollTimestampSecs {
			// re-query this feed
			pollFeeds = append(pollFeeds, u)
		}
	}

	if len(pollFeeds) == 0 {
		return
	}

	// Query each feed and send new items to subscribed rooms
	for _, u := range pollFeeds {
		feed, items, err := p.queryFeed(frService, u)
		if err != nil {
			logger.WithField("feed_url", u).WithError(err).Error("Failed to query feed")
			continue
		}
		// Loop backwards since [0] is the most recent and we want to send in chronological order
		for i := len(items) - 1; i >= 0; i-- {
			item := items[i]
			if err := p.sendToRooms(frService, cli, u, feed, item); err != nil {
				logger.WithFields(log.Fields{
					"feed_url":   u,
					log.ErrorKey: err,
					"item":       item,
				}).Error("Failed to send item to room")
			}
		}
	}

	// Persist the service to save the next poll times
	if _, err := database.GetServiceDB().StoreService(frService); err != nil {
		logger.WithError(err).Error("Failed to persist next poll times for service")
	}
}

// Query the given feed, update relevant timestamps and return NEW items
func (p *feedPoller) queryFeed(s *rssBotService, feedURL string) (*gofeed.Feed, []gofeed.Item, error) {
	log.WithField("feed_url", feedURL).Info("Querying feed")
	var items []gofeed.Item
	fp := gofeed.NewParser()
	feed, err := fp.ParseURL(feedURL)
	if err != nil {
		return nil, items, err
	}

	// Work out which items are new, if any (based on the last updated TS we have)
	// If the TS is 0 then this is the first ever poll, so let's not send 10s of events
	// into the room and just do new ones from this point onwards.
	if s.Feeds[feedURL].FeedUpdatedTimestampSecs != 0 {
		for _, i := range feed.Items {
			if i == nil || i.PublishedParsed == nil {
				continue
			}
			if i.PublishedParsed.Unix() > s.Feeds[feedURL].FeedUpdatedTimestampSecs {
				items = append(items, *i)
			}
		}
	}

	now := time.Now().Unix() // Second resolution

	// Work out when this feed was last updated
	var feedLastUpdatedTs int64
	if feed.UpdatedParsed != nil {
		feedLastUpdatedTs = feed.UpdatedParsed.Unix()
	} else if len(feed.Items) > 0 {
		i := feed.Items[0]
		if i != nil && i.PublishedParsed != nil {
			feedLastUpdatedTs = i.PublishedParsed.Unix()
		}
	}

	// Work out when to next poll this feed
	nextPollTsSec := now + minPollingIntervalSeconds
	if s.Feeds[feedURL].PollIntervalMins > int(minPollingIntervalSeconds/60) {
		nextPollTsSec = now + int64(s.Feeds[feedURL].PollIntervalMins*60)
	}
	// TODO: Handle the 'sy' Syndication extension to control update interval.
	// See http://www.feedforall.com/syndication.htm and http://web.resource.org/rss/1.0/modules/syndication/

	p.updateFeedInfo(s, feedURL, nextPollTsSec, feedLastUpdatedTs)
	return feed, items, nil
}

func (p *feedPoller) updateFeedInfo(s *rssBotService, feedURL string, nextPollTs, feedUpdatedTs int64) {
	for u := range s.Feeds {
		if u != feedURL {
			continue
		}
		f := s.Feeds[u]
		f.NextPollTimestampSecs = nextPollTs
		f.FeedUpdatedTimestampSecs = feedUpdatedTs
		s.Feeds[u] = f
	}
}

func (p *feedPoller) sendToRooms(s *rssBotService, cli *matrix.Client, feedURL string, feed *gofeed.Feed, item gofeed.Item) error {
	logger := log.WithField("feed_url", feedURL).WithField("title", item.Title)
	logger.Info("New feed item")
	for _, roomID := range s.Feeds[feedURL].Rooms {
		if _, err := cli.SendMessageEvent(roomID, "m.room.message", itemToHTML(feed, item)); err != nil {
			logger.WithError(err).WithField("room_id", roomID).Error("Failed to send to room")
		}
	}
	return nil
}

// SomeOne posted a new article: Title Of The Entry ( https://someurl.com/blag )
func itemToHTML(feed *gofeed.Feed, item gofeed.Item) matrix.HTMLMessage {
	return matrix.GetHTMLMessage("m.notice", fmt.Sprintf(
		"%s posted a new article: %s ( %s )",
		html.EscapeString(feed.Title), html.EscapeString(item.Title), html.EscapeString(item.Link),
	))
}

type rssBotService struct {
	types.DefaultService
	id            string
	serviceUserID string
	Feeds         map[string]struct { // feed_url => { }
		PollIntervalMins         int      `json:"poll_interval_mins"`
		Rooms                    []string `json:"rooms"`
		NextPollTimestampSecs    int64    // Internal: When we should poll again
		FeedUpdatedTimestampSecs int64    // Internal: The last time the feed was updated
	} `json:"feeds"`
}

func (s *rssBotService) ServiceUserID() string { return s.serviceUserID }
func (s *rssBotService) ServiceID() string     { return s.id }
func (s *rssBotService) ServiceType() string   { return "rssbot" }
func (s *rssBotService) Poller() types.Poller  { return &feedPoller{} }

// Register will check the liveness of each RSS feed given. If all feeds check out okay, no error is returned.
func (s *rssBotService) Register(oldService types.Service, client *matrix.Client) error {
	if len(s.Feeds) == 0 {
		// this is an error UNLESS the old service had some feeds in which case they are deleting us :(
		var numOldFeeds int
		oldFeedService, ok := oldService.(*rssBotService)
		if !ok {
			log.WithField("service", oldService).Error("Old service isn't a rssBotService")
		} else {
			numOldFeeds = len(oldFeedService.Feeds)
		}
		if numOldFeeds == 0 {
			return errors.New("An RSS feed must be specified.")
		}
		return nil
	}
	// Make sure we can parse the feed
	for feedURL, feedInfo := range s.Feeds {
		fp := gofeed.NewParser()
		if _, err := fp.ParseURL(feedURL); err != nil {
			return fmt.Errorf("Failed to read URL %s: %s", feedURL, err.Error())
		}
		if len(feedInfo.Rooms) == 0 {
			return fmt.Errorf("Feed %s has no rooms to send updates to", feedURL)
		}
	}
	return nil
}

func (s *rssBotService) PostRegister(oldService types.Service) {
	if len(s.Feeds) == 0 { // bye-bye :(
		logger := log.WithFields(log.Fields{
			"service_id":   s.ServiceID(),
			"service_type": s.ServiceType(),
		})
		logger.Info("Deleting service: No feeds remaining.")
		polling.StopPolling(s)
		if err := database.GetServiceDB().DeleteService(s.ServiceID()); err != nil {
			logger.WithError(err).Error("Failed to delete service")
		}
	}
}

func init() {
	types.RegisterService(func(serviceID, serviceUserID, webhookEndpointURL string) types.Service {
		r := &rssBotService{
			id:            serviceID,
			serviceUserID: serviceUserID,
		}
		return r
	})
}
<|endoftext|>


<commit_before>package androidbinary

import (
	"bytes"
	"io"
	"reflect"
	"testing"
)

var readStringPoolTests = []struct {
	input   []uint8
	strings []string
	styles  []string
}{
	{
		[]uint8{
			0x01, 0x00, // Type = RES_STRING_POOL_TYPE
			0x1C, 0x00, // HeaderSize = 28 bytes
			0x3C, 0x00, 0x00, 0x00, // Size = 60
			0x02, 0x00, 0x00, 0x00, // StringCount = 2
			0x02, 0x00, 0x00, 0x00, // StyleScount = 2
			0x00, 0x00, 0x00, 0x00, // Flags = 0x00
			0x2C, 0x00, 0x00, 0x00, // StringStart = 44
			0x34, 0x00, 0x00, 0x00, // StylesStart = 52

			// StringIndexes
			0x00, 0x00, 0x00, 0x00,
			0x04, 0x00, 0x00, 0x00,

			// StyleIndexes
			0x00, 0x00, 0x00, 0x00,
			0x04, 0x00, 0x00, 0x00,

			// Strings
			0x01, 0x00, 0x61, 0x00,
			0x01, 0x00, 0x42, 0x30,

			// Styles
			0x01, 0x00, 0x63, 0x00,
			0x01, 0x00, 0x43, 0x30,
		},
		[]string{"a", "\3042"},
		[]string{"b", "\3043"},
	},
}

func TestReadStringPool(t *testing.T) {
	for _, tt := range readStringPoolTests {
		buf := bytes.NewReader(tt.input)
		sr := io.NewSectionReader(buf, 0, int64(len(tt.input)))
		actual, err := readStringPool(sr)
		if err != nil {
			t.Errorf("got %v want no error", err)
		}
		if reflect.DeepEqual(actual.Strings, tt.strings) {
			t.Errorf("got %v want %v", actual.Strings, tt.strings)
		}
		if reflect.DeepEqual(actual.Styles, tt.styles) {
			t.Errorf("got %v want %v", actual.Styles, tt.styles)
		}
	}
}

var readUTF16Tests = []struct {
	input  []uint8
	output string
}{
	{
		[]uint8{0x00, 0x00, 0x61, 0x00},
		"",
	},
	{
		[]uint8{0x01, 0x00, 0x61, 0x00},
		"a",
	},
	{
		[]uint8{0x01, 0x00, 0x42, 0x30},
		"\u3042",
	},
	{
		[]uint8{0x00, 0x80, 0x01, 0x00, 0x61, 0x00},
		"a",
	},
}

func TestReadUTF16(t *testing.T) {
	for _, tt := range readUTF16Tests {
		buf := bytes.NewReader(tt.input)
		sr := io.NewSectionReader(buf, 0, int64(len(tt.input)))
		actual, err := readUTF16(sr)
		if err != nil {
			t.Errorf("got %v want no error", err)
		}
		if actual != tt.output {
			t.Errorf("got %v want %v", actual, tt.output)
		}
	}
}

var readUTF8Tests = []struct {
	input  []uint8
	output string
}{
	{
		[]uint8{0x00, 0x61},
		"",
	},
	{
		[]uint8{0x01, 0x61},
		"a",
	},
	{
		[]uint8{0x03, 0xE3, 0x81, 0x82},
		"\u3042",
	},
	{
		[]uint8{0x80, 0x01, 0x61},
		"a",
	},
}

func TestReadUTF8(t *testing.T) {
	for _, tt := range readUTF8Tests {
		buf := bytes.NewReader(tt.input)
		sr := io.NewSectionReader(buf, 0, int64(len(tt.input)))
		actual, err := readUTF8(sr)
		if err != nil {
			t.Errorf("got %v want no error", err)
		}
		if actual != tt.output {
			t.Errorf("got %v want %v", actual, tt.output)
		}
	}
}

var newZeroFilledReaderTests = []struct {
	input    []uint8
	actual   int64
	expected int64
	output   []uint8
}{
	{
		input:    []uint8{0x01, 0x23, 0x45, 0x67},
		actual:   4,
		expected: 4,
		output:   []uint8{0x01, 0x23, 0x45, 0x67},
	},
	{
		input:    []uint8{0x01, 0x23, 0x45, 0x67},
		actual:   4,
		expected: 8,
		output:   []uint8{0x01, 0x23, 0x45, 0x67, 0x00, 0x00, 0x00, 0x00},
	},
}

func TestNewZeroFilledReader(t *testing.T) {
	for _, tt := range newZeroFilledReaderTests {
		buf := bytes.NewReader(tt.input)
		r, err := newZeroFilledReader(buf, tt.actual, tt.expected)
		if err != nil {
			t.Errorf("got %v want no error", err)
		}
		actualBytes := make([]uint8, tt.expected)
		r.Read(actualBytes)
		for i, a := range tt.output {
			if actualBytes[i] != a {
				t.Errorf("got %v(position %d) wants %v", actualBytes[i], i, a)
			}
		}
	}
}
<commit_msg>utf8文字列長の変更にあわせてテスト修正<commit_after>package androidbinary

import (
	"bytes"
	"io"
	"reflect"
	"testing"
)

var readStringPoolTests = []struct {
	input   []uint8
	strings []string
	styles  []string
}{
	{
		[]uint8{
			0x01, 0x00, // Type = RES_STRING_POOL_TYPE
			0x1C, 0x00, // HeaderSize = 28 bytes
			0x3C, 0x00, 0x00, 0x00, // Size = 60
			0x02, 0x00, 0x00, 0x00, // StringCount = 2
			0x02, 0x00, 0x00, 0x00, // StyleScount = 2
			0x00, 0x00, 0x00, 0x00, // Flags = 0x00
			0x2C, 0x00, 0x00, 0x00, // StringStart = 44
			0x34, 0x00, 0x00, 0x00, // StylesStart = 52

			// StringIndexes
			0x00, 0x00, 0x00, 0x00,
			0x04, 0x00, 0x00, 0x00,

			// StyleIndexes
			0x00, 0x00, 0x00, 0x00,
			0x04, 0x00, 0x00, 0x00,

			// Strings
			0x01, 0x00, 0x61, 0x00,
			0x01, 0x00, 0x42, 0x30,

			// Styles
			0x01, 0x00, 0x63, 0x00,
			0x01, 0x00, 0x43, 0x30,
		},
		[]string{"a", "\3042"},
		[]string{"b", "\3043"},
	},
}

func TestReadStringPool(t *testing.T) {
	for _, tt := range readStringPoolTests {
		buf := bytes.NewReader(tt.input)
		sr := io.NewSectionReader(buf, 0, int64(len(tt.input)))
		actual, err := readStringPool(sr)
		if err != nil {
			t.Errorf("got %v want no error", err)
		}
		if reflect.DeepEqual(actual.Strings, tt.strings) {
			t.Errorf("got %v want %v", actual.Strings, tt.strings)
		}
		if reflect.DeepEqual(actual.Styles, tt.styles) {
			t.Errorf("got %v want %v", actual.Styles, tt.styles)
		}
	}
}

var readUTF16Tests = []struct {
	input  []uint8
	output string
}{
	{
		[]uint8{0x00, 0x00, 0x61, 0x00},
		"",
	},
	{
		[]uint8{0x01, 0x00, 0x61, 0x00},
		"a",
	},
	{
		[]uint8{0x01, 0x00, 0x42, 0x30},
		"\u3042",
	},
	{
		[]uint8{0x00, 0x80, 0x01, 0x00, 0x61, 0x00},
		"a",
	},
}

func TestReadUTF16(t *testing.T) {
	for _, tt := range readUTF16Tests {
		buf := bytes.NewReader(tt.input)
		sr := io.NewSectionReader(buf, 0, int64(len(tt.input)))
		actual, err := readUTF16(sr)
		if err != nil {
			t.Errorf("got %v want no error", err)
		}
		if actual != tt.output {
			t.Errorf("got %v want %v", actual, tt.output)
		}
	}
}

var readUTF8Tests = []struct {
	input  []uint8
	output string
}{
	{
		[]uint8{0x00, 0x00, 0x61},
		"",
	},
	{
		[]uint8{0x01, 0x01, 0x61},
		"a",
	},
	{
		[]uint8{0x01, 0x03, 0xE3, 0x81, 0x82},
		"\u3042",
	},
	{
		[]uint8{0x80, 0x01, 0x80, 0x01, 0x61},
		"a",
	},
}

func TestReadUTF8(t *testing.T) {
	for _, tt := range readUTF8Tests {
		buf := bytes.NewReader(tt.input)
		sr := io.NewSectionReader(buf, 0, int64(len(tt.input)))
		actual, err := readUTF8(sr)
		if err != nil {
			t.Errorf("got %v want no error", err)
		}
		if actual != tt.output {
			t.Errorf("got %v want %v", actual, tt.output)
		}
	}
}

var newZeroFilledReaderTests = []struct {
	input    []uint8
	actual   int64
	expected int64
	output   []uint8
}{
	{
		input:    []uint8{0x01, 0x23, 0x45, 0x67},
		actual:   4,
		expected: 4,
		output:   []uint8{0x01, 0x23, 0x45, 0x67},
	},
	{
		input:    []uint8{0x01, 0x23, 0x45, 0x67},
		actual:   4,
		expected: 8,
		output:   []uint8{0x01, 0x23, 0x45, 0x67, 0x00, 0x00, 0x00, 0x00},
	},
}

func TestNewZeroFilledReader(t *testing.T) {
	for _, tt := range newZeroFilledReaderTests {
		buf := bytes.NewReader(tt.input)
		r, err := newZeroFilledReader(buf, tt.actual, tt.expected)
		if err != nil {
			t.Errorf("got %v want no error", err)
		}
		actualBytes := make([]uint8, tt.expected)
		r.Read(actualBytes)
		for i, a := range tt.output {
			if actualBytes[i] != a {
				t.Errorf("got %v(position %d) wants %v", actualBytes[i], i, a)
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2016 Eric Wollesen <ericw at xmtp dot net>

package discord

import (
	"flag"
	"fmt"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/bwmarrin/discordgo"
	"github.com/spacemonkeygo/errors"
	"github.com/spacemonkeygo/spacelog"

	"xmtp.net/xmtpbot/dice"
	"xmtp.net/xmtpbot/dur"
	"xmtp.net/xmtpbot/html"
	"xmtp.net/xmtpbot/http_server"
	"xmtp.net/xmtpbot/mildred"
	"xmtp.net/xmtpbot/remind"
	"xmtp.net/xmtpbot/seen"
	"xmtp.net/xmtpbot/twitch"
	"xmtp.net/xmtpbot/urls"
	"xmtp.net/xmtpbot/util"
)

var (
	token  = flag.String("discord.token", "", "Discord API token")
	logger = spacelog.GetLogger()

	DiscordError = errors.NewClass("discord")
)

type bot struct {
	handler_callbacks []func()
	seen              seen.Store
	user_id           string
	urls              urls.Store
	mildred           mildred.Conn
	remind            remind.Remind
	twitch_client     twitch.Twitch
	http_server       http_server.Server
}

func New(urls_store urls.Store, seen_store seen.Store, mildred mildred.Conn,
	remind remind.Remind, twitch twitch.Twitch,
	http_server http_server.Server) *bot {

	return &bot{
		seen:          seen_store,
		urls:          urls_store,
		mildred:       mildred,
		remind:        remind,
		twitch_client: twitch,
		http_server:   http_server,
	}
}

func (b *bot) Run(shutdown chan bool, wg *sync.WaitGroup) (err error) {
	session, err := b.logIn(shutdown)
	if err != nil {
		return err
	}
	wg.Add(1)
	logger.Info("online")

	err = b.http_server.GiveRouter("twitch", b.twitch_client.ReceiveRouter)
	if err != nil {
		logger.Errore(err)
	}

	go b.http_server.Serve()

	go func() {
		<-shutdown
		logger.Infof("shutting down")
		b.logOut(session)
		wg.Done()
	}()

	return nil
}

func (b *bot) logIn(shutdown chan bool) (session *discordgo.Session, err error) {
	session, err = discordgo.New(getToken())
	if err != nil {
		return nil, DiscordError.Wrap(err)
	}

	err = b.addHandlers(session, shutdown)
	if err != nil {
		return nil, DiscordError.Wrap(err)
	}

	err = session.Open()
	if err != nil {
		return nil, DiscordError.Wrap(err)
	}

	return session, nil
}

func getToken() (discord_token string) {
	defer func() {
		if discord_token == "" {
			logger.Warnf("discord token is blank, connection will " +
				"silenty fail")
		}
	}()

	if *token == "" {
		return os.Getenv("DISCORD_TOKEN")
	}

	return *token
}

func (b *bot) logOut(session *discordgo.Session) {
	logger.Errore(session.Close())
	logger.Errore(b.removeHandlers())
	logger.Info("offline")
}

func (b *bot) removeHandlers() (err error) {
	for _, cb := range b.handler_callbacks {
		cb()
	}

	return nil
}

func (b *bot) addHandlers(session *discordgo.Session, shutdown chan bool) (
	err error) {

	b.handler_callbacks = append(b.handler_callbacks,
		session.AddHandler(b.messageHandler))
	b.handler_callbacks = append(b.handler_callbacks,
		session.AddHandler(b.presenceHandler))

	return nil
}

func (b *bot) handleCommand(cmd Command) string {
	switch cmd.cmd {
	case "dice":
		return dice.Roll(cmd.args)
	case "faq":
		return "No FAQs answered yet"
	case "idle":
		if cmd.args == "" {
			return "No name specified"
		}

		since, err := b.seen.Idle(cmd.args)
		if err != nil {
			return fmt.Sprintf("error retrieving idle for %q",
				cmd.args)
		}
		if since == nil {
			return fmt.Sprintf("No idle record for %q found",
				cmd.args)
		}
		return fmt.Sprintf("%s idle for %s", cmd.args, since)
	case "link":
		return b.lookupURL(cmd.args)
	case "np":
		return b.nowPlaying()
	case "ping":
		return "pong"
	case "remind":
		return b.setReminder(cmd)
	case "roll":
		return dice.Roll(cmd.args)
	case "seen":
		if cmd.args == "" {
			return "No name specified"
		}

		at, err := b.seen.LastSeen(cmd.args)
		if err != nil {
			return fmt.Sprintf("error retrieving last seen for %q",
				cmd.args)
		}
		if at == nil {
			return fmt.Sprintf("No seen record for %q found",
				cmd.args)
		}
		return fmt.Sprintf("%s was last seen %s", cmd.args, at)
	case "syn":
		return "ack"
	case "twitch":
		return b.twitch(cmd.args)
	case "url":
		return b.lookupURL(cmd.args)
	default:
		return fmt.Sprintf("unhandled command: %q", cmd.cmd)
	}
}

type Command struct {
	cmd     string
	args    string
	session *discordgo.Session
	message *discordgo.Message
}

func (b *bot) messageHandler(s *discordgo.Session, m *discordgo.MessageCreate) {
	logger.Debugf("<- %s", m.Content)
	logger.Warne(b.markSeen(m.Author.Username))

	if m.Author.ID == b.myDiscordUserId(s) {
		return
	}

	response := ""

	if strings.HasPrefix(m.Content, "!") {
		args := strings.SplitN(m.Content, " ", 2)
		new_args := ""
		if len(args) > 1 {
			new_args = args[1]
		}
		response = b.handleCommand(Command{
			cmd:     args[0][1:],
			args:    new_args,
			session: s,
			message: m.Message,
		})
	}

	urls := b.parseURLs(m.Content)
	if len(urls) > 0 {
		b.rememberURLs(urls...)
	}

	if response != "" {
		for _, chunk := range splitResponse(response, 10) {
			_, err := s.ChannelMessageSend(m.ChannelID, chunk)
			if err != nil {
				logger.Warne(err)
			}
		}
	}

}

func splitResponse(response string, num_lines int) (chunks []string) {
	lines := strings.Split(response, "\n")
	chunk := []string{lines[0]}
	for i := 1; i < len(lines); i++ {
		if len(chunk) >= num_lines {
			chunks = append(chunks, strings.Join(chunk, "\n"))
			chunk = []string{}
		}
		chunk = append(chunk, lines[i])
	}

	if len(chunk) > 0 {
		chunks = append(chunks, strings.Join(chunk, "\n"))
	}

	return chunks
}

func (b *bot) myDiscordUserId(s *discordgo.Session) string {
	if b.user_id != "" {
		return b.user_id
	}

	user, err := s.User("@me")
	if err != nil {
		logger.Warnf("unable to find my user id: %v", err)
		return ""
	}

	b.user_id = user.ID

	return user.ID
}

func (b *bot) presenceHandler(s *discordgo.Session, p *discordgo.PresenceUpdate) {
	logger.Warne(b.markSeen(p.User.Username))
}

func (b *bot) markSeen(name string) error {
	return b.seen.MarkSeen(name, nil)
}

func (b *bot) parseURLs(msg string) []string {
	return urls.Parse(msg)
}

func (b *bot) rememberURLs(urls ...string) {
	for _, url := range urls {
		title, err := html.ParseTitleFromURL(url)
		if err != nil {
			logger.Warne(err)
		}
		b.urls.Remember(url, title)
		logger.Debugf("remembered URL %q", url)
	}
}

func (b *bot) lookupURL(msg string) string {
	urls := b.urls.Lookup(msg)
	if len(urls) > 0 {
		lines := []string{}
		for _, url := range urls {
			lines = append(lines,
				fmt.Sprintf("%s - %s", url[0], url[1]))
		}
		sort.Strings(lines)
		return fmt.Sprintf("Matched %d URLs:\n%s",
			len(urls), strings.Join(lines, "\n"))
	}
	return "No matching URLs found"
}

func (b *bot) nowPlaying() string {
	cs := b.mildred.CurrentSong()
	if cs != nil {
		return cs.String()
	} else {
		return "error determining current song"
	}
}

func (b *bot) setReminder(cmd Command) string {
	duration, matched, err := dur.Parse(cmd.args)
	if err != nil {
		return "couldn't parse reminder duration"
	}
	msg := cmd.args[len(matched):]

	b.remind.Set(time.Now().Add(*duration), func() {
		msg := fmt.Sprintf("<@%s> reminder: %s",
			cmd.message.Author.ID, msg)
		_, err := cmd.session.ChannelMessageSend(cmd.message.ChannelID,
			msg)
		logger.Warne(err)
	})

	return "reminder set"
}

func (b *bot) twitch(args string) string {
	pieces := strings.SplitN(args, " ", 2)
	cmd := pieces[0]
	if len(pieces) > 1 {
		args = pieces[1]
	} else {
		args = ""
	}
	client := b.twitch_client

	switch cmd {
	case "auth":
		auth_url, err := client.Auth(args)
		if err != nil {
			return "error generating OAuth2 URL"
		}
		return fmt.Sprintf("Click here to authorize xMTP bot to access "+
			"your Twitch account: %s", auth_url)
	case "auth-follow":
		pieces := strings.Split(args, " ")
		err := client.AuthFollow(pieces[0], pieces[1:]...)
		if err != nil {
			logger.Errorf("error auth following: %v", err)
			return "error auth following"
		}
		return "Done"
	case "live":
		var response []string
		streams, err := client.Live()
		if err != nil {
			logger.Errore(err)
			return "error retrieving live twitch streams"
		}
		for _, stream := range streams {
			response = append(response,
				fmt.Sprintf("%s: %s",
					util.EscapeMarkdown(stream.Name()),
					stream.URL()))
		}
		if len(response) == 0 {
			return "no streams are live"
		}
		return strings.Join(response, "\n")
	case "follow":
		if args == "" {
			return "must specify channel name to follow"
		}
		logger.Debugf("args: %s", args)
		return client.Follow(args)
	case "unfollow":
		if args == "" {
			return "must specify channel name to unfollow"
		}
		logger.Debugf("args: %s", args)
		client.Unfollow(args)
		return "OK"
	case "following", "list":
		var response []string
		channels, err := client.Following(args)
		if err != nil {
			logger.Errore(err)
			return "error retrieving followed twitch channels"
		}
		for _, channel := range channels {
			// wrapping the link in parentheses seems to prevent the
			// discord client from expanding it?
			response = append(response,
				fmt.Sprintf("%s (%s)",
					util.EscapeMarkdown(channel.Name()), channel.URL()))
		}

		if len(channels) == 0 {
			return "no channels are followed"
		}
		return strings.Join(response, "\n")
	default:
		return fmt.Sprintf("unhandled twitch command: %q", cmd)
	}
}
<commit_msg>commands registration, plus a help command<commit_after>// Copyright (C) 2016 Eric Wollesen <ericw at xmtp dot net>

package discord

import (
	"flag"
	"fmt"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/bwmarrin/discordgo"
	"github.com/spacemonkeygo/errors"
	"github.com/spacemonkeygo/spacelog"

	"xmtp.net/xmtpbot/dice"
	"xmtp.net/xmtpbot/dur"
	"xmtp.net/xmtpbot/html"
	"xmtp.net/xmtpbot/http_server"
	"xmtp.net/xmtpbot/mildred"
	"xmtp.net/xmtpbot/remind"
	"xmtp.net/xmtpbot/seen"
	"xmtp.net/xmtpbot/twitch"
	"xmtp.net/xmtpbot/urls"
	"xmtp.net/xmtpbot/util"
)

var (
	token  = flag.String("discord.token", "", "Discord API token")
	logger = spacelog.GetLogger()

	DiscordError = errors.NewClass("discord")
)

type bot struct {
	handler_callbacks []func()
	seen              seen.Store
	user_id           string
	urls              urls.Store
	mildred           mildred.Conn
	remind            remind.Remind
	twitch_client     twitch.Twitch
	http_server       http_server.Server
	commands_mtx      sync.Mutex
	commands          map[string]CommandHandler
}

func New(urls_store urls.Store, seen_store seen.Store, mildred mildred.Conn,
	remind remind.Remind, twitch twitch.Twitch,
	http_server http_server.Server) *bot {

	b := &bot{
		seen:          seen_store,
		urls:          urls_store,
		mildred:       mildred,
		remind:        remind,
		twitch_client: twitch,
		http_server:   http_server,
		commands:      make(map[string]CommandHandler),
	}

	b.RegisterCommand("commands", simpleCommand(b.listCommands,
		"list available commands"))
	b.RegisterCommand("faq", staticCommand("No FAQs answered yet",
		"frequently answered questions"))
	b.RegisterCommand("help", simpleCommand(b.help,
		"list available commands"))
	b.RegisterCommand("idle", &commandHandler{
		help:    "reports a user's idle time",
		handler: b.idle,
	})
	b.RegisterCommand("link", simpleCommand(b.lookupURL,
		"search for a previously posted URL"))
	b.RegisterCommand("np", simpleCommand(b.nowPlaying,
		"report Mildred's currently playing track"))
	b.RegisterCommand("ping", staticCommand("pong", "pong"))
	b.RegisterCommand("remind", &commandHandler{
		help: "sets a reminder. " +
			"Example !remind 5 minutes take out the trash",
		handler: b.setReminder,
	})
	b.RegisterCommand("roll", simpleCommand(dice.Roll, "roll some dice"))
	b.RegisterCommand("seen", &commandHandler{
		help:    "reports when a user was last seen",
		handler: b.lastSeen,
	})
	b.RegisterCommand("syn", staticCommand("ack", "ack"))
	b.RegisterCommand("twitch", simpleCommand(b.twitch,
		"interact with twitch. Run \"!twitch help\" for more info"))
	b.RegisterCommand("url", simpleCommand(b.lookupURL,
		"search for a previously posted URL"))

	return b
}

func (b *bot) Run(shutdown chan bool, wg *sync.WaitGroup) (err error) {
	session, err := b.logIn(shutdown)
	if err != nil {
		return err
	}
	wg.Add(1)
	logger.Info("online")

	err = b.http_server.GiveRouter("twitch", b.twitch_client.ReceiveRouter)
	if err != nil {
		logger.Errore(err)
	}

	go b.http_server.Serve()

	go func() {
		<-shutdown
		logger.Infof("shutting down")
		b.logOut(session)
		wg.Done()
	}()

	return nil
}

func (b *bot) logIn(shutdown chan bool) (session *discordgo.Session, err error) {
	session, err = discordgo.New(getToken())
	if err != nil {
		return nil, DiscordError.Wrap(err)
	}

	err = b.addHandlers(session, shutdown)
	if err != nil {
		return nil, DiscordError.Wrap(err)
	}

	err = session.Open()
	if err != nil {
		return nil, DiscordError.Wrap(err)
	}

	return session, nil
}

func getToken() (discord_token string) {
	defer func() {
		if discord_token == "" {
			logger.Warnf("discord token is blank, connection will " +
				"silenty fail")
		}
	}()

	if *token == "" {
		return os.Getenv("DISCORD_TOKEN")
	}

	return *token
}

func (b *bot) logOut(session *discordgo.Session) {
	logger.Errore(session.Close())
	logger.Errore(b.removeHandlers())
	logger.Info("offline")
}

func (b *bot) removeHandlers() (err error) {
	for _, cb := range b.handler_callbacks {
		cb()
	}

	return nil
}

func (b *bot) addHandlers(session *discordgo.Session, shutdown chan bool) (
	err error) {

	b.handler_callbacks = append(b.handler_callbacks,
		session.AddHandler(b.messageHandler))
	b.handler_callbacks = append(b.handler_callbacks,
		session.AddHandler(b.presenceHandler))

	return nil
}

func (b *bot) handleCommand(cmd Command) string {
	handler, ok := b.commands[cmd.cmd]
	if ok {
		handler.Handle(cmd)
		return "" // TODO fixme ugly
	}

	return fmt.Sprintf("unhandled command: %q", cmd.cmd)
}

type Command struct {
	cmd     string
	args    string
	session *discordgo.Session
	message *discordgo.Message
}

func (c *Command) Reply(msg string) (err error) {
	_, err = c.session.ChannelMessageSend(c.message.ChannelID, msg)
	return err
}

func (b *bot) messageHandler(s *discordgo.Session, m *discordgo.MessageCreate) {
	logger.Debugf("<- %s", m.Content)
	logger.Warne(b.markSeen(m.Author.Username))

	if m.Author.ID == b.myDiscordUserId(s) {
		return
	}

	response := ""

	if strings.HasPrefix(m.Content, "!") {
		args := strings.SplitN(m.Content, " ", 2)
		new_args := ""
		if len(args) > 1 {
			new_args = args[1]
		}
		response = b.handleCommand(Command{
			cmd:     args[0][1:],
			args:    new_args,
			session: s,
			message: m.Message,
		})
	}

	urls := b.parseURLs(m.Content)
	if len(urls) > 0 {
		b.rememberURLs(urls...)
	}

	if response != "" {
		for _, chunk := range splitResponse(response, 10) {
			_, err := s.ChannelMessageSend(m.ChannelID, chunk)
			if err != nil {
				logger.Warne(err)
			}
		}
	}

}

func splitResponse(response string, num_lines int) (chunks []string) {
	lines := strings.Split(response, "\n")
	chunk := []string{lines[0]}
	for i := 1; i < len(lines); i++ {
		if len(chunk) >= num_lines {
			chunks = append(chunks, strings.Join(chunk, "\n"))
			chunk = []string{}
		}
		chunk = append(chunk, lines[i])
	}

	if len(chunk) > 0 {
		chunks = append(chunks, strings.Join(chunk, "\n"))
	}

	return chunks
}

func (b *bot) myDiscordUserId(s *discordgo.Session) string {
	if b.user_id != "" {
		return b.user_id
	}

	user, err := s.User("@me")
	if err != nil {
		logger.Warnf("unable to find my user id: %v", err)
		return ""
	}

	b.user_id = user.ID

	return user.ID
}

func (b *bot) presenceHandler(s *discordgo.Session, p *discordgo.PresenceUpdate) {
	logger.Warne(b.markSeen(p.User.Username))
}

func (b *bot) markSeen(name string) error {
	return b.seen.MarkSeen(name, nil)
}

func (b *bot) parseURLs(msg string) []string {
	return urls.Parse(msg)
}

func (b *bot) rememberURLs(urls ...string) {
	for _, url := range urls {
		title, err := html.ParseTitleFromURL(url)
		if err != nil {
			logger.Warne(err)
		}
		b.urls.Remember(url, title)
		logger.Debugf("remembered URL %q", url)
	}
}

func (b *bot) lookupURL(msg string) string {
	urls := b.urls.Lookup(msg)
	if len(urls) > 0 {
		lines := []string{}
		for _, url := range urls {
			lines = append(lines,
				fmt.Sprintf("%s - %s", url[0], url[1]))
		}
		sort.Strings(lines)
		return fmt.Sprintf("Matched %d URLs:\n%s",
			len(urls), strings.Join(lines, "\n"))
	}
	return "No matching URLs found"
}

func (b *bot) nowPlaying(args string) string {
	cs := b.mildred.CurrentSong()
	if cs != nil {
		return cs.String()
	} else {
		return "error determining current song"
	}
}

func (b *bot) idle(cmd Command) error {
	if cmd.args == "" {
		return cmd.Reply("No name specified")
	}

	since, err := b.seen.Idle(cmd.args)
	if err != nil {
		return cmd.Reply(fmt.Sprintf("Error retrieving idle for %q",
			cmd.args))
	}
	if since == nil {
		return cmd.Reply(fmt.Sprintf("No idle record for %q found",
			cmd.args))
	}

	return cmd.Reply(fmt.Sprintf("%s idle for %s", cmd.args, since))
}

func (b *bot) lastSeen(cmd Command) error {
	if cmd.args == "" {
		return cmd.Reply("No name specified")
	}

	at, err := b.seen.LastSeen(cmd.args)
	if err != nil {
		return cmd.Reply(fmt.Sprintf("Error retrieving last seen for %q",
			cmd.args))
	}
	if at == nil {
		return cmd.Reply(fmt.Sprintf("No seen record for %q found",
			cmd.args))
	}

	return cmd.Reply(fmt.Sprintf("%s was last seen %s", cmd.args, at))
}

func (b *bot) setReminder(cmd Command) error {
	duration, matched, err := dur.Parse(cmd.args)
	if err != nil {
		return cmd.Reply("couldn't parse reminder duration")
	}
	msg := cmd.args[len(matched):]

	b.remind.Set(time.Now().Add(*duration), func() {
		msg := fmt.Sprintf("<@%s> reminder: %s",
			cmd.message.Author.ID, msg)
		_, err := cmd.session.ChannelMessageSend(cmd.message.ChannelID,
			msg)
		logger.Warne(err)
	})

	return cmd.Reply("reminder set")
}

func (b *bot) twitch(args string) string {
	pieces := strings.SplitN(args, " ", 2)
	cmd := pieces[0]
	if len(pieces) > 1 {
		args = pieces[1]
	} else {
		args = ""
	}
	client := b.twitch_client

	switch cmd {
	case "auth":
		auth_url, err := client.Auth(args)
		if err != nil {
			return "error generating OAuth2 URL"
		}
		return fmt.Sprintf("Click here to authorize xMTP bot to access "+
			"your Twitch account: %s", auth_url)
	case "auth-follow":
		pieces := strings.Split(args, " ")
		err := client.AuthFollow(pieces[0], pieces[1:]...)
		if err != nil {
			logger.Errorf("error auth following: %v", err)
			return "error auth following"
		}
		return "Done"
	case "live":
		var response []string
		streams, err := client.Live()
		if err != nil {
			logger.Errore(err)
			return "error retrieving live twitch streams"
		}
		for _, stream := range streams {
			response = append(response,
				fmt.Sprintf("%s: %s",
					util.EscapeMarkdown(stream.Name()),
					stream.URL()))
		}
		if len(response) == 0 {
			return "no streams are live"
		}
		return strings.Join(response, "\n")
	case "follow":
		if args == "" {
			return "must specify channel name to follow"
		}
		logger.Debugf("args: %s", args)
		return client.Follow(args)
	case "unfollow":
		if args == "" {
			return "must specify channel name to unfollow"
		}
		logger.Debugf("args: %s", args)
		client.Unfollow(args)
		return "OK"
	case "following", "list":
		var response []string
		channels, err := client.Following(args)
		if err != nil {
			logger.Errore(err)
			return "error retrieving followed twitch channels"
		}
		for _, channel := range channels {
			// wrapping the link in parentheses seems to prevent the
			// discord client from expanding it?
			response = append(response,
				fmt.Sprintf("%s (%s)",
					util.EscapeMarkdown(channel.Name()), channel.URL()))
		}

		if len(channels) == 0 {
			return "no channels are followed"
		}
		return strings.Join(response, "\n")
	default:
		return fmt.Sprintf("unhandled twitch command: %q", cmd)
	}
}

type CommandHandler interface {
	Handle(Command) error
	Help() string
}

type commandHandler struct {
	help    string
	handler func(Command) error
}

func (h *commandHandler) Help() string {
	return h.help
}

func (h *commandHandler) Handle(cmd Command) (err error) {
	return h.handler(cmd)
}

func (b *bot) RegisterCommand(name string, handler CommandHandler) (err error) {
	b.commands_mtx.Lock()
	defer b.commands_mtx.Unlock()

	b.commands[name] = handler

	return nil
}

func staticCommand(response, help string) CommandHandler {
	return &commandHandler{
		help: help,
		handler: func(cmd Command) error {
			return cmd.Reply(response)
		},
	}
}

func simpleCommand(fn func(args string) string, help string) CommandHandler {
	return &commandHandler{
		help: help,
		handler: func(cmd Command) error {
			return cmd.Reply(fn(cmd.args))
		},
	}
}

func (b *bot) help(cmd string) string {
	if cmd == "" {
		return "usage: help <command>"
	}

	b.commands_mtx.Lock()
	ch, ok := b.commands[cmd]
	b.commands_mtx.Unlock()

	if !ok {
		return fmt.Sprintf("No help for %q found", cmd)
	}

	return fmt.Sprintf("%s: %s", cmd, ch.Help())
}

func (b *bot) listCommands(cmd string) string {
	b.commands_mtx.Lock()
	cmds := b.commands
	b.commands_mtx.Unlock()

	var strs []string
	for name, _ := range cmds {
		strs = append(strs, name)
	}
	sort.Strings(strs)

	return strings.Join(strs, ", ")
}
<|endoftext|>


<commit_before>package integration

import (
	"net/http"
	"os"
	"time"

	"github.com/containous/traefik/integration/try"
	"github.com/go-check/check"
	checker "github.com/vdemeester/shakers"
)

type TimeoutSuite struct{ BaseSuite }

func (s *TimeoutSuite) SetUpSuite(c *check.C) {
	s.createComposeProject(c, "timeout")
	s.composeProject.Start(c)
}

func (s *TimeoutSuite) TestForwardingTimeouts(c *check.C) {
	httpTimeoutEndpoint := s.composeProject.Container(c, "timeoutEndpoint").NetworkSettings.IPAddress
	file := s.adaptFile(c, "fixtures/timeout/forwarding_timeouts.toml", struct {
		TimeoutEndpoint string
	}{httpTimeoutEndpoint})
	defer os.Remove(file)

	cmd, display := s.traefikCmd(withConfigFile(file))
	defer display(c)
	err := cmd.Start()
	c.Assert(err, checker.IsNil)
	defer cmd.Process.Kill()

	err = try.GetRequest("http://127.0.0.1:8080/api/providers/file/routers", 60*time.Second, try.BodyContains("Path:/dialTimeout"))
	c.Assert(err, checker.IsNil)

	// This simulates a DialTimeout when connecting to the backend server.
	response, err := http.Get("http://127.0.0.1:8000/dialTimeout")
	c.Assert(err, checker.IsNil)
	c.Assert(response.StatusCode, checker.Equals, http.StatusGatewayTimeout)

	// This simulates a ResponseHeaderTimeout.
	response, err = http.Get("http://127.0.0.1:8000/responseHeaderTimeout?sleep=1000")
	c.Assert(err, checker.IsNil)
	c.Assert(response.StatusCode, checker.Equals, http.StatusGatewayTimeout)
}
<commit_msg>Assert that test timeout service is ready.<commit_after>package integration

import (
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/containous/traefik/integration/try"
	"github.com/go-check/check"
	checker "github.com/vdemeester/shakers"
)

type TimeoutSuite struct{ BaseSuite }

func (s *TimeoutSuite) SetUpSuite(c *check.C) {
	s.createComposeProject(c, "timeout")
	s.composeProject.Start(c)
}

func (s *TimeoutSuite) TestForwardingTimeouts(c *check.C) {
	httpTimeoutEndpoint := s.composeProject.Container(c, "timeoutEndpoint").NetworkSettings.IPAddress
	file := s.adaptFile(c, "fixtures/timeout/forwarding_timeouts.toml", struct {
		TimeoutEndpoint string
	}{httpTimeoutEndpoint})
	defer os.Remove(file)

	cmd, display := s.traefikCmd(withConfigFile(file))
	defer display(c)
	err := cmd.Start()
	c.Assert(err, checker.IsNil)
	defer cmd.Process.Kill()

	err = try.GetRequest("http://127.0.0.1:8080/api/providers/file/routers", 60*time.Second, try.BodyContains("Path:/dialTimeout"))
	c.Assert(err, checker.IsNil)

	// This simulates a DialTimeout when connecting to the backend server.
	response, err := http.Get("http://127.0.0.1:8000/dialTimeout")
	c.Assert(err, checker.IsNil)
	c.Assert(response.StatusCode, checker.Equals, http.StatusGatewayTimeout)

	// Check that timeout service is available
	statusURL := fmt.Sprintf("http://%s:9000/statusTest?status=200", httpTimeoutEndpoint)
	c.Assert(try.GetRequest(statusURL, 60*time.Second, try.StatusCodeIs(http.StatusOK)), checker.IsNil)

	// This simulates a ResponseHeaderTimeout.
	response, err = http.Get("http://127.0.0.1:8000/responseHeaderTimeout?sleep=1000")
	c.Assert(err, checker.IsNil)
	c.Assert(response.StatusCode, checker.Equals, http.StatusGatewayTimeout)
}
<|endoftext|>