<commit_before>package a

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"lutils/httplib"
	"lutils/logs"
	"lutils/mahonia"
	"reflect"
	"sort"
	"strings"
	"time"
)

//公共请求参数
type requestParams struct {
	AppId        string `ali:"app_id"`
	Method       string `ali:"method"`
	Format       string `ali:"format"`
	Charset      string `ali:"charset"`
	SignType     string `ali:"sign_type"`
	Sign         string `ali:"sign"`
	TimeStamp    string `ali:"timestamp"`
	Version      string `ali:"version"`
	AppAuthToken string `ali:"app_auth_token"`
	BizContent   string `ali:"biz_content"`
	AuthToken    string `ali:"auth_token"`
	Code         string `ali:"code"`
	GrantType    string `ali:"grant_type"`
}

func (params *requestParams) valid() error {
	if len(params.AppId) == 0 {
		return ErrAppIdNil
	}

	if len(params.Method) == 0 {
		return errors.New("method 不能为空")
	}

	if len(params.Format) == 0 {
		params.Format = "JSON"
	}

	if len(params.Format) > 0 && params.Format != "JSON" {
		return errors.New("format 仅支持JSON")
	}

	if len(params.Charset) == 0 {
		params.Charset = "GBK"
	}

	if len(params.Charset) > 0 && strings.ToUpper(params.Charset) != "GBK" {
		return errors.New("charset 目前仅支持GBK")
	}

	if len(params.SignType) == 0 {
		params.SignType = "RSA"
	}

	if len(params.SignType) > 0 && params.SignType != "RSA" {
		return errors.New("sign_type 仅支持RSA")
	}

	if len(params.TimeStamp) == 0 {
		params.TimeStamp = time.Now().Format("2006-01-02 15:04:05")
	}

	if len(params.TimeStamp) > 0 {
		if _, err := time.Parse("2006-01-02 15:04:05", params.TimeStamp); err != nil {
			return errors.New("timestamp 格式\"yyyy-MM-dd HH:mm:ss\"")
		}
	}

	if len(params.Version) == 0 {
		params.Version = "1.0"
	} else if params.Version != "1.0" {
		return errors.New("version 固定为：1.0")
	}

	return nil
}

type bizInterface interface {
	valid() error
}

type responseInterface interface{}

type Response struct {
	Code    string `json:"code,omitempty"`
	Msg     string `json:"msg,omitempty"`
	SubCode string `json:"sub_code,omitempty"`
	SubMsg  string `json:"sub_msg,omitempty"`
}

type ApiHander interface {
	apiMethod() string
	apiName() string
}

var (
	apiRegistry map[string]AlipayApi
)

func init() {
	apiRegistry = map[string]AlipayApi{}
}

func registerApi(handler ApiHander) {
	apiRegistry[handler.apiMethod()] = AlipayApi{
		apiname:   handler.apiName,
		apimethod: handler.apiMethod,
	}
}

func GetApi(method string) AlipayApi {
	return apiRegistry[method]
}

func GetSupportApis() string {
	lst := "\n=====================SUPPORTED ALIPAY API LIST=====================\n"
	for _, v := range apiRegistry {
		lst += "====[" + v.apiname() + ":" + v.apimethod() + "]====" + "\n"
	}
	return lst
}

type AlipayApi struct {
	params    requestParams
	apiname   func() string
	apimethod func() string
}

func (a *AlipayApi) SetAppId(app_id string) error {
	a.params.AppId = app_id
	if len(a.params.AppId) == 0 {
		return ErrAppIdNil
	}

	if _, ok := secretLst[a.params.AppId]; !ok {
		return ErrSecretNil
	}
	return nil
}

func (a *AlipayApi) SetBizContent(biz bizInterface) error {
	//有几个接口要独立处理
	if reflect.TypeOf(biz).Name() == reflect.TypeOf(Biz_alipay_system_oauth_token{}).Name() {
		b := biz.(Biz_alipay_system_oauth_token)
		a.params.Code = b.Code
		a.params.GrantType = b.GrantType
		a.params.BizContent = ""
		return nil
	} else if reflect.TypeOf(biz).Name() == reflect.TypeOf(Biz_alipay_user_info_share{}).Name() {
		b := biz.(Biz_alipay_user_info_share)
		a.params.AuthToken = b.AuthToken
		a.params.BizContent = ""
		return nil
	} else if reflect.TypeOf(biz).Name() == reflect.TypeOf(Biz_alipay_user_userinfo_share{}).Name() {
		b := biz.(Biz_alipay_user_userinfo_share)
		a.params.AuthToken = b.AuthToken
		a.params.BizContent = ""
		return nil
	} else if reflect.TypeOf(biz).Name() == reflect.TypeOf(Biz_alipay_open_auth_token_app{}).Name() {
		b := biz.(Biz_alipay_open_auth_token_app)
		a.params.Code = b.Code
		a.params.GrantType = b.GrantType
		a.params.BizContent = ""
		return nil
	}

	if v, err := a.biz_to_string(biz); err != nil {
		return err
	} else {
		a.params.BizContent = v
		return nil
	}
}

func (a *AlipayApi) struct_to_map() map[string]interface{} {
	t := reflect.TypeOf(a.params)
	v := reflect.ValueOf(a.params)

	var data = make(map[string]interface{})
	for i := 0; i < t.NumField(); i++ {
		key := t.Field(i).Name
		value := v.Field(i).Interface()
		tag := t.Field(i).Tag.Get("ali")
		if tag != "" {
			if strings.Contains(tag, ",") {
				ps := strings.Split(tag, ",")
				key = ps[0]
			} else {
				key = tag
			}
		}
		data[key] = value
	}
	return data
}

func (a *AlipayApi) utf_to_gbk(utf string) (gbk string) {
	gbk_enc := mahonia.NewEncoder("GBK")
	return gbk_enc.ConvertString(utf)
}

func (a *AlipayApi) gbk_to_utf(gbk string) (utf string) {
	enc := mahonia.NewDecoder("GBK")
	return enc.ConvertString(gbk)
}

func (a *AlipayApi) map_to_string(m map[string]interface{}) string {
	//对key进行升序排序.
	sorted_keys := make([]string, 0)
	for k, _ := range m {
		sorted_keys = append(sorted_keys, k)
	}
	sort.Strings(sorted_keys)

	//对key=value的键值对用&连接起来，略过空值
	var signStrings string
	for _, k := range sorted_keys {
		value := fmt.Sprintf("%v", m[k])
		if value != "" {
			signStrings = signStrings + k + "=" + value + "&"
		}
	}
	if len(signStrings) == 0 {
		return ""
	} else {
		signStrings = signStrings[:len(signStrings)-1]
	}
	return signStrings
}

func (a *AlipayApi) biz_to_string(b bizInterface) (string, error) {
	if err := b.valid(); err != nil {
		return "", err
	}
	content := ""
	if v, err := json.Marshal(&b); err != nil {
		return "", err
	} else {
		content = string(v)
	}

	temp_map := map[string]interface{}{
		"biz": content,
	}

	if v, err := json.Marshal(&temp_map); err != nil {
		return "", err
	} else {
		content = string(v)
	}
	return content[8 : len(content)-2], nil
}

func (a *AlipayApi) sign(c string) (sign string, err error) {
	//签名
	s := getSecret(a.params.AppId)
	block, _ := pem.Decode(s.PrivRSA)
	if block == nil {
		return "", errors.New("privateKey error!")
	}
	priv, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return "", err
	}
	_t := crypto.SHA1.New()
	_t.Write([]byte(c))
	digest := _t.Sum(nil)
	rsa_data, err := rsa.SignPKCS1v15(rand.Reader, priv, crypto.SHA1, digest)
	if err != nil {
		return "", err
	}
	result := string(rsa_data)
	result = base64.StdEncoding.EncodeToString(rsa_data)
	return result, nil
}

func (a *AlipayApi) verifySign(s, origin_sign, method_key string) bool {
	sign_start_index := strings.Index(s, ",\"sign\"")
	if sign_start_index == -1 {
		return false
	}
	tobe_signed := s[4+len(method_key) : sign_start_index]
	block, _ := pem.Decode(getSecret(a.params.AppId).AliPubRSA)
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		fmt.Printf("Failed to parse RSA public key: %s\n", err)
		return false
	}
	rsaPub, _ := pub.(*rsa.PublicKey)
	t := sha1.New()
	io.WriteString(t, tobe_signed)
	digest := t.Sum(nil)
	data, err := base64.StdEncoding.DecodeString(origin_sign)
	if err != nil {
		fmt.Println("DecodeString sig error, reason: ", err)
		return false
	}
	err = rsa.VerifyPKCS1v15(rsaPub, crypto.SHA1, digest, data)
	if err != nil {
		fmt.Println("Verify sig error, reason: ", err)
		return false
	}

	return true
}

func (a *AlipayApi) request(m map[string]interface{}) (string, error) {
	url_link := "https://openapi.alipay.com/gateway.do"
	if conf.SandBoxEnable {
		url_link = "https://openapi.alipaydev.com/gateway.do"
	}
	http_request := httplib.Post(url_link)
	tmp_string := ""
	for k, _ := range m {
		value := fmt.Sprintf("%v", m[k])
		if value != "" {
			http_request.Param(k, value)
			tmp_string = tmp_string + k + "=" + value + "\t"
		}
	}

	logs.DEBUG(fmt.Sprintf("==[请求参数]==[%s]", tmp_string))
	var string_result string
	if v, err := http_request.String(); err != nil {
		return "", err
	} else {
		string_result = v

	}
	return string_result, nil
}

func (a *AlipayApi) Run(resp responseInterface) error {
	defer logs.DEBUG("=====================ALIPAY REQUEST END=====================")
	logs.DEBUG("=====================ALIPAY REQUEST START=====================")
	logs.DEBUG(fmt.Sprintf("==[沙盒模式]==[%v]", conf.SandBoxEnable))
	logs.DEBUG(fmt.Sprintf("==[调用方法]==[%s]:[%s]", a.apiname(), a.apimethod()))
	a.params.Method = a.apimethod()
	if err := a.params.valid(); err != nil {
		fmt.Println(err.Error())
		return err
	}

	rv := reflect.ValueOf(resp)
	if rv.Kind() != reflect.Ptr || rv.IsNil() {
		if rv.Type() == nil {
			return errors.New("responseInterface is nil")
		}

		if rv.Type().Kind() != reflect.Ptr {
			return errors.New("responseInterface is non-pointer" + rv.Type().String())
		}
		return errors.New("responseInterface is nil" + rv.Type().String())
	}

	m := a.struct_to_map()
	//做请求参数的签名
	__sign := ""
	tobe_sign := a.map_to_string(m)
	logs.DEBUG(fmt.Sprintf("==[准备签名]==[%s]", tobe_sign))
	if v, err := a.sign(tobe_sign); err != nil {
		return err
	} else if len(v) == 0 {
		return ErrSign
	} else {
		__sign = v
	}
	logs.DEBUG(fmt.Sprintf("==[签名结果]==[%s]", __sign))
	m["sign"] = __sign
	//准备请求
	result_string := ""
	if v, err := a.request(m); err != nil {
		return err
	} else {
		result_string = v
		logs.DEBUG(fmt.Sprintf("==[响应结果]==[GBK 编码]:[%s]", result_string))
	}

	//把 method 转换为 key
	method_key := strings.Replace(a.params.Method, ".", "_", -1)
	method_key += "_response"

	result_string = strings.Replace(result_string, "error_response", method_key, -1)
	//解析结果
	resp_map := map[string]interface{}{}
	if err := json.Unmarshal([]byte(result_string), &resp_map); err != nil {
		return err
	}
	//原始签名
	if v, ok := resp_map["sign"].(string); ok && len(v) > 0 {
		if pass := a.verifySign(result_string, v, method_key); !pass {
			return ErrVerifySign
		}
	}

	//转码
	result_string = a.gbk_to_utf(result_string)
	logs.DEBUG(fmt.Sprintf("==[响应结果]==[UTF 编码]:[%s]", result_string))

	if err := json.Unmarshal([]byte(result_string), &resp_map); err != nil {
		return err
	}
	//把需要的内容再次换成string
	if v, err := json.Marshal(resp_map[method_key]); err != nil {
		return err
	} else {
		result_string = string(v)
	}

	if err := json.Unmarshal([]byte(result_string), resp); err != nil {
		return err
	}
	return nil
}
<commit_msg>加点日志<commit_after>package a

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"lutils/httplib"
	"lutils/logs"
	"lutils/mahonia"
	"reflect"
	"sort"
	"strings"
	"time"
)

//公共请求参数
type requestParams struct {
	AppId        string `ali:"app_id"`
	Method       string `ali:"method"`
	Format       string `ali:"format"`
	Charset      string `ali:"charset"`
	SignType     string `ali:"sign_type"`
	Sign         string `ali:"sign"`
	TimeStamp    string `ali:"timestamp"`
	Version      string `ali:"version"`
	AppAuthToken string `ali:"app_auth_token"`
	BizContent   string `ali:"biz_content"`
	AuthToken    string `ali:"auth_token"`
	Code         string `ali:"code"`
	GrantType    string `ali:"grant_type"`
}

func (params *requestParams) valid() error {
	if len(params.AppId) == 0 {
		return ErrAppIdNil
	}

	if len(params.Method) == 0 {
		return errors.New("method 不能为空")
	}

	if len(params.Format) == 0 {
		params.Format = "JSON"
	}

	if len(params.Format) > 0 && params.Format != "JSON" {
		return errors.New("format 仅支持JSON")
	}

	if len(params.Charset) == 0 {
		params.Charset = "GBK"
	}

	if len(params.Charset) > 0 && strings.ToUpper(params.Charset) != "GBK" {
		return errors.New("charset 目前仅支持GBK")
	}

	if len(params.SignType) == 0 {
		params.SignType = "RSA"
	}

	if len(params.SignType) > 0 && params.SignType != "RSA" {
		return errors.New("sign_type 仅支持RSA")
	}

	if len(params.TimeStamp) == 0 {
		params.TimeStamp = time.Now().Format("2006-01-02 15:04:05")
	}

	if len(params.TimeStamp) > 0 {
		if _, err := time.Parse("2006-01-02 15:04:05", params.TimeStamp); err != nil {
			return errors.New("timestamp 格式\"yyyy-MM-dd HH:mm:ss\"")
		}
	}

	if len(params.Version) == 0 {
		params.Version = "1.0"
	} else if params.Version != "1.0" {
		return errors.New("version 固定为：1.0")
	}

	return nil
}

