<commit_before>/********************************************************

       纯真IP数据库查询
       by:ying32
       2016/5/13
       QQ: 1444386932
--------------------------------------------------------

       未测试多种带有重定义或未重定义的IP
       如果查询出错，可发邮件给我，告知改进
       或者由您改进，但记得发我一份

       纯真网络：
         http://www.cz88.net/
       纯真IP数据格式：
         http://lumaqq.linuxsir.org/
                     article/qqwry_format_detail.html

********************************************************/

package qqwry

import (
	"bytes"
	//"code.google.com/p/mahonia"
	"github.com/axgle/mahonia"
	"encoding/binary"
	"fmt"
	"net"
	"os"
)

// 重定义模式
const (
	REDIRECT_MODE_1 = 1
	REDIRECT_MODE_2 = 2
)

var (
	// 全局数据，只一次
	dataReader *bytes.Reader
	// 标识"qqwry.dat"文件是否初始
	isInitReader bool
)

type QQWry struct {
	RecordCount uint32
	FirstRecord uint32
}

// 模拟构造函数
func NewQQWry(fileName string) *QQWry {
	// 只初始一次文件
	if !isInitReader {
		initFile(fileName)
		isInitReader = true
	}
	var rf, rl, rc uint32
	binary.Read(dataReader, binary.LittleEndian, &rf)
	binary.Read(dataReader, binary.LittleEndian, &rl)
	rc = (rl - rf) / 7
	return &QQWry{rc, rf}
}

// 获取版本信息
func (self *QQWry) Version() string {
	if self.RecordCount > 0 {
		offset, _, _ := self.getIPOffset(self.RecordCount, 1)
		if offset > 0 {
			return self.getIpLocationOf(offset)
		}
	}
	return "未知区域"
}

// 读取一个3字节并转为整数
func (self *QQWry) read3byteint() int64 {
	temp := make([]byte, 3)
	dataReader.Read(temp)
	return int64(uint32(temp[0]) | uint32(temp[1])<<8 | uint32(temp[2])<<16)
}

// 获取当前Reader位置
func (self *QQWry) dataReaderPosition() int64 {
	return dataReader.Size() - int64(dataReader.Len())
}

// AnsiToUtf8的转换
func (self *QQWry) ansiToUtf8(bytes []byte) string {
	decoder := mahonia.NewDecoder("gbk")
	if res, ok := decoder.ConvertStringOK(string(bytes)); ok {
		return res
	}
	return ""
}

// 字符串ip转整型
func (self *QQWry) StrToIP(ipstr string) uint32 {
	ip := net.ParseIP(ipstr)
	return binary.BigEndian.Uint32(ip.To4())
}

// 整型转ip字符，ipv4
func (self *QQWry) IPToStr(ip uint32) string {
	b := make([]byte, 4)
	binary.BigEndian.PutUint32(b, ip)
	return fmt.Sprintf("%d.%d.%d.%d", b[0], b[1], b[2], b[3])
}

// 获取ip偏移位置
func (self *QQWry) getIPOffset(recNum uint32, flag int) (int64, uint32, uint32) {
	var begin, end uint32
	var result int64 = int64(self.FirstRecord + recNum*7)
	dataReader.Seek(result, 0)
	binary.Read(dataReader, binary.LittleEndian, &begin)
	if flag == 1 {
		result = self.read3byteint()
		dataReader.Seek(result, 0)
		binary.Read(dataReader, binary.LittleEndian, &end)
	}
	return result, begin, end
}

// 读取区域信息
func (self *QQWry) readAreaStrOf(offset int64) string {
	dataReader.Seek(offset, 0)
	b, _ := dataReader.ReadByte()
	if b == REDIRECT_MODE_1 || b == REDIRECT_MODE_2 {
		dataReader.Seek(offset+1, 0)
		AreaOffset := self.read3byteint()
		if AreaOffset != 0 {
			return self.readStrOf(AreaOffset)
		}
	} else {
		return self.readStrOf(offset)
	}
	return "未知区域"
}

// 读指定位置的字符串
func (self *QQWry) readStrOf(offset int64) string {
	dataReader.Seek(offset, 0)
	strbytes := make([]byte, 0) // 字节数组
	b, _ := dataReader.ReadByte()
	for b != 0 {
		strbytes = append(strbytes, b)
		b, _ = dataReader.ReadByte()
	}
	return self.ansiToUtf8(strbytes)
}

// 读取指定位置ip的归属字符串
func (self *QQWry) getIpLocationOf(offset int64) string {
	dataReader.Seek(offset+4, 0)
	var CountryName, AreaName string
	b, _ := dataReader.ReadByte()

	switch b {
	case REDIRECT_MODE_1:
		Countryoffset := self.read3byteint()
		dataReader.Seek(Countryoffset, 0)
		b, _ := dataReader.ReadByte()
		if b == REDIRECT_MODE_2 {
			CountryName = self.readStrOf(self.read3byteint())
			dataReader.Seek(Countryoffset+4, 0)
		} else {
			CountryName = self.readStrOf(Countryoffset)
		}
		AreaName = self.readAreaStrOf(self.dataReaderPosition())

	case REDIRECT_MODE_2:
		CountryName = self.readStrOf(self.read3byteint())
		AreaName = self.readAreaStrOf(offset + 8)

	default:
		CountryName = self.readStrOf(self.dataReaderPosition() - 1)
		AreaName = self.readAreaStrOf(self.dataReaderPosition())

	}
	return CountryName + AreaName
}

// 读取指定ip字符串位置
func (self *QQWry) GetIPLocation(ip uint32) string {
	if self.RecordCount == 0 || ip <= 0 {
		return ""
	}
	var min, max, mid, beginIP, endIP uint32
	max = self.RecordCount - 1
	for min <= max {
		mid = (min + max) / 2
		_, beginIP, endIP = self.getIPOffset(mid, 0)
		if ip == beginIP {
			max = mid
			break
		} else if ip > beginIP {
			min = mid + 1
		} else {
			max = mid - 1
		}
	}
	offset, beginIP, endIP := self.getIPOffset(max, 1)
	if beginIP <= ip && endIP >= ip {
		return self.getIpLocationOf(offset)
	}
	return "***友情提示，未知IP***"
}

// 查找字符串ip
func (self *QQWry) GetIPLocationOfString(ipstr string) string {
	return self.GetIPLocation(self.StrToIP(ipstr))
}

// 初始文件, 只初始化一次
func initFile(fileName string) {
	if isInitReader {
		return
	}
	// init data
	dataStream, err := os.Open(fileName)
	if err != nil {
		panic("打开文件错误")
	}
	defer dataStream.Close()
	fileInfo, err := dataStream.Stat()
	if err != nil {
		panic("获取文件信息错误")
	}
	tempbuffer := make([]byte, fileInfo.Size())
	_, err = dataStream.Read(tempbuffer)
	if err != nil {
		panic("读取错误")
	}
	dataReader = bytes.NewReader(tempbuffer)
}
<commit_msg>因mahonia被克隆的版本函数改变了，所以要修改正确<commit_after>/********************************************************

       纯真IP数据库查询
       by:ying32
       2016/5/13
       QQ: 1444386932
--------------------------------------------------------

       未测试多种带有重定义或未重定义的IP
       如果查询出错，可发邮件给我，告知改进
       或者由您改进，但记得发我一份

       纯真网络：
         http://www.cz88.net/
       纯真IP数据格式：
         http://lumaqq.linuxsir.org/
                     article/qqwry_format_detail.html

********************************************************/

package qqwry

import (
	"bytes"
	//"code.google.com/p/mahonia"
	"encoding/binary"
	"fmt"
	"github.com/axgle/mahonia"
	"net"
	"os"
)

// 重定义模式
const (
	REDIRECT_MODE_1 = 1
	REDIRECT_MODE_2 = 2
)

var (
	// 全局数据，只一次
	dataReader *bytes.Reader
	// 标识"qqwry.dat"文件是否初始
	isInitReader bool
)

type QQWry struct {
	RecordCount uint32
	FirstRecord uint32
}

// 模拟构造函数
func NewQQWry(fileName string) *QQWry {
	// 只初始一次文件
	if !isInitReader {
		initFile(fileName)
		isInitReader = true
	}
	var rf, rl, rc uint32
	binary.Read(dataReader, binary.LittleEndian, &rf)
	binary.Read(dataReader, binary.LittleEndian, &rl)
	rc = (rl - rf) / 7
	return &QQWry{rc, rf}
}

// 获取版本信息
func (self *QQWry) Version() string {
	if self.RecordCount > 0 {
		offset, _, _ := self.getIPOffset(self.RecordCount, 1)
		if offset > 0 {
			return self.getIpLocationOf(offset)
		}
	}
	return "未知区域"
}

// 读取一个3字节并转为整数
func (self *QQWry) read3byteint() int64 {
	temp := make([]byte, 3)
	dataReader.Read(temp)
	return int64(uint32(temp[0]) | uint32(temp[1])<<8 | uint32(temp[2])<<16)
}

// 获取当前Reader位置
func (self *QQWry) dataReaderPosition() int64 {
	return dataReader.Size() - int64(dataReader.Len())
}

// AnsiToUtf8的转换
func (self *QQWry) ansiToUtf8(bytes []byte) string {
	decoder := mahonia.NewDecoder("gbk")
	return decoder.ConvertString(string(bytes))
}

// 字符串ip转整型
func (self *QQWry) StrToIP(ipstr string) uint32 {
	ip := net.ParseIP(ipstr)
	return binary.BigEndian.Uint32(ip.To4())
}

// 整型转ip字符，ipv4
func (self *QQWry) IPToStr(ip uint32) string {
	b := make([]byte, 4)
	binary.BigEndian.PutUint32(b, ip)
	return fmt.Sprintf("%d.%d.%d.%d", b[0], b[1], b[2], b[3])
}

// 获取ip偏移位置
func (self *QQWry) getIPOffset(recNum uint32, flag int) (int64, uint32, uint32) {
	var begin, end uint32
	var result int64 = int64(self.FirstRecord + recNum*7)
	dataReader.Seek(result, 0)
	binary.Read(dataReader, binary.LittleEndian, &begin)
	if flag == 1 {
		result = self.read3byteint()
		dataReader.Seek(result, 0)
		binary.Read(dataReader, binary.LittleEndian, &end)
	}
	return result, begin, end
}

// 读取区域信息
func (self *QQWry) readAreaStrOf(offset int64) string {
	dataReader.Seek(offset, 0)
	b, _ := dataReader.ReadByte()
	if b == REDIRECT_MODE_1 || b == REDIRECT_MODE_2 {
		dataReader.Seek(offset+1, 0)
		AreaOffset := self.read3byteint()
		if AreaOffset != 0 {
			return self.readStrOf(AreaOffset)
		}
	} else {
		return self.readStrOf(offset)
	}
	return "未知区域"
}

// 读指定位置的字符串
func (self *QQWry) readStrOf(offset int64) string {
	dataReader.Seek(offset, 0)
	strbytes := make([]byte, 0) // 字节数组
	b, _ := dataReader.ReadByte()
	for b != 0 {
		strbytes = append(strbytes, b)
		b, _ = dataReader.ReadByte()
	}
	return self.ansiToUtf8(strbytes)
}

