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

import (
	"fmt"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/henrylee2cn/goutil"
	"github.com/henrylee2cn/goutil/errors"
)

type (
	// Socket is a generic stream-oriented network connection.
	//
	// Multiple goroutines may invoke methods on a Socket simultaneously.
	Socket interface {
		// LocalAddr returns the local network address.
		LocalAddr() net.Addr

		// RemoteAddr returns the remote network address.
		RemoteAddr() net.Addr

		// SetDeadline sets the read and write deadlines associated
		// with the connection. It is equivalent to calling both
		// SetReadDeadline and SetWriteDeadline.
		//
		// A deadline is an absolute time after which I/O operations
		// fail with a timeout (see type Error) instead of
		// blocking. The deadline applies to all future and pending
		// I/O, not just the immediately following call to Read or
		// Write. After a deadline has been exceeded, the connection
		// can be refreshed by setting a deadline in the future.
		//
		// An idle timeout can be implemented by repeatedly extending
		// the deadline after successful Read or Write calls.
		//
		// A zero value for t means I/O operations will not time out.
		SetDeadline(t time.Time) error

		// SetReadDeadline sets the deadline for future Read calls
		// and any currently-blocked Read call.
		// A zero value for t means Read will not time out.
		SetReadDeadline(t time.Time) error

		// SetWriteDeadline sets the deadline for future Write calls
		// and any currently-blocked Write call.
		// Even if write times out, it may return n > 0, indicating that
		// some of the data was successfully written.
		// A zero value for t means Write will not time out.
		SetWriteDeadline(t time.Time) error

		// WritePacket writes header and body to the connection.
		// Note: must be safe for concurrent use by multiple goroutines.
		WritePacket(packet *Packet) error

		// ReadPacket reads header and body from the connection.
		// Note: must be safe for concurrent use by multiple goroutines.
		ReadPacket(packet *Packet) error

		// Read reads data from the connection.
		// Read can be made to time out and return an Error with Timeout() == true
		// after a fixed time limit; see SetDeadline and SetReadDeadline.
		Read(b []byte) (n int, err error)

		// Write writes data to the connection.
		// Write can be made to time out and return an Error with Timeout() == true
		// after a fixed time limit; see SetDeadline and SetWriteDeadline.
		Write(b []byte) (n int, err error)

		// Reset reset net.Conn
		// Reset(net.Conn)

		// Close closes the connection socket.
		// Any blocked Read or Write operations will be unblocked and return errors.
		Close() error

		// Public returns temporary public data of Socket.
		Public() goutil.Map
		// PublicLen returns the length of public data of Socket.
		PublicLen() int
		// Id returns the socket id.
		Id() string
		// SetId sets the socket id.
		SetId(string)
		// Reset reset net.Conn and ProtoFunc.
		Reset(netConn net.Conn, protoFunc ...ProtoFunc)
	}
	socket struct {
		net.Conn
		protocol  Proto
		id        string
		idMutex   sync.RWMutex
		ctxPublic goutil.Map
		mu        sync.RWMutex
		curState  int32
		fromPool  bool
	}
)

const (
	normal      int32 = 0
	activeClose int32 = 1
)

var _ net.Conn = Socket(nil)

// ErrProactivelyCloseSocket proactively close the socket error.
var ErrProactivelyCloseSocket = errors.New("socket is closed proactively")

// GetSocket gets a Socket from pool, and reset it.
func GetSocket(c net.Conn, protoFunc ...ProtoFunc) Socket {
	s := socketPool.Get().(*socket)
	s.Reset(c, protoFunc...)
	return s
}

var socketPool = sync.Pool{
	New: func() interface{} {
		s := newSocket(nil, nil)
		s.fromPool = true
		return s
	},
}

// NewSocket wraps a net.Conn as a Socket.
func NewSocket(c net.Conn, protoFunc ...ProtoFunc) Socket {
	return newSocket(c, protoFunc)
}

func newSocket(c net.Conn, protoFuncs []ProtoFunc) *socket {
	var s = &socket{
		protocol: getProto(protoFuncs, c),
		Conn:     c,
	}
	s.optimize()
	return s
}

// WritePacket writes header and body to the connection.
// WritePacket can be made to time out and return an Error with Timeout() == true
// after a fixed time limit; see SetDeadline and SetWriteDeadline.
// Note:
//  For the byte stream type of body, write directly, do not do any processing;
//  Must be safe for concurrent use by multiple goroutines.
func (s *socket) WritePacket(packet *Packet) error {
	s.mu.RLock()
	protocol := s.protocol
	s.mu.RUnlock()
	err := protocol.Pack(packet)
	if err != nil && s.isActiveClosed() {
		err = ErrProactivelyCloseSocket
	}
	return err
}

// ReadPacket reads header and body from the connection.
// Note:
//  For the byte stream type of body, read directly, do not do any processing;
//  Must be safe for concurrent use by multiple goroutines.
func (s *socket) ReadPacket(packet *Packet) error {
	s.mu.RLock()
	protocol := s.protocol
	s.mu.RUnlock()
	return protocol.Unpack(packet)
}

// Public returns temporary public data of Socket.
func (s *socket) Public() goutil.Map {
	if s.ctxPublic == nil {
		s.ctxPublic = goutil.RwMap()
	}
	return s.ctxPublic
}

// PublicLen returns the length of public data of Socket.
func (s *socket) PublicLen() int {
	if s.ctxPublic == nil {
		return 0
	}
	return s.ctxPublic.Len()
}

// Id returns the socket id.
func (s *socket) Id() string {
	s.idMutex.RLock()
	id := s.id
	if len(id) == 0 {
		id = s.RemoteAddr().String()
	}
	s.idMutex.RUnlock()
	return id
}

// SetId sets the socket id.
func (s *socket) SetId(id string) {
	s.idMutex.Lock()
	s.id = id
	s.idMutex.Unlock()
}

// Reset reset net.Conn and ProtoFunc.
func (s *socket) Reset(netConn net.Conn, protoFunc ...ProtoFunc) {
	atomic.StoreInt32(&s.curState, activeClose)
	if s.Conn != nil {
		s.Conn.Close()
	}
	s.mu.Lock()
	s.Conn = netConn
	s.SetId("")
	s.protocol = getProto(protoFunc, netConn)
	atomic.StoreInt32(&s.curState, normal)
	s.optimize()
	s.mu.Unlock()
}

