<commit_before>package terminal

import (
	"bytes"
	"errors"
	"fmt"
	"os/exec"
	"strings"
)

const (
	sessionPrefix     = "koding"
	defaultScreenPath = "/usr/bin/screen"
)

type Command struct {
	// Name is used for starting the terminal instance, it's the program path
	// usually
	Name string

	// Args is passed to the program name
	Args []string
}

var (
	ErrNoSession      = errors.New("ErrNoSession")
	ErrInvalidSession = errors.New("ErrInvalidSession")
)

// newCmd returns a new command instance that is used to start the terminal.
// The command line is created differently based on the incoming mode.
func newCommand(mode, session, username string) (*Command, error) {
	// let's assume by default its Screen
	name := defaultScreenPath
	args := []string{"-S"}

	switch mode {
	case "shared", "resume":
		if session == "" {
			return nil, errors.New("session is needed for 'shared' or 'resume' mode")
		}

		if !sessionExists(session, username) {
			return nil, ErrNoSession
		}

		args = append(args, sessionPrefix+"."+session)
		if mode == "shared" {
			args = append(args, "-x") // multiuser mode
		} else if mode == "resume" {
			args = append(args, "-raAd") // resume
		}
	case "noscreen":
		name = "/bin/bash"
		args = []string{}
	case "create":
		args = append(args, sessionPrefix+"."+randomString())
	default:
		return nil, fmt.Errorf("mode '%s' is unknown. Valid modes are:  [shared|noscreen|resume|create]", mode)
	}

	c := &Command{
		Name: name,
		Args: args,
	}

	fmt.Printf("c %#v\n", c)

	return c, nil
}

// screenSessions returns a list of sessions that belongs to the given
// username.  The sessions are in the form of ["k7sdjv12344", "askIj12sas12",
// ...]
// TODO: socket directory is different under darwin, it will not work probably
func screenSessions(username string) []string {
	// Do not include dead sessions in our result
	exec.Command(defaultScreenPath, "-wipe").Run()

	// We need to use ls here, because /var/run/screen mount is only
	// visible from inside of container. Errors are ignored.
	out, _ := exec.Command("ls", "/var/run/screen/S-"+username).Output()
	shellOut := string(bytes.TrimSpace(out))
	if shellOut == "" {
		return []string{}
	}

	names := strings.Split(shellOut, "\n")
	sessions := make([]string, len(names))

	prefix := sessionPrefix + "."
	for i, name := range names {
		segments := strings.SplitN(name, ".", 2)
		sessions[i] = strings.TrimPrefix(segments[1], prefix)
	}

	return sessions
}

// screenExists checks whether the given session exists in the running list of
// screen sessions.
func sessionExists(session, username string) bool {
	for _, s := range screenSessions(username) {
		if s == session {
			return true
		}
	}

	return false
}

// killSession kills the given SessionID
func killSession(session string) error {
	out, err := exec.Command(defaultScreenPath, "-X", "-S", sessionPrefix+"."+session, "kill").Output()
	if err != nil {
		return commandError("screen kill failed", err, out)
	}

	return nil
}

func commandError(message string, err error, out []byte) error {
	return fmt.Errorf("%s\n%s\n%s", message, err.Error(), string(out))
}
<commit_msg>kite-handler/terminal: start screen in bash shell<commit_after>package terminal

import (
	"bytes"
	"errors"
	"fmt"
	"os/exec"
	"strings"
)

const (
	sessionPrefix     = "koding"
	defaultScreenPath = "/usr/bin/screen"
)

type Command struct {
	// Name is used for starting the terminal instance, it's the program path
	// usually
	Name string

	// Args is passed to the program name
	Args []string
}

var (
	ErrNoSession      = errors.New("ErrNoSession")
	ErrInvalidSession = errors.New("ErrInvalidSession")
)

// newCmd returns a new command instance that is used to start the terminal.
// The command line is created differently based on the incoming mode.
func newCommand(mode, session, username string) (*Command, error) {
	// let's assume by default its Screen
	name := defaultScreenPath
	args := []string{"-e^Bb", "-s", "/bin/bash", "-S"}

	switch mode {
	case "shared", "resume":
		if session == "" {
			return nil, errors.New("session is needed for 'shared' or 'resume' mode")
		}

		if !sessionExists(session, username) {
			return nil, ErrNoSession
		}

		args = append(args, sessionPrefix+"."+session)
		if mode == "shared" {
			args = append(args, "-x") // multiuser mode
		} else if mode == "resume" {
			args = append(args, "-raAd") // resume
		}
	case "noscreen":
		name = "/bin/bash"
		args = []string{}
	case "create":
		args = append(args, sessionPrefix+"."+randomString())
	default:
		return nil, fmt.Errorf("mode '%s' is unknown. Valid modes are:  [shared|noscreen|resume|create]", mode)
	}

	c := &Command{
		Name: name,
		Args: args,
	}

	return c, nil
}

// screenSessions returns a list of sessions that belongs to the given
// username.  The sessions are in the form of ["k7sdjv12344", "askIj12sas12",
// ...]
// TODO: socket directory is different under darwin, it will not work probably
func screenSessions(username string) []string {
	// Do not include dead sessions in our result
	exec.Command(defaultScreenPath, "-wipe").Run()

	// We need to use ls here, because /var/run/screen mount is only
	// visible from inside of container. Errors are ignored.
	out, _ := exec.Command("ls", "/var/run/screen/S-"+username).Output()
	shellOut := string(bytes.TrimSpace(out))
	if shellOut == "" {
		return []string{}
	}

	names := strings.Split(shellOut, "\n")
	sessions := make([]string, len(names))

	prefix := sessionPrefix + "."
	for i, name := range names {
		segments := strings.SplitN(name, ".", 2)
		sessions[i] = strings.TrimPrefix(segments[1], prefix)
	}

	return sessions
}

// screenExists checks whether the given session exists in the running list of
// screen sessions.
func sessionExists(session, username string) bool {
	for _, s := range screenSessions(username) {
		if s == session {
			return true
		}
	}

	return false
}

// killSession kills the given SessionID
func killSession(session string) error {
	out, err := exec.Command(defaultScreenPath, "-X", "-S", sessionPrefix+"."+session, "kill").Output()
	if err != nil {
		return commandError("screen kill failed", err, out)
	}

	return nil
}

func commandError(message string, err error, out []byte) error {
	return fmt.Errorf("%s\n%s\n%s", message, err.Error(), string(out))
}
<|endoftext|>


<commit_before>package cache

import (
	"errors"
	"fmt"
	"gnd.la/cache/driver"
	"gnd.la/config"
	"gnd.la/encoding/codec"
	"gnd.la/encoding/pipe"
	"gnd.la/log"
	"reflect"
	"strings"
)

var (
	ErrNotFound = errors.New("item not found in cache")
	imports     = map[string]string{
		"memcache": "gnd.la/cache/driver/memcache",
		"redis":    "gnd.la/cache/driver/redis",
	}
)

type Cache struct {
	// The Logger to log debug messages and, more importantly, errors.
	// New() initialies the Logger to log.Std.
	Logger     *log.Logger
	prefix     string
	prefixLen  int
	driver     driver.Driver
	codec      *codec.Codec
	pipe       *pipe.Pipe
	numQueries int
}

func (c *Cache) backendKey(key string) string {
	if c.prefixLen > 0 {
		return c.prefix + key
	}
	return key
}

func (c *Cache) frontendKey(key string) string {
	return key[c.prefixLen:]
}

// Set stores the given object in the cache associated with the
// given key. Timeout is the number of seconds until the item
// expires. If the timeout is 0, the item never expires, but
// might be only purged from cache when running out of space.
func (c *Cache) Set(key string, object interface{}, timeout int) error {
	b, err := c.codec.Encode(object)
	if err != nil {
		eerr := &cacheError{
			op:    "encoding object",
			key:   key,
			codec: true,
			err:   err,
		}
		c.error(eerr)
		return eerr
	}
	return c.SetBytes(key, b, timeout)
}

// Get retrieves the requested item from the cache and decodes it
// into the passed interface{}, which must be addressable.
// If the item is not found in the cache, ErrNotFound is returned.
// Other errors mean that either there was a problem communicating
// with the cache or the item could not be decoded.
func (c *Cache) Get(key string, obj interface{}) error {
	b, err := c.GetBytes(key)
	if err != nil {
		return err
	}
	cerr := c.codec.Decode(b, obj)
	if cerr != nil {
		derr := &cacheError{
			op:    "decoding object",
			key:   key,
			codec: true,
			err:   cerr,
		}
		c.error(derr)
		return derr
	}
	return nil
}

// GetMulti returns several objects as a map[string]interface{}
// in only one roundtrip to the cache. The obj parameter is used only
// to initialize the objects before they're passed to the codec for decoding,
// since not all codecs include the object type in its serialization (e.g. JSON).
// If you're using a codec which encodes the object type (e.g. Gob) or you want
// to decode the objects into a raw interface{} you might pass nil as the second
// parameter to this function.
func (c *Cache) GetMulti(keys []string, obj interface{}) (map[string]interface{}, error) {
	c.numQueries++
	if c.prefixLen > 0 {
		k := make([]string, len(keys))
		for ii, v := range keys {
			k[ii] = c.backendKey(v)
		}
		keys = k
	}
	data, err := c.driver.GetMulti(keys)
	if err != nil {
		gerr := &cacheError{
			op:  "getting multiple keys",
			key: strings.Join(keys, ", "),
			err: err,
		}
		c.error(gerr)
		return nil, gerr
	}
	objects := make(map[string]interface{}, len(data))
	var typ *reflect.Type
	if obj != nil {
		t := reflect.TypeOf(obj)
		for t.Kind() == reflect.Ptr {
			t = t.Elem()
		}
		typ = &t
	}
	for k, v := range data {
		var object interface{}
		if typ != nil {
			p := reflect.New(*typ)
			err = c.codec.Decode(v, p.Interface())
			object = p.Elem().Interface()
		} else {
			err = c.codec.Decode(v, &object)
		}
		if err != nil {
			derr := &cacheError{
				op:  "decoding object",
				key: k,
				err: err,
			}
			c.error(derr)
			return nil, derr
		}
		objects[c.frontendKey(k)] = object
	}
	return objects, nil
}

// SetBytes stores the given byte array assocciated with
// the given key. See the documentation for Set for an
// explanation of the timeout parameter
func (c *Cache) SetBytes(key string, b []byte, timeout int) error {
	c.numQueries++
	if c.pipe != nil {
		var err error
		b, err = c.pipe.Encode(b)
		if err != nil {
			perr := &cacheError{
				op:  "encoding data with pipe",
				key: key,
				err: err,
			}
			c.error(perr)
			return perr
		}
	}
	k := c.backendKey(key)
	err := c.driver.Set(k, b, timeout)
	if err != nil {
		serr := &cacheError{
			op:  "setting key",
			key: key,
			err: err,
		}
		c.error(serr)
		return serr
	}
	c.debugf("Set key %s (%d bytes), expiring in %d", k, len(b), timeout)
	return nil
}

// GetBytes returns the byte array assocciated with the given key
func (c *Cache) GetBytes(key string) ([]byte, error) {
	c.numQueries++
	b, err := c.driver.Get(c.backendKey(key))
	if err != nil {
		gerr := &cacheError{
			op:  "getting key",
			key: key,
			err: err,
		}
		c.error(gerr)
		return nil, gerr
	}
	if b == nil {
		return nil, ErrNotFound
	}
	if c.pipe != nil {
		b, err = c.pipe.Decode(b)
		if err != nil {
			perr := &cacheError{
				op:  "decoding data with pipe",
				key: key,
				err: err,
			}
			c.error(perr)
			return nil, perr
		}
	}
	return b, nil
}