// 读取指定位置ip的归属字符串
func (self *QQWry) getIpLocationOf(offset int64) string {
	dataReader.Seek(offset+4, 0)
	var CountryName, AreaName string
	b, _ := dataReader.ReadByte()

	switch b {
	case REDIRECT_MODE_1:
		Countryoffset := self.read3byteint()
		dataReader.Seek(Countryoffset, 0)
		b, _ := dataReader.ReadByte()
		if b == REDIRECT_MODE_2 {
			CountryName = self.readStrOf(self.read3byteint())
			dataReader.Seek(Countryoffset+4, 0)
		} else {
			CountryName = self.readStrOf(Countryoffset)
		}
		AreaName = self.readAreaStrOf(self.dataReaderPosition())

	case REDIRECT_MODE_2:
		CountryName = self.readStrOf(self.read3byteint())
		AreaName = self.readAreaStrOf(offset + 8)

	default:
		CountryName = self.readStrOf(self.dataReaderPosition() - 1)
		AreaName = self.readAreaStrOf(self.dataReaderPosition())

	}
	return CountryName + AreaName
}

// 读取指定ip字符串位置
func (self *QQWry) GetIPLocation(ip uint32) string {
	if self.RecordCount == 0 || ip <= 0 {
		return ""
	}
	var min, max, mid, beginIP, endIP uint32
	max = self.RecordCount - 1
	for min <= max {
		mid = (min + max) / 2
		_, beginIP, endIP = self.getIPOffset(mid, 0)
		if ip == beginIP {
			max = mid
			break
		} else if ip > beginIP {
			min = mid + 1
		} else {
			max = mid - 1
		}
	}
	offset, beginIP, endIP := self.getIPOffset(max, 1)
	if beginIP <= ip && endIP >= ip {
		return self.getIpLocationOf(offset)
	}
	return "***友情提示，未知IP***"
}

// 查找字符串ip
func (self *QQWry) GetIPLocationOfString(ipstr string) string {
	return self.GetIPLocation(self.StrToIP(ipstr))
}

// 初始文件, 只初始化一次
func initFile(fileName string) {
	if isInitReader {
		return
	}
	// init data
	dataStream, err := os.Open(fileName)
	if err != nil {
		panic("打开文件错误")
	}
	defer dataStream.Close()
	fileInfo, err := dataStream.Stat()
	if err != nil {
		panic("获取文件信息错误")
	}
	tempbuffer := make([]byte, fileInfo.Size())
	_, err = dataStream.Read(tempbuffer)
	if err != nil {
		panic("读取错误")
	}
	dataReader = bytes.NewReader(tempbuffer)
}
<|endoftext|>


<commit_before>package jwt

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

	jwt "github.com/dgrijalva/jwt-go"
	"github.com/gkarlik/quark-go/logger"
	"golang.org/x/net/context"
)

// Credentials represents user credentials (username and password).
type Credentials struct {
	Username string `json:"username"` // user name
	Password string `json:"password"` // password
}

// Claims represents jwt claims.
type Claims struct {
	Username   string                 `json:"username"`   // user name
	Properties map[string]interface{} `json:"properties"` // additional jwt claims properties

	jwt.StandardClaims // standard jwt claims properties
}

const componentName = "JwtAuthenticationMiddleware"

// AuthenticationMiddleware represents HTTP middleware responsible for authentication (JWT based).
type AuthenticationMiddleware struct {
	Options Options // authentication middleware options
}

// AuthenticationFunc is a function used to authenticate user. Function receives user credentials and should return claims or an error.
type AuthenticationFunc func(credentials Credentials) (Claims, error)

// NewAuthenticationMiddleware creates instance of authentication middleware with options passed as argument.
// AuthenticationFunc and Secret options are required.
// Default context key value used to store jwt claims in request context is "Claims".
func NewAuthenticationMiddleware(opts ...Option) *AuthenticationMiddleware {
	am := &AuthenticationMiddleware{
		Options: Options{
			ContextKey: "Claims",
		},
	}

	for _, opt := range opts {
		opt(&am.Options)
	}

	if am.Options.Authenticate == nil {
		panic(fmt.Sprintf("[%s]: Cannot create instance - authentication function must be set!", componentName))
	}

	if am.Options.Secret == "" {
		panic(fmt.Sprintf("[%s]: Cannot create instance - secret must be set!", componentName))
	}

	return am
}

// Authenticate validates token using jwt specification. It parses token from 'Authorization' header which must be in form "bearer token".
func (am AuthenticationMiddleware) Authenticate(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// parse token from Authorization header
		ah := r.Header.Get("Authorization")
		if ah == "" {
			logger.Log().ErrorWithFields(logger.Fields{"component": componentName}, "No authorization header")
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		s := strings.Split(ah, " ")

		if len(s) != 2 || strings.ToUpper(s[0]) != "BEARER" {
			logger.Log().ErrorWithFields(logger.Fields{"component": componentName}, "Incorrect authorization header")
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		tokenString := s[1]
		if tokenString == "" {
			logger.Log().ErrorWithFields(logger.Fields{"component": componentName}, "TokenString is empty")
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(t *jwt.Token) (interface{}, error) {
			if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
				return nil, fmt.Errorf("[%s]: Unexpected signing method", componentName)
			}

			return []byte(am.Options.Secret), nil
		})

		if err != nil {
			logger.Log().ErrorWithFields(logger.Fields{
				"error":       err,
				"tokenString": tokenString,
				"component":   componentName,
			}, "Error parsing token string")
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		// get tokens claims and pass it into the original request
		if claims, ok := token.Claims.(*Claims); ok && token.Valid {
			ctx := context.WithValue(r.Context(), am.Options.ContextKey, *claims)
			next.ServeHTTP(w, r.WithContext(ctx))
		} else {
			logger.Log().ErrorWithFields(logger.Fields{"component": componentName}, "Token is invalid")
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}
	})
}

