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

// The version package implements version parsing.
// It also acts as guardian of the current client Juju version number.
package version

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"

	"labix.org/v2/mgo/bson"
)

// The presence and format of this constant is very important.
// The debian/rules build recipe uses this value for the version
// number of the release package.
const version = "1.15.0"

// Current gives the current version of the system.  If the file
// "FORCE-VERSION" is present in the same directory as the running
// binary, it will override this.
var Current = Binary{
	Number: MustParse(version),
	Series: readSeries("/etc/lsb-release"),
	Arch:   ubuntuArch(runtime.GOARCH),
}

func init() {
	toolsDir := filepath.Dir(os.Args[0])
	v, err := ioutil.ReadFile(filepath.Join(toolsDir, "FORCE-VERSION"))
	if err != nil {
		if !os.IsNotExist(err) {
			fmt.Fprintf(os.Stderr, "WARNING: cannot read forced version: %v\n", err)
		}
		return
	}
	Current.Number = MustParse(strings.TrimSpace(string(v)))
}

// Number represents a juju version.  When bugs are fixed the patch
// number is incremented; when new features are added the minor number
// is incremented and patch is reset; and when compatibility is broken
// the major version is incremented and minor and patch are reset.  The
// build number is automatically assigned and has no well defined
// sequence.  If the build number is greater than zero or the minor
// version is odd, it indicates that the release is still in
// development.
type Number struct {
	Major int
	Minor int
	Patch int
	Build int
}

// Zero is occasionally convenient and readable.
// Please don't change its value.
var Zero = Number{}

// Binary specifies a binary version of juju.
type Binary struct {
	Number
	Series string
	Arch   string
}

func (v Binary) String() string {
	return fmt.Sprintf("%v-%s-%s", v.Number, v.Series, v.Arch)
}

// GetBSON turns v into a bson.Getter so it can be saved directly
// on a MongoDB database with mgo.
func (v Binary) GetBSON() (interface{}, error) {
	return v.String(), nil
}

// SetBSON turns v into a bson.Setter so it can be loaded directly
// from a MongoDB database with mgo.
func (vp *Binary) SetBSON(raw bson.Raw) error {
	var s string
	err := raw.Unmarshal(&s)
	if err != nil {
		return err
	}
	v, err := ParseBinary(s)
	if err != nil {
		return err
	}
	*vp = v
	return nil
}

func (v Binary) MarshalJSON() ([]byte, error) {
	return json.Marshal(v.String())
}

func (vp *Binary) UnmarshalJSON(data []byte) error {
	var s string
	if err := json.Unmarshal(data, &s); err != nil {
		return err
	}
	v, err := ParseBinary(s)
	if err != nil {
		return err
	}
	*vp = v
	return nil
}

var (
	binaryPat = regexp.MustCompile(`^(\d{1,9})\.(\d{1,9})\.(\d{1,9})(\.\d{1,9})?-([^-]+)-([^-]+)$`)
	numberPat = regexp.MustCompile(`^(\d{1,9})\.(\d{1,9})\.(\d{1,9})(\.\d{1,9})?$`)
)

// MustParse parses a version and panics if it does
// not parse correctly.
func MustParse(s string) Number {
	v, err := Parse(s)
	if err != nil {
		panic(err)
	}
	return v
}

// MustParseBinary parses a binary version and panics if it does
// not parse correctly.
func MustParseBinary(s string) Binary {
	v, err := ParseBinary(s)
	if err != nil {
		panic(err)
	}
	return v
}

// ParseBinary parses a binary version of the form "1.2.3-series-arch".
func ParseBinary(s string) (Binary, error) {
	m := binaryPat.FindStringSubmatch(s)
	if m == nil {
		return Binary{}, fmt.Errorf("invalid binary version %q", s)
	}
	var v Binary
	v.Major = atoi(m[1])
	v.Minor = atoi(m[2])
	v.Patch = atoi(m[3])
	if m[4] != "" {
		v.Build = atoi(m[4][1:])
	}
	v.Series = m[5]
	v.Arch = m[6]
	return v, nil
}

// Parse parses the version, which is of the form 1.2.3
// giving the major, minor and release versions
// respectively.
func Parse(s string) (Number, error) {
	m := numberPat.FindStringSubmatch(s)
	if m == nil {
		return Number{}, fmt.Errorf("invalid version %q", s)
	}
	var v Number
	v.Major = atoi(m[1])
	v.Minor = atoi(m[2])
	v.Patch = atoi(m[3])
	if m[4] != "" {
		v.Build = atoi(m[4][1:])
	}
	return v, nil
}

// atoi is the same as strconv.Atoi but assumes that
// the string has been verified to be a valid integer.
func atoi(s string) int {
	n, err := strconv.Atoi(s)
	if err != nil {
		panic(err)
	}
	return n
}

func (v Number) String() string {
	s := fmt.Sprintf("%d.%d.%d", v.Major, v.Minor, v.Patch)
	if v.Build > 0 {
		s += fmt.Sprintf(".%d", v.Build)
	}
	return s
}

// Less returns whether v is semantically earlier in the
// version sequence than w.
func (v Number) Less(w Number) bool {
	switch {
	case v.Major != w.Major:
		return v.Major < w.Major
	case v.Minor != w.Minor:
		return v.Minor < w.Minor
	case v.Patch != w.Patch:
		return v.Patch < w.Patch
	case v.Build != w.Build:
		return v.Build < w.Build
	}
	return false
}

// GetBSON turns v into a bson.Getter so it can be saved directly
// on a MongoDB database with mgo.
func (v Number) GetBSON() (interface{}, error) {
	return v.String(), nil
}

// SetBSON turns v into a bson.Setter so it can be loaded directly
// from a MongoDB database with mgo.
func (vp *Number) SetBSON(raw bson.Raw) error {
	var s string
	err := raw.Unmarshal(&s)
	if err != nil {
		return err
	}
	v, err := Parse(s)
	if err != nil {
		return err
	}
	*vp = v
	return nil
}

func (v Number) MarshalJSON() ([]byte, error) {
	return json.Marshal(v.String())
}

func (vp *Number) UnmarshalJSON(data []byte) error {
	var s string
	if err := json.Unmarshal(data, &s); err != nil {
		return err
	}
	v, err := Parse(s)
	if err != nil {
		return err
	}
	*vp = v
	return nil
}

func isOdd(x int) bool {
	return x%2 != 0
}

// IsDev returns whether the version represents a development
// version. A version with an odd-numbered minor component or
// a nonzero build component is considered to be a development
// version.
func (v Number) IsDev() bool {
	return isOdd(v.Minor) || v.Build > 0
}

func readSeries(releaseFile string) string {
	data, err := ioutil.ReadFile(releaseFile)
	if err != nil {
		return "unknown"
	}
	for _, line := range strings.Split(string(data), "\n") {
		const p = "DISTRIB_CODENAME="
		if strings.HasPrefix(line, p) {
			return strings.Trim(line[len(p):], "\t '\"")
		}
	}
	return "unknown"
}

func ubuntuArch(arch string) string {
	if arch == "386" {
		arch = "i386"
	}
	return arch
}

// ParseMajorMinor takes an argument of the form "major.minor" and returns ints major and minor.
func ParseMajorMinor(vers string) (int, int, error) {
	parts := strings.Split(vers, ".")
	major, err := strconv.Atoi(parts[0])
	minor := -1
	if err != nil {
		return -1, -1, fmt.Errorf("invalid major version number %s: %v", parts[0], err)
	}
	if len(parts) == 2 {
		minor, err = strconv.Atoi(parts[1])
		if err != nil {
			return -1, -1, fmt.Errorf("invalid minor version number %s: %v", parts[1], err)
		}
	} else if len(parts) > 2 {
		return -1, -1, fmt.Errorf("invalid major.minor version number %s", vers)
	}
	return major, minor, nil
}
<commit_msg>[r=jameinel] Increment Juju to 1.15.1<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// The version package implements version parsing.
// It also acts as guardian of the current client Juju version number.
package version

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"

	"labix.org/v2/mgo/bson"
)

// The presence and format of this constant is very important.
// The debian/rules build recipe uses this value for the version
// number of the release package.
const version = "1.15.1"

// Current gives the current version of the system.  If the file
// "FORCE-VERSION" is present in the same directory as the running
// binary, it will override this.
var Current = Binary{
	Number: MustParse(version),
	Series: readSeries("/etc/lsb-release"),
	Arch:   ubuntuArch(runtime.GOARCH),
}

func init() {
	toolsDir := filepath.Dir(os.Args[0])
	v, err := ioutil.ReadFile(filepath.Join(toolsDir, "FORCE-VERSION"))
	if err != nil {
		if !os.IsNotExist(err) {
			fmt.Fprintf(os.Stderr, "WARNING: cannot read forced version: %v\n", err)
		}
		return
	}
	Current.Number = MustParse(strings.TrimSpace(string(v)))
}

// Number represents a juju version.  When bugs are fixed the patch
// number is incremented; when new features are added the minor number
// is incremented and patch is reset; and when compatibility is broken
// the major version is incremented and minor and patch are reset.  The
// build number is automatically assigned and has no well defined
// sequence.  If the build number is greater than zero or the minor
// version is odd, it indicates that the release is still in
// development.
type Number struct {
	Major int
	Minor int
	Patch int
	Build int
}

// Zero is occasionally convenient and readable.
// Please don't change its value.
var Zero = Number{}

// Binary specifies a binary version of juju.
type Binary struct {
	Number
	Series string
	Arch   string
}

func (v Binary) String() string {
	return fmt.Sprintf("%v-%s-%s", v.Number, v.Series, v.Arch)
}

// GetBSON turns v into a bson.Getter so it can be saved directly
// on a MongoDB database with mgo.
func (v Binary) GetBSON() (interface{}, error) {
	return v.String(), nil
}

// SetBSON turns v into a bson.Setter so it can be loaded directly
// from a MongoDB database with mgo.
func (vp *Binary) SetBSON(raw bson.Raw) error {
	var s string
	err := raw.Unmarshal(&s)
	if err != nil {
		return err
	}
	v, err := ParseBinary(s)
	if err != nil {
		return err
	}
	*vp = v
	return nil
}

func (v Binary) MarshalJSON() ([]byte, error) {
	return json.Marshal(v.String())
}

func (vp *Binary) UnmarshalJSON(data []byte) error {
	var s string
	if err := json.Unmarshal(data, &s); err != nil {
		return err
	}
	v, err := ParseBinary(s)
	if err != nil {
		return err
	}
	*vp = v
	return nil
}

var (
	binaryPat = regexp.MustCompile(`^(\d{1,9})\.(\d{1,9})\.(\d{1,9})(\.\d{1,9})?-([^-]+)-([^-]+)$`)
	numberPat = regexp.MustCompile(`^(\d{1,9})\.(\d{1,9})\.(\d{1,9})(\.\d{1,9})?$`)
)

// MustParse parses a version and panics if it does
// not parse correctly.
func MustParse(s string) Number {
	v, err := Parse(s)
	if err != nil {
		panic(err)
	}
	return v
}

// MustParseBinary parses a binary version and panics if it does
// not parse correctly.
func MustParseBinary(s string) Binary {
	v, err := ParseBinary(s)
	if err != nil {
		panic(err)
	}
	return v
}

// ParseBinary parses a binary version of the form "1.2.3-series-arch".
func ParseBinary(s string) (Binary, error) {
	m := binaryPat.FindStringSubmatch(s)
	if m == nil {
		return Binary{}, fmt.Errorf("invalid binary version %q", s)
	}
	var v Binary
	v.Major = atoi(m[1])
	v.Minor = atoi(m[2])
	v.Patch = atoi(m[3])
	if m[4] != "" {
		v.Build = atoi(m[4][1:])
	}
	v.Series = m[5]
	v.Arch = m[6]
	return v, nil
}

// Parse parses the version, which is of the form 1.2.3
// giving the major, minor and release versions
// respectively.
func Parse(s string) (Number, error) {
	m := numberPat.FindStringSubmatch(s)
	if m == nil {
		return Number{}, fmt.Errorf("invalid version %q", s)
	}
	var v Number
	v.Major = atoi(m[1])
	v.Minor = atoi(m[2])
	v.Patch = atoi(m[3])
	if m[4] != "" {
		v.Build = atoi(m[4][1:])
	}
	return v, nil
}