type bizInterface interface {
	valid() error
}

type responseInterface interface{}

type Response struct {
	Code    string `json:"code,omitempty"`
	Msg     string `json:"msg,omitempty"`
	SubCode string `json:"sub_code,omitempty"`
	SubMsg  string `json:"sub_msg,omitempty"`
}

type ApiHander interface {
	apiMethod() string
	apiName() string
}

var (
	apiRegistry map[string]AlipayApi
)

func init() {
	apiRegistry = map[string]AlipayApi{}
}

func registerApi(handler ApiHander) {
	apiRegistry[handler.apiMethod()] = AlipayApi{
		apiname:   handler.apiName,
		apimethod: handler.apiMethod,
	}
}

func GetApi(method string) AlipayApi {
	return apiRegistry[method]
}

func GetSupportApis() string {
	lst := "\n=====================SUPPORTED ALIPAY API LIST=====================\n"
	for _, v := range apiRegistry {
		lst += "====[" + v.apiname() + ":" + v.apimethod() + "]====" + "\n"
	}
	return lst
}

type AlipayApi struct {
	params    requestParams
	apiname   func() string
	apimethod func() string
}

func (a *AlipayApi) SetAppId(app_id string) error {
	a.params.AppId = app_id
	if len(a.params.AppId) == 0 {
		return ErrAppIdNil
	}

	if _, ok := secretLst[a.params.AppId]; !ok {
		return ErrSecretNil
	}
	return nil
}

func (a *AlipayApi) SetBizContent(biz bizInterface) error {
	//有几个接口要独立处理
	if reflect.TypeOf(biz).Name() == reflect.TypeOf(Biz_alipay_system_oauth_token{}).Name() {
		b := biz.(Biz_alipay_system_oauth_token)
		a.params.Code = b.Code
		a.params.GrantType = b.GrantType
		a.params.BizContent = ""
		return nil
	} else if reflect.TypeOf(biz).Name() == reflect.TypeOf(Biz_alipay_user_info_share{}).Name() {
		b := biz.(Biz_alipay_user_info_share)
		a.params.AuthToken = b.AuthToken
		a.params.BizContent = ""
		return nil
	} else if reflect.TypeOf(biz).Name() == reflect.TypeOf(Biz_alipay_user_userinfo_share{}).Name() {
		b := biz.(Biz_alipay_user_userinfo_share)
		a.params.AuthToken = b.AuthToken
		a.params.BizContent = ""
		return nil
	}

	if v, err := a.biz_to_string(biz); err != nil {
		return err
	} else {
		a.params.BizContent = v
		return nil
	}
}

func (a *AlipayApi) struct_to_map() map[string]interface{} {
	t := reflect.TypeOf(a.params)
	v := reflect.ValueOf(a.params)

	var data = make(map[string]interface{})
	for i := 0; i < t.NumField(); i++ {
		key := t.Field(i).Name
		value := v.Field(i).Interface()
		tag := t.Field(i).Tag.Get("ali")
		if tag != "" {
			if strings.Contains(tag, ",") {
				ps := strings.Split(tag, ",")
				key = ps[0]
			} else {
				key = tag
			}
		}
		data[key] = value
	}
	return data
}

func (a *AlipayApi) utf_to_gbk(utf string) (gbk string) {
	gbk_enc := mahonia.NewEncoder("GBK")
	return gbk_enc.ConvertString(utf)
}

func (a *AlipayApi) gbk_to_utf(gbk string) (utf string) {
	enc := mahonia.NewDecoder("GBK")
	return enc.ConvertString(gbk)
}

func (a *AlipayApi) map_to_string(m map[string]interface{}) string {
	//对key进行升序排序.
	sorted_keys := make([]string, 0)
	for k, _ := range m {
		sorted_keys = append(sorted_keys, k)
	}
	sort.Strings(sorted_keys)

	//对key=value的键值对用&连接起来，略过空值
	var signStrings string
	for _, k := range sorted_keys {
		value := fmt.Sprintf("%v", m[k])
		if value != "" {
			signStrings = signStrings + k + "=" + value + "&"
		}
	}
	if len(signStrings) == 0 {
		return ""
	} else {
		signStrings = signStrings[:len(signStrings)-1]
	}
	return signStrings
}

func (a *AlipayApi) biz_to_string(b bizInterface) (string, error) {
	if err := b.valid(); err != nil {
		return "", err
	}
	content := ""
	if v, err := json.Marshal(&b); err != nil {
		return "", err
	} else {
		content = string(v)
	}

	temp_map := map[string]interface{}{
		"biz": content,
	}

	if v, err := json.Marshal(&temp_map); err != nil {
		return "", err
	} else {
		content = string(v)
	}
	return content[8 : len(content)-2], nil
}

func (a *AlipayApi) sign(c string) (sign string, err error) {
	//签名
	s := getSecret(a.params.AppId)
	block, _ := pem.Decode(s.PrivRSA)
	if block == nil {
		return "", errors.New("privateKey error!")
	}
	priv, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return "", err
	}
	_t := crypto.SHA1.New()
	_t.Write([]byte(c))
	digest := _t.Sum(nil)
	rsa_data, err := rsa.SignPKCS1v15(rand.Reader, priv, crypto.SHA1, digest)
	if err != nil {
		return "", err
	}
	result := string(rsa_data)
	result = base64.StdEncoding.EncodeToString(rsa_data)
	return result, nil
}

func (a *AlipayApi) verifySign(s, origin_sign, method_key string) bool {
	sign_start_index := strings.Index(s, ",\"sign\"")
	if sign_start_index == -1 {
		return false
	}
	tobe_signed := s[4+len(method_key) : sign_start_index]
	block, _ := pem.Decode(getSecret(a.params.AppId).AliPubRSA)
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		fmt.Printf("Failed to parse RSA public key: %s\n", err)
		return false
	}
	rsaPub, _ := pub.(*rsa.PublicKey)
	t := sha1.New()
	io.WriteString(t, tobe_signed)
	digest := t.Sum(nil)
	data, err := base64.StdEncoding.DecodeString(origin_sign)
	if err != nil {
		fmt.Println("DecodeString sig error, reason: ", err)
		return false
	}
	err = rsa.VerifyPKCS1v15(rsaPub, crypto.SHA1, digest, data)
	if err != nil {
		fmt.Println("Verify sig error, reason: ", err)
		return false
	}

	return true
}

func (a *AlipayApi) request(m map[string]interface{}) (string, error) {
	url_link := "https://openapi.alipay.com/gateway.do"
	if conf.SandBoxEnable {
		url_link = "https://openapi.alipaydev.com/gateway.do"
	}
	http_request := httplib.Post(url_link)
	tmp_string := ""
	for k, _ := range m {
		value := fmt.Sprintf("%v", m[k])
		if value != "" {
			http_request.Param(k, value)
			logs.DEBUG(fmt.Sprintf("==[请求参数]==[%s]:[%s]", k, v))
			tmp_string = tmp_string + k + "=" + value + "\t"
		}
	}
	logs.DEBUG(fmt.Sprintf("==[请求参数]==[%s]", tmp_string))
	var string_result string
	if v, err := http_request.String(); err != nil {
		return "", err
	} else {
		string_result = v

	}
	return string_result, nil
}