// GenerateToken generates token using jwt specification, only HTTP POST method is allowed.
func (am AuthenticationMiddleware) GenerateToken(w http.ResponseWriter, r *http.Request) {
	// check if method is POST
	if r.Method != http.MethodPost {
		logger.Log().ErrorWithFields(logger.Fields{
			"method":    r.Method,
			"component": componentName,
		}, "Incorrect http method")
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// load credentials from request body
	var credentials Credentials
	decoder := json.NewDecoder(r.Body)
	if err := decoder.Decode(&credentials); err != nil {
		logger.Log().ErrorWithFields(logger.Fields{
			"error":       err,
			"credentials": credentials,
			"component":   componentName,
		}, "Cannot decode credentials")
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	// authenticate user using external function
	claims, err := am.Options.Authenticate(credentials)
	if err != nil {
		logger.Log().ErrorWithFields(logger.Fields{
			"error":     err,
			"component": componentName,
		}, "Cannot authenticate user")
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	// create token with claims
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	tokenString, err := token.SignedString([]byte(am.Options.Secret))
	if err != nil {
		logger.Log().ErrorWithFields(logger.Fields{
			"error":     err,
			"component": componentName,
		}, "Cannot sign token")
		http.Error(w, "Internal server error", http.StatusInternalServerError)
		return
	}

	logger.Log().InfoWithFields(logger.Fields{
		"token":     tokenString,
		"component": componentName,
	}, "Token generated - sending to client")

	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, "{ \"token\": %q }", tokenString)
}
<commit_msg>Credentials extension to support additional properties.<commit_after>package jwt

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

	jwt "github.com/dgrijalva/jwt-go"
	"github.com/gkarlik/quark-go/logger"
	"golang.org/x/net/context"
)

// Credentials represents user credentials (username and password).
type Credentials struct {
	Username   string            `json:"Username"`             // user name
	Password   string            `json:"Password"`             // password
	Properties map[string]string `json:"Properties,omitempty"` // additional properties
}

// Claims represents jwt claims.
type Claims struct {
	Username   string                 `json:"username"`   // user name
	Properties map[string]interface{} `json:"properties"` // additional jwt claims properties

	jwt.StandardClaims // standard jwt claims properties
}

const componentName = "JwtAuthenticationMiddleware"

// AuthenticationMiddleware represents HTTP middleware responsible for authentication (JWT based).
type AuthenticationMiddleware struct {
	Options Options // authentication middleware options
}

// AuthenticationFunc is a function used to authenticate user. Function receives user credentials and should return claims or an error.
type AuthenticationFunc func(credentials Credentials) (Claims, error)

// NewAuthenticationMiddleware creates instance of authentication middleware with options passed as argument.
// AuthenticationFunc and Secret options are required.
// Default context key value used to store jwt claims in request context is "Claims".
func NewAuthenticationMiddleware(opts ...Option) *AuthenticationMiddleware {
	am := &AuthenticationMiddleware{
		Options: Options{
			ContextKey: "Claims",
		},
	}

	for _, opt := range opts {
		opt(&am.Options)
	}

	if am.Options.Authenticate == nil {
		panic(fmt.Sprintf("[%s]: Cannot create instance - authentication function must be set!", componentName))
	}

	if am.Options.Secret == "" {
		panic(fmt.Sprintf("[%s]: Cannot create instance - secret must be set!", componentName))
	}

	return am
}

// Authenticate validates token using jwt specification. It parses token from 'Authorization' header which must be in form "bearer token".
func (am AuthenticationMiddleware) Authenticate(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// parse token from Authorization header
		ah := r.Header.Get("Authorization")
		if ah == "" {
			logger.Log().ErrorWithFields(logger.Fields{"component": componentName}, "No authorization header")
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		s := strings.Split(ah, " ")

		if len(s) != 2 || strings.ToUpper(s[0]) != "BEARER" {
			logger.Log().ErrorWithFields(logger.Fields{"component": componentName}, "Incorrect authorization header")
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		tokenString := s[1]
		if tokenString == "" {
			logger.Log().ErrorWithFields(logger.Fields{"component": componentName}, "TokenString is empty")
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(t *jwt.Token) (interface{}, error) {
			if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
				return nil, fmt.Errorf("[%s]: Unexpected signing method", componentName)
			}

			return []byte(am.Options.Secret), nil
		})

		if err != nil {
			logger.Log().ErrorWithFields(logger.Fields{
				"error":       err,
				"tokenString": tokenString,
				"component":   componentName,
			}, "Error parsing token string")
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		// get tokens claims and pass it into the original request
		if claims, ok := token.Claims.(*Claims); ok && token.Valid {
			ctx := context.WithValue(r.Context(), am.Options.ContextKey, *claims)
			next.ServeHTTP(w, r.WithContext(ctx))
		} else {
			logger.Log().ErrorWithFields(logger.Fields{"component": componentName}, "Token is invalid")
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}
	})
}

// GenerateToken generates token using jwt specification, only HTTP POST method is allowed.
func (am AuthenticationMiddleware) GenerateToken(w http.ResponseWriter, r *http.Request) {
	// check if method is POST
	if r.Method != http.MethodPost {
		logger.Log().ErrorWithFields(logger.Fields{
			"method":    r.Method,
			"component": componentName,
		}, "Incorrect http method")
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// load credentials from request body
	var credentials Credentials
	decoder := json.NewDecoder(r.Body)
	if err := decoder.Decode(&credentials); err != nil {
		logger.Log().ErrorWithFields(logger.Fields{
			"error":       err,
			"credentials": credentials,
			"component":   componentName,
		}, "Cannot decode credentials")
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	// authenticate user using external function
	claims, err := am.Options.Authenticate(credentials)
	if err != nil {
		logger.Log().ErrorWithFields(logger.Fields{
			"error":     err,
			"component": componentName,
		}, "Cannot authenticate user")
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	// create token with claims
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	tokenString, err := token.SignedString([]byte(am.Options.Secret))
	if err != nil {
		logger.Log().ErrorWithFields(logger.Fields{
			"error":     err,
			"component": componentName,
		}, "Cannot sign token")
		http.Error(w, "Internal server error", http.StatusInternalServerError)
		return
	}

	logger.Log().InfoWithFields(logger.Fields{
		"token":     tokenString,
		"component": componentName,
	}, "Token generated - sending to client")

	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, "{ \"token\": %q }", tokenString)
}
<|endoftext|>


<commit_before>package tdb

import (
	"bytes"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/Knetic/govaluate"
	"github.com/tecbot/gorocksdb"
	"gopkg.in/vmihailenco/msgpack.v2"
)

type query struct {
	table    string
	fields   []string
	filter   *govaluate.EvaluableExpression
	from     time.Time
	to       time.Time
	onValues func(key map[string]interface{}, field string, vals []float64)
}

type QueryStats struct {
	Scanned      int64
	FilterPass   int64
	FilterReject int64
	ReadValue    int64
	DataValid    int64
	InTimeRange  int64
}

func (db *DB) runQuery(q *query) (*QueryStats, error) {
	stats := &QueryStats{}

	if q.from.IsZero() {
		return stats, fmt.Errorf("Please specify a from")
	}
	if len(q.fields) == 0 {
		return stats, fmt.Errorf("Please specify at least one field")
	}
	t := db.getTable(q.table)
	if t == nil {
		return stats, fmt.Errorf("Unknown table %v", q.table)
	}
	fields := make([][]byte, 0, len(q.fields))
	for _, field := range q.fields {
		fieldBytes, err := msgpack.Marshal(strings.ToLower(field))
		if err != nil {
			return stats, fmt.Errorf("Unable to marshal field: %v", err)
		}
		fields = append(fields, fieldBytes)
	}
	sort.Sort(lexicographical(fields))
	if q.to.IsZero() {
		q.to = t.clock.Now()
	}
	q.from = roundTime(q.from, t.resolution)
	if q.to.IsZero() {
		q.to = t.clock.Now()
	}
	q.to = roundTime(q.to, t.resolution)
	numPeriods := int(q.to.Sub(q.from)/t.resolution) + 1
	log.Tracef("Query will return %d periods", numPeriods)

	ro := gorocksdb.NewDefaultReadOptions()
	// Go ahead and fill the cache
	ro.SetFillCache(true)
	it := t.archiveByKey.NewIterator(ro)
	defer it.Close()

	for _, fieldBytes := range fields {
		for it.Seek(fieldBytes); it.ValidForPrefix(fieldBytes); it.Next() {
			stats.Scanned++
			k := it.Key()
			kr := bytes.NewReader(k.Data())
			dec := msgpack.NewDecoder(kr)
			storedField, err := dec.DecodeString()
			if err != nil {
				k.Free()
				return stats, fmt.Errorf("Unable to decode field: %v", err)
			}
			key := make(map[string]interface{})
			err = dec.Decode(&key)
			if err != nil {
				k.Free()
				return stats, fmt.Errorf("Unable to decode key: %v", err)
			}
			k.Free()

			if q.filter != nil {
				include, err := q.filter.Evaluate(key)
				if err != nil {
					return stats, fmt.Errorf("Unable to apply filter: %v", err)
				}
				inc, ok := include.(bool)
				if !ok {
					return stats, fmt.Errorf("Filter expression returned something other than a boolean: %v", include)
				}
				if !inc {
					stats.FilterReject++
					continue
				}
				stats.FilterPass++
			}

			v := it.Value()
			stats.ReadValue++
			seq := sequence(v.Data())
			vals := make([]float64, numPeriods)
			if seq.isValid() {
				stats.DataValid++
				seqStart := seq.start()
				if log.IsTraceEnabled() {
					log.Tracef("Sequence starts at %v and has %d periods", seqStart.In(time.UTC), seq.numPeriods())
				}
				includeKey := false
				if !seqStart.Before(q.from) {
					to := q.to
					if to.After(seqStart) {
						to = seqStart
					}
					startOffset := int(seqStart.Sub(to) / t.resolution)
					log.Tracef("Start offset %d", startOffset)
					copyPeriods := seq.numPeriods()
					for i := 0; i+startOffset < copyPeriods && i < numPeriods; i++ {
						includeKey = true
						val := seq.valueAt(i + startOffset)
						log.Tracef("Grabbing value %f", val)
						vals[i] = val
					}
				}
				if includeKey {
					stats.InTimeRange++
					q.onValues(key, storedField, vals)
				}
			}
			v.Free()
		}
	}
	return stats, nil
}

type lexicographical [][]byte

func (a lexicographical) Len() int           { return len(a) }
func (a lexicographical) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a lexicographical) Less(i, j int) bool { return bytes.Compare(a[i], a[j]) < 0 }

func keysEqual(a map[string]interface{}, b map[string]interface{}) bool {
	if len(a) != len(b) {
		return false
	}
	for k, v := range a {
		if b[k] != v {
			return false
		}
	}
	return true
}
<commit_msg>Added query runtime to stats<commit_after>package tdb

import (
	"bytes"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/Knetic/govaluate"
	"github.com/tecbot/gorocksdb"
	"gopkg.in/vmihailenco/msgpack.v2"
)

type query struct {
	table    string
	fields   []string
	filter   *govaluate.EvaluableExpression
	from     time.Time
	to       time.Time
	onValues func(key map[string]interface{}, field string, vals []float64)
}

type QueryStats struct {
	Scanned      int64
	FilterPass   int64
	FilterReject int64
	ReadValue    int64
	DataValid    int64
	InTimeRange  int64
	Runtime      time.Duration
}

func (db *DB) runQuery(q *query) (*QueryStats, error) {
	start := time.Now()
	stats := &QueryStats{}

	if q.from.IsZero() {
		return stats, fmt.Errorf("Please specify a from")
	}
	if len(q.fields) == 0 {
		return stats, fmt.Errorf("Please specify at least one field")
	}
	t := db.getTable(q.table)
	if t == nil {
		return stats, fmt.Errorf("Unknown table %v", q.table)
	}
	fields := make([][]byte, 0, len(q.fields))
	for _, field := range q.fields {
		fieldBytes, err := msgpack.Marshal(strings.ToLower(field))
		if err != nil {
			return stats, fmt.Errorf("Unable to marshal field: %v", err)
		}
		fields = append(fields, fieldBytes)
	}
	sort.Sort(lexicographical(fields))
	if q.to.IsZero() {
		q.to = t.clock.Now()
	}
	q.from = roundTime(q.from, t.resolution)
	if q.to.IsZero() {
		q.to = t.clock.Now()
	}
	q.to = roundTime(q.to, t.resolution)
	numPeriods := int(q.to.Sub(q.from)/t.resolution) + 1
	log.Tracef("Query will return %d periods", numPeriods)

	ro := gorocksdb.NewDefaultReadOptions()
	// Go ahead and fill the cache
	ro.SetFillCache(true)
	it := t.archiveByKey.NewIterator(ro)
	defer it.Close()

	for _, fieldBytes := range fields {
		for it.Seek(fieldBytes); it.ValidForPrefix(fieldBytes); it.Next() {
			stats.Scanned++
			k := it.Key()
			kr := bytes.NewReader(k.Data())
			dec := msgpack.NewDecoder(kr)
			storedField, err := dec.DecodeString()
			if err != nil {
				k.Free()
				return stats, fmt.Errorf("Unable to decode field: %v", err)
			}
			key := make(map[string]interface{})
			err = dec.Decode(&key)
			if err != nil {
				k.Free()
				return stats, fmt.Errorf("Unable to decode key: %v", err)
			}
			k.Free()

			if q.filter != nil {
				include, err := q.filter.Evaluate(key)
				if err != nil {
					return stats, fmt.Errorf("Unable to apply filter: %v", err)
				}
				inc, ok := include.(bool)
				if !ok {
					return stats, fmt.Errorf("Filter expression returned something other than a boolean: %v", include)
				}
				if !inc {
					stats.FilterReject++
					continue
				}
				stats.FilterPass++
			}

			v := it.Value()
			stats.ReadValue++
			seq := sequence(v.Data())
			vals := make([]float64, numPeriods)
			if seq.isValid() {
				stats.DataValid++
				seqStart := seq.start()
				if log.IsTraceEnabled() {
					log.Tracef("Sequence starts at %v and has %d periods", seqStart.In(time.UTC), seq.numPeriods())
				}
				includeKey := false
				if !seqStart.Before(q.from) {
					to := q.to
					if to.After(seqStart) {
						to = seqStart
					}
					startOffset := int(seqStart.Sub(to) / t.resolution)
					log.Tracef("Start offset %d", startOffset)
					copyPeriods := seq.numPeriods()
					for i := 0; i+startOffset < copyPeriods && i < numPeriods; i++ {
						includeKey = true
						val := seq.valueAt(i + startOffset)
						log.Tracef("Grabbing value %f", val)
						vals[i] = val
					}
				}
				if includeKey {
					stats.InTimeRange++
					q.onValues(key, storedField, vals)
				}
			}
			v.Free()
		}
	}

	stats.Runtime = time.Now().Sub(start)
	return stats, nil
}

type lexicographical [][]byte

func (a lexicographical) Len() int           { return len(a) }
func (a lexicographical) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a lexicographical) Less(i, j int) bool { return bytes.Compare(a[i], a[j]) < 0 }

func keysEqual(a map[string]interface{}, b map[string]interface{}) bool {
	if len(a) != len(b) {
		return false
	}
	for k, v := range a {
		if b[k] != v {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>package gorethink

import (
	"fmt"
	"strconv"
	"strings"

	p "github.com/dancannon/gorethink/ql2"
)

type OptArgs interface {
	toMap() map[string]interface{}
}
type termsList []Term
type termsObj map[string]Term
type Term struct {
	name     string
	rootTerm bool
	termType p.Term_TermType
	data     interface{}
	args     []Term
	optArgs  map[string]Term
}

// build takes the query tree and prepares it to be sent as a JSON
// expression
func (t Term) build() interface{} {
	switch t.termType {
	case p.Term_DATUM:
		return t.data
	case p.Term_BINARY:
		if len(t.args) == 0 {
			return map[string]interface{}{
				"$reql_type$": "BINARY",
				"data":        t.data,
			}
		}
	}

	args := []interface{}{}
	optArgs := map[string]interface{}{}

	for _, v := range t.args {
		args = append(args, v.build())
	}

	for k, v := range t.optArgs {
		optArgs[k] = v.build()
	}

	return []interface{}{t.termType, args, optArgs}
}

// String returns a string representation of the query tree
func (t Term) String() string {
	switch t.termType {
	case p.Term_MAKE_ARRAY:
		return fmt.Sprintf("[%s]", strings.Join(argsToStringSlice(t.args), ", "))
	case p.Term_MAKE_OBJ:
		return fmt.Sprintf("{%s}", strings.Join(optArgsToStringSlice(t.optArgs), ", "))
	case p.Term_FUNC:
		// Get string representation of each argument
		args := []string{}
		for _, v := range t.args[0].args {
			args = append(args, fmt.Sprintf("var_%d", v.data))
		}

		return fmt.Sprintf("func(%s r.Term) r.Term { return %s }",
			strings.Join(args, ", "),
			t.args[1].String(),
		)
	case p.Term_VAR:
		return fmt.Sprintf("var_%s", t.args[0])
	case p.Term_IMPLICIT_VAR:
		return "r.Row"
	case p.Term_DATUM:
		switch v := t.data.(type) {
		case string:
			return strconv.Quote(v)
		default:
			return fmt.Sprintf("%v", v)
		}
	case p.Term_BINARY:
		if len(t.args) == 0 {
			return fmt.Sprintf("r.binary(<data>)")
		}
	}

	if t.rootTerm {
		return fmt.Sprintf("r.%s(%s)", t.name, strings.Join(allArgsToStringSlice(t.args, t.optArgs), ", "))
	} else {
		return fmt.Sprintf("%s.%s(%s)", t.args[0].String(), t.name, strings.Join(allArgsToStringSlice(t.args[1:], t.optArgs), ", "))
	}
}

type WriteResponse struct {
	Errors        int
	Created       int
	Inserted      int
	Updated       int
	Unchanged     int
	Replaced      int
	Deleted       int
	GeneratedKeys []string    `gorethink:"generated_keys"`
	FirstError    string      `gorethink:"first_error"` // populated if Errors > 0
	NewValue      interface{} `gorethink:"new_val"`
	OldValue      interface{} `gorethink:"old_val"`
}

type RunOpts struct {
	Db          interface{} `gorethink:"db,omitempty"`
	Profile     interface{} `gorethink:"profile,omitempty"`
	UseOutdated interface{} `gorethink:"use_outdated,omitempty"`
	NoReply     interface{} `gorethink:"noreply,omitempty"`
	TimeFormat  interface{} `gorethink:"time_format,omitempty"`

	// Unsupported options

	BatchConf interface{} `gorethink:"batch_conf,omitempty"`
}

func (o *RunOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// Run runs a query using the given connection.
//
//	rows, err := query.Run(sess)
//	if err != nil {
//		// error
//	}
//	for rows.Next() {
//		doc := MyDocumentType{}
//		err := r.Scan(&doc)
//		    // Do something with row
//	}
func (t Term) Run(s *Session, optArgs ...RunOpts) (*Cursor, error) {
	opts := map[string]interface{}{}
	if len(optArgs) >= 1 {
		opts = optArgs[0].toMap()
	}
	return s.startQuery(t, opts)
}

// RunWrite runs a query using the given connection but unlike Run automatically
// scans the result into a variable of type WriteResponse. This function should be used
// if you are running a write query (such as Insert,  Update, TableCreate, etc...)
//
// Optional arguments :
// "db", "use_outdated" (defaults to false), "noreply" (defaults to false) and "time_format".
//
//	res, err := r.Db("database").Table("table").Insert(doc).RunWrite(sess, r.RunOpts{
//		NoReply: true,
//	})
func (t Term) RunWrite(s *Session, optArgs ...RunOpts) (WriteResponse, error) {
	var response WriteResponse
	res, err := t.Run(s, optArgs...)
	if err == nil {
		err = res.One(&response)
	}
	return response, err
}

// Exec runs the query but does not return the result.
func (t Term) Exec(s *Session, optArgs ...RunOpts) error {
	res, err := t.Run(s, optArgs...)
	if err != nil {
		return err
	}
	err = res.Close()
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Added the binary_format option<commit_after>package gorethink

import (
	"fmt"
	"strconv"
	"strings"

	p "github.com/dancannon/gorethink/ql2"
)

type OptArgs interface {
	toMap() map[string]interface{}
}
type termsList []Term
type termsObj map[string]Term
type Term struct {
	name     string
	rootTerm bool
	termType p.Term_TermType
	data     interface{}
	args     []Term
	optArgs  map[string]Term
}

// build takes the query tree and prepares it to be sent as a JSON
// expression
func (t Term) build() interface{} {
	switch t.termType {
	case p.Term_DATUM:
		return t.data
	case p.Term_BINARY:
		if len(t.args) == 0 {
			return map[string]interface{}{
				"$reql_type$": "BINARY",
				"data":        t.data,
			}
		}
	}

	args := []interface{}{}
	optArgs := map[string]interface{}{}

	for _, v := range t.args {
		args = append(args, v.build())
	}

	for k, v := range t.optArgs {
		optArgs[k] = v.build()
	}

	return []interface{}{t.termType, args, optArgs}
}

// String returns a string representation of the query tree
func (t Term) String() string {
	switch t.termType {
	case p.Term_MAKE_ARRAY:
		return fmt.Sprintf("[%s]", strings.Join(argsToStringSlice(t.args), ", "))
	case p.Term_MAKE_OBJ:
		return fmt.Sprintf("{%s}", strings.Join(optArgsToStringSlice(t.optArgs), ", "))
	case p.Term_FUNC:
		// Get string representation of each argument
		args := []string{}
		for _, v := range t.args[0].args {
			args = append(args, fmt.Sprintf("var_%d", v.data))
		}

		return fmt.Sprintf("func(%s r.Term) r.Term { return %s }",
			strings.Join(args, ", "),
			t.args[1].String(),
		)
	case p.Term_VAR:
		return fmt.Sprintf("var_%s", t.args[0])
	case p.Term_IMPLICIT_VAR:
		return "r.Row"
	case p.Term_DATUM:
		switch v := t.data.(type) {
		case string:
			return strconv.Quote(v)
		default:
			return fmt.Sprintf("%v", v)
		}
	case p.Term_BINARY:
		if len(t.args) == 0 {
			return fmt.Sprintf("r.binary(<data>)")
		}
	}

	if t.rootTerm {
		return fmt.Sprintf("r.%s(%s)", t.name, strings.Join(allArgsToStringSlice(t.args, t.optArgs), ", "))
	} else {
		return fmt.Sprintf("%s.%s(%s)", t.args[0].String(), t.name, strings.Join(allArgsToStringSlice(t.args[1:], t.optArgs), ", "))
	}
}

type WriteResponse struct {
	Errors        int
	Created       int
	Inserted      int
	Updated       int
	Unchanged     int
	Replaced      int
	Deleted       int
	GeneratedKeys []string    `gorethink:"generated_keys"`
	FirstError    string      `gorethink:"first_error"` // populated if Errors > 0
	NewValue      interface{} `gorethink:"new_val"`
	OldValue      interface{} `gorethink:"old_val"`
}

type RunOpts struct {
	Db           interface{} `gorethink:"db,omitempty"`
	Profile      interface{} `gorethink:"profile,omitempty"`
	UseOutdated  interface{} `gorethink:"use_outdated,omitempty"`
	NoReply      interface{} `gorethink:"noreply,omitempty"`
	TimeFormat   interface{} `gorethink:"time_format,omitempty"`
	GroupFormat  interface{} `gorethink:"group_format,omitempty"`
	BinaryFormat interface{} `gorethink:"binary_format,omitempty"`

	// Unsupported options

	BatchConf interface{} `gorethink:"batch_conf,omitempty"`
}

func (o *RunOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// Run runs a query using the given connection.
//
//	rows, err := query.Run(sess)
//	if err != nil {
//		// error
//	}
//	for rows.Next() {
//		doc := MyDocumentType{}
//		err := r.Scan(&doc)
//		    // Do something with row
//	}
func (t Term) Run(s *Session, optArgs ...RunOpts) (*Cursor, error) {
	opts := map[string]interface{}{}
	if len(optArgs) >= 1 {
		opts = optArgs[0].toMap()
	}
	return s.startQuery(t, opts)
}

// RunWrite runs a query using the given connection but unlike Run automatically
// scans the result into a variable of type WriteResponse. This function should be used
// if you are running a write query (such as Insert,  Update, TableCreate, etc...)
//
//	res, err := r.Db("database").Table("table").Insert(doc).RunWrite(sess, r.RunOpts{
//		NoReply: true,
//	})
func (t Term) RunWrite(s *Session, optArgs ...RunOpts) (WriteResponse, error) {
	var response WriteResponse
	res, err := t.Run(s, optArgs...)
	if err == nil {
		err = res.One(&response)
	}
	return response, err
}

// Exec runs the query but does not return the result.
func (t Term) Exec(s *Session, optArgs ...RunOpts) error {
	res, err := t.Run(s, optArgs...)
	if err != nil {
		return err
	}
	err = res.Close()
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package dispatch
/*
 *  Filename:    queue.go
 *  Package:     dispatch
 *  Author:      Bryan Matsuo <bmatsuo@soe.ucsc.edu>
 *  Created:     Wed Jul  6 17:30:20 PDT 2011
 *  Description:
 */
import (
    "container/list"
)

//  A Queue is a queue for RegisteredTasks, used by a Dispatch.
type Queue interface {
    Enqueue(task RegisteredTask)  // Insert a DispatchTask
    Dequeue() RegisteredTask      // Remove the next task.
    Len() int                     // Number of items to be processed.
}

//  A simple linked-list FIFO satisfying the Queue interface.
type FIFO struct {
    waiting      *list.List  // A list with values of type func(int)
}
//  Create a new FIFO.
func NewFIFO() *FIFO {
    var q = new(FIFO)
    q.waiting = list.New()
    return q
}

//  See Queue.
func (dq *FIFO) Len() int {
    return dq.waiting.Len()
}
//  See Queue.
func (dq *FIFO) Enqueue(task RegisteredTask) {
    dq.waiting.PushBack(task)
}
//  See Queue.
func (dq *FIFO) Dequeue() RegisteredTask {
    var taskelm = dq.waiting.Front()
    dq.waiting.Remove(taskelm)
    return taskelm.Value.(RegisteredTask)
}
<commit_msg>Add a LIFO Queue object.<commit_after>package dispatch
/*
 *  Filename:    queue.go
 *  Package:     dispatch
 *  Author:      Bryan Matsuo <bmatsuo@soe.ucsc.edu>
 *  Created:     Wed Jul  6 17:30:20 PDT 2011
 *  Description:
 */
import (
    "container/list"
)

//  A Queue is a queue for RegisteredTasks, used by a Dispatch.
type Queue interface {
    Enqueue(task RegisteredTask)  // Insert a DispatchTask
    Dequeue() RegisteredTask      // Remove the next task.
    Len() int                     // Number of items to be processed.
}

//  A simple linked-list First In First Out (FIFO) Queue.
type FIFO struct {
    waiting      *list.List  // A list with values of type func(int)
}
//  Create a new FIFO.
func NewFIFO() *FIFO {
    var q = new(FIFO)
    q.waiting = list.New()
    return q
}

//  See Queue.
func (dq *FIFO) Len() int {
    return dq.waiting.Len()
}
//  See Queue.
func (dq *FIFO) Enqueue(task RegisteredTask) {
    dq.waiting.PushBack(task)
}
//  See Queue.
func (dq *FIFO) Dequeue() RegisteredTask {
    var taskelm = dq.waiting.Front()
    dq.waiting.Remove(taskelm)
    return taskelm.Value.(RegisteredTask)
}

//  A simple linked-list Last In First Out (LIFO) Queue.
type LIFO struct {
    waiting      *list.List  // A list with values of type func(int)
}
//  Create a new LIFO.
func NewLIFO() *LIFO {
    var q = new(LIFO)
    q.waiting = list.New()
    return q
}

//  See Queue.
func (dq *LIFO) Len() int {
    return dq.waiting.Len()
}
//  See Queue.
func (dq *LIFO) Enqueue(task RegisteredTask) {
    dq.waiting.PushFront(task)
}
//  See Queue.
func (dq *LIFO) Dequeue() RegisteredTask {
    var taskelm = dq.waiting.Front()
    dq.waiting.Remove(taskelm)
    return taskelm.Value.(RegisteredTask)
}
<|endoftext|>


<commit_before>// Rate limited io
package ratio

import (
	"io"
	"time"
)

type values struct {
	n   int
	err error
}

type operation struct {
	p       []byte
	values  chan values
	written int
}

type rateLimiter struct {
	stop      chan bool
	worker    chan *operation
	action    func(p []byte) (int, error)
	limit     int
	remaining int
	written   int
	op        *operation
}

func (rl *rateLimiter) reset() {
	rl.remaining = rl.limit
}

func (rl *rateLimiter) run(duration time.Duration) {
	for {
		select {
		case <-rl.stop:
			rl.close()
			return
		case <-time.Tick(duration):
			rl.reset()
			rl.write()
		case op := <-rl.worker:
			rl.op = op
			rl.write()
		}
	}
}

func (rl *rateLimiter) write() {
	if rl.op == nil {
		return
	}

	p := rl.op.p
	if rl.remaining < len(rl.op.p) {
		p = rl.op.p[:rl.remaining]
	}
	rl.remaining -= len(p)

	n, err := rl.action(p)
	rl.op.p = rl.op.p[n:]
	rl.op.written += n

	if len(rl.op.p) == 0 || err != nil {
		rl.op.values <- values{rl.op.written, err}
		rl.op = nil
	}
}

func (rl *rateLimiter) record(p []byte) (int, error) {
	op := &operation{
		p:      p,
		values: make(chan values, 1),
	}
	select {
	case <-rl.stop:
		return 0, io.EOF
	case rl.worker <- op:
		val := <-op.values
		return val.n, val.err
	}
}

func (rl *rateLimiter) close() {
	if rl.op != nil {
		rl.op.values <- values{rl.written, io.EOF}
	}
}

// Returns a rate limited io.Writer, allowing to write up to size bytes per duration.
func RateLimitedWriter(w io.Writer, size int, duration time.Duration) io.WriteCloser {
	rl := &rateLimiter{
		action: func(p []byte) (int, error) {
			return w.Write(p)
		},
		limit:     size,
		remaining: size,
		stop:      make(chan bool),
		worker:    make(chan *operation),
	}
	go rl.run(duration)
	return rl
}

// Returns a rate limited io.Reader, allowing to read up to size bytes per duration.
func RateLimitedReader(r io.Reader, size int, duration time.Duration) io.ReadCloser {
	rl := &rateLimiter{
		action: func(p []byte) (int, error) {
			return r.Read(p)
		},
		limit:     size,
		remaining: size,
		stop:      make(chan bool),
		worker:    make(chan *operation),
	}
	go rl.run(duration)
	return rl
}

func (rl *rateLimiter) Read(p []byte) (int, error) {
	return rl.record(p)
}

func (rl *rateLimiter) Write(p []byte) (int, error) {
	return rl.record(p)
}

func (rl *rateLimiter) Close() error {
	select {
	case <-rl.stop:
		return io.EOF
	default:
	}
	close(rl.stop)
	return nil
}
<commit_msg>don't EOF on Close<commit_after>// Rate limited io
package ratio

import (
	"io"
	"time"
)

type values struct {
	n   int
	err error
}

type operation struct {
	p       []byte
	values  chan values
	written int
}

type rateLimiter struct {
	stop      chan bool
	worker    chan *operation
	action    func(p []byte) (int, error)
	limit     int
	remaining int
	written   int
	op        *operation
}

func (rl *rateLimiter) reset() {
	rl.remaining = rl.limit
}

func (rl *rateLimiter) run(duration time.Duration) {
	for {
		select {
		case <-rl.stop:
			rl.close()
			return
		case <-time.Tick(duration):
			rl.reset()
			rl.write()
		case op := <-rl.worker:
			rl.op = op
			rl.write()
		}
	}
}

func (rl *rateLimiter) write() {
	if rl.op == nil {
		return
	}

	p := rl.op.p
	if rl.remaining < len(rl.op.p) {
		p = rl.op.p[:rl.remaining]
	}
	rl.remaining -= len(p)

	n, err := rl.action(p)
	rl.op.p = rl.op.p[n:]
	rl.op.written += n

	if len(rl.op.p) == 0 || err != nil {
		rl.op.values <- values{rl.op.written, err}
		rl.op = nil
	}
}

func (rl *rateLimiter) record(p []byte) (int, error) {
	op := &operation{
		p:      p,
		values: make(chan values, 1),
	}
	select {
	case <-rl.stop:
		return 0, io.EOF
	case rl.worker <- op:
		val := <-op.values
		return val.n, val.err
	}
}

func (rl *rateLimiter) close() {
	if rl.op != nil {
		rl.op.values <- values{rl.op.written, io.EOF}
	}
}

// Returns a rate limited io.Writer, allowing to write up to size bytes per duration.
func RateLimitedWriter(w io.Writer, size int, duration time.Duration) io.WriteCloser {
	rl := &rateLimiter{
		action: func(p []byte) (int, error) {
			return w.Write(p)
		},
		limit:     size,
		remaining: size,
		stop:      make(chan bool),
		worker:    make(chan *operation),
	}
	go rl.run(duration)
	return rl
}

// Returns a rate limited io.Reader, allowing to read up to size bytes per duration.
func RateLimitedReader(r io.Reader, size int, duration time.Duration) io.ReadCloser {
	rl := &rateLimiter{
		action: func(p []byte) (int, error) {
			return r.Read(p)
		},
		limit:     size,
		remaining: size,
		stop:      make(chan bool),
		worker:    make(chan *operation),
	}
	go rl.run(duration)
	return rl
}

func (rl *rateLimiter) Read(p []byte) (int, error) {
	return rl.record(p)
}

func (rl *rateLimiter) Write(p []byte) (int, error) {
	return rl.record(p)
}

func (rl *rateLimiter) Close() error {
	select {
	case <-rl.stop:
		return nil
	default:
	}
	close(rl.stop)
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"path"
	"strings"
	"strconv"
	"time"
)

type ServeConfig struct {
	Port string
	Dir  string
	Path string
}

type statusWriter struct {
	http.ResponseWriter
	status int
	length int
}

func (w *statusWriter) WriteHeader(status int) {
	w.status = status
	w.ResponseWriter.WriteHeader(status)
}

func (w *statusWriter) Write(b []byte) (int, error) {
	if w.status == 0 {
		w.status = 200
	}
	w.length = len(b)
	return w.ResponseWriter.Write(b)
}

// WriteLog Logs the Http Status for a request into fileHandler and returns a httphandler function which is a wrapper to log the requests.
func logger(handle http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, request *http.Request) {
		start := time.Now()
		writer := statusWriter{w, 0, 0}
		handle.ServeHTTP(&writer, request)
		end := time.Now()
		latency := end.Sub(start)
		statusCode := writer.status
		log.Println(request.Method, request.URL.Path, statusCode, latency)
	}
}

func getServeConfig(args []string) ServeConfig {

	dir, _ := os.Getwd()
	config := ServeConfig{"8000", dir, "/"}

	// if no args are passed, use default config
	if len(args) > 0 {
		for _, element := range args {

			_, err := strconv.Atoi(element)

			if err == nil {
				config.Port = element
			} else if(strings.Contains(element,":")) {
				s := strings.Split(element,":")
				config.Dir,config.Port = s[0],s[1]
			} else {
				config.Dir = element
			}
		}
	}

	return config
}

func serve(config ServeConfig) {

 	pathStat, err := os.Stat(config.Dir)
	if err != nil {
		log.Fatal(err)
	}

	if pathStat.IsDir() {
		http.Handle(config.Path, http.FileServer(http.Dir(config.Dir)))
	} else {
		http.HandleFunc(config.Path, func(w http.ResponseWriter, r *http.Request) {
			http.ServeFile(w, r, config.Dir)
		})
	}

	fmt.Println("\nServe Config \n Directory : " + path.Base(config.Dir) + " \n Path      : http://localhost:" + config.Port + config.Path + "\n")
	log.Println("Starting server on port: " + config.Port)

	log.Fatal(http.ListenAndServe(":"+config.Port, logger(http.DefaultServeMux)))
}

func main() {
	args := os.Args[1:]
	config := getServeConfig(args)
	serve(config)
}
<commit_msg>use package flag instead of os for args<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"
	"time"
)

type ServeConfig struct {
	Port string
	Dir  string
	Path string
}

type statusWriter struct {
	http.ResponseWriter
	status int
	length int
}

func (w *statusWriter) WriteHeader(status int) {
	w.status = status
	w.ResponseWriter.WriteHeader(status)
}

func (w *statusWriter) Write(b []byte) (int, error) {
	if w.status == 0 {
		w.status = http.StatusOK
	}
	w.length = len(b)
	return w.ResponseWriter.Write(b)
}

func logger(handle http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, request *http.Request) {
		start := time.Now()
		writer := statusWriter{w, 0, 0}
		handle.ServeHTTP(&writer, request)
		end := time.Now()
		latency := end.Sub(start)
		statusCode := writer.status
		log.Println(request.Method, request.URL.Path, statusCode, latency)
	}
}

func getServeConfig(args []string) ServeConfig {

	dir, _ := os.Getwd()
	config := ServeConfig{"8000", dir, "/"}

	// if no args are passed, use default config
	if len(args) > 0 {
		for _, element := range args {

			_, err := strconv.Atoi(element)

			if err == nil {
				config.Port = element
			} else if strings.Contains(element, ":") {
				s := strings.Split(element, ":")
				config.Dir, config.Port = s[0], s[1]
			} else {
				config.Dir = element
			}
		}
	}

	return config
}

func serve(config ServeConfig) {

	pathStat, err := os.Stat(config.Dir)
	if err != nil {
		log.Fatal(err)
	}

	if pathStat.IsDir() {
		http.Handle(config.Path, http.FileServer(http.Dir(config.Dir)))
	} else {
		http.HandleFunc(config.Path, func(w http.ResponseWriter, r *http.Request) {
			http.ServeFile(w, r, config.Dir)
		})
	}

	fmt.Println("\nServe Config \n Directory : " + path.Base(config.Dir) + " \n Path      : http://localhost:" + config.Port + config.Path + "\n")
	log.Println("Starting server on port: " + config.Port)

	log.Fatal(http.ListenAndServe(":"+config.Port, logger(http.DefaultServeMux)))
}

func main() {
	flag.Parse()
	args := flag.Args()
	config := getServeConfig(args)
	serve(config)
}
<|endoftext|>


<commit_before>package main

// @author Robin Verlangen
// Tool for logging data to CloudPelican directly from Go

// Imports
import (
    "cloudpelican"
    "log"
)

// Token
const TOKEN string = "123456"

// Example
func main() {
    var msg string = "This is a log message"

    // Write message and validate
    res := cloudpelican.LogMessage(TOKEN, msg)
    if !res {
        log.Println("Something went wrong")
    } else {
        log.Printf("Written %d bytes of data '%s' to backend.\n", len(msg), msg)
    }

    // Make sure any pending messages are written to the backend
    cloudpelican.Drain()
}<commit_msg>Test multiple messages<commit_after>package main

// @author Robin Verlangen
// Tool for logging data to CloudPelican directly from Go

// Imports
import (
    "cloudpelican"
    "log"
    "strconv"
)

// Token
const TOKEN string = "123456"

// Example
func main() {
    var msg string = "This is a log message"

    // Write message and validate
    for i := 0; i < 10; i++ {
        res := cloudpelican.LogMessage(TOKEN, msg + " " + strconv.Itoa(i))
        if !res {
            log.Println("Something went wrong")
        } else {
            log.Printf("Written %d bytes of data '%s' to backend.\n", len(msg), msg)
        }
    }

    // Make sure any pending messages are written to the backend
    cloudpelican.Drain()
}<|endoftext|>


<commit_before>package aws

import (
	"fmt"

	"github.com/MattAitchison/env"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"

	"github.com/gliderlabs/hostctl/providers"
)

var envSet = env.NewEnvSet("aws")

func init() {
	providers.Register(&awsProvider{}, "aws")
}

type awsProvider struct {
	client *ec2.EC2
}

func (p *awsProvider) Setup() error {
	p.client = ec2.New(nil)
	return nil
}

func (p *awsProvider) Env() *env.EnvSet {
	return envSet
}

func (p *awsProvider) Create(host providers.Host) error {
	res, err := p.client.RunInstances(&ec2.RunInstancesInput{
		ImageId:      aws.String(host.Image),
		MaxCount:     aws.Int64(1),
		MinCount:     aws.Int64(1),
		UserData:     aws.String(host.Userdata),
		KeyName:      aws.String(host.Keyname),
		InstanceType: aws.String(host.Flavor),
		Placement: &ec2.Placement{
			AvailabilityZone: aws.String(host.Region),
		},
	})
	if err != nil {
		return err
	}

	if res != nil && len(res.Instances) > 0 {
		return fmt.Errorf("no instances created")
	}
	_, err = p.client.CreateTags(&ec2.CreateTagsInput{
		Resources: []*string{
			res.Instances[0].InstanceId,
		},
		Tags: []*ec2.Tag{
			{
				Key:   aws.String("Name"),
				Value: aws.String(host.Name),
			},
		},
	})
	return err
}

func (p *awsProvider) Destroy(name string) error {
	for id, host := range p.list(name) {
		if host.Name == name {
			return p.destroy(id)
		}
	}
	return nil
}

func (p *awsProvider) destroy(id string) error {
	_, err := p.client.TerminateInstances(&ec2.TerminateInstancesInput{
		InstanceIds: []*string{aws.String(id)},
	})
	return err
}

func (p *awsProvider) List(pattern string) []providers.Host {
	var hosts []providers.Host
	for _, host := range p.list(pattern) {
		hosts = append(hosts, host)
	}
	return hosts
}

func (p *awsProvider) list(pattern string) map[string]providers.Host {
	res, err := p.client.DescribeInstances(&ec2.DescribeInstancesInput{
		Filters: []*ec2.Filter{
			{
				Name:   aws.String("tag:Name"),
				Values: []*string{aws.String(pattern)},
			},
		},
	})
	if err != nil {
		return nil
	}
	hosts := make(map[string]providers.Host)
	for i := range res.Reservations {
		for _, instance := range res.Reservations[i].Instances {
			if instance != nil && instance.PublicIpAddress != nil {
				id := *instance.InstanceId
				hosts[id] = providers.Host{
					Name: nametag(instance),
					IP:   *instance.PublicIpAddress,
				}
			}
		}
	}
	return hosts
}

func (p *awsProvider) Get(name string) *providers.Host {
	hosts := p.List(name)
	for _, host := range hosts {
		if host.Name == name {
			return &host
		}
	}
	return nil
}

func nametag(instance *ec2.Instance) string {
	for _, tag := range instance.Tags {
		if *tag.Key == "Name" {
			return *tag.Value
		}
	}
	return ""
}
<commit_msg>aws polling, env, comments, also fixed a bug in Create<commit_after>package aws

import (
	"fmt"
	"os"
	"time"

	"github.com/MattAitchison/env"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"

	"github.com/gliderlabs/hostctl/providers"
)

func init() {
	providers.Register(&awsProvider{}, "aws")
}

type awsProvider struct {
	client *ec2.EC2
	config *aws.Config
}

// Setup ec2.Client using aws supported credentials (env, credential file)
func (p *awsProvider) Setup() error {
	region := os.Getenv("HOSTCTL_REGION")
	if region == "" {
		return fmt.Errorf("HOSTCTL_REGION required")
	}
	p.config = aws.NewConfig().WithRegion(region)
	p.client = ec2.New(p.config)
	return nil
}

func (p *awsProvider) Env() *env.EnvSet {
	var envSet = env.NewEnvSet("aws")
	envSet.Secret("AWS_ACCESS_KEY", "access key for AWS")
	envSet.Secret("AWS_SECRET_KEY", "secret key for AWS")
	envSet.String("AWS_AVAILABILITY_ZONE", "", "availability zone for AWS; eg: us-west-2a")
	return envSet
}

// Create an instance based on a Host, poll until instance.Status=running
func (p *awsProvider) Create(host providers.Host) error {
	zone := os.Getenv("AWS_AVAILABILITY_ZONE")
	res, err := p.client.RunInstances(&ec2.RunInstancesInput{
		ImageId:      aws.String(host.Image),
		MaxCount:     aws.Int64(1),
		MinCount:     aws.Int64(1),
		UserData:     aws.String(host.Userdata),
		KeyName:      aws.String(host.Keyname),
		InstanceType: aws.String(host.Flavor),
		Placement: &ec2.Placement{
			AvailabilityZone: aws.String(zone),
		},
	})
	if err != nil {
		return err
	}
	if res == nil || len(res.Instances) == 0 {
		return fmt.Errorf("no instances created")
	}
	instanceID := res.Instances[0].InstanceId
	_, err = p.client.CreateTags(&ec2.CreateTagsInput{
		Resources: []*string{instanceID},
		Tags: []*ec2.Tag{
			{
				Key:   aws.String("Name"),
				Value: aws.String(host.Name),
			},
		},
	})
	if err != nil {
		return err
	}
	return p.pollFor(*instanceID, "running")
}

// Destroy the first instance with tag:Name=name
func (p *awsProvider) Destroy(name string) error {
	for id, host := range p.list(name) {
		if host.Name == name {
			return p.destroy(id)
		}
	}
	return nil
}

// destroy instance by ID polling until instance.State=terminated
func (p *awsProvider) destroy(id string) error {
	_, err := p.client.TerminateInstances(&ec2.TerminateInstancesInput{
		InstanceIds: []*string{aws.String(id)},
	})
	if err != nil {
		return err
	}
	return p.pollFor(id, "terminated")
}

func (p *awsProvider) List(pattern string) []providers.Host {
	var hosts []providers.Host
	for _, host := range p.list(pattern) {
		hosts = append(hosts, host)
	}
	return hosts
}

// list all instances with tag:Name matching pattern.
// Will NOT return any instances where PublicIpAddress is nil.
// Map key will be set to the instanceID.
func (p *awsProvider) list(pattern string) map[string]providers.Host {
	res, err := p.client.DescribeInstances(&ec2.DescribeInstancesInput{
		Filters: []*ec2.Filter{
			{
				Name:   aws.String("tag:Name"),
				Values: []*string{aws.String(pattern)},
			},
		},
	})
	if err != nil {
		return nil
	}
	hosts := make(map[string]providers.Host)
	for i := range res.Reservations {
		for _, instance := range res.Reservations[i].Instances {
			// Note: an instance that is terminated/terminating will not have a public IP.
			if instance != nil && instance.PublicIpAddress != nil {
				id := *instance.InstanceId
				hosts[id] = providers.Host{
					Name:    nametag(instance),
					IP:      *instance.PublicIpAddress,
					Region:  *instance.Placement.AvailabilityZone,
					Image:   *instance.ImageId,
					Keyname: *instance.KeyName,
					Flavor:  *instance.InstanceType,
				}
			}
		}
	}
	return hosts
}

// Get the first instance with tag:Name=name
func (p *awsProvider) Get(name string) *providers.Host {
	hosts := p.List(name)
	for _, host := range hosts {
		if host.Name == name {
			return &host
		}
	}
	return nil
}

// pollFor will poll every 2 seconds for an instance by ID,
// until instance.State=state, timeout or an error.
func (p *awsProvider) pollFor(id string, state string) error {
	timeout := time.After(1 * time.Minute)
	for {
		select {
		case <-timeout:
			return fmt.Errorf("aws provider: timed out wating for state: %s", state)
		default:
		}
		res, err := p.client.DescribeInstances(&ec2.DescribeInstancesInput{
			InstanceIds: []*string{aws.String(id)},
		})
		if err != nil {
			return err
		}
		for i := range res.Reservations {
			for _, instance := range res.Reservations[i].Instances {
				if *instance.State.Name == state {
					return nil
				}
			}
		}
		time.Sleep(2 * time.Second)
	}
}

func nametag(instance *ec2.Instance) string {
	for _, tag := range instance.Tags {
		if *tag.Key == "Name" {
			return *tag.Value
		}
	}
	return ""
}
<|endoftext|>


<commit_before>package libvirt

import (
	"errors"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/common"
	commonssh "github.com/mitchellh/packer/common/ssh"
	"github.com/mitchellh/packer/helper/communicator"
	"github.com/mitchellh/packer/helper/config"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/template/interpolate"
)

const BuilderId = "vtolstov.libvirt"

type Builder struct {
	config Config
	runner multistep.Runner
}

type Config struct {
	common.PackerConfig `mapstructure:",squash"`
	Comm                communicator.Config `mapstructure:",squash"`
	DomainType          string              `mapstructure:"domain_type"`
	BootCommand         []string            `mapstructure:"boot_command"`
	MemorySize          uint                `mapstructure:"memory_size"`
	Arch                string              `mapstructure:"arch"`
	DiskName            string              `mapstructure:"disk_name"`
	DiskType            string              `mapstructure:"disk_type"`
	DiskSize            uint                `mapstructure:"disk_size"`
	FloppyFiles         []string            `mapstructure:"floppy_files"`
	HTTPDir             string              `mapstructure:"http_directory"`
	HTTPPortMin         uint                `mapstructure:"http_port_min"`
	HTTPPortMax         uint                `mapstructure:"http_port_max"`
	ISOUrl              string              `mapstructure:"iso_url"`
	OutputDir           string              `mapstructure:"output_directory"`
	ShutdownCommand     string              `mapstructure:"shutdown_command"`
	SSHHostPortMin      uint                `mapstructure:"ssh_host_port_min"`
	SSHHostPortMax      uint                `mapstructure:"ssh_host_port_max"`
	SSHKeyPath          string              `mapstructure:"ssh_key_path"`
	SSHPassword         string              `mapstructure:"ssh_password"`
	SSHPort             uint                `mapstructure:"ssh_port"`
	SSHUser             string              `mapstructure:"ssh_username"`
	VMName              string              `mapstructure:"vm_name"`

	DomainXml   string `mapstructure:"domain_xml"`
	VolumeXml   string `mapstructure:"volume_xml"`
	PoolName    string `mapstructure:"pool_name"`
	PoolXml     string `mapstructure:"pool_xml"`
	NetworkName string `mapstructure:"network_name"`
	NetworkXml  string `mapstructure:"network_xml"`
	NetowrkType string `mapstructure:"network_type"`
	LibvirtUrl  string `mapstructure:"libvirt_url"`

	RawBootWait        string `mapstructure:"boot_wait"`
	RawShutdownTimeout string `mapstructure:"shutdown_timeout"`
	RawSSHWaitTimeout  string `mapstructure:"ssh_wait_timeout"`

	bootWait        time.Duration ``
	shutdownTimeout time.Duration ``
	sshWaitTimeout  time.Duration ``
	ctx             interpolate.Context
}

func (b *Builder) Prepare(raws ...interface{}) ([]string, error) {
	var errs *packer.MultiError
	warnings := make([]string, 0)

	err := config.Decode(&b.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &b.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"boot_command",
				"qemuargs",
			},
		},
	}, raws...)
	if err != nil {
		return nil, err
	}

	if b.config.DomainType == "" {
		b.config.DomainType = "kvm"
	}

	if b.config.DiskType == "" {
		b.config.DiskType = "raw"
	}

	if b.config.DiskSize == 0 {
		b.config.DiskSize = 5000
	}

	if b.config.PoolXml == "" {
		b.config.PoolXml = PackerPool
	}

	if b.config.PoolName == "" {
		b.config.PoolName = "packer"
	}

	if b.config.NetworkName == "" {
		b.config.NetworkName = "packer"
	}

	if b.config.VolumeXml == "" {
		b.config.VolumeXml = PackerVolume
	}

	if b.config.NetowrkType == "" {
		b.config.NetworkType = "user"
	}

	if b.config.NetworkType != "user" {
		if b.config.NetworkXml == "" {
			b.config.NetworkXml = PackerNetwork
		}
	}

	if b.config.DomainXml == "" {
		b.config.DomainXml = PackerQemuXML
	}

	if b.config.FloppyFiles == nil {
		b.config.FloppyFiles = make([]string, 0)
	}

	if b.config.HTTPPortMin == 0 {
		b.config.HTTPPortMin = 8000
	}

	if b.config.HTTPPortMax == 0 {
		b.config.HTTPPortMax = 9000
	}

	if b.config.OutputDir == "" {
		b.config.OutputDir = fmt.Sprintf("output-%s", b.config.PackerBuildName)
	}

	if b.config.RawBootWait == "" {
		b.config.RawBootWait = "10s"
	}

	if b.config.SSHHostPortMin == 0 {
		b.config.SSHHostPortMin = 2222
	}

	if b.config.SSHHostPortMax == 0 {
		b.config.SSHHostPortMax = 4444
	}

	if b.config.SSHPort == 0 {
		b.config.SSHPort = 22
	}

	if b.config.VMName == "" {
		b.config.VMName = fmt.Sprintf("packer-%s", b.config.PackerBuildName)
	}

	if b.config.DiskName == "" {
		b.config.DiskName = b.config.VMName
	}

	/*
		// Errors
		templates := map[string]*string{
			"disk_name":        &b.config.DiskName,
			"http_directory":   &b.config.HTTPDir,
			"output_directory": &b.config.OutputDir,
			"shutdown_command": &b.config.ShutdownCommand,
			"ssh_password":     &b.config.SSHPassword,
			"ssh_username":     &b.config.SSHUser,
			"vm_name":          &b.config.VMName,
			"boot_wait":        &b.config.RawBootWait,
			"shutdown_timeout": &b.config.RawShutdownTimeout,
			"ssh_wait_timeout": &b.config.RawSSHWaitTimeout,
			"domain_xml":       &b.config.DomainXml,
			"volume_xml":       &b.config.VolumeXml,
			"pool_name":        &b.config.PoolName,
			"pool_xml":         &b.config.PoolXml,
			"network_name":     &b.config.NetworkName,
			"network_xml":      &b.config.NetworkXml,
		}
	*/
	b.config.ISOUrl, err = common.DownloadableURL(b.config.ISOUrl)
	if err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Failed to parse iso_url: %s", err))
	}

	if b.config.HTTPPortMin > b.config.HTTPPortMax {
		errs = packer.MultiErrorAppend(
			errs, errors.New("http_port_min must be less than http_port_max"))
	}

	if !b.config.PackerForce {
		if _, err := os.Stat(b.config.OutputDir); err == nil {
			errs = packer.MultiErrorAppend(
				errs,
				fmt.Errorf("Output directory '%s' already exists. It must not exist.", b.config.OutputDir))
		}
	}

	b.config.bootWait, err = time.ParseDuration(b.config.RawBootWait)
	if err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Failed parsing boot_wait: %s", err))
	}

	if b.config.RawShutdownTimeout == "" {
		b.config.RawShutdownTimeout = "5m"
	}

	if b.config.RawSSHWaitTimeout == "" {
		b.config.RawSSHWaitTimeout = "20m"
	}

	b.config.shutdownTimeout, err = time.ParseDuration(b.config.RawShutdownTimeout)
	if err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Failed parsing shutdown_timeout: %s", err))
	}

	if b.config.SSHKeyPath != "" {
		if _, err := os.Stat(b.config.SSHKeyPath); err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("ssh_key_path is invalid: %s", err))
		} else if _, err := commonssh.FileSigner(b.config.SSHKeyPath); err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("ssh_key_path is invalid: %s", err))
		}
	}

	if b.config.SSHHostPortMin > b.config.SSHHostPortMax {
		errs = packer.MultiErrorAppend(
			errs, errors.New("ssh_host_port_min must be less than ssh_host_port_max"))
	}

	if b.config.SSHUser == "" && !strings.HasPrefix(b.config.LibvirtUrl, "lxc") {
		errs = packer.MultiErrorAppend(
			errs, errors.New("An ssh_username must be specified."))
	}

	b.config.sshWaitTimeout, err = time.ParseDuration(b.config.RawSSHWaitTimeout)
	if err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Failed parsing ssh_wait_timeout: %s", err))
	}

	// Warnings
	if b.config.ShutdownCommand == "" {
		warnings = append(warnings,
			"A shutdown_command was not specified. Without a shutdown command, Packer\n"+
				"will forcibly halt the virtual machine, which may result in data loss.")
	}

	if errs != nil && len(errs.Errors) > 0 {
		return warnings, errs
	}

	return warnings, nil
}