// Delete removes the key from the cache. An error is returned only
// if the item was found but couldn't be deleted. Deleting a non-existant
// item is always successful.
func (c *Cache) Delete(key string) error {
	c.numQueries++
	err := c.driver.Delete(c.backendKey(key))
	if err != nil {
		derr := &cacheError{
			op:  "deleting",
			key: key,
			err: err,
		}
		c.error(derr)
		return derr
	}
	return nil
}

// NumQueries returns the number of queries made to this
// cache since it was initialized.
func (c *Cache) NumQueries() int {
	return c.numQueries
}

// Close closes the cache connection. If you're using a cache
// using mux.Context helper methods, the cache will be closed
// for you.
func (c *Cache) Close() error {
	return c.driver.Close()
}

// Connection returns a interface{} wrapping the native connection
// type for the cache client (e.g. a memcache or redis connection).
// Some drivers might return a nil connection (like the fs or the
// dummy driver).
func (c *Cache) Connection() interface{} {
	return c.driver.Connection()
}

func (c *Cache) debugf(format string, arg ...interface{}) {
	if c.Logger != nil {
		c.Logger.Debugf(format, arg...)
	}
}

func (c *Cache) warningf(format string, arg ...interface{}) {
	if c.Logger != nil {
		c.Logger.Warningf(format, arg...)
	}
}

func (c *Cache) error(err *cacheError) {
	if c.Logger != nil {
		c.Logger.Error(err)
	}
}

func newConfig(conf *config.URL) (*Cache, error) {
	if conf.Options == nil {
		conf.Options = config.Options{}
	}
	cache := &Cache{
		Logger: log.Std,
	}

	if codecName := conf.Get("codec"); codecName != "" {
		cache.codec = codec.Get(codecName)
		if cache.codec == nil {
			if imp := codec.RequiredImport(codecName); imp != "" {
				return nil, fmt.Errorf("please import %q to use the codec %q", imp, codecName)
			}
			return nil, fmt.Errorf("unknown codec %q, maybe you forgot an import?", codecName)
		}
	} else {
		cache.codec = codec.Get("gob")
	}

	cache.prefix = conf.Get("prefix")
	cache.prefixLen = len(cache.prefix)
	if pipeName := conf.Get("pipe"); pipeName != "" {
		cache.pipe = pipe.Get(pipeName)
		if cache.pipe == nil {
			if imp := pipe.RequiredImport(pipeName); imp != "" {
				return nil, fmt.Errorf("please import %q to use the pipe %q", imp, pipeName)
			}
			return nil, fmt.Errorf("unknown pipe %q, maybe you forgot an import?", pipeName)
		}
	}
	var opener driver.Opener
	if conf.Scheme != "" {
		opener = driver.Get(conf.Scheme)
		if opener == nil {
			if imp := imports[conf.Scheme]; imp != "" {
				return nil, fmt.Errorf("please import %q to use the cache driver %q", imp, conf.Scheme)
			}
			return nil, fmt.Errorf("unknown cache driver %q, maybe you forgot an import?", conf.Scheme)
		}
	} else {
		opener = driver.Get("dummy")
	}
	var err error
	if cache.driver, err = opener(conf.Value, conf.Options); err != nil {
		return nil, err
	}
	return cache, nil
}

// New returns a new cache instance, using the given
// configuration URL. If the configuration is nil, a
// dummy cache is returned, which always returns that the
// object does not exists in the cache.
func New(url *config.URL) (*Cache, error) {
	if url == nil {
		// Use dummy cache
		url = &config.URL{}
	}
	return newConfig(url)
}
<commit_msg>Small change in comment<commit_after>package cache

import (
	"errors"
	"fmt"
	"gnd.la/cache/driver"
	"gnd.la/config"
	"gnd.la/encoding/codec"
	"gnd.la/encoding/pipe"
	"gnd.la/log"
	"reflect"
	"strings"
)

var (
	ErrNotFound = errors.New("item not found in cache")
	imports     = map[string]string{
		"memcache": "gnd.la/cache/driver/memcache",
		"redis":    "gnd.la/cache/driver/redis",
	}
)

type Cache struct {
	// The Logger to log debug messages and, more importantly, errors.
	// New() initialies the log.Logger to log.Std.
	Logger     *log.Logger
	prefix     string
	prefixLen  int
	driver     driver.Driver
	codec      *codec.Codec
	pipe       *pipe.Pipe
	numQueries int
}

func (c *Cache) backendKey(key string) string {
	if c.prefixLen > 0 {
		return c.prefix + key
	}
	return key
}

func (c *Cache) frontendKey(key string) string {
	return key[c.prefixLen:]
}

// Set stores the given object in the cache associated with the
// given key. Timeout is the number of seconds until the item
// expires. If the timeout is 0, the item never expires, but
// might be only purged from cache when running out of space.
func (c *Cache) Set(key string, object interface{}, timeout int) error {
	b, err := c.codec.Encode(object)
	if err != nil {
		eerr := &cacheError{
			op:    "encoding object",
			key:   key,
			codec: true,
			err:   err,
		}
		c.error(eerr)
		return eerr
	}
	return c.SetBytes(key, b, timeout)
}

// Get retrieves the requested item from the cache and decodes it
// into the passed interface{}, which must be addressable.
// If the item is not found in the cache, ErrNotFound is returned.
// Other errors mean that either there was a problem communicating
// with the cache or the item could not be decoded.
func (c *Cache) Get(key string, obj interface{}) error {
	b, err := c.GetBytes(key)
	if err != nil {
		return err
	}
	cerr := c.codec.Decode(b, obj)
	if cerr != nil {
		derr := &cacheError{
			op:    "decoding object",
			key:   key,
			codec: true,
			err:   cerr,
		}
		c.error(derr)
		return derr
	}
	return nil
}

// GetMulti returns several objects as a map[string]interface{}
// in only one roundtrip to the cache. The obj parameter is used only
// to initialize the objects before they're passed to the codec for decoding,
// since not all codecs include the object type in its serialization (e.g. JSON).
// If you're using a codec which encodes the object type (e.g. Gob) or you want
// to decode the objects into a raw interface{} you might pass nil as the second
// parameter to this function.
func (c *Cache) GetMulti(keys []string, obj interface{}) (map[string]interface{}, error) {
	c.numQueries++
	if c.prefixLen > 0 {
		k := make([]string, len(keys))
		for ii, v := range keys {
			k[ii] = c.backendKey(v)
		}
		keys = k
	}
	data, err := c.driver.GetMulti(keys)
	if err != nil {
		gerr := &cacheError{
			op:  "getting multiple keys",
			key: strings.Join(keys, ", "),
			err: err,
		}
		c.error(gerr)
		return nil, gerr
	}
	objects := make(map[string]interface{}, len(data))
	var typ *reflect.Type
	if obj != nil {
		t := reflect.TypeOf(obj)
		for t.Kind() == reflect.Ptr {
			t = t.Elem()
		}
		typ = &t
	}
	for k, v := range data {
		var object interface{}
		if typ != nil {
			p := reflect.New(*typ)
			err = c.codec.Decode(v, p.Interface())
			object = p.Elem().Interface()
		} else {
			err = c.codec.Decode(v, &object)
		}
		if err != nil {
			derr := &cacheError{
				op:  "decoding object",
				key: k,
				err: err,
			}
			c.error(derr)
			return nil, derr
		}
		objects[c.frontendKey(k)] = object
	}
	return objects, nil
}

// SetBytes stores the given byte array assocciated with
// the given key. See the documentation for Set for an
// explanation of the timeout parameter
func (c *Cache) SetBytes(key string, b []byte, timeout int) error {
	c.numQueries++
	if c.pipe != nil {
		var err error
		b, err = c.pipe.Encode(b)
		if err != nil {
			perr := &cacheError{
				op:  "encoding data with pipe",
				key: key,
				err: err,
			}
			c.error(perr)
			return perr
		}
	}
	k := c.backendKey(key)
	err := c.driver.Set(k, b, timeout)
	if err != nil {
		serr := &cacheError{
			op:  "setting key",
			key: key,
			err: err,
		}
		c.error(serr)
		return serr
	}
	c.debugf("Set key %s (%d bytes), expiring in %d", k, len(b), timeout)
	return nil
}

// GetBytes returns the byte array assocciated with the given key
func (c *Cache) GetBytes(key string) ([]byte, error) {
	c.numQueries++
	b, err := c.driver.Get(c.backendKey(key))
	if err != nil {
		gerr := &cacheError{
			op:  "getting key",
			key: key,
			err: err,
		}
		c.error(gerr)
		return nil, gerr
	}
	if b == nil {
		return nil, ErrNotFound
	}
	if c.pipe != nil {
		b, err = c.pipe.Decode(b)
		if err != nil {
			perr := &cacheError{
				op:  "decoding data with pipe",
				key: key,
				err: err,
			}
			c.error(perr)
			return nil, perr
		}
	}
	return b, nil
}

// Delete removes the key from the cache. An error is returned only
// if the item was found but couldn't be deleted. Deleting a non-existant
// item is always successful.
func (c *Cache) Delete(key string) error {
	c.numQueries++
	err := c.driver.Delete(c.backendKey(key))
	if err != nil {
		derr := &cacheError{
			op:  "deleting",
			key: key,
			err: err,
		}
		c.error(derr)
		return derr
	}
	return nil
}

// NumQueries returns the number of queries made to this
// cache since it was initialized.
func (c *Cache) NumQueries() int {
	return c.numQueries
}

// Close closes the cache connection. If you're using a cache
// using mux.Context helper methods, the cache will be closed
// for you.
func (c *Cache) Close() error {
	return c.driver.Close()
}

// Connection returns a interface{} wrapping the native connection
// type for the cache client (e.g. a memcache or redis connection).
// Some drivers might return a nil connection (like the fs or the
// dummy driver).
func (c *Cache) Connection() interface{} {
	return c.driver.Connection()
}

func (c *Cache) debugf(format string, arg ...interface{}) {
	if c.Logger != nil {
		c.Logger.Debugf(format, arg...)
	}
}

func (c *Cache) warningf(format string, arg ...interface{}) {
	if c.Logger != nil {
		c.Logger.Warningf(format, arg...)
	}
}

func (c *Cache) error(err *cacheError) {
	if c.Logger != nil {
		c.Logger.Error(err)
	}
}

func newConfig(conf *config.URL) (*Cache, error) {
	if conf.Options == nil {
		conf.Options = config.Options{}
	}
	cache := &Cache{
		Logger: log.Std,
	}

	if codecName := conf.Get("codec"); codecName != "" {
		cache.codec = codec.Get(codecName)
		if cache.codec == nil {
			if imp := codec.RequiredImport(codecName); imp != "" {
				return nil, fmt.Errorf("please import %q to use the codec %q", imp, codecName)
			}
			return nil, fmt.Errorf("unknown codec %q, maybe you forgot an import?", codecName)
		}
	} else {
		cache.codec = codec.Get("gob")
	}

	cache.prefix = conf.Get("prefix")
	cache.prefixLen = len(cache.prefix)
	if pipeName := conf.Get("pipe"); pipeName != "" {
		cache.pipe = pipe.Get(pipeName)
		if cache.pipe == nil {
			if imp := pipe.RequiredImport(pipeName); imp != "" {
				return nil, fmt.Errorf("please import %q to use the pipe %q", imp, pipeName)
			}
			return nil, fmt.Errorf("unknown pipe %q, maybe you forgot an import?", pipeName)
		}
	}
	var opener driver.Opener
	if conf.Scheme != "" {
		opener = driver.Get(conf.Scheme)
		if opener == nil {
			if imp := imports[conf.Scheme]; imp != "" {
				return nil, fmt.Errorf("please import %q to use the cache driver %q", imp, conf.Scheme)
			}
			return nil, fmt.Errorf("unknown cache driver %q, maybe you forgot an import?", conf.Scheme)
		}
	} else {
		opener = driver.Get("dummy")
	}
	var err error
	if cache.driver, err = opener(conf.Value, conf.Options); err != nil {
		return nil, err
	}
	return cache, nil
}

