<commit_before>package client_test

import (
    client "../src"
    "testing"
)

func TestValidateSendMessage(t *testing.T) {

    // test a basic valid message
    command := []byte("/msg 1234,2345,3456 Hi!")
    validatedCommand := client.ValidateSendMessage(command);
    if (string(validatedCommand) != "/msg 1234,2345,3456 Hi!") {
        t.Error("/msg command failed")
    }

    // test a message with spaces
    command = []byte("/msg 1234 This sentence contains many spaces.")
    validatedCommand = client.ValidateSendMessage(command);
    if (string(validatedCommand) != "/msg 1234 This sentence contains many spaces.") {
        t.Error("/msg command failed")
    }

    // test too many receivers
    command = []byte("/msg 1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127,128,129,130,131,132,133,134,135,136,137,138,139,140,141,142,143,144,145,146,147,148,149,150,151,152,153,154,155,156,157,158,159,160,161,162,163,164,165,166,167,168,169,170,171,172,173,174,175,176,177,178,179,180,181,182,183,184,185,186,187,188,189,190,191,192,193,194,195,196,197,198,199,200,201,202,203,204,205,206,207,208,209,210,211,212,213,214,215,216,217,218,219,220,221,222,223,224,225,226,227,228,229,230,231,232,233,234,235,236,237,238,239,240,241,242,243,244,245,246,247,248,249,250,251,252,253,254,255,256 This message has 256 receivers.")
    validatedCommand = client.ValidateSendMessage(command);
    if (string(validatedCommand) != "/msg 1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127,128,129,130,131,132,133,134,135,136,137,138,139,140,141,142,143,144,145,146,147,148,149,150,151,152,153,154,155,156,157,158,159,160,161,162,163,164,165,166,167,168,169,170,171,172,173,174,175,176,177,178,179,180,181,182,183,184,185,186,187,188,189,190,191,192,193,194,195,196,197,198,199,200,201,202,203,204,205,206,207,208,209,210,211,212,213,214,215,216,217,218,219,220,221,222,223,224,225,226,227,228,229,230,231,232,233,234,235,236,237,238,239,240,241,242,243,244,245,246,247,248,249,250,251,252,253,254,255 This message has 256 receivers.") {
        t.Error("/msg command failed")
    }

    // test too large message body
    data := []byte{}
    for {
        data = append(data, []byte{0}...)
        if (len(data) >= client.MAX_MESSAGE_BODY_SIZE + 1) {
            break;
        }
    }
    command = append([]byte("/msg 1,2,3 "), data...)
    validatedCommand = client.ValidateSendMessage(command);
    if (string(validatedCommand) == "/msg 1,2,3 " + string(data)) {
        t.Error("/msg message body too large")
    }

    // max message body
    data = []byte{}
    for {
        data = append(data, []byte{0}...)
        if (len(data) >= client.MAX_MESSAGE_BODY_SIZE - 1) {
            break;
        }
    }
    command = append([]byte("/msg 1,2,3 "), data...)
    validatedCommand = client.ValidateSendMessage(command);
    if (string(validatedCommand) != "/msg 1,2,3 " + string(data)) {
        t.Error("/msg message body size check failed")
    }

}<commit_msg>feature: test json and binary message body<commit_after>package client_test

import (
    client "../src"
    "testing"
)

func TestValidateSendMessage(t *testing.T) {

    // test a basic valid message
    command := []byte("/msg 1234,2345,3456 Hi!")
    validatedCommand := client.ValidateSendMessage(command);
    if (string(validatedCommand) != "/msg 1234,2345,3456 Hi!") {
        t.Error("/msg command failed")
    }

    // test a message with spaces
    command = []byte("/msg 1234 This sentence contains many spaces.")
    validatedCommand = client.ValidateSendMessage(command);
    if (string(validatedCommand) != "/msg 1234 This sentence contains many spaces.") {
        t.Error("/msg command failed")
    }

    // test too many receivers
    command = []byte("/msg 1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127,128,129,130,131,132,133,134,135,136,137,138,139,140,141,142,143,144,145,146,147,148,149,150,151,152,153,154,155,156,157,158,159,160,161,162,163,164,165,166,167,168,169,170,171,172,173,174,175,176,177,178,179,180,181,182,183,184,185,186,187,188,189,190,191,192,193,194,195,196,197,198,199,200,201,202,203,204,205,206,207,208,209,210,211,212,213,214,215,216,217,218,219,220,221,222,223,224,225,226,227,228,229,230,231,232,233,234,235,236,237,238,239,240,241,242,243,244,245,246,247,248,249,250,251,252,253,254,255,256 This message has 256 receivers.")
    validatedCommand = client.ValidateSendMessage(command);
    if (string(validatedCommand) != "/msg 1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127,128,129,130,131,132,133,134,135,136,137,138,139,140,141,142,143,144,145,146,147,148,149,150,151,152,153,154,155,156,157,158,159,160,161,162,163,164,165,166,167,168,169,170,171,172,173,174,175,176,177,178,179,180,181,182,183,184,185,186,187,188,189,190,191,192,193,194,195,196,197,198,199,200,201,202,203,204,205,206,207,208,209,210,211,212,213,214,215,216,217,218,219,220,221,222,223,224,225,226,227,228,229,230,231,232,233,234,235,236,237,238,239,240,241,242,243,244,245,246,247,248,249,250,251,252,253,254,255 This message has 256 receivers.") {
        t.Error("/msg command failed")
    }

    // test too large message body
    data := []byte{}
    for {
        data = append(data, []byte{0}...)
        if (len(data) >= client.MAX_MESSAGE_BODY_SIZE + 1) {
            break;
        }
    }
    command = append([]byte("/msg 1,2,3 "), data...)
    validatedCommand = client.ValidateSendMessage(command);
    if (string(validatedCommand) == "/msg 1,2,3 " + string(data)) {
        t.Error("/msg message body too large")
    }

    // max message body
    data = []byte{}
    for {
        data = append(data, []byte{0}...)
        if (len(data) >= client.MAX_MESSAGE_BODY_SIZE - 1) {
            break;
        }
    }
    command = append([]byte("/msg 1,2,3 "), data...)
    validatedCommand = client.ValidateSendMessage(command);
    if (string(validatedCommand) != "/msg 1,2,3 " + string(data)) {
        t.Error("/msg message body size check failed")
    }

    // test json message body
    json := "{\"books\":[{\"id\":1,\"language\":\"Java\",\"edition\":\"third\",\"author\":\"Herbert Schildt\"},{\"id\":7,\"language\":\"C++\",\"edition\":\"second\",\"author\":\"E.Balagurusamy\"}]}"
    command = []byte("/msg 1234,2345,3456 " + json)
    validatedCommand = client.ValidateSendMessage(command);
    if (string(validatedCommand) != "/msg 1234,2345,3456 " + json) {
        t.Error("/msg command failed")
    }

    // test binary message body
    binary := "1F332D01"
    command = []byte("/msg 1234,2345,3456 " + binary)
    validatedCommand = client.ValidateSendMessage(command);
    if (string(validatedCommand) != "/msg 1234,2345,3456 " + binary) {
        t.Error("/msg command failed")
    }

}<|endoftext|>


<commit_before>//package limitrequests contains tools for limiting the number of simulatenously launched requests.
package limitrequests

import (
	"sync/atomic"
	"time"
)

func min(a, b time.Duration) time.Duration {
	if a < b {
		return a
	}
	return b
}

//TotalLimiter limits the total number of current requests. Usage pattern is follows:
//suppose we want to run at most 500 of a func networkCall.
//const maxSimulataneousRequests, timeOut = 500, 30*time.Second
//var startRequest, endRequest = TotalRequestLimiter(maxSimulataneousRequests, timeOut)
//for request := range requests {
//	if startRequest() {
//		go func(request) {
//			defer endRequest
//			networkCall(request)
//		}(request)
//	} else {
//		log.Fatal(fmt.Errorf("could not launch networkCall(%v): timeout after %d seconds", maxSimulataneousRequests, timeout.Seconds())
//	}
//
//}
func TotalLimiter(maxRequests int64, timeout time.Duration) (startRequest func() bool, finishRequest func()) {
	var currentRequests *int64
	waitForRequests := func() bool {
		for wait, i := time.Duration(0), uint(0); wait < timeout; i++ {
			if atomic.LoadInt64(currentRequests) < maxRequests {
				return true
			}
			wait += min(10*time.Millisecond<<i, 250*time.Millisecond)
		}
		return false
	}

	startRequest = func() bool {
		if timeout := !waitForRequests(); timeout {
			return false
		}

		atomic.AddInt64(currentRequests, 1)
		return true
	}
	finishRequest = func() {
		atomic.AddInt64(currentRequests, -1)
	}
	return startRequest, finishRequest
}

//PerSecondLimiter is for situations where the total number of current requests is less important than the number of requests launched per second.
//Note that PerSecondLimiter cannot lock.
func PerSecondLimiter(maxRequestsPerSecond float64, pollingPeriod time.Duration) (startRequest func()) {
	if maxRequestsPerSecond <= 0 {
		panic("must enter a positive number for maxRequestsPerSecond")
	}
	if pollingPeriod <= 0 {
		panic("must enter a positive duration for pollingPeriod")
	}
	startTime := time.Now()
	var totalRequests *uint64
	return func() {
		for {
			elapsed := time.Now().Sub(startTime).Seconds()
			requestsPerSecond := float64(atomic.LoadUint64(totalRequests)) / elapsed
			if requestsPerSecond < maxRequestsPerSecond {
				atomic.AddUint64(totalRequests, 1)
				return
			}
			time.Sleep(pollingPeriod)
		}
	}
}
<commit_msg>limitrequests package!!! check it out :)<commit_after>//Package limitrequests contains tools for limiting the number of simulatenously launched requests.
package limitrequests

import (
	"sync/atomic"
	"time"
)

//ConcurrencyLimiter is an interface for limiting the number of simulataneous concurrent requests.
type ConcurrencyLimiter interface {
	add()
	Done()
	WaitForSlot() bool
	WaitAll() //equivalent to sync.WaitGroup.Wait()
	openSlot() bool
	Start() bool
	Total()
}

//Limiter is a concrete struct which implements the values & methods of ConcurrencyLimiter that should be shared across all implementations.
type Limiter struct {
	ConcurrencyLimiter
	CurrentRequests, TotalRequests *int64
	Timeout, PollingPeriod         time.Duration
}

//Start waits Timeout seconds for a for a slot to open. If a slot opens, it adds 1 to the current requests and returns true.
//otherwise, it returns false.
func (l *Limiter) Start() bool {
	if timeout := !l.WaitForSlot(); timeout {
		return false
	}

	l.add()
	return true
}

//Add a request.
func (l *Limiter) add() {
	atomic.AddInt64(l.CurrentRequests, 1)
	atomic.AddInt64(l.TotalRequests, 1)
}

//WaitForSlot returns true if a slot opens for requests before Timeout.
func (l *Limiter) WaitForSlot() bool {
	for wait, i := time.Duration(0), uint(0); wait < l.Timeout; i++ {
		if l.openSlot() {
			return true
		}
		wait += l.PollingPeriod
		time.Sleep(l.PollingPeriod)
	}
	return false
}

//WaitAll waits for all current requests to complete. This can wait forever if some requests never finish.
//Use WaitAllTimeOut if that's not what you want.
func (l *Limiter) WaitAll() {
	for atomic.LoadInt64(l.CurrentRequests) > 0 {
		time.Sleep(l.PollingPeriod)
	}
}

//WaitAllTimeout waits for timeout duration for all requests to complete. If you want to wait forever, use WaitAll instead.
func (l *Limiter) WaitAllTimeout(timeout time.Duration) bool {
	for wait := time.Duration(0); wait < l.Timeout; wait += l.PollingPeriod {
		if atomic.LoadInt64(l.CurrentRequests) == 0 {
			return true
		}
		time.Sleep(l.PollingPeriod)
	}
	return false
}

//TotalLimiter limits the total number of current requests. Build with NewTotalLimiter.
type TotalLimiter struct {
	*Limiter
	MaxRequests int64 `json:"max_requests,omitempty"`
}

func (psl *TotalLimiter) openSlot() bool {
	return atomic.LoadInt64(psl.CurrentRequests) < psl.MaxRequests
}

//NewTotalLimiter builds a TotalLimiter
func NewTotalLimiter(maxRequests int64, timeOut time.Duration) *TotalLimiter {
	limiter := Limiter{Timeout: timeOut}
	return &TotalLimiter{&limiter, maxRequests}
}