// atoi is the same as strconv.Atoi but assumes that
// the string has been verified to be a valid integer.
func atoi(s string) int {
	n, err := strconv.Atoi(s)
	if err != nil {
		panic(err)
	}
	return n
}

func (v Number) String() string {
	s := fmt.Sprintf("%d.%d.%d", v.Major, v.Minor, v.Patch)
	if v.Build > 0 {
		s += fmt.Sprintf(".%d", v.Build)
	}
	return s
}

// Less returns whether v is semantically earlier in the
// version sequence than w.
func (v Number) Less(w Number) bool {
	switch {
	case v.Major != w.Major:
		return v.Major < w.Major
	case v.Minor != w.Minor:
		return v.Minor < w.Minor
	case v.Patch != w.Patch:
		return v.Patch < w.Patch
	case v.Build != w.Build:
		return v.Build < w.Build
	}
	return false
}

// GetBSON turns v into a bson.Getter so it can be saved directly
// on a MongoDB database with mgo.
func (v Number) GetBSON() (interface{}, error) {
	return v.String(), nil
}

// SetBSON turns v into a bson.Setter so it can be loaded directly
// from a MongoDB database with mgo.
func (vp *Number) SetBSON(raw bson.Raw) error {
	var s string
	err := raw.Unmarshal(&s)
	if err != nil {
		return err
	}
	v, err := Parse(s)
	if err != nil {
		return err
	}
	*vp = v
	return nil
}

func (v Number) MarshalJSON() ([]byte, error) {
	return json.Marshal(v.String())
}

func (vp *Number) UnmarshalJSON(data []byte) error {
	var s string
	if err := json.Unmarshal(data, &s); err != nil {
		return err
	}
	v, err := Parse(s)
	if err != nil {
		return err
	}
	*vp = v
	return nil
}

func isOdd(x int) bool {
	return x%2 != 0
}

// IsDev returns whether the version represents a development
// version. A version with an odd-numbered minor component or
// a nonzero build component is considered to be a development
// version.
func (v Number) IsDev() bool {
	return isOdd(v.Minor) || v.Build > 0
}

func readSeries(releaseFile string) string {
	data, err := ioutil.ReadFile(releaseFile)
	if err != nil {
		return "unknown"
	}
	for _, line := range strings.Split(string(data), "\n") {
		const p = "DISTRIB_CODENAME="
		if strings.HasPrefix(line, p) {
			return strings.Trim(line[len(p):], "\t '\"")
		}
	}
	return "unknown"
}

func ubuntuArch(arch string) string {
	if arch == "386" {
		arch = "i386"
	}
	return arch
}

// ParseMajorMinor takes an argument of the form "major.minor" and returns ints major and minor.
func ParseMajorMinor(vers string) (int, int, error) {
	parts := strings.Split(vers, ".")
	major, err := strconv.Atoi(parts[0])
	minor := -1
	if err != nil {
		return -1, -1, fmt.Errorf("invalid major version number %s: %v", parts[0], err)
	}
	if len(parts) == 2 {
		minor, err = strconv.Atoi(parts[1])
		if err != nil {
			return -1, -1, fmt.Errorf("invalid minor version number %s: %v", parts[1], err)
		}
	} else if len(parts) > 2 {
		return -1, -1, fmt.Errorf("invalid major.minor version number %s", vers)
	}
	return major, minor, nil
}
<|endoftext|>


<commit_before>package restwebsocket

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// ConnectionType is the socket connection type, it can either be a serverside connection or a clientside connection
type ConnectionType int

const (
	// ServerSide means this connection is owned by a websocket-server
	ServerSide ConnectionType = iota
	// ClientSide means this connection is owned by a websocket-client
	ClientSide
)

const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// Time allowed to read the next pong message from the peer.
	pongWait = 20 * time.Second

	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10

	// Maximum message size allowed from peer.
	maxMessageSize = 512
)

// SocketConnection is a wrapper around the websocket connection to handle http
type SocketConnection struct {
	socketserver        *WebsocketServer
	socketclient        *WebsocketClient
	connType            ConnectionType
	conn                *websocket.Conn
	id                  string
	heartBeat           *heartbeat
	closeNotificationCh chan bool
	once                sync.Once
}

// NewSocketConnection creates a new socket connection
func NewSocketConnection(c *websocket.Conn, id string, keepAlive bool, pingHdlr, pongHdlr func(string) error, appData []byte) *SocketConnection {
	// Default ping handler is to send back a pong control message with the same application data
	// Default pong handler is to do nothing
	// websocket protocol mentions that the pong message should reply back with the exact appData recieved from the ping message
	if pingHdlr != nil {
		c.SetPingHandler(pingHdlr)
	}
	if pongHdlr != nil {
		c.SetPongHandler(pingHdlr)
	}

	sockconn := &SocketConnection{
		conn:                c,
		id:                  id,
		closeNotificationCh: make(chan bool),
	}
	if keepAlive {
		// create a new heartbeat object
		sockconn.heartBeat = newHeartBeat(sockconn, heartbeatPeriod, pingwriteWait, appData)
	}
	return sockconn
}

func (c *SocketConnection) setType(t ConnectionType) {
	c.connType = t
}

func (c *SocketConnection) setSocketServer(s *WebsocketServer) {
	c.socketserver = s
}

func (c *SocketConnection) setSocketClient(s *WebsocketClient) {
	c.socketclient = s
}

func (c *SocketConnection) SocketServer() *WebsocketServer {
	return c.socketserver
}

func (c *SocketConnection) SocketClient() *WebsocketClient {
	return c.socketclient
}

func (c *SocketConnection) Type() ConnectionType {
	return c.connType
}

func (c *SocketConnection) HeartBeat() *heartbeat {
	return c.heartBeat
}

func (c *SocketConnection) handleFailure() {
	log.Println("connection failure detected")
	if c.heartBeat != nil {
		c.heartBeat.stop()
	}
	fmt.Println("I am here")
	if c.connType == ServerSide {
		// ToAsk: is it OK to have this given that the channel might never be consumed if the handler has exited
		go c.once.Do(func() { c.closeNotificationCh <- true })
		log.Printf("removing connection with id %s from connection list\n", c.id)
		// remove this connection from the server connectionMap
		c.socketserver.unregister <- c
	} else {
		// try to reconnect with exponential backoff
		c.socketclient.Connect()
	}
}

// Close provides a graceful termination of the connection
func (c *SocketConnection) Close() error {
	// stop the heartbeat protocol
	if c.heartBeat != nil {
		c.heartBeat.stop()
	}
	// some more stuff to do before closing the connection
	return c.conn.Close()
}

func (c *SocketConnection) ID() string {
	return c.id
}

// WriteRequest writes a request to the underlying connection
func (c *SocketConnection) WriteRequest(req *http.Request) error {
	var err error
	var w io.WriteCloser
	if w, err = c.conn.NextWriter(websocket.TextMessage); err == nil {
		defer w.Close()
		if err = req.Write(w); err == nil {
			return nil
		}
	}
	log.Printf("error: %v", err)
	if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
		c.handleFailure()
	}
	return err
}

// readRequest is how the connection read an http request and returns an http response
func (c *SocketConnection) readRequest(ctx context.Context) (*response, error) {
	var reader io.Reader
	var err error
	var mt int
	var req *http.Request

	if mt, reader, err = c.conn.NextReader(); err == nil {
		if mt != websocket.TextMessage {
			log.Println("error: not a text message")
			return nil, errors.New("not a text message")
		}
		if req, err = http.ReadRequest(bufio.NewReader(reader)); err == nil {
			ctx, cancelCtx := context.WithCancel(ctx)
			req = req.WithContext(ctx)
			w := &response{
				conn:          c,
				cancelCtx:     cancelCtx,
				req:           req,
				reqBody:       req.Body,
				handlerHeader: make(http.Header),
				contentLength: -1,
			}
			w.cw.res = w
			bufw, err := c.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				// handle failure
			}
			w.cw.writer = bufw
			w.w = newBufioWriterSize(&w.cw, bufferBeforeChunkingSize)
			return w, nil
		}
	}
	return nil, err
}

// ResponseReader is a specialized reader that reads streams on websockets
type ResponseReader struct {
	c *SocketConnection
	r io.Reader
}

func (rr *ResponseReader) Read(p []byte) (int, error) {
	if rr.r == nil {
		_, reader, err := rr.c.conn.NextReader()
		if err != nil {
			return 0, err
		}
		rr.r = reader
	}
	count, err := rr.r.Read(p)
	// this is a fake EOF sent because of a flush at the server side
	if count == 0 && err == io.EOF {
		_, reader, err := rr.c.conn.NextReader()
		if err != nil {
			log.Printf("Err: %v", err)
			return 0, err
		}
		rr.r = reader
		// the correct count and EOF if any will be sent from here
		return rr.r.Read(p)
	}
	return count, err
}

func newResponseReader(c *SocketConnection) io.Reader {
	return &ResponseReader{
		c: c,
	}
}

// ReadResponse reads a response from the underlying connection
func (c *SocketConnection) ReadResponse() (*http.Response, error) {
	var err error
	var resp *http.Response
	if resp, err = http.ReadResponse(bufio.NewReader(newResponseReader(c)), nil); err == nil {
		return resp, nil
	}

	log.Printf("error: %v", err)
	if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
		c.handleFailure()
	}
	return nil, err
}

// WriteRaw writes generic bytes to the connection
func (c *SocketConnection) WriteRaw(b []byte) error {
	if err := c.conn.WriteMessage(websocket.TextMessage, b); err != nil {
		log.Printf("error: %v", err)
		if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
			c.handleFailure()
		}
		return err
	}
	return nil
}

func (c *SocketConnection) Serve(ctx context.Context, hdlr http.Handler) error {
	ctx, cancelCtx := context.WithCancel(ctx)
	defer cancelCtx()
	for {
		resp, err := c.readRequest(ctx)
		if err != nil {
			log.Println("error reading from connection")
			return err
		}
		hdlr.ServeHTTP(resp, resp.req)
		resp.cancelCtx()
		resp.finishRequest()
	}
}
<commit_msg>removing WriteRaw as it is not being used<commit_after>package restwebsocket

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// ConnectionType is the socket connection type, it can either be a serverside connection or a clientside connection
type ConnectionType int

const (
	// ServerSide means this connection is owned by a websocket-server
	ServerSide ConnectionType = iota
	// ClientSide means this connection is owned by a websocket-client
	ClientSide
)

const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// Time allowed to read the next pong message from the peer.
	pongWait = 20 * time.Second

	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10

	// Maximum message size allowed from peer.
	maxMessageSize = 512
)

// SocketConnection is a wrapper around the websocket connection to handle http
type SocketConnection struct {
	socketserver        *WebsocketServer
	socketclient        *WebsocketClient
	connType            ConnectionType
	conn                *websocket.Conn
	id                  string
	heartBeat           *heartbeat
	closeNotificationCh chan bool
	once                sync.Once
}

// NewSocketConnection creates a new socket connection
func NewSocketConnection(c *websocket.Conn, id string, keepAlive bool, pingHdlr, pongHdlr func(string) error, appData []byte) *SocketConnection {
	// Default ping handler is to send back a pong control message with the same application data
	// Default pong handler is to do nothing
	// websocket protocol mentions that the pong message should reply back with the exact appData recieved from the ping message
	if pingHdlr != nil {
		c.SetPingHandler(pingHdlr)
	}
	if pongHdlr != nil {
		c.SetPongHandler(pingHdlr)
	}

	sockconn := &SocketConnection{
		conn:                c,
		id:                  id,
		closeNotificationCh: make(chan bool),
	}
	if keepAlive {
		// create a new heartbeat object
		sockconn.heartBeat = newHeartBeat(sockconn, heartbeatPeriod, pingwriteWait, appData)
	}
	return sockconn
}

func (c *SocketConnection) setType(t ConnectionType) {
	c.connType = t
}

func (c *SocketConnection) setSocketServer(s *WebsocketServer) {
	c.socketserver = s
}