// New returns a new cache instance, using the given
// configuration URL. If the configuration is nil, a
// dummy cache is returned, which always returns that the
// object does not exists in the cache.
func New(url *config.URL) (*Cache, error) {
	if url == nil {
		// Use dummy cache
		url = &config.URL{}
	}
	return newConfig(url)
}
<|endoftext|>


<commit_before>package cache

import (
	"encoding/json"
	"time"

	"github.com/gomodule/redigo/redis"
)

//Redis redis cache
type Redis struct {
	conn *redis.Pool
}

//RedisOpts redis 连接属性
type RedisOpts struct {
	Host        string `yml:"host" json:"host"`
	Password    string `yml:"password" json:"password"`
	Database    int    `yml:"database" json:"database"`
	MaxIdle     int    `yml:"max_idle" json:"max_idle"`
	MaxActive   int    `yml:"max_active" json:"max_active"`
	IdleTimeout int    `yml:"idle_timeout" json:"idle_timeout"` //second
}

//NewRedis 实例化
func NewRedis(opts *RedisOpts) *Redis {
	pool := &redis.Pool{
		MaxActive:   opts.MaxActive,
		MaxIdle:     opts.MaxIdle,
		IdleTimeout: time.Second * time.Duration(opts.IdleTimeout),
		Dial: func() (redis.Conn, error) {
			return redis.Dial("tcp", opts.Host,
				redis.DialDatabase(opts.Database),
				redis.DialPassword(opts.Password),
			)
		},
		TestOnBorrow: func(conn redis.Conn, t time.Time) error {
			if time.Since(t) < time.Minute {
				return nil
			}
			_, err := conn.Do("PING")
			return err
		},
	}
	return &Redis{pool}
}

//SetConn 设置conn
func (r *Redis) SetConn(conn *redis.Pool) {
	r.conn = conn
}

//Get 获取一个值
func (r *Redis) Get(key string) interface{} {
	conn := r.conn.Get()
	defer conn.Close()

	var data []byte
	var err error
	if data, err = redis.Bytes(conn.Do("GET", key)); err != nil {
		return nil
	}
	var reply interface{}
	if err = json.Unmarshal(data, &reply); err != nil {
		return nil
	}

	return reply
}

//Set 设置一个值
func (r *Redis) Set(key string, val interface{}, timeout time.Duration) (err error) {
	conn := r.conn.Get()
	defer conn.Close()

	var data []byte
	if data, err = json.Marshal(val); err != nil {
		return
	}

	_, err = conn.Do("SETEX", key, int64(timeout/time.Second), data)

	return
}

//IsExist 判断key是否存在
func (r *Redis) IsExist(key string) bool {
	conn := r.conn.Get()
	defer conn.Close()

	a, _ := conn.Do("EXISTS", key)
	i := a.(int64)
	if i > 0 {
		return true
	}
	return false
}