func (a *AlipayApi) Run(resp responseInterface) error {
	defer logs.DEBUG("=====================ALIPAY REQUEST END=====================")
	logs.DEBUG("=====================ALIPAY REQUEST START=====================")
	logs.DEBUG(fmt.Sprintf("==[沙盒模式]==[%v]", conf.SandBoxEnable))
	logs.DEBUG(fmt.Sprintf("==[调用方法]==[%s]:[%s]", a.apiname(), a.apimethod()))
	a.params.Method = a.apimethod()
	if err := a.params.valid(); err != nil {
		fmt.Println(err.Error())
		return err
	}

	rv := reflect.ValueOf(resp)
	if rv.Kind() != reflect.Ptr || rv.IsNil() {
		if rv.Type() == nil {
			return errors.New("responseInterface is nil")
		}

		if rv.Type().Kind() != reflect.Ptr {
			return errors.New("responseInterface is non-pointer" + rv.Type().String())
		}
		return errors.New("responseInterface is nil" + rv.Type().String())
	}

	m := a.struct_to_map()
	//做请求参数的签名
	__sign := ""
	tobe_sign := a.map_to_string(m)
	logs.DEBUG(fmt.Sprintf("==[准备签名]==[%s]", tobe_sign))
	if v, err := a.sign(tobe_sign); err != nil {
		return err
	} else if len(v) == 0 {
		return ErrSign
	} else {
		__sign = v
	}
	logs.DEBUG(fmt.Sprintf("==[签名结果]==[%s]", __sign))
	m["sign"] = __sign
	//准备请求
	result_string := ""
	if v, err := a.request(m); err != nil {
		return err
	} else {
		result_string = v
		logs.DEBUG(fmt.Sprintf("==[响应结果]==[GBK 编码]:[%s]", result_string))
	}

	//把 method 转换为 key
	method_key := strings.Replace(a.params.Method, ".", "_", -1)
	method_key += "_response"

	result_string = strings.Replace(result_string, "error_response", method_key, -1)
	//解析结果
	resp_map := map[string]interface{}{}
	if err := json.Unmarshal([]byte(result_string), &resp_map); err != nil {
		return err
	}
	//原始签名
	if v, ok := resp_map["sign"].(string); ok && len(v) > 0 {
		if pass := a.verifySign(result_string, v, method_key); !pass {
			return ErrVerifySign
		}
	}

	//转码
	result_string = a.gbk_to_utf(result_string)
	logs.DEBUG(fmt.Sprintf("==[响应结果]==[UTF 编码]:[%s]", result_string))

	if err := json.Unmarshal([]byte(result_string), &resp_map); err != nil {
		return err
	}
	//把需要的内容再次换成string
	if v, err := json.Marshal(resp_map[method_key]); err != nil {
		return err
	} else {
		result_string = string(v)
	}

	if err := json.Unmarshal([]byte(result_string), resp); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package disk

import (
	"encoding/binary"
	"errors"
	"hash/crc32"
	"io"
	"os"
)

var (
	crc32cTable       = crc32.MakeTable(crc32.Castagnoli)
	crc32Len          = int64(4)
	ErrBadPayloadSize = errors.New("disk: bad payload size")
)

// writeBlock writes a full block into ws. len(p) must be (bs - crc32Len).
// Each block contains a crc32Len bytes crc32c of the payload and a (bs-crc32Len)
// bytes payload. Any error encountered is returned.
func writeBlock(ws io.WriteSeeker, index, bs int64, p []byte) error {
	if int64(len(p)) != bs-crc32Len {
		return ErrBadPayloadSize
	}

	// seek to the beginning of the block
	_, err := ws.Seek(index*bs, os.SEEK_SET)
	if err != nil {
		return nil
	}

	// write crc32c
	// TODO: reuse buffer
	b := make([]byte, crc32Len)
	binary.BigEndian.PutUint32(b, crc32.Checksum(p, crc32cTable))
	_, err = ws.Write(b)
	if err != nil {
		return err
	}

	// write payload
	_, err = ws.Write(p)
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>disk: loosen the checking in writeBlock<commit_after>package disk

import (
	"encoding/binary"
	"errors"
	"hash/crc32"
	"io"
	"os"
)

var (
	crc32cTable            = crc32.MakeTable(crc32.Castagnoli)
	crc32Len               = int64(4)
	ErrPayloadSizeTooLarge = errors.New("disk: payload size too large")
)

// writeBlock writes a full block into ws. len(p) must be smaller than (bs - crc32Len).
// Each block contains a crc32Len bytes crc32c of the payload and a (bs-crc32Len)
// bytes payload. Any error encountered is returned.
// It is caller's responsibility to ensure that only the last block in the file has
// len(p) < bs - crc32Len
func writeBlock(ws io.WriteSeeker, index, bs int64, p []byte) error {
	if int64(len(p)) > bs-crc32Len {
		return ErrPayloadSizeTooLarge
	}

	// seek to the beginning of the block
	_, err := ws.Seek(index*bs, os.SEEK_SET)
	if err != nil {
		return nil
	}

	// write crc32c
	// TODO: reuse buffer
	b := make([]byte, crc32Len)
	binary.BigEndian.PutUint32(b, crc32.Checksum(p, crc32cTable))
	_, err = ws.Write(b)
	if err != nil {
		return err
	}

	// write payload
	_, err = ws.Write(p)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package dispatcher

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

	"github.com/gorilla/websocket"
	"github.com/satori/go.uuid"
)

type MessageData struct {
	Event   string      `json:"e"`
	System  bool        `json:"s"`
	Message interface{} `json:"m"`
}

type DispatcherConfig struct {
	ReadBufferSize  int
	WriteBufferSize int
}

type Dispatcher struct {
	config   *DispatcherConfig
	upgrader websocket.Upgrader

	handlers map[string]map[string]Handler
	rooms    map[string]map[string]*Conn
}

var defaultDispatcherConfig = &DispatcherConfig{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

//NewDispatcher create dispatcher
func NewDispatcher(config *DispatcherConfig) *Dispatcher {

	if config == nil {
		config = defaultDispatcherConfig
	}

	d := &Dispatcher{
		upgrader: websocket.Upgrader{
			ReadBufferSize:  config.ReadBufferSize,
			WriteBufferSize: config.WriteBufferSize,
		},
		handlers: map[string]map[string]Handler{},
		rooms:    map[string]map[string]*Conn{},
	}

	return d
}

//Handle upgrades http connection to websocket connection
func (d *Dispatcher) Handle(w http.ResponseWriter, r *http.Request, h *http.Header) (*Conn, error) {

	var c *websocket.Conn
	var err error

	if h == nil {
		c, err = d.upgrader.Upgrade(w, r, nil)
	} else {
		c, err = d.upgrader.Upgrade(w, r, *h)
	}

	if err != nil {
		return nil, err
	}

	id := uuid.NewV4().String()

	conn := &Conn{
		dispatcher: d,
		conn:       c,
		ID:         id,
	}

	// ecah connection has it's seperate room with a room name as a connection ID
	d.rooms[id] = map[string]*Conn{}
	d.rooms[id][id] = conn

	d.handlers[id] = map[string]Handler{}

	go d.readMessages(id)

	return conn, nil
}

func (d *Dispatcher) readMessages(connectionID string) {

	c := d.rooms[connectionID][connectionID]

	for {
		_, message, err := c.conn.ReadMessage()
		if err != nil {
			log.Println("read:", err)
			break
		}

		msg := &MessageData{}
		json.Unmarshal(message, msg)
		if err != nil {
			log.Println("MessageData unmarshal:", err)
			break
		}

		// handle system events
		if msg.System {

			if msg.Event == "join" {
				d.join(c, msg.Message.(string))
			}

			if msg.Event == "leave" {
				d.leave(c, msg.Message.(string))
			}

			//all the system event handling ends here
			continue
		}

		if handler, ok := d.handlers[c.ID][msg.Event]; ok {
			handler(c.conn, msg.Message)
		}
	}

	//close connection if message reading stoped
	d.Close(connectionID)
}

//EmitTo send message to all connections in the room
func (d *Dispatcher) EmitTo(room string, event string, msg interface{}) error {

	connections, ok := d.rooms[room]
	if ok == false {
		return fmt.Errorf("Romm with %s name doesn't exist.", room)
	}

	// send message to each connection in the room
	for _, conn := range connections {
		err := conn.Emit(event, msg)
		if err != nil {
			return err
		}
	}

	return nil
}

//Close connection
func (d *Dispatcher) Close(ConnectionID string) error {

	conn, ok := d.rooms[ConnectionID][ConnectionID]
	if ok == false {
		return fmt.Errorf("Connection with '%s' ID not found", ConnectionID)
	}

	//todo: remove connection from each room to which it belongs to

	delete(d.rooms, ConnectionID)
	delete(d.handlers, ConnectionID)
	conn.conn.Close()

	return nil
}

func (d *Dispatcher) join(conn *Conn, roomName string) {

	_, ok := d.rooms[roomName]
	if ok == false {
		d.rooms[roomName] = map[string]*Conn{}
	}

	d.rooms[roomName][conn.ID] = conn

	//todo: add room name to conneciton object

}

func (d *Dispatcher) leave(conn *Conn, roomName string) {

	_, ok := d.rooms[roomName]
	if ok == false {
		return // can't find room, so there is nothing to leve
	}

	// check if room has more than one connection
	if len(d.rooms[roomName]) > 1 {
		// remove connection from the room
		delete(d.rooms[roomName], conn.ID)
	} else {
		// room has only one connection so we will remove the room
		delete(d.rooms, roomName)
	}

	//todo: remove rooms name from connection
}
<commit_msg>Add emitTo event handling<commit_after>package dispatcher

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

	"github.com/gorilla/websocket"
	"github.com/satori/go.uuid"
)

type MessageData struct {
	Event   string      `json:"e"`
	System  bool        `json:"s"`
	Message interface{} `json:"m"`
}

type DispatcherConfig struct {
	ReadBufferSize  int
	WriteBufferSize int
}

type Dispatcher struct {
	config   *DispatcherConfig
	upgrader websocket.Upgrader

	handlers map[string]map[string]Handler
	rooms    map[string]map[string]*Conn
}

var defaultDispatcherConfig = &DispatcherConfig{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

//NewDispatcher create dispatcher
func NewDispatcher(config *DispatcherConfig) *Dispatcher {

	if config == nil {
		config = defaultDispatcherConfig
	}

	d := &Dispatcher{
		upgrader: websocket.Upgrader{
			ReadBufferSize:  config.ReadBufferSize,
			WriteBufferSize: config.WriteBufferSize,
		},
		handlers: map[string]map[string]Handler{},
		rooms:    map[string]map[string]*Conn{},
	}

	return d
}

//Handle upgrades http connection to websocket connection
func (d *Dispatcher) Handle(w http.ResponseWriter, r *http.Request, h *http.Header) (*Conn, error) {

	var c *websocket.Conn
	var err error

	if h == nil {
		c, err = d.upgrader.Upgrade(w, r, nil)
	} else {
		c, err = d.upgrader.Upgrade(w, r, *h)
	}

	if err != nil {
		return nil, err
	}

	id := uuid.NewV4().String()

	conn := &Conn{
		dispatcher: d,
		conn:       c,
		ID:         id,
	}

	// ecah connection has it's seperate room with a room name as a connection ID
	d.rooms[id] = map[string]*Conn{}
	d.rooms[id][id] = conn

	d.handlers[id] = map[string]Handler{}

	go d.readMessages(id)

	return conn, nil
}

func (d *Dispatcher) readMessages(connectionID string) {

	c := d.rooms[connectionID][connectionID]

	for {
		_, message, err := c.conn.ReadMessage()
		if err != nil {
			log.Println("read:", err)
			break
		}

		msg := &MessageData{}
		json.Unmarshal(message, msg)
		if err != nil {
			log.Println("MessageData unmarshal:", err)
			break
		}

		// handle system events
		if msg.System {

			if msg.Event == "join" {

				d.join(c, msg.Message.(string))
			} else if msg.Event == "leave" {

				d.leave(c, msg.Message.(string))
			} else if msg.Event == "emitTo" {

				m := msg.Message.(map[string]interface{})

				room := m["r"].(string)
				event := m["e"].(string)

				err := d.EmitTo(room, event, m["m"])
				if err != nil {
					log.Println(err)
				}
			}

			//all the system event handling ends here
			continue
		}

		if handler, ok := d.handlers[c.ID][msg.Event]; ok {
			handler(c.conn, msg.Message)
		}
	}

	//close connection if message reading stoped
	d.Close(connectionID)
}

//EmitTo send message to all connections in the room
func (d *Dispatcher) EmitTo(room string, event string, msg interface{}) error {

	connections, ok := d.rooms[room]
	if ok == false {
		return fmt.Errorf("Room with name '%s' doesn't exist.", room)
	}

	// send message to each connection in the room
	for _, conn := range connections {
		err := conn.Emit(event, msg)
		if err != nil {
			return err
		}
	}

	return nil
}

//Close connection
func (d *Dispatcher) Close(ConnectionID string) error {

	conn, ok := d.rooms[ConnectionID][ConnectionID]
	if ok == false {
		return fmt.Errorf("Connection with '%s' ID not found", ConnectionID)
	}

	//todo: remove connection from each room to which it belongs to

	delete(d.rooms, ConnectionID)
	delete(d.handlers, ConnectionID)
	conn.conn.Close()

	return nil
}

func (d *Dispatcher) join(conn *Conn, roomName string) {

	_, ok := d.rooms[roomName]
	if ok == false {
		d.rooms[roomName] = map[string]*Conn{}
	}

	d.rooms[roomName][conn.ID] = conn

	//todo: add room name to conneciton object

}

func (d *Dispatcher) leave(conn *Conn, roomName string) {

	_, ok := d.rooms[roomName]
	if ok == false {
		return // can't find room, so there is nothing to leve
	}

	// check if room has more than one connection
	if len(d.rooms[roomName]) > 1 {
		// remove connection from the room
		delete(d.rooms[roomName], conn.ID)
	} else {
		// room has only one connection so we will remove the room
		delete(d.rooms, roomName)
	}

	//todo: remove rooms name from connection
}
<|endoftext|>


<commit_before>package net4g

import (
	"github.com/carsonsx/log4g"
	"reflect"
	"runtime/debug"
	"sync"
)

func Dispatch(dispatchers []*dispatcher, req NetReq, res NetRes) {
	found := false
	for _, p := range dispatchers {
		if p.running {
			p.dispatchChan <- &dispatchData{req: req, res: res}
			found = true
		}
	}
	if !found {
		log4g.Warn("not found any running dispatcher")
	}
}

func AddHandler(dispatcher *dispatcher, v interface{}, h func(req NetReq, res NetRes)) {
	dispatcher.AddHandler(h, reflect.TypeOf(v))
}

func NewDispatcher(name string) *dispatcher {
	p := new(dispatcher)
	p.Name = name
	p.dispatchChan = make(chan *dispatchData, 1000)
	p.sessionClosedChan = make(chan NetSession, 100)
	p.destroyChan = make(chan bool, 1)
	p.typeHandlers = make(map[reflect.Type]func(req NetReq, res NetRes))
	p.run()
	log4g.Info("new a %s dispatcher", name)
	return p
}

type dispatchData struct {
	req NetReq
	res NetRes
}

type dispatcher struct {
	Name                    string
	serializer              Serializer
	mgr                     *NetManager
	globalHandlers          []func(req NetReq, res NetRes)
	typeHandlers            map[reflect.Type]func(req NetReq, res NetRes)
	before_interceptors     []func(req NetReq, res NetRes)
	after_interceptors      []func(req NetReq, res NetRes)
	dispatchChan            chan *dispatchData
	sessionClosedChan       chan NetSession
	connectionClosedHandler func(session NetSession)
	destroyChan             chan bool
	destroyHandler          func()
	running                 bool
	wg                      sync.WaitGroup
}

func (p *dispatcher) AddHandler(h func(req NetReq, res NetRes), t ...reflect.Type) {
	if len(t) > 0 {
		p.typeHandlers[t[0]] = h
		log4g.Info("added a handler for %v", t[0])
	} else {
		p.globalHandlers = append(p.globalHandlers, h)
		log4g.Info("added global handler")
	}
}

func (p *dispatcher) OnConnectionClosed(h func(session NetSession)) {
	p.connectionClosedHandler = h
}

func (p *dispatcher) OnDestroy(h func()) {
	p.destroyHandler = h
}

func (p *dispatcher) dispatch(msg *dispatchData) {

	// safe the user handler to avoid the whole server down
	defer func() {
		if r := recover(); r != nil {
			log4g.Error("********************* Handler Panic *********************")
			log4g.Error(r)
			log4g.Error(string(debug.Stack()))
			msg.res.Close()
			log4g.Error("********************* Handler Panic *********************")
		}
	}()

	for _, i := range p.before_interceptors {
		i(msg.req, msg.res)
	}

	for _, h := range p.globalHandlers {
		h(msg.req, msg.res)
	}

	t := reflect.TypeOf(msg.req.Msg())
	if h, ok := p.typeHandlers[t]; ok {
		log4g.Trace("dispatcher[%s] is dispatching %v to handler ", p.Name, t)
		h(msg.req, msg.res)
	} else {
		log4g.Trace("dispatcher[%s] not found any handler ", p.Name)
	}

	for _, i := range p.after_interceptors {
		i(msg.req, msg.res)
	}
}

func (p *dispatcher) run() {
	p.wg.Add(1)
	// one dispatcher, one goroutine
	go func() {
		defer p.wg.Done()
	outer:
		for {
			select {
			case data := <-p.dispatchChan:
				p.dispatch(data)
			case session := <-p.sessionClosedChan:
				if p.connectionClosedHandler != nil {
					p.connectionClosedHandler(session)
				}
				session.Get("wg").(*sync.WaitGroup).Done()
			case <-p.destroyChan:
				if p.destroyHandler != nil {
					p.destroyHandler()
				}
				break outer
			}
		}
	}()

	p.running = true
}

func (p *dispatcher) Broadcast(v interface{}, filter func(session NetSession) bool) error {
	b, err := p.serializer.Serialize(v)
	if err != nil {
		log4g.Error(err)
		return err
	}
	p.mgr.Broadcast(b, filter)
	return nil
}

func (p *dispatcher) BroadcastAll(v interface{}) error {
	b, err := p.serializer.Serialize(v)
	if err != nil {
		log4g.Error(err)
		return err
	}
	p.mgr.BroadcastAll(b)
	return nil
}

func (p *dispatcher) BroadcastOthers(mySession NetSession, v interface{}) error {
	b, err := p.serializer.Serialize(v)
	if err != nil {
		log4g.Error(err)
		return err
	}
	p.mgr.BroadcastOthers(mySession, b)
	return nil
}

func (p *dispatcher) Someone(v interface{}, filter func(session NetSession) bool) error {
	b, err := p.serializer.Serialize(v)
	if err != nil {
		log4g.Error(err)
		return err
	}
	p.mgr.Someone(b, filter)
	return nil
}

func (p *dispatcher) handleConnectionClosed(session NetSession) {
	wg := new(sync.WaitGroup)
	wg.Add(1)
	session.Set("wg", wg)
	p.sessionClosedChan <- session
	wg.Wait()
	session.Remove("wg")
}

func (p *dispatcher) Destroy() {
	p.running = false
	p.destroyChan <- true
	p.wg.Wait()
	//how to close gracefully
	//close(p.dispatchChan)
	//close(p.sessionClosedChan)
	close(p.destroyChan)
}
<commit_msg>fix dispather close issue<commit_after>package net4g

import (
	"github.com/carsonsx/log4g"
	"reflect"
	"runtime/debug"
	"sync"
)

func Dispatch(dispatchers []*dispatcher, req NetReq, res NetRes) {
	found := false
	for _, p := range dispatchers {
		if p.running {
			p.dispatchChan <- &dispatchData{req: req, res: res}
			found = true
		}
	}
	if !found {
		log4g.Warn("not found any running dispatcher")
	}
}

func AddHandler(dispatcher *dispatcher, v interface{}, h func(req NetReq, res NetRes)) {
	dispatcher.AddHandler(h, reflect.TypeOf(v))
}

func NewDispatcher(name string) *dispatcher {
	p := new(dispatcher)
	p.Name = name
	p.dispatchChan = make(chan *dispatchData, 1000)
	p.sessionClosedChan = make(chan NetSession, 100)
	p.destroyChan = make(chan bool, 1)
	p.typeHandlers = make(map[reflect.Type]func(req NetReq, res NetRes))
	p.run()
	log4g.Info("new a %s dispatcher", name)
	return p
}

type dispatchData struct {
	req NetReq
	res NetRes
}

type dispatcher struct {
	Name                    string
	serializer              Serializer
	mgr                     *NetManager
	globalHandlers          []func(req NetReq, res NetRes)
	typeHandlers            map[reflect.Type]func(req NetReq, res NetRes)
	before_interceptors     []func(req NetReq, res NetRes)
	after_interceptors      []func(req NetReq, res NetRes)
	dispatchChan            chan *dispatchData
	sessionClosedChan       chan NetSession
	connectionClosedHandler func(session NetSession)
	destroyChan             chan bool
	destroyHandler          func()
	running                 bool
	wg                      sync.WaitGroup
}

func (p *dispatcher) AddHandler(h func(req NetReq, res NetRes), t ...reflect.Type) {
	if len(t) > 0 {
		p.typeHandlers[t[0]] = h
		log4g.Info("added a handler for %v", t[0])
	} else {
		p.globalHandlers = append(p.globalHandlers, h)
		log4g.Info("added global handler")
	}
}

func (p *dispatcher) OnConnectionClosed(h func(session NetSession)) {
	p.connectionClosedHandler = h
}

func (p *dispatcher) OnDestroy(h func()) {
	p.destroyHandler = h
}

func (p *dispatcher) dispatch(msg *dispatchData) {

	// safe the user handler to avoid the whole server down
	defer func() {
		if r := recover(); r != nil {
			log4g.Error("********************* Handler Panic *********************")
			log4g.Error(r)
			log4g.Error(string(debug.Stack()))
			msg.res.Close()
			log4g.Error("********************* Handler Panic *********************")
		}
	}()

	for _, i := range p.before_interceptors {
		i(msg.req, msg.res)
	}

	for _, h := range p.globalHandlers {
		h(msg.req, msg.res)
	}

	t := reflect.TypeOf(msg.req.Msg())
	if h, ok := p.typeHandlers[t]; ok {
		log4g.Trace("dispatcher[%s] is dispatching %v to handler ", p.Name, t)
		h(msg.req, msg.res)
	} else {
		log4g.Trace("dispatcher[%s] not found any handler ", p.Name)
	}

	for _, i := range p.after_interceptors {
		i(msg.req, msg.res)
	}
}

func (p *dispatcher) run() {
	p.wg.Add(1)
	// one dispatcher, one goroutine
	go func() {
		defer p.wg.Done()
	outer:
		for {
			select {
			case data := <-p.dispatchChan:
				p.dispatch(data)
			case session := <-p.sessionClosedChan:
				if p.connectionClosedHandler != nil {
					p.connectionClosedHandler(session)
				}
				session.Get("wg").(*sync.WaitGroup).Done()
			case <-p.destroyChan:
				if p.destroyHandler != nil {
					p.destroyHandler()
				}
				break outer
			}
		}
	}()

	p.running = true
}

func (p *dispatcher) Broadcast(v interface{}, filter func(session NetSession) bool) error {
	b, err := p.serializer.Serialize(v)
	if err != nil {
		log4g.Error(err)
		return err
	}
	p.mgr.Broadcast(b, filter)
	return nil
}

func (p *dispatcher) BroadcastAll(v interface{}) error {
	b, err := p.serializer.Serialize(v)
	if err != nil {
		log4g.Error(err)
		return err
	}
	p.mgr.BroadcastAll(b)
	return nil
}

func (p *dispatcher) BroadcastOthers(mySession NetSession, v interface{}) error {
	b, err := p.serializer.Serialize(v)
	if err != nil {
		log4g.Error(err)
		return err
	}
	p.mgr.BroadcastOthers(mySession, b)
	return nil
}

func (p *dispatcher) Someone(v interface{}, filter func(session NetSession) bool) error {
	b, err := p.serializer.Serialize(v)
	if err != nil {
		log4g.Error(err)
		return err
	}
	p.mgr.Someone(b, filter)
	return nil
}

func (p *dispatcher) handleConnectionClosed(session NetSession) {
	wg := new(sync.WaitGroup)
	wg.Add(1)
	session.Set("wg", wg)
	p.sessionClosedChan <- session
	wg.Wait()
	session.Remove("wg")
}

func (p *dispatcher) Destroy() {
	if p.running {
		p.destroyChan <- true
	}
	p.running = false
	p.wg.Wait()
	//how to close gracefully
	//close(p.dispatchChan)
	//close(p.sessionClosedChan)
	//close(p.destroyChan)
}
<|endoftext|>


<commit_before>package manifest

import (
	"bytes"
	"compress/gzip"
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"sync"
	"time"

	"github.com/kopia/kopia/block"
	"github.com/kopia/kopia/internal/kopialogging"
)

var log = kopialogging.Logger("kopia/manifest")

// ErrNotFound is returned when the metadata item is not found.
var ErrNotFound = errors.New("not found")

const manifestBlockPrefix = "m"
const autoCompactionBlockCount = 16

// Manager organizes JSON manifests of various kinds, including snapshot manifests
type Manager struct {
	mu sync.Mutex
	b  *block.Manager

	pendingEntries map[string]*manifestEntry

	committedEntries  map[string]*manifestEntry
	committedBlockIDs map[string]bool
}

// Put serializes the provided payload to JSON and persists it. Returns unique handle that represents the object.
func (m *Manager) Put(labels map[string]string, payload interface{}) (string, error) {
	if labels["type"] == "" {
		return "", fmt.Errorf("'type' label is required")
	}
	m.mu.Lock()
	defer m.mu.Unlock()

	random := make([]byte, 16)
	if _, err := rand.Read(random); err != nil {
		return "", fmt.Errorf("can't initialize randomness: %v", err)
	}

	b, err := json.Marshal(payload)
	if err != nil {
		return "", err
	}

	e := &manifestEntry{
		ID:      hex.EncodeToString(random),
		ModTime: time.Now().UTC(),
		Labels:  copyLabels(labels),
		Content: b,
	}

	m.pendingEntries[e.ID] = e

	return e.ID, nil
}

// GetMetadata returns metadata about provided manifest item or ErrNotFound if the item can't be found.
func (m *Manager) GetMetadata(id string) (*EntryMetadata, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	e := m.pendingEntries[id]
	if e == nil {
		e = m.committedEntries[id]
	}

	if e == nil || e.Deleted {
		return nil, ErrNotFound
	}

	return &EntryMetadata{
		ID:      id,
		ModTime: e.ModTime,
		Length:  len(e.Content),
		Labels:  copyLabels(e.Labels),
	}, nil
}

// Get retrieves the contents of the provided manifest item by deserializing it as JSON to provided object.
// If the manifest is not found, returns ErrNotFound.
func (m *Manager) Get(id string, data interface{}) error {
	b, err := m.GetRaw(id)
	if err != nil {
		return err
	}

	if err := json.Unmarshal(b, data); err != nil {
		return fmt.Errorf("unable to unmashal %q: %v", id, err)
	}

	return nil
}

// GetRaw returns raw contents of the provided manifest (JSON bytes) or ErrNotFound if not found.
func (m *Manager) GetRaw(id string) ([]byte, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	e := m.pendingEntries[id]
	if e == nil {
		e = m.committedEntries[id]
	}
	if e == nil || e.Deleted {
		return nil, ErrNotFound
	}

	return e.Content, nil
}

// Find returns the list of EntryMetadata for manifest entries matching all provided labels.
func (m *Manager) Find(labels map[string]string) []*EntryMetadata {
	m.mu.Lock()
	defer m.mu.Unlock()

	var matches []*EntryMetadata
	for _, e := range m.pendingEntries {
		if matchesLabels(e.Labels, labels) {
			matches = append(matches, cloneEntryMetadata(e))
		}
	}
	for _, e := range m.committedEntries {
		if m.pendingEntries[e.ID] != nil {
			// ignore committed that are also in pending
			continue
		}

		if matchesLabels(e.Labels, labels) {
			matches = append(matches, cloneEntryMetadata(e))
		}
	}

	sort.Slice(matches, func(i, j int) bool {
		return matches[i].ModTime.Before(matches[j].ModTime)
	})
	return matches
}

func cloneEntryMetadata(e *manifestEntry) *EntryMetadata {
	return &EntryMetadata{
		ID:      e.ID,
		Labels:  copyLabels(e.Labels),
		Length:  len(e.Content),
		ModTime: e.ModTime,
	}
}

// matchesLabels returns true when all entries in 'b' are found in the 'a'.
func matchesLabels(a, b map[string]string) bool {
	for k, v := range b {
		if a[k] != v {
			return false
		}
	}

	return true
}

// Flush persists changes to manifest manager.
func (m *Manager) Flush(ctx context.Context) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	_, err := m.flushPendingEntriesLocked(ctx)
	return err
}

func (m *Manager) flushPendingEntriesLocked(ctx context.Context) (string, error) {
	if len(m.pendingEntries) == 0 {
		return "", nil
	}

	man := manifest{}

	for _, e := range m.pendingEntries {
		man.Entries = append(man.Entries, e)
	}

	var buf bytes.Buffer

	gz := gzip.NewWriter(&buf)
	if err := json.NewEncoder(gz).Encode(man); err != nil {
		return "", fmt.Errorf("unable to marshal: %v", err)
	}
	if err := gz.Flush(); err != nil {
		return "", fmt.Errorf("unable to flush: %v", err)
	}
	if err := gz.Close(); err != nil {
		return "", fmt.Errorf("unable to close: %v", err)
	}

	blockID, err := m.b.WriteBlock(ctx, buf.Bytes(), manifestBlockPrefix)
	if err != nil {
		return "", err
	}

	for _, e := range m.pendingEntries {
		m.committedEntries[e.ID] = e
		delete(m.pendingEntries, e.ID)
	}

	m.committedBlockIDs[blockID] = true

	return blockID, nil
}

// Delete marks the specified manifest ID for deletion.
func (m *Manager) Delete(id string) {
	if m.pendingEntries[id] == nil && m.committedEntries[id] == nil {
		return
	}

	m.pendingEntries[id] = &manifestEntry{
		ID:      id,
		ModTime: time.Now().UTC(),
		Deleted: true,
	}
}

// Refresh updates the committed blocks from the underlying storage.
func (m *Manager) Refresh(ctx context.Context) error {
	return m.loadCommittedBlocks(ctx)
}

func (m *Manager) loadCommittedBlocks(ctx context.Context) error {
	log.Debugf("listing manifest blocks")
	blocks, err := m.b.ListBlocks(manifestBlockPrefix)
	if err != nil {
		return fmt.Errorf("unable to list manifest blocks: %v", err)
	}

	m.mu.Lock()
	defer m.mu.Unlock()

	m.committedEntries = map[string]*manifestEntry{}
	m.committedBlockIDs = map[string]bool{}

	log.Debugf("found %v manifest blocks", len(blocks))

	if err := m.loadManifestBlocks(ctx, blocks); err != nil {
		return fmt.Errorf("unable to load manifest blocks: %v", err)
	}

	if err := m.maybeCompactLocked(ctx); err != nil {
		return fmt.Errorf("error auto-compacting blocks")
	}

	return nil
}

func (m *Manager) loadManifestBlocks(ctx context.Context, blockIDs []string) error {
	t0 := time.Now()

	for _, b := range blockIDs {
		m.committedBlockIDs[b] = true
	}

	manifests, err := m.loadBlocksInParallel(ctx, blockIDs)
	if err != nil {
		return err
	}

	for _, man := range manifests {
		for _, e := range man.Entries {
			m.mergeEntry(e)
		}
	}

	// after merging, remove blocks marked as deleted.
	for k, e := range m.committedEntries {
		if e.Deleted {
			delete(m.committedEntries, k)
		}
	}

	log.Debugf("finished loading manifest blocks in %v.", time.Since(t0))

	return nil
}

func (m *Manager) loadBlocksInParallel(ctx context.Context, blockIDs []string) ([]manifest, error) {
	errors := make(chan error, len(blockIDs))
	manifests := make(chan manifest, len(blockIDs))
	ch := make(chan string, len(blockIDs))
	var wg sync.WaitGroup

	for i := 0; i < 8; i++ {
		wg.Add(1)
		go func(workerID int) {
			defer wg.Done()

			for blk := range ch {
				t1 := time.Now()
				man, err := m.loadManifestBlock(ctx, blk)
				log.Debugf("block %v loaded by worker %v in %v.", blk, workerID, time.Since(t1))

				if err != nil {
					errors <- err
				} else {
					manifests <- man
				}
			}
		}(i)
	}

	// feed block IDs for goroutines
	for _, b := range blockIDs {
		ch <- b
	}
	close(ch)

	// wait for workers to complete
	wg.Wait()
	close(errors)
	close(manifests)

	// if there was any error, forward it
	if err := <-errors; err != nil {
		return nil, err
	}

	var man []manifest
	for m := range manifests {
		man = append(man, m)
	}

	return man, nil
}

func (m *Manager) loadManifestBlock(ctx context.Context, blockID string) (manifest, error) {
	man := manifest{}
	blk, err := m.b.GetBlock(ctx, blockID)
	if err != nil {
		return man, fmt.Errorf("unable to read block %q: %v", blockID, err)
	}

	if len(blk) > 2 && blk[0] == '{' {
		if err := json.Unmarshal(blk, &man); err != nil {
			return man, fmt.Errorf("unable to parse block %q: %v", blockID, err)
		}
	} else {
		gz, err := gzip.NewReader(bytes.NewReader(blk))
		if err != nil {
			return man, fmt.Errorf("unable to unpack block %q: %v", blockID, err)
		}

		if err := json.NewDecoder(gz).Decode(&man); err != nil {
			return man, fmt.Errorf("unable to parse block %q: %v", blockID, err)
		}
	}

	return man, nil
}

// Compact performs compaction of manifest blocks.
func (m *Manager) Compact(ctx context.Context) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	return m.compactLocked(ctx)
}

func (m *Manager) maybeCompactLocked(ctx context.Context) error {
	if len(m.committedBlockIDs) < autoCompactionBlockCount {
		return nil
	}

	log.Debugf("performing automatic compaction of %v blocks", len(m.committedBlockIDs))
	if err := m.compactLocked(ctx); err != nil {
		return fmt.Errorf("unable to compact manifest blocks: %v", err)
	}

	if err := m.b.Flush(ctx); err != nil {
		return fmt.Errorf("unable to flush blocks after auto-compaction: %v", err)
	}

	return nil
}

func (m *Manager) compactLocked(ctx context.Context) error {
	log.Debugf("compactLocked: pendingEntries=%v blockIDs=%v", len(m.pendingEntries), len(m.committedBlockIDs))

	if len(m.committedBlockIDs) == 1 && len(m.pendingEntries) == 0 {
		return nil
	}

	// compaction needs to be atomic (deletes and rewrite should show up in one index block or not show up at all)
	// that's why we want to prevent index flushes while we're d.
	m.b.DisableIndexFlush()
	defer m.b.EnableIndexFlush()

	for _, e := range m.committedEntries {
		m.pendingEntries[e.ID] = e
	}

	blockID, err := m.flushPendingEntriesLocked(ctx)
	if err != nil {
		return err
	}

	// add the newly-created block to the list, could be duplicate
	for b := range m.committedBlockIDs {
		if b == blockID {
			// do not delete block that was just written.
			continue
		}

		if err := m.b.DeleteBlock(b); err != nil {
			return fmt.Errorf("unable to delete block %q: %v", b, err)
		}

		delete(m.committedBlockIDs, b)
	}

	return nil
}

func (m *Manager) mergeEntry(e *manifestEntry) {
	prev := m.committedEntries[e.ID]
	if prev == nil {
		m.committedEntries[e.ID] = e
		return
	}

	if e.ModTime.After(prev.ModTime) {
		m.committedEntries[e.ID] = e
	}
}

func copyLabels(m map[string]string) map[string]string {
	r := map[string]string{}
	for k, v := range m {
		r[k] = v
	}
	return r
}

// NewManager returns new manifest manager for the provided block manager.
func NewManager(ctx context.Context, b *block.Manager) (*Manager, error) {
	m := &Manager{
		b:                 b,
		pendingEntries:    map[string]*manifestEntry{},
		committedEntries:  map[string]*manifestEntry{},
		committedBlockIDs: map[string]bool{},
	}

	if err := m.loadCommittedBlocks(ctx); err != nil {
		return nil, err
	}

	return m, nil
}
<commit_msg>Fixed a race condition between two manifest managers, where the first one compacts blocks and second one is in the middle of listing and loading all blocks.<commit_after>package manifest

import (
	"bytes"
	"compress/gzip"
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"sync"
	"time"

	"github.com/kopia/kopia/block"
	"github.com/kopia/kopia/internal/kopialogging"
	"github.com/kopia/kopia/storage"
)

var log = kopialogging.Logger("kopia/manifest")

// ErrNotFound is returned when the metadata item is not found.
var ErrNotFound = errors.New("not found")

const manifestBlockPrefix = "m"
const autoCompactionBlockCount = 16

// Manager organizes JSON manifests of various kinds, including snapshot manifests
type Manager struct {
	mu sync.Mutex
	b  *block.Manager

	pendingEntries map[string]*manifestEntry

	committedEntries  map[string]*manifestEntry
	committedBlockIDs map[string]bool
}

// Put serializes the provided payload to JSON and persists it. Returns unique handle that represents the object.
func (m *Manager) Put(labels map[string]string, payload interface{}) (string, error) {
	if labels["type"] == "" {
		return "", fmt.Errorf("'type' label is required")
	}
	m.mu.Lock()
	defer m.mu.Unlock()

	random := make([]byte, 16)
	if _, err := rand.Read(random); err != nil {
		return "", fmt.Errorf("can't initialize randomness: %v", err)
	}

	b, err := json.Marshal(payload)
	if err != nil {
		return "", err
	}

	e := &manifestEntry{
		ID:      hex.EncodeToString(random),
		ModTime: time.Now().UTC(),
		Labels:  copyLabels(labels),
		Content: b,
	}

	m.pendingEntries[e.ID] = e

	return e.ID, nil
}

// GetMetadata returns metadata about provided manifest item or ErrNotFound if the item can't be found.
func (m *Manager) GetMetadata(id string) (*EntryMetadata, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	e := m.pendingEntries[id]
	if e == nil {
		e = m.committedEntries[id]
	}

	if e == nil || e.Deleted {
		return nil, ErrNotFound
	}

	return &EntryMetadata{
		ID:      id,
		ModTime: e.ModTime,
		Length:  len(e.Content),
		Labels:  copyLabels(e.Labels),
	}, nil
}

// Get retrieves the contents of the provided manifest item by deserializing it as JSON to provided object.
// If the manifest is not found, returns ErrNotFound.
func (m *Manager) Get(id string, data interface{}) error {
	b, err := m.GetRaw(id)
	if err != nil {
		return err
	}

	if err := json.Unmarshal(b, data); err != nil {
		return fmt.Errorf("unable to unmashal %q: %v", id, err)
	}

	return nil
}

// GetRaw returns raw contents of the provided manifest (JSON bytes) or ErrNotFound if not found.
func (m *Manager) GetRaw(id string) ([]byte, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	e := m.pendingEntries[id]
	if e == nil {
		e = m.committedEntries[id]
	}
	if e == nil || e.Deleted {
		return nil, ErrNotFound
	}

	return e.Content, nil
}

// Find returns the list of EntryMetadata for manifest entries matching all provided labels.
func (m *Manager) Find(labels map[string]string) []*EntryMetadata {
	m.mu.Lock()
	defer m.mu.Unlock()

	var matches []*EntryMetadata
	for _, e := range m.pendingEntries {
		if matchesLabels(e.Labels, labels) {
			matches = append(matches, cloneEntryMetadata(e))
		}
	}
	for _, e := range m.committedEntries {
		if m.pendingEntries[e.ID] != nil {
			// ignore committed that are also in pending
			continue
		}

		if matchesLabels(e.Labels, labels) {
			matches = append(matches, cloneEntryMetadata(e))
		}
	}

	sort.Slice(matches, func(i, j int) bool {
		return matches[i].ModTime.Before(matches[j].ModTime)
	})
	return matches
}

func cloneEntryMetadata(e *manifestEntry) *EntryMetadata {
	return &EntryMetadata{
		ID:      e.ID,
		Labels:  copyLabels(e.Labels),
		Length:  len(e.Content),
		ModTime: e.ModTime,
	}
}

// matchesLabels returns true when all entries in 'b' are found in the 'a'.
func matchesLabels(a, b map[string]string) bool {
	for k, v := range b {
		if a[k] != v {
			return false
		}
	}

	return true
}

// Flush persists changes to manifest manager.
func (m *Manager) Flush(ctx context.Context) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	_, err := m.flushPendingEntriesLocked(ctx)
	return err
}