func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packer.Artifact, error) {
	// Check we have the required tools

	steps := []multistep.Step{
		new(stepPrepareOutputDir),
		&common.StepCreateFloppy{
			Files: b.config.FloppyFiles,
		},
		new(stepHTTPServer),
		new(stepCreatePool),
		new(stepCreateNetwork),
		new(stepCreateVolume),
		new(stepCreateDomain),
		new(stepRun),
		new(stepTypeBootCommand),
		&communicator.StepConnect{
			Config:    &b.config.Comm,
			Host:      commHost,
			SSHConfig: sshConfig,
			SSHPort:   commPort,
		},
		// can we upload any guest helpers?
		new(common.StepProvision),
		new(stepShutdown),
		// compact disk?
	}
	/*
	   ▶       ▶       new(stepDownloadVolume),
	   ▶       ▶       new(stepDeleteVolume),

	*/
	// Setup the state bag
	state := &multistep.BasicStateBag{}
	state.Put("cache", cache)
	state.Put("config", &b.config)
	state.Put("hook", hook)
	state.Put("ui", ui)

	// Run
	if b.config.PackerDebug {
		b.runner = &multistep.DebugRunner{
			Steps:   steps,
			PauseFn: common.MultistepDebugFn(ui),
		}
	} else {
		b.runner = &multistep.BasicRunner{Steps: steps}
	}

	b.runner.Run(state)

	// If there was an error, return that
	if rawErr, ok := state.GetOk("error"); ok {
		return nil, rawErr.(error)
	}

	// If we were interrupted or cancelled, then just exit.
	if _, ok := state.GetOk(multistep.StateCancelled); ok {
		return nil, errors.New("Build was cancelled.")
	}

	if _, ok := state.GetOk(multistep.StateHalted); ok {
		return nil, errors.New("Build was halted.")
	}

	// Compile the artifact list
	files := make([]string, 0, 5)
	visit := func(path string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			files = append(files, path)
		}

		return err
	}

	if err := filepath.Walk(b.config.OutputDir, visit); err != nil {
		return nil, err
	}

	artifact := &Artifact{
		name:  b.config.VMName,
		dir:   b.config.OutputDir,
		f:     files,
		state: make(map[string]interface{}),
	}

	artifact.state["diskType"] = b.config.DiskType
	artifact.state["diskSize"] = b.config.DiskSize
	artifact.state["domainType"] = "kvm"

	return artifact, nil
}