//Delete 删除
func (r *Redis) Delete(key string) error {
	conn := r.conn.Get()
	defer conn.Close()

	if _, err := conn.Do("DEL", key); err != nil {
		return err
	}

	return nil
}
<commit_msg>增加SetRedisPool方法 (#265)<commit_after>package cache

import (
	"encoding/json"
	"time"

	"github.com/gomodule/redigo/redis"
)

//Redis redis cache
type Redis struct {
	conn *redis.Pool
}

//RedisOpts redis 连接属性
type RedisOpts struct {
	Host        string `yml:"host" json:"host"`
	Password    string `yml:"password" json:"password"`
	Database    int    `yml:"database" json:"database"`
	MaxIdle     int    `yml:"max_idle" json:"max_idle"`
	MaxActive   int    `yml:"max_active" json:"max_active"`
	IdleTimeout int    `yml:"idle_timeout" json:"idle_timeout"` //second
}

//NewRedis 实例化
func NewRedis(opts *RedisOpts) *Redis {
	pool := &redis.Pool{
		MaxActive:   opts.MaxActive,
		MaxIdle:     opts.MaxIdle,
		IdleTimeout: time.Second * time.Duration(opts.IdleTimeout),
		Dial: func() (redis.Conn, error) {
			return redis.Dial("tcp", opts.Host,
				redis.DialDatabase(opts.Database),
				redis.DialPassword(opts.Password),
			)
		},
		TestOnBorrow: func(conn redis.Conn, t time.Time) error {
			if time.Since(t) < time.Minute {
				return nil
			}
			_, err := conn.Do("PING")
			return err
		},
	}
	return &Redis{pool}
}

//SetRedisPool 设置redis连接池
func (r *Redis) SetRedisPool(pool *redis.Pool) {
	r.conn = pool
}

//SetConn 设置conn
func (r *Redis) SetConn(conn *redis.Pool) {
	r.conn = conn
}

//Get 获取一个值
func (r *Redis) Get(key string) interface{} {
	conn := r.conn.Get()
	defer conn.Close()

	var data []byte
	var err error
	if data, err = redis.Bytes(conn.Do("GET", key)); err != nil {
		return nil
	}
	var reply interface{}
	if err = json.Unmarshal(data, &reply); err != nil {
		return nil
	}

	return reply
}

//Set 设置一个值
func (r *Redis) Set(key string, val interface{}, timeout time.Duration) (err error) {
	conn := r.conn.Get()
	defer conn.Close()

	var data []byte
	if data, err = json.Marshal(val); err != nil {
		return
	}

	_, err = conn.Do("SETEX", key, int64(timeout/time.Second), data)

	return
}

//IsExist 判断key是否存在
func (r *Redis) IsExist(key string) bool {
	conn := r.conn.Get()
	defer conn.Close()

	a, _ := conn.Do("EXISTS", key)
	i := a.(int64)
	if i > 0 {
		return true
	}
	return false
}

//Delete 删除
func (r *Redis) Delete(key string) error {
	conn := r.conn.Get()
	defer conn.Close()

	if _, err := conn.Do("DEL", key); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package checkmysql

import (
	"fmt"
	"os"
	"strings"

	"github.com/mackerelio/checkers"
	"github.com/ziutek/mymysql/mysql"
	// MySQL Driver
	_ "github.com/ziutek/mymysql/native"
)

type mysqlSetting struct {
	Host string `short:"H" long:"host" default:"localhost" description:"Hostname"`
	Port string `short:"p" long:"port" default:"3306" description:"Port"`
	User string `short:"u" long:"user" default:"root" description:"Username"`
	Pass string `short:"P" long:"password" default:"" description:"Password" env:"MYSQL_PASSWORD"`
}

var commands = map[string](func([]string) *checkers.Checker){
	"replication": checkReplication,
	"connection":  checkConnection,
	"uptime":      checkUptime,
	"readonly":    checkReadOnly,
}

func separateSub(argv []string) (string, []string) {
	if len(argv) == 0 || strings.HasPrefix(argv[0], "-") {
		return "", argv
	}
	return argv[0], argv[1:]
}

// Do the plugin
func Do() {
	subCmd, argv := separateSub(os.Args[1:])
	fn, ok := commands[subCmd]
	if !ok {
		fmt.Println(`Usage:
  check-mysql [subcommand] [OPTIONS]

SubCommands:`)
		for k := range commands {
			fmt.Printf("  %s\n", k)
		}
		os.Exit(1)
	}
	ckr := fn(argv)
	ckr.Name = fmt.Sprintf("MySQL %s", strings.ToUpper(string(subCmd[0]))+subCmd[1:])
	ckr.Exit()
}

func newMySQL(m mysqlSetting) mysql.Conn {
	target := fmt.Sprintf("%s:%s", m.Host, m.Port)
	return mysql.New("tcp", "", target, m.User, m.Pass, "")
}
<commit_msg>add --socket option for check-mysql<commit_after>package checkmysql

import (
	"fmt"
	"os"
	"strings"

	"github.com/mackerelio/checkers"
	"github.com/ziutek/mymysql/mysql"
	// MySQL Driver
	_ "github.com/ziutek/mymysql/native"
)

type mysqlSetting struct {
	Host   string `short:"H" long:"host" default:"localhost" description:"Hostname"`
	Port   string `short:"p" long:"port" default:"3306" description:"Port"`
	Socket string `short:"S" long:"socket" default:"" description:"Path to unix socket"`
	User   string `short:"u" long:"user" default:"root" description:"Username"`
	Pass   string `short:"P" long:"password" default:"" description:"Password" env:"MYSQL_PASSWORD"`
}

var commands = map[string](func([]string) *checkers.Checker){
	"replication": checkReplication,
	"connection":  checkConnection,
	"uptime":      checkUptime,
	"readonly":    checkReadOnly,
}

func separateSub(argv []string) (string, []string) {
	if len(argv) == 0 || strings.HasPrefix(argv[0], "-") {
		return "", argv
	}
	return argv[0], argv[1:]
}

// Do the plugin
func Do() {
	subCmd, argv := separateSub(os.Args[1:])
	fn, ok := commands[subCmd]
	if !ok {
		fmt.Println(`Usage:
  check-mysql [subcommand] [OPTIONS]

SubCommands:`)
		for k := range commands {
			fmt.Printf("  %s\n", k)
		}
		os.Exit(1)
	}
	ckr := fn(argv)
	ckr.Name = fmt.Sprintf("MySQL %s", strings.ToUpper(string(subCmd[0]))+subCmd[1:])
	ckr.Exit()
}

func newMySQL(m mysqlSetting) mysql.Conn {
	proto := "tcp"
	target := fmt.Sprintf("%s:%s", m.Host, m.Port)
	if m.Socket != "" {
		proto = "unix"
		target = m.Socket
	}
	return mysql.New(proto, "", target, m.User, m.Pass, "")
}
<|endoftext|>


<commit_before>package hive

import (
	"sync"
	"time"
)

// manager model
type HiveManager struct {
	FreeWorkers  chan *HiveWorker
	Tasks        chan interface{}
	Done         chan bool
	TimeProcess  chan int64
	LastProcess  int64
	TotalTimeOut int64
}

// worker model
type HiveWorker struct {
	WorkerId    int
	TimeProcess chan int64
	FreeWorkers chan *HiveWorker
	Context     *Hive
}

// initiate new manager
func NewHiveManager(numWorkers int) HiveManager {
	var totaltimeout int64 = 10

	m := HiveManager{}
	m.FreeWorkers = make(chan *HiveWorker, numWorkers)
	m.Tasks = make(chan interface{})
	m.TimeProcess = make(chan int64)
	m.TotalTimeOut = totaltimeout
	m.LastProcess = time.Now().Unix()
	m.Done = make(chan bool, 1)
	return m
}

// do monitoring worker thats free or not
func (m *HiveManager) DoMonitor(wg *sync.WaitGroup) {
	for {
		select {
		case task := <-m.Tasks:
			wg.Add(1)
			go m.AssignTask(task, wg)
		case result := <-m.TimeProcess:
			wg.Add(1)
			go m.InProgress(result, wg)
		case <-m.Done:
			m.Done <- true
			return
		}
	}
	wg.Wait()
}

// assign task to free worker
func (m *HiveManager) AssignTask(task interface{}, wg *sync.WaitGroup) {
	defer wg.Done()
	select {
	case worker := <-m.FreeWorkers:
		wg.Add(1)
		go worker.Work(task, wg)
	case isDone := <-m.Done:
		m.Done <- isDone
		return
	}
}

// check if a task still in progress to wait it till finish
func (m *HiveManager) InProgress(result int64, wg *sync.WaitGroup) {
	defer wg.Done()
	m.LastProcess = int64(result)
}

// set the timeout to waiting for tasks execution
func (m *HiveManager) Timeout(seconds int, wg *sync.WaitGroup) {
	defer wg.Done()
	for {
		if time.Now().Unix()-m.LastProcess > int64(seconds) {
			m.Done <- true
			return
		} else {
			time.Sleep(time.Millisecond)
		}
	}
}

func (m *HiveManager) EndWorker() {
	for {
		select {
		case worker := <-m.FreeWorkers:
			worker.Context.Conn.Close()
		default:
			return
		}
	}
}

// do a task for worker
func (w *HiveWorker) Work(task interface{}, wg *sync.WaitGroup) {
	defer wg.Done()

	w.Context.fetch(task.(string))

	w.TimeProcess <- time.Now().Unix()
	w.FreeWorkers <- w
}
<commit_msg>Modify worker add log<commit_after>package hive

import (
	"log"
	"sync"
	"time"
)

// manager model
type HiveManager struct {
	FreeWorkers  chan *HiveWorker
	Tasks        chan interface{}
	Done         chan bool
	TimeProcess  chan int64
	LastProcess  int64
	TotalTimeOut int64
}

// worker model
type HiveWorker struct {
	WorkerId    int
	TimeProcess chan int64
	FreeWorkers chan *HiveWorker
	Context     *Hive
}

// initiate new manager
func NewHiveManager(numWorkers int) HiveManager {
	var totaltimeout int64 = 10

	m := HiveManager{}
	m.FreeWorkers = make(chan *HiveWorker, numWorkers)
	m.Tasks = make(chan interface{})
	m.TimeProcess = make(chan int64)
	m.TotalTimeOut = totaltimeout
	m.LastProcess = time.Now().Unix()
	m.Done = make(chan bool, 1)
	return m
}

// do monitoring worker thats free or not
func (m *HiveManager) DoMonitor(wg *sync.WaitGroup) {
	for {
		select {
		case task := <-m.Tasks:
			log.Println("Preparing do task", task.(string))
			wg.Add(1)
			go m.AssignTask(task, wg)
		case result := <-m.TimeProcess:
			wg.Add(1)
			go m.InProgress(result, wg)
		case <-m.Done:
			m.Done <- true
			m.EndWorker()
			return
		}
	}
	wg.Wait()
}

// assign task to free worker
func (m *HiveManager) AssignTask(task interface{}, wg *sync.WaitGroup) {
	defer wg.Done()
	select {
	case worker := <-m.FreeWorkers:
		log.Println("Assign task to worker", worker.WorkerId)
		wg.Add(1)
		go worker.Work(task, wg)
	case isDone := <-m.Done:
		m.Done <- isDone
		return
	}
}

// check if a task still in progress to wait it till finish
func (m *HiveManager) InProgress(result int64, wg *sync.WaitGroup) {
	defer wg.Done()
	m.LastProcess = int64(result)
}

// set the timeout to waiting for tasks execution
func (m *HiveManager) Timeout(seconds int, wg *sync.WaitGroup) {
	defer wg.Done()
	for {
		if time.Now().Unix()-m.LastProcess > int64(seconds) {
			m.Done <- true
			return
		} else {
			time.Sleep(time.Millisecond)
		}
	}
}

func (m *HiveManager) EndWorker() {
	for {
		select {
		case worker := <-m.FreeWorkers:
			worker.Context.Conn.Close()
		default:
			return
		}
	}
}

// do a task for worker
func (w *HiveWorker) Work(task interface{}, wg *sync.WaitGroup) {
	defer wg.Done()

	log.Println("Do task ", task.(string))
	w.Context.fetch(task.(string))

	w.TimeProcess <- time.Now().Unix()
	w.FreeWorkers <- w
}
<|endoftext|>


<commit_before>package hosts

import (
	"bufio"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
)

type Doer interface {
	Do(*http.Request) (*http.Response, error)
}

type Hosts struct {
	mu             sync.RWMutex
	dir            string
	fileName       string
	url            string
	stopCh         chan struct{}
	hosts          map[string]struct{}
	initialized    bool
	UpdateInterval time.Duration
	Parser         Parser
	Client         Doer
	Logger         *logrus.Logger
	AppName        string
	AppVersion     string
}

const (
	defaultUpdateInterval = 24 * time.Hour
	subDir                = "hosts"
)

func (h *Hosts) Init(u, cacheDir string) error {
	p, err := url.Parse(u)
	if err != nil {
		return err
	}

	dirName := path.Join(cacheDir, subDir)
	err = os.MkdirAll(dirName, 0700)
	if err != nil {
		return err
	}

	file := p.Path

	// strip leading '/'
	if len(file) > 0 && file[0] == '/' {
		file = file[1:]
	}

	// join host and path
	file = strings.Join([]string{p.Host, file}, "__")

	// replace '/' with '__'
	file = strings.Replace(file, "/", "__", -1)

	h.dir = dirName
	h.fileName = path.Join(dirName, file)
	h.url = u

	if h.UpdateInterval == 0 {
		h.UpdateInterval = defaultUpdateInterval
	}

	if h.Parser == nil {
		h.Parser = FileParser
	}

	if h.Logger == nil {
		h.Logger = logrus.New()
	}

	if h.Client == nil {
		h.Client = http.DefaultClient
	}

	return nil
}

func (h *Hosts) Start() (<-chan struct{}, <-chan error) {
	h.mu.Lock()
	defer h.mu.Unlock()

	if h.stopCh != nil {
		return nil, nil
	}

	h.stopCh = make(chan struct{})

	h.Logger.Debugf("starting hosts file %s", h.fileName)

	updatedCh := make(chan struct{})
	errCh := make(chan error)

	updater := func() {
		updated, err := h.Update()
		if err != nil {
			errCh <- err
		} else if updated {
			updatedCh <- struct{}{}
		}
	}

	go updater()

	ticker := time.NewTicker(h.UpdateInterval)
	defer ticker.Stop()

	go func() {
		for {
			select {
			case <-ticker.C:
				go updater()
			case <-h.stopCh:
				h.Logger.Debugf("stopped hosts file %s", h.fileName)
				close(h.stopCh)
				return
			}
		}
	}()

	return updatedCh, errCh
}

func (h *Hosts) Stop() {
	h.mu.Lock()
	defer h.mu.Unlock()

	if h.stopCh == nil {
		return
	}

	h.Logger.Debugf("stopping hosts file %s", h.fileName)
	h.stopCh <- struct{}{}
	<-h.stopCh
	h.stopCh = nil
}

func (h *Hosts) Update() (bool, error) {
	updated, err := h.doUpdate()
	if err != nil {
		return false, err
	}

	h.mu.RLock()
	initialized := h.initialized
	h.mu.RUnlock()

	if !updated && initialized {
		return false, nil
	}

	if err := h.parse(); err != nil {
		return false, err
	}

	return true, nil
}

func (h *Hosts) doUpdate() (updated bool, err error) {
	info, err := os.Stat(h.fileName)
	if err != nil && !os.IsNotExist(err) {
		return false, err
	}

	if !os.IsNotExist(err) {
		if info.IsDir() {
			return false, fmt.Errorf("%s is a directory", h.fileName)
		}

		if info.ModTime().After(time.Now().Add(-h.UpdateInterval)) {
			// file exists and does not need to be updated
			h.Logger.Debugf("%s does not need to be updated yet", h.fileName)
			return false, nil
		}

		// file exists and needs to be updated
	}

	h.Logger.Infof("updating %s", h.fileName)

	req, _ := http.NewRequest("GET", h.url, nil)
	req.Header.Set("User-Agent", fmt.Sprintf("%s/%s", h.AppName, h.AppVersion))
	h.debugRequestOut(req, false)

	res, err := h.Client.Do(req)
	if err != nil {
		return false, err
	}

	defer func() { _ = res.Body.Close() }()

	h.debugResponse(res, false)

	if res.StatusCode != http.StatusOK {
		err = NewErrStatusCode(res.StatusCode)
		h.Logger.
			WithError(err).
			WithFields(logrus.Fields{
				"fileName": h.fileName,
				"url":      h.url,
			}).
			Warn("error updating hosts")
		return false, err
	}

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

	// open the file and truncate it if it already exists
	f, err := os.Create(h.fileName)
	if err != nil {
		return false, nil
	}
	defer func() { _ = f.Close() }()

	_, err = io.Copy(f, res.Body)
	if err != nil {
		h.Logger.Warnf("error updating %s: %v", h.fileName, err)
		return false, err
	}

	h.Logger.Debugf("successfully updated %s", h.fileName)
	return true, nil
}

func (h *Hosts) Block(host string) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()

	_, found := h.hosts[host]
	return found
}

func (h *Hosts) parse() error {
	h.mu.Lock()
	defer h.mu.Unlock()

	h.Logger.Debugf("parsing %s", h.fileName)

	f, err := os.Open(h.fileName)
	if err != nil {
		return err
	}

	h.hosts = map[string]struct{}{}

	scanner := bufio.NewScanner(f)
	i := 0
	for scanner.Scan() {
		i++

		text := scanner.Text()
		text = stripComments(text)

		if len(text) == 0 {
			continue
		}

		host, valid := h.Parser.Parse(text)
		if !valid {
			continue
		}

		host = strings.ToLower(host)
		host = stripTrailing(host, ".")
		host = stripPort(host)

		if h.validateHost(i, host) {
			h.hosts[host] = struct{}{}
		}
	}

	if err := scanner.Err(); err != nil {
		return err
	}

	h.Logger.Debugf("parsed %d hosts in %s", len(h.hosts), h.fileName)

	h.initialized = true
	return nil
}

func stripComments(text string) string {
	i := strings.Index(text, "#")
	if i == 0 {
		return ""
	}

	if i != -1 {
		return text[:i-1]
	}

	return text
}

var portRe = regexp.MustCompile(`(\S+):\d+`)

func stripPort(text string) string {
	res := portRe.FindStringSubmatch(text)
	if len(res) < 2 {
		return text
	}
	return res[1]
}

func stripTrailing(text string, trail string) string {
	i := strings.LastIndex(text, trail)
	if i == -1 {
		return text
	}

	if len(text) != i+len(trail) {
		return text
	}

	return text[:i]
}

func (h *Hosts) validateHost(ln int, host string) bool {
	// the entire hostname (including the delimiting dots but not a trailing
	// dot) has a maximum of 253 ascii characters
	if len(host) > 253 {
		h.Logger.Warnf("%s:%d: hostname too long: %s", h.fileName, ln, host)
		return false
	}

	labels := strings.Split(host, ".")
	for _, label := range labels {
		// each label must be between 1 and 63 characters long
		l := len(label)

		if l == 0 {
			h.Logger.Warnf("%s:%d: empty label: %s", h.fileName, ln, host)
			return false
		}

		if l > 63 {
			h.Logger.Warnf("%s:%d: label too long: %s (%s)", h.fileName, ln, label, host)
			return false
		}

		// hostname labels may contain only the ASCII letters 'a' through 'z' (in a
		// case-insensitive manner), the digits '0' through '9', and the hyphen
		// ('-')
		for i := 0; i < l; i++ {
			c := label[i]

			if c >= 'a' && c <= 'z' {
				continue
			}

			// if c >= 'A' && c <= 'Z' {
			// 	continue
			// }

			if c >= '0' && c <= '9' {
				continue
			}

			if c == '_' {
				continue
			}

			// must not start or end with a hyphen
			if c == '-' && i > 0 && i < l-1 {
				continue
			}

			h.Logger.Warnf("%s:%d: invalid character \"%c\" (%s)", h.fileName, ln, c, host)
			return false
		}
	}

	return true
}
<commit_msg>clean up hosts updater<commit_after>package hosts

import (
	"bufio"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
)

type Doer interface {
	Do(*http.Request) (*http.Response, error)
}

type Hosts struct {
	mu                sync.RWMutex
	dir               string
	fileName          string
	url               string
	stopCh            chan struct{}
	hosts             map[string]struct{}
	initialized       bool
	UpdateInterval    time.Duration
	Parser            Parser
	Client            Doer
	Logger            *logrus.Logger
	AppName           string
	AppVersion        string
	NotifyUpdatedFunc func(file string, num int, updated bool, err error)
}

const (
	defaultUpdateInterval = 24 * time.Hour
	subDir                = "hosts"
)

func (h *Hosts) Init(u, cacheDir string) error {
	p, err := url.Parse(u)
	if err != nil {
		return err
	}

	dirName := path.Join(cacheDir, subDir)
	err = os.MkdirAll(dirName, 0700)
	if err != nil {
		return err
	}

	file := p.Path

	// strip leading '/'
	if len(file) > 0 && file[0] == '/' {
		file = file[1:]
	}

	// join host and path
	file = strings.Join([]string{p.Host, file}, "__")

	// replace '/' with '__'
	file = strings.Replace(file, "/", "__", -1)

	h.dir = dirName
	h.fileName = path.Join(dirName, file)
	h.url = u

	if h.UpdateInterval == 0 {
		h.UpdateInterval = defaultUpdateInterval
	}

	if h.Parser == nil {
		h.Parser = FileParser
	}

	if h.Logger == nil {
		h.Logger = logrus.New()
	}

	if h.Client == nil {
		h.Client = http.DefaultClient
	}

	return nil
}

func (h *Hosts) Start() {
	h.mu.Lock()
	defer h.mu.Unlock()

	if h.stopCh != nil {
		return
	}

	h.stopCh = make(chan struct{})

	h.Logger.Debugf("starting hosts file %s", h.fileName)

	go h.Update()

	ticker := time.NewTicker(h.UpdateInterval)
	defer ticker.Stop()

	go func() {
		for {
			select {
			case <-ticker.C:
				go h.Update()
			case <-h.stopCh:
				h.Logger.Debugf("stopped hosts file %s", h.fileName)
				close(h.stopCh)
				return
			}
		}
	}()
}

func (h *Hosts) Stop() {
	h.mu.Lock()
	defer h.mu.Unlock()

	if h.stopCh == nil {
		return
	}

	h.Logger.Debugf("stopping hosts file %s", h.fileName)
	h.stopCh <- struct{}{}
	<-h.stopCh
	h.stopCh = nil
}

func (h *Hosts) notify(updated bool, err error) {
	if (!updated && err == nil) || h.NotifyUpdatedFunc == nil {
		return
	}

	go h.NotifyUpdatedFunc(h.fileName, h.Len(), updated, err)
}

func (h *Hosts) Len() int {
	return len(h.hosts)
}

func (h *Hosts) Update() (updated bool, err error) {
	// wrap the defer in func so that it uses latest values of updated and err
	defer func() { h.notify(updated, err) }()

	updated, err = h.doUpdate()
	if err != nil {
		return
	}

	h.mu.RLock()
	initialized := h.initialized
	h.mu.RUnlock()

	if !updated && initialized {
		return
	}

	if err = h.parse(); err != nil {
		return
	}

	updated = true
	return
}

func (h *Hosts) doUpdate() (updated bool, err error) {
	info, err := os.Stat(h.fileName)
	if err != nil && !os.IsNotExist(err) {
		return false, err
	}

	if !os.IsNotExist(err) {
		if info.IsDir() {
			return false, fmt.Errorf("%s is a directory", h.fileName)
		}

		if info.ModTime().After(time.Now().Add(-h.UpdateInterval)) {
			// file exists and does not need to be updated
			h.Logger.Debugf("%s does not need to be updated yet", h.fileName)
			return false, nil
		}

		// file exists and needs to be updated
	}

	h.Logger.Infof("updating %s", h.fileName)

	req, _ := http.NewRequest("GET", h.url, nil)
	req.Header.Set("User-Agent", fmt.Sprintf("%s/%s", h.AppName, h.AppVersion))
	h.debugRequestOut(req, false)

	res, err := h.Client.Do(req)
	if err != nil {
		return false, err
	}

	defer func() { _ = res.Body.Close() }()

	h.debugResponse(res, false)

	if res.StatusCode != http.StatusOK {
		err = NewErrStatusCode(res.StatusCode)
		h.Logger.
			WithError(err).
			WithFields(logrus.Fields{
				"fileName": h.fileName,
				"url":      h.url,
			}).
			Warn("error updating hosts")
		return false, err
	}

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

	// open the file and truncate it if it already exists
	f, err := os.Create(h.fileName)
	if err != nil {
		return false, nil
	}
	defer func() { _ = f.Close() }()

	_, err = io.Copy(f, res.Body)
	if err != nil {
		h.Logger.Warnf("error updating %s: %v", h.fileName, err)
		return false, err
	}

	h.Logger.Debugf("successfully updated %s", h.fileName)
	return true, nil
}

func (h *Hosts) Block(host string) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()

	_, found := h.hosts[host]
	return found
}