func (m *Manager) flushPendingEntriesLocked(ctx context.Context) (string, error) {
	if len(m.pendingEntries) == 0 {
		return "", nil
	}

	man := manifest{}

	for _, e := range m.pendingEntries {
		man.Entries = append(man.Entries, e)
	}

	var buf bytes.Buffer

	gz := gzip.NewWriter(&buf)
	if err := json.NewEncoder(gz).Encode(man); err != nil {
		return "", fmt.Errorf("unable to marshal: %v", err)
	}
	if err := gz.Flush(); err != nil {
		return "", fmt.Errorf("unable to flush: %v", err)
	}
	if err := gz.Close(); err != nil {
		return "", fmt.Errorf("unable to close: %v", err)
	}

	blockID, err := m.b.WriteBlock(ctx, buf.Bytes(), manifestBlockPrefix)
	if err != nil {
		return "", err
	}

	for _, e := range m.pendingEntries {
		m.committedEntries[e.ID] = e
		delete(m.pendingEntries, e.ID)
	}

	m.committedBlockIDs[blockID] = true

	return blockID, nil
}

// Delete marks the specified manifest ID for deletion.
func (m *Manager) Delete(id string) {
	if m.pendingEntries[id] == nil && m.committedEntries[id] == nil {
		return
	}

	m.pendingEntries[id] = &manifestEntry{
		ID:      id,
		ModTime: time.Now().UTC(),
		Deleted: true,
	}
}