func (b *Builder) Cancel() {
	if b.runner != nil {
		log.Println("Cancelling the step runner...")
		b.runner.Cancel()
	}
}
<commit_msg>fix<commit_after>package libvirt

import (
	"errors"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/common"
	commonssh "github.com/mitchellh/packer/common/ssh"
	"github.com/mitchellh/packer/helper/communicator"
	"github.com/mitchellh/packer/helper/config"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/template/interpolate"
)

const BuilderId = "vtolstov.libvirt"

type Builder struct {
	config Config
	runner multistep.Runner
}

type Config struct {
	common.PackerConfig `mapstructure:",squash"`
	Comm                communicator.Config `mapstructure:",squash"`
	DomainType          string              `mapstructure:"domain_type"`
	BootCommand         []string            `mapstructure:"boot_command"`
	MemorySize          uint                `mapstructure:"memory_size"`
	Arch                string              `mapstructure:"arch"`
	DiskName            string              `mapstructure:"disk_name"`
	DiskType            string              `mapstructure:"disk_type"`
	DiskSize            uint                `mapstructure:"disk_size"`
	FloppyFiles         []string            `mapstructure:"floppy_files"`
	HTTPDir             string              `mapstructure:"http_directory"`
	HTTPPortMin         uint                `mapstructure:"http_port_min"`
	HTTPPortMax         uint                `mapstructure:"http_port_max"`
	ISOUrl              string              `mapstructure:"iso_url"`
	OutputDir           string              `mapstructure:"output_directory"`
	ShutdownCommand     string              `mapstructure:"shutdown_command"`
	SSHHostPortMin      uint                `mapstructure:"ssh_host_port_min"`
	SSHHostPortMax      uint                `mapstructure:"ssh_host_port_max"`
	SSHKeyPath          string              `mapstructure:"ssh_key_path"`
	SSHPassword         string              `mapstructure:"ssh_password"`
	SSHPort             uint                `mapstructure:"ssh_port"`
	SSHUser             string              `mapstructure:"ssh_username"`
	VMName              string              `mapstructure:"vm_name"`

	DomainXml   string `mapstructure:"domain_xml"`
	VolumeXml   string `mapstructure:"volume_xml"`
	PoolName    string `mapstructure:"pool_name"`
	PoolXml     string `mapstructure:"pool_xml"`
	NetworkName string `mapstructure:"network_name"`
	NetworkXml  string `mapstructure:"network_xml"`
	NetworkType string `mapstructure:"network_type"`
	LibvirtUrl  string `mapstructure:"libvirt_url"`

	RawBootWait        string `mapstructure:"boot_wait"`
	RawShutdownTimeout string `mapstructure:"shutdown_timeout"`
	RawSSHWaitTimeout  string `mapstructure:"ssh_wait_timeout"`

	bootWait        time.Duration ``
	shutdownTimeout time.Duration ``
	sshWaitTimeout  time.Duration ``
	ctx             interpolate.Context
}