func (h *Hosts) parse() error {
	h.mu.Lock()
	defer h.mu.Unlock()

	h.Logger.Debugf("parsing %s", h.fileName)

	f, err := os.Open(h.fileName)
	if err != nil {
		return err
	}

	h.hosts = map[string]struct{}{}

	scanner := bufio.NewScanner(f)
	i := 0
	for scanner.Scan() {
		i++

		text := scanner.Text()
		text = stripComments(text)

		if len(text) == 0 {
			continue
		}

		host, valid := h.Parser.Parse(text)
		if !valid {
			continue
		}

		host = strings.ToLower(host)
		host = stripTrailing(host, ".")
		host = stripPort(host)

		if h.validateHost(i, host) {
			h.hosts[host] = struct{}{}
		}
	}

	if err := scanner.Err(); err != nil {
		return err
	}

	h.Logger.Debugf("parsed %d hosts in %s", h.Len(), h.fileName)

	h.initialized = true
	return nil
}

func stripComments(text string) string {
	i := strings.Index(text, "#")
	if i == 0 {
		return ""
	}

	if i != -1 {
		return text[:i-1]
	}

	return text
}

var portRe = regexp.MustCompile(`(\S+):\d+`)

func stripPort(text string) string {
	res := portRe.FindStringSubmatch(text)
	if len(res) < 2 {
		return text
	}
	return res[1]
}

func stripTrailing(text string, trail string) string {
	i := strings.LastIndex(text, trail)
	if i == -1 {
		return text
	}

	if len(text) != i+len(trail) {
		return text
	}

	return text[:i]
}

func (h *Hosts) validateHost(ln int, host string) bool {
	// the entire hostname (including the delimiting dots but not a trailing
	// dot) has a maximum of 253 ascii characters
	if len(host) > 253 {
		h.Logger.Warnf("%s:%d: hostname too long: %s", h.fileName, ln, host)
		return false
	}

	labels := strings.Split(host, ".")
	for _, label := range labels {
		// each label must be between 1 and 63 characters long
		l := len(label)

		if l == 0 {
			h.Logger.Warnf("%s:%d: empty label: %s", h.fileName, ln, host)
			return false
		}

		if l > 63 {
			h.Logger.Warnf("%s:%d: label too long: %s (%s)", h.fileName, ln, label, host)
			return false
		}

		// hostname labels may contain only the ASCII letters 'a' through 'z' (in a
		// case-insensitive manner), the digits '0' through '9', and the hyphen
		// ('-')
		for i := 0; i < l; i++ {
			c := label[i]

			if c >= 'a' && c <= 'z' {
				continue
			}

			// if c >= 'A' && c <= 'Z' {
			// 	continue
			// }

			if c >= '0' && c <= '9' {
				continue
			}

			if c == '_' {
				continue
			}

			// must not start or end with a hyphen
			if c == '-' && i > 0 && i < l-1 {
				continue
			}

			h.Logger.Warnf("%s:%d: invalid character \"%c\" (%s)", h.fileName, ln, c, host)
			return false
		}
	}

	return true
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	daemon "github.com/tyranron/daemonigo"
	"net/http"
	"syscall"
	"time"
)

func init() {
	// Setting up daemon properties.
	daemon.AppName = "Graceful HTTP Server"
	daemon.PidFile = "graceful.pid"

	// Helper functions for printing.
	printStatusErr := func(e error) {
		fmt.Println("Checking status of " + daemon.AppName + " failed")
		fmt.Println("Details:", e.Error())
	}
	printFailed := func(e error) {
		fmt.Println("FAILED")
		fmt.Println("Details:", e.Error())
	}

	// Implementation of reload action.
	daemon.SetAction("reload", func() {
		isRunning, process, err := daemon.Status()
		if err != nil {
			printStatusErr(err)
			return
		}
		if !isRunning {
			fmt.Println(daemon.AppName + " is NOT running now")
			fmt.Printf("Starting %s...", daemon.AppName)
			if err := daemon.Start(1); err != nil {
				printFailed(err)
			} else {
				fmt.Println("OK")
			}
		} else {
			fmt.Printf("Reloading %s...", daemon.AppName)
			if err := process.Signal(syscall.SIGHUP); err != nil {
				printFailed(err)
				return
			}
			select {
			case <-func(prevPid int) chan bool {
				ch := make(chan bool)
				go func() {
					defer close(ch)
					for {
						time.Sleep(200 * time.Millisecond)
						switch isRunning, process, err := daemon.Status(); {
						case err != nil:
							printStatusErr(err)
							return
						case isRunning && (process.Pid != prevPid):
							fmt.Println("OK")
							return
						}
					}
				}()
				return ch
			}(process.Pid):
			case <-time.After(10 * time.Second):
				printStatusErr(fmt.Errorf("checking new process timed out, see application logs"))
			}
		}
	})

	// A simple program to test server during reloads.
	daemon.SetAction("test", func() {
		errs, oks := 0, 0
		for i := 0; i < 10000; i++ {
			if r, err := http.Get("http://127.0.0.1:8080/"); err != nil {
				print("E")
				errs++
			} else {
				print(".")
				oks++
				r.Body.Close()
			}
			time.Sleep(10 * time.Millisecond)
		}
		println("\n---------------------------")
		println("Succeed:", oks, "Errors:", errs)
	})
}
<commit_msg>graceful_http example: added -ms flag<commit_after>package main

import (
	"fmt"
	daemon "github.com/tyranron/daemonigo"
	"net/http"
	"syscall"
	"time"
	"flag"
)

func init() {
	// Setting up daemon properties.
	daemon.AppName = "Graceful HTTP Server"
	daemon.PidFile = "graceful.pid"

	// Helper functions for printing.
	printStatusErr := func(e error) {
		fmt.Println("Checking status of " + daemon.AppName + " failed")
		fmt.Println("Details:", e.Error())
	}
	printFailed := func(e error) {
		fmt.Println("FAILED")
		fmt.Println("Details:", e.Error())
	}

	// Implementation of reload action.
	daemon.SetAction("reload", func() {
		isRunning, process, err := daemon.Status()
		if err != nil {
			printStatusErr(err)
			return
		}
		if !isRunning {
			fmt.Println(daemon.AppName + " is NOT running now")
			fmt.Printf("Starting %s...", daemon.AppName)
			if err := daemon.Start(1); err != nil {
				printFailed(err)
			} else {
				fmt.Println("OK")
			}
		} else {
			fmt.Printf("Reloading %s...", daemon.AppName)
			if err := process.Signal(syscall.SIGHUP); err != nil {
				printFailed(err)
				return
			}
			select {
			case <-func(prevPid int) chan bool {
				ch := make(chan bool)
				go func() {
					defer close(ch)
					for {
						time.Sleep(200 * time.Millisecond)
						switch isRunning, process, err := daemon.Status(); {
						case err != nil:
							printStatusErr(err)
							return
						case isRunning && (process.Pid != prevPid):
							fmt.Println("OK")
							return
						}
					}
				}()
				return ch
			}(process.Pid):
			case <-time.After(10 * time.Second):
				printStatusErr(fmt.Errorf("checking new process timed out, see application logs"))
			}
		}
	})

	// A simple program to test server during reloads.
	var ms uint
	flag.UintVar(&ms, "ms", 10, "frequency of requests")
	daemon.SetAction("test", func() {
		if ms == 0 {
			ms = 10
		}
		errs, oks := 0, 0
		for i := 0; i < 10000; i++ {
			if r, err := http.Get("http://127.0.0.1:8080/"); err != nil {
				print("E")
				errs++
			} else {
				print(".")
				oks++
				r.Body.Close()
			}
			time.Sleep(time.Duration(ms) * time.Millisecond)
		}
		println("\n---------------------------")
		println("Succeed:", oks, "Errors:", errs)
	})
}
<|endoftext|>