// Refresh updates the committed blocks from the underlying storage.
func (m *Manager) Refresh(ctx context.Context) error {
	return m.loadCommittedBlocks(ctx)
}

func (m *Manager) loadCommittedBlocks(ctx context.Context) error {
	log.Debugf("listing manifest blocks")
	m.mu.Lock()
	defer m.mu.Unlock()

	for {
		blocks, err := m.b.ListBlocks(manifestBlockPrefix)
		if err != nil {
			return fmt.Errorf("unable to list manifest blocks: %v", err)
		}

		m.committedEntries = map[string]*manifestEntry{}
		m.committedBlockIDs = map[string]bool{}

		log.Debugf("found %v manifest blocks", len(blocks))
		err = m.loadManifestBlocks(ctx, blocks)
		if err == nil {
			// success
			break
		}
		if err == storage.ErrBlockNotFound {
			// try again, lost a race with another manifest manager which just did compaction
			continue
		}
		return fmt.Errorf("unable to load manifest blocks: %v", err)
	}

	if err := m.maybeCompactLocked(ctx); err != nil {
		return fmt.Errorf("error auto-compacting blocks")
	}

	return nil
}

func (m *Manager) loadManifestBlocks(ctx context.Context, blockIDs []string) error {
	t0 := time.Now()

	for _, b := range blockIDs {
		m.committedBlockIDs[b] = true
	}

	manifests, err := m.loadBlocksInParallel(ctx, blockIDs)
	if err != nil {
		return err
	}

	for _, man := range manifests {
		for _, e := range man.Entries {
			m.mergeEntry(e)
		}
	}

	// after merging, remove blocks marked as deleted.
	for k, e := range m.committedEntries {
		if e.Deleted {
			delete(m.committedEntries, k)
		}
	}

	log.Debugf("finished loading manifest blocks in %v.", time.Since(t0))

	return nil
}

func (m *Manager) loadBlocksInParallel(ctx context.Context, blockIDs []string) ([]manifest, error) {
	errors := make(chan error, len(blockIDs))
	manifests := make(chan manifest, len(blockIDs))
	ch := make(chan string, len(blockIDs))
	var wg sync.WaitGroup

	for i := 0; i < 8; i++ {
		wg.Add(1)
		go func(workerID int) {
			defer wg.Done()

			for blk := range ch {
				t1 := time.Now()
				man, err := m.loadManifestBlock(ctx, blk)

				if err != nil {
					errors <- err
					log.Debugf("block %v failed to be loaded by worker %v in %v: %v.", blk, workerID, time.Since(t1), err)
				} else {
					log.Debugf("block %v loaded by worker %v in %v.", blk, workerID, time.Since(t1))
					manifests <- man
				}
			}
		}(i)
	}

	// feed block IDs for goroutines
	for _, b := range blockIDs {
		ch <- b
	}
	close(ch)

	// wait for workers to complete
	wg.Wait()
	close(errors)
	close(manifests)

	// if there was any error, forward it
	if err := <-errors; err != nil {
		return nil, err
	}

	var man []manifest
	for m := range manifests {
		man = append(man, m)
	}

	return man, nil
}

func (m *Manager) loadManifestBlock(ctx context.Context, blockID string) (manifest, error) {
	man := manifest{}
	blk, err := m.b.GetBlock(ctx, blockID)
	if err != nil {
		// do not wrap the error here, we want to propagate original ErrBlockNotFound
		// which causes a retry if we lose list/delete race.
		return man, err
	}

	if len(blk) > 2 && blk[0] == '{' {
		if err := json.Unmarshal(blk, &man); err != nil {
			return man, fmt.Errorf("unable to parse block %q: %v", blockID, err)
		}
	} else {
		gz, err := gzip.NewReader(bytes.NewReader(blk))
		if err != nil {
			return man, fmt.Errorf("unable to unpack block %q: %v", blockID, err)
		}

		if err := json.NewDecoder(gz).Decode(&man); err != nil {
			return man, fmt.Errorf("unable to parse block %q: %v", blockID, err)
		}
	}

	return man, nil
}

// Compact performs compaction of manifest blocks.
func (m *Manager) Compact(ctx context.Context) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	return m.compactLocked(ctx)
}

func (m *Manager) maybeCompactLocked(ctx context.Context) error {
	if len(m.committedBlockIDs) < autoCompactionBlockCount {
		return nil
	}

	log.Debugf("performing automatic compaction of %v blocks", len(m.committedBlockIDs))
	if err := m.compactLocked(ctx); err != nil {
		return fmt.Errorf("unable to compact manifest blocks: %v", err)
	}

	if err := m.b.Flush(ctx); err != nil {
		return fmt.Errorf("unable to flush blocks after auto-compaction: %v", err)
	}

	return nil
}

func (m *Manager) compactLocked(ctx context.Context) error {
	log.Debugf("compactLocked: pendingEntries=%v blockIDs=%v", len(m.pendingEntries), len(m.committedBlockIDs))

	if len(m.committedBlockIDs) == 1 && len(m.pendingEntries) == 0 {
		return nil
	}

	// compaction needs to be atomic (deletes and rewrite should show up in one index block or not show up at all)
	// that's why we want to prevent index flushes while we're d.
	m.b.DisableIndexFlush()
	defer m.b.EnableIndexFlush()

	for _, e := range m.committedEntries {
		m.pendingEntries[e.ID] = e
	}

	blockID, err := m.flushPendingEntriesLocked(ctx)
	if err != nil {
		return err
	}

	// add the newly-created block to the list, could be duplicate
	for b := range m.committedBlockIDs {
		if b == blockID {
			// do not delete block that was just written.
			continue
		}

		if err := m.b.DeleteBlock(b); err != nil {
			return fmt.Errorf("unable to delete block %q: %v", b, err)
		}

		delete(m.committedBlockIDs, b)
	}

	return nil
}