// Close closes the connection socket.
// Any blocked Read or Write operations will be unblocked and return errors.
// If it is from 'GetSocket()' function(a pool), return itself to pool.
func (s *socket) Close() error {
	if s.isActiveClosed() {
		return nil
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.isActiveClosed() {
		return nil
	}
	atomic.StoreInt32(&s.curState, activeClose)

	var err error
	if s.Conn != nil {
		err = s.Conn.Close()
	}
	if s.fromPool {
		s.Conn = nil
		s.ctxPublic = nil
		s.protocol = nil
		socketPool.Put(s)
	}
	return err
}

func (s *socket) isActiveClosed() bool {
	return atomic.LoadInt32(&s.curState) == activeClose
}

func (s *socket) optimize() {
	if c, ok := s.Conn.(ifaceSetKeepAlive); ok {
		if changeKeepAlive {
			c.SetKeepAlive(keepAlive)
		}
		if keepAlivePeriod >= 0 && keepAlive {
			c.SetKeepAlivePeriod(keepAlivePeriod)
		}
	}
	if c, ok := s.Conn.(ifaceSetBuffer); ok {
		if readBuffer >= 0 {
			c.SetReadBuffer(readBuffer)
		}
		if writeBuffer >= 0 {
			c.SetWriteBuffer(writeBuffer)
		}
	}
	if c, ok := s.Conn.(ifaceSetNoDelay); ok {
		if !noDelay {
			c.SetNoDelay(noDelay)
		}
	}
}

type (
	ifaceSetKeepAlive interface {
		// SetKeepAlive sets whether the operating system should send
		// keepalive messages on the connection.
		SetKeepAlive(keepalive bool) error
		// SetKeepAlivePeriod sets period between keep alives.
		SetKeepAlivePeriod(d time.Duration) error
	}
	ifaceSetBuffer interface {
		// SetReadBuffer sets the size of the operating system's
		// receive buffer associated with the connection.
		SetReadBuffer(bytes int) error
		// SetWriteBuffer sets the size of the operating system's
		// transmit buffer associated with the connection.
		SetWriteBuffer(bytes int) error
	}
	ifaceSetNoDelay interface {
		// SetNoDelay controls whether the operating system should delay
		// packet transmission in hopes of sending fewer packets (Nagle's
		// algorithm).  The default is true (no delay), meaning that data is
		// sent as soon as possible after a Write.
		SetNoDelay(noDelay bool) error
	}
)

// Connection related system configuration
var (
	writeBuffer     int           = -1
	readBuffer      int           = -1
	changeKeepAlive bool          = false
	keepAlive       bool          = true
	keepAlivePeriod time.Duration = -1
	noDelay         bool          = true
)

// SetKeepAlive sets whether the operating system should send
// keepalive messages on the connection.
// Note: If have not called the function, the system defaults are used.
func SetKeepAlive(keepalive bool) {
	changeKeepAlive = true
	keepAlive = keepalive
}

// SetKeepAlivePeriod sets period between keep alives.
// Note: if d<0, don't change the value.
func SetKeepAlivePeriod(d time.Duration) {
	if d >= 0 {
		keepAlivePeriod = d
	} else {
		fmt.Println("socket: SetKeepAlivePeriod: invalid keepAlivePeriod:", d)
	}
}

// ReadBuffer returns the size of the operating system's
// receive buffer associated with the connection.
// Note: if using the system default value, bytes=-1 and isDefault=true.
func ReadBuffer() (bytes int, isDefault bool) {
	return readBuffer, readBuffer == -1
}

// SetReadBuffer sets the size of the operating system's
// receive buffer associated with the connection.
// Note: if bytes<0, don't change the value.
func SetReadBuffer(bytes int) {
	if bytes >= 0 {
		readBuffer = bytes
	} else {
		fmt.Println("socket: SetReadBuffer: invalid readBuffer size:", bytes)
	}
}

// WriteBuffer returns the size of the operating system's
// transmit buffer associated with the connection.
// Note: if using the system default value, bytes=-1 and isDefault=true.
func WriteBuffer() (bytes int, isDefault bool) {
	return writeBuffer, writeBuffer == -1
}

// SetWriteBuffer sets the size of the operating system's
// transmit buffer associated with the connection.
// Note: if bytes<0, don't change the value.
func SetWriteBuffer(bytes int) {
	if bytes >= 0 {
		writeBuffer = bytes
	} else {
		fmt.Println("socket: SetWriteBuffer: invalid writeBuffer size:", bytes)
	}
}

// SetNoDelay controls whether the operating system should delay
// packet transmission in hopes of sending fewer packets (Nagle's
// algorithm).  The default is true (no delay), meaning that data is
// sent as soon as possible after a Write.
func SetNoDelay(_noDelay bool) {
	noDelay = _noDelay
}

func getProto(protoFuncs []ProtoFunc, rw io.ReadWriter) Proto {
	if len(protoFuncs) > 0 {
		return protoFuncs[0](rw)
	}
	return defaultProtoFunc(rw)
}
<commit_msg>Optimize getProto<commit_after>// Package socket provides a concise, powerful and high-performance TCP.
//
// Copyright 2017 HenryLee. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package socket

import (
	"fmt"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/henrylee2cn/goutil"
	"github.com/henrylee2cn/goutil/errors"
)

type (
	// Socket is a generic stream-oriented network connection.
	//
	// Multiple goroutines may invoke methods on a Socket simultaneously.
	Socket interface {
		// LocalAddr returns the local network address.
		LocalAddr() net.Addr

		// RemoteAddr returns the remote network address.
		RemoteAddr() net.Addr

		// SetDeadline sets the read and write deadlines associated
		// with the connection. It is equivalent to calling both
		// SetReadDeadline and SetWriteDeadline.
		//
		// A deadline is an absolute time after which I/O operations
		// fail with a timeout (see type Error) instead of
		// blocking. The deadline applies to all future and pending
		// I/O, not just the immediately following call to Read or
		// Write. After a deadline has been exceeded, the connection
		// can be refreshed by setting a deadline in the future.
		//
		// An idle timeout can be implemented by repeatedly extending
		// the deadline after successful Read or Write calls.
		//
		// A zero value for t means I/O operations will not time out.
		SetDeadline(t time.Time) error

		// SetReadDeadline sets the deadline for future Read calls
		// and any currently-blocked Read call.
		// A zero value for t means Read will not time out.
		SetReadDeadline(t time.Time) error

		// SetWriteDeadline sets the deadline for future Write calls
		// and any currently-blocked Write call.
		// Even if write times out, it may return n > 0, indicating that
		// some of the data was successfully written.
		// A zero value for t means Write will not time out.
		SetWriteDeadline(t time.Time) error

		// WritePacket writes header and body to the connection.
		// Note: must be safe for concurrent use by multiple goroutines.
		WritePacket(packet *Packet) error

		// ReadPacket reads header and body from the connection.
		// Note: must be safe for concurrent use by multiple goroutines.
		ReadPacket(packet *Packet) error

		// Read reads data from the connection.
		// Read can be made to time out and return an Error with Timeout() == true
		// after a fixed time limit; see SetDeadline and SetReadDeadline.
		Read(b []byte) (n int, err error)

		// Write writes data to the connection.
		// Write can be made to time out and return an Error with Timeout() == true
		// after a fixed time limit; see SetDeadline and SetWriteDeadline.
		Write(b []byte) (n int, err error)

		// Reset reset net.Conn
		// Reset(net.Conn)

		// Close closes the connection socket.
		// Any blocked Read or Write operations will be unblocked and return errors.
		Close() error

		// Public returns temporary public data of Socket.
		Public() goutil.Map
		// PublicLen returns the length of public data of Socket.
		PublicLen() int
		// Id returns the socket id.
		Id() string
		// SetId sets the socket id.
		SetId(string)
		// Reset reset net.Conn and ProtoFunc.
		Reset(netConn net.Conn, protoFunc ...ProtoFunc)
	}
	socket struct {
		net.Conn
		protocol  Proto
		id        string
		idMutex   sync.RWMutex
		ctxPublic goutil.Map
		mu        sync.RWMutex
		curState  int32
		fromPool  bool
	}
)

const (
	normal      int32 = 0
	activeClose int32 = 1
)

var _ net.Conn = Socket(nil)

// ErrProactivelyCloseSocket proactively close the socket error.
var ErrProactivelyCloseSocket = errors.New("socket is closed proactively")

// GetSocket gets a Socket from pool, and reset it.
func GetSocket(c net.Conn, protoFunc ...ProtoFunc) Socket {
	s := socketPool.Get().(*socket)
	s.Reset(c, protoFunc...)
	return s
}

var socketPool = sync.Pool{
	New: func() interface{} {
		s := newSocket(nil, nil)
		s.fromPool = true
		return s
	},
}

// NewSocket wraps a net.Conn as a Socket.
func NewSocket(c net.Conn, protoFunc ...ProtoFunc) Socket {
	return newSocket(c, protoFunc)
}

func newSocket(c net.Conn, protoFuncs []ProtoFunc) *socket {
	var s = &socket{
		protocol: getProto(protoFuncs, c),
		Conn:     c,
	}
	s.optimize()
	return s
}

// WritePacket writes header and body to the connection.
// WritePacket can be made to time out and return an Error with Timeout() == true
// after a fixed time limit; see SetDeadline and SetWriteDeadline.
// Note:
//  For the byte stream type of body, write directly, do not do any processing;
//  Must be safe for concurrent use by multiple goroutines.
func (s *socket) WritePacket(packet *Packet) error {
	s.mu.RLock()
	protocol := s.protocol
	s.mu.RUnlock()
	err := protocol.Pack(packet)
	if err != nil && s.isActiveClosed() {
		err = ErrProactivelyCloseSocket
	}
	return err
}

// ReadPacket reads header and body from the connection.
// Note:
//  For the byte stream type of body, read directly, do not do any processing;
//  Must be safe for concurrent use by multiple goroutines.
func (s *socket) ReadPacket(packet *Packet) error {
	s.mu.RLock()
	protocol := s.protocol
	s.mu.RUnlock()
	return protocol.Unpack(packet)
}

// Public returns temporary public data of Socket.
func (s *socket) Public() goutil.Map {
	if s.ctxPublic == nil {
		s.ctxPublic = goutil.RwMap()
	}
	return s.ctxPublic
}

// PublicLen returns the length of public data of Socket.
func (s *socket) PublicLen() int {
	if s.ctxPublic == nil {
		return 0
	}
	return s.ctxPublic.Len()
}

// Id returns the socket id.
func (s *socket) Id() string {
	s.idMutex.RLock()
	id := s.id
	if len(id) == 0 {
		id = s.RemoteAddr().String()
	}
	s.idMutex.RUnlock()
	return id
}

// SetId sets the socket id.
func (s *socket) SetId(id string) {
	s.idMutex.Lock()
	s.id = id
	s.idMutex.Unlock()
}

// Reset reset net.Conn and ProtoFunc.
func (s *socket) Reset(netConn net.Conn, protoFunc ...ProtoFunc) {
	atomic.StoreInt32(&s.curState, activeClose)
	if s.Conn != nil {
		s.Conn.Close()
	}
	s.mu.Lock()
	s.Conn = netConn
	s.SetId("")
	s.protocol = getProto(protoFunc, netConn)
	atomic.StoreInt32(&s.curState, normal)
	s.optimize()
	s.mu.Unlock()
}

// Close closes the connection socket.
// Any blocked Read or Write operations will be unblocked and return errors.
// If it is from 'GetSocket()' function(a pool), return itself to pool.
func (s *socket) Close() error {
	if s.isActiveClosed() {
		return nil
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.isActiveClosed() {
		return nil
	}
	atomic.StoreInt32(&s.curState, activeClose)

	var err error
	if s.Conn != nil {
		err = s.Conn.Close()
	}
	if s.fromPool {
		s.Conn = nil
		s.ctxPublic = nil
		s.protocol = nil
		socketPool.Put(s)
	}
	return err
}

func (s *socket) isActiveClosed() bool {
	return atomic.LoadInt32(&s.curState) == activeClose
}

func (s *socket) optimize() {
	if c, ok := s.Conn.(ifaceSetKeepAlive); ok {
		if changeKeepAlive {
			c.SetKeepAlive(keepAlive)
		}
		if keepAlivePeriod >= 0 && keepAlive {
			c.SetKeepAlivePeriod(keepAlivePeriod)
		}
	}
	if c, ok := s.Conn.(ifaceSetBuffer); ok {
		if readBuffer >= 0 {
			c.SetReadBuffer(readBuffer)
		}
		if writeBuffer >= 0 {
			c.SetWriteBuffer(writeBuffer)
		}
	}
	if c, ok := s.Conn.(ifaceSetNoDelay); ok {
		if !noDelay {
			c.SetNoDelay(noDelay)
		}
	}
}

type (
	ifaceSetKeepAlive interface {
		// SetKeepAlive sets whether the operating system should send
		// keepalive messages on the connection.
		SetKeepAlive(keepalive bool) error
		// SetKeepAlivePeriod sets period between keep alives.
		SetKeepAlivePeriod(d time.Duration) error
	}
	ifaceSetBuffer interface {
		// SetReadBuffer sets the size of the operating system's
		// receive buffer associated with the connection.
		SetReadBuffer(bytes int) error
		// SetWriteBuffer sets the size of the operating system's
		// transmit buffer associated with the connection.
		SetWriteBuffer(bytes int) error
	}
	ifaceSetNoDelay interface {
		// SetNoDelay controls whether the operating system should delay
		// packet transmission in hopes of sending fewer packets (Nagle's
		// algorithm).  The default is true (no delay), meaning that data is
		// sent as soon as possible after a Write.
		SetNoDelay(noDelay bool) error
	}
)

// Connection related system configuration
var (
	writeBuffer     int           = -1
	readBuffer      int           = -1
	changeKeepAlive bool          = false
	keepAlive       bool          = true
	keepAlivePeriod time.Duration = -1
	noDelay         bool          = true
)

// SetKeepAlive sets whether the operating system should send
// keepalive messages on the connection.
// Note: If have not called the function, the system defaults are used.
func SetKeepAlive(keepalive bool) {
	changeKeepAlive = true
	keepAlive = keepalive
}

// SetKeepAlivePeriod sets period between keep alives.
// Note: if d<0, don't change the value.
func SetKeepAlivePeriod(d time.Duration) {
	if d >= 0 {
		keepAlivePeriod = d
	} else {
		fmt.Println("socket: SetKeepAlivePeriod: invalid keepAlivePeriod:", d)
	}
}

// ReadBuffer returns the size of the operating system's
// receive buffer associated with the connection.
// Note: if using the system default value, bytes=-1 and isDefault=true.
func ReadBuffer() (bytes int, isDefault bool) {
	return readBuffer, readBuffer == -1
}

// SetReadBuffer sets the size of the operating system's
// receive buffer associated with the connection.
// Note: if bytes<0, don't change the value.
func SetReadBuffer(bytes int) {
	if bytes >= 0 {
		readBuffer = bytes
	} else {
		fmt.Println("socket: SetReadBuffer: invalid readBuffer size:", bytes)
	}
}

// WriteBuffer returns the size of the operating system's
// transmit buffer associated with the connection.
// Note: if using the system default value, bytes=-1 and isDefault=true.
func WriteBuffer() (bytes int, isDefault bool) {
	return writeBuffer, writeBuffer == -1
}

// SetWriteBuffer sets the size of the operating system's
// transmit buffer associated with the connection.
// Note: if bytes<0, don't change the value.
func SetWriteBuffer(bytes int) {
	if bytes >= 0 {
		writeBuffer = bytes
	} else {
		fmt.Println("socket: SetWriteBuffer: invalid writeBuffer size:", bytes)
	}
}

// SetNoDelay controls whether the operating system should delay
// packet transmission in hopes of sending fewer packets (Nagle's
// algorithm).  The default is true (no delay), meaning that data is
// sent as soon as possible after a Write.
func SetNoDelay(_noDelay bool) {
	noDelay = _noDelay
}

func getProto(protoFuncs []ProtoFunc, rw io.ReadWriter) Proto {
	if len(protoFuncs) > 0 && protoFuncs[0] != nil {
		return protoFuncs[0](rw)
	}
	return defaultProtoFunc(rw)
}
<|endoftext|>


<commit_before>package controllers

import "github.com/robfig/revel"

func init() {
	revel.OnAppStart(Init)
	revel.InterceptMethod((*GorpController).Begin, revel.BEFORE)
	revel.InterceptMethod((*GorpController).Commit, revel.AFTER)
	revel.InterceptMethod((*GorpController).Rollback, revel.FINALLY)
}
<commit_msg>init.go: set GOMAXPROCS<commit_after>package controllers

import (
	"github.com/robfig/revel"
	"runtime"
)

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	revel.OnAppStart(Init)
	revel.InterceptMethod((*GorpController).Begin, revel.BEFORE)
	revel.InterceptMethod((*GorpController).Commit, revel.AFTER)
	revel.InterceptMethod((*GorpController).Rollback, revel.FINALLY)
}
<|endoftext|>


<commit_before>package app

import (
	"bytes"
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/go-martini/martini"
	"github.com/hashicorp/memberlist"
	"github.com/martini-contrib/binding"
	"github.com/martini-contrib/render"
	"net/http"
	"strconv"
	"time"
)

const API_VERSION = "v1"

// TODO Should this live in the config package?
// Using pointers lets us differentiate between a natural 0, and an int-default 0
// omitempty tells it to ignore anything where the name was provided, but an empty value
// inside of a request body.
type ConfigRequest struct {
	VisibilityTimeout *float64 `json:"visibility_timeout,omitempty"`
	MinPartitions     *int     `json:"min_partitions,omitempty"`
	MaxPartitions     *int     `json:"max_partitions,omitempty"`
	MaxPartitionAge   *float64 `json:"max_partition_age,omitempty"`
}

// TODO make message definitions more explicit

func logrusLogger() martini.Handler {
	return func(res http.ResponseWriter, req *http.Request, c martini.Context, log *logrus.Logger) {
		start := time.Now()

		log.WithFields(logrus.Fields{
			"method": req.Method,
			"path":   req.URL.Path,
			"time":   time.Since(start),
		}).Info("Started a request")

		c.Next()
		rw := res.(martini.ResponseWriter)

		log.WithFields(logrus.Fields{
			"method": req.Method,
			"path":   req.URL.Path,
			"status": rw.Status(),
			"time":   time.Since(start),
		}).Info("Completed a request")
	}
}

func dynamiqMartini() *martini.ClassicMartini {
	r := martini.NewRouter()
	m := martini.New()

	m.Map(logrus.New())
	m.Use(logrusLogger())
	m.Use(martini.Recovery())
	m.Use(martini.Static("public"))
	m.MapTo(r, (*martini.Routes)(nil))
	m.Action(r.Handle)
	return &martini.ClassicMartini{m, r}
}

type HTTP_API_V1 struct {
}

func (h HTTP_API_V1) InitWebserver(list *memberlist.Memberlist, cfg *Config) {
	// tieing our Queue to HTTP interface == bad we should move this somewhere else
	// Queues.Queues is dumb. Need a better name-chain
	queues := cfg.Queues
	// also tieing topics this is next for refactor
	topics := InitTopics(cfg, queues)
	m := dynamiqMartini()
	m.Use(render.Renderer())

	// Group the routes underneath their version
	m.Group("/v1", func(r martini.Router) {
		// STATUS / STATISTICS API BLOCK
		m.Get("/status/servers", func() string {
			return_string := ""
			for _, member := range list.Members() {
				return_string += fmt.Sprintf("Member: %s %s\n", member.Name, member.Addr)
			}
			return return_string
		})
		// END STATUS / STATISTICS API BLOCK

		// CONFIGURATION API BLOCK

		m.Delete("/topics/:topic", func(r render.Render, params martini.Params) {
			var present bool
			_, present = topics.TopicMap[params["topic"]]
			if present == true {
				deleted := topics.DeleteTopic(params["topic"])
				r.JSON(200, map[string]interface{}{"Deleted": deleted})
			} else {
				r.JSON(404, map[string]interface{}{"error": "Topic did not exist."})
			}
		})

		m.Put("/queues/:queue", func(r render.Render, params martini.Params) {
			var present bool
			_, present = queues.QueueMap[params["queue"]]
			if present != true {
				cfg.InitializeQueue(params["queue"])
				r.JSON(201, "created")
			} else {
				r.JSON(422, map[string]interface{}{"error": "Queue already exists."})
			}
		})

		m.Put("/topics/:topic", func(r render.Render, params martini.Params) {
			var present bool
			_, present = topics.TopicMap[params["topic"]]
			if present != true {
				topics.InitTopic(params["topic"])
				r.JSON(201, map[string]interface{}{"Queues": topics.TopicMap[params["topic"]].ListQueues()})
			} else {
				r.JSON(422, map[string]interface{}{"error": "Topic already exists."})
			}
		})

		m.Put("/topics/:topic/queues/:queue", func(r render.Render, params martini.Params) {
			var present bool
			_, present = topics.TopicMap[params["topic"]]
			if present != true {
				r.JSON(422, map[string]interface{}{"error": "Topic does not exist. Please create it first."})
			} else {
				_, present = queues.QueueMap[params["queue"]]
				if present != true {
					r.JSON(422, map[string]interface{}{"error": "Queue does not exist. Please create it first"})
				} else {
					topics.TopicMap[params["topic"]].AddQueue(cfg, params["queue"])
					r.JSON(200, map[string]interface{}{"Queues": topics.TopicMap[params["topic"]].ListQueues()})
				}
			}
		})

		// neeeds a little work....
		m.Delete("/topics/:topic/queues/:queue", func(r render.Render, params martini.Params) {
			var present bool
			_, present = topics.TopicMap[params["topic"]]
			if present != true {
				topics.InitTopic(params["topic"])
			}
			topics.TopicMap[params["topic"]].DeleteQueue(cfg, params["queue"])
			r.JSON(200, map[string]interface{}{"Queues": topics.TopicMap[params["topic"]].ListQueues()})
		})

		m.Patch("/queues/:queue", binding.Json(ConfigRequest{}), func(configRequest ConfigRequest, r render.Render, params martini.Params) {
			// Not sure of better way to get the queue name from the request, but would be good to not
			// Have to reach into params - better to bind it

			// There is probably an optimal order to set these in, depending on if the values for min/max
			// have shrunk or grown, so the system can self-adjust in the correct fashion. Or, maybe it's fine to do it however

			// Should we detect and throw errors on incorrect configuration names? Would make debugging simpler

			// Likely all of this belongs in config, where we just pass in an interface to the request object, and it
			// returns the first error it runs across. Would simplify the code here greatly.
			var err error
			if configRequest.VisibilityTimeout != nil {
				err = cfg.SetVisibilityTimeout(params["queue"], *configRequest.VisibilityTimeout)
				// We really need a proper way to generalize error handling
				// Writing this out every time is going to be silly
				if err != nil {
					logrus.Println(err)
					r.JSON(500, map[string]interface{}{"error": err.Error()})
					return
				}
			}

			if configRequest.MinPartitions != nil {
				err = cfg.SetMinPartitions(params["queue"], *configRequest.MinPartitions)
				if err != nil {
					logrus.Println(err)
					r.JSON(500, map[string]interface{}{"error": err.Error()})
					return
				}
			}

			if configRequest.MaxPartitions != nil {
				err = cfg.SetMaxPartitions(params["queue"], *configRequest.MaxPartitions)
				if err != nil {
					logrus.Println(err)
					r.JSON(500, map[string]interface{}{"error": err.Error()})
					return
				}
			}
			if configRequest.MaxPartitionAge != nil {
				err = cfg.SetMaxPartitionAge(params["queue"], *configRequest.MaxPartitionAge)
				if err != nil {
					logrus.Println(err)
					r.JSON(500, map[string]interface{}{"error": err.Error()})
					return
				}
			}

			r.JSON(200, "ok")
		})

		// END CONFIGURATION API BLOCK

		// DATA INTERACTION API BLOCK

		m.Get("/topics", func(r render.Render) {
			topicList := make([]string, 0, 10)
			for topicName, _ := range topics.TopicMap {
				topicList = append(topicList, topicName)
			}
			r.JSON(200, map[string]interface{}{"topics": topicList})
		})

		m.Get("/topics/:topic", func(r render.Render, params martini.Params) {
			var present bool
			_, present = topics.TopicMap[params["topic"]]
			if present != true {
				topics.InitTopic(params["topic"])
			}

			r.JSON(200, map[string]interface{}{"Queues": topics.TopicMap[params["topic"]].ListQueues()})
		})

		m.Put("/topics/:topic/message", func(r render.Render, params martini.Params, req *http.Request) {

			var present bool
			_, present = topics.TopicMap[params["topic"]]
			if present != true {
				topics.InitTopic(params["topic"])
			}
			var buf bytes.Buffer
			buf.ReadFrom(req.Body)

			response := topics.TopicMap[params["topic"]].Broadcast(cfg, buf.String())
			r.JSON(200, response)
		})

		m.Get("/queues", func(r render.Render, params martini.Params) {
			queueList := make([]string, 0, 10)
			for queueName, _ := range queues.QueueMap {
				queueList = append(queueList, queueName)
			}
			r.JSON(200, map[string]interface{}{"queues": queueList})
		})

		m.Get("/queues/:queue", func(r render.Render, params martini.Params) {
			//check if we've initialized this queue yet
			var present bool
			_, present = queues.QueueMap[params["queue"]]
			if present == true {
				queueReturn := make(map[string]interface{})
				queueReturn["visibility_timeout"], _ = cfg.GetVisibilityTimeout(params["queue"])
				queueReturn["min_partitions"], _ = cfg.GetMinPartitions(params["queue"])
				queueReturn["max_partitions"], _ = cfg.GetMaxPartitions(params["queue"])
				queueReturn["max_partition_age"], _ = cfg.GetMaxPartitionAge(params["queue"])
				queueReturn["partitions"] = queues.QueueMap[params["queue"]].Parts.PartitionCount()
				r.JSON(200, queueReturn)
			} else {
				r.JSON(404, fmt.Sprintf("There is no queue named %s", params["queue"]))
			}
		})

		m.Get("/queues/:queue/messages/:batchSize", func(r render.Render, params martini.Params) {
			//check if we've initialized this queue yet
			var present bool
			_, present = queues.QueueMap[params["queue"]]
			if present == true {
				batchSize, err := strconv.ParseUint(params["batchSize"], 10, 32)
				if err != nil {
					//log the error for unparsable input
					logrus.Error(err)
					r.JSON(422, err.Error())
				}
				messages, err := queues.QueueMap[params["queue"]].Get(cfg, list, uint32(batchSize))
				if err != nil {
					r.JSON(500, err.Error())
				}
				//TODO move this into the Queue.Get code
				messageList := make([]map[string]interface{}, 0, 10)
				//Format response
				for _, object := range messages {
					message := make(map[string]interface{})
					message["id"] = object.Key
					message["body"] = string(object.Data[:])
					messageList = append(messageList, message)
				}
				if err != nil {
					logrus.Error(err)
					r.JSON(500, err.Error())
				} else {
					r.JSON(200, messageList)
				}
			} else {
				// What is a sane result here?
				r.JSON(404, fmt.Sprintf("There is no queue named %s", params["queue"]))
			}
		})

		m.Put("/queues/:queue/message", func(params martini.Params, req *http.Request) string {
			var present bool
			_, present = queues.QueueMap[params["queue"]]
			if present == true {
				// parse the request body into a sting
				// TODO clean this up, full json api?
				var buf bytes.Buffer
				buf.ReadFrom(req.Body)
				uuid := queues.QueueMap[params["queue"]].Put(cfg, buf.String())

				return uuid
			} else {
				// What is a sane response to this?
				return ""
			}
		})

		m.Delete("/queues/:queue/message/:messageId", func(r render.Render, params martini.Params) {
			var present bool
			_, present = queues.QueueMap[params["queue"]]
			if present != true {
				cfg.InitializeQueue(params["queue"])
			}

			r.JSON(200, queues.QueueMap[params["queue"]].Delete(cfg, params["messageId"]))
		})
		// DATA INTERACTION API BLOCK
	})
	logrus.Fatal(http.ListenAndServe(":"+strconv.Itoa(cfg.Core.HttpPort), m))
}
<commit_msg>Removing dumb const<commit_after>package app

import (
	"bytes"
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/go-martini/martini"
	"github.com/hashicorp/memberlist"
	"github.com/martini-contrib/binding"
	"github.com/martini-contrib/render"
	"net/http"
	"strconv"
	"time"
)

// TODO Should this live in the config package?
// Using pointers lets us differentiate between a natural 0, and an int-default 0
// omitempty tells it to ignore anything where the name was provided, but an empty value
// inside of a request body.
type ConfigRequest struct {
	VisibilityTimeout *float64 `json:"visibility_timeout,omitempty"`
	MinPartitions     *int     `json:"min_partitions,omitempty"`
	MaxPartitions     *int     `json:"max_partitions,omitempty"`
	MaxPartitionAge   *float64 `json:"max_partition_age,omitempty"`
}

// TODO make message definitions more explicit

func logrusLogger() martini.Handler {
	return func(res http.ResponseWriter, req *http.Request, c martini.Context, log *logrus.Logger) {
		start := time.Now()

		log.WithFields(logrus.Fields{
			"method": req.Method,
			"path":   req.URL.Path,
			"time":   time.Since(start),
		}).Info("Started a request")

		c.Next()
		rw := res.(martini.ResponseWriter)

		log.WithFields(logrus.Fields{
			"method": req.Method,
			"path":   req.URL.Path,
			"status": rw.Status(),
			"time":   time.Since(start),
		}).Info("Completed a request")
	}
}

func dynamiqMartini() *martini.ClassicMartini {
	r := martini.NewRouter()
	m := martini.New()

	m.Map(logrus.New())
	m.Use(logrusLogger())
	m.Use(martini.Recovery())
	m.Use(martini.Static("public"))
	m.MapTo(r, (*martini.Routes)(nil))
	m.Action(r.Handle)
	return &martini.ClassicMartini{m, r}
}

type HTTP_API_V1 struct {
}

func (h HTTP_API_V1) InitWebserver(list *memberlist.Memberlist, cfg *Config) {
	// tieing our Queue to HTTP interface == bad we should move this somewhere else
	// Queues.Queues is dumb. Need a better name-chain
	queues := cfg.Queues
	// also tieing topics this is next for refactor
	topics := InitTopics(cfg, queues)
	m := dynamiqMartini()
	m.Use(render.Renderer())

	// Group the routes underneath their version
	m.Group("/v1", func(r martini.Router) {
		// STATUS / STATISTICS API BLOCK
		m.Get("/status/servers", func() string {
			return_string := ""
			for _, member := range list.Members() {
				return_string += fmt.Sprintf("Member: %s %s\n", member.Name, member.Addr)
			}
			return return_string
		})
		// END STATUS / STATISTICS API BLOCK

		// CONFIGURATION API BLOCK

		m.Delete("/topics/:topic", func(r render.Render, params martini.Params) {
			var present bool
			_, present = topics.TopicMap[params["topic"]]
			if present == true {
				deleted := topics.DeleteTopic(params["topic"])
				r.JSON(200, map[string]interface{}{"Deleted": deleted})
			} else {
				r.JSON(404, map[string]interface{}{"error": "Topic did not exist."})
			}
		})

		m.Put("/queues/:queue", func(r render.Render, params martini.Params) {
			var present bool
			_, present = queues.QueueMap[params["queue"]]
			if present != true {
				cfg.InitializeQueue(params["queue"])
				r.JSON(201, "created")
			} else {
				r.JSON(422, map[string]interface{}{"error": "Queue already exists."})
			}
		})

		m.Put("/topics/:topic", func(r render.Render, params martini.Params) {
			var present bool
			_, present = topics.TopicMap[params["topic"]]
			if present != true {
				topics.InitTopic(params["topic"])
				r.JSON(201, map[string]interface{}{"Queues": topics.TopicMap[params["topic"]].ListQueues()})
			} else {
				r.JSON(422, map[string]interface{}{"error": "Topic already exists."})
			}
		})

		m.Put("/topics/:topic/queues/:queue", func(r render.Render, params martini.Params) {
			var present bool
			_, present = topics.TopicMap[params["topic"]]
			if present != true {
				r.JSON(422, map[string]interface{}{"error": "Topic does not exist. Please create it first."})
			} else {
				_, present = queues.QueueMap[params["queue"]]
				if present != true {
					r.JSON(422, map[string]interface{}{"error": "Queue does not exist. Please create it first"})
				} else {
					topics.TopicMap[params["topic"]].AddQueue(cfg, params["queue"])
					r.JSON(200, map[string]interface{}{"Queues": topics.TopicMap[params["topic"]].ListQueues()})
				}
			}
		})

		// neeeds a little work....
		m.Delete("/topics/:topic/queues/:queue", func(r render.Render, params martini.Params) {
			var present bool
			_, present = topics.TopicMap[params["topic"]]
			if present != true {
				topics.InitTopic(params["topic"])
			}
			topics.TopicMap[params["topic"]].DeleteQueue(cfg, params["queue"])
			r.JSON(200, map[string]interface{}{"Queues": topics.TopicMap[params["topic"]].ListQueues()})
		})

		m.Patch("/queues/:queue", binding.Json(ConfigRequest{}), func(configRequest ConfigRequest, r render.Render, params martini.Params) {
			// Not sure of better way to get the queue name from the request, but would be good to not
			// Have to reach into params - better to bind it

			// There is probably an optimal order to set these in, depending on if the values for min/max
			// have shrunk or grown, so the system can self-adjust in the correct fashion. Or, maybe it's fine to do it however

			// Should we detect and throw errors on incorrect configuration names? Would make debugging simpler

			// Likely all of this belongs in config, where we just pass in an interface to the request object, and it
			// returns the first error it runs across. Would simplify the code here greatly.
			var err error
			if configRequest.VisibilityTimeout != nil {
				err = cfg.SetVisibilityTimeout(params["queue"], *configRequest.VisibilityTimeout)
				// We really need a proper way to generalize error handling
				// Writing this out every time is going to be silly
				if err != nil {
					logrus.Println(err)
					r.JSON(500, map[string]interface{}{"error": err.Error()})
					return
				}
			}

			if configRequest.MinPartitions != nil {
				err = cfg.SetMinPartitions(params["queue"], *configRequest.MinPartitions)
				if err != nil {
					logrus.Println(err)
					r.JSON(500, map[string]interface{}{"error": err.Error()})
					return
				}
			}

			if configRequest.MaxPartitions != nil {
				err = cfg.SetMaxPartitions(params["queue"], *configRequest.MaxPartitions)
				if err != nil {
					logrus.Println(err)
					r.JSON(500, map[string]interface{}{"error": err.Error()})
					return
				}
			}
			if configRequest.MaxPartitionAge != nil {
				err = cfg.SetMaxPartitionAge(params["queue"], *configRequest.MaxPartitionAge)
				if err != nil {
					logrus.Println(err)
					r.JSON(500, map[string]interface{}{"error": err.Error()})
					return
				}
			}

			r.JSON(200, "ok")
		})

		// END CONFIGURATION API BLOCK

		// DATA INTERACTION API BLOCK

		m.Get("/topics", func(r render.Render) {
			topicList := make([]string, 0, 10)
			for topicName, _ := range topics.TopicMap {
				topicList = append(topicList, topicName)
			}
			r.JSON(200, map[string]interface{}{"topics": topicList})
		})

		m.Get("/topics/:topic", func(r render.Render, params martini.Params) {
			var present bool
			_, present = topics.TopicMap[params["topic"]]
			if present != true {
				topics.InitTopic(params["topic"])
			}

			r.JSON(200, map[string]interface{}{"Queues": topics.TopicMap[params["topic"]].ListQueues()})
		})

		m.Put("/topics/:topic/message", func(r render.Render, params martini.Params, req *http.Request) {

			var present bool
			_, present = topics.TopicMap[params["topic"]]
			if present != true {
				topics.InitTopic(params["topic"])
			}
			var buf bytes.Buffer
			buf.ReadFrom(req.Body)

			response := topics.TopicMap[params["topic"]].Broadcast(cfg, buf.String())
			r.JSON(200, response)
		})

		m.Get("/queues", func(r render.Render, params martini.Params) {
			queueList := make([]string, 0, 10)
			for queueName, _ := range queues.QueueMap {
				queueList = append(queueList, queueName)
			}
			r.JSON(200, map[string]interface{}{"queues": queueList})
		})

		m.Get("/queues/:queue", func(r render.Render, params martini.Params) {
			//check if we've initialized this queue yet
			var present bool
			_, present = queues.QueueMap[params["queue"]]
			if present == true {
				queueReturn := make(map[string]interface{})
				queueReturn["visibility_timeout"], _ = cfg.GetVisibilityTimeout(params["queue"])
				queueReturn["min_partitions"], _ = cfg.GetMinPartitions(params["queue"])
				queueReturn["max_partitions"], _ = cfg.GetMaxPartitions(params["queue"])
				queueReturn["max_partition_age"], _ = cfg.GetMaxPartitionAge(params["queue"])
				queueReturn["partitions"] = queues.QueueMap[params["queue"]].Parts.PartitionCount()
				r.JSON(200, queueReturn)
			} else {
				r.JSON(404, fmt.Sprintf("There is no queue named %s", params["queue"]))
			}
		})

		m.Get("/queues/:queue/messages/:batchSize", func(r render.Render, params martini.Params) {
			//check if we've initialized this queue yet
			var present bool
			_, present = queues.QueueMap[params["queue"]]
			if present == true {
				batchSize, err := strconv.ParseUint(params["batchSize"], 10, 32)
				if err != nil {
					//log the error for unparsable input
					logrus.Error(err)
					r.JSON(422, err.Error())
				}
				messages, err := queues.QueueMap[params["queue"]].Get(cfg, list, uint32(batchSize))
				if err != nil {
					r.JSON(500, err.Error())
				}
				//TODO move this into the Queue.Get code
				messageList := make([]map[string]interface{}, 0, 10)
				//Format response
				for _, object := range messages {
					message := make(map[string]interface{})
					message["id"] = object.Key
					message["body"] = string(object.Data[:])
					messageList = append(messageList, message)
				}
				if err != nil {
					logrus.Error(err)
					r.JSON(500, err.Error())
				} else {
					r.JSON(200, messageList)
				}
			} else {
				// What is a sane result here?
				r.JSON(404, fmt.Sprintf("There is no queue named %s", params["queue"]))
			}
		})

		m.Put("/queues/:queue/message", func(params martini.Params, req *http.Request) string {
			var present bool
			_, present = queues.QueueMap[params["queue"]]
			if present == true {
				// parse the request body into a sting
				// TODO clean this up, full json api?
				var buf bytes.Buffer
				buf.ReadFrom(req.Body)
				uuid := queues.QueueMap[params["queue"]].Put(cfg, buf.String())

				return uuid
			} else {
				// What is a sane response to this?
				return ""
			}
		})

		m.Delete("/queues/:queue/message/:messageId", func(r render.Render, params martini.Params) {
			var present bool
			_, present = queues.QueueMap[params["queue"]]
			if present != true {
				cfg.InitializeQueue(params["queue"])
			}

			r.JSON(200, queues.QueueMap[params["queue"]].Delete(cfg, params["messageId"]))
		})
		// DATA INTERACTION API BLOCK
	})
	logrus.Fatal(http.ListenAndServe(":"+strconv.Itoa(cfg.Core.HttpPort), m))
}
<|endoftext|>


<commit_before>// Copyright 2011 The LevelDB-Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package record

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"strings"
	"testing"
)

