<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

// Dynamo DB cache layer

/* Depends on godynamo. Please refer to that package for how to configure
 * this package.
 *
 * The dynamo table we use has the following structure:
 *  Primary Hash key: hostname
 *  No range key
 *  No secondary indexes.
 *  columns:
 *      Mtime   - last access time
 *      Status  - status return value
 *      Data    - memory dump string
 *      Error   - error string value
 *
 */
package cache

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

	conf "github.com/smugmug/godynamo/conf"
	conf_file "github.com/smugmug/godynamo/conf_file"
	ep "github.com/smugmug/godynamo/endpoint"
	get "github.com/smugmug/godynamo/endpoints/get_item"
	put "github.com/smugmug/godynamo/endpoints/put_item"
	keepalive "github.com/smugmug/godynamo/keepalive"
)

var CACHE_TAB = "Heartbleed"
var expiry time.Duration

type CacheReply struct {
	Host       string
	LastUpdate int64
	Status     int64
	Data       string
	Error      string
}

/* Initialize the cache layer.
 * expiration - Go standard duration string, defaults to '10m'
 */
func Init(expiration string) {
	conf_file.Read()

	if conf.Vals.Initialized == false {
		panic("Uninitialized conf.Vals global")
	}

	if conf.Vals.Network.DynamoDB.KeepAlive {
		log.Printf("[cache] INFO: Launching background DynamoDB keepalive")
		go keepalive.KeepAlive([]string{conf.Vals.Network.DynamoDB.URL})
	}

	exp, err := time.ParseDuration(expiration)
	if err != nil {
		log.Printf("[cache] WARN: Could not parse expiry string [%s]. Expiring after 10m.", err.Error())
		expiry = time.Minute * 10
	} else {
		expiry = exp
	}
}

/* Fetch the record from the Cache.
 *
 * returns the record and if the record should be considered "OK" to use.
 * An OK record is a valid, non-expired cache stored entry.
 */
func Check(host string) (CacheReply, bool) {
	getr := get.Request{
		TableName: CACHE_TAB,
		Key:       make(ep.Item),
	}
	getr.Key["hostname"] = ep.AttributeValue{S: host}
	body, code, err := getr.EndpointReq()
	if err != nil {
		log.Printf("[cache] ERROR: %s", err.Error())
		return CacheReply{}, false
	}
	if code != http.StatusOK {
		return CacheReply{}, false
	}

	// get the time from the body
	// log.Printf("####CACHE_GET: %s %d\n", string(body), len(body))
	if len(body) < 3 {
		// record not found
		return CacheReply{}, false
	}

	var gr get.Response
	var reply CacheReply
	if err = json.Unmarshal([]byte(body), &gr); err == nil {
		reply.Status, err = strconv.ParseInt(gr.Item["Status"].N, 10, 64)
		if err != nil {
			// unparsable status
			log.Printf("[cache] ERROR: Bad Record %s, %s", host, err.Error())
			return CacheReply{}, false
		}

		reply.Host = gr.Item["hostname"].S
		reply.Error = gr.Item["Error"].S
		if reply.Error == "---" {
			reply.Error = ""
		}
		reply.Data = gr.Item["Data"].S
		if reply.Data == "---" {
			reply.Data = ""
		}

		reply.LastUpdate, err = strconv.ParseInt(gr.Item["Mtime"].N, 10, 64)
		if err != nil {
			// unparsable time
			log.Printf("[cache] ERROR: Bad Record %s, %s", host, err.Error())
			return CacheReply{}, false
		}
		if reply.LastUpdate < time.Now().UTC().Truncate(expiry).Unix() {
			// record has expired.
			return reply, false
		}
		// log.Printf("gr %+v", reply)
	}

	return reply, true
}

/* Set the state of the host in the cache
 *
 */
func Set(host string, state int, data, errS string) error {
	if data == "" {
		data = "---"
	}
	if errS == "" {
		errS = "---"
	}

	putr := put.Request{
		TableName: CACHE_TAB,
		Item:      make(ep.Item),
	}

	putr.Item["hostname"] = ep.AttributeValue{S: host}
	putr.Item["Mtime"] = ep.AttributeValue{N: strconv.FormatInt(time.Now().UTC().Unix(), 10)}
	putr.Item["Status"] = ep.AttributeValue{N: strconv.FormatInt(int64(state), 10)}
	putr.Item["Data"] = ep.AttributeValue{S: data}
	putr.Item["Error"] = ep.AttributeValue{S: errS}

	body, code, err := putr.EndpointReq()
	if err != nil {
		log.Printf("[cache] ERROR: put failed %v, %s", err, body)
		return err
	}
	if code != http.StatusOK {
		log.Printf("[cache] ERROR: put failed %v, %s", code, body)
		return fmt.Errorf("put failed: code %d", code)
	}

	// log.Printf("####CACHE_SET: %s\n", string(body))
	return nil
}
<commit_msg>Invalidate cache after "expiry", not every "expiry"<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

// Dynamo DB cache layer

/* Depends on godynamo. Please refer to that package for how to configure
 * this package.
 *
 * The dynamo table we use has the following structure:
 *  Primary Hash key: hostname
 *  No range key
 *  No secondary indexes.
 *  columns:
 *      Mtime   - last access time
 *      Status  - status return value
 *      Data    - memory dump string
 *      Error   - error string value
 *
 */
package cache

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

	conf "github.com/smugmug/godynamo/conf"
	conf_file "github.com/smugmug/godynamo/conf_file"
	ep "github.com/smugmug/godynamo/endpoint"
	get "github.com/smugmug/godynamo/endpoints/get_item"
	put "github.com/smugmug/godynamo/endpoints/put_item"
	keepalive "github.com/smugmug/godynamo/keepalive"
)

var CACHE_TAB = "Heartbleed"
var expiry time.Duration

type CacheReply struct {
	Host       string
	LastUpdate int64
	Status     int64
	Data       string
	Error      string
}

/* Initialize the cache layer.
 * expiration - Go standard duration string, defaults to '10m'
 */
func Init(expiration string) {
	conf_file.Read()

	if conf.Vals.Initialized == false {
		panic("Uninitialized conf.Vals global")
	}

	if conf.Vals.Network.DynamoDB.KeepAlive {
		log.Printf("[cache] INFO: Launching background DynamoDB keepalive")
		go keepalive.KeepAlive([]string{conf.Vals.Network.DynamoDB.URL})
	}

	exp, err := time.ParseDuration(expiration)
	if err != nil {
		log.Printf("[cache] WARN: Could not parse expiry string [%s]. Expiring after 10m.", err.Error())
		expiry = time.Minute * 10
	} else {
		expiry = exp
	}
}

/* Fetch the record from the Cache.
 *
 * returns the record and if the record should be considered "OK" to use.
 * An OK record is a valid, non-expired cache stored entry.
 */
func Check(host string) (CacheReply, bool) {
	getr := get.Request{
		TableName: CACHE_TAB,
		Key:       make(ep.Item),
	}
	getr.Key["hostname"] = ep.AttributeValue{S: host}
	body, code, err := getr.EndpointReq()
	if err != nil {
		log.Printf("[cache] ERROR: %s", err.Error())
		return CacheReply{}, false
	}
	if code != http.StatusOK {
		return CacheReply{}, false
	}

	// get the time from the body
	// log.Printf("####CACHE_GET: %s %d\n", string(body), len(body))
	if len(body) < 3 {
		// record not found
		return CacheReply{}, false
	}

	var gr get.Response
	var reply CacheReply
	if err = json.Unmarshal([]byte(body), &gr); err == nil {
		reply.Status, err = strconv.ParseInt(gr.Item["Status"].N, 10, 64)
		if err != nil {
			// unparsable status
			log.Printf("[cache] ERROR: Bad Record %s, %s", host, err.Error())
			return CacheReply{}, false
		}

		reply.Host = gr.Item["hostname"].S
		reply.Error = gr.Item["Error"].S
		if reply.Error == "---" {
			reply.Error = ""
		}
		reply.Data = gr.Item["Data"].S
		if reply.Data == "---" {
			reply.Data = ""
		}

		reply.LastUpdate, err = strconv.ParseInt(gr.Item["Mtime"].N, 10, 64)
		if err != nil {
			// unparsable time
			log.Printf("[cache] ERROR: Bad Record %s, %s", host, err.Error())
			return CacheReply{}, false
		}
		if reply.LastUpdate < time.Now().UTC().Add(-expiry).Unix() {
			// record has expired.
			return reply, false
		}
		// log.Printf("gr %+v", reply)
	}

	return reply, true
}

/* Set the state of the host in the cache
 *
 */
func Set(host string, state int, data, errS string) error {
	if data == "" {
		data = "---"
	}
	if errS == "" {
		errS = "---"
	}

	putr := put.Request{
		TableName: CACHE_TAB,
		Item:      make(ep.Item),
	}

	putr.Item["hostname"] = ep.AttributeValue{S: host}
	putr.Item["Mtime"] = ep.AttributeValue{N: strconv.FormatInt(time.Now().UTC().Unix(), 10)}
	putr.Item["Status"] = ep.AttributeValue{N: strconv.FormatInt(int64(state), 10)}
	putr.Item["Data"] = ep.AttributeValue{S: data}
	putr.Item["Error"] = ep.AttributeValue{S: errS}

	body, code, err := putr.EndpointReq()
	if err != nil {
		log.Printf("[cache] ERROR: put failed %v, %s", err, body)
		return err
	}
	if code != http.StatusOK {
		log.Printf("[cache] ERROR: put failed %v, %s", code, body)
		return fmt.Errorf("put failed: code %d", code)
	}

	// log.Printf("####CACHE_SET: %s\n", string(body))
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"reflect"
	"testing"
)

func Test_NewService(t *testing.T) {
	var cases = []struct {
		input Environment
		err   bool
	}{
		{
			input: Environment{
				Driver:    "mysql",
				Directory: "unkown",
			},
			err: true,
		},
		{
			input: Environment{
				Driver:    "unkown",
				Directory: "test",
			},
			err: true,
		},
		{
			input: Environment{
				Driver:    "mysql",
				Directory: "test",
			},
			err: false,
		},
	}

	for n, c := range cases {
		_, err := NewService(c.input)
		if (err != nil) != c.err {
			t.Error("case", n, "got unexpected error:", err)
		}
	}
}

func Test_Service_ListAvailableMigrations_ParseFilenames(t *testing.T) {
	var cases = []struct {
		directory string
		output    []uint64
	}{
		{
			directory: "test/",
			output:    []uint64{1, 2, 3},
		},
		{
			directory: "test/empty/",
			output:    nil,
		},
		{
			directory: "test/not_a_directory",
			output:    nil,
		},
	}

	for n, c := range cases {
		s := &Service{
			env: Environment{
				Directory: c.directory,
			},
		}
		versions := s.ListAvailableMigrations()

		if !reflect.DeepEqual(versions, c.output) {
			t.Error("case", n, "got unexpected outout:", versions)
		}
	}
}
<commit_msg>Fix Service unit tests name<commit_after>package main

import (
	"reflect"
	"testing"
)

func Test_NewService(t *testing.T) {
	var cases = []struct {
		input Environment
		err   bool
	}{
		{
			input: Environment{
				Driver:    "mysql",
				Directory: "unkown",
			},
			err: true,
		},
		{
			input: Environment{
				Driver:    "unkown",
				Directory: "test",
			},
			err: true,
		},
		{
			input: Environment{
				Driver:    "mysql",
				Directory: "test",
			},
			err: false,
		},
	}

	for n, c := range cases {
		_, err := NewService(c.input)
		if (err != nil) != c.err {
			t.Error("case", n, "got unexpected error:", err)
		}
	}
}

func Test_Service_ListAvailableMigrations(t *testing.T) {
	var cases = []struct {
		directory string
		output    []uint64
	}{
		{
			directory: "test/",
			output:    []uint64{1, 2, 3},
		},
		{
			directory: "test/empty/",
			output:    nil,
		},
		{
			directory: "test/not_a_directory",
			output:    nil,
		},
	}

	for n, c := range cases {
		s := &Service{
			env: Environment{
				Directory: c.directory,
			},
		}
		versions := s.ListAvailableMigrations()

		if !reflect.DeepEqual(versions, c.output) {
			t.Error("case", n, "got unexpected outout:", versions)
		}
	}
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"net"
	"reflect"

	"github.com/micro/go-micro/registry"
)