func (m *Manager) mergeEntry(e *manifestEntry) {
	prev := m.committedEntries[e.ID]
	if prev == nil {
		m.committedEntries[e.ID] = e
		return
	}

	if e.ModTime.After(prev.ModTime) {
		m.committedEntries[e.ID] = e
	}
}

func copyLabels(m map[string]string) map[string]string {
	r := map[string]string{}
	for k, v := range m {
		r[k] = v
	}
	return r
}

// NewManager returns new manifest manager for the provided block manager.
func NewManager(ctx context.Context, b *block.Manager) (*Manager, error) {
	m := &Manager{
		b:                 b,
		pendingEntries:    map[string]*manifestEntry{},
		committedEntries:  map[string]*manifestEntry{},
		committedBlockIDs: map[string]bool{},
	}

	if err := m.loadCommittedBlocks(ctx); err != nil {
		return nil, err
	}

	return m, nil
}
<|endoftext|>


<commit_before>// Copyright 2017 The Cockroach Authors.
//
// 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 apd

import (
	"bytes"
	"fmt"
	"math/rand"
	"testing"
)

// runBenches benchmarks a given function on random decimals on combinations of
// three parameters:
//
//   precision:    desired output precision
//   inScale:      the scale of the input decimal: the absolute value will be between
//                 10^inScale and 10^(inScale+1)
//   inNumDigits:  number of digits in the input decimal; if negative the number
//                 will be negative and the number of digits are the absolute value.
func runBenches(
	b *testing.B, precision, inScale, inNumDigits []int, fn func(*testing.B, *Context, *Decimal),
) {
	for _, p := range precision {
		ctx := BaseContext.WithPrecision(uint32(p))
		for _, s := range inScale {
			for _, d := range inNumDigits {
				numDigits := d
				negative := false
				if d < 0 {
					numDigits = -d
					negative = true
				}

				// Generate some random numbers with the given number of digits.
				nums := make([]Decimal, 20)
				for i := range nums {
					var buf bytes.Buffer
					if negative {
						buf.WriteByte('-')
					}
					buf.WriteByte('1' + byte(rand.Intn(9)))
					for j := 1; j < numDigits; j++ {
						buf.WriteByte('0' + byte(rand.Intn(10)))
					}
					if _, err := nums[i].SetString(buf.String()); err != nil {
						b.Fatal(err)
					}
					nums[i].SetExponent(int32(s - numDigits))
				}
				b.Run(
					fmt.Sprintf("P%d/S%d/D%d", p, s, d),
					func(b *testing.B) {
						for i := 0; i <= b.N; i++ {
							fn(b, ctx, &nums[i%len(nums)])
						}
					},
				)
			}
		}
	}
}

func BenchmarkExp(b *testing.B) {
	precision := []int{5, 10, 100}
	scale := []int{-5, -3, -2, -1, 1, 2}
	digits := []int{-100, -10, -2, 2, 10, 100}
	runBenches(
		b, precision, scale, digits,
		func(b *testing.B, ctx *Context, x *Decimal) {
			if _, err := ctx.Exp(&Decimal{}, x); err != nil {
				b.Fatal(err)
			}
		},
	)
}

func BenchmarkLn(b *testing.B) {
	precision := []int{2, 10, 100}
	scale := []int{100, -10, -2, 2, 10, 100}
	digits := []int{2, 10, 100}
	runBenches(
		b, precision, scale, digits,
		func(b *testing.B, ctx *Context, x *Decimal) {
			if _, err := ctx.Ln(&Decimal{}, x); err != nil {
				b.Fatal(err)
			}
		},
	)
}
<commit_msg>reduce # of benchmark cases<commit_after>// Copyright 2017 The Cockroach Authors.
//
// 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 apd

import (
	"bytes"
	"fmt"
	"math/rand"
	"testing"
)

// runBenches benchmarks a given function on random decimals on combinations of
// three parameters:
//
//   precision:    desired output precision
//   inScale:      the scale of the input decimal: the absolute value will be between
//                 10^inScale and 10^(inScale+1)
//   inNumDigits:  number of digits in the input decimal; if negative the number
//                 will be negative and the number of digits are the absolute value.
func runBenches(
	b *testing.B, precision, inScale, inNumDigits []int, fn func(*testing.B, *Context, *Decimal),
) {
	for _, p := range precision {
		ctx := BaseContext.WithPrecision(uint32(p))
		for _, s := range inScale {
			for _, d := range inNumDigits {
				numDigits := d
				negative := false
				if d < 0 {
					numDigits = -d
					negative = true
				}
				if numDigits > p {
					// Skip cases where we have more digits than the desired precision.
					continue
				}

				// Generate some random numbers with the given number of digits.
				nums := make([]Decimal, 20)
				for i := range nums {
					var buf bytes.Buffer
					if negative {
						buf.WriteByte('-')
					}
					buf.WriteByte('1' + byte(rand.Intn(9)))
					for j := 1; j < numDigits; j++ {
						buf.WriteByte('0' + byte(rand.Intn(10)))
					}
					if _, err := nums[i].SetString(buf.String()); err != nil {
						b.Fatal(err)
					}
					nums[i].SetExponent(int32(s - numDigits))
				}
				b.Run(
					fmt.Sprintf("P%d/S%d/D%d", p, s, d),
					func(b *testing.B) {
						for i := 0; i <= b.N; i++ {
							fn(b, ctx, &nums[i%len(nums)])
						}
					},
				)
			}
		}
	}
}

func BenchmarkExp(b *testing.B) {
	precision := []int{5, 10, 100}
	scale := []int{-4, -1, 2}
	digits := []int{-100, -10, -2, 2, 10, 100}
	runBenches(
		b, precision, scale, digits,
		func(b *testing.B, ctx *Context, x *Decimal) {
			if _, err := ctx.Exp(&Decimal{}, x); err != nil {
				b.Fatal(err)
			}
		},
	)
}

func BenchmarkLn(b *testing.B) {
	precision := []int{2, 10, 100}
	scale := []int{-100, -10, -2, 2, 10, 100}
	digits := []int{2, 10, 100}
	runBenches(
		b, precision, scale, digits,
		func(b *testing.B, ctx *Context, x *Decimal) {
			if _, err := ctx.Ln(&Decimal{}, x); err != nil {
				b.Fatal(err)
			}
		},
	)
}
<|endoftext|>


<commit_before>package dokugen

import (
	"strconv"
	"strings"
	"testing"
)

func TestCellCreation(t *testing.T) {

	number := 1
	data := strconv.Itoa(number)
	cell := NewCell(nil, 0, 0)
	cell.Load(data)
	if cell.Number() != number {
		t.Log("Number came back wrong")
		t.Fail()
	}
	if cell.Row != 0 {
		t.Log("Row came back wrong")
		t.Fail()
	}
	if cell.Col != 0 {
		t.Log("Cell came back wrong")
		t.Fail()
	}
	if cell.DataString() != data {
		t.Log("Cell round-tripped out with different string than data in")
		t.Fail()
	}
	//TODO: test failing for values that are too high.
	for i := 1; i <= DIM; i++ {
		if i == number {
			if !cell.Possible(i) {
				t.Log("We reported back a number we were explicitly set to was impossible")
				t.Fail()
			}
		} else if cell.Possible(i) {
			t.Log("We reported back that a number was possible when another number had been explicitly set")
			t.Fail()
		}
	}

	number = 2

	cell.SetNumber(number)

	if cell.Number() != number {
		t.Log("Number came back wrong after being set with SetNumber")
		t.Fail()
	}

	for i := 1; i <= DIM; i++ {
		if i == number {
			if !cell.Possible(i) {
				t.Log("We reported back a number we were explicitly set to was impossible (2nd set)")
				t.Fail()
			}
		} else if cell.Possible(i) {
			t.Log("We reported back that a number was possible when another number had been explicitly set (2nd set)")
			t.Fail()
		}
	}

	number = 0

	cell.SetNumber(number)

	if cell.Number() != number {
		t.Log("Number came back wrong after being set with SetNumber to 0")
		t.Fail()
	}

	for i := 1; i <= DIM; i++ {
		if !cell.Possible(i) {
			t.Log("We reported back a number was not possible when 0 had been explicitly set.")
			t.Fail()
		}
	}
	//TODO: test that a cell with nothing explicitly set operates like a counter for possibles.
}

func TestGridCreation(t *testing.T) {

	blockUpperLeftRow := make([]int, DIM)
	blockUpperLeftCol := make([]int, DIM)

	row := 0
	col := -1 * BLOCK_DIM

	for i := 0; i < DIM; i++ {
		col += BLOCK_DIM
		if col >= DIM {
			col = 0
			row += BLOCK_DIM
		}
		blockUpperLeftRow[i] = row
		blockUpperLeftCol[i] = col
	}

	cellData := "1"
	rowData := strings.Join(nCopies(cellData, DIM), COL_SEP)
	data := strings.Join(nCopies(rowData, DIM), ROW_SEP)
	grid := LoadGrid(data)
	if len(grid.cells) != DIM*DIM {
		t.Log("Didn't generate enough cells")
		t.Fail()
	}
	if grid.DataString() != data {
		t.Log("The grid round-tripped with different result than data in")
		t.Fail()
	}
	if grid.cells[10].Number() != 1 {
		t.Log("A random spot check of a cell had the wrong number: %s", grid.cells[10])
		t.Fail()
	}

	for count := 0; count < DIM; count++ {
		col := grid.Col(count)
		if num := len(col); num != DIM {
			t.Log("We got back a column but it had the wrong amount of items: ", num, "\n")
			t.Fail()
		}
		for i, cell := range col {
			if cell.Col != count {
				t.Log("One of the cells we got back when asking for column ", count, " was not in the right column.")
				t.Fail()
			}
			if cell.Row != i {
				t.Log("One of the cells we got back when asking for column ", count, " was not in the right row.")
				t.Fail()
			}
		}

		row := grid.Row(count)
		if len(row) != DIM {
			t.Log("We got back a row but it had the wrong number of items.")
			t.Fail()
		}
		for i, cell := range row {
			if cell.Row != count {
				t.Log("One of the cells we got back when asking for row ", count, " was not in the right rows.")
				t.Fail()
			}
			if cell.Col != i {
				t.Log("One of the cells we got back from row ", count, " was not in the right column.")
				t.Fail()
			}
		}

		block := grid.Block(count)
		if len(block) != DIM {
			t.Log("We got back a block but it had the wrong number of items.")
			t.Fail()
		}

		for _, cell := range block {
			if cell.Block != count {
				t.Log("We got a cell back in a block with the wrong block number")
				t.Fail()
			}
		}

		if block[0].Row != blockUpperLeftRow[count] || block[0].Col != blockUpperLeftCol[count] {
			t.Log("We got back the wrong first cell from block ", count, ": ", block[0])
			t.Fail()
		}

		if block[DIM-1].Row != blockUpperLeftRow[count]+BLOCK_DIM-1 || block[DIM-1].Col != blockUpperLeftCol[count]+BLOCK_DIM-1 {
			t.Log("We got back the wrong last cell from block ", count, ": ", block[0])
			t.Fail()
		}

	}

	cell := grid.Cell(2, 2)

	if cell.Row != 2 || cell.Col != 2 {
		t.Log("We grabbed a cell but what we got back was the wrong row and col.")
		t.Fail()
	}

	neighbors := cell.Neighbors()

	if len(neighbors) != (DIM-1)*3-(BLOCK_DIM-1)*2 {
		t.Log("We got a different number of neighbors than what we were expecting: ", len(neighbors))
		t.Fail()
	}
	neighborsMap := make(map[*Cell]bool)
	for _, neighbor := range neighbors {
		if neighbor == nil {
			t.Log("We found a nil neighbor")
			t.Fail()
		}
		if neighbor.Row != cell.Row && neighbor.Col != cell.Col && neighbor.Block != cell.Block {
			t.Log("We found a neighbor in ourselves that doesn't appear to be related: Neighbor: ", neighbor, " Cell: ", cell)
			t.Fail()
		}
		if _, ok := neighborsMap[neighbor]; ok {
			t.Log("We found a duplicate in the neighbors list")
			t.Fail()
		} else {
			neighborsMap[cell] = true
		}
	}

}

func nCopies(in string, copies int) (result []string) {
	for i := 0; i < copies; i++ {
		result = append(result, in)
	}
	return
}
<commit_msg>Explicitly tested that possible/impossible operates like a counter.<commit_after>package dokugen

import (
	"strconv"
	"strings"
	"testing"
)

func TestCellCreation(t *testing.T) {

	number := 1
	data := strconv.Itoa(number)
	cell := NewCell(nil, 0, 0)

	for i := 1; i <= DIM; i++ {
		if !cell.Possible(i) {
			t.Log("Cell reported ", i, " was impossible even though it hadn't been touched")
			t.Fail()
		}
		cell.setImpossible(i)
		if cell.Possible(i) {
			t.Log("Cell reported ", i, " was possible even though we explicitly set it as impossible")
			t.Fail()
		}
		cell.setImpossible(i)
		cell.setPossible(i)
		if cell.Possible(i) {
			t.Log("Cell reported ", i, " was possible even though we'd only called Possible 1x and Impossible 2x")
			t.Fail()
		}
		cell.setPossible(i)
		if !cell.Possible(i) {
			t.Log("Cell reported ", i, " was impossible even after matched calls to setPossible/setImpossible")
			t.Fail()
		}
	}

	cell.Load(data)
	if cell.Number() != number {
		t.Log("Number came back wrong")
		t.Fail()
	}
	if cell.Row != 0 {
		t.Log("Row came back wrong")
		t.Fail()
	}
	if cell.Col != 0 {
		t.Log("Cell came back wrong")
		t.Fail()
	}
	if cell.DataString() != data {
		t.Log("Cell round-tripped out with different string than data in")
		t.Fail()
	}
	//TODO: test failing for values that are too high.
	for i := 1; i <= DIM; i++ {
		if i == number {
			if !cell.Possible(i) {
				t.Log("We reported back a number we were explicitly set to was impossible")
				t.Fail()
			}
		} else if cell.Possible(i) {
			t.Log("We reported back that a number was possible when another number had been explicitly set")
			t.Fail()
		}
	}

	number = 2

	cell.SetNumber(number)

	if cell.Number() != number {
		t.Log("Number came back wrong after being set with SetNumber")
		t.Fail()
	}

	for i := 1; i <= DIM; i++ {
		if i == number {
			if !cell.Possible(i) {
				t.Log("We reported back a number we were explicitly set to was impossible (2nd set)")
				t.Fail()
			}
		} else if cell.Possible(i) {
			t.Log("We reported back that a number was possible when another number had been explicitly set (2nd set)")
			t.Fail()
		}
	}

	number = 0

	cell.SetNumber(number)

	if cell.Number() != number {
		t.Log("Number came back wrong after being set with SetNumber to 0")
		t.Fail()
	}

	for i := 1; i <= DIM; i++ {
		if !cell.Possible(i) {
			t.Log("We reported back a number was not possible when 0 had been explicitly set.")
			t.Fail()
		}
	}
}