func short(s string) string {
	if len(s) < 64 {
		return s
	}
	return fmt.Sprintf("%s...(skipping %d bytes)...%s", s[:20], len(s)-40, s[len(s)-20:])
}

// big returns a string of length n, composed of repetitions of partial.
func big(partial string, n int) string {
	return strings.Repeat(partial, n/len(partial)+1)[:n]
}

func TestEmpty(t *testing.T) {
	buf := new(bytes.Buffer)
	r := NewReader(buf)
	if _, err := r.Next(); err != io.EOF {
		t.Fatalf("got %v, want %v", err, io.EOF)
	}
}

func testGenerator(t *testing.T, reset func(), gen func() (string, bool)) {
	buf := new(bytes.Buffer)

	reset()
	w := NewWriter(buf)
	for {
		s, ok := gen()
		if !ok {
			break
		}
		ww, err := w.Next()
		if err != nil {
			t.Fatal(err)
		}
		if _, err := ww.Write([]byte(s)); err != nil {
			t.Fatal(err)
		}
	}
	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	reset()
	r := NewReader(buf)
	for {
		s, ok := gen()
		if !ok {
			break
		}
		rr, err := r.Next()
		if err != nil {
			t.Fatal(err)
		}
		x, err := ioutil.ReadAll(rr)
		if err != nil {
			t.Fatal(err)
		}
		if string(x) != s {
			t.Fatalf("got %q, want %q", short(string(x)), short(s))
		}
	}
	if _, err := r.Next(); err != io.EOF {
		t.Fatalf("got %v, want %v", err, io.EOF)
	}
}