//PerSecondLimiter limits the number of requests to an amount per second. Build with NewPerSecondLimiter
type PerSecondLimiter struct {
	*Limiter
	MaxRequestsPerSecond        float64 `json:"max_requests_per_second,omitempty"`
	startTimeOfLastRequestCycle time.Time
}

func (psl *PerSecondLimiter) openSlot() bool {
	if atomic.LoadInt64(psl.CurrentRequests) == 0 {
		// no live requests, so we need to reset the time from which we're counting requests per-second.
		psl.startTimeOfLastRequestCycle = time.Now()
		return true
	}
	return psl.requestsPerSecond() > psl.MaxRequestsPerSecond
}

func (psl *PerSecondLimiter) requestsPerSecond() float64 {
	elapsed := time.Now().Sub(psl.startTimeOfLastRequestCycle).Seconds()
	return elapsed / float64(atomic.LoadInt64(psl.CurrentRequests))
}

//NewPerSecondLimiter builds a PerSecondLimiter.
func NewPerSecondLimiter(maxRequestsPerSecond float64, timeOut time.Duration) *PerSecondLimiter {
	limiter := Limiter{Timeout: timeOut}
	return &PerSecondLimiter{&limiter, maxRequestsPerSecond, time.Now()}
}
<|endoftext|>


<commit_before>package dsig

import (
	"bytes"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"errors"
	"fmt"
	"regexp"

	"github.com/beevik/etree"
	"github.com/russellhaering/goxmldsig/etreeutils"
	"github.com/russellhaering/goxmldsig/types"
)

var uriRegexp = regexp.MustCompile("^#[a-zA-Z_][\\w.-]*$")

var (
	// ErrMissingSignature indicates that no enveloped signature was found referencing
	// the top level element passed for signature verification.
	ErrMissingSignature = errors.New("Missing signature referencing the top-level element")
)

type ValidationContext struct {
	CertificateStore X509CertificateStore
	IdAttribute      string
	Clock            *Clock
}

func NewDefaultValidationContext(certificateStore X509CertificateStore) *ValidationContext {
	return &ValidationContext{
		CertificateStore: certificateStore,
		IdAttribute:      DefaultIdAttr,
	}
}

// TODO(russell_h): More flexible namespace support. This might barely work.
func inNamespace(el *etree.Element, ns string) bool {
	for _, attr := range el.Attr {
		if attr.Value == ns {
			if attr.Space == "" && attr.Key == "xmlns" {
				return el.Space == ""
			} else if attr.Space == "xmlns" {
				return el.Space == attr.Key
			}
		}
	}

	return false
}

func childPath(space, tag string) string {
	if space == "" {
		return "./" + tag
	} else {
		return "./" + space + ":" + tag
	}
}

// The RemoveElement method on etree.Element isn't recursive...
func recursivelyRemoveElement(tree, el *etree.Element) bool {
	if tree.RemoveChild(el) != nil {
		return true
	}

	for _, child := range tree.Child {
		if childElement, ok := child.(*etree.Element); ok {
			if recursivelyRemoveElement(childElement, el) {
				return true
			}
		}
	}

	return false
}

// transform applies the passed set of transforms to the specified root element.
//
// The functionality of transform is currently very limited and purpose-specific.
//
// NOTE(russell_h): Ideally this wouldn't mutate the root passed to it, and would
// instead return a copy. Unfortunately copying the tree makes it difficult to
// correctly locate the signature. I'm opting, for now, to simply mutate the root
// parameter.
func (ctx *ValidationContext) transform(
	el *etree.Element,
	sig *types.Signature,
	ref *types.Reference) (*etree.Element, Canonicalizer, error) {
	transforms := ref.Transforms.Transforms

	if len(transforms) != 2 {
		return nil, nil, errors.New("Expected Enveloped and C14N transforms")
	}

	var canonicalizer Canonicalizer

	for _, transform := range transforms {
		algo := transform.Algorithm

		switch AlgorithmID(algo) {
		case EnvelopedSignatureAltorithmId:
			if !recursivelyRemoveElement(el, sig.UnderlyingElement()) {
				return nil, nil, errors.New("Error applying canonicalization transform: Signature not found")
			}

		case CanonicalXML10ExclusiveAlgorithmId:
			var prefixList string
			if transform.InclusiveNamespaces != nil {
				prefixList = transform.InclusiveNamespaces.PrefixList
			}

			canonicalizer = MakeC14N10ExclusiveCanonicalizerWithPrefixList(prefixList)

		case CanonicalXML11AlgorithmId:
			canonicalizer = MakeC14N11Canonicalizer()

		default:
			return nil, nil, errors.New("Unknown Transform Algorithm: " + algo)
		}
	}

	if canonicalizer == nil {
		return nil, nil, errors.New("Expected canonicalization transform")
	}

	return el, canonicalizer, nil
}

func (ctx *ValidationContext) digest(el *etree.Element, digestAlgorithmId string, canonicalizer Canonicalizer) ([]byte, error) {
	data, err := canonicalizer.Canonicalize(el)
	if err != nil {
		return nil, err
	}

	digestAlgorithm, ok := digestAlgorithmsByIdentifier[digestAlgorithmId]
	if !ok {
		return nil, errors.New("Unknown digest algorithm: " + digestAlgorithmId)
	}

	hash := digestAlgorithm.New()
	_, err = hash.Write(data)
	if err != nil {
		return nil, err
	}

	return hash.Sum(nil), nil
}

func (ctx *ValidationContext) verifySignedInfo(sig *types.Signature, canonicalizer Canonicalizer, signatureMethodId string, cert *x509.Certificate, decodedSignature []byte) error {
	signatureElement := sig.UnderlyingElement()

	signedInfo := signatureElement.FindElement(childPath(signatureElement.Space, SignedInfoTag))
	if signedInfo == nil {
		return errors.New("Missing SignedInfo")
	}

	// Canonicalize the xml
	canonical, err := canonicalSerialize(signedInfo)
	if err != nil {
		return err
	}

	signatureAlgorithm, ok := signatureMethodsByIdentifier[signatureMethodId]
	if !ok {
		return errors.New("Unknown signature method: " + signatureMethodId)
	}

	hash := signatureAlgorithm.New()
	_, err = hash.Write(canonical)
	if err != nil {
		return err
	}

	hashed := hash.Sum(nil)

	pubKey, ok := cert.PublicKey.(*rsa.PublicKey)
	if !ok {
		return errors.New("Invalid public key")
	}

	// Verify that the private key matching the public key from the cert was what was used to sign the 'SignedInfo' and produce the 'SignatureValue'
	err = rsa.VerifyPKCS1v15(pubKey, signatureAlgorithm, hashed[:], decodedSignature)
	if err != nil {
		return err
	}

	return nil
}

func (ctx *ValidationContext) validateSignature(el *etree.Element, sig *types.Signature, cert *x509.Certificate) (*etree.Element, error) {
	idAttr := el.SelectAttr(DefaultIdAttr)
	if idAttr == nil || idAttr.Value == "" {
		return nil, errors.New("Missing ID attribute")
	}

	var ref *types.Reference

	// Find the first reference which references the top-level element
	for _, _ref := range sig.SignedInfo.References {
		if _ref.URI == "" || _ref.URI[1:] == idAttr.Value {
			ref = &_ref
		}
	}

	// Perform all transformations listed in the 'SignedInfo'
	// Basically, this means removing the 'SignedInfo'
	transformed, canonicalizer, err := ctx.transform(el, sig, ref)
	if err != nil {
		return nil, err
	}

	digestAlgorithm := ref.DigestAlgo.Algorithm

	// Digest the transformed XML and compare it to the 'DigestValue' from the 'SignedInfo'
	digest, err := ctx.digest(transformed, digestAlgorithm, canonicalizer)
	if err != nil {
		return nil, err
	}

	decodedDigestValue, err := base64.StdEncoding.DecodeString(ref.DigestValue)
	if err != nil {
		return nil, err
	}

	if !bytes.Equal(digest, decodedDigestValue) {
		return nil, errors.New("Signature could not be verified")
	}

	// Decode the 'SignatureValue' so we can compare against it
	decodedSignature, err := base64.StdEncoding.DecodeString(sig.SignatureValue.Data)
	if err != nil {
		return nil, errors.New("Could not decode signature")
	}

	// Actually verify the 'SignedInfo' was signed by a trusted source
	signatureMethod := sig.SignedInfo.SignatureMethod.Algorithm
	err = ctx.verifySignedInfo(sig, canonicalizer, signatureMethod, cert, decodedSignature)
	if err != nil {
		return nil, err
	}

	return transformed, nil
}

func contains(roots []*x509.Certificate, cert *x509.Certificate) bool {
	for _, root := range roots {
		if root.Equal(cert) {
			return true
		}
	}
	return false
}

// findSignature searches for a Signature element referencing the passed root element.
func (ctx *ValidationContext) findSignature(el *etree.Element) (*types.Signature, error) {
	idAttr := el.SelectAttr(DefaultIdAttr)
	if idAttr == nil || idAttr.Value == "" {
		return nil, errors.New("Missing ID attribute")
	}

	var sig *types.Signature

	// Traverse the tree looking for a Signature element
	err := etreeutils.NSFindIterate(el, Namespace, SignatureTag, func(ctx etreeutils.NSContext, el *etree.Element) error {

		found := false
		err := etreeutils.NSFindIterateCtx(ctx, el, Namespace, SignedInfoTag,
			func(ctx etreeutils.NSContext, signedInfo *etree.Element) error {
				// Ignore any SignedInfo that isn't an immediate descendent of Signature.
				if signedInfo.Parent() != el {
					return nil
				}

				detachedSignedInfo, err := etreeutils.NSDetatch(ctx, signedInfo)
				if err != nil {
					return err
				}

				c14NMethod := detachedSignedInfo.FindElement(childPath(detachedSignedInfo.Space, CanonicalizationMethodTag))
				if c14NMethod == nil {
					return errors.New("missing CanonicalizationMethod on Signature")
				}

				c14NAlgorithm := c14NMethod.SelectAttrValue(AlgorithmAttr, "")

				var canonicalSignedInfo *etree.Element

				switch AlgorithmID(c14NAlgorithm) {
				case CanonicalXML10ExclusiveAlgorithmId:
					err := etreeutils.TransformExcC14n(detachedSignedInfo, "")
					if err != nil {
						return err
					}

					// NOTE: TransformExcC14n transforms the element in-place,
					// while canonicalPrep isn't meant to. Once we standardize
					// this behavior we can drop this, as well as the adding and
					// removing of elements below.
					canonicalSignedInfo = detachedSignedInfo

				case CanonicalXML11AlgorithmId:
					canonicalSignedInfo = canonicalPrep(detachedSignedInfo, map[string]struct{}{})

				default:
					return fmt.Errorf("invalid CanonicalizationMethod on Signature: %s", c14NAlgorithm)
				}

				el.RemoveChild(signedInfo)
				el.AddChild(canonicalSignedInfo)

				found = true

				return etreeutils.ErrTraversalHalted
			})
		if err != nil {
			return err
		}

		if !found {
			return errors.New("Missing SignedInfo")
		}

		// Unmarshal the signature into a structured Signature type
		_sig := &types.Signature{}
		err = etreeutils.NSUnmarshalElement(ctx, el, _sig)
		if err != nil {
			return err
		}

		// Traverse references in the signature to determine whether it has at least
		// one reference to the top level element. If so, conclude the search.
		for _, ref := range _sig.SignedInfo.References {
			if ref.URI == "" || ref.URI[1:] == idAttr.Value {
				sig = _sig
				return etreeutils.ErrTraversalHalted
			}
		}

		return nil
	})

	if err != nil {
		return nil, err
	}

	if sig == nil {
		return nil, ErrMissingSignature
	}

	return sig, nil
}

func (ctx *ValidationContext) verifyCertificate(sig *types.Signature) (*x509.Certificate, error) {
	now := ctx.Clock.Now()

	roots, err := ctx.CertificateStore.Certificates()
	if err != nil {
		return nil, err
	}

	var cert *x509.Certificate

	if sig.KeyInfo != nil {
		// If the Signature includes KeyInfo, extract the certificate from there
		if sig.KeyInfo.X509Data.X509Certificate.Data == "" {
			return nil, errors.New("missing X509Certificate within KeyInfo")
		}

		certData, err := base64.StdEncoding.DecodeString(sig.KeyInfo.X509Data.X509Certificate.Data)
		if err != nil {
			return nil, errors.New("Failed to parse certificate")
		}

		cert, err = x509.ParseCertificate(certData)
		if err != nil {
			return nil, err
		}
	} else {
		// If the Signature doesn't have KeyInfo, Use the root certificate if there is only one
		if len(roots) == 1 {
			cert = roots[0]
		} else {
			return nil, errors.New("Missing x509 Element")
		}
	}

	// Verify that the certificate is one we trust
	if !contains(roots, cert) {
		return nil, errors.New("Could not verify certificate against trusted certs")
	}

	if now.Before(cert.NotBefore) || now.After(cert.NotAfter) {
		return nil, errors.New("Cert is not valid at this time")
	}

	return cert, nil
}