func TestGridCreation(t *testing.T) {

	blockUpperLeftRow := make([]int, DIM)
	blockUpperLeftCol := make([]int, DIM)

	row := 0
	col := -1 * BLOCK_DIM

	for i := 0; i < DIM; i++ {
		col += BLOCK_DIM
		if col >= DIM {
			col = 0
			row += BLOCK_DIM
		}
		blockUpperLeftRow[i] = row
		blockUpperLeftCol[i] = col
	}

	cellData := "1"
	rowData := strings.Join(nCopies(cellData, DIM), COL_SEP)
	data := strings.Join(nCopies(rowData, DIM), ROW_SEP)
	grid := LoadGrid(data)
	if len(grid.cells) != DIM*DIM {
		t.Log("Didn't generate enough cells")
		t.Fail()
	}
	if grid.DataString() != data {
		t.Log("The grid round-tripped with different result than data in")
		t.Fail()
	}
	if grid.cells[10].Number() != 1 {
		t.Log("A random spot check of a cell had the wrong number: %s", grid.cells[10])
		t.Fail()
	}

	for count := 0; count < DIM; count++ {
		col := grid.Col(count)
		if num := len(col); num != DIM {
			t.Log("We got back a column but it had the wrong amount of items: ", num, "\n")
			t.Fail()
		}
		for i, cell := range col {
			if cell.Col != count {
				t.Log("One of the cells we got back when asking for column ", count, " was not in the right column.")
				t.Fail()
			}
			if cell.Row != i {
				t.Log("One of the cells we got back when asking for column ", count, " was not in the right row.")
				t.Fail()
			}
		}

		row := grid.Row(count)
		if len(row) != DIM {
			t.Log("We got back a row but it had the wrong number of items.")
			t.Fail()
		}
		for i, cell := range row {
			if cell.Row != count {
				t.Log("One of the cells we got back when asking for row ", count, " was not in the right rows.")
				t.Fail()
			}
			if cell.Col != i {
				t.Log("One of the cells we got back from row ", count, " was not in the right column.")
				t.Fail()
			}
		}

		block := grid.Block(count)
		if len(block) != DIM {
			t.Log("We got back a block but it had the wrong number of items.")
			t.Fail()
		}

		for _, cell := range block {
			if cell.Block != count {
				t.Log("We got a cell back in a block with the wrong block number")
				t.Fail()
			}
		}

		if block[0].Row != blockUpperLeftRow[count] || block[0].Col != blockUpperLeftCol[count] {
			t.Log("We got back the wrong first cell from block ", count, ": ", block[0])
			t.Fail()
		}

		if block[DIM-1].Row != blockUpperLeftRow[count]+BLOCK_DIM-1 || block[DIM-1].Col != blockUpperLeftCol[count]+BLOCK_DIM-1 {
			t.Log("We got back the wrong last cell from block ", count, ": ", block[0])
			t.Fail()
		}

	}

	cell := grid.Cell(2, 2)

	if cell.Row != 2 || cell.Col != 2 {
		t.Log("We grabbed a cell but what we got back was the wrong row and col.")
		t.Fail()
	}

	neighbors := cell.Neighbors()

	if len(neighbors) != (DIM-1)*3-(BLOCK_DIM-1)*2 {
		t.Log("We got a different number of neighbors than what we were expecting: ", len(neighbors))
		t.Fail()
	}
	neighborsMap := make(map[*Cell]bool)
	for _, neighbor := range neighbors {
		if neighbor == nil {
			t.Log("We found a nil neighbor")
			t.Fail()
		}
		if neighbor.Row != cell.Row && neighbor.Col != cell.Col && neighbor.Block != cell.Block {
			t.Log("We found a neighbor in ourselves that doesn't appear to be related: Neighbor: ", neighbor, " Cell: ", cell)
			t.Fail()
		}
		if _, ok := neighborsMap[neighbor]; ok {
			t.Log("We found a duplicate in the neighbors list")
			t.Fail()
		} else {
			neighborsMap[cell] = true
		}
	}

}

func nCopies(in string, copies int) (result []string) {
	for i := 0; i < copies; i++ {
		result = append(result, in)
	}
	return
}
<|endoftext|>


<commit_before>package sshkey

import (
	"crypto/dsa"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	cryptorand "crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/asn1"
	"encoding/pem"
	"fmt"
	"io"
	"math/big"

	"golang.org/x/crypto/ssh"
)

type Algorithm int

//go:generate enumer -type Algorithm -transform snake
const (
	RSA Algorithm = iota
	DSA
	ECDSA
	ED25519
)

var (
	ErrUnknownAlgorithm    = fmt.Errorf("sshkey: unknown private key algorithm")
	ErrInvalidRSAKeySize   = fmt.Errorf("sshkey: invalid private key rsa size: must be more than 1024")
	ErrInvalidECDSAKeySize = fmt.Errorf("sshkey: invalid private key ecdsa size, must be one of 256, 384 or 521")
	ErrInvalidDSAKeySize   = fmt.Errorf("sshkey: invalid private key dsa size, must be one of 1024, 2048 or 3072")
)

// Pair represents an ssh key pair, as in
type Pair struct {
	Private []byte
	Public  []byte
}

func NewPair(public, private interface{}) (*Pair, error) {
	kb, err := x509.MarshalPKCS8PrivateKey(private)
	if err != nil {
		return nil, err
	}

	privBlk := &pem.Block{
		Type:    "PRIVATE KEY",
		Headers: nil,
		Bytes:   kb,
	}

	publicKey, err := ssh.NewPublicKey(public)
	if err != nil {
		return nil, err
	}
	return &Pair{
		Private: pem.EncodeToMemory(privBlk),
		Public:  ssh.MarshalAuthorizedKey(publicKey),
	}, nil
}

func PairFromED25519(public ed25519.PublicKey, private ed25519.PrivateKey) (*Pair, error) {
	// see https://github.com/golang/crypto/blob/7f63de1d35b0f77fa2b9faea3e7deb402a2383c8/ssh/keys.go#L1273-L1443
	key := struct {
		Pub     []byte
		Priv    []byte
		Comment string
		Pad     []byte `ssh:"rest"`
	}{
		Pub:  public,
		Priv: private,
	}
	keyBytes := ssh.Marshal(key)

	pk1 := struct {
		Check1  uint32
		Check2  uint32
		Keytype string
		Rest    []byte `ssh:"rest"`
	}{
		Keytype: ssh.KeyAlgoED25519,
		Rest:    keyBytes,
	}
	pk1Bytes := ssh.Marshal(pk1)

	k := struct {
		CipherName   string
		KdfName      string
		KdfOpts      string
		NumKeys      uint32
		PubKey       []byte
		PrivKeyBlock []byte
	}{
		CipherName:   "none",
		KdfName:      "none",
		KdfOpts:      "",
		NumKeys:      1,
		PrivKeyBlock: pk1Bytes,
	}

	const opensshV1Magic = "openssh-key-v1\x00"

	privBlk := &pem.Block{
		Type:    "OPENSSH PRIVATE KEY",
		Headers: nil,
		Bytes:   append([]byte(opensshV1Magic), ssh.Marshal(k)...),
	}
	publicKey, err := ssh.NewPublicKey(public)
	if err != nil {
		return nil, err
	}
	return &Pair{
		Private: pem.EncodeToMemory(privBlk),
		Public:  ssh.MarshalAuthorizedKey(publicKey),
	}, nil
}

func PairFromDSA(key *dsa.PrivateKey) (*Pair, error) {
	// see https://github.com/golang/crypto/blob/7f63de1d35b0f77fa2b9faea3e7deb402a2383c8/ssh/keys.go#L1186-L1195
	// and https://linux.die.net/man/1/dsa
	k := struct {
		Version int
		P       *big.Int
		Q       *big.Int
		G       *big.Int
		Pub     *big.Int
		Priv    *big.Int
	}{
		Version: 0,
		P:       key.P,
		Q:       key.Q,
		G:       key.G,
		Pub:     key.Y,
		Priv:    key.X,
	}
	kb, err := asn1.Marshal(k)
	privBlk := &pem.Block{
		Type:    "DSA PRIVATE KEY",
		Headers: nil,
		Bytes:   kb,
	}
	publicKey, err := ssh.NewPublicKey(&key.PublicKey)
	if err != nil {
		return nil, err
	}
	return &Pair{
		Private: pem.EncodeToMemory(privBlk),
		Public:  ssh.MarshalAuthorizedKey(publicKey),
	}, nil
}

// GeneratePair generates a Private/Public key pair using algorithm t.
//
// When rand is nil "crypto/rand".Reader will be used.
//
// bits specifies the number of bits in the key to create. For RSA keys, the
// minimum size is 1024 bits and the default is 3072 bits. Generally, 3072 bits
// is considered sufficient. DSA keys must be exactly 1024 bits - or 2 or 3
// times that - as specified by FIPS 186-2. For ECDSA keys, bits determines the
// key length by selecting from one of three elliptic curve sizes: 256, 384 or
// 521 bits. Attempting to use bit lengths other than these three values for
// ECDSA keys will fail. Ed25519 keys have a fixed length and the bits will
// be ignored.
func GeneratePair(t Algorithm, rand io.Reader, bits int) (*Pair, error) {
	if rand == nil {
		rand = cryptorand.Reader
	}
	switch t {
	case DSA:
		if bits == 0 {
			bits = 1024
		}
		var sizes dsa.ParameterSizes
		switch bits {
		case 1024:
			sizes = dsa.L1024N160
		case 2048:
			sizes = dsa.L2048N256
		case 3072:
			sizes = dsa.L3072N256
		default:
			return nil, ErrInvalidDSAKeySize
		}

		params := dsa.Parameters{}
		if err := dsa.GenerateParameters(&params, rand, sizes); err != nil {
			return nil, err
		}

		dsakey := &dsa.PrivateKey{
			PublicKey: dsa.PublicKey{
				Parameters: params,
			},
		}
		if err := dsa.GenerateKey(dsakey, rand); err != nil {
			return nil, err
		}
		return PairFromDSA(dsakey)
	case ECDSA:
		if bits == 0 {
			bits = 521
		}
		var ecdsakey *ecdsa.PrivateKey
		var err error
		switch bits {
		case 256:
			ecdsakey, err = ecdsa.GenerateKey(elliptic.P256(), rand)
		case 384:
			ecdsakey, err = ecdsa.GenerateKey(elliptic.P384(), rand)
		case 521:
			ecdsakey, err = ecdsa.GenerateKey(elliptic.P521(), rand)
		default:
			ecdsakey, err = nil, ErrInvalidECDSAKeySize
		}
		if err != nil {
			return nil, err
		}
		return NewPair(&ecdsakey.PublicKey, ecdsakey)
	case ED25519:
		publicKey, privateKey, err := ed25519.GenerateKey(rand)
		if err != nil {
			return nil, err
		}
		return PairFromED25519(publicKey, privateKey)
	case RSA:
		if bits == 0 {
			bits = 4096
		}
		if bits < 1024 {
			return nil, ErrInvalidRSAKeySize
		}
		rsakey, err := rsa.GenerateKey(rand, bits)
		if err != nil {
			return nil, err
		}
		return NewPair(&rsakey.PublicKey, rsakey)
	default:
		return nil, ErrUnknownAlgorithm
	}
}
<commit_msg>add a comment for DSA default bits setting<commit_after>package sshkey

import (
	"crypto/dsa"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	cryptorand "crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/asn1"
	"encoding/pem"
	"fmt"
	"io"
	"math/big"

	"golang.org/x/crypto/ssh"
)

type Algorithm int

//go:generate enumer -type Algorithm -transform snake
const (
	RSA Algorithm = iota
	DSA
	ECDSA
	ED25519
)

var (
	ErrUnknownAlgorithm    = fmt.Errorf("sshkey: unknown private key algorithm")
	ErrInvalidRSAKeySize   = fmt.Errorf("sshkey: invalid private key rsa size: must be more than 1024")
	ErrInvalidECDSAKeySize = fmt.Errorf("sshkey: invalid private key ecdsa size, must be one of 256, 384 or 521")
	ErrInvalidDSAKeySize   = fmt.Errorf("sshkey: invalid private key dsa size, must be one of 1024, 2048 or 3072")
)

// Pair represents an ssh key pair, as in
type Pair struct {
	Private []byte
	Public  []byte
}

func NewPair(public, private interface{}) (*Pair, error) {
	kb, err := x509.MarshalPKCS8PrivateKey(private)
	if err != nil {
		return nil, err
	}

	privBlk := &pem.Block{
		Type:    "PRIVATE KEY",
		Headers: nil,
		Bytes:   kb,
	}

	publicKey, err := ssh.NewPublicKey(public)
	if err != nil {
		return nil, err
	}
	return &Pair{
		Private: pem.EncodeToMemory(privBlk),
		Public:  ssh.MarshalAuthorizedKey(publicKey),
	}, nil
}