func (b *Builder) Prepare(raws ...interface{}) ([]string, error) {
	var errs *packer.MultiError
	warnings := make([]string, 0)

	err := config.Decode(&b.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &b.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"boot_command",
				"qemuargs",
			},
		},
	}, raws...)
	if err != nil {
		return nil, err
	}

	if b.config.DomainType == "" {
		b.config.DomainType = "kvm"
	}

	if b.config.DiskType == "" {
		b.config.DiskType = "raw"
	}

	if b.config.DiskSize == 0 {
		b.config.DiskSize = 5000
	}

	if b.config.PoolXml == "" {
		b.config.PoolXml = PackerPool
	}

	if b.config.PoolName == "" {
		b.config.PoolName = "packer"
	}

	if b.config.NetworkName == "" {
		b.config.NetworkName = "packer"
	}

	if b.config.VolumeXml == "" {
		b.config.VolumeXml = PackerVolume
	}

	if b.config.NetworkType == "" {
		b.config.NetworkType = "user"
	}

	if b.config.NetworkType != "user" {
		if b.config.NetworkXml == "" {
			b.config.NetworkXml = PackerNetwork
		}
	}

	if b.config.DomainXml == "" {
		b.config.DomainXml = PackerQemuXML
	}

	if b.config.FloppyFiles == nil {
		b.config.FloppyFiles = make([]string, 0)
	}

	if b.config.HTTPPortMin == 0 {
		b.config.HTTPPortMin = 8000
	}

	if b.config.HTTPPortMax == 0 {
		b.config.HTTPPortMax = 9000
	}

	if b.config.OutputDir == "" {
		b.config.OutputDir = fmt.Sprintf("output-%s", b.config.PackerBuildName)
	}

	if b.config.RawBootWait == "" {
		b.config.RawBootWait = "10s"
	}

	if b.config.SSHHostPortMin == 0 {
		b.config.SSHHostPortMin = 2222
	}

	if b.config.SSHHostPortMax == 0 {
		b.config.SSHHostPortMax = 4444
	}

	if b.config.SSHPort == 0 {
		b.config.SSHPort = 22
	}

	if b.config.VMName == "" {
		b.config.VMName = fmt.Sprintf("packer-%s", b.config.PackerBuildName)
	}

	if b.config.DiskName == "" {
		b.config.DiskName = b.config.VMName
	}

	/*
		// Errors
		templates := map[string]*string{
			"disk_name":        &b.config.DiskName,
			"http_directory":   &b.config.HTTPDir,
			"output_directory": &b.config.OutputDir,
			"shutdown_command": &b.config.ShutdownCommand,
			"ssh_password":     &b.config.SSHPassword,
			"ssh_username":     &b.config.SSHUser,
			"vm_name":          &b.config.VMName,
			"boot_wait":        &b.config.RawBootWait,
			"shutdown_timeout": &b.config.RawShutdownTimeout,
			"ssh_wait_timeout": &b.config.RawSSHWaitTimeout,
			"domain_xml":       &b.config.DomainXml,
			"volume_xml":       &b.config.VolumeXml,
			"pool_name":        &b.config.PoolName,
			"pool_xml":         &b.config.PoolXml,
			"network_name":     &b.config.NetworkName,
			"network_xml":      &b.config.NetworkXml,
		}
	*/
	b.config.ISOUrl, err = common.DownloadableURL(b.config.ISOUrl)
	if err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Failed to parse iso_url: %s", err))
	}

	if b.config.HTTPPortMin > b.config.HTTPPortMax {
		errs = packer.MultiErrorAppend(
			errs, errors.New("http_port_min must be less than http_port_max"))
	}

	if !b.config.PackerForce {
		if _, err := os.Stat(b.config.OutputDir); err == nil {
			errs = packer.MultiErrorAppend(
				errs,
				fmt.Errorf("Output directory '%s' already exists. It must not exist.", b.config.OutputDir))
		}
	}

	b.config.bootWait, err = time.ParseDuration(b.config.RawBootWait)
	if err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Failed parsing boot_wait: %s", err))
	}

	if b.config.RawShutdownTimeout == "" {
		b.config.RawShutdownTimeout = "5m"
	}

	if b.config.RawSSHWaitTimeout == "" {
		b.config.RawSSHWaitTimeout = "20m"
	}

	b.config.shutdownTimeout, err = time.ParseDuration(b.config.RawShutdownTimeout)
	if err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Failed parsing shutdown_timeout: %s", err))
	}

	if b.config.SSHKeyPath != "" {
		if _, err := os.Stat(b.config.SSHKeyPath); err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("ssh_key_path is invalid: %s", err))
		} else if _, err := commonssh.FileSigner(b.config.SSHKeyPath); err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("ssh_key_path is invalid: %s", err))
		}
	}

	if b.config.SSHHostPortMin > b.config.SSHHostPortMax {
		errs = packer.MultiErrorAppend(
			errs, errors.New("ssh_host_port_min must be less than ssh_host_port_max"))
	}

	if b.config.SSHUser == "" && !strings.HasPrefix(b.config.LibvirtUrl, "lxc") {
		errs = packer.MultiErrorAppend(
			errs, errors.New("An ssh_username must be specified."))
	}

	b.config.sshWaitTimeout, err = time.ParseDuration(b.config.RawSSHWaitTimeout)
	if err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Failed parsing ssh_wait_timeout: %s", err))
	}

	// Warnings
	if b.config.ShutdownCommand == "" {
		warnings = append(warnings,
			"A shutdown_command was not specified. Without a shutdown command, Packer\n"+
				"will forcibly halt the virtual machine, which may result in data loss.")
	}

	if errs != nil && len(errs.Errors) > 0 {
		return warnings, errs
	}

	return warnings, nil
}