func testLiterals(t *testing.T, s []string) {
	var i int
	reset := func() {
		i = 0
	}
	gen := func() (string, bool) {
		if i == len(s) {
			return "", false
		}
		i++
		return s[i-1], true
	}
	testGenerator(t, reset, gen)
}

func TestMany(t *testing.T) {
	const n = 1e5
	var i int
	reset := func() {
		i = 0
	}
	gen := func() (string, bool) {
		if i == n {
			return "", false
		}
		i++
		return fmt.Sprintf("%d.", i-1), true
	}
	testGenerator(t, reset, gen)
}

func TestRandom(t *testing.T) {
	const n = 1e2
	var (
		i int
		r *rand.Rand
	)
	reset := func() {
		i, r = 0, rand.New(rand.NewSource(0))
	}
	gen := func() (string, bool) {
		if i == n {
			return "", false
		}
		i++
		return strings.Repeat(string(uint8(i)), r.Intn(2*blockSize+16)), true
	}
	testGenerator(t, reset, gen)
}

func TestBasic(t *testing.T) {
	testLiterals(t, []string{
		strings.Repeat("a", 1000),
		strings.Repeat("b", 97270),
		strings.Repeat("c", 8000),
	})
}

func TestBoundary(t *testing.T) {
	for i := blockSize - 16; i < blockSize+16; i++ {
		s0 := big("abcd", i)
		for j := blockSize - 16; j < blockSize+16; j++ {
			s1 := big("ABCDE", j)
			testLiterals(t, []string{s0, s1})
			testLiterals(t, []string{s0, "", s1})
			testLiterals(t, []string{s0, "x", s1})
		}
	}
}