func (c *SocketConnection) setSocketClient(s *WebsocketClient) {
	c.socketclient = s
}

func (c *SocketConnection) SocketServer() *WebsocketServer {
	return c.socketserver
}

func (c *SocketConnection) SocketClient() *WebsocketClient {
	return c.socketclient
}

func (c *SocketConnection) Type() ConnectionType {
	return c.connType
}

func (c *SocketConnection) HeartBeat() *heartbeat {
	return c.heartBeat
}

func (c *SocketConnection) handleFailure() {
	log.Println("connection failure detected")
	if c.heartBeat != nil {
		c.heartBeat.stop()
	}
	fmt.Println("I am here")
	if c.connType == ServerSide {
		// ToAsk: is it OK to have this given that the channel might never be consumed if the handler has exited
		go c.once.Do(func() { c.closeNotificationCh <- true })
		log.Printf("removing connection with id %s from connection list\n", c.id)
		// remove this connection from the server connectionMap
		c.socketserver.unregister <- c
	} else {
		// try to reconnect with exponential backoff
		c.socketclient.Connect()
	}
}

// Close provides a graceful termination of the connection
func (c *SocketConnection) Close() error {
	// stop the heartbeat protocol
	if c.heartBeat != nil {
		c.heartBeat.stop()
	}
	// some more stuff to do before closing the connection
	return c.conn.Close()
}

func (c *SocketConnection) ID() string {
	return c.id
}

// WriteRequest writes a request to the underlying connection
func (c *SocketConnection) WriteRequest(req *http.Request) error {
	var err error
	var w io.WriteCloser
	if w, err = c.conn.NextWriter(websocket.TextMessage); err == nil {
		defer w.Close()
		if err = req.Write(w); err == nil {
			return nil
		}
	}
	log.Printf("error: %v", err)
	if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
		c.handleFailure()
	}
	return err
}

// readRequest is how the connection read an http request and returns an http response
func (c *SocketConnection) readRequest(ctx context.Context) (*response, error) {
	var reader io.Reader
	var err error
	var mt int
	var req *http.Request

	if mt, reader, err = c.conn.NextReader(); err == nil {
		if mt != websocket.TextMessage {
			log.Println("error: not a text message")
			return nil, errors.New("not a text message")
		}
		if req, err = http.ReadRequest(bufio.NewReader(reader)); err == nil {
			ctx, cancelCtx := context.WithCancel(ctx)
			req = req.WithContext(ctx)
			w := &response{
				conn:          c,
				cancelCtx:     cancelCtx,
				req:           req,
				reqBody:       req.Body,
				handlerHeader: make(http.Header),
				contentLength: -1,
			}
			w.cw.res = w
			bufw, err := c.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				// handle failure
			}
			w.cw.writer = bufw
			w.w = newBufioWriterSize(&w.cw, bufferBeforeChunkingSize)
			return w, nil
		}
	}
	return nil, err
}

// ResponseReader is a specialized reader that reads streams on websockets
type ResponseReader struct {
	c *SocketConnection
	r io.Reader
}

func (rr *ResponseReader) Read(p []byte) (int, error) {
	if rr.r == nil {
		_, reader, err := rr.c.conn.NextReader()
		if err != nil {
			return 0, err
		}
		rr.r = reader
	}
	count, err := rr.r.Read(p)
	// this is a fake EOF sent because of a flush at the server side
	if count == 0 && err == io.EOF {
		_, reader, err := rr.c.conn.NextReader()
		if err != nil {
			log.Printf("Err: %v", err)
			return 0, err
		}
		rr.r = reader
		// the correct count and EOF if any will be sent from here
		return rr.r.Read(p)
	}
	return count, err
}

func newResponseReader(c *SocketConnection) io.Reader {
	return &ResponseReader{
		c: c,
	}
}

// ReadResponse reads a response from the underlying connection
func (c *SocketConnection) ReadResponse() (*http.Response, error) {
	var err error
	var resp *http.Response
	if resp, err = http.ReadResponse(bufio.NewReader(newResponseReader(c)), nil); err == nil {
		return resp, nil
	}

	log.Printf("error: %v", err)
	if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
		c.handleFailure()
	}
	return nil, err
}