// Validate verifies that the passed element contains a valid enveloped signature
// matching a currently-valid certificate in the context's CertificateStore.
func (ctx *ValidationContext) Validate(el *etree.Element) (*etree.Element, error) {
	// Make a copy of the element to avoid mutating the one we were passed.
	el = el.Copy()

	sig, err := ctx.findSignature(el)
	if err != nil {
		return nil, err
	}

	cert, err := ctx.verifyCertificate(sig)
	if err != nil {
		return nil, err
	}

	return ctx.validateSignature(el, sig, cert)
}
<commit_msg>Use ValidationContext.IdAttribute property<commit_after>package dsig

import (
	"bytes"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"errors"
	"fmt"
	"regexp"

	"github.com/beevik/etree"
	"github.com/russellhaering/goxmldsig/etreeutils"
	"github.com/russellhaering/goxmldsig/types"
)

var uriRegexp = regexp.MustCompile("^#[a-zA-Z_][\\w.-]*$")

var (
	// ErrMissingSignature indicates that no enveloped signature was found referencing
	// the top level element passed for signature verification.
	ErrMissingSignature = errors.New("Missing signature referencing the top-level element")
)

type ValidationContext struct {
	CertificateStore X509CertificateStore
	IdAttribute      string
	Clock            *Clock
}

func NewDefaultValidationContext(certificateStore X509CertificateStore) *ValidationContext {
	return &ValidationContext{
		CertificateStore: certificateStore,
		IdAttribute:      DefaultIdAttr,
	}
}

// TODO(russell_h): More flexible namespace support. This might barely work.
func inNamespace(el *etree.Element, ns string) bool {
	for _, attr := range el.Attr {
		if attr.Value == ns {
			if attr.Space == "" && attr.Key == "xmlns" {
				return el.Space == ""
			} else if attr.Space == "xmlns" {
				return el.Space == attr.Key
			}
		}
	}

	return false
}

func childPath(space, tag string) string {
	if space == "" {
		return "./" + tag
	} else {
		return "./" + space + ":" + tag
	}
}

// The RemoveElement method on etree.Element isn't recursive...
func recursivelyRemoveElement(tree, el *etree.Element) bool {
	if tree.RemoveChild(el) != nil {
		return true
	}

	for _, child := range tree.Child {
		if childElement, ok := child.(*etree.Element); ok {
			if recursivelyRemoveElement(childElement, el) {
				return true
			}
		}
	}

	return false
}

// transform applies the passed set of transforms to the specified root element.
//
// The functionality of transform is currently very limited and purpose-specific.
//
// NOTE(russell_h): Ideally this wouldn't mutate the root passed to it, and would
// instead return a copy. Unfortunately copying the tree makes it difficult to
// correctly locate the signature. I'm opting, for now, to simply mutate the root
// parameter.
func (ctx *ValidationContext) transform(
	el *etree.Element,
	sig *types.Signature,
	ref *types.Reference) (*etree.Element, Canonicalizer, error) {
	transforms := ref.Transforms.Transforms

	if len(transforms) != 2 {
		return nil, nil, errors.New("Expected Enveloped and C14N transforms")
	}

	var canonicalizer Canonicalizer

	for _, transform := range transforms {
		algo := transform.Algorithm

		switch AlgorithmID(algo) {
		case EnvelopedSignatureAltorithmId:
			if !recursivelyRemoveElement(el, sig.UnderlyingElement()) {
				return nil, nil, errors.New("Error applying canonicalization transform: Signature not found")
			}

		case CanonicalXML10ExclusiveAlgorithmId:
			var prefixList string
			if transform.InclusiveNamespaces != nil {
				prefixList = transform.InclusiveNamespaces.PrefixList
			}

			canonicalizer = MakeC14N10ExclusiveCanonicalizerWithPrefixList(prefixList)

		case CanonicalXML11AlgorithmId:
			canonicalizer = MakeC14N11Canonicalizer()

		default:
			return nil, nil, errors.New("Unknown Transform Algorithm: " + algo)
		}
	}

	if canonicalizer == nil {
		return nil, nil, errors.New("Expected canonicalization transform")
	}

	return el, canonicalizer, nil
}

func (ctx *ValidationContext) digest(el *etree.Element, digestAlgorithmId string, canonicalizer Canonicalizer) ([]byte, error) {
	data, err := canonicalizer.Canonicalize(el)
	if err != nil {
		return nil, err
	}

	digestAlgorithm, ok := digestAlgorithmsByIdentifier[digestAlgorithmId]
	if !ok {
		return nil, errors.New("Unknown digest algorithm: " + digestAlgorithmId)
	}

	hash := digestAlgorithm.New()
	_, err = hash.Write(data)
	if err != nil {
		return nil, err
	}

	return hash.Sum(nil), nil
}

func (ctx *ValidationContext) verifySignedInfo(sig *types.Signature, canonicalizer Canonicalizer, signatureMethodId string, cert *x509.Certificate, decodedSignature []byte) error {
	signatureElement := sig.UnderlyingElement()

	signedInfo := signatureElement.FindElement(childPath(signatureElement.Space, SignedInfoTag))
	if signedInfo == nil {
		return errors.New("Missing SignedInfo")
	}

	// Canonicalize the xml
	canonical, err := canonicalSerialize(signedInfo)
	if err != nil {
		return err
	}

	signatureAlgorithm, ok := signatureMethodsByIdentifier[signatureMethodId]
	if !ok {
		return errors.New("Unknown signature method: " + signatureMethodId)
	}

	hash := signatureAlgorithm.New()
	_, err = hash.Write(canonical)
	if err != nil {
		return err
	}

	hashed := hash.Sum(nil)

	pubKey, ok := cert.PublicKey.(*rsa.PublicKey)
	if !ok {
		return errors.New("Invalid public key")
	}

	// Verify that the private key matching the public key from the cert was what was used to sign the 'SignedInfo' and produce the 'SignatureValue'
	err = rsa.VerifyPKCS1v15(pubKey, signatureAlgorithm, hashed[:], decodedSignature)
	if err != nil {
		return err
	}

	return nil
}

func (ctx *ValidationContext) validateSignature(el *etree.Element, sig *types.Signature, cert *x509.Certificate) (*etree.Element, error) {
	idAttr := el.SelectAttr(ctx.IdAttribute)
	if idAttr == nil || idAttr.Value == "" {
		return nil, errors.New("Missing ID attribute")
	}

	var ref *types.Reference

	// Find the first reference which references the top-level element
	for _, _ref := range sig.SignedInfo.References {
		if _ref.URI == "" || _ref.URI[1:] == idAttr.Value {
			ref = &_ref
		}
	}

	// Perform all transformations listed in the 'SignedInfo'
	// Basically, this means removing the 'SignedInfo'
	transformed, canonicalizer, err := ctx.transform(el, sig, ref)
	if err != nil {
		return nil, err
	}

	digestAlgorithm := ref.DigestAlgo.Algorithm

	// Digest the transformed XML and compare it to the 'DigestValue' from the 'SignedInfo'
	digest, err := ctx.digest(transformed, digestAlgorithm, canonicalizer)
	if err != nil {
		return nil, err
	}

	decodedDigestValue, err := base64.StdEncoding.DecodeString(ref.DigestValue)
	if err != nil {
		return nil, err
	}

	if !bytes.Equal(digest, decodedDigestValue) {
		return nil, errors.New("Signature could not be verified")
	}

	// Decode the 'SignatureValue' so we can compare against it
	decodedSignature, err := base64.StdEncoding.DecodeString(sig.SignatureValue.Data)
	if err != nil {
		return nil, errors.New("Could not decode signature")
	}

	// Actually verify the 'SignedInfo' was signed by a trusted source
	signatureMethod := sig.SignedInfo.SignatureMethod.Algorithm
	err = ctx.verifySignedInfo(sig, canonicalizer, signatureMethod, cert, decodedSignature)
	if err != nil {
		return nil, err
	}

	return transformed, nil
}

func contains(roots []*x509.Certificate, cert *x509.Certificate) bool {
	for _, root := range roots {
		if root.Equal(cert) {
			return true
		}
	}
	return false
}

// findSignature searches for a Signature element referencing the passed root element.
func (ctx *ValidationContext) findSignature(el *etree.Element) (*types.Signature, error) {
	idAttr := el.SelectAttr(ctx.IdAttribute)
	if idAttr == nil || idAttr.Value == "" {
		return nil, errors.New("Missing ID attribute")
	}

	var sig *types.Signature

	// Traverse the tree looking for a Signature element
	err := etreeutils.NSFindIterate(el, Namespace, SignatureTag, func(ctx etreeutils.NSContext, el *etree.Element) error {

		found := false
		err := etreeutils.NSFindIterateCtx(ctx, el, Namespace, SignedInfoTag,
			func(ctx etreeutils.NSContext, signedInfo *etree.Element) error {
				// Ignore any SignedInfo that isn't an immediate descendent of Signature.
				if signedInfo.Parent() != el {
					return nil
				}

				detachedSignedInfo, err := etreeutils.NSDetatch(ctx, signedInfo)
				if err != nil {
					return err
				}

				c14NMethod := detachedSignedInfo.FindElement(childPath(detachedSignedInfo.Space, CanonicalizationMethodTag))
				if c14NMethod == nil {
					return errors.New("missing CanonicalizationMethod on Signature")
				}

				c14NAlgorithm := c14NMethod.SelectAttrValue(AlgorithmAttr, "")

				var canonicalSignedInfo *etree.Element

				switch AlgorithmID(c14NAlgorithm) {
				case CanonicalXML10ExclusiveAlgorithmId:
					err := etreeutils.TransformExcC14n(detachedSignedInfo, "")
					if err != nil {
						return err
					}

					// NOTE: TransformExcC14n transforms the element in-place,
					// while canonicalPrep isn't meant to. Once we standardize
					// this behavior we can drop this, as well as the adding and
					// removing of elements below.
					canonicalSignedInfo = detachedSignedInfo

				case CanonicalXML11AlgorithmId:
					canonicalSignedInfo = canonicalPrep(detachedSignedInfo, map[string]struct{}{})

				default:
					return fmt.Errorf("invalid CanonicalizationMethod on Signature: %s", c14NAlgorithm)
				}

				el.RemoveChild(signedInfo)
				el.AddChild(canonicalSignedInfo)

				found = true

				return etreeutils.ErrTraversalHalted
			})
		if err != nil {
			return err
		}

		if !found {
			return errors.New("Missing SignedInfo")
		}

		// Unmarshal the signature into a structured Signature type
		_sig := &types.Signature{}
		err = etreeutils.NSUnmarshalElement(ctx, el, _sig)
		if err != nil {
			return err
		}

		// Traverse references in the signature to determine whether it has at least
		// one reference to the top level element. If so, conclude the search.
		for _, ref := range _sig.SignedInfo.References {
			if ref.URI == "" || ref.URI[1:] == idAttr.Value {
				sig = _sig
				return etreeutils.ErrTraversalHalted
			}
		}

		return nil
	})

	if err != nil {
		return nil, err
	}

	if sig == nil {
		return nil, ErrMissingSignature
	}

	return sig, nil
}

func (ctx *ValidationContext) verifyCertificate(sig *types.Signature) (*x509.Certificate, error) {
	now := ctx.Clock.Now()

	roots, err := ctx.CertificateStore.Certificates()
	if err != nil {
		return nil, err
	}

	var cert *x509.Certificate

	if sig.KeyInfo != nil {
		// If the Signature includes KeyInfo, extract the certificate from there
		if sig.KeyInfo.X509Data.X509Certificate.Data == "" {
			return nil, errors.New("missing X509Certificate within KeyInfo")
		}

		certData, err := base64.StdEncoding.DecodeString(sig.KeyInfo.X509Data.X509Certificate.Data)
		if err != nil {
			return nil, errors.New("Failed to parse certificate")
		}

		cert, err = x509.ParseCertificate(certData)
		if err != nil {
			return nil, err
		}
	} else {
		// If the Signature doesn't have KeyInfo, Use the root certificate if there is only one
		if len(roots) == 1 {
			cert = roots[0]
		} else {
			return nil, errors.New("Missing x509 Element")
		}
	}

	// Verify that the certificate is one we trust
	if !contains(roots, cert) {
		return nil, errors.New("Could not verify certificate against trusted certs")
	}

	if now.Before(cert.NotBefore) || now.After(cert.NotAfter) {
		return nil, errors.New("Cert is not valid at this time")
	}

	return cert, nil
}