func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packer.Artifact, error) {
	// Check we have the required tools

	steps := []multistep.Step{
		new(stepPrepareOutputDir),
		&common.StepCreateFloppy{
			Files: b.config.FloppyFiles,
		},
		new(stepHTTPServer),
		new(stepCreatePool),
		new(stepCreateNetwork),
		new(stepCreateVolume),
		new(stepCreateDomain),
		new(stepRun),
		new(stepTypeBootCommand),
		&communicator.StepConnect{
			Config:    &b.config.Comm,
			Host:      commHost,
			SSHConfig: sshConfig,
			SSHPort:   commPort,
		},
		// can we upload any guest helpers?
		new(common.StepProvision),
		new(stepShutdown),
		// compact disk?
	}
	/*
	   ▶       ▶       new(stepDownloadVolume),
	   ▶       ▶       new(stepDeleteVolume),

	*/
	// Setup the state bag
	state := &multistep.BasicStateBag{}
	state.Put("cache", cache)
	state.Put("config", &b.config)
	state.Put("hook", hook)
	state.Put("ui", ui)

	// Run
	if b.config.PackerDebug {
		b.runner = &multistep.DebugRunner{
			Steps:   steps,
			PauseFn: common.MultistepDebugFn(ui),
		}
	} else {
		b.runner = &multistep.BasicRunner{Steps: steps}
	}

	b.runner.Run(state)

	// If there was an error, return that
	if rawErr, ok := state.GetOk("error"); ok {
		return nil, rawErr.(error)
	}

	// If we were interrupted or cancelled, then just exit.
	if _, ok := state.GetOk(multistep.StateCancelled); ok {
		return nil, errors.New("Build was cancelled.")
	}

	if _, ok := state.GetOk(multistep.StateHalted); ok {
		return nil, errors.New("Build was halted.")
	}

	// Compile the artifact list
	files := make([]string, 0, 5)
	visit := func(path string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			files = append(files, path)
		}

		return err
	}

	if err := filepath.Walk(b.config.OutputDir, visit); err != nil {
		return nil, err
	}

	artifact := &Artifact{
		name:  b.config.VMName,
		dir:   b.config.OutputDir,
		f:     files,
		state: make(map[string]interface{}),
	}

	artifact.state["diskType"] = b.config.DiskType
	artifact.state["diskSize"] = b.config.DiskSize
	artifact.state["domainType"] = "kvm"

	return artifact, nil
}

func (b *Builder) Cancel() {
	if b.runner != nil {
		log.Println("Cancelling the step runner...")
		b.runner.Cancel()
	}
}
<|endoftext|>


<commit_before>package main

func Factors(n uint64) ([]uint64) {
    factors := make([]uint64, 0, 2)

    for i := uint64(2); n > 1; {
        if n % i == 0 {
            factors = append(factors, i)
            n /= i
        } else {
            i++
        }
    }

    return factors
}

<commit_msg>go: add Primes util<commit_after>package main

func Factors(n uint64) ([]uint64) {
    factors := make([]uint64, 0, 2)

    for i := uint64(2); n > 1; {
        if n % i == 0 {
            factors = append(factors, i)
            n /= i
        } else {
            i++
        }
    }

    return factors
}

func Primes(n uint64) ([]uint64) {
    composites := make([]bool, n)
    composites[0] = true
    composites[1] = true

    for i := uint64(2); i < n; {
        for j := i * 2; j < n; j += i {
            composites[j] = true
        }
        for i++ ; i < n && composites[i]; i++ {}
    }

    primes := make([]uint64, 0)

    for i := uint64(2); i < n; i++ {
        if !composites[i] {
            primes = append(primes, i)
        }
    }

    return primes
}
<|endoftext|>