var (
	privateBlocks []*net.IPNet
)

func init() {
	for _, b := range []string{"10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16"} {
		if _, block, err := net.ParseCIDR(b); err == nil {
			privateBlocks = append(privateBlocks, block)
		}
	}
}

func extractValue(v reflect.Type) *registry.Value {
	if v == nil {
		return nil
	}

	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}

	arg := &registry.Value{
		Name: v.Name(),
		Type: v.Name(),
	}

	if v.Kind() == reflect.Struct {
		for i := 0; i < v.NumField(); i++ {
			val := extractValue(v.Field(i).Type)
			val.Name = v.Field(i).Name
			arg.Values = append(arg.Values, val)
		}
	}

	return arg
}

func extractEndpoint(method reflect.Method) *registry.Endpoint {
	if method.PkgPath != "" {
		return nil
	}

	var rspType, reqType reflect.Type
	var stream bool
	mt := method.Type

	switch mt.NumIn() {
	case 3:
		reqType = mt.In(1)
		rspType = mt.In(2)
	case 4:
		reqType = mt.In(2)
		rspType = mt.In(3)
	default:
		return nil
	}

	if rspType.Kind() == reflect.Func {
		stream = true
	}

	request := extractValue(reqType)
	response := extractValue(rspType)

	return &registry.Endpoint{
		Name:     method.Name,
		Request:  request,
		Response: response,
		Metadata: map[string]string{
			"stream": fmt.Sprintf("%v", stream),
		},
	}
}

func extractSubValue(typ reflect.Type) *registry.Value {
	var reqType reflect.Type
	switch typ.NumIn() {
	case 1:
		reqType = typ.In(0)
	case 2:
		reqType = typ.In(1)
	default:
		return nil
	}
	return extractValue(reqType)
}

func extractAddress(addr string) (string, error) {
	if len(addr) > 0 && (addr != "0.0.0.0" && addr != "[::]") {
		return addr, nil
	}

	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return "", fmt.Errorf("Failed to get interface addresses! Err: %v", err)
	}

	var ipAddr []byte

	for _, rawAddr := range addrs {
		var ip net.IP
		switch addr := rawAddr.(type) {
		case *net.IPAddr:
			ip = addr.IP
		case *net.IPNet:
			ip = addr.IP
		default:
			continue
		}

		if ip.To4() == nil {
			continue
		}

		if !isPrivateIP(ip.String()) {
			continue
		}

		ipAddr = ip
		break
	}

	if ipAddr == nil {
		return "", fmt.Errorf("No private IP address found, and explicit IP not provided")
	}

	return net.IP(ipAddr).String(), nil
}

func isPrivateIP(ipAddr string) bool {
	ip := net.ParseIP(ipAddr)
	for _, priv := range privateBlocks {
		if priv.Contains(ip) {
			return true
		}
	}
	return false
}
<commit_msg>Fix the sub extractor<commit_after>package server

import (
	"fmt"
	"net"
	"reflect"

	"github.com/micro/go-micro/registry"
)

var (
	privateBlocks []*net.IPNet
)

func init() {
	for _, b := range []string{"10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16"} {
		if _, block, err := net.ParseCIDR(b); err == nil {
			privateBlocks = append(privateBlocks, block)
		}
	}
}

func extractValue(v reflect.Type) *registry.Value {
	if v == nil {
		return nil
	}

	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}

	arg := &registry.Value{
		Name: v.Name(),
		Type: v.Name(),
	}

	if v.Kind() == reflect.Struct {
		for i := 0; i < v.NumField(); i++ {
			val := extractValue(v.Field(i).Type)
			val.Name = v.Field(i).Name
			arg.Values = append(arg.Values, val)
		}
	}

	return arg
}

func extractEndpoint(method reflect.Method) *registry.Endpoint {
	if method.PkgPath != "" {
		return nil
	}

	var rspType, reqType reflect.Type
	var stream bool
	mt := method.Type

	switch mt.NumIn() {
	case 3:
		reqType = mt.In(1)
		rspType = mt.In(2)
	case 4:
		reqType = mt.In(2)
		rspType = mt.In(3)
	default:
		return nil
	}

	if rspType.Kind() == reflect.Func {
		stream = true
	}

	request := extractValue(reqType)
	response := extractValue(rspType)

	return &registry.Endpoint{
		Name:     method.Name,
		Request:  request,
		Response: response,
		Metadata: map[string]string{
			"stream": fmt.Sprintf("%v", stream),
		},
	}
}

func extractSubValue(typ reflect.Type) *registry.Value {
	var reqType reflect.Type
	switch typ.NumIn() {
	case 1:
		reqType = typ.In(0)
	case 2:
		reqType = typ.In(1)
	case 3:
		reqType = typ.In(2)
	default:
		return nil
	}
	return extractValue(reqType)
}

func extractAddress(addr string) (string, error) {
	if len(addr) > 0 && (addr != "0.0.0.0" && addr != "[::]") {
		return addr, nil
	}

	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return "", fmt.Errorf("Failed to get interface addresses! Err: %v", err)
	}

	var ipAddr []byte

	for _, rawAddr := range addrs {
		var ip net.IP
		switch addr := rawAddr.(type) {
		case *net.IPAddr:
			ip = addr.IP
		case *net.IPNet:
			ip = addr.IP
		default:
			continue
		}

		if ip.To4() == nil {
			continue
		}

		if !isPrivateIP(ip.String()) {
			continue
		}

		ipAddr = ip
		break
	}

	if ipAddr == nil {
		return "", fmt.Errorf("No private IP address found, and explicit IP not provided")
	}

	return net.IP(ipAddr).String(), nil
}