// Validate verifies that the passed element contains a valid enveloped signature
// matching a currently-valid certificate in the context's CertificateStore.
func (ctx *ValidationContext) Validate(el *etree.Element) (*etree.Element, error) {
	// Make a copy of the element to avoid mutating the one we were passed.
	el = el.Copy()

	sig, err := ctx.findSignature(el)
	if err != nil {
		return nil, err
	}

	cert, err := ctx.verifyCertificate(sig)
	if err != nil {
		return nil, err
	}

	return ctx.validateSignature(el, sig, cert)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"net"
	"strconv"
)

func Reader(reader io.Reader, messageChannel chan []byte) {
	buf := make([]byte, 1024)
	for {
		n, err := reader.Read(buf[:])
		if err != nil {
			logger.Error("Can't read: %s", err.Error())
			return
		}

		if *debug {
			logger.Debug(fmt.Sprintf("Received from socket:  %s", buf[0:n]))
		}

		select {
		case messageChannel <- buf[0:n]:
		default:
		}
	}
}

func Sender(host string, port int, messageChannel chan []byte) {

	ServerAddress, err := net.ResolveUDPAddr("udp", host + ":" + strconv.Itoa(port))
	if err != nil {
		logger.Error("Can't resolve remote UDP address: %s", err.Error())
		return
	}

	LocalAddress, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
	if err != nil {
		logger.Error("Can't resolve local UDP address: %s", err.Error())
		return
	}

	Conn, err := net.DialUDP("udp", LocalAddress, ServerAddress)
	if err != nil {
		logger.Error("Can't dial up: %s", err.Error())
		return
	}

	defer Conn.Close()

	for {
		select {
		case msg := <-messageChannel:

			if *debug {
				logger.Debug(fmt.Sprintf("Writing to UDP socket:  %s", msg[:]))
			}

			_, err := Conn.Write(msg[:])
			if err != nil {
				logger.Error("Can't write: %s", err.Error())
			}
		}
	}
}

<commit_msg>fixng problem with UDP connection<commit_after>package main

import (
	"fmt"
	"io"
	"net"
	"strconv"
)

func Reader(reader io.Reader, messageChannel chan []byte) {
	buf := make([]byte, 4096)
	for {
		n, err := reader.Read(buf[:])
		if err != nil {
			logger.Error("Can't read: %s", err.Error())
			return
		}

		if *debug {
			logger.Debug(fmt.Sprintf("Received from socket:  %s", buf[0:n]))
		}

		select {
		case messageChannel <- buf[0:n]:
		default:
		}
	}
}

func Sender(host string, port int, messageChannel chan []byte) {

	ServerAddress, err := net.ResolveUDPAddr("udp", host + ":" + strconv.Itoa(port))
	if err != nil {
		logger.Error("Can't resolve remote UDP address: %s", err.Error())
		return
	}

	Conn, err := net.DialUDP("udp", nil, ServerAddress)
	if err != nil {
		logger.Error("Can't dial up: %s", err.Error())
		return
	}

	defer Conn.Close()

	for {
		select {
		case msg := <-messageChannel:

			if *debug {
				logger.Debug(fmt.Sprintf("Writing to UDP socket: %s", msg[:]))
			}

			_, err := Conn.Write(msg[:])
			if err != nil {
				logger.Error("Can't write: %s", err.Error())
			}
		}
	}
}

<|endoftext|>


<commit_before>// Package service implements the store service interface
package service

import (
	"context"
	"io"
	"time"

	"github.com/micro/go-micro/v2/client"
	"github.com/micro/go-micro/v2/errors"
	"github.com/micro/go-micro/v2/metadata"
	"github.com/micro/go-micro/v2/store"
	pb "github.com/micro/go-micro/v2/store/service/proto"
)

type serviceStore struct {
	options store.Options

	// Namespace to use
	Namespace string

	// Addresses of the nodes
	Nodes []string

	// Prefix to use
	Prefix string

	// Suffix to use
	Suffix string

	// store service client
	Client pb.StoreService
}

func (s *serviceStore) Close() error {
	return nil
}

func (s *serviceStore) Init(opts ...store.Option) error {
	for _, o := range opts {
		o(&s.options)
	}
	s.Namespace = s.options.Database
	s.Prefix = s.options.Table
	s.Nodes = s.options.Nodes

	return nil
}

func (s *serviceStore) Context() context.Context {
	ctx := context.Background()

	md := make(metadata.Metadata)

	if len(s.Namespace) > 0 {
		md["Micro-Namespace"] = s.Namespace
	}

	if len(s.Prefix) > 0 {
		md["Micro-Prefix"] = s.Prefix
	}

	return metadata.NewContext(ctx, md)
}

// Sync all the known records
func (s *serviceStore) List(opts ...store.ListOption) ([]string, error) {
	stream, err := s.Client.List(s.Context(), &pb.ListRequest{}, client.WithAddress(s.Nodes...))
	if err != nil && errors.Equal(err, errors.NotFound("", "")) {
		return nil, store.ErrNotFound
	} else if err != nil {
		return nil, err
	}
	defer stream.Close()

	var keys []string

	for {
		rsp, err := stream.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			return keys, err
		}

		for _, key := range rsp.Keys {
			keys = append(keys, key)
		}
	}

	return keys, nil
}

// Read a record with key
func (s *serviceStore) Read(key string, opts ...store.ReadOption) ([]*store.Record, error) {
	var options store.ReadOptions
	for _, o := range opts {
		o(&options)
	}

	rsp, err := s.Client.Read(s.Context(), &pb.ReadRequest{
		Key: key,
		Options: &pb.ReadOptions{
			Prefix: options.Prefix,
		},
	}, client.WithAddress(s.Nodes...))
	if err != nil && errors.Equal(err, errors.NotFound("", "")) {
		return nil, store.ErrNotFound
	} else if err != nil {
		return nil, err
	}

	records := make([]*store.Record, 0, len(rsp.Records))

	for _, val := range rsp.Records {
		records = append(records, &store.Record{
			Key:    val.Key,
			Value:  val.Value,
			Expiry: time.Duration(val.Expiry) * time.Second,
		})
	}

	return records, nil
}

// Write a record
func (s *serviceStore) Write(record *store.Record, opts ...store.WriteOption) error {
	_, err := s.Client.Write(s.Context(), &pb.WriteRequest{
		Record: &pb.Record{
			Key:    record.Key,
			Value:  record.Value,
			Expiry: int64(record.Expiry.Seconds()),
		},
	}, client.WithAddress(s.Nodes...))
	if err != nil && errors.Equal(err, errors.NotFound("", "")) {
		return store.ErrNotFound
	}

	return err
}

// Delete a record with key
func (s *serviceStore) Delete(key string, opts ...store.DeleteOption) error {
	_, err := s.Client.Delete(s.Context(), &pb.DeleteRequest{
		Key: key,
	}, client.WithAddress(s.Nodes...))
	if err != nil && errors.Equal(err, errors.NotFound("", "")) {
		return store.ErrNotFound
	}

	return err
}

func (s *serviceStore) String() string {
	return "service"
}

func (s *serviceStore) Options() store.Options {
	return s.options
}

// NewStore returns a new store service implementation
func NewStore(opts ...store.Option) store.Store {
	var options store.Options
	for _, o := range opts {
		o(&options)
	}

	service := &serviceStore{
		options:   options,
		Namespace: options.Database,
		Prefix:    options.Table,
		Nodes:     options.Nodes,
		Client:    pb.NewStoreService("go.micro.store", client.DefaultClient),
	}

	return service
}
<commit_msg>set database/table in header<commit_after>// Package service implements the store service interface
package service

import (
	"context"
	"io"
	"time"

	"github.com/micro/go-micro/v2/client"
	"github.com/micro/go-micro/v2/errors"
	"github.com/micro/go-micro/v2/metadata"
	"github.com/micro/go-micro/v2/store"
	pb "github.com/micro/go-micro/v2/store/service/proto"
)

type serviceStore struct {
	options store.Options

	// The database to use
	Database string

	// The table to use
	Table string

	// Addresses of the nodes
	Nodes []string

	// store service client
	Client pb.StoreService
}

func (s *serviceStore) Close() error {
	return nil
}

func (s *serviceStore) Init(opts ...store.Option) error {
	for _, o := range opts {
		o(&s.options)
	}
	s.Database = s.options.Database
	s.Table = s.options.Table
	s.Nodes = s.options.Nodes

	return nil
}

func (s *serviceStore) Context() context.Context {
	ctx := context.Background()

	md := make(metadata.Metadata)

	if len(s.Database) > 0 {
		md["Micro-Database"] = s.Database
	}

	if len(s.Table) > 0 {
		md["Micro-Table"] = s.Table
	}

	return metadata.NewContext(ctx, md)
}

// Sync all the known records
func (s *serviceStore) List(opts ...store.ListOption) ([]string, error) {
	stream, err := s.Client.List(s.Context(), &pb.ListRequest{}, client.WithAddress(s.Nodes...))
	if err != nil && errors.Equal(err, errors.NotFound("", "")) {
		return nil, store.ErrNotFound
	} else if err != nil {
		return nil, err
	}
	defer stream.Close()

	var keys []string

	for {
		rsp, err := stream.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			return keys, err
		}

		for _, key := range rsp.Keys {
			keys = append(keys, key)
		}
	}

	return keys, nil
}

// Read a record with key
func (s *serviceStore) Read(key string, opts ...store.ReadOption) ([]*store.Record, error) {
	var options store.ReadOptions
	for _, o := range opts {
		o(&options)
	}

	rsp, err := s.Client.Read(s.Context(), &pb.ReadRequest{
		Key: key,
		Options: &pb.ReadOptions{
			Prefix: options.Prefix,
		},
	}, client.WithAddress(s.Nodes...))
	if err != nil && errors.Equal(err, errors.NotFound("", "")) {
		return nil, store.ErrNotFound
	} else if err != nil {
		return nil, err
	}

	records := make([]*store.Record, 0, len(rsp.Records))

	for _, val := range rsp.Records {
		records = append(records, &store.Record{
			Key:    val.Key,
			Value:  val.Value,
			Expiry: time.Duration(val.Expiry) * time.Second,
		})
	}

	return records, nil
}

// Write a record
func (s *serviceStore) Write(record *store.Record, opts ...store.WriteOption) error {
	_, err := s.Client.Write(s.Context(), &pb.WriteRequest{
		Record: &pb.Record{
			Key:    record.Key,
			Value:  record.Value,
			Expiry: int64(record.Expiry.Seconds()),
		},
	}, client.WithAddress(s.Nodes...))
	if err != nil && errors.Equal(err, errors.NotFound("", "")) {
		return store.ErrNotFound
	}

	return err
}

// Delete a record with key
func (s *serviceStore) Delete(key string, opts ...store.DeleteOption) error {
	_, err := s.Client.Delete(s.Context(), &pb.DeleteRequest{
		Key: key,
	}, client.WithAddress(s.Nodes...))
	if err != nil && errors.Equal(err, errors.NotFound("", "")) {
		return store.ErrNotFound
	}

	return err
}

func (s *serviceStore) String() string {
	return "service"
}

func (s *serviceStore) Options() store.Options {
	return s.options
}

// NewStore returns a new store service implementation
func NewStore(opts ...store.Option) store.Store {
	var options store.Options
	for _, o := range opts {
		o(&options)
	}

	service := &serviceStore{
		options:  options,
		Database: options.Database,
		Table:    options.Table,
		Nodes:    options.Nodes,
		Client:   pb.NewStoreService("go.micro.store", client.DefaultClient),
	}

	return service
}
<|endoftext|>


<commit_before>package puppetmasterless

import (
	"io/ioutil"
	"os"
	"strings"
	"testing"

	"github.com/mitchellh/packer/packer"
)

func testConfig() map[string]interface{} {
	tf, err := ioutil.TempFile("", "packer")
	if err != nil {
		panic(err)
	}

	return map[string]interface{}{
		"manifest_file": tf.Name(),
	}
}