func (c *SocketConnection) Serve(ctx context.Context, hdlr http.Handler) error {
	ctx, cancelCtx := context.WithCancel(ctx)
	defer cancelCtx()
	for {
		resp, err := c.readRequest(ctx)
		if err != nil {
			log.Println("error reading from connection")
			return err
		}
		hdlr.ServeHTTP(resp, resp.req)
		resp.cancelCtx()
		resp.finishRequest()
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2013 CoreOS Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package dbus

import (
	"github.com/guelfey/go.dbus"
)

func (c *Conn) initJobs() {
	c.jobListener.jobs = make(map[dbus.ObjectPath]chan string)
}

func (c *Conn) jobComplete(signal *dbus.Signal) {
	var id uint32
	var job dbus.ObjectPath
	var unit string
	var result string
	dbus.Store(signal.Body, &id, &job, &unit, &result)
	c.jobListener.Lock()
	out, ok := c.jobListener.jobs[job]
	if ok {
		out <- result
	}
	c.jobListener.Unlock()
}

func (c *Conn) startJob(job string, args ...interface{}) (<-chan string, error) {
	c.jobListener.Lock()
	defer c.jobListener.Unlock()

	ch := make(chan string, 1)
	var path dbus.ObjectPath
	err := c.sysobj.Call(job, 0, args...).Store(&path)
	if err != nil {
		return nil, err
	}
	c.jobListener.jobs[path] = ch
	return ch, nil
}

func (c *Conn) runJob(job string, args ...interface{}) (string, error) {
	respCh, err := c.startJob(job, args...)
	if err != nil {
		return "", err
	}
	return <-respCh, nil
}

// StartUnit enqeues a start job and depending jobs, if any (unless otherwise
// specified by the mode string).
//
// Takes the unit to activate, plus a mode string. The mode needs to be one of
// replace, fail, isolate, ignore-dependencies, ignore-requirements. If
// "replace" the call will start the unit and its dependencies, possibly
// replacing already queued jobs that conflict with this. If "fail" the call
// will start the unit and its dependencies, but will fail if this would change
// an already queued job. If "isolate" the call will start the unit in question
// and terminate all units that aren't dependencies of it. If
// "ignore-dependencies" it will start a unit but ignore all its dependencies.
// If "ignore-requirements" it will start a unit but only ignore the
// requirement dependencies. It is not recommended to make use of the latter
// two options.
//
// Result string: one of done, canceled, timeout, failed, dependency, skipped.
// done indicates successful execution of a job. canceled indicates that a job
// has been canceled  before it finished execution. timeout indicates that the
// job timeout was reached. failed indicates that the job failed. dependency
// indicates that a job this job has been depending on failed and the job hence
// has been removed too. skipped indicates that a job was skipped because it
// didn't apply to the units current state.
func (c *Conn) StartUnit(name string, mode string) (string, error) {
	return c.runJob("StartUnit", name, mode)
}

// StopUnit is similar to StartUnit but stops the specified unit rather
// than starting it.
func (c *Conn) StopUnit(name string, mode string) (string, error) {
	return c.runJob("StopUnit", name, mode)
}

// ReloadUnit reloads a unit.  Reloading is done only if the unit is already running and fails otherwise.
func (c *Conn) ReloadUnit(name string, mode string) (string, error) {
	return c.runJob("ReloadUnit", name, mode)
}

// RestartUnit restarts a service.  If a service is restarted that isn't
// running it will be started.
func (c *Conn) RestartUnit(name string, mode string) (string, error) {
	return c.runJob("RestartUnit", name, mode)
}

// TryRestartUnit is like RestartUnit, except that a service that isn't running
// is not affected by the restart.
func (c *Conn) TryRestartUnit(name string, mode string) (string, error) {
	return c.runJob("TryRestartUnit", name, mode)
}

// ReloadOrRestart attempts a reload if the unit supports it and use a restart
// otherwise.
func (c *Conn) ReloadOrRestartUnit(name string, mode string) (string, error) {
	return c.runJob("ReloadOrRestartUnit", name, mode)
}

// ReloadOrTryRestart attempts a reload if the unit supports it and use a "Try"
// flavored restart otherwise.
func (c *Conn) ReloadOrTryRestartUnit(name string, mode string) (string, error) {
	return c.runJob("ReloadOrTryRestartUnit", name, mode)
}

// StartTransientUnit() may be used to create and start a transient unit, which
// will be released as soon as it is not running or referenced anymore or the
// system is rebooted. name is the unit name including suffix, and must be
// unique. mode is the same as in StartUnit(), properties contains properties
// of the unit.
func (c *Conn) StartTransientUnit(name string, mode string, properties ...Property) (string, error) {
	// the dbus interface for this method does not use the last argument and
	// should simply be given an empty list.  We use a concrete type here
	// (instead of the more appropriate interface{}) to satisfy the dbus library.
	return c.runJob("StartTransientUnit", name, mode, properties, make([]string, 0))
}

// KillUnit takes the unit name and a UNIX signal number to send.  All of the unit's
// processes are killed.
func (c *Conn) KillUnit(name string, signal int32) {
	c.sysobj.Call("KillUnit", 0, name, "all", signal).Store()
}

// GetUnitProperties takes the unit name and returns all of its dbus object properties.
func (c *Conn) GetUnitProperties(unit string) (map[string]interface{}, error) {
	var err error
	var props map[string]dbus.Variant

	path := ObjectPath("/org/freedesktop/systemd1/unit/" + unit)

	obj := c.sysconn.Object("org.freedesktop.systemd1", path)
	err = obj.Call("org.freedesktop.DBus.Properties.GetAll", 0, "org.freedesktop.systemd1.Unit").Store(&props)
	if err != nil {
		return nil, err
	}

	out := make(map[string]interface{}, len(props))
	for k, v := range props {
		out[k] = v.Value()
	}

	return out, nil
}

// ListUnits returns an array with all currently loaded units. Note that
// units may be known by multiple names at the same time, and hence there might
// be more unit names loaded than actual units behind them.
func (c *Conn) ListUnits() ([]UnitStatus, error) {
	result := make([][]interface{}, 0)
	err := c.sysobj.Call("ListUnits", 0).Store(&result)
	if err != nil {
		return nil, err
	}

	resultInterface := make([]interface{}, len(result))
	for i := range result {
		resultInterface[i] = result[i]
	}

	status := make([]UnitStatus, len(result))
	statusInterface := make([]interface{}, len(status))
	for i := range status {
		statusInterface[i] = &status[i]
	}

	err = dbus.Store(resultInterface, statusInterface...)
	if err != nil {
		return nil, err
	}

	return status, nil
}

type UnitStatus struct {
	Name        string          // The primary unit name as string
	Description string          // The human readable description string
	LoadState   string          // The load state (i.e. whether the unit file has been loaded successfully)
	ActiveState string          // The active state (i.e. whether the unit is currently started or not)
	SubState    string          // The sub state (a more fine-grained version of the active state that is specific to the unit type, which the active state is not)
	Followed    string          // A unit that is being followed in its state by this unit, if there is any, otherwise the empty string.
	Path        dbus.ObjectPath // The unit object path
	JobId       uint32          // If there is a job queued for the job unit the numeric job id, 0 otherwise
	JobType     string          // The job type as string
	JobPath     dbus.ObjectPath // The job object path
}

// EnableUnitFiles() may be used to enable one or more units in the system (by
// creating symlinks to them in /etc or /run).
//
// It takes a list of unit files to enable (either just file names or full
// absolute paths if the unit files are residing outside the usual unit
// search paths), and two booleans: the first controls whether the unit shall
// be enabled for runtime only (true, /run), or persistently (false, /etc).
// The second one controls whether symlinks pointing to other units shall
// be replaced if necessary.
//
// This call returns one boolean and an array with the changes made. The
// boolean signals whether the unit files contained any enablement
// information (i.e. an [Install]) section. The changes list consists of
// structures with three strings: the type of the change (one of symlink
// or unlink), the file name of the symlink and the destination of the
// symlink.
func (c *Conn) EnableUnitFiles(files []string, runtime bool, force bool) (bool, []EnableUnitFileChange, error) {
	var carries_install_info bool

	result := make([][]interface{}, 0)
	err := c.sysobj.Call("EnableUnitFiles", 0, files, runtime, force).Store(&carries_install_info, &result)
	if err != nil {
		return false, nil, err
	}

	resultInterface := make([]interface{}, len(result))
	for i := range result {
		resultInterface[i] = result[i]
	}

	changes := make([]EnableUnitFileChange, len(result))
	changesInterface := make([]interface{}, len(changes))
	for i := range changes {
		changesInterface[i] = &changes[i]
	}

	err = dbus.Store(resultInterface, changesInterface...)
	if err != nil {
		return false, nil, err
	}

	return carries_install_info, changes, nil
}

type EnableUnitFileChange struct {
	Type        string // Type of the change (one of symlink or unlink)
	Filename    string // File name of the symlink
	Destination string // Destination of the symlink
}
<commit_msg>Validate ObjectPath before use in Conn.GetUnitProperties Fixes issue #13.<commit_after>/*
Copyright 2013 CoreOS Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package dbus

import (
	"errors"
	"github.com/guelfey/go.dbus"
)

func (c *Conn) initJobs() {
	c.jobListener.jobs = make(map[dbus.ObjectPath]chan string)
}

func (c *Conn) jobComplete(signal *dbus.Signal) {
	var id uint32
	var job dbus.ObjectPath
	var unit string
	var result string
	dbus.Store(signal.Body, &id, &job, &unit, &result)
	c.jobListener.Lock()
	out, ok := c.jobListener.jobs[job]
	if ok {
		out <- result
	}
	c.jobListener.Unlock()
}

func (c *Conn) startJob(job string, args ...interface{}) (<-chan string, error) {
	c.jobListener.Lock()
	defer c.jobListener.Unlock()

	ch := make(chan string, 1)
	var path dbus.ObjectPath
	err := c.sysobj.Call(job, 0, args...).Store(&path)
	if err != nil {
		return nil, err
	}
	c.jobListener.jobs[path] = ch
	return ch, nil
}

func (c *Conn) runJob(job string, args ...interface{}) (string, error) {
	respCh, err := c.startJob(job, args...)
	if err != nil {
		return "", err
	}
	return <-respCh, nil
}

// StartUnit enqeues a start job and depending jobs, if any (unless otherwise
// specified by the mode string).
//
// Takes the unit to activate, plus a mode string. The mode needs to be one of
// replace, fail, isolate, ignore-dependencies, ignore-requirements. If
// "replace" the call will start the unit and its dependencies, possibly
// replacing already queued jobs that conflict with this. If "fail" the call
// will start the unit and its dependencies, but will fail if this would change
// an already queued job. If "isolate" the call will start the unit in question
// and terminate all units that aren't dependencies of it. If
// "ignore-dependencies" it will start a unit but ignore all its dependencies.
// If "ignore-requirements" it will start a unit but only ignore the
// requirement dependencies. It is not recommended to make use of the latter
// two options.
//
// Result string: one of done, canceled, timeout, failed, dependency, skipped.
// done indicates successful execution of a job. canceled indicates that a job
// has been canceled  before it finished execution. timeout indicates that the
// job timeout was reached. failed indicates that the job failed. dependency
// indicates that a job this job has been depending on failed and the job hence
// has been removed too. skipped indicates that a job was skipped because it
// didn't apply to the units current state.
func (c *Conn) StartUnit(name string, mode string) (string, error) {
	return c.runJob("StartUnit", name, mode)
}

// StopUnit is similar to StartUnit but stops the specified unit rather
// than starting it.
func (c *Conn) StopUnit(name string, mode string) (string, error) {
	return c.runJob("StopUnit", name, mode)
}

// ReloadUnit reloads a unit.  Reloading is done only if the unit is already running and fails otherwise.
func (c *Conn) ReloadUnit(name string, mode string) (string, error) {
	return c.runJob("ReloadUnit", name, mode)
}

// RestartUnit restarts a service.  If a service is restarted that isn't
// running it will be started.
func (c *Conn) RestartUnit(name string, mode string) (string, error) {
	return c.runJob("RestartUnit", name, mode)
}

// TryRestartUnit is like RestartUnit, except that a service that isn't running
// is not affected by the restart.
func (c *Conn) TryRestartUnit(name string, mode string) (string, error) {
	return c.runJob("TryRestartUnit", name, mode)
}

// ReloadOrRestart attempts a reload if the unit supports it and use a restart
// otherwise.
func (c *Conn) ReloadOrRestartUnit(name string, mode string) (string, error) {
	return c.runJob("ReloadOrRestartUnit", name, mode)
}

// ReloadOrTryRestart attempts a reload if the unit supports it and use a "Try"
// flavored restart otherwise.
func (c *Conn) ReloadOrTryRestartUnit(name string, mode string) (string, error) {
	return c.runJob("ReloadOrTryRestartUnit", name, mode)
}

// StartTransientUnit() may be used to create and start a transient unit, which
// will be released as soon as it is not running or referenced anymore or the
// system is rebooted. name is the unit name including suffix, and must be
// unique. mode is the same as in StartUnit(), properties contains properties
// of the unit.
func (c *Conn) StartTransientUnit(name string, mode string, properties ...Property) (string, error) {
	// the dbus interface for this method does not use the last argument and
	// should simply be given an empty list.  We use a concrete type here
	// (instead of the more appropriate interface{}) to satisfy the dbus library.
	return c.runJob("StartTransientUnit", name, mode, properties, make([]string, 0))
}

// KillUnit takes the unit name and a UNIX signal number to send.  All of the unit's
// processes are killed.
func (c *Conn) KillUnit(name string, signal int32) {
	c.sysobj.Call("KillUnit", 0, name, "all", signal).Store()
}

// GetUnitProperties takes the unit name and returns all of its dbus object properties.
func (c *Conn) GetUnitProperties(unit string) (map[string]interface{}, error) {
	var err error
	var props map[string]dbus.Variant

	path := ObjectPath("/org/freedesktop/systemd1/unit/" + unit)
	if !path.IsValid() {
		return nil, errors.New("invalid unit name: " + unit)
	}

	obj := c.sysconn.Object("org.freedesktop.systemd1", path)
	err = obj.Call("org.freedesktop.DBus.Properties.GetAll", 0, "org.freedesktop.systemd1.Unit").Store(&props)
	if err != nil {
		return nil, err
	}

	out := make(map[string]interface{}, len(props))
	for k, v := range props {
		out[k] = v.Value()
	}

	return out, nil
}

// ListUnits returns an array with all currently loaded units. Note that
// units may be known by multiple names at the same time, and hence there might
// be more unit names loaded than actual units behind them.
func (c *Conn) ListUnits() ([]UnitStatus, error) {
	result := make([][]interface{}, 0)
	err := c.sysobj.Call("ListUnits", 0).Store(&result)
	if err != nil {
		return nil, err
	}

	resultInterface := make([]interface{}, len(result))
	for i := range result {
		resultInterface[i] = result[i]
	}

	status := make([]UnitStatus, len(result))
	statusInterface := make([]interface{}, len(status))
	for i := range status {
		statusInterface[i] = &status[i]
	}

	err = dbus.Store(resultInterface, statusInterface...)
	if err != nil {
		return nil, err
	}

	return status, nil
}

type UnitStatus struct {
	Name        string          // The primary unit name as string
	Description string          // The human readable description string
	LoadState   string          // The load state (i.e. whether the unit file has been loaded successfully)
	ActiveState string          // The active state (i.e. whether the unit is currently started or not)
	SubState    string          // The sub state (a more fine-grained version of the active state that is specific to the unit type, which the active state is not)
	Followed    string          // A unit that is being followed in its state by this unit, if there is any, otherwise the empty string.
	Path        dbus.ObjectPath // The unit object path
	JobId       uint32          // If there is a job queued for the job unit the numeric job id, 0 otherwise
	JobType     string          // The job type as string
	JobPath     dbus.ObjectPath // The job object path
}

// EnableUnitFiles() may be used to enable one or more units in the system (by
// creating symlinks to them in /etc or /run).
//
// It takes a list of unit files to enable (either just file names or full
// absolute paths if the unit files are residing outside the usual unit
// search paths), and two booleans: the first controls whether the unit shall
// be enabled for runtime only (true, /run), or persistently (false, /etc).
// The second one controls whether symlinks pointing to other units shall
// be replaced if necessary.
//
// This call returns one boolean and an array with the changes made. The
// boolean signals whether the unit files contained any enablement
// information (i.e. an [Install]) section. The changes list consists of
// structures with three strings: the type of the change (one of symlink
// or unlink), the file name of the symlink and the destination of the
// symlink.
func (c *Conn) EnableUnitFiles(files []string, runtime bool, force bool) (bool, []EnableUnitFileChange, error) {
	var carries_install_info bool

	result := make([][]interface{}, 0)
	err := c.sysobj.Call("EnableUnitFiles", 0, files, runtime, force).Store(&carries_install_info, &result)
	if err != nil {
		return false, nil, err
	}

	resultInterface := make([]interface{}, len(result))
	for i := range result {
		resultInterface[i] = result[i]
	}

	changes := make([]EnableUnitFileChange, len(result))
	changesInterface := make([]interface{}, len(changes))
	for i := range changes {
		changesInterface[i] = &changes[i]
	}

	err = dbus.Store(resultInterface, changesInterface...)
	if err != nil {
		return false, nil, err
	}

	return carries_install_info, changes, nil
}

type EnableUnitFileChange struct {
	Type        string // Type of the change (one of symlink or unlink)
	Filename    string // File name of the symlink
	Destination string // Destination of the symlink
}
<|endoftext|>


<commit_before>package session

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/gob"
	"errors"
	"net/http"
)

func init() {
	gob.Register(M{})
}

type Storage interface {
	Clean(*Session) error
	Flush(*Session) error
	LoadTo(*http.Request, *Session) error
}

const (
	keySize    = 16
	aesKeySize = 32
)

var (
	defaultKey = genKey(keySize)
)

func SetKey(key []byte) {
	defaultKey = key[:keySize]
}

func GetKey() []byte {
	return defaultKey
}

func encrypt(key, value []byte) ([]byte, error) {
	if len(key) < aesKeySize-keySize {
		return nil, errTooShort
	}
	key = append(key[:aesKeySize-keySize], defaultKey...)
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	iv := make([]byte, block.BlockSize())
	rand.Read(iv)
	stream := cipher.NewCTR(block, iv)
	stream.XORKeyStream(value, value)
	return append(iv, value...), nil
}

var errTooShort = errors.New("Too short")

func decrypt(key, value []byte) ([]byte, error) {
	if len(key) < aesKeySize - keySize {
		return nil, errTooShort
	}
	key = append(key[:aesKeySize-keySize], defaultKey...)
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	if len(value) > block.BlockSize() {
		iv := value[:block.BlockSize()]
		value = value[block.BlockSize():]
		stream := cipher.NewCTR(block, iv)
		stream.XORKeyStream(value, value)
		return value, nil
	}
	return nil, errTooShort
}

func decoding(key []byte, src string, dst *M) error {
	// 1. base64 decoding
	buf, err := base64.StdEncoding.DecodeString(src)
	if err != nil {
		return err
	}
	// 2. cypto decoding
	buf, err = decrypt(key, buf)
	if err != nil {
		return err
	}
	// 3. gob decoding
	g := gob.NewDecoder(bytes.NewBuffer(buf))
	if err = g.Decode(&dst); err != nil {
		return err
	}
	return nil
}

func encoding(key []byte, src map[string]interface{}) (string, error){
	// 1. gob encoding
	var buf bytes.Buffer
	g := gob.NewEncoder(&buf)
	if err := g.Encode(src); err != nil {
		return "", err
	}
	// 2. cypto encoding
	ciphertext, err := encrypt(key, buf.Bytes())
	if err != nil {
		return "", err
	}
	// 3. base64 encoding
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}
<commit_msg>clearing errors<commit_after>package session

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/gob"
	"errors"
	"net/http"
)

func init() {
	gob.Register(M{})
}

type Storage interface {
	Clean(*Session) error
	Flush(*Session) error
	LoadTo(*http.Request, *Session) error
}