func isPrivateIP(ipAddr string) bool {
	ip := net.ParseIP(ipAddr)
	for _, priv := range privateBlocks {
		if priv.Contains(ip) {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package server

import (
	"errors"
	"github.com/sybrexsys/RapidKV/datamodel"
	"hash/crc32"
	"sync"
	"sync/atomic"
	"time"
)

const shardCount = 256

const ttlCheckPeriod = time.Millisecond * 500

type kvElementh struct {
	value       datamodel.CustomDataType
	ttl         time.Duration
	ttc         time.Duration
	lockSession time.Duration
	lockttc     time.Duration
	lock        int32
}

var startTime = time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)

const maxTime = time.Duration(0x7FFFFFFFFFFFFFFF)

type shardElem struct {
	sync.RWMutex
	mapkv             map[string]*kvElementh
	ttlMutex          sync.Mutex
	nextTTLProcessing int64
	quit              chan struct{}
}

type serverKV struct {
	baseList [shardCount]*shardElem
	group    *sync.WaitGroup
}

func crc(Key string) int {
	return int(crc32.ChecksumIEEE([]byte(Key)))
}

func (shard *shardElem) ttlChecker() {
	shard.Lock()
	defer shard.Unlock()
	ct := time.Since(startTime)
	nextttl := maxTime
	for k, v := range shard.mapkv {
		if (v.ttl < ct) && (v.lock == 0) {
			delete(shard.mapkv, k)
		} else {
			if nextttl > v.ttl {
				nextttl = v.ttl
			}
		}
	}
	shard.nextTTLProcessing = int64(nextttl)
}

func (shard *shardElem) ttlloop(group *sync.WaitGroup) {
	group.Add(1)
	ticker := time.NewTicker(ttlCheckPeriod).C
loop:
	for {
		select {
		case _ = <-ticker:
			next := atomic.LoadInt64(&shard.nextTTLProcessing)
			if next < int64(time.Since(startTime)) {
				shard.ttlChecker()
			}
		case _ = <-shard.quit:
			break loop
		}
	}
	group.Done()
}

func createServer() *serverKV {
	server := &serverKV{group: &sync.WaitGroup{}}
	for i := 0; i < shardCount; i++ {
		shred := &shardElem{
			mapkv:             make(map[string]*kvElementh, 256),
			nextTTLProcessing: 0x7FFFFFFFFFFFFFFF,
			quit:              make(chan struct{}),
		}
		server.baseList[i] = shred
		go shred.ttlloop(server.group)
	}
	return server
}

func (server *serverKV) Close() {
	for i := 0; i < shardCount; i++ {
		server.baseList[i].quit <- struct{}{}
	}
	server.group.Wait()
}

func (server *serverKV) keyToShard(Key string) *shardElem {
	shardidx := crc(Key) % shardCount
	return server.baseList[shardidx]
}

func (server *serverKV) SetValue(Key string, Value datamodel.CustomDataType, Session time.Duration) (*kvElementh, error) {
	elem, ok := server.GetValue(Key)
	if ok {
		if !atomic.CompareAndSwapInt32(&elem.lock, 0, 1) {
			if elem.lockSession != Session {
				return nil, errors.New("record busy")
			}
		} else {
			defer atomic.StoreInt32(&elem.lock, 0)
		}
		elem.value = Value
		elem.ttc = time.Since(startTime)
		elem.ttl = maxTime
		return elem, nil
	}
	shard := server.keyToShard(Key)
	shard.Lock()
	defer shard.Unlock()
	// check on append
	elem, ok = shard.mapkv[Key]
	if ok {
		if !atomic.CompareAndSwapInt32(&elem.lock, 0, 1) {
			return nil, errors.New("record busy")
		}
		defer atomic.StoreInt32(&elem.lock, 0)
		elem.value = Value
		elem.ttc = time.Since(startTime)
		elem.ttl = maxTime
		return elem, nil
	}
	elem = &kvElementh{
		value: Value,
		ttc:   time.Since(startTime),
		ttl:   maxTime,
	}
	shard.mapkv[Key] = elem
	return elem, nil
}

func (server *serverKV) GetValue(Key string) (*kvElementh, bool) {
	shard := server.keyToShard(Key)
	shard.RLock()
	defer shard.RUnlock()
	el, ok := shard.mapkv[Key]
	return el, ok
}

func (server *serverKV) LockValue(Key string, Session time.Duration) error {
	shard := server.keyToShard(Key)
	shard.RLock()
	defer shard.RUnlock()
	elem, ok := shard.mapkv[Key]
	if !ok {
		return errors.New("record not found")
	}
	if !atomic.CompareAndSwapInt32(&elem.lock, 0, 1) {
		if elem.lockSession != Session {
			return errors.New("record busy")
		}
	}
	return nil
}

func (server *serverKV) UnlockValue(Key string, Session time.Duration) error {
	shard := server.keyToShard(Key)
	shard.RLock()
	defer shard.RUnlock()
	elem, ok := shard.mapkv[Key]
	if !ok {
		return errors.New("record not found")
	}
	if !atomic.CompareAndSwapInt32(&elem.lock, 0, 1) {
		if elem.lockSession != Session {
			return errors.New("record busy")
		}
	}
	elem.ttc = time.Since(startTime)
	elem.ttl = maxTime
	atomic.StoreInt32(&elem.lock, 0)
	return nil
}

func (server *serverKV) SetTTL(Key string, TTL time.Duration, Session time.Duration) error {
	var needRLock bool
	shard := server.keyToShard(Key)
	shard.RLock()
	defer func() {
		if needRLock {
			shard.RUnlock()
		}
	}()
	elem, ok := shard.mapkv[Key]
	if !ok {
		return errors.New("record not found")
	}
	if !atomic.CompareAndSwapInt32(&elem.lock, 0, 1) {
		if elem.lockSession != Session {
			return errors.New("record busy")
		}
	}
	elem.ttl = elem.ttc + TTL
	if elem.ttl < time.Since(startTime) {
		shard.RUnlock()
		needRLock = false
		shard.Lock()
		defer shard.Unlock()
		delete(shard.mapkv, Key)
	} else {
		next := atomic.LoadInt64(&shard.nextTTLProcessing)
		if next > int64(elem.ttl) {
			shard.ttlMutex.Lock()
			defer shard.ttlMutex.Unlock()
			if shard.nextTTLProcessing > int64(elem.ttl) {
				shard.nextTTLProcessing = int64(elem.ttl)
			}
		}
	}
	return nil
}
<commit_msg>Changed hash 32 to hash 64<commit_after>package server

import (
	"errors"
	"github.com/sybrexsys/RapidKV/datamodel"
	"hash/crc64"
	"sync"
	"sync/atomic"
	"time"
)

const shardCount = 256

const ttlCheckPeriod = time.Millisecond * 500

type kvElementh struct {
	value       datamodel.CustomDataType
	ttl         time.Duration
	ttc         time.Duration
	lockSession time.Duration
	lockttc     time.Duration
	lock        int32
}

var startTime = time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)

const maxTime = time.Duration(0x7FFFFFFFFFFFFFFF)

var crc64Table = crc64.MakeTable(crc64.ECMA)

type shardElem struct {
	sync.RWMutex
	mapkv             map[string]*kvElementh
	ttlMutex          sync.Mutex
	nextTTLProcessing int64
	quit              chan struct{}
}

type serverKV struct {
	baseList [shardCount]*shardElem
	group    *sync.WaitGroup
}

func crc(Key string) int {
	return int(crc64.Checksum([]byte(Key), crc64Table))
}

func (shard *shardElem) ttlChecker() {
	shard.Lock()
	defer shard.Unlock()
	ct := time.Since(startTime)
	nextttl := maxTime
	for k, v := range shard.mapkv {
		if (v.ttl < ct) && (v.lock == 0) {
			delete(shard.mapkv, k)
		} else {
			if nextttl > v.ttl {
				nextttl = v.ttl
			}
		}
	}
	shard.nextTTLProcessing = int64(nextttl)
}

func (shard *shardElem) ttlloop(group *sync.WaitGroup) {
	group.Add(1)
	ticker := time.NewTicker(ttlCheckPeriod).C
loop:
	for {
		select {
		case _ = <-ticker:
			next := atomic.LoadInt64(&shard.nextTTLProcessing)
			if next < int64(time.Since(startTime)) {
				shard.ttlChecker()
			}
		case _ = <-shard.quit:
			break loop
		}
	}
	group.Done()
}

func createServer() *serverKV {
	server := &serverKV{group: &sync.WaitGroup{}}
	for i := 0; i < shardCount; i++ {
		shred := &shardElem{
			mapkv:             make(map[string]*kvElementh, 256),
			nextTTLProcessing: 0x7FFFFFFFFFFFFFFF,
			quit:              make(chan struct{}),
		}
		server.baseList[i] = shred
		go shred.ttlloop(server.group)
	}
	return server
}

func (server *serverKV) Close() {
	for i := 0; i < shardCount; i++ {
		server.baseList[i].quit <- struct{}{}
	}
	server.group.Wait()
}

func (server *serverKV) keyToShard(Key string) *shardElem {
	shardidx := crc(Key) % shardCount
	return server.baseList[shardidx]
}

func (server *serverKV) SetValue(Key string, Value datamodel.CustomDataType, Session time.Duration) (*kvElementh, error) {
	elem, ok := server.GetValue(Key)
	if ok {
		if !atomic.CompareAndSwapInt32(&elem.lock, 0, 1) {
			if elem.lockSession != Session {
				return nil, errors.New("record busy")
			}
		} else {
			defer atomic.StoreInt32(&elem.lock, 0)
		}
		elem.value = Value
		elem.ttc = time.Since(startTime)
		elem.ttl = maxTime
		return elem, nil
	}
	shard := server.keyToShard(Key)
	shard.Lock()
	defer shard.Unlock()
	// check on append
	elem, ok = shard.mapkv[Key]
	if ok {
		if !atomic.CompareAndSwapInt32(&elem.lock, 0, 1) {
			return nil, errors.New("record busy")
		}
		defer atomic.StoreInt32(&elem.lock, 0)
		elem.value = Value
		elem.ttc = time.Since(startTime)
		elem.ttl = maxTime
		return elem, nil
	}
	elem = &kvElementh{
		value: Value,
		ttc:   time.Since(startTime),
		ttl:   maxTime,
	}
	shard.mapkv[Key] = elem
	return elem, nil
}

func (server *serverKV) GetValue(Key string) (*kvElementh, bool) {
	shard := server.keyToShard(Key)
	shard.RLock()
	defer shard.RUnlock()
	el, ok := shard.mapkv[Key]
	return el, ok
}

func (server *serverKV) LockValue(Key string, Session time.Duration) error {
	shard := server.keyToShard(Key)
	shard.RLock()
	defer shard.RUnlock()
	elem, ok := shard.mapkv[Key]
	if !ok {
		return errors.New("record not found")
	}
	if !atomic.CompareAndSwapInt32(&elem.lock, 0, 1) {
		if elem.lockSession != Session {
			return errors.New("record busy")
		}
	}
	return nil
}

func (server *serverKV) UnlockValue(Key string, Session time.Duration) error {
	shard := server.keyToShard(Key)
	shard.RLock()
	defer shard.RUnlock()
	elem, ok := shard.mapkv[Key]
	if !ok {
		return errors.New("record not found")
	}
	if !atomic.CompareAndSwapInt32(&elem.lock, 0, 1) {
		if elem.lockSession != Session {
			return errors.New("record busy")
		}
	}
	elem.ttc = time.Since(startTime)
	elem.ttl = maxTime
	atomic.StoreInt32(&elem.lock, 0)
	return nil
}

func (server *serverKV) SetTTL(Key string, TTL time.Duration, Session time.Duration) error {
	var needRLock bool
	shard := server.keyToShard(Key)
	shard.RLock()
	defer func() {
		if needRLock {
			shard.RUnlock()
		}
	}()
	elem, ok := shard.mapkv[Key]
	if !ok {
		return errors.New("record not found")
	}
	if !atomic.CompareAndSwapInt32(&elem.lock, 0, 1) {
		if elem.lockSession != Session {
			return errors.New("record busy")
		}
	}
	elem.ttl = elem.ttc + TTL
	if elem.ttl < time.Since(startTime) {
		shard.RUnlock()
		needRLock = false
		shard.Lock()
		defer shard.Unlock()
		delete(shard.mapkv, Key)
	} else {
		next := atomic.LoadInt64(&shard.nextTTLProcessing)
		if next > int64(elem.ttl) {
			shard.ttlMutex.Lock()
			defer shard.ttlMutex.Unlock()
			if shard.nextTTLProcessing > int64(elem.ttl) {
				shard.nextTTLProcessing = int64(elem.ttl)
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package fsm

import (
	"fmt"
	"sort"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/swf"
	. "github.com/sclasen/swfsm/log"
	. "github.com/sclasen/swfsm/sugar"
)

const (
	FilterStatusAll                  = "ALL"                    // open + closed
	FilterStatusOpen                 = "OPEN"                   // open only
	FilterStatusOpenPriority         = "OPEN_PRIORITY"          // open (+ closed, only if open is totally empty)
	FilterStatusOpenPriorityWorkflow = "OPEN_PRIORITY_WORKFLOW" // open (+ closed, only if open not present workflow-by-workflow)
	FilterStatusClosed               = "CLOSED"                 // closed only
)

type Finder interface {
	FindAll(*FindInput) (*FindOutput, error)
	FindLatestByWorkflowID(workflowID string) (*swf.WorkflowExecution, error)
}

func NewFinder(domain string, c ClientSWFOps) Finder {
	return &finder{
		domain: domain,
		c:      c,
	}
}

type FindInput struct {
	MaximumPageSize *int64

	OpenNextPageToken   *string
	ClosedNextPageToken *string

	ReverseOrder *bool

	StatusFilter string

	StartTimeFilter *swf.ExecutionTimeFilter
	CloseTimeFilter *swf.ExecutionTimeFilter // only closed

	ExecutionFilter   *swf.WorkflowExecutionFilter
	TagFilter         *swf.TagFilter
	TypeFilter        *swf.WorkflowTypeFilter
	CloseStatusFilter *swf.CloseStatusFilter // only closed

	workflowIdIndex map[string]struct{} // private on input for passing for use on subsequent pages
}

type FindOutput struct {
	ExecutionInfos      []*swf.WorkflowExecutionInfo
	OpenNextPageToken   *string
	ClosedNextPageToken *string
}

type finder struct {
	domain string
	c      ClientSWFOps
}

func (f *finder) FindAll(input *FindInput) (output *FindOutput, err error) {
	if input.StatusFilter == "" {
		if (input.CloseStatusFilter != nil && input.CloseStatusFilter.Status != nil) ||
			(input.CloseTimeFilter != nil &&
				(input.CloseTimeFilter.OldestDate != nil || input.CloseTimeFilter.LatestDate != nil)) {
			input.StatusFilter = FilterStatusClosed
		} else {
			input.StatusFilter = FilterStatusOpenPriority
		}
	}

	if !stringsContain([]string{
		FilterStatusAll,
		FilterStatusOpen,
		FilterStatusOpenPriority,
		FilterStatusOpenPriorityWorkflow,
		FilterStatusClosed,
	},
		input.StatusFilter) {
		return nil, fmt.Errorf("Invalid status filter")
	}

	selectiveFilter := f.mostSelectiveFilters(input)
	output = &FindOutput{}

	if input.OpenNextPageToken != nil || (input.OpenNextPageToken == nil && input.ClosedNextPageToken == nil) {
		if stringsContain([]string{
			FilterStatusAll,
			FilterStatusOpen,
			FilterStatusOpenPriority,
			FilterStatusOpenPriorityWorkflow,
		}, input.StatusFilter) {
			openInput := &swf.ListOpenWorkflowExecutionsInput{
				Domain:          &f.domain,
				ReverseOrder:    input.ReverseOrder,
				MaximumPageSize: input.MaximumPageSize,
				NextPageToken:   input.OpenNextPageToken,
				StartTimeFilter: selectiveFilter.StartTimeFilter,
				ExecutionFilter: selectiveFilter.ExecutionFilter,
				TagFilter:       selectiveFilter.TagFilter,
				TypeFilter:      selectiveFilter.TypeFilter,
			}

			resp, err := f.c.ListOpenWorkflowExecutions(openInput)
			if err != nil {
				return nil, err
			}

			f.append(input, output, resp.ExecutionInfos)
			output.OpenNextPageToken = resp.NextPageToken
		}
	}

	if input.StatusFilter == FilterStatusOpenPriority {
		output.ExecutionInfos = f.applyInputLocally(input, output.ExecutionInfos)
		if len(output.ExecutionInfos) > 0 {
			return output, nil
		}
	}

	if input.ClosedNextPageToken != nil || (input.OpenNextPageToken == nil && input.ClosedNextPageToken == nil) {
		if stringsContain([]string{
			FilterStatusAll,
			FilterStatusClosed,
			FilterStatusOpenPriority,
			FilterStatusOpenPriorityWorkflow,
		}, input.StatusFilter) {
			closedInput := &swf.ListClosedWorkflowExecutionsInput{
				Domain:            &f.domain,
				ReverseOrder:      input.ReverseOrder,
				MaximumPageSize:   input.MaximumPageSize,
				NextPageToken:     input.ClosedNextPageToken,
				StartTimeFilter:   selectiveFilter.StartTimeFilter,
				CloseTimeFilter:   selectiveFilter.CloseTimeFilter,
				ExecutionFilter:   selectiveFilter.ExecutionFilter,
				TagFilter:         selectiveFilter.TagFilter,
				TypeFilter:        selectiveFilter.TypeFilter,
				CloseStatusFilter: selectiveFilter.CloseStatusFilter,
			}

			resp, err := f.c.ListClosedWorkflowExecutions(closedInput)

			if err != nil {
				return nil, err
			}

			f.append(input, output, resp.ExecutionInfos)
			output.ClosedNextPageToken = resp.NextPageToken
		}
	}

	output.ExecutionInfos = f.applyInputLocally(input, output.ExecutionInfos)

	return output, nil
}

type listAnyWorkflowExecutionsInput swf.ListClosedWorkflowExecutionsInput

func (f *finder) mostSelectiveFilters(input *FindInput) listAnyWorkflowExecutionsInput {
	filter := &listAnyWorkflowExecutionsInput{}
	f.setMostSelectiveMetadataFilter(input, filter)
	f.setMostSelectiveTimeFilter(input, filter)
	return *filter
}

func (f *finder) setMostSelectiveMetadataFilter(input *FindInput, output *listAnyWorkflowExecutionsInput) {
	if input.ExecutionFilter != nil {
		output.ExecutionFilter = input.ExecutionFilter
		return
	}

	if input.TagFilter != nil {
		output.TagFilter = input.TagFilter
		return
	}

	if input.TypeFilter != nil {
		output.TypeFilter = input.TypeFilter
		return
	}

	if input.CloseStatusFilter != nil {
		output.CloseStatusFilter = input.CloseStatusFilter
		return
	}
}

func (f *finder) setMostSelectiveTimeFilter(input *FindInput, output *listAnyWorkflowExecutionsInput) {
	if input.StartTimeFilter != nil {
		output.StartTimeFilter = input.StartTimeFilter
		return
	}

	if input.CloseTimeFilter != nil {
		output.CloseTimeFilter = input.CloseTimeFilter
		return
	}
}

func (f *finder) append(input *FindInput, output *FindOutput, infos []*swf.WorkflowExecutionInfo) {
	if input.workflowIdIndex == nil {
		input.workflowIdIndex = make(map[string]struct{})
	}

	for _, info := range infos {
		if input.StatusFilter == FilterStatusOpenPriorityWorkflow {
			if _, ok := input.workflowIdIndex[*info.Execution.WorkflowId]; ok {
				continue
			}
			input.workflowIdIndex[*info.Execution.WorkflowId] = struct{}{}
		}

		output.ExecutionInfos = append(output.ExecutionInfos, info)
	}
}

func (f *finder) applyInputLocally(input *FindInput, infos []*swf.WorkflowExecutionInfo) []*swf.WorkflowExecutionInfo {
	applied := []*swf.WorkflowExecutionInfo{}

	for _, info := range infos {
		if input.ExecutionFilter != nil {
			if *input.ExecutionFilter.WorkflowId != *info.Execution.WorkflowId {
				continue
			}
		}

		if input.TagFilter != nil {
			found := false
			for _, tag := range info.TagList {
				if *input.TagFilter.Tag == *tag {
					found = true
					break
				}
			}
			if !found {
				continue
			}
		}

		if input.TypeFilter != nil && input.TypeFilter.Name != nil {
			if *input.TypeFilter.Name != *info.WorkflowType.Name {
				continue
			}
			if input.TypeFilter.Version != nil && input.TypeFilter.Version != nil {
				if *input.TypeFilter.Version != *info.WorkflowType.Version {
					continue
				}
			}
		}

		if input.CloseStatusFilter != nil && input.CloseStatusFilter.Status != nil && info.CloseStatus != nil {
			if *input.CloseStatusFilter.Status != *info.CloseStatus {
				continue
			}
		}

		if input.StartTimeFilter != nil {
			if input.StartTimeFilter.OldestDate != nil && (*input.StartTimeFilter.OldestDate).After(*info.StartTimestamp) {
				continue
			}

			if input.StartTimeFilter.LatestDate != nil && (*input.StartTimeFilter.LatestDate).Before(*info.StartTimestamp) {
				continue
			}
		}

		if input.CloseTimeFilter != nil && info.CloseTimestamp != nil {
			if input.CloseTimeFilter.OldestDate != nil && (*input.CloseTimeFilter.OldestDate).After(*info.CloseTimestamp) {
				continue
			}

			if input.CloseTimeFilter.LatestDate != nil && (*input.CloseTimeFilter.LatestDate).Before(*info.CloseTimestamp) {
				continue
			}
		}

		applied = append(applied, info)
	}

	// swf default order is descending, so reverseOrder is ascending
	if input.ReverseOrder != nil && *input.ReverseOrder {
		sort.Sort(sortExecutionInfos(applied))
	} else {
		sort.Sort(sort.Reverse(sortExecutionInfos(applied)))
	}

	if input.MaximumPageSize != nil && *input.MaximumPageSize > 0 && int64(len(applied)) > *input.MaximumPageSize {
		applied = applied[:*input.MaximumPageSize]
	}

	return applied
}

func stringsContain(haystack []string, needle string) bool {
	for _, h := range haystack {
		if h == needle {
			return true
		}
	}
	return false
}

type sortExecutionInfos []*swf.WorkflowExecutionInfo

func (e sortExecutionInfos) Len() int      { return len(e) }
func (e sortExecutionInfos) Swap(i, j int) { e[i], e[j] = e[j], e[i] }
func (e sortExecutionInfos) Less(i, j int) bool {
	return (*e[i].StartTimestamp).Before(*e[j].StartTimestamp)
}

func (f *finder) FindLatestByWorkflowID(workflowID string) (exec *swf.WorkflowExecution, err error) {
	output, err := f.FindAll(&FindInput{
		StatusFilter:    FilterStatusOpenPriority,
		MaximumPageSize: I(1),
		ReverseOrder:    aws.Bool(false), // keep descending default
		StartTimeFilter: &swf.ExecutionTimeFilter{OldestDate: aws.Time(time.Unix(0, 0))},
		ExecutionFilter: &swf.WorkflowExecutionFilter{
			WorkflowId: S(workflowID),
		},
	})

	if err != nil {
		if ae, ok := err.(awserr.Error); ok {
			Log.Printf("component=client fn=findExecution at=list-open error-type=%s message=%s", ae.Code(), ae.Message())
		} else {
			Log.Printf("component=client fn=findExecution at=list-open error=%q", err)
		}
		return nil, err
	}

	if len(output.ExecutionInfos) == 1 {
		return output.ExecutionInfos[0].Execution, nil
	}
	return nil, nil
}
<commit_msg>Move f.workflowIdIndex into finder struct and add Reset() to clear it<commit_after>package fsm

import (
	"fmt"
	"sort"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/swf"
	. "github.com/sclasen/swfsm/log"
	. "github.com/sclasen/swfsm/sugar"
)

const (
	FilterStatusAll                  = "ALL"                    // open + closed
	FilterStatusOpen                 = "OPEN"                   // open only
	FilterStatusOpenPriority         = "OPEN_PRIORITY"          // open (+ closed, only if open is totally empty)
	FilterStatusOpenPriorityWorkflow = "OPEN_PRIORITY_WORKFLOW" // open (+ closed, only if open not present workflow-by-workflow)
	FilterStatusClosed               = "CLOSED"                 // closed only
)

type Finder interface {
	FindAll(*FindInput) (*FindOutput, error)
	FindLatestByWorkflowID(workflowID string) (*swf.WorkflowExecution, error)
	Reset()
}

func NewFinder(domain string, c ClientSWFOps) Finder {
	return &finder{
		domain:          domain,
		c:               c,
		workflowIdIndex: make(map[string]struct{}),
	}
}

type FindInput struct {
	MaximumPageSize *int64

	OpenNextPageToken   *string
	ClosedNextPageToken *string

	ReverseOrder *bool

	StatusFilter string

	StartTimeFilter *swf.ExecutionTimeFilter
	CloseTimeFilter *swf.ExecutionTimeFilter // only closed

	ExecutionFilter   *swf.WorkflowExecutionFilter
	TagFilter         *swf.TagFilter
	TypeFilter        *swf.WorkflowTypeFilter
	CloseStatusFilter *swf.CloseStatusFilter // only closed
}

type FindOutput struct {
	ExecutionInfos      []*swf.WorkflowExecutionInfo
	OpenNextPageToken   *string
	ClosedNextPageToken *string
}

type finder struct {
	domain          string
	c               ClientSWFOps
	workflowIdIndex map[string]struct{}
}

func (f *finder) FindAll(input *FindInput) (output *FindOutput, err error) {
	if input.StatusFilter == "" {
		if (input.CloseStatusFilter != nil && input.CloseStatusFilter.Status != nil) ||
			(input.CloseTimeFilter != nil &&
				(input.CloseTimeFilter.OldestDate != nil || input.CloseTimeFilter.LatestDate != nil)) {
			input.StatusFilter = FilterStatusClosed
		} else {
			input.StatusFilter = FilterStatusOpenPriority
		}
	}

	if !stringsContain([]string{
		FilterStatusAll,
		FilterStatusOpen,
		FilterStatusOpenPriority,
		FilterStatusOpenPriorityWorkflow,
		FilterStatusClosed,
	},
		input.StatusFilter) {
		return nil, fmt.Errorf("Invalid status filter")
	}

	selectiveFilter := f.mostSelectiveFilters(input)
	output = &FindOutput{}

	if input.OpenNextPageToken != nil || (input.OpenNextPageToken == nil && input.ClosedNextPageToken == nil) {
		if stringsContain([]string{
			FilterStatusAll,
			FilterStatusOpen,
			FilterStatusOpenPriority,
			FilterStatusOpenPriorityWorkflow,
		}, input.StatusFilter) {
			openInput := &swf.ListOpenWorkflowExecutionsInput{
				Domain:          &f.domain,
				ReverseOrder:    input.ReverseOrder,
				MaximumPageSize: input.MaximumPageSize,
				NextPageToken:   input.OpenNextPageToken,
				StartTimeFilter: selectiveFilter.StartTimeFilter,
				ExecutionFilter: selectiveFilter.ExecutionFilter,
				TagFilter:       selectiveFilter.TagFilter,
				TypeFilter:      selectiveFilter.TypeFilter,
			}

			resp, err := f.c.ListOpenWorkflowExecutions(openInput)
			if err != nil {
				return nil, err
			}

			f.append(input, output, resp.ExecutionInfos)
			output.OpenNextPageToken = resp.NextPageToken
		}
	}

	if input.StatusFilter == FilterStatusOpenPriority {
		output.ExecutionInfos = f.applyInputLocally(input, output.ExecutionInfos)
		if len(output.ExecutionInfos) > 0 {
			return output, nil
		}
	}

	if input.ClosedNextPageToken != nil || (input.OpenNextPageToken == nil && input.ClosedNextPageToken == nil) {
		if stringsContain([]string{
			FilterStatusAll,
			FilterStatusClosed,
			FilterStatusOpenPriority,
			FilterStatusOpenPriorityWorkflow,
		}, input.StatusFilter) {
			closedInput := &swf.ListClosedWorkflowExecutionsInput{
				Domain:            &f.domain,
				ReverseOrder:      input.ReverseOrder,
				MaximumPageSize:   input.MaximumPageSize,
				NextPageToken:     input.ClosedNextPageToken,
				StartTimeFilter:   selectiveFilter.StartTimeFilter,
				CloseTimeFilter:   selectiveFilter.CloseTimeFilter,
				ExecutionFilter:   selectiveFilter.ExecutionFilter,
				TagFilter:         selectiveFilter.TagFilter,
				TypeFilter:        selectiveFilter.TypeFilter,
				CloseStatusFilter: selectiveFilter.CloseStatusFilter,
			}

			resp, err := f.c.ListClosedWorkflowExecutions(closedInput)

			if err != nil {
				return nil, err
			}

			f.append(input, output, resp.ExecutionInfos)
			output.ClosedNextPageToken = resp.NextPageToken
		}
	}

	output.ExecutionInfos = f.applyInputLocally(input, output.ExecutionInfos)

	return output, nil
}

type listAnyWorkflowExecutionsInput swf.ListClosedWorkflowExecutionsInput

func (f *finder) mostSelectiveFilters(input *FindInput) listAnyWorkflowExecutionsInput {
	filter := &listAnyWorkflowExecutionsInput{}
	f.setMostSelectiveMetadataFilter(input, filter)
	f.setMostSelectiveTimeFilter(input, filter)
	return *filter
}

func (f *finder) setMostSelectiveMetadataFilter(input *FindInput, output *listAnyWorkflowExecutionsInput) {
	if input.ExecutionFilter != nil {
		output.ExecutionFilter = input.ExecutionFilter
		return
	}

	if input.TagFilter != nil {
		output.TagFilter = input.TagFilter
		return
	}

	if input.TypeFilter != nil {
		output.TypeFilter = input.TypeFilter
		return
	}

	if input.CloseStatusFilter != nil {
		output.CloseStatusFilter = input.CloseStatusFilter
		return
	}
}

func (f *finder) setMostSelectiveTimeFilter(input *FindInput, output *listAnyWorkflowExecutionsInput) {
	if input.StartTimeFilter != nil {
		output.StartTimeFilter = input.StartTimeFilter
		return
	}

	if input.CloseTimeFilter != nil {
		output.CloseTimeFilter = input.CloseTimeFilter
		return
	}
}

func (f *finder) append(input *FindInput, output *FindOutput, infos []*swf.WorkflowExecutionInfo) {
	for _, info := range infos {
		if input.StatusFilter == FilterStatusOpenPriorityWorkflow {
			if _, ok := f.workflowIdIndex[*info.Execution.WorkflowId]; ok {
				continue
			}
			f.workflowIdIndex[*info.Execution.WorkflowId] = struct{}{}
		}

		output.ExecutionInfos = append(output.ExecutionInfos, info)
	}
}

func (f *finder) applyInputLocally(input *FindInput, infos []*swf.WorkflowExecutionInfo) []*swf.WorkflowExecutionInfo {
	applied := []*swf.WorkflowExecutionInfo{}

	for _, info := range infos {
		if input.ExecutionFilter != nil {
			if *input.ExecutionFilter.WorkflowId != *info.Execution.WorkflowId {
				continue
			}
		}

		if input.TagFilter != nil {
			found := false
			for _, tag := range info.TagList {
				if *input.TagFilter.Tag == *tag {
					found = true
					break
				}
			}
			if !found {
				continue
			}
		}

		if input.TypeFilter != nil && input.TypeFilter.Name != nil {
			if *input.TypeFilter.Name != *info.WorkflowType.Name {
				continue
			}
			if input.TypeFilter.Version != nil && input.TypeFilter.Version != nil {
				if *input.TypeFilter.Version != *info.WorkflowType.Version {
					continue
				}
			}
		}

		if input.CloseStatusFilter != nil && input.CloseStatusFilter.Status != nil && info.CloseStatus != nil {
			if *input.CloseStatusFilter.Status != *info.CloseStatus {
				continue
			}
		}

		if input.StartTimeFilter != nil {
			if input.StartTimeFilter.OldestDate != nil && (*input.StartTimeFilter.OldestDate).After(*info.StartTimestamp) {
				continue
			}

			if input.StartTimeFilter.LatestDate != nil && (*input.StartTimeFilter.LatestDate).Before(*info.StartTimestamp) {
				continue
			}
		}

		if input.CloseTimeFilter != nil && info.CloseTimestamp != nil {
			if input.CloseTimeFilter.OldestDate != nil && (*input.CloseTimeFilter.OldestDate).After(*info.CloseTimestamp) {
				continue
			}

			if input.CloseTimeFilter.LatestDate != nil && (*input.CloseTimeFilter.LatestDate).Before(*info.CloseTimestamp) {
				continue
			}
		}

		applied = append(applied, info)
	}

	// swf default order is descending, so reverseOrder is ascending
	if input.ReverseOrder != nil && *input.ReverseOrder {
		sort.Sort(sortExecutionInfos(applied))
	} else {
		sort.Sort(sort.Reverse(sortExecutionInfos(applied)))
	}

	if input.MaximumPageSize != nil && *input.MaximumPageSize > 0 && int64(len(applied)) > *input.MaximumPageSize {
		applied = applied[:*input.MaximumPageSize]
	}

	return applied
}

func stringsContain(haystack []string, needle string) bool {
	for _, h := range haystack {
		if h == needle {
			return true
		}
	}
	return false
}

type sortExecutionInfos []*swf.WorkflowExecutionInfo

func (e sortExecutionInfos) Len() int      { return len(e) }
func (e sortExecutionInfos) Swap(i, j int) { e[i], e[j] = e[j], e[i] }
func (e sortExecutionInfos) Less(i, j int) bool {
	return (*e[i].StartTimestamp).Before(*e[j].StartTimestamp)
}

func (f *finder) FindLatestByWorkflowID(workflowID string) (exec *swf.WorkflowExecution, err error) {
	output, err := f.FindAll(&FindInput{
		StatusFilter:    FilterStatusOpenPriority,
		MaximumPageSize: I(1),
		ReverseOrder:    aws.Bool(false), // keep descending default
		StartTimeFilter: &swf.ExecutionTimeFilter{OldestDate: aws.Time(time.Unix(0, 0))},
		ExecutionFilter: &swf.WorkflowExecutionFilter{
			WorkflowId: S(workflowID),
		},
	})

	if err != nil {
		if ae, ok := err.(awserr.Error); ok {
			Log.Printf("component=client fn=findExecution at=list-open error-type=%s message=%s", ae.Code(), ae.Message())
		} else {
			Log.Printf("component=client fn=findExecution at=list-open error=%q", err)
		}
		return nil, err
	}

	if len(output.ExecutionInfos) == 1 {
		return output.ExecutionInfos[0].Execution, nil
	}
	return nil, nil
}

func (f *finder) Reset() {
	f.workflowIdIndex = make(map[string]struct{})
}
<|endoftext|>


<commit_before>/*

This provides Run for use in creating test suites

To use this declare a TestMain

// TestMain drives the tests
func TestMain(m *testing.M) {
	fstest.TestMain(m)
}

And then make and destroy a Run in each test

func TestMkdir(t *testing.T) {
	r := fstest.NewRun(t)
	defer r.Finalise()
	// test stuff
}

This will make r.Fremote and r.Flocal for a remote remote and a local
remote.  The remote is determined by the -remote flag passed in.

*/

package fstest

import (
	"bytes"
	"flag"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
	"sort"
	"testing"
	"time"

	"github.com/ncw/rclone/fs"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// Run holds the remotes for a test run
type Run struct {
	LocalName    string
	Flocal       fs.Fs
	Fremote      fs.Fs
	FremoteName  string
	cleanRemote  func()
	mkdir        map[string]bool // whether the remote has been made yet for the fs name
	Logf, Fatalf func(text string, args ...interface{})
}

// TestMain drives the tests
func TestMain(m *testing.M) {
	flag.Parse()
	if !*Individual {
		oneRun = newRun()
	}
	rc := m.Run()
	if !*Individual {
		oneRun.Finalise()
	}
	os.Exit(rc)
}

// oneRun holds the master run data if individual is not set
var oneRun *Run

// newRun initialise the remote and local for testing and returns a
// run object.
//
// r.Flocal is an empty local Fs
// r.Fremote is an empty remote Fs
//
// Finalise() will tidy them away when done.
func newRun() *Run {
	r := &Run{
		Logf:   log.Printf,
		Fatalf: log.Fatalf,
		mkdir:  make(map[string]bool),
	}

	Initialise()

	var err error
	r.Fremote, r.FremoteName, r.cleanRemote, err = RandomRemote(*RemoteName, *SubDir)
	if err != nil {
		r.Fatalf("Failed to open remote %q: %v", *RemoteName, err)
	}

	r.LocalName, err = ioutil.TempDir("", "rclone")
	if err != nil {
		r.Fatalf("Failed to create temp dir: %v", err)
	}
	r.LocalName = filepath.ToSlash(r.LocalName)
	r.Flocal, err = fs.NewFs(r.LocalName)
	if err != nil {
		r.Fatalf("Failed to make %q: %v", r.LocalName, err)
	}
	fs.CalculateModifyWindow(r.Fremote, r.Flocal)
	return r
}

// dirsToRemove sorts by string length
type dirsToRemove []string

func (d dirsToRemove) Len() int           { return len(d) }
func (d dirsToRemove) Swap(i, j int)      { d[i], d[j] = d[j], d[i] }
func (d dirsToRemove) Less(i, j int) bool { return len(d[i]) > len(d[j]) }

// NewRun initialise the remote and local for testing and returns a
// run object.  Call this from the tests.
//
// r.Flocal is an empty local Fs
// r.Fremote is an empty remote Fs
//
// Finalise() will tidy them away when done.
func NewRun(t *testing.T) *Run {
	var r *Run
	if *Individual {
		r = newRun()
	} else {
		// If not individual, use the global one with the clean method overridden
		r = new(Run)
		*r = *oneRun
		r.cleanRemote = func() {
			var toDelete dirsToRemove
			require.NoError(t, fs.Walk(r.Fremote, "", true, -1, func(dirPath string, entries fs.DirEntries, err error) error {
				if err != nil {
					if err == fs.ErrorDirNotFound {
						return nil
					}
					t.Fatalf("Error listing: %v", err)
				}
				for _, entry := range entries {
					switch x := entry.(type) {
					case fs.Object:
						err = x.Remove()
						if err != nil {
							t.Errorf("Error removing file %q: %v", x.Remote(), err)
						}
					case fs.Directory:
						toDelete = append(toDelete, x.Remote())
					}
				}
				return nil
			}))
			sort.Sort(toDelete)
			for _, dir := range toDelete {
				err := r.Fremote.Rmdir(dir)
				if err != nil {
					t.Errorf("Error removing dir %q: %v", dir, err)
				}
			}
			// Check remote is empty
			CheckItems(t, r.Fremote)
		}
	}
	r.Logf = t.Logf
	r.Fatalf = t.Fatalf
	r.Logf("Remote %q, Local %q, Modify Window %q", r.Fremote, r.Flocal, fs.Config.ModifyWindow)
	return r
}

// RenameFile renames a file in local
func (r *Run) RenameFile(item Item, newpath string) Item {
	oldFilepath := path.Join(r.LocalName, item.Path)
	newFilepath := path.Join(r.LocalName, newpath)
	if err := os.Rename(oldFilepath, newFilepath); err != nil {
		r.Fatalf("Failed to rename file from %q to %q: %v", item.Path, newpath, err)
	}

	item.Path = newpath

	return item
}

// WriteFile writes a file to local
func (r *Run) WriteFile(filePath, content string, t time.Time) Item {
	item := NewItem(filePath, content, t)
	// FIXME make directories?
	filePath = path.Join(r.LocalName, filePath)
	dirPath := path.Dir(filePath)
	err := os.MkdirAll(dirPath, 0770)
	if err != nil {
		r.Fatalf("Failed to make directories %q: %v", dirPath, err)
	}
	err = ioutil.WriteFile(filePath, []byte(content), 0600)
	if err != nil {
		r.Fatalf("Failed to write file %q: %v", filePath, err)
	}
	err = os.Chtimes(filePath, t, t)
	if err != nil {
		r.Fatalf("Failed to chtimes file %q: %v", filePath, err)
	}
	return item
}

// ForceMkdir creates the remote
func (r *Run) ForceMkdir(f fs.Fs) {
	err := f.Mkdir("")
	if err != nil {
		r.Fatalf("Failed to mkdir %q: %v", f, err)
	}
	r.mkdir[f.String()] = true
}

// Mkdir creates the remote if it hasn't been created already
func (r *Run) Mkdir(f fs.Fs) {
	if !r.mkdir[f.String()] {
		r.ForceMkdir(f)
	}
}

// WriteObjectTo writes an object to the fs, remote passed in
func (r *Run) WriteObjectTo(f fs.Fs, remote, content string, modTime time.Time, useUnchecked bool) Item {
	put := f.Put
	if useUnchecked {
		put = f.Features().PutUnchecked
		if put == nil {
			r.Fatalf("Fs doesn't support PutUnchecked")
		}
	}
	r.Mkdir(f)
	const maxTries = 10
	for tries := 1; ; tries++ {
		in := bytes.NewBufferString(content)
		objinfo := fs.NewStaticObjectInfo(remote, modTime, int64(len(content)), true, nil, nil)
		_, err := put(in, objinfo)
		if err == nil {
			break
		}
		// Retry if err returned a retry error
		if fs.IsRetryError(err) && tries < maxTries {
			r.Logf("Retry Put of %q to %v: %d/%d (%v)", remote, f, tries, maxTries, err)
			time.Sleep(2 * time.Second)
			continue
		}
		r.Fatalf("Failed to put %q to %q: %v", remote, f, err)
	}
	return NewItem(remote, content, modTime)
}

// WriteObject writes an object to the remote
func (r *Run) WriteObject(remote, content string, modTime time.Time) Item {
	return r.WriteObjectTo(r.Fremote, remote, content, modTime, false)
}

// WriteUncheckedObject writes an object to the remote not checking for duplicates
func (r *Run) WriteUncheckedObject(remote, content string, modTime time.Time) Item {
	return r.WriteObjectTo(r.Fremote, remote, content, modTime, true)
}

// WriteBoth calls WriteObject and WriteFile with the same arguments
func (r *Run) WriteBoth(remote, content string, modTime time.Time) Item {
	r.WriteFile(remote, content, modTime)
	return r.WriteObject(remote, content, modTime)
}

// CheckWithDuplicates does a test but allows duplicates
func (r *Run) CheckWithDuplicates(t *testing.T, items ...Item) {
	objects, size, err := fs.Count(r.Fremote)
	require.NoError(t, err)
	assert.Equal(t, int64(len(items)), objects)
	wantSize := int64(0)
	for _, item := range items {
		wantSize += item.Size
	}
	assert.Equal(t, wantSize, size)
}

// Clean the temporary directory
func (r *Run) cleanTempDir() {
	err := os.RemoveAll(r.LocalName)
	if err != nil {
		r.Logf("Failed to clean temporary directory %q: %v", r.LocalName, err)
	}
}

// Finalise cleans the remote and local
func (r *Run) Finalise() {
	// r.Logf("Cleaning remote %q", r.Fremote)
	r.cleanRemote()
	// r.Logf("Cleaning local %q", r.LocalName)
	r.cleanTempDir()
}
<commit_msg>fstest: check no files or directories between runs<commit_after>/*

This provides Run for use in creating test suites

To use this declare a TestMain

// TestMain drives the tests
func TestMain(m *testing.M) {
	fstest.TestMain(m)
}

And then make and destroy a Run in each test

func TestMkdir(t *testing.T) {
	r := fstest.NewRun(t)
	defer r.Finalise()
	// test stuff
}

This will make r.Fremote and r.Flocal for a remote remote and a local
remote.  The remote is determined by the -remote flag passed in.

*/

package fstest

import (
	"bytes"
	"flag"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
	"sort"
	"testing"
	"time"

	"github.com/ncw/rclone/fs"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// Run holds the remotes for a test run
type Run struct {
	LocalName    string
	Flocal       fs.Fs
	Fremote      fs.Fs
	FremoteName  string
	cleanRemote  func()
	mkdir        map[string]bool // whether the remote has been made yet for the fs name
	Logf, Fatalf func(text string, args ...interface{})
}

// TestMain drives the tests
func TestMain(m *testing.M) {
	flag.Parse()
	if !*Individual {
		oneRun = newRun()
	}
	rc := m.Run()
	if !*Individual {
		oneRun.Finalise()
	}
	os.Exit(rc)
}

// oneRun holds the master run data if individual is not set
var oneRun *Run

// newRun initialise the remote and local for testing and returns a
// run object.
//
// r.Flocal is an empty local Fs
// r.Fremote is an empty remote Fs
//
// Finalise() will tidy them away when done.
func newRun() *Run {
	r := &Run{
		Logf:   log.Printf,
		Fatalf: log.Fatalf,
		mkdir:  make(map[string]bool),
	}

	Initialise()

	var err error
	r.Fremote, r.FremoteName, r.cleanRemote, err = RandomRemote(*RemoteName, *SubDir)
	if err != nil {
		r.Fatalf("Failed to open remote %q: %v", *RemoteName, err)
	}

	r.LocalName, err = ioutil.TempDir("", "rclone")
	if err != nil {
		r.Fatalf("Failed to create temp dir: %v", err)
	}
	r.LocalName = filepath.ToSlash(r.LocalName)
	r.Flocal, err = fs.NewFs(r.LocalName)
	if err != nil {
		r.Fatalf("Failed to make %q: %v", r.LocalName, err)
	}
	fs.CalculateModifyWindow(r.Fremote, r.Flocal)
	return r
}

// dirsToRemove sorts by string length
type dirsToRemove []string

func (d dirsToRemove) Len() int           { return len(d) }
func (d dirsToRemove) Swap(i, j int)      { d[i], d[j] = d[j], d[i] }
func (d dirsToRemove) Less(i, j int) bool { return len(d[i]) > len(d[j]) }

// NewRun initialise the remote and local for testing and returns a
// run object.  Call this from the tests.
//
// r.Flocal is an empty local Fs
// r.Fremote is an empty remote Fs
//
// Finalise() will tidy them away when done.
func NewRun(t *testing.T) *Run {
	var r *Run
	if *Individual {
		r = newRun()
	} else {
		// If not individual, use the global one with the clean method overridden
		r = new(Run)
		*r = *oneRun
		r.cleanRemote = func() {
			var toDelete dirsToRemove
			require.NoError(t, fs.Walk(r.Fremote, "", true, -1, func(dirPath string, entries fs.DirEntries, err error) error {
				if err != nil {
					if err == fs.ErrorDirNotFound {
						return nil
					}
					t.Fatalf("Error listing: %v", err)
				}
				for _, entry := range entries {
					switch x := entry.(type) {
					case fs.Object:
						err = x.Remove()
						if err != nil {
							t.Errorf("Error removing file %q: %v", x.Remote(), err)
						}
					case fs.Directory:
						toDelete = append(toDelete, x.Remote())
					}
				}
				return nil
			}))
			sort.Sort(toDelete)
			for _, dir := range toDelete {
				err := r.Fremote.Rmdir(dir)
				if err != nil {
					t.Errorf("Error removing dir %q: %v", dir, err)
				}
			}
			// Check remote is empty
			CheckListingWithPrecision(t, r.Fremote, []Item{}, []string{}, r.Fremote.Precision())
		}
	}
	r.Logf = t.Logf
	r.Fatalf = t.Fatalf
	r.Logf("Remote %q, Local %q, Modify Window %q", r.Fremote, r.Flocal, fs.Config.ModifyWindow)
	return r
}

// RenameFile renames a file in local
func (r *Run) RenameFile(item Item, newpath string) Item {
	oldFilepath := path.Join(r.LocalName, item.Path)
	newFilepath := path.Join(r.LocalName, newpath)
	if err := os.Rename(oldFilepath, newFilepath); err != nil {
		r.Fatalf("Failed to rename file from %q to %q: %v", item.Path, newpath, err)
	}

	item.Path = newpath

	return item
}

// WriteFile writes a file to local
func (r *Run) WriteFile(filePath, content string, t time.Time) Item {
	item := NewItem(filePath, content, t)
	// FIXME make directories?
	filePath = path.Join(r.LocalName, filePath)
	dirPath := path.Dir(filePath)
	err := os.MkdirAll(dirPath, 0770)
	if err != nil {
		r.Fatalf("Failed to make directories %q: %v", dirPath, err)
	}
	err = ioutil.WriteFile(filePath, []byte(content), 0600)
	if err != nil {
		r.Fatalf("Failed to write file %q: %v", filePath, err)
	}
	err = os.Chtimes(filePath, t, t)
	if err != nil {
		r.Fatalf("Failed to chtimes file %q: %v", filePath, err)
	}
	return item
}

// ForceMkdir creates the remote
func (r *Run) ForceMkdir(f fs.Fs) {
	err := f.Mkdir("")
	if err != nil {
		r.Fatalf("Failed to mkdir %q: %v", f, err)
	}
	r.mkdir[f.String()] = true
}

// Mkdir creates the remote if it hasn't been created already
func (r *Run) Mkdir(f fs.Fs) {
	if !r.mkdir[f.String()] {
		r.ForceMkdir(f)
	}
}

// WriteObjectTo writes an object to the fs, remote passed in
func (r *Run) WriteObjectTo(f fs.Fs, remote, content string, modTime time.Time, useUnchecked bool) Item {
	put := f.Put
	if useUnchecked {
		put = f.Features().PutUnchecked
		if put == nil {
			r.Fatalf("Fs doesn't support PutUnchecked")
		}
	}
	r.Mkdir(f)
	const maxTries = 10
	for tries := 1; ; tries++ {
		in := bytes.NewBufferString(content)
		objinfo := fs.NewStaticObjectInfo(remote, modTime, int64(len(content)), true, nil, nil)
		_, err := put(in, objinfo)
		if err == nil {
			break
		}
		// Retry if err returned a retry error
		if fs.IsRetryError(err) && tries < maxTries {
			r.Logf("Retry Put of %q to %v: %d/%d (%v)", remote, f, tries, maxTries, err)
			time.Sleep(2 * time.Second)
			continue
		}
		r.Fatalf("Failed to put %q to %q: %v", remote, f, err)
	}
	return NewItem(remote, content, modTime)
}

// WriteObject writes an object to the remote
func (r *Run) WriteObject(remote, content string, modTime time.Time) Item {
	return r.WriteObjectTo(r.Fremote, remote, content, modTime, false)
}

// WriteUncheckedObject writes an object to the remote not checking for duplicates
func (r *Run) WriteUncheckedObject(remote, content string, modTime time.Time) Item {
	return r.WriteObjectTo(r.Fremote, remote, content, modTime, true)
}

// WriteBoth calls WriteObject and WriteFile with the same arguments
func (r *Run) WriteBoth(remote, content string, modTime time.Time) Item {
	r.WriteFile(remote, content, modTime)
	return r.WriteObject(remote, content, modTime)
}

// CheckWithDuplicates does a test but allows duplicates
func (r *Run) CheckWithDuplicates(t *testing.T, items ...Item) {
	objects, size, err := fs.Count(r.Fremote)
	require.NoError(t, err)
	assert.Equal(t, int64(len(items)), objects)
	wantSize := int64(0)
	for _, item := range items {
		wantSize += item.Size
	}
	assert.Equal(t, wantSize, size)
}

// Clean the temporary directory
func (r *Run) cleanTempDir() {
	err := os.RemoveAll(r.LocalName)
	if err != nil {
		r.Logf("Failed to clean temporary directory %q: %v", r.LocalName, err)
	}
}

// Finalise cleans the remote and local
func (r *Run) Finalise() {
	// r.Logf("Cleaning remote %q", r.Fremote)
	r.cleanRemote()
	// r.Logf("Cleaning local %q", r.LocalName)
	r.cleanTempDir()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
)

// 1, 如果接收者值未在方法体内使用，可以省略方法接收者标识符
// 2, 可变长参数的使用
// 3, 多返回值的使用
// 4, 返回值的声明和使用
func hello(prefix string, names ...string) (result string, err error) {
	result = prefix
	length := len(names)
	for index, name := range names {
		result += name
		if index < length-1 {
			result += ", "
		} else {
			result += " !"
		}
	}
	return
	// return result + " !", nil
}

func main() {
	result, _ := hello("Say hello to ", "Daniel", "Elim")
	fmt.Println(result)
}
<commit_msg>Modify comments<commit_after>package main

import (
	"fmt"
)

// 1, 可变长参数的使用
// 2, 多返回值的使用
// 3, 返回值的声明和使用
func hello(prefix string, names ...string) (result string, err error) {
	result = prefix
	length := len(names)
	for index, name := range names {
		result += name
		if index < length-1 {
			result += ", "
		} else {
			result += " !"
		}
	}
	return
	// return result + " !", nil
}

func main() {
	result, _ := hello("Say hello to ", "Daniel", "Elim")
	fmt.Println(result)
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"time"

	"github.com/gonfire/oauth2"
)

func tokenEndpoint(w http.ResponseWriter, r *http.Request) {
	// parse token request
	req, err := oauth2.ParseAccessTokenRequest(r)
	if err != nil {
		oauth2.WriteError(w, err)
		return
	}

	// validate token request
	err = req.Validate()
	if err != nil {
		oauth2.WriteError(w, err)
		return
	}

	// check if client is confidential
	if !req.Confidential() {
		oauth2.WriteError(w, oauth2.InvalidRequest(req.State, "Only cofidential clients are allowed"))
		return
	}

	// authenticate client
	client, found := clients[req.ClientID]
	if !found || !sameHash(client.secret, req.ClientSecret) {
		oauth2.WriteError(w, oauth2.InvalidClient(req.State, oauth2.NoDescription))
		return
	}

	// triage grant type
	if req.GrantType.Password() {
		handleResourceOwnerPasswordCredentialsGrant(w, req)
	} else if req.GrantType.ClientCredentials() {
		handleClientCredentialsGrant(w, req)
	} else if req.GrantType.AuthorizationCode() {
		handleAuthorizationCodeGrant(w, req)
	} else if req.GrantType.RefreshToken() {
		handleRefreshTokenGrant(w, req)
	}
}

func handleResourceOwnerPasswordCredentialsGrant(w http.ResponseWriter, req *oauth2.AccessTokenRequest) {
	// authenticate resource owner
	owner, found := users[req.Username]
	if !found || !sameHash(owner.secret, req.Password) {
		oauth2.WriteError(w, oauth2.AccessDenied(req.State, oauth2.NoDescription))
		return
	}

	// check scope
	if !allowedScope.Includes(req.Scope) {
		oauth2.WriteError(w, oauth2.InvalidScope(req.State, oauth2.NoDescription))
		return
	}

	// issue access and refresh token
	at, rt, res := createTokensAndResponse(req)

	// save tokens
	saveTokens(at, rt, req.Scope, req.ClientID, req.Username)

	// write response
	oauth2.WriteTokenResponse(w, res)
}

func handleClientCredentialsGrant(w http.ResponseWriter, req *oauth2.AccessTokenRequest) {
	// check scope
	if !allowedScope.Includes(req.Scope) {
		oauth2.WriteError(w, oauth2.InvalidScope(req.State, oauth2.NoDescription))
		return
	}

	// issue access and refresh token
	at, rt, res := createTokensAndResponse(req)

	// save tokens
	saveTokens(at, rt, req.Scope, req.ClientID, "")

	// write response
	oauth2.WriteTokenResponse(w, res)
}

func handleAuthorizationCodeGrant(w http.ResponseWriter, req *oauth2.AccessTokenRequest) {
	// parse authorization code
	authorizationCode, err := oauth2.ParseToken(secret, req.Code)
	if err != nil {
		oauth2.AddStateToError(err, req.State)
		oauth2.WriteError(w, err)
		return
	}

	// get stored authorization code by signature
	storedAuthorizationCode, found := authorizationCodes[authorizationCode.SignatureString()]
	if !found {
		oauth2.WriteError(w, oauth2.InvalidRequest(req.State, oauth2.NoDescription)) // TODO: Correct error?
		return
	}

	// validate ownership
	if storedAuthorizationCode.clientID != req.ClientID {
		oauth2.WriteError(w, oauth2.InvalidRequest(req.State, oauth2.NoDescription)) // TODO: Correct error?
		return
	}

	// validate scope and expiration
	if !storedAuthorizationCode.scope.Includes(req.Scope) || storedAuthorizationCode.expiresAt.Before(time.Now()) {
		oauth2.WriteError(w, oauth2.InvalidRequest(req.State, oauth2.NoDescription)) // TODO: Correct error?
		return
	}

	// validate redirect uri
	if req.RedirectURI != "" {
		if storedAuthorizationCode.redirectURI != req.RedirectURI {
			oauth2.WriteError(w, oauth2.InvalidRequest(req.State, oauth2.NoDescription)) // TODO: Correct error?
		}
	}

	// issue new access and refresh token
	at, rt, res := createTokensAndResponse(req)

	// save tokens
	saveTokens(at, rt, req.Scope, req.ClientID, "")

	// delete used authorization code
	delete(authorizationCodes, authorizationCode.SignatureString())

	// write response
	oauth2.WriteTokenResponse(w, res)
}

func handleRefreshTokenGrant(w http.ResponseWriter, req *oauth2.AccessTokenRequest) {
	// parse refresh token
	refreshToken, err := oauth2.ParseToken(secret, req.RefreshToken)
	if err != nil {
		oauth2.AddStateToError(err, req.State)
		oauth2.WriteError(w, err)
		return
	}

	// get stored refresh token by signature
	storedRefreshToken, found := refreshTokens[refreshToken.SignatureString()]
	if !found {
		oauth2.WriteError(w, oauth2.InvalidGrant(req.State, oauth2.NoDescription))
		return
	}

	// validate ownership
	if storedRefreshToken.clientID != req.ClientID {
		oauth2.WriteError(w, oauth2.InvalidGrant(req.State, oauth2.NoDescription))
		return
	}

	// validate scope and expiration
	if !storedRefreshToken.scope.Includes(req.Scope) || storedRefreshToken.expiresAt.Before(time.Now()) {
		oauth2.WriteError(w, oauth2.InvalidGrant(req.State, oauth2.NoDescription))
		return
	}

	// issue new access and refresh token
	at, rt, res := createTokensAndResponse(req)

	// save tokens
	saveTokens(at, rt, req.Scope, req.ClientID, storedRefreshToken.username)

	// delete used refresh token
	delete(refreshTokens, refreshToken.SignatureString())

	// write response
	oauth2.WriteTokenResponse(w, res)
}

func createTokensAndResponse(req *oauth2.AccessTokenRequest) (*oauth2.Token, *oauth2.Token, *oauth2.TokenResponse) {
	// generate new access token
	accessToken, err := oauth2.GenerateToken(secret, 32)
	if err != nil {
		panic(err)
	}

	// generate new refresh token
	refreshToken, err := oauth2.GenerateToken(secret, 32)
	if err != nil {
		panic(err)
	}

	// prepare response
	res := oauth2.NewBearerTokenResponse(accessToken.String(), int(tokenLifespan/time.Second))

	// set granted scope
	res.Scope = req.Scope

	// set refresh token
	res.RefreshToken = refreshToken.String()

	return accessToken, refreshToken, res
}

func saveTokens(accessToken, refreshToken *oauth2.Token, scope oauth2.Scope, clientID, username string) {
	// save access token
	accessTokens[accessToken.SignatureString()] = token{
		clientID:  clientID,
		username:  username,
		signature: accessToken.SignatureString(),
		expiresAt: time.Now().Add(tokenLifespan),
		scope:     scope,
	}

	// save refresh token
	refreshTokens[refreshToken.SignatureString()] = token{
		clientID:  clientID,
		username:  username,
		signature: refreshToken.SignatureString(),
		expiresAt: time.Now().Add(tokenLifespan),
		scope:     scope,
	}
}
<commit_msg>fix typo<commit_after>package main

import (
	"net/http"
	"time"

	"github.com/gonfire/oauth2"
)

func tokenEndpoint(w http.ResponseWriter, r *http.Request) {
	// parse token request
	req, err := oauth2.ParseAccessTokenRequest(r)
	if err != nil {
		oauth2.WriteError(w, err)
		return
	}

	// validate token request
	err = req.Validate()
	if err != nil {
		oauth2.WriteError(w, err)
		return
	}

	// check if client is confidential
	if !req.Confidential() {
		oauth2.WriteError(w, oauth2.InvalidRequest(req.State, "Only confidential clients are allowed"))
		return
	}

	// authenticate client
	client, found := clients[req.ClientID]
	if !found || !sameHash(client.secret, req.ClientSecret) {
		oauth2.WriteError(w, oauth2.InvalidClient(req.State, oauth2.NoDescription))
		return
	}

	// triage grant type
	if req.GrantType.Password() {
		handleResourceOwnerPasswordCredentialsGrant(w, req)
	} else if req.GrantType.ClientCredentials() {
		handleClientCredentialsGrant(w, req)
	} else if req.GrantType.AuthorizationCode() {
		handleAuthorizationCodeGrant(w, req)
	} else if req.GrantType.RefreshToken() {
		handleRefreshTokenGrant(w, req)
	}
}

func handleResourceOwnerPasswordCredentialsGrant(w http.ResponseWriter, req *oauth2.AccessTokenRequest) {
	// authenticate resource owner
	owner, found := users[req.Username]
	if !found || !sameHash(owner.secret, req.Password) {
		oauth2.WriteError(w, oauth2.AccessDenied(req.State, oauth2.NoDescription))
		return
	}

	// check scope
	if !allowedScope.Includes(req.Scope) {
		oauth2.WriteError(w, oauth2.InvalidScope(req.State, oauth2.NoDescription))
		return
	}

	// issue access and refresh token
	at, rt, res := createTokensAndResponse(req)

	// save tokens
	saveTokens(at, rt, req.Scope, req.ClientID, req.Username)

	// write response
	oauth2.WriteTokenResponse(w, res)
}

func handleClientCredentialsGrant(w http.ResponseWriter, req *oauth2.AccessTokenRequest) {
	// check scope
	if !allowedScope.Includes(req.Scope) {
		oauth2.WriteError(w, oauth2.InvalidScope(req.State, oauth2.NoDescription))
		return
	}

	// issue access and refresh token
	at, rt, res := createTokensAndResponse(req)

	// save tokens
	saveTokens(at, rt, req.Scope, req.ClientID, "")

	// write response
	oauth2.WriteTokenResponse(w, res)
}

func handleAuthorizationCodeGrant(w http.ResponseWriter, req *oauth2.AccessTokenRequest) {
	// parse authorization code
	authorizationCode, err := oauth2.ParseToken(secret, req.Code)
	if err != nil {
		oauth2.AddStateToError(err, req.State)
		oauth2.WriteError(w, err)
		return
	}

	// get stored authorization code by signature
	storedAuthorizationCode, found := authorizationCodes[authorizationCode.SignatureString()]
	if !found {
		oauth2.WriteError(w, oauth2.InvalidRequest(req.State, oauth2.NoDescription)) // TODO: Correct error?
		return
	}

	// validate ownership
	if storedAuthorizationCode.clientID != req.ClientID {
		oauth2.WriteError(w, oauth2.InvalidRequest(req.State, oauth2.NoDescription)) // TODO: Correct error?
		return
	}

	// validate scope and expiration
	if !storedAuthorizationCode.scope.Includes(req.Scope) || storedAuthorizationCode.expiresAt.Before(time.Now()) {
		oauth2.WriteError(w, oauth2.InvalidRequest(req.State, oauth2.NoDescription)) // TODO: Correct error?
		return
	}

	// validate redirect uri
	if req.RedirectURI != "" {
		if storedAuthorizationCode.redirectURI != req.RedirectURI {
			oauth2.WriteError(w, oauth2.InvalidRequest(req.State, oauth2.NoDescription)) // TODO: Correct error?
		}
	}

	// issue new access and refresh token
	at, rt, res := createTokensAndResponse(req)

	// save tokens
	saveTokens(at, rt, req.Scope, req.ClientID, "")

	// delete used authorization code
	delete(authorizationCodes, authorizationCode.SignatureString())

	// write response
	oauth2.WriteTokenResponse(w, res)
}

func handleRefreshTokenGrant(w http.ResponseWriter, req *oauth2.AccessTokenRequest) {
	// parse refresh token
	refreshToken, err := oauth2.ParseToken(secret, req.RefreshToken)
	if err != nil {
		oauth2.AddStateToError(err, req.State)
		oauth2.WriteError(w, err)
		return
	}

	// get stored refresh token by signature
	storedRefreshToken, found := refreshTokens[refreshToken.SignatureString()]
	if !found {
		oauth2.WriteError(w, oauth2.InvalidGrant(req.State, oauth2.NoDescription))
		return
	}

	// validate ownership
	if storedRefreshToken.clientID != req.ClientID {
		oauth2.WriteError(w, oauth2.InvalidGrant(req.State, oauth2.NoDescription))
		return
	}

	// validate scope and expiration
	if !storedRefreshToken.scope.Includes(req.Scope) || storedRefreshToken.expiresAt.Before(time.Now()) {
		oauth2.WriteError(w, oauth2.InvalidGrant(req.State, oauth2.NoDescription))
		return
	}

	// issue new access and refresh token
	at, rt, res := createTokensAndResponse(req)

	// save tokens
	saveTokens(at, rt, req.Scope, req.ClientID, storedRefreshToken.username)

	// delete used refresh token
	delete(refreshTokens, refreshToken.SignatureString())

	// write response
	oauth2.WriteTokenResponse(w, res)
}

func createTokensAndResponse(req *oauth2.AccessTokenRequest) (*oauth2.Token, *oauth2.Token, *oauth2.TokenResponse) {
	// generate new access token
	accessToken, err := oauth2.GenerateToken(secret, 32)
	if err != nil {
		panic(err)
	}

	// generate new refresh token
	refreshToken, err := oauth2.GenerateToken(secret, 32)
	if err != nil {
		panic(err)
	}

	// prepare response
	res := oauth2.NewBearerTokenResponse(accessToken.String(), int(tokenLifespan/time.Second))

	// set granted scope
	res.Scope = req.Scope

	// set refresh token
	res.RefreshToken = refreshToken.String()

	return accessToken, refreshToken, res
}

func saveTokens(accessToken, refreshToken *oauth2.Token, scope oauth2.Scope, clientID, username string) {
	// save access token
	accessTokens[accessToken.SignatureString()] = token{
		clientID:  clientID,
		username:  username,
		signature: accessToken.SignatureString(),
		expiresAt: time.Now().Add(tokenLifespan),
		scope:     scope,
	}

	// save refresh token
	refreshTokens[refreshToken.SignatureString()] = token{
		clientID:  clientID,
		username:  username,
		signature: refreshToken.SignatureString(),
		expiresAt: time.Now().Add(tokenLifespan),
		scope:     scope,
	}
}
<|endoftext|>


<commit_before>package permissions

import (
	"testing"

	"github.com/xyproto/db"
)

func TestPerm(t *testing.T) {
	//db.Verbose = true

	//userstate := NewUserStateSimple() // for localhost
	userstate := NewUserState("travis:@127.0.0.1/", true) // for travis-ci

	userstate.AddUser("bob", "hunter1", "bob@zombo.com")

	if !userstate.HasUser("bob") {
		t.Error("Error, user bob should exist")
	}

	if userstate.IsConfirmed("bob") {
		t.Error("Error, user bob should not be confirmed right now.")
	}

	userstate.MarkConfirmed("bob")

	if !userstate.IsConfirmed("bob") {
		t.Error("Error, user bob should be marked as confirmed right now.")
	}

	if userstate.IsAdmin("bob") {
		t.Error("Error, user bob should not have admin rights")
	}

	userstate.SetAdminStatus("bob")

	if !userstate.IsAdmin("bob") {
		t.Error("Error, user bob should have admin rights")
	}

	userstate.RemoveUser("bob")

	if userstate.HasUser("bob") {
		t.Error("Error, user bob should not exist")
	}
}

func TestPasswordBasic(t *testing.T) {
	//userstate := NewUserStateSimple() // for localhost
	userstate := NewUserState("travis:@127.0.0.1/", true) // for travis-ci

	// Assert that the default password algorithm is "bcrypt+"
	if userstate.PasswordAlgo() != "bcrypt+" {
		t.Error("Error, bcrypt+ should be the default password algorithm")
	}

	// Set password algorithm
	userstate.SetPasswordAlgo("sha256")

	// Assert that the algorithm is now sha256
	if userstate.PasswordAlgo() != "sha256" {
		t.Error("Error, setting password algorithm failed")
	}

}

// Check if the functionality for backwards compatible hashing works
func TestPasswordBackward(t *testing.T) {
	//userstate := NewUserStateSimple() // for localhost
	userstate := NewUserState("travis:@127.0.0.1/", true) // for travis-ci

	userstate.SetPasswordAlgo("sha256")
	userstate.AddUser("bob", "hunter1", "bob@zombo.com")
	if !userstate.HasUser("bob") {
		t.Error("Error, user bob should exist")
	}
	userstate.SetPasswordAlgo("sha256")
	if !userstate.CorrectPassword("bob", "hunter1") {
		t.Error("Error, the sha256 password really is correct")
	}

	userstate.SetPasswordAlgo("bcrypt")
	if userstate.CorrectPassword("bob", "hunter1") {
		t.Error("Error, the password as stored as sha256, not bcrypt")
	}

	userstate.SetPasswordAlgo("bcrypt+")
	if !userstate.CorrectPassword("bob", "hunter1") {
		t.Error("Error, the sha256 password is not correct when checking with bcrypt+")
	}

	userstate.RemoveUser("bob")
}

// Check if the functionality for backwards compatible hashing works
func TestPasswordNotBackward(t *testing.T) {
	//userstate := NewUserStateSimple() // for localhost
	userstate := NewUserState("travis:@127.0.0.1/", true) // for travis-ci

	userstate.SetPasswordAlgo("bcrypt")
	userstate.AddUser("bob", "hunter1", "bob@zombo.com")
	if !userstate.HasUser("bob") {
		t.Error("Error, user bob should exist")
	}
	userstate.SetPasswordAlgo("sha256")
	if userstate.CorrectPassword("bob", "hunter1") {
		t.Error("Error, the password is stored as bcrypt, should not be okay with sha256")
	}

	userstate.SetPasswordAlgo("bcrypt")
	if !userstate.CorrectPassword("bob", "hunter1") {
		t.Error("Error, the password should be correct when checking with bcrypt")
	}

	userstate.RemoveUser("bob")
}

func TestPasswordAlgoMatching(t *testing.T) {
	//userstate := NewUserStateSimple() // for localhost
	userstate := NewUserState("travis:@127.0.0.1/", true) // for travis-ci

	// generate two different password using the same credentials but different algos
	userstate.SetPasswordAlgo("sha256")
	sha256_hash := userstate.HashPassword("testuser@example.com", "textpassword")
	userstate.SetPasswordAlgo("bcrypt")
	bcrypt_hash := userstate.HashPassword("testuser@example.com", "textpassword")

	// they shouldn't match
	if sha256_hash == bcrypt_hash {
		t.Error("Error, different algorithms should not have a password match")
	}
}

func TestIUserState(t *testing.T) {
	//userstate := NewUserStateSimple() // for localhost
	userstate := NewUserState("travis:@127.0.0.1/", true) // for travis-ci

	// Check that the userstate qualifies for the IUserState interface
	var _ db.IUserState = userstate
}

func TestHostPassword(t *testing.T) {
	//userstate := NewUserStateSimple() // for localhost
	userstate := NewUserState("travis:@127.0.0.1/", true) // for travis-ci

	userstate.AddUser("bob", "hunter1", "bob@zombo.com")
	if !userstate.HasUser("bob") {
		t.Error("Error, user bob should exist")
	}

	userstate.RemoveUser("bob")
	if userstate.HasUser("bob") {
		t.Error("Error, user bob should not exist")
	}
}
<commit_msg>Refactoring the tests<commit_after>package permissions

import (
	"testing"

	"github.com/xyproto/db"
)

const connectionString = "travis:@127.0.0.1/" // for Travis-CI
//const connectionString = "username:password@host:port/database"

func TestPerm(t *testing.T) {
	//db.Verbose = true

	//userstate := NewUserStateSimple() // for localhost
	userstate := NewUserState(connectionString, true)

	userstate.AddUser("bob", "hunter1", "bob@zombo.com")

	if !userstate.HasUser("bob") {
		t.Error("Error, user bob should exist")
	}

	if userstate.IsConfirmed("bob") {
		t.Error("Error, user bob should not be confirmed right now.")
	}

	userstate.MarkConfirmed("bob")

	if !userstate.IsConfirmed("bob") {
		t.Error("Error, user bob should be marked as confirmed right now.")
	}

	if userstate.IsAdmin("bob") {
		t.Error("Error, user bob should not have admin rights")
	}

	userstate.SetAdminStatus("bob")

	if !userstate.IsAdmin("bob") {
		t.Error("Error, user bob should have admin rights")
	}

	userstate.RemoveUser("bob")

	if userstate.HasUser("bob") {
		t.Error("Error, user bob should not exist")
	}
}

func TestPasswordBasic(t *testing.T) {
	//userstate := NewUserStateSimple() // for localhost
	userstate := NewUserState(connectionString, true)

	// Assert that the default password algorithm is "bcrypt+"
	if userstate.PasswordAlgo() != "bcrypt+" {
		t.Error("Error, bcrypt+ should be the default password algorithm")
	}

	// Set password algorithm
	userstate.SetPasswordAlgo("sha256")

	// Assert that the algorithm is now sha256
	if userstate.PasswordAlgo() != "sha256" {
		t.Error("Error, setting password algorithm failed")
	}

}

// Check if the functionality for backwards compatible hashing works
func TestPasswordBackward(t *testing.T) {
	//userstate := NewUserStateSimple() // for localhost
	userstate := NewUserState(connectionString, true)

	userstate.SetPasswordAlgo("sha256")
	userstate.AddUser("bob", "hunter1", "bob@zombo.com")
	if !userstate.HasUser("bob") {
		t.Error("Error, user bob should exist")
	}
	userstate.SetPasswordAlgo("sha256")
	if !userstate.CorrectPassword("bob", "hunter1") {
		t.Error("Error, the sha256 password really is correct")
	}

	userstate.SetPasswordAlgo("bcrypt")
	if userstate.CorrectPassword("bob", "hunter1") {
		t.Error("Error, the password as stored as sha256, not bcrypt")
	}

	userstate.SetPasswordAlgo("bcrypt+")
	if !userstate.CorrectPassword("bob", "hunter1") {
		t.Error("Error, the sha256 password is not correct when checking with bcrypt+")
	}

	userstate.RemoveUser("bob")
}

// Check if the functionality for backwards compatible hashing works
func TestPasswordNotBackward(t *testing.T) {
	//userstate := NewUserStateSimple() // for localhost
	userstate := NewUserState(connectionString, true)

	userstate.SetPasswordAlgo("bcrypt")
	userstate.AddUser("bob", "hunter1", "bob@zombo.com")
	if !userstate.HasUser("bob") {
		t.Error("Error, user bob should exist")
	}
	userstate.SetPasswordAlgo("sha256")
	if userstate.CorrectPassword("bob", "hunter1") {
		t.Error("Error, the password is stored as bcrypt, should not be okay with sha256")
	}

	userstate.SetPasswordAlgo("bcrypt")
	if !userstate.CorrectPassword("bob", "hunter1") {
		t.Error("Error, the password should be correct when checking with bcrypt")
	}

	userstate.RemoveUser("bob")
}

func TestPasswordAlgoMatching(t *testing.T) {
	//userstate := NewUserStateSimple() // for localhost
	userstate := NewUserState(connectionString, true)

	// generate two different password using the same credentials but different algos
	userstate.SetPasswordAlgo("sha256")
	sha256_hash := userstate.HashPassword("testuser@example.com", "textpassword")
	userstate.SetPasswordAlgo("bcrypt")
	bcrypt_hash := userstate.HashPassword("testuser@example.com", "textpassword")

	// they shouldn't match
	if sha256_hash == bcrypt_hash {
		t.Error("Error, different algorithms should not have a password match")
	}
}

func TestIUserState(t *testing.T) {
	//userstate := NewUserStateSimple() // for localhost
	userstate := NewUserState(connectionString, true)

	// Check that the userstate qualifies for the IUserState interface
	var _ db.IUserState = userstate
}

func TestHostPassword(t *testing.T) {
	//userstate := NewUserStateSimple() // for localhost
	userstate := NewUserState(connectionString, true)

	userstate.AddUser("bob", "hunter1", "bob@zombo.com")
	if !userstate.HasUser("bob") {
		t.Error("Error, user bob should exist")
	}

	userstate.RemoveUser("bob")
	if userstate.HasUser("bob") {
		t.Error("Error, user bob should not exist")
	}
}
<|endoftext|>