<commit_before>// +build darwin dragonfly freebsd netbsd openbsd

// Copyright 2012 Google, Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

package bsdbpf

import (
	"github.com/google/gopacket"
	"golang.org/x/sys/unix"

	"fmt"
	"syscall"
	"time"
	"unsafe"
)

const wordSize = int(unsafe.Sizeof(uintptr(0)))

func bpfWordAlign(x int) int {
	return (((x) + (wordSize - 1)) &^ (wordSize - 1))
}

type TimedFrame struct {
	RawFrame  []byte
	Timestamp time.Time
}

type BPFSniffer struct {
	sniffDeviceName string
	bpfDeviceName string
	fd       int
	stopChan chan bool
	readChan chan TimedFrame
	readBuffer []byte
	readBufLen int
	lastReadLen int
	readBytesConsumed int
}

func NewBPFSniffer(sniffDeviceName, bpfDeviceName string, readBufLen int) *BPFSniffer {
	return &BPFSniffer{
		sniffDeviceName: sniffDeviceName,
		bpfDeviceName: bpfDeviceName,
		stopChan: make(chan bool, 0),
		readChan: make(chan TimedFrame, 0),
		readBufLen: readBufLen,
	}
}

func (b *BPFSniffer) Close() error {
	return syscall.Close(b.fd)
}

func (b *BPFSniffer) pickBpfDevice() {
	var err error
	for i := 0; i < 99; i++ {
		b.bpfDeviceName = fmt.Sprintf("/dev/bpf%d", i)
		b.fd, err = syscall.Open(b.bpfDeviceName, syscall.O_RDWR, 0)
		if err == nil {
			break
		}
	}
}

// Init is used to initialize a BPF device for promiscuous sniffing.
// It also starts a goroutine to continuously read frames.
func (b *BPFSniffer) Init() error {
	var err error
	enable := 1

	if b.bpfDeviceName == "" {
		b.pickBpfDevice()
	}

	// setup our read buffer
	if b.readBufLen == 0 {
		b.readBufLen, err = syscall.BpfBuflen(b.fd)
		if err != nil {
			panic(err)
		}
	} else {
		b.readBufLen, err = syscall.SetBpfBuflen(b.fd, b.readBufLen)
		if err != nil {
			panic(err)
		}
		var actualBufLen int
		actualBufLen, err = syscall.BpfBuflen(b.fd)
		if err != nil {
			panic(err)
		}
		fmt.Printf("actual buf len %d\n", actualBufLen)
	}
	b.readBuffer = make([]byte, b.readBufLen)

	err = syscall.SetBpfInterface(b.fd, b.sniffDeviceName)
	if err != nil {
		return err
	}

	// turning Immediate mode off should make the snffer
	// block when no packets are available.
	err = syscall.SetBpfImmediate(b.fd, enable)
	if err != nil {
		return err
	}

	// preserves the link level source address...
	// higher level protocol analyzers will not need this
	err = syscall.SetBpfHeadercmpl(b.fd, enable)
	if err != nil {
		return err
	}

	// forces the interface into promiscuous mode
	err = syscall.SetBpfPromisc(b.fd, enable)
	if err != nil {
		return err
	}

	return nil
}

func (b *BPFSniffer) ReadPacketData() ([]byte, gopacket.CaptureInfo, error) {
	var err error
	if b.readBytesConsumed >= b.lastReadLen {
		b.lastReadLen, err = syscall.Read(b.fd, b.readBuffer)
		if err != nil {
			return nil, gopacket.CaptureInfo{}, err
		}
		b.readBytesConsumed = 0
	}
	hdr := (*unix.BpfHdr)(unsafe.Pointer(&b.readBuffer[b.readBytesConsumed]))
	frameStart := b.readBytesConsumed + int(hdr.Hdrlen)
	b.readBytesConsumed += bpfWordAlign(int(hdr.Hdrlen) + int(hdr.Caplen))
	rawFrame := b.readBuffer[frameStart : frameStart+int(hdr.Caplen)]
	captureInfo := gopacket.CaptureInfo{
		Timestamp:     time.Unix(int64(hdr.Tstamp.Sec), int64(hdr.Tstamp.Usec)*1000),
		CaptureLength: len(rawFrame),
		Length:        len(rawFrame),
	}
	return rawFrame, captureInfo, nil
}
<commit_msg>Remove superfluous test code<commit_after>// +build darwin dragonfly freebsd netbsd openbsd

// Copyright 2012 Google, Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

package bsdbpf

import (
	"github.com/google/gopacket"
	"golang.org/x/sys/unix"

	"fmt"
	"syscall"
	"time"
	"unsafe"
)

const wordSize = int(unsafe.Sizeof(uintptr(0)))

func bpfWordAlign(x int) int {
	return (((x) + (wordSize - 1)) &^ (wordSize - 1))
}

type TimedFrame struct {
	RawFrame  []byte
	Timestamp time.Time
}

type BPFSniffer struct {
	sniffDeviceName string
	bpfDeviceName string
	fd       int
	stopChan chan bool
	readChan chan TimedFrame
	readBuffer []byte
	readBufLen int
	lastReadLen int
	readBytesConsumed int
}

func NewBPFSniffer(sniffDeviceName, bpfDeviceName string, readBufLen int) *BPFSniffer {
	return &BPFSniffer{
		sniffDeviceName: sniffDeviceName,
		bpfDeviceName: bpfDeviceName,
		stopChan: make(chan bool, 0),
		readChan: make(chan TimedFrame, 0),
		readBufLen: readBufLen,
	}
}

func (b *BPFSniffer) Close() error {
	return syscall.Close(b.fd)
}

func (b *BPFSniffer) pickBpfDevice() {
	var err error
	for i := 0; i < 99; i++ {
		b.bpfDeviceName = fmt.Sprintf("/dev/bpf%d", i)
		b.fd, err = syscall.Open(b.bpfDeviceName, syscall.O_RDWR, 0)
		if err == nil {
			break
		}
	}
}

// Init is used to initialize a BPF device for promiscuous sniffing.
// It also starts a goroutine to continuously read frames.
func (b *BPFSniffer) Init() error {
	var err error
	enable := 1

	if b.bpfDeviceName == "" {
		b.pickBpfDevice()
	}

	// setup our read buffer
	if b.readBufLen == 0 {
		b.readBufLen, err = syscall.BpfBuflen(b.fd)
		if err != nil {
			panic(err)
		}
	} else {
		b.readBufLen, err = syscall.SetBpfBuflen(b.fd, b.readBufLen)
		if err != nil {
			panic(err)
		}
	}
	b.readBuffer = make([]byte, b.readBufLen)

	err = syscall.SetBpfInterface(b.fd, b.sniffDeviceName)
	if err != nil {
		return err
	}

	// turning Immediate mode off should make the snffer
	// block when no packets are available.
	err = syscall.SetBpfImmediate(b.fd, enable)
	if err != nil {
		return err
	}

	// preserves the link level source address...
	// higher level protocol analyzers will not need this
	err = syscall.SetBpfHeadercmpl(b.fd, enable)
	if err != nil {
		return err
	}

	// forces the interface into promiscuous mode
	err = syscall.SetBpfPromisc(b.fd, enable)
	if err != nil {
		return err
	}

	return nil
}

func (b *BPFSniffer) ReadPacketData() ([]byte, gopacket.CaptureInfo, error) {
	var err error
	if b.readBytesConsumed >= b.lastReadLen {
		b.lastReadLen, err = syscall.Read(b.fd, b.readBuffer)
		if err != nil {
			return nil, gopacket.CaptureInfo{}, err
		}
		b.readBytesConsumed = 0
	}
	hdr := (*unix.BpfHdr)(unsafe.Pointer(&b.readBuffer[b.readBytesConsumed]))
	frameStart := b.readBytesConsumed + int(hdr.Hdrlen)
	b.readBytesConsumed += bpfWordAlign(int(hdr.Hdrlen) + int(hdr.Caplen))
	rawFrame := b.readBuffer[frameStart : frameStart+int(hdr.Caplen)]
	captureInfo := gopacket.CaptureInfo{
		Timestamp:     time.Unix(int64(hdr.Tstamp.Sec), int64(hdr.Tstamp.Usec)*1000),
		CaptureLength: len(rawFrame),
		Length:        len(rawFrame),
	}
	return rawFrame, captureInfo, nil
}
<|endoftext|>


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

// Package oauth2 contains Negroni middleware to provide
// user login via an OAuth 2.0 backend.

package oauth2

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"time"

	"github.com/codegangsta/negroni"
	sessions "github.com/goincremental/negroni-sessions"
	"golang.org/x/oauth2"
)

const (
	codeRedirect = 302
	keyToken     = "oauth2_token"
	keyNextPage  = "next"
	keyState     = "state"
)

var (
	// PathLogin sets the path to handle OAuth 2.0 logins.
	PathLogin = "/login"
	// PathLogout sets to handle OAuth 2.0 logouts.
	PathLogout = "/logout"
	// PathCallback sets the path to handle callback from OAuth 2.0 backend
	// to exchange credentials.
	PathCallback = "/oauth2callback"
	// PathError sets the path to handle error cases.
	PathError = "/oauth2error"
)

type Config oauth2.Config

// Tokens Represents a container that contains
// user's OAuth 2.0 access and refresh tokens.
type Tokens interface {
	Access() string
	Refresh() string
	Valid() bool
	ExpiryTime() time.Time
	ExtraData(string) interface{}
	Get() Token
}

type Token oauth2.Token

type token struct {
	oauth2.Token
}

func (t *token) ExtraData(key string) interface{} {
	return t.Extra(key)
}

// Returns the access token.
func (t *token) Access() string {
	return t.AccessToken
}

// Returns the refresh token.
func (t *token) Refresh() string {
	return t.RefreshToken
}

// Returns whether the access token is
// expired or not.
func (t *token) Valid() bool {
	if t == nil {
		return true
	}
	return t.Token.Valid()
}

// Returns the expiry time of the user's
// access token.
func (t *token) ExpiryTime() time.Time {
	return t.Expiry
}

// String returns the string representation of the token.
func (t *token) String() string {
	return fmt.Sprintf("tokens: %v", t)
}

// Returns oauth2.Token.
func (t *token) Get() Token {
	return (Token)(t.Token)
}