const (
	keySize    = 16
	aesKeySize = 32
)

var (
	defaultKey = genKey(keySize)
)

func SetKey(key []byte) {
	defaultKey = key[:keySize]
}

func GetKey() []byte {
	return defaultKey
}

var errKeyTooShort = errors.New("The key is too short")
var errValueTooShort = errors.New("The block is too short")

func encrypt(key, value []byte) ([]byte, error) {
	if len(key) < aesKeySize-keySize {
		return nil, errKeyTooShort
	}
	key = append(key[:aesKeySize-keySize], defaultKey...)
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	iv := make([]byte, block.BlockSize())
	rand.Read(iv)
	stream := cipher.NewCTR(block, iv)
	stream.XORKeyStream(value, value)
	return append(iv, value...), nil
}

func decrypt(key, value []byte) ([]byte, error) {
	if len(key) < aesKeySize - keySize {
		return nil, errKeyTooShort
	}
	key = append(key[:aesKeySize-keySize], defaultKey...)
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	if len(value) > block.BlockSize() {
		iv := value[:block.BlockSize()]
		value = value[block.BlockSize():]
		stream := cipher.NewCTR(block, iv)
		stream.XORKeyStream(value, value)
		return value, nil
	}
	return nil, errValueTooShort
}

func decoding(key []byte, src string, dst *M) error {
	// 1. base64 decoding
	buf, err := base64.StdEncoding.DecodeString(src)
	if err != nil {
		return err
	}
	// 2. cypto decoding
	buf, err = decrypt(key, buf)
	if err != nil {
		return err
	}
	// 3. gob decoding
	g := gob.NewDecoder(bytes.NewBuffer(buf))
	if err = g.Decode(&dst); err != nil {
		return err
	}
	return nil
}

func encoding(key []byte, src map[string]interface{}) (string, error){
	// 1. gob encoding
	var buf bytes.Buffer
	g := gob.NewEncoder(&buf)
	if err := g.Encode(src); err != nil {
		return "", err
	}
	// 2. cypto encoding
	ciphertext, err := encrypt(key, buf.Bytes())
	if err != nil {
		return "", err
	}
	// 3. base64 encoding
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}
<|endoftext|>


<commit_before>package consumergroup

import (
	"github.com/wvanbergen/kazoo-go"
	"testing"
)

func Test_PartitionDivision(t *testing.T) {

	consumers := kazoo.ConsumergroupInstanceList{
		&kazoo.ConsumergroupInstance{ID: "consumer1"},
		&kazoo.ConsumergroupInstance{ID: "consumer2"},
	}

	partitions := []partitionLeader{
		partitionLeader{id: 0, leader: 1, partition: &kazoo.Partition{ID: 0}},
		partitionLeader{id: 1, leader: 2, partition: &kazoo.Partition{ID: 1}},
		partitionLeader{id: 2, leader: 1, partition: &kazoo.Partition{ID: 2}},
		partitionLeader{id: 3, leader: 2, partition: &kazoo.Partition{ID: 3}},
		partitionLeader{id: 4, leader: 1, partition: &kazoo.Partition{ID: 4}},
	}

	division := dividePartitionsBetweenConsumers(consumers, partitions)

	if len(division["consumer1"]) != 3 || division["consumer1"][0].ID != 0 || division["consumer1"][1].ID != 2 || division["consumer1"][2].ID != 4 {
		t.Error("Consumer 1 should end up with partition 0, 2, and 4")
	}

	if len(division["consumer2"]) != 2 || division["consumer2"][0].ID != 1 || division["consumer2"][1].ID != 3 {
		t.Error("Consumer 2 should end up with partition 1 and 3")
	}
}
<commit_msg>Test to ensure that partitions are divided optimally among consumers. #61 / #63<commit_after>package consumergroup

import (
	"fmt"
	"github.com/wvanbergen/kazoo-go"
	"testing"
)

func createTestConsumerGroupInstanceList(size int) kazoo.ConsumergroupInstanceList {
	k := make(kazoo.ConsumergroupInstanceList, size)
	for i := range k {
		k[i] = &kazoo.ConsumergroupInstance{ID: fmt.Sprintf("consumer%d", i)}
	}
	return k
}

func createTestPartitions(count int) []partitionLeader {
	p := make([]partitionLeader, count)
	for i := range p {
		p[i] = partitionLeader{id: int32(i), leader: 1, partition: &kazoo.Partition{ID: int32(i)}}
	}
	return p
}

func Test_PartitionDivision(t *testing.T) {
	consumerPartitionTestCases := [][2]int{
		// {number of Consumers, number of Partitions}
		[2]int{2, 5},
		[2]int{5, 2},
		[2]int{9, 32},
		[2]int{10, 50},
	}
	for _, v := range consumerPartitionTestCases {
		consumers := createTestConsumerGroupInstanceList(v[0])
		partitions := createTestPartitions(v[1])
		division := dividePartitionsBetweenConsumers(consumers, partitions)

		// make sure every partition is used once
		grouping := make(map[int32]struct{})
		maxConsumed := 0
		minConsumed := len(partitions) + 1
		for _, v := range division {
			if len(v) > maxConsumed {
				maxConsumed = len(v)
			}
			if len(v) < minConsumed {
				minConsumed = len(v)
			}
			for _, partition := range v {
				if _, ok := grouping[partition.ID]; ok {
					t.Errorf("PartitionDivision: Partition %v was assigned more than once!", partition.ID)
				} else {
					grouping[partition.ID] = struct{}{}
				}
			}
		}
		if len(grouping) != len(partitions) {
			t.Errorf("PartitionDivision: Expected to divide %d partitions among consumers, but only %d partitions were consumed.", len(partitions), len(grouping))
		}
		if (maxConsumed - minConsumed) > 1 {
			t.Errorf("PartitionDivision: Partitions weren't divided evenly, consumers shouldn't have a difference of more than 1 in the number of partitions consumed (was %d).", maxConsumed-minConsumed)
		}
		if minConsumed > 1 && len(consumers) != len(division) {
			t.Errorf("PartitionDivision: Partitions weren't divided evenly, some consumers didn't get any paritions even though there were %d partitions and %d consumers.", len(partitions), len(consumers))
		}
	}
}
<|endoftext|>


<commit_before>package drone

type (
	// User represents a user account.
	User struct {
		ID     int64  `json:"id"`
		Login  string `json:"login"`
		Email  string `json:"email"`
		Avatar string `json:"avatar_url"`
		Active bool   `json:"active"`
		Admin  bool   `json:"admin"`
	}

	// Repo represents a repository.
	Repo struct {
		ID          int64  `json:"id,omitempty"`
		Owner       string `json:"owner"`
		Name        string `json:"name"`
		FullName    string `json:"full_name"`
		Avatar      string `json:"avatar_url,omitempty"`
		Link        string `json:"link_url,omitempty"`
		Kind        string `json:"scm,omitempty"`
		Clone       string `json:"clone_url,omitempty"`
		Branch      string `json:"default_branch,omitempty"`
		Timeout     int64  `json:"timeout,omitempty"`
		Visibility  string `json:"visibility"`
		IsPrivate   bool   `json:"private,omitempty"`
		IsTrusted   bool   `json:"trusted"`
		IsStarred   bool   `json:"starred,omitempty"`
		IsGated     bool   `json:"gated"`
		AllowPull   bool   `json:"allow_pr"`
		AllowPush   bool   `json:"allow_push"`
		AllowDeploy bool   `json:"allow_deploys"`
		AllowTag    bool   `json:"allow_tags"`
		Config      string `json:"config_file"`
	}

	// RepoPatch defines a repository patch request.
	RepoPatch struct {
		Config       *string `json:"config_file,omitempty"`
		IsTrusted    *bool   `json:"trusted,omitempty"`
		IsGated      *bool   `json:"gated,omitempty"`
		Timeout      *int64  `json:"timeout,omitempty"`
		Visibility   *string `json:"visibility"`
		AllowPull    *bool   `json:"allow_pr,omitempty"`
		AllowPush    *bool   `json:"allow_push,omitempty"`
		AllowDeploy  *bool   `json:"allow_deploy,omitempty"`
		AllowTag     *bool   `json:"allow_tag,omitempty"`
		BuildCounter *int    `json:"build_counter,omitempty"`
	}

	// Build defines a build object.
	Build struct {
		ID        int64   `json:"id"`
		Number    int     `json:"number"`
		Parent    int     `json:"parent"`
		Event     string  `json:"event"`
		Status    string  `json:"status"`
		Error     string  `json:"error"`
		Enqueued  int64   `json:"enqueued_at"`
		Created   int64   `json:"created_at"`
		Started   int64   `json:"started_at"`
		Finished  int64   `json:"finished_at"`
		Deploy    string  `json:"deploy_to"`
		Commit    string  `json:"commit"`
		Branch    string  `json:"branch"`
		Ref       string  `json:"ref"`
		Refspec   string  `json:"refspec"`
		Remote    string  `json:"remote"`
		Title     string  `json:"title"`
		Message   string  `json:"message"`
		Timestamp int64   `json:"timestamp"`
		Sender    string  `json:"sender"`
		Author    string  `json:"author"`
		Avatar    string  `json:"author_avatar"`
		Email     string  `json:"author_email"`
		Link      string  `json:"link_url"`
		Reviewer  string  `json:"reviewed_by"`
		Reviewed  int64   `json:"reviewed_at"`
		Procs     []*Proc `json:"procs,omitempty"`
	}

	// Proc represents a process in the build pipeline.
	Proc struct {
		ID       int64             `json:"id"`
		PID      int               `json:"pid"`
		PPID     int               `json:"ppid"`
		PGID     int               `json:"pgid"`
		Name     string            `json:"name"`
		State    string            `json:"state"`
		Error    string            `json:"error,omitempty"`
		ExitCode int               `json:"exit_code"`
		Started  int64             `json:"start_time,omitempty"`
		Stopped  int64             `json:"end_time,omitempty"`
		Machine  string            `json:"machine,omitempty"`
		Platform string            `json:"platform,omitempty"`
		Environ  map[string]string `json:"environ,omitempty"`
		Children []*Proc           `json:"children,omitempty"`
	}

	// Registry represents a docker registry with credentials.
	Registry struct {
		ID       int64  `json:"id"`
		Address  string `json:"address"`
		Username string `json:"username"`
		Password string `json:"password,omitempty"`
		Email    string `json:"email"`
		Token    string `json:"token"`
	}

	// Secret represents a secret variable, such as a password or token.
	Secret struct {
		ID     int64    `json:"id"`
		Name   string   `json:"name"`
		Value  string   `json:"value,omitempty"`
		Images []string `json:"image"`
		Events []string `json:"event"`
	}

	// Activity represents an item in the user's feed or timeline.
	Activity struct {
		Owner    string `json:"owner"`
		Name     string `json:"name"`
		FullName string `json:"full_name"`
		Number   int    `json:"number,omitempty"`
		Event    string `json:"event,omitempty"`
		Status   string `json:"status,omitempty"`
		Created  int64  `json:"created_at,omitempty"`
		Started  int64  `json:"started_at,omitempty"`
		Finished int64  `json:"finished_at,omitempty"`
		Commit   string `json:"commit,omitempty"`
		Branch   string `json:"branch,omitempty"`
		Ref      string `json:"ref,omitempty"`
		Refspec  string `json:"refspec,omitempty"`
		Remote   string `json:"remote,omitempty"`
		Title    string `json:"title,omitempty"`
		Message  string `json:"message,omitempty"`
		Author   string `json:"author,omitempty"`
		Avatar   string `json:"author_avatar,omitempty"`
		Email    string `json:"author_email,omitempty"`
	}
)
<commit_msg>adding patching of owner and name<commit_after>package drone