func TestStaleReader(t *testing.T) {
	buf := new(bytes.Buffer)

	w := NewWriter(buf)
	w0, err := w.Next()
	if err != nil {
		t.Fatal(err)
	}
	w0.Write([]byte("0"))
	w1, err := w.Next()
	if err != nil {
		t.Fatal(err)
	}
	w1.Write([]byte("11"))
	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	if got, want := buf.Len(), 2*headerSize+len("0")+len("11"); got != want {
		t.Fatalf("buffer length: got %d want %d", got, want)
	}

	r := NewReader(buf)
	r0, err := r.Next()
	if err != nil {
		t.Fatal(err)
	}
	r1, err := r.Next()
	if err != nil {
		t.Fatal(err)
	}
	p := make([]byte, 1)
	if _, err := r0.Read(p); err == nil || !strings.Contains(err.Error(), "stale") {
		t.Fatalf("stale read #0: unexpected error: %v", err)
	}
	if _, err := r1.Read(p); err != nil {
		t.Fatalf("fresh read #1: got %v want nil error", err)
	}
	if p[0] != '1' {
		t.Fatalf("fresh read #1: byte contents: got '%c' want '1'", p[0])
	}
}

func TestStaleWriter(t *testing.T) {
	buf := new(bytes.Buffer)

	w := NewWriter(buf)
	w0, err := w.Next()
	if err != nil {
		t.Fatal(err)
	}
	w1, err := w.Next()
	if err != nil {
		t.Fatal(err)
	}
	if _, err := w0.Write([]byte("0")); err == nil || !strings.Contains(err.Error(), "stale") {
		t.Fatalf("stale write #0: unexpected error: %v", err)
	}
	if _, err := w1.Write([]byte("11")); err != nil {
		t.Fatalf("fresh write #1: got %v want nil error", err)
	}
	if err := w.Flush(); err != nil {
		t.Fatalf("flush: %v", err)
	}
	if _, err := w1.Write([]byte("0")); err == nil || !strings.Contains(err.Error(), "stale") {
		t.Fatalf("stale write #1: unexpected error: %v", err)
	}
}

// TODO: test flush.
// TODO: test calling Next without exhausting the reader.
<commit_msg>leveldb/record: add more tests.<commit_after>// Copyright 2011 The LevelDB-Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package record

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"strings"
	"testing"
)

func short(s string) string {
	if len(s) < 64 {
		return s
	}
	return fmt.Sprintf("%s...(skipping %d bytes)...%s", s[:20], len(s)-40, s[len(s)-20:])
}

// big returns a string of length n, composed of repetitions of partial.
func big(partial string, n int) string {
	return strings.Repeat(partial, n/len(partial)+1)[:n]
}

func TestEmpty(t *testing.T) {
	buf := new(bytes.Buffer)
	r := NewReader(buf)
	if _, err := r.Next(); err != io.EOF {
		t.Fatalf("got %v, want %v", err, io.EOF)
	}
}

func testGenerator(t *testing.T, reset func(), gen func() (string, bool)) {
	buf := new(bytes.Buffer)

	reset()
	w := NewWriter(buf)
	for {
		s, ok := gen()
		if !ok {
			break
		}
		ww, err := w.Next()
		if err != nil {
			t.Fatal(err)
		}
		if _, err := ww.Write([]byte(s)); err != nil {
			t.Fatal(err)
		}
	}
	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	reset()
	r := NewReader(buf)
	for {
		s, ok := gen()
		if !ok {
			break
		}
		rr, err := r.Next()
		if err != nil {
			t.Fatal(err)
		}
		x, err := ioutil.ReadAll(rr)
		if err != nil {
			t.Fatal(err)
		}
		if string(x) != s {
			t.Fatalf("got %q, want %q", short(string(x)), short(s))
		}
	}
	if _, err := r.Next(); err != io.EOF {
		t.Fatalf("got %v, want %v", err, io.EOF)
	}
}

func testLiterals(t *testing.T, s []string) {
	var i int
	reset := func() {
		i = 0
	}
	gen := func() (string, bool) {
		if i == len(s) {
			return "", false
		}
		i++
		return s[i-1], true
	}
	testGenerator(t, reset, gen)
}

func TestMany(t *testing.T) {
	const n = 1e5
	var i int
	reset := func() {
		i = 0
	}
	gen := func() (string, bool) {
		if i == n {
			return "", false
		}
		i++
		return fmt.Sprintf("%d.", i-1), true
	}
	testGenerator(t, reset, gen)
}

func TestRandom(t *testing.T) {
	const n = 1e2
	var (
		i int
		r *rand.Rand
	)
	reset := func() {
		i, r = 0, rand.New(rand.NewSource(0))
	}
	gen := func() (string, bool) {
		if i == n {
			return "", false
		}
		i++
		return strings.Repeat(string(uint8(i)), r.Intn(2*blockSize+16)), true
	}
	testGenerator(t, reset, gen)
}

func TestBasic(t *testing.T) {
	testLiterals(t, []string{
		strings.Repeat("a", 1000),
		strings.Repeat("b", 97270),
		strings.Repeat("c", 8000),
	})
}

func TestBoundary(t *testing.T) {
	for i := blockSize - 16; i < blockSize+16; i++ {
		s0 := big("abcd", i)
		for j := blockSize - 16; j < blockSize+16; j++ {
			s1 := big("ABCDE", j)
			testLiterals(t, []string{s0, s1})
			testLiterals(t, []string{s0, "", s1})
			testLiterals(t, []string{s0, "x", s1})
		}
	}
}

func TestFlush(t *testing.T) {
	buf := new(bytes.Buffer)
	w := NewWriter(buf)
	// Write a couple of records. Everything should still be held
	// in the record.Writer buffer, so that buf.Len should be 0.
	w0, _ := w.Next()
	w0.Write([]byte("0"))
	w1, _ := w.Next()
	w1.Write([]byte("11"))
	if got, want := buf.Len(), 0; got != want {
		t.Fatalf("buffer length #0: got %d want %d", got, want)
	}
	// Flush the record.Writer buffer, which should yield 17 bytes.
	// 17 = 2*7 + 1 + 2, which is two headers and 1 + 2 payload bytes.
	if err := w.Flush(); err != nil {
		t.Fatal(err)
	}
	if got, want := buf.Len(), 17; got != want {
		t.Fatalf("buffer length #1: got %d want %d", got, want)
	}
	// Do another write, one that isn't large enough to complete the block.
	// The write should not have flowed through to buf.
	w2, _ := w.Next()
	w2.Write(bytes.Repeat([]byte("2"), 10000))
	if got, want := buf.Len(), 17; got != want {
		t.Fatalf("buffer length #2: got %d want %d", got, want)
	}
	// Flushing should get us up to 10024 bytes written.
	// 10024 = 17 + 7 + 10000.
	if err := w.Flush(); err != nil {
		t.Fatal(err)
	}
	if got, want := buf.Len(), 10024; got != want {
		t.Fatalf("buffer length #3: got %d want %d", got, want)
	}
	// Do a bigger write, one that completes the current block.
	// We should now have 32768 bytes (a complete block), without
	// an explicit flush.
	w3, _ := w.Next()
	w3.Write(bytes.Repeat([]byte("3"), 40000))
	if got, want := buf.Len(), 32768; got != want {
		t.Fatalf("buffer length #4: got %d want %d", got, want)
	}
	// Flushing should get us up to 50038 bytes written.
	// 50038 = 10024 + 2*7 + 40000. There are two headers because
	// the one record was split into two chunks.
	if err := w.Flush(); err != nil {
		t.Fatal(err)
	}
	if got, want := buf.Len(), 50038; got != want {
		t.Fatalf("buffer length #5: got %d want %d", got, want)
	}
	// Check that reading those records give the right lengths.
	r := NewReader(buf)
	wants := []int64{1, 2, 10000, 40000}
	for i, want := range wants {
		rr, _ := r.Next()
		n, err := io.Copy(ioutil.Discard, rr)
		if err != nil {
			t.Fatalf("read #%d: %v", i, err)
		}
		if n != want {
			t.Fatalf("read #%d: got %d bytes want %d", i, n, want)
		}
	}
}

func TestNonExhaustiveRead(t *testing.T) {
	const n = 100
	buf := new(bytes.Buffer)
	p := make([]byte, 10)
	rnd := rand.New(rand.NewSource(1))

	w := NewWriter(buf)
	for i := 0; i < n; i++ {
		length := len(p) + rnd.Intn(3*blockSize)
		s := string(uint8(i)) + "123456789abcdefgh"
		ww, _ := w.Next()
		ww.Write([]byte(big(s, length)))
	}
	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	r := NewReader(buf)
	for i := 0; i < n; i++ {
		rr, _ := r.Next()
		_, err := io.ReadFull(rr, p)
		if err != nil {
			t.Fatal(err)
		}
		want := string(uint8(i)) + "123456789"
		if got := string(p); got != want {
			t.Fatalf("read #%d: got %q want %q", i, got, want)
		}
	}
}

func TestStaleReader(t *testing.T) {
	buf := new(bytes.Buffer)

	w := NewWriter(buf)
	w0, err := w.Next()
	if err != nil {
		t.Fatal(err)
	}
	w0.Write([]byte("0"))
	w1, err := w.Next()
	if err != nil {
		t.Fatal(err)
	}
	w1.Write([]byte("11"))
	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	r := NewReader(buf)
	r0, err := r.Next()
	if err != nil {
		t.Fatal(err)
	}
	r1, err := r.Next()
	if err != nil {
		t.Fatal(err)
	}
	p := make([]byte, 1)
	if _, err := r0.Read(p); err == nil || !strings.Contains(err.Error(), "stale") {
		t.Fatalf("stale read #0: unexpected error: %v", err)
	}
	if _, err := r1.Read(p); err != nil {
		t.Fatalf("fresh read #1: got %v want nil error", err)
	}
	if p[0] != '1' {
		t.Fatalf("fresh read #1: byte contents: got '%c' want '1'", p[0])
	}
}

func TestStaleWriter(t *testing.T) {
	buf := new(bytes.Buffer)

	w := NewWriter(buf)
	w0, err := w.Next()
	if err != nil {
		t.Fatal(err)
	}
	w1, err := w.Next()
	if err != nil {
		t.Fatal(err)
	}
	if _, err := w0.Write([]byte("0")); err == nil || !strings.Contains(err.Error(), "stale") {
		t.Fatalf("stale write #0: unexpected error: %v", err)
	}
	if _, err := w1.Write([]byte("11")); err != nil {
		t.Fatalf("fresh write #1: got %v want nil error", err)
	}
	if err := w.Flush(); err != nil {
		t.Fatalf("flush: %v", err)
	}
	if _, err := w1.Write([]byte("0")); err == nil || !strings.Contains(err.Error(), "stale") {
		t.Fatalf("stale write #1: unexpected error: %v", err)
	}
}
<|endoftext|>