// Returns a new Google OAuth 2.0 backend endpoint.
func Google(config *Config) negroni.Handler {
	authUrl := "https://accounts.google.com/o/oauth2/auth"
	tokenUrl := "https://accounts.google.com/o/oauth2/token"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

// Returns a new Github OAuth 2.0 backend endpoint.
func Github(config *Config) negroni.Handler {
	authUrl := "https://github.com/login/oauth/authorize"
	tokenUrl := "https://github.com/login/oauth/access_token"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

func Facebook(config *Config) negroni.Handler {
	authUrl := "https://www.facebook.com/dialog/oauth"
	tokenUrl := "https://graph.facebook.com/oauth/access_token"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

func LinkedIn(config *Config) negroni.Handler {
	authUrl := "https://www.linkedin.com/uas/oauth2/authorization"
	tokenUrl := "https://www.linkedin.com/uas/oauth2/accessToken"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

// Returns a generic OAuth 2.0 backend endpoint.
func NewOAuth2Provider(config *Config, authUrl, tokenUrl string) negroni.HandlerFunc {
	c := &oauth2.Config{
		ClientID:     config.ClientID,
		ClientSecret: config.ClientSecret,
		Scopes:       config.Scopes,
		RedirectURL:  config.RedirectURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  authUrl,
			TokenURL: tokenUrl,
		},
	}

	return func(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
		s := sessions.GetSession(r)

		if r.Method == "GET" {
			switch r.URL.Path {
			case PathLogin:
				login(c, s, w, r)
			case PathLogout:
				logout(s, w, r)
			case PathCallback:
				handleOAuth2Callback(c, s, w, r)
			default:
				next(w, r)
			}
		} else {
			next(w, r)
		}

	}
}

func GetToken(r *http.Request) Tokens {
	s := sessions.GetSession(r)
	t := unmarshallToken(s)

	//not doing this doesn't pass through the
	//nil return, causing a test to fail - not sure why??
	if t == nil {
		return nil
	} else {
		return t
	}
}

func SetToken(r *http.Request, t interface{}) {
	s := sessions.GetSession(r)
	val, _ := json.Marshal(t)
	s.Set(keyToken, val)
	//Check immediately to see if the token is expired
	tk := unmarshallToken(s)
	if tk != nil {
		// check if the access token is expired
		if !tk.Valid() && tk.Refresh() == "" {
			s.Delete(keyToken)
			tk = nil
		}
	}
}

// Handler that redirects user to the login page
// if user is not logged in.
func LoginRequired() negroni.HandlerFunc {
	return func(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
		token := GetToken(r)
		if token == nil || !token.Valid() {
			// Set token to null to avoid redirection loop
			SetToken(r, nil)
			next := url.QueryEscape(r.URL.RequestURI())
			http.Redirect(rw, r, PathLogin+"?"+keyNextPage+"="+next, http.StatusFound)
		} else {
			next(rw, r)
		}
	}
}

func newState() string {
	var p [16]byte
	_, err := rand.Read(p[:])
	if err != nil {
		panic(err)
	}
	return hex.EncodeToString(p[:])
}

func login(config *oauth2.Config, s sessions.Session, w http.ResponseWriter, r *http.Request) {
	next := extractPath(r.URL.Query().Get(keyNextPage))

	if s.Get(keyToken) == nil {
		// User is not logged in.
		if next == "" {
			next = "/"
		}

		state := newState()
		// store the next url and state token in the session
		s.Set(keyState, state)
		s.Set(keyNextPage, next)
		http.Redirect(w, r, config.AuthCodeURL(state, oauth2.AccessTypeOffline), http.StatusFound)
		return
	}
	// No need to login, redirect to the next page.
	http.Redirect(w, r, next, http.StatusFound)
}

func logout(s sessions.Session, w http.ResponseWriter, r *http.Request) {
	next := extractPath(r.URL.Query().Get(keyNextPage))
	s.Delete(keyToken)
	http.Redirect(w, r, next, http.StatusFound)
}

func handleOAuth2Callback(config *oauth2.Config, s sessions.Session, w http.ResponseWriter, r *http.Request) {
	providedState := extractPath(r.URL.Query().Get("state"))
	fmt.Printf("Got state from request %s\n", providedState)

	//verify that the provided state is the state we generated
	//if it is not, then redirect to the error page
	originalState := s.Get(keyState)
	fmt.Printf("Got state from session %s\n", originalState)
	if providedState != originalState {
		http.Redirect(w, r, PathError, http.StatusFound)
		return
	}

	next := s.Get(keyNextPage).(string)
	fmt.Printf("Got a next page from the session: %s\n", next)
	code := r.URL.Query().Get("code")
	t, err := config.Exchange(oauth2.NoContext, code)
	if err != nil {
		// Pass the error message, or allow dev to provide its own
		// error handler.
		http.Redirect(w, r, PathError, http.StatusFound)
		return
	}
	// Store the credentials in the session.
	val, _ := json.Marshal(t)
	s.Set(keyToken, val)
	http.Redirect(w, r, next, http.StatusFound)
}

func unmarshallToken(s sessions.Session) *token {

	if s.Get(keyToken) == nil {
		return nil
	}

	data := s.Get(keyToken).([]byte)
	var tk oauth2.Token
	json.Unmarshal(data, &tk)
	return &token{tk}

}

func extractPath(next string) string {
	n, err := url.Parse(next)
	if err != nil {
		return "/"
	}
	return n.Path
}
<commit_msg>Removal of debug statements<commit_after>// Copyright 2014 GoIncremental Limited. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package oauth2 contains Negroni middleware to provide
// user login via an OAuth 2.0 backend.

package oauth2

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"time"

	"github.com/codegangsta/negroni"
	sessions "github.com/goincremental/negroni-sessions"
	"golang.org/x/oauth2"
)

const (
	codeRedirect = 302
	keyToken     = "oauth2_token"
	keyNextPage  = "next"
	keyState     = "state"
)

var (
	// PathLogin sets the path to handle OAuth 2.0 logins.
	PathLogin = "/login"
	// PathLogout sets to handle OAuth 2.0 logouts.
	PathLogout = "/logout"
	// PathCallback sets the path to handle callback from OAuth 2.0 backend
	// to exchange credentials.
	PathCallback = "/oauth2callback"
	// PathError sets the path to handle error cases.
	PathError = "/oauth2error"
)

type Config oauth2.Config

// Tokens Represents a container that contains
// user's OAuth 2.0 access and refresh tokens.
type Tokens interface {
	Access() string
	Refresh() string
	Valid() bool
	ExpiryTime() time.Time
	ExtraData(string) interface{}
	Get() Token
}

type Token oauth2.Token

type token struct {
	oauth2.Token
}

func (t *token) ExtraData(key string) interface{} {
	return t.Extra(key)
}

// Returns the access token.
func (t *token) Access() string {
	return t.AccessToken
}

// Returns the refresh token.
func (t *token) Refresh() string {
	return t.RefreshToken
}

// Returns whether the access token is
// expired or not.
func (t *token) Valid() bool {
	if t == nil {
		return true
	}
	return t.Token.Valid()
}

// Returns the expiry time of the user's
// access token.
func (t *token) ExpiryTime() time.Time {
	return t.Expiry
}

// String returns the string representation of the token.
func (t *token) String() string {
	return fmt.Sprintf("tokens: %v", t)
}

// Returns oauth2.Token.
func (t *token) Get() Token {
	return (Token)(t.Token)
}

// Returns a new Google OAuth 2.0 backend endpoint.
func Google(config *Config) negroni.Handler {
	authUrl := "https://accounts.google.com/o/oauth2/auth"
	tokenUrl := "https://accounts.google.com/o/oauth2/token"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

// Returns a new Github OAuth 2.0 backend endpoint.
func Github(config *Config) negroni.Handler {
	authUrl := "https://github.com/login/oauth/authorize"
	tokenUrl := "https://github.com/login/oauth/access_token"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

func Facebook(config *Config) negroni.Handler {
	authUrl := "https://www.facebook.com/dialog/oauth"
	tokenUrl := "https://graph.facebook.com/oauth/access_token"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

func LinkedIn(config *Config) negroni.Handler {
	authUrl := "https://www.linkedin.com/uas/oauth2/authorization"
	tokenUrl := "https://www.linkedin.com/uas/oauth2/accessToken"
	return NewOAuth2Provider(config, authUrl, tokenUrl)
}

// Returns a generic OAuth 2.0 backend endpoint.
func NewOAuth2Provider(config *Config, authUrl, tokenUrl string) negroni.HandlerFunc {
	c := &oauth2.Config{
		ClientID:     config.ClientID,
		ClientSecret: config.ClientSecret,
		Scopes:       config.Scopes,
		RedirectURL:  config.RedirectURL,
		Endpoint: oauth2.Endpoint{
			AuthURL:  authUrl,
			TokenURL: tokenUrl,
		},
	}

	return func(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
		s := sessions.GetSession(r)

		if r.Method == "GET" {
			switch r.URL.Path {
			case PathLogin:
				login(c, s, w, r)
			case PathLogout:
				logout(s, w, r)
			case PathCallback:
				handleOAuth2Callback(c, s, w, r)
			default:
				next(w, r)
			}
		} else {
			next(w, r)
		}

	}
}

func GetToken(r *http.Request) Tokens {
	s := sessions.GetSession(r)
	t := unmarshallToken(s)

	//not doing this doesn't pass through the
	//nil return, causing a test to fail - not sure why??
	if t == nil {
		return nil
	} else {
		return t
	}
}

func SetToken(r *http.Request, t interface{}) {
	s := sessions.GetSession(r)
	val, _ := json.Marshal(t)
	s.Set(keyToken, val)
	//Check immediately to see if the token is expired
	tk := unmarshallToken(s)
	if tk != nil {
		// check if the access token is expired
		if !tk.Valid() && tk.Refresh() == "" {
			s.Delete(keyToken)
			tk = nil
		}
	}
}

// Handler that redirects user to the login page
// if user is not logged in.
func LoginRequired() negroni.HandlerFunc {
	return func(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
		token := GetToken(r)
		if token == nil || !token.Valid() {
			// Set token to null to avoid redirection loop
			SetToken(r, nil)
			next := url.QueryEscape(r.URL.RequestURI())
			http.Redirect(rw, r, PathLogin+"?"+keyNextPage+"="+next, http.StatusFound)
		} else {
			next(rw, r)
		}
	}
}

func newState() string {
	var p [16]byte
	_, err := rand.Read(p[:])
	if err != nil {
		panic(err)
	}
	return hex.EncodeToString(p[:])
}

func login(config *oauth2.Config, s sessions.Session, w http.ResponseWriter, r *http.Request) {
	next := extractPath(r.URL.Query().Get(keyNextPage))

	if s.Get(keyToken) == nil {
		// User is not logged in.
		if next == "" {
			next = "/"
		}

		state := newState()
		// store the next url and state token in the session
		s.Set(keyState, state)
		s.Set(keyNextPage, next)
		http.Redirect(w, r, config.AuthCodeURL(state, oauth2.AccessTypeOffline), http.StatusFound)
		return
	}
	// No need to login, redirect to the next page.
	http.Redirect(w, r, next, http.StatusFound)
}

func logout(s sessions.Session, w http.ResponseWriter, r *http.Request) {
	next := extractPath(r.URL.Query().Get(keyNextPage))
	s.Delete(keyToken)
	http.Redirect(w, r, next, http.StatusFound)
}

func handleOAuth2Callback(config *oauth2.Config, s sessions.Session, w http.ResponseWriter, r *http.Request) {
	providedState := extractPath(r.URL.Query().Get("state"))

	//verify that the provided state is the state we generated
	//if it is not, then redirect to the error page
	originalState := s.Get(keyState)
	if providedState != originalState {
		http.Redirect(w, r, PathError, http.StatusFound)
		return
	}

	next := s.Get(keyNextPage).(string)
	code := r.URL.Query().Get("code")
	t, err := config.Exchange(oauth2.NoContext, code)
	if err != nil {
		// Pass the error message, or allow dev to provide its own
		// error handler.
		http.Redirect(w, r, PathError, http.StatusFound)
		return
	}
	// Store the credentials in the session.
	val, _ := json.Marshal(t)
	s.Set(keyToken, val)
	http.Redirect(w, r, next, http.StatusFound)
}

func unmarshallToken(s sessions.Session) *token {

	if s.Get(keyToken) == nil {
		return nil
	}

	data := s.Get(keyToken).([]byte)
	var tk oauth2.Token
	json.Unmarshal(data, &tk)
	return &token{tk}

}

func extractPath(next string) string {
	n, err := url.Parse(next)
	if err != nil {
		return "/"
	}
	return n.Path
}
<|endoftext|>


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

// Package oauth2 contains Martini handlers to provide
// user login via an OAuth 2.0 backend.
package oauth2

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"

	"code.google.com/p/goauth2/oauth"
	"github.com/codegangsta/martini"
	"github.com/martini-contrib/sessions"
)

const (
	codeRedirect = 302
	keyToken     = "oauth2_token"
	keyNextPage  = "next"
)

var (
	// Path to handle OAuth 2.0 logins.
	PathLogin = "/login"
	// Path to handle OAuth 2.0 logouts.
	PathLogout = "/logout"
	// Path to handle callback from OAuth 2.0 backend
	// to exchange credentials.
	PathCallback = "/oauth2callback"
	// Path to handle error cases.
	PathError = "/oauth2error"
)

// Represents OAuth2 backend options.
type Options struct {
	ClientId     string
	ClientSecret string
	RedirectURL  string
	Scopes       []string

	AuthUrl  string
	TokenUrl string
}

// Represents a container that contains
// user's OAuth 2.0 access and refresh tokens.
type Tokens interface {
	Access() string
	Refresh() string
	IsExpired() bool
	ExpiryTime() time.Time
}

type token struct {
	oauth.Token
}

// Returns the access token.
func (t *token) Access() string {
	return t.AccessToken
}

// Returns the refresh token.
func (t *token) Refresh() string {
	return t.RefreshToken
}

// Returns whether the access token is
// expired or not.
func (t *token) IsExpired() bool {
	if t == nil {
		return true
	}
	return t.Expired()
}

// Returns the expiry time of the user's
// access token.
func (t *token) ExpiryTime() time.Time {
	return t.Expiry
}

// Formats tokens into string.
func (t *token) String() string {
	return fmt.Sprintf("tokens: %v", t)
}

// Returns a new Google OAuth 2.0 backend endpoint.
func Google(opts *Options) martini.Handler {
	opts.AuthUrl = "https://accounts.google.com/o/oauth2/auth"
	opts.TokenUrl = "https://accounts.google.com/o/oauth2/token"
	return NewOAuth2Provider(opts)
}

// Returns a new Github OAuth 2.0 backend endpoint.
func Github(opts *Options) martini.Handler {
	opts.AuthUrl = "https://github.com/login/oauth/authorize"
	opts.TokenUrl = "https://github.com/login/oauth/access_token"
	return NewOAuth2Provider(opts)
}

func Facebook(opts *Options) martini.Handler {
	opts.AuthUrl = "https://www.facebook.com/dialog/oauth"
	opts.TokenUrl = "https://graph.facebook.com/oauth/access_token"
	return NewOAuth2Provider(opts)
}

// Returns a generic OAuth 2.0 backend endpoint.
func NewOAuth2Provider(opts *Options) martini.Handler {
	config := &oauth.Config{
		ClientId:     opts.ClientId,
		ClientSecret: opts.ClientSecret,
		RedirectURL:  opts.RedirectURL,
		Scope:        strings.Join(opts.Scopes, " "),
		AuthURL:      opts.AuthUrl,
		TokenURL:     opts.TokenUrl,
	}

	transport := &oauth.Transport{
		Config:    config,
		Transport: http.DefaultTransport,
	}

	return func(s sessions.Session, c martini.Context, w http.ResponseWriter, r *http.Request) {
		if r.Method == "GET" {
			switch r.URL.Path {
			case PathLogin:
				login(transport, s, w, r)
			case PathLogout:
				logout(transport, s, w, r)
			case PathCallback:
				handleOAuth2Callback(transport, s, w, r)
			}
		}

		tk := unmarshallToken(s)
		if tk != nil {
			// check if the access token is expired
			if tk.IsExpired() && tk.Refresh() == "" {
				s.Delete(keyToken)
				tk = nil
			}
		}
		// Inject tokens.
		c.MapTo(tk, (*Tokens)(nil))
	}
}

// Handler that redirects user to the login page
// if user is not logged in.
// Sample usage:
// m.Get("/login-required", oauth2.LoginRequired, func() ... {})
var LoginRequired martini.Handler = func() martini.Handler {
	return func(s sessions.Session, c martini.Context, w http.ResponseWriter, r *http.Request) {
		token := unmarshallToken(s)
		if token == nil || token.IsExpired() {
			next := url.QueryEscape(r.URL.RequestURI())
			http.Redirect(w, r, PathLogin+"?next="+next, codeRedirect)
		}
	}
}()

func login(t *oauth.Transport, s sessions.Session, w http.ResponseWriter, r *http.Request) {
	next := r.URL.Query().Get(keyNextPage)
	if s.Get(keyToken) == nil {
		// User is not logged in.
		http.Redirect(w, r, t.Config.AuthCodeURL(next), codeRedirect)
		return
	}
	// No need to login, redirect to the next page.
	http.Redirect(w, r, next, codeRedirect)
}

func logout(t *oauth.Transport, s sessions.Session, w http.ResponseWriter, r *http.Request) {
	next := r.URL.Query().Get(keyNextPage)
	s.Delete(keyToken)
	http.Redirect(w, r, next, codeRedirect)
}

func handleOAuth2Callback(t *oauth.Transport, s sessions.Session, w http.ResponseWriter, r *http.Request) {
	next := r.URL.Query().Get("state")
	code := r.URL.Query().Get("code")
	tk, err := t.Exchange(code)
	if err != nil {
		// Pass the error message, or allow dev to provide its own
		// error handler.
		http.Redirect(w, r, PathError, codeRedirect)
		return
	}
	// Store the credentials in the session.
	val, _ := json.Marshal(tk)
	s.Set(keyToken, val)
	http.Redirect(w, r, next, codeRedirect)
}

func unmarshallToken(s sessions.Session) (t *token) {
	if s.Get(keyToken) == nil {
		return
	}
	data := s.Get(keyToken).([]byte)
	var tk oauth.Token
	json.Unmarshal(data, &tk)
	return &token{tk}
}
<commit_msg>Fixes issue with unvalidated redirects in next parameter<commit_after>// Copyright 2014 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package oauth2 contains Martini handlers to provide
// user login via an OAuth 2.0 backend.
package oauth2

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"

	"code.google.com/p/goauth2/oauth"
	"github.com/codegangsta/martini"
	"github.com/martini-contrib/sessions"
)

const (
	codeRedirect = 302
	keyToken     = "oauth2_token"
	keyNextPage  = "next"
)

var (
	// Path to handle OAuth 2.0 logins.
	PathLogin = "/login"
	// Path to handle OAuth 2.0 logouts.
	PathLogout = "/logout"
	// Path to handle callback from OAuth 2.0 backend
	// to exchange credentials.
	PathCallback = "/oauth2callback"
	// Path to handle error cases.
	PathError = "/oauth2error"
)

// Represents OAuth2 backend options.
type Options struct {
	ClientId     string
	ClientSecret string
	RedirectURL  string
	Scopes       []string

	AuthUrl  string
	TokenUrl string
}

// Represents a container that contains
// user's OAuth 2.0 access and refresh tokens.
type Tokens interface {
	Access() string
	Refresh() string
	IsExpired() bool
	ExpiryTime() time.Time
}

type token struct {
	oauth.Token
}

// Returns the access token.
func (t *token) Access() string {
	return t.AccessToken
}

// Returns the refresh token.
func (t *token) Refresh() string {
	return t.RefreshToken
}

// Returns whether the access token is
// expired or not.
func (t *token) IsExpired() bool {
	if t == nil {
		return true
	}
	return t.Expired()
}

// Returns the expiry time of the user's
// access token.
func (t *token) ExpiryTime() time.Time {
	return t.Expiry
}

// Formats tokens into string.
func (t *token) String() string {
	return fmt.Sprintf("tokens: %v", t)
}

// Returns a new Google OAuth 2.0 backend endpoint.
func Google(opts *Options) martini.Handler {
	opts.AuthUrl = "https://accounts.google.com/o/oauth2/auth"
	opts.TokenUrl = "https://accounts.google.com/o/oauth2/token"
	return NewOAuth2Provider(opts)
}

// Returns a new Github OAuth 2.0 backend endpoint.
func Github(opts *Options) martini.Handler {
	opts.AuthUrl = "https://github.com/login/oauth/authorize"
	opts.TokenUrl = "https://github.com/login/oauth/access_token"
	return NewOAuth2Provider(opts)
}

func Facebook(opts *Options) martini.Handler {
	opts.AuthUrl = "https://www.facebook.com/dialog/oauth"
	opts.TokenUrl = "https://graph.facebook.com/oauth/access_token"
	return NewOAuth2Provider(opts)
}

// Returns a generic OAuth 2.0 backend endpoint.
func NewOAuth2Provider(opts *Options) martini.Handler {
	config := &oauth.Config{
		ClientId:     opts.ClientId,
		ClientSecret: opts.ClientSecret,
		RedirectURL:  opts.RedirectURL,
		Scope:        strings.Join(opts.Scopes, " "),
		AuthURL:      opts.AuthUrl,
		TokenURL:     opts.TokenUrl,
	}

	transport := &oauth.Transport{
		Config:    config,
		Transport: http.DefaultTransport,
	}

	return func(s sessions.Session, c martini.Context, w http.ResponseWriter, r *http.Request) {
		if r.Method == "GET" {
			switch r.URL.Path {
			case PathLogin:
				login(transport, s, w, r)
			case PathLogout:
				logout(transport, s, w, r)
			case PathCallback:
				handleOAuth2Callback(transport, s, w, r)
			}
		}

		tk := unmarshallToken(s)
		if tk != nil {
			// check if the access token is expired
			if tk.IsExpired() && tk.Refresh() == "" {
				s.Delete(keyToken)
				tk = nil
			}
		}
		// Inject tokens.
		c.MapTo(tk, (*Tokens)(nil))
	}
}

// Handler that redirects user to the login page
// if user is not logged in.
// Sample usage:
// m.Get("/login-required", oauth2.LoginRequired, func() ... {})
var LoginRequired martini.Handler = func() martini.Handler {
	return func(s sessions.Session, c martini.Context, w http.ResponseWriter, r *http.Request) {
		token := unmarshallToken(s)
		if token == nil || token.IsExpired() {
			next := url.QueryEscape(r.URL.RequestURI())
			http.Redirect(w, r, PathLogin+"?next="+next, codeRedirect)
		}
	}
}()

func login(t *oauth.Transport, s sessions.Session, w http.ResponseWriter, r *http.Request) {
	next := extractPath(r.URL.Query().Get(keyNextPage))
	if s.Get(keyToken) == nil {
		// User is not logged in.
		http.Redirect(w, r, t.Config.AuthCodeURL(next), codeRedirect)
		return
	}
	// No need to login, redirect to the next page.
	http.Redirect(w, r, next, codeRedirect)
}

func logout(t *oauth.Transport, s sessions.Session, w http.ResponseWriter, r *http.Request) {
	next := extractPath(r.URL.Query().Get(keyNextPage))
	s.Delete(keyToken)
	http.Redirect(w, r, next, codeRedirect)
}

func handleOAuth2Callback(t *oauth.Transport, s sessions.Session, w http.ResponseWriter, r *http.Request) {
	next := extractPath(r.URL.Query().Get("state"))
	code := r.URL.Query().Get("code")
	tk, err := t.Exchange(code)
	if err != nil {
		// Pass the error message, or allow dev to provide its own
		// error handler.
		http.Redirect(w, r, PathError, codeRedirect)
		return
	}
	// Store the credentials in the session.
	val, _ := json.Marshal(tk)
	s.Set(keyToken, val)
	http.Redirect(w, r, next, codeRedirect)
}

func unmarshallToken(s sessions.Session) (t *token) {
	if s.Get(keyToken) == nil {
		return
	}
	data := s.Get(keyToken).([]byte)
	var tk oauth.Token
	json.Unmarshal(data, &tk)
	return &token{tk}
}

func extractPath(next string) string {
	n, err := url.Parse(next)
	if err != nil {
		return "/"
	}
	return n.Path
}
<|endoftext|>