type (
	// User represents a user account.
	User struct {
		ID     int64  `json:"id"`
		Login  string `json:"login"`
		Email  string `json:"email"`
		Avatar string `json:"avatar_url"`
		Active bool   `json:"active"`
		Admin  bool   `json:"admin"`
	}

	// Repo represents a repository.
	Repo struct {
		ID          int64  `json:"id,omitempty"`
		Owner       string `json:"owner"`
		Name        string `json:"name"`
		FullName    string `json:"full_name"`
		Avatar      string `json:"avatar_url,omitempty"`
		Link        string `json:"link_url,omitempty"`
		Kind        string `json:"scm,omitempty"`
		Clone       string `json:"clone_url,omitempty"`
		Branch      string `json:"default_branch,omitempty"`
		Timeout     int64  `json:"timeout,omitempty"`
		Visibility  string `json:"visibility"`
		IsPrivate   bool   `json:"private,omitempty"`
		IsTrusted   bool   `json:"trusted"`
		IsStarred   bool   `json:"starred,omitempty"`
		IsGated     bool   `json:"gated"`
		AllowPull   bool   `json:"allow_pr"`
		AllowPush   bool   `json:"allow_push"`
		AllowDeploy bool   `json:"allow_deploys"`
		AllowTag    bool   `json:"allow_tags"`
		Config      string `json:"config_file"`
	}

	// RepoPatch defines a repository patch request.
	RepoPatch struct {
		Config       *string `json:"config_file,omitempty"`
		IsTrusted    *bool   `json:"trusted,omitempty"`
		IsGated      *bool   `json:"gated,omitempty"`
		Timeout      *int64  `json:"timeout,omitempty"`
		Visibility   *string `json:"visibility"`
		AllowPull    *bool   `json:"allow_pr,omitempty"`
		AllowPush    *bool   `json:"allow_push,omitempty"`
		AllowDeploy  *bool   `json:"allow_deploy,omitempty"`
		AllowTag     *bool   `json:"allow_tag,omitempty"`
		BuildCounter *int    `json:"build_counter,omitempty"`
		Name         *string `json:"name,omitempty"`
		Owner        *string `json:"owner,omitempty"`
	}

	// Build defines a build object.
	Build struct {
		ID        int64   `json:"id"`
		Number    int     `json:"number"`
		Parent    int     `json:"parent"`
		Event     string  `json:"event"`
		Status    string  `json:"status"`
		Error     string  `json:"error"`
		Enqueued  int64   `json:"enqueued_at"`
		Created   int64   `json:"created_at"`
		Started   int64   `json:"started_at"`
		Finished  int64   `json:"finished_at"`
		Deploy    string  `json:"deploy_to"`
		Commit    string  `json:"commit"`
		Branch    string  `json:"branch"`
		Ref       string  `json:"ref"`
		Refspec   string  `json:"refspec"`
		Remote    string  `json:"remote"`
		Title     string  `json:"title"`
		Message   string  `json:"message"`
		Timestamp int64   `json:"timestamp"`
		Sender    string  `json:"sender"`
		Author    string  `json:"author"`
		Avatar    string  `json:"author_avatar"`
		Email     string  `json:"author_email"`
		Link      string  `json:"link_url"`
		Reviewer  string  `json:"reviewed_by"`
		Reviewed  int64   `json:"reviewed_at"`
		Procs     []*Proc `json:"procs,omitempty"`
	}

	// Proc represents a process in the build pipeline.
	Proc struct {
		ID       int64             `json:"id"`
		PID      int               `json:"pid"`
		PPID     int               `json:"ppid"`
		PGID     int               `json:"pgid"`
		Name     string            `json:"name"`
		State    string            `json:"state"`
		Error    string            `json:"error,omitempty"`
		ExitCode int               `json:"exit_code"`
		Started  int64             `json:"start_time,omitempty"`
		Stopped  int64             `json:"end_time,omitempty"`
		Machine  string            `json:"machine,omitempty"`
		Platform string            `json:"platform,omitempty"`
		Environ  map[string]string `json:"environ,omitempty"`
		Children []*Proc           `json:"children,omitempty"`
	}

	// Registry represents a docker registry with credentials.
	Registry struct {
		ID       int64  `json:"id"`
		Address  string `json:"address"`
		Username string `json:"username"`
		Password string `json:"password,omitempty"`
		Email    string `json:"email"`
		Token    string `json:"token"`
	}

	// Secret represents a secret variable, such as a password or token.
	Secret struct {
		ID     int64    `json:"id"`
		Name   string   `json:"name"`
		Value  string   `json:"value,omitempty"`
		Images []string `json:"image"`
		Events []string `json:"event"`
	}

	// Activity represents an item in the user's feed or timeline.
	Activity struct {
		Owner    string `json:"owner"`
		Name     string `json:"name"`
		FullName string `json:"full_name"`
		Number   int    `json:"number,omitempty"`
		Event    string `json:"event,omitempty"`
		Status   string `json:"status,omitempty"`
		Created  int64  `json:"created_at,omitempty"`
		Started  int64  `json:"started_at,omitempty"`
		Finished int64  `json:"finished_at,omitempty"`
		Commit   string `json:"commit,omitempty"`
		Branch   string `json:"branch,omitempty"`
		Ref      string `json:"ref,omitempty"`
		Refspec  string `json:"refspec,omitempty"`
		Remote   string `json:"remote,omitempty"`
		Title    string `json:"title,omitempty"`
		Message  string `json:"message,omitempty"`
		Author   string `json:"author,omitempty"`
		Avatar   string `json:"author_avatar,omitempty"`
		Email    string `json:"author_email,omitempty"`
	}
)
<|endoftext|>


<commit_before>package logrus_sentry