func PairFromED25519(public ed25519.PublicKey, private ed25519.PrivateKey) (*Pair, error) {
	// see https://github.com/golang/crypto/blob/7f63de1d35b0f77fa2b9faea3e7deb402a2383c8/ssh/keys.go#L1273-L1443
	key := struct {
		Pub     []byte
		Priv    []byte
		Comment string
		Pad     []byte `ssh:"rest"`
	}{
		Pub:  public,
		Priv: private,
	}
	keyBytes := ssh.Marshal(key)

	pk1 := struct {
		Check1  uint32
		Check2  uint32
		Keytype string
		Rest    []byte `ssh:"rest"`
	}{
		Keytype: ssh.KeyAlgoED25519,
		Rest:    keyBytes,
	}
	pk1Bytes := ssh.Marshal(pk1)

	k := struct {
		CipherName   string
		KdfName      string
		KdfOpts      string
		NumKeys      uint32
		PubKey       []byte
		PrivKeyBlock []byte
	}{
		CipherName:   "none",
		KdfName:      "none",
		KdfOpts:      "",
		NumKeys:      1,
		PrivKeyBlock: pk1Bytes,
	}

	const opensshV1Magic = "openssh-key-v1\x00"

	privBlk := &pem.Block{
		Type:    "OPENSSH PRIVATE KEY",
		Headers: nil,
		Bytes:   append([]byte(opensshV1Magic), ssh.Marshal(k)...),
	}
	publicKey, err := ssh.NewPublicKey(public)
	if err != nil {
		return nil, err
	}
	return &Pair{
		Private: pem.EncodeToMemory(privBlk),
		Public:  ssh.MarshalAuthorizedKey(publicKey),
	}, nil
}

func PairFromDSA(key *dsa.PrivateKey) (*Pair, error) {
	// see https://github.com/golang/crypto/blob/7f63de1d35b0f77fa2b9faea3e7deb402a2383c8/ssh/keys.go#L1186-L1195
	// and https://linux.die.net/man/1/dsa
	k := struct {
		Version int
		P       *big.Int
		Q       *big.Int
		G       *big.Int
		Pub     *big.Int
		Priv    *big.Int
	}{
		Version: 0,
		P:       key.P,
		Q:       key.Q,
		G:       key.G,
		Pub:     key.Y,
		Priv:    key.X,
	}
	kb, err := asn1.Marshal(k)
	privBlk := &pem.Block{
		Type:    "DSA PRIVATE KEY",
		Headers: nil,
		Bytes:   kb,
	}
	publicKey, err := ssh.NewPublicKey(&key.PublicKey)
	if err != nil {
		return nil, err
	}
	return &Pair{
		Private: pem.EncodeToMemory(privBlk),
		Public:  ssh.MarshalAuthorizedKey(publicKey),
	}, nil
}

// GeneratePair generates a Private/Public key pair using algorithm t.
//
// When rand is nil "crypto/rand".Reader will be used.
//
// bits specifies the number of bits in the key to create. For RSA keys, the
// minimum size is 1024 bits and the default is 3072 bits. Generally, 3072 bits
// is considered sufficient. DSA keys must be exactly 1024 bits - or 2 or 3
// times that - as specified by FIPS 186-2. For ECDSA keys, bits determines the
// key length by selecting from one of three elliptic curve sizes: 256, 384 or
// 521 bits. Attempting to use bit lengths other than these three values for
// ECDSA keys will fail. Ed25519 keys have a fixed length and the bits will
// be ignored.
func GeneratePair(t Algorithm, rand io.Reader, bits int) (*Pair, error) {
	if rand == nil {
		rand = cryptorand.Reader
	}
	switch t {
	case DSA:
		if bits == 0 {
			// currently the ssh package can only decode 1024 bits dsa keys, so
			// that's going be the default for now see
			// https://github.com/golang/crypto/blob/7f63de1d35b0f77fa2b9faea3e7deb402a2383c8/ssh/keys.go#L411-L420
			bits = 1024
		}
		var sizes dsa.ParameterSizes
		switch bits {
		case 1024:
			sizes = dsa.L1024N160
		case 2048:
			sizes = dsa.L2048N256
		case 3072:
			sizes = dsa.L3072N256
		default:
			return nil, ErrInvalidDSAKeySize
		}

		params := dsa.Parameters{}
		if err := dsa.GenerateParameters(&params, rand, sizes); err != nil {
			return nil, err
		}

		dsakey := &dsa.PrivateKey{
			PublicKey: dsa.PublicKey{
				Parameters: params,
			},
		}
		if err := dsa.GenerateKey(dsakey, rand); err != nil {
			return nil, err
		}
		return PairFromDSA(dsakey)
	case ECDSA:
		if bits == 0 {
			bits = 521
		}
		var ecdsakey *ecdsa.PrivateKey
		var err error
		switch bits {
		case 256:
			ecdsakey, err = ecdsa.GenerateKey(elliptic.P256(), rand)
		case 384:
			ecdsakey, err = ecdsa.GenerateKey(elliptic.P384(), rand)
		case 521:
			ecdsakey, err = ecdsa.GenerateKey(elliptic.P521(), rand)
		default:
			ecdsakey, err = nil, ErrInvalidECDSAKeySize
		}
		if err != nil {
			return nil, err
		}
		return NewPair(&ecdsakey.PublicKey, ecdsakey)
	case ED25519:
		publicKey, privateKey, err := ed25519.GenerateKey(rand)
		if err != nil {
			return nil, err
		}
		return PairFromED25519(publicKey, privateKey)
	case RSA:
		if bits == 0 {
			bits = 4096
		}
		if bits < 1024 {
			return nil, ErrInvalidRSAKeySize
		}
		rsakey, err := rsa.GenerateKey(rand, bits)
		if err != nil {
			return nil, err
		}
		return NewPair(&rsakey.PublicKey, rsakey)
	default:
		return nil, ErrUnknownAlgorithm
	}
}
<|endoftext|>


<commit_before>/*
   Copyright The containerd Authors.

   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 shim

import (
	"context"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/containerd/containerd/namespaces"
	"github.com/pkg/errors"
)

var runtimePaths sync.Map

// Command returns the shim command with the provided args and configuration
func Command(ctx context.Context, runtime, containerdAddress, path string, cmdArgs ...string) (*exec.Cmd, error) {
	ns, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return nil, err
	}
	self, err := os.Executable()
	if err != nil {
		return nil, err
	}
	args := []string{
		"-namespace", ns,
		"-address", containerdAddress,
		"-publish-binary", self,
	}
	args = append(args, cmdArgs...)
	name := BinaryName(runtime)
	if name == "" {
		return nil, fmt.Errorf("invalid runtime name %s, correct runtime name should format like io.containerd.runc.v1", runtime)
	}

	var cmdPath string
	cmdPathI, cmdPathFound := runtimePaths.Load(name)
	if cmdPathFound {
		cmdPath = cmdPathI.(string)
	} else {
		var lerr error
		if cmdPath, lerr = exec.LookPath(name); lerr != nil {
			if eerr, ok := lerr.(*exec.Error); ok {
				if eerr.Err == exec.ErrNotFound {
					// LookPath only finds current directory matches based on
					// the callers current directory but the caller is not
					// likely in the same directory as the containerd
					// executables. Instead match the calling binaries path
					// (containerd) and see if they are side by side. If so
					// execute the shim found there.
					testPath := filepath.Join(filepath.Dir(self), name)
					if _, serr := os.Stat(testPath); serr == nil {
						path = testPath
					}
					if path == "" {
						return nil, errors.Wrapf(os.ErrNotExist, "runtime %q binary not installed %q", runtime, name)
					}
				}
			}
		}
		cmdPath, err = filepath.Abs(cmdPath)
		if err != nil {
			return nil, err
		}
		if cmdPathI, cmdPathFound = runtimePaths.LoadOrStore(name, cmdPath); cmdPathFound {
			// We didn't store cmdPath we loaded an already cached value. Use it.
			cmdPath = cmdPathI.(string)
		}
	}

	cmd := exec.Command(cmdPath, args...)
	cmd.Dir = path
	cmd.Env = append(os.Environ(), "GOMAXPROCS=2")
	cmd.SysProcAttr = getSysProcAttr()
	return cmd, nil
}

// BinaryName returns the shim binary name from the runtime name,
// empty string returns means runtime name is invalid
func BinaryName(runtime string) string {
	// runtime name should format like $prefix.name.version
	parts := strings.Split(runtime, ".")
	if len(parts) < 2 {
		return ""
	}

	return fmt.Sprintf(shimBinaryFormat, parts[len(parts)-2], parts[len(parts)-1])
}

// Connect to the provided address
func Connect(address string, d func(string, time.Duration) (net.Conn, error)) (net.Conn, error) {
	return d(address, 100*time.Second)
}

// WritePidFile writes a pid file atomically
func WritePidFile(path string, pid int) error {
	path, err := filepath.Abs(path)
	if err != nil {
		return err
	}
	tempPath := filepath.Join(filepath.Dir(path), fmt.Sprintf(".%s", filepath.Base(path)))
	f, err := os.OpenFile(tempPath, os.O_RDWR|os.O_CREATE|os.O_EXCL|os.O_SYNC, 0666)
	if err != nil {
		return err
	}
	_, err = fmt.Fprintf(f, "%d", pid)
	f.Close()
	if err != nil {
		return err
	}
	return os.Rename(tempPath, path)
}

// WriteAddress writes a address file atomically
func WriteAddress(path, address string) error {
	path, err := filepath.Abs(path)
	if err != nil {
		return err
	}
	tempPath := filepath.Join(filepath.Dir(path), fmt.Sprintf(".%s", filepath.Base(path)))
	f, err := os.OpenFile(tempPath, os.O_RDWR|os.O_CREATE|os.O_EXCL|os.O_SYNC, 0666)
	if err != nil {
		return err
	}
	_, err = f.WriteString(address)
	f.Close()
	if err != nil {
		return err
	}
	return os.Rename(tempPath, path)
}
<commit_msg>Fix bug in shim path lookup<commit_after>/*
   Copyright The containerd Authors.

   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 shim

import (
	"context"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/containerd/containerd/namespaces"
	"github.com/pkg/errors"
)

var runtimePaths sync.Map

// Command returns the shim command with the provided args and configuration
func Command(ctx context.Context, runtime, containerdAddress, path string, cmdArgs ...string) (*exec.Cmd, error) {
	ns, err := namespaces.NamespaceRequired(ctx)
	if err != nil {
		return nil, err
	}
	self, err := os.Executable()
	if err != nil {
		return nil, err
	}
	args := []string{
		"-namespace", ns,
		"-address", containerdAddress,
		"-publish-binary", self,
	}
	args = append(args, cmdArgs...)
	name := BinaryName(runtime)
	if name == "" {
		return nil, fmt.Errorf("invalid runtime name %s, correct runtime name should format like io.containerd.runc.v1", runtime)
	}

	var cmdPath string
	cmdPathI, cmdPathFound := runtimePaths.Load(name)
	if cmdPathFound {
		cmdPath = cmdPathI.(string)
	} else {
		var lerr error
		if cmdPath, lerr = exec.LookPath(name); lerr != nil {
			if eerr, ok := lerr.(*exec.Error); ok {
				if eerr.Err == exec.ErrNotFound {
					// LookPath only finds current directory matches based on
					// the callers current directory but the caller is not
					// likely in the same directory as the containerd
					// executables. Instead match the calling binaries path
					// (containerd) and see if they are side by side. If so
					// execute the shim found there.
					testPath := filepath.Join(filepath.Dir(self), name)
					if _, serr := os.Stat(testPath); serr == nil {
						cmdPath = testPath
					}
					if cmdPath == "" {
						return nil, errors.Wrapf(os.ErrNotExist, "runtime %q binary not installed %q", runtime, name)
					}
				}
			}
		}
		cmdPath, err = filepath.Abs(cmdPath)
		if err != nil {
			return nil, err
		}
		if cmdPathI, cmdPathFound = runtimePaths.LoadOrStore(name, cmdPath); cmdPathFound {
			// We didn't store cmdPath we loaded an already cached value. Use it.
			cmdPath = cmdPathI.(string)
		}
	}

	cmd := exec.Command(cmdPath, args...)
	cmd.Dir = path
	cmd.Env = append(os.Environ(), "GOMAXPROCS=2")
	cmd.SysProcAttr = getSysProcAttr()
	return cmd, nil
}

// BinaryName returns the shim binary name from the runtime name,
// empty string returns means runtime name is invalid
func BinaryName(runtime string) string {
	// runtime name should format like $prefix.name.version
	parts := strings.Split(runtime, ".")
	if len(parts) < 2 {
		return ""
	}

	return fmt.Sprintf(shimBinaryFormat, parts[len(parts)-2], parts[len(parts)-1])
}

// Connect to the provided address
func Connect(address string, d func(string, time.Duration) (net.Conn, error)) (net.Conn, error) {
	return d(address, 100*time.Second)
}

// WritePidFile writes a pid file atomically
func WritePidFile(path string, pid int) error {
	path, err := filepath.Abs(path)
	if err != nil {
		return err
	}
	tempPath := filepath.Join(filepath.Dir(path), fmt.Sprintf(".%s", filepath.Base(path)))
	f, err := os.OpenFile(tempPath, os.O_RDWR|os.O_CREATE|os.O_EXCL|os.O_SYNC, 0666)
	if err != nil {
		return err
	}
	_, err = fmt.Fprintf(f, "%d", pid)
	f.Close()
	if err != nil {
		return err
	}
	return os.Rename(tempPath, path)
}

// WriteAddress writes a address file atomically
func WriteAddress(path, address string) error {
	path, err := filepath.Abs(path)
	if err != nil {
		return err
	}
	tempPath := filepath.Join(filepath.Dir(path), fmt.Sprintf(".%s", filepath.Base(path)))
	f, err := os.OpenFile(tempPath, os.O_RDWR|os.O_CREATE|os.O_EXCL|os.O_SYNC, 0666)
	if err != nil {
		return err
	}
	_, err = f.WriteString(address)
	f.Close()
	if err != nil {
		return err
	}
	return os.Rename(tempPath, path)
}
<|endoftext|>