<commit_before>package filtergeoip2

import (
	"context"
	"net"

	geoip2 "github.com/oschwald/geoip2-golang"
	"github.com/tsaikd/gogstash/config"
	"github.com/tsaikd/gogstash/config/goglog"
	"github.com/tsaikd/gogstash/config/logevent"
)

// ModuleName is the name used in config file
const ModuleName = "geoip2"

// ErrorTag tag added to event when process geoip2 failed
const ErrorTag = "gogstash_filter_geoip2_error"

// FilterConfig holds the configuration json fields and internal objects
type FilterConfig struct {
	config.FilterConfig

	DBPath      string `json:"db_path"`      // geoip2 db file path, default: GeoLite2-City.mmdb
	IPField     string `json:"ip_field"`     // IP field to get geoip info
	Key         string `json:"key"`          // geoip destination field name, default: geoip
	QuietFail   bool   `json:"quiet"`        // fail quietly
	SkipPrivate bool   `json:"skip_private"` // skip private IP addresses
	FlatFormat  bool   `json:"flat_format"`  // flat format

	db *geoip2.Reader
}

// DefaultFilterConfig returns an FilterConfig struct with default values
func DefaultFilterConfig() FilterConfig {
	return FilterConfig{
		FilterConfig: config.FilterConfig{
			CommonConfig: config.CommonConfig{
				Type: ModuleName,
			},
		},
		DBPath:      "GeoLite2-City.mmdb",
		Key:         "geoip",
		QuietFail:   false, // backwards compatible
		SkipPrivate: false,
		FlatFormat:  false,
	}
}

// InitHandler initialize the filter plugin
func InitHandler(ctx context.Context, raw *config.ConfigRaw) (config.TypeFilterConfig, error) {
	conf := DefaultFilterConfig()
	err := config.ReflectConfig(raw, &conf)
	if err != nil {
		return nil, err
	}

	conf.db, err = geoip2.Open(conf.DBPath)
	if err != nil {
		return nil, err
	}

	return &conf, nil
}

// Event the main filter event
func (f *FilterConfig) Event(ctx context.Context, event logevent.LogEvent) logevent.LogEvent {
	ipstr := event.GetString(f.IPField)
	if ipstr == "" {
		// Passthru if empty
		return event
	}
	ip := net.ParseIP(ipstr)
	if f.SkipPrivate && f.privateIP(ip) {
		// Passthru
		return event
	}
	record, err := f.db.City(ip)
	if err != nil {
		if f.QuietFail {
			goglog.Logger.Error(err)
		}
		event.AddTag(ErrorTag)
		return event
	}
	if record == nil {
		event.AddTag(ErrorTag)
		return event
	}
	if record.Location.Latitude == 0 && record.Location.Longitude == 0 {
		event.AddTag(ErrorTag)
		return event
	}

	if f.FlatFormat {
		m := map[string]interface{}{
			"city_name":      record.City.Names["en"],
			"continent_code": record.Continent.Code,
			"country_code":   record.Country.IsoCode,
			"country_name":   record.Country.Names["en"],
			"ip":             ipstr,
			"latitude":       record.Location.Latitude,
			"location":       []float64{record.Location.Longitude, record.Location.Latitude},
			"longitude":      record.Location.Longitude,
			"postal_code":    record.Postal.Code,
			"timezone":       record.Location.TimeZone,
		}
		if len(record.Subdivisions) > 0 {
			m["region_code"] = record.Subdivisions[0].IsoCode
			m["region_name"] = record.Subdivisions[0].Names["en"]
		}
		event.SetValue(f.Key, m)
	} else {
		m := map[string]interface{}{
			"city": map[string]interface{}{
				"name": record.City.Names["en"],
			},
			"continent": map[string]interface{}{
				"code": record.Continent.Code,
				"name": record.Continent.Names["en"],
			},
			"country": map[string]interface{}{
				"code": record.Country.IsoCode,
				"name": record.Country.Names["en"],
			},
			"ip":        ipstr,
			"latitude":  record.Location.Latitude,
			"location":  []float64{record.Location.Longitude, record.Location.Latitude},
			"longitude": record.Location.Longitude,
			"timezone":  record.Location.TimeZone,
		}
		if len(record.Subdivisions) > 0 {
			m["region"] = map[string]interface{}{
				"code": record.Subdivisions[0].IsoCode,
				"name": record.Subdivisions[0].Names["en"],
			}
		}
		event.SetValue(f.Key, m)
	}

	return event
}

func (f *FilterConfig) privateIP(ip net.IP) bool {
	_, private24BitBlock, _ := net.ParseCIDR("10.0.0.0/8")
	_, private20BitBlock, _ := net.ParseCIDR("172.16.0.0/12")
	_, private16BitBlock, _ := net.ParseCIDR("192.168.0.0/16")
	return private24BitBlock.Contains(ip) || private20BitBlock.Contains(ip) || private16BitBlock.Contains(ip)
}
<commit_msg>Update geoip2 postal_code as optional field<commit_after>package filtergeoip2

import (
	"context"
	"net"

	geoip2 "github.com/oschwald/geoip2-golang"
	"github.com/tsaikd/gogstash/config"
	"github.com/tsaikd/gogstash/config/goglog"
	"github.com/tsaikd/gogstash/config/logevent"
)

// ModuleName is the name used in config file
const ModuleName = "geoip2"

// ErrorTag tag added to event when process geoip2 failed
const ErrorTag = "gogstash_filter_geoip2_error"

// FilterConfig holds the configuration json fields and internal objects
type FilterConfig struct {
	config.FilterConfig

	DBPath      string `json:"db_path"`      // geoip2 db file path, default: GeoLite2-City.mmdb
	IPField     string `json:"ip_field"`     // IP field to get geoip info
	Key         string `json:"key"`          // geoip destination field name, default: geoip
	QuietFail   bool   `json:"quiet"`        // fail quietly
	SkipPrivate bool   `json:"skip_private"` // skip private IP addresses
	FlatFormat  bool   `json:"flat_format"`  // flat format

	db *geoip2.Reader
}

// DefaultFilterConfig returns an FilterConfig struct with default values
func DefaultFilterConfig() FilterConfig {
	return FilterConfig{
		FilterConfig: config.FilterConfig{
			CommonConfig: config.CommonConfig{
				Type: ModuleName,
			},
		},
		DBPath:      "GeoLite2-City.mmdb",
		Key:         "geoip",
		QuietFail:   false, // backwards compatible
		SkipPrivate: false,
		FlatFormat:  false,
	}
}

// InitHandler initialize the filter plugin
func InitHandler(ctx context.Context, raw *config.ConfigRaw) (config.TypeFilterConfig, error) {
	conf := DefaultFilterConfig()
	err := config.ReflectConfig(raw, &conf)
	if err != nil {
		return nil, err
	}

	conf.db, err = geoip2.Open(conf.DBPath)
	if err != nil {
		return nil, err
	}

	return &conf, nil
}

// Event the main filter event
func (f *FilterConfig) Event(ctx context.Context, event logevent.LogEvent) logevent.LogEvent {
	ipstr := event.GetString(f.IPField)
	if ipstr == "" {
		// Passthru if empty
		return event
	}
	ip := net.ParseIP(ipstr)
	if f.SkipPrivate && f.privateIP(ip) {
		// Passthru
		return event
	}
	record, err := f.db.City(ip)
	if err != nil {
		if f.QuietFail {
			goglog.Logger.Error(err)
		}
		event.AddTag(ErrorTag)
		return event
	}
	if record == nil {
		event.AddTag(ErrorTag)
		return event
	}
	if record.Location.Latitude == 0 && record.Location.Longitude == 0 {
		event.AddTag(ErrorTag)
		return event
	}

	if f.FlatFormat {
		m := map[string]interface{}{
			"city_name":      record.City.Names["en"],
			"continent_code": record.Continent.Code,
			"country_code":   record.Country.IsoCode,
			"country_name":   record.Country.Names["en"],
			"ip":             ipstr,
			"latitude":       record.Location.Latitude,
			"location":       []float64{record.Location.Longitude, record.Location.Latitude},
			"longitude":      record.Location.Longitude,
			"timezone":       record.Location.TimeZone,
		}
		if record.Postal.Code != "" {
			m["postal_code"] = record.Postal.Code
		}
		if len(record.Subdivisions) > 0 {
			m["region_code"] = record.Subdivisions[0].IsoCode
			m["region_name"] = record.Subdivisions[0].Names["en"]
		}
		event.SetValue(f.Key, m)
	} else {
		m := map[string]interface{}{
			"city": map[string]interface{}{
				"name": record.City.Names["en"],
			},
			"continent": map[string]interface{}{
				"code": record.Continent.Code,
				"name": record.Continent.Names["en"],
			},
			"country": map[string]interface{}{
				"code": record.Country.IsoCode,
				"name": record.Country.Names["en"],
			},
			"ip":        ipstr,
			"latitude":  record.Location.Latitude,
			"location":  []float64{record.Location.Longitude, record.Location.Latitude},
			"longitude": record.Location.Longitude,
			"timezone":  record.Location.TimeZone,
		}
		if len(record.Subdivisions) > 0 {
			m["region"] = map[string]interface{}{
				"code": record.Subdivisions[0].IsoCode,
				"name": record.Subdivisions[0].Names["en"],
			}
		}
		event.SetValue(f.Key, m)
	}

	return event
}

func (f *FilterConfig) privateIP(ip net.IP) bool {
	_, private24BitBlock, _ := net.ParseCIDR("10.0.0.0/8")
	_, private20BitBlock, _ := net.ParseCIDR("172.16.0.0/12")
	_, private16BitBlock, _ := net.ParseCIDR("192.168.0.0/16")
	return private24BitBlock.Contains(ip) || private20BitBlock.Contains(ip) || private16BitBlock.Contains(ip)
}
<|endoftext|>


<commit_before>package air

import (
	"bytes"
	"errors"
	"fmt"
	"net/http"
	"path"
	"reflect"
	"runtime"
	"sync"

	"golang.org/x/net/context"
)

type (
	// Air is the top-level framework type.
	Air struct {
		pregases []GasFunc
		gases    []GasFunc
		maxParam *int
		pool     sync.Pool

		Router           *Router
		Binder           *Binder
		Renderer         *Renderer
		HTTPErrorHandler HTTPErrorHandler
		Logger           Logger
		Debug            bool
	}

	// HTTPError represents an error that occurred while handling a request.
	HTTPError struct {
		Code    int
		Message string
	}

	// HandlerFunc defines a function to server HTTP requests.
	HandlerFunc func(*Context) error

	// GasFunc defines a function to process gas.
	GasFunc func(HandlerFunc) HandlerFunc

	// HTTPErrorHandler is a centralized HTTP error handler.
	HTTPErrorHandler func(error, *Context)
)

// HTTP methods (which follows the REST principles)
const (
	GET    = "GET"
	POST   = "POST"
	PUT    = "PUT" // The Air advise you to forget the PATCH.
	DELETE = "DELETE"
)

// MIME types
const (
	MIMEApplicationJSON       = "application/json; charset=utf-8"
	MIMEApplicationJavaScript = "application/javascript; charset=utf-8"
	MIMEApplicationXML        = "application/xml; charset=utf-8"
	MIMEApplicationForm       = "application/x-www-form-urlencoded"
	MIMEApplicationProtobuf   = "application/protobuf"
	MIMEApplicationMsgpack    = "application/msgpack"
	MIMETextHTML              = "text/html; charset=utf-8"
	MIMETextPlain             = "text/plain; charset=utf-8"
	MIMEMultipartForm         = "multipart/form-data"
	MIMEOctetStream           = "application/octet-stream"
)