import (
	"encoding/json"
	"fmt"
	"runtime"
	"sync"
	"time"

	"github.com/getsentry/raven-go"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

var (
	severityMap = map[logrus.Level]raven.Severity{
		logrus.DebugLevel: raven.DEBUG,
		logrus.InfoLevel:  raven.INFO,
		logrus.WarnLevel:  raven.WARNING,
		logrus.ErrorLevel: raven.ERROR,
		logrus.FatalLevel: raven.FATAL,
		logrus.PanicLevel: raven.FATAL,
	}
)

// BufSize controls the number of logs that can be in progress before logging
// will start blocking. Set logrus_sentry.BufSize = <value> _before_ calling
// NewAsync*().
var BufSize uint = 8192

// SentryHook delivers logs to a sentry server.
type SentryHook struct {
	// Timeout sets the time to wait for a delivery error from the sentry server.
	// If this is set to zero the server will not wait for any response and will
	// consider the message correctly sent
	Timeout                 time.Duration
	StacktraceConfiguration StackTraceConfiguration

	client *raven.Client
	levels []logrus.Level

	ignoreFields map[string]struct{}
	extraFilters map[string]func(interface{}) interface{}

	asynchronous bool
	buf          chan *raven.Packet
	wg           sync.WaitGroup
	mu           sync.RWMutex
}

// The Stacktracer interface allows an error type to return a raven.Stacktrace.
type Stacktracer interface {
	GetStacktrace() *raven.Stacktrace
}

type causer interface {
	Cause() error
}

type pkgErrorStackTracer interface {
	StackTrace() errors.StackTrace
}

// StackTraceConfiguration allows for configuring stacktraces
type StackTraceConfiguration struct {
	// whether stacktraces should be enabled
	Enable bool
	// the level at which to start capturing stacktraces
	Level logrus.Level
	// how many stack frames to skip before stacktrace starts recording
	Skip int
	// the number of lines to include around a stack frame for context
	Context int
	// the prefixes that will be matched against the stack frame.
	// if the stack frame's package matches one of these prefixes
	// sentry will identify the stack frame as "in_app"
	InAppPrefixes []string
}

// NewSentryHook creates a hook to be added to an instance of logger
// and initializes the raven client.
// This method sets the timeout to 100 milliseconds.
func NewSentryHook(DSN string, levels []logrus.Level) (*SentryHook, error) {
	client, err := raven.New(DSN)
	if err != nil {
		return nil, err
	}
	return NewWithClientSentryHook(client, levels)
}

// NewWithTagsSentryHook creates a hook with tags to be added to an instance
// of logger and initializes the raven client. This method sets the timeout to
// 100 milliseconds.
func NewWithTagsSentryHook(DSN string, tags map[string]string, levels []logrus.Level) (*SentryHook, error) {
	client, err := raven.NewWithTags(DSN, tags)
	if err != nil {
		return nil, err
	}
	return NewWithClientSentryHook(client, levels)
}

// NewWithClientSentryHook creates a hook using an initialized raven client.
// This method sets the timeout to 100 milliseconds.
func NewWithClientSentryHook(client *raven.Client, levels []logrus.Level) (*SentryHook, error) {
	return &SentryHook{
		Timeout: 100 * time.Millisecond,
		StacktraceConfiguration: StackTraceConfiguration{
			Enable:        false,
			Level:         logrus.ErrorLevel,
			Skip:          5,
			Context:       0,
			InAppPrefixes: nil,
		},
		client:       client,
		levels:       levels,
		ignoreFields: make(map[string]struct{}),
		extraFilters: make(map[string]func(interface{}) interface{}),
	}, nil
}

// NewAsyncSentryHook creates a hook same as NewSentryHook, but in asynchronous
// mode. This method sets the timeout to 1000 milliseconds.
func NewAsyncSentryHook(DSN string, levels []logrus.Level) (*SentryHook, error) {
	hook, err := NewSentryHook(DSN, levels)
	return setAsync(hook), err
}

// NewAsyncWithTagsSentryHook creates a hook same as NewWithTagsSentryHook, but
// in asynchronous mode. This method sets the timeout to 1000 milliseconds.
func NewAsyncWithTagsSentryHook(DSN string, tags map[string]string, levels []logrus.Level) (*SentryHook, error) {
	hook, err := NewWithTagsSentryHook(DSN, tags, levels)
	return setAsync(hook), err
}

// NewAsyncWithClientSentryHook creates a hook same as NewWithClientSentryHook,
// but in asynchronous mode. This method sets the timeout to 1000 milliseconds.
func NewAsyncWithClientSentryHook(client *raven.Client, levels []logrus.Level) (*SentryHook, error) {
	hook, err := NewWithClientSentryHook(client, levels)
	return setAsync(hook), err
}

func setAsync(hook *SentryHook) *SentryHook {
	if hook == nil {
		return nil
	}
	hook.Timeout = 1 * time.Second
	hook.asynchronous = true
	hook.buf = make(chan *raven.Packet, BufSize)
	go hook.fire() // Log in background
	return hook
}

// Fire is called when an event should be sent to sentry
// Special fields that sentry uses to give more information to the server
// are extracted from entry.Data (if they are found)
// These fields are: error, logger, server_name, http_request, tags
func (hook *SentryHook) Fire(entry *logrus.Entry) error {
	hook.mu.RLock() // Allow multiple go routines to log simultaneously
	defer hook.mu.RUnlock()
	packet := raven.NewPacket(entry.Message)
	packet.Timestamp = raven.Timestamp(entry.Time)
	packet.Level = severityMap[entry.Level]
	packet.Platform = "go"

	df := newDataField(entry.Data)

	// set special fields
	if logger, ok := df.getLogger(); ok {
		packet.Logger = logger
	}
	if serverName, ok := df.getServerName(); ok {
		packet.ServerName = serverName
	}
	if eventID, ok := df.getEventID(); ok {
		packet.EventID = eventID
	}
	if tags, ok := df.getTags(); ok {
		packet.Tags = tags
	}
	if fingerprint, ok := df.getFingerprint(); ok {
		packet.Fingerprint = fingerprint
	}
	if req, ok := df.getHTTPRequest(); ok {
		packet.Interfaces = append(packet.Interfaces, req)
	}
	if user, ok := df.getUser(); ok {
		packet.Interfaces = append(packet.Interfaces, user)
	}

	// set stacktrace data
	stConfig := &hook.StacktraceConfiguration
	if stConfig.Enable && entry.Level <= stConfig.Level {
		if err, ok := df.getError(); ok {
			var currentStacktrace *raven.Stacktrace
			currentStacktrace = hook.findStacktrace(err)
			if currentStacktrace == nil {
				currentStacktrace = raven.NewStacktrace(stConfig.Skip, stConfig.Context, stConfig.InAppPrefixes)
			}
			err := errors.Cause(err)
			exc := raven.NewException(err, currentStacktrace)
			packet.Interfaces = append(packet.Interfaces, exc)
			packet.Culprit = err.Error()
		} else {
			currentStacktrace := raven.NewStacktrace(stConfig.Skip, stConfig.Context, stConfig.InAppPrefixes)
			if currentStacktrace != nil {
				packet.Interfaces = append(packet.Interfaces, currentStacktrace)
			}
		}
	} else {
		// set the culprit even when the stack trace is disabled, as long as we have an error
		if err, ok := df.getError(); ok {
			packet.Culprit = err.Error()
		}
	}

	// set other fields
	dataExtra := hook.formatExtraData(df)
	if packet.Extra == nil {
		packet.Extra = dataExtra
	} else {
		for k, v := range dataExtra {
			packet.Extra[k] = v
		}
	}

	if hook.asynchronous {
		hook.wg.Add(1)
		hook.buf <- packet
		return nil
	}
	return hook.sendPacket(packet)
}

func (hook *SentryHook) fire() {
	for {
		packet := <-hook.buf
		if err := hook.sendPacket(packet); err != nil {
			fmt.Println(err)
		}
		hook.wg.Done()
	}
}

// Flush waits for the log queue to empty. This function only does anything in
// asynchronous mode.
func (hook *SentryHook) Flush() {
	if !hook.asynchronous {
		return
	}
	hook.mu.Lock() // Claim exclusive access; any logging goroutines will block until the flush completes
	defer hook.mu.Unlock()

	hook.wg.Wait()
}

func (hook *SentryHook) sendPacket(packet *raven.Packet) error {
	_, errCh := hook.client.Capture(packet, nil)
	timeout := hook.Timeout
	if timeout != 0 {
		timeoutCh := time.After(timeout)
		select {
		case err := <-errCh:
			return err
		case <-timeoutCh:
			return fmt.Errorf("no response from sentry server in %s", timeout)
		}
	}
	return nil
}

func (hook *SentryHook) findStacktrace(err error) *raven.Stacktrace {
	var stacktrace *raven.Stacktrace
	var stackErr errors.StackTrace
	for err != nil {
		// Find the earliest *raven.Stacktrace, or error.StackTrace
		if tracer, ok := err.(Stacktracer); ok {
			stacktrace = tracer.GetStacktrace()
			stackErr = nil
		} else if tracer, ok := err.(pkgErrorStackTracer); ok {
			stacktrace = nil
			stackErr = tracer.StackTrace()
		}
		if cause, ok := err.(causer); ok {
			err = cause.Cause()
		} else {
			break
		}
	}
	if stackErr != nil {
		stacktrace = hook.convertStackTrace(stackErr)
	}
	return stacktrace
}

// convertStackTrace converts an errors.StackTrace into a natively consumable
// *raven.Stacktrace
func (hook *SentryHook) convertStackTrace(st errors.StackTrace) *raven.Stacktrace {
	stConfig := &hook.StacktraceConfiguration
	stFrames := []errors.Frame(st)
	frames := make([]*raven.StacktraceFrame, 0, len(stFrames))
	for i := range stFrames {
		pc := uintptr(stFrames[i])
		fn := runtime.FuncForPC(pc)
		file, line := fn.FileLine(pc)
		frame := raven.NewStacktraceFrame(pc, file, line, stConfig.Context, stConfig.InAppPrefixes)
		if frame != nil {
			frames = append(frames, frame)
		}
	}

	// Sentry wants the frames with the oldest first, so reverse them
	for i, j := 0, len(frames)-1; i < j; i, j = i+1, j-1 {
		frames[i], frames[j] = frames[j], frames[i]
	}
	return &raven.Stacktrace{Frames: frames}
}

// Levels returns the available logging levels.
func (hook *SentryHook) Levels() []logrus.Level {
	return hook.levels
}

// SetRelease sets release tag.
func (hook *SentryHook) SetRelease(release string) {
	hook.client.SetRelease(release)
}

// SetEnvironment sets environment tag.
func (hook *SentryHook) SetEnvironment(environment string) {
	hook.client.SetEnvironment(environment)
}

// AddIgnore adds field name to ignore.
func (hook *SentryHook) AddIgnore(name string) {
	hook.ignoreFields[name] = struct{}{}
}

// AddExtraFilter adds a custom filter function.
func (hook *SentryHook) AddExtraFilter(name string, fn func(interface{}) interface{}) {
	hook.extraFilters[name] = fn
}

func (hook *SentryHook) formatExtraData(df *dataField) (result map[string]interface{}) {
	// create a map for passing to Sentry's extra data
	result = make(map[string]interface{}, df.len())
	for k, v := range df.data {
		if df.isOmit(k) {
			continue // skip already used special fields
		}
		if _, ok := hook.ignoreFields[k]; ok {
			continue
		}

		if fn, ok := hook.extraFilters[k]; ok {
			v = fn(v) // apply custom filter
		} else {
			v = formatData(v) // use default formatter
		}
		result[k] = v
	}
	return result
}

// formatData returns value as a suitable format.
func formatData(value interface{}) (formatted interface{}) {
	switch value := value.(type) {
	case json.Marshaler:
		return value
	case error:
		return value.Error()
	case fmt.Stringer:
		return value.String()
	default:
		return value
	}
}
<commit_msg>Stop double-buffering in async hook (#47)<commit_after>package logrus_sentry

import (
	"encoding/json"
	"fmt"
	"runtime"
	"sync"
	"time"

	"github.com/getsentry/raven-go"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

var (
	severityMap = map[logrus.Level]raven.Severity{
		logrus.DebugLevel: raven.DEBUG,
		logrus.InfoLevel:  raven.INFO,
		logrus.WarnLevel:  raven.WARNING,
		logrus.ErrorLevel: raven.ERROR,
		logrus.FatalLevel: raven.FATAL,
		logrus.PanicLevel: raven.FATAL,
	}
)

// SentryHook delivers logs to a sentry server.
type SentryHook struct {
	// Timeout sets the time to wait for a delivery error from the sentry server.
	// If this is set to zero the server will not wait for any response and will
	// consider the message correctly sent.
	//
	// This is ignored for asynchronous hooks. If you want to set a timeout when
	// using an async hook (to bound the length of time that hook.Flush can take),
	// you probably want to create your own raven.Client and set
	// ravenClient.Transport.(*raven.HTTPTransport).Client.Timeout to set a
	// timeout on the underlying HTTP request instead.
	Timeout                 time.Duration
	StacktraceConfiguration StackTraceConfiguration

	client *raven.Client
	levels []logrus.Level

	ignoreFields map[string]struct{}
	extraFilters map[string]func(interface{}) interface{}

	asynchronous bool

	mu sync.RWMutex
	wg sync.WaitGroup
}

// The Stacktracer interface allows an error type to return a raven.Stacktrace.
type Stacktracer interface {
	GetStacktrace() *raven.Stacktrace
}

type causer interface {
	Cause() error
}

type pkgErrorStackTracer interface {
	StackTrace() errors.StackTrace
}

// StackTraceConfiguration allows for configuring stacktraces
type StackTraceConfiguration struct {
	// whether stacktraces should be enabled
	Enable bool
	// the level at which to start capturing stacktraces
	Level logrus.Level
	// how many stack frames to skip before stacktrace starts recording
	Skip int
	// the number of lines to include around a stack frame for context
	Context int
	// the prefixes that will be matched against the stack frame.
	// if the stack frame's package matches one of these prefixes
	// sentry will identify the stack frame as "in_app"
	InAppPrefixes []string
}

// NewSentryHook creates a hook to be added to an instance of logger
// and initializes the raven client.
// This method sets the timeout to 100 milliseconds.
func NewSentryHook(DSN string, levels []logrus.Level) (*SentryHook, error) {
	client, err := raven.New(DSN)
	if err != nil {
		return nil, err
	}
	return NewWithClientSentryHook(client, levels)
}

// NewWithTagsSentryHook creates a hook with tags to be added to an instance
// of logger and initializes the raven client. This method sets the timeout to
// 100 milliseconds.
func NewWithTagsSentryHook(DSN string, tags map[string]string, levels []logrus.Level) (*SentryHook, error) {
	client, err := raven.NewWithTags(DSN, tags)
	if err != nil {
		return nil, err
	}
	return NewWithClientSentryHook(client, levels)
}

// NewWithClientSentryHook creates a hook using an initialized raven client.
// This method sets the timeout to 100 milliseconds.
func NewWithClientSentryHook(client *raven.Client, levels []logrus.Level) (*SentryHook, error) {
	return &SentryHook{
		Timeout: 100 * time.Millisecond,
		StacktraceConfiguration: StackTraceConfiguration{
			Enable:        false,
			Level:         logrus.ErrorLevel,
			Skip:          5,
			Context:       0,
			InAppPrefixes: nil,
		},
		client:       client,
		levels:       levels,
		ignoreFields: make(map[string]struct{}),
		extraFilters: make(map[string]func(interface{}) interface{}),
	}, nil
}

// NewAsyncSentryHook creates a hook same as NewSentryHook, but in asynchronous
// mode.
func NewAsyncSentryHook(DSN string, levels []logrus.Level) (*SentryHook, error) {
	hook, err := NewSentryHook(DSN, levels)
	return setAsync(hook), err
}

// NewAsyncWithTagsSentryHook creates a hook same as NewWithTagsSentryHook, but
// in asynchronous mode.
func NewAsyncWithTagsSentryHook(DSN string, tags map[string]string, levels []logrus.Level) (*SentryHook, error) {
	hook, err := NewWithTagsSentryHook(DSN, tags, levels)
	return setAsync(hook), err
}

// NewAsyncWithClientSentryHook creates a hook same as NewWithClientSentryHook,
// but in asynchronous mode.
func NewAsyncWithClientSentryHook(client *raven.Client, levels []logrus.Level) (*SentryHook, error) {
	hook, err := NewWithClientSentryHook(client, levels)
	return setAsync(hook), err
}

func setAsync(hook *SentryHook) *SentryHook {
	if hook == nil {
		return nil
	}
	hook.asynchronous = true
	return hook
}

// Fire is called when an event should be sent to sentry
// Special fields that sentry uses to give more information to the server
// are extracted from entry.Data (if they are found)
// These fields are: error, logger, server_name, http_request, tags
func (hook *SentryHook) Fire(entry *logrus.Entry) error {
	hook.mu.RLock() // Allow multiple go routines to log simultaneously
	defer hook.mu.RUnlock()
	packet := raven.NewPacket(entry.Message)
	packet.Timestamp = raven.Timestamp(entry.Time)
	packet.Level = severityMap[entry.Level]
	packet.Platform = "go"

	df := newDataField(entry.Data)

	// set special fields
	if logger, ok := df.getLogger(); ok {
		packet.Logger = logger
	}
	if serverName, ok := df.getServerName(); ok {
		packet.ServerName = serverName
	}
	if eventID, ok := df.getEventID(); ok {
		packet.EventID = eventID
	}
	if tags, ok := df.getTags(); ok {
		packet.Tags = tags
	}
	if fingerprint, ok := df.getFingerprint(); ok {
		packet.Fingerprint = fingerprint
	}
	if req, ok := df.getHTTPRequest(); ok {
		packet.Interfaces = append(packet.Interfaces, req)
	}
	if user, ok := df.getUser(); ok {
		packet.Interfaces = append(packet.Interfaces, user)
	}

	// set stacktrace data
	stConfig := &hook.StacktraceConfiguration
	if stConfig.Enable && entry.Level <= stConfig.Level {
		if err, ok := df.getError(); ok {
			var currentStacktrace *raven.Stacktrace
			currentStacktrace = hook.findStacktrace(err)
			if currentStacktrace == nil {
				currentStacktrace = raven.NewStacktrace(stConfig.Skip, stConfig.Context, stConfig.InAppPrefixes)
			}
			err := errors.Cause(err)
			exc := raven.NewException(err, currentStacktrace)
			packet.Interfaces = append(packet.Interfaces, exc)
			packet.Culprit = err.Error()
		} else {
			currentStacktrace := raven.NewStacktrace(stConfig.Skip, stConfig.Context, stConfig.InAppPrefixes)
			if currentStacktrace != nil {
				packet.Interfaces = append(packet.Interfaces, currentStacktrace)
			}
		}
	} else {
		// set the culprit even when the stack trace is disabled, as long as we have an error
		if err, ok := df.getError(); ok {
			packet.Culprit = err.Error()
		}
	}

	// set other fields
	dataExtra := hook.formatExtraData(df)
	if packet.Extra == nil {
		packet.Extra = dataExtra
	} else {
		for k, v := range dataExtra {
			packet.Extra[k] = v
		}
	}

	_, errCh := hook.client.Capture(packet, nil)

	if hook.asynchronous {
		// Our use of hook.mu guarantees that we are following the WaitGroup rule of
		// not calling Add in parallel with Wait.
		hook.wg.Add(1)
		go func() {
			if err := <-errCh; err != nil {
				fmt.Println(err)
			}
			hook.wg.Done()
		}()
		return nil
	} else if timeout := hook.Timeout; timeout == 0 {
		return nil
	} else {
		timeoutCh := time.After(timeout)
		select {
		case err := <-errCh:
			return err
		case <-timeoutCh:
			return fmt.Errorf("no response from sentry server in %s", timeout)
		}
	}
}

// Flush waits for the log queue to empty. This function only does anything in
// asynchronous mode.
func (hook *SentryHook) Flush() {
	if !hook.asynchronous {
		return
	}
	hook.mu.Lock() // Claim exclusive access; any logging goroutines will block until the flush completes
	defer hook.mu.Unlock()

	hook.wg.Wait()
}

func (hook *SentryHook) findStacktrace(err error) *raven.Stacktrace {
	var stacktrace *raven.Stacktrace
	var stackErr errors.StackTrace
	for err != nil {
		// Find the earliest *raven.Stacktrace, or error.StackTrace
		if tracer, ok := err.(Stacktracer); ok {
			stacktrace = tracer.GetStacktrace()
			stackErr = nil
		} else if tracer, ok := err.(pkgErrorStackTracer); ok {
			stacktrace = nil
			stackErr = tracer.StackTrace()
		}
		if cause, ok := err.(causer); ok {
			err = cause.Cause()
		} else {
			break
		}
	}
	if stackErr != nil {
		stacktrace = hook.convertStackTrace(stackErr)
	}
	return stacktrace
}

// convertStackTrace converts an errors.StackTrace into a natively consumable
// *raven.Stacktrace
func (hook *SentryHook) convertStackTrace(st errors.StackTrace) *raven.Stacktrace {
	stConfig := &hook.StacktraceConfiguration
	stFrames := []errors.Frame(st)
	frames := make([]*raven.StacktraceFrame, 0, len(stFrames))
	for i := range stFrames {
		pc := uintptr(stFrames[i])
		fn := runtime.FuncForPC(pc)
		file, line := fn.FileLine(pc)
		frame := raven.NewStacktraceFrame(pc, file, line, stConfig.Context, stConfig.InAppPrefixes)
		if frame != nil {
			frames = append(frames, frame)
		}
	}

	// Sentry wants the frames with the oldest first, so reverse them
	for i, j := 0, len(frames)-1; i < j; i, j = i+1, j-1 {
		frames[i], frames[j] = frames[j], frames[i]
	}
	return &raven.Stacktrace{Frames: frames}
}

// Levels returns the available logging levels.
func (hook *SentryHook) Levels() []logrus.Level {
	return hook.levels
}

// SetRelease sets release tag.
func (hook *SentryHook) SetRelease(release string) {
	hook.client.SetRelease(release)
}

// SetEnvironment sets environment tag.
func (hook *SentryHook) SetEnvironment(environment string) {
	hook.client.SetEnvironment(environment)
}

// AddIgnore adds field name to ignore.
func (hook *SentryHook) AddIgnore(name string) {
	hook.ignoreFields[name] = struct{}{}
}

// AddExtraFilter adds a custom filter function.
func (hook *SentryHook) AddExtraFilter(name string, fn func(interface{}) interface{}) {
	hook.extraFilters[name] = fn
}

func (hook *SentryHook) formatExtraData(df *dataField) (result map[string]interface{}) {
	// create a map for passing to Sentry's extra data
	result = make(map[string]interface{}, df.len())
	for k, v := range df.data {
		if df.isOmit(k) {
			continue // skip already used special fields
		}
		if _, ok := hook.ignoreFields[k]; ok {
			continue
		}

		if fn, ok := hook.extraFilters[k]; ok {
			v = fn(v) // apply custom filter
		} else {
			v = formatData(v) // use default formatter
		}
		result[k] = v
	}
	return result
}

// formatData returns value as a suitable format.
func formatData(value interface{}) (formatted interface{}) {
	switch value := value.(type) {
	case json.Marshaler:
		return value
	case error:
		return value.Error()
	case fmt.Stringer:
		return value.String()
	default:
		return value
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 The StudyGolang Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// http://studygolang.com
// Author: polaris	polaris@studygolang.com

package middleware

import (
	"db"
	"logic"
	"model"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
	"util"

	. "http"

	"github.com/gorilla/context"
	"github.com/labstack/echo"
)

// AutoLogin 用于 echo 框架的自动登录和通过 cookie 获取用户信息
func AutoLogin() echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(ctx echo.Context) error {
			// github.com/gorilla/sessions 要求必须 Clear
			defer context.Clear(Request(ctx))

			ctx.Set("req_start_time", time.Now())

			var getCurrentUser = func(usernameOrId interface{}) {
				if db.MasterDB != nil {
					// TODO: 考虑缓存，或延迟查询，避免每次都查询
					user := logic.DefaultUser.FindCurrentUser(ctx, usernameOrId)
					if user.Uid != 0 {
						ctx.Set("user", user)

						if !util.IsAjax(ctx) && ctx.Path() != "/ws" {
							go logic.ViewObservable.NotifyObservers(user.Uid, 0, 0)
						}
					}
				}
			}

			session := GetCookieSession(ctx)
			username, ok := session.Values["username"]
			if ok {
				getCurrentUser(username)
			} else {
				// App（手机） 登录
				uid, ok := ParseToken(ctx.FormValue("token"))
				if ok {
					getCurrentUser(uid)
				}
			}

			if err := next(ctx); err != nil {
				return err
			}

			return nil
		}
	}
}

// NeedLogin 用于 echo 框架的验证必须登录的请求
func NeedLogin() echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(ctx echo.Context) error {
			user, ok := ctx.Get("user").(*model.Me)
			if !ok || user.Status != model.UserStatusAudit {
				method := ctx.Request().Method()
				if util.IsAjax(ctx) {
					if !strings.HasPrefix(ctx.Path(), "/account") {
						return ctx.JSON(http.StatusForbidden, `{"ok":0,"error":"403 Forbidden"}`)
					}
				} else {
					if method == "POST" {
						return ctx.HTML(http.StatusForbidden, `403 Forbidden`)
					}

					if !ok {
						reqURL := ctx.Request().URL()
						uri := reqURL.Path()
						if reqURL.QueryString() != "" {
							uri += "?" + reqURL.QueryString()
						}
						return ctx.Redirect(http.StatusSeeOther, "/account/login?redirect_uri="+url.QueryEscape(uri))
					} else {
						// 未激活可以查看账号信息
						if !strings.HasPrefix(ctx.Path(), "/account") {
							return echo.NewHTTPError(http.StatusForbidden, `您的邮箱未激活，<a href="/account/edit">去激活</a>`)
						}
					}
				}
			}

			if err := next(ctx); err != nil {
				return err
			}

			return nil
		}
	}
}

// AppNeedLogin 用于 echo 框架的验证必须登录的请求（APP 专用）
func AppNeedLogin() echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(ctx echo.Context) error {
			user, ok := ctx.Get("user").(*model.Me)
			if ok {
				// 校验 token 是否有效
				if !ValidateToken(ctx.QueryParam("token")) {
					return outputAppJSON(ctx, NeedReLoginCode, "token无效，请重新登录！")
				}

				if user.Status != model.UserStatusAudit {
					return outputAppJSON(ctx, 1, "账号未审核通过、被冻结或被停号，请联系我们")
				}
			} else {
				return outputAppJSON(ctx, NeedReLoginCode, "请先登录！")
			}

			if err := next(ctx); err != nil {
				return err
			}

			return nil
		}
	}
}