func TestProvisioner_Impl(t *testing.T) {
	var raw interface{}
	raw = &Provisioner{}
	if _, ok := raw.(packer.Provisioner); !ok {
		t.Fatalf("must be a Provisioner")
	}
}

func TestProvisionerPrepare_puppetBinDir(t *testing.T) {
	config := testConfig()

	delete(config, "puppet_bin_dir")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Test with a good one
	tf, err := ioutil.TempFile("", "packer")
	if err != nil {
		t.Fatalf("error tempfile: %s", err)
	}
	defer os.Remove(tf.Name())

	config["puppet_bin_dir"] = tf.Name()
	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestProvisionerPrepare_hieraConfigPath(t *testing.T) {
	config := testConfig()

	delete(config, "hiera_config_path")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Test with a good one
	tf, err := ioutil.TempFile("", "packer")
	if err != nil {
		t.Fatalf("error tempfile: %s", err)
	}
	defer os.Remove(tf.Name())

	config["hiera_config_path"] = tf.Name()
	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestProvisionerPrepare_manifestFile(t *testing.T) {
	config := testConfig()

	delete(config, "manifest_file")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err == nil {
		t.Fatal("should be an error")
	}

	// Test with a good one
	tf, err := ioutil.TempFile("", "packer")
	if err != nil {
		t.Fatalf("error tempfile: %s", err)
	}
	defer os.Remove(tf.Name())

	config["manifest_file"] = tf.Name()
	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestProvisionerPrepare_manifestDir(t *testing.T) {
	config := testConfig()

	delete(config, "manifestdir")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Test with a good one
	td, err := ioutil.TempDir("", "packer")
	if err != nil {
		t.Fatalf("error: %s", err)
	}
	defer os.RemoveAll(td)

	config["manifest_dir"] = td
	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestProvisionerPrepare_modulePaths(t *testing.T) {
	config := testConfig()

	delete(config, "module_paths")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Test with bad paths
	config["module_paths"] = []string{"i-should-not-exist"}
	p = new(Provisioner)
	err = p.Prepare(config)
	if err == nil {
		t.Fatal("should be an error")
	}

	// Test with a good one
	td, err := ioutil.TempDir("", "packer")
	if err != nil {
		t.Fatalf("error: %s", err)
	}
	defer os.RemoveAll(td)

	config["module_paths"] = []string{td}
	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestProvisionerPrepare_facterFacts(t *testing.T) {
	config := testConfig()

	delete(config, "facter")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Test with malformed fact
	config["facter"] = "fact=stringified"
	p = new(Provisioner)
	err = p.Prepare(config)
	if err == nil {
		t.Fatal("should be an error")
	}

	// Test with a good one
	td, err := ioutil.TempDir("", "packer")
	if err != nil {
		t.Fatalf("error: %s", err)
	}
	defer os.RemoveAll(td)

	facts := make(map[string]string)
	facts["fact_name"] = "fact_value"
	config["facter"] = facts

	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Make sure the default facts are present
	delete(config, "facter")
	p = new(Provisioner)
	err = p.Prepare(config)
	if p.config.Facter == nil {
		t.Fatalf("err: Default facts are not set in the Puppet provisioner!")
	}

	if _, ok := p.config.Facter["packer_build_name"]; !ok {
		t.Fatalf("err: packer_build_name fact not set in the Puppet provisioner!")
	}

	if _, ok := p.config.Facter["packer_builder_type"]; !ok {
		t.Fatalf("err: packer_builder_type fact not set in the Puppet provisioner!")
	}
}

func TestProvisionerPrepare_extraArguments(t *testing.T) {
	config := testConfig()

	// Test with missing parameter
	delete(config, "extra_arguments")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Test with malformed value
	config["extra_arguments"] = "{{}}"
	p = new(Provisioner)
	err = p.Prepare(config)
	if err == nil {
		t.Fatal("should be an error")
	}

	// Test with valid values
	config["extra_arguments"] = []string{
		"arg",
	}

	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestProvisionerProvision_extraArguments(t *testing.T) {
	config := testConfig()
	ui := &packer.MachineReadableUi{
		Writer: ioutil.Discard,
	}
	comm := new(packer.MockCommunicator)

	extraArguments := []string{
		"--some-arg=yup",
		"--some-other-arg",
	}
	config["extra_arguments"] = extraArguments

	// Test with valid values
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	err = p.Provision(ui, comm)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	expectedArgs := strings.Join(extraArguments, " ")

	if !strings.Contains(comm.StartCmd.Command, expectedArgs) {
		t.Fatalf("Command %q doesn't contain the expected arguments %q", comm.StartCmd.Command, expectedArgs)
	}

	// Test with missing parameter
	delete(config, "extra_arguments")

	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	err = p.Provision(ui, comm)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Check the expected `extra_arguments` position for an empty value
	splitCommand := strings.Split(comm.StartCmd.Command, " ")
	if "" == splitCommand[len(splitCommand)-2] {
		t.Fatalf("Command %q contains an extra-space which may cause arg parsing issues", comm.StartCmd.Command)
	}
}
<commit_msg>Add more tests for puppet-masterless provisioner<commit_after>package puppetmasterless

import (
	"io/ioutil"
	"os"
	"strings"
	"testing"

	"github.com/mitchellh/packer/packer"
)

func testConfig() map[string]interface{} {
	tf, err := ioutil.TempFile("", "packer")
	if err != nil {
		panic(err)
	}

	return map[string]interface{}{
		"manifest_file": tf.Name(),
	}
}

func TestProvisioner_Impl(t *testing.T) {
	var raw interface{}
	raw = &Provisioner{}
	if _, ok := raw.(packer.Provisioner); !ok {
		t.Fatalf("must be a Provisioner")
	}
}

func TestProvisionerPrepare_puppetBinDir(t *testing.T) {
	config := testConfig()

	delete(config, "puppet_bin_dir")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Test with a good one
	tf, err := ioutil.TempFile("", "packer")
	if err != nil {
		t.Fatalf("error tempfile: %s", err)
	}
	defer os.Remove(tf.Name())

	config["puppet_bin_dir"] = tf.Name()
	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestProvisionerPrepare_hieraConfigPath(t *testing.T) {
	config := testConfig()

	delete(config, "hiera_config_path")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Test with a good one
	tf, err := ioutil.TempFile("", "packer")
	if err != nil {
		t.Fatalf("error tempfile: %s", err)
	}
	defer os.Remove(tf.Name())

	config["hiera_config_path"] = tf.Name()
	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestProvisionerPrepare_manifestFile(t *testing.T) {
	config := testConfig()

	delete(config, "manifest_file")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err == nil {
		t.Fatal("should be an error")
	}

	// Test with a good one
	tf, err := ioutil.TempFile("", "packer")
	if err != nil {
		t.Fatalf("error tempfile: %s", err)
	}
	defer os.Remove(tf.Name())

	config["manifest_file"] = tf.Name()
	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestProvisionerPrepare_manifestDir(t *testing.T) {
	config := testConfig()

	delete(config, "manifestdir")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Test with a good one
	td, err := ioutil.TempDir("", "packer")
	if err != nil {
		t.Fatalf("error: %s", err)
	}
	defer os.RemoveAll(td)

	config["manifest_dir"] = td
	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestProvisionerPrepare_modulePaths(t *testing.T) {
	config := testConfig()

	delete(config, "module_paths")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Test with bad paths
	config["module_paths"] = []string{"i-should-not-exist"}
	p = new(Provisioner)
	err = p.Prepare(config)
	if err == nil {
		t.Fatal("should be an error")
	}

	// Test with a good one
	td, err := ioutil.TempDir("", "packer")
	if err != nil {
		t.Fatalf("error: %s", err)
	}
	defer os.RemoveAll(td)

	config["module_paths"] = []string{td}
	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestProvisionerPrepare_facterFacts(t *testing.T) {
	config := testConfig()

	delete(config, "facter")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Test with malformed fact
	config["facter"] = "fact=stringified"
	p = new(Provisioner)
	err = p.Prepare(config)
	if err == nil {
		t.Fatal("should be an error")
	}

	// Test with a good one
	td, err := ioutil.TempDir("", "packer")
	if err != nil {
		t.Fatalf("error: %s", err)
	}
	defer os.RemoveAll(td)

	facts := make(map[string]string)
	facts["fact_name"] = "fact_value"
	config["facter"] = facts

	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Make sure the default facts are present
	delete(config, "facter")
	p = new(Provisioner)
	err = p.Prepare(config)
	if p.config.Facter == nil {
		t.Fatalf("err: Default facts are not set in the Puppet provisioner!")
	}

	if _, ok := p.config.Facter["packer_build_name"]; !ok {
		t.Fatalf("err: packer_build_name fact not set in the Puppet provisioner!")
	}

	if _, ok := p.config.Facter["packer_builder_type"]; !ok {
		t.Fatalf("err: packer_builder_type fact not set in the Puppet provisioner!")
	}
}

func TestProvisionerPrepare_extraArguments(t *testing.T) {
	config := testConfig()

	// Test with missing parameter
	delete(config, "extra_arguments")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Test with malformed value
	config["extra_arguments"] = "{{}}"
	p = new(Provisioner)
	err = p.Prepare(config)
	if err == nil {
		t.Fatal("should be an error")
	}

	// Test with valid values
	config["extra_arguments"] = []string{
		"arg",
	}

	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
}

func TestProvisionerPrepare_stagingDir(t *testing.T) {
	config := testConfig()

	delete(config, "staging_directory")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Make sure the default staging directory is correct
	if p.config.StagingDir != "/tmp/packer-puppet-masterless" {
		t.Fatalf("err: Default staging_directory is not set in the Puppet provisioner!")
	}

	// Make sure default staging directory can be overridden
	config["staging_directory"] = "/tmp/override"
	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	if p.config.StagingDir != "/tmp/override" {
		t.Fatalf("err: Overridden staging_directory is not set correctly in the Puppet provisioner!")
	}
}

func TestProvisionerPrepare_workingDir(t *testing.T) {
	config := testConfig()

	delete(config, "working_directory")
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Make sure default working dir and staging dir are the same
	if p.config.WorkingDir != p.config.StagingDir {
		t.Fatalf("err: Default working_directory is not set to the same value as default staging_directory in the Puppet provisioner!")
	}

	// Make sure the default working directory is correct
	if p.config.WorkingDir != "/tmp/packer-puppet-masterless" {
		t.Fatalf("err: Default working_directory is not set in the Puppet provisioner!")
	}

	// Make sure default working directory can be overridden
	config["working_directory"] = "/tmp/override"
	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	if p.config.WorkingDir != "/tmp/override" {
		t.Fatalf("err: Overridden working_directory is not set correctly in the Puppet provisioner!")
	}
}

func TestProvisionerProvision_extraArguments(t *testing.T) {
	config := testConfig()
	ui := &packer.MachineReadableUi{
		Writer: ioutil.Discard,
	}
	comm := new(packer.MockCommunicator)

	extraArguments := []string{
		"--some-arg=yup",
		"--some-other-arg",
	}
	config["extra_arguments"] = extraArguments

	// Test with valid values
	p := new(Provisioner)
	err := p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	err = p.Provision(ui, comm)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	expectedArgs := strings.Join(extraArguments, " ")

	if !strings.Contains(comm.StartCmd.Command, expectedArgs) {
		t.Fatalf("Command %q doesn't contain the expected arguments %q", comm.StartCmd.Command, expectedArgs)
	}

	// Test with missing parameter
	delete(config, "extra_arguments")

	p = new(Provisioner)
	err = p.Prepare(config)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	err = p.Provision(ui, comm)
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	// Check the expected `extra_arguments` position for an empty value
	splitCommand := strings.Split(comm.StartCmd.Command, " ")
	if "" == splitCommand[len(splitCommand)-2] {
		t.Fatalf("Command %q contains an extra-space which may cause arg parsing issues", comm.StartCmd.Command)
	}
}
<|endoftext|>


<commit_before>package yorm

import (
	"bytes"
	"errors"
	"fmt"
	"reflect"
	"strings"
	"time"
)

//Insert  return lastInsertId and error if has
func (ex *executor) Insert(i interface{}, args ...string) (int64, error) {

	q, err := newTableSetter(reflect.ValueOf(i))
	if q == nil {
		return 0, err
	}

	clause := &bytes.Buffer{}
	clause.WriteString("INSERT INTO ")
	clause.WriteString(q.table)
	clause.WriteString(" SET ")

	fs := &bytes.Buffer{}
	dests := []interface{}{}

	e := reflect.ValueOf(i).Elem()

	var pk reflect.Value
	var columns []*column

	if len(args) == 0 {
		columns = q.columns

	} else {
		for _, arg := range args {
			arg = strings.ToLower(arg)
			for _, c := range q.columns {
				if strings.ToLower(c.fieldName) == arg || strings.ToLower(c.name) == arg {
					columns = append(columns, c)
				}
			}
		}
	}
	for _, c := range columns {
		v := e.FieldByName(c.fieldName)
		//todo how to handle a zero value ???
		//		if !v.IsValid() {
		//			continue
		//		}

		if c.isAuto {
			pk = v
			continue
		}
		vi := v.Interface()
		if v.Type() == TimeType {
			//zero time ,skip insert
			if vi.(time.Time).IsZero() {
				continue
			}
			vi = vi.(time.Time).Format(longSimpleTimeFormat)
		}

		fs.WriteString("," + c.name + "=?")
		dests = append(dests, fmt.Sprintf("%v", vi))
	}
	if fs.Len() == 0 {
		return 0, errors.New("no filed to insert")
	}

	clause.Write(fs.Bytes()[1:])

	r, err := ex.exec(clause.String(), dests...)
	if err != nil {
		return 0, err
	}
	id, err := r.LastInsertId()
	if id > 0 && pk.IsValid() {
		pk.SetInt(id)
	}
	return id, err
}

//Insert  return lastInsertId and error if has
func (ex *tranExecutor) Insert(i interface{}, args ...string) (int64, error) {

	q, err := newTableSetter(reflect.ValueOf(i))
	if q == nil {
		return 0, err
	}

	clause := &bytes.Buffer{}
	clause.WriteString("INSERT INTO ")
	clause.WriteString(q.table)
	clause.WriteString(" SET ")

	fs := &bytes.Buffer{}
	dests := []interface{}{}

	e := reflect.ValueOf(i).Elem()

	var pk reflect.Value
	var columns []*column

	if len(args) == 0 {
		columns = q.columns

	} else {
		for _, arg := range args {
			arg = strings.ToLower(arg)
			for _, c := range q.columns {
				if strings.ToLower(c.fieldName) == arg || strings.ToLower(c.name) == arg {
					columns = append(columns, c)
				}
			}
		}
	}
	for _, c := range columns {
		v := e.FieldByName(c.fieldName)
		//todo how to handle a zero value ???
		//		if !v.IsValid() {
		//			continue
		//		}

		if c.isAuto {
			pk = v
			continue
		}
		vi := v.Interface()
		if v.Type() == TimeType {
			//zero time ,skip insert
			if vi.(time.Time).IsZero() {
				continue
			}
			vi = vi.(time.Time).Format(longSimpleTimeFormat)
		}

		fs.WriteString("," + c.name + "=?")
		dests = append(dests, fmt.Sprintf("%v", vi))
	}
	if fs.Len() == 0 {
		return 0, errors.New("no filed to insert")
	}

	clause.Write(fs.Bytes()[1:])

	r, err := ex.exec(clause.String(), dests...)
	if err != nil {
		return 0, err
	}
	id, err := r.LastInsertId()
	if id > 0 && pk.IsValid() {
		pk.SetInt(id)
	}
	return id, err
}

//Insert insert a record.
func Insert(i interface{}, args ...string) (int64, error) {
	return defaultExecutor.Insert(i, args...)
}
<commit_msg>bool 自动变成int<commit_after>package yorm

import (
	"bytes"
	"errors"
	"fmt"
	"reflect"
	"strings"
	"time"
)

//Insert  return lastInsertId and error if has
func (ex *executor) Insert(i interface{}, args ...string) (int64, error) {

	q, err := newTableSetter(reflect.ValueOf(i))
	if q == nil {
		return 0, err
	}

	clause := &bytes.Buffer{}
	clause.WriteString("INSERT INTO ")
	clause.WriteString(q.table)
	clause.WriteString(" SET ")

	fs := &bytes.Buffer{}
	dests := []interface{}{}

	e := reflect.ValueOf(i).Elem()

	var pk reflect.Value
	var columns []*column

	if len(args) == 0 {
		columns = q.columns

	} else {
		for _, arg := range args {
			arg = strings.ToLower(arg)
			for _, c := range q.columns {
				if strings.ToLower(c.fieldName) == arg || strings.ToLower(c.name) == arg {
					columns = append(columns, c)
				}
			}
		}
	}
	for _, c := range columns {
		v := e.FieldByName(c.fieldName)
		//todo how to handle a zero value ???
		//		if !v.IsValid() {
		//			continue
		//		}

		if c.isAuto {
			pk = v
			continue
		}
		vi := v.Interface()
		switch v.Type() {

		case TimeType:
			//zero time ,skip insert
			if vi.(time.Time).IsZero() {
				continue
			}
			vi = vi.(time.Time).Format(longSimpleTimeFormat)

		case reflect.Bool:
			if vi.(bool) {
				vi = 1
			} else {
				vi = 0
			}
		}

		fs.WriteString("," + c.name + "=?")
		dests = append(dests, fmt.Sprintf("%v", vi))
	}
	if fs.Len() == 0 {
		return 0, errors.New("no filed to insert")
	}

	clause.Write(fs.Bytes()[1:])

	r, err := ex.exec(clause.String(), dests...)
	if err != nil {
		return 0, err
	}
	id, err := r.LastInsertId()
	if id > 0 && pk.IsValid() {
		pk.SetInt(id)
	}
	return id, err
}

//Insert  return lastInsertId and error if has
func (ex *tranExecutor) Insert(i interface{}, args ...string) (int64, error) {

	q, err := newTableSetter(reflect.ValueOf(i))
	if q == nil {
		return 0, err
	}

	clause := &bytes.Buffer{}
	clause.WriteString("INSERT INTO ")
	clause.WriteString(q.table)
	clause.WriteString(" SET ")

	fs := &bytes.Buffer{}
	dests := []interface{}{}

	e := reflect.ValueOf(i).Elem()

	var pk reflect.Value
	var columns []*column

	if len(args) == 0 {
		columns = q.columns

	} else {
		for _, arg := range args {
			arg = strings.ToLower(arg)
			for _, c := range q.columns {
				if strings.ToLower(c.fieldName) == arg || strings.ToLower(c.name) == arg {
					columns = append(columns, c)
				}
			}
		}
	}
	for _, c := range columns {
		v := e.FieldByName(c.fieldName)
		//todo how to handle a zero value ???
		//		if !v.IsValid() {
		//			continue
		//		}

		if c.isAuto {
			pk = v
			continue
		}
		vi := v.Interface()
		if v.Type() == TimeType {
			//zero time ,skip insert
			if vi.(time.Time).IsZero() {
				continue
			}
			vi = vi.(time.Time).Format(longSimpleTimeFormat)
		}

		fs.WriteString("," + c.name + "=?")
		dests = append(dests, fmt.Sprintf("%v", vi))
	}
	if fs.Len() == 0 {
		return 0, errors.New("no filed to insert")
	}

	clause.Write(fs.Bytes()[1:])

	r, err := ex.exec(clause.String(), dests...)
	if err != nil {
		return 0, err
	}
	id, err := r.LastInsertId()
	if id > 0 && pk.IsValid() {
		pk.SetInt(id)
	}
	return id, err
}

//Insert insert a record.
func Insert(i interface{}, args ...string) (int64, error) {
	return defaultExecutor.Insert(i, args...)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"io"
	"time"

	"github.com/abrander/gansoi/agents"
)

type (
	// Check defines a check to be conducted by Gansoi.
	Check struct {
		ID        string          `json:"id"`
		AgentID   string          `json:"agent"`
		Interval  time.Duration   `json:"interval"`
		Node      string          `json:"node"`
		Arguments json.RawMessage `json:"arguments"`
		Agent     agents.Agent    `json:"-"`
	}

	// CheckResult describes the result of one or more checks.
	CheckResult struct {
		ID      int         `json:"id"`
		CheckID string      `json:"check_id"`
		Error   string      `json:"error"`
		First   time.Time   `json:"first"`
		Last    time.Time   `json:"last"`
		Streak  int         `json:"streak"`
		Results interface{} `json:"results"`
	}
)

// NewCheckFromReader will instantiate a new check based on a JSON stream.
func NewCheckFromReader(r io.Reader) (*Check, error) {
	var check Check

	decoder := json.NewDecoder(r)
	err := decoder.Decode(&check)
	if err != nil {
		return nil, err
	}

	check.Agent = agents.GetAgent(check.AgentID)
	if check.Agent == nil {
		return nil, errors.New("Agent not found")
	}

	err = json.Unmarshal(check.Arguments, &check.Agent)
	if err != nil {
		return nil, err
	}

	return &check, nil
}

// NewCheckFromBytes will instantiate a new check from a byte array of JSON.
func NewCheckFromBytes(b []byte) (*Check, error) {
	return NewCheckFromReader(bytes.NewBuffer(b))
}

// Check implements agents.Agent. Will return true if the result changed.
func (c *Check) Check(result *CheckResult) bool {
	result.CheckID = c.ID

	var err error
	var errorString string

	result.Results, err = c.Agent.Check()
	if err != nil {
		errorString = err.Error()
	}

	result.Last = time.Now()
	if result.First == (time.Time{}) {
		result.First = result.Last
	}

	if errorString == result.Error {
		result.Streak++
		return false
	}

	// If we arrive here, we have witnessed a state change. Reset CheckResult
	// and save new error from check.
	result.ID = 0
	result.Streak = 1
	result.First = result.Last
	result.Error = errorString

	return true
}
<commit_msg>Implement json.Unmarshaler for Check.<commit_after>package main

import (
	"encoding/json"
	"errors"
	"time"

	"github.com/abrander/gansoi/agents"
	"github.com/abrander/gansoi/database"
)

type (
	// Check defines a check to be conducted by Gansoi.
	Check struct {
		ID        string          `json:"id"`
		AgentID   string          `json:"agent"`
		Interval  time.Duration   `json:"interval"`
		Node      string          `json:"node"`
		Arguments json.RawMessage `json:"arguments"`
		Agent     agents.Agent    `json:"-"`
	}

	checkProxy struct {
		ID        string          `json:"id"`
		AgentID   string          `json:"agent"`
		Interval  time.Duration   `json:"interval"`
		Node      string          `json:"node"`
		Arguments json.RawMessage `json:"arguments"`
	}

	// CheckResult describes the result of one or more checks.
	CheckResult struct {
		ID      int         `json:"id"`
		CheckID string      `json:"check_id"`
		Error   string      `json:"error"`
		First   time.Time   `json:"first"`
		Last    time.Time   `json:"last"`
		Streak  int         `json:"streak"`
		Results interface{} `json:"results"`
	}
)

func init() {
	database.RegisterType(Check{})
}

// UnmarshalJSON implements json.Unmarshaler.
func (c *Check) UnmarshalJSON(data []byte) error {
	proxy := checkProxy{}

	err := json.Unmarshal(data, &proxy)
	if err != nil {
		return err
	}

	c.ID = proxy.ID
	c.AgentID = proxy.AgentID
	c.Interval = proxy.Interval
	c.Node = proxy.Node
	c.Arguments = proxy.Arguments

	c.Agent = agents.GetAgent(c.AgentID)
	if c.Agent == nil {
		return errors.New("Agent not found")
	}

	err = json.Unmarshal(c.Arguments, &c.Agent)
	if err != nil {
		return err
	}

	return nil
}

// Check implements agents.Agent. Will return true if the result changed.
func (c *Check) Check(result *CheckResult) bool {
	result.CheckID = c.ID

	var err error
	var errorString string

	result.Results, err = c.Agent.Check()
	if err != nil {
		errorString = err.Error()
	}

	result.Last = time.Now()
	if result.First == (time.Time{}) {
		result.First = result.Last
	}

	if errorString == result.Error {
		result.Streak++
		return false
	}

	// If we arrive here, we have witnessed a state change. Reset CheckResult
	// and save new error from check.
	result.ID = 0
	result.Streak = 1
	result.First = result.Last
	result.Error = errorString

	return true
}
<|endoftext|>


<commit_before>package authority

import (
	"context"
	"io/ioutil"
	"os"
	"time"

	"github.com/gravitational/teleport"
	"github.com/gravitational/teleport/lib/defaults"
	kubeutils "github.com/gravitational/teleport/lib/kube/utils"

	"github.com/gravitational/trace"
	"github.com/pborman/uuid"
	log "github.com/sirupsen/logrus"
	certificates "k8s.io/api/certificates/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	watch "k8s.io/apimachinery/pkg/watch"
)

// Cert is a certificate
type Cert struct {
	// Cert is a signed certificate PEM block
	Cert []byte
	// CA is a PEM block with trusted CA
	CA []byte
}

// ProcessCSR processes CSR request with local k8s certificate authority
// and returns certificate PEM signed by CA
func ProcessCSR(csrPEM []byte, caCertPath string) (*Cert, error) {
	caPEM, err := ioutil.ReadFile(caCertPath)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	clt, _, err := kubeutils.GetKubeClient(os.Getenv(teleport.EnvKubeConfig))
	if err != nil {
		return nil, trace.Wrap(err)
	}
	id := "teleport-" + uuid.New()
	requests := clt.CertificatesV1beta1().CertificateSigningRequests()
	csr, err := requests.Create(&certificates.CertificateSigningRequest{
		ObjectMeta: metav1.ObjectMeta{
			Name: id,
		},
		Spec: certificates.CertificateSigningRequestSpec{
			Request: csrPEM,
			Usages: []certificates.KeyUsage{
				certificates.UsageDigitalSignature,
				certificates.UsageKeyEncipherment,
				certificates.UsageServerAuth,
			},
		},
	})
	if err != nil {
		return nil, trace.Wrap(err)
	}

	// Delete CSR as it seems to be hanging forever if not deleted manually.
	defer func() {
		if err := requests.Delete(id, &metav1.DeleteOptions{}); err != nil {
			log.Warningf("Failed to delete CSR: %v.", err)
		}
	}()
	csr.Status.Conditions = append(csr.Status.Conditions, certificates.CertificateSigningRequestCondition{
		Type:           certificates.CertificateApproved,
		Reason:         "TeleportApprove",
		Message:        "This CSR was approved by Teleport.",
		LastUpdateTime: metav1.Now(),
	})
	result, err := requests.UpdateApproval(csr)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	if result.Status.Certificate != nil {
		log.Debugf("Received certificate right after approval, returning.")
		return &Cert{Cert: result.Status.Certificate, CA: caPEM}, nil
	}

	ctx, cancel := context.WithTimeout(context.TODO(), defaults.CSRSignTimeout)
	defer cancel()

	watchForCert := func() ([]byte, error) {
		watcher, err := requests.Watch(metav1.ListOptions{
			TypeMeta: metav1.TypeMeta{
				Kind: teleport.KubeKindCSR,
			},
			FieldSelector: fields.Set{teleport.KubeMetadataNameSelector: id}.String(),
			Watch:         true,
		})
		if err != nil {
			return nil, trace.Wrap(err)
		}
		cert, err := waitForCSR(ctx, watcher.ResultChan())
		watcher.Stop()
		if err == nil {
			return cert, nil
		}
		return nil, trace.Wrap(err)
	}
	// this could be infinite loop, but is limited to certain number
	// of iterations just in case.
	for i := 0; i < int(defaults.CSRSignTimeout/time.Second); i++ {
		cert, err := watchForCert()
		if err == nil {
			return &Cert{Cert: cert, CA: caPEM}, nil
		}
		if !trace.IsRetryError(err) {
			return nil, trace.Wrap(err)
		}
		select {
		case <-time.After(time.Second):
			log.Debugf("Retry after network error: %v.", err)
		case <-ctx.Done():
			return nil, trace.BadParameter(timeoutCSRMessage)
		}
	}
	return nil, trace.BadParameter(timeoutCSRMessage)
}

const timeoutCSRMessage = "timeout while waiting for Kubernetes certificate"

func waitForCSR(ctx context.Context, eventsC <-chan watch.Event) ([]byte, error) {
	for {
		select {
		case event, ok := <-eventsC:
			if !ok {
				return nil, trace.Retry(nil, "events channel closed")
			}
			csr, ok := event.Object.(*certificates.CertificateSigningRequest)
			if !ok {
				log.Warnf("Unexpected resource type: %T, expected %T.", event.Object, &certificates.CertificateSigningRequest{})
				continue
			}
			if csr.Status.Certificate != nil {
				return csr.Status.Certificate, nil
			}
			log.Debugf("CSR got updated, but certificate is not ready yet: %v.", csr.Status.Conditions)
		case <-ctx.Done():
			return nil, trace.BadParameter(timeoutCSRMessage)
		}
	}
}
<commit_msg>Fix 401 Unauthorized errors for proxied K8S API requests<commit_after>package authority

import (
	"context"
	"io/ioutil"
	"os"
	"time"

	"github.com/gravitational/teleport"
	"github.com/gravitational/teleport/lib/defaults"
	kubeutils "github.com/gravitational/teleport/lib/kube/utils"

	"github.com/gravitational/trace"
	"github.com/pborman/uuid"
	log "github.com/sirupsen/logrus"
	certificates "k8s.io/api/certificates/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	watch "k8s.io/apimachinery/pkg/watch"
)

// Cert is a certificate
type Cert struct {
	// Cert is a signed certificate PEM block
	Cert []byte
	// CA is a PEM block with trusted CA
	CA []byte
}

// ProcessCSR processes CSR request with local k8s certificate authority
// and returns certificate PEM signed by CA
func ProcessCSR(csrPEM []byte, caCertPath string) (*Cert, error) {
	caPEM, err := ioutil.ReadFile(caCertPath)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	clt, _, err := kubeutils.GetKubeClient(os.Getenv(teleport.EnvKubeConfig))
	if err != nil {
		return nil, trace.Wrap(err)
	}
	id := "teleport-" + uuid.New()
	requests := clt.CertificatesV1beta1().CertificateSigningRequests()
	csr, err := requests.Create(&certificates.CertificateSigningRequest{
		ObjectMeta: metav1.ObjectMeta{
			Name: id,
		},
		Spec: certificates.CertificateSigningRequestSpec{
			Request: csrPEM,
			Usages: []certificates.KeyUsage{
				certificates.UsageDigitalSignature,
				certificates.UsageKeyEncipherment,
				certificates.UsageClientAuth,
			},
		},
	})
	if err != nil {
		return nil, trace.Wrap(err)
	}

	// Delete CSR as it seems to be hanging forever if not deleted manually.
	defer func() {
		if err := requests.Delete(id, &metav1.DeleteOptions{}); err != nil {
			log.Warningf("Failed to delete CSR: %v.", err)
		}
	}()
	csr.Status.Conditions = append(csr.Status.Conditions, certificates.CertificateSigningRequestCondition{
		Type:           certificates.CertificateApproved,
		Reason:         "TeleportApprove",
		Message:        "This CSR was approved by Teleport.",
		LastUpdateTime: metav1.Now(),
	})
	result, err := requests.UpdateApproval(csr)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	if result.Status.Certificate != nil {
		log.Debugf("Received certificate right after approval, returning.")
		return &Cert{Cert: result.Status.Certificate, CA: caPEM}, nil
	}

	ctx, cancel := context.WithTimeout(context.TODO(), defaults.CSRSignTimeout)
	defer cancel()

	watchForCert := func() ([]byte, error) {
		watcher, err := requests.Watch(metav1.ListOptions{
			TypeMeta: metav1.TypeMeta{
				Kind: teleport.KubeKindCSR,
			},
			FieldSelector: fields.Set{teleport.KubeMetadataNameSelector: id}.String(),
			Watch:         true,
		})
		if err != nil {
			return nil, trace.Wrap(err)
		}
		cert, err := waitForCSR(ctx, watcher.ResultChan())
		watcher.Stop()
		if err == nil {
			return cert, nil
		}
		return nil, trace.Wrap(err)
	}
	// this could be infinite loop, but is limited to certain number
	// of iterations just in case.
	for i := 0; i < int(defaults.CSRSignTimeout/time.Second); i++ {
		cert, err := watchForCert()
		if err == nil {
			return &Cert{Cert: cert, CA: caPEM}, nil
		}
		if !trace.IsRetryError(err) {
			return nil, trace.Wrap(err)
		}
		select {
		case <-time.After(time.Second):
			log.Debugf("Retry after network error: %v.", err)
		case <-ctx.Done():
			return nil, trace.BadParameter(timeoutCSRMessage)
		}
	}
	return nil, trace.BadParameter(timeoutCSRMessage)
}

const timeoutCSRMessage = "timeout while waiting for Kubernetes certificate"

func waitForCSR(ctx context.Context, eventsC <-chan watch.Event) ([]byte, error) {
	for {
		select {
		case event, ok := <-eventsC:
			if !ok {
				return nil, trace.Retry(nil, "events channel closed")
			}
			csr, ok := event.Object.(*certificates.CertificateSigningRequest)
			if !ok {
				log.Warnf("Unexpected resource type: %T, expected %T.", event.Object, &certificates.CertificateSigningRequest{})
				continue
			}
			if csr.Status.Certificate != nil {
				return csr.Status.Certificate, nil
			}
			log.Debugf("CSR got updated, but certificate is not ready yet: %v.", csr.Status.Conditions)
		case <-ctx.Done():
			return nil, trace.BadParameter(timeoutCSRMessage)
		}
	}
}
<|endoftext|>


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

package interceptor

import (
	"context"
	"errors"
	"fmt"
	"path/filepath"
	"time"

	"github.com/apigee/registry/log"
	"github.com/google/uuid"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type (
	resourceOperation   interface{ GetName() string }
	collectionOperation interface{ GetParent() string }
)

// CallLogger returns a gRPC server interceptor for logging API operations.
func CallLogger(opts ...log.Option) grpc.UnaryServerInterceptor {
	// Create a logger scoped to this interceptor, configured with the provided options.
	// Each request will share this logger as a base template.
	sharedLogger := log.NewLogger(opts...)
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		reqInfo := map[string]interface{}{
			"request_id": fmt.Sprintf("%.8s", uuid.New()),
			"method":     filepath.Base(info.FullMethod),
		}

		if r, ok := req.(resourceOperation); ok {
			reqInfo["resource"] = r.GetName()
		}

		if r, ok := req.(collectionOperation); ok {
			reqInfo["collection"] = r.GetParent()
		}

		// Bind request-scoped and inbound attributes to the context logger before handling the request.
		logger := log.WithInboundFields(ctx, sharedLogger).WithFields(reqInfo)
		ctx = log.NewContext(ctx, logger)

		logger.Info("Handling request.")
		start := time.Now()
		resp, err := handler(ctx, req)

		respInfo := map[string]interface{}{
			"duration":    time.Since(start),
			"status_code": status.Code(err),
		}

		if r, ok := resp.(resourceOperation); err == nil && ok {
			respInfo["resource"] = r.GetName()
		}

		// Bind response details before logging a response.
		logger = logger.WithFields(respInfo)

		// Error messages may include a status code, but we want to log messages and codes separately.
		if err != nil {
			st, _ := status.FromError(err)
			unwrapped := errors.New(st.Message())
			logger = logger.WithError(unwrapped)
		}

		switch status.Code(err) {
		case codes.OK:
			logger.Info("Success.")
		case codes.Internal:
			logger.Error("Internal error.")
		case codes.Unknown:
			logger.Error("Unknown error.")
		case codes.NotFound:
			logger.Info("Not found.")
		case codes.Unavailable:
			logger.Info("Unavailable.")
		default:
			logger.Info("User error.")
		}

		return resp, err
	}
}
<commit_msg>Update logging interceptor to print all gRPC status codes. (#713)<commit_after>// Copyright 2021 Google LLC. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package interceptor

import (
	"context"
	"errors"
	"fmt"
	"path/filepath"
	"time"

	"github.com/apigee/registry/log"
	"github.com/google/uuid"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type (
	resourceOperation   interface{ GetName() string }
	collectionOperation interface{ GetParent() string }
)

// CallLogger returns a gRPC server interceptor for logging API operations.
func CallLogger(opts ...log.Option) grpc.UnaryServerInterceptor {
	// Create a logger scoped to this interceptor, configured with the provided options.
	// Each request will share this logger as a base template.
	sharedLogger := log.NewLogger(opts...)
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		reqInfo := map[string]interface{}{
			"request_id": fmt.Sprintf("%.8s", uuid.New()),
			"method":     filepath.Base(info.FullMethod),
		}

		if r, ok := req.(resourceOperation); ok {
			reqInfo["resource"] = r.GetName()
		}

		if r, ok := req.(collectionOperation); ok {
			reqInfo["collection"] = r.GetParent()
		}

		// Bind request-scoped and inbound attributes to the context logger before handling the request.
		logger := log.WithInboundFields(ctx, sharedLogger).WithFields(reqInfo)
		ctx = log.NewContext(ctx, logger)

		logger.Info("Handling request.")
		start := time.Now()
		resp, err := handler(ctx, req)

		respInfo := map[string]interface{}{
			"duration":    time.Since(start),
			"status_code": status.Code(err),
		}

		if r, ok := resp.(resourceOperation); err == nil && ok {
			respInfo["resource"] = r.GetName()
		}

		// Bind response details before logging a response.
		logger = logger.WithFields(respInfo)

		// Error messages may include a status code, but we want to log messages and codes separately.
		if err != nil {
			st, _ := status.FromError(err)
			unwrapped := errors.New(st.Message())
			logger = logger.WithError(unwrapped)
		}

		switch status.Code(err) {
		case codes.OK:
			logger.Info("Success.")
		case codes.Canceled:
			logger.Info("Canceled.")
		case codes.Unknown:
			logger.Error("Unknown error.")
		case codes.InvalidArgument:
			logger.Error("Invalid argument.")
		case codes.DeadlineExceeded:
			logger.Error("Deadline exceeded.")
		case codes.NotFound:
			logger.Info("Not found.")
		case codes.AlreadyExists:
			logger.Error("Already exists.")
		case codes.PermissionDenied:
			logger.Error("Permission denied.")
		case codes.ResourceExhausted:
			logger.Error("Resource exhausted.")
		case codes.FailedPrecondition:
			logger.Error("Failed precondition.")
		case codes.Aborted:
			logger.Error("Aborted.")
		case codes.OutOfRange:
			logger.Error("Out of range.")
		case codes.Unimplemented:
			logger.Error("Unimplemented.")
		case codes.Internal:
			logger.Error("Internal error.")
		case codes.Unavailable:
			logger.Info("Unavailable.")
		case codes.DataLoss:
			logger.Info("Data loss.")
		case codes.Unauthenticated:
			logger.Info("Unauthenticated.")
		default:
			logger.Info("User error.")
		}

		return resp, err
	}
}
<|endoftext|>


<commit_before>package test

import (
	"io/ioutil"
	"os"
	"strconv"
	"testing"
	"time"

	"github.com/coreos/go-etcd/etcd"
)

// This test creates a single node and then set a value to it to trigger snapshot
func TestSimpleSnapshot(t *testing.T) {
	procAttr := new(os.ProcAttr)
	procAttr.Files = []*os.File{nil, os.Stdout, os.Stderr}
	args := []string{"etcd", "-n=node1", "-d=/tmp/node1", "-snapshot=true", "-snapshotCount=500"}

	process, err := os.StartProcess(EtcdBinPath, append(args, "-f"), procAttr)
	if err != nil {
		t.Fatal("start process failed:" + err.Error())
	}
	defer process.Kill()

	time.Sleep(time.Second)

	c := etcd.NewClient(nil)

	c.SyncCluster()
	// issue first 501 commands
	for i := 0; i < 501; i++ {
		result, err := c.Set("foo", "bar", 100)

		if err != nil || result.Key != "/foo" || result.Value != "bar" || result.TTL < 95 {
			if err != nil {
				t.Fatal(err)
			}

			t.Fatalf("Set failed with %s %s %v", result.Key, result.Value, result.TTL)
		}
	}

	// wait for a snapshot interval
	time.Sleep(3 * time.Second)

	snapshots, err := ioutil.ReadDir("/tmp/node1/snapshot")

	if err != nil {
		t.Fatal("list snapshot failed:" + err.Error())
	}

	if len(snapshots) != 1 {
		t.Fatal("wrong number of snapshot :[1/", len(snapshots), "]")
	}

	index, _ := strconv.Atoi(snapshots[0].Name()[2:5])

	if index < 507 || index > 510 {
		t.Fatal("wrong name of snapshot :", snapshots[0].Name())
	}

	// issue second 501 commands
	for i := 0; i < 501; i++ {
		result, err := c.Set("foo", "bar", 100)

		if err != nil || result.Key != "/foo" || result.Value != "bar" || result.TTL < 95 {
			if err != nil {
				t.Fatal(err)
			}

			t.Fatalf("Set failed with %s %s %v", result.Key, result.Value, result.TTL)
		}
	}

	// wait for a snapshot interval
	time.Sleep(3 * time.Second)

	snapshots, err = ioutil.ReadDir("/tmp/node1/snapshot")

	if err != nil {
		t.Fatal("list snapshot failed:" + err.Error())
	}

	if len(snapshots) != 1 {
		t.Fatal("wrong number of snapshot :[1/", len(snapshots), "]")
	}

	index, _ = strconv.Atoi(snapshots[0].Name()[2:6])

	if index < 1015 || index > 1018 {
		t.Fatal("wrong name of snapshot :", snapshots[0].Name())
	}
}
<commit_msg>test fix wrong assumption should be previous index + 501 + 6<commit_after>package test

import (
	"io/ioutil"
	"os"
	"strconv"
	"testing"
	"time"

	"github.com/coreos/go-etcd/etcd"
)

// This test creates a single node and then set a value to it to trigger snapshot
func TestSimpleSnapshot(t *testing.T) {
	procAttr := new(os.ProcAttr)
	procAttr.Files = []*os.File{nil, os.Stdout, os.Stderr}
	args := []string{"etcd", "-n=node1", "-d=/tmp/node1", "-snapshot=true", "-snapshotCount=500"}

	process, err := os.StartProcess(EtcdBinPath, append(args, "-f"), procAttr)
	if err != nil {
		t.Fatal("start process failed:" + err.Error())
	}
	defer process.Kill()

	time.Sleep(time.Second)

	c := etcd.NewClient(nil)

	c.SyncCluster()
	// issue first 501 commands
	for i := 0; i < 501; i++ {
		result, err := c.Set("foo", "bar", 100)

		if err != nil || result.Key != "/foo" || result.Value != "bar" || result.TTL < 95 {
			if err != nil {
				t.Fatal(err)
			}

			t.Fatalf("Set failed with %s %s %v", result.Key, result.Value, result.TTL)
		}
	}

	// wait for a snapshot interval
	time.Sleep(3 * time.Second)

	snapshots, err := ioutil.ReadDir("/tmp/node1/snapshot")

	if err != nil {
		t.Fatal("list snapshot failed:" + err.Error())
	}

	if len(snapshots) != 1 {
		t.Fatal("wrong number of snapshot :[1/", len(snapshots), "]")
	}

	index, _ := strconv.Atoi(snapshots[0].Name()[2:5])

	if index <= 507 || index >= 510 {
		t.Fatal("wrong name of snapshot :", snapshots[0].Name())
	}

	// issue second 501 commands
	for i := 0; i < 501; i++ {
		result, err := c.Set("foo", "bar", 100)

		if err != nil || result.Key != "/foo" || result.Value != "bar" || result.TTL < 95 {
			if err != nil {
				t.Fatal(err)
			}

			t.Fatalf("Set failed with %s %s %v", result.Key, result.Value, result.TTL)
		}
	}

	// wait for a snapshot interval
	time.Sleep(3 * time.Second)

	snapshots, err = ioutil.ReadDir("/tmp/node1/snapshot")

	if err != nil {
		t.Fatal("list snapshot failed:" + err.Error())
	}

	if len(snapshots) != 1 {
		t.Fatal("wrong number of snapshot :[1/", len(snapshots), "]")
	}

	index, _ = strconv.Atoi(snapshots[0].Name()[2:6])

	if index <= 1014 || index >= 1017 {
		t.Fatal("wrong name of snapshot :", snapshots[0].Name())
	}
}
<|endoftext|>


<commit_before>package scheduler

import (
	"sync"

	"github.com/chrislusf/glow/driver/plan"
	"github.com/chrislusf/glow/driver/scheduler/market"
	"github.com/chrislusf/glow/flow"
	"github.com/chrislusf/glow/resource"
)

type SubmitTaskGroup struct {
	FlowContext *flow.FlowContext
	TaskGroup   *plan.TaskGroup
	Bid         float64
	WaitGroup   *sync.WaitGroup
}

type ReleaseTaskGroupInputs struct {
	FlowContext *flow.FlowContext
	TaskGroups  []*plan.TaskGroup
	WaitGroup   *sync.WaitGroup
}

/*
resources are leased to driver, expires every X miniute unless renewed.
1. request resource
2. release resource
*/
func (s *Scheduler) EventLoop() {
	for {
		event := <-s.EventChan
		switch event := event.(type) {
		default:
		case SubmitTaskGroup:
			// fmt.Printf("processing %+v\n", event)
			taskGroup := event.TaskGroup
			pickedServerChan := make(chan market.Supply, 1)
			go func() {
				defer event.WaitGroup.Done()
				tasks := event.TaskGroup.Tasks

				// wait until inputs are registed
				s.shardLocator.waitForInputDatasetShardLocations(tasks[0])
				// fmt.Printf("inputs of %s is %s\n", tasks[0].Name(), s.allInputLocations(tasks[0]))

				s.Market.AddDemand(market.Requirement(taskGroup), event.Bid, pickedServerChan)

				// get assigned executor location
				supply := <-pickedServerChan
				allocation := supply.Object.(resource.Allocation)
				defer s.Market.ReturnSupply(supply)

				s.remoteExecuteOnLocation(event.FlowContext, taskGroup, allocation, event.WaitGroup)
			}()
		case ReleaseTaskGroupInputs:
			go func() {
				defer event.WaitGroup.Done()

				for _, taskGroup := range event.TaskGroups {
					tasks := taskGroup.Tasks
					for _, ds := range tasks[len(tasks)-1].Outputs {
						shardName := s.Option.ExecutableFileHash + "-" + ds.Name()
						location, _ := s.shardLocator.GetShardLocation(shardName)
						request := NewDeleteDatasetShardRequest(shardName)
						// println("deleting", ds.Name(), "on", location.URL())
						if err := RemoteDirectExecute(s.Option.TlsConfig, location.URL(), request); err != nil {
							println("Purging dataset error:", err.Error())
						}
					}
				}

			}()
		}
	}
}
<commit_msg>Move pickedServerChan definition before where it is used.<commit_after>package scheduler

import (
	"sync"

	"github.com/chrislusf/glow/driver/plan"
	"github.com/chrislusf/glow/driver/scheduler/market"
	"github.com/chrislusf/glow/flow"
	"github.com/chrislusf/glow/resource"
)

type SubmitTaskGroup struct {
	FlowContext *flow.FlowContext
	TaskGroup   *plan.TaskGroup
	Bid         float64
	WaitGroup   *sync.WaitGroup
}

type ReleaseTaskGroupInputs struct {
	FlowContext *flow.FlowContext
	TaskGroups  []*plan.TaskGroup
	WaitGroup   *sync.WaitGroup
}

/*
resources are leased to driver, expires every X miniute unless renewed.
1. request resource
2. release resource
*/
func (s *Scheduler) EventLoop() {
	for {
		event := <-s.EventChan
		switch event := event.(type) {
		default:
		case SubmitTaskGroup:
			// fmt.Printf("processing %+v\n", event)
			taskGroup := event.TaskGroup
			go func() {
				defer event.WaitGroup.Done()
				tasks := event.TaskGroup.Tasks

				// wait until inputs are registed
				s.shardLocator.waitForInputDatasetShardLocations(tasks[0])
				// fmt.Printf("inputs of %s is %s\n", tasks[0].Name(), s.allInputLocations(tasks[0]))

				pickedServerChan := make(chan market.Supply, 1)
				s.Market.AddDemand(market.Requirement(taskGroup), event.Bid, pickedServerChan)

				// get assigned executor location
				supply := <-pickedServerChan
				allocation := supply.Object.(resource.Allocation)
				defer s.Market.ReturnSupply(supply)

				s.remoteExecuteOnLocation(event.FlowContext, taskGroup, allocation, event.WaitGroup)
			}()
		case ReleaseTaskGroupInputs:
			go func() {
				defer event.WaitGroup.Done()

				for _, taskGroup := range event.TaskGroups {
					tasks := taskGroup.Tasks
					for _, ds := range tasks[len(tasks)-1].Outputs {
						shardName := s.Option.ExecutableFileHash + "-" + ds.Name()
						location, _ := s.shardLocator.GetShardLocation(shardName)
						request := NewDeleteDatasetShardRequest(shardName)
						// println("deleting", ds.Name(), "on", location.URL())
						if err := RemoteDirectExecute(s.Option.TlsConfig, location.URL(), request); err != nil {
							println("Purging dataset error:", err.Error())
						}
					}
				}

			}()
		}
	}
}
<|endoftext|>