// Headers
const (
	HeaderAcceptEncoding                = "Accept-Encoding"
	HeaderAllow                         = "Allow"
	HeaderAuthorization                 = "Authorization"
	HeaderContentDisposition            = "Content-Disposition"
	HeaderContentEncoding               = "Content-Encoding"
	HeaderContentLength                 = "Content-Length"
	HeaderContentType                   = "Content-Type"
	HeaderCookie                        = "Cookie"
	HeaderSetCookie                     = "Set-Cookie"
	HeaderIfModifiedSince               = "If-Modified-Since"
	HeaderLastModified                  = "Last-Modified"
	HeaderLocation                      = "Location"
	HeaderUpgrade                       = "Upgrade"
	HeaderVary                          = "Vary"
	HeaderWWWAuthenticate               = "WWW-Authenticate"
	HeaderXForwardedProto               = "X-Forwarded-Proto"
	HeaderXHTTPMethodOverride           = "X-HTTP-Method-Override"
	HeaderXForwardedFor                 = "X-Forwarded-For"
	HeaderXRealIP                       = "X-Real-IP"
	HeaderServer                        = "Server"
	HeaderOrigin                        = "Origin"
	HeaderAccessControlRequestMethod    = "Access-Control-Request-Method"
	HeaderAccessControlRequestHeaders   = "Access-Control-Request-Headers"
	HeaderAccessControlAllowOrigin      = "Access-Control-Allow-Origin"
	HeaderAccessControlAllowMethods     = "Access-Control-Allow-Methods"
	HeaderAccessControlAllowHeaders     = "Access-Control-Allow-Headers"
	HeaderAccessControlAllowCredentials = "Access-Control-Allow-Credentials"
	HeaderAccessControlExposeHeaders    = "Access-Control-Expose-Headers"
	HeaderAccessControlMaxAge           = "Access-Control-Max-Age"

	HeaderStrictTransportSecurity = "Strict-Transport-Security"
	HeaderXContentTypeOptions     = "X-Content-Type-Options"
	HeaderXXSSProtection          = "X-XSS-Protection"
	HeaderXFrameOptions           = "X-Frame-Options"
	HeaderContentSecurityPolicy   = "Content-Security-Policy"
	HeaderXCSRFToken              = "X-CSRF-Token"
)

// For easy for-range
var methods = [4]string{GET, POST, PUT, DELETE}

// Errors
var (
	ErrUnauthorized                = NewHTTPError(http.StatusUnauthorized)          // 401
	ErrNotFound                    = NewHTTPError(http.StatusNotFound)              // 404
	ErrMethodNotAllowed            = NewHTTPError(http.StatusMethodNotAllowed)      // 405
	ErrStatusRequestEntityTooLarge = NewHTTPError(http.StatusRequestEntityTooLarge) // 413
	ErrUnsupportedMediaType        = NewHTTPError(http.StatusUnsupportedMediaType)  // 415

	ErrInternalServerError = NewHTTPError(http.StatusInternalServerError) // 500
	ErrBadGateway          = NewHTTPError(http.StatusBadGateway)          // 502
	ErrServiceUnavailable  = NewHTTPError(http.StatusServiceUnavailable)  // 503
	ErrGatewayTimeout      = NewHTTPError(http.StatusGatewayTimeout)      // 504

	ErrInvalidRedirectCode = errors.New("Invalid Redirect Status Code")
	ErrCookieNotFound      = errors.New("Cookie Not Found")

	ErrDataTemplateNotSetted = errors.New("c.Data[\"template\"] Not Setted")
	ErrDataHTMLNotSetted     = errors.New("c.Data[\"html\"] Not Setted")
	ErrDataStringNotSetted   = errors.New("c.Data[\"string\"] Not Setted")
	ErrDataJSONNotSetted     = errors.New("c.Data[\"json\"] Not Setted")
	ErrDataJSONPNotSetted    = errors.New("c.Data[\"jsonp\"] Not Setted")
	ErrDataCallbackNotSetted = errors.New("c.Data[\"callback\"] Not Setted")
	ErrDataXMLNotSetted      = errors.New("c.Data[\"xml\"] Not Setted")
)

// Error handlers
var (
	NotFoundHandler = func(c *Context) error {
		return ErrNotFound
	}

	MethodNotAllowedHandler = func(c *Context) error {
		return ErrMethodNotAllowed
	}
)

// New creates an instance of Air.
func New() *Air {
	a := &Air{maxParam: new(int)}
	a.pool.New = func() interface{} {
		return a.NewContext(&Request{}, &Response{})
	}
	a.Router = NewRouter(a)

	// Defaults
	a.HTTPErrorHandler = a.DefaultHTTPErrorHandler
	a.Binder = &Binder{}
	a.Renderer = &Renderer{ViewsPath: "views"}
	a.Renderer.initDefaultTempleFuncMap()
	l := NewLogger("air")
	l.SetLevel(ERROR)
	a.Logger = l
	return a
}

// NewContext returns a Context instance.
func (a *Air) NewContext(req *Request, res *Response) *Context {
	return &Context{
		goContext: context.Background(),

		Request:     req,
		Response:    res,
		ParamValues: make([]string, *a.maxParam),
		Data:        make(map[string]interface{}),
		StatusCode:  http.StatusOK,
		Handler:     NotFoundHandler,
		Air:         a,
	}
}

// DefaultHTTPErrorHandler invokes the default HTTP error handler.
func (a *Air) DefaultHTTPErrorHandler(err error, c *Context) {
	code := http.StatusInternalServerError
	msg := http.StatusText(code)
	if he, ok := err.(*HTTPError); ok {
		code = he.Code
		msg = he.Message
	}
	if a.Debug {
		msg = err.Error()
	}
	if !c.Response.Committed {
		c.Data["string"] = msg
		c.StatusCode = code
		c.String()
	}
	a.Logger.Error(err)
}

// Precontain adds gases to the chain which is run before router.
func (a *Air) Precontain(gases ...GasFunc) {
	a.pregases = append(a.pregases, gases...)
}

// Contain adds gases to the chain which is run after router.
func (a *Air) Contain(gases ...GasFunc) {
	a.gases = append(a.gases, gases...)
}

// Get registers a new GET route for a path with matching handler in the router
// with optional route-level gases.
func (a *Air) Get(path string, handler HandlerFunc, gases ...GasFunc) {
	a.add(GET, path, handler, gases...)
}

// Post registers a new POST route for a path with matching handler in the
// router with optional route-level gases.
func (a *Air) Post(path string, handler HandlerFunc, gases ...GasFunc) {
	a.add(POST, path, handler, gases...)
}

// Put registers a new PUT route for a path with matching handler in the
// router with optional route-level gases.
func (a *Air) Put(path string, handler HandlerFunc, gases ...GasFunc) {
	a.add(PUT, path, handler, gases...)
}

// Delete registers a new DELETE route for a path with matching handler in the router
// with optional route-level gases.
func (a *Air) Delete(path string, handler HandlerFunc, gases ...GasFunc) {
	a.add(DELETE, path, handler, gases...)
}

// Static registers a new route with path prefix to serve static files from the
// provided root directory.
func (a *Air) Static(prefix, root string) {
	a.Get(prefix+"*", func(c *Context) error {
		return c.File(path.Join(root, c.P(0)))
	})
}

// File registers a new route with path to serve a static file.
func (a *Air) File(path, file string) {
	a.Get(path, func(c *Context) error {
		return c.File(file)
	})
}

func (a *Air) add(method, path string, handler HandlerFunc, gases ...GasFunc) {
	name := handlerName(handler)
	a.Router.Add(method, path, func(c *Context) error {
		h := handler
		// Chain gases
		for i := len(gases) - 1; i >= 0; i-- {
			h = gases[i](h)
		}
		return h(c)
	}, a)
	r := Route{
		Method:  method,
		Path:    path,
		Handler: name,
	}
	a.Router.Routes[method+path] = r
}

// Group creates a new router group with prefix and optional group-level gases.
func (a *Air) Group(prefix string, gases ...GasFunc) *Group {
	g := &Group{prefix: prefix, air: a}
	g.Contain(gases...)
	return g
}

// URI generates a URI from handler.
func (a *Air) URI(handler HandlerFunc, params ...interface{}) string {
	uri := new(bytes.Buffer)
	ln := len(params)
	n := 0
	name := handlerName(handler)
	for _, r := range a.Router.Routes {
		if r.Handler == name {
			for i, l := 0, len(r.Path); i < l; i++ {
				if r.Path[i] == ':' && n < ln {
					for ; i < l && r.Path[i] != '/'; i++ {
					}
					uri.WriteString(fmt.Sprintf("%v", params[n]))
					n++
				}
				if i < l {
					uri.WriteByte(r.Path[i])
				}
			}
			break
		}
	}
	return uri.String()
}

// AcquireContext returns an empty `Context` instance from the pool.
// You must be return the context by calling `ReleaseContext()`.
func (a *Air) AcquireContext() Context {
	return a.pool.Get().(Context)
}

// ReleaseContext returns the `Context` instance back to the pool.
// You must call it after `AcquireContext()`.
func (a *Air) ReleaseContext(c Context) {
	a.pool.Put(c)
}

func (a *Air) ServeHTTP(req *Request, res *Response) {
	c := a.pool.Get().(*Context)
	c.Reset(req, res)

	// Gases
	h := func(*Context) error {
		method := req.Method()
		path := req.URI.Path()
		a.Router.Find(method, path, c)
		h := c.Handler
		for i := len(a.gases) - 1; i >= 0; i-- {
			h = a.gases[i](h)
		}
		return h(c)
	}

	// Pregases
	for i := len(a.pregases) - 1; i >= 0; i-- {
		h = a.pregases[i](h)
	}

	// Execute chain
	if err := h(c); err != nil {
		a.HTTPErrorHandler(err, c)
	}

	a.pool.Put(c)
}

// Run starts the HTTP server.
func (a *Air) Run(addr string) {
	s := NewServer(addr)
	s.SetHandler(a)
	s.SetLogger(a.Logger)
	a.Renderer.parseTemplates()
	if a.Debug {
		a.Logger.SetLevel(DEBUG)
		a.Logger.Debug("Running In Debug Mode")
	}
	a.Logger.Error(s.Start())
}

// NewHTTPError creates a new HTTPError instance.
func NewHTTPError(code int, msg ...string) *HTTPError {
	he := &HTTPError{Code: code, Message: http.StatusText(code)}
	if len(msg) > 0 {
		he.Message = msg[0]
	}
	return he
}

// Error makes it compatible with `error` interface.
func (he *HTTPError) Error() string {
	return he.Message
}

// WrapGas wrap `HandlerFunc` into `GasFunc`.
func WrapGas(handler HandlerFunc) GasFunc {
	return func(next HandlerFunc) HandlerFunc {
		return func(c *Context) error {
			if err := handler(c); err != nil {
				return err
			}
			return next(c)
		}
	}
}

func handlerName(handler HandlerFunc) string {
	t := reflect.ValueOf(handler).Type()
	if t.Kind() == reflect.Func {
		return runtime.FuncForPC(reflect.ValueOf(handler).Pointer()).Name()
	}
	return t.String()
}
<commit_msg>style: beautify air.go<commit_after>package air

import (
	"bytes"
	"errors"
	"fmt"
	"net/http"
	"path"
	"reflect"
	"runtime"
	"sync"

	"golang.org/x/net/context"
)

type (
	// Air is the top-level framework struct.
	Air struct {
		pregases []GasFunc
		gases    []GasFunc
		maxParam *int
		pool     sync.Pool

		Router           *Router
		Binder           *Binder
		Renderer         *Renderer
		HTTPErrorHandler HTTPErrorHandler
		Logger           Logger
		Debug            bool
	}

	// HTTPError represents an error that occurred while handling a request.
	HTTPError struct {
		Code    int
		Message string
	}

	// HandlerFunc defines a function to server HTTP requests.
	HandlerFunc func(*Context) error

	// GasFunc defines a function to process gas.
	GasFunc func(HandlerFunc) HandlerFunc

	// HTTPErrorHandler is a centralized HTTP error handler.
	HTTPErrorHandler func(error, *Context)
)

// HTTP methods (which follows the REST principle)
const (
	GET    = "GET"
	POST   = "POST"
	PUT    = "PUT" // The Air advises you to forget the PATCH.
	DELETE = "DELETE"
)