func outputAppJSON(ctx echo.Context, code int, msg string) error {
	AccessControl(ctx)
	respJSON := `{"code":` + strconv.Itoa(code) + `,"msg":"` + msg + `}`
	return ctx.JSON(http.StatusForbidden, respJSON)
}
<commit_msg>bugfix:middleware.NeedLogin()<commit_after>// Copyright 2013 The StudyGolang Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// http://studygolang.com
// Author: polaris	polaris@studygolang.com

package middleware

import (
	"db"
	"logic"
	"model"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
	"util"

	. "http"

	"github.com/gorilla/context"
	"github.com/labstack/echo"
)

// AutoLogin 用于 echo 框架的自动登录和通过 cookie 获取用户信息
func AutoLogin() echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(ctx echo.Context) error {
			// github.com/gorilla/sessions 要求必须 Clear
			defer context.Clear(Request(ctx))

			ctx.Set("req_start_time", time.Now())

			var getCurrentUser = func(usernameOrId interface{}) {
				if db.MasterDB != nil {
					// TODO: 考虑缓存，或延迟查询，避免每次都查询
					user := logic.DefaultUser.FindCurrentUser(ctx, usernameOrId)
					if user.Uid != 0 {
						ctx.Set("user", user)

						if !util.IsAjax(ctx) && ctx.Path() != "/ws" {
							go logic.ViewObservable.NotifyObservers(user.Uid, 0, 0)
						}
					}
				}
			}

			session := GetCookieSession(ctx)
			username, ok := session.Values["username"]
			if ok {
				getCurrentUser(username)
			} else {
				// App（手机） 登录
				uid, ok := ParseToken(ctx.FormValue("token"))
				if ok {
					getCurrentUser(uid)
				}
			}

			if err := next(ctx); err != nil {
				return err
			}

			return nil
		}
	}
}

// NeedLogin 用于 echo 框架的验证必须登录的请求
func NeedLogin() echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(ctx echo.Context) error {
			user, ok := ctx.Get("user").(*model.Me)
			if !ok || user.Status != model.UserStatusAudit {
				method := ctx.Request().Method()
				if util.IsAjax(ctx) {
					if !strings.HasPrefix(ctx.Path(), "/account") {
						return ctx.JSON(http.StatusForbidden, map[string]interface{}{"ok": 0, "error": "403 Forbidden"})
					}
				} else {
					if method == "POST" {
						return ctx.HTML(http.StatusForbidden, `403 Forbidden`)
					}

					if !ok {
						reqURL := ctx.Request().URL()
						uri := reqURL.Path()
						if reqURL.QueryString() != "" {
							uri += "?" + reqURL.QueryString()
						}
						return ctx.Redirect(http.StatusSeeOther, "/account/login?redirect_uri="+url.QueryEscape(uri))
					} else {
						// 未激活可以查看账号信息
						if !strings.HasPrefix(ctx.Path(), "/account") {
							return echo.NewHTTPError(http.StatusForbidden, `您的邮箱未激活，<a href="/account/edit">去激活</a>`)
						}
					}
				}
			}

			if err := next(ctx); err != nil {
				return err
			}

			return nil
		}
	}
}

// AppNeedLogin 用于 echo 框架的验证必须登录的请求（APP 专用）
func AppNeedLogin() echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(ctx echo.Context) error {
			user, ok := ctx.Get("user").(*model.Me)
			if ok {
				// 校验 token 是否有效
				if !ValidateToken(ctx.QueryParam("token")) {
					return outputAppJSON(ctx, NeedReLoginCode, "token无效，请重新登录！")
				}

				if user.Status != model.UserStatusAudit {
					return outputAppJSON(ctx, 1, "账号未审核通过、被冻结或被停号，请联系我们")
				}
			} else {
				return outputAppJSON(ctx, NeedReLoginCode, "请先登录！")
			}

			if err := next(ctx); err != nil {
				return err
			}

			return nil
		}
	}
}

func outputAppJSON(ctx echo.Context, code int, msg string) error {
	AccessControl(ctx)
	return ctx.JSON(http.StatusForbidden, map[string]interface{}{"code": strconv.Itoa(code), "msg": msg})
}
<|endoftext|>