// For easy for-range
var methods = [4]string{GET, POST, PUT, DELETE}

// MIME types
const (
	MIMEApplicationJSON       = "application/json; charset=utf-8"
	MIMEApplicationJavaScript = "application/javascript; charset=utf-8"
	MIMEApplicationXML        = "application/xml; charset=utf-8"
	MIMEApplicationForm       = "application/x-www-form-urlencoded"
	MIMEApplicationProtobuf   = "application/protobuf"
	MIMEApplicationMsgpack    = "application/msgpack"
	MIMETextHTML              = "text/html; charset=utf-8"
	MIMETextPlain             = "text/plain; charset=utf-8"
	MIMEMultipartForm         = "multipart/form-data"
	MIMEOctetStream           = "application/octet-stream"
)

// Headers
const (
	HeaderAcceptEncoding                = "Accept-Encoding"
	HeaderAllow                         = "Allow"
	HeaderAuthorization                 = "Authorization"
	HeaderContentDisposition            = "Content-Disposition"
	HeaderContentEncoding               = "Content-Encoding"
	HeaderContentLength                 = "Content-Length"
	HeaderContentType                   = "Content-Type"
	HeaderCookie                        = "Cookie"
	HeaderSetCookie                     = "Set-Cookie"
	HeaderIfModifiedSince               = "If-Modified-Since"
	HeaderLastModified                  = "Last-Modified"
	HeaderLocation                      = "Location"
	HeaderUpgrade                       = "Upgrade"
	HeaderVary                          = "Vary"
	HeaderWWWAuthenticate               = "WWW-Authenticate"
	HeaderXForwardedProto               = "X-Forwarded-Proto"
	HeaderXHTTPMethodOverride           = "X-HTTP-Method-Override"
	HeaderXForwardedFor                 = "X-Forwarded-For"
	HeaderXRealIP                       = "X-Real-IP"
	HeaderServer                        = "Server"
	HeaderOrigin                        = "Origin"
	HeaderAccessControlRequestMethod    = "Access-Control-Request-Method"
	HeaderAccessControlRequestHeaders   = "Access-Control-Request-Headers"
	HeaderAccessControlAllowOrigin      = "Access-Control-Allow-Origin"
	HeaderAccessControlAllowMethods     = "Access-Control-Allow-Methods"
	HeaderAccessControlAllowHeaders     = "Access-Control-Allow-Headers"
	HeaderAccessControlAllowCredentials = "Access-Control-Allow-Credentials"
	HeaderAccessControlExposeHeaders    = "Access-Control-Expose-Headers"
	HeaderAccessControlMaxAge           = "Access-Control-Max-Age"

	HeaderStrictTransportSecurity = "Strict-Transport-Security"
	HeaderXContentTypeOptions     = "X-Content-Type-Options"
	HeaderXXSSProtection          = "X-XSS-Protection"
	HeaderXFrameOptions           = "X-Frame-Options"
	HeaderContentSecurityPolicy   = "Content-Security-Policy"
	HeaderXCSRFToken              = "X-CSRF-Token"
)

// Errors
var (
	ErrUnauthorized                = NewHTTPError(http.StatusUnauthorized)          // 401
	ErrNotFound                    = NewHTTPError(http.StatusNotFound)              // 404
	ErrMethodNotAllowed            = NewHTTPError(http.StatusMethodNotAllowed)      // 405
	ErrStatusRequestEntityTooLarge = NewHTTPError(http.StatusRequestEntityTooLarge) // 413
	ErrUnsupportedMediaType        = NewHTTPError(http.StatusUnsupportedMediaType)  // 415

	ErrInternalServerError = NewHTTPError(http.StatusInternalServerError) // 500
	ErrBadGateway          = NewHTTPError(http.StatusBadGateway)          // 502
	ErrServiceUnavailable  = NewHTTPError(http.StatusServiceUnavailable)  // 503
	ErrGatewayTimeout      = NewHTTPError(http.StatusGatewayTimeout)      // 504

	ErrInvalidRedirectCode = errors.New("Invalid Redirect Status Code")
	ErrCookieNotFound      = errors.New("Cookie Not Found")

	ErrDataTemplateNotSetted = errors.New("c.Data[\"template\"] Not Setted")
	ErrDataHTMLNotSetted     = errors.New("c.Data[\"html\"] Not Setted")
	ErrDataStringNotSetted   = errors.New("c.Data[\"string\"] Not Setted")
	ErrDataJSONNotSetted     = errors.New("c.Data[\"json\"] Not Setted")
	ErrDataJSONPNotSetted    = errors.New("c.Data[\"jsonp\"] Not Setted")
	ErrDataCallbackNotSetted = errors.New("c.Data[\"callback\"] Not Setted")
	ErrDataXMLNotSetted      = errors.New("c.Data[\"xml\"] Not Setted")
)

// HTTP error handlers
var (
	NotFoundHandler = func(c *Context) error {
		return ErrNotFound
	}

	MethodNotAllowedHandler = func(c *Context) error {
		return ErrMethodNotAllowed
	}
)

// New returns a new instance of `Air`.
func New() *Air {
	a := &Air{maxParam: new(int)}
	a.pool.New = func() interface{} {
		return a.NewContext(&Request{}, &Response{})
	}
	a.Router = NewRouter(a)

	// Defaults
	a.HTTPErrorHandler = a.DefaultHTTPErrorHandler
	a.Binder = &Binder{}
	a.Renderer = &Renderer{ViewsPath: "views"}
	a.Renderer.initDefaultTempleFuncMap()
	l := NewLogger("air")
	l.SetLevel(ERROR)
	a.Logger = l
	return a
}

// NewContext returns a new instance of `Context`.
func (a *Air) NewContext(req *Request, res *Response) *Context {
	return &Context{
		goContext: context.Background(),

		Request:     req,
		Response:    res,
		ParamValues: make([]string, *a.maxParam),
		Data:        make(map[string]interface{}),
		StatusCode:  http.StatusOK,
		Handler:     NotFoundHandler,
		Air:         a,
	}
}

// DefaultHTTPErrorHandler invokes the default HTTP error handler.
func (a *Air) DefaultHTTPErrorHandler(err error, c *Context) {
	code := http.StatusInternalServerError
	msg := http.StatusText(code)
	if he, ok := err.(*HTTPError); ok {
		code = he.Code
		msg = he.Message
	}
	if a.Debug {
		msg = err.Error()
	}
	if !c.Response.Committed {
		c.Data["string"] = msg
		c.StatusCode = code
		c.String()
	}
	a.Logger.Error(err)
}

// Precontain adds gases to the chain which is run before router.
func (a *Air) Precontain(gases ...GasFunc) {
	a.pregases = append(a.pregases, gases...)
}

// Contain adds gases to the chain which is run after router.
func (a *Air) Contain(gases ...GasFunc) {
	a.gases = append(a.gases, gases...)
}

// Get registers a new GET route for a path with matching handler in the router
// with optional route-level gases.
func (a *Air) Get(path string, handler HandlerFunc, gases ...GasFunc) {
	a.add(GET, path, handler, gases...)
}

// Post registers a new POST route for a path with matching handler in the
// router with optional route-level gases.
func (a *Air) Post(path string, handler HandlerFunc, gases ...GasFunc) {
	a.add(POST, path, handler, gases...)
}

// Put registers a new PUT route for a path with matching handler in the
// router with optional route-level gases.
func (a *Air) Put(path string, handler HandlerFunc, gases ...GasFunc) {
	a.add(PUT, path, handler, gases...)
}

// Delete registers a new DELETE route for a path with matching handler in
// the router with optional route-level gases.
func (a *Air) Delete(path string, handler HandlerFunc, gases ...GasFunc) {
	a.add(DELETE, path, handler, gases...)
}

// Static registers a new route with path prefix to serve static files from
// the provided root directory.
func (a *Air) Static(prefix, root string) {
	a.Get(prefix+"*", func(c *Context) error {
		return c.File(path.Join(root, c.P(0)))
	})
}

// File registers a new route with path to serve a static file.
func (a *Air) File(path, file string) {
	a.Get(path, func(c *Context) error {
		return c.File(file)
	})
}

// add registers a new route for a path with a HTTP method and matching handler
// in the router with optional route-level gases.
func (a *Air) add(method, path string, handler HandlerFunc, gases ...GasFunc) {
	name := handlerName(handler)
	a.Router.Add(method, path, func(c *Context) error {
		h := handler
		// Chain gases
		for i := len(gases) - 1; i >= 0; i-- {
			h = gases[i](h)
		}
		return h(c)
	}, a)
	r := Route{
		Method:  method,
		Path:    path,
		Handler: name,
	}
	a.Router.Routes[method+path] = r
}

// Group returns a new router group with prefix and optional group-level gases.
func (a *Air) Group(prefix string, gases ...GasFunc) *Group {
	g := &Group{prefix: prefix, air: a}
	g.Contain(gases...)
	return g
}

// URI returns a URI generated from handler.
func (a *Air) URI(handler HandlerFunc, params ...interface{}) string {
	uri := new(bytes.Buffer)
	ln := len(params)
	n := 0
	name := handlerName(handler)
	for _, r := range a.Router.Routes {
		if r.Handler == name {
			for i, l := 0, len(r.Path); i < l; i++ {
				if r.Path[i] == ':' && n < ln {
					for ; i < l && r.Path[i] != '/'; i++ {
					}
					uri.WriteString(fmt.Sprintf("%v", params[n]))
					n++
				}
				if i < l {
					uri.WriteByte(r.Path[i])
				}
			}
			break
		}
	}
	return uri.String()
}

// AcquireContext returns an empty `Context` instance from the pool.
// You must be return the context by calling `Air#ReleaseContext()`.
func (a *Air) AcquireContext() Context {
	return a.pool.Get().(Context)
}

// ReleaseContext returns the `Context` instance back to the pool.
// You must call it after `Air#AcquireContext()`.
func (a *Air) ReleaseContext(c Context) {
	a.pool.Put(c)
}

// ServeHTTP implements `Handler#ServeHTTP()`.
func (a *Air) ServeHTTP(req *Request, res *Response) {
	c := a.pool.Get().(*Context)
	c.Reset(req, res)

	// Gases
	h := func(*Context) error {
		method := req.Method()
		path := req.URI.Path()
		a.Router.Find(method, path, c)
		h := c.Handler
		for i := len(a.gases) - 1; i >= 0; i-- {
			h = a.gases[i](h)
		}
		return h(c)
	}

	// Pregases
	for i := len(a.pregases) - 1; i >= 0; i-- {
		h = a.pregases[i](h)
	}

	// Execute chain
	if err := h(c); err != nil {
		a.HTTPErrorHandler(err, c)
	}

	a.pool.Put(c)
}

// Run starts the HTTP server.
func (a *Air) Run(addr string) {
	s := NewServer(addr)
	s.SetHandler(a)
	s.SetLogger(a.Logger)
	a.Renderer.parseTemplates()
	if a.Debug {
		a.Logger.SetLevel(DEBUG)
		a.Logger.Debug("Running In Debug Mode")
	}
	a.Logger.Error(s.Start())
}

// NewHTTPError returns a new instance of `HTTPError`.
func NewHTTPError(code int, msg ...string) *HTTPError {
	he := &HTTPError{Code: code, Message: http.StatusText(code)}
	if len(msg) > 0 {
		he.Message = msg[0]
	}
	return he
}

// Error implements `error#Error()`.
func (he *HTTPError) Error() string {
	return he.Message
}

// WrapGas wraps `HandlerFunc` into `GasFunc`.
func WrapGas(handler HandlerFunc) GasFunc {
	return func(next HandlerFunc) HandlerFunc {
		return func(c *Context) error {
			if err := handler(c); err != nil {
				return err
			}
			return next(c)
		}
	}
}

// handlerName returns the handler's func name.
func handlerName(handler HandlerFunc) string {
	t := reflect.ValueOf(handler).Type()
	if t.Kind() == reflect.Func {
		return runtime.FuncForPC(reflect.ValueOf(handler).Pointer()).Name()
	}
	return t.String()
}
<|endoftext|>